test_kv_cache_utils.py · tduncan/Nvidia-Optimization-Stack-VLLM

@tduncan / Nvidia-Optimization-Stack-VLLM

Code Commits Issues Pulls Cloud Agent CI/CD

main

py 4,850 lines 177 KB

Raw

1	# SPDX-License-Identifier: Apache-2.0
2	# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3	import copy
4	import hashlib
5	import importlib
6	import mmap
7	import subprocess
8	import sys
9	from collections.abc import Callable
10	from dataclasses import replace
11	from pathlib import Path
12	from types import SimpleNamespace
13	from typing import Any, cast
14
15	import pytest
16	import torch
17
18	import vllm.v1.core.kv_cache_utils as kv_cache_utils
19	import vllm.v1.hisparse.runtime as hisparse_runtime_module
20	from vllm.config import (
21	CacheConfig,
22	KVTransferConfig,
23	ModelConfig,
24	SchedulerConfig,
25	VllmConfig,
26	)
27	from vllm.config.attention import HiSparseConfig
28	from vllm.config.kv_events import KVEventsConfig
29	from vllm.lora.request import LoRARequest
30	from vllm.multimodal.inputs import (
31	MultiModalFeatureSpec,
32	MultiModalKwargsItem,
33	PlaceholderRange,
34	)
35	from vllm.sampling_params import SamplingParams
36	from vllm.utils.hashing import sha256, sha256_cbor, xxhash, xxhash_cbor
37	from vllm.utils.mem_constants import GiB_bytes
38	from vllm.v1.core.kv_cache_manager import KVCacheBlocks, KVCacheManager
39	from vllm.v1.core.kv_cache_utils import (
40	BlockHash,
41	FreeKVCacheBlockQueue,
42	KVCacheBlock,
43	check_enough_kv_cache_memory,
44	estimate_max_model_len,
45	generate_block_hash_extra_keys,
46	generate_scheduler_kv_cache_config,
47	get_kv_cache_capacity,
48	get_kv_cache_configs,
49	get_kv_cache_groups,
50	get_max_concurrency_for_kv_cache_config,
51	get_request_block_hasher,
52	hash_block_tokens,
53	init_none_hash,
54	is_kv_cache_spec_uniform,
55	kv_cache_groups_tp_replicas,
56	make_block_hash_with_group_id,
57	tensor_data,
58	to_event_extra_keys,
59	)
60	from vllm.v1.hisparse.layout import (
61	create_hisparse_layout,
62	get_hisparse_kv_cache_groups,
63	)
64	from vllm.v1.kv_cache_interface import (
65	ChunkedLocalAttentionSpec,
66	CircularBufferSpec,
67	FullAttentionSpec,
68	HiddenStateCacheSpec,
69	HiSparseHotSpec,
70	HiSparseResidentSpec,
71	KpoolTailSpec,
72	KVCacheConfig,
73	KVCacheGroupSpec,
74	KVCacheSpec,
75	KVCacheSpecKind,
76	KVCacheTensor,
77	KVQuantMode,
78	MambaSpec,
79	MLAAttentionSpec,
80	SinkFullAttentionSpec,
81	SlidingWindowMLASpec,
82	SlidingWindowSpec,
83	SparseCacheRole,
84	UniformTypeKVCacheSpecs,
85	get_kv_cache_spec_kind,
86	get_kv_cache_spec_sliding_window,
87	is_full_attention_spec,
88	iter_layer_specs,
89	)
90	from vllm.v1.kv_cache_layout import KVCacheLayout
91	from vllm.v1.metrics.stats import CachingMetrics, PrefixCacheStats
92	from vllm.v1.request import Request
93
94	pytestmark = pytest.mark.cpu_test
95
96
97	@pytest.mark.parametrize("gpu_block_size", [32, 64])
98	@pytest.mark.parametrize("shared_host_pool", [False, True])
99	def test_hisparse_hma_uses_resolved_gpu_block_size(
100	monkeypatch, gpu_block_size, shared_host_pool
101	):
102	monkeypatch.setattr(
103	hisparse_runtime_module.current_platform, "is_cuda_alike", lambda: True
104	)
105	specs = {
106	"model.layers.0.self_attn": MLAAttentionSpec(
107	block_size=gpu_block_size,
108	num_kv_heads=1,
109	head_size=576,
110	dtype=torch.bfloat16,
111	is_index_group_leader=True,
112	),
113	"model.layers.0.self_attn.indexer": MLAAttentionSpec(
114	block_size=gpu_block_size,
115	num_kv_heads=1,
116	head_size=128,
117	dtype=torch.bfloat16,
118	cache_role=SparseCacheRole.INDEXER,
119	),
120	}
121	config = SimpleNamespace(
122	attention_config=SimpleNamespace(hisparse_config=HiSparseConfig()),
123	model_config=SimpleNamespace(
124	hf_config=SimpleNamespace(index_topk=128),
125	max_model_len=gpu_block_size,
126	),
127	parallel_config=SimpleNamespace(
128	tensor_parallel_size=2 if shared_host_pool else 1,
129	pipeline_parallel_size=1,
130	prefill_context_parallel_size=1,
131	decode_context_parallel_size=1,
132	world_size=2 if shared_host_pool else 1,
133	distributed_executor_backend="mp",
134	nnodes_within_dp=1,
135	),
136	cache_config=SimpleNamespace(
137	num_gpu_blocks_override=7,
138	prefix_cache_retention_interval=None,
139	get_resolved_kv_cache_layout=lambda: KVCacheLayout.BLHNC,
140	),
141	)
142	monkeypatch.setattr(kv_cache_utils, "get_hisparse_host_pool_bytes", lambda _: 2**30)
143	cache_config = kv_cache_utils.get_kv_cache_config_from_groups(
144	config, get_hisparse_kv_cache_groups(config, specs), available_memory=2**30
145	)
146	assert cache_config.num_blocks == 7
147	assert cache_config.hisparse_host_num_blocks is not None
148	assert cache_config.hisparse_host_num_blocks > 7
149
150	host_group, indexer_group, *auxiliary_groups = cache_config.kv_cache_groups
151	assert cache_config.hisparse_shared_host_pool is shared_host_pool
152	host_page = host_group.kv_cache_spec.page_size_bytes
153	alignment = mmap.PAGESIZE if shared_host_pool else 1
154	expected_host_stride = (host_page + alignment - 1) // alignment * alignment
155	assert cache_config.hisparse_host_block_stride == expected_host_stride
156	assert cache_config.hisparse_host_num_blocks == 2**30 // expected_host_stride
157	assert host_group.host_resident
158	assert not any(group.host_resident for group in [indexer_group, *auxiliary_groups])
159	host_layers = set(host_group.layer_names)
160	for tensor in cache_config.kv_cache_tensors:
161	assert all(
162	(name in host_layers) == tensor.host_resident for name in tensor.layers
163	)
164	host_tensor = next(t for t in cache_config.kv_cache_tensors if t.host_resident)
165	host_spec = host_group.kv_cache_spec.kv_cache_specs[host_tensor.layers[0]]
166	assert host_tensor.block_stride == host_spec.page_size_bytes
167	assert host_tensor.layer_stride == (
168	host_spec.page_size_bytes * cache_config.hisparse_host_num_blocks
169	)
170	assert host_group.kv_cache_spec.block_size == gpu_block_size
171	assert indexer_group.kv_cache_spec.block_size == gpu_block_size
172	host_specs = host_group.kv_cache_spec.kv_cache_specs
173	gpu_indexer_specs = indexer_group.kv_cache_spec.kv_cache_specs
174	assert set(host_specs) == {"model.layers.0.self_attn"}
175	assert set(gpu_indexer_specs) == {"model.layers.0.self_attn.indexer"}
176	assert indexer_group.kv_cache_spec.prefix_cacheable
177	assert host_group.enable_kv_transfer
178	auxiliary_specs = [group.kv_cache_spec for group in auxiliary_groups]
179	assert any(isinstance(spec, HiSparseResidentSpec) for spec in auxiliary_specs)
180	assert any(isinstance(spec, HiSparseHotSpec) for spec in auxiliary_specs)
181	assert all(
182	spec.block_size == gpu_block_size
183	for spec in auxiliary_specs
184	if isinstance(spec, (HiSparseResidentSpec, HiSparseHotSpec))
185	)
186	scheduler_block_size, hash_block_size = kv_cache_utils.resolve_kv_cache_block_sizes(
187	cache_config,
188	SimpleNamespace(
189	cache_config=SimpleNamespace(
190	block_size=16,
191	enable_prefix_caching=True,
192	prefix_match_unit=None,
193	),
194	parallel_config=SimpleNamespace(decode_context_parallel_size=1),
195	kv_transfer_config=object(),
196	),
197	)
198	assert scheduler_block_size == hash_block_size == gpu_block_size
199
200
201	@pytest.mark.parametrize("extra_blocks,ok", [(1, False), (2, True)])
202	def test_hisparse_host_pool_must_fit_max_model_len(monkeypatch, extra_blocks, ok):
203	"""Every page needs a host block, so one max_model_len request must fit."""
204	monkeypatch.setattr(
205	hisparse_runtime_module.current_platform, "is_cuda_alike", lambda: True
206	)
207	specs = {
208	"model.layers.0.self_attn": MLAAttentionSpec(
209	block_size=64,
210	num_kv_heads=1,
211	head_size=576,
212	dtype=torch.bfloat16,
213	is_index_group_leader=True,
214	),
215	"model.layers.0.self_attn.indexer": MLAAttentionSpec(
216	block_size=64,
217	num_kv_heads=1,
218	head_size=128,
219	dtype=torch.bfloat16,
220	cache_role=SparseCacheRole.INDEXER,
221	),
222	}
223	config = SimpleNamespace(
224	attention_config=SimpleNamespace(hisparse_config=HiSparseConfig()),
225	model_config=SimpleNamespace(
226	hf_config=SimpleNamespace(index_topk=128), max_model_len=64 * 4
227	),
228	parallel_config=SimpleNamespace(tensor_parallel_size=1, world_size=1),
229	cache_config=SimpleNamespace(num_gpu_blocks_override=7),
230	)
231	host_page = specs["model.layers.0.self_attn"].page_size_bytes
232	host_budget = (4 + extra_blocks) * host_page
233	groups = get_hisparse_kv_cache_groups(config, specs)
234
235	if ok:
236	layout = create_hisparse_layout(config, groups, host_budget=host_budget)
237	assert layout.host_num_blocks == 4 + extra_blocks
238	else:
239	with pytest.raises(ValueError, match="increase host_pool_gib"):
240	create_hisparse_layout(config, groups, host_budget=host_budget)
241
242
243	@pytest.mark.parametrize(
244	"max_model_len,num_gpu_blocks,ok",
245	[
246	(8192, 160, False),
247	(8192, 1000, True),
248	(32768, 656, False),
249	(32768, 657, True),
250	(-1, 657, True),
251	],
252	)
253	def test_hisparse_pool_must_fit_max_model_len(
254	monkeypatch, max_model_len, num_gpu_blocks, ok
255	):
256	"""Each resident group takes its own blocks from HiSparse's shared GPU pool,
257	so a pool that fits only the indexer pages must be rejected at startup.
258	Resident pages are capped at the in-flight window: at 32768 tokens with 2048
259	in flight, one request needs 512 indexer + 4 x 32 resident + 16 hot blocks,
260	plus the null block. Auto-fit (-1) must pick a length that passes the same
261	check."""
262	monkeypatch.setattr(
263	hisparse_runtime_module.current_platform, "is_cuda_alike", lambda: True
264	)
265	monkeypatch.setattr(kv_cache_utils, "get_hisparse_host_pool_bytes", lambda _: 2**30)
266	specs: dict[str, KVCacheSpec] = {}
267	for i in range(4):
268	specs[f"model.layers.{i}.self_attn"] = MLAAttentionSpec(
269	block_size=64,
270	num_kv_heads=1,
271	head_size=576,
272	dtype=torch.bfloat16,
273	is_index_group_leader=True,
274	)
275	specs[f"model.layers.{i}.self_attn.indexer"] = MLAAttentionSpec(
276	block_size=64,
277	num_kv_heads=1,
278	head_size=128,
279	dtype=torch.bfloat16,
280	cache_role=SparseCacheRole.INDEXER,
281	)
282	config = VllmConfig(model_config=ModelConfig(max_model_len=max_model_len))
283	config.scheduler_config.max_num_batched_tokens = 2048
284	# Async scheduling doubles the in-flight window, and with it the bound.
285	config.scheduler_config.async_scheduling = False
286	config.attention_config.hisparse_config = HiSparseConfig()
287	config.model_config.hf_config.index_topk = 128
288	config.cache_config.num_gpu_blocks_override = num_gpu_blocks
289	config.cache_config.kv_cache_layout = "BLHNC"
290	# Workers report the resident/hot caches HiSparse derives.
291	specs \|= {
292	name: group.kv_cache_spec
293	for group in get_hisparse_kv_cache_groups(config, specs) or []
294	if isinstance(group.kv_cache_spec, (HiSparseResidentSpec, HiSparseHotSpec))
295	for name in group.layer_names
296	}
297
298	if ok:
299	kv_cache_utils.get_kv_cache_configs(config, [specs], [2**34])
300	else:
301	with pytest.raises(ValueError, match="max seq len"):
302	kv_cache_utils.get_kv_cache_configs(config, [specs], [2**34])
303
304
305	def test_hisparse_derived_specs_do_not_affect_kv_cache_layout(monkeypatch):
306	"""Workers report HiSparse's resident/hot caches, which lay out their own raw
307	backing; KV cache layout resolution must skip them."""
308	from vllm.v1.attention.backends.utils import resolve_kv_cache_layout
309
310	monkeypatch.setattr(
311	hisparse_runtime_module.current_platform, "is_cuda_alike", lambda: True
312	)
313	monkeypatch.delenv("VLLM_KV_CACHE_LAYOUT", raising=False)
314	specs: dict[str, KVCacheSpec] = {
315	"model.layers.0.self_attn": MLAAttentionSpec(
316	block_size=64,
317	num_kv_heads=1,
318	head_size=576,
319	dtype=torch.bfloat16,
320	is_index_group_leader=True,
321	),
322	"model.layers.0.self_attn.indexer": MLAAttentionSpec(
323	block_size=64,
324	num_kv_heads=1,
325	head_size=128,
326	dtype=torch.bfloat16,
327	cache_role=SparseCacheRole.INDEXER,
328	),
329	}
330	config = VllmConfig(model_config=ModelConfig(max_model_len=1024))
331	config.attention_config.hisparse_config = HiSparseConfig()
332	config.model_config.hf_config.index_topk = 128
333	reported = specs \| {
334	name: group.kv_cache_spec
335	for group in get_hisparse_kv_cache_groups(config, specs) or []
336	if isinstance(group.kv_cache_spec, (HiSparseResidentSpec, HiSparseHotSpec))
337	for name in group.layer_names
338	}
339	assert len(reported) > len(specs)
340
341	layouts = []
342	for worker_specs in (specs, reported):
343	config.cache_config.kv_cache_layout = None
344	layouts.append(
345	resolve_kv_cache_layout(
346	config, [["BLHNC", "LBNHC"]], list(worker_specs.values())
347	)
348	)
349	assert layouts[0] == layouts[1]
350
351
352	def test_hisparse_rejects_deepseek_v4():
353	full_specs = {
354	"model.layers.0.attn": MLAAttentionSpec(
355	block_size=256,
356	num_kv_heads=1,
357	head_size=512,
358	dtype=torch.uint8,
359	model_version="deepseek_v4",
360	)
361	}
362	config = SimpleNamespace(
363	attention_config=SimpleNamespace(hisparse_config=HiSparseConfig()),
364	model_config=SimpleNamespace(hf_config=SimpleNamespace(index_topk=512)),
365	cache_config=SimpleNamespace(num_gpu_blocks_override=7),
366	)
367
368	with pytest.raises(ValueError, match="does not support DeepSeek V4"):
369	get_hisparse_kv_cache_groups(config, full_specs)
370
371
372	@pytest.fixture(autouse=True)
373	def _auto_init_hash_fn(request):
374	hash_fn: Callable
375	if "hash_fn" in request.fixturenames:
376	hash_fn = request.getfixturevalue("hash_fn")
377	else:
378	hash_fn = sha256
379	init_none_hash(hash_fn)
380
381
382	def make_request(
383	request_id: str,
384	prompt_token_ids: list[int] \| None,
385	block_size: int = 3,
386	hash_fn: Callable = hash,
387	mm_positions: list[PlaceholderRange] \| None = None,
388	mm_hashes: list[str] \| None = None,
389	cache_salt: str \| None = None,
390	prompt_embeds: torch.Tensor \| None = None,
391	):
392	mm_features = []
393	if mm_positions is not None:
394	for j, position in enumerate(mm_positions):
395	identifier = mm_hashes[j] if mm_hashes else f"hash_{j}"
396	mm_feature = MultiModalFeatureSpec(
397	data=MultiModalKwargsItem.dummy(),
398	mm_position=position,
399	identifier=identifier,
400	modality="image",
401	)
402	mm_features.append(mm_feature)
403
404	sampling_params = SamplingParams(max_tokens=17)
405	sampling_params.update_from_generation_config({}, eos_token_id=100)
406
407	return Request(
408	request_id=request_id,
409	prompt_token_ids=prompt_token_ids,
410	mm_features=mm_features if mm_features else None,
411	sampling_params=sampling_params,
412	pooling_params=None,
413	lora_request=None,
414	cache_salt=cache_salt,
415	block_hasher=get_request_block_hasher(block_size, hash_fn),
416	prompt_embeds=prompt_embeds,
417	)
418
419
420	@pytest.mark.parametrize("dcp", [1, 4])
421	def test_effective_attention_block_size_matches_events(dcp):
422	from vllm.distributed.kv_events import BlockStored
423	from vllm.v1.engine.core import EngineCore
424
425	config = KVCacheConfig(
426	num_blocks=32,
427	kv_cache_tensors=[],
428	kv_cache_groups=[
429	KVCacheGroupSpec(["attention"], new_kv_cache_spec()),
430	],
431	)
432	manager = KVCacheManager(
433	generate_scheduler_kv_cache_config([config]),
434	max_model_len=256,
435	scheduler_block_size=16 * dcp,
436	hash_block_size=16 * dcp,
437	dcp_world_size=dcp,
438	enable_kv_cache_events=True,
439	)
440	core = EngineCore.__new__(EngineCore)
441	core.vllm_config = SimpleNamespace(cache_config=CacheConfig(block_size=16))
442	core.scheduler = SimpleNamespace(kv_cache_manager=manager)
443	core._initialize_effective_attention_block_size()
444	block_size = core.vllm_config.cache_config.effective_attention_block_size
445	assert block_size == 16 * dcp
446
447	request = make_request(
448	"block-size", list(range(64)), block_size=16 * dcp, hash_fn=sha256
449	)
450	assert manager.allocate_slots(request, 64) is not None
451	assert [
452	event.block_size
453	for event in manager.take_events()
454	if isinstance(event, BlockStored)
455	] == [block_size]
456
457
458	def new_kv_cache_spec(
459	block_size=16,
460	num_kv_heads=2,
461	head_size=64,
462	dtype=torch.float32,
463	page_size_padded=None,
464	sliding_window=None,
465	attention_chunk_size=None,
466	kv_quant_mode=KVQuantMode.NONE,
467	):
468	return FullAttentionSpec(
469	block_size=block_size,
470	num_kv_heads=num_kv_heads,
471	head_size=head_size,
472	dtype=dtype,
473	page_size_padded=page_size_padded,
474	sliding_window=sliding_window,
475	attention_chunk_size=attention_chunk_size,
476	kv_quant_mode=kv_quant_mode,
477	)
478
479
480	def test_kv_cache_config_selects_only_transferable_groups():
481	"""Connectors must see a stable projection of transfer-eligible groups."""
482	groups = [
483	KVCacheGroupSpec(["layer.0"], new_kv_cache_spec()),
484	KVCacheGroupSpec(["layer.1"], new_kv_cache_spec(), enable_kv_transfer=False),
485	KVCacheGroupSpec(["layer.2"], new_kv_cache_spec()),
486	]
487	config = KVCacheConfig(
488	num_blocks=1,
489	kv_cache_tensors=[],
490	kv_cache_groups=groups,
491	)
492
493	assert config.transfer_group_ids == (0, 2)
494	assert config.transfer_groups == (groups[0], groups[2])
495	assert config.transfer_group_index_by_layer == {"layer.0": 0, "layer.2": 1}
496	first_blocks = [1, 2]
497	third_blocks = [4]
498	assert config.select_transfer_block_ids((first_blocks, [3], third_blocks)) == (
499	first_blocks,
500	third_blocks,
501	)
502
503
504	def test_kv_cache_config_selects_prefix_cacheable_groups():
505	"""Prefix stores exclude scratch state without changing transfer groups."""
506	full_group = KVCacheGroupSpec(["full"], new_kv_cache_spec())
507	qsa_group = KVCacheGroupSpec(
508	["qsa"],
509	CircularBufferSpec(
510	block_size=4,
511	num_kv_heads=1,
512	head_size=64,
513	head_size_v=0,
514	dtype=torch.float16,
515	),
516	)
517	disabled_group = KVCacheGroupSpec(
518	["disabled"], new_kv_cache_spec(), enable_kv_transfer=False
519	)
520	config = KVCacheConfig(
521	num_blocks=1,
522	kv_cache_tensors=[],
523	kv_cache_groups=[full_group, qsa_group, disabled_group],
524	)
525	assert config.transfer_group_ids == (0, 1)
526	assert config.select_transfer_block_ids(([1], [2], [3])) == ([1], [2])
527	assert config.prefix_cacheable_group_ids == (0,)
528	assert config.prefix_cacheable_groups == (full_group,)
529
530
531	def test_kv_cache_blocks_selects_requested_groups():
532	blocks = KVCacheBlocks(
533	(
534	[KVCacheBlock(1)],
535	[KVCacheBlock(2)],
536	[KVCacheBlock(3)],
537	)
538	)
539
540	assert blocks.get_block_ids(group_ids=(0, 2)) == ([1], [3])
541
542
543	def new_sliding_window_spec(
544	block_size=16,
545	num_kv_heads=2,
546	head_size=64,
547	dtype=torch.float32,
548	page_size_padded=None,
549	sliding_window=1,
550	):
551	return SlidingWindowSpec(
552	block_size=block_size,
553	num_kv_heads=num_kv_heads,
554	head_size=head_size,
555	dtype=dtype,
556	page_size_padded=page_size_padded,
557	sliding_window=sliding_window,
558	)
559
560
561	def new_chunked_local_attention_spec(
562	block_size=16,
563	num_kv_heads=2,
564	head_size=64,
565	dtype=torch.float32,
566	page_size_padded=None,
567	attention_chunk_size=4,
568	):
569	return ChunkedLocalAttentionSpec(
570	block_size=block_size,
571	num_kv_heads=num_kv_heads,
572	head_size=head_size,
573	dtype=dtype,
574	page_size_padded=page_size_padded,
575	attention_chunk_size=attention_chunk_size,
576	)
577
578
579	def new_mamba_spec(
580	block_size=16,
581	shapes=((2, 512), (3, 32, 32)),
582	dtypes=(torch.float32, torch.float32),
583	num_speculative_blocks=2,
584	mamba_cache_mode="none",
585	page_size_padded=None,
586	):
587	return MambaSpec(
588	block_size=block_size,
589	shapes=shapes,
590	dtypes=dtypes,
591	page_size_padded=page_size_padded,
592	mamba_cache_mode=mamba_cache_mode,
593	num_speculative_blocks=num_speculative_blocks,
594	)
595
596
597	@pytest.mark.parametrize("hash_fn", [sha256, sha256_cbor])
598	def test_none_hash(monkeypatch, hash_fn):
599	import vllm.v1.core.kv_cache_utils
600
601	# case 1: PYTHONHASHSEED is not set -> deterministic default seed so that
602	# independent processes compute identical block hashes for identical
603	# content (e.g. for KV cache reuse across nodes).
604	with monkeypatch.context() as m:
605	m.delenv("PYTHONHASHSEED", raising=False)
606	reloaded_kv_cache_utils = importlib.reload(vllm.v1.core.kv_cache_utils)
607	reloaded_kv_cache_utils.init_none_hash(hash_fn)
608	none_hash = reloaded_kv_cache_utils.NONE_HASH
609	assert isinstance(none_hash, bytes)
610	assert none_hash != b""
611	assert none_hash == hash_fn(reloaded_kv_cache_utils.DEFAULT_NONE_HASH_SEED)
612	# deterministic across re-initialization within the same environment
613	reloaded_kv_cache_utils.init_none_hash(hash_fn)
614	assert none_hash == reloaded_kv_cache_utils.NONE_HASH
615
616	# case 2: PYTHONHASHSEED is set, use the seed and hash_fn
617	with monkeypatch.context() as m:
618	m.setenv("PYTHONHASHSEED", "python hash seed")
619	reloaded_kv_cache_utils = importlib.reload(vllm.v1.core.kv_cache_utils)
620	reloaded_kv_cache_utils.init_none_hash(hash_fn)
621	assert reloaded_kv_cache_utils.NONE_HASH is not None
622	assert isinstance(reloaded_kv_cache_utils.NONE_HASH, bytes)
623	assert hash_fn("python hash seed") == reloaded_kv_cache_utils.NONE_HASH
624
625
626	@pytest.mark.parametrize("non_crypto_fn", [xxhash, xxhash_cbor])
627	def test_none_hash_seed_random_for_non_crypto(monkeypatch, non_crypto_fn):
628	"""Non-cryptographic algorithms keep the per-process random seed.
629
630	A deterministic seed is safe for SHA-256, whose collision resistance does
631	not depend on a secret, but xxHash is not collision resistant: a known seed
632	would let an attacker precompute colliding blocks offline. Keep the
633	unpredictable seed there unless the operator opts into a shared one.
634	"""
635	# PYTHONHASHSEED unset -> unpredictable, differs per resolution.
636	with monkeypatch.context() as m:
637	m.delenv("PYTHONHASHSEED", raising=False)
638	seeds = {kv_cache_utils.resolve_none_hash_seed(non_crypto_fn) for _ in range(5)}
639	assert len(seeds) == 5
640	assert kv_cache_utils.DEFAULT_NONE_HASH_SEED not in seeds
641
642	# PYTHONHASHSEED set -> operator opt-in wins, so peers can share a cache.
643	with monkeypatch.context() as m:
644	m.setenv("PYTHONHASHSEED", "12345")
645	assert kv_cache_utils.resolve_none_hash_seed(non_crypto_fn) == "12345"
646
647
648	@pytest.mark.parametrize("hash_fn", [sha256, sha256_cbor])
649	def test_none_hash_seed_deterministic_for_crypto(monkeypatch, hash_fn):
650	with monkeypatch.context() as m:
651	m.delenv("PYTHONHASHSEED", raising=False)
652	seed = kv_cache_utils.resolve_none_hash_seed(hash_fn)
653	assert seed == kv_cache_utils.DEFAULT_NONE_HASH_SEED
654	assert seed == kv_cache_utils.resolve_none_hash_seed(hash_fn)
655
656
657	def test_get_none_hash_seed_reports_effective_seed(monkeypatch):
658	"""P2P advertises the seed NONE_HASH was actually derived from.
659
660	The P2P tier is constructed before init_none_hash runs, so it must read the
661	resolved seed lazily rather than re-deriving it.
662	"""
663	import vllm.v1.core.kv_cache_utils
664
665	with monkeypatch.context() as m:
666	m.delenv("PYTHONHASHSEED", raising=False)
667	reloaded = importlib.reload(vllm.v1.core.kv_cache_utils)
668	reloaded.init_none_hash(sha256)
669	assert reloaded.get_none_hash_seed() == reloaded.DEFAULT_NONE_HASH_SEED
670
671	with monkeypatch.context() as m:
672	m.setenv("PYTHONHASHSEED", "12345")
673	reloaded = importlib.reload(vllm.v1.core.kv_cache_utils)
674	reloaded.init_none_hash(sha256)
675	assert reloaded.get_none_hash_seed() == "12345"
676
677
678	def test_kv_cache_block():
679	# Test KVCacheBlock initialization
680	block = KVCacheBlock(block_id=0)
681	assert block.block_id == 0
682	assert block.ref_cnt == 0
683	assert block.block_hash is None
684
685	# Test reference count manipulation
686	block.ref_cnt += 1
687	assert block.ref_cnt == 1
688	block.ref_cnt -= 1
689	assert block.ref_cnt == 0
690
691	# Test block hash setting and resetting
692	block_hash = make_block_hash_with_group_id(BlockHash(b"abc"), 0)
693	block.set_block_hash(block_hash)
694	assert block.block_hash == block_hash
695
696	block.reset_hash()
697	assert block.block_hash is None
698
699
700	def test_kv_cache_block_uses_slots():
701	block = KVCacheBlock(block_id=0)
702
703	# Slots eliminate per-instance __dict__, saving ~264 bytes per block.
704	# At 100K+ blocks this avoids tens of MB of overhead and GC pressure.
705	assert not hasattr(block, "__dict__")
706
707	# Verify that slots actually prevent dynamic attribute assignment.
708	with pytest.raises(AttributeError):
709	block.unexpected_field = True
710
711
712	def test_free_kv_cache_block_queue_initialization():
713	# Test with a single block
714	block = KVCacheBlock(block_id=0)
715	queue = FreeKVCacheBlockQueue([block])
716	assert queue.num_free_blocks == 1
717	assert queue.fake_free_list_head.next_free_block is block
718	assert queue.fake_free_list_tail.prev_free_block is block
719
720
721	def test_free_kv_cache_block_queue_operations():
722	# Create a list of KVCacheBlock objects
723	blocks = [KVCacheBlock(block_id=i) for i in range(5)]
724
725	# Create a FreeKVCacheBlockQueue with these blocks
726	queue = FreeKVCacheBlockQueue(blocks)
727
728	# Check initial state
729	assert queue.num_free_blocks == 5
730	assert queue.fake_free_list_head.next_free_block is blocks[0]
731	assert queue.fake_free_list_tail.prev_free_block is blocks[4]
732
733	# Pop the first block
734	block1 = queue.popleft()
735	assert block1 == blocks[0]
736	assert queue.num_free_blocks == 4
737	assert queue.fake_free_list_head.next_free_block is blocks[1]
738	assert queue.fake_free_list_tail.prev_free_block is blocks[4]
739
740	# Remove a block from the middle
741	block_to_remove = blocks[2]
742	queue.remove(block_to_remove)
743	assert queue.num_free_blocks == 3
744	assert blocks[1].next_free_block is blocks[3]
745	assert blocks[3].prev_free_block is blocks[1]
746
747	# Append a block back
748	queue.append(block_to_remove)
749	assert queue.num_free_blocks == 4
750	assert queue.fake_free_list_tail.prev_free_block is block_to_remove
751	assert block_to_remove.prev_free_block is blocks[4]
752	assert block_to_remove.next_free_block is queue.fake_free_list_tail
753
754	# Pop blocks until empty
755	for _ in range(4):
756	queue.popleft()
757	assert queue.num_free_blocks == 0
758	assert queue.fake_free_list_head.next_free_block is queue.fake_free_list_tail
759	assert queue.fake_free_list_tail.prev_free_block is queue.fake_free_list_head
760
761	# Attempt to pop from an empty queue
762	with pytest.raises(ValueError) as e:
763	queue.popleft()
764	assert str(e.value) == "No free blocks available"
765
766
767	def test_free_kv_cache_block_queue_append_n():
768	# Create an empty FreeKVCacheBlockQueue with these blocks
769	queue = FreeKVCacheBlockQueue([])
770	blocks = [KVCacheBlock(block_id=i) for i in range(6)]
771	# Append 0 block
772	# fake_head->fake_tail
773	queue.append_n([])
774	assert queue.num_free_blocks == 0
775	assert queue.fake_free_list_head.next_free_block is queue.fake_free_list_tail
776	assert queue.fake_free_list_tail.prev_free_block is queue.fake_free_list_head
777	# Append 1 block
778	# fake_head->b0->fake_tail
779	queue.append_n(blocks[0:1])
780	assert queue.num_free_blocks == 1
781	assert queue.fake_free_list_head.next_free_block is blocks[0]
782	assert blocks[0].prev_free_block is queue.fake_free_list_head
783	assert blocks[0].next_free_block is queue.fake_free_list_tail
784	assert queue.fake_free_list_tail.prev_free_block is blocks[0]
785	# Append 2 blocks
786	# fake_head->b0->b4->b5->fake_tail
787	queue.append_n(blocks[4:6])
788	assert queue.num_free_blocks == 3
789	assert queue.fake_free_list_head.next_free_block is blocks[0]
790	assert blocks[0].prev_free_block is queue.fake_free_list_head
791	assert blocks[0].next_free_block is blocks[4]
792	assert blocks[4].prev_free_block is blocks[0]
793	assert blocks[4].next_free_block is blocks[5]
794	assert blocks[5].prev_free_block is blocks[4]
795	assert blocks[5].next_free_block is queue.fake_free_list_tail
796	assert queue.fake_free_list_tail.prev_free_block is blocks[5]
797	# Append 3 blocks
798	# fake_head->b0->b4->b5->b1->b2->b3->fake_tail
799	queue.append_n(blocks[1:4])
800	assert queue.num_free_blocks == 6
801	assert queue.fake_free_list_head.next_free_block is blocks[0]
802	assert blocks[0].prev_free_block is queue.fake_free_list_head
803	assert blocks[0].next_free_block is blocks[4]
804	assert blocks[4].prev_free_block is blocks[0]
805	assert blocks[4].next_free_block is blocks[5]
806	assert blocks[5].prev_free_block is blocks[4]
807	assert blocks[5].next_free_block is blocks[1]
808	assert blocks[1].prev_free_block is blocks[5]
809	assert blocks[1].next_free_block is blocks[2]
810	assert blocks[2].prev_free_block is blocks[1]
811	assert blocks[2].next_free_block is blocks[3]
812	assert blocks[3].prev_free_block is blocks[2]
813	assert blocks[3].next_free_block is queue.fake_free_list_tail
814	assert queue.fake_free_list_tail.prev_free_block is blocks[3]
815
816	# Create an empty FreeKVCacheBlockQueue
817	invalid_queue = FreeKVCacheBlockQueue([])
818	# set prev_free_block to None and this will cause assertion in append_n
819	invalid_queue.fake_free_list_tail.prev_free_block = None
820	with pytest.raises(AssertionError):
821	# Append 1 block
822	# fake_head->fake_tail
823	invalid_queue.append_n(blocks[0:1])
824	assert invalid_queue.num_free_blocks == 0
825	assert (
826	invalid_queue.fake_free_list_head.next_free_block
827	== invalid_queue.fake_free_list_tail
828	)
829
830
831	def test_free_kv_cache_block_queue_prepend_n():
832	# Seed the queue with one block so prepend has an existing head to splice
833	# in front of (fake_head->b0->fake_tail).
834	blocks = [KVCacheBlock(block_id=i) for i in range(6)]
835	queue = FreeKVCacheBlockQueue(blocks[0:1])
836
837	# Prepend 0 blocks is a no-op.
838	queue.prepend_n([])
839	assert queue.num_free_blocks == 1
840	assert queue.fake_free_list_head.next_free_block is blocks[0]
841
842	# Prepend 2 blocks; they land in front of the existing head, in order.
843	# fake_head->b4->b5->b0->fake_tail
844	queue.prepend_n(blocks[4:6])
845	assert queue.num_free_blocks == 3
846	assert queue.fake_free_list_head.next_free_block is blocks[4]
847	assert blocks[4].prev_free_block is queue.fake_free_list_head
848	assert blocks[4].next_free_block is blocks[5]
849	assert blocks[5].prev_free_block is blocks[4]
850	assert blocks[5].next_free_block is blocks[0]
851	assert blocks[0].prev_free_block is blocks[5]
852	assert blocks[0].next_free_block is queue.fake_free_list_tail
853	assert queue.fake_free_list_tail.prev_free_block is blocks[0]
854
855	# A second prepend goes ahead of everything previously prepended.
856	# fake_head->b1->b2->b4->b5->b0->fake_tail
857	queue.prepend_n(blocks[1:3])
858	assert queue.num_free_blocks == 5
859	assert queue.fake_free_list_head.next_free_block is blocks[1]
860	assert blocks[1].next_free_block is blocks[2]
861	assert blocks[2].next_free_block is blocks[4]
862
863	# The popleft order reflects the front-to-back queue order.
864	assert [queue.popleft().block_id for _ in range(5)] == [1, 2, 4, 5, 0]
865	assert queue.num_free_blocks == 0
866
867
868	def test_free_kv_cache_block_queue_popleft_n():
869	blocks = [KVCacheBlock(block_id=i) for i in range(6)]
870	# Create an empty FreeKVCacheBlockQueue with these blocks
871	queue = FreeKVCacheBlockQueue(
872	[blocks[1], blocks[3], blocks[5], blocks[4], blocks[0], blocks[2]]
873	)
874	assert queue.num_free_blocks == 6
875	assert queue.fake_free_list_head.next_free_block is blocks[1]
876	assert blocks[1].prev_free_block is queue.fake_free_list_head
877	assert blocks[1].next_free_block is blocks[3]
878	assert blocks[3].prev_free_block is blocks[1]
879	assert blocks[3].next_free_block is blocks[5]
880	assert blocks[5].prev_free_block is blocks[3]
881	assert blocks[5].next_free_block is blocks[4]
882	assert blocks[4].prev_free_block is blocks[5]
883	assert blocks[4].next_free_block is blocks[0]
884	assert blocks[0].prev_free_block is blocks[4]
885	assert blocks[0].next_free_block is blocks[2]
886	assert blocks[2].prev_free_block is blocks[0]
887	assert blocks[2].next_free_block is queue.fake_free_list_tail
888	assert queue.fake_free_list_tail.prev_free_block is blocks[2]
889
890	# Pop 0 block
891	# fake_head->b1->b3->b5->b4->b0->b2->fake_tail
892	assert len(queue.popleft_n(0)) == 0
893	assert queue.num_free_blocks == 6
894	# Pop 1 block
895	# fake_head->b3->b5->b4->b0->b2->fake_tail
896	result_blocks = queue.popleft_n(1)
897	assert queue.num_free_blocks == 5
898	assert len(result_blocks) == 1
899	assert result_blocks[0] is blocks[1]
900	for block in result_blocks:
901	assert block.prev_free_block is None
902	assert block.next_free_block is None
903	# Pop 2 blocks
904	# fake_head->b4->b0->b2->fake_tail
905	result_blocks = queue.popleft_n(2)
906	assert len(result_blocks) == 2
907	assert queue.num_free_blocks == 3
908	assert result_blocks[0] is blocks[3]
909	assert result_blocks[1] is blocks[5]
910	for block in result_blocks:
911	assert block.prev_free_block is None
912	assert block.next_free_block is None
913	# Pop 3 blocks
914	# fake_head->fake_tail
915	result_blocks = queue.popleft_n(3)
916	assert len(result_blocks) == 3
917	assert queue.num_free_blocks == 0
918	assert result_blocks[0] is blocks[4]
919	assert result_blocks[1] is blocks[0]
920	assert result_blocks[2] is blocks[2]
921	for block in result_blocks:
922	assert block.prev_free_block is None
923	assert block.next_free_block is None
924
925
926	def test_free_kv_cache_block_queue_get_all_free_blocks():
927	# Create a list of KVCacheBlock objects
928	blocks = [KVCacheBlock(block_id=i) for i in range(5)]
929
930	# Create a FreeKVCacheBlockQueue with these blocks
931	queue = FreeKVCacheBlockQueue(blocks)
932
933	# Check all blocks are correctly retrieved
934	assert queue.get_all_free_blocks() == blocks
935
936	# Pop a block and check again
937	queue.popleft()
938	assert queue.get_all_free_blocks() == blocks[1:]
939
940	# Remove a block and check again
941	block_to_remove = blocks[2]
942	queue.remove(block_to_remove)
943	assert queue.get_all_free_blocks() == blocks[1:2] + blocks[3:]
944
945	# Append a block back and check again
946	queue.append(block_to_remove)
947	assert queue.get_all_free_blocks() == blocks[1:2] + blocks[3:] + [block_to_remove]
948
949
950	def test_generate_block_hash_extra_keys():
951	request = make_request(
952	request_id="0",
953	prompt_token_ids=[_ for _ in range(20)],
954	mm_positions=[
955	PlaceholderRange(offset=0, length=5),
956	PlaceholderRange(offset=10, length=5),
957	],
958	mm_hashes=["hash1", "hash2"],
959	)
960
961	# Test with no extra keys
962	extra_keys, next_mm_idx = generate_block_hash_extra_keys(request, 0, 5, 0)
963	assert extra_keys == (("mm", "hash1", 0),)
964	assert next_mm_idx == 1
965
966	# Test with partial overlap
967	extra_keys, next_mm_idx = generate_block_hash_extra_keys(request, 3, 8, 0)
968	assert extra_keys == (("mm", "hash1", -3),)
969	assert next_mm_idx == 1
970
971	# Test with no overlap
972	extra_keys, next_mm_idx = generate_block_hash_extra_keys(request, 6, 10, 0)
973	assert extra_keys is None
974	assert next_mm_idx == 1
975
976	# Test with multiple extra keys
977	extra_keys, next_mm_idx = generate_block_hash_extra_keys(request, 0, 15, 0)
978	assert extra_keys == (("mm", "hash1", 0), ("mm", "hash2", 10))
979	assert next_mm_idx == 2
980
981
982	def test_generate_block_hash_extra_keys_no_mm_inputs():
983	request = make_request(
984	request_id="0",
985	prompt_token_ids=[_ for _ in range(6)],
986	mm_positions=None,
987	mm_hashes=None,
988	)
989
990	extra_keys, next_mm_idx = generate_block_hash_extra_keys(request, 0, 5, 0)
991	assert extra_keys is None
992	assert next_mm_idx == 0
993
994
995	def test_generate_block_hash_extra_keys_cache_salt():
996	request = make_request(
997	request_id="0",
998	prompt_token_ids=[_ for _ in range(6)],
999	mm_positions=None,
1000	mm_hashes=None,
1001	cache_salt="salt",
1002	)
1003
1004	# salt is added for the first token
1005	extra_keys, _ = generate_block_hash_extra_keys(request, 0, 1, 0)
1006	assert extra_keys == (("cache_salt", "salt"),)
1007	extra_keys, _ = generate_block_hash_extra_keys(request, 0, 10, 0)
1008	assert extra_keys == (("cache_salt", "salt"),)
1009
1010	# no salt added for other tokens
1011	extra_keys, _ = generate_block_hash_extra_keys(request, 1, 2, 0)
1012	assert extra_keys is None
1013	extra_keys, _ = generate_block_hash_extra_keys(request, 6, 10, 0)
1014	assert extra_keys is None
1015
1016	# works together with other extra keys
1017	request_mm = make_request(
1018	request_id="0",
1019	prompt_token_ids=[_ for _ in range(20)],
1020	mm_positions=[
1021	PlaceholderRange(offset=0, length=5),
1022	],
1023	mm_hashes=["hash1"],
1024	cache_salt="salt",
1025	)
1026
1027	# Test with no extra keys
1028	extra_keys, next_mm_idx = generate_block_hash_extra_keys(request_mm, 0, 5, 0)
1029	assert extra_keys == (("mm", "hash1", 0), ("cache_salt", "salt"))
1030	assert next_mm_idx == 1
1031
1032
1033	def test_generate_block_hash_extra_keys_prompt_embeds():
1034	prompt_embeds = torch.randn(10, 3)
1035	request = make_request(
1036	request_id="0",
1037	prompt_token_ids=None,
1038	mm_positions=None,
1039	mm_hashes=None,
1040	prompt_embeds=prompt_embeds,
1041	)
1042
1043	# Test with prompt embeds for the first block
1044	extra_keys, _ = generate_block_hash_extra_keys(request, 0, 5, 0)
1045	expected_embeds = prompt_embeds[0:5]
1046	expected_hash = hashlib.sha256(kv_cache_utils.tensor_data(expected_embeds)).digest()
1047	assert extra_keys == (("prompt_embeds", expected_hash),)
1048
1049	# Test with prompt embeds for the second block
1050	extra_keys, _ = generate_block_hash_extra_keys(request, 5, 10, 0)
1051	expected_embeds = prompt_embeds[5:10]
1052	expected_hash = hashlib.sha256(kv_cache_utils.tensor_data(expected_embeds)).digest()
1053	assert extra_keys == (("prompt_embeds", expected_hash),)
1054
1055
1056	def test_generate_block_hash_extra_keys_prompt_embeds_cached(monkeypatch):
1057	prompt_embeds = torch.randn(10, 3)
1058	request = make_request(
1059	request_id="0",
1060	prompt_token_ids=None,
1061	mm_positions=None,
1062	mm_hashes=None,
1063	prompt_embeds=prompt_embeds,
1064	block_size=20,
1065	)
1066
1067	num_tensor_data_calls = 0
1068	original_tensor_data = kv_cache_utils.tensor_data
1069
1070	def counting_tensor_data(tensor: torch.Tensor):
1071	nonlocal num_tensor_data_calls
1072	num_tensor_data_calls += 1
1073	return original_tensor_data(tensor)
1074
1075	monkeypatch.setattr(kv_cache_utils, "tensor_data", counting_tensor_data)
1076
1077	extra_keys_1, _ = generate_block_hash_extra_keys(request, 0, 5, 0)
1078	extra_keys_2, _ = generate_block_hash_extra_keys(request, 0, 5, 0)
1079	assert extra_keys_1 == extra_keys_2
1080	assert num_tensor_data_calls == 1
1081
1082
1083	def test_generate_block_hash_extra_keys_different_prompt_embeds():
1084	prompt_embeds1 = torch.randn(10, 3)
1085	prompt_embeds2 = torch.randn(10, 3)
1086	request1 = make_request(
1087	request_id="0",
1088	prompt_token_ids=None,
1089	mm_positions=None,
1090	mm_hashes=None,
1091	prompt_embeds=prompt_embeds1,
1092	)
1093	request2 = make_request(
1094	request_id="1",
1095	prompt_token_ids=None,
1096	mm_positions=None,
1097	mm_hashes=None,
1098	prompt_embeds=prompt_embeds2,
1099	)
1100
1101	extra_keys1, _ = generate_block_hash_extra_keys(request1, 0, 5, 0)
1102	extra_keys2, _ = generate_block_hash_extra_keys(request2, 0, 5, 0)
1103	assert extra_keys1 != extra_keys2
1104
1105
1106	def test_generate_block_hash_extra_keys_lora():
1107	request = make_request(
1108	request_id="0",
1109	prompt_token_ids=[_ for _ in range(6)],
1110	)
1111
1112	request.lora_request = LoRARequest(
1113	lora_name="test_lora_adapter", lora_int_id=1, lora_path="/path/to/lora"
1114	)
1115
1116	extra_keys, _ = generate_block_hash_extra_keys(request, 0, 3, 0)
1117	assert extra_keys == (("lora", "test_lora_adapter", "/path/to/lora"),)
1118
1119	request.lora_request = None
1120	extra_keys, _ = generate_block_hash_extra_keys(request, 0, 3, 0)
1121	assert extra_keys is None
1122
1123
1124	@pytest.mark.parametrize("hash_fn", [sha256, sha256_cbor])
1125	def test_lora_name_and_cache_salt_block_hashes_do_not_collide(hash_fn):
1126	"""A cache_salt equal to a LoRA name must not share that LoRA's blocks."""
1127	lora_request = LoRARequest(
1128	lora_name="foo", lora_int_id=1, lora_path="/path/to/lora"
1129	)
1130	lora_req = Request(
1131	request_id="lora",
1132	prompt_token_ids=[0, 1, 2],
1133	sampling_params=SamplingParams(max_tokens=1),
1134	pooling_params=None,
1135	lora_request=lora_request,
1136	block_hasher=get_request_block_hasher(3, hash_fn),
1137	)
1138	salted_req = make_request("salted", [0, 1, 2], hash_fn=hash_fn, cache_salt="foo")
1139
1140	assert lora_req.block_hashes[0] != salted_req.block_hashes[0]
1141
1142
1143	@pytest.mark.parametrize("hash_fn", [sha256, sha256_cbor])
1144	def test_lora_path_change_changes_block_hashes(hash_fn):
1145	"""Re-pointing a LoRA name at another adapter must not reuse its blocks."""
1146
1147	def make_lora_request(lora_path: str) -> Request:
1148	return Request(
1149	request_id=lora_path,
1150	prompt_token_ids=[0, 1, 2],
1151	sampling_params=SamplingParams(max_tokens=1),
1152	pooling_params=None,
1153	lora_request=LoRARequest(
1154	lora_name="foo", lora_int_id=1, lora_path=lora_path
1155	),
1156	block_hasher=get_request_block_hasher(3, hash_fn),
1157	)
1158
1159	first = make_lora_request("/path/to/lora_a")
1160	second = make_lora_request("/path/to/lora_b")
1161
1162	assert first.block_hashes[0] != second.block_hashes[0]
1163
1164
1165	def test_to_event_extra_keys_keeps_untagged_event_format():
1166	"""KV events keep publishing the extra-key shapes consumers already parse."""
1167	request = make_request(
1168	"0",
1169	list(range(10)),
1170	mm_positions=[PlaceholderRange(offset=2, length=3)],
1171	mm_hashes=["hash1"],
1172	cache_salt="salt",
1173	)
1174	request.lora_request = LoRARequest(
1175	lora_name="adapter", lora_int_id=1, lora_path="/path/to/lora"
1176	)
1177
1178	extra_keys, _ = generate_block_hash_extra_keys(request, 0, 5, 0)
1179
1180	assert extra_keys == (
1181	("lora", "adapter", "/path/to/lora"),
1182	("mm", "hash1", 2),
1183	("cache_salt", "salt"),
1184	)
1185	event_keys = ("adapter", ("hash1", 2), "salt")
1186	assert to_event_extra_keys([extra_keys, None]) == [event_keys, None]
1187	assert to_event_extra_keys([]) is None
1188	assert to_event_extra_keys(None) is None
1189
1190
1191	@pytest.mark.parametrize("hash_fn", [sha256, sha256_cbor])
1192	def test_hash_block_tokens(hash_fn):
1193	parent_block_hash = BlockHash(b"123")
1194	curr_block_token_ids = (1, 2, 3)
1195	extra_keys = ("key1", "key2")
1196
1197	block_hash = hash_block_tokens(
1198	hash_fn, parent_block_hash, curr_block_token_ids, extra_keys
1199	)
1200	expected = hash_fn((parent_block_hash, curr_block_token_ids, extra_keys))
1201	assert block_hash == expected
1202
1203
1204	@pytest.mark.parametrize("hash_fn", [sha256, sha256_cbor])
1205	def test_request_block_hasher(hash_fn):
1206	request = make_request(
1207	request_id="0",
1208	prompt_token_ids=[_ for _ in range(6)],
1209	block_size=3,
1210	hash_fn=hash_fn,
1211	mm_positions=[
1212	PlaceholderRange(offset=0, length=3),
1213	PlaceholderRange(offset=3, length=3),
1214	],
1215	mm_hashes=["hash1", "hash2"],
1216	)
1217
1218	block_hashes = request.block_hashes
1219	assert len(block_hashes) == 2
1220	assert block_hashes[0] == hash_fn(
1221	(kv_cache_utils.NONE_HASH, (0, 1, 2), (("mm", "hash1", 0),))
1222	)
1223	assert block_hashes[1] == hash_fn(
1224	(block_hashes[0], (3, 4, 5), (("mm", "hash2", 0),))
1225	)
1226
1227
1228	@pytest.mark.parametrize("hash_fn", [sha256, sha256_cbor])
1229	def test_request_block_hasher_incremental_append_with_multiple_mm_features(hash_fn):
1230	mm_positions = [
1231	PlaceholderRange(offset=4, length=2),
1232	PlaceholderRange(offset=6, length=1),
1233	]
1234	incremental = make_request(
1235	request_id="incremental",
1236	prompt_token_ids=list(range(7)),
1237	block_size=4,
1238	hash_fn=hash_fn,
1239	mm_positions=mm_positions,
1240	mm_hashes=["A", "B"],
1241	)
1242	incremental.append_output_token_ids(7)
1243	fresh = make_request(
1244	request_id="fresh",
1245	prompt_token_ids=list(range(8)),
1246	block_size=4,
1247	hash_fn=hash_fn,
1248	mm_positions=mm_positions,
1249	mm_hashes=["A", "B"],
1250	)
1251
1252	expected_second_hash = hash_fn(
1253	(incremental.block_hashes[0], (4, 5, 6, 7), (("mm", "A", 0), ("mm", "B", 2)))
1254	)
1255	assert incremental.block_hashes[1] == expected_second_hash
1256	assert incremental.block_hashes == fresh.block_hashes
1257
1258
1259	@pytest.mark.parametrize("hash_fn", [sha256, sha256_cbor])
1260	def test_hash_tokens_different_mm_input(hash_fn):
1261	request1 = make_request(
1262	request_id="0",
1263	prompt_token_ids=[_ for _ in range(6)],
1264	block_size=3,
1265	hash_fn=hash_fn,
1266	mm_positions=[
1267	PlaceholderRange(offset=0, length=3),
1268	PlaceholderRange(offset=3, length=3),
1269	],
1270	mm_hashes=["hash1", "hash2"],
1271	)
1272	request2 = make_request(
1273	request_id="1",
1274	prompt_token_ids=[_ for _ in range(6)],
1275	mm_positions=[
1276	PlaceholderRange(offset=0, length=3),
1277	PlaceholderRange(offset=3, length=3),
1278	],
1279	mm_hashes=["hash3", "hash2"],
1280	)
1281	block_hashes1 = request1.block_hashes
1282	block_hashes2 = request2.block_hashes
1283	assert block_hashes1[0] != block_hashes2[0]
1284	assert block_hashes1[1] != block_hashes2[1]
1285
1286
1287	@pytest.mark.parametrize("hash_fn", [sha256, sha256_cbor])
1288	def test_hash_request_tokens_no_mm_inputs(hash_fn):
1289	request = make_request(
1290	request_id="0",
1291	prompt_token_ids=[_ for _ in range(6)],
1292	block_size=3,
1293	hash_fn=hash_fn,
1294	mm_positions=None,
1295	mm_hashes=None,
1296	)
1297
1298	block_hashes = request.block_hashes
1299
1300	assert len(block_hashes) == 2
1301	assert block_hashes[0] == hash_fn((kv_cache_utils.NONE_HASH, (0, 1, 2), None))
1302	assert block_hashes[1] == hash_fn((block_hashes[0], (3, 4, 5), None))
1303
1304
1305	def _stats(requests: int, queries: int, hits: int) -> PrefixCacheStats:
1306	return PrefixCacheStats(requests=requests, queries=queries, hits=hits)
1307
1308
1309	def test_metrics_empty_distinguishes_no_queries_from_no_hits():
1310	"""`hit_rate` alone cannot tell the two apart; `empty` can.
1311
1312	Both an unobserved window and a genuine all-miss window report a hit
1313	rate of 0.0, so anything surfacing that number to a human has to check
1314	`empty` first - which is what the prefix-cache log line does.
1315	"""
1316	metrics = CachingMetrics(max_recent_requests=5)
1317	assert metrics.empty
1318	assert metrics.hit_rate == 0.0
1319
1320	metrics.observe(_stats(1, 20, 0))
1321	assert not metrics.empty
1322	assert metrics.hit_rate == 0.0
1323
1324
1325	def test_metrics():
1326	"""Test the prefix caching metrics."""
1327	metrics = CachingMetrics(max_recent_requests=5)
1328	assert metrics.hit_rate == 0.0
1329
1330	metrics.observe(_stats(1, 20, 9))
1331	# 9 / 20 = 0.45
1332	assert metrics.hit_rate == 0.45
1333
1334	metrics.observe(_stats(4, 80, 16))
1335
1336	# 25 / 100 = 0.25
1337	assert metrics.hit_rate == 0.25
1338
1339	metrics.observe(_stats(1, 10, 2))
1340
1341	# Remove (20, 9) and add (10, 2): 18 / 90 = 0.2
1342	assert metrics.aggregated_requests == 5
1343	assert metrics.aggregated_query_total == 90
1344	assert metrics.aggregated_query_hit == 18
1345	assert metrics.hit_rate == 0.2
1346
1347	metrics.reset()
1348	assert metrics.hit_rate == 0.0
1349	assert metrics.aggregated_requests == 0
1350	assert metrics.aggregated_query_total == 0
1351	assert metrics.aggregated_query_hit == 0
1352	assert not metrics.query_queue
1353
1354
1355	def test_metrics_empty_stats():
1356	"""Test the prefix caching metrics with empty stats."""
1357	metrics = CachingMetrics(max_recent_requests=5)
1358	metrics.observe(_stats(0, 0, 0))
1359	metrics.observe(_stats(1, 20, 9))
1360	metrics.observe(_stats(0, 0, 0))
1361	metrics.observe(_stats(4, 80, 16))
1362	metrics.observe(_stats(0, 0, 0))
1363	metrics.observe(_stats(1, 10, 2))
1364	# Remove (20, 9) and add (10, 2): 18 / 90 = 0.2
1365	assert metrics.aggregated_requests == 5
1366	assert metrics.aggregated_query_total == 90
1367	assert metrics.aggregated_query_hit == 18
1368	assert metrics.hit_rate == 0.2
1369
1370	# Only the latest added stats preserved 10 / 20 = 0.5
1371	metrics.observe(_stats(11, 20, 10))
1372	assert metrics.aggregated_requests == 11
1373	assert metrics.aggregated_query_total == 20
1374	assert metrics.aggregated_query_hit == 10
1375	assert metrics.hit_rate == 0.5
1376
1377	# Only the latest added stats preserved 30 / 40 = 0.75
1378	metrics.observe(_stats(22, 40, 30))
1379	assert metrics.aggregated_requests == 22
1380	assert metrics.aggregated_query_total == 40
1381	assert metrics.aggregated_query_hit == 30
1382	assert metrics.hit_rate == 0.75
1383
1384
1385	def test_get_kv_cache_configs_multiple_workers():
1386	model_config = ModelConfig(max_model_len=16)
1387	vllm_config = VllmConfig(model_config=model_config)
1388	vllm_config.cache_config.kv_cache_layout = "LBNHC"
1389	vllm_config.cache_config.prefix_cache_retention_interval = None
1390
1391	ref_kv_cache_spec = new_kv_cache_spec()
1392	same_kv_cache_specs = [
1393	{
1394	"layer1": new_kv_cache_spec(),
1395	"layer2": new_kv_cache_spec(),
1396	},
1397	{
1398	"layer1": new_kv_cache_spec(),
1399	"layer2": new_kv_cache_spec(),
1400	},
1401	]
1402
1403	# Basic case. All things are the same.
1404	kv_cache_configs = get_kv_cache_configs(
1405	vllm_config,
1406	same_kv_cache_specs,
1407	[
1408	ref_kv_cache_spec.page_size_bytes * 2 * 10,
1409	ref_kv_cache_spec.page_size_bytes * 2 * 10,
1410	],
1411	)
1412	expected = KVCacheConfig(
1413	num_blocks=10,
1414	kv_cache_tensors=[
1415	KVCacheTensor(
1416	size=ref_kv_cache_spec.page_size_bytes * 10 * 2,
1417	layers=["layer1", "layer2"],
1418	layer_stride=ref_kv_cache_spec.page_size_bytes * 10,
1419	block_stride=ref_kv_cache_spec.page_size_bytes,
1420	),
1421	],
1422	kv_cache_groups=[
1423	KVCacheGroupSpec(["layer1", "layer2"], ref_kv_cache_spec),
1424	],
1425	)
1426	assert kv_cache_configs == [expected, expected]
1427
1428	# Different available memory. This is the case for TP.
1429	# Use the smallest memory available.
1430	kv_cache_configs = get_kv_cache_configs(
1431	vllm_config,
1432	same_kv_cache_specs,
1433	[
1434	ref_kv_cache_spec.page_size_bytes * 2 * 10,
1435	ref_kv_cache_spec.page_size_bytes * 2 * 20,
1436	],
1437	)
1438	assert kv_cache_configs == [expected, expected]
1439
1440	# Different KV cache specs. This is the case for PP.
1441	different_layer_specs = [
1442	{
1443	"layer1": new_kv_cache_spec(),
1444	},
1445	{
1446	"layer2": new_kv_cache_spec(),
1447	"layer3": new_kv_cache_spec(),
1448	},
1449	]
1450
1451	# Different workers have different layers.
1452	kv_cache_configs = get_kv_cache_configs(
1453	vllm_config,
1454	different_layer_specs,
1455	[
1456	ref_kv_cache_spec.page_size_bytes * 2 * 10,
1457	ref_kv_cache_spec.page_size_bytes * 2 * 10,
1458	],
1459	)
1460	assert kv_cache_configs == [
1461	KVCacheConfig(
1462	num_blocks=10,
1463	kv_cache_tensors=[
1464	KVCacheTensor(
1465	size=ref_kv_cache_spec.page_size_bytes * 10,
1466	layers=["layer1"],
1467	layer_stride=ref_kv_cache_spec.page_size_bytes * 10,
1468	block_stride=ref_kv_cache_spec.page_size_bytes,
1469	),
1470	],
1471	kv_cache_groups=[
1472	KVCacheGroupSpec(["layer1"], new_kv_cache_spec()),
1473	],
1474	),
1475	KVCacheConfig(
1476	num_blocks=10,
1477	kv_cache_tensors=[
1478	KVCacheTensor(
1479	size=ref_kv_cache_spec.page_size_bytes * 10 * 2,
1480	layers=["layer2", "layer3"],
1481	layer_stride=ref_kv_cache_spec.page_size_bytes * 10,
1482	block_stride=ref_kv_cache_spec.page_size_bytes,
1483	),
1484	],
1485	kv_cache_groups=[
1486	KVCacheGroupSpec(["layer2", "layer3"], new_kv_cache_spec()),
1487	],
1488	),
1489	]
1490
1491	# Some layers are the same, some are different. This is the case for TP+PP
1492	tp_pp_kv_cache_specs = [
1493	{
1494	"layer1": new_kv_cache_spec(),
1495	"layer2": new_kv_cache_spec(),
1496	},
1497	{
1498	"layer1": new_kv_cache_spec(),
1499	"layer2": new_kv_cache_spec(),
1500	},
1501	{
1502	"layer3": new_kv_cache_spec(),
1503	},
1504	{
1505	"layer3": new_kv_cache_spec(),
1506	},
1507	]
1508
1509	kv_cache_configs = get_kv_cache_configs(
1510	vllm_config,
1511	tp_pp_kv_cache_specs,
1512	[
1513	ref_kv_cache_spec.page_size_bytes * 2 * 10,
1514	ref_kv_cache_spec.page_size_bytes * 2 * 10,
1515	ref_kv_cache_spec.page_size_bytes * 2 * 10,
1516	ref_kv_cache_spec.page_size_bytes * 2 * 10,
1517	],
1518	)
1519	expected_12 = KVCacheConfig(
1520	num_blocks=10,
1521	kv_cache_tensors=[
1522	KVCacheTensor(
1523	size=ref_kv_cache_spec.page_size_bytes * 10 * 2,
1524	layers=["layer1", "layer2"],
1525	layer_stride=ref_kv_cache_spec.page_size_bytes * 10,
1526	block_stride=ref_kv_cache_spec.page_size_bytes,
1527	),
1528	],
1529	kv_cache_groups=[
1530	KVCacheGroupSpec(["layer1", "layer2"], ref_kv_cache_spec),
1531	],
1532	)
1533	expected_3 = KVCacheConfig(
1534	num_blocks=10,
1535	kv_cache_tensors=[
1536	KVCacheTensor(
1537	size=ref_kv_cache_spec.page_size_bytes * 10,
1538	layers=["layer3"],
1539	layer_stride=ref_kv_cache_spec.page_size_bytes * 10,
1540	block_stride=ref_kv_cache_spec.page_size_bytes,
1541	),
1542	],
1543	kv_cache_groups=[
1544	KVCacheGroupSpec(["layer3"], ref_kv_cache_spec),
1545	],
1546	)
1547	assert kv_cache_configs == [
1548	expected_12,
1549	expected_12,
1550	expected_3,
1551	expected_3,
1552	]
1553
1554	# Different workers have different types of layers. This is the case for
1555	# hybrid models + PP.
1556	different_type_layer_specs = [
1557	{
1558	"layer1": new_kv_cache_spec(),
1559	"layer2": new_kv_cache_spec(),
1560	},
1561	{
1562	"layer3": new_sliding_window_spec(),
1563	"layer4": new_sliding_window_spec(),
1564	},
1565	]
1566	kv_cache_configs = get_kv_cache_configs(
1567	vllm_config,
1568	different_type_layer_specs,
1569	[
1570	ref_kv_cache_spec.page_size_bytes * 2 * 10,
1571	ref_kv_cache_spec.page_size_bytes * 2 * 10,
1572	],
1573	)
1574	assert kv_cache_configs == [
1575	KVCacheConfig(
1576	num_blocks=10,
1577	kv_cache_tensors=[
1578	KVCacheTensor(
1579	size=ref_kv_cache_spec.page_size_bytes * 10 * 2,
1580	layers=["layer1", "layer2"],
1581	layer_stride=ref_kv_cache_spec.page_size_bytes * 10,
1582	block_stride=ref_kv_cache_spec.page_size_bytes,
1583	),
1584	],
1585	kv_cache_groups=[
1586	KVCacheGroupSpec(["layer1", "layer2"], ref_kv_cache_spec),
1587	KVCacheGroupSpec([], new_sliding_window_spec()),
1588	],
1589	),
1590	KVCacheConfig(
1591	num_blocks=10,
1592	kv_cache_tensors=[
1593	KVCacheTensor(
1594	size=ref_kv_cache_spec.page_size_bytes * 10 * 2,
1595	layers=["layer3", "layer4"],
1596	layer_stride=ref_kv_cache_spec.page_size_bytes * 10,
1597	block_stride=ref_kv_cache_spec.page_size_bytes,
1598	),
1599	],
1600	kv_cache_groups=[
1601	KVCacheGroupSpec([], ref_kv_cache_spec),
1602	KVCacheGroupSpec(["layer3", "layer4"], new_sliding_window_spec()),
1603	],
1604	),
1605	]
1606
1607	# When divided into multiple KVCacheGroups, need to ensure the number of
1608	# layers per group is similar.
1609	different_type_layer_specs = [
1610	{
1611	"layer1": new_kv_cache_spec(),
1612	"layer2": new_sliding_window_spec(),
1613	"layer3": new_sliding_window_spec(),
1614	},
1615	{
1616	"layer4": new_kv_cache_spec(),
1617	"layer5": new_sliding_window_spec(),
1618	"layer6": new_sliding_window_spec(),
1619	},
1620	]
1621	kv_cache_configs = get_kv_cache_configs(
1622	vllm_config,
1623	different_type_layer_specs,
1624	[
1625	ref_kv_cache_spec.page_size_bytes * 10,
1626	ref_kv_cache_spec.page_size_bytes * 10,
1627	],
1628	)
1629	assert kv_cache_configs == [
1630	KVCacheConfig(
1631	num_blocks=10,
1632	kv_cache_tensors=[
1633	KVCacheTensor(
1634	size=ref_kv_cache_spec.page_size_bytes * 10,
1635	layers=["layer1"],
1636	layer_stride=ref_kv_cache_spec.page_size_bytes * 10,
1637	block_stride=ref_kv_cache_spec.page_size_bytes,
1638	),
1639	KVCacheTensor(
1640	size=ref_kv_cache_spec.page_size_bytes * 10,
1641	layers=["layer2"],
1642	layer_stride=ref_kv_cache_spec.page_size_bytes * 10,
1643	block_stride=ref_kv_cache_spec.page_size_bytes,
1644	),
1645	KVCacheTensor(
1646	size=ref_kv_cache_spec.page_size_bytes * 10,
1647	layers=["layer3"],
1648	layer_stride=ref_kv_cache_spec.page_size_bytes * 10,
1649	block_stride=ref_kv_cache_spec.page_size_bytes,
1650	),
1651	],
1652	kv_cache_groups=[
1653	KVCacheGroupSpec(["layer1"], ref_kv_cache_spec),
1654	KVCacheGroupSpec(["layer2"], new_sliding_window_spec()),
1655	KVCacheGroupSpec(["layer3"], new_sliding_window_spec()),
1656	],
1657	),
1658	KVCacheConfig(
1659	num_blocks=10,
1660	kv_cache_tensors=[
1661	KVCacheTensor(
1662	size=ref_kv_cache_spec.page_size_bytes * 10,
1663	layers=["layer4"],
1664	layer_stride=ref_kv_cache_spec.page_size_bytes * 10,
1665	block_stride=ref_kv_cache_spec.page_size_bytes,
1666	),
1667	KVCacheTensor(
1668	size=ref_kv_cache_spec.page_size_bytes * 10,
1669	layers=["layer5"],
1670	layer_stride=ref_kv_cache_spec.page_size_bytes * 10,
1671	block_stride=ref_kv_cache_spec.page_size_bytes,
1672	),
1673	KVCacheTensor(
1674	size=ref_kv_cache_spec.page_size_bytes * 10,
1675	layers=["layer6"],
1676	layer_stride=ref_kv_cache_spec.page_size_bytes * 10,
1677	block_stride=ref_kv_cache_spec.page_size_bytes,
1678	),
1679	],
1680	kv_cache_groups=[
1681	KVCacheGroupSpec(["layer4"], ref_kv_cache_spec),
1682	KVCacheGroupSpec(["layer5"], new_sliding_window_spec()),
1683	KVCacheGroupSpec(["layer6"], new_sliding_window_spec()),
1684	],
1685	),
1686	]
1687
1688	# Have conflicting layers. Need to raise an error.
1689	conflicting_layer_specs = [
1690	{
1691	"layer1": new_kv_cache_spec(),
1692	},
1693	{
1694	"layer1": new_sliding_window_spec(),
1695	},
1696	]
1697	with pytest.raises(AssertionError):
1698	get_kv_cache_configs(
1699	vllm_config,
1700	conflicting_layer_specs,
1701	[
1702	ref_kv_cache_spec.page_size_bytes * 2 * 10,
1703	ref_kv_cache_spec.page_size_bytes * 2 * 10,
1704	],
1705	)
1706
1707
1708	@pytest.mark.parametrize(
1709	"asymmetric_memory",
1710	[False, True],
1711	ids=["symmetric", "asymmetric"],
1712	)
1713	def test_get_kv_cache_configs_pp_sharding(asymmetric_memory):
1714	model_config = ModelConfig(max_model_len=512)
1715	vllm_config = VllmConfig(model_config=model_config)
1716	vllm_config.cache_config.kv_cache_layout = "LBNHC"
1717	vllm_config.cache_config.prefix_cache_retention_interval = None
1718
1719	ref_kv_cache_spec = new_kv_cache_spec()
1720	pp_kv_cache_specs = [
1721	{"layer1": ref_kv_cache_spec},
1722	{"layer2": ref_kv_cache_spec},
1723	]
1724
1725	expected_num_blocks = model_config.max_model_len // ref_kv_cache_spec.block_size + 1
1726	avail_memory = ref_kv_cache_spec.page_size_bytes * expected_num_blocks
1727
1728	# With per-worker validation, each worker only needs memory for its own
1729	# layers. Worker 2 having more memory shouldn't affect worker 1's config.
1730	available_memory = (
1731	[avail_memory, avail_memory * 2] if asymmetric_memory else [avail_memory] * 2
1732	)
1733
1734	kv_cache_configs = get_kv_cache_configs(
1735	vllm_config,
1736	pp_kv_cache_specs,
1737	available_memory,
1738	)
1739
1740	assert kv_cache_configs == [
1741	KVCacheConfig(
1742	num_blocks=expected_num_blocks,
1743	kv_cache_tensors=[
1744	KVCacheTensor(
1745	size=ref_kv_cache_spec.page_size_bytes * expected_num_blocks,
1746	layers=["layer1"],
1747	layer_stride=ref_kv_cache_spec.page_size_bytes
1748	* expected_num_blocks,
1749	block_stride=ref_kv_cache_spec.page_size_bytes,
1750	),
1751	],
1752	kv_cache_groups=[KVCacheGroupSpec(["layer1"], ref_kv_cache_spec)],
1753	),
1754	KVCacheConfig(
1755	num_blocks=expected_num_blocks,
1756	kv_cache_tensors=[
1757	KVCacheTensor(
1758	size=ref_kv_cache_spec.page_size_bytes * expected_num_blocks,
1759	layers=["layer2"],
1760	layer_stride=ref_kv_cache_spec.page_size_bytes
1761	* expected_num_blocks,
1762	block_stride=ref_kv_cache_spec.page_size_bytes,
1763	),
1764	],
1765	kv_cache_groups=[KVCacheGroupSpec(["layer2"], ref_kv_cache_spec)],
1766	),
1767	]
1768
1769
1770	def test_project_kv_cache_groups_to_worker():
1771	spec_a = new_kv_cache_spec()
1772	spec_b = new_kv_cache_spec(num_kv_heads=4)
1773
1774	global_groups = [
1775	KVCacheGroupSpec(["layer1", "layer2", "layer3"], spec_a),
1776	]
1777	worker_spec = {"layer1": spec_a, "layer2": spec_a}
1778	projected = kv_cache_utils._project_kv_cache_groups_to_worker(
1779	global_groups, worker_spec
1780	)
1781	assert len(projected) == 1
1782	assert projected[0].layer_names == ["layer1", "layer2"]
1783	assert projected[0].kv_cache_spec is spec_a
1784
1785	projected = kv_cache_utils._project_kv_cache_groups_to_worker(
1786	global_groups, {"layer4": spec_a}
1787	)
1788	assert len(projected) == 1
1789	assert projected[0].layer_names == []
1790	assert projected[0].kv_cache_spec is spec_a
1791
1792	uniform_spec = UniformTypeKVCacheSpecs(
1793	block_size=16,
1794	kv_cache_specs={"layer1": spec_a, "layer2": spec_b, "layer3": spec_a},
1795	)
1796	global_groups_uniform = [
1797	KVCacheGroupSpec(["layer1", "layer2", "layer3"], uniform_spec),
1798	]
1799	projected = kv_cache_utils._project_kv_cache_groups_to_worker(
1800	global_groups_uniform, {"layer1": spec_a, "layer3": spec_a}
1801	)
1802	assert len(projected) == 1
1803	assert projected[0].layer_names == ["layer1", "layer3"]
1804	proj_spec = projected[0].kv_cache_spec
1805	assert isinstance(proj_spec, UniformTypeKVCacheSpecs)
1806	assert set(proj_spec.kv_cache_specs.keys()) == {"layer1", "layer3"}
1807
1808
1809	@pytest.mark.parametrize("sliding_window", [None, 256])
1810	@pytest.mark.parametrize("disable_hybrid", [False, True])
1811	@pytest.mark.parametrize("pcp_size", [1, 4])
1812	def test_dcp_target_allocates_replicated_draft_independently(
1813	monkeypatch, sliding_window, disable_hybrid, pcp_size
1814	):
1815	"""A draft must retain all positions even when the target shards them."""
1816	from vllm.v1.attention.backends.utils import resolve_kv_cache_layout
1817
1818	monkeypatch.delenv("VLLM_KV_CACHE_LAYOUT", raising=False)
1819	config = VllmConfig(model_config=ModelConfig(max_model_len=1024))
1820	config.parallel_config.decode_context_parallel_size = 4
1821	config.parallel_config.prefill_context_parallel_size = pcp_size
1822	config.scheduler_config.disable_hybrid_kv_cache_manager = disable_hybrid
1823	config.cache_config.block_size = 16
1824	config.cache_config.kv_cache_layout = None
1825	draft_args = dict(
1826	block_size=16,
1827	num_kv_heads=8,
1828	head_size=64,
1829	dtype=torch.bfloat16,
1830	dcp_sharded=False,
1831	)
1832	draft = (
1833	FullAttentionSpec(**draft_args)
1834	if sliding_window is None
1835	else SlidingWindowSpec(**draft_args, sliding_window=sliding_window)
1836	)
1837	specs = {"target": new_mla_spec(), "draft": draft}
1838	layout = resolve_kv_cache_layout(config, [["LBHNC", "BLHNC"]], specs.values())
1839	assert layout == KVCacheLayout.BLHNC
1840	groups = get_kv_cache_groups(config, specs)
1841	assert [group.layer_names for group in groups] == [["target"], ["draft"]]
1842	widths = [g.kv_cache_spec.max_num_blocks_per_req(config, 1024) for g in groups]
1843	assert widths == [16, 64]
1844
1845	cache_config = kv_cache_utils.get_kv_cache_config_from_groups(
1846	config, groups, available_memory=16 * 1024 * 1024
1847	)
1848	scheduler_config = generate_scheduler_kv_cache_config([cache_config])
1849	assert get_max_concurrency_for_kv_cache_config(config, cache_config) > 0
1850	assert get_max_concurrency_for_kv_cache_config(config, cache_config) == (
1851	get_max_concurrency_for_kv_cache_config(config, scheduler_config)
1852	)
1853	manager = KVCacheManager(
1854	scheduler_config,
1855	max_model_len=1024,
1856	hash_block_size=16,
1857	scheduler_block_size=64,
1858	dcp_world_size=4,
1859	pcp_world_size=pcp_size,
1860	enable_caching=True,
1861	)
1862	assert manager.coordinator.group_block_sizes == (64, 16)
1863	request = make_request("replicated-draft", [1] * 65, block_size=16, hash_fn=sha256)
1864	blocks = manager.allocate_slots(request, 65)
1865	assert blocks is not None
1866	assert [len(group) for group in blocks.blocks] == [2, 5]
1867	manager.cache_blocks(request, 64)
1868	cached_request = make_request(
1869	"cached-draft", [1] * 65, block_size=16, hash_fn=sha256
1870	)
1871	cached_blocks, num_cached, _ = manager.get_computed_blocks(cached_request)
1872	assert num_cached == 64
1873	assert [len(group) for group in cached_blocks.blocks] == [1, 4]
1874
1875
1876	@pytest.mark.parametrize("use_mla", [False, True])
1877	def test_full_attention_merge_preserves_replicated_cache_geometry(use_mla):
1878	config = VllmConfig(model_config=ModelConfig(max_model_len=1024))
1879	config.parallel_config.decode_context_parallel_size = 4
1880	draft = replace(
1881	new_mla_spec() if use_mla else new_kv_cache_spec(),
1882	block_size=16,
1883	dcp_sharded=False,
1884	)
1885	merged = type(draft).merge([draft, draft])
1886	wrapped = UniformTypeKVCacheSpecs.from_specs({"draft": merged})
1887	assert wrapped is not None and not wrapped.dcp_sharded
1888	assert merged.max_num_blocks_per_req(config, 1024) == 64
1889	assert merged.max_memory_usage_bytes(config) == 64 * draft.page_size_bytes
1890	assert kv_cache_utils.resolve_dcp_kv_block_size(merged, 4) == 16
1891	with pytest.raises(AssertionError):
1892	type(draft).merge([draft, replace(draft, dcp_sharded=True)])
1893
1894
1895	@pytest.mark.parametrize("with_draft", [False, True])
1896	@pytest.mark.parametrize("indexer_alignment", [None, 512])
1897	def test_sparse_mla_preserves_physical_row_addressing(with_draft, indexer_alignment):
1898	from vllm.v1.attention.backends.mla.sparse_utils import flat_kv_row_view
1899	from vllm.v1.worker.utils import allocate_kv_cache
1900
1901	config = VllmConfig(model_config=ModelConfig(max_model_len=1024))
1902	config.cache_config.kv_cache_layout = "BLHNC"
1903	config.parallel_config.decode_context_parallel_size = 4
1904	common = dict(block_size=64, num_kv_heads=1, dtype=torch.uint8)
1905	specs = {
1906	"target": MLAAttentionSpec(
1907	**common,
1908	head_size=576,
1909	state_content_bytes=656,
1910	block_stride_alignment=656,
1911	cache_dtype_str="fp8_ds_mla",
1912	is_index_group_leader=True,
1913	),
1914	"indexer": MLAAttentionSpec(
1915	**common,
1916	head_size=132,
1917	cache_role=SparseCacheRole.INDEXER,
1918	block_stride_alignment=indexer_alignment,
1919	),
1920	"draft": SlidingWindowSpec(
1921	block_size=64,
1922	num_kv_heads=64,
1923	head_size=64,
1924	dtype=torch.bfloat16,
1925	sliding_window=2048,
1926	dcp_sharded=False,
1927	),
1928	}
1929	if not with_draft:
1930	del specs["draft"]
1931	groups = get_kv_cache_groups(config, specs)
1932	cache_config = kv_cache_utils.get_kv_cache_config_from_groups(
1933	config, groups, available_memory=8 * 1024 * 1024
1934	)
1935	caches = allocate_kv_cache(
1936	cache_config, torch.device("cpu"), KVCacheLayout.BLHNC, [64] * len(groups)
1937	)
1938	cache = caches["target"].squeeze(1)
1939	if indexer_alignment is not None:
1940	assert caches["indexer"].stride(0) % indexer_alignment == 0
1941	rows, stride_rows = flat_kv_row_view(cache, 64)
1942	cache[1, 0].fill_(7)
1943	torch.testing.assert_close(rows[stride_rows], cache[1, 0])
1944	assert torch.all(rows[stride_rows] == 7)
1945
1946
1947	@pytest.mark.parametrize(
1948	"layer_type,dcp_size,expected_width",
1949	[
1950	("mla", 1, 64),
1951	("mla", 2, 32),
1952	# Mamba state is replicated, not DCP-sharded, and its width is the
1953	# resident state block count rather than cdiv(max_len, block_size).
1954	("mamba", 2, 3),
1955	],
1956	)
1957	def test_uniform_type_spec_block_table_width_matches_layer_spec(
1958	layer_type, dcp_size, expected_width
1959	):
1960	# The runner sizes the block table from the group spec while the metadata
1961	# builders are constructed from the per-layer spec, so the aggregate must
1962	# report the same width as the layers it wraps.
1963	vllm_config = VllmConfig(model_config=ModelConfig(max_model_len=1024))
1964	vllm_config.parallel_config.decode_context_parallel_size = dcp_size
1965	layer_spec = new_mla_spec() if layer_type == "mla" else new_mamba_spec()
1966	uniform_spec = UniformTypeKVCacheSpecs(
1967	block_size=layer_spec.block_size,
1968	kv_cache_specs={"layer1": layer_spec, "layer2": layer_spec},
1969	)
1970
1971	assert layer_spec.max_num_blocks_per_req(vllm_config, 1024) == expected_width
1972	assert uniform_spec.max_num_blocks_per_req(vllm_config, 1024) == expected_width
1973
1974
1975	def test_merge_kv_cache_spec():
1976	same_layer_specs = [
1977	new_kv_cache_spec(num_kv_heads=32),
1978	new_kv_cache_spec(num_kv_heads=32),
1979	]
1980	merged_layer_spec = same_layer_specs[0].merge(same_layer_specs)
1981	assert merged_layer_spec.block_size == 16
1982	assert merged_layer_spec.num_kv_heads == 32
1983	assert merged_layer_spec.head_size == 64
1984	assert merged_layer_spec.dtype == torch.float32
1985	assert merged_layer_spec.sliding_window is None
1986
1987	different_layer_specs = [
1988	new_kv_cache_spec(num_kv_heads=32),
1989	new_kv_cache_spec(num_kv_heads=16),
1990	]
1991	with pytest.raises(AssertionError):
1992	different_layer_specs[0].merge(different_layer_specs)
1993
1994	full_spec = new_kv_cache_spec(num_kv_heads=32)
1995	different_type_layer_specs = [
1996	full_spec,
1997	SlidingWindowSpec(
1998	block_size=full_spec.block_size,
1999	num_kv_heads=full_spec.num_kv_heads,
2000	head_size=full_spec.head_size,
2001	dtype=full_spec.dtype,
2002	sliding_window=1,
2003	),
2004	]
2005	with pytest.raises(AssertionError):
2006	different_type_layer_specs[0].merge(different_type_layer_specs)
2007	with pytest.raises(AssertionError):
2008	different_type_layer_specs[1].merge(different_type_layer_specs)
2009
2010	different_sliding_window_layer_specs = [
2011	new_kv_cache_spec(num_kv_heads=32),
2012	new_kv_cache_spec(num_kv_heads=32, sliding_window=1),
2013	new_kv_cache_spec(num_kv_heads=32, sliding_window=2),
2014	]
2015	with pytest.raises(ValueError):
2016	different_sliding_window_layer_specs[0].merge(
2017	different_sliding_window_layer_specs
2018	)
2019
2020	same_sliding_window_layer_specs = [
2021	new_kv_cache_spec(num_kv_heads=32, sliding_window=1),
2022	new_kv_cache_spec(num_kv_heads=32, sliding_window=1),
2023	]
2024	merged_layer_spec = same_sliding_window_layer_specs[0].merge(
2025	same_sliding_window_layer_specs
2026	)
2027	assert merged_layer_spec.sliding_window == 1
2028
2029	same_sliding_window_layer_spec_with_none = [
2030	new_kv_cache_spec(num_kv_heads=32, sliding_window=1),
2031	new_kv_cache_spec(num_kv_heads=32, sliding_window=None),
2032	]
2033	merged_layer_spec = same_sliding_window_layer_spec_with_none[0].merge(
2034	same_sliding_window_layer_spec_with_none
2035	)
2036	assert merged_layer_spec.sliding_window == 1
2037
2038
2039	def test_is_kv_cache_spec_uniform():
2040	kv_cache_spec = {
2041	"layer_1": new_kv_cache_spec(num_kv_heads=32),
2042	"layer_2": new_kv_cache_spec(num_kv_heads=32),
2043	}
2044	assert is_kv_cache_spec_uniform(kv_cache_spec)
2045
2046	kv_cache_spec = {
2047	"layer_1": new_kv_cache_spec(num_kv_heads=32),
2048	"layer_2": new_kv_cache_spec(num_kv_heads=32, sliding_window=1),
2049	}
2050	assert is_kv_cache_spec_uniform(kv_cache_spec)
2051
2052	kv_cache_spec = {
2053	"layer_1": new_kv_cache_spec(num_kv_heads=32),
2054	"layer_2": new_sliding_window_spec(num_kv_heads=32, sliding_window=1),
2055	}
2056	assert not is_kv_cache_spec_uniform(kv_cache_spec)
2057
2058	kv_cache_spec = {
2059	"layer_1": new_sliding_window_spec(num_kv_heads=32, sliding_window=1),
2060	"layer_2": new_sliding_window_spec(num_kv_heads=32, sliding_window=1),
2061	}
2062	assert is_kv_cache_spec_uniform(kv_cache_spec)
2063
2064	kv_cache_spec = {
2065	"layer_1": new_sliding_window_spec(num_kv_heads=32, sliding_window=1),
2066	"layer_2": new_sliding_window_spec(num_kv_heads=32, sliding_window=2),
2067	}
2068	assert not is_kv_cache_spec_uniform(kv_cache_spec)
2069
2070	script = """
2071	import sys
2072
2073	from vllm.v1.core.kv_cache_utils import is_kv_cache_spec_uniform
2074	from vllm.v1.kv_cache_interface import KVCacheSpec
2075
2076	if sys.flags.optimize < 1:
2077	raise RuntimeError("subprocess is not running with optimization enabled")
2078	specs = {
2079	"a": KVCacheSpec(block_size=1),
2080	"b": KVCacheSpec(block_size=2),
2081	}
2082	if is_kv_cache_spec_uniform(specs):
2083	raise RuntimeError("different specs were treated as uniform")
2084	"""
2085	subprocess.run(
2086	[sys.executable, "-O", "-c", script],
2087	cwd=Path(__file__).resolve().parents[3],
2088	check=True,
2089	timeout=60,
2090	)
2091
2092
2093	@pytest.mark.parametrize(
2094	("model_id", "max_model_len", "want_estimated_max_len"),
2095	[
2096	("Qwen/Qwen1.5-7B", 16385, 16384),
2097	("Qwen/Qwen1.5-7B", 16383, 16383),
2098	],
2099	)
2100	def test_estimate_max_model_len(model_id, max_model_len, want_estimated_max_len):
2101	# Create a VllmConfig
2102	model_config = ModelConfig(
2103	model_id,
2104	runner="generate",
2105	dtype="float16",
2106	max_model_len=max_model_len,
2107	)
2108	scheduler_config = SchedulerConfig(
2109	max_num_batched_tokens=32768,
2110	max_model_len=model_config.max_model_len,
2111	is_encoder_decoder=model_config.is_encoder_decoder,
2112	)
2113
2114	vllm_config = VllmConfig(
2115	model_config=model_config,
2116	scheduler_config=scheduler_config,
2117	)
2118
2119	# Create KV cache specs
2120	kv_cache_spec = {}
2121	for i in range(32):
2122	layer_name = f"layer_{i}"
2123	kv_cache_spec[layer_name] = FullAttentionSpec(
2124	block_size=16,
2125	num_kv_heads=32,
2126	head_size=128,
2127	dtype=torch.float16,
2128	)
2129	# Estimate the maximum model length, 16384 model_len need 8GB
2130	estimated_max_len = estimate_max_model_len(
2131	vllm_config, kv_cache_spec, 8 * GiB_bytes
2132	)
2133	assert estimated_max_len == want_estimated_max_len
2134
2135
2136	def test_get_max_concurrency_for_kv_cache_config():
2137	# Create a VllmConfig
2138	model_id = "Qwen/Qwen1.5-7B"
2139	max_model_len = 16384
2140	model_config = ModelConfig(
2141	model_id,
2142	runner="generate",
2143	dtype="float16",
2144	max_model_len=max_model_len,
2145	)
2146	scheduler_config = SchedulerConfig(
2147	max_num_batched_tokens=1024,
2148	enable_chunked_prefill=True,
2149	max_model_len=model_config.max_model_len,
2150	is_encoder_decoder=model_config.is_encoder_decoder,
2151	# Pin to sync: SWA per-request bounds grow with overlapping batches.
2152	async_scheduling=False,
2153	)
2154
2155	vllm_config = VllmConfig(
2156	model_config=model_config,
2157	scheduler_config=scheduler_config,
2158	)
2159	assert vllm_config.max_concurrent_batches == 1
2160
2161	full_attention_spec = FullAttentionSpec(
2162	block_size=16,
2163	num_kv_heads=32,
2164	head_size=128,
2165	dtype=torch.float16,
2166	)
2167
2168	sliding_window_spec = SlidingWindowSpec(
2169	block_size=16,
2170	num_kv_heads=32,
2171	head_size=128,
2172	dtype=torch.float16,
2173	sliding_window=1024,
2174	)
2175
2176	kv_cache_config_full_attention = KVCacheConfig(
2177	num_blocks=int(1024 * 1.5),
2178	kv_cache_tensors=[],
2179	kv_cache_groups=[
2180	KVCacheGroupSpec([f"layer_{i}" for i in range(32)], full_attention_spec),
2181	],
2182	)
2183	max_concurrency_full_attention = get_max_concurrency_for_kv_cache_config(
2184	vllm_config, kv_cache_config_full_attention
2185	)
2186	assert max_concurrency_full_attention == 1.5
2187
2188	kv_cache_config_sliding_window = KVCacheConfig(
2189	num_blocks=129 * 3,
2190	kv_cache_tensors=[],
2191	kv_cache_groups=[
2192	KVCacheGroupSpec([f"layer_{i}" for i in range(32)], sliding_window_spec),
2193	],
2194	)
2195	max_concurrency_sliding_window = get_max_concurrency_for_kv_cache_config(
2196	vllm_config, kv_cache_config_sliding_window
2197	)
2198	assert max_concurrency_sliding_window == 3
2199
2200	kv_cache_config_hybrid_model = KVCacheConfig(
2201	num_blocks=(1024 + 129) * 3,
2202	kv_cache_tensors=[],
2203	kv_cache_groups=[
2204	KVCacheGroupSpec([f"layer_{i}" for i in range(32)], full_attention_spec),
2205	KVCacheGroupSpec(
2206	[f"layer_{i}" for i in range(32, 64)], sliding_window_spec
2207	),
2208	],
2209	)
2210	max_concurrency_hybrid_model = get_max_concurrency_for_kv_cache_config(
2211	vllm_config, kv_cache_config_hybrid_model
2212	)
2213	assert max_concurrency_hybrid_model == 3
2214	num_tokens, max_concurrency = get_kv_cache_capacity(
2215	vllm_config, kv_cache_config_hybrid_model
2216	)
2217	assert num_tokens == max_concurrency_hybrid_model * max_model_len
2218	assert max_concurrency == max_concurrency_hybrid_model
2219
2220	# Unequal group sizes in the standard layout: each group's pages cost
2221	# whole pool blocks, so a request needs 1024 + 129 = 1153 blocks — the
2222	# same as the equal-hybrid case above, regardless of the second group
2223	# holding only 2 layers.
2224	kv_cache_config_unequal_groups = KVCacheConfig(
2225	num_blocks=1153 * 3,
2226	kv_cache_tensors=[],
2227	kv_cache_groups=[
2228	KVCacheGroupSpec([f"layer_{i}" for i in range(32)], full_attention_spec),
2229	KVCacheGroupSpec(["layer_32", "layer_33"], sliding_window_spec),
2230	],
2231	)
2232	assert (
2233	get_max_concurrency_for_kv_cache_config(
2234	vllm_config, kv_cache_config_unequal_groups
2235	)
2236	== 3
2237	)
2238
2239	# UniformTypeKVCacheSpecs group (worker config shape): the aggregated
2240	# spec's memory/page ratio equals a single layer's page count, so the
2241	# group needs 1024 blocks and the request 1153 in total. The previous
2242	# formula normalized both groups' memory by the first group's page size,
2243	# reporting 3459/1057 = 3.27 here instead of 3 — and a different value
2244	# again for the scheduler-config shape below.
2245	uniform_full_spec = UniformTypeKVCacheSpecs(
2246	block_size=full_attention_spec.block_size,
2247	kv_cache_specs={f"layer_{i}": full_attention_spec for i in range(4)},
2248	)
2249	kv_cache_config_uniform_group = KVCacheConfig(
2250	num_blocks=1153 * 3,
2251	kv_cache_tensors=[],
2252	kv_cache_groups=[
2253	KVCacheGroupSpec([f"layer_{i}" for i in range(4)], uniform_full_spec),
2254	KVCacheGroupSpec(["layer_4", "layer_5"], sliding_window_spec),
2255	],
2256	)
2257	assert (
2258	get_max_concurrency_for_kv_cache_config(
2259	vllm_config, kv_cache_config_uniform_group
2260	)
2261	== 3
2262	)
2263
2264	# Scheduler-config shape: generate_scheduler_kv_cache_config replaces the
2265	# uniform-type group's spec with a representative per-layer spec.
2266	# Capacity must not change between the two shapes (the engine computes
2267	# on the scheduler config, the worker loop on the worker config).
2268	kv_cache_config_scheduler_shape = generate_scheduler_kv_cache_config(
2269	[copy.deepcopy(kv_cache_config_uniform_group)]
2270	)
2271	assert get_max_concurrency_for_kv_cache_config(
2272	vllm_config, kv_cache_config_scheduler_shape
2273	) == get_max_concurrency_for_kv_cache_config(
2274	vllm_config, kv_cache_config_uniform_group
2275	)
2276
2277
2278	def test_allocate_with_lookahead():
2279	"""Verify that lookahead tokens correctly affect block allocation."""
2280	block_size = 4
2281	config = KVCacheConfig(
2282	num_blocks=10,
2283	kv_cache_tensors=[
2284	KVCacheTensor(
2285	size=100,
2286	layers=["layer1"],
2287	layer_stride=100,
2288	block_stride=10,
2289	),
2290	],
2291	kv_cache_groups=[
2292	KVCacheGroupSpec(["layer1"], new_kv_cache_spec(block_size=block_size)),
2293	],
2294	)
2295
2296	request = make_request(
2297	request_id="0",
2298	prompt_token_ids=[],
2299	block_size=block_size,
2300	mm_positions=None,
2301	mm_hashes=None,
2302	)
2303
2304	# Test case 1: Requires additional lookahead tokens
2305	kv_cache_manager = KVCacheManager(
2306	kv_cache_config=config,
2307	max_model_len=100,
2308	scheduler_block_size=block_size,
2309	hash_block_size=block_size,
2310	)
2311	blocks = kv_cache_manager.allocate_slots(
2312	request,
2313	num_new_tokens=3,
2314	num_lookahead_tokens=2, # Total required: 3+2=5 tokens
2315	)
2316	assert len(blocks.get_block_ids()[0]) == 2 # ceil(5/4)=2 blocks
2317
2318	# Test case 2: With precomputed blocks
2319	kv_cache_manager = KVCacheManager(
2320	kv_cache_config=config,
2321	max_model_len=100,
2322	scheduler_block_size=block_size,
2323	hash_block_size=block_size,
2324	)
2325	# required_blocks = ceil((3 + 2) /4) = 2
2326	blocks = kv_cache_manager.allocate_slots(
2327	request,
2328	num_new_tokens=3,
2329	num_lookahead_tokens=2,
2330	)
2331	assert len(blocks.get_block_ids()[0]) == 2
2332
2333	# Test case 3: With precomputed blocks
2334	# required_blocks = ceil((3 + 4) / 4) = 2
2335	kv_cache_manager = KVCacheManager(
2336	kv_cache_config=config,
2337	max_model_len=100,
2338	scheduler_block_size=block_size,
2339	hash_block_size=block_size,
2340	)
2341	blocks = kv_cache_manager.allocate_slots(
2342	request,
2343	num_new_tokens=3,
2344	num_lookahead_tokens=4,
2345	)
2346	assert len(blocks.get_block_ids()[0]) == 2
2347
2348
2349	def test_get_kv_cache_config_one_worker():
2350	# pass max_model_len to pass check_enough_kv_cache_memory
2351	model_config = ModelConfig(max_model_len=16)
2352	vllm_config = VllmConfig(model_config=model_config)
2353	vllm_config.cache_config.kv_cache_layout = "LBNHC"
2354	vllm_config.cache_config.prefix_cache_retention_interval = None
2355
2356	mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2
2357	# all layers are full attention -> single group
2358	kv_cache_specs_full = {
2359	"layer_1": new_kv_cache_spec(),
2360	"layer_2": new_kv_cache_spec(),
2361	}
2362	kv_cache_config_full = get_kv_cache_configs(
2363	vllm_config, [kv_cache_specs_full], [mem_per_block_per_layer * 2 * 32]
2364	)[0]
2365	print(kv_cache_config_full)
2366	assert kv_cache_config_full == KVCacheConfig(
2367	num_blocks=32,
2368	kv_cache_tensors=[
2369	KVCacheTensor(
2370	size=mem_per_block_per_layer * 32 * 2,
2371	layers=["layer_1", "layer_2"],
2372	layer_stride=mem_per_block_per_layer * 32,
2373	block_stride=mem_per_block_per_layer,
2374	),
2375	],
2376	kv_cache_groups=[KVCacheGroupSpec(["layer_1", "layer_2"], new_kv_cache_spec())],
2377	)
2378
2379	# all layers are sliding window -> single group
2380	kv_cache_specs_sliding = {
2381	"layer_1": new_sliding_window_spec(),
2382	"layer_2": new_sliding_window_spec(),
2383	}
2384	kv_cache_config_sliding = get_kv_cache_configs(
2385	vllm_config, [kv_cache_specs_sliding], [mem_per_block_per_layer * 2 * 32]
2386	)[0]
2387	assert kv_cache_config_sliding == KVCacheConfig(
2388	num_blocks=32,
2389	kv_cache_tensors=[
2390	KVCacheTensor(
2391	size=mem_per_block_per_layer * 32 * 2,
2392	layers=["layer_1", "layer_2"],
2393	layer_stride=mem_per_block_per_layer * 32,
2394	block_stride=mem_per_block_per_layer,
2395	),
2396	],
2397	kv_cache_groups=[
2398	KVCacheGroupSpec(["layer_1", "layer_2"], new_sliding_window_spec())
2399	],
2400	)
2401
2402	# full + sliding, but disable_hybrid_kv_cache_manager
2403	vllm_config.scheduler_config.disable_hybrid_kv_cache_manager = True
2404	kv_cache_specs_hybrid = {
2405	"layer_1": new_kv_cache_spec(),
2406	"layer_2": new_sliding_window_spec(),
2407	}
2408	kv_cache_config_hybrid = get_kv_cache_configs(
2409	vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 2 * 32]
2410	)[0]
2411	assert kv_cache_config_hybrid == KVCacheConfig(
2412	num_blocks=32,
2413	kv_cache_tensors=[
2414	KVCacheTensor(
2415	size=mem_per_block_per_layer * 32 * 2,
2416	layers=["layer_1", "layer_2"],
2417	layer_stride=mem_per_block_per_layer * 32,
2418	block_stride=mem_per_block_per_layer,
2419	),
2420	],
2421	kv_cache_groups=[
2422	KVCacheGroupSpec(
2423	["layer_1", "layer_2"], new_kv_cache_spec(sliding_window=1)
2424	),
2425	],
2426	)
2427	vllm_config.scheduler_config.disable_hybrid_kv_cache_manager = False
2428
2429	# full + sliding, with hybrid_kv_cache_manager
2430	kv_cache_specs_hybrid = {
2431	"layer_1": new_kv_cache_spec(),
2432	"layer_2": new_sliding_window_spec(),
2433	}
2434	kv_cache_config_hybrid = get_kv_cache_configs(
2435	vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 2 * 32]
2436	)[0]
2437	assert kv_cache_config_hybrid == KVCacheConfig(
2438	num_blocks=64,
2439	kv_cache_tensors=[
2440	KVCacheTensor(
2441	size=mem_per_block_per_layer * 64,
2442	layers=["layer_1"],
2443	layer_stride=mem_per_block_per_layer * 64,
2444	block_stride=mem_per_block_per_layer,
2445	),
2446	KVCacheTensor(
2447	size=mem_per_block_per_layer * 64,
2448	layers=["layer_2"],
2449	layer_stride=mem_per_block_per_layer * 64,
2450	block_stride=mem_per_block_per_layer,
2451	),
2452	],
2453	kv_cache_groups=[
2454	KVCacheGroupSpec(["layer_1"], new_kv_cache_spec()),
2455	KVCacheGroupSpec(["layer_2"], new_sliding_window_spec()),
2456	],
2457	)
2458
2459	# 2 full + 4 sliding, 2 layers per group
2460	kv_cache_specs_hybrid = {
2461	"layer_1": new_kv_cache_spec(),
2462	"layer_2": new_kv_cache_spec(),
2463	"layer_3": new_sliding_window_spec(),
2464	"layer_4": new_sliding_window_spec(),
2465	"layer_5": new_sliding_window_spec(),
2466	"layer_6": new_sliding_window_spec(),
2467	}
2468	kv_cache_config_hybrid = get_kv_cache_configs(
2469	vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 2 * 32]
2470	)[0]
2471	assert kv_cache_config_hybrid == KVCacheConfig(
2472	num_blocks=32,
2473	kv_cache_tensors=[
2474	KVCacheTensor(
2475	size=mem_per_block_per_layer * 32 * 2,
2476	layers=["layer_1", "layer_2"],
2477	layer_stride=mem_per_block_per_layer * 32,
2478	block_stride=mem_per_block_per_layer,
2479	),
2480	KVCacheTensor(
2481	size=mem_per_block_per_layer * 32 * 2,
2482	layers=["layer_3", "layer_5"],
2483	layer_stride=mem_per_block_per_layer * 32,
2484	block_stride=mem_per_block_per_layer,
2485	),
2486	KVCacheTensor(
2487	size=mem_per_block_per_layer * 32 * 2,
2488	layers=["layer_4", "layer_6"],
2489	layer_stride=mem_per_block_per_layer * 32,
2490	block_stride=mem_per_block_per_layer,
2491	),
2492	],
2493	kv_cache_groups=[
2494	KVCacheGroupSpec(["layer_1", "layer_2"], new_kv_cache_spec()),
2495	KVCacheGroupSpec(["layer_3", "layer_5"], new_sliding_window_spec()),
2496	KVCacheGroupSpec(["layer_4", "layer_6"], new_sliding_window_spec()),
2497	],
2498	)
2499
2500	# 3 full + 7 sliding, pad to 3 full + 9 sliding
2501	kv_cache_specs_hybrid = {
2502	"layer_1": new_kv_cache_spec(),
2503	"layer_2": new_kv_cache_spec(),
2504	"layer_3": new_kv_cache_spec(),
2505	"layer_4": new_sliding_window_spec(),
2506	"layer_5": new_sliding_window_spec(),
2507	"layer_6": new_sliding_window_spec(),
2508	"layer_7": new_sliding_window_spec(),
2509	"layer_8": new_sliding_window_spec(),
2510	"layer_9": new_sliding_window_spec(),
2511	"layer_10": new_sliding_window_spec(),
2512	}
2513	kv_cache_config_hybrid = get_kv_cache_configs(
2514	vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 3 * 32]
2515	)[0]
2516	assert kv_cache_config_hybrid == KVCacheConfig(
2517	num_blocks=32,
2518	kv_cache_tensors=[
2519	KVCacheTensor(
2520	size=mem_per_block_per_layer * 32 * 3,
2521	layers=["layer_1", "layer_2", "layer_3"],
2522	layer_stride=mem_per_block_per_layer * 32,
2523	block_stride=mem_per_block_per_layer,
2524	),
2525	KVCacheTensor(
2526	size=mem_per_block_per_layer * 32 * 3,
2527	layers=["layer_4", "layer_7", "layer_10"],
2528	layer_stride=mem_per_block_per_layer * 32,
2529	block_stride=mem_per_block_per_layer,
2530	),
2531	KVCacheTensor(
2532	size=mem_per_block_per_layer * 32 * 3,
2533	layers=["layer_5", "layer_8"],
2534	layer_stride=mem_per_block_per_layer * 32,
2535	block_stride=mem_per_block_per_layer,
2536	),
2537	KVCacheTensor(
2538	size=mem_per_block_per_layer * 32 * 3,
2539	layers=["layer_6", "layer_9"],
2540	layer_stride=mem_per_block_per_layer * 32,
2541	block_stride=mem_per_block_per_layer,
2542	),
2543	],
2544	kv_cache_groups=[
2545	KVCacheGroupSpec(["layer_1", "layer_2", "layer_3"], new_kv_cache_spec()),
2546	KVCacheGroupSpec(
2547	["layer_4", "layer_7", "layer_10"], new_sliding_window_spec()
2548	),
2549	KVCacheGroupSpec(["layer_5", "layer_8"], new_sliding_window_spec()),
2550	KVCacheGroupSpec(["layer_6", "layer_9"], new_sliding_window_spec()),
2551	],
2552	)
2553
2554	# 6 full + 5 sliding, pad to 6 full + 6 sliding. This is a typical case for gpt-oss
2555	# eagle where there is only one more full attention layer than sliding window layers
2556	kv_cache_specs_hybrid = {
2557	"layer_1": new_kv_cache_spec(),
2558	"layer_2": new_kv_cache_spec(),
2559	"layer_3": new_kv_cache_spec(),
2560	"layer_4": new_kv_cache_spec(),
2561	"layer_5": new_kv_cache_spec(),
2562	"layer_6": new_kv_cache_spec(),
2563	"layer_7": new_sliding_window_spec(),
2564	"layer_8": new_sliding_window_spec(),
2565	"layer_9": new_sliding_window_spec(),
2566	"layer_10": new_sliding_window_spec(),
2567	"layer_11": new_sliding_window_spec(),
2568	}
2569
2570	kv_cache_config_hybrid = get_kv_cache_configs(
2571	vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 6 * 32]
2572	)[0]
2573	print(kv_cache_config_hybrid)
2574	assert kv_cache_config_hybrid == KVCacheConfig(
2575	num_blocks=32,
2576	kv_cache_tensors=[
2577	KVCacheTensor(
2578	size=mem_per_block_per_layer * 32 * 6,
2579	layers=[
2580	"layer_1",
2581	"layer_2",
2582	"layer_3",
2583	"layer_4",
2584	"layer_5",
2585	"layer_6",
2586	],
2587	layer_stride=mem_per_block_per_layer * 32,
2588	block_stride=mem_per_block_per_layer,
2589	),
2590	KVCacheTensor(
2591	size=mem_per_block_per_layer * 32 * 6,
2592	layers=["layer_7", "layer_8", "layer_9", "layer_10", "layer_11"],
2593	layer_stride=mem_per_block_per_layer * 32,
2594	block_stride=mem_per_block_per_layer,
2595	),
2596	],
2597	kv_cache_groups=[
2598	KVCacheGroupSpec(
2599	["layer_1", "layer_2", "layer_3", "layer_4", "layer_5", "layer_6"],
2600	new_kv_cache_spec(),
2601	),
2602	KVCacheGroupSpec(
2603	["layer_7", "layer_8", "layer_9", "layer_10", "layer_11"],
2604	new_sliding_window_spec(),
2605	),
2606	],
2607	)
2608
2609	# different hidden size but same type, use UniformTypeKVCacheSpecs
2610	kv_cache_specs_hybrid = {
2611	"layer_1": new_kv_cache_spec(head_size=128),
2612	"layer_2": new_kv_cache_spec(head_size=64),
2613	}
2614	kv_cache_config_hybrid = get_kv_cache_configs(
2615	vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 3 * 32]
2616	)[0]
2617	# Layers of different page sizes pack densely into one allocation: the
2618	# 2x-larger layer_1 takes the first region, layer_2 the next.
2619	assert kv_cache_config_hybrid == KVCacheConfig(
2620	num_blocks=32,
2621	kv_cache_tensors=[
2622	KVCacheTensor(
2623	size=mem_per_block_per_layer * 32 * 3,
2624	layers=["layer_1"],
2625	layer_stride=mem_per_block_per_layer * 2 * 32,
2626	block_stride=mem_per_block_per_layer * 2,
2627	),
2628	KVCacheTensor(
2629	size=mem_per_block_per_layer * 32 * 3,
2630	layers=["layer_2"],
2631	layer_stride=mem_per_block_per_layer * 32,
2632	block_stride=mem_per_block_per_layer,
2633	offset=mem_per_block_per_layer * 2 * 32,
2634	),
2635	],
2636	kv_cache_groups=[
2637	KVCacheGroupSpec(
2638	["layer_1", "layer_2"],
2639	UniformTypeKVCacheSpecs(
2640	block_size=16, kv_cache_specs=kv_cache_specs_hybrid
2641	),
2642	)
2643	],
2644	)
2645
2646	# Different hidden size and different type, align by different block size
2647	kv_cache_specs_hybrid = {
2648	"layer_1": new_kv_cache_spec(head_size=64),
2649	"layer_2": new_sliding_window_spec(head_size=32),
2650	}
2651	kv_cache_config_hybrid = get_kv_cache_configs(
2652	vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 32]
2653	)[0]
2654	assert kv_cache_config_hybrid == KVCacheConfig(
2655	num_blocks=32,
2656	kv_cache_tensors=[
2657	KVCacheTensor(
2658	size=mem_per_block_per_layer * 32,
2659	layers=["layer_1"],
2660	layer_stride=mem_per_block_per_layer * 32,
2661	block_stride=mem_per_block_per_layer,
2662	),
2663	KVCacheTensor(
2664	size=mem_per_block_per_layer * 32,
2665	layers=["layer_2"],
2666	layer_stride=mem_per_block_per_layer * 32,
2667	block_stride=mem_per_block_per_layer,
2668	),
2669	],
2670	kv_cache_groups=[
2671	KVCacheGroupSpec(["layer_1"], new_kv_cache_spec(head_size=64)),
2672	KVCacheGroupSpec(
2673	["layer_2"], new_sliding_window_spec(head_size=32, block_size=32)
2674	),
2675	],
2676	)
2677
2678	# different hidden size that cannot be aligned by using different block size,
2679	# but can be aligned by padding the smaller physical page.
2680	swa_spec = new_sliding_window_spec(head_size=96)
2681	kv_cache_specs_hybrid = {
2682	"layer_1": new_kv_cache_spec(head_size=64),
2683	"layer_2": swa_spec,
2684	}
2685
2686	kv_cache_config_hybrid = get_kv_cache_configs(
2687	vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 2 * 32]
2688	)[0]
2689	padded_page_size = swa_spec.page_size_bytes
2690	assert kv_cache_config_hybrid == KVCacheConfig(
2691	num_blocks=42,
2692	kv_cache_tensors=[
2693	KVCacheTensor(
2694	size=padded_page_size * 42,
2695	layers=["layer_1"],
2696	layer_stride=padded_page_size * 42,
2697	block_stride=padded_page_size,
2698	),
2699	KVCacheTensor(
2700	size=padded_page_size * 42,
2701	layers=["layer_2"],
2702	layer_stride=padded_page_size * 42,
2703	block_stride=padded_page_size,
2704	),
2705	],
2706	kv_cache_groups=[
2707	KVCacheGroupSpec(
2708	["layer_1"],
2709	new_kv_cache_spec(
2710	head_size=64,
2711	page_size_padded=padded_page_size,
2712	),
2713	),
2714	KVCacheGroupSpec(
2715	["layer_2"],
2716	new_sliding_window_spec(head_size=96),
2717	),
2718	],
2719	)
2720
2721	# Test num_gpu_blocks_override
2722	vllm_config.cache_config.num_gpu_blocks_override = 16
2723	kv_cache_config_override_blocks = get_kv_cache_configs(
2724	vllm_config, [kv_cache_specs_full], [mem_per_block_per_layer * 2 * 32]
2725	)[0]
2726	assert kv_cache_config_override_blocks == KVCacheConfig(
2727	num_blocks=16,
2728	kv_cache_tensors=[
2729	KVCacheTensor(
2730	size=mem_per_block_per_layer * 16 * 2,
2731	layers=["layer_1", "layer_2"],
2732	layer_stride=mem_per_block_per_layer * 16,
2733	block_stride=mem_per_block_per_layer,
2734	),
2735	],
2736	kv_cache_groups=[KVCacheGroupSpec(["layer_1", "layer_2"], new_kv_cache_spec())],
2737	)
2738
2739
2740	def test_get_kv_cache_configs_attention_free():
2741	kv_cache_specs: dict[str, KVCacheSpec] = {}
2742	vllm_config = VllmConfig(model_config=ModelConfig(max_model_len=16))
2743	vllm_config.cache_config.prefix_cache_retention_interval = None
2744	kv_cache_configs = get_kv_cache_configs(vllm_config, [kv_cache_specs], [0])
2745	assert kv_cache_configs == [
2746	KVCacheConfig(
2747	num_blocks=1,
2748	kv_cache_tensors=[],
2749	kv_cache_groups=[],
2750	)
2751	]
2752
2753
2754	def test_generate_uniform_type_kv_cache_specs():
2755	# All layers are full attention, can be merged
2756	kv_cache_specs = {
2757	"layer_1": new_kv_cache_spec(),
2758	"layer_2": new_kv_cache_spec(head_size=128),
2759	}
2760	uniform_spec = UniformTypeKVCacheSpecs.from_specs(kv_cache_specs)
2761	assert uniform_spec == UniformTypeKVCacheSpecs(
2762	block_size=16, kv_cache_specs=kv_cache_specs
2763	)
2764
2765	# Full attention + sliding window, cannot be merged
2766	kv_cache_specs = {
2767	"layer_1": new_kv_cache_spec(),
2768	"layer_2": new_sliding_window_spec(sliding_window=1),
2769	}
2770	uniform_spec = UniformTypeKVCacheSpecs.from_specs(kv_cache_specs)
2771	assert uniform_spec is None
2772
2773	# different order of full attention + sliding window, cannot be merged
2774	kv_cache_specs = {
2775	"layer_1": new_sliding_window_spec(sliding_window=1),
2776	"layer_2": new_kv_cache_spec(),
2777	}
2778	uniform_spec = UniformTypeKVCacheSpecs.from_specs(kv_cache_specs)
2779	assert uniform_spec is None
2780
2781	# Same-size sliding window, can be merged
2782	kv_cache_specs = {
2783	"layer_1": new_sliding_window_spec(sliding_window=1),
2784	"layer_2": new_sliding_window_spec(sliding_window=1, head_size=128),
2785	}
2786	uniform_spec = UniformTypeKVCacheSpecs.from_specs(kv_cache_specs)
2787	assert uniform_spec == UniformTypeKVCacheSpecs(
2788	block_size=16, kv_cache_specs=kv_cache_specs
2789	)
2790
2791	# different block sizes, cannot be merged
2792	kv_cache_specs = {
2793	"layer_1": new_kv_cache_spec(block_size=16),
2794	"layer_2": new_kv_cache_spec(block_size=32),
2795	}
2796	uniform_spec = UniformTypeKVCacheSpecs.from_specs(kv_cache_specs)
2797	assert uniform_spec is None
2798
2799
2800	def test_generate_scheduler_kv_cache_config():
2801	kv_cache_specs = {
2802	"layer_1": new_kv_cache_spec(),
2803	"layer_2": new_kv_cache_spec(head_size=128),
2804	}
2805	kv_cache_configs = [
2806	KVCacheConfig(
2807	num_blocks=10,
2808	kv_cache_tensors=[],
2809	kv_cache_groups=[
2810	KVCacheGroupSpec(
2811	["layer_1", "layer_2"],
2812	UniformTypeKVCacheSpecs(
2813	block_size=16, kv_cache_specs=kv_cache_specs
2814	),
2815	),
2816	],
2817	)
2818	]
2819	scheduler_kv_cache_config = generate_scheduler_kv_cache_config(kv_cache_configs)
2820	assert scheduler_kv_cache_config == KVCacheConfig(
2821	num_blocks=10,
2822	kv_cache_tensors=[],
2823	kv_cache_groups=[KVCacheGroupSpec(["layer_1", "layer_2"], new_kv_cache_spec())],
2824	)
2825
2826
2827	def _glm5_like_kv_cache_spec(
2828	mamba_spec_factory=new_mamba_spec,
2829	) -> tuple[dict[str, KVCacheSpec], list[str]]:
2830	"""(mamba, mamba, mamba, MLA + indexer) * 11 plus a trailing mamba."""
2831	kv_cache_spec: dict[str, KVCacheSpec] = {}
2832	mamba_layers = []
2833	for i in range(45):
2834	if i % 4 == 3:
2835	kv_cache_spec[f"layers.{i}.attn"] = MLAAttentionSpec(
2836	block_size=1024,
2837	num_kv_heads=1,
2838	head_size=576,
2839	dtype=torch.bfloat16,
2840	)
2841	kv_cache_spec[f"layers.{i}.indexer"] = MLAAttentionSpec(
2842	block_size=1024,
2843	num_kv_heads=1,
2844	head_size=132,
2845	dtype=torch.uint8,
2846	tokens_per_state=16,
2847	)
2848	else:
2849	name = f"layers.{i}.linear_attn"
2850	kv_cache_spec[name] = mamba_spec_factory()
2851	mamba_layers.append(name)
2852	return kv_cache_spec, mamba_layers
2853
2854
2855	def _glm5_like_kv_cache_spec_with_tail(
2856	mamba_spec_factory=new_mamba_spec,
2857	) -> dict[str, KVCacheSpec]:
2858	"""Production kpool=4 proportions: (mamba3, MLA + indexer + tail) 11.
2859
2860	The tail's logical page (2 * kpool * 2indexer_head_dim bf16) must fit
2861	inside the indexer page it parasitizes (block_size//kpool * 132 B), which
2862	is why the fixture drops the tokens_per_state=16 shape used above.
2863	"""
2864	kv_cache_spec, _ = _glm5_like_kv_cache_spec(mamba_spec_factory)
2865	for i in range(3, 45, 4):
2866	kpool = 4
2867	kv_cache_spec[f"layers.{i}.indexer"] = replace(
2868	cast(MLAAttentionSpec, kv_cache_spec[f"layers.{i}.indexer"]),
2869	tokens_per_state=kpool,
2870	)
2871	kv_cache_spec[f"layers.{i}.tail"] = KpoolTailSpec(
2872	block_size=kpool,
2873	num_kv_heads=2,
2874	head_size=128,
2875	head_size_v=0,
2876	dtype=torch.bfloat16,
2877	sliding_window=kpool,
2878	)
2879	return kv_cache_spec
2880
2881
2882	def _tensor_by_layer(kv_cache_config: KVCacheConfig) -> dict[str, KVCacheTensor]:
2883	return {
2884	layer_name: tensor
2885	for tensor in kv_cache_config.kv_cache_tensors
2886	for layer_name in tensor.layers
2887	}
2888
2889
2890	def _layer_offset(tensor: KVCacheTensor, layer_name: str) -> int:
2891	return tensor.offset + tensor.layers.index(layer_name) * tensor.layer_stride
2892
2893
2894	def test_get_kv_cache_config_balanced_mamba_hybrid():
2895	"""Hybrid slot sharing: mamba layers co-own the MLA slot tensors."""
2896	model_config = ModelConfig(max_model_len=8192)
2897	vllm_config = VllmConfig(model_config=model_config)
2898
2899	kv_cache_spec, mamba_layers = _glm5_like_kv_cache_spec()
2900	mla_page = kv_cache_spec["layers.3.attn"].page_size_bytes
2901	idx_page = kv_cache_spec["layers.3.indexer"].page_size_bytes
2902
2903	groups = kv_cache_utils.get_kv_cache_groups(vllm_config, kv_cache_spec)
2904	uniform_groups = [
2905	group
2906	for group in groups
2907	if isinstance(group.kv_cache_spec, UniformTypeKVCacheSpecs)
2908	]
2909	mamba_groups = [
2910	group for group in groups if isinstance(group.kv_cache_spec, MambaSpec)
2911	]
2912
2913	# Round-robin into G = ceil(34 / 11) = 4 groups so every mamba layer
2914	# gets an MLA slot.
2915	assert len(groups) == 5
2916	assert len(uniform_groups) == 1
2917	assert len(uniform_groups[0].layer_names) == 22
2918	assert uniform_groups[0].kv_cache_spec.get_max_layers_per_page_size() == 11
2919	assert [len(group.layer_names) for group in mamba_groups] == [9, 9, 8, 8]
2920	for k, name in enumerate(mamba_layers):
2921	assert name in mamba_groups[k % 4].layer_names
2922	# Mamba pages are padded up to the MLA page.
2923	for group in mamba_groups:
2924	assert group.kv_cache_spec.page_size_padded == mla_page
2925	assert group.kv_cache_spec.page_size_bytes == mla_page
2926
2927	bytes_per_block = kv_cache_utils._pool_bytes_per_block(groups)
2928	assert bytes_per_block == 11 * mla_page + 11 * idx_page
2929
2930	# Every block id is charged the full per-block sum.
2931	attn_blocks = uniform_groups[0].kv_cache_spec.max_memory_usage_pages(vllm_config)
2932	mamba_blocks_per_group = 1 + new_mamba_spec().num_speculative_blocks
2933	blocks_per_request = attn_blocks + 4 * mamba_blocks_per_group
2934	assert (
2935	kv_cache_utils._max_memory_usage_bytes_from_groups(vllm_config, groups)
2936	== blocks_per_request * bytes_per_block
2937	)
2938
2939	available_memory = bytes_per_block * 100 + 1
2940	kv_cache_config = kv_cache_utils.get_kv_cache_config_from_groups(
2941	vllm_config, groups, available_memory
2942	)
2943	assert kv_cache_config.num_blocks == 100
2944
2945	# Every logical layer has a view into one backing allocation. Mamba views
2946	# alias their corresponding MLA slot by using the same byte offset.
2947	assert len(kv_cache_config.kv_cache_tensors) == 56
2948	assert {t.size for t in kv_cache_config.kv_cache_tensors} == {bytes_per_block * 100}
2949	tensors = _tensor_by_layer(kv_cache_config)
2950	mla_layer_names = [f"layers.{i}.attn" for i in range(3, 45, 4)]
2951	for i, mla_name in enumerate(mla_layer_names):
2952	mla_tensor = tensors[mla_name]
2953	assert mla_tensor.block_stride == mla_page
2954	for group in mamba_groups:
2955	if i < len(group.layer_names):
2956	assert tensors[group.layer_names[i]].offset == mla_tensor.offset
2957	for i in range(11):
2958	assert tensors[f"layers.{4 * i + 3}.indexer"].block_stride == idx_page
2959
2960	total_allocated = next(iter({t.size for t in kv_cache_config.kv_cache_tensors}))
2961	assert total_allocated == bytes_per_block * 100
2962	assert 0 <= available_memory - total_allocated < bytes_per_block
2963
2964	assert get_max_concurrency_for_kv_cache_config(
2965	vllm_config, kv_cache_config
2966	) == pytest.approx(100 / blocks_per_request)
2967
2968
2969	def test_get_kv_cache_config_kpool_tail_coowns_indexer_tensor():
2970	"""The kpool tail parasitizes the indexer tensors instead of getting its
2971	own: sibling idx/tail tensors paired by layer order, zero standalone tail
2972	bytes, one shared block per request, and no prefix-caching leakage from
2973	the kpool-sized scratch group."""
2974	model_config = ModelConfig(max_model_len=8192)
2975	vllm_config = VllmConfig(model_config=model_config)
2976
2977	kv_cache_spec = _glm5_like_kv_cache_spec_with_tail()
2978	mla_page = kv_cache_spec["layers.3.attn"].page_size_bytes
2979	idx_page = kv_cache_spec["layers.3.indexer"].page_size_bytes
2980	tail_logical_page = kv_cache_spec["layers.3.tail"].page_size_bytes
2981	assert tail_logical_page < idx_page
2982
2983	groups = kv_cache_utils.get_kv_cache_groups(vllm_config, kv_cache_spec)
2984	tail_group = next(
2985	group
2986	for group in groups
2987	if isinstance(group.kv_cache_spec, UniformTypeKVCacheSpecs)
2988	and all(
2989	isinstance(spec, KpoolTailSpec)
2990	for spec in group.kv_cache_spec.kv_cache_specs.values()
2991	)
2992	)
2993	# The tail never prefix-caches; its page is padded up to the indexer page
2994	# so the runner's strided view rides the indexer storage.
2995	assert not tail_group.kv_cache_spec.prefix_cacheable
2996	tail_inner = cast(
2997	KpoolTailSpec, tail_group.kv_cache_spec.kv_cache_specs["layers.3.tail"]
2998	)
2999	assert tail_inner.page_size_padded == idx_page
3000
3001	bytes_per_block = kv_cache_utils._pool_bytes_per_block(groups)
3002	assert bytes_per_block == 11 * mla_page + 11 * idx_page
3003
3004	kv_cache_config = kv_cache_utils.get_kv_cache_config_from_groups(
3005	vllm_config, groups, bytes_per_block * 100 + 1
3006	)
3007	assert kv_cache_config.num_blocks == 100
3008	# Tail layers get logical views but no additional storage: each view aliases
3009	# its sibling indexer at the same offset in the shared allocation.
3010	assert len(kv_cache_config.kv_cache_tensors) == 67
3011	tensors = _tensor_by_layer(kv_cache_config)
3012	for i in range(11):
3013	idx_name = f"layers.{4 * i + 3}.indexer"
3014	tail_name = f"layers.{4 * i + 3}.tail"
3015	assert tensors[idx_name].offset == tensors[tail_name].offset
3016	assert {t.size for t in kv_cache_config.kv_cache_tensors} == {bytes_per_block * 100}
3017
3018	# The layout detector surfaces the sibling names in layer order for the
3019	# accounting and connector paths.
3020	layout = kv_cache_utils._glm5_next_tensor_layout(kv_cache_config.kv_cache_groups)
3021	assert layout is not None
3022	assert layout[3] == [f"layers.{4 * i + 3}.indexer" for i in range(11)]
3023	assert layout[6] == [f"layers.{4 * i + 3}.tail" for i in range(11)]
3024
3025	# Accounting charges exactly one extra shared block per request for the
3026	# tail, not a per-sequence allocation.
3027	attn_group = next(
3028	group
3029	for group in groups
3030	if isinstance(group.kv_cache_spec, UniformTypeKVCacheSpecs)
3031	and group is not tail_group
3032	)
3033	attn_blocks = attn_group.kv_cache_spec.max_memory_usage_pages(vllm_config)
3034	mamba_blocks_per_group = 1 + new_mamba_spec().num_speculative_blocks
3035	blocks_per_request = attn_blocks + 4 * mamba_blocks_per_group + 1
3036	assert (
3037	kv_cache_utils._max_memory_usage_bytes_from_groups(vllm_config, groups)
3038	== blocks_per_request * bytes_per_block
3039	)
3040
3041
3042	def test_glm5_kpool_tail_does_not_drag_hash_block_size():
3043	"""The tail's kpool-sized scratch block (4 tokens) must not constrain the
3044	prefix-cache hash granularity: participating groups alone decide it."""
3045	model_config = ModelConfig(max_model_len=8192)
3046	vllm_config = VllmConfig(model_config=model_config)
3047
3048	def align_mamba():
3049	return new_mamba_spec(mamba_cache_mode="align")
3050
3051	groups = kv_cache_utils.get_kv_cache_groups(
3052	vllm_config, _glm5_like_kv_cache_spec_with_tail(align_mamba)
3053	)
3054	kv_cache_config = KVCacheConfig(
3055	num_blocks=1,
3056	kv_cache_tensors=[],
3057	kv_cache_groups=groups,
3058	)
3059	hash_vllm_config = SimpleNamespace(
3060	cache_config=SimpleNamespace(
3061	block_size=16,
3062	enable_prefix_caching=True,
3063	prefix_match_unit=None,
3064	),
3065	parallel_config=SimpleNamespace(decode_context_parallel_size=1),
3066	kv_transfer_config=object(),
3067	)
3068	# gcd(attn 1024, mamba 16) with the tail's 4 excluded; scheduler size is
3069	# the lcm including the tail (1024 % 4 == 0, so it coincides).
3070	assert kv_cache_utils.resolve_kv_cache_block_sizes(
3071	kv_cache_config, hash_vllm_config
3072	) == (1024, 16)
3073
3074
3075	def test_get_kv_cache_config_mamba_hybrid_sharing_infeasible():
3076	"""Reject GLM-5.3-Flash layouts whose Mamba page exceeds the MLA page."""
3077	model_config = ModelConfig(max_model_len=8192)
3078	vllm_config = VllmConfig(model_config=model_config)
3079
3080	# (512, 1024) fp32 state = 2 MiB per page > 1.125 MiB MLA page.
3081	def big_mamba_spec():
3082	return new_mamba_spec(shapes=((512, 1024),), dtypes=(torch.float32,))
3083
3084	kv_cache_spec, _ = _glm5_like_kv_cache_spec(big_mamba_spec)
3085	mla_page = kv_cache_spec["layers.3.attn"].page_size_bytes
3086	assert big_mamba_spec().page_size_bytes > mla_page
3087
3088	with pytest.raises(ValueError, match="does not fit the MLA page"):
3089	kv_cache_utils.get_kv_cache_groups(vllm_config, kv_cache_spec)
3090
3091
3092	def test_get_kv_cache_config_mamba_hybrid_sharing_infeasible_no_indexer():
3093	"""Use the generic layout error when no kpool indexer is present."""
3094	model_config = ModelConfig(max_model_len=8192)
3095	vllm_config = VllmConfig(model_config=model_config)
3096	vllm_config.cache_config.kv_cache_layout = "LBHNC"
3097
3098	kv_cache_spec: dict[str, KVCacheSpec] = {}
3099	for i in range(27):
3100	if i % 4 == 3 or i == 26:
3101	kv_cache_spec[f"layers.{i}.attn"] = MLAAttentionSpec(
3102	block_size=1024,
3103	num_kv_heads=1,
3104	head_size=576,
3105	dtype=torch.bfloat16,
3106	)
3107	else:
3108	kv_cache_spec[f"layers.{i}.linear_attn"] = new_mamba_spec(
3109	shapes=((512, 1024),), dtypes=(torch.float32,)
3110	)
3111	mla_page = kv_cache_spec["layers.3.attn"].page_size_bytes
3112	assert kv_cache_spec["layers.0.linear_attn"].page_size_bytes > mla_page
3113
3114	with pytest.raises(NotImplementedError, match="page size"):
3115	kv_cache_utils.get_kv_cache_groups(vllm_config, kv_cache_spec)
3116
3117
3118	def test_get_kv_cache_config_mamba_hybrid_sharing_prepadded_mamba():
3119	"""Platform-prepadded mamba pages (mamba_page_size_padded hint) must not
3120	disable slot sharing; the layout re-pads them to the MLA page."""
3121	model_config = ModelConfig(max_model_len=8192)
3122	vllm_config = VllmConfig(model_config=model_config)
3123
3124	def prepadded_mamba_spec():
3125	return new_mamba_spec(page_size_padded=294_912)
3126
3127	kv_cache_spec, _ = _glm5_like_kv_cache_spec(prepadded_mamba_spec)
3128	mla_page = kv_cache_spec["layers.3.attn"].page_size_bytes
3129	assert prepadded_mamba_spec().page_size_bytes < mla_page
3130
3131	groups = kv_cache_utils.get_kv_cache_groups(vllm_config, kv_cache_spec)
3132	mamba_groups = [
3133	group for group in groups if isinstance(group.kv_cache_spec, MambaSpec)
3134	]
3135	assert len(groups) == 5
3136	assert [len(group.layer_names) for group in mamba_groups] == [9, 9, 8, 8]
3137	for group in mamba_groups:
3138	assert group.kv_cache_spec.page_size_bytes == mla_page
3139
3140	bytes_per_block = kv_cache_utils._pool_bytes_per_block(groups)
3141	kv_cache_config = kv_cache_utils.get_kv_cache_config_from_groups(
3142	vllm_config, groups, bytes_per_block * 100 + 1
3143	)
3144	assert kv_cache_config.num_blocks == 100
3145	assert len(kv_cache_config.kv_cache_tensors) == 56
3146
3147
3148	def test_get_kv_cache_config_mamba_hybrid_sharing_pp_balanced_projection():
3149	"""Round-robin mamba grouping keeps practical PP splits balanced: every
3150	stage's largest projected mamba group slice fits its projected MLA
3151	layers, so per-stage slot tensors all have an MLA owner."""
3152	model_config = ModelConfig(max_model_len=8192)
3153	vllm_config = VllmConfig(model_config=model_config)
3154
3155	kv_cache_spec, _ = _glm5_like_kv_cache_spec()
3156	global_groups = kv_cache_utils.get_kv_cache_groups(vllm_config, kv_cache_spec)
3157	mla_page = kv_cache_spec["layers.3.attn"].page_size_bytes
3158	idx_page = kv_cache_spec["layers.3.indexer"].page_size_bytes
3159
3160	# PP=2 split at transformer layer 22/23.
3161	# Stage 0 of a PP=2 split at transformer layer 22/23: 5 MLA(+indexer)
3162	# layers, projected mamba groups [5, 5, 4, 4].
3163	worker_spec = {n: s for n, s in kv_cache_spec.items() if int(n.split(".")[1]) <= 22}
3164	groups = kv_cache_utils._project_kv_cache_groups_to_worker(
3165	global_groups, worker_spec
3166	)
3167	mamba_groups = [
3168	group for group in groups if isinstance(group.kv_cache_spec, MambaSpec)
3169	]
3170	assert [len(group.layer_names) for group in mamba_groups] == [5, 5, 4, 4]
3171
3172	bytes_per_block = kv_cache_utils._pool_bytes_per_block(groups)
3173	assert bytes_per_block == 5 * mla_page + 5 * idx_page
3174
3175	kv_cache_config = kv_cache_utils.get_kv_cache_config_from_groups(
3176	vllm_config, groups, bytes_per_block * 100 + 1
3177	)
3178	assert kv_cache_config.num_blocks == 100
3179	tensors = _tensor_by_layer(kv_cache_config)
3180	# Every projected Mamba slot aliases an MLA owner.
3181	assert tensors["layers.3.attn"].offset == tensors["layers.0.linear_attn"].offset
3182	assert tensors["layers.3.attn"].offset == tensors["layers.1.linear_attn"].offset
3183	assert tensors["layers.3.attn"].offset == tensors["layers.2.linear_attn"].offset
3184	assert tensors["layers.3.attn"].offset == tensors["layers.4.linear_attn"].offset
3185	# The last slot only hosts the two groups with a 5th projected layer.
3186	assert tensors["layers.19.attn"].offset == tensors["layers.21.linear_attn"].offset
3187	assert tensors["layers.19.attn"].offset == tensors["layers.22.linear_attn"].offset
3188	assert {t.size for t in kv_cache_config.kv_cache_tensors} == {bytes_per_block * 100}
3189
3190
3191	def test_get_kv_cache_config_mamba_hybrid_sharing_pp_group_count_bump(monkeypatch):
3192	"""PP=4's default partition [11,11,12,11] leaves stage 0 with 2 MLA
3193	layers but a round-robin slice of 3 under the minimum 4 mamba groups;
3194	grouping bumps to 5 groups so every stage's projection keeps sharing
3195	on instead of silently falling back on stage 0."""
3196	from vllm.config import ParallelConfig
3197	from vllm.distributed.utils import get_pp_indices
3198
3199	monkeypatch.setattr(ModelConfig, "get_total_num_hidden_layers", lambda self: 45)
3200	vllm_config = VllmConfig(
3201	model_config=ModelConfig(max_model_len=8192),
3202	parallel_config=ParallelConfig(pipeline_parallel_size=4),
3203	)
3204
3205	kv_cache_spec, mamba_layers = _glm5_like_kv_cache_spec()
3206	groups = kv_cache_utils.get_kv_cache_groups(vllm_config, kv_cache_spec)
3207	mamba_groups = [g for g in groups if isinstance(g.kv_cache_spec, MambaSpec)]
3208	assert [len(g.layer_names) for g in mamba_groups] == [7, 7, 7, 7, 6]
3209	for k, name in enumerate(mamba_layers):
3210	assert name in mamba_groups[k % 5].layer_names
3211
3212	for rank in range(4):
3213	start, end = get_pp_indices(45, rank, 4)
3214	worker_spec = {
3215	n: s
3216	for n, s in kv_cache_spec.items()
3217	if start <= int(n.split(".")[1]) < end
3218	}
3219	projected = kv_cache_utils._project_kv_cache_groups_to_worker(
3220	groups, worker_spec
3221	)
3222	assert kv_cache_utils._glm5_next_tensor_layout(projected) is not None
3223
3224
3225	def test_get_kv_cache_config_mamba_hybrid_sharing_pp_starved_stage(monkeypatch):
3226	"""Reject PP stages whose Mamba layers have no MLA slot to share."""
3227	from vllm.config import ParallelConfig
3228
3229	monkeypatch.setattr(ModelConfig, "get_total_num_hidden_layers", lambda self: 45)
3230	monkeypatch.setenv("VLLM_PP_LAYER_PARTITION", "3,42")
3231	vllm_config = VllmConfig(
3232	model_config=ModelConfig(max_model_len=8192),
3233	parallel_config=ParallelConfig(pipeline_parallel_size=2),
3234	)
3235
3236	kv_cache_spec, _ = _glm5_like_kv_cache_spec()
3237	with pytest.raises(ValueError, match="VLLM_PP_LAYER_PARTITION"):
3238	kv_cache_utils.get_kv_cache_groups(vllm_config, kv_cache_spec)
3239
3240
3241	def test_get_kv_cache_config_mamba_hybrid_sharing_beats_cross_layers_flag():
3242	"""Hybrid slot sharing must take precedence over the experimental
3243	enable_cross_layers_blocks packed layout: generic packing would give the
3244	MLA slots a strided view, breaking the contiguous virtual split."""
3245	model_config = ModelConfig(max_model_len=8192)
3246	kv_transfer_config = KVTransferConfig(
3247	kv_connector="NixlConnector",
3248	kv_role="kv_both",
3249	kv_connector_extra_config={"enable_cross_layers_blocks": "true"},
3250	)
3251	vllm_config = VllmConfig(
3252	model_config=model_config, kv_transfer_config=kv_transfer_config
3253	)
3254
3255	kv_cache_spec, _ = _glm5_like_kv_cache_spec()
3256	groups = kv_cache_utils.get_kv_cache_groups(vllm_config, kv_cache_spec)
3257
3258	mla_page = kv_cache_spec["layers.3.attn"].page_size_bytes
3259	idx_page = kv_cache_spec["layers.3.indexer"].page_size_bytes
3260	bytes_per_block = kv_cache_utils._pool_bytes_per_block(groups)
3261	assert bytes_per_block == 11 * mla_page + 11 * idx_page
3262
3263	kv_cache_config = kv_cache_utils.get_kv_cache_config_from_groups(
3264	vllm_config, groups, bytes_per_block * 100 + 1
3265	)
3266	assert kv_cache_config.num_blocks == 100
3267	assert len(kv_cache_config.kv_cache_tensors) == 56
3268	assert all(
3269	tensor.block_stride in (mla_page, idx_page)
3270	for tensor in kv_cache_config.kv_cache_tensors
3271	)
3272	assert {t.size for t in kv_cache_config.kv_cache_tensors} == {bytes_per_block * 100}
3273
3274
3275	def test_get_kv_cache_config_mamba_hybrid_sharing_no_indexer():
3276	"""Kimi-Linear-like: MLA without indexer layers, idx_stride == 0."""
3277	model_config = ModelConfig(max_model_len=8192)
3278	vllm_config = VllmConfig(model_config=model_config)
3279	vllm_config.cache_config.kv_cache_layout = "LBNHC"
3280
3281	# 20 mamba + 7 MLA layers -> G = ceil(20 / 7) = 3 groups of [7, 7, 6].
3282	kv_cache_spec: dict[str, KVCacheSpec] = {}
3283	for i in range(27):
3284	if i % 4 == 3 or i == 26:
3285	kv_cache_spec[f"layers.{i}.attn"] = MLAAttentionSpec(
3286	block_size=1024,
3287	num_kv_heads=1,
3288	head_size=576,
3289	dtype=torch.bfloat16,
3290	)
3291	else:
3292	kv_cache_spec[f"layers.{i}.linear_attn"] = new_mamba_spec()
3293	mla_page = kv_cache_spec["layers.3.attn"].page_size_bytes
3294
3295	groups = kv_cache_utils.get_kv_cache_groups(vllm_config, kv_cache_spec)
3296	mamba_groups = [
3297	group for group in groups if isinstance(group.kv_cache_spec, MambaSpec)
3298	]
3299	assert len(groups) == 4
3300	assert [len(group.layer_names) for group in mamba_groups] == [7, 7, 6]
3301	for group in mamba_groups:
3302	assert group.kv_cache_spec.page_size_bytes == mla_page
3303
3304	bytes_per_block = kv_cache_utils._pool_bytes_per_block(groups)
3305	assert bytes_per_block == 7 * mla_page
3306
3307	kv_cache_config = kv_cache_utils.get_kv_cache_config_from_groups(
3308	vllm_config, groups, bytes_per_block * 100 + 1
3309	)
3310	assert kv_cache_config.num_blocks == 100
3311	# Each group has one tensor descriptor and all groups alias the same slots.
3312	assert len(kv_cache_config.kv_cache_tensors) == 4
3313	tensors = _tensor_by_layer(kv_cache_config)
3314	mla_names = [f"layers.{i}.attn" for i in (*range(3, 27, 4), 26)]
3315	for index, mla_name in enumerate(mla_names):
3316	mla_offset = _layer_offset(tensors[mla_name], mla_name)
3317	for group in mamba_groups:
3318	if index < len(group.layer_names):
3319	mamba_name = group.layer_names[index]
3320	assert _layer_offset(tensors[mamba_name], mamba_name) == mla_offset
3321	assert {t.size for t in kv_cache_config.kv_cache_tensors} == {bytes_per_block * 100}
3322
3323
3324	def test_get_kv_cache_capacity_after_scheduler_unwrap():
3325	"""max_concurrency must survive the scheduler-config unwrap.
3326
3327	Regression for the balanced-mamba hybrid layout:
3328	``generate_scheduler_kv_cache_config`` flattens the MLA
3329	``UniformTypeKVCacheSpecs`` group into a single ``MLAAttentionSpec``, so the
3330	scheduler config holds an MLA spec (page size A) next to ``MambaSpec``
3331	groups (page size B). ``EngineCore._initialize_kv_caches`` calls
3332	``get_kv_cache_capacity`` on that unwrapped config; the uniform-page-size
3333	path used to assert-fail (``assert len(page_sizes) == 1``) on this topology.
3334	"""
3335	model_config = ModelConfig(max_model_len=8192)
3336	vllm_config = VllmConfig(model_config=model_config)
3337
3338	kv_cache_spec: dict[str, KVCacheSpec] = {}
3339	for i in range(45):
3340	if i % 4 == 3:
3341	kv_cache_spec[f"layers.{i}.attn"] = MLAAttentionSpec(
3342	block_size=1024,
3343	num_kv_heads=1,
3344	head_size=576,
3345	dtype=torch.bfloat16,
3346	)
3347	kv_cache_spec[f"layers.{i}.indexer"] = MLAAttentionSpec(
3348	block_size=1024,
3349	num_kv_heads=1,
3350	head_size=132,
3351	dtype=torch.uint8,
3352	tokens_per_state=16,
3353	)
3354	else:
3355	kv_cache_spec[f"layers.{i}.linear_attn"] = new_mamba_spec()
3356
3357	groups = kv_cache_utils.get_kv_cache_groups(vllm_config, kv_cache_spec)
3358	bytes_per_block = kv_cache_utils._pool_bytes_per_block(groups)
3359	kv_cache_config = kv_cache_utils.get_kv_cache_config_from_groups(
3360	vllm_config, groups, bytes_per_block * 100 + 1
3361	)
3362
3363	scheduler_config = generate_scheduler_kv_cache_config([kv_cache_config])
3364	# Confirm we are exercising the regression path: the MLA group is no longer
3365	# UniformTypeKVCacheSpecs, so sibling specs have differing page sizes.
3366	assert not any(
3367	isinstance(g.kv_cache_spec, UniformTypeKVCacheSpecs)
3368	for g in scheduler_config.kv_cache_groups
3369	)
3370
3371	unwrapped_groups = scheduler_config.kv_cache_groups
3372	expected_max_mem = kv_cache_utils._max_memory_usage_bytes_from_groups(
3373	vllm_config, unwrapped_groups
3374	)
3375	expected_pool = kv_cache_utils._pool_bytes_per_block(unwrapped_groups)
3376	expected_blocks_per_request = (
3377	expected_max_mem + expected_pool - 1
3378	) // expected_pool
3379
3380	_, max_concurrency = get_kv_cache_capacity(vllm_config, scheduler_config)
3381	assert max_concurrency > 0
3382	assert max_concurrency == pytest.approx(
3383	scheduler_config.num_blocks / expected_blocks_per_request
3384	)
3385
3386
3387	def new_mla_spec(cache_dtype_str=None, block_size: int = 16):
3388	# head_size = kv_lora_rank(512) + qk_rope_head_dim(64) = 576
3389	return MLAAttentionSpec(
3390	block_size=block_size,
3391	num_kv_heads=1,
3392	max_tp_shards=1,
3393	head_size=576,
3394	dtype=torch.float32,
3395	cache_dtype_str=cache_dtype_str,
3396	)
3397
3398
3399	def new_swa_mla_spec(head_size=576, sliding_window=128, model_version=None):
3400	return SlidingWindowMLASpec(
3401	block_size=16,
3402	num_kv_heads=1,
3403	max_tp_shards=1,
3404	head_size=head_size,
3405	dtype=torch.float32,
3406	sliding_window=sliding_window,
3407	model_version=model_version,
3408	)
3409
3410
3411	def new_indexer_mla_spec(block_size=16):
3412	# Sparse-attention indexer k_cache: an MLAAttentionSpec with a much smaller
3413	# page size than the main MLA attention (uint8, small head), so their pages
3414	# cannot be unified.
3415	return MLAAttentionSpec(
3416	block_size=block_size,
3417	num_kv_heads=1,
3418	head_size=132,
3419	dtype=torch.uint8,
3420	)
3421
3422
3423	def test_mixed_page_size_groups_use_spec_compatibility():
3424	specs = {}
3425	for i in range(3):
3426	specs[f"mla.{i}"] = new_mla_spec()
3427	specs[f"indexer.{i}"] = new_indexer_mla_spec()
3428	specs.update({f"swa.{i}": new_swa_mla_spec(head_size=1024) for i in range(5)})
3429
3430	config = _grouping_config()
3431	config.cache_config = CacheConfig()
3432	config.cache_config.kv_cache_layout = "BLNHC"
3433	groups = get_kv_cache_groups(config, specs)
3434
3435	assert len(groups) == 3
3436	assert {name for group in groups for name in group.layer_names} == set(specs)
3437	assert sorted(len(group.layer_names) for group in groups) == [2, 3, 6]
3438
3439
3440	def _grouping_config():
3441	cache_config = CacheConfig()
3442	cache_config.kv_cache_layout = "LBNHC"
3443	return SimpleNamespace(
3444	scheduler_config=SimpleNamespace(disable_hybrid_kv_cache_manager=False),
3445	speculative_config=None,
3446	cache_config=cache_config,
3447	)
3448
3449
3450	def test_hidden_state_group_preserves_hybrid_prefix_cache_granularity():
3451	block_size = 544
3452	full_spec = FullAttentionSpec(
3453	block_size=block_size,
3454	num_kv_heads=1,
3455	# FullAttentionSpec stores both K and V, so this produces a
3456	# 544 * 1 * (512 + 512) * 2 = 1,114,112-byte page.
3457	head_size=512,
3458	dtype=torch.bfloat16,
3459	)
3460	mamba_spec = MambaSpec(
3461	block_size=block_size,
3462	shapes=((557056,),),
3463	dtypes=(torch.bfloat16,),
3464	mamba_cache_mode="align",
3465	)
3466	hidden_spec = HiddenStateCacheSpec(
3467	block_size=block_size,
3468	num_kv_heads=3,
3469	head_size=1024,
3470	dtype=torch.bfloat16,
3471	)
3472	assert full_spec.page_size_bytes == mamba_spec.page_size_bytes
3473
3474	groups = get_kv_cache_groups(
3475	_grouping_config(),
3476	{
3477	"model.full_attn": full_spec,
3478	"model.mamba": mamba_spec,
3479	"cache_only_layers.0": hidden_spec,
3480	},
3481	)
3482
3483	hidden_group = next(
3484	group
3485	for group in groups
3486	if isinstance(group.kv_cache_spec, HiddenStateCacheSpec)
3487	)
3488	assert hidden_group.kv_cache_spec.block_size == 136
3489
3490	kv_cache_config = KVCacheConfig(
3491	num_blocks=1,
3492	kv_cache_tensors=[],
3493	kv_cache_groups=groups,
3494	)
3495	vllm_config = SimpleNamespace(
3496	cache_config=SimpleNamespace(
3497	block_size=16,
3498	enable_prefix_caching=True,
3499	prefix_match_unit=None,
3500	),
3501	parallel_config=SimpleNamespace(decode_context_parallel_size=1),
3502	kv_transfer_config=object(),
3503	)
3504	assert kv_cache_utils.resolve_kv_cache_block_sizes(
3505	kv_cache_config, vllm_config
3506	) == (544, 136)
3507
3508
3509	def test_resolve_dcp_kv_block_size_unwraps_uniform_type_specs():
3510	attention = FullAttentionSpec(
3511	block_size=16,
3512	num_kv_heads=1,
3513	head_size=64,
3514	dtype=torch.float16,
3515	)
3516	mamba = MambaSpec(
3517	block_size=16,
3518	shapes=((1, 1),),
3519	dtypes=(torch.float16,),
3520	)
3521	wrapped_attention = UniformTypeKVCacheSpecs(
3522	block_size=16,
3523	kv_cache_specs={"attention.0": attention, "attention.1": attention},
3524	)
3525	wrapped_mamba = UniformTypeKVCacheSpecs(
3526	block_size=16,
3527	kv_cache_specs={"mamba.0": mamba, "mamba.1": mamba},
3528	)
3529
3530	assert kv_cache_utils.resolve_dcp_kv_block_size(attention, 4) == 64
3531	assert kv_cache_utils.resolve_dcp_kv_block_size(wrapped_attention, 4) == 64
3532	assert kv_cache_utils.resolve_dcp_kv_block_size(mamba, 4) == 16
3533	assert kv_cache_utils.resolve_dcp_kv_block_size(wrapped_mamba, 4) == 16
3534
3535	scaled_attention = kv_cache_utils.resolve_dcp_kv_cache_spec(wrapped_attention, 4)
3536	assert scaled_attention.block_size == 64
3537	assert isinstance(scaled_attention, UniformTypeKVCacheSpecs)
3538	assert all(
3539	spec.block_size == 64 for spec in scaled_attention.kv_cache_specs.values()
3540	)
3541	assert kv_cache_utils.resolve_dcp_kv_cache_spec(wrapped_mamba, 4) is wrapped_mamba
3542
3543
3544	def test_multi_run_layer_compact_strides_place_hoisted_heads():
3545	"""A layer-compact run region is its own dense allocation: under LHBNC the head
3546	groups sit between the layers and the blocks, so a run's block stride is one head
3547	group's slice, not the whole page (regression test for setStorage out-of-bounds
3548	on ROCm hybrid models)."""
3549	full = new_kv_cache_spec()
3550	swa = new_sliding_window_spec(sliding_window=full.block_size * 2)
3551	page = full.page_size_bytes
3552	assert swa.page_size_bytes == page
3553
3554	vllm_config = VllmConfig(model_config=ModelConfig(max_model_len=16))
3555	vllm_config.cache_config.kv_cache_layout = "LHBNC"
3556	config = kv_cache_utils.get_kv_cache_config_from_groups(
3557	vllm_config,
3558	[
3559	KVCacheGroupSpec(["full.0", "full.1"], full),
3560	KVCacheGroupSpec(["swa.0", "swa.1"], swa),
3561	],
3562	available_memory=4 * page * 8,
3563	)
3564
3565	num_blocks = config.num_blocks
3566	head_group = full.block_size * full.state_content_size_bytes
3567	for tensor in config.kv_cache_tensors:
3568	# One head group per block, head groups between the layers and blocks.
3569	assert tensor.block_stride == head_group
3570	assert tensor.layer_stride == full.num_heads * num_blocks * head_group
3571	# The last block of the last layer's last head group must stay within the
3572	# allocation (the old page-based stride overran it).
3573	end = (
3574	tensor.offset
3575	+ tensor.layer_stride * (len(tensor.layers) - 1)
3576	+ (full.num_heads - 1) * num_blocks * head_group
3577	+ tensor.block_stride * (num_blocks - 1)
3578	+ head_group
3579	)
3580	assert end <= tensor.size
3581
3582
3583	def test_mla_draft_prefers_standard_layout_when_pages_can_be_unified():
3584	specs = {
3585	"target.0.attn": new_mla_spec(),
3586	"draft.0": new_sliding_window_spec(num_kv_heads=1, head_size=288),
3587	}
3588	assert len({spec.page_size_bytes for spec in specs.values()}) == 1
3589
3590	groups = get_kv_cache_groups(_grouping_config(), specs)
3591
3592	assert len(groups) == 2
3593	assert all(
3594	not isinstance(group.kv_cache_spec, UniformTypeKVCacheSpecs) for group in groups
3595	)
3596
3597
3598	def test_mla_with_incompatible_swa_uses_one_full_allocation_group(caplog_vllm):
3599	# Sparse MLA pages cannot be padded safely. Keeping the draft's attention
3600	# compute sliding-window while promoting only its allocation semantics lets
3601	# every layer share the target's block table and remain contiguous.
3602	draft = new_sliding_window_spec(block_size=16)
3603	specs = {
3604	"target.0.attn": new_mla_spec(block_size=64),
3605	"target.0.indexer": new_indexer_mla_spec(block_size=64),
3606	"draft.0": draft,
3607	}
3608
3609	groups = get_kv_cache_groups(_grouping_config(), specs)
3610	assert len(groups) == 1
3611	assert set(groups[0].layer_names) == set(specs)
3612	group_spec = groups[0].kv_cache_spec
3613	assert isinstance(group_spec, UniformTypeKVCacheSpecs)
3614	assert group_spec.block_size == 64
3615	promoted_draft = group_spec.kv_cache_specs["draft.0"]
3616	assert isinstance(promoted_draft, FullAttentionSpec)
3617	assert not isinstance(promoted_draft, SlidingWindowSpec)
3618	assert promoted_draft.block_size == 64
3619	assert promoted_draft.sliding_window == draft.sliding_window
3620	assert specs["draft.0"] is draft
3621	assert "attention compute is unchanged" in caplog_vllm.text
3622
3623
3624	def test_hidden_states_with_tp_scales_page_size():
3625	"""When TP shrinks KV pages below the hidden-state per-token cost,
3626	get_kv_cache_groups must scale up target block sizes so that the
3627	common page accommodates the unsharded hidden states."""
3628	# Simulate TP=4 sharding a model with 8 KV heads → 2 per rank.
3629	# KV page = block_size(16) * num_kv_heads(2) * head_size(64) * dtype(2)
3630	# = 16 * 2 * 64 * 2 = 4096 bytes.
3631	kv_spec = new_kv_cache_spec(
3632	block_size=16,
3633	num_kv_heads=2,
3634	head_size=64,
3635	dtype=torch.bfloat16,
3636	)
3637	# Hidden-state per-token cost = num_hidden_states(6) * hidden_size(512)
3638	# * dtype(2) = 6144 bytes, which exceeds the 4096-byte KV page.
3639	hs_spec = HiddenStateCacheSpec(
3640	block_size=16,
3641	num_kv_heads=6,
3642	head_size=512,
3643	dtype=torch.bfloat16,
3644	)
3645	specs = {
3646	"target.0.attn": kv_spec,
3647	"target.1.attn": kv_spec,
3648	"cache_only_layers.48": hs_spec,
3649	}
3650
3651	groups = get_kv_cache_groups(_grouping_config(), specs)
3652
3653	# The hidden-state layer should be present and no assertion should fire.
3654	all_layers = {name for g in groups for name in g.layer_names}
3655	assert "cache_only_layers.48" in all_layers
3656
3657	# The target group block sizes must have been scaled up.
3658	for g in groups:
3659	if "cache_only_layers.48" not in g.layer_names:
3660	assert g.kv_cache_spec.block_size > kv_spec.block_size
3661
3662
3663	def test_get_kv_cache_spec_kind_prefers_specific_attention_subclasses():
3664	assert get_kv_cache_spec_kind(new_mla_spec()) == KVCacheSpecKind.MLA_ATTENTION
3665
3666	sliding_window_mla_spec = SlidingWindowMLASpec(
3667	block_size=16,
3668	num_kv_heads=1,
3669	head_size=576,
3670	dtype=torch.float32,
3671	sliding_window=128,
3672	)
3673	assert (
3674	get_kv_cache_spec_kind(sliding_window_mla_spec)
3675	== KVCacheSpecKind.SLIDING_WINDOW_MLA
3676	)
3677
3678	sink_full_attention_spec = SinkFullAttentionSpec(
3679	block_size=16,
3680	num_kv_heads=1,
3681	head_size=64,
3682	dtype=torch.float32,
3683	sink_len=4,
3684	)
3685	assert (
3686	get_kv_cache_spec_kind(sink_full_attention_spec)
3687	== KVCacheSpecKind.SINK_FULL_ATTENTION
3688	)
3689
3690
3691	def test_get_kv_cache_spec_kind_unwraps_uniform_type_specs():
3692	uniform_mla_spec = UniformTypeKVCacheSpecs(
3693	block_size=16,
3694	kv_cache_specs={
3695	"layer_1": new_mla_spec(),
3696	"layer_2": new_mla_spec(cache_dtype_str="fp8"),
3697	},
3698	)
3699	assert get_kv_cache_spec_kind(uniform_mla_spec) == KVCacheSpecKind.MLA_ATTENTION
3700
3701	uniform_swa_mla_spec = UniformTypeKVCacheSpecs(
3702	block_size=16,
3703	kv_cache_specs={
3704	"layer_1": SlidingWindowMLASpec(
3705	block_size=16,
3706	num_kv_heads=1,
3707	head_size=576,
3708	dtype=torch.float32,
3709	sliding_window=128,
3710	),
3711	"layer_2": SlidingWindowMLASpec(
3712	block_size=16,
3713	num_kv_heads=1,
3714	head_size=1024,
3715	dtype=torch.float32,
3716	sliding_window=128,
3717	),
3718	},
3719	)
3720	assert (
3721	get_kv_cache_spec_kind(uniform_swa_mla_spec)
3722	== KVCacheSpecKind.SLIDING_WINDOW_MLA
3723	)
3724
3725
3726	def test_get_kv_cache_spec_kind_unknown_for_mixed_uniform_type_specs():
3727	uniform_mixed_spec = UniformTypeKVCacheSpecs(
3728	block_size=16,
3729	kv_cache_specs={
3730	"layer_1": new_mla_spec(),
3731	"layer_2": SlidingWindowMLASpec(
3732	block_size=16,
3733	num_kv_heads=1,
3734	head_size=576,
3735	dtype=torch.float32,
3736	sliding_window=128,
3737	),
3738	},
3739	)
3740	assert get_kv_cache_spec_kind(uniform_mixed_spec) == KVCacheSpecKind.UNKNOWN
3741
3742
3743	def test_get_kv_cache_spec_sliding_window_reads_windowed_specs():
3744	full_attention_spec = FullAttentionSpec(
3745	block_size=16,
3746	num_kv_heads=1,
3747	head_size=64,
3748	dtype=torch.float32,
3749	)
3750	sliding_window_spec = SlidingWindowSpec(
3751	block_size=16,
3752	num_kv_heads=1,
3753	head_size=64,
3754	dtype=torch.float32,
3755	sliding_window=128,
3756	)
3757
3758	assert get_kv_cache_spec_sliding_window(full_attention_spec) is None
3759	assert get_kv_cache_spec_sliding_window(sliding_window_spec) == 128
3760
3761
3762	def test_get_kv_cache_spec_sliding_window_unwraps_uniform_type_specs():
3763	uniform_window_spec = UniformTypeKVCacheSpecs(
3764	block_size=16,
3765	kv_cache_specs={
3766	"layer_1": SlidingWindowSpec(
3767	block_size=16,
3768	num_kv_heads=1,
3769	head_size=64,
3770	dtype=torch.float32,
3771	sliding_window=128,
3772	),
3773	"layer_2": SlidingWindowSpec(
3774	block_size=16,
3775	num_kv_heads=2,
3776	head_size=64,
3777	dtype=torch.float32,
3778	sliding_window=128,
3779	),
3780	},
3781	)
3782	mixed_window_spec = UniformTypeKVCacheSpecs(
3783	block_size=16,
3784	kv_cache_specs={
3785	"layer_1": SlidingWindowSpec(
3786	block_size=16,
3787	num_kv_heads=1,
3788	head_size=64,
3789	dtype=torch.float32,
3790	sliding_window=128,
3791	),
3792	"layer_2": SlidingWindowSpec(
3793	block_size=16,
3794	num_kv_heads=1,
3795	head_size=64,
3796	dtype=torch.float32,
3797	sliding_window=256,
3798	),
3799	},
3800	)
3801
3802	assert get_kv_cache_spec_sliding_window(uniform_window_spec) == 128
3803	assert get_kv_cache_spec_sliding_window(mixed_window_spec) is None
3804
3805
3806	def test_merge_mla_spec():
3807	kv_cache_specs = [
3808	new_mla_spec(),
3809	new_mla_spec(),
3810	]
3811	mla_spec = kv_cache_specs[0].merge(kv_cache_specs)
3812	assert mla_spec == new_mla_spec()
3813
3814	kv_cache_specs = [
3815	new_mla_spec(cache_dtype_str="fp8_ds_mla"),
3816	new_mla_spec(cache_dtype_str="fp8_ds_mla"),
3817	]
3818	mla_spec = kv_cache_specs[0].merge(kv_cache_specs)
3819	assert mla_spec == new_mla_spec(cache_dtype_str="fp8_ds_mla")
3820
3821	kv_cache_specs = [
3822	new_mla_spec(cache_dtype_str="fp8_ds_mla"),
3823	new_mla_spec(cache_dtype_str=None),
3824	]
3825	with pytest.raises(AssertionError):
3826	kv_cache_specs[0].merge(kv_cache_specs)
3827
3828	kv_cache_specs = [
3829	new_kv_cache_spec(),
3830	new_mla_spec(),
3831	]
3832	with pytest.raises(AssertionError):
3833	kv_cache_specs[0].merge(kv_cache_specs)
3834
3835	kv_cache_specs = [
3836	new_mla_spec(cache_dtype_str="fp8_ds_mla"),
3837	new_kv_cache_spec(),
3838	]
3839	with pytest.raises(AssertionError):
3840	kv_cache_specs[0].merge(kv_cache_specs)
3841
3842
3843	@pytest.mark.parametrize("hash_fn", [sha256, sha256_cbor])
3844	def test_request_block_hasher_with_prompt_embeds(hash_fn: Callable[[Any], bytes]):
3845	block_size = 3
3846	num_tokens = 2 * block_size
3847	prompt_token_ids = [_ for _ in range(num_tokens)]
3848	hidden_size = 5
3849	prompt_embeds = torch.randn((num_tokens, hidden_size))
3850
3851	request = make_request(
3852	request_id="0",
3853	prompt_token_ids=prompt_token_ids,
3854	block_size=block_size,
3855	hash_fn=hash_fn,
3856	prompt_embeds=prompt_embeds,
3857	)
3858
3859	block_hashes = request.block_hashes
3860	assert len(block_hashes) == 2
3861
3862	block1_embeds_hash = hashlib.sha256(
3863	tensor_data(prompt_embeds[:block_size])
3864	).digest()
3865	expected_hash1 = hash_fn(
3866	(
3867	kv_cache_utils.NONE_HASH,
3868	tuple(prompt_token_ids[:block_size]),
3869	(("prompt_embeds", block1_embeds_hash),),
3870	)
3871	)
3872	assert block_hashes[0] == expected_hash1
3873
3874	block2_embeds_hash = hashlib.sha256(
3875	tensor_data(prompt_embeds[block_size:num_tokens])
3876	).digest()
3877	expected_hash2 = hash_fn(
3878	(
3879	block_hashes[0],
3880	tuple(prompt_token_ids[block_size:num_tokens]),
3881	(("prompt_embeds", block2_embeds_hash),),
3882	)
3883	)
3884	assert block_hashes[1] == expected_hash2
3885
3886
3887	@pytest.mark.parametrize("hash_fn", [sha256, sha256_cbor])
3888	def test_request_with_prompt_embeds_and_mm_inputs(hash_fn: Callable[[Any], bytes]):
3889	block_size = 3
3890	num_tokens = 2 * block_size
3891	prompt_token_ids = [_ for _ in range(num_tokens)]
3892	hidden_size = 5
3893	prompt_embeds = torch.randn((num_tokens, hidden_size))
3894
3895	request = make_request(
3896	request_id="0",
3897	prompt_token_ids=prompt_token_ids,
3898	block_size=block_size,
3899	hash_fn=hash_fn,
3900	mm_positions=[
3901	PlaceholderRange(offset=0, length=3),
3902	PlaceholderRange(offset=3, length=3),
3903	],
3904	mm_hashes=["hash1", "hash2"],
3905	prompt_embeds=prompt_embeds,
3906	)
3907
3908	block_hashes = request.block_hashes
3909	assert len(block_hashes) == 2
3910
3911	block1_embeds_hash = hashlib.sha256(
3912	tensor_data(prompt_embeds[:block_size])
3913	).digest()
3914	expected_hash1 = hash_fn(
3915	(
3916	kv_cache_utils.NONE_HASH,
3917	tuple(prompt_token_ids[:block_size]),
3918	(("mm", "hash1", 0), ("prompt_embeds", block1_embeds_hash)),
3919	)
3920	)
3921	assert block_hashes[0] == expected_hash1
3922
3923	block2_embeds_hash = hashlib.sha256(
3924	tensor_data(prompt_embeds[block_size:num_tokens])
3925	).digest()
3926	expected_hash2 = hash_fn(
3927	(
3928	block_hashes[0],
3929	tuple(prompt_token_ids[block_size:num_tokens]),
3930	(("mm", "hash2", 0), ("prompt_embeds", block2_embeds_hash)),
3931	)
3932	)
3933	assert block_hashes[1] == expected_hash2
3934
3935
3936	def test_auto_fit_max_model_len():
3937	"""Test that max_model_len=-1 auto-fits to available GPU memory."""
3938	# Create config with original_max_model_len=-1 to trigger auto-fit
3939	model_config = ModelConfig(max_model_len=1024)
3940	# Simulate the user passing -1 by setting original_max_model_len
3941	model_config.original_max_model_len = -1
3942	vllm_config = VllmConfig(model_config=model_config)
3943	vllm_config.cache_config.kv_cache_layout = "LBNHC"
3944
3945	mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2 # 16KB per block per layer
3946	kv_cache_specs = {
3947	"layer_1": new_kv_cache_spec(),
3948	"layer_2": new_kv_cache_spec(),
3949	}
3950
3951	# With enough memory, max_model_len stays at the derived max
3952	large_available_memory = mem_per_block_per_layer * 2 * 1024 # plenty of memory
3953	_kv_cache_configs = get_kv_cache_configs(
3954	vllm_config, [kv_cache_specs], [large_available_memory]
3955	)
3956	assert vllm_config.model_config.max_model_len == 1024
3957
3958	# Reset for next test
3959	model_config = ModelConfig(max_model_len=1024)
3960	model_config.original_max_model_len = -1
3961	vllm_config = VllmConfig(model_config=model_config)
3962	vllm_config.cache_config.kv_cache_layout = "LBNHC"
3963
3964	# With limited memory, max_model_len should be reduced
3965	# Need memory for at least max_model_len tokens
3966	# 32 blocks worth of memory for 2 layers = can fit 32*16=512 tokens
3967	limited_memory = mem_per_block_per_layer * 2 * 32
3968	_kv_cache_configs = get_kv_cache_configs(
3969	vllm_config, [kv_cache_specs], [limited_memory]
3970	)
3971	# Should be reduced to fit in memory
3972	assert vllm_config.model_config.max_model_len < 1024
3973	assert vllm_config.model_config.max_model_len > 0
3974
3975
3976	def test_auto_fit_max_model_len_with_hybrid():
3977	"""Test that auto-fit works with hybrid KV cache specs."""
3978	# Create config with original_max_model_len=-1 to trigger auto-fit
3979	model_config = ModelConfig(max_model_len=8192)
3980	# Simulate the user passing -1 by setting original_max_model_len
3981	model_config.original_max_model_len = -1
3982	vllm_config = VllmConfig(model_config=model_config)
3983	vllm_config.cache_config.kv_cache_layout = "LBNHC"
3984
3985	mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2 # 16KB per block per layer
3986	gamma = 2
3987	kv_cache_specs = {
3988	"layer_1": new_mamba_spec(num_speculative_blocks=gamma),
3989	"layer_2": new_kv_cache_spec(),
3990	}
3991
3992	# One extra block on top of what a 1024-token request needs: the pool
3993	# reserves one block as the null block.
3994	available_memory = mem_per_block_per_layer * (1024 // 16 + 1 + gamma + 1)
3995	_kv_cache_configs = get_kv_cache_configs(
3996	vllm_config, [kv_cache_specs], [available_memory]
3997	)
3998	assert vllm_config.model_config.max_model_len == 1024
3999
4000
4001	def test_auto_fit_max_model_len_not_triggered():
4002	"""Test that auto-fit is not triggered when original_max_model_len is not -1."""
4003	model_config = ModelConfig(max_model_len=16)
4004	# original_max_model_len should be None by default, not -1
4005	vllm_config = VllmConfig(model_config=model_config)
4006	vllm_config.cache_config.kv_cache_layout = "LBNHC"
4007
4008	mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2
4009	kv_cache_specs = {
4010	"layer_1": new_kv_cache_spec(),
4011	"layer_2": new_kv_cache_spec(),
4012	}
4013
4014	# This should work normally without auto-fit
4015	_kv_cache_configs = get_kv_cache_configs(
4016	vllm_config, [kv_cache_specs], [mem_per_block_per_layer * 2 * 32]
4017	)
4018	assert vllm_config.model_config.max_model_len == 16
4019
4020
4021	def test_auto_fit_max_model_len_respects_num_gpu_blocks_override():
4022	"""Auto-fit must size max_model_len against the override-clamped pool, not
4023	the raw `available_memory`. Without this, auto-fit could pick a
4024	max_model_len that no longer fits once `num_gpu_blocks_override` is applied.
4025	"""
4026	model_config = ModelConfig(max_model_len=16384)
4027	model_config.original_max_model_len = -1 # request auto-fit
4028	vllm_config = VllmConfig(model_config=model_config)
4029	vllm_config.cache_config.kv_cache_layout = "LBNHC"
4030	# Cap the cache to 32 blocks regardless of available memory.
4031	vllm_config.cache_config.num_gpu_blocks_override = 32
4032
4033	mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2
4034	kv_cache_specs = {
4035	"layer_1": new_kv_cache_spec(), # block_size=16
4036	"layer_2": new_kv_cache_spec(),
4037	}
4038	# Plenty of raw memory (1024 blocks per layer would fit max_model_len=16384).
4039	large_available_memory = mem_per_block_per_layer * 2 * 1024
4040
4041	get_kv_cache_configs(vllm_config, [kv_cache_specs], [large_available_memory])
4042
4043	# 32 blocks * block_size 16 = 512 token slots, so max_model_len must
4044	# auto-fit at or below that.
4045	assert 0 < vllm_config.model_config.max_model_len <= 32 * 16
4046
4047
4048	def test_check_enough_kv_cache_memory_respects_num_gpu_blocks_override():
4049	"""Admission check must use the override-clamped pool size, not raw
4050	`available_memory`. Without this, startup could accept a max_model_len
4051	that does not actually fit in `num_gpu_blocks_override` blocks.
4052	"""
4053	model_config = ModelConfig(max_model_len=16384)
4054	vllm_config = VllmConfig(model_config=model_config)
4055	vllm_config.cache_config.kv_cache_layout = "LBNHC"
4056	# 32 blocks is far too small for max_model_len=16384 (would need 1024).
4057	vllm_config.cache_config.num_gpu_blocks_override = 32
4058
4059	mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2
4060	kv_cache_specs = {
4061	"layer_1": new_kv_cache_spec(),
4062	"layer_2": new_kv_cache_spec(),
4063	}
4064	# Plenty of raw memory: a bytes-only check against this would pass.
4065	large_available_memory = mem_per_block_per_layer * 2 * 1024
4066
4067	with pytest.raises(ValueError, match="max seq len"):
4068	get_kv_cache_configs(vllm_config, [kv_cache_specs], [large_available_memory])
4069
4070
4071	def test_unify_kv_cache_page_size_uses_padding_for_non_divisible_sizes():
4072	"""DFlash drafters can have a smaller head size than the target model.
4073
4074	For example, MiMo uses 192-dim target KV heads while its DFlash draft uses
4075	128-dim KV heads. The resulting page sizes are 3:2 rather than an integer
4076	block-size multiple, so the smaller page must be padded instead.
4077	"""
4078	# Attention layers read padded pages through a strided view, so padding is
4079	# allowed.
4080	target_spec = new_kv_cache_spec(
4081	block_size=16,
4082	num_kv_heads=1,
4083	head_size=192,
4084	dtype=torch.bfloat16,
4085	)
4086	draft_spec = new_sliding_window_spec(
4087	block_size=16,
4088	num_kv_heads=1,
4089	head_size=128,
4090	dtype=torch.bfloat16,
4091	sliding_window=1024,
4092	)
4093
4094	unified_specs = kv_cache_utils.unify_kv_cache_spec_page_size(
4095	{
4096	"target_attn": target_spec,
4097	"draft_attn": draft_spec,
4098	}
4099	)
4100
4101	assert unified_specs["target_attn"] == target_spec
4102	unified_draft_spec = unified_specs["draft_attn"]
4103	assert unified_draft_spec.block_size == draft_spec.block_size
4104	assert unified_draft_spec.real_page_size_bytes == draft_spec.real_page_size_bytes
4105	assert unified_draft_spec.page_size_padded == target_spec.page_size_bytes
4106	assert unified_draft_spec.page_size_bytes == target_spec.page_size_bytes
4107
4108
4109	def test_unpadded_page_size_includes_per_token_head_scales():
4110	# Per-token-head quant carries inline fp32 scales that are carved from the
4111	# raw KV allocation, so they must be budgeted into the offload width. The
4112	# packing is published by the owning backend's customize_spec hook.
4113	from vllm.v1.attention.backends.triton_attn import TritonAttentionBackend
4114
4115	dense = new_kv_cache_spec(dtype=torch.uint8)
4116	spec = TritonAttentionBackend.customize_spec(
4117	new_kv_cache_spec(
4118	dtype=torch.uint8, kv_quant_mode=KVQuantMode.FP8_PER_TOKEN_HEAD
4119	)
4120	)
4121	scales = 2 * spec.block_size * spec.num_kv_heads * 4
4122	assert spec.unpadded_page_size_bytes == dense.unpadded_page_size_bytes + scales
4123	assert spec.page_size_bytes == spec.unpadded_page_size_bytes
4124
4125
4126	def test_page_size_padded_wins():
4127	# An explicit padded page size takes precedence over the unpadded size.
4128	from vllm.v1.attention.backends.triton_attn import TritonAttentionBackend
4129
4130	spec = TritonAttentionBackend.customize_spec(
4131	new_kv_cache_spec(
4132	dtype=torch.uint8,
4133	kv_quant_mode=KVQuantMode.FP8_PER_TOKEN_HEAD,
4134	page_size_padded=65536,
4135	)
4136	)
4137	assert spec.page_size_bytes == 65536
4138
4139
4140	def test_unify_hybrid_kv_cache_specs():
4141	# 1. has_full_attention and has_sliding_window
4142	before_spec_1 = new_kv_cache_spec(block_size=64)
4143	before_spec_2 = new_sliding_window_spec(
4144	block_size=16, page_size_padded=32 * 1024, sliding_window=1024
4145	)
4146	kv_cache_spec = {
4147	"layer_1": before_spec_1,
4148	"layer_2": before_spec_2,
4149	}
4150	kv_cache_spec["draft_layer_1"] = replace(before_spec_1, dcp_sharded=False)
4151	kv_cache_spec["draft_layer_2"] = replace(before_spec_2, dcp_sharded=False)
4152	kv_cache_utils.unify_hybrid_kv_cache_specs(kv_cache_spec)
4153	expected_spec_1 = new_kv_cache_spec(block_size=64)
4154	expected_spec_2 = new_kv_cache_spec(
4155	block_size=64, page_size_padded=64 * 1024, sliding_window=1024
4156	)
4157	assert kv_cache_spec["layer_1"] == expected_spec_1
4158	assert kv_cache_spec["layer_2"] == expected_spec_2
4159	assert kv_cache_spec["layer_2"].page_size_bytes == 64 * 1024
4160	assert kv_cache_spec["draft_layer_2"] == replace(expected_spec_2, dcp_sharded=False)
4161
4162	# 2. has_full_attention and has_chunked_local_attention
4163	before_spec_1 = new_kv_cache_spec()
4164	before_spec_2 = new_chunked_local_attention_spec(
4165	page_size_padded=32 * 1024, attention_chunk_size=512
4166	)
4167	kv_cache_spec = {
4168	"layer_1": before_spec_1,
4169	"layer_2": before_spec_2,
4170	}
4171	kv_cache_utils.unify_hybrid_kv_cache_specs(kv_cache_spec)
4172	expected_spec_1 = new_kv_cache_spec()
4173	expected_spec_2 = new_kv_cache_spec(
4174	page_size_padded=32 * 1024, attention_chunk_size=512
4175	)
4176
4177	assert kv_cache_spec["layer_1"] == expected_spec_1
4178	assert kv_cache_spec["layer_2"] == expected_spec_2
4179
4180	# 3. has_full_attention, has_sliding_window and has_chunked_local_attention
4181	before_spec_1 = new_kv_cache_spec()
4182	before_spec_2 = new_sliding_window_spec(
4183	page_size_padded=32 * 1024, sliding_window=1024
4184	)
4185	before_spec_3 = new_chunked_local_attention_spec(
4186	page_size_padded=32 * 1024, attention_chunk_size=512
4187	)
4188	kv_cache_spec = {
4189	"layer_1": before_spec_1,
4190	"layer_2": before_spec_2,
4191	"layer_3": before_spec_3,
4192	}
4193	kv_cache_utils.unify_hybrid_kv_cache_specs(kv_cache_spec)
4194	expected_spec_1 = new_kv_cache_spec()
4195	expected_spec_2 = new_kv_cache_spec(page_size_padded=32 * 1024, sliding_window=1024)
4196	expected_spec_3 = new_kv_cache_spec(
4197	page_size_padded=32 * 1024, attention_chunk_size=512
4198	)
4199	assert kv_cache_spec["layer_1"] == expected_spec_1
4200	assert kv_cache_spec["layer_2"] == expected_spec_2
4201	assert kv_cache_spec["layer_3"] == expected_spec_3
4202
4203	# 4. No FullAttentionSpec, should not convert
4204	kv_cache_spec = {
4205	"layer_1": new_sliding_window_spec(sliding_window=1024),
4206	"layer_2": new_chunked_local_attention_spec(attention_chunk_size=512),
4207	}
4208
4209	with pytest.raises(ValueError):
4210	kv_cache_utils.unify_hybrid_kv_cache_specs(kv_cache_spec)
4211
4212	# Replicated Mamba state still requires the hybrid cache manager.
4213	kv_cache_spec = {"attention": new_kv_cache_spec(), "mamba": new_mamba_spec()}
4214	with pytest.raises(ValueError):
4215	kv_cache_utils.unify_hybrid_kv_cache_specs(kv_cache_spec)
4216
4217
4218	def test_unify_kv_cache_spec_page_size_mamba():
4219	"""Regression test for https://github.com/vllm-project/vllm/issues/43626.
4220
4221	MambaSpec's page_size_bytes is determined by its state shapes and does not
4222	change with block_size, so unify_kv_cache_spec_page_size must pad the Mamba
4223	page instead of scaling its block_size. This situation arises when a layer
4224	with a page larger than the (already platform-aligned) Mamba page joins the
4225	specs, e.g. a dense draft model with more KV heads than the hybrid main
4226	model.
4227	"""
4228	# 1. Hybrid main model (Mamba + full attention, pages already aligned at
4229	# 16KB) plus a dense draft model layer with a 2x larger page (32KB).
4230	# Reproduces the bare AssertionError from #43626: 32768 % 16384 == 0, so
4231	# the old code scaled the Mamba block_size, which left page_size_bytes
4232	# unchanged at 16384.
4233	mamba_spec = new_mamba_spec() # page_size_bytes = 16384
4234	main_attn_spec = new_kv_cache_spec() # page_size_bytes = 16384
4235	draft_attn_spec = new_kv_cache_spec(num_kv_heads=4) # page_size_bytes = 32768
4236	assert mamba_spec.page_size_bytes == main_attn_spec.page_size_bytes == 16384
4237	assert draft_attn_spec.page_size_bytes == 32768
4238
4239	unified = kv_cache_utils.unify_kv_cache_spec_page_size(
4240	{
4241	"mamba_layer": mamba_spec,
4242	"main_attn_layer": main_attn_spec,
4243	"draft_attn_layer": draft_attn_spec,
4244	}
4245	)
4246	# Mamba page is padded; block_size (caching granularity) is unchanged.
4247	assert unified["mamba_layer"].page_size_bytes == 32768
4248	assert unified["mamba_layer"].page_size_padded == 32768
4249	assert unified["mamba_layer"].block_size == mamba_spec.block_size
4250	# Attention layer with smaller page still unifies by scaling block_size.
4251	assert unified["main_attn_layer"].page_size_bytes == 32768
4252	assert unified["main_attn_layer"].block_size == 2 * main_attn_spec.block_size
4253	assert unified["main_attn_layer"].page_size_padded is None
4254	# Layer already at max page size is unchanged.
4255	assert unified["draft_attn_layer"] == draft_attn_spec
4256
4257	# 2. Mamba page already padded by the platform (state smaller than the
4258	# padded page); the padding is re-applied at the new maximum.
4259	padded_mamba_spec = new_mamba_spec(
4260	shapes=((2, 256), (3, 32, 32)), page_size_padded=16384
4261	)
4262	assert padded_mamba_spec.page_size_bytes == 16384
4263	unified = kv_cache_utils.unify_kv_cache_spec_page_size(
4264	{
4265	"mamba_layer": padded_mamba_spec,
4266	"draft_attn_layer": draft_attn_spec,
4267	}
4268	)
4269	assert unified["mamba_layer"].page_size_bytes == 32768
4270	assert unified["mamba_layer"].page_size_padded == 32768
4271
4272	# 3. Mamba page that does not evenly divide the maximum page size is
4273	# padded as well (the divisibility constraint only applies to block_size
4274	# scaling).
4275	odd_mamba_spec = new_mamba_spec(shapes=((6144,),))
4276	assert odd_mamba_spec.page_size_bytes == 24576
4277	assert 32768 % odd_mamba_spec.page_size_bytes != 0
4278	unified = kv_cache_utils.unify_kv_cache_spec_page_size(
4279	{
4280	"mamba_layer": odd_mamba_spec,
4281	"draft_attn_layer": draft_attn_spec,
4282	}
4283	)
4284	assert unified["mamba_layer"].page_size_bytes == 32768
4285
4286	# 4. Attention layers with non-divisible page sizes are padded too: every
4287	# backend reads a padded page through the view's block stride, so there is
4288	# no longer a case that must raise.
4289	unified = kv_cache_utils.unify_kv_cache_spec_page_size(
4290	{
4291	"attn_layer": new_kv_cache_spec(block_size=24), # 24576
4292	"draft_attn_layer": draft_attn_spec, # 32768
4293	}
4294	)
4295	assert unified["attn_layer"].page_size_padded == 32768
4296	assert unified["attn_layer"].page_size_bytes == 32768
4297	assert unified["attn_layer"].block_size == 24
4298
4299	# 5. Uniform page sizes are returned unchanged.
4300	specs = {
4301	"mamba_layer": new_mamba_spec(),
4302	"attn_layer": new_kv_cache_spec(),
4303	}
4304	assert kv_cache_utils.unify_kv_cache_spec_page_size(specs) == specs
4305
4306
4307	def test_hma_not_disabled_when_kv_events_enabled():
4308	"""Test enabling KV events must not force disable_hybrid_kv_cache_manager to True.
4309
4310	This test guards against that regression by verifying that a VllmConfig
4311	with kv_events_config set still resolves disable_hybrid_kv_cache_manager
4312	to False (i.e. HMA remains enabled) when no other condition requires it
4313	to be disabled.
4314	"""
4315	model_config = ModelConfig(max_model_len=16)
4316	kv_events_config = KVEventsConfig(
4317	enable_kv_cache_events=True,
4318	publisher="null",
4319	)
4320
4321	# Leave disable_hybrid_kv_cache_manager as None (the default) so that
4322	# VllmConfig.__post_init__ resolves it automatically.
4323	vllm_config = VllmConfig(
4324	model_config=model_config,
4325	kv_events_config=kv_events_config,
4326	)
4327
4328	assert vllm_config.scheduler_config.disable_hybrid_kv_cache_manager is False, (
4329	"kv_events_config must not force-disable the hybrid KV cache manager."
4330	)
4331
4332
4333	def test_resolve_block_hashes_gate():
4334	# Resolve symbols through the module so they stay consistent with each other
4335	# even after other tests reload ``kv_cache_utils``.
4336	resolve_block_hashes = kv_cache_utils.resolve_block_hashes
4337	BlockHashListWithBlockSize = kv_cache_utils.BlockHashListWithBlockSize
4338	# Raw, hash_block_size-granularity hashes (contents are opaque here).
4339	raw = [BlockHash(bytes([i])) for i in range(8)]
4340
4341	# block_size == hash_block_size: always reuse the raw hashes.
4342	assert resolve_block_hashes(raw, 2, 2, alignment_tokens=2) is raw
4343	assert (
4344	resolve_block_hashes(
4345	raw, 2, 2, supports_fine_grained_hash_lookup=True, alignment_tokens=2
4346	)
4347	is raw
4348	)
4349
4350	# Fine-grained manager, partial hits ON (alignment_tokens == hash_block_size
4351	# < block_size): keep raw hashes so the manager can scan at hash granularity.
4352	assert (
4353	resolve_block_hashes(
4354	raw, 2, 4, supports_fine_grained_hash_lookup=True, alignment_tokens=2
4355	)
4356	is raw
4357	)
4358
4359	# Fine-grained manager, partial hits OFF (alignment_tokens ==
4360	# scheduler_block_size >= block_size): must fall back to a block-size view,
4361	# exactly like the pre-refactor coordinator did.
4362	off = resolve_block_hashes(
4363	raw, 2, 4, supports_fine_grained_hash_lookup=True, alignment_tokens=4
4364	)
4365	assert isinstance(off, BlockHashListWithBlockSize)
4366	assert off.scale_factor == 2
4367
4368	# Non-fine-grained manager (e.g. sliding window): always a block-size view
4369	# when block_size != hash_block_size, regardless of alignment_tokens.
4370	swa = resolve_block_hashes(
4371	raw, 2, 4, supports_fine_grained_hash_lookup=False, alignment_tokens=2
4372	)
4373	assert isinstance(swa, BlockHashListWithBlockSize)
4374	assert swa.scale_factor == 2
4375
4376
4377	def test_resolve_block_hashes_rejects_mismatched_view():
4378	resolve_block_hashes = kv_cache_utils.resolve_block_hashes
4379	BlockHashListWithBlockSize = kv_cache_utils.BlockHashListWithBlockSize
4380	raw = [BlockHash(bytes([i])) for i in range(8)]
4381
4382	# A view built at this block_size is returned as-is (idempotent).
4383	view = BlockHashListWithBlockSize(raw, 2, 4)
4384	assert resolve_block_hashes(view, 2, 4) is view
4385
4386	# A view built at a different block_size must fail loudly rather than be
4387	# silently reinterpreted at the wrong granularity.
4388	mismatched = BlockHashListWithBlockSize(raw, 2, 8)
4389	with pytest.raises(AssertionError):
4390	resolve_block_hashes(mismatched, 2, 4)
4391
4392
4393	@pytest.mark.parametrize("use_override", [True, False])
4394	def test_kv_cache_reserves_null_block_for_max_model_len(use_override):
4395	"""A KV cache sized to exactly the blocks a max_model_len request needs must
4396	be rejected. BlockPool keeps one block as the null block, so only
4397	num_blocks - 1 are usable; accepting num_blocks == needed would leave the
4398	request unschedulable and hang the engine. Covers both the
4399	num_gpu_blocks_override path and the memory-derived block count.
4400	"""
4401	block_size = 16
4402	max_model_len = 512 # needs 512 / 16 = 32 blocks
4403	vllm_config = VllmConfig(model_config=ModelConfig(max_model_len=max_model_len))
4404	spec = new_kv_cache_spec(block_size=block_size)
4405
4406	# 32 blocks -> only 31 usable after the null block: one short -> reject.
4407	if use_override:
4408	# Ample raw memory; the override alone constrains the block count.
4409	vllm_config.cache_config.num_gpu_blocks_override = 32
4410	available_memory = [spec.page_size_bytes * 1024]
4411	else:
4412	available_memory = [spec.page_size_bytes * 32]
4413
4414	with pytest.raises(ValueError, match="max seq len"):
4415	get_kv_cache_configs(vllm_config, [{"layer1": spec}], available_memory)
4416
4417
4418	def test_auto_fit_max_model_len_reserves_null_block():
4419	"""Auto-fit (max_model_len=-1) must size max_model_len against usable
4420	blocks, not the total pool. With memory for exactly 64 blocks, one is the
4421	null block, so auto-fit must settle on 63 * block_size; picking the full
4422	pool would leave a max-length request unschedulable and hang the engine.
4423	"""
4424	block_size = 16
4425	model_config = ModelConfig(max_model_len=1024)
4426	model_config.original_max_model_len = -1
4427	vllm_config = VllmConfig(model_config=model_config)
4428	vllm_config.cache_config.kv_cache_layout = "LBNHC"
4429	spec = new_kv_cache_spec(block_size=block_size)
4430
4431	# Exactly the 1024 / 16 = 64 blocks a full-length request would need.
4432	available_memory = [spec.page_size_bytes * 64]
4433
4434	get_kv_cache_configs(vllm_config, [{"layer1": spec}], available_memory)
4435
4436	assert vllm_config.model_config.max_model_len == 63 * block_size
4437
4438
4439	def test_check_enough_kv_cache_memory_reserves_null_block():
4440	"""The public admission check must reject a KV cache sized to exactly the
4441	blocks a max_model_len request needs. BlockPool keeps one block as the
4442	null block, so only num_blocks - 1 are usable; accepting
4443	num_blocks == needed would leave the request unschedulable and hang the
4444	engine.
4445	"""
4446	block_size = 16
4447	max_model_len = 512 # needs 512 / 16 = 32 blocks
4448	vllm_config = VllmConfig(model_config=ModelConfig(max_model_len=max_model_len))
4449	spec = new_kv_cache_spec(block_size=block_size)
4450
4451	# 32 blocks -> only 31 usable after the null block: one short -> reject.
4452	with pytest.raises(ValueError, match="max seq len"):
4453	check_enough_kv_cache_memory(
4454	vllm_config, {"layer1": spec}, spec.page_size_bytes * 32
4455	)
4456
4457	# 33 blocks -> 32 usable after the null block -> accept.
4458	check_enough_kv_cache_memory(
4459	vllm_config, {"layer1": spec}, spec.page_size_bytes * 33
4460	)
4461
4462
4463	def test_is_full_attention_spec_unwraps_uniform_type_specs():
4464	"""``UniformTypeKVCacheSpecs`` is not a ``FullAttentionSpec``, so callers
4465	scanning groups with a bare isinstance miss DeepSeek-V4-shaped configs."""
4466	common = dict(num_kv_heads=1, head_size=1, dtype=torch.float32)
4467	full = FullAttentionSpec(block_size=4, **common)
4468	mla = MLAAttentionSpec(block_size=4, **common)
4469	swa = SlidingWindowMLASpec(block_size=4, sliding_window=8, **common)
4470
4471	def wrap(**specs):
4472	return UniformTypeKVCacheSpecs(block_size=4, kv_cache_specs=specs)
4473
4474	assert is_full_attention_spec(full)
4475	assert is_full_attention_spec(mla)
4476	assert not is_full_attention_spec(swa)
4477
4478	assert is_full_attention_spec(wrap(a=mla, b=mla))
4479	assert not is_full_attention_spec(wrap(a=swa, b=swa))
4480
4481	# Every layer must qualify: a group holding a recycling sliding-window layer
4482	# has no stable slot layout. Grouping forbids mixing today, but the wrapper
4483	# is also built directly (e.g. projecting groups onto a PP worker).
4484	assert not is_full_attention_spec(wrap(a=mla, b=swa))
4485	assert not is_full_attention_spec(wrap())
4486
4487
4488	def test_iter_layer_specs_returns_group_members():
4489	common = dict(num_kv_heads=1, head_size=1, dtype=torch.float32)
4490	full = FullAttentionSpec(block_size=4, **common)
4491	mla = MLAAttentionSpec(block_size=4, **common)
4492
4493	assert list(iter_layer_specs(full)) == [full]
4494	wrapped = UniformTypeKVCacheSpecs(
4495	block_size=4, kv_cache_specs={"a": full, "b": mla}
4496	)
4497	assert list(iter_layer_specs(wrapped)) == [full, mla]
4498
4499
4500	def test_wrapped_mamba_group_requires_block_zeroing():
4501	mamba = MambaSpec(
4502	block_size=4,
4503	shapes=((4, 1),),
4504	dtypes=(torch.float32,),
4505	)
4506	wrapped = UniformTypeKVCacheSpecs(block_size=4, kv_cache_specs={"mamba": mamba})
4507	config = KVCacheConfig(
4508	num_blocks=4,
4509	kv_cache_tensors=[],
4510	kv_cache_groups=[KVCacheGroupSpec(["mamba"], wrapped)],
4511	)
4512
4513	assert config.needs_kv_cache_zeroing
4514
4515
4516	def _spec_decode_grouping_config(method="dspark", model_type=None):
4517	"""Grouping config with an EAGLE-family speculative method enabled."""
4518	return SimpleNamespace(
4519	scheduler_config=SimpleNamespace(disable_hybrid_kv_cache_manager=False),
4520	cache_config=SimpleNamespace(
4521	get_resolved_kv_cache_layout=lambda: SimpleNamespace(
4522	is_block_outermost=True
4523	)
4524	),
4525	model_config=SimpleNamespace(hf_config=SimpleNamespace(model_type=model_type)),
4526	speculative_config=SimpleNamespace(
4527	method=method,
4528	use_eagle=lambda: True,
4529	use_eagle_block_drop=lambda: True,
4530	),
4531	)
4532
4533
4534	def _hybrid_specs_with_draft(draft: bool, draft_shares_target_spec: bool = False):
4535	"""A K3-shaped hybrid: MLA full attention + Mamba, optionally plus a
4536	DSpark-style draft MLA layer marked non_causal_multi_token_decode.
4537
4538	The target's fp8 KV dtype is one difference from the draft on Kimi-K3
4539	(target `--kv-cache-dtype fp8_e4m3`, draft `auto`). Pass
4540	draft_shares_target_spec to match that dtype while retaining the draft
4541	marker as a group boundary.
4542	"""
4543	target_dtype = None if draft_shares_target_spec else "fp8_e4m3"
4544	specs = {
4545	"target.attn.0": new_mla_spec(block_size=64, cache_dtype_str=target_dtype),
4546	"target.attn.1": new_mla_spec(block_size=64, cache_dtype_str=target_dtype),
4547	"target.mamba.0": new_mamba_spec(block_size=64, mamba_cache_mode="align"),
4548	"target.mamba.1": new_mamba_spec(block_size=64, mamba_cache_mode="align"),
4549	}
4550	if draft:
4551	draft_spec = new_mla_spec(block_size=64)
4552	specs["draft.attn.0"] = replace(draft_spec, non_causal_multi_token_decode=True)
4553	return specs
4554
4555
4556	def test_draft_group_annotated_on_hybrid_general_path():
4557	# A drafter's MLA layer carries non_causal_multi_token_decode, so its group
4558	# is identifiable without keying off a model version.
4559	groups = get_kv_cache_groups(
4560	_spec_decode_grouping_config(), _hybrid_specs_with_draft(draft=True)
4561	)
4562
4563	flagged = [g for g in groups if g.is_eagle_group]
4564	assert len(flagged) == 1
4565	assert "draft.attn.0" in flagged[0].layer_names
4566
4567
4568	def test_mamba_groups_never_flagged_even_when_draft_shares_a_group():
4569	# Packed uniform-type groups can contain distinct target and draft layer
4570	# specs; the combined group still holds volatile draft KV and must be
4571	# flagged. What must never happen is a Mamba group being flagged: that
4572	# widens its lookup window to two consecutive chunks, which align-mode
4573	# checkpointing never produces, zeroing every lookup.
4574	groups = get_kv_cache_groups(
4575	_spec_decode_grouping_config(),
4576	_hybrid_specs_with_draft(draft=True, draft_shares_target_spec=True),
4577	)
4578
4579	for group in groups:
4580	if "draft.attn.0" in group.layer_names:
4581	assert group.is_eagle_group
4582	if isinstance(group.kv_cache_spec, MambaSpec):
4583	assert not group.is_eagle_group
4584
4585
4586	def test_draft_group_not_annotated_without_spec_decode():
4587	# The marker alone must not flag anything; the eagle semantics only apply
4588	# when a speculative method is actually enabled.
4589	config = _spec_decode_grouping_config()
4590	config.speculative_config = None
4591	groups = get_kv_cache_groups(config, _hybrid_specs_with_draft(draft=True))
4592
4593	assert not any(g.is_eagle_group for g in groups)
4594
4595
4596	def test_unidentifiable_draft_with_mamba_warns(caplog_vllm):
4597	# No group carries the draft marker, so consumers fall back to
4598	# conservative behavior that silently breaks reuse for Mamba groups.
4599	# That must at least be visible.
4600	groups = get_kv_cache_groups(
4601	_spec_decode_grouping_config(), _hybrid_specs_with_draft(draft=False)
4602	)
4603
4604	assert not any(g.is_eagle_group for g in groups)
4605	assert "no KV cache group could be identified as the draft model's" in (
4606	caplog_vllm.text
4607	)
4608
4609
4610	def test_unidentifiable_draft_without_mamba_does_not_warn(caplog_vllm):
4611	# Pure-attention models degrade gracefully under the consumers'
4612	# conservative fallback (a one-block hit drop at most), so the warning
4613	# stays silent to avoid noise on every unannotated EAGLE deployment.
4614	specs = {
4615	"target.attn.0": new_mla_spec(block_size=64),
4616	"target.attn.1": new_mla_spec(block_size=64),
4617	}
4618	groups = get_kv_cache_groups(_spec_decode_grouping_config(), specs)
4619
4620	assert not any(g.is_eagle_group for g in groups)
4621	assert "could be identified as the draft model's" not in caplog_vllm.text
4622
4623
4624	def test_no_warning_when_draft_group_is_identified(caplog_vllm):
4625	get_kv_cache_groups(
4626	_spec_decode_grouping_config(), _hybrid_specs_with_draft(draft=True)
4627	)
4628
4629	assert "could be identified as the draft model's" not in caplog_vllm.text
4630
4631
4632	def _deepseek_v4_specs(model_version="deepseek_v4"):
4633	"""DeepseekV4-shaped specs: full MLA layers plus sliding-window MLA layers
4634	at differing page sizes, with the MTP draft layer registered last."""
4635	return {
4636	"model.layers.0.self_attn.attn": new_mla_spec(),
4637	"model.layers.1.self_attn.attn": new_mla_spec(),
4638	"model.layers.2.self_attn.attn": new_swa_mla_spec(
4639	head_size=1024, model_version=model_version
4640	),
4641	# The MTP block registers last, and its sliding-window size differs, so
4642	# it lands in a group of its own.
4643	"model.layers.3.self_attn.attn": new_swa_mla_spec(
4644	head_size=1024, sliding_window=256, model_version=model_version
4645	),
4646	}
4647
4648
4649	@pytest.mark.parametrize(
4650	("method", "model_type"),
4651	[("mtp", "deepseek_v4"), ("dspark", "deepseek_v4"), ("dspark", "deepseek_v41")],
4652	)
4653	def test_deepseek_v4_draft_group_annotated_on_packed_path(method, model_type):
4654	# DeepseekV4's MTP block reuses the target's decoder layer, so its spec
4655	# carries no draft marker and only the positional rule can find it. This
4656	# pins the pre-existing behaviour that the unified annotator must preserve.
4657	groups = get_kv_cache_groups(
4658	_spec_decode_grouping_config(method=method, model_type=model_type),
4659	_deepseek_v4_specs(model_version=None),
4660	)
4661
4662	flagged = [g for g in groups if g.is_eagle_group]
4663	assert len(flagged) == 1
4664	assert "model.layers.3.self_attn.attn" in flagged[0].layer_names
4665
4666
4667	def test_trailing_layer_fallback_applies_to_any_mtp_model():
4668	# The positional rule is sound for every MTP drafter, not just DeepseekV4:
4669	# MTP blocks reuse the target's decoder layer (no spec marker) and always
4670	# register after every target layer. The model_type must not gate it.
4671	groups = get_kv_cache_groups(
4672	_spec_decode_grouping_config(method="mtp", model_type="other"),
4673	_deepseek_v4_specs(),
4674	)
4675
4676	flagged = [g for g in groups if g.is_eagle_group]
4677	assert len(flagged) == 1
4678	assert "model.layers.3.self_attn.attn" in flagged[0].layer_names
4679
4680
4681	def _qwen3_5_hybrid_specs(with_mtp_layer: bool):
4682	"""Qwen3.5-shaped hybrid: repeating [GDN x3, full-attn x1] blocks, with
4683	the MTP drafter's full-attn layer (spec-identical to the target's)
4684	registered last."""
4685	specs = {}
4686	idx = 0
4687	for _ in range(2):
4688	for _ in range(3):
4689	specs[f"model.layers.{idx}.linear_attn"] = new_mamba_spec(
4690	mamba_cache_mode="align"
4691	)
4692	idx += 1
4693	specs[f"model.layers.{idx}.self_attn.attn"] = new_kv_cache_spec()
4694	idx += 1
4695	if with_mtp_layer:
4696	specs["mtp.layers.0.self_attn.attn"] = new_kv_cache_spec()
4697	return specs
4698
4699
4700	def test_qwen3_5_mtp_draft_group_annotated_on_hybrid_path(caplog_vllm):
4701	# A hybrid mamba + full-attention model with an MTP drafter that is
4702	# spec-indistinguishable from the target reaches the general multi-group
4703	# path. The trailing-layer rule must locate the draft group there so the
4704	# Mamba groups are not swept up by the flag-all consumer fallback.
4705	groups = get_kv_cache_groups(
4706	_spec_decode_grouping_config(method="mtp", model_type="qwen3_5"),
4707	_qwen3_5_hybrid_specs(with_mtp_layer=True),
4708	)
4709
4710	flagged = [g for g in groups if g.is_eagle_group]
4711	assert len(flagged) == 1
4712	assert "mtp.layers.0.self_attn.attn" in flagged[0].layer_names
4713	for group in groups:
4714	if any(
4715	isinstance(spec, MambaSpec)
4716	for spec in iter_layer_specs(group.kv_cache_spec)
4717	):
4718	assert not group.is_eagle_group
4719	assert "could be identified as the draft model's" not in caplog_vllm.text
4720
4721
4722	@pytest.mark.parametrize("method", ["eagle", "eagle3", "dspark"])
4723	def test_non_mtp_eagle_hybrid_still_warns(caplog_vllm, method):
4724	# Other drafters are not covered by the trailing-layer rule, so an
4725	# unidentifiable hybrid draft must still warn.
4726	groups = get_kv_cache_groups(
4727	_spec_decode_grouping_config(method=method, model_type="qwen3_5"),
4728	_qwen3_5_hybrid_specs(with_mtp_layer=True),
4729	)
4730
4731	assert not any(g.is_eagle_group for g in groups)
4732	assert "could be identified as the draft model's" in caplog_vllm.text
4733
4734
4735	def test_trailing_layer_fallback_requires_exact_partition():
4736	# If the groups do not partition the layers exactly (e.g. a caller that
4737	# dropped or duplicated layers), the positional rule is meaningless and
4738	# must not fire.
4739	from vllm.v1.core.kv_cache_utils import _annotate_eagle_groups
4740
4741	specs = _qwen3_5_hybrid_specs(with_mtp_layer=True)
4742	config = _spec_decode_grouping_config(method="mtp", model_type="qwen3_5")
4743	groups = get_kv_cache_groups(config, specs)
4744	for g in groups:
4745	g.is_eagle_group = False
4746	# Remove one layer from its group: no longer an exact partition.
4747	trimmed = [
4748	KVCacheGroupSpec(
4749	[n for n in g.layer_names if n != "model.layers.0.linear_attn"],
4750	g.kv_cache_spec,
4751	)
4752	for g in groups
4753	]
4754	_annotate_eagle_groups(config, specs, trimmed, use_trailing_layer_fallback=True)
4755
4756	assert not any(g.is_eagle_group for g in trimmed)
4757
4758
4759	_GQA_SPEC = FullAttentionSpec(
4760	block_size=16, num_kv_heads=8, head_size=64, dtype=torch.float32
4761	)
4762	_GQA_SWA_SPEC = SlidingWindowSpec(
4763	block_size=16, num_kv_heads=8, head_size=64, dtype=torch.float32, sliding_window=128
4764	)
4765
4766
4767	@pytest.mark.parametrize(
4768	"specs,expected",
4769	[
4770	pytest.param({"l.0": new_mla_spec(), "l.1": new_mla_spec()}, 4, id="mla"),
4771	pytest.param(
4772	{
4773	"l.0": new_mla_spec(),
4774	"l.1": new_mla_spec(),
4775	"s.0": new_swa_mla_spec(),
4776	"s.1": new_swa_mla_spec(),
4777	},
4778	4,
4779	id="mla-and-swa-mla",
4780	),
4781	pytest.param({"l.0": _GQA_SPEC, "l.1": _GQA_SPEC}, 1, id="gqa"),
4782	pytest.param(
4783	{"l.0": new_mla_spec(), "l.1": _GQA_SPEC},
4784	1,
4785	id="mla-and-gqa-uniform-group",
4786	),
4787	pytest.param(
4788	{
4789	"l.0": new_mla_spec(),
4790	"l.1": new_mla_spec(),
4791	"s.0": _GQA_SWA_SPEC,
4792	"s.1": _GQA_SWA_SPEC,
4793	},
4794	1,
4795	id="mla-and-gqa-swa",
4796	),
4797	],
4798	)
4799	def test_kv_tp_replicas(monkeypatch, specs, expected):
4800	"""Resolved per layer before scheduler flattening hides mixed groups."""
4801	from vllm.config import ParallelConfig
4802	from vllm.platforms import current_platform
4803
4804	monkeypatch.setattr(current_platform, "device_count", lambda: 4)
4805	parallel_config = ParallelConfig(tensor_parallel_size=4)
4806	vllm_config = VllmConfig(
4807	model_config=ModelConfig(max_model_len=16), parallel_config=parallel_config
4808	)
4809	vllm_config.cache_config.kv_cache_layout = "LBNHC"
4810	mem = sum(s.page_size_bytes for s in specs.values()) * 10
4811	configs = get_kv_cache_configs(vllm_config, [specs], [mem])
4812	assert configs[0].kv_tp_replicas == expected
4813	scheduler = generate_scheduler_kv_cache_config(configs)
4814	assert scheduler.kv_tp_replicas == expected
4815
4816
4817	@pytest.mark.parametrize(
4818	"specs,tp_size,dcp_size,expected",
4819	[
4820	pytest.param([new_mla_spec()], 8, 1, 8, id="mla"),
4821	pytest.param([new_mla_spec()], 8, 2, 1, id="mla-dcp"),
4822	pytest.param(
4823	[replace(new_mla_spec(), max_tp_shards=None)], 8, 1, 1, id="mla-unset"
4824	),
4825	pytest.param([replace(_GQA_SPEC, max_tp_shards=2)], 8, 1, 4, id="gqa-partial"),
4826	pytest.param([replace(_GQA_SPEC, max_tp_shards=8)], 4, 1, 1, id="gqa-sharded"),
4827	pytest.param(
4828	[new_mla_spec(), replace(_GQA_SWA_SPEC, max_tp_shards=4)],
4829	8,
4830	1,
4831	2,
4832	id="mla-and-gqa-partial",
4833	),
4834	pytest.param(
4835	[
4836	HiddenStateCacheSpec(
4837	block_size=16, num_kv_heads=1, head_size=64, dtype=torch.float32
4838	)
4839	],
4840	8,
4841	1,
4842	1,
4843	id="hidden-state",
4844	),
4845	],
4846	)
4847	def test_kv_cache_groups_tp_replicas(specs, tp_size, dcp_size, expected):
4848	"""Replicas are the gcd of each layer's tp_size // max_tp_shards."""
4849	groups = [KVCacheGroupSpec([f"l.{i}"], spec) for i, spec in enumerate(specs)]
4850	assert kv_cache_groups_tp_replicas(groups, tp_size, dcp_size) == expected