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 14, 05:30 UTC
[Performance] Remove input pads in cutlass_mla and optimize v_proj output handling (#25184)
Signed-off-by: Alexander Matveev <amatveev@redhat.com>
Alexander Matveev committed
Sep 22, 2025 at 21:20 UTC
0b7bed9c386d41945b990d3ac5311ac7d9c8b499
2 files changed
+55
-20
vllm/v1/attention/backends/mla/common.py
+39
-6
@@ -942,6 +942,7 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
942
qk_head_dim: int,
943
v_head_dim: int,
944
kv_b_proj: ColumnParallelLinear,
945
+ q_pad_num_heads: Optional[int] = None,
946
) -> None:
947
if kv_sharing_target_layer_name is not None:
948
raise NotImplementedError("KV sharing is not supported for MLA")
@@ -959,6 +960,7 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
960
self.qk_head_dim = qk_head_dim
961
self.v_head_dim = v_head_dim
962
self.kv_b_proj = kv_b_proj
963
+ self.q_pad_num_heads = q_pad_num_heads
964
965
if use_flashinfer_prefill():
966
logger.debug_once("Using FlashInfer prefill for MLA")
@@ -1134,7 +1136,7 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
1136
True, #Indicates actual_seq_lens are on GPU or CPU.
1137
)
1138
1137
- def _v_up_proj(self, x):
1139
+ def _v_up_proj(self, x: torch.Tensor, out: torch.Tensor):
1140
# Convert from (B, N, L) to (N, B, L)
1141
x = x.view(-1, self.num_heads, self.kv_lora_rank).transpose(0, 1)
1142
if is_rocm_aiter_fp8bmm_enabled():
@@ -1146,12 +1148,23 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
1148
transpose_bm=True)
1149
# Convert from (B, N, V) to (B, N * V)
1150
x = x.reshape(-1, self.num_heads * self.v_head_dim)
1151
+ # Copy result
1152
+ out.copy_(x)
1153
else:
1154
+ # Convert from (B, N * V) to (N, B, V)
1155
+ out = out.view(-1, self.num_heads, self.v_head_dim).transpose(0, 1)
1156
+
1157
# Multiply (N, B, L) x (N, L, V) -> (N, B, V)
1151
- x = torch.bmm(x, self.W_UV)
1158
+ torch.bmm(x, self.W_UV, out=out) # Reuse "out" to make it "hot"
1159
+
1160
# Convert from (N, B, V) to (B, N * V)
1153
- x = x.transpose(0, 1).reshape(-1, self.num_heads * self.v_head_dim)
1154
- return x
1161
+ out_new = out.transpose(0, 1).reshape(
1162
+ -1, self.num_heads * self.v_head_dim)
1163
+
1164
+ # Adjust output buffer shape back to the original (B, N * V)
1165
+ N, B, V = out.shape
1166
+ out.resize_((B, N * V))
1167
+ out.copy_(out_new) # Copy result
1168
1169
def process_weights_after_loading(self, act_dtype: torch.dtype):
1170
@@ -1559,6 +1572,15 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
1572
# Convert from (B, N, P) to (N, B, P)
1573
decode_q_nope = decode_q_nope.transpose(0, 1)
1574
1575
+ # Pads the head_dim if necessary (for the underlying kernel)
1576
+ if self.q_pad_num_heads is not None:
1577
+ B, N, L = decode_q_pe.shape
1578
+ decode_pe_padded = decode_q_pe.new_empty(
1579
+ (B, self.q_pad_num_heads, L))
1580
+ decode_pe_padded.resize_((B, N, L))
1581
+ decode_pe_padded.copy_(decode_q_pe)
1582
+ decode_q_pe = decode_pe_padded
1583
+
1584
if is_rocm_aiter_fp8bmm_enabled():
1585
# Multiply+Transpose (N, B, P)x(N, P, L)->(N, B, L)->(B, N, L)
1586
decode_ql_nope = aiter_triton_fp8_bmm(decode_q_nope,
@@ -1567,8 +1589,19 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
1589
group_size=128,
1590
transpose_bm=True)
1591
else:
1592
+ # Pads the head_dim if necessary (for the underlying kernel)
1593
+ N, B, P = decode_q_nope.shape
1594
+ _, _, L = self.W_UK_T.shape
1595
+ if self.q_pad_num_heads is not None:
1596
+ decode_ql_nope = decode_q_nope.new_empty(
1597
+ (self.q_pad_num_heads, B, L))
1598
+ decode_ql_nope.resize_((N, B, L))
1599
+
1600
+ else:
1601
+ decode_ql_nope = decode_q_nope.new_empty((N, B, L))
1602
+
1603
# Multiply (N, B, P) x (N, P, L) -> (N, B, L)
1571
- decode_ql_nope = torch.bmm(decode_q_nope, self.W_UK_T)
1604
+ torch.bmm(decode_q_nope, self.W_UK_T, out=decode_ql_nope)
1605
# Convert from (N, B, L) to (B, N, L)
1606
decode_ql_nope = decode_ql_nope.transpose(0, 1)
1607
@@ -1603,5 +1636,5 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
1636
attn_out = cp_lse_ag_out_rs(attn_out, lse, get_dcp_group())
1637
1638
# v_up projection
1606
- output[:num_decode_tokens] = self._v_up_proj(attn_out)
1639
+ self._v_up_proj(attn_out, out=output[:num_decode_tokens])
1640
return output_padded
vllm/v1/attention/backends/mla/cutlass_mla.py
+16
-14
@@ -74,6 +74,8 @@ class SM100Workspace:
74
75
g_sm100_workspace = SM100Workspace(128 * 1024 * 1024) # 128MB
76
77
+MAX_HEADS = 128
78
+
79
80
class CutlassMLAImpl(MLACommonImpl[MLACommonMetadata]):
81
can_return_lse_for_decode: bool = True
@@ -92,10 +94,18 @@ class CutlassMLAImpl(MLACommonImpl[MLACommonMetadata]):
94
kv_sharing_target_layer_name: Optional[str],
95
# MLA Specific Arguments
96
**mla_args) -> None:
95
- super().__init__(num_heads, head_size, scale, num_kv_heads,
96
- alibi_slopes, sliding_window, kv_cache_dtype,
97
- logits_soft_cap, attn_type,
98
- kv_sharing_target_layer_name, **mla_args)
97
+ super().__init__(num_heads,
98
+ head_size,
99
+ scale,
100
+ num_kv_heads,
101
+ alibi_slopes,
102
+ sliding_window,
103
+ kv_cache_dtype,
104
+ logits_soft_cap,
105
+ attn_type,
106
+ kv_sharing_target_layer_name,
107
+ q_pad_num_heads=MAX_HEADS,
108
+ **mla_args)
109
110
unsupported_features = [alibi_slopes, sliding_window, logits_soft_cap]
111
if any(unsupported_features):
@@ -157,14 +167,6 @@ class CutlassMLAImpl(MLACommonImpl[MLACommonMetadata]):
167
168
MAX_HEADS = 128
169
assert H <= MAX_HEADS, f"H must be <= {MAX_HEADS}, but got {H}"
160
- if H < MAX_HEADS:
161
- q_nope_padded = q_nope.new_empty((B_q, MAX_HEADS, D_q_nope))
162
- q_nope_padded[:, :H] = q_nope
163
- q_nope = q_nope_padded
164
-
165
- q_pe_padded = q_pe.new_empty((B_q, MAX_HEADS, D_q_pe))
166
- q_pe_padded[:, :H] = q_pe
167
- q_pe = q_pe_padded
170
171
assert len(page_table.shape) == 2
172
B_block_table, block_num = page_table.shape
@@ -206,9 +208,9 @@ class CutlassMLAImpl(MLACommonImpl[MLACommonMetadata]):
208
)
209
210
if H < MAX_HEADS:
211
+ # Extract the subsets of the outputs
212
+ lse = lse[:, :H] if self.need_to_return_lse_for_decode else lse
213
out = out[:, :H]
210
- if self.need_to_return_lse_for_decode:
211
- lse = lse[:, :H].contiguous()
214
215
return out, lse
216