K-文化的一切——从回归到 K-美妆,发送到您的邮箱订阅邮件

METAL MEDIA

PTXBench: Benchmark and Adapt LLMs for GPU Kernel Optimization with Architecture-specific PTX

arXiv:2608.173792026-08-17

AI能写GPU底层代码,但表现还是半吊子

PTXBench是一个新的评测基准,用来测试AI模型能否直接使用英伟达GPU最底层的可编程接口PTX,在H100和B200显卡上写出高效的计算核心代码。研究团队对多个前沿模型在矩阵乘法和注意力机制任务上进行测试,发现模型在较简单的前向计算中表现不错,但在复杂的反向注意力计算中成功率大幅下降,而且即使模型成功执行了目标底层指令,也不代表速度就快。团队还对开源模型Qwen3.6-27B做了微调实验,发现效果并不稳定,数据的均衡性和教师模型的水平比单纯堆数据量更重要。

METAL MEDIA 解读图

AI能写GPU底层代码,但表现还是半吊子

  1. 01推出PTXBench评测基准,用来衡量大语言模型能否使用PTX(英伟达GPU最底层的可编程指令接口)写出高效的GPU计算核心
  2. 02给模型提供GPU架构知识,让它多轮修改CUDA代码,同时分别检测代码是否正确、目标底层指令是否真正在运行时执行、速度是否超过cuBLAS等顶尖库
  3. 03测试了Gemini 3.1 Pro、Claude Opus 4.8、GLM-5.2和Qwen3.6-27B等模型,发现前向计算任务成功率较高,但复杂的反向注意力任务失败率明显上升,而且执行了目标指令不代表速度就有竞争力
  4. 04用名为Fixit的方法,收集模型的失败案例并配上教师模型给出的修正版本和解释,用来微调Qwen3.6-27B,结果在部分任务上有提升,但对新问题的泛化能力不稳定
  5. 05训练数据是否覆盖多种问题类型、是否均衡,以及负责给出修正答案的教师模型水平,比单纯增加数据量对结果影响更大
这是 METAL MEDIA 制作的解读图,并非论文作者提供的原图。

他们做了什么

  1. 推出PTXBench评测基准,用来衡量大语言模型能否使用PTX(英伟达GPU最底层的可编程指令接口)写出高效的GPU计算核心
  2. 给模型提供GPU架构知识,让它多轮修改CUDA代码,同时分别检测代码是否正确、目标底层指令是否真正在运行时执行、速度是否超过cuBLAS等顶尖库
  3. 测试了Gemini 3.1 Pro、Claude Opus 4.8、GLM-5.2和Qwen3.6-27B等模型,发现前向计算任务成功率较高,但复杂的反向注意力任务失败率明显上升,而且执行了目标指令不代表速度就有竞争力
  4. 用名为Fixit的方法,收集模型的失败案例并配上教师模型给出的修正版本和解释,用来微调Qwen3.6-27B,结果在部分任务上有提升,但对新问题的泛化能力不稳定
  5. 训练数据是否覆盖多种问题类型、是否均衡,以及负责给出修正答案的教师模型水平,比单纯增加数据量对结果影响更大
