Connect with us

人工智能

SGLang:结构化语言模型程序的高效执行

mm
SGLang: Efficient Execution of Structured Language Model Programs

大型语言模型(LLM)越来越多地被用于需要多次生成调用、先进的提示技术、控制流和结构化输入/输出的复杂任务。然而,用于编程和执行这些应用程序的高效系统却缺乏。SGLang是一个新引入的系统,旨在通过提供复杂语言模型程序的高效执行来解决这个问题。SGLang由一个前端语言和一个运行时组成。前端通过提供用于生成和并行控制的原语来简化编程,而运行时通过像RadixAttention用于KV缓存重用和压缩有限状态机用于更快的结构化输出解码这样的新型优化来加速执行。实验表明,SGLang在各种大型语言和多模态模型上实现了最高6.4×的吞吐量,相比之下,包括agent控制、逻辑推理、少次学习基准、JSON解码、检索增强生成管道和多回合聊天在内的任务都有所涉及。

最近在LLM能力方面的进展已经扩大了其实用性,使其能够处理更广泛的通用任务,并作为自主代理运行。在这些应用中,LLM参与多轮规划、推理和与外部环境的交互。这是通过工具使用、多个输入模式和各种提示技术(如少次学习、自一致性、思维框架和思维树)来实现的。这些新用例需要多个、通常依赖的LLM生成调用,表明使用多调用结构来完成复杂任务的趋势。

这种转变标志着从简单聊天到更复杂的程序化使用LLM的转变,在这种转变中,程序调度和控制LLM的生成过程。这些程序被称为“语言模型程序”(LM程序)。高级提示技术和代理工作流属于LM程序的范畴。LM程序有两个共同的特性:(1)LM程序通常涉及多个LLM调用,伴随着控制流来完成复杂任务并提高整体质量。(2)LM程序接收结构化输入并产生结构化输出,使得LM程序可以组合和集成到现有的软件系统中。

在本文中,我们将更深入地探讨SGLang框架,探索其架构,分析其性能,并将其与最先进的框架进行比较。那么让我们开始吧。

SGLang简介

尽管LM程序被广泛使用,但当前用于表达和执行它们的系统仍然效率低下。SGLang确定了两大挑战,主要与LM程序的高效使用相关联:

  • 编程复杂性:开发LM程序很繁琐且困难,因为LLM的非确定性性质。这涉及广泛的字符串操作、实验性提示调优、脆弱的输出解析、处理多个输入模式以及实现并行机制。这种复杂性大大降低了甚至简单程序的可读性。
  • 执行效率低下:执行LM程序效率低下是由于冗余计算和内存使用。最先进的推理引擎,旨在减少延迟和提高吞吐量,缺乏对工作负载的直接了解,导致了显著的低效率。一个值得注意的例子是Key-Value(KV)缓存的重用,KV缓存由可用于生成推理的可重用中间张量组成。当前系统缺乏有效的机制来促进KV缓存在共享公共前缀的多个LLM调用之间的重用,导致不必要的计算和浪费的内存。此外,对于结构化输出(如JSON模式)的约束解码是次优的,因为现有系统只解码一个令牌一次。

为了解决这些挑战,SGLang引入了一种用于LLM的结构化生成语言。核心思想是系统地利用LM程序中的多调用结构以实现高效执行。如以下图所示,SGLang由两个部分组成:前端语言和后端运行时。

前端简化了LM程序的编程,而运行时则加速了其执行。这些部分可以协同工作以获得更好的性能,也可以独立运行。

SGLang是嵌入在Python中的特定领域语言,提供用于生成(例如extend、gen、select)和并行控制(例如fork、join)的原语。它与Python的控制流和库兼容,允许用户使用本地Python语法轻松开发高级提示工作流。SGLang包括一个解释器和一个编译器。解释器管理提示状态作为一个流,并将原语操作提交给流以异步执行,确保适当的控制和程序内并行。另外,SGLang程序可以被跟踪和编译以进一步优化。SGLang的运行时提出几种新型优化以加速LM程序的执行:

  • RadixAttention:该技术实现了在多个生成调用之间自动重用KV缓存。在现有的推理引擎中,请求的KV缓存在处理后被丢弃,防止在多个调用之间重用并减慢执行。SGLang在一个基数树中维护一个KV缓存的LRU缓存,管理KV缓存作为一个传统缓存,并使用基数树进行高效匹配、插入和驱逐。这使得运行时能够高效地处理各种重用模式。
  • 压缩有限状态机:该技术实现了更快的结构化输出约束解码。现有的系统仅遵循下一个令牌的约束,使它们只能一次解码一个令牌。相反,SGLang分析约束并构建一个压缩的有限状态机来表示它们,将多令牌路径压缩为单步路径,允许一次解码多个令牌以获得更快的速度。
  • API推测执行:对于仅API模型(如OpenAI的GPT-4),SGLang引入了API推测执行来优化多调用程序。

