From e578919c3f1f4d18140dc6c0bb8dc852cbb57184 Mon Sep 17 00:00:00 2001 From: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> Date: Mon, 20 Jul 2026 04:19:27 -0700 Subject: [PATCH] [nvbugs/6473161][fix] Wire auto_functionalized_v2 in remove_copy_pass Torch 2.11 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 'query'/'key' are absent. Previously remove_copy_for_mutates_args never set is_v2=True at the call site, and the v2 branch itself keyed into _all_bases by the mutates_args positional index (k-1) rather than by the __base_index kwarg, which is unrelated to op-arg position. Both bugs together caused KeyError('query') / IndexError on flashinfer_apply_rope_with_cos_sin_cache_inplace and similar ops under torch.compile+FLASHINFER. Detect auto_functionalized_v2 at the call site 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. Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> --- tensorrt_llm/_torch/compilation/remove_copy_pass.py | 11 +++++++---- tests/integration/test_lists/waives.txt | 4 ---- 2 files changed, 7 insertions(+), 8 deletions(-) diff --git a/tensorrt_llm/_torch/compilation/remove_copy_pass.py b/tensorrt_llm/_torch/compilation/remove_copy_pass.py index 8e5eb7a81148..efca38f91218 100644 --- a/tensorrt_llm/_torch/compilation/remove_copy_pass.py +++ b/tensorrt_llm/_torch/compilation/remove_copy_pass.py @@ -17,15 +17,18 @@ def remove_copy_for_mutates_args(graph: Graph): nodes_to_remove: list[Node] = [] - def remove_functionalize_inner(node: Node, mutates_args: dict, is_v2=False): + def remove_functionalize_inner(node: Node, mutates_args: dict): getitem_nodes = [ user for user in node.users if is_call_function(user, getitem) ] kwargs = {k: v for k, v in node.kwargs.items() if not k.startswith("_")} - if is_v2: - for k, v in mutates_args.items(): - kwargs[v] = node.kwargs["_all_bases"][k - 1] + if is_call_function(node, auto_functionalized_v2): + all_bases = node.kwargs["_all_bases"] + for arg_name in mutates_args.values(): + base_index = node.kwargs.get(f"_{arg_name}_base_index") + kwargs[arg_name] = (None if base_index is None else + all_bases[base_index]) for getitem_node in getitem_nodes: idx = getitem_node.args[1] diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index 264002846315..d27c2766c221 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -304,10 +304,6 @@ full:L40S/accuracy/test_llm_api_autodeploy.py::TestLlama3_1_8B::test_auto_dtype[ full:L40S/accuracy/test_llm_api_autodeploy.py::TestLlama3_1_8B::test_auto_dtype[trtllm-False-1] SKIP (https://nvbugs/6322045) full:L40S/accuracy/test_llm_api_autodeploy.py::TestNemotronH::test_auto_dtype[trtllm-flashinfer_ssm-False] SKIP (https://nvbugs/6327147) full:L40S/accuracy/test_llm_api_autodeploy.py::TestNemotronH::test_auto_dtype[trtllm-triton_ssm-False] SKIP (https://nvbugs/6327147) -full:L40S/accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_fp8[fp8kv=False-attn_backend=FLASHINFER-torch_compile=True] SKIP (https://nvbugs/6473161) -full:L40S/accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_fp8[fp8kv=True-attn_backend=FLASHINFER-torch_compile=True] SKIP (https://nvbugs/6473161) -full:L40S/accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_fp8_4gpus[tp4-fp8kv=False-attn_backend=FLASHINFER-torch_compile=True] SKIP (https://nvbugs/6473161) -full:L40S/accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_fp8_4gpus[tp4-fp8kv=True-attn_backend=FLASHINFER-torch_compile=True] SKIP (https://nvbugs/6473161) full:L40S/accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B::test_fp8[latency-torch_compile=True] SKIP (https://nvbugs/6276841) full:L40S/accuracy/test_llm_api_pytorch_multimodal.py::TestExaone4_5_33B::test_auto_dtype[forced_chunked_prefill] SKIP (https://nvbugs/6327149) full:L40S/accuracy/test_llm_api_pytorch_multimodal.py::TestExaone4_5_33B::test_auto_dtype[full_budget] SKIP (https://nvbugs/6327149)