[Bugfix][Attention] Guard sparse MLA masked MHA workspace (#50906)
Signed-off-by: yimdev <5779256+yimdev@users.noreply.github.com> Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> Co-authored-by: OpenAI Codex <codex@openai.com> Co-authored-by: Matthew Bonanni <mbonanni@redhat.com>
yimdev committed
Aug 4, 2026 at 21:59 UTC
199644d410dbfd94255d6b0cde08479858032266
3 files changed
+115
-13
tests/v1/attention/test_sparse_mla_backends.py
+37
@@ -21,6 +21,11 @@ from tests.v1.attention.utils import (
21
)
22
from vllm import _custom_ops as ops
23
from vllm.config import set_current_vllm_config
24
+from vllm.model_executor.layers.attention.sparse_mla_attention import (
25
+ GLOBAL_TOPK_MASK_MAX_BYTES,
26
+ _masked_mha_workspace_fits,
27
+ _topk_mask_shape,
28
+)
29
from vllm.model_executor.layers.linear import ColumnParallelLinear
30
from vllm.platforms import current_platform
31
@@ -859,6 +864,38 @@ def test_split_prefill_chunks(seq_lens, max_buf, expected):
864
assert out == expected
865
866
867
+@pytest.mark.parametrize(
868
+ ("max_query_len", "expected"),
869
+ [(32768, True), (33024, False)],
870
+)
871
+def test_masked_mha_workspace_fits_single_request_boundary(max_query_len, expected):
872
+ """A 32K prefill needs the default workspace exactly; shrinking it would
873
+ push a supported request onto MQA."""
874
+ assert (
875
+ _masked_mha_workspace_fits(
876
+ batch_size=1,
877
+ max_query_len=max_query_len,
878
+ context_chunk_max_seq_lens=None,
879
+ workspace_numel=GLOBAL_TOPK_MASK_MAX_BYTES // torch.int32.itemsize,
880
+ )
881
+ is expected
882
+ )
883
+
884
+
885
+def test_masked_mha_workspace_fits_accounts_for_batch_and_context():
886
+ """Request count and context chunk length are independent multipliers."""
887
+ base = dict(batch_size=2, max_query_len=2048, context_chunk_max_seq_lens=[2048])
888
+ exact = math.prod(_topk_mask_shape(2, 2048, 2048))
889
+
890
+ assert _masked_mha_workspace_fits(**base, workspace_numel=exact)
891
+ assert not _masked_mha_workspace_fits(
892
+ **{**base, "batch_size": 3}, workspace_numel=exact
893
+ )
894
+ assert not _masked_mha_workspace_fits(
895
+ **{**base, "context_chunk_max_seq_lens": [4096]}, workspace_numel=exact
896
+ )
897
+
898
+
899
PREFILL_BATCH_SPECS = {
900
"short_dense_mha": BatchSpec(seq_lens=[64, 128], query_lens=[64, 128]),
901
"short_context_dense_mha": BatchSpec(seq_lens=[128, 160], query_lens=[64, 32]),
vllm/model_executor/layers/attention/mla_attention.py
+1
@@ -769,6 +769,7 @@ class MLAAttention(nn.Module, AttentionLayerBase):
769
query_len=prefill.max_query_len,
770
seq_len=prefill_max_seq_len,
771
)
772
+ and self.impl.masked_mha_workspace_fits(prefill) # type: ignore[attr-defined]
773
)
774
use_mha = (use_dense_mha or use_masked_mha) and not (
775
self._vllm_config.attention_config.sparse_mla_force_mqa
vllm/model_executor/layers/attention/sparse_mla_attention.py
+77
-13
@@ -2,6 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
"""Shared MHA implementation and metadata builder for sparse MLA backends."""
4
5
+import math
6
from shutil import which
7
from typing import TYPE_CHECKING, ClassVar, Generic, TypeVar, cast
8
@@ -40,7 +41,38 @@ logger = init_logger(__name__)
41
42
T = TypeVar("T", bound=AttentionMetadata)
43
43
-GLOBAL_TOPK_MASK_MAX_BYTES = 64 * 1024 * 1024 # 64 MiB
44
+GLOBAL_TOPK_MASK_MAX_BYTES = 128 * 1024 * 1024 # 128 MiB
45
+
46
+
47
+def _topk_mask_shape(
48
+ batch_size: int,
49
+ max_query_len: int,
50
+ max_key_len: int,
51
+ reserve_key_starts_word: bool = False,
52
+) -> tuple[int, int, int]:
53
+ """Shape of a bit-packed top-k mask, shared by every site that builds one."""
54
+ tile_m = 128 if max_query_len <= 128 else 256
55
+ padded_q_len = triton.cdiv(max_query_len, tile_m) * tile_m
56
+ num_words = triton.cdiv(max_key_len, 32) + int(reserve_key_starts_word)
57
+ return batch_size, padded_q_len, num_words
58
+
59
+
60
+def _masked_mha_workspace_fits(
61
+ batch_size: int,
62
+ max_query_len: int,
63
+ context_chunk_max_seq_lens: list[int] | None,
64
+ workspace_numel: int,
65
+) -> bool:
66
+ """Return whether the suffix and per-context-chunk masks fit the workspace.
67
+
68
+ The global mask is excluded: it always needs more, and has its own check.
69
+ """
70
+ max_key_len = max(
71
+ max_query_len,
72
+ max(context_chunk_max_seq_lens or (), default=0),
73
+ )
74
+ needed = math.prod(_topk_mask_shape(batch_size, max_query_len, max_key_len))
75
+ return needed <= workspace_numel
76
77
78
def _is_masked_mha_available(
@@ -504,6 +536,30 @@ class SparseMLACommonImpl(MLACommonBaseImpl[T], Generic[T]):
536
kv_cache_dtype=kv_cache_dtype,
537
)
538
539
+ @staticmethod
540
+ def masked_mha_workspace_fits(prefill: MLACommonPrefillMetadata) -> bool:
541
+ """Whether this prefill batch's top-k masks fit the workspace."""
542
+ workspace = prefill.topk_mask_workspace
543
+ if workspace is None or prefill.query_lens_cpu is None:
544
+ return False
545
+ fits = _masked_mha_workspace_fits(
546
+ batch_size=len(prefill.query_lens_cpu),
547
+ max_query_len=prefill.max_query_len,
548
+ context_chunk_max_seq_lens=(
549
+ None
550
+ if prefill.chunked_context is None
551
+ else prefill.chunked_context.max_seq_lens
552
+ ),
553
+ workspace_numel=workspace.numel(),
554
+ )
555
+ if not fits:
556
+ logger.warning_once(
557
+ "Sparse MLA top-k mask workspace (%d MiB) is too small for some "
558
+ "prefill batches; those fall back to slower sparse MQA.",
559
+ workspace.numel() * torch.int32.itemsize // (1024 * 1024),
560
+ )
561
+ return fits
562
+
563
@staticmethod
564
def _slice_topk_per_req(
565
topk_all: torch.Tensor,
@@ -554,12 +610,14 @@ class SparseMLACommonImpl(MLACommonBaseImpl[T], Generic[T]):
610
mask is too large, signalling the caller to fall back to per-chunk
611
index remapping.
612
"""
557
- batch_size = len(q_lens)
558
- tile_m = 128 if max_query_len <= 128 else 256
559
- padded_q_len = triton.cdiv(max_query_len, tile_m) * tile_m
560
- num_words_padded = (max_seq_len + 31) // 32 + 1
613
+ batch_size, padded_q_len, num_words_padded = _topk_mask_shape(
614
+ len(q_lens),
615
+ max_query_len,
616
+ max_seq_len,
617
+ reserve_key_starts_word=True,
618
+ )
619
needed = batch_size * padded_q_len * num_words_padded
562
- if needed * torch.int32.itemsize > GLOBAL_TOPK_MASK_MAX_BYTES:
620
+ if needed > topk_mask_workspace.numel():
621
return None
622
623
mask = topk_mask_workspace[:needed].view(
@@ -597,15 +655,21 @@ class SparseMLACommonImpl(MLACommonBaseImpl[T], Generic[T]):
655
)
656
from vllm.vllm_flash_attn import flash_attn_varlen_func
657
600
- tile_m = 128 if max_seqlen_q <= 128 else 256
601
- padded_q_len = triton.cdiv(max_seqlen_q, tile_m) * tile_m
658
if dense_mask is None:
603
- batch_size = len(q_lens)
604
- num_words = (max_seqlen_k + 31) // 32
659
assert topk_mask_workspace is not None
606
- workspace_3d = topk_mask_workspace[
607
- : batch_size * padded_q_len * num_words
608
- ].view(batch_size, padded_q_len, num_words)
660
+ batch_size, padded_q_len, num_words = _topk_mask_shape(
661
+ len(q_lens), max_seqlen_q, max_seqlen_k
662
+ )
663
+ words_needed = batch_size * padded_q_len * num_words
664
+ if words_needed > topk_mask_workspace.numel():
665
+ raise ValueError(
666
+ f"Sparse MLA top-k mask needs {words_needed} int32 words (batch="
667
+ f"{len(q_lens)}, q={max_seqlen_q}, k={max_seqlen_k}) but the "
668
+ f"workspace holds {topk_mask_workspace.numel()}."
669
+ )
670
+ workspace_3d = topk_mask_workspace[:words_needed].view(
671
+ batch_size, padded_q_len, num_words
672
+ )
673
dense_mask = _build_topk_mask(
674
topk_per_req,
675
q_lens,