# atrex-kernel-agent

**Repository Path**: alibaba/atrex-kernel-agent

## Basic Information

- **Project Name**: atrex-kernel-agent
- **Description**: No description available
- **Primary Language**: Unknown
- **License**: Apache-2.0
- **Default Branch**: main
- **Homepage**: None
- **GVP Project**: No

## Statistics

- **Stars**: 0
- **Forks**: 0
- **Created**: 2026-06-24
- **Last Updated**: 2026-07-07

## Categories & Tags

**Categories**: Uncategorized

**Tags**: None

## README

# Atrex Kernel Agent (AKA)

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.

![Atrex architecture](atrex-architecture.png)

![Atrex optimization loop](atrex-optimization-loop.png)

## What It Does

- Creates an isolated optimization workspace under `/tmp/kernel_opt_<name>/`.
- Looks up target hardware specs from the local `gpu-wiki` knowledge base.
- Runs Roofline analysis and sets auditable performance targets.
- Implements a correct baseline kernel before entering optimization.
- Runs the profile-driven optimization loop: profile with `ncu` or `rocprofv3`, extract bottleneck evidence, query `gpu-wiki` / reference projects / web sources for relevant optimization knowledge, write an evidence-based plan, apply one optimization category, validate correctness and performance, record memory, commit, then repeat until Stop Conditions are met.
- Records plans, profile artifacts, structured memory, reports, and Git commits for every accepted iteration.

For the full architecture and workflow design, see [`docs/design.md`](docs/design.md).

## Requirements

Installation requires:

- `bash`
- `git`
- `jq`
- Codex or Claude Code installed

Running optimization tasks also requires platform-specific profiling tools:

- NVIDIA: `ncu`
- AMD: `rocprofv3`, wrapped by `tools/profile_kernel.sh`

## Installation

```bash
./install.sh
```

Common options:

```bash
./install.sh --hooks-only          # Install or update hooks only
./install.sh --without-github      # Skip GitHub reference repositories listed by gpu-wiki
./install.sh --max-iterations N    # Configure hook stop behavior after memory/vN.json exceeds N
./install.sh --uninstall           # Remove hooks installed by this script
```

The installer detects:

- Codex: `$CODEX_HOME` or `~/.codex`
- Claude Code: `$CLAUDE_HOME` or `~/.claude`

It also prepares the default local knowledge base at `/tmp/gpu-wiki/` and optional reference projects at `/tmp/reference-projects/`.

After installation, restart Codex / Claude Code or open a new session so the hooks and Skills are loaded.

## Quick Start

Ask the Agent to optimize a kernel with at least:

- `platform`: target hardware platform, such as `H20` or `MI308X`.
- `framework`: target implementation framework, such as `CuteDSL` or `FlyDSL`.
- `kernel_demo`: path to the initial PyTorch logic or kernel implementation file.

Example:

```text
/gpu-kernel-optimizer Optimize /path/to/kernel_demo.py on MI308X with FlyDSL, dtype bf16, rel_err < 0.01.
```

The Agent will initialize a workspace, source hardware specs from `gpu-wiki`, write the workspace configuration, build a baseline, profile the kernel, and iterate until the configured Stop Conditions are met.

## Main Files

```text
.
├── SKILL.md                         # Top-level gpu-kernel-optimizer Skill router
├── install.sh                       # Installer / uninstaller
├── docs/                            # Detailed project design docs
├── reference/                       # Workspace, plan, memory, and profiling templates
├── skills/                          # Baseline, optimizer, restart, and output-contract Skills
├── tools/                           # Profiling, utilization, memory, and measurement tools
└── gpu-wiki/                        # Local GPU knowledge base
```

## License

Licensed under the [Apache License 2.0](LICENSE).