SGLang:为 LM 程序而生的前端 DSL + 协同设计运行时 —— 阅读笔记

SGLang:为 LM 程序而生的前端 DSL + 协同设计运行时 —— 阅读笔记

笔记日期: 2026-05-21 笔记作者: Zhongzhu Zhou 论文标题: SGLang: Efficient Execution of Structured Language Model Programs 作者: Lianmin Zheng、Liangsheng Yin、Zhiqiang Xie、Chuyue Sun、Jeff Huang、Cody Hao Yu、Shiyi Cao、Christos Kozyrakis、Ion Stoica、Joseph E. Gonzalez、Clark Barrett、Ying Sheng(Stanford / UC Berkeley / SJTU / Texas A&M) arXiv: 2312.07104v2,2024-06-06 状态: NeurIPS 2024 正式录用 代码: sgl-project/sglang

一句话总结

凡是把 LLM 应用从 demo 推过工程门槛的人,大概都撞过同一面墙:一次 gen() 调用很轻,但当你把两三次、二十次调用串成一个程序,加上分支、并行 fork、结构化输出、共享 system prompt 之后,你会突然发现 推理引擎把每一次调用都当成”这辈子第一次遇到”在做。KV 缓存被重新算一遍,system prompt 被重新 tokenize 再 prefill 一遍,few-shot 样例从头喂一遍,几个 fork 之间在做本该共享的工作。SGLang 就是冲着这堵墙来的。

SGLang 最干净的读法是把它看成被同一个观点黏在一起的两半:

  1. 一个 Python 嵌入式 DSL,专门服务”语言模型程序(LM program)“:gen()select()extend()fork()join(),再外加图像/视频原语。前端小、跟 Python 同源、可组合 —— branch-solve-merge、tree-of-thought、agent 循环、JSON-schema 约束解码,全都可以用原生 Python 控制流写。
  2. 一个协同设计的运行时,它知道 哪些调用与哪些调用共享前缀,并把这件事用到底。主力技术 RadixAttention 把 KV 缓存当成一棵 LRU 化的树形结构来管理,凡是出现的公共前缀 —— 不管是跨调用、跨实例、还是跨 TP rank —— 一律免费复用。

再加上两件锦上添花的优化:压缩有限状态机(compressed FSM) 把 regex 约束里被强制确定的多个 token 在一次 forward 里全部解码完;API 推测执行(API speculative execution) 对 GPT-4 这种黑盒接口顺手多生成几 token 来节省 round-trip。

实验数字也很有说服力:在 Llama-7B 上对比 vLLM、LMQL、Guidance 拿到最多 6.4× 吞吐3.7× 延迟降低;Mixtral-8x7B 和 Llama-70B 加上张量并行依旧能拿到同量级的提升;LLaVA 图像/视频负载上 6× 加速。缓存命中率在所有 benchmark 上落在 50%–99% 之间,缓存感知调度逼近离线最优解的 96%。Chatbot Arena 真实流量下,LLaVA-NeXT-34B 命中率 52.4%,Vicuna-33B 命中率 74.1%,首 token 延迟降低 1.7×。

这份笔记我会:(a) 给已经熟悉 vLLM PagedAttention 但还没认真想过”前缀感知调度”的读者把背景铺平;(b) 把 RadixAttention 讲透 —— 数据结构、LRU 策略、缓存感知调度定理、前端 hint;(c) 把压缩 FSM 解码和 API 推测执行写清楚;(d) 复盘实验,包括我觉得最值得看的几个消融;(e) 把 SGLang 放到 LLM 服务系统和 LM 程序框架文献的坐标系里。

1. 前置知识

SGLang 横跨三个子领域:LM 程序框架、LLM 服务系统、结构化生成。背景不到位的话,贡献看上去都不大;背景到位之后,这篇论文就变成”一个小心翼翼的观察 + 它必然的推论链”。下面把需要的几块讲清楚。

1.1 什么是”语言模型程序(LM program)”?

DSPy 和这篇论文都在用这个词,大致是同一个意思。LM program 是任何”运行时行为涉及多次 LLM 调用、并用控制流把它们串起来”的可执行制品。最典型的例子包括:

  • Chain-of-thought + self-consistency。 给同一个问题采样 kk 条不同的推理链,然后对答案多数投票。
  • Branch-solve-merge。 在不同子 prompt 上并行采样 kk 个偏题解,然后合并。
  • Tree-of-thought。 搭一棵搜索树,每个节点都是模型生成的局部推理步骤,扩展有希望的分支,剪掉死分支。
  • ReAct 等 agent。 把模型调用和工具调用(搜索、计算器、代码执行)交替起来。
  • 检索增强生成(RAG)管线。 Embed → retrieve → re-rank → generate,有时再带一个自批评循环。
  • JSON 模式 / 结构化输出。 单次 gen() 调用,加上 regex(或者 JSON schema 编译出的 regex)约束。

