TLDR: Astra is the first LLM-based multi-agent system designed to optimize existing GPU CUDA kernels, specifically from the SGLang LLM serving framework. It uses specialized AI agents (coding, testing, profiling, planning) that collaborate iteratively to generate, test, and refine code. Astra achieved an average speedup of 1.32x on SGLang kernels using zero-shot prompting, demonstrating its ability to autonomously apply advanced optimization techniques like loop transformations, memory access improvements, and CUDA intrinsics, significantly enhancing performance and correctness.
In the world of high-performance computing and machine learning, optimizing GPU (Graphics Processing Unit) kernels is a persistent and significant challenge. These kernels are the fundamental building blocks that accelerate tasks like training and serving large language models (LLMs). Traditionally, achieving peak performance for these kernels has demanded extensive manual tuning by expert engineers, a process that is both time-consuming and resource-intensive. While compiler-based systems have helped, they still require substantial human effort in design and engineering.
Recent advancements have seen researchers explore the use of large language models (LLMs) for generating GPU kernels. However, much of this prior work has focused on translating high-level PyTorch modules into CUDA code. A new system, Astra, takes a different approach, introducing the first LLM-based multi-agent system specifically designed for GPU kernel optimization.
Astra distinguishes itself by starting with existing CUDA implementations. Instead of translating from PyTorch, it extracts and optimizes code directly from SGLang, a widely used framework for serving LLMs in production environments. This focus on optimizing existing, production-grade code means that even modest improvements can have a substantial real-world impact, given that SGLang is responsible for generating trillions of tokens daily across major enterprises.
How Astra Works: A Collaborative AI System
Astra operates as a multi-agent system, where specialized LLM agents work together in an iterative feedback loop. This collaborative structure is key, as kernel optimization involves multiple stages that a single LLM agent might struggle to master comprehensively. The system comprises four main agents:
- Testing Agent: This agent is responsible for creating a suite of test cases from the baseline kernel and then validating the correctness of any new, optimized kernel candidates.
- Profiling Agent: It measures the execution time of kernels on the generated test suite, providing crucial performance feedback.
- Planning Agent: Combining the correctness and performance data, this agent proposes targeted modifications and optimization strategies.
- Coding Agent: This agent takes the suggestions from the planning agent and applies them to generate new, optimized kernel implementations.
This iterative cycle of code generation, testing, profiling, and planning allows Astra to systematically explore the optimization space, ensuring that the resulting kernels are not only high-performance but also functionally correct.
Impressive Performance Gains
Astra was evaluated on three critical kernels extracted from SGLang: silu_and_mul, fused_add_rmsnorm, and merge_attn_states_lse. Using zero-shot prompting with OpenAI’s o4-mini model, Astra achieved an average speedup of 1.32 times, with some kernels seeing improvements up to 1.46 times. Importantly, these results were obtained without any additional training, such as supervised fine-tuning or reinforcement learning, highlighting the inherent effectiveness of the multi-agent approach in a pure prompting setting.
A comparison with a single-agent approach further underscored Astra’s advantages. While the single-agent method achieved an average speedup of 1.08 times, Astra’s multi-agent system consistently delivered better results, especially for more complex kernels. The dedicated roles of the agents in Astra helped avoid issues like biased profiling results that could occur in a single-agent setup.
Under the Hood: Optimization Strategies
Detailed analysis of the optimized kernels revealed several sophisticated strategies autonomously applied by the LLM agents:
- Loop Transformations: Astra identified and hoisted loop-invariant computations out of inner loops, significantly reducing redundant calculations and increasing throughput.
- Memory Access Improvements: The system restructured memory access patterns, for instance, by using vectorized loads (grouping two half-precision values into a single instruction) to reduce memory transactions and boost effective memory bandwidth.
- CUDA Intrinsics: Astra made extensive use of CUDA intrinsics, which are specialized functions that map directly to hardware-level operations, leading to greater efficiency.
- Fast Math Operations: The optimized kernels replaced standard math library calls with faster CUDA device intrinsics for operations like exponentiation and reciprocal calculations, improving arithmetic pipeline utilization.
These strategies collectively contribute to the observed speedups, demonstrating the LLMs’ capability to apply advanced, expert-level optimization techniques.
Also Read:
- AI Agents Achieve Breakthrough in Autonomous Code Evolution for NP-Complete Problems
- Enhancing Language Model Accuracy Through User Feedback and Adaptive Decoding
Future Directions
While Astra represents a significant leap in GPU kernel optimization, the researchers acknowledge certain limitations. Currently, the pre-processing (extracting and simplifying kernels) and post-processing (integrating optimized kernels back into SGLang) steps are performed manually. Automating these complex steps is a key area for future research, which would allow Astra to scale to an even broader range of kernels and frameworks beyond SGLang, such as vLLM, PyTorch, and TorchTitan.
Astra marks a promising new paradigm for GPU kernel optimization, showcasing the power of multi-agent LLM systems in tackling complex, real-world engineering challenges. For more details, you can read the full research paper: Astra: A Multi-Agent System for GPU Kernel Performance Optimization.
