vLLM kv cache 管理

题目描述

实现 vLLM 的 PagedAttention 显存管理：把 KV Cache 切成固定大小的页（block），按需分配、跨请求共享、按 LRU 回收。

传统实现按 max_model_len 给每个请求预留连续显存，实际序列往往远短于上限，浪费极大。vLLM 借用操作系统的分页思想：物理页固定装 block_size 个 token 的 KV，请求持有一张页表（block table）把逻辑位置映射到物理页号，于是

内部碎片被限制在每个请求的最后一页之内，而不是按上限预留。

本题只做管理这一半（分配 / 共享 / 回收 / 前缀缓存），不涉及 attention 计算。

你要实现的类

• KVCacheBlock —— 一个物理页，带 ref_cnt 与 block_hash
• FreeKVCacheBlockQueue —— 空闲页队列，队头先被分配 = 队头最该被逐出
• BlockPool —— 页池 + prefix cache 哈希表
• Request —— 一个请求的 token 与已算进度
• KVCacheManager —— 调度器视角的 get_computed_blocks / allocate_slots / cache_blocks / free

签名与字段以右侧模板为准，不要改动，评测按这套接口调用。

五个关键规则

1. 只有写满的页才能进 prefix cache。半满的页内容还会变，登记进去会让后来者读到脏数据。
2. hash 必须串成链。块 hash 由 (父块 hash, 本块 token) 计算，表达的是"相同前缀"而不只是"相同内容"——[1,1,1,1][9,9,9,9] 与 [2,2,2,2][9,9,9,9] 的第二块内容一样，但前缀不同，不允许命中。
3. 前缀命中是零拷贝共享。命中的页 ref_cnt += 1，若它当前躺在空闲队列里要先摘出来；ref_cnt 归零才允许回收。
4. 整个请求全部命中时，必须留最后一块重算，否则没有 logits 可以采样。
5. 释放时分两档排序：没有 hash 的页插队头（永远不可能被前缀命中，优先被抢走），有 hash 的页挂队尾（尽量留给下一个同前缀请求）。页表要逆序释放 —— 越靠后的页前缀越长、越难复用，应该越早被逐出。

另外：页被重新分配时，它此前承载的前缀缓存立即失效，必须把哈希表里的对应项一起删掉。这正是 vLLM 的取舍 —— prefix cache 是机会性的，显存紧张时缓存页随时被抢走。

输入输出约定

评测脚手架会用你的类跑一串操作，每步产出一条记录：

• prefill：新请求进来，先查前缀缓存 → 设置 num_computed_tokens → allocate_slots 补页 → cache_blocks 登记。显存不足时记 {"oom": true}。
• decode：追加 1 个 token，只在跨页时才需要新页。
• free：请求结束，归还页表。
• state：快照当前空闲队列顺序、空闲页数、prefix cache 里的页、各请求页表、非零 ref_cnt。

提示

• 空闲队列用带 fake head/tail 的双向链表，是为了支持 O(1) 从链表中间摘除（前缀命中时 touch 一个仍在空闲队列里的页）。
• allocate_slots 要按 ceil((num_computed_tokens + num_new_tokens) / block_size) 算总需求，只申请与现有页表的差额；不够时返回 None，让调度器去排队或抢占。
• 对照真实代码：vllm/v1/core/kv_cache_utils.py、block_pool.py、kv_cache_manager.py。