SGLang：结构化语言模型程序的高效执行
单位：
¹ 斯坦福大学（Stanford University）
² 加州大学伯克利分校（UC Berkeley）
³ 上海交通大学（Shanghai Jiao Tong University）
⁴ 德州农工大学（Texas A&M University）
⁵ 独立研究员（Independent Researcher）
摘要
大语言模型（LLM）正越来越多地用于需要多次生成调用、高级提示技术、控制流以及结构化输入/输出的复杂任务。然而，目前缺乏能够高效编程和执行此类应用的系统。我们提出了 SGLang ，一个用于高效执行复杂语言模型程序的系统。SGLang 由一门前端语言和一个运行时组成。前端语言通过用于生成与并行控制的原语简化了编程；运行时通过若干新颖的优化加速执行，例如用于 KV 缓存复用的 RadixAttention，以及用于加速结构化输出解码的压缩有限状态机（compressed finite state machine）。实验表明，在智能体控制、逻辑推理、少样本学习基准、JSON 解码、检索增强生成流水线以及多轮对话等任务上，SGLang 在多种大语言模型和多模态模型上取得了比最先进推理系统高达 6.4 倍 的吞吐量。代码已开源：https://github.com/sgl-project/sglang。
1 引言
近年来，LLM 能力的提升拓宽了它们的用途，使其能够处理更广泛的通用任务，并充当自主智能体（autonomous agents）36, 11, 47, 63, 55。在这类应用中，LLM 会参与多轮规划、推理以及与外部环境的交互。这通过工具使用 50, 45、多种输入模态 56, 2 以及种类繁多的提示技术 38 来实现，例如少样本学习（few-shot learning）10、自洽性（self-consistency）62、思维骨架（skeleton-of-thought）43 和思维树（tree-of-thought）66 等。所有这些新用例都需要多次、且往往相互依赖的 LLM 生成调用，呈现出用多调用结构完成复杂任务的趋势 67, 29。
这些模式的出现标志着我们与 LLM 交互方式的转变：从简单的对话转向一种更复杂的、对 LLM 进行编程化使用的形式，即使用程序来调度和控制 LLM 的生成过程。我们将这类程序称为"语言模型程序"（Language Model Programs，简称 LM 程序）7, 28。上述高级提示技术和智能体工作流都属于 LM 程序的范畴。LM 程序具有两个共同特性：
- LM 程序通常包含多个与控制流交织在一起的 LLM 调用。这是完成复杂任务、提升整体质量所必需的。
- LM 程序接收结构化输入并产生结构化输出。这是实现 LM 程序组合、以及将 LM 程序集成到现有软件系统中所必需的。
尽管 LM 程序已被广泛使用，但当前用于表达和执行它们的系统仍然低效。我们识别出高效使用 LM 程序所面临的两大主要挑战：
第一，由于 LLM 的非确定性本质，编写 LM 程序既繁琐又困难。 开发一个 LM 程序往往需要大量的字符串操作、对提示进行实验性调优、脆弱的输出解析、处理多种输入模态，以及实现并行机制。这种复杂性显著降低了即使是简单程序的可读性（见第 2 节）。
第二，也是更重要的一点，由于存在冗余计算和内存浪费，执行 LM 程序是低效的。 最先进的推理引擎（如 vLLM 33、TGI 22 和 TensorRT-LLM 44）在不了解工作负载的情况下进行了降低延迟、提升吞吐量的优化。这使得这些系统具有通用性和鲁棒性，但对任何给定的工作负载也会带来显著的低效。一个突出的例子是键值（Key-Value，KV）缓存的复用（见第 3 节）。KV 缓存由可复用的中间张量组成，对生成式推理至关重要。在 LM 程序典型的批量执行过程中，共享公共前缀的多个不同 LLM 调用之间存在大量复用 KV 缓存的机会。然而，当前系统缺乏实现这种复用的有效机制，导致不必要的计算和内存浪费。另一个例子是面向结构化输出（如 JSON 模式）的受限解码（constrained decoding），其中 LLM 的输出被限制为遵循由正则表达式定义的特定语法规则（见第 4 节）。在这些约束下，往往可以一次解码多个 token。然而，现有系统一次只能解码一个 token，导致解码速度不佳。
图 1：系统架构：解释器执行语言原语，并配以优化后的运行时。
为应对这些挑战，我们提出了 SGLang ------一种面向 LLM 的结构化生成语言 （S tructured G eneration Language）。其核心思想是系统性地利用 LM 程序中的多调用结构来实现高效执行。如图 1 所示，它由两部分组成：前端语言和后端运行时。前端简化了 LM 程序的编程，运行时则加速其执行。这两部分可以协同工作以获得更好的性能，也可以独立运作。
我们将 SGLang 介绍为一门嵌入在 Python 中的领域特定语言（DSL）。它提供了用于生成的原语（如 extend、gen、select）和用于并行控制的原语（如 fork、join）。SGLang 与 Python 的控制流和库兼容，因此用户可以使用原生 Python 语法轻松开发高级提示工作流。我们为 SGLang 提供了解释器和编译器。解释器将提示状态作为流（stream）来管理，并将原语操作提交到流中异步执行，从而确保对同步和程序内并行的正确控制。此外，SGLang 程序还可以被追踪（trace）并编译，以获得更多优化。
在运行时方面，我们提出了若干新颖的优化来加速 SGLang 程序的执行。第一项技术 是 RadixAttention，它实现了跨多次生成调用的 KV 缓存自动复用。在现有的推理引擎中，请求的 KV 缓存在处理完成后即被丢弃，导致 KV 缓存无法跨多次调用复用，显著拖慢执行速度。与之不同，我们的系统在一棵基数树（radix tree）中为所有请求维护 KV 缓存的 LRU 缓存。该方法将 KV 缓存像传统缓存一样管理，并利用基数树实现高效的匹配、插入和驱逐。它使运行时能够借助缓存感知（cache-aware）的调度策略高效处理各种复用模式。第二项技术 是压缩有限状态机，它为结构化输出实现更快的受限解码。现有系统仅针对下一个 token 遵循约束，通过屏蔽不允许 token 的概率来实现，因此一次只能解码一个 token。与之不同，我们的系统分析约束并构建压缩有限状态机来表示约束。该方法在可能的情况下将多 token 路径压缩为单步路径，从而允许一次解码多个 token，实现更快的解码速度。最后，SGLang 还支持 OpenAI GPT-4 等仅提供 API 的模型，为此我们引入了第三项技术------API 投机执行（API speculative execution），以优化面向仅 API 模型的多调用程序。
使用 SGLang，我们实现了多种 LLM 应用，包括智能体控制、逻辑推理、少样本学习基准、JSON 解码、检索增强生成流水线、多轮对话以及多模态处理。我们在 NVIDIA A10G 和 A100 GPU 上测试了包括 Llama-7B/70B 58、Mistral-8x7B 25、LLaVA-v1.5-7B（图像）39 和 LLaVA-NeXT-34B（视频）69 在内的模型。实验结果表明，与现有的编程和推理系统（包括 Guidance 19、vLLM 33 和 LMQL 7）相比，SGLang 在广泛的工作负载、模型和硬件配置上取得了高达 6.4 倍的吞吐量提升。
2 编程模型
本节通过一个贯穿全文的示例介绍 SGLang 编程模型，描述其语言原语与执行模式，并概述运行时优化机会。该编程模型通过提供灵活、可组合的原语，简化了多调用工作流中繁琐的操作（如字符串操作、API 调用、约束指定、并行化）。
图 2：在 SGLang 中实现的多维度文章评判程序，使用了 branch-solve-merge 提示技术 51。SGLang 提供的原语以红色标出。
一个贯穿全文的示例。 该语言是嵌入在 Python 中的领域特定语言。图 2 展示了一个使用 branch-solve-merge 提示方法 51 评估一篇关于某张图片的文章的程序。函数 multi_dimensional_judge 接收三个参数：s、path 和 essay。s 管理提示状态，path 是图像文件路径，essay 是文章文本。可以使用 += 运算符将新字符串和 SGLang 原语追加到状态 s 中以供执行。该函数首先将图像和文章加入提示，然后使用 select 检查文章是否与图像相关，并将结果存入 s["related"]。如果相关，提示会被分叉（fork）成三个副本，从不同维度并行评估，并使用 gen 将结果存入 f["judgment"]。接着，程序合并各评判结果、生成摘要并给出字母等级评分。最后，程序以 JSON 格式返回结果，遵循由正则表达式约束 regex 定义的模式（schema）。SGLang 极大地简化了这个程序------若使用类似 OpenAI API 的接口编写等价程序，由于需要手动进行字符串操作和并行控制，代码行数将是前者的 2.1 倍。
语言原语。 SGLang 提供用于控制提示状态、生成和并行的原语。它们可以与 Python 语法和库一起使用。这些原语如下：
gen：调用模型进行生成，并将结果存入以其第一个参数指定的名称命名的变量中。它支持regex参数，用于约束输出遵循由正则表达式定义的语法（如某个 JSON 模式）。select：调用模型从一个列表中选择概率最高的选项。- 运算符
+=或extend：向提示追加一个字符串。 - 运算符
[variable_name]：获取某次生成的结果。 fork：创建提示状态的并行分叉。join：重新汇合提示状态。image和video：接收图像和视频输入。
执行模式。 执行 SGLang 程序最简单的方式是通过解释器，其中提示被当作异步流处理。extend、gen 和 select 等原语被提交到流中异步执行。这些非阻塞调用允许 Python 代码继续运行而无需等待生成完成，类似于异步启动 CUDA 内核。每个提示由后台线程中的流执行器管理，从而实现程序内并行。获取生成结果时会阻塞直至结果就绪，以确保正确的同步。另一种方式是将 SGLang 程序编译为计算图，并用图执行器执行，从而允许更多优化。本文默认使用解释器模式，并在附录 C 节（第 9 节）中讨论编译器模式的结果。SGLang 通过自己的 SGLang Runtime（SRT）支持开放权重模型，也支持 OpenAI 和 Anthropic 等 API 模型。
对比。 面向 LLM 的编程系统可分为高层系统（如 LangChain、DSPy）和低层系统（如 LMQL、Guidance、SGLang）。高层系统提供预定义或自动生成的提示，例如 DSPy 的提示优化器。低层系统通常不修改提示，但允许直接操作提示和原语。SGLang 是类似 LMQL 和 Guidance 的低层系统。表 1 比较了它们的特性。SGLang 更关注运行时效率，并自带协同设计的运行时，从而允许后文介绍的新颖优化。高层语言（如 DSPy）可以被编译到低层语言（如 SGLang）。我们在第 6 节（评测）中演示了将 SGLang 作为后端集成到 DSPy 中，以获得更好的运行时效率。
运行时优化。 图 2 展示了三个运行时优化机会：KV 缓存复用、快速受限解码、API 投机执行。我们将在接下来的章节中讨论它们。
表 1：LMQL、Guidance 与 SGLang 之间的比较。
| 系统 | 语法 | 语言原语 | 运行时后端 |
|---|---|---|---|
| LMQL | 自定义 | extend、gen、select | HF Transformers、llama.cpp、OpenAI |
| Guidance | Python | extend、gen、select、image | HF Transformers、llama.cpp、OpenAI |
| SGLang | Python | extend、gen、select、image、video、fork、join | SGLang Runtime（SRT）、OpenAI |
3 基于 RadixAttention 的高效 KV 缓存复用
SGLang 程序可以串联多次生成调用，并可用 fork 原语创建并行副本。此外，不同的程序实例往往共享一些公共部分（如系统提示）。这些场景在执行过程中产生了大量共享的提示前缀，带来众多 KV 缓存复用的机会。
在 LLM 推理过程中，KV 缓存存储前向传播的中间张量，以供解码后续 token 时复用。它们以自注意力机制 60 中的键值对命名。KV 缓存的计算仅依赖于前缀 token。因此，具有相同提示前缀的请求可以复用 KV 缓存，从而减少冗余计算和内存占用。更多背景介绍和示例见附录 A（第 8 节）。
鉴于 KV 缓存的复用机会，优化 SGLang 程序的一个关键挑战是跨多次调用和多个实例复用 KV 缓存。虽然一些系统探索了某些 KV 缓存复用场景 33, 65, 24, 17，但它们往往需要手动配置，且无法处理所有复用模式（如动态树结构）。因此，大多数最先进的推理系统对每个请求都重新计算 KV 缓存。我们将在第 7 节（相关工作）讨论它们的局限性以及与我们的区别。
本节介绍 RadixAttention------一种在运行时自动、系统地复用 KV 缓存的新颖技术。与现有系统在生成请求完成后即丢弃 KV 缓存不同，我们的系统将提示和生成结果的缓存保留在一棵基数树中，从而实现高效的前缀搜索、复用、插入和驱逐。我们实现了 LRU 驱逐策略和缓存感知调度策略，以提升缓存命中率。RadixAttention 与连续批处理（continuous batching）68、分页注意力（paged attention）33 和张量并行（tensor parallelism）52 等技术兼容。此外，在没有缓存命中时，它只引入可忽略不计的内存和时间开销。
RadixAttention。 基数树是一种数据结构，是经典字典树（trie，前缀树）的一种节省空间的替代方案。与典型的树不同，基数树的边不仅可以标注单个元素，还可以标注不同长度的元素序列，从而显著提高效率。在我们的系统中，我们利用基数树管理 token 序列与其对应 KV 缓存张量之间的映射。这些 KV 缓存张量以非连续的分页布局存储，每页大小相当于一个 token。由于 GPU 显存很快会被 KV 缓存填满，我们引入了一个简单的 LRU 驱逐策略，优先驱逐最近最少使用的叶子节点。通过优先驱逐叶子节点，我们使其公共祖先能够被继续复用，直到这些祖先自身也变成叶子并被驱逐。
在连续批处理场景下，我们不能驱逐当前运行批次正在使用的节点。因此，每个节点维护一个引用计数器，指示有多少正在运行的请求在使用它。引用计数为零的节点才可被驱逐。注意，我们并不预先分配固定大小的内存池作为缓存，而是让缓存的 token 与当前运行的请求共享同一内存池。因此，系统会为缓存和运行中的请求动态分配内存。当有足够多的等待请求开始运行时，系统将驱逐所有缓存的 token，以换取更大的批次大小。
图 3 展示了面对若干传入请求时基数树的维护过程。前端解释器将完整提示发送给运行时，运行时执行前缀匹配与复用。树结构存储在 CPU 上，维护开销可以忽略不计。在执行 fork 原语期间，前端会先发送前缀作为提示（hint），确保前缀被正确插入树中，然后再发送其余提示。这种"前端提示"简化了运行时的调度和匹配，体现了前端与运行时协同设计的好处。
图 3：采用 LRU 驱逐策略的 RadixAttention 操作示例，横跨九个时间点进行展示。 该图演示了基数树响应各种请求的动态演化过程。这些请求包括两个聊天会话、一批少样本学习查询，以及一次自洽性采样。每条树边带有一个标签，表示一个子字符串或一个 token 序列。节点以颜色编码以反映不同状态：绿色表示新添加的节点，蓝色表示在该时间点被访问的缓存节点，红色表示已被驱逐的节点。
- 步骤 (1)：基数树初始为空。
- 步骤 (2)：服务器处理传入的用户消息 "Hello"，并以 LLM 输出 "Hi" 回应。系统提示 "You are a helpful assistant"、用户消息 "Hello!" 和 LLM 回复 "Hi!" 被合并为一条边，链接到一个新节点并入树。
- 步骤 (3)：一个新提示到达，服务器在基数树中找到该提示的前缀（即对话的第一轮）并复用其 KV 缓存。新的一轮对话作为新节点追加到树中。
- 步骤 (4)：一个新的聊天会话开始。来自 (3) 的节点 "b" 被拆分为两个节点，使两个聊天会话能够共享系统提示。
- 步骤 (5)：第二个聊天会话继续。然而，由于内存限制，来自 (4) 的节点 "c" 必须被驱逐。新的一轮追加在 (4) 中节点 "d" 之后。
- 步骤 (6)：服务器收到一个少样本学习查询，处理它并将其插入树中。根节点被拆分，因为新查询与现有节点不共享任何前缀。
- 步骤 (7)：服务器收到一批额外的少样本学习查询。这些查询共享同一组少样本示例，因此我们拆分来自 (6) 的节点 "e" 以实现共享。
- 步骤 (8)：服务器收到来自第一个聊天会话的新消息。它驱逐了第二个聊天会话的所有节点（节点 "g" 和 "h"），因为它们最近最少被使用。
- 步骤 (9)：服务器收到一个请求，要求为 (8) 中节点 "j" 的问题采样更多答案，很可能是用于自洽性提示。为给这些请求腾出空间，我们驱逐了 (8) 中的节点 "i"、"k" 和 "l"。
缓存感知调度。 我们将缓存命中率定义为：
缓存命中率 = 被缓存的提示 token 数 ÷ 提示 token 总数
当等待队列中有许多请求时，它们的执行顺序会显著影响缓存命中率。例如，如果请求调度器频繁地在不同的、不相关的请求之间切换，就会导致缓存抖动（thrashing）和较低的命中率。我们设计了一种缓存感知调度算法来提高缓存命中率。在批处理场景下，我们按匹配前缀长度对请求排序，优先处理匹配前缀较长的请求，而不是采用先来先服务（FCFS）的调度方式。附录中的算法 1 给出了连续批处理下缓存感知调度的伪代码，该算法采用最长共享前缀优先的顺序。在延迟更敏感的场景中，我们仍然可以容忍有限的批次重排序，以改善缓存复用。此外，我们针对离线情形的最优调度证明了以下定理。（注：在实践中，实际计算与定理 1 证明中描述的并不完全相同，因为输出 token 数量不可预测，可能导致 KV 缓存的重计算。）
定理 1 对于一批请求，通过以深度优先搜索（DFS）的顺序访问由这些请求构成的基数树，并令缓存大小 ≥ 最大请求长度，即可达到最优的缓存命中率。最长共享前缀优先的顺序等价于深度优先搜索顺序。
证明见附录 A.3 节。在在线情形下，DFS 顺序会被打破，但我们的调度仍然在完整基数树的增量部分上近似 DFS 行为，如附录 A.3 所述。虽然贪心的缓存感知调度可以实现高吞吐量，但它可能导致饥饿（starvation）问题。我们将其与其他公平调度方法 53 的结合留作未来工作。
分布式场景。 RadixAttention 可以扩展到多块 GPU。对于张量并行，每块 GPU 维护一个分片的 KV 缓存。由于树操作完全相同，无需额外同步。多工作节点的数据并行在附录 A.4 节中讨论。
4 基于压缩有限状态机的高效受限解码
在 LM 程序中，用户常常希望约束模型的输出遵循特定格式，例如 JSON 模式。这可以提升可控性和鲁棒性，并使输出更易于解析。SGLang 提供 regex 参数，使用正则表达式来强制施加此类约束，正则表达式对许多实际场景而言已具有足够的表达力。现有系统通过将正则表达式转换为有限状态机（FSM）来支持这一功能 64。在解码过程中，它们维护当前 FSM 状态，从后继状态中检索允许的 token，并将无效 token 的概率置为零，逐 token 解码。然而，当存在一次解码多个 token 的机会时，这种逐 token 的方式是低效的。例如，图 2 中的常量序列 {"summary": " 在正常解码过程中跨越多个 token（如图 4© 所示），需要多个解码阶段，尽管解码该序列时每一步只有一个有效的后继 token。因此，整个序列可以在单步（即一次前向传播）中完成解码。然而，现有系统一次只能解码一个 token，因为现有系统中 FSM 与模型执行器之间缺乏集成，无法进行多 token 处理，导致解码缓慢。
图 4：普通 FSM 与压缩 FSM 的解码过程（下划线 _ 表示空格）。
SGLang 通过构建一个基于压缩 FSM 的快速受限解码运行时克服了这一局限。该运行时分析 FSM，并将 FSM 中相邻的单一转移边压缩为单条边（如图 4(b) 所示），使其能够识别出何时可以一起解码多个 token。在图 4(d) 中，压缩转移边上的多个 token 可以在一次前向传播中完成解码，极大地加速了解码过程。该方法具有通用性，适用于所有正则表达式。更多背景和实现细节见附录 B（第 9 节）。
5 基于 API 投机执行的高效端点调用
前面几节介绍了针对开放权重模型的优化，这些优化需要修改模型推理过程。此外，SGLang 也适用于仅提供 API 访问的模型，例如 OpenAI 的 GPT-4。然而，对于这些模型，我们只能调用一个黑盒 API 端点。
本节介绍一种针对黑盒 API 模型的新优化，它通过投机执行（speculative execution）加速多调用 SGLang 程序的执行并降低 API 成本。例如，某个程序可能要求模型以多调用模式生成对一个人物的描述：
s += context + "name:" + gen("name", stop="\n") + "job:" + gen("job", stop="\n")朴素地执行时，两个 gen 原语对应两次 API 调用，这意味着用户需要为 context 的输入 token 费用支付两次。在 SGLang 中，我们可以在第一次调用时启用投机执行，通过忽略停止条件让它继续多生成一些 token。解释器保留这些额外的生成输出，并将它们与后续原语进行匹配和复用。在某些情况下，通过精心的提示工程，模型可以以很高的准确率正确匹配模板，从而为我们节省一次 API 调用的延迟和输入成本。
6 评测
我们在多种 LLM 工作负载上评测了 SGLang 的性能。随后，我们进行了消融实验和案例研究，以展示特定组件的有效性。SGLang 基于 PyTorch 46 实现，使用了来自 FlashInfer 70 和 Triton 57 的自定义 CUDA 内核。
6.1 实验设置
模型。 我们测试了稠密的 Llama-2 模型 58、稀疏混合专家 Mixtral 模型 25、多模态的 LLaVA 图像 37 与视频 69 模型，以及 API 模型 OpenAI GPT-3.5。对于开放权重模型，参数量从 70 亿到 700 亿不等，我们使用 float16 精度。
硬件。 我们在配备 NVIDIA A10G GPU（24GB）的 AWS EC2 G5 实例上运行了大部分实验。我们在单块 A10G GPU 上运行 7B 模型，并使用张量并行 52 在多块 A10G GPU 上运行更大的模型。我们还在 A100（80GB）GPU 上运行了一些补充实验。
基线。 我们将 SGLang 与两类基线进行比较：一是带有各自语言和默认运行时的高层编程系统，二是提供标准 OpenAI 风格 Completion API 的低层推理引擎。除非另有说明，我们不开启会改变计算结果的优化，以便所有系统计算相同的结果。基线包括：
- Guidance 19：一种控制 LLM 的语言。我们使用 Guidance v0.1.8 及 llama.cpp 后端。
- vLLM 33：一个高吞吐量推理引擎。我们使用 vLLM v0.2.5 及其默认 API 服务器。（注：RadixAttention 已作为可选的实验特性被部分集成到最新版本的 vLLM 中；因此我们使用了更早的版本进行对比。）
- LMQL 7：一种查询语言。我们使用 LMQL v0.7.3 及 Hugging Face Transformers 后端。
工作负载。 我们测试了以下内容：5-shot MMLU 21 和 20-shot HellaSwag 71 基准。对 MMLU 我们解码一个 token；对 HellaSwag 我们使用 select 原语选择概率最高的答案。对于 ReAct 智能体 67 和生成式智能体 47，我们从原论文中提取轨迹（traces）并重放。我们对 GSM-8K 问题使用思维树（Tree-of-thought）66，对提示语生成使用思维骨架（Skeleton-of-thought）43。我们使用结合 branch-solve-merge 51 技术的 LLM 评判；使用由正则表达式指定模式的 JSON 解码；进行 4 轮的多轮对话，其中每轮输入在 256--512 token 之间随机采样------多轮对话（短）表示短输出（4--8 token），多轮对话（长）表示长输出（256--512 token）；以及 DSPy 检索增强生成（RAG）流水线 28 的官方示例。
指标。 我们报告两个性能指标：吞吐量和延迟。对于吞吐量，我们运行足够大的程序实例批次以计算最大吞吐量，比较每秒执行的程序实例数（程序数/秒，p/s）。对于延迟，我们一次执行单个程序（不做批处理），并报告多个实例的平均延迟。
图 5：Llama-7B 模型上的归一化吞吐量。越高越好。
图 6：Llama-7B 模型上的归一化延迟。越低越好。
6.2 端到端性能
开放权重模型上的结果。 延迟和吞吐量结果如图 5 和图 6 所示。SGLang 将吞吐量最高提升了 6.4 倍，并将延迟最高降低了 3.7 倍。这些改进来自 KV 缓存复用、对单个程序内并行性的利用，以及更快的受限解码。接下来，我们解释各基准上加速的原因。
在 MMLU 上，SGLang 可以通过 RadixAttention 复用 5-shot 示例的 KV 缓存。RadixAttention 同时有益于吞吐量和延迟：它通过共享 KV 缓存减少总内存占用，从而允许更大的批次以提升最大吞吐量；它还减少了 prefill 的计算量，从而降低首 token 延迟。在 HellaSwag 上，SGLang 既复用少样本示例的 KV 缓存，也复用多个选项共用的问题前缀的 KV 缓存，形成两级共享。对于 ReAct 智能体和生成式智能体，SGLang 复用智能体模板和此前调用的 KV 缓存。在思维树和思维骨架上，SGLang 并行化单个程序内的生成调用，并尽可能多地复用 KV 缓存。在 JSON 解码上，SGLang 通过压缩有限状态机一次解码多个 token 来加速解码。在多轮对话中，SGLang 复用聊天历史的 KV 缓存。对于短输出，加速更为明显，因为 KV 缓存复用主要有助于减少前缀时间；对于长输出，由于不同聊天会话之间没有太多共享，且解码时间占主导，几乎没有加速。在 DSPy RAG 流水线中，SGLang 复用公共上下文示例的 KV 缓存。在这些基准上，缓存命中率范围为 50% 到 99%。附录中的图 12 列出了所有基准上达到的和最优的缓存命中率，显示我们的缓存感知调度平均达到最优命中率的 96%。
由于性能缓慢和功能缺失，我们在最后五个基准中的部分测试里排除了 LMQL 和 Guidance。LMQL 的问题源于缓慢的 token 级处理和未优化的后端，而 Guidance 缺少批处理和并行支持。
图 7：采用张量并行的 Mixtral-8x7B 模型上的归一化吞吐量。越高越好。
采用张量并行的更大模型上的结果。 我们在同一组基准上运行了更大的模型 Mixtral-8x7B 和 Llama-70B（采用张量并行），并在图 7 和附录图 11 中报告结果。更大模型上的加速呈现出与较小模型类似的趋势，表明我们的优化可以很好地泛化到更大的模型。此处省略了 Guidance 和 LMQL，因为它们缺乏高效的张量并行实现。
多模态模型上的结果。 SGLang 通过 image 和 video 原语原生支持多模态模型。本文中的优化与多模态模型兼容。对于 RadixAttention，我们计算输入图像的哈希，并将其作为基数树中的键，从而可以复用来自同一图像的图像 token 的 KV 缓存。我们在 llava-bench-in-the-wild 上运行 LLaVA-v1.5-7B（图像），在 ActivityNet 上运行 LLaVA-NeXT-34B（视频）。由于这些模型未被其他基线系统很好地支持，我们使用模型作者在 Hugging Face Transformers 中的原始实现作为基线。如表 2 所示，SGLang 在这些基准上提供了高达 6 倍的吞吐量。在 llava-bench-in-the-wild 中，针对同一张图像有多个问题，SGLang 运行时在这种情况下复用了 KV 缓存。
表 2：多模态 LLaVA 图像与视频模型上的吞吐量比较。
| 模型 | LLaVA-v1.5-7B（图像） | LLaVA-NeXT-34B（视频） |
|---|---|---|
| 作者原始实现 | 0.18 图像/秒 | 0.02 帧/秒 |
| SGLang | 1.15 图像/秒 | 0.10 帧/秒 |
生产部署。 SGLang 已在 Chatbot Arena 14 中部署，用于服务开放权重模型。由于某些模型流量较低，每个模型仅由一个 SGLang 工作节点服务。一个月后，我们观察到 LLaVA-Next-34B 39 的 RadixAttention 缓存命中率为 52.4%，Vicuna-33B 13 为 74.1%。缓存命中来自公共系统消息、频繁复用的示例图像以及多轮聊天历史。这使 Vicuna-33B 的首 token 延迟平均降低了 1.7 倍。
API 模型上的结果。 我们测试了一个使用 OpenAI GPT-3.5 模型从维基百科页面提取三个字段的提示。通过使用少样本提示，API 投机执行的准确率很高；由于提取了三个字段，它将输入 token 的成本降低了约三分之二（约为原来的三分之一）。
图 8：(a)(b) 缓存命中率消融实验。© RadixAttention 消融实验。
6.3 消融实验
缓存命中率与延迟/吞吐量的关系。 图 8(a)(b) 展示了在思维树基准上缓存命中率与性能指标（首 token 延迟、总延迟、批次大小和吞吐量）之间的关系。该图通过在运行时部分禁用匹配的 token 获得。图中显示，更高的缓存命中率带来更大的批次大小、更高的吞吐量和更低的延迟。
RadixAttention 的有效性。 我们在若干代表性基准上测试了 RadixAttention 及其组件的有效性。如图 8© 所示，"No Cache" 表示不使用任何缓存；"No Tree-Structure" 表示使用简单的基于表的缓存而非树结构缓存；"FCFS Schedule" 表示使用先来先服务策略而非我们的缓存感知调度；"Random Schedule" 表示使用随机顺序调度请求；"No Frontend Parallelism" 表示禁用解释器中的并行；"No Frontend Hint" 表示禁用从解释器发送 fork 提示；"Full optimizations" 表示开启所有优化。实验结果表明，这些组件中的每一个都是达到最佳性能所必需的。禁用来自前端解释器的并行和提示也会导致运行时性能不佳，凸显了前端语言与运行时协同设计的重要性。
RadixAttention 的开销。 我们在一个没有任何 KV 缓存复用机会的基准上测试了 RadixAttention 的开销。该基准测量 ShareGPT 数据集上的吞吐量。运行 100 个请求需要 74.3 秒，而用于管理 RadixAttention 数据结构的时间仅为 0.2 秒，即不到 0.3% 的可忽略开销。这是因为树操作的复杂度是线性且很小的。因此，我们可以默认开启 RadixAttention。
压缩有限状态机的有效性。 我们在 JSON 解码基准上测试了压缩有限状态机及其组件的有效性。实验结果表明，压缩有限状态机将吞吐量提高了 1.6 倍，因为它可以一次解码多个 token。此外，我们需要对状态机进行预处理，并将其复用于一批请求；否则，为每个请求重新进行预处理会使吞吐量降低 2.4 倍。
7 相关工作
已有多项工作探索了 KV 缓存的复用，其中许多与我们的工作同期进行。独特的是，我们的 RadixAttention 首次提出将 KV 缓存视为基于树的 LRU 缓存。它是首个支持多级共享、缓存感知调度、前端-运行时协同调度以及分布式场景的解决方案。vLLM 33 和 ChunkedAttention 65 探索了一些简单的复用场景（如系统提示共享），但未涵盖多级树结构共享或 LRU 缓存。PromptCache 17 提出了超出前缀范围的 KV 缓存模块化复用，但可能影响准确率，最高下降 43%。HydraGen 24、FlashInfer 70 和 ChunkedAttention 65 专注于 CUDA 内核优化，不包含 LRU 缓存的概念。API Serve 1 和 LLM-SQL 40 研究了特定应用（如与外部 API 调用交织、关系数据库）中的 KV 缓存复用，但它们没有我们的基数树或缓存感知调度。
现已有若干 LLM 编程与智能体框架，例如 Guidance 19、LMQL 7、DSPy 28、LangChain 34、AutoGen 64b 和 LLM Compiler 29。Guidance 和 LMQL 与 SGLang 最为相似，我们在第 2 节中对它们进行了比较。我们的创新在于为所提出的编程模型提供新颖的运行时优化以加速执行。SGLang 与其他框架兼容，并能加速它们（例如我们评测中的 DSPy 示例）。此外，SGLang 还与许多其他常见的推理优化技术兼容 68, 48, 3, 33, 70, 12, 35, 23, 26, 42, 41, 16。
8 未来方向与结论
未来方向。 尽管 SGLang 取得了进展，但仍存在一些局限，同时也揭示了有前景的未来研究方向。这些包括：扩展 SGLang 以支持更多输出模态；将 RadixAttention 适配到跨越多个内存层级（如 DRAM、磁盘）运行 54；在 RadixAttention 中实现模糊语义匹配；在 SGLang 之上提供更高层的原语；解决缓存感知调度中的饥饿问题 53；以及增强 SGLang 编译器以执行高级静态优化，如调度和内存规划。
结论。 我们提出了 SGLang------一个用于高效编程和执行结构化语言模型程序的框架。SGLang 通过 RadixAttention、压缩有限状态机和语言解释器等新颖优化，显著提升了复杂 LM 程序的吞吐量和延迟。它是开发高级提示技术和智能体工作流的有价值工具。源代码已公开。
致谢
本项目得到了斯坦福自动推理中心（Stanford Center for Automated Reasoning）以及来自 Astronomer、Google、IBM、Intel、Lacework、Microsoft、穆罕默德·本·扎耶德人工智能大学（Mohamed Bin Zayed University of Artificial Intelligence）、Nexla、Samsung SDS、Uber 和 VMware 的捐赠支持。Lianmin Zheng 由 Meta 博士奖学金资助。我们感谢 Yuanhan Zhang 和 Bo Li 对 LLaVA-NeXT（视频）的支持。
参考文献
1 Reyna Abhyankar, Zijian He, Vikranth Srivatsa, Hao Zhang, Yiying Zhang. Apiserve: Efficient api support for large-language model inferencing. arXiv preprint arXiv:2402.01869, 2024.
2 Jean-Baptiste Alayrac, Jeff Donahue, Pauline Luc, Antoine Miech, Iain Barr, Yana Hasson, Karel Lenc, Arthur Mensch, Katherine Millican, Malcolm Reynolds, et al. Flamingo: a visual language model for few-shot learning. Advances in Neural Information Processing Systems, 35:23716--23736, 2022.
3 Reza Yazdani Aminabadi, Samyam Rajbhandari, Ammar Ahmad Awan, Cheng Li, Du Li, Elton Zheng, Olatunji Ruwase, Shaden Smith, Minjia Zhang, Jeff Rasley, et al. Deepspeed-inference: enabling efficient inference of transformer models at unprecedented scale. In SC22: International Conference for High Performance Computing, Networking, Storage and Analysis, pages 1--15. IEEE, 2022.
7 Luca Beurer-Kellner, Marc Fischer, Martin Vechev. Prompting is programming: A query language for large language models. Proceedings of the ACM on Programming Languages, 7(PLDI):1946--1969, 2023.
10 Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877--1901, 2020.
11 Sébastien Bubeck, Varun Chandrasekaran, Ronen Eldan, Johannes Gehrke, Eric Horvitz, Ece Kamar, Peter Lee, Yin Tat Lee, Yuanzhi Li, Scott Lundberg, et al. Sparks of artificial general intelligence: Early experiments with gpt-4. arXiv preprint arXiv:2303.12712, 2023.
12 Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, Christopher Ré. Flashattention: Fast and memory-efficient exact attention with io-awareness. Advances in Neural Information Processing Systems, 35:16344--16359, 2022.
13 Wei-Lin Chiang, Zhuohan Li, Zi Lin, Ying Sheng, Zhanghao Wu, Hao Zhang, Lianmin Zheng, Siyuan Zhuang, Yonghao Zhuang, Joseph E. Gonzalez, Ion Stoica, Eric P. Xing. Vicuna: An open-source chatbot impressing gpt-4 with 90%* chatgpt quality, March 2023.
14 Wei-Lin Chiang, Lianmin Zheng, Ying Sheng, Anastasios Nikolas Angelopoulos, Tianle Li, Dacheng Li, Hao Zhang, Banghua Zhu, Michael Jordan, Joseph E Gonzalez, et al. Chatbot arena: An open platform for evaluating llms by human preference. arXiv preprint arXiv:2403.04132, 2024.
16 Suyu Ge, Yunan Zhang, Liyuan Liu, Minjia Zhang, Jiawei Han, Jianfeng Gao. Model tells you what to discard: Adaptive kv cache compression for llms. In The Twelfth International Conference on Learning Representations, 2023.
17 In Gim, Guojun Chen, Seung-seob Lee, Nikhil Sarda, Anurag Khandelwal, Lin Zhong. Prompt cache: Modular attention reuse for low-latency inference. arXiv preprint arXiv:2311.04934, 2023.
19 Guidance AI. A guidance language for controlling large language models. https://github.com/guidance-ai/guidance. Accessed: 2023-11.
21 Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, Jacob Steinhardt. Measuring massive multitask language understanding. In International Conference on Learning Representations, 2020.
22 Hugging Face. Text generation inference. https://github.com/huggingface/text-generation-inference. Accessed: 2023-11.
23 Coleman Hooper, Sehoon Kim, Hiva Mohammadzadeh, Michael W Mahoney, Yakun Sophia Shao, Kurt Keutzer, Amir Gholami. Kvquant: Towards 10 million context length llm inference with kv cache quantization. arXiv preprint arXiv:2401.18079, 2024.
24 Jordan Juravsky, Bradley Brown, Ryan Ehrlich, Daniel Y Fu, Christopher Ré, Azalia Mirhoseini. Hydragen: High-throughput llm inference with shared prefixes. arXiv preprint arXiv:2402.05099, 2024.
25 Albert Q Jiang, Alexandre Sablayrolles, Antoine Roux, Arthur Mensch, Blanche Savary, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Emma Bou Hanna, Florian Bressand, et al. Mixtral of experts. arXiv preprint arXiv:2401.04088, 2024.
26 Hao Kang, Qingru Zhang, Souvik Kundu, Geonhwa Jeong, Zaoxing Liu, Tushar Krishna, Tuo Zhao. Gear: An efficient kv cache compression recipe for near-lossless generative inference of llm. arXiv preprint arXiv:2403.05527, 2024.
28 Omar Khattab, Arnav Singhvi, Paridhi Maheshwari, Zhiyuan Zhang, Keshav Santhanam, Sri Vardhamanan, Saiful Haq, Ashutosh Sharma, Thomas T. Joshi, Hanna Moazam, Heather Miller, Matei Zaharia, Christopher Potts. Dspy: Compiling declarative language model calls into self-improving pipelines. arXiv preprint arXiv:2310.03714, 2023.
29 Sehoon Kim, Suhong Moon, Ryan Tabrizi, Nicholas Lee, Michael W Mahoney, Kurt Keutzer, Amir Gholami. An llm compiler for parallel function calling. arXiv preprint arXiv:2312.04511, 2023.
33 Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph Gonzalez, Hao Zhang, Ion Stoica. Efficient memory management for large language model serving with pagedattention. In Proceedings of the 29th Symposium on Operating Systems Principles, pages 611--626, 2023.
34 LangChain AI. Langchain. https://github.com/langchain-ai/langchain. Accessed: 2023-11.
35 Ji Lin, Jiaming Tang, Haotian Tang, Shang Yang, Xingyu Dang, Song Han. Awq: Activation-aware weight quantization for llm compression and acceleration. arXiv preprint arXiv:2306.00978, 2023.
36 OpenAI. Gpt-4 technical report, 2023.
37 Haotian Liu, Chunyuan Li, Yuheng Li, Yong Jae Lee. Improved baselines with visual instruction tuning. arXiv preprint arXiv:2310.03744, 2023.
38 Xiaoxia Liu, Jingyi Wang, Jun Sun, Xiaohan Yuan, Guoliang Dong, Peng Di, Wenhai Wang, Dongxia Wang. Prompting frameworks for large language models: A survey. arXiv preprint arXiv:2311.12785, 2023.
39 Haotian Liu, Chunyuan Li, Yuheng Li, Bo Li, Yuanhan Zhang, Sheng Shen, Yong Jae Lee. Llava-next: Improved reasoning, ocr, and world knowledge, January 2024.
40 Shu Liu, Asim Biswal, Audrey Cheng, Xiangxi Mo, Shiyi Cao, Joseph E Gonzalez, Ion Stoica, Matei Zaharia. Optimizing llm queries in relational workloads. arXiv preprint arXiv:2403.05821, 2024.
41 Zichang Liu, Aditya Desai, Fangshuo Liao, Weitao Wang, Victor Xie, Zhaozhuo Xu, Anastasios Kyrillidis, Anshumali Shrivastava. Scissorhands: Exploiting the persistence of importance hypothesis for llm kv cache compression at test time. Advances in Neural Information Processing Systems, 36, 2024.
42 Zirui Liu, Jiayi Yuan, Hongye Jin, Shaochen Zhong, Zhaozhuo Xu, Vladimir Braverman, Beidi Chen, Xia Hu. Kivi: A tuning-free asymmetric 2bit quantization for kv cache. arXiv preprint arXiv:2402.02750, 2024.
43 Xuefei Ning, Zinan Lin, Zixuan Zhou, Zifu Wang, Huazhong Yang, Yu Wang. Skeleton-of-thought: Prompting LLMs for efficient parallel generation. In The Twelfth International Conference on Learning Representations, 2024.
44 NVIDIA. Tensorrt-llm. https://github.com/NVIDIA/TensorRT-LLM. Accessed: 2023-11.
45 Shishir G. Patil, Tianjun Zhang, Xin Wang, Joseph E. Gonzalez. Gorilla: Large language model connected with massive apis. arXiv preprint arXiv:2305.15334, 2023.
46 Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high-performance deep learning library. Advances in neural information processing systems, 32, 2019.
47 Joon Sung Park, Joseph C. O'Brien, Carrie J. Cai, Meredith Ringel Morris, Percy Liang, Michael S. Bernstein. Generative agents: Interactive simulacra of human behavior. In the 36th Annual ACM Symposium on User Interface Software and Technology (UIST '23), 2023.
48 Reiner Pope, Sholto Douglas, Aakanksha Chowdhery, Jacob Devlin, James Bradbury, Jonathan Heek, Kefan Xiao, Shivani Agrawal, Jeff Dean. Efficiently scaling transformer inference. Proceedings of Machine Learning and Systems, 5, 2023.
50 Timo Schick, Jane Dwivedi-Yu, Roberto Dessì, Roberta Raileanu, Maria Lomeli, Luke Zettlemoyer, Nicola Cancedda, Thomas Scialom. Toolformer: Language models can teach themselves to use tools. arXiv preprint arXiv:2302.04761, 2023.
51 Swarnadeep Saha, Omer Levy, Asli Celikyilmaz, Mohit Bansal, Jason Weston, Xian Li. Branch-solve-merge improves large language model evaluation and generation. arXiv preprint arXiv:2310.15123, 2023.
52 Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, Bryan Catanzaro. Megatron-lm: Training multi-billion parameter language models using model parallelism. arXiv preprint arXiv:1909.08053, 2019.
53 Ying Sheng, Shiyi Cao, Dacheng Li, Banghua Zhu, Zhuohan Li, Danyang Zhuo, Joseph E Gonzalez, Ion Stoica. Fairness in serving large language models. arXiv preprint arXiv:2401.00588, 2023.
54 Ying Sheng, Lianmin Zheng, Binhang Yuan, Zhuohan Li, Max Ryabinin, Beidi Chen, Percy Liang, Christopher Ré, Ion Stoica, Ce Zhang. Flexgen: high-throughput generative inference of large language models with a single gpu. In International Conference on Machine Learning, pages 31094--31116. PMLR, 2023.
55 Theodore Sumers, Shunyu Yao, Karthik Narasimhan, Thomas Griffiths. Cognitive architectures for language agents. Transactions on Machine Learning Research, 2023.
56 Gemini Team, Rohan Anil, Sebastian Borgeaud, Yonghui Wu, Jean-Baptiste Alayrac, Jiahui Yu, Radu Soricut, Johan Schalkwyk, Andrew M Dai, Anja Hauth, et al. Gemini: a family of highly capable multimodal models. arXiv preprint arXiv:2312.11805, 2023.
57 Philippe Tillet, Hsiang-Tsung Kung, David Cox. Triton: an intermediate language and compiler for tiled neural network computations. In Proceedings of the 3rd ACM SIGPLAN International Workshop on Machine Learning and Programming Languages, pages 10--19, 2019.
58 Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09288, 2023.
60 Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
62 Xuezhi Wang, Jason Wei, Dale Schuurmans, Quoc V Le, Ed H Chi, Sharan Narang, Aakanksha Chowdhery, Denny Zhou. Self-consistency improves chain of thought reasoning in language models. In The Eleventh International Conference on Learning Representations, 2022.
63 Guanzhi Wang, Yuqi Xie, Yunfan Jiang, Ajay Mandlekar, Chaowei Xiao, Yuke Zhu, Linxi Fan, Anima Anandkumar. Voyager: An open-ended embodied agent with large language models. arXiv preprint arXiv:2305.16291, 2023.
64 Brandon T Willard, Rémi Louf. Efficient guided generation for large language models. arXiv e-prints, arXiv:2307, 2023.
64b Qingyun Wu, Gagan Bansal, Jieyu Zhang, Yiran Wu, Shaokun Zhang, Erkang Zhu, Beibin Li, Li Jiang, Xiaoyun Zhang, Chi Wang. Autogen: Enabling next-gen llm applications via multi-agent conversation framework. arXiv preprint arXiv:2308.08155, 2023.
65 Lu Ye, Ze Tao, Yong Huang, Yang Li. Chunkattention: Efficient self-attention with prefix-aware kv cache and two-phase partition. arXiv preprint arXiv:2402.15220, 2024.
66 Shunyu Yao, Dian Yu, Jeffrey Zhao, Izhak Shafran, Thomas L Griffiths, Yuan Cao, Karthik Narasimhan. Tree of thoughts: Deliberate problem solving with large language models. arXiv preprint arXiv:2305.10601, 2023.
67 Shunyu Yao, Jeffrey Zhao, Dian Yu, Nan Du, Izhak Shafran, Karthik R Narasimhan, Yuan Cao. React: Synergizing reasoning and acting in language models. In The Eleventh International Conference on Learning Representations, 2022.
68 Gyeong-In Yu, Joo Seong Jeong, Geon-Woo Kim, Soojeong Kim, Byung-Gon Chun. Orca: A distributed serving system for Transformer-Based generative models. In 16th USENIX Symposium on Operating Systems Design and Implementation (OSDI 22), pages 521--538, 2022.
69 Yuanhan Zhang, Bo Li, Haotian Liu, Yong-jae Lee, Liangke Gui, Di Fu, Jiashi Feng, Ziwei Liu, Chunyuan Li. Llava-next: A strong zero-shot video understanding model, April 2024.
70 Zihao Ye, Lequn Chen, Ruihang Lai, Yilong Zhao, Size Zheng, Junru Shao, Bohan Hou, Hongyi Jin, Yifei Zuo, Liangsheng Yin, Tianqi Chen, Luis Ceze. Accelerating self-attentions for llm serving with flashinfer, February 2024.
71 Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, Yejin Choi. Hellaswag: Can a machine really finish your sentence? In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pages 4791--4800, 2019.
其余参考文献（按原论文顺序）：Vivien Tran-Thien. Fast, high-fidelity llm decoding with regex constraints, 2024.；Michael Kuchnik, Virginia Smith, George Amvrosiadis. Validating large language models with relm. Proceedings of Machine Learning and Systems , 5, 2023.；Yujia Li et al. Competition-level code generation with alphacode. Science , 378(6624):1092--1097, 2022.；Vikranth Srivatsa et al. Preble: Efficient distributed prompt scheduling for llm serving, 2024.；Aakanksha Chowdhery et al. Palm: Scaling language modeling with pathways. arXiv preprint arXiv:2204.02311, 2022.
附录
附录 A：RadixAttention 补充细节
A.1 KV 缓存背景
目前使用的大多数 LLM，如 GPT-3 10、PaLM 15 和 LLaMA 58，都基于自回归 Transformer 架构 60。这些模型基于前面的 token 预测序列中下一个 token 的概率。在推理过程中，模型首先通过一次前向传播处理输入 token 序列（此过程称为 "prefill"），然后逐个顺序解码输出 token，每个 token 都依赖于之前的 token（此过程称为 "decoding"）。我们将接收一个输入 token 序列并生成一个输出 token 序列的过程称为一次生成调用（generation call）。在整个过程中，每个 token 都会产生一些中间张量，用于解码后续的 token。这些中间张量被称为"KV 缓存"，以自注意力机制中的键值对命名。
讨论本文优化时的一个重要观察是：KV 缓存的计算仅依赖于之前的所有 token，因此具有相同前缀的不同序列可以复用前缀 token 的 KV 缓存，避免冗余计算。
在 LLM 程序中，常常有多个文本片段和生成调用被追加到同一个提示中。在多个串联调用之间缓存已计算的先前 token 的 KV 缓存，可以减少冗余计算。然而，这种优化既不是免费的，也并非微不足道，因为它需要额外的存储和更复杂的内存管理。此外，在 LLM 程序中，从单个提示生成多个输出，或从当前状态分叉出新提示，都是常见做法 32。基本的前缀共享已在 vLLM 33 中研究过。更高级的共享模式（如不规则的树结构共享）也可以被采用。图 9 展示了跨多次调用共享 KV 缓存的四种典型模式；现有系统都无法自动处理所有这些模式。相反，第 3 节介绍的我们的 RadixAttention 可以在运行时自动处理所有这些模式。
图 9：KV 缓存共享示例。 蓝色框表示可共享的提示部分，绿色框表示不可共享的部分，黄色框标记不可共享的模型输出。可共享的元素包括少样本学习示例、自洽性 62 中的问题、多轮对话中的聊天历史，以及思维树 66 中的搜索历史。
A.2 缓存感知调度伪代码
算法 1 展示了连续批处理下 RadixAttention 缓存感知调度的伪代码。
算法 1：连续批处理下 RadixAttention 的缓存感知调度
输入：基数树 T，内存池 P，当前运行批次 B，等待队列 Q。
输出：已完成的请求和更新后的系统状态。
// 从等待队列获取所有请求
requests ← Q.get_all_requests()
// 为所有等待请求搜索前缀匹配
for req ∈ requests do
req.prefix_node, req.prefix_len ← T.match_prefix(req.input_tokens)
end for
// 按匹配的前缀长度对请求排序
requests.sort()
// 为下一批次选择请求
available_size ← T.evictable_size() + P.available_size()
current_size ← 0
new_batch ← []
for req ∈ requests do
if req.size() + current_size < available_size then
new_batch.append(req)
delta ← T.increase_ref_counter(req.prefix_node)
available_size ← available_size + delta
end if
end for
Q.remove_requests(new_batch)
// 将请求插入当前运行批次
B.merge(new_batch)
// 分配新内存，必要时进行驱逐
needed_size ← B.needed_size()
success, buffer ← P.alloc(needed_size)
if not success then
T.evict(needed_size)
success, buffer ← P.alloc(needed_size)
end if
B.run(buffer)
// 处理已完成的请求
finished_requests ← B.drop_finished_requests()
for req ∈ finished_requests do
T.decrease_ref_counter(req.prefix_node)
T.insert(req)
end for
return finished_requestsA.3 定理 1 的证明
定理 1（重述） 对于一批请求，通过以深度优先搜索（DFS）的顺序访问由这些请求构成的基数树，并令缓存大小 ≥ 最大请求长度，即可达到最优的缓存命中率。最长共享前缀优先的顺序等价于深度优先搜索顺序。
证明： 首先，我们证明深度优先搜索（DFS）顺序可以达到最优的缓存命中率。令 R 表示批次中的请求集合，T 表示由 R 构建的基数树。对于 T 的每条边 e，与 e 关联的 KV 缓存至少需要计算一次。令 |e| 表示与 e 关联的 KV 缓存的大小，令 C 表示 R 的 KV 缓存计算复杂度。我们得到下界：
C ≥ Σ_{e ∈ edges(T)} |e|
考虑以 DFS 顺序访问基数树 T。对于 T 的每条边 e，当我们第一次计算与 e 关联的 KV 缓存时，接着会计算 e 的整个子树。在计算 e 的子树期间，边 e 会被持续命中，因此不会发生额外计算。在以 e 为根的子树计算完成后，边 e 将不会再被访问。注意到，当缓存大小 ≥ 最大请求长度（等于基数树 T 中的最长路径）时，在计算 e 的子树期间，边 e 不会被驱逐，因为该子树包含 e 的公共前缀会被持续命中。因此，与每条边 e 关联的 KV 缓存只会被计算一次。于是我们达到下界：
C = Σ_{e ∈ edges(T)} |e|
缓存命中率定义为：
（Σ_{r∈R} r 中被缓存的 prefill token 数）÷（Σ_{r∈R} r 中的 prefill token 数）
它等于 1 − C ÷（Σ_{r∈R} prefill token 数），达到了其上界，从而实现了最优性。
接下来，我们通过归纳法证明最长共享前缀优先顺序等价于 DFS 顺序：
- （基础情形） 开始时，由于没有任何缓存，某个对应于 T 中节点 x 的随机请求将被处理。所有对应于从根到 x 路径上节点 {v₁, ..., vₙ} 的请求都不需要重新计算。对应于节点 {v₁, ..., vₙ, x} 的请求的计算复杂度与某个合法的 DFS 一致。从根到 x 的路径被缓存。
- （归纳步骤） 假设我们刚访问了 T 中的节点 y，且已访问的节点符合 DFS 顺序。令 P 表示从根到 y 的路径。那么每个尚未访问的节点与已访问节点的最低公共祖先都在 P 上。由于 P 上的节点已被缓存，一个尚未访问、且其最低公共祖先在 P 上的节点 z 将具有最长的共享前缀。最长共享前缀优先顺序将选择 z，这是一个合法的 DFS 顺序。从根到 z 的路径将被缓存，因为它是最近使用的。
证毕。
在在线情形下，DFS 顺序会被打破，但最长共享前缀调度仍然在隐含基数树的增量部分上近似 DFS 行为。我们通过考虑添加一批新请求的步骤来说明这一点。
令 T 表示迄今为止已访问的基数树部分，T′ 表示添加新批次请求后的整个新基数树。令 C 表示 T 中已缓存节点的集合。令 longest© 表示 C 中从根出发路径最长、且在 T′ 中的子树尚未被完全访问的节点。
最长共享前缀调度随后将以 DFS 顺序处理 T′ 中以 longest© 为根的子树。在此过程中可能发生驱逐，C 中剩余的缓存节点变为 C⁽¹⁾ ⊆ C。随后将对 T′ 中以 longest(C⁽¹⁾) 为根的子树进行 DFS。
类似地，我们会得到 C⁽²⁾, ..., C⁽ᵏ⁾，直到 C⁽ᵏ⁾ 中只剩下一个叶节点，其在 T′ 中的子树尚未被完全访问。此时，我们达到了一个合法的 DFS 状态。T′ 的剩余部分将按照定理 1 证明中描述的 DFS 顺序被访问。
A.4 数据并行分布式 RadixAttention
为了将 RadixAttention 适配到具有多个副本工作节点（即数据并行）的分布式场景，我们开发了一种机制：每个工作节点维护自己的子树，而路由器（router）监督一棵元树（meta-tree）。这棵元树充当一棵字典树，跟踪所有子树及其关联设备。当一批新请求到达路由器时，在元树上执行前缀匹配。我们基于每个请求的亲和性（以与特定工作节点及同组其他请求的共享前缀长度衡量）实现了多种策略，以做出最小化冗余计算的高效分发决策。每次处理新请求时，路由器和工作节点都会独立更新各自的树。如果某个工作节点发生驱逐，它会将该驱逐提交到一个队列，路由器随后在空闲时段处理该队列以更新元树。
我们使用四个工作节点和 MMLU 数据集对这种分布式配置进行了基准测试，观察到它实现了线性扩展和最优的缓存命中率，且这种弱一致性分布式缓存设计带来的开销很小。在最大化数据局部性和并行处理效率之间存在权衡。探索高级调度策略以优化这一权衡，被指定为未来的研究方向。此外，来自 Preble 49 的同期工作基于 SGLang 的早期版本研究了数据并行调度。
附录 B：压缩有限状态机补充细节
本节讨论用于更快受限解码的压缩有限状态机的背景和实现细节。我们的目标是让 LLM 遵循正则表达式（regex），正则表达式提供了更强的表达力，可用于表示 JSON 模式等常见格式。为实现这一目标，我们将正则表达式转换为有限状态机（FSM），以在解码过程中引导生成过程 64。FSM 本质上是一个由节点（状态）和边（带有字符串/字符的转移）组成的图。从初始状态开始，每次转移将边上的字符串追加进来以移动到下一状态，并使用一组终止状态来结束整个过程。该机制引导 LLM 的解码，根据 FSM 的当前状态转移过滤掉无效 token，如图 10 所示。解码可能涉及 FSM 中的多次转移，直至到达终止状态。
图 10：正则表达式如何转换为 FSM 以及 FSM 如何引导解码过程的示例。
受限解码的挑战来自于这样一个事实：约束通常以自然语言格式表达，即正则表达式由字符/字符串描述，而 LLM 被设计为以 token 来理解和处理这些内容。这造成了复杂的局面，因为字符串与 token 之间的映射错综复杂，不存在一一对应关系 31。
本节内容源自一篇早期博客文章（https://lmsys.org/blog/2024-02-05-compressed-fsm/）。我们也鼓励读者阅读该博客文章，以获得更多背景知识和更容易的理解。
B.1 压缩有限状态机的实现细节
为简化压缩 FSM 的构建，我们在字符/字符串（而非 token）上构建原始 FSM。我们正式定义单一转移边和压缩边的概念如下：
- 单一转移边（singular transition edge）： 一条边是单一转移边，当且仅当：1）其源节点只有一个后继节点；2）其上只有一个可接受的字符/字符串。
- 压缩边（compressed edge）： 一条边由若干条连续相邻的边 (e₀, e₁, ..., eₖ) 压缩而成，当且仅当 e₁, ..., eₖ 均为单一转移边。压缩边的文本是 e₀, e₁, ..., eₖ 文本的拼接。
从基于字符的 FSM 出发，我们递归地将单一转移边合并到其前驱边中，直到无法继续压缩为止，最终得到压缩 FSM。该方法加速了解码过程，SGLang 运行时配合压缩 FSM 的效率如图 11（见原文图 compressed_fsm）所示。
图 11：使用压缩 FSM 与普通 FSM 解码的对比：左子图描绘了每次前向传播的解码过程，右子图解释了各结果组成部分的来源。
B.2 通过重分词处理分词伪影
当生成一个新 token 时，我们获取该 token 的字符串，并在当前状态的所有出边中搜索以刚解码字符串开头的那条边，然后向前移动。然而，当转移边是一条压缩得很好、包含很长字符串的边时，我们也可以预判接下来几轮解码出的字符串。加速正是发生在这里，我们将此过程称为向前跳转（Jump Forward） 。然而，我们仍需将该字符串转换为 token 以供后续解码阶段使用，而由于 LLM 特定的预训练和分词方法，直接划分可能会改变其本意 59，因此这一步并不直接。例如，图 2 中正则表达式的压缩文本是 {"summary": "，根据分词器，它只能被分词为 {"、summary、": 和 "，而不能随机划分为 {"、summa、ry 和 ":_"。为解决这一问题，我们使用原始分词器对所有先前文本以及压缩边的文本进行重分词（retokenize），确保与 LLM 的原始输入格式对齐。这只会带来很小的重分词开销。
B.3 未来扩展：解决概率失真问题
字符串与 token 之间的鸿沟还带来了概率分布偏斜的问题 59。例如，在图 2 中，正则表达式 "[ABCD][+-]?" 表示从 A+ 到 D− 的等级，但如果换成更宽泛的词汇，如 Excellent|Above Average|Fair|Below Average，运行时可能会不准确地将 A 映射到 Above Average------因为 Above Average 位于一条压缩转移边上------从而曲解了等级体系。发生这种情况是因为 LLM 无法识别具体的选项范围，导致产生不恰当的 token 序列。要为每个选项计算准确的概率，需要将产生该选项的所有 token 序列的概率求和，这使解码复杂化并增加开销。一种变通方法是将选项或正则表达式直接包含在 prefill 提示中，引导 LLM 意识到其选项并以恰当的 token 序列输出决策。然而，这种方法并没有解决概率失真的根本问题，凸显了进一步研究以提高压缩 FSM 准确性的必要性。
附录 C：补充实验设置与结果
补充实验设置。 图 5 和图 6 通过在单块 A10G（24GB）GPU 上运行 Llama-7B 获得。图 7 通过在 8 块 A10G（24GB）GPU 上以张量并行运行 Mixtral-8x7B 获得。图 8© 通过在单块 A10G（24GB）GPU 上运行 Llama-7B 获得。附录图 12（e2e_throughput_llama_70b）通过在 4 块 A100（80GB）GPU 上以张量并行运行 Llama-70B 获得。表 2 通过在单块 A10G（24GB）GPU 上运行 LLaVA-v1.5-7B、在单块 A100（80GB）GPU 上运行 LLaVA-Next-34B 获得。基准图中的每个柱需要运行几分钟到一小时。
补充实验结果。 图 12 展示了图 5 所列基准上达到的和最优的缓存命中率。图 13 展示了采用张量并行的 Llama-2-70B 上的吞吐量。
图 12：采用张量并行的 Llama-2-70B 模型上的归一化吞吐量。越高越好。
图 13：各基准上达到的缓存命中率与最优缓存命中率。
图 14（左）：使用 skeleton-of-thought 提示进行并行提示语建议的 SGLang 程序。
图 14（右）：图 14（左）中程序对应的计算图。三条流对应三次函数调用。
附录 D：编译器模式
除了论文主体中使用的解释器模式之外，运行 SGLang 程序的另一种方式是将它们编译为计算图，并用图执行器执行。这为更多编译优化打开了机会，因为我们可以重写图并进行更多静态规划。
D.1 设计与实现
我们为 SGLang 设计了一种中间表示（IR），将 SGLang 程序结构和操作表示为计算图。该图包含表示原语操作符的节点和表示依赖关系的边。图 14（右）展示了与图 14（左）中程序对应的图。在程序中，每次调用被装饰的函数或 fork 都会创建一个新的提示状态或一条流。
节点有若干种类型。SGLang 程序中 += 和 + 操作符的每个操作数都由一个 IR 节点表示，包括 ConstantText、Argument、Gen、Select、Variable、Fork、GetForkItem 和 Join。依赖关系有两种类型：流内依赖（intra-stream dependency），即使用 += 提交到一条流中的操作必须在该流中所有先前操作之后执行；以及流间依赖（inter-stream dependency），即当一条流需要从另一条流获取变量值时产生，需要同步。fork 等操作会操纵多条流，因此引入流间依赖。
为生成图，我们使用追踪（tracing）技术，以抽象参数运行程序并动态构建图。该方法仅限于没有数据依赖控制流的程序，我们计划在未来工作中解决这一限制。图构建完成后，即可通过图执行器执行，无需重新解释原始 Python 程序。这带来了图重写优化、降低运行时开销和程序序列化等好处。执行时，为每条数据流启动流执行器，按拓扑顺序将 IR 节点分发到各条流中。
D.2 编译器优化案例研究：用于改善前缀共享的代码移动
我们探索了针对 SGLang IR 的一项编译优化：用于改善前缀共享的代码移动（code movement）。我们预计更多经典的编译器技术也可以被应用，如自动调优（auto-tuning）和指令选择。
该优化旨在通过对图中的节点重排序来增加常量前缀的长度，从而改善前缀共享。它不严格保持原始计算，因此被归类为激进优化（aggressive optimization）。例如，将提示从 "Here is a question + {question}. Please act as a math expert and solve the given question."（这里有一个问题 + {问题}。请扮演一位数学专家并解答给定的问题。）改为 "Please act as a math expert and solve the given question. Here is a question + {question}."（请扮演一位数学专家并解答给定的问题。这里有一个问题 + {问题}。），会产生更长的可共享前缀。这项优化很有意思，因为由于 SGLang 中存在自然语言指令，传统的程序分析无法实现它。作为替代，我们提示 GPT-4 对图节点重排序。我们编写了一个包含若干示例的提示，向 GPT-4 讲授 SGLang IR 的概念，并发现 GPT-4 能够成功地为一些简单的 SGLang 程序应用此优化。
我们评估了该优化的有效性。我们从互联网上收集了 20 个提示模板，并用 SGLang 实现它们。我们使用其中 5 个模板作为少样本训练示例，其余 15 个作为测试用例。评估表明，对于 15 个模板中的 12 个，GPT-4 成功地对图节点进行了重排序且没有改变语义------这一点通过对修改后提示的人工检查得到确认。平均而言，该优化使可共享前缀长度增加了 60 个 token，展示了 GPT-4 的有效性。生成优化提示顺序的失败源于对图节点背后语义的错误理解------它过于激进，即使改变原有语义也要把所有常量放在前面。本案例研究旨在探索使用 GPT-4 进行编译优化。未来需要更多工作来使这类优化变得可靠。