[PyTorch] NCCL EP eager mode and drop-on-overflow policy

examples/pytorch/ep/bench/ep_bench.py

@@ -181,8 +181,9 @@ def main():
         ep_group,
         num_experts=E,
         max_tokens_per_rank=T,
-        recv_capacity_per_rank=recv_pr,
         hidden_dim=H,
+        num_topk=K,
+        recv_capacity_per_rank=recv_pr,
         max_num_sms=args.max_num_sms,
     )
 
@@ -207,8 +208,6 @@ def main():
         recv_capacity_per_rank=recv_pr,
         hidden_dim=H,
         num_local_experts=num_local_experts,
-        dispatch_recv_tokens=caller_recv_tokens,
-        combine_grad_expert_out=caller_grad_expert_out,
     )
 
     tokens = tokens_hbm
@@ -235,8 +234,13 @@ def main():
     eo_p = recv_tokens.detach().clone().requires_grad_(True)
 
     # Stand-in callables; the cuda-graph branch below swaps in graphed versions.
-    fwd_bwd_dispatch_fn = lambda x: ep_dispatch(buffer, x, topk_idx, topk_w)[0]  # noqa: E731
-    fwd_bwd_combine_fn = lambda expert_out: ep_combine(buffer, expert_out)  # noqa: E731
+    # caller_* are None unless the caller opted in, matching ep_dispatch/ep_combine defaults.
+    fwd_bwd_dispatch_fn = lambda x: ep_dispatch(  # noqa: E731
+        buffer, x, topk_idx, topk_w, recv_tokens=caller_recv_tokens
+    )[0]
+    fwd_bwd_combine_fn = lambda expert_out: ep_combine(  # noqa: E731
+        buffer, expert_out, grad_out=caller_grad_expert_out
+    )
 
     def _dispatch_raw():
         _ep_dispatch_raw(buffer, topk_idx, tokens, topk_w, recv_tokens, recv_w)
@@ -246,7 +250,7 @@ def _combine_raw():
         _ep_combine_raw(buffer, expert_out, out_buf)
 
     def _ep_dispatch_fwd():
-        ep_dispatch(buffer, tokens.detach(), topk_idx, topk_w)
+        ep_dispatch(buffer, tokens.detach(), topk_idx, topk_w, recv_tokens=caller_recv_tokens)
 
     def _ep_dispatch_fwd_bwd():
         tokens_p.grad = None
@@ -288,11 +292,11 @@ def _ep_combine_fwd_bwd():
         # Graph fwd+bwd of the autograd-wrapped ops via make_graphed_callables.
         class _DispatchMod(torch.nn.Module):
             def forward(self, x):
-                return ep_dispatch(buffer, x, topk_idx, topk_w)[0]
+                return ep_dispatch(buffer, x, topk_idx, topk_w, recv_tokens=caller_recv_tokens)[0]
 
         class _CombineMod(torch.nn.Module):
             def forward(self, expert_out):
-                return ep_combine(buffer, expert_out)
+                return ep_combine(buffer, expert_out, grad_out=caller_grad_expert_out)
 
         disp_mod = _DispatchMod().cuda()
         comb_mod = _CombineMod().cuda()

examples/pytorch/ep/ep_moe.py

@@ -131,8 +131,9 @@ def main():
         ep_group,
         num_experts=num_experts,
         max_tokens_per_rank=T,
-        recv_capacity_per_rank=recv_pr,
         hidden_dim=args.hidden,
