TritonGen: Inference-Time Control Strategies Improve GPU Kernel Generation Quality

The TritonGen framework uses inference-time control strategies such as grammar-constrained decoding, correctness feedback, and compiler repair loops to significantly improve the effectiveness, correctness, and performance of Triton GPU kernel generation without fine-tuning the model. This thread will introduce the background, core methods, experimental evidence, and future directions in separate floors.

Challenges in Code Generation

Large language models excel in code generation, but generating functionally correct and high-performance GPU kernels still faces significant challenges (involving complex memory models, parallel execution semantics, and hardware-specific optimization techniques).

Introduction to the Triton Language

Triton is a Python-like programming language developed by OpenAI, specifically designed for writing high-performance GPU kernels. It has a high level of abstraction and performance close to handwritten CUDA, allowing developers to focus on algorithm logic while the compiler handles low-level optimizations.

Grammar-constrained decoding is one of the core technologies of TritonGen. Traditional autoregressive generation does not consider syntax and easily produces syntactic errors; this strategy introduces context-free grammar (CFG) constraints, selecting only syntactically valid tokens at each step, fundamentally eliminating syntactic errors and improving the compile rate of generated code.

Even syntactically correct code may have logical errors. TritonGen verifies correctness by executing the generated kernel, collects error information (such as value mismatches, segmentation faults, etc.) and feeds it back to the model, simulating the human debugging process. It converges to the correct implementation through multiple iterations and operates entirely at inference time without updating model parameters.

TritonGen uses compiler error messages and profiler outputs to optimize generated results: when compilation fails, it parses error feedback and sends it to the model; when performance is poor, it uses profiling data to identify bottlenecks. This tool-augmented generation strategy leverages existing toolchain capabilities, enabling collaboration between AI and tools to improve kernel performance.

Experimental results show that the system with grammar constraints and feedback loops has significant improvements in code validity, functional correctness, and execution performance compared to the baseline model. Moreover, these improvements do not require modifying model parameters, are generalizable and transferable, and are highly attractive to those with limited resources.

The core idea of TritonGen (using inference-time control strategies to improve generation quality) can be extended to fields such as structured data generation and formal proof. Future directions include designing more fine-grained constraint mechanisms, exploring multimodal feedback, and combining control strategies with fine-tuning methods to further unleash the model's potential.