Jürgen SchmiedandClaude
a80e41d690
RecoverSSM's commit-plan and conv-window compaction kernels are model independent. Move them from models/kimi_k3 into model_executor/layers/mamba/ops/recoverssm.py (Kimi-K3 keeps its private names as import aliases; kernel source unchanged), add mamba/recoverssm_utils.py (activation predicate, argument validation) and the generic spec-row/request mapping helpers to v1/attention/backends/recoverssm_metadata.py. CacheConfig.use_kda_recoverssm becomes use_recoverssm: the switch now selects RecoverSSM for any model whose recurrent layers implement it, not only KDA. No functional change for Kimi-K3. Signed-off-by: Jürgen Schmied <juergenschmied70@gmail.com> Co-authored-by: Claude <noreply@anthropic.com>
380 lines
18 KiB
Python
380 lines
18 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
from collections.abc import Callable
|
|
from dataclasses import field
|
|
from functools import cache
|
|
from typing import Any, ClassVar, Literal
|
|
|
|
from pydantic import Field, field_validator, model_validator
|
|
|
|
from vllm.config.utils import config
|
|
from vllm.logger import init_logger
|
|
from vllm.utils.torch_utils import (
|
|
is_quantized_kv_cache,
|
|
kv_cache_uses_per_token_head_scales,
|
|
)
|
|
from vllm.v1.kv_cache_layout import KVCacheLayout
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
_LAYOUT_COMPAT_ALIASES = {
|
|
"NHD": "LBNHC",
|
|
"HND": "LBHNC",
|
|
}
|
|
|
|
|
|
@cache
|
|
def _layout_from_name(layout_name: str) -> KVCacheLayout:
|
|
layout_name = _LAYOUT_COMPAT_ALIASES.get(layout_name, layout_name)
|
|
try:
|
|
return KVCacheLayout[layout_name]
|
|
except KeyError:
|
|
raise ValueError(
|
|
f"Unknown KV cache layout {layout_name!r}. "
|
|
f"Valid layouts: {[m.name for m in KVCacheLayout]}"
|
|
) from None
|
|
|
|
|
|
CacheDType = Literal[
|
|
"auto",
|
|
"float16",
|
|
"bfloat16",
|
|
"fp8",
|
|
"fp8_e4m3",
|
|
"fp8_e5m2",
|
|
"fp8_inc",
|
|
"fp8_ds_mla",
|
|
"nvfp4_ds_mla",
|
|
"turboquant_k8v4",
|
|
"turboquant_4bit_nc",
|
|
"turboquant_k3v4_nc",
|
|
"turboquant_3bit_nc",
|
|
"int4_per_token_head",
|
|
"int8_per_token_head",
|
|
"fp8_per_token_head",
|
|
"nvfp4",
|
|
"nvfp4_4over6",
|
|
]
|
|
|
|
|
|
MambaDType = Literal["auto", "float32", "float16", "bfloat16"]
|
|
MambaCacheMode = Literal["all", "align", "none"]
|
|
PrefixCachingHashAlgo = Literal["sha256", "sha256_cbor", "xxhash", "xxhash_cbor"]
|
|
KVOffloadingBackend = Literal["native", "lmcache"]
|
|
|
|
|
|
@config
|
|
class CacheConfig:
|
|
"""Configuration for the KV cache."""
|
|
|
|
DEFAULT_BLOCK_SIZE: ClassVar[int] = 16
|
|
|
|
block_size: int = Field(default=None, gt=0) # type: ignore[assignment]
|
|
"""Size of a contiguous cache block in number of tokens.
|
|
Accepts None (meaning "use default"). After construction, always int."""
|
|
user_specified_block_size: bool = field(default=False, init=False)
|
|
"""Whether block_size was explicitly provided. Derived automatically."""
|
|
user_specified_mamba_block_size: bool = field(default=False, init=False)
|
|
"""Whether mamba_block_size was explicitly provided. Derived automatically."""
|
|
kv_cache_layout: str | None = field(default=None, init=False)
|
|
"""Resolved physical KV cache layout name (a ``KVCacheLayout`` member).
|
|
|
|
``None`` means the layout has not been resolved yet. The engine core
|
|
resolves it once (``resolve_kv_cache_layout``) before memory profiling, and
|
|
every worker process adopts the resolved name — via the
|
|
``set_kv_cache_layout`` RPC, or ``KVCacheConfig.kv_cache_layout`` for
|
|
workers spawned after resolution — before the KV cache is allocated. Once
|
|
set the value is final and read with ``get_resolved_kv_cache_layout``,
|
|
which raises on ``None``. Tests and standalone tools may pre-set a value,
|
|
which resolution then honors as-is."""
|
|
prefix_match_unit: int | None = Field(default=None, gt=0)
|
|
"""The finest token boundary (in tokens) a prefix-cache hit can land on.
|
|
|
|
Prefix-cache keys are computed every `prefix_match_unit` tokens. It can
|
|
be set finer than the physical KV cache block sizes (e.g. 32 vs a
|
|
1024-token hybrid-model block) as long as every KV cache group's
|
|
`block_size` is divisible by it, enabling cache hits at boundaries
|
|
inside a physical block. It controls matching granularity only, not how
|
|
often states are stored.
|
|
|
|
This equals to the `hash_block_size` used throughout the KV cache code.
|
|
"""
|
|
gpu_memory_utilization: float = Field(default=0.92, gt=0, le=1)
|
|
"""The fraction of GPU memory to be used for the model executor, which can
|
|
range from 0 to 1. For example, a value of 0.5 would imply 50% GPU memory
|
|
utilization. If unspecified, will use the default value of 0.92. This is a
|
|
per-instance limit, and only applies to the current vLLM instance. It does
|
|
not matter if you have another vLLM instance running on the same GPU. For
|
|
example, if you have two vLLM instances running on the same GPU, you can
|
|
set the GPU memory utilization to 0.5 for each instance. On non-GPU
|
|
installs, this value controls the corresponding device memory utilization.
|
|
"""
|
|
|
|
@property
|
|
def device_memory_utilization(self) -> float:
|
|
"""Device-neutral alias for ``gpu_memory_utilization``."""
|
|
return self.gpu_memory_utilization
|
|
|
|
@device_memory_utilization.setter
|
|
def device_memory_utilization(self, value: float) -> None:
|
|
self.gpu_memory_utilization = value
|
|
|
|
cache_dtype: CacheDType = "auto"
|
|
"""Data type for kv cache storage. If "auto", will use model data type.
|
|
CUDA 11.8+ supports fp8 (=fp8_e4m3) and fp8_e5m2. ROCm (AMD GPU) supports
|
|
fp8 (=fp8_e4m3). Intel Gaudi (HPU) supports fp8 (using fp8_inc).
|
|
Some models (namely DeepSeekV3.2) default to fp8, set to bfloat16 to use
|
|
bfloat16 instead, this is an invalid option for models that do not default
|
|
to fp8.
|
|
"nvfp4_4over6" uses the NVFP4 layout and selects between max/6 and max/4
|
|
scales per 16 values by minimizing squared reconstruction error.
|
|
"""
|
|
is_attention_free: bool = False
|
|
"""Whether the model is attention-free. This is primarily set in
|
|
`ModelConfig` and that value should be manually duplicated here."""
|
|
num_gpu_blocks_override: int | None = None
|
|
"""Number of GPU blocks to use. This overrides the profiled `num_gpu_blocks`
|
|
if specified. Does nothing if `None`. Used for testing preemption."""
|
|
sliding_window: int | None = None
|
|
"""Sliding window size for the KV cache. This is primarily set in
|
|
`ModelConfig` and that value should be manually duplicated here."""
|
|
enable_prefix_caching: bool = True
|
|
"""Whether to enable prefix caching."""
|
|
prefix_caching_hash_algo: PrefixCachingHashAlgo = "sha256"
|
|
"""Set the hash algorithm for prefix caching:
|
|
|
|
- "sha256" uses Pickle for object serialization before hashing. This is the current
|
|
default, as SHA256 is the most secure choice to avoid potential hash collisions.
|
|
- "sha256_cbor" provides a reproducible, cross-language compatible hash. It
|
|
serializes objects using canonical CBOR and hashes them with SHA-256.
|
|
- "xxhash" uses Pickle serialization with xxHash (128-bit) for faster,
|
|
non-cryptographic hashing. Requires the optional ``xxhash`` package.
|
|
IMPORTANT: Use of a hashing algorithm that is not considered cryptographically
|
|
secure theoretically increases the risk of hash collisions, which can cause
|
|
undefined behavior or even leak private information in multi-tenant environments.
|
|
Even if collisions are still very unlikely, it is important to consider your
|
|
security risk tolerance against the performance benefits before turning this on.
|
|
- "xxhash_cbor" combines canonical CBOR serialization with xxHash for
|
|
reproducible hashing. Requires the optional ``xxhash`` package."""
|
|
prefix_cache_retention_interval: int | None = Field(default=0, ge=0)
|
|
"""Token interval between retained sliding-window and Mamba prefix-cache
|
|
checkpoints. ``0`` retains only semantic checkpoints, including the latest
|
|
replay boundary and shared-prefix junctions. Positive values additionally
|
|
retain periodic checkpoints at the specified interval, which must be a
|
|
multiple of the scheduler block size. ``None`` retains checkpoints densely.
|
|
Applies only to sliding-window and Mamba cache groups."""
|
|
kv_cache_dtype_skip_layers: list[str] = field(default_factory=list)
|
|
"""Layer patterns to skip KV cache quantization. Accepts layer indices
|
|
(e.g., '0', '2', '4') or attention type names (e.g., 'sliding_window')."""
|
|
mamba_page_size_padded: int | None = None
|
|
""" Optional override for mamba page size; used by hybrid mamba/attention
|
|
models to ensure exact alignment with attention page size."""
|
|
skip_page_size_padded: int | None = None
|
|
"""Optional override for the page size of layers skipped from KV cache
|
|
quantization (``--kv-cache-dtype-skip-layers``); set during block-size
|
|
alignment so unquantized skip layers pad up to the quantized primary's
|
|
page."""
|
|
mamba_block_size: int | None = Field(default=None, gt=0)
|
|
"""Size of a contiguous cache block in number of tokens for mamba cache.
|
|
Can be set only when prefix caching is enabled.
|
|
Value must be a multiple of 8 to align with causal_conv1d kernel."""
|
|
mamba_cache_dtype: MambaDType = "auto"
|
|
"""The data type to use for the Mamba cache (both the conv as well as the
|
|
ssm state). If set to 'auto', the data type will be inferred from the model
|
|
config."""
|
|
mamba_ssm_cache_dtype: MambaDType = "auto"
|
|
"""The data type to use for the Mamba cache (ssm state only, conv state will
|
|
still be controlled by mamba_cache_dtype). If set to 'auto', the data type
|
|
for the ssm state will be determined by mamba_cache_dtype."""
|
|
mamba_cache_mode: MambaCacheMode = "none"
|
|
"""The cache strategy for Mamba layers:
|
|
|
|
- "none": set when prefix caching is disabled.
|
|
- "all": cache the mamba state of all tokens at position i * block_size.
|
|
- "align": only cache the mamba state of the last token of each scheduler step and
|
|
when the token is at position i * block_size. This is the default when prefix
|
|
caching is enabled.
|
|
"""
|
|
enable_mamba_shared_prefix_checkpoint: bool = False
|
|
"""Also register a Mamba "align" checkpoint at the shared-prefix junction --
|
|
where an EAGLE/MTP sibling was observed to resume -- instead of only at the
|
|
prompt tail. Off by default; only takes effect with `mamba_cache_mode`
|
|
"align", EAGLE on the Mamba group, and a prefix match unit smaller than the
|
|
Mamba block size."""
|
|
replayssm_buffer_len: int = Field(default=16, gt=0)
|
|
"""ReplaySSM logical history length B for Mamba2. Triton uses B physical
|
|
rows and FlashInfer uses B+1. Kimi-K3 speculative decode does not use B.
|
|
Default 16."""
|
|
use_replayssm: bool = False
|
|
"""Use the ReplaySSM Mamba2 decode kernel: cache recent SSM inputs and skip
|
|
the per-step full-state store, writing the checkpoint back only on flush.
|
|
Requires mamba_cache_mode 'none' or 'align' (prefix caching) and the Triton
|
|
or FlashInfer mamba backend; standard (non-speculative) decode only. In align
|
|
mode flushes are most efficient when mamba_block_size is a multiple of
|
|
replayssm_buffer_len, but this is not required."""
|
|
use_recoverssm: bool = field(default=False, init=False)
|
|
"""Whether Kimi-K3 KDA uses RecoverSSM speculative decode."""
|
|
|
|
# Will be set after profiling.
|
|
num_gpu_blocks: int | None = field(default=None, init=False)
|
|
"""The number of blocks to allocate for GPU memory."""
|
|
num_cpu_blocks: int | None = field(default=None, init=False)
|
|
"""The number of blocks to allocate for CPU memory."""
|
|
|
|
# Set after KV cache initialization.
|
|
effective_attention_block_size: int | None = field(default=None, init=False)
|
|
"""Full-attention block size in tokens, including DCP, or None if unavailable."""
|
|
kv_cache_size_tokens: int | None = field(default=None, init=False)
|
|
"""Per-DP-engine KV cache capacity in tokens (group-aware). Uses
|
|
group-aware capacity since num_gpu_blocks * block_size can be wrong
|
|
for hybrid models where requests occupy multiple KV cache groups."""
|
|
kv_cache_max_concurrency: float | None = field(default=None, init=False)
|
|
"""Per-DP-engine maximum concurrency at max_model_len tokens."""
|
|
|
|
kv_sharing_fast_prefill: bool = False
|
|
"""In some KV sharing setups, e.g. YOCO (https://arxiv.org/abs/2405.05254),
|
|
some layers can skip tokens corresponding to prefill. This flag enables
|
|
attention metadata for eligible layers to be overridden with metadata
|
|
necessary for implementing this optimization in some models (e.g. Gemma3n)
|
|
"""
|
|
|
|
swa_bounded_replay: bool = True
|
|
"""Keep the sliding-window KV of models that support it (DeepSeek-V4.1)
|
|
out of prefix caching and rebuild it after a prefix hit by recomputing the
|
|
hit's last window. Requires model runner V2."""
|
|
|
|
kv_cache_memory_bytes: int | None = None
|
|
"""Size of KV Cache per GPU in bytes. By default, this is set to None
|
|
and vllm can automatically infer the kv cache size based on
|
|
gpu_memory_utilization. However, users may want to manually specify
|
|
the kv cache memory size. kv_cache_memory_bytes allows more fine-grain
|
|
control of how much memory gets used when compared with using
|
|
gpu_memory_utilization. Note that kv_cache_memory_bytes
|
|
(when not-None) ignores gpu_memory_utilization"""
|
|
|
|
kv_offloading_size: float | None = None
|
|
"""Size of the KV cache offloading buffer in GiB. When TP > 1, this is
|
|
the total buffer size summed across all TP ranks. By default, this is set
|
|
to None, which means no KV offloading is enabled. When set, vLLM will
|
|
enable KV cache offloading to CPU using the kv_offloading_backend."""
|
|
|
|
kv_offloading_backend: KVOffloadingBackend = "native"
|
|
"""The backend to use for KV cache offloading. Supported backends include
|
|
'native' (vLLM native CPU offloading), 'lmcache'.
|
|
KV offloading is only activated when kv_offloading_size is set."""
|
|
|
|
def compute_hash(self) -> str:
|
|
"""WARNING: Whenever a new field is added to this config,
|
|
ensure that it is included in the factors list if
|
|
it affects the computation graph.
|
|
|
|
Provide a hash that uniquely identifies all the configs
|
|
that affect the structure of the computation
|
|
graph from input ids/embeddings to the final hidden states,
|
|
excluding anything before input ids/embeddings and after
|
|
the final hidden states.
|
|
"""
|
|
ignored_factors = {
|
|
# Runtime/derived knobs that don't affect compiled graph shape
|
|
"gpu_memory_utilization",
|
|
"kv_cache_memory_bytes",
|
|
"is_attention_free",
|
|
"num_gpu_blocks_override",
|
|
"enable_prefix_caching",
|
|
"prefix_caching_hash_algo",
|
|
"prefix_cache_retention_interval",
|
|
# Prefix-caching implementation detail (doesn't affect compiled graph).
|
|
"prefix_match_unit",
|
|
"enable_mamba_shared_prefix_checkpoint",
|
|
"mamba_page_size_padded",
|
|
"skip_page_size_padded",
|
|
"user_specified_block_size",
|
|
"user_specified_mamba_block_size",
|
|
"_block_size_resolved",
|
|
# Post-init/derived counters
|
|
"num_gpu_blocks",
|
|
"num_cpu_blocks",
|
|
"effective_attention_block_size",
|
|
"kv_cache_size_tokens",
|
|
"kv_cache_max_concurrency",
|
|
# Feature toggles not impacting compiled graph shape
|
|
"kv_sharing_fast_prefill",
|
|
"swa_bounded_replay",
|
|
}
|
|
|
|
from vllm.config.utils import get_hash_factors, hash_factors
|
|
|
|
factors = get_hash_factors(self, ignored_factors)
|
|
return hash_factors(factors)
|
|
|
|
def metrics_info(self):
|
|
# convert cache_config to dict(key: str, value: str) for prometheus
|
|
# metrics info
|
|
return {key: str(value) for key, value in self.__dict__.items()}
|
|
|
|
_block_size_resolved: bool = field(default=False, init=False)
|
|
"""Guard against pydantic re-running _apply_block_size_default."""
|
|
|
|
@field_validator("block_size", mode="wrap")
|
|
@classmethod
|
|
def _skip_none_validation(cls, value: Any, handler: Callable) -> Any:
|
|
if value is None:
|
|
return value
|
|
return handler(value)
|
|
|
|
@model_validator(mode="after")
|
|
def _apply_block_size_default(self) -> "CacheConfig":
|
|
# Pydantic re-runs validators when CacheConfig is nested inside
|
|
# another pydantic model (e.g. VllmConfig). Guard against that.
|
|
if self._block_size_resolved:
|
|
return self
|
|
self._block_size_resolved = True
|
|
if self.block_size is None:
|
|
self.block_size = self.DEFAULT_BLOCK_SIZE
|
|
else:
|
|
self.user_specified_block_size = True
|
|
if self.mamba_block_size is not None:
|
|
self.user_specified_mamba_block_size = True
|
|
return self
|
|
|
|
@field_validator("mamba_cache_mode", mode="after")
|
|
@classmethod
|
|
def _validate_mamba_cache_mode(cls, mode: MambaCacheMode) -> MambaCacheMode:
|
|
if mode == "all":
|
|
logger.warning_once(
|
|
"Mamba cache mode 'all' is deprecated and will be removed in an "
|
|
"upcoming release. If this is a problem, please open an issue "
|
|
"at https://github.com/vllm-project/vllm/issues."
|
|
)
|
|
return mode
|
|
|
|
@field_validator("cache_dtype", mode="after")
|
|
@classmethod
|
|
def _validate_cache_dtype(cls, cache_dtype: CacheDType) -> CacheDType:
|
|
if kv_cache_uses_per_token_head_scales(cache_dtype):
|
|
logger.info(
|
|
"Using %s data type to store kv cache. It reduces the GPU "
|
|
"memory footprint and boosts the performance. "
|
|
"Dynamic per-token-head scales will be computed at runtime.",
|
|
str(cache_dtype),
|
|
)
|
|
elif is_quantized_kv_cache(cache_dtype):
|
|
logger.info(
|
|
"Using %s data type to store kv cache. It reduces the GPU "
|
|
"memory footprint and boosts the performance. "
|
|
"Meanwhile, it may cause accuracy drop without a proper "
|
|
"scaling factor",
|
|
str(cache_dtype),
|
|
)
|
|
return cache_dtype
|
|
|
|
def get_resolved_kv_cache_layout(self) -> KVCacheLayout:
|
|
if self.kv_cache_layout is None:
|
|
raise ValueError(
|
|
"KV cache layout has not been resolved yet; it is resolved once "
|
|
"by the engine core (resolve_kv_cache_layout) unless explicitly "
|
|
"set by the user."
|
|
)
|
|
return _layout_from_name(self.kv_cache_layout)
|