[Bugfix][KV Cache][MLA] Align packed block strides for V3.2 sparse MLA (#55528)
Signed-off-by: lz <145014769+200lz@users.noreply.github.com> Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Co-authored-by: OpenAI Codex <codex@openai.com> Co-authored-by: Lucas Wilkinson <lwilkins@redhat.com> Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
lz committed
Sep 25, 2026 at 14:33 UTC
29468dde8b515031dc6d4d9d06bf0a2fa0442098
6 files changed
+187
-27
tests/v1/attention/test_mla_backends.py
+3
-1
@@ -9,7 +9,7 @@ Known Issues:
9
"""
10
11
import sys
12
-from types import SimpleNamespace
12
+from types import MethodType, SimpleNamespace
13
14
import pytest
15
import torch
@@ -218,12 +218,14 @@ def test_mla_kv_cache_spec_uses_layer_cache_dtype(
218
cache_dtype: str, expected_quant_mode: KVQuantMode
219
):
220
layer = SimpleNamespace(
221
+ attn_backend=flashmla_module.FlashMLABackend,
222
kv_cache_dtype=cache_dtype,
223
head_size=576,
224
indexer=None,
225
non_causal_multi_token_decode=False,
226
sliding_window=None,
227
)
228
+ layer._uses_flat_kv_cache = MethodType(MLAAttention._uses_flat_kv_cache, layer)
229
vllm_config = SimpleNamespace(
230
cache_config=SimpleNamespace(block_size=64), model_config=None
231
)
tests/v1/attention/test_sparse_mla_kv_cache_layout.py
new
+119
@@ -0,0 +1,119 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Sparse MLA allocation and warmup must agree on physical addressing."""
4
+
5
+from types import MethodType, SimpleNamespace
6
+from unittest.mock import MagicMock
7
+
8
+import pytest
9
+import torch
10
+
11
+from vllm.config import CacheConfig
12
+from vllm.model_executor.layers.attention.mla_attention import MLAAttention
13
+from vllm.model_executor.models.deepseek_v2 import DeepseekV32IndexerCache
14
+from vllm.platforms import current_platform
15
+from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import (
16
+ FlashInferMLASparseTRTLLMBackend,
17
+)
18
+from vllm.v1.attention.backends.mla.flashmla_sparse import FlashMLASparseBackend
19
+from vllm.v1.attention.backends.mla.indexer import DeepseekV32IndexerBackend
20
+from vllm.v1.attention.backends.mla.sparse_utils import (
21
+ _CONVERT_REQ_INDEX_TO_GLOBAL_INDEX_KERNEL,
22
+ flat_kv_row_view,
23
+)
24
+from vllm.v1.attention.backends.utils import get_supported_kv_cache_layouts
25
+from vllm.v1.core.kv_cache_utils import get_kv_cache_config_from_groups
26
+from vllm.v1.kv_cache_interface import KVCacheGroupSpec, UniformTypeKVCacheSpecs
27
+from vllm.v1.kv_cache_layout import KVCacheLayout
28
+from vllm.v1.worker.utils import allocate_kv_cache
29
+
30
+pytestmark = pytest.mark.skip_global_cleanup
31
+
32
+
33
+@pytest.mark.parametrize(
34
+ "device",
35
+ [
36
+ "cpu",
37
+ pytest.param(
38
+ "cuda",
39
+ marks=pytest.mark.skipif(
40
+ not current_platform.is_cuda(), reason="CUDA only"
41
+ ),
42
+ ),
43
+ ],
44
+)
45
+@pytest.mark.parametrize(
46
+ "layout", [KVCacheLayout.BLHNC, KVCacheLayout.BLNHC, KVCacheLayout.LBNHC]
47
+)
48
+@pytest.mark.parametrize(
49
+ "backend,cache_dtype,sm100,stride_alignment",
50
+ [
51
+ (FlashInferMLASparseTRTLLMBackend, "auto", True, 1152),
52
+ (FlashInferMLASparseTRTLLMBackend, "fp8", True, 576),
53
+ (FlashMLASparseBackend, "auto", True, 1152),
54
+ (FlashMLASparseBackend, "fp8_ds_mla", False, None),
55
+ (FlashMLASparseBackend, "fp8_ds_mla", True, 656),
56
+ (FlashMLASparseBackend, "nvfp4_ds_mla", True, 352),
57
+ ],
58
+)
59
+def test_allocation_and_warmup_follow_addressing_mode(
60
+ monkeypatch, layout, backend, cache_dtype, sm100, stride_alignment, device
61
+):
62
+ monkeypatch.setattr(
63
+ current_platform, "is_device_capability_family", lambda family: sm100
64
+ )
65
+ config = SimpleNamespace(
66
+ cache_config=CacheConfig(block_size=64),
67
+ model_config=SimpleNamespace(dtype=torch.bfloat16),
68
+ attention_config=SimpleNamespace(hisparse_config=None),
69
+ kernel_config=SimpleNamespace(enable_jit_warmup=True),
70
+ )
71
+ config.cache_config.kv_cache_layout = layout.name
72
+ layer = SimpleNamespace(
73
+ _vllm_config=config,
74
+ attn_backend=backend,
75
+ kv_cache_dtype=cache_dtype,
76
+ head_size=576,
77
+ sliding_window=None,
78
+ indexer=None,
79
+ non_causal_multi_token_decode=False,
80
+ )
81
+ layer._uses_flat_kv_cache = MethodType(MLAAttention._uses_flat_kv_cache, layer)
82
+ spec = MLAAttention.get_kv_cache_spec(layer, config)
83
+ indexer = SimpleNamespace(
84
+ cache_config=config.cache_config, head_dim=132, dtype=torch.uint8
85
+ )
86
+ index_spec = DeepseekV32IndexerCache.get_kv_cache_spec(indexer, config)
87
+ assert spec.block_stride_alignment == stride_alignment
88
+ assert index_spec.block_stride_alignment is None
89
+ assert layout in get_supported_kv_cache_layouts(
90
+ [backend, DeepseekV32IndexerBackend]
91
+ )
92
+ specs = {"mla": spec, "indexer": index_spec}
93
+ group = KVCacheGroupSpec(
94
+ list(specs), UniformTypeKVCacheSpecs(block_size=64, kv_cache_specs=specs)
95
+ )
96
+ cache = get_kv_cache_config_from_groups(config, [group], 2**20)
97
+ views = allocate_kv_cache(cache, torch.device(device), layout)
98
+ register = MagicMock()
99
+ monkeypatch.setattr(
100
+ _CONVERT_REQ_INDEX_TO_GLOBAL_INDEX_KERNEL, "register_warmup", register
101
+ )
102
+ MLAAttention.bind_kv_cache(layer, views["mla"])
103
+ DeepseekV32IndexerCache.bind_kv_cache(indexer, views["indexer"])
104
+ assert indexer.kv_cache.stride(0) == views["indexer"].stride(0)
105
+ if stride_alignment is not None:
106
+ assert (
107
+ views["mla"].stride(0) * views["mla"].element_size() % stride_alignment == 0
108
+ )
109
+ if not layer._uses_flat_kv_cache():
110
+ register.assert_not_called()
111
+ else:
112
+ rows, stride = flat_kv_row_view(layer.kv_cache, 64)
113
+ register.assert_called_once_with(config, block_stride_rows=stride)
114
+ layer.kv_cache[1, 0].fill_(3)
115
+ torch.testing.assert_close(rows[stride], layer.kv_cache[1, 0])
116
+ if stride_alignment is None:
117
+ assert cache.kv_cache_tensors[0].size == cache.num_blocks * sum(
118
+ item.page_size_bytes for item in specs.values()
119
+ )
tests/v1/core/test_contiguous_kv_packing.py
+27
@@ -10,6 +10,7 @@ block); the allocation is the same either way.
10
"""
11
12
from dataclasses import replace
13
+from math import lcm
14
from unittest.mock import MagicMock
15
16
import pytest
@@ -26,9 +27,11 @@ from vllm.v1.core.kv_cache_utils import (
27
get_kv_cache_groups,
28
resolve_kv_cache_block_sizes,
29
)
30
+from vllm.v1.hisparse.layout import _build_hisparse_kv_cache_tensors
31
from vllm.v1.kv_cache_interface import (
32
CircularBufferSpec,
33
FullAttentionSpec,
34
+ HiSparseHotSpec,
35
KVCacheGroupSpec,
36
KVCacheLayout,
37
KVCacheSpec,
@@ -79,6 +82,30 @@ def _mock_vllm_config(layout: str | None):
82
return config
83
84
85
+def test_packed_alignment_preserves_hot_pages_and_generic_stride():
86
+ """One shared stride must satisfy both constraints, regardless of rounding order."""
87
+ spec = replace(_mla(128), block_stride_alignment=1152)
88
+ hot = HiSparseHotSpec(block_size=16, page_size=5120, blocks_per_request=2)
89
+ groups = [KVCacheGroupSpec(["attention"], spec), KVCacheGroupSpec(["hot"], hot)]
90
+ alignment = lcm(1152, hot.page_size_bytes)
91
+ stride = _pool_bytes_per_block(groups)
92
+ assert stride == alignment
93
+ layout = KVCacheLayout.BLHNC
94
+ tensors = _build_hisparse_kv_cache_tensors(groups, 3, 3 * stride, layout, stride)
95
+ assert {tensor.block_stride for tensor in tensors} == {alignment}
96
+ assert {tensor.size for tensor in tensors} == {3 * alignment}
97
+
98
+
99
+@pytest.mark.parametrize("spec", [_full(), _mla(128)])
100
+def test_block_alignment_survives_merge_and_rejects_incompatible_specs(spec):
101
+ aligned = replace(spec, block_stride_alignment=512)
102
+ assert type(spec).merge([aligned, aligned]).block_stride_alignment == 512
103
+ with pytest.raises(AssertionError):
104
+ type(spec).merge([aligned, spec])
105
+ with pytest.raises(ValueError, match="must be positive"):
106
+ replace(spec, block_stride_alignment=0)
107
+
108
+
109
def _pages(groups) -> dict[str, int]:
110
return {
111
name: group.kv_cache_spec.kv_cache_specs[name].page_size_bytes
vllm/model_executor/layers/attention/mla_attention.py
+18
-8
@@ -735,13 +735,7 @@ class MLAAttention(nn.Module, AttentionLayerBase):
735
self.kv_cache = kv_cache.squeeze(1)
736
if (
737
self._vllm_config.kernel_config.enable_jit_warmup
738
- and self.attn_backend.get_name()
739
- in (
740
- "FLASHMLA_SPARSE",
741
- "FLASHINFER_MLA_SPARSE",
742
- "FLASHINFER_MLA_SPARSE_SM120",
743
- "DEEPSEEK_V32_INDEXER",
744
- )
738
+ and self._uses_flat_kv_cache()
739
):
740
from vllm.v1.attention.backends.mla.sparse_utils import (
741
_CONVERT_REQ_INDEX_TO_GLOBAL_INDEX_KERNEL,
@@ -754,6 +748,13 @@ class MLAAttention(nn.Module, AttentionLayerBase):
748
block_stride_rows=self.kv_cache.stride(0) // row_width,
749
)
750
751
+ def _uses_flat_kv_cache(self) -> bool:
752
+ backend = self.attn_backend.get_name()
753
+ return backend == "FLASHINFER_MLA_SPARSE" or (
754
+ backend == "FLASHMLA_SPARSE"
755
+ and self.kv_cache_dtype not in ("fp8_ds_mla", "nvfp4_ds_mla")
756
+ )
757
+
758
@property
759
def chunked_prefill_workspace_size(self) -> int:
760
if self._chunked_prefill_workspace_size is None:
@@ -1365,11 +1366,20 @@ class MLAAttention(nn.Module, AttentionLayerBase):
1366
**common_kwargs,
1367
sliding_window=self.sliding_window,
1368
)
1368
- return MLAAttentionSpec(
1369
+ spec = MLAAttentionSpec(
1370
**common_kwargs,
1371
is_index_group_leader=self.indexer is not None,
1372
non_causal_multi_token_decode=self.non_causal_multi_token_decode,
1373
)
1374
+ # SM100 FlashMLA paged kernels also express TMA coordinates in token rows.
1375
+ uses_tma_rows = (
1376
+ self.attn_backend.get_name() == "FLASHMLA_SPARSE"
1377
+ and self.kv_cache_dtype in ("fp8_ds_mla", "nvfp4_ds_mla")
1378
+ and current_platform.is_device_capability_family(100)
1379
+ )
1380
+ if self._uses_flat_kv_cache() or uses_tma_rows:
1381
+ spec = replace(spec, block_stride_alignment=spec.state_content_size_bytes)
1382
+ return spec
1383
1384
def _v_up_proj(self, x: torch.Tensor, out: torch.Tensor):
1385
# Convert from (B, N, L) to (N, B, L)
vllm/v1/core/kv_cache_utils.py
+5
-11
@@ -1608,23 +1608,17 @@ def _get_kv_cache_bytes_per_block(
1608
for group in kv_cache_groups
1609
)
1610
assert bytes_per_block > 0
1611
- hot_page_sizes = [
1611
+ alignments = [
1612
group.kv_cache_spec.page_size_bytes
1613
for group in kv_cache_groups
1614
if isinstance(group.kv_cache_spec, HiSparseHotSpec)
1615
]
1616
- if hot_page_sizes:
1617
- bytes_per_block = round_up(bytes_per_block, math.lcm(*hot_page_sizes))
1618
- stride_alignments = [
1619
- spec.block_stride_alignment
1616
+ alignments.extend(
1617
+ _get_per_layer_spec(group, layer_name).block_stride_alignment or 1
1618
for group in kv_cache_groups
1619
for layer_name in group.layer_names
1622
- if isinstance(spec := _get_per_layer_spec(group, layer_name), MLAAttentionSpec)
1623
- and spec.block_stride_alignment
1624
- ]
1625
- if stride_alignments:
1626
- bytes_per_block = round_up(bytes_per_block, math.lcm(*stride_alignments))
1627
- return bytes_per_block
1620
+ )
1621
+ return round_up(bytes_per_block, math.lcm(*alignments))
1622
1623
1624
def validate_kv_cache_layout(
vllm/v1/kv_cache_interface.py
+15
-7
@@ -6,7 +6,7 @@ from __future__ import annotations
6
import copy
7
from collections import Counter
8
from collections.abc import Collection, Sequence
9
-from dataclasses import dataclass, fields, replace
9
+from dataclasses import dataclass, field, fields, replace
10
from enum import Enum, IntEnum
11
from fractions import Fraction
12
from functools import cached_property
@@ -159,6 +159,13 @@ class KVCacheSpec:
159
# number of tokens in a block
160
block_size: int
161
162
+ block_stride_alignment: int | None = field(default=None, kw_only=True)
163
+ """Required byte alignment between physical blocks, including packed layers."""
164
+
165
+ def __post_init__(self):
166
+ if self.block_stride_alignment is not None and self.block_stride_alignment <= 0:
167
+ raise ValueError("block_stride_alignment must be positive")
168
+
169
@property
170
def prefix_cacheable(self) -> bool:
171
"""Whether this spec's group participates in prefix caching."""
@@ -492,6 +499,7 @@ class AttentionSpec(KVCacheSpec):
499
token (Whisper block pooling: ``Fraction(1, block_pool_size)``)."""
500
501
def __post_init__(self):
502
+ super().__post_init__()
503
if self.head_size_v is None:
504
object.__setattr__(self, "head_size_v", self.head_size)
505
@@ -598,6 +606,7 @@ class FullAttentionSpec(AttentionSpec):
606
)
607
merged_spec = cls(
608
block_size=specs[0].block_size,
609
+ block_stride_alignment=specs[0].block_stride_alignment,
610
num_kv_heads=specs[0].num_kv_heads,
611
head_size=specs[0].head_size,
612
head_size_v=specs[0].head_size_v,
@@ -648,12 +657,6 @@ class MLAAttentionSpec(FullAttentionSpec):
657
is_index_group_leader: bool = False
658
storage_block_size: int | None = None
659
"""Token width used to view storage when it differs from the kernel block."""
651
- block_stride_alignment: int | None = None
652
- """Required alignment, in bytes, of the distance between consecutive
653
- blocks of this cache. In block-major layouts that distance is the whole
654
- block (all layers' pages), so the allocator rounds the block up to it.
655
- DeepGEMM's paged sparse MQA-logits kernels address pages as
656
- ``base + page * stride`` and need it 512B-aligned."""
660
# Group capability enabled when any member flattens a non-causal query block
661
# into decode rows. Runtime metadata still selects causal vs. non-causal mode.
662
non_causal_multi_token_decode: bool = False
@@ -751,6 +754,7 @@ class RSWASpec(FullAttentionSpec):
754
base = FullAttentionSpec.merge(specs) # type: ignore[arg-type]
755
return cls(
756
block_size=base.block_size,
757
+ block_stride_alignment=base.block_stride_alignment,
758
num_kv_heads=base.num_kv_heads,
759
head_size=base.head_size,
760
head_size_v=base.head_size_v,
@@ -945,6 +949,7 @@ class SlidingWindowMLASpec(SlidingWindowSpec):
949
sliding_window_set = set(spec.sliding_window for spec in specs)
950
extra_retained_set = set(spec.extra_retained_tokens for spec in specs)
951
bounded_replay_set = set(spec.bounded_replay for spec in specs)
952
+ block_stride_alignment_set = {spec.block_stride_alignment for spec in specs}
953
assert (
954
len(cache_dtype_str_set) == 1
955
and len(tokens_per_state_set) == 1
@@ -952,6 +957,7 @@ class SlidingWindowMLASpec(SlidingWindowSpec):
957
and len(sliding_window_set) == 1
958
and len(extra_retained_set) == 1
959
and len(bounded_replay_set) == 1
960
+ and len(block_stride_alignment_set) == 1
961
), (
962
"All attention layers in the same KV cache group must use the same "
963
"quantization method, tokens per state, model version, sliding "
@@ -959,6 +965,7 @@ class SlidingWindowMLASpec(SlidingWindowSpec):
965
)
966
return cls(
967
block_size=specs[0].block_size,
968
+ block_stride_alignment=block_stride_alignment_set.pop(),
969
num_kv_heads=specs[0].num_kv_heads,
970
head_size=specs[0].head_size,
971
dtype=specs[0].dtype,
@@ -1175,6 +1182,7 @@ class SinkFullAttentionSpec(FullAttentionSpec):
1182
head_size=specs[0].head_size,
1183
head_size_v=specs[0].head_size_v,
1184
sink_len=specs[0].sink_len,
1185
+ block_stride_alignment=specs[0].block_stride_alignment,
1186
dtype=specs[0].dtype,
1187
kv_quant_mode=specs[0].kv_quant_mode,
1188
page_size_padded=specs[0].page_size_padded,