[None][fix] Fix Qwen3Next MoE expert-quant probe and GDN verify tensor alignment - #17051
Conversation
e6689c6 to
8d567e1
Compare
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (4)
WalkthroughThe change centralizes Qwen3Next expert exclusion detection for quantization and MTP backend selection, adds pattern and layer coverage tests, registers the test, and enforces tensor alignment before the FlashInfer GDN verification kernel. ChangesQwen3Next expert quantization
FlashInfer input alignment
Estimated code review effort: 3 (Moderate) | ~20 minutes Suggested reviewers: Sequence Diagram(s)sequenceDiagram
participant ModelConfig
participant Qwen3NextSparseMoeBlock
participant Qwen3NextMTP
ModelConfig->>Qwen3NextSparseMoeBlock: provide exclusion configuration and layer index
Qwen3NextSparseMoeBlock->>ModelConfig: detect excluded routed experts
ModelConfig->>Qwen3NextSparseMoeBlock: provide kv_cache_quant_algo override
ModelConfig->>Qwen3NextMTP: provide per-layer exclusion result
Qwen3NextMTP->>Qwen3NextMTP: select CUTLASS fallback when experts are excluded
🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (3 passed)
✨ 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.
Actionable comments posted: 1
🧹 Nitpick comments (1)
tensorrt_llm/_torch/modules/fla/fused_sigmoid_gating_recurrent.py (1)
388-396: 🎯 Functional Correctness | 🔵 Trivial | 🏗️ Heavy liftAdd regression coverage for misaligned
a/bslices.Please add or confirm a CUDA/FlashInfer test that exercises 16-byte-offset fused-output views and verifies
_flashinfer_gdn_verifycompletes successfully.🤖 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/fla/fused_sigmoid_gating_recurrent.py` around lines 388 - 396, Add regression coverage for the fused-output views consumed by _flashinfer_gdn_verify, constructing CUDA tensors whose a and b slices have a 16-byte offset and therefore misaligned data pointers. Run the test through the FlashInfer verification path and assert it completes successfully, preserving existing aligned-input coverage if present.
🤖 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.
Inline comments:
In `@tests/unittest/_torch/models/test_qwen3_next_moe_quant.py`:
- Around line 1-183: Add
tests/unittest/_torch/models/test_qwen3_next_moe_quant.py to the appropriate
scheduled test-list files under tests/integration/test_lists/test-db/ and
tests/integration/test_lists/qa/. Ensure the module is included so all listed
tests—especially test_real_checkpoint_shapes and the exclusion-coverage
tests—run in both relevant test suites.
---
Nitpick comments:
In `@tensorrt_llm/_torch/modules/fla/fused_sigmoid_gating_recurrent.py`:
- Around line 388-396: Add regression coverage for the fused-output views
consumed by _flashinfer_gdn_verify, constructing CUDA tensors whose a and b
slices have a 16-byte offset and therefore misaligned data pointers. Run the
test through the FlashInfer verification path and assert it completes
successfully, preserving existing aligned-input coverage if present.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: df97bad3-6a62-45b8-8331-bf5a038ee48d
📒 Files selected for processing (3)
tensorrt_llm/_torch/models/modeling_qwen3_next.pytensorrt_llm/_torch/modules/fla/fused_sigmoid_gating_recurrent.pytests/unittest/_torch/models/test_qwen3_next_moe_quant.py
8d567e1 to
5944cd9
Compare
…from quant Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
…fy kernel Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
5944cd9 to
6b348f7
Compare
|
/bot run --disable-fail-fast |
|
Matcher swap is a clear improvement — One thing worth a look, not blocking: |
|
PR_Github #62730 [ run ] triggered by Bot. Commit: |
|
PR_Github #62730 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #62771 [ run ] triggered by Bot. Commit: |
fredricz-20070104
left a comment
fredricz-20070104
left a comment
There was a problem hiding this comment.
Review summary - CONCERNS
Verdict: The exclusion-probe refactor is clean and well unit-tested, but two production paths ship with no test on the code/arch they change, so I would not merge before those gaps are closed or verified.
Concerns
-
[MAJOR]
tensorrt_llm/_torch/modules/fla/fused_sigmoid_gating_recurrent.py:391- GDN a/b alignment fix untested and arch-specific- What is wrong: the new
if a.data_ptr() % 32 != 0: a = a.clone(...)(and same forb) guards the FlashInfer CuTe-DSL kernel's 32-byte alignment assertion, motivated by the Qwen3.5 TEP16 case (128/16=8, misaligned). - How it fails: the CuTe-DSL verify kernel runs on Blackwell-class GPUs, but the only CI test added (
l0_a10.yml) is a CPU-only Python unit test of the exclusion probe on an A10 (sm86). If the clone does not realign, or another tensor argument is misaligned, the kernel assert fires only on the untested arch and never in pre-merge CI. CodeRabbit flagged the same gap. - Suggested fix: add a CUDA/FlashInfer regression that builds
a/bas 16-byte-offset views of a fusedin_proj_babuffer and runs_flashinfer_gdn_verifyto completion, gated with the appropriate arch marker (e.g.skip_pre_blackwell).
- What is wrong: the new
-
[MAJOR]
tensorrt_llm/_torch/models/modeling_qwen3_next.py:245- regular-layer bf16 override without a backend fallback- What is wrong:
Qwen3NextSparseMoeBlock.__init__now setsexpert_quant_config = QuantConfig(kv_cache_quant_algo=...)(bf16) for any excluded layer, including regular layers 0..N-1 (test_regular_layers_are_covered_too). OnlyQwen3NextMTP.__init__also forcesmoe_backend="CUTLASS". - How it fails: the diff's own comment states DEEPGEMM and WIDEEP do not auto-switch to CUTLASS for a bf16 config. A regular layer whose experts are excluded while
moe_backendis DEEPGEMM/WIDEEP would handcreate_moea bf16 config to a backend that rejects bf16 experts -> create_weights/forward abort, the exact failure the MTP fallback exists to prevent. No test exercisescreate_moebackend selection; the unit test only asserts the probe boolean. - Suggested fix: confirm
create_moeauto-switches DEEPGEMM/WIDEEP to CUTLASS for a bf16 regular layer, or extend the CUTLASS fallback to regular excluded layers as well; add a test that asserts the resolved backend/expert config, not just the probe result.
- What is wrong:
Minor notes (non-blocking)
tensorrt_llm/_torch/models/modeling_qwen3_next.py:803- the MTP comment says only DEEPGEMM/WIDEEP need the CUTLASS switch, but the guard is stillmoe_backend != "CUTLASS", firing for every non-CUTLASS backend. Harmless, but the comment misleads; align it with the code.- Register
test_qwen3_next_moe_quant.pyin theqa/test lists too, not justl0_a10, so scheduled suites pick it up.
QA view
- Test coverage: partial - the new unit test covers
_experts_excluded_from_quantthoroughly, but the GDN alignment clone, the SparseMoeBlock bf16 override reachingcreate_moe, and the MTP CUTLASS fallback are all uncovered. - SM coverage: the GDN fix targets a Blackwell (sm100) CuTe-DSL kernel; the only added test runs CPU-only on A10 (sm86). Real arch coverage gap on the path being fixed. The exclusion probe itself is architecture-independent.
- Test code: clean CPU-only test, but scoped to the probe function only; does not touch its two consumers.
- Test time: small - one lightweight parametrized CPU-only unit file added to l0_a10 pre_merge; no model load or GPU.
- Needs
/qa-verify: yes - arch-specific alignment fix with no Blackwell test, and backend-selection changes validated only at the probe level. QA should run the Qwen3.5 NVFP4/TEP16 MoE path on Blackwell and a DEEPGEMM/WIDEEP excluded-layer config before trusting this.
Does this actually fix the reported issue?
Partial. No specific bug ID is attached. The expert-quant probe half is fixed and well covered. The GDN verify alignment fix and the two probe consumers are plausible but unverified: the alignment path has no test on the arch it targets, and the regular-layer/backend-selection path is only probe-tested. The still-uncovered paths are the FlashInfer verify kernel on Blackwell and create_moe backend resolution for excluded layers.
Possible new issues
- Regular excluded layer on DEEPGEMM/WIDEEP: bf16 config set but backend not forced to CUTLASS -> possible MoE weight-creation abort.
- The clone fix assumes the caching allocator yields 32-byte-aligned storage and that
a/bare the only misaligned kernel arguments; other arguments are not re-checked. - Broadening the override to all excluded layers can change backend selection for a previously-working regular-layer configuration.
What I could not verify
create_moe's actual auto-switch behaviour for bf16 configs across DEEPGEMM/WIDEEP/other backends (not shown in the diff), whether clone() guarantees 32-byte alignment here, and whether any tensor argument other than a/b/initial_state_indices can be misaligned in the verify kernel.
Automated review by NVCortex Lite, run by @fredricz-20070104.
|
PR_Github #62771 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #62895 [ run ] triggered by Bot. Commit: |
|
PR_Github #62895 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
@BowenFu I agree that the regular-layer backend selection and the missing construction-level test are valid gaps. This PR is intended to fix the reproduced Qwen3.5 MTP failure and the FlashInfer GDN alignment failure. The affected checkpoint excludes the MTP experts rather than regular-layer experts, so the reported path is covered by the existing MTP CUTLASS fallback. To keep this fix moving, I would like to merge this PR as-is. I will address the broader regular-layer behavior in a follow-up PR by making the backend selection consistent for all excluded experts and adding tests that exercise I will also add the misaligned |
|
PR_Github #62985 [ run ] triggered by Bot. Commit: |
|
PR_Github #62985 [ run ] completed with state |
8 participants
Dev Engineer Review
_experts_excluded_from_quant(model_config, layer_idx)that checksquant_config.exclude_modulesagainst both:layer_idxoffset (MTP layer index =NUM_HIDDEN_LAYERS).Qwen3NextSparseMoeBlock’s per-layerexpert_quant_configto aQuantConfigderived frommodel_config.quant_config.kv_cache_quant_algowhen that layer’s experts are excluded, andQwen3NextMTP.__init__to choosemoe_backend="CUTLASS"on a per-layer_idxbasis when MTP experts are excluded.Qwen3NextMTPprivate static heuristics used for pattern matching/exclusion decisions (_mtp_pattern_covers_experts,_is_mtp_excluded_from_quant).a/binto 32-byte-aligned contiguous storage whendata_ptr()is misaligned (and preserving existinginitial_state_indiceshandling withdisable_state_update=True).QA Engineer Review
tests/unittest/_torch/models/test_qwen3_next_moe_quant.pytest_experts_covered_by_exclusion(pattern)test_experts_not_covered_by_exclusion(pattern)test_known_exclude_module_shapes(shape, exclude_modules, quant_algo, expected)test_no_exclusions_needs_no_fallback()test_regular_layers_are_covered_too(pattern, layer_idx, expected)test_missing_layer_idx_is_a_noop()tests/integration/test_lists/test-db/l0_a10.yml: addedunittest/_torch/models/test_qwen3_next_moe_quant.pyunderl0_a10pre_merge/pytorch.Description
Test Coverage
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.