CUDA Kernel 面试主线

17 课 GPU 底层算子与性能优化

CUDA Kernel 核心主线

从 GPU 线程映射出发，依次学习 Warp 归约、共享内存 Tiling、Bank Conflict、数值稳定算子（RMSNorm/LayerNorm/Online Softmax）和 Tiled GEMM；之后拆解 DeepSeek 的 DeepGEMM（Hopper TMA + WGMMA 的 FP8 GEMM）与 DeepEP（MoE 专家并行通信），最后覆盖具身与自动驾驶常用算子：Flow Matching 去噪求解器、3D 点云体素化和 QAT 伪量化。

学习目标

  1. 熟练掌握 Thread / Warp / Block / Grid 坐标系统与 Grid-Stride Loop；

  2. 掌握 Warp Shuffle 指令与 Block 级双阶段树状归约；

  3. 掌握 Shared Memory Tiling、Bank Conflict 避让（Padding 技巧）与合并访存（Coalesced Access）；

  4. 掌握 Online Softmax 与 FlashAttention 前置算法的数学推导与数值稳定设计；

  5. 深入理解 Hopper (SM90) 硬件级 TMA、WGMMA 异步流水线与 DeepSeek 细粒度分块量化；

  6. 掌握 MoE 专家并行通信拓扑、SM 分区，以及 Mega-MoE 的计算-通信重叠；

  7. 掌握具身智能 Flow Matching 去噪积分步进、3D 点云体素化与 QAT 伪量化 STE 底层算子实现。

课程目录

课程名称 模块归属
第 1 课：线程映射与 grid-stride loop CUDA Kernel 基础
第 2 课：Warp/Block 归约与两阶段求和 CUDA Kernel 基础
第 3 课：GEMV 与行级并行 线性代数算子
第 4 课：RMSNorm：统计量与两遍访问 归一化算子
第 5 课：LayerNorm：均值、方差与数值稳定 归一化算子
第 6 课：稳定 Softmax 与朴素复杂度 注意力前置算子
第 7 课：Online Softmax 与可结合状态 注意力前置算子
第 8 课：朴素转置与合并访存 访存优化算子
第 9 课：共享内存转置与 bank conflict 访存优化算子
第 10 课：共享内存 Tiled GEMM 核心矩阵乘算子
第 11 课：SiLU 与算子融合 激活融合算子
第 12 课：ArgMax 对归约与确定性 极值归约算子
第 13 课：DeepSeek DeepGEMM 核心架构与 Hopper FP8 算子拆解 工业前沿进阶
第 14 课：DeepSeek DeepEP 核心架构与 MoE 专家并行通信算子拆解 工业前沿进阶
第 15 课：具身智能连续动作生成算子：Diffusion Policy 与 Flow Matching 去噪求解器 空间具身算子进阶
第 16 课：自动驾驶与具身感知：3D 点云体素化（Voxelization）与空间聚合算子 空间具身算子进阶
第 17 课：硬件感知量化：QAT 伪量化（FakeQuant）、STE 与对称/非对称算子实战 空间具身算子进阶

专题演进与架构图稿

FlashAttention v1 IO-aware tiling + online softmax
目标：不物化 N x N 的 S/P 矩阵，用 tile 在 SRAM 内完成局部注意力并在线更新输出。
核心瓶颈与解法
标准 attention：
1. S = Q K^T 写 HBM
2. softmax(S) 写 HBM
3. P V 再读 HBM

v1 改写为：
按 Q tile 和 K/V tile 流式扫描
只在 SRAM/register 中保留局部 S、P
每行维护 m、l、acc，最后 acc/l 得 O

