PaperlayerView exact revision on arXiv

Source-linked AI summary

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

François Costa, Charly Castes, Thomas Bourgeat, Azalia Mirhoseini

arXiv:2609.36800v1cs.AI

TL;DR

Compiler backends are costly to build and maintain as hardware and programming models evolve. TAIC uses an LLM to generate Triton-to-PTX code with evaluation feedback; across twelve common kernels, it reaches 0.83×–2.23× Triton performance and also beats Triton on recent kernels.

  • Problem

    Traditional compiler backends are costly to build and maintain as workloads and accelerators evolve, while new hardware features can exceed existing compiler abstractions.

    Source: §1 · p. 1

  • Method

    TAIC and TCENV combine LLM-based PTX generation with compilation feedback, correctness testing, benchmarking, profiling, and iterative optimization.

    Source: §1 · p. 3

  • Results

    Generated PTX achieves 0.83×–2.23× Triton performance across twelve kernels on Ada, Hopper, and Blackwell, and beats recent kernels including BitDelta (3.34×) on B200.

    Source: §1 · p. 3

  • Takeaways & Limitations

    The results suggest direct high-level-to-low-level AI compilation could reduce the time and engineering effort needed to bring up software for new chips.

    Source: §1 · p. 3

    The extended verifier does not support atomics, input-dependent control flow, or bit-level value manipulation.

    Source: §4.5 · p. 9

Abstract

from arXiv · show

Compiler backends are expensive to build and maintain as programming models, workloads, and accelerators evolve. We investigate whether large language models can replace the conventional optimizing and lowering pipeline, a process that we call AI lowering. We study AI lowering from Triton to NVIDIA PTX: an LLM agent translates Triton kernels directly into PTX. We build an environment that evaluates candidate PTX, and an agentic harness in which an LLM translates Triton kernels into PTX. Across twelve common kernels on Ada, Hopper, and Blackwell GPUs and ten kernels from recent ML papers, AI lowering achieves 0.83x-3.34x the performance of autotuned Triton. The largest gains come from transformations that Triton's lowering pipeline does not perform, such as decoding packed binary weights directly into Tensor Core operands (3.34x on BitDelta), assigning each thread a complete softmax row in tensor memory (1.37x on FlashAttention), and reusing overlapping convolution windows (up to 2.23x). These results rely on a robust evaluation harness with comprehensive verification support. We build on Volta, an existing PTX verifier, and substantially extend it to support modern GPU architectures by introducing support for Blackwell's tcgen05 Tensor Core interface. This requires modeling three architectural features: managed tensor memory, descriptor-based operand layouts, and asynchronous execution coordinated through commits, waits, memory barriers, and proxy fences. We discuss the challenges involved in formalizing them, as well as the current limitations. Our results suggest an emerging future in which AI compilers replace custom-written intermediate representations and checkers, reducing the time and engineering effort required to bring up software for new general-purpose and custom chips.

1 INTRODUCTION

Rapidly evolving, diverse accelerators make compiler backends costly to build and maintain. The paper introduces direct Triton-to-PTX generation with an evaluation harness and extends PTX verification to modern GPU features.

1 INTRODUCTION

  • Diverse, rapidly changing accelerator features strain existing compiler abstractions and make conventional IR-centric backend development expensive to maintain.The paper notes that performance-critical implementations sometimes bypass IR abstractions using target-specific intrinsics or assembly.Source: §1 · pp. 1–2
  • The Volta extension adds Hopper and Blackwell support, including modeling Blackwell tcgen05 tensor memory, descriptor-based layouts, and asynchronous execution ordering.The paper identifies formalizing these features and proving verifier soundness as remaining challenges.Source: §1 · p. 2
  • TAIC and TCENV pair LLM-based PTX generation with compilation feedback, correctness testing, benchmarking, profiling, and iterative optimization under the kernel contract.TAIC generates PTX directly from Triton source, bypassing the conventional optimizing and lowering backend.Source: §1 · pp. 2–3
  • Generated PTX achieves 0.83×–2.23× Triton performance on the twelve-kernel suite across Ada, Hopper, and Blackwell, and beats recent kernels on B200, including BitDelta (3.34×).Other B200 comparisons include FlashAttention (1.37×) and a Mamba-2 primitive (1.10×).Source: §1 · p. 3
  • The results suggest AI compilers could replace custom-written intermediate representations and reduce the time and engineering effort needed to bring up software for new chips.The paper ties this direction to end-to-end translation from high-level source to low-level code.Source: §1 · p. 3

