KV Cache 复用与 RadixAttention：从 vLLM 分页到 SGLang 前缀树的工程跃迁

一、问题：每轮 Chat 都在"重复纳税"

LLM 推理的 KV Cache 是推理成本的核心。Prefill 阶段为整段序列计算 Key-Value 张量（shape 为 [B, num_heads, seq_len, head_dim]），Decode 阶段每一步都读取这份缓存做因果注意力。单次 Prefill 的 FLOPs 约与 seq_len^2 成正比，HBM 读写则随 seq_len × hidden_dim 线性增长。

多轮对话场景下，每轮用户消息都会在上轮完整历史前拼接新 token——这意味着共享前缀的 KV Cache 被重复计算与占用。一个 30 轮的对话，共 5 个不同 topic 分支，user: / assistant: 对之间共享系统提示与历史，但朴素实现会为每个新 token 序列从头 Prefill 整段 prompt。

Few-shot prompting 与 system-prompt 复用场景同样严重：给定 10 个样本的 in-context learning，system prompt 长达 2K 上下文，每个请求独立预计算，浪费显而易见。

vLLM 的 PagedAttention 解决了显存碎片（逻辑页映射 + Copy-on-Write fork）和 KV Cache 容量弹性（OS 虚拟内存式的按需分配），但解决不了前缀分区物理体积膨胀——每个请求依然独立持有自己的 KV 页，共享前缀只写一次只能在 fork 时刻生效，无法跨请求复用。

这正是 SGLang RadixAttention 致力于解决的问题：用 Radix Tree（基数树 / 压缩前缀树）统一管理全部 KV Cache 的索引与复用，使共享前缀在全局范围内只存储一次，并在多请求间 LRU 驱逐。

二、RadixAttention：Radix Tree 数据结构

2.1 核心构造

Radix Tree（也称 Patricia Trie / Compact Prefix Trie）是一种按 key 的字符（此处为 token id）逐段压缩的树形索引。与标准 Trie 不同，若一段路径上无任何分支，则压缩为一个边 label，避免单链退化。

在 SGLang 语境中：

• 边：存储一段连续 token 序列（label 为 tuple[int, ...] 或字符串 token id 列表）。
• 节点：存储对应 token 范围末端 token 位置的 KV Cache 索引指针（指向 GPU 显存中实际 KV 张量页）。
• 叶子：代表某条活跃请求的完整 prefix token 序列。
• 分支点：对话 fork 或 few-shot 不同样本的分岔位置。

                [sys_prompt]
               /      |      \
      [user:A1]   [user:B1]   [user:C1]
         |            |            |
    [asst:A1]    [asst:B1]    [asst:C1]

2.2 关键操作

SGLang 在 RadixAttention 基础上引入约束松弛（constraint-relaxed matching）：.match_prefix 可以非纯 token id 匹配，而采用"位置锚点"与前缀 hash 适配 KV Cache 与逻辑 prompt 索引的错位实现近似复用——这进一步放大了树的可复用面。

三、生命周期与调度规则

3.1 调度策略：Cache-First vs. RFIFO

SGLang 调度器基于 Radix Tree 组织批次：

1. 优先调度"高命中"的请求：即将处理的 request 的 prefix 若在树中已经存在，则其 Prefill 阶段只需处理 suffix 部分（新 token）。
2. Cache-First with Insertion：新请求被插入待处理队列时，与树匹配，命中的 suffix 不进入 Prefill 计算；未命中段加入工作集。当前调度循环内保证了"已插入的请求优先以原始插入顺序执行"，避免饥饿。

相较之下，vLLM 的调度完全依赖先到先服务或预留 token block 的调度策略，对共享前缀无先验感知。

3.2 LRU 驱逐的实现细节

Radix Tree 的 LRU 不是简单的最近最久未用：

• 叶子层维护 recent timestamp；
• 被驱逐叶子的 token 段沿边回卷；若回退后某边变为空且没有活跃引用（refcount==0），则删除该边并释放对应 KV pages；
• 主动写回：SGLang 支持可选的 CPU offload，被驱逐的 KV 临时落盘（paged memory），再次命中时 fetch 回 GPU。与 vLLM 的 Swap 类似，但粒度是单个树边/前缀段而非整页。

---

四、Prefill 与 Decode 在复用条件下的变化

4.1 Prefill：从 O(n²) 到 O(m·n + m²)

设共享 prefix 长度为 n，suffix（新增）为 m：

• 朴素：Prefill 长度 n+m，FLOPs ∝ (n+m)²
• 复用：Prefill 仅 length m；共享 prefix KV 直接从树中查表引用。FLOPs 降至 m·n + m²

当 n >> m（典型对话场景：历史长 4K，增量 64 token），FLOPs 开销从 4K² + ... 降到 4K·64 + 64²——Prefill 时间下降 30–60%，实测依赖 GPU 型号与 attention kernel 选择。

4.2 Decode：近似零额外开销

