GPU Kernel 入门：Triton / FlashAttention

一句话定位：并行训练把活分给一群卡、推理优化与 MLA把显存管省，本篇再下探一层——一张卡内部怎么把单个算子跑快：GPU 的执行模型、Triton 编程语言，以及注意力算子的标准答案 FlashAttention（省的是搬运不是计算）。前置：Attention 与 Transformer——被优化的对象就是那篇的注意力计算；建议先读推理优化，看看显存被 KV Cache 吃得多紧——本篇把镜头从显存容量推进到卡内的搬运。

先看一个反直觉的现象：算力常常不是瓶颈

直觉里 GPU 慢是因为算不动。

真相是：现代 GPU 的算力增速远超显存带宽增速——多数时候计算单元在等数据到位，而不是在计算。衡量一个算子快不快，先要看它是「计算受限」还是「搬运受限」：矩阵乘有大比例的乘加复用，偏计算受限；注意力要生成并搬运一个随序列长度 N 平方增长的中间矩阵（N×N，N 为序列长度），搬运常常占大头。

那「跑快一个算子」的正确姿势是什么？

答案不在算法书里，而在 GPU 的存储结构里——数据放的位置，决定了它被送到计算单元的速度。先看全景：

图：GPU 显存层级与算子优化全景——SRAM 快而小、HBM 大而慢；FlashAttention 靠分块让中间矩阵不落地，搬运省了、结果精确

渐进拆解：从硬件到算法

GPU 执行模型：一群计算单元与一张存储地形图

一块 GPU 上有上百个流式多处理器（SM，计算单元的「车间」），每个车间里线程按 32 个一组（warp）整齐调度。

存储则是一座山：山顶是寄存器（每个计算单元私有的最小号存储），往下一级是片上共享内存（SRAM，整卡合计约 20 MiB、带宽比 HBM 快约一个数量级），山脚是 HBM 显存（几十 GB、容量大但相对片上慢）。写算子就是在这张地形图上搬东西：数据从 HBM 搬上 SRAM、算完搬回去——搬运的趟数，往往比乘加的次数更决定耗时。

判断一个算子偏向哪边的粗口径是「每搬运一个数，能摊上几次乘加」：

算子	每搬一数摊几次乘加	受限侧
大矩阵乘	摊得多（大比例乘加复用）	偏计算受限
注意力	摊得少（还要搬运 N×N 中间矩阵）	搬运受限
逐元素操作	摊得少	搬运受限

这就是「IO 受限」算子的真相，也是本篇一切优化的出发点。

Triton：用「块」的视角写算子

传统路线（CUDA）要开发者亲自管理上千线程的排布，门槛高、易出错；Triton（由 Tillet 等提出，OpenAI 2021 年开源 1.0 版并主导发展）把编程单位升级成「数据块」——而需要指令级极致手工控制（榨干硬件特性的特化内核）时，CUDA 仍占优势，这是两者的分工线。

你写「这个块的数据怎么算」，编译器负责块内的线程排布与共享内存优化——至于块铺到哪些 SM 上，仍由开发者以网格形式指定。一行块级代码背后是编译器生成的底层排布——写算子的心智负担从「指挥每个工人」降到「给每个车间发一张工序卡」。

Triton 因此成为 AI 算子圈的通用方言：PyTorch 官方的编译方案用它生成内核，FlashAttention、各模型的专用算子都有 Triton 版本。内核源码离开 GPU 环境读着抽象，本篇不贴大段源码——Triton 官方教程的融合注意力内核函数主体约百行，值得读的是它的骨架：动手段提炼了一段块级最小例（四件套：程序号、块下标、搬运、写回），完整注意力内核见官方教程 06。

FlashAttention 与在线 softmax：难点与解法

难点在 softmax 是「全行统计」。注意力打分后要过 softmax（见注意力篇），而 softmax 的定义要求知道整行的最大值与总和——可分块计算时每块只看到行的一部分，怎么算出精确的整行 softmax？这是注意力分块化的真正拦路虎。FlashAttention 论文拿来攻克它的「在线 softmax」技巧，出自更早的在线归一化工作（2018，NVIDIA），FA 的贡献是把它与分块整合成 IO 感知的完整方案。

