diff --git a/tests/pytorch/distributed/run_ep.py b/tests/pytorch/distributed/run_ep.py index 6ab5f901f7..ff62249c5d 100644 --- a/tests/pytorch/distributed/run_ep.py +++ b/tests/pytorch/distributed/run_ep.py @@ -16,7 +16,7 @@ from ep_reference import BlockScaledTensor, MoeEpReference, MoeFormat, quantize_blockwise import transformer_engine.pytorch as te from transformer_engine.pytorch import ops as te_ops -from transformer_engine.common.recipe import MXFP8BlockScaling +from transformer_engine.common.recipe import MXFP8BlockScaling, NVFP4BlockScaling from transformer_engine.pytorch.ep import ( EpBuffer, EpConfig, @@ -732,6 +732,437 @@ def test_combine_bwd_mxfp8(self): _degroup_mxfp8(g_mx).float(), src_bf.grad.float()[:n], atol=5e-2, rtol=5e-2 ) + # NVFP4 dispatch / combine backward + + def _require_nvfp4_shapes(self): + if HIDDEN_DIM % 512 != 0: + self.skipTest("NVFP4 EP needs HIDDEN_DIM % 512 == 0 (set NVTE_EP_HIDDEN_DIM)") + if torch.cuda.get_device_capability() < (10, 0): + self.skipTest("NVFP4 EP tests require Blackwell (SM 10.0) or newer") + + @staticmethod + def _nvfp4_decoded(tokens): + """Tokens as the receiver decodes them: RN pack then decode.""" + return tex.ep_nvfp4_unpack(tex.ep_nvfp4_pack(tokens, False), HIDDEN_DIM) + + class _RecordPack: + """Records the packed rows of every stochastic NVFP4 pack the EP module runs.""" + + def __enter__(self): + self.orig, self.packed = tex.ep_nvfp4_pack, [] + + def record(tokens, stochastic): + out = self.orig(tokens, stochastic) + if stochastic: + self.packed.append(out) + return out + + tex.ep_nvfp4_pack = record + return self + + def __exit__(self, *exc): + tex.ep_nvfp4_pack = self.orig + + def _assert_rows_match_per_expert(self, got, ref, counts, max_mismatch=0.0): + """Row-set equality inside each expert block (row order can differ between dispatch + runs). Rows are keyed by a fixed random projection: FP4-decoded rows can share a row + sum, which would make a sum-sorted order ambiguous. ``max_mismatch`` is the largest + fraction of differing elements allowed per expert.""" + cum = [0] + counts.cumsum(0).tolist() + generator = torch.Generator(device=got.device) + generator.manual_seed(7) + proj = torch.randn(HIDDEN_DIM, generator=generator, device=got.device, dtype=torch.float64) + worst = 0.0 + for lo, hi in zip(cum[:-1], cum[1:]): + a, b = got[lo:hi].double(), ref[lo:hi].double() + a = a[torch.argsort(a @ proj)] + b = b[torch.argsort(b @ proj)] + worst = max(worst, (a != b).float().mean().item() if hi > lo else 0.0) + if worst > max_mismatch: + print(f"[rank {self.cfg.rank}] NVFP4 rows differ in {worst:.5f}", flush=True) + self.assertLessEqual(worst, max_mismatch) + + def _nvfp4_dispatch_ref(self, tokens, topk_idx, w, alignment=0): + """BF16 dispatch of the receiver-decoded tokens: (recv, counts).""" + ref_recv, _rw, ref_tc = ep_dispatch( + self._make_buffer(alignment=alignment), self._nvfp4_decoded(tokens), topk_idx, w + ) + torch.cuda.synchronize() + return ref_recv, ref_tc + + @staticmethod + def _decode_nvfp4_rowwise(grouped, counts, hidden): + """Decode the rowwise part of a swizzled NVFP4 GroupedTensor to FP32 [rows, hidden].""" + e2m1 = torch.tensor([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0], device=counts.device) + rows = int(counts.sum().item()) + data = grouped.rowwise_data.view(torch.uint8)[: rows * hidden // 2].view(rows, hidden // 2) + codes = torch.stack([data & 15, data >> 4], dim=-1).reshape(rows, hidden).long() + vals = e2m1[codes & 7] * torch.where((codes & 8) != 0, -1.0, 1.0) + ks = hidden // 16 + r = torch.arange(rows, device=counts.device)[:, None] + c = torch.arange(ks, device=counts.device)[None, :] + idx = ((r // 128) * ((ks + 3) // 4) + c // 4) * 512 + (r % 32) * 16 + ((r // 32) % 4) * 4 + idx = idx + c % 4 + scales = grouped.scale_inv.view(torch.uint8)[idx].view(torch.float8_e4m3fn).float() + group = torch.repeat_interleave(torch.arange(counts.numel(), device=counts.device), counts) + dec = grouped.amax.view(-1)[group].float() / (6.0 * 448.0) + return vals * scales.repeat_interleave(16, dim=1) * dec[:, None] + + @_eager_test_include + @_zero_copy_test_include + def test_dispatch_nvfp4(self): + """NVFP4 dispatch outside NVFP4 autocast returns BF16 recv tokens equal to dispatching + the sender-decoded tokens; recv weights are not quantized.""" + self._require_nvfp4_shapes() + topk_idx, tokens, w = _make_moe_inputs(self.cfg.rank, self.cfg.ep_size) + buf = self._make_buffer(dispatch_fwd_quant_recipe=NVFP4BlockScaling()) + recv, rw, tc = ep_dispatch(buf, tokens, topk_idx, w) + torch.cuda.synchronize() + self.assertEqual(recv.dtype, torch.bfloat16) + self.assertEqual(recv.shape[-1], HIDDEN_DIM) + n = int(tc.sum().item()) + ref_recv, ref_tc = self._nvfp4_dispatch_ref(tokens, topk_idx, w) + self.assertEqual(int(ref_tc.sum().item()), n) + self._assert_rows_match_per_expert(recv, ref_recv, ref_tc) + _ref, ref_rw, _ = ep_dispatch(self._make_buffer(), tokens, topk_idx, w) + torch.cuda.synchronize() + torch.testing.assert_close(rw[:n].sort().values, ref_rw[:n].sort().values, atol=0, rtol=0) + + @_eager_test_include + def test_dispatch_nvfp4_autograd(self): + """The NVFP4 dispatch backward is a straight-through BF16 gradient: seeding the recv grad + with ones scatters TOP_K back to each token under identity routing.""" + self._require_nvfp4_shapes() + topk_idx, tokens, w = _make_identity_inputs(self.cfg.rank, self.cfg.ep_size) + tokens_p = tokens.detach().clone().requires_grad_(True) + buf = self._make_buffer(dispatch_fwd_quant_recipe=NVFP4BlockScaling()) + recv, _rw, _tc = ep_dispatch(buf, tokens_p, topk_idx, w) + torch.autograd.backward(recv, grad_tensors=torch.ones_like(recv)) + torch.cuda.synchronize() + torch.testing.assert_close( + tokens_p.grad.float(), torch.full_like(tokens_p, float(TOP_K)).float(), atol=0, rtol=0 + ) + + @_zero_copy_test_include + def test_dispatch_nvfp4_capture(self): + """NVFP4 dispatch captured in a CUDA graph with caller recv buffers: a replay with new + tokens decodes the new tokens into the caller buffer, bounded by the replay's counts. + Under zero-copy the caller buffers are symm-mem and the packed rows land in the buffer's + persistent symm-mem scratch, so capture allocates nothing from the symm-mem pool.""" + if EAGER: + self.skipTest("capture requires non-eager static recv capacity") + self._require_nvfp4_shapes() + recv_buf, rw_buf = self._caller_recv_symm() if ZERO_COPY else self._caller_recv() + buf = self._make_buffer(dispatch_fwd_quant_recipe=NVFP4BlockScaling()) + topk_idx, tokens, w = _make_moe_inputs(self.cfg.rank, self.cfg.ep_size) + static_tokens = tokens.clone() + out = {} + + def step(): + out["recv"], out["rw"], _tc = ep_dispatch( + buf, static_tokens, topk_idx, w, recv_tokens=recv_buf, recv_topk_weights=rw_buf + ) + + graph = self._capture(step) + self._assert_nvfp4_scratch(buf) + new_tokens = (tokens.float() * -3.0 + 0.125).to(torch.bfloat16) + static_tokens.copy_(new_tokens) + recv_buf.zero_() + graph.replay() + torch.cuda.synchronize() + self.assertEqual(out["recv"].data_ptr(), recv_buf.data_ptr()) + ref_recv, ref_tc = self._nvfp4_dispatch_ref(new_tokens, topk_idx, w) + self.assertEqual(int(buf.tokens_per_expert.sum().item()), int(ref_tc.sum().item())) + self._assert_rows_match_per_expert(recv_buf, ref_recv, ref_tc) + + def _caller_recv_symm(self): + """Persistent symm-mem recv token + weight buffers for a zero-copy capture.""" + rc = self.cfg.recv_capacity_per_rank + return ( + symm_mem_alloc((rc, HIDDEN_DIM), torch.bfloat16, self.ep_group), + symm_mem_alloc((rc,), torch.float32, self.ep_group), + ) + + def _assert_nvfp4_scratch(self, buf): + """Under zero-copy the BF16-decode path keeps its packed landing as a persistent + symm-mem scratch on the buffer; otherwise it allocates per call and keeps none.""" + if ZERO_COPY: + self.assertIsNotNone(buf.nvfp4_packed_scratch) + self.assertTrue(is_symm_backed(buf.nvfp4_packed_scratch)) + else: + self.assertIsNone(buf.nvfp4_packed_scratch) + + @_eager_test_include + @_zero_copy_test_include + @_mxfp8_align_test + def test_dispatch_nvfp4_direct_receive(self): + """NVFP4 dispatch under NVFP4 autocast returns a swizzled NVFP4 GroupedTensor whose + rowwise rows decode to the sender's packed tokens (per expert, order-tolerant), with + finite per-expert scales; the backward is the straight-through BF16 gradient. Outside + autocast, and in eager mode, recv is the decoded BF16 tokens.""" + self._require_nvfp4_shapes() + from transformer_engine.pytorch.tensor.grouped_tensor import GroupedTensor + + topk_idx, tokens, w = _make_moe_inputs(self.cfg.rank, self.cfg.ep_size) + tokens_p = tokens.detach().clone().requires_grad_(True) + buf = self._make_buffer(dispatch_fwd_quant_recipe=NVFP4BlockScaling(), alignment=256) + recv_bf16, _, _ = ep_dispatch(buf, tokens, topk_idx, w) + self.assertEqual(recv_bf16.dtype, torch.bfloat16) + with te.autocast(enabled=True, recipe=NVFP4BlockScaling()): + recv, _rw, tc = ep_dispatch(buf, tokens_p, topk_idx, w) + counts = tc.clone() + if EAGER: + self.assertEqual(recv.dtype, torch.bfloat16) + return + self.assertIsInstance(recv, GroupedTensor) + self.assertTrue(recv._with_gemm_swizzled_scales) + torch.cuda.synchronize() + self.assertTrue(torch.isfinite(recv.amax).all() and (recv.amax > 0).all()) + self.assertTrue(torch.isfinite(recv.columnwise_amax).all()) + + ref_recv, ref_tc = self._nvfp4_dispatch_ref(tokens, topk_idx, w, alignment=256) + torch.testing.assert_close(counts, ref_tc, atol=0, rtol=0) + got = self._decode_nvfp4_rowwise(recv, counts, HIDDEN_DIM) + # Block scales re-expressed against the expert scale can underflow for tokens far + # smaller than the expert's largest; allow a tiny fraction of such elements. + self._assert_rows_match_per_expert(got, ref_recv[: got.shape[0]], counts, max_mismatch=1e-3) + + if ZERO_COPY: + return # zero-copy dispatch_bwd needs a symm-mem grad; checked in the other passes + g_recv = torch.ones(recv.shape, dtype=torch.bfloat16, device=self.cfg.device) + torch.autograd.backward(recv, grad_tensors=g_recv) + torch.cuda.synchronize() + torch.testing.assert_close( + tokens_p.grad.float(), torch.full_like(tokens_p, float(TOP_K)).float(), atol=0, rtol=0 + ) + + def _combine_bwd_bf16_ref(self, tokens, topk_idx, w, rows, grad, alignment=0): + """BF16 combine backward of ``grad`` on the same routing: the expert_out grad.""" + buf = self._make_buffer(alignment=alignment) + ep_dispatch(buf, tokens, topk_idx, w) + src = torch.zeros(rows, HIDDEN_DIM, dtype=torch.bfloat16, device=self.cfg.device) + src.requires_grad_(True) + out = ep_combine(buf, self._expert_out(src)) + torch.autograd.backward(out, grad_tensors=grad) + torch.cuda.synchronize() + return src.grad + + def _combine_grad(self, num_tokens): + generator = torch.Generator(device=self.cfg.device) + generator.manual_seed(100 + self.cfg.rank) + g = torch.randn(num_tokens, HIDDEN_DIM, generator=generator, device=self.cfg.device) + g = g * torch.logspace(-2, 1, num_tokens, device=self.cfg.device)[:, None] + return g.to(torch.bfloat16) + + @_eager_test_include + @_zero_copy_test_include + def test_combine_bwd_nvfp4(self): + """NVFP4 combine backward outside NVFP4 autocast returns a BF16 expert_out grad equal to + a BF16 combine backward of the sender-decoded (stochastically rounded) result-grad.""" + self._require_nvfp4_shapes() + topk_idx, tokens, w = _make_moe_inputs(self.cfg.rank, self.cfg.ep_size) + buf = self._make_buffer(combine_bwd_quant_recipe=NVFP4BlockScaling()) + recv, _rw, tc = ep_dispatch(buf, tokens, topk_idx, w) + rows = recv.shape[0] + src = torch.zeros(rows, HIDDEN_DIM, dtype=torch.bfloat16, device=self.cfg.device) + src.requires_grad_(True) + out = ep_combine(buf, self._expert_out(src)) + grad = self._combine_grad(out.shape[0]) + with self._RecordPack() as rec: + torch.autograd.backward(out, grad_tensors=grad) + torch.cuda.synchronize() + self.assertEqual(len(rec.packed), 1) + self.assertEqual(src.grad.dtype, torch.bfloat16) + self.assertEqual(tuple(src.grad.shape), (rows, HIDDEN_DIM)) + decoded = tex.ep_nvfp4_unpack(rec.packed[0], HIDDEN_DIM) + n = int(tc.sum().item()) + ref = self._combine_bwd_bf16_ref(tokens, topk_idx, w, rows, decoded) + self._assert_rows_match_per_expert(src.grad[:n], ref[:n], tc) + # The only quantization is on the sender: the SR-decoded grad stays close to the grad. + rel = (decoded.float() - grad.float()).norm(dim=1) / grad.float().norm(dim=1) + self.assertLess(rel.max().item(), 0.25) + + def test_combine_bwd_nvfp4_caller_grad_out(self): + """NVFP4 combine backward decodes into the caller grad_out and leaves rows past the live + recv count untouched.""" + if EAGER: + self.skipTest("eager mode forbids a caller grad_out") + self._require_nvfp4_shapes() + rc = self.cfg.recv_capacity_per_rank + topk_idx, tokens, w = _make_moe_inputs(self.cfg.rank, self.cfg.ep_size) + buf = self._make_buffer(combine_bwd_quant_recipe=NVFP4BlockScaling()) + _recv, _rw, tc = ep_dispatch(buf, tokens, topk_idx, w) + gbuf = torch.full((rc, HIDDEN_DIM), 7.0, dtype=torch.bfloat16, device=self.cfg.device) + src = torch.zeros(rc, HIDDEN_DIM, dtype=torch.bfloat16, device=self.cfg.device) + src.requires_grad_(True) + out = ep_combine(buf, self._expert_out(src), grad_out=gbuf) + grad = self._combine_grad(out.shape[0]) + with self._RecordPack() as rec: + torch.autograd.backward(out, grad_tensors=grad) + torch.cuda.synchronize() + self.assertEqual(src.grad.data_ptr(), gbuf.data_ptr()) + n = int(tc.sum().item()) + self.assertTrue(bool((gbuf[n:] == 7.0).all())) + decoded = tex.ep_nvfp4_unpack(rec.packed[0], HIDDEN_DIM) + ref = self._combine_bwd_bf16_ref(tokens, topk_idx, w, rc, decoded) + self._assert_rows_match_per_expert(gbuf[:n], ref[:n], tc) + + @_zero_copy_test_include + def test_combine_bwd_nvfp4_capture(self): + """NVFP4 combine fwd+bwd captured in a CUDA graph: each replay decodes the replay's grad + exactly into the caller grad_out, and draws fresh stochastic-rounding bits. Under + zero-copy the combine input and grad_out are symm-mem and the packed rows land in the + buffer's persistent symm-mem scratch, so capture allocates nothing from the pool.""" + if EAGER: + self.skipTest("capture requires non-eager static recv capacity") + self._require_nvfp4_shapes() + rc = self.cfg.recv_capacity_per_rank + topk_idx, tokens, w = _make_moe_inputs(self.cfg.rank, self.cfg.ep_size) + buf = self._make_buffer(combine_bwd_quant_recipe=NVFP4BlockScaling()) + _recv, _rw, tc = ep_dispatch(buf, tokens, topk_idx, w) # fixes the routing + torch.cuda.synchronize() + if ZERO_COPY: + gbuf = symm_mem_alloc((rc, HIDDEN_DIM), torch.bfloat16, self.ep_group) + src = symm_mem_alloc((rc, HIDDEN_DIM), torch.bfloat16, self.ep_group) + gbuf.zero_() + src.zero_() + else: + gbuf = torch.zeros(rc, HIDDEN_DIM, dtype=torch.bfloat16, device=self.cfg.device) + src = torch.zeros(rc, HIDDEN_DIM, dtype=torch.bfloat16, device=self.cfg.device) + src.requires_grad_(True) + static_grad = self._combine_grad(buf.max_tokens_per_rank) + + def step(): + out = ep_combine(buf, src, grad_out=gbuf) + torch.autograd.grad(out, src, grad_outputs=static_grad) + + with self._RecordPack() as rec: + graph = self._capture(step) + self._assert_nvfp4_scratch(buf) + static_grad.copy_((static_grad.float() * -2.0 + 0.0625).to(torch.bfloat16)) + gbuf.zero_() + graph.replay() + torch.cuda.synchronize() + first_packed = rec.packed[-1].clone() + first = gbuf.clone() + graph.replay() + torch.cuda.synchronize() + second_packed = rec.packed[-1].clone() + n = int(tc.sum().item()) + self.assertFalse(torch.equal(first_packed, second_packed)) + for packed, got in ((first_packed, first), (second_packed, gbuf)): + decoded = tex.ep_nvfp4_unpack(packed, HIDDEN_DIM) + ref = self._combine_bwd_bf16_ref(tokens, topk_idx, w, rc, decoded) + self._assert_rows_match_per_expert(got[:n], ref[:n], tc) + + def _assert_direct_grad_matches(self, g, packed, tokens, topk_idx, w, counts, rows): + """Rowwise rows of the direct-receive grad decode to a BF16 combine backward of the + sender's packed (stochastically rounded) grad, per expert and order-tolerant; per-expert + scales are finite.""" + from transformer_engine.pytorch.tensor.grouped_tensor import GroupedTensor + + self.assertIsInstance(g, GroupedTensor) + self.assertTrue(g._with_gemm_swizzled_scales) + self.assertTrue(g.quantizer.stochastic_rounding) + self.assertTrue(torch.isfinite(g.amax).all() and torch.isfinite(g.columnwise_amax).all()) + decoded = tex.ep_nvfp4_unpack(packed, HIDDEN_DIM) + ref = self._combine_bwd_bf16_ref(tokens, topk_idx, w, rows, decoded, alignment=256) + got = self._decode_nvfp4_rowwise(g, counts, HIDDEN_DIM) + self._assert_rows_match_per_expert(got, ref[: got.shape[0]], counts, max_mismatch=1e-3) + + @_eager_test_include + @_zero_copy_test_include + @_mxfp8_align_test + def test_combine_bwd_nvfp4_direct(self): + """NVFP4 combine backward under NVFP4 autocast returns FC2's grad_output as a swizzled + NVFP4 GroupedTensor (stochastic rounding) matching a BF16 combine backward of the + sender-packed grad. In eager mode it decodes to BF16.""" + self._require_nvfp4_shapes() + topk_idx, tokens, w = _make_moe_inputs(self.cfg.rank, self.cfg.ep_size) + buf = self._make_buffer(combine_bwd_quant_recipe=NVFP4BlockScaling(), alignment=256) + recv, _rw, tc = ep_dispatch(buf, tokens, topk_idx, w) + rows = recv.shape[0] + counts = tc.clone() + src = torch.zeros(rows, HIDDEN_DIM, dtype=torch.bfloat16, device=self.cfg.device) + src.requires_grad_(True) + with self._RecordPack() as rec: + with te.autocast(enabled=True, recipe=NVFP4BlockScaling()): + out = ep_combine(buf, self._expert_out(src)) + grad = self._combine_grad(out.shape[0]) + torch.autograd.backward(out, grad_tensors=grad) + torch.cuda.synchronize() + self.assertEqual(len(rec.packed), 1) + if EAGER: + self.assertEqual(src.grad.dtype, torch.bfloat16) + return + self._assert_direct_grad_matches(src.grad, rec.packed[0], tokens, topk_idx, w, counts, rows) + + @_mxfp8_align_test + def test_combine_bwd_nvfp4_direct_receive_opt_out(self): + """With nvfp4_grad_direct_receive=False the NVFP4 combine backward returns the BF16 + expert_out grad even under NVFP4 autocast, for consumers that need the decoded grad + (a grouped linear with a bias): it equals a BF16 combine backward of the sender-decoded + result-grad.""" + self._require_nvfp4_shapes() + topk_idx, tokens, w = _make_moe_inputs(self.cfg.rank, self.cfg.ep_size) + buf = self._make_buffer(combine_bwd_quant_recipe=NVFP4BlockScaling(), alignment=256) + recv, _rw, tc = ep_dispatch(buf, tokens, topk_idx, w) + rows = recv.shape[0] + src = torch.zeros(rows, HIDDEN_DIM, dtype=torch.bfloat16, device=self.cfg.device) + src.requires_grad_(True) + with self._RecordPack() as rec: + with te.autocast(enabled=True, recipe=NVFP4BlockScaling()): + out = ep_combine(buf, self._expert_out(src), nvfp4_grad_direct_receive=False) + grad = self._combine_grad(out.shape[0]) + torch.autograd.backward(out, grad_tensors=grad) + torch.cuda.synchronize() + self.assertEqual(len(rec.packed), 1) + self.assertIsInstance(src.grad, torch.Tensor) + self.assertEqual(src.grad.dtype, torch.bfloat16) + self.assertEqual(tuple(src.grad.shape), (rows, HIDDEN_DIM)) + decoded = tex.ep_nvfp4_unpack(rec.packed[0], HIDDEN_DIM) + n = int(tc.sum().item()) + ref = self._combine_bwd_bf16_ref(tokens, topk_idx, w, rows, decoded, alignment=256) + self._assert_rows_match_per_expert(src.grad[:n], ref[:n], tc) + + @_mxfp8_align_test + def test_combine_bwd_nvfp4_direct_capture(self): + """Combine fwd+bwd (direct NVFP4 grad, packed rows landing in the caller grad_out) + captured in a CUDA graph: each replay matches a BF16 combine backward of that replay's + packed grad, and replays draw fresh stochastic-rounding bits.""" + if EAGER: + self.skipTest("capture requires non-eager static recv capacity") + self._require_nvfp4_shapes() + rc = self.cfg.recv_capacity_per_rank + topk_idx, tokens, w = _make_moe_inputs(self.cfg.rank, self.cfg.ep_size) + buf = self._make_buffer(combine_bwd_quant_recipe=NVFP4BlockScaling(), alignment=256) + _recv, _rw, tc = ep_dispatch(buf, tokens, topk_idx, w) + torch.cuda.synchronize() + counts = tc.clone() + gbuf = torch.zeros(rc, HIDDEN_DIM, dtype=torch.bfloat16, device=self.cfg.device) + src = torch.zeros(rc, HIDDEN_DIM, dtype=torch.bfloat16, device=self.cfg.device) + src.requires_grad_(True) + static_grad = self._combine_grad(buf.max_tokens_per_rank) + res = {} + + def step(): + with te.autocast(enabled=True, recipe=NVFP4BlockScaling()): + out = ep_combine(buf, src, grad_out=gbuf) + (res["g"],) = torch.autograd.grad(out, src, grad_outputs=static_grad) + + with self._RecordPack() as rec: + graph = self._capture(step) + static_grad.copy_((static_grad.float() * -2.0 + 0.0625).to(torch.bfloat16)) + packs = [] + for _ in range(2): + graph.replay() + torch.cuda.synchronize() + packs.append(rec.packed[-1].clone()) + self._assert_direct_grad_matches( + res["g"], packs[-1], tokens, topk_idx, w, counts, rc + ) + self.assertFalse(torch.equal(packs[0], packs[1]), "replays must draw fresh bits") + @_zero_copy_test_include def test_zero_copy_pool_auto_alloc(self): """Zero-copy with recv/grad left None: ep_dispatch/ep_combine allocate their IO diff --git a/tests/pytorch/nvfp4/test_nvfp4_ep_rows.py b/tests/pytorch/nvfp4/test_nvfp4_ep_rows.py new file mode 100644 index 0000000000..99971146f8 --- /dev/null +++ b/tests/pytorch/nvfp4/test_nvfp4_ep_rows.py @@ -0,0 +1,594 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. +"""Single-GPU tests for the NVFP4 EP token rows: the pack (RN and stochastic rounding, +power-of-two token scales), the BF16 decode, and the direct receive that builds the expert GEMM's +NVFP4 GroupedTensor from packed rows (rowwise codes reused, columnwise RHT quantized in one pass). +""" + +import os +import random + +import pytest +import torch + +import transformer_engine.pytorch as te +import transformer_engine_torch as tex +from transformer_engine.common.recipe import NVFP4BlockScaling +from transformer_engine.pytorch.ep import _nvfp4_direct_receive, _nvfp4_row_bytes +from transformer_engine.pytorch.tensor.nvfp4_tensor import ( + NVFP4Quantizer, + get_hadamard_matrix, + get_wgrad_sign_vector, +) + +recipe_available, reason_for_no_recipe = te.is_nvfp4_available(return_reason=True) +pytestmark = pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) + +E2M1 = torch.tensor([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0]) +LUT = torch.cat([E2M1, -E2M1]) +NVFP4_MAX = 6.0 * 448.0 +FAST_MATH = os.environ.get("NVTE_USE_FAST_MATH", "0") == "1" + + +def _tokens(total, hidden, seed, spread=(-2, 1)): + """BF16 tokens whose magnitudes are log-spaced over ``spread`` decades, in random order.""" + gen = torch.Generator(device="cuda") + gen.manual_seed(seed) + x = torch.randn(total, hidden, generator=gen, device="cuda") + x = x * torch.logspace(*spread, total, device="cuda")[:, None] + perm = torch.randperm(total, generator=gen, device="cuda") + return x[perm].to(torch.bfloat16) + + +def _pack(x, stochastic=False, seed=None): + if seed is not None: + torch.cuda.manual_seed(seed) + return tex.ep_nvfp4_pack(x, stochastic) + + +def _unpack(packed, hidden): + return tex.ep_nvfp4_unpack(packed, hidden) + + +def _receive(packed, hidden, counts, stochastic=False): + return _nvfp4_direct_receive(packed, hidden, counts, torch.bfloat16, stochastic) + + +def _swz(r, c, ncols): + return ( + ((r // 128) * ((ncols + 3) // 4) + c // 4) * 512 + + (r % 32) * 16 + + ((r // 32) % 4) * 4 + + (c % 4) + ) + + +def _codes(data_bytes): + d = data_bytes.to(torch.int64) + return torch.stack([d & 15, d >> 4], dim=-1).flatten(-2) + + +def _token_scales(packed, hidden): + off = hidden // 2 + hidden // 16 + return packed[:, off : off + 4].contiguous().view(torch.float32).flatten() + + +def _ref_pack(x): + """Reference RN pack in PyTorch: codes [T, H], block scales (uint8) [T, H/16], token scales.""" + x = x.float() + num_tokens, hidden = x.shape + amax = x.abs().amax(dim=1) + bits = (amax / NVFP4_MAX).view(torch.int32) + exp = bits & 0x7F800000 + dec = torch.where((bits & 0x7FFFFF) == 0, exp, exp + 0x800000).clamp_min(0x800000) + decode = torch.where(amax == 0, 0.0, dec.view(torch.float32)) + encode = torch.where(amax == 0, 1.0, 1.0 / decode) + blocks = x.view(num_tokens, hidden // 16, 16) + bmax = blocks.abs().amax(dim=-1) + bscale = (bmax * (encode * (1.0 / 6.0))[:, None]).clamp_max(448.0).to(torch.float8_e4m3fn) + rcp = (1.0 / (bscale.float() * decode[:, None])).clamp_max(torch.finfo(torch.float32).max) + q = (blocks * rcp[..., None]).view(num_tokens, hidden) + a = q.abs() + code = torch.zeros_like(q, dtype=torch.int64) + for thr, inclusive in ((0.25, True), (0.75, False), (1.25, True), (1.75, False), (2.5, True)): + code += (a > thr) if inclusive else (a >= thr) + code += (a >= 3.5).long() + (a > 5.0).long() + code += torch.where(q < 0, 8, 0) + return code, bscale.view(torch.uint8), decode + + +def _dequant_rowwise(g, total, hidden): + lut = LUT.to("cuda") + data = g.rowwise_data[: total * hidden // 2].view(total, hidden // 2) + vals = lut[_codes(data)] + r = torch.arange(total, device="cuda")[:, None] + c = torch.arange(hidden // 16, device="cuda")[None, :] + sb = g.scale_inv[_swz(r, c, hidden // 16)].view(torch.float8_e4m3fn).float() + counts = g.first_dims + group_of_row = torch.repeat_interleave(torch.arange(counts.numel(), device="cuda"), counts) + gscale = (g.amax / NVFP4_MAX)[group_of_row][:, None] + return vals * sb.repeat_interleave(16, dim=1) * gscale + + +def _dequant_columnwise_inverse_rht(g, hidden): + """Columnwise operand dequantized and mapped back through the inverse RHT: [total, hidden].""" + lut = LUT.to("cuda") + signs = get_wgrad_sign_vector("cuda") + had = get_hadamard_matrix(16, "cuda") # symmetric, orthonormal (1/4 scaled) + out = [] + off = 0 + for gi, n in enumerate(g.first_dims.tolist()): + if n == 0: + continue + data = g.columnwise_data[off * hidden // 2 : (off + n) * hidden // 2].view(hidden, n // 2) + vals = lut[_codes(data)] # [hidden, n] + k = torch.arange(hidden, device="cuda")[:, None] + b = torch.arange(n // 16, device="cuda")[None, :] + sc = g.columnwise_scale_inv[off * hidden // 16 + _swz(k, b, n // 16)] + sc = sc.view(torch.float8_e4m3fn).float().repeat_interleave(16, dim=1) + y = vals * sc * (g.columnwise_amax[gi] / NVFP4_MAX) + # y = H (s * v) per 16-token block => v = s * (H y) + yb = y.view(hidden, n // 16, 16) + v = torch.einsum("ij,kbj->kbi", had, yb) * signs + out.append(v.reshape(hidden, n).t()) + off += n + return torch.cat(out, dim=0) + + +def _te_group_quantize(x, counts, hidden, stochastic=False): + q = NVFP4Quantizer( + rowwise=True, + columnwise=True, + with_rht=True, + with_post_rht_amax=True, + stochastic_rounding=stochastic, + ) + q.optimize_for_gemm = True + q.internal = True + return tex.group_quantize( + x, q, counts.numel(), counts, tensor_offsets=tex.splits_to_offsets(counts, hidden) + ) + + +def _rel(a, b): + return ((a - b).norm() / b.norm().clamp_min(1e-30)).item() + + +def test_row_bytes_layout(): + assert _nvfp4_row_bytes(7168) == 4048 + x = _tokens(4, 7168, seed=0) + assert tuple(_pack(x).shape) == (4, 4048) + + +@pytest.mark.skipif(FAST_MATH, reason="the reference uses IEEE division") +@pytest.mark.parametrize("hidden", [512, 1024, 7168]) +def test_pack_matches_reference(hidden): + x = _tokens(300, hidden, seed=hidden) + x[5] = 0 + x[7, :16] = 0 # an all-zero block + packed = _pack(x) + code, bscale, decode = _ref_pack(x) + got_code = _codes(packed[:, : hidden // 2]) + # Zero magnitudes may carry either sign. + assert torch.equal( + torch.where(got_code & 7 == 0, 0, got_code), torch.where(code & 7 == 0, 0, code) + ) + assert torch.equal(packed[:, hidden // 2 : hidden // 2 + hidden // 16], bscale) + assert torch.equal(_token_scales(packed, hidden), decode) + assert (packed[:, hidden // 2 + hidden // 16 + 4 :] == 0).all() + + +def test_pack_pow2_scales_and_decode(): + hidden = 1024 + x = _tokens(64, hidden, seed=0) + x[3] = 0 + packed = _pack(x) + ts = _token_scales(packed, hidden) + assert ts[3].item() == 0.0 + nz = ts[ts > 0] + assert torch.equal(torch.frexp(nz).mantissa, torch.full_like(nz, 0.5)), "not powers of two" + amax = x.float().abs().amax(dim=1) + ok = amax > 0 + assert (amax[ok] <= NVFP4_MAX * ts[ok]).all() + assert (amax[ok] > NVFP4_MAX * ts[ok] / 2).all() + dec = _unpack(packed, hidden).float() + assert dec[3].abs().max().item() == 0.0 + # Per-token scales: at least as accurate as TE's per-tensor NVFP4 cast of the same tokens. + te_dec = NVFP4Quantizer(rowwise=True, columnwise=False)(x).dequantize().float() + assert _rel(dec, x.float()) <= 1.02 * _rel(te_dec, x.float()) + + +def test_unpack_bounded_by_counts(): + hidden = 512 + packed = _pack(_tokens(256, hidden, seed=1)) + counts = torch.tensor([64, 0, 32], dtype=torch.int64, device="cuda") + out = torch.full((256, hidden), 7.0, dtype=torch.bfloat16, device="cuda") + tex.ep_nvfp4_unpack(packed, hidden, counts, out) + assert torch.equal(out[:96], _unpack(packed, hidden)[:96]) + assert (out[96:] == 7.0).all() + + +@pytest.mark.parametrize("hidden", [512, 7168]) +def test_pack_sr_scales_neighbors_and_seed(hidden): + x = _tokens(128, hidden, seed=3) + x[2] = 0 + rn = _pack(x) + a = _pack(x, stochastic=True, seed=11) + b = _pack(x, stochastic=True, seed=11) + c = _pack(x, stochastic=True, seed=12) + assert torch.equal(a, b), "same seed must give the same bits" + assert not torch.equal(a, c), "different seeds must draw different bits" + # Block scales, token scales and padding are those of the RN pack. + assert torch.equal(a[:, hidden // 2 :], rn[:, hidden // 2 :]) + # Every value is one of the two grid neighbours of the scaled input. + dec = _unpack(a, hidden).float() + dec_rn = _unpack(rn, hidden).float() + assert dec[2].abs().max().item() == 0.0 + off = hidden // 2 + sb = a[:, off : off + hidden // 16].contiguous().view(torch.float8_e4m3fn).float() + ts = _token_scales(a, hidden)[:, None] + step_scale = (sb * ts).repeat_interleave(16, dim=1) # value of one E2M1 unit (1.0) + # E2M1 spacing is <= 2 units, so |SR - x| <= 2 * scale and |SR - RN| <= 2 * scale. + assert ((dec - x.float()).abs() <= 2 * step_scale + 1e-30).all() + assert ((dec - dec_rn).abs() <= 2 * step_scale + 1e-30).all() + + +def test_pack_sr_unbiased(): + hidden = 512 + x = _tokens(64, hidden, seed=5, spread=(-0.3, 0.3)) + n = 512 + acc = torch.zeros(64, hidden, dtype=torch.float64, device="cuda") + torch.cuda.manual_seed(1000) + for _ in range(n): + acc += _unpack(_pack(x, stochastic=True), hidden).double() + mean = acc / n + rn = _unpack(_pack(x), hidden).double() + xf = x.double() + bias = _rel(mean, xf) + rn_err = _rel(rn, xf) + signed = ((mean - xf).sum() / xf.abs().sum()).item() + print(f"SR pack: rel(mean of {n}) = {bias:.5f}, RN rel err = {rn_err:.5f}, signed {signed:.2e}") + assert bias < 0.25 * rn_err + assert abs(signed) < 1e-3 + + +@pytest.mark.parametrize("hidden", [512, 7168]) +@pytest.mark.parametrize("counts_list", [[256, 0, 512, 256], [768], [0, 256], [128, 384]]) +def test_direct_receive_rowwise_exact_and_columnwise_quality(hidden, counts_list): + counts = torch.tensor(counts_list, dtype=torch.int64, device="cuda") + total = sum(counts_list) + capacity = total + 256 # rows past the counts must be ignored + # Token magnitudes within ~8x of each other (like normalized activations), so re-expressed + # block scales stay in E4M3's normal range and the rowwise reuse is exact. + x = _tokens(capacity, hidden, seed=total + hidden, spread=(-0.45, 0.45)) + x[1] = 0 # a zero token + packed = _pack(x) + decoded = _unpack(packed, hidden) + g = _receive(packed, hidden, counts) + torch.cuda.synchronize() + assert g._with_gemm_swizzled_scales + assert not g.quantizer.stochastic_rounding + + # Per-expert second-level scales: 6*448*G and 4x that, G = max token scale of the expert. + ts = _token_scales(packed[:total], hidden) + starts = [0] + torch.tensor(counts_list).cumsum(0).tolist() + for gi in range(len(counts_list)): + seg = ts[starts[gi] : starts[gi + 1]] + G = seg.max().item() if seg.numel() and seg.max().item() > 0 else 1.0 + assert g.amax[gi].item() == NVFP4_MAX * G + assert g.columnwise_amax[gi].item() == 4 * NVFP4_MAX * G + + # Rowwise: the sender's codes and values are reused exactly. + n_data = total * hidden // 2 + assert torch.equal(g.rowwise_data[:n_data], packed[:total, : hidden // 2].reshape(-1)) + assert torch.equal(_dequant_rowwise(g, total, hidden), decoded[:total].float()) + + # Columnwise: as close to the decoded tokens as TE's own RHT quantization of them. + ref = _te_group_quantize(decoded[:total].contiguous(), counts, hidden) + got_err = _rel(_dequant_columnwise_inverse_rht(g, hidden), decoded[:total].float()) + te_err = _rel(_dequant_columnwise_inverse_rht(ref, hidden), decoded[:total].float()) + print(f"columnwise rel err direct={got_err:.5f} te={te_err:.5f}") + assert got_err < 1.05 * te_err + 1e-4 + + +def test_direct_receive_wide_token_range_rowwise_quality(): + """With token magnitudes spread over 1000x, small tokens' block scales underflow against the + expert scale; the rowwise error stays at TE's per-tensor level.""" + hidden = 1024 + counts = torch.tensor([512, 512], dtype=torch.int64, device="cuda") + x = _tokens(1024, hidden, seed=7, spread=(-2, 1)) + packed = _pack(x) + decoded = _unpack(packed, hidden) + g = _receive(packed, hidden, counts) + ref = _te_group_quantize(decoded.contiguous(), counts, hidden) + got_err = _rel(_dequant_rowwise(g, 1024, hidden), x.float()) + te_err = _rel(_dequant_rowwise(ref, 1024, hidden), x.float()) + print(f"rowwise rel err vs original: direct={got_err:.5f} decode+TE={te_err:.5f}") + assert got_err <= te_err * 1.02 + + +@pytest.mark.parametrize("hidden", [512, 7168]) +@pytest.mark.parametrize("counts_list", [[256, 0, 512, 256], [768], [0, 256]]) +def test_direct_receive_stochastic_rounding(hidden, counts_list): + counts = torch.tensor(counts_list, dtype=torch.int64, device="cuda") + total = sum(counts_list) + capacity = total + 256 + x = _tokens(capacity, hidden, seed=total + hidden, spread=(-0.45, 0.45)) + x[1] = 0 + packed = _pack(x, stochastic=True, seed=77) + decoded = _unpack(packed, hidden) + g = _receive(packed, hidden, counts, stochastic=True) + g_rn = _receive(packed, hidden, counts) + torch.cuda.synchronize() + assert g._with_gemm_swizzled_scales + assert g.quantizer.stochastic_rounding + # Rowwise and the per-expert scales do not depend on the rounding mode of the receive. + n_data = total * hidden // 2 + assert torch.equal(g.rowwise_data[:n_data], g_rn.rowwise_data[:n_data]) + assert torch.equal(g.amax, g_rn.amax) and torch.equal(g.columnwise_amax, g_rn.columnwise_amax) + assert torch.equal(_dequant_rowwise(g, total, hidden), decoded[:total].float()) + # Columnwise: a single SR draw is about as close to the decoded rows as TE's own RHT + SR + # quantization of them, and SR is unbiased (the mean over draws converges). + ref = _te_group_quantize(decoded[:total].contiguous(), counts, hidden, stochastic=True) + target = decoded[:total].float() + got_err = _rel(_dequant_columnwise_inverse_rht(g, hidden), target) + te_err = _rel(_dequant_columnwise_inverse_rht(ref, hidden), target) + rn_err = _rel(_dequant_columnwise_inverse_rht(g_rn, hidden), target) + acc = torch.zeros_like(target, dtype=torch.float64) + draws = 32 + for _ in range(draws): + acc += _dequant_columnwise_inverse_rht( + _receive(packed, hidden, counts, stochastic=True), hidden + ).double() + mean_err = _rel(acc / draws, target.double()) + print( + f"columnwise SR rel err: direct={got_err:.5f} te_sr={te_err:.5f} direct_rn={rn_err:.5f}" + f" mean_of_{draws}={mean_err:.5f}" + ) + assert got_err < 1.15 * te_err + 1e-4 + assert mean_err < 0.4 * got_err + + +def test_direct_receive_cuda_graph_reads_live_counts(): + hidden = 512 + counts = torch.tensor([256, 256], dtype=torch.int64, device="cuda") + packed = _pack(_tokens(768, hidden, seed=1, spread=(-0.45, 0.45))) + out = {} + + def step(): + out["g"] = _receive(packed, hidden, counts) + + step() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + step() + counts.copy_(torch.tensor([512, 256], device="cuda")) + packed.copy_(_pack(_tokens(768, hidden, seed=2, spread=(-0.45, 0.45)))) + graph.replay() + torch.cuda.synchronize() + g = out["g"] + decoded = _unpack(packed, hidden) + assert torch.equal(_dequant_rowwise(g, 768, hidden), decoded.float()) + eager = _receive(packed, hidden, counts) + n = 768 * hidden // 2 + assert torch.equal(g.columnwise_data[:n], eager.columnwise_data[:n]) + assert torch.equal(g.columnwise_amax, eager.columnwise_amax) + + +def test_sr_pack_and_receive_cuda_graph_fresh_bits(): + hidden = 512 + counts = torch.tensor([256, 256], dtype=torch.int64, device="cuda") + x = _tokens(512, hidden, seed=9, spread=(-0.45, 0.45)) + out = {} + + def step(): + packed = tex.ep_nvfp4_pack(x, True) + out["packed"] = packed + out["g"] = _receive(packed, hidden, counts, stochastic=True) + + step() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + step() + snaps = [] + for _ in range(2): + graph.replay() + torch.cuda.synchronize() + g = out["g"] + snaps.append((out["packed"].clone(), g.columnwise_data[: 512 * hidden // 2].clone())) + dec = _unpack(out["packed"], hidden) + assert torch.equal(_dequant_rowwise(g, 512, hidden), dec.float()) + assert not torch.equal(snaps[0][0], snaps[1][0]), "pack must draw fresh bits per replay" + assert not torch.equal(snaps[0][1], snaps[1][1]), "receive must draw fresh bits per replay" + + +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="needs 2 GPUs") +def test_ops_run_on_input_device(): + """The pack, decode and direct receive run on their input's device, with that device's RNG, + not on the current device: tensors on cuda:1 while cuda:0 is current give the same result + as with cuda:1 current, and only cuda:1's generator advances.""" + hidden, sign_mask, seed = 512, 0x5A5A, 1234 + dev = torch.device("cuda", 1) + gens = torch.cuda.default_generators + + def run(): + torch.cuda.manual_seed_all(seed) + offsets = [gens[0].get_offset(), gens[1].get_offset()] + packed = tex.ep_nvfp4_pack(x, False) + out = [packed, tex.ep_nvfp4_pack(x, True), tex.ep_nvfp4_unpack(packed, hidden, counts)] + out += tex.ep_nvfp4_direct_receive(packed, hidden, counts, sign_mask, False) + out += tex.ep_nvfp4_direct_receive(packed, hidden, counts, sign_mask, True) + torch.cuda.synchronize(dev) + advanced = [gens[0].get_offset() - offsets[0], gens[1].get_offset() - offsets[1]] + return out, advanced + + with torch.cuda.device(dev): + x = _tokens(512, hidden, seed=3) + counts = torch.tensor([256, 256], dtype=torch.int64, device=dev) + ref, ref_advanced = run() + with torch.cuda.device(0): + got, advanced = run() + assert ref_advanced[0] == 0 and ref_advanced[1] > 0 + assert advanced == ref_advanced, "stochastic calls must draw from the input device's RNG" + for g, r in zip(got, ref): + assert g.device == dev + assert torch.equal(g, r) + + +def _grouped_mlp(group_size, hidden, recipe, weight_requires_grad): + with te.quantized_model_init(enabled=True, recipe=recipe): + fc1 = te.ops.GroupedLinear( + group_size, + hidden, + 2 * hidden, + bias=False, + device="cuda", + dtype=torch.bfloat16, + split_alignment=256, + ) + fc2 = te.ops.GroupedLinear( + group_size, + hidden, + hidden, + bias=False, + device="cuda", + dtype=torch.bfloat16, + split_alignment=256, + ) + module = te.ops.Sequential(fc1, te.ops.ScaledSwiGLU(glu_interleave_size=32), fc2) + with torch.no_grad(): + for param in module.parameters(): + param.requires_grad_(requires_grad=weight_requires_grad) + return module, fc1, fc2 + + +def _split_sizes(group_size): + split = [256 * i for i in range(group_size)] + random.Random(0).shuffle(split) + return torch.tensor(split, dtype=torch.int64, device="cuda"), sum(split) + + +@pytest.mark.parametrize("weight_requires_grad", (False, True)) +def test_fused_grouped_mlp_with_direct_receive(weight_requires_grad): + """The fused grouped MLP fed the direct-receive GroupedTensor is at least as accurate (vs an + unquantized reference) as when it group-quantizes the decoded BF16 rows itself.""" + if not te.ops.fused.GroupedMLP_CuTeGEMMGLU.is_supported(): + pytest.skip("Fused grouped MLP (CuTeDSL) is not supported on this system") + group_size, hidden = 4, 512 + split_sizes, total = _split_sizes(group_size) + packed = _pack(_tokens(total, hidden, seed=3)) + x_dec = _unpack(packed, hidden) + probs = torch.rand((total,), dtype=torch.bfloat16, device="cuda") + dy = torch.rand((total, hidden), dtype=torch.bfloat16, device="cuda") - 0.5 + recipe = NVFP4BlockScaling(disable_stochastic_rounding=True) + module, fc1, fc2 = _grouped_mlp(group_size, hidden, recipe, weight_requires_grad) + + def _run(x, quantized=True): + with te.autocast(enabled=quantized, recipe=recipe): + y = module(x, split_sizes, probs, split_sizes) + grads = [] + if weight_requires_grad: + y.backward(dy) + for lin in (fc1, fc2): + for idx in range(group_size): + w = getattr(lin, f"weight{idx}") + grads.append(w.grad.detach().float().clone()) + w.grad = None + return y.detach().float(), grads + + y_dec, g_dec = _run(x_dec) + y_dir, g_dir = _run(_receive(packed, hidden, split_sizes)) + forward_ops = module._module_groups[0]._forward_ops + assert isinstance(forward_ops[0][0], te.ops.fused.GroupedMLP_CuTeGEMMGLU) + + # Both runs share weights and all other quantization; they differ only in how the FC1 input + # is quantized (direct: the sender's codes and a bound-scaled RHT operand; decode path: TE + # requantizes the decoded BF16 rows). Judge each against the same module run unquantized on + # the decoded rows, and require the direct path to be no less accurate. + y_ref, g_ref = _run(x_dec, quantized=False) + names = ["y"] + [f"fc{i}.w{j}" for i in (1, 2) for j in range(group_size)] + err_dir = {"y": _rel(y_dir, y_ref)} + err_dec = {"y": _rel(y_dec, y_ref)} + err_dir.update({n: _rel(a, b) for n, a, b in zip(names[1:], g_dir, g_ref)}) + err_dec.update({n: _rel(a, b) for n, a, b in zip(names[1:], g_dec, g_ref)}) + print("rel err vs unquantized: direct", {k: round(v, 4) for k, v in err_dir.items()}) + print("rel err vs unquantized: decode", {k: round(v, 4) for k, v in err_dec.items()}) + for n in err_dir: + assert err_dir[n] <= 1.1 * err_dec[n] + 1e-3, (n, err_dir[n], err_dec[n]) + + +class _InjectGrad(torch.autograd.Function): + """Identity whose backward returns a prepared grad (the combine backward's GroupedTensor).""" + + @staticmethod + def forward(ctx, y, holder): # type: ignore[override] + ctx.holder = holder + return y.view_as(y) + + @staticmethod + def backward(ctx, _g): # type: ignore[override] + return ctx.holder["grad"], None + + +@pytest.mark.parametrize("weight_requires_grad", (False, True)) +def test_fused_grouped_mlp_backward_with_direct_grad(weight_requires_grad): + """The fused grouped MLP backward fed FC2's grad_output as the stochastically rounded + direct-receive GroupedTensor is as accurate (vs an unquantized reference) as when it + SR-quantizes the BF16 grad itself, except for FC2's wgrad, whose doubly rounded columnwise + operand adds variance.""" + if not te.ops.fused.GroupedMLP_CuTeGEMMGLU.is_supported(): + pytest.skip("Fused grouped MLP (CuTeDSL) is not supported on this system") + group_size, hidden = 4, 512 + split_sizes, total = _split_sizes(group_size) + x = _tokens(total, hidden, seed=3, spread=(-0.3, 0.3)) + probs = torch.rand((total,), dtype=torch.bfloat16, device="cuda") + dy = _tokens(total, hidden, seed=4, spread=(-0.3, 0.3)) + recipe = NVFP4BlockScaling() + module, fc1, fc2 = _grouped_mlp(group_size, hidden, recipe, weight_requires_grad) + + def _run(direct, quantized=True, seed=0): + torch.manual_seed(seed) + xin = x.detach().clone().requires_grad_(True) + with te.autocast(enabled=quantized, recipe=recipe): + y = module(xin, split_sizes, probs, split_sizes) + if direct: + grad = _receive( + _pack(dy, stochastic=True, seed=seed + 1), hidden, split_sizes, stochastic=True + ) + _InjectGrad.apply(y, {"grad": grad}).backward(torch.zeros_like(y)) + else: + y.backward(dy) + grads = [xin.grad.detach().float().clone()] + if weight_requires_grad: + for lin in (fc1, fc2): + for idx in range(group_size): + w = getattr(lin, f"weight{idx}") + grads.append(w.grad.detach().float().clone()) + w.grad = None + return grads + + ref = _run(False, quantized=False) + names = ["dx"] + [f"fc{i}.w{j}" for i in (1, 2) for j in range(group_size)] + names = names[: len(ref)] + # Average errors over a few SR draws to compare the two noisy paths fairly. + err_dir = {n: 0.0 for n in names} + err_te = {n: 0.0 for n in names} + draws = 3 + for s in range(draws): + gd = _run(True, seed=10 + s) + gt = _run(False, seed=10 + s) + for n, a, b, r in zip(names, gd, gt, ref): + assert torch.isfinite(a).all(), n + err_dir[n] += _rel(a, r) / draws + err_te[n] += _rel(b, r) / draws + forward_ops = module._module_groups[0]._forward_ops + assert isinstance(forward_ops[0][0], te.ops.fused.GroupedMLP_CuTeGEMMGLU) + print("rel err vs unquantized: direct", {k: round(v, 4) for k, v in err_dir.items()}) + print("rel err vs unquantized: te ", {k: round(v, 4) for k, v in err_te.items()}) + for n in names: + # FC2's wgrad reads the columnwise operand, which is rounded twice (sender SR, then RHT + + # SR on the receiver) where TE rounds once: both unbiased, but the variances add up. + # dgrad and FC1's wgrad see a single rounding. + factor = 1.25 if n.startswith("fc2") else 1.05 + assert err_dir[n] <= factor * err_te[n] + 2e-3, (n, err_dir[n], err_te[n]) diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index f90f74449c..c91bb3e88c 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -364,6 +364,7 @@ list(APPEND transformer_engine_cuda_arch_specific_sources cast/cast_grouped.cu cast/cast_grouped_dbias.cu cast/mxfp8/requantize_mxfp8.cu + ep/nvfp4_ep_rows.cu fused_rope/fused_mla_kv_rope_mxfp8.cu gemm/cutlass_grouped_gemm.cu hadamard_transform/group_hadamard_transform.cu diff --git a/transformer_engine/common/ep/nvfp4_ep_rows.cu b/transformer_engine/common/ep/nvfp4_ep_rows.cu new file mode 100644 index 0000000000..814844322b --- /dev/null +++ b/transformer_engine/common/ep/nvfp4_ep_rows.cu @@ -0,0 +1,715 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +/*! \file nvfp4_ep_rows.cu + * \brief NVFP4 token rows for the EP dispatch and combine backward. + * + * The sender packs each token into one NVFP4 row (E2M1 codes, E4M3 block scales and a + * power-of-two FP32 token scale); the rows travel as an unscaled payload. The receiver either + * decodes them to BF16 or, with the direct receive, builds the grouped GEMM's NVFP4 operands + * from them in one pass: + * + * - Rowwise: the sender's E2M1 codes are reused unchanged. With power-of-two token scales each + * block scale is only re-expressed against the per-expert scale G = max token scale, an exact + * power-of-two shift while it stays in E4M3's normal range. + * - Columnwise: the 16-point RHT along tokens is decoded and quantized in the same pass. Its + * second-level scale is the bound 4 * G: |RHT(x)| <= 4 max|x| and max|x| <= 6 * 448 * G, so + * no amax pass over the data is needed. + * + * The same receive serves the combine backward (FC2's grad_output): the sender packs with + * stochastic rounding and the receiver quantizes the columnwise operand with stochastic + * rounding too, as the NVFP4 recipe does for gradients. + */ + +#include +#include +#include +#include + +#include + +#include "../common.h" +#include "../util/cuda_runtime.h" +#include "../util/curanddx.hpp" +#include "../util/logging.h" +#include "../util/ptx.cuh" +#include "../utils.cuh" + +namespace transformer_engine { +namespace ep_nvfp4 { +namespace { + +constexpr int kTileM = 128; // tokens per receive tile (8 RHT blocks of 16) +constexpr int kTileK = 256; // hidden columns per receive tile (16 scale blocks of 16) +constexpr int kThreads = 256; +constexpr int kMaxGroups = 128; +constexpr float kNvfp4Max = 6.0f * 448.0f; + +using Philox = curanddx::detail::philox4x32_native_state; + +size_t min_row_bytes(size_t hidden) { return hidden / 2 + hidden / 16 + 4; } + +__device__ __forceinline__ float e4m3_to_float(uint8_t b) { + __nv_fp8_e4m3 v; + v.__x = b; + return static_cast(v); +} + +__device__ __forceinline__ uint8_t float_to_e4m3_satfinite(float x) { + return static_cast(__nv_cvt_float_to_fp8(x, __NV_SATFINITE, __NV_E4M3)); +} + +#if FP4_TYPE_SUPPORTED + +// Low nibble = lo; RN, satfinite. +__device__ __forceinline__ uint8_t float2_to_e2m1x2(float lo, float hi) { + return static_cast( + __nv_cvt_float2_to_fp4x2(make_float2(lo, hi), __NV_E2M1, cudaRoundNearest)); +} + +// Low nibble = v.x; RN, satfinite. +__device__ __forceinline__ uint8_t half2_to_e2m1x2(__half2 v) { + return static_cast( + __nv_cvt_halfraw2_to_fp4x2(static_cast<__half2_raw>(v), __NV_E2M1, cudaRoundNearest)); +} + +// Stochastic rounding of 4 floats to E2M1 (satfinite); x0 lands in the low nibble. +__device__ __forceinline__ uint16_t float4_to_e2m1x4_sr(float x0, float x1, float x2, float x3, + uint32_t rbits) { + uint16_t out = 0; + constexpr bool has_rs = ARCH_HAS_STOCHASTIC_ROUNDING; + if constexpr (has_rs) { + asm volatile("cvt.rs.satfinite.e2m1x4.f32 %0, {%1, %2, %3, %4}, %5;\n" + : "=h"(out) + : "f"(x3), "f"(x2), "f"(x1), "f"(x0), "r"(rbits)); + } else { + NVTE_DEVICE_ERROR("Stochastic rounding to E2M1 requires SM 10.0, 10.3 or 10.7."); + } + return out; +} + +#endif // FP4_TYPE_SUPPORTED + +__device__ __forceinline__ Philox make_rng(const int64_t *rng_state, uint64_t subsequence) { + Philox rng; + rng.init(static_cast(rng_state[0]), subsequence, static_cast(rng_state[1])); + return rng; +} + +template +__device__ __forceinline__ float reciprocal(float x) { + float r; + if constexpr (kFastMath) { + asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(r) : "f"(x)); + } else { + r = 1.0f / x; + } + return fminf(r, 3.4028234663852886e38f); +} + +// Offset of scale (row r, scale col c) in the GEMM-swizzled 128x4 tile layout. +__device__ __forceinline__ int64_t swizzled_scale_offset(int64_t r, int64_t c, + int64_t num_scale_cols) { + return ((r / 128) * ((num_scale_cols + 3) / 4) + c / 4) * 512 + (r % 32) * 16 + + ((r / 32) % 4) * 4 + (c % 4); +} + +__device__ __forceinline__ void cp_async16(void *smem, const void *gmem) { + const uint32_t addr = static_cast(__cvta_generic_to_shared(smem)); + asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n" ::"r"(addr), "l"(gmem)); +} + +__device__ __forceinline__ void cp_async4(void *smem, const void *gmem) { + const uint32_t addr = static_cast(__cvta_generic_to_shared(smem)); + asm volatile("cp.async.ca.shared.global [%0], [%1], 4;\n" ::"r"(addr), "l"(gmem)); +} + +__device__ __forceinline__ void cp_async_commit() { asm volatile("cp.async.commit_group;\n" ::); } + +template +__device__ __forceinline__ void cp_async_wait() { + asm volatile("cp.async.wait_group %0;\n" ::"n"(N)); +} + +struct Groups { + int num; + int64_t ends[kMaxGroups]; // inclusive prefix sums of the counts +}; + +__device__ __forceinline__ void load_groups(const int64_t *counts, int num_groups, Groups &s) { + if (threadIdx.x == 0) { + int64_t acc = 0; + for (int g = 0; g < num_groups; ++g) { + acc += counts[g]; + s.ends[g] = acc; + } + s.num = num_groups; + } + __syncthreads(); +} + +__device__ __forceinline__ int group_of_row(const Groups &s, int64_t r) { + int g = 0; + while (g < s.num - 1 && s.ends[g] <= r) ++g; + return g; +} + +__device__ __forceinline__ float group_scale(const int *gmax_bits, int g) { + const float G = __int_as_float(gmax_bits[g]); + return G > 0.0f ? G : 1.0f; +} + +// Sender pack: one block per token row, one thread per 16-element scale block (blockDim = H/16). +template +__global__ void pack_kernel(const __nv_bfloat16 *__restrict__ x, uint8_t *__restrict__ out, int H, + int64_t row_bytes, const int64_t *__restrict__ rng_state) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) && FP4_TYPE_SUPPORTED + __shared__ float s_warp_max[32]; + const int64_t row = blockIdx.x; + const int b = threadIdx.x; + const __nv_bfloat16 *xr = x + row * H + b * 16; + uint8_t *orow = out + row * row_bytes; + float v[16]; + { + const uint4 a = *reinterpret_cast(xr); + const uint4 c = *reinterpret_cast(xr + 8); + const __nv_bfloat162 *pa = reinterpret_cast(&a); + const __nv_bfloat162 *pc = reinterpret_cast(&c); +#pragma unroll + for (int i = 0; i < 4; ++i) { + const float2 fa = __bfloat1622float2(pa[i]); + const float2 fc = __bfloat1622float2(pc[i]); + v[2 * i] = fa.x; + v[2 * i + 1] = fa.y; + v[8 + 2 * i] = fc.x; + v[8 + 2 * i + 1] = fc.y; + } + } + float bmax = 0.0f; +#pragma unroll + for (int i = 0; i < 16; ++i) bmax = fmaxf(bmax, fabsf(v[i])); + float m = bmax; +#pragma unroll + for (int o = 16; o > 0; o >>= 1) m = fmaxf(m, __shfl_xor_sync(0xffffffffu, m, o)); + const int lane = threadIdx.x & 31, warp = threadIdx.x >> 5; + if (lane == 0) s_warp_max[warp] = m; + __syncthreads(); + if (warp == 0) { + const int nw = (blockDim.x + 31) >> 5; + float mm = lane < nw ? s_warp_max[lane] : 0.0f; +#pragma unroll + for (int o = 16; o > 0; o >>= 1) mm = fmaxf(mm, __shfl_xor_sync(0xffffffffu, mm, o)); + if (lane == 0) s_warp_max[0] = mm; + } + __syncthreads(); + const float row_amax = s_warp_max[0]; + // Power-of-two token scale: the smallest 2^e >= amax / (6 * 448). + const uint32_t ratio_bits = __float_as_uint(__fdiv_rn(row_amax, kNvfp4Max)); + const uint32_t exp_bits = ratio_bits & 0x7F800000u; + uint32_t dec_bits = (ratio_bits & 0x7FFFFFu) == 0 ? exp_bits : exp_bits + 0x800000u; + dec_bits = max(dec_bits, 0x800000u); + const float global_decode = row_amax == 0.0f ? 0.0f : __uint_as_float(dec_bits); + const float global_encode = row_amax == 0.0f ? 1.0f : __fdiv_rn(1.0f, global_decode); + // Same block-scale arithmetic as TE's NVFP4 cast (core_nvfp4.cuh). + const float bs_f = fminf(bmax * (global_encode * (1.0f / 6.0f)), 448.0f); + const uint8_t bs = float_to_e4m3_satfinite(bs_f); + const float rcp = reciprocal(e4m3_to_float(bs) * global_decode); + uint32_t lo = 0, hi = 0; + if constexpr (kStochastic) { + Philox rng = make_rng(rng_state, static_cast(row) * blockDim.x + b); + const uint4 rb = rng.generate4(); + lo = static_cast( + float4_to_e2m1x4_sr(v[0] * rcp, v[1] * rcp, v[2] * rcp, v[3] * rcp, rb.x)) | + (static_cast( + float4_to_e2m1x4_sr(v[4] * rcp, v[5] * rcp, v[6] * rcp, v[7] * rcp, rb.y)) + << 16); + hi = static_cast( + float4_to_e2m1x4_sr(v[8] * rcp, v[9] * rcp, v[10] * rcp, v[11] * rcp, rb.z)) | + (static_cast( + float4_to_e2m1x4_sr(v[12] * rcp, v[13] * rcp, v[14] * rcp, v[15] * rcp, rb.w)) + << 16); + } else { +#pragma unroll + for (int p = 0; p < 4; ++p) { + lo |= static_cast(float2_to_e2m1x2(v[2 * p] * rcp, v[2 * p + 1] * rcp)) << (8 * p); + hi |= static_cast(float2_to_e2m1x2(v[8 + 2 * p] * rcp, v[8 + 2 * p + 1] * rcp)) + << (8 * p); + } + } + *reinterpret_cast(orow + b * 8) = make_uint2(lo, hi); + orow[H / 2 + b] = bs; + const int64_t tail = H / 2 + H / 16; + if (b == 0) *reinterpret_cast(orow + tail) = global_decode; + for (int64_t i = tail + 4 + b; i < row_bytes; i += blockDim.x) orow[i] = 0; +#else + NVTE_DEVICE_ERROR("NVFP4 EP pack requires SM 10.0+."); +#endif +} + +// Decode: one thread per 16-element block of a row, grid-stride, bounded by sum(counts). +__global__ void __launch_bounds__(kThreads) + unpack_kernel(const uint8_t *__restrict__ packed, const int64_t *__restrict__ counts, + int num_counts, int64_t num_rows, int H, int64_t row_bytes, + __nv_bfloat16 *__restrict__ out) { + constexpr float kE2M1[8] = {0.0f, 0.5f, 1.0f, 1.5f, 2.0f, 3.0f, 4.0f, 6.0f}; + __shared__ int64_t s_bound; + if (threadIdx.x == 0) { + int64_t bound = num_rows; + if (counts != nullptr) { + int64_t total = 0; + for (int g = 0; g < num_counts; ++g) total += counts[g]; + bound = min(total, num_rows); + } + s_bound = bound; + } + __syncthreads(); + const int blocks_per_row = H / 16; + const int64_t num_items = s_bound * blocks_per_row; + for (int64_t i = blockIdx.x * static_cast(blockDim.x) + threadIdx.x; i < num_items; + i += static_cast(gridDim.x) * blockDim.x) { + const int64_t r = i / blocks_per_row; + const int b = static_cast(i % blocks_per_row); + const uint8_t *row = packed + r * row_bytes; + const uint2 codes = *reinterpret_cast(row + b * 8); + const float token_scale = *reinterpret_cast(row + H / 2 + H / 16); + const float scale = e4m3_to_float(row[H / 2 + b]) * token_scale; + const uint32_t words[2] = {codes.x, codes.y}; + __nv_bfloat162 vals[8]; +#pragma unroll + for (int w = 0; w < 2; ++w) { +#pragma unroll + for (int p = 0; p < 4; ++p) { + const uint32_t byte = (words[w] >> (8 * p)) & 0xFFu; + const uint32_t c0 = byte & 0xFu, c1 = byte >> 4; + const float v0 = (c0 & 8u ? -kE2M1[c0 & 7u] : kE2M1[c0 & 7u]) * scale; + const float v1 = (c1 & 8u ? -kE2M1[c1 & 7u] : kE2M1[c1 & 7u]) * scale; + vals[w * 4 + p] = __floats2bfloat162_rn(v0, v1); + } + } + uint4 *dst = reinterpret_cast(out + r * H + b * 16); + dst[0] = *reinterpret_cast(&vals[0]); + dst[1] = *reinterpret_cast(&vals[4]); + } +} + +// Receive pass 1: per-group max of the (non-negative) token decode scales, as float bits. +__global__ void group_max_kernel(const uint8_t *__restrict__ packed, const int64_t *counts, + int num_groups, int64_t capacity, int64_t row_bytes, + int64_t scale_offset, int *__restrict__ gmax_bits) { + __shared__ Groups s; + __shared__ int local_max[kMaxGroups]; + load_groups(counts, num_groups, s); + for (int g = threadIdx.x; g < num_groups; g += blockDim.x) local_max[g] = 0; + __syncthreads(); + const int64_t total = min(s.ends[num_groups - 1], capacity); + for (int64_t r = blockIdx.x * static_cast(blockDim.x) + threadIdx.x; r < total; + r += static_cast(gridDim.x) * blockDim.x) { + const float ts = *reinterpret_cast(packed + r * row_bytes + scale_offset); + atomicMax(&local_max[group_of_row(s, r)], __float_as_int(fmaxf(ts, 0.0f))); + } + __syncthreads(); + for (int g = threadIdx.x; g < num_groups; g += blockDim.x) { + if (local_max[g] != 0) atomicMax(&gmax_bits[g], local_max[g]); + } +} + +// Receive pass 2: rowwise copy + scale re-expression, and columnwise RHT quantization, per +// 128x256 tile. Values are handled relative to the expert scale G (x / G): the codes times the +// re-expressed block scales. The columnwise math runs in half2 (two adjacent hidden columns share +// every E2M1 byte, block scale and token scale), with the RHT's 1/4 and random signs folded into +// the per-token scale. Relative values are <= 4 * 6 * 448, well inside FP16. +template +__global__ void __launch_bounds__(kThreads) + direct_receive_kernel(const uint8_t *__restrict__ packed, const int64_t *counts, int num_groups, + int64_t capacity, int H, int64_t row_bytes, + const int *__restrict__ gmax_bits, uint32_t sign_mask, + const int64_t *__restrict__ rng_state, uint8_t *__restrict__ row_data, + uint8_t *__restrict__ row_scale, uint8_t *__restrict__ col_data, + uint8_t *__restrict__ col_scale, float *__restrict__ row_amax, + float *__restrict__ col_amax) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) && FP4_TYPE_SUPPORTED + __shared__ Groups s; + // Double-buffered tile inputs, filled with cp.async one tile ahead. + __shared__ __align__(16) uint8_t s_codes_buf[2][kTileM][kTileK / 2]; + __shared__ __align__(16) uint8_t s_bscale_buf[2][kTileM][kTileK / 16]; + __shared__ __align__(16) float s_tscale_buf[2][kTileM]; + __shared__ __align__(16) __half2 s_hscale[kTileM][kTileK / 16]; // b * ts / G * sign / 4, x2 + load_groups(counts, num_groups, s); + if (blockIdx.x == 0) { + for (int g = threadIdx.x; g < num_groups; g += blockDim.x) { + const float G = group_scale(gmax_bits, g); + row_amax[g] = kNvfp4Max * G; + col_amax[g] = 4.0f * kNvfp4Max * G; + } + } + + const int64_t total = min(s.ends[num_groups - 1], capacity); + const int col_tiles = H / kTileK; + const int64_t num_tiles = (total / kTileM) * col_tiles; + const int64_t num_scale_cols_row = H / 16; + const int64_t tscale_offset = H / 2 + H / 16; + + auto issue_loads = [&](int64_t tile, int buf) { + const int64_t lr0 = (tile / col_tiles) * kTileM; + const int lc0 = static_cast(tile % col_tiles) * kTileK; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const int idx = threadIdx.x + i * kThreads; // 0..1023 + const int rr = idx >> 3; // 8 x 16 B per row + const int cb = (idx & 7) * 16; + cp_async16(&s_codes_buf[buf][rr][cb], packed + (lr0 + rr) * row_bytes + lc0 / 2 + cb); + } + if (threadIdx.x < kTileM) { + const int rr = threadIdx.x; + cp_async16(&s_bscale_buf[buf][rr][0], packed + (lr0 + rr) * row_bytes + H / 2 + lc0 / 16); + } else { + const int rr = threadIdx.x - kTileM; + cp_async4(&s_tscale_buf[buf][rr], packed + (lr0 + rr) * row_bytes + tscale_offset); + } + cp_async_commit(); + }; + + int buf = 0; + if (blockIdx.x < num_tiles) issue_loads(blockIdx.x, 0); + for (int64_t t = blockIdx.x; t < num_tiles; t += gridDim.x, buf ^= 1) { + if (t + gridDim.x < num_tiles) { + issue_loads(t + gridDim.x, buf ^ 1); + cp_async_wait<1>(); + } else { + cp_async_wait<0>(); + } + __syncthreads(); + auto &s_codes = s_codes_buf[buf]; + auto &s_bscale = s_bscale_buf[buf]; + auto &s_tscale = s_tscale_buf[buf]; + const int64_t r0 = (t / col_tiles) * kTileM; + const int c0 = static_cast(t % col_tiles) * kTileK; + const int g = group_of_row(s, r0); + const int64_t g_end = s.ends[g]; + const int64_t g_rows = g == 0 ? g_end : g_end - s.ends[g - 1]; + const int64_t g_off = g_end - g_rows; + const float inv_G = 1.0f / group_scale(gmax_bits, g); // exact: G is a power of two + + // The rowwise data is the sender's codes, unchanged. +#pragma unroll + for (int i = 0; i < 4; ++i) { + const int idx = threadIdx.x + i * kThreads; + const int rr = idx >> 3; + const int cb = (idx & 7) * 16; + *reinterpret_cast(row_data + (r0 + rr) * (H / 2) + c0 / 2 + cb) = + *reinterpret_cast(&s_codes[rr][cb]); + } + + // Re-expressed block scales: rowwise E4M3 (swizzled, one 4-byte word per item) and the FP16 + // columnwise factors. +#pragma unroll + for (int i = 0; i < 2; ++i) { + const int idx = threadIdx.x + i * kThreads; // 0..511: 128 rows x 4 words + const int rr = idx & (kTileM - 1); + const int m = idx >> 7; + const float ratio = s_tscale[rr] * inv_G; + const float col_factor = (((sign_mask >> (rr & 15)) & 1) ? -0.25f : 0.25f); + uint32_t word = 0; +#pragma unroll + for (int j = 0; j < 4; ++j) { + const float b = e4m3_to_float(s_bscale[rr][m * 4 + j]) * ratio; + word |= static_cast(float_to_e4m3_satfinite(b)) << (8 * j); + s_hscale[rr][m * 4 + j] = __float2half2_rn(b * col_factor); + } + const int64_t off = swizzled_scale_offset(r0 + rr, c0 / 16 + m * 4, num_scale_cols_row); + *reinterpret_cast(row_scale + off) = word; + } + __syncthreads(); + + // Columnwise: thread = 2 adjacent hidden columns x 4 RHT blocks of 16 tokens. + { + const int cp = threadIdx.x & 127; + const int th = threadIdx.x >> 7; + const int sc = cp >> 3; + uint32_t out_x[8], out_y[8]; + uint8_t scl_x[4], scl_y[4]; + Philox rng; + if constexpr (kStochastic) { + rng = make_rng(rng_state, static_cast(t) * kThreads + threadIdx.x); + } +#pragma unroll + for (int b4 = 0; b4 < 4; ++b4) { + const int rbase = (th * 4 + b4) * 16; + __half2 v[16]; +#pragma unroll + for (int i = 0; i < 16; ++i) { + const int rr = rbase + i; + const __half2_raw raw = __nv_cvt_fp4x2_to_halfraw2( + static_cast<__nv_fp4x2_storage_t>(s_codes[rr][cp]), __NV_E2M1); + v[i] = __hmul2(__half2(raw), s_hscale[rr][sc]); + } + // In-place 16-point Walsh-Hadamard butterfly. +#pragma unroll + for (int h = 1; h < 16; h <<= 1) { +#pragma unroll + for (int i = 0; i < 16; ++i) { + if ((i & h) == 0) { + const __half2 a = v[i], b = v[i + h]; + v[i] = __hadd2(a, b); + v[i + h] = __hsub2(a, b); + } + } + } + __half2 m2 = __habs2(v[0]); +#pragma unroll + for (int i = 1; i < 16; ++i) m2 = __hmax2(m2, __habs2(v[i])); + const float2 bmax = __half22float2(m2); + // Relative to G: block scale = bmax_rel / 24, value = y_rel / (4 * block scale). + const uint8_t sbx = float_to_e4m3_satfinite(fminf(bmax.x * (1.0f / 24.0f), 448.0f)); + const uint8_t sby = float_to_e4m3_satfinite(fminf(bmax.y * (1.0f / 24.0f), 448.0f)); + scl_x[b4] = sbx; + scl_y[b4] = sby; + const __half2 rx = __float2half2_rn(reciprocal(4.0f * e4m3_to_float(sbx))); + const __half2 ry = __float2half2_rn(reciprocal(4.0f * e4m3_to_float(sby))); + if constexpr (kStochastic) { + // 4 random words per 8 tokens of both columns. + const __half2 r2 = __halves2half2(__low2half(rx), __low2half(ry)); +#pragma unroll + for (int w = 0; w < 2; ++w) { + const uint4 rb = rng.generate4(); + float fx[8], fy[8]; +#pragma unroll + for (int j = 0; j < 8; ++j) { + const float2 f = __half22float2(__hmul2(v[w * 8 + j], r2)); + fx[j] = f.x; + fy[j] = f.y; + } + out_x[b4 * 2 + w] = + static_cast(float4_to_e2m1x4_sr(fx[0], fx[1], fx[2], fx[3], rb.x)) | + (static_cast(float4_to_e2m1x4_sr(fx[4], fx[5], fx[6], fx[7], rb.y)) + << 16); + out_y[b4 * 2 + w] = + static_cast(float4_to_e2m1x4_sr(fy[0], fy[1], fy[2], fy[3], rb.z)) | + (static_cast(float4_to_e2m1x4_sr(fy[4], fy[5], fy[6], fy[7], rb.w)) + << 16); + } + } else { +#pragma unroll + for (int w = 0; w < 2; ++w) { + uint32_t wx = 0, wy = 0; +#pragma unroll + for (int p = 0; p < 4; ++p) { + const int i = w * 8 + p * 2; + // Regroup (col c, col c+1) pairs into (token i, token i+1) pairs per column. + const __half2 qx = __hmul2(__lows2half2(v[i], v[i + 1]), rx); + const __half2 qy = __hmul2(__highs2half2(v[i], v[i + 1]), ry); + wx |= static_cast(half2_to_e2m1x2(qx)) << (8 * p); + wy |= static_cast(half2_to_e2m1x2(qy)) << (8 * p); + } + out_x[b4 * 2 + w] = wx; + out_y[b4 * 2 + w] = wy; + } + } + } + const int64_t local = r0 - g_off; + const int64_t k = c0 + 2 * cp; + uint8_t *dst = col_data + g_off * (H / 2) + k * (g_rows / 2) + local / 2 + th * 32; + *reinterpret_cast(dst) = make_uint4(out_x[0], out_x[1], out_x[2], out_x[3]); + *reinterpret_cast(dst + 16) = make_uint4(out_x[4], out_x[5], out_x[6], out_x[7]); + dst += g_rows / 2; + *reinterpret_cast(dst) = make_uint4(out_y[0], out_y[1], out_y[2], out_y[3]); + *reinterpret_cast(dst + 16) = make_uint4(out_y[4], out_y[5], out_y[6], out_y[7]); + const int64_t blk = local / 16 + th * 4; + const int64_t ncols = g_rows / 16; + const uint32_t sx = + scl_x[0] | (scl_x[1] << 8) | (scl_x[2] << 16) | (static_cast(scl_x[3]) << 24); + const uint32_t sy = + scl_y[0] | (scl_y[1] << 8) | (scl_y[2] << 16) | (static_cast(scl_y[3]) << 24); + *reinterpret_cast(col_scale + g_off * (H / 16) + + swizzled_scale_offset(k, blk, ncols)) = sx; + *reinterpret_cast(col_scale + g_off * (H / 16) + + swizzled_scale_offset(k + 1, blk, ncols)) = sy; + } + __syncthreads(); + } +#else + NVTE_DEVICE_ERROR("NVFP4 EP direct receive requires SM 10.0+."); +#endif +} + +const int64_t *optional_rng_state(const Tensor *rng_state) { + if (rng_state == nullptr || rng_state->data.dptr == nullptr) return nullptr; + NVTE_CHECK(rng_state->dtype() == DType::kInt64 && rng_state->numel() == 2, + "rng_state must be int64 [2] (seed, offset)."); + const int arch = cuda::sm_arch(); + NVTE_CHECK(arch == 100 || arch == 103 || arch == 107, + "Stochastic rounding to E2M1 requires SM 10.0, 10.3 or 10.7 (got SM ", arch, ")."); + return reinterpret_cast(rng_state->data.dptr); +} + +void check_blackwell() { + NVTE_CHECK(cuda::sm_arch() >= 100, "NVFP4 EP token rows require SM 10.0+."); +} + +} // namespace + +void pack(const Tensor &input, Tensor &packed, const Tensor *rng_state, bool use_fast_math, + cudaStream_t stream) { + check_blackwell(); + NVTE_CHECK(input.dtype() == DType::kBFloat16, "NVFP4 EP pack input must be BF16."); + NVTE_CHECK(packed.dtype() == DType::kByte, "NVFP4 EP packed rows must be kByte."); + NVTE_CHECK(input.dim() == 2 && packed.dim() == 2, "NVFP4 EP pack expects 2D tensors."); + const size_t T = input.flat_first_dim(); + const size_t H = input.flat_last_dim(); + NVTE_CHECK(H % 512 == 0 && H <= 16384, "NVFP4 EP pack requires hidden % 512 == 0 and ", + "hidden <= 16384 (got ", H, ")."); + const size_t row_bytes = packed.flat_last_dim(); + NVTE_CHECK(packed.flat_first_dim() == T && row_bytes % 16 == 0 && row_bytes >= min_row_bytes(H), + "NVFP4 EP packed rows must be [", T, ", >= ", min_row_bytes(H), + " (multiple of 16)], got ", packed.shape(), "."); + const int64_t *rng = optional_rng_state(rng_state); + if (T == 0) return; + const auto *x = reinterpret_cast(input.data.dptr); + auto *out = reinterpret_cast(packed.data.dptr); + auto launch = [&](auto kernel) { + kernel<<(T), static_cast(H / 16), 0, stream>>>( + x, out, static_cast(H), static_cast(row_bytes), rng); + }; + if (use_fast_math) { + rng != nullptr ? launch(pack_kernel) : launch(pack_kernel); + } else { + rng != nullptr ? launch(pack_kernel) : launch(pack_kernel); + } + NVTE_CHECK_CUDA(cudaGetLastError()); +} + +void unpack(const Tensor &packed, const Tensor *counts, Tensor &output, cudaStream_t stream) { + NVTE_CHECK(packed.dtype() == DType::kByte, "NVFP4 EP packed rows must be kByte."); + NVTE_CHECK(output.dtype() == DType::kBFloat16, "NVFP4 EP unpack output must be BF16."); + const size_t R = packed.flat_first_dim(); + const size_t H = output.flat_last_dim(); + const size_t row_bytes = packed.flat_last_dim(); + NVTE_CHECK(H % 64 == 0, "NVFP4 EP unpack requires hidden % 64 == 0 (got ", H, ")."); + NVTE_CHECK(output.flat_first_dim() == R, "NVFP4 EP unpack output must have ", R, " rows."); + NVTE_CHECK(row_bytes % 16 == 0 && row_bytes >= min_row_bytes(H), + "NVFP4 EP packed rows are too short for hidden ", H, "."); + const int64_t *counts_ptr = nullptr; + int num_counts = 0; + if (counts != nullptr && counts->data.dptr != nullptr) { + NVTE_CHECK(counts->dtype() == DType::kInt64, "counts must be int64."); + counts_ptr = reinterpret_cast(counts->data.dptr); + num_counts = static_cast(counts->numel()); + } + if (R == 0) return; + const size_t items = R * (H / 16); + const size_t max_grid = 4 * static_cast(cuda::sm_count()); + const int grid = + static_cast(std::min(DIVUP(items, static_cast(kThreads)), max_grid)); + unpack_kernel<<>>(reinterpret_cast(packed.data.dptr), + counts_ptr, num_counts, static_cast(R), + static_cast(H), static_cast(row_bytes), + reinterpret_cast<__nv_bfloat16 *>(output.data.dptr)); + NVTE_CHECK_CUDA(cudaGetLastError()); +} + +void direct_receive(const Tensor &packed, const Tensor &counts, Tensor &workspace, + uint32_t sign_mask, Tensor &row_data, Tensor &row_scale, Tensor &col_data, + Tensor &col_scale, Tensor &row_amax, Tensor &col_amax, const Tensor *rng_state, + size_t hidden, bool use_fast_math, cudaStream_t stream) { + check_blackwell(); + NVTE_CHECK(packed.dtype() == DType::kByte && packed.dim() == 2, + "NVFP4 EP packed rows must be 2D kByte."); + NVTE_CHECK(counts.dtype() == DType::kInt64, "counts must be int64."); + const int num_groups = static_cast(counts.numel()); + NVTE_CHECK(num_groups >= 1 && num_groups <= kMaxGroups, "NVFP4 EP direct receive supports 1..", + kMaxGroups, " groups (got ", num_groups, ")."); + NVTE_CHECK(workspace.dtype() == DType::kInt32 && workspace.numel() >= num_groups, + "workspace must be int32 with at least ", num_groups, " elements."); + NVTE_CHECK(hidden % kTileK == 0, "NVFP4 EP direct receive requires hidden % ", kTileK, + " == 0 (got ", hidden, ")."); + const size_t capacity = packed.flat_first_dim(); + NVTE_CHECK(capacity % kTileM == 0, "NVFP4 EP direct receive requires the row capacity to be a ", + "multiple of ", kTileM, " (got ", capacity, ")."); + const size_t row_bytes = packed.flat_last_dim(); + NVTE_CHECK(row_bytes % 16 == 0 && row_bytes >= min_row_bytes(hidden), + "NVFP4 EP packed rows are too short for hidden ", hidden, "."); + const size_t data_bytes = capacity * hidden / 2, scale_bytes = capacity * hidden / 16; + for (const Tensor *t : {&row_data, &col_data}) { + NVTE_CHECK(t->dtype() == DType::kByte && t->numel() >= data_bytes, + "NVFP4 EP direct receive data outputs must be kByte with ", data_bytes, " bytes."); + } + for (const Tensor *t : {&row_scale, &col_scale}) { + NVTE_CHECK(t->dtype() == DType::kByte && t->numel() >= scale_bytes, + "NVFP4 EP direct receive scale outputs must be kByte with ", scale_bytes, " bytes."); + } + for (const Tensor *t : {&row_amax, &col_amax}) { + NVTE_CHECK(t->dtype() == DType::kFloat32 && t->numel() >= num_groups, + "NVFP4 EP direct receive amax outputs must be float32 [", num_groups, "]."); + } + const int64_t *rng = optional_rng_state(rng_state); + + const auto *packed_ptr = reinterpret_cast(packed.data.dptr); + const auto *counts_ptr = reinterpret_cast(counts.data.dptr); + auto *gmax = reinterpret_cast(workspace.data.dptr); + NVTE_CHECK_CUDA(cudaMemsetAsync(gmax, 0, num_groups * sizeof(int), stream)); + const int num_sms = cuda::sm_count(); + group_max_kernel<<<2 * num_sms, kThreads, 0, stream>>>( + packed_ptr, counts_ptr, num_groups, static_cast(capacity), + static_cast(row_bytes), static_cast(hidden / 2 + hidden / 16), gmax); + NVTE_CHECK_CUDA(cudaGetLastError()); + auto launch = [&](auto kernel) { + int per_sm = 1; + NVTE_CHECK_CUDA(cudaOccupancyMaxActiveBlocksPerMultiprocessor(&per_sm, kernel, kThreads, 0)); + const int grid = std::max(1, per_sm) * num_sms; + kernel<<>>(packed_ptr, counts_ptr, num_groups, + static_cast(capacity), static_cast(hidden), + static_cast(row_bytes), gmax, sign_mask, rng, + reinterpret_cast(row_data.data.dptr), + reinterpret_cast(row_scale.data.dptr), + reinterpret_cast(col_data.data.dptr), + reinterpret_cast(col_scale.data.dptr), + reinterpret_cast(row_amax.data.dptr), + reinterpret_cast(col_amax.data.dptr)); + }; + if (use_fast_math) { + rng != nullptr ? launch(direct_receive_kernel) + : launch(direct_receive_kernel); + } else { + rng != nullptr ? launch(direct_receive_kernel) + : launch(direct_receive_kernel); + } + NVTE_CHECK_CUDA(cudaGetLastError()); +} + +} // namespace ep_nvfp4 +} // namespace transformer_engine + +void nvte_ep_nvfp4_pack(const NVTETensor input, NVTETensor packed, const NVTETensor rng_state, + int use_fast_math, cudaStream_t stream) { + NVTE_API_CALL(nvte_ep_nvfp4_pack); + using namespace transformer_engine; + ep_nvfp4::pack(*convertNVTETensorCheck(input), *convertNVTETensorCheck(packed), + convertNVTETensor(rng_state), use_fast_math != 0, stream); +} + +void nvte_ep_nvfp4_unpack(const NVTETensor packed, const NVTETensor counts, NVTETensor output, + cudaStream_t stream) { + NVTE_API_CALL(nvte_ep_nvfp4_unpack); + using namespace transformer_engine; + ep_nvfp4::unpack(*convertNVTETensorCheck(packed), convertNVTETensor(counts), + *convertNVTETensorCheck(output), stream); +} + +void nvte_ep_nvfp4_direct_receive(const NVTETensor packed, const NVTETensor counts, + NVTETensor workspace, uint32_t sign_mask, NVTETensor row_data, + NVTETensor row_scale, NVTETensor col_data, NVTETensor col_scale, + NVTETensor row_amax, NVTETensor col_amax, + const NVTETensor rng_state, size_t hidden, int use_fast_math, + cudaStream_t stream) { + NVTE_API_CALL(nvte_ep_nvfp4_direct_receive); + using namespace transformer_engine; + ep_nvfp4::direct_receive(*convertNVTETensorCheck(packed), *convertNVTETensorCheck(counts), + *convertNVTETensorCheck(workspace), sign_mask, + *convertNVTETensorCheck(row_data), *convertNVTETensorCheck(row_scale), + *convertNVTETensorCheck(col_data), *convertNVTETensorCheck(col_scale), + *convertNVTETensorCheck(row_amax), *convertNVTETensorCheck(col_amax), + convertNVTETensor(rng_state), hidden, use_fast_math != 0, stream); +} diff --git a/transformer_engine/common/include/transformer_engine/ep.h b/transformer_engine/common/include/transformer_engine/ep.h index 1cd432a3a0..e492276f0d 100644 --- a/transformer_engine/common/include/transformer_engine/ep.h +++ b/transformer_engine/common/include/transformer_engine/ep.h @@ -309,6 +309,75 @@ void nvte_ep_combine_bwd_v2(NVTETensor handle_mem, NVTETensor grad, NVTECommWind NVTETensor grad_expert_out, NVTECommWindow grad_expert_out_win, const NVTEEpLayerConfig* layer_cfg, cudaStream_t stream); +/* -- NVFP4 token rows ------------------------------------------------------ */ +/* Pack, decode and receive kernels for moving NVFP4 tokens through the + * unscaled EP dispatch / combine-backward path. They do not need NCCL EP and + * are available in every build. Pack and receive need Blackwell (SM 10.0+). + * + * Packed row layout, row_bytes per token (a multiple of 16 and at least + * H/2 + H/16 + 4): [H/2 bytes of E2M1 pairs (low nibble first) | H/16 E4M3 + * block scales (one per 16 elements) | FP32 token decode scale | zero pad]. + * The token decode scale is a power of two: the smallest 2^e with + * amax <= 6 * 448 * 2^e (0 for an all-zero token). */ + +/*! \brief Quantize BF16 tokens to packed NVFP4 rows. + * + * \param[in] input BF16 tokens [T, H]; H % 512 == 0 and H <= 16384. + * \param[out] packed Packed rows, kByte [T, row_bytes]. + * \param[in] rng_state Optional (may be NULL) int64 [2] (seed, offset) on device. + * When given, E2M1 values are stochastically rounded; else RN. + * \param[in] use_fast_math Use the approximate reciprocal, as NVTE_USE_FAST_MATH does for + * TE's NVFP4 casts. + * \param[in] stream CUDA stream. + */ +void nvte_ep_nvfp4_pack(const NVTETensor input, NVTETensor packed, const NVTETensor rng_state, + int use_fast_math, cudaStream_t stream); + +/*! \brief Decode packed NVFP4 rows to BF16. + * + * \param[in] packed Packed rows, kByte [R, row_bytes]. + * \param[in] counts Optional (may be NULL) int64 [G] on device. When given, only the + * first min(sum(counts), R) rows are decoded and the rest of + * ``output`` is left untouched (CUDA-graph safe). + * \param[out] output BF16 [R, H]; H % 64 == 0. + * \param[in] stream CUDA stream. + */ +void nvte_ep_nvfp4_unpack(const NVTETensor packed, const NVTETensor counts, NVTETensor output, + cudaStream_t stream); + +/*! \brief Build the per-expert NVFP4 grouped-GEMM operands straight from packed rows. + * + * Rows are grouped by expert: group g owns counts[g] consecutive rows, each count a multiple + * of 128. The rowwise operand reuses the sender's E2M1 codes; its block scales are re-expressed + * against the group scale G = max token decode scale of the group (exact while they stay in + * E4M3's normal range). The columnwise operand is the 16-point random Hadamard transform along + * tokens, quantized in the same pass. Second-level amaxes are row_amax[g] = 6 * 448 * G and + * col_amax[g] = 4 * 6 * 448 * G (a bound of the post-RHT amax, so no amax pass is needed). + * Scales use the GEMM-swizzled layout. Bounded on device by ``counts`` (CUDA-graph safe). + * + * \param[in] packed Packed rows (see above), kByte [R, row_bytes]; R % 128 == 0. + * \param[in] counts int64 [G] on device, 1 <= G <= 128. + * \param[in,out] workspace int32 [G] scratch. + * \param[in] sign_mask RHT random signs: bit i set when sign i is -1. + * \param[out] row_data kByte [R * H / 2]. + * \param[out] row_scale kByte [R * H / 16]. + * \param[out] col_data kByte [R * H / 2]; group g is [H, counts[g]]. + * \param[out] col_scale kByte [R * H / 16]. + * \param[out] row_amax float32 [G]. + * \param[out] col_amax float32 [G]. + * \param[in] rng_state Optional (may be NULL) int64 [2] (seed, offset) on device: + * stochastic rounding of the columnwise operand when given. + * \param[in] hidden H; H % 256 == 0. + * \param[in] use_fast_math Use the approximate reciprocal. + * \param[in] stream CUDA stream. + */ +void nvte_ep_nvfp4_direct_receive(const NVTETensor packed, const NVTETensor counts, + NVTETensor workspace, uint32_t sign_mask, NVTETensor row_data, + NVTETensor row_scale, NVTETensor col_data, NVTETensor col_scale, + NVTETensor row_amax, NVTETensor col_amax, + const NVTETensor rng_state, size_t hidden, int use_fast_math, + cudaStream_t stream); + #ifdef __cplusplus } #endif diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index 02bf677344..1128f7afa2 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -824,6 +824,17 @@ void ep_combine_bwd(at::Tensor handle_mem, at::Tensor grad, at::Tensor grad_expe // Registers the EP pybind functions on `m`. Defined under NVTE_WITH_NCCL_EP. void register_ep_bindings(pybind11::module_ &m); +// NVFP4 token rows for the EP dispatch / combine backward (independent of NCCL EP). +at::Tensor ep_nvfp4_pack(const at::Tensor &tokens, bool stochastic); + +at::Tensor ep_nvfp4_unpack(const at::Tensor &packed, int64_t hidden, + std::optional counts, std::optional out); + +// Returns {row_data, row_scale, col_data, col_scale, row_amax, col_amax}. +std::vector ep_nvfp4_direct_receive(const at::Tensor &packed, int64_t hidden, + const at::Tensor &counts, int64_t sign_mask, + bool stochastic); + /*************************************************************************************************** * NVSHMEM APIs **************************************************************************************************/ diff --git a/transformer_engine/pytorch/csrc/extensions/ep_nvfp4.cpp b/transformer_engine/pytorch/csrc/extensions/ep_nvfp4.cpp new file mode 100644 index 0000000000..086cd52cd8 --- /dev/null +++ b/transformer_engine/pytorch/csrc/extensions/ep_nvfp4.cpp @@ -0,0 +1,136 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +// NVFP4 token rows for the EP dispatch and combine backward. Independent of NCCL EP. + +#include +#include + +#include +#include + +#include "../common.h" +#include "../extensions.h" +#include "../pybind.h" +#include "common/util/system.h" +#include "transformer_engine/ep.h" + +namespace transformer_engine::pytorch { + +namespace { + +int64_t nvfp4_row_bytes(int64_t hidden) { + return static_cast(roundup(static_cast(hidden / 2 + hidden / 16 + 4), 16)); +} + +int use_fast_math() { return transformer_engine::getenv("NVTE_USE_FAST_MATH") ? 1 : 0; } + +// (seed, offset) from torch's CUDA generator; under graph capture the offset is read on device, +// so every replay draws fresh bits. +at::Tensor make_rng_state(size_t elts_per_thread) { + auto gen = at::get_generator_or_default( + std::nullopt, at::cuda::detail::getDefaultCUDAGenerator()); + at::PhiloxCudaState philox_args = init_philox_state(gen, elts_per_thread); + auto rng_state = at::empty({2}, at::TensorOptions().dtype(at::kLong).device(at::kCUDA)); + philox_unpack(philox_args, static_cast(rng_state.data_ptr())); + return rng_state; +} + +} // namespace + +at::Tensor ep_nvfp4_pack(const at::Tensor &tokens, bool stochastic) { + NVTE_CHECK(tokens.is_cuda() && tokens.dim() == 2, "tokens must be a 2D CUDA tensor."); + // Launch, allocate and draw RNG state on the input's device, not the current one. + const at::cuda::CUDAGuard device_guard(tokens.device()); + const auto x = tokens.contiguous(); + auto packed = at::empty({x.size(0), nvfp4_row_bytes(x.size(1))}, x.options().dtype(at::kByte)); + std::optional rng_state; + TensorWrapper rng_state_nvte; + if (stochastic) { + rng_state = make_rng_state(4); + rng_state_nvte = makeTransformerEngineTensor(*rng_state); + } + auto x_nvte = makeTransformerEngineTensor(x); + auto packed_nvte = makeTransformerEngineTensor(packed); + NVTE_SCOPED_GIL_RELEASE({ + nvte_ep_nvfp4_pack(x_nvte.data(), packed_nvte.data(), rng_state_nvte.data(), use_fast_math(), + at::cuda::getCurrentCUDAStream()); + }); + return packed; +} + +at::Tensor ep_nvfp4_unpack(const at::Tensor &packed, int64_t hidden, + std::optional counts, std::optional out) { + NVTE_CHECK(packed.is_cuda() && packed.dim() == 2 && packed.scalar_type() == at::kByte, + "packed must be a 2D uint8 CUDA tensor."); + const at::cuda::CUDAGuard device_guard(packed.device()); + const auto p = packed.contiguous(); + if (!out) { + out = at::empty({p.size(0), hidden}, p.options().dtype(at::kBFloat16)); + } + NVTE_CHECK(out->is_contiguous() && out->numel() == p.size(0) * hidden, + "out must be a contiguous [", p.size(0), ", ", hidden, "] tensor."); + auto out_2d = out->view({p.size(0), hidden}); + TensorWrapper counts_nvte; + if (counts) { + counts = counts->contiguous(); + counts_nvte = makeTransformerEngineTensor(*counts); + } + auto p_nvte = makeTransformerEngineTensor(p); + auto out_nvte = makeTransformerEngineTensor(out_2d); + NVTE_SCOPED_GIL_RELEASE({ + nvte_ep_nvfp4_unpack(p_nvte.data(), counts_nvte.data(), out_nvte.data(), + at::cuda::getCurrentCUDAStream()); + }); + return *out; +} + +std::vector ep_nvfp4_direct_receive(const at::Tensor &packed, int64_t hidden, + const at::Tensor &counts, int64_t sign_mask, + bool stochastic) { + NVTE_CHECK(packed.is_cuda() && packed.dim() == 2 && packed.scalar_type() == at::kByte, + "packed must be a 2D uint8 CUDA tensor."); + NVTE_CHECK(counts.scalar_type() == at::kLong, "counts must be int64."); + const at::cuda::CUDAGuard device_guard(packed.device()); + const auto p = packed.contiguous(); + const auto c = counts.contiguous(); + const int64_t capacity = p.size(0); + const int64_t num_groups = c.numel(); + const auto u8 = p.options().dtype(at::kByte); + const auto f32 = p.options().dtype(at::kFloat); + auto row_data = at::empty({capacity * hidden / 2}, u8); + auto row_scale = at::empty({capacity * hidden / 16}, u8); + auto col_data = at::empty({capacity * hidden / 2}, u8); + auto col_scale = at::empty({capacity * hidden / 16}, u8); + auto row_amax = at::empty({num_groups}, f32); + auto col_amax = at::empty({num_groups}, f32); + auto workspace = at::empty({num_groups}, p.options().dtype(at::kInt)); + std::optional rng_state; + TensorWrapper rng_state_nvte; + if (stochastic) { + rng_state = make_rng_state(32); + rng_state_nvte = makeTransformerEngineTensor(*rng_state); + } + auto p_nvte = makeTransformerEngineTensor(p); + auto c_nvte = makeTransformerEngineTensor(c); + auto workspace_nvte = makeTransformerEngineTensor(workspace); + auto row_data_nvte = makeTransformerEngineTensor(row_data); + auto row_scale_nvte = makeTransformerEngineTensor(row_scale); + auto col_data_nvte = makeTransformerEngineTensor(col_data); + auto col_scale_nvte = makeTransformerEngineTensor(col_scale); + auto row_amax_nvte = makeTransformerEngineTensor(row_amax); + auto col_amax_nvte = makeTransformerEngineTensor(col_amax); + NVTE_SCOPED_GIL_RELEASE({ + nvte_ep_nvfp4_direct_receive( + p_nvte.data(), c_nvte.data(), workspace_nvte.data(), static_cast(sign_mask), + row_data_nvte.data(), row_scale_nvte.data(), col_data_nvte.data(), col_scale_nvte.data(), + row_amax_nvte.data(), col_amax_nvte.data(), rng_state_nvte.data(), + static_cast(hidden), use_fast_math(), at::cuda::getCurrentCUDAStream()); + }); + return {row_data, row_scale, col_data, col_scale, row_amax, col_amax}; +} + +} // namespace transformer_engine::pytorch diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp index df6d16782d..540b821d0d 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -372,6 +372,15 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { #ifdef NVTE_WITH_NCCL_EP transformer_engine::pytorch::register_ep_bindings(m); #endif // NVTE_WITH_NCCL_EP + m.def("ep_nvfp4_pack", &transformer_engine::pytorch::ep_nvfp4_pack, + "Pack BF16 tokens into NVFP4 EP rows (power-of-two token scales)", py::arg("tokens"), + py::arg("stochastic")); + m.def("ep_nvfp4_unpack", &transformer_engine::pytorch::ep_nvfp4_unpack, + "Decode NVFP4 EP rows to BF16", py::arg("packed"), py::arg("hidden"), + py::arg("counts") = std::nullopt, py::arg("out") = std::nullopt); + m.def("ep_nvfp4_direct_receive", &transformer_engine::pytorch::ep_nvfp4_direct_receive, + "Build per-expert NVFP4 GEMM operands from NVFP4 EP rows", py::arg("packed"), + py::arg("hidden"), py::arg("counts"), py::arg("sign_mask"), py::arg("stochastic")); // Permutation functions m.def("moe_permute_fwd", transformer_engine::pytorch::moe_permute_fwd, "MOE permute FWD", diff --git a/transformer_engine/pytorch/ep.py b/transformer_engine/pytorch/ep.py index cf4c7b1e51..80705bddbf 100644 --- a/transformer_engine/pytorch/ep.py +++ b/transformer_engine/pytorch/ep.py @@ -332,6 +332,7 @@ class EpBuffer: "total_recv_tokens", "dispatch_fwd_quant_recipe", "combine_bwd_quant_recipe", + "nvfp4_packed_scratch", ) def __init__( @@ -374,6 +375,8 @@ def __init__( self.zero_copy = bool(tex.ep_get_zero_copy()) self.dispatch_fwd_quant_recipe = dispatch_fwd_quant_recipe self.combine_bwd_quant_recipe = combine_bwd_quant_recipe + # Zero-copy landing for NVFP4 rows that are decoded to BF16 (see _nvfp4_landing). + self.nvfp4_packed_scratch = None size_bytes = tex.ep_handle_mem_size(self.top_k, self.alignment) self.handle_mem = torch.empty(int(size_bytes), dtype=torch.uint8, device=device) @@ -642,6 +645,162 @@ def _ep_combine_raw(buffer: "EpBuffer", expert_out: torch.Tensor, result: torch. ) +# NVFP4 token rows +# +# One packed row per token: [H/2 bytes of E2M1 pairs | H/16 E4M3 block scales | FP32 token +# scale], zero-padded to 16 bytes. The second-level scale is per token (a power of two) instead of +# per tensor, so every token decodes on its own and ranks need no scale agreement. The rows travel +# as an unscaled payload (a BF16 view: dispatch and combine backward only copy bytes). The receiver +# either decodes them to BF16 or builds the expert GEMM's NVFP4 GroupedTensor directly: the +# rowwise operand reuses the sender's E2M1 codes and the columnwise RHT operand is quantized in +# the same pass (see transformer_engine/common/ep/nvfp4_ep_rows.cu). + +_NVFP4_DIRECT_RECEIVE_ALIGNMENT = 128 +_NVFP4_DIRECT_RECEIVE_MAX_EXPERTS = 128 + + +def _nvfp4_row_bytes(hidden: int) -> int: + """Bytes of one packed NVFP4 token row, padded to the 16-byte EP row alignment.""" + return (hidden // 2 + hidden // 16 + 4 + 15) // 16 * 16 + + +def _is_nvfp4_recipe(recipe) -> bool: + if recipe is None: + return False + from ..common.recipe import NVFP4BlockScaling + + return isinstance(recipe, NVFP4BlockScaling) + + +def _check_quant_recipe(recipe, what: str) -> None: + """EP quantizes on the wire with MXFP8BlockScaling or NVFP4BlockScaling.""" + from ..common.recipe import MXFP8BlockScaling + + if not isinstance(recipe, MXFP8BlockScaling) and not _is_nvfp4_recipe(recipe): + raise NotImplementedError( + f"{what} supports MXFP8BlockScaling and NVFP4BlockScaling only; got " + f"{type(recipe).__name__}." + ) + if _is_nvfp4_recipe(recipe): + from ..common.recipe import Format + + if recipe.fp8_format != Format.E4M3: + raise NotImplementedError(f"{what} with NVFP4 supports E4M3 block scales only.") + + +def _check_nvfp4_hidden(hidden: int) -> None: + if hidden % 512 != 0 or hidden > 16384: + raise ValueError( + f"NVFP4 EP quantization requires hidden % 512 == 0 and hidden <= 16384, got {hidden}." + ) + + +def _nvfp4_gemm_recipe(buffer: "EpBuffer", rows: int, grad: bool): + """The active NVFP4 autocast recipe when the receiver can hand the expert GEMMs an NVFP4 + GroupedTensor directly, else None (the receiver decodes to BF16). ``grad`` selects the + gradient quantization (combine backward) instead of the forward input (dispatch). + + Needs graph mode (static recv capacity), NVFP4 autocast with the RHT columnwise operand that + the direct receive produces, per-expert groups aligned to 128 rows, and at most 128 local + experts. Evaluated in the forward pass, where the autocast state is visible.""" + from .quantization import FP8GlobalStateManager + from ..common.recipe import Format + + if buffer.eager or not FP8GlobalStateManager.is_fp8_enabled(): + return None + recipe = FP8GlobalStateManager.get_fp8_recipe() + if not recipe.nvfp4() or recipe.fp8_format != Format.E4M3: + return None + qparams = recipe.fp4_quant_bwd_grad if grad else recipe.fp4_quant_fwd_inp + if not qparams.random_hadamard_transform or qparams.fp4_2d_quantization: + return None + if not grad and (recipe.row_scaled_activation or recipe.nvfp4_4over6 in ("activations", "all")): + return None + if ( + buffer.alignment < _NVFP4_DIRECT_RECEIVE_ALIGNMENT + or buffer.num_local_experts > _NVFP4_DIRECT_RECEIVE_MAX_EXPERTS + or rows % _NVFP4_DIRECT_RECEIVE_ALIGNMENT != 0 + ): + return None + return recipe + + +def _nvfp4_landing( + buf, rows: int, row_bytes: int, device, zero_copy: bool, buffer: Optional["EpBuffer"] = None +) -> torch.Tensor: + """Receive buffer for ``rows`` packed NVFP4 rows, as the BF16 view the EP ops copy. Carved from + a caller buffer when given (so a persistent or symm-mem buffer backs it). Otherwise, under + zero-copy with ``buffer`` given (the rows are decoded to BF16 right after the receive), a + persistent symm-mem scratch kept on ``buffer``: allocated on first use outside capture and + reused after, so a captured decode path allocates nothing from the symm-mem pool. Else + allocated.""" + if buf is None: + if not zero_copy or buffer is None: + return _alloc_io((rows, row_bytes // 2), torch.bfloat16, device, zero_copy) + need = rows * row_bytes + scratch = buffer.nvfp4_packed_scratch + if scratch is None or scratch.numel() < need: + capacity = max(rows, buffer.recv_capacity_per_rank or 0) * row_bytes + scratch = _alloc_io((capacity,), torch.uint8, device, zero_copy) + buffer.nvfp4_packed_scratch = scratch + return scratch[:need].view(torch.bfloat16).view(rows, row_bytes // 2) + if not buf.is_contiguous(): + raise ValueError("NVFP4 receive buffer must be contiguous.") + flat = buf.view(-1).view(torch.uint8) + need = rows * row_bytes + if flat.numel() < need: + raise ValueError(f"NVFP4 receive buffer too small: need {need} bytes, got {flat.numel()}.") + return flat[:need].view(torch.bfloat16).view(rows, row_bytes // 2) + + +def _nvfp4_direct_receive( + recv_packed: torch.Tensor, + hidden: int, + counts: torch.Tensor, + fake_dtype: torch.dtype, + stochastic: bool, +): + """Per-expert NVFP4 GroupedTensor (rowwise + columnwise RHT operands, GEMM-swizzled scales) + built in one pass from packed rows (uint8 [capacity, row_bytes]) grouped by ``counts``. + Second-level scales are per expert: rowwise ``6 * 448 * G`` and columnwise ``4 * 6 * 448 * G`` + with ``G`` the largest token scale of the expert. ``stochastic`` rounds the columnwise + operand stochastically (gradients).""" + from .tensor.grouped_tensor import GroupedTensor + from .tensor.nvfp4_tensor import NVFP4Quantizer, get_random_sign_mask_for_rht + + row_data, row_scale, col_data, col_scale, row_amax, col_amax = tex.ep_nvfp4_direct_receive( + recv_packed, + hidden, + counts, + get_random_sign_mask_for_rht(True, recv_packed.device.index), + stochastic, + ) + quantizer = NVFP4Quantizer( + rowwise=True, + columnwise=True, + with_rht=True, + with_post_rht_amax=True, + stochastic_rounding=stochastic, + ) + quantizer.optimize_for_gemm = True + quantizer.internal = True + return GroupedTensor( + shape=(recv_packed.shape[0], hidden), + dtype=fake_dtype, + num_tensors=counts.numel(), + quantizer=quantizer, + data=row_data, + columnwise_data=col_data, + scale_inv=row_scale, + columnwise_scale_inv=col_scale, + amax=row_amax, + columnwise_amax=col_amax, + first_dims=counts, + tensor_offsets=tex.splits_to_offsets(counts, hidden), + with_gemm_swizzled_scales=True, + ) + + # autograd.Function wrappers @@ -677,6 +836,10 @@ def _ep_prepare_and_dispatch_fwd( output (a per-expert GroupedTensor for MXFP8, else the raw recv tokens), the recv topk weights, and a _DispatchState for backward. No autograd; the caller owns context handling.""" _require_bf16("dispatch input", tokens) + if _is_nvfp4_recipe(buffer.dispatch_fwd_quant_recipe): + return _ep_prepare_and_dispatch_nvfp4_fwd( + tokens, topk_weights, topk_idx, buffer, recv_tokens, recv_topk_weights + ) tokens_scale_inv = None if buffer.dispatch_fwd_quant_recipe is not None: from .tensor.mxfp8_tensor import MXFP8Quantizer @@ -784,6 +947,98 @@ def _ep_prepare_and_dispatch_fwd( return recv_tokens, recv_topk_weights, state +def _ep_prepare_and_dispatch_nvfp4_fwd( + tokens: torch.Tensor, + topk_weights: torch.Tensor, + topk_idx: torch.Tensor, + buffer: "EpBuffer", + recv_tokens: Optional[torch.Tensor], + recv_topk_weights: Optional[torch.Tensor], +): + """NVFP4 dispatch: pack each token into an NVFP4 row and move the rows as an unscaled payload. + Under NVFP4 autocast (graph mode, see ``_nvfp4_gemm_recipe``) the receiver returns the expert + GEMM's NVFP4 GroupedTensor, with the packed rows landing in ``recv_tokens`` when given; + otherwise it decodes to BF16 recv tokens (into ``recv_tokens`` when given). The backward is a + straight-through BF16 gradient.""" + tokens_2d = tokens.reshape(-1, tokens.shape[-1]) + hidden = tokens_2d.shape[-1] + packed = tex.ep_nvfp4_pack(tokens_2d, False) + row_bytes = packed.shape[-1] + handle_mem = buffer.handle_mem + tokens_per_expert = buffer.tokens_per_expert + direct = False + if buffer.eager: + recv_packed, recv_topk_weights = tex.ep_prepare_and_dispatch( + handle_mem, + topk_idx, + packed.view(torch.bfloat16), + topk_weights, + tokens_per_expert, + buffer.total_recv_tokens, + buffer.top_k, + buffer.alignment, + ) + recv_out = tex.ep_nvfp4_unpack(recv_packed.view(torch.uint8), hidden) + else: + device = tokens_2d.device + num_recv_tokens = buffer.recv_capacity_per_rank + gemm_recipe = _nvfp4_gemm_recipe(buffer, num_recv_tokens, grad=False) + direct = gemm_recipe is not None + recv_packed = _nvfp4_landing( + recv_tokens if direct else None, + num_recv_tokens, + row_bytes, + device, + buffer.zero_copy, + buffer=None if direct else buffer, + ) + if recv_topk_weights is None: + recv_topk_weights = _alloc_io( + (num_recv_tokens,), torch.float32, device, buffer.zero_copy + ) + torch.ops.transformer_engine_ep.prepare_and_dispatch( + handle_mem, + topk_idx, + packed.view(torch.bfloat16), + topk_weights, + tokens_per_expert, + buffer.total_recv_tokens, + buffer.top_k, + buffer.alignment, + recv_packed, + recv_topk_weights, + None, + None, + ) + if direct: + recv_out = _nvfp4_direct_receive( + recv_packed.view(torch.uint8), + hidden, + tokens_per_expert, + tokens.dtype, + stochastic=False, + ) + else: + # Decode only the live rows (device-side count, so CUDA graph replays see it). + recv_out = tex.ep_nvfp4_unpack( + recv_packed.view(torch.uint8), hidden, tokens_per_expert, recv_tokens + ) + state = _DispatchState( + handle_mem=handle_mem, + tokens_shape=tokens.shape, + topk_weights_shape=topk_weights.shape, + tokens_T_flat=tokens_2d.shape[0], + topk_T_flat=topk_weights.numel() // topk_weights.shape[-1], + top_k=topk_weights.shape[-1], + layer_args=(buffer.top_k, buffer.alignment, topk_idx.dtype), + hidden_dim=hidden, + eager=buffer.eager, + # The GroupedTensor is returned like the MXFP8 one (not detached). + is_scaled=direct, + ) + return recv_out, recv_topk_weights, state + + def _ep_dispatch_bwd( state: "_DispatchState", g_recv_tokens: torch.Tensor, @@ -902,6 +1157,12 @@ class _CombineState: device: torch.device eager: bool zero_copy: bool + # NVFP4 combine backward: the NVFP4 autocast recipe when the backward returns FC2's + # grad_output as a GroupedTensor (decided in the forward, where autocast is visible). + nvfp4_gemm_recipe: object = None + # NVFP4 combine backward: the EpBuffer whose persistent zero-copy scratch receives the packed + # rows when they are decoded to BF16 (see _nvfp4_landing). + nvfp4_buffer: object = None def _ep_combine_fwd( @@ -910,6 +1171,7 @@ def _ep_combine_fwd( buffer: "EpBuffer", num_local_tokens: int, bwd_quant_recipe, + nvfp4_grad_direct_receive: bool = True, ): """Run combine and return ``(result, _CombineState)``. Eager mode is not graph-capturable, so it calls the backend op directly and skips the torch.library dispatch. No autograd; the caller @@ -936,19 +1198,25 @@ def _ep_combine_fwd( eager=eager, zero_copy=buffer.zero_copy, ) + if _is_nvfp4_recipe(bwd_quant_recipe): + if nvfp4_grad_direct_receive: + state.nvfp4_gemm_recipe = _nvfp4_gemm_recipe(buffer, expert_out.shape[0], grad=True) + state.nvfp4_buffer = buffer return result, state def _ep_combine_bwd(state: "_CombineState", g_result: torch.Tensor): """Scatter the result-grad to expert positions and return the expert_out grad. High-precision - sends the grad as-is; a quantized recipe (MXFP8 today) quantizes it and returns a per-expert - GroupedTensor. No autograd; the caller owns context handling.""" + sends the grad as-is; MXFP8 quantizes it and returns a per-expert GroupedTensor; NVFP4 see + ``_ep_combine_bwd_nvfp4``. No autograd; the caller owns context handling.""" _require_bf16("combine backward grad_output", g_result) if not g_result.is_contiguous(): g_result = g_result.contiguous() handle_mem = state.handle_mem - if state.bwd_quant_recipe is None: + if _is_nvfp4_recipe(state.bwd_quant_recipe): + grad_expert_out = _ep_combine_bwd_nvfp4(state, g_result) + elif state.bwd_quant_recipe is None: grad_expert_out = state.grad_out if grad_expert_out is None: grad_expert_out = _alloc_io( @@ -1009,6 +1277,63 @@ def _ep_combine_bwd(state: "_CombineState", g_result: torch.Tensor): return grad_expert_out +def _ep_combine_bwd_nvfp4(state: "_CombineState", g_result: torch.Tensor): + """NVFP4 combine backward: pack the result-grad rows to NVFP4 (stochastic rounding per the + recipe) and scatter them as an unscaled payload. Under NVFP4 autocast the receiver returns FC2's + grad_output as an NVFP4 GroupedTensor, with the packed rows landing in ``grad_out`` when given; + otherwise it decodes them to the BF16 expert_out grad (into ``grad_out`` when given), bounded + by the device-side per-expert counts so CUDA graph replays see the live count.""" + recv_rows, hidden = state.expert_out_shape[0], state.expert_out_shape[-1] + stochastic = state.bwd_quant_recipe.fp4_quant_bwd_grad.stochastic_rounding + packed = tex.ep_nvfp4_pack(g_result.reshape(-1, hidden), stochastic) + gemm_recipe = state.nvfp4_gemm_recipe + ge_packed = _nvfp4_landing( + state.grad_out if gemm_recipe is not None else None, + recv_rows, + packed.shape[-1], + state.device, + state.zero_copy, + buffer=None if gemm_recipe is not None else state.nvfp4_buffer, + ) + if state.eager: + tex.ep_combine_bwd( + state.handle_mem, + packed.view(torch.bfloat16), + ge_packed, + *state.layer_args, + None, + None, + ) + else: + torch.ops.transformer_engine_ep.combine_bwd( + state.handle_mem, + packed.view(torch.bfloat16), + ge_packed, + *state.layer_args, + None, + None, + ) + if gemm_recipe is not None: + return _nvfp4_direct_receive( + ge_packed.view(torch.uint8), + hidden, + state.token_counts, + state.expert_out_dtype, + stochastic=gemm_recipe.fp4_quant_bwd_grad.stochastic_rounding, + ) + grad_expert_out = state.grad_out + if grad_expert_out is None: + grad_expert_out = torch.empty( + state.expert_out_shape, dtype=state.expert_out_dtype, device=state.device + ) + return tex.ep_nvfp4_unpack( + ge_packed.view(torch.uint8), + hidden, + None if state.eager else state.token_counts, + grad_expert_out.view(-1, hidden), + ).view(state.expert_out_shape) + + class _EpCombine(torch.autograd.Function): """Autograd combine; bwd scatters the expert_out grad into ``grad_out``. When the caller supplies it that buffer is used as-is; otherwise it is allocated in the backward from the @@ -1027,11 +1352,17 @@ def forward( # type: ignore[override] buffer: "EpBuffer", num_local_tokens: int, bwd_quant_recipe=None, + nvfp4_grad_direct_receive: bool = True, ): """Combine fwd; stashes the backward state on ctx. When ``bwd_quant_recipe`` is set, the - backward sends the result-grad as MXFP8.""" + backward sends the result-grad as MXFP8 or NVFP4.""" result, ctx.state = _ep_combine_fwd( - expert_out, grad_out, buffer, num_local_tokens, bwd_quant_recipe + expert_out, + grad_out, + buffer, + num_local_tokens, + bwd_quant_recipe, + nvfp4_grad_direct_receive, ) return result @@ -1045,6 +1376,7 @@ def backward(ctx, g_result): # type: ignore[override] None, # buffer None, # num_local_tokens None, # bwd_quant_recipe + None, # nvfp4_grad_direct_receive ) @@ -1213,16 +1545,23 @@ def ep_dispatch( ): """Prepare + dispatch with autograd. ``tokens`` is bfloat16; ``topk_idx`` is int32 or int64. - When the buffer's ``dispatch_fwd_quant_recipe`` is set (``MXFP8BlockScaling`` only for now), tokens - are quantized internally and recv is returned as a per-expert ``GroupedTensor``; otherwise recv - stays bfloat16. A pre-quantized ``tokens`` is not accepted. + When the buffer's ``dispatch_fwd_quant_recipe`` is ``MXFP8BlockScaling``, tokens are quantized + internally and recv is returned as a per-expert ``GroupedTensor``; otherwise recv stays + bfloat16. With ``NVFP4BlockScaling`` tokens travel as packed NVFP4 rows (hidden % 512 == 0); + under NVFP4 autocast in graph mode, with an alignment of at least 128 and at most 128 local + experts, recv is the expert GEMM's NVFP4 ``GroupedTensor`` (rowwise and columnwise RHT operands, + built in one pass from the received rows), else the decoded bfloat16 tokens. A pre-quantized + ``tokens`` is not accepted. ``recv_tokens`` / ``recv_topk_weights`` are the recv outputs: pass caller-owned buffers (symm-mem-backed under zero-copy) or leave them None to allocate. For MXFP8 the recv data and scales share ``recv_tokens`` (data then scales), so size it to at least - ``recv_capacity_per_rank * (hidden + hidden/block)`` bytes. Eager mode sizes the recv outputs + ``recv_capacity_per_rank * (hidden + hidden/block)`` bytes; an NVFP4 ``GroupedTensor`` recv + lands the packed rows in it. Eager mode sizes the recv outputs per step and forbids caller-supplied buffers. Under zero-copy, leaving them None allocates from the symm-mem pool, which is not CUDA-graph capturable; pass persistent buffers to capture a graph. + When NVFP4 recv is decoded to bfloat16 under zero-copy, the packed rows land in a persistent + symm-mem scratch on ``buffer``, allocated by the first call, so run one step before capture. Returns (recv_tokens, recv_topk_weights, tokens_per_expert); tokens_per_expert is non-diff. See ``buffer.total_recv_tokens`` for the per-step recv total. @@ -1245,13 +1584,9 @@ def ep_dispatch( ) if buffer.dispatch_fwd_quant_recipe is not None: - from ..common.recipe import MXFP8BlockScaling - - if not isinstance(buffer.dispatch_fwd_quant_recipe, MXFP8BlockScaling): - raise NotImplementedError( - "EP block-scaled dispatch supports MXFP8BlockScaling only; got " - f"{type(buffer.dispatch_fwd_quant_recipe).__name__}." - ) + _check_quant_recipe(buffer.dispatch_fwd_quant_recipe, "EP block-scaled dispatch") + if _is_nvfp4_recipe(buffer.dispatch_fwd_quant_recipe): + _check_nvfp4_hidden(tokens.shape[-1]) # Fused prepare + dispatch in one C++ op. Eager sizes the recv outputs from the per-step host # recv-count (allocated in C++, no caller buffers); non-eager sizes them to the static recv @@ -1273,16 +1608,23 @@ def ep_combine( *, num_local_tokens: Optional[int] = None, grad_out: Optional[torch.Tensor] = None, + nvfp4_grad_direct_receive: bool = True, ): """Combine with autograd; caller pre-applies topk weighting. ``expert_out`` is the combine input (symm-mem-backed under zero-copy). ``grad_out`` is the backward's grad target: pass a caller-owned buffer or leave it None to allocate. For MXFP8 the grad data and scales share ``grad_out`` (data then scales), so size it to at least - ``recv_capacity_per_rank * (hidden + hidden/block)`` bytes (non-zero-copy only). Eager mode sizes - the grad target per step and forbids a caller-supplied buffer. Under zero-copy, leaving it None + ``recv_capacity_per_rank * (hidden + hidden/block)`` bytes (non-zero-copy only). For NVFP4 the + backward returns FC2's grad_output as an NVFP4 ``GroupedTensor`` under NVFP4 autocast (see + ``ep_dispatch``; the packed rows land in ``grad_out``), else a bfloat16 grad decoded into + ``grad_out``. The NVFP4 ``GroupedTensor`` holds only GEMM operands (swizzled scales), so a + consumer that needs the decoded gradient, such as a grouped linear with a bias (bias gradient), + must pass ``nvfp4_grad_direct_receive=False`` to get the bfloat16 grad instead. Eager mode + sizes the grad target per step and forbids a caller-supplied buffer. Under zero-copy, leaving it None allocates from the symm-mem pool, which is not CUDA-graph capturable; pass a persistent buffer to - capture a graph. Result shape is (num_local_tokens, hidden_dim); num_local_tokens defaults to + capture a graph (the NVFP4 bfloat16 decode path also needs one step before capture, see + ``ep_dispatch``). Result shape is (num_local_tokens, hidden_dim); num_local_tokens defaults to buffer.max_tokens_per_rank. """ _require_bf16("expert_out", expert_out) @@ -1294,16 +1636,13 @@ def ep_combine( if num_local_tokens is None: num_local_tokens = buffer.max_tokens_per_rank # When combine_bwd_quant_recipe is set the combine backward sends the result-grad over the - # wire as MXFP8 and returns the expert_out grad as a GroupedTensor. + # wire quantized: MXFP8 returns the expert_out grad as a GroupedTensor; NVFP4 returns it as an + # NVFP4 GroupedTensor under NVFP4 autocast and as BF16 otherwise. bwd_quant_recipe = None if buffer.combine_bwd_quant_recipe is not None: - from ..common.recipe import MXFP8BlockScaling - - if not isinstance(buffer.combine_bwd_quant_recipe, MXFP8BlockScaling): - raise NotImplementedError( - "EP combine backward supports MXFP8BlockScaling only; got " - f"{type(buffer.combine_bwd_quant_recipe).__name__}." - ) + _check_quant_recipe(buffer.combine_bwd_quant_recipe, "EP combine backward") + if _is_nvfp4_recipe(buffer.combine_bwd_quant_recipe): + _check_nvfp4_hidden(expert_out.shape[-1]) bwd_quant_recipe = buffer.combine_bwd_quant_recipe return _EpCombine.apply( expert_out, @@ -1311,4 +1650,5 @@ def ep_combine( buffer, num_local_tokens, bwd_quant_recipe, + nvfp4_grad_direct_receive, )