Skip to content

ARGUS: Agentic GPU Optimization Guided by Data-Flow Invariants

🕒 Published (v1): 2026-04-16 15:49 UTC · Source: Arxiv · link

Why this paper was selected

Agentic GPU kernel optimization via data-flow invariants; closes performance gap

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

💬 Ask ChatGPT✦ Ask Claude

TL;DR

ARGUS is an agentic GPU kernel optimization framework that introduces compile-time data-flow invariants—encoded as tag functions and tag assertions in a tile-based Pythonic DSL—to provide dense, structured feedback to an LLM agent navigating complex, coupled hardware optimizations. On AMD MI300X, generated kernels for GEMM, flash attention, and MoE reach 99–104% of hand-optimized assembly library throughput and run 2–1543× faster than existing agentic baselines.

Problem

LLM-based coding agents produce functionally correct GPU kernels but fail to reach competitive performance on critical workloads (GEMM, attention, MoE). The root cause is that peak GPU throughput requires coordinated reasoning over tiling, shared-memory staging, software pipelining, register allocation, and instruction scheduling—a globally constrained optimization space—while existing agents receive only sparse pass/fail signals from unit tests and cannot diagnose which global constraint was violated or where.

Method

ARGUS combines three mutually reinforcing components:

  1. Tile-based DSL with data-flow invariants. Kernels are written in a Pythonic, CuTe/TileLang-inspired DSL. Tag functions attach symbolic annotations (e.g., logical coordinates) to tensor elements and propagate them through data and control flow. Tag assertions (conformity: tags must match; non-conformity: tags must differ) encode global correctness properties at use sites, e.g., asserting that paired MFMA operands always carry compatible coordinate tags.

  2. Compile-time invariant validation at zero runtime cost. The compiler lowers DSL programs to MLIR and runs a flow-sensitive, path-insensitive abstract interpretation over a layout algebra. Tags are purely symbolic; they are never materialized at runtime. At control-flow joins, tags are merged under the partial order \(\bot < t < \top\). Assertion failures are encoded as SMT integer constraints and discharged by a solver; violations yield concrete counterexamples (thread ID, data element, program point).

  3. Agentic optimization harness with ICRL. Five subcomponents operate in a closed loop: (i) a persistent knowledge base of GPU-specific optimization skills and their expected invariants, curated by domain experts; (ii) a learnable LLM planner that maps the current DSL kernel, compile-time violation reports, and runtime profiles to ranked proposals \(\langle\text{optimization}, \text{context}, \text{score}\rangle\)—planner prompts are treated as mutable parameters updated via text gradients; (iii) an optimization selector that samples from proposals; (iv) a lowering agent that implements transformations guided by invariants; and (v) a validator agent that runs compile-time checks, unit tests, and profiling.

Key Contributions

  • Introduces data-flow invariants (tag functions + tag assertions) into agentic GPU kernel generation, turning global correctness constraints into dense, actionable compiler feedback.
  • A tile-based Pythonic DSL exposing hardware instructions and compiler scheduling policies (e.g., eager materializations, s_barrier placement, AGPR register selection) while hiding low-level IR.
  • Compile-time invariant validation via abstract interpretation + SMT solving with zero runtime overhead.
  • An agentic optimization harness using ICRL to jointly learn optimization selection and invariant synthesis across tasks.
  • Demonstrates near-parity with hand-optimized assembly on the three kernel families (GEMM, flash attention, MoE) that account for >90% of LLM inference GPU time.

Results

  • GEMM, flash attention, MoE on AMD MI300X: 99–104% of effective TFLOPS of state-of-the-art hand-optimized assembly libraries (AITER, rocBLAS).
  • vs. agentic baselines: 2–1543Ă— geometric-mean throughput gains across evaluated workload families.
  • vs. KernelFalcon (strongest agentic baseline) on flash attention: 2.4Ă— throughput improvement.
  • KernelBench generalization: 100% of Level 1 tasks correct, 90% of Level 2 tasks correct.

Limitations

  • Path-insensitive abstract interpretation: branch-dependent invariant violations may be missed or over-approximated with \(\top\).
  • Prototype requires thread-block dimensions to be compile-time constants; all loops are unrolled before analysis, which may not scale to dynamically-shaped kernels.
  • Knowledge base of optimization skills and invariant templates is curated by domain experts, limiting fully automated bootstrapping for new hardware.
  • Evaluation is restricted to AMD MI300X; generalization to NVIDIA architectures is claimed but not empirically validated in the paper.
  • The ICRL planner updates planner prompts via text gradients, which may be sensitive to prompt initialization and reward signal quality.

Relevance to Harnesses / Meta-Harnesses

ARGUS is a concrete, domain-specific instance of a closed-loop agentic harness: it wraps a DSL compiler, a knowledge base, and multiple specialized LLM agents (planner, optimizer, lowering, validator) into a single iterative pipeline where compile-time and runtime signals serve as structured reward. The use of ICRL to update the planner's prompts as mutable parameters directly demonstrates how a harness can self-improve across tasks—a design pattern directly applicable to meta-harnesses that orchestrate LLM agents over multi-step optimization loops. The knowledge base as persistent long-term memory shared across tasks parallels skill libraries in general-purpose agent harnesses. The key insight—that dense, structured feedback (concrete counterexamples) is strictly more useful than sparse pass/fail signals for agent-guided search—is a transferable principle for harness design in any domain where global invariants govern correctness.