Token Attention：细粒度 KV 管理 vs Page 级别

谁该读这一篇？ 想搞清 SGLang 如何在 token-level 与 paged KV 分配之间切换的工程师；要做 attention backend / memory pool 改造的开发者。 前置阅读： 01-radix-attention.md、00-prerequisites.md §4。 耗时： 30 分钟 学完能： 1. 把 KV cache 的物理布局画出来（slot / page / block）； 2. 对比 token-level（SGLang）和 page-level（vLLM）的代价/收益矩阵； 3. 解释常见 page_size=1 路径与大页 backend 的取舍； 4. 看到 KV 显存使用情况，能判断碎片是不是问题； 5. 在 memory_pool.py 和 allocator/ 里找到对应的实现。

1. 物理布局回顾

KV cache 不是按"请求"组织的，是按"slot / page"组织的——所有请求的 KV 都挤在同一块 GPU 显存里。

GPU 显存 (KV pool):
  ┌────────────────────────────────────────────┐
  │ slot 0 │ slot 1 │ slot 2 │ ... │ slot N    │
  └────────────────────────────────────────────┘
    K[N, num_layers, num_kv_heads, head_dim]
    V[N, num_layers, num_kv_heads, head_dim]

每个 slot 容纳"1 个 token 的 K/V"（page_size=1），或 "page_size 个 token 的 K/V"（page_size>1）。 每个请求有一张"逻辑 token id → 物理 slot id"的索引表（在 SGLang 是 req_to_token 矩阵）。

2. 两种粒度

3. page_size=1 的代价

3.1 索引表更大

seq_len = 4096:
  page=16: 索引表 256 项
  page=1:  索引表 4096 项

每个请求都要把 4096 个 slot id 存下来。但这是 int32 张量、4096 × 4 byte = 16 KB，可以忽略。

3.2 Attention kernel 必须支持非连续

经典 attention 算 $\text{softmax}(QK^T) V$ 期望 K、V 是连续矩阵。 token-level 时 K、V 是按 slot id 散落的，必须用支持 indirection（gather）的 kernel。 FlashInfer 提供这类 kernel：

• BatchPrefillWithPagedKVCacheWrapper —— 按 page table gather K/V。
• BatchDecodeWithPagedKVCacheWrapper
• MultiLevelCascadeAttentionWrapper —— 处理 RadixTree 的共享 prefix。

SGLang 会根据平台、模型与显式 backend 解析 attention backend 和 page size；常见 CUDA 配置可走 FlashInfer，但 Blackwell、ROCm、NPU 和特定模型可能解析到其他 backend 或强制 16 / 64 / 128 等页大小。 因此不要把“默认永远是 FlashInfer + page=1”当成跨平台契约，应以 /server_info 和启动日志为准。

3.3 一些 kernel 的内部最小粒度

FlashInfer 内部 tile size 通常是 16 或 32（warp 大小相关）。 即使外面 page_size=1，内部还是会以 16 token 为单位访问—— SGLang 在分配 slot 时把同请求的 slot 尽量连续（reduce 实际 indirection），所以性能没掉太多。

4. SGLang 也支持 page_size > 1

sglang serve MODEL_PATH --page-size 16

什么时候开：

• 极长 context（128k+）：page=16 让索引表小一截。
• 特定 attention kernel 要求 page_size 对齐（一些 quantization 后端）。
• KV cache 主要装"几大请求"而不是"很多小请求"（碎片不显著）。

代价：

• 命中粒度变 16，可能丢一些 token 级匹配机会。
• RadixCache 内部要"page 对齐"，分裂操作变粗。

不要只凭 workload 名称选页大小；先接受当前 backend 的解析结果，仅在命中粒度、索引开销或 kernel 约束明确时覆盖，并对 TTFT、TPOT、吞吐与可用 KV tokens 做 A/B benchmark。

5. MHATokenToKVPool 类（1804+）

源码：memory_pool.py:1804。

class MHATokenToKVPool(KVCache):
    """Token-level KV pool for multi-head attention."""
    def __init__(self, size, dtype, head_num, head_dim, layer_num, device, ...):
        # 显存一次性预分配：
        self.k_buffer = [torch.empty((size, head_num, head_dim), dtype=dtype)
                          for _ in range(layer_num)]
        self.v_buffer = [torch.empty((size, head_num, head_dim), dtype=dtype)
                          for _ in range(layer_num)]
        ...

    def set_kv_buffer(self, layer_id, loc, cache_k, cache_v):
        """把新算的 K/V 写入指定 slot."""
        self.k_buffer[layer_id][loc] = cache_k
        self.v_buffer[layer_id][loc] = cache_v

