StreamSparkAI

概览/核心摘要 (Executive Summary)

本讲座深入探讨了语言模型推理（Inference）的关键概念、挑战与优化策略。推理是将训练好的固定模型应用于实际任务（如聊天机器人、代码补全、模型评估、强化学习中的样本生成）以根据提示生成响应的过程。其效率至关重要，因为与一次性的训练成本相比，推理会重复多次。核心挑战在于，与训练时可以并行处理整个序列不同，推理（尤其是Transformer模型）通常需要按顺序生成token，导致计算资源利用率低，且主要受内存带宽限制（Memory-bound），特别是在token生成阶段。

讲座详细分析了Transformer模型在推理时的计算强度，区分了“Prefill”（编码提示，计算受限）和“Generation”（逐个生成token，内存受限）两个阶段。MLP层的计算强度与批处理大小（B）相关，而Attention层的计算强度在生成阶段近似为1，无法通过批处理改善。因此，减少KV缓存大小成为提升推理效率的关键，方法包括分组查询注意力（GQA）、多头潜在注意力（MLA）、跨层注意力（CLA）和局部注意力。

此外，讲座还探讨了Transformer的替代架构，如状态空间模型（SSM，如Mamba、Jamba）、线性注意力模型（BASED、MiniMax-01）和扩散模型，这些架构通过改变根本的计算方式（如将O(T)的KV缓存变为O(1)的状态，或并行生成token）来寻求推理效率的巨大提升。其他优化技术包括量化（降低数字精度，如fp8, int8, LLM.int8()）、模型剪枝（移除不重要部分并蒸馏）和推测采样（使用小模型预生成，大模型验证，保证无损）。最后，讨论了处理动态工作负载的系统级优化，如连续批处理和PagedAttention（借鉴操作系统分页思想管理KV缓存，减少碎片化）。核心目标是在保证模型准确性的前提下，通过架构创新、系统优化和算法技巧，最大限度地提升推理的速度和吞吐量。

理解推理工作负载

概览

推理在多种场景中都扮演着核心角色：
* 实际应用：如聊天机器人、代码补全、批量数据处理。
* 模型评估：例如，评估模型的指令遵循能力。
* 测试时计算：指在输出最终答案前进行更多的“思考”或推理。
* 通过强化学习进行训练：需要生成样本，然后根据奖励进行评分。

效率的重要性：与训练的一次性成本不同，推理会重复执行多次，因此效率至关重要。
* Sam Altman曾提到，OpenAI 当时每天生成约1000亿个单词。
* Aman Sanger提到，Cursor 当时每天编写近10亿行被接受的代码。

指标

衡量推理性能的关键指标包括：
* 首个token生成时间 (Time To First Token, TTFT)：用户在看到任何生成内容之前需要等待的时间，对交互式应用非常重要。
* 延迟 (Latency)：以秒/token为单位，表示token向用户显示的速度，同样对交互式应用很重要。
* 吞吐量 (Throughput)：以tokens/秒为单位，衡量单位时间内生成的token总数，对批量处理应用尤其有用。

效率的关键考量

• 训练 (监督式)：可以访问所有token，允许在序列上进行并行化（例如Transformer中的矩阵乘法）。
• 推理：必须按顺序生成token，因为每个token的生成依赖于之前所有token，这使得并行化变得困难，难以充分利用计算资源，并导致其通常是内存受限 (memory limited)。

进行推理的公司与开源包

许多公司和组织都在关注并优化推理效率：
* 封闭模型提供商：OpenAI, Anthropic, Google等。
* 开放权重模型提供商：Together, Fireworks, Deepinfra等。
* 开源软件包：
* VLLM (Berkeley)
* TensorRT (NVIDIA)
* TGI (Hugging Face)

Transformer 回顾

讲座回顾了Transformer模型中的符号及其维度含义，以便后续分析：
* B: 批处理大小 (batch size)
* L: 层数 (number of layers)
* S: 序列长度 (查询) (sequence length - query)
* T: 序列长度 (键值) / 生成token数 (sequence length - key/value or tokens to generate/query)
* V: 词汇表大小 (vocab size)
* D: d_model，模型嵌入维度 (embedding dimension)
* F: MLP 隐藏层维度 (通常为 $4 \cdot D$)
* H: 注意力头维度 (attention head dimension)
* N: 查询头数量 (number of query heads, $N \cdot H = D$)
* K: 键/值头数量 (number of key/value heads)
* G: 每个键值头对应的查询头数量 ($N = K \cdot G$)

