From 608c12473f3f93b5452ac3c4fa13fb280b7da674 Mon Sep 17 00:00:00 2001 From: Wei Zhao <51183510+wzhao18@users.noreply.github.com> Date: Tue, 11 Aug 2026 01:34:30 -0400 Subject: [PATCH] [Attention] Fix MLA prefill workspace allocation size (#51733) Signed-off-by: wzhao18 --- tests/distributed/test_dcp_direct_a2a_lse_reduce.py | 7 +++++-- vllm/model_executor/layers/attention/mla_attention.py | 6 +----- .../layers/attention/sparse_mla_attention.py | 5 +---- 3 files changed, 7 insertions(+), 11 deletions(-) diff --git a/tests/distributed/test_dcp_direct_a2a_lse_reduce.py b/tests/distributed/test_dcp_direct_a2a_lse_reduce.py index b05a54825d..13561bac78 100644 --- a/tests/distributed/test_dcp_direct_a2a_lse_reduce.py +++ b/tests/distributed/test_dcp_direct_a2a_lse_reduce.py @@ -494,9 +494,12 @@ def test_dcp_chunk_workspace_alignment_covers_interleave(): config.cache_config.block_size = 32 config.parallel_config.decode_context_parallel_size = 8 config.parallel_config.cp_kv_cache_interleave_size = 8 - config.scheduler_config.max_num_seqs = 3 - assert align_mla_chunked_context_workspace_size(config, 100) == 192 + # Alignment is lcm(block_size, dcp_size * interleave_size) = 64, and the + # workspace only has to hold a single aligned chunk step, independent of + # max_num_seqs. + assert align_mla_chunked_context_workspace_size(config, 100) == 128 + assert align_mla_chunked_context_workspace_size(config, 8) == 64 def test_sparse_mla_builder_initializes_dcp_manager(monkeypatch): diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 1e7647d163..0b59956543 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -1705,11 +1705,7 @@ def align_mla_chunked_context_workspace_size( parallel_config.decode_context_parallel_size * parallel_config.cp_kv_cache_interleave_size, ) - workspace_size = max( - workspace_size, - vllm_config.scheduler_config.max_num_seqs * alignment, - ) - return round_up(workspace_size, alignment) + return round_up(max(workspace_size, alignment), alignment) def build_mla_chunked_context_metadata( diff --git a/vllm/model_executor/layers/attention/sparse_mla_attention.py b/vllm/model_executor/layers/attention/sparse_mla_attention.py index 132bb5f920..a619b74035 100644 --- a/vllm/model_executor/layers/attention/sparse_mla_attention.py +++ b/vllm/model_executor/layers/attention/sparse_mla_attention.py @@ -196,10 +196,7 @@ class SparseMLACommonMetadataBuilder(AttentionMetadataBuilder[T]): 64 * 1024, scheduler_config.max_num_seqs * topk_tokens, ) - workspace_size = max( - workspace_size, - scheduler_config.max_num_seqs * cache_config.block_size, - ) + workspace_size = max(workspace_size, cache_config.block_size) if vllm_config.parallel_config.decode_context_parallel_size > 1: return align_mla_chunked_context_workspace_size(vllm_config, workspace_size) return workspace_size