2 RELATED WORK

The related work spans scheduling search, high-level GPU programming, kernel generation, assembly optimization, and verification. AI lowering differs by generating hardware-specific PTX directly from a fixed Triton program.

2 RELATED WORK

  • Unlike searching a hand-engineered schedule space or selecting a library implementation, AI lowering synthesizes hardware-specific PTX directly from a fixed Triton kernel.The paper contrasts this with systems that separate tensor computations from schedules or compile higher-level programming models.Source: §2 · p. 3
  • This task requires PTX-only output for a fixed, self-contained Triton program and leaves instruction selection unconstrained within the compilation contract.The cited GPU-kernel generation systems instead evaluate CUDA or Triton generation from PyTorch-level tasks, or prescribe architecture-specific PTX instructions.Source: §2 · p. 3
  • Unlike assembly superoptimization or approaches built on explicit intermediate representations, AI lowering starts from Triton and jointly performs instruction selection, allocation, and scheduling in PTX.The passage contrasts these approaches with STOKE, SuperCoder, MLIR, CompilerGym, LLM Compiler, and CAKE.Source: §2 · p. 3
  • Finite test suites are used by KernelBench and related systems, while Volta symbolically executes PTX but excludes Hopper- and Blackwell-introduced instructions.The passage distinguishes testing-based correctness approaches from symbolic PTX verification.Source: §2 · p. 3

3 AN ALTERNATIVE COMPILATION ROUTE

The paper frames direct Triton-to-PTX lowering as finding a fast, correct implementation under a fixed kernel contract, then pairs it with an environment and iterative LLM harness to generate and evaluate candidates.

3 AN ALTERNATIVE COMPILATION ROUTE

  • Triton provides a kernel-level boundary: it specifies computation and interface while leaving instruction selection, thread mapping, data movement, synchronization, and scheduling to the compiler.Starting from PyTorch would add graph-level optimization concerns, while CUDA already exposes many low-level choices explicitly.Source: §3 · pp. 3–4

3.1 PROBLEM FORMULATION

  • Valid PTX lowerings must preserve the fixed specialization, launch interface, ABI, and target architecture while matching Triton's semantics within numerical tolerances.The objective is to minimize latency among programs that compile for the target and satisfy safety, ABI, and correctness requirements.Source: §3.1 · p. 4

3.2 THE TRITON COMPILER ENVIRONMENT: AN ENVIRONMENT FOR AI LOWERING

  • TCENV fixes compilation and launch metadata, preserves the PTX interface, and provides infrastructure to assemble, validate, benchmark, and profile generated PTX.It combines randomized differential tests with optional symbolic verification; its Triton reference kernels are autotuned on the target GPU.Source: §3.2 · pp. 4–5

3.3 THE TRITON AI COMPILER: AN AGENT PERFORMING AI LOWERING

  • TAIC directly translates Triton to PTX and iteratively refines candidates using compilation, correctness, benchmark, sanitizer, and profiler feedback.TCENV evaluates branched patches under the same contract, and only candidates that pass gates and improve the measured objective can seed another round.Source: §3.3 · pp. 5–6

4 EVALUATION

Across GPU generations, TAIC generates architecture-specific PTX and matches or exceeds autotuned Triton on many workloads, while evaluation reveals verification limits for modern GPU features.

4 EVALUATION

  • Evaluation uses GPT-6 Astra on L40S, H100, and B200 GPUs, comparing against the fastest correct autotuned standard-Triton configuration for each kernel and GPU.The evaluated kernels span attention, GEMMs, activations, and normalization or regularization at fixed LLM-representative problem sizes.Source: §4 · p. 6

