drkernel-rl-data

DR.Kernel RL 数据集是一个用于强化学习（RL）训练的数据集，专门针对 Triton 内核生成任务。该数据集主要包含一个单轮查询池及相关元数据，用于在 KernelGYM 中进行多轮在线 rollout。数据集以 Parquet 表格形式存储，包含 71,996 行数据，每行代表一个优化任务。主要字段包括：`data_source`（数据源标签）、`prompt`（用于生成的聊天提示，当前版本为单轮用户输入）、`ability`（任务能力标签，如 `kernel_optimization`）、`reward_model`（奖励元数据，包含参考 PyTorch 代码）和 `extra_info`（用于 rollout 和奖励跟踪的辅助元数据）。数据集适用于 RL 训练，模型通过接收初始提示生成候选代码，并通过 KernelGYM 进行多轮反馈和评估。数据集基于 ByteDance-Seed/cudaLLM-data 构建，并添加了 RL 相关的元数据。

搜集汇总

数据集介绍

构建方式

在强化学习驱动的代码生成领域，DR.Kernel RL数据集通过精心设计的数据结构，为模型训练提供了高质量的查询池与评估基准。该数据集源自ByteDance-Seed/cudaLLM-data，经过重构与增强，形成了包含71,996条优化任务的Parquet文件。每条记录均包含单轮用户提示、任务能力标签、奖励模型元数据及辅助信息，其中奖励模型部分嵌入了参考PyTorch代码作为评估基准，确保了训练过程中奖励信号的可靠性与一致性。

特点

该数据集的核心特征在于其专为在线多轮强化学习训练而优化，所有提示均为单轮用户查询，便于在KernelGYM环境中发起交互式轨迹生成。数据集结构高度统一，每条记录均标注为kernel_optimization能力，奖励模型风格均为rule，数据源标识为cuda_llm，保证了训练任务的专注性与评估标准的一致性。此外，数据集提供了丰富的元数据字段，如入口点、操作符列表等，为代码执行与性能分析提供了必要的上下文信息。

使用方法

在强化学习训练流程中，该数据集作为初始查询池，驱动模型在KernelGYM环境中进行多轮代码生成与评估。训练时，模型接收提示字段作为上下文，生成候选Triton内核代码，随后系统依据奖励模型中的参考代码进行编译、正确性及性能分析，并将反馈融入后续轮次。用户可通过Hugging Face Datasets库直接加载数据，或结合DR.Kernel项目提供的训练脚本，配置TRLOO、MRS等算法进行端到端的模型优化。

背景与挑战

背景概述

在深度学习与高性能计算交叉领域，Triton内核生成作为优化GPU计算性能的关键技术，长期面临自动化与效率提升的挑战。DR.Kernel RL数据集由香港科技大学自然语言处理实验室于2026年创建，旨在通过强化学习范式推进Triton内核的自动生成与优化。该数据集作为DR.Kernel研究项目的核心组成部分，聚焦于内核优化任务，通过提供结构化的查询池与参考代码元数据，为在线多轮交互式训练奠定数据基础。其构建依托于ByteDance-Seed的cudaLLM-data，并经过重新组织以适配KernelGYM训练框架，显著推动了面向专用硬件的高效代码生成研究。

当前挑战

该数据集致力于解决Triton内核自动生成这一复杂领域问题，其核心挑战在于如何确保生成内核在功能正确性、编译通过性与运行效率之间达到最优平衡。构建过程中的挑战则体现在多维度：其一，需将原始代码优化任务转化为适合强化学习交互的序列化提示与奖励信号；其二，必须设计精确的元数据结构以封装参考代码、评估入口点及操作符信息，支撑在线执行与奖励计算；其三，需保证数据规模与质量足以驱动多轮策略优化，同时维持提示与真实代码间的一致性。这些挑战共同构成了数据集构建与效用发挥的关键瓶颈。

常用场景

经典使用场景

在深度学习编译器优化领域，DR.Kernel RL数据集为强化学习训练提供了核心的查询池与参考元数据。该数据集通过单轮用户提示启动在线多轮推演，模型根据提示生成Triton内核优化代码，随后在KernelGYM环境中执行评估，获取编译正确性、性能剖析等多维度反馈，进而驱动TRLOO、MRS等先进强化学习算法进行迭代优化。这一流程精准模拟了实际内核开发中的交互式调试场景，为自动化代码生成研究奠定了实验基础。

衍生相关工作

围绕该数据集衍生的经典工作主要包括DR.Kernel论文中提出的KernelGYM训练框架，以及基于TRLOO、MRS、PR、PRS等强化学习算法的内核生成模型。这些研究系统性探索了多轮交互奖励机制对代码生成质量的影响，并为后续工作如CUDA内核自动优化、跨硬件适配的编译器设计提供了方法论借鉴。数据集本身亦构建于ByteDance-Seed/cudaLLM-data等开源资源之上，促进了领域内数据资源的迭代与共享。

数据集最近研究

最新研究方向

在强化学习驱动的代码生成领域，DR.Kernel RL数据集正推动着前沿研究向多轮交互式优化方向深化。该数据集专注于Triton内核生成任务，通过提供大规模的单轮查询池与参考代码元数据，支撑在线多轮rollout训练。当前研究热点集中于利用KernelGYM平台，结合TRLOO、MRS等先进算法，实现模型在编译正确性、执行效率等多维度奖励信号下的迭代优化。这一方向不仅呼应了人工智能辅助高性能计算的需求，也为自动化硬件感知代码生成提供了可扩展的评估框架，对提升GPU编程效率与降低专家依赖具有显著意义。