From 613e5ebccd0261c8c6a730bbbb82458d9bea85d7 Mon Sep 17 00:00:00 2001 From: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> Date: Mon, 20 Jul 2026 04:37:10 -0700 Subject: [PATCH] [nvbugs/6478707][fix] Wire auto_functionalized_v2 in remove_copy_pass 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 __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 __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 __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> --- tensorrt_llm/_torch/compilation/utils.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/compilation/utils.py b/tensorrt_llm/_torch/compilation/utils.py index 963c0cb1b1a4..99e0cbb63789 100644 --- a/tensorrt_llm/_torch/compilation/utils.py +++ b/tensorrt_llm/_torch/compilation/utils.py @@ -183,7 +183,9 @@ def inplace_info(): 2: "output_sf" }, "mla_custom_op_inplace": { - 1: "output" + 1: "output", + 2: "dsv4_output", + 3: "dsv4_output_sf", }, "mla_dsa_attn_inplace": { 1: "output"