Skip to content
BranchesTags

Repository files navigation

DataKernelBench

Can LLMs optimize database queries on GPUs?

DataKernelBench evaluates LLMs on a novel task: optimizing analytical database queries as GPU kernels. It first represents each SQL query as a validated PyTorch program called a TorchPlan. It then evaluates LLMs by asking them to optimize either the tensor-intensive core (core) or the full query implementation (full) using CUDA or Triton, with execution-guided repair. The benchmark covers all 22 TPC-H queries.

On TPC-H SF10 using one NVIDIA H100, the strongest configuration—GPT-5.5 with CUDA at the full-query level—achieves 2.11× overall speedup over compiled TorchPlan with a 100% pass rate. The paper evaluates ten proprietary and open-weight models and analyzes how model strength, backend, optimization scope, and prompt context affect correctness and performance.

Links: Project page · Paper · Poster · Slides · Code · Data: Hugging Face · Zenodo

Benchmark artifacts

TPC-H SF10 (H100). Layout: artifacts/README.md.

Hands-on tutorial (self-contained, runs on Colab): artifacts/datakernelbench_tutorial.ipynb

Dask-cuDF / SF100 preview (data larger than one GPU; not a release): artifacts-dask/README.md.

Results

Column definitions for both CSV types: results/README.md.

Regenerate: poetry run python kernel_plans/prepare_results.py

Coding harness

1. Setup Code

git clone <datakernelbench repo URL>
cd datakernelbench
curl -sSL https://install.python-poetry.org | python3 -
poetry install

Copy .env.example as .env and fill the values.

2. Setup Data

utils/extract_benchmark.py uses DuckDB’s tpch / tpcds extensions (dbgen / dsdgen) to generate data at a chosen scale factor, export table CSVs, write query SQL, execute each query and save its result as CSV, emit logical plans as JSON (EXPLAIN (FORMAT JSON)), and save table metadata via SUMMARIZE and PRAGMA storage_info.

poetry run python utils/extract_benchmark.py
poetry run python utils/extract_benchmark.py scale_factor=0.01 benchmarks=[tpch]

Defaults live in utils/conf/extract_benchmark.yaml (benchmarks: [tpch, tpcds], scale_factor: 1).

Output layout (under benchmarks/<tpch|tpcds>/)

benchmarks/<benchmark>/
├── data/
│   └── sf<scale>/              # one folder per scale factor
│       ├── <table>.csv
│       └── …
├── meta_data/
│   └── sf<scale>/
│       ├── summarize/        # per-table: <table>.csv
│       └── storage_info/     # per-table: <table>.csv
├── sql_queries/              # shared across scale factors
│   ├── q01.sql
│   └── …
├── sql_query_results/
│   └── sf<scale>/            # result rows from executing each SQL query
│       ├── q01.csv
│       └── …
└── json_plans/
    └── sf<scale>/            # logical plan JSON for this SF’s statistics
        ├── q01.json
        └── …

sf<scale> uses dots as underscores (e.g. 0.01 → sf0_01).

3. Get Baseline Torch Plans

Generate

baseline/generate_torch_plans.py reads SQL from benchmarks/<benchmark>/sql_queries/, fills a prompt template, calls LiteLLM, and writes one .py plan per query under outputs/baseline/<benchmark>/ (shared across scale factors).

Defaults live in baseline/conf/generate_torch_plans.yaml. Logs: logs/baseline/<benchmark>/generate_torch_plans.log

poetry run python baseline/generate_torch_plans.py
poetry run python baseline/generate_torch_plans.py benchmark=tpcds scale_factor=1 max_llm_rounds=10
poetry run python baseline/generate_torch_plans.py prompt_path=prompts/baseline/base_with_json_plan.md
Setting Meaning
llm_provider litellm (default) or rits
model Model id (LiteLLM or RITS, e.g. deepseek-ai/DeepSeek-V3.2)
prompt_path Template (default prompts/baseline/base.md)
max_llm_rounds Repair loop with DuckDB validation when > 1
benchmark, scale_factor JSON plans and validation data use benchmarks/.../<sf>/

Use prompts/baseline/base_with_json_plan.md to inline DuckDB JSON plans from benchmarks/<benchmark>/json_plans/<sf>/ when the template contains {{INSERT_JSON_PLAN_HERE}}.

LLM credentials in .env: LiteLLM (LITELLM_BASE_URL, LITELLM_API_KEY) or RITS (RITS_BASE_URL, RITS_API_KEY).

poetry run python baseline/generate_torch_plans.py llm_provider=rits model=deepseek-ai/DeepSeek-V3.2

Validate and benchmark (shared config slice)

Validate and benchmark use the same Hydra keys for grouping artifacts: benchmark, scale_factor, gpu (auto-detected when null), and model (torch_eager or torch_compile). Use the same overrides on both commands so validation.csv and timing.csv land in one directory.

