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

Anjiang Wei; Tianran Sun; Yogesh Seenichamy; Hang Song; Anne Ouyang; Azalia Mirhoseini; Ke Wang; Alex Aiken

Source-linked AI summary

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

Anjiang Wei, Tianran Sun, Yogesh Seenichamy, Hang Song, Anne Ouyang, Azalia Mirhoseini, Ke Wang, Alex Aiken

arXiv:2509.07506v2Sep 9, 2025cs.DCcs.AIcs.CLcs.LGcs.SE

TL;DR

GPU kernel optimization remains a difficult, labor-intensive problem, and prior LLM approaches have largely generated CUDA from high-level PyTorch specifications. Astra instead uses specialized agents to iteratively optimize existing SGLang CUDA kernels, achieving a 1.32× average speedup with zero-shot prompting while exposing concrete optimization strategies.

Problem
GPU kernel optimization requires substantial manual effort, while prior LLM approaches have largely focused on translating high-level PyTorch modules into CUDA code.
Method
Astra coordinates specialized agents for code generation, testing, profiling, and planning to iteratively optimize existing CUDA kernels extracted from SGLang.
Results
1.32× average speedup is achieved on three SGLang kernels using zero-shot prompting with OpenAI o4-mini, versus 1.08× for a single-agent baseline.
Takeaways & Limitations
The case analysis shows that LLMs can apply loop transformations, restructure memory access, use CUDA intrinsics, and exploit fast math operations for kernel optimization.
Takeaways & Limitations
The evaluation covers three CUDA kernels and is tailored to SGLang, while preprocessing and postprocessing remain fully manual.

Abstract

from arXiv · show

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.

1 Introduction

GPU kernel optimization is difficult and labor-intensive, while existing manual, compiler-based, and LLM approaches leave important gaps. Astra addresses this challenge by coordinating specialized agents to optimize existing CUDA kernels, achieving measurable speedups on SGLang kernels.

Motivation: GPU kernel optimization remains difficult because hardware evolution, new architectures, and dynamic workloads require extensive manual tuning.Many implementations operate below hardware peak, limiting performance, cost, and energy efficiency.
Prior Approaches: Existing approaches span manual tuning, compiler-based optimization, and LLM-driven kernel generation, but each still involves substantial design or engineering effort.Prior LLM work includes kernel generation and training-based methods, while compiler systems use abstractions and autotuning.
Astra: Astra introduces a multi-agent system that iteratively coordinates code generation, testing, profiling, and planning for GPU kernel optimization.The system targets existing CUDA implementations rather than translating high-level PyTorch models into CUDA code.
Astra: Unlike KernelBench, Astra optimizes existing CUDA kernels extracted from SGLang, reflecting production settings where performance improvements are needed after kernels already exist.The optimized kernels can be reintegrated into SGLang, a production-grade LLM serving framework.
Results: 1.32× average speedup is achieved on three SGLang kernels with zero-shot OpenAI o4-mini prompting, compared with 1.08× for a single-agent baseline.The evaluation uses no additional supervised fine-tuning or reinforcement learning.
Results: Astra’s generated kernels gain performance through loop transformations, memory-access restructuring, CUDA intrinsics, and fast math operations.These strategies were identified through detailed manual analysis of the optimized kernels.

2 Related Work

Related work covers compiler and DSL systems, multi-agent frameworks, and LLM-driven high-performance code generation. For GPU kernels, iterative feedback combines compilation, correctness validation, profiling, and self-reflection while preserving functional equivalence.

Multi-Agent Systems: Multi-agent systems decompose complex programming workflows into subtasks such as planning, implementation, testing, and profiling.Frameworks including AutoGen, Trace, and MetaGPT exemplify this collaborative paradigm.
Compiler and Learning-Based Approaches: Compiler and DSL systems such as Halide, TVM, MLIR, XLA, and CUTLASS provide abstractions for expressing tensor computations and applying compiler optimizations.Autotuning frameworks further improve performance by searching optimization choices.
LLM-Driven Approaches: LLM research has expanded from general code generation to vectorization, assembly optimization, DSL programming, tensor optimization, and performant GPU-kernel generation.Iterative refinement is especially natural when optimization provides verifiable rewards.
LLM-Driven Approaches: Kernel optimization methods use feedback loops involving compilation checks, correctness validation, runtime profiling, or self-reflection to refine candidate code.Correctness requires equivalence to the original program for all inputs, in addition to high performance.

3 Method

Astra formalizes CUDA optimization as maximizing geometric-mean speedup while preserving correctness, then applies four specialized agents in iterative optimization rounds. The workflow tests, profiles, plans, codes, and records candidate kernels throughout the process.

