From 02a8f520303df858c0430c49c275f9d68710ec12 Mon Sep 17 00:00:00 2001 From: Siddharth Singh Date: Fri, 12 Jun 2026 11:08:34 -0700 Subject: [PATCH 1/6] no route dummy/pad tokens to experts --- .../inference/contexts/dynamic_context.py | 65 +++++++++-- megatron/core/inference/contexts/gpu_view.py | 10 ++ megatron/core/transformer/attention.py | 1 + .../moe/inference_routing_mask_kernel.py | 106 ++++++++++++++++++ .../moe/token_dispatcher_inference.py | 31 +++++ 5 files changed, 204 insertions(+), 9 deletions(-) create mode 100644 megatron/core/transformer/moe/inference_routing_mask_kernel.py diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index 90add4c0632..13d0acdf6d2 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -629,6 +629,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 = ( @@ -673,12 +679,17 @@ 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, @@ -686,10 +697,6 @@ def __init__(self, model_config: TransformerConfig, inference_config: InferenceC 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 @@ -708,8 +715,19 @@ 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 @@ -1015,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 @@ -1026,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 @@ -1152,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( @@ -2375,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 @@ -2451,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 diff --git a/megatron/core/inference/contexts/gpu_view.py b/megatron/core/inference/contexts/gpu_view.py index 651d95055da..2066375d19e 100644 --- a/megatron/core/inference/contexts/gpu_view.py +++ b/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) diff --git a/megatron/core/transformer/attention.py b/megatron/core/transformer/attention.py index 8fad62c60c5..840c7e73baa 100644 --- a/megatron/core/transformer/attention.py +++ b/megatron/core/transformer/attention.py @@ -89,6 +89,7 @@ except ImportError: HAVE_FA4 = False + try: from flash_mla import flash_mla_with_kvcache, get_mla_metadata diff --git a/megatron/core/transformer/moe/inference_routing_mask_kernel.py b/megatron/core/transformer/moe/inference_routing_mask_kernel.py new file mode 100644 index 00000000000..0bda004ceb8 --- /dev/null +++ b/megatron/core/transformer/moe/inference_routing_mask_kernel.py @@ -0,0 +1,106 @@ +# 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.autotune( + configs=[ + triton.Config({"BLOCK_M": 8}), + triton.Config({"BLOCK_M": 128}) + ], + key=["total_rows", "TOPK"], +) +@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 (autotuned) + 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_TOPK = triton.next_power_of_2(topk) + grid = lambda META: (triton.cdiv(total_rows, META["BLOCK_M"]),) # noqa: E731 + + _mask_routing_padding_kernel[grid]( + routing_map, + real_token_count_tensor, + total_rows=total_rows, + tp_rank=tp_rank, + TOPK=topk, + BLOCK_TOPK=BLOCK_TOPK, + ) diff --git a/megatron/core/transformer/moe/token_dispatcher_inference.py b/megatron/core/transformer/moe/token_dispatcher_inference.py index 081497f734c..47802265d7f 100644 --- a/megatron/core/transformer/moe/token_dispatcher_inference.py +++ b/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,16 @@ 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 _rank_token_offset(cls) -> torch.Tensor: return cls._step_metadata[1:2] @@ -343,6 +360,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 +484,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 +546,15 @@ 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. + 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 From ed2d4176a91808b037e78c14f914c8f180dc4935 Mon Sep 17 00:00:00 2001 From: Siddharth Singh Date: Wed, 1 Jul 2026 11:12:20 -0700 Subject: [PATCH 2/6] Address review: fix MTP token count masking, guard None dispatcher, add kernel tests Signed-off-by: Siddharth Singh --- .../text_generation_controller.py | 9 ++ megatron/core/transformer/attention.py | 1 - .../moe/token_dispatcher_inference.py | 26 +++++- .../test_moe_dispatching_and_routing.py | 82 +++++++++++++++++++ 4 files changed, 113 insertions(+), 5 deletions(-) diff --git a/megatron/core/inference/text_generation_controllers/text_generation_controller.py b/megatron/core/inference/text_generation_controllers/text_generation_controller.py index 6b75c4685ac..664a7d61bbf 100644 --- a/megatron/core/inference/text_generation_controllers/text_generation_controller.py +++ b/megatron/core/inference/text_generation_controllers/text_generation_controller.py @@ -39,6 +39,7 @@ from megatron.core.transformer.enums import InferenceCudaGraphScope from megatron.core.transformer.moe.moe_layer import BaseMoELayer from megatron.core.transformer.moe.router_replay import RouterReplay, RouterReplayAction +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, @@ -836,6 +837,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}") diff --git a/megatron/core/transformer/attention.py b/megatron/core/transformer/attention.py index 840c7e73baa..8fad62c60c5 100644 --- a/megatron/core/transformer/attention.py +++ b/megatron/core/transformer/attention.py @@ -89,7 +89,6 @@ except ImportError: HAVE_FA4 = False - try: from flash_mla import flash_mla_with_kvcache, get_mla_metadata diff --git a/megatron/core/transformer/moe/token_dispatcher_inference.py b/megatron/core/transformer/moe/token_dispatcher_inference.py index 47802265d7f..b08d88f2641 100644 --- a/megatron/core/transformer/moe/token_dispatcher_inference.py +++ b/megatron/core/transformer/moe/token_dispatcher_inference.py @@ -343,6 +343,21 @@ def set_real_token_count_tensor(cls, tensor: torch.Tensor) -> None: """ 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] @@ -550,10 +565,13 @@ def token_dispatch(self, hidden_states, probs): # 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. - mask_routing_padding( - self.routing_map, self.__class__._real_token_count_tensor, self.sp_rank - ) + # 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 diff --git a/tests/unit_tests/inference/test_moe_dispatching_and_routing.py b/tests/unit_tests/inference/test_moe_dispatching_and_routing.py index 49b5df613f7..5b21ab4c364 100644 --- a/tests/unit_tests/inference/test_moe_dispatching_and_routing.py +++ b/tests/unit_tests/inference/test_moe_dispatching_and_routing.py @@ -439,3 +439,85 @@ def test_cuda_graph_dispatch_combine(self, max_rank_tokens, seed): expected_combined = (global_hidden[start:end].float() * ep_size).bfloat16() torch.testing.assert_close(graph_combined, expected_combined, atol=0, rtol=0) + + +# ────────────────────────────────────────────────────────────────────── +# mask_routing_padding kernel +# ────────────────────────────────────────────────────────────────────── + +from megatron.core.transformer.moe.inference_routing_mask_kernel import ( # noqa: E402 + HAVE_TRITON, + mask_routing_padding, +) + +requires_triton_cuda = pytest.mark.skipif( + not HAVE_TRITON or not torch.cuda.is_available(), + reason="mask_routing_padding requires triton and CUDA", +) + + +@pytest.mark.internal +@requires_triton_cuda +class TestMaskRoutingPadding: + """Unit tests for the CUDA-graph padding-row routing mask. + + ``mask_routing_padding`` fills ``routing_map[real_token_count:, :]`` with -1 so + the NVLS dispatcher routes padding rows to no expert. ``real_token_count`` is in + the global (pre-SP-shard) frame; a non-zero ``tp_rank`` shifts local rows into + that frame before the comparison. Runs standalone — no context or NVLS hardware. + """ + + TOPK = 6 + + def _routing_map(self, n_rows, fill=3): + # All entries non-negative so masked (-1) slots are unambiguous. + return torch.full((n_rows, self.TOPK), fill, dtype=torch.int64, device="cuda") + + def _real_token_count(self, count): + return torch.tensor([count], dtype=torch.int32, device="cuda") + + @pytest.mark.parametrize("n_rows, real_count", [(16, 10), (128, 1), (7, 7), (64, 0)]) + def test_masks_rows_past_real_count(self, n_rows, real_count): + """Rows >= real_count become -1; rows < real_count are untouched (tp_rank=0).""" + routing_map = self._routing_map(n_rows) + original = routing_map.clone() + + mask_routing_padding(routing_map, self._real_token_count(real_count), tp_rank=0) + + torch.testing.assert_close(routing_map[:real_count], original[:real_count]) + assert torch.all(routing_map[real_count:] == -1) + + def test_real_count_equal_rows_is_noop(self): + """real_count == n_rows masks nothing (the unpadded decode case).""" + routing_map = self._routing_map(32) + original = routing_map.clone() + + mask_routing_padding(routing_map, self._real_token_count(32), tp_rank=0) + + torch.testing.assert_close(routing_map, original) + + def test_sp_rank_offset(self): + """Local rows are shifted by tp_rank * n_rows into the global frame. + + With 8 local rows on SP rank 1, local row r is global row r + 8. A global + real_count of 11 keeps global rows [8, 11) real (local [0, 3)) and masks + global rows [11, 16) (local [3, 8)). + """ + routing_map = self._routing_map(8) + original = routing_map.clone() + + mask_routing_padding(routing_map, self._real_token_count(11), tp_rank=1) + + torch.testing.assert_close(routing_map[:3], original[:3]) + assert torch.all(routing_map[3:] == -1) + + def test_sp_rank_fully_masked(self): + """An SP rank entirely beyond real_count is fully masked. + + 8 local rows on rank 1 cover global rows [8, 16); real_count=8 masks all. + """ + routing_map = self._routing_map(8) + + mask_routing_padding(routing_map, self._real_token_count(8), tp_rank=1) + + assert torch.all(routing_map == -1) From 1288d0bd5116040ae75fdf2ef7e6656ab95b985c Mon Sep 17 00:00:00 2001 From: Siddharth Singh Date: Wed, 1 Jul 2026 11:27:12 -0700 Subject: [PATCH 3/6] Replace autotune with BLOCK_M heuristic to avoid recompilation Signed-off-by: Siddharth Singh --- .../moe/inference_routing_mask_kernel.py | 13 ++++--------- 1 file changed, 4 insertions(+), 9 deletions(-) diff --git a/megatron/core/transformer/moe/inference_routing_mask_kernel.py b/megatron/core/transformer/moe/inference_routing_mask_kernel.py index 0bda004ceb8..ef38fbaf8db 100644 --- a/megatron/core/transformer/moe/inference_routing_mask_kernel.py +++ b/megatron/core/transformer/moe/inference_routing_mask_kernel.py @@ -32,13 +32,6 @@ HAVE_TRITON = False -@triton.autotune( - configs=[ - triton.Config({"BLOCK_M": 8}), - triton.Config({"BLOCK_M": 128}) - ], - key=["total_rows", "TOPK"], -) @triton.jit def _mask_routing_padding_kernel( routing_map_ptr, # int64* [total_rows, topk] @@ -46,7 +39,7 @@ def _mask_routing_padding_kernel( 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 (autotuned) + 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.""" @@ -93,8 +86,9 @@ def mask_routing_padding( if total_rows == 0: return + BLOCK_M = 8 if total_rows < 64 else 128 BLOCK_TOPK = triton.next_power_of_2(topk) - grid = lambda META: (triton.cdiv(total_rows, META["BLOCK_M"]),) # noqa: E731 + grid = (triton.cdiv(total_rows, BLOCK_M),) _mask_routing_padding_kernel[grid]( routing_map, @@ -102,5 +96,6 @@ def mask_routing_padding( total_rows=total_rows, tp_rank=tp_rank, TOPK=topk, + BLOCK_M=BLOCK_M, BLOCK_TOPK=BLOCK_TOPK, ) From 7ae32ed58c458f2428b89816b74f7a918deb4c13 Mon Sep 17 00:00:00 2001 From: Siddharth Singh Date: Wed, 1 Jul 2026 11:44:29 -0700 Subject: [PATCH 4/6] format --- megatron/core/inference/contexts/dynamic_context.py | 4 +--- .../transformer/moe/inference_routing_mask_kernel.py | 12 ++++++------ 2 files changed, 7 insertions(+), 9 deletions(-) diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index 13d0acdf6d2..0473e394c2c 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -724,9 +724,7 @@ def __init__(self, model_config: TransformerConfig, inference_config: InferenceC # 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 - ) + NVLSAllGatherVDispatcher.set_real_token_count_tensor(self.gpu_view.real_token_count) # Print info. active_blocks = self.kv_block_allocator.active_count diff --git a/megatron/core/transformer/moe/inference_routing_mask_kernel.py b/megatron/core/transformer/moe/inference_routing_mask_kernel.py index ef38fbaf8db..e38869a1f6d 100644 --- a/megatron/core/transformer/moe/inference_routing_mask_kernel.py +++ b/megatron/core/transformer/moe/inference_routing_mask_kernel.py @@ -34,13 +34,13 @@ @triton.jit def _mask_routing_padding_kernel( - routing_map_ptr, # int64* [total_rows, topk] - real_token_count_ptr, # int32* [1] + 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 + 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) From 7a24757cba6088d559e3a09b4a51bd8fe320e90f Mon Sep 17 00:00:00 2001 From: Siddharth Singh Date: Mon, 13 Jul 2026 13:04:11 -0700 Subject: [PATCH 5/6] lint --- .../text_generation_controllers/text_generation_controller.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/megatron/core/inference/text_generation_controllers/text_generation_controller.py b/megatron/core/inference/text_generation_controllers/text_generation_controller.py index e4314d429b2..a9c5a6a92f7 100644 --- a/megatron/core/inference/text_generation_controllers/text_generation_controller.py +++ b/megatron/core/inference/text_generation_controllers/text_generation_controller.py @@ -41,8 +41,8 @@ from megatron.core.transformer.enums import InferenceCudaGraphScope from megatron.core.transformer.moe.moe_layer import BaseMoELayer from megatron.core.transformer.moe.router_replay import RouterReplay, RouterReplayAction -from megatron.core.transformer.moe.token_dispatcher_inference import NVLSAllGatherVDispatcher 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, From fd6eca29036a684ca442dd93707c2f286c9b36f1 Mon Sep 17 00:00:00 2001 From: Siddharth Singh Date: Mon, 13 Jul 2026 14:41:42 -0700 Subject: [PATCH 6/6] DCO Remediation Commit for Siddharth Singh I, Siddharth Singh , hereby add my Signed-off-by to this commit: 02a8f520303df858c0430c49c275f9d68710ec12 I, Siddharth Singh , hereby add my Signed-off-by to this commit: 7ae32ed58c458f2428b89816b74f7a918deb4c13 Signed-off-by: Siddharth Singh