PTXBench: Benchmark and Adapt LLMs for GPU Kernel Optimization with Architecture-specific PTX 요약 설명

요약
2 PTXBENCH
- 2.1 BENCHMARK TASKS AND CONTROLLED CONTEXT
- 2.2 MULTI-TURN EVALUATION PROTOCOL
3 ADAPTING LLMS TO ARCHITECTURE-SPECIFIC PTX PROGRAMMING
4 BENCHMARK RESULTS
- 4.2 ARCHITECTURE-SPECIFIC PTX CAPABILITY REMAINS UNEVEN
- 4.4 EXPLICIT ARCHITECTURE KNOWLEDGE IMPROVES TARGET INSTRUCTION EXECUTION
5 ADAPTATION RESULTS
- 5.1 TRAINING FORMAT AND DATA RECIPE RESULTS
- 5.2 EFFECTS AND GENERALIZATION OF FIXIT SFT
부록
짧은 생각

이번 글에서는 PTXBench: Benchmark and Adapt LLMs for GPU Kernel Optimization with Architecture-specific PTX 논문의 핵심 포인트만 간단히 정리한다.

2026년 8월 18일(Arxiv), arXiv
Zhang, Genghan, Dong, Yixin, Fan, Chengze, Zeng, Zhichen, Yuan, Yueming, Zhu, Shaowei, Olukotun, Kunle.
Stanford University, Carnegie Mellon University, RadixArk, Independent Researcher
논문 링크
Github
Project Page

요약

PTXBench는 H100(Hopper)과 B200(Blackwell)에서 LLM이 특정 아키텍처 전용 PTX 명령어 계열을 포함한 CUDA–PTX kernel을 생성하고 활용하는 능력을 평가·적응하기 위한 benchmark다. 기존 GPU kernel benchmark가 기능 정확도와 종합 효율을 주로 평가한 것과 달리, 이 연구는 고정 workload에서 요구한 명령어 계열이 runtime에 실제 실행되는지를 별도 조건으로 검증하고 frontier library 대비 speedup도 측정한다. MiniPTXAgent는 통제된 아키텍처 지식, kernel 생성, 컴파일·실행 피드백, 수정을 여러 turn에 걸쳐 연결한다. 저자들은 forward workload에서는 성공 사례가 나타나지만 복잡한 attention backward에서는 성공률이 낮고, target instruction 실행 자체가 경쟁력 있는 성능을 보장하지 않는다고 보고한다. 또한 Qwen3.6-27B에 Fixit repair-conditioned SFT를 적용해 일부 task를 개선했으나, held-out shape·attention variant로의 일반화는 고르지 않으며 data coverage·balance와 reasoning teacher의 품질이 중요하다고 결론짓는다.
- 평가 대상은 BF16 GEMM과 네 가지 MHA workload이며, GEMM의 기준선은 cuBLAS v13.1.0, 일반 attention의 기준선은 cuDNN v9.20.0, GQA의 기준선은 FlashInfer v0.6.14다.
- PTXBench는 일반 CUDA 코드나 vendor library 호출로 성능을 내는 경우와 구분하기 위해, 지정된 inline PTX 명령어 계열이 평가 workload에서 runtime에 실행되는지를 능력 조건으로 둔다.

2 PTXBENCH

PTXBench의 각 instance는 cuBLAS·cuDNN·FlashInfer 등의 frontier library로 구성한 reference operator, 고정 workload, target GPU architecture, 요구하는 architecture-specific PTX instruction family를 결합한다. 모델은 초기 구현을 고치는 방식이 아니라 inline PTX를 포함한 CUDA kernel을 처음부터 생성한다. workload·solution·correctness check는 FlashInfer-Trace schema로 표현하므로, 호환되는 task collection에도 같은 workflow를 적용할 수 있다. benchmark는 architecture parameter, PTX instruction용 CUDA wrapper, layout·synchronization·memory consistency contract와 일부 연산자의 전문가 검증 scheduling 원칙을 prompt에 제공해, 외부 문서 검색보다 주어진 저수준 제약을 정확히 해석·조합하는 능력을 측정한다.
- 아키텍처별 knowledge pack은 20k–30k tokens를 차지하며, Table 3 ablation은 이 문맥이 없을 때 평가 모델이 요청된 PTX를 실행하지 못했음을 보인다.
- 핵심 측정 축은 기능 정확도, target instruction의 runtime 실행, frontier library 대비 효율성이다.

