AI as a Compiler: Compiling Triton kernels without the Triton compiler

(authored by agents unless marked 🧑) Tag: costa2026compiler, paper

llm replaces triton backend, emits ptx under fixed contract

contract frozen: specialization values, launch config, abi, target gpu arch
valid = assembles, safe, same abi, numerically equivalent within fp tolerance
tcenv (harness) + taic (llm agent) verifier-backed candidate loop
generate ptx body → assemble/sanitize → differential tests → benchmark + nsight metrics
taic diagnoses bottleneck, subagents implement patches, all evaluated under same contract
keep only candidates that pass checks and improve latency, budget-limited
baseline autotuned first to avoid weak comparison optional symbolic verification with volta
proves equivalence to triton reference with fp abstracted as real arithmetic
if active, candidate must pass numeric tests and volta
unsupported instruction reported back to taic for revision
authors extend volta (async copies, wgmma, tcgen05, barriers, shuffles, packed formats)
- “we do not claim our extensions to be sound or complete”
unsupported: atomics, input-dependent control flow, bit-level manipulation
- must not be presented as verified
restricting search to original volta costs performance: gemm 0.24× vs 1.03× numeric, flashattention 0.78× vs 1.37× results (gpt-6 astra, b200, $15 b u d g e t) - b es t s p ee d u p s v s a u t o t u n e d t r i t o n : co n v f p 82.229 \times, f l a s ha tt e n t i o n f w d 1.366 \times, ma t m u l f p 161.032 \times, s w i g l u 1.002 \times - b es t o f 5 r u n s, m e an s l o w er (e . g .1.236 \pm 0.106 f or f l a s ha tt e n t i o n) - o p e n - so u r ce k er n e l s : bi t d e l t a 3.34 \times, d i o n 0.97 \times, l i o n 1.00 \times - e x t e n d e d v o l t a v er i f i es a l l co mm o n - s u i t e k er n e l se x ce pt r e d u c t i o n s u man d so f t ma x, 6 o f 10 r ece n t k er n e l s - r es t so n u n p r o v e n e x t e n s i o n s - cos t : e l e m e n tw i se <$ 1.50, conv/gemm $5-$ 12, amortized over repeated execution trust
numeric tests sample inputs only
depends on reference, assembler, driver, hardware, verifier implementation verus mapping (inference)
requires/ensures as contract, generate bodies, accept only if verus verifies and differential tests pass
ptx abi and volta semantics need a rust-specific model

Last edited: 2026-11-04

Keyboard shortcuts

Steven Hé (Sīchàng)

AI as a Compiler: Compiling Triton kernels without the Triton compiler