为Trellis引入RadixAttention

# 向 Trellis 引入 RadixAttention 来源：https://trellis.unfoldml.com/blog/radix-attention-intro 我们创建了 Trellis，旨在普及大语言模型（LLM）推理，同时不牺牲用户的数据隐私。为此，我们构建了一套系统，让用户能够在自己拥有和操作的硬件（如笔记本电脑、工作站和服务器）上进行部署。 为了贴近用户的实际环境，我们必须支持各种性能水平的硬件，因此尽可能抓住优化机会至关重要。 在这篇文章中，我们将重点介绍 Trellis 如何优化 LLM 推理中的 *预填充（prefill）* 阶段。首先简要介绍问题背景，然后讨论我们实现的一种适用于基于聊天和代理式 LLM 会话的缓存策略，最后通过基准测试结果进行总结。 ## LLM 预填充与 KV 缓存 Trellis 集群中的计算节点执行 LLM Transformer 模块（注意力机制、线性层等）的数值运算。在数值计算得出结果（即采样下一个 token）之前，我们必须为所有操作数填充系数；这就是 LLM 推理的 *预填充阶段*，它受计算能力限制。具体来说，要生成 token \\(t+1\\)，注意力机制需要所有之前 token（tokens \\(0 \ldots t\\)）的键和值（即投影后的嵌入向量）。 第一个关键观察：在自回归生成中，采样得到的 token 会被追加到已生成的 token 序列末尾，然后整个序列作为输入再次送入模型。这一点结合因果注意力掩码，意味着位置 \\(i < j\\) 的嵌入（因此也包括 K 和 V 矩阵）不依赖于位置 \\(k \geq j\\) 的后续值。因此，在生成所有后续 token 时，这些嵌入可以被**缓存**，并且缓存是只追加的。 ## 使用 RadixAttention 进行前缀缓存 第二个关键观察来自 LLM 的实际使用场景，即系统提示词先于用户请求出现的聊天会话：基于相同模板的重复 LLM 请求必须用公共前缀字符串预填充模型，这意味着我们可以预计算并缓存系统提示词的所有 token 嵌入。不仅如此，原则上我们还可以缓存提示词的**所有前缀**，只要知道这些前缀会在**跨请求**时被重用。RadixAttention (https://www.lmsys.org/blog/2024-01-17-sglang/) 的核心洞见是使用一种专门的数据结构——**基数树 (radix tree)** (https://en.wikipedia.org/wiki/Radix_tree)——它能在字符串前缀共享完整子串的常见情况下最小化存储（例如，对于字符串 "hello my name is Alice" 和 "hello my name is Bob"，基数树仅存储 3 个条目："hello my name is "、"Alice"、"Bob"）。 ## Trellis 中的 RadixAttention 在 Trellis 集群中，可以有多个节点发出并发推理请求，这些请求由多个节点提供服务，每个节点持有一个或多个 Transformer 模块。因此，KV 缓存机制也充当了并发机制，因为不相关的请求可能共享一个前缀（例如，如果团队复用相同的系统提示词）。 我们实现的 RadixAttention KV 缓存是**块式分页（block-paged）** 的：K/V 激活值以固定大小的**块**（token 嵌入）存储，这些块从共享的**池**中分配，每个会话维护一个有序的块 ID 列表，而不是连续的数组。当两个会话共享一个前缀时，第二个会话**获取（acquire）** 已有的前缀块（通过引用计数递增），而不是**分配（allocate）** 并重新计算新的块。 ## 基准测试 ## 设置与定义 以下所有基准测试均在单个 Mac M1 处理器上运行 **TinyLlama 1.1B Q4_0** 模型，使用合成提示。重复出现的参数包括： - **L** — 提示长度（token 数）。 - **D** — *解码预算（decode budget）*：每个会话生成的 token 数。 - **N** — 并发解码会话的数量。我们使用一个共享的“系统提示词”进行预热，再加上 N 个分叉的用户后缀，每个会话一个。 - **φ（phi）** — *共享前缀比例*：每个提示中公共前缀所占的比例。φ=0.5 表示每个提示的前半部分在所有调用中都是相同的。注意，φ 是**按块量化的**：内部会将共享区域四舍五入到整数个 32 token 的块（`prefixBlocks = round(φ·L/32)`），因此实际的 φ 会被约束为 32/L 的倍数。 - **op** — 一次完整的基准测试**迭代**：构建一个新的推理管道，预热共享前缀，并运行所有 N 个会话直至完成。因此“每次迭代的分配量”是整个迭代期间的堆流量，**不是**单个 token 的。 ## 吞吐量与共享前缀比例 当两个提示的公共前缀比例增大时，生成吞吐量和内存分配都会稳步改善。 提示长度 L: 512 tokens，解码预算 D: 每个会话 64 tokens 使用 RadixAttention 处理共享公共前缀的提示，能在生成吞吐量和内存分配方面获得多大提升？事实证明，提升相当可观。 **分配**曲线（红色轨迹）显示缓存按预期工作：随着 \\(\phi\\) 增加，每次迭代的堆分配次数单调下降，因为更大比例的前缀块是从池中**获取**的（引用计数递增），而不是**重新分配**的，而且这些块所覆盖的预填充位置不再需要重新计算。 令人欣慰的是，**吞吐量**曲线（蓝色）也随着 \\(\phi\\) 增加而线性改善，因为已缓存的 KV 位置无需预填充（\\(\phi\\) 为 0.75 和 0.9 时的波动可能是由于基准测试环境噪声引起的）。 ## 首 token 生成时间 后续共享公共前缀的请求的延迟也随着相同 token 比例的增加而线性改善。 这里我们想了解，当提示前缀跨 LLM 调用共享时，RadixAttention 在消除冗余计算方面的有效性。 TTFT 的测量方法是：先用一个提示预热缓存，然后为单个请求（并发数为 1）解码一个 token。TTFT 随着 \\(\phi\\) 增加稳步下降，因为更大的共享前缀意味着更多的位置由缓存块提供，更少的位置需要从头预填充。 阅读此图时需注意两点。首先，在 \\(\phi\\) 值较低时方差较大，这是首次迭代的伪像：第一次迭代会承受一次性冷启动成本。注意，即使在网格的最右端，可达到的最大命中率也是 `(L/32 − 1)/(L/32)`，而不是 1.0；探测提示始终在最后一个块上与预热提示分叉，因此有一个块永远不会共享。 ## 结论 在这篇文章中，我们展示了 RadixAttention 如何让 Trellis 节点运行速度提升 30-40%，同时内存使用量也大幅减少；随着 LLM 会话越来越长的趋势持续，这些进步很可能会产生更大的复合效应。 Trellis 仍处于早期阶段，而分布式、并发的 LLM 运行时的设计空间非常广阔。如果您以前思考过此类问题，并希望分享反馈，我们期待收到您的来信 (https://trellis.unfoldml.com/contact)！ ## 附录：选择缓存块大小 我们启发式地选择块大小，以平衡簿记开销和延迟。下图帮助我们选择了默认块大小 B = 32。 基数树缓存的叶子节点数与 TTFT 随缓存块大小的变化。误差线考虑了不同的目标序列长度。