Source note
ARGUS: Agentic GPU Optimization Guided by Data-Flow Invariants
Summary
Argus is an agentic system for generating high-performance GPU kernels by giving the coding agent compile-time data-flow invariants instead of relying only on pass/fail tests. On AMD MI300X, it reports kernels for GEMM, flash attention, and MoE that match hand-optimized assembly throughput and beat prior agentic systems by large margins.
Problem
- Existing LLM coding agents can write functionally correct GPU kernels, but their performance is often far below hand-optimized libraries on workloads such as GEMM, attention, and MoE.
- Peak GPU performance depends on many coupled choices at once: tiling, shared-memory staging, software pipelining, instruction scheduling, register use, and data layout transformations. Unit tests do not show which global constraint a kernel violated.
- This matters because these kernels account for most LLM inference time, and even small efficiency gains can cut large GPU costs at deployment scale.
Approach
- Argus adds data-flow invariants to GPU kernel generation. These are compile-time rules that specify how data elements should stay aligned as the kernel tiles, reorders, stages, and computes on them.
- It uses a tile-based Python DSL with tag functions and tag assertions. Tag functions attach symbolic labels such as logical coordinates to tensor elements. Tag assertions check that the right elements meet at use sites, for example in matrix-core instructions.
- The compiler verifies these invariants at compile time with abstract interpretation over layout algebra plus SMT solving. If a check fails, it returns a concrete counterexample with the thread, data element, and program point. The paper says this adds zero runtime overhead because tags exist only in compilation.
- An agentic optimization harness uses an LLM planner plus in-context reinforcement learning to choose optimizations, write invariants, compile, test, profile, and revise the kernel using invariant violations and runtime performance as feedback.
Results
- On AMD MI300X, Argus reports 99–104% of the effective throughput of state-of-the-art hand-optimized assembly for GEMM, flash attention, and MoE kernels.
- Against existing agentic baselines, it reports 2–1543× higher geometric-mean throughput across the evaluated workload families.
- For flash attention, it reports a 2.4× speedup over KernelFalcon, named as the strongest baseline.
- The evaluated kernel families account for over 90% of GPU execution time in LLM inference on the target platform, so the reported gains target major serving bottlenecks.
- On KernelBench, Argus reports correct kernels for 100% of Level 1 tasks and 90% of Level 2 tasks.
- The excerpt does not provide deeper per-dataset breakdowns, ablations, or exact baseline numbers for every workload family beyond the headline comparisons above.