跳到主要内容

数据截至 (上游 commit f1e2ace65149)

02 · llmc/ 的 CUDA kernel 层

这一章讲什么： 第 1 章的数学在 GPU 上怎么写才快。llmc/ 目录一层一个头文件，每个文件里 是 kernel + launcher。本章按「基础设施 → 各层」的顺序走：先认识四件反复出现的工具，再看 LayerNorm、attention、matmul、fused classifier、encoder 各自的关键技巧。

1. 它要解决的小问题​

同样的数学，写成 CUDA kernel 时多出哪些全新的问题？

CPU 版只需关心「算得对」；GPU 版还要回答：一行 768 个数由谁算（线程/warp/block 怎么分工）、数 据怎么搬（全局/共享内存、cache 策略）、归约怎么做（一个 warp 内怎么求和）、多 kernel 能不能合 并。llm.c 的答案里最有营养的不是某个 kernel，而是这些决策反复使用的一小套工具。

2. 基础设施：先认识四件工具​

2.1 floatX：编译期选精度​

所有 kernel 的激活/权重类型都是 floatX，在 llmc/cuda_common.h:75-92 按编译宏三选一：

编译宏floatXPRECISION_MODE
ENABLE_FP32floatFP32
ENABLE_FP16halfFP16
默认__nv_bfloat16BF16

默认 bf16——这就是「混合精度主线」的类型基础（第 4 章展开）。同一类型宏也传给 cuBLAS （llmc/cublas_common.h:15-19 的 CUBLAS_LOWP）。

2.2 Packed128：逼编译器用 128 位访存​

GPU 一次 128 位（16 字节）的向量加载比 4 次 32 位加载省指令。Packed128 （llmc/cuda_utils.cuh:15-48）是一个 16 字节对齐的小 struct，里面装 size = 16/sizeof(T) 个元 素：fp32 装 4 个，bf16 装 8 个。配套函数：

函数作用
load128 / store128普通 128 位读/写（llmc/cuda_utils.cuh:53-72）
load128cs / store128cs带 .cs（streaming）cache 提示：这数据很快不再用，别占 L1
store128cg缓存进 L2 但绕过 L1

别名 f128 = Packed128<float>、x128 = Packed128<floatX>（llmc/cuda_utils.cuh:78-79）。翻 llmc/ 任意 kernel 都会看到：循环步长是 WARP_SIZE * x128::size，因为每个线程一口吃 8 个 bf16。

2.3 warp/block 归约​

归一化、softmax、求范数都要「把 32 个线程的值合成一个」。两个原语：

  • warpReduceSum / warpReduceMax（llmc/cuda_utils.cuh:147-160）：__shfl_xor_sync 蝶形归 约，5 步把 32 个值加成一个。
  • blockReduce（llmc/cuda_utils.cuh:166-187）：先 warp 内归约，各 warp 结果写 shared memory，再由 warp 0 归约一次——整个 block 一个值。

2.4 错误检查与 NVTX​

cudaCheck（llmc/cuda_common.h:52-58）把每个 CUDA 调用包成「出错即打印文件行号并退出」； NVTX_RANGE_FN()（llmc/cuda_common.h:124）给每个 launcher 打上 profiler 可见的范围标签——第 5 章的剖析全靠它。

3. LayerNorm：从「能跑」到「融合」​

3.1 两个前向 kernel​

llmc/layernorm.cuh 里留着两代实现，恰好演示优化思路：

  • kernel3（llmc/layernorm.cuh:20-65）：一个 warp 负责一行（一个 token 的 C 维）。lane 间 隔读、两次 warpReduceSum 求均值和方差，写出时用 __ldcs/__stcs 流式提示。
  • kernel6（llmc/layernorm.cuh:67-140）：先把 weight/bias 搬进 shared memory，输入按 x128 向量化读入并暂存 smem，均值方差照算，写出时从 smem 读——全局内存只碰一次。

launcher（layernorm_forward，llmc/layernorm.cuh:433-456）有个实用细节：kernel6 的 smem 需 求可能超过默认 48KB，于是先 cudaFuncSetAttribute 尝试提额，失败就退回 kernel3——同一个 接口，按硬件能力自动降级（llmc/layernorm.cuh:444-454）。