使用SGLang,各种LLM应用程序都被实现,包括agent控制、逻辑推理、少次学习基准、JSON解码、检索增强生成管道、多回合聊天和多模态处理。性能在包括Llama-7B/70B、Mistral-8x7B、LLaVA-v1.5-7B(图像)和LLaVA-NeXT-34B(视频)在NVIDIA A10G和A100 GPU上的模型上进行了测试。实验结果表明,SGLang在广泛的工作负载、模型和硬件设置中实现了最高6.4×的吞吐量,相比之下,包括Guidance、vLLM和LMQL在内的现有编程和推理系统都有所涉及。

SGLang:编程模型和方法论

SGLang编程模型通过一个运行示例介绍,描述其语言原语和执行模式,并概述运行时优化机会。该模型简化了多调用工作流中繁琐的操作(例如字符串操作、API调用、约束规范和并行),并提供了灵活且可组合的原语。SGLang是嵌入在Python中的特定领域语言。以下图显示了一个使用分支-求解-合并提示方法评估图像文章的程序。

函数multi_dimensional_judge需要三个参数:`s`、`path`和`essay`。s管理提示状态,path是图像文件路径,essay是文章文本。新的字符串和SGLang原语可以使用+=操作符追加到状态s以进行执行。首先,函数将图像和文章添加到提示中。然后,它使用select检查文章是否与图像相关,并将结果存储在s[“related”]中。如果相关,则提示被分叉为三个副本以进行并行评估,使用gen将结果存储在f[“judgment”]中。接下来,它合并判断、生成摘要并分配字母成绩。最后,它以JSON格式返回结果,遵循由正则表达式约束regex定义的模式。SGLang极大地简化了这个程序,因为使用OpenAI API接口的等效程序需要2.1×的代码行数,主要是由于手动字符串操作和并行控制。

SGLang提供了控制提示状态、生成和并行的原语,可以与Python语法和库一起使用。以下是原语:

gen:调用模型进行生成,并将结果存储在变量中,变量名称指定在其第一个参数中。它支持一个`regex`参数来约束输出以遵循由正则表达式定义的语法(例如JSON模式)。

  • select:调用模型从列表中选择最高概率选项。
  • +=或extend:将字符串追加到提示中。
  • [变量名称]:检索生成的结果。
  • fork:创建提示状态的并行分支。
  • join:重新连接提示状态。
  • image和video:接受图像和视频输入。

执行SGLang程序的最简单方法是通过解释器,其中提示被视为异步流。像extend、gen和select这样的原语被提交给流以进行异步执行。这些非阻塞调用允许Python代码在生成完成之前继续运行,类似于异步启动CUDA内核。每个提示由一个流执行器在后台线程中管理,实现程序内并行。检索生成结果将阻塞直到它们准备好,确保正确的同步。或者,SGLang程序可以被编译为计算图并使用图执行器执行,允许更多优化。本文默认使用解释器模式,并在附录D中讨论编译器模式结果。SGLang支持带有其自身SGLang运行时(SRT)的开放权重模型,以及API模型,如OpenAI和Anthropic模型。

LLM的编程系统可以分为高级(例如LangChain、DSPy)和低级(例如LMQL、Guidance、SGLang)。高级系统提供预定义或自动生成的提示,例如DSPy的提示优化器。低级系统通常不会修改提示,但允许直接操作提示和原语。SGLang是一个低级系统,类似于LMQL和Guidance。以下表格比较了它们的功能。

SGLang更侧重于运行时效率,并带有自己的协同设计运行时,允许进行新型优化。高级语言(例如DSPy)可以被编译为低级语言(例如SGLang)。在DSPy中将SGLang作为后端用于更好的运行时效率的集成在后面进行了演示。

上面的示例演示了RadixAttention在九个时间点上的操作,展示了基数树在响应不同请求时的动态演变。这些请求包括两个聊天会话、一个少次学习查询批次和自一致性采样。每个树边都带有一个标签,表示子字符串或令牌序列。节点被着色以反映不同的状态:绿色表示新添加的节点,蓝色表示在时间点访问的缓存节点,红色表示已被驱逐的节点。

