single wg fixes

Signed-off-by: Lucas Wilkinson <[email protected]>

Author: LucasWilkinson

cmake/external_projects/vllm_flash_attn.cmake

@@ -38,7 +38,7 @@ else()
   FetchContent_Declare(
           vllm-flash-attn
           GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git
-          GIT_TAG e93779c59ba4905e56e5c39dc2c1904ada71fa21
+          GIT_TAG e46f09441a8ee3231e54551cc8994ca768178e69
           GIT_PROGRESS TRUE
           # Don't share the vllm-flash-attn build between build types
           BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn

vllm/v1/attention/backends/flash_attn.py

@@ -286,7 +286,9 @@ def __init__(self, runner: "GPUModelRunner"):
 
         self.runner = runner
         self.aot_schedule = (get_flash_attn_version() == 3)
-        self.num_heads = model_config.get_num_attention_heads(
+        self.num_heads_q = model_config.get_num_attention_heads(
+            runner.parallel_config)
+        self.num_heads_kv = model_config.get_num_kv_heads(
             runner.parallel_config)
         self.headdim = model_config.get_head_size()
         self.page_size = self.runner.block_size
@@ -340,8 +342,8 @@ def schedule(cu_query_lens, max_query_len, seqlens, max_seq_len,
                     max_seqlen_q=max_query_len,
                     max_seqlen_k=max_seq_len,
                     cache_seqlens=seqlens,
-                    num_heads_q=self.num_heads,
-                    num_heads_kv=self.num_heads,
+                    num_heads_q=self.num_heads_q,
+                    num_heads_kv=self.num_heads_kv,
                     headdim=self.headdim,
                     page_size=self.page_size,
                     cu_seqlens_q=cu_query_lens,