| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from types import SimpleNamespace |
| 5 | |
| 6 | import pytest |
| 7 | import torch |
| 8 | |
| 9 | from vllm.v1.core.kv_cache_manager import KVCacheManager |
| 10 | from vllm.v1.engine.core import EngineCore |
| 11 | from vllm.v1.kv_cache_interface import ( |
| 12 | KVCacheConfig, |
| 13 | KVCacheGroupSpec, |
| 14 | MambaSpec, |
| 15 | MLAAttentionSpec, |
| 16 | ) |
| 17 | |
| 18 | pytestmark = pytest.mark.cpu_test |
| 19 | |
| 20 | BASE_BLOCK_SIZE = 1536 |
| 21 | DCP_WORLD_SIZE = 8 |
| 22 | |
| 23 | |
| 24 | def _make_engine_core_with_dcp_manager() -> EngineCore: |
| 25 | mla_spec = MLAAttentionSpec( |
| 26 | block_size=BASE_BLOCK_SIZE, |
| 27 | num_kv_heads=1, |
| 28 | head_size=576, |
| 29 | dtype=torch.bfloat16, |
| 30 | ) |
| 31 | mamba_spec = MambaSpec( |
| 32 | block_size=BASE_BLOCK_SIZE, |
| 33 | shapes=((1, 1),), |
| 34 | dtypes=(torch.float32,), |
| 35 | mamba_cache_mode="align", |
| 36 | ) |
| 37 | config = KVCacheConfig( |
| 38 | num_blocks=32, |
| 39 | kv_cache_tensors=[], |
| 40 | kv_cache_groups=[ |
| 41 | KVCacheGroupSpec(["mla"], mla_spec), |
| 42 | KVCacheGroupSpec(["mamba"], mamba_spec), |
| 43 | ], |
| 44 | ) |
| 45 | manager = KVCacheManager( |
| 46 | config, |
| 47 | max_model_len=BASE_BLOCK_SIZE * DCP_WORLD_SIZE * 2, |
| 48 | scheduler_block_size=BASE_BLOCK_SIZE * DCP_WORLD_SIZE, |
| 49 | hash_block_size=BASE_BLOCK_SIZE, |
| 50 | dcp_world_size=DCP_WORLD_SIZE, |
| 51 | ) |
| 52 | scheduler = SimpleNamespace(kv_cache_config=config, kv_cache_manager=manager) |
| 53 | engine_core = object.__new__(EngineCore) |
| 54 | engine_core.scheduler = scheduler |
| 55 | return engine_core |
| 56 | |
| 57 | |
| 58 | def test_kv_cache_group_metadata_uses_effective_dcp_block_sizes() -> None: |
| 59 | engine_core = _make_engine_core_with_dcp_manager() |
| 60 | |
| 61 | metadata = engine_core.get_kv_cache_group_metadata() |
| 62 | |
| 63 | # DCP scales attention blocks but does not scale replicated Mamba state. |
| 64 | assert [item["block_size"] for item in metadata] == [12288, 1536] |
| 65 | assert [item["kind"] for item in metadata] == ["mla_attention", "mamba"] |
| 66 | |
| 67 | managers = engine_core.scheduler.kv_cache_manager.coordinator.single_type_managers |
| 68 | assert [item["block_size"] for item in metadata] == [ |
| 69 | manager.block_size for manager in managers |
| 70 | ] |