用模块化迭代优化生成高性能硬件内核，提升正确率与效率
MKEvolve: A Modular Multi-Agent Framework for Kernel Code Generation
- 分步拆解复杂模块，逐块生成并优化子内核
- 在Triton上比端到端方法提升正确率，减少35%令牌消耗
- 生成的内核可替换、可追踪，适合模型迁移与调试
尽管基于大模型的代码生成进展迅速，为硬件加速器编写正确且高效的内核仍是扩展现代机器学习工作负载的关键瓶颈。我们提出MKEvolve（Modular Kernel Evolve）框架，通过迭代式协同进化复杂PyTorch模块的模块化分解及其对应的子模块内核，利用大模型驱动的束搜索独立优化每个子内核，并在迭代中通过拆分与合并来优化整体分解结构。最终生成的内核是独立验证的子内核程序组合，具备可配置性（子内核可替换）、可解释性（错误与性能提升可追溯至特定子内核），并易于适配同类模型架构。在Triton上对KernelBench L2和L3的实验显示，该方法在多算子序列及完整模型架构上，相较端到端直接合成基线，在提升正确率与加速比的同时，将大模型令牌使用量减少高达35%。
原文摘要 · Abstract (English)
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%.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。