JustGoal method speeds up AI training kernels without expert input
JustGoal: With Only a Goal, the AI Baseline Is Above the Expert Kernels
Distributed, Parallel, and Cluster Computing
Summary
Training large AI language models requires running some key computer programs, called kernels, very efficiently. Usually, experts spend a lot of time manually improving these kernels using special knowledge and tools. The authors tested a new method called JustGoal, where an AI agent only knows the goals and limits it must meet, without extra expert knowledge or workflows. JustGoal was faster than common open-source kernels and often beat specialized expert methods, showing that simple goal-focused approaches can work very well as AI models improve.
What this means in practice
- •For ai infrastructure teams: Use goal-driven AI agents to optimize GPU kernels for faster large language model training on NVIDIA hardware, reducing reliance on manual expert tuning.
- •For high performance computing teams: Improve efficiency of GPU-accelerated computations by adopting method that automatically designs kernels from objectives alone without predefined knowledge.
Authors
Yaojian Chen, Dazhi Peng, Zhui Zhu, Leshan Li, Haojun Sun, Yiyun Zheng, Qingfeng He, Shangzhan Li, Xianglong Deng, Yifei Shen, Shengye Xiang, Zhen Li, Lin Gan, Yuxuan Li, Xu Han, Zhiyuan Liu
Abstract
A small number of GPU kernels often determines the efficiency of large language model training and inference. High-performance kernel development typically relies on expert knowledge through manual implementation, collaboration with general-purpose models, or specialized agents equipped with domain knowledge, evaluation methods, and search workflows. We investigate how much value this knowledge provides as foundation models become stronger. We introduce JustGoal, a goal-only baseline in which a general-purpose agent receives only the optimization objective and constraints, without kernel-specific infrastructure, knowledge bases, or predefined workflows, and builds its own evaluation loop. On NVIDIA B200 GPUs, JustGoal outperforms open-source libraries on 8 of 10 operators, achieving per-operator geometric mean speedups of 1.14x to 2.34x. It also surpasses the fastest publicly released specialized-agent kernel on 5 of 9 operators with available comparisons. In a controlled same-model comparison using Fable 5.1 on single shapes, the KDA-0.5 harness reaches 0.95x of JustGoal's performance, while single-mechanism ablations show no consistent benefit from domain knowledge or predefined workflows. These results indicate that specialized machinery should be justified by clear gains over a same-model, same-scope goal-only baseline and reevaluated whenever the foundation model changes.