perf(moe): sync-free HybridEP dispatch (capacity mode, equal token counts)

nemo_automodel/components/models/common/utils.py

@@ -442,6 +442,15 @@ class BackendConfig:
             manager instance across MoE layers.
         dispatcher_async_dispatch: Whether DeepEP/UCCL-EP dispatch and combine should return
             asynchronously and allocate their outputs on the communication stream.
+        dispatcher_capacity_factor: HybridEP only, dynamic routing. Run HybridEP dispatch in its
+            non-blocking mode with output buffers sized to the first microbatch's permuted row
+            count times this factor (EP-group max, 4-token aligned) instead of letting every
+            dispatch drain the compute stream to read the exact count. Overflow trips a
+            device-side assert. None (default) keeps the blocking reference path.
+        dispatcher_equal_token_counts: HybridEP only. Declare that every EP rank dispatches the
+            same row count, so the per-dispatch EP-group max all-reduce (and its host sync)
+            that derives the pad size is skipped and the count is only aligned. Default False;
+            keep it False for variable-length or in-batch-packed inputs.
         enable_deepep: Removed and ignored. Logs a warning if set; configure "dispatcher"
             and "experts" explicitly instead.
         fake_balanced_gate: If True, replace the learned Gate with FakeBalancedGate
@@ -501,6 +510,17 @@ class BackendConfig:
     dispatcher_num_sms: int = 20
     dispatcher_share_token_dispatcher: bool = True
     dispatcher_async_dispatch: bool = False
+    # HybridEP only, dynamic routing: after one blocking calibration dispatch per MoE layer, size every
+    # later dispatch's output buffers to ceil(calibrated rows x factor) (EP-group max, aligned) and run
+    # HybridEP in its non-blocking mode. Removes the per-dispatch compute-stream drain that HybridEP's
+    # blocking mode needs to learn the permuted row count (and the per-layer barrier it implies); an
+    # overflow of the capacity trips a device-side assert instead of silently truncating. None = blocking.
+    dispatcher_capacity_factor: float | None = None
+    # HybridEP only: every EP rank dispatches the same number of rows (fixed-shape batches, which is
+    # every batch that is not variable-length / in-batch packed), so the per-dispatch EP-group MAX
+    # all-reduce + int() host sync that derives the pad size is skipped and the local count is only
+    # aligned. Leave False for variable-length inputs: unequal counts abort the HybridEP collective.
+    dispatcher_equal_token_counts: bool = False
     mok: MoKBackendConfig = field(default_factory=MoKBackendConfig)
     enable_deepep: bool | None = None  # Removed: ignored with a warning; set dispatcher/experts explicitly
     fake_balanced_gate: bool = False

nemo_automodel/components/moe/experts.py

@@ -959,6 +959,15 @@ def __init__(
         self.dispatcher_num_sms = dispatcher_num_sms
         self.dispatcher_share_token_dispatcher = dispatcher_share_token_dispatcher
         self.dispatcher_async_dispatch = dispatcher_async_dispatch
+        # HybridEP capacity mode (BackendConfig.dispatcher_capacity_factor): the dispatcher returns
+        # device-side tokens_per_expert and buffers of a fixed capacity, so the per-microbatch
+        # count_nonzero host read below is skipped as well (rows are never empty).
+        self.dispatcher_capacity_factor = (
+            getattr(backend, "dispatcher_capacity_factor", None) if backend is not None else None
+        )
+        self.dispatcher_equal_token_counts = (
+            bool(getattr(backend, "dispatcher_equal_token_counts", False)) if backend is not None else False
+        )
 
         # Allocate projection tensor - size depends on whether activation is gated
         # Gated (SwiGLU, Quick-GEGLU): [n_experts, dim, 2*inter_dim]
@@ -992,6 +1001,8 @@ def init_token_dispatcher(self, ep_mesh: DeviceMesh):
             moe_hybridep_num_sms=self.dispatcher_num_sms,
             moe_share_token_dispatcher=self.dispatcher_share_token_dispatcher,
             moe_deepep_async_dispatch=self.dispatcher_async_dispatch,
+            moe_hybridep_capacity_factor=self.dispatcher_capacity_factor,
+            moe_hybridep_equal_token_counts=self.dispatcher_equal_token_counts,
             moe_benchmark_static_routing=self.static_routing,
         )
 
