Add AllGather-V token dispatcher for inference

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
@@ -313,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.
@@ -378,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
@@ -500,6 +522,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 +538,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 +556,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:

megatron/core/inference/contexts/dynamic_context.py

@@ -579,21 +579,40 @@ 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
         )
+        # With the allgather_v dispatcher, CUDA graphs can handle up to
+        # 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 = 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 = (
             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,
             )
         )
 
@@ -1608,6 +1627,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
 

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

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,