Inference: Do not route pad/dummy tokens to any expert

megatron/core/inference/contexts/dynamic_context.py

@@ -631,6 +631,12 @@ def __init__(self, model_config: TransformerConfig, inference_config: InferenceC
             and model_config.inference_moe_token_dispatcher_type == 'nccl'
         )
 
+        # are we using the inference_optimized nvls ep dispatcher for MoEs?
+        self._nvls_dispatcher = (
+            get_pg_size(self.expert_model_parallel_group) > 1
+            and model_config.inference_moe_token_dispatcher_type == 'nvls'
+        )
+
         # are we using the training a2a dispatcher for MoEs?
         # Note that this is not optimal for speed.
         self._training_ep_dispatcher = (
@@ -675,23 +681,24 @@ def __init__(self, model_config: TransformerConfig, inference_config: InferenceC
 
         # Allocate per-step dispatcher buffers upfront so update_metadata never
         # triggers an allocation inside a captured CUDA graph.
-        # Both dispatchers need _valid_tokens_tensor initialized even at EP=1:
-        # mcore_fused_moe's Triton kernel reads it as a pointer regardless of EP size.
-        if model_config.inference_moe_token_dispatcher_type == 'nccl':
+        #
+        # The shared _valid_tokens_tensor scalar is read as a pointer by both fused
+        # MoE backends (mcore_fused_moe and vllm_fused_moe) regardless of EP size, so
+        # allocate it unconditionally (covers EP=1, where no dispatcher comm buffers
+        # exist). The EP>1 dispatchers below reallocate it as part of their own buffer
+        # setup, which is harmless.
+        InferenceAllGatherDispatcherBase.allocate_valid_tokens_tensor()
+        if self._nccl_ep_dispatcher:
             NCCLAllGatherDispatcher.allocate_buffers()
-        elif get_pg_size(self.expert_model_parallel_group) > 1:
-            # Use moe_latent_size if set, else hidden_size.
+        elif self._nvls_dispatcher:
+            # Use moe_latent_size if set (latent MoE: SuperV3, UltraV3), else hidden_size.
             moe_hidden_size = model_config.moe_latent_size or model_config.hidden_size
             NVLSAllGatherVDispatcher.allocate_buffers(
                 per_rank_worst_case_token_count=self.round_up_tokens(self.max_tokens) // tp_size,
                 topk=model_config.moe_router_topk,
                 hidden_size=moe_hidden_size,
                 ep_group=self.expert_model_parallel_group,
             )
-        else:
-            # EP=1 with nvls: skip symmetric memory init (requires NVLink between
-            # multiple GPUs) and just initialize the shared valid_tokens scalar.
-            InferenceAllGatherDispatcherBase.allocate_valid_tokens_tensor()
 
         # Deal with chunked prefill
         self.enable_chunked_prefill = inference_config.enable_chunked_prefill
@@ -710,8 +717,17 @@ def __init__(self, model_config: TransformerConfig, inference_config: InferenceC
 
         # Allocate GPU state.
         self.is_tensor_state_allocated = False
+        self._bookkeeping_no_real_work = False
         self.initialize_all_tensors()
 
+        # Bind the GPU real-token-count tensor onto the NVLS dispatcher class
+        # so it can mask out CUDA-graph padding tokens during routing. The
+        # tensor lives inside gpu_view._buf (fixed address) and is refreshed
+        # each step by transfer_bookkeeping_to_gpu(). NVLS-only — the NCCL
+        # dispatcher requires equal token counts across ranks already.
+        if self._nvls_dispatcher:
+            NVLSAllGatherVDispatcher.set_real_token_count_tensor(self.gpu_view.real_token_count)
+
         # Print info.
         active_blocks = self.kv_block_allocator.active_count
         total_blocks = self.kv_block_allocator.total_count
@@ -1017,6 +1033,10 @@ def initialize_all_tensors(self) -> None:
         _tok_int32_bytes = self.max_tokens * 4
         # Request-level fields are all 4 bytes wide (5 int32 + 2 float32 = 7 fields).
         _req_4byte_bytes = self.max_requests * 4
+        # Scalar: real (unpadded) token count for the current step. Refreshed
+        # in transfer_bookkeeping_to_gpu(); read on GPU via
+        # `gpu_view.real_token_count` (MoE routing masks padding tokens).
+        _real_token_count_bytes = 4
         # MHA section: 5 fields (int32) shared between GraphedMHAMetadata and
         # NonGraphedMHAMetadata. max_bs == max_requests.
         _mha_query_lengths_bytes = self.max_requests * 4
@@ -1028,6 +1048,7 @@ def initialize_all_tensors(self) -> None:
             3 * _tok_int64_bytes
             + 3 * _tok_int32_bytes
             + 7 * _req_4byte_bytes
+            + _real_token_count_bytes
             + _mha_query_lengths_bytes
             + _mha_cu_query_seq_lengths_bytes
             + _mha_kv_seq_lengths_bytes
@@ -1154,6 +1175,14 @@ def initialize_all_tensors(self) -> None:
         ].view(torch.int32)
         _off += _req_4byte_bytes
 
+        # Scalar staging slot for the real (unpadded) token count. Refreshed
+        # from `self.batch_dimensions.token_count` in transfer_bookkeeping_to_gpu()
+        # and read on GPU via `gpu_view.real_token_count`.
+        self._staging_real_token_count = self._cpu_bookkeeping_buf[
+            _off : _off + _real_token_count_bytes
+        ].view(torch.int32)
+        _off += _real_token_count_bytes
+
         # Static tensor addresses to make `last_token_logits` graphable with speculative decoding.
         max_logit_idxs = self.max_requests * (self.num_speculative_tokens + 1)
         self.active_logit_idxs = torch.zeros(
@@ -2377,6 +2406,13 @@ def initialize_attention_state(
         # No-op when the queue is already empty (regular non-warmup steps).
         self._execute_pending_mamba_ops()
 
+        # Record whether this step produces real output — false on CUDA-graph
+        # capture (warmup) or dummy EP steps. Used by transfer_bookkeeping_to_gpu
+        # to publish real_token_count=0 so MoE routing masks all padding tokens.
+        self._bookkeeping_no_real_work = (
+            construct_graph_dimensions is not None or is_expert_parallel_dummy_cuda_graph_step
+        )
+
         # Run the H2D transfer here so callers that bypass the controller
         # (e.g. unit tests that call `model.forward()` directly after
         # `initialize_attention_state()`) see populated GPU bookkeeping. The
@@ -2453,6 +2489,15 @@ def transfer_bookkeeping_to_gpu(self) -> None:
             self._staging_request_query_lengths[n_active:padded_active] = 0
             self._staging_request_kv_length_offsets[n_active:padded_active] = 0
 
+        # Real (unpadded) token count for this step. CUDA-graph replay pads
+        # the token dim to a captured size; MoE routing reads this on GPU and
+        # rewrites padding rows' routing entries to -1 so they don't go to
+        # any expert. Set to 0 on CUDA-graph capture / dummy EP steps so
+        # every row gets masked out.
+        self._staging_real_token_count[0] = (
+            0 if self._bookkeeping_no_real_work else self.batch_dimensions.token_count
+        )
+
         # Coalesced H2D: one cudaMemcpyAsync for the entire bookkeeping buffer.
         # Copying the whole (max_tokens + max_requests)-sized buffer including
         # unused slots is cheap (~71 KB total, ~3-5 us on PCIe Gen4) and saves

megatron/core/inference/contexts/gpu_view.py

@@ -43,6 +43,9 @@ def __init__(
         # query_lengths, kv_length_offsets) + 1 int32 (top_k) + 2 float32
         # (temperature, top_p) + 1 int32 (active_request_last_token_idxs) = 7 fields.
         req_4byte_bytes = max_requests * 4
+        # Scalar: real (unpadded) token count for the current step. Used by
+        # MoE routing to mask out CUDA-graph padding tokens.
+        real_token_count_bytes = 4
 
         # MHA section: 5 fields shared by both graphed and non-graphed MHAMetadata
         # (only one is active per step, so sharing storage is fine).
@@ -74,6 +77,7 @@ def __init__(
             3 * tok_int64_bytes
             + 3 * tok_int32_bytes
             + 7 * req_4byte_bytes
+            + real_token_count_bytes
             + mha_query_lengths_bytes
             + mha_cu_query_seq_lengths_bytes
             + mha_kv_seq_lengths_bytes
@@ -163,6 +167,12 @@ def __init__(
         )
         off += req_4byte_bytes
 
+        # Real (unpadded) token count for the current step. Scalar int32 view.
+        # MoE routing reads this to skip routing CUDA-graph padding tokens to
+        # experts. Refreshed each step by transfer_bookkeeping_to_gpu().
+        self.real_token_count = self._buf[off : off + real_token_count_bytes].view(torch.int32)
+        off += real_token_count_bytes
+
         # MHA flash-attention metadata (shared between GraphedMHAMetadata and
         # NonGraphedMHAMetadata — only one is active per step).
         self.mha_query_lengths = self._buf[off : off + mha_query_lengths_bytes].view(torch.int32)

megatron/core/inference/text_generation_controllers/text_generation_controller.py

@@ -42,6 +42,7 @@
 from megatron.core.transformer.moe.moe_layer import BaseMoELayer
 from megatron.core.transformer.moe.router_replay import RouterReplay, RouterReplayAction
 from megatron.core.transformer.moe.router_trace import get_moe_router_tracer
+from megatron.core.transformer.moe.token_dispatcher_inference import NVLSAllGatherVDispatcher
 from megatron.core.transformer.utils import set_model_to_sequence_parallel
 from megatron.core.utils import (
     accepts_parameter,
@@ -962,6 +963,14 @@ def _compute_serial_mtp_and_sample(self):
         position_ids_buf[0, active_request_count:] = 0
 
         nvtx_range_pop("mtp-spec-decoding/serial-mtp-init")
+
+        # MTP MoE forwards are request-count shaped: the routing map holds
+        # active_request_count real rows followed by padding up to padded_count.
+        # The NVLS routing mask defaults to the main step's token count, so point
+        # it at the MTP row count instead, else padding rows route to experts.
+        if context._nvls_dispatcher:
+            NVLSAllGatherVDispatcher.modify_real_token_count_for_mtp(active_request_count)
+
         for depth in range(self.num_mtp_depths):
             nvtx_range_push(f"mtp-spec-decoding/depth-{depth}")
 

megatron/core/transformer/moe/inference_routing_mask_kernel.py

@@ -0,0 +1,101 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+
+"""Triton kernel for masking CUDA-graph padding rows of a local routing map.
+
+Under CUDA-graph capture the local token count is padded up to a captured
+graph size; those padding rows have garbage routing indices and, if left
+alone, would dispatch padding tokens to real experts. This kernel zeroes
+that out by writing ``-1`` into every topk slot of rows in
+``[real_token_count, local_tokens)``.
+
+The kernel reads ``real_token_count`` from a fixed-address ``int32[1]`` GPU
+tensor, so it is safe to call from inside a captured graph: only the value
+behind the pointer changes between replays.
+"""
+
+from torch import Tensor
+
+try:
+    import triton
+    import triton.language as tl
+
+    HAVE_TRITON = True
+except ImportError:
+    from unittest.mock import MagicMock
+
+    from megatron.core.utils import null_decorator
+
+    triton = MagicMock()
+    triton.jit = null_decorator
+    triton.autotune = null_decorator
+    tl = MagicMock()
+    HAVE_TRITON = False
+
+
+@triton.jit
+def _mask_routing_padding_kernel(
+    routing_map_ptr,  # int64* [total_rows, topk]
+    real_token_count_ptr,  # int32* [1]
+    total_rows: tl.int32,
+    tp_rank: tl.int32,  # SP/TP rank — local row r maps to global row r + tp_rank*total_rows
+    TOPK: tl.constexpr,  # actual topk
+    BLOCK_M: tl.constexpr,  # rows per program
+    BLOCK_TOPK: tl.constexpr,  # next_power_of_2(TOPK), column block
+):
+    """Fill `routing_map[real_token_count:, :]` with -1, BLOCK_M rows per program."""
+    pid = tl.program_id(0)
+    rows = pid * BLOCK_M + tl.arange(0, BLOCK_M)
+
+    real_count = tl.load(real_token_count_ptr).to(tl.int32)
+
+    # real_count is in the global (pre-SP-shard) frame; rows is local to this SP rank.
+    global_rows = rows + tp_rank * total_rows
+    row_mask = (global_rows >= real_count) & (rows < total_rows)
+
+    cols = tl.arange(0, BLOCK_TOPK)
+    col_mask = cols < TOPK
+
+    offs = rows[:, None].to(tl.int64) * TOPK + cols[None, :].to(tl.int64)
+    mask = row_mask[:, None] & col_mask[None, :]
+
+    neg_one = tl.full((BLOCK_M, BLOCK_TOPK), -1, dtype=tl.int64)
+    tl.store(routing_map_ptr + offs, neg_one, mask=mask)
+
+
+def mask_routing_padding(
+    routing_map: Tensor, real_token_count_tensor: Tensor, tp_rank: int = 0
+) -> None:
+    """In-place fill -1 into ``routing_map[real_token_count:, :]``.
+
+    Args:
+        routing_map: ``[N, topk]`` int64 local routing map. ``N`` is the
+            (possibly CUDA-graph-padded) local token count.
+        real_token_count_tensor: ``[1]`` int32 GPU tensor holding the real
+            (unpadded) token count for this step, in the global (pre-SP-shard)
+            frame. Read inside the kernel so the mask boundary moves correctly
+            across CUDA-graph replays.
+        tp_rank: This rank's index in the SP/TP group. Local row ``r`` is
+            row ``r + tp_rank * N`` in the global frame; the kernel uses this
+            offset to compare against ``real_token_count_tensor``.
+    """
+    assert routing_map.is_cuda, "routing_map must be on CUDA"
+    assert routing_map.dim() == 2, f"expected 2D routing_map, got {routing_map.shape}"
+    assert routing_map.dtype.is_floating_point is False, "routing_map must be integer"
+
+    total_rows, topk = routing_map.shape
+    if total_rows == 0:
+        return
+
+    BLOCK_M = 8 if total_rows < 64 else 128
+    BLOCK_TOPK = triton.next_power_of_2(topk)
+    grid = (triton.cdiv(total_rows, BLOCK_M),)
+
+    _mask_routing_padding_kernel[grid](
+        routing_map,
+        real_token_count_tensor,
+        total_rows=total_rows,
+        tp_rank=tp_rank,
+        TOPK=topk,
+        BLOCK_M=BLOCK_M,
+        BLOCK_TOPK=BLOCK_TOPK,
+    )

megatron/core/transformer/moe/token_dispatcher_inference.py

@@ -40,6 +40,7 @@
     gather_from_sequence_parallel_region,
     reduce_scatter_to_sequence_parallel_region,
 )
+from megatron.core.transformer.moe.inference_routing_mask_kernel import mask_routing_padding
 from megatron.core.transformer.moe.shared_experts import SharedExpertMLP
 from megatron.core.transformer.moe.token_dispatcher import MoEAllGatherTokenDispatcher
 from megatron.core.transformer.transformer_config import TransformerConfig
@@ -313,6 +314,12 @@ class NVLSAllGatherVDispatcher(InferenceAllGatherDispatcherBase):
     _step_metadata: Optional[torch.Tensor] = None  # [3] int32
     _per_rank_worst_case_token_count: int = 2048  # round_up_tokens(max_tokens) // tp_size
 
+    # [1] int32 view onto context.gpu_view.real_token_count. Fixed GPU address;
+    # written each step by the context's transfer_bookkeeping_to_gpu(). Holds the
+    # real (unpadded) local token count so the dispatcher can mask routing for
+    # CUDA-graph padding tokens. Wired once by the context after gpu_view init.
+    _real_token_count_tensor: Optional[torch.Tensor] = None
+
     # ── Class-level symmetric buffer handles (allocated once at model init) ───────
     # Dtypes: hidden=bf16, routing=int64, probs=fp32, rsv=fp32.
     _symm_agv_hidden: Optional[dict] = None  # {"tensor": ..., "handle": ...}
@@ -326,6 +333,31 @@ def _get_rsv_tensor(cls) -> Optional[torch.Tensor]:
         unpermute output directly into it, avoiding a copy before RSV."""
         return cls._symm_rsv["tensor"] if cls._symm_rsv is not None else None
 
+    @classmethod
+    def set_real_token_count_tensor(cls, tensor: torch.Tensor) -> None:
+        """Bind the context's GPU real-token-count tensor on the dispatcher class.
+
+        Called once by DynamicInferenceContext after gpu_view is initialised.
+        The tensor is a fixed-address int32[1] view whose value is refreshed
+        each step by transfer_bookkeeping_to_gpu().
+        """
+        cls._real_token_count_tensor = tensor
+
+    @classmethod
+    def modify_real_token_count_for_mtp(cls, mtp_token_count: int) -> None:
+        """Override the routing-mask token count for an MTP forward.
+
+        Each step the context publishes batch_dimensions.token_count into the
+        bound tensor. MTP forwards are request-count shaped, so the controller
+        calls this before an MTP forward to point the mask at the MTP row count
+        instead.
+        """
+        assert cls._real_token_count_tensor is not None, (
+            "real-token-count tensor not wired; DynamicInferenceContext must "
+            "call set_real_token_count_tensor first"
+        )
+        cls._real_token_count_tensor.fill_(mtp_token_count)
+
     @classmethod
     def _rank_token_offset(cls) -> torch.Tensor:
         return cls._step_metadata[1:2]
@@ -343,6 +375,7 @@ def _delete_buffers(cls):
         cls._symm_agv_probs = None
         cls._symm_rsv = None
         cls._symm_metadata = None
+        cls._real_token_count_tensor = None
 
     @classmethod
     def allocate_buffers(
@@ -466,6 +499,10 @@ def __init__(
             runs_metadata_sync=runs_metadata_sync,
         )
         self.topk = config.moe_router_topk
+        # Rank inside pg_collection.tp — the *standard* TP group that SP shards
+        # the routing map along. Base class self.tp_rank is the expt_tp rank,
+        # which is not what we want for the SP padding offset.
+        self.sp_rank = get_pg_rank(pg_collection.tp)
         # Set in dispatch_preprocess; consumed by token_dispatch and token_combine.
         self._local_tokens: int = 0
         # When shared_expert_overlap is enabled, the shared expert forward is launched
@@ -524,6 +561,18 @@ def token_dispatch(self, hidden_states, probs):
         if self._runs_metadata_sync:
             self.update_metadata(hidden_states.shape[0])
 
+        # Mask out CUDA-graph padding rows of the local routing map so the AGV
+        # propagates -1 into agv_r for those slots; padding tokens then route
+        # to no expert. _real_token_count_tensor is wired by the context and
+        # holds the *global* unpadded token count, so we pass self.sp_rank to
+        # shift local rows into the global frame for the comparison. When unset
+        # (standalone dispatcher use without a context) all rows are real, so
+        # skip the mask.
+        if self.__class__._real_token_count_tensor is not None:
+            mask_routing_padding(
+                self.routing_map, self.__class__._real_token_count_tensor, self.sp_rank
+            )
+
         agv_h = self.__class__._symm_agv_hidden
         agv_r = self.__class__._symm_agv_routing
         agv_p = self.__class__._symm_agv_probs