flatten qkv

Signed-off-by: Isotr0py <2037008807@qq.com>
vllm-project · Jan 26, 2025 · 159f0f2 · 159f0f2
1 parent 0d5228d
commit 159f0f2
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/vllm/attention/layer.py b/vllm/attention/layer.py
@@ -263,15 +263,16 @@ def forward(
                                         device=key.device)
 
             out = flash_attn_varlen_func(
-                query,
-                key,
-                value,
+                query.flatten(0, 1),
+                key.flatten(0, 1),
+                value.flatten(0, 1),
                 cu_seqlens_q=cu_seqlens_q,
                 cu_seqlens_k=cu_seqlens_k,
                 max_seqlen_q=q_len,
                 max_seqlen_k=kv_len,
                 softmax_scale=self.scale,
             )
+            out = out.reshape(bsz, q_len, -1)
         elif self.attn_backend == _Backend.XFORMERS:
             from xformers import ops as xops