文章标题：SGLang: 高效执行结构化语言模型程序
作者/机构：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
所属机构：斯坦福大学, 加州大学伯克利分校, 上海交通大学, 德州农工大学, 独立研究员

A1 主要贡献

本文针对日益复杂的语言模型（LLM）应用中存在的编程和执行效率低下问题，提出了SGLang，一个用于高效执行复杂语言模型程序的系统。

核心问题：
随着LLM越来越多地用于需要多次生成调用、高级提示技术、控制流和结构化输入/输出的复杂任务（如自主代理、多轮规划与推理），现有的系统在编程和执行这些"语言模型程序"（LM Programs）方面效率低下。主要挑战有两个：
1. 编程复杂性：开发LM程序通常需要大量的字符串操作、繁琐的提示调整、脆弱的输出解析以及并行机制的实现，导致程序可读性差且开发困难。
2. 执行效率低下：先进的推理引擎（如vLLM）虽然通用，但在执行具体的LM程序时存在大量冗余计算和内存使用。例如，它们缺乏跨多个LLM调用有效重用键值缓存（KV cache）的机制，并且在生成结构化输出（如JSON）时，一次只解码一个token，速度较慢。

研究目标与主要贡献：
为了解决上述挑战，本文提出了SGLang，其核心思想是系统性地利用LM程序中的多调用结构来提高执行效率。SGLang包含一个前端语言和一个后端运行时，两者可以协同工作，也可以独立运行。

创新点：
1. SGLang前端语言：一个嵌入在Python中的领域特定语言，提供了用于生成（gen, select）、并行控制（fork, join）等高级原语。它简化了LM程序的编程，提高了代码的可读性和开发效率。
2. RadixAttention技术：一种新颖的运行时优化技术，通过将所有请求的KV cache维护在一个基数树（radix tree）结构的LRU缓存中，实现了跨多个生成调用的KV cache自动重用。这显著减少了冗余计算和内存浪费，并配合缓存感知调度策略进一步提升了缓存命中率。
3. 压缩有限状态机（Compressed FSM）：一种用于加速约束解码（如生成JSON）的技术。通过分析语法约束并构建压缩的有限状态机，系统能够将多token路径压缩为单步路径，从而在一次前向传播中解码多个token，大幅提升解码速度。
4. API投机执行（API Speculative Execution）：针对仅提供API访问的模型（如GPT-4），SGLang提出了一种优化策略，通过在一次API调用中投机性地生成更多内容，并将其与后续的生成原语进行匹配和重用，从而减少API调用次数、延迟和成本。

实验表明，SGLang在多种LLM和多模态模型上，相较于现有系统，在代理控制、逻辑推理、少样本学习等任务上实现了高达6.4倍的吞吐量提升。

图1：系统架构：一个解释器使用优化的运行时执行语言原语。
图1：系统架构：一个解释器使用优化的运行时执行语言原语。

A3 背景知识/关键Observation/设计原则

第2章 编程模型

本章通过一个运行示例介绍SGLang编程模型，描述其语言原语和执行模式，并概述了运行时的优化机会。该编程模型通过提供灵活且可组合的原语，简化了多调用工作流中的繁琐操作，例如字符串操作、API调用、约束指定和并行处理。

一个运行中的例子。SGLang是一个嵌入在Python中的领域特定语言。图2展示了一个使用“分支-解决-合并”提示技术【索引40，Branch-solve-merge improves large language model evaluation and generation，2023，arXiv】来评估一篇关于图像的文章的程序。函数multi_dimensional_judge接收三个参数：s（管理提示状态）、path（图像文件路径）和essay（文章文本）。新的字符串和SGLang原语可以通过+=操作符附加到状态s上执行。首先，函数将图像和文章添加到提示中。然后，它使用select原语检查文章是否与图像相关，并将结果存储在s["related"]中。如果相关，提示将被fork成三个副本，用于从不同维度进行并行评估，并使用gen原语将结果存储在f["judgment"]中。接下来，程序合并这些判断，生成一个总结，并给出一个字母等级。最后，它遵循由正则表达式约束regex定义的模式，以JSON格式返回结果。SGLang极大地简化了这个程序，相比使用类似OpenAI API接口的实现，由于避免了手动的字符串操作和并行控制，代码行数减少了2.1倍。

图2：在SGLang中实现的多维度文章评判器，利用了分支-解决-合并的提示技术【40】。SGLang提供的原语以红色显示。
图2：在SGLang中实现的多维度文章评判器，利用了分支-解决-合并的提示技术【40】。SGLang提供的原语以红色显示。

语言原语。SGLang提供了用于控制提示状态、生成和并行性的原语，它们可以与Python语法和库一起使用。这些原语包括：
* gen：调用模型进行生成，并将结果存储在其第一个参数指定的变量名中。它支持一个regex参数，用以约束输出遵循由正则表达式定义的语法（例如JSON模式）。
* select：调用模型从一个列表中选择概率最高的选项。
* += 或 extend：将一个字符串附加到提示中。
* [variable_name]：获取生成任务的结果。
* fork：创建提示状态的并行分支。
* join：重新合并提示状态。
* image 和 video：接收图像和视频输入。