3.2 跨 block 边界的融合：fused_residual_forward5​

Transformer 里每个残差相加后面紧跟一个 LayerNorm。两个 kernel 意味着残差结果要在全局内存绕一 圈。fused_residual_forward_kernel5（llmc/layernorm.cuh:142-219）把两者合一：

inp1 + inp2 ──► residual（写出）──► 顺手在寄存器/smem 里求均值方差 ──► normed（写出）

更有意思的是它的调用位置：在 gpt2_forward 里，这个 kernel 被调用时吃的是「本层 MLP 输出」 和「上一站残差」，吐出的归一化结果却是下一层的 ln1（train_gpt2.cu:736-750）——融合跨过 了 Transformer block 的边界；最后一层的归一化结果则直接写成最终的 lnf （train_gpt2.cu:746-749）。

3.3 反向：两级归约 + 「最后一个 block 收尾」​

layernorm_backward_kernel10（llmc/layernorm.cuh:233-431）要算三种梯度：dinp（逐元素，好 并行）、dweight/dbias（要对全部 B×T 求和，难并行）。对后者：

  1. grid 只开「刚好填满 GPU」的 block 数，每个 block 在 shared memory 里累加自己那份；
  2. 各 block 把部分和写到全局 scratch（llmc/layernorm.cuh:369-373）；
  3. 每个 block 完成时 atomicInc 一个计数器（llmc/layernorm.cuh:378），最后一个完成的 block（*tmp_flag == gridDim.x-1，llmc/layernorm.cuh:381）顺手把所有部分和归约完。

用一次原子计数换掉了「再 launch 一个归约 kernel」，这是 GPU 编程里经典的 persistent-kernel 手 法。

4. Attention：不手写 flash，也能拼出来​

4.1 五步流水线（非 cuDNN 路径）​

默认路径不手写 flash attention，而是把注意力拆成已有零件（attention_forward， llmc/attention.cuh:195-237）：

inp (B,T,3C)
  │ ① permute_kernel        拆成 q/k/v 三份 (B,NH,T,hs)
  ▼
  │ ② matmul_cublaslt       preatt = q @ k^T，batched GEMM（B*NH 批）
  ▼
  │ ③ softmax_forward_kernel5  行内 online softmax（含因果掩码与 scale）
  ▼
  │ ④ matmul_cublaslt       vaccum = att @ v，再一次 batched GEMM
  ▼
  │ ⑤ unpermute_kernel      (B,NH,T,hs) 拼回 (B,T,C)
  ▼
out (B,T,C)

②④ 两次矩阵乘直接调 cuBLASLt 的 strided batched GEMM——「手写 kernel 负责数据重排，极致优化 的 GEMM 交给库」是 llm.c 的一贯分工。inp 缓冲区同时被复用为 preatt/vaccum 的 scratch （llmc/attention.cuh:197-199 注释）。

4.2 online softmax kernel​

softmax_forward_kernel5（llmc/attention.cuh:85-150）一个 warp 处理一行注意力分数，有两个 细节值得抄：

  • online（在线）softmax：一趟扫描里同时维护 running max 和 running sum，发现更大的 max 就 把已有和按比率折算（llmc/attention.cuh:117-129）——比「先求 max 再求和」省一趟读写。
  • 反向序遍历 block：idx = (gridDim.x - blockIdx.x - 1) * num_warps + warp_id （llmc/attention.cuh:96-101）。kernel 结束后紧接着是 att @ v 的 GEMM，反向序让 cache 里留 下的是矩阵左上角——正是下一次乘法最先读的部分。

4.3 反向与 cuDNN 路径​

attention_backward（llmc/attention.cuh:239-268）是 6 个调用的镜像：unpermute 反向 → 两次 GEMM（datt、dv）→ softmax 反向（就地把 datt 改成 dpreatt， softmax_autoregressive_backward_inplace_kernel，llmc/attention.cuh:152-191）→ 两次 GEMM（dq、dk）→ permute 反向。

编译时加 USE_CUDNN=1 则整段换成 cuDNN 的 flash attention（train_gpt2.cu:721-724 分流到 attention_forward_cudnn，实现在 llmc/cudnn_att.cpp）。代价是编译时间从几秒涨到约一分钟 （Makefile:24-25 注释），所以默认关闭。

