arXiv:2604.18616cs.DCcs.AI2026-04

用数据流不变量指导智能体优化GPU代码,性能接近人工调优。

ARGUS: Agentic GPU Optimization Guided by Data-Flow Invariants

  • 基于数据流不变量的编译时约束检查,实现精准错误定位。
  • 生成代码达顶尖手写汇编99%-104%性能,比现有智能体快2-1543倍。
  • 适用于大模型推理核心算子,可推广至200个基准任务。

基于大语言模型的编码智能体虽能生成功能正确的GPU内核,但在矩阵乘法、注意力和专家混合(MoE)等关键计算上性能远低于手工优化库。峰值性能需协同优化分块、共享内存调度、软件流水和指令调度,而现有智能体仅依赖稀疏的通过/失败反馈,无法诊断全局约束违例。我们提出Argus,一种基于数据流不变量的智能体框架:在编译时定义数据执行过程中的 choreography 规则。Argus采用基于分块的PythonDSL,暴露硬件指令与编译策略,隐藏底层表示;通过标签函数传播符号注解,标签断言在使用点强制关系约束。违反时,编译器返回具体反例(线程、数据元素、程序点),提供密集结构化反馈以精准修复。不变量通过布局代数与SMT求解进行抽象解释验证,零运行时开销。上下文强化学习规划器学习选择优化策略并合成有效不变量,依托精心构建的GPU优化技术知识库。我们在AMD MI300X GPU上评估,涵盖占大模型推理90%以上时间的GEMM、flash attention和MoE内核。生成内核达到状态领先手写汇编99%-104%吞吐率,比现有智能体系统快2-1543倍。同时在200个KernelBench任务中,解决100%的Level 1问题和90%的Level 2问题。

原文摘要 · Abstract (English)

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.

GPU优化智能体不变量验证大模型推理

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。