Harness Engineering for LLM-Driven GPU Kernel Generation
Harness-based LLM-assisted GPU kernel system achieves 1.62-29.68× speedups, integrating human skills and profile-backed search.
Key Findings
Methodology
The system decouples evaluation harness from an optimization controller, leveraging human-designed skills to encode operator constraints, references, and promotion rules. Codex and Claude Code generate candidate kernels within these constraints. The controller converts profiler and workload evidence into bounded candidate generation, using profile-backed states and trajectory memory for reproducibility. Multi-round iterative optimization incorporates human-guided plateau recovery, combining multi-language support (CUDA, Triton, CuTe) and latency-first promotion. This approach ensures correctness and performance improvements, validated across five operator definitions with mean speedups from 1.62× to 29.68×, demonstrating effective human-AI collaboration.
Key Results
- Retained artifacts achieved average speedups of 1.62×, 18.05×, 29.68×, 1.12×, and 13.70× over FlashInfer baselines across five operator definitions. Agent-assisted kernels outperformed full-agent artifacts, emphasizing the importance of expert guidance and workload context.
- Experiments conducted on NVIDIA Blackwell B200 with CUDA 13.2, PyTorch 2.12, Triton 3.6, validating correctness and latency improvements. Multi-round iterative refinement effectively avoided optimization stagnation.
- Analysis of final kernels revealed shape-aware dispatch and workload-grounded promotion as key strategies, especially in sparse attention and gated networks, leading to significant performance gains.
Significance
This work advances GPU kernel automation by integrating human expertise with LLM-generated code, addressing the engineering challenges of correctness, performance, and reproducibility. It provides a practical system framework for reliable kernel generation, crucial for deploying large models efficiently. The approach bridges the gap between theoretical optimization and real-world engineering constraints, promising broad impact on AI inference acceleration and hardware design.
Technical Contribution
The paper introduces a harness-based architecture that separates evaluation and optimization, utilizing profile-backed states and trajectory memory for reproducibility. It supports multi-language kernels, employs a multi-stage search with latency-first promotion, and incorporates human-guided plateau recovery. These innovations enable stable, high-performance kernel generation aligned with engineering constraints, surpassing prior automatic methods that lacked such systematic controls.
Novelty
This is the first comprehensive system integrating harness engineering principles with LLM-assisted kernel generation, emphasizing engineering constraints, reproducibility, and human-in-the-loop control. Unlike previous works focusing solely on search algorithms, this approach ensures correctness and performance within a controllable, engineer-friendly framework, setting a new standard for practical AI-driven kernel optimization.
Limitations
- The system heavily relies on human-designed skills and references, limiting full automation. Its effectiveness depends on the quality of expert input and profile data, which may not generalize across all operators or hardware.
- Multi-round optimization is time-consuming, especially for complex operators, potentially hindering rapid deployment in production environments.
- The approach's transferability to different hardware platforms remains untested, requiring further validation for broader applicability.
Future Work
Future directions include automating skill generation via reinforcement learning, reducing human intervention, and enhancing system scalability. Extending support to diverse hardware architectures and workloads will improve generalization. Incorporating more sophisticated profiling and adaptive strategies could further accelerate convergence and robustness, making the system more autonomous and widely applicable.
AI Executive Summary
As deep learning models grow in size and complexity, optimizing GPU kernels has become a critical bottleneck for inference efficiency. Traditional manual tuning methods are labor-intensive and often inconsistent, especially when dealing with irregular operators like sparse attention or mixture-of-experts (MoE). To address this challenge, this work introduces a harness-based system that leverages large language models (LLMs) to assist in GPU kernel generation, grounded in engineering principles and human expertise.
The core idea is to decouple the evaluation process from the optimization controller. The evaluation harness handles compilation, correctness verification, and latency measurement, while the controller manages candidate generation based on profiler evidence and workload characteristics. Human experts define reusable skills capturing operator constraints, references, and promotion rules, which guide the LLMs—Codex and Claude Code—in generating candidate kernels within a constrained, verifiable environment.
This multi-stage, profile-backed approach ensures that only candidates improving latency without sacrificing correctness are promoted. The system employs multi-language support (CUDA, Triton, CuTe), multi-round iterative refinement, and human-in-the-loop interventions such as plateau recovery. Experimental results on NVIDIA Blackwell B200 GPUs demonstrate speedups ranging from 1.62× to 29.68× across five operator definitions, significantly outperforming baseline FlashInfer kernels.
The significance of this work lies in its practical, engineering-oriented framework that combines AI's generative power with rigorous validation, addressing longstanding issues of reliability and reproducibility in kernel optimization. It paves the way for automated, yet controllable, kernel design that can adapt to complex workloads and hardware constraints. Future research aims to enhance automation, reduce human dependency, and extend applicability across diverse hardware platforms, ultimately transforming AI inference deployment and hardware acceleration.
Deep Analysis
Background
深度学习模型不断扩大,GPU硬件成为推理的核心平台。传统核优化依赖手工调优,效率低且难以应对复杂操作符。早期工作如KernelBench、KernelEvolve通过生成-评估-选择循环探索优化空间,但缺乏系统工程支持。随着大语言模型(LLM)崛起,利用其生成能力辅助核设计成为新趋势。现有研究多集中于自动化搜索,但缺乏系统工程支持,导致生成核在实际应用中难以保证正确性和性能。本文在此背景下提出Harness工程思想,结合Profile-backed状态管理和人类专家的技能定义,构建了一个可控、可追溯的核生成系统,填补了自动化与工程实践的空白。
Core Problem
GPU核优化面临多重挑战:一是核的正确性必须严格保证,否则会引发错误;二是性能提升需在多样化工作负载中实现,而非单一指标;三是工程实现复杂,容易出现基线漂移、验证不充分等问题。传统方法难以兼顾这些需求,导致优化效果不稳定。如何在保证正确性的同时,系统性地提升核性能,成为亟待解决的问题。特别是在多操作符、多硬件环境下,如何实现高效、可靠的自动化核生成,仍是研究难点。
Innovation
核心创新包括:1)引入Harness工程思想,将人类设计的技能封装为可复用的约束和操作规则,确保生成核符合工程要求;2)采用Profile-backed状态管理,将Profiler信息作为搜索的基础,提升搜索的可追溯性和稳定性;3)支持多语言(CUDA、Triton、CuTe),满足不同硬件和微调需求;4)结合多轮迭代和Latency-first promotion策略,有效平衡性能与正确性,避免优化停滞。这些创新使系统在保证核正确性的基础上,显著提升了性能,突破了传统自动优化的局限。
Methodology
- �� 设计评估工具(Harness)与优化控制器(Controller),将工作负载、Profile信息、候选核封装为结构化上下文;
- �� 利用人类专家定义的技能捕获操作符约束、参考代码、推广规则,确保核的工程可控性;
- �� 采用Profile-backed状态管理,记录每轮候选核的性能、正确性和Profiler证据,作为下一轮搜索的基础;
- �� 支持多语言(CUDA C++、Triton、CuTe),满足不同硬件需求;
- �� 结合二阶段搜索策略:候选生成(模型/模型新)、验证(正确性、性能)、推广(Latency-first);
- �� 多轮迭代流程:识别瓶颈、生成候选、编译验证、Profiler测量、轨迹存储、人工引导plateau recovery;
- �� 通过多任务、多工作负载、多轮优化,确保核在不同场景下的性能提升与正确性。
Experiments
在NVIDIA Blackwell B200 GPU上,使用CUDA 13.2、PyTorch 2.12、Triton 3.6环境,评估五个操作符定义(MoE FP8、DSA top-k、DSA attention、GDN decode、GDN prefill),对比基础FlashInfer基线。每个定义采用不同工作负载(如19、128、23、54、100个),验证核的正确性和性能提升。多轮优化中,结合Profiler信息和人工引导,筛选出最优候选核。通过对比Agent-Assisted与Full-Agent方法,验证了系统的有效性。指标包括平均延迟、速度提升(最高达29.68倍)和正确性保证。还进行了微调和高试验轮次验证,确保核的稳健性。
Results
在五个操作符定义中,Retained artifacts的平均速度提升分别为1.62x、18.05x、29.68x、1.12x和13.70x,显著优于基础基线。Agent-Assisted核在所有定义中优于Full-Agent核,特别是在稀疏注意力和门控网络中表现出极大优势。多轮优化结合Profiler信息,有效避免了微调中的错误,验证了Shape-aware dispatch和Workload-grounded promotion的有效性。实验结果显示,结合人类专家引导的系统架构,核生成的性能提升具有高度稳定性和可控性。
Applications
该系统适用于深度学习模型部署中的GPU核优化,特别是在大规模模型推理、稀疏注意力、门控网络等场景。通过结合Profile信息和工程化技能,能显著缩短核开发周期,提升推理速度,降低能耗。未来可扩展到多硬件平台,支持自动调优,推动深度学习硬件产业升级。
Limitations & Outlook
系统依赖人类专家设计的技能,自动化程度有限,难以完全自主。多轮优化耗时较长,实际部署需权衡时间成本。对极端或新颖操作符的适应性不足,未来需增强泛化能力和迁移性。
Plain Language Accessible to non-experts
想象你在一个工厂里,工人们要制造一台复杂的机器。每个机器都需要不同的零件和组装步骤,工人们需要确保每个步骤都正确无误,才能保证机器正常工作。以前,工人们自己设计每个机器的制造流程,既费时又容易出错。现在,有了智能助手(就像大语言模型),它可以根据工人的指导,快速生成符合要求的制造流程。工人提供一些基本规则和参考样例,助手根据这些信息,自动设计出多种方案。然后,工人会检查这些方案,挑选最优的那一个。这个过程不断循环,直到找到既快又正确的制造流程。这个系统就像一个智能工厂助手,结合人类经验和AI的计算能力,大大提高了生产效率,保证每台机器都能高质量制造出来。
ELI14 Explained like you're 14
想象你在学校的厨房里做饭。以前,你需要自己想办法,把所有材料和步骤都安排好,才能做出美味的菜肴。这很难,也容易出错。现在,有个聪明的机器人助手,它可以帮你设计菜谱。你告诉它你想做什么菜,它会根据你的要求,给出几种不同的做法。你可以试试这些方案,看看哪个最合你口味。你还可以告诉它你喜欢吃辣还是不辣,它会调整菜谱。你们一起反复试验,最后找到最棒的做法。这个机器人就像一个厨房助手,帮你节省时间,又保证菜做得好吃。这个系统用在GPU核优化上也是一样的:人类提供一些基本规则和参考,AI帮忙设计不同的核方案,然后人类挑选最好的,反复优化,最终让GPU跑得更快、更稳定。
Glossary
Harness工程 (Harness Engineering)
一种系统设计思想,将工程约束和验证机制融入自动化流程,确保AI生成的核符合性能和正确性要求。
论文中用于确保GPU核生成的工程可控性与可追溯性。
Profile-backed状态 (Profile-backed State)
利用Profiler收集的性能和正确性信息,作为搜索和优化的基础状态。
用于多轮优化中记录核的性能指标。
Latency-first promotion (延迟优先推广)
一种优先提升核性能的策略,确保优化目标是降低平均延迟。
在多轮优化中用以筛选候选核。
Shape-aware dispatch (形状感知调度)
根据工作负载的形状特征,动态调整核调度策略。
优化稀疏注意力和门控网络中的调度路径。
Agent-Assisted (代理辅助)
结合人类专家和AI模型共同进行核设计的系统架构。
论文的核心方法之一。
Open Questions Unanswered questions from this research
- 1 如何进一步减少多轮优化所需时间,提升系统的自动化程度,成为未来研究的重要方向。当前系统仍依赖人类技能设计,自动化水平有限。
Applications
Immediate Applications
深度学习模型推理加速
利用该系统优化GPU核,提升大模型推理速度,降低能耗,适用于云端推理平台和边缘设备。
硬件加速器设计
为新一代GPU或专用加速器设计定制核,提升硬件性能和能效,推动AI硬件产业升级。
Long-term Vision
自动化深度学习硬件生态
实现从模型到硬件的全链路自动化优化,推动深度学习硬件的普及与定制化,缩短研发周期。
Abstract
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.