| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from __future__ import annotations |
| 5 | |
| 6 | import copy |
| 7 | from collections import Counter |
| 8 | from collections.abc import Collection, Sequence |
| 9 | from dataclasses import dataclass, field, fields, replace |
| 10 | from enum import Enum, IntEnum |
| 11 | from fractions import Fraction |
| 12 | from functools import cached_property |
| 13 | from math import prod |
| 14 | from typing import TYPE_CHECKING, TypeVar |
| 15 | |
| 16 | import torch |
| 17 | from typing_extensions import Self |
| 18 | |
| 19 | from vllm.logger import init_logger |
| 20 | from vllm.utils.math_utils import cdiv, round_up |
| 21 | from vllm.utils.torch_utils import get_dtype_size |
| 22 | from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum |
| 23 | from vllm.v1.kv_cache_layout import _DIM_B, _DIM_L, KVCacheLayout |
| 24 | from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry |
| 25 | |
| 26 | if TYPE_CHECKING: |
| 27 | from vllm.config import VllmConfig |
| 28 | |
| 29 | logger = init_logger(__name__) |
| 30 | |
| 31 | _SpecT = TypeVar("_SpecT", bound="KVCacheSpec") |
| 32 | |
| 33 | |
| 34 | # --------------------------------------------------------------------------- |
| 35 | # KV cache quantization mode |
| 36 | # --------------------------------------------------------------------------- |
| 37 | |
| 38 | |
| 39 | class KVQuantMode(IntEnum): |
| 40 | """KV cache quantization mode. |
| 41 | |
| 42 | Used by attention backends and kernels to dispatch quantization logic |
| 43 | without string matching on ``kv_cache_dtype``. |
| 44 | """ |
| 45 | |
| 46 | NONE = 0 |
| 47 | FP8_PER_TENSOR = 1 # per-tensor scales (current fp8 path) |
| 48 | INT8_PER_TOKEN_HEAD = 2 # per-token-head dynamic scales for int8 |
| 49 | FP8_PER_TOKEN_HEAD = 3 # per-token-head dynamic scales for fp8 |
| 50 | INT4_PER_TOKEN_HEAD = 4 # packed 2×int4/byte, RHT + asymmetric zp |
| 51 | NVFP4 = 5 # packed fp4 data + fp8 block scales |
| 52 | # Hadamard-rotated Lloyd-Max quant, packed K+V per slot. |
| 53 | TURBOQUANT_K8V4 = 6 |
| 54 | TURBOQUANT_4BIT_NC = 7 |
| 55 | TURBOQUANT_K3V4_NC = 8 |
| 56 | TURBOQUANT_3BIT_NC = 9 |
| 57 | NVFP4_DS_MLA = 10 # opaque-bytes NVFP4 DS-MLA layouts (FlashMLA sparse) |
| 58 | |
| 59 | @property |
| 60 | def is_per_token_head(self) -> bool: |
| 61 | """True for any per-token-head quantization mode.""" |
| 62 | return self in ( |
| 63 | KVQuantMode.INT8_PER_TOKEN_HEAD, |
| 64 | KVQuantMode.FP8_PER_TOKEN_HEAD, |
| 65 | KVQuantMode.INT4_PER_TOKEN_HEAD, |
| 66 | ) |
| 67 | |
| 68 | @property |
| 69 | def is_nvfp4(self) -> bool: |
| 70 | """True for NVFP4 packed quantization mode.""" |
| 71 | return self == KVQuantMode.NVFP4 |
| 72 | |
| 73 | @property |
| 74 | def is_turboquant(self) -> bool: |
| 75 | """True for any turboquant quantization mode.""" |
| 76 | return self in ( |
| 77 | KVQuantMode.TURBOQUANT_K8V4, |
| 78 | KVQuantMode.TURBOQUANT_4BIT_NC, |
| 79 | KVQuantMode.TURBOQUANT_K3V4_NC, |
| 80 | KVQuantMode.TURBOQUANT_3BIT_NC, |
| 81 | ) |
| 82 | |
| 83 | |
| 84 | def get_kv_quant_mode(kv_cache_dtype: str) -> KVQuantMode: |
| 85 | """Map a ``kv_cache_dtype`` string to a :class:`KVQuantMode`.""" |
| 86 | if kv_cache_dtype == "int4_per_token_head": |
| 87 | return KVQuantMode.INT4_PER_TOKEN_HEAD |
| 88 | if kv_cache_dtype == "int8_per_token_head": |
| 89 | return KVQuantMode.INT8_PER_TOKEN_HEAD |
| 90 | if kv_cache_dtype == "fp8_per_token_head": |
| 91 | return KVQuantMode.FP8_PER_TOKEN_HEAD |
| 92 | # Must precede the ``nvfp4`` prefix test below, which would otherwise match. |
| 93 | if kv_cache_dtype == "nvfp4_ds_mla": |
| 94 | # Page size is keyed on cache_dtype_str in the MLA specs, not |
| 95 | # nvfp4_kv_cache_full_dim. |
| 96 | return KVQuantMode.NVFP4_DS_MLA |
| 97 | if kv_cache_dtype.startswith("nvfp4"): |
| 98 | return KVQuantMode.NVFP4 |
| 99 | if isinstance(kv_cache_dtype, str) and kv_cache_dtype.startswith("turboquant_"): |
| 100 | return KVQuantMode[kv_cache_dtype.upper()] |
| 101 | if isinstance(kv_cache_dtype, str) and kv_cache_dtype.startswith("fp8"): |
| 102 | return KVQuantMode.FP8_PER_TENSOR |
| 103 | return KVQuantMode.NONE |
| 104 | |
| 105 | |
| 106 | def is_quantized_kv_cache(kv_cache_dtype: str) -> bool: |
| 107 | return get_kv_quant_mode(kv_cache_dtype) != KVQuantMode.NONE |
| 108 | |
| 109 | |
| 110 | def replace_as( |
| 111 | spec: KVCacheSpec, |
| 112 | target_cls: type[_SpecT], |
| 113 | *, |
| 114 | drop: Collection[str] = (), |
| 115 | **changes, |
| 116 | ) -> _SpecT: |
| 117 | """``dataclasses.replace``, but rebuilding *spec* as *target_cls* |
| 118 | e.g. ``SlidingWindowSpec`` -> ``FullAttentionSpec`` |
| 119 | |
| 120 | Every field of *spec* must exist on *target_cls* unless named in *drop*; |
| 121 | fields only *target_cls* has keep their default values. |
| 122 | """ |
| 123 | kwargs = { |
| 124 | f.name: getattr(spec, f.name) |
| 125 | for f in fields(spec) |
| 126 | if f.init and f.name not in drop |
| 127 | } |
| 128 | kwargs.update(changes) |
| 129 | return target_cls(**kwargs) |
| 130 | |
| 131 | |
| 132 | def kv_cache_uses_per_token_head_scales(kv_cache_dtype: str) -> bool: |
| 133 | """Return True if *kv_cache_dtype* needs per-token-head scales.""" |
| 134 | return get_kv_quant_mode(kv_cache_dtype).is_per_token_head |
| 135 | |
| 136 | |
| 137 | class KVCacheSpecKind(str, Enum): |
| 138 | FULL_ATTENTION = "full_attention" |
| 139 | MLA_ATTENTION = "mla_attention" |
| 140 | SLIDING_WINDOW = "sliding_window" |
| 141 | SLIDING_WINDOW_MLA = "sliding_window_mla" |
| 142 | MAMBA = "mamba" |
| 143 | CHUNKED_LOCAL_ATTENTION = "chunked_local_attention" |
| 144 | SINK_FULL_ATTENTION = "sink_full_attention" |
| 145 | ENCODER_ONLY_ATTENTION = "encoder_only_attention" |
| 146 | CROSS_ATTENTION = "cross_attention" |
| 147 | UNKNOWN = "unknown" |
| 148 | |
| 149 | |
| 150 | class SparseCacheRole(str, Enum): |
| 151 | SPARSE = "sparse" |
| 152 | INDEXER = "indexer" |
| 153 | |
| 154 | |
| 155 | @dataclass(frozen=True) |
| 156 | class KVCacheSpec: |
| 157 | """A base class for specifying the KV cache format of one layer.""" |
| 158 | |
| 159 | # number of tokens in a block |
| 160 | block_size: int |
| 161 | dcp_sharded: bool = field(default=False, kw_only=True) |
| 162 | """Whether DCP shards this cache's token positions across ranks.""" |
| 163 | |
| 164 | block_stride_alignment: int | None = field(default=None, kw_only=True) |
| 165 | """Required byte alignment between physical blocks, including packed layers.""" |
| 166 | |
| 167 | def __post_init__(self): |
| 168 | if self.block_stride_alignment is not None and self.block_stride_alignment <= 0: |
| 169 | raise ValueError("block_stride_alignment must be positive") |
| 170 | |
| 171 | @property |
| 172 | def prefix_cacheable(self) -> bool: |
| 173 | """Whether this spec's group participates in prefix caching.""" |
| 174 | return True |
| 175 | |
| 176 | @property |
| 177 | def prefix_replay_tokens(self) -> int: |
| 178 | """DeepSeek-V4.1 only: bounded replay. Currently only for DSV41 SWA.""" |
| 179 | return 0 |
| 180 | |
| 181 | @property |
| 182 | def num_heads(self) -> int: |
| 183 | raise NotImplementedError |
| 184 | |
| 185 | @property |
| 186 | def tokens_per_state(self) -> int | Fraction: |
| 187 | raise NotImplementedError |
| 188 | |
| 189 | @property |
| 190 | def state_content_size_bytes(self) -> int: |
| 191 | raise NotImplementedError |
| 192 | |
| 193 | @property |
| 194 | def page_size_bytes(self) -> int: |
| 195 | """The size of a page with `block_size` tokens in bytes. |
| 196 | |
| 197 | Returns: |
| 198 | The page size |
| 199 | |
| 200 | """ |
| 201 | raise NotImplementedError |
| 202 | |
| 203 | @property |
| 204 | def num_states(self) -> int: |
| 205 | return self.get_num_kernel_states(self.block_size) |
| 206 | |
| 207 | def get_num_kernel_states(self, kernel_block_size: int) -> int: |
| 208 | if self.tokens_per_state > 0: |
| 209 | return kernel_block_size // self.tokens_per_state |
| 210 | return 1 |
| 211 | |
| 212 | @property |
| 213 | def block_table_token_alignment(self) -> int | None: |
| 214 | return 128 |
| 215 | |
| 216 | def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int: |
| 217 | """The maximum possible memory usage of this KV cache in bytes. |
| 218 | |
| 219 | Returns: |
| 220 | The KV cache size in bytes |
| 221 | |
| 222 | """ |
| 223 | raise NotImplementedError |
| 224 | |
| 225 | def max_num_blocks_per_req(self, vllm_config: VllmConfig, max_len: int) -> int: |
| 226 | """The number of block table entries needed per request, i.e. the row |
| 227 | length of the worker-side block table for this cache group. |
| 228 | |
| 229 | Args: |
| 230 | vllm_config: The vllm config. |
| 231 | max_len: The maximum sequence length to size for, including the |
| 232 | encoder length for encoder-decoder models. |
| 233 | |
| 234 | """ |
| 235 | return cdiv(max_len, self.block_size) |
| 236 | |
| 237 | def copy_with_new_block_size(self, block_size: int) -> Self: |
| 238 | """Create a new KVCacheSpec from self but replacing the block size.""" |
| 239 | return replace(self, block_size=block_size) |
| 240 | |
| 241 | @classmethod |
| 242 | def merge(cls, specs: list[Self]) -> Self: |
| 243 | """Merge a list of KVCacheSpec objects into a single KVCacheSpec object.""" |
| 244 | if not all(spec == specs[0] for spec in specs[1:]): |
| 245 | raise AssertionError( |
| 246 | "All layers in the same KV cache group must be the same." |
| 247 | ) |
| 248 | return copy.deepcopy(specs[0]) |
| 249 | |
| 250 | def is_uniform_with_collection( |
| 251 | self, kv_cache_specs: dict[str, KVCacheSpec] |
| 252 | ) -> bool: |
| 253 | """Whether this KVCacheSpec is uniform with all specs of all layers.""" |
| 254 | uniform_type_base_spec = KVCacheSpecRegistry.get_uniform_type_base_spec(self) |
| 255 | assert uniform_type_base_spec is not None, ( |
| 256 | f"Unsupported KV cache spec type: {type(self)}. " |
| 257 | "Please register it using @register_kv_cache_spec decorator." |
| 258 | ) |
| 259 | return all( |
| 260 | isinstance(spec, uniform_type_base_spec) for spec in kv_cache_specs.values() |
| 261 | ) |
| 262 | |
| 263 | @property |
| 264 | def has_layer_views(self) -> bool: |
| 265 | """Whether generic allocation creates per-layer cache views. |
| 266 | |
| 267 | Specs without a per-layer shape keep the raw backing tensor and lay it |
| 268 | out themselves. They also have no attention module of their own, so |
| 269 | they take no part in attention-backend or kernel-block selection. |
| 270 | """ |
| 271 | return True |
| 272 | |
| 273 | @property |
| 274 | def uses_slot_mapping(self) -> bool: |
| 275 | """Whether the worker computes a per-token slot mapping for this spec. |
| 276 | |
| 277 | Specs that address their pages themselves (raw storage, ring buffers) |
| 278 | take no slot mapping row. |
| 279 | """ |
| 280 | return True |
| 281 | |
| 282 | |
| 283 | def group_kernel_blocks(cache: torch.Tensor, num_blocks: int) -> torch.Tensor: |
| 284 | """View a kernel-block-granular layer cache with manager blocks as dim 0. |
| 285 | |
| 286 | Kernel block splitting subdivides each manager block into uniformly strided |
| 287 | kernel blocks, so grouping is a pure view: ``(num_blocks * ratio, ...)`` |
| 288 | """ |
| 289 | if cache.shape[0] == num_blocks: |
| 290 | return cache |
| 291 | assert cache.shape[0] % num_blocks == 0 |
| 292 | return cache.unflatten(0, (num_blocks, -1)) |
| 293 | |
| 294 | |
| 295 | def compute_layer_kv_cache_shape_bytes( |
| 296 | spec: KVCacheSpec, |
| 297 | num_blocks: int, |
| 298 | kernel_block_size: int | None = None, |
| 299 | ) -> tuple[int, ...]: |
| 300 | """Return the 4D logical shape ``(B, H, N, C)`` where C is in bytes.""" |
| 301 | bs = kernel_block_size if kernel_block_size is not None else spec.block_size |
| 302 | assert spec.block_size % bs == 0, ( |
| 303 | f"Kernel block size {bs} must divide KV cache block size {spec.block_size}." |
| 304 | ) |
| 305 | blocks_per_page = spec.block_size // bs |
| 306 | return ( |
| 307 | num_blocks * blocks_per_page, |
| 308 | spec.num_heads, |
| 309 | spec.get_num_kernel_states(bs), |
| 310 | spec.state_content_size_bytes, |
| 311 | ) |
| 312 | |
| 313 | |
| 314 | def compute_layout_strides( |
| 315 | spec: KVCacheSpec, |
| 316 | num_blocks: int, |
| 317 | num_layers: int, |
| 318 | layout: KVCacheLayout, |
| 319 | kernel_block_size: int | None = None, |
| 320 | fixed_strides: tuple[int | None, ...] = (None,) * 5, |
| 321 | ) -> tuple[int, ...]: |
| 322 | """Byte strides in logical ``[L, B, H, N, C]`` axis order.""" |
| 323 | assert len(fixed_strides) == 5 |
| 324 | assert all(stride is None or stride > 0 for stride in fixed_strides) |
| 325 | shape = ( |
| 326 | num_layers, |
| 327 | *compute_layer_kv_cache_shape_bytes(spec, num_blocks, kernel_block_size), |
| 328 | ) |
| 329 | order = layout.stride_order |
| 330 | padded_page_size = getattr(spec, "page_size_padded", None) |
| 331 | if padded_page_size is not None: |
| 332 | assert kernel_block_size is None or kernel_block_size == spec.block_size, ( |
| 333 | "Padded KV pages do not support kernel block splitting." |
| 334 | ) |
| 335 | page_grid_end = max(order.index(_DIM_L), order.index(_DIM_B)) + 1 |
| 336 | page_grid_shape = tuple(shape[dim] for dim in order[:page_grid_end]) |
| 337 | assert prod(page_grid_shape) == num_layers * num_blocks, ( |
| 338 | "Page padding requires dimensions outside the page tail to be L, B, " |
| 339 | f"or singleton; got {layout.name} with shape {shape}." |
| 340 | ) |
| 341 | |
| 342 | strides = [0] * 5 |
| 343 | current_stride = 1 |
| 344 | for physical_idx, dim in reversed(tuple(enumerate(order))): |
| 345 | if padded_page_size is not None and physical_idx == page_grid_end - 1: |
| 346 | current_stride = max(current_stride, padded_page_size) |
| 347 | assert current_stride % padded_page_size == 0 |
| 348 | strides[dim] = fixed_strides[dim] or current_stride |
| 349 | current_stride = strides[dim] * shape[dim] |
| 350 | return tuple(strides) |
| 351 | |
| 352 | |
| 353 | def create_kv_cache_views( |
| 354 | raw: torch.Tensor, |
| 355 | spec: KVCacheSpec, |
| 356 | num_blocks: int, |
| 357 | layout: KVCacheLayout, |
| 358 | kv_cache_tensor: KVCacheTensor, |
| 359 | kernel_block_size: int | None = None, |
| 360 | ) -> list[torch.Tensor]: |
| 361 | """View a flat int8 buffer as one 4D ``[B, H, N, C]`` view per layer. |
| 362 | |
| 363 | Block ``b`` of layer ``l`` starts at the tensor offset plus its layer and |
| 364 | block stride contributions. |
| 365 | """ |
| 366 | num_layers = len(kv_cache_tensor.layers) |
| 367 | layer_stride = kv_cache_tensor.layer_stride |
| 368 | block_stride = kv_cache_tensor.block_stride |
| 369 | shape_bytes = compute_layer_kv_cache_shape_bytes( |
| 370 | spec, num_blocks, kernel_block_size |
| 371 | ) |
| 372 | ratio = shape_bytes[0] // num_blocks |
| 373 | if ratio > 1: |
| 374 | # Kernel blocks subdivide a manager block into `ratio` equal pieces, so |
| 375 | # they sit a constant stride apart only if a block is one dense page: no |
| 376 | # padding at its end, and no other layer's page before the next block. |
| 377 | dense_page_size = prod(compute_layer_kv_cache_shape_bytes(spec, 1)[1:]) |
| 378 | if block_stride != dense_page_size: |
| 379 | raise ValueError( |
| 380 | f"The resolved KV cache layout ({layout.name}) does not store " |
| 381 | "blocks as dense, unpadded pages (block stride " |
| 382 | f"{block_stride} != page {dense_page_size}), so a manager " |
| 383 | f"block cannot be split into {ratio} kernel blocks of " |
| 384 | f"{kernel_block_size} tokens. Reduce --block-size to " |
| 385 | f"{kernel_block_size} or set VLLM_KV_CACHE_LAYOUT to a " |
| 386 | "layer-compact layout (e.g. LBNHC)." |
| 387 | ) |
| 388 | assert block_stride % ratio == 0, ( |
| 389 | f"Block stride {block_stride} must divide into {ratio} equal kernel blocks." |
| 390 | ) |
| 391 | block_stride //= ratio |
| 392 | |
| 393 | logical_shape = (num_layers, *shape_bytes) |
| 394 | strides = compute_layout_strides( |
| 395 | spec, |
| 396 | num_blocks, |
| 397 | num_layers, |
| 398 | layout, |
| 399 | kernel_block_size, |
| 400 | fixed_strides=(layer_stride, block_stride, None, None, None), |
| 401 | ) |
| 402 | dtype = getattr(spec, "dtype", None) |
| 403 | |
| 404 | view_5d = torch.as_strided( |
| 405 | raw, |
| 406 | size=logical_shape, |
| 407 | stride=strides, |
| 408 | storage_offset=raw.storage_offset() + kv_cache_tensor.offset, |
| 409 | ) |
| 410 | |
| 411 | views = [] |
| 412 | for layer_idx in range(num_layers): |
| 413 | cache_logical = view_5d[layer_idx] |
| 414 | if dtype is not None: |
| 415 | cache_logical = cache_logical.view(dtype) |
| 416 | views.append(cache_logical) |
| 417 | return views |
| 418 | |
| 419 | |
| 420 | @dataclass(frozen=True, kw_only=True) |
| 421 | class HiSparseHotSpec(KVCacheSpec): |
| 422 | """Ephemeral per-request HiSparse hot-cache allocation.""" |
| 423 | |
| 424 | page_size: int |
| 425 | blocks_per_request: int |
| 426 | |
| 427 | @property |
| 428 | def prefix_cacheable(self) -> bool: |
| 429 | return False |
| 430 | |
| 431 | @property |
| 432 | def page_size_bytes(self) -> int: |
| 433 | return self.page_size |
| 434 | |
| 435 | @property |
| 436 | def block_table_token_alignment(self) -> int | None: |
| 437 | return None |
| 438 | |
| 439 | def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int: |
| 440 | return self.blocks_per_request * self.page_size |
| 441 | |
| 442 | def max_num_blocks_per_req(self, vllm_config: VllmConfig, max_len: int) -> int: |
| 443 | return self.blocks_per_request |
| 444 | |
| 445 | @property |
| 446 | def has_layer_views(self) -> bool: |
| 447 | return False |
| 448 | |
| 449 | @property |
| 450 | def uses_slot_mapping(self) -> bool: |
| 451 | return False |
| 452 | |
| 453 | |
| 454 | @dataclass(frozen=True, kw_only=True) |
| 455 | class HiSparseResidentSpec(KVCacheSpec): |
| 456 | """Reclaimable GPU-resident pages for host-backed HiSparse KV.""" |
| 457 | |
| 458 | page_size: int |
| 459 | |
| 460 | @property |
| 461 | def prefix_cacheable(self) -> bool: |
| 462 | return False |
| 463 | |
| 464 | @property |
| 465 | def page_size_bytes(self) -> int: |
| 466 | return self.page_size |
| 467 | |
| 468 | def max_admission_blocks_per_request( |
| 469 | self, max_in_flight_tokens: int, max_model_len: int |
| 470 | ) -> int: |
| 471 | num_tokens = min(max_in_flight_tokens, max_model_len) |
| 472 | return cdiv(num_tokens, self.block_size) |
| 473 | |
| 474 | def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int: |
| 475 | max_blocks = self.max_admission_blocks_per_request( |
| 476 | max_in_flight_tokens=vllm_config.max_in_flight_tokens, |
| 477 | max_model_len=vllm_config.model_config.max_model_len, |
| 478 | ) |
| 479 | return max_blocks * self.page_size |
| 480 | |
| 481 | @property |
| 482 | def has_layer_views(self) -> bool: |
| 483 | return False |
| 484 | |
| 485 | |
| 486 | @dataclass(frozen=True, kw_only=True) |
| 487 | class AttentionSpec(KVCacheSpec): |
| 488 | dcp_sharded: bool = True |
| 489 | max_tp_shards: int | None = None |
| 490 | """Distinct shards this cache splits into across TP; TP ranks beyond this |
| 491 | hold replicas. None: unknown, treated as not replicated.""" |
| 492 | num_kv_heads: int |
| 493 | head_size: int |
| 494 | dtype: torch.dtype |
| 495 | head_size_v: int = None # type: ignore[assignment] |
| 496 | kv_quant_mode: KVQuantMode = KVQuantMode.NONE |
| 497 | page_size_padded: int | None = None |
| 498 | num_head_slots: int | None = None |
| 499 | """H of the logical ``[B, H, N, C]`` page when packing diverges from one |
| 500 | slot per KV head. None means one slot per KV head. Published by the backend. |
| 501 | """ |
| 502 | state_content_bytes: int | None = None |
| 503 | """C in bytes when packed; None means dense K/V content.""" |
| 504 | tokens_per_state: int | Fraction = 1 |
| 505 | """Tokens covered by one stored state. Ints > 1 compress multiple tokens |
| 506 | into one state (DSv4 sparse MLA); fractions < 1 store multiple states per |
| 507 | token (Whisper block pooling: ``Fraction(1, block_pool_size)``).""" |
| 508 | |
| 509 | def __post_init__(self): |
| 510 | super().__post_init__() |
| 511 | if self.head_size_v is None: |
| 512 | object.__setattr__(self, "head_size_v", self.head_size) |
| 513 | |
| 514 | @property |
| 515 | def num_heads(self) -> int: |
| 516 | if self.num_head_slots is not None: |
| 517 | return self.num_head_slots |
| 518 | return self.num_kv_heads |
| 519 | |
| 520 | @property |
| 521 | def state_content_size_bytes(self) -> int: |
| 522 | """Bytes per (head slot, stored state) cell of the page.""" |
| 523 | if self.state_content_bytes is not None: |
| 524 | return self.state_content_bytes |
| 525 | return (self.head_size + self.head_size_v) * get_dtype_size(self.dtype) |
| 526 | |
| 527 | @property |
| 528 | def unpadded_page_size_bytes(self) -> int: |
| 529 | return self.num_heads * self.num_states * self.state_content_size_bytes |
| 530 | |
| 531 | @property |
| 532 | def page_size_bytes(self) -> int: |
| 533 | if self.page_size_padded is not None: |
| 534 | assert self.page_size_padded >= self.unpadded_page_size_bytes |
| 535 | return self.page_size_padded |
| 536 | return self.unpadded_page_size_bytes |
| 537 | |
| 538 | @property |
| 539 | def real_page_size_bytes(self) -> int: |
| 540 | """Alias of ``unpadded_page_size_bytes`` |
| 541 | TODO(lucas): follow up with TPU backend to see if we can remove this property. |
| 542 | """ |
| 543 | return self.unpadded_page_size_bytes |
| 544 | |
| 545 | def max_num_blocks_per_req(self, vllm_config: VllmConfig, max_len: int) -> int: |
| 546 | parallel_config = vllm_config.parallel_config |
| 547 | kv_shard_count = ( |
| 548 | parallel_config.decode_context_parallel_size if self.dcp_sharded else 1 |
| 549 | ) |
| 550 | return cdiv(max_len, self.block_size * kv_shard_count) |
| 551 | |
| 552 | |
| 553 | @dataclass(frozen=True, kw_only=True) |
| 554 | class FullAttentionSpec(AttentionSpec): |
| 555 | """When hybrid allocator is disabled and the model contains both full |
| 556 | attention layers and sliding window attention layers, sliding |
| 557 | window attention are regarded as full attention in KV cache manager |
| 558 | (blocks are allocated for all tokens), while computed as sliding window |
| 559 | attention in model runner. |
| 560 | In this case, we use FullAttentionSpec and record the sliding window size. |
| 561 | """ |
| 562 | |
| 563 | sliding_window: int | None = None |
| 564 | """ |
| 565 | Default to None for not using sliding window attention. |
| 566 | """ |
| 567 | attention_chunk_size: int | None = None |
| 568 | |
| 569 | non_causal: bool = False |
| 570 | """ |
| 571 | Whether the layer attends non-causally (e.g. Prefix LM). Carried on the |
| 572 | spec so the engine core, which collects specs from all workers before the |
| 573 | scheduler is built, can adjust scheduling policy (chunked prefill / prefix |
| 574 | caching) regardless of tensor-parallel layout. It does not affect the KV |
| 575 | cache layout itself. |
| 576 | """ |
| 577 | |
| 578 | def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int: |
| 579 | max_model_len = vllm_config.model_config.max_model_len |
| 580 | dcp_world_size = vllm_config.parallel_config.decode_context_parallel_size |
| 581 | if self.dcp_sharded and dcp_world_size > 1: |
| 582 | max_model_len = cdiv(max_model_len, dcp_world_size) |
| 583 | return cdiv(max_model_len, self.block_size) * self.page_size_bytes |
| 584 | |
| 585 | @classmethod |
| 586 | def merge_window_sizes(cls, window_sizes: set[int]) -> int | None: |
| 587 | if len(window_sizes) == 0: |
| 588 | return None |
| 589 | elif len(window_sizes) == 1: |
| 590 | return window_sizes.pop() |
| 591 | else: |
| 592 | raise ValueError( |
| 593 | "All attention layers in the same KV cache group must have the " |
| 594 | "same window size." |
| 595 | ) |
| 596 | |
| 597 | @classmethod |
| 598 | def merge(cls, specs: list[Self]) -> Self: |
| 599 | """Merge a list of FullAttentionSpec objects into a single |
| 600 | FullAttentionSpec object. |
| 601 | """ |
| 602 | assert all(isinstance(spec, FullAttentionSpec) for spec in specs), ( |
| 603 | "All attention layers in the same KV cache group must be FullAttentionSpec." |
| 604 | ) |
| 605 | |
| 606 | sliding_window = set( |
| 607 | spec.sliding_window for spec in specs if spec.sliding_window is not None |
| 608 | ) |
| 609 | attention_chunk_size = set( |
| 610 | spec.attention_chunk_size |
| 611 | for spec in specs |
| 612 | if spec.attention_chunk_size is not None |
| 613 | ) |
| 614 | assert not any(isinstance(spec, MLAAttentionSpec) for spec in specs), ( |
| 615 | "MLAAttentionSpec should be merged in MLAAttentionSpec.merge" |
| 616 | ) |
| 617 | merged_spec = cls( |
| 618 | block_size=specs[0].block_size, |
| 619 | block_stride_alignment=specs[0].block_stride_alignment, |
| 620 | num_kv_heads=specs[0].num_kv_heads, |
| 621 | head_size=specs[0].head_size, |
| 622 | head_size_v=specs[0].head_size_v, |
| 623 | dtype=specs[0].dtype, |
| 624 | kv_quant_mode=specs[0].kv_quant_mode, |
| 625 | dcp_sharded=specs[0].dcp_sharded, |
| 626 | max_tp_shards=specs[0].max_tp_shards, |
| 627 | page_size_padded=specs[0].page_size_padded, |
| 628 | num_head_slots=specs[0].num_head_slots, |
| 629 | state_content_bytes=specs[0].state_content_bytes, |
| 630 | tokens_per_state=specs[0].tokens_per_state, |
| 631 | sliding_window=cls.merge_window_sizes(sliding_window), |
| 632 | attention_chunk_size=cls.merge_window_sizes(attention_chunk_size), |
| 633 | # If any layer in the group is non-causal, treat the group as |
| 634 | # non-causal so the engine core disables incompatible scheduling. |
| 635 | non_causal=any(spec.non_causal for spec in specs), |
| 636 | ) |
| 637 | for spec in specs: |
| 638 | for f in fields(AttentionSpec): |
| 639 | assert getattr(spec, f.name) == getattr(merged_spec, f.name), ( |
| 640 | "All attention layers in the same KV cache group must have " |
| 641 | "the same attention spec." |
| 642 | ) |
| 643 | assert (merged_spec.sliding_window is not None) + ( |
| 644 | merged_spec.attention_chunk_size is not None |
| 645 | ) <= 1, ( |
| 646 | "Model with both sliding window layers and chunked local attention " |
| 647 | "layers is not supported." |
| 648 | ) |
| 649 | return merged_spec |
| 650 | |
| 651 | |
| 652 | def _apply_alignment_padding(spec: MLAAttentionSpec | SlidingWindowMLASpec): |
| 653 | if spec.alignment is None: |
| 654 | return |
| 655 | actual_page_size = spec.real_page_size_bytes |
| 656 | padded_page_size = round_up(actual_page_size, spec.alignment) |
| 657 | if padded_page_size != actual_page_size: |
| 658 | object.__setattr__(spec, "page_size_padded", padded_page_size) |
| 659 | |
| 660 | |
| 661 | @dataclass(frozen=True, kw_only=True) |
| 662 | class MLAAttentionSpec(FullAttentionSpec): |
| 663 | # TODO(Lucas/Chen): less hacky way to do this |
| 664 | cache_dtype_str: str | None = None |
| 665 | # DeepseekV4 only fields. Non-DeepseekV4 MLA models leave these at defaults. |
| 666 | alignment: int | None = None # Default to None for no padding. |
| 667 | model_version: str | None = None |
| 668 | cache_role: SparseCacheRole = SparseCacheRole.SPARSE |
| 669 | is_index_group_leader: bool = False |
| 670 | storage_block_size: int | None = None |
| 671 | """Token width used to view storage when it differs from the kernel block.""" |
| 672 | # Group capability enabled when any member flattens a non-causal query block |
| 673 | # into decode rows. Runtime metadata still selects causal vs. non-causal mode. |
| 674 | non_causal_multi_token_decode: bool = False |
| 675 | # MLA stores a single latent vector per state; there is no separate V. |
| 676 | head_size_v: int = 0 |
| 677 | |
| 678 | def __post_init__(self): |
| 679 | super().__post_init__() |
| 680 | _apply_alignment_padding(self) |
| 681 | |
| 682 | @classmethod |
| 683 | def merge(cls, specs: list[Self]) -> Self: |
| 684 | assert all(isinstance(spec, MLAAttentionSpec) for spec in specs), ( |
| 685 | "All attention layers in the same KV cache group must be MLAAttentionSpec." |
| 686 | ) |
| 687 | cache_dtype_str_set = set(spec.cache_dtype_str for spec in specs) |
| 688 | tokens_per_state_set = set(spec.tokens_per_state for spec in specs) |
| 689 | model_version_set = set(spec.model_version for spec in specs) |
| 690 | cache_role_set = {spec.cache_role for spec in specs} |
| 691 | index_group_leader_set = {spec.is_index_group_leader for spec in specs} |
| 692 | storage_block_size_set = set(spec.storage_block_size for spec in specs) |
| 693 | block_stride_alignment_set = {spec.block_stride_alignment for spec in specs} |
| 694 | assert ( |
| 695 | len(cache_dtype_str_set) == 1 |
| 696 | and len(tokens_per_state_set) == 1 |
| 697 | and len(model_version_set) == 1 |
| 698 | and len(cache_role_set) == 1 |
| 699 | and len(index_group_leader_set) == 1 |
| 700 | and len(storage_block_size_set) == 1 |
| 701 | and len(block_stride_alignment_set) == 1 |
| 702 | ), ( |
| 703 | "All attention layers in the same KV cache group must use the same " |
| 704 | "quantization method, tokens per state, model version, cache role, " |
| 705 | "index-sharing role, storage block size and block stride alignment." |
| 706 | ) |
| 707 | merged_spec = cls( |
| 708 | block_size=specs[0].block_size, |
| 709 | num_kv_heads=specs[0].num_kv_heads, |
| 710 | head_size=specs[0].head_size, |
| 711 | dtype=specs[0].dtype, |
| 712 | kv_quant_mode=specs[0].kv_quant_mode, |
| 713 | dcp_sharded=specs[0].dcp_sharded, |
| 714 | max_tp_shards=specs[0].max_tp_shards, |
| 715 | page_size_padded=specs[0].page_size_padded, |
| 716 | num_head_slots=specs[0].num_head_slots, |
| 717 | state_content_bytes=specs[0].state_content_bytes, |
| 718 | cache_dtype_str=cache_dtype_str_set.pop(), |
| 719 | tokens_per_state=tokens_per_state_set.pop(), |
| 720 | model_version=model_version_set.pop(), |
| 721 | cache_role=cache_role_set.pop(), |
| 722 | is_index_group_leader=index_group_leader_set.pop(), |
| 723 | storage_block_size=storage_block_size_set.pop(), |
| 724 | block_stride_alignment=block_stride_alignment_set.pop(), |
| 725 | non_causal_multi_token_decode=any( |
| 726 | spec.non_causal_multi_token_decode for spec in specs |
| 727 | ), |
| 728 | ) |
| 729 | for spec in specs: |
| 730 | for f in fields(AttentionSpec): |
| 731 | assert getattr(spec, f.name) == getattr(merged_spec, f.name), ( |
| 732 | "All attention layers in the same KV cache group must have " |
| 733 | "the same attention spec." |
| 734 | ) |
| 735 | return merged_spec |
| 736 | |
| 737 | |
| 738 | @dataclass(frozen=True, kw_only=True) |
| 739 | class HiddenStateCacheSpec(MLAAttentionSpec): |
| 740 | """Marker for hidden-state cache layers used by extract_hidden_states.""" |
| 741 | |
| 742 | pass |
| 743 | |
| 744 | |
| 745 | @dataclass(frozen=True, kw_only=True) |
| 746 | class RSWASpec(FullAttentionSpec): |
| 747 | """KV cache spec for Reference Sliding Window Attention (R-SWA). |
| 748 | |
| 749 | Prefill (image + text prompt) tokens are always globally visible. |
| 750 | Only the last ``rswa_window`` generated tokens are kept in the KV cache; |
| 751 | gap blocks (between the prefill tail and the current decode window) are |
| 752 | evicted during each decode step to bound memory at |
| 753 | O(prefix_blocks + window_blocks). |
| 754 | """ |
| 755 | |
| 756 | rswa_window: int |
| 757 | |
| 758 | @classmethod |
| 759 | def merge(cls, specs: list[RSWASpec]) -> RSWASpec: |
| 760 | assert all(isinstance(spec, RSWASpec) for spec in specs), ( |
| 761 | "All attention layers in the same KV cache group must be RSWASpec." |
| 762 | ) |
| 763 | rswa_windows = {spec.rswa_window for spec in specs} |
| 764 | assert len(rswa_windows) == 1, ( |
| 765 | f"All R-SWA layers must share the same rswa_window, got {rswa_windows}" |
| 766 | ) |
| 767 | # Delegate common field merging to the parent, then reattach rswa_window. |
| 768 | base = FullAttentionSpec.merge(specs) # type: ignore[arg-type] |
| 769 | return cls( |
| 770 | block_size=base.block_size, |
| 771 | block_stride_alignment=base.block_stride_alignment, |
| 772 | num_kv_heads=base.num_kv_heads, |
| 773 | head_size=base.head_size, |
| 774 | head_size_v=base.head_size_v, |
| 775 | dtype=base.dtype, |
| 776 | kv_quant_mode=base.kv_quant_mode, |
| 777 | dcp_sharded=base.dcp_sharded, |
| 778 | max_tp_shards=base.max_tp_shards, |
| 779 | page_size_padded=base.page_size_padded, |
| 780 | num_head_slots=base.num_head_slots, |
| 781 | state_content_bytes=base.state_content_bytes, |
| 782 | tokens_per_state=base.tokens_per_state, |
| 783 | sliding_window=base.sliding_window, |
| 784 | attention_chunk_size=base.attention_chunk_size, |
| 785 | non_causal=base.non_causal, |
| 786 | rswa_window=rswa_windows.pop(), |
| 787 | ) |
| 788 | |
| 789 | |
| 790 | @dataclass(frozen=True, kw_only=True) |
| 791 | class ChunkedLocalAttentionSpec(AttentionSpec): |
| 792 | attention_chunk_size: int |
| 793 | |
| 794 | def max_admission_blocks_per_request( |
| 795 | self, max_in_flight_tokens: int, max_model_len: int |
| 796 | ) -> int: |
| 797 | """Per-request admission cap, in blocks. |
| 798 | |
| 799 | Single source of truth for both startup pool sizing |
| 800 | (`max_memory_usage_bytes`) and the runtime admission gate, so requests |
| 801 | admitted by startup can also be admitted at runtime. |
| 802 | |
| 803 | `max_in_flight_tokens` is the max tokens scheduled but not yet settled |
| 804 | (one batch per concurrent step); see `VllmConfig.max_in_flight_tokens`. |
| 805 | """ |
| 806 | # During chunked prefill, we hold KV for at most one chunk window plus |
| 807 | # the in-flight tokens, since frees happen on the processed-token basis. |
| 808 | num_tokens = min( |
| 809 | self.attention_chunk_size + max_in_flight_tokens, max_model_len |
| 810 | ) |
| 811 | return cdiv(num_tokens, self.block_size) |
| 812 | |
| 813 | def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int: |
| 814 | max_blocks = self.max_admission_blocks_per_request( |
| 815 | max_in_flight_tokens=vllm_config.max_in_flight_tokens, |
| 816 | max_model_len=vllm_config.model_config.max_model_len, |
| 817 | ) |
| 818 | return max_blocks * self.page_size_bytes |
| 819 | |
| 820 | def is_uniform_with_collection( |
| 821 | self, kv_cache_specs: dict[str, KVCacheSpec] |
| 822 | ) -> bool: |
| 823 | return all( |
| 824 | isinstance(spec, ChunkedLocalAttentionSpec) |
| 825 | and spec.attention_chunk_size == self.attention_chunk_size |
| 826 | for spec in kv_cache_specs.values() |
| 827 | ) |
| 828 | |
| 829 | |
| 830 | @dataclass(frozen=True, kw_only=True) |
| 831 | class SlidingWindowSpec(AttentionSpec): |
| 832 | sliding_window: int |
| 833 | # The trailing edge of the window is extended by ``extra_retained_tokens`` |
| 834 | # so that those extra trailing tokens' blocks are retained (but not |
| 835 | # attended). This is needed for multi-module spec decoding which can |
| 836 | # re-prefill the last num_spec_prefill_tokens - 1 tokens from the end |
| 837 | # of the sequence, and thus needs to delay freeing/caching of blocks. |
| 838 | extra_retained_tokens: int = 0 |
| 839 | |
| 840 | def max_admission_blocks_per_request( |
| 841 | self, max_in_flight_tokens: int, max_model_len: int |
| 842 | ) -> int: |
| 843 | """Per-request admission cap, in blocks. |
| 844 | |
| 845 | Single source of truth for both startup pool sizing |
| 846 | (`max_memory_usage_bytes`) and the runtime admission gate. Per-request |
| 847 | real-held blocks plateau at this bound because |
| 848 | `SlidingWindowManager.remove_skipped_blocks` runs from `allocate_slots` |
| 849 | before each chunk's `get_num_blocks_to_allocate`. |
| 850 | |
| 851 | `max_in_flight_tokens` is the max tokens scheduled but not yet settled |
| 852 | (one batch per concurrent step); see `VllmConfig.max_in_flight_tokens`. |
| 853 | """ |
| 854 | # During chunked prefill, we hold KV for the last `sliding_window-1` |
| 855 | # computed tokens plus the in-flight tokens (frees happen on the |
| 856 | # processed-token basis); never more than `max_model_len`. An additional |
| 857 | # `extra_retained_tokens` trailing tokens are kept alive below the |
| 858 | # window for multi-module spec decoding, and must be accounted here too. |
| 859 | num_tokens = min( |
| 860 | self.sliding_window - 1 + self.extra_retained_tokens + max_in_flight_tokens, |
| 861 | max_model_len, |
| 862 | ) |
| 863 | # +1 because the sliding window may not start from the beginning of |
| 864 | # the block. E.g. block size 4 and num_token 4 needs two blocks |
| 865 | # [XXCD][EF] to store the 6-token window [CDEF]. |
| 866 | return cdiv(num_tokens, self.block_size) + 1 |
| 867 | |
| 868 | def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int: |
| 869 | assert ( |
| 870 | not self.dcp_sharded |
| 871 | or vllm_config.parallel_config.decode_context_parallel_size == 1 |
| 872 | ), "DCP not support sliding window." |
| 873 | max_blocks = self.max_admission_blocks_per_request( |
| 874 | max_in_flight_tokens=vllm_config.max_in_flight_tokens, |
| 875 | max_model_len=vllm_config.model_config.max_model_len, |
| 876 | ) |
| 877 | return max_blocks * self.page_size_bytes |
| 878 | |
| 879 | def is_uniform_with_collection( |
| 880 | self, kv_cache_specs: dict[str, KVCacheSpec] |
| 881 | ) -> bool: |
| 882 | return all( |
| 883 | isinstance(spec, SlidingWindowSpec) |
| 884 | and spec.sliding_window == self.sliding_window |
| 885 | for spec in kv_cache_specs.values() |
| 886 | ) |
| 887 | |
| 888 | |
| 889 | @dataclass(frozen=True, kw_only=True) |
| 890 | class CircularBufferSpec(AttentionSpec): |
| 891 | """One block per request holding the raw keys of the token group that |
| 892 | is still being compressed. |
| 893 | |
| 894 | ``block_size`` is the ring capacity. It must exceed the compression ratio |
| 895 | by the speculative lookahead: a speculative step stores all of its rows, |
| 896 | drafts included, before acceptance is known, while the next step still |
| 897 | reads the open group's committed keys from the ring. |
| 898 | """ |
| 899 | |
| 900 | @property |
| 901 | def block_table_token_alignment(self) -> int | None: |
| 902 | return None |
| 903 | |
| 904 | def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int: |
| 905 | # The ring occupies one block per request for its whole lifetime. |
| 906 | del vllm_config |
| 907 | return self.page_size_bytes |
| 908 | |
| 909 | def max_num_blocks_per_req(self, vllm_config: VllmConfig, max_len: int) -> int: |
| 910 | del vllm_config, max_len |
| 911 | return 1 |
| 912 | |
| 913 | def is_uniform_with_collection( |
| 914 | self, kv_cache_specs: dict[str, KVCacheSpec] |
| 915 | ) -> bool: |
| 916 | return all( |
| 917 | isinstance(spec, CircularBufferSpec) for spec in kv_cache_specs.values() |
| 918 | ) |
| 919 | |
| 920 | @property |
| 921 | def prefix_cacheable(self) -> bool: |
| 922 | return False |
| 923 | |
| 924 | @property |
| 925 | def uses_slot_mapping(self) -> bool: |
| 926 | return False |
| 927 | |
| 928 | |
| 929 | @dataclass(frozen=True, kw_only=True) |
| 930 | class SlidingWindowMLASpec(SlidingWindowSpec): |
| 931 | """Sliding window attention with MLA cache format.""" |
| 932 | |
| 933 | cache_dtype_str: str | None = None |
| 934 | # DeepseekV4-only: see MLAAttentionSpec.model_version. |
| 935 | alignment: int | None = None # Default to None for no padding. |
| 936 | model_version: str | None = None |
| 937 | bounded_replay: bool = False |
| 938 | |
| 939 | # MLA stores a single latent vector per state; there is no separate V. |
| 940 | head_size_v: int = 0 |
| 941 | |
| 942 | def __post_init__(self): |
| 943 | assert self.model_version in (None, "deepseek_v4"), ( |
| 944 | f"Unsupported model version: {self.model_version}" |
| 945 | ) |
| 946 | super().__post_init__() |
| 947 | _apply_alignment_padding(self) |
| 948 | |
| 949 | @property |
| 950 | def prefix_cacheable(self) -> bool: |
| 951 | return not self.bounded_replay |
| 952 | |
| 953 | @property |
| 954 | def prefix_replay_tokens(self) -> int: |
| 955 | return self.sliding_window if self.bounded_replay else 0 |
| 956 | |
| 957 | @classmethod |
| 958 | def merge(cls, specs: list[Self]) -> Self: |
| 959 | assert all(isinstance(spec, SlidingWindowMLASpec) for spec in specs), ( |
| 960 | "All attention layers in the same KV cache group must be " |
| 961 | "SlidingWindowMLASpec." |
| 962 | ) |
| 963 | cache_dtype_str_set = set(spec.cache_dtype_str for spec in specs) |
| 964 | tokens_per_state_set = set(spec.tokens_per_state for spec in specs) |
| 965 | model_version_set = set(spec.model_version for spec in specs) |
| 966 | sliding_window_set = set(spec.sliding_window for spec in specs) |
| 967 | extra_retained_set = set(spec.extra_retained_tokens for spec in specs) |
| 968 | bounded_replay_set = set(spec.bounded_replay for spec in specs) |
| 969 | block_stride_alignment_set = {spec.block_stride_alignment for spec in specs} |
| 970 | assert len({spec.dcp_sharded for spec in specs}) == 1 |
| 971 | assert len({spec.max_tp_shards for spec in specs}) == 1 |
| 972 | assert ( |
| 973 | len(cache_dtype_str_set) == 1 |
| 974 | and len(tokens_per_state_set) == 1 |
| 975 | and len(model_version_set) == 1 |
| 976 | and len(sliding_window_set) == 1 |
| 977 | and len(extra_retained_set) == 1 |
| 978 | and len(bounded_replay_set) == 1 |
| 979 | and len(block_stride_alignment_set) == 1 |
| 980 | ), ( |
| 981 | "All attention layers in the same KV cache group must use the same " |
| 982 | "quantization method, tokens per state, model version, sliding " |
| 983 | "window size, retained token count, and replay policy." |
| 984 | ) |
| 985 | return cls( |
| 986 | block_size=specs[0].block_size, |
| 987 | block_stride_alignment=block_stride_alignment_set.pop(), |
| 988 | num_kv_heads=specs[0].num_kv_heads, |
| 989 | head_size=specs[0].head_size, |
| 990 | dtype=specs[0].dtype, |
| 991 | dcp_sharded=specs[0].dcp_sharded, |
| 992 | max_tp_shards=specs[0].max_tp_shards, |
| 993 | page_size_padded=specs[0].page_size_padded, |
| 994 | num_head_slots=specs[0].num_head_slots, |
| 995 | state_content_bytes=specs[0].state_content_bytes, |
| 996 | sliding_window=sliding_window_set.pop(), |
| 997 | extra_retained_tokens=extra_retained_set.pop(), |
| 998 | cache_dtype_str=cache_dtype_str_set.pop(), |
| 999 | tokens_per_state=tokens_per_state_set.pop(), |
| 1000 | model_version=model_version_set.pop(), |
| 1001 | bounded_replay=bounded_replay_set.pop(), |
| 1002 | ) |
| 1003 | |
| 1004 | def is_uniform_with_collection( |
| 1005 | self, kv_cache_specs: dict[str, KVCacheSpec] |
| 1006 | ) -> bool: |
| 1007 | return all( |
| 1008 | isinstance(spec, SlidingWindowMLASpec) |
| 1009 | and spec.sliding_window == self.sliding_window |
| 1010 | and spec.bounded_replay == self.bounded_replay |
| 1011 | for spec in kv_cache_specs.values() |
| 1012 | ) |
| 1013 | |
| 1014 | |
| 1015 | @dataclass(frozen=True, kw_only=True) |
| 1016 | class KpoolTailSpec(SlidingWindowSpec): |
| 1017 | """One-block circular scratch cache for a kpool indexer's raw tail.""" |
| 1018 | |
| 1019 | def max_admission_blocks_per_request( |
| 1020 | self, max_in_flight_tokens: int, max_model_len: int |
| 1021 | ) -> int: |
| 1022 | return 1 |
| 1023 | |
| 1024 | def max_num_blocks_per_req(self, vllm_config: VllmConfig, max_len: int) -> int: |
| 1025 | return 1 |
| 1026 | |
| 1027 | def is_uniform_with_collection( |
| 1028 | self, kv_cache_specs: dict[str, KVCacheSpec] |
| 1029 | ) -> bool: |
| 1030 | return all(isinstance(spec, KpoolTailSpec) for spec in kv_cache_specs.values()) |
| 1031 | |
| 1032 | @property |
| 1033 | def prefix_cacheable(self) -> bool: |
| 1034 | return False |
| 1035 | |
| 1036 | @property |
| 1037 | def uses_slot_mapping(self) -> bool: |
| 1038 | return False |
| 1039 | |
| 1040 | |
| 1041 | @dataclass(frozen=True) |
| 1042 | class MambaSpec(KVCacheSpec): |
| 1043 | shapes: tuple[tuple[int, ...], ...] |
| 1044 | dtypes: tuple[torch.dtype, ...] |
| 1045 | page_size_padded: int | None = None |
| 1046 | mamba_type: MambaAttentionBackendEnum = MambaAttentionBackendEnum.MAMBA2 |
| 1047 | mamba_cache_mode: str = "none" |
| 1048 | num_speculative_blocks: int = 0 |
| 1049 | num_prefill_checkpoint_blocks: int = 0 |
| 1050 | prefill_checkpoint_alignment: int | None = None |
| 1051 | num_heads: int = 1 |
| 1052 | tokens_per_state: int = -1 |
| 1053 | # False: the state is sharded across TP ranks (e.g. GDN). True: every TP |
| 1054 | # rank holds the full state (e.g. the replicated PLE conv state). |
| 1055 | tp_replicated: bool = False |
| 1056 | |
| 1057 | @property |
| 1058 | def state_content_size_bytes(self) -> int: |
| 1059 | return sum( |
| 1060 | prod(shape) * get_dtype_size(dtype) |
| 1061 | for (shape, dtype) in zip(self.shapes, self.dtypes) |
| 1062 | ) |
| 1063 | |
| 1064 | @property |
| 1065 | def real_page_size_bytes(self) -> int: |
| 1066 | return self.state_content_size_bytes |
| 1067 | |
| 1068 | @property |
| 1069 | def page_size_bytes(self) -> int: |
| 1070 | page_size = sum( |
| 1071 | prod(shape) * get_dtype_size(dtype) |
| 1072 | for (shape, dtype) in zip(self.shapes, self.dtypes) |
| 1073 | ) |
| 1074 | if self.page_size_padded is not None: |
| 1075 | assert self.page_size_padded >= page_size |
| 1076 | return self.page_size_padded |
| 1077 | return page_size |
| 1078 | |
| 1079 | @property |
| 1080 | def block_table_token_alignment(self) -> int | None: |
| 1081 | return None |
| 1082 | |
| 1083 | def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int: |
| 1084 | if vllm_config.cache_config.mamba_cache_mode == "align": |
| 1085 | return self.page_size_bytes * ( |
| 1086 | 2 + self.num_speculative_blocks + self.num_prefill_checkpoint_blocks |
| 1087 | ) |
| 1088 | else: |
| 1089 | return self.page_size_bytes * (1 + self.num_speculative_blocks) |
| 1090 | |
| 1091 | def max_num_blocks_per_req(self, vllm_config: VllmConfig, max_len: int) -> int: |
| 1092 | # Mamba state is replicated across DCP/PCP ranks, never sharded, so |
| 1093 | # no CP scaling applies. |
| 1094 | if vllm_config.cache_config.mamba_cache_mode == "align": |
| 1095 | # Block table rows are position-indexed over the full sequence |
| 1096 | # even though only 2 + num_speculative_blocks state blocks are |
| 1097 | # resident at a time (earlier states are nulled out by |
| 1098 | # remove_skipped_blocks), so the row length must cover max_len |
| 1099 | # rather than max_memory_usage_bytes. |
| 1100 | return cdiv(max_len, self.block_size) + self.num_speculative_blocks |
| 1101 | return cdiv(self.max_memory_usage_bytes(vllm_config), self.page_size_bytes) |
| 1102 | |
| 1103 | def is_uniform_with_collection( |
| 1104 | self, kv_cache_specs: dict[str, KVCacheSpec] |
| 1105 | ) -> bool: |
| 1106 | return all( |
| 1107 | isinstance(spec, MambaSpec) |
| 1108 | and spec.num_speculative_blocks == self.num_speculative_blocks |
| 1109 | and spec.num_prefill_checkpoint_blocks == self.num_prefill_checkpoint_blocks |
| 1110 | and spec.prefill_checkpoint_alignment == self.prefill_checkpoint_alignment |
| 1111 | and spec.page_size_bytes == self.page_size_bytes |
| 1112 | and spec.tp_replicated == self.tp_replicated |
| 1113 | for spec in kv_cache_specs.values() |
| 1114 | ) |
| 1115 | |
| 1116 | |
| 1117 | def get_mamba_prefill_checkpoint_position( |
| 1118 | num_tokens: int, |
| 1119 | hash_block_size: int, |
| 1120 | drop_eagle_block: bool, |
| 1121 | ) -> int: |
| 1122 | """Return the reusable Mamba checkpoint boundary for a prefill.""" |
| 1123 | checkpoint_position = (num_tokens - 1) // hash_block_size * hash_block_size |
| 1124 | if drop_eagle_block: |
| 1125 | checkpoint_position -= hash_block_size |
| 1126 | return max(checkpoint_position, 0) |
| 1127 | |
| 1128 | |
| 1129 | def is_mamba_prefill_checkpoint_valid( |
| 1130 | query_start: int, |
| 1131 | query_end: int, |
| 1132 | checkpoint_position: int, |
| 1133 | hash_block_size: int, |
| 1134 | mamba_block_size: int, |
| 1135 | checkpoint_alignment: int | None, |
| 1136 | ) -> bool: |
| 1137 | """Whether a backend can export the checkpoint in this query.""" |
| 1138 | if checkpoint_alignment is None: |
| 1139 | return False |
| 1140 | assert checkpoint_alignment > 0 |
| 1141 | |
| 1142 | initial_state_col = (query_start - 1) // mamba_block_size |
| 1143 | checkpoint_col = cdiv(query_end, mamba_block_size) - 2 |
| 1144 | return ( |
| 1145 | query_start % hash_block_size == 0 |
| 1146 | and checkpoint_col > initial_state_col |
| 1147 | and query_start + hash_block_size <= checkpoint_position |
| 1148 | and query_start < checkpoint_position < query_end |
| 1149 | and (checkpoint_position - query_start) % checkpoint_alignment == 0 |
| 1150 | ) |
| 1151 | |
| 1152 | |
| 1153 | @dataclass(frozen=True) |
| 1154 | class EncoderOnlyAttentionSpec(AttentionSpec): |
| 1155 | def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int: |
| 1156 | # Encoder-only layers do not need KV cache |
| 1157 | return 0 |
| 1158 | |
| 1159 | |
| 1160 | @dataclass(frozen=True) |
| 1161 | class CrossAttentionSpec(AttentionSpec): |
| 1162 | """KV cache spec for cross-attention layers in encoder-decoder models.""" |
| 1163 | |
| 1164 | def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int: |
| 1165 | # For cross-attention, we need to cache encoder states |
| 1166 | # Get encoder length (e.g., 1500 for Whisper). |
| 1167 | max_encoder_len = vllm_config.scheduler_config.max_num_encoder_input_tokens |
| 1168 | return cdiv(max_encoder_len, self.block_size) * self.page_size_bytes |
| 1169 | |
| 1170 | |
| 1171 | @dataclass(frozen=True) |
| 1172 | class SinkFullAttentionSpec(FullAttentionSpec): |
| 1173 | sink_len: int | None = None |
| 1174 | |
| 1175 | @classmethod |
| 1176 | def merge(cls, specs: list[Self]) -> Self: |
| 1177 | """Merge a list of FullAttentionSpec objects into a single |
| 1178 | FullAttentionSpec object. |
| 1179 | """ |
| 1180 | assert all(isinstance(spec, FullAttentionSpec) for spec in specs), ( |
| 1181 | "All attention layers in the same KV cache group must be FullAttentionSpec." |
| 1182 | ) |
| 1183 | |
| 1184 | sliding_window = set( |
| 1185 | spec.sliding_window for spec in specs if spec.sliding_window is not None |
| 1186 | ) |
| 1187 | attention_chunk_size = set( |
| 1188 | spec.attention_chunk_size |
| 1189 | for spec in specs |
| 1190 | if spec.attention_chunk_size is not None |
| 1191 | ) |
| 1192 | assert not any(isinstance(spec, MLAAttentionSpec) for spec in specs), ( |
| 1193 | "MLAAttentionSpec should be merged in MLAAttentionSpec.merge" |
| 1194 | ) |
| 1195 | merged_spec = cls( |
| 1196 | block_size=specs[0].block_size, |
| 1197 | num_kv_heads=specs[0].num_kv_heads, |
| 1198 | head_size=specs[0].head_size, |
| 1199 | head_size_v=specs[0].head_size_v, |
| 1200 | sink_len=specs[0].sink_len, |
| 1201 | block_stride_alignment=specs[0].block_stride_alignment, |
| 1202 | dtype=specs[0].dtype, |
| 1203 | kv_quant_mode=specs[0].kv_quant_mode, |
| 1204 | dcp_sharded=specs[0].dcp_sharded, |
| 1205 | max_tp_shards=specs[0].max_tp_shards, |
| 1206 | page_size_padded=specs[0].page_size_padded, |
| 1207 | num_head_slots=specs[0].num_head_slots, |
| 1208 | state_content_bytes=specs[0].state_content_bytes, |
| 1209 | sliding_window=cls.merge_window_sizes(sliding_window), |
| 1210 | attention_chunk_size=cls.merge_window_sizes(attention_chunk_size), |
| 1211 | non_causal=any(spec.non_causal for spec in specs), |
| 1212 | ) |
| 1213 | for spec in specs: |
| 1214 | for f in fields(AttentionSpec): |
| 1215 | assert getattr(spec, f.name) == getattr(merged_spec, f.name), ( |
| 1216 | "All attention layers in the same KV cache group must have " |
| 1217 | "the same attention spec." |
| 1218 | ) |
| 1219 | assert (merged_spec.sliding_window is not None) + ( |
| 1220 | merged_spec.attention_chunk_size is not None |
| 1221 | ) <= 1, ( |
| 1222 | "Model with both sliding window layers and chunked local attention " |
| 1223 | "layers is not supported." |
| 1224 | ) |
| 1225 | return merged_spec |
| 1226 | |
| 1227 | |
| 1228 | @dataclass(frozen=True) |
| 1229 | class UniformTypeKVCacheSpecs(KVCacheSpec): |
| 1230 | """A KV cache spec for multiple layers with the same type of attention. Here, |
| 1231 | same types means always need the same number of token slots. For example, |
| 1232 | sliding window attentions with different window sizes are not the same type |
| 1233 | and should not be merged into one UniformTypeKVCacheSpecs. |
| 1234 | """ |
| 1235 | |
| 1236 | kv_cache_specs: dict[str, KVCacheSpec] |
| 1237 | |
| 1238 | def __post_init__(self): |
| 1239 | super().__post_init__() |
| 1240 | if self.kv_cache_specs: |
| 1241 | object.__setattr__(self, "dcp_sharded", self.first_spec.dcp_sharded) |
| 1242 | |
| 1243 | @property |
| 1244 | def prefix_cacheable(self) -> bool: |
| 1245 | return all(spec.prefix_cacheable for spec in self.kv_cache_specs.values()) |
| 1246 | |
| 1247 | @property |
| 1248 | def prefix_replay_tokens(self) -> int: |
| 1249 | return max(spec.prefix_replay_tokens for spec in self.kv_cache_specs.values()) |
| 1250 | |
| 1251 | @property |
| 1252 | def first_spec(self) -> KVCacheSpec: |
| 1253 | """Return the first spec in the group.""" |
| 1254 | return next(iter(self.kv_cache_specs.values())) |
| 1255 | |
| 1256 | @property |
| 1257 | def page_size_bytes(self) -> int: |
| 1258 | return sum(spec.page_size_bytes for spec in self.kv_cache_specs.values()) |
| 1259 | |
| 1260 | @property |
| 1261 | def block_table_token_alignment(self) -> int | None: |
| 1262 | return self.first_spec.block_table_token_alignment |
| 1263 | |
| 1264 | def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int: |
| 1265 | max_num_pages = max( |
| 1266 | cdiv(spec.max_memory_usage_bytes(vllm_config), spec.page_size_bytes) |
| 1267 | for spec in self.kv_cache_specs.values() |
| 1268 | ) |
| 1269 | return max_num_pages * self.page_size_bytes |
| 1270 | |
| 1271 | def max_num_blocks_per_req(self, vllm_config: VllmConfig, max_len: int) -> int: |
| 1272 | # Metadata builders are constructed from the per-layer spec, so the base |
| 1273 | # cdiv(max_len, block_size) would drop its DCP sharding and size the |
| 1274 | # block table wider than those builders expect. |
| 1275 | widths = { |
| 1276 | spec.max_num_blocks_per_req(vllm_config, max_len) |
| 1277 | for spec in self.kv_cache_specs.values() |
| 1278 | } |
| 1279 | assert len(widths) == 1, ( |
| 1280 | "All layers in the same KV cache group must need the same number " |
| 1281 | f"of block table entries, got {sorted(widths)}." |
| 1282 | ) |
| 1283 | return next(iter(widths)) |
| 1284 | |
| 1285 | @classmethod |
| 1286 | def is_uniform_type(cls, kv_cache_specs: dict[str, KVCacheSpec]) -> bool: |
| 1287 | """Whether all layers have the same type of KV cache spec. |
| 1288 | |
| 1289 | Uses the registry to determine grouping base classes, so custom specs |
| 1290 | that inherit from FullAttentionSpec are treated as full attention. |
| 1291 | """ |
| 1292 | block_sizes = set(spec.block_size for spec in kv_cache_specs.values()) |
| 1293 | if len(block_sizes) > 1: |
| 1294 | # Different block sizes, not uniform. |
| 1295 | return False |
| 1296 | if len({spec.dcp_sharded for spec in kv_cache_specs.values()}) > 1: |
| 1297 | return False |
| 1298 | first_spec = next(iter(kv_cache_specs.values())) |
| 1299 | return first_spec.is_uniform_with_collection(kv_cache_specs) |
| 1300 | |
| 1301 | @classmethod |
| 1302 | def from_specs(cls, kv_cache_specs: dict[str, KVCacheSpec]) -> Self | None: |
| 1303 | """Return a SameTypeKVCacheSpecs object if all layers have the same type |
| 1304 | of KV cache spec. Return None if not. |
| 1305 | """ |
| 1306 | if cls.is_uniform_type(kv_cache_specs): |
| 1307 | block_size = next(iter(kv_cache_specs.values())).block_size |
| 1308 | return cls(block_size=block_size, kv_cache_specs=kv_cache_specs) |
| 1309 | else: |
| 1310 | return None |
| 1311 | |
| 1312 | def get_max_layers_per_page_size(self) -> int: |
| 1313 | """Max number of layers sharing a page size. For a balanced bucket |
| 1314 | this equals the number of repetitions of the layer pattern.""" |
| 1315 | return Counter( |
| 1316 | spec.page_size_bytes for spec in self.kv_cache_specs.values() |
| 1317 | ).most_common(1)[0][1] |
| 1318 | |
| 1319 | def max_memory_usage_pages(self, vllm_config: VllmConfig) -> int: |
| 1320 | return max( |
| 1321 | cdiv(spec.max_memory_usage_bytes(vllm_config), spec.page_size_bytes) |
| 1322 | for spec in self.kv_cache_specs.values() |
| 1323 | ) |
| 1324 | |
| 1325 | @property |
| 1326 | def uses_slot_mapping(self) -> bool: |
| 1327 | return self.first_spec.uses_slot_mapping |
| 1328 | |
| 1329 | |
| 1330 | def iter_layer_specs(kv_cache_spec: KVCacheSpec) -> Collection[KVCacheSpec]: |
| 1331 | """The per-layer specs a KV cache group spec covers. |
| 1332 | |
| 1333 | ``UniformTypeKVCacheSpecs`` groups keep one spec per layer; every other |
| 1334 | spec describes its group on its own. Returns the layer specs either way so |
| 1335 | callers do not have to special-case the wrapper. |
| 1336 | """ |
| 1337 | if isinstance(kv_cache_spec, UniformTypeKVCacheSpecs): |
| 1338 | return kv_cache_spec.kv_cache_specs.values() |
| 1339 | return (kv_cache_spec,) |
| 1340 | |
| 1341 | |
| 1342 | def is_full_attention_spec(kv_cache_spec: KVCacheSpec) -> bool: |
| 1343 | """Whether a KV cache group spec is (or wraps) full attention. |
| 1344 | |
| 1345 | ``UniformTypeKVCacheSpecs`` is not itself a ``FullAttentionSpec``, so a bare |
| 1346 | isinstance check misses groups that carry the wrapper -- DeepSeek-V4's MLA |
| 1347 | layers, or any model taking the ``UniformTypeKVCacheSpecs.from_specs`` path. |
| 1348 | |
| 1349 | Every layer must be full attention: a group holding a recycling |
| 1350 | (sliding-window) layer has no stable slot layout, so callers that key data |
| 1351 | by slot cannot use it. |
| 1352 | """ |
| 1353 | layer_specs = iter_layer_specs(kv_cache_spec) |
| 1354 | return len(layer_specs) > 0 and all( |
| 1355 | isinstance(spec, FullAttentionSpec) for spec in layer_specs |
| 1356 | ) |
| 1357 | |
| 1358 | |
| 1359 | def get_kv_cache_spec_kind(kv_cache_spec: KVCacheSpec) -> KVCacheSpecKind: |
| 1360 | if isinstance(kv_cache_spec, UniformTypeKVCacheSpecs): |
| 1361 | inner_kinds = { |
| 1362 | get_kv_cache_spec_kind(spec) |
| 1363 | for spec in kv_cache_spec.kv_cache_specs.values() |
| 1364 | } |
| 1365 | if len(inner_kinds) == 1: |
| 1366 | return next(iter(inner_kinds)) |
| 1367 | # A group is only formed when all members share one registered |
| 1368 | # uniform_type_base_spec, so UNKNOWN would discard what the merge |
| 1369 | # already established. |
| 1370 | base_specs = { |
| 1371 | KVCacheSpecRegistry.get_uniform_type_base_spec(spec) |
| 1372 | for spec in kv_cache_spec.kv_cache_specs.values() |
| 1373 | } |
| 1374 | if len(base_specs) == 1 and next(iter(base_specs)) is FullAttentionSpec: |
| 1375 | return KVCacheSpecKind.FULL_ATTENTION |
| 1376 | return KVCacheSpecKind.UNKNOWN |
| 1377 | # Keep subclass checks before base classes so specialized specs keep their |
| 1378 | # more precise kind. |
| 1379 | if isinstance(kv_cache_spec, SlidingWindowMLASpec): |
| 1380 | return KVCacheSpecKind.SLIDING_WINDOW_MLA |
| 1381 | if isinstance(kv_cache_spec, MLAAttentionSpec): |
| 1382 | return KVCacheSpecKind.MLA_ATTENTION |
| 1383 | if isinstance(kv_cache_spec, SinkFullAttentionSpec): |
| 1384 | return KVCacheSpecKind.SINK_FULL_ATTENTION |
| 1385 | if isinstance(kv_cache_spec, FullAttentionSpec): |
| 1386 | return KVCacheSpecKind.FULL_ATTENTION |
| 1387 | if isinstance(kv_cache_spec, ChunkedLocalAttentionSpec): |
| 1388 | return KVCacheSpecKind.CHUNKED_LOCAL_ATTENTION |
| 1389 | if isinstance(kv_cache_spec, SlidingWindowSpec): |
| 1390 | return KVCacheSpecKind.SLIDING_WINDOW |
| 1391 | if isinstance(kv_cache_spec, MambaSpec): |
| 1392 | return KVCacheSpecKind.MAMBA |
| 1393 | if isinstance(kv_cache_spec, EncoderOnlyAttentionSpec): |
| 1394 | return KVCacheSpecKind.ENCODER_ONLY_ATTENTION |
| 1395 | if isinstance(kv_cache_spec, CrossAttentionSpec): |
| 1396 | return KVCacheSpecKind.CROSS_ATTENTION |
| 1397 | return KVCacheSpecKind.UNKNOWN |
| 1398 | |
| 1399 | |
| 1400 | def get_kv_cache_spec_sliding_window(kv_cache_spec: KVCacheSpec) -> int | None: |
| 1401 | if isinstance(kv_cache_spec, UniformTypeKVCacheSpecs): |
| 1402 | inner_windows = { |
| 1403 | get_kv_cache_spec_sliding_window(spec) |
| 1404 | for spec in kv_cache_spec.kv_cache_specs.values() |
| 1405 | } |
| 1406 | return next(iter(inner_windows)) if len(inner_windows) == 1 else None |
| 1407 | if isinstance(kv_cache_spec, SlidingWindowSpec): |
| 1408 | return kv_cache_spec.sliding_window |
| 1409 | return None |
| 1410 | |
| 1411 | |
| 1412 | @dataclass |
| 1413 | class KVCacheTensor: |
| 1414 | """A class for specifying how the workers should initialize the KV cache. |
| 1415 | |
| 1416 | Placement of a set of same-shaped layers in the KV cache allocation. |
| 1417 | Layer ``layers[l]``'s page for block ``b`` starts at |
| 1418 | ``offset + l * layer_stride + b * block_stride`` bytes into the backing |
| 1419 | allocation of ``size`` bytes. Layer-outermost layouts give each layer a |
| 1420 | contiguous region (``layer_stride = page * num_blocks``, |
| 1421 | ``block_stride = page``); block-outermost layouts make each block a |
| 1422 | block of all layers' pages (``layer_stride = page``, ``block_stride`` = |
| 1423 | the packed block). Tensors whose address ranges overlap |
| 1424 | alias the same bytes: cache groups overlay each other, which is sound |
| 1425 | because a block ID is owned by one group at a time. |
| 1426 | """ |
| 1427 | |
| 1428 | size: int # total size of the backing allocation in bytes |
| 1429 | layers: list[str] # layer names in L order |
| 1430 | layer_stride: int |
| 1431 | block_stride: int |
| 1432 | offset: int = 0 # byte offset of layers[0]'s block 0 |
| 1433 | host_resident: bool = False |
| 1434 | |
| 1435 | |
| 1436 | class KVCacheGroupRole(str, Enum): |
| 1437 | DEFAULT = "default" |
| 1438 | HISPARSE_SOURCE = "hisparse_source" |
| 1439 | HISPARSE_INDEXER = "hisparse_indexer" |
| 1440 | |
| 1441 | |
| 1442 | @dataclass |
| 1443 | class KVCacheGroupSpec: |
| 1444 | """Represents a group of model layers that share the same KV cache block table. |
| 1445 | These layers are regarded as one layer in the KV cache manager. |
| 1446 | """ |
| 1447 | |
| 1448 | # The names of model layers in this group |
| 1449 | layer_names: list[str] |
| 1450 | # The KV cache spec of this manager layer |
| 1451 | kv_cache_spec: KVCacheSpec |
| 1452 | # Whether this group contains EAGLE/MTP draft attention layers. |
| 1453 | is_eagle_group: bool = False |
| 1454 | # Host groups use the dedicated HiSparse pool; others share the device pool. |
| 1455 | host_resident: bool = False |
| 1456 | # Whether this group is part of the externally transferable KV state. |
| 1457 | # Ephemeral accelerator-side replicas are rebuilt from their transferable |
| 1458 | # source group after a connector load. |
| 1459 | enable_kv_transfer: bool = True |
| 1460 | role: KVCacheGroupRole = KVCacheGroupRole.DEFAULT |
| 1461 | |
| 1462 | |
| 1463 | @dataclass |
| 1464 | class KVCacheConfig: |
| 1465 | """The KV cache configuration of a model.""" |
| 1466 | |
| 1467 | num_blocks: int |
| 1468 | """The number of KV cache blocks""" |
| 1469 | kv_cache_tensors: list[KVCacheTensor] |
| 1470 | """How should model runner initialize the KV cache tensors for each layer""" |
| 1471 | kv_cache_groups: list[KVCacheGroupSpec] |
| 1472 | """ |
| 1473 | The kv cache groups of the model. |
| 1474 | For models with only one type of attention, there is only one group that |
| 1475 | contains all layers. |
| 1476 | For models with multiple types of attention, there will be multiple groups, |
| 1477 | see `_get_kv_cache_config_uniform_page_size` for more details. |
| 1478 | """ |
| 1479 | prefix_cache_retention_interval: int | None = None |
| 1480 | """Resolved retention policy for local prefix-cache checkpoints.""" |
| 1481 | kv_cache_layout: str | None = None |
| 1482 | """The KV cache layout resolved by the engine core, adopted by all workers.""" |
| 1483 | hisparse_host_num_blocks: int | None = None |
| 1484 | """Capacity of the dedicated HiSparse host-block manager, when enabled.""" |
| 1485 | |
| 1486 | hisparse_host_block_stride: int | None = None |
| 1487 | """Physical bytes between consecutive HiSparse host blocks.""" |
| 1488 | |
| 1489 | hisparse_shared_host_pool: bool = False |
| 1490 | """Whether local TP ranks share one physical HiSparse host pool.""" |
| 1491 | |
| 1492 | kv_tp_replicas: int = 1 |
| 1493 | """Consecutive TP ranks holding identical KV for every layer (1: none).""" |
| 1494 | |
| 1495 | @cached_property |
| 1496 | def transfer_group_ids(self) -> tuple[int, ...]: |
| 1497 | """IDs of cache groups that participate in external KV transfer.""" |
| 1498 | return tuple( |
| 1499 | group_id |
| 1500 | for group_id, group in enumerate(self.kv_cache_groups) |
| 1501 | if group.enable_kv_transfer |
| 1502 | ) |
| 1503 | |
| 1504 | @cached_property |
| 1505 | def prefix_cacheable_group_ids(self) -> tuple[int, ...]: |
| 1506 | """IDs of transferable groups eligible for hash-addressed stores.""" |
| 1507 | return tuple( |
| 1508 | group_id |
| 1509 | for group_id in self.transfer_group_ids |
| 1510 | if self.kv_cache_groups[group_id].kv_cache_spec.prefix_cacheable |
| 1511 | ) |
| 1512 | |
| 1513 | @cached_property |
| 1514 | def prefix_cacheable_groups(self) -> tuple[KVCacheGroupSpec, ...]: |
| 1515 | """Transferable groups eligible for hash-addressed stores.""" |
| 1516 | return tuple( |
| 1517 | self.kv_cache_groups[group_id] |
| 1518 | for group_id in self.prefix_cacheable_group_ids |
| 1519 | ) |
| 1520 | |
| 1521 | @cached_property |
| 1522 | def transfer_groups(self) -> tuple[KVCacheGroupSpec, ...]: |
| 1523 | """Cache groups that participate in external KV transfer.""" |
| 1524 | return tuple( |
| 1525 | self.kv_cache_groups[group_id] for group_id in self.transfer_group_ids |
| 1526 | ) |
| 1527 | |
| 1528 | @cached_property |
| 1529 | def transfer_group_index_by_layer(self) -> dict[str, int]: |
| 1530 | """Transfer-group tuple index for each participating layer.""" |
| 1531 | return { |
| 1532 | layer_name: group_index |
| 1533 | for group_index, group in enumerate(self.transfer_groups) |
| 1534 | for layer_name in group.layer_names |
| 1535 | } |
| 1536 | |
| 1537 | def select_transfer_block_ids( |
| 1538 | self, block_ids: Sequence[list[int]] |
| 1539 | ) -> tuple[list[int], ...]: |
| 1540 | """Select block IDs for externally transferable cache groups.""" |
| 1541 | if len(block_ids) != len(self.kv_cache_groups): |
| 1542 | raise ValueError( |
| 1543 | f"Expected {len(self.kv_cache_groups)} KV cache groups, " |
| 1544 | f"got {len(block_ids)}." |
| 1545 | ) |
| 1546 | return tuple(block_ids[group_id] for group_id in self.transfer_group_ids) |
| 1547 | |
| 1548 | @property |
| 1549 | def host_group_ids(self) -> tuple[int, ...]: |
| 1550 | return tuple( |
| 1551 | group_id |
| 1552 | for group_id, group in enumerate(self.kv_cache_groups) |
| 1553 | if group.host_resident |
| 1554 | ) |
| 1555 | |
| 1556 | def num_blocks_of(self, tensor: KVCacheTensor) -> int: |
| 1557 | """Number of blocks addressable by the pool backing ``tensor``.""" |
| 1558 | if not tensor.host_resident: |
| 1559 | return self.num_blocks |
| 1560 | assert self.hisparse_host_num_blocks is not None |
| 1561 | return self.hisparse_host_num_blocks |
| 1562 | |
| 1563 | @property |
| 1564 | def has_mamba_layers(self) -> bool: |
| 1565 | return any( |
| 1566 | isinstance(spec, MambaSpec) |
| 1567 | for group in self.kv_cache_groups |
| 1568 | for spec in iter_layer_specs(group.kv_cache_spec) |
| 1569 | ) |
| 1570 | |
| 1571 | @property |
| 1572 | def has_mixed_precision_kv_cache(self) -> bool: |
| 1573 | """Whether device attention caches use more than one precision.""" |
| 1574 | kv_cache_precisions: set[tuple[torch.dtype, KVQuantMode]] = set() |
| 1575 | for group in self.kv_cache_groups: |
| 1576 | if group.host_resident: |
| 1577 | continue |
| 1578 | kv_cache_precisions.update( |
| 1579 | (spec.dtype, spec.kv_quant_mode) |
| 1580 | for spec in iter_layer_specs(group.kv_cache_spec) |
| 1581 | if isinstance(spec, AttentionSpec) |
| 1582 | ) |
| 1583 | return len(kv_cache_precisions) > 1 |
| 1584 | |
| 1585 | @property |
| 1586 | def needs_kv_cache_zeroing(self) -> bool: |
| 1587 | """Whether newly allocated KV cache blocks must be zeroed before use. |
| 1588 | |
| 1589 | Required for Mamba layers, whose state is read before it is fully written |
| 1590 | (#35219), and for mixed-precision caches, where a block reused across |
| 1591 | groups can be reinterpreted under a different precision and decode stale |
| 1592 | bytes to NaN/Inf. Uniform-precision caches skip zeroing. |
| 1593 | """ |
| 1594 | return self.has_mamba_layers or self.has_mixed_precision_kv_cache |