Update argument

nod-ai · Oct 16, 2024 · e2eb98c · e2eb98c
1 parent 4af53c3
commit e2eb98c
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/sharktank/tests/evaluate/perplexity_test.py b/sharktank/tests/evaluate/perplexity_test.py
@@ -62,7 +62,7 @@ def test_llama3_8B_f16_non_decomposed(self):
             [
                 f"--gguf-file={llama_8b_f16_gguf_path}",
                 f"--tokenizer-config-json={llama_8b_f16_tokenizer_path}",
-                f"--attn-kernel=torch_sdpa",
+                f"--attention-kernel=torch_sdpa",
             ]
         )
 
@@ -141,7 +141,7 @@ def test_llama3_8B_fp8_non_decomposed(self):
             [
                 f"--gguf-file={llama_8b_fp8_gguf_path}",
                 f"--tokenizer-config-json={llama_8b_fp8_tokenizer_path}",
-                f"--attn-kernel=torch_sdpa",
+                f"--attention-kernel=torch_sdpa",
             ]
         )
 
@@ -229,7 +229,7 @@ def test_llama3_405B_f16_non_decomposed(self):
                 f"--gguf-file={llama_405b_f16_gguf_path}",
                 f"--tokenizer-config-json={llama_405b_f16_tokenizer_path}",
                 f"--tensor-parallelism-size={tensor_parallelism_size}",
-                f"--attn-kernel=torch_sdpa",
+                f"--attention-kernel=torch_sdpa",
             ]
         )
 
@@ -314,7 +314,7 @@ def test_llama3_405B_fp8_non_decomposed(self):
                 f"--gguf-file={llama_405b_fp8_gguf_path}",
                 f"--tokenizer-config-json={llama_405b_fp8_tokenizer_path}",
                 f"--tensor-parallelism-size={tensor_parallelism_size}",
-                f"--attn-kernel=torch_sdpa",
+                f"--attention-kernel=torch_sdpa",
             ]
         )