PTXBench benchmark와 adaptation workflow의 개요

2.1 BENCHMARK TASKS AND CONTROLLED CONTEXT

통제된 문맥을 위해 모든 trajectory는 동일한 architecture-specific knowledge pack을 base prompt로 받는다. pack에는 hardware parameter, PTX wrapper, layout·synchronization·memory consistency contract가 포함되며, FlashAttention처럼 잘 연구된 연산자에는 고정된 전문가 검증 scheduling 원칙도 제공한다. 이는 모델이 외부 문서를 찾는 능력보다 같은 명세 아래에서 architecture-specific PTX를 올바르게 조합하는 능력을 비교하기 위한 설계다. Table 6에 따르면 H100 pack은 22,314 tokens이고 B200 pack은 28,815 tokens이며, B200에는 shared-memory descriptor, peer CTA, tensor memory 등을 위한 추가 contract가 포함돼 더 길다.
- reference는 cuBLAS·cuDNN·FlashInfer 등으로 구성하지만, 생성 source가 cuDNN 또는 cuBLAS header를 포함하면 실행 결과와 무관하게 오답으로 처리한다.
- 모든 모델이 같은 long-context 명세를 받으므로, 결과에는 제공된 아키텍처 지식을 해석해 kernel에 반영하는 능력이 포함된다.

2.2 MULTI-TURN EVALUATION PROTOCOL

MiniPTXAgent는 고정된 호출 횟수 안에서 이전 kernel과 execution feedback을 보존하며 수정한다. 먼저 CPU container에서 `nvcc`로 candidate를 컴파일하고, 컴파일에 성공한 candidate만 격리된 profiling service로 보내 memory safety, runtime error, 기능 정확도, latency를 검사한다. 기능 정확도는 output shape·data type을 검사한 뒤 `torch.allclose`의 `atol=rtol=1e-2`로 값을 비교한다. 정확한 candidate의 latency는 10회 warmup 후 50회 timed iteration의 중앙값으로 측정하며, speedup은 reference latency를 candidate latency로 나눈 값이다.
- target instruction correctness는 기능 정확도와 runtime target instruction 실행을 모두 만족하는 조건이다. Hopper에서는 GMMA compute 또는 UTMA payload movement를, Blackwell에서는 TCGEN05 tensor path를 대상으로 하며, Hopper에서 계승된 TMA만 사용한 Blackwell kernel은 인정하지 않는다.
- 먼저 최종 SASS에 대상 명령어 계열이 있는지 정적으로 검사하고, 존재할 때만 Nsight Compute로 matching instruction의 predicate-enabled thread count가 양수인지 동적으로 검사한다. 따라서 dead code나 launch되지 않은 kernel에만 대상 명령어가 있는 경우를 제외한다.
- 이 지표는 지정 명령어가 실행됐음을 뜻할 뿐, 그 명령어가 유용한 일을 했는지 또는 측정한 speedup의 원인인지를 증명하지는 않는다.

3 ADAPTING LLMS TO ARCHITECTURE-SPECIFIC PTX PROGRAMMING

