ReqToTokenPool、KV 索引表与释放生命周期#

这章解决什么问题#

前面的调度与内存章节已经解释了 batch 怎样成形、allocator 怎样分配页、prefix cache 怎样复用，但还有一层非常关键的“账本逻辑”没有单独展开：请求在逻辑上是一串 token，在物理上是一组 KV 槽位，那么系统究竟用什么结构把这两者连起来，又在什么时候把这层映射回收掉？

这正是 ReqToTokenPool 在解决的问题。

如果不把这一层单独讲清楚，读者会看到很多 allocator、page、cache node 和 req_pool_idx，但很难建立一个稳定判断框架：某个请求当前到底占了哪些物理位置，这些位置在 extend、decode、speculative over-allocation、session 复用和释放时又会发生什么。

为什么 ReqToTokenPool 值得单独成章#

TokenToKVPoolAllocator 负责“哪里还有空闲 KV 槽位”，RadixCache 负责“哪些前缀已经被缓存”，但这两者都不能直接回答“某个请求的第 N 个 token 现在映射到哪一个物理槽位”。这层桥正是 ReqToTokenPool：

  • 行：请求槽位，也就是 req_pool_idx
  • 列：逻辑 token 位置
  • 值：物理 KV index

这不是一个可有可无的辅助表，而是调度层、cache 层和 attention backend 之间共同依赖的中介结构。forward_batch.req_to_token_pool.req_to_token[...] 之所以会在不同 backend 里反复出现，就是因为执行层最终仍然要靠这张表把“序列位置”翻译成“真正的 KV 位置”。

一张图：请求槽位怎样变成物理 KV 索引#

这张图解决的理解障碍是：很多读者会把 ReqToTokenPool 误以为只是“请求 id 到缓存对象”的普通映射，但它其实是一张二维索引表。

flowchart LR
    Req["Req / req_pool_idx"] --> Table["ReqToTokenPool.req_to_token[row, col]"]
    Table --> KV["physical KV indices"]
    KV --> Pool["TokenToKVPoolAllocator / KV pool"]
    KV --> Cache["RadixCache / SessionAwareCache"]
    Table --> Attn["attention backend reads page table"]

图比纯文字多解释的一点是：ReqToTokenPool 既服务分配，也服务执行；它不是只在 scheduler 里短暂存在。

ReqToTokenPool 的数据形状是什么#

python/sglang/srt/mem_cache/memory_pool.py 里，ReqToTokenPool 的核心状态非常直接：

  • req_to_token: 一个形状为 (size, max_context_len) 的 torch.int32 张量
  • free_slots: 可复用的请求槽位列表

这说明它不是 dictionary，而是一张预先分配好的稠密表。这样设计的收益很直接：

  • scheduler 和 backend 都可以按 tensor 方式批量访问
  • page-table 类操作可以直接落在 GPU 侧张量上
  • req_pool_idx 变成稳定、便宜的行索引

代价也很明确：你要为 max_running_requests * max_context_len 预留一整张表，所以 ReqToTokenPool 不是零成本元数据，而是运行时拓扑的一部分。

req_pool_idx 是怎样拿到的#

在 python/sglang/srt/mem_cache/common.py 里，alloc_req_slots(...) 会先调用 req_to_token_pool.alloc(reqs)。而 ReqToTokenPool.alloc(...) 并不是无条件给每个请求分配新槽位，它会先识别哪些请求可以复用已有槽位：

  • 已经有 req_pool_idx
  • 并且是 chunked prefill 的继续段，或者已经有 committed KV

这说明 req_pool_idx 的生命周期并不总是“一次请求一生只分一次、最后一次性释放”。在 chunked prefill、session 继续写入、部分 decode 续跑这些路径里，请求槽位可以跨 batch 复用。

这一点很重要，因为它解释了为什么 ReqToTokenPool 更像“请求的物理身份”，而不是“某轮 batch 的临时变量”。

extend 阶段怎样写这张表#

alloc_for_extend(...) 的工作顺序可以压成四步：

  1. 先通过 alloc_req_slots(...) 确定每个请求的 req_pool_idx
  2. 再让 allocator 为新的 extend token 分配物理槽位
  3. 把 prefix 部分已有的 prefix_indices 和新分配出来的 out_cache_loc 拼起来
  4. 通过 write_cache_indices(...) 写回 req_to_token_pool

这里最关键的不是“会写”，而是“写的时候会同时保留 prefix 和新增 extend”。这说明 ReqToTokenPool 不是只保存最新一段，而是在逻辑上维护整条序列到物理位置的映射。

decode 阶段为什么只追加一列#

alloc_for_decode(...) 走的是另一条更短的路径：

  • 先从 req_to_token_pool.req_to_token[req_pool_indices, seq_lens - 1] 取出每个请求当前最后一个 token 的物理位置
  • 让 paged allocator 基于 last_loc 决定 decode token 落在哪
  • 再把新位置写到 (req_pool_indices, locs) 这一列