Per config slice:

results/<benchmark>/<sf>/<gpu>/baseline_<model>/
  validation.csv    # validate: id, pass, failure
  timing.csv        # benchmark: timings + source_path
  metadata.json     # benchmark run config

logs/baseline/<benchmark>/<sf>/<gpu>/baseline_<model>/
  validate_torch_plans.log

Join on id (e.g. pandas merge). Plans live at outputs/baseline/<benchmark>/qNN.py (no SF in path); table data from benchmarks/<benchmark>/data/sf<tag>/.

Validate — baseline/validate_torch_plans.py compares DuckDB SQL to each torch plan. Defaults: baseline/conf/validate_torch_plans.yaml.

poetry run python baseline/validate_torch_plans.py
poetry run python baseline/validate_torch_plans.py model=torch_eager benchmark=tpch scale_factor=1 limit=5
Setting Meaning
model torch_eager or torch_compile (same execution mode as benchmark: compile wraps _query_core when set)
benchmark, scale_factor Validation data from benchmarks/.../data/sf<tag>/
gpu Auto-detected when null (H100 / A100 / V100)
max_diffs Max differing rows printed per failure

Benchmark — baseline/benchmark_torch_plans.py times each plan (torch.compile when model=torch_compile). Defaults: baseline/conf/benchmark_torch_plans.yaml. No log files under results/; benchmark output is tqdm-only.

poetry run python baseline/benchmark_torch_plans.py
poetry run python baseline/benchmark_torch_plans.py model=torch_compile benchmark=tpch scale_factor=1
poetry run python baseline/benchmark_torch_plans.py model=torch_eager benchmark=tpcds scale_factor=10 limit=5
Setting Meaning
warmup, runs Timed repeats (defaults match kernel_plans)

Each benchmark run replaces timing.csv and metadata.json for that slice and purges __pycache__ under the repo first. Each validate run replaces validation.csv for that slice.

4. Get Optimized Kernel Plans

Prompt templates: prompts/kernel_plans/ (composition details in that folder’s README).

Before a run on a new host, refresh environment versions (included in composed prompts when environment.md exists):

poetry run python utils/write_environment_versions.py

Requires validated baseline torch plans at outputs/baseline/<benchmark>/.

Generate

kernel_plans/generate_kernel_plans.py composes prompts from prompts/kernel_plans/, calls LiteLLM, and runs a repair loop per query: correctness vs the baseline torch plan, speedup vs baseline torch.compile(_query_core) at run_query level (median after warmup). GPU is auto-detected (H100 / A100 / V100).

Defaults live in kernel_plans/conf/generate_kernel_plans.yaml. Override on the CLI with Hydra:

poetry run python kernel_plans/generate_kernel_plans.py
poetry run python kernel_plans/generate_kernel_plans.py benchmark=tpcds scale_factor=1 framework=triton level=core
poetry run python kernel_plans/generate_kernel_plans.py limit=5 max_rounds=10 min_speedup=1.05 gpu=h100
Key Default Description
framework triton triton or pytorch or cuda
level core core or full
benchmark tpch e.g. tpch, tpcds
scale_factor 1.0 dataset scale factor
gpu auto h100 / a100 / v100, or null to detect
llm_provider litellm litellm or rits
model (see yaml) LLM model id (LiteLLM or RITS)
max_rounds 10 LLM repair rounds per query
max_output_tokens 16384 LLM response token limit (API default is 4096)
min_speedup 1.05 vs baseline torch.compile (in-loop bench during generate)
warmup / runs 2 / 5 timing repeats (in-loop bench during generate)
load_baseline_timing true use results/.../timing.csv and validation.csv; fall back on demand if missing
warmup_timeout_coef / warmup_timeout_floor_min 20 / 1 correctness + bench warmup: max(1 min, 20× baseline_compile_ms); wall clock kills hung run_query
run_timeout_coef 3 bench timed runs: coef × baseline_compile_ms; -1 = off
query_wall_timeout_min 15 kill isolated per-query worker after this many minutes (LLM + nvcc + bench)
resume false keep outputs/ and logs/; skip any query with terminal results.json (pass, no_speedup, fail, …); re-run only incomplete queries, append logs
limit all max queries (null = all)
debug false tracebacks in terminal/log; per-round round_NN.debug.txt with LLM I/O
isolate_queries true one subprocess per query (fresh CUDA after bad kernels)

Only the fastest correct passing plan per query is written under outputs/kernel_plans/. Each generate run clears that run’s output and log directories first.

Per config slice (same layout as baseline results/, plus framework and level):

outputs/kernel_plans/<benchmark>/<sf>/<gpu>/<model>/<framework>/<level>/
  qNN.py                                  # best correct plan (or latest if none correct)

