[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.
The change preserves FP4 scale-factor handling and avoids the FX graph and synchronization change that caused the SM120 asynchronous IMA.
The change does not alter public APIs.
waives.txt removes seven RTX PRO 6000 Blackwell waiver entries.
The waiver format remains consistent, with no duplicate or unrelated entries.

QA Engineer Review

No test-db/ or qa/ files were modified.
The removed waivers cover five DeepSeekV3Lite BF16 mtp_nextn=2 cases and two DeepSeekV3Lite NVFP4 mtp_nextn=2 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

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

@coderabbitai resume to resume automatic reviews.
@coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

▶️ Resume reviews
🔍 Trigger review

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: 0c0ffbbc-8f4c-4703-a33f-153299475a24

📥 Commits

Reviewing files that changed from the base of the PR and between e189237 and a8b4a3c.

📒 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

Included review availability: Your plan provides up to 12 included reviews per hour; 9 remain after this review.

Walkthrough

MLA._forward_impl now uses type-specific slicing for FP4 quantized and plain tensors. The integration waiver list removes seven DeepSeekV3Lite mtp_nextn=2 entries for RTX PRO 6000 configurations.

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 seven related RTX PRO 6000 integration-test waivers.

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

Merge Risk: ⚪ Minimal · up to a8b4a

This PR restores the intended plain-tensor slicing path while preserving the specialized path for Fp4QuantizedTensor inputs; no actionable merge-blocking risk remains beyond normal checks and review.

Suggested reviewers: kris1025, tburt-nv

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly identifies bug 6473374 and summarizes the main fix in MLA.forward_impl.
Description check	✅ Passed	The description explains the root cause, fix, test plan, and bug link, but it does not include the template's full checklist.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 1 functions across 1 files.
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>

coderabbitai · 2026-08-21T00:38:23Z

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 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 force-pushed the repair-bot-bug6473374 branch from f26152e to 8781a2d Compare August 11, 2026 11:59

coderabbitai Bot mentioned this pull request Aug 12, 2026

[https://nvbugs/6140408][chore] Unwaive DeepSeekV3Lite test_nvfp4_4gpus CUTLASS pp4 and CUTEDSL tp2pp2 #17569

Merged

1 task

trtllm-agent and others added 2 commits August 20, 2026 17:34

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

a8b4a3c

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

trtllm-agent force-pushed the repair-bot-bug6473374 branch from 8781a2d to a8b4a3c Compare August 21, 2026 00:38

JunyiXu-nv closed this Sep 3, 2026

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!

Reviews paused

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!

coderabbitai Bot commented Aug 21, 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