Fixit은 적응 대상 모델이 pre-adaptation checkpoint에서 실제로 생성한 실패를 학습 입력으로 바꾸는 repair-conditioned SFT 방식이다. 문제와 architecture-specific context를 포함한 prompt $x$에서 기준 모델 $\pi_0$가 실패 kernel $k^-$를 생성하면, feedback 함수 $H$가 compilation 또는 execution feedback $e$를 수집한다. repair teacher $\pi_R$는 같은 문제·실패 kernel·feedback을 조건으로 수정 kernel을 생성하고, 기능 정확도를 통과한 $k^+$만 남긴다. 이어 reasoning teacher $\pi_T$가 실패에서 repair로 이어지는 rationale $r$을 합성하며, student는 $(x,k^-,e)$를 조건으로 받고 $(r,k^+)$를 목표로 학습한다.
- 실패는 adaptation 전 고정 checkpoint에서 한 번 수집하므로, 학습은 해당 checkpoint가 실제로 방문한 error state를 대상으로 한다.
- repair와 rationale은 교사가 생성하며, correctness-filtered repair만 사용해 오류가 있는 수정안을 감독 신호로 쓰지 않는다.
- Qwen3.6-27B의 LoRA adaptation에서는 최종 assistant message만 최적화한다. 앞선 prompt·실패 kernel·feedback은 조건 문맥으로 고정되고, 최종 assistant message에만 supervised loss의 gradient가 흐른다.

4 BENCHMARK RESULTS

평가는 turn correctness rate, target instruction turn correctness rate, correct turn 중 best speedup, target instruction을 실행한 correct turn으로 한정한 best speedup을 함께 사용한다. `Fast_p`는 모든 평가 turn 가운데 correct kernel이 speedup threshold $p$를 넘는 비율이고, `Fast^{Inst.}_p`는 여기에 runtime target instruction 실행 조건을 추가한 비율이다. 실선은 세 prompt variant의 평균이며, variant마다 네 개의 8-turn trajectory를 수행해 총 $N=32$ turns를 평가하고 음영은 prompt별 최소·최대를 나타낸다. NCU profiling이 실패한 turn은 target instruction 분자에 넣지 않아 해당 수치는 보수적으로 보고된다.
- Table 1은 target-instruction 지표를 첫 줄에, 제한 없는 correctness를 괄호 줄에 제시하며, 각 세 수는 각각 ≤1, ≤4, ≤8 turns의 누적 결과다.
- Figure 2와 Appendix Table 7에서 B200 GEMM의 8-turn target-instruction best speedup은 Claude Opus 4.8이 1.012×, Gemini 3.1 Pro가 0.892×다. Table 2는 Gemini 3.1 Pro의 knowledge cutoff가 Blackwell PTX ISA 8.7 공개와 같은 달이라고 제시하지만, 이 수치만으로 knowledge cutoff가 성능을 결정한다고 판단할 수는 없다.
- Gemini 3.1 Pro의 Triton과 CUDA–PTX를 비교한 Figure 3에서 H100 causal MHA forward의 peak은 CUDA–PTX가 0.768×, Triton이 0.759×다. 반면 B200의 두 backward workload에서는 Triton이 0.484×·0.436×, CUDA–PTX가 0.133×·0.015×다. 이는 동일 모델·언어·workload·GPU 조건에서의 비교이며, 저자들은 새 아키텍처에서 고수준 kernel language가 correct kernel 생성에 여전히 유용하되 최고 성능을 보장하지는 않는다고 해석한다.

H100 및 B200의 target-instruction·제한 없는 turn correctness rate

모델 release·knowledge cutoff와 PTX ISA 공개 이후 경과 기간

H100 및 B200에서 Gemini 3.1 Pro의 Triton·CUDA–PTX Fast 분포

4.2 ARCHITECTURE-SPECIFIC PTX CAPABILITY REMAINS UNEVEN