所有这些用例都共享两个结构性特征:

  1. 多次 LLM 调用 + 控制流。 任务不可能在一次 forward 里答完,后面的调用依赖前面调用的输出。
  2. 结构化输入输出。 调用收到的是结构化 prompt(system prompt + few-shot + 动态上下文),吐出的是结构化输出(JSON、分数、枚举里的某个选项)。

“两个特征”是这篇论文的道德基础。每个特征都打开一类优化空间,而这类空间是单次调用引擎拿不到的:特征 (1) 在调用之间制造 前缀共享;特征 (2) 在一次调用内部制造 多 token 确定性

1.2 KV 缓存,60 秒讲完

decoder-only Transformer 在推理时,对每一层、每一个位置,都会算出一对向量:key KKvalue VV。计算第 tt 个 token 的 forward pass 需要前面 0,1,,t10, 1, \dots, t-1 所有位置的 keys 和 values。每一步都重新算这些 K/V 在算力和带宽上都是 O(t2)O(t^2),所以我们缓存它们:每次 forward 后,新算出来的 Kt,VtK_t, V_t 被追加到每请求的 KV 缓存 里,放在 HBM 上,下一步只算 Kt+1,Vt+1K_{t+1}, V_{t+1},前面用缓存。

对 SGLang 来说关键的一点是:KV 缓存只依赖前缀 token。如果两个不同请求共享前 LL 个 token,它们在这些位置上的 KV 缓存就是 逐 bit 相同 的。复用是数学上无损的 —— 不是近似、不是量化、不是学到的压缩 —— 是精确复用。

像最早期的 vLLM 这种单调用引擎,把每个请求当成一个自包含单位:admission 时分配 KV 缓存,completion 时释放,跨请求复用根本不在桌面上。SGLang 填的就是这条缝。

1.3 vLLM 与 PagedAttention

vLLM(Kwon 等,SOSP 2023)解决的是另一类 KV 缓存问题:单请求内部的碎片化。最朴素的方式是给每个请求按最坏情况生成长度预留一块连续的 KV 缓存,先结束的请求就把 HBM 浪费掉。PagedAttention 把 KV 缓存切成固定大小的块,像操作系统给虚拟内存分页一样管理,消灭了内部碎片,实测能让 batch size 大体翻一倍。

SGLang 也用 paged 存储 —— 它跟 PagedAttention 兼容,不是替代。但 SGLang 在 跨请求 维度加了一层:这些 paged block 现在可以在多个共享前缀的请求之间被引用。每一个 page 在某种意义上变成了一个内容寻址的 extent。

1.4 Radix 树

Radix 树(也叫 Patricia trie、压缩前缀树)就是把所有”只有一个孩子的内部节点”和它的父亲合并掉的 trie。边上挂的不是单个元素,而是变长序列。对我们要做的事来说:它是存一组字符串(或 token 序列)的天然数据结构,可以在 与集合大小无关O(new string)O(|\text{new string}|) 时间内,找到任何新字符串与集合的最长公共前缀。

如果你写过路由表,你就用过 radix 树;如果你写过 tokenizer 的词表查询,你也用过它。SGLang 的贡献是把它套到 KV 缓存 extent 上 —— 边上是 token-id 序列,叶子上挂的是逐 token 的 KV 缓存 page 指针。

1.5 LMQL 与 Guidance —— 两位最近的前辈

之前已经有两种针对同一个生态位的语言:

  • LMQL(Beurer-Kellner 等,PLDI 2023):自定义语法的声明式语言,有 where 子句作为约束系统。约束/类型系统很强,但后端是 Hugging Face Transformers,缺少连续批处理和张量并行。
  • Guidance(Microsoft,2023):Python 嵌入式 DSL,原语类似(genselect)。后端是 llama.cpp,同样缺少批处理和并行支持。

两者都早于 SGLang 前端,论文也很认真地致谢。SGLang 真正的不同就在 前端和运行时的协同设计:LMQL/Guidance 可以被重定向到任何后端,但 没有一个后端真正理解它们的程序结构来吃这块红利。SGLang 的运行时(SRT,SGLang Runtime)是为前端的前缀共享和并行语义专门写的,而前端会主动给运行时发”提示(hint)“,运行时按 hint 行事。

1.6 连续批处理(Continuous batching)

迭代级别批处理(Yu 等,OSDI 2022)是现在的事实标准:每一次迭代,batch 都可以加入一个新请求或者赶走一个结束的请求。所有现代引擎(vLLM、TGI、TensorRT-LLM、SGLang)用的都是某种变体。有意思的问题是 调度顺序:当队列里的请求比 batch 能装的更多时,我们按什么顺序去 admit?

SGLang 的回答是:把”已经在缓存里”的前缀重叠最多的那一个请求拉进来。这就是论文里的 cache-aware scheduling(缓存感知调度)。

1.7 受约束 / 结构化解码

