[https://nvbugs/6478707][fix] detect v2 inside remove_functionalize_inner, resolve each mutates_args via… by trtllm-agent · Pull Request #16629 · NVIDIA/TensorRT-LLM

trtllm-agent · 2026-07-20T13:46:42Z

Summary

Root cause: remove_copy_pass never treated auto_functionalized_v2 HOP nodes as v2, so kwargs["output"] raised KeyError since v2 stores mutable args in _all_bases
Fix: detect v2 inside remove_functionalize_inner, resolve each mutates_args via _<arg>_base_index into _all_bases (None-safe), and register the two additional mla_custom_op_inplace mutable args (dsv4_output, dsv4_output_sf) in inplace_info
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/6478707

Dev Engineer Review

Updated remove_copy_pass handling for auto_functionalized_v2 nodes to resolve mutable arguments through _all_bases using _<arg>_base_index, including None handling.
Registered dsv4_output and dsv4_output_sf as mutable outputs for mla_custom_op_inplace.
Addresses KeyError failures caused by v2 nodes storing mutable arguments indirectly.
No configuration, public API, or test-list changes identified.

QA Engineer Review

No test changes.

coderabbitai · 2026-07-20T13:51: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: fa000fff-29f0-4f13-8592-ba79e49700d3

📥 Commits

Reviewing files that changed from the base of the PR and between d7dfd7f and 613e5eb.

📒 Files selected for processing (1)

tensorrt_llm/_torch/compilation/utils.py

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

tensorrt_llm/_torch/compilation/utils.py

Walkthrough

The MLA custom operation’s inplace metadata now records dsv4_output and dsv4_output_sf alongside the existing output entry.

Changes

MLA inplace metadata

Layer / File(s)	Summary
Additional mutated output mappings `tensorrt_llm/_torch/compilation/utils.py`	`mla_custom_op_inplace` now maps indices 2 and 3 to `dsv4_output` and `dsv4_output_sf`.

Estimated code review effort: 1 (Trivial) | ~3 minutes

Suggested reviewers: brnguyen2, mikeiovine, tburt-nv

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title is specific and matches the main fix: v2 detection and mutable-arg resolution in remove_functionalize_inner.
Description check	✅ Passed	The description covers summary, test plan, and bug link, but it does not follow the full template sections exactly.
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.}

lori-ren · 2026-07-22T02:13:18Z

Some changes are already done by #16624, should we merge #16624 before this one?

JunyiXu-nv · 2026-07-22T02:36:45Z

Some changes are already done by #16624, should we merge #16624 before this one?

I see. I'll hold this PR before #16624 get merged.

Torch 2.12 AOT-autograd emits auto_functionalized_v2 HOP nodes for mutable-arg ops even when inductor's enable_auto_functionalized_v2 is False (that flag only affects inductor lowering). In v2 nodes, mutable tensors live in _all_bases and each mutable arg has its own _<name>_base_index kwarg; direct kwargs like 'output' are absent. Previously remove_copy_for_mutates_args never detected v2 nodes at the call site, so kwargs["output"] raised KeyError. The v2 branch itself keyed into _all_bases by mutates_args positional index (k-1) rather than by the _<name>_base_index kwarg, which is unrelated to op-arg position. Detect auto_functionalized_v2 inside the inner helper and, for each mutable arg, resolve its base tensor via _<arg>_base_index into _all_bases, mapping a missing/None index to None so optional mutable args are preserved. Also register the two additional mutable args (dsv4_output, dsv4_output_sf) for mla_custom_op_inplace in inplace_info() so the reconstructed op call carries them. This unblocks TestDeepSeekV3Lite::test_nvfp4_4gpus with torch_compile under moe_backend=TRTLLM on B200. Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>

lori-ren · 2026-07-27T08:52:28Z

@JunyiXu-nv #16629 may have fixed this issue, should we unwaive these tests?

trtllm-agent · 2026-08-07T06:15:35Z

