Fix mla prefill context performance (#13897)

Signed-off-by: ZhongYingMatrix <zhongyingmatrix@gmail.com>
2025-03-07 01:35:49 +08:00
parent e642ec962c
commit 9f1710f1ac
2 changed files with 2 additions and 2 deletions
--- a/vllm/v1/attention/backends/mla/common.py
+++ b/vllm/v1/attention/backends/mla/common.py
@@ -874,7 +874,7 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
            )

            kv_c_normed = workspace[:toks]\
-                [..., :self.kv_lora_rank].unsqueeze(1)
+                [..., :self.kv_lora_rank]
            k_pe = workspace[:toks]\
                [..., self.kv_lora_rank:].unsqueeze(1)