前馈传递的FLOPs (浮点运算次数) 约为：$6 \cdot (B \cdot T) \cdot (\text{num_params} + O(T))$。

计算强度回顾

计算强度 (Arithmetic Intensity) 定义为每字节内存传输所做的计算量（FLOPs/Byte），是衡量计算受限与内存受限的关键。
* 矩阵乘法示例 ($X (B \times D) \cdot W (D \times F)$):
* FLOPs: $2 \cdot B \cdot D \cdot F$
* 传输字节数: $2 \cdot B \cdot D + 2 \cdot D \cdot F + 2 \cdot B \cdot F$ (假设bf16，每元素2字节)
* 强度: $\frac{B \cdot D \cdot F}{B \cdot D + D \cdot F + B \cdot F}$
* 若B远小于D和F，则强度简化为 $\approx B$。
* 加速器计算强度 (H100)：
* flops_per_second $\approx 989 \times 10^{12}$
* memory_bandwidth $\approx 3.35 \times 10^{12}$ bytes/sec
* accelerator_intensity $\approx 295$
* 结论：当矩阵乘法的强度 $B > 295$ 时，计算受限（理想情况）；否则，内存受限（不理想）。
* 极端情况 ($B=1$)：对应矩阵-向量乘积，计算强度为1，严重内存受限。这与推理中逐个token生成的情况类似。

推理的计算强度

朴素采样 (Naive Sampling)

为生成每个token，将包括历史记录在内的整个序列输入Transformer。复杂度高，生成T个token需要 $O(T^3)$ FLOPs (因每次前馈是 $O(T^2)$)。存在大量冗余计算。

带 KV 缓存的采样

通过存储和复用先前token的键（Key）和值（Value）向量（即KV缓存），可以显著减少冗余计算。KV缓存通常存储在高带宽内存（HBM）中。
推理包含两个阶段：
1. Prefill (预填充)：处理输入的提示（prompt），将其编码为向量。此阶段可以像训练一样并行化处理序列中的所有token，通常是计算受限的。
2. Generation (生成)：逐个生成新的响应token。此阶段是串行的，因为每个新token的生成都依赖于之前所有token（包括已生成的），通常是内存受限的。

讲座中用 S 表示条件依赖的token数量（如prompt长度），用 T 表示查询或生成的token数量。在Prefill阶段，$T$ 可以理解为prompt的长度；在Generation阶段，$T=1$。

MLP 层 (仅关注矩阵乘法)

• FLOPs: $6 \cdot B \cdot T \cdot D \cdot F$
• 传输字节数: $4 \cdot B \cdot T \cdot D + 4 \cdot B \cdot T \cdot F + 6 \cdot D \cdot F$
• 强度: $\frac{6 \cdot B \cdot T \cdot D \cdot F}{4 \cdot B \cdot T \cdot D + 4 \cdot B \cdot T \cdot F + 6 \cdot D \cdot F}$。若 $B \cdot T$ 远小于 D 和 F，则强度 $\approx B \cdot T$。
  • Prefill ($T$ 较大): 强度 $B \cdot T$ 可以很大，易于实现计算受限。
  • Generation ($T=1$): 强度 $\approx B$。需要较大的并发请求数（B）才能提高计算强度。

Attention 层 (关注使用 FlashAttention 的矩阵乘法)

• FLOPs: $4 \cdot B \cdot S \cdot T \cdot D$
• 传输字节数: $4 \cdot B \cdot T \cdot D + 4 \cdot B \cdot S \cdot D$
• 强度: $\frac{4 \cdot B \cdot S \cdot T \cdot D}{4 \cdot B \cdot T \cdot D + 4 \cdot B \cdot S \cdot D} = \frac{S \cdot T}{S+T}$。
  • Prefill ($T=S$): 强度 $\approx S/2$。如果序列长度S足够长，可以是计算受限的。
  • Generation ($T=1$): 强度 $\approx S/(S+1) \approx 1$。这是非常低的强度，导致内存受限。
• 与MLP层的关键区别：Attention层的强度不依赖于批处理大小B。
  • 原因：MLP层中，每个序列共享相同的MLP权重；而在Attention层中，每个序列拥有其独立的KV缓存，导致内存访问随B线性增加，抵消了FLOPs随B增加带来的强度提升潜力。