代价：反向传播重算局部 S/P
收益：中间矩阵显存从 O(N^2) 降到 O(Nd)
线程与数据搬运过程
HBM
Q, K, V, O
SMEM/SRAM
Qi, Kj, Vj tile
Registers
m, l, acc
线程视角：CTA 负责一个 query tile；warp 协作加载 K/V tile；Tensor Core 做 QK 和 PV；CUDA core 做 max、exp、sum。
v1 典型 sliced-K：不同 warp 分 K/V 切片，Q 对所有 warp 可见；每个 warp 产生局部 acc 后，需要写 SMEM、同步、再归并。
概念伪代码：一个 CTA 处理一个 Q block
parallel for q_block i:
Qi = load Q[i] to SMEM/register
m = -inf; l = 0; acc = 0

for k_block j in 0..T:
Kj,Vj = load K[j],V[j] to SMEM
S = Qi @ Kj.T / sqrt(d)
if causal: S += mask(i,j)

m_new = max(m, rowmax(S))
alpha = exp(m - m_new)
P = exp(S - m_new)

acc = acc * alpha + P @ Vj
l = l * alpha + rowsum(P)
m = m_new

O[i] = acc / l
store O[i], m, l
Tensor Core
QK 和 PV MMA
v1 的讲解重点：为什么快
少写 HBM
S/P 不落全局内存；HBM 流量是主要收益来源。
online softmax
只要 m 和 l，就能分块得到精确 softmax。
kernel fusion
QK、softmax、PV 融合，减少 kernel 间读写。
重算换内存
反向不存 P，重算局部分数，节省显存。
不足
并行划分还粗：小 batch/少 head 时 SM 吃不满；warp 间通信多。
视频讲解顺序
1. 先画出标准 attention 的 S/P 两个 N x N 中间矩阵，强调不是算不动，而是搬不动。
2. 画 HBM 到 SRAM 的 tile 搬运：K/V 一块一块进来，Q block 固定，输出累加。
3. 推 online softmax：m_new、alpha、l、acc 四个状态变量解释清楚。
4. 最后指出 v1 的短板：CTA/warp 划分还不够好，这就自然过渡到 v2。
读图方法：蓝色表示 HBM/全局内存，绿色表示 SMEM/SRAM，紫色表示寄存器/累加器，橙色表示 Tensor Core 或 CUDA Core 计算。箭头表示数据搬运或流水依赖。
来源：FlashAttention 论文 arXiv:2205.14135；NVIDIA online softmax/tile 说明。

FlashAttention v1 IO-Aware Tiling 原理 · Excalidraw 源文件

FlashAttention v2 better parallelism + work partitioning
目标：数学和 v1 基本相同，把 GPU 利用率从 IO-aware 进一步推向接近 GEMM。
v2 相对 v1 改了什么
v1 已经少搬 HBM，但还不够快：
1. 非 matmul 操作占比高
2. 只按 batch/head 并行，长序列小 batch 时 SM 不满
3. sliced-K 需要 warp 间写 SMEM、同步、归并

v2 的三板斧：
减少 rescale/mask/bound check 等非 matmul
沿 sequence 维度增加 CTA 并行
从 sliced-K 改为 sliced-Q，减少 warp 间通信
v1 sliced-K vs v2 sliced-Q
v1 sliced-K
warp0..3 分别拿 K/V 的不同切片
每个 warp 得到局部 O
写 SMEM -> sync -> reduce
痛点：共享内存读写和同步开销
v2 sliced-Q
warp0..3 分别拿 Q 的不同 row 切片
K/V 对所有 warp 共享
每个 warp 直接得到自己的 O 行
收益：基本不需要 warp 间归并
线程与搬运节奏
warp0: Q rows 0..31
warp1: Q rows 32..63
warp2: Q rows 64..95
warp3: Q rows 96..127
K/V tile 只加载一份到 SMEM，所有 warp 读取；每个 warp 维护自己 Q 行的 m/l/acc，最后直接写对应 O 行。
概念伪代码：更细的 query block 并行
parallel for batch, head, q_block i:
split Qi rows across warps
keep K/V tile visible to all warps
m,l,acc are per Q-row slice

for k_block j:
load Kj,Vj once to SMEM
each warp handles its Q slice:
S_w = Q_w @ Kj.T
apply mask only when needed
online_softmax_update
acc_w += P_w @ Vj