architecture-specific PTX 능력은 모델·GPU·workload에 따라 고르지 않았다. 저자들은 Claude Opus 4.8이 Blackwell에서 Gemini 3.1 Pro보다 substantially higher target-instruction correctness rate를 보였다고 서술하며, 이 비교는 B200의 target-instruction correctness 지표에 한정된다. Table 1의 B200 GEMM 8-turn 값도 Claude Opus 4.8은 80.2%, Gemini 3.1 Pro는 45.8%로 이를 뒷받침한다. 같은 workload의 target-instruction best speedup은 각각 1.012×와 0.892×이며, 이 값은 GEMM·B200·qualifying kernel 조건의 peak 비교다. GLM-5.2는 H100·B200 GEMM과 H100 non-causal attention에서 Gemini 3.1 Pro와 경쟁적이지만 B200 attention에서는 뒤처졌고, Qwen3.6-27B는 B200에서 correct GEMM kernel 하나를 만들었으나 selected Blackwell instruction은 실행하지 못했으며 H100 workload에서는 correct kernel을 만들지 못했다.
- Figure 2의 speedup distribution은 복잡한 attention backward일수록 qualifying kernel 생성과 speedup이 제한되는 양상을 보인다. 저자들은 target instruction 실행, correct kernel 생성, 경쟁력 있는 성능 달성을 구분되는 능력으로 해석한다.
- Claude Opus 4.8과 GLM-5.2는 새 Blackwell architecture를 대상으로 할 때 architecture-specific PTX 대신 generic CUDA로 후퇴하는 경향을 보였다고 저자들은 설명한다.
- Gemini 3.1 Pro는 Blackwell instruction 실행 성공률이 Claude Opus 4.8보다 낮음에도 일부 workload에서 가까운 peak speedup을 보였다는 저자 서술이 있다. 다만 각 workload의 peak 수치는 별도로 해석해야 하며, 이를 전체 성능의 동등성으로 일반화할 수는 없다.

4.4 EXPLICIT ARCHITECTURE KNOWLEDGE IMPROVES TARGET INSTRUCTION EXECUTION

Gemini 3.1 Pro와 세 prompt 구성으로 명시적 architecture knowledge를 ablation한 결과, architecture parameter만 준 조건은 8 turns에서 turn correctness 26.0%였지만 target instruction success는 없었다. PTX template function을 추가하면 target instruction 실행이 나타났고, template function과 architecture contract를 모두 넣은 조건의 8-turn target-instruction correctness는 38.5%였다. 이는 template function 조건의 19.8%보다 18.7%p 높다. 그러나 8-turn target-instruction best speedup은 contract 포함 조건이 0.515×, template function 조건이 0.542×이므로, 저자들은 contract의 주된 효과를 가장 빠른 kernel 생성이 아니라 target instruction의 신뢰성 있고 correct한 실행으로 한정한다.
- Table 3에서 contract 포함 조건의 ≤1/≤4/≤8-turn target-instruction correctness는 8.3%/33.3%/38.5%이고, target-instruction best speedup은 0.206×/0.375×/0.515×다.
- architecture parameter와 template function만 넣은 조건의 대응 target-instruction correctness는 –/20.8%/19.8%, target-instruction best speedup은 –/0.542×/0.542×다. 따라서 8-turn peak speedup만으로 contract 포함 조건이 더 빠르다고 말할 수 없으며, 원문 주장은 correctness 우위에 관한 것이다.

Gemini 3.1 Pro의 architecture-specific prompt knowledge ablation

5 ADAPTATION RESULTS

