2. 大模型推理基础

大模型推理服务部署的核心挑战,源于模型架构本身的计算与访存特性。无论是部署一个7B参数的对话模型,还是面向高并发场景优化吞吐,工程师都必须深入理解Transformer的推理行为。本章将从Transformer架构的推理视角出发,剖析自回归生成机制、KV Cache的演进逻辑,以及推理阶段区别于训练阶段的关键计算特性,为后续的性能调优与工程部署奠定理论基础。

Transformer 推理阶段核心机制总览 输入 Token 序列 Token Embedding + 位置编码 Decoder 层 × N(如 32 层) 自注意力(Masked)→ 前馈网络 → LayerNorm 残差连接贯穿每层 输出 Token 概率 LM Head(词表映射) Softmax 采样 自回归循环:新 Token 拼接至输入序列,重复前向计算 KV Cache 机制 无 KV Cache 每步重算全部 K/V 计算量 O(n²) 有 KV Cache 缓存历史 K/V 矩阵 每步仅计算新 Token 推理阶段关键计算特性 Prefill 阶段 并行处理全部输入 Token 计算密集(GEMM) 生成首个 Token Decode 阶段 逐 Token 串行生成 访存密集(GEMV) 受限于内存带宽 Batch 影响 增大 Batch 摊薄权重读取 提升计算利用率 但受显存容量限制 显存瓶颈 KV Cache 随序列增长 权重 + 激活 + Cache 需精细显存规划

上图展示了Transformer推理时的完整数据流。输入Token经过嵌入层后,进入堆叠的Decoder层进行自注意力与前馈计算,最终通过输出头产生下一个Token的概率分布。自回归循环将新生成的Token重新拼接至输入序列,而KV Cache则在这一循环中扮演关键角色——它缓存了历史Token的Key和Value矩阵,避免了重复计算。下方三个色块分别对应Prefill阶段、Decode阶段以及Batch与显存的影响,这些正是推理服务部署时最核心的优化切入点。

2.1 Transformer架构的推理视角

核心原理:从训练到推理的范式转换

Transformer架构在训练阶段采用并行方式处理整个序列,通过Teacher Forcing机制一次性计算所有位置的输出。但在推理阶段,模型只能以自回归方式逐个生成Token——每一步的输出依赖于之前所有步骤的生成结果。这种串行依赖关系决定了推理无法像训练那样完全并行化。

从计算图角度看,推理时的Decoder层包含两个核心子模块:Masked Multi-Head Self-Attention和Feed-Forward Network(FFN)。自注意力机制通过Query、Key、Value三个线性变换,计算序列内部Token之间的关联权重;FFN则对每个Token独立进行两次线性变换和一次非线性激活。在推理时,这两个模块的计算模式截然不同:自注意力需要处理序列维度的交互,而FFN则是纯粹的元素级变换。

模块训练阶段推理阶段(Prefill)推理阶段(Decode)
自注意力并行计算全部位置并行计算全部输入Token仅计算新Token的注意力
FFN批量矩阵乘法(GEMM)大矩阵乘法单向量矩阵乘法(GEMV)
数据复用高(Batch大)中(Batch=1)低(每次仅1个Token)

判断方法:识别推理阶段的架构特征

在实际部署中,工程师需要快速判断一个模型在推理时的行为特征。关键观察点包括:

  • 层数(L)与隐藏维度(H):决定模型参数量与计算量,例如7B模型通常为32层、H=4096。
  • 注意力头数(A):影响并行计算粒度,多头机制允许在头维度上并行。
  • 词表大小(V):LM Head的映射矩阵维度为H×V,在输出阶段可能成为瓶颈。
  • 是否使用GQA/MQA:分组查询注意力(Grouped Query Attention)通过共享Key/Value头减少KV Cache占用,是近年推理优化的关键架构选择。
经验提醒:在评估一个模型的可服务性时,不要只看参数量。两个同为7B的模型,如果一个是MHA(Multi-Head Attention),另一个是GQA,其KV Cache占用可能相差数倍,直接影响最大并发数和上下文长度。

案例拆解:LLaMA-2-7B的推理架构分析

以LLaMA-2-7B为例,其架构参数为:32层Decoder、隐藏维度4096、32个注意力头、词表大小32000。在推理时,每个Token的生成需要经过32层的完整前向计算。假设输入序列长度为2048,Prefill阶段需要并行处理2048个Token的注意力计算,计算量约为2×L×H²×序列长度(约2.2 TFLOPs);而Decode阶段每生成一个Token,计算量仅为Prefill的1/2048,但需要从显存中读取全部7B参数(约14GB),这导致Decode阶段严重受限于显存带宽而非计算能力。

这一分析直接指导了部署策略:对于长对话场景,应优先考虑KV Cache的显存优化;对于高并发场景,应通过增大Batch来摊薄权重读取开销。

常见误区与总结

