Re-store evicted blocks to CPU to improve lookup hit rate

问题描述

一个请求有 10 个 block（1～10），如果 block 3 被 CPU LRU 淘汰了，下次 lookup 时：

lookup 机制：最长连续前缀匹配，从 block 1 开始找
block 1、2 命中 → block 3 不在 CPU → 连续命中被截断
结果：只能命中 2 个 block（512 tokens），需要从 token 512 开始 prefill

Fix 的作用：在调度时检查之前 store 过的 block 是否还在 CPU，如果被淘汰了就 re-store（GPU → CPU）。这样下次 lookup 时 block 3 又回到 CPU，连续命中不被截断，避免 prefill。

为什么淘汰的 block 可以 re-store？ 因为请求的 GPU block 还没释放（请求还在运行），KV 数据在 GPU 里，可以从 GPU 重新 copy 一份到 CPU。

技术实现

检查的时机

Scheduler 每个 step 的完整流程：

1. Lookup (get_num_new_matched_tokens)
   → 检查 CPU cache，决定有多少 token 可以命中
   → 返回 num_external_computed_tokens

2. Allocate (allocate_slots)
   → 分配 GPU block 给请求

3. Store check (build_kv_connector_meta → _prepare_eager_store_specs)  ← Fix 在这里
   → 扫描请求的 block，检查 CPU cache 是否还有
   → 如果被淘汰 → re-store
   → 返回 SimpleCPUOffloadMetadata

4. Worker 执行
   → 根据 metadata 执行 GPU → CPU 传输

检查的原理

初始化

获取 GPU block pool 引用
获取 CPU block pool 引用和剩余空间（num_free，每次分配 -1，到 0 停止）
获取 attention group 列表（Hybrid 模型有多个 group）
获取已经在传输中的 GPU block 集合（in_flight），用于去重

主循环：遍历每个 request

获取 request 的 store 状态（_reqs_to_store）。如果 request 不在里面（新请求首次调度），或者已经 finished，跳过。
如果 request 被抢占（GPU 满了，被踢回 waiting queue）：清空 block 列表和游标（num_stored_blocks 归零）。下次运行时从头开始 store。因为之前的 GPU block 已经被释放了，新分配的是不同的物理 block。
把新分配的 GPU block 追加到 state.block_ids[g]（group g 的所有 GPU block ID）。
如果这个 request 没有新 token 计算（num_scheduled_tokens == 0）：跳过。因为没有新 confirmed 的 KV 数据，Phase 1b 不会扫描到新 block。注意：这也意味着 Phase 1a（re-store 淘汰 block）不会执行——这是 fix 的一个限制。
如果 request 没有分配任何 GPU block（异常情况），跳过。

Phase 1：扫描 block，分类

Phase 1a：重新扫描已 store 的 block（0 到 already_stored_g-1）

这些是之前已经 store 过的 block，但可能被 CPU LRU 淘汰了。逐个检查：

block 是否有效：is_null → null block 没有 KV 数据，跳过
block 有 hash 吗：block_hash is None → SWA 的某些位置可能没有 hash（永远不会被 prefix cache 命中），跳过
已经在传输中吗：in in_flight → 避免重复 store
CPU cache 里还有吗：cached_block_hash_to_block.get_one_block(hash) is not None → 还在就不需要 re-store
CPU 没空间了：num_free <= 0 → 停止扫描，否则预留一个空间

如果以上检查都通过 → block 被 CPU 淘汰了，加入 re-store 列表。

Phase 1b：扫描新增的 block（already_stored_g 到 ready_blocks_g-1）

这些是”新计算完成但还没 store”的 block。检查逻辑相同：

is_null → 跳过
block_hash is None → 跳过（SWA 位置）
in in_flight 或 CPU cache 有 → 跳过
num_free <= 0 → 停止

如果通过 → 新 block，第一次 store 到 CPU。

Phase 2：批量分配 CPU block

批量分配 CPU block（cpu_pool.get_new_blocks(n)）
Stamp hash：cpu_blk._block_hash = bhash → 把 hash 注册到 CPU block 上
更新全局结果：merged_gpu_block_ids.extend(...)、in_flight.update(...)
Touch GPU block：防止传输过程中被释放

游标更新

state.num_stored_blocks[g] = max(state.num_stored_blocks[g], advanced_per_group[g])

num_stored_blocks：游标，记录之前已经处理到第几个 block
advanced_per_group：本次扫描了多少个 block
为什么用 max()：Phase 1a 扫描的是旧 block（已算在旧游标里），用 += 会过度推进。用 max() 确保游标 = “已扫描的最大位置”

代价/收益分析

指标	值	说明
单次 re-store 代价	~0.5ms + 2MB CPU	DMA 传输 + 一个 block 空间
避免 prefill 收益	~25ms/block	GPU 计算 256 tokens 的代价
收益/代价比	50x	命中时的收益是代价的 50 倍
对话 agent 期望收益	25x	命中率 > 50%（连续对话）

对话 Agent 场景

对话 agent 特征：

多轮对话：用户连续追问 10-50 轮
相同 prefix：每轮包含历史
中等并发：10-100 个同时活跃对话

在这个场景下：

re-store 的 block 会被反复访问（对话历史）
相同 prefix 的请求共享 block
CPU cache 压力相对可控

结论：re-store fix 在对话 agent 场景下非常合理，ROI > 25x。