大模型FCN算子（GEMM）CUDA实现

1. 性能目标到底是什么#

单次 GEMM 的工作量为：

运行时间为 $t$ 时，吞吐为 $2MNK/t$。但只报 TFLOP/s 不足以判断 kernel 好坏，因为它还受矩阵是否足够大、是否对齐、是否发生额外读写、launch 开销和数值精度影响。

一个有用的最小模型是 Roofline：

其中 $I$ 是算术强度。GEMM 通过 tile 复用 A/B 的元素，通常最终受计算单元约束；而小矩阵、skinny GEMM、拆得过细的 epilogue 则可能重新落回带宽或 launch 受限。

因此，“极致优化”不是只看一项指标，而是同时做到：

• A/B 的全局访存合并且对齐；
• shared memory 的布局没有可避免的 bank conflict；
• Tensor Core 或 CUDA Core 持续有数据可算；
• 寄存器、shared memory、CTA 数之间保持可接受的 occupancy；
• 输出的 bias、activation、residual、量化尽量在同一次写回中完成；
• 对不适合主路径的形状及时换 kernel，而不是硬撑。

2. 从朴素实现出发：为什么它离峰值很远#

标量 baseline 的逻辑如下：

1
template <typename T>
2
__global__ void gemm_naive(int M, int N, int K,
3
                           const T* A, const T* B, T* D) {
4
  int m = blockIdx.y * blockDim.y + threadIdx.y;
5
  int n = blockIdx.x * blockDim.x + threadIdx.x;
6
  if (m >= M || n >= N) return;
7
  float acc = 0.f;
8
  for (int k = 0; k < K; ++k)
9
    acc += float(A[m * K + k]) * float(B[k * N + n]);
10
  D[m * N + n] = T(acc);
11
}

它正确，却几乎没有复用：相邻 thread 会重复读取同一行 A 或同一列 B。即使 L1/L2 能帮上一部分忙，数据搬运仍由缓存偶然性决定，无法稳定接近峰值。优化的第一原则就是把“会被多个输出元素使用的数据”显式搬到片上。

3. 第一级：CTA tile 与层次化分块#

一个 thread block 计算 $D$ 的 BM × BN tile，并沿 K 每次推进 BK。数据层级是：Global A/B → shared-memory A/B tile → warp tile → registers / MMA fragments → D tile。

在每个 K tile 中，CTA 协作将 A/B 连续搬入 shared memory；之后每个 warp 读取其中一小块，多个输出元素复用同一份数据。关键参数不是“越大越好”：

• BM, BN 越大，A/B 的复用越好，但累加器和 shared memory 也更多；
• BK 太小，循环与同步开销明显；太大则挤占 shared memory、降低常驻 CTA；
• warp tile 决定寄存器 fragment 大小与并行度；
• instruction tile 必须匹配实际架构和 dtype。

一个正确的调优流程是先枚举少量合理 tile，而不是从无穷多数字中盲猜。以常见矩阵乘数据通路为例：BM/BN 取 64、128、256 一类，BK 对 FP16/BF16 常取 32/64，warp tile 与 MMA 指令 tile 对齐；实际值必须让编译器资源报告和 profile 说话。

4. 第二级：访存布局比算术更早决定成败#

5. 第三级：让搬运和计算真正重叠#

同步式 tile loop 是“搬 A/B → 等待 → 计算 → 再搬”的串行流程。现代 GPU 的关键是多阶段流水：当 stage s 的 fragment 正在 MMA 时，stage s+1 的 A/B tile 已经在路上。

1
// stage 0: async copy A0/B0 -> shared
2
// stage 1: async copy A1/B1 -> shared, compute A0 * B0
3
// stage 2: async copy A2/B2 -> shared, compute A1 * B1
4
// ...

在支持的架构上，可用异步 copy 建立双缓冲或多缓冲；再新的数据通路可以由专门的搬运机制和 warp-specialized producer/consumer 组织。核心不在于记住某个 intrinsic 名称，而在于三个不变量：

图 2 中的 C_acc 是同一输出 tile 的寄存器累加器；它只在所有 K tile 完成后才交给 epilogue。双缓冲的目的不是增加计算量，而是让 load(k+1) 与 MMA(k) 重叠。

