DeepGEMM 完全拆解：300 行 FP8 GEMM 内核，如何逼近 Hopper 的算力峰值

FP8 的诱惑：一倍算力，一半显存#

FP8 是 8 位浮点格式，分为两种子格式：E4M3（4 位指数 + 3 位尾数，最大可表示值 448）和 E5M2（5 位指数 + 2 位尾数，动态范围更大但精度更低）。在《浮点数与数值稳定性》那篇里我们讲过 IEEE 754 的位布局与舍入机制，FP8 只是把同样的结构压缩到 8 位。这里只说结论：在 Hopper 架构上，FP8 的张量核算力是 BF16 的两倍——H800 SXM5 的 FP8 稠密峰值约 1979 TFLOPS（稀疏约 3959），而 BF16 只有约 989 TFLOPS。对推理来说，把权重和激活换成 FP8，等于免费拿到一倍的 GEMM 算力，还能让权重显存占用减半、访存量减半。

所以”能不能用 FP8”从来不是性能问题，而是精度问题。

离群值如何摧毁 per-tensor 量化#

FP8 的精度问题来自它极窄的动态范围。传统做法是 per-tensor 量化：整张矩阵共用一个缩放因子，把矩阵的最大绝对值映射到 FP8 的可表示上限（E4M3 为 448）：

这个公式对分布均匀的矩阵没问题，但 LLM 的激活和权重中总存在少量离群值（outlier）——个别元素比绝大多数元素大几个数量级。举个直观的例子：假设一个 128 元素的块里，127 个元素都在 $[-0.1, 0.1]$ 区间，只有一个元素是 40。per-tensor 方案下 $s = 40/448 \approx 0.089$，其余 127 个正常元素量化后的量级只有 $0.1/0.089 \approx 1.1$——也就是在 E4M3 的可表示范围内只用了最底部的两个刻度。E4M3 在接近 0 的区域（非规格化数附近）精度最差，这些正常值几乎全部被压成相同的粗粒度量化值，误差被离群值”传染”了整个块。

DeepSeek-V3 技术报告（arXiv:2412.19437）提出的解法是细粒度量化（fine-grained quantization）：不再整张矩阵共用一个缩放因子，而是把矩阵切成小块，每块独立计算自己的缩放因子。具体规则是：

• 激活（A）：按 $1 \times 128$ 的 tile 分组——即每个 token、每 128 个通道一组；
• 权重（B）：按 $128 \times 128$ 的 block 分组——即每 128 个输入通道 × 每 128 个输出通道一组。

每组有自己的缩放因子 $s_A$（激活）和 $s_W$（权重），离群值的影响被限制在它所在的小组内，不会污染整张矩阵。

为什么激活和权重的分组形状不一样？论文附录 B.2 记录了实验：把激活也按 $128\times128$ 的方块分组会导致训练不稳定，而按行切分的 $1\times128$（每组的形状是”一个 token 的连续 128 通道”）在实验中是稳定的。一个合理的解释是：权重是静态矩阵，按方块分组没有副作用；而激活是动态变化的，方块分组意味着同一 token 的不同 K 区段用不同缩放，等价于对特征维度做了分段增益调整，干扰了训练动力学。

细粒度量化还让一个此前不现实的选择变得可行：全统一采用 E4M3。E4M3 动态范围小（最大 448，最小非规格化数 $2^{-9}$ 量级），FP8 训练时代的主流做法是前向用 E4M3、反向用 E5M2（范围大但精度低）。DeepSeek-V3 发现，缩放因子按小组切分后，指数位实际上由组内元素”共享”——每组只需要覆盖本组元素的范围，E4M3 的窄动态范围不再致命，于是全部张量统一 E4M3，精度更高、实现更简单。缩放因子在线计算：量化时实时求组内 $\max(|x|)$，不依赖任何历史统计。

这套方案后来被 NVIDIA 官方化——Blackwell 第五代张量核的 MXFP8 微缩放指令集原生支持这种细粒度缩放（只是要求缩放因子是 2 的幂的 UE8M0 格式），而 DeepSeek 在 Hopper 上的实现正是先行者。

