wehub-resource-sync
982 lines
37 KiB
Python
982 lines
37 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||
|
||
from __future__ import annotations
|
||
|
||
import copy
|
||
from collections import Counter
|
||
from dataclasses import dataclass, fields, replace
|
||
from enum import Enum, IntEnum
|
||
from math import prod
|
||
from typing import TYPE_CHECKING
|
||
|
||
import torch
|
||
from typing_extensions import Self
|
||
|
||
from vllm.logger import init_logger
|
||
from vllm.utils.math_utils import cdiv, round_up
|
||
from vllm.utils.torch_utils import get_dtype_size, nvfp4_kv_cache_full_dim
|
||
from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
|
||
from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry
|
||
|
||
if TYPE_CHECKING:
|
||
from vllm.config import VllmConfig
|
||
|
||
logger = init_logger(__name__)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# KV cache quantization mode
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
class KVQuantMode(IntEnum):
|
||
"""KV cache quantization mode.
|
||
|
||
Used by attention backends and kernels to dispatch quantization logic
|
||
without string matching on ``kv_cache_dtype``.
|
||
"""
|
||
|
||
NONE = 0
|
||
FP8_PER_TENSOR = 1 # per-tensor scales (current fp8 path)
|
||
INT8_PER_TOKEN_HEAD = 2 # per-token-head dynamic scales for int8
|
||
FP8_PER_TOKEN_HEAD = 3 # per-token-head dynamic scales for fp8
|
||
INT4_PER_TOKEN_HEAD = 4 # packed 2×int4/byte, RHT + asymmetric zp
|
||
NVFP4 = 5 # packed fp4 data + fp8 block scales
|
||
|
||
@property
|
||
def is_per_token_head(self) -> bool:
|
||
"""True for any per-token-head quantization mode."""
|
||
return self in (
|
||
KVQuantMode.INT8_PER_TOKEN_HEAD,
|
||
KVQuantMode.FP8_PER_TOKEN_HEAD,
|
||
KVQuantMode.INT4_PER_TOKEN_HEAD,
|
||
)
|
||
|
||
@property
|
||
def is_nvfp4(self) -> bool:
|
||
"""True for NVFP4 packed quantization mode."""
|
||
return self == KVQuantMode.NVFP4
|
||
|
||
|
||
def get_kv_quant_mode(kv_cache_dtype: str) -> KVQuantMode:
|
||
"""Map a ``kv_cache_dtype`` string to a :class:`KVQuantMode`."""
|
||
if kv_cache_dtype == "int4_per_token_head":
|
||
return KVQuantMode.INT4_PER_TOKEN_HEAD
|
||
if kv_cache_dtype == "int8_per_token_head":
|
||
return KVQuantMode.INT8_PER_TOKEN_HEAD
|
||
if kv_cache_dtype == "fp8_per_token_head":
|
||
return KVQuantMode.FP8_PER_TOKEN_HEAD
|
||
if kv_cache_dtype == "nvfp4":
|
||
return KVQuantMode.NVFP4
|
||
if isinstance(kv_cache_dtype, str) and kv_cache_dtype.startswith("fp8"):
|
||
return KVQuantMode.FP8_PER_TENSOR
|
||
return KVQuantMode.NONE
|
||
|
||
|
||
def is_quantized_kv_cache(kv_cache_dtype: str) -> bool:
|
||
return get_kv_quant_mode(kv_cache_dtype) != KVQuantMode.NONE
|
||
|
||
|
||
def kv_cache_uses_per_token_head_scales(kv_cache_dtype: str) -> bool:
|
||
"""Return True if *kv_cache_dtype* needs per-token-head scales."""
|
||
return get_kv_quant_mode(kv_cache_dtype).is_per_token_head
|
||
|
||
|
||
class KVCacheSpecKind(str, Enum):
|
||
FULL_ATTENTION = "full_attention"
|
||
MLA_ATTENTION = "mla_attention"
|
||
SLIDING_WINDOW = "sliding_window"
|
||
SLIDING_WINDOW_MLA = "sliding_window_mla"
|
||
MAMBA = "mamba"
|
||
CHUNKED_LOCAL_ATTENTION = "chunked_local_attention"
|
||
SINK_FULL_ATTENTION = "sink_full_attention"
|
||
ENCODER_ONLY_ATTENTION = "encoder_only_attention"
|
||
CROSS_ATTENTION = "cross_attention"
|
||
UNKNOWN = "unknown"
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class KVCacheSpec:
|
||
"""
|
||
A base class for specifying the KV cache format of one layer.
|
||
"""
|
||
|
||
# number of tokens in a block
|
||
block_size: int
|
||
|
||
@property
|
||
def page_size_bytes(self) -> int:
|
||
"""
|
||
The size of a page with `block_size` tokens in bytes.
|
||
|
||
Returns:
|
||
The page size
|
||
"""
|
||
raise NotImplementedError
|
||
|
||
@property
|
||
def storage_block_size(self) -> int:
|
||
return self.block_size
|
||
|
||
def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:
|
||
"""
|
||
The maximum possible memory usage of this KV cache in bytes.
|
||
|
||
Returns:
|
||
The KV cache size in bytes
|
||
"""
|
||
raise NotImplementedError
|
||
|
||
def max_num_blocks_per_req(self, vllm_config: VllmConfig, max_len: int) -> int:
|
||
"""
|
||
The number of block table entries needed per request, i.e. the row
|
||
length of the worker-side block table for this cache group.
|
||
|
||
Args:
|
||
vllm_config: The vllm config.
|
||
max_len: The maximum sequence length to size for, including the
|
||
encoder length for encoder-decoder models.
|
||
"""
|
||
return cdiv(max_len, self.block_size)
|
||
|
||
def copy_with_new_block_size(self, block_size: int) -> Self:
|
||
"""
|
||
Create a new KVCacheSpec from self but replacing the block size.
|
||
"""
|
||
return replace(self, block_size=block_size)
|
||
|
||
@classmethod
|
||
def merge(cls, specs: list[Self]) -> Self:
|
||
"""
|
||
Merge a list of KVCacheSpec objects into a single KVCacheSpec object.
|
||
"""
|
||
assert all(spec == specs[0] for spec in specs[1:]), (
|
||
"All layers in the same KV cache group must be the same."
|
||
)
|
||
return copy.deepcopy(specs[0])
|
||
|
||
def is_uniform_with_collection(
|
||
self, kv_cache_specs: dict[str, KVCacheSpec]
|
||
) -> bool:
|
||
"""
|
||
Whether this KVCacheSpec is uniform with all specs of all layers.
|
||
"""
|
||
uniform_type_base_spec = KVCacheSpecRegistry.get_uniform_type_base_spec(self)
|
||
assert uniform_type_base_spec is not None, (
|
||
f"Unsupported KV cache spec type: {type(self)}. "
|
||
"Please register it using @register_kv_cache_spec decorator."
|
||
)
|
||
return all(
|
||
isinstance(spec, uniform_type_base_spec) for spec in kv_cache_specs.values()
|
||
)
|
||
|
||
|
||
@dataclass(frozen=True, kw_only=True)
|
||
class AttentionSpec(KVCacheSpec):
|
||
num_kv_heads: int
|
||
head_size: int
|
||
dtype: torch.dtype
|
||
kv_quant_mode: KVQuantMode = KVQuantMode.NONE
|
||
page_size_padded: int | None = None
|
||
indexes_kv_by_block_stride: bool = False
|
||
|
||
@property
|
||
def page_size_bytes(self) -> int:
|
||
real_page_size = self.real_page_size_bytes
|
||
# Per-token-head scales are stored in separate tensors managed
|
||
# by the attention backend, but the memory is carved from the
|
||
# raw KV cache allocation so it must be budgeted here.
|
||
if self.kv_quant_mode.is_per_token_head:
|
||
real_page_size += (
|
||
2 * self.block_size * self.num_kv_heads * get_dtype_size(torch.float32)
|
||
)
|
||
if self.page_size_padded is not None:
|
||
assert self.page_size_padded >= real_page_size
|
||
return self.page_size_padded
|
||
return real_page_size
|
||
|
||
@property
|
||
def real_page_size_bytes(self) -> int:
|
||
if self.kv_quant_mode.is_nvfp4:
|
||
# Packed layout: fp4 data + fp8 block scales per head.
|
||
head_dim = nvfp4_kv_cache_full_dim(self.head_size)
|
||
elif self.kv_quant_mode == KVQuantMode.INT4_PER_TOKEN_HEAD:
|
||
head_dim = self.head_size // 2
|
||
else:
|
||
head_dim = self.head_size
|
||
return (
|
||
2
|
||
* self.block_size
|
||
* self.num_kv_heads
|
||
* head_dim
|
||
* get_dtype_size(self.dtype)
|
||
)
|
||
|
||
def max_num_blocks_per_req(self, vllm_config: VllmConfig, max_len: int) -> int:
|
||
# Attention KV is token-interleaved across DCP/PCP ranks, so each rank
|
||
# only stores max_len // (dcp * pcp) tokens per request.
|
||
parallel_config = vllm_config.parallel_config
|
||
total_cp_size = (
|
||
parallel_config.decode_context_parallel_size
|
||
* parallel_config.prefill_context_parallel_size
|
||
)
|
||
return cdiv(max_len, self.block_size * total_cp_size)
|
||
|
||
|
||
@dataclass(frozen=True, kw_only=True)
|
||
class FullAttentionSpec(AttentionSpec):
|
||
"""
|
||
When hybrid allocator is disabled and the model contains both full
|
||
attention layers and sliding window attention layers, sliding
|
||
window attention are regarded as full attention in KV cache manager
|
||
(blocks are allocated for all tokens), while computed as sliding window
|
||
attention in model runner.
|
||
In this case, we use FullAttentionSpec and record the sliding window size.
|
||
"""
|
||
|
||
head_size_v: int = None # type: ignore[assignment]
|
||
|
||
sliding_window: int | None = None
|
||
"""
|
||
Default to None for not using sliding window attention.
|
||
"""
|
||
attention_chunk_size: int | None = None
|
||
|
||
non_causal: bool = False
|
||
"""
|
||
Whether the layer attends non-causally (e.g. Prefix LM). Carried on the
|
||
spec so the engine core, which collects specs from all workers before the
|
||
scheduler is built, can adjust scheduling policy (chunked prefill / prefix
|
||
caching) regardless of tensor-parallel layout. It does not affect the KV
|
||
cache layout itself.
|
||
"""
|
||
|
||
def __post_init__(self):
|
||
if self.head_size_v is None:
|
||
object.__setattr__(self, "head_size_v", self.head_size)
|
||
|
||
def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:
|
||
max_model_len = vllm_config.model_config.max_model_len
|
||
dcp_world_size = vllm_config.parallel_config.decode_context_parallel_size
|
||
pcp_world_size = vllm_config.parallel_config.prefill_context_parallel_size
|
||
# Note(hc): each dcp rank only need save
|
||
# (max_model_len//dcp_world_size) tokens locally.
|
||
if dcp_world_size * pcp_world_size > 1:
|
||
max_model_len = cdiv(max_model_len, dcp_world_size * pcp_world_size)
|
||
return cdiv(max_model_len, self.block_size) * self.page_size_bytes
|
||
|
||
@classmethod
|
||
def merge_window_sizes(cls, window_sizes: set[int]) -> int | None:
|
||
if len(window_sizes) == 0:
|
||
return None
|
||
elif len(window_sizes) == 1:
|
||
return window_sizes.pop()
|
||
else:
|
||
raise ValueError(
|
||
"All attention layers in the same KV cache group must have the "
|
||
"same window size."
|
||
)
|
||
|
||
@classmethod
|
||
def merge(cls, specs: list[Self]) -> Self:
|
||
"""
|
||
Merge a list of FullAttentionSpec objects into a single
|
||
FullAttentionSpec object.
|
||
"""
|
||
assert all(isinstance(spec, FullAttentionSpec) for spec in specs), (
|
||
"All attention layers in the same KV cache group must be FullAttentionSpec."
|
||
)
|
||
|
||
sliding_window = set(
|
||
spec.sliding_window for spec in specs if spec.sliding_window is not None
|
||
)
|
||
attention_chunk_size = set(
|
||
spec.attention_chunk_size
|
||
for spec in specs
|
||
if spec.attention_chunk_size is not None
|
||
)
|
||
assert not any(isinstance(spec, MLAAttentionSpec) for spec in specs), (
|
||
"MLAAttentionSpec should be merged in MLAAttentionSpec.merge"
|
||
)
|
||
merged_spec = cls(
|
||
block_size=specs[0].block_size,
|
||
num_kv_heads=specs[0].num_kv_heads,
|
||
head_size=specs[0].head_size,
|
||
head_size_v=specs[0].head_size_v,
|
||
dtype=specs[0].dtype,
|
||
kv_quant_mode=specs[0].kv_quant_mode,
|
||
page_size_padded=specs[0].page_size_padded,
|
||
indexes_kv_by_block_stride=specs[0].indexes_kv_by_block_stride,
|
||
sliding_window=cls.merge_window_sizes(sliding_window),
|
||
attention_chunk_size=cls.merge_window_sizes(attention_chunk_size),
|
||
# If any layer in the group is non-causal, treat the group as
|
||
# non-causal so the engine core disables incompatible scheduling.
|
||
non_causal=any(spec.non_causal for spec in specs),
|
||
)
|
||
for spec in specs:
|
||
for f in fields(AttentionSpec):
|
||
assert getattr(spec, f.name) == getattr(merged_spec, f.name), (
|
||
"All attention layers in the same KV cache group must have "
|
||
"the same attention spec."
|
||
)
|
||
assert (merged_spec.sliding_window is not None) + (
|
||
merged_spec.attention_chunk_size is not None
|
||
) <= 1, (
|
||
"Model with both sliding window layers and chunked local attention "
|
||
"layers is not supported."
|
||
)
|
||
return merged_spec
|
||
|
||
@property
|
||
def real_page_size_bytes(self) -> int:
|
||
if self.kv_quant_mode.is_nvfp4:
|
||
# Packed layout per head: fp4 data + fp8 block scales.
|
||
# fp4 data: head_size//2 bytes (2 fp4 values per byte)
|
||
# fp8 block scale: head_size//16 bytes (1 scale per 16 elements)
|
||
last_dim = nvfp4_kv_cache_full_dim(
|
||
self.head_size
|
||
) + nvfp4_kv_cache_full_dim(self.head_size_v)
|
||
elif self.kv_quant_mode == KVQuantMode.INT4_PER_TOKEN_HEAD:
|
||
last_dim = self.head_size // 2 + self.head_size_v // 2
|
||
else:
|
||
last_dim = self.head_size + self.head_size_v
|
||
return (
|
||
self.block_size * self.num_kv_heads * last_dim * get_dtype_size(self.dtype)
|
||
)
|
||
|
||
|
||
def _apply_alignment_padding(spec: MLAAttentionSpec | SlidingWindowMLASpec):
|
||
if spec.alignment is None:
|
||
return
|
||
actual_page_size = spec.real_page_size_bytes
|
||
padded_page_size = round_up(actual_page_size, spec.alignment)
|
||
if padded_page_size != actual_page_size:
|
||
object.__setattr__(spec, "page_size_padded", padded_page_size)
|
||
|
||
|
||
@dataclass(frozen=True, kw_only=True)
|
||
class TQFullAttentionSpec(FullAttentionSpec):
|
||
"""FullAttentionSpec with TQ-aware page size.
|
||
|
||
Python equivalent of the C++ TQ4FullAttentionSpec. Overrides
|
||
real_page_size_bytes to use TQ slot bytes instead of the raw
|
||
head_size * dtype formula.
|
||
"""
|
||
|
||
tq_slot_size: int = 0
|
||
|
||
@property
|
||
def real_page_size_bytes(self) -> int:
|
||
if self.tq_slot_size > 0:
|
||
return self.block_size * self.num_kv_heads * self.tq_slot_size
|
||
return super().real_page_size_bytes
|
||
|
||
@classmethod
|
||
def merge(cls, specs: list[Self]) -> Self:
|
||
merged = super().merge(specs)
|
||
assert all(s.tq_slot_size == specs[0].tq_slot_size for s in specs), (
|
||
"All TQ layers in the same KV cache group must use the same tq_slot_size."
|
||
)
|
||
return replace(merged, tq_slot_size=specs[0].tq_slot_size)
|
||
|
||
|
||
@dataclass(frozen=True, kw_only=True)
|
||
class MLAAttentionSpec(FullAttentionSpec):
|
||
# TODO(Lucas/Chen): less hacky way to do this
|
||
cache_dtype_str: str | None = None
|
||
# DeepseekV4 only fields. Non-DeepseekV4 MLA models leave these at defaults.
|
||
alignment: int | None = None # Default to None for no padding.
|
||
compress_ratio: int = 1 # Default to 1 for no compression.
|
||
model_version: str | None = None
|
||
|
||
def __post_init__(self):
|
||
super().__post_init__()
|
||
_apply_alignment_padding(self)
|
||
|
||
@property
|
||
def storage_block_size(self) -> int:
|
||
return self.block_size // self.compress_ratio
|
||
|
||
@property
|
||
def real_page_size_bytes(self) -> int:
|
||
if self.cache_dtype_str == "fp8_ds_mla":
|
||
if self.model_version == "deepseek_v4":
|
||
# DeepseekV4: 448B NoPE + 128B RoPE + 8B fp8 scale = 584B per token.
|
||
# head_size stays semantic (512); bytes are determined here.
|
||
return self.storage_block_size * 584
|
||
# V3.2 main MLA: 656-byte custom layout (kv_lora_rank=512 +
|
||
# qk_rope_head_dim=64, head_size=576). See flashmla_sparse.py.
|
||
return self.block_size * 656
|
||
if self.kv_quant_mode == KVQuantMode.INT4_PER_TOKEN_HEAD:
|
||
head_dim = self.head_size // 2
|
||
else:
|
||
head_dim = self.head_size
|
||
return (
|
||
self.storage_block_size
|
||
* self.num_kv_heads
|
||
* head_dim
|
||
* get_dtype_size(self.dtype)
|
||
)
|
||
|
||
@classmethod
|
||
def merge(cls, specs: list[Self]) -> Self:
|
||
assert all(isinstance(spec, MLAAttentionSpec) for spec in specs), (
|
||
"All attention layers in the same KV cache group must be MLAAttentionSpec."
|
||
)
|
||
cache_dtype_str_set = set(spec.cache_dtype_str for spec in specs)
|
||
compress_ratio_set = set(spec.compress_ratio for spec in specs)
|
||
model_version_set = set(spec.model_version for spec in specs)
|
||
block_stride_set = set(spec.indexes_kv_by_block_stride for spec in specs)
|
||
assert (
|
||
len(cache_dtype_str_set) == 1
|
||
and len(compress_ratio_set) == 1
|
||
and len(model_version_set) == 1
|
||
and len(block_stride_set) == 1
|
||
), (
|
||
"All attention layers in the same KV cache group must use the same "
|
||
"quantization method, compress ratio, model version, and KV block "
|
||
"stride indexing."
|
||
)
|
||
return cls(
|
||
block_size=specs[0].block_size,
|
||
num_kv_heads=specs[0].num_kv_heads,
|
||
head_size=specs[0].head_size,
|
||
dtype=specs[0].dtype,
|
||
kv_quant_mode=specs[0].kv_quant_mode,
|
||
page_size_padded=specs[0].page_size_padded,
|
||
indexes_kv_by_block_stride=block_stride_set.pop(),
|
||
cache_dtype_str=cache_dtype_str_set.pop(),
|
||
compress_ratio=compress_ratio_set.pop(),
|
||
model_version=model_version_set.pop(),
|
||
)
|
||
|
||
|
||
@dataclass(frozen=True, kw_only=True)
|
||
class HiddenStateCacheSpec(MLAAttentionSpec):
|
||
"""Marker for hidden-state cache layers used by extract_hidden_states."""
|
||
|
||
pass
|
||
|
||
|
||
@dataclass(frozen=True, kw_only=True)
|
||
class RSWASpec(FullAttentionSpec):
|
||
"""KV cache spec for Reference Sliding Window Attention (R-SWA).
|
||
|
||
Prefill (image + text prompt) tokens are always globally visible.
|
||
Only the last ``rswa_window`` generated tokens are kept in the KV cache;
|
||
gap blocks (between the prefill tail and the current decode window) are
|
||
evicted during each decode step to bound memory at
|
||
O(prefix_blocks + window_blocks).
|
||
"""
|
||
|
||
rswa_window: int
|
||
|
||
@classmethod
|
||
def merge(cls, specs: list[RSWASpec]) -> RSWASpec:
|
||
assert all(isinstance(spec, RSWASpec) for spec in specs), (
|
||
"All attention layers in the same KV cache group must be RSWASpec."
|
||
)
|
||
rswa_windows = {spec.rswa_window for spec in specs}
|
||
assert len(rswa_windows) == 1, (
|
||
f"All R-SWA layers must share the same rswa_window, got {rswa_windows}"
|
||
)
|
||
# Delegate common field merging to the parent, then reattach rswa_window.
|
||
base = FullAttentionSpec.merge(specs) # type: ignore[arg-type]
|
||
return cls(
|
||
block_size=base.block_size,
|
||
num_kv_heads=base.num_kv_heads,
|
||
head_size=base.head_size,
|
||
head_size_v=base.head_size_v,
|
||
dtype=base.dtype,
|
||
kv_quant_mode=base.kv_quant_mode,
|
||
page_size_padded=base.page_size_padded,
|
||
indexes_kv_by_block_stride=base.indexes_kv_by_block_stride,
|
||
sliding_window=base.sliding_window,
|
||
attention_chunk_size=base.attention_chunk_size,
|
||
non_causal=base.non_causal,
|
||
rswa_window=rswa_windows.pop(),
|
||
)
|
||
|
||
|
||
@dataclass(frozen=True, kw_only=True)
|
||
class ChunkedLocalAttentionSpec(AttentionSpec):
|
||
attention_chunk_size: int
|
||
|
||
def max_admission_blocks_per_request(
|
||
self, max_in_flight_tokens: int, max_model_len: int
|
||
) -> int:
|
||
"""Per-request admission cap, in blocks.
|
||
|
||
Single source of truth for both startup pool sizing
|
||
(`max_memory_usage_bytes`) and the runtime admission gate, so requests
|
||
admitted by startup can also be admitted at runtime.
|
||
|
||
`max_in_flight_tokens` is the max tokens scheduled but not yet settled
|
||
(one batch per concurrent step); see `VllmConfig.max_in_flight_tokens`.
|
||
"""
|
||
# During chunked prefill, we hold KV for at most one chunk window plus
|
||
# the in-flight tokens, since frees happen on the processed-token basis.
|
||
num_tokens = min(
|
||
self.attention_chunk_size + max_in_flight_tokens, max_model_len
|
||
)
|
||
return cdiv(num_tokens, self.block_size)
|
||
|
||
def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:
|
||
max_blocks = self.max_admission_blocks_per_request(
|
||
max_in_flight_tokens=vllm_config.max_in_flight_tokens,
|
||
max_model_len=vllm_config.model_config.max_model_len,
|
||
)
|
||
return max_blocks * self.page_size_bytes
|
||
|
||
def is_uniform_with_collection(
|
||
self, kv_cache_specs: dict[str, KVCacheSpec]
|
||
) -> bool:
|
||
return all(
|
||
isinstance(spec, ChunkedLocalAttentionSpec)
|
||
and spec.attention_chunk_size == self.attention_chunk_size
|
||
for spec in kv_cache_specs.values()
|
||
)
|
||
|
||
|
||
@dataclass(frozen=True, kw_only=True)
|
||
class SlidingWindowSpec(AttentionSpec):
|
||
sliding_window: int
|
||
head_size_v: int = None # type: ignore[assignment]
|
||
|
||
def __post_init__(self):
|
||
if self.head_size_v is None:
|
||
object.__setattr__(self, "head_size_v", self.head_size)
|
||
|
||
@property
|
||
def real_page_size_bytes(self) -> int:
|
||
# Mirror ``FullAttentionSpec.real_page_size_bytes`` for NVFP4 KV cache.
|
||
if self.kv_quant_mode.is_nvfp4:
|
||
last_dim = nvfp4_kv_cache_full_dim(
|
||
self.head_size
|
||
) + nvfp4_kv_cache_full_dim(self.head_size_v)
|
||
return (
|
||
self.block_size
|
||
* self.num_kv_heads
|
||
* last_dim
|
||
* get_dtype_size(self.dtype)
|
||
)
|
||
return (
|
||
self.block_size
|
||
* self.num_kv_heads
|
||
* (self.head_size + self.head_size_v)
|
||
* get_dtype_size(self.dtype)
|
||
)
|
||
|
||
def max_admission_blocks_per_request(
|
||
self, max_in_flight_tokens: int, max_model_len: int
|
||
) -> int:
|
||
"""Per-request admission cap, in blocks.
|
||
|
||
Single source of truth for both startup pool sizing
|
||
(`max_memory_usage_bytes`) and the runtime admission gate. Per-request
|
||
real-held blocks plateau at this bound because
|
||
`SlidingWindowManager.remove_skipped_blocks` runs from `allocate_slots`
|
||
before each chunk's `get_num_blocks_to_allocate`.
|
||
|
||
`max_in_flight_tokens` is the max tokens scheduled but not yet settled
|
||
(one batch per concurrent step); see `VllmConfig.max_in_flight_tokens`.
|
||
"""
|
||
# During chunked prefill, we hold KV for the last `sliding_window-1`
|
||
# computed tokens plus the in-flight tokens (frees happen on the
|
||
# processed-token basis); never more than `max_model_len`.
|
||
num_tokens = min(self.sliding_window - 1 + max_in_flight_tokens, max_model_len)
|
||
# +1 because the sliding window may not start from the beginning of
|
||
# the block. E.g. block size 4 and num_token 4 needs two blocks
|
||
# [XXCD][EF] to store the 6-token window [CDEF].
|
||
return cdiv(num_tokens, self.block_size) + 1
|
||
|
||
def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:
|
||
assert vllm_config.parallel_config.decode_context_parallel_size == 1, (
|
||
"DCP not support sliding window."
|
||
)
|
||
max_blocks = self.max_admission_blocks_per_request(
|
||
max_in_flight_tokens=vllm_config.max_in_flight_tokens,
|
||
max_model_len=vllm_config.model_config.max_model_len,
|
||
)
|
||
return max_blocks * self.page_size_bytes
|
||
|
||
def is_uniform_with_collection(
|
||
self, kv_cache_specs: dict[str, KVCacheSpec]
|
||
) -> bool:
|
||
return all(
|
||
isinstance(spec, SlidingWindowSpec)
|
||
and spec.sliding_window == self.sliding_window
|
||
for spec in kv_cache_specs.values()
|
||
)
|
||
|
||
|
||
@dataclass(frozen=True, kw_only=True)
|
||
class SlidingWindowMLASpec(SlidingWindowSpec):
|
||
"""Sliding window attention with MLA cache format."""
|
||
|
||
cache_dtype_str: str | None = None
|
||
# DeepseekV4-only: see MLAAttentionSpec.model_version.
|
||
alignment: int | None = None # Default to None for no padding.
|
||
compress_ratio: int = 1
|
||
model_version: str | None = None
|
||
|
||
def __post_init__(self):
|
||
_apply_alignment_padding(self)
|
||
|
||
@property
|
||
def storage_block_size(self) -> int:
|
||
return self.block_size // self.compress_ratio
|
||
|
||
@property
|
||
def real_page_size_bytes(self) -> int:
|
||
if self.model_version == "deepseek_v4" and self.cache_dtype_str == "fp8_ds_mla":
|
||
# DeepseekV4 FlashMLA: 448B NoPE + 128B RoPE + 8B fp8 scale = 584B
|
||
# per token. FlashInfer's contiguous bf16/fp8 cache falls through to
|
||
# the element-size formula below.
|
||
return self.storage_block_size * 584
|
||
assert self.model_version in (None, "deepseek_v4"), (
|
||
f"Unsupported model version: {self.model_version}"
|
||
)
|
||
return (
|
||
self.storage_block_size
|
||
* self.num_kv_heads
|
||
* self.head_size
|
||
* get_dtype_size(self.dtype)
|
||
)
|
||
|
||
@classmethod
|
||
def merge(cls, specs: list[Self]) -> Self:
|
||
assert all(isinstance(spec, SlidingWindowMLASpec) for spec in specs), (
|
||
"All attention layers in the same KV cache group must be "
|
||
"SlidingWindowMLASpec."
|
||
)
|
||
cache_dtype_str_set = set(spec.cache_dtype_str for spec in specs)
|
||
compress_ratio_set = set(spec.compress_ratio for spec in specs)
|
||
model_version_set = set(spec.model_version for spec in specs)
|
||
sliding_window_set = set(spec.sliding_window for spec in specs)
|
||
block_stride_set = set(spec.indexes_kv_by_block_stride for spec in specs)
|
||
assert (
|
||
len(cache_dtype_str_set) == 1
|
||
and len(compress_ratio_set) == 1
|
||
and len(model_version_set) == 1
|
||
and len(sliding_window_set) == 1
|
||
and len(block_stride_set) == 1
|
||
), (
|
||
"All attention layers in the same KV cache group must use the same "
|
||
"quantization method, compress ratio, model version, sliding "
|
||
"window size, and KV block stride indexing."
|
||
)
|
||
return cls(
|
||
block_size=specs[0].block_size,
|
||
num_kv_heads=specs[0].num_kv_heads,
|
||
head_size=specs[0].head_size,
|
||
dtype=specs[0].dtype,
|
||
page_size_padded=specs[0].page_size_padded,
|
||
indexes_kv_by_block_stride=block_stride_set.pop(),
|
||
sliding_window=sliding_window_set.pop(),
|
||
cache_dtype_str=cache_dtype_str_set.pop(),
|
||
compress_ratio=compress_ratio_set.pop(),
|
||
model_version=model_version_set.pop(),
|
||
)
|
||
|
||
def is_uniform_with_collection(
|
||
self, kv_cache_specs: dict[str, KVCacheSpec]
|
||
) -> bool:
|
||
return all(
|
||
isinstance(spec, SlidingWindowMLASpec)
|
||
and spec.sliding_window == self.sliding_window
|
||
for spec in kv_cache_specs.values()
|
||
)
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class MambaSpec(KVCacheSpec):
|
||
shapes: tuple[tuple[int, ...], ...]
|
||
dtypes: tuple[torch.dtype]
|
||
page_size_padded: int | None = None
|
||
mamba_type: MambaAttentionBackendEnum = MambaAttentionBackendEnum.MAMBA2
|
||
mamba_cache_mode: str = "none"
|
||
num_speculative_blocks: int = 0
|
||
|
||
@property
|
||
def page_size_bytes(self) -> int:
|
||
page_size = sum(
|
||
prod(shape) * get_dtype_size(dtype)
|
||
for (shape, dtype) in zip(self.shapes, self.dtypes)
|
||
)
|
||
if self.page_size_padded is not None:
|
||
assert self.page_size_padded >= page_size
|
||
return self.page_size_padded
|
||
return page_size
|
||
|
||
def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:
|
||
if vllm_config.cache_config.mamba_cache_mode == "all":
|
||
max_model_len = vllm_config.model_config.max_model_len
|
||
return (
|
||
cdiv(max_model_len, self.block_size) + self.num_speculative_blocks
|
||
) * self.page_size_bytes
|
||
elif vllm_config.cache_config.mamba_cache_mode == "align":
|
||
return self.page_size_bytes * (2 + self.num_speculative_blocks)
|
||
else:
|
||
return self.page_size_bytes * (1 + self.num_speculative_blocks)
|
||
|
||
def max_num_blocks_per_req(self, vllm_config: VllmConfig, max_len: int) -> int:
|
||
# Mamba state is replicated across DCP/PCP ranks, never sharded, so
|
||
# no CP scaling applies.
|
||
if vllm_config.cache_config.mamba_cache_mode == "align":
|
||
# Block table rows are position-indexed over the full sequence
|
||
# even though only 2 + num_speculative_blocks state blocks are
|
||
# resident at a time (earlier states are nulled out by
|
||
# remove_skipped_blocks), so the row length must cover max_len
|
||
# rather than max_memory_usage_bytes.
|
||
return cdiv(max_len, self.block_size) + self.num_speculative_blocks
|
||
return cdiv(self.max_memory_usage_bytes(vllm_config), self.page_size_bytes)
|
||
|
||
def is_uniform_with_collection(
|
||
self, kv_cache_specs: dict[str, KVCacheSpec]
|
||
) -> bool:
|
||
return all(
|
||
isinstance(spec, MambaSpec)
|
||
and spec.num_speculative_blocks == self.num_speculative_blocks
|
||
for spec in kv_cache_specs.values()
|
||
)
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class EncoderOnlyAttentionSpec(AttentionSpec):
|
||
def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:
|
||
# Encoder-only layers do not need KV cache
|
||
return 0
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class CrossAttentionSpec(AttentionSpec):
|
||
"""
|
||
KV cache spec for cross-attention layers in encoder-decoder models.
|
||
"""
|
||
|
||
def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:
|
||
# For cross-attention, we need to cache encoder states
|
||
# Get encoder length (e.g., 1500 for Whisper).
|
||
max_encoder_len = vllm_config.scheduler_config.max_num_encoder_input_tokens
|
||
return cdiv(max_encoder_len, self.block_size) * self.page_size_bytes
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class SinkFullAttentionSpec(FullAttentionSpec):
|
||
sink_len: int | None = None
|
||
|
||
@classmethod
|
||
def merge(cls, specs: list[Self]) -> Self:
|
||
"""
|
||
Merge a list of FullAttentionSpec objects into a single
|
||
FullAttentionSpec object.
|
||
"""
|
||
assert all(isinstance(spec, FullAttentionSpec) for spec in specs), (
|
||
"All attention layers in the same KV cache group must be FullAttentionSpec."
|
||
)
|
||
|
||
sliding_window = set(
|
||
spec.sliding_window for spec in specs if spec.sliding_window is not None
|
||
)
|
||
attention_chunk_size = set(
|
||
spec.attention_chunk_size
|
||
for spec in specs
|
||
if spec.attention_chunk_size is not None
|
||
)
|
||
assert not any(isinstance(spec, MLAAttentionSpec) for spec in specs), (
|
||
"MLAAttentionSpec should be merged in MLAAttentionSpec.merge"
|
||
)
|
||
merged_spec = cls(
|
||
block_size=specs[0].block_size,
|
||
num_kv_heads=specs[0].num_kv_heads,
|
||
head_size=specs[0].head_size,
|
||
head_size_v=specs[0].head_size_v,
|
||
sink_len=specs[0].sink_len,
|
||
dtype=specs[0].dtype,
|
||
kv_quant_mode=specs[0].kv_quant_mode,
|
||
page_size_padded=specs[0].page_size_padded,
|
||
indexes_kv_by_block_stride=specs[0].indexes_kv_by_block_stride,
|
||
sliding_window=cls.merge_window_sizes(sliding_window),
|
||
attention_chunk_size=cls.merge_window_sizes(attention_chunk_size),
|
||
non_causal=any(spec.non_causal for spec in specs),
|
||
)
|
||
for spec in specs:
|
||
for f in fields(AttentionSpec):
|
||
assert getattr(spec, f.name) == getattr(merged_spec, f.name), (
|
||
"All attention layers in the same KV cache group must have "
|
||
"the same attention spec."
|
||
)
|
||
assert (merged_spec.sliding_window is not None) + (
|
||
merged_spec.attention_chunk_size is not None
|
||
) <= 1, (
|
||
"Model with both sliding window layers and chunked local attention "
|
||
"layers is not supported."
|
||
)
|
||
return merged_spec
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class UniformTypeKVCacheSpecs(KVCacheSpec):
|
||
"""
|
||
A KV cache spec for multiple layers with the same type of attention. Here,
|
||
same types means always need the same number of token slots. For example,
|
||
sliding window attentions with different window sizes are not the same type
|
||
and should not be merged into one UniformTypeKVCacheSpecs.
|
||
"""
|
||
|
||
kv_cache_specs: dict[str, KVCacheSpec]
|
||
|
||
@property
|
||
def page_size_bytes(self) -> int:
|
||
return sum(spec.page_size_bytes for spec in self.kv_cache_specs.values())
|
||
|
||
def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:
|
||
max_num_pages = max(
|
||
cdiv(spec.max_memory_usage_bytes(vllm_config), spec.page_size_bytes)
|
||
for spec in self.kv_cache_specs.values()
|
||
)
|
||
return max_num_pages * self.page_size_bytes
|
||
|
||
@classmethod
|
||
def is_uniform_type(cls, kv_cache_specs: dict[str, KVCacheSpec]) -> bool:
|
||
"""
|
||
Whether all layers have the same type of KV cache spec.
|
||
|
||
Uses the registry to determine grouping base classes, so custom specs
|
||
that inherit from FullAttentionSpec are treated as full attention.
|
||
"""
|
||
block_sizes = set(spec.block_size for spec in kv_cache_specs.values())
|
||
if len(block_sizes) > 1:
|
||
# Different block sizes, not uniform.
|
||
return False
|
||
first_spec = next(iter(kv_cache_specs.values()))
|
||
return first_spec.is_uniform_with_collection(kv_cache_specs)
|
||
|
||
@classmethod
|
||
def from_specs(cls, kv_cache_specs: dict[str, KVCacheSpec]) -> Self | None:
|
||
"""
|
||
Return a SameTypeKVCacheSpecs object if all layers have the same type
|
||
of KV cache spec. Return None if not.
|
||
"""
|
||
if cls.is_uniform_type(kv_cache_specs):
|
||
block_size = next(iter(kv_cache_specs.values())).block_size
|
||
return cls(block_size=block_size, kv_cache_specs=kv_cache_specs)
|
||
else:
|
||
return None
|
||
|
||
# NOTE: below util functions are only used by DeepseekV4 for now.
|
||
def get_page_sizes(self) -> list[int]:
|
||
return list(set(spec.page_size_bytes for spec in self.kv_cache_specs.values()))
|
||
|
||
def get_num_layer_tuples(self) -> int:
|
||
return Counter(
|
||
spec.page_size_bytes for spec in self.kv_cache_specs.values()
|
||
).most_common(1)[0][1]
|
||
|
||
def max_memory_usage_pages(self, vllm_config: VllmConfig) -> int:
|
||
return max(
|
||
cdiv(spec.max_memory_usage_bytes(vllm_config), spec.page_size_bytes)
|
||
for spec in self.kv_cache_specs.values()
|
||
)
|
||
|
||
|
||
def get_kv_cache_spec_kind(kv_cache_spec: KVCacheSpec) -> KVCacheSpecKind:
|
||
if isinstance(kv_cache_spec, UniformTypeKVCacheSpecs):
|
||
inner_kinds = {
|
||
get_kv_cache_spec_kind(spec)
|
||
for spec in kv_cache_spec.kv_cache_specs.values()
|
||
}
|
||
if len(inner_kinds) == 1:
|
||
return next(iter(inner_kinds))
|
||
return KVCacheSpecKind.UNKNOWN
|
||
# Keep subclass checks before base classes so specialized specs keep their
|
||
# more precise kind.
|
||
if isinstance(kv_cache_spec, SlidingWindowMLASpec):
|
||
return KVCacheSpecKind.SLIDING_WINDOW_MLA
|
||
if isinstance(kv_cache_spec, MLAAttentionSpec):
|
||
return KVCacheSpecKind.MLA_ATTENTION
|
||
if isinstance(kv_cache_spec, SinkFullAttentionSpec):
|
||
return KVCacheSpecKind.SINK_FULL_ATTENTION
|
||
if isinstance(kv_cache_spec, FullAttentionSpec):
|
||
return KVCacheSpecKind.FULL_ATTENTION
|
||
if isinstance(kv_cache_spec, ChunkedLocalAttentionSpec):
|
||
return KVCacheSpecKind.CHUNKED_LOCAL_ATTENTION
|
||
if isinstance(kv_cache_spec, SlidingWindowSpec):
|
||
return KVCacheSpecKind.SLIDING_WINDOW
|
||
if isinstance(kv_cache_spec, MambaSpec):
|
||
return KVCacheSpecKind.MAMBA
|
||
if isinstance(kv_cache_spec, EncoderOnlyAttentionSpec):
|
||
return KVCacheSpecKind.ENCODER_ONLY_ATTENTION
|
||
if isinstance(kv_cache_spec, CrossAttentionSpec):
|
||
return KVCacheSpecKind.CROSS_ATTENTION
|
||
return KVCacheSpecKind.UNKNOWN
|
||
|
||
|
||
def get_kv_cache_spec_sliding_window(kv_cache_spec: KVCacheSpec) -> int | None:
|
||
if isinstance(kv_cache_spec, UniformTypeKVCacheSpecs):
|
||
inner_windows = {
|
||
get_kv_cache_spec_sliding_window(spec)
|
||
for spec in kv_cache_spec.kv_cache_specs.values()
|
||
}
|
||
return next(iter(inner_windows)) if len(inner_windows) == 1 else None
|
||
if isinstance(kv_cache_spec, SlidingWindowSpec):
|
||
return kv_cache_spec.sliding_window
|
||
return None
|
||
|
||
|
||
@dataclass
|
||
class KVCacheTensor:
|
||
"""
|
||
A class for specifying how the workers should initialize the KV cache.
|
||
"""
|
||
|
||
size: int # size of the KV cache tensor in bytes
|
||
shared_by: list[str] # layer names that share the same KV cache tensor
|
||
offset: int = 0 # byte offset of this layer within a contiguous block
|
||
block_stride: int = 0 # total bytes per block in a packed layout (0 = not packed)
|
||
|
||
|
||
@dataclass
|
||
class KVCacheGroupSpec:
|
||
"""
|
||
Represents a group of model layers that share the same KV cache block table.
|
||
These layers are regarded as one layer in the KV cache manager.
|
||
"""
|
||
|
||
# The names of model layers in this group
|
||
layer_names: list[str]
|
||
# The KV cache spec of this manager layer
|
||
kv_cache_spec: KVCacheSpec
|
||
# Whether this group contains EAGLE/MTP draft attention layers.
|
||
is_eagle_group: bool = False
|
||
|
||
|
||
@dataclass
|
||
class KVCacheConfig:
|
||
"""
|
||
The KV cache configuration of a model.
|
||
"""
|
||
|
||
num_blocks: int
|
||
"""The number of KV cache blocks"""
|
||
kv_cache_tensors: list[KVCacheTensor]
|
||
"""How should model runner initialize the KV cache tensors for each layer"""
|
||
kv_cache_groups: list[KVCacheGroupSpec]
|
||
"""
|
||
The kv cache groups of the model.
|
||
For models with only one type of attention, there is only one group that
|
||
contains all layers.
|
||
For models with multiple types of attention, there will be multiple groups,
|
||
see `_get_kv_cache_config_uniform_page_size` for more details.
|
||
"""
|
||
|
||
@property
|
||
def has_mamba_layers(self) -> bool:
|
||
return any(isinstance(g.kv_cache_spec, MambaSpec) for g in self.kv_cache_groups)
|
||
|
||
@property
|
||
def needs_kv_cache_zeroing(self) -> bool:
|
||
return self.has_mamba_layers
|