步骤1:基数树最初是空的。

步骤2:服务器处理传入的用户消息“Hello”并用LLM输出“Hi”进行响应。系统提示“你是一个有帮助的助手”,用户消息“Hello!”和LLM回复“Hi!”被合并到树中作为一个单一的边,链接到一个新节点。

步骤3:到达一个新的提示,服务器在基数树中找到提示的前缀(即对话的第一轮),并重用其KV缓存。新的轮次被追加到树中作为一个新节点。

步骤4:开始一个新的聊天会话。步骤3中的节点被分割为两个节点,以便两个聊天会话可以共享系统提示。

步骤5:第二个聊天会话继续。然而,由于内存限制,步骤4中的一个节点必须被驱逐。新的轮次被追加在步骤4中剩余的节点之后。

步骤6:服务器接收到一个少次学习查询,处理它,并将其插入到树中。根节点被分割,因为新的查询与现有节点没有共享任何前缀。

步骤7:服务器接收到一批额外的少次学习查询。这些查询共享相同的一组少次学习示例,因此步骤6中的一个节点被分割以实现共享。

步骤8:服务器接收到来自第一个聊天会话的新消息。它驱逐所有来自第二个聊天会话的节点,因为它们是最近最少使用的。

步骤9:服务器接收到一个请求,用于对步骤8中的问题进行更多答案采样,可能是用于自一致性提示。为了为这些请求腾出空间,多个节点被驱逐。

这个例子演示了RadixAttention如何处理不同请求对节点的动态分配和驱逐,以确保KV缓存重用和内存管理的效率。

SGLang:评估和结果

开放权重模型的结果

延迟和吞吐量结果如下图所示。SGLang提高了最高6.4×的吞吐量,并降低了最高3.7×的延迟。这些改进是由于KV缓存重用、单个程序内并行的利用以及更快的约束解码所致。

在这些基准测试中,缓存命中率从50%到99%不等。图13(附录)列出了所有基准测试的实现和最优缓存命中率,显示SGLang的缓存感知调度在平均上接近96%的最优命中率。

大型模型的结果

更大的模型Mixtral-8x7B和Llama-70B在相同的基准测试中使用张量并行进行了测试,结果如下图所示。在更大的模型上,SGLang的优化表现出类似的趋势,表明其优化对更大的模型具有良好的泛化性。由于Guidance和LMQL缺乏对张量并行的高效实现,因此它们被省略。

多模态模型的结果

SGLang本身支持多模态模型,具有图像和视频原语。该论文中的优化与多模态模型兼容。对于RadixAttention,计算输入图像的哈希并将其用作基数树中的键,允许重用来自同一图像的图像令牌的KV缓存。LLaVA-v1.5-7B(图像)在llava-bench-in-the-wild上运行,LLaVA-NeXT-34B(视频)在ActivityNet上运行。由于这些模型没有被其他基准系统很好地支持,因此使用了Hugging Face Transformers中模型作者的原始实现作为基准。如以下表格所示,SGLang在这些基准测试中提供了最高6×的吞吐量。在llava-bench-in-the-wild中,同一图像的多个问题被处理,SGLang运行时重用了KV缓存。

生产部署

SGLang已被部署在Chatbot Arena中,以服务开放权重模型。由于某些模型的流量较低,因此每个模型仅由一个SGLang工作者提供服务。在一个月后,LLaVA-Next-34B的RadixAttention缓存命中率为52.4%,Vicuna-33B的缓存命中率为74.1%。缓存命中来自公共系统消息、频繁重用的示例图像和多轮聊天历史。这种情况降低了Vicuna-33B的平均首令牌延迟1.7×。

最后的思考

在本文中,我们讨论了SGLang,这是一个新引入的系统,旨在通过提供复杂语言模型程序的高效执行来解决这个问题。SGLang由一个前端语言和一个运行时组成。前端简化了编程,提供了用于生成和并行控制的原语,而运行时则通过新型优化加速了执行,例如RadixAttention用于KV缓存重用和压缩有限状态机用于更快的结构化输出解码。实验表明,SGLang实现了最高6.4×的吞吐量,相比之下,包括agent控制、逻辑推理、少次学习基准、JSON解码、检索增强生成管道和多回合聊天在内的任务都有所涉及。

专业为工程师,心为作家。 Kunal是一名技术作家,对AI和ML有着深厚的热爱和理解,致力于通过其引人入胜和信息丰富的文档来简化这些领域中的复杂概念。