5. Matmul：一切交给 cuBLASLt 的 epilogue​

5.1 一个包装函数打天下​

所有前向/反向的矩阵乘都走 matmul_cublaslt（llmc/matmul.cuh:109-228）。它每次调用做的事： 建 operation descriptor、按 transA/transB 建 layout、（可选）配 strided batch、选 epilogue、 cublasLtMatmulAlgoGetHeuristic 选算法、执行、销毁。入口先检查所有指针 16 字节对齐 （llmc/matmul.cuh:118-122）——不对齐直接退出，因为向量化访存全建立在这个假设上。

5.2 epilogue：免费的逐元素运算​

cuBLASLt 允许在 GEMM 输出时顺带做一个逐元素操作，llm.c 用到了四种（llmc/matmul.cuh:172-184）：

epilogue用在哪效果
CUBLASLT_EPILOGUE_BIAS前向 qkv/proj/fc加 bias 不单独开 kernel
CUBLASLT_EPILOGUE_GELU_AUX_BIAS前向 FC + GELU 融合GEMM 输出直接过 GELU，并把 GELU 前值存到 AUX 供反向
CUBLASLT_EPILOGUE_DGELU反向 dinp顺带算 GELU 导数
CUBLASLT_EPILOGUE_BGRADB反向顺带归约出 bias 梯度

matmul_forward_cublaslt（llmc/matmul.cuh:231-242）按 gelu_fusion 开关决定走融合还是先 GEMM 再单独 gelu_forward；注释记录了取舍：GELU 融合只在 H100+ 上划算（llmc/matmul.cuh:235）。

5.3 反向的「=、+=」约定在 GEMM 层的落实​

matmul_backward（llmc/matmul.cuh:244-289）把第 1 章的约定翻译成 cuBLASLt 参数：

  • dinp：每个 micro-step 只有一个来源，beta=0 直接写；
  • dweight：要跨 micro-step 累积，accumulate=true 即 beta=1，D += C （llmc/matmul.cuh:287-288）；
  • dbias：走自研的 matmul_backward_bias_kernel9（llmc/matmul.cuh:17-81），warp 分工 + 必要 时跨 block 部分和（reduce_add_sum_kernel 收尾）。

6. fused classifier：全仓库最省显存的一个 kernel​

6.1 问题：B×T×V 实在太大了​

V = 50257。B=64、T=1024 时，probs 一个张量就超过 30 亿元素。朴素流程要把 logits → softmax 概率 → 交叉熵 → 反向 dlogits 各物化一遍。fused_classifier_kernel5 （llmc/fused_classifier.cuh:70-136）把整条链融成一个 kernel，probs 从不落盘：

  1. prepare_softmax_blockwide3（llmc/fused_classifier.cuh:19-60）：block 级 online softmax， 只算出该行的 max 与 sum（两个标量）；
  2. 单线程算 target 位置的概率，写出 loss（llmc/fused_classifier.cuh:82-85）；
  3. 同一 kernel 里直接把 logits 原地改写成梯度 (p - indicator) * dloss （llmc/fused_classifier.cuh:96-113）——反向传播的第一步顺带完成。

6.2 一个真实踩过的坑：race condition​

kernel 中部有一句普通但救命的 __syncthreads()，旁边注释记录了 bug 本身 （llmc/fused_classifier.cuh:89-93）：算 loss 读 logits 与「把 logits 改写成梯度」是并发 的；一旦改写先发生，exp 的参数会暴涨，loss 变成 inf。先同步，再改写。

7. encoder 反向：确定性是设计出来的​

7.1 问题：词嵌入反向天然想上 atomicAdd​

wte 反向要把梯度按 token id scatter 累加：同一个 token 在 batch 里出现多次，最直接的写法是 atomicAdd——但原子加顺序不定，浮点结果不可复现。

7.2 解法：CPU 先分桶，GPU 每桶一个 block​

