Basics
Tokenizer → Transformer → Training
CS336 · 第 19 课
规模化 ML 系统、硬件感知算法与工业实践。Dan Fu 分享从 FlashAttention 到构建高效 ML 系统的经验，以及对课程的整体回顾。
来源状态：Spring 2026 官方日程确认本讲嘉宾为 Dan Fu，但当前官方仓库未提供可核对的 handout 或 trace。
以下内容仅作为主题背景导读，不视为官方客座讲义的逐段重建。
本讲问题：课程最后用 FlashAttention 和 GPU kernel 工程回看一条共同主线：系统优化不是记住技巧，而是先找到数据移动瓶颈，再共同设计算法、实现和测量。
workload 与 bottleneck → IO-aware algorithm → kernel mapping → benchmark / profiler → 规模化系统约束 → 把同一分析方法迁移到新硬件与新模型。
Spring 2026 官方日程只确认 Dan Fu guest lecture；无本地 handout 的细节均应视作基于公开研究的补充框架。本课不是收集“加速技巧”，而是练习闭环：定义 workload，列 FLOPs/bytes，定位 bottleneck，改变 dataflow，再验证端到端收益。
model workload + hardware hierarchy
↓ resource accounting
hypothesized bottleneck
↓ algorithm / layout / kernel change
microbenchmark + profiler
↓
end-to-end quality / latency / throughput / cost
└──────── new bottleneck ────────┘朴素 attention 物化 \(T×T\) score/probability matrices 到 HBM；FlashAttention 按 tiles 流式维护 softmax statistics，让 Q/K/V blocks 在片上复用，避免写回完整中间矩阵。输出仍是 exact attention（在浮点舍入允许范围内），所以收益来自减少 HBM IO，不是把 \(O(T^2)\) 数学 interactions 变成 \(O(T)\)。
若该 kernel 只占 step 的 20%，即使自身变成 2×，Amdahl 上总 step 也只从 \(1\) 降到 \(0.9\)，约 1.11×；随后通信或输入管线可能成为新瓶颈。算法-硬件协同不是针对某张 GPU 背参数，而是让对象、数据路径与测量保持一致。
Dan Fu 是 UC San Diego 计算机科学与工程系助理教授，SandyResearch 实验室负责人，同时担任 Together AI 的 VP of Kernels。他在 Stanford 获得博士学位，师从 Percy Liang 和 Christopher Manning。
| 项目 | 贡献 | 影响 |
|---|---|---|
| FlashAttention | IO-aware exact attention | 成为 LLM 训练/推理标准 |
| FlashFFTConv | 高效长序列卷积 | 长上下文建模 |
| ThunderKittens | GPU kernel DSL | 简化 kernel 开发 |
| Monarch Mixer | 硬件高效架构 | 替代 Transformer 的探索 |
核心主题：算法和硬件不应分开设计。最好的 ML 系统是那些从硬件特性出发设计的算法，以及为算法优化的硬件。
标准 Attention 的 O(n²) 内存是长序列的瓶颈：
标准 Attention:
1. 计算 S = Q @ K^T [n × n]
2. 计算 P = softmax(S) [n × n]
3. 计算 O = P @ V [n × d]
问题: 需要存储 n × n 的中间矩阵 S 和 P
n = 100K 时, n² = 10^10 个元素 → 几十 GB 内存
不存储完整的 attention matrix，而是分块计算：
FlashAttention 的关键洞察:
1. softmax 可以分块计算 (online softmax)
2. 不需要存储完整的 n × n attention matrix
3. 只需要 O(n) 的额外内存
算法:
1. 将 Q, K, V 分成小块 (blocks)
2. 对每个 Q 块:
a. 加载 Q 块到 SRAM
b. 遍历所有 K, V 块:
- 加载 K, V 块到 SRAM
- 计算局部 attention
- 用 online softmax 合并结果
c. 写回最终输出
结果:
- 内存: O(n²) → O(n)
- 速度: 2-4× 加速（减少 HBM 访问）
- 精确: 与标准 attention 完全一致（无近似）
关键在于利用 GPU 的内存层次结构：
| 内存层级 | 大小 | 带宽 | 用途 |
|---|---|---|---|
| HBM (High Bandwidth Memory) | 40-80 GB | 3 TB/s | 存储模型参数、activations |
| SRAM | ~200 KB/SM | 19 TB/s | 计算时的临时存储 |
| Registers | ~256 KB/SM | ~100 TB/s | 当前计算的数据 |
标准 attention 的问题：n×n 矩阵存在 HBM 中，每次计算都需要从 HBM 读取。FlashAttention 让数据在 SRAM 中完成计算，大幅减少 HBM 访问。
关键洞察：FlashAttention 没有改变 attention 的数学定义，只是改变了计算顺序。这是"硬件感知算法"的典范——同样的数学，不同的实现，巨大的性能差异。
Dan Fu 的研究强调算法和硬件不应分开考虑：
传统方式:
1. 设计算法（不考虑硬件）
2. 实现算法
3. 优化实现（如果时间允许）
硬件感知方式:
1. 理解硬件特性（内存层次、并行度、带宽）
2. 从硬件出发设计算法
3. 算法和实现同时优化
例: FlashAttention
- 从 GPU 内存层次出发
- 设计分块策略匹配 SRAM 大小
- 算法天然高效
理解 kernel 性能的基本工具：
Roofline 模型:
性能 = min(峰值算力, 带宽 × 算术强度)
算术强度 = FLOPs / bytes transferred
例: H100
- 峰值 FP16 算力: 989 TFLOPS
- HBM 带宽: 3.35 TB/s
- Ridge point: 989e12 / 3.35e12 = 295 FLOPs/byte
如果算术强度 < 295 → memory-bound
如果算术强度 > 295 → compute-bound
FlashAttention 的目标:
减少 bytes transferred（不改变 FLOPs）
让 kernel 更接近 compute-bound
大规模模型训练需要多 GPU 协作：
| 并行维度 | 切分什么 | 通信 |
|---|---|---|
| Data Parallel | 数据 | 梯度 AllReduce |
| Tensor Parallel | 权重矩阵 | activation AllReduce |
| Pipeline Parallel | 层 | activation 点对点 |
| Sequence Parallel | 序列 | activation AllReduce |
| Expert Parallel | MoE experts | token routing |
同样的算法，不同的 kernel 实现可以有数倍的性能差异：
例: 矩阵乘法
朴素实现: 10 GFLOPS
优化实现: 800+ GFLOPS (接近峰值)
优化技术:
1. Tiling: 分块匹配内存层次
2. Vectorization: 使用 SIMD 指令
3. Shared memory: 减少 global memory 访问
4. Register blocking: 最大化寄存器使用
5. Software pipelining: 隐藏延迟
6. Warp-level primitives: 利用 warp 内通信
Dan Fu 参与开发的 GPU kernel DSL：
ThunderKittens 风格 (伪代码):
用数据流图描述 kernel
编译器自动处理:
- 内存分配
- 数据搬运
- 同步
开发者关注: 计算逻辑
编译器关注: 性能优化
工业界的 ML 系统研究与学术界的区别：
| 维度 | 学术 | 工业 |
|---|---|---|
| 目标 | 新算法/新发现 | 可靠、可扩展、成本效益 |
| 规模 | 单节点/小集群 | 数千 GPU 集群 |
| 容错 | 可以失败 | 必须高可靠 |
| 迭代速度 | 月/周 | 天/小时 |
| 指标 | 精度 | 精度 + 吞吐量 + 成本 |
训练之后，推理是持续的成本：
Tokenizer → Transformer → Training
Kernels → Parallelism → Inference
小规模实验 → 预测大规模配置
来源 → 过滤 → 去重
以及：Alignment（SFT → RLHF → DPO → GRPO）
The Bitter Lesson（苦涩教训）：历史上，利用计算量的通用方法最终胜出。但这不意味着算法不重要——而是能随规模扩展的算法才是关键。效率在大规模时更加重要——你承受不起浪费。
公式：accuracy = efficiency × resources
不要只读论文。从第一行代码开始构建一个语言模型。理解每个组件的细节。
效率不是"锦上添花"，而是核心。在大规模时，效率差异决定项目是否可行。
现代 ML 离不开硬件。理解 GPU 内存层次、带宽、并行度，才能设计高效算法。
在数据上的投入往往比在模型架构上的投入回报更高。花时间理解和改进数据。
这个领域发展极快。保持学习、保持质疑、保持开放。
问题 1：FlashAttention 的核心贡献是什么？
问题 2：Roofline 模型的含义是什么？
问题 3：硬件感知算法的核心思想是什么？
问题 4：CS336 的核心公式 accuracy = efficiency × resources 的含义是什么？
某 kernel 占 step 20%，自身加速 2×，其余不变。端到端约加速多少？
答案：新时间 \(0.8+0.2/2=0.9\)，speedup \(1/0.9≈1.11×\)。
若答 2×，你把局部 benchmark improvement 直接当成了整条 critical path。
面对一个新模型或新硬件，不先问“用了什么名词”，而按顺序问：对象是什么、shape 是什么、数据依赖是什么、算多少 FLOPs、搬多少 bytes、跨哪些 links、优化目标怎样被测量。能独立重建这条链，才真正完成 CS336。
硬件感知算法的典范，IO-aware 精确 attention
算法-硬件协同设计，Roofline 模型
可靠性、可扩展性、成本效益
理解源于构建，效率驱动一切