1. consumer 在读取前确认对应 stage 已完成；
2. producer 不覆盖仍被消费的 stage；
3. stage 数增加带来的 latency hiding，必须大于它对 shared memory 和 occupancy 的代价。

CUTLASS 把这些策略封装为可组合的 mainloop、collective 和 epilogue，是阅读工业级 GEMM 流水线的好入口；其近代实现还覆盖了 grouped GEMM、FP8、block-scaled 类型与架构专用调度。NVIDIA CUTLASS

6. Tensor Core 不是一个 dtype 开关#

Tensor Core 路径的本质是 warp/warp-group 使用特定矩阵乘累加指令，操作 fragment 而非单个标量。要让它工作，至少要同时满足：dtype、矩阵布局、指令 tile、K 对齐和目标计算能力的约束。

典型策略是：主路径处理满足 tile 对齐的大块，边界使用 predicated load/store 或单独的 residue kernel；不要让每一次 MMA 都包在复杂的逐元素边界判断里。对于常见混合精度，输入可以是 FP16/BF16/TF32/FP8，而 accumulator 通常是 FP32；CUDA Programming Guide 明确说明 BF16 WMMA fragment 以 float 为累加类型，TF32 需要从 FP32 显式转换得到相应语义。CUDA C++ Programming Guide

7. 精度矩阵：从 INT4/FP4 到 FP64#

“支持所有精度”不应理解为一个模板把所有类型硬塞进一个指令序列，而是每种精度都有明确的输入、累加、缩放、输出与 fallback 规则。

这里有四个容易被忽略的细节：

1. 存储与累加分离。 FP16 输入、FP32 累加不是“可选优化”，而应是通用训练的默认数值合同。
2. 缩放粒度属于 ABI。 INT4/FP4/FP8 的 scale 按 tensor、channel、group 还是 block 存储，会同时决定精度、访存和 kernel layout；必须在 tensor descriptor 中明确。
3. K 方向分割会影响数值。 split-K 改变累加顺序；高精度或可复现模式应采用固定归约树，必要时使用更高精度 partial。
4. 低比特不等于总更快。 若 shape 太小、scale 读取占主导，或需要单独 dequant kernel，端到端速度可能反而下降；理想做法是把 dequant/scale 融合在加载和 epilogue 中。

当前 CUTLASS 的功能与变更记录展示了 FP8、NVFP4、MX 格式、block scaling 和 sparse/block-scaled GEMM 等路线；它们说明低比特 GEMM 是一套数据布局与调度问题，而不只是把 float 换成更短的类型。CUTLASS changelog

8. GEMM 不是一种形状：分段 dispatch 才是通用解#

和 OneFlow 的分段 Softmax 一样，GEMM 也不存在一个在所有形状上最优的 kernel。一个实用 dispatcher 至少把以下信息作为 key：

1
{arch, dtypeA, dtypeB, dtypeD, accumulator,
2
 layoutA, layoutB, layoutD, M, N, K,
3
 alignmentA/B/D, alpha/beta, epilogue, batch/group count,
4
 deterministic, workspace limit}

对应的主路径可这样划分：

split-K 并不是免费午餐：它增加 workspace 或原子/归约开销，只有原始网格并行度不足时才值得使用。Stream-K 的意义也正是在 tile 间重新分配 K 工作，缓解固定 tile 调度的不平衡；它应作为 benchmark 候选，而非默认开关。

9. 一个可维护的 kernel 骨架#

生产代码通常不把所有细节塞进 __global__ 函数，而是把策略变为类型和常量。下面是结构示意：

1
template <class ElementA, class ElementB, class ElementD,
2
          class Accumulator, class Tile, class Layout,
3
          class Mainloop, class Epilogue>
4
struct GemmKernel {
5
  __device__ void operator()(GemmArgs args) {
6
    // 1. 根据 CTA 坐标取得 D 的 Tile。
7
    // 2. Mainloop: 预取 A/B，管理 pipeline stage，执行 MMA。
8
    // 3. Epilogue: alpha/beta、bias、residual、activation、量化/转换。
9
    // 4. 只写回最终 D。
10
  }
11
};
12