在线 softmax：带着「运行统计量」流式更新。 解法是让每一块计算都携带三个运行统计量——目前为止见过的最大值 m、指数和 l、加权输出累加 acc。每来一个新块：更新最大值，用新旧最大值之差把已有的 l 和 acc 按比例缩放，再并入新块的贡献。

数学上可以证明（本篇动手实算）：这样流式算出的结果与「整行一次算」的 softmax 精确等价——分块不引入任何近似（浮点实现因求和顺序不同会有 1e-16 量级的舍入差，见动手一节）。至于第二处难关「中间矩阵必须落地」：既然块的统计量随身携带，N×N 打分矩阵就在 SRAM 里算完即弃、从不写回 HBM——省搬运与省显存，在这一刻成了同一件事。

图：在线 softmax 分块流程——三个运行统计量流过每一块，N×N 从不落地；精确、省显存、省 IO 三合一

大白话：普通注意力像「把 500 页草稿纸搬上下楼、算完再对答案」，搬得人仰马翻；FlashAttention 是「草稿在桌上（SRAM）写完就撕」——省的全是上下楼的力气，题还是那些题，答案一字不差。Triton 则是不用学指挥千人的口令，写好每张工序卡就行。

技术定义：Triton 是以数据块为编程单位的算子语言——编译器负责块内的线程排布（warp 级）与访存优化，块到 SM 的铺排由开发者以网格形式指定。FlashAttention 是 IO 感知的精确注意力算法：分块（tiling）＋在线 softmax（m/l/acc 三统计量随块更新）保证与全行 softmax 精确等价；N×N 矩阵不写入 HBM，显存由 O(N²) 降为 O(N)，注意力算子提速 2～4 倍（端到端 15% 到 3 倍量级）。

给从业者的对照：FlashAttention 已是注意力算子的默认事实标准——PyTorch 的融合注意力入口 scaled_dot_product_attention（按形状自动分派内核，FlashAttention 系是首选后端）与各推理框架的注意力路径，背后接的基本都是它（或其 v2/v3 演化）；Triton 也是 AMD/Intel 等非 NVIDIA 阵营写算子的主要选择。

什么时候需要自己动手写内核？当你的模型有一个自定义结构（新注意力变体、特殊门控）且它成为吞吐热点时——先用 profiler（性能剖析工具，能量出每步耗时花在哪）确认热点真的在这个算子，再考虑 Triton 重写；绝大多数场景，「组合现成算子」仍是最优解。

本篇与训练并行篇的分界：并行管「卡与卡之间」，内核管「一张卡里面」。

动手：CPU 上验证「在线 softmax 是精确等式」

Triton 内核需要 GPU 环境才能跑（完整可运行的内核源码见 Triton 官方教程 06）；无需 GPU 也能验证的，是在线 softmax 的数学等价性——它是纯逻辑。本节三段：

分块流式 vs 全行一次算，验证在线 softmax 是精确等式。
算一笔 N=4096 的搬运账，感受「省的是 IO」的量级。
Triton 逐元素相加最小例，把「工序卡」心智落成代码。

第一段：分块流式 vs 全行一次算，结果应完全一致：

import numpy as np
rng = np.random.default_rng(7)
scores = rng.normal(size=(4, 16)) * 3        # 4 个查询 × 16 个键的打分

def softmax_row(x):                          # 基准：整行一次算
    e = np.exp(x - x.max(axis=1, keepdims=True))
    return e / e.sum(axis=1, keepdims=True)

full = softmax_row(scores)
acc = np.zeros((4, 16)); m = np.full(4, -np.inf); l = np.zeros(4)
for s in range(0, 16, 4):                    # 在线 softmax：每块 4 个键
    blk = scores[:, s:s+4]
    m_new = np.maximum(m, blk.max(axis=1))
    l = l * np.exp(m - m_new) + np.exp(blk - m_new[:, None]).sum(axis=1)
    acc = acc * np.exp(m - m_new)[:, None]   # 旧贡献按最大值变化缩放
    for j in range(blk.shape[1]):
        acc[:, s+j] = np.exp(blk[:, j] - m_new)
    m = m_new
