Harness Engineering for LLM-Driven GPU Kernel Generation
作者: Yue Shui, Chenyu Ma, Hangfei Xu, Shengzhao Wen, Yanpeng Wang
分类: cs.LG, cs.AI
发布日期: 2026-07-20
备注: 24 pages, 6 figures. Extended technical report on our submission to the MLSys 2026 FlashInfer AI Kernel Generation Contest. Code: https://github.com/syhya/mlsys26-flashinfer-contest
💡 一句话要点
提出基于测试框架的LLM驱动GPU内核生成优化方法
🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 大型语言模型 GPU内核生成 性能优化 测试框架 机器学习系统
📋 核心要点
- 现有方法在生成GPU内核时缺乏有效的约束和验证机制，导致生成代码的可靠性不足。
- 本文提出了一种以测试框架为中心的系统，分离评估框架与优化控制器，以提高生成内核的质量和性能。
- 实验结果显示，优化后的内核在多个操作符定义上相较于基线有显著的延迟提升，表明该方法的有效性。
📝 摘要（中文）
大型语言模型（LLMs）可以辅助生成GPU内核，但其实际效果依赖于生成代码的可靠约束、验证、分析和选择。本文提出了一种以测试框架为中心的系统，用于在NVIDIA Blackwell B200 GPU上进行LLM驱动的GPU内核优化。该系统将评估框架与基于分析的优化控制器分离，框架负责编译、正确性、官方对齐时序和文档归档，而控制器则将分析器和工作负载证据转化为有界候选生成决策。通过五个操作符定义，保留的官方对齐文档在提供的FlashInfer基线之上实现了平均延迟加速，分别为1.62倍、18.05倍、29.68倍、1.12倍和13.70倍。实验结果表明，专家提供的优化方向和高质量参考对可靠的AI驱动内核优化至关重要。
🔬 方法详解
问题定义：本文旨在解决大型语言模型生成的GPU内核在可靠性和性能上的不足，现有方法缺乏有效的约束和验证机制，导致生成的代码难以满足实际应用需求。
核心思路：论文提出的系统通过将评估框架与优化控制器分离，利用人类专家提供的技能和LLM生成的候选内核，确保生成代码的正确性和性能。这样的设计使得生成过程更加可控和可靠。
技术框架：该系统主要包括两个模块：评估框架和优化控制器。评估框架负责代码的编译、正确性验证、性能计时和文档归档，而优化控制器则利用分析器和工作负载证据来生成有界的候选内核。
关键创新：最重要的创新在于引入了以测试框架为中心的系统架构，确保了生成内核的高质量和高性能，同时结合了人类专家的知识与LLM的生成能力，形成了有效的协同工作机制。
关键设计：系统中关键的设计包括人类专家编写的技能集，涵盖操作符约束、参考文献、分析程序和提升规则，以及Codex和Claude Code代理在这些约束内生成候选内核的能力。
🖼️ 关键图片
📊 实验亮点
实验结果显示，优化后的内核在五个操作符定义上实现了显著的性能提升，具体为1.62倍、18.05倍、29.68倍、1.12倍和13.70倍的平均延迟加速，表明该方法在实际应用中的有效性和优势。
🎯 应用场景
该研究的潜在应用领域包括高性能计算、深度学习模型优化和图形处理等。通过提高GPU内核生成的可靠性和性能，能够显著提升相关应用的运行效率，推动AI和计算机视觉等领域的发展。
📄 摘要（原文）
Large language models (LLMs) can assist GPU kernel generation, but their practical effectiveness depends on whether generated code can be reliably constrained, validated, profiled, and selected. This paper presents a harness-centered system for LLM-driven GPU kernel optimization in the MLSys 2026 FlashInfer AI Kernel Generation Contest on NVIDIA Blackwell B200 GPUs. The system separates an evaluation harness from a profile-backed optimization controller: the harness enforces compilation, correctness, official-aligned timing, and artifact archival, while the controller turns profiler and workload evidence into bounded candidate-generation decisions. Human-authored skills capture operator constraints, references, profiling procedures, and promotion rules, while Codex and Claude Code agents generate candidate kernels inside those constraints. Across five operator definitions, the retained official-aligned artifacts achieved mean-latency speedups over supplied FlashInfer baselines of 1.62x, 18.05x, 29.68x, 1.12x, and 13.70x. The Agent-Assisted kernels outperform the Full-Agent artifacts across the evaluated definitions, indicating that expert-provided optimization directions, high-quality references, and workload context remain critical for reliable AI-driven kernel optimization.