AVO Makes the Coding Agent the Variation Operator, Beating FA4 by 10.5% on B200

AVO: Agentic Variation Operators for Autonomous Evolutionary Search

Terry Chen, Zhifan Ye, Bing Xu, Zihao Ye, Timmy Liu, Ali Hassani, Tianqi Chen, Andrew Kerr, Haicheng Wu, Yang Xu, Yu-Jung Chen, Hanfeng Chen, Aditya Kane, Ronny Krashinsky, Ming-Yu Liu, Vinod Grover, Luis Ceze, Roger Bringmann, John Tran, Wei Liu, Fung Xie, Michael Lightstone, Humphrey Shi

cs.LG

2026-03-26

NVIDIA's AVO uses a coding agent as the variation operator. In 7 days on B200, MHA kernels hit 1668 TFLOPS, beating cuDNN by 3.5% and FA4 by 10.5% on causal attention.

What problem this solves

FunSearch and AlphaEvolve already showed that an LLM can sit inside an evolutionary loop and write better programs. In those systems the model still does one job: look at sampled parents and emit a candidate. Parent selection, evaluation, and population bookkeeping stay as fixed heuristics. The model cannot open a manual, run a profiler, read a failed compile, and revise before it is allowed to submit.

Attention kernels sit on that gap. FlashAttention-4 and cuDNN already took months of human work on Blackwell. What remains is PTX-level pipeline bubbles, barrier cost, and register spills, the kind of work that only yields if you iterate against hardware. A single-shot LLM mutation does not iterate at that grain.

Method

AVO replaces the whole variation operator Vary with a coding agent that plans, uses tools, and keeps memory. The lineage of kernels and scores, a domain knowledge base, and the scoring function are all in its reach. It decides which prior kernels to inspect, which PTX docs to open, and when to compile and bench.

The knowledge base includes CUDA guides, the PTX ISA, Blackwell specs, and FlashAttention-4 source. Each candidate is a CUDA kernel with inline PTX. Scoring is a vector over configs: fail the numerical check against a reference and that config scores 0, no matter how fast it runs. A version is git-committed into the lineage only when it is correct and at least as fast as the current best. Failed attempts stay in the agent's private trajectory.

This study uses a single lineage, no islands and no MAP-Elites archive, to isolate the operator. The agent is NVIDIA's internal general-purpose coding agent; the paper says its architecture was not special-cased for kernels. Long runs stall or loop, so a supervisor rereads the trajectory when progress dies and redirects search. The 7-day MHA run committed 40 versions.

Results

All numbers are forward prefilling on NVIDIA B200, CUDA 13.1, PyTorch 2.10.0. Baselines are cuDNN 9.19.1 and official FA4 (commit 71bf77c). Head dim 128, BF16, sequence lengths 4K/8K/16K/32K, with total tokens fixed at 32768 (batch 8 at 4K, batch 1 at 32K). MHA uses 16 heads, causal and non-causal. Timing follows the FA4 paper's script, 10 repeats.

Seven days with no human in the loop. The agent tried more than 500 optimization directions internally and committed 40 kernels. Causal MHA geometric mean: 1520 TFLOPS, against cuDNN 1488 and FA4 1426. Non-causal geometric mean: 1630, against cuDNN 1611 and FA4 1620. Peak single config is non-causal 32K at 1668 TFLOPS.

Settingvs cuDNNvs FA4
Causal MHA+0.4% to +3.5%+5.0% to +10.5%
Non-causal MHA, seq > 16K+1.8% to +2.4%Short sequences sit inside measurement noise of both baselines
Causal GQAup to +7.0%up to +9.3%
Non-causal GQAup to +6.0%up to +4.5%

The largest single-config gap is causal 4K: AVO 1392 TFLOPS, cuDNN 1344 (+3.5%), FA4 1259 (+10.5%). GQA was not re-evolved. The same agent adapted the finished MHA kernel in about 30 minutes, using Qwen3 layouts (32 query heads, 4 or 8 KV heads).

Committed history is a few jumps plus plateaus, not a smooth climb. Named inflections: QK-PV interleaving with bitmask causal masking (v8), a single-pass softmax (v13), branchless accumulator rescaling (v20), correction/MMA pipeline overlap (v30), register rebalancing across warp groups (v33). Ablations versus the previous commit, geomean TFLOPS:

OptimizationVersionsNon-causalCausal
Branchless accumulator rescalingv19→v20+8.1%+1.6%
Correction/MMA pipeline overlapv29→v30+1.1%+0.4%
Register rebalancingv32→v33+2.1%0%

v20 is the biggest single step. The kernel used to branch on whether any thread needed to rescale the output accumulator, paying warp sync on every K-block iteration. The agent always computes the scale factor and predicate-selects 1.0 when rescaling is unnecessary, which also lets it swap a blocking fence for a lighter ordering fence. That path covers fully unmasked iterations only, so non-causal gains +8.1% while causal still uses the old branch on masked blocks. On registers, Blackwell gives 2048 warp-registers per SM. FA4-style split was 192/80/48 for softmax/correction/rest; the agent saw correction spilling to local memory and moved 8 registers from softmax to each of the other groups (184/88/56).

Why it matters

For people who write GPU kernels, the concrete signal is small but real: on Blackwell attention that humans already spent months on, giving the agent the right to inspect, measure, and change strategy still wrings out a few percent. For people who build LLM evolutionary search, the claim is about the interface. The bottleneck may not be a stronger Generate model. It may be that Sample, evaluation, and strategy switches were never under the model's control.

Against cuDNN this is incremental. The causal MHA ceiling is 3.5%, and short non-causal configs do not separate from the baselines. The knowledge base includes FA4 source, so the agent could read an expert kernel from day one. What travels is the loop (correctness gate, commit-only-if-not-worse, supervisor on stall), not this particular kernel.

Limitations

The study covers a single lineage, forward prefilling, head dim 128, BF16. No decode, no backward, no other head sizes or precisions. The agent is an internal product. The paper does not name the frontier model, and it does not report token spend or GPU-hours for the 7-day run. The seed is labeled x0 with no description, while the trajectory text says v1–v20 closed the gap from a naive implementation to the baselines. Those two statements do not line up, so outsiders cannot judge how weak the start really was.

The supervisor has almost no interface detail: trigger and intervention are unspecified. GQA was a prompted adaptation after the fact, not a claim that search generalized on its own. v33 is 0% on causal, so some finds are config-bound. The appendix notes that absolute TFLOPS moves with driver, thermals, and clocks, and re-compares against the FA4 paper's published baseline numbers. Neither the kernel nor the agent is released, so reproducing the same setup is hard.

Terms

Source

What people are saying

Related papers

All paper explainers