13
GemmPlan make_plan(const GemmArgs& args) {
14
  // 静态规则先过滤不合法的 dtype/layout/alignment；
15
  // 查 shape+arch 的 autotune 缓存；缺失时仅 benchmark 少量候选。
16
}

这种分层使主循环只关心“如何高复用地计算”，epilogue 只关心“如何生成最终输出”，dispatcher 则对形状做选择。它也让你能把框架的 bias、GELU/SwiGLU、残差、输出量化等融合进去，而不污染 MMA 内层循环。

10. 融合 epilogue：节省的往往比 MMA 更真实#

以 FCN 为例，常见路径是：

1
D = A @ W
2
Y = activation(D + bias)

若拆成两个 kernel，D 至少经历一次完整的 global write 与一次 global read。把 alpha/beta + bias + activation + cast/quantize 放在 GEMM epilogue 中，可以让 accumulator fragment 在寄存器中完成最后变换后只写一次 Y。

但融合有边界：复杂 activation、大量额外输入或分支会提高寄存器占用，可能伤害主 GEMM。判断标准永远是端到端 benchmark，而不是“kernel 数量越少越好”。

11. 如何验证“快”与“对”#

基准至少应包含 LLM 常见 shape：prefill 的大 M、decode 的 M=1..16、MLP 的大 N/K、QKV 的非方阵、MoE 的 grouped GEMM，以及难看的不对齐尾部。每项都要记录：

1
environment: GPU/SM、CUDA/驱动、时钟状态
2
problem: 矩阵 layout、M/N/K、dtype、accumulator
3
timing: warmup、重复次数、CUDA event 时间、总时间、TFLOP/s、workspace
4
correctness: 最大绝对/相对误差、随机种子、确定性模式

正确性基准要独立于被测实现：低精度输入可用 FP32 或 FP64 reference 比较，比较容忍度随 dtype 和 K 长度设置；对 FP8/FP4/INT4，额外测 scale 边界、NaN/Inf、饱和、非整齐 K 和 beta 非零。性能分析使用 Nsight Compute 观察 Tensor Core 活跃度、DRAM 吞吐、shared bank conflict、寄存器 spill、occupancy 和 warp stall 原因。

不要只对一个漂亮的 8192×8192×8192 报数。真正的大模型推理，往往被 decode、短 batch、量化 scale、融合边界和 grouped workload 决定。

12. 常见的“优化”陷阱#

1. 用一个 tile 打天下。 大矩阵最优 tile 在 skinny GEMM 上可能只会浪费线程。
2. 无条件加深 pipeline。 stage 多会占用更多 shared memory，可能使常驻 CTA 下降。
3. 只测 kernel，不测 epilogue/量化/转换。 用户感知的是端到端时间。
4. 默认 split-K。 没有并行度问题时，额外规约就是纯负担。
5. 将 TF32/FP8 当作无损 FP32。 数值合同必须由模型和测试决定。
6. 把未对齐指针强转成 vector 类型。 这不是优化，是未定义行为风险。
7. 忽略 launch 开销。 对 decode 或小 GEMM，CUDA Graph、grouped/persistent 调度常比更大的 MMA tile 重要。

总结#

高性能 GEMM 的路径可以概括为：用 CTA/warp/instruction 三层 tile 建立复用，用对齐向量化和无冲突 shared layout 喂饱计算单元，用异步多阶段流水隐藏搬运延迟，用 Tensor Core 与混合精度匹配硬件能力，最后用 shape-aware dispatch 覆盖大矩阵、skinny、split-K 与 grouped workload。

对于大模型 FCN，真正可上线的版本还必须把 FP4/INT4、FP8、FP16/BF16、TF32、FP32、FP64 的数值合同和 fallback 写清楚，并把 bias/activation/quantize 融进合适的 epilogue。所谓“极致优化”，从来不是某个花哨 intrinsic，而是一套在真实输入、真实数值约束和真实硬件上持续测得最快的系统工程。

参考资料#

• dlog：CUDA学习之路[11]：矩阵乘法计算详解
• OneFlow：如何实现一个高效的 Softmax CUDA kernel？
• NVIDIA CUTLASS
• CUTLASS 变更记录
• CUDA C++ Programming Guide