A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 21:00 UTC
Fix mla prefill context performance (#13897)
Signed-off-by: ZhongYingMatrix <zhongyingmatrix@gmail.com>
Ying Zhong committed
Mar 7, 2025 at 01:35 UTC
9f1710f1ace3535920c0bb6d4cc329c36289080e
2 files changed
+2
-2
vllm/attention/backends/mla/common.py
+1
-1
@@ -1308,7 +1308,7 @@ class MLACommonImpl(MLAAttentionImpl[T], Generic[T]):
1308
)
1309
1310
kv_c_normed = workspace[:toks]\
1311
- [..., :self.kv_lora_rank].unsqueeze(1)
1311
+ [..., :self.kv_lora_rank]
1312
k_pe = workspace[:toks]\
1313
[..., self.kv_lora_rank:].unsqueeze(1)
1314
vllm/v1/attention/backends/mla/common.py
+1
-1
@@ -874,7 +874,7 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
874
)
875
876
kv_c_normed = workspace[:toks]\
877
- [..., :self.kv_lora_rank].unsqueeze(1)
877
+ [..., :self.kv_lora_rank]
878
k_pe = workspace[:toks]\
879
[..., self.kv_lora_rank:].unsqueeze(1)
880