[https://nvbugs/6473374][fix] Hoist the isinstance check out of the helper at the single call site in… by trtllm-agent · Pull Request #16637 · NVIDIA/TensorRT-LLM

trtllm-agent · 2026-07-20T20:51:33Z

Summary

Root cause: On SM120 the fused NVFP4 kernel is disabled by RMSNorm's SM guard so no Fp4QuantizedTensor is ever produced, but _slice_hidden_states_to_num_tokens's outlined function-call boundary in MLA.forward_impl still reshapes the traced FX graph, shifting multi-stream sync-event insertion around the piecewise CUDA graph regions and producing an async IMA at spec sampler event synchronize.
Fix: Hoist the isinstance check out of the helper at the single call site in MLA.forward_impl so the plain-tensor path takes an inline hidden_states[:num_tokens, ...] (restoring the pre-501777ac89 FX-node shape) while Fp4QuantizedTensor callers still go through the swizzled-SF-aware helper; also remove the now-fixed waiver.
Automated fix generated by repair-bot

Test plan

Verify fix on the same GPU type as the original failure
Check for regressions in related tests

Links

Bug: https://nvbugs/6473374

Dev Engineer Review

MLA.forward_impl now uses an inline slice for plain tensors.
Fp4QuantizedTensor inputs still use _slice_hidden_states_to_num_tokens.
This preserves FP4 scale-factor handling and avoids the FX graph and synchronization change that caused the SM120 asynchronous IMA.
The change is limited to the sole helper call site and does not alter public APIs.
tests/integration/test_lists/waives.txt removes eight RTX PRO 6000 Blackwell waiver entries.
The waiver file changes use the existing format and do not add duplicates or unrelated scope.

QA Engineer Review

No test-db/ or qa/ files were modified.
Removed waivers cover one disaggregated guided-decoding case, five DeepSeekV3Lite BF16 mtp_nextn=2 cases, and two DeepSeekV3Lite NVFP4 cases.
The requested DeepSeekV3Lite NVFP4 accuracy test and associated LLM function QA test completed successfully.
Verdict: needs follow-up because CBTS coverage data is unavailable.

coderabbitai · 2026-07-20T20:53:49Z

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 1f296a70-6b5a-454b-b320-598c7def0db2

📥 Commits

Reviewing files that changed from the base of the PR and between 8e588da and c61dacb.

📒 Files selected for processing (2)

tensorrt_llm/_torch/modules/mla.py
tests/integration/test_lists/waives.txt

💤 Files with no reviewable changes (1)

tests/integration/test_lists/waives.txt

🚧 Files skipped from review as they are similar to previous changes (1)

tensorrt_llm/_torch/modules/mla.py

Walkthrough

Changes

MLA slicing behavior

Layer / File(s)	Summary
Conditional hidden-state slicing and validation waiver `tensorrt_llm/_torch/modules/mla.py`, `tests/integration/test_lists/waives.txt`	`MLA.forward_impl` uses specialized slicing for `Fp4QuantizedTensor`, directly slices other tensors, and removes eight related RTX PRO 6000 integration-test waivers.

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

Possibly related PRs

NVIDIA/TensorRT-LLM#17155: Removes related DeepSeekV3Lite waiver entries.
NVIDIA/TensorRT-LLM#17191: Addresses Fp4QuantizedTensor handling in a Torch execution path.
NVIDIA/TensorRT-LLM#17294: Removes integration-test waiver entries.

Suggested labels: ci: full pre-merge approved

Suggested reviewers: kris1025, brnguyen2, tburt-nv, pengbowang-nv

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly identifies the bug fix: hoisting the isinstance check at the single helper call site.
Description check	✅ Passed	The description explains the root cause, fix, affected behavior, test plan, and bug reference, although it uses alternate section headings.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

_{Comment @coderabbitai help to get the list of available commands.}

xinhe-nv · 2026-07-21T01:57:37Z

/bot run --only-qa-verify test accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4[moe_backend=CUTLASS-mtp_nextn=2-fp8kv=False-attention_dp=False-cuda_graph=False-overlap_scheduler=False-torch_compile=True]

tensorrt-cicd · 2026-07-21T02:03:17Z

PR_Github #60527 [ run ] triggered by Bot. Commit: 422d191 Link to invocation

tensorrt-cicd · 2026-07-21T02:43:49Z

PR_Github #60527 [ run ] completed with state SUCCESS. Commit: 422d191
LLM_FUNCTION_AUTO_V2C #291 completed with status: 'SUCCESS'
QA verify test: accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4[moe_backend=CUTLASS-mtp_nextn=2-fp8kv=False-attention_dp=False-cuda_graph=False-overlap_scheduler=False-torch_compile=True] (NVBug 6473374, branch repair-bot-bug6473374, fork tensorrt-cicd, dry_run_close=true)

Link to invocation

BowenFu

LGTM. The hoisted plain-tensor path (mla.py:1422) is byte-identical to the helper's non-Fp4 branch (mla.py:95); FP4 callers still route through the helper, and the change only moves the FX-graph call boundary. Waiver change is a removal. Codex-confirmed no other executable callers of the helper.

liji-nv · 2026-07-27T07:57:14Z

+        # boundary otherwise shifts multi-stream sync-event insertion around the
+        # piecewise CUDA graph regions, producing an async CUDA IMA at sampler
+        # event synchronize on SM120 + torch_compile + MTP=2 (nvbugs/6473374).
+        if isinstance(hidden_states, Fp4QuantizedTensor):


The explanation does not make sense. The _slice_hidden_states_to_num_tokens already shortcut the non-FP4 tensor path. And the forward_impl is inside the MLA custom op and not visible for fx graph.

ZhanruiSunCh

LGTM for infra part.

…forward_impl slice helper On SM120 (RTX PRO 6000 Blackwell Server Edition), the fused NVFP4 kernel is disabled by RMSNorm's SM guard and no Fp4QuantizedTensor is ever produced. But MLA.forward_impl still calls _slice_hidden_states_to_num_tokens, whose outlined function-call boundary shifts multi-stream sync-event insertion around piecewise CUDA graph regions under torch.compile (enable_piecewise_cuda_graph=True, max_num_streams=3), producing an async CUDA IMA at spec sampler event synchronize on the mtp_nextn=2 + torch_compile=True config. Hoist the isinstance check out of the helper so the plain-tensor path takes an inline slice (matching the pre-501777ac89 FX-node shape) while the Fp4QuantizedTensor path still goes through the swizzled-SF-aware helper. No behavior change for the fused NVFP4 boundary fold on SM 10.x. Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>

Signed-off-by: handongl <handongl@nvidia.com>

coderabbitai · 2026-08-06T14:01:53Z

Note

GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer.

trtllm-agent requested review from a team as code owners July 20, 2026 20:51

trtllm-agent assigned JunyiXu-nv Jul 20, 2026

trtllm-agent requested review from BowenFu, LarryXFly, ZhanruiSunCh, brnguyen2, kris1025, pengbowang-nv, tburt-nv and xinhe-nv July 20, 2026 20:51

github-actions Bot assigned trtllm-agent Jul 20, 2026

xinhe-nv approved these changes Jul 21, 2026

View reviewed changes

BowenFu approved these changes Jul 21, 2026

View reviewed changes

liji-nv requested changes Jul 27, 2026

View reviewed changes

trtllm-agent force-pushed the repair-bot-bug6473374 branch from 422d191 to c9a9257 Compare July 28, 2026 06:35

ZhanruiSunCh approved these changes Jul 28, 2026

View reviewed changes

trtllm-agent force-pushed the repair-bot-bug6473374 branch from c9a9257 to c5638a1 Compare July 28, 2026 14:43

trtllm-agent requested a review from a team as a code owner July 28, 2026 14:43

LarryXFly approved these changes Jul 30, 2026

View reviewed changes

trtllm-agent force-pushed the repair-bot-bug6473374 branch from c5638a1 to 3098f1d Compare August 3, 2026 10:41

trtllm-agent and others added 2 commits August 6, 2026 06:58

[nvbugs/6473374][chore] Remove stale waiver after fix

c61dacb

Signed-off-by: handongl <handongl@nvidia.com>

trtllm-agent force-pushed the repair-bot-bug6473374 branch from 3098f1d to c61dacb Compare August 6, 2026 14:01

coderabbitai Bot mentioned this pull request Aug 6, 2026

fix: guard PRE_MLP NVFP4 fusion for dense layers #17191

Open

Navigation Menu

Uh oh!

[https://nvbugs/6473374][fix] Hoist the isinstance check out of the helper at the single call site in… - #16637

[https://nvbugs/6473374][fix] Hoist the isinstance check out of the helper at the single call site in…#16637
trtllm-agent wants to merge 2 commits into
NVIDIA:mainfrom
tensorrt-cicd:repair-bot-bug6473374

trtllm-agent commented Jul 20, 2026 •

edited by coderabbitai Bot

Loading

Uh oh!

coderabbitai Bot commented Jul 20, 2026 •

edited

Loading

Uh oh!

xinhe-nv commented Jul 21, 2026

Uh oh!

tensorrt-cicd commented Jul 21, 2026

Uh oh!

tensorrt-cicd commented Jul 21, 2026

Uh oh!

BowenFu left a comment

Uh oh!

liji-nv Jul 27, 2026

Uh oh!

ZhanruiSunCh left a comment

Uh oh!

coderabbitai Bot commented Aug 6, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

9 participants

Uh oh!

Conversation

trtllm-agent commented Jul 20, 2026 • edited by coderabbitai Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Test plan

Links

Dev Engineer Review

QA Engineer Review

Uh oh!

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

Uh oh!

Walkthrough

Changes

Uh oh!

xinhe-nv commented Jul 21, 2026

Uh oh!

tensorrt-cicd commented Jul 21, 2026

Uh oh!

tensorrt-cicd commented Jul 21, 2026

Uh oh!

BowenFu left a comment

Choose a reason for hiding this comment

Uh oh!

liji-nv Jul 27, 2026

Choose a reason for hiding this comment

Uh oh!

ZhanruiSunCh left a comment

Choose a reason for hiding this comment

Uh oh!

coderabbitai Bot commented Aug 6, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

9 participants

trtllm-agent commented Jul 20, 2026 •

edited by coderabbitai Bot

Loading

coderabbitai Bot commented Jul 20, 2026 •

edited

Loading