总结：Prefill阶段通常是计算受限的。Generation阶段是内存受限的。在Generation阶段，MLP层的强度为B（依赖并发请求数），而Attention层的强度始终接近1（无法通过批处理改善）。

计算 Transformer 统计数据

分析Transformer模型的内存占用、延迟和吞吐量。
* 参数数量 (num_params)：基于 V, D, L, F, N, K, H 计算。讲座中引用了一个Python函数中的公式：num_params = 2*V*D + L*( D*F*3 + 2*D*K*H + 2*D*N*H ) (此公式可能包含词嵌入和输出层，以及每层MLP和Attention的参数，但具体细节可能因模型实现而异)。
* 参数所需内存 (parameter_size)：num_params * 2 (假设使用bf16，每个参数2字节)。推理时不需要梯度和优化器状态。
* KV缓存大小 (kv_cache_size)：每个序列需要存储的KV缓存大小。讲座中公式为 S * D * L * 2 * 2 (S:序列长度, D:模型维度, L:层数, 第一个2代表K和V, 第二个2代表bf16的字节数)。
* 总内存使用 (memory)：B * kv_cache_size + parameter_size。
* 延迟 (Latency)：由于生成阶段是内存受限的，延迟主要由内存I/O决定：memory / memory_bandwidth。
* 吞吐量 (Throughput)：B / Latency (因为并行生成B个token)。

示例：Llama 2 13B on an H100

使用具体配置 ($S=1024, D=5120, F=13824, N=40, K=40, H=128, L=40, V=32000$, memory_bandwidth=$3.35 \times 10^{12}$ bytes/sec) 分析不同批处理大小（B）的影响：
* Batch size = 1：延迟约 80ms/token，吞吐量约 124 tokens/sec。
* Batch size = 64：内存增加，延迟增加，但吞吐量显著提高。
* Batch size = 256：内存需求可能超出单个H100的容量（如示例中240GB > 80GB HBM）。即使能装下，吞吐量增益也会出现递减。

延迟与吞吐量的权衡：
* 小批处理B：低延迟，低吞吐量。
* 大批处理B：高延迟，高吞吐量。
* 简单并行性：启动M个模型副本，延迟不变，吞吐量增加M倍。
* 更难的并行性：对模型和KV缓存进行分片。

首个token生成时间 (TTFT) 主要由Prefill阶段决定。通常在Prefill阶段使用较小B以优化TTFT，在Generation阶段使用较大B以优化吞吐量。

减少 KV 缓存大小

由于内存是推理瓶颈，减少KV缓存大小是关键优化方向，需注意不显著牺牲准确性。

分组查询注意力 (Grouped-Query Attention, GQA)

• 多头注意力 (MHA)：N个查询头，N个键头，N个值头。
• 多查询注意力 (MQA)：所有N个查询头共享单个键头和值头 ($K=1$)。表达能力可能不足。
• 分组查询注意力 (GQA)：介于MHA和MQA之间。N个查询头，K个键/值头 ($K<N$)，每 $G=N/K$ 个查询头共享一组键/值头。
• 效果：GQA将KV缓存大小减少了 $N/K$ 倍，从而可以提高吞吐量或使用更大的批处理大小。Ainslie+ 2023研究表明，GQA能在保持与MHA相当准确性的同时显著减少推理时间。Llama 2的70B模型和Llama 3均采用了GQA。

多头潜在注意力 (Multi-Head Latent Attention, MLA)

• 核心思想：将每个键向量和值向量从 $N \cdot H$ 维度投影到一个更低的C维度。
• 示例 (DeepSeek v2)：将 $N \cdot H = 16384$ 降至 $C=512$。由于MLA与RoPE不兼容，需额外增加维度支持RoPE (如 $512+64=576$)。
• 效果：显著减少KV缓存，带来类似GQA的延迟/吞吐量改进。DeepSeek v2论文声称MLA在准确性上可能略优于MHA，且成本更低。

跨层注意力 (Cross-Layer Attention, CLA)

• 核心思想：在不同层之间共享KV缓存，类似于GQA在头之间共享。
• 效果 (Brandon+ 2024)：在准确性和KV缓存大小（进而影响延迟和吞吐量）的帕累托前沿上有所改进。

局部注意力 (Local Attention)

