[None][fix] Fix Qwen3.5 MoE fallback, GDN alignment, FP8 activation, and draft KV cache - #17120
Open
Wanli-Jiang wants to merge 4 commits into
Open
[None][fix] Fix Qwen3.5 MoE fallback, GDN alignment, FP8 activation, and draft KV cache#17120Wanli-Jiang wants to merge 4 commits into
Wanli-Jiang wants to merge 4 commits into
Conversation
Wanli-Jiang
requested review from
allisonlim-nv,
eopXD,
leslie-fang25,
lowsfer,
schetlur-nv and
xxi-nv
July 31, 2026 09:26
Collaborator
Author
|
/bot run --disable-fail-fast |
Contributor
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 |
Contributor
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
Comment on lines
+233
to
+315
| @skip_unsupported | ||
| def test_fi_mtp_verify_misaligned_ab_slices(): | ||
| """Non-32B-aligned ``a``/``b`` column slices must be realigned. | ||
|
|
||
| ``gdn_mixer._compute_tokenwise_inputs`` splits the fused ``in_proj_ba`` | ||
| projection into ``b = projected_states_ba[:, :num_v_heads_per_tp]`` and | ||
| ``a = projected_states_ba[:, num_v_heads_per_tp:]``. ``a`` therefore starts | ||
| ``num_v_heads_per_tp`` elements in, and that byte offset is not a multiple | ||
| of 32 for many TP splits (in bf16, whenever the per-rank v-head count is not | ||
| a multiple of 16 -- e.g. 128 v-heads over TEP16). The FI kernel then asserts | ||
| ``Misaligned Tensor data on argument``. Attention-DP hides this because | ||
| v-heads are not sharded there. | ||
|
|
||
| The gating inputs are fp32, matching the other tests in this file. ``HV=4`` | ||
| gives the requested 16-byte offset. Two fused buffers exercise each guard: | ||
| the ordinary split misaligns ``a``; shifting the fused buffer misaligns | ||
| ``b`` instead. | ||
|
|
||
| ``.contiguous()`` is not a fix: at ``draft_token_num == 1`` the token dim is | ||
| size 1, so the offset view already reports as contiguous. | ||
| """ | ||
| from tensorrt_llm._torch.modules.fla.fused_sigmoid_gating_recurrent import ( | ||
| _flashinfer_gdn_verify, | ||
| ) | ||
|
|
||
| torch.manual_seed(0) | ||
| dev = "cuda" | ||
| N, T, H, HV, K, V = 2, 3, 4, 4, 128, 128 | ||
| q = (torch.randn(N, T, H, K, device=dev) * 0.1).to(torch.bfloat16) | ||
| k = (torch.randn(N, T, H, K, device=dev) * 0.1).to(torch.bfloat16) | ||
| v = (torch.randn(N, T, HV, V, device=dev) * 0.1).to(torch.bfloat16) | ||
| A_log = torch.empty(HV, device=dev).uniform_(1.0, 16.0).log() | ||
| dt_bias = torch.randn(HV, device=dev) * 0.1 | ||
| state_pool = (torch.randn(N, HV, V, K, device=dev) * 0.1).to(torch.bfloat16) | ||
| idx = torch.arange(N, device=dev, dtype=torch.int32) | ||
|
|
||
| # Same column split gdn_mixer performs on the fused in_proj_ba output: | ||
| # the base is aligned, while a starts 16 bytes into each row. | ||
| ba = torch.randn(N * T, 2 * HV, device=dev) * 0.1 | ||
| b_aligned = ba[:, :HV].view(N, T, HV) | ||
| a_misaligned = ba[:, HV:].view(N, T, HV) | ||
| assert b_aligned.data_ptr() % 32 == 0 | ||
| assert a_misaligned.data_ptr() % 32 == 16 | ||
|
|
||
| # Shift a second fused buffer by 16 bytes so b is misaligned and a is | ||
| # aligned after the additional 16-byte column offset. | ||
| shifted_storage = torch.randn(N * T * 2 * HV + HV, device=dev) * 0.1 | ||
| shifted_ba = shifted_storage[HV:].view(N * T, 2 * HV) | ||
| b_misaligned = shifted_ba[:, :HV].view(N, T, HV) | ||
| a_aligned = shifted_ba[:, HV:].view(N, T, HV) | ||
| assert b_misaligned.data_ptr() % 32 == 16 | ||
| assert a_aligned.data_ptr() % 32 == 0 | ||
|
|
||
| def _run(a_in, b_in): | ||
| return _flashinfer_gdn_verify( | ||
| A_log=A_log, | ||
| a=a_in, | ||
| dt_bias=dt_bias, | ||
| softplus_beta=1.0, | ||
| softplus_threshold=20.0, | ||
| q=q, | ||
| k=k, | ||
| v=v, | ||
| b=b_in, | ||
| initial_state_source=state_pool, | ||
| initial_state_indices=idx, | ||
| intermediate_states_buffer=torch.zeros( | ||
| N, T, HV, V, K, device=dev, dtype=torch.bfloat16 | ||
| ), | ||
| scale=K**-0.5, | ||
| use_qk_l2norm_in_kernel=True, | ||
| ) | ||
|
|
||
| for a, b in [ | ||
| (a_misaligned, b_aligned), | ||
| (a_aligned, b_misaligned), | ||
| ]: | ||
| out_misaligned = _run(a, b) | ||
| out_aligned = _run( | ||
| a.clone(memory_format=torch.contiguous_format), | ||
| b.clone(memory_format=torch.contiguous_format), | ||
| ) | ||
| torch.testing.assert_close(out_misaligned.float(), out_aligned.float()) |
Contributor
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
rg -n -i -C2 \
'test_flashinfer_gdn_verify|test_fi_mtp_verify_misaligned_ab_slices|flashinfer.*gdn' \
tests/integration/test_lists/test-db \
tests/integration/test_lists/qa || trueRepository: NVIDIA/TensorRT-LLM
Length of output: 157
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- matching test-list files ---'
rg --files tests/integration/test_lists | rg '(^|/)(test-db|qa)(/|$)|cbts_touchmap|coverage' | head -200
printf '%s\n' '--- test module references across test lists ---'
rg -n -i -C3 \
'test_flashinfer_gdn_verify|flashinfer.*gdn|mamba.*flashinfer|gdn.*verify' \
tests/integration/test_lists || true
printf '%s\n' '--- test module and nearby test definitions ---'
rg -n -C2 \
'def test_fi_mtp_verify_misaligned_ab_slices|^def test_|^class ' \
tests/unittest/_torch/modules/mamba/test_flashinfer_gdn_verify.py | head -240Repository: NVIDIA/TensorRT-LLM
Length of output: 7779
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- all references to the test module basename ---'
rg -n -i -C2 \
'test_flashinfer_gdn_verify\.py|_torch/modules/mamba|modules/mamba|gdn' \
tests/integration/test_lists/test-db \
tests/integration/test_lists/qa || true
printf '%s\n' '--- test-db entries near relevant GPU coverage ---'
sed -n '40,90p' tests/integration/test_lists/test-db/l0_h100.yml
sed -n '350,405p' tests/integration/test_lists/test-db/l0_b200.yml
printf '%s\n' '--- QA list syntax and possible unittest entries ---'
rg -n -i -C2 \
'unittest|mamba|flashinfer|gdn' \
tests/integration/test_lists/qa | head -240
printf '%s\n' '--- list metadata and database references ---'
rg -n -i \
'cbts_touchmap\.sqlite|coverage report|test-db|qa|pytest' \
tests/integration/test_lists/test-db/README.md \
tests/integration/test_lists/qa/README.md | head -160Repository: NVIDIA/TensorRT-LLM
Length of output: 31915
Add QA test-list coverage.
Test coverage summary:
- Added:
test_fi_mtp_verify_misaligned_ab_slices. - CI coverage: present through
unittest/_torch/modules/mambainl0_h100.yml,l0_b200.yml,l0_b300.yml, andl0_gb300_multi_gpus.yml. - QA coverage: no entry for
test_flashinfer_gdn_verify.py. - Verdict: needs follow-up.
Add unittest/_torch/modules/mamba/test_flashinfer_gdn_verify.py to the applicable QA list, such as tests/integration/test_lists/qa/llm_function_core.txt.
🤖 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/modules/mamba/test_flashinfer_gdn_verify.py` around
lines 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.
Source: Path instructions
Comment on lines
+251
to
+253
| ``.contiguous()`` is not a fix: at ``draft_token_num == 1`` the token dim is | ||
| size 1, so the offset view already reports as contiguous. | ||
| """ |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Cover the contiguous misaligned view.
Lines 251-252 identify the case where .contiguous() does not copy a misaligned view. Line 260 uses N=2, T=3, where the sliced views are non-contiguous and .contiguous() would allocate. A future replacement of clone() with .contiguous() would pass this test but fail the single-request, single-draft-token case.
Run both offset cases with N=1, T=1. Assert that the misaligned view is contiguous and that its data pointer remains 16 bytes off a 32-byte boundary.
Also applies to: 260-260
🤖 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/modules/mamba/test_flashinfer_gdn_verify.py` around
lines 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().
Comment on lines
+85
to
+98
| def test_fp8_block_scale_moe_fallback_tactic_is_explicit_and_deterministic(): | ||
| valid_tactics = [ | ||
| [8, 0], | ||
| [32, 4], | ||
| [16, 0], | ||
| [32, 0], | ||
| ] | ||
|
|
||
| assert _select_explicit_fallback_tactic(valid_tactics) == [32, 0] | ||
|
|
||
| with pytest.raises(RuntimeError, match="no valid fallback tactic"): | ||
| _select_explicit_fallback_tactic([]) | ||
|
|
||
|
|
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
test_name='test_fp8_block_scale_moe_fallback_tactic_is_explicit_and_deterministic'
echo 'Integration-path references:'
rg -n -C 4 \
'fp8_block_scale_moe_runner|choose_one|best_tactic|_select_explicit_fallback_tactic' \
tests/unittest/_torch/modules/moe
echo 'Test-list registration:'
fd --type f --extension yml --extension yaml . tests/integration/test_lists |
while IFS= read -r list_file; do
rg -n -C 2 "$test_name|test_moe_backend.py" "$list_file" || true
doneRepository: NVIDIA/TensorRT-LLM
Length of output: 6151
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo 'Helper and runner definitions:'
rg -n -C 12 \
'def _select_explicit_fallback_tactic|def fp8_block_scale_moe_runner|best_tactic == -1|_select_explicit_fallback_tactic' \
tensorrt_llm tests
echo 'Changed test-file context:'
git diff -- tests/unittest/_torch/modules/moe/test_moe_backend.py
echo 'All test-list references to the changed test file:'
rg -n -C 2 \
'(^|/)(unittest/)?_torch/modules/moe/test_moe_backend\.py|test_fp8_block_scale_moe_fallback_tactic_is_explicit_and_deterministic' \
tests/integration/test_lists
echo 'Relevant list files:'
fd --type f . tests/integration/test_lists | rg '(^|/)(test-db|qa)(/|$)|\.ya?ml$'Repository: NVIDIA/TensorRT-LLM
Length of output: 36844
Add runner-path coverage and register the test.
Test coverage summary — verdict: insufficient.
- Added:
test_fp8_block_scale_moe_fallback_tactic_is_explicit_and_deterministic. - Modified: none.
- Removed: none.
- Covered: helper selection and empty-tactic handling.
- Not covered:
fp8_block_scale_moe_runnerwhenbest_tactic == -1, including fallback tactic delivery tokernel_runner. - Registration: no
test-db/orqa/entry selects this test. Add an explicit entry to the relevant list, such astests/integration/test_lists/test-db/l0_b200.yml.
🤖 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/modules/moe/test_moe_backend.py` around lines 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.
Source: Path instructions
Collaborator
|
PR_Github #63011 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #63011 [ run ] completed with state
|
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #63139 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #63139 [ run ] completed with state |
leslie-fang25
approved these changes
Aug 1, 2026
xxi-nv
reviewed
Aug 2, 2026
|
|
||
| assert _select_explicit_fallback_tactic(valid_tactics) == [32, 0] | ||
|
|
||
| with pytest.raises(RuntimeError, match="no valid fallback tactic"): |
Collaborator
There was a problem hiding this comment.
@sunnyqgg Could you help to rethink about the test of the kernel-bound logic?
It is a general issue; we don't have a well-defined file to do this kind of test for now.
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>
Wanli-Jiang
force-pushed
the
user/williamj/qwen35-bug-fix3
branch
from
August 3, 2026 06:16
ac5e2f3 to
b30713c
Compare
Contributor
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
VALLIS-NERIA
approved these changes
Aug 3, 2026
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #63378 [ run ] triggered by Bot. Commit: |
BowenFu
approved these changes
Aug 3, 2026
Collaborator
|
PR_Github #63378 [ run ] completed with state
|
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #63449 [ run ] triggered by Bot. Commit: |
2 tasks
Collaborator
|
PR_Github #63449 [ run ] completed with state
|
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #63590 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #63590 [ run ] completed with state
|
1 task
nv-guomingz
approved these changes
Aug 4, 2026
sunnyqgg
requested changes
Aug 4, 2026
sunnyqgg
left a comment
Collaborator
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).
1 task
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #63709 [ run ] triggered by Bot. Commit: |
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
8 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.
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.