makora-ai/triton-gpu-latency

Triton GPU延迟数据集是一个大型数据集，包含大量PyTorch问题（主要来自KernelBench），与候选Triton内核实现及其测量的GPU运行时配对，由MakoraGenerate生成。每一行是一个自包含的Python程序，定义了一个参考的PyTorch模型（使用普通PyTorch操作编写）和一个新模型（通过手写或生成的Triton内核重新实现相同的前向传递），标签是新模型的运行时。该数据集专为训练和评估生成快速GPU内核的模型而构建，例如用于内核合成大型语言模型的监督微调/强化学习数据、代码到延迟奖励模型的监督，或内核自动调优器的基准。数据集包括训练和测试分割，总行数约601k，程序大小在5k-15k字符之间，偶尔可达300k，运行时标签为浮点数（可为空表示失败），并包含问题ID和保留组等测试列。The Triton GPU Latency Dataset is a large dataset of PyTorch problems (mostly from KernelBench) paired with candidate Triton-kernel implementations and their measured GPU runtimes, generated by MakoraGenerate. Each row is a self-contained Python program that defines a reference Model written with plain PyTorch ops and a ModelNew that re-implements the same forward pass with a hand-written or generated Triton kernel, with the label being the runtime of executing ModelNew. It is built for training and evaluating models that generate fast GPU kernels, such as SFT/RL data for kernel-synthesis LLMs, supervision for code→latency reward models, or a benchmark for kernel autotuners. The dataset includes train and test splits with approximately 601k rows, programs typically 5k–15k characters (occasionally up to ~300k), latency labels as float64 (nullable for failures), and test-only columns like problem_id and holdout_group.

搜集汇总

数据集介绍

构建方式

该数据集旨在服务于高性能GPU内核生成领域，通过将PyTorch定义的参考问题与候选Triton内核实现及其在GPU上的运行时延配对构建而成。数据生成流程中，每一行均包含一个自包含的Python程序，该程序以标准PyTorch操作定义参考模型，并以手写或自动生成的Triton内核重新实现相同的前向传播，其运行耗时作为标签。数据集源自KernelBench等来源的PyTorch问题，借助MakoraGenerate流水线批量产出，总计约60万条记录，存储为Parquet格式，兼顾大规模与高效存取。

使用方法

该数据集灵活支持多种下游任务：通过过滤成功案例（延迟非空）可构建用于监督微调内核生成模型的数据；直接以延迟为回归目标可训练延迟预测器或奖励模型；将是否成功编译作为二分类标签可训练正确性验证器。使用时需注意，延迟值仅适用于同一问题内部的相对比较，跨问题对比意义有限。用户可通过分隔符提取候选内核代码，或利用测试集中的保留组划分评估模型在不同泛化难度下的表现。对于训练，建议对参考问题部分去重以避免数据泄漏。

背景与挑战

背景概述

在大规模深度学习模型加速与编译优化的前沿探索中，GPU内核自动生成与调优已成为提升计算效率的关键技术。近年来，Triton作为一种专为高效GPU内核设计的高级编程语言，凭借其简化CUDA开发流程的优势迅速崛起。由Makora AI团队于2023年创建并发布的Triton GPU Latency数据集，旨在为训练能够自动生成高性能Triton内核的模型提供大规模监督信号。该数据集涵盖约60万条PyTorch算子与对应Triton内核实现的配对样本，并记录了内核在GPU上的实际运行延迟，从而支撑了内核合成语言模型的微调与强化学习、代码到延迟奖励模型的训练，以及内核自动调优器的基准测试。作为一个跨代码生成与系统优化两大领域的重要桥梁，该数据集对推动高效GPU程序自动化合成的研究具有显著影响力。

当前挑战

该数据集核心解决的领域挑战在于，从高层算子描述到高性能GPU内核的自动映射仍高度依赖专家手工优化，缺乏大规模、可复用的训练数据来驱动自动化方法。数据集构建过程中面临多重困难：首先是延迟标签的获取，需在统一基准平台上执行大量候选内核并精确测量其运行时，而测量结果受硬件、批处理策略与预热机制影响，难以跨问题直接比较；其次是高失败率（约33%–38%），大量候选内核因编译错误、数值不匹配或运行异常而无法产生有效延迟，带来了稀疏标注与不均衡样本的挑战；此外，每个问题对应的参考模型在训练集中重复出现，需进行去重以避免数据泄漏，而测试集设计了三种不同重叠程度的分组以衡量泛化能力，进一步增加了数据集划分与评估的复杂性。

常用场景

经典使用场景

Triton GPU Latency数据集的核心设计理念在于为GPU内核合成领域提供大规模监督信号。该数据集收录了超过60万个从KernelBench衍生的PyTorch问题及其对应的Triton内核实现，每一条数据都包含了完整可执行的Python程序与实测的GPU运行延迟。借助这种（程序，延迟）的配对结构，研究者能够将其直接用作指令微调（SFT）素材，仅筛选出成功运行的核函数用于训练内核生成模型，也可作为训练延迟预测器或奖励模型的回归任务数据，更可通过缺失标签来训练正确性分类器。数据集中精心划分的训练集与测试集，以及全留出、重度留出等不同评估模式，为衡量模型在不同泛化难度下的表现提供了坚实基础。

解决学术问题

该数据集精准回应了高性能计算与机器学习交叉领域中一个核心挑战：如何自动化地生成高效GPU内核。传统依赖手工调优的CUDA内核编写方式既耗时又需要深厚专家经验，而自动内核生成技术则长期受困于缺乏大规模、高质量的训练数据。Triton GPU Latency数据集通过提供数十万个带有真实硬件延迟标签的Triton内核候选方案，使研究者得以训练模型理解代码结构与运行时性能之间的复杂映射关系。这不仅解决了代码生成领域长期存在的性能反馈缺失问题，更为内核自动调优器提供了标准化评估基准，推动了从启发式搜索向数据驱动优化范式的学术转变。该数据集的开放发布显著降低了GPU内核合成研究的入门门槛，其包含的失败案例数据也为鲁棒性验证器研究开辟了新方向。

实际应用

在实际工程应用中，基于该数据集训练的模型可被集成到深度学习编译器流水线中，实现从PyTorch计算图到高效Triton内核的自动生成与优化。对于需要频繁部署新型模型架构的AI公司，该技术能够大幅缩短算子开发周期，替代手动编写CUDA内核的传统流程。在云GPU服务场景中，延迟预测模型可协助调度系统选择最优内核变体，提升资源利用率并降低用户等待时间。此外，该数据集衍生出的错误检测分类器可嵌入到代码审核环节，自动识别潜在的内核编译失败或数值精度问题，为生产环境的稳定性提供保障。这些应用共同指向一个愿景：让非专家级的深度学习工程师也能获得专业级的GPU内核优化能力。

数据集最近研究

最新研究方向

triton-gpu-latency数据集聚焦于GPU内核生成与编译优化的前沿方向，通过提供大规模PyTorch问题与Triton内核实现及其运行时延迟的配对样本，为内核合成大语言模型的监督微调与强化学习、代码到延迟奖励模型的训练，以及内核自动调优器的基准测试提供了关键数据支撑。该数据集反映了深度学习领域对定制化GPU内核以提升模型推理效率的迫切需求，特别是在大模型部署场景中，自动生成高效内核已成为突破性能瓶颈的关键路径。其核心价值在于将内核生成转化为可量化优化的机器学习问题，通过海量成功与失败案例的双重标注，既推动了生成模型对高性能内核的模仿能力，又为验证器提供了负样本学习机制，从而加速了从硬件感知代码生成到自动化性能优化的闭环生态建设。