[Attention] Fix MLA prefill workspace allocation size (#51733)
Signed-off-by: wzhao18 <wzhao18.sz@gmail.com>
Wei Zhao committed
Aug 11, 2026 at 01:34 UTC
608c12473f3f93b5452ac3c4fa13fb280b7da674
3 files changed
+7
-11
tests/distributed/test_dcp_direct_a2a_lse_reduce.py
+5
-2
@@ -494,9 +494,12 @@ def test_dcp_chunk_workspace_alignment_covers_interleave():
494
config.cache_config.block_size = 32
495
config.parallel_config.decode_context_parallel_size = 8
496
config.parallel_config.cp_kv_cache_interleave_size = 8
497
- config.scheduler_config.max_num_seqs = 3
497
499
- assert align_mla_chunked_context_workspace_size(config, 100) == 192
498
+ # Alignment is lcm(block_size, dcp_size * interleave_size) = 64, and the
499
+ # workspace only has to hold a single aligned chunk step, independent of
500
+ # max_num_seqs.
501
+ assert align_mla_chunked_context_workspace_size(config, 100) == 128
502
+ assert align_mla_chunked_context_workspace_size(config, 8) == 64
503
504
505
def test_sparse_mla_builder_initializes_dcp_manager(monkeypatch):
vllm/model_executor/layers/attention/mla_attention.py
+1
-5
@@ -1705,11 +1705,7 @@ def align_mla_chunked_context_workspace_size(
1705
parallel_config.decode_context_parallel_size
1706
* parallel_config.cp_kv_cache_interleave_size,
1707
)
1708
- workspace_size = max(
1709
- workspace_size,
1710
- vllm_config.scheduler_config.max_num_seqs * alignment,
1711
- )
1712
- return round_up(workspace_size, alignment)
1708
+ return round_up(max(workspace_size, alignment), alignment)
1709
1710
1711
def build_mla_chunked_context_metadata(
vllm/model_executor/layers/attention/sparse_mla_attention.py
+1
-4
@@ -196,10 +196,7 @@ class SparseMLACommonMetadataBuilder(AttentionMetadataBuilder[T]):
196
64 * 1024,
197
scheduler_config.max_num_seqs * topk_tokens,
198
)
199
- workspace_size = max(
200
- workspace_size,
201
- scheduler_config.max_num_seqs * cache_config.block_size,
202
- )
199
+ workspace_size = max(workspace_size, cache_config.block_size)
200
if vllm_config.parallel_config.decode_context_parallel_size > 1:
201
return align_mla_chunked_context_workspace_size(vllm_config, workspace_size)
202
return workspace_size