diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index c3e8b4f5e1c..712bf4cada6 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/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 diff --git a/megatron/core/transformer/moe/token_dispatcher_inference.py b/megatron/core/transformer/moe/token_dispatcher_inference.py index e85115528b8..081497f734c 100644 --- a/megatron/core/transformer/moe/token_dispatcher_inference.py +++ b/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: