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

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
 
@@ -147,6 +165,9 @@ def setUpClass(cls):
             recv_capacity_per_rank=cls.cfg.recv_capacity_per_rank,
             hidden_dim=HIDDEN_DIM,
             zero_copy=ZERO_COPY,
+            eager=EAGER,
+            max_num_topk=TOP_K,
+            drop_on_overflow=OVERFLOW,
         )
 
     def setUp(self):
@@ -155,6 +176,14 @@ 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(
@@ -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,14 @@ 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;
+  // Upper bound on per-token top-k; required for eager mode with the
+  // expert-major layout, where NCCL EP sizes internal buffers from it.
+  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 +332,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,9 @@ 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 library derives its internal bound and 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 +60,14 @@ typedef struct {
    *  by NVTECommWindow handles and transfer in place (no staging copies);
    *  0 (default) = staged. */
   int zero_copy;
+  /*! Upper bound on per-token top-k across the group's handles. Required for
+   *  eager mode (max_recv_tokens_per_rank == 0) with the expert-major layout,
+   *  where it sizes internal buffers; 0 = unset. */
+  int num_topk;
+  /*! Recv overflow policy. When nonzero, tokens past max_recv_tokens_per_rank
+   *  are dropped and dispatch continues; 0 (default) traps on overflow.
+   *  Not supported in eager mode. */
+  int drop_on_overflow;
 } NVTEEpGroupConfig;
 
 /*! \brief Per-layer configuration consumed by nvte_ep_handle_mem_size and
@@ -121,13 +131,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.
  */

transformer_engine/pytorch/csrc/extensions.h

@@ -676,7 +676,7 @@ void grouped_swizzle_for_gemm(py::handle &tensor, bool rowwise, bool columnwise)
 void ep_initialize(uintptr_t comm_ptr, const std::string &group_name, int64_t num_experts,
                    int64_t max_tokens_per_rank, int64_t max_recv_tokens_per_rank,
                    int64_t hidden_dim, int64_t max_num_sms, pybind11::object max_token_dtype,
-                   bool zero_copy);
+                   bool zero_copy, int64_t max_num_topk, bool drop_on_overflow);
 
 void ep_finalize();
 
@@ -686,8 +686,9 @@ bool ep_get_zero_copy();
 // Returns the handle_mem byte size for the given layer config.
 int64_t ep_handle_mem_size(int64_t top_k, int64_t dispatch_output_per_expert_alignment);
 
-void ep_prepare(at::Tensor handle_mem, at::Tensor topk_idx, at::Tensor token_counts, int64_t top_k,
-                int64_t dispatch_output_per_expert_alignment);
+void ep_prepare(at::Tensor handle_mem, at::Tensor topk_idx, at::Tensor tokens_per_expert,
+                int64_t top_k, int64_t dispatch_output_per_expert_alignment,
+                at::Tensor total_recv_tokens);
 
 void ep_dispatch(at::Tensor handle_mem, at::Tensor topk_idx, at::Tensor tokens,
                  at::Tensor topk_weights, at::Tensor recv_tokens, at::Tensor recv_topk_weights);