ARGUS: Agentic GPU Optimization Guided by Data-Flow Invariants

1. Problem Statement

Large Language Model (LLM) based coding agents have demonstrated the ability to generate functionally correct GPU kernels. However, they fail to achieve peak performance comparable to manually optimized libraries (e.g., cuBLAS, AITER) for critical workloads like Matrix Multiplication (GEMM), Flash Attention, and Mixture-of-Experts (MoE).

• The Gap: On the AMD MI300X, agent-generated kernels are 2–600× slower than hand-optimized assembly.
• Root Cause: High-performance GPU kernels require tightly coupled, global optimizations (tiling, shared-memory staging, software pipelining, instruction scheduling, and register allocation). Existing agents rely on sparse pass/fail feedback from unit tests. When a kernel fails to meet performance targets or violates global constraints (e.g., data misalignment in shared memory), agents cannot diagnose the specific violation. They lack "dense" feedback to guide the complex, non-local reasoning required to fix these issues without breaking correctness.

2. Methodology: The Argus Framework

Argus is an agentic framework designed to bridge this gap by introducing Data-Flow Invariants as a mechanism for compile-time verification and dense feedback. It consists of three core components:

A. The Argus DSL (Domain Specific Language)

Argus uses a tile-based, Pythonic DSL (inspired by CuTe and TileLang) that exposes hardware instructions and compiler policies while hiding low-level representations.

• Tag Functions: These attach symbolic tags (e.g., logical coordinates) to tensor elements. Tags propagate through control and data flow, allowing the system to track the "identity" and "location" of data elements as they move from global memory to shared memory and registers.
• Tag Assertions: These are compile-time checks (e.g., assert tag(A) == tag(B)) that enforce relational constraints at use sites. They verify that data is correctly paired (e.g., operands for a matrix multiply match) or separated (e.g., concurrent producers do not overwrite the same memory).
• Expressiveness: The DSL supports vectorized loads, explicit instruction scheduling, and register class selection (e.g., AGPRs), enabling agents to explore low-level optimizations.

B. Compile-Time Static Analysis

The Argus compiler validates data-flow invariants entirely at compile time, incurring zero runtime overhead.

• Mechanism: It performs a flow-sensitive, path-insensitive analysis using abstract interpretation over a layout algebra.
• SMT Solving: Tag assertions are encoded as constraints and solved using an SMT solver (Z3).
• Feedback Loop: If an invariant is violated, the compiler does not just return "False." It generates a concrete counterexample identifying the specific thread, data element, and program point causing the violation. This provides the LLM agent with actionable, structured feedback to target fixes precisely.

C. Agentic Optimization Harness (ICRL)

Argus employs an In-Context Reinforcement Learning (ICRL) planner to iteratively optimize kernels.

• Workflow: The agent inspects the kernel, invariant violation reports, and runtime profiles. It retrieves optimization strategies from a curated Knowledge Base (containing expert-verified patterns like software pipelining or bank conflict mitigation).
• Learning: The planner treats its prompt as a mutable policy. It uses invariant violations and runtime performance as reward signals. Through text gradients, the agent learns to select effective optimizations and synthesize the necessary invariants to guide intrusive transformations (like reorganizing data layouts) without breaking correctness.

3. Key Contributions

1. Data-Flow Invariants for Agents: Introduced the concept of using data-flow invariants (adapted from information flow security) to guide LLM agents in GPU kernel optimization, providing dense, structured feedback rather than sparse test results.
2. Tag-Based DSL & Static Analysis: Designed a tile-based DSL with tag functions and assertions, backed by a zero-overhead static analysis that verifies global correctness properties at compile time.
3. ICRL Optimization Harness: Built an agentic system that uses ICRL to learn how to bind optimization knowledge to specific kernels and synthesize invariants, creating a closed loop between code generation, verification, and optimization.
4. State-of-the-Art Performance: Demonstrated that Argus can generate kernels that match or exceed the performance of hand-optimized assembly libraries.

4. Experimental Results

The evaluation was conducted on the AMD MI300X GPU across three dominant LLM inference workloads: GEMM, Flash Attention, and MoE.

• Performance vs. Hand-Optimized Libraries:
  • Argus-generated kernels achieved 99–104% of the effective throughput (TFLOPS) of state-of-the-art hand-optimized assembly libraries (e.g., AITER, HipBlasLt).
  • Specifically, Argus outperformed the hand-written AITER Flash Attention kernel by 3.5%.
• Performance vs. Existing Agents:
  • Argus delivered geometric-mean throughput gains of 2× to 1543× over existing agentic systems (e.g., KernelFalcon, KSearch, CUDAForge).
  • For Flash Attention, Argus was 2.4× faster than KernelFalcon.
  • For Fused MoE, Argus was 1543× faster than CUDAForge (which failed to utilize hardware matrix cores correctly).
• Generalization (KernelBench):
  • Argus successfully generated correct kernels for 100% of Level 1 and 90% of Level 2 problems in the KernelBench suite.
• Ablation Study:
  • Removing data-flow invariants reduced the success rate (Pass@1) from 75% to 60% (Level 1) and 57% to 40% (Level 2), while increasing token usage. This confirms that invariants are critical for guiding complex, intrusive optimizations.

5. Significance

• Bridging the Performance Gap: Argus proves that LLM agents can reach the performance of human experts in GPU kernel optimization, a domain previously thought to require months of manual engineering.
• Shift in Verification Strategy: It moves the paradigm from "test-after-generation" to "verify-during-generation." By providing concrete counterexamples for global constraints, it allows agents to reason about correctness at an abstract level rather than relying on trial-and-error.
• Scalability and Adaptability: The framework is not limited to AMD hardware; the abstractions (DSL, invariants, ICRL) are generalizable to other GPU architectures (e.g., NVIDIA) and other domains where global data-flow properties are critical.
• Efficiency: By automating the generation of highly optimized kernels, Argus has the potential to significantly reduce the billions of dollars in infrastructure costs associated with inefficient LLM inference.