[https://nvbugs/6517842][fix] Handle mutable tensor lists in remove c…

tensorrt_llm/_torch/compilation/remove_copy_pass.py

@@ -37,20 +37,45 @@ def remove_functionalize_inner(node: Node, mutates_args: dict, is_v2=False):
         ]
 
         kwargs = {k: v for k, v in node.kwargs.items() if not k.startswith("_")}
+        tensor_list_replacements = {}
         if is_v2:
             all_bases = node.kwargs["_all_bases"]
+            # Mutated bases are flattened into consecutive outputs after the
+            # operator's regular outputs.
+            first_mutated_output = min(mutates_args)
             for arg in inplace_func._schema.arguments:
                 if arg.alias_info is None or not arg.alias_info.is_write:
                     continue
-                base_index_key = f"_{arg.name}_base_index"
-                base_index = node.kwargs[base_index_key]
-                kwargs[arg.name] = (None if base_index is None else
-                                    all_bases[base_index])
+                length_key = f"_{arg.name}_length"
+                if length_key in node.kwargs:
+                    length = node.kwargs[length_key]
+                    if length is None:
+                        kwargs[arg.name] = None
+                        continue
+
+                    kwargs[arg.name] = []
+                    for index in range(length):
+                        base_index = node.kwargs[
+                            f"_{arg.name}_{index}_base_index"]
+                        base = (None if base_index is None else
+                                all_bases[base_index])
+                        kwargs[arg.name].append(base)
+                        if base_index is not None:
+                            tensor_list_replacements[first_mutated_output +
+                                                     base_index] = (arg.name,
+                                                                    base)
+                else:
+                    base_index = node.kwargs[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]
-            mutated_arg = mutates_args[idx]
-            replacement = kwargs[mutated_arg]
+            if idx in tensor_list_replacements:
+                mutated_arg, replacement = tensor_list_replacements[idx]
+            else:
+                mutated_arg = mutates_args[idx]
+                replacement = kwargs[mutated_arg]
             assert replacement is not None, (
                 f"getitem user for optional output '{mutated_arg}' "
                 "has no base tensor -- graph is malformed")

tests/integration/test_lists/waives.txt