4.1 ARCHITECTURAL COVERAGE AND SOURCE SEMANTICS UNDERSTANDING

  • Generated PTX uses mma.sync and cp.async on Ada, wgmma and TMA on Hopper, and tcgen05 with tensor-memory accumulators on Blackwell.Hopper and Blackwell features fall outside Volta’s originally validated semantics, requiring the emulator extension for verification.Source: §4.1 · pp. 6–7, Figure 3 · p. 6

4.2 RESULTS ON COMMON KERNELS

  • 1.17-2.23× convolution speedup and 1.08-1.29× softmax speedup contrast with near-parity elementwise kernels; FP16 GEMM reaches only 0.83× on H100.Convolution reuses inputs across overlapping windows and overlaps packed data movement with Tensor Core execution.Source: Figure 4 · p. 7, §4.2 · pp. 7–8

4.1 ARCHITECTURAL COVERAGE AND SOURCE SEMANTICS UNDERSTANDING

  • TAIC produces semantically equivalent PTX for random kernels and misleadingly named or commented examples, and respects explicit IEEE-754 semantics by avoiding Tensor Core operations.The tests include unfamiliar integer kernels and deliberate changes to familiar algorithmic cues.Source: §4.1 · p. 7

4.3 REPEATABILITY UNDER A FIXED SEARCH BUDGET

  • Across five independent B200 runs, TAIC matches or outperforms Triton, with moderate variance on kernels showing significant gains.Runs used configured API-cost gates of $15 for convolution, matrix multiplication, and FlashAttention, and $1 for SwiGLU; an in-flight request could exceed the gate.Source: Table 1 · p. 8, §4.3 · p. 8

4.4 RESULTS ON RECENTLY PUBLISHED KERNELS

  • 3.34× and 2.51× are the BitDelta speedups; FlashSinkhorn reaches 1.55×, SageAttention 1.41×, and FlashAttention 1.37× over Triton on B200.Memory-bound and recurrent workloads remain near parity, including Lion at 1.00× and Dion at 0.97×. The largest gains involve reorganizing computation and data movement.Source: §4.4 · p. 8, Figure 5 · p. 9

4.5 FORMAL VERIFICATION IN THE LOOP

  • Original Volta-in-the-loop filtering yields less competitive kernels because modern GPU performance features are not all supported.The extension grows Volta’s trusted computing base from 42 kLoC to 60 kLoC; unrestricted TAIC output is verified for all common-suite kernels except reduction sum and softmax, and six recent-paper kernels. Atomics, input-dependent control flow, and bit-level value manipulation remain unsupported.Source: §4.5 · p. 9

5 CONCLUSION

Large language models can lower specialized Triton kernels directly to competitive PTX that passes randomized differential testing. The one-time generation cost can be amortized across repeated executions, and AI lowering may be especially compelling on new or lightly supported hardware.

5 CONCLUSION

  • AI lowering may be compelling for new or lightly supported hardware, where synthesis could cost substantially less than developing and maintaining a dedicated optimizing compiler backend.The paper notes that generation can be orders of magnitude more expensive than millisecond-scale Triton compilation, but that cost is incurred once and can be amortized over repeated execution.Source: §5 · p. 10

A REPRESENTATIVE OPTIMIZATION EXAMPLES

The section examines six kernels where generated PTX outperforms Triton, spanning common kernels and recently published workloads evaluated on B200.

A REPRESENTATIVE OPTIMIZATION EXAMPLES

A.1 COMMON KERNELS

The common-kernel examples show generated PTX reworking reduction algorithms, reusing overlapping convolution inputs, and reducing WGMMA and epilogue overhead. The implementations tailor instruction choice, data movement, and synchronization to the target GPU.

