Add missing iree_hal_target_device flags

nod-ai · Jan 6, 2025 · 69997e6 · 69997e6
1 parent afc8cff
commit 69997e6
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/sharktank/tests/models/llama/benchmark_amdgpu_test.py b/sharktank/tests/models/llama/benchmark_amdgpu_test.py
@@ -339,6 +339,7 @@ def setUp(self):
             irpa_path=str(self.irpa_path),
             batch_size=4,
             iree_hip_target="gfx942",
+            iree_hal_target_device="hip",
             attention_kernel="torch",
             tensor_parallelism_size=self.tensor_parallelism_size,
             block_seq_stride=32,
@@ -347,6 +348,7 @@ def setUp(self):
             irpa_path=str(self.irpa_path_fp8),
             batch_size=4,
             iree_hip_target="gfx942",
+            iree_hal_target_device="hip",
             attention_kernel="torch",
             tensor_parallelism_size=self.tensor_parallelism_size,
             block_seq_stride=32,
@@ -355,6 +357,7 @@ def setUp(self):
             irpa_path=str(self.irpa_path_fp8),
             batch_size=4,
             iree_hip_target="gfx942",
+            iree_hal_target_device="hip",
             attention_kernel="torch",
             tensor_parallelism_size=self.tensor_parallelism_size,
             block_seq_stride=32,
@@ -664,6 +667,7 @@ def setUp(self):
             irpa_path=str(self.irpa_path),
             batch_size=4,
             iree_hip_target="gfx942",
+            iree_hal_target_device="hip",
             attention_kernel="torch",
             tensor_parallelism_size=self.tensor_parallelism_size,
             block_seq_stride=32,
@@ -672,6 +676,7 @@ def setUp(self):
             irpa_path=str(self.irpa_path_fp8),
             batch_size=4,
             iree_hip_target="gfx942",
+            iree_hal_target_device="hip",
             attention_kernel="torch",
             tensor_parallelism_size=self.tensor_parallelism_size,
             block_seq_stride=32,