Astra: A Multi-Agent System for GPU Kernel Performance Optimization

TL;DR

<3-5 sentence high-level summary> Astra presents a novel multi-agent, LLM-driven framework for optimizing existing CUDA kernels sourced from SGLang, addressing the challenge of achieving high performance without extensive manual tuning. By coordinating specialized agents for code generation, testing, profiling, and planning, Astra explores the optimization space and produces kernels that are both correct and faster, achieving an average 1.32× speedup across three kernels in zero-shot prompting. The work demonstrates concrete optimization strategies such as loop-invariant hoisting, warp-level reductions with intrinsics, and vectorized memory accesses, validating the feasibility and practical impact of multi-agent LLM systems in GPU kernel optimization. Limitations include the need for manual pre-/post-processing and a small kernel set, pointing to future work on automation and broader framework support to scale the approach.

Abstract

GPU kernel optimization has long been a central challenge at the intersection of high-performance computing and machine learning. Efficient kernels are crucial for accelerating large language model (LLM) training and serving, yet attaining high performance typically requires extensive manual tuning. Compiler-based systems reduce some of this burden, but still demand substantial manual design and engineering effort. Recently, researchers have explored using LLMs for GPU kernel generation, though prior work has largely focused on translating high-level PyTorch modules into CUDA code. In this work, we introduce Astra, the first LLM-based multi-agent system for GPU kernel optimization. Unlike previous approaches, Astra starts from existing CUDA implementations extracted from SGLang, a widely deployed framework for serving LLMs, rather than treating PyTorch modules as the specification. Within Astra, specialized LLM agents collaborate through iterative code generation, testing, profiling, and planning to produce kernels that are both correct and high-performance. On kernels from SGLang, Astra achieves an average speedup of 1.32x using zero-shot prompting with OpenAI o4-mini. A detailed case study further demonstrates that LLMs can autonomously apply loop transformations, optimize memory access patterns, exploit CUDA intrinsics, and leverage fast math operations to yield substantial performance gains. Our work highlights multi-agent LLM systems as a promising new paradigm for GPU kernel optimization. Our code is publicly available at https://github.com/Anjiang-Wei/Astra.

Figures (5)

• Figure 1: Overview of Astra. Given an existing GPU kernel extracted from SGLang, Astra employs a multi-agent approach for kernel optimization, where specialized LLM agents collaborate through iterative code generation, testing, profiling, and planning to produce correct and high-performance kernels.
• Figure 2: Hoisting loop-invariant computation in merge_attn_states_lse.
• Figure 3: Reduction strategies in fused_add_rmsnorm. \ref{['fig:red-baseline']}: block-level tree reduction in shared memory with synchronization each step. \ref{['fig:red-optimized']}: intra-warp reduction in registers using __shfl_down_sync, followed by a short inter-warp aggregation in shared memory.
• Figure 4: Comparison of global-memory loads in the baseline and optimized kernels. The baseline uses a scalar half-precision load, while the optimized version employs a vectorized half2 load for improved efficiency.
• Figure 5: Side-by-side SiLU implementations. The optimized kernel replaces a division with a reciprocal–multiply sequence and uses the fast exponential intrinsic, improving compute throughput.