细粒度缩放在 GEMM 层面引爆的两个问题#

细粒度量化在算法上很漂亮，但落到 GEMM kernel 上，立刻撞上两道墙——这两道墙定义了 DeepGEMM 的全部设计空间：

1. 缩放必须逐块乘，张量核做不到。反量化时，每个输出元素 $d_{ij}$ 需要乘上它对应激活块和权重块的缩放因子。K 方向每 128 个元素换一组缩放，这个”每 K 块变一次”的缩放无法用张量核的矩阵乘指令表达。DeepSeek-V3 论文原话是：“along the inner dimension of GEMM operations… This functionality is not directly supported in the standard FP8 GEMM”。标准 FP8 GEMM（CUDA 库和 CUTLASS 提供的）只支持 per-tensor 或 per-row/per-column 缩放，per-K-block 缩放需要自己写 kernel。
2. 张量核的 FP32 累加并不精确。DeepSeek-V3 论文实测发现，H800 上 FP8 GEMM 的张量核累加精度”仅保留约 14 位”（“the accumulation precision of FP8 GEMM on NVIDIA H800 GPUs is limited to retaining around 14 bits”），远低于完整 FP32 累加的精度。K 维度动辄 7168（DeepSeek-V3 的 hidden size）甚至 16384，截断误差会在长累加链里不断累积，最终反量化结果偏离真实值。

顺带一提，乘法本身不是误差来源：两个 E4M3 的乘积最多需要 8 位有效数字（4 位 + 4 位），FP32 的 24 位尾数可以精确容纳，所以 FP8 GEMM 的误差几乎全部来自累加环节。这解释了为什么所有 FP8 方案都把”高精度累加”作为头等大事。

DeepGEMM 的全部设计，就是围绕这两条约束展开的。

为什么只做 NT 布局#

因为线性层在推理中就是 $Y = X W^\top$ 的形态：激活 $X$ 是 $(M, K)$ 行主存储，权重 $W$ 在显存里按 $(N, K)$ 列主存（等价于 $W^\top$ 是 $(K, N)$ 行主）。把权重视为”转置的 RHS”直接命中 NT 布局，不需要任何额外搬移。DeepSeek-V3 的 attention 投影、MoE 的 gate/up/down 投影全是这个形态。至于转置、FP8 转换这些杂活，DeepGEMM 的态度很明确：不提供，请在前一个 kernel 里融合掉——“our primary focus is on optimizing the GEMM kernels themselves”（我们只专注优化 GEMM kernel 本身）。它甚至要求激活的缩放因子以TMA 对齐且转置（列主）的布局给出：$1\times128$ 的激活缩放形状是 $(M, K/128)$，每 token 一行；转置成 $(K/128, M)$ 列主后，一个 M 块的一列恰好是 $BLOCK\_M$ 个连续 FP32，TMA 引擎可以一次搬进共享内存。

分组 GEMM 的两种布局：MoE 推理的两个阶段#

分组 GEMM 的两种布局对应 MoE 推理的两个阶段，值得展开：

• 连续布局（contiguous）：MoE 里每个专家处理不同数量的 token。prefill 阶段 token 数量已知，把去往同一个专家的 token 在 M 维上拼接成一个连续段，整个段作为一个大 GEMM 计算。与 CUTLASS 的传统分组 GEMM 不同，DeepGEMM 的组只沿 M 轴切分，N 和 K 在所有组间固定——正好对应同层专家形状相同这一事实。每个专家段的 M 需要按 GEMM 的 M 块大小对齐（有专门的 get_mk_alignment_for_contiguous_layout() 查询接口）。
• 掩码布局（masked）：decode 阶段开启 CUDA Graph 后，CPU 不再知道每个专家将收到多少 token（token 数是运行期动态决定的）。这时显存里的 token 排布是”碎片化”的，无法拼成连续段。DeepGEMM 用一个 mask 张量描述每个位置是否有效，kernel 只计算有效部分。DeepEP（DeepSeek 的 MoE 通信库）的低延迟 kernel 输出正是这种布局，二者在 V3/R1 的解码链路里是配对的——DeepEP 做 all-to-all 通信，DeepGEMM 吃它的输出做专家计算。

