From 519194249015bc5523ee7b8e85210fc3ee39bffc Mon Sep 17 00:00:00 2001 From: Keshav Santhanam Date: Wed, 8 Apr 2026 12:50:18 -0700 Subject: [PATCH 1/7] Add AllGatherV dispatcher for MoE inference Signed-off-by: Keshav Santhanam --- .../core/inference/batch_dimensions_utils.py | 22 + .../inference/contexts/dynamic_context.py | 8 + megatron/core/inference/moe/fused_moe.py | 53 ++- megatron/core/transformer/moe/experts.py | 43 +- megatron/core/transformer/moe/moe_layer.py | 60 ++- .../moe/token_dispatcher_inference_v.py | 430 ++++++++++++++++++ .../core/transformer/transformer_config.py | 36 ++ 7 files changed, 630 insertions(+), 22 deletions(-) create mode 100644 megatron/core/transformer/moe/token_dispatcher_inference_v.py diff --git a/megatron/core/inference/batch_dimensions_utils.py b/megatron/core/inference/batch_dimensions_utils.py index e27438e63d0..ea74c052126 100644 --- a/megatron/core/inference/batch_dimensions_utils.py +++ b/megatron/core/inference/batch_dimensions_utils.py @@ -142,6 +142,7 @@ def adjust_batch_dims_for_expert_parallelism( decode_only_cuda_graphs: bool, smallest_non_decode_cuda_graph_size: int, ep_group: Optional[torch.distributed.ProcessGroup] = None, + skip_ep_sync: bool = False, ) -> Optional["InferenceBatchDimensions"]: """Adjusted cuda graph batch dimensions for expert parallelism. We take the max token count across expert model parallel group. @@ -153,6 +154,13 @@ def adjust_batch_dims_for_expert_parallelism( ep_group: Optional expert parallel process group. If None, uses global parallel state. When using different EP sizes for inference vs training, pass the inference EP group explicitly. + skip_ep_sync: If True, skip the all-reduce across EP ranks entirely and + return local_batch_dims unchanged. Use this when the MoE dispatcher + pins AllGather/ReduceScatter to a fixed max buffer size (e.g. + ``inference_moe_cuda_graph_dispatcher='allgather_v'`` with + ``inference_moe_max_tokens`` set), so that NCCL collectives are + identical across all CUDA graphs and EP ranks can independently + select different graphs. Return: (InferenceBatchDimensions) A new InferenceBatchDimensions object with @@ -161,6 +169,14 @@ def adjust_batch_dims_for_expert_parallelism( ep_size = get_pg_size(ep_group) if ep_size <= 1: return local_batch_dims + + # When the dispatcher uses fixed-max-size AllGather/ReduceScatter buffers, + # every CUDA graph embeds the same NCCL collective regardless of actual + # batch size. EP ranks can independently match their own graph — no + # cross-rank synchronization is needed. + if skip_ep_sync: + return local_batch_dims + # all reduce local work across expert model parallel group is_non_decode = local_batch_dims.prefill_req_count > 0 @@ -500,6 +516,7 @@ def match_graph_config( strict: bool = False, decode_only_cuda_graphs: bool = False, ep_group: Optional[torch.distributed.ProcessGroup] = None, + skip_ep_sync: bool = False, ) -> Optional[InferenceBatchDimensions]: """ Matches the best CUDA graph batch dimension for the given real batch dimension. @@ -515,6 +532,10 @@ def match_graph_config( ep_group: Optional expert parallel process group. If None, uses global parallel state. When using different EP sizes for inference vs training, pass the inference EP group explicitly. + skip_ep_sync: If True, skip the all-reduce across EP ranks in + adjust_batch_dims_for_expert_parallelism. Use when the MoE dispatcher + uses fixed-max-size AllGather/ReduceScatter buffers so that EP ranks + can independently select different CUDA graphs. Returns: The best matching CUDA graph batch dimension, or None if no applicable match is found """ @@ -529,6 +550,7 @@ def match_graph_config( decode_only_cuda_graphs=decode_only_cuda_graphs, ep_group=ep_group, smallest_non_decode_cuda_graph_size=smallest_non_decode_cuda_graph_size, + skip_ep_sync=skip_ep_sync, ) if adjusted_batch_dim is None: diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index 383d0ecbd95..f5157cb2d76 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -579,6 +579,13 @@ def __init__(self, model_config: TransformerConfig, inference_config: InferenceC ), "Router recording/replay requested but no MoE experts specified!" self.moe_routing_metadata = RoutingMetadata(self, model_config.moe_router_topk) + # When the allgather_v dispatcher is active with fixed-max-buffer mode, + # EP ranks can independently select CUDA graphs — no cross-rank sync needed. + self._skip_ep_sync = ( + getattr(model_config, "inference_moe_cuda_graph_dispatcher", "fused") == "allgather_v" + and getattr(model_config, "inference_moe_max_tokens", None) is not None + ) + # CUDA graph config list self.use_cuda_graphs_for_non_decode_steps = ( inference_config.use_cuda_graphs_for_non_decode_steps @@ -1608,6 +1615,7 @@ def initialize_attention_state( strict=self.is_hybrid_model, decode_only_cuda_graphs=(not self.use_cuda_graphs_for_non_decode_steps), ep_group=self.expert_model_parallel_group, + skip_ep_sync=self._skip_ep_sync, ) self._using_cuda_graph_this_step = best_graph is not None diff --git a/megatron/core/inference/moe/fused_moe.py b/megatron/core/inference/moe/fused_moe.py index 39382eee079..e83bf7856b3 100644 --- a/megatron/core/inference/moe/fused_moe.py +++ b/megatron/core/inference/moe/fused_moe.py @@ -90,15 +90,22 @@ def mcore_fused_moe( tokens_per_expert: Optional[torch.Tensor] = None, skip_permute: bool = False, disable_fused_quant_kernels: bool = False, + expert_offsets: Optional[torch.Tensor] = None, + permutation_map: Optional[torch.Tensor] = None, ) -> torch.Tensor: """Fused MoE: [permute ->] pad -> FC1 -> activation -> FC2 -> unpad [-> unpermute]. - Two modes: + Three modes: - skip_permute=False (default): tokens are unpermuted. Requires routing_map. Performs full permute -> compute -> unpermute. - skip_permute=True: tokens are already permuted by the dispatcher. Requires tokens_per_expert. Pads to alignment, computes, then unpads. Probs are applied during unpad. + - expert_offsets is not None: tokens are already permuted AND padded by the + dispatcher (InferenceAllGatherVTokenDispatcher). Requires expert_offsets + (inclusive prefix sums for grouped_mm) and permutation_map (for padding-aware + activation). Performs only FC1 -> activation -> FC2. The caller handles + unpermute. This mode is fully CUDA-graphable with no host-device sync. Unless disable_fused_quant_kernels=True, when weights are MXFP8, uses fused kernels that combine permute/activation with MXFP8 quantization into single @@ -108,7 +115,7 @@ def mcore_fused_moe( hidden_states: [num_tokens, hidden_size] BF16 input. probs: routing probabilities. Shape is [num_tokens, topk] when skip_permute=False, or [num_tokens] (already gathered) when - skip_permute=True. + skip_permute=True. Ignored when expert_offsets is provided. fc1_weight: stacked weight for FC1 (torch.Tensor for BF16, MXFP8Tensor for MXFP8). fc2_weight: stacked weight for FC2 (same type as fc1_weight). activation_type: ActivationType enum (SQUARED_RELU). @@ -120,9 +127,18 @@ def mcore_fused_moe( disable_fused_quant_kernels: if True, disable fused permute+quantize and activation+quantize kernels for MXFP8, using separate launches instead. Useful for debugging. Ignored when weights are BF16. + expert_offsets: [num_local_experts] int32 inclusive prefix sums of aligned + token counts, as produced by Triton compute_expert_offsets. When provided, + hidden_states must already be permuted and padded. FC1 -> activation -> FC2 + is performed using these offsets and the raw output is returned (no unpermute). + permutation_map: [output_size] int32, original token index or -1 for padding. + Required when expert_offsets is provided. Used by the activation kernel to + skip padding rows. Returns: - [num_tokens, hidden_size] BF16 output. + [num_tokens, hidden_size] BF16 output (when expert_offsets is None), or + raw [output_size, hidden_size] expert output (when expert_offsets is provided, + caller handles unpermute). """ assert ( hidden_states.dtype == torch.bfloat16 @@ -133,6 +149,37 @@ def mcore_fused_moe( # Fused quant kernels only apply to MXFP8 path use_fused_quant = use_mxfp8 and not disable_fused_quant_kernels + # --- Pre-permuted + pre-padded path (AllGatherV dispatcher) --- + # Tokens are already in expert-grouped order with alignment padding. + # Just do FC1 -> activation -> FC2 using the provided offsets. + if expert_offsets is not None: + assert permutation_map is not None, ( + "permutation_map is required when expert_offsets is provided" + ) + offs = expert_offsets + if use_mxfp8: + assert ( + HAVE_SCALED_GMM + ), "torch.nn.functional.scaled_grouped_mm not available. Install PyTorch 2.10+." + mm_fn_local = _mxfp8_grouped_mm + activation_fn = _get_activation_func(activation_type, fused_quant=use_fused_quant) + # Quantize input for MXFP8 path + hidden_states = MXFP8Tensor.from_bf16(hidden_states, backend="triton") + else: + assert ( + HAVE_GROUPED_MM + ), "torch.nn.functional.grouped_mm not available. Install PyTorch 2.10+." + mm_fn_local = _bf16_grouped_mm + activation_fn = _get_activation_func(activation_type, fused_quant=False) + + fc1_output = mm_fn_local(hidden_states, fc1_weight, offs) + activation_out = activation_fn(fc1_output, permutation_map) + if use_mxfp8 and not isinstance(activation_out, MXFP8Tensor): + activation_out = MXFP8Tensor.from_bf16(activation_out, backend="triton") + fc2_output = mm_fn_local(activation_out, fc2_weight, offs) + # Return raw output — the dispatcher handles unpermute + prob weighting. + return fc2_output + if use_mxfp8: assert ( HAVE_SCALED_GMM diff --git a/megatron/core/transformer/moe/experts.py b/megatron/core/transformer/moe/experts.py index 34e9fb17a02..db79fffd77c 100644 --- a/megatron/core/transformer/moe/experts.py +++ b/megatron/core/transformer/moe/experts.py @@ -638,7 +638,14 @@ def _flashinfer_forward(self, hidden_states, routing_map, probs): return output, None def _mcore_fused_moe_forward( - self, hidden_states, probs, routing_map=None, tokens_per_expert=None, skip_permute=False + self, + hidden_states, + probs, + routing_map=None, + tokens_per_expert=None, + skip_permute=False, + expert_offsets=None, + permutation_map=None, ): """Torch grouped_mm fused MoE forward via mcore_fused_moe.""" local_expert_start = self.ep_group.rank() * self.num_local_experts @@ -654,6 +661,8 @@ def _mcore_fused_moe_forward( tokens_per_expert=tokens_per_expert, skip_permute=skip_permute, disable_fused_quant_kernels=self.config.inference_moe_disable_fused_quant_kernels, + expert_offsets=expert_offsets, + permutation_map=permutation_map, ) return output, None @@ -663,22 +672,32 @@ def forward( tokens_per_expert: Optional[torch.Tensor], permuted_probs: torch.Tensor, routing_map: Optional[torch.Tensor] = None, + expert_offsets: Optional[torch.Tensor] = None, + permutation_map: Optional[torch.Tensor] = None, ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: - """Forward pass with three modes: + """Forward pass with four modes: - Training: delegates to parent TEGroupedMLP. - - Inference + CUDA graphed: FlashInfer cutlass_fused_moe. tokens_per_expert - is not used in this path; the FlashInfer kernel operates directly on - routing_map. - - Inference + eager: torch.nn.functional.grouped_mm with GPU-resident cumsum offsets. + - Inference + CUDA graphed + AllGatherV dispatcher: tokens are already + permuted and padded by Triton kernels. Uses grouped_mm with the + provided expert_offsets. No FlashInfer dependency. + - Inference + CUDA graphed + fused dispatcher: FlashInfer cutlass_fused_moe. + tokens_per_expert is not used; the kernel operates directly on routing_map. + - Inference + eager: torch.nn.functional.grouped_mm with GPU-resident + cumsum offsets. Args: permuted_local_hidden_states: [num_tokens, hidden_size] input hidden states. tokens_per_expert: [num_experts] number of tokens routed to each expert. - None when using the CUDA-graphed FlashInfer path. + None when using CUDA-graphed paths. permuted_probs: [num_tokens, topk] routing probabilities. routing_map: [num_tokens, topk] token-to-expert assignment indices. Required for the FlashInfer CUDA-graphed path, None otherwise. + expert_offsets: [num_local_experts] int32 inclusive prefix sums of aligned + token counts. When provided, hidden_states are already permuted and + padded by the AllGatherV dispatcher's Triton kernels. + permutation_map: [output_size] int32, original token index or -1 for + padding. Required when expert_offsets is provided. """ if self.training: @@ -698,6 +717,16 @@ def forward( self._build_concatenated_weights() self._concatenated_weights_built = True + # Pre-permuted path: tokens already grouped by expert with padding. + # Directly use grouped_mm with the provided offsets — no FlashInfer needed. + if expert_offsets is not None: + return self._mcore_fused_moe_forward( + permuted_local_hidden_states, + permuted_probs, + expert_offsets=expert_offsets, + permutation_map=permutation_map, + ) + resolved_backend = resolve_inference_grouped_gemm_backend( self.inference_grouped_gemm_backend, self.is_inference_cuda_graphed_iteration, diff --git a/megatron/core/transformer/moe/moe_layer.py b/megatron/core/transformer/moe/moe_layer.py index 35b567679fe..98248a1466f 100644 --- a/megatron/core/transformer/moe/moe_layer.py +++ b/megatron/core/transformer/moe/moe_layer.py @@ -28,6 +28,9 @@ from megatron.core.transformer.moe.token_dispatcher_inference import ( InferenceCUDAGraphTokenDispatcher, ) +from megatron.core.transformer.moe.token_dispatcher_inference_v import ( + InferenceAllGatherVTokenDispatcher, +) from megatron.core.transformer.transformer_config import TransformerConfig from megatron.core.typed_torch import apply_module, not_none from megatron.core.utils import internal_api @@ -352,8 +355,16 @@ def _setup_inference_mode(self, pg_collection): """Set up inference-optimized token dispatcher and state. Called from __init__ when config.transformer_impl == "inference_optimized". - Creates an InferenceCUDAGraphTokenDispatcher alongside the standard dispatcher, - which is swapped in during CUDA-graphed forward passes. + Creates an inference CUDA-graph token dispatcher alongside the standard + dispatcher, which is swapped in during CUDA-graphed forward passes. + + The dispatcher type is selected by ``inference_moe_cuda_graph_dispatcher``: + - ``"fused"`` (default): AllGather + FlashInfer/CUTLASS fused MoE kernel. + Requires FlashInfer. The fused kernel handles permutation internally. + - ``"allgather_v"``: AllGather + Triton permute + grouped_mm + Triton + unpermute + ReduceScatter. No FlashInfer dependency. Token counts + are computed on-device by Triton kernels — fully CUDA-graphable + with no host-device synchronization. """ assert self.config.moe_token_dispatcher_type == "alltoall", ( @@ -361,12 +372,22 @@ def _setup_inference_mode(self, pg_collection): f"got '{self.config.moe_token_dispatcher_type}'" ) self.is_inference_cuda_graphed_iteration = False - self._inference_token_dispatcher = InferenceCUDAGraphTokenDispatcher( - self.num_local_experts, - self.local_expert_indices, - config=self.config, - pg_collection=pg_collection, - ) + + dispatcher_type = self.config.inference_moe_cuda_graph_dispatcher + if dispatcher_type == "allgather_v": + self._inference_token_dispatcher = InferenceAllGatherVTokenDispatcher( + self.num_local_experts, + self.local_expert_indices, + config=self.config, + pg_collection=pg_collection, + ) + else: + self._inference_token_dispatcher = InferenceCUDAGraphTokenDispatcher( + self.num_local_experts, + self.local_expert_indices, + config=self.config, + pg_collection=pg_collection, + ) def setup_delayed_wgrad_for_dispatch_backward_overlap(self): """Initializes CUDA events and streams for overlapping expert @@ -499,10 +520,25 @@ def routed_experts_compute(self, hidden_states: torch.Tensor, probs: torch.Tenso hasattr(self, "_inference_token_dispatcher") and self.is_inference_cuda_graphed_iteration ): - routing_map = self.token_dispatcher.routing_map - expert_output, mlp_bias = apply_module(self.experts)( - dispatched_input, tokens_per_expert, permuted_probs, routing_map=routing_map - ) + # AllGatherV dispatcher: tokens are pre-permuted with Triton — + # pass expert_offsets and permutation_map for grouped_mm. + if isinstance(self.token_dispatcher, InferenceAllGatherVTokenDispatcher): + expert_output, mlp_bias = apply_module(self.experts)( + dispatched_input, + tokens_per_expert, + permuted_probs, + expert_offsets=self.token_dispatcher.expert_offsets, + permutation_map=self.token_dispatcher.permutation_map, + ) + else: + # Fused dispatcher: pass routing_map for FlashInfer. + routing_map = self.token_dispatcher.routing_map + expert_output, mlp_bias = apply_module(self.experts)( + dispatched_input, + tokens_per_expert, + permuted_probs, + routing_map=routing_map, + ) else: expert_output, mlp_bias = apply_module(self.experts)( dispatched_input, tokens_per_expert, permuted_probs diff --git a/megatron/core/transformer/moe/token_dispatcher_inference_v.py b/megatron/core/transformer/moe/token_dispatcher_inference_v.py new file mode 100644 index 00000000000..0dd8a481f33 --- /dev/null +++ b/megatron/core/transformer/moe/token_dispatcher_inference_v.py @@ -0,0 +1,430 @@ +# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +""" +CUDA-graph-compatible AllGather-V token dispatcher for inference. + +This dispatcher replaces the FlashInfer-based InferenceCUDAGraphTokenDispatcher +with a Triton-based permutation pipeline that works with standard +torch.nn.functional.grouped_mm. It avoids any host-device synchronization: + + 1. AllGather routing_map, probs, and hidden_states across EP ranks. + 2. Triton kernels compute per-expert token counts and aligned prefix-sum + offsets entirely on-device (no cudaMemcpy D->H for token counts). + 3. Triton permute kernel groups tokens by local expert with alignment padding. + 4. Expert compute uses grouped_mm with GPU-resident offsets. + 5. Triton unpermute kernel scatters weighted expert outputs back. + 6. ReduceScatter combines contributions across EP ranks. + +When ``inference_moe_max_tokens`` is set, AllGather/ReduceScatter buffers are +pinned to a fixed maximum size. This makes the NCCL collectives identical +across every CUDA graph, so different EP ranks can independently select +different graphs without any cross-rank synchronization -- the all-reduce in +``adjust_batch_dims_for_expert_parallelism`` is no longer required. + +The "V" (variable) refers to the fact that each expert receives a data-dependent +number of tokens, computed on-device without cross-rank synchronization. +""" + +from typing import List, Optional + +import torch + +from megatron.core.inference.communication.torch_symm_triton import ( + are_tensors_nvls_eligible, + multimem_all_gather_fused, + multimem_reduce_scatter, +) +from megatron.core.inference.moe.permute import permute_tokens, unpermute_tokens +from megatron.core.inference.symmetric_memory import SymmetricMemoryManager +from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.tensor_parallel import ( + gather_from_sequence_parallel_region, + reduce_scatter_to_sequence_parallel_region, +) +from megatron.core.transformer.moe.token_dispatcher import MoEAllGatherTokenDispatcher +from megatron.core.transformer.transformer_config import TransformerConfig + + +class InferenceAllGatherVTokenDispatcher(MoEAllGatherTokenDispatcher): + """CUDA-graph-compatible AllGather-V token dispatcher with Triton permutation. + + Unlike the FlashInfer-based InferenceCUDAGraphTokenDispatcher, this dispatcher + explicitly permutes tokens into expert-grouped order using Triton kernels, + making it compatible with torch.nn.functional.grouped_mm / scaled_grouped_mm. + + Token counts per expert are computed on-device by Triton kernels -- no + host-device synchronization is needed, so the full dispatch/combine pipeline + is CUDA-graphable. + + When ``config.inference_moe_max_tokens`` is set and EP > 1, all + AllGather/ReduceScatter buffers are pinned to a fixed maximum size so that + the embedded NCCL collectives are identical across every CUDA graph. This + allows different EP ranks to capture and replay **different** graphs + independently -- no cross-rank batch-dimension synchronization is required. + + Key features: + - AllGather/ReduceScatter for EP communication (CUDA-graph safe) + - NVLS collectives on Hopper+ with automatic NCCL fallback + - Triton-based permute/unpermute (no FlashInfer dependency) + - GPU-resident token counts and expert offsets (no D->H sync) + - Fixed-max-buffer mode eliminates EP rank synchronization + """ + + def __init__( + self, + num_local_experts: int, + local_expert_indices: List[int], + config: TransformerConfig, + pg_collection: Optional[ProcessGroupCollection] = None, + ) -> None: + super().__init__( + num_local_experts=num_local_experts, + local_expert_indices=local_expert_indices, + config=config, + pg_collection=pg_collection, + ) + self.topk = config.moe_router_topk + self.local_expert_start = local_expert_indices[0] + self.triton_nvls_kernels_allowed = not self.config.inference_disable_triton_nvls_kernels + + # Alignment for grouped_mm / scaled_grouped_mm. + # MXFP8 swizzle requires 128; BF16 grouped_mm requires 16. + self._expert_alignment = 128 if config.fp8_recipe == "mxfp8" else 16 + + # Fixed-max-buffer mode: pin AllGather/ReduceScatter to this per-rank + # token count so every CUDA graph embeds the same collective. + self._max_tokens_per_rank: Optional[int] = config.inference_moe_max_tokens + if self.ep_size > 1 and self._max_tokens_per_rank is None: + raise ValueError( + "inference_moe_max_tokens must be set when using the 'allgather_v' " + "dispatcher with EP > 1. Set it to cuda_graph_max_tokens " + "(typically max_requests * (num_speculative_tokens + 1))." + ) + + # Cached between dispatch_postprocess and combine_preprocess. + self.expert_offsets: Optional[torch.Tensor] = None + self.permutation_map: Optional[torch.Tensor] = None + self._permuted_probs: Optional[torch.Tensor] = None + self._num_global_tokens: int = 0 + self._actual_local_tokens: int = 0 + + # ------------------------------------------------------------------ + # Padding helpers for fixed-max-buffer mode + # ------------------------------------------------------------------ + + def _pad_to_max( + self, + hidden_states: torch.Tensor, + routing_map: torch.Tensor, + probs: torch.Tensor, + ): + """Embed actual tokens into fixed-max-size buffers. + + Creates tensors of size [max_tokens_per_rank, ...] and copies actual + data into the leading rows. Padding rows get zeros for hidden_states + and probs, and -1 for routing_map (so Triton permute skips them). + + All output sizes are fixed across CUDA graph replays. + """ + max_tokens = self._max_tokens_per_rank + actual = hidden_states.shape[0] + hidden_dim = hidden_states.shape[1] + topk = probs.shape[1] + device = hidden_states.device + + padded_hidden = torch.zeros( + max_tokens, hidden_dim, dtype=hidden_states.dtype, device=device + ) + padded_hidden[:actual] = hidden_states + + padded_routing_map = torch.full( + (max_tokens, topk), -1, dtype=routing_map.dtype, device=device + ) + padded_routing_map[:actual] = routing_map + + padded_probs = torch.zeros( + max_tokens, topk, dtype=probs.dtype, device=device + ) + padded_probs[:actual] = probs + + return padded_hidden, padded_routing_map, padded_probs + + # ------------------------------------------------------------------ + # AllGather helpers (symmetric memory / NVLS) + # ------------------------------------------------------------------ + + def _maybe_allocate_ag_buffers( + self, routing_map: torch.Tensor, probs: torch.Tensor, hidden_states: torch.Tensor + ) -> dict: + """Allocate a single symmetric memory output buffer for fused all-gather. + + Returns sliced views for routing_map, probs, and hidden_states, or all-None + when symmetric memory is unavailable. + """ + _NONE = { + "handle": None, + "routing_map": None, + "routing_map_offset": 0, + "probs": None, + "probs_offset": 0, + "hidden_states": None, + "hidden_states_offset": 0, + } + + local_tokens = probs.size(0) + global_tokens = local_tokens * self.ep_size + topk = probs.size(-1) + hidden_dim = hidden_states.size(-1) + + result = SymmetricMemoryManager.get_buffer( + "ep", process_group=self.ep_group + ).maybe_get_tensors( + [ + (global_tokens * topk, routing_map.dtype), + (global_tokens * topk, probs.dtype), + (global_tokens * hidden_dim, hidden_states.dtype), + ] + ) + + if result["handle"] is None: + return _NONE + + (rm_buf, rm_off), (p_buf, p_off), (hs_buf, hs_off) = result["tensors"] + return { + "handle": result["handle"], + "routing_map": rm_buf, + "routing_map_offset": rm_off, + "probs": p_buf, + "probs_offset": p_off, + "hidden_states": hs_buf, + "hidden_states_offset": hs_off, + } + + def _maybe_allocate_rs_buffer(self, x: torch.Tensor) -> dict: + """Allocate a symmetric memory buffer for reduce-scatter input.""" + return SymmetricMemoryManager.get_buffer( + "ep", process_group=self.ep_group + ).maybe_get_tensor(list(x.size()), dtype=x.dtype) + + # ------------------------------------------------------------------ + # Dispatch: pad -> AllGather -> Triton permute + # ------------------------------------------------------------------ + + def dispatch_preprocess(self, hidden_states, routing_map, probs): + """Cache routing_map, reshape, and optionally pad to fixed max size. + + When ``_max_tokens_per_rank`` is set and EP > 1, embeds the actual + tokens into fixed-max-size buffers so that every CUDA graph produces + the same AllGather buffer size. + + Overrides the base class to insert the padding step. + """ + self.hidden_shape = hidden_states.shape + hidden_states = hidden_states.view(-1, self.hidden_shape[-1]) + self._actual_local_tokens = hidden_states.shape[0] + + if self._max_tokens_per_rank is not None and self.ep_size > 1: + hidden_states, routing_map, probs = self._pad_to_max( + hidden_states, routing_map, probs + ) + + self.routing_map = routing_map + return hidden_states, probs + + def token_dispatch(self, hidden_states, probs): + """Gather tokens from all EP ranks using AllGather. + + After ``dispatch_preprocess`` padding, every rank sends exactly + ``max_tokens_per_rank`` tokens (if configured). The resulting + AllGather is the same size across all CUDA graphs, allowing + different EP ranks to replay different graphs independently. + + Uses fused NVLS multimem_all_gather on Hopper+ GPUs when available, + with NCCL fallback. + + Args: + hidden_states: [tokens_per_rank, hidden_dim] (may be padded). + probs: [tokens_per_rank, topk] (may be padded). + + Returns: + (hidden_states, probs) gathered across all EP ranks. + Also updates self.routing_map in-place to the gathered shape. + """ + if self.ep_size == 1: + return hidden_states, probs + + nvls_eligible = self.triton_nvls_kernels_allowed and are_tensors_nvls_eligible( + hidden_states, probs, self.routing_map + ) + ag_buffers = None + + if nvls_eligible: + ag_buffers = self._maybe_allocate_ag_buffers(self.routing_map, probs, hidden_states) + + can_use_nvls = nvls_eligible and ag_buffers["handle"] is not None + + if can_use_nvls: + local_tokens = probs.size(0) + global_tokens = local_tokens * self.ep_size + topk = probs.size(1) + hidden_dim = hidden_states.size(1) + routing_map_dtype = self.routing_map.dtype + probs_dtype = probs.dtype + hidden_dtype = hidden_states.dtype + + multimem_all_gather_fused( + ag_buffers["routing_map"].view(torch.bfloat16), + self.routing_map.view(torch.bfloat16), + ag_buffers["routing_map_offset"], + ag_buffers["probs"].view(torch.bfloat16), + probs.view(torch.bfloat16), + ag_buffers["probs_offset"], + ag_buffers["hidden_states"].view(torch.bfloat16), + hidden_states.view(torch.bfloat16), + ag_buffers["hidden_states_offset"], + ag_buffers["handle"], + ) + self.routing_map = ( + ag_buffers["routing_map"].view(routing_map_dtype).view(global_tokens, topk) + ) + probs = ag_buffers["probs"].view(probs_dtype).view(global_tokens, topk) + hidden_states = ( + ag_buffers["hidden_states"].view(hidden_dtype).view(global_tokens, hidden_dim) + ) + else: + with torch.no_grad(): + self.routing_map = gather_from_sequence_parallel_region( + self.routing_map, group=self.tp_ep_group + ) + probs = gather_from_sequence_parallel_region(probs, group=self.tp_ep_group) + hidden_states = gather_from_sequence_parallel_region( + hidden_states, group=self.tp_ep_group + ) + + return hidden_states, probs + + def dispatch_postprocess(self, hidden_states, probs): + """Permute gathered tokens into expert-grouped order using Triton. + + Uses on-device Triton kernels to: + 1. Count tokens per local expert (atomic histogram). + 2. Compute aligned prefix-sum offsets for grouped_mm. + 3. Permute tokens + probs into expert-contiguous layout with + alignment padding. + + No host-device synchronization occurs -- all metadata stays GPU-resident. + Padding tokens (routing_map == -1) are automatically skipped by the + Triton permute kernel. + + Args: + hidden_states: [global_tokens, hidden_dim] gathered hidden states. + probs: [global_tokens, topk] gathered routing probabilities. + + Returns: + (permuted_hidden, tokens_per_expert, permuted_probs) + - permuted_hidden: [output_size, hidden_dim] expert-grouped tokens. + - tokens_per_expert: None (offsets are stored on self.expert_offsets + instead -- grouped_mm uses those directly). + - permuted_probs: [output_size] flat routing probabilities matching + the permuted token order. + """ + self._num_global_tokens = hidden_states.shape[0] + + permuted_hidden, permuted_probs, permutation_map, inclusive_offsets = permute_tokens( + hidden_states, + probs, + self.routing_map, + self.local_expert_start, + self.num_local_experts, + alignment=self._expert_alignment, + ) + + # Cache for combine_preprocess (Triton unpermute) and expert compute. + self.expert_offsets = inclusive_offsets + self.permutation_map = permutation_map + self._permuted_probs = permuted_probs + + # tokens_per_expert = None: the expert uses self.expert_offsets directly. + return permuted_hidden, None, permuted_probs + + # ------------------------------------------------------------------ + # Combine: Triton unpermute -> ReduceScatter -> unpad + # ------------------------------------------------------------------ + + def combine_preprocess(self, expert_output): + """Scatter weighted expert outputs back to original token positions. + + Uses the Triton unpermute kernel which performs weighted (by routing + probability) atomic scatter-add in fp32, then casts to bf16 for + the subsequent ReduceScatter. + + Args: + expert_output: [output_size, hidden_dim] raw FC2 output in + expert-grouped order (no probability weighting applied yet). + + Returns: + [global_tokens, hidden_dim] bf16 tensor with each token's output + equal to the sum of its weighted expert contributions on this rank. + """ + output = unpermute_tokens( + expert_output, + self._permuted_probs, + self.permutation_map, + self._num_global_tokens, + ) + return output.to(torch.bfloat16) + + def token_combine(self, hidden_states): + """Reduce-scatter expert outputs back to local token slices. + + Sums contributions across EP ranks (each rank contributes non-zero + values only for tokens routed to its local experts) and scatters + the result so each rank receives its local portion. + + Uses NVLS multimem_reduce_scatter on Hopper+ when available. + + Args: + hidden_states: [global_tokens, hidden_dim] combined expert output. + + Returns: + [tokens_per_rank, hidden_dim] bf16 (may include max-buffer padding). + """ + if self.ep_size == 1: + return hidden_states + + output_shape = list(hidden_states.size()) + output_shape[0] = hidden_states.size(0) // self.ep_size + output = torch.empty(output_shape, dtype=hidden_states.dtype, device=hidden_states.device) + + nvls_eligible = ( + self.triton_nvls_kernels_allowed + and output.dtype in (torch.bfloat16, torch.float32) + and are_tensors_nvls_eligible(output) + ) + rs_buffer = None + + if nvls_eligible: + rs_buffer = self._maybe_allocate_rs_buffer(hidden_states) + + can_use_nvls = nvls_eligible and rs_buffer["handle"] is not None + + if can_use_nvls: + rs_buffer["tensor"].copy_(hidden_states) + multimem_reduce_scatter(output, rs_buffer["tensor"], rs_buffer["handle"]) + return output.to(torch.bfloat16) + else: + hidden_states = reduce_scatter_to_sequence_parallel_region( + hidden_states, group=self.tp_ep_group + ) + return hidden_states.to(torch.bfloat16) + + def combine_postprocess(self, hidden_states): + """Strip max-buffer padding and restore original tensor shape. + + When fixed-max-buffer mode is active, the ReduceScatter output has + ``max_tokens_per_rank`` rows. This method slices back to the actual + local token count, then reshapes to the original ``[S/TP, B, H]``. + """ + if self._max_tokens_per_rank is not None and self.ep_size > 1: + hidden_states = hidden_states[: self._actual_local_tokens] + return hidden_states.view(self.hidden_shape) diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 975f971fbc9..9c27340587e 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -941,6 +941,37 @@ class TransformerConfig(ModelParallelConfig): fp8_recipe='mxfp8'. Set to True to disable fusion and use separate kernel launches (useful for debugging).""" + inference_moe_cuda_graph_dispatcher: Literal['fused', 'allgather_v'] = "fused" + """Selects the CUDA-graph-compatible token dispatcher for inference MoE. + Options: + - 'fused' (default): AllGather + FlashInfer/CUTLASS fused MoE kernel. + The fused kernel handles token permutation internally. Requires FlashInfer. + - 'allgather_v': AllGather + Triton permute + grouped_mm + Triton unpermute + + ReduceScatter. Token counts per expert are computed on-device by Triton + kernels with no host-device synchronization, making the full pipeline + CUDA-graphable without FlashInfer. Supports BF16 and MXFP8 weights via + torch.nn.functional.grouped_mm / scaled_grouped_mm. + + When EP > 1, the AllGather/ReduceScatter buffers are pinned to a fixed + maximum size (inference_moe_max_tokens * ep_size) so that all CUDA graphs + embed the same collective regardless of actual batch size. This allows + different EP ranks to independently select different CUDA graphs without + any cross-rank synchronization (no all-reduce for batch dimension matching). + """ + + inference_moe_max_tokens: Optional[int] = None + """Maximum tokens per EP rank for the 'allgather_v' dispatcher. + + When set, the AllGather/ReduceScatter buffers are pinned to this size + (times ep_size), making the NCCL collectives identical across all CUDA + graphs. This eliminates the need for the all-reduce in + adjust_batch_dims_for_expert_parallelism — each EP rank can independently + select its own CUDA graph. + + Should be set to cuda_graph_max_tokens (typically max_requests * + (num_speculative_tokens + 1)). Required when using 'allgather_v' with EP > 1. + """ + mrope_section: Optional[List[int]] = None """ Multimodal rope section is for channel dimension of temporal, height and width in rope calculation. """ @@ -1207,6 +1238,11 @@ def __post_init__(self): f"got '{self.inference_grouped_gemm_backend}'" ) + assert self.inference_moe_cuda_graph_dispatcher in ('fused', 'allgather_v'), ( + f"inference_moe_cuda_graph_dispatcher must be 'fused' or 'allgather_v', " + f"got '{self.inference_moe_cuda_graph_dispatcher}'" + ) + if self.cuda_graph_impl == "local": if self.inference_grouped_gemm_backend == "te": raise ValueError( From 438446a1da5305034c1a41857905390b6a04447b Mon Sep 17 00:00:00 2001 From: Keshav Santhanam Date: Wed, 8 Apr 2026 13:06:12 -0700 Subject: [PATCH 2/7] Fix CLI args Signed-off-by: Keshav Santhanam --- megatron/core/transformer/moe/moe_layer.py | 21 ++++++++---- .../moe/token_dispatcher_inference_v.py | 12 ++++--- .../core/transformer/transformer_config.py | 34 +++++++++++-------- megatron/training/arguments.py | 18 ++++++++++ 4 files changed, 59 insertions(+), 26 deletions(-) diff --git a/megatron/core/transformer/moe/moe_layer.py b/megatron/core/transformer/moe/moe_layer.py index 98248a1466f..43dcf5e6b09 100644 --- a/megatron/core/transformer/moe/moe_layer.py +++ b/megatron/core/transformer/moe/moe_layer.py @@ -275,22 +275,28 @@ def __init__( is_expert=False, ) - # Initialize token dispatcher - if config.moe_token_dispatcher_type == "allgather": + # Initialize token dispatcher. + # 'allgather_v' is inference-only; for the base (training) dispatcher + # we use alltoall since _setup_inference_mode swaps in the real + # InferenceAllGatherVTokenDispatcher during CUDA-graphed iterations. + base_dispatcher_type = config.moe_token_dispatcher_type + if base_dispatcher_type == "allgather_v": + base_dispatcher_type = "alltoall" + if base_dispatcher_type == "allgather": self.token_dispatcher = MoEAllGatherTokenDispatcher( self.num_local_experts, self.local_expert_indices, config=self.config, pg_collection=pg_collection, ) - elif config.moe_token_dispatcher_type == "alltoall": + elif base_dispatcher_type == "alltoall": self.token_dispatcher = MoEAlltoAllTokenDispatcher( self.num_local_experts, self.local_expert_indices, config=self.config, pg_collection=pg_collection, ) - elif config.moe_token_dispatcher_type == "flex": + elif base_dispatcher_type == "flex": self.token_dispatcher = MoEFlexTokenDispatcher( self.num_local_experts, self.local_expert_indices, @@ -358,7 +364,8 @@ def _setup_inference_mode(self, pg_collection): Creates an inference CUDA-graph token dispatcher alongside the standard dispatcher, which is swapped in during CUDA-graphed forward passes. - The dispatcher type is selected by ``inference_moe_cuda_graph_dispatcher``: + The dispatcher type is selected by ``inference_moe_cuda_graph_dispatcher`` + (auto-derived from ``moe_token_dispatcher_type`` in TransformerConfig): - ``"fused"`` (default): AllGather + FlashInfer/CUTLASS fused MoE kernel. Requires FlashInfer. The fused kernel handles permutation internally. - ``"allgather_v"``: AllGather + Triton permute + grouped_mm + Triton @@ -367,8 +374,8 @@ def _setup_inference_mode(self, pg_collection): with no host-device synchronization. """ - assert self.config.moe_token_dispatcher_type == "alltoall", ( - f"Inference-optimized MoE requires 'alltoall' dispatcher, " + assert self.config.moe_token_dispatcher_type in ("alltoall", "allgather_v"), ( + f"Inference-optimized MoE requires 'alltoall' or 'allgather_v' dispatcher, " f"got '{self.config.moe_token_dispatcher_type}'" ) self.is_inference_cuda_graphed_iteration = False diff --git a/megatron/core/transformer/moe/token_dispatcher_inference_v.py b/megatron/core/transformer/moe/token_dispatcher_inference_v.py index 0dd8a481f33..cc1887204c4 100644 --- a/megatron/core/transformer/moe/token_dispatcher_inference_v.py +++ b/megatron/core/transformer/moe/token_dispatcher_inference_v.py @@ -15,12 +15,15 @@ 5. Triton unpermute kernel scatters weighted expert outputs back. 6. ReduceScatter combines contributions across EP ranks. -When ``inference_moe_max_tokens`` is set, AllGather/ReduceScatter buffers are -pinned to a fixed maximum size. This makes the NCCL collectives identical +When ``inference_moe_max_tokens`` is set (automatically derived from +``--inference-dynamic-batching-max-tokens``), AllGather/ReduceScatter buffers +are pinned to a fixed maximum size. This makes the NCCL collectives identical across every CUDA graph, so different EP ranks can independently select different graphs without any cross-rank synchronization -- the all-reduce in ``adjust_batch_dims_for_expert_parallelism`` is no longer required. +Selected via ``--moe-token-dispatcher-type allgather_v``. + The "V" (variable) refers to the fact that each expert receives a data-dependent number of tokens, computed on-device without cross-rank synchronization. """ @@ -97,8 +100,9 @@ def __init__( if self.ep_size > 1 and self._max_tokens_per_rank is None: raise ValueError( "inference_moe_max_tokens must be set when using the 'allgather_v' " - "dispatcher with EP > 1. Set it to cuda_graph_max_tokens " - "(typically max_requests * (num_speculative_tokens + 1))." + "dispatcher with EP > 1. It is automatically derived from " + "--inference-dynamic-batching-max-tokens (or max_requests * " + "(num_speculative_tokens + 1)). Make sure one of these is set." ) # Cached between dispatch_postprocess and combine_preprocess. diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 9c27340587e..8a69dafd285 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -747,9 +747,12 @@ class TransformerConfig(ModelParallelConfig): specified capacity, similar to GShard, Switch-Transformer, and DeepSpeed-MoE. Note that this is currently unsupported so should remain False.""" - moe_token_dispatcher_type: Literal['allgather', 'alltoall', 'flex'] = "allgather" + moe_token_dispatcher_type: Literal['allgather', 'alltoall', 'flex', 'allgather_v'] = "allgather" """The type of token dispatcher to use. The default is 'allgather'. - Options are 'allgather','alltoall' and 'flex'.""" + Options are 'allgather', 'alltoall', 'flex', and 'allgather_v'. + 'allgather_v' is an inference-optimized dispatcher that uses Triton-based + permutation with fixed-max-buffer AllGather/ReduceScatter, eliminating + EP rank synchronization for CUDA graph matching.""" moe_enable_deepep: bool = False """[Experimental] Enable DeepEP for efficient token dispatching and combine in MoE models.""" @@ -943,20 +946,15 @@ class TransformerConfig(ModelParallelConfig): inference_moe_cuda_graph_dispatcher: Literal['fused', 'allgather_v'] = "fused" """Selects the CUDA-graph-compatible token dispatcher for inference MoE. + + Auto-derived from ``moe_token_dispatcher_type`` in ``__post_init__``: + - ``moe_token_dispatcher_type='allgather_v'`` → ``'allgather_v'`` + - anything else → ``'fused'`` + Options: - 'fused' (default): AllGather + FlashInfer/CUTLASS fused MoE kernel. - The fused kernel handles token permutation internally. Requires FlashInfer. - 'allgather_v': AllGather + Triton permute + grouped_mm + Triton unpermute - + ReduceScatter. Token counts per expert are computed on-device by Triton - kernels with no host-device synchronization, making the full pipeline - CUDA-graphable without FlashInfer. Supports BF16 and MXFP8 weights via - torch.nn.functional.grouped_mm / scaled_grouped_mm. - - When EP > 1, the AllGather/ReduceScatter buffers are pinned to a fixed - maximum size (inference_moe_max_tokens * ep_size) so that all CUDA graphs - embed the same collective regardless of actual batch size. This allows - different EP ranks to independently select different CUDA graphs without - any cross-rank synchronization (no all-reduce for batch dimension matching). + + ReduceScatter. No FlashInfer dependency. Fully CUDA-graphable. """ inference_moe_max_tokens: Optional[int] = None @@ -968,8 +966,9 @@ class TransformerConfig(ModelParallelConfig): adjust_batch_dims_for_expert_parallelism — each EP rank can independently select its own CUDA graph. - Should be set to cuda_graph_max_tokens (typically max_requests * - (num_speculative_tokens + 1)). Required when using 'allgather_v' with EP > 1. + Automatically populated from --inference-dynamic-batching-max-tokens (or + max_requests * (num_speculative_tokens + 1)) when --moe-token-dispatcher-type + is 'allgather_v'. Required when using 'allgather_v' with EP > 1. """ mrope_section: Optional[List[int]] = None @@ -1238,6 +1237,11 @@ def __post_init__(self): f"got '{self.inference_grouped_gemm_backend}'" ) + # Auto-derive inference_moe_cuda_graph_dispatcher from + # moe_token_dispatcher_type so users only need one CLI flag. + if self.moe_token_dispatcher_type == "allgather_v": + self.inference_moe_cuda_graph_dispatcher = "allgather_v" + assert self.inference_moe_cuda_graph_dispatcher in ('fused', 'allgather_v'), ( f"inference_moe_cuda_graph_dispatcher must be 'fused' or 'allgather_v', " f"got '{self.inference_moe_cuda_graph_dispatcher}'" diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index 0bfe6142d01..f45db993e13 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -1718,6 +1718,17 @@ def core_transformer_config_from_args(args, config_class=None): kw_args['moe_latent_size'] = args.moe_latent_size + # For the allgather_v dispatcher, derive inference_moe_max_tokens from the + # dynamic batching max_tokens so users don't need a separate flag. + if getattr(args, 'moe_token_dispatcher_type', None) == 'allgather_v': + max_tokens = getattr(args, 'inference_dynamic_batching_max_tokens', None) + max_requests = getattr(args, 'inference_dynamic_batching_max_requests', None) + if max_tokens is not None: + kw_args['inference_moe_max_tokens'] = max_tokens + elif max_requests is not None: + num_spec = getattr(args, 'num_speculative_tokens', 0) or 0 + kw_args['inference_moe_max_tokens'] = max_requests * (num_spec + 1) + if args.te_precision_config_file: assert not 'quant_recipe' in kw_args, "Quantization recipe already configured." # TODO(kwyss): Prohibit fp8_params or fp4_params with this flexibility @@ -3107,6 +3118,13 @@ def _add_moe_args(parser): group.add_argument('--moe-aux-loss-coeff', type=float, nargs='+', default=0.0, help='Scaling coefficient for the aux loss: a starting value of 1e-2 is recommended.') # Token dispatcher arguments + group.add_argument('--moe-token-dispatcher-type', type=str, + choices=['allgather', 'alltoall', 'flex', 'allgather_v'], + default='allgather', + help='Token dispatcher type for MoE layers. ' + '"allgather_v" is an inference-optimized dispatcher that uses ' + 'Triton-based permutation with fixed-max-buffer AllGather/ReduceScatter, ' + 'eliminating EP rank synchronization for CUDA graph matching.') # MoE communication overlap arguments group.add_argument('--moe-upcycling-granularity', type=int, default=1, From 8bae04fece38cb19456ff9266b916d5c966b482f Mon Sep 17 00:00:00 2001 From: Keshav Santhanam Date: Wed, 8 Apr 2026 13:16:16 -0700 Subject: [PATCH 3/7] Fix args Signed-off-by: Keshav Santhanam --- megatron/training/arguments.py | 8 -------- 1 file changed, 8 deletions(-) diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index f45db993e13..3b9e0abcd75 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -3117,14 +3117,6 @@ def _add_moe_args(parser): help='Determines the load balancing strategy for the router. "aux_loss" corresponds to the load balancing loss used in GShard and SwitchTransformer; "seq_aux_loss" corresponds to the load balancing loss used in DeepSeekV2, which computes the loss for each individual sample; "sinkhorn" corresponds to the balancing algorithm used in S-BASE, and "none" implies no load balancing. The default is "aux_loss".') group.add_argument('--moe-aux-loss-coeff', type=float, nargs='+', default=0.0, help='Scaling coefficient for the aux loss: a starting value of 1e-2 is recommended.') - # Token dispatcher arguments - group.add_argument('--moe-token-dispatcher-type', type=str, - choices=['allgather', 'alltoall', 'flex', 'allgather_v'], - default='allgather', - help='Token dispatcher type for MoE layers. ' - '"allgather_v" is an inference-optimized dispatcher that uses ' - 'Triton-based permutation with fixed-max-buffer AllGather/ReduceScatter, ' - 'eliminating EP rank synchronization for CUDA graph matching.') # MoE communication overlap arguments group.add_argument('--moe-upcycling-granularity', type=int, default=1, From 3aa1dbba7bff9d4ab078ccb47b43c84afd34edb5 Mon Sep 17 00:00:00 2001 From: Keshav Santhanam Date: Wed, 8 Apr 2026 14:47:00 -0700 Subject: [PATCH 4/7] Remove max_requests limitation on cuda graph token count Signed-off-by: Keshav Santhanam --- .../core/inference/batch_dimensions_utils.py | 18 ++++++++++++------ .../core/inference/contexts/dynamic_context.py | 11 ++++++++++- 2 files changed, 22 insertions(+), 7 deletions(-) diff --git a/megatron/core/inference/batch_dimensions_utils.py b/megatron/core/inference/batch_dimensions_utils.py index ea74c052126..bfba6bd35f5 100644 --- a/megatron/core/inference/batch_dimensions_utils.py +++ b/megatron/core/inference/batch_dimensions_utils.py @@ -329,6 +329,7 @@ def generate_cuda_graph_batch_dimensions_list( max_sequence_length: int, use_cuda_graphs_for_non_decode_steps: bool, num_speculative_tokens: int = 0, + skip_ep_sync: bool = False, ) -> Tuple[List[InferenceBatchDimensions], Optional[List[int]]]: """ Generate CUDA graph batch dimensions. @@ -394,12 +395,17 @@ def add_if_valid(token_count: int, prefill_req_count: int, decode_req_count: int ): cuda_graph_max_tokens = max_tokens - assert cuda_graph_max_tokens == max_requests * (num_speculative_tokens + 1), ( - f"cuda_graph_max_tokens ({cuda_graph_max_tokens}) must equal max_requests *" - f"(num_speculative_tokens + 1) ({max_requests * (num_speculative_tokens + 1)}). " - "This is required for correctly syncing EP ranks: " - f"prefill and decode graph pools must have the same token count granularity." - ) + # When EP ranks sync batch dimensions, prefill and decode graph + # pools must use the same token granularity. With skip_ep_sync + # (allgather_v dispatcher), ranks select graphs independently so + # cuda_graph_max_tokens can be larger (e.g. max_tokens). + if not skip_ep_sync: + assert cuda_graph_max_tokens == max_requests * (num_speculative_tokens + 1), ( + f"cuda_graph_max_tokens ({cuda_graph_max_tokens}) must equal max_requests *" + f"(num_speculative_tokens + 1) ({max_requests * (num_speculative_tokens + 1)}). " + "This is required for correctly syncing EP ranks: " + f"prefill and decode graph pools must have the same token count granularity." + ) if num_cuda_graphs != -1: # if -1, no need to adjust. This will be taken care of in diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index f5157cb2d76..4f3901bd726 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -590,17 +590,26 @@ def __init__(self, model_config: TransformerConfig, inference_config: InferenceC self.use_cuda_graphs_for_non_decode_steps = ( inference_config.use_cuda_graphs_for_non_decode_steps ) + # With the allgather_v dispatcher, CUDA graphs can handle up to + # max_tokens because EP ranks select graphs independently (no sync). + # Otherwise, cuda_graph_max_tokens is capped at max_requests to keep + # prefill and decode graph pools at the same token granularity. + if self._skip_ep_sync: + cuda_graph_max_tokens = self.max_tokens + else: + cuda_graph_max_tokens = self.max_requests * (self.num_speculative_tokens + 1) self.cuda_graph_batch_dimensions_list, self.cuda_graph_token_counts = ( CUDAGraphBatchDimensionBuilder.generate_cuda_graph_batch_dimensions_list( tp_size=tp_size, num_cuda_graphs=inference_config.num_cuda_graphs, - cuda_graph_max_tokens=self.max_requests * (self.num_speculative_tokens + 1), + cuda_graph_max_tokens=cuda_graph_max_tokens, cuda_graph_mixed_prefill_request_count=inference_config.cuda_graph_mixed_prefill_count, max_requests=self.max_requests, max_tokens=self.max_tokens, max_sequence_length=self.max_sequence_length, use_cuda_graphs_for_non_decode_steps=self.use_cuda_graphs_for_non_decode_steps, num_speculative_tokens=self.num_speculative_tokens, + skip_ep_sync=self._skip_ep_sync, ) ) From e52033e332a7b9d8780aafcb847942ba1bbd75da Mon Sep 17 00:00:00 2001 From: Keshav Santhanam Date: Wed, 8 Apr 2026 15:25:37 -0700 Subject: [PATCH 5/7] Fix dummy rank Signed-off-by: Keshav Santhanam --- megatron/core/transformer/moe/moe_layer.py | 74 ++++++++++++---------- 1 file changed, 39 insertions(+), 35 deletions(-) diff --git a/megatron/core/transformer/moe/moe_layer.py b/megatron/core/transformer/moe/moe_layer.py index 43dcf5e6b09..485ea87c20c 100644 --- a/megatron/core/transformer/moe/moe_layer.py +++ b/megatron/core/transformer/moe/moe_layer.py @@ -276,33 +276,38 @@ def __init__( ) # Initialize token dispatcher. - # 'allgather_v' is inference-only; for the base (training) dispatcher - # we use alltoall since _setup_inference_mode swaps in the real - # InferenceAllGatherVTokenDispatcher during CUDA-graphed iterations. - base_dispatcher_type = config.moe_token_dispatcher_type - if base_dispatcher_type == "allgather_v": - base_dispatcher_type = "alltoall" - if base_dispatcher_type == "allgather": + # 'allgather_v' uses the same InferenceAllGatherVTokenDispatcher for both + # eager and CUDA-graph modes — this ensures EP collectives (AllGather / + # ReduceScatter) are consistent regardless of whether a rank replays a + # graph or falls back to eager (e.g. dummy EP ranks). + if config.moe_token_dispatcher_type == "allgather": self.token_dispatcher = MoEAllGatherTokenDispatcher( self.num_local_experts, self.local_expert_indices, config=self.config, pg_collection=pg_collection, ) - elif base_dispatcher_type == "alltoall": + elif config.moe_token_dispatcher_type == "alltoall": self.token_dispatcher = MoEAlltoAllTokenDispatcher( self.num_local_experts, self.local_expert_indices, config=self.config, pg_collection=pg_collection, ) - elif base_dispatcher_type == "flex": + elif config.moe_token_dispatcher_type == "flex": self.token_dispatcher = MoEFlexTokenDispatcher( self.num_local_experts, self.local_expert_indices, config=self.config, pg_collection=pg_collection, ) + elif config.moe_token_dispatcher_type == "allgather_v": + self.token_dispatcher = InferenceAllGatherVTokenDispatcher( + self.num_local_experts, + self.local_expert_indices, + config=self.config, + pg_collection=pg_collection, + ) else: raise ValueError( f"Unsupported token dispatcher type: {config.moe_token_dispatcher_type}" @@ -382,12 +387,12 @@ def _setup_inference_mode(self, pg_collection): dispatcher_type = self.config.inference_moe_cuda_graph_dispatcher if dispatcher_type == "allgather_v": - self._inference_token_dispatcher = InferenceAllGatherVTokenDispatcher( - self.num_local_experts, - self.local_expert_indices, - config=self.config, - pg_collection=pg_collection, - ) + # The base dispatcher is already InferenceAllGatherVTokenDispatcher + # (set in __init__), which uses the same AllGather/ReduceScatter + # collectives in both eager and graph modes. No swap needed — + # set _inference_token_dispatcher to None so + # set_inference_cuda_graphed_iteration skips the swap. + self._inference_token_dispatcher = None else: self._inference_token_dispatcher = InferenceCUDAGraphTokenDispatcher( self.num_local_experts, @@ -523,29 +528,28 @@ def routed_experts_compute(self, hidden_states: torch.Tensor, probs: torch.Tenso dispatched_input, tokens_per_expert, permuted_probs = ( self.token_dispatcher.dispatch_postprocess(hidden_states, probs) ) - if ( + if isinstance(self.token_dispatcher, InferenceAllGatherVTokenDispatcher): + # AllGatherV dispatcher: tokens are pre-permuted with Triton — + # pass expert_offsets and permutation_map for grouped_mm. + expert_output, mlp_bias = apply_module(self.experts)( + dispatched_input, + tokens_per_expert, + permuted_probs, + expert_offsets=self.token_dispatcher.expert_offsets, + permutation_map=self.token_dispatcher.permutation_map, + ) + elif ( hasattr(self, "_inference_token_dispatcher") and self.is_inference_cuda_graphed_iteration ): - # AllGatherV dispatcher: tokens are pre-permuted with Triton — - # pass expert_offsets and permutation_map for grouped_mm. - if isinstance(self.token_dispatcher, InferenceAllGatherVTokenDispatcher): - expert_output, mlp_bias = apply_module(self.experts)( - dispatched_input, - tokens_per_expert, - permuted_probs, - expert_offsets=self.token_dispatcher.expert_offsets, - permutation_map=self.token_dispatcher.permutation_map, - ) - else: - # Fused dispatcher: pass routing_map for FlashInfer. - routing_map = self.token_dispatcher.routing_map - expert_output, mlp_bias = apply_module(self.experts)( - dispatched_input, - tokens_per_expert, - permuted_probs, - routing_map=routing_map, - ) + # Fused dispatcher: pass routing_map for FlashInfer. + routing_map = self.token_dispatcher.routing_map + expert_output, mlp_bias = apply_module(self.experts)( + dispatched_input, + tokens_per_expert, + permuted_probs, + routing_map=routing_map, + ) else: expert_output, mlp_bias = apply_module(self.experts)( dispatched_input, tokens_per_expert, permuted_probs From 365767601a1972ae118ee07c06650e6ebc70f4d8 Mon Sep 17 00:00:00 2001 From: Keshav Santhanam Date: Wed, 8 Apr 2026 15:40:54 -0700 Subject: [PATCH 6/7] Bug fix Signed-off-by: Keshav Santhanam --- megatron/core/inference/contexts/dynamic_context.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index 4f3901bd726..dca366ecd20 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -591,11 +591,14 @@ def __init__(self, model_config: TransformerConfig, inference_config: InferenceC inference_config.use_cuda_graphs_for_non_decode_steps ) # With the allgather_v dispatcher, CUDA graphs can handle up to - # max_tokens because EP ranks select graphs independently (no sync). + # inference_moe_max_tokens per rank because EP ranks select graphs + # independently (no sync) and AllGather buffers are pinned to that size. # Otherwise, cuda_graph_max_tokens is capped at max_requests to keep # prefill and decode graph pools at the same token granularity. if self._skip_ep_sync: - cuda_graph_max_tokens = self.max_tokens + cuda_graph_max_tokens = getattr( + model_config, "inference_moe_max_tokens", None + ) or self.max_requests * (self.num_speculative_tokens + 1) else: cuda_graph_max_tokens = self.max_requests * (self.num_speculative_tokens + 1) self.cuda_graph_batch_dimensions_list, self.cuda_graph_token_counts = ( From bb7b1287945e6e4f260313805348b53af18d96a7 Mon Sep 17 00:00:00 2001 From: Keshav Santhanam Date: Wed, 8 Apr 2026 15:54:00 -0700 Subject: [PATCH 7/7] Bug fix Signed-off-by: Keshav Santhanam --- megatron/core/transformer/moe/router.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/megatron/core/transformer/moe/router.py b/megatron/core/transformer/moe/router.py index a773775a299..d4362e5cd0e 100644 --- a/megatron/core/transformer/moe/router.py +++ b/megatron/core/transformer/moe/router.py @@ -825,7 +825,14 @@ def forward(self, input: torch.Tensor, padding_mask: Optional[torch.Tensor] = No - top_indices: Selected expert indices [num_tokens, topk] """ - if self.training or not self.is_inference_cuda_graphed_iteration: + # The allgather_v dispatcher always needs dense [num_tokens, topk] output + # (expert indices), not the sparse [num_tokens, num_experts] boolean mask. + # Use dense output in both eager and graph modes so collectives stay consistent. + use_dense = ( + self.is_inference_cuda_graphed_iteration + or self.config.inference_moe_cuda_graph_dispatcher == "allgather_v" + ) + if self.training or not use_dense: return super().forward(input, padding_mask) return self._forward(input, padding_mask)