[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 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 changes use the existing format and do not add duplicates or unrelated scope.

QA Engineer Review

No test-db/ or qa/ files were modified.
The 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: 269bf22e-8f8e-4a1e-8786-0ed7c2ca71f8

📥 Commits

Reviewing files that changed from the base of the PR and between f26152e and 8781a2d.

📒 Files selected for processing (1)

tests/integration/test_lists/waives.txt

💤 Files with no reviewable changes (1)

tests/integration/test_lists/waives.txt

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 plain 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#16940: Related to MLA slicing and DeepSeek-V4 FlashMLA/FP4 handling.

Suggested reviewers: kris1025, tburt-nv

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly identifies the bug fix and the main implementation change: hoisting the type check at the single call site.
Description check	✅ Passed	The description explains the root cause, fix, affected behavior, test plan, and bug link, so it is sufficiently complete.
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.

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.

brnguyen2

The code change is a no-op semantically — the helper's non-FP4 arm is already hidden_states[:num_tokens, ...], and Dynamo inlines plain Python functions, so "the function-call boundary changes FX node identity" needs evidence (a before/after graph dump or event ordering) rather than assertion. More concretely, the waiver removals are broader than the stated mechanism: several removed entries are torch_compile=False (and one cuda_graph=False) configs, which a torch.compile/piecewise-graph node-ordering explanation cannot account for. Please either show those configs passing on RTX PRO 6000, or keep their waivers and narrow this to the torch_compile=True cases.

brnguyen2 · 2026-08-08T19:37:56Z

@@ -298,24 +298,16 @@ full:RTX_6000D/accuracy/test_llm_api_pytorch.py::TestQwen3_5_35B_A3B::test_fp8[e
 full:RTX_6000D/accuracy/test_llm_api_pytorch.py::TestQwen3_5_4B::test_dflash SKIP (https://nvbugs/6273850)
 full:RTX_6000D/accuracy/test_llm_api_pytorch.py::TestQwen3_5_4B::test_fp8 SKIP (https://nvbugs/6273850)
 full:RTX_PRO_6000_Blackwell_Server_Edition/accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype[mtp_nextn=2-overlap_scheduler=True] SKIP (https://nvbugs/6400067)


The removals below this line include configurations the stated root cause doesn't cover — e.g. test_bfloat16[mtp_nextn=2-attention_dp=False-cuda_graph=True-overlap_scheduler=False-torch_compile=False-...] and test_nvfp4[...-cuda_graph=False-overlap_scheduler=False-torch_compile=False]. If the IMA comes from torch.compile FX node ordering around piecewise CUDA graphs, the torch_compile=False cases were failing for some other reason and un-waiving them will just re-open the CI failure. Restrict the removal to configs you actually re-ran green.

pengbowang-nv

Review from attention perspective: the change to MLA is not meaningful and is a no-op.

coderabbitai · 2026-08-11T03:53:20Z

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.

…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>

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 2 times, most recently 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

brnguyen2 reviewed Aug 8, 2026

View reviewed changes

pengbowang-nv requested changes Aug 11, 2026

View reviewed changes

trtllm-agent force-pushed the repair-bot-bug6473374 branch from c61dacb to f26152e Compare August 11, 2026 03:53

trtllm-agent and others added 2 commits August 11, 2026 04:56

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

8781a2d

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

trtllm-agent force-pushed the repair-bot-bug6473374 branch from f26152e to 8781a2d Compare August 11, 2026 11:59

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!

brnguyen2 left a comment

Uh oh!

brnguyen2 Aug 8, 2026

Uh oh!

pengbowang-nv left a comment

Uh oh!

coderabbitai Bot commented Aug 11, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

11 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!

brnguyen2 left a comment

Choose a reason for hiding this comment

Uh oh!

brnguyen2 Aug 8, 2026

Choose a reason for hiding this comment

Uh oh!

pengbowang-nv left a comment

Choose a reason for hiding this comment

Uh oh!

coderabbitai Bot commented Aug 11, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

11 participants

trtllm-agent commented Jul 20, 2026 •

edited by coderabbitai Bot

Loading

coderabbitai Bot commented Jul 20, 2026 •

edited

Loading