这说明 decode 阶段并不会重写整行，而是沿着“当前序列长度”继续向后追加。这也是 attention backend 能在 decode 路径里快速拿到 page table 的原因：整条历史索引已经在这张表里累积好了。

prefix cache 与这张表怎样互相咬合#

RadixCache、SessionAwareCache、LMCache offload 这些路径里，系统经常会从 req_to_token_pool.req_to_token[...] 读出某一段 kv_indices，再决定：

  • 哪一段前缀已经 committed
  • 哪一段可以插入 cache tree
  • 哪一段 over-allocated token 需要回收
  • 某个 session slot 是否接管了这组 KV

这说明 prefix cache 并不自己维护一套完全独立的“请求到物理 token”索引。它依然要通过 ReqToTokenPool 这层账本去拿真实 index。

释放为什么不是一句 free(req) 就结束#

release_kv_cache(req, tree_cache, is_insert=True) 把释放路径分成了几层：

  1. 先调用 tree_cache.cache_finished_req(req, is_insert=is_insert)，决定是否把已完成前缀留进 cache
  2. 处理 speculative decoding 之类路径留下的 over-allocated KV
  3. 在需要时调用 tree_cache.token_to_kv_pool_allocator.free(indices_to_free)
  4. 最后才 tree_cache.req_to_token_pool.free(req)

这条顺序说明一件事：请求结束不等于它占用的所有 KV 都应当立刻消失。系统会先决定哪些 KV 要转移给 cache，哪些只是临时多分了必须回收，最后才回收请求槽位本身。

如果把释放理解成单纯的“请求跑完 -> free”，你就会看不懂：

  • 为什么 prefix cache 能在请求结束后继续复用 KV
  • 为什么 session 路径会让 req_pool_idx 的所有权转移
  • 为什么 speculative 路径还要额外处理 over-allocation

为什么 SessionAwareCache 会让释放语义更复杂#

common.py 里有一句非常值得技术书明确写出来的注释：SessionAwareCache.cache_finished_req 可能会把 req_pool_idx 的所有权转移给 session slot，于是后续 cleanup 会被跳过。

这意味着一个常见误判需要被主动纠正：

  • 误判：请求结束后，req_pool_idx 和所有 KV 一定马上回收
  • 实际：在 session-aware 路径里，请求的物理身份可能被更长寿命的 session slot 接管

这就是为什么“请求生命周期结束”与“物理索引生命周期结束”不是一回事。

执行层怎样消费这张表#

从 hardware_backend/*/attention/*、forward_batch_info.py 和多种 speculative worker 可以看到，执行层会频繁读取：

  • forward_batch.req_to_token_pool.req_to_token
  • forward_batch.token_to_kv_pool

也就是说，ReqToTokenPool 不是 scheduler 专用内部状态，而是 attention page-table、sparse attention、speculative gather 和 cache migration 都会读取的底层材料。

换句话说，调度层写进去的不是“管理信息”，而是之后每一轮 forward 真正会消费的执行证据。

这一层最容易出现的误解#

1. 以为 req_pool_idx 就是请求 id#

不是。它更像运行时给请求分配的物理行号，用来索引 ReqToTokenPool 这张表。

2. 以为 allocator 一分配完，映射关系就结束了#

真正稳定保存“这个请求有哪些物理 KV index”的仍然是 ReqToTokenPool。

3. 以为请求结束后所有物理槽位都会立即释放#

prefix cache、session-aware cache 和 speculative over-allocation 都会改变释放顺序。

如果你怀疑 KV index 写坏了，先看哪里#

建议按这个顺序查：

  1. memory_pool.py::ReqToTokenPool.alloc/write/free
  2. common.py::alloc_for_extend(...)
  3. common.py::alloc_for_decode(...)
  4. common.py::release_kv_cache(...)
  5. 对应 cache 实现里如何读取 req_to_token_pool.req_to_token[...]
  6. attention backend 是否按预期消费 req_to_token

如果现象是“decode 阶段 page table 错位”，重点看 alloc_for_decode(...) 怎样拿 last_loc；如果现象是“请求结束后显存没回落”，重点看 release_kv_cache(...) 里 prefix cache 和 session-aware 路径有没有接管 KV 所有权。

小结#

这一章真正要补齐的，是调度与内存部分里最底层但最稳定的一条账本主线：

  • ReqToTokenPool 把逻辑 token 位置映射到物理 KV index
  • extend 和 decode 都会写这张表，只是写法不同
  • prefix cache、session cache 和执行 backend 都会消费这张表
  • 请求结束时，系统先处理 KV 所有权和 over-allocation，再释放请求槽位

把这一层读清楚之后，你再回头看 allocator、cache 和 attention backend，很多“为什么这里又在读 req_to_token”的问题就会自然收敛。