From db39dda8fa59eddb2c895b42b67f5045b58b7b88 Mon Sep 17 00:00:00 2001 From: HaochenYuan Date: Mon, 6 Jul 2026 04:18:00 -0700 Subject: [PATCH 1/3] Fix HybridEP token equalization under torch.compile without CUDA graphs Signed-off-by: HaochenYuan --- .../core/transformer/moe/token_dispatcher.py | 29 ++++++++++--------- 1 file changed, 16 insertions(+), 13 deletions(-) diff --git a/megatron/core/transformer/moe/token_dispatcher.py b/megatron/core/transformer/moe/token_dispatcher.py index 61bd7a6f94c..e5af99ab70a 100644 --- a/megatron/core/transformer/moe/token_dispatcher.py +++ b/megatron/core/transformer/moe/token_dispatcher.py @@ -1068,20 +1068,23 @@ def setup_metadata(self, routing_map: torch.Tensor, probs: torch.Tensor): self._original_num_tokens = num_tokens padded_num_tokens = num_tokens - equalize_thd_token_counts = ( + if ( self.config.sequence_packing_scheduler is not None or self.config.moe_hybridep_pad_variable_tokens - ) - if equalize_thd_token_counts: - if self.config.sequence_packing_scheduler is not None and ( - torch.cuda.is_current_stream_capturing() or torch.compiler.is_compiling() - ): - # CUDA graph path: routing_map has already been padded to a static - # length upstream (CUDA graph + sequence packing implies - # cu_seqlens_q_padded -> max_seqlen_per_dp_cp_rank), so num_tokens - # is identical across the EP communication group. Skip the - # all_reduce + .item() during both dynamo tracing and stream - # capture, and use the local value directly. + ): + pad_alignment = self.config.pad_packed_seq_alignment + has_static_token_count = ( + self.config.sequence_packing_scheduler is not None + and pad_alignment is not None + and ( + pad_alignment == "max" + or pad_alignment == self.config.max_seqlen_per_dp_cp_rank + ) + ) + if has_static_token_count: + # The sequence-packing path has already padded every rank to the + # same configured maximum. This is a data invariant independent + # of Dynamo tracing or CUDA graph capture state. padded_num_tokens = num_tokens else: # Use the actual tp_ep max so all ranks in the MoE communication @@ -1098,7 +1101,7 @@ def setup_metadata(self, routing_map: torch.Tensor, probs: torch.Tensor): routing_map = routing_map.reshape(num_tokens, self.num_experts) probs = probs.reshape(num_tokens, self.num_experts) - if equalize_thd_token_counts and padded_num_tokens > num_tokens: + if padded_num_tokens > num_tokens: pad_rows = padded_num_tokens - num_tokens routing_map = torch.cat( [routing_map, routing_map.new_zeros((pad_rows, self.num_experts))], dim=0 From 91b25eff360c13b208b257628b33c674dbefa7bd Mon Sep 17 00:00:00 2001 From: HaochenYuan Date: Mon, 6 Jul 2026 04:22:55 -0700 Subject: [PATCH 2/3] refactor Signed-off-by: HaochenYuan --- megatron/core/transformer/moe/token_dispatcher.py | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/megatron/core/transformer/moe/token_dispatcher.py b/megatron/core/transformer/moe/token_dispatcher.py index e5af99ab70a..52fdcdccfd4 100644 --- a/megatron/core/transformer/moe/token_dispatcher.py +++ b/megatron/core/transformer/moe/token_dispatcher.py @@ -1081,12 +1081,7 @@ def setup_metadata(self, routing_map: torch.Tensor, probs: torch.Tensor): or pad_alignment == self.config.max_seqlen_per_dp_cp_rank ) ) - if has_static_token_count: - # The sequence-packing path has already padded every rank to the - # same configured maximum. This is a data invariant independent - # of Dynamo tracing or CUDA graph capture state. - padded_num_tokens = num_tokens - else: + if not has_static_token_count: # Use the actual tp_ep max so all ranks in the MoE communication # group pass the same token count to HybridEP. max_num_tokens_across_ep = torch.tensor( @@ -1096,6 +1091,8 @@ def setup_metadata(self, routing_map: torch.Tensor, probs: torch.Tensor): max_num_tokens_across_ep, op=torch.distributed.ReduceOp.MAX, group=self.group ) padded_num_tokens = int(max_num_tokens_across_ep.item()) + # Otherwise, the sequence-packing path has already padded every rank + # to the same configured maximum, so the local token count can be used. padded_num_tokens += -padded_num_tokens % HYBRIDEP_TOKEN_ALIGNMENT self._padded_num_tokens = padded_num_tokens From 7de9a00d4ff17eb40840977673fd57d385d56b33 Mon Sep 17 00:00:00 2001 From: HaochenYuan Date: Tue, 7 Jul 2026 05:19:47 -0700 Subject: [PATCH 3/3] fix linting Signed-off-by: HaochenYuan --- megatron/core/transformer/moe/token_dispatcher.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/megatron/core/transformer/moe/token_dispatcher.py b/megatron/core/transformer/moe/token_dispatcher.py index 52fdcdccfd4..6cd45faa221 100644 --- a/megatron/core/transformer/moe/token_dispatcher.py +++ b/megatron/core/transformer/moe/token_dispatcher.py @@ -1077,8 +1077,7 @@ def setup_metadata(self, routing_map: torch.Tensor, probs: torch.Tensor): self.config.sequence_packing_scheduler is not None and pad_alignment is not None and ( - pad_alignment == "max" - or pad_alignment == self.config.max_seqlen_per_dp_cp_rank + pad_alignment == "max" or pad_alignment == self.config.max_seqlen_per_dp_cp_rank ) ) if not has_static_token_count: