[https://nvbugs/6166097][fix] Fix CuteDSL NVFP4 EPLB weight layout#15538
Conversation
|
/bot run |
📝 WalkthroughWalkthroughAdds a new overridable hook ChangesCuteDSL NVFP4 shared weight-scale finalization hook
Perf benchmark YAML config updates
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (2)
tests/scripts/perf/disaggregated/gb200_wideep_stress-deepseek-r1-fp4_8k1k_ctx2_gen1_dep32_bs128_eplb288_mtp3_ccb-NIXL.yaml (1)
49-49: 📐 Maintainability & Code Quality | 🔵 TrivialCoverage split is sufficient for this PR layer.
Coverage is appropriate with accuracy still exercised by:
tests/scripts/perf/disaggregated/gb200_wideep_accuracy-deepseek-r1-fp4_1k1k_ctx2_gen1_dep16_bs128_eplb288_mtp3_ccb-NIXL.yamltests/scripts/perf/disaggregated/gb200_wideep_accuracy-deepseek-r1-fp4_gpqa_diamond_1k1k_ctx2_gen1_dep16_bs128_eplb288_mtp3_ccb-NIXL.yamland stress-only runs now explicitly skipping accuracy in:
tests/scripts/perf/disaggregated/gb200_wideep_stress-deepseek-r1-fp4_1k1k_ctx2_gen1_dep16_bs128_eplb288_mtp3_ccb-NIXL.yamltests/scripts/perf/disaggregated/gb200_wideep_stress-deepseek-r1-fp4_8k1k_ctx2_gen1_dep32_bs128_eplb288_mtp3_ccb-NIXL.yamlNo additional follow-up test file is needed for this config-only cohort.
As per path instructions,
tests/**: "Act as a QA engineer reviewing test changes and coverage for TensorRT-LLM."🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/scripts/perf/disaggregated/gb200_wideep_stress-deepseek-r1-fp4_8k1k_ctx2_gen1_dep32_bs128_eplb288_mtp3_ccb-NIXL.yaml` at line 49, The review comment confirms that setting enable_accuracy_test to false in the gb200_wideep_stress-deepseek-r1-fp4_8k1k_ctx2_gen1_dep32_bs128_eplb288_mtp3_ccb-NIXL.yaml configuration file is appropriate because accuracy testing coverage is already provided by dedicated accuracy test configuration files. Ensure this setting remains in place in the stress test file, as the test coverage strategy intentionally separates accuracy validation (handled by dedicated accuracy test configs) from stress and performance testing (handled by stress-only test configs that skip accuracy).Source: Path instructions
tensorrt_llm/_torch/modules/fused_moe/quantization.py (1)
2543-2544: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAnnotate the new hooks as returning
None.These hooks only mutate module state; add explicit return annotations to keep the changed Python API aligned with the repo typing rule. As per coding guidelines, “Always annotate functions with return types (use
Noneif no return).”Proposed typing fix
- def _prepare_shared_weight_scales_for_finalization(self, - module: torch.nn.Module): + def _prepare_shared_weight_scales_for_finalization( + self, module: torch.nn.Module) -> None: @@ - def _prepare_shared_weights_for_finalization(self, module: torch.nn.Module): + def _prepare_shared_weights_for_finalization( + self, module: torch.nn.Module) -> None: @@ - def _prepare_shared_weight_scales_for_finalization(self, - module: torch.nn.Module): + def _prepare_shared_weight_scales_for_finalization( + self, module: torch.nn.Module) -> None:Also applies to: 3073-3073, 3088-3089
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/modules/fused_moe/quantization.py` around lines 2543 - 2544, Add explicit return type annotations to the hook methods that lack them. The method _prepare_shared_weight_scales_for_finalization and the other hook methods mentioned in the comment (at lines 3073 and 3088-3089) should be annotated with -> None since they only mutate module state and do not return values. This aligns with the repository's typing guidelines that require all functions to have explicit return type annotations.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@tensorrt_llm/_torch/modules/fused_moe/quantization.py`:
- Around line 2543-2544: Add explicit return type annotations to the hook
methods that lack them. The method
_prepare_shared_weight_scales_for_finalization and the other hook methods
mentioned in the comment (at lines 3073 and 3088-3089) should be annotated with
-> None since they only mutate module state and do not return values. This
aligns with the repository's typing guidelines that require all functions to
have explicit return type annotations.
In
`@tests/scripts/perf/disaggregated/gb200_wideep_stress-deepseek-r1-fp4_8k1k_ctx2_gen1_dep32_bs128_eplb288_mtp3_ccb-NIXL.yaml`:
- Line 49: The review comment confirms that setting enable_accuracy_test to
false in the
gb200_wideep_stress-deepseek-r1-fp4_8k1k_ctx2_gen1_dep32_bs128_eplb288_mtp3_ccb-NIXL.yaml
configuration file is appropriate because accuracy testing coverage is already
provided by dedicated accuracy test configuration files. Ensure this setting
remains in place in the stress test file, as the test coverage strategy
intentionally separates accuracy validation (handled by dedicated accuracy test
configs) from stress and performance testing (handled by stress-only test
configs that skip accuracy).
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 36f890c7-d76d-40c4-83aa-c4d5a203f3dc
📒 Files selected for processing (6)
tensorrt_llm/_torch/modules/fused_moe/quantization.pytests/scripts/perf/disaggregated/gb200_wideep_accuracy-deepseek-r1-fp4_1k1k_ctx2_gen1_dep16_bs128_eplb288_mtp3_ccb-NIXL.yamltests/scripts/perf/disaggregated/gb200_wideep_accuracy-deepseek-r1-fp4_gpqa_diamond_1k1k_ctx2_gen1_dep16_bs128_eplb288_mtp3_ccb-NIXL.yamltests/scripts/perf/disaggregated/gb200_wideep_stress-deepseek-r1-fp4_1k1k_ctx2_gen1_dep16_bs128_eplb288_mtp3_ccb-NIXL.yamltests/scripts/perf/disaggregated/gb200_wideep_stress-deepseek-r1-fp4_8k1k_ctx2_gen1_dep32_bs128_eplb288_mtp3_ccb-NIXL.yamltests/unittest/_torch/modules/moe/test_moe_module.py
|
PR_Github #55215 [ run ] triggered by Bot. Commit: |
|
PR_Github #55215 [ run ] completed with state |
2158d4d to
1807c22
Compare
|
/bot run |
|
PR_Github #55219 [ run ] triggered by Bot. Commit: |
|
PR_Github #55219 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #55264 [ run ] triggered by Bot. Commit: |
|
PR_Github #55264 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #55354 [ run ] triggered by Bot. Commit: |
|
PR_Github #55354 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #55369 [ run ] triggered by Bot. Commit: |
|
PR_Github #55369 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #55387 [ run ] triggered by Bot. Commit: |
|
PR_Github #55387 [ run ] completed with state
|
…red EPLB weights and scales before registration, so dynamically migrated expert slots match the device layout. Update EPLB tests and clean up wideep accuracy/stress configs. Signed-off-by: Tingfeng Xian <289617005+nv-xtf@users.noreply.github.com>
Signed-off-by: Tingfeng Xian <289617005+nv-xtf@users.noreply.github.com>
1807c22 to
e610171
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #55395 [ run ] triggered by Bot. Commit: |
|
/bot run --disable-fail-fast |
|
PR_Github #55395 [ run ] completed with state
|
|
PR_Github #55475 [ run ] triggered by Bot. Commit: |
|
PR_Github #55475 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #55514 [ run ] triggered by Bot. Commit: |
|
PR_Github #55514 [ run ] completed with state |
…VIDIA#15538) Signed-off-by: Tingfeng Xian <289617005+nv-xtf@users.noreply.github.com>
4 participants
Description
Dynamic EPLB can migrate MoE experts by copying shared host weights into device expert slots. For CuteDSL NVFP4 MoE, the device
w3/w1weights and scales are transformed into the CuteDSL gate/up interleaved layout after loading, but the shared EPLB weights and scales were not prepared with the same layout before registration.As a result, dynamically migrated expert slots could receive weights in the wrong layout, causing corrupted logits and near-zero DeepSeek-R1 FP4 wideep accuracy.
This PR applies the CuteDSL
w3/w1interleave to shared EPLB weights and scales before registration, so migrated expert slots match the device layout. It also updates EPLB test coverage and cleans up wideep accuracy/stress configs:removing the force-accept speculative-decoding override (
TLLM_SPEC_DECODE_FORCE_NUM_ACCEPTED_TOKENS) from accuracy configs (MTP itself is kept), and disabling accuracy gating in stress-only configs.Test Coverage
test_configurable_moe_multi_gpu_eplb) to use the CuteDSL NVFP4 quant method for CUTEDSL coverage.flexible-extract exact_match: 0.9568strict-match exact_match: 0.9568PR 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.Summary by CodeRabbit
New Features
Tests