Fix triton compilation issue (#3984)

Co-authored-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>

Fix triton compilation issue (#3984)
Co-authored-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
d04973ad · Bellk17 · GitHub · fbb9d9ee · d04973ad
Unverified Commit d04973ad authored Apr 12, 2024 by Bellk17 Committed by GitHub Apr 12, 2024
Show whitespace changes
Inline Side-by-side

Showing with 5 additions and 1 deletion

vllm/attention/ops/triton_flash_attention.py vllm/attention/ops/triton_flash_attention.py +5 -1

No files found.
--- a/vllm/attention/ops/triton_flash_attention.py
+++ b/vllm/attention/ops/triton_flash_attention.py
@@ -415,7 +415,11 @@ def attn_fwd(
            return

    is_mqa = hq != hk
-    off_h_k = off_h_q % hk if is_mqa else off_h_q
+    if is_mqa:  # noqa: SIM108
+        off_h_k = off_h_q % hk
+    else:
+        off_h_k = off_h_q
+
    n_extra_tokens = 0
    if seqlen_k < BLOCK_N:
        n_extra_tokens = BLOCK_N - seqlen_k