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 3, 14:04 UTC
[Perf] Optimize MLA attention `_v_up_proj` bmm by removing additional copy (#42561)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
May 15, 2026 at 11:14 UTC
491e8d8539ba35045f3e0ab7954085564304aba5
1 file changed
+2
-13
vllm/model_executor/layers/attention/mla_attention.py
+2
-13
@@ -989,19 +989,8 @@ class MLAAttention(nn.Module, AttentionLayerBase):
989
x, self.W_V, self.W_V_scale, group_size=128, transpose_bm=True, YQ=out
990
)
991
else:
992
- # Convert from (B, N * V) to (N, B, V)
993
- out = out.transpose(0, 1)
994
-
995
- # Multiply (N, B, L) x (N, L, V) -> (N, B, V)
996
- torch.bmm(x, self.W_UV, out=out) # Reuse "out" to make it "hot"
997
-
998
- # Convert from (N, B, V) to (B, N * V)
999
- out_new = out.transpose(0, 1).reshape(-1, self.num_heads * self.v_head_dim)
1000
-
1001
- # Adjust output buffer shape back to the original (B, N * V)
1002
- N, B, V = out.shape
1003
- out.resize_((B, N * V))
1004
- out.copy_(out_new) # Copy result
992
+ # Multiply + Transpose (N, B, L) x (N, L, V)->(N, B, V)->(B, N, V)
993
+ torch.bmm(x, self.W_UV, out=out.transpose(0, 1))
994
995
996
def unified_mla_kv_cache_update(