GPUMODE/KernelBook

这是一个由PyTorch程序和等价的Triton代码对组成的精选数据集，用于训练模型以实现PyTorch代码到Triton代码的转换。数据集包含了从GitHub上收集的PyTorch代码仓库，并通过一系列处理步骤生成了对应的Triton代码。每个数据条目都包含了仓库名称、许可证信息、GitHub星标数、提交SHA散列和指向特定提交的仓库链接等元数据。

This is a curated collection of pairs of PyTorch programs and equivalent Triton code, which can be used to train models to translate PyTorch code to Triton code. The dataset includes PyTorch code repositories collected from GitHub, with corresponding Triton code generated through a series of processing steps. Each entry in the dataset contains metadata such as repository name, license information, GitHub star count, commit SHA, and a direct link to the repository at the specific commit.

搜集汇总

数据集介绍

构建方式

在深度学习编译器与代码生成领域，PyTorch到Triton的自动翻译是提升计算效率的关键技术路径。KernelBook数据集通过系统化的流水线构建而成：首先从Stack v1数据集中筛选出涵盖大量PyTorch代码的GitHub仓库，随后提取其中的torch.nn模块并为其创建单元测试以验证代码正确性。利用PyTorch 2.5.0的torch.compile工具，将验证后的代码编译生成对应的Triton内核，并参照KernelBench格式进行规范化处理。最后，为每个条目补充许可证、星标数、提交哈希等元数据，形成结构化的高质量配对数据集。

特点

该数据集以PyTorch程序与等价Triton代码的配对为核心，具有鲜明的实用性与规范性。每个条目包含仓库名称、许可证列表、星标数、提交哈希及仓库链接等丰富元数据，便于研究者追溯代码来源与验证其可靠性。数据仅收录采用MIT、Apache 2.0等宽松许可证的仓库，确保在研究与开发场景中的合规使用。提供JSON与Parquet两种格式，兼顾人类可读性与高效分析需求，尤其适合大规模深度学习编译器的训练与评估。

使用方法

研究者可便捷地加载该数据集以训练或评估PyTorch到Triton的代码翻译模型。通过Python的json库可直接读取JSON版本，获取包含代码配对与元数据的列表；或使用pandas的read_parquet方法加载Parquet格式，利用DataFrame进行高效筛选与分析。例如，可基于星标数阈值筛选高质量仓库，或按许可证类型过滤数据。推荐在PyTorch 2.5.0环境下运行生成的Triton代码以复现最佳效果，同时建议在使用前核实个别仓库的许可证条款。

背景与挑战

背景概述

KernelBook数据集由Sahan Paliskara和Mark Saroufim于2025年创建，依托于PyTorch社区与GPU模式研究团队，旨在解决深度学习领域内PyTorch代码到高性能Triton内核的自动翻译问题。随着深度学习模型规模的急剧膨胀，手动编写和优化GPU内核成为性能瓶颈，而Triton作为一种领域特定语言，能够简化内核开发并提升计算效率。该数据集从GitHub上收集了经过许可筛选的PyTorch仓库，通过提取独立模块、构建单元测试、利用torch.compile生成Triton代码，并参照KernelBench格式进行转化，最终形成了一对一的代码映射集合。其发布为模型训练提供了标准化基准，推动了代码生成与编译器优化技术的交叉融合，对降低高性能计算门槛具有显著影响力。

当前挑战

KernelBook面临的核心挑战在于所解决的领域问题与构建过程中的双重复杂性。领域层面，PyTorch到Triton的自动翻译需应对深度学习操作的高度多样化，包括动态形状、控制流及非标准算子，现有编译器难以保证生成代码的正确性与性能最优性，且Triton生态尚不成熟，缺乏统一优化策略。构建过程中，数据采集面临代码质量参差不齐的难题，大量仓库缺乏可复现的测试环境，导致单元测试创建失败而需剔除；同时，torch.compile生成的Triton代码可能包含冗余或依赖特定PyTorch版本（2.5.0），限制了跨版本兼容性。此外，许可筛选虽聚焦于宽松许可证，但许可证变体众多，增加了元数据标注的复杂性，确保数据合法可用仍需人工核查。

常用场景

经典使用场景

在深度学习编译器与高性能计算领域，KernelBook数据集为PyTorch程序到Triton内核的自动翻译提供了宝贵的训练资源。该数据集精心收集了来自GitHub开源仓库中经过许可的PyTorch模块，并通过torch.compile工具链生成对应的Triton代码，形成了一一对应的程序对。研究者可将其用于训练序列到序列模型或代码翻译模型，旨在自动将PyTorch算子实现转换为高效的Triton内核，从而降低手动编写GPU内核的门槛，加速深度学习模型的推理与训练。

衍生相关工作

KernelBook数据集的发布催生了一系列富有影响力的后续工作。研究者基于该数据集探索了基于Transformer的代码翻译架构，如CodeT5与StarCoder的微调版本，在PyTorch到Triton的翻译任务上取得了突破性进展。同时，该数据集也被用于评估和对比不同编译器后端（如Inductor与XLA）生成内核的质量，推动了编译器自动化评估基准的建立。此外，部分工作利用该数据集进行迁移学习，将知识泛化至其他框架（如JAX）的内核生成任务，展示了其在跨框架代码优化中的广泛适用性。

数据集最近研究

最新研究方向

在深度学习编译器与代码生成领域，GPUMODE/KernelBook数据集的出现标志着PyTorch到Triton代码自动翻译研究迈入新阶段。该数据集通过系统化收集GitHub上的PyTorch仓库，提取独立模块并利用torch.compile生成等价Triton内核，构建了首个大规模、高保真的代码配对资源。当前前沿研究方向聚焦于利用该数据集训练神经翻译模型，实现PyTorch算子到Triton的高效自动转换，以突破手动编写高性能GPU内核的瓶颈。该工作与KernelBench等基准测试紧密关联，为GPU内核优化、编译器自动调优以及端到端深度学习加速提供了标准化训练与评估平台。其开源许可特性与详尽的元数据标注（如星标数、许可证类型）进一步推动了可复现研究与跨仓库代码迁移的实践，有望显著降低Triton编程门槛，加速AI基础设施的底层创新。