no cross-warp reduction for O slice
store each warp output slice

extra parallelism:
more CTAs cover sequence blocks
improves occupancy when B*H is small
面试讲法：v2 是工程调度升级
一句话：
v2 没有改变 exact attention 的数学，主要是让计算更像高效 GEMM。

三个关键词：
1. fewer non-matmul FLOPs
Tensor Core 做 matmul 很快，FP32 exp/max/mask 相对昂贵。

2. sequence parallelism
长序列、小 batch、少 head 时，不能只靠 B*H 个 CTA。

3. sliced-Q work partitioning
Q 行天然独立，按 Q 切能减少 shared memory 通信。

过渡到 v3：
A100 上 v2 已经不错，但 H100 有 TMA/WGMMA/FP8，新硬件能力没吃满。
读图方法：蓝色表示 HBM/全局内存，绿色表示 SMEM/SRAM，紫色表示寄存器/累加器，橙色表示 Tensor Core 或 CUDA Core 计算。箭头表示数据搬运或流水依赖。
来源：FlashAttention-2 arXiv:2307.08691；Princeton NLP Group FA2 博客。

FlashAttention v2 Work Partitioning 原理 · Excalidraw 源文件

FlashAttention v3：Hopper 上的 TMA + WGMMA + ping-pong
目标：在 H100 上把搬运、QK、softmax、PV 交错起来，让 Tensor Core 尽量不停。
Hopper 上的新瓶颈
Tensor Core 变快后，v2 的同步主循环里，数据搬运和 softmax 开始挡路。
Hopper 提供三样新能力：TMA 异步搬 tile 到 SMEM；WGMMA 以 warpgroup 为单位发 Tensor Core MMA；FP8 吞吐翻倍但要控误差。
v3 的分工
Producer warpgroup：只用 TMA 把下一块 K/V 搬进 SMEM 环形缓冲，搬完发信号。
Consumer warpgroup：等信号，算 S=QKT、online softmax、O+=PV。
两个 consumer 处理不同的 Q 块，轮流占用 Tensor Core。
流水线时间轴（ping-pong 调度）
t0 t1 t2 t3 t4
Producer warpgroup TMA load K/V0 TMA load K/V1 TMA load K/V2 TMA load K/V3 …
Consumer A（Q 块 0） 等待 GEMM S=QK0T softmax0 GEMM PV0 + QK1T softmax1
Consumer B（Q 块 1） 等待 等待 GEMM S=QK0T softmax0 GEMM PV0 + QK1T
绿 TMA 搬运   紫 WGMMA（Tensor Core）   黄 softmax（CUDA core）
softmax 不在 Tensor Core 上执行。一个 warpgroup 做 softmax 时，另一个正好在跑 GEMM，于是 exp/max/sum 的耗时被遮住。每个 warpgroup 内部还会把下一块的 QK 与本块的 softmax 重叠（intra-warpgroup overlap）。
数据路径
HBM（K/V tiles）→ TMA 异步拷贝 → SMEM 环形缓冲（stage 0/1/2）→ WGMMA 读 SMEM → 寄存器累加 O
FP8 路径
吞吐更高，但误差更大。v3 用 block quantization（按块取 scale）和 incoherent processing（先乘随机正交矩阵打散 outlier）保精度。
讲解顺序：先回顾 v2 的同步主循环，再画 producer / consumer，然后用时间轴解释 TMA 预取、WGMMA 与 softmax 交错，最后补 FP8。
来源：FlashAttention-3，arXiv:2407.08608。

FlashAttention v3 Hopper 异步流水 · Excalidraw 源文件

FlashAttention v4 Blackwell pipeline co-design: fully async MMA + TMEM + softmax rethink
目标：B200/GB200 上 Tensor Core 增长更快，瓶颈转向 softmax、SMEM 带宽和流水线调度。
为什么 v4 还要重写
Blackwell 的非对称扩展：
Tensor Core BF16 吞吐大幅提升
但 shared memory 带宽、exp 单元等没有同幅提升

