PTXBench

PTXBench is an open-source environment for evaluating and building kernel agents that use architecture-specifc CUDA/PTX. It contains:

• mini-ptx-agent: the reusable agent, prompt, trajectory inspection, and benchmark implementation.
• FIBServe: a GPU profiling service derived from FlashInfer-Bench.

🚧 This repository is still under construction.

System overview

Get a first result

The smallest real PTXBench run asks one model to optimize one GEMM for three turns. It needs:

• an NVIDIA Hopper/H100-class FIBServe instance loaded with the AccRL/accrl-training FlashInfer Trace dataset;
• the ptxbench-eval:dev Docker image; and
• either an OpenAI-compatible Qwen endpoint or credentials for one of the hosted models supported by mini-ptx-agent.

Set up the Python environment and evaluator image:

uv sync --all-packages --group dev
docker build -f docker/Dockerfile.eval -t ptxbench-eval:dev .

Start FIBServe as described below, then select a model. For a Qwen endpoint, the served ID must exactly match MODEL_NAME:

export MODEL_NAME=Qwen3.6-27B
export ACCRL_MODEL_HOST=localhost:30062
export SERVICE_URL=http://localhost:10000

uv run ptxbench quickstart --check
uv run ptxbench quickstart --run

For example, a hosted OpenAI model can be used without ACCRL_MODEL_HOST:

export MODEL_NAME=GPT-5.4
export OPENAI_API_KEY=...
uv run ptxbench quickstart --run

Triton quickstart

The same quickstart can generate and evaluate Triton kernels through the standard evaluator and FIBServe setup; no Triton-specific Docker image is required. Pass --language triton to select the checked-in Hopper Triton prompt and configuration:

uv run ptxbench quickstart --language triton --check
uv run ptxbench quickstart --language triton --run

CUDA remains the default when --language is omitted.

Every run leaves the full trajectories/exp_000.json, evaluator logs, and a concise quickstart-result.json under data/eval_runs/quickstart-...-gemm/. CUDA candidates are saved as exp_000/kernel.cu, while Triton candidates are saved as exp_000/kernel.py; correctness-passing versions are additionally saved under success/exp_000/ as kernel_vN.cu or kernel_vN.py, respectively. The report deliberately distinguishes “the runner completed” from “the kernel was correct” and “the 1.0x target was achieved.” mini-ptx-agent writes trajectories in a JSON format compatible with mini-swe-agent and its trajectory tooling. Reprint any run with:

uv run ptxbench quickstart --report data/eval_runs/quickstart-...-gemm

Local paths

For local development, configure the project/data paths and the Compose trace-set collection:

export PTXBENCH_ROOT=/home/ubuntu/PTXBench
export PTXBENCH_DATA_ROOT="$PTXBENCH_ROOT/data"
export PTXBENCH_TRACESETS_ROOT="$PTXBENCH_ROOT/data/datasets"
export DATASET_ROOTS=/workspace/trace-sets/accrl-training
export MINI_PTX_AGENT_ROOT="$PTXBENCH_ROOT/packages/mini-ptx-agent"

PTXBENCH_DATA_ROOT may point at an existing directory to hold experiment artifacts.

PTXBENCH_TRACESETS_ROOT is a shared parent directory on the host. Each child is a complete FlashInfer Trace dataset with its own definitions/ and workloads/ directories. For example:

$PTXBENCH_TRACESETS_ROOT/
├── accrl-training/
│   ├── definitions/
│   └── workloads/
└── another-trace-set/
    ├── definitions/
    └── workloads/

Compose bind-mounts that parent directory read-only at /workspace/trace-sets:

host:      /home/ubuntu/PTXBench/data/datasets
container: /workspace/trace-sets

DATASET_ROOTS is always a colon-separated list selecting dataset directories inside the mount; a single dataset is simply a one-item list. Compose uses the shared parent because it cannot expand one environment variable into a variable number of bind mounts.

Download the quickstart trace set from AccRL/accrl-training into $PTXBENCH_TRACESETS_ROOT/accrl-training. To load another trace set, place it alongside that directory and extend the list:

export DATASET_ROOTS=/workspace/trace-sets/accrl-training:/workspace/trace-sets/another-trace-set

