课程简介
CS336 Lecture 10：推理（Inference）
学习目标
完成本讲学习后，你应该能够：
- 理解推理（inference）与训练在计算特性上的根本差异，尤其是内存受限（memory-bound）与计算受限（compute-bound）的区别
- 使用算术强度（Arithmetic Intensity）框架分析 Transformer 各组件在预填充（prefill）和生成（generation）阶段的效率
- 解释 KV 缓存（KV cache）的作用、存储结构及其对延迟和吞吐量的影响
- 掌握推理性能指标：TTFT、延迟（latency）和吞吐量（throughput），并理解它们之间的权衡关系
- 列举并理解减少 KV 缓存大小的多种技术：GQA、MLA、CLA、滑动窗口注意力等
- 理解量化、剪枝、投机采样（speculative sampling）等推理加速方法的基本原理
- 理解连续批处理（continuous batching）和分页注意力（PagedAttention）如何应对动态推理工作负载
课程主线
本讲围绕一个核心问题展开：如何让已训练好的模型尽可能快、尽可能准确地生成输出？
课程从推理效率的重要性出发，依次展开以下内容：
- 推理工作负载的数学分析——为什么推理与训练的工作负载性质完全不同
- 延迟与吞吐量的权衡——不同场景下"快"的不同含义
- 有损优化技术——通过减少 KV 缓存、量化、剪枝等降低推理成本
- 无损优化技术——投机采样（speculative sampling）保证采样质量的同时加速推理
- 动态工作负载处理——面向真实服务场景的连续批处理和分页注意力
第一部分：推理工作负载的数学分析
推理的重要性
推理（inference）在多个场景中都有重要应用：
- 实际使用：聊天机器人、代码补全、智能体（agents）、批量数据处理
- 模型评估：特别是需要生成式输出的评估任务
- 强化学习：生成轨迹（trajectories）后进行评分并更新权重
关键观察：训练是一次性成本，而推理是重复性成本。推理 token 的消耗量可以达到惊人的规模。更重要的是，随着智能体时代的到来，大部分生成的 token 并非供人类阅读，而是用于模型自身的思考、工具调用和反思过程，因此推理的算力消耗没有自然的上限。
推理与训练的根本差异
训练（以监督微调为例）：可以一次性看到所有 token，并在序列维度上并行化。在 Transformer 中，注意力、MLP 的计算都将序列作为一个维度处理，因此可以一次性处理所有 token。
推理：由于自回归（autoregressive）特性，必须顺序生成 token，无法在序列维度上并行化。这使得获得高计算强度（compute intensity）或充分利用算力变得更加困难。
符号约定
本讲采用类似 einops 的图示符号来表示张量维度：
- B：批次维度（batch），也代表序列数量
- T / S：序列维度（sequence），也代表 token 数量。S 代表输入 token 数（条件 token），T 代表输出 token 数（生成 token）
- D：模型维度（model dimension）
- F：MLP 投影维度，约定 F = 4D
- N：查询头数量（query heads）
- K：键/值头数量（key/value heads），也是 GQA 中的组数
- H：每个头的维度（head dimension），D = N × H
- L：层数（layers）
- V：词表大小（vocabulary size）
维度符号规则：
- 红色（contracting）：同时出现在两个操作数中，从结果中消失（被收缩）
- 黑色（regular）：只出现在一个操作数中，保留在结果中
- 蓝色（attaching / batching）：同时出现在两个操作数中，保留在结果中
算术强度（Arithmetic Intensity）回顾
算术强度 = 浮点运算次数（FLOPs）/ 传输字节数（bytes transferred）
目标是让算术强度尽可能高。
以矩阵乘法 X(B×D) @ W(D×F) 为例：
内存读写：
- 读取 X（B×D）：2BD 字节（bf16 格式）
- 读取 W（D×F）：2DF 字节
- 写入 Y（B×F）：2BF 字节
FLOPs：2BDF
算术强度 = 2BDF / (2BD + 2DF + 2BF)
假设 B 远小于 D 和 F，强度简化为 B。
加速器强度（accelerator intensity）= 每秒 FLOPs / 内存带宽。
对于 H100：加速器强度约 295。这意味着：
- 算术强度 > 295 → 计算受限（compute-bound），好
- 算术强度 < 295 → 内存受限（memory-bound），差
极端情况：B = 1 时（矩阵向量乘法），算术强度 = 1，属于内存受限。这正是推理中常见的场景。
KV 缓存（KV Cache）
朴素推理：每次生成一个 token，把整个历史序列重新喂给模型。每次前向传播的复杂度是 O(T²)，生成 T 个 token 总复杂度为 O(T³)。
关键观察：由于因果（causal）Transformer 的激活只依赖过去的 token，前缀的计算可以在多次生成之间共享。因此，将计算出的键（keys）和值（values）存储在 HBM 中，称为 KV 缓存。
KV 缓存结构：对于每个序列（B）、每个 token（S）、每一层（L）、每个 KV 头（K），存储一个 H 维向量。每个 token 需要存储 key 和 value 两个向量。
推理的两个阶段：
- 预填充（Prefill）：给定提示词，并行编码所有 token，计算并填充 KV 缓存。这类似训练，可以并行化。
- 生成（Generation）：顺序生成新的响应 token，每个 step 只生成一个 token。
MLP 层的算术强度分析
内存读写（仅关注矩阵乘法）：
- 读取 X（B×T×D）：2BTD
- 读取 Wup、Wgate、Wdown（各 D×F）：3 × 2DF
- 写入 U（B×T×F）：2BTF
- 写入 G（B×T×F）：2BTF
- 写入 Y（B×T×D）：2BTD
FLOPs：6BTDF（上投影、门控、下投影各 2BTDF）
算术强度 = 6BTDF / (4BTD + 4BTF + 6DF)
假设 BT 远小于 D 和 F，强度简化为 BT。
结论：
- 预填充阶段：BT 可以很大（大批次、长序列），容易达到计算受限
- 生成阶段：T = 1，强度降为 B。B 代表并发请求数，对于交互式应用不可预测且可能很低
注意力层的算术强度分析
内存读写：
- 读取 Q（B×T×D）、K（B×S×D）、V（B×S×D）：2BTD + 4BSD
- 写入 Y（B×T×D）：2BTD
FLOPs：4BSTD（QK^T 和 softmax(A)@V 各 2BSTD）
算术强度 = ST/(S+T)
关键观察：强度不依赖 B！
预填充阶段（T = S）：强度 = S/2，还不错。
生成阶段（T = 1）：强度 = S/(S+1) < 1，很差。
为什么批次大小对注意力没有帮助？
- MLP：所有序列共享相同的 MLP 权重。增加 B 只需要加载一次权重，就能处理所有序列，因此算术强度随 B 提高。
- 注意力：每个序列有自己的 KV 缓存（Q、K、V 都依赖 B）。增加 B 相当于给每个序列做独立的矩阵向量乘法，并不能分摊参数加载成本。
总结：
| 组件 | 预填充强度 | 生成强度 |
|---|---|---|
| MLP | B×S（很好） | B（需要并发请求） |
| 注意力 | S/2（可用） | <1（根本瓶颈） |
核心结论：预填充是计算受限的，生成是内存受限的。生成阶段的注意力是 Transformer 推理的根本性瓶颈。
第二部分：延迟与吞吐量
推理性能指标
- 首 token 延迟（TTFT，Time-to-First-Token）：用户从提交查询到收到第一个 token 的等待时间。对交互式应用很重要。
- 延迟（Latency，秒/token）：单个查询中 token 流式输出的速度。对交互式应用重要。
- 吞吐量（Throughput，tokens/秒）：多个查询整体上 token 产生的速度。对批处理任务重要。
关键洞察：推理是内存受限的，因此延迟和吞吐量主要由需要传输的内存量决定。假设计算和通信可以完美重叠，瓶颈就是内存 I/O。
实例：Llama 2 13B 在 H100 上的性能分析
模型配置：S=1024，D=5120，F=13824，N=40，K=40（MHA），H=128，L=40，V=32000
参数数量：约 130 亿参数，与宣称一致
内存占用：
- 参数（bf16）：约 26 GB
- KV 缓存（每个序列）：2 × S × K × H × L × 2 字节
- 总内存 = B × KV缓存每序列大小 + 参数大小，是 B 的线性函数
延迟 = 总内存 / 内存带宽
吞吐量 = B / 延迟
关键结果：
| 批次大小 B | 延迟（秒/token） | 吞吐量（tokens/秒） | 说明 |
|---|---|---|---|
| 1 | 0.008 | 124 | 最低延迟 |
| 64 | 增加 | 增加 | 延迟变差，吞吐量改善 |
| 256 | 进一步增加 | 进一步增加 | 内存溢出（超过 80GB H100 内存） |
增加批次大小的效果：
- 延迟变差：KV 缓存随 B 线性增长，需要读写更大的内存
- 吞吐量改善：参数加载成本在更大批次中分摊
- 收益递减：吞吐量渐近于某个上限，且最终受限于内存容量
延迟与吞吐量的权衡：
- 小批次 → 更好延迟，更差吞吐量
- 大批次 → 更好吞吐量，更差延迟
实际含义：
- TTFT 本质上由预填充时间决定。想要更快的 TTFT，应使用较小批次；想要更高吞吐量，应使用较大批次
- 简单并行化：启动 M 个模型副本，延迟不变，吞吐量增加 M 倍
第三部分：有损优化——减少 KV 缓存
为什么减少 KV 缓存很重要
推理受内存限制，而 KV 缓存占据大量内存（大批次时甚至可能超过参数大小）。减少 KV 缓存大小可直接转化为延迟和吞吐量的提升。关键是在减少内存的同时不损失太多精度。
分组查询注意力（GQA，Grouped-Query Attention）
背景：
- 多头注意力（MHA）：每个 token 有 N 个查询头，每个头有对应的 key 和 value（K = N）
- 多查询注意力（MQA）：K = 1，所有查询头共享一组 KV，效果很差
- 分组查询注意力（GQA）：K 介于 1 和 N 之间，每个 KV 头服务 N/K 个查询头
效果：KV 缓存减少因子 N/K
验证：在 Llama 2 13B 配置中，将 K 从 40 降为 8（1:5 稀疏度）：
- 延迟和吞吐量都得到改善（减少内存使用同时改善两者）
- 之前批次 256 会内存溢出，现在可以放入内存
精度影响：GQA 论文显示在多种评估中表现良好，但应持保留态度——DeepSeek 论文表明 GQA 实际有精度损失。
多潜变量注意力（MLA，Multi-head Latent Attention）
核心思想：不是直接存储 K 和 V，而是将激活压缩到低维潜空间，需要时再投影回来。
实现方式：
- 正常注意力：K = W_K·h，V = W_V·h，维度为 N×H
- MLA：存储压缩向量 c = W_c·h（C 维），需要时计算 K = W_K·c，V = W_V·c
- DeepSeek-V2：将 16384 维压缩到 512 维（+64 维用于 RoPE 兼容，总计 576）
注意：MLA 与 RoPE 不兼容，需要添加额外维度处理。
精度：MLA 比 MHA 略好（且更便宜），而 MHA 优于 GQA。这意味着"越激进压缩不一定越差"，关键在于压缩方式。
跨层注意力（CLA，Cross-Layer Attention）
核心思想：像 GQA 在头之间共享 KV 一样，在层之间共享 KV。只为一小部分层计算 KV，其他层直接使用相邻层的 KV。
效果：经验上改善精度与 KV 缓存大小的帕累托前沿。
局部（滑动窗口）注意力
核心思想：生成 token 时只看最后 K 个 token，而非全部历史。
优势：
- KV 缓存独立于序列长度（只依赖批次大小和其他变量）
- 有效上下文长度可随层数线性增长
问题：损害精度，降低表达能力。
改进方案：将局部注意力与全局注意力交错（混合模型），在部分层使用完整注意力，其余层使用局部注意力。
线性注意力与状态空间模型
不是存储 KV 缓存，而是将整个历史压缩为某种表示（如状态向量）。包括：
- 朴素线性注意力：将 KV 值求和为单一向量
- 门控网络（gated nets）、delta 网络（delta nets）、Mamba 等更复杂的压缩方法
权衡：线性注意力对长上下文的整体总结更高效，但会丢失细节。对于"大海捞针"式检索，滑动窗口注意力可能更好。可组合使用不同类型的注意力以捕捉不同方面。
DeepSeek 稀疏注意力
- 压缩稀疏注意力（CSA）：每 M 个 token 压缩为 1 个
- DeepSeek 稀疏注意力（DSA）：选择保留 token 子集（通过计算轻量级注意力分数确定）
- 重度压缩注意力（HCA）：进一步压缩
第四部分：有损优化——量化与剪枝
量化（Quantization）
核心思想：降低数字精度，减少内存占用，从而提升延迟和吞吐量。
精度选择：
- fp32（4 字节）：训练时参数和优化器状态
- bf16（2 字节）：推理默认
- fp8（1 字节）：H100 支持，可以训练但需谨慎
- int8（1 字节）：便宜但精度较低，仅推理
- int4（0.5 字节）：更便宜，精度更低
量化感知训练（QAT，Quantization-Aware Training）：
- 训练前向传播中模拟量化和反量化，让权重适应量化误差
- 优点：效果更好；缺点：需要昂贵的大规模训练
训练后量化（PTQ，Post-Training Quantization）：
- 训练完成后进行，成本低
- 简单方法：对每层或张量确定缩放因子和零点
- GPTQ：利用 Hessian 信息逐层量化，跟踪误差并传播到未量化权重中
激活感知量化（AWQ，Activation-aware Weight Quantization）：
- 观察：某些激活通道值较大，与其交互的权重更重要
- 方法：基于激活为少数重要通道分配更高精度
- 效果：fp16 到 int3 降低 4 倍内存，3.2 倍加速
模型剪枝（Model Pruning）
核心思想：移除模型中不重要的部分，然后修复它。
算法：
- 在小的校准集（约 1024 样本）上估计不同组件（层、头、隐藏维度）的重要性
- 移除不重要的部分，得到更小的模型
- 将原始模型蒸馏到剪枝后的模型中
效果：可将 15B 模型剪到 8B，精度损失不大，训练成本远低于从头训练。
第五部分：投机采样（Speculative Sampling）
核心洞察
预填充（同时编码所有 token）比生成（一次一个 token）快得多。"检查"比"生成"更快——给定一个序列评估它比逐步生成快得多。
算法
- 使用廉价的草稿模型（draft model）p 猜测 K 个 token
- 使用目标模型 q 并行计算这些 token 的对数几率
- 以概率 min(1, q/p) 接受每个 token
- 拒绝时从残差分布采样，然后退出
关键性质：保证从目标模型中精确采样（无损）！
简要证明（两元素词表例）：
- 假设 p(A) > q(A)，则 p(B) < q(B)
- P[采样 A] = p(A)·(q(A)/p(A)) + p(B)·1·0 = q(A)
- P[采样 B] = p(B)·1 + p(A)·(1 - q(A)/p(A))·1 = q(B)
性能与调优
- 草稿 token 数太少：无法充分利用目标模型的批处理能力
- 草稿 token 数太多：被拒绝概率增加
- 最佳值通常约 3-4 个 token
- 草稿模型应尽量接近目标模型（通过蒸馏）
扩展
- Medusa：草稿模型并行生成多个 token
- EAGLE：草稿模型利用目标模型的高层特征
实用观点：你在减少 KV 缓存方面的所有努力都可以用于构建更好的草稿模型——如果模型好到可以直接服务，就不需要投机采样；否则可以作为投机采样的草稿模型。
第六部分：动态工作负载处理
问题描述
实时网站场景：请求在不同时间到达，具有共享前缀和不同长度，远不同于训练中整齐划一的 token 块。
连续批处理（Continuous Batching）
问题：静态批处理等待慢请求，浪费算力。
方案（Orca 系统引入）：
- 迭代级调度：每一步为所有序列解码一个 token
- 序列完成时从批次移除
- 新请求到达时加入批次
选择性批处理：处理不同序列长度时的策略：
- 注意力计算：因长度不同无法有效共享张量，分别处理
- 非注意力部分（如 MLP，占大量 FLOPs）：拼接所有序列为一个大序列处理
分页注意力（PagedAttention）
问题：KV 缓存的碎片化：
- 内部碎片化：为最大 token 长度（如 1024）分配连续内存，但实际生成更短，浪费空间
- 外部碎片化：不同序列之间的空隙太小，无法有效利用
方案（借鉴操作系统分页）：
- 将序列的 KV 缓存划分为非连续块（blocks）
- 通过索引跟踪块的位置
- 块级对齐保证统一性
共享机制：
- 系统提示词可缓存一次，多个查询共享
- 多个生成（如程序合成中多个候选响应）可共享前缀 KV 缓存
- 写时复制（copy-on-write）语义：采样到相同 token 时继续共享，分叉时拆分块
其他优化：
- 融合块读取和注意力的内核
- 使用最新内核（FlashAttention、FlashDecoding）
- CUDA graphs 避免内核启动开销
课程作业关联
- 算术强度计算：完成 Lecture 2 的算术强度推导后，本讲扩展了 MLP 和注意力层的具体计算
- KV 缓存实现：课程项目中的推理实现需要处理预填充和生成两个阶段
- GQA 实践：在模型配置中调整 K 值，观察对内存、延迟和吞吐量的影响
- 性能分析：使用 TransformerPerformanceStats 类计算不同配置下的性能指标
常见误区
"推理和训练差不多快"——错。推理因自回归特性无法在序列维度并行化，生成阶段尤其受内存限制。
"增加批次总是更好"——不一定。增加批次改善吞吐量但恶化延迟，且受内存容量限制。
"GQA 没有精度损失"——论文中的结论需持保留态度。DeepSeek 论文表明 GQA 实际有精度损失。经验性结果需要具体模型具体验证。
"减少 KV 缓存就是减少模型维度"——不对。关键是找到模型中可压缩的位置，不加区分地减少一切通常更差。
"投机采样会有损精度"——不会。投机采样的数学保证它精确采样自目标模型。
"滑动窗口注意力有效上下文只是窗口大小"——不对。信息可以在层间传播，有效上下文长度随层数线性增长。
"内存受限意味着加速器完全闲置"——确实让人沮丧，但这是推理的本质特性。
复习清单
概念理解
- [ ] 能解释为什么推理是内存受限的，而训练是计算受限的
- [ ] 能说明预填充和生成阶段的区别及其对并行化的影响
- [ ] 能推导 MLP 层和注意力层的算术强度表达式
- [ ] 能解释为什么注意力算术强度不依赖批次大小
- [ ] 能说明 TTFT、延迟、吞吐量的定义和适用场景
KV 缓存
- [ ] 能解释 KV 缓存的结构和存储内容
- [ ] 能计算给定配置下的 KV 缓存大小
- [ ] 能列出减少 KV 缓存的各种方法（GQA、MLA、CLA、滑动窗口注意力）
优化技术
- [ ] 能解释 GQA 的组数与 KV 缓存缩减的关系
- [ ] 能理解 MLA 的压缩机制及其 RoPE 处理
- [ ] 能区分量化感知训练和训练后量化
- [ ] 能理解 AWQ 的核心观察和做法
- [ ] 能描述模型剪枝的基本流程
投机采样
- [ ] 能解释投机采样的算法步骤
- [ ] 能说明为什么它是无损的
- [ ] 能理解草稿 token 数量的权衡
系统层面
- [ ] 能理解连续批处理的工作原理
- [ ] 能解释分页注意力解决的碎片化问题
- [ ] 能说明 KV 缓存共享和写时复制的应用场景
核心公式
- [ ] 算术强度 = FLOPs / 字节数
- [ ] MLP 生成阶段强度 ≈ B
- [ ] 注意力生成阶段强度 ≈ S/(S+1) < 1
- [ ] GQA KV 缓存缩减因子 = N/K
课程总结
本讲的中心思想可以概括为：
推理非常重要，与训练在计算特性上完全不同。推理受内存限制，生成阶段尤其如此。减少 KV 缓存是提升推理效率的关键方向，可以通过架构设计（GQA、MLA、CLA）、量化、剪枝、投机采样等多种方式实现。同时，面对动态的实时工作负载，需要借鉴操作系统思想（分页、连续批处理）来管理推理资源。
未来方向：新架构（状态空间模型、线性注意力、扩散模型）可能在推理效率上有更大潜力，因为注意力机制和 KV 缓存的构建方式从根本上使其对推理不友好。如果能够设计出为推理而生的新架构，或许能解锁更多可能。
推理效率：从内存瓶颈到投机采样与动态批处理
引言
训练一个大型语言模型是一次性的巨额投资，但真正决定一个模型能否被广泛使用的，往往是它在推理阶段的效率。在斯坦福大学 CS336《Language Modeling from Scratch》课程的第十讲中，我们深入探讨了推理这一关键环节——它不仅是模型部署后的日常运营成本，更是在智能体（agent）时代成为决定算力消耗的核心因素。
本讲的核心观点是：推理与训练在计算特性上存在根本性差异。训练时我们可以同时看到所有 token 并在序列维度上进行并行化，而推理必须顺序生成 token，这使得推理变得高度依赖内存带宽（memory-bound）。理解这一点后，我们才能系统地分析各种优化技术背后的原理。
本文将依据课堂内容，从算术强度（Arithmetic Intensity）的数学框架出发，逐步解析 KV 缓存、量化、剪枝、投机采样（Speculative Sampling）以及连续批处理（Continuous Batching）等技术，并说明它们在课程整体学习路径中的位置。
从计算强度看推理的本质
为什么推理如此重要？
在课程开始，讲师给出了一个令人印象深刻的对比：OpenAI 每天生成约 8.6 万亿个 token，而 DeepSeek-V4 的训练仅用了 32 万亿个 token。这意味着，不到四天时间，OpenAI 的推理输出量就等同于一次前沿模型的全规模训练量。
更关键的是，随着智能体（agent）时代的到来，推理的 token 消耗正在发生质变。在传统聊天机器人场景中，生成的内容最终要由人类阅读，而人的阅读速度是有限的；但在智能体场景中，模型需要自主思考、调用工具、进行反思，大部分生成的 token 并非直接呈现给用户，而是用于内部的推理过程。这意味着 token 生成量可以无限增长，推理效率的提升永远具有实际价值。
算术强度：衡量效率的标尺
要理解推理为什么是内存受限的，我们需要借用算术强度（Arithmetic Intensity）这一概念——它表示每传输一个字节的数据能完成多少次浮点运算。理想情况下，我们希望这个值越高越好。
课程的推导从最简单的矩阵乘法开始。假设我们要计算 X 与 W 的乘积，其中 X 的维度是 B×D，W 的维度是 D×F。需要从 HBM（高带宽内存）中读取这两个矩阵，计算浮点运算次数，然后将结果写回。假设使用 bf16 精度：
- 读取 X：2×B×D bytes
- 读取 W：2×D×F bytes
- 计算 FLOPs：2×B×D×F
- 写入结果：2×B×F bytes
计算算术强度并进行简化（假设 B 远小于 D 和 F）后，我们得到算术强度等于 B。对于一个 GPU 加速器，其"加速器强度"等于每秒浮点运算次数除以内存带宽。课程中以 H100 为例，这个值大约是 295。当我们的算法算术强度大于 295 时，它是计算受限（compute-bound）的；反之则是内存受限（memory-bound）的。
极端情况：当 B=1 时，算术强度降为 1。每次只处理一个序列，读取整个参数矩阵却只做很少的计算——这正是推理中最常见的情形。
预填充与生成：两个截然不同的阶段
推导完基础框架后，我们进入 Transformer 的层间分析。课程对 MLP 层和注意力层分别计算了算术强度，关键结论如下：
MLP 层：其计算本质是大型矩阵乘法，算术强度为 B×T（B 为批大小，T 为序列长度）。在预填充阶段（prefill），B×T 通常足够大，容易达到计算受限状态。在生成阶段（generation），T=1，算术强度退化为 B——这依赖于并发请求数量。
注意力层：情况更为复杂。注意力层的算术强度公式为 S×T/(S+T)，其中 S 是已生成的先前 token 数量，T 是本次生成的 token 数量。在预填充阶段（T=S），强度为 S/2，这是可以接受的。但在生成阶段（T=1），强度变为 S/(S+1)，小于 1。这个值远低于任何现代加速器的强度阈值，意味着注意力层在生成阶段永远是内存瓶颈。
为什么注意力层的算术强度不随批大小 B 扩展？这涉及到权重共享的基本差异。MLP 层的权重对所有序列是共享的，只需加载一次即可为整个批次服务；而注意力层中每个序列都有自己独立的 KV 缓存，增加 B 并不会摊销参数加载成本——每个序列都在做独立的小规模矩阵-向量乘法。
结论：预填充阶段是计算受限的，生成阶段是内存受限的。而生成阶段的注意力计算是真正的"根本性瓶颈"。
延迟与吞吐量的权衡
在上述数学框架基础上，课程引入了一个具体的例子：在一台 H100 上运行 Llama 2 13B 模型。这个模型具有一定的特定架构配置，我们暂不需要记住所有参数细节，只需关注性能计算的逻辑。
内存占用的构成：首先是参数内存——13B 参数以 bf16 格式存储，占用约 26GB。其次是 KV 缓存——它随批大小 B 线性增长，包含每个序列的键值对。总内存使用量是两者之和。
延迟的决定因素：由于推理是内存受限的，延迟主要由内存 I/O 决定。计算公式为（参数大小 + KV 缓存大小）除以内存带宽。这意味着延迟随 B 线性增长——更大的批处理意味着需要读写更大的 KV 缓存。
吞吐量的非线性特性：吞吐量约为 B 除以延迟，但存在渐近线。增加 B 确实能提升吞吐量，因为参数加载成本被更多人分摊，但 KV 缓存随之增大，边际收益递减。更重要的是，最终会遇到硬件内存上限——内存不足（OOM），无法继续扩大批次。
课程提供了一个形象的比喻：批处理就像坐公交车。单个人（批大小=1）随叫随走，延迟很低；但公交车（大批次）一次能运送很多人，吞吐量更高。你不可能同时获得低延迟和高吞吐量，必须在两者之间做出权衡。
对 TTFT 的影响：首 token 延迟（TTFT）本质上由预填充阶段决定，较小的批大小能改善首 token 延迟，而较大的批大小则有利于提升整体吞吐量。
减少 KV 缓存：多种无损与有损策略
既然确立了"内存是推理瓶颈，KV 缓存占用大量内存"这一前提，那么减少 KV 缓存顺理成章地成为核心优化方向。课程的策略介绍由浅入深，从架构到系统层面逐步展开。
分组查询注意力（GQA）
这是一个我们已经熟悉的概念。在多头注意力（MHA）中，每个查询头都有对应的键和值头；而分组查询注意力（GQA）保持相同数量的查询头，但只使用更少的键值头组，每个组服务多个查询头。通过减少键值头的数量（从 N 减少到 K），KV 缓存大小缩减为原来的 K/N。
课程特别强调了这一技术的历史背景：2023 年 GQA 论文证明了在保持良好性能的同时，时间开销大幅下降。值得注意的细节是，当我们在 GQA 配置上进一步增大批大小时，之前因为在批量 256 时内存溢出的情况现在变得可行了——这说明减少 KV 缓存不仅直接降低延迟，还能间接通过允许更大批次来提升吞吐量。
多潜变量注意力（MLA）
这是 DeepSeek 提出的技术，其思路比 GQA 更进一步。MLA 不再简单地减少键值头的数量，而是将它们压缩成低维向量存储。具体而言，不再直接计算完整的 K 和 V，而是先将激活投影到一个压缩维度 C（例如 512 维）中，再从压缩表示中恢复出实际的键和值。这样只需要存储较小的压缩向量，而不是完整的键值向量。课程提醒我们，MLA 与 RoPE 位置编码不完全兼容，因此需要额外添加维度来处理这个问题。这个技术说明，架构层面的巧妙设计可以大幅降低 KV 缓存大小，同时保持模型表现。
跨层注意力（CLA）
另一个共享思路是跨层共享 KV。GQA 在头之间共享，而 CLA 在层之间共享——只对部分层计算 KV 缓存，后续层直接复用前一层的 KV。这同样改善了准确率与 KV 缓存大小的帕累托前沿。
滑动窗口注意力与混合架构
对于长上下文场景，一个自然的想法是只注视最近的 K 个 token（滑动窗口注意力）。这样 KV 缓存就不再依赖序列长度，但代价是降低了表达能力。实践中人们倾向于将局部注意力与全局注意力交错，在部分层使用完整注意力，在其余层使用局部注意力。课程也提及线性注意力和状态空间模型（如 Mamba）作为另一种压缩历史的方式，但它们无法保证与滑动窗口注意力的对比结论，因为存在"大海捞针"式的任务——过度压缩历史必然损失信息。
量化与剪枝：系统层面的挤压
量化是系统优化中的常见手段，核心思想是降低数值精度。从 bf16 到 int4，每一步都减少内存占用。量化感知训练（QAT）在训练期间就模拟量化误差，让权重适应量化，但需要高昂的大量训练成本。更常见的是训练后量化（PTQ），它依赖校准集来确定缩放因子和零点，其中 GPTQ 使用 Hessian 信息修正误差，而 AWQ（激活感知量化）则观察到某些激活通道值特别大，与其交互的权重更为关键，因此为这些少数通道保留更高精度。
模型剪枝则走另一条路——直接移除不重要的部分。NVIDIA 的论文展示了如何用校准集估计不同部分的重要性（如隐藏单元激活值的大小），剪除不重要的层或单元后，用蒸馏来修复模型。这样能将 15B 的模型压缩到 8B 且保持良好效果，而训练成本远低于从头训练。
投机采样：利用"检查"与"生成"的不对称性
上述方法大多有损，而投机采样（Speculative Sampling）是课程介绍的一种无损加速方法。
方法的核心洞察是：检查一个序列是否好，远比逐 token 生成一个序列要快。预填充阶段可以并行处理大量 token，而生成阶段必须串行。因此，可以引入一个廉价的小型草稿模型（draft model）快速生成几个候选 token，然后用目标大模型并行地对这些候选进行验证。验证通过则接受，否则在残差分布上继续采样。
数学上，这本质上是修改版的拒绝采样。课程给出了一个简要的证明思路：只要接受概率设定为 min(1, q/p)，即目标模型概率与草稿模型概率的比值，那么即使草稿模型的分布与目标模型不完全一致，最终的采样结果也能保证是从目标分布中精确采样。这意味着输出质量与逐 token 生成完全一致，但速度显著提升。
实际操作中，草稿 token 数量需要选择恰当——太少则无法充分利用目标模型的批处理能力，太多则被拒绝的概率升高。课程中演示的最佳数值约为 3 或 4 个 token。此外，如果我们将前面讨论的各种技术融合到草稿模型中，效果会更好——剪枝、量化和蒸馏都可以用于构建更贴近目标模型的草稿模型。
动态工作负载的处理
真实的推理服务远非理想化的静态批处理。请求在不同时间到达，具有共享前缀，长度也不同。这种"参差不齐"的工作负载给批处理带来了极大挑战。
连续批处理（Continuous Batching）由 Orca 系统引入，其思路是迭代级调度：每一步为所有活动序列解码一个 token；序列完成即移除，新请求到达即加入，批次始终动态更新。即使请求长度不同，也可以通过"选择性批处理"来应对——对于 MLP 这类不依赖序列长度的操作，可以直接将不同长度的序列拼接成一个巨型批次处理。
分页注意力（PagedAttention）则从操作系统借鉴了内存管理的思想。KV 缓存不再需要连续的内存区域，而是划分为大小相同的块（blocks），通过索引管理。这解决了两个碎片化问题：内部碎片化（分配了超出实际需要的缓冲区）和外部碎片化（序列之间的空隙无法有效利用）。更重要的是，分页使得共享成为可能——多个请求可以共享相同的前缀（如系统提示词）的 KV 缓存，甚至可以在一个提示词上生成多个不同响应时实现写时复制语义。
本讲在课程中的位置
这门课从语言建模的基础开始（第一讲到第五讲），涉及缩放定律（第九讲），再进入推理（第十讲），最后回到缩放定律的深入讨论。推理看似独立，实际上是把先前所有内容串联起来的关键环节：
- 它与训练并行计算的基础设施（如批量矩阵乘法）形成鲜明对比，正是因为推理无法像训练那样在序列维度上并行化，才需要专门的设计。
- 它在缩放定律的背景下理解：推理的效率直接影响训练模型的规模选择和经济性。
- 它为后续的强化学习（RL）打下基础——RL 需要大量生成轨迹进行评估，推理效率决定 RL 的上限。
推理也是连接理论与实际应用的桥梁，将 Transformer 架构、注意力机制、内存带宽等底层知识与具体的系统优化技术（如 vLLM、SGLang 等）联系起来。
结语
本讲的核心思想可以总结为一句话：减少 KV 缓存，但不要让准确性损失太多。无论是架构创新（GQA、MLA、CLA）、数值精度调整（量化），还是模型压缩（剪枝、蒸馏），所有技术都遵循这个原则。
展望未来，讲师提到一个值得深思的观点：注意力机制和 KV 缓存的架构设计从根本上决定了它对推理并不友好，而新架构（如状态空间模型、线性注意力或扩散模型）可能蕴含更大的突破潜力。Transformer 当初也不是为推理效率而设计，而是为训练效率设计的——设计一个专门为推理而生的新架构，也许正是下一代语言模型的方向。
同时，系统的视角也至关重要。分页、连续批处理、投机执行这些来自操作系统和体系结构的概念，在推理场景中焕发出新的生命力。推理效率的提升不只靠单一技术，而是架构设计、训练策略和系统工程的综合优化。
官方视频与讲义入口：
- 课程主页：CS336 Language Modeling from Scratch
- 本讲视频：YouTube
- 大部分图表和推导灵感参考了 JAX-ML Scaling Book 的 Transformers 与 Inference 章节
CS336 2026 Lecture 10：完整中文逐字稿
- 视频：官方视频
- 来源：英文人工字幕提取后翻译为中文；保留段落起始时间戳，只修正明显的断句、重复和语序。
- 说明：这是逐字稿，不是讲义摘要；相邻且语义连续的片段已合并为阅读段落；术语保留必要英文。
正文
[00:05.320] 好的，我们开始吧。上次，Tatsu 讲了缩放定律，我们今天先暂时放下这个话题，来谈谈推理。推理这个问题本身非常简单。你已经训练好了一个模型，然后收到一个提示词，你希望尽可能准确、尽可能快地生成回复。推理虽然只占一节课，但它的重要性其实在日益增长。它出现在很多地方。很明显，一旦你训练好了模型，你不可能就把它晾在那儿。如果你是研究者，你会在论文里放一个图来展示你的模型效果如何，但对其他所有人来说，你实际上是想要使用这个模型的。那么“使用”是什么样的呢？可能是与 AI 助手或聊天机器人对话，也可能是在做代码补全。如今，智能体非常流行，这需要推理。批处理数据也同样需要。它还被用于评估，特别是那些需要生成的评估任务。
[01:19.670] 它也被用在训练过程中。如果你在做强化学习，你需要有一个模型，让它生成轨迹，然后进行评分，再相应地更新权重。所以推理无处不在。而效率，正如这门课的主题一样，在这里比以往任何时候都更重要。即使你只从实际应用的角度来看，训练是一次性成本。它可能非常昂贵，也确实很昂贵。但一旦你完成了训练，事情就结束了。但推理是重复性成本。你基本上每天都要承担这个成本。据估计，OpenAI 每天要生成 8.6 万亿个 token。作为参考，今年早些时候发布的 DeepSeek-V4 是在 32 万亿个 token 上训练的。所以，不到四天的时间里，OpenAI 需要生成的 token 数量，以及相应的算力，就至少相当于训练 DeepSeek-V4 的规模。
[02:26.540] 当然，即使是前沿模型可能在更多 token 上训练，但推理所生成的 token 数量仍然会更高。所以这只是一个视角。此外，我认为推理的重要性在过去一年里有所增长，因为曾经我们主要把语言模型看作是聊天机器人或助手，你输入一个提示词，然后得到一个回复，假设人类的目的是去阅读这个回复。但现在，随着我们越来越多地进入一个智能体时代，实际情况看起来更像是，一个查询进来，智能体会去做一系列的事情。它会思考，会推理，会调用一些工具，会进行反思。最后，它会为人类生成一些可读的输出。所以，智能体生成的大部分 token 实际上并不是为了让人阅读的。因此，你真的应该把生成的 token 数量看作是算力消耗。而且这个消耗是没有上限的。
[03:34.110] 如果你的问题足够宏大，你就需要大量的算力和海量的 token。如果说在聊天机器人时代，推理在某个程度后就已经足够快了，因为人阅读速度有限。但对于智能体来说，通过进一步压榨推理来获取价值是没有止境的。所以有很多人在商业方面做推理。当然，所有闭源 API 提供商都必须服务他们的模型，所以推理对他们来说是大事。也有很多提供商在服务开源权重模型，并提供推理服务。在开源社区，有很多软件包。vLLM 可能是一个比较流行的，算是首选。SGLang 是另一个，特别适合智能体工作负载，但可能还没那么流行。还有 NVIDIA 的 TensorRT，它真的非常快，但适用范围更窄。
[04:42.190] 然后是 llama.cpp，如果你想在 CPU 上运行推理，这是一个流行的包。所以，希望我已经论证了推理非常重要。如果你能让推理速度提高一倍，哪怕只提高 10%，那都是件大事。那么现在的问题是，快到底意味着什么？有几个指标可以衡量“快”这个概念。它们在不同的场景下适用，并且有不同的权衡取舍。一个衡量“快”的指标是首 token 延迟（TTFT）。这基本上就是用户在生成开始前等待的时间。你把一个查询输入到 ChatGPT，过了若干毫秒，然后从第一个 token 出现开始计算，这就是这个时间。这对交互式应用通常很有用，因为那段你只是等待、什么也做不了的延迟，越久用户体验就越差。而一旦 token 开始出现，也许速度
[05:54.210] 就不必那么快了，因为你本来也要阅读，你不可能读得那么快。第二个指标是延迟。这是从单个用户的角度来看的。对一个查询来说，token 出现得有多快？这对交互式应用也很重要。基本上，这就是 token 流式输出的速度。然后第二个相关概念是吞吐量，衡量每秒的 token 数，这是对于许多查询来说 token 出现的速度。所以延迟，或者说单个延迟和吞吐量显然密切相关。它们通常是一致的——总的来说，许多干预措施会同时改善延迟和吞吐量。但正如我们稍后将看到的，这里实际上存在一个权衡。而吞吐量对于批处理任务很有用。比如你有一 PB 的数据，想用语言模型处理。
[06:57.000] 你只希望这个任务能完成。至于单个查询是更快、更早还是更晚地完成，并不重要。好的。那么，是什么决定了这些推理指标下的效率呢？总体的要点如下。这可能是本课需要记住的一个高层次的要点，那就是在训练时，你会同时看到所有 token，因为在监督微调中，你会看到所有的 token，并且可以在序列维度上进行并行化。回想一下，在 Transformer 中，注意力、MLP 的计算，序列只是一个维度。所以它是一个被相乘的大张量。因此，你基本上是一次性处理所有 token。在推理中，你做不到这一点。由于推理的自回归特性，你必须顺序地生成 token，一次一个。所以这就是为什么推理是一个非常不同于训练的工作负载的根本问题，因为你无法在序列维度上并行化。
[08:07.030] 因此，正如我们稍后将看到的，要获得高计算强度或充分利用算力会更加困难。好的。那我们继续。首先，我——只是预览一下这堂课的内容。首先，我会做一些数学推导，来理解 Transformer 的计算强度、吞吐量和延迟，以及如何思考这些问题。然后我会讨论各种通过减小 KV 缓存、量化、模型剪枝来降低成本的技术，最后，我会谈谈投机采样，以及一些实际的考虑。好的。这部分算是复习，主要是交换一下，确保我们在符号上达成一致。这堂课很多内容基于 Google 关于 Transformer 和推理的缩放定律一书。所以我建议大家去看看。它写得非常好。有些图表也是大量借鉴自那本书。好了。为了理解这个图，我们先把符号定义一下。
[09:30.380] 所以我们要用符号来表示。用这种图示方式来指代维度，但它们同时也代表长度。所以 B 是批次维度，也代表序列的数量。T 是序列维度，也代表词元的数量。D 是模型维度，H 是维度。所以如果我要写这个，就是将一个维度为 B、T、D 的张量与另一个维度为 D 和 H 的矩阵相乘，然后得到 BTH。这里发生的是，有一些维度被标记为红色，这些维度将会被收缩（contracting）。这些维度同时出现在两个操作数中，然后从结果中消失。然后还有其他的常规维度用黑色表示，它们只出现在一个操作数中，但保留在结果里。所以这是一般形式。另一种情况是，如果你有蓝色的维度，那意味着它们是附加维度（attaching dimensions），它们同时出现在两个操作数中，但保留在结果里。
[10:36.630] 它们不会被收缩或约减。好的。所以这里的解释是，我们稍后会看到，对于每个序列和每个词元位置，我们有一个张量。你有一个向量，你想把它乘以一个矩阵。然后这里，你基本上是在两个张量之间做一系列的点积。好的。那么有了这个基础，这就是一个 Transformer 块的全貌描述。它看起来有点像电路图。这里面有很多细节，但我实际上觉得这可能是关于什么是 Transformer 最清晰的定义。你会看到很多关于 Transformer 的描述，坦白说，其中很多都很难准确理解到底发生了什么。这个定义基本上精确地告诉了你张量的形状，并让你能够推理依赖关系。所以对于 Transformer 块，你取 X，也就是某一层的激活值，先通过注意力（attention）机制，然后再通过 MLP。
[11:44.020] 简单回顾一下，你使用查询矩阵、键矩阵和值矩阵进行投影。你在这里会看到我们有 N，这是头的数量，而 H 是每个头的维度。所以查询矩阵的维度是批次乘以序列长度乘以头的数量再乘以每个头的维度。而 K 和 V，这里不是 N，而是 K，我们稍后会看到，我们可能有一个更小的键值头数量。然后这是注意力操作，注意这里有批次维度。所以 B 同时出现在两者中。我们是在头维度上进行收缩。我们稍后会回到为什么这个 B 同时出现在两者中会让推理变得困难。为什么注意力在那里会成为瓶颈。所以先记住这个想法。而 MLP 非常简单直接。你只需要做一些数学运算。这是门控矩阵。这是上投影和下投影。所以谢谢大家。我已经实现了这个。
[13:07.470] 所以我不想过多赘述，只是换一下符号。按照惯例，我们假设在 MLP 中，MLP 将 D 维的模型维度上投影到其四倍大小。所以 F 始终是——每当你看到 F 时，就把它当作 4D。然后模型维度会被分配到 N 个头中。所以模型维度总是等于头的数量乘以每个头的维度。然后头的数量会被分割，在分组查询注意力（Group Query Attention）的情况下，分割为组的数量和每组中的头的数量。最后，我们有两个变量，S 和 T，它们都代表序列维度。区别在于 S 代表输入词元的数量，而 T 代表输出词元的数量。所以在训练时，这两个是相同的，因为我们预测的输出和输入是相同的。而在推理时，T 将是 1，S 就是输入。
[14:22.960] 好的。再回顾一下算术强度（Arithmetic Intensity）。我在第二讲中讲过这个。所以简单热个身。假设我要乘两个矩阵；一个 B 乘 D 的矩阵和一个 D 乘 F 的矩阵。直观上，B 是批次维度，D 是隐藏维度或模型维度，F 是 MLP 中的投影维度。记住，算术强度是多少？我必须计算浮点运算次数（flops），也必须计算移动的内存量。我们可以这样做。那么，要乘这些矩阵，你需要做什么？记住，在系统讲座中，你必须从 HBM 中读取 X。如果你以 bf16 存储所有内容，推理时就是这样，我们会有 2 乘以 B 乘以 D 的读取量。你还要读取 W，那又是 2 乘以 D 乘以 F。然后你做矩阵乘法。所以浮点运算次数是 2 乘以 B 乘以 D 乘以 F。然后你要把结果写回去。
[15:33.180] 明白了吗？所以你每次读取和写入，基本上就是元素的数量。这有点像二次项，而矩阵乘法是三次的。记住，这一点很重要，因为这是你获得高算术强度的方式。所以总的浮点运算次数是三次方的，而传输的字节数是读取和写入的总和。记住，算术强度就是我们每传输一个字节能做多少计算。我们希望这个值很高。所以强度等于浮点运算次数除以传输字节数。我用符号来表示所有这些，因为这样比用数字更容易看清和处理。我们可以做一个简化，假设这里的批次维度远小于 D 和 F，那么我们可以简化这个式子。具体来说，这意味着我们设 D 等于 C 乘以 D，F 等于 C 乘以 B，并让 C 趋向无穷大。这样 D 和 F 都非常大，而 B 较小，这个式子就简化为 B 的强度。
[16:51.590] 其要点是，记住我们之前计算矩阵乘法的算术强度时，大约是 N/3。这与之类似，只不过不是方阵，而是非方阵。好的。为了比较，一个硬件的加速器强度（accelerator intensity），你看规格表里的每秒浮点运算次数，再看内存带宽，也就是HBM 之间数据移动的速度，两者相除，那就是加速器强度。现在，你把你的计算强度与加速器强度进行比较。如果前者大于后者，那么你是计算密集型（compute bound），这很好；如果你小于它，那么你是内存密集型（memory bound），这很糟糕。所以在这种情况下，对于 H100，对于这个特定的矩阵乘法操作，如果批次大小大于 295，你就是计算密集型的。这里有个极端情况。假设你只有一个样本。会发生什么？如果你只有一个样本，那么你的算术强度就是 1。
[18:16.320] 这将是内存密集型的。你可以看到这里发生了什么。你在读取这个 D 乘 F 的矩阵，但 B 只是 1。所以你实际上只做了 2 乘以 D 乘以 F 的浮点运算。而这基本上就是你在推理中会遇到的工作负载。你不会得到这些完整的矩阵，你得到的是这些非常薄的矩阵或张量。好的，我先停一下，看看有没有问题。是的。在 Transformer 那张幻灯片上，我想确认一下[听不清]。你说，跨过[听不清]可以搜索[听不清]？是的，那应该是跨过头（heads）而不是立方（cubes）吗？让我看看。所以这里的组数是——等等，我确认一下。
[19:31.790] 所以 K 是键值头的数量。所以每个 KV 头有 g 个。我的理解是[听不清]每个键路径会有一个[听不清]与之对应。然后[听不清]那么 g 就是[听不清] p 的数量。那意味着我们有 k。这就是我的问题。是的，好的。是的，我认为你说得对。所以k应该是组数，而g应该是每组中的头数。对，是的，好的，谢谢指正，我稍后会改过来。好的，那我们继续多聊一些关于推理的算术强度。这里是如何理解推理过程中发生的事情。假设你只是朴素地去做。这里是一个提示词，永远不会给你……“Never gonna give you up”。我把它输入给Transformer。Transformer生成键、值和激活。最终，它在输出词表上生成logits。
[21:02.720] 然后你对一个词元进行采样，这个词元被拼接回提示词，然后你就再重复这个过程。你生成“never”，把它附加到提示词上，如此类推。所以这是最朴素的方法。如果你有一个黑盒，输入一个序列并输出词元上的分布——这正是Transformer所做的——你可以直接反复应用它。好的，这样是可行的，但这样真的很糟糕，因为每次你生成一个词元，实际上需要的时间是T的平方量级，其中T是你到目前为止生成的词元数量。所以生成T个词元实际上是T的三次方，因为注意力本身已经是T的平方，而且你必须对每个生成的词元都做一次这样的计算。
[22:00.910] 所以这相当糟糕，但关键是，你实际上不必这样做。很多计算可以在前缀之间共享。例如，如果你正在生成“never gonna give you up”而不是生成“up”，实际上你对“never gonna give you”计算了许多相同的键和值。所以这些词元是不应该改变的，明白吗？这是因为这是一个因果Transformer。如果它是双向的，那么如果你附加一个词元，一切都会改变。但如果它是因果的，这些激活不会因为你追加的任何词元而改变。基于这个观察，第一个显而易见的事情是在HBM中存储一个KV缓存，这样在连续生成词元之间，你可以直接重用KV缓存。这意味着当你试图生成“gonna”时，你实际上不必为之前的这些词元计算所有的键和激活。
[23:10.460] 这就是使用KV缓存的样子。会有两个阶段：预填充。你得到提示词，你填充KV缓存，也就是Transformer计算出的键值对的集合，然后你生成logits——和我们之前做的计算一样。现在你有了这个KV缓存，然后还有这个分布，你从中采样一个词元。现在你把它输入Transformer，它会使用这个缓存，然后产生下一个词元的分布，以及一个新的KV对，这基本上是对应于“up”这个词元的激活。然后这被输入进去。现在你有了累积的KV缓存。你有这些logits，你进行采样，然后输入到Transformer。你输出下一个词元的分布以及新词元，这个新词元被添加到KV缓存中，如此类推。
[24:19.960] 所以KV缓存的形式是：对于每个序列——有B个；对于每个词元，有S个；并且对于每一层和每一个头，你存储一个H维向量。并且再次重申，对于推理，有预填充阶段。你给定一个提示词，你把提示词编码成这些向量。这就像在训练中一样可以并行化，因为你看到了完整的提示词，你可以计算KV缓存。然后在生成阶段，你顺序地生成新的响应词元。但至少你不需要为你已经看过的词元支付生成KV缓存的计算成本。那大家都对KV缓存清楚了吗？有什么问题吗？好的，那让我们尝试计算一下MLP和注意力层的FLOPs和内存I/O。记住，S是我们正在“条件化”的词元数量，T是我们正在生成的词元数量。
[25:39.260] 我们要抽象地进行这个计算。然后之后，我们会具体到预填充阶段，那时T等于S，以及生成阶段，那时T等于1。对于MLP层——在所有这些计算中，我只关注矩阵乘法，因为矩阵乘法才是真正需要大量计算的东西。其他所有东西的FLOPs不多，并且可能可以融合到矩阵乘法中。好的，让我们来做这个计算。这和矩阵计算是一样的。我不打算在这里一步步详细展开每个细节，但形式就是：你从高带宽内存中读取X。你读取所有的参数。你计算上投影，然后把它写入HBM。你计算门控，并写入HBM，然后你基本上计算其下投影，并写入HBM。所以FLOPs的数量取决于B、T、D和F。也就是批量大小、序列长度、模型维度和前馈维度，即MLP维度。
[27:08.450] 传输的字节数就是这个表达式。现在你可以计算算术强度了，重申一下，就是FLOPs除以传输的字节数。这是一个表达式。我们假设，和之前一样，B乘以T远小于D和F。这样，我们就看到强度是B乘以T。这类似于矩阵乘法的情况。这里有更多的矩阵，更多的维度，但本质上是一样的。这是有道理的。MLP基本上就是一个大的矩阵乘法。这也说得通，因为批量维度和序列长度维度——一切都是独立的。对于MLP来说，它们不相交互。注意力则是另一回事。所以是B乘以T。所以只要在预填充阶段，你把B乘以T做得足够大——大批量、长序列——就没问题。现在，让我们看看生成阶段。记住，生成阶段有两个问题。
[28:24.430] 一个问题是生成阶段T等于1。你一次只生成一个词元。这意味着你的算术强度将是B，但在生成阶段，B是并发请求的数量。如果你在批处理环境中，你可以控制这一点。但如果你，比如说，在服务一个聊天机器人，并发请求的数量基本上就是有多少个并发用户。这可能很高，也可能很低。所以这有点不可预测，并且会随时间变化。这就是我们讨论连续批处理时要解决的问题。但总体而言，这并不算糟糕，因为只要你有大批量，序列长度不会真正帮到你。但如果你有大批量，你应该就没事了。现在让我们看看注意力层。再次强调，S是已经生成的先前词元。T是你想为其生成logits的词元数量。在注意力中，你从HBM中读取QKV矩阵。你计算注意力。你计算softmax，这实际上不太重要，然后计算值矩阵，然后把结果写入HBM。
[29:54.460] 所以FLOPs是B乘以S乘以T乘以D，传输的字节数就是这个表达式。所有这些都是矩阵乘法，对吧？所以这应该总是比传输的字节数高一个多项式阶，因为它是矩阵乘法。唯一的问题是，那个因子看起来是什么样的？注意力的那个因子是S乘以T除以S加T。让我们看看预填充。预填充阶段，这意味着预填充的强度是当T等于S时的S除以2。这很好，对吧？因为只要你在注意力上有足够长的序列长度，你就能保持高算术强度。注意这里批处理维度并没有参与计算。我稍后会解释一下原因。但对于生成来说，这是个坏消息。所以生成的计算强度是 S 除以 S 加 1，这个值小于 1，或者我们干脆就说它是 1。而算术强度，记住 1 是很差的。
[31:19.340] 我们希望 H100 的算术强度能达到大约 295 才能饱和计算能力。所以这确实是瓶颈。我们为 MLP 注意力、预填充和生成做了所有这些分析，然后发现这就是瓶颈。那么让我们试着思考一下为什么这是瓶颈。与 MLP 不同——MLP 在生成阶段其实是可以的，只要你的序列——抱歉，只要你的批大小足够大。所以 ML 注意力的问题在于——让我们看看 MLP。每个序列都使用相同的 MLP 权重。所以这些不依赖于 B。而在注意力层中，每个序列都有自己的 KV 缓存。所以这些都依赖于 MLP。你可以这样理解这里发生的事情：在 MLP 的情况下，B 变大实际上是有帮助的，因为你只需要加载一次这些 MLP 权重，就可以把它用于所有的序列。这就是为什么你能获得高算术强度，因为——简化一点说，你只加载一次。
[32:58.510] 然后你完成所有的批处理，这样就很好。而相比之下，在注意力机制中，所有这些都依赖于 B。所以增加 B 并没有帮助。就好像每个序列你基本上都在做矩阵向量乘法。它们都是独立的。所以做更多的矩阵向量乘法并没有帮助。如果你还记得在最开始的时候，我给你们展示了这个例子。这个例子也有相当差的算术强度。它不是矩阵向量乘法，因为我们本质上是在按坐标进行批处理。这本质上和做点积是一样的，而点积的算术强度是非常糟糕的。记住，这是在注意力机制中。这个蓝色的 B 就是为什么注意力的算术强度不随 B 扩展的原因，也是为什么这是一个瓶颈。所以这里总结一下。预填充是计算密集型的。生成是内存密集型的。所以如果你看预填充的 MLP 计算强度，它是 B 乘以 S，非常好。
[34:12.480] 预填充的注意力计算强度是 S 除以 2——不是那么好，但还可以用。生成的 MLP 计算强度也是可行的，但需要较长的并发请求。但真正算得上是根本性瓶颈的是生成的注意力计算强度，就是这样。如果你坚持使用 Transformer，你确实无法改善这一点。好了。这里暂停一下，有问题吗？所以现在每当你听到人们说，哦，推理是内存密集型的，你就知道为什么了。好了。那么现在让我们用这些直觉和计算来思考我们的推理指标，吞吐量和延迟，以及 TFRT。好的。所以推理是内存密集型的。现在主要的——我是说，在某些方面这简化了很多事情，因为当我们考虑事情需要多长时间时，你只需要看需要传输多少内存，对吧？因为假设你重叠了通信和计算，那么瓶颈将只是你需要处理的内存量。
[35:49.470] 所以从某种意义上说，这很好，因为它更简单。但从另一方面来说，你的加速器闲置在那里什么都不做，这又让人沮丧。那么让我们来看一个例子。比如说 Llama 2 130 亿参数模型在 H100 上。那么延迟和吞吐量是多少？记住 Llama 2 130 亿参数模型有一个特定的形状。这是序列长度、模型维度、前馈维度、查询头数量、键值头数量。这里没有 GQA。所以 N 等于 k，头维度、层数、词表大小，以及 H100 的内存带宽。那么现在，使用这个配置，让我们计算一下 Transformer 的性能统计数据，好吗？
[36:50.960] 这些是我们的输入。我要计算的统计数据包括参数量、内存使用、延迟和吞吐量。首先，什么占用内存？参数占用内存。所以我们计算参数数量。然后你查看嵌入层、MLP 层、KQV 的投影层。最终你会得到某个参数数量。假设我们使用 bf16，这在推理中总是如此，那么参数占这么多字节。内存中还有 KV 缓存。KV 缓存的大小是序列中的令牌数乘以头数和 KV 头数，再乘以头维度，再乘以层数。键有一个，值有一个，然后 bf16 有多个二倍的因子。所以这就是 KV 缓存的大小。总内存使用量是——这是每个序列的量，一共有 B 个序列。所以就是 B 乘以那个加上参数大小。这就是你所需的内存量。现在，延迟是多少？延迟是由内存 I/O 决定的，对吧？因为推理是内存密集型的，大多数情况下假设你重叠了通信和计算。
[38:49.720] 所有内存都将在 HBM 和 SRAM 之间来回搬运参数。这就是移动内存所需的时间。吞吐量是延迟的倒数，但我们同时生成 B 个令牌。所以这是每秒令牌数。这是每个令牌的秒数。这是每秒令牌数，同时还有一个额外的 B 因子，因为你正在处理大小为 B 的批。那么现在，让我们计算一下对于这个配置，这些实际值是多少？参数量是 130 亿。所以一个好的合理性检查——这确实是宣称的 130 亿参数模型。内存是这个项，大约是 8.38 亿乘以 B。所以它基本上是关于 B 的线性函数加上一些其他项。所以这是 KV 缓存，它随 B 增长。这是参数，是参数数量的两倍。现在，延迟就是这个除以内存带宽，所以它和内存有相同的形式。
[40:15.050] 吞吐量是 B 除以那个。所以注意到延迟，当你增加 B 时，它会增长，因为 KV 缓存会变大。为了处理数据，你必须来回复制 KV 缓存。现在，吞吐量更有趣，因为当你增加 B 时，吞吐量确实会改善，但会达到一个极限。吞吐量改善是因为你现在在更大的批次上分摊成本。但同时，你处理这些数据的速度显然也在增加。所以它会渐近到一个极限。它不会——吞吐量不可能趋向无穷大。到目前为止有什么问题吗？所以基本上，你只需要记住延迟是关于 B 的线性函数。这里的常数是 KV 缓存的大小，然后这是参数量。然后吞吐量大致正比于 B 除以 B 加某个数。那么现在，让我们把这个具体化到各种情况中，好吗？
[41:40.140] 所以如果你有批大小为 1，这就是你会得到的。你得到每个令牌 0.008 秒的延迟。吞吐量是每秒 124 个令牌。那么如果你增加批大小会发生什么，你会看到延迟上升，但吞吐量也会上升。所以延迟变差了，但吞吐量改善了。这是件有趣的事情，因为我们想的是，哦，我们只是想让速度变快。但“快”在这里其实有两种含义，而这取决于你关心哪一个，是完全相反的。如果你想调整批大小，这真的会决定你是想要低延迟还是高吞吐量。那么如果你增加——比如说，增加——我们真的想要高吞吐量，因为我们要处理大量文档。让我们把批大小再增大一些，好吗？
[42:42.970] 所以延迟会变得更糟。我是说，是的，它会变得更糟。吞吐量会变得更好。但这里的主要问题在于你的内存，你会耗尽内存，因为用于存储 KV 缓存的内存将超过你的 H100 内存。如果你有 B200，你可以进一步增大批大小，但最终你总会达到某个极限。所以你能提升吞吐量的程度是有限制的。你永远无法达到渐近线，因为你会撞上内存瓶颈。而且，你的吞吐量提升——收益实际上也在递减。因此，增大批大小会恶化延迟，因为现在你需要读写更大的 KV 缓存。记住，这是批处理。如果你是单个查询，你必须等待所有人都完成。所以你基本上就是——你在等公共汽车，延迟相当高。你等着，然后出发。而公共汽车的吞吐量则相当不错，因为你可以一次运送所有人。
[43:59.930] 所以随着批大小增加，吞吐量会提升，因为记住，参数是共享的，你只需要将它们加载到内存中一次。然后你就可以处理大量的序列。所以这是在延迟和吞吐量之间做权衡，为了确保大家理解一致——较小的批大小带来更好的延迟，但吞吐量更差。较大的批大小带来更好的吞吐量，但延迟更差。好的。我不打算深入讨论并行计算。推理还有另一个维度，就是你可以将模型切分到多个设备上。如果你想了解更多，可以查阅关于推理的 Scaling 一书章节。举个非常简单的例子，如果你启动模型的 M 个副本，延迟不变，吞吐量增加 M 倍。还有一个我们没提到的指标是首 token 时间（TTFT）。这本质上就是预填充（prefill）所需的时间，因为完成预填充后，你基本上就可以开始生成了。
[45:11.840] 所以如果你想要——如果你想要更快的 TTFT，你应该使用较小的批大小，而要提升吞吐量则要用较大的批大小。希望这已经说清楚了。关于吞吐量和延迟以及它们之间的相互作用，有什么问题吗？好的。现在我们已经有了一个概念框架，用以从算术强度、吞吐量和延迟的角度思考推理效率。接下来让我们试着让它更快。我们如何让推理更快？这里有各种不同的技术，种类繁多，从改变模型架构到系统优化，以及介于两者之间的各种方法。所以推理在某种意义上是一个非常丰富的跨领域话题。现在回过头来看，最显而易见的，也是希望已经深深印在你脑中的一点是，内存是推理的瓶颈，而 KV 缓存占据了大量的内存。
[46:39.050] 而且，如果批大小足够大，KV 缓存甚至可能比参数数量还要大。所以让我们尝试减少 KV 缓存的大小。不过，你需要注意如何做这件事，因为你要确保在此过程中不会损失太多精度。这里有一件你可以做的事情，我们之前已经讨论过了，那就是分组查询注意力（GQA）。提醒一下，多头注意力（MHA）——基本上对于每个 token，你都有一个键（key）、一个值（value）和一个查询（query）。如果你使用分组查询注意力，你基本计算相同数量的查询，但只有更少的组，你为每个组计算键和值。这里的 K 就是组的数量。在多头注意力中，K 等于 N，没有缩减。还有一种叫做多查询注意力（MQA）的东西，没人用它，因为它效果真的很差。K 等于 1。而在两者之间的某个位置，我们希望能在精度和速度之间找到平衡。
[47:54.920] 这就是 2023 年引入 GQA 的那篇论文。他们展示的是，如果你看每个样本的时间，这与延迟和吞吐量都有关系，你会发现 MHA——即多头注意力，完整的注意力机制，时间很高。而当你开始让 K 等于 1，它会快得多。然后你实际上可以继续增加 K 到 K 等于 8，效果仍然相当不错。而最终，你的时间会大幅上升。好的，那为什么 GQA 会改善延迟和吞吐量呢？它将 KV 缓存减少了 N 除以 K 的因子。顺便友善地提醒一下，减少内存使用能带来加速，因为我们是内存受限的。让我们重新看看我们熟悉的 Llama 模型。记住在初始配置中，我们使用的是 K 等于 N。所以那里没有多头——那是多头注意力，键和值的数量没有减少。对于这种配置，记住使用批大小 64 时，我们得到这样的吞吐量和延迟。
[49:27.510] 现在如果使用 GQA，我们加入一个，比如说，1/5 的稀疏度（键值组数）。这相当大地减少了内存占用，进而改善了延迟也提升了吞吐量。所以延迟和吞吐量并不总是一致的。如果你减少内存使用量，这会让两者都得到改善。这主要是批大小这个维度造成的。那是矛盾的关键点。好的，这很棒。让我们实际上继续增大批大小，看看会发生什么。之前，如果我们使用批大小 256，我们会内存溢出。但现在，它能够放入内存了。我们可以看到，延迟受到一些影响，因为我们增大了批大小，但吞吐量按比例上升了。所以有时你可以联合调整这些参数。你可以减少 KV 缓存，但那样你就能增大批大小，从而允许进行其他权衡。
[50:52.810] 所以每当你做一些有损的更改时，最后必须做的事情是确保精度不会下降。这篇论文表明，对于 GQA，时间方面更好。而且在各种评估中，它基本都表现良好。那么关于这些精度数值，我认为你始终需要持保留态度，因为这是针对特定模型的。后来，DeepSeek 的论文应该会表明它实际上确实有损害。所以我想，这里所有那些不只是数学推导的东西，都要持保留态度。好的，说到 DeepSeek，这里有另一个减少 KV 缓存的想法。主题就是减少 KV 缓存，延迟和吞吐量就会提升。这是多头注意力，每个 token 都有相同数量的查询、键和值。而 GQA，记住，我们减少了每个 token 的键和值的数量。抱歉，不完全是针对每个 token。我们基本上是减少了值和键的总数量。
[52:07.280] 而现在，DeepSeek 的多潜变量注意力（MLA）说，我们实际上要让每个 token 的键和值的数量保持相同，本质上。但我要对它们进行参数化——我要压缩它们。通常，你有某种——你是如何计算你的键和值的？你有激活值，你将它们乘以某个矩阵得到 K，乘以另一个矩阵得到 V。这些通常是 N 乘以 H 的维度，就像你的模型维度一样，那是很大的。所以 MLA 说，我将把这些激活值投影到一个 C 维空间中。DeepSeek-V2 将它从 16000 减少到 512，所以这是相当激进的压缩。然后我将从这个压缩表示中计算出键和 V。现在，我可以只存储 C。这个小得多。而当我需要我的键和值时，我可以直接将它们物化出来。
[53:14.390] 不过这里有一个小问题，即 MLA 与直接作用于键和值的 RoPE 不兼容。所以他们的做法是增加额外的维度来处理 RoPE。但大致来说，这仍然是一个相当大的缩减。因此，延迟和吞吐量的提升只是简单的数学计算。KV 缓存越小，速度就越快。在某一点之前几乎是线性扩展的。然后，记住，你需要检查你的模型是否准确。所以首先，这个结果与 GQA 论文的结果有些矛盾或存在张力。他们表明 GQA 实际上并没有那么好。这是 MHA。这是 GQA。这些数字比这些数字小。但他们表明他们的方法，即 MLA（Multi-Latent Attention，多潜变量注意力）甚至比 MHA 表现得更好一点。但我们可以说它大致相同。
[54:35.820] 所以这一列和这一列比——我猜这个表上没有 GQA。你得在这里比较。所以我会给你们看一个——是的。这相比于减少模型维度相比如何？所以问题是，这与减少模型维度相比如何？这是个好问题。这些消融实验没有展示这一点。我的猜测是，减少模型维度只会让事情变得更糟，因为你不加区分地减少了一切。我认为这类方法的诀窍在于找到模型中你可以压缩的地方。而且这种经验性的做法，我认为你不一定能确切知道。你只能做大量的实验，看看什么有效。所以这里有一个减少 KV 缓存的另一个想法。这被称为跨层注意力（cross-layer attention）。这个想法是，通常每一层都有 KV，即 K 和 V。
[55:43.790] 但假设我们不做这个，我们只为一部分层计算 KV。然后，对于这一层，我就直接使用前一层的 KV 缓存。所以这是另一种共享方式。就像 GQA 在头之间共享 K 一样，现在我是在层之间共享 KV。而且根据经验，这篇论文表明这样做改善了帕累托前沿（Pareto frontier）。所以这些模型中的每一个都更好——给定一种方法，你总是可以通过改变 K 和头的维度来扫描 KV 缓存的大小，我认为这与你关于改变模型维度的观点有些关联。但如果你采用这种 CLA（Cross-Layer Attention，跨层注意力），效果更好。继续往下，快速浏览一下减少 KV 缓存的不同技术。有局部或滑动窗口注意力（local or sliding window attention）。这是一个相当古老且非常自然的想法。如果你看完整的注意力矩阵，它是 N 的平方。
[57:14.260] 而不是那样做，如果你要生成一个 token，你只需看最后 K 个 token。所以你本质上有一个滑动窗口。对于你生成的每一个 token，你只依赖于最后 K 个。现在，如果你这样做，有效——KV 缓存现在就独立于序列长度了。它只是批次大小乘以其他变量，这很棒。这尤其适用于长上下文。现在，由于层数的原因，有效上下文长度实际上大于声明的上下文长度，因为信息可以在层间向下传播得更远。现在，你可以做更花哨的事情。你可以不采用密集的层选择，而是将其间隔开。你也可以做这种全局加滑动窗口的方式，即你对一个固定的不同 token 点网格进行注意力，再加上一个局部滑动窗口。
[58:25.820] 所以你可以做各种事情。现在，这样做的问题在于它实际上仍然会损害准确性。这降低了表达能力。这里没有免费的午餐，或者至少这是一顿昂贵的午餐。所以人们想出的解决方案是将局部注意力与全局注意力交错。这些混合模型在部分层使用完全注意力，在另一些层使用局部注意力。你基本上总是在尝试——这样就可以让你在一定程度上减少 KV 缓存。而且你总是在尝试平衡准确性和速度。是的。就混合模型而言，使用线性时不变（linear time invariant）与使用滑动窗口之间的权衡有区别吗？是的。所以问题是，线性注意力与滑动窗口注意力相比如何？我不打算讨论线性注意力。
[59:28.530] 但很快地说一下，有一大堆方法，它们不是存储 KV 缓存，而是基本上计算整个历史的某种压缩表示。所以线性，最朴素的线性注意力就是你将 KV 值求和成一个单一的向量。这绝对是独立于序列长度的。你可以做更花哨的事情。有门控网络（gated net）、delta 网络（delta nets）和 Mamba，这些方法试图压缩但又不忘记太多。现在的问题是，这些方法如何比较？这些也被用来替代滑动窗口注意力，人们用它们取得了不错的结果。你也可以组合使用完全注意力、滑动窗口注意力和线性注意力，因为它们捕捉不同的方面。如果你关心局部高分辨率的内容，那么滑动注意力更好。如果你只想要对过去广泛的总结，那么其他线性注意力可能更好。
[01:00:38.460] 那么对于长上下文的句子，你会说你的注意力会更适合这种场景吗？所以问题是，对于长上下文，线性注意力会更好吗？没有免费的午餐，对吧？假设你有一个非常长的上下文，并且你在解决一个大海捞针的问题。如果你必须将整个历史压缩成一个小上下文，你就会丢失信息，可能就无法检索到它。我猜我觉得人们似乎倾向于采用混合架构，你总是需要一些更长距离的注意力，但我只是想在我脑子里理解，使用滑动窗口与使用一些 Mamba、delta net 层之间的权衡是什么。使用 delta net 层是否总是比使用滑动窗口更好？或者实际有代表性的权衡是什么？是的。也许我会说 Mamba 和 delta net 比滑动窗口注意力更强大。
[01:01:46.805] 也许你可以把 Mamba 想成可能——它当然可以表示滑动窗口注意力的某些方面，因为你可以——当你在做递归的时候，它可以只看最后的状态。所以是的，也许你可以认为线性注意力或其扩展更好。它们有更多的空间。一旦你做了滑动窗口注意力，你就做完了。没有其他你能做的了——[听不清] 是的。好的。那么让我快速过一下这个。重点提一下这个 DeepSeek。DeepSeek 持续在创新不同类型的注意力机制。
[01:02:30.670] 记住，他们提出了多潜变量注意力，它压缩了键值。还有这个叫做——现在他们有压缩稀疏注意力（compressed sparse attention），DeepSeek 稀疏注意力（DeepSeek sparse attention），和重度压缩注意力（heavily compressed attention）。我一直记不住所有这些缩写以及它们的含义，但让我们看看这张图。通常，你有你的 KV token 和你的查询 token。压缩注意力基本上会将每 M 个 token 压缩成一个 token。然后还有一个叫做 DeepSeek 稀疏注意力的东西，它基本上会选择保留这些 token 的一个子集。你选择子集的方式是，你实际计算一些 LiDAR 权重查询和键，然后你做一个较小的注意力。
[01:03:24.340] 从而得到这些索引分数，你就知道要保留哪些。这是一种极快的方法来确定你需要保留哪些 token。然后你使用这些。接着，还会有一些进一步的压缩。为了节省时间，我就不展开讲了。这一节的目标是减少 KV 缓存，因为 KV 缓存与内存相关。而且我们之前看到推理是受内存限制的。所以这直接转化为吞吐量和延迟的改善。关键在于不损害准确率。你可以在不同层之间、不同头之间以及头维度上使用更低维度的 KV 缓存。你可以使用局部注意力。你可以使用线性注意力，这在前面已经讨论过了，还有更多方法。还有扩散模型，这是一种非自回归的生成方式，可以快得多。好的，让我们谈谈其他一些重要的想法。
[01:04:33.270] 量化与其说是一种架构视角，不如说是一种系统层面的视角，用来让东西变得更小。核心思想就是降低数字的精度。更少的内存意味着更低的延迟和更高的吞吐量。显然，你需要注意准确率。量化有很多选择，从 bf16 一直到 int4。如果你担心量化会搞砸，你可以做的一件事是，在训练模型时就考虑到量化。这就是量化感知训练。在训练期间的前馈传播过程中，你进行量化和反量化，本质上就是在训练时模拟这些量化误差。然后，通常情况下，权重会适应量化。效果会更好。但缺点是它需要昂贵的大规模训练。所以人们通常的做法是，先训练模型，事后再进行量化。这就是训练后量化。它通常便宜得多。一种简单的方法是对每一层或每个张量，确定每个张量的缩放因子和零点。
[01:06:11.390] 然后你分别量化它。这种方法通常效果不太好。你可以使用一种叫做 GPTQ 的想法，它利用一些 Hessian 信息来逐层量化。然后你跟踪误差，这些误差会传播到未量化的权重中。这样你就可以纠正这些误差。现在，激活感知量化是一种更复杂的方法，它观察到一些激活通道的值很大，与这些通道交互的权重更重要。所以让我们为这些权重分配更高的精度。让我们看看这张图。通常情况下，如果你拿一个 fp16 的权重矩阵并量化它，假设你量化到 int3。但你要做的是找出哪些是激活的？这些可能是激活通道。其中一些激活通道的值很大。所以如果它们普遍很大，那么你基本上为这个通道分配 fp16。其他所有通道保持 3 bit。对于少数重要的通道，你使用更高的精度。另一个思路是做模型剪枝。
[01:07:45.220] 在这里，你拿一个大模型，撕掉它的一部分，然后修复它。这是一种很粗糙的方法，但事实证明它有效。NVIDIA 有一篇论文，基本上——让我实际上——让我看看。首先你需要估计模型不同部分的重要性，选择最重要的部分。然后你基本上移除不同的隐藏单元，甚至不同的层。现在，我们有一个模型，它不会很好。所以你做的事情是进行后训练。你在你关心的数据或任务上再训一些，来修复它。从某种意义上说，这是一种减少 KV 缓存的训练方式，但你是用一个好模型的部分来初始化它。这似乎效果还不错。他们能够把一个 15B 的模型剪到一个 8B 的模型。而且准确率损失不大。
[01:09:06.390] 而且你用于训练模型或完成这个过程的训练量要少得多。好的，这里总结一下，目标是在不损害准确率的情况下降低推理复杂度。你可以把这主要看作是减少参数数量或 KV 缓存。你可以定义一个更快的模型架构来从头训练它，或者定义一个更快的模型架构，从原始模型初始化权重，原始模型可能有不同的架构，但你只是做出这个弗兰肯斯坦式的东西，然后通过蒸馏来修复这个更快的模型。好的。是的。你能详细说说如何区分重要和不重要的层吗？那么如何区分重要层和不重要层？一般来说，你有一个校准集，你把输入传入模型。你基本上看激活值的幅度。
[01:10:20.180] 其中一些，特别是如果是死单元，会接近 0。而那些值大的，你想要保留。这是一个高层次的想法。我想问的是，为什么激活值高就重要？可能是——比如说如果它总是很高呢？还是说那——所以问题是，如果所有激活值都很高怎么办？一般来说，这是一种经验观察：一些通道会比其他通道高得多。如果情况不是这样，这些技术就不一定有效，但事实恰好如此，因为这就是这些模型最终被训练出来的方式。然后你可以利用这一点。假设一个神经元在所有样本上值总是 100。那是否意味着它必然有意义？或者它可能只是训练的产物，我只是……？是的，我明白了。所以问题是，如果一个神经元总是 100 怎么办？
[01:11:24.850] 我的意思是，如果是这种情况，你也可以看方差相关的问题。如果它是 100，你不能直接移除它，因为那样一切都会出问题。如果均值高但方差低，也许有另一种方法可以直接把偏置吸收进去。好的，让我快速讲讲另一个想法。到目前为止，我们看的都是有损方法，这些方法能大幅压缩 KV 缓存，但可能损害准确率。有一种非常优雅的无损方法。这叫做投机采样或投机解码。记住，如果你在做预填充，你可以并行编码所有 token。这也会给你概率。这很快。这是计算受限的，一切都很好。而在生成阶段，是一次一个 token。所以“检查”比“生成”更快。如果我给你一个序列，快速告诉我它有多好，比一次一个地生成要快得多。
[01:12:26.320] 你可以利用这种不对称性，思路如下。我们要做的是用一个廉价的草稿模型从猜测中生成几个 token，比如说四个 token。然后我们用真正关心的模型，即目标模型 Q，来审查这些 token，接受或不接受它们。这些东西的选取是有所权衡的。草稿模型更小更便宜。所以即使它受内存限制，必须一次生成一个，也不会太糟糕。而目标模型又大又贵，但我们让它做的就是并行处理一批 token，这样它也不会太糟糕。好的，这里有一个视频展示了这个过程是如何运作的。如果你使用一个大型模型逐个token地生成，速度会非常慢。但如果你使用投机解码，你可以看到小模型先生成一批token，然后大模型基本上是在对它们进行评判。
[01:13:42.230] 然后，你基本上可以一次性获得这么一批token，接着可能又是另一批，如此反复。好的，这是算法——关于投机解码，有几篇论文是差不多同一时间发表的，这是其中之一。这里的核心思路是，为了生成内容，我们要生成K个token。我们直接从草稿模型中进行采样。所以p是草稿模型。我们采样K个token。然后，我们并行地使用q来计算这些草稿token的逻辑值。接下来，我们需要决定是否接受这些token。这就是需要用到一些数学、概率和统计知识的地方。基本上，我们会以概率Min 1（即q除以p这个比值的下限）来决定接受与否。如果q远大于p，那么q越大，就越倾向于接受它。
[01:14:57.690] 否则，我们从残差分布中进行采样，然后退出。这基本上就是拒绝采样，但不同于拒绝采样的是，有时拒绝后你什么也得不到。但在这里，我们总是能保证从目标模型中获得一个精确的样本。我会跳过这个简单的证明。它基本上和拒绝采样的论证方式相同，用来证明从目标模型中得到的概率是精确的。最初的论文表明这种方法速度很快。而且一般来说，如果草稿token数太少，你就无法充分利用目标模型侧的批处理能力。如果草稿token数太多，那么你被拒绝的概率会更高。所以会有一个最佳点，在这个例子中，大约是3或4个。通常，草稿模型比目标模型小得多。理想情况下，你希望草稿模型尽可能接近目标模型，这意味着你需要对它进行蒸馏。实际上，这意味着我们刚刚讨论的很多想法
[01:16:20.670] 同样适用于投机解码。所以基本上，思路是尝试通过各种不同的技巧来减少你的KV缓存。如果你最终得到了一个满意的模型，那就直接用它来提供服务。如果你不满意，那么它至少可以作为一个草稿模型，然后你可以用你的主模型来修复问题。目前关于投机解码有大量的文献，它们在原始版本的基础上进行了改进，这些我暂时跳过。现在很快来讲讲动态工作负载。这是一个用例场景：你在运行一个实时网站，用户过来与你的模型聊天。请求在不同时间到达，它们有共享的前缀，也有不同的长度。所以这种情况相当混乱。这远远不同于那种非常简单的一次性处理相同数量token块的训练过程。那么在这种情况下该怎么办呢？
[01:17:30.680] 这里有一个思路——有一个叫Orca的系统被开发了出来。这实际上是很早以前的事了，它引入了连续批处理这个概念。思路是你收到了这样一堆请求。这里是第一个请求的前缀，你正在生成这个token。这是第二个。这是第三个、第四个。它是参差不齐的，因为每个前缀的长度都不一样。我们要做的是逐步解码。每一步，你都为所有序列解码一个token。下一步，你再为所有序列解码另一个token。然后，如果一个序列结束了，你就把它从这个批次中移除。当新请求到达时，你又把它加入这个批次，然后继续处理。这就是为什么它被称为连续批处理，因为这个批次是动态更新的，要么移除已经完成的旧序列，要么加入新序列。
[01:18:34.770] 现在，这里的一个问题是，我们之前看到的所有批处理工作都假设所有序列具有相同的维度。你处理的是张量，每个切片都有相同的维度。但如果你的请求长度不同，你该怎么办呢？这里有一个叫做选择性批处理的概念。假设你有一个长度为3、一个长度为9和一个长度为5的请求。在注意力计算中，你对此无能为力，因为注意力取决于你的序列长度。所以如果你有一个3乘3和一个9乘9的计算，你无法有效地共享张量。但是对于非注意力部分，比如MLP层，它占据了大量的FLOPs，你可以简单地将所有序列拼接在一起，形成一个巨型序列，然后进行处理。最后一个想法是分页注意力。这是在BLM论文中提出的。当然，现在的BLM已经有很多其他的功能了，但这是当时的核心思想。
[01:19:49.030] 那么问题是，KV缓存是如何存储的？如果你考虑传入的请求，你必须把它们放到某个地方的内存中。一般来说，存在一个碎片化的问题。这就是硬盘上曾经发生过的——或者过去经常发生的事情。在过去，你不得不对你的硬盘进行碎片整理。所以有两种类型的碎片化。一种是，你必须分配足够的缓冲区，因为你不知道序列什么时候会停止。所以你可能有一个1,024的最大token限制，你就必须分配所有这些内存。而且你不能在里面放入其他任何东西，因为你只会一直生成直到达到最大token数，这非常浪费。这是内部碎片化。另外，在多个请求之间也可能存在空隙。这些空隙可能太小，无法被有效利用。
[01:20:50.040] 所以这就是纯粹被浪费的空间。这里的解决方案是——这些是搞系统的人，所以他们懂操作系统。他们说，好吧，我们以前解决过这个问题，所以我们在这里用同样的思路。我们将一个序列的KV缓存划分为非连续的块，好吗？所以如果你有这样一个序列，“Four score and seven years ago our fathers brought forth”，我们只是把它切分成这些大小为4的块。块的去向并不重要，但它们会按照块进行对齐。所以这里面有一些统一性。这样，当两个请求共享相同的——实际上可以共享相同的KV缓存。你可能会有这个块和那个块。这个块可能放在这里和这里，那个块可能放在那边。它们是交错放置的。但只要你有索引，并且跟踪所有内容的位置，就没问题。
[01:21:58.410] 特别是，如果你有系统提示词，那么你实际上可以只缓存一次这些系统提示词及其KV缓存。这对于所有查询都是有用的，好吗？这非常有用，因为如果很多人使用相同的系统提示词，你就不必为每个请求都计算缓存了。此外，在许多应用中，你会有相同的提示词，但你实际上想要生成多个不同的响应。在这种情况下，你也可以共享提示词的KV缓存，只生成各自独特的响应。好的，举个例子，如果你要生成，比如说，从“Four score and seven years ago are______”开始的多个生成结果，会发生的情况是你会得到“Four score and seven”作为前缀。然后你可能会从“years ago”或“our”开始生成。这就叫做写时复制语义。你保留这个。然后你有这两个采样结果。如果它们碰巧采样到了相同的token，你就继续使用那个token。
[01:23:17.010] 但如果它们采样了不同的 token，你就会拆分这个块，然后就可以在那里继续下去。所以基本上你是在尽可能多地共享前缀缓存。还有很多其他优化，比如内核优化，但我没有时间一一覆盖了。但总体思路是，你是在用这些操作系统的隐喻来管理你的推理过程。所以这里总结一下——推理真的、真的非常重要。它和训练非常不同。即使是同一个模型，但你在让模型做非常不同的事情，最终会变得非常受内存限制。而且如果你是在实时聊天机器人场景中，它也是动态的。我们看到了各种不同的技术来改进推理。你可以做量化。你可以设计新的架构。你可以做剪枝和蒸馏。你还可以使用投机采样。但所有这些都遵循着同一个原则，那就是减少你的 KV 缓存，但不要让准确性损失太多。
[01:24:28.540] 然后还有一些来自系统领域的思想，比如分页和投机执行，这些可以被实际应用于实时推理服务器。有一件事我们实际上没有太多机会讨论，只是简单提了一下，那就是我认为新架构在改进方面有着巨大的潜力，比如状态空间模型、线性注意力或扩散模型。在某种程度上，KV 缓存和注意力机制的构建方式从根本上使它成为一种对推理不太友好的架构。所以如果你能提出一种为推理而设计的新架构，就像 Transformer 当初不是为了推理而设计的那样，这或许能解锁很多东西。好，那我就讲到这里。下次课，Tatsu 会回来讲缩放定律的第二部分。
CS336 2026 · Lecture 10 · Inference
- 视频：官方视频
- 字幕：en-US · manual captions
- 处理：连续字幕合并为自然段；每段保留首条字幕的起始时间戳
- 状态：原语言字幕稿，待校对
正文
[00:05.320] All right. Let's get started. So last time, Tatsu talked about scaling laws, and we're going to take a little break from that and talk about an inference. So the problem of inference is very simple. You've trained a model, and you're given a prompt, and you want to produce the response usually as accurately and as quickly as you can. So inference turns out to be only one lecture, but it's actually of growing importance. And it shows up in many places. So clearly, once you've trained a model, you don't just sit there. And if you're a researcher, you put a plot in your paper of here's how my model works, but for everyone else, you actually want to use this model. So what does use look like? It could be chatting with an AI assistant or chatbot, it could be doing code completion. These days, agents are very popular, and that requires inference. Batch data processing as well. It's also used for evaluation. For evaluation that requires generation in particular.
[01:19.670] And it's also used inside training. If you're doing reinforcement learning, you need to have a model, you generate rollouts, and you score them, and you update the weights appropriately. So inference shows up all over the place. And efficiency, as is the theme of this class, really matters more than ever here. And even if you just look at the actual use app of vantage point, training is a one-time cost. It could be very expensive. It is very expensive. But once you're done with it, that's it. But inference is a repeated cost. You kind of incur them every single day. So OpenAI is estimated to produce 8.6 trillion tokens a day. And just for reference, DeepSeek-V4, which came out earlier this year, was trained on 32 trillion tokens. So in less than four days, the number of tokens that OpenAI has to produce, and therefore the compute is at least DeepSeek-V4.
[02:26.540] And of course, even the frontier models might be trained on more tokens, but still, inference, the number of tokens generated is going to be higher. So that's just some perspective. And moreover, I think the importance of inference has grown in the last year because once upon a time, we thought of language models as primarily as chatbots or assistants, where you put in a prompt, you get back a response, and presumably, the goal of the human is to read the response. But now, as we move more into a genetic world, what it looks like is that a query goes in, and the agent is going to do a bunch of stuff. It's going to think, it's going to reason, it's going to call some tools, it's going to introspect. And at the end of the day, it will produce some output for a human to read. So most of the tokens that the agent produces is actually not for reading. And so you should really think about the number of tokens generated as really the compute spend. And there's no limit.
[03:34.110] If you have an ambitious enough problem, you're going to need much compute and a lot of tokens. So if you were in the chatbot world, maybe inference after a certain point, it's fast enough, because humans can only read so fast. But for an agent, there's no limit to how much a value you can get out of squeezing more out of inference. So there's a lot of people doing inference on the commercial side. Of course, all the closed API providers have to serve their models, so inference is a big deal for them. And there is also a bunch of providers serving open weight models and providing inference as well. And in the open source community, there's a bunch of packages. So VLM is probably a popular one. Kind of a go-to. SG Lang is another one that's particularly good for agentic workloads, but maybe not as popular yet. And then there's TensorRT from NVIDIA, which is really fast, but it's more narrow.
[04:42.190] And then Llama CPP, if you want to run inference on CPU, this is a popular package. So hopefully, I've argued that inference is huge. If you can make inference twice as fast or even 10% faster, that is a big deal. So now, the question is, what does fast mean? So there's a few metrics that capture the notion of fast. And they're going to be applicable in different settings and have different trade offs. So one metric of fast is time to first token or TTFT. This is essentially how long a user waits before any generation happens. So you put in a query into ChatGPT, and some number of milliseconds passes by, and then from the point the first token comes in, that is the time. And this is generally useful for interactive applications because that latency where you're just waiting and doing nothing is-- the longer that is, the worse the user experience, and as soon as tokens come in, it doesn't maybe
[05:54.210] have to be that fast, because if you're going to read it anyway, you can't read that fast. So the second metric is latency. And this is from the standpoint of an individual user. How fast are tokens appearing for one query? This is also important for interactive applications. So basically, this is how fast the tokens are kind of streaming through. And then the second related concept is throughput, measuring tokens per second, this is how fast tokens are appearing for many queries. So latency or one latency and throughput are clearly very related. It's very kind of aligned-- in general, many interventions will make latency and throughput better. But as we'll see later, there's actually a trade off here. And throughput is useful for if you're doing a bunch of batch processing. You have a petabyte of data, and you want to process it with a language model.
[06:57.000] You just want that job to get done. It doesn't matter if one query is coming in faster, sooner, or later. OK. So what determines the efficiency by either any of these metrics of inference. So the high level bit is as follows. And this is-- maybe one high level bit to remember from the lecture is that when you're doing training, you see all the tokens at once, because in supervised fine-tuning, you see all the tokens, and you can parallelize over the sequence. Remember, think about in the transformer, the calculation for attention, MLP, the sequence is just a dimension. So it's a big tensor that gets multiplied. And so you, essentially, process all the tokens at once. In inference, you can't do that. You have to generate because of the autoregressive nature of inference. Tokens sequentially. One at a time. So this is the fundamental problem of why inference is a very different workload than training, because you can't parallelize across the sequence dimension.
[08:07.030] And therefore, as we'll see later, it's going to be harder to have high arithmetic intensity or fully utilize the compute. OK. All right. So let's go on. So first, I'm going to-- just maybe a preview of what's in this lecture. First, I'm going to do some math to understand the arithmetic intensity and throughput and latency, how to think about that for a transformer. And then I'm going to talk about various techniques to reduce the cost by making the kv cache smaller, quantizing model pruning, and then finally, I'll talk about speculative decoding, and then some kind of practical concerns. All right. So this is going to be a bit of a review, mostly kind of swapping in, making sure we're on the same page with respect to notation. Much of this lecture is based on the scaling book from Google on transformers and inference. So I do recommend that people go check it out. It's a very nicely written. And some of the figures are generously lifted from that book. OK. So just to establish a bit of notation to understand this diagram.
[09:30.380] So we're going to use symbols. Think about it in ienops to denote dimensions, but also they're kind of length. So B is going to be the batch dimension and also the number of sequences. T is going to be the sequence dimension, which is also the number of tokens. D is the model dimension, H is the dimension. So if I'm going to write this, this is taking a product between a tensor with dimensions B, T, and D, and another matrix of D and H, and I produce BTH. So what is going on here is that there are some dimensions marked red, which are going to be contracting. These dimensions appear in both the operands and disappear from the results. And then there's going to be other regular dimensions in black that appear only in one operand and stay in the result. So this is the general form. Another case is if you have blue dimensions, that means they're attaching dimensions, and they appear in both, but stay in the result.
[10:36.630] They do not get contracted or reduced. OK. So the interpretation here, as we'll see later, is that we have this tensor for every sequence and token position. You have a vector and you want to multiply that by a matrix. And then here, you're basically taking a bunch of dot products between two tensors. OK. So with that in mind, this in all its full glory, is the description of a transformer block. So it looks a little bit like a circuit diagram. There's a lot of details here, but I actually find this probably the most crisp definition of what a transformer is. You see a lot of description of transformer, and frankly, a lot of them are-- it's hard to understand what is exactly happening. This one basically tells you exactly the shapes of the tensors and allows you to reason about the dependency. So for transformer block, you take the X, which is the activations of one layer, you feed it through a tension, and then you feed it through MLP.
[11:44.020] And just as a kind of a quick refresher, you projected using a query matrix, a key matrix and a value matrix, you'll see here that we have N, which is the number of heads, and H is the dimension of each head. So the query matrix is batch times sequence times number of heads times the head dimension, and K and V, instead of having N, we have K, which we'll see later, which we have a potentially smaller number key value heads. And then this is the attention operation, where notice that there's batching dimensions. So B appears in both. And we're contracting over the head dimension. And we'll come back to why this B, B in both, makes inference hard. Why attention is kind of a bottleneck there. So just hold on to that thought. And the MLP is very simple and straightforward. You just do some math models. This is the gating matrix. This is the upper direction and a down projection. So thank you guys. I have implemented this.
[13:07.470] So I don't want to belabor it, but just to swap in the notation. By convention, we're going to assume that in the MLP, the MLP up projects the D dimensional, model dimension into four times that. So F is always-- think about it as 4D whenever you see F. And then the model dimension gets split across the N heads. So the model dimension is always equal to number of heads times the head dimension. And then the number of heads get split, in the case of group query attention, into the number of groups and the number of heads per group. And finally, we have two variables, S and T, which both represents the sequence dimension. And the difference is that S is going to represent the number of input tokens, and T is going to represent the number of output tokens. So in training time, these two are the same, because we're predicting all the same inputs as outputs. And inference time, T is going to be 1, and S is going to be the input.
[14:22.960] OK. So another review for arithmetic intensity. So I talked about this in the second lecture. So just as a warm up. Suppose I'm multiplying two matrices; a B by D matrix and a D by F matrix. So intuition. B is the batch dimension, D is the hidden dimension or the model dimension, and F is the projection dimension in the MLP. So remember, what are the arithmetic intensity? I have to count flops, and I have to count the amount of memory moved. So we can do this as follows. So what do you have to do to multiply these matrices? Now, remember, the systems, lectures, you have to read X from HBM. And if you're storing everything in bf16, which is the case for inference, we're going to have 2 times B times D, you're going to read W. That's going to be another 2 times D times F. You're going to do the matmul. So that's number of flops is 2 times B times D times F And then you're going to write it back.
[15:33.180] OK? So every time you read and write, it's basically the number of entries. It's kind of a quadratic-like term, and matmuls are qubit. And remember, this is going to be important because that's how you're going to get high arithmetic intensity. So total number of flops is the cubic, and then the bytes transferred is the sum of the reads and writes. So remember, arithmetic intensity is how much compute we do per byte transfer. And we want this to be high. So the intensity is going to be equal to flops over bytes transfer. And I'm representing all these things symbolically, because it will be easier to see and work with rather than numbers. One just simplification we can make is that suppose that the batch dimension here is much less than D and F, then we can simplify this. So in particular, what this is doing is saying D equals C times D and F equals C times B and let C goes to infinity, so both D and F are really large and B is smaller, and this just reduces to the intensity of B.
[16:51.590] So the upshot of that is that remember when we did arithmetic intensity for matmul, it was something like N/3. And this is the kind of analog, where instead of just having square matrices, we have non-square matrices. OK. So just to compare, so the accelerator intensity of a hardware is you look at the flops per second in the spec sheet, you look at how much memory bandwidth, how fast memory gets moved between HBM, you divide, and that is the accelerated intensity. And now, you compare your computations intensity with your accelerator intensity. And if it's greater then your compute bound, which is good, if you're less, then you're memory bound, and that's bad. So in this case, for H100, for this particular matmul operation, you're compute bound if the batch size is greater than 295. So here's an extreme case. So suppose you just have one example. What happens is you have one example, then your arithmetic intensity is going to be one.
[18:16.320] And this is going to be memory bound. And you can see what's going on here. You're reading this DF matrix, but B is only one. So you're effectively doing only 2 times D times F flops. And this is basically the workload that you'll see in inference. You don't get these full matrices, you get these very thin matrices or tensors. OK. Let me stop there in case there's any questions. Yeah. In the transformer slide, I just want to confirm [INAUDIBLE]. You said, across [INAUDIBLE] can search across [INAUDIBLE]? Yeah, should that be across tables instead of cubes? Let's see. So the number of groups here is-- wait, I'm just making sure.
[19:31.790] So K is the number key value heads. And so you have g per KV head. My understanding is [INAUDIBLE] each key path will have one [INAUDIBLE] correspond to it. And [INAUDIBLE] then g is the number of [INAUDIBLE] p. That means we have k. That's my question. Yeah, OK. Yeah, I think you're right. So k should be the number of groups, and g should be the number of the heads per one of those groups. Yes. Yeah, OK. Thanks for that. I'll fix that later. OK. So let's talk a little bit more about the arithmetic intensity of inference. So here's how to think about what's happening in inference. Let's say you just do it naively. So here is a prompt never going to give you. I'm going to the transformer. And the transformer does generates keys and values and activations. And at the end of the day, it generates logits over the output vocabulary.
[21:02.720] And you sample a token, and that token gets concatenated with the prompt, and then you just do this again. You generate never. You attach that to the prompt and so on and so forth. So this is the most naive thing. If you have a black box that takes a sequence and outputs a distribution over tokens, which is exactly what a transformer does, you can just apply this repeatedly. OK, so this works, but this is really bad because each time you generate one token, you actually take order T squared time where T is the number of tokens that you've generated so far. So generate T tokens is actually T cubed because the attention is already T squared, and you have to do that once for every token.
[22:00.910] So this is pretty bad, but the observation is that you don't actually have to do this. A lot of the work can actually be shared across prefixes. So for example, if you're over generating network versus you're generating up, you're actually computing a lot of the same key values for never going to give. So these tokens, they shouldn't change, OK? And this is because it's a causal transformer. If it was bidirectional, then if you attach a token, then everything changes. But if it's causal, then the activations here don't change based on any tokens you append. So with that observation, the first obvious thing is to store a KV cache in HPM, so that, between successive token generations, you can just reuse the KV cache. So that means when you are trying to generate gonna, you don't actually have to compute all the keys and activations of these previous tokens.
[23:10.460] So this is what it looks like with a KV cache. So there's going to be two stages pre-fill. So you get the prompt, and you populate the KV cache, which is the set key and value pairs that the transformer computes, and then you generate the logit-- just the same computation as we did before. And now you have the KV cache here and then this distribution which you sample at a token. And now you feed that through the transformer, which uses this cache and then produces both the distribution of the next tokens, but a new KB, which is basically activations corresponding to the token up. And then this gets fed in. So now you have the command to KV cache. You have this logits you sample, and then that feeds to the transformer. You output the distribution over next tokens as well as the new token which you add to the KB cache and so on and so forth.
[24:19.960] So the KV cache formula is for every sequence-- so there's B of them. For every token, there's S of them. And for every layer and every head, you store H dimensional vector. And just to reiterate for inference, there's prefill. You given a prompt. You encode the prompt into these vectors. And this is parallelizable just like in training because you see the entire prompt. You can compute the KV cache. And then in generation, you generate the new response tokens sequentially. But at least you don't have to pay for generating the KV cache of the tokens that you already looked at. So is everyone comfortable with the KV cache? Any questions about this? All right, so let's try to figure out the FLOPs and memory I/O for both MLP and attention layers. So remember S is the number of tokens we're kind of conditioning on, and T is the number of tokens we're generating.
[25:39.260] So we're going to do this abstractly. And then later, we'll specify to the prefill where T equals S, and generation T equals 1. So for the MLP layers-- and through all of this, I'm just going to look at the MAmmoTH because the MAmmoTH are the thing that actually require a lot of work. Everything else can is not that many FLOPs and can be potentially fused into the MAmmoTH too. OK. So let's do this calculation. So it's the same as the matrix calculations. I'm not going to maybe belabored step through every single detail here, but just the form is you read X from high bandwidth memory. You read all the parameters. You compute the upper projection, and then you write it to HBM. You compute the gate, and you write HBM, and then you compute basically the down projection of that, and you write it to HBM. So the number of flops depends on B and T and D and F. So batch size, sequence length, model dimension, and the feedforward dimension, MLP dimension.
[27:08.450] And the bytes transferred is this expression. So now you can compute the arithmetic intensity which, again, is FLOPs divided by bytes transferred. So this is some expression. We're going to assume, that just like before, B times T is much smaller than D and F. And with that, then we see that intensity is B times T. So this is analogous to just the MAmmoTH case. There's more matrices and there's more dimensions, but fundamentally it's the same thing. And that makes sense. An MLP is basically a big MAmmoTH. And it also makes sense because the batch dimension and the sequence length dimension are-- everything is independent. And for the MLP, they don't interact. Attention is different story. And so B times T. So as long as you have-- in the pre-fill, if you make B times T large enough-- large batches, long sequences-- you'll be fine. Now, let's look at generation. So remember, generation, there's two problems here.
[28:24.430] One is that generation T equals 1. You're only generating one token at a time. And so that means your arithmetic intensity is going to be B, but B in generation is the number of concurrent requests. And so if you're in a batch setting, you can control that. But if you're, let's say, serving a chatbot, it's the number of concurrent requests is essentially how many users, concurrent users there are? Which can be high. It be low. So it's a bit unpredictable. It can be changing over time. So that's something we're going to address when we talk about continuous batching. But overall, this is not bad because as long as you have large batches, the sequence length isn't going to really help you. But if you have large batches, you should be good. So now let's look at the attention layer. So again, S is in the previous tokens already generated. T is the number of tokens you want to generate logits for. So in attention, you read the QKV matrices from HPM. You compute the attention. You compute the softmax which doesn't really matter, and compute the value matrix, and then you write the result to HBM.
[29:54.460] So the FLOPs is B times S times T times D, and the bytes transferred is this expression. Everything's is a MAmmoTH, right? So this should be always one degree higher of a polynomial than the bytes transferred because it's a MAmmoTH. The only question is, what is that factor look like? So that factor for attention is S times T over S plus T. So let's look at the prefill. So prefill, what this is saying is that the prefill intensity is S over 2 when T equals S. So this is good, right? Because as long as you have long sequence lengths for attention, you're going to maintain high arithmetic intensity. Notice that the batching dimension doesn't happen here. I'll explain a bit why later. But for generation, this is bad news. So the generated nation intensity is S over S plus 1, and that's less than 1, or even let's just call it 1. And aromatic intensity, remember 1 is bad.
[31:19.340] We want it to be something like 295 for H100 to be to saturate the compute. And so this is really the bottleneck. So we did all this analysis for MLP attention, prefilled generation, and we've found this to be a bottleneck. So let's try to contemplate why this is a bottleneck. So unlike MLPs-- so MLP is for generation is actually OK as long as your sequence is-- sorry, your batch is large enough. So the problem with this ML attention is that-- let's look at MLPs. Every sequence hits the same MLP weights. So these don't depend on B. Whereas, in attention layer, each sequence has its own KV cache. So these all depend on MLP. And so you can think about this as what's going on here is that in the MLP case, having B being big actually is helpful because you get to load these MLP weights once. And you can use it for all your sequences. That's how you get high arithmetic intensity because you use it-- simplifying a bit, you load it once.
[32:58.510] You do all your batch processing, and then that is good. Whereas, in attention, all these depend on B. So increasing B doesn't help. It's like for every sequence, you're basically doing a MAmmoTH. So they're all independent. So doing more MAmmoTHs isn't helpful. Just if you remember in the very beginning, I showed you this example. This also has pretty bad arithmetic intensity. It's not a MAmmoTH because we're essentially batching by a coordinate. And this is essentially basically the same as doing a dot product, which has horrible arithmetic intensity. And remember, this is in the attention. This blue B is the cause of why the attention is arithmetic intensity doesn't scale with B, and why that is a bottleneck. So just to summarize here. So prefill is compute-bound. Generation is memory-bound. So if you look at the MLP intensity for prefill, it's B times S, great.
[34:12.480] Prefill for attention, intensity S over 2-- not as good, but workable. Generation MLP intensity, also workable, requires long concurrent requests. But it is really the generation attention intensity which is a fundamental bottleneck, and that's it. If you're sticking with a transformer, you can't really improve this. OK. Pause there for questions. So now whenever you hear people say, oh, inference is memory bound, you know why. OK. So let's now use these intuitions and calculations to think about our inference metrics throughput and latency and also TFRT. All right. So inference is memory bound. Now the main-- I mean, in some ways this simplifies a lot of things because when we think about how long things take, you just look at how much memory needs to be transferred, OK? Because assuming you overlap communication and computation, the bottleneck is going to be just the amount of memory that you have to deal with.
[35:49.470] So in some ways, it's nice because it's simpler. But in other ways, it's frustrating that your accelerators are sitting there not doing anything. So let's walk through an example. So for Llama 2.13B on an H100. So what is the latency and the throughput? So remember Llama 2.13B has a particular shape. So this is a sequence length, the model dimension, feedforward dimension, number of query, number key value heads. There's no GQA here. So the N equals k, head dimension, number of layers of vocab size, and the memory bandwidth for H100. So now, using this config, let's compute the transformer performance statistics, OK?
[36:50.960] So these are our inputs here. And just the statistics I'm going to compute are going to be number of parameters, the memory, usage latency, and throughput. So first of all, what takes memory? So the parameters take memory. So we compute the number of parameters. And so you look at the embeddings, the MLP layers, the projections of the KQV. At the end of the day, you get some number of parameters. Assuming that we're doing a bf16, which is always the case for inference, then parameters take up this many bytes. Also in memory is the KV cache. And the KV cache is the number of tokens in your sequence times the number of heads and the KV heads times the head dimension times the number of layers. You have one for the key and one for the value, and then you have multiple two for bf16. So that's the size of the KV cache. And the total memory usage is-- this is for each sequence you have B sequences. So that's B times that plus the parameter size. And that's the amount of memory you need. And now, what's the latency? The latency is determined by the memory I/O, OK? Because inference is memory-bound, most assuming you overlap communication and compute.
[38:49.720] All the memory, it's going to be on shuffling parameters back and forth between HPM and SRAM. So that's how long it takes to move memory around. Throughput is the inverse of latency, but we're generating B tokens in parallel. So this is tokens per second. This is seconds per token. This is tokens per second, as well as having an additional B because you're processing a batch of B. So now, let's compute for this config what are these actual values? So num parameters is 13 billion. So a good sanity check-- that is advertised as a 13 billion parameter model. Memory is this term, which is about 838 million times B. So it's basically a linear function of B plus some other term. So this is a KV cache which grows as B. This is the parameters, which is double the number of params. Now, latency is just this scale by the memory bandwidth, so it has the same form as the memory.
[40:15.050] And throughput is B over that. So notice that latency, as you increase B, grows because the KV cache grows. And in order to process stuff, you have to copy the KV cache back and forth. And now, throughput is more interesting because, as you increase B, throughput does improve but up to a limit. Throughput improves because now you're amortizing the cost over a larger batch. But also, the speed at which you're processing this obviously is also increasing over time. So this asymptotes. It's not going to-- throughput can't possibly go to infinity. Any questions about this so far? So basically, you just have to remember latency is a linear function in B. This is where the constant is the size of KV cache, and then this is the number of parameters. And then throughput is some proportional to B over B plus something. So now, let's instantiate this for a bunch of situations, OK?
[41:40.140] So if you have batch size 1, so this is what you would get. You get a latency of 0.008 seconds per token. And the throughput is 124 tokens per second. So what happens if you increase the batch size, you'll see that the latency goes up, but the throughput also goes up. So latency gets worse, but the throughput improves. So this is an interesting thing because we think about, oh, we just want to make it go fast. But fast actually has two meanings here, which actually, depending on which one you care about, is completely opposite. If you want to tune your batch size, that's going to really determine if you want latency or throughput. And what happens if you increase-- let's say, increase-- we really want high throughput because we're processing lots of documents. Let's increase the batch size even more, OK?
[42:42.970] So the latency gets even worse. I mean, yeah, it gets worse. Throughput gets even better. But the main problem here is that your memory, you run out of memory because the memory, for storing this KV cache, is going to exceed your H100 memory. And if you have B200s, you can increase the batch size more, but eventually you hit some of limit. So there's limitations to how much you can increase your throughput. You'll never get to the asymptote because you'll hit the memory. And also, your throughput is getting-- your gains are kind of diminishing as well. So increasing batch size worsens the latency because now you have a larger KV cache to read and write. And remember, it's batched. So if you're an individual query, you have to wait for everyone to finish. So you have to basically-- you're waiting for a bus, and the latency is pretty high. You wait and then you go. Whereas, the throughput of a bus is pretty good because you can move everyone at once.
[43:59.930] So the throughput improves as batch size increases because, remember, the parameters are shared, and you load that once into memory. And you can process a lot of sequences. So trade-off between latency and throughput, just to make sure everyone's aligned-- smaller batch sizes use better latency but worse throughput. Large batch sizes yield better throughput but worse latency. OK. I'm not going to really talk about parallelism too much. There's also another dimension of inference, which is you can shard your model across multiple all devices. You can look at the scaling book chapter on inference if you want to know more. Just as a very trivial example, if you launch M copies of the model, the latency is the same, and the throughput increases by M. And then the other metric we didn't talk about is time to first token. And this is essentially the time it takes to do prefill because after you finish prefill, then you can basically start generating.
[45:11.840] So if you want to fast-- if you want faster TTFT, you should use smaller batch sizes, and you want larger batch sizes to improve throughput. So hopefully, that is clear. Any questions about throughput and latency and how their intention with each other? OK. So now that we have a conceptual framework for thinking about how efficient inference is in terms of arithmetic intensity and throughput and latency, now let's try to make it faster. How do we make inference faster? There's a bunch of different techniques which are quite varied, ranging from changing the model architecture to doing systems optimization and everything in between. So inference in some sense is a fairly rich cross-cutting topic. So the most maybe now in hindsight obvious thing you can think about is hopefully beat it into your head that memory is the bottleneck for inference, and the KV cache takes up a lot of memory.
[46:39.050] And it could even be larger than the number of parameters if for a large enough batch size. So let's just try to reduce the size of the KV cache. Now, you have to be careful about how you do this because you want to make sure you don't lose too much accuracy in the process. So here is one thing you can do, which we already talked about, which is grouped query attention. And just as a reminder here, so multi-headed attention-- basically for every token, you have a key and a value and a query. And if you do grouped query attention, you basically compute the same number of queries, but you have only a smaller number of groups, and you compute key and value for each group. So K is the number of groups here. And so in the multi-headed attention, K is N, no reduction. There's something called multi query attention, which no one uses because it's really bad. K equals 1. And somewhere in between is hopefully where we'll find a balance between accuracy and speed.
[47:54.920] So this is the paper that introduces GQA from 2023. And they show that if you look at time per sample, which is related to both latency and throughput, you see that the MHA-- so Multi-Head Attention. Full attention has this high time. Whereas, if you start K equals 1, it's much faster. And then you can actually keep on increasing the K to K equals 8, and it's still pretty good. And eventually, your time goes up quite a bit. OK. So why does GQA improve latency and throughput? Well, it reduces the KV cache by a factor of N over K. And just as a friendly reminder, reducing memory usage leads to speed up because we're memory-bound. So let's revisit our friendly Llama model here. Remember in the initial configuration, we're just using K equals N. So there's no multi-headed-- there's multi-head attention, no reduction in the number key and values. And so for this one, remember using a batch size of 64, we get this throughput and latency.
[49:27.510] Now if you do GQA, we're going to put in, let's say, a kind of a sparsity of 1 one five. That reduces the memory quite a bit, which, in turn, improves the latency and also improves the throughput. So it's not that latency and throughput are always at odds. If you reduce the amount of memory, then it improves both. It's mainly the batch dimension that allows. That is the point of tension. OK. So this is great. And let's actually just increase the batch size even more and see what happens. So before, if we had a batch size of 256, we ran out of memory. But now, it fits in memory. And we see that, the latency suffers a bit because we're increasing the batch size, but the throughput goes up proportionally. So sometimes you play with these parameters jointly. You can reduce the KV cache, but that allows you to increase the batch size and allow making other trade-offs.
[50:52.810] So the final thing you have to do whenever you do some lossy change is that you make sure the accuracy doesn't drop. And this paper shows that for a GQA, the time is better. But across a bunch of evals, basically it works well. So now with these accuracy values, I think you always have to take it with a grain of salt because this is for a particular model. Later, the deep sea paper should show that it actually does hurt. So I guess, take everything that's not just math with a grain of salt here. OK. So speaking of DeepSeek, here's another idea to reduce the KV cache. So the theme is reduce the KV cache and latency and throughput improve. So this is multi-headed attention, same number of queries and keys and values for every token. And GQA, remember, we have reduced the number of keys and values for every token. Sorry, not for every token. We basically have reduced the number of values and keys.
[52:07.280] And now, the multi-latent attention for DeepSeek says we're actually going to leave the number of keys and values the same one for every token essentially. But I'm going to parameterize-- I'm going to compress these. So normally, you have some-- how do you compute your keys and values? You have your activations, and you multiply them by some matrix to get K and some other matrix to get V. And these are generally N times H dimensions, like your model dim, which is big. So MLA says I'm going to actually project these activations down into C dimensions. So DeepSeek-V2 reduced it from 16,000 to 512, so this is quite aggressive compression here. And then I'm going to compute the k and the V from this compressed representation. So now, I can just store C. This is much smaller. And then when I need my keys and values, I can just materialize them.
[53:14.390] So there is one wrinkle here, which is that MLA is not compatible with rope, which operates directly on the keys and values. So what they do is add additional dimensions for handling the rope. But more or less, it's still a pretty big reduction. So the latency and throughput improvements follow by just simple math. The smaller the KV cache, the faster you go. It's almost kind of linear scaling up until at some point. And then, remember, you need to check whether your model is accurate. So first of all, this is the result that kind of contradicts or is in tension with the GQA paper. They show that GQA actually isn't that great. So this is MHA. This is GQA. These numbers are smaller than these numbers. But they show that their method, MLA, Multi-Latent Attention works even a little bit better than MHA. But let's just say it's about the same.
[54:35.820] So this column and this column are much better than-- I guess there's no GQA on this table. You have to compare over here. So I'll show you an-- yeah. How does this compare to reducing the number of dimensions, the model? So the question is, how does this compare with reducing the dimension of the model? So that's a good question. These ablations don't show that. My guess is that reducing the model dimension just makes things worse because you're indiscriminately just reducing everything. I think the trick in all of this kind of is to find places in the model where you can squeeze. And this apiary, I don't think you can necessarily know for sure. You just have to do a bunch of experimentation and see what works. So here's another idea to reduce your KV cache. This is called cross-layer attention. So the idea is that normally, every layer has KV, K's and V's.
[55:43.790] But let's say instead of doing that, we're just going to compute KVs for a subset of the layers. And then, just for this layer, I'm just going to use the previous layers on KV cache. So this is another way of sharing. Just like GQA shares K across heads, now I'm sharing KVs across layers. And empirically, this paper shows that doing this improves the Pareto frontier. So each of these models is better than-- given a method you can always sweep this the size of the KV cache by changing the K and the head dimension, which, I guess, kind of relates to your point about changing model dimension. But if you do this CLA cross layer attention, it's better. So moving on, whirlwind tour of different techniques for reducing KV cache. There's a local or sliding window attention. This is a fairly kind of old idea and a very kind of natural idea. So if you look at the full attention matrix, it's N squared.
[57:14.260] And instead of doing that, if you're going to generate a token, you just look at the last K tokens. So you essentially have a sliding window. For every token you generate, you just depend on the last K. And now, if you do that, the effective-- so the KV cache is now independent of the sequence length. It's just the number of the batch times the other variables, which is great. And this is especially great for long context. Now, because of the number of layers, the effective context length is actually larger than the number of the stated context length because information can propagate farther in if you go down the layers. Now, you can do fancier things. You can maybe not do a dense selection of layers, but you can space it out. You can also do this global plus sliding window where aware you have attention to a fixed grid of different token points, plus a local sliding window.
[58:25.820] So you can do various things. Now, the problem with this is that it actually still hurts accuracy. So this reduces expressivity. There's no free lunch here, or at least this was an expensive lunch. So the solution that people come up with is that they interleave local attention with global attention. So these hybrid models have full attention for some of the layers and some local layers for some of the other layers. And you're basically always trying to-- and then so that allows you to essentially reduce the KV cache a little bit. And you're always trying to balance accuracy with the speed. Yeah. Hybrid model-wise, is there a difference between the trade-off in using a linear time invariant versus a sliding window? Yeah. So the question is, what about linear attention versus sliding window attention? So I'm not going to talk about linear attention.
[59:28.530] But very quickly, there's a bunch of methods that where, instead of storing KV cache, since you basically compute some of compressed representation of all the history. So linear, the most naive linear attention is you just sum the KV values up into a single vector. So that's definitely independent of the sequence length. You can do fancier things. There's gated net, delta nets, and Mamba which allows you to try to compress but not forget as much. Now the question is, how do those compare? Those have also been used in place of sliding window attention, and people have gotten good results with them. You can also use a combination of full attention, sliding window intention, and the linear attention because they capture different aspects. If you care about local high-resolution stuff, then sliding attentions better. If you just want broad summaries of the past, then the other linear attention might be better.
[01:00:38.460] So then for a long context sentence, would you say your attention would be a better setting for that? So the question is, for long context with linear attention be better? So there's no free lunch, right? Let's say you have a very long context, and you're solving a needle in a haystack problem. If you have to compress your entire history into a small context, you're just going to lose information, and you might just not be able to retrieve it. I guess I feel like what it seems like people go with hybrid architectures that you'll always need some of longer attention, but I'm just trying to understand in my head, what the trade-off between using a sliding window versus some of a Mamba, delta net layer would be. Is using a delta net layer better than using sliding window consistently? Or what is the actual maybe representative trade-offs? Yeah. Maybe I'll say that the Mamba and delta net are more powerful than the sliding window.
[01:01:46.805] Attention maybe you can think about the Mamba as probably can-- it certainly can represent some of the aspects of sliding window attention because you can just-- as you're doing the recurrence, it can just look at the last state. So yeah, maybe you can think about the linear attention or its extensions as being better. They have more room. Once you do sliding window attention, you're done. There's nothing else you can-- [INAUDIBLE] Yeah. OK. So let me just quickly go through this. Just highlight this DeepSeek. DeepSeek continues to innovate different types of attention mechanisms.
[01:02:30.670] So remember, they came up with the multi-latent attention, which compresses the key values. There is this thing called-- now they have compressed sparse attention, DeepSeek sparse attention, and heavily compressed attention. I never remember all these acronyms and what they mean, but let's look at this diagram. So normally, have your KV tokens and your query token. So compressed attention is going to basically compress every M tokens into one token. And then there's this thing called DeepSeek sparse attention, which basically selects a subset of those to keep. And the way you select a subset is that you actually compute some LiDAR weight queries and keys, and then you do a smaller attention
[01:03:24.340] to get these index scores so you know to keep. So a lightning fast way to figure out what tokens you need to keep. And then you use those. And then, there's some more compression that happens. So in the interest of time, I'll move on. The goal of this section is to reduce the KV cache because KV cache is related to memory. And we saw that inference is memory-bound. So that directly translates to improvements in throughput and latency. And the key is to do this without hurting accuracy. So you can do lower dimensional KV caches across layers by across heads and across head dimension. You can do local attention. You can do linear attention, which was discussed earlier, and there's much more. And there's also diffusion models, which is a non-autoregressive way to generate which can be much faster. OK. So let's talk about a few other ideas which are important.
[01:04:33.270] So quantization is more of a, I would say, less of a architecture and much more of a systems of perspective on how to make things smaller. So the key idea here is just reduce the precision of numbers. Less memory means higher latency and throughput. And obviously, you have to worry about accuracy. So quantization is-- there's many options here all the way from bf16 all the way down to int4. And so one thing you can do, if you're scared that quantization is going to mess you up, is that you train a model with quantization in mind. So this is quantization aware training. And during the feedforward pass during training, you quantize and dequantize, and you're basically simulating these quantization errors as you train. So then, generally, now the weights are adapted towards quantization. Things will work better. But the con is that it requires expensive large-scale training. And so what typically people do is they train models, and then they quantize after the fact. So this is post-training quantization. It's much cheaper often. There's a naive way to do it is that for every layer or tensor, you basically determine the scale and the 0 point for each, let's say, tensor.
[01:06:11.390] And then you quantize that separately. This generally doesn't work as well. You can use this idea called GPTQ, which uses some Hessian information to quantize layer by layer. And then you keep track of the errors which get propagated into the non-quantized weights. And so it allows you to correct for the errors. And now, activation-aware quantizing is more of a sophisticated way where the observation is that some activation channels are large and those the weights that interact with those matter more. So let's allocate more precision to these weights. So let's look at this picture. So normally, if you take fp16 weight matrix, and you quantize it, you get-- let's say you quantize it down to int3. But what you're going to do is you figure out which of these are active? So these are maybe activation channels. And some of these activation channels are large. So if they're large in general, then you basically allocate like, let's say, fp16 for this channel. You keep everything else as in three. And for a few important channels, you use higher precision. So another idea here is to do model pruning.
[01:07:45.220] So here you just take a large model, and you rip out pieces of it, and you fix it up. It's a very crude way, but it turns out to work. So there's this paper from NVIDIA where, essentially-- let me actually-- let's see. So you first have to estimate the importance of the different parts of the model and choose the most important parts. And then you basically remove different hidden units and different even layers. And now, we have a model. It's not going to be very good. And so what you do is you post-train it. You train it some more on the data or the tasks that you care about to heal it. So this is, in some sense, a training way to reduce the KV cache, but where you initialize it with parts of a good model. So this seems to work pretty well. So they were able to take a 15B model and reduce it to an 8B model. And it doesn't really hurt accuracy by too much.
[01:09:06.390] And the amount that you use to train the model or through this process is much less. OK. So just to summarize here, the game is to reduce the inference complexity without hurting accuracy. You can think about this as mostly reducing the number of parameters or CPU cache. You can define a faster model architecture to train it, or you can define a faster model architecture. Initialize the weights from original model, which might have a different architecture, but you just make this Frankenstein thing, and then you repair the faster model with distillation. OK. Yeah. Could you say more about you distinguish the important from the unimportant layers? So how do you distinguish the important layers from the unimportant layers? So in general, you have a calibration set, and you pass the inputs through the model. And you're basically looking at the magnitude of the activations.
[01:10:20.180] And the ones that are-- some of them, especially if they're dead units, will be kind of close to 0. And the ones that are large, you want to keep. So that's a high-level idea. I guess why does it matter if the activation is high? It could just be like-- what if it's just always high, for instance? Or is that-- So the question is, what if all the activations are high? So in general, this is a kind of an empirical observation that some of the channels will be much higher than others. If this weren't true, then these techniques wouldn't necessarily work, but it happens to be true because that's how these models ended up being trained. And then you can exploit that. Say a neuron, which is always value 100 across all the samples. Would that mean that it's necessarily meaningful? Or maybe it's just an artifact of training that I'm just like [INAUDIBLE]? Yeah, I see. So the question is, what if a neuron is always 100?
[01:11:24.850] I mean, if that's the case, you can also look at variance-related questions. If it's 100, you can't just remove it because then everything is going to be broken. If it's high mean and low variance, maybe there's another a way to just incorporate the bias essentially. OK. Let me quickly go through this other idea. So far, we've looked at lossy methods which basically really crunch down the KV cache, but it could hurt accuracy. This is a very elegant way of doing this in a lossless way. This is called speculative sampling or speculative decoding. So remember that if you're doing prefill, you can encode all the tokens in parallel. This also gives you probabilities. This is fast. This is compute-bound and all nice things. And in generation, it's one at a time. So checking is faster than generation. If I give you a sequence, it's fast to tell me how good it is, much faster than it is to generate one at a time.
[01:12:26.320] So you can exploit this asymmetry using the following idea. So what we're going to do is use a cheap draft model to basically generate from the guest a few tokens, let's say four tokens. And then we're going to use the model we actually care about, the target model Q, which to basically review these tokens and accept or not accept them. So things are kind of chosen to be balanced. The draft model is just smaller and cheaper. And so even though it's memory-bound, and it has to generate one at a time, is not too bad. Whereas, the target model is big and expensive, but what we're doing is asking it to process a batch of tokens in parallel so it won't also be too bad. OK, so here's the video that shows how things work. So if you use a big model token by token, it's going to be pretty slow. But if you are doing speculative decoding, then you can see the small model generating a bunch of tokens, and then the large model basically critiquing them.
[01:13:42.230] And then, so you basically can get this burst of tokens and then maybe another burst of tokens and so on and so forth. OK. So here is the algorithm for-- there's a few papers that came out with speculative decoding around the same time. This is one of them. And so the idea here is that in order to generate, we're going to generate K tokens. And we're just going to sample from this draft model. So p is the draft model. We're going to sample K tokens. And then in parallel, we're going to compute the logits of these draft tokens using q. And then now we have to determine whether we accept or not. And this is where you do a bit of math and probability and statistics. You basically are going to accept with probability Min 1 over this ratio of q over p. So if q is much larger than p, the larger the q is, the more likely you want to accept it.
[01:14:57.690] And otherwise you sample from this residual distribution and exit. So this is basically rejection sampling except for rejection sampling. Sometimes when you reject you get nothing. But here we always are guaranteed to get an exact sample from the target model. I'm going to skip this simple proof. It's basically the same arguments as rejection sampling to show that it's the exact probabilities from the target model. And the initial paper shows that this fast. And generally, if you have too few draft tokens, you're not really leveraging the batching on the target model side. And if you have too many, then you're going to reject more often. So there's a sweet spot around, in this case, three or four. And in general, the draft model is much smaller than the target model. And ideally, you want the model draft model to be as close to the target, which means that you want to distill it, which means that actually a lot of the same ideas that we just
[01:16:20.670] talked about are applicable to speculative decoding as well. So basically, the idea is that let's try to reduce your KV cache via all the different shenanigans. And if you end up with a model you're happy with, just serve that. If you're not happy with it, then it at least can be a draft model and you can use your main model to fix things up. There's a bunch of whole literature on speculative decoding right now that improve over the original, which I'll skip for now. So very quickly now, dynamic workloads. So this is the use case, is you're serving a live website, and users come and chat with your model. The requests arrive at different times. They have different shared prefixes, and they have different lengths. So it's kind of pretty messy. It's far from this very simple training where you have these blocks of the same number of tokens all at once. So what do you do in this case?
[01:17:30.680] So there's this idea-- there's a system called Orca that was built. This is actually very early on which introduced this idea of continuous batching. So the idea is that you get a bunch of requests that look like this. So here's a prefix of the first request, and you're generating this token. Here's the second one. Here's the third one, the fourth one. It's jagged because every prefix has a different length. And what we're going to do is we're going to decode step by step. So every step you decode one token for all the sequences. Next step, you decode another token for all the sequences. And then if you end, you just eject that sequence. And then as new requests arrive to the batch, then you basically put it in the batch, and then you just continue. So that's why it's called continuous batching because it's this batch is dynamically being updated with either old finished sequences being evicted and new ones coming in.
[01:18:34.770] So now, one problem here is that everything we've seen batching works when all the sequences have the same dimensionality. You have tensors. Every slice has the same dimensionality. But if you request here has a different length, so what do you do about that? So there's this idea called selective batching where let's say you have a length 3, length 9, and a length 5. So here in the attention computation, you can't really do anything about this because attention depends on the length of your sequence. So if you have a 3 by 3, a 9 by 9 computation, you can't really share the tensor effectively. But for the non-attention, the MLP layers, which takes up a lot of FLOPs, you can actually just concatenate all the sequences together to form a mega sequence, and you process that. So final idea is page attention. This was introduced in the BLM paper. Of course, BLM has many other bells and whistles now, but this is the core idea at that time.
[01:19:49.030] So the question is, how is the KV cache stored? So if you think about requests coming in, you have to put them in memory somewhere. And in general, there's this problem that you get fragmentation. This is what happens to-- or used to happen to your hard drive. And you have to defrag your hard drive back in the day. So there's two types of fragmentation. One is that you have to allocate enough buffer so that-- because you don't when you're going to stop. So you might have a max token limit of 1,024. So you have to allocate all this memory. And you can't put anything in there because you're going to just generate until you hit the max tokens, and that's very wasteful. That's internal fragmentation. And then there's also you have-- there could be space between different requests. And that space is maybe too small to get used effectively.
[01:20:50.040] So that's just wasted space. So the solution here is just these are systems people. So they know their operating systems. They said, OK, well we've solved this problem once before, so let's just use the same idea here. So we're going to divide the KV cache of a sequence into non-contiguous blocks, OK? So if you have this sequence, four score and seven years ago our fathers brought forth, we're just going to chunk it up into these blocks of size 4. Doesn't matter where the blocks go, but they're going to be aligned according to the block. So there's some uniformity there. So when two requests share the same-- can actually share the same KV cache. So you might have this block and this block. So this block might go here and here, and this block might be over there. So they're interspersed. But as long as you have the indices and keep track of where everything is, it's fine.
[01:21:58.410] So in particular, if you have, system prompts, then you can actually just cache these system prompts, the KV cache and the system prompts, once. And that can be useful for all the queries, OK? So this is very useful because if a lot of people are using the same system prompt, then you don't have to compute the cache for every request. Also there are many applications where you have the same prompt, and you actually want to generate multiple responses. So in that case, you can also just share the KB cache for the prompt and just have unique responses coming out. OK, So for example, if you were to generate, let's say, multiple generations from score and seven years ago are, blank, and so what would happen here is that you would have score and seven. And you would start by having years ago or our. And then this is called copy on write semantics. You keep this. And then we have these two samples. And if they had happened to sample let's say the same token, you just continue with that.
[01:23:17.010] But if they sample different tokens, you split the block, and then you can continue there. So you're basically sharing as much of the prefix cache as you can. There's a bunch of other optimizations like kernels that I'm not going to have time to go over. But the general idea is that you're using these operating systems metaphors to manage your inference. So summary here-- inference is really, really important. It's very different from training. Even though it's the same model, but you're asking the model to do something very different, ends up being very memory-bound. And it's also dynamic if you're in a live chatbot use case. We saw a variety of different techniques to improve inference. You can quantize. You can come up with new architectures. You can prune and distill. You can also use speculative sampling. But all of these are driven by the same principle here, which is reduce your KV cache but don't hurt accuracy too much.
[01:24:28.540] And then there's ideas from systems like paging and speculative execution that can be brought to bear for actually live inference servers. One thing we didn't really get a chance to talk about, which was discussed briefly, is that I think that new architectures have actually a huge potential for improvement, things like state space models or linear attention or diffusion. At some level, the KV cache and the way that attention is built fundamentally makes it an inference unfriendly kind of architecture. So if you can come up with a new architecture that is designed for inference in the way that the transformer was not, this can maybe unlock a lot. OK, so I will stop there. And next class, Tatsu will return and talk about scaling laws part 2.