+        num_topk=args.top_k,
+        recv_capacity_per_rank=recv_pr,
     )
     try:
         _run_layer(
@@ -182,15 +183,13 @@ def _run_layer(args, rank, world_size, ep_size, num_experts, num_local_experts,
         recv_capacity_per_rank=recv_pr,
         hidden_dim=args.hidden,
         num_local_experts=num_local_experts,
-        dispatch_recv_tokens=recv_tokens,
-        combine_grad_expert_out=grad_expert_out,
     )
 
-    recv_t, recv_w_out, _tc = ep_dispatch(buffer, tokens, topk_idx, topk_w)
+    recv_t, recv_w_out, _tc = ep_dispatch(buffer, tokens, topk_idx, topk_w, recv_tokens=recv_tokens)
     expert_out = _batched_expert_linear(recv_t, kernels_local, num_local_experts)
     # Apply per-slot topk weighting before combine.
     expert_out = expert_out * recv_w_out.unsqueeze(-1).to(expert_out.dtype)
-    out = ep_combine(buffer, expert_out)
+    out = ep_combine(buffer, expert_out, grad_out=grad_expert_out)
 
     loss = 0.5 * (out.float() ** 2).sum()
     loss.backward()

tests/pytorch/distributed/run_ep.py

@@ -26,6 +26,8 @@
 
 
 ZERO_COPY = os.environ.get("NVTE_EP_ZERO_COPY", "0") == "1"
+EAGER = os.environ.get("NVTE_EP_EAGER", "0") == "1"
+OVERFLOW = os.environ.get("NVTE_EP_OVERFLOW", "0") == "1"
 
 # Must come after the transformer_engine import so libtransformer_engine.so is loaded.
 import transformer_engine_torch as tex  # noqa: F401
@@ -43,6 +45,18 @@ def _zero_copy_test_include(fn):
     return fn
 
 
+def _eager_test_include(fn):
+    """Mark a test to run in the eager pass; others skip there."""
+    fn._eager_test_include = True
+    return fn
+
+
+def _overflow_test_include(fn):
+    """Mark a test to run in the overflow (drop-on-overflow) pass; others skip there."""
+    fn._overflow_test_include = True
+    return fn
+
+
 class _StageToSymm(torch.autograd.Function):
     """Identity op that stages ``src`` into a symm-mem buffer; grad passes through.
     Lets a test feed a symm-mem-backed, autograd-tracked tensor into ep_combine.
@@ -126,6 +140,10 @@ def _make_cfg() -> _Cfg:
     active = min(cfg.num_experts, T * cfg.ep_size * TOP_K)
     overconc = cfg.num_experts // active
     cfg.recv_capacity_per_rank = NUM_LOCAL_EXPERTS * max(T * cfg.ep_size * TOP_K, 16) * overconc * 2
+    if OVERFLOW:
+        # Undersize recv capacity so identity routing overflows a rank's budget;
+        # HT requires capacity >= max_tokens_per_rank.
+        cfg.recv_capacity_per_rank = TOKENS_PER_RANK
     cfg.device = torch.device("cuda", torch.cuda.current_device())
     return cfg
 
@@ -144,9 +162,12 @@ def setUpClass(cls):
             cls.ep_group,
             num_experts=cls.cfg.num_experts,
             max_tokens_per_rank=TOKENS_PER_RANK,
-            recv_capacity_per_rank=cls.cfg.recv_capacity_per_rank,
             hidden_dim=HIDDEN_DIM,
+            num_topk=TOP_K,
+            # Omit recv_capacity_per_rank to select eager mode.
+            recv_capacity_per_rank=None if EAGER else cls.cfg.recv_capacity_per_rank,
             zero_copy=ZERO_COPY,
+            drop_on_overflow=OVERFLOW,
         )
 
     def setUp(self):
@@ -155,14 +176,22 @@ def setUp(self):
             getattr(self, self._testMethodName), "_zero_copy_test_include", False
         ):
             self.skipTest("not exercised in zero-copy mode")
+        # Only the eager-capable tests run in the eager pass.
+        if EAGER and not getattr(getattr(self, self._testMethodName), "_eager_test_include", False):
+            self.skipTest("not exercised in eager mode")
+        # Only the overflow-capable tests run in the overflow pass.
+        if OVERFLOW and not getattr(
+            getattr(self, self._testMethodName), "_overflow_test_include", False
+        ):
+            self.skipTest("not exercised in overflow mode")
 
     def _make_buffer(self, alignment=0, top_k=TOP_K):
         return EpBuffer(
             top_k=top_k,
             max_tokens_per_rank=TOKENS_PER_RANK,
-            recv_capacity_per_rank=self.cfg.recv_capacity_per_rank,
             hidden_dim=HIDDEN_DIM,
             num_local_experts=NUM_LOCAL_EXPERTS,
+            recv_capacity_per_rank=None if EAGER else self.cfg.recv_capacity_per_rank,
             alignment=alignment,
         )
 
@@ -184,12 +213,12 @@ def _stage_grad_symm(self, x, symm_buf=None):
         return _GradToSymm.apply(x, symm_buf)
 
     def _make_raw_recv(self, dtype=torch.bfloat16):
-        """Raw recv tensors + token_counts for the primitive tests."""
+        """Raw recv tensors + tokens_per_expert for the primitive tests."""
         rc = self.cfg.recv_capacity_per_rank
         return (
             torch.empty(rc, HIDDEN_DIM, dtype=dtype, device=self.cfg.device),
             torch.empty(rc, dtype=torch.float32, device=self.cfg.device),
-            torch.empty(NUM_LOCAL_EXPERTS, dtype=torch.int32, device=self.cfg.device),
+            torch.empty(NUM_LOCAL_EXPERTS, dtype=torch.int64, device=self.cfg.device),
         )
 
     @staticmethod
@@ -205,17 +234,63 @@ def _moe_step(self, buffer, topk_idx, tokens, w):
 
     # Prepare
 
+    @_eager_test_include
     def test_primitive_prepare(self):
         buf = self._make_buffer()
         topk_idx, _toks, _w = _make_identity_inputs(self.cfg.rank, self.cfg.ep_size)
-        token_counts = ep_prepare(buf, topk_idx)
+        tokens_per_expert = ep_prepare(buf, topk_idx)
         torch.cuda.synchronize()
-        self.assertEqual(token_counts.shape, (NUM_LOCAL_EXPERTS,))
-        local = int(token_counts.sum().item())
+        self.assertEqual(tokens_per_expert.shape, (NUM_LOCAL_EXPERTS,))
+        local = int(tokens_per_expert.sum().item())
         total = torch.tensor([local], dtype=torch.int64, device=self.cfg.device)
         dist.all_reduce(total, op=dist.ReduceOp.SUM, group=self.ep_group)
         self.assertEqual(int(total.item()), self.cfg.world_size * TOKENS_PER_RANK * TOP_K)
 
+    @_eager_test_include
+    def test_eager_recv_sizing(self):
+        """Eager mode sizes dispatch outputs to the exact per-step recv-token total."""
+        if not EAGER:
+            self.skipTest("eager-only assertions")
+        buf = self._make_buffer()
+        topk_idx, tokens, w = _make_identity_inputs(self.cfg.rank, self.cfg.ep_size)
+        recv_t, recv_w, tokens_per_expert = ep_dispatch(buf, tokens, topk_idx, w)
+        torch.cuda.synchronize()
+        # The per-step recv-token total is exposed on the buffer (int64 [1]).
+        self.assertEqual(buf.total_recv_tokens.dtype, torch.int64)
+        total = int(buf.total_recv_tokens.item())
+        # recv outputs are sized to the recv total, not recv_capacity_per_rank.
+        self.assertEqual(recv_t.shape[0], total)
+        self.assertEqual(recv_w.shape[0], total)
+        # padded total is at least the unpadded per-expert sum and within capacity.
+        self.assertGreaterEqual(total, int(tokens_per_expert.sum().item()))
+        self.assertLessEqual(total, self.cfg.recv_capacity_per_rank)
+
+    @_overflow_test_include
+    def test_overflow_drop(self):
+        """drop_on_overflow: recv past capacity is dropped and dispatch continues
+        instead of trapping; the pre-drop recv total exceeds recv_capacity."""
+        if not OVERFLOW:
+            self.skipTest("overflow-only assertions")
+        buf = self._make_buffer()
+        topk_idx, tokens, w = _make_identity_inputs(self.cfg.rank, self.cfg.ep_size)
+        # Identity routing sends TOKENS_PER_RANK * TOP_K tokens to each rank, which
+        # overflows the deliberately undersized capacity.
+        expected_recv = TOKENS_PER_RANK * TOP_K
+        self.assertGreater(expected_recv, self.cfg.recv_capacity_per_rank)
+        # total_recv_tokens reports the true (pre-drop) recv total, counting the
+        # tokens that will be dropped; the per-expert counts exclude them and sum
+        # to the kept tokens (capped at recv_capacity_per_rank).
+        tokens_per_expert = ep_prepare(buf, topk_idx)
+        torch.cuda.synchronize()
+        self.assertEqual(int(buf.total_recv_tokens.item()), expected_recv)
+        self.assertEqual(int(tokens_per_expert.sum().item()), self.cfg.recv_capacity_per_rank)
+        # Dispatch drops overflowing tokens and completes (no trap); recv outputs
+        # stay capped at recv_capacity_per_rank.
+        recv_t, recv_w, _ = ep_dispatch(buf, tokens, topk_idx, w)
+        torch.cuda.synchronize()
+        self.assertEqual(recv_t.shape[0], self.cfg.recv_capacity_per_rank)
+        self.assertEqual(recv_w.shape[0], self.cfg.recv_capacity_per_rank)
+
     # Identity round-trip via raw primitives
 
     def test_primitive_dispatch_combine_identity(self):
@@ -330,6 +405,7 @@ def test_zero_copy_pool_auto_alloc(self):
 
     # Multi-iter stability
 
+    @_eager_test_include
     def test_dispatch_autograd_multiple_iterations(self):
         """5 fwd+bwd iters on the same EpBuffer must be bit-stable."""
         buf = self._make_buffer()
@@ -479,6 +555,7 @@ def zero_grads():
             )
 
     @_zero_copy_test_include
+    @_eager_test_include
     def test_combine_autograd(self):
         """ep_combine fwd+bwd; bwd grad target is the EpBuffer symm buffer (zc) or in-flight."""
         buf = self._make_buffer()

tests/pytorch/distributed/run_test_ep.sh

@@ -45,11 +45,13 @@ RET=0
 run_pass() {
   local label="$1"
   local zc="$2"
+  local eager="${3:-0}"
+  local overflow="${4:-0}"
   local log="stdout_ep_${label}.txt"
   echo "=== Running ${SCRIPT} [${label}] on ${NUM_RANKS} GPUs (timeout=${TEST_TIMEOUT_S}s) ==="
   # setsid + kill-after so SIGKILL takes down the whole process group, not just torchrun.
-  NVTE_EP_ZERO_COPY="${zc}" setsid timeout --foreground --kill-after=10 --signal=TERM \
-    "${TEST_TIMEOUT_S}" \
+  NVTE_EP_ZERO_COPY="${zc}" NVTE_EP_EAGER="${eager}" NVTE_EP_OVERFLOW="${overflow}" setsid timeout --foreground \
+    --kill-after=10 --signal=TERM "${TEST_TIMEOUT_S}" \
     torchrun --standalone --nnodes=1 --nproc-per-node="${NUM_RANKS}" \
     "${SCRIPT}" 2>&1 | tee "${log}"
   local rc=${PIPESTATUS[0]}
@@ -68,5 +70,7 @@ run_pass() {
 
 run_pass "default" 0
 run_pass "zero_copy" 1
+run_pass "eager" 0 1
+run_pass "overflow" 0 0 1
 
 exit $RET

transformer_engine/common/ep/ep_backend.cpp

@@ -94,8 +94,15 @@ void EPBackend::validate_config(const NVTEEpGroupConfig& config) {
   NVTE_CHECK(config.num_experts > 0, "num_experts must be positive, got ", config.num_experts);
   NVTE_CHECK(config.max_tokens_per_rank > 0, "max_tokens_per_rank must be positive, got ",
              config.max_tokens_per_rank);
-  NVTE_CHECK(config.max_recv_tokens_per_rank > 0, "max_recv_tokens_per_rank must be positive, got ",
+  // 0 selects eager mode (NCCL_EP_AUTO); any explicit budget must be positive.
+  NVTE_CHECK(config.max_recv_tokens_per_rank >= 0,
+             "max_recv_tokens_per_rank must be non-negative, got ",
              config.max_recv_tokens_per_rank);
+  NVTE_CHECK(!(config.zero_copy && config.max_recv_tokens_per_rank == 0),
+             "zero-copy and eager (max_recv_tokens_per_rank = 0) modes are mutually exclusive");
+  NVTE_CHECK(!(config.drop_on_overflow && config.max_recv_tokens_per_rank == 0),
+             "drop_on_overflow (overflow drop) is not supported in eager mode "
+             "(max_recv_tokens_per_rank = 0)");
   NVTE_CHECK(config.hidden_dim > 0, "hidden_dim must be positive, got ", config.hidden_dim);
   NVTE_CHECK(config.max_token_dtype >= 0 && config.max_token_dtype < kNVTENumTypes,
              "max_token_dtype out of range, got ", static_cast<int>(config.max_token_dtype));
@@ -210,9 +217,13 @@ void EPBackend::init(ncclComm_t ep_comm, NVTEEpGroupConfig group_config) {
   cfg.max_num_sms = group_config.num_comm_sms > 0
                         ? static_cast<unsigned int>(group_config.num_comm_sms)
                         : NCCL_EP_AUTO;
-  // Must be > 0; NCCL EP errors out on 0.
+  // 0 = NCCL_EP_AUTO, which enables eager mode (recv buffers sized per routing).
   cfg.max_recv_tokens_per_rank = static_cast<unsigned int>(group_config.max_recv_tokens_per_rank);
   cfg.zero_copy = group_config.zero_copy ? NCCL_EP_ZERO_COPY_ON : NCCL_EP_ZERO_COPY_OFF;
+  // Per-token top-k; NCCL EP sizes internal buffers from it in eager mode.
+  cfg.num_topk = static_cast<unsigned int>(group_config.num_topk);
+  cfg.overflow_policy =
+      group_config.drop_on_overflow ? NCCL_EP_OVERFLOW_DROP : NCCL_EP_OVERFLOW_AUTO;
 
   NVTE_CHECK_NCCL(ncclEpCreateGroup(&ep_group_, ep_comm, &cfg));
 
@@ -320,27 +331,33 @@ size_t EPBackend::handle_mem_size(NVTEEpLayerConfig layer_cfg) {
 }
 
 void EPBackend::prepare(void* handle_mem, const NVTETensor topk_idx,
-                        NVTETensor recv_tokens_per_expert,
-                        NVTETensor /*total_recv_tokens_per_rank*/, NVTEEpLayerConfig layer_cfg,
-                        cudaStream_t stream) {
-  // total_recv_tokens_per_rank is a reserved placeholder; not yet populated.
+                        NVTETensor recv_tokens_per_expert, NVTETensor total_recv_tokens_per_rank,
+                        NVTEEpLayerConfig layer_cfg, cudaStream_t stream) {
   NVTE_CHECK(handle_mem != nullptr, "handle_mem must not be null");
   NVTE_CHECK(layer_cfg.top_k > 0, "top_k must be > 0, got ", layer_cfg.top_k);
   NVTE_CHECK(nvte_tensor_shape(topk_idx).ndim == 2, "topk_idx must be 2D [T, top_k]");
 
   NVTEShape topk_idx_shape;
   ncclEpTensor_t nccl_topk_idx = make_nccl_ep_tensor(topk_idx, topk_idx_shape);
 
-  // ncclEpUpdateHandle writes per-expert counts via expert_counters.
+  // ncclEpUpdateHandle writes per-expert counts via expert_counters and, when
+  // provided, the scalar padded recv-slot total via recv_total_counter.
   NVTEShape recv_tokens_per_expert_shape;
   ncclEpTensor_t recv_tokens_per_expert_desc;
   if (recv_tokens_per_expert != nullptr) {
     recv_tokens_per_expert_desc =
         make_nccl_ep_tensor(recv_tokens_per_expert, recv_tokens_per_expert_shape);
   }
+  NVTEShape total_recv_shape;
+  ncclEpTensor_t total_recv_desc;
+  if (total_recv_tokens_per_rank != nullptr) {
+    total_recv_desc = make_nccl_ep_tensor(total_recv_tokens_per_rank, total_recv_shape);
+  }
   ncclEpLayoutInfo_t layout_info = NCCL_EP_LAYOUT_INFO_INIT;
   layout_info.expert_counters =
       (recv_tokens_per_expert != nullptr) ? &recv_tokens_per_expert_desc : nullptr;
+  layout_info.recv_total_counter =
+      (total_recv_tokens_per_rank != nullptr) ? &total_recv_desc : nullptr;
 
   std::lock_guard<std::mutex> lock(mutex_);
   NVTE_CHECK(initialized_, "EPBackend not initialized");

transformer_engine/common/include/transformer_engine/ep.h

@@ -44,7 +44,8 @@ typedef struct {
   int num_experts;
   /*! Upper bound on tokens this rank sends per dispatch. */
   int max_tokens_per_rank;
-  /*! Upper bound on tokens this rank receives per dispatch (must be > 0). */
+  /*! Upper bound on tokens this rank receives per dispatch. 0 selects eager
+   *  mode: the caller sizes recv buffers to the per-routing recv count. */
   int max_recv_tokens_per_rank;
   /*! Token hidden dimension. */
   int hidden_dim;
@@ -58,6 +59,12 @@ typedef struct {
    *  by NVTECommWindow handles and transfer in place (no staging copies);
    *  0 (default) = staged. */
   int zero_copy;
+  /*! Per-token top-k; sizes NCCL EP internal buffers. Required in eager mode
+   *  (max_recv_tokens_per_rank == 0); 0 = unset. */
+  int num_topk;
+  /*! Recv overflow policy. Nonzero drops tokens past max_recv_tokens_per_rank
+   *  and continues; 0 (default) traps. Not supported in eager mode. */
+  int drop_on_overflow;
 } NVTEEpGroupConfig;
 
 /*! \brief Per-layer configuration consumed by nvte_ep_handle_mem_size and
@@ -121,13 +128,14 @@ size_t nvte_ep_handle_mem_size(const NVTEEpLayerConfig* layer_cfg);
  *  AllGathers topk_idx across the EP group and stages per-expert offsets and
  *  counts into handle_mem so the matching dispatch/combine/_bwd can run with
  *  no further routing computation. Must precede every dispatch/combine/_bwd
- *  that uses this handle_mem. recv_tokens_per_expert becomes host-valid after a
- *  stream sync.
+ *  that uses this handle_mem. recv_tokens_per_expert and total_recv_tokens_per_rank
+ *  become host-valid after a stream sync.
  *
  *  \param[in]     handle_mem                 uint8 routing-state buffer.
  *  \param[in]     topk_idx                   [T, top_k] int64 routing indices.
- *  \param[out]    recv_tokens_per_expert     [num_local_experts] int32 counts.
- *  \param[out]    total_recv_tokens_per_rank Reserved placeholder; may be null. Unused for now.
+ *  \param[out]    recv_tokens_per_expert     [num_local_experts] int32/int64 counts.
+ *  \param[out]    total_recv_tokens_per_rank Optional [1] int32/int64 scalar: padded
+ *                                            recv-slot total for this rank. May be null.
  *  \param[in]     layer_cfg                  Per-call layer configuration (struct_size set).
  *  \param[in]     stream                     CUDA stream.
  */