Figure 1: Overview of the PTXBench benchmark and adaptation workflow.
Figure 1: Overview of the PTXBench benchmark and adaptation workflow.
Table 1: Target-instruction and unrestricted turn correctness rates (%) on H100 and B200.
GPUModelGEMMMHA-FwdMHA-Fwd-CausalMHA-BwdMHA-Bwd-Causal
H100Gemini 3.1 Pro33.3 / 60.4 / 56.2 (33.3 / 62.5 / 59.4)33.3 / 39.6 / 45.8 (33.3 / 39.6 / 45.8)25.0 / 52.1 / 55.2 (25.0 / 52.1 / 55.2)8.3 / 33.3 / 38.5 (8.3 / 33.3 / 38.5)– / 22.9 / 25.0 (8.3 / 25.0 / 26.0)
Claude Opus 4.891.7 / 95.8 / 94.8 (91.7 / 95.8 / 94.8)50.0 / 81.2 / 90.6 (66.7 / 89.6 / 94.8)50.0 / 77.1 / 75.0 (83.3 / 89.6 / 82.3)8.3 / 62.5 / 79.2 (50.0 / 83.3 / 89.6)– / 22.9 / 44.8 (25.0 / 41.7 / 59.4)
GLM-5.233.3 / 60.4 / 61.5 (33.3 / 62.5 / 62.5)16.7 / 8.3 / 8.3 (16.7 / 14.6 / 15.6)– / 8.3 / 10.4 (8.3 / 16.7 / 17.7)8.3 / 6.2 / 5.2 (8.3 / 18.8 / 16.7)– (– / 22.9 / 15.6)
Qwen3.6-27B
B200Gemini 3.1 Pro8.3 / 47.9 / 45.8 (8.3 / 47.9 / 45.8)– / 12.5 / 15.6 (– / 12.5 / 15.6)– / 8.3 / 10.4 (8.3 / 12.5 / 12.5)– / 2.1 / 2.1 (– / 2.1 / 2.1)– (– / 2.1 / 1.0)
Claude Opus 4.825.0 / 64.6 / 80.2 (75.0 / 81.2 / 88.5)– / 20.8 / 38.5 (83.3 / 87.5 / 86.5)– / 16.7 / 32.3 (91.7 / 77.1 / 83.3)– / 6.2 / 10.4 (83.3 / 89.6 / 91.7)– (25.0 / 52.1 / 68.8)
GLM-5.216.7 / 27.1 / 28.1 (33.3 / 37.5 / 34.4)– / – / 1.0 (33.3 / 22.9 / 26.0)8.3 / 2.1 / 1.0 (33.3 / 22.9 / 25.0)– (8.3 / 25.0 / 30.2)– (16.7 / 18.8 / 22.9)
Qwen3.6-27B– (– / – / 1.0)
Figure 2: FastpInst. on H100 (top) and B200 (bottom).
Figure 2: FastpInst. on H100 (top) and B200 (bottom).
Table 2: Model release dates, knowledge cutoffs, and estimated calendar lag from the release of Hopper PTX ISA 8.0 (Dec. 2022) and Blackwell PTX ISA 8.7 (Jan. 2025) [26, 29]. Lags use monthly granularity, from PTX release to disclosed cutoff; otherwise, model release dates provide upper bounds.
ModelModel releaseKnowledge cutoffLag after PTX release (months)
HopperBlackwell
Gemini 3.1 ProFeb. 2026Jan. 202525≈0
Claude Opus 4.8May 2026Jan. 20263712
GLM-5.2June 2026Not disclosed≤42≤17
Qwen3.6-27BApr. 2026Not disclosed≤40≤15
Figure 3: Gemini 3.1 Pro Fastp distributions for Triton and CUDA-PTX on H100 and B200.
Figure 3: Gemini 3.1 Pro Fastp distributions for Triton and CUDA-PTX on H100 and B200.
Table 3: Ablation of architecture-specific prompt knowledge for Gemini 3.1 Pro.
Prompt knowledgeTarget inst. correctness (%) (turn correctness)Target inst. best speedup (best speedup)
Architecture parameters– (50.0 / 29.2 / 26.0)– (0.056 / 0.119 / 0.273)
Architecture parameters + PTX template functions– / 20.8 / 19.8 (– / 20.8 / 19.8)– / 0.542 / 0.542 (– / 0.542 / 0.542)
Architecture parameters + PTX template functions + architecture contract8.3 / 33.3 / 38.5 (8.3 / 33.3 / 38.5)0.206 / 0.375 / 0.515 (0.206 / 0.375 / 0.515)
Figure 4: Training data recipes. Pie area is proportional to record count, and slices show the fraction drawn from each problem. The top shows labels for the checkpoints; the bottom line lists training formats and reasoning teachers.
Figure 4: Training data recipes. Pie area is proportional to record count, and slices show the fraction drawn from each problem. The top shows labels for the checkpoints; the bottom line lists training formats and reasoning teachers.
Table 4: Target-instruction and unrestricted turn correctness rates (%) under SFT and prompt-time supervision.
ConditionMHA-FwdMHA-Fwd-CausalMHA-BwdMHA-Bwd-Causal
Qwen3.6-27B w/o expert guidance
Qwen3.6-27B
Qwen3.6-27B-s1 w/o expert guidance– / – / 4.2 (– / – / 4.2)– / – / 4.2 (– / – / 4.2)– / 4.2 / 3.1 (– / 4.2 / 3.1)
Qwen3.6-27B-s116.7 / 16.7 / 19.8 (16.7 / 16.7 / 19.8)– / – / 3.1 (– / – / 3.1)– / 2.1 / 4.2 (– / 2.1 / 4.2)– / 2.1 / 4.2 (– / 2.1 / 5.2)
Qwen3.6-27B + retrieved repair notes
Qwen3.6-27B + retrieved repair notes and fixed kernel– / 29.2 / 29.2 (– / 29.2 / 29.2)– / 20.8 / 27.1 (– / 20.8 / 27.1)– / 14.6 / 18.8 (– / 14.6 / 20.8)– / 20.8 / 33.3 (– / 25.0 / 37.5)
Figure 5: SFT training data recipe comparison (complete results in Appendix Tables 8 and 9).
Figure 5: SFT training data recipe comparison (complete results in Appendix Tables 8 and 9).
Table 5: Selected SASS families for target instruction measurement.
GPUTagSelected SASS family
H100 (Hopper)H*GMMA tensor compute; UTMALDG, UTMASTG, and UTMAREDG TMA payload transfer
B200 (Blackwell)BUTC*, LDTM, and STTM from the TCGEN05 tensor pathway
Figure 6: Correctness and speedup of Qwen3.6-27B-s1 on the training and held-out problems.
Figure 6: Correctness and speedup of Qwen3.6-27B-s1 on the training and held-out problems.
Table 6: Token counts for architecture-specific prompt components.
GPUComponentTokens
H100Architecture parameter259
Template functions18,601
Architecture contract3,454
Total22,314
B200Architecture parameter413
Template functions18,871
Architecture contract9,531
Total28,815
Figure 7: How Fixit SFT changes reasoning length and error types across turns.
Figure 7: How Fixit SFT changes reasoning length and error types across turns.
Table 7: Best speedup with target instruction execution on H100 and B200.
GPUModelGEMMMHA-FwdMHA-Fwd-CausalMHA-BwdMHA-Bwd-Causal
H100Gemini 3.1 Pro0.687 / 0.934 / 0.962 (0.687 / 0.934 / 0.962)0.555 / 0.730 / 0.730 (0.555 / 0.730 / 0.730)0.614 / 0.651 / 0.768 (0.614 / 0.651 / 0.768)0.206 / 0.375 / 0.515 (0.206 / 0.375 / 0.515)– / 0.634 / 0.639 (0.065 / 0.634 / 0.639)
Claude Opus 4.80.770 / 0.968 / 0.976 (0.770 / 0.968 / 0.976)0.759 / 0.770 / 0.839 (0.759 / 0.770 / 0.839)0.758 / 0.806 / 0.806 (0.758 / 0.806 / 0.806)0.300 / 0.440 / 0.489 (0.300 / 0.440 / 0.489)– / 0.499 / 0.499 (0.058 / 0.499 / 0.499)
GLM-5.20.447 / 0.692 / 0.692 (0.447 / 0.692 / 0.692)0.407 / 0.470 / 0.607 (0.407 / 0.470 / 0.607)– / 0.471 / 0.533 (0.015 / 0.471 / 0.533)0.316 / 0.437 / 0.437 (0.316 / 0.437 / 0.437)– (– / 0.101 / 0.101)
Qwen3.6-27B
B200Gemini 3.1 Pro0.273 / 0.680 / 0.892 (0.273 / 0.680 / 0.892)– / 0.280 / 0.280 (– / 0.280 / 0.280)– / 0.206 / 0.248 (0.013 / 0.206 / 0.248)– / 0.087 / 0.133 (– / 0.087 / 0.133)– (– / 0.015 / 0.015)
Claude Opus 4.80.782 / 1.012 / 1.012 (0.782 / 1.012 / 1.012)– / 0.253 / 0.300 (0.110 / 0.253 / 0.300)– / 0.232 / 0.269 (0.042 / 0.232 / 0.269)– / 0.069 / 0.149 (0.022 / 0.136 / 0.155)– (0.023 / 0.088 / 0.116)
GLM-5.20.162 / 0.632 / 0.632 (0.162 / 0.632 / 0.632)– / – / 0.027 (0.024 / 0.024 / 0.035)0.098 / 0.098 / 0.098 (0.098 / 0.098 / 0.098)– (0.019 / 0.030 / 0.040)– (0.018 / 0.034 / 0.034)
Qwen3.6-27B– (– / – / 0.006)
Figure 8: Cross-language transfer of Fixit SFT from CUDA-PTX to Triton on Hopper.
Figure 8: Cross-language transfer of Fixit SFT from CUDA-PTX to Triton on Hopper.
Table 8: Turn correctness rates for the five-problem SFT evaluation.
SFT-ed Model LabelGEMMMHA-FwdMHA-Fwd-CausalMHA-BwdMHA-Bwd-Causal
Qwen3.6-27B-s0– / 6.2 / 5.2 (– / 6.2 / 5.2)– / 2.1 / 1.0 (– / 2.1 / 1.0)– / – / 1.0 (– / – / 1.0)
Qwen3.6-27B-s125.0 / 14.6 / 13.5 (25.0 / 14.6 / 13.5)16.7 / 16.7 / 19.8 (16.7 / 16.7 / 19.8)– / – / 3.1 (– / – / 3.1)– / 2.1 / 4.2 (– / 2.1 / 4.2)– / 2.1 / 4.2 (– / 2.1 / 5.2)
Qwen3.6-27B-s2– / 2.1 / 2.1 (– / 2.1 / 2.1)– / 2.1 / 5.2 (– / 2.1 / 5.2)– / – / 1.0 (– / – / 1.0)– / – / 1.0 (– / – / 1.0)
Qwen3.6-27B-s3– / 2.1 / 3.1 (– / 2.1 / 3.1)– / 4.2 / 2.1 (– / 4.2 / 2.1)– / 6.2 / 4.2 (– / 6.2 / 4.2)– / 4.2 / 4.2 (– / 4.2 / 4.2)
Qwen3.6-27B-s48.3 / 8.3 / 10.4 (8.3 / 8.3 / 10.4)– / – / 4.2 (– / – / 4.2)– / 4.2 / 2.1 (– / 4.2 / 2.1)– / – / 1.0 (– / – / 1.0)
Qwen3.6-27B-s516.7 / 14.6 / 12.5 (16.7 / 14.6 / 12.5)– / 2.1 / 4.2 (– / 2.1 / 4.2)– / 2.1 / 3.1 (– / 2.1 / 3.1)– / – / 5.2 (– / – / 5.2)– / – / 4.2 (– / – / 4.2)
Qwen3.6-27B-s68.3 / 2.1 / 1.0 (8.3 / 2.1 / 1.0)
Figure 9: Turn-level error-state transitions for GEMM.
Figure 9: Turn-level error-state transitions for GEMM.
Table 9: Best speedups for the five-problem SFT evaluation.
SFT-ed Model LabelGEMMMHA-FwdMHA-Fwd-CausalMHA-BwdMHA-Bwd-Causal
Qwen3.6-27B-s0– / 0.589 / 0.589 (– / 0.589 / 0.589)– / 0.644 / 0.644 (– / 0.644 / 0.644)– / – / 0.209 (– / – / 0.209)
Qwen3.6-27B-s10.303 / 0.303 / 0.340 (0.303 / 0.303 / 0.340)0.556 / 0.556 / 0.565 (0.556 / 0.556 / 0.565)– / – / 0.395 (– / – / 0.395)– / 0.380 / 0.389 (– / 0.380 / 0.389)– / 0.192 / 0.199 (– / 0.192 / 0.199)
Qwen3.6-27B-s2– / 0.209 / 0.211 (– / 0.209 / 0.211)– / 0.548 / 0.548 (– / 0.548 / 0.548)– / – / 0.315 (– / – / 0.315)– / – / 0.296 (– / – / 0.296)
Qwen3.6-27B-s3– / 0.274 / 0.280 (– / 0.274 / 0.280)– / 0.465 / 0.465 (– / 0.465 / 0.465)– / 0.388 / 0.388 (– / 0.388 / 0.388)– / 0.494 / 0.494 (– / 0.494 / 0.494)
Qwen3.6-27B-s40.276 / 0.373 / 0.373 (0.276 / 0.373 / 0.373)– / – / 0.452 (– / – / 0.452)– / 0.383 / 0.383 (– / 0.383 / 0.383)– / – / 0.199 (– / – / 0.199)
Qwen3.6-27B-s50.325 / 0.446 / 0.446 (0.325 / 0.446 / 0.446)– / 0.425 / 0.573 (– / 0.425 / 0.573)– / 0.241 / 0.246 (– / 0.241 / 0.246)– / – / 0.295 (– / – / 0.295)– / – / 0.246 (– / – / 0.246)
Qwen3.6-27B-s60.073 / 0.073 / 0.073 (0.073 / 0.073 / 0.073)
Figure 10: FastpInst. under SFT and prompt-time supervision.
Figure 10: FastpInst. under SFT and prompt-time supervision.
Table 10: Training data recipes.
SFT-ed Model LabelConfigTemplateReasoning SynthesizerRecord Count
Qwen3.6-27B-s08ops-ExtendedKernelGenGLM-5.2494
Qwen3.6-27B-s14opsFixitGLM-5.2158
Qwen3.6-27B-s24ops-ExtendedFixitGLM-5.2259
Qwen3.6-27B-s38ops-ExtendedFixitGLM-5.2406
Qwen3.6-27B-s48ops-Post-balancedFixitGLM-5.2170
Qwen3.6-27B-s58ops-Pre-balancedFixitGLM-5.2258
Qwen3.6-27B-s68ops-Pre-balancedFixitQwen3.6-27B258
Figure 11: Detailed execution infrastructure supporting PTXBench. MiniPTXAgent compiles generated CUDA–PTX locally and sends successfully compiled candidates to an isolated GPU profiling service for sanitization, evaluation, and optional diagnostics. A reliability monitor pauses dispatch, restarts unhealthy service state, and excludes affected turns from trajectories. After a restart, it also excludes thermally abnormal GPUs because throttling distorted latency by over 15% in our measurements [35].
Figure 11: Detailed execution infrastructure supporting PTXBench. MiniPTXAgent compiles generated CUDA–PTX locally and sends successfully compiled candidates to an isolated GPU profiling service for sanitization, evaluation, and optional diagnostics. A reliability monitor pauses dispatch, restarts unhealthy service state, and excludes affected turns from trajectories. After a restart, it also excludes thermally abnormal GPUs because throttling distorted latency by over 15% in our measurements [35].

