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, 19:00 UTC
[Perf] Optimize MLA `compute_prefill_context` memory allocation (#42460)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
May 12, 2026 at 19:23 UTC
3d635c58c058bc3362c116298b4fa643abaf9138
1 file changed
+16
-12
vllm/model_executor/layers/attention/mla_attention.py
+16
-12
@@ -2051,6 +2051,7 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
2051
use_fp8_prefill = prefill_metadata.q_data_type == current_platform.fp8_dtype()
2052
2053
output = None
2054
+ merge_output = None
2055
iters = len(prefill_metadata.chunked_context.seq_tot)
2056
workspace = prefill_metadata.chunked_context.workspace
2057
@@ -2126,18 +2127,19 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
2127
output = attn_output
2128
output_lse = attn_softmax_lse
2129
else:
2129
- output_tmp = torch.empty_like(output)
2130
- output_lse_tmp = torch.empty_like(output_lse)
2130
+ if merge_output is None:
2131
+ merge_output = torch.empty_like(output)
2132
+ merge_output_lse = torch.empty_like(output_lse)
2133
merge_attn_states(
2132
- output=output_tmp,
2133
- output_lse=output_lse_tmp,
2134
+ output=merge_output,
2135
+ output_lse=merge_output_lse,
2136
prefix_output=output,
2137
prefix_lse=output_lse,
2138
suffix_output=attn_output,
2139
suffix_lse=attn_softmax_lse,
2140
)
2139
- output = output_tmp
2140
- output_lse = output_lse_tmp
2141
+ output, merge_output = merge_output, output
2142
+ output_lse, merge_output_lse = merge_output_lse, output_lse
2143
2144
return output, output_lse
2145
@@ -2161,6 +2163,7 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
2163
assert prefill_metadata.chunked_context.chunk_size is not None
2164
2165
output = None
2166
+ merge_output = None
2167
iters = len(prefill_metadata.chunked_context.seq_tot)
2168
workspace = prefill_metadata.chunked_context.workspace
2169
@@ -2232,18 +2235,19 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
2235
output = attn_output
2236
output_lse = attn_softmax_lse
2237
else:
2235
- output_tmp = torch.empty_like(output)
2236
- output_lse_tmp = torch.empty_like(output_lse)
2238
+ if merge_output is None:
2239
+ merge_output = torch.empty_like(output)
2240
+ merge_output_lse = torch.empty_like(output_lse)
2241
merge_attn_states(
2238
- output=output_tmp,
2239
- output_lse=output_lse_tmp,
2242
+ output=merge_output,
2243
+ output_lse=merge_output_lse,
2244
prefix_output=output,
2245
prefix_lse=output_lse,
2246
suffix_output=attn_output,
2247
suffix_lse=attn_softmax_lse,
2248
)
2245
- output = output_tmp
2246
- output_lse = output_lse_tmp
2249
+ output, merge_output = merge_output, output
2250
+ output_lse, merge_output_lse = merge_output_lse, output_lse
2251
2252
return output, output_lse
2253