斯坦福大学 CME295 课程《Transformer 与大语言模型》2026 秋季学期第 5 讲：

（原标题：Stanford CME295 Transformers & LLMs | Autumn 2026 | Lecture 5 - LLM Systems）

🎯 核心概述

大语言模型（LLM）的规模化部署面临硬件显存容量不足与计算速度受限的双重挑战。Flash Attention 作为一种精确算法，通过融合 Softmax 与矩阵乘法操作，以迭代方式计算注意力机制，旨在减少高带宽内存（HBM）的读写次数并提升单 GPU 效率。解决 LLM 规模化问题需回答两个核心问题：如何让模型适配硬件内存限制，以及如何最大化计算速度。Flash Attention 通过避免存储完整的 $N \times N$ 注意力矩阵，将操作融合在片上高速缓存中执行，从而在不牺牲精度的前提下显著降低内存带宽压力并加速推理与训练过程。

💻 硬件瓶颈与数学原理

• 显存容量限制：以 NVIDIA H200 GPU 为例，其拥有 141GB HBM 显存和 TB/s 级传输速度，但存储百亿参数模型（FP16 格式需 200GB）已超出其容量。此外，HBM 比片上 SRAM 慢约一个数量级，导致传统方法必须将巨大的中间矩阵写入并读回 HBM，造成严重的带宽瓶颈。
• Softmax 计算特性：Softmax 分母需遍历整行向量，这要求传统实现存储完整的注意力矩阵。Flash Attention 利用指数函数性质，通过维护累积分子 $N_t$ 和分母 $D_t$ 实现迭代计算，避免了中间矩阵的显存占用。
• 数值稳定性机制：为防止指数运算溢出，算法动态追踪当前最大值 $m_t$。在每次迭代中，使用 $e^{-(m_{t+1}-m_t)}$ 对历史累积值进行重缩放，确保计算在有限精度下既稳定又精确，实现了自注意力计算的精确结果而非近似。

🧠 Flash Attention 核心机制

• 内存优化策略：算法利用 SRAM 的高速特性处理小块数据，避免在 HBM 中存储巨大的 $N \times N$ Softmax 矩阵。例如，对于长度为 1000 的序列，传统方法需处理百万级条目，而 Flash Attention 通过分块计算显著减少了读写次数。
• 训练加速与重计算：在反向传播阶段，Flash Attention 不存储中间激活值，而是通过重计算（Recomputation）获取梯度。虽然这导致 GFLOPs 从 66.6 增至 75.2，但整体运行时间因内存带宽压力的降低而大幅缩短。
• 精确性保证：通过迭代更新分子与分母并引入重缩放因子，算法避免了数值不稳定问题。这种机制确保了在减少内存访问的同时，计算结果与传统高精度方法保持一致，适用于对精度要求严格的训练场景。

🚀 多 GPU 并行策略

• 数据并行优化：将数据分片至不同设备独立处理。Zero Redundancy Optimizer (ZeRO) 通过分片参数、梯度和优化器状态消除冗余存储，以通信成本换取内存节省，有效解决了单卡显存不足的问题。
• 模型并行技术：
• 张量并行：在矩阵乘法层面拆分计算（如按注意力头或 FFN 层切分），各设备独立执行后合并结果，适用于降低单层计算延迟。
• 流水线并行：按层划分模型至不同设备，需处理前向与反向传播中的“气泡”空闲时间问题，通过微批次调度优化吞吐量。
• 上下文并行：针对长序列场景，将 Query 分组保留在本地，Key/Value 在各设备间流转以完成全局注意力计算，突破单卡序列长度限制。
• 专家并行：在 MoE 模型中，Token 被路由至存放特定专家的 GPU 进行计算后返回，实现了稀疏激活的高效利用。

⚡ 推理延迟与成本分析

• 延迟构成要素：请求处理包含排队时间（TQ）、预填充时间（t-prefill）和自回归生成时间。预填充阶段因输入并行处理较快，而逐词生成阶段耗时较长，是整体延迟的主要来源。
• 关键性能指标：基础设施监控重点为 TTFT（首字延迟）和 TPOT（每输出 token 时间），以优化队列与并行策略；TTLT（末字延迟）受模型推理长度影响较大，非纯系统瓶颈。
• 场景对比分析：短输入长输出场景更慢，因为自回归解码需逐 Token 串行执行；而长输入短输出可并行处理前向传播，仅少量步骤需串行生成。市场数据显示输出价格约为输入价格的 5 倍，反映了生成阶段计算密集的特性。

🖥️ Roofline 模型与硬件权衡

• 算术强度定义：定义为计算量与内存移动量的比值，用于判断系统是受限于内存带宽还是计算吞吐量。高算术强度意味着计算受限，低算术强度意味着内存受限。
• 运行状态分析：解码阶段因需加载全部权重且仅预测单 token，处于内存受限区；预填充阶段若优化得当（如使用 Flash Attention），则处于计算受限区，能更充分利用 GPU 算力。
• 硬件架构取舍：GPU 设计中 SRAM 容量小但速度快，HBM 容量大但带宽低，存在物理尺寸与性能的取舍。理解这一权衡有助于针对不同推理阶段选择最优的算法实现。

💾 显存管理与注意力优化

• PagedAttention 机制：通过块状分配（典型块大小 B=16）和元数据映射解决显存碎片化问题，避免为最大输出长度预留连续内存导致的资源浪费，提高了显存利用率。
• GQA 机制应用：将多头注意力的 Key/Value 投影头数从 H 减少至 G（如 G=2, H=24），可节省约 91% 的 KV 缓存显存，在表达力与效率间取得平衡，显著降低了长序列推理的内存占用。
• MLA 技术压缩：通过低秩分解共享潜在表示，进一步压缩每个 token 存储的向量数量，降低内存占用。这种技术在保持模型性能的同时，大幅减少了 KV 缓存的大小，适用于超长上下文场景。

🚀 推测解码原理与优化

• 加速逻辑机制：利用小模型（Draft）并行生成多个候选 Token，再由目标模型一次性验证，将串行生成转化为并行校验。这种方法有效缓解了自回归解码的串行瓶颈，提升了整体吞吐量。
• 分布一致性保证：通过拒绝采样算法证明，该过程能精确复现目标模型的概率分布，确保输出质量不降。这种数学上的严谨性使得推测解码成为无损加速的重要手段。
• 优化策略实施：采用知识蒸馏对齐草稿与目标模型分布，或使用多 Token 预测头减少验证失败率，从而最大化加速效果并避免无效计算。通过调整草稿模型的长度和接受阈值，可以在延迟与吞吐量之间找到最佳平衡点。