Vibecoding GPU Kernels

@maharshii
英语2026年8月09日
118K
408
30
11
593

TL;DR

Maharshi 介绍了一种利用 LLM 生成 GPU kernel 的工作流,该工作流通过编译检查、参考实现对比测试以及迭代优化,构建了一个可验证的反馈循环。

手写 GPU 内核需要耐心、努力,还得搭上你的理智。我一直在尝试用现代 LLM(Claude Opus 5、GPT 5.6 Sol)直接帮我把内核生成出来。我想聊聊我从这件事中学到的东西。

说实话,编写 GPU 内核是一个验证起来出奇容易的任务。

首先,你要确定要为哪些操作编写内核。其次,写出第一个版本,确保它能编译通过、没有明显错误。第三步,用一个速度较慢的参考实现来检查正确性。如果对不上,就尝试修复正确性问题。完成之后,再对内核的执行时间做基准测试。后续的版本都建立在这个流程之上,不断优化,直到达到 Roofline 指标或者自己满意为止。

maharshi - inline image

可验证循环:GPU 内核开发

上图把这个过程展示成一个可验证循环:

  • 编译检查是一个紧凑的本地循环(B <-> C),在考虑正确性之前就要先跑通。
  • 正确性检查(D <-> E <-> F)是核心的可验证回报循环。正是这部分让内核编写成为自动化验证的理想问题,因为你有一个可以对照的基准真值(ground truth)。
  • 优化(G -> H -> 回到 D)对每个新版本都复用同一套正确性循环,因为一个快但错误的内核毫无价值,每个版本都必须验证正确性。
  • Roofline/满意度检查(I)是外层循环,决定是继续优化还是就此打住。

第一个版本

上图仍然是高层视角,魔鬼藏在细节里。我们需要确保我们的 LLM Agent 拥有所需要的全部上下文,才能开始写一个好的初版。

这时候,CUDA DSL 就登场了。Triton、CuTeDSL 和 Tilelang 都是非常容易上手、直接用 Python 就能写的 DSL。相比 CUDA C++,它们的学习曲线要平缓不少,但这些 DSL 的抽象概念反而更容易把我们的 Agent 搞糊涂。我们需要一种方法,把这些抽象概念作为上下文传给 Agent。

现代 LLM 已经知道怎么写“好”的 Triton 了。即使没有任何上下文,它们也能很好地处理 Triton 的抽象概念。不过,对于像 CuTeDSL 这样提供比 Triton 更多控制力的 DSL,我发现在上下文目录里准备一份 Agent 可以查阅的 DSL 抽象说明,会带来很大帮助。

例如,在上下文目录里克隆

NVIDIA cutlass

仓库就是一个好办法,可以让 Agent 在编写 CuTeDSL 内核时查找与

Layout Algebra、Copy/GEMM atoms、内存层次结构、示例内核

等相关的抽象概念。

以我的经验,一个好的初版内核应该能顺利编译、没有明显错误,并且通过下面要讲的正确性测试。

测试、基准测试与性能分析

在给 Agent 足够上下文之后,真正的瓶颈就转移到了验证上。参考实现本身,以及对照它进行验证,变得越来越重要。我把这个阶段称为正确性测试,或者干脆叫测试。参考实现的速度并不重要,重要的是它的意图。你打算测量和验证什么,你的 Agent 就会去优化什么。

通常,如果计算不需要在低精度下进行,我会测量最大绝对/相对误差(MAE)、均方误差(MSE/RMSE)和PSNR(峰值信噪比)。当涉及低精度时,我倾向于测量PSNR余弦相似度(cossim)。

如何让内核版本真正在 GPU 上跑起来,取决于 GPU 是在本地还是通过云端访问。不管怎样,我们的 Agent 都应该能以某种方式访问到它的输出。

我觉得下面这个 rung 方法 是组织 N 个测试函数的好办法:

python
1def rung(name):
2 def deco(fn):
3 try:
4 out = fn()
5 results[name] = {"ok": True, **(out or {})}
6 print(f"[{name}] ok " + " ".join(f"{k}={v}" for k, v in (out or {}).items()))
7 except Exception as e:
8 results[name] = {"ok": False, "err": f"{type(e).__name__}: {e}"}
9 print(f"[{name}] FAILED {type(e).__name__}: {e}")
10 traceback.print_exc()
11 return fn
12 return deco

可以这样调用:

python
1out = {}
2
3@rung("pre-checks")
4def _():
5 run_pure_checks()
6 run_dsl_checks()
7
8@rung("run")
9def _():
10 out["o"] = custom_kernel(*inputs)
11 torch.cuda.synchronize()
12 return {"shape": tuple(out["o"].shape),
13 "finite": bool(torch.isfinite(out["o"]).all())}

在 benchmark rung 里,你可以做几件事:

  • 对内核做端到端执行时间的基准测试,统计整体耗时
  • 使用内核内追踪(intra-kernel tracing)来基准测试内核内部的各个部分,并把结果 dump 到输出中(使用自定义 tracer 或 CUPTI)
  • 把生成的 IR、PTX、SASS 和 CUBIN dump 到 dumps 目录,让 Agent 自己去翻

最后一点可以再展开说说。有时候,DSL 在降级(lower)后可能会生成次优的 PTX(最终是 SASS),而你能找到更好的指令或 shape 来替代。我们的 Agent 可以通读 PTX/SASS 文本文件,直接内联更底层的代码,而不是让 DSL 去处理那个次优的部分。同样,把“可搜索的”PTX 文档作为上下文传进去,在这里也非常有帮助。

最后,把所有环节串起来的是性能分析(Profiling)。如果你的 Agent 能访问 NCU(Nsight Compute Systems)CLI,你可以让它对内核做性能分析并生成报告,作为上面那个可验证反馈循环的一部分。

写在最后

“那 GPU 内核开发是不是已经死了?”

“嗯,是,也不是。”

说是,是因为布局、索引

一键保存

使用 YouMind AI 深度阅读爆款文章

保存原文、追问细节、总结观点,并在一个 AI 工作空间里把爆款文章沉淀成可复用笔记。

了解 YouMind
写给创作者

把你的 Markdown 变成干净的 𝕏 文章

图片上传、表格、代码块,往 𝕏 上手动重排太痛苦。YouMind 把整篇 Markdown 一键转成干净、可直接发布的 𝕏 文章草稿。

试试 Markdown 转 𝕏

更多可拆解样本

近期爆款文章

探索更多爆款文章