吞吐量提升5倍，联合设计后端系统和前端语言的LLM接口来了

原文来源：机器之心

图片来源：由无界 AI生成

大型语言模型 (LLM) 越来越多地用于需要多个链式生成调用、高级 prompt 技术、控制流以及与外部环境交互的复杂任务。然而，用于编程和执行这些应用程序的现有高效系统存在着明显的缺陷。

现在，开源社区的研究者们面向 LLM 提出了一种结构化生成语言（Structured Generation Language）——SGLang。SGLang 能够增强与 LLM 的交互，通过联合设计后端运行时系统和前端语言，使 LLM 更快、更可控。机器学习领域知名学者、CMU 助理教授陈天奇还转发了这项研究。

总的来说，SGLang 的贡献主要包括：

在后端，研究团队提出了 RadixAttention，这是一种跨多个 LLM 生成调用的 KV 缓存（KV cache）复用技术，自动且高效。 在前端，研究团队开发了一种嵌入 Python 的、灵活的域指定（domain-specific）语言来控制生成过程。该语言可以在解释器模式或编译器模式下执行。

后端前端组件协同工作，可提高复杂 LLM 程序的执行和编程效率。

该研究使用 SGLang 实现了常见的 LLM 工作负载，包括智能体、推理、提取、对话和小样本学习任务，并在 NVIDIA A10G GPU 上采用 Llama-7B 和 Mixtral-8x7B 模型。如下图 1 、图 2 表明，与现有系统（即 Guidance 和 vLLM）相比，SGLang 的吞吐量提高了 5 倍。

图 1：不同系统在 LLM 任务上的吞吐量（A10G、FP16 上的 Llama-7B、张量并行度 = 1）

图 2：不同系统在 LLM 任务上的吞吐量（A10G、FP16 上的 Mixtral-8x7B，张量并行度 = 8）

后端：使用 RadixAttention 自动 KV 缓存复用

在 SGLang 运行时的开发过程中，该研究发现了复杂 LLM 程序的优化关键 ——KV 缓存复用，当前系统对此处理不佳。KV 缓存复用意味着具有相同前缀的不同 prompt 可以共享中间 KV 缓存，避免冗余的内存和计算。在涉及多个 LLM 调用的复杂程序中，可能存在各种 KV 缓存复用模式。下图 3 说明了 LLM 工作负载中常见的四种此类模式。虽然某些系统能够在某些场景下处理 KV 缓存复用，但通常需要手动配置和临时调整。此外，由于可能的复用模式的多样性，即使通过手动配置，现有系统也无法自动适应所有场景。

图 3：KV 缓存共享示例。蓝色框是可共享的 prompt 部分，绿色框是不可共享的部分，黄色框是不可共享的模型输出。可共享的部分包括小样本学习示例、自洽（self-consistency）问题、多轮对话中的对话历史以及思维树（tree-of-thought）中的搜索历史。

为了系统地利用这些复用机会，该研究提出了一种在运行时自动 KV 缓存复用的新方法 —— RadixAttention。该方法不是在完成生成请求后丢弃 KV 缓存，而是在基数树（radix tree）中保留 prompt 和生成结果的 KV 缓存。这种数据结构可以实现高效的前缀搜索、插入和驱逐。该研究实现了最近最少使用（LRU）驱逐策略，并辅以缓存感知调度策略，以提高缓存命中率。

基数树可作为 trie（前缀树）节省空间的替代方案。与典型的树不同，基数树的边缘不仅可以用单个元素来标记，还可以用不同长度的元素序列来标记，这提高了基数树的效率。

该研究利用基数树来管理映射，这种映射是在充当键的 token 序列和充当值的相应 KV 缓存张量之间进行的。这些 KV 缓存张量以分页布局存储在 GPU 上，其中每个页的大小相当于一个 token。

考虑到 GPU 内存容量有限，无法重新训练无限的 KV 缓存张量，这就需要驱逐策略。该研究采用 LRU 驱逐策略，递归地驱逐叶节点。此外，RadixAttention 与连续批处理和分页注意力等现有技术兼容。对于多模态模型，RadixAttention 可以轻松扩展以处理图像 token。

下图说明了在处理多个传入请求时如何维护基数树。前端总是向运行时发送完整的 prompt，运行时会自动进行前缀匹配、复用和缓存。树形结构存储在 CPU 上，维护开销较小。

图 4. 采用 LRU 驱逐策略的 RadixAttention 操作示例，分九个步骤进行说明。