From 418e5038cc142ba0ea00a0264f6493c9506586e6 Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Thu, 13 Aug 2026 09:48:27 -0700 Subject: [PATCH] [PyTorch] Fuse eager EP prepare and dispatch into a single op Signed-off-by: Phuong Nguyen --- tests/pytorch/distributed/run_ep.py | 21 ++ transformer_engine/pytorch/csrc/extensions.h | 13 ++ .../pytorch/csrc/extensions/ep.cpp | 41 ++++ transformer_engine/pytorch/ep.py | 188 +++++++++++++----- 4 files changed, 218 insertions(+), 45 deletions(-) diff --git a/tests/pytorch/distributed/run_ep.py b/tests/pytorch/distributed/run_ep.py index 6e09ed316f..c1ac784f7a 100644 --- a/tests/pytorch/distributed/run_ep.py +++ b/tests/pytorch/distributed/run_ep.py @@ -458,6 +458,27 @@ def test_dispatch_mxfp8(self): self.assertTrue(is_symm_backed(recv_mx.scale_inv)) self._assert_mxfp8_matches_bf16(recv_mx, tokens, topk_idx, w, tc) + @_eager_test_include + @_mxfp8_align_test + def test_dispatch_mxfp8_autograd(self): + """MXFP8 dispatch fwd+bwd. Seeding the recv grad with ones scatters TOP_K back to each token + under identity routing, so grad_tokens equals TOP_K through the dispatch backward and the + input quantizer STE. Exercises the fused eager MXFP8 dispatch backward.""" + self._require_mxfp8_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=MXFP8BlockScaling(), alignment=128) + recv_mx, _rw, _tc = ep_dispatch(buf, tokens_p, topk_idx, w) + g_recv = torch.ones(recv_mx.shape, dtype=torch.bfloat16, device=self.cfg.device) + torch.autograd.backward(recv_mx, 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=5e-2, + rtol=5e-2, + ) + @_zero_copy_test_include @_mxfp8_align_test def test_caller_provides_dispatch_recv_mxfp8(self): diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index 9114b1e453..3a4760e566 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -725,6 +725,19 @@ void ep_dispatch(at::Tensor handle_mem, at::Tensor topk_idx, at::Tensor tokens, std::optional tokens_scale_inv = std::nullopt, std::optional recv_scale_inv = std::nullopt); +// Fused eager-mode prepare + dispatch in a single call so no Python runs between +// the host recv-count read and the dispatch launch. Prepare writes the per-step +// recv total directly into pinned-host total_recv_tokens (UVA); a stream sync then +// makes it host-readable with no D2H copy, and the recv outputs are sized from it. +// Passing tokens_scale_inv selects the MXFP8 path and also allocates + returns the +// recv scale-inverse. Eager forbids symm-mem zero-copy IO. Returns {recv_tokens, +// recv_topk_weights} (bf16), or {recv_tokens, recv_topk_weights, recv_scale_inv} (MXFP8). +std::vector ep_prepare_and_dispatch_eager( + at::Tensor handle_mem, at::Tensor topk_idx, at::Tensor tokens, at::Tensor topk_weights, + at::Tensor tokens_per_expert, at::Tensor total_recv_tokens, int64_t top_k, + int64_t dispatch_output_per_expert_alignment, + std::optional tokens_scale_inv = std::nullopt); + void ep_combine(at::Tensor handle_mem, at::Tensor expert_out, at::Tensor result); void ep_dispatch_bwd(at::Tensor handle_mem, at::Tensor grad, at::Tensor g_recv_topk_weights, diff --git a/transformer_engine/pytorch/csrc/extensions/ep.cpp b/transformer_engine/pytorch/csrc/extensions/ep.cpp index cef489dbbb..c1766cf175 100644 --- a/transformer_engine/pytorch/csrc/extensions/ep.cpp +++ b/transformer_engine/pytorch/csrc/extensions/ep.cpp @@ -347,6 +347,41 @@ void ep_dispatch(at::Tensor handle_mem, at::Tensor topk_idx, at::Tensor tokens, recv_topk_w_te.data(), recv_topk_w_win, stream); } +std::vector ep_prepare_and_dispatch_eager(at::Tensor handle_mem, at::Tensor topk_idx, + at::Tensor tokens, at::Tensor topk_weights, + at::Tensor tokens_per_expert, + at::Tensor total_recv_tokens, int64_t top_k, + int64_t dispatch_output_per_expert_alignment, + std::optional tokens_scale_inv) { + auto stream = at::cuda::getCurrentCUDAStream().stream(); + NVTE_CHECK(tokens.dim() == 2, "eager dispatch tokens must be 2D [T, H]"); + NVTE_CHECK(total_recv_tokens.is_pinned(), "eager total_recv_tokens must be pinned host memory"); + + // Prepare writes the per-step recv total straight into pinned host total_recv_tokens + // (UVA), so a stream sync makes it host-readable with no D2H copy. Done here so no + // Python runs between the count read and the dispatch launch below. + ep_prepare(handle_mem, topk_idx, tokens_per_expert, top_k, dispatch_output_per_expert_alignment, + total_recv_tokens); + NVTE_CHECK_CUDA(cudaStreamSynchronize(stream)); + const int64_t num_recv = *total_recv_tokens.data_ptr(); + + // Recv outputs sized to the host count (eager forbids zero-copy, so plain allocations). + const int64_t H = tokens.size(-1); + auto recv_tokens = at::empty({num_recv, H}, tokens.options()); + auto recv_topk_weights = at::empty({num_recv}, topk_weights.options()); + if (tokens_scale_inv.has_value()) { + // MXFP8: allocate the recv scale-inverse alongside the fp8 data and route both. + auto recv_scale_inv = + at::empty({num_recv, tokens_scale_inv->size(-1)}, tokens_scale_inv->options()); + ep_dispatch(handle_mem, topk_idx, tokens, topk_weights, recv_tokens, recv_topk_weights, + tokens_scale_inv, recv_scale_inv); + return {recv_tokens, recv_topk_weights, recv_scale_inv}; + } + ep_dispatch(handle_mem, topk_idx, tokens, topk_weights, recv_tokens, recv_topk_weights, + std::nullopt, std::nullopt); + return {recv_tokens, recv_topk_weights}; +} + void ep_combine(at::Tensor handle_mem, at::Tensor expert_out, at::Tensor result) { auto stream = at::cuda::getCurrentCUDAStream().stream(); NVTE_CHECK(expert_out.dim() >= 2, "expert_out must be at least 2D [..., recv_pr, H]"); @@ -509,6 +544,12 @@ void register_ep_bindings(pybind11::module_& m) { py::arg("tokens"), py::arg("topk_weights"), py::arg("recv_tokens"), py::arg("recv_topk_weights"), py::arg("tokens_scale_inv") = std::nullopt, py::arg("recv_scale_inv") = std::nullopt, py::call_guard()); + m.def("ep_prepare_and_dispatch_eager", &ep_prepare_and_dispatch_eager, + "Fused eager EP prepare + dispatch", py::arg("handle_mem"), py::arg("topk_idx"), + py::arg("tokens"), py::arg("topk_weights"), py::arg("tokens_per_expert"), + py::arg("total_recv_tokens"), py::arg("top_k"), + py::arg("dispatch_output_per_expert_alignment"), py::arg("tokens_scale_inv") = std::nullopt, + py::call_guard()); m.def("ep_combine", &ep_combine, "EP combine", py::call_guard()); m.def("ep_dispatch_bwd", &ep_dispatch_bwd, "EP dispatch backward", py::call_guard()); diff --git a/transformer_engine/pytorch/ep.py b/transformer_engine/pytorch/ep.py index b0c66df4cc..1a841e2fb2 100644 --- a/transformer_engine/pytorch/ep.py +++ b/transformer_engine/pytorch/ep.py @@ -248,7 +248,6 @@ class EpBuffer: "zero_copy", "eager", "total_recv_tokens", - "_host_total_recv_tokens", "dispatch_fwd_quant_recipe", "combine_bwd_quant_recipe", ) @@ -300,13 +299,15 @@ def __init__( ) # Persistent tensor; keep resident if activation CPU offloading is on. mark_not_offload(self.handle_mem) - # Per-step recv-token total (device int64 [1]), written by ep_prepare. - # Eager mode uses it to size the recv outputs; graph mode reads it after - # replay to detect overflow past recv_capacity_per_rank. - self.total_recv_tokens = torch.empty(1, dtype=torch.int64, device=device) - mark_not_offload(self.total_recv_tokens) - # Host mirror of total_recv_tokens, set by ep_prepare in eager mode. - self._host_total_recv_tokens: Optional[int] = None + # Per-step recv-token total (int64 [1]), written by ep_prepare. Eager reads it + # host-side to size the recv outputs, so it lives in pinned host memory the prepare + # kernel writes directly (UVA) — no D2H copy, just a stream sync. Graph mode keeps + # it on device for the backend's post-replay overflow check. + if self.eager: + self.total_recv_tokens = torch.empty(1, dtype=torch.int64, pin_memory=True) + else: + self.total_recv_tokens = torch.empty(1, dtype=torch.int64, device=device) + mark_not_offload(self.total_recv_tokens) # torch.library custom ops (so they don't graph-break under torch.compile) @@ -433,9 +434,9 @@ def ep_prepare(buffer: "EpBuffer", topk_idx: torch.Tensor) -> torch.Tensor: ``buffer.tokens_per_expert`` (int64, shape [num_local_experts]). topk_idx must be int32 or int64. - Also fills ``buffer.total_recv_tokens`` (device int64 [1]) with the per-step - recv total; eager mode mirrors it to the host to size the recv outputs, graph - mode reads it device-side to detect overflow. + Also fills ``buffer.total_recv_tokens`` (int64 [1]; pinned host in eager mode, + device otherwise) with the per-step recv total; graph mode reads it device-side + to detect overflow. """ torch.ops.transformer_engine_ep.prepare( buffer.handle_mem, @@ -445,8 +446,6 @@ def ep_prepare(buffer: "EpBuffer", topk_idx: torch.Tensor) -> torch.Tensor: buffer.alignment, buffer.total_recv_tokens, ) - if buffer.eager: - buffer._host_total_recv_tokens = int(buffer.total_recv_tokens.item()) return buffer.tokens_per_expert @@ -472,6 +471,28 @@ def _ep_combine_raw(buffer: "EpBuffer", expert_out: torch.Tensor, result: torch. # autograd.Function wrappers +def _dispatch_backward(ctx, g_recv_tokens, g_recv_topk_weights): + """Shared dispatch bwd: run dispatch_bwd and reshape grads to the fwd input layout.""" + (handle_mem,) = ctx.saved_tensors + device = handle_mem.device + g_recv_tokens = g_recv_tokens.contiguous() + g_recv_topk_weights = g_recv_topk_weights.contiguous() + # Dispatch grad follows the recv grad's (high-precision) dtype; the quantizer's STE + # owns the fp8 boundary for scaled inputs. + grad_tokens = torch.empty( + ctx.tokens_T_flat, ctx.hidden_dim, dtype=g_recv_tokens.dtype, device=device + ) + grad_topk_weights = torch.empty(ctx.topk_T_flat, ctx.top_k, dtype=torch.float32, device=device) + torch.ops.transformer_engine_ep.dispatch_bwd( + handle_mem, + g_recv_tokens, + g_recv_topk_weights, + grad_tokens, + grad_topk_weights, + ) + return grad_tokens.view(ctx.tokens_shape), grad_topk_weights.view(ctx.topk_weights_shape) + + class _EpDispatch(torch.autograd.Function): """Autograd dispatch; caller runs prepare first. bwd uses user-supplied grad inputs as-is.""" @@ -538,7 +559,7 @@ def forward( # type: ignore[override] ctx.save_for_backward(handle_mem) ctx.tokens_shape = tokens.shape ctx.topk_weights_shape = topk_weights.shape - ctx.num_tokens = tokens_data.shape[0] + ctx.tokens_T_flat = tokens_data.shape[0] ctx.topk_T_flat = topk_weights.numel() // topk_weights.shape[-1] ctx.top_k = topk_weights.shape[-1] ctx.hidden_dim = hidden @@ -561,32 +582,14 @@ def forward( # type: ignore[override] @staticmethod def backward(ctx, g_recv_tokens, g_recv_topk_weights): # type: ignore[override] """Dispatch bwd; normalizes grad-input layout, otherwise passes through.""" - (handle_mem,) = ctx.saved_tensors - device = handle_mem.device - g_recv_tokens = g_recv_tokens.contiguous() - g_recv_topk_weights = g_recv_topk_weights.contiguous() - # Dispatch grad follows the recv grad's (high-precision) dtype; the quantizer's STE - # owns the fp8 boundary for scaled inputs. - grad_tokens = torch.empty( - ctx.num_tokens, ctx.hidden_dim, dtype=g_recv_tokens.dtype, device=device - ) - grad_topk_weights = torch.empty( - ctx.topk_T_flat, ctx.top_k, dtype=torch.float32, device=device - ) - torch.ops.transformer_engine_ep.dispatch_bwd( - handle_mem, - g_recv_tokens, - g_recv_topk_weights, - grad_tokens, - grad_topk_weights, - ) + grad_tokens, grad_topk_weights = _dispatch_backward(ctx, g_recv_tokens, g_recv_topk_weights) return ( None, # handle_mem None, # recv_tokens None, # recv_topk_weights None, # topk_idx - grad_tokens.view(ctx.tokens_shape), - grad_topk_weights.view(ctx.topk_weights_shape), + grad_tokens, + grad_topk_weights, None, # tokens_scale_inv (scales; non-differentiable) None, # token_counts (per-expert counts; non-differentiable) None, # num_recv_tokens (sizing scalar) @@ -594,6 +597,90 @@ def backward(ctx, g_recv_tokens, g_recv_topk_weights): # type: ignore[override] ) +class _EpPrepareAndDispatchEager(torch.autograd.Function): + """Fused eager prepare + dispatch in one C++ op, so no Python runs between the host recv-count + read and the dispatch launch. Eager sizes and allocates the recv outputs from the host count, + so it cannot take caller-supplied buffers. When tokens_scale_inv is set (MXFP8), tokens is the + quantized operand and recv is returned as a per-expert GroupedTensor. Backward is shared.""" + + @staticmethod + def forward( # type: ignore[override] + ctx, + handle_mem: torch.Tensor, + tokens_per_expert: torch.Tensor, + total_recv_tokens: torch.Tensor, + topk_idx: torch.Tensor, + tokens: torch.Tensor, + topk_weights: torch.Tensor, + top_k: int, + alignment: int, + tokens_scale_inv: Optional[torch.Tensor] = None, + ): + """Fused prepare + dispatch fwd; recv outputs are sized from the host recv-count.""" + is_scaled = tokens_scale_inv is not None + tokens_data = tokens._rowwise_data if isinstance(tokens, QuantizedTensor) else tokens + if is_scaled: + if tokens._fp8_dtype != tex.DType.kFloat8E4M3: + raise NotImplementedError("EP dispatch supports only E4M3 MXFP8 tokens for now.") + # Reinterpret the byte-backed fp8 data so the backend sees a scaled tensor. + recv_tokens, recv_topk_weights, recv_scale_inv = tex.ep_prepare_and_dispatch_eager( + handle_mem, + topk_idx, + tokens_data.view(torch.float8_e4m3fn), + topk_weights, + tokens_per_expert, + total_recv_tokens, + top_k, + alignment, + tokens_scale_inv, + ) + else: + recv_tokens, recv_topk_weights = tex.ep_prepare_and_dispatch_eager( + handle_mem, + topk_idx, + tokens, + topk_weights, + tokens_per_expert, + total_recv_tokens, + top_k, + alignment, + ) + ctx.save_for_backward(handle_mem) + ctx.tokens_shape = tokens.shape + ctx.topk_weights_shape = topk_weights.shape + ctx.tokens_T_flat = tokens_data.shape[0] + ctx.topk_T_flat = topk_weights.numel() // topk_weights.shape[-1] + ctx.top_k = topk_weights.shape[-1] + ctx.hidden_dim = tokens_data.shape[-1] + if is_scaled: + # Wrap expert-major recv data + scales into a per-expert GroupedTensor (post-launch). + recv_out = _make_grouped_mxfp8( + recv_tokens.view(tokens._rowwise_data.dtype), + recv_scale_inv, + tokens_per_expert, + tokens._fp8_dtype, + tokens.dtype, + ) + return recv_out, recv_topk_weights.detach() + return recv_tokens.detach(), recv_topk_weights.detach() + + @staticmethod + def backward(ctx, g_recv_tokens, g_recv_topk_weights): # type: ignore[override] + """Dispatch bwd; normalizes grad-input layout, otherwise passes through.""" + grad_tokens, grad_topk_weights = _dispatch_backward(ctx, g_recv_tokens, g_recv_topk_weights) + return ( + None, # handle_mem + None, # tokens_per_expert + None, # total_recv_tokens + None, # topk_idx + grad_tokens, + grad_topk_weights, + None, # top_k + None, # alignment + None, # tokens_scale_inv + ) + + 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 @@ -853,13 +940,8 @@ def ep_dispatch( "and cannot use caller-supplied recv_tokens / recv_topk_weights" ) - # Prepare (routing AllGather) up front so the recv outputs can be sized; in - # eager mode ep_prepare also host-syncs this step's recv-token total. - tokens_per_expert = ep_prepare(buffer, topk_idx) - num_recv_tokens = ( - buffer._host_total_recv_tokens if buffer.eager else buffer.recv_capacity_per_rank - ) - + # Quantize up front (before prepare) so the quant kernels overlap the eager count sync and the + # quantized tensor stays the autograd operand; grad reaches the pre-quant input. tokens_scale_inv = None if buffer.dispatch_fwd_quant_recipe is not None: from ..common.recipe import MXFP8BlockScaling @@ -869,10 +951,26 @@ def ep_dispatch( "EP block-scaled dispatch supports MXFP8BlockScaling only; got " f"{type(buffer.dispatch_fwd_quant_recipe).__name__}." ) - # Quantize here (not in forward) so the quantized tensor stays the autograd operand and grad - # reaches the pre-quant input; forward then carves the recv buffers and routes. tokens, tokens_scale_inv = _quantize_mxfp8(tokens) + # Eager: fused prepare + count read + recv alloc + dispatch in one C++ op so no Python runs + # between the host count read and the dispatch launch. bf16 and MXFP8. + if buffer.eager: + recv_out, recv_topk_weights = _EpPrepareAndDispatchEager.apply( + buffer.handle_mem, + buffer.tokens_per_expert, + buffer.total_recv_tokens, + topk_idx, + tokens, + topk_weights, + buffer.top_k, + buffer.alignment, + tokens_scale_inv, + ) + return recv_out, recv_topk_weights, buffer.tokens_per_expert + + # Non-eager (capacity): un-fused prepare + dispatch, recv sized to recv_capacity_per_rank. + tokens_per_expert = ep_prepare(buffer, topk_idx) recv_tokens, recv_topk_weights = _EpDispatch.apply( buffer.handle_mem, recv_tokens, @@ -882,7 +980,7 @@ def ep_dispatch( topk_weights, tokens_scale_inv, tokens_per_expert, - num_recv_tokens, + buffer.recv_capacity_per_rank, buffer.payload_dtype, ) return recv_tokens, recv_topk_weights, tokens_per_expert