为什么重要

GPU底层代码优化直接决定AI训练和推理的速度与成本,但手写这类代码需要极其稀缺的专业知识。这项研究首次系统检验了AI是否真能利用最新硬件特性,以及该如何训练才能真正提升这种能力,对相关行业具有参考价值。

Figure 12: Evaluation of workload state caching, generation-profiling pipelining, and GPU sharing.
Figure 12: Evaluation of workload state caching, generation-profiling pipelining, and GPU sharing.

本文术语

  • PTX · 英伟达GPU最底层的可编程指令接口,可直接控制硬件
  • GEMM · 矩阵乘法运算,常用于衡量GPU性能的标准任务
  • SFT(监督微调) · 用带有正确答案的样本继续训练模型的方法
  • LoRA · 一种轻量级微调方法,只更新模型的一小部分参数而非整体
  • SASS · PTX指令在GPU上实际运行时被编译成的英伟达原生机器码

论文原文摘要(英文)

We introduce PTXBench, a benchmark for evaluating and adapting large language models (LLMs) to use architecture-specific PTX for GPU kernel optimization. PTXBench measures functional correctness, whether selected target instructions execute at runtime, and speedup over frontier libraries across GEMM and attention workloads on H100 and B200 GPUs. Our evaluation shows that architecture-specific PTX capability remains uneven: success rates fall substantially on complex attention backward workloads, and executing the target instructions does not necessarily translate into competitive performance. No evaluated model consistently matches frontier libraries across the suite. We further adapt Qwen3.6-27B using supervised fine-tuning. Repair-conditioned training improves several tasks, but generalization remains uneven; data coverage, balance, and the quality of the reasoning teacher matter in addition to dataset size. PTXBench provides an auditable testbed for measuring and improving LLMs' ability to exploit evolving GPU architectures.

作者 · Genghan Zhang

在 arXiv 阅读

最新论文

全部论文 →

METAL MEDIA 最新报道

图片来源: Genghan Zhang et al., arXiv:2608.17379, CC BY 4.0