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

@tduncan / Nvidia-Optimization-Stack-VLLM

Code Commits Issues Pulls Cloud Agent CI/CD

main

py 2,964 lines 115 KB

Raw

1	# SPDX-License-Identifier: Apache-2.0
2	# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3	"""KV-Cache Utilities."""
4
5	import copy
6	import hashlib
7	import math
8	import os
9	from collections import defaultdict
10	from collections.abc import Callable, Iterable, Iterator, Sequence
11	from dataclasses import dataclass, field, replace
12	from functools import partial, reduce
13	from typing import TYPE_CHECKING, Any, NamedTuple, NewType, TypeAlias, cast, overload
14
15	from vllm import envs
16	from vllm.config import VllmConfig
17	from vllm.logger import init_logger
18	from vllm.multimodal.utils import get_mm_features_in_window
19	from vllm.utils.hashing import xxhash, xxhash_cbor
20	from vllm.utils.math_utils import cdiv, round_up
21	from vllm.utils.mem_utils import format_gib
22	from vllm.utils.torch_utils import get_dtype_size
23	from vllm.v1.hisparse.layout import (
24	get_hisparse_host_pool_bytes,
25	get_hisparse_kv_cache_config,
26	get_hisparse_kv_cache_groups,
27	)
28	from vllm.v1.kv_cache_interface import (
29	AttentionSpec,
30	ChunkedLocalAttentionSpec,
31	CircularBufferSpec,
32	FullAttentionSpec,
33	HiddenStateCacheSpec,
34	HiSparseHotSpec,
35	KpoolTailSpec,
36	KVCacheConfig,
37	KVCacheGroupSpec,
38	KVCacheLayout,
39	KVCacheSpec,
40	KVCacheTensor,
41	MambaSpec,
42	MLAAttentionSpec,
43	SlidingWindowMLASpec,
44	SlidingWindowSpec,
45	UniformTypeKVCacheSpecs,
46	compute_layout_strides,
47	iter_layer_specs,
48	replace_as,
49	)
50	from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry
51	from vllm.v1.request import Request
52	from vllm.v1.utils import tensor_data
53
54	if TYPE_CHECKING:
55	from vllm.v1.core.block_pool import BlockPool
56
57
58	# BlockHash represents the hash of a single KV-cache block used for
59	# prefix caching. Treating it as a distinct type from `bytes` helps
60	# catch accidental misuse when passing around raw byte strings.
61	BlockHash = NewType("BlockHash", bytes)
62
63	# `BlockHashWithGroupId` combines a `BlockHash` with its KV cache group ID.
64	# It is represented as raw bytes for compactness and efficiency. The helper
65	# functions below pack/unpack the `BlockHash` and group id into/from the key.
66	BlockHashWithGroupId = NewType("BlockHashWithGroupId", bytes)
67
68	# ExternalBlockHash is used for reproducible prefix-cache block hashing.
69	# It's a union of `bytes` and `int` to keep backward compatibility
70	# after we default block hashing to use sha256 bytes.
71	ExternalBlockHash: TypeAlias = bytes \| int
72
73
74	def make_block_hash_with_group_id(
75	block_hash: BlockHash, group_id: int
76	) -> BlockHashWithGroupId:
77	"""Pack a `BlockHash` and group id into a `BlockHashWithGroupId`.
78
79	The group id is encoded using 4 bytes in big-endian order and appended to
80	the block hash bytes. This representation avoids creating tuples while
81	still allowing us to recover both components when needed.
82	"""
83	return BlockHashWithGroupId(block_hash + group_id.to_bytes(4, "big", signed=False))
84
85
86	def get_block_hash(key: BlockHashWithGroupId) -> BlockHash:
87	"""Extract the `BlockHash` from a `BlockHashWithGroupId`."""
88	return BlockHash(key[:-4])
89
90
91	def get_group_id(key: BlockHashWithGroupId) -> int:
92	"""Extract the group id from a `BlockHashWithGroupId`."""
93	return int.from_bytes(key[-4:], "big", signed=False)
94
95
96	def maybe_convert_block_hash(hash_bytes: BlockHash) -> ExternalBlockHash:
97	if not envs.VLLM_KV_EVENTS_USE_INT_BLOCK_HASHES:
98	return hash_bytes
99	return int.from_bytes(hash_bytes, byteorder="big") & ((1 << 64) - 1)
100
101
102	logger = init_logger(__name__)
103
104	# The hash seed for the first block of any prefix block sequence.
105	#
106	# For cryptographic hash algorithms it is derived deterministically from a fixed
107	# default seed, so independent vLLM processes compute identical block hashes for
108	# identical content and can share a prefix cache (e.g. KV cache reuse across
109	# nodes) without extra configuration. This does not weaken collision resistance,
110	# which for SHA-256 does not depend on keeping the seed secret; ``cache_salt``
111	# remains the mechanism for intentional cache isolation.
112	#
113	# Non-cryptographic algorithms keep a per-process random seed, because a
114	# predictable seed would let an attacker precompute colliding blocks offline
115	# (see #12621). Setting PYTHONHASHSEED overrides the seed in both cases.
116	#
117	# The function `init_none_hash` initializes this variable globally.
118	NONE_HASH: BlockHash
119
120	# Fixed seed used when the PYTHONHASHSEED environment variable is not set and
121	# the hash algorithm is cryptographic.
122	DEFAULT_NONE_HASH_SEED = "vllm-none-hash"
123
124	# Algorithms that are not collision resistant, so the seed must stay secret.
125	_NON_CRYPTO_HASH_FUNCTIONS = frozenset({xxhash, xxhash_cbor})
126
127	# The seed NONE_HASH was derived from, set by init_none_hash.
128	_NONE_HASH_SEED: str \| None = None
129
130
131	def resolve_none_hash_seed(hash_fn: Callable[[Any], bytes]) -> str:
132	"""Resolve the seed to derive NONE_HASH from.
133
134	PYTHONHASHSEED wins if set. Otherwise cryptographic algorithms get the
135	fixed default (shareable across processes) and non-cryptographic ones get
136	fresh random bytes, keeping the seed unpredictable where collision
137	resistance depends on it.
138	"""
139	hash_seed = os.getenv("PYTHONHASHSEED")
140	if hash_seed is not None:
141	return hash_seed
142	if hash_fn in _NON_CRYPTO_HASH_FUNCTIONS:
143	return os.urandom(32).hex()
144	return DEFAULT_NONE_HASH_SEED
145
146
147	def get_none_hash_seed() -> str:
148	"""Return the seed NONE_HASH was derived from.
149
150	Components that must agree on NONE_HASH across processes (the P2P tier
151	advertises this during its connect handshake) read the resolved seed here
152	instead of re-deriving it, so they observe the random seed too. Falls back
153	to the deterministic seed before ``init_none_hash`` has run.
154	"""
155	if _NONE_HASH_SEED is None:
156	return DEFAULT_NONE_HASH_SEED
157	return _NONE_HASH_SEED
158
159
160	def init_none_hash(hash_fn: Callable[[Any], bytes]):
161	global NONE_HASH, _NONE_HASH_SEED
162
163	_NONE_HASH_SEED = resolve_none_hash_seed(hash_fn)
164	if hash_fn in _NON_CRYPTO_HASH_FUNCTIONS and os.getenv("PYTHONHASHSEED") is None:
165	logger.warning(
166	"Using a random per-process NONE_HASH seed because %s is not "
167	"collision resistant. Block hashes are therefore not reproducible "
168	"across processes; set PYTHONHASHSEED to a shared value to reuse "
169	"the prefix cache across instances, or use sha256.",
170	hash_fn.__name__,
171	)
172	NONE_HASH = BlockHash(hash_fn(_NONE_HASH_SEED))
173
174
175	@dataclass(slots=True)
176	class KVCacheBlock:
177	"""KV-cache block metadata."""
178
179	# Block ID, ranging from 0 to num_gpu_blocks - 1.
180	block_id: int
181	# Reference count.
182	ref_cnt: int = 0
183	# Deferred frees may contain blocks from multiple pools.
184	pool: "BlockPool \| None" = field(
185	default=None, repr=False, compare=False, kw_only=True
186	)
187	# The hash key (block hash + group id) of the block, only available
188	# when the block is full and cached.
189	_block_hash: BlockHashWithGroupId \| None = None
190	# Number of prefix tokens covered by _block_hash. For full blocks this is
191	# the full block boundary; partial entries can end inside a cache block.
192	_block_hash_num_tokens: int \| None = None
193
194	# Used to construct a doubly linked list for free blocks.
195	# These two attributes should only be manipulated by FreeKVCacheBlockQueue.
196	prev_free_block: "KVCacheBlock \| None" = None
197	next_free_block: "KVCacheBlock \| None" = None
198
199	# Whether the block is a null block that should never be cached.
200	is_null: bool = False
201
202	@property
203	def block_hash(self) -> BlockHashWithGroupId \| None:
204	return self._block_hash
205
206	@property
207	def block_hash_num_tokens(self) -> int \| None:
208	return self._block_hash_num_tokens
209
210	def set_block_hash(
211	self,
212	block_hash: BlockHashWithGroupId,
213	num_tokens: int \| None = None,
214	) -> None:
215	assert self.block_hash is None and self._block_hash_num_tokens is None, (
216	"The block already has a hash. This should not happen."
217	)
218	self._block_hash = block_hash
219	self._block_hash_num_tokens = num_tokens
220
221	def reset_hash(self):
222	"""Reset the block hash when the block is evicted."""
223	self._block_hash = None
224	self._block_hash_num_tokens = None
225
226	def __repr__(self) -> str:
227	# Use block_id instead of KVCacheBlock object to avoid calling __repr__
228	# on KVCacheBlock object recursively.
229	prev_block_id = self.prev_free_block.block_id if self.prev_free_block else None
230	next_block_id = self.next_free_block.block_id if self.next_free_block else None
231	return (
232	f"KVCacheBlock(block_id={self.block_id}, "
233	f"ref_cnt={self.ref_cnt}, "
234	f"_block_hash={self._block_hash!r}, "
235	f"_block_hash_num_tokens={self._block_hash_num_tokens}, "
236	f"prev_free_block={prev_block_id}, "
237	f"next_free_block={next_block_id})"
238	)
239
240
241	class KVCacheBlockCopy(NamedTuple):
242	src_block_id: int
243	dst_block_id: int
244
245
246	class FreeKVCacheBlockQueue:
247	"""This class organizes a list of KVCacheBlock objects to a doubly linked
248	list of free blocks. We implement this class instead of using Python
249	builtin deque to support removing a block in the middle of the queue
250	in O(1) time. To close the performance gap to the builtin deque which is
251	implemented in C++, this class does not allocate any Python objects when
252	manipulating the linked list. Instead, this class manipulates the
253	prev_free_block and next_free_block attributes of the given blocks.
254
255	The queue is ordered by block ID in the beginning. When a block is allocated
256	and then freed, it will be appended back with the eviction order:
257	1. The least recent used block is at the front (LRU).
258	2. If two blocks have the same last accessed time (allocated by the
259	same sequence), the one with more hash tokens (the tail of a block
260	chain) is at the front.
261	Note that we maintain this order by reversing the block order when free
262	blocks of a request. This operation is outside of this class.
263
264	Args:
265	blocks: A list of KVCacheBlock objects.
266
267	"""
268
269	def __init__(self, blocks: list[KVCacheBlock]) -> None:
270	self.num_free_blocks = len(blocks)
271
272	# Initialize doubly links of consecutive blocks
273	for i in range(self.num_free_blocks):
274	if i > 0:
275	blocks[i].prev_free_block = blocks[i - 1]
276	if i < self.num_free_blocks - 1:
277	blocks[i].next_free_block = blocks[i + 1]
278
279	# Create a fake head and a tail block for the doubly linked list to
280	# reduce branching in the code
281	#
282	# The implementation guaranteed that the fake head and tail
283	# are NEVER got popped, so we could safely assume each real blocks
284	# in the queue has prev and next blocks.
285	self.fake_free_list_head = KVCacheBlock(block_id=-1)
286	self.fake_free_list_tail = KVCacheBlock(block_id=-1)
287	if self.num_free_blocks > 0:
288	# Connect fake_head and fake_tail to the first and last block
289	# respectively.
290	self.fake_free_list_head.next_free_block = blocks[0]
291	blocks[0].prev_free_block = self.fake_free_list_head
292	self.fake_free_list_tail.prev_free_block = blocks[-1]
293	blocks[-1].next_free_block = self.fake_free_list_tail
294	else:
295	# For empty list, simply connect the fake head and tail.
296	self.fake_free_list_head.next_free_block = self.fake_free_list_tail
297	self.fake_free_list_tail.prev_free_block = self.fake_free_list_head
298
299	def popleft(self) -> KVCacheBlock:
300	"""Pop the first free block and reduce num_free_blocks by 1.
301
302	Returns:
303	The first free block.
304
305	"""
306	if (
307	self.fake_free_list_head.next_free_block is self.fake_free_list_tail
308	or self.fake_free_list_head.next_free_block is None
309	):
310	assert self.num_free_blocks == 0, (
311	f"num_free_blocks ({self.num_free_blocks}) is out of sync "
312	"with the free list."
313	)
314	raise ValueError("No free blocks available")
315
316	first_block: KVCacheBlock = self.fake_free_list_head.next_free_block
317
318	if first_block.next_free_block is None:
319	# This should not happen if the block is from the free list.
320	# It indicates a bug in the caller's logic.
321	raise RuntimeError(
322	"Invalid block found in popleft() "
323	"which doesn't have a valid next_free_block"
324	)
325
326	# Connect fake_head and the next block of first_block (i.e. second block
327	# or fake tail).
328	self.fake_free_list_head.next_free_block = first_block.next_free_block
329	first_block.next_free_block.prev_free_block = self.fake_free_list_head
330
331	# Remove the block from the linked list.
332	first_block.prev_free_block = first_block.next_free_block = None
333
334	self.num_free_blocks -= 1
335	return first_block
336
337	def popleft_n(self, n: int) -> list[KVCacheBlock]:
338	"""Pop the first n free blocks and reduce num_free_blocks by n.
339
340	Args:
341	n: The number of blocks to pop.
342
343	Returns:
344	A list of n free blocks.
345
346	"""
347	if n == 0:
348	return []
349	assert self.num_free_blocks >= n
350	self.num_free_blocks -= n
351
352	curr_block = self.fake_free_list_head.next_free_block
353	# Pop n blocks from the head of the list
354	ret = []
355	for _ in range(n):
356	assert curr_block is not None
357	ret.append(curr_block)
358	last_block = curr_block
359	curr_block = curr_block.next_free_block
360	# Reset prev_free_block and next_free_block of all popped blocks
361	last_block.prev_free_block = None
362	last_block.next_free_block = None
363
364	if curr_block is not None:
365	# The queue is not empty, connect the fake head to
366	# the new first block.
367	self.fake_free_list_head.next_free_block = curr_block
368	curr_block.prev_free_block = self.fake_free_list_head
369	return ret
370
371	def remove(self, block: KVCacheBlock) -> None:
372	"""Remove a block in the free list and reduce num_free_blocks by 1.
373
374	Args:
375	block: The block to remove.
376
377	"""
378	if block.prev_free_block is None or block.next_free_block is None:
379	# This should not happen if the block is from the free list.
380	# It indicates a bug in the caller's logic.
381	raise RuntimeError(f"remove() called on an invalid block: {block}")
382
383	# Link the previous block to the next block.
384	block.prev_free_block.next_free_block = block.next_free_block
385	# Link the next block to the previous block.
386	block.next_free_block.prev_free_block = block.prev_free_block
387
388	# Remove the block from the linked list.
389	block.prev_free_block = block.next_free_block = None
390	self.num_free_blocks -= 1
391
392	def append(self, block: KVCacheBlock) -> None:
393	"""Put a block back into the free list and increase
394	num_free_blocks by 1.
395
396	Args:
397	block: The block to append.
398
399	"""
400	if self.fake_free_list_tail.prev_free_block is None:
401	raise RuntimeError(
402	"prev_free_block of fake_free_list_tail should always exist"
403	)
404	last_block: KVCacheBlock = self.fake_free_list_tail.prev_free_block
405
406	# Connect the new block after the last block.
407	last_block.next_free_block = block
408	block.prev_free_block = last_block
409
410	# Connect the fake tail after the new block.
411	block.next_free_block = self.fake_free_list_tail
412	self.fake_free_list_tail.prev_free_block = block
413
414	self.num_free_blocks += 1
415
416	def prepend_n(self, blocks: list[KVCacheBlock]) -> None:
417	"""Put a list of blocks at the front of the free list."""
418	if len(blocks) == 0:
419	return
420
421	first_block = self.fake_free_list_head.next_free_block
422	assert first_block is not None, (
423	"next_free_block of fake_free_list_head should always exist"
424	)
425
426	prev_block = self.fake_free_list_head
427	for block in blocks:
428	block.prev_free_block = prev_block
429	prev_block.next_free_block = block
430	prev_block = block
431
432	prev_block.next_free_block = first_block
433	first_block.prev_free_block = prev_block
434
435	self.num_free_blocks += len(blocks)
436
437	def append_n(self, blocks: list[KVCacheBlock]) -> None:
438	"""Put a list of blocks back into the free list
439
440	Args:
441	blocks: The blocks to append.
442
443	"""
444	if len(blocks) == 0:
445	return
446
447	last_block = self.fake_free_list_tail.prev_free_block
448	assert last_block is not None, (
449	"prev_free_block of fake_free_list_tail should always exist"
450	)
451	# Add inter-connections between consecutive blocks
452	for block in blocks:
453	block.prev_free_block = last_block
454	last_block.next_free_block = block
455	last_block = block
456
457	# Connect the last block of <blocks> to the fake tail
458	last_block.next_free_block = self.fake_free_list_tail
459	self.fake_free_list_tail.prev_free_block = last_block
460
461	self.num_free_blocks += len(blocks)
462
463	def get_all_free_blocks(self) -> list[KVCacheBlock]:
464	"""Get all free blocks in the free list. Mainly used for testing.
465
466	Returns:
467	A list of free blocks.
468
469	"""
470	ret = []
471	if self.fake_free_list_head.next_free_block is None:
472	raise RuntimeError(
473	"next_free_block of fake_free_list_head should always exist"
474	)
475	# Start from the first block
476	curr_block: KVCacheBlock = self.fake_free_list_head.next_free_block
477	# As long as next_free_block is available, we haven't reached to
478	# the fake tail yet.
479	while curr_block.next_free_block is not None:
480	ret.append(curr_block)
481	curr_block = curr_block.next_free_block
482	return ret
483
484	def iter_blocks_after(
485	self,
486	cursor: KVCacheBlock \| None,
487	) -> Iterator[KVCacheBlock]:
488	"""Iterate free blocks in eviction order after the cursor."""
489	if cursor is None:
490	curr_block = self.fake_free_list_head.next_free_block
491	else:
492	curr_block = cursor.next_free_block
493
494	while curr_block is not None and curr_block is not self.fake_free_list_tail:
495	yield curr_block
496	curr_block = curr_block.next_free_block
497
498
499	def _gen_mm_extra_hash_keys(
500	request: Request, start_token_idx: int, end_token_idx: int, start_mm_idx: int
501	) -> tuple[list[Any], int]:
502	"""Generate extra keys related to MultiModal request for block hash
503	computation. For multi-modal inputs, the extra keys are
504	("mm", mm_hash, start_offset) tuples that indicate a mm input contained
505	in the block and its starting offset in the block tokens.
506
507	Args:
508	request: The request object.
509	start_token_idx: The start token index of the block.
510	end_token_idx: The end token index of the block.
511	start_mm_idx: The start multi-modal index of the block.
512
513	Returns:
514	A tuple of extra keys and the next multi-modal index.
515
516	"""
517	extra_keys: list[Any] = []
518
519	mm_features = request.mm_features
520	if not mm_features:
521	return extra_keys, start_mm_idx
522
523	# Note that we assume mm_features are sorted by mm_position.offset.
524	# We do not need to check all mm inputs if the start token index is out of
525	# range. This usually happens in the late prefill phase and decoding phase.
526	last_pos = mm_features[-1].mm_position
527	if last_pos.offset + last_pos.length <= start_token_idx:
528	return extra_keys, start_mm_idx
529
530	# Support start_mm_idx == -1 to indicate the last mm input.
531	if start_mm_idx < 0:
532	assert -start_mm_idx <= len(mm_features)
533	start_mm_idx = len(mm_features) + start_mm_idx
534
535	curr_mm_idx = start_mm_idx
536	while mm_features and curr_mm_idx < len(mm_features):
537	mm_feature = mm_features[curr_mm_idx]
538	assert mm_feature.identifier is not None
539	offset = mm_feature.mm_position.offset
540	length = mm_feature.mm_position.length
541	if end_token_idx > offset:
542	if start_token_idx >= offset + length:
543	# This block has passed the current mm input.
544	curr_mm_idx += 1
545	continue
546
547	# The block contains the current mm input. Include its offset
548	# relative to the start of the block so prefix-cache keys stay
549	# distinct when the same MM item appears at different positions
550	# within otherwise-identical placeholder blocks.
551	extra_keys.append(("mm", mm_feature.identifier, offset - start_token_idx))
552
553	if end_token_idx >= offset + length:
554	# If this block contains the end of the current mm input,
555	# move to the next mm input as this block may also contain
556	# the next mm input.
557	curr_mm_idx += 1
558	else:
559	# Otherwise this block is done with mm inputs.
560	break
561	else:
562	# This block has not reached the current mm input.
563	break
564	return extra_keys, curr_mm_idx
565
566
567	def _gen_lora_extra_hash_keys(request: Request) -> list[tuple[str, str, str]]:
568	"""Generate extra keys related to LoRA for block hash computation.
569
570	The adapter path is included so that re-pointing a LoRA name at a different
571	adapter does not reuse KV computed with the previous one.
572
573	Args:
574	request: The request object.
575
576	Returns:
577	Return the LoRA name and path of the request if it is a LoRA request.
578	Return empty list otherwise.
579
580	"""
581	lora_request = request.lora_request
582	if not lora_request:
583	return []
584	return [("lora", lora_request.lora_name, lora_request.lora_path)]
585
586
587	def _gen_prompt_embeds_extra_hash_keys(
588	request: Request, start_token_idx: int, end_token_idx: int
589	) -> list[tuple[str, bytes]]:
590	"""Generate extra keys related to prompt embeds for block hash computation.
591
592	Args:
593	request: The request object.
594	start_token_idx: The start token index of the block.
595	end_token_idx: The end token index of the block.
596
597	Returns:
598	Return a stable hash of the block prompt embeddings if prompt embeds
599	are present. Return empty list otherwise.
600
601	"""
602	if request.prompt_embeds is None:
603	return []
604	block_range = (start_token_idx, end_token_idx)
605	embeds_hash = request._prompt_embeds_per_block_hashes.get(block_range)
606	if embeds_hash is None:
607	block_prompt_embeds = request.prompt_embeds[start_token_idx:end_token_idx]
608	# Hash prompt embeds once per block and cache on request
609	embeds_hash = hashlib.sha256(tensor_data(block_prompt_embeds)).digest()
610	request._prompt_embeds_per_block_hashes[block_range] = embeds_hash
611	return [("prompt_embeds", embeds_hash)]
612
613
614	def generate_block_hash_extra_keys(
615	request: Request, start_token_idx: int, end_token_idx: int, start_mm_idx: int
616	) -> tuple[tuple[Any, ...] \| None, int]:
617	"""Generate extra keys for the block hash. The extra keys can come from
618	the multi-modal inputs, request specific metadata (e.g., LoRA names), and
619	hashed data from prompt embeddings.
620
621	Args:
622	request: The request object.
623	start_token_idx: The start token index of the block.
624	end_token_idx: The end token index of the block.
625	start_mm_idx: The start multi-modal index of the block.
626
627	Returns:
628	A tuple of extra keys and the next multi-modal index.
629
630	"""
631	mm_extra_keys: list[Any]
632	mm_extra_keys, new_start_mm_idx = _gen_mm_extra_hash_keys(
633	request, start_token_idx, end_token_idx, start_mm_idx
634	)
635	lora_extra_keys = _gen_lora_extra_hash_keys(request)
636	cache_salt_keys: list[tuple[str, str]] = (
637	[("cache_salt", request.cache_salt)]
638	if (start_token_idx == 0 and request.cache_salt)
639	else []
640	)
641	prompt_embeds_keys = _gen_prompt_embeds_extra_hash_keys(
642	request, start_token_idx, end_token_idx
643	)
644
645	extra_keys: list[Any] = (
646	lora_extra_keys + mm_extra_keys + cache_salt_keys + prompt_embeds_keys
647	)
648
649	if not extra_keys:
650	return None, new_start_mm_idx
651
652	return tuple(extra_keys), new_start_mm_idx
653
654
655	def to_event_extra_keys(
656	extra_keys: Iterable[tuple[Any, ...] \| None] \| None,
657	) -> list[tuple[Any, ...] \| None] \| None:
658	"""Convert block-hash extra keys to the untagged per-block list published
659	in KV events.
660
661	External KV event consumers parse the pre-tagging shapes: bare LoRA names
662	and cache salts, ``(mm_identifier, offset)`` pairs and bare prompt-embeds
663	digests. Events keep that format until consumers handle the tagged keys.
664
665	Args:
666	extra_keys: One entry per block, each as returned by
667	`generate_block_hash_extra_keys`, or None.
668
669	Returns:
670	One untagged entry per block, or None if there are no entries.
671
672	"""
673	if not extra_keys:
674	return None
675	event_keys = [
676	None
677	if keys is None
678	else tuple(key[1:] if key[0] == "mm" else key[1] for key in keys)
679	for keys in extra_keys
680	]
681	return event_keys or None
682
683
684	def hash_block_tokens(
685	hash_function: Callable[[Any], bytes],
686	parent_block_hash: BlockHash \| None,
687	curr_block_token_ids: Sequence[int],
688	extra_keys: tuple[Any, ...] \| None = None,
689	) -> BlockHash:
690	"""Computes a hash value corresponding to the contents of a block and
691	the contents of the preceding block(s). The hash value is used for
692	prefix caching. We use LRU cache for this function to avoid recomputing
693	hash values for the same block contents.
694
695	Args:
696	hash_function: The hash function used to compute block hash.
697	parent_block_hash: The hash of the parent block. None
698	if this is the first block.
699	curr_block_token_ids: A list of token ids in the current
700	block. The current block is assumed to be full.
701	extra_keys: Extra keys for the block.
702
703	Returns:
704	The hash value of the block and the token ids in the block.
705	The entire tuple is used as the hash key of the block.
706
707	"""
708	if not parent_block_hash:
709	parent_block_hash = NONE_HASH
710
711	curr_block_token_ids_tuple = tuple(curr_block_token_ids)
712	return BlockHash(
713	hash_function((parent_block_hash, curr_block_token_ids_tuple, extra_keys))
714	)
715
716
717	def resolve_dcp_kv_block_size(spec: KVCacheSpec, dcp_world_size: int) -> int:
718	"""Return the token span of a cache block under DCP."""
719	return spec.block_size * (dcp_world_size if spec.dcp_sharded else 1)
720
721
722	def resolve_dcp_kv_cache_spec(spec: KVCacheSpec, dcp_world_size: int) -> KVCacheSpec:
723	"""Return a KV cache spec with block sizes adjusted for DCP."""
724	block_size = resolve_dcp_kv_block_size(spec, dcp_world_size)
725	if block_size == spec.block_size:
726	return spec
727	if isinstance(spec, UniformTypeKVCacheSpecs):
728	return replace(
729	spec,
730	block_size=block_size,
731	kv_cache_specs={
732	name: resolve_dcp_kv_cache_spec(layer_spec, dcp_world_size)
733	for name, layer_spec in spec.kv_cache_specs.items()
734	},
735	)
736	return replace(spec, block_size=block_size)
737
738
739	def resolve_kv_cache_block_sizes(
740	kv_cache_config: KVCacheConfig,
741	vllm_config: VllmConfig,
742	) -> tuple[int, int]:
743	"""Resolve (scheduler_block_size, hash_block_size).
744
745	- ``scheduler_block_size`` is the token-alignment invariant used by the
746	scheduler (e.g. for ``num_computed_tokens`` rounding). Single group:
747	``cache_config.block_size * dcp``. Multiple groups: LCM of every
748	group's effective block size. Attention groups are scaled by DCP;
749	Mamba groups keep their full per-rank state and are not scaled.
750	- ``hash_block_size`` is the granularity at which ``Request.block_hashes``
751	is computed. Single group: equals scheduler block size, and any other
752	``cache_config.prefix_match_unit`` is rejected while block hashing is
753	active. Multiple groups:
754	``cache_config.prefix_match_unit`` override if set, else the GCD of
755	group block sizes; every group's block size must be divisible by it.
756	Returns the scheduler block size (i.e. disables finer hashing) if block
757	hashing is inactive or a mamba group is not using cache mode "align".
758	"""
759	cache_config = vllm_config.cache_config
760	dcp = vllm_config.parallel_config.decode_context_parallel_size
761	groups = kv_cache_config.kv_cache_groups
762
763	if len(groups) <= 1:
764	if groups and not groups[0].kv_cache_spec.dcp_sharded:
765	dcp = 1
766	bs = cache_config.block_size * dcp
767	# The Mamba prefill checkpoint builder reads prefix_match_unit directly,
768	# so a value dropped here puts its checkpoints off the scheduler's grid.
769	if (
770	cache_config.prefix_match_unit not in (None, bs)
771	and len(groups) == 1
772	and (
773	cache_config.enable_prefix_caching
774	or vllm_config.kv_transfer_config is not None
775	)
776	):
777	raise ValueError(
778	f"Invalid prefix_match_unit={cache_config.prefix_match_unit}; "
779	"with a single KV cache group, prefix-cache hits land on the "
780	f"block size ({bs}). Unset it or set it to {bs}."
781	)
782	return bs, bs
783
784	group_block_sizes = [
785	resolve_dcp_kv_block_size(g.kv_cache_spec, dcp) for g in groups
786	]
787	scheduler_block_size = math.lcm(*group_block_sizes)
788
789	logger.info("kv cache group sizes %s", group_block_sizes)
790	logger.info("kv lcm block sizes %s", scheduler_block_size)
791
792	# Block hashes are only consumed by prefix caching and KV connectors
793	# (P/D, offloading); when neither is active, keep hash_block_size equal
794	# to the scheduler block size.
795	connector_enabled = vllm_config.kv_transfer_config is not None
796	if not (cache_config.enable_prefix_caching or connector_enabled):
797	return scheduler_block_size, scheduler_block_size
798
799	# Mamba groups outside align mode break divisibility; back off to the
800	# scheduler block size. Read the mode from the resolved group spec because
801	# its block size may have been updated independently of cache_config.
802	if any(
803	isinstance(spec, MambaSpec) and spec.mamba_cache_mode != "align"
804	for group in groups
805	for spec in iter_layer_specs(group.kv_cache_spec)
806	):
807	return scheduler_block_size, scheduler_block_size
808
809	hashing_sizes = [
810	block_size
811	for group, block_size in zip(groups, group_block_sizes)
812	if group.kv_cache_spec.prefix_cacheable
813	] or group_block_sizes
814	requested = cache_config.prefix_match_unit
815	hash_block_size = requested if requested is not None else math.gcd(*hashing_sizes)
816	if any(bs % hash_block_size != 0 for bs in hashing_sizes):
817	raise ValueError(
818	f"Invalid prefix_match_unit={hash_block_size}; prefix-cacheable "
819	"KV cache group block sizes must be divisible by prefix_match_unit. "
820	f"Got group block sizes={group_block_sizes}, "
821	f"prefix-cacheable={hashing_sizes}."
822	)
823	prefix_alignments = {
824	spec.tokens_per_state
825	for group in groups
826	for spec in iter_layer_specs(group.kv_cache_spec)
827	if spec.prefix_cacheable
828	and isinstance(spec.tokens_per_state, int)
829	and spec.tokens_per_state > 1
830	}
831	has_partial_mamba_group = any(
832	isinstance(spec, MambaSpec)
833	and spec.mamba_cache_mode == "align"
834	and (
835	(dcp == 1 and block_size > hash_block_size)
836	or (dcp > 1 and block_size >= hash_block_size)
837	)
838	for group, block_size in zip(groups, group_block_sizes)
839	for spec in iter_layer_specs(group.kv_cache_spec)
840	)
841	cache_hit_alignment = (
842	hash_block_size if has_partial_mamba_group else scheduler_block_size
843	)
844	if any(cache_hit_alignment % alignment for alignment in prefix_alignments):
845	raise ValueError(
846	f"Invalid prefix_match_unit={hash_block_size}; prefix-cache boundaries "
847	"must align with each spec's per-state compression. "
848	f"Got alignments={sorted(prefix_alignments)}."
849	)
850	return scheduler_block_size, hash_block_size
851
852
853	def get_request_block_hasher(
854	hash_block_size: int,
855	caching_hash_fn: Callable[[Any], bytes],
856	) -> Callable[[Request], list[BlockHash]]:
857	"""Returns a function which computes the list of un-computed block hashes
858	of a request.
859
860	Hashes are computed at ``hash_block_size`` granularity and chained over the
861	full prefix, so each hash uniquely fingerprints the prefix ending at its
862	boundary. Coarser group block sizes and partial-cache boundaries reuse
863	these hashes directly (see ``BlockHashListWithBlockSize``).
864	"""
865
866	def request_block_hasher(request: Request) -> list[BlockHash]:
867	start_token_idx = len(request.block_hashes) * hash_block_size
868	num_tokens = request.num_tokens
869
870	if start_token_idx + hash_block_size > num_tokens:
871	# Early stop when there no new full blocks created.
872	return []
873
874	curr_mm_idx = 0
875	mm_features = request.mm_features
876	if start_token_idx > 0 and mm_features:
877	last_mm_pos = mm_features[-1].mm_position
878	if last_mm_pos.offset + last_mm_pos.length > start_token_idx:
879	curr_mm_idx, _ = get_mm_features_in_window(
880	mm_features,
881	start_token_idx,
882	start_token_idx + hash_block_size,
883	)
884
885	prev_block_hash_value = (
886	request.block_hashes[-1] if request.block_hashes else None
887	)
888	new_block_hashes: list[BlockHash] = []
889	while True:
890	end_token_idx = start_token_idx + hash_block_size
891	if end_token_idx > num_tokens:
892	# We only hash full blocks
893	break
894
895	# MM and LoRA requests need extra keys for block-hash computation.
896	extra_keys, curr_mm_idx = generate_block_hash_extra_keys(
897	request, start_token_idx, end_token_idx, curr_mm_idx
898	)
899
900	# Compute the hash of the current block
901	block_tokens = request.all_token_ids[start_token_idx:end_token_idx]
902	block_hash = hash_block_tokens(
903	caching_hash_fn, prev_block_hash_value, block_tokens, extra_keys
904	)
905
906	new_block_hashes.append(block_hash)
907	start_token_idx += hash_block_size
908	prev_block_hash_value = block_hash
909
910	return new_block_hashes
911
912	return request_block_hasher
913
914
915	def _check_enough_kv_cache_memory(
916	available_memory: int,
917	get_needed_memory: Callable[[], int],
918	max_model_len: int,
919	estimate_max_model_len: Callable[[int], int],
920	):
921	if available_memory <= 0:
922	raise ValueError(
923	"No available memory for the cache blocks. "
924	"Try increasing `gpu_memory_utilization` when initializing the engine "
925	"(this flag also controls CPU memory reservation on the CPU "
926	"backend, despite its name). "
927	"See https://docs.vllm.ai/en/latest/configuration/conserving_memory/ "
928	"for more details."
929	)
930
931	needed_memory = get_needed_memory()
932
933	if needed_memory > available_memory:
934	estimated_max_len = estimate_max_model_len(available_memory)
935	estimated_msg = ""
936	if estimated_max_len > 0:
937	estimated_msg = (
938	"Based on the available memory, "
939	f"the estimated maximum model length is {estimated_max_len}. "
940	)
941
942	raise ValueError(
943	f"To serve at least one request with the model's max seq len "
944	f"({max_model_len}), ({format_gib(needed_memory)} GiB KV "
945	f"cache is needed, which is larger than the available KV cache "
946	f"memory ({format_gib(available_memory)} GiB). {estimated_msg}"
947	f"Try increasing `gpu_memory_utilization` (which also controls "
948	f"CPU memory on the CPU backend) or decreasing `max_model_len` "
949	f"when initializing the engine. "
950	f"See https://docs.vllm.ai/en/latest/configuration/conserving_memory/ "
951	f"for more details."
952	)
953
954
955	def max_memory_usage_bytes(
956	vllm_config: VllmConfig, kv_cache_specs: Iterable[KVCacheSpec]
957	) -> int:
958	"""Get the maximum memory usage in bytes for the given KV cache specs."""
959	return sum(spec.max_memory_usage_bytes(vllm_config) for spec in kv_cache_specs)
960
961
962	def estimate_max_model_len(
963	vllm_config: VllmConfig,
964	kv_cache_spec: dict[str, KVCacheSpec],
965	available_memory: int,
966	) -> int:
967	"""Estimates the maximum model length that can fit in the available memory
968	using binary search.
969
970	This function temporarily modifies max_model_len during estimation but
971	restores the original value before returning, ensuring no side effects.
972
973	Args:
974	vllm_config: The global VllmConfig
975	kv_cache_spec: The kv cache spec of each attention layer in the model
976	available_memory: Memory available for KV cache in bytes.
977
978	Returns:
979	The estimated maximum model length that can fit in the available memory.
980
981	"""
982	# Save the original max_model_len to restore after estimation
983	original_max_model_len = vllm_config.model_config.max_model_len
984
985	# Define a function to check if a given model length fits in memory
986	def fits_in_memory(model_len: int) -> bool:
987	# Temporarily modify the max_model_len for this calculation
988	vllm_config.model_config.max_model_len = model_len
989	# Calculate memory needed for the given model length
990	memory_needed = max_memory_usage_bytes(vllm_config, kv_cache_spec.values())
991	return memory_needed <= available_memory
992
993	try:
994	# Binary search for the maximum model length
995	left, right = 1, original_max_model_len
996
997	# If even the smallest model length doesn't fit, return 0
998	if not fits_in_memory(left):
999	return 0
1000
1001	# Binary search for the maximum model length that fits
1002	result = 1
1003	while left <= right:
1004	mid = (left + right) // 2
1005	if fits_in_memory(mid):
1006	result = mid
1007	left = mid + 1
1008	else:
1009	right = mid - 1
1010	return result
1011	finally:
1012	# Always restore the original max_model_len to avoid side effects
1013	vllm_config.model_config.max_model_len = original_max_model_len
1014
1015
1016	def check_enough_kv_cache_memory(
1017	vllm_config: VllmConfig,
1018	kv_cache_spec: dict[str, KVCacheSpec],
1019	available_memory: int,
1020	):
1021	"""Checks whether `available_memory` is enough for the KV cache to hold at
1022	least one request with the model's max_model_len.
1023
1024	Args:
1025	vllm_config: The global VllmConfig
1026	kv_cache_spec: The kv cache spec of each attention layer in the model
1027	available_memory: Memory available for KV cache in bytes.
1028
1029	Raises:
1030	ValueError: If there is not enough memory available for the KV cache.
1031
1032	"""
1033	# No need to check for available memory if the kv_cache_spec is empty
1034	if kv_cache_spec:
1035	# Reserve the null block BlockPool permanently holds back, so the check
1036	# plans against usable blocks, as in get_kv_cache_configs. Group a copy
1037	# of the specs since grouping may unify them in-place.
1038	groups = get_kv_cache_groups(vllm_config, dict(kv_cache_spec))
1039	check_memory = (
1040	available_memory - _pool_bytes_per_block(groups)
1041	if groups
1042	else available_memory
1043	)
1044	_check_enough_kv_cache_memory(
1045	check_memory,
1046	lambda: max_memory_usage_bytes(vllm_config, kv_cache_spec.values()),
1047	vllm_config.model_config.max_model_len,
1048	lambda am: estimate_max_model_len(vllm_config, kv_cache_spec, am),
1049	)
1050
1051
1052	def create_kv_cache_group_specs(
1053	kv_cache_spec: dict[str, KVCacheSpec], grouped_layer_names: list[list[str]]
1054	) -> list[KVCacheGroupSpec]:
1055	"""Create KVCacheGroupSpec object for each kv cache group layer.
1056	The layers in the same group should share the same
1057	KVCacheSpec.
1058
1059	Args:
1060	kv_cache_spec:
1061	A mapping from each layer name to its corresponding KVCacheSpec.
1062	grouped_layer_names:
1063	A list of kv cache groups, where each element is a list of layer
1064	names that belong to the same group and should share the same
1065	KVCacheSpec.
1066
1067	Returns:
1068	A list of KVCacheGroupSpec objects, one for each group.
1069
1070	"""
1071	kv_cache_groups = []
1072	for layer_names_one_group in grouped_layer_names:
1073	layer_specs = [
1074	kv_cache_spec[layer_name] for layer_name in layer_names_one_group
1075	]
1076	merged_layer_spec = layer_specs[0].merge(layer_specs)
1077	kv_cache_groups.append(
1078	KVCacheGroupSpec(layer_names_one_group, merged_layer_spec)
1079	)
1080	return kv_cache_groups
1081
1082
1083	def is_kv_cache_spec_uniform(kv_cache_spec: dict[str, KVCacheSpec]) -> bool:
1084	"""Whether all layers in the given KVCacheSpec have the same KV cache spec.
1085	Note that we regard FullAttentionSpec with and without sliding window as
1086	the same type.
1087
1088	Args:
1089	kv_cache_spec: The kv cache spec of each attention layer in the model
1090
1091	Returns:
1092	True if all layers have the same type, False otherwise.
1093
1094	"""
1095	if not kv_cache_spec:
1096	# Encoder-only models do not have KV cache, kv_cache_type can be
1097	# regarded as uniform.
1098	return True
1099	try:
1100	kv_cache_spec_values = list(kv_cache_spec.values())
1101	_ = kv_cache_spec_values[0].merge(kv_cache_spec_values)
1102	except AssertionError:
1103	return False
1104	return True
1105
1106
1107	def get_max_concurrency_for_kv_cache_config(
1108	vllm_config: VllmConfig, kv_cache_config: KVCacheConfig
1109	) -> float:
1110	"""Get the maximum concurrency for the given KV cache configuration.
1111
1112	A request at max_model_len consumes whole blocks from each group's block
1113	table — cdiv(per-request bytes, page bytes) of the group's spec — and all
1114	device groups draw those block ids from one shared pool, so the per-request
1115	total is the sum over groups. The memory/page ratio is identical whether
1116	a group carries an aggregated UniformTypeKVCacheSpecs (worker config) or
1117	a representative per-layer spec (scheduler config), so both capacity
1118	call sites agree.
1119
1120	Host groups use a separate pool; the smaller concurrency limit applies.
1121	"""
1122	num_blocks_per_request = 0
1123	host_blocks_per_request = 0
1124	for group in kv_cache_config.kv_cache_groups:
1125	required = cdiv(
1126	group.kv_cache_spec.max_memory_usage_bytes(vllm_config),
1127	group.kv_cache_spec.page_size_bytes,
1128	)
1129	if group.host_resident:
1130	host_blocks_per_request += required
1131	else:
1132	num_blocks_per_request += required
1133	limits = [kv_cache_config.num_blocks / num_blocks_per_request]
1134	if host_blocks_per_request:
1135	assert kv_cache_config.hisparse_host_num_blocks is not None
1136	limits.append(
1137	kv_cache_config.hisparse_host_num_blocks / host_blocks_per_request
1138	)
1139	return min(limits)
1140
1141
1142	def may_override_num_blocks(vllm_config: VllmConfig, num_blocks: int) -> int:
1143	"""Override the number of kv cache blocks if `num_gpu_blocks_override` is set.
1144	The override is logged once, at the call site in `get_kv_cache_configs`.
1145	"""
1146	if vllm_config.cache_config.num_gpu_blocks_override is not None:
1147	num_blocks = vllm_config.cache_config.num_gpu_blocks_override
1148	return num_blocks
1149
1150
1151	def _pool_bytes_per_block(kv_cache_groups: list[KVCacheGroupSpec]) -> int:
1152	"""Bytes consumed by one block in the worker's shared KV cache pool, mirroring
1153	the divisor used by `get_kv_cache_config_from_groups` to convert
1154	`available_memory` into `num_blocks`. Used to compute the effective KV cache
1155	capacity once `num_gpu_blocks_override` is applied.
1156	"""
1157	return _get_kv_cache_bytes_per_block(
1158	[group for group in kv_cache_groups if not group.host_resident]
1159	)
1160
1161
1162	def get_uniform_page_size(kv_cache_specs: Iterable[KVCacheSpec]) -> int:
1163	"""Get the page size of the KV cache."""
1164	page_sizes = {layer.page_size_bytes for layer in kv_cache_specs}
1165	assert len(page_sizes) == 1
1166	return page_sizes.pop()
1167
1168
1169	def _get_kv_cache_groups_uniform_spec(
1170	kv_cache_specs: dict[str, KVCacheSpec],
1171	) -> list[KVCacheGroupSpec]:
1172	"""Generates the KV cache configuration for a model with the same KV cache
1173	spec for all layers.
1174
1175	Args:
1176	kv_cache_specs: The kv cache spec of each attention layer in the model
1177
1178	Returns:
1179	The generated KVCacheGroupSpecs
1180
1181	"""
1182	return create_kv_cache_group_specs(kv_cache_specs, [list(kv_cache_specs.keys())])
1183
1184
1185	def _get_kv_cache_groups_uniform_type(
1186	spec: UniformTypeKVCacheSpecs,
1187	) -> list[KVCacheGroupSpec]:
1188	"""Generates the KV cache configuration for a model with one type of KV cache
1189	but different hidden sizes. All layers are merged into one group.
1190
1191	Args:
1192	spec: The UniformTypeKVCacheSpecs of the model
1193
1194	Returns:
1195	The generated KVCacheGroupSpecs
1196
1197	"""
1198	return [KVCacheGroupSpec(list(spec.kv_cache_specs.keys()), spec)]
1199
1200
1201	def _pp_balanced_mamba_group_count(
1202	vllm_config: VllmConfig,
1203	mamba_layer_names: list[str],
1204	mla_layer_names: list[str],
1205	) -> int \| None:
1206	"""Return a Mamba group count whose PP projections fit the MLA slots."""
1207	num_groups = cdiv(len(mamba_layer_names), len(mla_layer_names))
1208	pp_size = vllm_config.parallel_config.pipeline_parallel_size
1209	if pp_size == 1:
1210	return num_groups
1211
1212	from vllm.distributed.utils import get_pp_indices
1213	from vllm.model_executor.models.utils import extract_layer_index
1214
1215	total_layers = vllm_config.model_config.get_total_num_hidden_layers()
1216	mamba_indices = [extract_layer_index(name) for name in mamba_layer_names]
1217	mla_indices = [extract_layer_index(name) for name in mla_layer_names]
1218	for rank in range(pp_size):
1219	start, end = get_pp_indices(total_layers, rank, pp_size)
1220	num_mamba = sum(start <= index < end for index in mamba_indices)
1221	num_mla = sum(start <= index < end for index in mla_indices)
1222	if not num_mamba:
1223	continue
1224	if not num_mla:
1225	return None
1226	num_groups = max(num_groups, cdiv(num_mamba, num_mla))
1227	return num_groups
1228
1229
1230	def _get_kv_cache_groups_glm5_next(
1231	vllm_config: VllmConfig,
1232	kv_cache_spec: dict[str, KVCacheSpec],
1233	) -> list[KVCacheGroupSpec] \| None:
1234	"""Build GLM-5.3-Flash groups with Mamba/MLA and tail/indexer aliasing."""
1235	mamba_specs = {
1236	name: spec
1237	for name, spec in kv_cache_spec.items()
1238	if isinstance(spec, MambaSpec)
1239	}
1240	tail_specs = {
1241	name: spec
1242	for name, spec in kv_cache_spec.items()
1243	if isinstance(spec, KpoolTailSpec)
1244	}
1245	attn_specs = {
1246	name: spec
1247	for name, spec in kv_cache_spec.items()
1248	if not isinstance(spec, (MambaSpec, KpoolTailSpec))
1249	}
1250	if not mamba_specs or not all(
1251	type(spec) is MLAAttentionSpec for spec in attn_specs.values()
1252	):
1253	return None
1254
1255	mla_specs = cast(dict[str, MLAAttentionSpec], attn_specs)
1256	idx_pages = {
1257	spec.page_size_bytes for spec in mla_specs.values() if spec.tokens_per_state > 1
1258	}
1259	if not idx_pages:
1260	return None
1261
1262	assert all(spec.page_size_padded is None for spec in mla_specs.values())
1263	assert len(idx_pages) == 1
1264	mla_names = [name for name, spec in mla_specs.items() if spec.tokens_per_state == 1]
1265	mla_pages = {mla_specs[name].page_size_bytes for name in mla_names}
1266	assert len(mla_pages) == 1
1267	mla_page = mla_pages.pop()
1268	uniform_spec = UniformTypeKVCacheSpecs.from_specs(attn_specs)
1269	assert uniform_spec is not None
1270
1271	tail_group: KVCacheGroupSpec \| None = None
1272	if tail_specs:
1273	idx_page = next(iter(idx_pages))
1274	padded_tail_specs: dict[str, KVCacheSpec] = {
1275	name: replace(spec, page_size_padded=idx_page)
1276	for name, spec in tail_specs.items()
1277	}
1278	tail_uniform = UniformTypeKVCacheSpecs.from_specs(padded_tail_specs)
1279	assert tail_uniform is not None
1280	tail_group = KVCacheGroupSpec(list(padded_tail_specs), tail_uniform)
1281
1282	any_mamba = next(iter(mamba_specs.values()))
1283	assert all(spec == any_mamba for spec in mamba_specs.values())
1284	if any_mamba.real_page_size_bytes > mla_page:
1285	raise ValueError(
1286	f"the mamba state page ({any_mamba.real_page_size_bytes} bytes) "
1287	f"does not fit the MLA page ({mla_page} bytes); increase tensor "
1288	"parallelism or use a wider KV cache dtype"
1289	)
1290	padded_specs: dict[str, KVCacheSpec] = {
1291	name: replace(any_mamba, page_size_padded=mla_page) for name in mamba_specs
1292	}
1293	num_groups = _pp_balanced_mamba_group_count(
1294	vllm_config, list(mamba_specs), mla_names
1295	)
1296	if num_groups is None:
1297	raise ValueError(
1298	"a pipeline stage has mamba layers but no MLA layer to share "
1299	"slots with; realign the stage boundaries (VLLM_PP_LAYER_PARTITION)"
1300	)
1301	mamba_grouped_names: list[list[str]] = [[] for _ in range(num_groups)]
1302	for index, name in enumerate(mamba_specs):
1303	mamba_grouped_names[index % num_groups].append(name)
1304
1305	return (
1306	[KVCacheGroupSpec(list(attn_specs), uniform_spec)]
1307	+ ([tail_group] if tail_group is not None else [])
1308	+ create_kv_cache_group_specs(padded_specs, mamba_grouped_names)
1309	)
1310
1311
1312	def _glm5_next_tensor_layout(
1313	kv_cache_groups: list[KVCacheGroupSpec],
1314	) -> (
1315	tuple[
1316	KVCacheGroupSpec,
1317	list[KVCacheGroupSpec],
1318	list[str],
1319	list[str],
1320	int,
1321	int,
1322	list[str],
1323	int,
1324	]
1325	\| None
1326	):
1327	"""Recognize the GLM-5.3-Flash grouping after optional PP projection."""
1328	uniform_groups = [
1329	group
1330	for group in kv_cache_groups
1331	if isinstance(group.kv_cache_spec, UniformTypeKVCacheSpecs)
1332	]
1333	mamba_groups = [
1334	group for group in kv_cache_groups if isinstance(group.kv_cache_spec, MambaSpec)
1335	]
1336	attn_group: KVCacheGroupSpec \| None = None
1337	tail_group: KVCacheGroupSpec \| None = None
1338	for group in uniform_groups:
1339	inner = cast(UniformTypeKVCacheSpecs, group.kv_cache_spec).kv_cache_specs
1340	if all(type(spec) is MLAAttentionSpec for spec in inner.values()):
1341	attn_group = group
1342	elif all(isinstance(spec, KpoolTailSpec) for spec in inner.values()):
1343	tail_group = group
1344	if attn_group is None or not mamba_groups:
1345	return None
1346	if len(uniform_groups) + len(mamba_groups) != len(kv_cache_groups):
1347	return None
1348
1349	attn_uniform = cast(UniformTypeKVCacheSpecs, attn_group.kv_cache_spec)
1350	mla_inner = cast(dict[str, MLAAttentionSpec], attn_uniform.kv_cache_specs)
1351	if not all(
1352	type(spec) is MLAAttentionSpec and spec.page_size_padded is None
1353	for spec in mla_inner.values()
1354	):
1355	return None
1356	mla_names = [
1357	name for name in attn_group.layer_names if mla_inner[name].tokens_per_state == 1
1358	]
1359	idx_names = [
1360	name for name in attn_group.layer_names if mla_inner[name].tokens_per_state > 1
1361	]
1362	mla_pages = {mla_inner[name].page_size_bytes for name in mla_names}
1363	idx_pages = {mla_inner[name].page_size_bytes for name in idx_names}
1364	if len(mla_pages) != 1 or len(idx_pages) != 1:
1365	return None
1366	mla_page = mla_pages.pop()
1367	idx_page = idx_pages.pop()
1368	if any(group.kv_cache_spec.page_size_bytes != mla_page for group in mamba_groups):
1369	return None
1370
1371	tail_names: list[str] = []
1372	tail_page = 0
1373	if tail_group is not None:
1374	tail_names = list(tail_group.layer_names)
1375	tail_inner = cast(
1376	UniformTypeKVCacheSpecs, tail_group.kv_cache_spec
1377	).kv_cache_specs
1378	tail_pages = {
1379	cast(KpoolTailSpec, spec).unpadded_page_size_bytes
1380	for spec in tail_inner.values()
1381	}
1382	if len(tail_pages) != 1 or len(tail_names) != len(idx_names):
1383	return None
1384	tail_page = tail_pages.pop()
1385	if tail_page > idx_page:
1386	return None
1387
1388	return (
1389	attn_group,
1390	mamba_groups,
1391	mla_names,
1392	idx_names,
1393	mla_page,
1394	idx_page,
1395	tail_names,
1396	tail_page,
1397	)
1398
1399
1400	def unify_kv_cache_spec_page_size(
1401	kv_cache_spec: dict[str, KVCacheSpec],
1402	) -> dict[str, KVCacheSpec]:
1403	"""Unify the page size of the given KVCacheSpec. If the page size of all layers
1404	are the same, return the original KVCacheSpec. If not same, unify the page
1405	size by increasing the block size of layers with smaller page size. Two
1406	cases cannot be unified by block size alone and pad their physical page to
1407	the maximum instead: Mamba layers, whose page size comes from state shapes
1408	and is independent of block size; and non-MLA attention layers whose page
1409	does not evenly divide the maximum (the padded page is read through a
1410	strided view). MLA is excluded because sparse MLA indexes the cache in
1411	whole token rows (see ``flat_kv_row_view``), so its block stride can only
1412	be padded by its own row-aligned ``alignment``, not to an arbitrary page
1413	size. Raise NotImplementedError if failed to unify the page size;
1414	``get_kv_cache_groups`` catches it to try the full-allocation fallback
1415	(e.g. MLA next to an incompatible sliding-window draft).
1416
1417	Args:
1418	kv_cache_spec: The KVCacheSpec of each attention layer in the model
1419
1420	Returns:
1421	The updated KVCacheSpec with the same page_size_bytes.
1422
1423	"""
1424	page_sizes = {layer.page_size_bytes for layer in kv_cache_spec.values()}
1425	if len(page_sizes) <= 1:
1426	# All layers have the same page size, no need to unify.
1427	return kv_cache_spec
1428
1429	max_page_size = max(page_sizes)
1430	new_kv_cache_spec = {}
1431	for layer_name, layer_spec in kv_cache_spec.items():
1432	if layer_spec.page_size_bytes == max_page_size:
1433	new_kv_cache_spec[layer_name] = layer_spec
1434	elif isinstance(layer_spec, MambaSpec):
1435	# MambaSpec's page size is determined by its state shapes and does
1436	# not scale with block_size, so pad the page instead. This is the
1437	# same padding mechanism the platform uses to align Mamba pages
1438	# with the main model's attention page size; it is needed here
1439	# when another layer (e.g. from a draft model) has a larger page
1440	# than the already-aligned Mamba page.
1441	new_spec: KVCacheSpec = replace(layer_spec, page_size_padded=max_page_size)
1442	assert new_spec.page_size_bytes == max_page_size
1443	new_kv_cache_spec[layer_name] = new_spec
1444	else:
1445	layer_page_size = layer_spec.page_size_bytes
1446	if max_page_size % layer_page_size == 0:
1447	ratio = max_page_size // layer_page_size
1448	new_block_size = layer_spec.block_size * ratio
1449	new_spec = replace(layer_spec, block_size=new_block_size)
1450	elif isinstance(layer_spec, AttentionSpec) and not isinstance(
1451	layer_spec, MLAAttentionSpec
1452	):
1453	new_spec = replace(layer_spec, page_size_padded=max_page_size)
1454	else:
1455	raise NotImplementedError(
1456	f"Layer {layer_name}: page size is not divisible by the "
1457	"maximum page size and cannot be padded. Padding is only "
1458	"supported for non-MLA attention layers."
1459	)
1460	assert new_spec.page_size_bytes == max_page_size
1461	new_kv_cache_spec[layer_name] = new_spec
1462	return new_kv_cache_spec
1463
1464
1465	def is_kv_cache_type_attention_free(kv_cache_spec: dict[str, KVCacheSpec]) -> bool:
1466	# kv_cache_spec is an empty dict for attention free models
1467	return not kv_cache_spec
1468
1469
1470	def _get_kv_cache_groups_uniform_page_size(
1471	kv_cache_spec: dict[str, KVCacheSpec],
1472	) -> list[KVCacheGroupSpec]:
1473	"""Generates the KV cache groups for hybrid models with multiple
1474	attention types but still with a uniform page size (physical memory per
1475	block per layer) for all layers.
1476
1477	Detailed explanation about kv cache management of hybrid models:
1478	The layers in the models are repeated with some patterns, e.g., a model
1479	with 10 full attention layers and 20 sliding window attention layers can be
1480	regarded as repeating the pattern (1 * full, 2 * sw) 10 times.
1481	The KVCacheManager allocates different block tables for each of the 3 layers
1482	in the pattern, and repeats each of them 10 times to generate the
1483	block_table for the 30 layers in the model.
1484	Therefore, we can group the layers in the model into 3 kv_cache_groups, each
1485	of which contains 10 layers in the model.
1486	The KVCacheManager allocates the block_table for each group based on its
1487	kv_cache spec, and the model runner applies the block table to each layer
1488	in the group.
1489	For example:
1490	1. A model only uses full attention. The pattern is
1491	(num_hidden_layers * full), so there is only one group and the block table
1492	is shared by all layers. It is already handled by
1493	`_get_kv_cache_config_uniform_type`.
1494	2. A model with 10 full attention layers and 20 sliding window
1495	attention layers. There are 3 layers in the pattern (1 * full, 2 * sw), so
1496	there are 3 kv_cache_groups, each of which represents 10 layers.
1497
1498	To simplify the implementation, we make the following assumptions:
1499	1. Physical memory per block: Must be the same across all KV cache groups.
1500	Breaking this assumption is non-trivial due to memory fragmentation concerns
1501	when allocating blocks of different sizes.
1502	2. Tokens per block (block_size): Currently, we directly use
1503	`CacheConfig.block_size` for all layers. It can be extended to vary by KV
1504	cache group, but within each KV cache group, all layers must share the same
1505	block size.
1506	3. Physical memory per token per layer: This property is decided by model
1507	config. Currently we only support models that have the same physical memory
1508	per token per layer for all layers. Can be relaxed with a simple extension,
1509	but still need to keep physical memory per block the same for all groups.
1510	4. Number of layers per group: Currently assumed the same for all layers.
1511	Can be relaxed with a simple extension, but still need to keep physical
1512	memory per block the same for all groups.
1513	5. Attention type within groups: All layers in a group must share the same
1514	attention type. One exception is that, when
1515	`--disable-hybrid-kv-cache-manager` is true, the single group for full
1516	attention layers may also include attention layers using sliding window or
1517	LLaMA 4 local attention. See `unify_hybrid_kv_cache_specs` for more details.
1518	6. Support for multiple attention types: The design for most components is
1519	general to an arbitrary number of attention types. But
1520	`find_longest_cache_hit` only supports one attention type or two
1521	types of full-attention plus exactly one another type. The general
1522	implementation of this function is feasible but we don't know how to
1523	implement it cleanly yet.
1524
1525	As we assume tokens per block, physical memory per token per layer, and
1526	number of layers per group are the same now, we can ensure that physical
1527	memory per block is the same for all groups.
1528
1529	Args:
1530	kv_cache_spec: The KVCacheSpec of each attention layer in the model
1531	Returns:
1532	The generated KVCacheGroupSpecs
1533
1534	"""
1535	# Group all layers by kv_cache_spec.
1536	# E.g., 2 full attention layers and 3 sliding window attention layers,
1537	# -> (full.0, full.1), (sw.0, sw.1, sw.2).
1538	same_type_layers: dict[KVCacheSpec, list[str]] = defaultdict(list)
1539	for layer_name, layer_spec in kv_cache_spec.items():
1540	same_type_layers[layer_spec].append(layer_name)
1541
1542	# Attempt to further merge same-type layers based on whether their KV
1543	# cache specs can be merged, to minimize the group count. This benefits
1544	# situations where specs share a block layout and differ only in a
1545	# property it can reconcile (e.g. full attention layers differing only in
1546	# sliding window / attention chunk size).
1547	layer_buckets: list[list[str]] = []
1548	spec_buckets: list[list[KVCacheSpec]] = []
1549	for layer_spec, layer_names in same_type_layers.items():
1550	for names, specs in zip(layer_buckets, spec_buckets):
1551	try:
1552	# A raise means that the specs are incompatible.
1553	type(specs[0]).merge([*specs, layer_spec])
1554	except (AssertionError, ValueError):
1555	continue
1556	names.extend(layer_names)
1557	specs.append(layer_spec)
1558	break
1559	else:
1560	layer_buckets.append(list(layer_names))
1561	spec_buckets.append([layer_spec])
1562
1563	# Split each group into smaller groups, to make the number of layers in each
1564	# group identical. Add padding to the last group of each type if necessary.
1565	# E.g., (full.0, full.1), (sw.0, sw.1, sw.2)
1566	# split to 3 groups with 2 layers each:
1567	# (full.0, full.1), (sw.0, sw.2), (sw.1, padding).
1568	# FIXME(Chen): At the moment of writing this code (2025-06-02), all
1569	# open-source hybrid model follows a n:1 pattern between different attention
1570	# types (e.g., Gemma3 5:1 between sw and full, LLaMA4 3:1 between local and
1571	# full), so we can use the "1" in the n:1 pattern as the group size, which
1572	# is the minimum number of layers among all attention types. Need a better
1573	# strategy if we want to support more complex patterns (e.g., 20 full + 30
1574	# sw, where the group size should be 10).
1575	min_num_layers = min([len(layers) for layers in layer_buckets])
1576	group_size = min_num_layers
1577	max_num_layers = max([len(layers) for layers in layer_buckets])
1578	if max_num_layers < min_num_layers * 1.5:
1579	# If the number of layers is not much larger than the minimum number of
1580	# layers, use the maximum number of layers as the group size to avoid
1581	# too many padding layers. A typical example is gpt-oss-20b + eagle,
1582	# with 12 sw + 13 full. We pad it to (13 sw, 13 full) instead of
1583	# (12 sw, 24 full). 1.5 is a heuristic to avoid too many padding
1584	# layers while accommodating speculative decoding drafters that add
1585	# extra layers to one attention type.
1586	group_size = max_num_layers
1587	grouped_layers = []
1588	for layers in layer_buckets:
1589	num_padding_layers = group_size - len(layers) % group_size
1590	if num_padding_layers != group_size:
1591	logger.warning(
1592	"Add %d padding layers, may waste at most %.2f%% KV cache memory", # noqa
1593	num_padding_layers,
1594	num_padding_layers / len(layers) * 100,
1595	)
1596	num_groups = cdiv(len(layers), group_size)
1597	# In PP case, say if we have
1598	# - stage 0: full.0, sw.0, sw.1
1599	# - stage 1: full.1, sw.2, sw.3
1600	# We should have 3 groups: (full.0, full.1), (sw.0, sw.2), (sw.1, sw.3)
1601	# It can't be (full.0, full.1), (sw.0, sw.1), (sw.2, sw.3) because
1602	# the 3 groups in stage 0 will be (full.0), (sw.0, sw.1), (empty group)
1603	# and it will be padded to (full.0, padding), (sw.0, sw.1),
1604	# (padding, padding) to ensure the number of layers in each group is
1605	# the same and will cause memory waste.
1606	# To avoid this, we assign layers[i::num_groups] to the i-th group
1607	# instead of layers[i * group_size: (i + 1) * group_size]
1608	for i in range(num_groups):
1609	grouped_layers.append(layers[i::num_groups])
1610	return create_kv_cache_group_specs(kv_cache_spec, grouped_layers)
1611
1612
1613	def _get_per_layer_spec(
1614	group: KVCacheGroupSpec,
1615	layer_name: str,
1616	) -> KVCacheSpec:
1617	spec = group.kv_cache_spec
1618	if isinstance(spec, UniformTypeKVCacheSpecs):
1619	return spec.kv_cache_specs[layer_name]
1620	return spec
1621
1622
1623	def _get_kv_cache_bytes_per_block(
1624	kv_cache_groups: list[KVCacheGroupSpec],
1625	) -> int:
1626	"""Return the largest cache group's bytes per block."""
1627	if (glm5_layout := _glm5_next_tensor_layout(kv_cache_groups)) is not None:
1628	_, _, mla_names, idx_names, mla_page, idx_page, _, _ = glm5_layout
1629	return len(mla_names) * mla_page + len(idx_names) * idx_page
1630
1631	bytes_per_block = max(
1632	sum(
1633	_get_per_layer_spec(group, layer_name).page_size_bytes
1634	for layer_name in group.layer_names
1635	)
1636	for group in kv_cache_groups
1637	)
1638	assert bytes_per_block > 0
1639	alignments = [
1640	group.kv_cache_spec.page_size_bytes
1641	for group in kv_cache_groups
1642	if isinstance(group.kv_cache_spec, HiSparseHotSpec)
1643	]
1644	alignments.extend(
1645	_get_per_layer_spec(group, layer_name).block_stride_alignment or 1
1646	for group in kv_cache_groups
1647	for layer_name in group.layer_names
1648	)
1649	return round_up(bytes_per_block, math.lcm(*alignments))
1650
1651
1652	def validate_kv_cache_layout(
1653	layout: KVCacheLayout,
1654	kv_cache_groups: list[KVCacheGroupSpec],
1655	) -> None:
1656	"""Validate that the resolved layout can express this model's packing.
1657
1658	The layout was chosen once in the engine core from the backends' supported
1659	sets; a backend whose model packs pages side by side (e.g. the DeepSeek-V4
1660	indexer) declares block-outermost layouts there, so an inexpressible
1661	layout reaching this point is an error.
1662	"""
1663	page_sizes = {
1664	_get_per_layer_spec(group, layer_name).page_size_bytes
1665	for group in kv_cache_groups
1666	for layer_name in group.layer_names
1667	}
1668	if len(page_sizes) == 1:
1669	# A rectangular layer dim exists; every layout can express it.
1670	return
1671
1672	# Mixed page sizes pack pages side by side within a block, which needs each page
1673	# to be one contiguous chunk inside its block (a block-compact layout) and, with
1674	# multiple KV cache groups, the layer dim inside the block dim.
1675	if not layout.is_block_compact or (
1676	len(kv_cache_groups) > 1 and layout.is_layer_compact
1677	):
1678	raise ValueError(
1679	f"KV cache layout {layout.name} cannot express this model's "
1680	f"mixed page sizes ({sorted(page_sizes)}); a backend should "
1681	"declare block-outermost supported layouts (e.g. BLHNC), or "
1682	"set VLLM_KV_CACHE_LAYOUT=BLHNC."
1683	)
1684
1685
1686	def get_kv_cache_config_from_groups(
1687	vllm_config: VllmConfig,
1688	kv_cache_groups: list[KVCacheGroupSpec],
1689	available_memory: int,
1690	) -> KVCacheConfig:
1691	"""Generate the KV cache configuration from the KV cache groups and spec
1692	of each layer.
1693
1694	Args:
1695	vllm_config: The global VllmConfig
1696	kv_cache_groups: The KV cache groups
1697	available_memory: Memory available for KV cache in bytes
1698	Returns:
1699	The generated KVCacheConfig
1700
1701	"""
1702	if len(kv_cache_groups) == 0:
1703	# Attention free models do not have KV cache.
1704	# Return num_blocks=1 as BlockPool always needs a null_block.
1705	return KVCacheConfig(
1706	num_blocks=1,
1707	kv_cache_tensors=[],
1708	kv_cache_groups=kv_cache_groups,
1709	prefix_cache_retention_interval=(
1710	vllm_config.cache_config.prefix_cache_retention_interval
1711	),
1712	)
1713
1714	if vllm_config.attention_config.hisparse_config is not None:
1715	host_budget = get_hisparse_host_pool_bytes(vllm_config)
1716	return get_hisparse_kv_cache_config(
1717	vllm_config, kv_cache_groups, available_memory, host_budget
1718	)
1719
1720	if (glm5_layout := _glm5_next_tensor_layout(kv_cache_groups)) is not None:
1721	(
1722	attn_group,
1723	mamba_groups,
1724	mla_names,
1725	idx_names,
1726	mla_page,
1727	idx_page,
1728	tail_names,
1729	_,
1730	) = glm5_layout
1731	bytes_per_block = len(mla_names) * mla_page + len(idx_names) * idx_page
1732	num_blocks = may_override_num_blocks(
1733	vllm_config, available_memory // bytes_per_block
1734	)
1735	size = bytes_per_block * num_blocks
1736	attn_specs = cast(
1737	UniformTypeKVCacheSpecs, attn_group.kv_cache_spec
1738	).kv_cache_specs
1739
1740	kv_cache_tensors: list[KVCacheTensor] = []
1741
1742	def add_tensor(layer_name: str, spec: KVCacheSpec, offset: int) -> None:
1743	kv_cache_tensors.append(
1744	KVCacheTensor(
1745	size=size,
1746	layers=[layer_name],
1747	layer_stride=spec.page_size_bytes * num_blocks,
1748	block_stride=spec.page_size_bytes,
1749	offset=offset,
1750	)
1751	)
1752
1753	for index, mla_name in enumerate(mla_names):
1754	offset = index * mla_page * num_blocks
1755	add_tensor(mla_name, attn_specs[mla_name], offset)
1756	for group in mamba_groups:
1757	if index < len(group.layer_names):
1758	add_tensor(group.layer_names[index], group.kv_cache_spec, offset)
1759
1760	idx_base = len(mla_names) * mla_page * num_blocks
1761	for index, idx_name in enumerate(idx_names):
1762	offset = idx_base + index * idx_page * num_blocks
1763	add_tensor(idx_name, attn_specs[idx_name], offset)
1764	if tail_names:
1765	tail_name = tail_names[index]
1766	tail_group = next(
1767	group for group in kv_cache_groups if tail_name in group.layer_names
1768	)
1769	tail_specs = cast(
1770	UniformTypeKVCacheSpecs, tail_group.kv_cache_spec
1771	).kv_cache_specs
1772	add_tensor(tail_name, tail_specs[tail_name], offset)
1773
1774	return KVCacheConfig(
1775	num_blocks=num_blocks,
1776	kv_cache_tensors=kv_cache_tensors,
1777	kv_cache_groups=kv_cache_groups,
1778	prefix_cache_retention_interval=(
1779	vllm_config.cache_config.prefix_cache_retention_interval
1780	),
1781	)
1782
1783	layout = vllm_config.cache_config.get_resolved_kv_cache_layout()
1784	validate_kv_cache_layout(layout, kv_cache_groups)
1785	bytes_per_block = _get_kv_cache_bytes_per_block(kv_cache_groups)
1786	interleaved_block_stride = bytes_per_block if layout.is_block_outermost else None
1787
1788	num_blocks = available_memory // bytes_per_block
1789	num_blocks = may_override_num_blocks(vllm_config, num_blocks)
1790	size = bytes_per_block * num_blocks
1791
1792	# Groups alias from byte 0. Spec regions are laid out differently:
1793	#
1794	# block-outer (the same packing repeats for every block):
1795	# group 0: \| blk 0 [ A \| B \| pad ] \| blk 1 [ A \| B \| pad ] \| ...
1796	# group 1: \| blk 0 [ C \| D ] \| blk 1 [ C \| D ] \| ...
1797	# \|<--- bytes_per_block -->\|
1798	#
1799	# layer-outer (only supported for uniform page sizes or single-group models):
1800	# group 0: \| A [ blk 0 \| blk 1 \| ... ] \| B [ blk 0 \| blk 1 \| ... ] \|
1801	# group 1: \| C [ blk 0 \| blk 1 \| ... ] \| D [ blk 0 \| blk 1 \| ... ] \|
1802
1803	kv_cache_tensors = []
1804	for group in kv_cache_groups:
1805	group_spec = group.kv_cache_spec
1806	layers_by_spec: defaultdict[KVCacheSpec, list[str]] = defaultdict(list)
1807	if isinstance(group_spec, UniformTypeKVCacheSpecs):
1808	for layer_name, spec in group_spec.kv_cache_specs.items():
1809	layers_by_spec[spec].append(layer_name)
1810	elif group.layer_names:
1811	layers_by_spec[group_spec].extend(group.layer_names)
1812
1813	byte_offset = 0
1814	for spec, layer_names in layers_by_spec.items():
1815	layer_stride, block_stride, _, _, _ = compute_layout_strides(
1816	spec,
1817	num_blocks,
1818	len(layer_names),
1819	layout,
1820	fixed_strides=(None, interleaved_block_stride, None, None, None),
1821	)
1822	offset = (
1823	byte_offset
1824	* max(layer_stride, spec.page_size_bytes)
1825	// spec.page_size_bytes
1826	)
1827	kv_cache_tensors.append(
1828	KVCacheTensor(
1829	size=size,
1830	layers=layer_names,
1831	layer_stride=layer_stride,
1832	block_stride=block_stride,
1833	offset=offset,
1834	)
1835	)
1836	byte_offset += len(layer_names) * spec.page_size_bytes
1837
1838	return KVCacheConfig(
1839	num_blocks=num_blocks,
1840	kv_cache_tensors=kv_cache_tensors,
1841	kv_cache_groups=kv_cache_groups,
1842	prefix_cache_retention_interval=(
1843	vllm_config.cache_config.prefix_cache_retention_interval
1844	),
1845	)
1846
1847
1848	def _promote_local_kv_cache_specs(
1849	kv_cache_spec: dict[str, KVCacheSpec],
1850	) -> dict[str, KVCacheSpec]:
1851	"""Use full-attention allocation for local-attention cache specs.
1852
1853	The returned specs affect KV cache management only. Attention modules keep
1854	their original sliding-window or chunked-local compute behavior.
1855	"""
1856	promoted_specs = kv_cache_spec.copy()
1857
1858	if is_kv_cache_spec_uniform(
1859	promoted_specs
1860	) or UniformTypeKVCacheSpecs.is_uniform_type(promoted_specs):
1861	return promoted_specs
1862
1863	has_full_attention = any(
1864	isinstance(spec, FullAttentionSpec) for spec in promoted_specs.values()
1865	)
1866	has_sliding_window = any(
1867	isinstance(spec, SlidingWindowSpec) for spec in promoted_specs.values()
1868	)
1869	has_chunked_local_attention = any(
1870	isinstance(spec, ChunkedLocalAttentionSpec) for spec in promoted_specs.values()
1871	)
1872	full_block_sizes = {
1873	spec.block_size
1874	for spec in promoted_specs.values()
1875	if isinstance(spec, FullAttentionSpec)
1876	}
1877	full_attention_block_size = (
1878	next(iter(full_block_sizes)) if len(full_block_sizes) == 1 else None
1879	)
1880
1881	def promoted_page_size_padded(spec: AttentionSpec, block_size: int) -> int \| None:
1882	if spec.page_size_padded is None:
1883	return None
1884	unpadded_page_size = (
1885	spec.unpadded_page_size_bytes * block_size // spec.block_size
1886	)
1887	return max(spec.page_size_padded, unpadded_page_size)
1888
1889	promotions: dict[type[AttentionSpec], type[AttentionSpec]] = {
1890	SlidingWindowMLASpec: MLAAttentionSpec,
1891	SlidingWindowSpec: FullAttentionSpec,
1892	ChunkedLocalAttentionSpec: FullAttentionSpec,
1893	}
1894
1895	if has_full_attention and (has_sliding_window or has_chunked_local_attention):
1896	for layer_name, spec in kv_cache_spec.items():
1897	target_cls = next(
1898	(promotions[c] for c in type(spec).__mro__ if c in promotions), None
1899	)
1900	if target_cls is None:
1901	continue
1902	assert isinstance(spec, AttentionSpec)
1903	block_size = full_attention_block_size or spec.block_size
1904	promoted_specs[layer_name] = replace_as(
1905	spec,
1906	target_cls,
1907	# Promoted specs allocate blocks for all tokens and never free
1908	# below the window, so the trailing-edge extension is moot.
1909	drop=("extra_retained_tokens",),
1910	block_size=block_size,
1911	page_size_padded=promoted_page_size_padded(spec, block_size),
1912	)
1913
1914	if not (
1915	is_kv_cache_spec_uniform(promoted_specs)
1916	or UniformTypeKVCacheSpecs.is_uniform_type(promoted_specs)
1917	):
1918	raise ValueError("Failed to promote local KV cache specs to one unified type.")
1919
1920	return promoted_specs
1921
1922
1923	def _try_get_full_allocation_fallback_groups(
1924	kv_cache_spec: dict[str, KVCacheSpec],
1925	) -> list[KVCacheGroupSpec] \| None:
1926	"""Try a supported full-allocation fallback for local-attention layers."""
1927	if any(isinstance(spec, HiddenStateCacheSpec) for spec in kv_cache_spec.values()):
1928	return None
1929	if any(
1930	isinstance(spec, (SlidingWindowMLASpec, ChunkedLocalAttentionSpec))
1931	for spec in kv_cache_spec.values()
1932	):
1933	return None
1934
1935	has_mla = any(isinstance(spec, MLAAttentionSpec) for spec in kv_cache_spec.values())
1936	has_regular_swa = any(
1937	isinstance(spec, SlidingWindowSpec) for spec in kv_cache_spec.values()
1938	)
1939	if not (has_mla and has_regular_swa):
1940	return None
1941
1942	try:
1943	promoted_specs = _promote_local_kv_cache_specs(kv_cache_spec)
1944	except ValueError:
1945	return None
1946	uniform_spec = UniformTypeKVCacheSpecs.from_specs(promoted_specs)
1947	if uniform_spec is None:
1948	return None
1949	logger.warning(
1950	"KV cache page sizes cannot be unified; treating sliding-window "
1951	"layers as full attention for cache allocation. Sliding-window "
1952	"attention compute is unchanged."
1953	)
1954	return _get_kv_cache_groups_uniform_type(uniform_spec)
1955
1956
1957	def unify_hybrid_kv_cache_specs(kv_cache_spec: dict[str, KVCacheSpec]):
1958	"""This function tries to convert the KV cache specs to one type if the model
1959	is a hybrid model with multiple type of KV cache. It will convert all
1960	SlidingWindowSpec to FullAttentionSpec if both types are present.
1961
1962	Args:
1963	kv_cache_spec: The kv cache spec of each attention layer in the model
1964
1965	"""
1966	groups: defaultdict[bool, dict[str, KVCacheSpec]] = defaultdict(dict)
1967	for name, spec in kv_cache_spec.items():
1968	replicated = isinstance(spec, AttentionSpec) and not spec.dcp_sharded
1969	groups[replicated][name] = spec
1970	for specs in groups.values():
1971	promoted_specs = _promote_local_kv_cache_specs(specs)
1972	if promoted_specs == specs:
1973	continue
1974	logger.warning(
1975	"Hybrid KV cache manager is disabled for this hybrid model, "
1976	"This means we do not enable any optimizations for saving KV cache "
1977	"memory (e.g., dropping the KV cache outside the sliding window). "
1978	"The compute of layers like sliding window is still saved."
1979	)
1980	kv_cache_spec.update(promoted_specs)
1981
1982
1983	def _approximate_gcd(values: Sequence[int], *, lower_bound: int \| None = None) -> int:
1984	"""Pick a chunk size that minimizes total upward padding.
1985
1986	Each x is rounded up to a multiple of d:
1987
1988	x -> ceil(x / d) * d
1989
1990	Total padding is:
1991
1992	pad(d) = sum_i (ceil(x_i / d) * d - x_i)
1993
1994	We brute-force d in [lower_bound, max(values)] (fine for small lists / small
1995	maxima) and return the d with minimum padding. Ties prefer larger d.
1996	"""
1997	if not values:
1998	raise ValueError("values must be non-empty")
1999	if any(x <= 0 for x in values):
2000	raise ValueError(f"values must be positive, got: {list(values)!r}")
2001
2002	min_d = max(1, lower_bound if lower_bound is not None else 1)
2003	max_d = max(values)
2004	if min_d > max_d:
2005	return min_d
2006
2007	best_d = min_d
2008	best_pad: int \| None = None
2009	for d in range(min_d, max_d + 1):
2010	pad = sum((d - (x % d)) % d for x in values)
2011	if best_pad is None or pad < best_pad or (pad == best_pad and d > best_d):
2012	best_pad = pad
2013	best_d = d
2014
2015	return best_d
2016
2017
2018	def _get_packed_kv_cache_groups(
2019	vllm_config: VllmConfig,
2020	kv_cache_spec: dict[str, KVCacheSpec],
2021	) -> list[KVCacheGroupSpec] \| None:
2022	"""Group mixed-page-size layers for contiguous block-outermost packing.
2023
2024	Greedily buckets layers into uniform-type specs. Buckets with equal layer
2025	counts per page size are treated as a repeating layer pattern (one layer
2026	per page size) and split into groups covering the same number of pattern
2027	repeats (picked by ``_approximate_gcd`` to minimize padding), so all
2028	groups pack into the same per-block layout. Mamba buckets are additionally
2029	split to fit the block the attention buckets already need.
2030	Returns None when the layout is not block-outermost or all layers already
2031	share one page size.
2032	"""
2033	layout = vllm_config.cache_config.get_resolved_kv_cache_layout()
2034	page_sizes = {spec.page_size_bytes for spec in kv_cache_spec.values()}
2035	if not layout.is_block_outermost or len(page_sizes) <= 1:
2036	return None
2037
2038	buckets: list[dict[str, KVCacheSpec]] = []
2039	for name, spec in kv_cache_spec.items():
2040	for bucket in buckets:
2041	candidate = {**bucket, name: spec}
2042	if UniformTypeKVCacheSpecs.is_uniform_type(candidate):
2043	bucket[name] = spec
2044	break
2045	else:
2046	buckets.append({name: spec})
2047
2048	bucketed = []
2049	for bucket in buckets:
2050	uniform_spec = UniformTypeKVCacheSpecs.from_specs(bucket)
2051	assert uniform_spec is not None
2052	page_size_layers: dict[int, list[str]] = defaultdict(list)
2053	for layer_name, layer_spec in bucket.items():
2054	page_size_layers[layer_spec.page_size_bytes].append(layer_name)
2055	# Only 1:1 patterns (one layer of each page size per repeat) are
2056	# supported; counts sharing a gcd > 1 (e.g. 2:1) could in principle
2057	# repeat too, but such buckets are emitted whole instead.
2058	balanced = len(set(map(len, page_size_layers.values()))) == 1
2059	bucketed.append((uniform_spec, page_size_layers, balanced))
2060
2061	# Balanced buckets that mix page sizes must stay whole, so the largest one
2062	# sets a floor on the repeats per group; larger single-size buckets are
2063	# split down toward it. No such bucket means nothing needs packing.
2064	min_repeats_per_group = max(
2065	(
2066	spec.get_max_layers_per_page_size()
2067	for spec, page_size_layers, balanced in bucketed
2068	if balanced and len(page_size_layers) > 1
2069	),
2070	default=0,
2071	)
2072	repeats_per_group = (
2073	_approximate_gcd(
2074	[
2075	spec.get_max_layers_per_page_size()
2076	for spec, _, balanced in bucketed
2077	if balanced
2078	],
2079	lower_bound=min_repeats_per_group,
2080	)
2081	if min_repeats_per_group
2082	else None
2083	)
2084
2085	def num_groups_for(spec: UniformTypeKVCacheSpecs, balanced: bool) -> int:
2086	if balanced and repeats_per_group is not None:
2087	return cdiv(spec.get_max_layers_per_page_size(), repeats_per_group)
2088	return 1
2089
2090	def widest_group_bytes(page_size_layers: dict[int, list[str]], n: int) -> int:
2091	"""Page bytes of the largest of the n groups a bucket splits into."""
2092	return sum(
2093	cdiv(len(names), n) * page for page, names in page_size_layers.items()
2094	)
2095
2096	# Bytes a block must hold however the state buckets end up split: mamba,
2097	# circular-buffer and unsplit sliding-window buckets can go down to one
2098	# state per group, every other bucket's split is fixed by the repeat pattern.
2099	def is_state_bucket(spec: UniformTypeKVCacheSpecs) -> bool:
2100	if isinstance(spec.first_spec, (MambaSpec, CircularBufferSpec)):
2101	return True
2102	return repeats_per_group is None and isinstance(
2103	spec.first_spec, SlidingWindowSpec
2104	)
2105
2106	anchor_bytes = max(
2107	(
2108	widest_group_bytes(
2109	page_size_layers,
2110	len(spec.kv_cache_specs)
2111	if is_state_bucket(spec)
2112	else num_groups_for(spec, balanced),
2113	)
2114	for spec, page_size_layers, balanced in bucketed
2115	),
2116	default=0,
2117	)
2118
2119	groups = []
2120	for spec, page_size_layers, balanced in bucketed:
2121	num_groups = num_groups_for(spec, balanced)
2122	# Cap a state group at the states a block already fits rather than let
2123	# it widen the block.
2124	if anchor_bytes and is_state_bucket(spec):
2125	states_per_block = max(anchor_bytes // spec.first_spec.page_size_bytes, 1)
2126	num_groups = max(
2127	num_groups, cdiv(len(spec.kv_cache_specs), states_per_block)
2128	)
2129	if num_groups == 1:
2130	groups.append(KVCacheGroupSpec(list(spec.kv_cache_specs), spec))
2131	continue
2132
2133	pattern_repeats = list(zip(*page_size_layers.values()))
2134	for i in range(num_groups):
2135	group_layer_names = [
2136	name for repeat in pattern_repeats[i::num_groups] for name in repeat
2137	]
2138	group_layer_specs = {
2139	name: spec.kv_cache_specs[name] for name in group_layer_names
2140	}
2141	group_spec = UniformTypeKVCacheSpecs.from_specs(group_layer_specs)
2142	assert group_spec is not None
2143	groups.append(KVCacheGroupSpec(group_layer_names, group_spec))
2144
2145	_annotate_eagle_groups(
2146	vllm_config,
2147	kv_cache_spec,
2148	groups,
2149	use_trailing_layer_fallback=_uses_trailing_mtp_layers(vllm_config),
2150	)
2151	_warn_if_unannotated_eagle_mamba(vllm_config, groups)
2152	return groups
2153
2154
2155	def _uses_trailing_mtp_layers(vllm_config: VllmConfig) -> bool:
2156	"""Whether the drafter's KV layers can be located positionally.
2157
2158	MTP drafters register their KV layers after the target layers, without a
2159	distinguishing spec marker. DeepseekV4/V4.1 do this with ``dspark`` too.
2160	The annotator separately checks that the groups partition the layers exactly.
2161	"""
2162	spec_config = vllm_config.speculative_config
2163	if spec_config is None or not spec_config.use_eagle():
2164	return False
2165	if spec_config.method == "mtp":
2166	return True
2167	model_config = vllm_config.model_config
2168	return model_config is not None and model_config.hf_config.model_type in (
2169	"deepseek_v4",
2170	"deepseek_v41",
2171	)
2172
2173
2174	def _groups_partition_layers_exactly(
2175	kv_cache_spec: dict[str, KVCacheSpec],
2176	kv_cache_groups: list[KVCacheGroupSpec],
2177	) -> bool:
2178	"""Whether the groups cover every layer of ``kv_cache_spec`` exactly once.
2179
2180	The trailing-layer draft fallback is only meaningful when the last
2181	registered layer is guaranteed to appear in exactly one group.
2182	"""
2183	names = [name for group in kv_cache_groups for name in group.layer_names]
2184	return len(names) == len(kv_cache_spec) and set(names) == set(kv_cache_spec)
2185
2186
2187	def _annotate_eagle_groups(
2188	vllm_config: VllmConfig,
2189	kv_cache_spec: dict[str, KVCacheSpec],
2190	kv_cache_groups: list[KVCacheGroupSpec],
2191	use_trailing_layer_fallback: bool = False,
2192	) -> None:
2193	"""Flag the KV cache groups that hold drafter attention layers.
2194
2195	Two detection rules, in order of preference:
2196
2197	1. Spec-driven. ``non_causal_multi_token_decode`` is declared on
2198	MLAAttentionSpec and set by drafter attention layers that run a
2199	non-causal multi-token decode (today only Kimi-K3 DSpark). It survives
2200	MLAAttentionSpec.merge, so it still identifies a group after per-group
2201	spec merging, wherever grouping happens to land. It is sufficient but
2202	not necessary: a drafter whose spec is indistinguishable from the
2203	target's cannot be found this way.
2204	2. Positional fallback for MTP drafters (including DeepseekV4/V4.1 DSpark), whose
2205	MTP block reuses the target's own decoder layer and so carries no spec
2206	marker. Their draft attention layers always register after every
2207	target layer, so flag whichever group holds the last registered layer.
2208	This rule is only valid where the groups partition exactly the layers
2209	of ``kv_cache_spec``, which is re-checked here before applying it.
2210	When a drafter's trailing caches span several groups, this rule flags
2211	only the group holding the very last layer and must be generalized.
2212	FIXME(yifan): avoid/generalize this hacky check.
2213
2214	Args:
2215	vllm_config: Config supplying the speculative method, if any.
2216	kv_cache_spec: The kv cache spec of each attention layer, in layer
2217	registration order. Only read by rule 2.
2218	kv_cache_groups: Groups to annotate in place.
2219	use_trailing_layer_fallback: Enable rule 2. Callers gate this on
2220	``_uses_trailing_mtp_layers``.
2221
2222	"""
2223	spec_config = vllm_config.speculative_config
2224	if spec_config is None or not spec_config.use_eagle_block_drop():
2225	return
2226
2227	for group in kv_cache_groups:
2228	if any(
2229	getattr(spec, "non_causal_multi_token_decode", False)
2230	for spec in iter_layer_specs(group.kv_cache_spec)
2231	):
2232	group.is_eagle_group = True
2233
2234	if not use_trailing_layer_fallback:
2235	return
2236	if not _groups_partition_layers_exactly(kv_cache_spec, kv_cache_groups):
2237	return
2238	last_layer = next(reversed(kv_cache_spec))
2239	for group in kv_cache_groups:
2240	if last_layer in group.layer_names:
2241	group.is_eagle_group = True
2242	break
2243
2244
2245	def _warn_if_unannotated_eagle_mamba(
2246	vllm_config: VllmConfig,
2247	kv_cache_groups: list[KVCacheGroupSpec],
2248	) -> None:
2249	"""Warn when no KV cache group could be identified as the draft model's.
2250
2251	Args:
2252	vllm_config: Config supplying the speculative method, if any.
2253	kv_cache_groups: Groups as they will be handed to consumers.
2254
2255	"""
2256	spec_config = vllm_config.speculative_config
2257	if spec_config is None or not spec_config.use_eagle():
2258	return
2259	if any(group.is_eagle_group for group in kv_cache_groups):
2260	return
2261	mamba_groups = [
2262	idx
2263	for idx, group in enumerate(kv_cache_groups)
2264	if any(
2265	isinstance(spec, MambaSpec)
2266	for spec in iter_layer_specs(group.kv_cache_spec)
2267	)
2268	]
2269	if not mamba_groups:
2270	return
2271	logger.warning(
2272	"Speculative decoding (method=%s) is enabled but no KV cache group "
2273	"could be identified as the draft model's.",
2274	spec_config.method,
2275	)
2276
2277
2278	def _largest_divisor_at_most(value: int, limit: int) -> int:
2279	for candidate in range(min(value, limit), 0, -1):
2280	if value % candidate == 0:
2281	return candidate
2282	return 1
2283
2284
2285	def _ensure_min_page_size(
2286	groups: list[KVCacheGroupSpec],
2287	common_page: int,
2288	hidden_specs: dict[str, HiddenStateCacheSpec],
2289	) -> tuple[list[KVCacheGroupSpec], int]:
2290	"""Scale up group block sizes so the common page is at least as large as
2291	the biggest hidden-state per-token cost.
2292
2293	This protects the hidden state extraction feature, where we store hidden states in
2294	the KV cache, and return them to clients using a custom hidden state connector.
2295
2296	Returns the (possibly scaled) groups and updated common page size.
2297	"""
2298	max_per_token = max(
2299	spec.num_kv_heads * spec.head_size * get_dtype_size(spec.dtype)
2300	for spec in hidden_specs.values()
2301	)
2302	if max_per_token <= common_page:
2303	return groups, common_page
2304
2305	scale = cdiv(max_per_token, common_page)
2306	common_page *= scale
2307	scaled: list[KVCacheGroupSpec] = []
2308	for g in groups:
2309	s = g.kv_cache_spec
2310	kw: dict[str, Any] = {"block_size": s.block_size * scale}
2311	if isinstance(s, (AttentionSpec, MambaSpec)) and s.page_size_padded is not None:
2312	kw["page_size_padded"] = s.page_size_padded * scale
2313	scaled.append(KVCacheGroupSpec(g.layer_names, replace(s, **kw)))
2314	return scaled, common_page
2315
2316
2317	def get_kv_cache_groups(
2318	vllm_config: VllmConfig,
2319	kv_cache_spec: dict[str, KVCacheSpec],
2320	) -> list[KVCacheGroupSpec]:
2321	"""Split the layers in the model into groups with the same KV cache spec.
2322
2323	Args:
2324	vllm_config: The global VllmConfig
2325	kv_cache_spec: The kv cache spec of each attention layer in the model
2326
2327	Returns:
2328	The generated KVCacheGroups
2329
2330	"""
2331	if vllm_config.scheduler_config.disable_hybrid_kv_cache_manager:
2332	unify_hybrid_kv_cache_specs(kv_cache_spec)
2333
2334	if is_kv_cache_type_attention_free(kv_cache_spec):
2335	# This returns an empty list to allow for the KVCacheManager to handle
2336	# attention free models.
2337	return []
2338
2339	if hisparse_groups := get_hisparse_kv_cache_groups(vllm_config, kv_cache_spec):
2340	return hisparse_groups
2341
2342	if is_kv_cache_spec_uniform(kv_cache_spec):
2343	# KV cache of all layers are the same, which is true for
2344	# most models. Allocate the same amount of memory for
2345	# each layer.
2346	return _get_kv_cache_groups_uniform_spec(kv_cache_spec)
2347	elif uniform_spec := UniformTypeKVCacheSpecs.from_specs(kv_cache_spec):
2348	# All layers need the same number of token slots (e.g., all layers are
2349	# full attention, or all layers are sliding window attention with the
2350	# same window size). Put all layers into one group.
2351	return _get_kv_cache_groups_uniform_type(uniform_spec)
2352	elif glm5_groups := _get_kv_cache_groups_glm5_next(vllm_config, kv_cache_spec):
2353	return glm5_groups
2354
2355	# Hidden-state layers use their own block table and must not be absorbed
2356	# into a compatible attention bucket.
2357	hidden_specs = {
2358	k: v for k, v in kv_cache_spec.items() if isinstance(v, HiddenStateCacheSpec)
2359	}
2360	filtered_spec = {
2361	k: v
2362	for k, v in kv_cache_spec.items()
2363	if not isinstance(v, HiddenStateCacheSpec)
2364	}
2365
2366	if packed_groups := _get_packed_kv_cache_groups(vllm_config, filtered_spec):
2367	# Block-outermost blocks are strided by the widest group, so hidden
2368	# groups need no page alignment.
2369	packed_groups += [
2370	KVCacheGroupSpec([name], spec) for name, spec in hidden_specs.items()
2371	]
2372	return packed_groups
2373
2374	# Prefer preserving each layer's cache semantics. If physical pages cannot
2375	# be unified, try a supported allocation-only fallback before failing.
2376	try:
2377	filtered_spec = unify_kv_cache_spec_page_size(filtered_spec)
2378	except NotImplementedError:
2379	fallback_groups = _try_get_full_allocation_fallback_groups(kv_cache_spec)
2380	if fallback_groups is None:
2381	raise
2382	return fallback_groups
2383	groups = _get_kv_cache_groups_uniform_page_size(filtered_spec)
2384
2385	# Add hidden-state layers back with page aligned to the common page.
2386	if hidden_specs:
2387	common_page = get_uniform_page_size([g.kv_cache_spec for g in groups])
2388	# TP may shrink the common page below hidden-state per-token cost.
2389	groups, common_page = _ensure_min_page_size(groups, common_page, hidden_specs)
2390	group_block_size = math.gcd(*(g.kv_cache_spec.block_size for g in groups))
2391	for name, spec in hidden_specs.items():
2392	per_token = spec.num_kv_heads * spec.head_size * get_dtype_size(spec.dtype)
2393	max_block_size = max(common_page // per_token, 1)
2394	new_bs = _largest_divisor_at_most(group_block_size, max_block_size)
2395	wasted_bytes = common_page - new_bs * per_token
2396	logger.info(
2397	"Using block size %d for hidden-state cache layer %s; "
2398	"page alignment wastes %d bytes (%.2f%%) per block",
2399	new_bs,
2400	name,
2401	wasted_bytes,
2402	wasted_bytes / common_page * 100,
2403	)
2404	aligned = replace(spec, block_size=new_bs, page_size_padded=common_page)
2405	groups.append(KVCacheGroupSpec([name], aligned))
2406
2407	_annotate_eagle_groups(
2408	vllm_config,
2409	kv_cache_spec,
2410	groups,
2411	use_trailing_layer_fallback=_uses_trailing_mtp_layers(vllm_config),
2412	)
2413	_warn_if_unannotated_eagle_mamba(vllm_config, groups)
2414	return groups
2415
2416
2417	def _layer_tp_replicas(spec: KVCacheSpec, tp_size: int, dcp_size: int) -> int:
2418	if not isinstance(spec, AttentionSpec) or spec.max_tp_shards is None:
2419	return 1
2420	if spec.dcp_sharded and dcp_size > 1:
2421	return 1
2422	return max(1, tp_size // spec.max_tp_shards)
2423
2424
2425	def kv_cache_groups_tp_replicas(
2426	groups: list[KVCacheGroupSpec], tp_size: int, dcp_size: int = 1
2427	) -> int:
2428	"""Consecutive TP ranks holding identical KV for every layer."""
2429	specs = [spec for g in groups for spec in iter_layer_specs(g.kv_cache_spec)]
2430	if not specs:
2431	return 1
2432	return reduce(
2433	math.gcd, (_layer_tp_replicas(s, tp_size, dcp_size) for s in specs), tp_size
2434	)
2435
2436
2437	def generate_scheduler_kv_cache_config(
2438	kv_cache_configs: list[KVCacheConfig],
2439	) -> KVCacheConfig:
2440	"""Generate the KV cache configuration for the scheduler."""
2441	assert all(
2442	[cfg.num_blocks == kv_cache_configs[0].num_blocks for cfg in kv_cache_configs]
2443	)
2444	assert all(
2445	cfg.hisparse_host_num_blocks == kv_cache_configs[0].hisparse_host_num_blocks
2446	for cfg in kv_cache_configs
2447	)
2448	# All workers have the same kv_cache_config except layer names, so use
2449	# an arbitrary one to initialize the scheduler.
2450	cfg = copy.deepcopy(kv_cache_configs[0])
2451	for group in cfg.kv_cache_groups:
2452	if isinstance(group.kv_cache_spec, UniformTypeKVCacheSpecs):
2453	# All layers in the UniformTypeKVCacheSpecs have the same type,
2454	# so use an arbitrary one to initialize the scheduler.
2455	group.kv_cache_spec = next(
2456	iter(group.kv_cache_spec.kv_cache_specs.values())
2457	)
2458	return cfg
2459
2460
2461	def get_kv_cache_capacity(
2462	vllm_config: VllmConfig, kv_cache_config: KVCacheConfig
2463	) -> tuple[int, float]:
2464	"""Get the group-aware KV cache token capacity and max concurrency."""
2465	max_model_len = vllm_config.model_config.max_model_len
2466	max_concurrency = get_max_concurrency_for_kv_cache_config(
2467	vllm_config, kv_cache_config
2468	)
2469	return int(max_concurrency * max_model_len), max_concurrency
2470
2471
2472	def update_kv_cache_capacity(
2473	vllm_config: VllmConfig, kv_cache_config: KVCacheConfig
2474	) -> None:
2475	"""Store and log the resolved KV cache capacity."""
2476	num_tokens, max_concurrency = get_kv_cache_capacity(vllm_config, kv_cache_config)
2477	vllm_config.cache_config.kv_cache_size_tokens = num_tokens
2478	vllm_config.cache_config.kv_cache_max_concurrency = max_concurrency
2479	max_model_len = vllm_config.model_config.max_model_len
2480	device_type = vllm_config.device_config.device_type
2481	logger.info_once(
2482	"%s KV cache size: %s tokens, "
2483	"Maximum concurrency for %s tokens per request: %.2fx",
2484	"GPU" if device_type == "cuda" else device_type.upper(),
2485	f"{num_tokens:,}",
2486	f"{max_model_len:,}",
2487	max_concurrency,
2488	)
2489
2490
2491	def _max_memory_usage_bytes_from_groups(
2492	vllm_config: VllmConfig,
2493	kv_cache_groups: list[KVCacheGroupSpec],
2494	) -> int:
2495	"""Calculate maximum memory usage in bytes from KV cache groups.
2496
2497	This correctly accounts for padding in hybrid models. For example, if a
2498	model has 8 full attention layers and 9 sliding window layers, they will
2499	be padded to 9 full + 9 sliding window for uniform group sizes.
2500
2501	Each group independently claims blocks from the shared pool, so a request consumes
2502	the sum of the per-group block counts, i.e. ``bytes_per_block * total_blocks``.
2503	"""
2504	kv_cache_groups = [group for group in kv_cache_groups if not group.host_resident]
2505	if not kv_cache_groups:
2506	return 0
2507
2508	if (glm5_layout := _glm5_next_tensor_layout(kv_cache_groups)) is not None:
2509	(
2510	attn_group,
2511	mamba_groups,
2512	mla_names,
2513	idx_names,
2514	mla_page,
2515	idx_page,
2516	tail_names,
2517	_,
2518	) = glm5_layout
2519	uniform_spec = cast(UniformTypeKVCacheSpecs, attn_group.kv_cache_spec)
2520	total_blocks = uniform_spec.max_memory_usage_pages(vllm_config)
2521	total_blocks += sum(
2522	cdiv(
2523	group.kv_cache_spec.max_memory_usage_bytes(vllm_config),
2524	group.kv_cache_spec.page_size_bytes,
2525	)
2526	for group in mamba_groups
2527	)
2528	if tail_names:
2529	total_blocks += 1
2530	return total_blocks * (len(mla_names) * mla_page + len(idx_names) * idx_page)
2531
2532	bytes_per_block = _pool_bytes_per_block(kv_cache_groups)
2533	total_blocks = 0
2534	for group in kv_cache_groups:
2535	spec = group.kv_cache_spec
2536	if isinstance(spec, UniformTypeKVCacheSpecs):
2537	total_blocks += spec.max_memory_usage_pages(vllm_config)
2538	else:
2539	total_blocks += cdiv(
2540	spec.max_memory_usage_bytes(vllm_config),
2541	spec.page_size_bytes,
2542	)
2543
2544	return bytes_per_block * total_blocks
2545
2546
2547	def _estimate_max_model_len_from_groups(
2548	vllm_config: VllmConfig,
2549	kv_cache_groups: list[KVCacheGroupSpec],
2550	available_memory: int,
2551	) -> int:
2552	"""Binary search for the maximum model length that fits in available memory.
2553	Returns 0 if even 1 token doesn't fit.
2554	"""
2555	original_max = vllm_config.model_config.max_model_len
2556	hisparse_enabled = (
2557	vllm_config.attention_config.hisparse_config is not None
2558	and bool(kv_cache_groups)
2559	)
2560
2561	def fits(model_len: int) -> bool:
2562	vllm_config.model_config.max_model_len = model_len
2563	if hisparse_enabled:
2564	try:
2565	get_kv_cache_config_from_groups(
2566	vllm_config, kv_cache_groups, available_memory
2567	)
2568	except ValueError:
2569	return False
2570	return (
2571	_max_memory_usage_bytes_from_groups(vllm_config, kv_cache_groups)
2572	<= available_memory
2573	)
2574
2575	try:
2576	left, right = 1, original_max
2577	if not fits(left):
2578	return 0
2579	result = 1
2580	while left <= right:
2581	mid = (left + right) // 2
2582	if fits(mid):
2583	result = mid
2584	left = mid + 1
2585	else:
2586	right = mid - 1
2587	return result
2588	finally:
2589	vllm_config.model_config.max_model_len = original_max
2590
2591
2592	def _auto_fit_max_model_len(
2593	vllm_config: VllmConfig,
2594	projected_groups_per_worker: list[list[KVCacheGroupSpec]],
2595	available_memory: list[int],
2596	) -> None:
2597	"""When max_model_len is set to -1, this function estimates the largest
2598	context length that can be supported with the available GPU memory.
2599	It uses binary search to find the maximum length that fits across all
2600	workers.
2601
2602	Args:
2603	vllm_config: The global VllmConfig (will be modified in-place)
2604	projected_groups_per_worker: KV cache groups projected to each worker.
2605	available_memory: Memory available for KV cache in bytes for each
2606	worker.
2607
2608	"""
2609	original_max = vllm_config.model_config.max_model_len
2610
2611	if all(not groups for groups in projected_groups_per_worker):
2612	# All workers have empty specs (attention-free model)
2613	logger.info_once(
2614	"Auto-fit max_model_len: attention-free model, "
2615	"using derived max_model_len=%d",
2616	original_max,
2617	)
2618	return
2619
2620	# Find the max_model_len that fits across all workers.
2621	auto_fit_max = original_max
2622	limiting_worker_mem = available_memory[0]
2623	for groups, avail_mem in zip(projected_groups_per_worker, available_memory):
2624	if not groups:
2625	continue
2626	worker_max = _estimate_max_model_len_from_groups(vllm_config, groups, avail_mem)
2627	if worker_max < auto_fit_max:
2628	auto_fit_max = worker_max
2629	limiting_worker_mem = avail_mem
2630
2631	if auto_fit_max <= 0:
2632	raise ValueError(
2633	"Cannot auto-fit max_model_len: not enough GPU memory available "
2634	"to serve even a single token. Try increasing `gpu_memory_utilization`."
2635	)
2636
2637	if auto_fit_max >= original_max:
2638	# The model's full context length fits in memory
2639	logger.info_once(
2640	"Auto-fit max_model_len: full model context length %d fits in "
2641	"available GPU memory",
2642	original_max,
2643	)
2644	else:
2645	# Need to reduce max_model_len to fit in memory
2646	vllm_config.model_config.max_model_len = auto_fit_max
2647	logger.info_once(
2648	"Auto-fit max_model_len: reduced from %d to %d to fit in "
2649	"available GPU memory (%s GiB available for KV cache)",
2650	original_max,
2651	auto_fit_max,
2652	format_gib(limiting_worker_mem),
2653	)
2654
2655
2656	def _project_kv_cache_groups_to_worker(
2657	global_kv_cache_groups: list[KVCacheGroupSpec],
2658	worker_spec: dict[str, KVCacheSpec],
2659	) -> list[KVCacheGroupSpec]:
2660	"""Projects global KV cache groups onto a single worker's assigned layers.
2661
2662	In pipeline parallelism, each worker only owns a subset of layers. This
2663	function filters the global groups to include only layers present on the
2664	given worker, adjusting UniformTypeKVCacheSpecs accordingly.
2665
2666	Args:
2667	global_kv_cache_groups: The global KV cache groups for the whole model.
2668	worker_spec: The KV cache spec of each layer on this worker.
2669
2670	Returns:
2671	The projected KV cache groups containing only this worker's layers.
2672
2673	"""
2674	projected_groups: list[KVCacheGroupSpec] = []
2675	for group in global_kv_cache_groups:
2676	worker_layer_names = [
2677	layer_name for layer_name in group.layer_names if layer_name in worker_spec
2678	]
2679	if len(worker_layer_names) == len(group.layer_names):
2680	projected_groups.append(group)
2681	continue
2682	group_spec = group.kv_cache_spec
2683	if worker_layer_names and isinstance(group_spec, UniformTypeKVCacheSpecs):
2684	group_spec = UniformTypeKVCacheSpecs(
2685	block_size=group_spec.block_size,
2686	kv_cache_specs={
2687	layer_name: group_spec.kv_cache_specs[layer_name]
2688	for layer_name in worker_layer_names
2689	},
2690	)
2691	projected_groups.append(
2692	KVCacheGroupSpec(
2693	worker_layer_names,
2694	group_spec,
2695	is_eagle_group=group.is_eagle_group and bool(worker_layer_names),
2696	)
2697	)
2698	return projected_groups
2699
2700
2701	def get_kv_cache_configs(
2702	vllm_config: VllmConfig,
2703	kv_cache_specs: list[dict[str, KVCacheSpec]],
2704	available_memory: list[int],
2705	) -> list[KVCacheConfig]:
2706	"""Generates the KV cache configurations for a model.
2707	Since we use a shared centralized controller for all workers, we need the
2708	`kv_cache_config` to be consistent across all workers to make sure
2709	the KV cache allocation can be applied to all workers. However, different
2710	workers may have different memory available, and different type of layers
2711	(when pipeline parallel is enabled). To handle the difference between
2712	workers, the current implementation is:
2713	1. Merge the KV cache specs of all workers to get the KVCacheSpecs for
2714	the whole model.
2715	2. Generate the KV cache groups based on the layer ratio of the whole model.
2716	This also handles spec unification for hybrid models.
2717	3. Handle auto-fit max_model_len and memory checks using per-worker
2718	projected groups to account for PP sharding.
2719	4. Generate the KV cache configs for each worker based on the KV cache
2720	grouping strategy. (This is reasonable because the layer ratio of
2721	different PP stages are similar.)
2722	5. Change the num_blocks of each worker to the smallest among all workers
2723	and shrink tensor sizes proportionally to avoid allocating unused memory.
2724
2725	Args:
2726	vllm_config: The global VllmConfig
2727	kv_cache_specs: List of dict[layer_name, KVCacheSpec] for each worker.
2728	available_memory: Memory available for KV cache in bytes for each
2729	worker.
2730
2731	Returns:
2732	The generated KVCacheConfigs for each worker.
2733
2734	"""
2735	# Merge the KV cache specs of all workers. Different PP stages may have
2736	# different layer names, and different TP ranks of the same PP stage should
2737	# have the same KV cache spec.
2738	merged_kv_cache_specs: dict[str, KVCacheSpec] = {}
2739	for kv_cache_spec_one_worker in kv_cache_specs:
2740	for layer_name, layer_spec in kv_cache_spec_one_worker.items():
2741	if layer_name not in merged_kv_cache_specs:
2742	merged_kv_cache_specs[layer_name] = layer_spec
2743	else:
2744	assert merged_kv_cache_specs[layer_name] == layer_spec, (
2745	"The KV cache specs for the same layer are different "
2746	"across workers. This is not supported yet."
2747	)
2748
2749	# Check if the KV cache specs are registered correctly.
2750	# This is to prevent that some layers are initialized with unregistered specs.
2751	KVCacheSpecRegistry.check_kv_cache_spec_registry(merged_kv_cache_specs)
2752
2753	# When speculating with more than 1 speculative module (e.g. multi-layered MTP)
2754	# tag every SlidingWindowSpec with how many extra tokens to retain in the window.
2755	extra_retained_tokens = max(0, vllm_config.num_prefill_lookahead_tokens - 1)
2756	for layer_name, layer_spec in merged_kv_cache_specs.items():
2757	if isinstance(layer_spec, SlidingWindowSpec):
2758	merged_kv_cache_specs[layer_name] = replace(
2759	layer_spec, extra_retained_tokens=extra_retained_tokens
2760	)
2761
2762	# Get global KV cache groups. This also handles spec unification for
2763	# hybrid models when disable_hybrid_kv_cache_manager is enabled.
2764	# After this call, merged_kv_cache_specs may be modified in-place.
2765	global_kv_cache_groups = get_kv_cache_groups(vllm_config, merged_kv_cache_specs)
2766
2767	# If original_max_model_len was -1, automatically
2768	# determine the maximum model length that fits in available GPU memory.
2769	# We use per-worker projected groups to account for PP sharding.
2770	projected_groups_per_worker = [
2771	_project_kv_cache_groups_to_worker(global_kv_cache_groups, worker_spec)
2772	for worker_spec in kv_cache_specs
2773	]
2774
2775	# If `num_gpu_blocks_override` is set, the cache size that will actually
2776	# be allocated is decoupled from the profiled `available_memory`:
2777	# `may_override_num_blocks` in `get_kv_cache_config_from_groups` clamps
2778	# `num_blocks` to the override. Reflect that in `available_memory` here so
2779	# auto-fit, the admission check, and the per-worker config builder all
2780	# plan against the same effective capacity.
2781	override = vllm_config.cache_config.num_gpu_blocks_override
2782	if override is not None:
2783	adjusted_memory: list[int] = []
2784	for groups, avail_mem in zip(projected_groups_per_worker, available_memory):
2785	if not groups:
2786	adjusted_memory.append(avail_mem)
2787	continue
2788	bytes_per_block = _pool_bytes_per_block(groups)
2789	logger.info(
2790	"Overriding num_gpu_blocks=%d with num_gpu_blocks_override=%d",
2791	avail_mem // bytes_per_block,
2792	override,
2793	)
2794	adjusted_memory.append(override * bytes_per_block)
2795	available_memory = adjusted_memory
2796
2797	# Reserve the null block BlockPool permanently holds back, so auto-fit and
2798	# the capacity check both plan against usable blocks. Allocation below
2799	# still uses the full memory.
2800	check_memory = [
2801	avail_mem - _pool_bytes_per_block(groups) if groups else avail_mem
2802	for groups, avail_mem in zip(projected_groups_per_worker, available_memory)
2803	]
2804
2805	if vllm_config.model_config.original_max_model_len == -1:
2806	_auto_fit_max_model_len(vllm_config, projected_groups_per_worker, check_memory)
2807
2808	# Check if the available memory is enough per worker.
2809	for groups, avail_mem in zip(projected_groups_per_worker, check_memory):
2810	if not groups:
2811	continue
2812	_check_enough_kv_cache_memory(
2813	avail_mem,
2814	partial(_max_memory_usage_bytes_from_groups, vllm_config, groups),
2815	vllm_config.model_config.max_model_len,
2816	partial(_estimate_max_model_len_from_groups, vllm_config, groups),
2817	)
2818
2819	kv_cache_configs: list[KVCacheConfig] = []
2820	for projected_groups, kv_cache_spec_one_worker, available_memory_one_worker in zip(
2821	projected_groups_per_worker, kv_cache_specs, available_memory
2822	):
2823	assert sum(len(group.layer_names) for group in projected_groups) == len(
2824	kv_cache_spec_one_worker
2825	), "Some layers are not assigned to any group."
2826	kv_cache_configs.append(
2827	get_kv_cache_config_from_groups(
2828	vllm_config, projected_groups, available_memory_one_worker
2829	)
2830	)
2831
2832	# Change the num_blocks of each rank to the smallest among all ranks.
2833	# We also need to shrink the tensor size proportionally to avoid
2834	# allocating unused memory.
2835	min_num_blocks = min(
2836	kv_cache_config.num_blocks for kv_cache_config in kv_cache_configs
2837	)
2838	for i, kv_cache_config in enumerate(kv_cache_configs):
2839	if kv_cache_config.num_blocks == min_num_blocks:
2840	continue
2841	# Re-plan with exactly the memory the smallest rank can afford, so
2842	# strides and offsets stay consistent with the shrunken allocation.
2843	groups = kv_cache_config.kv_cache_groups
2844	kv_cache_configs[i] = get_kv_cache_config_from_groups(
2845	vllm_config, groups, min_num_blocks * _pool_bytes_per_block(groups)
2846	)
2847
2848	for kv_cache_config in kv_cache_configs:
2849	kv_cache_config.kv_tp_replicas = kv_cache_groups_tp_replicas(
2850	kv_cache_config.kv_cache_groups,
2851	vllm_config.parallel_config.tensor_parallel_size,
2852	vllm_config.parallel_config.decode_context_parallel_size,
2853	)
2854
2855	return kv_cache_configs
2856
2857
2858	class BlockHashListWithBlockSize:
2859	"""Convert block-hash granularity from `hash_block_size` to `target_block_size`.
2860	Used when KV cache groups have different block sizes: `hash_block_size`
2861	is the size used to compute the original `block_hashes`; `target_block_size`
2862	is the group's actual block size.
2863
2864	Currently, only scaling up by an integer factor is supported (i.e.,
2865	`target_block_size` is a multiple of `hash_block_size`). Conversion is
2866	performed lazily on access for efficiency. Each `hash_block_size` hash is
2867	already chained over its entire prefix, so the hash at the last
2868	`hash_block_size` boundary of a `target_block_size` block uniquely
2869	fingerprints that block's prefix; we use it directly.
2870
2871	Example (`hash_block_size` = 16, `target_block_size` = 32):
2872	the second 16-size hash already covers tokens 0-31, so it is the 32-size
2873	hash:
2874
2875	Block hashes with block_size 16:
2876	\| Token Range \| 0-15 \| 16-31 \| 32-47 \| 48-63 \|
2877	\|-------------\|------\|-------\|-------\|-------\|
2878	\| Hash \| A \| B \| C \| D \|
2879
2880	Block hashes with block_size 32:
2881	\| Token Range \| 0-31 \| 32-63 \|
2882	\|-------------\|------\|-------\|
2883	\| Hash \| B \| D \|
2884
2885	Args:
2886	block_hashes: Block hashes to convert, computed at `hash_block_size`.
2887	hash_block_size: Block size at which `block_hashes` were computed.
2888	target_block_size: Desired block size; must be a multiple of `hash_block_size`.
2889
2890	"""
2891
2892	def __init__(
2893	self,
2894	block_hashes: list[BlockHash],
2895	hash_block_size: int,
2896	target_block_size: int,
2897	):
2898	self.block_hashes = block_hashes
2899	assert target_block_size % hash_block_size == 0
2900	self.scale_factor = target_block_size // hash_block_size
2901
2902	def __len__(self) -> int:
2903	return len(self.block_hashes) // self.scale_factor
2904
2905	@overload
2906	def __getitem__(self, idx: int) -> BlockHash: ...
2907
2908	@overload
2909	def __getitem__(self, idx: slice) -> list[BlockHash]: ...
2910
2911	def __getitem__(self, idx):
2912	if isinstance(idx, int):
2913	return self._get_value_at(idx)
2914
2915	if isinstance(idx, slice):
2916	start, stop, step = idx.indices(len(self))
2917	return [self._get_value_at(i) for i in range(start, stop, step)]
2918
2919	raise TypeError(f"Invalid index type: {type(idx)!r}")
2920
2921	def __iter__(self) -> Iterator[BlockHash]:
2922	for i in range(len(self)):
2923	yield self._get_value_at(i)
2924
2925	def _get_value_at(self, idx: int) -> BlockHash:
2926	# The last hash_block_size hash within the target block already chains
2927	# over the whole prefix, so it is the target block's hash.
2928	return self.block_hashes[(idx + 1) * self.scale_factor - 1]
2929
2930
2931	BlockHashList = list[BlockHash] \| BlockHashListWithBlockSize
2932
2933
2934	def resolve_block_hashes(
2935	block_hashes: BlockHashList,
2936	hash_block_size: int,
2937	block_size: int,
2938	*,
2939	supports_fine_grained_hash_lookup: bool = False,
2940	alignment_tokens: int \| None = None,
2941	) -> BlockHashList:
2942	"""Resolve the block-hash view at ``block_size``.
2943
2944	When ``block_size`` equals ``hash_block_size``, reuse the precomputed block
2945	hashes directly; otherwise view them at ``block_size`` granularity.
2946	Fine-grained lookup keeps the original hashes for partial cache hits.
2947	"""
2948	if block_size == hash_block_size:
2949	return block_hashes
2950	if isinstance(block_hashes, BlockHashListWithBlockSize):
2951	# Already a block-size view
2952	assert block_hashes.scale_factor == block_size // hash_block_size
2953	return block_hashes
2954	# Fine-grained partial hits keep the raw hashes. The caller passes
2955	# alignment_tokens = hash_block_size to enable them, else >= block_size.
2956	if (
2957	supports_fine_grained_hash_lookup
2958	and alignment_tokens is not None
2959	and alignment_tokens < block_size
2960	and block_size % alignment_tokens == 0
2961	):
2962	return block_hashes
2963	assert block_size % hash_block_size == 0
2964	return BlockHashListWithBlockSize(block_hashes, hash_block_size, block_size)