全 JIT：安装零编译#

另一个贯穿全局的设计是全 JIT（Just-In-Time）：安装时零编译，所有 kernel 在运行时由内置的轻量 JIT 模块用 NVCC 现场编译。形状、块大小、流水线级数全部作为模板参数写死在编译期。这一点后面单独开一节讲，它是 DeepGEMM 性能的另一半来源。

线程组织：TMA 线程组与数学 warpgroup#

• kernel 以 持久化（persistent） 方式启动：gridDim = num_sms，每个 SM 恰好一个 CTA，CTA 内部用一个块调度器循环领取 (M 块, N 块) 任务，而不是每个任务启动一次 kernel——省掉了任务间反复启动、反复建立流水线的开销。
• 线程分成两类。一类是 TMA 线程组：固定 128 个线程（4 个 warp），只负责发 TMA 拷贝指令和等待屏障。为什么是恰好 4 个 warp？源码注释给出了直接原因：寄存器重新配置指令 setmaxnreg.aligned 要求以整个 warpgroup（4 个 warp）为单位操作。TMA 线程组通过 warpgroup_reg_dealloc<40> 把寄存器让渡到 40 个——它们只跑地址计算和屏障，40 个寄存器绰绰有余，省下的寄存器空间留给数学线程。
• 另一类是数学 warpgroup：128 个线程一组，通过 warpgroup_reg_alloc<232> 独占 232 个寄存器，专职执行 WGMMA 和提升的 FFMA。当 BLOCK_M = 128 时数学部分有 2 个 warpgroup（256 线程），当 BLOCK_M = 64 时只有 1 个。232 个寄存器是 Hopper 上单线程的寄存器上限（每个 SM 64K 寄存器 / 256 线程），数学线程把几乎全部寄存器都用在了累加器阵列上。
• 计算块参数：BLOCK_K 固定为 128——源码里有一句静态断言 “Only support per-128-channel FP8 scaling”，因为缩放粒度就是 128，K 块必须与之对齐；BLOCK_M 取 64 或 128；BLOCK_N 在 16 到 128 之间以 8 为步长搜索（选参规则见 JIT 一节）。每个 K 块由 4 条 WGMMA 指令完成（Hopper 的 FP8 WGMMA 是 $m64 \times n \times k32$ 形状，$128/32=4$）。

共享内存布局：每一字节都有去处#

共享内存的分配完全由编译期常量决定（Hopper 每 SM 227 KB，即 232448 字节）：

• 输出缓冲 smem_d：$BLOCK\_M \times BLOCK\_N \times 2$ 字节（BF16 结果暂存，128 字节 swizzle 对齐）；
• A 数据每流水线级：$BLOCK\_M \times 128$ 字节（E4M3，一字节一个元素）；
• A 缩放每级只有 $BLOCK\_M \times 4$ 字节——因为 $1\times128$ 缩放每行每 K 块只有一个 FP32，128 行的 A 块每级只需 512 字节，是 A 数据本身的 1/128。这正是细粒度缩放”压缩”带来的红利：缩放矩阵小到可以全程住在共享内存里；
• B 数据每级：$BLOCK\_N \times 128$ 字节；
• B 缩放整块常驻：$\lceil K/128 \rceil \times 4$ 字节（B 的 $128\times128$ 缩放沿 K 方向每块一个 FP32；BLOCK_N 不能整除 128 时翻倍）；
• 两级屏障数组：$2 \times stages \times 8$ 字节（每级一个 full barrier + 一个 empty barrier）。

A 和 B 的数据走 128B swizzle（TMA 的典型共享内存布局，配合银行冲突消除），缩放和屏障用朴素布局。

生产者-消费者同步：TMA 与数学线程的握手#

