自回归解码中的挑战

大型语言模型，尤其是那些基于 Transformer 架构的模型，通过自回归过程生成文本。这表示它们基于目前已生成的所有标记来预测序列中的下一个标记。尽管功能强大，但这种顺序的、步进式生成方法带来了显著的计算和内存瓶颈，特别是在处理长序列或为实时应用程序提供服务时。

生成的顺序特性

基本的挑战源于自回归解码的核心定义。为了生成第 $i$ 个标记，模型需要将前 $i-1$ 个已生成的标记序列作为输入。这就产生了固有的顺序依赖性：在已知当前标记之前，无法计算下一个标记。

请看一个 PyTorch 中的简化生成循环：

import torch
import torch.nn.functional as F

def generate_next_token(model, input_ids):
    """给定当前序列，生成下一个标记。"""
    with torch.no_grad():
        # 获取模型在最后一个标记位置的对数几率
        outputs = model(input_ids)
        next_token_logits = outputs.logits[:, -1, :]

        # 应用 softmax 获取概率
        probs = F.softmax(next_token_logits, dim=-1)

        # 采样下一个标记（例如，使用贪婪解码）
        next_token_id = torch.argmax(probs, dim=-1).unsqueeze(-1)
    return next_token_id

# 使用示例
model = ... # 您的预训练自回归大型语言模型
tokenizer = ... # 您的分词器

prompt = "The quick brown fox"
input_ids = tokenizer.encode(prompt, return_tensors="pt")

max_new_tokens = 50
for _ in range(max_new_tokens):
    next_token = generate_next_token(model, input_ids)
    # 检查是否是序列结束标记
    if next_token.item() == tokenizer.eos_token_id:
        break
    # 添加新标记并继续
    input_ids = torch.cat([input_ids, next_token], dim=-1)

generated_text = tokenizer.decode(input_ids[0])
print(generated_text)

此循环展示了迭代过程。每次调用 generate_next_token 都会对模型执行一次完整的正向传播（或其很大一部分）。对于长度为 $N$ 的序列，生成 $M$ 个新标记需要 $M$ 次单独的正向传播。这种顺序执行直接导致了延迟；总时间大致与生成的标记数量成正比。

内存带宽限制

Transformer 中的一个主要瓶颈，尤其是在推理期间，是内存带宽。在自回归生成的每个步骤中，模型都需要计算注意力分数。标准的自注意力机制需要计算当前标记的查询向量与序列中所有先前标记的键向量之间的相互关系。

为了有效进行此操作，为所有先前标记计算的键 ($K$) 和值 ($V$) 通常保存在 GPU 内存中。随着序列变长，这个 $K/V$ 缓存的大小会随序列长度 ($n$) 线性增长。在步骤 $i$ 的注意力计算需要将这些缓存的 $K$ 和 $V$ 张量（大小与 $i$ 成比例）从高带宽内存 (HBM) 加载到计算单元（如 Tensor Cores）中。

$$\text{Attention}(Q_i, K_{1..i}, V_{1..i}) = \text{softmax}\left(\frac{Q_i K_{1..i}^T}{\sqrt{d_k}}\right) V_{1..i}$$

在这里，$Q_i$ 是在步骤 $i$ 处生成的标记的查询，而 $K_{1..i}$ 和 $V_{1..i}$ 代表从 1 到 $i$ 的所有标记的键和值。在每个步骤中从内存中获取 $K_{1..i}$ 和 $V_{1..i}$ 会消耗大量的内存带宽。对于大型模型和长序列，数据在 HBM 和计算单元之间传输所花费的时间可能会主导整体延迟，使得该过程受内存限制而非计算限制。

自回归生成中的顺序依赖和不断增加的内存访问。每一步都需要加载先前计算的键和值。

计算成本的扩展

尽管内存带宽通常是单标记生成延迟的主要瓶颈，但计算成本本身也不容忽视。自注意力计算每层的计算复杂度为 $O(n^2 d)$，$n$ 为序列长度，$d$ 为模型的隐藏维度。尽管像 KV 缓存（下文讨论）这样的技术避免了对过去标记重新计算键和值，但在步骤 $i$ 的注意力分数计算仍然涉及到当前查询 $Q_i$ 与所有先前键 $K_{1..i}$ 之间的点积计算，使得该部分计算的成本与 $i \times d$ 成正比。随后与值向量 $V_{1..i}$ 的乘法也以类似方式扩展。

将所有层和前馈网络的计算相加，每个生成步骤都需要大量的浮点运算 (FLOPs)，并且这个数量会随着当前序列长度的增加而增加。对于非常长的序列，这种计算需求增加了整体延迟。

吞吐量方面的挑战

自回归解码的顺序特性也给吞吐量带来了挑战，吞吐量是指所有并发请求在单位时间内生成的输出标记数量。尽管可以在并行批次中处理多个独立的生成请求，但每个序列内部的生成仍然是顺序的。这限制了硬件能够被充分利用的程度，特别是当序列长度差异很大或单个请求需要非常长的输出时。优化批处理策略变得重要，但它并没有从根本上改变每个序列的顺序约束。

这些是与标准自回归解码过程相关的固有挑战。顺序依赖性、内存带宽限制、扩展的计算成本和吞吐量约束使得本章后续部分论述的专门优化技术成为必需。KV 缓存、优化的注意力实现以及推测解码等策略直接针对这些瓶颈，以使大型语言模型推理更快、更有效。