End-to-End Operator Generation Benchmark for Multiple DSLs
Production-trace driven · Executable PyTorch references · Source-aware metadata
Atrex-Bench measures whether an agent platform + model stack can turn a PyTorch reference into a DSL kernel that compiles, runs correctly, and approaches the achievable hardware peak (speed-of-light, SOL).
Every operator ships as a small, self-contained directory: an executable PyTorch Model, an
input generator, a shape spec, and hidden source-aware metadata. The evaluator runs a candidate
through three stages — compile → correctness → performance — and reports SOL efficiency against
a cached roofline.
- [2026/07] The preliminary technical report is released, introducing Atrex-Bench and Atrex-Kernel-Agent, a profile-driven agent for iterative GPU kernel optimization.
- [2026/06] The first version of Atrex-Bench is released: a production-trace-driven benchmark focused on compute-limited, memory-rich GPU fleets such as XPU-A and H20.
- 30 operators, mostly derived from real production traces (vLLM / SGLang / AITER / rtp-llm).
- 4 DSL backends out of the box: Triton, Gluon, FlyDSL, CuteDSL.
- Multi-vendor GPUs: the same operators, references, and evaluator run on both AMD (ROCm) and NVIDIA (CUDA), with prebuilt images for each.
- Three-stage evaluator: compile, numerical correctness, and performance vs. roofline SOL.
- Generation harness that drives an LLM CLI (Claude Code / Codex) inside a scoped workspace.
- Optional optimization agent: Atrex Kernel Agent provides a profile-driven workflow for GPU kernel implementation and iterative optimization.
- Leak-resistant by design: a one-shot cleanup strips the checkout to the agent-visible surface, and provenance / roofline numbers are never staged for the agent.
Everything runs inside a container — a clean runtime with PyTorch, GPU-accelerated kernels, Node 22, and the Claude Code / Codex CLIs, and no project code.
You git clone the repo inside the container (see Usage), so the code you run is
always your own checkout — no host mounts.
We provide prebuilt images for both AMD and NVIDIA GPUs:
| Platform | Image |
|---|---|
| AMD (ROCm) | treinfra/atrex-bench:rocm7.2_ubuntu22.04_py3.10_pytorch_release_2.10.0 |
| NVIDIA (CUDA) | treinfra/atrex-bench:2.10.0-cuda12.8-cudnn9-devel |
1 · Pull or build the image for your GPU platform:
# AMD GPU (ROCm)
docker pull treinfra/atrex-bench:rocm7.2_ubuntu22.04_py3.10_pytorch_release_2.10.0
# or build from source:
docker build -t atrex-bench:rocm -f docker/Dockerfile.rocm .
# NVIDIA GPU (CUDA)
docker pull treinfra/atrex-bench:2.10.0-cuda12.8-cudnn9-devel
# or build from source:
docker build -t atrex-bench:cuda -f docker/Dockerfile.cuda .2 · Start the container. Everything under Usage — including your API key — runs inside this shell:
# AMD GPU
docker run -it --rm \
--device=/dev/kfd --device=/dev/dri --group-add video \
treinfra/atrex-bench:rocm7.2_ubuntu22.04_py3.10_pytorch_release_2.10.0
# NVIDIA GPU
docker run -it --rm \
--gpus all \
treinfra/atrex-bench:2.10.0-cuda12.8-cudnn9-develEverything below runs inside the container you started above (Python ≥ 3.10 and PyTorch
already ship in the image). The example uses the ROCm image; on the CUDA image the commands are
the same, except you skip source /opt/venv/bin/activate (it ships PyTorch in the system Python).
Point Claude Code at your API key. Append this to ~/.bashrc — fill in your own
ANTHROPIC_API_KEY (and set ANTHROPIC_BASE_URL only if you go through a gateway instead of
the official API):
case ":$PATH:" in
*":$HOME/.local/bin:"*) ;;
*) export PATH="$HOME/.local/bin:$PATH" ;;
esac
export CLAUDE_CODE_EXPERIMENTAL_AGENT_TEAMS=1
export IS_SANDBOX=1
export ANTHROPIC_API_KEY="<YOUR_ANTHROPIC_API_KEY>"
# export ANTHROPIC_BASE_URL="<YOUR_API_BASE_URL>" # optional: custom gateway
export ANTHROPIC_MODEL="<YOUR_MODEL>"
claude() {
command claude --dangerously-skip-permissions --effort max "$@"
}Reload it with source ~/.bashrc. (Prefer Codex? Export OPENAI_API_KEY instead and pass
--cli codex below.)
Generation and evaluation run as separate sessions (e.g. separate pods). The generation
agent has full filesystem access, so its pod must never hold the answer files (metadata.json,
roofline.json, the evaluator) — otherwise it can just read the targets. So you clone the repo
and strip those files in place before generating; evaluation uses its own fresh, full clone.
git clone the repo, run the cleanup script, then generate:
git clone https://github.com/alibaba/atrex-bench.git && cd atrex-bench
source /opt/venv/bin/activate # activate the bundled venv (ROCm image)
python scripts/cleanup_for_generation.py # strip the answer files in place
pip install -e . --no-deps --no-build-isolation # deps already ship in the image
python scripts/run_generate.py \
--operator attention_forward \
--backend flydsl \
--cli claude \
--output-dir outputs/attention_forward_flydsl \
--mirror-tracecleanup_for_generation.py keeps only each operator's reference.py / input.py /
shapes.json, the prompt/ templates, and the generation runner; it deletes
metadata.json, roofline.json, configs/, the evaluator, tests/, the README, and the
.git history, so provenance and SOL targets cannot leak. It rewrites the checkout in
place — copy out your generated kernel before evaluating, and never run cleanup in an
evaluation session.
- Backends:
triton,gluon,flydsl,cutedsl. CLI:--cli claude(default) or--cli codex. - Each run writes
generated_kernel.pyplus ageneration.jsonbundle and a trace sidecar under the output directory.
Atrex Kernel Agent provides the
gpu-kernel-optimizer skill for profile-driven kernel implementation and optimization.
Follow its
Interactive Skill installation guide
before running the command below.
The default AKA installation base is ~/aka_kernel_opt. If you installed AKA with a custom
prefix, set AKA_HOME to that prefix. run_generate.py does not take a skill-directory
argument; it lets the selected CLI discover skills from its config home. Point that config
home at the matching AKA subdirectory when launching generation.
The optimizer requires an explicit target platform. Set the platform in the selected backend's optimizer prompt before generating:
AKA_HOME="$HOME/aka_kernel_opt"
TARGET_GPU="<TARGET_GPU_MODEL>"
sed -i "s/for the target GPU/for the ${TARGET_GPU} GPU/" \
prompt/flydsl/generate_kernel_with_optimizer.md
CLAUDE_CONFIG_DIR="$AKA_HOME/.claude" \
python scripts/run_generate.py \
--operator attention_forward \
--backend flydsl \
--template prompt/flydsl/generate_kernel_with_optimizer.md \
--cli claude \
--output-dir outputs/attention_forward_flydsl_agent \
--mirror-trace \
--skillFor Codex, set CODEX_HOME="$AKA_HOME/.codex" and use --cli codex. Ensure
that $CODEX_HOME/config.toml contains a top-level model = "..." entry.
The optimizer template explicitly invokes gpu-kernel-optimizer; --skill enables the
tools required by its workflow. Runs generated with --skill use a different agent workflow
and must not be compared directly with non-skill runs.
Evaluation needs the full repo (metadata.json, roofline.json, configs/, and the
evaluator), so run it in a separate session with a fresh clone you do not strip:
git clone https://github.com/alibaba/atrex-bench.git && cd atrex-bench
source /opt/venv/bin/activate # activate the bundled venv (ROCm image)
pip install -e . --no-deps --no-build-isolation # deps already ship in the image
python scripts/run_eval.py \
--input path/to/attention_forward_flydsl/generated_kernel.py \
--reference-dir data/attention_forward \
--output results/attention_forward_flydsl \
--num-correctness-cases 5 \
--bench-iters 3--input— the candidate from the generate stage (a single file exposingclass Model). Bring thegenerated_kernel.pyyou produced into this session and point--inputat it.--reference-dir— an operator directory underdata/.--num-correctness-cases/--bench-iters— correctness samples per shape (default 1) and timed perf iterations (default 100).- The output directory archives every input artifact plus
eval_result.json.
Measure the torch.compile baseline instead of a candidate:
python scripts/run_eval.py --torch-compile --reference-dir data/attention_forward --output results/torch_compileCompare a candidate (B) with a baseline (A) using a counterbalanced A-B-B-A schedule in one evaluator invocation:
python scripts/run_eval.py \
--baseline-input path/to/baseline_kernel.py \
--input path/to/candidate_kernel.py \
--reference-dir data/attention_forward \
--output results/attention_forward_abbaEach of the four runs performs the full compile, correctness, and performance
pipeline in an isolated worker. The top-level eval_result.json reports the
per-shape and aggregate candidate speedup over the baseline. With managed clock
locking, one lock and monitoring session covers the complete A-B-B-A sequence.
To evaluate an explicit shape subset, add --shape-ids 0,2 or
--shape-ids @ids.json (a JSON array of original shape IDs). The SDK config field
is shape_ids. Selection keeps the full reference bundle and the same RNG
seeds, tolerances, and timing settings as a serial run. Candidate, Torch compile,
and ABBA modes support it. Results mark the selected IDs and partial coverage in
shape_selection.
Run disjoint shards on independent GPUs, then merge their result dictionaries
or JSON paths with atrex_bench.merge_evaluation_shards(shards, output="merged.json").
Merging requires complete coverage and matching input/evaluator hashes,
settings, hardware, toolchain, and seeds. See the
shape shard contract and Gateway examples
for agate eval --shape-shards and manual subset usage.
Inside an existing Ray cluster, Atrex-Bench can split the complete shape set and merge the report automatically:
python scripts/run_eval.py --distribution ray_data --ray-max-workers 4 \
--ray-shapes-per-shard 1 --input candidate.py --reference-dir data/attention_forward \
--output results/distributedThe same options work in evaluate(config). The CPU driver snapshots the full
reference and evaluator sources. The Ray Data pool defaults to one GPU worker per
shard; an explicit positive worker limit caps that count. Each actor reserves one
GPU and runs the ordinary
evaluator in fresh subprocesses. The report retains driver placement and per-shard
actor/node IDs, interpreter, GPU assignment, and execution times so callers can
verify worker overlap instead of inferring it from the configured maximum. These
times cover shard evaluations, not individual CUDA kernels. Node-local driver
versions are retained per shard; the merged report uses driver_version: null
because its work spans independently managed nodes.
GPU model/architecture and software runtime must still match. Install
atrex-bench[ray] with matching native dependencies in the driver and worker
images. See the Ray Data runtime and cleanup contract
for file transport, clock locking, ABBA, and lifecycle requirements.
Worker/node failures retry only the affected shard (up to three retries by default, four total attempts). Successful
shard reports are committed on the driver node and reused even if their worker reply
is lost. Candidate failures are not retried. The admitted GPU worker limit remains
unchanged; a driver/registry loss fails the run instead of replaying successful work.
See the architecture and failure-path diagrams for the bounded worker pool, result commits, retries and cleanup mapped to code.
atrex-bench/
├── configs/ # hardware SKU profiles (roofline peaks)
├── data/ # one directory per operator (+ data/README.md)
├── prompt/ # generation-stage prompt templates
├── scripts/ # CLI entrypoints: generate, cleanup, eval, roofline, trace
├── src/atrex_bench/ # package: generation runner, evaluator, roofline
└── tests/ # unit + end-to-end tests
Each operator directory is self-describing:
data/<operator>/
├── reference.py # class Model(nn.Module) — definition only
├── input.py # _make_inputs(**kwargs) -> dict[str, Tensor]
├── shapes.json # shape spec keyed by id (init_kwargs + input_kwargs)
├── metadata.json # id, dtype, input/output dtypes, origin (hidden from agent)
└── roofline.json # cached W / Q / SOL_time_ms per device (hidden from agent)
See data/README.md for the full schema and the data-maintenance workflow.
data/ ships 30 operator directories, each following the same five-file contract above.
Most are trace-derived (status: "trace_reference"); the rest are curated.
data/operator_importance.json holds trace-based prioritization scores.
Full operator list
When adding or reshaping operator data, keep shapes.json and roofline.json aligned and
refresh SOL_time_ms through scripts/roofline.py.
See SOL semantics for the optional dispatch-latency
floor. Use scripts/static_roofline_check.py to validate saved SOL values and
compare them with production latency on the matching device without running kernels.
Please cite our paper if it is helpful to your research.
@misc{atrex2026,
title = {Are LLM-Generated GPU Kernels Production-Ready? A Trace-Driven Benchmark and Optimization Agent},
author = {Lingyun Yang and Yuxiao Wang and Shenghao Liang and Linfeng Yang and Daocheng Ying and Chunbo You and Rui Zhang and Luping Wang and Yinghao Yu and Guodong Yang and Liping Zhang},
year = {2026},
eprint = {2607.14541},
archivePrefix = {arXiv},
primaryClass = {cs.AI},
url = {https://arxiv.org/abs/2607.14541}
}Licensed under the Apache License 2.0.
Copyright 2026 Alibaba Group.