D2K-BENCH

A GPU-kernel co-design benchmark: 26 tasks with reference implementations, per-task evaluation environments, packed task tables, example agent workspaces, and a design/code capability scorer.

This repository accompanies the paper D2K-Bench: Can LLM Agents Turn Expert Designs into Efficient GPU Kernels?.

Directory contents

• dataset/ — one directory per task: the naive reference, the per-task evaluator config and grader inputs, the model-visible task spec and starter, and the design rubric, schema, and expert design guidance.
• tasks/ — packed Parquet task tables (frontier_sol_triton.parquet, frontier_sol_cuda.parquet, frontier_sol_cutlass.parquet), 26 rows each.
• example_workspace/ — prebuilt agent workspaces per task (baseline, design guidance, and the cumulative L1 and L1+L2 ablations).
• docker/ — the Dockerfile and build script for the runtime image (shared base, per-task environments, and the sandbox server).
• script/ — packing, workspace build, batch launch, timing, grading, sandbox server/client, API proxies, and offline validation.
• capability/ — design authoring and design/code scoring (scorer/) with the prompt templates they load (prompt/).
• requirements-control.txt — Python dependencies for the control host.

Design contract

Each task carries expert design guidance split into three cumulative levels:

• L1 — High-level algorithmic insights (level_1_high_level_algorithmic_insights)
• L2 — Dataflow design (level_2_dataflow_design)
• L3 — Low-level optimization tricks (level_3_low_level_optimization_tricks)

The shared level definitions are in script/design_levels.json; cross-level reasoning is carried in cross_level_rationale. The capability scorer reports JS_design and JS_impl, each the arithmetic mean of the three level scores (L1 + L2 + L3) / 3, using the shared 0/25/50/75/100 anchors.

Install

Use Python 3.10+ on the control host.

See SETUP.md for task-table generation, configurable paths and API credentials. Start with .env.example; keep local keys and run records out of releases. Run python script/audit_public_release.py --history before publishing changes.

python -m pip install -r requirements-control.txt

Build the GPU runtime either with Docker (bash docker/build.sh up) or docker-free on a CUDA-matched host (python script/setup_envs.py). Rebuild the packed task tables from the dataset when needed:

python script/pack/convert_frontier_sol.py --all --framework triton \
  --dataset-root dataset --output tasks/frontier_sol_triton.parquet

Start

Validate the package offline (no GPU or network):

python script/validate_design_levels.py
python script/validate_submission.py

Start the GPU sandbox, run a task, and score a solution:

python script/sandbox/deploy.py local start
python script/run_batch.py setup --name run1 --tasks muon_optimizer_step__triton --run
python script/bench.py run --task muon_optimizer_step --targets expert,naive,solution --solution-dir <dir>
python capability/scorer/score.py --release-root <release> --out <out>

See the README.md in each subdirectory for its own contents and commands, and SETUP.md for deployment details.

Citation

If you use D2K-Bench in your research, please cite our paper:

@misc{li2026d2kbench,
  title         = {{D2K-Bench}: Can {LLM} Agents Turn Expert Designs into Efficient {GPU} Kernels?},
  author        = {Daifeng Li and Huiqiang Jiang and Chengruidong Zhang and Wei Wu and Xudong Guo and Jianhong Tu and Jianwei Zhang and Binhang Yuan and Dayiheng Liu},
  year          = {2026},
  eprint        = {2610.03226},
  archivePrefix = {arXiv},
  primaryClass  = {cs.LG},
  url           = {https://arxiv.org/abs/2610.03226}
}