想强迫 LLM 输出合法 JSON(或合法 SQL/XML/任意正则语言),标准套路是:每一步根据当前自动机状态算一个词表掩码,把不允许的 token logit 设成 -\infty,正常采样;自动机状态根据采样到的 token 推进。Outlines、lm-format-enforcer 是 2023 年把这种做法做火的两个库。

约束通常以 regex 形式给出。绝大多数情况下都能用 JSON-schema → regex 的编译器。regex → NFA → DFA,常常还会被极小化成一个状态数可控的 FSM,每一步去查 FSM。

SGLang 攻击的是这种低效:如果从某个状态出发,FSM 只有一条出边,标签是序列 “summary”:“,那么这一串 token 是 被强制 的 —— LLM 没有自由。逐 token 解码白白浪费了好几次 forward。我们应该把它们一次性解出来。

带着这七块前置,论文的剩下部分基本可以总结成一句:“把 KV 缓存装进 radix 树,按它来调度,把被强制的子序列一次性解码完。“

2. 编程模型

论文第 2 节介绍前端。我在这节走得快一点,因为运行时才是更厚的贡献。但语言本身很重要:第 3–5 节每一个运行时优化,都依赖前端把”对的信息”喂上去。

2.1 跑通示例

论文开篇就给了一个 multi_dimensional_judge 函数:一个 branch-solve-merge 的作文打分器,输入是图像和一篇作文,先用 select 让模型判断作文跟图是否相关,再 fork 出三个并行评估器(清晰度、原创性、证据),最后产出 JSON 摘要带 letter grade。

如果用原生 OpenAI 风格的代码写,大概要 50 行 asyncio 胶水 + 一堆脆弱的字符串解析;SGLang 写出来大约 25 行。省下的行数不是重点,可读性才是 —— 你可以一眼看到控制流,因为控制流就是 Python。

@function
def multi_dimensional_judge(s, path, essay):
    s += system("Evaluate an essay about an image.")
    s += user(image(path) + "Essay:" + essay)
    s += assistant("Sure!")

    s += user("Is the essay related to the image?")
    s += assistant(select("related", choices=["yes", "no"]))
    if s["related"] == "no": return

    forks = s.fork(len(dimensions))
    for f, dim in zip(forks, dimensions):
        f += user("Evaluate based on the following dimension:" + dim + ".")
        f += assistant("Judgment:" + gen("judgment", stop="END"))

    judgment = "\n".join(f["judgment"] for f in forks)
    s += user("Provide the judgment, summary, and a letter grade")
    s += assistant(judgment + "In summary," + gen("summary", stop=".") +
                   "The grade of it is" + gen("grade"))

    schema = r'\{"summary": "[\w\d\s]+\.", "grade": "[ABCD][+-]?"\}'
    s += assistant(gen("output", regex=schema))

需要注意五件事:

  1. system / user / assistant 会按当前模型的 chat template 渲染 —— 程序与具体模型解耦。
  2. select 调一次模型,把输出约束到给定枚举里;运行时通过对每个选项算 logprob 实现。
  3. s.fork(k) 把当前 prompt 状态复制 kk 份并行 —— 每个 fork 继承共同前缀,后面的 append 才发散。
  4. gen("name", stop=..., regex=...) 是主力,regex 参数直接接上压缩 FSM 解码器。
  5. s["name"] 取出之前某个 gen 的结果,这才是把 Python 控制流变成”真的可以 if 分支”的关键 —— 没有它,你没办法基于模型输出做分支。

2.2 原语,正式版

  • gen(name, stop=, regex=, max_tokens=) —— 采样若干 token,以 name 命名存下;可选 regex 约束。
  • select(name, choices=) —— 在给定选项里挑 logprob 最高的那个。
  • extend 或者 += —— 在状态后面追加一段(带 role 标记)字符串。
  • [name] —— 取一个之前 gen 的结果;尚未就绪则阻塞。
  • fork(k) —— 把状态分成 kk 个并行副本。
  • join —— 同步回主流;s["name"]for f in forks 会隐式触发。
  • image(path) / video(path) —— 挂多模态内容。

集合刻意压得很小。论文 Table 1 跟 LMQL、Guidance 对比:SGLang 显式多了 videoforkjoin —— 别人没有一级的并行原语。

2.3 执行模式 —— interpreter vs compiler

interpreter(默认)把 SGLang 函数当成 Python 跑。每个原语被丢到一个 per-program 的 stream(运行在后台线程),原语都是 non-blocking,除非程序回读了某个结果。这让 程序内部并行 几乎免费:forks = s.fork(3) 调度三条独立 stream,每个 fork 内部的 gen 都是异步发起。

compiler 模式(论文 Appendix D)把程序 trace 成静态计算图,做图级优化:死代码消除、公共子表达式消除、提前规划前缀。论文报告 interpreter 与 compiler 性能相当,主要结果都用 interpreter;compiler 主要帮短程序,因为短程序里 Python 解释开销占比大。

一个好用的心智模型:interpreter 之于 SGLang runtime,就像 CUDA stream 之于 GPU。程序调 gen() 立刻返回,运行时在后台做活儿,read 时隐式同步。