@@ -1072,7 +1083,11 @@ def forward(
         # With static routing (forced balance, no noise) every expert receives tokens by
         # construction, so the count_nonzero device-to-host read (one per microbatch, and
         # again per activation-checkpoint recompute) can be skipped.
-        if self.static_routing or torch.count_nonzero(tokens_per_expert) > 0:
+        if (
+            self.static_routing
+            or self.dispatcher_capacity_factor is not None
+            or torch.count_nonzero(tokens_per_expert) > 0
+        ):
             if self.use_torch_mm:
                 tokens_per_expert_gpu = tokens_per_expert.to(
                     device=permuted_local_hidden_states.device, non_blocking=True

nemo_automodel/components/moe/megatron/fused_a2a.py

@@ -140,8 +140,15 @@ def __init__(self) -> None:
         self._cursor = 0
         self.replay_misses = 0
 
-    def record(self, handle, tokens_per_expert) -> None:
-        self._records.append([handle, tokens_per_expert, None])
+    def record(self, handle, tokens_per_expert, num_permuted_tokens=None) -> None:
+        """Log one checkpoint-forward dispatch.
+
+        When the forward already ran with a host-side extent (capacity mode, static-routing pin),
+        that integer is the extent its output was sized to and the one the replay must reuse;
+        recording it also spares ``finalize`` the device-to-host reduction for that dispatch.
+        """
+        extent = num_permuted_tokens if isinstance(num_permuted_tokens, int) else None
+        self._records.append([handle, tokens_per_expert, extent])
 
     def finalize(self) -> None:
         """Cache each layout extent after the checkpoint-forward op context exits."""
@@ -643,7 +650,7 @@ def forward(
         if recorder is not None and _hybridep_dispatch_replay_state.mode == "record":
             # Keep only the reusable layout and its output extent. Recomputed
             # activations and probabilities are still redispatched through it.
-            recorder.record(handle, tokens_per_expert)
+            recorder.record(handle, tokens_per_expert, num_permuted_tokens)
         return (
             dispatched_hidden,
             dispatched_probs,
@@ -686,6 +693,8 @@ def backward(ctx, grad_x):
             handle=handle,
             pad_multiple=ctx.pad_multiple,
             num_permuted_tokens=ctx.num_permuted_tokens,
+            # Capacity mode hands a host int: the backward dispatch then needs no stream drain either.
+            non_blocking=isinstance(ctx.num_permuted_tokens, int),
         )
         return dispatched_hidden, None, None, None
 

nemo_automodel/components/moe/megatron/token_dispatcher.py

@@ -14,6 +14,7 @@
 # limitations under the License.
 
 import logging
+import math
 import os
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
@@ -361,6 +362,17 @@ def forward(self, token_indices: torch.Tensor, token_probs: torch.Tensor) -> tup
 _STATIC_ROUTING_PAD_PIN = os.environ.get("NEMO_STATIC_ROUTING_PAD_PIN", "1") != "0"
 
 
+def _assert_no_hybridep_overflow(handle, capacity: int) -> None:
+    """Device-side guard: HybridEP truncates silently when the capacity is exceeded and sets
+    ``overflow_flag`` (handle item 10); fail loudly instead of training on dropped tokens."""
+    flag = handle[10] if isinstance(handle, (tuple, list)) and len(handle) > 10 else None
+    if torch.is_tensor(flag):
+        torch._assert_async(
+            (flag == 0).reshape(()),
+            f"HybridEP dispatch overflowed its capacity of {capacity} permuted rows; raise BackendConfig.dispatcher_capacity_factor",
+        )
+
+
 class _HybridEPManager(_DispatchManager):
     """
     A manager class to handle fused all-to-all communication processes for MoE models using
@@ -385,13 +397,21 @@ def __init__(
         permute_fusion: bool = False,
         moe_hybridep_num_sms: int = 24,
         benchmark_static_routing: bool = False,
+        moe_hybridep_capacity_factor: float | None = None,
+        moe_hybridep_equal_token_counts: bool = False,
     ):
         self.group = group
         self.num_local_experts = num_local_experts
         self.num_experts = num_experts
         self.router_topk = router_topk
         self.permute_fusion = permute_fusion
         self.moe_hybridep_num_sms = moe_hybridep_num_sms
+        # Capacity mode (dynamic routing): after one blocking calibration dispatch, later dispatches
+        # pass this many rows as num_permuted_tokens and run non-blocking (see dispatch()).
+        self.hybridep_capacity_factor = moe_hybridep_capacity_factor
+        self._hybridep_capacity: int | None = None
+        # Equal row counts across the EP group: the pad size is the aligned local count, no collective.
+        self.equal_token_counts = moe_hybridep_equal_token_counts
         # Benchmark-only (TokenDispatcherConfig.moe_benchmark_static_routing):
         # persist num_permuted_tokens across dispatches, see dispatch()/reset.
         self.benchmark_static_routing = benchmark_static_routing
@@ -481,7 +501,11 @@ def dispatch(
         if torch.distributed.is_initialized() and torch.distributed.get_world_size(self.group) > 1:
             num_tokens = hidden_states.shape[0]
             pin = self.benchmark_static_routing and _STATIC_ROUTING_PAD_PIN
-            if pin and self._static_target_tokens is not None and self._static_target_tokens >= num_tokens:
+            if self.equal_token_counts:
+                # Every rank holds the same row count by construction (fixed-shape batches): the
+                # group-wide maximum is the local count, so only align it. No collective, no host sync.
+                target_tokens = -(-num_tokens // _HYBRIDEP_TOKEN_ALIGNMENT) * _HYBRIDEP_TOKEN_ALIGNMENT
+            elif pin and self._static_target_tokens is not None and self._static_target_tokens >= num_tokens:
                 target_tokens = self._static_target_tokens
             else:
                 group_max = torch.tensor(num_tokens, device=hidden_states.device)
@@ -502,6 +526,12 @@ def dispatch(
                 self.routing_map = nn.functional.pad(self.routing_map, (0, 0, 0, pad_tokens))
                 self.token_probs = nn.functional.pad(self.token_probs, (0, 0, 0, pad_tokens))
 
+        capacity_mode = self.hybridep_capacity_factor is not None and not self.benchmark_static_routing
+        if capacity_mode and self._hybridep_capacity is not None:
+            # Non-blocking HybridEP dispatch: the buffers are sized to the calibrated capacity, the
+            # real counts stay on the device, and no compute-stream drain happens on this call.
+            self.num_permuted_tokens = self._hybridep_capacity
+
         dispatched_hidden, self.dispatched_probs, _, tokens_per_expert, self.handle = hybrid_ep_dispatch(
             x=hidden_states,
             routing_map=self.routing_map,
@@ -515,12 +545,46 @@ def dispatch(
         )
 
         self.tokens_per_expert = tokens_per_expert
+        if capacity_mode:
+            if self._hybridep_capacity is None:
+                self._calibrate_hybridep_capacity(tokens_per_expert, dispatched_hidden.device)
+            else:
+                _assert_no_hybridep_overflow(self.handle, self._hybridep_capacity)
+                self.num_permuted_tokens = self._hybridep_capacity
+            return dispatched_hidden
         self.num_permuted_tokens = self.tokens_per_expert.sum()
         if self.benchmark_static_routing and getattr(self, "_static_num_permuted_tokens", None) is None:
             self._static_num_permuted_tokens = self.num_permuted_tokens
 
         return dispatched_hidden
 
+    def _calibrate_hybridep_capacity(self, tokens_per_expert: torch.Tensor, device: torch.device) -> None:
+        """Turn the one blocking calibration dispatch into the capacity every later dispatch uses.
+
+        The blocking dispatch already brought ``tokens_per_expert`` to the host, so reading its sum is
+        free; the capacity is that count times ``hybridep_capacity_factor``, aligned to the HybridEP
+        token alignment (and ``pad_multiple``), taken as the EP-group maximum once so every rank
+        allocates the same buffers.
+        """
+        actual = int(tokens_per_expert.sum())
+        align = max(int(self.pad_multiple or 0), _HYBRIDEP_TOKEN_ALIGNMENT)
+        cap = -(-int(math.ceil(actual * self.hybridep_capacity_factor)) // align) * align
+        if torch.distributed.is_initialized() and torch.distributed.get_world_size(self.group) > 1:
+            cap_t = torch.tensor(cap, device=device)
+            torch.distributed.all_reduce(cap_t, op=torch.distributed.ReduceOp.MAX, group=self.group)
+            cap = int(cap_t)
+        self._hybridep_capacity = cap
+        # This dispatch itself was blocking: its combine (and the combine's backward dispatch) can
+        # use the exact host-side count.
+        self.num_permuted_tokens = actual
+        if not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0:
+            logging.getLogger(__name__).info(
+                "HybridEP capacity mode: calibrated %d permuted rows x %.2f -> capacity %d; later dispatches run non-blocking",
+                actual,
+                self.hybridep_capacity_factor,
+                cap,
+            )
+
     def combine(
         self,
         hidden_states: torch.Tensor,
@@ -583,6 +647,12 @@ class TokenDispatcherConfig:
     None means no changes for dtype."""
 
     moe_flex_dispatcher_backend: Literal["deepep", "hybridep", "uccl_ep"] = "deepep"
+
+    # HybridEP capacity mode, see BackendConfig.dispatcher_capacity_factor
+
+    moe_hybridep_capacity_factor: float | None = None
+    # HybridEP: skip the per-dispatch pad-size all-reduce, see BackendConfig.dispatcher_equal_token_counts
+    moe_hybridep_equal_token_counts: bool = False
     """Backend for the flex token dispatcher. Options: 'deepep', 'hybridep', or 'uccl_ep'."""
 
     moe_deepep_num_sms: int = 20
@@ -705,6 +775,8 @@ def __init__(
                         permute_fusion=self.config.moe_permute_fusion,
                         moe_hybridep_num_sms=self.config.moe_hybridep_num_sms,
                         benchmark_static_routing=self.config.moe_benchmark_static_routing,
+                        moe_hybridep_capacity_factor=self.config.moe_hybridep_capacity_factor,
+                        moe_hybridep_equal_token_counts=self.config.moe_hybridep_equal_token_counts,
                     )
                 self._comm_manager = MoEFlexTokenDispatcher.shared_hybridep_manager
             else:
@@ -716,6 +788,8 @@ def __init__(
                     permute_fusion=self.config.moe_permute_fusion,
                     moe_hybridep_num_sms=self.config.moe_hybridep_num_sms,
                     benchmark_static_routing=self.config.moe_benchmark_static_routing,
+                    moe_hybridep_capacity_factor=self.config.moe_hybridep_capacity_factor,
+                    moe_hybridep_equal_token_counts=self.config.moe_hybridep_equal_token_counts,
                 )
             self.hybridep_metadata_processor = _HybridEPMetadataProcessor(
                 num_experts=self.tp_size * self.config.num_moe_experts,

tests/unit_tests/moe/test_backend_config.py

@@ -492,3 +492,17 @@ def test_preprocess_requires_router_and_hybridep(self):
                 dispatcher="deepep",
                 cuda_graph=CudaGraphConfig(modules=["moe_router", "moe_preprocess"]),
             )
+
+
+def test_dispatcher_capacity_factor_defaults_off():
+    from nemo_automodel.components.models.common.utils import BackendConfig
+
+    assert BackendConfig().dispatcher_capacity_factor is None
+    assert BackendConfig(dispatcher_capacity_factor=1.5).dispatcher_capacity_factor == 1.5
+
+
+def test_dispatcher_equal_token_counts_defaults_off():
+    from nemo_automodel.components.models.common.utils import BackendConfig
+
+    assert BackendConfig().dispatcher_equal_token_counts is False
+    assert BackendConfig(dispatcher_equal_token_counts=True).dispatcher_equal_token_counts is True

tests/unit_tests/moe/test_fused_a2a.py

@@ -112,7 +112,7 @@ def combine_with_unpermute(self, *, hidden, probs=None, **kwargs):
         return combined_hidden, combined_probs
 
 
-def _run_checkpointed_hybridep(context_fn):
+def _run_checkpointed_hybridep(context_fn, num_permuted_tokens=None):
     x = torch.randn(4, 3, requires_grad=True)
     routing_map = torch.ones(4, 2, dtype=torch.bool)
     probs = torch.full((4, 2), 0.5, requires_grad=True)
@@ -126,7 +126,7 @@ def block(hidden, token_probs):
             1,
             24,
             24,
-            None,
+            num_permuted_tokens,
             None,
         )
         return dispatched_hidden.sin().sum() + dispatched_probs.square().sum()
@@ -181,3 +181,32 @@ def save_replay_sensitive_ops(ctx, func, *args, **kwargs):
     assert buffer.cached_dispatches == 1
     assert buffer.replayed_num_permuted_tokens == 5
     assert isinstance(buffer.replayed_num_permuted_tokens, int)
+
+
+def test_hybridep_recorder_keeps_a_host_extent_without_reducing_tokens_per_expert():
+    recorder = fused_a2a.HybridEPDispatchReplayRecorder()
+    tokens_per_expert = mock.MagicMock(spec=torch.Tensor)
+    # capacity mode / static pin: the forward already ran with a host-side extent
+    recorder.record("layout", tokens_per_expert, 24)
+    # blocking dispatch: the extent comes from the reduction, after the checkpoint context exits
+    recorder.record("layout", torch.tensor([2, 3]))
+    recorder.finalize()
+
+    assert recorder.take() == ["layout", tokens_per_expert, 24]
+    assert recorder.take()[2] == 5
+    tokens_per_expert.sum.assert_not_called()
+
+
+def test_hybridep_checkpoint_replay_reuses_the_forward_capacity_extent():
+    from nemo_automodel.components.moe.parallelizer import _replay_hybridep_dispatch_on_recompute
+
+    buffer = _DriftingHybridEPBuffer()
+    fused_a2a._hybrid_ep_buffer = buffer
+    context_fn = _replay_hybridep_dispatch_on_recompute(lambda: (nullcontext(), nullcontext()))
+
+    _run_checkpointed_hybridep(context_fn, num_permuted_tokens=24)
+
+    assert buffer.full_dispatches == 1
+    assert buffer.cached_dispatches == 1
+    # the recompute output must be sized like the forward's (capacity rows), not to this dispatch's token count
+    assert buffer.replayed_num_permuted_tokens == 24