Merge pull request opendatahub-io#17 from ROCm/triton-config-fix

[ROCm] adding a missing triton autotune config
prarit · May 16, 2024 · bebcbe6 · bebcbe6
2 parents 4b39609 + dfef216
commit bebcbe6
Showing 1 changed file with 10 additions and 0 deletions.
diff --git a/vllm/attention/ops/triton_flash_attention.py b/vllm/attention/ops/triton_flash_attention.py
@@ -239,6 +239,16 @@ def _attn_fwd_inner(
             num_stages=1,
             num_warps=8,
         ),
+        triton.Config(
+            {
+                "BLOCK_M": 128,
+                "BLOCK_N": 64,
+                "waves_per_eu": 1,
+                "PRE_LOAD_V": False,
+            },
+            num_stages=1,
+            num_warps=4,
+        ),
         triton.Config(
             {
                 "BLOCK_M": 128,