Task Definition: The objective is to maximize geometric-mean speedup while preserving kernel correctness.Correctness is assessed over diverse test inputs, while geometric means aggregate speedup ratios across the test suite.
Multi-Agent System: Astra assigns testing, profiling, planning, and coding to specialized agents in the CUDA optimization pipeline.Testing generates and validates cases, profiling measures runtime, planning proposes modifications, and coding applies them.
Multi-Agent System: The setup begins by constructing a test suite and profiling the baseline, then repeats planning, coding, testing, and profiling for R rounds.The baseline is represented as S0, and each subsequent candidate is derived from the previous kernel.
Multi-Agent System: Each optimization round generates a candidate kernel, validates correctness, measures performance, and records the result in an optimization log.The log stores the round, code, correctness indicator, and performance, enabling systematic tracking of the optimization trajectory.
Implementation: Astra uses manually simplified standalone CUDA kernels as inputs and reintegrates optimized kernels into SGLang for validation against the original implementation.This preprocessing and post-processing addresses internal dependencies in the raw SGLang kernels.

4 Experimental Setup

The evaluation measures correctness and execution speed for three SGLang kernels across representative LLM-related tensor shapes. Experiments use repeated timing runs on NVIDIA H100 GPUs with OpenAI o4-mini and five optimization rounds.

Metrics: The evaluation compares generated-kernel outputs with the original SGLang implementation and reports execution-time speedup.Correctness uses diverse tensor shapes, while performance compares baseline and optimized runtimes on the same shapes.
Tables: Table 1 organizes the evaluated kernels by name and computation, while Table 2 compares baseline and optimized LoC and execution time.Table 2 also indicates that all optimized kernels are correct.
Kernels: The benchmark includes silu_and_mul, fused_add_rmsnorm, and merge_attn_states_lse from SGLang.These are the three kernels evaluated in the experimental setup.
Performance Measurement: Performance is measured across representative input shapes using 20 warm-up runs followed by 100 repetitions per shape.The shapes reflect dimensions used in modern LLaMA-7B, 13B, and 70B models.
Implementation: The experiment uses NVIDIA H100 GPUs, OpenAI o4-mini agents, and five optimization rounds.Astra is implemented with the OpenAI Agents SDK framework.

5 Results

Astra produces correct, faster kernels across three SGLang workloads, with multi-agent coordination outperforming a single-agent setup and case studies identifying concrete optimization mechanisms.

Correctness: All three Astra-optimized kernels were validated against the original SGLang implementations and confirmed correct.Validation used manually constructed test cases rather than tests generated by the testing agent.
Performance: 1.32× average speedup and up to 1.46× were achieved across representative tensor shapes after five optimization rounds.The individual speedups were 1.26× for merge_attn_states_lse, 1.25× for fused_add_rmsnorm, and 1.46× for silu_and_mul.
Single-Agent Comparison: 1.32× vs. 1.08× speedup shows that the multi-agent approach outperformed the single-agent approach while both generated correct kernels.The multi-agent advantage became more pronounced for more complex kernels.
Case Studies: Hoisting loop-invariant computations out of the inner loop reduced repeated exponentials and division, leaving memory loads, multiply–add, and a store.The transformation lowered instruction count and increased throughput.
Case Studies: Warp-level reduction kept partial sums in registers and reduced synchronization overhead before a short inter-warp shared-memory aggregation.This register-resident intra-warp phase yielded higher arithmetic throughput and lower memory traffic than a shared-memory-only approach.
Case Studies: Vectorized __half2 loads reduced memory transactions, while CUDA intrinsics and reciprocal–multiplication improved SiLU compute throughput.The optimized SiLU used __expf, __frcp_rn, and __fmul_rn instead of standard math calls and floating-point division.

6 Discussion

Astra is evaluated across tensor shapes and aims for general tensor-computation improvements rather than shape-specific tuning, but the current study remains limited in scope and manual integration work.

Evaluation across shapes: Astra achieves consistent speedups across four representative tensor shapes for each of three kernels.The evaluated shapes cover sequence length, number of heads, head dimension, batch size, and hidden size, depending on the kernel.
Evaluation across shapes: Astra does not prompt agents to optimize for a particular tensor shape, targeting general tensor computations instead.This differs from tensor compiler approaches that perform shape-specific tuning.
Scope: The evaluation currently covers three CUDA kernels and is tailored to SGLang.The authors plan to extend support to additional kernels and frameworks such as vLLM, PyTorch, and TorchTitan.
Limitations: Pre-processing and post-processing remain fully manual, including extracting stand-alone kernels, reintegrating optimized kernels into SGLang, and validating them against the original implementation.The paper identifies automation, potentially with human-in-the-loop guidance, as future work for scaling Astra.

7 Conclusion

Astra applies specialized LLM agents to existing CUDA kernels from SGLang, addressing GPU kernel optimization through coordinated generation, testing, profiling, and planning. It achieves an average speedup of 1.32×, while case studies identify several concrete optimization strategies.

Contribution: Astra is the first LLM-based multi-agent system designed specifically for GPU kernel optimization.It operates directly on existing CUDA kernels from SGLang rather than translating high-level PyTorch modules into CUDA.
Method: Astra coordinates specialized agents for code generation, testing, profiling, and planning to produce correct and high-performance kernels.
Results: 1.32× average speedup is achieved on the evaluated SGLang kernels.
Optimization strategies: Case studies show LLMs applying loop transformations, restructuring memory access, exploiting CUDA intrinsics, and using fast math operations.