Skip to content

Releases: alibaba/atrex-kernel-agent

Atrex Kernel Agent v0.2.0 — Orchestrated optimization loop

Choose a tag to compare

@smallmou smallmou released this 18 Jul 12:03

Atrex Kernel Agent(AKA) v0.2.0

This release grows AKA from a single interactive skill into a dual-route, profile-driven kernel optimization system, adds a full NVIDIA profiling toolchain, and significantly expands the GPU knowledge base.

Highlights

  • Two optimization routes. The same profile-driven workflow can now be driven either interactively (Route 1: gpu-kernel-optimizer Skill + hooks) or unattended (Route 2: orchestrator/optimize.py). Both share the same knowledge base, tools, and structured memory format.
  • Orchestrated loop (Route 2). orchestrator/optimize.py owns the outer loop, spawning a fresh clean session per iteration over a single git workspace. State crosses sessions only through disk, HEAD is always the best kernel, and termination is mechanical (max iterations / token budget / target-utilization short-circuit).
  • SOL-ExecBench native. Drive a single-operator campaign directly from a SOL-ExecBench op directory (definition.json, reference.py, workload.jsonl) via --op-dir; ground-truth files are used verbatim and never edited.

Features

  • Clean-session-per-iteration optimization loop with git-native, restart-safe stall tracking.
  • Automated Triton→Gluon convert session (runs inline, gated by a stall counter via --convert-after).
  • Subagent architecture under agents/ (baseline, profiler, research decompose, convert, partial-restart) shared by both routes.
  • Decompose step for fusion ops.
  • Stall-count gated search strategy in the Stage-2 research phase.
  • NVIDIA profiling toolchain: tools/profile_nvidia.sh, tools/ncu_helpers/, tools/classify_ncu.py, tools/extract_nvidia_asm.py, tools/extract_ttgir.py.
  • Submodule-based reference management for reference-projects/ and 3rdparty knowledge/tooling.

Atrex Kernel Agent v0.1.0 — Initial Open-Source Release

Choose a tag to compare

@smallmou smallmou released this 16 Jul 07:47

Atrex Kernel Agent (AKA) v0.1.0

Initial open-source release. AKA is an end-to-end Agent project for GPU kernel implementation, analysis, profiling, and iterative optimization. It helps an Agent turn PyTorch logic or an existing kernel into a high-performance GPU kernel through a structured, profile-driven workflow.

✨ Highlights

  • Isolated optimization workspace: Automatically creates a dedicated workspace under /tmp/kernel_opt_<name>/.
  • Hardware-aware targets: Looks up target hardware specs from the local gpu-wiki knowledge base and runs Roofline analysis to set auditable performance targets.
  • Profile-driven optimization loop: Builds a correct baseline kernel first, then iterates — profile with ncu (NVIDIA) or rocprofv3 (AMD), extract bottleneck evidence, query gpu-wiki / reference projects / web sources, write an evidence-based plan, apply one optimization category, validate correctness and performance, record memory, and commit — until the Stop Conditions are met.
  • Fully traceable: Records plans, profile artifacts, structured memory, reports, and Git commits for every accepted iteration.
  • Multi-platform: Supports NVIDIA (ncu) and AMD (rocprofv3, wrapped by tools/profile_kernel.sh).
  • Multiple Agent runtimes: Works with both Codex and Claude Code.

📦 Installation

./install.sh