Avoid unnecessary MoE router host synchronization - #6432
Merged
Conversation
svcnvidia-nemo-ci
marked this pull request as draft
August 11, 2026 06:51
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:
See the contribution guide for more details. |
Closed
JF-D
marked this pull request as ready for review
August 11, 2026 06:52
JF-D
force-pushed
the
jiangfeid/moe_router_async_host_sync
branch
from
August 11, 2026 06:57
e4c74fa to
dd7d892
Compare
Contributor
Author
|
@claude strict-review |
Contributor
Author
|
/ok to test dd7d892 |
jiemingz
approved these changes
Aug 11, 2026
fanshiqing
approved these changes
Aug 12, 2026
Partial MoE router CUDA graphs currently block the host after every replay, even when all token-dispatcher outputs remain on the GPU. Synchronize only when a graph output resides on the host, and lazily reuse the D2H completion event. Add unit coverage for both CUDA-only and host-output paths. Signed-off-by: Jiangfei Duan <jiangfeid@nvidia.com>
JF-D
force-pushed
the
jiangfeid/moe_router_async_host_sync
branch
from
August 12, 2026 15:37
dd7d892 to
3e5fbbf
Compare
ericharper
approved these changes
Aug 12, 2026
svcnvidia-nemo-ci
added
Approved
JF-D
enabled auto-merge
August 12, 2026 15:56
yaoyu-33
approved these changes
Aug 12, 2026
Contributor
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/31662952600 |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
6 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do?
Avoids blocking the host after partial MoE router CUDA graph replay when all
token-dispatcher outputs remain on the GPU.
The router path previously recorded and synchronized a CUDA event after every
graph replay. This synchronization is only required when asynchronous D2H
outputs must be consumed by eager host code.
This change:
Issue tracking
For PRs from open-source community contributors:
Linked issue:
Contribution process
Pre-checks