用工程化系统提升大模型生成GPU核函数的可靠性与性能。
Harness Engineering for LLM-Driven GPU Kernel Generation
- 构建分离式评估框架，约束代码生成并自动验证正确性。
- 在5个算子上平均提速13.7倍，最高达29.68倍。
- 专家经验+高质量参考仍比纯自动化更可靠，适合优化工程师。
大型语言模型可辅助生成GPU核函数，但其实际效果依赖于生成代码能否被可靠约束、验证、分析和选择。本文提出一种以评估工装为核心的系统，用于MLSys 2026 FlashInfer AI核函数生成竞赛中NVIDIA Blackwell B200 GPU上的大模型驱动核函数优化。系统将评估工装与基于性能分析的优化控制器分离：工装强制执行编译、正确性、官方对齐时间及产物归档；控制器则根据剖析器数据与工作负载证据做出有界候选生成决策。人类专家提供操作符约束、参考代码、剖析流程和采纳规则，而Codex与Claude Code代理在这些约束内生成候选核函数。在五个算子定义中，保留的官方对齐产物相较于原始FlashInfer基线，平均延迟加速比分别为1.62x、18.05x、29.68x、1.12x和13.70x。代理辅助生成的核函数优于全自动化产物，表明专家提供的优化方向、高质量参考和工作负载上下文对可靠AI驱动的核函数优化依然至关重要。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。