9.0
重磅
AI SCORE
编程提效2026-08-16 10:39
Codex自动化研究：CUDA内核232倍加速实战
dev.to · AI#AI编程#CUDA优化#Codex
Editor brief · 编辑速览
用OpenAI Codex自动化研究循环将CUDA内核从2.3ms优化至9.9µs，主要手段包括内存合并、warp级归约和冗余读取消除。
使用 OpenAI 的 Codex 构建自动化研究循环，我将一个自定义 CUDA kernel 从每次操作 2.3 ms 的基准优化到了仅 9.9 µs——提升了 232 倍——期间没有手动编写一行汇编代码，也没有手动调优内存访问模式。本文将详细拆解这个自动化研究 pipeline 的工作原理、令我意外的地方、它失败的地方，以及如何将同样的方法论应用到自己的计算密集型工作负载上。
使用 Codex 进行自动化研究循环，可以通过自动化假设生成、基准测试和迭代，将数周的 kernel 优化压缩到几个小时内完成。
232 倍的性能提升主要来自三个改动：内存合并（memory coalescing）、warp 级归约（warp-level reduction），以及消除冗余的全局内存读取。
Codex 并不会在第一次就做对——真正的力量在于反馈循环，而不是单次提示。
你不需要成为 CUDA 专家才能开始。但你需要理解自己的性能瓶颈。
这种方法最适合具有明确、可衡量性能目标的操作（推理 kernel、数据预处理、线性代数原语）。
在深入基准测试之前，让我们先定义"自动化研究"在这个语境下究竟意味着什么，因为这个术语被用得太随意了。
使用 Codex 进行自动化研究，指的是一个结构化的 pipeline——你将 Codex（或其他 Codex 级别的模型）不仅仅用作代码自动补全工具，而是作为一个自主研究智能体，能够：
提出优化假设
编写相应的实现
针对已知基准进行基准测试
分析 profiler 输出
根据所学到的东西用新假设进行迭代
这与"让 Codex 写一个 CUDA kernel"有本质区别。区别在于反馈循环——性能数据被反馈到模型的上下文中，模型会推理为什么某个方案比另一个更快或更慢。
可以把它想象成一个初级 GPU 工程师，能在你睡觉时编写和测试 50 个实验。
[INTERNAL_LINK: Introduction to CUDA Kernel Optimization]
目标是一个用于 transformer 推理 pipeline 的自定义 softmax kernel，运行在 NVIDIA H100 GPU 上。基准是一个为正确性而非速度编写的朴素实现：
# Baseline: naive softmax in CUDA (simplified)
# - No shared memory usage
# - Uncoalesced global memory reads
# - Atomic operations for reduction
# Baseline benchmark: 2.3 ms per forward pass (batch_size=512, seq_len=2048)
作为参考，PyTorch 内置的 F.softmax 在相同配置下运行时间约为 0.87 ms——所以即便是"生产级"基准也不怎么样。这个朴素实现是我们在研究原型中快速编写后从未重新审视的代码。
硬件环境：
GPU: NVIDIA H100 SXM5 (80GB)
Profiling 工具: NVIDIA Nsight Compute
这里是方法论变得有趣的地方。我构建的 pipeline 有五个组件：
一个管理研究循环的 Python 脚本。它调用 Codex API、将生成的代码写入磁盘、编译、用 nvprof 或 Nsight Compute 运行、捕获输出，然后传回给模型。
# Simplified orchestrator pseudocode
for iteration in range(max_iterations):
hypothesis = codex.generate(
context=previous_results,
prompt=f"Given this profiler output: {profiler_data},
propose one specific optimization and implement it."
)
kernel_code = extract_code(hypothesis)
benchmark_result = compile_and_benchmark(kernel_code)
previous_results.append({
"hypothesis": hypothesis,
"result": benchmark_result,
"speedup": calculate_speedup(benchmark_result, baseline)
})
提示词至关重要。通用的提示词产生通用的结果。以下是效果最好的结构：
You are optimizing a CUDA kernel for [operation].
Current performance: [X ms]
Target performance: [Y ms]
Hardware: [GPU model]
Previous attempts and their results: [structured list]
Profiler output from the last run:
[nsight compute output]
Identify ONE specific bottleneck visible in this profiler data.
Propose a concrete optimization. Write the complete modified kernel.
Do not change the function signature.
关键约束是一次只做一个优化。当我让 Codex 同时提出多个改动时，根本无法将性能提升归因到具体的决策上。
每个生成的 kernel 都会用 nvcc 编译，并在基准测试前根据正确性测试套件运行。这一步是硬性要求——Codex 经常产生快速但数值结果不正确的 kernel。
nvcc -O3 -arch=sm_90 kernel.cu -o kernel_test
./validate_correctness kernel_test # Compare against PyTorch reference
./benchmark kernel_test --iterations=1000 --warmup=100
NVIDIA Nsight Compute 提供了结构化的 JSON 输出，比原始文本对 LLM 更有用。我配置它输出：
内存吞吐量（GB/s）
L1/L2 缓存命中率
内存访问模式（合并分析）
每次迭代都会记录到一个结构化的 JSON 文件中。这有两个目的：它为 Codex 提供了丰富的历史记录，告诉它尝试过什么；同时也为你提供了审计跟踪，让你理解什么真正有效。
以下是自动化研究循环实际发现的内容，按顺序排列：
Profiler 立即标记出了未合并的全局内存读取。基准 kernel 以列主序模式访问元素，导致了 32 次独立的内存事务，而本应只需一次。
Codex 的修复方案：重构线程索引方式，让相邻线程访问相邻的内存地址。教科书式的优化，但模型从 profiler 数据中正确地识别出了这个问题。
结果：2.3 ms → 0.19 ms（12 倍）
修复了内存合并后，profiler 现在显示 L2 缓存抖动。Codex 提议将输入 tile 加载到共享内存中并在那里执行归约，从而大幅减少全局内存带宽消耗。
这需要三次迭代才能做对——前两次尝试由于共享内存写入中的竞态条件产生了错误结果。验证套件捕获了这两个失败。
结果：0.19 ms → 0.040 ms（4.8 倍）
这里变得真正有趣了。Codex 提议用 warp shuffle 指令（__shfl_down_sync）替换共享内存归约，这种指令允许 warp 内的线程直接通过寄存器交换数据，完全不触碰共享内存。
坦白说：我知道 warp shuffle 的存在，但在我自己的优化过程中可能不会这么早就想到用它。模型从 profiler 的 warp 效率指标中识别出了这个模式。
结果：0.040 ms → 0.012 ms（3.2 倍）
最后一个显著的提升来自调整线程块配置。Codex 分析了 profiler 输出中的寄存器使用量和共享内存占用，然后推荐了特定的 <<<gridDim, blockDim>>> 参数以最大化 SM 占用率。
结果：0.012 ms → 0.0099 ms（1.6 倍）
最终 kernel 在该配置下比 PyTorch 的生产级 softmax 快 2.6 倍。这是部署时真正重要的数字。
没有一篇自动化研究的文章能不谈失败。以下是 pipeline 挣扎的地方：
正确性问题很频繁。大约 35% 的生成 kernel 没有通过验证套件。大多数失败很微妙：边界条件中的差一错误（off-by-one）、缺少 __syncthreads() 调用，或对非 2 次幂序列长度处理不当。
模型有时会退化。在第 9 和第 11 次迭代中，Codex 提议的改动理论上合理但产生了更慢的 kernel。自动化研究循环正确地将这些识别为退化并丢弃了它们，但这提醒我们：模型是在启发式地推理性能，而不是分析性地推理。
硬件特定知识存在差距。Codex 的训练数据中 Volta/Ampere CUDA 代码可能比 H100 特定优化更多。一些 Hopper 特定功能（如 Tensor Memory Accelerator）需要在提示词中明确提示后模型才会考虑。
它停滞了。在第 15 次迭代之后，又进行了 12 次迭代都没有产生有意义的改进。模型不断提议对已尝试过的方法进行微调。此时，确实需要人类专业知识来识别下一个前沿。
[INTERNAL_LINK: When to Use AI-Assisted Optimization vs. Manual Tuning]
以下是诚实的工具链分解：
一个诚实的警告：你需要 GPU 访问权限。在 H100 上运行 15+ 次基准迭代会累积成本。如果你没有本地硬件，Lambda Labs GPU Cloud 提供具有竞争力的 H100 实例价格，这就是我用于长时间运行的平台。
如果你想复制这个方法论，以下是实用的检查清单：
在研究循环之前：
[ ] 建立一个正确的基准并彻底验证
[ ] Profile 你的基准以识别前 2–3 个瓶颈
[ ] 定义明确的性能目标（不要盲目优化）
[ ] 设置自动化正确性验证——这是最重要的一步
在研究循环期间：
[ ] 每次迭代只做一个优化
[ ] 始终将 profiler 输出传回给模型，而不仅仅是 timing 数字
[ ] 记录所有内容——你之后会想回顾什么起了作用
[ ] 在开始前设置最大迭代预算
在以下情况下停止：
当每次迭代的收益低于约 5% 时
当模型开始重复之前尝试过的方法时
当达到硬件理论极限时（检查 roofline model）
[INTERNAL_LINK: Roofline Model Analysis for GPU Kernels]
Q: 我需要懂 CUDA 才能使用这种方法吗？
你需要足够的 CUDA 知识来验证输出并理解 profiler 数据——大约"中级"水平。你不需要成为 kernel 专家，但如果你无法阅读 CUDA kernel 并发现明显错误，验证步骤就会变得不可靠。我建议在尝试在生产代码上使用此方法之前先过一遍 CUDA Programming Guide。
Q: 这在 AMD GPU 上配合 ROCm 也能用吗？
我没有在 ROCm 上系统地测试过。Codex 的训练数据严重偏向 CUDA，所以在 HIP kernel 上你会看到更多的正确性失败。方法论是合理的，但预计需要在验证上花费更多时间。
Q: 运行这个 pipeline 的成本是多少？
我 15 次迭代的运行成本约为 4.20 美元的 Codex API 调用（由于 profiler 输出，提示词很长）。Lambda Labs 上 GPU 计算成本约为 18 美元用于完整的基准测试套件。总计：不到 25 美元就获得了 232 倍的提升。这是一个非凡的投资回报。
Q: 这种方法能用于 CPU kernel（AVX、NEON）吗？
可以，但需要修改。将 Nsight Compute 替换为 perf 或 Intel VTune，并调整提示词模板以适应 SIMD intrinsics。反馈循环原理相同。Codex 对 AVX-512 intrinsics 的了解还不错，但不如对 CUDA 的了解。
Q: 232 倍的提升在不同 batch size 下可复现吗？
提升因配置而异。在较小的 batch size（batch_size=64）下，我们看到约 180 倍。在较大的 size（batch_size=1024）下，提升压缩到约 95 倍，因为基准变得不那么病态了。232 倍这个数字是针对所述基准测试配置的。始终在你的实际工作负载上测量。
使用 Codex 进行自动化研究不是魔法，也不是深度专业知识的替代品。它是一个真正的力量倍增器，让你能够比任何人类手动更快地探索优化搜索空间。
232 倍的提升是真实的、可复现的，并且已部署在生产中。但更重要的是成果是方法论：一个结构化的自动化循环，将 GPU profiler 输出转化为可操作的代码变更，只需最少的人工干预。
如果你正在从事推理优化、数据 pipeline 加速，或任何具有可衡量性能目标的计算密集型问题，这种方法值得认真考虑。
准备好自己尝试了吗？从代码库中最小、最独立的 kernel 开始。先构建验证套件。Profile 后再写提示词。分享你的结果——围绕自动化研究方法论的社区基准还很薄弱，真实数据很有价值。
[INTERNAL_LINK: Getting Started with CUDA Kernel Profiling]
对 pipeline 设置有疑问或想分享自己的结果？欢迎在下方留言或直接联系我。每条回复我都会看。