[Attention] Vectorize sparse MLA mask loads (#52217)
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Matthew Bonanni committed
Aug 18, 2026 at 23:32 UTC
8e46accab22bc864f3ddbe8276d7c7cd2eb71239
3 files changed
+46
-11
tests/v1/attention/test_sparse_mla_mask.py
+22
@@ -8,6 +8,7 @@ pytest.importorskip("cutlass")
8
9
from vllm.model_executor.layers.attention.sparse_mla_attention import (
10
_build_topk_mask,
11
+ _topk_mask_shape,
12
)
13
14
@@ -27,3 +28,24 @@ def test_build_topk_mask_single_request_matches_generic_path() -> None:
28
29
torch.testing.assert_close(single_req[0, 0], generic[0, 0])
30
torch.testing.assert_close(single_req[0, 1], generic[1, 0])
31
+
32
+
33
+def test_topk_mask_rows_are_aligned_for_vectorized_loads() -> None:
34
+ assert _topk_mask_shape(2, 129, 129) == (2, 256, 8)
35
+ assert _topk_mask_shape(2, 129, 129, reserve_key_starts_word=True) == (
36
+ 2,
37
+ 256,
38
+ 8,
39
+ )
40
+
41
+
42
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
43
+def test_build_topk_mask_preserves_aligned_row_storage() -> None:
44
+ shape = _topk_mask_shape(1, 1, 129)
45
+ out = torch.zeros(shape, dtype=torch.int32, device="cuda")
46
+ topk = torch.tensor([[0, 128]], dtype=torch.int32, device="cuda")
47
+
48
+ mask = _build_topk_mask([topk], [1], 1, 129, out)
49
+
50
+ assert mask.shape == (1, 1, 8)
51
+ assert mask[0, 0].tolist() == [1, 0, 0, 0, 1, 0, 0, 0]
vllm/model_executor/layers/attention/sparse_mla_attention.py
+5
-4
@@ -58,6 +58,7 @@ def _topk_mask_shape(
58
tile_m = 128 if max_query_len <= 128 else 256
59
padded_q_len = triton.cdiv(max_query_len, tile_m) * tile_m
60
num_words = triton.cdiv(max_key_len, 32) + int(reserve_key_starts_word)
61
+ num_words = triton.cdiv(num_words, 4) * 4
62
return batch_size, padded_q_len, num_words
63
64
@@ -432,12 +433,12 @@ def _build_topk_mask(
433
max_seq_len: int,
434
out: torch.Tensor,
435
) -> torch.Tensor:
435
- """Build a bit-packed top-k mask into ``out[:B, :max_Q, :num_words]``."""
436
+ """Build a bit-packed top-k mask while preserving padded row storage."""
437
batch_size = len(q_lens)
438
num_words = (max_seq_len + 31) // 32
439
total_rows = batch_size * max_q_len
440
if total_rows == 0:
440
- return out[:batch_size, :max_q_len, :num_words]
441
+ return out[:batch_size, :max_q_len]
442
443
total_q = sum(q_lens)
444
mask_row_stride = out.stride(-2)
@@ -457,7 +458,7 @@ def _build_topk_mask(
458
BLOCK_TOPK=triton.next_power_of_2(num_topk),
459
BLOCK_WORDS=block_words,
460
)
460
- return out[:1, :max_q_len, :num_words]
461
+ return out[:1, :max_q_len]
462
463
topk_packed = torch.cat(topk_indices_per_req, dim=0)
464
num_topk = topk_packed.shape[1]
@@ -478,7 +479,7 @@ def _build_topk_mask(
479
BLOCK_TOPK=triton.next_power_of_2(num_topk),
480
BLOCK_WORDS=block_words,
481
)
481
- return out[:batch_size, :max_q_len, :num_words]
482
+ return out[:batch_size, :max_q_len]
483
484
485
class SparseMLACommonImpl(MLACommonBaseImpl[T], Generic[T]):
vllm/model_executor/layers/attention/sparse_mla_mask.py
+19
-7
@@ -20,15 +20,27 @@ def dense_mask_mod(
20
batch_idx = utils.ssa_to_scalar(batch)
21
q_idx = utils.ssa_to_scalar(q_idx)
22
kv_idx = utils.ssa_to_scalar(kv_idx)
23
- word_idx = kv_idx >> 5
24
- bit_idx = cutlass.Uint32(kv_idx & 31)
25
- word = dense_mask[batch_idx, q_idx, word_idx]
26
- result = cute.make_rmem_tensor(1, dtype=cutlass.Uint32)
27
- result[0] = utils.shr_u32(cutlass.Uint32(word), bit_idx)
28
- return result.load()
23
+ batch_stride, query_stride, _ = dense_mask.stride
24
+ aligned_mask = cute.make_tensor(
25
+ dense_mask.iterator,
26
+ cute.make_layout(
27
+ dense_mask.shape,
28
+ stride=(
29
+ cute.assume(batch_stride, divby=4),
30
+ cute.assume(query_stride, divby=4),
31
+ 1,
32
+ ),
33
+ ),
34
+ )
35
+ mask_row = aligned_mask[batch_idx, q_idx, None]
36
+ mask_chunks = cute.flat_divide(mask_row, (4,))
37
+ mask_chunk = mask_chunks[None, (kv_idx >> 5) >> 2]
38
+ loaded = cute.make_rmem_tensor_like(mask_chunk)
39
+ cute.autovec_copy(mask_chunk, loaded)
40
+ return cute.recast_tensor(loaded, cutlass.Uint32).load()
41
42
31
-dense_mask_mod.__vec_size__ = 32
43
+dense_mask_mod.__vec_size__ = 128
44
45
46
@cute.jit