Ave: Guiding Agentic GPU Optimization Using Data-Flow Invariants

arXiv:2604.18616 · cs.DC, cs.AI, cs.PL · Submitted 2026-04-16 · Read on arXiv

Listen

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: Next we'll be talking about the paper "Ave: Guiding Agentic GPU Optimization Using Data-Flow Invariants".

Jane: The paper was written by Haohui Mai, Xiaoyan Guo, Xiangyun Ding, Daifeng Li, Qiuchu Yu et al. from CausalFlow Inc. and Hong Kong University of Science and Technology and Tsinghua University and Stanford University and University of Chinese Academy of Sciences and University of California, Riverside.

Tom: Stay tuned as we take you through the paper and discuss its implications.

Title: Tom: We are starting our show with a real powerhouse of a paper today.

Jane: It really is, Tom, and it's called "ARGUS: Agentic GPU Optimization Guided by Data-Flow Invariants."

Tom: That title is quite a mouthful for our listeners!

Jane: It's because the researchers are doing something incredibly complex with hardware.

Tom: Are we talking about making computers faster?

Jane: Exactly, but specifically making the chips that run AI much more efficient.

Tom: I see a huge list of names here from places like HKUST and Tsinghua University.

Jane: They've pulled together experts from Stanford and UC Riverside too to tackle this.

Tom: It sounds like a massive global effort to fix how we use GPUs.

Jane: It is, because they want to stop relying on humans to write these perfect little instructions for the hardware.

Lu: This is such a beautiful way to think about it!

Tom: What do you mean by beautiful, Lu?

Lu: Instead of just telling an AI to "write code," they are giving it a set of logical rules that act like a compass.

Jane: That's a great way to put it, Lu.

Meng: I have to ask though, is this actually going to work in a real data center?

Tom: That's the million-dollar question, Meng.

Meng: Because right now, if an AI makes one tiny mistake in a GPU kernel, the whole system might just crash or run incredibly slowly.

Jane: That's exactly why they added those "invariants" to the title.

Lu: It's like giving a chef a recipe that also tells them exactly how the heat should move through the pan!

Lalam: And when we think about that, it changes our relationship with technology.

Tom: How so, Lalam?

Lalam: We are moving from a world where humans have to micromanage every single electron to a world where AI understands the fundamental laws of its own playground.

Jane: That's a very profound way to look at it.

Tom: It's definitely more than just a simple coding tool.

Jane: We should probably look at what they actually built to see if that vision holds up.

Summary: Tom: Now that we've cleared up the name, let's get into the actual meat of how ARGUS works.

Jane: Right, it isn't just a standard AI writing code; it uses these "data-flow invariants."

Tom: Can you break that down for us without using too much jargon, Jane?

Jane: Sure! Think of an invariant as a rule that says "this piece of data must always stay paired with this other piece."

Tom: So it's like a constant check to make sure nothing gets lost in transit?

Jane: Precisely, and they use a special language, or DSL, to make these rules easy for the AI to follow.

Tom: And I see they use something called an SMT solver in the process.

Jane: That's the part that double-checks everything at compile time so there's no slowdown when you actually run it.

Tom: It seems like a very tight loop between planning and checking.

Jane: It is, with an "in-context reinforcement learning" planner that learns from its own mistakes.

Lu: I love the idea of the AI having its own internal feedback loop like this!

Tom: What's so special about that, Lu?

Lu: Most AIs just guess and check if it works, but ARGUS actually understands *why* a mistake happened because the compiler tells it exactly where the rule was broken.

Meng: That sounds much more reliable than what I'm seeing in most AI coding tools today.

Jane: It really is, Meng, because it gives the AI "dense feedback" instead of just a simple pass or fail.

Meng: So, if it's trying to optimize a complex math operation and fails, it actually knows which specific thread caused the problem?

Jane: Yes, that's exactly what the paper says.

Tom: That sounds like it could save engineers an incredible amount of time.

Lalam: It also means we are teaching AI to have a sense of structural integrity.

Tom: How does that impact the way we view AI progress, Lalam?

Lalam: It moves us past the era of "probabilistic guessing" and into an era of "verifiable reasoning" in software.

Jane: That's a huge distinction.

Tom: Let's see if those reasoning skills actually lead to any massive speed improvements.

Summary: Tom: We've seen the mechanics, now let's talk about the actual performance leaps they found.

Jane: The numbers in this paper are honestly hard to wrap your head around.

Tom: They aren't just saying it's "faster," right?

Jane: No, they achieved between ninety-nine percent and one hundred four percent of the speed of hand-optimized assembly code.

Tom: Wait, so an AI can almost match what the world's best human engineers do by hand?

Jane: Almost exactly, and they did it on the AMD MI300X GPU.

Tom: And compared to other AI agents, the speedup is even crazier.

Jane: It can be up to one thousand five hundred forty-three times faster than existing agentic systems in some cases!

Tom: That's not just a small improvement; that's a total shift in the landscape.

Jane: They tested it on the big stuff like GEMM, flash attention, and MoE kernels.