执行模式。执行SGLang程序最简单的方式是通过解释器，其中提示被视为一个异步流。像extend、gen和select这样的原语被提交到流中进行异步执行。这些非阻塞调用允许Python代码在不等待生成完成的情况下继续运行，这类似于异步启动CUDA内核。每个提示由一个后台线程中的流执行器管理，从而实现了程序内的并行性。获取生成结果的操作将会阻塞，直到结果准备就绪，以确保正确的同步。另外，SGLang程序也可以被编译成计算图，并由图执行器执行，从而允许更多的优化。本文默认使用解释器模式，并在附录D中讨论了编译器模式的结果。SGLang支持使用其自有的SGLang运行时（SRT）的开源权重模型，以及像OpenAI和Anthropic模型这样的API模型。

对比。针对LLM的编程系统可以分为高级（如LangChain, DSPy）和低级（如LMQL, Guidance, SGLang）。高级系统提供预定义或自动生成的提示，例如DSPy的提示优化器。低级系统通常不改变提示，但允许直接操作提示和原语。SGLang是一个类似于LMQL和Guidance的低级系统。表1比较了它们的特性。SGLang更侧重于运行时效率，并附带了其协同设计的运行时，从而实现了后续将介绍的新颖优化。高级语言（如DSPy）可以被编译成低级语言（如SGLang）。我们在第6节中演示了将SGLang作为DSPy的后端集成以获得更好运行时效率的案例。


表1：LMQL、Guidance和SGLang之间的比较。

运行时优化。图2展示了三种运行时优化机会：KV缓存重用、快速约束解码和API投机执行。我们将在接下来的章节中讨论它们。

A2 方法细节

第3章 使用RadixAttention实现高效的KV缓存重用

KV缓存重用机会。SGLang程序可以链接多个生成调用，并使用 "fork" 原语创建并行副本。此外，不同的程序实例通常共享一些共同的部分（例如，系统提示）。这些场景在执行过程中会产生许多共享的提示前缀，从而带来了大量重用KV缓存的机会。在LLM推理过程中，KV缓存存储了前向传播过程中的中间张量，这些张量被用于解码未来的token。它们以自注意力机制【索引51，Attention is all you need，2017，NeurIPS】中的键值对命名。KV缓存的计算仅依赖于前缀token。因此，具有相同提示前缀的请求可以重用KV缓存，从而减少冗余计算和内存使用。更多背景和一些示例在附录A中提供。

RadixAttention技术提出。鉴于KV缓存的重用机会，优化SGLang程序的关键挑战是在多个调用和实例之间重用KV缓存。虽然一些系统探索了某些KV缓存的重用案例【索引23，Efficient memory management for large language model serving with pagedattention，2023，SOSP】、【索引58，Chunkattention: Efficient self-attention with prefix-aware kv cache and two-phase partition，2024，arXiv】、【索引18，Hydragen: High-throughput llm inference with shared prefixes，2024，arXiv】、【索引12，Prompt cache: Modular attention reuse for low-latency inference，2023，arXiv】，但它们通常需要手动配置，并且无法处理所有的重用模式（例如，动态树结构）。因此，大多数先进的推理系统会为每个请求重新计算KV缓存。我们将在第7节讨论它们的局限性以及我们的不同之处。

RadixAttention核心机制。本节介绍RadixAttention，一种在运行时自动且系统地重用KV缓存的新技术。与现有系统在生成请求完成后丢弃KV缓存不同，我们的系统将提示和生成结果的缓存保留在一个基数树（radix tree）中，从而实现了高效的前缀搜索、重用、插入和驱逐。我们实现了一个LRU（最近最少使用）驱逐策略和一个缓存感知调度策略来提高缓存命中率。RadixAttention与连续批处理【索引60，Orca: A distributed serving system for {Transformer-Based} generative models，2022，OSDI】、分页注意力【索引23，Efficient memory management for large language model serving with pagedattention，2023，SOSP】和张量并行【索引44，Megatron-lm: Training multi-billion parameter language models using model parallelism，2019，arXiv】等技术兼容。此外，当没有缓存命中时，它引入的内存和时间开销可以忽略不计。

RadixAttention实现细节。基数树是一种数据结构，作为经典trie（前缀树）的空间高效替代方案。与典型的树不同，基数树的边不仅可以标记单个元素，还可以标记不同长度的元素序列，从而显著提高效率。在我们的系统中，我们利用基数树来管理token序列与其对应的KV缓存张量之间的映射。这些KV缓存张量存储在非连续的分页布局中，每个页的大小相当于一个token。由于GPU内存很快会被KV缓存填满，我们引入了一个简单的LRU驱逐策略，该策略首先驱逐最近最少使用的叶节点。通过首先驱逐叶节点，我们使得它们的共同祖先能够被重用，直到这些祖先也变成叶节点并被驱逐。在连续批处理的设置中，我们不能驱逐当前正在运行的批次所使用的节点。因此，每个节点维护一个引用计数器，指示有多少正在运行的请求正在使用它。如果一个节点的引用计数器为零，则该节点是可驱逐的。请注意，我们没有预先分配一个固定大小的内存池作为缓存。相反，我们让缓存的token和当前运行的请求共享同一个内存池。因此，系统会动态地为缓存和运行中的请求分配内存。当有足够多的等待请求运行时，系统将驱逐所有缓存的token，以支持更大的批次大小。图3展示了在处理几个传入请求时，基数树是如何维护的。前端解释器将完整的提示发送到运行时，运行时执行前缀匹配和重用。树结构存储在CPU上，维护开销可以忽略不计。在执行fork原语期间，前端首先发送前缀作为提示，以确保前缀被正确插入到树中，然后再发送剩余的提示。这种“前端提示”简化了运行时的调度和匹配，体现了前端与运行时协同设计的优势。

