[#20365][fix] Use the fused DSpark RMSNorm/RoPE kernels on SM107 by lucifer1004 · Pull Request #20366 · NVIDIA/TensorRT-LLM

lucifer1004 · 2026-10-15T04:23:37Z

Description

Fixes #20365.

is_fused_dspark_rmsnorm_rope_supported() gated on _get_dspark_arch_str(), the SM100/SM103 table of the DSV4 DSpark attention kernel, so on SM107 the MLA-backboned DSpark drafter always fell back to the eager RMSNorm/RoPE chain, although the ops' own error messages list SM107.

The RMSNorm/RoPE kernels are plain SIMT CuTe kernels, so this PR gives them their own architecture set, _DSPARK_RMSNORM_ROPE_SM_VERSIONS = {100, 103, 107}, checked by _is_dspark_rmsnorm_rope_sm_supported(). _get_dspark_arch_str() and every DSV4 DSpark attention gate are unchanged, so the tcgen05 attention kernel stays SM100/SM103-only; its attention-preparation ops (*_cache_write, *_draft_block) are only reached when the attention gate also passes.

Test Coverage

tests/unittest/_torch/speculative/test_dspark_cute_dsl_rmsnorm_rope.py: the module now also runs on SM107, and the new test_fused_dspark_rmsnorm_rope_supported_architectures checks the architecture set and that _get_dspark_arch_str(107) stays None. The file is in l0_b200.yml, so the new test runs there; there is no SM107 test list to add it to.
Validated on SM107 with an equivalent patch: the three drafter call shapes (H = 8 / 16, N = 1 / 3 / 8) match the eager chain bit for bit, and Kimi K3 + DSpark decode is ~0.3 ms/step (~2%) faster. The updated test file passed 21/21 on SM107 (branch diff applied on a build from an earlier main; see the test report comment).
SM100/SM103: behaviour unchanged (same set membership); not re-run.
pre-commit on the changed files: clean.

PR Checklist

PR description clearly explains what and why.
PR follows the TRT-LLM coding guidelines to the best of my knowledge.
Test cases are provided for new code paths.
No API changes, no new dependencies, no CODEOWNERS changes.
Please check this after reviewing the above items as appropriate for this PR.

Dev Engineer Review

The RMSNorm/RoPE support gate now uses a separate SM100/SM103/SM107 set. This enables the fused SIMT CuTe kernels on SM107 without widening the DSV4 attention gate, which remains SM100/SM103-only. The change preserves the distinction between the RMSNorm/RoPE kernels and the tcgen05 attention kernel.

The updated test file has not been run on SM107, and SM100/SM103 behavior was not re-run. GPU validation remains a follow-up.

QA Engineer Review

The test module adds test_fused_dspark_rmsnorm_rope_supported_architectures. It checks support for SM100/SM103/SM107, rejection of SM90/SM101/SM110/SM120, the default architecture check, and that the DSV4 architecture helper returns None for SM107. Existing tests cover invalid inputs, output against a reference, CUDA graph replay, compilation caching, and norm_dim behavior.

The test module is listed in tests/integration/test_lists/test-db/l0_b200.yml. Coverage verdict: needs follow-up because the updated tests have not been run on SM107, and SM100/SM103 behavior was not re-run.

Per-File QA Perspective

tensorrt_llm/_torch/custom_ops/dspark_rmsnorm_rope_custom_op.py: The fused RMSNorm/RoPE support check now accepts SM107 when tensor constraints pass. Verify fused dispatch on SM107 and confirm DSV4 attention remains unavailable there.
tests/unittest/_torch/speculative/test_dspark_cute_dsl_rmsnorm_rope.py: Adds architecture-gate coverage and retains tests for input validation, numerical behavior, graph replay, caching, and norm dimensions. The module is listed in the CI test list tests/integration/test_lists/test-db/l0_b200.yml; SM107 execution remains unverified.

The RMSNorm/RoPE support gate reused the SM100/SM103 table of the DSV4 DSpark attention kernel, so the drafter always took the eager path on SM107. Give the SIMT RMSNorm/RoPE kernels their own architecture set that includes SM107 and leave the attention gate unchanged. Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com>

coderabbitai · 2026-10-15T04:30:48Z

Important

Review skipped

Review was skipped as selected files did not have any reviewable changes.

⚙️ Run configuration

Configuration used: Repository: NVIDIA/TensorRT-LLM/.coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 3f2cb1c6-0a4a-4ea8-af86-25c8be61d38e

📥 Commits

Reviewing files that changed from the base of the PR and between 203f8fa and b3cf32d.

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

🔍 Trigger review

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Repository: NVIDIA/TensorRT-LLM/.coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 1556f570-4208-4138-a4a6-518743a37ce4

📥 Commits

Reviewing files that changed from the base of the PR and between 15f060c and 203f8fa.

📒 Files selected for processing: 2

tensorrt_llm/_torch/custom_ops/dspark_rmsnorm_rope_custom_op.py
tests/unittest/_torch/speculative/test_dspark_cute_dsl_rmsnorm_rope.py

Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 10 remain after this review.

Walkthrough

The fused DSpark RMSNorm/RoPE GPU support gate now accepts SM107 using an architecture set separate from the DSpark attention-kernel map. Tests cover the supported architectures, rejected architectures, and the SM107 architecture-string result.

Changes

RMSNorm/RoPE GPU support

Layer / File(s)	Summary
RMSNorm/RoPE architecture gate `tensorrt_llm/_torch/custom_ops/dspark_rmsnorm_rope_custom_op.py`, `tests/unittest/_torch/speculative/test_dspark_cute_dsl_rmsnorm_rope.py`	The fused-op validator uses a cached support check for SM100, SM103, and SM107 instead of the attention-kernel architecture map. Tests cover supported and rejected architectures and confirm that the DSV4 architecture-string helper returns `None` for SM107.

Priority: ➖ Normal

Estimated code review effort: 2 (Simple) | ~10 minutes

Change: Bug fix · Severity of issue fixed: Medium

Suggested reviewers: bowenfu

Merge Risk: ⚪ Minimal · up to 203f8

The change enables fused RMSNorm/RoPE support on SM107 while keeping the DSV4 attention limit separate. An SM107 test run is not reported, but no actionable merge-blocking risk is established.

Pre-merge checks: 4 · 1

Failed checks: 1 warning

Check name	Status	Explanation	Resolution
Docstring Coverage		Docstring coverage is 66.67% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 6 functions across 2 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

Passed checks: 4 passed

Check name	Status	Explanation
Linked Issues check		Issue `#20365` requires the fused DSpark RMSNorm/RoPE path on SM107 while keeping the DSV4 attention gate separate. The source change adds `_DSPARK_RMSNORM_ROPE_SM_VERSIONS = frozenset((100, 103, 107))`…
Out of Scope Changes check		The changes are limited to the RMSNorm/RoPE support gate and its focused architecture tests. The test also verifies that the separate DSV4 attention gate remains unchanged. No unrelated implementation…
Title check		The title clearly identifies the fix and the affected DSpark RMSNorm/RoPE kernels on SM107. It follows the required ticket/type format.
Description check		The description explains the issue, the architecture-gating change, scope limitations, test coverage, validation results, and checklist status. It is sufficiently complete for review.

✨ 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.}