2.4 SGLang 在栈里的位置

论文画了一张干净的二层图:高层框架(LangChain、DSPy)负责 prompt 管理与 pipeline 编排;低层语言(LMQL、Guidance、SGLang)提供直接原语。SGLang 同时是高层框架的 编译目标 —— 论文演示了 DSPy → SGLang 的下沉,RAG pipeline 拿到了对应的加速。这是正确的分工:高层管 prompt 设计与质量,低层管吞吐。

3. RadixAttention

整篇论文的核心。第 3 节和 Appendix A。

3.1 机会在哪

任何非 trivial 的 LM 程序里,前缀都以三种方式被共享:

  • 程序内部跨调用。 system prompt 在每次调用里都出现;few-shot 样例不变;fork 之后,fork 前的前缀在所有 fork 之间相同。
  • 跨程序实例。 同一个聊天应用的两个用户共享 system prompt;同一个 chatbot 的两个用户共享模型 persona;self-consistency 采样从同一个前缀扇出。
  • 跨张量并行(TP) rank。 每个 rank 只持有模型的一片,但 前缀(以 token id 形式)在所有 rank 上是同一个。每个 rank 复用自己那一片的 KV 缓存,完全不需要 rank 间额外通信。

没有复用,每一个共享前缀都被每个请求重新算一遍。有了 RadixAttention,它就算一次、存进树、下次查表。

3.2 数据结构

radix 树的每个节点持有:

  • 边标签 —— 从父节点过来的一段 token id 序列;
  • 指向一段 KV 缓存 page 的指针,这些 page 就是上面那段 token 对应的 KV;
  • 引用计数 —— 当前有多少在飞的请求正在用这段 token(用了就不能 evict);
  • 最后访问时间戳,用于 LRU 淘汰。

树根是空前缀。新进来一个 token 序列 t0t1t2tnt_0 t_1 t_2 \dots t_n 的请求,运行时从根开始往下走,在每个节点上贪心地沿着标签匹配请求前缀,能走多远走多远。匹配到的前缀长度 mm 告诉 prefill kernel 可以跳过哪段:t0tm1t_0 \dots t_{m-1} 的 KV 缓存已经在 GPU 上,prefill 只需要算 tmtnt_m \dots t_n 的。

如果匹配在某条边的中间停下来 —— 也就是请求前缀只匹配了某条 ll 长边的前 k<lk < l 个 token —— 运行时会把那个节点 拆成两个:一段父边持有前 kk 个 token(同时被旧请求和新请求引用),一段子边持有剩下的 lkl-k 个 token(仍属旧请求)。

新请求结束时,它新生成的 token —— 加上 prefill 阶段新算出来的前缀 token —— 作为新边插回树。沿途经过的节点引用计数减一。

这是一棵典型的”带引用计数的 copy-on-write radix 树”,只不过 payload 是 GPU 内存 page 而不是磁盘 block。论文指出:这棵树维护在 CPU 端,开销可以忽略 —— 树的遍历和分裂代价都被模型 forward 完全淹没。

3.3 LRU 淘汰与引用计数

HBM 很快就会满:7B 模型的 KV 缓存大约是 0.5 GB / 1000 tokens,24 GB 的 A10 大概能装总共 3 万 token 的缓存。淘汰是必然的。

RadixAttention 的策略是 在叶子上做 LRU。两个重要细节:

  1. 引用计数。 当前正在被某次 forward 读的节点不能 evict。引用计数在 admission 时 +1,completion 时 -1。这让”缓存”与”活跃 batch”可以共享一个 HBM 池 —— HBM 里没有”缓存区”和”活跃区”两个固定区域,所有 KV 内存都是同一池,一个节点引用计数归零它就变成可 evict 对象。
  2. 叶子优先淘汰。 永远先 evict LRU 叶子,内部节点只有在它所有孩子都被 evict 之后才能被 evict。这保住了”公共祖先” —— 也就是未来最可能被命中的那段前缀。

论文 Figure 3 有一张漂亮的 9 步示意图,展示 radix 树在两次 chat session、一次 few-shot batch、一次 self-consistency 采样过程中的演化。这张图就算别的不读也值得看 —— 它把”树天然刻画了 LM 程序的共享模式”这个抽象 claim 变得具体可感。

3.4 前端 hint

一个微妙的实现细节。当 SGLang interpreter 执行 s.fork(k) 时,它可以朴素地把 kk 份完整 prompt(每一份直到 fork 点都相同)甩给运行时,让运行时在 admission 阶段自己发现公共前缀。这做法能用,但有竞态:如果这 kk 份 fork 在前缀尚未完整插入树之前就到达运行时,可能每个 fork 都会触发一次自己的公共前缀 prefill,白白浪费计算。

前端给运行时发了一个 hint:“我下面要发 kk 个请求,它们共享这段前缀”。运行时先把前缀作为一次 no-completion warmup 插入树,再 admit 这 kk 个孩子,保证 fork 到达时前缀一定在树里。论文 Figure 8(c) 的消融显示:关掉 hint,tree-of-thought 上吞吐有可测下降。

