English version can be found here

Introduction

我去年一直没有认真 follow SOTA 的 kernel 技术, 主要有两个原因, 一是我做的 work 基本用 triton 就够了, 二是我感觉 claude/gpt 一年之内就可以征服 kernel. 很遗憾, 原因一今天已经不成立了, 同时原因二还尚未实现; 所以我只好自己学了.

之前在国内 lab 的时候还是在 ampere 的 GPU 上学习的怎样写 kernel, 感觉难度并不是很大, 如果当时就有 opus 的话. 说到底就是因为黑魔法并没有很多, 本质上我觉得只要掌握了 NVIDIA GPU 的 调度模型和 memory hierarchy, 那写 CUDA 就是写 tensor-core extended 的 C++, 那似乎只需要知道 ldmatrix 和 mma 咋写就够了? 但是在 hopper 上 TMA 和 wgmma 的加入让情况变得复杂了很多, 而为了让 instruction issue/epilogue/wave quantization 不要成为 bottleneck, 又有了 kernel 设计上的一堆黑魔法, 例如 warp specialization 和各种形式的 persistent kernel. 总的来讲, 当代 kernel 和相关 ptx instruction 的复杂程度已经完全不是一个 ADHD 能认真读文档学会的了.

幸运的是, 我们有 opus 4.7, 让我们再次歌颂 claude! 虽然他还没有学会怎样 design 一个 kernel, 但是还是知道怎样基于一个 design 去实现 kernel 的, 而且至少他读各种 ptx 文档还有抄 CUTLASS 的速度比我快多了. 这样以来, 我的工作就变成了: 让 claude 教我怎样写 kernel, 然后我来教 claude 怎样 profile 一个 kernel. 下面用写一个最简单的 GEMM kernel 的过程来展示一下这个 tour.

Road to 989

这个小标题是受 frank 启发的, 989 是 hopper GPU 的 fp16 tensor core peak performance. 下面从一个最基本的 SIMT kernel 开始讲起.

Step 1: non-tiled SIMT kernel

这个应该是所有人学 CUDA 的第一课, opus 显然直接秒了. 写完之后我用了 cuda-skill 让 opus 开始用 NCU 做 profiling, 这里就开始幽默搞笑了. 他一看到 NCU speed of light section 写了个 memory throughput 96%, 然后 NCU 的愚蠢 comment 又告诉他这个 kernel 已经在一个 resource 上遇到瓶颈了, 他就觉得 ok that’s all, 想直接 move on 了. 一个 kernel implementation 尽力的标准, 不在于他让哪个 hardware resource fully utilized, 而是他达到了他的 design 应该达到的 throughput. 例如, 对于一个正常的 MxNxK GEMM kernel, 没做 tiling, 是 memory bound; 那应该去分析的就是, 这个 kernel 需要做多少 global memory read, 有多少能 cache, 考虑了这两个之后, theoretical throughput 就出来了, 然后看实际的 throughput 和理论的差距, 以及差距的原因 (例如看 L1 cache access 有多少个 sector, 过多了就说明 memory coalescing 没做好)