# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import copy import hashlib import importlib import mmap import subprocess import sys from collections.abc import Callable from dataclasses import replace from pathlib import Path from types import SimpleNamespace from typing import Any, cast import pytest import torch import vllm.v1.core.kv_cache_utils as kv_cache_utils import vllm.v1.hisparse.runtime as hisparse_runtime_module from vllm.config import ( CacheConfig, KVTransferConfig, ModelConfig, SchedulerConfig, VllmConfig, ) from vllm.config.attention import HiSparseConfig from vllm.config.kv_events import KVEventsConfig from vllm.lora.request import LoRARequest from vllm.multimodal.inputs import ( MultiModalFeatureSpec, MultiModalKwargsItem, PlaceholderRange, ) from vllm.sampling_params import SamplingParams from vllm.utils.hashing import sha256, sha256_cbor, xxhash, xxhash_cbor from vllm.utils.mem_constants import GiB_bytes from vllm.v1.core.kv_cache_manager import KVCacheBlocks, KVCacheManager from vllm.v1.core.kv_cache_utils import ( BlockHash, FreeKVCacheBlockQueue, KVCacheBlock, check_enough_kv_cache_memory, estimate_max_model_len, generate_block_hash_extra_keys, generate_scheduler_kv_cache_config, get_kv_cache_capacity, get_kv_cache_configs, get_kv_cache_groups, get_max_concurrency_for_kv_cache_config, get_request_block_hasher, hash_block_tokens, init_none_hash, is_kv_cache_spec_uniform, kv_cache_groups_tp_replicas, make_block_hash_with_group_id, tensor_data, to_event_extra_keys, ) from vllm.v1.hisparse.layout import ( create_hisparse_layout, get_hisparse_kv_cache_groups, ) from vllm.v1.kv_cache_interface import ( ChunkedLocalAttentionSpec, CircularBufferSpec, FullAttentionSpec, HiddenStateCacheSpec, HiSparseHotSpec, HiSparseResidentSpec, KpoolTailSpec, KVCacheConfig, KVCacheGroupSpec, KVCacheSpec, KVCacheSpecKind, KVCacheTensor, KVQuantMode, MambaSpec, MLAAttentionSpec, SinkFullAttentionSpec, SlidingWindowMLASpec, SlidingWindowSpec, SparseCacheRole, UniformTypeKVCacheSpecs, get_kv_cache_spec_kind, get_kv_cache_spec_sliding_window, is_full_attention_spec, iter_layer_specs, ) from vllm.v1.kv_cache_layout import KVCacheLayout from vllm.v1.metrics.stats import CachingMetrics, PrefixCacheStats from vllm.v1.request import Request pytestmark = pytest.mark.cpu_test @pytest.mark.parametrize("gpu_block_size", [32, 64]) @pytest.mark.parametrize("shared_host_pool", [False, True]) def test_hisparse_hma_uses_resolved_gpu_block_size( monkeypatch, gpu_block_size, shared_host_pool ): monkeypatch.setattr( hisparse_runtime_module.current_platform, "is_cuda_alike", lambda: True ) specs = { "model.layers.0.self_attn": MLAAttentionSpec( block_size=gpu_block_size, num_kv_heads=1, head_size=576, dtype=torch.bfloat16, is_index_group_leader=True, ), "model.layers.0.self_attn.indexer": MLAAttentionSpec( block_size=gpu_block_size, num_kv_heads=1, head_size=128, dtype=torch.bfloat16, cache_role=SparseCacheRole.INDEXER, ), } config = SimpleNamespace( attention_config=SimpleNamespace(hisparse_config=HiSparseConfig()), model_config=SimpleNamespace( hf_config=SimpleNamespace(index_topk=128), max_model_len=gpu_block_size, ), parallel_config=SimpleNamespace( tensor_parallel_size=2 if shared_host_pool else 1, pipeline_parallel_size=1, prefill_context_parallel_size=1, decode_context_parallel_size=1, world_size=2 if shared_host_pool else 1, distributed_executor_backend="mp", nnodes_within_dp=1, ), cache_config=SimpleNamespace( num_gpu_blocks_override=7, prefix_cache_retention_interval=None, get_resolved_kv_cache_layout=lambda: KVCacheLayout.BLHNC, ), ) monkeypatch.setattr(kv_cache_utils, "get_hisparse_host_pool_bytes", lambda _: 2**30) cache_config = kv_cache_utils.get_kv_cache_config_from_groups( config, get_hisparse_kv_cache_groups(config, specs), available_memory=2**30 ) assert cache_config.num_blocks == 7 assert cache_config.hisparse_host_num_blocks is not None assert cache_config.hisparse_host_num_blocks > 7 host_group, indexer_group, *auxiliary_groups = cache_config.kv_cache_groups assert cache_config.hisparse_shared_host_pool is shared_host_pool host_page = host_group.kv_cache_spec.page_size_bytes alignment = mmap.PAGESIZE if shared_host_pool else 1 expected_host_stride = (host_page + alignment - 1) // alignment * alignment assert cache_config.hisparse_host_block_stride == expected_host_stride assert cache_config.hisparse_host_num_blocks == 2**30 // expected_host_stride assert host_group.host_resident assert not any(group.host_resident for group in [indexer_group, *auxiliary_groups]) host_layers = set(host_group.layer_names) for tensor in cache_config.kv_cache_tensors: assert all( (name in host_layers) == tensor.host_resident for name in tensor.layers ) host_tensor = next(t for t in cache_config.kv_cache_tensors if t.host_resident) host_spec = host_group.kv_cache_spec.kv_cache_specs[host_tensor.layers[0]] assert host_tensor.block_stride == host_spec.page_size_bytes assert host_tensor.layer_stride == ( host_spec.page_size_bytes * cache_config.hisparse_host_num_blocks ) assert host_group.kv_cache_spec.block_size == gpu_block_size assert indexer_group.kv_cache_spec.block_size == gpu_block_size host_specs = host_group.kv_cache_spec.kv_cache_specs gpu_indexer_specs = indexer_group.kv_cache_spec.kv_cache_specs assert set(host_specs) == {"model.layers.0.self_attn"} assert set(gpu_indexer_specs) == {"model.layers.0.self_attn.indexer"} assert indexer_group.kv_cache_spec.prefix_cacheable assert host_group.enable_kv_transfer auxiliary_specs = [group.kv_cache_spec for group in auxiliary_groups] assert any(isinstance(spec, HiSparseResidentSpec) for spec in auxiliary_specs) assert any(isinstance(spec, HiSparseHotSpec) for spec in auxiliary_specs) assert all( spec.block_size == gpu_block_size for spec in auxiliary_specs if isinstance(spec, (HiSparseResidentSpec, HiSparseHotSpec)) ) scheduler_block_size, hash_block_size = kv_cache_utils.resolve_kv_cache_block_sizes( cache_config, SimpleNamespace( cache_config=SimpleNamespace( block_size=16, enable_prefix_caching=True, prefix_match_unit=None, ), parallel_config=SimpleNamespace(decode_context_parallel_size=1), kv_transfer_config=object(), ), ) assert scheduler_block_size == hash_block_size == gpu_block_size @pytest.mark.parametrize("extra_blocks,ok", [(1, False), (2, True)]) def test_hisparse_host_pool_must_fit_max_model_len(monkeypatch, extra_blocks, ok): """Every page needs a host block, so one max_model_len request must fit.""" monkeypatch.setattr( hisparse_runtime_module.current_platform, "is_cuda_alike", lambda: True ) specs = { "model.layers.0.self_attn": MLAAttentionSpec( block_size=64, num_kv_heads=1, head_size=576, dtype=torch.bfloat16, is_index_group_leader=True, ), "model.layers.0.self_attn.indexer": MLAAttentionSpec( block_size=64, num_kv_heads=1, head_size=128, dtype=torch.bfloat16, cache_role=SparseCacheRole.INDEXER, ), } config = SimpleNamespace( attention_config=SimpleNamespace(hisparse_config=HiSparseConfig()), model_config=SimpleNamespace( hf_config=SimpleNamespace(index_topk=128), max_model_len=64 * 4 ), parallel_config=SimpleNamespace(tensor_parallel_size=1, world_size=1), cache_config=SimpleNamespace(num_gpu_blocks_override=7), ) host_page = specs["model.layers.0.self_attn"].page_size_bytes host_budget = (4 + extra_blocks) * host_page groups = get_hisparse_kv_cache_groups(config, specs) if ok: layout = create_hisparse_layout(config, groups, host_budget=host_budget) assert layout.host_num_blocks == 4 + extra_blocks else: with pytest.raises(ValueError, match="increase host_pool_gib"): create_hisparse_layout(config, groups, host_budget=host_budget) @pytest.mark.parametrize( "max_model_len,num_gpu_blocks,ok", [ (8192, 160, False), (8192, 1000, True), (32768, 656, False), (32768, 657, True), (-1, 657, True), ], ) def test_hisparse_pool_must_fit_max_model_len( monkeypatch, max_model_len, num_gpu_blocks, ok ): """Each resident group takes its own blocks from HiSparse's shared GPU pool, so a pool that fits only the indexer pages must be rejected at startup. Resident pages are capped at the in-flight window: at 32768 tokens with 2048 in flight, one request needs 512 indexer + 4 x 32 resident + 16 hot blocks, plus the null block. Auto-fit (-1) must pick a length that passes the same check.""" monkeypatch.setattr( hisparse_runtime_module.current_platform, "is_cuda_alike", lambda: True ) monkeypatch.setattr(kv_cache_utils, "get_hisparse_host_pool_bytes", lambda _: 2**30) specs: dict[str, KVCacheSpec] = {} for i in range(4): specs[f"model.layers.{i}.self_attn"] = MLAAttentionSpec( block_size=64, num_kv_heads=1, head_size=576, dtype=torch.bfloat16, is_index_group_leader=True, ) specs[f"model.layers.{i}.self_attn.indexer"] = MLAAttentionSpec( block_size=64, num_kv_heads=1, head_size=128, dtype=torch.bfloat16, cache_role=SparseCacheRole.INDEXER, ) config = VllmConfig(model_config=ModelConfig(max_model_len=max_model_len)) config.scheduler_config.max_num_batched_tokens = 2048 # Async scheduling doubles the in-flight window, and with it the bound. config.scheduler_config.async_scheduling = False config.attention_config.hisparse_config = HiSparseConfig() config.model_config.hf_config.index_topk = 128 config.cache_config.num_gpu_blocks_override = num_gpu_blocks config.cache_config.kv_cache_layout = "BLHNC" # Workers report the resident/hot caches HiSparse derives. specs |= { name: group.kv_cache_spec for group in get_hisparse_kv_cache_groups(config, specs) or [] if isinstance(group.kv_cache_spec, (HiSparseResidentSpec, HiSparseHotSpec)) for name in group.layer_names } if ok: kv_cache_utils.get_kv_cache_configs(config, [specs], [2**34]) else: with pytest.raises(ValueError, match="max seq len"): kv_cache_utils.get_kv_cache_configs(config, [specs], [2**34]) def test_hisparse_derived_specs_do_not_affect_kv_cache_layout(monkeypatch): """Workers report HiSparse's resident/hot caches, which lay out their own raw backing; KV cache layout resolution must skip them.""" from vllm.v1.attention.backends.utils import resolve_kv_cache_layout monkeypatch.setattr( hisparse_runtime_module.current_platform, "is_cuda_alike", lambda: True ) monkeypatch.delenv("VLLM_KV_CACHE_LAYOUT", raising=False) specs: dict[str, KVCacheSpec] = { "model.layers.0.self_attn": MLAAttentionSpec( block_size=64, num_kv_heads=1, head_size=576, dtype=torch.bfloat16, is_index_group_leader=True, ), "model.layers.0.self_attn.indexer": MLAAttentionSpec( block_size=64, num_kv_heads=1, head_size=128, dtype=torch.bfloat16, cache_role=SparseCacheRole.INDEXER, ), } config = VllmConfig(model_config=ModelConfig(max_model_len=1024)) config.attention_config.hisparse_config = HiSparseConfig() config.model_config.hf_config.index_topk = 128 reported = specs | { name: group.kv_cache_spec for group in get_hisparse_kv_cache_groups(config, specs) or [] if isinstance(group.kv_cache_spec, (HiSparseResidentSpec, HiSparseHotSpec)) for name in group.layer_names } assert len(reported) > len(specs) layouts = [] for worker_specs in (specs, reported): config.cache_config.kv_cache_layout = None layouts.append( resolve_kv_cache_layout( config, [["BLHNC", "LBNHC"]], list(worker_specs.values()) ) ) assert layouts[0] == layouts[1] def test_hisparse_rejects_deepseek_v4(): full_specs = { "model.layers.0.attn": MLAAttentionSpec( block_size=256, num_kv_heads=1, head_size=512, dtype=torch.uint8, model_version="deepseek_v4", ) } config = SimpleNamespace( attention_config=SimpleNamespace(hisparse_config=HiSparseConfig()), model_config=SimpleNamespace(hf_config=SimpleNamespace(index_topk=512)), cache_config=SimpleNamespace(num_gpu_blocks_override=7), ) with pytest.raises(ValueError, match="does not support DeepSeek V4"): get_hisparse_kv_cache_groups(config, full_specs) @pytest.fixture(autouse=True) def _auto_init_hash_fn(request): hash_fn: Callable if "hash_fn" in request.fixturenames: hash_fn = request.getfixturevalue("hash_fn") else: hash_fn = sha256 init_none_hash(hash_fn) def make_request( request_id: str, prompt_token_ids: list[int] | None, block_size: int = 3, hash_fn: Callable = hash, mm_positions: list[PlaceholderRange] | None = None, mm_hashes: list[str] | None = None, cache_salt: str | None = None, prompt_embeds: torch.Tensor | None = None, ): mm_features = [] if mm_positions is not None: for j, position in enumerate(mm_positions): identifier = mm_hashes[j] if mm_hashes else f"hash_{j}" mm_feature = MultiModalFeatureSpec( data=MultiModalKwargsItem.dummy(), mm_position=position, identifier=identifier, modality="image", ) mm_features.append(mm_feature) sampling_params = SamplingParams(max_tokens=17) sampling_params.update_from_generation_config({}, eos_token_id=100) return Request( request_id=request_id, prompt_token_ids=prompt_token_ids, mm_features=mm_features if mm_features else None, sampling_params=sampling_params, pooling_params=None, lora_request=None, cache_salt=cache_salt, block_hasher=get_request_block_hasher(block_size, hash_fn), prompt_embeds=prompt_embeds, ) @pytest.mark.parametrize("dcp", [1, 4]) def test_effective_attention_block_size_matches_events(dcp): from vllm.distributed.kv_events import BlockStored from vllm.v1.engine.core import EngineCore config = KVCacheConfig( num_blocks=32, kv_cache_tensors=[], kv_cache_groups=[ KVCacheGroupSpec(["attention"], new_kv_cache_spec()), ], ) manager = KVCacheManager( generate_scheduler_kv_cache_config([config]), max_model_len=256, scheduler_block_size=16 * dcp, hash_block_size=16 * dcp, dcp_world_size=dcp, enable_kv_cache_events=True, ) core = EngineCore.__new__(EngineCore) core.vllm_config = SimpleNamespace(cache_config=CacheConfig(block_size=16)) core.scheduler = SimpleNamespace(kv_cache_manager=manager) core._initialize_effective_attention_block_size() block_size = core.vllm_config.cache_config.effective_attention_block_size assert block_size == 16 * dcp request = make_request( "block-size", list(range(64)), block_size=16 * dcp, hash_fn=sha256 ) assert manager.allocate_slots(request, 64) is not None assert [ event.block_size for event in manager.take_events() if isinstance(event, BlockStored) ] == [block_size] def new_kv_cache_spec( block_size=16, num_kv_heads=2, head_size=64, dtype=torch.float32, page_size_padded=None, sliding_window=None, attention_chunk_size=None, kv_quant_mode=KVQuantMode.NONE, ): return FullAttentionSpec( block_size=block_size, num_kv_heads=num_kv_heads, head_size=head_size, dtype=dtype, page_size_padded=page_size_padded, sliding_window=sliding_window, attention_chunk_size=attention_chunk_size, kv_quant_mode=kv_quant_mode, ) def test_kv_cache_config_selects_only_transferable_groups(): """Connectors must see a stable projection of transfer-eligible groups.""" groups = [ KVCacheGroupSpec(["layer.0"], new_kv_cache_spec()), KVCacheGroupSpec(["layer.1"], new_kv_cache_spec(), enable_kv_transfer=False), KVCacheGroupSpec(["layer.2"], new_kv_cache_spec()), ] config = KVCacheConfig( num_blocks=1, kv_cache_tensors=[], kv_cache_groups=groups, ) assert config.transfer_group_ids == (0, 2) assert config.transfer_groups == (groups[0], groups[2]) assert config.transfer_group_index_by_layer == {"layer.0": 0, "layer.2": 1} first_blocks = [1, 2] third_blocks = [4] assert config.select_transfer_block_ids((first_blocks, [3], third_blocks)) == ( first_blocks, third_blocks, ) def test_kv_cache_config_selects_prefix_cacheable_groups(): """Prefix stores exclude scratch state without changing transfer groups.""" full_group = KVCacheGroupSpec(["full"], new_kv_cache_spec()) qsa_group = KVCacheGroupSpec( ["qsa"], CircularBufferSpec( block_size=4, num_kv_heads=1, head_size=64, head_size_v=0, dtype=torch.float16, ), ) disabled_group = KVCacheGroupSpec( ["disabled"], new_kv_cache_spec(), enable_kv_transfer=False ) config = KVCacheConfig( num_blocks=1, kv_cache_tensors=[], kv_cache_groups=[full_group, qsa_group, disabled_group], ) assert config.transfer_group_ids == (0, 1) assert config.select_transfer_block_ids(([1], [2], [3])) == ([1], [2]) assert config.prefix_cacheable_group_ids == (0,) assert config.prefix_cacheable_groups == (full_group,) def test_kv_cache_blocks_selects_requested_groups(): blocks = KVCacheBlocks( ( [KVCacheBlock(1)], [KVCacheBlock(2)], [KVCacheBlock(3)], ) ) assert blocks.get_block_ids(group_ids=(0, 2)) == ([1], [3]) def new_sliding_window_spec( block_size=16, num_kv_heads=2, head_size=64, dtype=torch.float32, page_size_padded=None, sliding_window=1, ): return SlidingWindowSpec( block_size=block_size, num_kv_heads=num_kv_heads, head_size=head_size, dtype=dtype, page_size_padded=page_size_padded, sliding_window=sliding_window, ) def new_chunked_local_attention_spec( block_size=16, num_kv_heads=2, head_size=64, dtype=torch.float32, page_size_padded=None, attention_chunk_size=4, ): return ChunkedLocalAttentionSpec( block_size=block_size, num_kv_heads=num_kv_heads, head_size=head_size, dtype=dtype, page_size_padded=page_size_padded, attention_chunk_size=attention_chunk_size, ) def new_mamba_spec( block_size=16, shapes=((2, 512), (3, 32, 32)), dtypes=(torch.float32, torch.float32), num_speculative_blocks=2, mamba_cache_mode="none", page_size_padded=None, ): return MambaSpec( block_size=block_size, shapes=shapes, dtypes=dtypes, page_size_padded=page_size_padded, mamba_cache_mode=mamba_cache_mode, num_speculative_blocks=num_speculative_blocks, ) @pytest.mark.parametrize("hash_fn", [sha256, sha256_cbor]) def test_none_hash(monkeypatch, hash_fn): import vllm.v1.core.kv_cache_utils # case 1: PYTHONHASHSEED is not set -> deterministic default seed so that # independent processes compute identical block hashes for identical # content (e.g. for KV cache reuse across nodes). with monkeypatch.context() as m: m.delenv("PYTHONHASHSEED", raising=False) reloaded_kv_cache_utils = importlib.reload(vllm.v1.core.kv_cache_utils) reloaded_kv_cache_utils.init_none_hash(hash_fn) none_hash = reloaded_kv_cache_utils.NONE_HASH assert isinstance(none_hash, bytes) assert none_hash != b"" assert none_hash == hash_fn(reloaded_kv_cache_utils.DEFAULT_NONE_HASH_SEED) # deterministic across re-initialization within the same environment reloaded_kv_cache_utils.init_none_hash(hash_fn) assert none_hash == reloaded_kv_cache_utils.NONE_HASH # case 2: PYTHONHASHSEED is set, use the seed and hash_fn with monkeypatch.context() as m: m.setenv("PYTHONHASHSEED", "python hash seed") reloaded_kv_cache_utils = importlib.reload(vllm.v1.core.kv_cache_utils) reloaded_kv_cache_utils.init_none_hash(hash_fn) assert reloaded_kv_cache_utils.NONE_HASH is not None assert isinstance(reloaded_kv_cache_utils.NONE_HASH, bytes) assert hash_fn("python hash seed") == reloaded_kv_cache_utils.NONE_HASH @pytest.mark.parametrize("non_crypto_fn", [xxhash, xxhash_cbor]) def test_none_hash_seed_random_for_non_crypto(monkeypatch, non_crypto_fn): """Non-cryptographic algorithms keep the per-process random seed. A deterministic seed is safe for SHA-256, whose collision resistance does not depend on a secret, but xxHash is not collision resistant: a known seed would let an attacker precompute colliding blocks offline. Keep the unpredictable seed there unless the operator opts into a shared one. """ # PYTHONHASHSEED unset -> unpredictable, differs per resolution. with monkeypatch.context() as m: m.delenv("PYTHONHASHSEED", raising=False) seeds = {kv_cache_utils.resolve_none_hash_seed(non_crypto_fn) for _ in range(5)} assert len(seeds) == 5 assert kv_cache_utils.DEFAULT_NONE_HASH_SEED not in seeds # PYTHONHASHSEED set -> operator opt-in wins, so peers can share a cache. with monkeypatch.context() as m: m.setenv("PYTHONHASHSEED", "12345") assert kv_cache_utils.resolve_none_hash_seed(non_crypto_fn) == "12345" @pytest.mark.parametrize("hash_fn", [sha256, sha256_cbor]) def test_none_hash_seed_deterministic_for_crypto(monkeypatch, hash_fn): with monkeypatch.context() as m: m.delenv("PYTHONHASHSEED", raising=False) seed = kv_cache_utils.resolve_none_hash_seed(hash_fn) assert seed == kv_cache_utils.DEFAULT_NONE_HASH_SEED assert seed == kv_cache_utils.resolve_none_hash_seed(hash_fn) def test_get_none_hash_seed_reports_effective_seed(monkeypatch): """P2P advertises the seed NONE_HASH was actually derived from. The P2P tier is constructed before init_none_hash runs, so it must read the resolved seed lazily rather than re-deriving it. """ import vllm.v1.core.kv_cache_utils with monkeypatch.context() as m: m.delenv("PYTHONHASHSEED", raising=False) reloaded = importlib.reload(vllm.v1.core.kv_cache_utils) reloaded.init_none_hash(sha256) assert reloaded.get_none_hash_seed() == reloaded.DEFAULT_NONE_HASH_SEED with monkeypatch.context() as m: m.setenv("PYTHONHASHSEED", "12345") reloaded = importlib.reload(vllm.v1.core.kv_cache_utils) reloaded.init_none_hash(sha256) assert reloaded.get_none_hash_seed() == "12345" def test_kv_cache_block(): # Test KVCacheBlock initialization block = KVCacheBlock(block_id=0) assert block.block_id == 0 assert block.ref_cnt == 0 assert block.block_hash is None # Test reference count manipulation block.ref_cnt += 1 assert block.ref_cnt == 1 block.ref_cnt -= 1 assert block.ref_cnt == 0 # Test block hash setting and resetting block_hash = make_block_hash_with_group_id(BlockHash(b"abc"), 0) block.set_block_hash(block_hash) assert block.block_hash == block_hash block.reset_hash() assert block.block_hash is None def test_kv_cache_block_uses_slots(): block = KVCacheBlock(block_id=0) # Slots eliminate per-instance __dict__, saving ~264 bytes per block. # At 100K+ blocks this avoids tens of MB of overhead and GC pressure. assert not hasattr(block, "__dict__") # Verify that slots actually prevent dynamic attribute assignment. with pytest.raises(AttributeError): block.unexpected_field = True def test_free_kv_cache_block_queue_initialization(): # Test with a single block block = KVCacheBlock(block_id=0) queue = FreeKVCacheBlockQueue([block]) assert queue.num_free_blocks == 1 assert queue.fake_free_list_head.next_free_block is block assert queue.fake_free_list_tail.prev_free_block is block def test_free_kv_cache_block_queue_operations(): # Create a list of KVCacheBlock objects blocks = [KVCacheBlock(block_id=i) for i in range(5)] # Create a FreeKVCacheBlockQueue with these blocks queue = FreeKVCacheBlockQueue(blocks) # Check initial state assert queue.num_free_blocks == 5 assert queue.fake_free_list_head.next_free_block is blocks[0] assert queue.fake_free_list_tail.prev_free_block is blocks[4] # Pop the first block block1 = queue.popleft() assert block1 == blocks[0] assert queue.num_free_blocks == 4 assert queue.fake_free_list_head.next_free_block is blocks[1] assert queue.fake_free_list_tail.prev_free_block is blocks[4] # Remove a block from the middle block_to_remove = blocks[2] queue.remove(block_to_remove) assert queue.num_free_blocks == 3 assert blocks[1].next_free_block is blocks[3] assert blocks[3].prev_free_block is blocks[1] # Append a block back queue.append(block_to_remove) assert queue.num_free_blocks == 4 assert queue.fake_free_list_tail.prev_free_block is block_to_remove assert block_to_remove.prev_free_block is blocks[4] assert block_to_remove.next_free_block is queue.fake_free_list_tail # Pop blocks until empty for _ in range(4): queue.popleft() assert queue.num_free_blocks == 0 assert queue.fake_free_list_head.next_free_block is queue.fake_free_list_tail assert queue.fake_free_list_tail.prev_free_block is queue.fake_free_list_head # Attempt to pop from an empty queue with pytest.raises(ValueError) as e: queue.popleft() assert str(e.value) == "No free blocks available" def test_free_kv_cache_block_queue_append_n(): # Create an empty FreeKVCacheBlockQueue with these blocks queue = FreeKVCacheBlockQueue([]) blocks = [KVCacheBlock(block_id=i) for i in range(6)] # Append 0 block # fake_head->fake_tail queue.append_n([]) assert queue.num_free_blocks == 0 assert queue.fake_free_list_head.next_free_block is queue.fake_free_list_tail assert queue.fake_free_list_tail.prev_free_block is queue.fake_free_list_head # Append 1 block # fake_head->b0->fake_tail queue.append_n(blocks[0:1]) assert queue.num_free_blocks == 1 assert queue.fake_free_list_head.next_free_block is blocks[0] assert blocks[0].prev_free_block is queue.fake_free_list_head assert blocks[0].next_free_block is queue.fake_free_list_tail assert queue.fake_free_list_tail.prev_free_block is blocks[0] # Append 2 blocks # fake_head->b0->b4->b5->fake_tail queue.append_n(blocks[4:6]) assert queue.num_free_blocks == 3 assert queue.fake_free_list_head.next_free_block is blocks[0] assert blocks[0].prev_free_block is queue.fake_free_list_head assert blocks[0].next_free_block is blocks[4] assert blocks[4].prev_free_block is blocks[0] assert blocks[4].next_free_block is blocks[5] assert blocks[5].prev_free_block is blocks[4] assert blocks[5].next_free_block is queue.fake_free_list_tail assert queue.fake_free_list_tail.prev_free_block is blocks[5] # Append 3 blocks # fake_head->b0->b4->b5->b1->b2->b3->fake_tail queue.append_n(blocks[1:4]) assert queue.num_free_blocks == 6 assert queue.fake_free_list_head.next_free_block is blocks[0] assert blocks[0].prev_free_block is queue.fake_free_list_head assert blocks[0].next_free_block is blocks[4] assert blocks[4].prev_free_block is blocks[0] assert blocks[4].next_free_block is blocks[5] assert blocks[5].prev_free_block is blocks[4] assert blocks[5].next_free_block is blocks[1] assert blocks[1].prev_free_block is blocks[5] assert blocks[1].next_free_block is blocks[2] assert blocks[2].prev_free_block is blocks[1] assert blocks[2].next_free_block is blocks[3] assert blocks[3].prev_free_block is blocks[2] assert blocks[3].next_free_block is queue.fake_free_list_tail assert queue.fake_free_list_tail.prev_free_block is blocks[3] # Create an empty FreeKVCacheBlockQueue invalid_queue = FreeKVCacheBlockQueue([]) # set prev_free_block to None and this will cause assertion in append_n invalid_queue.fake_free_list_tail.prev_free_block = None with pytest.raises(AssertionError): # Append 1 block # fake_head->fake_tail invalid_queue.append_n(blocks[0:1]) assert invalid_queue.num_free_blocks == 0 assert ( invalid_queue.fake_free_list_head.next_free_block == invalid_queue.fake_free_list_tail ) def test_free_kv_cache_block_queue_prepend_n(): # Seed the queue with one block so prepend has an existing head to splice # in front of (fake_head->b0->fake_tail). blocks = [KVCacheBlock(block_id=i) for i in range(6)] queue = FreeKVCacheBlockQueue(blocks[0:1]) # Prepend 0 blocks is a no-op. queue.prepend_n([]) assert queue.num_free_blocks == 1 assert queue.fake_free_list_head.next_free_block is blocks[0] # Prepend 2 blocks; they land in front of the existing head, in order. # fake_head->b4->b5->b0->fake_tail queue.prepend_n(blocks[4:6]) assert queue.num_free_blocks == 3 assert queue.fake_free_list_head.next_free_block is blocks[4] assert blocks[4].prev_free_block is queue.fake_free_list_head assert blocks[4].next_free_block is blocks[5] assert blocks[5].prev_free_block is blocks[4] assert blocks[5].next_free_block is blocks[0] assert blocks[0].prev_free_block is blocks[5] assert blocks[0].next_free_block is queue.fake_free_list_tail assert queue.fake_free_list_tail.prev_free_block is blocks[0] # A second prepend goes ahead of everything previously prepended. # fake_head->b1->b2->b4->b5->b0->fake_tail queue.prepend_n(blocks[1:3]) assert queue.num_free_blocks == 5 assert queue.fake_free_list_head.next_free_block is blocks[1] assert blocks[1].next_free_block is blocks[2] assert blocks[2].next_free_block is blocks[4] # The popleft order reflects the front-to-back queue order. assert [queue.popleft().block_id for _ in range(5)] == [1, 2, 4, 5, 0] assert queue.num_free_blocks == 0 def test_free_kv_cache_block_queue_popleft_n(): blocks = [KVCacheBlock(block_id=i) for i in range(6)] # Create an empty FreeKVCacheBlockQueue with these blocks queue = FreeKVCacheBlockQueue( [blocks[1], blocks[3], blocks[5], blocks[4], blocks[0], blocks[2]] ) assert queue.num_free_blocks == 6 assert queue.fake_free_list_head.next_free_block is blocks[1] assert blocks[1].prev_free_block is queue.fake_free_list_head assert blocks[1].next_free_block is blocks[3] assert blocks[3].prev_free_block is blocks[1] assert blocks[3].next_free_block is blocks[5] assert blocks[5].prev_free_block is blocks[3] assert blocks[5].next_free_block is blocks[4] assert blocks[4].prev_free_block is blocks[5] assert blocks[4].next_free_block is blocks[0] assert blocks[0].prev_free_block is blocks[4] assert blocks[0].next_free_block is blocks[2] assert blocks[2].prev_free_block is blocks[0] assert blocks[2].next_free_block is queue.fake_free_list_tail assert queue.fake_free_list_tail.prev_free_block is blocks[2] # Pop 0 block # fake_head->b1->b3->b5->b4->b0->b2->fake_tail assert len(queue.popleft_n(0)) == 0 assert queue.num_free_blocks == 6 # Pop 1 block # fake_head->b3->b5->b4->b0->b2->fake_tail result_blocks = queue.popleft_n(1) assert queue.num_free_blocks == 5 assert len(result_blocks) == 1 assert result_blocks[0] is blocks[1] for block in result_blocks: assert block.prev_free_block is None assert block.next_free_block is None # Pop 2 blocks # fake_head->b4->b0->b2->fake_tail result_blocks = queue.popleft_n(2) assert len(result_blocks) == 2 assert queue.num_free_blocks == 3 assert result_blocks[0] is blocks[3] assert result_blocks[1] is blocks[5] for block in result_blocks: assert block.prev_free_block is None assert block.next_free_block is None # Pop 3 blocks # fake_head->fake_tail result_blocks = queue.popleft_n(3) assert len(result_blocks) == 3 assert queue.num_free_blocks == 0 assert result_blocks[0] is blocks[4] assert result_blocks[1] is blocks[0] assert result_blocks[2] is blocks[2] for block in result_blocks: assert block.prev_free_block is None assert block.next_free_block is None def test_free_kv_cache_block_queue_get_all_free_blocks(): # Create a list of KVCacheBlock objects blocks = [KVCacheBlock(block_id=i) for i in range(5)] # Create a FreeKVCacheBlockQueue with these blocks queue = FreeKVCacheBlockQueue(blocks) # Check all blocks are correctly retrieved assert queue.get_all_free_blocks() == blocks # Pop a block and check again queue.popleft() assert queue.get_all_free_blocks() == blocks[1:] # Remove a block and check again block_to_remove = blocks[2] queue.remove(block_to_remove) assert queue.get_all_free_blocks() == blocks[1:2] + blocks[3:] # Append a block back and check again queue.append(block_to_remove) assert queue.get_all_free_blocks() == blocks[1:2] + blocks[3:] + [block_to_remove] def test_generate_block_hash_extra_keys(): request = make_request( request_id="0", prompt_token_ids=[_ for _ in range(20)], mm_positions=[ PlaceholderRange(offset=0, length=5), PlaceholderRange(offset=10, length=5), ], mm_hashes=["hash1", "hash2"], ) # Test with no extra keys extra_keys, next_mm_idx = generate_block_hash_extra_keys(request, 0, 5, 0) assert extra_keys == (("mm", "hash1", 0),) assert next_mm_idx == 1 # Test with partial overlap extra_keys, next_mm_idx = generate_block_hash_extra_keys(request, 3, 8, 0) assert extra_keys == (("mm", "hash1", -3),) assert next_mm_idx == 1 # Test with no overlap extra_keys, next_mm_idx = generate_block_hash_extra_keys(request, 6, 10, 0) assert extra_keys is None assert next_mm_idx == 1 # Test with multiple extra keys extra_keys, next_mm_idx = generate_block_hash_extra_keys(request, 0, 15, 0) assert extra_keys == (("mm", "hash1", 0), ("mm", "hash2", 10)) assert next_mm_idx == 2 def test_generate_block_hash_extra_keys_no_mm_inputs(): request = make_request( request_id="0", prompt_token_ids=[_ for _ in range(6)], mm_positions=None, mm_hashes=None, ) extra_keys, next_mm_idx = generate_block_hash_extra_keys(request, 0, 5, 0) assert extra_keys is None assert next_mm_idx == 0 def test_generate_block_hash_extra_keys_cache_salt(): request = make_request( request_id="0", prompt_token_ids=[_ for _ in range(6)], mm_positions=None, mm_hashes=None, cache_salt="salt", ) # salt is added for the first token extra_keys, _ = generate_block_hash_extra_keys(request, 0, 1, 0) assert extra_keys == (("cache_salt", "salt"),) extra_keys, _ = generate_block_hash_extra_keys(request, 0, 10, 0) assert extra_keys == (("cache_salt", "salt"),) # no salt added for other tokens extra_keys, _ = generate_block_hash_extra_keys(request, 1, 2, 0) assert extra_keys is None extra_keys, _ = generate_block_hash_extra_keys(request, 6, 10, 0) assert extra_keys is None # works together with other extra keys request_mm = make_request( request_id="0", prompt_token_ids=[_ for _ in range(20)], mm_positions=[ PlaceholderRange(offset=0, length=5), ], mm_hashes=["hash1"], cache_salt="salt", ) # Test with no extra keys extra_keys, next_mm_idx = generate_block_hash_extra_keys(request_mm, 0, 5, 0) assert extra_keys == (("mm", "hash1", 0), ("cache_salt", "salt")) assert next_mm_idx == 1 def test_generate_block_hash_extra_keys_prompt_embeds(): prompt_embeds = torch.randn(10, 3) request = make_request( request_id="0", prompt_token_ids=None, mm_positions=None, mm_hashes=None, prompt_embeds=prompt_embeds, ) # Test with prompt embeds for the first block extra_keys, _ = generate_block_hash_extra_keys(request, 0, 5, 0) expected_embeds = prompt_embeds[0:5] expected_hash = hashlib.sha256(kv_cache_utils.tensor_data(expected_embeds)).digest() assert extra_keys == (("prompt_embeds", expected_hash),) # Test with prompt embeds for the second block extra_keys, _ = generate_block_hash_extra_keys(request, 5, 10, 0) expected_embeds = prompt_embeds[5:10] expected_hash = hashlib.sha256(kv_cache_utils.tensor_data(expected_embeds)).digest() assert extra_keys == (("prompt_embeds", expected_hash),) def test_generate_block_hash_extra_keys_prompt_embeds_cached(monkeypatch): prompt_embeds = torch.randn(10, 3) request = make_request( request_id="0", prompt_token_ids=None, mm_positions=None, mm_hashes=None, prompt_embeds=prompt_embeds, block_size=20, ) num_tensor_data_calls = 0 original_tensor_data = kv_cache_utils.tensor_data def counting_tensor_data(tensor: torch.Tensor): nonlocal num_tensor_data_calls num_tensor_data_calls += 1 return original_tensor_data(tensor) monkeypatch.setattr(kv_cache_utils, "tensor_data", counting_tensor_data) extra_keys_1, _ = generate_block_hash_extra_keys(request, 0, 5, 0) extra_keys_2, _ = generate_block_hash_extra_keys(request, 0, 5, 0) assert extra_keys_1 == extra_keys_2 assert num_tensor_data_calls == 1 def test_generate_block_hash_extra_keys_different_prompt_embeds(): prompt_embeds1 = torch.randn(10, 3) prompt_embeds2 = torch.randn(10, 3) request1 = make_request( request_id="0", prompt_token_ids=None, mm_positions=None, mm_hashes=None, prompt_embeds=prompt_embeds1, ) request2 = make_request( request_id="1", prompt_token_ids=None, mm_positions=None, mm_hashes=None, prompt_embeds=prompt_embeds2, ) extra_keys1, _ = generate_block_hash_extra_keys(request1, 0, 5, 0) extra_keys2, _ = generate_block_hash_extra_keys(request2, 0, 5, 0) assert extra_keys1 != extra_keys2 def test_generate_block_hash_extra_keys_lora(): request = make_request( request_id="0", prompt_token_ids=[_ for _ in range(6)], ) request.lora_request = LoRARequest( lora_name="test_lora_adapter", lora_int_id=1, lora_path="/path/to/lora" ) extra_keys, _ = generate_block_hash_extra_keys(request, 0, 3, 0) assert extra_keys == (("lora", "test_lora_adapter", "/path/to/lora"),) request.lora_request = None extra_keys, _ = generate_block_hash_extra_keys(request, 0, 3, 0) assert extra_keys is None @pytest.mark.parametrize("hash_fn", [sha256, sha256_cbor]) def test_lora_name_and_cache_salt_block_hashes_do_not_collide(hash_fn): """A cache_salt equal to a LoRA name must not share that LoRA's blocks.""" lora_request = LoRARequest( lora_name="foo", lora_int_id=1, lora_path="/path/to/lora" ) lora_req = Request( request_id="lora", prompt_token_ids=[0, 1, 2], sampling_params=SamplingParams(max_tokens=1), pooling_params=None, lora_request=lora_request, block_hasher=get_request_block_hasher(3, hash_fn), ) salted_req = make_request("salted", [0, 1, 2], hash_fn=hash_fn, cache_salt="foo") assert lora_req.block_hashes[0] != salted_req.block_hashes[0] @pytest.mark.parametrize("hash_fn", [sha256, sha256_cbor]) def test_lora_path_change_changes_block_hashes(hash_fn): """Re-pointing a LoRA name at another adapter must not reuse its blocks.""" def make_lora_request(lora_path: str) -> Request: return Request( request_id=lora_path, prompt_token_ids=[0, 1, 2], sampling_params=SamplingParams(max_tokens=1), pooling_params=None, lora_request=LoRARequest( lora_name="foo", lora_int_id=1, lora_path=lora_path ), block_hasher=get_request_block_hasher(3, hash_fn), ) first = make_lora_request("/path/to/lora_a") second = make_lora_request("/path/to/lora_b") assert first.block_hashes[0] != second.block_hashes[0] def test_to_event_extra_keys_keeps_untagged_event_format(): """KV events keep publishing the extra-key shapes consumers already parse.""" request = make_request( "0", list(range(10)), mm_positions=[PlaceholderRange(offset=2, length=3)], mm_hashes=["hash1"], cache_salt="salt", ) request.lora_request = LoRARequest( lora_name="adapter", lora_int_id=1, lora_path="/path/to/lora" ) extra_keys, _ = generate_block_hash_extra_keys(request, 0, 5, 0) assert extra_keys == ( ("lora", "adapter", "/path/to/lora"), ("mm", "hash1", 2), ("cache_salt", "salt"), ) event_keys = ("adapter", ("hash1", 2), "salt") assert to_event_extra_keys([extra_keys, None]) == [event_keys, None] assert to_event_extra_keys([]) is None assert to_event_extra_keys(None) is None @pytest.mark.parametrize("hash_fn", [sha256, sha256_cbor]) def test_hash_block_tokens(hash_fn): parent_block_hash = BlockHash(b"123") curr_block_token_ids = (1, 2, 3) extra_keys = ("key1", "key2") block_hash = hash_block_tokens( hash_fn, parent_block_hash, curr_block_token_ids, extra_keys ) expected = hash_fn((parent_block_hash, curr_block_token_ids, extra_keys)) assert block_hash == expected @pytest.mark.parametrize("hash_fn", [sha256, sha256_cbor]) def test_request_block_hasher(hash_fn): request = make_request( request_id="0", prompt_token_ids=[_ for _ in range(6)], block_size=3, hash_fn=hash_fn, mm_positions=[ PlaceholderRange(offset=0, length=3), PlaceholderRange(offset=3, length=3), ], mm_hashes=["hash1", "hash2"], ) block_hashes = request.block_hashes assert len(block_hashes) == 2 assert block_hashes[0] == hash_fn( (kv_cache_utils.NONE_HASH, (0, 1, 2), (("mm", "hash1", 0),)) ) assert block_hashes[1] == hash_fn( (block_hashes[0], (3, 4, 5), (("mm", "hash2", 0),)) ) @pytest.mark.parametrize("hash_fn", [sha256, sha256_cbor]) def test_request_block_hasher_incremental_append_with_multiple_mm_features(hash_fn): mm_positions = [ PlaceholderRange(offset=4, length=2), PlaceholderRange(offset=6, length=1), ] incremental = make_request( request_id="incremental", prompt_token_ids=list(range(7)), block_size=4, hash_fn=hash_fn, mm_positions=mm_positions, mm_hashes=["A", "B"], ) incremental.append_output_token_ids(7) fresh = make_request( request_id="fresh", prompt_token_ids=list(range(8)), block_size=4, hash_fn=hash_fn, mm_positions=mm_positions, mm_hashes=["A", "B"], ) expected_second_hash = hash_fn( (incremental.block_hashes[0], (4, 5, 6, 7), (("mm", "A", 0), ("mm", "B", 2))) ) assert incremental.block_hashes[1] == expected_second_hash assert incremental.block_hashes == fresh.block_hashes @pytest.mark.parametrize("hash_fn", [sha256, sha256_cbor]) def test_hash_tokens_different_mm_input(hash_fn): request1 = make_request( request_id="0", prompt_token_ids=[_ for _ in range(6)], block_size=3, hash_fn=hash_fn, mm_positions=[ PlaceholderRange(offset=0, length=3), PlaceholderRange(offset=3, length=3), ], mm_hashes=["hash1", "hash2"], ) request2 = make_request( request_id="1", prompt_token_ids=[_ for _ in range(6)], mm_positions=[ PlaceholderRange(offset=0, length=3), PlaceholderRange(offset=3, length=3), ], mm_hashes=["hash3", "hash2"], ) block_hashes1 = request1.block_hashes block_hashes2 = request2.block_hashes assert block_hashes1[0] != block_hashes2[0] assert block_hashes1[1] != block_hashes2[1] @pytest.mark.parametrize("hash_fn", [sha256, sha256_cbor]) def test_hash_request_tokens_no_mm_inputs(hash_fn): request = make_request( request_id="0", prompt_token_ids=[_ for _ in range(6)], block_size=3, hash_fn=hash_fn, mm_positions=None, mm_hashes=None, ) block_hashes = request.block_hashes assert len(block_hashes) == 2 assert block_hashes[0] == hash_fn((kv_cache_utils.NONE_HASH, (0, 1, 2), None)) assert block_hashes[1] == hash_fn((block_hashes[0], (3, 4, 5), None)) def _stats(requests: int, queries: int, hits: int) -> PrefixCacheStats: return PrefixCacheStats(requests=requests, queries=queries, hits=hits) def test_metrics_empty_distinguishes_no_queries_from_no_hits(): """`hit_rate` alone cannot tell the two apart; `empty` can. Both an unobserved window and a genuine all-miss window report a hit rate of 0.0, so anything surfacing that number to a human has to check `empty` first - which is what the prefix-cache log line does. """ metrics = CachingMetrics(max_recent_requests=5) assert metrics.empty assert metrics.hit_rate == 0.0 metrics.observe(_stats(1, 20, 0)) assert not metrics.empty assert metrics.hit_rate == 0.0 def test_metrics(): """Test the prefix caching metrics.""" metrics = CachingMetrics(max_recent_requests=5) assert metrics.hit_rate == 0.0 metrics.observe(_stats(1, 20, 9)) # 9 / 20 = 0.45 assert metrics.hit_rate == 0.45 metrics.observe(_stats(4, 80, 16)) # 25 / 100 = 0.25 assert metrics.hit_rate == 0.25 metrics.observe(_stats(1, 10, 2)) # Remove (20, 9) and add (10, 2): 18 / 90 = 0.2 assert metrics.aggregated_requests == 5 assert metrics.aggregated_query_total == 90 assert metrics.aggregated_query_hit == 18 assert metrics.hit_rate == 0.2 metrics.reset() assert metrics.hit_rate == 0.0 assert metrics.aggregated_requests == 0 assert metrics.aggregated_query_total == 0 assert metrics.aggregated_query_hit == 0 assert not metrics.query_queue def test_metrics_empty_stats(): """Test the prefix caching metrics with empty stats.""" metrics = CachingMetrics(max_recent_requests=5) metrics.observe(_stats(0, 0, 0)) metrics.observe(_stats(1, 20, 9)) metrics.observe(_stats(0, 0, 0)) metrics.observe(_stats(4, 80, 16)) metrics.observe(_stats(0, 0, 0)) metrics.observe(_stats(1, 10, 2)) # Remove (20, 9) and add (10, 2): 18 / 90 = 0.2 assert metrics.aggregated_requests == 5 assert metrics.aggregated_query_total == 90 assert metrics.aggregated_query_hit == 18 assert metrics.hit_rate == 0.2 # Only the latest added stats preserved 10 / 20 = 0.5 metrics.observe(_stats(11, 20, 10)) assert metrics.aggregated_requests == 11 assert metrics.aggregated_query_total == 20 assert metrics.aggregated_query_hit == 10 assert metrics.hit_rate == 0.5 # Only the latest added stats preserved 30 / 40 = 0.75 metrics.observe(_stats(22, 40, 30)) assert metrics.aggregated_requests == 22 assert metrics.aggregated_query_total == 40 assert metrics.aggregated_query_hit == 30 assert metrics.hit_rate == 0.75 def test_get_kv_cache_configs_multiple_workers(): model_config = ModelConfig(max_model_len=16) vllm_config = VllmConfig(model_config=model_config) vllm_config.cache_config.kv_cache_layout = "LBNHC" vllm_config.cache_config.prefix_cache_retention_interval = None ref_kv_cache_spec = new_kv_cache_spec() same_kv_cache_specs = [ { "layer1": new_kv_cache_spec(), "layer2": new_kv_cache_spec(), }, { "layer1": new_kv_cache_spec(), "layer2": new_kv_cache_spec(), }, ] # Basic case. All things are the same. kv_cache_configs = get_kv_cache_configs( vllm_config, same_kv_cache_specs, [ ref_kv_cache_spec.page_size_bytes * 2 * 10, ref_kv_cache_spec.page_size_bytes * 2 * 10, ], ) expected = KVCacheConfig( num_blocks=10, kv_cache_tensors=[ KVCacheTensor( size=ref_kv_cache_spec.page_size_bytes * 10 * 2, layers=["layer1", "layer2"], layer_stride=ref_kv_cache_spec.page_size_bytes * 10, block_stride=ref_kv_cache_spec.page_size_bytes, ), ], kv_cache_groups=[ KVCacheGroupSpec(["layer1", "layer2"], ref_kv_cache_spec), ], ) assert kv_cache_configs == [expected, expected] # Different available memory. This is the case for TP. # Use the smallest memory available. kv_cache_configs = get_kv_cache_configs( vllm_config, same_kv_cache_specs, [ ref_kv_cache_spec.page_size_bytes * 2 * 10, ref_kv_cache_spec.page_size_bytes * 2 * 20, ], ) assert kv_cache_configs == [expected, expected] # Different KV cache specs. This is the case for PP. different_layer_specs = [ { "layer1": new_kv_cache_spec(), }, { "layer2": new_kv_cache_spec(), "layer3": new_kv_cache_spec(), }, ] # Different workers have different layers. kv_cache_configs = get_kv_cache_configs( vllm_config, different_layer_specs, [ ref_kv_cache_spec.page_size_bytes * 2 * 10, ref_kv_cache_spec.page_size_bytes * 2 * 10, ], ) assert kv_cache_configs == [ KVCacheConfig( num_blocks=10, kv_cache_tensors=[ KVCacheTensor( size=ref_kv_cache_spec.page_size_bytes * 10, layers=["layer1"], layer_stride=ref_kv_cache_spec.page_size_bytes * 10, block_stride=ref_kv_cache_spec.page_size_bytes, ), ], kv_cache_groups=[ KVCacheGroupSpec(["layer1"], new_kv_cache_spec()), ], ), KVCacheConfig( num_blocks=10, kv_cache_tensors=[ KVCacheTensor( size=ref_kv_cache_spec.page_size_bytes * 10 * 2, layers=["layer2", "layer3"], layer_stride=ref_kv_cache_spec.page_size_bytes * 10, block_stride=ref_kv_cache_spec.page_size_bytes, ), ], kv_cache_groups=[ KVCacheGroupSpec(["layer2", "layer3"], new_kv_cache_spec()), ], ), ] # Some layers are the same, some are different. This is the case for TP+PP tp_pp_kv_cache_specs = [ { "layer1": new_kv_cache_spec(), "layer2": new_kv_cache_spec(), }, { "layer1": new_kv_cache_spec(), "layer2": new_kv_cache_spec(), }, { "layer3": new_kv_cache_spec(), }, { "layer3": new_kv_cache_spec(), }, ] kv_cache_configs = get_kv_cache_configs( vllm_config, tp_pp_kv_cache_specs, [ ref_kv_cache_spec.page_size_bytes * 2 * 10, ref_kv_cache_spec.page_size_bytes * 2 * 10, ref_kv_cache_spec.page_size_bytes * 2 * 10, ref_kv_cache_spec.page_size_bytes * 2 * 10, ], ) expected_12 = KVCacheConfig( num_blocks=10, kv_cache_tensors=[ KVCacheTensor( size=ref_kv_cache_spec.page_size_bytes * 10 * 2, layers=["layer1", "layer2"], layer_stride=ref_kv_cache_spec.page_size_bytes * 10, block_stride=ref_kv_cache_spec.page_size_bytes, ), ], kv_cache_groups=[ KVCacheGroupSpec(["layer1", "layer2"], ref_kv_cache_spec), ], ) expected_3 = KVCacheConfig( num_blocks=10, kv_cache_tensors=[ KVCacheTensor( size=ref_kv_cache_spec.page_size_bytes * 10, layers=["layer3"], layer_stride=ref_kv_cache_spec.page_size_bytes * 10, block_stride=ref_kv_cache_spec.page_size_bytes, ), ], kv_cache_groups=[ KVCacheGroupSpec(["layer3"], ref_kv_cache_spec), ], ) assert kv_cache_configs == [ expected_12, expected_12, expected_3, expected_3, ] # Different workers have different types of layers. This is the case for # hybrid models + PP. different_type_layer_specs = [ { "layer1": new_kv_cache_spec(), "layer2": new_kv_cache_spec(), }, { "layer3": new_sliding_window_spec(), "layer4": new_sliding_window_spec(), }, ] kv_cache_configs = get_kv_cache_configs( vllm_config, different_type_layer_specs, [ ref_kv_cache_spec.page_size_bytes * 2 * 10, ref_kv_cache_spec.page_size_bytes * 2 * 10, ], ) assert kv_cache_configs == [ KVCacheConfig( num_blocks=10, kv_cache_tensors=[ KVCacheTensor( size=ref_kv_cache_spec.page_size_bytes * 10 * 2, layers=["layer1", "layer2"], layer_stride=ref_kv_cache_spec.page_size_bytes * 10, block_stride=ref_kv_cache_spec.page_size_bytes, ), ], kv_cache_groups=[ KVCacheGroupSpec(["layer1", "layer2"], ref_kv_cache_spec), KVCacheGroupSpec([], new_sliding_window_spec()), ], ), KVCacheConfig( num_blocks=10, kv_cache_tensors=[ KVCacheTensor( size=ref_kv_cache_spec.page_size_bytes * 10 * 2, layers=["layer3", "layer4"], layer_stride=ref_kv_cache_spec.page_size_bytes * 10, block_stride=ref_kv_cache_spec.page_size_bytes, ), ], kv_cache_groups=[ KVCacheGroupSpec([], ref_kv_cache_spec), KVCacheGroupSpec(["layer3", "layer4"], new_sliding_window_spec()), ], ), ] # When divided into multiple KVCacheGroups, need to ensure the number of # layers per group is similar. different_type_layer_specs = [ { "layer1": new_kv_cache_spec(), "layer2": new_sliding_window_spec(), "layer3": new_sliding_window_spec(), }, { "layer4": new_kv_cache_spec(), "layer5": new_sliding_window_spec(), "layer6": new_sliding_window_spec(), }, ] kv_cache_configs = get_kv_cache_configs( vllm_config, different_type_layer_specs, [ ref_kv_cache_spec.page_size_bytes * 10, ref_kv_cache_spec.page_size_bytes * 10, ], ) assert kv_cache_configs == [ KVCacheConfig( num_blocks=10, kv_cache_tensors=[ KVCacheTensor( size=ref_kv_cache_spec.page_size_bytes * 10, layers=["layer1"], layer_stride=ref_kv_cache_spec.page_size_bytes * 10, block_stride=ref_kv_cache_spec.page_size_bytes, ), KVCacheTensor( size=ref_kv_cache_spec.page_size_bytes * 10, layers=["layer2"], layer_stride=ref_kv_cache_spec.page_size_bytes * 10, block_stride=ref_kv_cache_spec.page_size_bytes, ), KVCacheTensor( size=ref_kv_cache_spec.page_size_bytes * 10, layers=["layer3"], layer_stride=ref_kv_cache_spec.page_size_bytes * 10, block_stride=ref_kv_cache_spec.page_size_bytes, ), ], kv_cache_groups=[ KVCacheGroupSpec(["layer1"], ref_kv_cache_spec), KVCacheGroupSpec(["layer2"], new_sliding_window_spec()), KVCacheGroupSpec(["layer3"], new_sliding_window_spec()), ], ), KVCacheConfig( num_blocks=10, kv_cache_tensors=[ KVCacheTensor( size=ref_kv_cache_spec.page_size_bytes * 10, layers=["layer4"], layer_stride=ref_kv_cache_spec.page_size_bytes * 10, block_stride=ref_kv_cache_spec.page_size_bytes, ), KVCacheTensor( size=ref_kv_cache_spec.page_size_bytes * 10, layers=["layer5"], layer_stride=ref_kv_cache_spec.page_size_bytes * 10, block_stride=ref_kv_cache_spec.page_size_bytes, ), KVCacheTensor( size=ref_kv_cache_spec.page_size_bytes * 10, layers=["layer6"], layer_stride=ref_kv_cache_spec.page_size_bytes * 10, block_stride=ref_kv_cache_spec.page_size_bytes, ), ], kv_cache_groups=[ KVCacheGroupSpec(["layer4"], ref_kv_cache_spec), KVCacheGroupSpec(["layer5"], new_sliding_window_spec()), KVCacheGroupSpec(["layer6"], new_sliding_window_spec()), ], ), ] # Have conflicting layers. Need to raise an error. conflicting_layer_specs = [ { "layer1": new_kv_cache_spec(), }, { "layer1": new_sliding_window_spec(), }, ] with pytest.raises(AssertionError): get_kv_cache_configs( vllm_config, conflicting_layer_specs, [ ref_kv_cache_spec.page_size_bytes * 2 * 10, ref_kv_cache_spec.page_size_bytes * 2 * 10, ], ) @pytest.mark.parametrize( "asymmetric_memory", [False, True], ids=["symmetric", "asymmetric"], ) def test_get_kv_cache_configs_pp_sharding(asymmetric_memory): model_config = ModelConfig(max_model_len=512) vllm_config = VllmConfig(model_config=model_config) vllm_config.cache_config.kv_cache_layout = "LBNHC" vllm_config.cache_config.prefix_cache_retention_interval = None ref_kv_cache_spec = new_kv_cache_spec() pp_kv_cache_specs = [ {"layer1": ref_kv_cache_spec}, {"layer2": ref_kv_cache_spec}, ] expected_num_blocks = model_config.max_model_len // ref_kv_cache_spec.block_size + 1 avail_memory = ref_kv_cache_spec.page_size_bytes * expected_num_blocks # With per-worker validation, each worker only needs memory for its own # layers. Worker 2 having more memory shouldn't affect worker 1's config. available_memory = ( [avail_memory, avail_memory * 2] if asymmetric_memory else [avail_memory] * 2 ) kv_cache_configs = get_kv_cache_configs( vllm_config, pp_kv_cache_specs, available_memory, ) assert kv_cache_configs == [ KVCacheConfig( num_blocks=expected_num_blocks, kv_cache_tensors=[ KVCacheTensor( size=ref_kv_cache_spec.page_size_bytes * expected_num_blocks, layers=["layer1"], layer_stride=ref_kv_cache_spec.page_size_bytes * expected_num_blocks, block_stride=ref_kv_cache_spec.page_size_bytes, ), ], kv_cache_groups=[KVCacheGroupSpec(["layer1"], ref_kv_cache_spec)], ), KVCacheConfig( num_blocks=expected_num_blocks, kv_cache_tensors=[ KVCacheTensor( size=ref_kv_cache_spec.page_size_bytes * expected_num_blocks, layers=["layer2"], layer_stride=ref_kv_cache_spec.page_size_bytes * expected_num_blocks, block_stride=ref_kv_cache_spec.page_size_bytes, ), ], kv_cache_groups=[KVCacheGroupSpec(["layer2"], ref_kv_cache_spec)], ), ] def test_project_kv_cache_groups_to_worker(): spec_a = new_kv_cache_spec() spec_b = new_kv_cache_spec(num_kv_heads=4) global_groups = [ KVCacheGroupSpec(["layer1", "layer2", "layer3"], spec_a), ] worker_spec = {"layer1": spec_a, "layer2": spec_a} projected = kv_cache_utils._project_kv_cache_groups_to_worker( global_groups, worker_spec ) assert len(projected) == 1 assert projected[0].layer_names == ["layer1", "layer2"] assert projected[0].kv_cache_spec is spec_a projected = kv_cache_utils._project_kv_cache_groups_to_worker( global_groups, {"layer4": spec_a} ) assert len(projected) == 1 assert projected[0].layer_names == [] assert projected[0].kv_cache_spec is spec_a uniform_spec = UniformTypeKVCacheSpecs( block_size=16, kv_cache_specs={"layer1": spec_a, "layer2": spec_b, "layer3": spec_a}, ) global_groups_uniform = [ KVCacheGroupSpec(["layer1", "layer2", "layer3"], uniform_spec), ] projected = kv_cache_utils._project_kv_cache_groups_to_worker( global_groups_uniform, {"layer1": spec_a, "layer3": spec_a} ) assert len(projected) == 1 assert projected[0].layer_names == ["layer1", "layer3"] proj_spec = projected[0].kv_cache_spec assert isinstance(proj_spec, UniformTypeKVCacheSpecs) assert set(proj_spec.kv_cache_specs.keys()) == {"layer1", "layer3"} @pytest.mark.parametrize("sliding_window", [None, 256]) @pytest.mark.parametrize("disable_hybrid", [False, True]) @pytest.mark.parametrize("pcp_size", [1, 4]) def test_dcp_target_allocates_replicated_draft_independently( monkeypatch, sliding_window, disable_hybrid, pcp_size ): """A draft must retain all positions even when the target shards them.""" from vllm.v1.attention.backends.utils import resolve_kv_cache_layout monkeypatch.delenv("VLLM_KV_CACHE_LAYOUT", raising=False) config = VllmConfig(model_config=ModelConfig(max_model_len=1024)) config.parallel_config.decode_context_parallel_size = 4 config.parallel_config.prefill_context_parallel_size = pcp_size config.scheduler_config.disable_hybrid_kv_cache_manager = disable_hybrid config.cache_config.block_size = 16 config.cache_config.kv_cache_layout = None draft_args = dict( block_size=16, num_kv_heads=8, head_size=64, dtype=torch.bfloat16, dcp_sharded=False, ) draft = ( FullAttentionSpec(**draft_args) if sliding_window is None else SlidingWindowSpec(**draft_args, sliding_window=sliding_window) ) specs = {"target": new_mla_spec(), "draft": draft} layout = resolve_kv_cache_layout(config, [["LBHNC", "BLHNC"]], specs.values()) assert layout == KVCacheLayout.BLHNC groups = get_kv_cache_groups(config, specs) assert [group.layer_names for group in groups] == [["target"], ["draft"]] widths = [g.kv_cache_spec.max_num_blocks_per_req(config, 1024) for g in groups] assert widths == [16, 64] cache_config = kv_cache_utils.get_kv_cache_config_from_groups( config, groups, available_memory=16 * 1024 * 1024 ) scheduler_config = generate_scheduler_kv_cache_config([cache_config]) assert get_max_concurrency_for_kv_cache_config(config, cache_config) > 0 assert get_max_concurrency_for_kv_cache_config(config, cache_config) == ( get_max_concurrency_for_kv_cache_config(config, scheduler_config) ) manager = KVCacheManager( scheduler_config, max_model_len=1024, hash_block_size=16, scheduler_block_size=64, dcp_world_size=4, pcp_world_size=pcp_size, enable_caching=True, ) assert manager.coordinator.group_block_sizes == (64, 16) request = make_request("replicated-draft", [1] * 65, block_size=16, hash_fn=sha256) blocks = manager.allocate_slots(request, 65) assert blocks is not None assert [len(group) for group in blocks.blocks] == [2, 5] manager.cache_blocks(request, 64) cached_request = make_request( "cached-draft", [1] * 65, block_size=16, hash_fn=sha256 ) cached_blocks, num_cached, _ = manager.get_computed_blocks(cached_request) assert num_cached == 64 assert [len(group) for group in cached_blocks.blocks] == [1, 4] @pytest.mark.parametrize("use_mla", [False, True]) def test_full_attention_merge_preserves_replicated_cache_geometry(use_mla): config = VllmConfig(model_config=ModelConfig(max_model_len=1024)) config.parallel_config.decode_context_parallel_size = 4 draft = replace( new_mla_spec() if use_mla else new_kv_cache_spec(), block_size=16, dcp_sharded=False, ) merged = type(draft).merge([draft, draft]) wrapped = UniformTypeKVCacheSpecs.from_specs({"draft": merged}) assert wrapped is not None and not wrapped.dcp_sharded assert merged.max_num_blocks_per_req(config, 1024) == 64 assert merged.max_memory_usage_bytes(config) == 64 * draft.page_size_bytes assert kv_cache_utils.resolve_dcp_kv_block_size(merged, 4) == 16 with pytest.raises(AssertionError): type(draft).merge([draft, replace(draft, dcp_sharded=True)]) @pytest.mark.parametrize("with_draft", [False, True]) @pytest.mark.parametrize("indexer_alignment", [None, 512]) def test_sparse_mla_preserves_physical_row_addressing(with_draft, indexer_alignment): from vllm.v1.attention.backends.mla.sparse_utils import flat_kv_row_view from vllm.v1.worker.utils import allocate_kv_cache config = VllmConfig(model_config=ModelConfig(max_model_len=1024)) config.cache_config.kv_cache_layout = "BLHNC" config.parallel_config.decode_context_parallel_size = 4 common = dict(block_size=64, num_kv_heads=1, dtype=torch.uint8) specs = { "target": MLAAttentionSpec( **common, head_size=576, state_content_bytes=656, block_stride_alignment=656, cache_dtype_str="fp8_ds_mla", is_index_group_leader=True, ), "indexer": MLAAttentionSpec( **common, head_size=132, cache_role=SparseCacheRole.INDEXER, block_stride_alignment=indexer_alignment, ), "draft": SlidingWindowSpec( block_size=64, num_kv_heads=64, head_size=64, dtype=torch.bfloat16, sliding_window=2048, dcp_sharded=False, ), } if not with_draft: del specs["draft"] groups = get_kv_cache_groups(config, specs) cache_config = kv_cache_utils.get_kv_cache_config_from_groups( config, groups, available_memory=8 * 1024 * 1024 ) caches = allocate_kv_cache( cache_config, torch.device("cpu"), KVCacheLayout.BLHNC, [64] * len(groups) ) cache = caches["target"].squeeze(1) if indexer_alignment is not None: assert caches["indexer"].stride(0) % indexer_alignment == 0 rows, stride_rows = flat_kv_row_view(cache, 64) cache[1, 0].fill_(7) torch.testing.assert_close(rows[stride_rows], cache[1, 0]) assert torch.all(rows[stride_rows] == 7) @pytest.mark.parametrize( "layer_type,dcp_size,expected_width", [ ("mla", 1, 64), ("mla", 2, 32), # Mamba state is replicated, not DCP-sharded, and its width is the # resident state block count rather than cdiv(max_len, block_size). ("mamba", 2, 3), ], ) def test_uniform_type_spec_block_table_width_matches_layer_spec( layer_type, dcp_size, expected_width ): # The runner sizes the block table from the group spec while the metadata # builders are constructed from the per-layer spec, so the aggregate must # report the same width as the layers it wraps. vllm_config = VllmConfig(model_config=ModelConfig(max_model_len=1024)) vllm_config.parallel_config.decode_context_parallel_size = dcp_size layer_spec = new_mla_spec() if layer_type == "mla" else new_mamba_spec() uniform_spec = UniformTypeKVCacheSpecs( block_size=layer_spec.block_size, kv_cache_specs={"layer1": layer_spec, "layer2": layer_spec}, ) assert layer_spec.max_num_blocks_per_req(vllm_config, 1024) == expected_width assert uniform_spec.max_num_blocks_per_req(vllm_config, 1024) == expected_width def test_merge_kv_cache_spec(): same_layer_specs = [ new_kv_cache_spec(num_kv_heads=32), new_kv_cache_spec(num_kv_heads=32), ] merged_layer_spec = same_layer_specs[0].merge(same_layer_specs) assert merged_layer_spec.block_size == 16 assert merged_layer_spec.num_kv_heads == 32 assert merged_layer_spec.head_size == 64 assert merged_layer_spec.dtype == torch.float32 assert merged_layer_spec.sliding_window is None different_layer_specs = [ new_kv_cache_spec(num_kv_heads=32), new_kv_cache_spec(num_kv_heads=16), ] with pytest.raises(AssertionError): different_layer_specs[0].merge(different_layer_specs) full_spec = new_kv_cache_spec(num_kv_heads=32) different_type_layer_specs = [ full_spec, SlidingWindowSpec( block_size=full_spec.block_size, num_kv_heads=full_spec.num_kv_heads, head_size=full_spec.head_size, dtype=full_spec.dtype, sliding_window=1, ), ] with pytest.raises(AssertionError): different_type_layer_specs[0].merge(different_type_layer_specs) with pytest.raises(AssertionError): different_type_layer_specs[1].merge(different_type_layer_specs) different_sliding_window_layer_specs = [ new_kv_cache_spec(num_kv_heads=32), new_kv_cache_spec(num_kv_heads=32, sliding_window=1), new_kv_cache_spec(num_kv_heads=32, sliding_window=2), ] with pytest.raises(ValueError): different_sliding_window_layer_specs[0].merge( different_sliding_window_layer_specs ) same_sliding_window_layer_specs = [ new_kv_cache_spec(num_kv_heads=32, sliding_window=1), new_kv_cache_spec(num_kv_heads=32, sliding_window=1), ] merged_layer_spec = same_sliding_window_layer_specs[0].merge( same_sliding_window_layer_specs ) assert merged_layer_spec.sliding_window == 1 same_sliding_window_layer_spec_with_none = [ new_kv_cache_spec(num_kv_heads=32, sliding_window=1), new_kv_cache_spec(num_kv_heads=32, sliding_window=None), ] merged_layer_spec = same_sliding_window_layer_spec_with_none[0].merge( same_sliding_window_layer_spec_with_none ) assert merged_layer_spec.sliding_window == 1 def test_is_kv_cache_spec_uniform(): kv_cache_spec = { "layer_1": new_kv_cache_spec(num_kv_heads=32), "layer_2": new_kv_cache_spec(num_kv_heads=32), } assert is_kv_cache_spec_uniform(kv_cache_spec) kv_cache_spec = { "layer_1": new_kv_cache_spec(num_kv_heads=32), "layer_2": new_kv_cache_spec(num_kv_heads=32, sliding_window=1), } assert is_kv_cache_spec_uniform(kv_cache_spec) kv_cache_spec = { "layer_1": new_kv_cache_spec(num_kv_heads=32), "layer_2": new_sliding_window_spec(num_kv_heads=32, sliding_window=1), } assert not is_kv_cache_spec_uniform(kv_cache_spec) kv_cache_spec = { "layer_1": new_sliding_window_spec(num_kv_heads=32, sliding_window=1), "layer_2": new_sliding_window_spec(num_kv_heads=32, sliding_window=1), } assert is_kv_cache_spec_uniform(kv_cache_spec) kv_cache_spec = { "layer_1": new_sliding_window_spec(num_kv_heads=32, sliding_window=1), "layer_2": new_sliding_window_spec(num_kv_heads=32, sliding_window=2), } assert not is_kv_cache_spec_uniform(kv_cache_spec) script = """ import sys from vllm.v1.core.kv_cache_utils import is_kv_cache_spec_uniform from vllm.v1.kv_cache_interface import KVCacheSpec if sys.flags.optimize < 1: raise RuntimeError("subprocess is not running with optimization enabled") specs = { "a": KVCacheSpec(block_size=1), "b": KVCacheSpec(block_size=2), } if is_kv_cache_spec_uniform(specs): raise RuntimeError("different specs were treated as uniform") """ subprocess.run( [sys.executable, "-O", "-c", script], cwd=Path(__file__).resolve().parents[3], check=True, timeout=60, ) @pytest.mark.parametrize( ("model_id", "max_model_len", "want_estimated_max_len"), [ ("Qwen/Qwen1.5-7B", 16385, 16384), ("Qwen/Qwen1.5-7B", 16383, 16383), ], ) def test_estimate_max_model_len(model_id, max_model_len, want_estimated_max_len): # Create a VllmConfig model_config = ModelConfig( model_id, runner="generate", dtype="float16", max_model_len=max_model_len, ) scheduler_config = SchedulerConfig( max_num_batched_tokens=32768, max_model_len=model_config.max_model_len, is_encoder_decoder=model_config.is_encoder_decoder, ) vllm_config = VllmConfig( model_config=model_config, scheduler_config=scheduler_config, ) # Create KV cache specs kv_cache_spec = {} for i in range(32): layer_name = f"layer_{i}" kv_cache_spec[layer_name] = FullAttentionSpec( block_size=16, num_kv_heads=32, head_size=128, dtype=torch.float16, ) # Estimate the maximum model length, 16384 model_len need 8GB estimated_max_len = estimate_max_model_len( vllm_config, kv_cache_spec, 8 * GiB_bytes ) assert estimated_max_len == want_estimated_max_len def test_get_max_concurrency_for_kv_cache_config(): # Create a VllmConfig model_id = "Qwen/Qwen1.5-7B" max_model_len = 16384 model_config = ModelConfig( model_id, runner="generate", dtype="float16", max_model_len=max_model_len, ) scheduler_config = SchedulerConfig( max_num_batched_tokens=1024, enable_chunked_prefill=True, max_model_len=model_config.max_model_len, is_encoder_decoder=model_config.is_encoder_decoder, # Pin to sync: SWA per-request bounds grow with overlapping batches. async_scheduling=False, ) vllm_config = VllmConfig( model_config=model_config, scheduler_config=scheduler_config, ) assert vllm_config.max_concurrent_batches == 1 full_attention_spec = FullAttentionSpec( block_size=16, num_kv_heads=32, head_size=128, dtype=torch.float16, ) sliding_window_spec = SlidingWindowSpec( block_size=16, num_kv_heads=32, head_size=128, dtype=torch.float16, sliding_window=1024, ) kv_cache_config_full_attention = KVCacheConfig( num_blocks=int(1024 * 1.5), kv_cache_tensors=[], kv_cache_groups=[ KVCacheGroupSpec([f"layer_{i}" for i in range(32)], full_attention_spec), ], ) max_concurrency_full_attention = get_max_concurrency_for_kv_cache_config( vllm_config, kv_cache_config_full_attention ) assert max_concurrency_full_attention == 1.5 kv_cache_config_sliding_window = KVCacheConfig( num_blocks=129 * 3, kv_cache_tensors=[], kv_cache_groups=[ KVCacheGroupSpec([f"layer_{i}" for i in range(32)], sliding_window_spec), ], ) max_concurrency_sliding_window = get_max_concurrency_for_kv_cache_config( vllm_config, kv_cache_config_sliding_window ) assert max_concurrency_sliding_window == 3 kv_cache_config_hybrid_model = KVCacheConfig( num_blocks=(1024 + 129) * 3, kv_cache_tensors=[], kv_cache_groups=[ KVCacheGroupSpec([f"layer_{i}" for i in range(32)], full_attention_spec), KVCacheGroupSpec( [f"layer_{i}" for i in range(32, 64)], sliding_window_spec ), ], ) max_concurrency_hybrid_model = get_max_concurrency_for_kv_cache_config( vllm_config, kv_cache_config_hybrid_model ) assert max_concurrency_hybrid_model == 3 num_tokens, max_concurrency = get_kv_cache_capacity( vllm_config, kv_cache_config_hybrid_model ) assert num_tokens == max_concurrency_hybrid_model * max_model_len assert max_concurrency == max_concurrency_hybrid_model # Unequal group sizes in the standard layout: each group's pages cost # whole pool blocks, so a request needs 1024 + 129 = 1153 blocks — the # same as the equal-hybrid case above, regardless of the second group # holding only 2 layers. kv_cache_config_unequal_groups = KVCacheConfig( num_blocks=1153 * 3, kv_cache_tensors=[], kv_cache_groups=[ KVCacheGroupSpec([f"layer_{i}" for i in range(32)], full_attention_spec), KVCacheGroupSpec(["layer_32", "layer_33"], sliding_window_spec), ], ) assert ( get_max_concurrency_for_kv_cache_config( vllm_config, kv_cache_config_unequal_groups ) == 3 ) # UniformTypeKVCacheSpecs group (worker config shape): the aggregated # spec's memory/page ratio equals a single layer's page count, so the # group needs 1024 blocks and the request 1153 in total. The previous # formula normalized both groups' memory by the first group's page size, # reporting 3459/1057 = 3.27 here instead of 3 — and a different value # again for the scheduler-config shape below. uniform_full_spec = UniformTypeKVCacheSpecs( block_size=full_attention_spec.block_size, kv_cache_specs={f"layer_{i}": full_attention_spec for i in range(4)}, ) kv_cache_config_uniform_group = KVCacheConfig( num_blocks=1153 * 3, kv_cache_tensors=[], kv_cache_groups=[ KVCacheGroupSpec([f"layer_{i}" for i in range(4)], uniform_full_spec), KVCacheGroupSpec(["layer_4", "layer_5"], sliding_window_spec), ], ) assert ( get_max_concurrency_for_kv_cache_config( vllm_config, kv_cache_config_uniform_group ) == 3 ) # Scheduler-config shape: generate_scheduler_kv_cache_config replaces the # uniform-type group's spec with a representative per-layer spec. # Capacity must not change between the two shapes (the engine computes # on the scheduler config, the worker loop on the worker config). kv_cache_config_scheduler_shape = generate_scheduler_kv_cache_config( [copy.deepcopy(kv_cache_config_uniform_group)] ) assert get_max_concurrency_for_kv_cache_config( vllm_config, kv_cache_config_scheduler_shape ) == get_max_concurrency_for_kv_cache_config( vllm_config, kv_cache_config_uniform_group ) def test_allocate_with_lookahead(): """Verify that lookahead tokens correctly affect block allocation.""" block_size = 4 config = KVCacheConfig( num_blocks=10, kv_cache_tensors=[ KVCacheTensor( size=100, layers=["layer1"], layer_stride=100, block_stride=10, ), ], kv_cache_groups=[ KVCacheGroupSpec(["layer1"], new_kv_cache_spec(block_size=block_size)), ], ) request = make_request( request_id="0", prompt_token_ids=[], block_size=block_size, mm_positions=None, mm_hashes=None, ) # Test case 1: Requires additional lookahead tokens kv_cache_manager = KVCacheManager( kv_cache_config=config, max_model_len=100, scheduler_block_size=block_size, hash_block_size=block_size, ) blocks = kv_cache_manager.allocate_slots( request, num_new_tokens=3, num_lookahead_tokens=2, # Total required: 3+2=5 tokens ) assert len(blocks.get_block_ids()[0]) == 2 # ceil(5/4)=2 blocks # Test case 2: With precomputed blocks kv_cache_manager = KVCacheManager( kv_cache_config=config, max_model_len=100, scheduler_block_size=block_size, hash_block_size=block_size, ) # required_blocks = ceil((3 + 2) /4) = 2 blocks = kv_cache_manager.allocate_slots( request, num_new_tokens=3, num_lookahead_tokens=2, ) assert len(blocks.get_block_ids()[0]) == 2 # Test case 3: With precomputed blocks # required_blocks = ceil((3 + 4) / 4) = 2 kv_cache_manager = KVCacheManager( kv_cache_config=config, max_model_len=100, scheduler_block_size=block_size, hash_block_size=block_size, ) blocks = kv_cache_manager.allocate_slots( request, num_new_tokens=3, num_lookahead_tokens=4, ) assert len(blocks.get_block_ids()[0]) == 2 def test_get_kv_cache_config_one_worker(): # pass max_model_len to pass check_enough_kv_cache_memory model_config = ModelConfig(max_model_len=16) vllm_config = VllmConfig(model_config=model_config) vllm_config.cache_config.kv_cache_layout = "LBNHC" vllm_config.cache_config.prefix_cache_retention_interval = None mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2 # all layers are full attention -> single group kv_cache_specs_full = { "layer_1": new_kv_cache_spec(), "layer_2": new_kv_cache_spec(), } kv_cache_config_full = get_kv_cache_configs( vllm_config, [kv_cache_specs_full], [mem_per_block_per_layer * 2 * 32] )[0] print(kv_cache_config_full) assert kv_cache_config_full == KVCacheConfig( num_blocks=32, kv_cache_tensors=[ KVCacheTensor( size=mem_per_block_per_layer * 32 * 2, layers=["layer_1", "layer_2"], layer_stride=mem_per_block_per_layer * 32, block_stride=mem_per_block_per_layer, ), ], kv_cache_groups=[KVCacheGroupSpec(["layer_1", "layer_2"], new_kv_cache_spec())], ) # all layers are sliding window -> single group kv_cache_specs_sliding = { "layer_1": new_sliding_window_spec(), "layer_2": new_sliding_window_spec(), } kv_cache_config_sliding = get_kv_cache_configs( vllm_config, [kv_cache_specs_sliding], [mem_per_block_per_layer * 2 * 32] )[0] assert kv_cache_config_sliding == KVCacheConfig( num_blocks=32, kv_cache_tensors=[ KVCacheTensor( size=mem_per_block_per_layer * 32 * 2, layers=["layer_1", "layer_2"], layer_stride=mem_per_block_per_layer * 32, block_stride=mem_per_block_per_layer, ), ], kv_cache_groups=[ KVCacheGroupSpec(["layer_1", "layer_2"], new_sliding_window_spec()) ], ) # full + sliding, but disable_hybrid_kv_cache_manager vllm_config.scheduler_config.disable_hybrid_kv_cache_manager = True kv_cache_specs_hybrid = { "layer_1": new_kv_cache_spec(), "layer_2": new_sliding_window_spec(), } kv_cache_config_hybrid = get_kv_cache_configs( vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 2 * 32] )[0] assert kv_cache_config_hybrid == KVCacheConfig( num_blocks=32, kv_cache_tensors=[ KVCacheTensor( size=mem_per_block_per_layer * 32 * 2, layers=["layer_1", "layer_2"], layer_stride=mem_per_block_per_layer * 32, block_stride=mem_per_block_per_layer, ), ], kv_cache_groups=[ KVCacheGroupSpec( ["layer_1", "layer_2"], new_kv_cache_spec(sliding_window=1) ), ], ) vllm_config.scheduler_config.disable_hybrid_kv_cache_manager = False # full + sliding, with hybrid_kv_cache_manager kv_cache_specs_hybrid = { "layer_1": new_kv_cache_spec(), "layer_2": new_sliding_window_spec(), } kv_cache_config_hybrid = get_kv_cache_configs( vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 2 * 32] )[0] assert kv_cache_config_hybrid == KVCacheConfig( num_blocks=64, kv_cache_tensors=[ KVCacheTensor( size=mem_per_block_per_layer * 64, layers=["layer_1"], layer_stride=mem_per_block_per_layer * 64, block_stride=mem_per_block_per_layer, ), KVCacheTensor( size=mem_per_block_per_layer * 64, layers=["layer_2"], layer_stride=mem_per_block_per_layer * 64, block_stride=mem_per_block_per_layer, ), ], kv_cache_groups=[ KVCacheGroupSpec(["layer_1"], new_kv_cache_spec()), KVCacheGroupSpec(["layer_2"], new_sliding_window_spec()), ], ) # 2 full + 4 sliding, 2 layers per group kv_cache_specs_hybrid = { "layer_1": new_kv_cache_spec(), "layer_2": new_kv_cache_spec(), "layer_3": new_sliding_window_spec(), "layer_4": new_sliding_window_spec(), "layer_5": new_sliding_window_spec(), "layer_6": new_sliding_window_spec(), } kv_cache_config_hybrid = get_kv_cache_configs( vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 2 * 32] )[0] assert kv_cache_config_hybrid == KVCacheConfig( num_blocks=32, kv_cache_tensors=[ KVCacheTensor( size=mem_per_block_per_layer * 32 * 2, layers=["layer_1", "layer_2"], layer_stride=mem_per_block_per_layer * 32, block_stride=mem_per_block_per_layer, ), KVCacheTensor( size=mem_per_block_per_layer * 32 * 2, layers=["layer_3", "layer_5"], layer_stride=mem_per_block_per_layer * 32, block_stride=mem_per_block_per_layer, ), KVCacheTensor( size=mem_per_block_per_layer * 32 * 2, layers=["layer_4", "layer_6"], layer_stride=mem_per_block_per_layer * 32, block_stride=mem_per_block_per_layer, ), ], kv_cache_groups=[ KVCacheGroupSpec(["layer_1", "layer_2"], new_kv_cache_spec()), KVCacheGroupSpec(["layer_3", "layer_5"], new_sliding_window_spec()), KVCacheGroupSpec(["layer_4", "layer_6"], new_sliding_window_spec()), ], ) # 3 full + 7 sliding, pad to 3 full + 9 sliding kv_cache_specs_hybrid = { "layer_1": new_kv_cache_spec(), "layer_2": new_kv_cache_spec(), "layer_3": new_kv_cache_spec(), "layer_4": new_sliding_window_spec(), "layer_5": new_sliding_window_spec(), "layer_6": new_sliding_window_spec(), "layer_7": new_sliding_window_spec(), "layer_8": new_sliding_window_spec(), "layer_9": new_sliding_window_spec(), "layer_10": new_sliding_window_spec(), } kv_cache_config_hybrid = get_kv_cache_configs( vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 3 * 32] )[0] assert kv_cache_config_hybrid == KVCacheConfig( num_blocks=32, kv_cache_tensors=[ KVCacheTensor( size=mem_per_block_per_layer * 32 * 3, layers=["layer_1", "layer_2", "layer_3"], layer_stride=mem_per_block_per_layer * 32, block_stride=mem_per_block_per_layer, ), KVCacheTensor( size=mem_per_block_per_layer * 32 * 3, layers=["layer_4", "layer_7", "layer_10"], layer_stride=mem_per_block_per_layer * 32, block_stride=mem_per_block_per_layer, ), KVCacheTensor( size=mem_per_block_per_layer * 32 * 3, layers=["layer_5", "layer_8"], layer_stride=mem_per_block_per_layer * 32, block_stride=mem_per_block_per_layer, ), KVCacheTensor( size=mem_per_block_per_layer * 32 * 3, layers=["layer_6", "layer_9"], layer_stride=mem_per_block_per_layer * 32, block_stride=mem_per_block_per_layer, ), ], kv_cache_groups=[ KVCacheGroupSpec(["layer_1", "layer_2", "layer_3"], new_kv_cache_spec()), KVCacheGroupSpec( ["layer_4", "layer_7", "layer_10"], new_sliding_window_spec() ), KVCacheGroupSpec(["layer_5", "layer_8"], new_sliding_window_spec()), KVCacheGroupSpec(["layer_6", "layer_9"], new_sliding_window_spec()), ], ) # 6 full + 5 sliding, pad to 6 full + 6 sliding. This is a typical case for gpt-oss # eagle where there is only one more full attention layer than sliding window layers kv_cache_specs_hybrid = { "layer_1": new_kv_cache_spec(), "layer_2": new_kv_cache_spec(), "layer_3": new_kv_cache_spec(), "layer_4": new_kv_cache_spec(), "layer_5": new_kv_cache_spec(), "layer_6": new_kv_cache_spec(), "layer_7": new_sliding_window_spec(), "layer_8": new_sliding_window_spec(), "layer_9": new_sliding_window_spec(), "layer_10": new_sliding_window_spec(), "layer_11": new_sliding_window_spec(), } kv_cache_config_hybrid = get_kv_cache_configs( vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 6 * 32] )[0] print(kv_cache_config_hybrid) assert kv_cache_config_hybrid == KVCacheConfig( num_blocks=32, kv_cache_tensors=[ KVCacheTensor( size=mem_per_block_per_layer * 32 * 6, layers=[ "layer_1", "layer_2", "layer_3", "layer_4", "layer_5", "layer_6", ], layer_stride=mem_per_block_per_layer * 32, block_stride=mem_per_block_per_layer, ), KVCacheTensor( size=mem_per_block_per_layer * 32 * 6, layers=["layer_7", "layer_8", "layer_9", "layer_10", "layer_11"], layer_stride=mem_per_block_per_layer * 32, block_stride=mem_per_block_per_layer, ), ], kv_cache_groups=[ KVCacheGroupSpec( ["layer_1", "layer_2", "layer_3", "layer_4", "layer_5", "layer_6"], new_kv_cache_spec(), ), KVCacheGroupSpec( ["layer_7", "layer_8", "layer_9", "layer_10", "layer_11"], new_sliding_window_spec(), ), ], ) # different hidden size but same type, use UniformTypeKVCacheSpecs kv_cache_specs_hybrid = { "layer_1": new_kv_cache_spec(head_size=128), "layer_2": new_kv_cache_spec(head_size=64), } kv_cache_config_hybrid = get_kv_cache_configs( vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 3 * 32] )[0] # Layers of different page sizes pack densely into one allocation: the # 2x-larger layer_1 takes the first region, layer_2 the next. assert kv_cache_config_hybrid == KVCacheConfig( num_blocks=32, kv_cache_tensors=[ KVCacheTensor( size=mem_per_block_per_layer * 32 * 3, layers=["layer_1"], layer_stride=mem_per_block_per_layer * 2 * 32, block_stride=mem_per_block_per_layer * 2, ), KVCacheTensor( size=mem_per_block_per_layer * 32 * 3, layers=["layer_2"], layer_stride=mem_per_block_per_layer * 32, block_stride=mem_per_block_per_layer, offset=mem_per_block_per_layer * 2 * 32, ), ], kv_cache_groups=[ KVCacheGroupSpec( ["layer_1", "layer_2"], UniformTypeKVCacheSpecs( block_size=16, kv_cache_specs=kv_cache_specs_hybrid ), ) ], ) # Different hidden size and different type, align by different block size kv_cache_specs_hybrid = { "layer_1": new_kv_cache_spec(head_size=64), "layer_2": new_sliding_window_spec(head_size=32), } kv_cache_config_hybrid = get_kv_cache_configs( vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 32] )[0] assert kv_cache_config_hybrid == KVCacheConfig( num_blocks=32, kv_cache_tensors=[ KVCacheTensor( size=mem_per_block_per_layer * 32, layers=["layer_1"], layer_stride=mem_per_block_per_layer * 32, block_stride=mem_per_block_per_layer, ), KVCacheTensor( size=mem_per_block_per_layer * 32, layers=["layer_2"], layer_stride=mem_per_block_per_layer * 32, block_stride=mem_per_block_per_layer, ), ], kv_cache_groups=[ KVCacheGroupSpec(["layer_1"], new_kv_cache_spec(head_size=64)), KVCacheGroupSpec( ["layer_2"], new_sliding_window_spec(head_size=32, block_size=32) ), ], ) # different hidden size that cannot be aligned by using different block size, # but can be aligned by padding the smaller physical page. swa_spec = new_sliding_window_spec(head_size=96) kv_cache_specs_hybrid = { "layer_1": new_kv_cache_spec(head_size=64), "layer_2": swa_spec, } kv_cache_config_hybrid = get_kv_cache_configs( vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 2 * 32] )[0] padded_page_size = swa_spec.page_size_bytes assert kv_cache_config_hybrid == KVCacheConfig( num_blocks=42, kv_cache_tensors=[ KVCacheTensor( size=padded_page_size * 42, layers=["layer_1"], layer_stride=padded_page_size * 42, block_stride=padded_page_size, ), KVCacheTensor( size=padded_page_size * 42, layers=["layer_2"], layer_stride=padded_page_size * 42, block_stride=padded_page_size, ), ], kv_cache_groups=[ KVCacheGroupSpec( ["layer_1"], new_kv_cache_spec( head_size=64, page_size_padded=padded_page_size, ), ), KVCacheGroupSpec( ["layer_2"], new_sliding_window_spec(head_size=96), ), ], ) # Test num_gpu_blocks_override vllm_config.cache_config.num_gpu_blocks_override = 16 kv_cache_config_override_blocks = get_kv_cache_configs( vllm_config, [kv_cache_specs_full], [mem_per_block_per_layer * 2 * 32] )[0] assert kv_cache_config_override_blocks == KVCacheConfig( num_blocks=16, kv_cache_tensors=[ KVCacheTensor( size=mem_per_block_per_layer * 16 * 2, layers=["layer_1", "layer_2"], layer_stride=mem_per_block_per_layer * 16, block_stride=mem_per_block_per_layer, ), ], kv_cache_groups=[KVCacheGroupSpec(["layer_1", "layer_2"], new_kv_cache_spec())], ) def test_get_kv_cache_configs_attention_free(): kv_cache_specs: dict[str, KVCacheSpec] = {} vllm_config = VllmConfig(model_config=ModelConfig(max_model_len=16)) vllm_config.cache_config.prefix_cache_retention_interval = None kv_cache_configs = get_kv_cache_configs(vllm_config, [kv_cache_specs], [0]) assert kv_cache_configs == [ KVCacheConfig( num_blocks=1, kv_cache_tensors=[], kv_cache_groups=[], ) ] def test_generate_uniform_type_kv_cache_specs(): # All layers are full attention, can be merged kv_cache_specs = { "layer_1": new_kv_cache_spec(), "layer_2": new_kv_cache_spec(head_size=128), } uniform_spec = UniformTypeKVCacheSpecs.from_specs(kv_cache_specs) assert uniform_spec == UniformTypeKVCacheSpecs( block_size=16, kv_cache_specs=kv_cache_specs ) # Full attention + sliding window, cannot be merged kv_cache_specs = { "layer_1": new_kv_cache_spec(), "layer_2": new_sliding_window_spec(sliding_window=1), } uniform_spec = UniformTypeKVCacheSpecs.from_specs(kv_cache_specs) assert uniform_spec is None # different order of full attention + sliding window, cannot be merged kv_cache_specs = { "layer_1": new_sliding_window_spec(sliding_window=1), "layer_2": new_kv_cache_spec(), } uniform_spec = UniformTypeKVCacheSpecs.from_specs(kv_cache_specs) assert uniform_spec is None # Same-size sliding window, can be merged kv_cache_specs = { "layer_1": new_sliding_window_spec(sliding_window=1), "layer_2": new_sliding_window_spec(sliding_window=1, head_size=128), } uniform_spec = UniformTypeKVCacheSpecs.from_specs(kv_cache_specs) assert uniform_spec == UniformTypeKVCacheSpecs( block_size=16, kv_cache_specs=kv_cache_specs ) # different block sizes, cannot be merged kv_cache_specs = { "layer_1": new_kv_cache_spec(block_size=16), "layer_2": new_kv_cache_spec(block_size=32), } uniform_spec = UniformTypeKVCacheSpecs.from_specs(kv_cache_specs) assert uniform_spec is None def test_generate_scheduler_kv_cache_config(): kv_cache_specs = { "layer_1": new_kv_cache_spec(), "layer_2": new_kv_cache_spec(head_size=128), } kv_cache_configs = [ KVCacheConfig( num_blocks=10, kv_cache_tensors=[], kv_cache_groups=[ KVCacheGroupSpec( ["layer_1", "layer_2"], UniformTypeKVCacheSpecs( block_size=16, kv_cache_specs=kv_cache_specs ), ), ], ) ] scheduler_kv_cache_config = generate_scheduler_kv_cache_config(kv_cache_configs) assert scheduler_kv_cache_config == KVCacheConfig( num_blocks=10, kv_cache_tensors=[], kv_cache_groups=[KVCacheGroupSpec(["layer_1", "layer_2"], new_kv_cache_spec())], ) def _glm5_like_kv_cache_spec( mamba_spec_factory=new_mamba_spec, ) -> tuple[dict[str, KVCacheSpec], list[str]]: """(mamba, mamba, mamba, MLA + indexer) * 11 plus a trailing mamba.""" kv_cache_spec: dict[str, KVCacheSpec] = {} mamba_layers = [] for i in range(45): if i % 4 == 3: kv_cache_spec[f"layers.{i}.attn"] = MLAAttentionSpec( block_size=1024, num_kv_heads=1, head_size=576, dtype=torch.bfloat16, ) kv_cache_spec[f"layers.{i}.indexer"] = MLAAttentionSpec( block_size=1024, num_kv_heads=1, head_size=132, dtype=torch.uint8, tokens_per_state=16, ) else: name = f"layers.{i}.linear_attn" kv_cache_spec[name] = mamba_spec_factory() mamba_layers.append(name) return kv_cache_spec, mamba_layers def _glm5_like_kv_cache_spec_with_tail( mamba_spec_factory=new_mamba_spec, ) -> dict[str, KVCacheSpec]: """Production kpool=4 proportions: (mamba*3, MLA + indexer + tail) * 11. The tail's logical page (2 * kpool * 2*indexer_head_dim * bf16) must fit inside the indexer page it parasitizes (block_size//kpool * 132 B), which is why the fixture drops the tokens_per_state=16 shape used above. """ kv_cache_spec, _ = _glm5_like_kv_cache_spec(mamba_spec_factory) for i in range(3, 45, 4): kpool = 4 kv_cache_spec[f"layers.{i}.indexer"] = replace( cast(MLAAttentionSpec, kv_cache_spec[f"layers.{i}.indexer"]), tokens_per_state=kpool, ) kv_cache_spec[f"layers.{i}.tail"] = KpoolTailSpec( block_size=kpool, num_kv_heads=2, head_size=128, head_size_v=0, dtype=torch.bfloat16, sliding_window=kpool, ) return kv_cache_spec def _tensor_by_layer(kv_cache_config: KVCacheConfig) -> dict[str, KVCacheTensor]: return { layer_name: tensor for tensor in kv_cache_config.kv_cache_tensors for layer_name in tensor.layers } def _layer_offset(tensor: KVCacheTensor, layer_name: str) -> int: return tensor.offset + tensor.layers.index(layer_name) * tensor.layer_stride def test_get_kv_cache_config_balanced_mamba_hybrid(): """Hybrid slot sharing: mamba layers co-own the MLA slot tensors.""" model_config = ModelConfig(max_model_len=8192) vllm_config = VllmConfig(model_config=model_config) kv_cache_spec, mamba_layers = _glm5_like_kv_cache_spec() mla_page = kv_cache_spec["layers.3.attn"].page_size_bytes idx_page = kv_cache_spec["layers.3.indexer"].page_size_bytes groups = kv_cache_utils.get_kv_cache_groups(vllm_config, kv_cache_spec) uniform_groups = [ group for group in groups if isinstance(group.kv_cache_spec, UniformTypeKVCacheSpecs) ] mamba_groups = [ group for group in groups if isinstance(group.kv_cache_spec, MambaSpec) ] # Round-robin into G = ceil(34 / 11) = 4 groups so every mamba layer # gets an MLA slot. assert len(groups) == 5 assert len(uniform_groups) == 1 assert len(uniform_groups[0].layer_names) == 22 assert uniform_groups[0].kv_cache_spec.get_max_layers_per_page_size() == 11 assert [len(group.layer_names) for group in mamba_groups] == [9, 9, 8, 8] for k, name in enumerate(mamba_layers): assert name in mamba_groups[k % 4].layer_names # Mamba pages are padded up to the MLA page. for group in mamba_groups: assert group.kv_cache_spec.page_size_padded == mla_page assert group.kv_cache_spec.page_size_bytes == mla_page bytes_per_block = kv_cache_utils._pool_bytes_per_block(groups) assert bytes_per_block == 11 * mla_page + 11 * idx_page # Every block id is charged the full per-block sum. attn_blocks = uniform_groups[0].kv_cache_spec.max_memory_usage_pages(vllm_config) mamba_blocks_per_group = 1 + new_mamba_spec().num_speculative_blocks blocks_per_request = attn_blocks + 4 * mamba_blocks_per_group assert ( kv_cache_utils._max_memory_usage_bytes_from_groups(vllm_config, groups) == blocks_per_request * bytes_per_block ) available_memory = bytes_per_block * 100 + 1 kv_cache_config = kv_cache_utils.get_kv_cache_config_from_groups( vllm_config, groups, available_memory ) assert kv_cache_config.num_blocks == 100 # Every logical layer has a view into one backing allocation. Mamba views # alias their corresponding MLA slot by using the same byte offset. assert len(kv_cache_config.kv_cache_tensors) == 56 assert {t.size for t in kv_cache_config.kv_cache_tensors} == {bytes_per_block * 100} tensors = _tensor_by_layer(kv_cache_config) mla_layer_names = [f"layers.{i}.attn" for i in range(3, 45, 4)] for i, mla_name in enumerate(mla_layer_names): mla_tensor = tensors[mla_name] assert mla_tensor.block_stride == mla_page for group in mamba_groups: if i < len(group.layer_names): assert tensors[group.layer_names[i]].offset == mla_tensor.offset for i in range(11): assert tensors[f"layers.{4 * i + 3}.indexer"].block_stride == idx_page total_allocated = next(iter({t.size for t in kv_cache_config.kv_cache_tensors})) assert total_allocated == bytes_per_block * 100 assert 0 <= available_memory - total_allocated < bytes_per_block assert get_max_concurrency_for_kv_cache_config( vllm_config, kv_cache_config ) == pytest.approx(100 / blocks_per_request) def test_get_kv_cache_config_kpool_tail_coowns_indexer_tensor(): """The kpool tail parasitizes the indexer tensors instead of getting its own: sibling idx/tail tensors paired by layer order, zero standalone tail bytes, one shared block per request, and no prefix-caching leakage from the kpool-sized scratch group.""" model_config = ModelConfig(max_model_len=8192) vllm_config = VllmConfig(model_config=model_config) kv_cache_spec = _glm5_like_kv_cache_spec_with_tail() mla_page = kv_cache_spec["layers.3.attn"].page_size_bytes idx_page = kv_cache_spec["layers.3.indexer"].page_size_bytes tail_logical_page = kv_cache_spec["layers.3.tail"].page_size_bytes assert tail_logical_page < idx_page groups = kv_cache_utils.get_kv_cache_groups(vllm_config, kv_cache_spec) tail_group = next( group for group in groups if isinstance(group.kv_cache_spec, UniformTypeKVCacheSpecs) and all( isinstance(spec, KpoolTailSpec) for spec in group.kv_cache_spec.kv_cache_specs.values() ) ) # The tail never prefix-caches; its page is padded up to the indexer page # so the runner's strided view rides the indexer storage. assert not tail_group.kv_cache_spec.prefix_cacheable tail_inner = cast( KpoolTailSpec, tail_group.kv_cache_spec.kv_cache_specs["layers.3.tail"] ) assert tail_inner.page_size_padded == idx_page bytes_per_block = kv_cache_utils._pool_bytes_per_block(groups) assert bytes_per_block == 11 * mla_page + 11 * idx_page kv_cache_config = kv_cache_utils.get_kv_cache_config_from_groups( vllm_config, groups, bytes_per_block * 100 + 1 ) assert kv_cache_config.num_blocks == 100 # Tail layers get logical views but no additional storage: each view aliases # its sibling indexer at the same offset in the shared allocation. assert len(kv_cache_config.kv_cache_tensors) == 67 tensors = _tensor_by_layer(kv_cache_config) for i in range(11): idx_name = f"layers.{4 * i + 3}.indexer" tail_name = f"layers.{4 * i + 3}.tail" assert tensors[idx_name].offset == tensors[tail_name].offset assert {t.size for t in kv_cache_config.kv_cache_tensors} == {bytes_per_block * 100} # The layout detector surfaces the sibling names in layer order for the # accounting and connector paths. layout = kv_cache_utils._glm5_next_tensor_layout(kv_cache_config.kv_cache_groups) assert layout is not None assert layout[3] == [f"layers.{4 * i + 3}.indexer" for i in range(11)] assert layout[6] == [f"layers.{4 * i + 3}.tail" for i in range(11)] # Accounting charges exactly one extra shared block per request for the # tail, not a per-sequence allocation. attn_group = next( group for group in groups if isinstance(group.kv_cache_spec, UniformTypeKVCacheSpecs) and group is not tail_group ) attn_blocks = attn_group.kv_cache_spec.max_memory_usage_pages(vllm_config) mamba_blocks_per_group = 1 + new_mamba_spec().num_speculative_blocks blocks_per_request = attn_blocks + 4 * mamba_blocks_per_group + 1 assert ( kv_cache_utils._max_memory_usage_bytes_from_groups(vllm_config, groups) == blocks_per_request * bytes_per_block ) def test_glm5_kpool_tail_does_not_drag_hash_block_size(): """The tail's kpool-sized scratch block (4 tokens) must not constrain the prefix-cache hash granularity: participating groups alone decide it.""" model_config = ModelConfig(max_model_len=8192) vllm_config = VllmConfig(model_config=model_config) def align_mamba(): return new_mamba_spec(mamba_cache_mode="align") groups = kv_cache_utils.get_kv_cache_groups( vllm_config, _glm5_like_kv_cache_spec_with_tail(align_mamba) ) kv_cache_config = KVCacheConfig( num_blocks=1, kv_cache_tensors=[], kv_cache_groups=groups, ) hash_vllm_config = SimpleNamespace( cache_config=SimpleNamespace( block_size=16, enable_prefix_caching=True, prefix_match_unit=None, ), parallel_config=SimpleNamespace(decode_context_parallel_size=1), kv_transfer_config=object(), ) # gcd(attn 1024, mamba 16) with the tail's 4 excluded; scheduler size is # the lcm including the tail (1024 % 4 == 0, so it coincides). assert kv_cache_utils.resolve_kv_cache_block_sizes( kv_cache_config, hash_vllm_config ) == (1024, 16) def test_get_kv_cache_config_mamba_hybrid_sharing_infeasible(): """Reject GLM-5.3-Flash layouts whose Mamba page exceeds the MLA page.""" model_config = ModelConfig(max_model_len=8192) vllm_config = VllmConfig(model_config=model_config) # (512, 1024) fp32 state = 2 MiB per page > 1.125 MiB MLA page. def big_mamba_spec(): return new_mamba_spec(shapes=((512, 1024),), dtypes=(torch.float32,)) kv_cache_spec, _ = _glm5_like_kv_cache_spec(big_mamba_spec) mla_page = kv_cache_spec["layers.3.attn"].page_size_bytes assert big_mamba_spec().page_size_bytes > mla_page with pytest.raises(ValueError, match="does not fit the MLA page"): kv_cache_utils.get_kv_cache_groups(vllm_config, kv_cache_spec) def test_get_kv_cache_config_mamba_hybrid_sharing_infeasible_no_indexer(): """Use the generic layout error when no kpool indexer is present.""" model_config = ModelConfig(max_model_len=8192) vllm_config = VllmConfig(model_config=model_config) vllm_config.cache_config.kv_cache_layout = "LBHNC" kv_cache_spec: dict[str, KVCacheSpec] = {} for i in range(27): if i % 4 == 3 or i == 26: kv_cache_spec[f"layers.{i}.attn"] = MLAAttentionSpec( block_size=1024, num_kv_heads=1, head_size=576, dtype=torch.bfloat16, ) else: kv_cache_spec[f"layers.{i}.linear_attn"] = new_mamba_spec( shapes=((512, 1024),), dtypes=(torch.float32,) ) mla_page = kv_cache_spec["layers.3.attn"].page_size_bytes assert kv_cache_spec["layers.0.linear_attn"].page_size_bytes > mla_page with pytest.raises(NotImplementedError, match="page size"): kv_cache_utils.get_kv_cache_groups(vllm_config, kv_cache_spec) def test_get_kv_cache_config_mamba_hybrid_sharing_prepadded_mamba(): """Platform-prepadded mamba pages (mamba_page_size_padded hint) must not disable slot sharing; the layout re-pads them to the MLA page.""" model_config = ModelConfig(max_model_len=8192) vllm_config = VllmConfig(model_config=model_config) def prepadded_mamba_spec(): return new_mamba_spec(page_size_padded=294_912) kv_cache_spec, _ = _glm5_like_kv_cache_spec(prepadded_mamba_spec) mla_page = kv_cache_spec["layers.3.attn"].page_size_bytes assert prepadded_mamba_spec().page_size_bytes < mla_page groups = kv_cache_utils.get_kv_cache_groups(vllm_config, kv_cache_spec) mamba_groups = [ group for group in groups if isinstance(group.kv_cache_spec, MambaSpec) ] assert len(groups) == 5 assert [len(group.layer_names) for group in mamba_groups] == [9, 9, 8, 8] for group in mamba_groups: assert group.kv_cache_spec.page_size_bytes == mla_page bytes_per_block = kv_cache_utils._pool_bytes_per_block(groups) kv_cache_config = kv_cache_utils.get_kv_cache_config_from_groups( vllm_config, groups, bytes_per_block * 100 + 1 ) assert kv_cache_config.num_blocks == 100 assert len(kv_cache_config.kv_cache_tensors) == 56 def test_get_kv_cache_config_mamba_hybrid_sharing_pp_balanced_projection(): """Round-robin mamba grouping keeps practical PP splits balanced: every stage's largest projected mamba group slice fits its projected MLA layers, so per-stage slot tensors all have an MLA owner.""" model_config = ModelConfig(max_model_len=8192) vllm_config = VllmConfig(model_config=model_config) kv_cache_spec, _ = _glm5_like_kv_cache_spec() global_groups = kv_cache_utils.get_kv_cache_groups(vllm_config, kv_cache_spec) mla_page = kv_cache_spec["layers.3.attn"].page_size_bytes idx_page = kv_cache_spec["layers.3.indexer"].page_size_bytes # PP=2 split at transformer layer 22/23. # Stage 0 of a PP=2 split at transformer layer 22/23: 5 MLA(+indexer) # layers, projected mamba groups [5, 5, 4, 4]. worker_spec = {n: s for n, s in kv_cache_spec.items() if int(n.split(".")[1]) <= 22} groups = kv_cache_utils._project_kv_cache_groups_to_worker( global_groups, worker_spec ) mamba_groups = [ group for group in groups if isinstance(group.kv_cache_spec, MambaSpec) ] assert [len(group.layer_names) for group in mamba_groups] == [5, 5, 4, 4] bytes_per_block = kv_cache_utils._pool_bytes_per_block(groups) assert bytes_per_block == 5 * mla_page + 5 * idx_page kv_cache_config = kv_cache_utils.get_kv_cache_config_from_groups( vllm_config, groups, bytes_per_block * 100 + 1 ) assert kv_cache_config.num_blocks == 100 tensors = _tensor_by_layer(kv_cache_config) # Every projected Mamba slot aliases an MLA owner. assert tensors["layers.3.attn"].offset == tensors["layers.0.linear_attn"].offset assert tensors["layers.3.attn"].offset == tensors["layers.1.linear_attn"].offset assert tensors["layers.3.attn"].offset == tensors["layers.2.linear_attn"].offset assert tensors["layers.3.attn"].offset == tensors["layers.4.linear_attn"].offset # The last slot only hosts the two groups with a 5th projected layer. assert tensors["layers.19.attn"].offset == tensors["layers.21.linear_attn"].offset assert tensors["layers.19.attn"].offset == tensors["layers.22.linear_attn"].offset assert {t.size for t in kv_cache_config.kv_cache_tensors} == {bytes_per_block * 100} def test_get_kv_cache_config_mamba_hybrid_sharing_pp_group_count_bump(monkeypatch): """PP=4's default partition [11,11,12,11] leaves stage 0 with 2 MLA layers but a round-robin slice of 3 under the minimum 4 mamba groups; grouping bumps to 5 groups so every stage's projection keeps sharing on instead of silently falling back on stage 0.""" from vllm.config import ParallelConfig from vllm.distributed.utils import get_pp_indices monkeypatch.setattr(ModelConfig, "get_total_num_hidden_layers", lambda self: 45) vllm_config = VllmConfig( model_config=ModelConfig(max_model_len=8192), parallel_config=ParallelConfig(pipeline_parallel_size=4), ) kv_cache_spec, mamba_layers = _glm5_like_kv_cache_spec() groups = kv_cache_utils.get_kv_cache_groups(vllm_config, kv_cache_spec) mamba_groups = [g for g in groups if isinstance(g.kv_cache_spec, MambaSpec)] assert [len(g.layer_names) for g in mamba_groups] == [7, 7, 7, 7, 6] for k, name in enumerate(mamba_layers): assert name in mamba_groups[k % 5].layer_names for rank in range(4): start, end = get_pp_indices(45, rank, 4) worker_spec = { n: s for n, s in kv_cache_spec.items() if start <= int(n.split(".")[1]) < end } projected = kv_cache_utils._project_kv_cache_groups_to_worker( groups, worker_spec ) assert kv_cache_utils._glm5_next_tensor_layout(projected) is not None def test_get_kv_cache_config_mamba_hybrid_sharing_pp_starved_stage(monkeypatch): """Reject PP stages whose Mamba layers have no MLA slot to share.""" from vllm.config import ParallelConfig monkeypatch.setattr(ModelConfig, "get_total_num_hidden_layers", lambda self: 45) monkeypatch.setenv("VLLM_PP_LAYER_PARTITION", "3,42") vllm_config = VllmConfig( model_config=ModelConfig(max_model_len=8192), parallel_config=ParallelConfig(pipeline_parallel_size=2), ) kv_cache_spec, _ = _glm5_like_kv_cache_spec() with pytest.raises(ValueError, match="VLLM_PP_LAYER_PARTITION"): kv_cache_utils.get_kv_cache_groups(vllm_config, kv_cache_spec) def test_get_kv_cache_config_mamba_hybrid_sharing_beats_cross_layers_flag(): """Hybrid slot sharing must take precedence over the experimental enable_cross_layers_blocks packed layout: generic packing would give the MLA slots a strided view, breaking the contiguous virtual split.""" model_config = ModelConfig(max_model_len=8192) kv_transfer_config = KVTransferConfig( kv_connector="NixlConnector", kv_role="kv_both", kv_connector_extra_config={"enable_cross_layers_blocks": "true"}, ) vllm_config = VllmConfig( model_config=model_config, kv_transfer_config=kv_transfer_config ) kv_cache_spec, _ = _glm5_like_kv_cache_spec() groups = kv_cache_utils.get_kv_cache_groups(vllm_config, kv_cache_spec) mla_page = kv_cache_spec["layers.3.attn"].page_size_bytes idx_page = kv_cache_spec["layers.3.indexer"].page_size_bytes bytes_per_block = kv_cache_utils._pool_bytes_per_block(groups) assert bytes_per_block == 11 * mla_page + 11 * idx_page kv_cache_config = kv_cache_utils.get_kv_cache_config_from_groups( vllm_config, groups, bytes_per_block * 100 + 1 ) assert kv_cache_config.num_blocks == 100 assert len(kv_cache_config.kv_cache_tensors) == 56 assert all( tensor.block_stride in (mla_page, idx_page) for tensor in kv_cache_config.kv_cache_tensors ) assert {t.size for t in kv_cache_config.kv_cache_tensors} == {bytes_per_block * 100} def test_get_kv_cache_config_mamba_hybrid_sharing_no_indexer(): """Kimi-Linear-like: MLA without indexer layers, idx_stride == 0.""" model_config = ModelConfig(max_model_len=8192) vllm_config = VllmConfig(model_config=model_config) vllm_config.cache_config.kv_cache_layout = "LBNHC" # 20 mamba + 7 MLA layers -> G = ceil(20 / 7) = 3 groups of [7, 7, 6]. kv_cache_spec: dict[str, KVCacheSpec] = {} for i in range(27): if i % 4 == 3 or i == 26: kv_cache_spec[f"layers.{i}.attn"] = MLAAttentionSpec( block_size=1024, num_kv_heads=1, head_size=576, dtype=torch.bfloat16, ) else: kv_cache_spec[f"layers.{i}.linear_attn"] = new_mamba_spec() mla_page = kv_cache_spec["layers.3.attn"].page_size_bytes groups = kv_cache_utils.get_kv_cache_groups(vllm_config, kv_cache_spec) mamba_groups = [ group for group in groups if isinstance(group.kv_cache_spec, MambaSpec) ] assert len(groups) == 4 assert [len(group.layer_names) for group in mamba_groups] == [7, 7, 6] for group in mamba_groups: assert group.kv_cache_spec.page_size_bytes == mla_page bytes_per_block = kv_cache_utils._pool_bytes_per_block(groups) assert bytes_per_block == 7 * mla_page kv_cache_config = kv_cache_utils.get_kv_cache_config_from_groups( vllm_config, groups, bytes_per_block * 100 + 1 ) assert kv_cache_config.num_blocks == 100 # Each group has one tensor descriptor and all groups alias the same slots. assert len(kv_cache_config.kv_cache_tensors) == 4 tensors = _tensor_by_layer(kv_cache_config) mla_names = [f"layers.{i}.attn" for i in (*range(3, 27, 4), 26)] for index, mla_name in enumerate(mla_names): mla_offset = _layer_offset(tensors[mla_name], mla_name) for group in mamba_groups: if index < len(group.layer_names): mamba_name = group.layer_names[index] assert _layer_offset(tensors[mamba_name], mamba_name) == mla_offset assert {t.size for t in kv_cache_config.kv_cache_tensors} == {bytes_per_block * 100} def test_get_kv_cache_capacity_after_scheduler_unwrap(): """max_concurrency must survive the scheduler-config unwrap. Regression for the balanced-mamba hybrid layout: ``generate_scheduler_kv_cache_config`` flattens the MLA ``UniformTypeKVCacheSpecs`` group into a single ``MLAAttentionSpec``, so the scheduler config holds an MLA spec (page size A) next to ``MambaSpec`` groups (page size B). ``EngineCore._initialize_kv_caches`` calls ``get_kv_cache_capacity`` on that unwrapped config; the uniform-page-size path used to assert-fail (``assert len(page_sizes) == 1``) on this topology. """ model_config = ModelConfig(max_model_len=8192) vllm_config = VllmConfig(model_config=model_config) kv_cache_spec: dict[str, KVCacheSpec] = {} for i in range(45): if i % 4 == 3: kv_cache_spec[f"layers.{i}.attn"] = MLAAttentionSpec( block_size=1024, num_kv_heads=1, head_size=576, dtype=torch.bfloat16, ) kv_cache_spec[f"layers.{i}.indexer"] = MLAAttentionSpec( block_size=1024, num_kv_heads=1, head_size=132, dtype=torch.uint8, tokens_per_state=16, ) else: kv_cache_spec[f"layers.{i}.linear_attn"] = new_mamba_spec() groups = kv_cache_utils.get_kv_cache_groups(vllm_config, kv_cache_spec) bytes_per_block = kv_cache_utils._pool_bytes_per_block(groups) kv_cache_config = kv_cache_utils.get_kv_cache_config_from_groups( vllm_config, groups, bytes_per_block * 100 + 1 ) scheduler_config = generate_scheduler_kv_cache_config([kv_cache_config]) # Confirm we are exercising the regression path: the MLA group is no longer # UniformTypeKVCacheSpecs, so sibling specs have differing page sizes. assert not any( isinstance(g.kv_cache_spec, UniformTypeKVCacheSpecs) for g in scheduler_config.kv_cache_groups ) unwrapped_groups = scheduler_config.kv_cache_groups expected_max_mem = kv_cache_utils._max_memory_usage_bytes_from_groups( vllm_config, unwrapped_groups ) expected_pool = kv_cache_utils._pool_bytes_per_block(unwrapped_groups) expected_blocks_per_request = ( expected_max_mem + expected_pool - 1 ) // expected_pool _, max_concurrency = get_kv_cache_capacity(vllm_config, scheduler_config) assert max_concurrency > 0 assert max_concurrency == pytest.approx( scheduler_config.num_blocks / expected_blocks_per_request ) def new_mla_spec(cache_dtype_str=None, block_size: int = 16): # head_size = kv_lora_rank(512) + qk_rope_head_dim(64) = 576 return MLAAttentionSpec( block_size=block_size, num_kv_heads=1, max_tp_shards=1, head_size=576, dtype=torch.float32, cache_dtype_str=cache_dtype_str, ) def new_swa_mla_spec(head_size=576, sliding_window=128, model_version=None): return SlidingWindowMLASpec( block_size=16, num_kv_heads=1, max_tp_shards=1, head_size=head_size, dtype=torch.float32, sliding_window=sliding_window, model_version=model_version, ) def new_indexer_mla_spec(block_size=16): # Sparse-attention indexer k_cache: an MLAAttentionSpec with a much smaller # page size than the main MLA attention (uint8, small head), so their pages # cannot be unified. return MLAAttentionSpec( block_size=block_size, num_kv_heads=1, head_size=132, dtype=torch.uint8, ) def test_mixed_page_size_groups_use_spec_compatibility(): specs = {} for i in range(3): specs[f"mla.{i}"] = new_mla_spec() specs[f"indexer.{i}"] = new_indexer_mla_spec() specs.update({f"swa.{i}": new_swa_mla_spec(head_size=1024) for i in range(5)}) config = _grouping_config() config.cache_config = CacheConfig() config.cache_config.kv_cache_layout = "BLNHC" groups = get_kv_cache_groups(config, specs) assert len(groups) == 3 assert {name for group in groups for name in group.layer_names} == set(specs) assert sorted(len(group.layer_names) for group in groups) == [2, 3, 6] def _grouping_config(): cache_config = CacheConfig() cache_config.kv_cache_layout = "LBNHC" return SimpleNamespace( scheduler_config=SimpleNamespace(disable_hybrid_kv_cache_manager=False), speculative_config=None, cache_config=cache_config, ) def test_hidden_state_group_preserves_hybrid_prefix_cache_granularity(): block_size = 544 full_spec = FullAttentionSpec( block_size=block_size, num_kv_heads=1, # FullAttentionSpec stores both K and V, so this produces a # 544 * 1 * (512 + 512) * 2 = 1,114,112-byte page. head_size=512, dtype=torch.bfloat16, ) mamba_spec = MambaSpec( block_size=block_size, shapes=((557056,),), dtypes=(torch.bfloat16,), mamba_cache_mode="align", ) hidden_spec = HiddenStateCacheSpec( block_size=block_size, num_kv_heads=3, head_size=1024, dtype=torch.bfloat16, ) assert full_spec.page_size_bytes == mamba_spec.page_size_bytes groups = get_kv_cache_groups( _grouping_config(), { "model.full_attn": full_spec, "model.mamba": mamba_spec, "cache_only_layers.0": hidden_spec, }, ) hidden_group = next( group for group in groups if isinstance(group.kv_cache_spec, HiddenStateCacheSpec) ) assert hidden_group.kv_cache_spec.block_size == 136 kv_cache_config = KVCacheConfig( num_blocks=1, kv_cache_tensors=[], kv_cache_groups=groups, ) vllm_config = SimpleNamespace( cache_config=SimpleNamespace( block_size=16, enable_prefix_caching=True, prefix_match_unit=None, ), parallel_config=SimpleNamespace(decode_context_parallel_size=1), kv_transfer_config=object(), ) assert kv_cache_utils.resolve_kv_cache_block_sizes( kv_cache_config, vllm_config ) == (544, 136) def test_resolve_dcp_kv_block_size_unwraps_uniform_type_specs(): attention = FullAttentionSpec( block_size=16, num_kv_heads=1, head_size=64, dtype=torch.float16, ) mamba = MambaSpec( block_size=16, shapes=((1, 1),), dtypes=(torch.float16,), ) wrapped_attention = UniformTypeKVCacheSpecs( block_size=16, kv_cache_specs={"attention.0": attention, "attention.1": attention}, ) wrapped_mamba = UniformTypeKVCacheSpecs( block_size=16, kv_cache_specs={"mamba.0": mamba, "mamba.1": mamba}, ) assert kv_cache_utils.resolve_dcp_kv_block_size(attention, 4) == 64 assert kv_cache_utils.resolve_dcp_kv_block_size(wrapped_attention, 4) == 64 assert kv_cache_utils.resolve_dcp_kv_block_size(mamba, 4) == 16 assert kv_cache_utils.resolve_dcp_kv_block_size(wrapped_mamba, 4) == 16 scaled_attention = kv_cache_utils.resolve_dcp_kv_cache_spec(wrapped_attention, 4) assert scaled_attention.block_size == 64 assert isinstance(scaled_attention, UniformTypeKVCacheSpecs) assert all( spec.block_size == 64 for spec in scaled_attention.kv_cache_specs.values() ) assert kv_cache_utils.resolve_dcp_kv_cache_spec(wrapped_mamba, 4) is wrapped_mamba def test_multi_run_layer_compact_strides_place_hoisted_heads(): """A layer-compact run region is its own dense allocation: under LHBNC the head groups sit between the layers and the blocks, so a run's block stride is one head group's slice, not the whole page (regression test for setStorage out-of-bounds on ROCm hybrid models).""" full = new_kv_cache_spec() swa = new_sliding_window_spec(sliding_window=full.block_size * 2) page = full.page_size_bytes assert swa.page_size_bytes == page vllm_config = VllmConfig(model_config=ModelConfig(max_model_len=16)) vllm_config.cache_config.kv_cache_layout = "LHBNC" config = kv_cache_utils.get_kv_cache_config_from_groups( vllm_config, [ KVCacheGroupSpec(["full.0", "full.1"], full), KVCacheGroupSpec(["swa.0", "swa.1"], swa), ], available_memory=4 * page * 8, ) num_blocks = config.num_blocks head_group = full.block_size * full.state_content_size_bytes for tensor in config.kv_cache_tensors: # One head group per block, head groups between the layers and blocks. assert tensor.block_stride == head_group assert tensor.layer_stride == full.num_heads * num_blocks * head_group # The last block of the last layer's last head group must stay within the # allocation (the old page-based stride overran it). end = ( tensor.offset + tensor.layer_stride * (len(tensor.layers) - 1) + (full.num_heads - 1) * num_blocks * head_group + tensor.block_stride * (num_blocks - 1) + head_group ) assert end <= tensor.size def test_mla_draft_prefers_standard_layout_when_pages_can_be_unified(): specs = { "target.0.attn": new_mla_spec(), "draft.0": new_sliding_window_spec(num_kv_heads=1, head_size=288), } assert len({spec.page_size_bytes for spec in specs.values()}) == 1 groups = get_kv_cache_groups(_grouping_config(), specs) assert len(groups) == 2 assert all( not isinstance(group.kv_cache_spec, UniformTypeKVCacheSpecs) for group in groups ) def test_mla_with_incompatible_swa_uses_one_full_allocation_group(caplog_vllm): # Sparse MLA pages cannot be padded safely. Keeping the draft's attention # compute sliding-window while promoting only its allocation semantics lets # every layer share the target's block table and remain contiguous. draft = new_sliding_window_spec(block_size=16) specs = { "target.0.attn": new_mla_spec(block_size=64), "target.0.indexer": new_indexer_mla_spec(block_size=64), "draft.0": draft, } groups = get_kv_cache_groups(_grouping_config(), specs) assert len(groups) == 1 assert set(groups[0].layer_names) == set(specs) group_spec = groups[0].kv_cache_spec assert isinstance(group_spec, UniformTypeKVCacheSpecs) assert group_spec.block_size == 64 promoted_draft = group_spec.kv_cache_specs["draft.0"] assert isinstance(promoted_draft, FullAttentionSpec) assert not isinstance(promoted_draft, SlidingWindowSpec) assert promoted_draft.block_size == 64 assert promoted_draft.sliding_window == draft.sliding_window assert specs["draft.0"] is draft assert "attention compute is unchanged" in caplog_vllm.text def test_hidden_states_with_tp_scales_page_size(): """When TP shrinks KV pages below the hidden-state per-token cost, get_kv_cache_groups must scale up target block sizes so that the common page accommodates the unsharded hidden states.""" # Simulate TP=4 sharding a model with 8 KV heads → 2 per rank. # KV page = block_size(16) * num_kv_heads(2) * head_size(64) * dtype(2) # = 16 * 2 * 64 * 2 = 4096 bytes. kv_spec = new_kv_cache_spec( block_size=16, num_kv_heads=2, head_size=64, dtype=torch.bfloat16, ) # Hidden-state per-token cost = num_hidden_states(6) * hidden_size(512) # * dtype(2) = 6144 bytes, which exceeds the 4096-byte KV page. hs_spec = HiddenStateCacheSpec( block_size=16, num_kv_heads=6, head_size=512, dtype=torch.bfloat16, ) specs = { "target.0.attn": kv_spec, "target.1.attn": kv_spec, "cache_only_layers.48": hs_spec, } groups = get_kv_cache_groups(_grouping_config(), specs) # The hidden-state layer should be present and no assertion should fire. all_layers = {name for g in groups for name in g.layer_names} assert "cache_only_layers.48" in all_layers # The target group block sizes must have been scaled up. for g in groups: if "cache_only_layers.48" not in g.layer_names: assert g.kv_cache_spec.block_size > kv_spec.block_size def test_get_kv_cache_spec_kind_prefers_specific_attention_subclasses(): assert get_kv_cache_spec_kind(new_mla_spec()) == KVCacheSpecKind.MLA_ATTENTION sliding_window_mla_spec = SlidingWindowMLASpec( block_size=16, num_kv_heads=1, head_size=576, dtype=torch.float32, sliding_window=128, ) assert ( get_kv_cache_spec_kind(sliding_window_mla_spec) == KVCacheSpecKind.SLIDING_WINDOW_MLA ) sink_full_attention_spec = SinkFullAttentionSpec( block_size=16, num_kv_heads=1, head_size=64, dtype=torch.float32, sink_len=4, ) assert ( get_kv_cache_spec_kind(sink_full_attention_spec) == KVCacheSpecKind.SINK_FULL_ATTENTION ) def test_get_kv_cache_spec_kind_unwraps_uniform_type_specs(): uniform_mla_spec = UniformTypeKVCacheSpecs( block_size=16, kv_cache_specs={ "layer_1": new_mla_spec(), "layer_2": new_mla_spec(cache_dtype_str="fp8"), }, ) assert get_kv_cache_spec_kind(uniform_mla_spec) == KVCacheSpecKind.MLA_ATTENTION uniform_swa_mla_spec = UniformTypeKVCacheSpecs( block_size=16, kv_cache_specs={ "layer_1": SlidingWindowMLASpec( block_size=16, num_kv_heads=1, head_size=576, dtype=torch.float32, sliding_window=128, ), "layer_2": SlidingWindowMLASpec( block_size=16, num_kv_heads=1, head_size=1024, dtype=torch.float32, sliding_window=128, ), }, ) assert ( get_kv_cache_spec_kind(uniform_swa_mla_spec) == KVCacheSpecKind.SLIDING_WINDOW_MLA ) def test_get_kv_cache_spec_kind_unknown_for_mixed_uniform_type_specs(): uniform_mixed_spec = UniformTypeKVCacheSpecs( block_size=16, kv_cache_specs={ "layer_1": new_mla_spec(), "layer_2": SlidingWindowMLASpec( block_size=16, num_kv_heads=1, head_size=576, dtype=torch.float32, sliding_window=128, ), }, ) assert get_kv_cache_spec_kind(uniform_mixed_spec) == KVCacheSpecKind.UNKNOWN def test_get_kv_cache_spec_sliding_window_reads_windowed_specs(): full_attention_spec = FullAttentionSpec( block_size=16, num_kv_heads=1, head_size=64, dtype=torch.float32, ) sliding_window_spec = SlidingWindowSpec( block_size=16, num_kv_heads=1, head_size=64, dtype=torch.float32, sliding_window=128, ) assert get_kv_cache_spec_sliding_window(full_attention_spec) is None assert get_kv_cache_spec_sliding_window(sliding_window_spec) == 128 def test_get_kv_cache_spec_sliding_window_unwraps_uniform_type_specs(): uniform_window_spec = UniformTypeKVCacheSpecs( block_size=16, kv_cache_specs={ "layer_1": SlidingWindowSpec( block_size=16, num_kv_heads=1, head_size=64, dtype=torch.float32, sliding_window=128, ), "layer_2": SlidingWindowSpec( block_size=16, num_kv_heads=2, head_size=64, dtype=torch.float32, sliding_window=128, ), }, ) mixed_window_spec = UniformTypeKVCacheSpecs( block_size=16, kv_cache_specs={ "layer_1": SlidingWindowSpec( block_size=16, num_kv_heads=1, head_size=64, dtype=torch.float32, sliding_window=128, ), "layer_2": SlidingWindowSpec( block_size=16, num_kv_heads=1, head_size=64, dtype=torch.float32, sliding_window=256, ), }, ) assert get_kv_cache_spec_sliding_window(uniform_window_spec) == 128 assert get_kv_cache_spec_sliding_window(mixed_window_spec) is None def test_merge_mla_spec(): kv_cache_specs = [ new_mla_spec(), new_mla_spec(), ] mla_spec = kv_cache_specs[0].merge(kv_cache_specs) assert mla_spec == new_mla_spec() kv_cache_specs = [ new_mla_spec(cache_dtype_str="fp8_ds_mla"), new_mla_spec(cache_dtype_str="fp8_ds_mla"), ] mla_spec = kv_cache_specs[0].merge(kv_cache_specs) assert mla_spec == new_mla_spec(cache_dtype_str="fp8_ds_mla") kv_cache_specs = [ new_mla_spec(cache_dtype_str="fp8_ds_mla"), new_mla_spec(cache_dtype_str=None), ] with pytest.raises(AssertionError): kv_cache_specs[0].merge(kv_cache_specs) kv_cache_specs = [ new_kv_cache_spec(), new_mla_spec(), ] with pytest.raises(AssertionError): kv_cache_specs[0].merge(kv_cache_specs) kv_cache_specs = [ new_mla_spec(cache_dtype_str="fp8_ds_mla"), new_kv_cache_spec(), ] with pytest.raises(AssertionError): kv_cache_specs[0].merge(kv_cache_specs) @pytest.mark.parametrize("hash_fn", [sha256, sha256_cbor]) def test_request_block_hasher_with_prompt_embeds(hash_fn: Callable[[Any], bytes]): block_size = 3 num_tokens = 2 * block_size prompt_token_ids = [_ for _ in range(num_tokens)] hidden_size = 5 prompt_embeds = torch.randn((num_tokens, hidden_size)) request = make_request( request_id="0", prompt_token_ids=prompt_token_ids, block_size=block_size, hash_fn=hash_fn, prompt_embeds=prompt_embeds, ) block_hashes = request.block_hashes assert len(block_hashes) == 2 block1_embeds_hash = hashlib.sha256( tensor_data(prompt_embeds[:block_size]) ).digest() expected_hash1 = hash_fn( ( kv_cache_utils.NONE_HASH, tuple(prompt_token_ids[:block_size]), (("prompt_embeds", block1_embeds_hash),), ) ) assert block_hashes[0] == expected_hash1 block2_embeds_hash = hashlib.sha256( tensor_data(prompt_embeds[block_size:num_tokens]) ).digest() expected_hash2 = hash_fn( ( block_hashes[0], tuple(prompt_token_ids[block_size:num_tokens]), (("prompt_embeds", block2_embeds_hash),), ) ) assert block_hashes[1] == expected_hash2 @pytest.mark.parametrize("hash_fn", [sha256, sha256_cbor]) def test_request_with_prompt_embeds_and_mm_inputs(hash_fn: Callable[[Any], bytes]): block_size = 3 num_tokens = 2 * block_size prompt_token_ids = [_ for _ in range(num_tokens)] hidden_size = 5 prompt_embeds = torch.randn((num_tokens, hidden_size)) request = make_request( request_id="0", prompt_token_ids=prompt_token_ids, block_size=block_size, hash_fn=hash_fn, mm_positions=[ PlaceholderRange(offset=0, length=3), PlaceholderRange(offset=3, length=3), ], mm_hashes=["hash1", "hash2"], prompt_embeds=prompt_embeds, ) block_hashes = request.block_hashes assert len(block_hashes) == 2 block1_embeds_hash = hashlib.sha256( tensor_data(prompt_embeds[:block_size]) ).digest() expected_hash1 = hash_fn( ( kv_cache_utils.NONE_HASH, tuple(prompt_token_ids[:block_size]), (("mm", "hash1", 0), ("prompt_embeds", block1_embeds_hash)), ) ) assert block_hashes[0] == expected_hash1 block2_embeds_hash = hashlib.sha256( tensor_data(prompt_embeds[block_size:num_tokens]) ).digest() expected_hash2 = hash_fn( ( block_hashes[0], tuple(prompt_token_ids[block_size:num_tokens]), (("mm", "hash2", 0), ("prompt_embeds", block2_embeds_hash)), ) ) assert block_hashes[1] == expected_hash2 def test_auto_fit_max_model_len(): """Test that max_model_len=-1 auto-fits to available GPU memory.""" # Create config with original_max_model_len=-1 to trigger auto-fit model_config = ModelConfig(max_model_len=1024) # Simulate the user passing -1 by setting original_max_model_len model_config.original_max_model_len = -1 vllm_config = VllmConfig(model_config=model_config) vllm_config.cache_config.kv_cache_layout = "LBNHC" mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2 # 16KB per block per layer kv_cache_specs = { "layer_1": new_kv_cache_spec(), "layer_2": new_kv_cache_spec(), } # With enough memory, max_model_len stays at the derived max large_available_memory = mem_per_block_per_layer * 2 * 1024 # plenty of memory _kv_cache_configs = get_kv_cache_configs( vllm_config, [kv_cache_specs], [large_available_memory] ) assert vllm_config.model_config.max_model_len == 1024 # Reset for next test model_config = ModelConfig(max_model_len=1024) model_config.original_max_model_len = -1 vllm_config = VllmConfig(model_config=model_config) vllm_config.cache_config.kv_cache_layout = "LBNHC" # With limited memory, max_model_len should be reduced # Need memory for at least max_model_len tokens # 32 blocks worth of memory for 2 layers = can fit 32*16=512 tokens limited_memory = mem_per_block_per_layer * 2 * 32 _kv_cache_configs = get_kv_cache_configs( vllm_config, [kv_cache_specs], [limited_memory] ) # Should be reduced to fit in memory assert vllm_config.model_config.max_model_len < 1024 assert vllm_config.model_config.max_model_len > 0 def test_auto_fit_max_model_len_with_hybrid(): """Test that auto-fit works with hybrid KV cache specs.""" # Create config with original_max_model_len=-1 to trigger auto-fit model_config = ModelConfig(max_model_len=8192) # Simulate the user passing -1 by setting original_max_model_len model_config.original_max_model_len = -1 vllm_config = VllmConfig(model_config=model_config) vllm_config.cache_config.kv_cache_layout = "LBNHC" mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2 # 16KB per block per layer gamma = 2 kv_cache_specs = { "layer_1": new_mamba_spec(num_speculative_blocks=gamma), "layer_2": new_kv_cache_spec(), } # One extra block on top of what a 1024-token request needs: the pool # reserves one block as the null block. available_memory = mem_per_block_per_layer * (1024 // 16 + 1 + gamma + 1) _kv_cache_configs = get_kv_cache_configs( vllm_config, [kv_cache_specs], [available_memory] ) assert vllm_config.model_config.max_model_len == 1024 def test_auto_fit_max_model_len_not_triggered(): """Test that auto-fit is not triggered when original_max_model_len is not -1.""" model_config = ModelConfig(max_model_len=16) # original_max_model_len should be None by default, not -1 vllm_config = VllmConfig(model_config=model_config) vllm_config.cache_config.kv_cache_layout = "LBNHC" mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2 kv_cache_specs = { "layer_1": new_kv_cache_spec(), "layer_2": new_kv_cache_spec(), } # This should work normally without auto-fit _kv_cache_configs = get_kv_cache_configs( vllm_config, [kv_cache_specs], [mem_per_block_per_layer * 2 * 32] ) assert vllm_config.model_config.max_model_len == 16 def test_auto_fit_max_model_len_respects_num_gpu_blocks_override(): """Auto-fit must size max_model_len against the override-clamped pool, not the raw `available_memory`. Without this, auto-fit could pick a max_model_len that no longer fits once `num_gpu_blocks_override` is applied. """ model_config = ModelConfig(max_model_len=16384) model_config.original_max_model_len = -1 # request auto-fit vllm_config = VllmConfig(model_config=model_config) vllm_config.cache_config.kv_cache_layout = "LBNHC" # Cap the cache to 32 blocks regardless of available memory. vllm_config.cache_config.num_gpu_blocks_override = 32 mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2 kv_cache_specs = { "layer_1": new_kv_cache_spec(), # block_size=16 "layer_2": new_kv_cache_spec(), } # Plenty of raw memory (1024 blocks per layer would fit max_model_len=16384). large_available_memory = mem_per_block_per_layer * 2 * 1024 get_kv_cache_configs(vllm_config, [kv_cache_specs], [large_available_memory]) # 32 blocks * block_size 16 = 512 token slots, so max_model_len must # auto-fit at or below that. assert 0 < vllm_config.model_config.max_model_len <= 32 * 16 def test_check_enough_kv_cache_memory_respects_num_gpu_blocks_override(): """Admission check must use the override-clamped pool size, not raw `available_memory`. Without this, startup could accept a max_model_len that does not actually fit in `num_gpu_blocks_override` blocks. """ model_config = ModelConfig(max_model_len=16384) vllm_config = VllmConfig(model_config=model_config) vllm_config.cache_config.kv_cache_layout = "LBNHC" # 32 blocks is far too small for max_model_len=16384 (would need 1024). vllm_config.cache_config.num_gpu_blocks_override = 32 mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2 kv_cache_specs = { "layer_1": new_kv_cache_spec(), "layer_2": new_kv_cache_spec(), } # Plenty of raw memory: a bytes-only check against this would pass. large_available_memory = mem_per_block_per_layer * 2 * 1024 with pytest.raises(ValueError, match="max seq len"): get_kv_cache_configs(vllm_config, [kv_cache_specs], [large_available_memory]) def test_unify_kv_cache_page_size_uses_padding_for_non_divisible_sizes(): """DFlash drafters can have a smaller head size than the target model. For example, MiMo uses 192-dim target KV heads while its DFlash draft uses 128-dim KV heads. The resulting page sizes are 3:2 rather than an integer block-size multiple, so the smaller page must be padded instead. """ # Attention layers read padded pages through a strided view, so padding is # allowed. target_spec = new_kv_cache_spec( block_size=16, num_kv_heads=1, head_size=192, dtype=torch.bfloat16, ) draft_spec = new_sliding_window_spec( block_size=16, num_kv_heads=1, head_size=128, dtype=torch.bfloat16, sliding_window=1024, ) unified_specs = kv_cache_utils.unify_kv_cache_spec_page_size( { "target_attn": target_spec, "draft_attn": draft_spec, } ) assert unified_specs["target_attn"] == target_spec unified_draft_spec = unified_specs["draft_attn"] assert unified_draft_spec.block_size == draft_spec.block_size assert unified_draft_spec.real_page_size_bytes == draft_spec.real_page_size_bytes assert unified_draft_spec.page_size_padded == target_spec.page_size_bytes assert unified_draft_spec.page_size_bytes == target_spec.page_size_bytes def test_unpadded_page_size_includes_per_token_head_scales(): # Per-token-head quant carries inline fp32 scales that are carved from the # raw KV allocation, so they must be budgeted into the offload width. The # packing is published by the owning backend's customize_spec hook. from vllm.v1.attention.backends.triton_attn import TritonAttentionBackend dense = new_kv_cache_spec(dtype=torch.uint8) spec = TritonAttentionBackend.customize_spec( new_kv_cache_spec( dtype=torch.uint8, kv_quant_mode=KVQuantMode.FP8_PER_TOKEN_HEAD ) ) scales = 2 * spec.block_size * spec.num_kv_heads * 4 assert spec.unpadded_page_size_bytes == dense.unpadded_page_size_bytes + scales assert spec.page_size_bytes == spec.unpadded_page_size_bytes def test_page_size_padded_wins(): # An explicit padded page size takes precedence over the unpadded size. from vllm.v1.attention.backends.triton_attn import TritonAttentionBackend spec = TritonAttentionBackend.customize_spec( new_kv_cache_spec( dtype=torch.uint8, kv_quant_mode=KVQuantMode.FP8_PER_TOKEN_HEAD, page_size_padded=65536, ) ) assert spec.page_size_bytes == 65536 def test_unify_hybrid_kv_cache_specs(): # 1. has_full_attention and has_sliding_window before_spec_1 = new_kv_cache_spec(block_size=64) before_spec_2 = new_sliding_window_spec( block_size=16, page_size_padded=32 * 1024, sliding_window=1024 ) kv_cache_spec = { "layer_1": before_spec_1, "layer_2": before_spec_2, } kv_cache_spec["draft_layer_1"] = replace(before_spec_1, dcp_sharded=False) kv_cache_spec["draft_layer_2"] = replace(before_spec_2, dcp_sharded=False) kv_cache_utils.unify_hybrid_kv_cache_specs(kv_cache_spec) expected_spec_1 = new_kv_cache_spec(block_size=64) expected_spec_2 = new_kv_cache_spec( block_size=64, page_size_padded=64 * 1024, sliding_window=1024 ) assert kv_cache_spec["layer_1"] == expected_spec_1 assert kv_cache_spec["layer_2"] == expected_spec_2 assert kv_cache_spec["layer_2"].page_size_bytes == 64 * 1024 assert kv_cache_spec["draft_layer_2"] == replace(expected_spec_2, dcp_sharded=False) # 2. has_full_attention and has_chunked_local_attention before_spec_1 = new_kv_cache_spec() before_spec_2 = new_chunked_local_attention_spec( page_size_padded=32 * 1024, attention_chunk_size=512 ) kv_cache_spec = { "layer_1": before_spec_1, "layer_2": before_spec_2, } kv_cache_utils.unify_hybrid_kv_cache_specs(kv_cache_spec) expected_spec_1 = new_kv_cache_spec() expected_spec_2 = new_kv_cache_spec( page_size_padded=32 * 1024, attention_chunk_size=512 ) assert kv_cache_spec["layer_1"] == expected_spec_1 assert kv_cache_spec["layer_2"] == expected_spec_2 # 3. has_full_attention, has_sliding_window and has_chunked_local_attention before_spec_1 = new_kv_cache_spec() before_spec_2 = new_sliding_window_spec( page_size_padded=32 * 1024, sliding_window=1024 ) before_spec_3 = new_chunked_local_attention_spec( page_size_padded=32 * 1024, attention_chunk_size=512 ) kv_cache_spec = { "layer_1": before_spec_1, "layer_2": before_spec_2, "layer_3": before_spec_3, } kv_cache_utils.unify_hybrid_kv_cache_specs(kv_cache_spec) expected_spec_1 = new_kv_cache_spec() expected_spec_2 = new_kv_cache_spec(page_size_padded=32 * 1024, sliding_window=1024) expected_spec_3 = new_kv_cache_spec( page_size_padded=32 * 1024, attention_chunk_size=512 ) assert kv_cache_spec["layer_1"] == expected_spec_1 assert kv_cache_spec["layer_2"] == expected_spec_2 assert kv_cache_spec["layer_3"] == expected_spec_3 # 4. No FullAttentionSpec, should not convert kv_cache_spec = { "layer_1": new_sliding_window_spec(sliding_window=1024), "layer_2": new_chunked_local_attention_spec(attention_chunk_size=512), } with pytest.raises(ValueError): kv_cache_utils.unify_hybrid_kv_cache_specs(kv_cache_spec) # Replicated Mamba state still requires the hybrid cache manager. kv_cache_spec = {"attention": new_kv_cache_spec(), "mamba": new_mamba_spec()} with pytest.raises(ValueError): kv_cache_utils.unify_hybrid_kv_cache_specs(kv_cache_spec) def test_unify_kv_cache_spec_page_size_mamba(): """Regression test for https://github.com/vllm-project/vllm/issues/43626. MambaSpec's page_size_bytes is determined by its state shapes and does not change with block_size, so unify_kv_cache_spec_page_size must pad the Mamba page instead of scaling its block_size. This situation arises when a layer with a page larger than the (already platform-aligned) Mamba page joins the specs, e.g. a dense draft model with more KV heads than the hybrid main model. """ # 1. Hybrid main model (Mamba + full attention, pages already aligned at # 16KB) plus a dense draft model layer with a 2x larger page (32KB). # Reproduces the bare AssertionError from #43626: 32768 % 16384 == 0, so # the old code scaled the Mamba block_size, which left page_size_bytes # unchanged at 16384. mamba_spec = new_mamba_spec() # page_size_bytes = 16384 main_attn_spec = new_kv_cache_spec() # page_size_bytes = 16384 draft_attn_spec = new_kv_cache_spec(num_kv_heads=4) # page_size_bytes = 32768 assert mamba_spec.page_size_bytes == main_attn_spec.page_size_bytes == 16384 assert draft_attn_spec.page_size_bytes == 32768 unified = kv_cache_utils.unify_kv_cache_spec_page_size( { "mamba_layer": mamba_spec, "main_attn_layer": main_attn_spec, "draft_attn_layer": draft_attn_spec, } ) # Mamba page is padded; block_size (caching granularity) is unchanged. assert unified["mamba_layer"].page_size_bytes == 32768 assert unified["mamba_layer"].page_size_padded == 32768 assert unified["mamba_layer"].block_size == mamba_spec.block_size # Attention layer with smaller page still unifies by scaling block_size. assert unified["main_attn_layer"].page_size_bytes == 32768 assert unified["main_attn_layer"].block_size == 2 * main_attn_spec.block_size assert unified["main_attn_layer"].page_size_padded is None # Layer already at max page size is unchanged. assert unified["draft_attn_layer"] == draft_attn_spec # 2. Mamba page already padded by the platform (state smaller than the # padded page); the padding is re-applied at the new maximum. padded_mamba_spec = new_mamba_spec( shapes=((2, 256), (3, 32, 32)), page_size_padded=16384 ) assert padded_mamba_spec.page_size_bytes == 16384 unified = kv_cache_utils.unify_kv_cache_spec_page_size( { "mamba_layer": padded_mamba_spec, "draft_attn_layer": draft_attn_spec, } ) assert unified["mamba_layer"].page_size_bytes == 32768 assert unified["mamba_layer"].page_size_padded == 32768 # 3. Mamba page that does not evenly divide the maximum page size is # padded as well (the divisibility constraint only applies to block_size # scaling). odd_mamba_spec = new_mamba_spec(shapes=((6144,),)) assert odd_mamba_spec.page_size_bytes == 24576 assert 32768 % odd_mamba_spec.page_size_bytes != 0 unified = kv_cache_utils.unify_kv_cache_spec_page_size( { "mamba_layer": odd_mamba_spec, "draft_attn_layer": draft_attn_spec, } ) assert unified["mamba_layer"].page_size_bytes == 32768 # 4. Attention layers with non-divisible page sizes are padded too: every # backend reads a padded page through the view's block stride, so there is # no longer a case that must raise. unified = kv_cache_utils.unify_kv_cache_spec_page_size( { "attn_layer": new_kv_cache_spec(block_size=24), # 24576 "draft_attn_layer": draft_attn_spec, # 32768 } ) assert unified["attn_layer"].page_size_padded == 32768 assert unified["attn_layer"].page_size_bytes == 32768 assert unified["attn_layer"].block_size == 24 # 5. Uniform page sizes are returned unchanged. specs = { "mamba_layer": new_mamba_spec(), "attn_layer": new_kv_cache_spec(), } assert kv_cache_utils.unify_kv_cache_spec_page_size(specs) == specs def test_hma_not_disabled_when_kv_events_enabled(): """Test enabling KV events must not force disable_hybrid_kv_cache_manager to True. This test guards against that regression by verifying that a VllmConfig with kv_events_config set still resolves disable_hybrid_kv_cache_manager to False (i.e. HMA remains enabled) when no other condition requires it to be disabled. """ model_config = ModelConfig(max_model_len=16) kv_events_config = KVEventsConfig( enable_kv_cache_events=True, publisher="null", ) # Leave disable_hybrid_kv_cache_manager as None (the default) so that # VllmConfig.__post_init__ resolves it automatically. vllm_config = VllmConfig( model_config=model_config, kv_events_config=kv_events_config, ) assert vllm_config.scheduler_config.disable_hybrid_kv_cache_manager is False, ( "kv_events_config must not force-disable the hybrid KV cache manager." ) def test_resolve_block_hashes_gate(): # Resolve symbols through the module so they stay consistent with each other # even after other tests reload ``kv_cache_utils``. resolve_block_hashes = kv_cache_utils.resolve_block_hashes BlockHashListWithBlockSize = kv_cache_utils.BlockHashListWithBlockSize # Raw, hash_block_size-granularity hashes (contents are opaque here). raw = [BlockHash(bytes([i])) for i in range(8)] # block_size == hash_block_size: always reuse the raw hashes. assert resolve_block_hashes(raw, 2, 2, alignment_tokens=2) is raw assert ( resolve_block_hashes( raw, 2, 2, supports_fine_grained_hash_lookup=True, alignment_tokens=2 ) is raw ) # Fine-grained manager, partial hits ON (alignment_tokens == hash_block_size # < block_size): keep raw hashes so the manager can scan at hash granularity. assert ( resolve_block_hashes( raw, 2, 4, supports_fine_grained_hash_lookup=True, alignment_tokens=2 ) is raw ) # Fine-grained manager, partial hits OFF (alignment_tokens == # scheduler_block_size >= block_size): must fall back to a block-size view, # exactly like the pre-refactor coordinator did. off = resolve_block_hashes( raw, 2, 4, supports_fine_grained_hash_lookup=True, alignment_tokens=4 ) assert isinstance(off, BlockHashListWithBlockSize) assert off.scale_factor == 2 # Non-fine-grained manager (e.g. sliding window): always a block-size view # when block_size != hash_block_size, regardless of alignment_tokens. swa = resolve_block_hashes( raw, 2, 4, supports_fine_grained_hash_lookup=False, alignment_tokens=2 ) assert isinstance(swa, BlockHashListWithBlockSize) assert swa.scale_factor == 2 def test_resolve_block_hashes_rejects_mismatched_view(): resolve_block_hashes = kv_cache_utils.resolve_block_hashes BlockHashListWithBlockSize = kv_cache_utils.BlockHashListWithBlockSize raw = [BlockHash(bytes([i])) for i in range(8)] # A view built at this block_size is returned as-is (idempotent). view = BlockHashListWithBlockSize(raw, 2, 4) assert resolve_block_hashes(view, 2, 4) is view # A view built at a different block_size must fail loudly rather than be # silently reinterpreted at the wrong granularity. mismatched = BlockHashListWithBlockSize(raw, 2, 8) with pytest.raises(AssertionError): resolve_block_hashes(mismatched, 2, 4) @pytest.mark.parametrize("use_override", [True, False]) def test_kv_cache_reserves_null_block_for_max_model_len(use_override): """A KV cache sized to exactly the blocks a max_model_len request needs must be rejected. BlockPool keeps one block as the null block, so only num_blocks - 1 are usable; accepting num_blocks == needed would leave the request unschedulable and hang the engine. Covers both the num_gpu_blocks_override path and the memory-derived block count. """ block_size = 16 max_model_len = 512 # needs 512 / 16 = 32 blocks vllm_config = VllmConfig(model_config=ModelConfig(max_model_len=max_model_len)) spec = new_kv_cache_spec(block_size=block_size) # 32 blocks -> only 31 usable after the null block: one short -> reject. if use_override: # Ample raw memory; the override alone constrains the block count. vllm_config.cache_config.num_gpu_blocks_override = 32 available_memory = [spec.page_size_bytes * 1024] else: available_memory = [spec.page_size_bytes * 32] with pytest.raises(ValueError, match="max seq len"): get_kv_cache_configs(vllm_config, [{"layer1": spec}], available_memory) def test_auto_fit_max_model_len_reserves_null_block(): """Auto-fit (max_model_len=-1) must size max_model_len against usable blocks, not the total pool. With memory for exactly 64 blocks, one is the null block, so auto-fit must settle on 63 * block_size; picking the full pool would leave a max-length request unschedulable and hang the engine. """ block_size = 16 model_config = ModelConfig(max_model_len=1024) model_config.original_max_model_len = -1 vllm_config = VllmConfig(model_config=model_config) vllm_config.cache_config.kv_cache_layout = "LBNHC" spec = new_kv_cache_spec(block_size=block_size) # Exactly the 1024 / 16 = 64 blocks a full-length request would need. available_memory = [spec.page_size_bytes * 64] get_kv_cache_configs(vllm_config, [{"layer1": spec}], available_memory) assert vllm_config.model_config.max_model_len == 63 * block_size def test_check_enough_kv_cache_memory_reserves_null_block(): """The public admission check must reject a KV cache sized to exactly the blocks a max_model_len request needs. BlockPool keeps one block as the null block, so only num_blocks - 1 are usable; accepting num_blocks == needed would leave the request unschedulable and hang the engine. """ block_size = 16 max_model_len = 512 # needs 512 / 16 = 32 blocks vllm_config = VllmConfig(model_config=ModelConfig(max_model_len=max_model_len)) spec = new_kv_cache_spec(block_size=block_size) # 32 blocks -> only 31 usable after the null block: one short -> reject. with pytest.raises(ValueError, match="max seq len"): check_enough_kv_cache_memory( vllm_config, {"layer1": spec}, spec.page_size_bytes * 32 ) # 33 blocks -> 32 usable after the null block -> accept. check_enough_kv_cache_memory( vllm_config, {"layer1": spec}, spec.page_size_bytes * 33 ) def test_is_full_attention_spec_unwraps_uniform_type_specs(): """``UniformTypeKVCacheSpecs`` is not a ``FullAttentionSpec``, so callers scanning groups with a bare isinstance miss DeepSeek-V4-shaped configs.""" common = dict(num_kv_heads=1, head_size=1, dtype=torch.float32) full = FullAttentionSpec(block_size=4, **common) mla = MLAAttentionSpec(block_size=4, **common) swa = SlidingWindowMLASpec(block_size=4, sliding_window=8, **common) def wrap(**specs): return UniformTypeKVCacheSpecs(block_size=4, kv_cache_specs=specs) assert is_full_attention_spec(full) assert is_full_attention_spec(mla) assert not is_full_attention_spec(swa) assert is_full_attention_spec(wrap(a=mla, b=mla)) assert not is_full_attention_spec(wrap(a=swa, b=swa)) # Every layer must qualify: a group holding a recycling sliding-window layer # has no stable slot layout. Grouping forbids mixing today, but the wrapper # is also built directly (e.g. projecting groups onto a PP worker). assert not is_full_attention_spec(wrap(a=mla, b=swa)) assert not is_full_attention_spec(wrap()) def test_iter_layer_specs_returns_group_members(): common = dict(num_kv_heads=1, head_size=1, dtype=torch.float32) full = FullAttentionSpec(block_size=4, **common) mla = MLAAttentionSpec(block_size=4, **common) assert list(iter_layer_specs(full)) == [full] wrapped = UniformTypeKVCacheSpecs( block_size=4, kv_cache_specs={"a": full, "b": mla} ) assert list(iter_layer_specs(wrapped)) == [full, mla] def test_wrapped_mamba_group_requires_block_zeroing(): mamba = MambaSpec( block_size=4, shapes=((4, 1),), dtypes=(torch.float32,), ) wrapped = UniformTypeKVCacheSpecs(block_size=4, kv_cache_specs={"mamba": mamba}) config = KVCacheConfig( num_blocks=4, kv_cache_tensors=[], kv_cache_groups=[KVCacheGroupSpec(["mamba"], wrapped)], ) assert config.needs_kv_cache_zeroing def _spec_decode_grouping_config(method="dspark", model_type=None): """Grouping config with an EAGLE-family speculative method enabled.""" return SimpleNamespace( scheduler_config=SimpleNamespace(disable_hybrid_kv_cache_manager=False), cache_config=SimpleNamespace( get_resolved_kv_cache_layout=lambda: SimpleNamespace( is_block_outermost=True ) ), model_config=SimpleNamespace(hf_config=SimpleNamespace(model_type=model_type)), speculative_config=SimpleNamespace( method=method, use_eagle=lambda: True, use_eagle_block_drop=lambda: True, ), ) def _hybrid_specs_with_draft(draft: bool, draft_shares_target_spec: bool = False): """A K3-shaped hybrid: MLA full attention + Mamba, optionally plus a DSpark-style draft MLA layer marked non_causal_multi_token_decode. The target's fp8 KV dtype is one difference from the draft on Kimi-K3 (target `--kv-cache-dtype fp8_e4m3`, draft `auto`). Pass draft_shares_target_spec to match that dtype while retaining the draft marker as a group boundary. """ target_dtype = None if draft_shares_target_spec else "fp8_e4m3" specs = { "target.attn.0": new_mla_spec(block_size=64, cache_dtype_str=target_dtype), "target.attn.1": new_mla_spec(block_size=64, cache_dtype_str=target_dtype), "target.mamba.0": new_mamba_spec(block_size=64, mamba_cache_mode="align"), "target.mamba.1": new_mamba_spec(block_size=64, mamba_cache_mode="align"), } if draft: draft_spec = new_mla_spec(block_size=64) specs["draft.attn.0"] = replace(draft_spec, non_causal_multi_token_decode=True) return specs def test_draft_group_annotated_on_hybrid_general_path(): # A drafter's MLA layer carries non_causal_multi_token_decode, so its group # is identifiable without keying off a model version. groups = get_kv_cache_groups( _spec_decode_grouping_config(), _hybrid_specs_with_draft(draft=True) ) flagged = [g for g in groups if g.is_eagle_group] assert len(flagged) == 1 assert "draft.attn.0" in flagged[0].layer_names def test_mamba_groups_never_flagged_even_when_draft_shares_a_group(): # Packed uniform-type groups can contain distinct target and draft layer # specs; the combined group still holds volatile draft KV and must be # flagged. What must never happen is a Mamba group being flagged: that # widens its lookup window to two consecutive chunks, which align-mode # checkpointing never produces, zeroing every lookup. groups = get_kv_cache_groups( _spec_decode_grouping_config(), _hybrid_specs_with_draft(draft=True, draft_shares_target_spec=True), ) for group in groups: if "draft.attn.0" in group.layer_names: assert group.is_eagle_group if isinstance(group.kv_cache_spec, MambaSpec): assert not group.is_eagle_group def test_draft_group_not_annotated_without_spec_decode(): # The marker alone must not flag anything; the eagle semantics only apply # when a speculative method is actually enabled. config = _spec_decode_grouping_config() config.speculative_config = None groups = get_kv_cache_groups(config, _hybrid_specs_with_draft(draft=True)) assert not any(g.is_eagle_group for g in groups) def test_unidentifiable_draft_with_mamba_warns(caplog_vllm): # No group carries the draft marker, so consumers fall back to # conservative behavior that silently breaks reuse for Mamba groups. # That must at least be visible. groups = get_kv_cache_groups( _spec_decode_grouping_config(), _hybrid_specs_with_draft(draft=False) ) assert not any(g.is_eagle_group for g in groups) assert "no KV cache group could be identified as the draft model's" in ( caplog_vllm.text ) def test_unidentifiable_draft_without_mamba_does_not_warn(caplog_vllm): # Pure-attention models degrade gracefully under the consumers' # conservative fallback (a one-block hit drop at most), so the warning # stays silent to avoid noise on every unannotated EAGLE deployment. specs = { "target.attn.0": new_mla_spec(block_size=64), "target.attn.1": new_mla_spec(block_size=64), } groups = get_kv_cache_groups(_spec_decode_grouping_config(), specs) assert not any(g.is_eagle_group for g in groups) assert "could be identified as the draft model's" not in caplog_vllm.text def test_no_warning_when_draft_group_is_identified(caplog_vllm): get_kv_cache_groups( _spec_decode_grouping_config(), _hybrid_specs_with_draft(draft=True) ) assert "could be identified as the draft model's" not in caplog_vllm.text def _deepseek_v4_specs(model_version="deepseek_v4"): """DeepseekV4-shaped specs: full MLA layers plus sliding-window MLA layers at differing page sizes, with the MTP draft layer registered last.""" return { "model.layers.0.self_attn.attn": new_mla_spec(), "model.layers.1.self_attn.attn": new_mla_spec(), "model.layers.2.self_attn.attn": new_swa_mla_spec( head_size=1024, model_version=model_version ), # The MTP block registers last, and its sliding-window size differs, so # it lands in a group of its own. "model.layers.3.self_attn.attn": new_swa_mla_spec( head_size=1024, sliding_window=256, model_version=model_version ), } @pytest.mark.parametrize( ("method", "model_type"), [("mtp", "deepseek_v4"), ("dspark", "deepseek_v4"), ("dspark", "deepseek_v41")], ) def test_deepseek_v4_draft_group_annotated_on_packed_path(method, model_type): # DeepseekV4's MTP block reuses the target's decoder layer, so its spec # carries no draft marker and only the positional rule can find it. This # pins the pre-existing behaviour that the unified annotator must preserve. groups = get_kv_cache_groups( _spec_decode_grouping_config(method=method, model_type=model_type), _deepseek_v4_specs(model_version=None), ) flagged = [g for g in groups if g.is_eagle_group] assert len(flagged) == 1 assert "model.layers.3.self_attn.attn" in flagged[0].layer_names def test_trailing_layer_fallback_applies_to_any_mtp_model(): # The positional rule is sound for every MTP drafter, not just DeepseekV4: # MTP blocks reuse the target's decoder layer (no spec marker) and always # register after every target layer. The model_type must not gate it. groups = get_kv_cache_groups( _spec_decode_grouping_config(method="mtp", model_type="other"), _deepseek_v4_specs(), ) flagged = [g for g in groups if g.is_eagle_group] assert len(flagged) == 1 assert "model.layers.3.self_attn.attn" in flagged[0].layer_names def _qwen3_5_hybrid_specs(with_mtp_layer: bool): """Qwen3.5-shaped hybrid: repeating [GDN x3, full-attn x1] blocks, with the MTP drafter's full-attn layer (spec-identical to the target's) registered last.""" specs = {} idx = 0 for _ in range(2): for _ in range(3): specs[f"model.layers.{idx}.linear_attn"] = new_mamba_spec( mamba_cache_mode="align" ) idx += 1 specs[f"model.layers.{idx}.self_attn.attn"] = new_kv_cache_spec() idx += 1 if with_mtp_layer: specs["mtp.layers.0.self_attn.attn"] = new_kv_cache_spec() return specs def test_qwen3_5_mtp_draft_group_annotated_on_hybrid_path(caplog_vllm): # A hybrid mamba + full-attention model with an MTP drafter that is # spec-indistinguishable from the target reaches the general multi-group # path. The trailing-layer rule must locate the draft group there so the # Mamba groups are not swept up by the flag-all consumer fallback. groups = get_kv_cache_groups( _spec_decode_grouping_config(method="mtp", model_type="qwen3_5"), _qwen3_5_hybrid_specs(with_mtp_layer=True), ) flagged = [g for g in groups if g.is_eagle_group] assert len(flagged) == 1 assert "mtp.layers.0.self_attn.attn" in flagged[0].layer_names for group in groups: if any( isinstance(spec, MambaSpec) for spec in iter_layer_specs(group.kv_cache_spec) ): assert not group.is_eagle_group assert "could be identified as the draft model's" not in caplog_vllm.text @pytest.mark.parametrize("method", ["eagle", "eagle3", "dspark"]) def test_non_mtp_eagle_hybrid_still_warns(caplog_vllm, method): # Other drafters are not covered by the trailing-layer rule, so an # unidentifiable hybrid draft must still warn. groups = get_kv_cache_groups( _spec_decode_grouping_config(method=method, model_type="qwen3_5"), _qwen3_5_hybrid_specs(with_mtp_layer=True), ) assert not any(g.is_eagle_group for g in groups) assert "could be identified as the draft model's" in caplog_vllm.text def test_trailing_layer_fallback_requires_exact_partition(): # If the groups do not partition the layers exactly (e.g. a caller that # dropped or duplicated layers), the positional rule is meaningless and # must not fire. from vllm.v1.core.kv_cache_utils import _annotate_eagle_groups specs = _qwen3_5_hybrid_specs(with_mtp_layer=True) config = _spec_decode_grouping_config(method="mtp", model_type="qwen3_5") groups = get_kv_cache_groups(config, specs) for g in groups: g.is_eagle_group = False # Remove one layer from its group: no longer an exact partition. trimmed = [ KVCacheGroupSpec( [n for n in g.layer_names if n != "model.layers.0.linear_attn"], g.kv_cache_spec, ) for g in groups ] _annotate_eagle_groups(config, specs, trimmed, use_trailing_layer_fallback=True) assert not any(g.is_eagle_group for g in trimmed) _GQA_SPEC = FullAttentionSpec( block_size=16, num_kv_heads=8, head_size=64, dtype=torch.float32 ) _GQA_SWA_SPEC = SlidingWindowSpec( block_size=16, num_kv_heads=8, head_size=64, dtype=torch.float32, sliding_window=128 ) @pytest.mark.parametrize( "specs,expected", [ pytest.param({"l.0": new_mla_spec(), "l.1": new_mla_spec()}, 4, id="mla"), pytest.param( { "l.0": new_mla_spec(), "l.1": new_mla_spec(), "s.0": new_swa_mla_spec(), "s.1": new_swa_mla_spec(), }, 4, id="mla-and-swa-mla", ), pytest.param({"l.0": _GQA_SPEC, "l.1": _GQA_SPEC}, 1, id="gqa"), pytest.param( {"l.0": new_mla_spec(), "l.1": _GQA_SPEC}, 1, id="mla-and-gqa-uniform-group", ), pytest.param( { "l.0": new_mla_spec(), "l.1": new_mla_spec(), "s.0": _GQA_SWA_SPEC, "s.1": _GQA_SWA_SPEC, }, 1, id="mla-and-gqa-swa", ), ], ) def test_kv_tp_replicas(monkeypatch, specs, expected): """Resolved per layer before scheduler flattening hides mixed groups.""" from vllm.config import ParallelConfig from vllm.platforms import current_platform monkeypatch.setattr(current_platform, "device_count", lambda: 4) parallel_config = ParallelConfig(tensor_parallel_size=4) vllm_config = VllmConfig( model_config=ModelConfig(max_model_len=16), parallel_config=parallel_config ) vllm_config.cache_config.kv_cache_layout = "LBNHC" mem = sum(s.page_size_bytes for s in specs.values()) * 10 configs = get_kv_cache_configs(vllm_config, [specs], [mem]) assert configs[0].kv_tp_replicas == expected scheduler = generate_scheduler_kv_cache_config(configs) assert scheduler.kv_tp_replicas == expected @pytest.mark.parametrize( "specs,tp_size,dcp_size,expected", [ pytest.param([new_mla_spec()], 8, 1, 8, id="mla"), pytest.param([new_mla_spec()], 8, 2, 1, id="mla-dcp"), pytest.param( [replace(new_mla_spec(), max_tp_shards=None)], 8, 1, 1, id="mla-unset" ), pytest.param([replace(_GQA_SPEC, max_tp_shards=2)], 8, 1, 4, id="gqa-partial"), pytest.param([replace(_GQA_SPEC, max_tp_shards=8)], 4, 1, 1, id="gqa-sharded"), pytest.param( [new_mla_spec(), replace(_GQA_SWA_SPEC, max_tp_shards=4)], 8, 1, 2, id="mla-and-gqa-partial", ), pytest.param( [ HiddenStateCacheSpec( block_size=16, num_kv_heads=1, head_size=64, dtype=torch.float32 ) ], 8, 1, 1, id="hidden-state", ), ], ) def test_kv_cache_groups_tp_replicas(specs, tp_size, dcp_size, expected): """Replicas are the gcd of each layer's tp_size // max_tp_shards.""" groups = [KVCacheGroupSpec([f"l.{i}"], spec) for i, spec in enumerate(specs)] assert kv_cache_groups_tp_replicas(groups, tp_size, dcp_size) == expected