[None][fix] Fix Qwen3.5 MoE fallback, GDN alignment, FP8 activation, and draft KV cache - #17120
Conversation
|
/bot run --disable-fail-fast |
WalkthroughThe change adds a persistent DeepSeek FP8 activation kernel with runtime dispatch, deterministic FP8 MoE tactic fallback, per-layer Qwen3Next MoE configuration, draft KV-cache pool-ratio normalization, and FlashInfer GDN alignment coverage. ChangesFP8 MoE execution
Qwen3Next MoE configuration
One-model draft KV cache
FlashInfer GDN alignment
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant MoERunner
participant ActivationLauncher
participant DeepSeekKernel
MoERunner->>ActivationLauncher: pass numExperts and tileTokensDim
ActivationLauncher->>ActivationLauncher: check layout eligibility
ActivationLauncher->>DeepSeekKernel: launch selected activation kernel
DeepSeekKernel-->>ActivationLauncher: write scaled FP8 outputs
Suggested reviewers: 🚥 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: 4
🧹 Nitpick comments (3)
tests/unittest/_torch/models/test_qwen3_next_moe_quant.py (1)
169-223: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd type annotations to the new helper functions.
_build_moe_blockand nested_capturehave no parameter or return annotations. Add precise annotations for the test inputs, captured values, and return types.As per coding guidelines, “Annotate every function, use
Nonefor non-returning functions, avoidAnyand unnecessary type ignores.”🤖 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/unittest/_torch/models/test_qwen3_next_moe_quant.py` around lines 169 - 223, Add precise type annotations to the test helper `_build_moe_block` for all parameters and its captured-result return value, and annotate nested `_capture` parameters and its non-returning behavior with `None`. Use concrete existing types for backend, module exclusions, layer index, quantization configuration, and the captured dictionary values; avoid `Any` and unnecessary type ignores.Source: Coding guidelines
cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/DevKernel.cu (2)
270-277: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winReplace the literal 8 with a named constant.
innerDim % 8 == 0hides the actual requirement. The packed path needs each row start to be 4-byte aligned, so the derived condition is onkDsActEltsPerThread. Name the value so the constraint stays tied to the packing width.♻️ Proposed refactor
+constexpr int kDsActInnerDimAlignment = 2 * kDsActEltsPerThread; + constexpr bool shouldUsePermutedActivation( int outputDim, int innerDim, int numTokens, int topK, int numExperts, int tileTokensDim) { - bool const layoutEligible = outputDim >= kDsActEltsPerSf && outputDim % kDsActEltsPerSf == 0 && innerDim % 8 == 0; + bool const layoutEligible = outputDim >= kDsActEltsPerSf && outputDim % kDsActEltsPerSf == 0 + && innerDim % kDsActInnerDimAlignment == 0;As per coding guidelines: "Avoid magic literals except
0,nullptr,true, andfalse; initialize named constants instead, usingk-prefixed camelCase names."🤖 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 `@cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/DevKernel.cu` around lines 270 - 277, Update shouldUsePermutedActivation to replace the literal 8 in the innerDim alignment check with a k-prefixed named constant representing the required kDsActEltsPerThread packing width, and use that constant in the modulo condition.Source: Coding guidelines
601-613: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRename the shadowing
numCtas, and note thatmaxTasksis not an upper bound.Two points on the grid sizing:
- Line 612 declares
numCtas, which shadows thenumCtascomputed at line 559 and still used by the else branch. Use a distinct name.maxTasksusesnumTokens * topK, buttotalNumPaddedTokenscan exceed that value, because each local expert adds up totileTokensDim - 1padding rows. The grid-stride loop still covers every task, so this is not a correctness defect, but the cap can launch fewer CTAs than one wave of real work. State that in the comment so the bound is not read as exact.♻️ Proposed refactor
- int const numCtas = static_cast<int>(std::min<int64_t>(ctasForAllTasks, int64_t{numSms} * 32)); - dim3 const permutedGrid(std::max(numCtas, 1), 1, 1); + // Note: ctasForAllTasks is an estimate, not an upper bound; per-expert + // tile padding can push totalNumPaddedTokens above numTokens * topK. + // The grid-stride loop absorbs the difference. + int const numPermutedCtas = static_cast<int>(std::min<int64_t>(ctasForAllTasks, int64_t{numSms} * 32)); + dim3 const permutedGrid(std::max(numPermutedCtas, 1), 1, 1);🤖 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 `@cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/DevKernel.cu` around lines 601 - 613, Rename the inner `numCtas` in the `usePermuted` branch to avoid shadowing the outer `numCtas` used by the alternate path. Update the adjacent `maxTasks` comment to state that `numTokens * topK` is only an estimate and may be below `totalNumPaddedTokens` because expert padding adds rows; retain the grid-stride coverage and existing cap behavior.
🤖 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 226-266: The existing tests only validate
_build_moe_block/create_moe and do not exercise Qwen3NextMTP.__init__. Add a
parametrized regression test that directly constructs Qwen3NextMTP with an
excluded MTP experts layer for both enable_attention_dp=False and
enable_attention_dp=True, verifying the constructor completes and preserves the
expected excluded-layer BF16/CUTLASS behavior.
In `@tests/unittest/_torch/modules/mamba/test_flashinfer_gdn_verify.py`:
- Around line 251-253: Update the contiguous-view test around the offset cases
to use N=1 and T=1, covering both misaligned views. Assert that each misaligned
view reports contiguous while its data pointer remains 16 bytes off a 32-byte
boundary, ensuring the test distinguishes clone() from .contiguous().
- Around line 233-315: Add
tests/unittest/_torch/modules/mamba/test_flashinfer_gdn_verify.py to the
applicable QA test-list file, such as llm_function_core.txt, so
test_fi_mtp_verify_misaligned_ab_slices receives QA coverage. Preserve the
existing CI test-list entries.
In `@tests/unittest/_torch/modules/moe/test_moe_backend.py`:
- Around line 85-98: Add coverage for the best_tactic == -1 branch in
fp8_block_scale_moe_runner, verifying the selected fallback tactic is passed to
kernel_runner, while retaining the existing helper tests. Register the new or
updated test explicitly in the appropriate test-db/qa list, such as l0_b200.yml,
so the test runner selects it.
---
Nitpick comments:
In `@cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/DevKernel.cu`:
- Around line 270-277: Update shouldUsePermutedActivation to replace the literal
8 in the innerDim alignment check with a k-prefixed named constant representing
the required kDsActEltsPerThread packing width, and use that constant in the
modulo condition.
- Around line 601-613: Rename the inner `numCtas` in the `usePermuted` branch to
avoid shadowing the outer `numCtas` used by the alternate path. Update the
adjacent `maxTasks` comment to state that `numTokens * topK` is only an estimate
and may be below `totalNumPaddedTokens` because expert padding adds rows; retain
the grid-stride coverage and existing cap behavior.
In `@tests/unittest/_torch/models/test_qwen3_next_moe_quant.py`:
- Around line 169-223: Add precise type annotations to the test helper
`_build_moe_block` for all parameters and its captured-result return value, and
annotate nested `_capture` parameters and its non-returning behavior with
`None`. Use concrete existing types for backend, module exclusions, layer index,
quantization configuration, and the captured dictionary values; avoid `Any` and
unnecessary type ignores.
🪄 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: e5c23c0c-2bbc-4458-aace-b4feb6746204
📒 Files selected for processing (11)
cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/DevKernel.cucpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/DevKernel.hcpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.cucpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.htensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.pytensorrt_llm/_torch/models/modeling_qwen3_next.pytensorrt_llm/_torch/pyexecutor/_util.pytests/unittest/_torch/executor/test_kv_cache_estimation.pytests/unittest/_torch/models/test_qwen3_next_moe_quant.pytests/unittest/_torch/modules/mamba/test_flashinfer_gdn_verify.pytests/unittest/_torch/modules/moe/test_moe_backend.py
|
PR_Github #63011 [ run ] triggered by Bot. Commit: |
|
PR_Github #63011 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #63139 [ run ] triggered by Bot. Commit: |
|
PR_Github #63139 [ run ] completed with state |
Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
ac5e2f3 to
b30713c
Compare
There was a problem hiding this comment.
🧹 Nitpick comments (3)
cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/DevKernel.cu (1)
600-612: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRename the inner
numCtasto avoid shadowing.Line 560 declares
auto numCtasin the same function scope. Line 608 declaresint const numCtasinside the new branch. The two values mean different things. Rename the inner one, for example tonumPermutedCtas.♻️ Proposed rename
- int const numCtas = static_cast<int>(std::min<int64_t>(ctasForAllTasks, int64_t{numSms} * 32)); - dim3 const permutedGrid(std::max(numCtas, 1), 1, 1); + int const numPermutedCtas = static_cast<int>(std::min<int64_t>(ctasForAllTasks, int64_t{numSms} * 32)); + dim3 const permutedGrid(std::max(numPermutedCtas, 1), 1, 1);🤖 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 `@cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/DevKernel.cu` around lines 600 - 612, Rename the branch-local `numCtas` used to construct `permutedGrid` in the `shouldUsePermutedActivation` path to a distinct name such as `numPermutedCtas`, and update the corresponding grid initialization while leaving the outer `numCtas` unchanged.tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py (2)
70-71: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUse built-in generic annotations.
Replace
List[List[int]]andList[int]withlist[list[int]]andlist[int]in the new annotations. This follows the repository’s Python 3.10 typing style.As per coding guidelines, prefer built-in generic types and the
|union syntax. Based on learnings, this repository supports Python 3.10+ features such as PEP 585 generics.Proposed annotation update
def _select_explicit_fallback_tactic( - valid_tactics: List[List[int]]) -> List[int]: + valid_tactics: list[list[int]]) -> list[int]: ... - def get_fallback_tactic(self, hidden_size: int, - num_tokens: int) -> List[int]: + def get_fallback_tactic(self, hidden_size: int, + num_tokens: int) -> list[int]:Also applies to: 914-915
🤖 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/custom_ops/trtllm_gen_custom_ops.py` around lines 70 - 71, Update the annotations on _select_explicit_fallback_tactic and the additionally affected declarations around the referenced locations to use Python 3.10 built-in generics: replace List[List[int]] with list[list[int]] and List[int] with list[int].Sources: Coding guidelines, Learnings
861-861: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winMark the fallback cache as a private typed class variable.
Ruff reports RUF012 for the mutable class attribute at Line 861. Keep the cache shared, but annotate it with
ClassVar[...], rename it to_fallback_tactic_dict, and update the references at Lines 925 and 930. AddClassVarto the existingtypingimport if needed.As per coding guidelines, non-public names must use a leading underscore.
Proposed cache declaration and reference update
class FP8BlockScaleMoERunner(TunableRunner): - fallback_tactic_dict = dict() + _fallback_tactic_dict: ClassVar[ + dict[tuple[int, int, int, int, int], tuple[int, ...]] + ] = {} ... - tactic = FP8BlockScaleMoERunner.fallback_tactic_dict.get(key) + tactic = FP8BlockScaleMoERunner._fallback_tactic_dict.get(key) ... - FP8BlockScaleMoERunner.fallback_tactic_dict[key] = tactic + FP8BlockScaleMoERunner._fallback_tactic_dict[key] = tacticAlso applies to: 925-930
🤖 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/custom_ops/trtllm_gen_custom_ops.py` at line 861, Update the shared cache declaration to a private, typed class variable named _fallback_tactic_dict using ClassVar, adding ClassVar to the existing typing imports if necessary. Replace all references in the surrounding cache logic, including the usages near lines 925 and 930, while preserving the cache’s shared behavior.Sources: Coding guidelines, Linters/SAST tools
🤖 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 `@cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/DevKernel.cu`:
- Around line 600-612: Rename the branch-local `numCtas` used to construct
`permutedGrid` in the `shouldUsePermutedActivation` path to a distinct name such
as `numPermutedCtas`, and update the corresponding grid initialization while
leaving the outer `numCtas` unchanged.
In `@tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py`:
- Around line 70-71: Update the annotations on _select_explicit_fallback_tactic
and the additionally affected declarations around the referenced locations to
use Python 3.10 built-in generics: replace List[List[int]] with list[list[int]]
and List[int] with list[int].
- Line 861: Update the shared cache declaration to a private, typed class
variable named _fallback_tactic_dict using ClassVar, adding ClassVar to the
existing typing imports if necessary. Replace all references in the surrounding
cache logic, including the usages near lines 925 and 930, while preserving the
cache’s shared behavior.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 4d3ecbc4-b481-4796-87d5-19ef86fdfaa8
📒 Files selected for processing (13)
cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/DevKernel.cucpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/DevKernel.hcpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.cucpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.hcpp/tests/unit_tests/kernels/CMakeLists.txtcpp/tests/unit_tests/kernels/blockScaleMoeActivationTest.cutensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.pytensorrt_llm/_torch/models/modeling_qwen3_next.pytensorrt_llm/_torch/pyexecutor/_util.pytests/unittest/_torch/executor/test_kv_cache_estimation.pytests/unittest/_torch/models/test_qwen3_next_moe_quant.pytests/unittest/_torch/modules/mamba/test_flashinfer_gdn_verify.pytests/unittest/_torch/modules/moe/test_moe_backend.py
🚧 Files skipped from review as they are similar to previous changes (9)
- tests/unittest/_torch/modules/moe/test_moe_backend.py
- cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.h
- tests/unittest/_torch/executor/test_kv_cache_estimation.py
- cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.cu
- tensorrt_llm/_torch/models/modeling_qwen3_next.py
- tests/unittest/_torch/modules/mamba/test_flashinfer_gdn_verify.py
- tensorrt_llm/_torch/pyexecutor/_util.py
- tests/unittest/_torch/models/test_qwen3_next_moe_quant.py
- cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/DevKernel.h
|
/bot run --disable-fail-fast |
|
PR_Github #63590 [ run ] triggered by Bot. Commit: |
|
PR_Github #63590 [ run ] completed with state
|
sunnyqgg
left a comment
sunnyqgg
left a comment
There was a problem hiding this comment.
trtllm-gen already supports the fusion of dsFp8 + swiGlu. Please:
- trtllm-gen side: Export a set of dsFp8 + swiGlu batched GEMM cubins for the actual DeepSeek model shapes (hidden/intermediate sizes and tile
configurations). - TensorRT-LLM side: Update the vendored KernelMetaInfo.h, change runner.cu:414 to fusedAct = true (or decide based on actType), and remove the
standalone activation kernel call path; the autotuner's config space will pick up the new cubins automatically. - Accuracy validation + performance comparison (expected gain: the activation kernel's entire global-memory round trip is eliminated, at a slightly higher GEMM1 epilogue cost).
|
/bot run --disable-fail-fast |
|
PR_Github #63709 [ run ] triggered by Bot. Commit: |
|
PR_Github #63709 [ run ] completed with state
|
|
/bot help |
GitHub Bot Help
Provide a user friendly way for developers to interact with a Jenkins server. Run See details below for each supported subcommand. Details
Launch build/test pipelines. All previously running jobs will be killed.
kill
Kill all running builds associated with pull request. skip
Skip testing for latest commit on pull request. reuse-pipeline
Reuse a previous pipeline to validate current commit. This action will also kill all currently running builds associated with the pull request. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break. |
|
/bot run --only-multi-gpu-test --disable-fail-fast |
|
PR_Github #63887 [ run ] triggered by Bot. Commit: |
sunnyqgg
left a comment
sunnyqgg
left a comment
There was a problem hiding this comment.
We should use trtllm-gen moe kernel directly, since the time is tight we can use the current implementation temporarily
|
PR_Github #63887 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #63948 [ run ] triggered by Bot. Commit: |
|
PR_Github #63948 [ run ] completed with state |
tongyuantongyu
left a comment
tongyuantongyu
left a comment
There was a problem hiding this comment.
Approve for changes to runtime-owned files.
test_qwen3_next_moe_quant.py landed on main (NVIDIA#17120) after this branch was cut, and parametrizes its two backend-resolution probes over WIDEEP. test_unexcluded_layer_keeps_configured_backend_and_layer_quant_config reaches resolve_moe_cls() -> get_moe_cls() with moe_backend=WIDEEP, which this branch turns into a ValueError, so that parametrization now fails. test_excluded_layer_builds_bf16_on_cutlass only passes by accident: Qwen3NextSparseMoeBlock rewrites moe_backend to CUTLASS before the call. Remove the WIDEEP entries from both parametrize lists, matching how the rest of this branch handles WIDEEP in unit tests. The other backends keep full coverage of both probes. Signed-off-by: xxi <xxi@nvidia.com>
test_qwen3_next_moe_quant.py landed on main (NVIDIA#17120) after this branch was cut, and parametrizes its two backend-resolution probes over WIDEEP. test_unexcluded_layer_keeps_configured_backend_and_layer_quant_config reaches resolve_moe_cls() -> get_moe_cls() with moe_backend=WIDEEP, which this branch turns into a ValueError, so that parametrization now fails. test_excluded_layer_builds_bf16_on_cutlass only passes by accident: Qwen3NextSparseMoeBlock rewrites moe_backend to CUTLASS before the call. Remove the WIDEEP entries from both parametrize lists, matching how the rest of this branch handles WIDEEP in unit tests. The other backends keep full coverage of both probes. Signed-off-by: xxi <xxi@nvidia.com>
10 participants
Dev Engineer Review
[1.0].blockScaleMoeActivationTestCUDA target.QA Engineer Review
Added test coverage:
test_kv_cache_estimation.pytest_qwen3_next_moe_quant.pytest_excluded_layer_builds_bf16_on_cutlasstest_unexcluded_layer_keeps_configured_backend_and_layer_quant_configtest_flashinfer_gdn_verify.pytest_fi_mtp_verify_misaligned_ab_slicestest_moe_backend.pytest_fp8_block_scale_moe_fallback_tactic_is_explicit_and_deterministicblockScaleMoeActivationTest.cuThe test functions are not registered in the provided
tests/integration/test_lists/,test-db/, orqa/changes.Verdict: needs follow-up.
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.