[PyTorch] Fuse eager EP prepare and dispatch into a single op

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):

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<at::Tensor> tokens_scale_inv = std::nullopt,
                  std::optional<at::Tensor> 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<at::Tensor> 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<at::Tensor> 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,

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<at::Tensor> 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<at::Tensor> 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<int64_t>();
+
+  // 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<py::gil_scoped_release>());
+  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<py::gil_scoped_release>());
   m.def("ep_combine", &ep_combine, "EP combine", py::call_guard<py::gil_scoped_release>());
   m.def("ep_dispatch_bwd", &ep_dispatch_bwd, "EP dispatch backward",
         py::call_guard<py::gil_scoped_release>());

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,39 +582,105 @@ 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)
             None,  # payload_dtype (sizing scalar)
         )
 
 
+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