第 4 课：SGLang RadixAttention、内存池与 Overlap Scheduling

本课目标与通过标准

本课产出：能逐行解释 RadixCache.match_prefix/insert/cache_*_req、请求池与 token-to-KV pool 的引用关系，并说明 overlap event loop 的读写屏障。

通过要求：唯一代码填空题通过全部断言；Q1～Q3 都能沿源码对象给出因果链；能指出一个正确性不变量、一个性能边界和一个需要 benchmark 才能确认的结论；总分至少 8/10。

源码版本与阅读方法

本课程使用固定 commit 的永久链接保证行号可复现，同时链接 current docs 供核对最新变化。阅读时先画调用图和状态所有权，再进入分支；不要从大文件第一行机械顺读。

源码阅读顺序：

1. mem_cache/radix_cache.py::RadixKey：token ids 之外为何还有 extra_key、page alignment 与 EAGLE bigram view。

2. RadixCache.match_prefix、_match_prefix_helper、_split_node：压缩 radix edge 怎样在中间命中时分裂。

3. insert、cache_finished_req、cache_unfinished_req：请求 KV 怎样移交给 cache 引用。

4. inc_lock_ref/dec_lock_ref/evict：protected 与 evictable 计数。

5. Scheduler.event_loop_overlap：何时处理上一批结果、何时 launch 当前批、WAR barrier 为什么存在。

源码快照：SGLang 9deb6952afa4。

核心概念与运行机制

1. 它是什么，解决什么问题

在多轮对话、Few-Shot 提示工程、智能体（Agent）多步推理以及结构化 JSON/代码生成等高频场景中，大量请求天然共享极其庞大的公共 Prompt 前缀（如系统提示词 System Prompt、历史上下文或 Few-Shot 示例）。传统推理架构在每次接收到新请求时，若不进行跨请求前缀缓存，就会无谓地重复计算这些长前缀的 Key-Value 张量，导致计算资源极大浪费、首字时延（TTFT, Time To First Token）显著激增，且多份完全相同的 KV Cache 会迅速耗尽宝贵的 GPU 显存。

SGLang RadixAttention 颠覆了静态预设前缀缓存的局限，创造性地将计算机经典数据结构——基数树（Radix Tree / 压缩前缀树）直接引入大模型推理引擎的 KV Cache 管理层： - 动态自动前缀共享：无需开发者显式声明前缀边界，系统自动将所有历史与进行中请求的 Token 序列组织成一棵按 Token 压缩的基数树。只要新请求包含已有路径的前缀，即可实现零计算开销的 KV 复用。 - 与底层物理内存池解耦：如图 4-1 所示，

根节点: []

[Hello, World] (Len=2, Ref=3)

[How are you?] (Len=4, Ref=1)

[Explain CUDA] (Len=3, Ref=2)

可驱逐叶子堆 (按 last_access_time)

堆顶: 最早访问的叶子 (Ref=0)

较近访问的叶子

Ref降为0

图 4-1：SGLang RadixAttention 基数树与按访问时间的叶子淘汰机制

Radix 树节点本身并不存储海量的 KV 权重张量，而是持有指向底层分页内存池（Memory Pool / Paged KV Blocks）的引用指针与引用计数。 - 动态生命周期与自动驱逐：当显存告急时，系统把未被活跃请求锁定的叶子节点放入以 last_access_time 为键的堆，优先驱逐最久未访问的叶子（效果近似 LRU），将物理块归还内存池，实现了逻辑前缀索引与物理显存容量的完美自适应。

2. 它如何工作

RadixAttention 与其调度子系统的高性能运转依赖于严密的数据流与异步流水线设计：

1. 树节点分裂与前缀匹配（Prefix Matching & Node Splitting）：
   • 当新请求携带有序列 $T = [t_1, t_2, \dots, t_m]$ 进入调度器时，调用 RadixCache.match_prefix 从根节点（Root）向下遍历匹配。
   • 若匹配路径在某个已有压缩边的中间中断（例如已有边表示 [1, 2, 3, 4]，而新请求在该位置为 [1, 2, 5]），系统调用 _split_node 将原边原子地分裂为一个公共父节点 [1, 2] 和两个子边分支 [3, 4] 与 [5]。分裂操作仅调整树结构指针与物理 Block 偏移量，绝对不触发 GPU 显存内 KV 数据的物理搬运与拷贝。
