[ROCm][AMD][Bugfix] adding a missing triton autotune config (#4845)

b5853f99 · Hongxia Yang · GitHub · f09edd8a · b5853f99
Unverified Commit b5853f99 authored May 16, 2024 by Hongxia Yang Committed by GitHub May 16, 2024
Show whitespace changes
Inline Side-by-side

Showing with 10 additions and 0 deletions

vllm/attention/ops/triton_flash_attention.py vllm/attention/ops/triton_flash_attention.py +10 -0

No files found.
--- 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,