为什么 Karpathy 的 llm.c 坚决不用 atomicAdd?扒开 llm.c 源码看顶级 CUDA 工程美学
打开 llm.c 的llmc/mfu.h,你会看到一件奇怪的事情:一个训练大语言模型的项目,居然用一个 95 行的静态数组手写了一份"GPU 性能数据库"——从 2017 年的 Tesla V100 到 2024 年的 H100,涵盖 38 款 GPU 的 Tensor Core 数量和频率。没有任何外部依赖,不查任何数据库,不调任何系统 API——就是一个硬编码的 C struct 数组。你的第一反应可能是"这也太原始了吧",但当你看到get_flops_promised()函数仅用一行乘法就能从这个数据库中推算出任意 GPU 在任意精度下的理论峰值 TFLOPS 时,你会意识到这个"原始"的设计背后隐藏着一个极其实用的工程判断:在训练大模型时,你真正需要的不是精确到小数点后三位的理论峰值,而是一个足够快、足够稳定、不依赖任何运行时环境的 MFU 估算基线。但更让我震撼的不是 MFU 计算本身,而是在阅读llmc/encoder.cuh的反向传播时发现的一个设计:为了让 WTE(Word Token Embedding)的梯度累加变成确定性的,llm.c 在 CPU 端做了一次完整的 bucket 排序——把所有需要更新同一个 embedding 行的 token 按照词表索引分桶,然后按桶大小降序排列再发给 GPU。这意味着每次反向传播,CPU 都要执行一次 O(B×T) 的 hash 分桶 + 排序操作。为什么要付出这个代价?因为如果你用atomicAdd把多个 token 的梯度直接累加到同一个 embedding 行上,不同的执行顺序会产生不同的浮点舍入结果——你的训练就变得不可重现了。