GEMM、cuBLAS/CUTLASS 与算子库

一个 transformer 前向的 FLOPs 里，矩阵乘占 90% 以上（QKV/输出投影、FFN 的两个大矩阵乘）。矩阵乘的实现质量就是训练与推理速度的第一决定因素。本章讲清楚 GEMM 为什么值得用一整章、cuBLAS 与 CUTLASS 的分工、以及"为什么我的 GEMM 慢"的排查方法。

GEMM 的显存账：为什么它特殊

$C = AB$ ， $M \times K$ 乘 $K \times N$ ：

FLOPs： $2MNK$ ——随尺寸立方增长
读写： $2(MK + KN + MN)$ ——只随尺寸平方增长
算术强度： $I = \frac{MNK}{MK + KN + MN}$ ——随尺寸线性增长

这个特性让 GEMM 在所有算子里独树一帜（见算术强度）：

尺寸	算术强度	H100 上归类
4096³	~1365	计算受限，能打满 Tensor Core
1024 级	~340	刚过平衡点（295）
256 级	~85	访存受限，利用率上限 29%
64 级	~21	严重访存受限

GEMM 的速度由尺寸决定，不由实现决定——小矩阵乘写得再好也打不满算力，这是数学决定的。所以"把小 GEMM 合并成大 GEMM"（batch 合并、算子融合、更大 batch）比"换更好的 GEMM 库"收益大得多。

cuBLAS 与 CUTLASS：分工

cuBLAS：NVIDIA 的闭源 GEMM 库，黑盒、稳定、快。日常调用的 torch.matmul/nn.Linear 底层就是它。

CUTLASS：NVIDIA 开源的 GEMM 模板库，C++ 模板，给你积木（tile 调度、Tensor Core 抽象、epilogue 融合）自己拼：

	cuBLAS	CUTLASS
形态	闭源库，函数调用	开源模板，代码生成
适用	标准 GEMM，直接用	自定义 GEMM（融合 epilogue、特殊 shape、FP8 特殊路径）
控制力	无（黑盒）	全部（tile、stage、调度策略）
成本	零	高（模板元编程陡峭）

选型：标准 GEMM 用 cuBLAS（PyTorch 默认），自定义融合（GEMM + 偏置 + 激活 + 量化）用 CUTLASS。cuBLASLt 提供 GEMM+融合的受限接口，覆盖常见组合。

GEMM 内核在做什么

一个高性能 GEMM 的完整结构（理解它，才能理解为什么慢）：

1. 分块：M×N×K 切成 tile（如 128×256×64），分给各 SM
2. 流水：TMA 异步搬运下一块 tile ↔ wgmma 计算当前块（多级流水）
3. 喂数据：共享内存 tile → Tensor Core 的固定指令形状
4. epilogue：输出 tile 的 epilogue 融合（bias、激活、量化）

每个环节都是可调参数（CUTLASS 里显式暴露，cuBLAS 里自动选择）。GEMM 的调优空间在 tile/stage/调度策略上能差 20–40%——这就是 autotune 存在的原因（torch.compile 的 max-autotune、CUTLASS profiler 都在扫这个空间）。

为什么我的 GEMM 慢

按出现频率排序的排查清单：

原因	症状	检测/修复
尺寸太小/维度不齐	利用率 <30%	Roofline 分析；合并 batch、pad 维度到 64/128 倍数
TF32/精度没开	FP32 GEMM 慢 8 倍	`torch.backends.cuda.matmul.allow_tf32=True`（或用 BF16）
频繁的 GEMM 切换	小 kernel 密集，launch 开销主导	算子融合（torch.compile）、CUDA Graph
没用对库路径	`nn.Linear` 意外走了朴素实现	profiler 看 kernel 名（应有 cutlass/gemm 字样）
非连续内存	触发隐式拷贝/转置	`.contiguous()` 检查、内存格式对齐
硬件没喂饱	stream 太少、异步没开	多流并行、异步搬运

第一行的"尺寸问题"占真实世界的 80%：GEMM 慢的第一反应应该是"矩阵多大"，而不是"库怎么调"。

Autotune：让机器帮你选

GEMM 参数空间（tile 形状 × stage 数 × 调度策略 × 精度路径）巨大，手工选不现实，标准做法是自动调优：

torch.compile(mode="max-autotune")：对每个 GEMM 扫多个实现选最快的，编译时间换运行时间
CUTLASS profiler：离线扫参数空间，产出最优配置
shape bucketing：动态 shape 的问题——每个新 shape 都要重新 tune，所以把 shape 归并到有限的桶里

注意：autotune 的结果是shape 敏感的——训练时 batch 固定，tune 一次全程有效；推理时 batch 变化大，要 bucketing 或按 batch 档位缓存。

与其他卷的连接

GEMM 的算术强度与平衡点 → Roofline 一章的核心例
FP8/FP4 让 Tensor Core 算力翻倍 → 混合精度与 FP8
GEMM 是训练 step 的大头 → MFU 的分母
推理的 GEMV（解码时权重读一遍）→ KV Cache 与批处理

本章要点

矩阵乘占 transformer FLOPs 的 90%+，它的实现质量就是速度的第一决定因素。
GEMM 算术强度随尺寸线性增长：小 GEMM 数学上就打不满，"合并成大 GEMM"比"换库"收益大。
cuBLAS 是黑盒标准路径，CUTLASS 是自定义融合的积木；torch.matmul 底层就是 cuBLAS。
"GEMM 慢"80% 是尺寸/维度问题：先 Roofline，再 pad 维度、合并 batch，最后才调库。
autotune（max-autotune/CUTLASS profiler）扫 tile/stage/调度空间，能差 20–40%；结果是 shape 敏感的。