No files are merged on disk, and no dataset is copied into a Docker image.

Development setup

uv sync --all-packages --group dev

Install the optional Tinker training stack when reproducing an AccRL training run:

uv sync --all-packages --extra training --group dev
uv run python -c "import tinker, tinker_cookbook"

The training commands require a valid TINKER_API_KEY; keep it in the environment or a local .env file, which is ignored by Git.

Build the isolated kernel evaluator:

docker build -f docker/Dockerfile.eval -t ptxbench-eval:dev .

Build and start FIBServe:

cp .env.example docker/.env
docker compose --env-file docker/.env -f docker/compose.yaml up --build fibserve

The public experiment index starts at experiments/README.md. The source/data boundary and release procedure are documented in RELEASING.md. For Fixit, use experiments/fixit/README.md as the single start page and scripts/reproduce_fixit.sh as the runnable entrypoint.

Use scripts/smoke_fixit.sh --check for a non-mutating dependency and configuration preflight.

Use scripts/reproduce_fixit.sh --check to validate the complete source closure and from-scratch to run failure mining, Gemini repair, SFT, and the paper's expert-guided evaluation in order. KernelGen starts at experiments/kernelgen/README.md and uses scripts/reproduce_kernelgen.sh. ptxbench-inspect remains part of the supported CLI.

For a live run, select a dedicated OpenAI-compatible model endpoint. PTXBench does not default to a shared model serve:

MODEL_NAME=Qwen3.6-27B \
ACCRL_MODEL_HOST=localhost:30062 \
SERVICE_URL=http://localhost:11000 \
scripts/smoke_fixit.sh --run

To run the same orchestration from the agent container, set PTXBENCH_HOST_ROOT to the checkout's absolute host path. The path is mounted at the same location in the agent so sibling eval containers launched through the Docker socket can mount trajectory workspaces correctly:

export PTXBENCH_HOST_ROOT="$(pwd)"
docker compose -f docker/compose.yaml run --rm \
  --entrypoint bash agent scripts/smoke_fixit.sh --check

Harbor

PTXBench includes Harbor-compatible tasks that run on an unmodified Harbor checkout while using ptxbench eval and FIBServe for compilation, sanitization, and H100 benchmarking. The agent creates its implementation from scratch; before launching a run, the checked-in instruction can be refreshed from the live FIBServe definition to keep its task encoding aligned with AccRL. The integration guide covers image builds, prompt rendering, Harbor launch options, and the resulting ATIF trajectory.

Huggingface

Other example datasets include non-4096 sequence-length attention workload records at Genghan/accrl-training-heavy and a more diverse flashinfer-ai/flashinfer-trace. The byte-exact historical s0-s6 training parquets are also retained in Genghan/PTXBench-Qwen3.6-27B-SFT.

Major dependencies

• mini-swe-agent 2.4.6: an external, exactly pinned Python dependency used as the agent runtime and Docker environment abstraction.
• tinker-cookbook: an optional, locked dependency for reproducing the Tinker SFT training and checkpoint-export workflow.

Citation

If you use PTXBench in your work, please cite the paper:

@misc{ptxbench2026,
  title = {{PTXBench}: Benchmark and Adapt {LLMs} for {GPU} Kernel Optimization with Architecture-specific {PTX}},
  author = {Genghan Zhang and Yixin Dong and Chengze Fan and Zhichen Zeng and Yueming Yuan and Shaowei Zhu and Kunle Olukotun},
  year = {2026},
  eprint = {2608.17379},
  archivePrefix = {arXiv},
  primaryClass = {cs.CL},
  url = {http://arxiv.org/abs/2608.17379},
}

Licensing and provenance

PTXBench and mini-ptx-agent are Apache-2.0. FIBServe is derived from FlashInfer Bench and preserves its Apache-2.0 license and NOTICE. mini-swe-agent is consumed under its MIT license and is not vendored into this repository. tinker-cookbook is consumed under Apache-2.0 and is also not vendored. See NOTICE for the NVIDIA reference-material rationale, hosted-model workflow terms, and licenses for code adapted from Helion and FlashAttention.