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

谁该读这一篇？ 想搞清"为什么 SGLang 默认 page_size=1，vLLM 默认 16"的工程师；要做 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. 解释为什么 SGLang 选 page_size=1 仍能跑得快（FlashInfer cascade kernel）； 4. 看到 KV 显存使用情况，能判断碎片是不是问题； 5. 在 memory_pool.py 和 allocator.py 里找到对应的实现。

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. 两种粒度

维度	Page-level（vLLM 默认 16）	Token-level（SGLang 默认 1）
一个物理单元包多少 token	16	1
索引表项数	ceil(seq_len / 16)	seq_len
分配单元	一次分 16 token 的 page	一次分 1 token 的 slot
碎片	少（只在尾部不到 16 token）	略多（每个请求都有零头）
前缀命中精度	必须 16 倍数	任意 token
Attention kernel 适配	标准 paged kernel	FlashInfer 的 cascade 系列

3. SGLang 默认 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 默认 attention backend 就是 FlashInfer，所以 token-level 不是问题。

3.3 一些 kernel 的内部最小粒度

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

4. SGLang 也支持 page_size > 1

python -m sglang.launch_server --page-size 16

什么时候开：

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

代价：

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

默认 page=1 在大多数 chat / agent 场景最优。

5. `MHATokenToKVPool` 类（797+）

源码：memory_pool.py:797。

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 分配（121+）

源码：allocator.py:121。

职责：维护一个 "哪些 slot 空着" 的位图，分配/释放。

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."""
        ...

两种关键 alloc 路径：

alloc_extend：prefill 时一次分配新 prompt 长度个 slot（alloc_extend_kernel，241）。
alloc_decode：decode 时每步分 1 个 slot（alloc_decode_kernel，327）。

两种路径都是 Triton kernel，因为分配次数高、Python 实现会成瓶颈。

7. 碎片问题

Token-level + 高 churn（请求频繁来去）会让 free pool 变得"零碎"：

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

如果你需要分配 64 个连续 slot 给一个新 prompt，可能找不到。

SGLang 的处理：

alloc_extend_kernel 允许非连续分配——不要求 slot 连续，只要够数就行。
真正算 attention 时通过 page table 间接访问，连续不连续不影响正确性。

但有性能影响：

内存访问 locality 变差，attention kernel 的 L2 cache 命中率掉。
实际测得在重碎片场景下 TPOT 涨 5-10%。

对策：

定期 flush_cache（flush_cache.py）。
监控 free pool fragmentation rate。
严重时考虑加大 page_size（牺牲粒度换 locality）。

8. PagedAttention vs Token Attention 性能对比（粗略）

workload	vLLM (page=16)	SGLang (page=1)
单轮独立请求	baseline	-1% ~ +1%
短 prompt 高 QPS	baseline	+3% (碎片代价显现)
长 system prompt 共享	baseline	+20% (token-level 命中精确)
Agent 多分支	baseline	+50% (cascade kernel + fork)

实际跑出来差距大小取决于具体 workload。但一个稳定的判断：前缀复用 + 分支越多，SGLang token-level 优势越大。

9. 怎么观察这一切

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

会看到：

sglang:kv_pool_used_size{...}        # 占用 slot 数
sglang:kv_pool_total_size{...}       # 总 slot 数
sglang:cache_hit_rate{...}           # 命中率

used_size / total_size 是占用率，正常 50-90%。持续接近 100% 说明 KV pool 偏紧、evict 频繁。

10. 小结

SGLang 默认 token-level（page_size=1）；vLLM 默认 16。
Token-level：命中粒度精确，但索引表大、对 attention kernel 要求高。
FlashInfer 的 paged + cascade kernel 是 SGLang 选 token-level 的底气。
Allocator 用 Triton kernel 加速分配，支持非连续 slot。
实际性能取决于 workload，前缀复用 + 分支多则 SGLang 优势明显。

11. 自检

SGLang 默认 page=1 比 vLLM 默认 page=16 的代价 / 收益矩阵？

答案

**收益**：(a) 命中粒度精确到 1 token，前缀复用率高的业务能 100% 命中；(b) 显存碎片少（不需要为凑整 16 token 浪费尾部）；(c) fork 共享精确到 token 级。 **代价**：(a) 索引表项数 16×（4096 token = 4096 项 vs 256 项），但只占 KB 级显存可忽略；(b) Attention kernel 必须支持非连续 K/V indirection（FlashInfer 是这条路径的现实依赖）；(c) 内存访问局部性略差，attention kernel 实际仍按 16 token tile 访问，但 SGLang 分配时尽量让同请求 slot 连续，减少损失。实测多分支 / 高复用场景下 page=1 的 token-level 命中能比 page=16 快 20-50%。

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

答案

token-level 时 K/V 物理布局是按 slot id 散落的（非连续），经典 attention kernel 要求 K/V 是连续矩阵，需要先 gather 拼起来，慢且费显存。 FlashInfer 的 paged attention kernel 内置 page table indirection（按 indices 间接访问 K/V），不需要 gather；cascade kernel 进一步把 shared prefix 的 K/V 只读一次给 N 个 Q 用，token-level 才能高效跑。退路：用 Triton 后端也能跑，但比 FlashInfer 慢 15-30%。

KV pool 碎片严重时会怎么表现？怎么处理？

答案

表现：(a) `available_size()` 看似足够但 `alloc_extend` 失败（找不到足够槽位）；(b) attention kernel L2 cache 命中率掉，TPOT 涨 5-10%；(c) 长时间运行后吞吐缓降。处理：(a) 调用 `flush_cache`（[`flush_cache.py`](../sglang/python/sglang/srt/mem_cache/flush_cache.py)）强制清空 RadixCache + 重整 free pool；(b) 监控 `fragmentation_rate` 指标；(c) 严重时改用 `page_size > 1` 牺牲粒度换 locality；(d) 业务流量平稳后定期重启副本（K8s rollout 配合 PDB）。

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

答案

`MLATokenToKVPool`（[`memory_pool.py:1631`](../sglang/python/sglang/srt/mem_cache/memory_pool.py)）。原因：MLA（Multi-head Latent Attention）把 K 和 V 压缩到一个**共享低维 latent 空间**，物理上不再是 `[num_kv_heads, head_dim]` 而是 `[latent_dim]`。普通 MHATokenToKVPool 的张量形状不适配。 DeepSeek-V3 用 `DSATokenToKVPool`（MLA 的 sparse 变体）。所有 MLA 系列都得用 `flashinfer-mla` 或 `flashmla` attention backend 配套。

alloc_extend_kernel 为什么用 Triton 实现而不是 Python？

答案

分配调用频率极高：prefill batch 每请求 1 次 + decode 每步 batch_size 次。批量大时 Python 实现里的循环、tensor 索引、bitmap 操作会成为 CPU 瓶颈（per-call 0.5 ms 看似小，但 1000 QPS 下累积成 ms 级延迟）。 Triton kernel 把"在 free 位图里找 N 个连续 / 散落 slot"这件事 fuse 成单个 GPU kernel，几 μs 完成，无 CPU-GPU 同步。 PyTorch 调度时 Python 调用 Triton 一次，结果直接在 GPU 上可用，不必拷回 CPU。

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.py。

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