QwenLM/D2K-Bench

QwenLM/D2K-Bench

Language: Python

Stars: 0

Forks: 0

Open issues: 0

Created: 2026-09-29T05:27:57Z

Pushed: 2026-10-05T02:51:59Z

Default branch: main

Fork: no

Archived: no

README:

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](SETUP.md) for task-table generation, configurable paths and API credentials. Start with [.env.example](.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
python capability/scorer/score.py --release-root --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}
}