TritonForge: Profiling-Guided Framework for Automated Triton Kernel Optimization

2025-12-09

AI Summary

Triton GPU kernel optimization heavily relies on expert knowledge and is labor-intensive. Method: This paper proposes the first automated optimization framework that integrates dynamic performance profiling with large language model (LLM)-driven semantic reasoning. It employs a modular architecture to close the loop among bottleneck identification, data-driven code refactoring, and automatic correctness and performance validation—enabling architecture-agnostic, end-to-end tuning without human intervention. Innovatively, it unifies Nsight runtime analysis, static kernel feature extraction, and LLM-guided semantic-level rewriting to jointly optimize memory access patterns, arithmetic intensity, and parallelism. Results: Evaluated across diverse Triton kernels and GPU architectures, the framework achieves an average speedup of 1.76×, with peak improvements up to 5×, significantly lowering the barrier to high-performance kernel development.

Technology Category

• Search and Optimization: Learning to Search
• Machine Learning: Optimization
• Natural Language Processing: Learning & Optimization for NLP

Application Category

• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

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 data-driven feedback from profiling results, the system identifies performance bottlenecks, proposes targeted code modifications, and evaluates their impact automatically. While our prototype leverages large language models (LLMs) to assist in code reasoning and transformation, the framework remains modular and model-agnostic. Across diverse kernel types and GPU architectures, 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.

Problem

• Automates optimization of GPU kernels written in Triton language
• Reduces manual effort in achieving expert-level GPU performance
• Identifies performance bottlenecks through profiling-guided analysis

Innovation

• Profiling-guided automated framework for Triton kernel optimization
• Integrates kernel analysis, profiling, and iterative code transformation
• Uses data-driven feedback to identify bottlenecks and modify code