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