风险提示:初学者常误以为推理时模型会“记住”之前的计算结果。实际上,Transformer本身没有记忆能力,所有历史信息都通过KV Cache或重新计算来保留。如果KV Cache管理不当(如未及时清理),会导致显存碎片化或OOM。

理解Transformer的推理视角,核心在于认识到串行生成与并行训练的本质差异。这种差异决定了推理优化的所有方向:要么减少串行步数(如投机采样),要么降低单步开销(如KV Cache、量化),要么提高并行效率(如Continuous Batching)。

2.2 自回归生成机制

核心原理:逐Token的条件概率分解

自回归生成将序列生成问题分解为条件概率的连乘形式:P(y₁, y₂, ..., yₜ) = Π P(yₜ | y₁, ..., yₜ₋₁)。每一步,模型接收已生成的全部Token作为输入,输出下一个Token的概率分布。这种机制保证了生成序列的连贯性,但也带来了推理延迟的线性增长——生成N个Token需要N次前向计算。

在工程实现中,自回归循环包含三个关键步骤:(1)输入预处理:将新生成的Token进行Embedding并拼接位置编码;(2)前向计算:通过全部Decoder层得到logits;(3)采样策略:从概率分布中选择下一个Token。采样策略(如贪心、Top-K、Top-P、温度缩放)直接影响生成质量和计算开销。

操作流程:自回归生成的标准循环

# 伪代码:自回归生成循环
generated_tokens = [bos_token_id]
for step in range(max_new_tokens):
    # 1. 将当前序列转换为模型输入
    input_ids = torch.tensor([generated_tokens])
    # 2. 前向传播(使用KV Cache加速)
    logits = model(input_ids, past_key_values=kv_cache)
    # 3. 取最后一个位置的logits
    next_token_logits = logits[:, -1, :]
    # 4. 应用采样策略
    next_token = sample(next_token_logits, temperature=0.8, top_p=0.9)
    # 5. 更新序列和KV Cache
    generated_tokens.append(next_token)
    kv_cache = update_kv_cache(kv_cache, next_token)
    # 6. 终止条件判断
    if next_token == eos_token_id:
        break

案例拆解:对话场景中的自回归行为

考虑一个多轮对话场景:用户输入“介绍一下Transformer”,模型需要生成约200字的回答。假设模型生成速度为20 Token/s,则总延迟为10秒。在这10秒内,模型执行了200次前向计算,每次计算仅处理1个新Token,但需要读取全部模型权重。如果使用KV Cache,每次前向计算只需读取新Token对应的Key/Value和全部权重;如果不使用KV Cache,则每次都要重新计算所有历史Token的Key/Value,计算量将随序列长度平方增长。

实际部署中,我们通常设置最大生成长度(max_new_tokens)来限制循环次数,同时通过早停机制(遇到EOS Token)提前终止。对于长对话,还需考虑上下文窗口限制——当历史Token数超过模型最大长度时,需要截断或摘要压缩。

常见误区与优化方向

关键认知:自回归生成的延迟由两部分组成:固定开销(模型权重读取)和增长开销(KV Cache读取)。当序列较短时,固定开销占主导;当序列较长时,KV Cache读取成为瓶颈。因此,优化策略需要根据平均序列长度动态调整。

常见的自回归优化手段包括:投机采样(用小模型草稿+大模型验证)、并行解码(同时生成多个候选Token)、Medusa多头解码等。这些方法的核心思想都是打破严格的串行依赖,以计算换延迟。

2.3 KV Cache:推理加速的核心机制

核心原理:缓存历史Key/Value矩阵

在自注意力计算中,每个Token的Query会与所有历史Token的Key计算点积,再与对应的Value加权求和。对于第t步生成,需要用到第1到t-1步的Key和Value。如果不做缓存,每一步都要重新计算所有历史Token的Key和Value,计算复杂度为O(t²)。KV Cache的核心思想是:将已计算过的Key和Value矩阵存储在显存中,后续步骤直接复用,从而将单步计算复杂度降为O(t)。

KV Cache的存储结构为:每层、每个注意力头、每个Token对应一个Key向量和一个Value向量。对于LLaMA-2-7B(32层、32头、头维度128),每个Token的KV Cache大小为:2(K和V)× 32层 × 32头 × 128维 × 2字节(FP16)= 512KB。这意味着2048个Token的上下文需要1GB显存。

模型层数头数头维度每Token KV大小(FP16)4K上下文KV占用
LLaMA-2-7B3232128512 KB2 GB
LLaMA-2-13B4040128800 KB3.2 GB
Mistral-7B (GQA)328 (KV头)128128 KB512 MB

操作流程:KV Cache的分配与更新

在推理服务中,KV Cache的管理是性能优化的核心环节。标准流程如下:

  1. 预分配:根据最大序列长度和Batch大小,预先分配连续的显存块。
  2. Prefill填充:在首次前向计算时,计算所有输入Token的K/V并写入缓存。
  3. Decode追加:每生成一个新Token,计算其K/V并追加到缓存末尾。
  4. 释放与复用:序列生成结束后,释放对应缓存块供其他请求使用。

