TritonForge: Profiling-Guided Framework for Automated Triton Kernel Optimization

TL;DR

TritonForge is an automated framework that uses profiling and iterative code transformation to optimize GPU kernels written in Triton, significantly improving performance with minimal manual effort.

Contribution

It introduces a profiling-guided, automated optimization system for Triton GPU kernels, reducing manual tuning and achieving substantial performance gains.

Findings

Up to 5x performance improvement over baseline kernels

Average success rate of 1.76x across diverse kernel types

Effective identification and mitigation of performance bottlenecks

Abstract

High-performance GPU kernel optimization remains a critical yet labor-intensive task in modern machine learning workloads. Although Triton, a domain-specific language for GPU programming, enables developers to write efficient kernels with concise code, achieving expert-level performance still requires deep understanding of GPU architectures and low-level performance trade-offs. We present TritonForge, a profiling-guided framework for automated Triton kernel optimization. TritonForge integrates kernel analysis, runtime profiling, and iterative code transformation to streamline the optimization process. By incorporating feedback from profiling results, the system identifies performance bottlenecks, proposes targeted code modifications, and evaluates their impact automatically. Across diverse kernel types, TritonForge achieves up to 5x performance improvement over baseline implementations and on average 1.76x of the cases are successful, providing a foundation for future research in automated GPU performance optimization.