Reduction sum

  • 1.125× on L40S FP8 and 1.09× on H100 FP8 are achieved with distinct reduction strategies, while H100 FP16 uses Tensor Core MMA operations for row sums.The H100 FP16 path replaces scalar conversion and reduction chains with MMA instructions, shuffle-and-add stages, and one atomic update per warp, without shared-memory reduction or a block barrier.Source: §A.1 · p. 13

Convolution

  • 2.23× on B200 FP8 convolution, 1.90× on H100 FP16, and 1.17× on L40S FP8 result from reusing overlapping input-window values.Packed accesses and deeper pipelining reduce redundant gathers while retaining the Tensor Core computation.Source: §A.1 · p. 13

H100 FP8 GEMM

  • Larger WGMMA tiles, explicit double buffering, and a transposing stmatrix epilogue reduce instruction, synchronization, and layout-conversion overhead on H100 FP8 GEMM.The gain is reported on Hopper and absent on the other architectures.Source: §A.1 · p. 13

A.2 RECENTLY PUBLISHED KERNELS

The recent-kernel examples show that generated PTX can jointly reconsider instruction sequences, data movement, synchronization, and decomposition. The examples include direct binary-weight decoding, softmax ownership aligned with tensor memory, and reduction partitioning across independent warps.

A.2 RECENTLY PUBLISHED KERNELS

  • The examples attribute gains to jointly reconsidering instruction sequence, data movement, synchronization, and decomposition at the target level rather than merely emitting PTX directly.Source: §A.2 · p. 13

BitDelta

  • 3.34× Triton performance on BitDelta comes from decoding binary weights directly into Tensor Core operands; the batched variant reaches 2.51× using the same strategy.Packed sign-bit manipulation constructs FP16 ±1 values, and direct tensor-memory placement eliminates a shared-memory store-reload cycle.Source: §A.2 · pp. 13–14

FlashAttention

  • 1.37× speedup on FlashAttention comes from assigning each thread a complete 32-score row, making softmax reductions register-local and reducing block barriers from three to one per loop iteration.Combining two 64-row query tiles into one 128-row tile lets each key/value tile serve twice as many queries.Source: §A.2 · p. 14

FlashSinkhorn

  • 1.55× Triton performance on FlashSinkhorn comes from assigning independent 128-point slices to warps and combining their partial results only once at the end.The kernel passes evaluation tolerances while avoiding block barriers and shuffles during each warp's eight loop iterations.Source: §A.2 · p. 14

B VERIFICATION AND VOLTA EXTENSIONS

The paper extends Volta to verify more modern PTX constructs, including asynchronous operations and newer instruction families, but presents the extensions as a proof of concept rather than a sound or complete verifier.

B VERIFICATION AND VOLTA EXTENSIONS

  • The extensions support common PTX from recent GPU architectures, but are not claimed to be sound or complete, and their formalization is left for future work.The existing Volta feature set is insufficient for verifying PTX optimized for recent architectures.Source: B VERIFICATION AND VOLTA EXTENSIONS · p. 14
  • Asynchronous memory movement is tracked from initiation until an appropriate barrier, with hazards reported when it overlaps synchronous or asynchronous accesses.The extension tracks asynchronous operations at both source and destination.Source: B VERIFICATION AND VOLTA EXTENSIONS · p. 14
  • Support is added for complex Tensor Core instruction families, math operators, packed values, and their arithmetic, conversion, and load/store operations.Supported Tensor Core families include tcgen05, wgmma, and ldmatrix.Source: B VERIFICATION AND VOLTA EXTENSIONS · p. 14
  • For shuffle reductions, concurrent writes of the same value to one address are treated as harmless by checking value equality when race detection runs.The equality check invokes the equivalence checker during symbolic execution.Source: B VERIFICATION AND VOLTA EXTENSIONS · p. 14
  • Data-dependent control flow and atomics remain unsupported, and real-arithmetic modeling can reject kernels that depend on floating-point bit layouts despite expected finite-precision results.Atomics are excluded to avoid considering all possible thread interleavings.Source: B VERIFICATION AND VOLTA EXTENSIONS · pp. 14–15
Loading 2609.36800v1…