从 4 到 285 TFLOP/s:Blackwell 极致 GEMM 内核系列开篇

HanGuo97 · x · 2026-09-05

Luke D. Huang 发布系列博客第一部:用 CuTeDSL 在 B200 上写 9 个逐步优化的 BF16 GEMM 内核,最终在多个矩阵维度达到 cuBLAS 的 99% 性能。本篇(19 分钟阅读)覆盖:

Kernel 1(朴素 matmul):每线程处理一个输出元素,遍历整个 K 维,基线 4 TFLOP/s。
CuTeDSL 心智模型:如何描述数据布局、切分 tile、映射到 GPU 操作。
Kernel 2(tensor core + TMA + TMEM):每 CTA 处理一块 C,用 TMA 和 barrier 把 A/B tile 搬进 shared memory,发射 tcgen05 MMA 指令累加到 TMEM,经寄存器写回全局内存,提升到 155 TFLOP/s。
Kernel 3(swizzling):发现 shared memory bank 冲突成为瓶颈,仅改变内存布局让并发读分散到不同 bank,提升到 285 TFLOP/s。

作者过去半年深入 GPU 与性能工程,这是记录「写出光速 GEMM 内核」全过程的第一部分。

「Infra」频道最新