diff --git a/sharktank/tests/models/llama/benchmark_amdgpu_test.py b/sharktank/tests/models/llama/benchmark_amdgpu_test.py
index 166b0b78a..13a2c35e4 100644
--- a/sharktank/tests/models/llama/benchmark_amdgpu_test.py
+++ b/sharktank/tests/models/llama/benchmark_amdgpu_test.py
@@ -270,7 +270,7 @@ def testBenchmark8B_f16_Non_Decomposed_Prefill_Input_Len_2048(self):
             hip_device_id=self.iree_device,
             vmfb_name=output_vmfb,
             irpa_path=self.irpa_path,
-            args=self.iree_run_prefill_nondecomposed_args_fp16,
+            args=self.iree_run_prefill_nondecomposed_args_fp16_2048,
             cwd=self.repo_root,
         )