图3：RadixAttention操作的示例，采用LRU驱逐策略，跨越九个时间点进行说明。该图展示了基数树为响应各种请求而发生的动态演变。这些请求包括两个聊天会话、一批少样本学习查询和一次自洽性抽样。每个树边都带有一个表示子字符串或token序列的标签。节点根据不同状态进行颜色编码：绿色代表新添加的节点，蓝色代表在该时间点被访问的缓存节点，红色代表已被驱逐的节点。在步骤（1）中，基数树初始为空。在步骤（2）中，服务器处理一个传入的用户消息“Hello”并以LLM输出“Hi”作为回应。系统提示“你是一个有帮助的助手”、用户消息“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”。
图3：RadixAttention操作的示例，采用LRU驱逐策略，跨越九个时间点进行说明。该图展示了基数树为响应各种请求而发生的动态演变。这些请求包括两个聊天会话、一批少样本学习查询和一次自洽性抽样。每个树边都带有一个表示子字符串或token序列的标签。节点根据不同状态进行颜色编码：绿色代表新添加的节点，蓝色代表在该时间点被访问的缓存节点，红色代表已被驱逐的节点。在步骤（1）中，基数树初始为空。在步骤（2）中，服务器处理一个传入的用户消息“Hello”并以LLM输出“Hi”作为回应。系统提示“你是一个有帮助的助手”、用户消息“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”。

缓存感知调度。我们将缓存命中率定义为 $\frac{\text{缓存的提示token数量}}{\text{提示token总数}}$。当等待队列中有许多请求时，它们的执行顺序会显著影响缓存命中率。例如，如果请求调度器频繁地在不同且不相关的请求之间切换，可能导致缓存颠簸（cache thrashing）和低命中率。我们设计了一种缓存感知的调度算法来提高缓存命中率。在批处理设置中，我们根据匹配前缀的长度对请求进行排序，并优先处理具有更长匹配前缀的请求，而不是采用先到先得的调度方式。附录中的算法1展示了使用连续批处理的缓存感知调度的伪代码。该算法使用最长共享前缀优先的顺序。在对延迟更敏感的设置中，我们可能仍然能够容忍有限的批次重新排序以提高缓存重用。此外，我们证明了以下关于离线情况下最优调度的定理。

图4：正常和压缩FSM的解码过程（下划线_表示空格）。
图4：正常和压缩FSM的解码过程（下划线_表示空格）。

最优调度定理。定理 3.1：对于一批请求，当缓存大小 ≥ 最大请求长度时，通过以深度优先搜索（DFS）的顺序访问请求的基数树，我们可以实现最优的缓存命中率。最长共享前缀优先顺序等价于深度优先搜索顺序。

定理证明与在线情况。证明在附录A.3节中。在在线情况下，DFS顺序会被打乱，但我们的调度仍然能在完整的基数树的增量部分上近似DFS行为，如附录A.3节所述。虽然贪心的缓存感知调度可以实现高吞吐量，但可能导致饥饿问题。我们将其与其他公平调度方法【索引42，Fairness in serving large language models，2024，arXiv】的集成作为未来工作。

分布式情况。RadixAttention可以扩展到多个GPU。对于张量并行，每个GPU维护一个分片的KV缓存。由于树操作是相同的，因此不需要额外的同步。对于有多个工作节点的数据并行情况在附录A.4节中讨论。

第4章 使用压缩有限状态机实现高效的约束解码