print(np.abs(full - acc / l[:, None]).max())
# 实算输出：1.11e-16 —— 数学上精确等价、零近似；
# 1e-16 是浮点求和顺序不同带来的舍入差（非逐位一致）

第二段：算一笔搬运账，感受「省的是 IO」的量级（N=4096、bf16 两字节）：

N = 4096
io_naive = 2 * N * N * 2 / 2**20             # N×N 打分矩阵写出+读回
mem = N * N * 2 / 2**20                      # 它占的显存
print(f"中间矩阵往返 {io_naive:.0f} MiB ｜ 显存 {mem:.0f} MiB")
# 实算输出：中间矩阵往返 64 MiB ｜ 显存 32 MiB —— FlashAttention 两项全免

账本读法：第一段的示意代码为了逐位对照保留了完整打分行，真实内核的输出累加只有 N×d 大小、从不存打分行——这正是「N×N 不落地」的实现形态。第一段证明「分块零近似」（省 IO 的前提是不算错），第二段证明「不落地省真金白银」——序列越长两项收益越大（都是 N² 增长）。这个「先证等价、再算搬运」的顺序，就是读任何内核优化的标准姿势。

第三段看 Triton 编程模型本尊——用最小例（逐元素相加，Triton 世界的 hello world；需 GPU 环境运行，结构即官方教程同款骨架）把「工序卡」心智落成代码：

import triton
import triton.language as tl

@triton.jit
def add_kernel(x_ptr, y_ptr, out_ptr, n, BLOCK: tl.constexpr):
    pid = tl.program_id(0)                     # 工序卡号：本程序负责第几个数据块
    offs = pid * BLOCK + tl.arange(0, BLOCK)   # 这张卡名下的元素下标
    mask = offs < n                            # 末块可能不满，越界位裁掉
    x = tl.load(x_ptr + offs, mask=mask)       # 搬运上片（HBM → 片上）
    y = tl.load(y_ptr + offs, mask=mask)
    tl.store(out_ptr + offs, x + y, mask=mask) # 算完写回显存

# out = add_kernel[grid](x, y, out, n, BLOCK=1024)  # grid=(块数,)，如 (triton.cdiv(n, BLOCK),)——发车间

骨架读法：开发者只写「一张工序卡干什么」（本块的下标、搬什么、算什么、写回哪），块内线程排布与片上内存交给编译器——这就是「块为编程单位」的全部含义；tl.load/tl.store 两个搬运动作，正是上文「搬运的趟数决定耗时」发生在代码里的位置。注意力内核（官方教程 06）比它多的是循环与在线 softmax 统计量，骨架不变。

交互工具 · 原理在线 softmax 单步 → 一行 16 分数分 4 块逐块流入，看新最大值到来时旧贡献按 exp(m_old−m_new) 缩放、l/acc 实时变化，末步与全行一次 softmax 对照（差 ~1e-16 精确等式）；附 N×N 搬运账——FlashAttention 省的是搬运不是计算

边界与取舍

什么时候需要自己写内核？

判据是「热点是否真的卡在这个算子」，按顺序走：

先用 profiler 定位热点；热点在数据加载或网络 → 写内核是南辕北辙，到此为止。
热点在注意力、且标准实现满足不了（自定义注意力变体、特殊掩码结构）→ 才值得动手。
用 Triton 重写该算子。

多数场景的最优解仍是组合现成算子——PyTorch 的融合注意力入口（scaled_dot_product_attention）背后就是 FlashAttention 系内核，先查它有没有覆盖你的形状。

FlashAttention 不是万能的

它优化的是注意力的搬运，不减少注意力本身的计算量（乘加次数一次没少）——序列不太长、瓶颈在别处时收益有限；它的分块依赖 SRAM 容量约束，块大小的选择与硬件绑定，跨硬件移植要重新调参（v2/v3 的演化很大程度就是适配新一代卡）。与本篇开头的判断呼应：它是「IO 受限」时代的答案，若未来算力与带宽的剪刀差收敛，下一代的优化目标会换位置。