lucifer1004 · 2026-10-15T15:26:32Z

Test report (branch diff applied on a build from an earlier main, SM107, single node):

tests/unittest/_torch/speculative/test_dspark_cute_dsl_rmsnorm_rope.py with this branch: 21/21 passed.
Same test file against the unpatched build: 11 failed / 10 passed, as expected (SM107 not accepted there).

Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com>

lucifer1004 · 2026-10-19T08:30:24Z

Thanks for the review. Merged current main (0372799), which has not touched either file in this PR, so the diff is unchanged. Starting CI.

lucifer1004 · 2026-10-19T08:30:26Z

/bot run

tensorrt-cicd · 2026-10-19T08:36:26Z

PR_Github #78864 [ run ] triggered by Bot. Commit: b3cf32d Link to invocation

tensorrt-cicd · 2026-10-19T11:59:21Z

PR_Github #78864 [ run ] completed with state SUCCESS. Commit: b3cf32d
/LLM/main/L0_MergeRequest_PR pipeline #64985 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

lucifer1004 · 2026-10-19T12:33:26Z

CI #78864 failed outside this PR:

L40S-PyTorch-2: llmapi/test_llm_api_pytorch_bart.py::test_bart_pytorch_generate_encoder_decoder_end_to_end[bf16-kv-v2-cuda-graph-on-greedy-overlap-bart-large-cnn] timed out in llm.generate ("Test terminated unexpectedly"). It is a BART encoder-decoder test and does not use the DSpark RMSNorm/RoPE op; it passed in the other recent runs.
DGX_H100-PyTorch-3: unittest/_torch/speculative/test_eagle3.py::test_llama_eagle3_long_prompt[False] was cut off when fail-fast aborted the stage; the same test passed on DGX_B200-PyTorch-4 in this run.
The remaining stages were aborted by fail-fast. Build stages passed.

