[Attention] Fix MLA prefill workspace allocation size (#51733)

Signed-off-by: wzhao18 <wzhao18.sz@gmail.com>

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):

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(

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