Hugging Face: 从第一原则理解连续批处理

连续批处理是大型语言模型（LLM）服务的一项关键优化，通过并行处理多个对话并在它们完成后立即交换来最大化吞吐量。通过消除填充的需求并允许同时处理预填充和解码阶段，连续批处理使 AI 服务能够高效地处理数千个并发用户。

注意力机制和预填充阶段的作用

注意力机制是 LLM 生成的核心，它是网络中不同 token 相互作用的唯一部分。在标准前向传播中，模型处理输入序列（预填充阶段）以预测下一个 token。

此过程涉及将输入 token 投影到查询（$Q$）、键（$K$）和值（$V$）状态。然后模型通过 $Q K^{T}$ 计算相似度得分，其相对于序列长度 $n$ 的复杂度为二次方 $\mathcal{O}(n^{2} d)$。因果注意力掩码被应用以确保每个 token 仅与其之前的 token 相互作用，防止未来 token 对过去产生影响。

通过 KV 缓存优化生成

如果模型为每个新生成的 token 重新计算序列中每个之前 token 的 $K$ 和 $V$ 状态，那么逐个生成 token（解码阶段）在计算上会是浪费的。

KV 缓存 通过存储在预填充和之前解码步骤中创建的键和值状态来解决此问题。这将生成 token $n + 1$ 的计算成本从 $\mathcal{O}(n^{2})$ 降低到 $\mathcal{O}(n)$，用内存换取计算。例如，在 Llama-2-7B 中，存储一个 token 需要大约 16 KB 的内存，在 32 层和 32 头的 float16 精度下。

通过分块预填充处理长提示

当初始提示异常长时，激活所需的内存可能超过可用的 GPU 内存，使得无法在一次前向传播中完成预填充。

分块预填充 通过将初始提示分割成更小、易管理的块来解决此问题。通过利用 KV 缓存，模型存储第一个块的状态，并将其预置到后续块的状态之前，使得提示可以逐步处理而不丢失信息。

从批处理生成到连续批处理

为了提高吞吐量（每秒生成的 token 数），模型会并行处理多个提示。然而，传统方法会引入显著的低效率：

批处理生成： 要求所有提示具有相同的长度，因而需要填充标记（<pad>）。当提示长度不同时，这会造成浪费。
动态调度： 允许用新提示替换已完成的提示，但仍然需要大量填充，因为预填充阶段的新提示必须与解码阶段的其他提示一起批处理。

不规则批处理

为了消除填充浪费，不规则批处理 完全移除批处理轴，而是将所有提示连接成一个序列。通过注意力掩码严格控制 token 交互，确保一个提示的 token 不会与另一个提示的 token 相互作用。

连续批处理算法

连续批处理将不规则批处理与动态调度相结合，以充分利用 GPU。过程遵循以下逻辑步骤：

最大化内存预算： 系统尝试填充每批次 $m$ 个 token 的 GPU 内存预算。
优先处理解码： 当前处于解码阶段的所有提示首先被加入批次（每个占用一个 token）。
用预填充填充： 剩余空间由预填充阶段的提示填充，必要时使用分块预填充来拆分输入。
动态替换： 生成 <eos> token 的已完成提示会被立即移除，并用新的传入请求替换。

关键技术总结

| | --- | --- | | KV 缓存 | 避免在解码过程中重新计算过去的 token 表示。 | | 分块预填充 | 使得可以处理超过 GPU 内存限制的提示。 | | 不规则批处理 | 通过连接序列并使用掩码来消除填充浪费。 | | 动态调度 | 通过实时交换请求来保持高吞吐量。 |