Skip to content

AKG kernel Agent: A Multi-Agent Framework for Cross-Platform Kernel Synthesis

πŸ•’ Published (v1): 2025-12-29 12:42 UTC Β· Source: Arxiv Β· link

Ask a follow-up

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

πŸ’¬ Ask ChatGPT✦ Ask Claude

TL;DR

AKG kernel Agent is a multi-agent system (Designer, Coder, Verifier, Conductor) that automates GPU/NPU/CPU kernel generation across multiple DSLs (Triton, CUDA-C, TileLang, CPP) via a hardware-agnostic intermediate representation called the Unified Sketch. The system couples document-driven knowledge ingestion with hierarchical RAG and iterative island-model search to achieve correctness and performance without manual tuning. On KernelBench Level 1 with Triton-Ascend it reaches 1.46Γ— geometric mean speedup over PyTorch Eager.

Problem

Manual kernel optimization cannot keep pace with rapidly diversifying hardware (GPUs, NPUs, CPUs) and model architectures. Single-LLM approaches conflate high-level strategy with low-level DSL syntax, producing code with correctness or performance failures, while fine-tuned models suffer from data scarcity. Existing benchmarks also exhibit reward-hacking loopholes and test only static shapes, masking real-world brittleness.

Method

Four specialized agents collaborate in a closed loop:

  • Designer produces a Unified Sketch β€” a minimalist IR S = (D, O, C, H) capturing declarations, operations (alloc/load/store/compute), control flow, and @llm_hint optimization decorators β€” that is DSL-agnostic and reusable across backends.
  • Coder translates the Sketch into a target DSL (Triton, CUDA-C, TileLang, AscendC, CPP) using dynamically loaded API docs and retrieved examples.
  • Verifier checks numerical correctness via element-wise relative error (threshold Ο„ per dtype) and profiles speedup; computes pass@k and fast_p metrics.
  • Conductor performs LLM-based error classification and routes failures: syntax/API/runtime errors β†’ Coder; algorithm/memory-pattern errors β†’ Designer.

Supporting mechanisms: - Document-Driven Integration (DDI): structured DocSpec (concepts, APIs, expert hints, examples) ingested at runtime; LLM-based compression selects relevant API subsets per task. - Hierarchical RAG: (1) embedding cosine search on computation logic β†’ (2) hard filter on DSL/backend/operator type β†’ (3) shape-compatible semantic matching. - Iterative island-model search: K=2 semi-independent islands, P=4 parallel candidates/round, R=3 rounds; stratified contrastive sampling of top/bottom performers drives LLM analysis toward effective optimization directions; elite migration prevents premature convergence.

Key Contributions

  • Unified Sketch IR: hardware-agnostic intermediate representation decoupling optimization strategy from DSL syntax.
  • Document-Driven Integration framework enabling zero-code extension to new DSLs and hardware targets.
  • Three-stage hierarchical retrieval (semantic β†’ structural β†’ shape) for kernel-specific RAG.
  • Island-model multi-round search with LLM-based contrastive analysis for performance tuning.
  • New benchmark: 198–214 operators across 8 categories, static + dynamic shapes, with reward-hacking vulnerabilities from KernelBench fixed.

Results

  • KernelBench Level 1, Triton-Ascend: 100% pass@4 overall; 1.46Γ— geometric mean speedup; 77.6% fastβ‚€.β‚ˆ, 65.5% fast₁.β‚€.
  • KernelBench Level 1, Triton-CUDA: 100% pass@4; 1.06Γ— geometric mean speedup.
  • KernelBench Level 1, CPP-CPU: 91% pass@4; 1.04Γ— geometric mean speedup; Scan & Loss at 9.00Γ— geometric mean.
  • AKG benchmark (dynamic), Triton-CUDA: 90.9% pass@4; element-wise 100%, reduction 93.3%, normalization 92.6%.
  • AKG benchmark (dynamic), Triton-Ascend: 85.4% pass@4.
  • TileLang-CUDA and CUDA-C show significantly lower pass@4 (44% and 59% respectively) on KernelBench, indicating maturity gaps.
  • MatMul correctness is hardest across all backends (63.6–72.7% on AKG benchmark); convolution reaches 0% on TileLang/CUDA-C.

Limitations

  • TileLang and CUDA-C backends achieve substantially lower correctness (44–59%) versus Triton, indicating uneven DSL maturity.
  • Convolution operators on TileLang and CUDA-C are entirely unsupported/untested (0% or excluded).
  • MatMul correctness is consistently the weakest operator category, suggesting the Unified Sketch abstraction struggles with tiled matmul complexity.
  • Sorting operators on Ascend achieve 0% correctness in both dynamic and static settings.
  • Performance evaluation uses DeepSeek V3.1 as the LLM backbone; results are not decomposed to isolate agentic contributions from model capability.
  • Island-model search runs only R=3 rounds with K=2 islands, which may under-explore on harder operators.
  • No comparison against other multi-agent kernel systems (Astra, QiMeng-Kernel) in the quantitative tables.

Relevance to Harnesses / Meta-Harnesses

AKG kernel Agent is a concrete instantiation of a domain-specific multi-agent meta-harness: a Conductor orchestrator routes subtasks between specialized sub-agents (Designer, Coder, Verifier) based on runtime error classification, closely mirroring the harness pattern of adaptive task routing with feedback loops. The Document-Driven Integration framework is effectively a plugin/extension API for the harness β€” adding a new DSL or backend requires only a conforming DocSpec, not code changes β€” a key meta-harness design principle. The iterative island-model search loop demonstrates how a harness can govern a population of parallel agent invocations and synthesize cross-run learning without manual intervention. This work is relevant as a case study in how harness-level orchestration (closed-loop, error-type-aware routing, retrieval-augmented context injection) unlocks capabilities that single-model pipelines cannot achieve, and shows how to evaluate such harnesses rigorously with correctness + performance metrics.