Skip to content

[dev] Fix HybridEP token equalization under torch.compile without CUDA graphs #5668

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
HaochenYuan wants to merge 4 commits into
base: dev
Choose a base branch
from
+14 −15
Open

[dev] Fix HybridEP token equalization under torch.compile without CUDA graphs #5668

Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
db39dda
Fix HybridEP token equalization under torch.compile without CUDA graphs
HaochenYuan Jul 6, 2026
91b25ef
refactor
HaochenYuan Jul 6, 2026
7de9a00
fix linting
HaochenYuan Jul 7, 2026
2800534
Merge branch 'dev' into fix/hybridep-static-padding-detection
HaochenYuan Jul 7, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
29 changes: 14 additions & 15 deletions megatron/core/transformer/moe/token_dispatcher.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1068,22 +1068,19 @@ 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.
padded_num_tokens = num_tokens
else:
):
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 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(
Expand All @@ -1093,12 +1090,14 @@ 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

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
Expand Down
Loading