• 核心思想：只关注最近的k个token（滑动窗口注意力），使得KV缓存大小与序列长度无关（保持固定）。
• 问题：可能损害长距离依赖建模的准确性。
• 解决方案：将局部注意力层与全局注意力层交错使用（混合层）。例如，Character AI 每6层使用1个全局注意力层。

总结：减少KV缓存的目标是通过低维KV缓存（GQA, MLA）、共享KV缓存（CLA）或在部分层使用局部注意力来实现，同时尽量不损害模型准确性。

Transformer 的替代方案

探索超越标准Transformer架构以大幅改进推理效率的可能性。

状态空间模型 (State-Space Models, SSM)

• S4：基于经典状态空间模型 ($ \dot{x} = Ax + Bu, y = Cx + Du $)，通过离散化实现快速计算。擅长处理合成的长上下文任务，但在对语言重要的联想回忆任务上表现不佳。
• Mamba：允许SSM参数依赖于输入，改进了SSM在语言任务上的表现，在1B参数规模上能与Transformer相当。
• Jamba (AI21 Labs)：混合Transformer和Mamba层（如1:7比例），并使用MoE（混合专家模型）。

基于线性注意力的模型

• 核心思想：通过泰勒展开等方式将Attention的复杂度从二次降为线性。
• BASED：使用线性Attention + 局部Attention。
• MiniMax-01：使用线性Attention + 少量完全Attention的MoE模型。
• 要点：线性+局部Attention（通常仍需少量完全Attention）已产生SOTA模型。将 $O(T)$ 的KV缓存替换为 $O(1)$ 的状态，极大提升推理效率。

扩散模型 (Diffusion Models)

• 核心思想：并行生成所有token（非自回归），然后通过多步迭代优化来修正整个序列。从随机噪声开始，逐步去噪生成文本。
• 优势：并行生成，可能带来极高的token输出速度。
• 示例 (Inception Labs - Mercury Coder)：在编码基准测试中显示出远超Transformer类模型的速度，同时保持有竞争力的准确性。
• 挑战：在文本生成上应用比图像生成更复杂，但潜力巨大。

量化 (Quantization)

核心思想：降低模型权重和/或激活值的数字精度，以减少内存占用和传输，从而提升内存受限的推理速度。

数字格式比较

• fp32 (4字节)：通常用于训练。
• bf16 (2字节)：推理的默认格式。
• fp8 (1字节)：如H100上的e4m3格式，范围约(-240, 240)。可用于训练（需谨慎）和推理。
• int8 (1字节)：范围(-128, 127)。比fp8便宜，但精度可能稍逊，主要用于推理。
• int4 (0.5字节)：范围(-8, 7)。更便宜，但精度损失更大。

量化方法：
* 量化感知训练 (QAT)：在训练过程中引入量化操作，通常效果较好但实现复杂，需重新训练。
* 训练后量化 (PTQ)：对已训练好的模型进行量化，通常在少量校准数据上确定缩放因子和零点。

LLM.int8() (Dettmers+ 2022)

• 问题：标准量化（按绝对值最大值缩放）易受异常值（outliers）影响，导致精度损失。
• 解决方案：识别并分离出异常值，这些异常值以fp16精度处理，其余大部分值进行int8量化。
• 效果：在保持较小准确性损失的情况下实现量化，但可能比纯fp16推理慢15-23%（因混合精度处理）。

激活感知量化 (Activation-aware Quantization)

• 思想：根据激活值的大小来选择哪些权重（通常0.1-1%）保留高精度，其余进行更低精度量化。
• 示例：fp16 -> int3 可带来显著内存降低和加速。

模型剪枝 (Model Pruning)

核心思想：移除模型中不重要的部分（如层、头、隐藏单元），然后通过蒸馏等方法修复模型性能，以达到更小、更快的目的。

NVIDIA 算法 (Muralidharan+ 2024)

1. 在小型校准数据集上识别重要的模型组件（层、头、隐藏维度）。
2. 移除不重要的组件，得到一个更小的剪枝模型。
3. 将原始大模型的知识蒸馏到剪枝后的小模型中。
4. 结果示例：Minitron 8B（由Nemotron-4 15B剪枝而来）在MMLU上优于Gemma 7B和Mistral 7B，且训练成本大幅降低。

推测采样 (Speculative Sampling)