2. 引用计数保护与按访问时间淘汰（Lock Reference & Eviction）：
   • 处于活动推理状态（Prefill 或 Decode）的请求所引用的所有树节点，其 lock_ref 计数均大于 0，禁止被任何驱逐策略回收。
   • 当请求完成全部生成并退出时，调用 dec_lock_ref 释放锁定；当节点的 lock_ref 降为 0 时，该节点成为可驱逐（evictable）对象，其可驱逐大小计入 evictable_size。
   • 当系统需要分配新物理 Block 但可用配额耗尽时，RadixCache.evict 收集当前所有 lock_ref == 0 的叶子，按 last_access_time 建堆（heapq），反复弹出最早访问的叶子并释放其物理 KV 块；若父节点因此变为无锁叶子，再把父节点压入堆，直到释放量满足需求（如图 4-1 所示）。SGLang 并不维护 LRU 双向链表。
3. CPU 调度与 GPU 前向的重叠掩盖（Overlap Scheduling）：
   • 如图 4-2 所示，在 SGLang SRT 的 overlap 事件循环中，

CPU 调度线程 (Thread 1)

Batch N+1: Radix 匹配与显存块分配

准备输入张量与元数据

GPU 计算流 (Stream 1)

Batch N: 模型前向计算与 Attention 计算

Batch N+1: 立即进入 GPU 计算

重叠并行 (Overlap)流水衔接

图 4-2：SGLang CPU 调度准备与 GPU 执行重叠掩盖 (Overlap) 工作流

CPU 侧的 Python 调度器与 GPU 侧的模型前向计算被解耦为两条异步流水线。 - 当 GPU 正在 CUDA Stream 1 上全速执行 Batch $N$ 的矩阵乘法与 Attention 算子时，CPU 调度线程已经在 Stream 0 侧提前进行 Batch $N+1$ 的 Radix Tree 前缀查找、物理 Block 索引分配、张量组装（Tensorization）与元数据打包。 - 这一流水线消除了传统同步架构下每个 Step 之间高达数百微秒的 CPU 调度气泡，实现了极高的 GPU 计算占空比。

3. 正确性条件与常见误区

RadixAttention 的工程实现极其精巧，任何细微的不变量破坏都会导致不可挽回的推理错误甚至显存损坏：

• 命名空间与 extra_key 严格隔离：
  • 常见误区：误以为只要 Token ID 序列完全相同，其计算出的 KV Cache 就一定能够通用。
  • 正确性要求：在多租户系统、动态 LoRA 适配器加载、不同 RoPE 缩放基数（$\theta$）、多模态 Vision-Language 模型（相同的图像占位符 Token 对应完全不同的图像特征向量）场景下，相同的 Token 序列产生截然不同的隐层表达。SGLang 必须通过 RadixKey(token_ids, extra_key) 将 LoRA ID、多模态特征哈希、模型版本等元数据编码进命名空间，防止跨租户越权访问或错误状态复用。
• Page 对齐不变量（Page Alignment Invariant）：
  • PagedAttention 等底层内核以固定大小的 Block（如 16 或 32 个 tokens）为物理调度单元。若匹配到的公共前缀长度非 Block Size 的整数倍（如 Block Size 为 16，前缀长 25），未满一个完整 Page 的尾部 9 个 token 不能被盲目作为只读前缀直接复用，必须严格遵循向下截断对齐原则： $$\text{aligned\_len} = \lfloor \text{matched\_len} / \text{page\_size} \rfloor \times \text{page\_size}$$ 否则后续 Decode Step 在原地追加 KV 时会直接覆盖同一 Page 内共享前缀的其他请求的数据。
• 写后读（WAR, Write-After-Read）数据冒险与异步屏障：
  • 在图 4-2 的 Overlap 模式下，Batch $N$ 的 GPU Kernel 正在异步读取上一批的输入缓冲区或采样元数据，而 CPU 线程已经在流水准备 Batch $N+1$ 的写操作。
  • 必须使用 CUDA Event 记录 Batch $N$ 在 GPU 上的读取完成时机（read_finished_event），并在 Batch $N+1$ 写入覆盖共享 Buffer 之前强制等待该事件完成（WAR Barrier）。绝对不能仅凭 Python 代码的书写次序来假定硬件时序。

4. 性能与工程取舍

• 前缀命中率与树管理开销的权衡：
  • 在高并发对话代理或搜索检索增强（RAG）场景下，RadixAttention 可将 Prefill 吞吐提升 3~5 倍，TTFT 降低 80% 以上。
  • 然而在完全随机、无任何共享提示词的独立长文档生成负载下，频繁的树节点匹配、分裂与淘汰堆维护反而会带来无谓的 CPU 开销。
• Block Size 大小的物理取舍：
  • 较大的 Block Size（如 32 或 64）：显存连续性好，GPU Memory Coalescing 效率高，Page Table 元数据极小；但会导致较大的内部碎片，且降低 Page 对齐下的前缀复用颗粒度。
  • 较小的 Block Size（如 4 或 8）：前缀复用率达到极致，几乎不浪费任何公共前缀；但会导致 Page Table 尺寸膨胀数倍，增加寻址开销与 TLB/Cache 缺失率。