这是前端-后端协同设计的教科书例子:任何一边单独都做不出这个优化。

3.5 缓存感知调度

论文里在理论上最让人舒服的贡献是 cache-aware scheduling。问题是:当等待队列大于 batch 能装下的时候,我们以什么顺序 admit 它们?

直觉:挑那个已匹配前缀最长的请求。它 prefill 得最少,腾出周期给别人。等价地说,挑那个对缓存命中最多的。

算法:按”对当前 radix 树的最长匹配前缀长度”降序排队,把队头拉进下一个 batch。结束后(树可能长大了)重算前缀长度再来。

论文证了一个干净的离线最优结果:

定理 3.1。 对一批请求,要达到最优的缓存命中率,只需要按 DFS 顺序遍历它们的请求 radix 树,且缓存大小至少为最长请求长度。“最长共享前缀优先”顺序就是一种 DFS 顺序。

证明(Appendix A.3)是一个小的”交换论证”:在任何一步,把下一个要服务的请求和它在 DFS 顺序中的兄弟换一下,不可能增加命中率。“缓存大小 ≥ 最长请求长度”这一条件是为了保证任何单个请求的前缀,在它自己被服务期间,不会被淘汰掉。

在线情况下,请求到达时间不可预期,DFS 被打断。但”最长前缀优先”的启发式在动态生长的树上,仍然是 DFS 的一个近似。论文实测命中率(Figure 13)在 11 个 benchmark 上 平均达到离线最优的 96%

一个 caveat:贪心的缓存感知调度可能 饿死 一些前缀始终匹配不到的请求。论文承认这一点,把公平性扩展指向后续工作 [42]。

3.6 分布式扩展

对张量并行,每个 rank 持有模型一片和 KV 缓存的一片。radix 树本身 不需要 分片,因为它的 key 是 token-id 序列,所有 rank 都看到同一段 token,每个 rank 只是把自己那片 page 取出来用。同步只通过模型 forward 发生,跟任何其他 TP 优化没差别。

对数据并行(多 worker),每个 worker 一棵独立的树。论文(Appendix A.4)讨论了一种路由策略:把相同前缀的 prompt 钉到同一个 worker,以一点点 load imbalance 换更高的 per-worker 命中率。

3.7 为什么它比 vLLM 的前缀缓存更好?

vLLM 之前就已经有实验性的 prefix caching。论文很仔细地把差异写清楚:

  • 共享范围。 vLLM 只缓存”system prompt 那一段公共前缀”;RadixAttention 缓存 任何深度任何 前缀。
  • 树结构。 vLLM 的表式缓存表达不了 fork 程序的分支结构。Figure 8(c) 的 “No Tree-Structure” 消融显示:退到只用表,吞吐会掉一大块。
  • 调度耦合。 vLLM admit 是 FCFS;RadixAttention 是前缀感知的。“FCFS Schedule” 消融把”放弃前缀感知排序”的代价单独剥出来。
  • 淘汰策略。 vLLM 在 cache extent 上没有 LRU;RadixAttention 有。

最戏剧化的是 “No Cache” 消融:完全关掉 RadixAttention,在那几个有利负载(LLM Judge、Tree-of-Thought、MMLU)上能掉掉 50%–75% 吞吐。只关掉树结构、退回表,大概抢回一半;再退回 FCFS,再抢回四分之一。每个组件都重要。

3.8 没命中时的开销

一个合理担忧:如果我的负载根本没有共享前缀 —— 比如 ShareGPT 那种独立 chat 流 —— RadixAttention 会不会反而变慢?

论文回答(第 6.3 节):在 100 个 ShareGPT 请求上,总运行时间 74.3 s,其中 RadixAttention 的簿记开销 0.2 s,占 0.27%,论文说”可忽略”,确实是。

这是 RadixAttention 能 默认开启 的实操理由:碰到共享前缀的负载就拿 6× 加速,没有共享的负载几乎不付任何代价。

4. 用于受约束解码的压缩 FSM

论文第 4 节,比 RadixAttention 小但工程感很干净。

4.1 背景

用户给一个 regex(或者 JSON schema 编译出的 regex),运行时构造一个确定性 FSM。每一步,FSM 当前状态决定一个允许 token 的集合,logit mask 把其他清零,模型采样,FSM 根据采样到的 token 推进。

