From 3697d0b420b629961288edb08eef8c0d19a6ed8a Mon Sep 17 00:00:00 2001 From: Helen Ngo Date: Tue, 9 Jun 2026 12:29:23 -0700 Subject: [PATCH 1/4] fix EP=1 inference Signed-off-by: Helen Ngo --- .../inference/contexts/dynamic_context.py | 26 +++++++++---------- 1 file changed, 13 insertions(+), 13 deletions(-) diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index c3e8b4f5e1c..c25512823e2 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -668,19 +668,19 @@ 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, - ) + # NCCLAllGatherDispatcher is also needed at EP=1: mcore_fused_moe's Triton + # kernel reads _valid_tokens_tensor 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 (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, + ) # Deal with chunked prefill self.enable_chunked_prefill = inference_config.enable_chunked_prefill From fbfeef9eb9ebc4b0255bc3eec3c5b2fa7796a7c3 Mon Sep 17 00:00:00 2001 From: Helen Ngo Date: Tue, 9 Jun 2026 12:56:41 -0700 Subject: [PATCH 2/4] fix Signed-off-by: Helen Ngo --- megatron/core/inference/contexts/dynamic_context.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index c25512823e2..a4d2d65eb20 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -668,12 +668,12 @@ 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. - # NCCLAllGatherDispatcher is also needed at EP=1: mcore_fused_moe's Triton + # NCCLAllGatherDispatcher is still needed at EP=1: mcore_fused_moe's Triton # kernel reads _valid_tokens_tensor 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 (latent MoE: SuperV3, UltraV3), else hidden_size. + # 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, From 4ec84b2e6c8b5bafc1daad347add3a4a7ba8aed6 Mon Sep 17 00:00:00 2001 From: Helen Ngo Date: Tue, 9 Jun 2026 14:26:33 -0700 Subject: [PATCH 3/4] fix garbage generations for EP=1 Signed-off-by: Helen Ngo --- megatron/core/inference/contexts/dynamic_context.py | 8 ++++++-- .../core/transformer/moe/token_dispatcher_inference.py | 2 ++ 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index a4d2d65eb20..34af25f6b37 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -668,8 +668,8 @@ 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. - # NCCLAllGatherDispatcher is still needed at EP=1: mcore_fused_moe's Triton - # kernel reads _valid_tokens_tensor as a pointer regardless of EP size. + # 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: @@ -681,6 +681,10 @@ 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. + NCCLAllGatherDispatcher.allocate_buffers() # 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..d46831d0183 100644 --- a/megatron/core/transformer/moe/token_dispatcher_inference.py +++ b/megatron/core/transformer/moe/token_dispatcher_inference.py @@ -178,6 +178,7 @@ def token_dispatch(self, hidden_states, probs): Also updates self.routing_map to [total_tokens, topk]. """ if self.ep_size == 1: + InferenceAllGatherDispatcherBase._valid_tokens_tensor.fill_(hidden_states.shape[0]) return hidden_states, probs if self._runs_metadata_sync: @@ -505,6 +506,7 @@ def token_dispatch(self, hidden_states, probs): Also updates self.routing_map to [global_max, topk] int64. """ if self.ep_size == 1: + InferenceAllGatherDispatcherBase._valid_tokens_tensor.fill_(hidden_states.shape[0]) return hidden_states, probs if self._runs_metadata_sync: From 0bb287261179e474fbcabdbc5d92f26e72e5b60c Mon Sep 17 00:00:00 2001 From: Helen Ngo Date: Thu, 11 Jun 2026 06:37:03 -0700 Subject: [PATCH 4/4] address comments Signed-off-by: Helen Ngo --- .../core/inference/contexts/dynamic_context.py | 3 ++- .../moe/token_dispatcher_inference.py | 16 ++++++++++++++-- 2 files changed, 16 insertions(+), 3 deletions(-) diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index 34af25f6b37..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, ) @@ -684,7 +685,7 @@ def __init__(self, model_config: TransformerConfig, inference_config: InferenceC else: # EP=1 with nvls: skip symmetric memory init (requires NVLink between # multiple GPUs) and just initialize the shared valid_tokens scalar. - NCCLAllGatherDispatcher.allocate_buffers() + 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 d46831d0183..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,7 +188,8 @@ def token_dispatch(self, hidden_states, probs): Also updates self.routing_map to [total_tokens, topk]. """ if self.ep_size == 1: - InferenceAllGatherDispatcherBase._valid_tokens_tensor.fill_(hidden_states.shape[0]) + if self._runs_metadata_sync: + InferenceAllGatherDispatcherBase._valid_tokens_tensor.fill_(hidden_states.shape[0]) return hidden_states, probs if self._runs_metadata_sync: @@ -506,7 +517,8 @@ def token_dispatch(self, hidden_states, probs): Also updates self.routing_map to [global_max, topk] int64. """ if self.ep_size == 1: - InferenceAllGatherDispatcherBase._valid_tokens_tensor.fill_(hidden_states.shape[0]) + if self._runs_metadata_sync: + InferenceAllGatherDispatcherBase._valid_tokens_tensor.fill_(hidden_states.shape[0]) return hidden_states, probs if self._runs_metadata_sync: