arXiv:2511.20100cs.DCcs.CL2025-11AAAI被引 17

用分层策略让大模型高效生成高性能GPU代码,准确率超现有方法。

QiMeng-Kernel: Macro-Thinking Micro-Coding Paradigm for LLM-Based High-Performance GPU Kernel Generation

  • 分两阶段:先选优化策略,再逐步实现代码,避免直接生成全量错误。
  • 在KernelBench上准确率达99%(Level 1-2),比顶尖模型高50%以上。
  • 适合想自动化生成高性能计算代码的研究者和工程师。

高性能GPU内核开发对人工智能与科学计算至关重要,但依赖专家手工设计且可移植性差。尽管大语言模型(LLM)有望实现自动化,通用型和微调后的LLM仍存在正确性与效率的双重矛盾。根源在于现有方法直接生成完整优化的底层程序,需探索涵盖优化策略与代码实现的庞大空间。为此,本文提出宏观思考微观编码(MTMC)框架,受人类专家分阶段优化启发,将优化策略与实现细节解耦:宏观思考通过强化学习引导轻量级LLM高效探索最大化硬件利用率的语义优化策略;微观编码利用通用型LLM逐步实现宏观提出的优化方案,避免全内核生成错误。二者协同有效覆盖广阔优化空间与复杂实现细节,实现大模型驱动的高性能GPU内核生成。在广泛使用的基准测试中,MTMC表现优异:在KernelBench上,1-2级精度接近100%,3级达70%,较现有通用与领域微调模型提升超50%,运行速度最高快7.3倍,较专家优化的PyTorch Eager内核快2.2倍;在更具挑战性的TritonBench上,精度达59.64%,速度提升34倍。

原文摘要 · Abstract (English)

Developing high-performance GPU kernels is critical for AI and scientific computing, but remains challenging due to its reliance on expert crafting and poor portability. While LLMs offer promise for automation, both general-purpose and finetuned LLMs suffer from two fundamental and conflicting limitations: correctness and efficiency. The key reason is that existing LLM-based approaches directly generate the entire optimized low-level programs, requiring exploration of an extremely vast space encompassing both optimization policies and implementation codes. To address the challenge of exploring an intractable space, we propose Macro Thinking Micro Coding (MTMC), a hierarchical framework inspired by the staged optimization strategy of human experts. It decouples optimization strategy from implementation details, ensuring efficiency through high-level strategy and correctness through low-level implementation. Specifically, Macro Thinking employs reinforcement learning to guide lightweight LLMs in efficiently exploring and learning semantic optimization strategies that maximize hardware utilization. Micro Coding leverages general-purpose LLMs to incrementally implement the stepwise optimization proposals from Macro Thinking, avoiding full-kernel generation errors. Together, they effectively navigate the vast optimization space and intricate implementation details, enabling LLMs for high-performance GPU kernel generation. Comprehensive results on widely adopted benchmarks demonstrate the superior performance of MTMC on GPU kernel generation in both accuracy and running time. On KernelBench, MTMC achieves near 100% and 70% accuracy at Levels 1-2 and 3, over 50% than SOTA general-purpose and domain-finetuned LLMs, with up to 7.3x speedup over LLMs, and 2.2x over expert-optimized PyTorch Eager kernels. On the more challenging TritonBench, MTMC attains up to 59.64% accuracy and 34x speedup.

GPU优化大模型生成代码自动化高性能计算

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。