Lu: It’s like watching a student suddenly start performing at a professional level overnight!

Tom: What do you think is the secret sauce there, Lu?

Lu: It's that they aren't just optimizing the code; they are optimizing the way data flows through every single layer of the hardware.

Meng: I'm looking at these results and wondering about the cost-to-benefit ratio.

Jane: Are you worried about how much computing power it takes to run this agent, Meng?

Meng: Not exactly, but I want to know if this can be applied to different types of chips easily.

Jane: The authors say the design is generalizable, even though they used AMD for the tests.

Meng: If it can move from AMD to NVIDIA without a total rewrite, that's a game changer for us.

Tom: It sounds like it could make AI training much cheaper in the long run.

Lalam: And that has massive cultural implications for how accessible AI becomes.

Tom: How do you see that playing out, Lalam?

Lalam: If we can make the most expensive part of AI—the hardware efficiency—much more affordable, then computing power becomes a utility like water or electricity.

Jane: That's a powerful vision for the future.

Tom: We should probably wrap this up and see what the final verdict is.

Conclusion: Tom: We are coming to the end of our deep dive into "ARGUS: Agentic GPU Optimization Guided by Data-Flow Invariants."

Jane: It has been a fascinating look at how we can bridge the gap between AI and high-performance hardware.

Tom: We've seen how these data-flow invariants turn a guessing game into a precise science.

Jane: And we've seen that the performance gains are absolutely massive compared to previous methods.

Lu: I just think we are standing on the edge of a new era where software builds itself with perfect logic!

Meng: From my side, if this scales, it's going to drastically change how we deploy models in production.

Lalam: It's really about democratizing the ability to run intelligence efficiently across the globe.

Tom: Well, thank you all for joining us today.

Jane: Thanks for listening, and we'll see you next time with another incredible paper!

Tom: Goodbye everyone!

CausalFlow Inc. · Hong Kong University of Science and Technology · Tsinghua University · Stanford University · University of Chinese Academy of Sciences · University of California, Riverside

cs.DC, cs.AI, cs.PL

Submitted: 2026-04-16

Updated: 2026-09-14

DOI: 10.1145/3830418.3843902

Code: https://github.com/ROCm/aiter

License: http://creativecommons.org/licenses/by/4.0/

Importance score: 85/100

The gist: Argus is an agentic framework designed to automate complex, multi-level GPU kernel optimizations by addressing the performance gap between LLM-generated code and hand-optimized libraries.

Key concepts

Data-flow invariants: These are logical rules that ensure specific pieces of data remain correctly paired or handled throughout a process. They act like a compass for the AI, providing structural integrity and ensuring the code follows fundamental laws during complex hardware optimization tasks instead of relying on probabilistic guessing.
SMT solver: This tool is used to double-check all rules at compile time to prevent performance slowdowns when the code actually runs. By verifying that the AI's instructions adhere to required invariants before execution, it ensures the resulting GPU kernels are both correct and highly efficient.
In-context reinforcement learning: This is a planning mechanism where the AI learns from its own mistakes using "dense feedback." Unlike standard AIs that simply guess if code works, this system understands exactly why a mistake happened because the compiler identifies precisely where a specific rule was broken.

Terminology

Core — 21× in this paper · explained in 2 episodes
Reinforcement Learning — 11× in this paper · explained in 4 episodes
Agentic Framework — 10× in this paper · explained in 6 episodes
Abstraction — 8× in this paper · explained in 1 episode
Verification — 8× in this paper · explained in 2 episodes
Agentic Optimization — 5× in this paper · explained in 1 episode
Formal Verification — 5× in this paper · explained in 2 episodes
SMT solver — 5× in this paper · explained in 2 episodes
Agentic System — 4× in this paper · explained in 2 episodes
Adam — 3× in this paper · explained in 2 episodes
Benchmarking — 3× in this paper · explained in 4 episodes
Ablation — 2× in this paper · explained in 1 episode
Fine-Tuning — 2× in this paper · explained in 5 episodes
Generalizability — 2× in this paper · explained in 6 episodes
Generalization — 2× in this paper · explained in 13 episodes
In-Context Learning — 2× in this paper · explained in 4 episodes
Mixture-of-Experts (MoE) — 2× in this paper · explained in 7 episodes
Reward Hacking — 2× in this paper · explained in 8 episodes
World Model — 2× in this paper · explained in 7 episodes
AI Agents — 1× in this paper · explained in 2 episodes
Adaptability — 1× in this paper · explained in 1 episode
Agentic Coding — 1× in this paper · explained in 1 episode
Agentic Loop — 1× in this paper · explained in 1 episode
Agentic Workflows — 1× in this paper · explained in 1 episode
Alignment — 1× in this paper · explained in 1 episode
Attention Mechanism — 1× in this paper · explained in 3 episodes
Benchmark — 1× in this paper · explained in 2 episodes
Downstream Tasks — 1× in this paper · explained in 2 episodes
In-Context Learning (ICL) — 1× in this paper · explained in 5 episodes
Quantization — 1× in this paper · explained in 12 episodes
Robustness — 1× in this paper · explained in 5 episodes
Sparsity — 1× in this paper · explained in 4 episodes
Structural Integrity — 1× in this paper · explained in 3 episodes
Verifier — 1× in this paper · explained in 2 episodes

