Update vllm/attention/ops/triton_flash_attention.py

Bellk17 · WoosukKwon · web-flow · commit fae4f82cb5d8 · 2024-04-11T16:16:06.000-07:00
Co-authored-by: Woosuk Kwon &lt;woosuk.kwon@berkeley.edu&gt;
diff --git a/vllm/attention/ops/triton_flash_attention.py b/vllm/attention/ops/triton_flash_attention.py
@@ -415,7 +415,7 @@ def attn_fwd(
             return
 
     is_mqa = hq != hk
-    if is_mqa:
+    if is_mqa:  # noqa: SIM108
         off_h_k = off_h_q % hk
     else:
         off_h_k = off_h_q