关键设计：

• k_buffer / v_buffer 是 List[Tensor]，每层一个张量。
• 每个张量第 0 维是 slot id。
• "插入"操作就是按 slot id 索引写入。

变体：

• MLA（DeepSeek 系）：MLA 压缩了 KV，用 MLATokenToKVPool。
• FP4 / FP8 KV：用 MHATokenToKVPoolFP4 等子类。
• Hybrid linear（如 Mamba 混合架构）：用 HybridLinearKVPool。
• DSA（DeepSeek Sparse Attention）：用 DSATokenToKVPool。

每种变体只是写入/读取格式不同，外部 RadixCache 接口完全一致。

6. TokenToKVPoolAllocator：slot 分配（28+）

源码：allocator/token.py:28。抽象接口在 allocator/base.py:27。

职责：维护空闲 slot / page id 张量并完成分配、释放。

class TokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
    def __init__(self, size, dtype, device, kvcache):
        self.size = size
        self.free_pages = ...           # 自由 slot 列表

    def alloc(self, need_size: int) -> torch.Tensor:
        """从 free pool 里拿 need_size 个 slot id."""
        ...

    def free(self, free_index: torch.Tensor):
        """把 slot id 还回 free pool."""
        ...

两种关键使用场景：

• extend / prefill：一次需要一批新 token slot。
• decode：每个运行请求通常再需要一个 token slot。

TokenToKVPoolAllocator.alloc() 直接从 free_pages 张量切出任意 slot id； PagedTokenToKVPoolAllocator 才有 page-aware 的 alloc_extend / alloc_decode Triton kernel。

7. 空闲表与分配失败

请求频繁来去会让空闲 id 顺序变化：

free slot 分布: 100, 101, 105, 109, 200-250, 400, ...

token allocator 不要求这 64 个 slot 物理连续：只要 available_size() 足够，单纯的外部“不连续” 不会导致分配失败。启用排序时，释放项先进入 release_pages，需要时由 merge_and_sort_free() 合并排序；attention 通过索引间接访问。

真正的分配压力来自可用 token / page 总量不足，或某个 page-aware / backend 约束不满足。 当前没有 fragmentation_rate Prometheus 指标；联合观察 sglang:kv_available_tokens、 sglang:kv_evictable_tokens、sglang:kv_used_tokens、retraction / eviction 日志和 TPOT。 不要把定期 flush_cache 或重启当 allocator 正常运行的必要条件。

8. PagedAttention vs Token Attention 选型

实际差距取决于 backend、模型、序列长度和复用边界，必须实测。

9. 怎么观察这一切

curl http://localhost:30000/metrics | grep -E "kv_pool|memory_pool|cache_hit"

会看到：

sglang:kv_used_tokens{...}           # 活跃 KV slot 数
sglang:max_total_num_tokens{...}     # 配置后的设备 KV token 上限
sglang:kv_available_tokens{...}      # 当前可用 KV token
sglang:cache_hit_rate{...}           # 命中率

kv_used_tokens / max_total_num_tokens 可作为占用率近似；具体健康区间取决于 workload。 若占用持续接近上限且 kv_evictable_tokens 很少，说明 KV pool 偏紧或活动请求锁定了大量页。

10. 小结

• SGLang 的 page size 由平台、模型与 backend 解析；常见路径可为 1，也有 16 / 64 / 128 等强制值。
• Token-level：命中粒度精确，但索引表大、对 attention kernel 要求高。
• FlashInfer 的 paged + cascade kernel 是 SGLang 选 token-level 的底气。
• Token allocator 用 GPU tensor free list 支持非连续 slot；page-aware extend 分配等批量路径使用 Triton kernel。
• 实际性能取决于 workload，前缀复用 + 分支多则 SGLang 优势明显。

11. 自检

1. page=1 与 page=16 的代价 / 收益矩阵？

2. 为什么 token-level 需要 FlashInfer cascade kernel？

3. free slot 不连续时会不会出现“总量够但分配失败”？怎么观察？

4. MLA 模型在 SGLang 里用哪个 KV pool 子类？为什么？

5. 什么时候会使用 alloc_extend_kernel？

12. 下一步

• 03-code-walkthrough/05-radix-tree.md — 上面 RadixCache 的实现细节。
• 03-code-walkthrough/06-attention-backends.md — FlashInfer 怎么吃 token-level KV。
• 04-optimizations/01-flashinfer.md — Cascade Inference 原理。
• 源码：memory_pool.py、allocator/。

上游源码：sglang/python/sglang/srt/mem_cache/。