• 纯 Python 树遍历 vs. C++ 核心加速：
  • 随着并发请求数突破数千，基数树节点数膨胀到数万量级，Python 解释器的 GIL 与对象访问时延会成为瓶颈。高并发工业部署中通常将 RadixCache 核心下沉至 C++ 或 Rust 扩展中实现无 GIL 锁的高性能并发查找。

具体演示

设系统配置 page_size = 4，当前缓存树中已注册一条路径： $$\text{Path 1}: \quad [\text{tok}_1, \text{tok}_2, \text{tok}_3, \text{tok}_4, \text{tok}_5, \text{tok}_6, \text{tok}_7, \text{tok}_8], \quad \text{namespace} = \text{"lora\_finance"}$$ 新请求携带序列： $$\text{Query}: \quad [\text{tok}_1, \text{tok}_2, \text{tok}_3, \text{tok}_4, \text{tok}_5, \text{tok}_6, \text{tok}_7, \text{tok}_9], \quad \text{namespace} = \text{"lora\_finance"}$$

1. 字符匹配长度为 7（前 7 个 token 完全一致）。

2. 应用页面对齐截断：$\lfloor 7 / 4 \rfloor \times 4 = 4$。

3. 调度器仅安全复用前 4 个 Token 的物理 KV 块，并将请求的有效 Prefill 起始位置推进到第 5 个 Token，后续 4 个 Token（即 $\text{tok}_5 \sim \text{tok}_9$）交由 ModelRunner 重新执行轻量 Prefill 计算。

4. 若新请求的 namespace 为 "lora\_medical"，即使前 7 个 Token 完全相同，匹配结果也立即强制归零，安全隔离不同权重的隐层计算。

实践任务：唯一代码填空题

实现 page-aligned、namespace-aware 的最长前缀查询。该练习是 RadixCache.match_prefix 的可运行语义模型，不要求实现压缩树。

只能修改 TODO/______ 位置，不得删除断言或放宽通过条件。

def longest_cached_prefix(query_tokens, query_namespace, cached_paths, page_size):
    if page_size <= 0:
        raise ValueError("page_size must be positive")
    best = 0
    for cached_tokens, cached_namespace in cached_paths:
        if cached_namespace != query_namespace:
            continue
        common = 0
        for left, right in zip(query_tokens, cached_tokens):
            if left != right:
                break
            common += 1
        # TODO：只有完整 page 能作为可复用 prefix 返回。
        aligned = ______
        best = max(best, aligned)
    return best

cache = [
    ([1, 2, 3, 4, 5, 9], "base"),
    ([1, 2, 3, 4, 6, 7, 8, 9], "adapter-A"),
]
assert longest_cached_prefix([1, 2, 3, 4, 5, 8], "base", cache, 4) == 4
assert longest_cached_prefix([1, 2, 3, 4, 6, 7], "adapter-A", cache, 2) == 6
assert longest_cached_prefix([1, 2, 3, 4], "adapter-B", cache, 2) == 0

检查方法

运行断言；增加 query 比 cache 更短的用例。解释真实 radix tree 为什么还需要 split node 和引用计数。

Q1

RadixAttention 与 PagedAttention 分别解决哪一层问题？

你的答案：

Q2

为什么 extra_key 是正确性和安全边界，而不只是优化标签？列出至少三项应进入 namespace 的状态。

你的答案：

Q3

Overlap scheduling 中为什么需要 WAR barrier？只在 Python 中先处理上一批再写下一批不行吗？

你的答案：

评分规则

• 代码 4 分：主路径 2 分、边界条件 1 分、能映射回源码对象 1 分；
• Q1～Q3 各 2 分：必须包含对象、状态变化、正确性或成本链；
• 一票否决：把论文峰值写成普遍生产结论；把源码支持写成所有模型/硬件可用；混淆算法正确性与性能；只背类名而说不清状态所有权。

Tip💡 点击展开参考答案与详细代码实现

参考答案（仅 answer 分支）

先独立完成。核对后请更换一个 batch、cache 或硬件条件重新推演，不能只复制。

def longest_cached_prefix(query_tokens, query_namespace, cached_paths, page_size):
    if page_size <= 0:
        raise ValueError("page_size must be positive")
    best = 0
    for cached_tokens, cached_namespace in cached_paths:
        if cached_namespace != query_namespace:
            continue
        common = 0
        for left, right in zip(query_tokens, cached_tokens):
            if left != right:
                break
            common += 1
        aligned = common // page_size * page_size
        best = max(best, aligned)
    return best

