Source-linked AI summary
Ave: Guiding Agentic GPU Optimization Using Data-Flow Invariants
Haohui Mai, Xiaoyan Guo, Xiangyun Ding, Daifeng Li, Qiuchu Yu, Chenzhun Guo, Cong Wang, Jiacheng Zhao, Christos Kozyrakis, Binhang Yuan
TL;DR
LLM agents can generate correct GPU kernels, but their performance trails expert libraries, and sparse test and profile feedback makes optimization constraints difficult to diagnose. Ave combines a tile-based DSL, compile-time data-flow invariant checks, and a learning-guided agentic workflow. On AMD MI300X, it approaches hand-optimized library throughput and outperforms agentic baselines across evaluated workloads.
Problem
LLM-generated GPU kernels can be functionally correct yet remain slower than expert libraries, while unit tests and profiles provide sparse feedback that does not identify the responsible transformation or ordering.
Source: §1 · p. 1Method
Ave combines a tile-based DSL with tag functions and assertions, compile-time invariant analysis, and an ICRL planner and lowering agent that implement optimization proposals and instantiate invariants.
Source: §1 · p. 3Results
89–99% of the effective throughput of state-of-the-art hand-optimized libraries, and 1.62–1176× geometric-mean throughput gains over uncontaminated agentic baselines across evaluated workloads.
Source: §1 · p. 3Takeaways & Limitations
Compile-time invariant witnesses provide actionable correctness feedback for optimizing kernels through intrusive transformations without runtime overhead.
Source: §11 · p. 13Ave does not model the full heap or track tags written to global memory, so it cannot express invariants depending on heap-resident values, including the described MoE sorting property.
Source: §5 · p. 6
Abstract
from arXiv · showhide
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.
1 Introduction
Ave addresses the gap between functionally correct LLM-generated kernels and production performance by giving agents compile-time data-flow invariants as guardrails. Its DSL, static analysis, and learning-guided optimization loop support complex transformations, with results approaching hand-optimized libraries and outperforming agentic baselines.
1 Introduction
- Existing agentic harnesses produce GEMM kernels 1.8–14.3× slower than hipBLASLt, while unit tests and profiles reveal failure or slowness without identifying the responsible transformation or ordering.Higher-level Triton narrows the gap to approximately 2× but does not expose fine-grained instruction scheduling or wavefront specialization.Source: §1 · pp. 1–2
- Ave expresses correctness constraints as tag functions and assertions that track data relationships through transformations such as shared-memory staging and software pipelining.Conformity assertions check alignment, while non-conformity assertions check separation, such as preventing concurrent producers from overwriting the same shared-memory region.Source: §1 · p. 2
- Compile-time, flow-sensitive and path-insensitive analysis checks invariants with an SMT solver and reports concrete violating assignments, without runtime overhead.Reports identify the thread, data element, and program point witnessing an abstract-state violation.Source: §1 · p. 2
- An ICRL planner retrieves optimization proposals and invariant templates, while a lowering agent implements them and instantiates tags and assertions.Invariant violations and runtime performance serve as reward signals for improving subsequent plans.Source: §1 · p. 2
- Ave combines invariant checks with agent feedback because tags alone only specify properties, static analysis cannot repair violations, and an ICRL agent lacks structured feedback.The authors leave discovering new algorithmic optimizations and fully verifying implementation correctness outside Ave’s scope.Source: §1 · p. 2
- 89–99% of the effective throughput of state-of-the-art hand-optimized libraries, with 1.62–1176× geometric-mean throughput gains over uncontaminated agentic baselines.Ave also solves all 100 KernelBench Level 1 and 88 Level 2 tasks within three attempts.Source: §1 · p. 3
2 Motivating Example
Flash attention illustrates why GPU kernel optimization requires coordinating algorithmic choices with hardware-specific layouts, memory use, and scheduling. Ave uses data-flow invariants to check correctness properties that become error-prone across these transformations.
2 Motivating Example
- Flash attention avoids materializing the N×N score matrix, reducing memory from O(N^2) to O(N) by tiling Q, K, and V and using online softmax.The tile sizes are chosen so the working set fits in registers and shared memory.Source: §2 · p. 3
- On MI300X, peak performance requires balancing register and shared-memory use against occupancy, while vectorized access, reuse, and software pipelining address its ∼250 arithmetic-to-bandwidth ratio.Instruction scheduling and register allocation also matter because poor schedules stall in-order pipelines and excess register use causes spills.Source: §2 · p. 3
- The flash-attention kernel uses MFMA for both GEMMs, but its row-major attention operations require operand reordering and V transposition for hardware-specific swizzled layouts.The kernel also pipelines global loads with computation, partitions work across disjoint wavefront groups, and schedules instructions to avoid stalls and spills.Source: §2 · p. 3
- Because layout transformations and data shuffling can mismatch expected and actual layouts, Ave’s tag assertions check that values remain correctly paired at each use site.The illustrated kernel uses 512 threads, d=128, Br=256, and Bc=64; the figure elides the full pipeline schedule and wavefront-group control flow.Source: §2 · p. 3
3 Overview of Ave
Ave combines a tile-based DSL, compile-time invariant validation, and an agentic harness to generate and optimize GPU kernels. Its workflow uses invariant and test validation before profiling, then feeds violations and performance back into planning.
3 Overview of Ave
- Ave’s three components are a tile-based Pythonic DSL with tags and assertions, compile-time invariant analysis, and an agentic harness guided by invariant feedback.The implementation primarily targets AMD MI300X and generalizes to AMD MI350X and NVIDIA H200.Source: §3 · p. 3
- The DSL exposes hardware instructions and compiler policies while abstracting complex memory-layout encodings as tiles.Figure 1 illustrates a flash-attention implementation and visualizes tag propagation for V across memory levels.Source: §3 · p. 4, Figure 1 · p. 4
- The compiler lowers the DSL to MLIR, propagates tags through assignments and shared memory, and uses an SMT solver to return concrete witnesses for failed assertions at compile time.The analysis is flow-sensitive and path-insensitive to keep solving tractable, with zero runtime overhead.Source: §3 · p. 4
- A planner retrieves optimization knowledge, a selector orders compatible proposals, and a lowering agent implements them; validation precedes profiling.Invariant violations and performance update the benchmark-local planner through ICRL.Source: §3 · p. 4
4 The Ave DSL
The Ave DSL represents memory accesses as tiles whose layouts map logical coordinates to physical offsets, and it provides operations for allocation, reinterpretation, and dispatch. Tag functions and assertions track operand relationships, while hardware and compiler primitives expose low-level optimization controls.
4 The Ave DSL
- A tile pairs a memory region with a layout function that maps multidimensional coordinates to linear offsets.For integer shapes and strides, the layout computes a coordinate–stride sum; nested layouts also represent hardware-swizzled layouts and dispatch across threads or wavefronts.Source: §4 · pp. 4–5
- The DSL allocates register- and shared-memory-backed tiles and can reinterpret a tile through view(), which requires source and destination layouts to cover identical memory sizes.The example partitions Q by Br and accesses it in 16-byte chunks, a common way to use GPU memory bandwidth.Source: §4 · p. 5
- For the illustrated V transpose, each thread loads a 4×2 subtile and scatters halves into transposed shared-memory positions, feeding the PV loop with swizzled loads.The organization uses 4-byte global loads, 8-byte shared-memory stores, and 16-byte shared-memory loads.Source: §4 · p. 5
- Tag functions attach symbolic values to tensor elements and propagate them through assignments; assertions require conformity or non-conformity between tags.In the flash-attention example, tag matches check that Q and K elements belong to the same head and reduction position, and that P and V align on their shared reduction dimension.Source: §4 · p. 5
- The DSL exposes thread-level control, synchronization, eager materialization, and compiler scheduling barriers so agents can tune topology and instruction placement.It also exposes vectorized instructions and register-class choices, including AGPRs, to give explicit control over hardware features.Source: §4 · pp. 5–6
5 Validating Data-Flow Invariants
Ave propagates symbolic tags through data and control flow, merging them at branches and checking invariant assertions with SMT constraints. Its bounded memory model and compile-time assumptions limit which properties it can express.
5 Validating Data-Flow Invariants
- At control-flow joins, identical concrete tags remain precise, ⊥ conforms with any concrete tag, and differing concrete tags merge to ⊤.Constants carry ⊥; Ave can also reset shared-memory tags to ⊥ when storage is reused across pipeline stages.Source: §5 · p. 6
- Byte-granularity tracking supports evaluated byte-aligned accesses and packed storage, but treats packed FP4 values within a byte as one unit.The paper does not evaluate FP4 kernels, and individual nibble permutations are invisible to the invariants.Source: §5 · p. 6
- Ave does not model the full heap or track global-memory writes, so it cannot express invariants that depend on heap-resident values or currently tracked multi-GPU buffers.For example, it cannot assert that MoE tuples reach the correct shared-memory bucket; producer–consumer queue modeling remains future work for cross-GPU buffers.Source: §5 · p. 6
- The flow-sensitive, path-insensitive analysis encodes tags and assertions as SMT constraints and returns concrete assignments witnessing failed checks.Witnesses can include thread IDs, loop indices, and tensor coordinates; path-insensitive merging avoids enumerating branch conditions.Source: §5 · p. 6
- Tags and validation run entirely at compile time with no runtime overhead, while the prototype requires compile-time-known workgroup dimensions and unrolls loops before analysis.Fixed workgroup dimensions let Ave model registers as a fixed-size array indexed by thread ID.Source: §5 · p. 6
6 Agentic Optimization Harness
Ave repeatedly proposes GPU optimizations, lowers selected proposals into code with instantiated invariants, validates candidates, and profiles those that pass. Its benchmark-local planner learns from optimization outcomes while the harness retains the fastest validated kernel.
6 Agentic Optimization Harness
- The harness retains the fastest candidate that passes invariant checks and unit tests, replacing the current kernel only when profiling measures a speedup.Each round proposes candidates, constructs an ordered plan, lowers it, validates it, and profiles passing candidates.Source: §6 · pp. 6–7
- The persistent expert-curated knowledge base stores hardware facts, DSL references, and recipes with applicability conditions, transformations, hardware controls, and invariant templates.Entries are curated from documentation, vendor examples, public pull requests, and microbenchmarks, then filtered for the request, GPU, and current kernel.Source: §6 · p. 6
- The benchmark-local planner observes the current best kernel, retrieved knowledge, and feedback history, then produces ranked proposals with optimizations, preconditions, target regions, invariant templates, and expected benefits.A selector orders compatible proposals into a plan for the lowering agent.Source: §6 · pp. 6–7
- The lowering agent binds selected instructions to code, specializes parameters, and materializes invariant templates as tag functions and assertions; validation compiles the candidate, checks invariants, and runs unit tests.Passing candidates are profiled, and their transformation and outcome are appended to the next round's feedback history.Source: §6 · p. 7
- The reward combines compilation status, invariant score, test status, and runtime improvement; runtime improvement is measured only after compilation, all invariant checks, and unit tests pass.An invariant score is computed only when compilation succeeds and assertions exist; otherwise it is zero. A validated slowdown receives a negative performance term.Source: §6 · p. 7
- After each round, PolicyEval summarizes the record, Analyze produces a textual critique, and ParameterUpdate edits the benchmark-local prompt without changing the LLM weights or knowledge base.The example uses an invariant witness to trace mismatched MFMA operands to a flattened LDS access, after which a revised binding passed checks and was profiled.Source: §6 · pp. 7–8
7 Implementation
Ave combines a Python DSL front end with an extended MLIR compiler, custom tag-analysis passes, SMT solving, and a JIT layer.
7 Implementation
- The prototype includes about 22,000 lines of C++ and 7,000 lines of Python across its DSL compiler, analysis, harness, and JIT infrastructure.Its front end parses DSL programs into an AST, then lowers through an extended MLIR dialect and LLVM IR to GPU executables.Source: §7 · p. 8
- Custom MLIR passes propagate tags and emit SMT constraints for data-flow invariants, which Z3 solves.Source: §7 · p. 8
- The JIT propagates compile-time constants, caches compiled binaries, and exposes kernel invocation through Python bindings.The JIT layer follows Triton's model.Source: §7 · p. 8
8 Discussion
Ave's invariants target low-level transformations rather than algorithm selection, with additional limits from memory modeling, analysis scope, and reward-hacking defenses. The discussion also describes low-level compiler controls and extensions to other GPU architectures.
8 Discussion
- Data-flow invariants validate concrete transformations but do not propose algorithmic changes or search-space choices such as tile sizes and wavefront specialization.The paper suggests applying invariants after algorithm selection and initial kernel generation, because using them while the algorithm changes could constrain exploration.Source: §8 · p. 8
- Ave checks shared-memory layouts and operand pairing with tractable analysis, while properties involving global-memory writes, including MoE routing and multi-GPU communication, require new abstractions.The analysis tracks byte-aligned accesses and treats two packed FP4 values within a byte as one unit, making individual nibble permutations invisible.Source: §8 · p. 8
- Higher token use alone does not close the expert-library gap without invariants, as lowering agents can exhaust their contexts on trial-and-error repairs.Stronger models often improve generated GPU kernels.Source: §8 · p. 8
- Manual reward-hacking inspection and data-flow invariants are both non-exhaustive: inspection may miss shortcuts outside known patterns, and invariants cannot detect behavior outside their vocabulary.Inspection targets shortcuts such as PyTorch fallbacks, precomputed or zeroed outputs, and branches on benchmark values.Source: §8 · p. 8
- Ave exposes low-level controls such as inline assembly, register-class selection, and scheduling barriers as policies chosen by LLMs and implemented by the compiler.The discussion notes that higher-level DSL automation can eliminate manual optimizations, while peak performance still requires low-level controls.Source: §8 · p. 8
- Ave remains syntactically close to existing DSLs so LLMs can generalize through in-context learning and adopt new frontier models without retraining.The paper notes that better compiler automation could absorb more low-level decisions.Source: §8 · p. 8
- Porting Ave to a new GPU requires toolchain and intrinsic integration plus hardware-specific knowledge-base entries; experts must encode new techniques.The paper reports initial NVIDIA Hopper and AMD CDNA4 support took approximately one person-week, with localized changes and architecture-specific entries.Source: §8 · p. 8
- The paper plans to model TMEM as shared memory using a lifetime heuristic, propagating tags through reads and writes without changing the core analysis.The heuristic reflects optimized kernels typically allocating TMEM near entry and releasing it near exit.Source: §8 · p. 9
9 Evaluation
Ave is evaluated on MI300X production kernels, flash-attention optimization, cross-architecture and KernelBench generality, and invariant effectiveness. Across these studies, results compare throughput, optimization behavior, and reliability.
9 Evaluation
- The evaluation uses MI300X and H200 servers and evaluates agentic frameworks with GPT-5.6 Sol, DeepSeek V4 Pro, and DeepSeek V4 Flash.Only valid kernels—those compiling, passing invariant checks and correctness tests, and passing manual reward-hacking review—contribute to performance.Source: §9 · p. 9
9.1 Kernel Performance
- With GPT, Ave reaches 0.89–0.99× the expert references and 1.62–1176× the uncontaminated agentic baselines in geometric-mean throughput.The references are hipBLASLt for GEMM and AITER for attention and MoE.Source: §9.1 · p. 10
- For Flash Attention, Ave matches AITER’s handwritten assembly and reaches normalized throughput of 6.6× KSearch, 1176× CUDAForge, and 3× KernelFalcon; it reaches 0.89× hipBLASLt.On fused MoE, Ave reaches 8× KernelFalcon and 51× CUDAForge; attention invariant checks take only 10–40 ms.Source: §9.1 · p. 11
9.2 Ablation Study of Ave’s Optimization
- Across five cumulative flash-attention stages, throughput rises monotonically; +Pipe contributes 75–116% across sequence lengths and 100% by geometric mean.The final configuration reaches 2.7–3.6× Naïve throughput.Source: §9.2 · p. 11, Figure 3 · p. 11
9.4 Effectiveness and Cost of Data-Flow Invariants
- Invariants improve success in every dataset–model pair and generally reduce turns, tokens, and API cost, though their effect on wall time is mixed.The study selects KernelBench GEMM and convolution tasks requiring MFMA, software pipelining, and hardware out-of-bounds-guarded loads.Source: §9.4 · p. 12
10 Related Work
Prior work spans agentic kernel-generation workflows, GPU DSLs and compilers, and information-flow invariants; Ave combines these directions with low-level control and compile-time data-flow feedback.
Synthesizing GPU Kernels with LLMs
- Ave uses low-level hardware and compiler intrinsics as dense process rewards, unlike systems relying only on end-to-end correctness and performance rewards.The DSL exposes control for optimizations such as instruction scheduling.Source: §10 · p. 13
ICL and ICRL
- Ave uses ICRL to learn planner prompts that bind optimization knowledge to kernels, with invariant violations and runtime performance serving as process rewards.KernelBlaster also applies ICRL to CUDA kernel optimization, using past optimization trajectories for cross-task generalization [26].Source: §10 · p. 13
DSL Designs for GPU Kernels
- Ave combines tile abstractions with explicit compiler controls, while its restricted memory model keeps static analysis tractable and compile-time invariants provide structural feedback.The DSL's Triton-like syntax lets LLMs reuse knowledge from training.Source: §10 · p. 13
Optimizing Compilers for Tensor Programs
- Ave complements tensor compilers by exposing the SIMT programming model and low-level intrinsics, while delegating schedule exploration to LLMs.Tensor compilers use heuristics and exhaustive search for schedules.Source: §10 · p. 13
Invariants and Information Flow Systems
- Ave adapts information-flow tracking and enforcement from security systems to GPU optimization, using tag propagation to check operand pairings and expose cross-layer constraints.Prior systems such as HiStar and JFlow enforce security policies.Source: §10 · p. 13
- Ave uses lightweight static analysis rather than formal verification to report actionable invariant failures for automatic agent iteration.The prototype uses Z3 for layout algebra; the passage notes that formal verification can multiply sequential-program development effort by 2–6×.Source: §10 · p. 13
11 Conclusion
Ave uses compile-time data-flow invariants to make LLM-generated GPU optimizations more reliable. Its evaluation reaches near-library performance on AMD MI300X and generalizes to NVIDIA H200.
11 Conclusion
- Compile-time data-flow invariants make LLM-generated GPU optimizations more reliable.The tile-based DSL and static analysis turn violations into concrete witnesses that guide agents through transformations without runtime overhead.Source: §11 · p. 13