Quick Start

Get up and running with Atrex-Bench in minutes using Docker and the built-in CLI tools.

Run with Docker

The ROCm runtime image bundles PyTorch, AITER, FlyDSL, and evaluation tools. 

docker run -it --rm \
    --device=/dev/kfd --device=/dev/dri --group-add video \
    -e ANTHROPIC_API_KEY -e OPENAI_API_KEY \
    -v $PWD:/workspace \
    <YOUR_REGISTRY>/atrex-bench:rocm7.2

Run Evaluation

Evaluate a candidate kernel against a reference operator. 

python scripts/run_eval.py \
  --input path/to/candidate.py \
  --reference-dir data/fused_moe \
  --output results/run

Generate Kernel

Use an LLM CLI to generate a candidate kernel from a reference. 

python scripts/run_generate.py \
  --operator fused_moe \
  --backend triton \
  --cli codex