Harness Engineering for LLM-Driven GPU Kernel Generation
2026-07-20 • Machine Learning
Machine LearningArtificial Intelligence
AI summaryⓘ
The authors built a system to help AI models create better GPU code by organizing how the code is checked and improved. They use two parts: one that tests and times the code, and another that decides what new code to try based on those tests. Human experts provide important rules and examples to guide the AI, which helps make the code much faster than the original versions. Their results show that combining expert guidance with AI leads to better GPU programs than using AI alone.
large language modelsGPU kernelscode optimizationprofilingNVIDIA Blackwell B200compilationAI-assisted programmingperformance benchmarkingoperator constraintskernel generation
Authors
Yue Shui, Chenyu Ma, Hangfei Xu, Shengzhao Wen, Yanpeng Wang
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.