logs/kernel_plans/<benchmark>/<sf>/<gpu>/<model>/<framework>/<level>/
  generate_kernel_plans.log               # mirrors terminal output
  generate_kernel_plans.json              # run summary + per-query aggregates
  queries/qNN/
    results.json                          # round-wise status, tokens, speedup
    round_01.py
    round_01.debug.txt                    # debug=true: LLM input/output/status

<model> is the LiteLLM id with / → __ (e.g. aws/claude-opus-4-7 → aws__claude-opus-4-7). Use the same benchmark, scale_factor, gpu, model, framework, and level on generate and benchmark so paths align.

Benchmark

Re-benchmark existing kernel plans in the output folder (no LLM). Defaults: kernel_plans/conf/benchmark_kernel_plans.yaml. Use the same path keys as generate so it reads the right outputs/kernel_plans/.../ slice.

poetry run python kernel_plans/benchmark_kernel_plans.py
poetry run python kernel_plans/benchmark_kernel_plans.py benchmark=tpch scale_factor=1 framework=triton level=core
poetry run python kernel_plans/benchmark_kernel_plans.py limit=5 min_speedup=1.05 gpu=h100 model=aws/claude-opus-4-7
Key Default Description
framework, level, benchmark, scale_factor, gpu, model (see yaml) same run dimensions as generate
load_baseline_timing true same as generate — read baseline medians from results/ when available
min_speedup 1.05 vs baseline torch.compile
warmup / runs 2 / 5 timing repeats
max_diffs 5 max cell diffs shown on correctness mismatch
limit all max queries (null = all)
debug false verbose logging
isolate_queries true one subprocess per query (fresh CUDA after bad kernels)

Logs: logs/kernel_plans/<benchmark>/<sf>/<gpu>/<model>/<framework>/<level>/benchmark_kernel_plans.[json|log].

Results (same hierarchy as outputs/logs):

results/<benchmark>/<sf>/<gpu>/<model>/<framework>/<level>/
  validation.csv                          # id, pass, failure (vs baseline torch plan)
  timing.csv                              # candidate run_query medians + raw repeats
  metadata.json                           # run config + resolved_gpu

Prepare results

kernel_plans/prepare_results.py merges baseline and kernel-plan runs into one query-level CSV per benchmark. Defaults: kernel_plans/conf/prepare_results.yaml.

poetry run python kernel_plans/prepare_results.py
poetry run python kernel_plans/prepare_results.py benchmarks=[tpch]
poetry run python kernel_plans/prepare_results.py benchmarks=[tpch,tpcds] output_dir=results results_prefix=query_timing
Setting Meaning
benchmarks Which benchmarks to export (default [tpch, tpcds])
output_dir Where CSVs are written (default results)
results_prefix Output filename prefix (default query_timing → query_timing_tpch.csv)
logs_dir Root for kernel-plan logs (default logs/kernel_plans)
speedup_pass_levels Speedup thresholds in percent vs torch.compile for speedup_pass_rate_*pct columns (default [5, 10, 20])

Sources

  • Baselines — results/<benchmark>/<sf>/<gpu>/baseline_torch_*/ (timing.csv + validation.csv); tagged run_kind=baseline, framework=torch, level=full.
  • Model runs — logs/kernel_plans/<benchmark>/.../generate_kernel_plans.json; tagged run_kind=model.

Rows are sorted by hyperparameters then query_id. See results/README.md for column definitions.

Leaderboards — also writes output_dir/leaderboards/<benchmark>_sf<N>_<gpu>.csv (one file per benchmark × scale factor × gpu). Metrics are over queries where baseline torch.compile passed validation and has timing. Rows are ranked by sort_metric (default: overall_speedup_vs_torch_compile, higher is better).

Citation

If you use DataKernelBench, please cite:

@misc{datakernelbench,
      title={DataKernelBench: Can LLMs Optimize Database Queries on GPUs?}, 
      author={Gokul Karthik Kumar and Yotam Perlitz and Corey Lammie and Andrea Giovannini and Katja Hose},
      year={2026},
      eprint={2608.25061},
      archivePrefix={arXiv},
      primaryClass={cs.CL},
      url={https://arxiv.org/abs/2608.25061}, 
}

IBM Public Repository Disclosure All content in this repository including code has been provided by IBM under the associated open source software license and IBM is under no obligation to provide enhancements, updates, or support. IBM developers produced this code as an open source project (not as an IBM product), and IBM makes no assertions as to the level of quality nor security, and will not be maintaining this code going forward.

About

DataKernelBench evaluates LLMs on a novel task: translating SQL queries into optimized CUDA and Triton kernels via a PyTorch intermediate representation, then measuring whether they execute faster on GPUs.

kerneldf.github.io/datakernelbench

Topics

cudadatabasegpukernelsllmpytorchsqltriton

Resources

Readme
Apache-2.0 license
Activity
Custom properties

Stars

1 star

Watchers

0 watching

Forks

1 fork
Report repository

Releases

Packages

Contributors

Languages