结果：
单纯把 v3 搬过来，matmul 更快了，softmax 和数据路径更显眼。

v4 的方向：
更深的 fully async MMA pipeline
更大的 tile，提高 Tensor Core 连续工作量
TMEM 存中间矩阵，减少 SMEM 往返
软件 exp + 条件 rescale，压低非 matmul 开销
反向用 2-CTA MMA 降低 dQ 原子累加和 SMEM 流量
Blackwell 数据路径：减少 SMEM 压力，把 Tensor Core 喂满
HBM
Q/K/V/O
SMEM
staging
TMEM
MMA accum
Tensor Core
fully async MMA
讲解重点：v4 不只是 tile 更大。Blackwell 上 MMA、load 更异步，TMEM 靠近 Tensor Core 存中间结果；这样避免很多中间量反复进出 SMEM。
软 softmax 的意义：exp 和 rescale 变成瓶颈后，用近似/软件化分担特殊函数单元压力；条件 rescale 则跳过不必要的旧 acc/l 重缩放。
概念伪代码：显式调度 main / softmax / correction
for q_tile i:
init m,l,acc in registers/TMEM
prefetch Q,K,V stages

for k_tile j in pipelined order:
launch async MMA: S_j = Q @ K_j.T
while MMA runs:
schedule softmax for previous S
maybe skip rescale if max change small

softmax stage:
m_new = max(m, rowmax(S_j))
if need_rescale(m, m_new):
acc *= exp_approx(m - m_new)
l *= exp_approx(m - m_new)
P = exp_approx(S_j - m_new)

launch async MMA: acc += P @ V_j
keep intermediates in TMEM when possible

final normalize O = acc / l

backward idea:
2-CTA MMA cooperates on larger tile
reduce SMEM traffic and dQ atomics
线程/CTA 视角
Producer / loader
异步搬 Q/K/V tile，维护 staging buffer。
MMA warpgroup
启动 fully async QK 和 PV MMA，结果优先留在 TMEM。
Softmax warpgroup
处理 rowmax、exp、sum、条件 rescale。
Correction stage
当 max 变化需要校正时，再缩放旧 acc/l。
Backward 2-CTA
两个 CTA 协同更大 MMA tile，减少 dQ atomic add。
从 v3 到 v4 的面试对比句
v3 的关键词是 Hopper 异步化：TMA/WGMMA + warp specialization，解决搬运和计算重叠。
v4 的关键词是 Blackwell 非对称瓶颈：Tensor Core 太快后，softmax、SMEM、原子累加和调度成为主要问题，所以要 TMEM、更深 pipeline、软件 exp、条件 rescale、2-CTA MMA。
读图方法：蓝色表示 HBM/全局内存，绿色表示 SMEM/SRAM，紫色表示寄存器/累加器，橙色表示 Tensor Core 或 CUDA Core 计算。箭头表示数据搬运或流水依赖。
来源：FlashAttention-4 arXiv:2603.05451；PyTorch FlexAttention/FA4 博客。

FlashAttention v4 Blackwell 架构演进 · Excalidraw 源文件