@@ -39,7 +39,6 @@ accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[pp4-mt
 accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[pp4-mtp_nextn=2-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=True] SKIP (https://nvbugs/6388153)
 accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[tp2pp2-mtp_nextn=0-attention_dp=False-cuda_graph=False-overlap_scheduler=False-torch_compile=False] SKIP (https://nvbugs/6428094)
 accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[tp2pp2-mtp_nextn=0-attention_dp=False-cuda_graph=False-overlap_scheduler=False-torch_compile=True] SKIP (https://nvbugs/6428096)
-accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[tp2pp2-mtp_nextn=2-attention_dp=False-cuda_graph=False-overlap_scheduler=False-torch_compile=True] SKIP (https://nvbugs/6517842)
 accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[tp2pp2-mtp_nextn=2-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=True] SKIP (https://nvbugs/6198774)
 accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[tp4-mtp_nextn=2-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=True] SKIP (https://nvbugs/6198774)
 accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus_online_eplb[mtp_nextn=2-moe_backend=WIDEEP] SKIP (https://nvbugs/6313993)
@@ -60,7 +59,6 @@ accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backe
 accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTLASS-mtp_nextn=0-pp4-fp8kv=True-attention_dp=True-cuda_graph=True-overlap_scheduler=True-low_precision_combine=False-torch_compile=True] SKIP (https://nvbugs/6428063)
 accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTLASS-mtp_nextn=0-tp2pp2-fp8kv=False-attention_dp=False-cuda_graph=False-overlap_scheduler=False-low_precision_combine=False-torch_compile=False] SKIP (https://nvbugs/6384625)
 accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTLASS-mtp_nextn=0-tp2pp2-fp8kv=False-attention_dp=False-cuda_graph=False-overlap_scheduler=False-low_precision_combine=False-torch_compile=True] SKIP (https://nvbugs/6384625)
-accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTLASS-mtp_nextn=0-tp2pp2-fp8kv=True-attention_dp=True-cuda_graph=True-overlap_scheduler=True-low_precision_combine=False-torch_compile=True] SKIP (https://nvbugs/6445472)
 accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTLASS-mtp_nextn=2-pp4-fp8kv=False-attention_dp=False-cuda_graph=False-overlap_scheduler=False-low_precision_combine=False-torch_compile=False] SKIP (https://nvbugs/6422432)
 accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTLASS-mtp_nextn=2-pp4-fp8kv=True-attention_dp=True-cuda_graph=True-overlap_scheduler=True-low_precision_combine=False-torch_compile=False] SKIP (https://nvbugs/6245394)
 accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTLASS-mtp_nextn=2-tp2pp2-fp8kv=False-attention_dp=False-cuda_graph=False-overlap_scheduler=False-low_precision_combine=False-torch_compile=False] SKIP (https://nvbugs/6384625)
@@ -146,7 +144,6 @@ full:B200/accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16[mt
 full:B200/accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16[mtp_nextn=2-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=True-enable_chunked_prefill=True-v2_kv_cache=True] SKIP (https://nvbugs/6422343)
 full:B200/accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[tp4-mtp_nextn=0-attention_dp=False-cuda_graph=False-overlap_scheduler=False-torch_compile=False] SKIP (https://nvbugs/6525007)
 full:B200/accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_cute_dsl_nvfp4_4gpus[tp4-fp8kv=False-attention_dp=False-cuda_graph=False-overlap_scheduler=False-torch_compile=False] SKIP (https://nvbugs/6526186)
-full:B200/accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_cute_dsl_nvfp4_4gpus[tp4-fp8kv=False-attention_dp=False-cuda_graph=False-overlap_scheduler=False-torch_compile=True] SKIP (https://nvbugs/6474888)
 full:B200/accuracy/test_llm_api_pytorch.py::TestLagunaXS::test_fp8 SKIP (https://nvbugs/6525011)
 full:B200/accuracy/test_llm_api_pytorch.py::TestLlama3_3_70BInstruct::test_nvfp4_tp4[torch_compile=True] SKIP (https://nvbugs/6525010)
 full:B200/accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_auto_dtype[tp_size=8-ep_size=8] SKIP (https://nvbugs/6384747)
@@ -172,7 +169,6 @@ full:B300/accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_nvfp4_multi_gp
 full:B300/accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_nvfp4_multi_gpus_chunked_prefill[latency] SKIP (https://nvbugs/6483369)
 full:B300/accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_nvfp4_multi_gpus_chunked_prefill[latency_qsplit] SKIP (https://nvbugs/6423866)
 full:B300/accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16[mtp_nextn=2-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=True-enable_chunked_prefill=True-v2_kv_cache=True] SKIP (https://nvbugs/6422343)
-full:B300/accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_cute_dsl_nvfp4_4gpus[tp4-fp8kv=False-attention_dp=False-cuda_graph=False-overlap_scheduler=False-torch_compile=True] SKIP (https://nvbugs/6474888)
 full:B300/accuracy/test_llm_api_pytorch.py::TestGemma3_1BInstruct::test_fp8_prequantized[torch_compile=True] SKIP (https://nvbugs/6475346)
 full:B300/accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_auto_dtype[tp_size=8-ep_size=8] SKIP (https://nvbugs/6445375)
 full:B300/accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_mxfp8[use_msa=False] SKIP (https://nvbugs/6424188)
@@ -194,7 +190,6 @@ full:GB200/accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16[m
 full:GB200/accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16[mtp_nextn=2-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=True-enable_chunked_prefill=True-v2_kv_cache=False] SKIP (https://nvbugs/6400067)
 full:GB200/accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16[mtp_nextn=2-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=True-enable_chunked_prefill=True-v2_kv_cache=True] SKIP (https://nvbugs/6400067)
 full:GB200/accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_cute_dsl_bf16_gemm[cuda_graph=True] SKIP (https://nvbugs/6525897)
-full:GB200/accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_cute_dsl_nvfp4_4gpus[tp4-fp8kv=False-attention_dp=False-cuda_graph=False-overlap_scheduler=False-torch_compile=True] SKIP (https://nvbugs/6474888)
 full:GB200/accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_mxfp8[use_msa=False] SKIP (https://nvbugs/6479471)
 full:GB200/accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4[use_msa=False] SKIP (https://nvbugs/6479471)
 full:GB200/accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_fp8_4gpus[attention_dp_on-python_mamba_cache] SKIP (https://nvbugs/6528742)
@@ -218,7 +213,6 @@ full:GB300/accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_gen_
 full:GB300/accuracy/test_llm_api_autodeploy.py::TestNemotronNanoV3::test_accuracy[nvfp4-1-attn_dp_off-trtllm] SKIP (https://nvbugs/6329165)
 full:GB300/accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16[mtp_nextn=2-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=True-enable_chunked_prefill=True-v2_kv_cache=True] SKIP (https://nvbugs/6422343)
 full:GB300/accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[pp4-mtp_nextn=2-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=False] SKIP (https://nvbugs/6525057)
-full:GB300/accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_cute_dsl_nvfp4_4gpus[tp4-fp8kv=False-attention_dp=False-cuda_graph=False-overlap_scheduler=False-torch_compile=True] SKIP (https://nvbugs/6474888)
 full:GB300/accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_fp8_block_scales_4gpus[pp4-mtp_nextn=0-fp8kv=True-attention_dp=False-cuda_graph=True-overlap_scheduler=True-torch_compile=True-sampler_async_worker=False] SKIP (https://nvbugs/6388153)
 full:GB300/accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTLASS-mtp_nextn=0-pp4-fp8kv=True-attention_dp=True-cuda_graph=True-overlap_scheduler=True-low_precision_combine=False-torch_compile=False] SKIP (https://nvbugs/6388153)
 full:GB300/accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_chunked_prefill[trtllm-fp8] SKIP (https://nvbugs/6529692)

tests/unittest/_torch/compilation/test_remove_copy_pass.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from collections.abc import Callable
 from operator import getitem
 
 import pytest
@@ -57,6 +58,46 @@ def test_remove_copy_for_mutates_args_auto_functionalized_v2(
     graph.lint()
 
 
+@pytest.mark.parametrize(
+    "inplace_func",
+    [
+        torch.ops.trtllm.pp_recv_tensors.default,
+        torch.ops.trtllm.pp_send_tensors.default,
+    ],
+)
+def test_remove_copy_for_mutates_tensor_list(
+    inplace_func: Callable[..., object],
+) -> None:
+    graph = Graph()
+    tensor_0 = graph.placeholder("tensor_0")
+    tensor_1 = graph.placeholder("tensor_1")
+    functionalized = graph.call_function(
+        auto_functionalized_v2,
+        args=(inplace_func,),
+        kwargs={
+            "_all_bases": (tensor_0, tensor_1),
+            "_tensors_length": 2,
+            "_tensors_0_base_index": 0,
+            "_tensors_1_base_index": 1,
+        },
+    )
+    mutated_0 = graph.call_function(getitem, args=(functionalized, 1))
+    mutated_1 = graph.call_function(getitem, args=(functionalized, 2))
+    clone_0 = graph.call_function(torch.ops.aten.clone.default, args=(mutated_0,))
+    clone_1 = graph.call_function(torch.ops.aten.clone.default, args=(mutated_1,))
+    graph.output((clone_0, clone_1))
+
+    remove_copy_pass.remove_copy_for_mutates_args(graph)
+
+    inplace_nodes = [node for node in graph.nodes if node.target == inplace_func]
+    assert len(inplace_nodes) == 1
+    assert inplace_nodes[0].kwargs == {"tensors": [tensor_0, tensor_1]}
+    assert clone_0.args[0] is tensor_0
+    assert clone_1.args[0] is tensor_1
+    assert all(node.target != auto_functionalized_v2 for node in graph.nodes)
+    graph.lint()
+
+
 def test_remove_copy_for_mutates_args_restores_optional_none() -> None:
     graph = Graph()
     hidden_states = graph.placeholder("hidden_states")