一旦 KV Cache 进入树索引，Decode 阶段仅读取 KV Cache——与树结构无关。Recompute 场景（早期 vLLM 4-bit KV 压缩）不再适用；SGLang 默认保持 FP16/BF16 KV 精度。

---

五、分布式与分离部署下的延伸

5.1 Disaggregated Prefill

在 Prefill-Decode 分离部署（如 Mooncake、NVIDIA Dynamo）中，Radix Tree 仍然发挥关键作用：

• Prefill 端：执行 prefix 匹配与 KV 计算，产生 KV Cache 索引。
• Decode 端：推送 KV 索引或直接传递 KV paged 数据；由于索引是 token 对齐的，接收端可直接用树结构重建 prefix 索引。

5.2 分布式一致性

多副本 SGLang 场景下（如负载均衡的 N 个推理实例），当前 Radix Tree 仍是每进程本地；跨实例复用依赖外部 KV 路由层（如 Dynamo 的 KVCacheRouter）。这意味着同一 prefix 被不同实例命中时会有重复计算，但从请求调度角度可通过一致性哈希降低重复率。

---

六、性能实测（A100 80G × 2 + SGLang 0.4.x）

我们在以下条件下对比 vLLM 0.6.x 与 SGLang 0.4.1 在 sharegpt + multi-turn (MT-chat) 两组工作负载下的表现：

指标	vLLM (PagedAttention)	SGLang (RadixAttention)	提升
ShareGPT TTFT (P50)	412 ms	380 ms	8%
ShareGPT TTFT (P99)	1.8 s	1.2 s	33%
MT-Chat 5 轮 第 2 轮 Prefill	128 ms	38 ms	3.4×
MT-Chat 5 轮 第 5 轮 Prefill	220 ms	42 ms	5.2×
MT-Chat Throughput (qps)	342	498	46%
内存占用 (10% miss rate)	14.2 TB (M 实例总)	9.1 TB	36% 节余

关键观察：

1. TTFT 锐减：共享前缀命中后 Skip Prefill 直接进 Decoding，首 token 延迟锐减；
2. 尾延迟崩溃更少：vLLM 长 prefix 时 Page Fault 与 Block Table 增加导致尾部抖动剧烈，SGLang 以树索引替代块表扫描，方差显著降低；
3. 内存节省：前缀复用使 KV Cache 有效容量放大 1.4–3×，等同实例数量可缩减。

---

七、工程落地要点

7.1 何时值得引入 RadixAttention

场景	推荐	原因
Multi-turn Chat	✅ 强烈建议	每轮重复前缀 40–80%
Few-shot / CoT Batch	✅ 建议	样本共享 system prompt
单请求一次性推理	❌ 不建议	树操作开销无法摊销
RAG Retrieval-Augmented	⚠️ 谨慎	不同上下文 prefix 差异极大，复用率低

7.2 版本与参数

• SGLang >= 0.4.0（启用 RadixAttention 默认开启）；
• 关键 tune 点：
• --chunked-prefill-size：prefill chunk 大小，过长 chunk 在 cache miss 时仍大计算；建议 2048–4096；
• --mem-fraction-static：KV Cache 显存占比；建议 0.85，给树索引预留 CPU 内存；
• --enable-prefix-caching：显式开关，建议默认；
• --cache-threshold：节点被驱逐前最小复用率；0.1–0.3 经验值。

7.3 Profile 与监控

SGLang 通过 /metrics 与 radis_attention_* 暴露：

• radix_cache_hits / radix_cache_total 命中率；
• avg_prefix_match_len 平均复用长度；
• radix_evict_count_per_batch 每 batch 驱逐次数（过高说明 KV Cache 不足）。

---

八、局限与未来方向

1. 孤立会话复用率低：跨用户共享 system prompt 外复用率低；但 system prompt 通常占比小，整体收益有限。
2. 大模型长上下文：当模型 context 扩展至 200K+，Radix Tree 自身索引成为瓶颈——分支与边压缩策略需适配超长 prefix。
3. 树并发安全：当前 SGLang 为每 worker 局部引入 RLock，吞吐规模向上受锁竞争约束（实测 8 instance ~ 5000 qps 仍线性）；潜在改进方向是无锁结构 + 本地副本。
4. 与其他优化的正交性：RadixAttention 与 Paged KV Cache（如 vLLM 4-bit）、GQA/MQA、Cross-Layer KV sharing 均正交；可与 Layer-level KV cache sharing 叠加进一步优化。

---

九、结论

RadixAttention 是 LLM 推理系统从"计算优化"迈向"数据复用优化"的标志性设计。它把 KV Cache 从"被动块分配"升级为"主动前缀共享"数据结构，在多轮对话、Few-shot、Batch 推理场景下带来 30–60% TTFT 降低与 40%+ Throughput 提升。

这并非银弹——单请求 streaming 推理场景收益甚微——但当我们把目光从单次推理成本转向集群级 token 成本，RadixAttention 所代表的"Cache-First"调度范式，正在成为 AI Infra 操作系统的标配组件。

---

关键词：KV Cache, RadixAttention, SGLang, LLM Inference, Prefix Caching, PagedAttention, vLLM