一、先讲清楚它到底解决什么
标准 Attention
S = QK^T 生成 N x N 分数矩阵
P = softmax(S) 再乘 V
问题：S 和 P 写回 HBM，长序列显存和带宽都爆
FlashAttention 的一句话答案
不近似，不改变数学结果
把 K/V 分块搬进 SRAM
边算 QK，边做 online softmax，边累加 PV
避免物化完整 N x N 注意力矩阵
online softmax 关键状态
m：当前块最大值，用于数值稳定
l：exp 和的缩放累加
O：输出累加器
新块到来时重缩放旧 O 和 l，再加新块贡献
面试先避坑
不是 sparse attention，也不是 low-rank 近似
时间复杂度仍是 O(N^2 d)
显存占用从存 N^2 中间矩阵转向 O(Nd)
优势来自 IO-aware kernel fusion 和 tiling
FlashAttention 面试讲解画布：从 v1 到 v4 的优化主线
核心问题：标准 Attention 不缺 FLOPs，先卡在显存 IO，后卡在如何持续喂满新 GPU 的 Tensor Core
二、v1 到 v4 的优化演进
2022
IO-aware tiling
HBM <-> SRAM 少搬数据
2023
更好并行划分
减少非 matmul 和共享内存通信
2024
Hopper 专用
TMA + WGMMA 异步流水
FP8
2026
Blackwell 专用强化
全异步 MMA + TMEM
更激进 softmax
v1：IO-aware exact attention
瓶颈：HBM 读写远慢于 SRAM
核心：Q 分块，K/V 分块，块内计算 S、P、O
不保存 S/P，反向只保存少量统计量再重算
面试句：用更多计算重算换少写显存，精确结果不变
v2：更好的并行度
瓶颈：v1 还没接近 GEMM，occupancy 和 warp 通信不理想
核心：跨 thread block 并行单个 head
warp 内重新分工，减少 shared memory 往返
面试句：算法思想不变，主要优化 work partitioning
v3：Hopper 异步流水
瓶颈：H100 上 Tensor Core 很快，load 和 softmax 会拖后腿
核心：TMA 搬运 + WGMMA 计算异步化
producer-consumer warp specialization
matmul 与 softmax 交错，FP8 低精度并控误差
v4：Blackwell 非对称扩展
瓶颈：B200 Tensor Core 翻倍，但 shared memory 和 exp 单元没同步翻倍
核心：更深 fully async MMA pipeline，大 tile
TMEM 存中间结果，2-CTA MMA 优化反向
软件 exp 与条件 softmax rescale 降低非 matmul 开销
版本主线记忆法
v1：少搬数据 -> v2：多开并行 -> v3：异步重叠 -> v4：面向 Blackwell 重排 pipeline，削掉非矩阵瓶颈
三、面试常考点速答
Q1 为什么显存降了但复杂度没变？
仍要算所有 QK^T 点积，FLOPs 还是二次；省的是不把 S/P 写回 HBM。
Q2 online softmax 怎么保证精确？
每块维护 max 和 sum，遇到更大 max 时把旧累加按 exp(old_m-new_m) 重缩放。
Q3 反向为什么能省内存？
前向只存输出和每行统计量；反向重算局部分数和概率，避免保存完整注意力矩阵。
Q4 v2 相比 v1 本质提升是什么？
不是换数学公式，而是让 block、warp、head 的任务划分更适合 GPU 并行。
Q5 v3 为什么和 Hopper 绑定？
依赖 Hopper 的 TMA、WGMMA、FP8 等能力，把搬运、矩阵乘、softmax 尽量重叠。
Q6 v4 为什么又要改？
Blackwell Tensor Core 更快，但 shared memory 和 exp 不是等比例更快，瓶颈转向 pipeline 和非 matmul。
Q7 什么时候收益最大？
长序列、大 batch 或显存带宽敏感时最明显；很短序列、特殊 mask 或不支持硬件时收益会变小。
四、讲解视频分镜顺序
1
先画标准 Attention：S 和 P 为什么占 N^2
2
解释 GPU 内存层级：HBM 慢，SRAM 快但小
3
v1：tiling + online softmax，精确但不物化矩阵
4
v2：work partitioning，让 GPU occupancy 上去
5
v3：Hopper 异步流水，搬运和计算重叠
6
v4：Blackwell 非对称瓶颈，TMEM 与 softmax 重排
可在视频口播的来源提示
主要依据：
arXiv 2205.14135 FlashAttention v1
arXiv 2307.08691 FlashAttention-2
arXiv 2407.08608 FlashAttention-3
arXiv 2603.05451 FlashAttention-4
NVIDIA 2024 技术博客与 PyTorch 2026 FlexAttention/FA4 博客

FlashAttention v1→v4 面试讲解时间线 · Excalidraw 源文件