[NeurIPS'24] SGLang: Efficient Execution of Structured Language Model Programs 论文阅读

发表于 分类于 阅读次数： 本文字数： 1.5k 阅读时长 ≈ 5 分钟

#Efficient KV Cache Reuse with RadixAttention

SGLang 程序可以通过 “fork” 原语串联多个生成调用并创建并行副本。此外，不同的程序实例通常共享某些公共部分（例如，系统 Prompt）。这些场景在执行过程中会生成许多共享的 Prompt 前缀，从而为重用 KVCache 提供了大量机会。在 LLM 推理过程中，KVCache 存储前向传递的中间 Tensor，用于解码未来的 Token。它们以自注意力机制中的键值对命名 [51]。 KVCache 的计算仅取决于前缀 Token。因此，具有相同 Prompt 前缀的请求可以重用 KVCache，从而减少冗余计算和内存使用。更多背景和一些示例在附录 A 中提供。

鉴于 KVCache 重用的机会，优化 SGLang 程序的一个关键挑战是在多个调用和实例之间重用 KVCache。虽然一些系统探索了某些 KVCache 重用案例 [23, 58, 18, 12]，但它们通常需要手动配置，并且无法处理所有重用模式（例如，动态树结构）。因此，大多数最先进的推理系统为每个请求重新计算 KVCache。我们将在第 7 节讨论它们的局限性以及我们的差异。

本节介绍 RadixAttention，这是一种在运行时自动且系统地复用 KVCache 的新技术。与在生成请求完成后丢弃 KVCache 的现有系统不同，我们的系统将 Prompt 和生成结果的缓存保留在 Radix Tree 中，从而实现高效的前缀搜索、复用、插入和逐出。我们实现了一种 LRU 逐出策略和一种缓存感知调度策略，以提高缓存命中率。 RadixAttention 与 Continuous Batching [ 60 ]、分页注意力 [23 ] 和 TP [ 44 ] 等技术兼容。此外，在未发生缓存命中时，它仅引入可忽略的内存和时间开销。

RadixAttention。Radix Tree 是一种数据结构，可作为经典 trie（前缀树）的空间高效替代方案。与典型的树结构不同，Radix Tree 的边不仅可以用单个元素 Token，还可以用长度不等的元素序列 Token，从而显著提升效率。在我们的系统中，我们利用 Radix Tree 管理 Token 序列与其对应 KVCache Tensor 之间的映射。这些 KVCache Tensor 以非连续、分页布局存储，其中每个页的大小等同于一个 Token。由于 KVCache 会迅速占满 GPU 内存，我们引入一种简单的 LRU 逐出策略，首先逐出最近最少使用的叶节点。通过首先逐出叶节点，我们使其共同祖先得以复用，直到这些祖先也成为叶节点并被逐出。

在 Continuous Batching 设置中，我们不能驱逐当前正在运行的批次所使用的节点。因此，每个节点维护一个引用计数器，指示有多少正在运行的任务正在使用它。如果节点的引用计数器为零，则该节点可以被驱逐。请注意，我们不预先分配一个固定大小的内存池作为缓存。相反，我们让缓存的 Token 和当前正在运行的任务共享同一个内存池。因此，系统动态地为缓存和正在运行的任务分配内存。当足够多的等待任务运行时，系统将驱逐所有缓存的 Token，以支持更大的批次大小。图 3 展示了对于多个传入请求如何维护 Radix Tree。

前端解释器将完整 Prompt 发送至运行时，由运行时执行前缀匹配与复用。树结构存储在 CPU 上，维护开销可忽略不计。在执行 fork 原语期间，前端首先将前缀作为 Prompt 信息发送，以确保前缀被正确插入树中。随后，前端发送其余 Prompt。这种 “前端 Hint” 简化了运行时调度与匹配，体现了前端-运行时协同设计的优势。

缓存感知调度。我们将缓存命中率定义为 cached prompt tokens / prompt tokens。当等待队列中有大量请求时，它们的执行顺序会显著影响缓存命中率。例如，如果请求调度器频繁在不同且不相关的请求之间切换，可能导致缓存抖动和较低的命中率。我们设计了一种缓存感知调度算法以提高缓存命中率。 在批处理设置中，我们对按匹配前缀长度对请求进行排序，并优先处理匹配前缀更长的请求，而不是采用先来先服务调度。 Alg. 1（附录）给出了采用 Continuous Batching 的缓存感知调度的伪代码。该算法采用最长共享前缀优先顺序。在对延迟更敏感的设置中，我们可能仍能容忍有限的批次重排序，以提高缓存复用率。此外，我们针对离线情形下的最优调度证明了以下定理。

定理 3.1。对于一批请求，我们可以通过以 DFS 顺序访问请求的 Radix Tree，并确保缓存大小 ≥ 最大请求长度，从而实现最优的缓存命中率。

证明见第 A.3 节（附录）。在线情况下，DFS 顺序将被破坏，但我们的调度仍然在扩展的全 Radix Tree 的部分上近似于 DFS 行为，如第 A.3 节所述。虽然贪婪缓存感知调度可以实现高吞吐量，但它可能导致饥饿。我们将与其他公平调度方法[42]的集成留作未来工作。

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