LogoCookLLM文档
LogoCookLLM文档
首页CookLLM

原理精讲

词元化
Tokenization 基础BPE 算法详解GPT 系列 TokenizerBPE 训练工程化
模型架构
Transformer LM
从 token ids 到 logitsEmbedding 与 LM Head
Attention 机制
Self-Attention 到 GQAAttention Sink
位置编码
位置编码基础RoPE 数学推导RoPE 代码实现长度外推
GPU 编程基础
GPU 架构基础张量布局Triton 入门:向量加法
FlashAttention
Flash Attention 原理详解从朴素实现到 Auto-TuningBlock Pointer 与多维支持Causal Masking 优化Grouped Query Attention反向传播实现
分布式训练
数据并行ZeRO 优化器全分片数据并行张量并行流水线并行多维混合并行
推理优化
KV CacheContinuous BatchingPagedAttention

动手训练

概述
预训练
预训练数据Tokenizer 训练模型架构数据流水线训练循环监控与验证
X (Twitter)
系统工程FlashAttention

从朴素实现到 Auto-Tuning

会员专享

编写第一个 Flash Attention Kernel,并利用 Auto-Tune 进行性能优化。

在权益中心获取代码

我们在上一章推导出了 Flash Attention 的数学公式(Tiling + Online Softmax)。现在,是时候把数学变成代码了。

核心循环结构

从公式到 Kernel,中间隔着三个必须先想清楚的问题。

第一,两层循环谁并行、谁串行。 上一章的推导写成了 "Outer Q, Inner K":外层遍历 QQQ 的分块,内层遍历 K,VK, VK,V 的分块。但在 Triton 里你不会看到外层循环,因为它被 SPMD 模型隐式并行化了:每个 Program 认领一个 QQQ 块,tl.program_id(0) 就是它的编号。真正写成 for 的只有内循环。这一点如果没建立起来,读 Kernel 时会一直找不到外层循环在哪。

第二,哪些数据在循环中不变。 QQQ 块一旦载入就固定不动,整个内循环都在复用它;K,VK, VK,V 则每轮重新加载。这个区分决定了 tl.load 写在循环外还是循环内,也是 Flash Attention 能把 HBM 访问从 O(N2)O(N^2)O(N2) 降到 O(N2d/M)O(N^2 d / M)O(N2d/M) 的直接原因。

第三,累加器要带几个状态。 朴素 softmax 可以先算完整行再归一化,分块之后不行:每读入一个新的 KKK 块,最大值可能被刷新,此前累加的结果需要按比例修正。所以除了输出累加器 acc,还要维护 running max 和 running sum 两个统计量,共三个跨迭代状态。

理清这三点,下面的 Kernel 就只是把它们逐条翻译成 Triton 语法:

登录以继续阅读

这是一篇付费内容,请登录您的账户以访问完整内容。

Flash Attention 原理详解

通过交互式可视化,深入理解 Flash Attention 的核心技术:内存瓶颈、Online Softmax、与分块矩阵乘法。

Block Pointer 与多维支持

从单序列扩展到 Batch/Head 并行,并使用 Block Pointer 简化指针管理。

目录

核心循环结构
理解 tl.constexpr 的必要性
理解指针运算 (Pointer Arithmetic)
互动指南
验证数值正确性
使用 Auto-Tuning 寻找最佳配置
引入 @triton.autotune
流水线并行 (Pipeline Parallelism)
关键参数解析
总结