约束解码的挑战。在LM程序中，用户通常希望约束模型的输出以遵循特定格式，例如JSON模式。这可以提高可控性和鲁棒性，并使输出更容易解析。SGLang提供了一个regex参数，通过使用正则表达式来强制执行此类约束，正则表达式的表达能力足以应对许多实际场景。现有系统通过将正则表达式转换为有限状态机（FSM）【索引54，Efficient guided generation for large language models，2023，arXiv】来支持这一点。在解码过程中，它们维护当前的FSM状态，从下一个状态中检索允许的token，并将无效token的概率设置为零，从而逐个token进行解码。然而，当有机会一次解码多个token时，这种逐个token解码的方法是低效的。例如，图2中的常量序列{"summary": "在正常的解码过程中跨越多个token，如图4(c)所示，需要多个解码阶段，尽管在解码它时只有一个有效的下一个token。因此，整个序列可以在一个步骤（即一次前向传播）中被解码。然而，现有系统一次只能解码一个token，因为现有系统中FSM和模型运行器之间缺乏集成，阻碍了多token处理，导致解码速度缓慢。

SGLang的压缩FSM方案。SGLang通过创建一个带有压缩FSM的快速约束解码运行时来克服这一限制。该运行时分析FSM，并将FSM中相邻的单一转换边压缩成单一边，如图4(b)所示，使其能够识别何时可以一起解码多个token。在图4(d)中，压缩转换边上的多个token可以在一次前向传播中被解码，这极大地加速了解码过程。该方法是通用的，适用于所有正则表达式。更多关于背景和实现的细节在附录B中。

第5章 使用API投机执行实现高效的端点调用

API模型的优化背景。前面的章节介绍了针对开源权重模型的优化，这些优化需要修改模型推理过程。此外，SGLang也适用于仅能通过API访问的模型，例如OpenAI的GPT-4。然而，对于这些模型，我们只能调用一个黑盒API端点。

API投机执行。本节介绍一种针对黑盒API模型的新优化，该优化使用投机执行来加速多调用SGLang程序的执行并降低其API成本。例如，一个程序可能会通过多调用模式要求模型生成一个角色的描述：s += context + "name:" + gen("name", stop="\n") + "job:" + gen("job", stop="\n")。简单地，两个gen原语对应两次API调用，这意味着用户需要为上下文（context）的输入token费用支付两次。在SGLang中，我们可以在第一次调用时启用投机执行，让它忽略停止条件继续生成几个token。解释器会保留这些额外的生成输出，并与后续的原语进行匹配和重用。在某些情况下，通过精心的提示工程，模型可以高精度地正确匹配模板，从而为我们节省一次API调用的延迟和输入成本。

A4 实验环境

  • 模型架构：

    • 稠密模型：Llama-2系列（7B, 70B）【索引49，Llama 2: Open foundation and fine-tuned chat models，2023，arXiv】
    • 稀疏混合专家模型：Mixtral-8x7B【索引17，Mixtral of experts，2024，arXiv】
    • 多模态模型：LLaVA-v1.5-7B（图像）【索引27，Improved baselines with visual instruction tuning，2023，arXiv】，LLaVA-NeXT-34B（视频）【索引62，Llava-next: A strong zero-shot video understanding model，2024，arXiv】
    • API模型：OpenAI GPT-3.5
    • 所有开源权重模型均使用float16精度。

  • 硬件配置：

    • 主要实验平台：AWS EC2 G5实例，配备NVIDIA A10G GPU（24GB）。
    • 7B模型在单个A10G GPU上运行。
    • 更大模型（如Mixtral-8x7B, Llama-70B）在多个A10G GPU上使用张量并行运行。
    • 部分额外实验在NVIDIA A100G GPU（80GB）上运行。

  • 软件配置：

    • 实现：SGLang使用PyTorch【索引37，Pytorch: An imperative style, high-performance deep learning library，2019，NeurIPS】实现。
    • 依赖库：使用了来自FlashInfer【索引59，Accelerating self-attentions for llm serving with flashinfer，2024，arXiv】和Triton【索引48，Triton: an intermediate language and compiler for tiled neural network computations，2019，MLPL】的自定义CUDA内核。

  • 基线系统：

    • Guidance v0.1.8【索引13，A guidance language for controlling large language models，2023，GitHub】，使用llama.cpp后端。
    • vLLM v0.2.5【索引23，Efficient memory management for large language model serving with pagedattention，2023，SOSP】，使用其默认API服务器。
    • LMQL v0.7.3【索引4，Prompting is programming: A query language for large language models，2023，PLDI】，使用Hugging Face Transformers后端。

  • 数据集/工作负载：

    • 少样本学习：5-shot MMLU【索引14，Measuring massive multitask language understanding，2020，ICLR】和20-shot HellaSwag【索引61，Hellaswag: Can a machine really finish your sentence?，2019，ACL】。
    • Agent：ReAct Agent【索引57，React: Synergizing reasoning and acting in language models，2022，ICLR】和Generative Agents【索引36，Generative agents: Interactive simulacra of human behavior，2023，UIST】的轨迹回放。
    • 复杂推理：用于GSM-8K问题的Tree-of-thought【索引56，Tree of thoughts: Deliberate problem solving with large language models，2023，arXiv】和用于提示生成的Skeleton-of-thought【索引33，Skeleton-of-thought: Prompting LLMs for efficient parallel generation，2024，ICLR】。
    • 其他：使用branch-solve-merge技术的LLM评判【索引40，Branch-solve-merge improves large language model evaluation and generation，2023，arXiv】、JSON解码、多轮对话（4轮，长短两种输出）、DSPy RAG流水线【索引20，Dspy: Compiling declarative language model calls into self-improving pipelines，2023，arXiv】。
    • 多模态：llava-bench-in-the-wild（图像），ActivityNet（视频）。

A4 实验结果

端到端性能
* 开源权重模型（Llama-7B）：如图5和图6所示，SGLang在吞吐量上最高提升了6.4倍，延迟最多降低了3.7倍。
* MMLU和HellaSwag：RadixAttention通过重用少样本示例和问题前缀的KV缓存，提升了吞吐量和延迟。
* ReAct和Generative Agents：SGLang重用了代理模板和历史调用的KV缓存。
* Tree-of-thought和Skeleton-of-thought：SGLang利用了程序内并行性并尽可能重用KV缓存。
* JSON解码：压缩有限状态机通过一次解码多个token加速了解码过程。
* 多轮对话：SGLang重用了聊天历史的KV缓存，在短输出场景下加速更明显。长输出场景由于解码时间占主导，加速不明显。
* DSPy RAG流水线：SGLang重用了通用上下文示例的KV缓存。
* 在这些基准测试中，缓存命中率在50%到99%之间。缓存感知调度策略平均达到了最优命中率的96%（图13）。

图5：在Llama-7B模型上的归一化吞吐量。越高越好。
图5：在Llama-7B模型上的归一化吞吐量。越高越好。

图6：在Llama-7B模型上的归一化延迟。越低越好。
图6：在Llama-7B模型上的归一化延迟。越低越好。

  • 更大规模模型（Mixtral-8x7B, Llama-70B）：如图7和图12所示，在大模型和张量并行设置下，SGLang的性能提升趋势与小模型相似，表明其优化具有良好的泛化性。

图7：在Mixtral-8x7B模型上使用张量并行的归一化吞吐量。越高越好。
图7：在Mixtral-8x7B模型上使用张量并行的归一化吞吐量。越高越好。

  • 多模态模型：如表2所示，SGLang在LLaVA图像和视频模型上的吞吐量比基线（Hugging Face实现）高出最多6倍。这是通过对输入图像哈希并将其作为基数树的键，从而重用相同图像的KV缓存实现的。

  • 生产环境部署：SGLang已部署在Chatbot Arena【索引8，Chatbot arena: An open platform for evaluating llms by human preference，2024，arXiv】。在一个月的观察期内，LLaVA-Next-34B和Vicuna-33B的RadixAttention缓存命中率分别达到了52.4%和74.1%，主要来自系统消息、常用示例图片和多轮对话历史。这使得Vicuna-33B的平均首token延迟降低了1.7倍。

  • API模型：在使用GPT-3.5从维基百科页面提取三个字段的测试中，API投机执行的准确率很高，并将输入token的成本降低了约三倍。

消融研究
* 缓存命中率与性能：如图8(a)(b)所示，更高的缓存命中率直接导致了更大的批处理大小、更高的吞吐量和更低的首token延迟与总延迟。
* RadixAttention有效性：如图8(c)所示，与"无缓存"、"无树结构缓存"、"FCFS调度"等变体相比，完整的RadixAttention（包含所有组件，如树结构、缓存感知调度、前端并行和提示）性能最佳，证明了每个组件的必要性以及前后端协同设计的重要性。
* RadixAttention开销：在没有任何KV缓存重用机会的基准测试中，管理RadixAttention数据结构的开销不到0.3%，可以忽略不计。
* 压缩有限状态机有效性：在JSON解码基准测试中，压缩FSM将吞吐量提高了1.6倍。此外，对状态机进行预处理并批量重用至关重要，否则性能会下降2.4倍。

图8：(a)(b) 缓存命中率消融研究。(c) RadixAttention消融研究。
图8：(a)(b) 缓存命中率消融研究。(c) RadixAttention消融研究。

A5 结论

结论：
本文介绍了SGLang，一个用于高效编程和执行结构化语言模型程序的框架。通过引入RadixAttention、压缩有限状态机和语言解释器等新颖优化，SGLang显著提升了复杂LM程序的吞吐量和延迟。它为开发高级提示技术和代理工作流提供了有价值的工具。

未来方向：
尽管SGLang取得了显著进展，但仍存在一些局限性，为未来的研究指明了方向：
1. 扩展输出模态：将SGLang的支持扩展到除文本外的其他输出模态。
2. 多层级内存RadixAttention：使RadixAttention能够跨越多级内存层次结构（如DRAM、磁盘）运行。
3. 模糊语义匹配：在RadixAttention中实现模糊语义匹配，而不仅仅是精确前缀匹配。
4. 更高级别的原语：在SGLang之上提供更高级别的编程原语。
5. 公平性调度：解决缓存感知调度中可能出现的饥饿问题。
6. 增强编译器：增强SGLang编译器以执行更高级的静态优化，如调度和内存规划。

A6 附录

附录A RadixAttention的额外细节

A.1 KV缓存背景。当前多数LLM（如GPT-3【索引5，Language models are few-shot learners，2020，NeurIPS】、PaLM【索引9，Palm: Scaling language modeling with pathways，2022，arXiv】、LLaMA【索引49，Llama 2: Open foundation and fine-tuned chat models，2023，arXiv】）基于自回归Transformer架构【索引51，Attention is all you need，2017，NeurIPS】。推理过程分为“prefill”（处理输入token）和“decoding”（逐个生成输出token）。在此过程中，每个token生成的中间张量，即“KV缓存”，被用于解码后续token。KV缓存的计算仅依赖于其之前的所有token，因此具有相同前缀的不同序列可以重用前缀的KV缓存以避免冗余计算。在LLM程序中，多个文本片段和生成调用被附加到单个提示中，跨调用缓存KV缓存可以减少冗余。此外，从单个提示生成多个输出或从当前状态分叉新提示的情况很常见【索引25，Competition-level code generation with alphacode，2022，Science】。虽然vLLM【索引23，Efficient memory management for large language model serving with pagedattention，2023，SOSP】已探索了基本的前缀共享，但更高级的共享模式如不规则的树状结构共享仍是挑战。图9展示了四种典型的KV缓存共享模式，现有系统都无法自动处理所有这些模式，而RadixAttention可以。

图9：KV缓存共享示例。蓝色框代表可共享的提示部分，绿色框表示不可共享的部分，黄色框标记不可共享的模型输出。可共享的元素包括少样本学习示例、自洽性中的问题【53】、多轮聊天中的聊天历史以及思想树中的搜索历史【56】。
图9：KV缓存共享示例。蓝色框代表可共享的提示部分，绿色框表示不可共享的部分，黄色框标记不可共享的模型输出。可共享的元素包括少样本学习示例、自洽性中的问题【53】、多轮聊天中的聊天历史以及思想树中的搜索历史【56】。

A.2 缓存感知调度的伪代码。算法1展示了使用连续批处理的RadixAttention缓存感知调度的伪代码。

A.3 定理3.1的证明。定理 3.1：对于一批请求，当缓存大小 ≥ 最大请求长度时，通过以深度优先搜索（DFS）的顺序访问请求的基数树，可以实现最优的缓存命中率。最长共享前缀优先顺序等价于深度优先搜索顺序。
证明：首先证明DFS顺序达到最优缓存命中率。设$R$为批次中的请求集，$T$为由$R$构建的基数树。对于$T$的每条边$e$，其关联的KV缓存至少需要计算一次。设$|e|$为边$e$的KV缓存大小，$C$为$R$的KV缓存计算复杂度。我们得到计算复杂度的下界为 $C \ge \sum_{e \in T} |e|$。
当我们以DFS顺序访问树$T$时，首次计算边$e$的KV缓存后，将继续计算$e$的整个子树。在此过程中，边$e$会被持续命中，不会产生额外计算。完成子树计算后，边$e$不会再被访问。只要缓存大小不小于最大请求长度（即$T$中最长路径），边$e$在其子树计算期间不会被驱逐。因此，每条边$e$的KV缓存只计算一次，达到了计算复杂度的下界 $C = \sum_{e \in T} |e|$。缓存命中率定义为 $\frac{\sum_{r \in R} \text{number of cached prefill tokens}}{\sum_{r \in R} \text{number of prefill tokens}}$，即 $1 - \frac{C}{\sum_{r \in R} \text{number of prefill tokens}}$，此时达到其上界，实现最优性。
接着，通过归纳法证明最长共享前缀优先顺序等价于DFS顺序。
* 基础步骤：开始时缓存为空，任意选择一个请求（对应节点$x$），其从根到$x$的路径上的节点{v1, ..., vn}的计算符合DFS顺序，该路径被缓存。
* 归纳步骤：假设已访问节点$y$，且访问顺序符合DFS。设$P$为从根到$y$的路径。未访问节点中，与已访问节点在$P$上有最低共同祖先的节点$z$将具有最长共享前缀。因此，最长共享前缀优先策略会选择$z$，这符合DFS顺序。
在在线场景中，DFS顺序会被打断，但最长共享前缀调度仍能在增量部分近似DFS行为。

A.4 数据并行分布式RadixAttention。为了将RadixAttention应用于具有多个副本工作节点（即数据并行）的分布式设置，我们开发了一种机制：每个工作节点维护自己的子树，而路由器则监管一个元树（meta-tree）。这个元树作为一个trie，跟踪所有子树及其关联的设备。当新一批请求到达路由器时，会在元树上执行前缀匹配。我们根据每个请求的亲和度（affinity）——通过与特定工作节点和同组其他请求的共享前缀长度来衡量——来实施各种策略，以做出高效的调度决策，从而最小化冗余计算。每次处理新请求时，路由器和工作节点都会独立更新各自的树。如果工作节点发生驱逐，它会将此驱逐操作提交到一个队列中，路由器在低活动期间处理该队列以更新元树。我们使用四个工作节点和MMLU数据集对这种分布式配置进行了基准测试，观察到它实现了线性扩展和最优的缓存命中率，且这种弱一致性分布式缓存设计的开销很小。在最大化数据局部性和并行处理效率之间存在一个权衡，探索更先进的调度策略来优化这一权衡是未来的研究方向。此外，Preble的同期工作【索引45，Preble: Efficient distributed prompt scheduling for llm serving，2024】基于早期版本的SGLang研究了数据并行调度。

附录B 压缩有限状态机的额外细节

背景。本节讨论用于更快约束解码的压缩有限状态机的背景和实现细节。我们的目标是让LLM遵循正则表达式（regex），它具有更强的表达能力，可以用来表示JSON模式等常见格式。为实现此目的，我们将regex转换为有限状态机（FSM）来指导解码过程【索引54，Efficient guided generation for large language models，2023，arXiv】。FSM本质上是一个带有节点（状态）和边（带字符串/字符的转换）的图。从初始状态开始，每次转换都会附加边上的字符串以移动到下一个状态，并使用一组最终状态来结束过程。该机制通过根据FSM当前状态的转换过滤掉无效token来指导LLM的解码，如图10所示。


图10：正则表达式如何转换为FSM以及FSM如何指导解码过程的示例。

挑战。约束解码的挑战在于约束通常以自然语言格式表达，即regex由字符/字符串描述，而LLM则被设计为将这些作为token来解释和处理。这造成了一个复杂的情况，因为字符串和token之间的映射是复杂的，并且缺乏一一对应的关系【索引22，Validating large language models with relm，2023，MLSys】。

B.1 压缩有限状态机的实现细节。为了简化压缩FSM的构建，我们在字符/字符串级别而不是token级别构建原始FSM。我们正式定义单一转换边和压缩边的概念如下：
* 单一转换边 (Singular transition edge)：一条边是单一转换边，如果 1) 其源节点只有一个后继节点，并且 2) 其上只有一个可接受的字符/字符串。
* 压缩边 (Compressed edge)：一条由多个连续相邻边 $(e_0, e_1, \dots, e_k)$ 压缩而成的边，当且仅当 $e_1, \dots, e_k$ 都是单一转换边。压缩边的文本是 $e_0, e_1, \dots, e_k$ 文本的拼接。
从一个基于字符的FSM开始，我们递归地将单一转换边合并到它们的前驱边中，直到无法进一步压缩，从而得到一个压缩FSM。这种方法加速了解码过程，SGLang运行时的效率如图11所示。

