From 5be7bee724954ca3783f740f93064eaa86346cfc Mon Sep 17 00:00:00 2001 From: Jin Li <59594262+liji-nv@users.noreply.github.com> Date: Mon, 27 Jul 2026 20:20:00 -0700 Subject: [PATCH] [https://nvbugs/6517842][fix] Handle mutable tensor lists in remove copy pass auto_functionalized_v2 represents mutable Tensor[] arguments with a list length and per-element base indices. The remove-copy pass assumed every mutable argument had a single base index, causing torch.compile to look up the nonexistent _tensors_base_index key for pipeline-parallel send and receive operations. Reconstruct mutable tensor lists from their flattened bases, map each functionalized output back to the corresponding input tensor, and retain the existing scalar and optional-tensor behavior. Add regression coverage for both pp_send_tensors and pp_recv_tensors. Unwaive the cases tracked by NVBUG 6517842, NVBUG 6445472, and NVBUG 6474888. Tested on B200: - 5 remove_copy_pass unit tests passed - DeepSeekV3-Lite TP2/PP2 torch.compile accuracy case passed - TensorRT-LLM build completed successfully Signed-off-by: Jin Li <59594262+liji-nv@users.noreply.github.com> --- .../_torch/compilation/remove_copy_pass.py | 37 ++++++++++++++--- tests/integration/test_lists/waives.txt | 6 --- .../compilation/test_remove_copy_pass.py | 41 +++++++++++++++++++ 3 files changed, 72 insertions(+), 12 deletions(-) diff --git a/tensorrt_llm/_torch/compilation/remove_copy_pass.py b/tensorrt_llm/_torch/compilation/remove_copy_pass.py index 0c8d3ede8b17..f9bb925a1e1c 100644 --- a/tensorrt_llm/_torch/compilation/remove_copy_pass.py +++ b/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") diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index e0da833bc85b..ffb09cd12204 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -38,7 +38,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=0-attention_dp=False-cuda_graph=True-overlap_scheduler=False-torch_compile=False] SKIP (https://nvbugs/6198785) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[tp4-mtp_nextn=0-attention_dp=False-cuda_graph=True-overlap_scheduler=True-torch_compile=True] SKIP (https://nvbugs/6198785) @@ -61,7 +60,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) @@ -145,7 +143,6 @@ full:A100X/llmapi/test_llm_examples.py::test_llmapi_speculative_decoding_mtp SKI full:B200/accuracy/test_epd_disagg_multimodal.py::TestVideoMMEEPD::test_disaggregated_videomme[nemotron_nano_v3_omni_fp8] SKIP (https://nvbugs/6327718) 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=False] SKIP (https://nvbugs/6422343) 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_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::TestMiniMaxM3::test_auto_dtype[tp_size=8-ep_size=8] SKIP (https://nvbugs/6384747) full:B200/accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_mxfp8[use_msa=False] SKIP (https://nvbugs/6424188) full:B200/accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4[use_msa=False] SKIP (https://nvbugs/6424188) @@ -167,7 +164,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) @@ -185,7 +181,6 @@ full:GB200/accuracy/test_dwdp_disaggregated_serving.py::TestDwdpDeepSeekV3Lite:: full:GB200/accuracy/test_dwdp_disaggregated_serving.py::TestDwdpDeepSeekV3Lite::test_dwdp_accuracy_mode_b_overlap SKIP (https://nvbugs/6276923) 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_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::TestNemotronV3Ultra::test_nvfp4_4gpus_block_reuse[ADP4] SKIP (https://nvbugs/6474894) @@ -203,7 +198,6 @@ full:GB200/test_e2e.py::test_qwen_e2e_cpprunner_large_new_tokens[DeepSeek-R1-Dis full:GB300/accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_gen_first[adp-mtp2] SKIP (https://nvbugs/6295740) 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_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::TestGemma3_1BInstruct::test_fp8_prequantized[torch_compile=True] SKIP (https://nvbugs/6475346) diff --git a/tests/unittest/_torch/compilation/test_remove_copy_pass.py b/tests/unittest/_torch/compilation/test_remove_copy_pass.py index ab5d7afad17a..ef78351e9453 100644 --- a/tests/unittest/_torch/compilation/test_remove_copy_pass.py +++ b/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")