【vllm】(三)vLLM v1 Core — 模块超深度逐行分析之三

第九章:KV Cache 工具体系逐行解析 --- kv_cache_utils.py

9.1 模块总览

kv_cache_utils.py(1693行)是 vLLM v1 KV Cache 管理的基础设施层,定义了:

  • 类型系统:BlockHash、BlockHashWithGroupId 等核心类型
  • 数据结构:KVCacheBlock 元数据、FreeKVCacheBlockQueue 双向链表
  • Hash 计算:block hash 生成、extra keys 机制
  • 内存估算:KV cache 内存需求计算、max model len 估算
  • 分组策略:KV cache group 划分与配置生成
  • 辅助工具:BlockHashListWithBlockSize 惰性转换

9.2 导入区(1-30行)

python 复制代码
import copy
import hashlib
import os
from collections import defaultdict
from collections.abc import Callable, Iterable, Iterator, Sequence
from dataclasses import dataclass, replace
from functools import partial
from typing import Any, NewType, TypeAlias, overload
  • copy:用于深拷贝 KV cache 配置(generate_scheduler_kv_cache_config中)
  • hashlib:SHA-256 哈希用于 prompt embeds 额外键
  • os:读取 PYTHONHASHSEED 环境变量
  • defaultdict:构建 same_type_layers 分组映射
  • NewType:创建 BlockHash 等语义类型别名
  • replace:dataclass 不可变更新(unify_kv_cache_spec_page_size
  • overload:为 get_block_ids__getitem__ 提供精确类型签名

vllm.utils.hashing 导入 sha256_cborxxhash_cbor------这是两种 CBOR 序列化哈希函数,用于 block hash 计算。从 vllm.v1.kv_cache_interface 导入各种 KVCacheSpec 子类,涵盖 FullAttention、SlidingWindow、ChunkedLocal、Mamba 等不同注意力机制的 spec。

9.3 BlockHash 类型体系(32-58行)

python 复制代码
BlockHash = NewType("BlockHash", bytes)

设计意图BlockHash 本质是 bytes,但通过 NewType 赋予独立类型身份。这样在函数签名中,编译器(mypy/pyright)可以捕获将普通 bytes 误传为 BlockHash 的错误。底层是 32 字节的 SHA-256 或 xxhash 输出。

python 复制代码
BlockHashWithGroupId = NewType("BlockHashWithGroupId", bytes)

设计意图 :将 BlockHash 与 KV cache group ID 打包为单个 bytes 对象。原始设计可能用 (BlockHash, int) 元组,但元组创建/解构有 Python 对象分配开销。将 group_id 编码为 4 字节大端整数追加到 block_hash 后面,形成紧凑的 bytes 键,既省内存又利于 dict 查找性能。

python 复制代码
ExternalBlockHash: TypeAlias = bytes | int

设计意图 :外部(KV events)使用的 block hash 类型,兼容旧版 int 和新版 bytes。向后兼容桥接------旧系统用 64-bit int 表示 hash,新系统用 SHA-256 bytes。

打包/解包函数

python 复制代码
def make_block_hash_with_group_id(block_hash: BlockHash, group_id: int) -> BlockHashWithGroupId:
    return BlockHashWithGroupId(block_hash + group_id.to_bytes(4, "big", signed=False))

block_hash(32字节)与 group_id(4字节大端)拼接为 36 字节。大端字节序保证跨平台一致性。signed=False 断言 group_id 非负。

python 复制代码
def get_block_hash(key: BlockHashWithGroupId) -> BlockHash:
    return BlockHash(key[:-4])

def get_group_id(key: BlockHashWithGroupId) -> int:
    return int.from_bytes(key[-4:], "big", signed=False)

通过固定偏移量切片实现 O(1) 解包------无需解析分隔符。

python 复制代码
def maybe_convert_block_hash(hash_bytes: BlockHash) -> ExternalBlockHash:
    if not envs.VLLM_KV_EVENTS_USE_INT_BLOCK_HASHES:
        return hash_bytes
    return int.from_bytes(hash_bytes, byteorder="big") & ((1 << 64) - 1)

设计意图 :条件性将 bytes hash 转为 64-bit int(截断高字节),用于 KV events 系统的向后兼容。环境变量 VLLM_KV_EVENTS_USE_INT_BLOCK_HASHES 控制是否使用旧格式。

9.4 NONE_HASH 全局种子(60-84行)

python 复制代码
NONE_HASH: BlockHash
_CBOR_HASH_FUNCTIONS = frozenset({sha256_cbor, xxhash_cbor})

def init_none_hash(hash_fn: Callable[[Any], bytes]):
    global NONE_HASH
    hash_seed = os.getenv("PYTHONHASHSEED")
    if hash_seed is None and hash_fn in _CBOR_HASH_FUNCTIONS:
        logger.warning(...)
    if hash_seed is None:
        NONE_HASH = BlockHash(os.urandom(32))
    else:
        NONE_HASH = BlockHash(hash_fn(hash_seed))

设计意图NONE_HASH 是 block hash 链的"创世种子"------第一个 block 的 parent_block_hash 为 None 时,用 NONE_HASH 替代。这确保了:

  1. 确定性 :若设了 PYTHONHASHSEED,同一种子产生相同 hash 链,跨进程共享 prefix cache
  2. 随机性 :若未设种子,用 os.urandom(32) 随机生成,避免 hash 碰撞攻击
  3. CBOR 警告:CBOR hash 函数依赖 Python dict 迭代顺序(受 PYTHONHASHSEED 影响),不设种子则 hash 不可复现

_CBOR_HASH_FUNCTIONS 冻结集合用于快速判断当前 hash 函数是否是 CBOR 系列。

9.5 KVCacheBlock 数据类(86-136行)

python 复制代码
@dataclass(slots=True)
class KVCacheBlock:
    block_id: int
    ref_cnt: int = 0
    _block_hash: BlockHashWithGroupId | None = None
    prev_free_block: "KVCacheBlock | None" = None
    next_free_block: "KVCacheBlock | None" = None
    is_null: bool = False

逐字段解析

  • block_id:物理 block 编号,范围 [0, num_gpu_blocks-1],是 GPU KV cache tensor 中的索引
  • ref_cnt:引用计数------有多少请求正在使用此 block。=0 时 block 在 free queue 中可被淘汰;>0 时不可淘汰
  • _block_hash:block 的 hash 键(含 group_id),仅当 block 满且被缓存时设置。下划线前缀表示通过 property 访问
  • prev_free_block/next_free_block:双向链表指针,仅由 FreeKVCacheBlockQueue 操作
  • is_null:null block 标记------用于 sliding window/chunked local attention 中"跳过"的位置,永不缓存

slots=True 避免了 __dict__ 开销------在数万 block 的场景下,每个对象节省约 100+ 字节。

block_hash property

python 复制代码
@property
def block_hash(self) -> BlockHashWithGroupId | None:
    return self._block_hash

@block_hash.setter
def block_hash(self, block_hash: BlockHashWithGroupId):
    assert self.block_hash is None, "The block already has a hash."
    self._block_hash = block_hash

setter 断言 block 只能被 hash 一次------如果已有 hash 说明逻辑错误(重复缓存)。

python 复制代码
def reset_hash(self):
    self._block_hash = None

淘汰时调用,清除 hash 使 block 可重新使用。

repr:用 block_id 替代对象引用避免递归打印链表。

9.6 FreeKVCacheBlockQueue 双向链表(138-300行)

python 复制代码
class FreeKVCacheBlockQueue:
    def __init__(self, blocks: list[KVCacheBlock]) -> None:
        self.num_free_blocks = len(blocks)
        # 初始化连续 block 间的双向链接
        for i in range(self.num_free_blocks):
            if i > 0:
                blocks[i].prev_free_block = blocks[i - 1]
            if i < self.num_free_blocks - 1:
                blocks[i].next_free_block = blocks[i + 1]
        # 哨兵头尾
        self.fake_free_list_head = KVCacheBlock(block_id=-1)
        self.fake_free_list_tail = KVCacheBlock(block_id=-1)

设计意图 :为什么不直接用 Python deque

  1. O(1) 中间删除deque 删除中间元素是 O(n)。当 block 被 touch(prefix cache 命中)需要从 free list 移除时,必须 O(1)
  2. 零额外分配 :利用 KVCacheBlock 自身的 prev/next 指针,不创建额外 Node 对象
  3. 哨兵节点fake_head/fake_tail(block_id=-1)永远不被弹出,保证每个真实 block 都有非 None 的 prev/next,减少空指针判断分支

LRU 语义:队首是最久未使用的 block(最先被淘汰),队尾是最近释放的。释放时逆序追加保证同一请求的 block 按尾部优先淘汰。

popleft()(约 210-230行):

python 复制代码
def popleft(self) -> KVCacheBlock:
    first_block = self.fake_free_list_head.next_free_block
    # 断言非空
    self.fake_free_list_head.next_free_block = first_block.next_free_block
    first_block.next_free_block.prev_free_block = self.fake_free_list_head
    first_block.prev_free_block = first_block.next_free_block = None
    self.num_free_blocks -= 1
    return first_block

标准双向链表头删:head → first → second 变为 head → second。断开 first 的两个指针防止悬空引用。

popleft_n()(约 232-255行):

批量弹出 n 个 block------遍历链表收集 n 个节点,断开与后续链表的连接。单次遍历 O(n)。

remove()(约 257-275行):

python 复制代码
def remove(self, block: KVCacheBlock) -> None:
    block.prev_free_block.next_free_block = block.next_free_block
    block.next_free_block.prev_free_block = block.prev_free_block
    block.prev_free_block = block.next_free_block = None
    self.num_free_blocks -= 1

O(1) 中间删除------这是选择自定义链表而非 deque 的核心原因。当 prefix cache 命中时,block 从 free list 中移除(因为现在有引用了)。

append()/append_n() :尾部追加,维护哨兵指针。append_n 批量操作减少循环内判断。

get_all_free_blocks():从 head 遍历到 tail 收集所有 block,主要用于测试。

9.7 Extra Keys 生成机制(302-450行)

9.7.1 need_extra_keys()
python 复制代码
def need_extra_keys(request: Request) -> bool:
    return (
        bool(request.mm_features)
        or (request.lora_request is not None)
        or (request.cache_salt is not None)
    )

设计意图:判断是否需要额外 hash 键。三种场景:

  • 多模态:相同 token ids + 不同图片 → 不同 KV cache(图片占位符相同但内容不同)
  • LoRA:相同 prompt + 不同 LoRA → 不同 KV cache
  • Cache Salt:用户指定的缓存隔离键
9.7.2 _gen_mm_extra_hash_keys()
python 复制代码
def _gen_mm_extra_hash_keys(
    request: Request, start_token_idx: int, end_token_idx: int, start_mm_idx: int
) -> tuple[list[Any], int]:

算法核心 :遍历 request 的 mm_features(按 offset 排序),找出与当前 block token 范围 [start_token_idx, end_token_idx) 有重叠的 MM 输入。

关键逻辑

  1. 快速跳过:若最后一个 MM 输入的结束位置 ≤ start_token_idx,直接返回
  2. start_mm_idx 参数实现增量遍历------从前一个 block 的 MM 索引继续,避免从头扫描
  3. 对于每个命中的 MM 输入,额外键为 (mm_identifier, offset - start_token_idx)------同一张图在不同位置的 block 产生不同 hash

为什么需要 offset 差值? 因为图片 placeholder tokens 可能跨越多个 block,同一张图在 block A(起始)和 block B(续接)的 token ids 完全相同(都是 placeholder id),但 KV cache 内容不同(因为 attention 上下文不同)。offset 差值区分了"图片开始于本 block 开头"和"图片开始于本 block 中间"。

9.7.3 _gen_lora_extra_hash_keys()
python 复制代码
def _gen_lora_extra_hash_keys(request: Request) -> list[str]:
    if not request.lora_request:
        return []
    return [request.lora_request.lora_name]

LoRA 额外键是 LoRA 名称字符串------同一 prompt 在不同 LoRA 下产生不同 KV cache。

9.7.4 _gen_prompt_embeds_extra_hash_keys()
python 复制代码
def _gen_prompt_embeds_extra_hash_keys(
    request: Request, start_token_idx: int, end_token_idx: int
) -> list[bytes]:
    if request.prompt_embeds is None:
        return []
    block_range = (start_token_idx, end_token_idx)
    embeds_hash = request._prompt_embeds_per_block_hashes.get(block_range)
    if embeds_hash is None:
        block_prompt_embeds = request.prompt_embeds[start_token_idx:end_token_idx]
        embeds_hash = hashlib.sha256(tensor_data(block_prompt_embeds)).digest()
        request._prompt_embeds_per_block_hashes[block_range] = embeds_hash
    return [embeds_hash]

设计意图:prompt embeddings 替代 token ids 时,相同 token ids 可能对应不同 embeddings。对每个 block 的 embedding 做 SHA-256 哈希并缓存在 request 上,避免重复计算。

9.7.5 generate_block_hash_extra_keys() --- 总入口
python 复制代码
def generate_block_hash_extra_keys(
    request: Request, start_token_idx: int, end_token_idx: int, start_mm_idx: int
) -> tuple[tuple[Any, ...] | None, int]:

聚合四种额外键:LoRA → MM → Cache Salt → Prompt Embeds。

Cache Salt 特殊处理

python 复制代码
cache_salt_keys: list[str] = (
    [request.cache_salt] if (start_token_idx == 0 and request.cache_salt) else []
)

仅第一个 block(start_token_idx==0)包含 salt------因为 salt 是请求级别的,不是 block 级别的,只需要在 hash 链起点注入一次就能区分整个链。

9.8 hash_block_tokens()(452-480行)

python 复制代码
def hash_block_tokens(
    hash_function: Callable[[Any], bytes],
    parent_block_hash: BlockHash | None,
    curr_block_token_ids: Sequence[int],
    extra_keys: tuple[Any, ...] | None = None,
) -> BlockHash:
    if not parent_block_hash:
        parent_block_hash = NONE_HASH
    curr_block_token_ids_tuple = tuple(curr_block_token_ids)
    return BlockHash(
        hash_function((parent_block_hash, curr_block_token_ids_tuple, extra_keys))
    )

设计意图:block hash 是链式哈希------当前 block 的 hash 依赖父 block 的 hash。这保证了 prefix 的唯一性:相同 token 序列前缀必然产生相同 hash 链,不同前缀必然不同(hash 抗碰撞性)。

三元组输入(parent_hash, token_ids_tuple, extra_keys) 作为哈希函数输入。tuple() 转换使序列可哈希。

链式哈希的重要性:Block N 的 hash = H(Block N-1 的 hash, Block N 的 tokens, extra_keys)。这意味着:

  • 若 Block 0 的 tokens 不同,所有后续 block hash 都不同
  • 若 Block K 的 tokens 相同且前缀完全一致,Block K 的 hash 相同 → 前缀缓存命中

9.9 get_request_block_hasher()(482-530行)

python 复制代码
def get_request_block_hasher(
    block_size: int,
    caching_hash_fn: Callable[[Any], bytes],
) -> Callable[[Request], list[BlockHash]]:

工厂模式 :返回一个闭包函数 request_block_hasher,捕获 block_sizecaching_hash_fn。这个闭包被设置到 Request 对象上,在 request 生命周期中增量计算 block hash。

增量计算

python 复制代码
start_token_idx = len(request.block_hashes) * block_size

从已有 hash 的位置继续计算,只处理新增的完整 block。

MM 索引优化

python 复制代码
if start_token_idx > 0:
    curr_mm_idx = -1  # 指向最后一个 MM 输入

当已有 hash 时,新 block 只可能与最后一个 MM 输入重叠(因为生成阶段不会有新的 MM 输入),设为 -1 直接跳到末尾。

9.10 内存估算函数(532-650行)

9.10.1 _check_enough_kv_cache_memory()

检查可用内存是否足够容纳一个最大长度请求的 KV cache。不足时尝试估算可支持的最大长度并抛出详细错误信息。

9.10.2 estimate_max_model_len()
python 复制代码
def estimate_max_model_len(vllm_config, kv_cache_spec, available_memory) -> int:
    # 二分搜索
    left, right = 1, original_max_model_len
    while left <= right:
        mid = (left + right) // 2
        if fits_in_memory(mid):
            result = mid
            left = mid + 1
        else:
            right = mid - 1
    return result

关键细节 :二分搜索过程中临时修改 vllm_config.model_config.max_model_len,在 finally 块中恢复------确保副作用不泄露。fits_in_memory() 调用 max_memory_usage_bytes() 计算给定长度下的内存需求。

9.10.3 max_memory_usage_bytes()
python 复制代码
def max_memory_usage_bytes(vllm_config, kv_cache_specs) -> int:
    return sum(spec.max_memory_usage_bytes(vllm_config) for spec in kv_cache_specs)

每个 KVCacheSpec 子类实现了自己的内存计算逻辑,此处简单求和。

9.10.4 check_enough_kv_cache_memory()

_check_enough_kv_cache_memory 的封装,处理空 kv_cache_spec(attention-free 模型)。

9.11 KV Cache Group 分组策略(660-1100行)

9.11.1 create_kv_cache_group_specs()
python 复制代码
def create_kv_cache_group_specs(kv_cache_spec, grouped_layer_names) -> list[KVCacheGroupSpec]:
    for layer_names_one_group in grouped_layer_names:
        layer_specs = [kv_cache_spec[layer_name] for layer_name in layer_names_one_group]
        merged_layer_spec = layer_specs[0].merge(layer_specs)
        kv_cache_groups.append(KVCacheGroupSpec(layer_names_one_group, merged_layer_spec))

设计意图 :同一 group 内的 layers 共享 block table 和 KV cache spec。merge() 合并同一 group 内各层 spec,验证兼容性。

9.11.2 is_kv_cache_spec_uniform()

判断所有层是否使用相同类型的 KV cache spec(FullAttention 和带 sliding window 的 FullAttention 视为同类型)。通过尝试 merge() 来验证------合并失败则说明不兼容。

9.11.3 _get_kv_cache_groups_uniform_page_size()(约 880-1000行)

这是混合模型(hybrid model)的核心分组逻辑

混合模型的分组原则(注释非常详细):

  1. 物理内存 per block 必须相同:不同大小的 block 会导致内存碎片
  2. block_size 在 group 内一致:不同 group 可以不同 block_size,但同一 group 内必须相同
  3. 每层每 token 的物理内存由模型配置决定:当前假设所有层相同
  4. 每组层数相同:简化 block table 管理

算法

  1. 按 spec 类型分组:same_type_layers: dict[KVCacheSpec, list[str]]
  2. 确定组大小 group_size = min_num_layers(最小类型中层的数量)
    • 启发式:若最大类型层数 < 最小类型 × 1.5,用最大值(避免过多 padding)
  3. 交错分配层到各组(Pipeline Parallelism 优化)

PP 交错分配

python 复制代码
for i in range(num_groups):
    grouped_layers.append(layers[i::num_groups])

例如 2 个 PP stage 各有 [full.0, sw.0, sw.1][full.1, sw.2, sw.3],交错后得到 (full.0, full.1), (sw.0, sw.2), (sw.1, sw.3)------确保每个 PP stage 的每个 group 都有层,避免空 group 导致的内存浪费。

9.11.4 unify_kv_cache_spec_page_size()

当不同层有不同 page_size 时,通过增大较小 page_size 层的 block_size 来统一。前提是大 page_size 必须是小 page_size 的整数倍。

9.11.5 unify_hybrid_kv_cache_specs()

disable_hybrid_kv_cache_manager=True 时,将 SlidingWindowSpec 和 ChunkedLocalAttentionSpec 强制转为 FullAttentionSpec------放弃内存优化(不丢弃窗口外的 KV cache),但简化管理。

9.12 KV Cache 配置生成(1100-1500行)

9.12.1 get_kv_cache_groups()

决策树

  1. 若禁用 hybrid manager → 统一 spec
  2. 若 attention-free → 返回空列表
  3. 若 spec 完全一致 → 单组
  4. 若类型一致但 hidden_size 不同 → UniformTypeKVCacheSpecs
  5. 否则 → 统一 page_size 后按混合模型分组
9.12.2 get_kv_cache_config_from_groups()

内存布局

单组情况 :每层独立 tensor,大小 = page_size_bytes * num_blocks

多组情况group_size 个共享 tensor,每个由各 group 的一层共享:

复制代码
Tensor 0: full.0, sw.0, sw.1 共享(size = available_memory // group_size)
Tensor 1: full.1, sw.2, sw.3 共享

不同 group 使用不同 block table,同一 tensor 的不同区域由不同 group 的 block table 索引。

9.12.3 get_kv_cache_configs() --- 跨 Worker 配置生成

Pipeline Parallelism 处理

  1. 合并所有 worker 的 KV cache spec
  2. 基于全局层比例生成分组
  3. 投射到每个 worker 的层子集(_project_kv_cache_groups_to_worker
  4. 自动适配 max_model_len(若设为 -1)
  5. 对每个 worker 生成配置
  6. 关键:将所有 worker 的 num_blocks 统一为最小值,并按比例缩减 tensor 大小------确保不同 PP stage 的 block table 一致
9.12.4 _auto_fit_max_model_len()

original_max_model_len == -1 时,二分搜索找到所有 worker 都能容纳的最大长度。遍历每个 worker 的投影组,取最小值。

9.13 BlockHashListWithBlockSize(约 1620-1693行)

python 复制代码
class BlockHashListWithBlockSize:
    def __init__(self, block_hashes, hash_block_size, target_block_size):
        self.scale_factor = target_block_size // hash_block_size

设计意图 :当不同 KV cache group 有不同 block_size 时,hash 以最小 block_size(hash_block_size)计算,但某些 group 需要以更大 block_size 访问 hash。此类提供惰性转换------将 scale_factor 个小 hash 拼接为一个大 hash。

例如:hash_block_size=16,target_block_size=32,scale_factor=2。

  • 原始 hash:[A(0-15), B(16-31), C(32-47), D(48-63)]
  • 转换后:[AB(0-31), CD(32-63)]
python 复制代码
def _get_value_at(self, idx: int) -> BlockHash:
    base = idx * self.scale_factor
    merged_hash: bytes = self.block_hashes[base]
    for i in range(base + 1, base + self.scale_factor):
        merged_hash += self.block_hashes[i]
    return BlockHash(merged_hash)

拼接 bytes 而非重新哈希------保证前缀匹配的确定性。

python 复制代码
BlockHashList = list[BlockHash] | BlockHashListWithBlockSize

类型别名,允许代码统一处理普通 hash 列表和需要缩放的列表。

9.14 辅助函数

  • may_override_num_blocks() :若设了 num_gpu_blocks_override,强制覆盖自动计算的 block 数
  • get_num_blocks()available_memory // page_size // num_layers 计算可用 block 数
  • get_uniform_page_size():断言所有层 page_size 一致并返回
  • get_max_concurrency_for_kv_cache_config():估算最大并发请求数
  • generate_scheduler_kv_cache_config():为调度器生成简化版配置(用任意一个 worker 的配置,去除 UniformTypeKVCacheSpecs 包装)
  • _report_kv_cache_config():日志输出 KV cache 大小和最大并发数

第十章:Block Pool 逐行解析 --- block_pool.py

10.1 模块总览

block_pool.py(509行)实现了 KV cache block 的物理内存池管理,核心职责:

  • 分配:从 free queue 获取 block
  • 释放:归还 block 到 free queue
  • 缓存查找:通过 hash 查找已缓存的 block(prefix cache hit)
  • 缓存写入:将满 block 的 hash 元数据写入缓存
  • 淘汰:释放被缓存但无引用的 block
  • 重置:清空所有 prefix cache

10.2 BlockHashToBlockMap(14-90行)

python 复制代码
class BlockHashToBlockMap:
    def __init__(self):
        self._cache: dict[
            BlockHashWithGroupId, KVCacheBlock | dict[int, KVCacheBlock]
        ] = {}

Union 类型设计------核心优化:

  • 常见情况 :一个 hash 只映射到一个 block → 直接存 KVCacheBlock 对象
  • 罕见情况 :多个 block 有相同 hash → 存 dict[int, KVCacheBlock]

为什么不总是用 dict? 因为大多数 hash 只对应一个 block,dict 创建有额外内存开销(每个 dict ~200字节 + hash table),Union 类型避免了 99% 情况下的 dict 分配,降低 GC 压力。

NOTE #1:不去重------即使已存在相同 hash 的 block,也保留两个。原因:block table 是 append-only 的,改变已分配的 block_id 会破坏 block table 的一致性。

get_one_block()

python 复制代码
def get_one_block(self, key: BlockHashWithGroupId) -> KVCacheBlock | None:
    blocks = self._cache.get(key)
    if blocks is not None:
        if isinstance(blocks, KVCacheBlock):
            return blocks
        if isinstance(blocks, dict):
            return next(iter(blocks.values()))

返回任意一个匹配 block------prefix cache 命中时不需要特定 block_id,任何内容相同的都行。

insert()

python 复制代码
def insert(self, key: BlockHashWithGroupId, block: KVCacheBlock) -> None:
    blocks = self._cache.get(key)
    if blocks is None:
        self._cache[key] = block  # 单 block
    elif isinstance(blocks, KVCacheBlock):
        self._cache[key] = {blocks.block_id: blocks, block.block_id: block}  # 升级为 dict
    elif isinstance(blocks, dict):
        blocks[block.block_id] = block  # 追加到 dict

三级渐进:None → 单对象 → dict。升级时创建新 dict 而非原地修改,因为替换 dict 引用是原子的。

pop()

python 复制代码
def pop(self, key: BlockHashWithGroupId, block_id: int) -> KVCacheBlock | None:
    blocks = self._cache.pop(key, None)
    ...
    if isinstance(blocks, KVCacheBlock):
        if blocks.block_id == block_id:
            return blocks
        self._cache[key] = blocks  # 放回
        return None
    if isinstance(blocks, dict):
        block = blocks.pop(block_id, None)
        if len(blocks) > 0:
            self._cache[key] = blocks  # dict 还有条目,放回
        return block  # dict 为空则不放回(自动清理)

按 block_id 精确弹出------淘汰时必须淘汰特定 block。若 dict 只剩一个条目,不降级为单对象(简化实现,TODO 注释表明未来可能优化)。

10.3 BlockPool 类(92-509行)

10.3.1 构造函数
python 复制代码
def __init__(self, num_gpu_blocks, enable_caching, hash_block_size,
             enable_kv_cache_events=False, metrics_collector=None):
    self.blocks: list[KVCacheBlock] = [KVCacheBlock(idx) for idx in range(num_gpu_blocks)]
    self.free_block_queue = FreeKVCacheBlockQueue(self.blocks)
    self.cached_block_hash_to_block: BlockHashToBlockMap = BlockHashToBlockMap()
    self.null_block = self.free_block_queue.popleft()
    self.null_block.is_null = True

初始化流程

  1. 创建 num_gpu_blocks 个 KVCacheBlock 对象
  2. 将所有 block 放入 free queue
  3. 创建 hash 缓存映射
  4. 弹出 block_id=0 作为 null_block------这是全局唯一的占位 block,用于 sliding window 等场景中表示"此处无有效 KV cache"。null_block 永不被释放、永不被缓存。
10.3.2 get_cached_block()
python 复制代码
def get_cached_block(self, block_hash, kv_cache_group_ids) -> list[KVCacheBlock] | None:
    cached_blocks = []
    for group_id in kv_cache_group_ids:
        block_hash_with_group_id = make_block_hash_with_group_id(block_hash, group_id)
        block = self.cached_block_hash_to_block.get_one_block(block_hash_with_group_id)
        if not block:
            return None
        cached_blocks.append(block)
    return cached_blocks

设计意图:一个逻辑 block hash 对应多个 KV cache group(混合模型),必须所有 group 都有缓存才算命中。任何一个 group 缺失则整个 cache miss。

10.3.3 cache_full_blocks()

这是 prefix caching 的核心写入逻辑:

python 复制代码
def cache_full_blocks(self, request, blocks, num_cached_blocks, num_full_blocks, block_size, kv_cache_group_id):

参数

  • blocks:请求的所有 block
  • num_cached_blocks:已缓存的 block 数(跳过这些)
  • num_full_blocks:满 block 总数(缓存到这里)

block_size vs hash_block_size 处理

python 复制代码
if block_size == self.hash_block_size:
    block_hashes: BlockHashList = request.block_hashes
else:
    block_hashes = BlockHashListWithBlockSize(request.block_hashes, self.hash_block_size, block_size)

当 block_size > hash_block_size 时,使用 BlockHashListWithBlockSize 惰性缩放。

缓存逻辑:遍历新增的满 block,跳过 null block,设置 block_hash 并插入缓存映射。

KV Events 支持 :当 enable_kv_cache_events=True 时,生成 BlockStored 事件记录:

  • block_hashes 列表
  • 父 block hash
  • token_ids 范围
  • LoRA 信息
  • extra_keys 列表(每个 block 可能不同)
10.3.4 get_new_blocks()
python 复制代码
def get_new_blocks(self, num_blocks: int) -> list[KVCacheBlock]:
    ret = self.free_block_queue.popleft_n(num_blocks)
    if self.enable_caching:
        for block in ret:
            self._maybe_evict_cached_block(block)  # 淘汰旧缓存
            block.ref_cnt += 1
    else:
        for block in ret:
            block.ref_cnt += 1
    return ret

关键:分配时若 block 有缓存 hash,必须先淘汰(清除 hash 映射)再分配。因为从 free queue 取出的 block 可能仍保留之前的 hash------它被释放回 free queue 时 hash 未被清除(因为其他请求可能还在引用同 hash 的 block)。

10.3.5 _maybe_evict_cached_block()
python 复制代码
def _maybe_evict_cached_block(self, block: KVCacheBlock) -> bool:
    block_hash = block.block_hash
    if block_hash is None:
        return False
    if self.cached_block_hash_to_block.pop(block_hash, block.block_id) is None:
        return False
    block.reset_hash()
    ...
    return True

三步淘汰:

  1. 检查 block 是否有 hash
  2. 从缓存映射中弹出(精确匹配 block_id)
  3. 重置 block 的 hash

为什么 block 可能在 free queue 但仍有 hash? 因为 block 释放时(ref_cnt 降为 0)不立即清除 hash------同 hash 的其他 block 仍在缓存中,保留映射可以让未来请求命中。

10.3.6 touch()
python 复制代码
def touch(self, blocks: Sequence[KVCacheBlock]) -> None:
    for block in blocks:
        if block.ref_cnt == 0 and not block.is_null:
            self.free_block_queue.remove(block)  # O(1) 中间删除
        block.ref_cnt += 1

设计意图:prefix cache 命中时"触摸"已有 block------增加引用计数,从 free queue 移除(防止被淘汰)。这是自定义链表而非 deque 的直接受益场景。

10.3.7 free_blocks()
python 复制代码
def free_blocks(self, ordered_blocks: Iterable[KVCacheBlock]) -> None:
    blocks_list = list(ordered_blocks)
    for block in blocks_list:
        block.ref_cnt -= 1
    self.free_block_queue.append_n(
        [block for block in blocks_list if block.ref_cnt == 0 and not block.is_null]
    )

两遍处理:先减少所有引用计数,再收集 ref_cnt==0 的 block 追加到 free queue。两遍是必要的------一个 block 可能有多个引用,只在最后一个引用释放时才归还。

不重置 hash:释放后 hash 保留在缓存映射中,允许未来请求复用。

10.3.8 evict_blocks()
python 复制代码
def evict_blocks(self, block_ids: set[int]) -> None:
    for block_id in block_ids:
        block = self.blocks[block_id]
        self._maybe_evict_cached_block(block)

外部触发的淘汰------由 KV connector 报告需要淘汰的 block ID。只清除 hash 映射,不释放 block(ref_cnt > 0 的 block 不能释放)。

10.3.9 reset_prefix_cache()
python 复制代码
def reset_prefix_cache(self) -> bool:
    num_used_blocks = self.num_gpu_blocks - self.get_num_free_blocks()
    if num_used_blocks != 1:  # 只有 null_block 在用
        return False
    self.cached_block_hash_to_block = BlockHashToBlockMap()
    for block in self.blocks:
        block.reset_hash()
    return True

安全检查:只有所有请求都完成(只剩 null_block 在用)时才允许重置。RLHF 场景下权重更新后调用,确保旧 prefix cache 失效。

10.3.10 get_usage()
python 复制代码
def get_usage(self) -> float:
    total_gpu_blocks = self.num_gpu_blocks - 1  # 减去 null_block
    return 1.0 - (self.get_num_free_blocks() / total_gpu_blocks)

null_block 永远被占用,从总量中排除。

10.3.11 take_events()

原子性取出并清空 KV events 队列------事件在 cache_full_blocks 和 _maybe_evict_cached_block 中累积,在这里一次性消费。


第十一章:KV Cache Manager 体系逐行解析

11.1 kv_cache_manager.py(552行)--- 顶层管理器

11.1.1 KVCacheBlocks 数据类(15-100行)
python 复制代码
@dataclass
class KVCacheBlocks:
    blocks: tuple[Sequence[KVCacheBlock], ...]

设计意图:KVCacheManager 与 Scheduler 之间的接口数据类型,隐藏内部实现细节。

维度选择blocks[i][j] → 第 i 个 KV cache group 的第 j 个 block。外层是 group 而非 block------因为未来可能不同 group 有不同 block 数量。

空 block 优化

python 复制代码
def new_empty(self) -> "KVCacheBlocks":
    return KVCacheBlocks(tuple(() for _ in range(len(self.blocks))))

每个 group 一个空元组。KVCacheManager 预构造了 empty_kv_cache_blocks 避免重复创建------GC 压力优化。

get_block_ids()

python 复制代码
def get_block_ids(self, allow_none=False) -> tuple[list[int], ...] | None:
    if allow_none and all(len(group) == 0 for group in self.blocks):
        return None
    return tuple([blk.block_id for blk in group] for group in self.blocks)

allow_none=True 时,若所有 group 为空则返回 None(避免传输空 block table)。

get_unhashed_block_ids_all_groups()

python 复制代码
return [
    [block.block_id for block in group if block.block_hash is None and not block.is_null]
    for group in self.blocks
]

获取所有未缓存 block 的 ID------这些 block 的 GPU 内存需要清零,防止残留数据污染。

**add()**:

python 复制代码
def __add__(self, other):
    return KVCacheBlocks(
        tuple(list(itertools.chain(blk1, blk2)) for blk1, blk2 in zip(self.blocks, other.blocks))
    )

合并两组 blocks(compute blocks + new blocks)。

11.1.2 KVCacheManager 类(约 100-552行)

构造函数

python 复制代码
class KVCacheManager:
    def __init__(self, kv_cache_config, max_model_len, hash_block_size,
                 enable_caching=True, use_eagle=False, log_stats=False, ...):
        self.coordinator = get_kv_cache_coordinator(...)
        self.block_pool = self.coordinator.block_pool
        self.empty_kv_cache_blocks = KVCacheBlocks(
            tuple(() for _ in range(self.num_kv_cache_groups))
        )

核心设计:KVCacheManager 是外观模式(Facade)------将实际工作委托给 KVCacheCoordinator,自身只处理调度层面的逻辑。

空 blocks 预分配empty_kv_cache_blocks 用嵌套元组确保不可变性,避免 GC 开销。

11.1.3 get_computed_blocks()
python 复制代码
def get_computed_blocks(self, request: Request) -> tuple[KVCacheBlocks, int]:
    if not self.enable_caching or request.skip_reading_prefix_cache:
        return self.empty_kv_cache_blocks, 0

    max_cache_hit_length = request.num_tokens - 1
    computed_blocks, num_new_computed_tokens = (
        self.coordinator.find_longest_cache_hit(request.block_hashes, max_cache_hit_length)
    )
    ...
    return self.create_kv_cache_blocks(computed_blocks), num_new_computed_tokens

num_tokens - 1 限制:必须至少重算最后一个 token 以获得 logits------即使全部命中,也不能跳过最终 token 的计算。

skip_reading_prefix_cache:需要 prompt logprobs 的请求或 pooling 模型请求跳过 prefix cache------因为需要完整计算所有 token。

11.1.4 can_fit_full_sequence()

准入控制------检查是否有足够 block 容纳请求的完整序列。防止 chunked prefill 只检查第一个 chunk 就准入,后续 chunk 无 block 可用。

11.1.5 allocate_slots()(约 180-400行)--- 最复杂的方法

Block 布局图(代码注释中的 ASCII 图):

复制代码
| < comp > | < new_comp > | < ext_comp > | < new > | < lookahead > |
  • comp:已计算的 token(本地 prefix cache 命中)
  • new_comp:新发现的 prefix cache 命中
  • ext_comp:外部 connector 缓存的 token
  • new:新 token(含未验证的 draft token)
  • lookahead:spec decode 的前瞻 token

三阶段分配

  1. 释放跳过的 block:sliding window 场景下,窗口外的 block 不再需要
  2. 处理 prefix token:touch 缓存命中的 block,分配外部计算 token 的 block
  3. 分配新 block:为待计算 token 分配新 block

关键细节

python 复制代码
num_tokens_to_cache = min(
    total_computed_tokens + num_new_tokens,
    request.num_tokens,
)

缓存上限为 request.num_tokens------排除 draft token(可能被拒绝),只缓存已确认的 token。

delay_cache_blocks:P/D(Prefill/Decode disaggregation)场景下,block 的 KV 数据将从远程加载,延迟缓存直到加载完成。

11.1.6 free()
python 复制代码
def free(self, request: Request) -> None:
    self.coordinator.free(request.request_id)

释放委托给 coordinator,coordinator 遍历所有 single_type_manager 执行释放。

11.1.7 get_num_common_prefix_blocks()

遍历请求的 block,统计 ref_cnt == 活跃请求数 的连续 block 数------这些 block 被所有请求共享,可用于 cascade attention 优化。

11.1.8 create_kv_cache_blocks()
python 复制代码
def create_kv_cache_blocks(self, blocks: tuple[list[KVCacheBlock], ...]) -> KVCacheBlocks:
    return KVCacheBlocks(blocks) if any(blocks) else self.empty_kv_cache_blocks

空 block 复用预构造对象------每次调度可能产生大量空结果,避免 GC。

11.2 single_type_kv_cache_manager.py(1133行)--- 类型特定管理器

11.2.1 SingleTypeKVCacheManager 抽象基类(约 1-400行)
python 复制代码
class SingleTypeKVCacheManager(ABC):
    def __init__(self, kv_cache_spec, block_pool, enable_caching, kv_cache_group_id,
                 dcp_world_size=1, pcp_world_size=1):
        self.block_size = kv_cache_spec.block_size
        if dcp_world_size * pcp_world_size > 1:
            self.block_size *= dcp_world_size * pcp_world_size

DCP/PCP block_size 调整:Decode Context Parallelism 和 Prefill Context Parallelism 将序列切分到多个 GPU,每个 GPU 只处理一部分 token,因此 block_size 需要乘以 world_size------一个逻辑 block 跨多个 GPU 的物理 block。

python 复制代码
self.req_to_blocks: defaultdict[str, list[KVCacheBlock]] = defaultdict(list)
self.num_cached_block: dict[str, int] = {}
  • req_to_blocks:请求 → 已分配 block 列表
  • num_cached_block:请求 → 已缓存 block 数(跟踪缓存进度,避免重复缓存)
11.2.2 get_num_blocks_to_allocate()

计算需要新分配多少 block:

python 复制代码
if request_id in self.num_cached_block:
    # Running 请求无新 cache hit
    assert len(new_computed_blocks) == 0
    return max(num_required_blocks - num_req_blocks, 0)

Running 请求快速路径:已在运行的请求不可能有新的 prefix cache hit(因为所有新 token 都是生成 token),只需检查是否需要更多 block(spec decode 可能导致 block 缩减)。

新请求计算

python 复制代码
num_new_blocks = max(
    num_required_blocks - max(num_skipped_blocks, num_local_computed_blocks),
    0,
)
num_evictable_blocks = self._get_num_evictable_blocks(new_computed_blocks[num_skipped_new_computed_blocks:])
return num_new_blocks + num_evictable_blocks

evictable block 计数:cache hit 的 block 如果 ref_cnt==0(在 free queue 中),被 touch 时会从 free queue 移除,等效于"占用"了一个 free block 的位置。因此需要纳入分配计算。

11.2.3 allocate_new_computed_blocks()

三步处理新计算的 block:

  1. Touch:增加引用计数,从 free queue 移除
  2. Skip padding:跳过的位置用 null_block 填充
  3. 添加 computed blocks:追加到请求的 block 列表
  4. External tokens:为外部计算 token 分配新 block
python 复制代码
self.num_cached_block[request_id] = len(req_blocks)

标记所有当前 block 为已缓存------这些 block 的 hash 已经在 prefix cache 中,cache_blocks() 会跳过它们。

11.2.4 allocate_new_blocks()
python 复制代码
num_required_blocks = cdiv(num_tokens, self.block_size)
num_new_blocks = num_required_blocks - len(req_blocks)
if num_new_blocks <= 0:
    return []
new_blocks = self.block_pool.get_new_blocks(num_new_blocks)
req_blocks.extend(new_blocks)
if type(self.kv_cache_spec) in (FullAttentionSpec, TQFullAttentionSpec):
    self.new_block_ids.extend(b.block_id for b in new_blocks)

new_block_ids 追踪:只对 FullAttention/TQFullAttention 类型的 spec 记录新 block ID------这些 block 的 GPU 内存需要清零。

11.2.5 cache_blocks()
python 复制代码
def cache_blocks(self, request, num_tokens):
    num_cached_blocks = self.num_cached_block.get(request.request_id, 0)
    num_full_blocks = num_tokens // self.block_size
    if num_cached_blocks >= num_full_blocks:
        return
    self.block_pool.cache_full_blocks(
        request=request,
        blocks=self.req_to_blocks[request.request_id],
        num_cached_blocks=num_cached_blocks,
        num_full_blocks=num_full_blocks,
        ...
    )
    self.num_cached_block[request.request_id] = num_full_blocks

增量缓存------只缓存从 num_cached_blocksnum_full_blocks 的新满 block。

11.2.6 free()
python 复制代码
def free(self, request_id):
    req_blocks = self.req_to_blocks.pop(request_id, [])
    ordered_blocks = reversed(req_blocks)  # 逆序:尾部优先淘汰
    self.block_pool.free_blocks(ordered_blocks)
    self.num_cached_block.pop(request_id, None)

逆序释放:尾部 block 是最近计算的,淘汰价值最低(LRU 语义下最不希望保留的)。

11.2.7 find_longest_cache_hit() --- 抽象方法

每个子类必须实现自己的 cache hit 查找逻辑,因为不同注意力类型有截然不同的缓存语义。

11.2.8 remove_skipped_blocks()
python 复制代码
def remove_skipped_blocks(self, request_id, total_computed_tokens):
    num_skipped_tokens = self.get_num_skipped_tokens(total_computed_tokens)
    if num_skipped_tokens <= 0:
        return
    blocks = self.req_to_blocks[request_id]
    num_skipped_blocks = num_skipped_tokens // self.block_size
    num_skipped_blocks = min(num_skipped_blocks, len(blocks))
    removed_blocks = []
    for i in range(num_skipped_blocks - 1, -1, -1):
        if blocks[i] == self._null_block:
            break
        removed_blocks.append(blocks[i])
        blocks[i] = self._null_block
    self.block_pool.free_blocks(removed_blocks)

从后向前遍历:一旦遇到 null_block 就停止------之前的 block 必然已经被替换过。

提前释放:即使请求未调度,也释放窗口外的 block------减少淘汰压力。

11.2.9 FullAttentionManager(约 401-460行)

python 复制代码
class FullAttentionManager(SingleTypeKVCacheManager):
    @classmethod
    def find_longest_cache_hit(cls, block_hashes, max_length, ...) -> tuple[list[KVCacheBlock], ...]:
        computed_blocks = tuple([] for _ in range(len(kv_cache_group_ids)))
        for block_hash in itertools.islice(block_hashes, max_num_blocks):
            if cached_block := block_pool.get_cached_block(block_hash, kv_cache_group_ids):
                for computed, cached in zip(computed_blocks, cached_block):
                    computed.append(cached)
            else:
                break

算法:从左到右线性扫描 block hash 链,遇到第一个 cache miss 就停止。Full attention 的"向下封闭性"(downward-closed property)保证了这一点------如果 Block K 在缓存中,Block 0 到 K-1 也必然在缓存中(因为 hash 链式依赖)。

EAGLE 处理

python 复制代码
if use_eagle and computed_blocks[0]:
    for computed in computed_blocks:
        computed.pop()

丢弃最后一个命中 block------EAGLE(推测解码)需要最后一个 block 的隐藏状态作为 draft head 的输入,必须重算。

alignment_tokens 对齐

python 复制代码
while (block_size != alignment_tokens
       and len(computed_blocks[0]) * block_size % alignment_tokens != 0):
    for computed in computed_blocks:
        computed.pop()

混合模型中不同 group 有不同 block_size,cache hit 长度必须是 LCM 的倍数。

get_num_common_prefix_blocks()

python 复制代码
def get_num_common_prefix_blocks(self, running_request_id):
    blocks = self.req_to_blocks[running_request_id]
    for block in blocks:
        if block.ref_cnt == len(self.req_to_blocks):
            num_common_blocks += 1
        else:
            break
    return num_common_blocks

前缀共享计数:从第一个 block 开始,统计 ref_cnt 等于总请求数的连续 block 数。

11.2.10 SlidingWindowManager(约 461-600行)

python 复制代码
class SlidingWindowManager(SingleTypeKVCacheManager):
    def __init__(self, kv_cache_spec, **kwargs):
        super().__init__(kv_cache_spec, **kwargs)
        self.sliding_window = kv_cache_spec.sliding_window
find_longest_cache_hit()

算法与 FullAttention 完全不同------sliding window 不满足向下封闭性:

python 复制代码
# 从右向左扫描
for i in range(max_num_blocks - 1, -1, -1):
    if cached_block := block_pool.get_cached_block(block_hashes[i], kv_cache_group_ids):
        num_contiguous_blocks += 1
        if num_contiguous_blocks >= sliding_window_contiguous_blocks:
            # 找到足够长的连续命中
            match_found = True
            break
    else:
        num_contiguous_blocks = 0

设计原理 :Sliding window attention 只关注最近 sliding_window 个 token。要使一个 block 可复用,需要保证从它到当前位置的连续 block 都在缓存中(足够覆盖窗口)。从右向左扫描可以最早找到满足条件的组合。

初始化为 null blocks

python 复制代码
computed_blocks = tuple(
    [block_pool.null_block] * max_num_blocks
    for _ in range(len(kv_cache_group_ids))
)

所有位置先填充 null_block,只有实际命中的位置替换为真实 block------窗口外的位置保持 null。

get_num_skipped_tokens()
python 复制代码
def get_num_skipped_tokens(self, num_computed_tokens):
    return max(0, num_computed_tokens - self.sliding_window + 1)

直觉:若 sliding_window=4, num_computed_tokens=7,下一个 token 需要 tokens 4-7(窗口内的),tokens 0-3 被跳过。

11.2.11 ChunkedLocalAttentionManager(约 601-750行)

python 复制代码
class ChunkedLocalAttentionManager(SingleTypeKVCacheManager):
    def __init__(self, kv_cache_spec, **kwargs):
        super().__init__(kv_cache_spec, **kwargs)
        self.attention_chunk_size = kv_cache_spec.attention_chunk_size
find_longest_cache_hit()

分块本地注意力:attention 被切分为固定大小的 chunk,每个 token 只关注当前 chunk 内的 token。

算法

  1. 计算本地窗口起始位置:local_attention_start_idx = (max_length // chunk_size) * chunk_size
  2. 窗口之前的 block 全部标记为 null(已计算,不需要 attention)
  3. 窗口内的 block 从左到右扫描缓存
python 复制代码
local_attention_start_block_idx = local_attention_start_idx // kv_cache_spec.block_size
computed_blocks = tuple(
    [block_pool.null_block] * local_attention_start_block_idx
    for _ in range(len(kv_cache_group_ids))
)
for i in range(local_attention_start_block_idx, max_num_blocks):
    if cached_block := block_pool.get_cached_block(block_hashes[i], kv_cache_group_ids):
        ...
    else:
        break
get_num_skipped_tokens()
python 复制代码
def get_num_skipped_tokens(self, num_computed_tokens):
    return (num_computed_tokens // self.attention_chunk_size) * self.attention_chunk_size

跳过所有完整 chunk 的 token------当前 chunk 左侧的 token 不需要 attention。

11.2.12 MambaManager(约 751-950行)

python 复制代码
class MambaManager(SingleTypeKVCacheManager):
    def __init__(self, kv_cache_spec, block_pool, **kwargs):
        super().__init__(kv_cache_spec, block_pool, **kwargs)
        self.cached_blocks_this_step: set[BlockHashWithGroupId] = set()
        self.mamba_cache_mode = kv_cache_spec.mamba_cache_mode
        self.num_speculative_blocks = kv_cache_spec.num_speculative_blocks
        if self.mamba_cache_mode == "align":
            self.last_state_block_idx: dict[str, int] = {}
            self._allocated_block_reqs: set[str] = set()

Mamba 特殊性

  • Mamba 是状态空间模型(SSM),不需要完整前缀的 KV cache,只需要最后一个状态的 SSM state
  • cached_blocks_this_step:跟踪当前 step 缓存的 block hash,防止同一 step 内的 block 被其他请求复用(Mamba 的状态计算是增量的,不能跨请求共享增量结果)
  • mamba_cache_mode"align" 模式下 block 对齐分配,"none" 模式下类似 FullAttention
find_longest_cache_hit()
python 复制代码
# 从右向左扫描,只需最后一个命中
for i in range(max_num_blocks - 1, -1, -1):
    if cached_block := block_pool.get_cached_block(...):
        # 前面插入 dummy null blocks
        computed.extend([block_pool.null_block] * i)
        computed.append(cached_block)
        break

只需最后一个 block:Mamba 的递推性质意味着只需要最后一个状态的 block,前面的都是 null。

get_num_blocks_to_allocate() --- Mamba 重写

align 模式

python 复制代码
num_required_blocks = cdiv(num_tokens, self.block_size) + self.num_speculative_blocks
if request_id in self._allocated_block_reqs:
    num_new_blocks = 1  # 旧请求:复用 speculative blocks
else:
    num_new_blocks = 1 + self.num_speculative_blocks  # 新请求

为什么只需要 1 + speculative_blocks? Mamba align 模式下,每个 step 只推进一个 block 的 SSM state,speculative blocks 用于投机解码的状态空间预留。

防重入

python 复制代码
if (len(new_computed_blocks) > 0
    and new_computed_blocks[-1].block_hash in self.cached_blocks_this_step):
    return self.block_pool.num_gpu_blocks + 1  # 强制延迟调度

同一 step 内缓存的 block 不能被同 step 的其他请求使用------返回超大的分配需求使调度器认为无足够空间,推迟到下一个 step。

allocate_new_blocks() --- Mamba align 模式

复杂的 block 复用逻辑

  1. 记录 last_state_block_idx------两个 step 前分配的 block 可以释放(SSM state 已拷贝到新 block)
  2. 填充 null blocks 到 num_skipped_blocks
  3. 复用上一步的 speculative blocks------避免重复分配
  4. 只分配 1 个新 block 作为当前 step 的运行状态
get_num_skipped_tokens()
python 复制代码
def get_num_skipped_tokens(self, num_computed_tokens):
    return num_computed_tokens - 1

Mamba 只保留最后一个 token 的状态,跳过所有之前的 token。

remove_skipped_blocks() --- Mamba 重写
python 复制代码
def remove_skipped_blocks(self, request_id, num_computed_tokens):
    # 异步调度安全:减去 speculative blocks
    num_computed_tokens = max(0, num_computed_tokens - self.num_speculative_blocks)
    super().remove_skipped_blocks(request_id, num_computed_tokens)
    if self.mamba_cache_mode == "align":
        last_state_block_idx = self.last_state_block_idx.get(request_id)
        if (last_state_block_idx is not None
            and last_state_block_idx < cdiv(num_computed_tokens, self.block_size) - 1):
            # 释放旧的 SSM state block
            ...

释放两个 step 前的 SSM state block------已拷贝到新 block,不再需要。

cache_blocks() --- Mamba 重写
python 复制代码
def cache_blocks(self, request, num_tokens):
    num_cached_blocks_before = self.num_cached_block.get(request.request_id, 0)
    super().cache_blocks(request, num_tokens)
    num_cached_blocks_after = self.num_cached_block.get(request.request_id, 0)
    if num_cached_blocks_after > num_cached_blocks_before:
        for block in self.req_to_blocks[request.request_id][num_cached_blocks_before:num_cached_blocks_after]:
            if block.is_null:
                continue
            self.cached_blocks_this_step.add(block.block_hash)

将新缓存的 block hash 加入 cached_blocks_this_step------防止同 step 内跨请求复用。

new_step_starts()
python 复制代码
def new_step_starts(self) -> None:
    self.cached_blocks_this_step.clear()

每个新 step 清空------上一步缓存的 block 现在可以被新请求安全复用。

11.2.13 CrossAttentionManager(约 951-1000行)

python 复制代码
class CrossAttentionManager(SingleTypeKVCacheManager):
    def allocate_new_computed_blocks(self, ...):
        assert len(new_computed_blocks) == 0  # 不共享 prefix cache

    def cache_blocks(self, request, num_tokens):
        raise ValueError("Should not be called as prefix caching is disabled.")

    def find_longest_cache_hit(cls, ...):
        raise NotImplementedError("CrossAttentionManager does not support caching")

设计意图:Cross-attention(编码器-解码器模型如 Whisper)的 KV cache 是请求特定的------编码器状态(音频/图像)在不同请求间不共享。因此:

  • 不参与 prefix cache
  • 不缓存 block
  • find_longest_cache_hit 抛出异常

11.2.14 SinkFullAttentionManager(约 1001-1015行)

python 复制代码
class SinkFullAttentionManager(FullAttentionManager):
    def __init__(self, kv_cache_spec, block_pool, enable_caching, kv_cache_group_id, ...):
        super().__init__(...)
        sink_len = kv_cache_spec.sink_len
        num_sink_block = sink_len // self.block_size
        self.sink_blocks = self.block_pool.free_block_queue.popleft_n(num_sink_block)

Sink Attention(如 StreamingLLM):保留序列开头的"sink" token + 最近窗口。在初始化时预分配 sink blocks------这些 block 永不被释放,始终保留序列开头的 KV cache。

11.2.15 工厂函数

python 复制代码
spec_manager_map: dict[type[KVCacheSpec], type[SingleTypeKVCacheManager]] = {
    FullAttentionSpec: FullAttentionManager,
    TQFullAttentionSpec: FullAttentionManager,
    MLAAttentionSpec: FullAttentionManager,
    SlidingWindowSpec: SlidingWindowManager,
    ChunkedLocalAttentionSpec: ChunkedLocalAttentionManager,
    MambaSpec: MambaManager,
    CrossAttentionSpec: CrossAttentionManager,
    SinkFullAttentionSpec: SinkFullAttentionManager,
}

def get_manager_for_kv_cache_spec(kv_cache_spec, **kwargs):
    manager_class = spec_manager_map[type(kv_cache_spec)]
    return manager_class(kv_cache_spec, **kwargs)

策略模式:根据 KVCacheSpec 类型选择对应的管理器。TQFullAttention 和 MLAAttention 复用 FullAttentionManager------它们的 block 管理逻辑与标准 FullAttention 相同。

11.3 kv_cache_coordinator.py(591行)--- 协调器

11.3.1 KVCacheCoordinator 抽象基类(1-200行)
python 复制代码
class KVCacheCoordinator(ABC):
    def __init__(self, kv_cache_config, max_model_len, use_eagle, enable_caching, ...):
        self.block_pool = BlockPool(...)
        self.single_type_managers = tuple(
            get_manager_for_kv_cache_spec(
                kv_cache_spec=kv_cache_group.kv_cache_spec,
                block_pool=self.block_pool,
                enable_caching=enable_caching,
                kv_cache_group_id=i,
                ...
            )
            for i, kv_cache_group in enumerate(self.kv_cache_config.kv_cache_groups)
        )

设计意图:Coordinator 持有一个 BlockPool 和一组 SingleTypeKVCacheManager,将操作分发到每个 manager。

所有 manager 共享同一个 BlockPool------这是关键设计。不同 KV cache group 的 block 来自同一个物理内存池,确保 block ID 全局唯一且内存不浪费。

核心方法
  • get_num_blocks_to_allocate():遍历所有 manager,对 CrossAttention 特殊处理(用 encoder token 数而非 decoder token 数),累加所需 block 数
  • allocate_new_computed_blocks():分发到每个 manager
  • allocate_new_blocks():分发到每个 manager,返回各组新 block 的元组
  • cache_blocks():分发到每个 manager
  • free():分发到每个 manager
  • remove_skipped_blocks():分发到每个 manager
  • get_num_common_prefix_blocks():收集每个 manager 的结果
  • new_step_starts():分发到每个 manager
11.3.2 KVCacheCoordinatorNoPrefixCache(约 200-240行)
python 复制代码
class KVCacheCoordinatorNoPrefixCache(KVCacheCoordinator):
    def find_longest_cache_hit(self, block_hashes, max_cache_hit_length):
        blocks = tuple([] for _ in range(self.num_single_type_manager))
        return blocks, 0

禁用 prefix cache 时的空实现:始终返回空命中。支持任意数量的 KV cache group(包括 0)。

11.3.3 UnitaryKVCacheCoordinator(约 242-290行)
python 复制代码
class UnitaryKVCacheCoordinator(KVCacheCoordinator):
    def find_longest_cache_hit(self, block_hashes, max_cache_hit_length):
        hit_blocks = self.single_type_managers[0].find_longest_cache_hit(
            block_hashes=block_hashes,
            max_length=max_cache_hit_length,
            kv_cache_group_ids=[0],
            ...
        )
        return hit_blocks, len(hit_blocks[0]) * self.block_size

单一 KV cache group:直接委托给唯一的 manager,计算命中 token 数 = 命中 block 数 × block_size。

11.3.4 HybridKVCacheCoordinator(约 292-530行)--- 最复杂的协调器
python 复制代码
class HybridKVCacheCoordinator(KVCacheCoordinator):
    def verify_and_split_kv_cache_groups(self):
        # 按 spec 类型分组
        attention_groups = []
        for i, g in enumerate(self.kv_cache_config.kv_cache_groups):
            manager_cls = self.single_type_managers[i].__class__
            spec = g.kv_cache_spec
            # 查找相同 spec 的已有组
            for existing_spec, group_ids, existing_cls in attention_groups:
                if existing_spec == spec:
                    group_ids.append(i)
                    break
            else:
                attention_groups.append((spec, [i], manager_cls))

        # Full attention 排最前------其左到右扫描提供更紧的初始边界
        self.attention_groups = sorted(
            attention_groups,
            key=lambda x: not isinstance(x[0], FullAttentionSpec),
        )

        # LCM block size
        self.lcm_block_size = lcm(*block_sizes)

分组逻辑:相同 spec 的 group 合并处理------它们共享 find_longest_cache_hit 逻辑。

Full attention 优先排序:Full attention 的 cache hit 是向下封闭的(从左到右连续),先扫描它得到一个较紧的上界,减少后续非 full attention 类型的扫描工作量。

LCM block size:混合模型中不同 group 的 block_size 可能不同,cache hit 长度必须是所有 block_size 的最小公倍数,确保每个 group 都能对齐到 block 边界。

find_longest_cache_hit() --- 不动点迭代算法
python 复制代码
def find_longest_cache_hit(self, block_hashes, max_cache_hit_length):
    hit_length = max_cache_hit_length
    while True:
        curr_hit_length = hit_length
        for spec, group_ids, manager_cls in self.attention_groups:
            if is_full_attn and cached_blocks is not None:
                # Full attention:直接截断到新长度
                num_blocks = curr_hit_length // spec.block_size
                curr_hit_length = num_blocks * spec.block_size
            else:
                hit_blocks = manager_cls.find_longest_cache_hit(...)
                curr_hit_length = len(hit_blocks[0]) * spec.block_size
                for group_id, blocks in zip(group_ids, hit_blocks):
                    hit_blocks_by_group[group_id] = blocks

        if curr_hit_length >= hit_length:
            break  # 收敛
        hit_length = curr_hit_length
        if is_simple_hybrid:
            break  # 简单混合:一次迭代
    ...

不动点算法原理

  1. 从最大可能长度开始
  2. 每个注意力类型可能将长度缩短(因为它的 cache hit 不够长)
  3. 如果任何类型缩短了长度,重启所有类型的检查(因为更短的长度可能影响其他类型的判断)
  4. 长度单调递减且有下界 0,必然收敛

Full attention 优化:利用向下封闭性,第一次计算后,后续迭代只需截断 block 列表,无需重新扫描。

Simple hybrid 优化:1 full + 1 other 的组合只需要一次迭代------full attention 先确定上界,other 类型可能进一步缩短,但缩短后不会再影响 full attention(向下封闭)。

11.3.5 get_kv_cache_coordinator() 工厂函数

python 复制代码
def get_kv_cache_coordinator(...) -> KVCacheCoordinator:
    if not enable_caching:
        return KVCacheCoordinatorNoPrefixCache(...)
    if len(kv_cache_config.kv_cache_groups) == 1:
        return UnitaryKVCacheCoordinator(...)
    return HybridKVCacheCoordinator(...)

三级选择:

  1. 无缓存 → NoPrefixCache
  2. 单组 → Unitary
  3. 多组 → Hybrid

第十二章:辅助模块逐行解析

12.1 encoder_cache_manager.py(381行)

12.1.1 EncoderCacheManager

核心设计:管理多模态编码器输出(如视觉 embedding)的缓存,实现跨请求共享。

粒度:按单个多模态输入项(由 mm_hash 标识)缓存,而非按 token。

数据结构

python 复制代码
self.cached: dict[str, set[str]] = {}  # mm_hash → 引用该数据的请求 ID 集合
self.freeable: OrderedDict[str, int] = OrderedDict()  # mm_hash → num_encoder_embeds(可回收)
self.freed: list[str] = []  # 已淘汰的 mm_hash 列表(等待 worker 清理)

三层状态

  1. 活跃cached[hash] 非空(有请求引用)→ 不可回收
  2. 可回收cached[hash] 为空且在 freeable 中 → 可淘汰释放空间
  3. 已淘汰 :在 freed 列表中 → 等待 worker 清理物理内存
check_and_update_cache()
python 复制代码
def check_and_update_cache(self, request, input_id) -> bool:
    mm_hash = request.mm_features[input_id].identifier
    if mm_hash not in self.cached:
        return False
    if not self.cached[mm_hash]:  # 缓存但无引用
        num_encoder_embeds = self.freeable.pop(mm_hash)
        self.num_freeable_slots -= num_encoder_embeds
    self.cached[mm_hash].add(request.request_id)
    return True

状态转换 :可回收 → 活跃。从 freeable 中移除,减少 num_freeable_slots

can_allocate()
python 复制代码
def can_allocate(self, request, input_id, encoder_compute_budget, num_embeds_to_schedule) -> bool:
    num_embeds = request.get_num_encoder_embeds(input_id)
    if num_embeds > encoder_compute_budget:
        return False
    num_embeds += num_embeds_to_schedule
    if num_embeds <= self.num_free_slots:
        return True
    if num_embeds > self.num_freeable_slots:
        return False
    # 淘汰可回收条目
    while num_embeds > self.num_free_slots:
        mm_hash, num_free_embeds = self.freeable.popitem(last=False)  # FIFO
        del self.cached[mm_hash]
        self.freed.append(mm_hash)
        self.num_free_slots += num_free_embeds
    return True

淘汰策略 :FIFO(OrderedDict 的 popitem(last=False) 取最早插入的)。最早缓存的 embedding 最不可能被未来请求使用。

注意 :淘汰只更新管理器状态,物理内存清理由 worker 在收到 freed 列表后执行。

allocate()
python 复制代码
def allocate(self, request, input_id):
    mm_hash = request.mm_features[input_id].identifier
    if mm_hash not in self.cached:
        self.cached[mm_hash] = set()
    num_encoder_embeds = request.get_num_encoder_embeds(input_id)
    self.cached[mm_hash].add(request_id)
    self.num_free_slots -= num_encoder_embeds
    self.num_freeable_slots -= num_encoder_embeds

同时减少 free 和 freeable:新分配的空间既不可用也不可回收。

free_encoder_input()
python 复制代码
def free_encoder_input(self, request, input_id):
    mm_hash = request.mm_features[input_id].identifier
    if not self.cached.get(mm_hash, None):
        return
    self.cached[mm_hash].discard(req_id)
    if not self.cached[mm_hash]:  # 无引用了
        num_encoder_embeds = request.get_num_encoder_embeds(input_id)
        self.freeable[mm_hash] = num_encoder_embeds
        self.num_freeable_slots += num_encoder_embeds

状态转换:活跃 → 可回收。不立即释放,等需要空间时再淘汰(延迟释放优化)。

compute_mm_encoder_budget()
python 复制代码
def compute_mm_encoder_budget(scheduler_config, mm_max_toks_per_item):
    max_tokens_per_mm_item = max(mm_max_toks_per_item.values())
    encoder_compute_budget = max(
        scheduler_config.max_num_encoder_input_tokens, max_tokens_per_mm_item)
    encoder_cache_size = max(
        scheduler_config.encoder_cache_size, max_tokens_per_mm_item)
    return encoder_compute_budget, encoder_cache_size

确保计算预算和缓存空间至少能容纳最大的单个多模态输入项。

12.1.2 EncoderDecoderCacheManager

编码器-解码器模型的简化缓存管理器,不实现跨请求共享,仅跟踪空间分配。

python 复制代码
class EncoderDecoderCacheManager(EncoderCacheManager):
    def get_freed_mm_hashes(self):
        to_free = self.to_free
        self.to_free = self.allocated
        self.allocated = []
        return to_free

双缓冲allocated 收集当前 step 分配的条目,to_free 是上一步分配的条目(现在可以释放)。模拟 EncoderCacheManager 的延迟释放语义------模型执行后才释放物理内存。

12.2 kv_cache_metrics.py(96行)

12.2.1 BlockMetricsState
python 复制代码
class BlockMetricsState:
    def __init__(self):
        now_ns = time.monotonic_ns()
        self.birth_time_ns = now_ns
        self.last_access_ns = now_ns
        self.access_history: deque[int] = deque(maxlen=4)

跟踪单个 block 的生命周期指标:

  • birth_time_ns:分配时间
  • last_access_ns:最后访问时间
  • access_history:最近 4 次访问时间(有界防内存增长)

方法

  • get_lifetime_seconds():当前时间 - 分配时间
  • get_idle_time_seconds():当前时间 - 最后访问时间
  • get_reuse_gaps_seconds():相邻访问间的时间间隔列表
12.2.2 KVCacheMetricsCollector
python 复制代码
class KVCacheMetricsCollector:
    def __init__(self, sample_rate: float = 0.01):
        self.sample_rate = sample_rate
        self.block_metrics: dict[int, BlockMetricsState] = {}
        self._eviction_events: list[KVCacheEvictionEvent] = []

采样机制:默认 1% 采样率------不是每个 block 都跟踪指标,而是随机采样。在数万 block 的场景下,1% 采样足以获得统计显著的指标,同时将内存和 CPU 开销降低两个数量级。

python 复制代码
def should_sample_block(self) -> bool:
    return random.random() < self.sample_rate

事件流

  1. on_block_allocated():采样命中时创建 BlockMetricsState
  2. on_block_accessed():更新 last_access_ns 和 access_history
  3. on_block_evicted():记录淘汰事件(生命周期、空闲时间、重用间隔)

drain_events():原子性取出并清空淘汰事件列表。

12.3 sched/output.py(261行)--- 调度器输出数据结构

12.3.1 NewRequestData
python 复制代码
@dataclass
class NewRequestData:
    req_id: str
    prompt_token_ids: list[int] | None
    mm_features: list[MultiModalFeatureSpec]
    sampling_params: SamplingParams | None
    pooling_params: PoolingParams | None
    block_ids: tuple[list[int], ...]
    num_computed_tokens: int
    lora_request: LoRARequest | None
    prompt_embeds: torch.Tensor | None = None
    prefill_token_ids: list[int] | None = None  # v2 model runner only

设计意图:首次调度时发送完整请求数据到 worker。Worker 缓存这些数据,后续步骤只发送增量。

from_request() 工厂方法:从 Request 对象提取必要字段,避免手动构造。

anon_repr():脱敏版本------用长度替代实际 token ids,用于日志记录避免泄露 prompt 内容。

12.3.2 CachedRequestData
python 复制代码
@dataclass
class CachedRequestData:
    req_ids: list[str]
    resumed_req_ids: set[str]  # 被抢占后恢复的请求
    new_token_ids: list[list[int]]  # PP only
    all_token_ids: dict[str, list[int]]  # connector 传播
    new_block_ids: list[tuple[list[int], ...] | None]
    num_computed_tokens: list[int]
    num_output_tokens: list[int]

设计意图:已缓存请求的增量更新。Worker 已有这些请求的基础数据,只需发送变化部分。

resumed_req_ids:被抢占(preemption)后重新调度的请求------其 block_ids 需要替换而非追加(因为抢占时 block 被释放了)。

is_context_phase()

python 复制代码
def is_context_phase(self, req_id: str) -> bool:
    return num_output_tokens is not None and num_output_tokens == 0

判断请求是否仍在 prefill 阶段(尚未产出任何 decode token)。

12.3.3 SchedulerOutput
python 复制代码
@dataclass
class SchedulerOutput:
    scheduled_new_reqs: list[NewRequestData]
    scheduled_cached_reqs: CachedRequestData
    num_scheduled_tokens: dict[str, int]
    total_num_scheduled_tokens: int
    scheduled_spec_decode_tokens: dict[str, list[int]]
    scheduled_encoder_inputs: dict[str, list[int]]
    num_common_prefix_blocks: list[int]
    finished_req_ids: set[str]
    free_encoder_mm_hashes: list[str]
    preempted_req_ids: set[str] | None = None
    has_structured_output_requests: bool = False
    pending_structured_output_tokens: bool = False
    num_invalid_spec_tokens: dict[str, int] | None = None
    kv_connector_metadata: KVConnectorMetadata | None = None
    ec_connector_metadata: ECConnectorMetadata | None = None
    new_block_ids_to_zero: list[int] | None = None

字段解析

  • num_scheduled_tokens:每个请求本次调度的 token 数------worker 用此确定 prefill 长度
  • scheduled_spec_decode_tokens:推测解码的 draft token ids------按请求 ID 索引
  • scheduled_encoder_inputs:需要编码器处理的输入索引------如 [0, 1] 表示第 0 和第 1 张图片需要编码
  • num_common_prefix_blocks:每组 KV cache 的共享前缀 block 数------cascade attention 优化
  • finished_req_ids:上一步完成的请求------worker 需要清理其缓存状态
  • free_encoder_mm_hashes:需要释放的编码器缓存条目
  • new_block_ids_to_zero:新分配 block 的 ID 列表------worker 在使用前清零,防止残留 NaN/脏数据

make_empty() 工厂方法:创建空输出,用于无请求需要调度时。

12.3.4 GrammarOutput
python 复制代码
@dataclass
class GrammarOutput:
    structured_output_request_ids: list[str]
    grammar_bitmask: npt.NDArray[np.int32]

结构化输出的语法位掩码------用于约束解码(JSON schema、regex 等)。

12.4 sched/interface.py(243行)--- 调度器接口

12.4.1 PauseState
python 复制代码
class PauseState(enum.IntEnum):
    UNPAUSED = 0
    PAUSED_NEW = 1   # 不调度新请求,只调度已 running 的
    PAUSED_ALL = 2   # 不调度任何请求

设计意图:DP attention(Data Parallel attention)场景下,需要暂停调度以同步不同 DP rank 的状态。

12.4.2 SchedulerInterface 抽象基类

定义调度器的完整接口契约:

核心方法

  • schedule()SchedulerOutput:每次调度迭代调用,决定哪些请求处理多少 token
  • update_from_output():接收 model runner 输出,更新调度器状态
  • add_request():添加新请求
  • finish_requests():完成/中止请求

辅助方法

  • get_grammar_bitmask():获取结构化输出的语法位掩码
  • update_draft_token_ids():更新 draft token ids 并验证语法
  • has_finished_requests():检查是否有待清理的完成请求
  • pause_state/set_pause_state:暂停/恢复调度
  • reset_prefix_cache():重置前缀缓存
  • reset_encoder_cache():重置编码器缓存
  • get_request_counts():返回 (running, waiting) 请求数
  • make_stats():生成统计信息
  • shutdown():关闭调度器
  • get_kv_connector():获取 KV connector(可选)

12.5 sched/request_queue.py(208行)--- 请求队列

12.5.1 SchedulingPolicy
python 复制代码
class SchedulingPolicy(Enum):
    FCFS = "fcfs"      # 先来先服务
    PRIORITY = "priority"  # 优先级调度
12.5.2 RequestQueue 抽象基类

定义队列接口:add_request, pop_request, peek_request, prepend_request, prepend_requests, remove_request, remove_requests, bool , len , iter

12.5.3 FCFSRequestQueue
python 复制代码
class FCFSRequestQueue(deque[Request], RequestQueue):

继承 deque :直接继承 Python 的 collections.deque,获得 O(1) 的左端弹出/右端追加。

方法实现

  • add_request()append():右端追加
  • pop_request()popleft():左端弹出
  • peek_request()[0]:查看左端
  • prepend_request()appendleft():左端插入(抢占请求重新入队)
  • remove_requests():过滤后重建 deque(deque 不支持原地过滤)
12.5.4 PriorityRequestQueue
python 复制代码
class PriorityRequestQueue(RequestQueue):
    def __init__(self):
        self._heap: list[Request] = []

基于堆 :使用 heapq 实现优先级队列。Request 类的 __lt__ 方法按 (priority, arrival_time) 排序。

关键差异

  • prepend_request()add_request():优先级队列无"前端"概念,按优先级插入
  • remove_request()heap.remove() + heapify():O(n) 删除 + O(n) 重堆化
  • __iter__():复制堆后逐个弹出------不破坏原始堆
12.5.5 create_request_queue()
python 复制代码
def create_request_queue(policy: SchedulingPolicy) -> RequestQueue:
    if policy == SchedulingPolicy.PRIORITY:
        return PriorityRequestQueue()
    elif policy == SchedulingPolicy.FCFS:
        return FCFSRequestQueue()

简单工厂。

12.6 sched/utils.py(130行)--- 调度工具函数

12.6.1 重复检测
python 复制代码
def _has_repeating_pattern(token_ids, pattern_len, repetition_min_count) -> bool:
    for n in range(1, pattern_len + 1):
        target_token = token_ids[-n]
        for m in range(1, repetition_min_count):
            if token_ids[-(pattern_len * m + n)] != target_token:
                return False
    return True

算法 :检查 token_ids 末尾是否存在长度为 pattern_len、重复 repetition_min_count 次的模式。对于模式中的每个位置 n,检查前 repetition_min_count - 1 次重复是否与最后一次匹配。

python 复制代码
def check_sequence_repetition(token_ids, params) -> bool:
    for pattern_len in range(min_pattern_size, max_pattern_size + 1):
        if pattern_len * min_count > len(token_ids):
            return False
        if _has_repeating_pattern(token_ids, pattern_len, min_count):
            return True
    return False

遍历所有可能模式长度,检查是否存在重复。

12.6.2 remove_all()
python 复制代码
def remove_all(lst, items_to_remove) -> list:
    if len(items_to_remove) == 1:
        item = next(iter(items_to_remove))
        with contextlib.suppress(ValueError):
            lst.remove(item)
        return lst
    return [item for item in lst if item not in items_to_remove]

优化 :单个元素时用 list.remove()(原地,O(n)),多个元素时用列表推导(创建新列表)。单个元素是最常见情况(完成一个请求时)。

12.6.3 check_stop()
python 复制代码
def check_stop(request, max_model_len) -> bool:
    if request.num_output_tokens < sampling_params.min_tokens:
        return False
    # 检查 EOS
    if last_token_id == sampling_params.eos_token_id:
        request.status = RequestStatus.FINISHED_STOPPED
        return True
    # 检查 stop token ids
    if last_token_id in (sampling_params.stop_token_ids or ()):
        ...
    # 检查长度上限
    if request.num_tokens >= max_model_len or request.num_output_tokens >= request.max_tokens:
        request.status = RequestStatus.FINISHED_LENGTH_CAPPED
        return True
    # 检查重复模式
    if repetition_detection is not None and check_sequence_repetition(...):
        request.status = RequestStatus.FINISHED_REPETITION
        return True
    return False

四级停止条件

  1. 最小 token 数:低于 min_tokens 不停止
  2. EOS / stop tokens:显式停止信号
  3. 长度上限:达到 max_model_len 或 max_tokens
  4. 重复检测:检测到无限循环模式

12.7 sched/async_scheduler.py(60行)--- 异步调度器

python 复制代码
class AsyncScheduler(Scheduler):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._spec_token_placeholders: list[int] = [-1] * self.num_spec_tokens

设计意图:异步调度的核心差异------在调度时就为 spec decode 预分配 placeholder,而非等到 model runner 输出后。

_update_after_schedule()
python 复制代码
def _update_after_schedule(self, scheduler_output):
    super()._update_after_schedule(scheduler_output)
    for req_id in scheduler_output.num_scheduled_tokens:
        request = self.requests[req_id]
        if request.is_prefill_chunk:
            continue
        # 异步场景:每个 decode step 产出 1 + num_spec_tokens 个 token
        request.num_output_placeholders += 1 + cur_num_spec_tokens
        request.spec_token_ids = self._spec_token_placeholders

异步调度语义 :调度时不知道实际的 spec token ids,用 placeholder [-1, -1, ...] 代替。Worker 进程会填充实际值。

num_output_placeholders :跟踪已调度但尚未确认的 token 数------在 _update_request_with_output 中逐步消耗。

_update_request_with_output()
python 复制代码
def _update_request_with_output(self, request, new_token_ids):
    if request.discard_latest_async_tokens:
        # 强制抢占时丢弃最新异步 token
        request.discard_latest_async_tokens = False
        return [], False

    status_before_update = request.status
    new_token_ids, stopped = super()._update_request_with_output(request, new_token_ids)

    request.num_output_placeholders -= len(new_token_ids)

    # 缓存新 token(排除 placeholder)
    if status_before_update == RequestStatus.RUNNING:
        self.kv_cache_manager.cache_blocks(
            request, request.num_computed_tokens - request.num_output_placeholders
        )
    return new_token_ids, stopped

缓存边界计算num_computed_tokens - num_output_placeholders------排除尚未确认的 placeholder token,只缓存已确认的有效 token。

discard_latest_async_tokensreset_prefix_cache 强制抢占时设置------丢弃最近的异步 token,因为它们可能基于过时的 prefix cache。


架构总结

整体设计模式

  1. 分层委托:KVCacheManager → KVCacheCoordinator → SingleTypeKVCacheManager → BlockPool,每层只关注自己的职责
  2. 策略模式spec_manager_map 将 KVCacheSpec 类型映射到对应管理器,支持灵活扩展
  3. 外观模式:KVCacheManager 是 Scheduler 的简洁接口,隐藏 coordinator/block pool 的复杂性
  4. 共享资源池:所有 SingleTypeKVCacheManager 共享一个 BlockPool,避免内存碎片和重复管理
  5. 增量计算:block hash 增量计算、cache_blocks 增量缓存、CachedRequestData 增量更新

关键性能优化

  1. FreeKVCacheBlockQueue:自定义双向链表替代 deque,实现 O(1) 中间删除
  2. BlockHashToBlockMap Union 类型:避免 99% 场景下的 dict 分配
  3. empty_kv_cache_blocks 复用:避免每次调度创建空对象
  4. 采样指标:1% 采样率大幅降低指标收集开销
  5. 惰性 BlockHashListWithBlockSize:hash 缩放延迟到访问时
  6. 增量 MM 索引:start_mm_idx 参数避免每次从头扫描 mm_features

混合模型支持

HybridKVCacheCoordinator 的不动点算法是处理多注意力类型 prefix cache 的关键创新------通过迭代收敛找到所有类型都满足的最长公共前缀,同时利用 Full attention 的向下封闭性优化迭代效率。

相关推荐
踩着两条虫1 小时前
VTJ.PRO 企业级应用开发实战指南
前端·人工智能·低代码·重构·架构
青槿吖1 小时前
告别RestTemplate!Feign让微服务调用像点外卖一样简单
java·开发语言·分布式·spring cloud·微服务·云原生·架构
一碗白开水一2 小时前
【技术探索】解码Mamba:从SSM到革命性序列建模架构的前世今生
架构
黑金IT2 小时前
通过“套壳”架构打造工业级 AI 视频生成流水线
人工智能·架构·ai视频
努力成为一个程序猿.2 小时前
Flink运行时架构
大数据·架构·flink
搬砖的前端2 小时前
本地模型+TRAE CN 打造最优模型组合实测:开源主模型+本地辅模型,对标GPT5.2/5.3/Gemini-3-Flash
前端·ai·mac·ai编程·qwen·trae·qwen3.6
懂AI的老郑2 小时前
人工智能手机的构建思路:从架构到实现
人工智能·智能手机·架构
DavidSoCool2 小时前
Dify使用ChatFlow实现调用数据库问答
数据库·ai·知识库·dify
ofoxcoding2 小时前
GPT-5.4 API 怎么低延迟调用?2026 年 5 种接入方案实测对比
python·gpt·ai·flask