cache = [
    ([1, 2, 3, 4, 5, 9], "base"),
    ([1, 2, 3, 4, 6, 7, 8, 9], "adapter-A"),
]
assert longest_cached_prefix([1, 2, 3, 4, 5, 8], "base", cache, 4) == 4
assert longest_cached_prefix([1, 2, 3, 4, 6, 7], "adapter-A", cache, 2) == 6
assert longest_cached_prefix([1, 2, 3, 4], "adapter-B", cache, 2) == 0

Q1 参考答案

RadixAttention 与 PagedAttention 处于大模型推理系统完全不同的抽象层次，解决不同维度的系统瓶颈：

1. RadixAttention 作用于逻辑命名空间与前缀索引层（Logical Namespace & Prefix Indexing Layer）：
   • 解决的核心问题是跨请求、跨对话轮次的逻辑前缀冗余计算。
   • 它通过维护一棵按 Token 压缩的基数树（Radix Tree），建立从逻辑 Token 序列到物理显存块引用的多对一映射关系，并结合引用计数与按 last_access_time 的叶子堆淘汰实现动态复用与自动驱逐。它不负责显存物理页面的直接分配与底层注意力算子的寻址。
2. PagedAttention 作用于物理显存分配与计算内核层（Physical Memory Allocation & Kernel Execution Layer）：
   • 解决的核心问题是单请求内部序列不规则动态增长带来的显存碎片化与预分配浪费。
   • 它借鉴操作系统的虚拟内存分页机制，将连续的逻辑 KV Cache 离散化映射到非连续的物理显存块（Block/Page），消除外部碎片，并把内部浪费限制在每个序列最后一个未填满的块内，并在 GPU 端提供支持 Page Table 间接寻址的高性能 Flash-style Attention 算子。
3. 二者的协同共存：
   • 在 SGLang SRT 架构中，Radix 树的叶子/中间节点所持有的 value，正是 Paged Memory Pool 中的物理块索引（Block Indices）。RadixAttention 负责在上层提供前缀匹配与逻辑共享决策，PagedAttention 负责在下层执行物理块管理与高并发内核计算，二者不仅不冲突，反而是现代高性能推理系统的标准组合。

Q2 参考答案

token 相同不等于 KV 相同。LoRA/权重版本、多模态特征哈希、租户 ID 都要进 extra_key，否则会复用错误 KV。共享前缀还会形成 TTFT 计时侧信道，泄露其他租户的 prompt。

Q3 参考答案

只依靠 Python 代码的书写次序完全无法保证时序正确，必须显式引入写后读（WAR, Write-After-Read）屏障：

1. GPU 异步执行模型与假象：
   • PyTorch 和 CUDA Driver 的 API 具有天然的非阻塞异步性。在 CPU 线程中调用 model.forward() 或算子发射，仅仅是将指令压入 CUDA Compute Stream 的硬件环形缓冲区（Ring Buffer），CPU 线程瞬间返回。
2. 写后读（WAR）数据竞争的本质：
   • 在 Overlap 调度流水线中，当 GPU 仍在 Stream 1 上全速执行 Batch $N$ 的矩阵乘与 Attention 计算（持续读取共享的输入张量 input_ids、KV 块索引表 req_to_token、采样元数据等）时，CPU 调度线程已经在流水线上准备 Batch $N+1$。
   • 若直接在 CPU 上复用或覆盖写入下一批次的共享缓冲内存，CPU 的写操作将与 GPU 尚未完成的读操作发生严重的并发冲突（WAR Hazard），导致 GPU Batch $N$ 读入被污染的新数据，产生不可复现的数值畸变或 CUDA 非法内存越界访问（Illegal Memory Access）。
3. 为什么不能简单采用全局同步（torch.cuda.synchronize()）：
   • 若在 Python 层面直接调用全局同步等待上轮完成，虽然能够避免 WAR 冒险，但这会让 CPU 和 GPU 的执行退化为严格串行交替，重新引入高达数百微秒的 GPU 调度空洞，使 Overlap 流水线完全失效。
4. WAR Barrier 的最优工程解法：
   • 在 Batch $N$ 发射 GPU 读取内核后，在计算流上记录一个专用的轻量级事件 read_finished_event = torch.cuda.Event()；
   • 当 CPU 或复制流准备写入覆盖共享 Buffer 前，通过 write_stream.wait_event(read_finished_event) 让硬件写入流精准等待先前的读取完成，以极低开销实现无 CPU 阻塞的严格时序保证。

参考资料

• SGLang RadixCache source
• SGLang overlap scheduler source
• SGLang memory pool source
• SGLang schedule policy source
• SGLang/RadixAttention paper

源码链接固定到本课审阅 commit；current docs、支持矩阵和默认参数会变化，面试或部署前必须按目标版本重新核对。

← 上一课
第 3 课：SGLang SRT 源码地图 下一课 →
第 5 课：稀疏推理