AlphaEvolve: A Coding Agent for Scientific and Algorithmic Discoveries
Optimizing the Gemini Core
Training large models such as Gemini requires substantial computational resources. Gemini is built on the JAX framework, while Pallas is its extension that makes it possible to write highly specialized programs (so-called kernels) tailored for optimal execution on hardware accelerators. Efficient Pallas kernels are therefore essential for optimizing the performance of Gemini model training. An important aspect of kernel optimization is tuning the tiling strategy for matrix multiplication operations. This technique involves dividing the computation of a large matrix multiplication into smaller subproblems to better balance computation with data movement, which is crucial for accelerating the overall computation. Traditionally, kernel developers rely either on search-based autotuning or manually created heuristics to determine near-optimal tiling configurations for different input shapes. However, autotuning interrupts the research workflow and requires retuning for every change in input shape. Conversely, manually creating effective tiling heuristics represents a major engineering challenge due to its complexity, requiring a deep understanding of both kernel functionality and hardware specifics. The main advantage of a powerful heuristic is its ability to deliver high performance across arbitrary input shapes. Therefore, to accelerate the design of high-performance kernels for new hardware and simplify their use by model developers, our goal is to facilitate the process of generating heuristics.
That challenge was addressed by deploying AlphaEvolve to optimize tiling heuristics for an important matrix multiplication kernel used in Gemini training. The goal was to minimize the kernel's actual runtime. AlphaEvolve iteratively explores and refines tiling heuristics for this kernel by proposing candidate code, aiming to minimize this runtime across various input shapes on real TPU accelerators. Kernel correctness is preserved because AlphaEvolve optimizes the tiling strategy for this kernel rather than changing its underlying mathematical operation. To create training and evaluation datasets for AlphaEvolve, realistic kernel input shapes were automatically collected from its users. Half of these input shapes constitute the training set, which provides optimization objectives during the evolutionary process. The remaining input shapes constitute the evaluation set, which is used to test the general applicability of the resulting heuristic.
This automated approach enabled AlphaEvolve to discover a heuristic that provides an average 23% speedup across all kernels compared with the existing expert-designed heuristic, corresponding to a 1% reduction in Gemini's overall training time. In addition, using AlphaEvolve significantly shortened the kernel optimization period, from several months of specialized engineering work to just days of automated experiments. This acceleration speeds up the deployment of optimized kernels, allowing kernel engineers to devote their expertise to more strategic, higher-level optimization problems. Furthermore, AlphaEvolve offers a path toward automating the manual tuning process and improving the ergonomics of using Gemini kernels. The tiling heuristic discovered using AlphaEvolve was deployed in production, directly increasing the efficiency of Gemini training and the research and development speed of the Gemini team. This deployment also represents a new case in which Gemini, through AlphaEvolve's capabilities, optimizes its own training process.
Assistance with Hardware Circuit Design
Specialized hardware, such as Google's Tensor Processing Units (TPUs), is essential for achieving the resource efficiency required to operate modern AI systems at scale. However, designing new computer chips is a complex and time-consuming process that often takes years. Register-Transfer Level (RTL) optimization, a critical step in this process, involves manually rewriting hardware descriptions to improve metrics such as performance, power consumption, and area, requiring months of iteration by highly skilled engineers.
In this work, AlphaEvolve was tasked with optimizing an already highly optimized implementation of a key TPU arithmetic circuit in Verilog within the matrix multiplication unit. The goal of the optimization was to reduce both area and power consumption while preserving the component's core functionality. Crucially, the final design had to pass thorough verification methods confirming that the modified circuit retained functional correctness. AlphaEvolve managed to find a simple code rewrite that removed unnecessary bits—a change whose correctness was verified by TPU designers. Although this particular improvement was later independently captured by synthesis tools as well, AlphaEvolve's contribution at the RTL stage demonstrates its ability to improve RTL source code and provide optimizations early in the design process.
This improvement, integrated into the upcoming TPU, represents Gemini's first direct contribution to TPU arithmetic circuits, achieved through AlphaEvolve, and opens the way for similar contributions in the future. A key advantage of AlphaEvolve is that it communicates proposed changes directly in Verilog, the standard language used by hardware engineers, fostering trust and simplifying adoption. This initial exploration demonstrates a new approach in which language-model-driven code evolution assists with hardware design and potentially shortens time to market.
Direct Optimization of Compiler-Generated Code
The transformer architecture is used in most modern neural networks, from large language models to AlphaFold. The fundamental computational mechanism of transformers is so-called attention, which is most commonly implemented using FlashAttention. In our computational stack, FlashAttention is implemented as an acceleration kernel in Pallas, wrapped in a higher-level layer of JAX code that handles input preparation and subsequent output processing. The machine learning compiler (XLA) then translates this implementation into a sequence of intermediate representations (IR), each of which adds more details for execution on specific hardware. At these stages, improved decisions about memory access orchestration or computation scheduling can significantly reduce runtime on specific hardware.
AlphaEvolve was tasked with directly optimizing the IR generated by the XLA compiler, which encapsulates the FlashAttention kernel together with preprocessing and postprocessing code. We optimized the configuration corresponding to a highly important transformer model used for large-scale inference on GPUs, with the goal of minimizing the module's total execution time. This was an exceptionally challenging task because (1) the IR is designed primarily for debugging purposes rather than direct editing by developers, and (2) it is generated by the compiler and already highly optimized. Every modification proposed by AlphaEvolve was checked against reference (unmodified) code on random inputs to ensure numerical correctness throughout the optimization process. The final version of the code was thoroughly verified by human experts and confirmed to be correct for all possible inputs.
AlphaEvolve was able to provide meaningful optimizations at both levels of abstraction exposed by the IR. First, the FlashAttention kernel for the relevant configuration was sped up by 32%. Second, AlphaEvolve found improvements in the preprocessing and postprocessing of the kernel's inputs and outputs, resulting in a 15% speedup in this part. These results demonstrate AlphaEvolve's ability to optimize compiler-generated code, offering the potential to incorporate discovered optimizations into existing compilers for specific use cases or, in the longer term, to include AlphaEvolve directly in the compilation workflow.
Although AlphaEvolve is a relatively new tool, it is already delivering significant benefits in three key areas of AI and hardware development and optimization. Its ability to automatically optimize complex kernels, contribute to hardware circuit design, and improve compiler-generated code opens up new possibilities for more efficient use of computational resources and faster research in AI. As this tool continues to develop, its impact on scientific and algorithmic discoveries can be expected to become even more significant.
