Skip to content

fix(moe): exclude padding tokens from dropless HybridEP routing - #5515

Open
seonjinn wants to merge 6 commits into
NVIDIA:mainfrom
seonjinn:sna/hybridep-seqpack-padding-mask-20260626
Open

fix(moe): exclude padding tokens from dropless HybridEP routing#5515
seonjinn wants to merge 6 commits into
NVIDIA:mainfrom
seonjinn:sna/hybridep-seqpack-padding-mask-20260626

Conversation

@seonjinn

@seonjinn seonjinn commented Jun 26, 2026

Copy link
Copy Markdown
Contributor
  • I, the PR author, have personally reviewed every line of this PR.

What does this PR do?

Exclude physical padding rows from MoE routing before dropless HybridEP dispatch.

This ports the routing-mask behavior already merged to dev in #5542 onto main:

  • zero probabilities and clear routing assignments for padding rows;
  • apply the behavior only to flex + hybridep with no expert or rank capacity factor;
  • preserve existing routing behavior for other dispatchers, backends, and capacity modes.

Why?

#6114 fixes padding-mask broadcasting for expert-bias token counts, but it does not remove physical padding rows from HybridEP dispatch. NeMo-RL sequence-packed THD workloads construct and pass this mask in NVIDIA-NeMo/RL#2964; MCore must consume it before dispatch to prevent padding rows from reaching experts.

After merging current main, this PR drops its older changes that are now owned by #6114 and the current sequence-parallel mask-alignment paths. The final diff is limited to the router behavior and focused unit coverage.

Validation

Validation Result
H100, 8-rank focused test on unpatched main Expected RED: 2 failed, 8 passed; fused and unfused HybridEP retained 64 nonzero padding routes
H100, 8-rank focused test with this change GREEN: 10 passed, 54 deselected
Non-HybridEP and capacity-mode controls All 6 parameterized cases passed
#6114 expert-bias masked/unmasked controls Both cases passed
Static checks py_compile, Black, isort, Ruff, and git diff --check passed

The focused H100 run used the final router and test contents. The branch then merged the latest upstream main; that upstream commit did not touch either changed file.

Related

Signed-off-by: Seonjin Na <sna@nvidia.com>
@seonjinn
seonjinn requested review from a team as code owners June 26, 2026 22:23
@copy-pr-bot

copy-pr-bot Bot commented Jun 26, 2026

Copy link
Copy Markdown

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

@svcnvidia-nemo-ci
svcnvidia-nemo-ci marked this pull request as draft June 26, 2026 22:26
@github-actions

Copy link
Copy Markdown
Contributor

This PR has been automatically converted to draft because all PRs must start as drafts.

When you are ready for review, click Ready for Review to begin the review process. This will:

  1. Add the oncall reviewer (optional reviewer)
  2. Add required review teams based on your changes

See the contribution guide for more details.

Signed-off-by: Seonjin Na <sna@nvidia.com>
@yaox12
yaox12 requested a review from Autumn1998 June 29, 2026 02:18
Signed-off-by: seonjinn <sna@nvidia.com>
Signed-off-by: seonjinn <sna@nvidia.com>
Signed-off-by: seonjinn <sna@nvidia.com>
Signed-off-by: seonjinn <sna@nvidia.com>
@seonjinn seonjinn changed the title Handle HybridEP packed padding masks fix(moe): exclude padding tokens from dropless HybridEP routing Aug 22, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant