From 249313e7b5666eaf3350dfb64590d485de4ee914 Mon Sep 17 00:00:00 2001 From: YangFei1990 Date: Wed, 12 Aug 2026 10:09:50 -0700 Subject: [PATCH] let mxfp8 return plain torch tensor Signed-off-by: YangFei1990 --- tests/pytorch/distributed/run_ep.py | 80 ++++++++-------- transformer_engine/pytorch/ep.py | 96 ++++++++++--------- .../pytorch/ops/fused/grouped_mlp.py | 17 ++++ 3 files changed, 110 insertions(+), 83 deletions(-) diff --git a/tests/pytorch/distributed/run_ep.py b/tests/pytorch/distributed/run_ep.py index 6e09ed316f..d85c035fd3 100644 --- a/tests/pytorch/distributed/run_ep.py +++ b/tests/pytorch/distributed/run_ep.py @@ -22,6 +22,7 @@ symm_mem_alloc, release_symm_mem_pool, is_symm_backed, + mxfp8_carrier_to_grouped, _ep_combine_raw, _ep_dispatch_raw, ) @@ -136,11 +137,12 @@ def _make_identity_inputs(rank, ep_size, device="cuda"): ) -def _degroup_mxfp8(recv_grouped, valid_counts=None): - """Dequantize a per-expert MXFP8 GroupedTensor to a dense tensor in expert-major order. - With ``valid_counts`` keep only the first ``valid_counts[e]`` rows of each padded expert - slot; otherwise return every (padded) row.""" - parts = recv_grouped.split_into_quantized_tensors() +def _degroup_mxfp8(recv_carrier, token_counts, valid_counts=None): + """Rebuild the per-expert MXFP8 grouped view of an opaque EP carrier tensor and dequantize + it to a dense tensor in expert-major order. ``token_counts`` is the padded per-expert row + counts the carrier was routed with. With ``valid_counts`` keep only the first + ``valid_counts[e]`` rows of each padded expert slot; otherwise return every (padded) row.""" + parts = mxfp8_carrier_to_grouped(recv_carrier, token_counts).split_into_quantized_tensors() if valid_counts is None: return torch.cat([p.dequantize() for p in parts], dim=0) return torch.cat([p.dequantize()[:v] for p, v in zip(parts, valid_counts)], dim=0) @@ -439,46 +441,49 @@ def _assert_mxfp8_matches_bf16(self, recv_mx, tokens, topk_idx, w, tc): torch.cuda.synchronize() n = int(tc.sum()) torch.testing.assert_close( - _degroup_mxfp8(recv_mx).float(), ref_recv.float()[:n], atol=1e-2, rtol=1e-2 + _degroup_mxfp8(recv_mx, tc).float(), ref_recv.float()[:n], atol=1e-2, rtol=1e-2 ) @_eager_test_include @_zero_copy_test_include @_mxfp8_align_test def test_dispatch_mxfp8(self): - """MXFP8 dispatch quantizes bf16 tokens internally; recv (a per-expert GroupedTensor) - dequantized matches a bf16 dispatch of the same tokens. Under zero-copy the recv data and - scales are symm-mem backed.""" + """MXFP8 dispatch quantizes bf16 tokens internally; recv is an opaque plain-tensor + carrier (payload-dtype [rows, hidden], storage = E4M3 data then compact e8m0 scales) + whose rebuilt grouped view, dequantized, matches a bf16 dispatch of the same tokens. + Under zero-copy the carrier is symm-mem backed.""" self._require_mxfp8_shapes() topk_idx, tokens, w = _make_identity_inputs(self.cfg.rank, self.cfg.ep_size) buf = self._make_buffer(dispatch_fwd_quant_recipe=MXFP8BlockScaling(), alignment=128) recv_mx, _rw, tc = ep_dispatch(buf, tokens, topk_idx, w) + # The recv is a plain torch tensor with the payload dtype, not a GroupedTensor. + self.assertIs(type(recv_mx), torch.Tensor) + self.assertEqual(recv_mx.dtype, torch.bfloat16) if ZERO_COPY: - self.assertTrue(is_symm_backed(recv_mx.rowwise_data)) - self.assertTrue(is_symm_backed(recv_mx.scale_inv)) + self.assertTrue(is_symm_backed(recv_mx)) self._assert_mxfp8_matches_bf16(recv_mx, tokens, topk_idx, w, tc) @_zero_copy_test_include @_mxfp8_align_test def test_caller_provides_dispatch_recv_mxfp8(self): - """One caller-supplied buffer holds the recv data followed by the e8m0 scales; ep_dispatch - slices it and the returned GroupedTensor views the data and scale regions of that buffer.""" + """The caller-supplied recv_tokens buffer becomes the MXFP8 carrier: ep_dispatch carves + its storage into recv data followed by the e8m0 scales, returns it as-is, and the + rebuilt grouped view points at those regions.""" self._require_mxfp8_shapes() - from transformer_engine.pytorch.constants import MXFP8_BLOCK_SCALING_SIZE - rc = self.cfg.recv_capacity_per_rank - cols = HIDDEN_DIM // MXFP8_BLOCK_SCALING_SIZE - nbytes = rc * (HIDDEN_DIM + cols) # fp8 data + e8m0 scales, one byte per element + # Payload-dtype [rows, hidden] carrier: bf16's 2 bytes/elem >= 1 (data) + 1/32 (scales). if ZERO_COPY: - recv_buf = symm_mem_alloc((nbytes,), torch.uint8, self.ep_group) + recv_buf = symm_mem_alloc((rc, HIDDEN_DIM), torch.bfloat16, self.ep_group) else: - recv_buf = torch.empty(nbytes, dtype=torch.uint8, device=self.cfg.device) + recv_buf = torch.empty(rc, HIDDEN_DIM, dtype=torch.bfloat16, device=self.cfg.device) buf = self._make_buffer(dispatch_fwd_quant_recipe=MXFP8BlockScaling(), alignment=128) topk_idx, tokens, w = _make_identity_inputs(self.cfg.rank, self.cfg.ep_size) recv_mx, _rw, tc = ep_dispatch(buf, tokens, topk_idx, w, recv_tokens=recv_buf) - # the returned GroupedTensor views the caller buffer's data then scale regions - self.assertEqual(recv_mx.rowwise_data.data_ptr(), recv_buf.data_ptr()) - self.assertEqual(recv_mx.scale_inv.data_ptr(), recv_buf.data_ptr() + rc * HIDDEN_DIM) + # the returned carrier is the caller's buffer; its grouped view carves data then scales + self.assertEqual(recv_mx.data_ptr(), recv_buf.data_ptr()) + recv_grouped = mxfp8_carrier_to_grouped(recv_mx, tc) + self.assertEqual(recv_grouped.rowwise_data.data_ptr(), recv_buf.data_ptr()) + self.assertEqual(recv_grouped.scale_inv.data_ptr(), recv_buf.data_ptr() + rc * HIDDEN_DIM) self._assert_mxfp8_matches_bf16(recv_mx, tokens, topk_idx, w, tc) @_zero_copy_test_include @@ -531,35 +536,34 @@ def test_caller_provides_grad_expert_out(self): @_zero_copy_test_include @_mxfp8_align_test def test_combine_bwd_mxfp8_caller_grad_out(self): - """MXFP8 combine backward into a single caller buffer sliced into data + e8m0 scales: the - returned per-expert GroupedTensor views those regions and, dequantized, matches a bf16 - combine backward reference on the same routing. Under zero-copy the caller buffer and combine - input are symm-mem backed.""" + """MXFP8 combine backward into a caller-supplied carrier (expert_out-shaped bf16, storage + carved into data + e8m0 scales): the grad IS that buffer, and its rebuilt grouped view, + dequantized, matches a bf16 combine backward reference on the same routing. Under + zero-copy the caller buffer and combine input are symm-mem backed.""" self._require_mxfp8_shapes() - from transformer_engine.pytorch.constants import MXFP8_BLOCK_SCALING_SIZE - rc = self.cfg.recv_capacity_per_rank - cols = HIDDEN_DIM // MXFP8_BLOCK_SCALING_SIZE topk_idx, tokens, w = _make_identity_inputs(self.cfg.rank, self.cfg.ep_size) eo_vals = ( torch.linspace(-0.5, 0.5, rc * HIDDEN_DIM, device=self.cfg.device) .reshape(rc, HIDDEN_DIM) .to(torch.bfloat16) ) - # MXFP8 combine backward writes into one caller buffer (data then e8m0 scales) + # MXFP8 combine backward writes into one caller carrier; autograd hands it back as the + # grad, so it must match expert_out's shape and dtype. buf_mx = self._make_buffer(combine_bwd_quant_recipe=MXFP8BlockScaling(), alignment=128) _recv, _rw, tc = ep_dispatch(buf_mx, tokens, topk_idx, w) # seeds the routing - nbytes = rc * (HIDDEN_DIM + cols) if ZERO_COPY: - grad_buf = symm_mem_alloc((nbytes,), torch.uint8, self.ep_group) + grad_buf = symm_mem_alloc((rc, HIDDEN_DIM), torch.bfloat16, self.ep_group) else: - grad_buf = torch.empty(nbytes, dtype=torch.uint8, device=self.cfg.device) + grad_buf = torch.empty(rc, HIDDEN_DIM, dtype=torch.bfloat16, device=self.cfg.device) src_mx = eo_vals.detach().clone().requires_grad_(True) out_mx = ep_combine(buf_mx, self._expert_out(src_mx), grad_out=grad_buf) (0.5 * (out_mx.float() ** 2).sum()).backward() - g_mx = src_mx.grad # per-expert GroupedTensor viewing grad_buf - self.assertEqual(g_mx.rowwise_data.data_ptr(), grad_buf.data_ptr()) - self.assertEqual(g_mx.scale_inv.data_ptr(), grad_buf.data_ptr() + rc * HIDDEN_DIM) + g_mx = src_mx.grad # the caller carrier, returned as-is + self.assertEqual(g_mx.data_ptr(), grad_buf.data_ptr()) + g_grouped = mxfp8_carrier_to_grouped(g_mx, tc) + self.assertEqual(g_grouped.rowwise_data.data_ptr(), grad_buf.data_ptr()) + self.assertEqual(g_grouped.scale_inv.data_ptr(), grad_buf.data_ptr() + rc * HIDDEN_DIM) # bf16 reference combine backward on the same routing buf_bf = self._make_buffer(alignment=128) ep_dispatch(buf_bf, tokens, topk_idx, w) @@ -569,7 +573,7 @@ def test_combine_bwd_mxfp8_caller_grad_out(self): torch.cuda.synchronize() n = int(tc.sum()) torch.testing.assert_close( - _degroup_mxfp8(g_mx).float(), src_bf.grad.float()[:n], atol=5e-2, rtol=5e-2 + _degroup_mxfp8(g_mx, tc).float(), src_bf.grad.float()[:n], atol=5e-2, rtol=5e-2 ) @_eager_test_include @@ -604,7 +608,7 @@ def test_combine_bwd_mxfp8(self): torch.cuda.synchronize() n = int(tc.sum()) torch.testing.assert_close( - _degroup_mxfp8(g_mx).float(), src_bf.grad.float()[:n], atol=5e-2, rtol=5e-2 + _degroup_mxfp8(g_mx, tc).float(), src_bf.grad.float()[:n], atol=5e-2, rtol=5e-2 ) @_zero_copy_test_include diff --git a/transformer_engine/pytorch/ep.py b/transformer_engine/pytorch/ep.py index b0c66df4cc..9afef43f6f 100644 --- a/transformer_engine/pytorch/ep.py +++ b/transformer_engine/pytorch/ep.py @@ -485,7 +485,6 @@ def forward( # type: ignore[override] tokens: torch.Tensor, topk_weights: torch.Tensor, tokens_scale_inv: Optional[torch.Tensor] = None, - token_counts: Optional[torch.Tensor] = None, num_recv_tokens: Optional[int] = None, payload_dtype: torch.dtype = torch.bfloat16, ): @@ -504,8 +503,13 @@ def forward( # type: ignore[override] if is_scaled: if tokens._fp8_dtype != tex.DType.kFloat8E4M3: raise NotImplementedError("EP dispatch supports only E4M3 MXFP8 tokens for now.") - # recv data + scales share one buffer (data then scales); carve or allocate it here. - recv_tokens, recv_scale_inv = _scale_alloc_io( + # The recv is an opaque MXFP8 carrier: one payload-dtype-shaped [rows, hidden] + # tensor whose storage holds [E4M3 data | compact e8m0 scales | slack]. It gives + # the caller plain-tensor semantics (views, record_stream, offload, storage + # release); the consumer rebuilds the grouped view via mxfp8_carrier_to_grouped. + if recv_tokens is None: + recv_tokens = _alloc_io((num_recv_tokens, hidden), payload_dtype, device, zero_copy) + recv_data, recv_scale_inv = _scale_alloc_io( recv_tokens, num_recv_tokens, hidden, @@ -517,7 +521,7 @@ def forward( # type: ignore[override] ) # Reinterpret byte-backed FP8 data as the fp8 dtype so the backend sees a scaled tensor. dispatch_tokens = tokens_data.view(torch.float8_e4m3fn) - dispatch_recv = recv_tokens.view(torch.float8_e4m3fn) + dispatch_recv = recv_data.view(torch.float8_e4m3fn) else: if recv_tokens is None: recv_tokens = _alloc_io((num_recv_tokens, hidden), payload_dtype, device, zero_copy) @@ -544,19 +548,9 @@ def forward( # type: ignore[override] ctx.hidden_dim = hidden # Detach so the long-lived buffers aren't tracked as differentiable outputs; # autograd re-attaches grad_fn pointing back at this Function. For scaled inputs - # the expert-major recv data + scales are wrapped into a per-expert GroupedTensor - # so downstream grouped GEMM and autograd see a proper quantized grouped tensor. - if is_scaled: - recv_out = _make_grouped_mxfp8( - recv_tokens.view(tokens._rowwise_data.dtype), - recv_scale_inv, - token_counts, - tokens._fp8_dtype, - tokens.dtype, - ) - else: - recv_out = recv_tokens.detach() - return recv_out, recv_topk_weights.detach() + # recv_tokens is the opaque MXFP8 carrier (data + scales packed in its storage); + # its grad is the plain high-precision recv grad either way. + return recv_tokens.detach(), recv_topk_weights.detach() @staticmethod def backward(ctx, g_recv_tokens, g_recv_topk_weights): # type: ignore[override] @@ -588,7 +582,6 @@ def backward(ctx, g_recv_tokens, g_recv_topk_weights): # type: ignore[override] grad_tokens.view(ctx.tokens_shape), grad_topk_weights.view(ctx.topk_weights_shape), None, # tokens_scale_inv (scales; non-differentiable) - None, # token_counts (per-expert counts; non-differentiable) None, # num_recv_tokens (sizing scalar) None, # payload_dtype (sizing scalar) ) @@ -613,7 +606,6 @@ def forward( # type: ignore[override] grad_out: Optional[torch.Tensor], expert_out: torch.Tensor, bwd_quant_recipe=None, - token_counts: Optional[torch.Tensor] = None, ): """Combine fwd; stashes the bwd grad target or expert_out shape to size it. When ``bwd_quant_recipe`` is set, the backward sends the result-grad as MXFP8.""" @@ -623,7 +615,6 @@ def forward( # type: ignore[override] ctx.save_for_backward(handle_mem) ctx.grad_out = grad_out ctx.bwd_quant_recipe = bwd_quant_recipe - ctx.token_counts = token_counts ctx.expert_out_shape = expert_out.shape ctx.expert_out_dtype = expert_out.dtype ctx.device = device @@ -632,8 +623,8 @@ def forward( # type: ignore[override] @staticmethod def backward(ctx, g_result): # type: ignore[override] """Combine bwd; scatters the result-grad to expert positions. High-precision sends the grad - as-is; a quantized recipe (MXFP8 today) quantizes it and returns the expert_out grad as a - per-expert GroupedTensor.""" + as-is; a quantized recipe (MXFP8 today) quantizes it and returns the expert_out grad as an + opaque MXFP8 carrier tensor.""" if not g_result.is_contiguous(): g_result = g_result.contiguous() (handle_mem,) = ctx.saved_tensors @@ -649,8 +640,16 @@ def backward(ctx, g_result): # type: ignore[override] mx, g_scale_inv = _quantize_mxfp8(g_result) g_data = mx._rowwise_data recv_pr, hidden = ctx.expert_out_shape[0], ctx.expert_out_shape[-1] + # The expert_out grad is an opaque MXFP8 carrier (see _EpDispatch.forward): one + # expert_out-shaped tensor whose storage holds [E4M3 data | e8m0 scales | slack], + # matching autograd's shape/dtype expectations with plain-tensor semantics. + carrier = ctx.grad_out + if carrier is None: + carrier = _alloc_io( + ctx.expert_out_shape, ctx.expert_out_dtype, ctx.device, tex.ep_get_zero_copy() + ) ge_data, ge_scale_inv = _scale_alloc_io( - ctx.grad_out, + carrier, recv_pr, hidden, g_scale_inv.shape[-1], @@ -667,9 +666,7 @@ def backward(ctx, g_result): # type: ignore[override] g_scale_inv, ge_scale_inv, ) - grad_expert_out = _make_grouped_mxfp8( - ge_data, ge_scale_inv, ctx.token_counts, mx._fp8_dtype, ctx.expert_out_dtype - ) + grad_expert_out = carrier return ( None, # handle_mem @@ -678,7 +675,6 @@ def backward(ctx, g_result): # type: ignore[override] None, # grad_out grad_expert_out, None, # bwd_quant_recipe - None, # token_counts ) @@ -786,26 +782,35 @@ def _scale_alloc_io(buf, rows, data_cols, scale_cols, data_dtype, scale_dtype, d return data, scale_inv -def _make_grouped_mxfp8(data, scale_inv, token_counts, fp8_dtype, fake_dtype): - """Wrap expert-major MXFP8 recv data + compact e8m0 scales as a per-expert ``GroupedTensor``. +def mxfp8_carrier_to_grouped(carrier: torch.Tensor, token_counts: torch.Tensor): + """Rebuild the per-expert MXFP8 ``GroupedTensor`` from an opaque EP carrier tensor. - ``token_counts`` (int64 [num_local_experts]) is the padded per-expert row counts (128-aligned), - used as the group sizes. Grouping is device-side (first_dims/tensor_offsets), so the counts never - sync to host; the outer shape is the static recv capacity, bounded per expert by first_dims. + The carrier is the plain ``[rows, hidden]`` tensor EP dispatch forward (and combine + backward) return when quantization is on: its storage holds ``rows*hidden`` bytes of + expert-major E4M3 data followed by ``rows*(hidden/block)`` bytes of compact e8m0 scales. + ``token_counts`` (int64 [num_local_experts]) is the padded per-expert row counts + (alignment-padded), used as the group sizes; grouping is device-side + (first_dims/tensor_offsets), so the counts never sync to host. """ + from .constants import MXFP8_BLOCK_SCALING_SIZE from .tensor.grouped_tensor import GroupedTensor from .tensor.mxfp8_tensor import MXFP8Quantizer - assert data.dim() == 2, "recv data must be 2D [capacity_rows, hidden]" - capacity_rows, hidden = data.shape - quantizer = MXFP8Quantizer(fp8_dtype, rowwise=True, columnwise=False) + assert carrier.dim() == 2, "EP carrier must be 2D [rows, hidden]" + if not carrier.is_contiguous(): + raise ValueError("EP carrier must be contiguous; got a strided tensor.") + rows, hidden = carrier.shape + scale_cols = hidden // MXFP8_BLOCK_SCALING_SIZE + flat = carrier.detach().reshape(-1).view(torch.uint8) + data_bytes = rows * hidden + quantizer = MXFP8Quantizer(tex.DType.kFloat8E4M3, rowwise=True, columnwise=False) return GroupedTensor( - shape=(capacity_rows, hidden), - dtype=fake_dtype, + shape=(rows, hidden), + dtype=carrier.dtype, num_tensors=token_counts.numel(), quantizer=quantizer, - data=data.reshape(-1).detach(), - scale_inv=scale_inv.reshape(-1).detach(), + data=flat[:data_bytes], + scale_inv=flat[data_bytes : data_bytes + rows * scale_cols], first_dims=token_counts, tensor_offsets=tex.splits_to_offsets(token_counts, hidden), ) @@ -822,9 +827,12 @@ 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 set (``MXFP8BlockScaling`` only for now), + tokens are quantized internally and recv is returned as an opaque MXFP8 carrier: a plain + ``[rows, hidden]`` payload-dtype tensor whose storage holds the E4M3 data then the compact + e8m0 scales (rebuild the grouped view with ``mxfp8_carrier_to_grouped``; the contents are + not valid payload-dtype data). Otherwise recv stays bfloat16. 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 @@ -881,7 +889,6 @@ def ep_dispatch( tokens, topk_weights, tokens_scale_inv, - tokens_per_expert, num_recv_tokens, buffer.payload_dtype, ) @@ -915,7 +922,7 @@ 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 as MXFP8 and returns the expert_out grad as an opaque MXFP8 carrier (see ep_dispatch). bwd_quant_recipe = None if buffer.combine_bwd_quant_recipe is not None: from ..common.recipe import MXFP8BlockScaling @@ -933,5 +940,4 @@ def ep_combine( grad_out, expert_out, bwd_quant_recipe, - buffer.tokens_per_expert, ) diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index 909e5a8a9b..f5504291de 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -1231,6 +1231,14 @@ def fuser_forward( ) fc1_input_quantizer.optimize_for_gemm = True fc1_input_quantizer.internal = True + if getattr(fc1_op, "ep_mxfp8_carrier_input", False) and not isinstance( + input_, GroupedTensor + ): + # Input is an opaque MXFP8 EP-dispatch carrier (data + scales packed in a plain + # tensor's storage); rebuild the per-expert grouped view before the prequant path. + from ...ep import mxfp8_carrier_to_grouped + + input_ = mxfp8_carrier_to_grouped(input_, split_sizes) if isinstance(input_, GroupedTensor): # Input arrived already quantized (e.g. FP8 token dispatch): reuse its rowwise data # for the GEMM and let the helper supply whatever else the GEMMs need. An input that @@ -1877,6 +1885,15 @@ def fuser_backward( output_fc2_dbias = fc2_op.has_bias fc2_dbias_packed = None fc2_dy = None + if getattr(fc2_op, "ep_mxfp8_carrier_grad", False) and not isinstance( + grad_output, GroupedTensor + ): + # Grad arrived as an opaque MXFP8 EP combine-backward carrier (data + scales packed + # in a plain tensor's storage); rebuild the per-expert grouped view before the + # prequant grad path. + from ...ep import mxfp8_carrier_to_grouped + + grad_output = mxfp8_carrier_to_grouped(grad_output.contiguous(), split_sizes) if isinstance(grad_output, GroupedTensor): # Grad output arrived already quantized (e.g. FP8 token dispatch): reuse its rowwise # data for the dgrad GEMM. Bias grads are reduced from the dequantized grad, which is