📚 课程目录第4章大模型推理原理

大模型推理原理：自回归生成、KV缓存与束搜索

大模型（如GPT、LLaMA、Qwen等）的推理过程与传统的分类模型或嵌入模型有本质区别。它并非一次前向计算直接输出最终答案，而是通过自回归（Autoregressive）方式逐token生成文本。理解这一底层机制，是后续学习TensorRT-LLM中连续批处理（In-flight Batching）、KV缓存复用、PagedAttention等高级优化的前提。本节将拆解自回归生成的核心循环、KV缓存为何是性能关键，以及束搜索如何在“质量”与“开销”之间做权衡。

4.1 自回归生成：逐Token的循环

自回归模型的核心假设是：下一个token的概率分布仅依赖于此前已生成的所有token。用公式表达即：

P(y_1, y_2, ..., y_T) = ∏_{t=1}^{T} P(y_t | y_1, y_2, ..., y_{t-1})

在推理时，模型将已生成的token序列作为输入，前向计算一次，得到词表上每个token的logits，经Softmax后取概率最高（或按采样策略选择）的token作为新输出，然后将其拼接到输入序列末尾，进入下一轮。这个过程反复执行，直到遇到结束符<EOS>或达到最大长度限制。

关键点：每一步生成都依赖上一步的输出，因此推理过程是串行的，无法像训练那样对序列内部做并行化。这决定了单次请求的延迟下限——至少需要执行T次前向计算（T为输出token数）。

下面这张图展示了自回归生成的前三步循环，注意每一步输入序列都在增长，而模型权重保持不变。

从图中可以看到，第2步的输入包含了第1步新生成的"on"。这意味着每一步的Attention计算量都在增加——因为序列长度变长了。如果不做任何优化，生成第100个token时，需要对前99个token的Key和Value重新计算一遍Attention，这造成了巨大的算力浪费。

4.2 KV缓存：避免重复计算的工程基石

仔细观察Attention公式：Attention(Q,K,V) = softmax(QK^T/√d)V。在自回归生成的第t步，Q来自当前最新token，而K和V来自所有历史token。历史token的K、V在之前的步骤中已经计算过，且它们不会随新token的到来而改变（因为前面的token内容已固定）。

因此，一个自然的优化是：在每一步计算后，将当前token的K和V保存到缓存中。下一步只需计算新token的K、V，然后与缓存中已有的K、V拼接，即可完成Attention计算。这就是KV缓存（KV Cache）。

对比项	无KV缓存	有KV缓存
第t步计算量	对前t个token全部重新计算K、V	仅计算第t个token的K、V
时间复杂度	O(t²)（累计）	O(t)（累计）
显存占用	无额外缓存	每层需存储2×t×d×h字节
实际推理速度	极慢，几乎不可用	提升数十倍以上

KV缓存的大小计算方式如下（以单层为例）：

KV缓存大小 = 2 (K和V) × 序列长度 × 隐藏层维度 × 层数 × 字节数(如FP16为2字节)

以LLaMA-7B为例：隐藏维度4096，层数32，若生成序列长度为2048，则单请求KV缓存占用约为：

2 × 2048 × 4096 × 32 × 2 = 1,073,741,824 字节 ≈ 1GB

重要误区：KV缓存是显存消耗大户，而非模型权重。在TensorRT-LLM中，KV缓存通常需要单独分配显存池，并通过PagedAttention等机制动态管理。如果并发请求数较多，KV缓存总占用可能远超模型权重本身。

4.3 束搜索：在质量与开销间权衡

自回归生成时，每一步若只取概率最高的token（贪心解码），容易陷入局部最优，产生重复或平庸的文本。束搜索（Beam Search）通过维护多个候选序列（束），在每一步扩展时保留概率最高的前K个序列，最终选择整体概率最高的序列作为输出。

束搜索的流程如下：

初始化：以起始token开始，计算第一个token的概率分布，取前K个token作为K个初始束。
扩展：对当前每个束，分别计算下一个token的概率分布，每个束扩展出V个候选（V为词表大小），共得到K×V个候选序列。
剪枝：计算所有候选序列的累积对数概率（log概率相加），保留最高的K个序列作为新的束。
终止：当某个束生成<EOS>或达到最大长度时，将其放入完成列表；当完成列表中有K个序列或达到最大束数时停止。

束宽度K是核心超参数：K越大，搜索空间越大，生成质量通常越好，但计算开销也成倍增加（每一步需计算K×V次前向）。实际应用中，K=4或K=8是常见选择。

案例：假设K=2，词表V=5。第1步得到概率最高的两个token为"我"(0.4)和"你"(0.3)。第2步对"我"扩展出5个候选，对"你"也扩展出5个候选，共10个候选序列。计算各自累积概率后，保留最高的2个序列，例如"我爱"(0.4×0.5=0.2)和"你是"(0.3×0.6=0.18)。后续步骤重复此过程。

束搜索与贪心解码的对比总结如下：

贪心解码：每步只取最高概率token，速度快，但可能错过全局最优组合。
束搜索：保留多个候选路径，生成质量更高，适合翻译、摘要等对准确性要求高的任务。
采样（Sampling）：按概率分布随机采样，增加多样性，适合对话、故事生成等创造性任务。

4.4 常见误区与复盘

在实际部署和优化中，以下几个问题值得特别关注：

误区一：认为KV缓存是“可选优化”。实际上，没有KV缓存的Transformer自回归生成在长序列下计算复杂度为O(n²)，几乎无法实际使用。TensorRT-LLM中KV缓存是默认且强制的基础设施。

误区二：混淆“模型参数量”与“推理显存需求”。推理时显存占用 = 模型权重 + KV缓存 + 激活值 + 临时缓冲区。对于长序列、大并发场景，KV缓存往往占主导地位。例如，服务32个并发请求、每个序列长度2048的LLaMA-7B，仅KV缓存就需要约32GB显存。

复盘建议：在评估推理性能时，不要只看“每秒生成token数”，还要关注首token延迟（TTFT）和每token延迟（TPOT）。TTFT主要受Prompt处理（Prefill）影响，TPOT主要受KV缓存读取和Attention计算影响。TensorRT-LLM的调优目标通常是降低TPOT，同时通过连续批处理提高整体吞吐。

束搜索虽然能提升质量，但需注意：在TensorRT-LLM中，束搜索需要为每个束单独维护KV缓存，因此显存开销会乘以束宽度K。如果服务端显存有限，可考虑使用贪心解码+温度采样的折中方案，或使用长度惩罚来减少重复。