利用“检查（verifying）已生成的token序列比从头生成（generating）单个token更快”这一特性来加速推理，同时保证采样结果与目标大模型完全一致（无损）。
* 回顾推理阶段：Prefill（并行编码，计算受限，可输出概率） vs. Generation（串行生成，内存受限）。
* 核心思想 (Leviathan+ 2022, Chen+ 2023)：
1. 使用一个更小、更便宜的草稿模型 (draft model) $p$ 快速生成一小段（如k个）候选token。
2. 使用目标大模型 (target model) $q$ 并行地评估这k个候选token的概率。
3. 根据一定的接受/拒绝准则（类似于重要性采样或Metropolis-Hastings中的步骤，如以 $q(t_i|H)/p(t_i|H)$ 的概率接受token）来决定接受哪些草稿token。如果某个token被拒绝，则从目标模型 $q$ 中重新采样该位置的token。
* 关键特性：数学上保证最终生成的序列精确采样自目标模型 $q$。
* 实践：
* 目标模型可以是70B，草稿模型可以是几B（如6B或1B）。
* 通常希望草稿模型尽可能接近目标模型（可通过蒸馏实现）。
* 效果：在XSum、HumanEval等基准上显示出显著加速（如1.9-2倍）。
* 扩展：
* Medusa：草稿模型并行生成多个候选token序列。
* EAGLE：草稿模型利用目标模型的高级特征来辅助生成。
* 总结：精确采样，利用检查与生成的不对称性，草稿模型设计有很大创新空间。

处理动态工作负载

实际服务中，请求到达时间不一、序列长度各异、可能共享前缀，给高效批处理带来挑战。

连续批处理 (Continuous Batching)

• 问题：传统批处理等待一批请求满才开始，或因序列长度不一导致填充（padding）浪费。
• 解决方案：迭代级调度。
  • 逐步解码，每生成一个token后，调度器检查是否有新请求到达，可将其加入当前批处理，无需等待整个批次完成。

选择性批处理

• 问题：批处理通常要求张量维度一致，但不同请求长度不同。
• 解决方案：
  • Attention计算通常需要按序列分别处理（因KV缓存是序列特定的）。
  • 非Attention计算（如MLP层）可以将不同长度序列的激活值连接（concatenate）成一个大张量进行批处理，因为这些权重是共享的。

PagedAttention (VLLM, Kwon+ 2023)

借鉴操作系统中的虚拟内存和分页思想来管理KV缓存，以解决内存碎片化问题。
* 先前状态的问题：为每个请求预分配最大长度的KV缓存，导致：
* 内部碎片化：若实际生成长度远小于预分配长度，则浪费空间。
* 外部碎片化：已分配块之间可能存在无法利用的小块空闲空间。
* PagedAttention解决方案：
* 将序列的KV缓存划分为固定大小的、非连续的块 (blocks)。
* 逻辑上连续的KV序列可以映射到物理内存中不连续的块。
* 共享KV缓存：
* 共享系统提示（System Prompt）。
* 一个提示采样多个响应时，共享提示部分的KV缓存。
* 通过写时复制 (copy-on-write)机制在块级别实现高效共享和修改。当需要修改共享块时，才为其创建副本。
* 其他VLLM优化：
* 融合Kernel（如融合块读取和Attention计算）以减少Kernel启动开销。
* 使用最新的优化Kernel（如FlashAttention, FlashDecoding）。
* 使用CUDA Graphs进一步减少Kernel启动开销。
* 总结：利用操作系统思想（分页）高效管理内存，处理动态工作负载。

总结 (讲座核心观点)

• 推理对于语言模型的实际应用、评估和部分训练过程至关重要。
• 推理与训练在特性上存在显著差异：推理通常是内存受限的，且面临动态工作负载的挑战。
• 提升推理效率的技术多种多样，包括：
  • 新模型架构：如GQA, MLA, CLA, 局部注意力, SSM, 线性注意力, 扩散模型。
  • 量化。
  • 模型剪枝与蒸馏。
  • 推测解码。
• 借鉴系统领域思想（如推测执行对应推测采样，分页对应PagedAttention）能有效优化。
• 新架构的探索为推理性能的提升带来了巨大潜力。
  • 从头开始方法：定义更快的模型架构并训练。
  • 蒸馏方法：定义更快架构，用原始模型初始化，然后蒸馏修复。
  • 无损捷径：如推测采样，在不牺牲准确性的前提下加速。