[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>

tests/v1/attention/test_mla_backends.py

@@ -9,7 +9,7 @@ Known Issues:
 """
 
 import sys
-from types import SimpleNamespace
+from types import MethodType, SimpleNamespace
 
 import pytest
 import torch
@@ -218,12 +218,14 @@ def test_mla_kv_cache_spec_uses_layer_cache_dtype(
     cache_dtype: str, expected_quant_mode: KVQuantMode
 ):
     layer = SimpleNamespace(
+        attn_backend=flashmla_module.FlashMLABackend,
         kv_cache_dtype=cache_dtype,
         head_size=576,
         indexer=None,
         non_causal_multi_token_decode=False,
         sliding_window=None,
     )
+    layer._uses_flat_kv_cache = MethodType(MLAAttention._uses_flat_kv_cache, layer)
     vllm_config = SimpleNamespace(
         cache_config=SimpleNamespace(block_size=64), model_config=None
     )

tests/v1/attention/test_sparse_mla_kv_cache_layout.py

@@ -0,0 +1,119 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Sparse MLA allocation and warmup must agree on physical addressing."""
+
+from types import MethodType, SimpleNamespace
+from unittest.mock import MagicMock
+
+import pytest
+import torch
+
+from vllm.config import CacheConfig
+from vllm.model_executor.layers.attention.mla_attention import MLAAttention
+from vllm.model_executor.models.deepseek_v2 import DeepseekV32IndexerCache
+from vllm.platforms import current_platform
+from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import (
+    FlashInferMLASparseTRTLLMBackend,
+)
+from vllm.v1.attention.backends.mla.flashmla_sparse import FlashMLASparseBackend
+from vllm.v1.attention.backends.mla.indexer import DeepseekV32IndexerBackend
+from vllm.v1.attention.backends.mla.sparse_utils import (
+    _CONVERT_REQ_INDEX_TO_GLOBAL_INDEX_KERNEL,
+    flat_kv_row_view,
+)
+from vllm.v1.attention.backends.utils import get_supported_kv_cache_layouts
+from vllm.v1.core.kv_cache_utils import get_kv_cache_config_from_groups
+from vllm.v1.kv_cache_interface import KVCacheGroupSpec, UniformTypeKVCacheSpecs
+from vllm.v1.kv_cache_layout import KVCacheLayout
+from vllm.v1.worker.utils import allocate_kv_cache
+
+pytestmark = pytest.mark.skip_global_cleanup
+
+
+@pytest.mark.parametrize(
+    "device",
+    [
+        "cpu",
+        pytest.param(
+            "cuda",
+            marks=pytest.mark.skipif(
+                not current_platform.is_cuda(), reason="CUDA only"
+            ),
+        ),
+    ],
+)
+@pytest.mark.parametrize(
+    "layout", [KVCacheLayout.BLHNC, KVCacheLayout.BLNHC, KVCacheLayout.LBNHC]
+)
+@pytest.mark.parametrize(
+    "backend,cache_dtype,sm100,stride_alignment",
+    [
+        (FlashInferMLASparseTRTLLMBackend, "auto", True, 1152),
+        (FlashInferMLASparseTRTLLMBackend, "fp8", True, 576),
+        (FlashMLASparseBackend, "auto", True, 1152),
+        (FlashMLASparseBackend, "fp8_ds_mla", False, None),
+        (FlashMLASparseBackend, "fp8_ds_mla", True, 656),
+        (FlashMLASparseBackend, "nvfp4_ds_mla", True, 352),
+    ],
+)
+def test_allocation_and_warmup_follow_addressing_mode(
+    monkeypatch, layout, backend, cache_dtype, sm100, stride_alignment, device
+):
+    monkeypatch.setattr(
+        current_platform, "is_device_capability_family", lambda family: sm100
+    )
+    config = SimpleNamespace(
+        cache_config=CacheConfig(block_size=64),
+        model_config=SimpleNamespace(dtype=torch.bfloat16),
+        attention_config=SimpleNamespace(hisparse_config=None),
+        kernel_config=SimpleNamespace(enable_jit_warmup=True),
+    )
+    config.cache_config.kv_cache_layout = layout.name
+    layer = SimpleNamespace(
+        _vllm_config=config,
+        attn_backend=backend,
+        kv_cache_dtype=cache_dtype,
+        head_size=576,
+        sliding_window=None,
+        indexer=None,
+        non_causal_multi_token_decode=False,
+    )
+    layer._uses_flat_kv_cache = MethodType(MLAAttention._uses_flat_kv_cache, layer)
+    spec = MLAAttention.get_kv_cache_spec(layer, config)
+    indexer = SimpleNamespace(
+        cache_config=config.cache_config, head_dim=132, dtype=torch.uint8
+    )
+    index_spec = DeepseekV32IndexerCache.get_kv_cache_spec(indexer, config)
+    assert spec.block_stride_alignment == stride_alignment
+    assert index_spec.block_stride_alignment is None
+    assert layout in get_supported_kv_cache_layouts(
+        [backend, DeepseekV32IndexerBackend]
+    )
+    specs = {"mla": spec, "indexer": index_spec}
+    group = KVCacheGroupSpec(
+        list(specs), UniformTypeKVCacheSpecs(block_size=64, kv_cache_specs=specs)
+    )
+    cache = get_kv_cache_config_from_groups(config, [group], 2**20)
+    views = allocate_kv_cache(cache, torch.device(device), layout)
+    register = MagicMock()
+    monkeypatch.setattr(
+        _CONVERT_REQ_INDEX_TO_GLOBAL_INDEX_KERNEL, "register_warmup", register
+    )
+    MLAAttention.bind_kv_cache(layer, views["mla"])
+    DeepseekV32IndexerCache.bind_kv_cache(indexer, views["indexer"])
+    assert indexer.kv_cache.stride(0) == views["indexer"].stride(0)
+    if stride_alignment is not None:
+        assert (
+            views["mla"].stride(0) * views["mla"].element_size() % stride_alignment == 0
+        )
+    if not layer._uses_flat_kv_cache():
+        register.assert_not_called()
+    else:
+        rows, stride = flat_kv_row_view(layer.kv_cache, 64)
+        register.assert_called_once_with(config, block_stride_rows=stride)
+        layer.kv_cache[1, 0].fill_(3)
+        torch.testing.assert_close(rows[stride], layer.kv_cache[1, 0])
+    if stride_alignment is None:
+        assert cache.kv_cache_tensors[0].size == cache.num_blocks * sum(
+            item.page_size_bytes for item in specs.values()
+        )

tests/v1/core/test_contiguous_kv_packing.py

@@ -10,6 +10,7 @@ block); the allocation is the same either way.
 """
 
 from dataclasses import replace
+from math import lcm
 from unittest.mock import MagicMock
 
 import pytest
@@ -26,9 +27,11 @@ from vllm.v1.core.kv_cache_utils import (
     get_kv_cache_groups,
     resolve_kv_cache_block_sizes,
 )
+from vllm.v1.hisparse.layout import _build_hisparse_kv_cache_tensors
 from vllm.v1.kv_cache_interface import (
     CircularBufferSpec,
     FullAttentionSpec,
+    HiSparseHotSpec,
     KVCacheGroupSpec,
     KVCacheLayout,
     KVCacheSpec,
@@ -79,6 +82,30 @@ def _mock_vllm_config(layout: str | None):
     return config
 
 
+def test_packed_alignment_preserves_hot_pages_and_generic_stride():
+    """One shared stride must satisfy both constraints, regardless of rounding order."""
+    spec = replace(_mla(128), block_stride_alignment=1152)
+    hot = HiSparseHotSpec(block_size=16, page_size=5120, blocks_per_request=2)
+    groups = [KVCacheGroupSpec(["attention"], spec), KVCacheGroupSpec(["hot"], hot)]
+    alignment = lcm(1152, hot.page_size_bytes)
+    stride = _pool_bytes_per_block(groups)
+    assert stride == alignment
+    layout = KVCacheLayout.BLHNC
+    tensors = _build_hisparse_kv_cache_tensors(groups, 3, 3 * stride, layout, stride)
+    assert {tensor.block_stride for tensor in tensors} == {alignment}
+    assert {tensor.size for tensor in tensors} == {3 * alignment}
+
+
+@pytest.mark.parametrize("spec", [_full(), _mla(128)])
+def test_block_alignment_survives_merge_and_rejects_incompatible_specs(spec):
+    aligned = replace(spec, block_stride_alignment=512)
+    assert type(spec).merge([aligned, aligned]).block_stride_alignment == 512
+    with pytest.raises(AssertionError):
+        type(spec).merge([aligned, spec])
+    with pytest.raises(ValueError, match="must be positive"):
+        replace(spec, block_stride_alignment=0)
+
+
 def _pages(groups) -> dict[str, int]:
     return {
         name: group.kv_cache_spec.kv_cache_specs[name].page_size_bytes

vllm/model_executor/layers/attention/mla_attention.py

@@ -735,13 +735,7 @@ class MLAAttention(nn.Module, AttentionLayerBase):
         self.kv_cache = kv_cache.squeeze(1)
         if (
             self._vllm_config.kernel_config.enable_jit_warmup
-            and self.attn_backend.get_name()
-            in (
-                "FLASHMLA_SPARSE",
-                "FLASHINFER_MLA_SPARSE",
-                "FLASHINFER_MLA_SPARSE_SM120",
-                "DEEPSEEK_V32_INDEXER",
-            )
+            and self._uses_flat_kv_cache()
         ):
             from vllm.v1.attention.backends.mla.sparse_utils import (
                 _CONVERT_REQ_INDEX_TO_GLOBAL_INDEX_KERNEL,
@@ -754,6 +748,13 @@ class MLAAttention(nn.Module, AttentionLayerBase):
                 block_stride_rows=self.kv_cache.stride(0) // row_width,
             )
 
+    def _uses_flat_kv_cache(self) -> bool:
+        backend = self.attn_backend.get_name()
+        return backend == "FLASHINFER_MLA_SPARSE" or (
+            backend == "FLASHMLA_SPARSE"
+            and self.kv_cache_dtype not in ("fp8_ds_mla", "nvfp4_ds_mla")
+        )
+
     @property
     def chunked_prefill_workspace_size(self) -> int:
         if self._chunked_prefill_workspace_size is None:
@@ -1365,11 +1366,20 @@ class MLAAttention(nn.Module, AttentionLayerBase):
                 **common_kwargs,
                 sliding_window=self.sliding_window,
             )
-        return MLAAttentionSpec(
+        spec = MLAAttentionSpec(
             **common_kwargs,
             is_index_group_leader=self.indexer is not None,
             non_causal_multi_token_decode=self.non_causal_multi_token_decode,
         )
+        # SM100 FlashMLA paged kernels also express TMA coordinates in token rows.
+        uses_tma_rows = (
+            self.attn_backend.get_name() == "FLASHMLA_SPARSE"
+            and self.kv_cache_dtype in ("fp8_ds_mla", "nvfp4_ds_mla")
+            and current_platform.is_device_capability_family(100)
+        )
+        if self._uses_flat_kv_cache() or uses_tma_rows:
+            spec = replace(spec, block_stride_alignment=spec.state_content_size_bytes)
+        return spec
 
     def _v_up_proj(self, x: torch.Tensor, out: torch.Tensor):
         # Convert from (B, N, L) to (N, B, L)

vllm/v1/core/kv_cache_utils.py

@@ -1608,23 +1608,17 @@ def _get_kv_cache_bytes_per_block(
         for group in kv_cache_groups
     )
     assert bytes_per_block > 0
-    hot_page_sizes = [
+    alignments = [
         group.kv_cache_spec.page_size_bytes
         for group in kv_cache_groups
         if isinstance(group.kv_cache_spec, HiSparseHotSpec)
     ]
-    if hot_page_sizes:
-        bytes_per_block = round_up(bytes_per_block, math.lcm(*hot_page_sizes))
-    stride_alignments = [
-        spec.block_stride_alignment
+    alignments.extend(
+        _get_per_layer_spec(group, layer_name).block_stride_alignment or 1
         for group in kv_cache_groups
         for layer_name in group.layer_names
-        if isinstance(spec := _get_per_layer_spec(group, layer_name), MLAAttentionSpec)
-        and spec.block_stride_alignment
-    ]
-    if stride_alignments:
-        bytes_per_block = round_up(bytes_per_block, math.lcm(*stride_alignments))
-    return bytes_per_block
+    )
+    return round_up(bytes_per_block, math.lcm(*alignments))
 
 
 def validate_kv_cache_layout(

vllm/v1/kv_cache_interface.py

@@ -6,7 +6,7 @@ from __future__ import annotations
 import copy
 from collections import Counter
 from collections.abc import Collection, Sequence
-from dataclasses import dataclass, fields, replace
+from dataclasses import dataclass, field, fields, replace
 from enum import Enum, IntEnum
 from fractions import Fraction
 from functools import cached_property
@@ -159,6 +159,13 @@ class KVCacheSpec:
     # number of tokens in a block
     block_size: int
 
+    block_stride_alignment: int | None = field(default=None, kw_only=True)
+    """Required byte alignment between physical blocks, including packed layers."""
+
+    def __post_init__(self):
+        if self.block_stride_alignment is not None and self.block_stride_alignment <= 0:
+            raise ValueError("block_stride_alignment must be positive")
+
     @property
     def prefix_cacheable(self) -> bool:
         """Whether this spec's group participates in prefix caching."""
