MKEvolve: A Modular Multi-Agent Framework for Kernel Code Generation

Jason Yoo, Rajarshi Saha, Tao Yu, Shaowei Zhu, Wei Tang, Youngsuk Park

Keywords: LLM, Code Generation, Kernel, Triton, Multi-Agent, Test-Time Scaling

TL;DR: MKEvolve uses multi-agent decomposition and optimization to generate modular GPU kernels that improve correctness and performance while reducing token usage.

Abstract: Despite rapid progress in LLM-based code generation, writing correct and performant kernels for hardware accelerators remains a key bottleneck in scaling modern ML workloads. We present MKEvolve (Modular Kernel Evolve), a framework that iteratively co-evolves a modular decomposition of complex PyTorch modules and the LLM-generated kernel for each submodule, refining the decomposition by splitting and fusing across iterations while independently improving each subkernel via LLM-driven beam search. The resulting kernels are programmatic compositions of independently verified subkernels, making them configurable (subkernel implementations are swappable), interpretable (errors and speedups are traceable to specific subkernels), and readily adaptable to related model architectures. Experiments with Triton on KernelBench L2 and L3, spanning multi-operator sequences and full model architectures, show that MKEvolve improves both correctness and speedup over end-to-end direct synthesis baselines while reducing LLM token usage by up to 35%.

Email Sharing: We authorize the sharing of all author emails with Program Chairs.

Data Release: We authorize the release of our submission and author names to the public in the event of acceptance.

Submission Number: 14