From f938cee374fe0e14c219f13e301eba569183ed75 Mon Sep 17 00:00:00 2001 From: Siddhartha Raman S Date: Tue, 27 Oct 2026 11:57:59 -0700 Subject: [PATCH 1/6] [Common][PyTorch] Add NVFP4 EP dispatch and combine backward with direct receive Accept NVFP4BlockScaling as the EP dispatch_fwd_quant_recipe and combine_bwd_quant_recipe. Tokens travel as packed NVFP4 rows (E2M1 codes, E4M3 block scales and a power-of-two FP32 token scale) through the unscaled dispatch / combine-backward path. Under NVFP4 autocast in graph mode, the receiver builds the expert GEMM's NVFP4 GroupedTensor straight from the rows in one pass: the rowwise operand reuses the sender's codes and the columnwise RHT operand is quantized with a bound-derived per-expert scale, so no BF16 round trip or amax pass is needed. The combine backward uses stochastic rounding per the recipe. Otherwise the rows are decoded to BF16. Add nvte_ep_nvfp4_pack / _unpack / _direct_receive to the common EP API, the tex bindings, single-GPU tests and multi-GPU EP harness tests. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Siddhartha Raman S --- tests/pytorch/distributed/run_ep.py | 374 ++++++++- tests/pytorch/nvfp4/test_nvfp4_ep_rows.py | 561 ++++++++++++++ transformer_engine/common/CMakeLists.txt | 1 + transformer_engine/common/ep/nvfp4_ep_rows.cu | 715 ++++++++++++++++++ .../common/include/transformer_engine/ep.h | 69 ++ transformer_engine/pytorch/csrc/extensions.h | 11 + .../pytorch/csrc/extensions/ep_nvfp4.cpp | 131 ++++ .../pytorch/csrc/extensions/pybind.cpp | 9 + transformer_engine/pytorch/ep.py | 333 +++++++- 9 files changed, 2178 insertions(+), 26 deletions(-) create mode 100644 tests/pytorch/nvfp4/test_nvfp4_ep_rows.py create mode 100644 transformer_engine/common/ep/nvfp4_ep_rows.cu create mode 100644 transformer_engine/pytorch/csrc/extensions/ep_nvfp4.cpp diff --git a/tests/pytorch/distributed/run_ep.py b/tests/pytorch/distributed/run_ep.py index 6ab5f901f73..7002c0efd1d 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,378 @@ 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 + ) + + 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.""" + if EAGER: + self.skipTest("capture requires non-eager static recv capacity") + self._require_nvfp4_shapes() + recv_buf, rw_buf = 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) + 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) + + @_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) + + 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.""" + 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() + 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) + 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_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 00000000000..6bf8b1f5bf5 --- /dev/null +++ b/tests/pytorch/nvfp4/test_nvfp4_ep_rows.py @@ -0,0 +1,561 @@ +# 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" + + +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 077a9a2b3a1..e8935c4b3f6 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -363,6 +363,7 @@ list(APPEND transformer_engine_cuda_arch_specific_sources cast/cast_grouped.cu cast/cast_grouped_dbias.cu cast/fused_group_requantize.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 00000000000..814844322ba --- /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 1eb84b73a6a..58c8bac78b0 100644 --- a/transformer_engine/common/include/transformer_engine/ep.h +++ b/transformer_engine/common/include/transformer_engine/ep.h @@ -297,6 +297,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 f6cf7d2fa77..815e4b4176b 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -798,6 +798,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 00000000000..ffe7b4d03b5 --- /dev/null +++ b/transformer_engine/pytorch/csrc/extensions/ep_nvfp4.cpp @@ -0,0 +1,131 @@ +/************************************************************************* + * 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 "../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."); + 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 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 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 89e7d92aca9..4e86736bf77 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -364,6 +364,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 9ed10fb4721..b8d02c61ffe 100644 --- a/transformer_engine/pytorch/ep.py +++ b/transformer_engine/pytorch/ep.py @@ -599,6 +599,148 @@ def _ep_combine_raw(buffer: "EpBuffer", expert_out: torch.Tensor, result: torch. tex.ep_combine(buffer.handle_mem, expert_out, result) +# 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) -> 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), else allocated.""" + if buf is None: + return _alloc_io((rows, row_bytes // 2), torch.bfloat16, device, zero_copy) + 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 @@ -633,6 +775,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 @@ -738,6 +884,92 @@ 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 + ) + 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], + 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, @@ -849,6 +1081,9 @@ 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 def _ep_combine_fwd( @@ -881,19 +1116,23 @@ def _ep_combine_fwd( eager=eager, zero_copy=buffer.zero_copy, ) + if _is_nvfp4_recipe(bwd_quant_recipe): + state.nvfp4_gemm_recipe = _nvfp4_gemm_recipe(buffer, expert_out.shape[0], grad=True) 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( @@ -940,6 +1179,50 @@ 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, + ) + if state.eager: + tex.ep_combine_bwd(state.handle_mem, packed.view(torch.bfloat16), ge_packed, None, None) + else: + torch.ops.transformer_engine_ep.combine_bwd( + state.handle_mem, packed.view(torch.bfloat16), ge_packed + ) + 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 @@ -960,7 +1243,7 @@ def forward( # type: ignore[override] bwd_quant_recipe=None, ): """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 ) @@ -1144,14 +1427,19 @@ 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. @@ -1176,13 +1464,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 @@ -1210,8 +1494,10 @@ def ep_combine( ``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``. 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 buffer.max_tokens_per_rank. @@ -1225,16 +1511,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, From 347e62480f83e70a568dda9306df7f434a475756 Mon Sep 17 00:00:00 2001 From: Siddhartha Raman S Date: Tue, 10 Nov 2026 08:38:34 -0800 Subject: [PATCH 2/6] [PyTorch] Fix NVFP4 EP device selection and zero-copy capture of the decode path Guard the NVFP4 EP bindings (pack, unpack, direct receive) with the input tensor's device so the launch stream, allocations and the stochastic rounding generator follow the input rather than the current device. Under zero-copy, land packed NVFP4 rows that are decoded to BF16 (dispatch and combine backward outside NVFP4 autocast) in a persistent symm-mem scratch on the EpBuffer, allocated on first use, instead of allocating from the symm-mem pool on every call. This makes the decode path CUDA-graph capturable with caller-provided persistent buffers. Add a cross-device test for the bindings and run the NVFP4 dispatch and combine-backward capture tests in the zero-copy EP pass. Signed-off-by: Siddhartha Raman S Co-Authored-By: Claude Opus 5.5 --- tests/pytorch/distributed/run_ep.py | 41 ++++++++++++++++--- tests/pytorch/nvfp4/test_nvfp4_ep_rows.py | 33 +++++++++++++++ .../pytorch/csrc/extensions/ep_nvfp4.cpp | 5 +++ transformer_engine/pytorch/ep.py | 40 +++++++++++++++--- 4 files changed, 109 insertions(+), 10 deletions(-) diff --git a/tests/pytorch/distributed/run_ep.py b/tests/pytorch/distributed/run_ep.py index 7002c0efd1d..dc952a10cc6 100644 --- a/tests/pytorch/distributed/run_ep.py +++ b/tests/pytorch/distributed/run_ep.py @@ -843,13 +843,16 @@ def test_dispatch_nvfp4_autograd(self): 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.""" + 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() + 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() @@ -861,6 +864,7 @@ def step(): ) 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_() @@ -871,6 +875,23 @@ def step(): 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 @@ -986,9 +1007,12 @@ def test_combine_bwd_nvfp4_caller_grad_out(self): 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.""" + 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() @@ -997,8 +1021,14 @@ def test_combine_bwd_nvfp4_capture(self): 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() - 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) + 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) @@ -1008,6 +1038,7 @@ def step(): 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() diff --git a/tests/pytorch/nvfp4/test_nvfp4_ep_rows.py b/tests/pytorch/nvfp4/test_nvfp4_ep_rows.py index 6bf8b1f5bf5..99971146f87 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_ep_rows.py +++ b/tests/pytorch/nvfp4/test_nvfp4_ep_rows.py @@ -403,6 +403,39 @@ def step(): 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( diff --git a/transformer_engine/pytorch/csrc/extensions/ep_nvfp4.cpp b/transformer_engine/pytorch/csrc/extensions/ep_nvfp4.cpp index ffe7b4d03b5..086cd52cd80 100644 --- a/transformer_engine/pytorch/csrc/extensions/ep_nvfp4.cpp +++ b/transformer_engine/pytorch/csrc/extensions/ep_nvfp4.cpp @@ -7,6 +7,7 @@ // NVFP4 token rows for the EP dispatch and combine backward. Independent of NCCL EP. #include +#include #include #include @@ -42,6 +43,8 @@ at::Tensor make_rng_state(size_t elts_per_thread) { 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; @@ -63,6 +66,7 @@ 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)); @@ -90,6 +94,7 @@ std::vector ep_nvfp4_direct_receive(const at::Tensor &packed, int64_ 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); diff --git a/transformer_engine/pytorch/ep.py b/transformer_engine/pytorch/ep.py index b8d02c61ffe..cd43e6b0e9b 100644 --- a/transformer_engine/pytorch/ep.py +++ b/transformer_engine/pytorch/ep.py @@ -331,6 +331,7 @@ class EpBuffer: "total_recv_tokens", "dispatch_fwd_quant_recipe", "combine_bwd_quant_recipe", + "nvfp4_packed_scratch", ) def __init__( @@ -372,6 +373,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) @@ -679,11 +682,25 @@ def _nvfp4_gemm_recipe(buffer: "EpBuffer", rows: int, grad: bool): return recipe -def _nvfp4_landing(buf, rows: int, row_bytes: int, device, zero_copy: bool) -> torch.Tensor: +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), else allocated.""" + 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: - return _alloc_io((rows, row_bytes // 2), torch.bfloat16, device, zero_copy) + 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) @@ -922,7 +939,12 @@ def _ep_prepare_and_dispatch_nvfp4_fwd( 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 + 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( @@ -1084,6 +1106,9 @@ class _CombineState: # 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( @@ -1118,6 +1143,7 @@ def _ep_combine_fwd( ) if _is_nvfp4_recipe(bwd_quant_recipe): state.nvfp4_gemm_recipe = _nvfp4_gemm_recipe(buffer, expert_out.shape[0], grad=True) + state.nvfp4_buffer = buffer return result, state @@ -1195,6 +1221,7 @@ def _ep_combine_bwd_nvfp4(state: "_CombineState", g_result: torch.Tensor): 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, None, None) @@ -1442,6 +1469,8 @@ def ep_dispatch( 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. @@ -1499,7 +1528,8 @@ def ep_combine( ``ep_dispatch``; the packed rows land in ``grad_out``), else a bfloat16 grad decoded into ``grad_out``. 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) From 2804cc871869918ba674267aea33d5e024b5d402 Mon Sep 17 00:00:00 2001 From: Siddhartha Raman S Date: Fri, 13 Nov 2026 12:30:18 -0800 Subject: [PATCH 3/6] [PyTorch] EP combine: option to return the BF16 grad under NVFP4 autocast Review feedback: under NVFP4 autocast the NVFP4 combine backward returns FC2's grad_output as a swizzled NVFP4 GroupedTensor, which holds only GEMM operands. A consumer that needs the decoded gradient, such as a grouped linear with a bias (bias gradient), could not get it. Add ep_combine(..., nvfp4_grad_direct_receive=True); with False the backward decodes the received rows to the BF16 expert_out grad, as outside NVFP4 autocast. Test: the opt-out returns a BF16 grad equal to a BF16 combine backward of the sender-decoded result-grad. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Siddhartha Raman S --- tests/pytorch/distributed/run_ep.py | 27 +++++++++++++++++++++++++++ transformer_engine/pytorch/ep.py | 20 +++++++++++++++++--- 2 files changed, 44 insertions(+), 3 deletions(-) diff --git a/tests/pytorch/distributed/run_ep.py b/tests/pytorch/distributed/run_ep.py index dc952a10cc6..ce1cfecd575 100644 --- a/tests/pytorch/distributed/run_ep.py +++ b/tests/pytorch/distributed/run_ep.py @@ -1098,6 +1098,33 @@ def test_combine_bwd_nvfp4_direct(self): 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) + self._assert_rows_match_per_expert(src.grad[:n], ref[:n], tc) + 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 diff --git a/transformer_engine/pytorch/ep.py b/transformer_engine/pytorch/ep.py index 862d6c89f19..7a2d28c8496 100644 --- a/transformer_engine/pytorch/ep.py +++ b/transformer_engine/pytorch/ep.py @@ -1170,6 +1170,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 @@ -1197,7 +1198,8 @@ def _ep_combine_fwd( zero_copy=buffer.zero_copy, ) if _is_nvfp4_recipe(bwd_quant_recipe): - state.nvfp4_gemm_recipe = _nvfp4_gemm_recipe(buffer, expert_out.shape[0], grad=True) + 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 @@ -1337,11 +1339,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 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 @@ -1355,6 +1363,7 @@ def backward(ctx, g_result): # type: ignore[override] None, # buffer None, # num_local_tokens None, # bwd_quant_recipe + None, # nvfp4_grad_direct_receive ) @@ -1586,6 +1595,7 @@ 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. @@ -1595,7 +1605,10 @@ def ep_combine( ``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``. Eager mode sizes the grad target per step and forbids a caller-supplied buffer. Under zero-copy, leaving it None + ``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 (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 @@ -1624,4 +1637,5 @@ def ep_combine( buffer, num_local_tokens, bwd_quant_recipe, + nvfp4_grad_direct_receive, ) From a3c6deab1fd6b73cac2219a78b7b7001f8bddc9c Mon Sep 17 00:00:00 2001 From: Siddhartha Raman S Date: Fri, 13 Nov 2026 13:03:03 -0800 Subject: [PATCH 4/6] [PyTorch] EP NVFP4: pass the layer arguments that upstream main now requires The upstream merge brought _DispatchState.layer_args and the (top_k, alignment, topk_idx_dtype) arguments of the combine backward; the NVFP4 dispatch and combine-backward paths of this PR did not pass them (merged without a textual conflict), so every NVFP4 EP path failed. Pass them as the BF16/MXFP8 paths do. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Siddhartha Raman S --- transformer_engine/pytorch/ep.py | 17 +++++++++++++++-- 1 file changed, 15 insertions(+), 2 deletions(-) diff --git a/transformer_engine/pytorch/ep.py b/transformer_engine/pytorch/ep.py index 7a2d28c8496..80705bddbfc 100644 --- a/transformer_engine/pytorch/ep.py +++ b/transformer_engine/pytorch/ep.py @@ -1030,6 +1030,7 @@ def _ep_prepare_and_dispatch_nvfp4_fwd( 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). @@ -1295,10 +1296,22 @@ def _ep_combine_bwd_nvfp4(state: "_CombineState", g_result: torch.Tensor): 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, None, None) + 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.handle_mem, + packed.view(torch.bfloat16), + ge_packed, + *state.layer_args, + None, + None, ) if gemm_recipe is not None: return _nvfp4_direct_receive( From d5d95574963d78e9d0500cc9c97764a9d8b2c40e Mon Sep 17 00:00:00 2001 From: Siddhartha Raman S Date: Fri, 13 Nov 2026 13:23:01 -0800 Subject: [PATCH 5/6] [PyTorch] EP test: compare the opt-out combine grad with an alignment=256 reference Co-Authored-By: Claude Opus 5.5 Signed-off-by: Siddhartha Raman S --- tests/pytorch/distributed/run_ep.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/pytorch/distributed/run_ep.py b/tests/pytorch/distributed/run_ep.py index ce1cfecd575..b3b267039b7 100644 --- a/tests/pytorch/distributed/run_ep.py +++ b/tests/pytorch/distributed/run_ep.py @@ -1122,7 +1122,7 @@ def test_combine_bwd_nvfp4_direct_receive_opt_out(self): 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) + 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) def test_combine_bwd_nvfp4_direct_capture(self): From 6d83d2aae14cf37b28cd01f80c1d3d5a01e5f48c Mon Sep 17 00:00:00 2001 From: Siddhartha Raman S Date: Fri, 13 Nov 2026 15:10:08 -0800 Subject: [PATCH 6/6] [PyTorch] EP test: keep the alignment-pass decorator on the direct-capture test Adding the opt-out test before test_combine_bwd_nvfp4_direct_capture moved that test's @_mxfp8_align_test decorator onto the new test, so the capture test was skipped in the alignment=256 passes. Both tests use alignment=256 and carry the decorator. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Siddhartha Raman S --- tests/pytorch/distributed/run_ep.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/pytorch/distributed/run_ep.py b/tests/pytorch/distributed/run_ep.py index b3b267039b7..ff62249c5d9 100644 --- a/tests/pytorch/distributed/run_ep.py +++ b/tests/pytorch/distributed/run_ep.py @@ -1125,6 +1125,7 @@ def test_combine_bwd_nvfp4_direct_receive_opt_out(self): 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