Fix EP=1 inference by allocating buffers anyway

megatron/core/inference/contexts/dynamic_context.py

@@ -36,6 +36,7 @@
 from megatron.core.package_info import __version__ as mcore_version
 from megatron.core.transformer import MLATransformerConfig, TransformerConfig
 from megatron.core.transformer.moe.token_dispatcher_inference import (
+    InferenceAllGatherDispatcherBase,
     NCCLAllGatherDispatcher,
     NVLSAllGatherVDispatcher,
 )
@@ -668,19 +669,23 @@ 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.
-        if get_pg_size(self.expert_model_parallel_group) > 1:
-            if self._nccl_ep_dispatcher:
-                NCCLAllGatherDispatcher.allocate_buffers()
-            else:
-                # Use moe_latent_size if set (latent MoE: SuperV3, UltraV3), else hidden_size.
-                moe_hidden_size = model_config.moe_latent_size or model_config.hidden_size
-                NVLSAllGatherVDispatcher.allocate_buffers(
-                    per_rank_worst_case_token_count=self.round_up_tokens(self.max_tokens)
-                    // tp_size,
-                    topk=model_config.moe_router_topk,
-                    hidden_size=moe_hidden_size,
-                    ep_group=self.expert_model_parallel_group,
-                )
+        # 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':
+            NCCLAllGatherDispatcher.allocate_buffers()
+        elif get_pg_size(self.expert_model_parallel_group) > 1:
+            # Use moe_latent_size if set, else hidden_size.
+            moe_hidden_size = model_config.moe_latent_size or model_config.hidden_size
+            NVLSAllGatherVDispatcher.allocate_buffers(
+                per_rank_worst_case_token_count=self.round_up_tokens(self.max_tokens) // tp_size,
+                topk=model_config.moe_router_topk,
+                hidden_size=moe_hidden_size,
+                ep_group=self.expert_model_parallel_group,
+            )
+        else:
+            # EP=1 with nvls: skip symmetric memory init (requires NVLink between
+            # multiple GPUs) and just initialize the shared valid_tokens scalar.
+            InferenceAllGatherDispatcherBase.allocate_valid_tokens_tensor()
 
         # Deal with chunked prefill
         self.enable_chunked_prefill = inference_config.enable_chunked_prefill

megatron/core/transformer/moe/token_dispatcher_inference.py

@@ -80,6 +80,16 @@ def _valid_tokens(cls) -> torch.Tensor:
     def _get_host_valid_tokens_estimate(cls) -> Optional[int]:
         return cls._host_valid_tokens_estimate
 
+    @classmethod
+    def allocate_valid_tokens_tensor(cls) -> None:
+        """Allocate the per-step valid-tokens scalar shared across all dispatcher subclasses.
+
+        Called at model init from the dynamic context to ensure the buffer receives a valid pointer.
+        Must run outside CUDA graph capture so the stable address is available during replay.
+        """
+        device = torch.cuda.current_device()
+        cls._valid_tokens_tensor = torch.zeros(1, dtype=torch.int32, device=device)
+
     def update_metadata(self, local_tokens: int) -> None:
         """Per-step metadata refresh fired from the first instance's token_dispatch.
 
@@ -178,6 +188,8 @@ def token_dispatch(self, hidden_states, probs):
             Also updates self.routing_map to [total_tokens, topk].
         """
         if self.ep_size == 1:
+            if self._runs_metadata_sync:
+                InferenceAllGatherDispatcherBase._valid_tokens_tensor.fill_(hidden_states.shape[0])
             return hidden_states, probs
 
         if self._runs_metadata_sync:
@@ -505,6 +517,8 @@ def token_dispatch(self, hidden_states, probs):
             Also updates self.routing_map to [global_max, topk] int64.
         """
         if self.ep_size == 1:
+            if self._runs_metadata_sync:
+                InferenceAllGatherDispatcherBase._valid_tokens_tensor.fill_(hidden_states.shape[0])
             return hidden_states, probs
 
         if self._runs_metadata_sync: