Source-linked AI summary
DataKernelBench: Can LLMs Optimize Database Queries on GPUs?
Gokul Karthik Kumar, Yotam Perlitz, Corey Lammie, Andrea Giovannini, Katja Hose
TL;DR
Generic ML compilation and existing kernel benchmarks leave irregular, heterogeneous database-style operators insufficiently tested, despite the potential value of query-specific GPU optimization. DataKernelBench translates SQL into validated TorchPlan programs and evaluates LLM-generated CUDA or Triton rewrites through execution-guided repair. Across TPC-H workloads, it reports strong gains for selected queries, including 2.11× over compiled TorchPlan at 100% pass rate, while extending to partitioned multi-GPU execution.
Problem
Existing LLM kernel benchmarks focus on regular ML operators, leaving irregular, heterogeneous database-style operators and their query-specific optimization insufficiently evaluated.
Method
DataKernelBench translates SQL into validated PyTorch TorchPlan programs and evaluates core or full CUDA and Triton rewrites through execution-guided repair.
Results
Across ten LLMs on TPC-H SF10, GPT-5.5 with CUDA-full achieves 2.11× over compiled TorchPlan with 100% pass rate; SF100 partitioned execution reaches 2.54× over its partitioned TorchPlan baseline.
Takeaways & Limitations
LLM-generated kernels are a useful optimization path for selected recurring queries, with gains associated with kernel fusion, fused aggregation, and broader execution-strategy changes.
Takeaways & Limitations
The primary evaluation covers TPC-H SF10 on one H100, while broader workloads, hardware generations, and data distributions remain unevaluated; the framework does not cover arbitrary SQL-to-kernel generation.
Abstract
from arXiv · showhide
GPUs increasingly accelerate database systems, but query-specific peak performance still often relies on hand-written kernels. Existing LLM kernel benchmarks focus on machine learning operators, leaving irregular, heterogeneous, data-movement-heavy database-style operators untested. We introduce DataKernelBench, which translates SQL into validated PyTorch TorchPlan programs and evaluates LLMs that optimize either the core tensor-bounded snippet or the full query in CUDA or Triton through execution-guided repair. Across ten proprietary and open-weight models on TPC-H SF10 with an H100 GPU, the strongest full-query CUDA configuration achieves $2.11\times$ speedup over torch.compile at full pass rate. We find that higher-performing implementations commonly use kernel fusion and execution-strategy changes, stronger models benefit most from full-query specialization, and workload context matters more than hardware context. To handle data larger than GPU memory, we extend TorchPlan with Dask-cuDF for on-demand partition loading on TPC-H SF100 with four H100 GPUs, achieving $2.54\times$ speedup
1 Introduction
Analytical workloads expose irregular, heterogeneous computation that can exceed generic compiler performance, motivating LLM-generated query-specific GPU kernels. DataKernelBench addresses this gap with validated TorchPlan programs and controlled CUDA or Triton specialization.
- Motivation: Analytical queries combine irregular memory access, complex predicates, heterogeneous operators, and query-specific fusion opportunities that challenge generic compilation.Peak performance often still requires bespoke GPU kernels, whose development and maintenance are expensive.
- Research gap: The benchmark fills a gap left by ML-focused kernel benchmarks, which do not test database-style operators with irregular access and strict output semantics.Prior benchmarks primarily evaluate dense, regular ML operators.
- Approach: The framework translates SQL into validated PyTorch TorchPlan programs that expose stable optimization targets while preserving query semantics.TorchPlan separates table handling in run_query from the tensor-intensive _query_core hot path.
- Motivation: DataKernelBench targets recurring queries whose repeated execution can justify query-specific specialization after validation and cost-benefit assessment.Deployments can fall back to compiled TorchPlan or a general-purpose GPU query engine when specialization is not worthwhile.
- Approach: DataKernelBench evaluates ten LLMs across optimization scope, programming interface, prompt context, and query structure using plan-level performance comparisons.The benchmark supports core and full specialization in CUDA or Triton under execution-guided repair.
2 Related Work
Prior GPU and AI-for-data-systems research provides libraries, database engines, tensor query processing, and LLM kernel synthesis, but existing kernel benchmarks largely target ML operators. DataKernelBench shifts this evaluation toward heterogeneous analytical queries.
- GPU data systems: GPU analytical processing spans dataframe libraries, composable engines, and GPU-specialized database systems.Examples include cuDF, Velox, Sirius, Kinetica, and SQream.
- Tensor Query Processing: Tensor Query Processing maps relational operators to tensor computations on ML runtimes and has demonstrated terabyte-scale multi-GPU analytics.Later work narrows the mismatch between SQL operators and tensor operations and supports compressed-data execution.
- LLM kernel synthesis: LLM kernel-generation systems iteratively synthesize CUDA or Triton kernels, while benchmarks such as KernelBench and TritonBench evaluate this capability on ML operators.These benchmarks provide the methodological foundation for DataKernelBench.
- Research gap: Analytical query processing differs from ML kernels through heterogeneous relational operators, irregular data access, and strict output semantics.DataKernelBench adapts prior evaluation styles to database-style workloads.
3 The DataKernelBench Framework
DataKernelBench constructs validated TorchPlan tasks from SQL, then measures whether LLMs can rewrite query implementations with correct and faster CUDA or Triton kernels. It separates table handling from tensor-hot-path specialization and evaluates both core and full scopes.
- Framework: The benchmark first constructs and validates a baseline TorchPlan, then asks LLMs to inject optimized Triton or CUDA kernels while preserving semantics.This two-stage design separates benchmark construction from model evaluation.
- Workload: TPC-H SF10 supplies 22 analytical queries over eight relational tables, with runtime execution using GPU-resident cuDF dataframes.DuckDB generates the data, SQL queries, and reference outputs.
- TorchPlan: TorchPlan acts as an executable intermediate representation and verified contract between declarative SQL logic and imperative GPU kernel generation.Each plan is differentially validated against DuckDB before entering the benchmark.
- TorchPlan: run_query handles table access, joins, projections, parameters, and output formatting, while _query_core isolates the tensor-intensive hot path over aligned GPU tensors.This decomposition enables specialization without changing the benchmark-facing query interface.
- Optimization scope: Core optimization restricts changes to _query_core and helpers, whereas full optimization permits rewriting internal query code while preserving the external run_query API.The two levels distinguish tensor-kernel-like optimization from broader query specialization.
- Evaluation: Correctness checks compare generated outputs hierarchically with the baseline, requiring exact non-floating-point matches and tolerance-based floating-point agreement.Validation failures, mismatches, tracebacks, and measured speedups feed execution-guided repair.
- Evaluation: Speedup is defined as Tbase/Tcand, and a candidate is accepted only when functionally correct and at least smin faster than compiled TorchPlan.Median runtimes are measured after warmup, with Tbase from torch.compile and Tcand from timed generated-module runs.
4 Experiments
On TPC-H SF10 with an H100, DataKernelBench evaluates ten LLMs generating CUDA or Triton kernels at core and full optimization levels against TorchPlan with torch.compile. Results show strong but query- and model-dependent gains, with full-query specialization, kernel fusion, workload context, and partitioned execution shaping performance.
- Comparative Evaluation: GPT-5.5 CUDA-full achieves 2.11× overall speedup over TorchPlan-compile at 100% pass rate, outperforming Sirius at 1.37×.Qwen3.5-397B-A17B Triton-full is the strongest open-weight configuration at 1.26× with 100% pass rate.
- Comparative Evaluation: Top-ranked models generally need fewer repair rounds and tokens; GPT-5.5 reaches the top result with 1.4 mean rounds, while lower-ranked open models require 3.5–5.9.Token cost alone is insufficient: Gemini 3.1 Pro reaches 1.44× speedup but passes only 54.5% of queries.
- CUDA vs. Triton Trade-offs: Triton is best for six models and CUDA for four, with backend preference varying by model rather than being universal.GPT-5.5, Gemini 3.1 Pro, and Claude Opus 4.7 fall in the CUDA-preferred region, while several open-weight models favor Triton.
- Understanding Performance Gains: GPT-5.5 beats Sirius on 8/13 queries referencing at most three tables but only 2/9 queries referencing more than three, despite winning overall runtime.Sirius is fastest on many queries while GPT-5.5-generated kernels dominate a different subset, indicating strong per-query non-uniformity.
- Understanding Performance Gains: Kernel fusion is dominant: GPT-5.5 combines filtering, projection, and aggregation into fewer passes, while 21/22 implementations retain cuDF for string or DataFrame operations.This hybrid execution uses CUDA mainly for the numeric hot path and avoids intermediate masks, gathers, and temporary tensors.
- Understanding Performance Gains: Q14 reaches 11.2× speedup by replacing a materialized join with a compact promotional-part lookup probed during a lineitem scan.Alternatives retaining the cuDF merge remain near 1.07×, illustrating the impact of broader execution-plan changes.
- Effect of Optimization Scope: core vs. full: For top-5 models, full over core adds mean speedup gains of +0.31 for CUDA and +0.29 for Triton, versus only +0.02 and +0.05 for bottom-5 models.The results suggest end-to-end specialization is an emergent capability of stronger code LLMs.
- Prompt Ablation: Removing workload data drops speedup by −0.40 under CUDA and −0.22 under Triton, larger than the corresponding hardware-context drops of −0.27 and −0.06.The ablation indicates that table sizes, types, and other concrete data properties guide optimization more strongly than raw GPU specifications.
5 Conclusion
DataKernelBench evaluates LLM-generated GPU kernels for analytical queries and finds strong gains on selected recurring workloads, including beyond-single-GPU execution.
- 2.11× over compiled TorchPlan with 100% pass rate is achieved by GPT-5.5 with CUDA-full across ten LLMs on TPC-H SF10.
- Kernel fusion, fused aggregation, and broader execution-strategy changes explain gains identified through plan inspection.
- 2.54× speedup over the partitioned TorchPlan baseline is achieved on TPC-H SF100 across four H100 GPUs, with all 22 queries passing.
- The benchmark positions LLM-generated kernels as an optimization path for selected recurring queries and a testbed for LLM-driven database and GPU optimization.
6 Limitations
The evaluation has a deliberately controlled scope spanning workloads, software stacks, memory assumptions, and the SQL-to-kernel pipeline.
- The primary evaluation covers TPC-H scale factor 10 on one H100 GPU, while broader workloads, hardware, and data distributions remain unevaluated.
- The framework centers on PyTorch TorchPlan with Triton and CUDA, excluding interfaces such as Numba, CuTe DSL, and emerging domain-specific GPU frameworks.
- The primary pipeline assumes the working set fits device memory, while the multi-GPU proof of concept does not systematically evaluate spilling, unified memory, or cluster-size scaling.
- Validated baseline TorchPlans condition kernel synthesis on a fixed tensor program, so the benchmark does not cover arbitrary SQL through final optimized-kernel generation.
A Fused CUDA Kernel Example
The examples separate table handling from the tensor hot path, while partitioned execution changes data loading and multi-join queries expose strategy choices across functions.
- The GPT-5.5 CUDA Q6 variant keeps cuDF I/O and parameters in run_query while _query_core launches a compiled fused device kernel.
- Dask-cuDF changes Q6 table handling from eager cuDF to partitioned execution across GPUs while leaving the tensor hot path unchanged.
- Q11 keeps joins in run_query, while _query_core performs grouping, thresholding, and sorting after joined columns become tensors.
- Full optimization may change the Q11 join strategy and how work is divided between run_query and _query_core, unlike core optimization.
- Table 4 compares total workload runtimes for two TorchPlan generators, with optimized rows using CUDA-full and lower runtime preferred.
D TorchPlan Generator Sensitivity
Sensitivity analysis shows that optimizer conclusions persist across TorchPlans generated by different strong models, despite differences in baseline runtime and execution environments.
- GPT-5.5 and Claude Opus 4.7 each generated 22/22 valid TorchPlans, but their compiled runtimes were 1.74 s and 1.51 s, respectively.
- Both optimized implementations remain faster than their corresponding baselines, and GPT-5.5 remains the strongest optimizer across the two generator sets.
- Each optimizer is fastest on TorchPlans generated by the other strong model, suggesting cross-model diversity may be useful.
- The benchmark uses GPU acceleration for LLM-generated CUDA/Triton kernels and Sirius, while DuckDB provides the CPU-only database baseline.
F Repair Behavior
Execution-guided repair substantially improves LLM-generated kernel outcomes, but success varies by query, backend, and model. Failures arise mainly from compilation/import problems or execution crashes and timeouts.
- Repair effectiveness: 77.5% of trajectories eventually produce a correct kernel with at least 1.05× speedup.Cumulative success increases from 40% after round 1 to 57% after round 2 and 65% after round 3.
- Failure categories: Approximately 38% of round-1 failures fail to compile or import, 52% crash or time out, and 9% return incorrect output.Column mismatches account for less than 1%.
- Backend differences: Approximately 74% of CUDA failures are compilation or import errors, whereas 79% of Triton failures are execution errors.
- Query difficulty: Q13 is the hardest query, with 13/40 successful trajectories, while Q6 and Q19 succeed in all 40 trajectories.Q20, Q3, and Q1 are also difficult.
- Prompt improvements: Pinning runtime versions and specifying kernel-source and date-representation guidance reduces repair rounds without materially changing final speedup.
G Choke-Point Coverage Analysis
Choke-point coverage measures the number of explicitly annotated relational categories in each query. LLM specialization is relatively stronger on narrower query patterns, whereas Sirius becomes relatively stronger as queries combine more choke points.
- Coverage measure: Choke-point coverage counts distinct choke-point categories explicitly associated with each TPC-H query.The measure is conservative because unannotated relational patterns are not counted.
- Correlation results: All ten LLMs show negative correlation between coverage and per-query speedup, while Sirius shows positive correlation.
- Correlation results: ρ = −0.267 across pooled LLM-query pairs, with n = 220.
- Comparative pattern: The head-to-head LLM/Sirius speedup ratio decreases with coverage for all ten LLMs.The authors interpret this as LLM specialization being most effective for narrower query patterns, with Sirius relatively stronger on more relationally complex queries.
H Multi-GPU Scaling
Dask-cuDF scales TorchPlan execution across up to four H100 GPUs by concurrently processing table partitions. CUDA kernel plans retain an advantage over TorchPlan-compile, while the current partition layout limits scaling beyond four GPUs.
- Execution layout: Dask-cuDF runs independent table partitions concurrently through map_partitions, queuing remaining partitions until workers become available.Lineitem has four partitions and every other table has one, so four is the maximum concurrent partition count in this setting.
- Scaling results: Runtime over all 22 queries falls from 77.5 s to 36.44 s for CUDA kernel plans and from 273 s to 92.6 s for TorchPlan-compile when moving from one to four GPUs.
- Scaling boundary: At higher scale factors, additional GPUs could run more partitions in parallel.