坏味道预警：把「融合」当万能药

融合算子省的是中间结果的落地搬运，前提是「确有中间结果在搬运」——有人对已经融合好的库函数再手写一遍「融合版」，结果更慢（库版本的块调度与硬件适配更深）。写内核前先读 profiler 的搬运计数，确认「中间结果落地」真的在耗时间里，再动手。

自测

合上文章先自己答，再展开对照：

FlashAttention 快在省显存还是省 IO？为什么这两件事是一件事？

快在省 IO——它一次乘加都没少算，省的是 N×N 中间矩阵在 HBM 上的往返搬运（N=4096、bf16 口径约 64 MiB 往返）。「省显存」与「省 IO」是一件事的两面：在线 softmax 让分块计算随身携带运行统计量（最大值/指数和/累加输出），N×N 打分矩阵在片上 SRAM 内算完即弃、从不写入 HBM——同一件事的两份收益：不落地所以不占显存（O(N²) 降为 O(N)，论文 4.3 节注意力算子显存基准：相对精确注意力基线最高省 20 倍），不落地所以不用搬运（注意力算子提速 2～4 倍）。它的「快」精确成立的前提也在这：分块结果与全行 softmax 精确等价（本篇动手实算差值 1.11e-16，浮点舍入量级），没有任何近似损失。

为什么注意力分块化的拦路虎是 softmax？在线 softmax 怎么解决的？

拦路虎在 softmax 的「全行统计」性质：它的定义需要整行的最大值（防指数溢出）与整行的指数和（分母），而分块计算时每块只看到行的一部分——拿不到这两样，就没法归一化。在线 softmax 的解法是给每个查询行配三个运行统计量：目前最大值 m、运行指数和 l、输出累加 acc；每来一个新块就更新 m，并利用 m 的新旧差把已有的 l 与 acc 按比例缩放（数学上等价于「如果早知道这个最大值，之前的部分和应该是什么」），再并入新块贡献；流到行尾时 acc÷l 就是精确结果。本篇动手段的 CPU 实算显示与全行 softmax 最大差 1.11e-16（浮点舍入量级）——分块零近似。

Triton 与手写 CUDA 的分工是什么？什么时候轮到写内核？

Triton 把编程单位从「线程」升级为「数据块」：开发者写每块数据怎么算，块内的线程排布与访存优化交给编译器，块铺到哪些 SM 由开发者指定——心智负担从指挥每个工人降到给车间发工序卡，因此成为 AI 算子的通用方言（PyTorch 编译方案与 FlashAttention 的 Triton 版都在其上）。CUDA 仍在需要极致手工控制（指令级排布、特殊硬件特性）的场合占优。轮到写内核的判据链：先用 profiler 确认热点真的在某个算子（而非数据加载或网络）→ 该算子是自定义结构、现成库没覆盖 → 再用 Triton 重写，并顺手核对「中间结果是否真的在落地搬运」（融合的收益来源）——多数场景组合现成算子仍是最优解。

参考与延伸

Dao 等《FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness》（NeurIPS 2022）：在线 softmax 与 IO 感知分块的原论文 — arxiv.org/abs/2205.14135
Triton 官方教程 06《Fused Attention》：FlashAttention v2 的 Triton 实现（GPU 环境可跑的完整源码） — triton-lang.org/main/getting-started/tutorials/06-fused-attention.html
华盛顿大学《From Online Softmax to FlashAttention》（课程讲义，2023）：从在线 softmax 到 FlashAttention 的逐步推导 — courses.cs.washington.edu/courses/cse599m/23sp/notes/flashattn.pdf
Dao-AILab《flash-attention》（官方仓库）：CUDA 参考实现 — github.com/dao-ailab/flash-attention
Tillet 等《Triton: An Intermediate Language and Compiler for Tiled Neural Network Computations》（2019）：Triton 的提出论文 — openai.com/index/triton

下一站进入 Harness 工程：模型与算子都已就位，怎么让模型稳定地替人干活——Agent 运行时的设计与纪律，从「什么是 Harness」讲起。

学习舱 · 知识库