B.2 通过重分词处理分词伪影。当生成一个新token时，我们获取该token的字符串，并在当前状态的所有出边中搜索以该刚解码的字符串开头的边，然后向前推进。然而，当转换边是一个良好压缩且包含很长字符串的边时，我们也可以预测下一轮解码的字符串。这就是加速发生的地方，我们称之为“跳跃前进”（Jump Forward）。然而，由于LLM特定的预训练和分词方法，我们仍然需要将字符串转换为token用于后续解码阶段，这并非易事；直接划分可能会改变预期的含义【索引50，Fast, high-fidelity llm decoding with regex constraints，2024】。例如，图2的正则表达式中的压缩文本是 {"summary": "，根据分词器只能被分词为 {", summary, : 和 _"，而不能随机划分。为了解决这个问题，我们使用原始的分词器对所有先前的文本以及压缩边的文本进行重新分词，确保与LLM的原始输入格式对齐。这只带来了轻微的重分词开销。


图11：使用压缩FSM与普通FSM进行解码的比较：左子图描绘了每次前向传播的解码过程，而右子图解释了各种结果组成部分的来源。

B.3 未来扩展：解决概率扭曲问题。字符串和token之间的鸿沟也带来了概率分布倾斜的问题【索引50，Fast, high-fidelity llm decoding with regex constraints，2024】。例如，图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(c)是在单个A10G (24GB) GPU上运行Llama-7B获得。图12是在4个A100G (80GB) GPU上使用张量并行运行Llama-70B获得。表2是在单个A10G (24GB) GPU上运行LLaVA-v1.5-7B和在单个A100G (80GB) GPU上运行LLaVA-Next-34B获得。

额外实验结果。图13显示了在图5所列基准测试上实现的缓存命中率和最优缓存命中率。图12显示了在Llama-2-70B上使用张量并行的吞吐量。

图12：在Llama-2-70B模型上使用张量并行的归一化吞吐量。越高越好。
图12：在Llama-2-70B模型上使用张量并行的归一化吞吐量。越高越好。

图13：在各种基准测试上实现的缓存命中率和最优缓存命中率。
图13：在各种基准测试上实现的缓存命中率和最优缓存命中率。

附录D 编译器模式

设计与实现。除了论文主体中使用的解释器模式，运行SGLang程序的另一种方式是将其编译为计算图，并由图执行器执行。这为更多的编译优化提供了机会，因为我们可以重写图并执行更多的静态规划。我们为SGLang设计了一种中间表示（IR），它将SGLang程序的结构和操作表示为一个计算图。该图包括原语操作符的节点和表示依赖关系的边。图14b展示了对应于图14a程序的图。在程序中，每次调用被装饰的函数或fork都会创建一个新的提示状态或流。节点类型包括ConstantText, Argument, Gen, Select, Variable, Fork, GetForkItem, 和 Join。依赖关系分为两种：流内依赖（同一流内的操作按顺序执行）和流间依赖（一个流需要从另一个流获取变量值，需要同步）。为了生成图，我们使用追踪技术，用抽象参数运行程序并动态构建图。此方法目前仅限于没有数据依赖控制流的程序。

(a) 使用思想骨架提示进行并行建议的SGLang程序。

图14：一个SGLang程序及其对应的数据流图。
图14：一个SGLang程序及其对应的数据流图。

图14b：图14a程序的计算图。三个流对应三个函数调用。
图14b：图14a程序的计算图。三个流对应三个函数调用。

编译器优化案例研究：通过代码移动改善前缀共享。我们探索了一种SGLang IR的编译优化：通过代码移动改善前缀共享。该优化旨在通过重新排序图中的节点来增加常量前缀的长度，从而改善前缀共享。这是一种激进的优化，因为它不严格保留原始计算。例如，将提示从“这是一个问题 + {问题}。请扮演数学专家并解决给定的问题。”更改为“请扮演数学专家并解决给定的问题。这是一个问题 + {问题}。”会产生更长的可共享前缀。我们尝试使用GPT-4来重新排序图节点。通过提供带有几个示例的提示来教GPT-4 SGLang IR的概念，我们发现GPT-4可以成功地为一些简单的SGLang程序应用此优化。在对20个提示模板的评估中，GPT-4在15个测试用例中的12个成功地重新排序了图节点而没有改变语义，平均使可共享前缀长度增加了60个token。失败的原因在于对图节点背后语义的错误理解。这项案例研究旨在探索使用GPT-4进行编译器优化，未来需要更多工作使其可靠。

方法细节中的引用汇总

以下是在方法细节章节（第2、3、4、5章）及相关附录中引用的参考文献及其上下文描述：

  • [4] Luca Beurer-Kellner, Marc Fischer, and Martin Vechev. Prompting is programming: A query language for large language models. (PLDI, 2023).

    • 引用段落: 第1章，引言部分将"Language Model Programs"的概念归因于此文献及[20]。
    • 引用段落: 第2章，对比部分将SGLang与LMQL和Guidance并列为低级LLM编程系统。

  • [12] In Gim, et al. Prompt cache: Modular attention reuse for low-latency inference. (arXiv, 2023).

    • 引用段落: 第3章，在讨论KV缓存重用时，提到现有系统（包括此文献）需要手动配置或无法处理所有重用模式。

  • [13] Guidance AI. A guidance language for controlling large language models. (GitHub, 2023).

    • 引用段落: 第2章，对比部分将SGLang与LMQL和Guidance并列为低级LLM编程系统。

  • [18] Jordan Juravsky, et al. Hydragen: High-throughput llm inference with shared prefixes. (arXiv, 2024).

    • 引用段落: 第3章，在讨论KV缓存重用时，提到现有系统（包括此文献）需要手动配置或无法处理所有重用模式。

  • [22] Michael Kuchnik, et al. Validating large language models with relm. (MLSys, 2023).

    • 引用段落: 附录B，解释约束解码的挑战时引用，说明字符串和token之间缺乏一一对应关系。

  • [23] Woosuk Kwon, et al. Efficient memory management for large language model serving with pagedattention. (SOSP, 2023).

    • 引用段落: 第3章，在讨论KV缓存重用时，提到vLLM等系统探索了简单的重用案例，但有限制。同时提到RadixAttention与PagedAttention兼容。
    • 引用段落: 附录A，在背景介绍中引用vLLM对基本前缀共享的研究。

  • [40] Swarnadeep Saha, et al. Branch-solve-merge improves large language model evaluation and generation. (arXiv, 2023).

    • 引用段落: 第2章，在介绍运行示例时，说明该示例使用了“分支-解决-合并”的提示方法。

  • [42] Ying Sheng, et al. Fairness in serving large language models. (arXiv, 2024).

    • 引用段落: 第3章，在讨论缓存感知调度时，承认贪心策略可能导致饥饿问题，并指出与公平调度方法的集成为未来工作。

  • [44] Mohammad Shoeybi, et al. Megatron-lm: Training multi-billion parameter language models using model parallelism. (arXiv, 2019).

    • 引用段落: 第3章，提到RadixAttention与张量并行等技术兼容。

  • [50] Vivien Tran-Thien. Fast, high-fidelity llm decoding with regex constraints. (2024).

    • 引用段落: 附录B，在讨论分词伪影和概率扭曲问题时引用，说明直接划分字符串可能改变语义，并指出概率分布倾斜的问题。

  • [51] Ashish Vaswani, et al. Attention is all you need. (NeurIPS, 2017).

    • 引用段落: 第3章及附录A，在解释KV缓存时，将其追溯到Transformer架构中的自注意力机制。

  • [54] Brandon T Willard and Rémi Louf. Efficient guided generation for large language models. (arXiv, 2023).

    • 引用段落: 第4章及附录B，在介绍约束解码时，引用其将正则表达式转换为FSM来指导生成的方法。

  • [58] Lu Ye, et al. Chunkattention: Efficient self-attention with prefix-aware kv cache and two-phase partition. (arXiv, 2024).

    • 引用段落: 第3章，在讨论KV缓存重用时，提到现有系统（包括此文献）需要手动配置或无法处理所有重用模式。

  • [60] Gyeong-In Yu, et al. Orca: A distributed serving system for {Transformer-Based} generative models. (OSDI, 2022).

    • 引用段落: 第3章，提到RadixAttention与连续批处理等技术兼容。