JustGoal: With Only a Goal, the AI Baseline Is Above the Expert Kernels

GPU kernels written on NVIDIA B200 by a general-purpose agent given only an optimization goal

English | 中文

_{Speedup of single-shape JustGoal and the KDA-0.5 reproduction (whole definition) over the reference kernel across model generations on four operators; each bar is the geometric mean over the operator's shapes, the dashed line is the reference kernel.}

This repository releases the operator (kernel) code from the paper JustGoal: With Only a Goal, the AI Baseline Is Above the Expert Kernels. JustGoal gives a general-purpose agent only an optimization objective and constraints, without kernel-specific evaluation infrastructure, knowledge bases, or search workflows. The agent builds its own evaluation loop and iterates on its implementations. On NVIDIA B200, JustGoal outperforms open-source libraries on 8 of the 10 operators entering the comparison, with per-operator geometric mean speedups of 1.14–2.34×, and surpasses the fastest publicly released specialized-agent kernel on 5 of the 9 operators that have one. In a same-model comparison with Fable 5.1 on single shapes, the KDA-0.5 harness reaches 0.95× the performance of JustGoal.

Paper

JustGoal: With Only a Goal, the AI Baseline Is Above the Expert Kernels

arXiv: https://arxiv.org/abs/2610.36434

Repository Layout

Directory	Contents
goal_single_shape/	/goal runs on a single specified shape
goal_full_shape/	/goal runs on all shapes of an operator definition
kda05_reproduction_full_shape/	Runs under the KDA-0.5 setup
kda05_reproduction_single_shape/	Runs under the KDA-0.5 setup on a single specified shape
ablation/	Ablation runs

Every folder is organized as <model>/<operator>/<shape>/. In the full-shape folders the shape level is the flashinfer-bench definition; kda05_reproduction_full_shape/ adds a level for each released version and ablation/ one for each variant. The goal_full_shape/, kda05_reproduction_full_shape/ and kda05_reproduction_single_shape/ kernels come with a flashinfer-bench solution.json.

Environment

• GPU: NVIDIA B200 (sm_100a)
• CUDA: 13.0 (nvcc 13.0.88, driver 580.105.08). Two kernels point at a CUDA 13.1 toolkit: goal_single_shape/opus-5/mla_paged_prefill/b1_q33_kv34 and goal_single_shape/opus-5/gqa_paged_decode/b16_kv2841.
• CuTe DSL (nvidia-cutlass-dsl): 4.6.2. The DSA, GDN and MoE kernels in kda05_reproduction_full_shape/ and kda05_reproduction_single_shape/ use 4.6.0, the version pinned by the KDA-0.5 setup.

Contributing

Issues and pull requests are welcome, in particular bug reports on wrong results or build failures, and retest results from other B200 setups.

Citation

BibTeX

@misc{chen2026justgoal,
  title  = {JustGoal: With Only a Goal, the AI Baseline Is Above the Expert Kernels},
  author = {Chen, Yaojian and Peng, Dazhi and Zhu, Zhui and Li, Leshan and Sun, Haojun
            and Zheng, Yiyun and He, Qingfeng and Li, Shangzhan and Deng, Xianglong
            and Shen, Yifei and Xiang, Shengye and Li, Zhen and Gan, Lin and Li, Yuxuan
            and Han, Xu and Liu, Zhiyuan},
  year   = {2026},
  eprint = {2610.36434},
  archivePrefix = {arXiv},
  url    = {https://arxiv.org/abs/2610.36434}
}

License

Licensed under the Apache License, Version 2.0. Files adapted from NVIDIA CUTLASS examples keep their BSD-3-Clause headers and license; they are listed in NOTICE.

Acknowledgments

We build on FlashInfer, flashinfer-bench and the MLSys 2026 contest workloads, FlashAttention, and CUTLASS / CuTe DSL.

This work is from Tsinghua University, ModelBest, Harbin Institute of Technology and the University of Chinese Academy of Sciences.