从Jam项目看GPU内核优化实战
A primer on GPUs and kernel optimization
2025年12月，我着手用强化学习训练大语言模型，发现现有框架无法满足需求，于是决定从零构建Jam项目。历经10个月，这个起初看似愚蠢的想法已演变为一个极速、灵活且处于前沿的框架。本文作为Jam系列首篇，不直接谈论架构，而是从GPU基础架构、分布式计算切入，详解数据精度、线程块、warp调度及SM资源占用等核心概念。我们将结合Jam中的真实案例，探讨如何在JAX环境下进行kernel优化，并分享利用LLM辅助开发kernel的最新工作流。
实现这一目标绝非易事，它需要大量阅读、学习新领域、挖掘未文档化的特性，以及研读那些在网上几乎找不到文档的库源码。
HN 评论区
11- augment_me
居然没提一句关于内核基准测试/评估的事，简直离谱。
在我看来，在这个所有高性能 GPU 内核都是生成或爬山算法出来的时代，这是我们看到的最主要的错误来源。
实际运行时，内核应该用热缓存还是冷缓存？你的内核是否必须与参考实现按位等价？相对于参考实现，你的内核应该允许多大的误差范围？为了避免热噪声，你要运行多少次内核？你是否在不同的 GPU 上运行以避免过拟合当前硬件？你能否对内核使用 cudagraph 编译以消除启动开销？你是否在运行内核和参考实现时采用 A/B 或 B/A 交错的方式，以避免热偏差？
- saagarjha
> 在用于大部分这项工作的 H200 上，GPU 的主内存是 HBM（高带宽内存）。讽刺的是，这反而是芯片上最慢的内存。
较真一下：HBM 是片外的，不像共享内存和寄存器那样在片内。
- brcmthrowaway
得有人写本关于这个的书。