adaptation 실험은 long PTX prompt와 kernel trace를 수용하는 262K-token context 및 통제 가능한 LoRA scale을 이유로 Qwen3.6-27B를 대상으로 하고, Gemini 3.1 Pro를 repair teacher로 사용한다. KernelGen은 원래 문제 prompt에서 Gemini 3.1 Pro가 direct solution을 만들고 GLM-5.2가 retained correct kernel의 rationale을 합성하는 기준선이다. Fixit recipe는 s6를 제외하고 GLM-5.2를 reasoning teacher로 사용하며, s6는 Qwen3.6-27B가 자기 repair rationale을 합성한다. 연구는 training format, problem coverage·balance, reasoning-teacher 선택을 분리해 평가하고 held-out task 및 prompt-time supervision과도 비교한다.
- 모든 SFT model은 `Qwen/Qwen3.6-27B`에서 독립적으로 시작하며 LoRA rank 32, 5 epochs, batch size 2, learning rate $4.65\times10^{-4}$의 선형 schedule을 사용한다. Adam은 $\beta_1=0.9$, $\beta_2=0.95$, $\epsilon=10^{-8}$이며, 65,536 tokens를 넘는 example은 Tinker API token limit 때문에 제거한다.
- Figure 4와 Appendix Table 10은 recipe별 problem mix, format, rationale synthesizer, filtering 뒤 record count를 제시한다. s1은 4ops Fixit·GLM-5.2 rationale·158 records이고, s5와 s6는 같은 8ops-Pre-balanced Fixit 구성 및 258 records를 사용하되 reasoning synthesizer만 다르다.
- Table 4와 Figure 10의 prompt-time supervision 비교에서는 error trajectory에서 Codex가 정제한 expert guidance와, BM25로 선택한 repair 경험의 GPT 5.4 요약을 사용한다. retrieval 조건은 repair note만 제공하거나 corrected kernel까지 함께 제공하는 경우를 구분한다.

SFT training data recipe의 problem 구성·format·reasoning teacher

5.1 TRAINING FORMAT AND DATA RECIPE RESULTS

동일한 여덟 problem class, GLM-5.2 reasoning synthesis, 비슷한 record 수를 맞춘 KernelGen s0와 Fixit s3를 비교하면, 8 turns에서 s3는 GEMM·MHA-Fwd-Causal·MHA-Bwd correctness를 개선했지만 MHA-Fwd·MHA-Bwd-Causal에서는 s0보다 낮았다. 따라서 pre-adaptation checkpoint에서 한 번 수집한 failure를 조건으로 하는 Fixit이 일부 task에는 도움이 되지만 direct-solution supervision을 일관되게 앞서지는 못한다. s3의 더 긴 reasoning도 kernel performance를 안정적으로 예측하지 못했다. recipe 간에는 단순 record 수보다 coverage와 balance가 중요했으며, 상대적으로 균형 잡힌 s1과 s5만 다섯 evaluation problem 모두에서 correct kernel을 생성했다.
- 원문은 s2가 s1보다 1.6× 많은 records이고 s3가 s4보다 2.4× 많지만, 두 recipe 모두 MHA-Bwd-Causal을 풀지 못했다고 보고한다. 반대로 s4에서 1.5× 더 큰 balanced s5로 바꾸면 8-turn correctness는 네 문제에서 개선되고 MHA-Fwd에서는 같았으며 peak speedup은 네 문제에서 향상됐다.
- 고정된 Fixit example에서 reasoning synthesizer만 GLM-5.2인 s5와 Qwen3.6-27B인 s6로 바꾼 비교에서 s5는 다섯 문제를 모두 풀었지만 s6는 GEMM만 풀었다. 저자들은 target-model failure가 유용한 입력이지만 repair rationale 생성에는 더 강한 teacher가 도움이 된다고 해석한다.

SFT training data recipe별 correctness·target instruction·speedup 비교

5.2 EFFECTS AND GENERALIZATION OF FIXIT SFT

