Paper deep dive
ARGUS: Agentic GPU Optimization Guided by Data-Flow Invariants
Haohui Mai, Xiaoyan Guo, Xiangyun Ding, Daifeng Li, Qiuchu Yu, Chenzhun Guo, Cong Wang, Jiacheng Zhao, Christos Kozyrakis, Binhang Yuan
Intelligence
Status: succeeded | Model: Gemma-4-26B-A4B | Prompt: intel-v1 | Confidence: 98%
Last extracted: 4/26/2026, 9:28:52 PM
Summary
Argus is an agentic framework designed to optimize GPU kernels (such as GEMM, Flash Attention, and MoE) by using data-flow invariants. It introduces a tile-based, Pythonic DSL that utilizes tag functions and tag assertions to encode and verify compile-time specifications. These invariants are verified via abstract interpretation and SMT solving, providing dense, structured feedback to an in-context reinforcement learning (ICRL) planner. The framework achieves near-peak performance (99-104% of hand-optimized assembly) on the AMD MI300X GPU and significantly outperforms existing agentic systems.
Entities (7)
Relation Signals (5)
Argus → evaluateson → KernelBench
confidence 100% · Argus further generalizes to 200 KernelBench tasks
Argus → targets → AMD MI300X
confidence 100% · We evaluate Argus on the AMD MI300X GPU
Argus → uses → Data-flow Invariants
confidence 100% · Argus, an agentic framework that addresses this through data-flow invariants
Argus → utilizes → Tag Functions
confidence 100% · The DSL provides tag functions to propagate symbolic annotations
Argus → implements → ICRL Planner
confidence 90% · An in-context reinforcement learning (ICRL) planner learns to select optimizations
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
Abstract:LLM-based coding agents can generate functionally correct GPU kernels, yet their performance remains far below hand-optimized libraries on critical computations such as matrix multiplication, attention, and Mixture-of-Experts (MoE). Peak GPU performance requires coordinated reasoning over tightly coupled optimizations, including tiling, shared-memory staging, software pipelining, and instruction scheduling, while existing agents rely on sparse pass/fail feedback, leaving them unable to diagnose global constraint violations. We present Argus, an agentic framework that addresses this through data-flow invariants: compile-time specifications encoding how data must be choreographed throughout kernel execution. Argus introduces a tile-based, Pythonic DSL exposing hardware instructions and compiler policies while hiding low-level representations. The DSL provides tag functions to propagate symbolic annotations through data and control flow, and tag assertions to enforce relational constraints at use sites. When violations occur, the compiler returns concrete counterexamples identifying the thread, data element, and program point, enabling dense, structured feedback for targeted fixes. Invariants are verified at compile time via abstract interpretation over a layout algebra and SMT solving, with zero runtime overhead. An in-context reinforcement learning planner learns to select optimizations and synthesize effective invariants, supported by a curated knowledge base of GPU optimization techniques. We evaluate Argus on the AMD MI300X GPU across GEMM, flash attention, and MoE kernels accounting for over 90% of GPU time in LLM inference. Generated kernels achieve 99-104% of state-of-the-art hand-optimized assembly throughput and are 2-1543x faster than existing agentic systems. Argus further generalizes to 200 KernelBench tasks, solving 100% of Level 1 and 90% of Level 2 problems.
Tags
Links
- Source: https://arxiv.org/abs/2604.18616v1
- Canonical: https://arxiv.org/abs/2604.18616v1
Trouble viewing inline? Open PDF directly →
Full Text
74,569 characters extracted from source content.
Expand or collapse full text
ARGUS: Agentic GPU Optimization Guided by Data-Flow Invariants Haohui Mai ˇ “ ( Xiaoyan Guo 4 Xiangyun Ding ˇ “ ( ,5 Daifeng Li 1 Qiuchu Yu 4 Chenzhun Guo 4 Cong Wang 2 Jiacheng Zhao 4 Christos Kozyrakis 3 Binhang Yuan 1 CausalFlow Inc. ˇ “ ( HKUST 1 Tsinghua University 2 Stanford University 3 UCAS 4 UC Riverside 5 Abstract Recent LLM-based coding agents can generate functionally correct GPU kernels across diverse workloads, yet their per- formance remains far below that of manually optimized li- braries on critical computations such as matrix multiplica- tion, attention, and Mixture-of-Experts (MoE). This gap is fundamental: peak GPU performance requires coordinated reasoning over tightly coupled optimizations, including tiling, shared-memory staging, software pipelining, and instruc- tion scheduling, while existing agents rely on sparse pass/fail feedback from unit tests, leaving them unable to diagnose violations of global constraints. We present Argus, an agentic framework that addresses this limitation through data-flow invariants, i.e., compile- time specifications that encode how data must be chore- ographed throughout a kernel’s execution. Argus introduces a tile-based, Pythonic DSL that exposes hardware instruc- tions and compiler policies while hiding low-level repre- sentations, maintaining both expressivity and learnability for LLMs. The DSL introduces tag functions to propagate symbolic annotations through data and control flow, and tag assertions to enforce relational constraints at use sites, centralizing global correctness properties. When violations occur, the compiler returns concrete counterexamples that identify the thread, data element, and program point, provid- ing dense, structured feedback for targeted fixes. Invariants are verified at compile time via abstract interpretation over a layout algebra and SMT solving, incurring zero runtime over- head. An in-context reinforcement learning (ICRL) planner further learns to select optimizations and synthesize effec- tive invariants, supported by a curated knowledge base of GPU-specific optimization techniques. We evaluate Argus on the AMD MI300X GPU across GEMM, flash attention, and MoE kernels that together ac- count for over 90% of GPU time in LLM inference. The gen- erated kernels achieve 99–104% of the effective throughput of state-of-the-art hand-optimized assembly implementa- tions and are 2–1543×faster than existing agentic systems in geometric-mean throughput across the evaluated work- load families. Argus further generalizes to 200 KernelBench tasks, producing correct kernels for 100% of Level 1 and 90% of Level 2 problems. 1 Introduction Efficient GPU kernel implementation is critical to large-scale LLM deployment. Given the hundreds of billions of dollars invested in GPU infrastructure, even single-digit improve- ments in kernel efficiency can translate into savings on the order of billions of dollars [61]. A well-optimized matrix mul- tiplication (GEMM) on an NVIDIA A100 GPU can be 50× faster than a naïve implementation [8], but achieving this speedup requires coordinated optimizations across the full software stack: hardware matrix cores [41], software pipelin- ing [34], memory hierarchy exploitation [37], and instruction scheduling [31]. Tensor compilers [15,26,42,59,63,64,72] automate subsets of this space through heuristic and evolu- tionary search, but for peak performance, hardware vendors ship hand-optimized assembly libraries [1,4] that require months of engineering and struggle to keep pace with yearly hardware refreshes and evolving workloads [57]. Recent work applies LLMs and coding agents to GPU ker- nel generation [6,10,17,22,23,29,38,57,62,69], producing functionally correct kernels across a range of workloads [49]. However, the generated kernels are far from competitive compared with hand-optimized code on performance-critical workloads: on our AMD MI300X platform, the best agent- generated CUDA/HIP GEMM kernels are 2–600×slower than the state-of-the-art library. The gap is fundamental: correctness can be verified with tests, but high performance requires coordinated reasoning across software pipelining, instruction scheduling, register allocation, and NUMA-aware memory access, while accounting for the effects of memory aliasing [53] and synchronization [7]. Training data for such low-level optimizations is scarce, and the full kernel optimiza- tion trajectory often exceeds LLM context windows, where performance degrades [39]. Targeting Triton [59] narrows the gap to approximately 2×by offloading optimizations like software pipelining to the compiler, but the remaining 2×, including fine-grained instruction scheduling and warp specialization, lies below Triton’s abstraction level. In this paper, we present Argus, an agentic framework that automates complex, multi-level GPU kernel optimiza- tions. Argus makes these optimizations robust through data- flow invariants: compile-time assertions that specify how 1 arXiv:2604.18616v1 [cs.DC] 16 Apr 2026 input data must be correlated throughout a kernel’s compu- tation. Data-flow invariants are an established concept in in- formation flow security [46,65,67]; Argus adapts such tech- niques to GPU kernel optimization. For example, a GEMM kernel tiles and stages matrices A and B through shared memory before feeding them into hardware MFMA instruc- tions that expect operands in a specific swizzled layout. A data-flow invariant can assert that the MFMA instructions al- ways receive correctly paired elements of A and B, regardless of the transformations applied. By making these invariants explicit, Argus provides the agent with correctness con- straints that guide it through intrusive transformations such as software pipelining and data layout reorganization. Three essential challenges arise in making this approach applicable. First, the invariant programming language must be succinct enough to express global properties in one cen- tral place, yet be able to produce concrete violation reports (e.g., which thread holds a mismatched element, at which program point) to steer the agent toward fixes in subsequent iterations. Second, tracking and enforcing invariants must add zero runtime overhead. A fully optimized kernel satu- rates hardware function units, tailors access patterns to the memory hierarchy, and relies on precise instruction sched- ules; the analysis should operate entirely at compile time across tiling, staging, and layout transformations without perturbing the generated code. Third, the invariants them- selves must be synthesized automatically, since they are application-specific. For instance, how tiles of A and B are assigned to warps, staged in shared memory, and eventu- ally consumed by MFMA instructions varies across kernels. Argus must instantiate these invariants automatically to optimize kernels without human intervention. Argus addresses these challenges with three ideas: tag functions, static analysis, and an in-context RL coding agent. First, Argus introduces tag functions and tag assertions as a concise language for expressing data-flow invariants. Argus generates and optimizes kernels in a tile-based, Pythonic DSL inspired by CuTe [47] and TileLang [63]. The DSL augments tensor operations with tag functions, which attach symbolic tags (such as logical coordinates) to tensor elements and propagate them through control and data flow, and with tag assertions, which specify whether tags at a use site must match or must differ. Conformity assertions express align- ment constraints, for example requiring the paired operands of a matrix multiplication to carry compatible tags, while non-conformity assertions express separation constraints, such as ensuring that concurrent producers do not overwrite the same shared-memory region. Second, Argus applies static analysis that enforces these invariants entirely at compile time and turns failures into ac- tionable feedback for the LLM coding agent. Tags are purely symbolic and are never materialized at runtime. The Ar- gus compiler tracks their propagation via abstract inter- pretation [16] over a layout algebra [47], and uses an SMT solver [19] to discharge the resulting integer constraints. To remain tractable on heavily optimized kernels, the analysis is path-insensitive. When an assertion fails, the compiler emits a concrete counterexample identifying the offending thread, data element, and program point, giving the agent precise guidance for the next revision. Third, Argus leverages in-context reinforcement learning (ICRL) [23,45] to learn to select optimizations and draft ef- fective data-flow invariants across tasks. An LLM-powered agentic coding system first inspects the kernel implementa- tion, compile-time invariant violation reports, and runtime profiles, then retrieves applicable optimizations from a per- sistent knowledge base and drafts concrete data-flow invari- ants together with optimization instructions (Figure 1). The agentic system then implements the corresponding transfor- mations guided by those invariants. Argus treats planner prompts as mutable parameters, uses invariant violations and runtime performance as reward signals, and applies text gradients [66] to update them. We highlight that the combination of these three ideas pro- duces value greater than the sum of their parts. Tag functions alone are declarative specifications; static analysis alone can verify properties but cannot fix violations; and the RL agent alone lacks the dense, structured feedback needed to navigate the vast optimization space. Together, they form a closed loop: tag functions express correctness properties that static analysis verifies at compile time, producing concrete coun- terexamples that serve as dense reward signals for the agent to reason about correctness at an abstract level rather than patching implementations by trial and error. We extensively evaluate Argus on three production ker- nel families: GEMM [4], flash attention [18], and Mixture-of- Experts (MoE) feed-forward networks [55]. Together, these kernels account for over 90% of GPU execution time in LLM inference [40] on the AMD MI300X platform. The kernels generated by Argus achieve 99–104% of the effective mathe- matical throughput (TFLOPS) of state-of-the-art manually optimized assembly libraries [1,4]. Compared with existing agentic baselines, Argus delivers geometric-mean through- put gains of 2–1543×across the evaluated workload families. For flash attention, Argus outperforms the strongest base- line, KernelFalcon [62], by 2.4×. Beyond these kernel families, Argus generalizes to KernelBench [49], completing 100% of Level 1 and 90% of Level 2 tasks. We summarize our key contributions below: •We introduce data-flow invariants into agentic GPU kernel generation, improving both robustness and per- formance of the generated kernels. •We design a tile-based DSL with tag functions and tag assertions, backed by a static analysis that validates data-flow invariants at compile time with zero runtime overhead. 2 •We build an agentic optimization harness that uses ICRL to select optimizations and synthesize invariants for kernels written in the Argus DSL. • We evaluate Argus on GEMM, flash attention, and MoE kernels as well as KernelBench tasks, demon- strating substantial performance gains over existing agentic baselines. 2 Motivating Example We use flash attention [18] on the AMD MI300X GPU as a running example to illustrate why high-performance kernel development demands tightly coupled algorithmic and hardware-level reasoning. Together with general matrix multiplications (GEMM) and Mixture-of-Experts (MoE) lay- ers, attention accounts for up to 90% of GPU execution time in LLM serving [40], making its optimization critical. The standard attention mechanism operates on three input matrices Q,K,V∈ R 푁×푑 , where푁is the sequence length and 푑is the head dimension, and computes O= softmax(QK 푇 )V. A naïve implementation materializes the푁×푁score matrix S= QK 푇 , requiring(푁 2 )memory, prohibitive for long contexts (e.g.,푁=256K). Flash attention avoids this by tiling: it loads a퐵 푟 ×푑block of Q into registers, then streams through K and V in퐵 푐 ×푑blocks staged in shared memory, accumulating into O via online softmax [44] to achieve(푁) memory. The tile sizes퐵 푟 and퐵 푐 are chosen so the working set fits in registers and shared memory. GPU Architecture. Argus targets the AMD MI300X GPU, which implements the CDNA3 architecture [2]. Like other modern GPUs, the MI300X follows a Single Instruction, Mul- tiple Threads (SIMT) design: thousands of die-efficient, in- order cores are organized into streaming multiprocessors (SMs). 1 Each SM contains execution cores, a shared regis- ter pool, and a high-speed scratchpad called shared mem- ory. Threads are organized hierarchically: 64 threads form a warp, executed in lockstep on a 16-lane-wide vector unit; several warps form a threadgroup sharing a register pool and shared memory within one SM. The MI300X includes dedi- cated Matrix Cores for MFMA operations. At the ISA level, it supports 2-wide vectorized fused multiply-add instructions and provides Accumulator Vector General Purpose Registers (AGPRs) that double the register pool. Unlike hardware that automatically manages concurrent execution, both features require explicit compiler support to exploit effectively. Achieving peak performance on a GPU requires exploit- ing this architecture along three axes. First, the kernel must maximize occupancy by utilizing all available threads, regis- ters, shared memory, and dedicated compute units such as the Matrix Cores. Second, memory hierarchy must be man- aged efficiently: e.g., MI300X has an arithmetic-to-bandwidth ratio of∼250, so the kernel must perform at least 250 op- erations per byte transferred from global memory to stay 1 AMD refers to SMs as Compute Units. compute-bound. Kernels vectorize memory accesses, stage data in shared memory for reuse, and implement software pipelining to overlap computation with data transfers. Third, efficient instruction scheduling and register allocation are demanded at the assembly level. On in-order cores, poor scheduling causes pipeline stalls, and exceeding the register budget causes spills [51]. This requires both structuring the kernel to expose sufficient instruction-level parallelism and relying on the compiler to produce a good schedule. Optimizing Flash Attention on MI300X. Turning flash atten- tion into a kernel reaching peak throughput requires a cas- cade of hardware-specific optimizations. Figure 1 shows a simplified Argus DSL implementation with 512 threads, 푑=128, 퐵 푟 =256, and 퐵 푐 =64, embodying these optimizations. The kernel executes two GEMMs S=QK 푇 and O=PV (where P=softmax(QK 푇 )) on Matrix Cores via MFMA in- structions. Each 32×32×8 MFMA consumes operands in a hardware-specific swizzled layout: each thread in a 64-thread warp contributes 4 elements from row-major A∈ R 32×8 and 4 from column-major B∈ R 8×32 , producing a column-major result C. Reaching peak throughput requires staging tiles in shared memory for reuse, transposing V on the fly to match the MFMA operand layout, software pipelining to overlap global memory loads with MFMA computation, partitioning dataflow across disjoint warp groups, and careful instruction scheduling to keep the in-order cores busy without pipeline stalls while staying within the register budget to avoid spills. For example, AITER’s hand-written assembly kernel uses 11 pipeline stages across two warp groups. Section 4 details how the DSL expresses these optimizations. The data shuffling and layout transformations are error- prone: mismatches between expected and actual data layouts cause correctness bugs. Ensuring correctness of these global properties is time-consuming when drafting and debugging optimizations. Argus introduces tag functions and data-flow invariants, expressed asassertstatements in Figure 1, to verify that values are correctly paired at each use site. Sec- tion 4 describes these mechanisms in detail. 3 Overview of Argus Argus is an agentic framework for generating and opti- mizing GPU kernels. It combines three components: (i) a tile-based, Pythonic DSL for expressing data-flow invariants via tag functions and tag assertions (§4), (i) a compile-time analysis that validates these invariants at zero runtime cost (§5), and (i) an agentic harness that iteratively applies opti- mizations guided by invariant feedback (§6). The prototype targets AMD MI300X, but the design generalizes to NVIDIA GPUs and other hardware [70]. Argus DSL. Argus generates and optimizes kernels in a Pythonic, tile-based DSL syntactically similar to CuTe [47] and Triton [59]. Its distinguishing feature is support for tag functions and tag assertions that express data-flow invariants. 3 1 def attn(d:const, gqa:const, Q:Tensor((sq,h q ,d),bf16), 2 K:Tensor((sq,h kv ,d),bf16), V:Tensor((sq,h kv ,d),bf16), . . .): 3 # Assign tags to QKV 4 T Q ,T K = Q[sq,h q ,d]→(sq%32,h q /gqa,d), K[sq,h kv ,d]→(sq%32,h kv ,d) 5 T V = V[sq,h kv ,d]→(sq,h kv ,d%32) 6 # Vectorize and tile memory access for MFMA 7 idx q ,idx h ,tid = blockIdx.x, blockIdx.y, threadIdx.x 8 wid,wtid,gQ = tid/64, tid%64, Q.view((sq/Br,Br,h q ,d/8),u128) 9 gK,sK = K.view((sq/Bc,Bc,h kv ,d/8),u128), make_shared((512,2),u128) 10 gV,sV = V.view((sq/Bc,Bc,h kv ,d/2),u32), make_shared((2,2,4,4,32),u64) 11 rQ,tU = make_local((8,),u128), make_local((2,4,),u32) 12 f sK, f sV = sK.view((2,32,8,2),u128), sV.view((2,2,4,4,16),u128) 13 ∀푖 ∈ [0, 8): rQ[i] = gQ[idx q ,wid*32+wtid%32,idx h ,i*2+wtid/32] 14 for i in range(sq/Bc): 15 for j in range(2): 16 sK[tid,j] = gK[i,j*32+tid/16,idx h /gqa,tid%16] 17 # Vectorize V T via bit rearrangements 18 ∀푘 ∈ [0, 4): tU[j,k] = gV[i,j*32+wid*4+k,idx h /gqa,wtid] 19 nw = (wid&2)*2+wid/4*2+wid%2 20 sV[0,j,wtid/16,nw%4,(wtid%16)*2+nw/4] = concat 3 푘=0 (tU[j,k] lo ) 21sV[1,j,wtid/16,nw%4,(wtid%16)*2+nw/4] = concat 3 푘=0 (tU[j,k] hi ) 22 syncthreads() 23 rS, f rS,rO = make_local((2,),u256),rS.view((4,8),bf16),make_local((4,),u256) 24 T e rS = f rS[x,y]→(i*Bc+ x 2 *32+x%2*16+ y 4 *8+ wtid 32 *4+y%4, idx h gqa ,wtid%32) 25 for j in range(16): 26 tK,tQ = f sK[j/8,wtid%32,j%8,wtid/32], rQ[j%8] 27 f tK, f tQ = tK.view((8,), bf16), tQ.view((8,), bf16) 28 assert tag( f tK[. . .]) == tag( f tQ[. . .]) # f tQ and f tK should match 29 rS[j/8] = matmul( f tK, f tQ,rS[j/8]) # S=QK T 30. . . 31 for j in range(16): # PV, V is transposed due to MFMA 32 tV = f sV[wtid%2,j/8,j%4,j%8/4*2+wtid/32,(wtid/2)%16] 33 f tV = tV.view((8,),bf16) 34 assert tag[ f tV[. . .]] == tag[ f rS[. . .]] # f tV and f rS should match 35 rO[j%4] = matmul( f tV, f rS[j/4],rO[j%4]) Agentic Harness Knowledge Base Initial kernel candidates Kernel descriptionHW specifications Runtime Libraries Output kernels Learnable Planner (via ICRL) Targeted instructions with invariants Lowering Invariant validations Testing & profiling Compile-time & Runtime feedbacks Optimization candidates Optimization Selector Argus DSL . . . r 0 r 8 sq d Global memory wid=0wid=2 lo hi Shared memory Register tid Figure 1. Overview of Argus. Left: Simplified DSL implementation of flash attention (푑=128,퐵 푟 =256,퐵 푐 =64, 512 threads). Top right: Agentic kernel generation workflow. Bottom right: Tag propagation for푉across memory levels; each color–shape combination represents a unique tag, and background shading links code regions to memory access patterns. r 0 and r 8 represent row 0 and row 8. The DSL also exposes hardware instructions and compiler policies (e.g., instruction scheduling) for low-level optimiza- tions, surfacing policy interfaces while hiding mechanisms and intermediate representations to keep context concise for effective LLM reasoning [39] (§4). Invariant Validation. Argus validates data-flow invariants entirely at compile time with zero runtime overhead. The compiler lowers DSL programs to MLIR [36] and runs a flow-sensitive, path-insensitive analysis that propagates tags across assignments and shared memory accesses. The analy- sis targets invariant validation rather than full formal proofs: path-insensitive merging avoids exploring all branch condi- tions, keeping solving tractable. Tag assertions are encoded as SMT constraints and delegated to a solver [19]. Viola- tions are reported with concrete counterexamples, giving the agent actionable feedback (§5). Agentic Optimization Harness. Argus adopts an agentic work- flow to iteratively optimize kernel implementation. The top right of Figure 1 depicts the overall flow. Starting from a description of the desired functionality and target hardware, Argus prepares an initial kernel in its DSL, either manually or via existing agentic generation tools [10,22,23,29,62,69]. The kernel is accompanied by unit tests that serve as correct- ness oracles. The harness comprises five components: (i) a persistent knowledge base documenting generic and GPU- specific optimizations, (i) a learnable, LLM-powered planner 4 that binds optimizations to the current kernel and proposes contextual data-flow invariants, (i) an optimization selec- tor that samples from the planner’s proposals and produces concrete optimization plans, (iv) a lowering agent that im- plements the selected transformations guided by data-flow invariants, and (v) a validator agent that checks invariants at compile time, runs unit tests, and collects runtime profiles. Argus leverages ICRL and use invariant violations and run- time performance as reward signals to enable the planner to learn from past optimization outcomes and to improve its invariant proposals across tasks (§6). 4 The Argus DSL Layouts and Tiles. The Argus DSL models all memory ac- cesses via tiles [21,47,63]. A tile is a memory region paired with a layout function [11] that maps multi-dimensional logical coordinates to one-dimensional physical offsets. A layout function (s,t) is parameterized by two tuples: shapes s= (푠 1 ,· ,푠 푛 )and strides t= (푡 1 ,· ,푡 푛 ). It maps an푛- element coordinate c=(푐 1 ,· ,푐 푛 ) to a linear index. In the simplest case, all elements in s and t are integers, and the layout function computes (푐 1 ,· ,푐 푛 )= Í 푛 1 푐 푖 푡 푖 . A contiguous tensor is a tile whose strides are the cu- mulative product of preceding dimension sizes. For brevity, strides are omitted for contiguous tensors. For example, the Q:Tensor(sq,h q ,d,bf16) in Figure 1 has the layout: Q (푐 1 ,푐 2 ,푐 3 )=푐 1 · h q · d+푐 2 · d+푐 3 The subscript operator in Figure 1 maps coordinates to a lin- ear index and accesses the corresponding memory location. A nested layout contains elements푠 푖 and푡 푖 that are them- selves tuples, mapping coordinates by wrapping around di- mensions. For example, to process the first nested tuple c 1 , 1 can be parameterized by ((푠 1,1 ,푠 1,2 ),(푡 1,1 ,푡 1,2 )), giving 1 (c 1 )= (푐 1,1 %푠 1,1 ) · 푡 1,1 + (푐 1,1 /푠 1,1 ) · 푡 1,2 . Nested layouts primarily represent hardware-swizzled layouts used by ten- sor cores and matrix cores. Layout algebras further define composition, logical division, and logical projection over lay- out functions to represent operations such as data shuffling and dispatching disjoint work to warps and threads; we refer the reader to [11] for details. DSL programs may use theview()function to reinterpret a tile under a different logical layout. For example, line 8 of Figure 1 derivesgQto logically partitionQby퐵 푟 and access data in 16-byte chunks: gQ = Q.view((sq/Br, Br, h q , d/8), u128) Vectorizing loads in 16-byte chunks is a common optimiza- tion to fully utilize GPU memory bandwidth. Argus requires tiles to express all memory accesses. The DSL is memory-safe: only safe layout operations are permit- ted (e.g.,view()requires source and destination layouts to cover identical memory sizes). These constraints simplify program analysis, eliminating full-scale alias analysis on unsafe languages like CUDA. Continuing with the flash attention example in Figure 1, the kernel maps work onto 8 warps. Each warp owns 32 rows of Q and produces the corresponding 32 rows of O, while all warps share the퐵 푐 ×푑tiles of K and V staged in shared memory. Thematmulfunction performs a 32×32×16 fused multiply-accumulate per warp via two 32×32×8 MFMA instructions: each thread holds 8 bf16 elements from A, 8 from B, and 16 accumulator elements from C. Because MFMA expects B in column-major layout, computing O=PVrequires transposing V. The kernel performs this transpose during shared-memory staging: each thread loads a 4×2 sub-tile of V into registerstU(line 18), then scatters the low and high halves via bit-level rearrangements into transposed positions insV(lines 19–20). This approach vectorizes global reads as 4-byte loads, writes tosVas 8-byte stores, and reads back as 16-byte loads swizzled to match the MFMA operand layout, making the result directly consumable by thePVloop (lines 31–35, highlighted in gray). Tag Functions and Assertions. A tag function annotates tensor elements with symbolic tags: tuples of constants, symbolic expressions of coordinates, or the special symbols⊥ and⊤. Tags propagate through assignments. For example,T Q ,T K , andT V in Figure 1 assign tags to el- ements ofQ,K, andVbased on their coordinates (the type bf16 is omitted for brevity): T Q = Q[sq,h q ,d] → (sq%32, h q /gqa, d) T K = K[sq,h kv ,d]→ (sq%32, h kv , d) T V = V[sq,h kv ,d]→ (sq, h kv , d%32) Tag assertions require conformity or non-conformity be- tween tags. The assertion on line 28 in Figure 1 requires matching tags for elements of Q and K: assert tag( f tQ[. . .]) == tag( f tK[. . .]) Computing S=QK 푇 requires each warp to multiply its 32 rows of Q with two 32×128 blocks of K 푇 per iteration. The tag functionsT Q andT K returnsq%32together with the head and dimension indices, so their tags match only when the two operands originate from the same row within a tile, head, and dimension. A passing assertion confirms that the imple- mentation preserves the semantics of S=QK 푇 regardless of optimizations such as shared-memory staging or software pipelining. Similarly, tag functionsT V andT e rS with the as- sertion on line 34 ensure matching elements for O=PV. Figure 1 (bottom right) visualizesV’s tags as data propagate from global memory to shared memory and registers. Compiler and Hardware Primitives. The DSL exposes com- piler and hardware primitives to enable agents to explore low-level optimizations during code generation. The DSL operates at the thread level: kernels manage thread indices and synchronization (e.g.,syncthreads()) directly, enabling topology-aware dispatch optimizations [63]. 5 Argus also exposes eager materializations and scheduling barriers for fine-tuning instruction scheduling in the com- piler backend. It extends the AMD GPU backend’s scheduling barriers to support fine-grained placement of the synchro- nization primitive (i.e.,s_barrier). These interfaces provide LLMs explicit control over hardware features that are critical for reaching peak performance on GPUs [30, 60]. Additionally, Argus exposes vectorized instructions for explicit use in the DSL. For example, the AMD MI300X issues one 2-way vectorized FMA per cycle on packed registers, dou- bling scalar FMA throughput. The AMD GPU backend can emit these instructions, but suboptimal register allocation often prevents it; explicit vectorization in the DSL sidesteps this limitation. Argus further exposes instructions such as buffer_load_dwordx4for branchless memory access with hardware out-of-bounds protection, and allows selecting the accumulator register class (AGPR) directly in buffer load and MFMA instructions, practically doubling available vector registers over off-the-shelf toolchains. 5 Validating Data-Flow Invariants Argus validates data-flow invariants by tracking tags along control and data flows at compile time. The compiler lowers DSL programs to MLIR and performs a flow-sensitive, path- insensitive analysis that propagates tags across assignments and shared memory accesses. At control-flow join points, Argus merges tags using⊥and⊤, with the partial order ⊥< 푡<⊤ for any concrete tag 푡 : merge(푡 1 ,푡 2 )= 푡 1 if 푡 2 ≤ 푡 1 푡 2 if 푡 1 < 푡 2 ⊤ otherwise Constants carry the tag⊥. In the selection patternrQ = 0 if isOOB(...) else Q[. . .], the tag ofrQis the tag ofQ[. . .], sincemerge(⊥,푡)= 푡. Conversely, if two writes to the same shared memory location carry different tags, the location is tagged⊤. Argus supports resetting shared memory tags to ⊥, enabling safe reuse of shared memory segments across pipeline stages while maintaining tag tracking. The minimum tracking granularity is one byte, which suf- fices because popular quantization schemes (e.g.,mxfp4and nvfp4) and hardware-accelerated sub-byte matrix-multiply instructions are byte-aligned. Global and shared memory are modeled as typed memory chunks. Argus propagates tags automatically when accesses use different widths (e.g., vectorizing loads to 16-byte chunks). Argus does not model the full heap and does not track tags for global writes. Argus encodes propagated tags and tag assertions as con- straints in SMT theories and delegates layout-algebra reason- ing to an SMT solver. By design, Argus validates data-flow invariants rather than attempting full formal proofs: path- insensitive merging avoids exploring all branch conditions, keeping the analysis tractable. Tags are purely compile-time constructs: tracking and anal- ysis occur entirely within the compiler and incur no runtime overhead. The current prototype requires thread-block di- mensions to be compile-time constants, allowing Argus to model registers as a fixed-size array indexed by thread ID. Loops are unrolled before analysis. 6 Agentic Optimization Harness The agentic harness in Argus iteratively optimizes a kernel implementation written in the Argus DSL, retaining the best-performing correct candidate across iterations. Persistent Knowledge Base. The harness is grounded by a persistent knowledge base that stores reusable optimization skills spanning (i) global, intrusive restructurings (e.g., tiling, shared-memory staging, and software pipelining) and (i) low-level, hardware-specific tactics (e.g., vectorized mem- ory access patterns and out-of-bounds guards such as AMD buffer load operations). Each entry records the intended code transformation pattern in the DSL and the data-flow invari- ants that must hold after the rewrite. This knowledge base serves as long-term memory across tasks and is curated by domain experts for auditability and stability. Learnable Planner. The planner is an LLM agent binding ab- stract optimization knowledge to the current kernel instance. Given the current DSL implementation, structured compile- time feedback from invariant validation, and runtime profil- ing summaries, it produces a ranked list of proposals of the form⟨optimization, context, score⟩: the optimization skill to apply, an instantiation context describing where and how to apply it, and a confidence/utility score. The context also specializes generic invariant templates into concrete tag func- tions and tag assertions at the appropriate use sites. Optimization Selector. The optimization selector turns the planner’s ranked proposal list into a concrete optimization trajectory. Rather than deterministically choosing the top- ranked option, it samples from the proposal distribution to maintain exploration in a tightly coupled optimization space where locally plausible steps may interact. The output is an executable plan that resolves decisions left underspecified by the planner while preserving the invariants and intent carried in the proposal context. Lowering Agent. Given a concrete plan, the lowering agent implements the corresponding transformations directly in the Argus DSL. It applies rewrites at the level of tiles, layouts, and thread-level control, and inserts or updates tag functions and tag assertions that encode the contextual invariants required by the chosen optimization. The agent retrieves relevant kernel implementations from the knowledge base as few-shot examples to ground the LLM in Argus DSL syntax and optimization patterns. Validator Agent. The validator provides the feedback sig- nals that make iterative optimization reliable. It first invokes 6 compile-time invariant validation, ensuring that newly in- serted or modified tag assertions are discharged by the com- piler analysis and SMT solving. Failures are surfaced as ac- tionable invariant violations. The validator then runs the accompanying unit tests to confirm functional correctness. For candidates that pass both static invariant checks and unit tests, the validator collects runtime profiling summaries and performance metrics, which are fed back to the planner in the next iteration. Algorithm 1: ICRL algorithm to learn the planner. Input: Initial planner 휃 0 , LLM 휋 휃 for 푘= 0, 1, . . . do Sample initial kernel instance 푠 0 ∼ ; Initialize trajectory buffer 휏 ←∅; for 푡= 0, 1, . . .,푇 − 1 do 푃 푡 ← 휋 휃 푘 (푠 푡 ); 푎 푡 ← Select(푃 푡 ); 푠 푡+1 ← Lower(푠 푡 ,푎 푡 ); 푟 푘,푡 ← (푠 푡+1 ,휏 푘 ); 퐵 푘 ← 퐵 푘 ∪(푠 푡 ,푎 푡 ,푟 푘,푡 ); end 퐸 푘 ← PolicyEval(퐵 푘 ); 푔 푘 ← Analyze(퐸 푘 ); 휃 푘+1 ← ParameterUpdate(휃 푘 ,푔 푘 ); end In-Context Reinforcement Learning Loop. Argus improves the planner across tasks via ICRL [23,45]. Following prior ICRL formulations (Algorithm 1), Argus treats the plan- ner prompt (parameterized by휃) as a mutable policy휋 휃 and updates it based on optimization trajectories. Two de- sign choices are notable. First, Argus keeps the knowledge base fixed and learns휃to more effectively bind its entries to each concrete kernel instance, rather than learning new skills. Second, the rewardcombines runtime performance with process rewards derived from invariant validation: viola- tions provide dense, structured negative signals that identify globally inconsistent transformations. Practically, Argus uses the LLM to evaluate optimization validity and perfor- mance (PolicyEval), compute approximate text-level gra- dients (Analyze), and update휃via text-gradient [66]. Over time, these updates steer the planner towards proposals that both (i) select effective optimizations, and (i) synthesize in- variants sufficiently strong to guide the intrusive rewrites (e.g., pipelining and layout reorganization) without degener- ating into trial-and-error debugging. 7 Implementation The Argus prototype comprises approximately 22,000 lines of C++ and 7,000 lines of Python, spanning the DSL compiler, program analysis, agentic harness, and JIT infrastructure. The Python front-end parses DSL programs into an AST and passes it to the C++ compiler, which lowers the AST to an ex- tended MLIR dialect encoding layout algebras. The dialect is then lowered to LLVM IR and compiled to GPU executables. We extend the AMDGPU backend in LLVM 20 to support explicit scheduling of synchronization barriers, selection of register classes for buffer load and MFMA instructions, and eager materialization via phantom uses expressed as inline assembly. Tag propagation and invariant checking are im- plemented as standard MLIR passes. The pass traverses the IR and emits SMT constraints encoding data-flow invari- ants; Argus delegates the constraint optimization problem solving to Z3 [19]. The JIT layer follows the same model as Triton, which supports generics to propagate compile-time constants into GPU kernels, caches compiled binaries, and exposes Python bindings for invocation. 8 Discussion Limitations of data-flow invariants. Data-flow invariants do not address algorithmic optimizations: replacing standard attention with online softmax [18] changes the algorithm itself, not how data flows through it. Tile sizes and warp spe- cialization are search-space decisions handled by the planner and knowledge base. Data-flow invariants verify that ele- ments are correctly paired, routed, and synchronized after tiling, staging, and layout transformations, guarding the cor- rectness of each concrete implementation. Directly generating CUDA or assembly. Rapid improvements in LLM coding abilities [48] suggest training LLMs to directly generate high-performance CUDA or assembly via reinforce- ment learning [20]. Two factors work against this approach. Production GPU kernels demand correctness across tiling, shared-memory staging, and software pipelining; Argus’s deterministic analysis catches data-flow violations that prob- abilistic generation cannot guarantee. Moreover, a single optimized flash attention kernel in Argus’s DSL spans hun- dreds of lines; the equivalent CUDA or assembly would be an order of magnitude larger, exhausting the context window and degrading the generation quality [39]. Investing in better compilers. Higher-level DSLs such as Triton automate layouts and thread scheduling, saving a handful of optimizations. Yet production kernels still require low-level knobs such as inlined assembly, explicit register-class selec- tion, and scheduling barriers for peak performance [1,60]. Compiler stacks are complex artifacts that lag behind hard- ware releases. Vendors such as AMD ship hand-optimized assembly libraries to bridge the gap for critical workloads like GEMM and attention. Argus exposes these low-level primitives as codegen policy: LLMs choose policies while the compiler handles mechanisms, offering a more adaptable path as frontier models improve. Supporting new hardware. Porting Argus to a new GPU re- quires (i) integrating with the target toolchain and exposing 7 its low-level intrinsics, and (i) extending the knowledge base with hardware-specific optimizations. The ICRL loop learns to bind existing knowledge entries to new kernels, so gen- eralizability is preserved without retraining. The trade-off is coverage: new optimization techniques require a domain expert to author knowledge base entries before Argus can ex- ploit them. Letting the agent propose entries autonomously could close this gap but risks propagating under-tested op- timizations, undermining the stability and auditability that expert-curated entries provide. Invariants vs. formal verification. Full formal verification could prove data-flow properties with higher precision but requires significant proof engineering [7,28,33], and it is unclear how to generalize proofs across diverse optimiza- tions tightly coupled with implementation details. Argus’s lightweight program analysis trades precision for a practical benefit: it steers LLMs to reason in terms of data-flow in- variants during generation, providing denser feedback than end-to-end test failures and guiding LLMs toward globally correct plans rather than local trial-and-error patches. Language design for LLM-generated code. Argus remains syn- tactically close to existing DSLs so that LLMs can generalize from their training data via in-context learning [45], enabling adoption of new frontier models without retraining. The DSL need not be human-readable, however; a more radical de- sign could be more token-efficient and enable more powerful generations. 9 Evaluation Our evaluation answers the following research questions: Q1.How effectively does Argus generate and optimize performant kernel implementations? (§9.1) Q2.How well does the Argus DSL capture application- and compiler-level optimizations? (§9.2) Q3. How does Argus generalize to a diverse set of kernel implementation tasks? (§9.3) Q4.How effectively do data-flow invariants guide trans- formations and reduce cost? (§9.4) Experiment Environment. We evaluate Argus on a server equipped with two 64-core AMD EPYC 9554 CPUs, 2 TB DDR5 memory, and eight AMD Instinct MI300X GPUs, each with eight XCDs and 192 GB HBM3 memory. The machine runs Ubuntu 22.04.5 LTS, Linux 5.15.0, and ROCm 7.1.1 [5]. Our compiler augmentations build on ROCm’s LLVM 20.0 fork. For timing, we warm up for at least one second, then re- port the average over the larger of 100 runs or a five-second window. All kernels run inside CUDA graphs, and all pro- cesses are pinned to the CPU socket attached to the target GPU. 9.1 Kernel Performances Setup. To answerQ1, we quantitatively evaluate three kernel families that dominate LLM inference GPU time: matrix mul- tiplication (GEMM), flash attention, and Mixture-of-Experts (MoE) dispatch and combine fused with the SwiGLU acti- vations. For GEMM, we multiply two square bf16 matrices of varying sizes with random data drawn from a normal distribution. For flash attention, we measure bf16 general- ized multi-query attention (GQA) with causal masking at varying sequence lengths, using batch size 16, 8 query heads, 1 key-value head, and head dimension 128, matching the prefill stage of LLaMA-3 70B with tensor parallelism of 8. For MoE, we measure an fp8 feed-forward network with fp8 block-quantized weights following the DeepSeek-V3.2 deployment configuration [20] (dim=7168, inter_dim=2048, 32 experts per GPU under expert parallelism of 8). Baselines. We compare Argus against two categories of strong baselines. The first category is state-of-the-art man- ually optimized libraries, including HipBlasLt [4], HipKit- tens [30], AITER [1], and Triton [59]. The second is ker- nels generated by open-source agentic GPU kernel optimiza- tion frameworks, including KernelBench [49], KSearch [10], CUDAForge [69], and KernelFalcon [62]. For hand-optimized libraries, we use HipBlasLt from ROCm 7.1.1, AITER 0.1.10, and the ROCm fork of Triton (3.5.1+rocm7.1.1.gita272dfa8). We adopt well-optimized Triton implementations of GEMM, GQA, and MoE from HipKittens and AITER, and skip HipKit- tens’ GQA as it is unavailable on our AMD MI300X platform. We enable auto-tuning for all libraries and report the best per- formance for each configuration. For agentic frameworks, all systems including Argus are powered by GLM-5 [24] with prompt templates adjusted to target AMD MI300X GPUs. For each framework that performs iterative optimization, we run 10 iterations and report the best result. For KernelBench, whose pipeline does not iterate, we sample 10 times and re- port the best. Table 2 reports performance in effective FLOPS for all three kernel families. Speedups are aggregated as geo- metric means over the five configurations of each workload family. GEMM. Table 1 breaks down the optimizations applied to GEMM across all baselines. We categorize optimizations by the scope of code changes they require, since larger, more invasive changes are harder for LLMs to produce correctly. The first category comprises global intrusive changes that restructure the kernel, including software pipelining, dis- patching multiplications along the퐾dimension across mul- tiple threadgroups (Split퐾), adopting swizzled layouts to use hardware-accelerated MFMA instructions (MFMA Matmul), staggering workloads along the퐾dimension to mitigate memory controller hot spots [3] (Stagger퐾), and asynchro- nous data loading (Async memcpy). 8 ArgusHipBlasLtHipKittensAITERTritonKernelFalconKSearchKernelBenchCUDAForge Global intrusive changes Software pipelining✓--- Split K---✓----- MFMA Matmul✓-- Stagger K✓---- Async memcpy✓----- Local source changes Bank conflict mitigation ✓---✕ Vectorized loads✓--- Loop unrolling✓-- Workgroup swizzling✓--- ISA-specific optimizations HW OOB-guarded loads ✓----- Use AGPRs✓----- Instruction scheduling✓---- Table 1. GEMM optimizations across Argus, hand-optimized libraries and multiple kernel generation frameworks.✓, -, and✕denote implemented, absent, and incorrectly implemented optimizations. The second category consists of local source changes ex- pressible as small, localized patches in high-level languages: mitigating shared memory bank conflicts (bank conflict mit- igation), vectorizing memory accesses (vectorized loads), ag- gressively unrolling loops (loop unrolling), and rearranging workgroups for better L2 and chiplet locality (workgroup swizzling). The compiler can perform the heavy lifting for several of these (e.g., loop unrolling), so the source-level change is often localized and minimal. The third category covers ISA-specific optimizations im- plemented via inline assembly or compiler intrinsics: us- ing AMD GPU buffer load instructions that guard against out-of-bounds accesses in hardware rather than explicit branches (HW OOB-guarded loads), using MI300X accu- mulator GPR registers to reduce register pressure (AGPRs), and overlapping MFMA, ALU, and memory instructions to hide latency (instruction scheduling). HW OOB-guarded loads and instruction scheduling can be partially expressed through AMD GPU-specific compiler intrinsics, while ex- plicit AGPR use requires either manually written assembly or whole-kernel register allocation management in C++, as done by HipKittens. As shown in Table 2, Argus and the hand-optimized li- braries (HipBlasLt, HipKittens, AITER) apply the most per- formant optimizations across all three categories, while the agentic frameworks are not competitive. Argus consistently outperforms every agentic baseline across all configurations, with speedups ranging from roughly 2×over KernelFal- con on GEMM to 1543×over CUDAForge on fused MoE; KernelFalcon is the strongest agentic baseline throughout. Among the three CUDA-emitting frameworks (KernelBench, CUDAForge, and KSearch), only KSearch uses MFMA in- structions, because it is the only one that explicitly requests them in its prompts. KernelFalcon generates Triton, which allows the Triton compiler to apply pipelining, vectoriza- tion, and loop unrolling internally, but lacks both high-level optimizations like Stagger K and low-level instruction sched- uling, reaching about 2×lower throughput than Argus. Sep- arately, Argus’s GEMM kernel is about 1.8×faster than the hand-optimized Triton kernel listed in Table 2. The Triton kernel includes all global intrusive and local source changes but cannot express Async memcpy or the ISA-specific op- timizations in Table 1, as they fall below Triton’s language abstractions. These account for∼50% of the gap between the Triton kernel and Argus. AITER outperforms HipBlasLt at MNK=1024 but falls behind at larger sizes because the release we benchmarked implements only SplitK partitioning for bf16 GEMM, which helps at small sizes, whereas grid-based partitioning is more efficient at larger ones. Flash Attention. Argus delivers 3.5% higher throughput than AITER (Table 2), a hand-written assembly library from AMD that targets this exact GQA configuration. KernelBench, K- Search, and CUDAForge are 59–120×slower than Argus in throughput because none of them use hardware matrix cores for theQK 푇 andPVmultiplications; without matrix core us- age, other optimizations yield marginal gains. KernelFalcon generates a Triton kernel that is 2.4×slower than Argus in throughput across sequence lengths. Its kernel lacks the application-level optimization that shuffles access patterns of matrix V to improve global memory load efficiency and 9 ArgusHipBlasLtHipKittensAITERTritonKernelFalconKSearchKernelBenchCUDAForge BF16 Square GEMM, Matrix Size (M×N×K) 1024146.91144.74133.47208.4366.3759.4344.5114.0913.71 2048393.35380.19359.16247.22272.36279.1063.5015.0014.50 4096612.57605.00605.95329.07335.57286.8769.3313.0112.67 8192650.54629.45612.00304.96341.35284.3968.6612.8012.59 16384573.70626.55567.66226.40331.66270.6660.5711.5011.39 Speedup1.001.000.950.620.550.490.140.030.03 GQA Forward Flash Attention, Sequence length 1024403.36--344.23352.78208.428.274.589.67 2048524.55--528.39460.07244.386.684.959.68 4096614.41--643.85565.15249.596.835.0510.00 8192720.56--687.28625.54247.626.815.1010.13 16384722.34--706.86679.22246.324.684.589.94 Speedup1.00--0.970.900.410.010.010.02 Fused MoE, Sequence length 1024437.84--416.47225.388.784.940.810.30 2048484.23--465.58266.098.645.170.830.32 4096515.56--528.22298.628.285.130.810.33 8192529.67--547.88328.548.135.150.820.33 16384526.45--567.65341.848.045.080.820.33 Speedup1.00--1.010.580.020.010.0020.001 Table 2. Effective throughputs (TFLOPS) and speedups relative to Argus for GEMM, flash attention, and fused MoE kernels. Speedups are geometric means over the five configurations of each workload family. stages V for on-the-fly transpose. The Triton compiler maps computations onto matrix cores but does not apply software pipelining or instruction interleaving. The warp specializa- tion and fine-grained scheduling within pipelining stages fall below Triton’s abstraction level and cannot be expressed in the generated kernel. Fused MoE. Both Argus and AITER partition work across tokens and experts, use matrix cores for the projections, and quantize intermediate results of the up projection to fp8 to leverage fp8 matrix cores, which provide 16×the through- put of scalar units on MI300X. Fused MoE is the most chal- lenging kernel for agentic frameworks, which are substan- tially slower than Argus. Measured in effective throughput, the slowdown ranges from roughly 59×for KernelFalcon to 1543×for CUDAForge. The specification is complex: inputs and weights are block-quantized to fp8, expert assignments use sorted maps with indirection through token IDs, and tokens must be grouped for coalesced access. These map- pings are prerequisites for correctness, and optimizing the kernel requires preserving them across transformations. No agentic framework uses MFMA instructions, as the under- lying LLMs cannot maintain the required mappings while generating correct matrix core code. Consequently, all omit fp8 requantization of intermediate results consumed by the down projection, preventing it from using fp8 matrix cores and forgoing their 16×throughput advantage. Most frame- works (excluding KernelBench) also lack expert-dimension parallelism. All frameworks also miss that the fp8 tensors are 16-byte aligned, failing to exploit vectorized loads. 9.2 Effects on Different Optimization Levels To answerQ2, we selectively disable optimizations of the flash attention kernel to quantify the contribution of each optimization level. We apply six optimizations sequentially: (i) the naïve implementation (Naive), (i) on-demand trans- posing V (TransV), (i) asynchronous loads for Q and K (Async memcpy), (iv) mitigating shared memory bank con- flicts (Bank conflict), (v) software pipelining with warp spe- cialization (Pipeline+WS), and (vi) fine-grained instruction scheduling (All). 10 Figure 2 shows the performance ablation results. We ob- serve that the most significant single contributor is bank conflict mitigation, improving throughput by roughly 30%. Adding pipelining and warp specialization alone causes a slight regression due to increased instruction and branch counts; instruction scheduling recovers and extends the gains, yielding a 2.4–2.8×overall speedup over the naïve baseline. Applying instruction scheduling directly on top of bank conflict mitigation (without pipelining) yields mar- ginal improvement, as the scheduler has limited opportunity to reorder an unpipelined instruction stream. Bank conflict mitigation thus represents a local optimum: it delivers a large speedup on its own, but further progress requires the non-trivial step of adding pipelining and warp specialization, an optimization that initially degrades performance before instruction scheduling can exploit it. 0 50 100 150 200 250 300 350 400 102420484096819216384 Naive Transpose V +Async memcpy +Bank conflict +Pipeline+WS +Inst. schedule (All) Performance (TFLOPS) Sequence length Figure 2. Ablation of optimizations on flash attention kernel. 9.3 Generality of Argus The previous sections focus on three kernel families. To test whether the approach generalizes (Q3), we apply Argus to 200 Level-1 and Level-2 problems from KernelBench. As shown in Table 3, Argus produces correct kernels for all 100 Level-1 problems and 90 of 100 Level-2 problems, achieving geometric-mean speedups of 0.74×and 0.88×over the Ker- nelBench reference implementations, respectively. Of these, 49 Level-1 and 42 Level-2 kernels outperform the PyTorch reference, benefiting from memory-layout transforms, oper- ator fusion, vectorized loads, shared-memory staging, and improved instruction scheduling. For 4 Level-2 problems Ar- gus falls back to calling the PyTorch operator directly; we count these as failures. The Argus DSL is absent from LLM training corpora. To evaluate how effectively LLMs can generate Argus code without fine-tuning, we use GLM-5 to generate initial kernels for the same 200 problems under two prompt strategies: (i) three-shot in-context examples and (i) compact syntax summaries only. Table 3 shows that in-context examples have a marginal effect on correctness: Pass@1 improves by 1 and 17 percentage points for Level-1 and Level-2, respectively; Pass@3 improves by 5 percentage points for Level-1 and remains comparable for Level-2. We attribute this to the syntactic similarity between the Argus DSL and existing DSLs. In-context examples do reduce the average number of files read by 13% (Level 1) and 8% (Level 2), suggesting that the LLM makes fewer exploratory tool calls when provided with examples. 9.4 Effectiveness of Data-Flow Invariants To isolate the contribution of data-flow invariants (i.e., an- sweringQ4), we modify the planner to remove all data-flow invariants and disable compiler feedback in the agentic har- ness. From KernelBench, we select 25 Level-1 and 35 Level-2 GEMM- and convolution-related problems, excluding trivial problems like ReLU, to form a representative set of 60 prob- lems. We apply three optimizations from Table 1: MFMA Matmul, Software pipelining, and HW OOB-guarded loads. Table 3 shows the effect of data-flow invariants on guiding optimizations. Invariants raise Pass@1 from 60% to 75% on Level-1 and from 40% to 57% on Level-2, while also reducing average token usage from 18.1M to 15.0M on Level-1 and from 14.2M to 13.5M on Level-2. These results suggest that data-flow invariants make intrusive optimization both more effective and more efficient. 10 Related work Synthesizing GPU Kernels with LLMs. Frontier coding mod- els [24,48] combined with agentic workflows [6,10,14,17,22, 23,27,29,49,62,68,69] have made generating GPU kernels from natural language specifications practical. These work- flows differ along three axes: reward signals, optimization knowledge, and execution quality. Most use correctness and relative performance gains as rewards, which are susceptible to reward hacking [35]. Concurrent work휇CUTLASS [27] substitutes the relative gains with Speed-of-Light guidance, penalizing claims beyond hardware capability to reduce re- ward hacking. Optimization knowledge can come from fine- tuning and reinforcement learning [57], world models [10], or evolutionary search [38]. Reliable execution of complex optimizations is essential for meaningful performance gains on production kernels. Argus incorporates data-flow invariants as dense process rewards that guide both generation and subsequent optimiza- tion. Exposing low-level hardware and compiler intrinsics further enables fine-grained optimizations such as instruc- tion scheduling. To our knowledge, Argus is the first agentic framework capable of implementing all optimizations re- quired to match expert-tuned production GPU kernels; our evaluation shows that such kernels are orders of magnitude faster than those produced by generation-based approaches. ICL and ICRL. In-context learning (ICL) enables LLMs to adapt to new downstream tasks from examples provided in 11 PerformanceAdding in-context examplesAdding Invariants Valid%GeoMeanMinMax> 1×Pass@1%Pass@3%Avg. files readPass@1%Avg. token usage Level 1100%0.740.0111.494982→8395→10010.28→8.9560→7518.1M→15.0M Level 290%0.880.017.844239→5686→8519.81→18.1940→5714.2M→13.5M Table 3. Validity and performance of Argus-generated kernels over 200 KernelBench problems, and the effectiveness of in-context examples and invariants on 60 representative problems. A kernel is valid if it passes functionality tests and human verification. The right arrow shows the change after adding in-context examples or invariants. the prompt without model parameter updates [45]. Further- more, in-context reinforcement learning (ICRL) extends ICL by treating the prompt as a dynamic learnable policy and up- dating it based on reward signals. For example, TextGrad [66] provides a general framework for computing text-based gra- dients over LLM outputs. In terms of kernel code generation, KernelBlaster [23] applies ICRL to CUDA kernel optimiza- tion, maintaining a memory of past optimization trajectories and using them to improve cross-task generalization. Argus builds on these techniques: we use ICRL to learn planner prompts that bind optimization knowledge to concrete ker- nels, with data-flow invariant violations serving as process rewards alongside runtime performance. DSL Designs for GPU Kernels. A central theme in GPU DSL design is determining the appropriate level of abstraction to improve productivity. Triton describes GPU workloads at the block level, eliminating explicit management of memory lay- outs and synchronization. Halide [52] and Exo [32] separate algorithm specification from scheduling policy for tensor workloads. Tile-based DSLs [21,27,47,63] further structure programs around tiles as primitives for data movement and matrix multiplication, offloading implementation complex- ity to compilers. The raised abstractions, however, incur an abstraction tax: these DSLs occasionally resort to leaky low- level abstractions [60] to recover performance by providing direct control over memory layouts, warp specialization, reg- ister allocation, and instruction scheduling. The Argus DSL is a tile-based DSL that targets LLMs rather than human developers. It is memory safe to facilitate static analysis, and it incorporates data-flow invariants with compile-time checking to provide dense structural feedback for kernel gen- eration. It stays syntactically close to Triton so that LLMs can leverage existing training knowledge. Optimizing Compilers for Tensor Programs. Tensor compil- ers [15,26,37,42,54,64,71,72] lower DSLs and tensor programs to GPU binaries. Their optimizations target high tensor core utilization, efficient memory layout, software pipelining, operator fusion, and overlapping computation with memory access. They rely on heuristics and exhaus- tive search to find efficient schedules. Argus complements these compilers: we expose the hardware SIMT program- ming model along with low-level intrinsics in the DSL, and delegate schedule exploration to LLMs. Invariants and Information Flow Systems. Invariants provide a concise way to specify and reason about end-to-end prop- erties across multi-layered systems [43,58,65]. Information flow analysis enforces such invariants at the system level. HiStar [67] assigns security labels to data, defines a partial order on labels, tracks their propagation, and permits infor- mation flow only along that order. JFlow [46] augments Java with information labels and enforces information flows at the language level. These prior systems focus on enforcing security policies; Argus adapts the same tracking and en- forcement mechanisms to GPU kernel optimization, using tag propagation to verify correct operand pairings and giving LLMs a global view of correctness. Formal verification [12,13,33,56] can detect invariant violations precisely, but multiplies development effort by 2– 6×[33] for sequential programs; verifying concurrent GPU kernels demands additional modeling effort [7,28]. Argus in- stead adopts lightweight static analysis to generate warnings, enabling the agentic loop to iterate automatically. The cur- rent prototype uses Z3 [19] to reason about layout algebra; a refined version could leverage polyhedral analysis [9,25,50] for faster validation. 11 Conclusion We introduced Argus, which shows that data-flow invari- ants can enable robust optimizations that bridge the per- formance gap between LLM-generated and hand-optimized GPU kernels. By reporting compile-time invariant violations with concrete counterexamples, the system provides dense structural feedback that guides LLM coding agents through intrusive transformations such as software pipelining and layout reorganization. The DSL complements this feedback by exposing optimization policy while hiding compiler mech- anisms, keeping the language learnable for LLMs yet capable of reaching assembly-level performance. A curated knowl- edge base grounds these optimizations for the ICRL planner to learn to bind them to concrete kernels, combining au- ditability with adaptability. Although we evaluate on the AMD MI300X, the abstractions are not AMD-specific: the DSL, invariant checking, and ICRL planner generalize to other GPU architectures and, more broadly, to any domain where correctness depends on global data-flow properties across complex transformations. 12 References [1]AMD. Aiter: AI Tensor Engine for ROCm. https://github.com/ROCm/ aiter, 2025. [2] AMD. AMD instinct MI300 instruction set architecture. https:// w.amd.com/content/dam/amd/en/documents/instinct-tech- docs/instruction-set-architectures/amd-instinct-mi300-cdna3- instruction-set-architecture.pdf, August 2025. [3] AMD. AMD instinct MI300X workload optimization. https://rocm.d ocs.amd.com/en/latest/how-to/rocm-for-ai/inference-optimization /workload.html, 2025. [4]AMD. hipBLASLt: General matrix-matrix operations library for amd gpus. https://rocm.docs.amd.com/projects/hipBLASLt/en/latest/index .html, 2026. [5] AMD. Rocm software. https://github.com/ROCm/ROCm, 2026. [6]Carlo Baronio, Pietro Marsella, Ben Pan, Simon Guo, and Silas Alberti. Kevin: Multi-turn RL for generating CUDA kernels. In The Fourteenth International Conference on Learning Representations, 2026. [7]Adam Betts, Nathan Chong, Alastair Donaldson, Shaz Qadeer, and Paul Thomson. GPUVerify: a verifier for GPU kernels. OOPSLA ’12, 2012. [8] Simon Boehm. How to optimize a CUDA matmul kernel for cuBLAS- like performance: a worklog. https://siboehm.com/articles/22/CUDA- M, 2022. [9]Uday Bondhugula, Albert Hartono, J. Ramanujam, and P. Sadayappan. A practical automatic polyhedral parallelizer and locality optimizer. PLDI ’08, 2008. [10]Shiyi Cao, Ziming Mao, Joseph E. Gonzalez, and Ion Stoica. K-Search: LLM kernel generation via co-evolving intrinsic world model, 2026. [11]Jack Carlisle, Jay Shah, Reuben Stern, and Paul VanKoughnett. Cate- gorical foundations for CuTe layouts. 2026. [12]Bodhisatwa Chatterjee, Drew Zagieboylo, Sana Damani, Siva Hari, and Christos Kozyrakis. Proofwright: Towards agentic formal verification of CUDA. CoRR, abs/2511.12294, 2025. [13]Haogang Chen, Daniel Ziegler, Tej Chajed, Adam Chlipala, M. Frans Kaashoek, and Nickolai Zeldovich. Using crash hoare logic for certify- ing the FSCQ file system. In USENIX ATC 16, 2016. [14]Terry Chen, Zhifan Ye, Bing Xu, Zihao Ye, Timmy Liu, Ali Hassani, Tianqi Chen, Andrew Kerr, Haicheng Wu, Yang Xu, Yu-Jung Chen, Hanfeng Chen, Aditya Kane, Ronny Krashinsky, Ming-Yu Liu, Vinod Grover, Luis Ceze, Roger Bringmann, John Tran, Wei Liu, Fung Xie, Michael Lightstone, and Humphrey Shi. AVO: Agentic variation oper- ators for autonomous evolutionary search, 2026. [15]Tianqi Chen, Thierry Moreau, Ziheng Jiang, Lianmin Zheng, Eddie Yan, Haichen Shen, Meghan Cowan, Leyuan Wang, Yuwei Hu, Luis Ceze, Carlos Guestrin, and Arvind Krishnamurthy. TVM: An automated end-to-end optimizing compiler for deep learning. OSDI’18, 2018. [16] Patrick Cousot and Radhia Cousot. Abstract interpretation: a unified lattice model for static analysis of programs by construction or approx- imation of fixpoints. In Proceedings of the 4th ACM SIGACT-SIGPLAN Symposium on Principles of Programming Languages, POPL ’77, 1977. [17]Weinan Dai, Hanlin Wu, Qiying Yu, Huan ang Gao, Jiahao Li, Chengquan Jiang, Weiqiang Lou, Yufan Song, Hongli Yu, Jiaze Chen, Wei-Ying Ma, Ya-Qin Zhang, Jingjing Liu, Mingxuan Wang, Xin Liu, and Hao Zhou. CUDA agent: Large-scale agentic RL for high- performance CUDA kernel generation, 2026. [18]Tri Dao. Flashattention-2: Faster attention with better parallelism and work partitioning. arXiv preprint arXiv:2307.08691, 2023. [19]Leonardo de Moura and Nikolaj Bjørner. Z3: An efficient SMT solver. In TACAS, 2008. [20]DeepSeek-AI. Deepseek-v3.2: Pushing the frontier of open large lan- guage models, 2025. [21]Yaoyao Ding, Bohan Hou, Xiao Zhang, Allan Lin, Tianqi Chen, Cody Hao Yu, Yida Wang, and Gennady Pekhimenko. Tilus: A tile- level GPGPU programming language for low-precision computation. ASPLOS ’26, 2025. [22]Juncheng Dong, Yang Yang, Tao Liu, Yang Wang, Feng Qi, Vahid Tarokh, Kaushik Rangadurai, and Shuang Yang. STARK: Strategic team of agents for refining kernels. ICLR’26, 2026. [23]Shengjun Kris Dong, Sahil Modi, Dima Nikiforov, Sana Damani, Ed- ward Lin, Siva Kumar Sastry Hari, and Christos Kozyrakis. Ker- nelBlaster: Continual cross-task CUDA optimization via memory- augmented in-context reinforcement learning. CoRR, abs/2602.14293, 2026. [24]GLM-5-Team. GLM-5: from vibe coding to agentic engineering, 2026. [25]Tobias Grosser, Armin Groesslinger, and Christian Lengauer. Polly - performing polyhedral optimizations on a low-level intermediate representation. Parallel Processing Letters, 2012. [26]Yue Guan, Xinwei Qiang, Zaifeng Pan, Daniels Johnson, Yuanwei Fang, Keren Zhou, Yuke Wang, Wanlu Li, Yufei Ding, and Adnan Aziz. Mercury: Unlocking multi-GPU operator optimization for LLMs via remote memory scheduling. SOSP ’25, 2025. [27]Siva Kumar Sastry Hari, Vignesh Balaji, Sana Damani, Qijing Huang, and Christos Kozyrakis. Improving efficiency of GPU kernel opti- mization agents using a domain-specific language and speed-of-light guidance, 2026. [28]Chris Hawblitzel, Jon Howell, Manos Kapritsos, Jacob R. Lorch, Bryan Parno, Michael L. Roberts, Srinath Setty, and Brian Zill. IronFleet: proving practical distributed systems correct. SOSP ’15, 2015. [29] Bodun Hu, Yoga Sri Varshan V, Saurabh Agarwal, and Aditya Akella. CUCo: An agentic framework for compute and communication co- design, 2026. [30]William Hu, Drew Wadsworth, Sean Siddens, Stanley Winata, Daniel Y. Fu, Ryann Swann, Muhammad Osama, Christopher Ré, and Simran Arora. HipKittens: Fast and furious AMD kernels, 2025. [31] Rodrigo Huerta, Mojtaba Abaie Shoushtary, José-Lorenzo Cruz, and Antonio Gonzalez. Dissecting and modeling the architecture of modern GPU cores. MICRO ’25, 2025. [32]Yuka Ikarashi, Kevin Qian, Samir Droubi, Alex Reinking, Gilbert Louis Bernstein, and Jonathan Ragan-Kelley. Exo 2: Growing a scheduling language. ASPLOS ’25, 2025. [33]Gerwin Klein, Kevin Elphinstone, Gernot Heiser, June Andronick, David Cock, Philip Derrin, Dhammika Elkaduwe, Kai Engelhardt, Rafal Kolanski, Michael Norrish, Thomas Sewell, Harvey Tuch, and Simon Winwood. seL4: formal verification of an OS kernel. SOSP ’09, 2009. [34] M. Lam. Software pipelining: an effective scheduling technique for vliw machines. PLDI ’88, 1988. [35]Robert Tjarko Lange, Qi Sun, Aaditya Prasad, Maxence Faldor, Yujin Tang, and David Ha. Towards robust agentic CUDA kernel bench- marking, verification, and optimization, 2025. [36]Chris Lattner, Mehdi Amini, Uday Bondhugula, Albert Cohen, Andy Davis, Jacques Pienaar, River Riddle, Tatiana Shpeisman, Nicolas Vasi- lache, and Oleksandr Zinenko. MLIR: Scaling compiler infrastructure for domain specific computation. CGO’21, 2021. [37]Yijin Li, Jiacheng Zhao, Sun Qianqi, Haohui Mai, Lei Chen, Wanlu Cao, Yanfan Chen, Li zhicheng, Ying Liu, Xinyuan Zhang, Xiyu Shi, Jie Zhao, Jingling Xue, Huimin Cui, and XiaoBing Feng. SIRIUS: Harvesting whole-program optimization opportunities for DNNs. In Proceedings of Machine Learning and Systems, 2023. [38]Gang Liao, Hongsen Qin, Ying Wang, Alicia Golden, Michael Kuchnik, Yavuz Yetim, Jia Jiunn Ang, Chunli Fu, Yihan He, Samuel Hsia, Zewei Jiang, Dianshi Li, Uladzimir Pashkevich, Varna Puvvada, Feng Shi, Matt Steiner, Ruichao Xiao, Nathan Yan, Xiayu Yu, Zhou Fang, Ro- man Levenstein, Kunming Ho, Haishan Zhu, Alec Hammond, Richard Li, Ajit Mathews, Kaustubh Gondkar, Abdul Zainul-Abedin, Ketan Singh, Hongtao Yu, Wenyuan Chi, Barney Huang, Sean Zhang, Noah Weller, Zach Marine, Wyatt Cook, Carole-Jean Wu, and Gaoxiang Liu. KernelEvolve: Scaling agentic kernel coding for heterogeneous AI accelerators at Meta, 2026. 13 [39]Nelson F. Liu, Kevin Lin, John Hewitt, Ashwin Paranjape, Michele Bevilacqua, Fabio Petroni, and Percy Liang. Lost in the middle: How language models use long contexts. Transactions of the Association for Computational Linguistics, 12:157–173, 2024. [40]Zichang Liu, Jue Wang, Tri Dao, Tianyi Zhou, Binhang Yuan, Zhao Song, Anshumali Shrivastava, Ce Zhang, Yuandong Tian, Christopher Ré, and Beidi Chen. Deja Vu: contextual sparsity for efficient LLMs at inference time. ICML’23, 2023. [41]Weile Luo, Ruibo Fan, Zeyu Li, Dayou Du, Qiang Wang, and Xiaowen Chu. Benchmarking and dissecting the nvidia hopper GPU architec- ture. IPDPS’24, 2024. [42]Lingxiao Ma, Zhiqiang Xie, Zhi Yang, Jilong Xue, Youshan Miao, Wei Cui, Wenxiang Hu, Fan Yang, Lintao Zhang, and Lidong Zhou. Rammer: Enabling holistic deep learning compiler optimizations with rTasks. OSDI’20, 2020. [43]Haohui Mai, Edgar Pek, Hui Xue, Samuel Talmadge King, and Parthasarathy Madhusudan. Verifying security invariants in Expres- sOS. ASPLOS ’13, 2013. [44]Maxim Milakov and Natalia Gimelshein. Online normalizer calculation for softmax. arXiv preprint arXiv:1805.02867, 2018. [45] Giovanni Monea, Antoine Bosselut, Kianté Brantley, and Yoav Artzi. LLMs are in-context bandit reinforcement learners. In Second Confer- ence on Language Modeling, 2025. [46]Andrew C. Myers. JFlow: practical mostly-static information flow control. POPL’19, pages 228–241, January 1999. [47]NVIDIA. CUTLASS: CUDA templates for linear algebra subroutines and solvers. https://github.com/NVIDIA/cutlass, 2026. [48]OpenAI. GPT-5.3-Codex system card. Technical report, OpenAI, February 2026. [49] Anne Ouyang, Simon Guo, Simran Arora, Alex L Zhang, William Hu, Christopher Re, and Azalia Mirhoseini. KernelBench: Can LLMs write efficient GPU kernels? ICML’25, 2025. [50]Benoît Pradelle, Benoît Meister, Muthu Baskaran, Jonathan Springer, and Richard Lethin. Polyhedral optimization of TensorFlow compu- tation graphs. In Abhinav Bhatele, David Boehme, Joshua A. Levine, Allen D. Malony, and Martin Schulz, editors, ProTools, pages 74–89, Cham, 2019. Springer International Publishing. [51]Fernando Magno Quintão Pereira and Jens Palsberg. Register alloca- tion by puzzle solving. PLDI ’08, 2008. [52] Jonathan Ragan-Kelley, Connelly Barnes, Andrew Adams, Sylvain Paris, Frédo Durand, and Saman Amarasinghe. Halide: a language and compiler for optimizing parallelism, locality, and recomputation in image processing pipelines. PLDI ’13, 2013. [53]G. Ramalingam. The undecidability of aliasing. ACM Trans. Program. Lang. Syst., 16(5):1467–1471, September 1994. [54]Amit Sabne. XLA : Compiling machine learning for peak performance, 2020. [55]Noam Shazeer, *Azalia Mirhoseini, *Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. In International Conference on Learning Representations, 2017. [56]Helgi Sigurbjarnarson, James Bornholt, Emina Torlak, and Xi Wang. Push-Button verification of file systems via crash refinement. OSDI’16, 2016. [57]Songqiao Su, Xiaofei Sun, Xiaoya Li, Albert Wang, Jiwei Li, and Chris Shum. CUDA-L2: Surpassing cublas performance for matrix multipli- cation through reinforcement learning, 2025. [58]Shuo Tang, Haohui Mai, and Samuel T. King. Trust and protection in the illinois browser operating system. OSDI’10, 2010. [59]Philippe Tillet, H. T. Kung, and David Cox. Triton: an intermediate language and compiler for tiled neural network computations. In Pro- ceedings of the 3rd ACM SIGPLAN International Workshop on Machine Learning and Programming Languages, 2019. [60]Triton. Gluon. https://github.com/triton-lang/triton/tree/main/lib/Di alect/Gluon, 2026. [61] Gael Varoquaux, Sasha Luccioni, and Meredith Whittaker. Hype, sustainability, and the price of the bigger-is-better paradigm in AI. FAccT ’25, 2025. [62]Laura Wang and PyTorch Team. KernelFalcon: Autonomous gpu kernel generation via deep agents. https://pytorch.org/blog/kernelfalcon- autonomous-gpu-kernel-generation-via-deep-agents, 2025. [63]Lei Wang, Yu Cheng, Yining Shi, Zhiwen Mo, Zhengju Tang, Wenhao Xie, Tong Wu, Lingxiao Ma, Yuqing Xia, Jilong Xue, Fan Yang, and Zhi Yang. Tilelang: Bridge programmability and performance in modern neural kernels. ICLR’26, 2026. [64]Mengdi Wu, Xinhao Cheng, Shengyu Liu, Chunan Shi, Jianan Ji, Man Kit Ao, Praveen Velliengiri, Xupeng Miao, Oded Padon, and Zhihao Jia. Mirage: a multi-level superoptimizer for tensor programs. OSDI ’25, 2025. [65]Alexander Yip, Xi Wang, Nickolai Zeldovich, and M. Frans Kaashoek. Improving application security with data flow assertions. SOSP ’09, 2009. [66] Mert Yuksekgonul, Federico Bianchi, Joseph Boen, Sheng Liu, Zhi Huang, Carlos Guestrin, and James Zou. TextGrad: Automatic "differ- entiation" via text, 2024. [67]Nickolai Zeldovich, Silas Boyd-Wickizer, Eddie Kohler, and David Mazières. Making information flow explicit in HiStar. OSDI’06, 2006. [68]Yi Zhai, Sijia Yang, Keyu Pan, Renwei Zhang, Shuo Liu, Chao Liu, Zichun Ye, Jianmin Ji, Jie Zhao, Yu Zhang, and Yanyong Zhang. En- abling tensor language model to assist in generating High-Performance tensor programs for deep learning. OSDI’24, 2024. [69] Zijian Zhang, Rong Wang, Shiyang Li, Yuebo Luo, Mingyi Hong, and Caiwen Ding. CudaForge: An agent framework with hardware feed- back for CUDA kernel optimization, 2025. [70]Zuoning Zhang, Dhruv Parikh, Youning Zhang, and Viktor Prasanna. Benchmarking the performance of large language models on the cere- bras wafer scale engine, 2024. [71]Jie Zhao, Bojie Li, Wang Nie, Zhen Geng, Renwei Zhang, Xiong Gao, Bin Cheng, Chen Wu, Yun Cheng, Zheng Li, Peng Di, Kun Zhang, and Xuefeng Jin. AKG: automatic kernel generation for neural processing units using polyhedral transformations. PLDI’21, 2021. [72]Lianmin Zheng, Chengfan Jia, Minmin Sun, Zhao Wu, Cody Hao Yu, Ameer Haj-Ali, Yida Wang, Jun Yang, Danyang Zhuo, Koushik Sen, Joseph E. Gonzalez, and Ion Stoica. Ansor: Generating High- Performance tensor programs for deep learning. OSDI’20, 2020. 14