两级流水线的同步用 CUDA 的 named barrier 实现：每级流水有一个 full_barrier（数据就绪）和一个 empty_barrier（缓冲区空闲）。握手协议是：

1. TMA 线程组等待 empty_barrier[s]（第 s 级缓冲区已空），然后发出 3 条 TMA 拷贝（A、A 缩放、B），每条 TMA 的 arrive_and_expect_tx 告诉屏障”我承诺搬运的字节数”；
2. 数学 warpgroup 等待 full_barrier[s]（TMA 搬完，屏障自动计数归零），执行 WGMMA 和提升；
3. 数学线程 arrive 到 empty_barrier[s]，通知 TMA 线程组”第 s 级用完了，可以覆盖”。

multicast 开启时（cluster 模式），屏障还带分布式语义：full_barrier 的初始计数按簇大小调整，数学线程用 arrive(lane_idx) 按簇内编号抵达，加上 fence_view_async_shared 和 fence_barrier_init 保证异步代理（TMA）与正常代理（数学线程）之间的可见性。这些细节是 warp 特化 kernel 最容易写错的地方，DeepGEMM 把它们收敛在 fp8_gemm.cuh 一个文件里。

主循环骨架#

数学 warpgroup 的主循环可以浓缩成下面的伪代码（与真实源码一一对应）：

1
// 数学 warpgroup 的主循环（简化）
2
while (scheduler.get_next_block(m_block_idx, n_block_idx)) {
3
    // 预取本块用到的 B 缩放因子到 smem（与上一个任务的 TMA store 重叠）
4
    load_scales_b(smem_scales_b, n_block_idx);
5

6
    for (int t = 0; t < K / (kNumStages * BLOCK_K); ++t) {      // K 迭代
7
        for (int s = 0; s < kNumStages; ++s) {                   // 流水线阶段，全展开
8
            full_barriers[s]->wait(...);                         // 等 TMA 把第 s 级数据搬好
9
            float scale_a_0 = ld_shared(smem_scales_a[s] + r_0); // 读 A 缩放（每行一个）
10
            float scale_a_1 = ld_shared(smem_scales_a[s] + r_1);
11
            float scale_b   = ld_shared(smem_scales_b + t * kNumStages + s);
12

13
            wgmma(smem_a[s], smem_b[s], accum);                  // 张量核：累加当前 K 块
14
            wgmma_wait<0>();
15

16
            // CUDA 核提升：乘缩放，加进最终累加器（每 K 块一次）
17
            float scale_0 = scale_a_0 * scale_b;
18
            float scale_1 = scale_a_1 * scale_b;
19
            #pragma unroll
20
            for (int i = 0; i < kNumAccum / 4; ++i) {
21
                final_accum[i*4+0] += scale_0 * accum[i*4+0];
22
                final_accum[i*4+1] += scale_0 * accum[i*4+1];
23
                final_accum[i*4+2] += scale_1 * accum[i*4+2];
24
                final_accum[i*4+3] += scale_1 * accum[i*4+3];
25
            }
26
            empty_barriers[s]->arrive();                         // 通知 TMA 线程组：本级已用完
27
        }
28
    }
29
    // 结果写回：FP32 → BF16 转换后经 stmatrix 写入 smem，再由 TMA store 落回全局内存
30
    stmatrix(smem_d, final_accum);
31
    tma_store_d(smem_d, m_block_idx, n_block_idx);
32
}

几个容易忽略、但决定性能的细节：

