Fix mla prefill context performance

Signed-off-by: ZhongYingMatrix <zhongyingmatrix@gmail.com>

Author: ZhongYingMatrix

vllm/attention/backends/mla/common.py

@@ -1308,7 +1308,7 @@ def _compute_prefill_context(
             )
 
             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)
 

vllm/v1/attention/backends/mla/common.py

@@ -874,7 +874,7 @@ def _compute_prefill_context(
             )
 
             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)