세부 분석용 s1은 head dimension 128의 네 MHA task로 학습한, 다섯 evaluation problem을 푼 가장 작은 recipe다. Figure 6에서 s1은 GEMM, 모든 d64 MHA task, 두 d96 forward task로 일부 전이하지만 d96 backward 두 task와 GQA에서는 correct kernel을 만들지 못했다. 저자들은 d96이 H100 WGMMA의 m64 tile과 맞지 않는 점을 d96 backward가 특히 어려운 이유로 든다. CUDA–PTX에서 학습한 s1을 같은 Hopper workload의 Triton 생성에 전이하면 모든 workload에서 turn-level correctness는 base보다 낮지만, causal MHA forward의 best correct speedup은 0.238×에서 0.632×로, causal MHA backward는 0.043×에서 0.331×로 상승했다. 따라서 이 조건에서는 correct kernel을 만들 가능성은 낮아졌으나 일부 causal variant의 peak 성능은 높아졌다.
- Figure 7과 Figure 9에서 s1은 매 turn의 reasoning을 더 길게 만들고 실패를 compilation error에서 runtime·numerical error 쪽으로 이동시킨다. GEMM에서 s1은 turn 0에 correct kernel 세 개를 생성하고 뒤의 일곱 turn 중 여섯 turn에서 적어도 하나를 생성했으나, base model은 어떤 turn에서도 correct kernel을 만들지 못했다.
- SFT와 prompt-time supervision을 비교하면, expert guidance가 있어도 base Qwen3.6-27B는 correct MHA kernel을 만들지 못한 반면 s1은 guidance 없이도 일부 MHA kernel을 만들었다. repair note만 retrieval한 base model도 correct kernel을 만들지 못했고, fixed kernel까지 함께 제공할 때 correctness가 상승했다. 저자들은 답이 prompt에 거의 포함되는 solution-bearing 조건이므로 후자의 상승을 독립적 추론 능력 향상으로 해석하지 않는다.
- Table 4에서 expert guidance를 준 Qwen3.6-27B-s1의 MHA-Fwd target-instruction correctness는 ≤1/≤4/≤8 turns에 16.7%/16.7%/19.8%였고, MHA-Fwd-Causal·MHA-Bwd·MHA-Bwd-Causal도 8 turns에서 각각 3.1%·4.2%·4.2%의 nonzero target-instruction correctness를 기록했다.
- 연구 범위는 단일 27B base model, modest LoRA dataset, H100·B200의 BF16 GEMM·attention에 한정된다. 따라서 repair conditioning·data balance·teacher quality 효과가 다른 model family나 industry-scale post-training에 그대로 이전되는지는 이 논문만으로 확인할 수 없다.

Qwen3.6-27B-s1의 training·held-out problem correctness와 speedup

Fixit SFT 전후 turn별 reasoning 길이와 workload별 오류 유형

Hopper에서 CUDA–PTX Fixit SFT의 Triton cross-language transfer

SFT와 prompt-time supervision 조건의 turn correctness rate

SFT와 prompt-time supervision 조건의 target-instruction Fast 분포

부록

부록은 H100 80 GB와 B200에서의 컴파일·trajectory·추론·SFT 설정, SASS와 Nsight Compute를 결합한 target instruction execution 측정 및 profiling service의 구성, architecture-specific prompt token 수를 설명한다. 또한 controlled context에 사용한 FlashAttention-3 pseudocode의 두 불일치와 수정 사항, benchmark·SFT의 보조 결과를 제시한다.

짧은 생각

runtime instruction 검증을 static SASS 존재 확인에서 멈추지 않고 predicate-enabled thread count 확인으로 이어 간 설계는, 명령어가 source나 binary에 포함된 사실과 평가 workload에서 실행된 사실을 구분한다는 점에서 중요하다. 부록의 미실행 dummy kernel 사례는 정적 검사가 왜 충분하지 않은지를 보이며, undefined behavior로 runtime profiling을 분류할 수 없는 turn을 target-instruction 분자에서 보수적으로 제외한 처리도 지표의 해석 범위를 분명하게 한다. 다만 Blackwell의 `UTC*` 매칭은 TCGEN05 pathway의 control instruction까지 넓게 포괄하므로, target-instruction success를 fifth-generation MMA가 유의미한 작업을 수행했다는 좁은 주장으로 읽을 수는 없다. 앞으로는 동일한 correct kernel에서 해당 명령어의 dynamic count나 instruction-level 기여까지 대조한다면, 실행 여부와 성능 향상의 연결을 더 직접적으로 평가할 수 있을 것이다.

Gorio Tech Blog