encoder_backward（llmc/encoder.cuh:169-229）的步骤：

  1. 先发 wpe kernel（wpe_backward_kernel，llmc/encoder.cuh:119-154）：位置嵌入每个 (t,c) 元素只被一个线程写，天然确定；
  2. CPU 把 (token, 通道组) 分桶：用 hashmap 把 B×T 个位置按 token id × 通道组聚类 （llmc/encoder.cuh:186-196）；
  3. 桶按大小降序排序——大桶先跑，避免最后剩一个大桶拖尾（llmc/encoder.cuh:200-204）；
  4. 每个 block 负责一桶内的一段，wte_backward_kernel（llmc/encoder.cuh:47-116）内各 warp 各 算各的、warp 0 汇总，每个梯度元素只写一次——顺序固定，结果确定。

写回时用 stochastic_rounding（见第 4 章）把 fp32 累加结果舍入成 bf16，种子由 bucket/thread 下标唯一确定，同样保证可复现（llmc/encoder.cuh:103-112）。

8. 其余小 kernel 一页速览​

kernel文件一句话
encoder_forward_kernel3llmc/encoder.cuh:19wte+wpe 向量加，x128 读写
gelu_forward_kernel2 / gelu_backward_inplace_kernelllmc/gelu.cuh:13、:28逐元素 tanh 近似；反向就地改写
residual_forward_kernelllmc/layernorm.cuh:221逐元素加，通常被 kernel5 融合替代
global_norm_squared_kernel + global_norm_aggregate_kernelllmc/global_norm.cuh:27、:38梯度范数两级归约（第 4 章）
adamw_kernel3llmc/adamw.cuh:50AdamW + 随机舍入（第 4 章）

9. 关键细节与坑​

  • 对齐是硬假设。 所有 load128/store128 要求 16 字节对齐；cuBLASLt 包装入口直接检查并退出 （llmc/matmul.cuh:118-122）。词表填充到 50304、通道数取 64 的倍数，都是为此服务。
  • 非 cuDNN 注意力要求 att 缓冲区预先清零。 当实际 T 小于分配时的 seq_len，未用区域必须 memset，否则因果掩码外读到脏数据（train_gpt2.cu:722-725）；整个激活内存在分配时统一清零 过一次（train_gpt2.cu:267-270）。
  • shared memory 超额要显式申请。 kernel6/kernel5 的 smem 可能超 48KB，launcher 用 cudaFuncSetAttribute 申请、失败自动降级（llmc/layernorm.cuh:444-454、:477-488）。
  • 这些 kernel 都有「前世」。 每个主力 kernel 在 dev/cuda/ 里都有从 v1 到 vN 的演进史和基 准（第 5 章 §4）；llmc/ 里函数名的数字后缀（kernel3/kernel5/kernel6/kernel10）就是那段历 史的年轮。
  • block reduce 的 smem 复用陷阱。 blockReduce 的注释提醒：循环里复用时必须 final_sync，否则共享内存被下一轮提前覆写（llmc/cuda_utils.cuh:166-187）。

10. 代码地图​

主题文件路径符号名
精度与类型llmc/cuda_common.h、llmc/cublas_common.hfloatX、PRECISION_MODE、CUBLAS_LOWP
向量访存与归约llmc/cuda_utils.cuhPacked128、load128cs、store128cg、warpReduceSum、blockReduce
LayerNorm 前向llmc/layernorm.cuhlayernorm_forward_kernel3、layernorm_forward_kernel6、layernorm_forward
残差融合llmc/layernorm.cuhfused_residual_forward_kernel5、fused_residual_forward5
LayerNorm 反向llmc/layernorm.cuhlayernorm_backward_kernel10、layernorm_backward
注意力（默认）llmc/attention.cuhattention_forward、permute_kernel、softmax_forward_kernel5、unpermute_kernel、attention_backward
注意力（cuDNN）llmc/cudnn_att.cpp、llmc/cudnn_att.hattention_forward_cudnn、attention_backward_cudnn
矩阵乘llmc/matmul.cuhmatmul_cublaslt、matmul_forward_cublaslt、matmul_backward、matmul_backward_bias_kernel9
分类器融合llmc/fused_classifier.cuhfused_classifier_kernel5、prepare_softmax_blockwide3、fused_classifier
编码器llmc/encoder.cuhencoder_forward、encoder_backward、wte_backward_kernel、wpe_backward_kernel
GELUllmc/gelu.cuhgelu_forward、gelu_backward_inplace