Rerunning.

lucifer1004 · 2026-10-19T12:33:28Z

/bot run

tensorrt-cicd · 2026-10-19T12:41:21Z

PR_Github #78918 [ run ] triggered by Bot. Commit: b3cf32d Link to invocation

tensorrt-cicd · 2026-10-19T15:24:36Z

PR_Github #78918 [ run ] completed with state SUCCESS. Commit: b3cf32d
/LLM/main/L0_MergeRequest_PR pipeline #65017 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

lucifer1004 · 2026-10-19T15:29:58Z

CI #78918 failed outside this PR: CPU-Generic-arm-1 timed out in unittest/llmapi/test_mpi_session.py::test_remote_mpi_session[submit_sync] (the mpirun launch hit its 60 s timeout, exit status 124), and fail-fast then aborted the x86_64 stages. This is the same arm-only MPI session timeout seen on other PRs today (#20368, #20374); this PR does not touch the MPI session code. Rerunning.

lucifer1004 · 2026-10-19T15:29:59Z

/bot run

tensorrt-cicd · 2026-10-19T15:52:28Z

PR_Github #78964 [ run ] triggered by Bot. Commit: b3cf32d Link to invocation

tensorrt-cicd · 2026-10-19T20:24:02Z

PR_Github #78964 [ run ] completed with state SUCCESS. Commit: b3cf32d
/LLM/main/L0_MergeRequest_PR pipeline #65056 completed with status: 'SUCCESS'

CI Report

Link to invocation

lucifer1004 · 2026-10-19T20:30:18Z

@NVIDIA/trt-llm-runtime-devs could one of you take a look when you have a moment? CI is green and dc3671 has approved; the remaining requirement is a code-owner review for tensorrt_llm/_torch/custom_ops/dspark_rmsnorm_rope_custom_op.py and its test under tests/unittest/_torch/speculative/. Thanks!

github-actions · 2026-10-20T01:09:13Z

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.

lucifer1004 requested a review from a team as a code owner October 15, 2026 04:23

lucifer1004 requested review from tongyuantongyu and zhaoyuanh-nvidia October 15, 2026 04:23

github-actions Bot assigned lucifer1004 Oct 15, 2026

lucifer1004 mentioned this pull request Oct 19, 2026

[#20551][perf] MLA DSpark: run each drafter layer's kv path beside its q path #20552

Draft

5 tasks done

dc3671 self-requested a review October 19, 2026 08:17

dc3671 approved these changes Oct 19, 2026

View reviewed changes

Merge main into upstream/dspark-rmsnorm-rope-sm107

b3cf32d

Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com>

lori-ren approved these changes Oct 20, 2026

View reviewed changes

trtllm-agent added the ci: full pre-merge approved label Oct 20, 2026

lori-ren merged commit 083447c into NVIDIA:main Oct 20, 2026
12 checks passed

Conversation

lucifer1004 commented Oct 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Test Coverage

PR Checklist

Dev Engineer Review

QA Engineer Review

Per-File QA Perspective

Uh oh!

coderabbitai Bot commented Oct 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Review skipped

Walkthrough

Changes

Uh oh!

lucifer1004 commented Oct 15, 2026

Uh oh!

lucifer1004 commented Oct 19, 2026

Uh oh!

lucifer1004 commented Oct 19, 2026

Uh oh!

tensorrt-cicd commented Oct 19, 2026

Uh oh!

tensorrt-cicd commented Oct 19, 2026

Uh oh!

lucifer1004 commented Oct 19, 2026

Uh oh!

lucifer1004 commented Oct 19, 2026

Uh oh!

tensorrt-cicd commented Oct 19, 2026

Uh oh!

tensorrt-cicd commented Oct 19, 2026

Uh oh!

lucifer1004 commented Oct 19, 2026

Uh oh!

lucifer1004 commented Oct 19, 2026

Uh oh!

tensorrt-cicd commented Oct 19, 2026

Uh oh!

tensorrt-cicd commented Oct 19, 2026

Uh oh!

lucifer1004 commented Oct 19, 2026

Uh oh!

github-actions Bot commented Oct 20, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

lucifer1004 commented Oct 15, 2026 •

edited

Loading

coderabbitai Bot commented Oct 15, 2026 •

edited

Loading