论文

MKEvolve:内核代码生成的模块化多智能体框架

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

摘要

尽管基于LLM的代码生成进展迅速,为硬件加速器编写正确且高性能的内核仍是扩展现代ML工作负载的关键瓶颈。我们提出MKEvolve(模块化内核进化),一个框架:迭代地共同演化复杂PyTorch模块的模块化分解与每个子模块的LLM生成内核,跨迭代经拆分与融合精炼分解,同时经LLM驱动的束搜索独立改进每个子内核。所得内核是独立验证子内核的程序化组合:可配置(子内核实现可替换)、可解释(错误与加速可追溯到具体子内核)、且易适配相关模型架构。在KernelBench L2与L3上、以Triton、横跨多算子序列与完整模型架构的实验显示:MKEvolve在正确性与加速上都超过端到端直接合成基线,同时把LLM token用量最多降低35%。