# Ave: Guiding Agentic GPU Optimization Using Data-Flow Invariants

arXiv: https://arxiv.org/abs/2604.18616 · HTML: https://mlsanity.com/paper/2604.18616/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Software engineering coding agents, one of 16 directions in LLM agents and tool use.
- Builds on 25 of its 83 references that are in the map.
- First posted on arXiv 16 Apr 2026.

## Placed in

- LLM agents and tool use / [Software engineering coding agents](https://mlsanity.com/agents/software-engineering-coding-agents/)

## Builds on (most cited first, 25 of 25)

- [Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer](https://mlsanity.com/paper/1701.06538/) (2017)
- [Lost in the Middle: How Language Models Use Long Contexts](https://mlsanity.com/paper/2307.03172/) (2023)
- [FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning](https://mlsanity.com/paper/2307.08691/) (2023)
- [DeepSeek-V3.2: Pushing the Frontier of Open Large Language Models](https://mlsanity.com/paper/2512.02556/) (2025)
- [TextGrad: Automatic "Differentiation" via Text](https://mlsanity.com/paper/2406.07496/) (2024)
- [KernelBench: Can LLMs Write Efficient GPU Kernels?](https://mlsanity.com/paper/2502.10517/) (2025)
- [Ansor: Generating High-Performance Tensor Programs for Deep Learning](https://mlsanity.com/paper/2006.06762/) (2020)
- [Online normalizer calculation for softmax](https://mlsanity.com/paper/1805.02867/) (2018)
- [Towards Robust Agentic CUDA Kernel Benchmarking, Verification, and Optimization](https://mlsanity.com/paper/2509.14279/) (2025)
- [Comet: Fine-grained Computation-communication Overlapping for Mixture-of-Experts](https://mlsanity.com/paper/2502.19811/) (2025)
- [CUDA Agent: Large-Scale Agentic RL for High-Performance CUDA Kernel Generation](https://mlsanity.com/paper/2602.24286/) (2026)
- [CudaForge: An Agent Framework with Hardware Feedback for CUDA Kernel Optimization](https://mlsanity.com/paper/2511.01884/) (2025)
- [Mirage: A Multi-Level Superoptimizer for Tensor Programs](https://mlsanity.com/paper/2405.05751/) (2024)
- [LLMs Are In-Context Bandit Reinforcement Learners](https://mlsanity.com/paper/2410.05362/) (2024)
- [Benchmarking and Dissecting the Nvidia Hopper GPU Architecture](https://mlsanity.com/paper/2402.13499/) (2024)
- [K-Search: LLM Kernel Generation via Co-Evolving Intrinsic World Model](https://mlsanity.com/paper/2602.19128/) (2026)
- [AVO: Agentic Variation Operators for Autonomous Evolutionary Search](https://mlsanity.com/paper/2603.24517/) (2026)
- [MPK: A Compiler and Runtime for Mega-Kernelizing Tensor Programs](https://mlsanity.com/paper/2512.22219/) (2025)
- [KernelBlaster: Continual Cross-Task CUDA Optimization via Memory-Augmented In-Context Reinforcement Learning](https://mlsanity.com/paper/2602.14293/) (2026)
- [ProofWright: Towards Agentic Formal Verification of CUDA](https://mlsanity.com/paper/2511.12294/) (2025)
- [Improving Efficiency of GPU Kernel Optimization Agents using a Domain-Specific Language and Speed-of-Light Guidance](https://mlsanity.com/paper/2603.29010/) (2026)
- [Tilus: A Tile-Level GPGPU Programming Language for Low-Precision Computation](https://mlsanity.com/paper/2504.12984/) (2025)
- [CUCo: An Agentic Framework for Compute and Communication Co-design](https://mlsanity.com/paper/2603.02376/) (2026)
- [Analyzing Modern NVIDIA GPU cores](https://mlsanity.com/paper/2503.20481/) (2025)
- [HipKittens: Fast and Furious AMD Kernels](https://mlsanity.com/paper/2511.08083/) (2025)

## Abstract (arXiv)

LLM coding agents can generate correct GPU kernels, but their performance still trails expert libraries. Reaching peak throughput requires coordinating low-level optimizations such as shared-memory staging, software pipelining, and instruction scheduling. Yet unit tests and profiles provide only sparse end-to-end feedback, making it difficult for agents to identify which global constraints an optimization violates. We present Ave, an agentic framework that uses data-flow invariants as compile-time guardrails for GPU kernel optimization. Ave provides a tile-based Pythonic DSL that exposes hardware instructions and compiler policies while abstracting complex memory layouts. Tag functions assign symbolic labels to data, the compiler propagates them through data and control flow, and tag assertions enforce required relationships at use sites. A flow-sensitive, path-insensitive analysis with an SMT solver checks these assertions and returns concrete counterexamples for violations, with no runtime overhead. An in-context reinforcement learning planner proposes optimizations from a curated knowledge base, while a lowering agent implements them and instantiates the required invariants. We evaluate Ave on AMD MI300X across GEMM, flash attention, and MoE, which together account for up to 90% of GPU time in LLM inference. With GPT-5.6 Sol, Ave achieves 89-99% of the effective throughput of state-of-the-art hand-optimized libraries and improves geometric-mean throughput by 1.62-1176x over uncontaminated agentic baselines. On 200 KernelBench tasks, Ave produces valid kernels within three attempts for 100% of Level 1 and 88% of Level 2 problems.