Repository navigation
[https://nvbugs/6892493][fix] Support MegaMoE CuteDSL TEP for all models and fix the Kimi-K3 CUDA-graph hang (cherry-pick #20055 + part of #19960) - #20619
Conversation
|
/bot run --extra-stage "GB300-8_GPUs-2_Nodes-PyTorch-SingleNvlinkDomain-Post-Merge-1" |
|
PR_Github #78995 [ run ] triggered by Bot. Commit: |
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info
WalkthroughThe changes define a global routed-output contract for fused-communication MoE backends. The scheduler and model callers use that contract to control routed reductions. The Kimi runtime changes shared-expert all-reduce ordering for a specific MegaMoE configuration. The changes also expand MegaMoE TEP eligibility and token-metadata validation. ChangesFused MoE routed-output handling
Priority: ➖ Normal Estimated code review effort: 3 (Moderate) | ~25 minutes Change: Bug fix Sequence Diagram(s)sequenceDiagram
participant Runtime as KimiK3MoERuntime
participant Routed as Routed branch
participant Event as CUDA event
participant Shared as Shared branch
Runtime->>Routed: Run routed work
Routed->>Event: Record completion before norm/up tail
Shared->>Event: Wait for routed completion
Shared->>Shared: Perform shared all-reduce
Possibly related issues
Possibly related PRs
Suggested reviewers: Merge Risk: ⚪ Minimal · up to
|
| Check name | Status | Explanation | Resolution |
|---|---|---|---|
| Docstring Coverage |  |
Docstring coverage is 42.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 50 functions across 11 files. (1 skipped:… | Write docstrings for the functions missing them to satisfy the coverage threshold. |
-
Check name Status Explanation Title check The title follows the required NVBugs and fix format and clearly summarizes both main changes: MegaMoE CuteDSL TEP support and the Kimi-K3 CUDA-graph hang fix. Description check The description includes the required sections, explains the problem and solution, lists relevant test coverage, and marks the checklist as reviewed. Linked Issues check Check skipped because no linked issues were found for this pull request. Out of Scope Changes check Check skipped because no linked issues were found for this pull request.
Full details: Docstring Coverage
-
Explanation
Docstring coverage is 42.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 50 functions across 11 files. (1 skipped: 1 unsupported.)
Resolution
Write docstrings for the functions missing them to satisfy the coverage threshold.
✨ Finishing Touches
-
🧪 Generate unit tests (beta)
-
- Create a new PR
-
- Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
Comment @coderabbitai help to get the list of available commands.
|
PR_Github #78995 [ run ] completed with state
|
|
/bot run --disable-fail-fast --extra-stage "GB300-8_GPUs-2_Nodes-PyTorch-SingleNvlinkDomain-Post-Merge-1" |
|
PR_Github #79014 [ run ] triggered by Bot. Commit: |
|
PR_Github #79014 [ run ] completed with state
|
xxi-nv
left a comment
xxi-nv
left a comment
There was a problem hiding this comment.
Approved based on [B] 6892493: [TensorRT-LLM][release/1.3]: accuracy/test_kimi3.py::TestKimiK3DSpark::test_gsm8k_tep8 got value error
|
Automatically added "ci: full pre-merge approved" because this PR has satisfied the required GitHub review approvals. Unresolved review conversations and other required checks remain independent merge requirements. |
|
/bot run --disable-fail-fast |
|
PR_Github #79112 [ run ] triggered by Bot. Commit: |
|
PR_Github #79112 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #79197 [ run ] triggered by Bot. Commit: |
|
PR_Github #79197 [ run ] completed with state |
|
/bot run --disable-fail-fast --extra-stage "GB300-8_GPUs-2_Nodes-PyTorch-SingleNvlinkDomain-Post-Merge-1" |
3ec3da3 to
c854a3f
Compare
|
/bot run --disable-fail-fast --extra-stage "GB300-8_GPUs-2_Nodes-PyTorch-SingleNvlinkDomain-Post-Merge-1" |
|
PR_Github #79275 [ run ] triggered by Bot. Commit: |
|
/bot run --disable-fail-fast |
|
PR_Github #79642 [ run ] triggered by Bot. Commit: |
|
PR_Github #79615 [ run ] completed with state |
schetlur-nv
left a comment
schetlur-nv
left a comment
There was a problem hiding this comment.
rubber stamp for devs since this is a cherry pick from main
|
PR_Github #79642 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #79708 [ run ] triggered by Bot. Commit: |
|
PR_Github #79708 [ run ] completed with state |
…uce past fused-comm MegaMoE Backport the shared-expert all-reduce ordering from NVIDIA#19960 (e88d19c) to release/1.3. With MEGAMOE_CUTEDSL under TEP, the fused-comm kernel runs its own cross-rank exchange while the shared-expert all-reduce ran concurrently on the aux stream, which deadlocks at CUDA-graph warmup. The shared all-reduce now waits for the routed MoE. Only this ordering change and its unit tests are taken; the rest of NVIDIA#19960 is Rubin performance work and is not needed on release/1.3. Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com>
Cherry-pick of NVIDIA#20055 to release/1.3. MegaMoE CuteDSL under TEP (attention TP + MoE EP) was gated to MiniMax-style models because the fused kernel returns a globally combined routed output that other model wrappers would all-reduce again. ConfigurableMoE now exposes returns_global_routed_output, callers opt in with accepts_global_routed_output, and FusedCommMoEScheduler hands a rank partial to callers that reduce externally, so the gate is removed. Kimi-K3 and MiniMax-M3 opt in. Strict MoE resolution failures now include each candidate's rejection reason. Unwaive the Kimi-K3 DSpark TEP8 accuracy test. Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com>
…adata in FusedCommMoEScheduler With MegaMoE CuteDSL now allowed under TEP, a DSpark drafter without attention DP passes all_rank_num_tokens=[num_tokens]. _strip_adp_padding treated that as per-EP-rank metadata and indexed it by moe_ep_rank, so ranks 1+ raised IndexError. Treat a single count as local metadata and reject lists that do not hold one count per EP rank. Backport of the FusedCommMoEScheduler fix and tests from NVIDIA#19959. Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com>
cb49e39 to
73fcc67
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #79949 [ run ] triggered by Bot. Commit: |
|
PR_Github #79949 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #80041 [ run ] triggered by Bot. Commit: |
|
PR_Github #80041 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #80177 [ run ] triggered by Bot. Commit: |
|
PR_Github #80177 [ run ] completed with state |
Dev Engineer Review
ConfigurableMoEadds a default-false global routed-output opt-in. Opted-in callers can skip their routed all-reduce when the backend returns globally combined output; other callers retain the rank-partial behavior.#479, succeeded; failure details and the result of any later rerun are unavailable.QA Engineer Review
tests/integration/test_lists/test-db/l0_b300.ymladds CI selectors for singleton token metadata, full-EP indexing, and rejection of partial EP metadata.tests/integration/test_lists/waives.txtremoves the waiver forTestKimiK3DSpark::test_gsm8k_tep8. The test remains selected intests/integration/test_lists/qa/llm_function_core.txtand is listed withTIMEOUT (60)intest-db/l0_gb300_multi_nodes_node2_gpu8.yml.test_kimi_k3_mlp.pyinl0_h100.ymlandl0_b200.yml;test_moe_backend.pyin severaltest-dbfiles, includingl0_b300.yml; andtest_minimax_m3.pyinl0_cpu.yml.Per-File QA Perspective
tensorrt_llm/_torch/models/modeling_kimi_linear.py— Verify shared-expert reduction ordering for serial, parallel, compiled, and multi-stream execution.tensorrt_llm/_torch/models/modeling_minimaxm3.py— Verify that MiniMax-M3 skips routed reduction only when the MoE reports globally combined output.tensorrt_llm/_torch/moe/fused_moe/configurable_moe.py— Verify the default-false opt-in and backend-dependent global-output contract.tensorrt_llm/_torch/moe/fused_moe/create_moe.py— Verify option forwarding and unchanged behavior for callers that do not opt in.tensorrt_llm/_torch/moe/fused_moe/interface.py— Verify that the new global-output contract defaults to false.tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_cute_dsl.py— Verify broader TEP acceptance, continued MoE-TP rejection, and the repeated-compute warning.tensorrt_llm/_torch/moe/fused_moe/moe_resolution.py— Verify strict errors include backend rejection details.tensorrt_llm/_torch/moe/fused_moe/moe_scheduler.py— Verify rank-partial versus global output behavior and token-metadata validation.tests/unittest/_torch/moe/test_kimi_k3_mlp.py— Covers shared-expert construction and reduction ordering. The file is selected inl0_h100.ymlandl0_b200.yml.tests/unittest/_torch/moe/test_moe_backend.py— Covers backend eligibility, output handling, token metadata, and strict-resolution errors. Selectors appear in severaltest-dbfiles, includingl0_b300.yml.tests/unittest/_torch/models/test_minimax_m3.py— Covers MiniMax-M3 routed-output behavior. The file is selected inl0_cpu.yml.tests/integration/test_lists/test-db/l0_b300.yml— Adds CI entries for singleton, full-EP, and rejected partial token metadata cases.tests/integration/test_lists/waives.txt— Removes the Kimi-K3 accuracy-test waiver; verify the now-unwaived test in the applicable integration run.Description
This brings two fixes to release/1.3 so Kimi-K3 can run MegaMoE CuteDSL under TEP (attention TP + MoE EP). It fixes https://nvbugs/6892493 and https://nvbugs/6946611.
_defer_shared_allreduce), with its unit tests. [None][feat] Kimi K3 optimizations on Rubin #19960 is otherwise a large Rubin optimization PR, so only this ordering change is taken.ConfigurableMoEnow reports this throughreturns_global_routed_output, so callers that opt in (Kimi-K3, MiniMax-M3) skip their own all-reduce, and other callers get a rank partial. The TEP gate is removed. Strict MoE resolution errors now include each backend's rejection reason.The #20055 conflicts came only from main-only fields (
unfinalized_combine_fn,supports_routing_aux), which are left out.Test Coverage
tests/unittest/_torch/moe/test_kimi_k3_mlp.py: shared-expert construction andtest_kimi_k3_shared_allreduce_ordering(from [None][feat] Kimi K3 optimizations on Rubin #19960).tests/unittest/_torch/moe/test_moe_backend.pyandtests/unittest/_torch/models/test_minimax_m3.py: TEP resolution, routed-output contract, and strict resolution errors (from [https://nvbugs/6892493][fix] Support MegaMoE CuteDSL TEP for all models #20055).accuracy/test_kimi3.py::TestKimiK3DSpark::test_gsm8k_tep8is unwaived (GB300 2-node post-merge stage), with CUDA graphs on.PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
If PR introduces API changes, an appropriate PR label is added - either
api-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin the PR title.Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.