用LLM优化数据库查询，实现在GPU上2.5倍提速
DataKernelBench: Can LLMs Optimize Database Queries on GPUs?
- 将SQL转为可验证的TorchPlan程序，用LLM修复并优化核心计算或全查询
- 在H100上对TPC-H SF10实现2.11倍加速，大模型数据集达2.54倍
- 适合关注数据库与LLM结合、GPU算子优化的研究者和工程团队
GPU正日益加速数据库系统，但查询级性能巅峰仍依赖人工编写内核。现有基于LLM的内核基准测试聚焦于机器学习操作，忽略了不规则、异构且数据移动密集的数据库式操作。我们提出DataKernelBench，将SQL转化为经验证的PyTorch TorchPlan程序，并通过执行引导修复评估LLM在CUDA或Triton中优化核心张量绑定片段或完整查询的能力。在十种专有及开源模型上，使用H100 GPU运行TPC-H SF10时，最优全查询CUDA配置相较TorchPlan基线实现2.11倍加速，且达到全通过率。研究发现，高性能实现普遍采用内核融合与执行策略变更；更强模型从全查询特化中获益最大；工作负载上下文比硬件上下文影响更大。针对超出GPU内存的数据，我们扩展TorchPlan支持Dask-cuDF，实现按需分区加载，在四张H100 GPU上的TPC-H SF100上达成2.54倍加速。
原文摘要 · Abstract (English)
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 the TorchPlan baseline 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. Project page: https://kerneldf.github.io/datakernelbench
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。