现代推理框架(如vLLM、TensorRT-LLM)采用PagedAttention技术,将KV Cache划分为固定大小的块(如16个Token/块),通过页表映射实现非连续显存管理,显著减少碎片化。

案例拆解:KV Cache对并发能力的影响

假设部署环境为单张A100-80GB,模型为LLaMA-2-7B(权重14GB)。可用显存约66GB。若每个请求的上下文长度为2048 Token,则每个请求的KV Cache占用1GB。理论上可同时处理约66个请求。但如果使用Mistral-7B(GQA架构),每Token KV Cache仅128KB,同样显存可支持约256个请求。这一差异直接决定了服务提供商的成本结构。

在实际调优中,我们发现当Batch大小从1增加到32时,Decode阶段的吞吐量提升约20倍,但KV Cache占用也线性增长。因此,Batch大小与KV Cache容量之间存在直接权衡,需要通过监控显存使用率动态调整。

常见误区与风险

风险提示:KV Cache并非越大越好。过大的KV Cache会挤占计算所需的显存空间,导致Batch被迫减小,反而降低整体吞吐。此外,KV Cache的分配和释放如果过于频繁,会产生显存碎片,长期运行后可用显存逐渐减少。

KV Cache的优化方向包括:量化(将FP16压缩为INT8,减少50%占用)、GQA/MQA架构(共享KV头)、窗口注意力(只保留最近N个Token的KV)、滑动窗口等。选择哪种方案取决于应用场景:长文档摘要需要完整上下文,而多轮对话可能只需要最近几轮的信息。

2.4 推理阶段的关键计算特性

核心原理:Prefill与Decode的算力鸿沟

推理过程可明确分为两个阶段:Prefill(预填充)和Decode(解码)。Prefill阶段并行处理用户输入的Prompt,计算量大但次数少(仅一次);Decode阶段逐Token生成,计算量小但次数多(等于生成Token数)。这两个阶段的计算特性截然不同:

  • Prefill:计算密集型。输入Token多,矩阵乘法维度大,GPU计算单元利用率高。例如处理2048 Token的Prompt,计算量相当于生成2048个Token的Decode总量。
  • Decode:访存密集型。每次仅处理1个Token,矩阵乘法退化为矩阵向量乘(GEMV),计算单元利用率极低(通常低于5%),主要瓶颈是显存带宽。

这一差异导致一个反直觉的现象:增加Batch大小对Decode阶段的吞吐提升远大于Prefill阶段。因为Batch增大后,多个请求的GEMV可以合并为GEMM,计算利用率显著提升。

判断方法:识别推理瓶颈

在实际部署中,可以通过以下指标判断当前系统处于计算瓶颈还是访存瓶颈:

指标计算瓶颈特征访存瓶颈特征
GPU利用率接近100%通常低于30%
显存带宽利用率较低接近峰值
增加Batch吞吐线性提升吞吐提升有限
延长序列长度延迟线性增长延迟超线性增长

案例拆解:Batch调优的实测数据

我们在A100-80GB上对LLaMA-2-7B进行Batch调优实验。固定输入Prompt为512 Token,生成256 Token。实验结果如下:

Batch=1:  延迟 380ms, 吞吐 2.6 Token/s, GPU利用率 8%
Batch=8:  延迟 520ms, 吞吐 15.4 Token/s, GPU利用率 35%
Batch=32: 延迟 780ms, 吞吐 41.0 Token/s, GPU利用率 72%
Batch=64: 延迟 1100ms, 吞吐 58.2 Token/s, GPU利用率 88%
Batch=128: 延迟 1800ms, 吞吐 71.1 Token/s, GPU利用率 95%

从数据可见,Batch从1增至32时,吞吐提升约15倍;但从32增至128时,仅提升1.7倍。这是因为当Batch足够大时,计算逐渐成为瓶颈,继续增大Batch只会线性增加延迟。实际部署中,我们选择Batch=32作为平衡点,兼顾延迟和吞吐。

常见误区与总结复盘

核心结论:推理优化的本质是将访存密集的Decode计算转化为计算密集的GEMM。实现手段包括:增大Batch、使用Continuous Batching(动态Batch)、融合内核(减少内存访问次数)、量化权重(减少访存量)。

另一个常见误区是忽略显存带宽的物理限制。即使计算利用率达到100%,Decode阶段的Token生成速度仍受限于权重读取速度。例如,7B模型FP16权重14GB,A100的显存带宽为2TB/s,理论最大Token生成速度为2TB/s ÷ 14GB ≈ 143 Token/s。这解释了为什么量化(如INT8)能将速度提升近一倍——权重减半,带宽瓶颈减半。

总结本章,Transformer推理的核心挑战在于:自回归的串行特性决定了延迟的下限,KV Cache的显存占用决定了并发能力的上限,Prefill与Decode的计算特性差异决定了优化策略的着力点。工程师需要在这三个维度之间寻找平衡,才能构建高效、稳定的推理服务。