NVBug 6478707 is closed as Bug - Fixed. The linked bug appears resolved elsewhere or for a reason that does not prove this PR is redundant. This PR should be judged on its own merits; repair-bot is not auto-closing it.

brnguyen2

Approving — the comments below are optional touch-ups, not blockers.

Indices match the op's declared mutates_args=("output", "dsv4_output", "dsv4_output_sf") at tensorrt_llm/_torch/modules/mla.py:161, so the mapping is right.

Two things:

The description doesn't match the diff. The auto_functionalized_v2 detection and _all_bases resolution it claims are already on main (commit e96c330, PR #16410), and nvbugs/6478707 is already closed as fixed by that PR. This diff is only the two inplace_info entries — please retitle/rewrite so the change stands on its own (registering the DSv4 mutable args).
No test covers the new entries; see the inline note.

brnguyen2 · 2026-08-09T13:34:17Z

        "mla_custom_op_inplace": {
-            1: "output"
+            1: "output",
+            2: "dsv4_output",


Please add a case to tests/unittest/_torch/compilation/test_remove_copy_pass.py covering the DSv4-enabled graph — _all_bases=(output, dsv4_output, dsv4_output_sf) with all three _<arg>_base_index set and getitem users at 1/2/3. The existing tests only exercise the all-None variant, so index 2/3 here are untested; under auto_functionalized_v2 the getitem index is a position in _all_bases, not in the declared arg list, and a wrong entry silently rewires a getitem to the wrong buffer rather than failing loudly.

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

trtllm-agent assigned JunyiXu-nv Jul 20, 2026

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

trtllm-agent requested review from asfiyab-nvidia, brnguyen2, pcastonguay, schetlur-nv, tburt-nv, yingguo-trt and yuanjingx87 July 20, 2026 13:46

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

mikeiovine requested review from mikeiovine and removed request for pcastonguay and schetlur-nv July 20, 2026 15:22

yingguo-trt approved these changes Jul 21, 2026

View reviewed changes

JunyiXu-nv approved these changes Jul 22, 2026

View reviewed changes

asfiyab-nvidia approved these changes Jul 23, 2026

View reviewed changes

trtllm-agent force-pushed the repair-bot-bug6478707 branch from d7dfd7f to 613e5eb Compare July 24, 2026 08:48

tburt-nv removed request for a team, tburt-nv and yuanjingx87 July 27, 2026 18:47

brnguyen2 approved these changes Aug 9, 2026

View reviewed changes

Navigation Menu

Uh oh!

[https://nvbugs/6478707][fix] detect v2 inside remove_functionalize_inner, resolve each mutates_args via… - #16629

[https://nvbugs/6478707][fix] detect v2 inside remove_functionalize_inner, resolve each mutates_args via…#16629
trtllm-agent wants to merge 1 commit into
NVIDIA:mainfrom
tensorrt-cicd:repair-bot-bug6478707

trtllm-agent commented Jul 20, 2026 •

edited by coderabbitai Bot

Loading

Uh oh!

coderabbitai Bot commented Jul 20, 2026 •

edited

Loading

Uh oh!

lori-ren commented Jul 22, 2026

Uh oh!

JunyiXu-nv commented Jul 22, 2026

Uh oh!

lori-ren commented Jul 27, 2026

Uh oh!

trtllm-agent commented Aug 7, 2026

Uh oh!

brnguyen2 left a comment

Uh oh!

brnguyen2 Aug 9, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

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

lori-ren commented Jul 22, 2026

Uh oh!

JunyiXu-nv commented Jul 22, 2026

Uh oh!

lori-ren commented Jul 27, 2026

Uh oh!

trtllm-agent commented Aug 7, 2026

Uh oh!

brnguyen2 left a comment

Choose a reason for hiding this comment

Uh oh!

brnguyen2 Aug 9, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

trtllm-agent commented Jul 20, 2026 •

edited by coderabbitai Bot

Loading

coderabbitai Bot commented Jul 20, 2026 •

edited

Loading