低效点:如果 FSM 在某个状态只有一条出边,标签是 summary":_" 这一串 token(就是上面例子里的 {"summary": "),模型根本没有自由。我们却花了 11 次独立 forward 在解码 11 个一进入此状态就已经被定死的 token。

4.2 压缩 FSM

预处理 FSM:任何”链状的、出度均为 1 的状态序列”被 熔合 成一条超级边,标签是这些边的 token 拼接结果。这在结构上跟”基于 FSM 建一棵 radix 树”完全一样(我猜这是有意呼应第 3 节,但论文没明显点破)。

运行时,当模型进入这种超级边时,运行时在 一次 forward 里把整条边上的 token 全部解码。机制是:复用 prefill kernel,带上 forced_tokens = [...];下一次采样从超级边末尾继续。

4.3 为什么这不是 trivial 的?

vLLM / TGI / TensorRT-LLM 没有它的两个理由:

  1. FSM 和模型 runner 通常在代码库的不同部分。 mask 由约束库(Outlines、lm-format-enforcer)算出,runner 只看到 mask、看不到约束结构。要知道”下一段 kk 个 token 被强制”,runner 必须理解 FSM 本身,而不是每步消费一个 mask。
  2. FSM 是 per-request 的。 每次结构化生成有自己的 regex,FSM 预处理代价不为零。SGLang 对每个 regex 预处理一次,然后在 batch 里凡是用到同一个 regex 的请求都共享这棵压缩 FSM。论文提到:不做 batch 内 FSM 复用,吞吐会跌 2.4×。

效果:同样硬件、同样模型,JSON 解码吞吐 1.6×。

4.4 局限

技巧只对 强制子序列 有效 —— regex 里没有分支的部分。像 "name": "<anything>" 这种自由字段,该一次 forward 一个 token 还是一次一个。收益集中在长 literal 段:JSON 分隔符、schema key、固定 boilerplate。

如果负载是 大部分自由文本(比如自然语言摘要 + 一个小 JSON 包壳),每请求收益不会很大;如果负载有大量结构骨架(深嵌 JSON、有固定脚手架的代码生成),收益就大。论文的 JSON 解码 benchmark 正好夹在中间。

5. API 推测执行

第 5 节,全文最短,只适用于 API-only 的模型(GPT-4 等)。概念干净。

假设程序长这样:

s += context + "name:" + gen("name", stop="\n") + "job:" + gen("job", stop="\n")

朴素做法是 两次 API 调用,每次都被收费 input 包含 context。带上 API 推测执行,第一次调用 忽略 stop="\n" 参数,让模型在 name 之后继续生成。Interpreter 把那段延续与下一段模板(job:)对一下,若模型猜对了,就把第二次 gen 合并进第一次。两次调用变一次,context 的 input token 费用只交一次,不是两次。

speculation 准不准依赖 prompt 工程:如果 context 强暗示模型在 name 后接 job:,命中率高。论文报告:在 Wikipedia 字段抽取任务上,加上 few-shot,input token 成本降低 3×。

本质是把”针对模板的推测解码”从模型内部(那里需要拿 logit)挪到了 API/模板层(那里只需要响应文本)。贡献不大,但有用。

6. 实验

第 6 节。SGLang 用 PyTorch 实现,关键 CUDA kernel 来自 FlashInfer 和 Triton。

6.1 设置

  • 模型: Llama-2 7B / 13B / 70B(dense)、Mixtral-8x7B(MoE)、LLaVA-v1.5-7B(图像)、LLaVA-NeXT-34B(视频)、GPT-3.5(API)。
  • 硬件: 主要是 AWS EC2 g5(A10G,24 GB),部分实验在 A100(80 GB);≥13B 模型用张量并行。
  • 基线: Guidance v0.1.8(llama.cpp 后端)、vLLM v0.2.5(默认 API server)、LMQL v0.7.3(HF Transformers 后端)。
  • 负载: 共 11 个 —— 5-shot MMLU、20-shot HellaSwag、ReAct agents、generative agents、Tree-of-Thought(GSM8K)、Skeleton-of-Thought(tip 生成)、LLM Judge(branch-solve-merge)、JSON decoding、Multi-Turn Chat(short / long output)、DSPy RAG 管线。

6.2 Llama-7B 结果

论文 Figure 5(吞吐)和 Figure 6(延迟)。SGLang 在 所有有前缀共享的负载 上都赢:

  • MMLU。 few-shot 前缀共享,命中率非常高;SGLang 比 vLLM 快约 3.4×。
  • HellaSwag。 两级共享 —— few-shot 样例 + 每题的多选前缀;命中率最强,SGLang ~5–6×。
  • ReAct / generative agents。 agent 模板 + 上一轮上下文共享;收益来自跨轮 KV 复用。
  • Tree-of-Thought / Skeleton-of-Thought。 程序内部 fork;收益来自 RadixAttention + 程序内并行。
  • JSON decoding。 压缩 FSM 挑大梁,比逐 token FSM baseline 快 1.6×。
  • Multi-Turn Chat(短输出)。 聊天历史跨轮共享;KV 复用降低首 token 延迟。
  • Multi-Turn Chat(长输出)。 几乎无加速 —— decode 占主导,共享只能影响 prefill。
  • DSPy RAG 管线。 上下文样例跨 pipeline 调用共享;SGLang 透明地加速 DSPy 后端。

“长输出 chat 几乎没加速”是诚实的负面披露,它精确告诉你这种技术的失效边界:任何 prefill 占端到端时间比重小的地方,前缀共享都救不了你。

6.3 Mixtral-8x7B 与 Llama-70B

Figure 7 与附录 Figure 12 —— 多 A10G 上的张量并行。加速幅度 1.5–4×。论文把 LMQL 和 Guidance 从这组对比里剔除了,因为它们对 TP 支持不够。SGLang 的 TP scaling 与底层引擎一致(vLLM 是唯一同档基线),所以这组对比在该范围内是公平的。

6.4 多模态 LLaVA

Table 2:LLaVA-v1.5-7B 图像 benchmark 上 6.4× 吞吐,LLaVA-NeXT-34B 视频 5×。机制:把输入图像 hash,以 hash 作为 radix 树的 key,凡是同一张图被多次询问(在 llava-bench-in-the-wild 里很常见),视觉编码器对应的图像 token 的 KV 缓存就被复用。

这是一次漂亮的泛化:“前缀”不一定是语言 token,视觉 token 也是可缓存的前缀。

6.5 Chatbot Arena 生产数据

最强的外部验证:一个月的真实 Chatbot Arena 流量。

  • LLaVA-Next-34B:52.4% 缓存命中率
  • Vicuna-33B:74.1% 缓存命中率,平均首 token 延迟降低 1.7×

这些命中率不是 controlled benchmark 出来的,而是真实用户 prompt 杂乱分布下的结果。“超过一半的 token 位置都命中”这件事,本身就说明了:前缀共享是真实 LM 负载的基本性质,不是 benchmark 设计的人造现象。

6.6 消融(Figure 8)

三个子图,各讲一个故事:

  • 命中率 vs 延迟 / 吞吐。 几乎线性:命中率↑ → batch size↑ → 吞吐↑ → 首 token 延迟↓。这张图是 人为节流命中率 画出来的,所以把”缓存命中的效应”从其他变量里干净地剥出来了。
  • 命中率 vs 首 token 延迟。 同一个故事从延迟侧讲。tree-of-thought 上,每多 10 个百分点命中率,TTFT 大概省 25 ms。
  • 组件消融。 分别关闭 tree structure、cache-aware scheduling、frontend parallelism、frontend hint,每一个单独都让 Tree-of-Thought / LLM-Judge / MMLU / Multi-Turn Chat(short)的吞吐掉 10%–30%。四个加起来才有头条数字。

组件消融是最重要的一张图。否则很容易把全部收益归到”加了 prefix caching”,而消融证明每一块的贡献都不可替代。

6.7 没复用时的开销

前面提过:ShareGPT 上 0.27% 开销。这就是”RadixAttention 默认开启”在工程上能成立的关键结果。

7. 讨论

7.1 这篇论文做对了什么

  1. 协同设计。 这是少有的”前端语言和运行时被一起设计、而非互相 retrofit”的论文。没有前端的 fork() 和共享前缀语义,RadixAttention 仍能用,但 hint 机制和并行性会丢;没有这个运行时,DSL 还是能跑,但不会比 LMQL 快。
  2. 与其他引擎优化的可组合。 连续批处理、PagedAttention、张量并行、FlashAttention —— RadixAttention 跟它们都正交,论文也明说能组合。
  3. 诚实的负面结果。 Multi-Turn Chat(长输出)上 SGLang 与 vLLM 基本打平这件事,跟正面结果摆在同一张 Figure 5 里,没有藏。这在论文里挺少见。
  4. radix 树这个抽象。 一旦你看到它,就回不去了:LM 程序的前缀共享问题和路由表的最长匹配问题是同一个问题。radix 树对这两个问题都是对的数据结构。

7.2 我觉得被低估的部分

  • 公平性故事。 论文说缓存感知调度可能饿死请求,这是真的,但没给一个公平 baseline。生产环境里这件事很重要 —— Chatbot Arena 不可能上一个会饿死用户的调度器。后续工作 [42] 在补这一块。
  • compiler 模式。 论文提了 compiler,但结果只放在 Appendix D。对延迟敏感的单程序负载,compiler 应该能帮很多;论文低估了它的潜力。
  • radix 树在图像/视频里的角色。 Table 2 那段只占一小段,但它把这个技术泛化到了非文本模态 —— 多数读者第一遍读会漏掉这条线索。

7.3 局限

  • 没有模糊 / 语义匹配。 两段”意思相同但差一个 token”的 prompt 不能共享 KV。论文把这一点列为后续工作;最近的 CacheBlend、RAG-Cache 等用 embedding 代理在 attack 它。
  • 没有 DRAM / 磁盘扩展。 树完全活在 HBM,大型 prompt 语料(比如 200k 上下文的 RAG 存储)会希望能 spill 到主机内存或 NVMe。论文 §8 指向了这一方向。
  • 贪心带来的饿死风险。 同上。
  • 对长输出负载收益有限。 decode 主导的地方,RadixAttention 只能影响小小的 prefill 那一段。缓解方案(speculative decoding 等)在论文范围之外。
  • API 推测执行依赖 prompt 工程。 模型不按模板走时,speculation 是白花的字节。论文没细测失败模式。
  • 前端只支持 Python。 嵌入 JavaScript、Rust 等宿主语言不平凡,因为 interpreter 依赖 Python 的线程模型。

7.4 它改变了什么

论文出来 18 个月后,SGLang 已经和 vLLM、TensorRT-LLM 一起成为主流开源服务栈之一。RadixAttention(以各种形式)被 port 到了 vLLM 的 prefix caching;压缩 FSM 想法被 port 到 Outlines 和 TensorRT-LLM 的结构化输出模式;前端 DSL 在多家实验室(Together、DeepSeek、Mistral)被当成 agent workflow 的编排层。

更广的概念影响是:这篇论文帮人们把”LLM serving 不只是推理,而是对一类程序的解释,只不过 forward pass 是神经网络”这件事讲清楚了。一旦你接受这个 framing,正确的抽象就不是”请求”,而是”调用图”;正确的原语就不是”KV 缓存”,而是”共享 extent 的树”。这两个抽象已经在它之后几乎每篇 serving 论文里显式出现。

8. 在文献里的位置

为了自己整理方便,我把 SGLang 放到我最近在博客系列里复盘过的 LLM-systems 论文坐标里:

  • vLLM(PagedAttention,2023) —— 修单请求 KV 缓存碎片。SGLang 修跨请求 KV 缓存复用。两者互补,产线同时使用。
  • Orca(2022) —— 引入迭代级别批处理。SGLang 沿用 Orca 的批处理循环,在上面叠加前缀感知调度。
  • Sarathi-Serve(OSDI 2024) —— chunked prefill + 无停顿混合 batch 解决吞吐-延迟权衡。与 SGLang 正交:可以同时做 chunked prefill 和 RadixAttention。最近版 SGLang 已经把 Sarathi 风格调度合进去了。
  • DistServe(OSDI 2024) —— prefill 和 decode 解耦到两组 worker。设计点不同;SGLang 是单池,DistServe 是双池。权衡是 interconnect 带宽 vs 调度器复杂度。
  • HydraGen(2024)ChunkedAttention(2024) —— 共享前缀 attention 的 CUDA kernel 优化。比 SGLang 更底层;SGLang 在它们之上,决定 哪些 该共享,而不是 怎么 算共享后的 attention。
  • Prompt Cache(2023) —— 模块化 prompt 缓存,刻意牺牲精度。SGLang 的精确复用无损;Prompt Cache 拿精度换更高复用率。
  • LMQL(2023)、Guidance(2023) —— 前端 DSL 的前辈。SGLang 借用了它们的原语,但把重投入压在了运行时上,这是别人没动的杠杆。
  • DSPy(2023) —— 高层框架。SGLang 是后端;DSPy 下沉到 SGLang。
  • FlexGen(ICML 2023) —— 单 GPU 高吞吐推理,靠 offloading。正交;SGLang 假设全部装得下 HBM,FlexGen 不假设。

我能给的最短一行小结是:SGLang 之于 LM 程序,就像 JIT 编译器之于动态语言。前端捕获程序结构,运行时挖掘结构去跳过工作。 一旦这个类比对上,每个具体贡献 —— RadixAttention、压缩 FSM、API 推测执行、前端 hint —— 都自然变成”利用已知结构跳过冗余工作”的某种实例。

9. 收尾感想

把这篇论文跟 Sarathi-Serve、DistServe 连着读之后,我对 LLM-serving 栈的方向有了更清晰的图像。三层正在结晶:

  1. Kernel 层。 FlashAttention、FlashInfer、HydraGen —— 尽可能快的原语。
  2. Engine 层。 vLLM、TensorRT-LLM、SGLang Runtime —— 把 kernel 编排成迭代级别批处理、配 paged 内存。
  3. Program / Orchestrator 层。 SGLang 前端、DSPy、LangChain、LlamaIndex —— 描述 engine 应该做什么工作。

SGLang 是我读过的、唯一同时对 (2) 和 (3) 都做出实质贡献,并且把它们当成一个统一设计问题的论文。不管 SGLang 代码库长期会不会赢下 engine 战争(vLLM 也跑得很快),它的概念结论是永久的:程序是一等公民,前缀会被共享,按树来调度

我自己最关心的下一个未答问题是:“前缀”被泛化到”语义前缀”会发生什么? 两段意思相同但差一个同义词的 prompt,今天 token 上没有公共前缀,也就没有 KV 共享。如果我们能对语义相似的前缀算出可接受精度损失的近似 KV 复用(像 Prompt Cache 那样),真实 chatbot 流量上的命中率有可能从 74% 涨到 >90%。这就是下一个 2× 的吞吐,而 SGLang 的 radix 树是显然该被插入这个能力的地方。


代码与可复现 benchmark:https://github.com/sgl-project/sglang。论文:https://arxiv.org/abs/2312.07104