@@ -492,6 +499,7 @@ class AttentionSpec(KVCacheSpec):
     token (Whisper block pooling: ``Fraction(1, block_pool_size)``)."""
 
     def __post_init__(self):
+        super().__post_init__()
         if self.head_size_v is None:
             object.__setattr__(self, "head_size_v", self.head_size)
 
@@ -598,6 +606,7 @@ class FullAttentionSpec(AttentionSpec):
         )
         merged_spec = cls(
             block_size=specs[0].block_size,
+            block_stride_alignment=specs[0].block_stride_alignment,
             num_kv_heads=specs[0].num_kv_heads,
             head_size=specs[0].head_size,
             head_size_v=specs[0].head_size_v,
@@ -648,12 +657,6 @@ class MLAAttentionSpec(FullAttentionSpec):
     is_index_group_leader: bool = False
     storage_block_size: int | None = None
     """Token width used to view storage when it differs from the kernel block."""
-    block_stride_alignment: int | None = None
-    """Required alignment, in bytes, of the distance between consecutive
-    blocks of this cache. In block-major layouts that distance is the whole
-    block (all layers' pages), so the allocator rounds the block up to it.
-    DeepGEMM's paged sparse MQA-logits kernels address pages as
-    ``base + page * stride`` and need it 512B-aligned."""
     # Group capability enabled when any member flattens a non-causal query block
     # into decode rows. Runtime metadata still selects causal vs. non-causal mode.
     non_causal_multi_token_decode: bool = False
@@ -751,6 +754,7 @@ class RSWASpec(FullAttentionSpec):
         base = FullAttentionSpec.merge(specs)  # type: ignore[arg-type]
         return cls(
             block_size=base.block_size,
+            block_stride_alignment=base.block_stride_alignment,
             num_kv_heads=base.num_kv_heads,
             head_size=base.head_size,
             head_size_v=base.head_size_v,
@@ -945,6 +949,7 @@ class SlidingWindowMLASpec(SlidingWindowSpec):
         sliding_window_set = set(spec.sliding_window for spec in specs)
         extra_retained_set = set(spec.extra_retained_tokens for spec in specs)
         bounded_replay_set = set(spec.bounded_replay for spec in specs)
+        block_stride_alignment_set = {spec.block_stride_alignment for spec in specs}
         assert (
             len(cache_dtype_str_set) == 1
             and len(tokens_per_state_set) == 1
@@ -952,6 +957,7 @@ class SlidingWindowMLASpec(SlidingWindowSpec):
             and len(sliding_window_set) == 1
             and len(extra_retained_set) == 1
             and len(bounded_replay_set) == 1
+            and len(block_stride_alignment_set) == 1
         ), (
             "All attention layers in the same KV cache group must use the same "
             "quantization method, tokens per state, model version, sliding "
@@ -959,6 +965,7 @@ class SlidingWindowMLASpec(SlidingWindowSpec):
         )
         return cls(
             block_size=specs[0].block_size,
+            block_stride_alignment=block_stride_alignment_set.pop(),
             num_kv_heads=specs[0].num_kv_heads,
             head_size=specs[0].head_size,
             dtype=specs[0].dtype,
@@ -1175,6 +1182,7 @@ class SinkFullAttentionSpec(FullAttentionSpec):
             head_size=specs[0].head_size,
             head_size_v=specs[0].head_size_v,
             sink_len=specs[0].sink_len,
+            block_stride_alignment=specs[0].block_stride_alignment,
             dtype=specs[0].dtype,
             kv_quant_mode=specs[0].kv_quant_mode,
             page_size_padded=specs[0].page_size_padded,