摘要
尽管基于大语言模型（LLM）的代码生成取得了快速进展，但为硬件加速器编写正确且高效的内核仍然是现代机器学习工作负载扩展的关键瓶颈。我们提出了 MKEvolve（模块化内核演进）框架，该框架通过迭代地共同演进复杂 PyTorch 模块的模块化分解和每个子模块的 LLM 生成内核，来解决这一问题。
在迭代过程中，通过拆分和融合来细化分解，同时通过 LLM 驱动的束搜索独立改进每个子内核。最终生成的内核是独立验证的子内核的程序化组合，使其具有可配置性（子内核实现可互换）、可解释性（错误和加速可追溯到特定子内核）以及易于适应相关模型架构的特点。
在 KernelBench L2 和 L3 上使用 Triton 进行的实验表明，MKEvolve 在正确性和加速性能上均优于端到端直接合成的基线，同时减少了 LLM 的令牌使用量，最多可减少 35%。
博主点评： MKEvolve 提供了一种创新的框架，利用模块化设计不断优化内核代码生成，极大地提升了性能和可维护性。这种方法不仅提高了生成代码的正确性，还有效降低了资源消耗，展示了现代机器学习工具与技术的深度融合潜力。