PTXBench: Benchmark and Adapt LLMs for GPU Kernel Optimization with Architecture-specific PTX
Genghan Zhang, Yixin Dong, Chengze Fan, Zhichen Zeng, Yueming Yuan, Shaowei Zhu, Kunle Olukotun
cs.CL, cs.AI
2026-08-18
斯坦福用运行时SASS核查测LLM写架构专用PTX。Opus 4.8在H100前向正确率高，反向注意力掉队；B200上常写出正确但不走新指令的核。没有模型全套打平厂商库。
KernelBench一类基准问的是：模型能不能写出比PyTorch更快的正确GPU核。更快可能来自普通CUDA，也可能来自直接调用cuBLAS。新一代GPU的价值在GMMA、TMA、Blackwell的TCGEN05这些指令上。开发者真正要问的是：模型会不会按指定架构把这些指令写进核，并且在评测负载里真的执行到。
PTXBench把一道题钉死成：参考实现（cuBLAS / cuDNN / FlashInfer）+ 固定负载 + 目标GPU + 必须用的PTX族。模型从零写带inline PTX的CUDA，禁止include厂商库。每条轨迹共享一份2万到3万token的架构知识包（参数、指令包装、同步与内存契约）。MiniPTXAgent最多8轮：先在CPU容器里过nvcc，再做正确性（atol=rtol=1e-2）和CUPTI计时。
「目标指令正确」比功能正确更严：SASS里要出现指定族，且Nsight Compute给出正的谓词使能线程数，避免死代码里的假阳性。Hopper认GMMA或UTMA，Blackwell认TCGEN05，光有从Hopper继承的TMA不算。指标是按轮次的正确率和Fastp（正确且加速比超过p的轮次占比），每题3个提示×4条8轮轨迹（N=32）。评测模型包括Gemini 3.1 Pro、Claude Opus 4.8、GLM-5.2、Qwen3.6-27B，卡是H100和B200。
能力很不均匀。H100上8轮，Opus 4.8的GEMM目标指令正确率94.8%，多头注意力前向90.6%，因果反向掉到44.8%。Qwen3.6-27B在Hopper上没有写出任何正确核。B200上差距换成另一种：Opus反向注意力功能正确率91.7%，目标指令正确率只有10.4%，模型会写对，但不走新架构路径。没有模型在整套题目上稳定打平厂商库。Gemini在Blackwell GEMM上到0.892× cuBLAS，Opus到1.012×，注意力仍远差于Hopper。
Triton在新架构上更稳。Hopper上CUDA-PTX与Triton接近，因果前向峰值0.768×对0.759×；Blackwell两个反向负载上，Triton是0.484×和0.436×，直接写PTX只有0.133×和0.015×。去掉知识包，Gemini 8轮正确率还有26.0%，目标指令成功率为零；补上模板和架构契约后，目标指令正确率到38.5%。
用Qwen3.6-27B做LoRA，Gemini当修复教师，Fixit配方并不处处压过直接生成。数据覆盖和类别平衡比单纯加条数更重要，推理教师换成学生自己，五道题只剩GEMM能做。最小的全覆盖配方s1能迁到部分d=64和d=96前向，d=96反向和GQA仍是零；迁到Triton后正确率下降，但因果前向峰值从0.238×升到0.632×。
「核是对的」和「用上了这代硬件」是两件事。做GPU核智能体，该把运行时指令探针和厂商库加速比拆开报，否则排行榜会被普通CUDA或暗中调库撑起来。知识包不是可选项：不喂契约，模型几乎不会主动用新PTX。对27B做修复式SFT能打开能力，但配方不平衡就会掉任务，不能当成通用训练处方。
适应实验只有一个27B加小LoRA集，换到工业级训练或其他模型族，修复条件、平衡、教师质量的效应可能变号。题目目前限于H100/B200上的BF16 GEMM和注意力，覆盖不了完整算子面。正确性阈值1e-2偏松，数值敏感核可能被放过去。目标指令探针只证明指令跑到了，不证明它贡献了加速。