从输入提示到看到第一个词出现之间发生了什么

本文是大语言模型工作原理系列的第四部分，重点介绍训练好的模型在实际推理时发生了什么。当用户输入一个提示并按下回车后，模型需要逐步生成回复，而不是一次性输出全部内容。

自回归生成

与训练时模型能看到整个序列并并行预测每个位置的下一词元不同，推理时模型只能一个一个地生成词元。具体过程如下：用户输入提示，分词器将其转换为词元ID；然后运行一次前向传播，模型输出整个提示之后下一个词元的logits；从这些logits中采样一个词元，将其追加到序列中；再运行一次前向传播，这次输入是原始提示加上新生成的词元；如此重复，直到达到所需长度或遇到停止词元。这就是自回归生成——每一步都依赖前一步的输出，因此必须串行进行。

例如，用户输入“What is gravity?”，分词后为[What, is, gravity, ?]（4个词元）。第一步输入这4个词元，前向传播后得到下一个词元的概率分布，“Gravity”概率最高，采样得到“Gravity”。第二步输入[What, is, gravity, ?, Gravity]（5个词元），得到“is”，依此类推，逐步生成“Gravity is the force…”直到停止。每一步模型都要重新处理整个已有序列，这导致大量冗余计算。

预填充与解码阶段

实际推理分为两个阶段。

预填充阶段：一次性并行处理整个提示。例如50个词元的提示只需一次前向传播，GPU同时计算所有词元通过每一层的注意力模式（因果掩码确保只关注前面的词元）。预填充的真正目的是计算注意力，为每个词元生成键（Key）和值（Value）向量，这些向量直接存入KV缓存供解码阶段复用。短提示很快，但超长提示（数万词元）可能因注意力计算复杂度与序列长度成平方关系而变慢。

解码阶段：逐个生成新词元，每次前向传播只处理一个新词元。GPU每次仅做少量矩阵乘法，瓶颈从计算转向内存带宽——大部分时间花在从内存加载模型权重而非做数学运算。这就是为什么GPU内存带宽（数据从显存读取的速度）对解码性能比原始算力更重要。典型的聊天场景中，100词元提示+500词元回复，预填充很快，解码的500次顺序前向传播才是用户感知到的延迟。这也是回复流式输出一个词元一个词元出现的原因——每个词元生成后立即发送，无需等待全部完成。

KV缓存：跨步骤节省工作量

回忆解码阶段每次都要处理到目前为止的整个序列。在第100步，词元1到99与第99步时完全一样，注意力机制却重复计算它们的键向量和值向量，造成浪费。KV缓存通过存储每个已计算词元的K和V向量来避免重复：在第100步，只计算新词元的K和V，复用缓存中词元1到99的向量。

预填充时，模型处理完整提示并将所有词元的K和V向量存入缓存。解码时，生成词元51只需运行一次前向传播，计算其K和V，然后与缓存的词元1-50的向量拼接，再执行注意力。如此反复，缓存不断增长但旧词元的向量永不重算。

KV缓存的内存占用

以Llama 3.1-8B为例：32层、8个KV头（分组查询注意力）、每个头维度128、FP16精度（每个值2字节），支持最长128,000词元的上下文。计算单个词元在单个层单个头的KV：一个Key向量128个数×2字节=256字节，Value同样，所以一个词元在单个头单个层需512字节。每层有8个头，因此一个词元在一层需512×8=4,096字节。32层共4,096×32=131,072字节≈128KB。对于4,096词元的序列：128KB×4,096=512MB。若达到128,000词元上下文，则约16GB。服务多个用户时需乘以并发序列数，例如32用户各4,096词元，KV缓存需16GB，与模型权重本身（约16GB）相当。这就是分组查询注意力（GQA）的意义——将KV头数从32减至8，缓存大小降至1/4。虽然可能会损失一些细粒度注意力模式，但实际内存节省远大于质量差异。

KV缓存节省的工作量

无缓存时生成1000个词元需从头计算K/V向量：第1步计算1个，第2步2个，……总计约500,000次K/V计算。有缓存后每个词元的K/V只计算一次，总共1000次。但注意力点积的数量仍然是1+2+…+1000≈500,000，这一O(n²)成本不变。KV缓存主要节省了计算K/V向量的矩阵乘法，而无缓存时模型需要为每个历史词元在每一层重复执行昂贵的矩阵运算。

解码策略

模型输出每个词元的概率后，还需要通过解码策略来决定实际选择哪个词元。常见的策略包括贪婪解码（总是选概率最高的词元）、温度采样（调整概率分布的尖锐程度）、top-k采样（只从概率最高的k个词元中采样）和top-p（nucleus）采样（从累积概率超过p的最小集合中采样）。这些策略平衡了回复的创造性与连贯性。在实际应用中，开发者会根据任务需求调整这些参数，例如聊天场景可能使用较高的温度以增加多样性，而事实性问答则倾向于低温度或贪婪解码。

总之，大语言模型的推理是一个复杂但精妙的过程，从输入提示到最终回复的每一步都涉及大量工程优化。理解预填充、解码、KV缓存和解码策略，有助于开发者更好地利用和优化这些模型。