diff --git a/tests/v1/attention/test_mla_backends.py b/tests/v1/attention/test_mla_backends.py index 2a9cd7736b..9ae0fc216a 100644 --- a/tests/v1/attention/test_mla_backends.py +++ b/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 ) diff --git a/tests/v1/attention/test_sparse_mla_kv_cache_layout.py b/tests/v1/attention/test_sparse_mla_kv_cache_layout.py new file mode 100644 index 0000000000..cd3c9b2181 --- /dev/null +++ b/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() + ) diff --git a/tests/v1/core/test_contiguous_kv_packing.py b/tests/v1/core/test_contiguous_kv_packing.py index 58d4e20bf7..638deeedc5 100644 --- a/tests/v1/core/test_contiguous_kv_packing.py +++ b/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 diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 4ebf3ef40d..2ae111c1c3 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/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) diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index 8d60b09fa3..ee0957839e 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/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( diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index fc8b88921a..499b50a3a2 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/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,