diff --git a/tensorrt_llm/_torch/compilation/remove_copy_pass.py b/tensorrt_llm/_torch/compilation/remove_copy_pass.py
index 0c8d3ede8b17..0857b8ca8c60 100644
--- a/tensorrt_llm/_torch/compilation/remove_copy_pass.py
+++ b/tensorrt_llm/_torch/compilation/remove_copy_pass.py
@@ -31,20 +31,17 @@ 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:
+        if is_call_function(node, auto_functionalized_v2):
             all_bases = node.kwargs["_all_bases"]
-            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
+            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:
diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt
index 80e697faf3be..1cb6ee95825b 100644
--- a/tests/integration/test_lists/waives.txt
+++ b/tests/integration/test_lists/waives.txt
@@ -239,10 +239,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)