• 提升与 MMA 的重叠。提升跑在 CUDA 核（FFMA 单元）上，WGMMA 跑在张量核上，二者是不同执行单元，天然可以并行。warp 特化的意义就在这里：当数学线程忙于 FFMA 提升时，张量核的下一条 WGMMA 已经排入队列；当张量核在算时，FFMA 流水线也没闲着。三级流水线互相隐藏延迟，是这套 kernel 逼近峰值的结构基础。
• 非对齐 BLOCK_N 的双缩放问题。当 BLOCK_N 能整除 128 时（128、64、32、16），一个 B 块恰好落在权重的一个 $128\times128$ 缩放方块内，每个 K 块只需要一个 $s_W$（源码里的 kMustUseUniformedScaleB 常量）。但当 BLOCK_N 取 112、96、80 这类非整除值时，一个 B 块会横跨两个权重缩放方块，每个 K 块需要两个 $s_W$，提升时用 predicate 在两组缩放间选择——代码注释特别强调”用 predicate 而不是两个循环，对性能非常重要”（predicate 不引入分支跳转，GPU 的指令发射流水线不受惩罚）。
• 非整除 K 的处理。K 不一定能整除 stages × 128，最后一轮 K 迭代可能只剩不足 stages 个有效级。launch_k_iterations 用 DivisibleK/NotDivisibleK 两个编译期分支区分：无效级不搬数据，只做 full_barrier->arrive() 空抵达，保证屏障计数正确。这个分支也是编译期常量，不影响展开质量。
• stmatrix + TMA store 的两段式写回。累加器在寄存器里，不能直接交给 TMA。先由 stmatrix 指令（warp 级矩阵存储，一次写 32×2 个 BF16）把最终累加器从寄存器写进 smem_d，然后由单个线程发 TMA store 把 smem_d 写回全局内存。TMA store 是异步的，tma_store_wait<0> 保证在 SM 复用 smem_d 前完成。输出走了”寄存器→smem→全局”两跳，但两跳都是专用硬件路径，比逐线程全局写回（还要处理合并访问）快得多。

TMA 的全方位使用#

Hopper 引入的 TMA（Tensor Memory Accelerator）是异步数据搬移引擎：CPU 侧配好 CUtensorMap 描述符后，GPU 上一条指令就能把任意形状的 2D tile 从全局内存搬进共享内存，全程不占寄存器、不占计算线程。DeepGEMM 把 TMA 用到了极致：

• TMA load：A、B、A 的缩放因子全部走 TMA 进 smem。注意缩放因子也走 TMA——因为缩放矩阵转置后每列连续，天然满足 TMA 的 2D tile 加载条件（A 缩放的 tile 形状是 $BLOCK\_M \times 1$）；
• TMA store：输出 D 由 TMA 直接写回全局内存（省掉常规的寄存器→全局的写回路径）；
• TMA multicast：仅在 A 上开启。A 按 M 行切块，同一行 M 块会被 N 方向相邻的多个 CTA 共享；开启 cluster（线程块簇）后，一份 A tile 经 NVLink 广播到簇内所有 CTA 的共享内存，每个 CTA 省一次全局内存读取。B 没有这个待遇——B 按 N 切块，不同 CTA 之间不共享。multicast 的启用条件比较苛刻：$M \ge 1024$、N 能被 BLOCK_N × 2 整除、SM 数能被 2 整除，cluster 尺寸取 2（源码注释明确警告：cluster ≥ 4 会造成性能回退，因为簇间同步和分布式屏障的开销超过了广播的收益）；
• TMA 描述符预取：kernel 启动第一件事就是 prefetch_tma_descriptor，把四个描述符（A、B、scales_a、D）提前载入缓存，避免首次 TMA 指令的冷启动延迟。

统一块调度器与 rasterization#

调度器（scheduler.cuh）是另一个被反复打磨的部件。CTA 内唯一的调度线程按 next_block_idx = (++iter) * gridDim.x + blockIdx.x 领取全局块编号——这是持久化 kernel 的标准做法，保证所有 CTA 以”波浪”方式扫过整个任务。真正有讲究的是块的遍历顺序：不是简单的 M 主序或 N 主序，而是栅格化（rasterization）——每 16 个 N 块组成一个小组，小组内 N 快扫，然后跳到下一个 M 块。这样安排后，同一时刻处于同一 N 带内的 CTA 会共享 A 的 L2 缓存行，显著提高 L2 命中率（CUTLASS 的 threadblock rasterization 文档讲的是同一件事，DeepGEMM 把它内联进了自己的调度器，少一层抽象）。

