LEGO: An LLM-Enabled Hierarchical Optimizer for Tensor Computation Graphs with Structure-Aware Search and Compositional Synthesis

Automating end-to-end GPU kernel generation with Large Language Models (LLMs) faces a critical tension between global performance and exploration efficiency. We present LEGO, a hierarchical framework that resolves this trade-off via a parallel multi-agent search over a recursive AND-OR FusionTree. LEGO synergizes two complementary flows: Top-Down Construction decomposes complex graphs into valid, context-isolated sub-problems to guarantee correctness and enable parallel exploration, while Bottom-Up Mutation speculatively fuses verified sub-plans to recover global locality for peak performance. This bi-directional mechanism effectively prunes the search space to avoid repetitive unguided sampling, while naturally parallelizing exploration and enabling the discovery of sophisticated fusion strategies. Evaluations demonstrate that LEGO achieves 2.18×–13.48× speedups over PyTorch Eager and, compared to monolithic baselines, reduces end-to-end exploration time by up to 2.47× and token consumption by up to 7×, across diverse end-to-end models.

Modern AI models like ChatGPT are expensive and energy-hungry to run, and squeezing maximum speed from the graphics chips they use requires hand-crafting tiny low-level programs called "kernels," painstaking work that normally demands expert engineers. Researchers have started using AI language models to write these programs automatically, but they hit a wall: handed a whole model at once, the AI gets overwhelmed and makes mistakes; split into small pieces, it handles each but loses speed where the pieces meet. We built LEGO, a system that works in both directions, much like assembling something from bricks and then smoothing the joints. Because AI models are built from repeated blocks, LEGO splits a model along its natural seams and hands each block to a separate AI "worker" running in parallel, keeping every task small enough to get right. It then fuses neighboring blocks back together to win back the speed lost at the seams. The result: models that run 2 to 13 times faster than standard tools, found at up to 7 times lower search cost than letting one AI tackle the whole model at once. LEGO can even accelerate brand-new model designs that no human has yet optimized by hand.