Summary

Argus is an agentic framework designed to automate complex, multi-level GPU kernel optimizations by addressing the performance gap between LLM-generated code and hand-optimized libraries. It matters because even single-digit improvements in kernel efficiency can translate into savings on the order of billions of dollars in large-scale LLM deployment.

The Argus DSL and Data-Flow Invariants

Argus introduces a tile-based, Pythonic DSL that exposes hardware instructions and compiler policies while hiding low-level representations to maintain expressivity and learnability for LLMs. The core innovation is the use of data-flow invariants, which are compile-time specifications that encode how data must be choreographed throughout a kernel’s execution. This allows the agent to navigate intrusive transformations such as software pipelining and data layout reorganization with confidence.

The DSL utilizes two primary mechanisms:

Tag functions: These attach symbolic tags (such as logical coordinates) to tensor elements and propagate them through control and data flow.
Tag assertions: These specify whether tags at a use site must match or must differ, allowing for the enforcement of relational constraints like conformity or non-conformity.

Static Analysis and Feedback

To ensure correctness without performance penalties, Argus employs static analysis that enforces these invariants entirely at compile time with zero runtime overhead. The compiler tracks tag propagation via abstract interpretation over a layout algebra and utilizes an SMT solver to discharge integer constraints. To remain tractable on heavily optimized kernels, the analysis is path-insensitive.

When an assertion fails, the system provides dense, structured feedback by emitting a concrete counterexample that identifies the thread, data element, and program point. This mechanism allows the agent to move beyond patching implementations by trial and error toward reasoning about correctness at an abstract level.

The Agentic Optimization Harness

The framework utilizes an agentic optimization harness that iteratively refines kernel implementations through an in-context reinforcement learning (ICRL) planner. This planner learns to select optimizations and draft effective invariants by interacting with a curated knowledge base of GPU-specific optimization techniques. The system treats planner prompts as mutable parameters and uses invariant violations and runtime performance as reward signals.

The harness comprises five essential components:

A persistent knowledge base documenting generic and GPU-specific optimizations.
An LLM-powered planner that produces ranked proposals of optimization, context, and score.
An optimization selector that samples from the planner’s proposals to maintain exploration in a tightly coupled optimization space.
A lowering agent that implements transformations guided by data-flow invariants and few-shot examples.
A validator agent that checks invariants at compile time, runs unit tests, and collects runtime profiles.

Performance and Generalization

Evaluations on the AMD MI300X GPU demonstrate that Argus-generated kernels achieve 99–104% of the effective throughput of state-of-the-art hand-optimized assembly implementations. In critical workloads like GEMM, flash attention, and MoE—which account for over 90% of GPU time in LLM inference—Argus is 2–1543× faster than existing agentic systems in geometric-mean throughput.

The framework also demonstrates significant generalizability, producing correct kernels for 100% of Level 1 and 90% of Level 2 problems on the KernelBench benchmark.

Improvements for AI systems

1. Integration of a Symbolic Data-Flow Invariant Feedback Loop

Capability: Instead of relying on sparse, binary pass/fail unit test feedback, the AI agent receives dense, structured, and mathematically grounded counterexamples. When an optimization (e.g., software pipelining or layout reorganization) violates a global constraint, the system identifies the exact thread, data element, and program point of the violation. This allows the agent to perform surgical, single-step code corrections rather than inefficient, trial-and-error debugging.

2. Implementation of an In-Context Reinforcement Learning (ICRL) Planner with a Persistent Optimization Knowledge Base

Capability: The agent can autonomously select and instantiate complex, multi-level optimization strategies (such as Split-K partitioning, asynchronous memory copies, or bank conflict mitigation) by retrieving proven patterns from a curated knowledge base. By treating the optimization prompt as a mutable, learnable parameter, the system learns to map high-level architectural knowledge to specific kernel topologies, significantly improving cross-task generalization.

3. Transition to a Tile-Based Symbolic DSL for Agentic Reasoning

Capability: The AI operates on a high-level, tile-based language that uses symbolic tags (logical coordinates) rather than raw memory offsets. This prevents context window exhaustion by drastically reducing the token count required to describe complex kernels. It also enables the agent to maintain global correctness properties—such as ensuring operands are correctly swizzled for Matrix Cores—even through highly intrusive transformations like register-level data shuffling.

4. Exposure of Hardware-Specific Policy Interfaces in the Agent's Action Space

Capability: The agent gains direct, fine-grained control over low-level hardware knobs, including accumulator register class (AGPR) selection, software-defined scheduling barriers, and hardware-guarded out-of-bounds memory loads. This allows the AI to bridge the abstraction gap inherent in high-level compilers, enabling it to achieve 99–104% of the effective throughput of hand-optimized assembly libraries for critical workloads like GEMM, Flash Attention, and Mixture-of-Experts (MoE).

Abstract

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.