对掩码布局，调度器额外维护一个组间累计和（cumsum）：get_next_block 时先跳过没有有效任务的组，只给有 token 的专家分配 M 块，且每个 M 块只覆盖实际有效的范围——解码时专家间负载严重不均衡（有的专家收 3 个 token，有的收 300 个），这个”动态跳过”是必须的。调度器在 current_iter 上做增量推进，还保证了 CUDA Graph 重放时行为完全确定。

非对齐块大小：把最后 20 个 SM 用起来#

这是 DeepGEMM 少有的”为了极端利用率而做”的设计，也是 README 里用 🐳 标出的独有优化之一。常规 kernel 的 BLOCK_N 必须是 2 的幂（128、64……），但有些形状用 2 的幂会浪费 SM。README 举的例子：$M=256, N=7168$ 时，BLOCK_M=128、BLOCK_N=128 的常规分块产生 $(256/128) \times (7168/128) = 112$ 个块，而 H800 有 132 个 SM——20 个 SM 完全闲着。把 BLOCK_N 放宽到 112 这种非对齐值后，块数变成 $(256/128) \times (7168/112) = 128$，全部 SM 都在干活。

这个优化和细粒度缩放叠加是很拧巴的：非对齐 BLOCK_N 会让 B 块横跨两个缩放方块（前面讲的双缩放问题），但 DeepGEMM 选择用 predicate 和第二条缩放行把它扛下来，换取 10% 量级的 SM 利用率提升。选参代码里 BLOCK_N 的候选集是 range(16, 129, 8)，即 16、24、32、…、128——故意包含非 2 的幂，正是为了这种形状。

控制发射的 SM 数：省 L2、保频率#

调度器还有一处反直觉的设计：num_min_sms——即使任务能占满所有 SM，也故意少用一些。源码的注释写得很直白：“less L2 cache usage and less GPU frequency drop”（减少 L2 占用、避免 GPU 降频）。GPU 的 boost 频率取决于功耗与温度，SM 全开 + 高访存强度会把整卡推向功耗墙导致降频，降频对全 kernel 的伤害远大于闲置几个 SM；少用几个 SM、把任务波浪数控制住，反而能稳定在高频。具体规则是 num_min_sms = ceil(max(波浪所需SM数, num_sms - 8) / multicast) * multicast——留 8 个 SM 做功耗缓冲，同时保证 multicast 的簇对齐。这个细节在 README 里没有，是读源码才能发现的”隐藏调参”。

FFMA SASS 交错：改二进制比特的”歪门邪道”#

这是 DeepGEMM 最出名的野路子，值得单独讲讲。作者对比 NVCC 12.2 与 12.3 编译出的 CUTLASS FP8 kernel，发现 12.3 的性能提升来自 SASS 里一串 FADD 指令中某个比特被翻转——翻的是 yield 位，效果是让当前 warp 让出执行槽，给其他 warp 更多发射机会（作者自嘲说”just a guess”，这只是他们的推测）。顺着这个线索，他们写了一个脚本（interleave_ffma.py）在编译产物上直接改 FFMA 指令的二进制：除了翻 yield 位，还把 reuse 位也翻转——reuse 位表示操作数寄存器可以复用，warp 让出执行槽后复用语义失效，必须一并清掉，否则结果不正确。这套 SASS 修补让细粒度缩放的 FP8 GEMM 在某些形状上再提升 10%+，因为它创造了更多让提升 FFMA 与 WGMMA 重叠的机会。

这条技巧的谱系可以追溯到 2015 年 Scott Gray 的 MaxAs 项目——一个为 Maxwell 架构手写 SASS 汇编器的开源项目，在官方工具链之外直接操作指令二进制。DeepGEMM 只是把同样的思路用到了 FP8 时代：当编译器调度不完美时，直接改产物。随着 NVCC 12.9 开始自动做 FFMA 交错，这个补丁在 2025 年 7 月的重构中被移除——编译器追上了手工作业，这也解释了为什么 DeepGEMM 强烈建议使用 CUDA 12.8+：性能不仅取决于 kernel 代码，还取决于编译器版本。