LLM推理和优化（1）：基本概念介绍

一、LLM推理的核心过程：自回归生成

LLM（如DeepSeek、ChatGPT、LLaMA系列等）的推理本质是自回归生成：从初始输入（如[CLS]或用户prompt）开始，逐token预测下一个词，直到生成结束符（如[EOS]）。其核心分为两个阶段：

1. Initialization阶段（初始化）

• 目标：准备第一个token的生成条件。
• 关键步骤：
  • 输入编码：将初始prompt转换为token序列（如[CLS]你好），嵌入为向量x_0。
  • 初始隐藏状态：通过Transformer的编码器（或直接使用预训练参数）生成第一层的隐藏状态h_0。
  • KV Cache初始化：为每一层的每个注意力头创建空的Key/Value缓存（形状：[batch, heads, seq_len, head_dim]）。此时seq_len=0，因为尚无历史token。

示例：生成首词“今天”时，输入为[CLS]，初始化后仅计算第一层的h_0，KV Cache为空。

在LLM推理中，Initialization阶段（初始化阶段）又称“预填充阶段”（Prefill Stage）。这一命名源于其核心功能：为后续的逐token生成预填充（Prefill）KV Cache和初始隐藏状态。

工程实现

Hugging Face的transformers库、NVIDIA的FasterTransformer均采用prefill和generation区分这两个阶段。例如：

# 伪代码：Hugging Face生成逻辑
outputs = model.prefill(prompt)  # 预填充KV Cache（Initialization）
for _ in range(max_new_tokens):
    outputs = model.generate_step(outputs)  # 解码阶段，逐token生成

术语对比：Initialization vs Prefill

场景	常用术语	含义侧重
学术描述	Initialization	强调“初始化隐藏状态和缓存”
工程实践	Prefill	强调“预填充固定长度的输入”
用户视角	输入处理阶段	对应“用户输入的prompt处理”

本质是同一阶段，但“Prefill”更直观反映了其“为生成提前准备历史KV”的工程目标。

2. Decoding阶段（解码）

• 目标：逐token生成，每步复用历史计算结果。
• 核心逻辑（以生成第t个token为例）：
  1. 当前token处理：将第t-1步生成的token嵌入x_t，与前一步隐藏状态拼接，输入Transformer层。
  2. 注意力计算优化：
     • 查询（Query）：仅计算当前token的Query向量Q_t（因为只关注当前位置）。
     • 键值（Key/Value）：复用KV Cache中的历史Key/Value，并追加当前token的Key_t、Value_t。
     • 注意力得分：计算Q_t与所有历史Key的相似度（仅需一次矩阵乘法，而非重复全量计算）。
  3. 更新KV Cache：将当前层的Key_t、Value_t追加到缓存中（seq_len += 1）。
  4. 生成概率：通过LM头输出第t个token的概率分布，选择下一词（贪心/采样）。

3. 举个栗子

• 输入：用户prompt“请写一首诗：”（4个token）。
• Prefill阶段：
  1. 计算这4个token的所有层Key/Value，填充到KV Cache（此时缓存长度=4）。
  2. 生成第一个待扩展的隐藏状态（对应第4个token的输出）。
• Decoding阶段：
  逐句生成诗句，每步：
  1. 计算当前token的Q（仅1个token）。
  2. 复用Prefill的4个KV + 之前生成的KV，计算注意力。
  3. 追加当前token的KV到缓存（缓存长度逐步增加到4+N）。

通过“预填充”，避免了每次生成新token时重复计算prompt的KV，这正是LLM实现高效推理的关键优化之一。

二、原始Transformer的效率瓶颈：O(n²)的重复计算

• 时间复杂度：训练时并行计算所有token的注意力（O(n²)），但推理时需自回归生成，每步需重新计算所有历史token的Key/Value，导致总复杂度为O(n³)（n为序列长度）。
• 空间复杂度：每次推理需保存所有中间层的Key/Value，内存占用随n线性增长，长文本（如n=4k）时显存爆炸。
• 现实痛点：生成1000字的文章需重复计算百万次注意力，传统Transformer无法支持实时交互。

三、KV Cache：用空间换时间的核心优化

1. 方法本质

缓存历史层的Key/Value，避免重复计算。每个Transformer层维护独立的KV Cache，存储该层所有已生成token的Key/Value向量。

2. 具体实现步骤（以单batch为例）

1. 初始化缓存（t=0）：

   • 每层创建空缓存：K_cache = [], V_cache = []（形状：[num_layers, heads, 0, head_dim]）。

2. 第t步生成（t≥1）：

   • 前向传播：输入当前token嵌入，通过Transformer层计算当前层的Q_t, K_t, V_t。
   • 拼接缓存：

     K_cache[t_layer] = torch.cat([K_cache[t_layer], K_t], dim=2)  # 在seq_len维度追加
     V_cache[t_layer] = torch.cat([V_cache[t_layer], V_t], dim=2)
     

   • 注意力计算：

     attn_scores = Q_t @ K_cache[t_layer].transpose(-2, -1)  # Q_t: [1, heads, 1, d], K_cache: [1, heads, t, d]
     attn_probs = softmax(attn_scores / sqrt(d)) @ V_cache[t_layer]  # 仅需O(t)计算
     

   • 更新隐藏状态：将注意力输出传入下一层，直到LM头生成token。

3. 循环：重复步骤2，直到生成[EOS]或达到最大长度。

3. 优化效果

• 时间：每步注意力从O(n²)→O(n)，总复杂度O(n²)（接近线性）。
• 空间：缓存占用O(n)（每层存储历史K/V），但避免了重复计算的中间变量，实际显存节省50%+。
• 典型案例：LLaMA-2 70B在4k序列长度下，KV Cache使推理速度提升4倍（NVIDIA官方数据）。

四、延伸：KV Cache的局限性与改进

• 显存瓶颈：长上下文（如100k token）的KV Cache占用巨大（每层约4k token×4byte×2（KV）≈32KB，64层×100k≈2GB）。
• 优化方向：
  • 分页缓存（Paged Attention）：NVIDIA提出，用非连续内存存储KV，减少碎片化（2023年突破）。
  • 动态缓存：仅保留最近相关token的KV（如检索增强LLM）。

KV Cache是LLM落地的基石，其设计思想（复用历史计算）贯穿现代推理优化（如FlashAttention、QLoRA），最终实现了从“实验室模型”到“实时对话”的跨越。