drkernel-coldstart-8k

DR.Kernel冷启动数据集用于在DR.Kernel中进行监督微调（SFT）以初始化内核生成能力（Triton编码+迭代优化），然后再进行TRLOO/MRS/PR/PRS强化学习。数据集包含8,920条多轮对话轨迹，每条轨迹由5轮对话组成，共10条消息，角色顺序固定为用户和助理交替。数据以Parquet表格形式存储，包含字段如messages（完整的多轮聊天历史）、uuid（样本ID）、entry_point（入口类/函数名）、final_speedup（最终加速比）等。数据集通过KernelGYM反馈收集，每轮对话包含服务器反馈和改进请求，支持多轮SFT训练。数据来源于ByteDance-Seed/cudaLLM-data，并在此基础上构建了多轮轨迹（迭代反馈+改进）。使用该数据集时，请引用DR.Kernel和原始查询源。The DR.Kernel Cold-start Dataset is intended for supervised fine-tuning (SFT) within the DR.Kernel framework to initialize the kernel generation capability (Triton coding + iterative optimization), prior to reinforcement learning with TRLOO/MRS/PR/PRS. The dataset comprises 8,920 multi-turn dialogue trajectories, each consisting of 5 conversation rounds with a total of 10 messages, where the speaker sequence strictly alternates between user and assistant. The data is stored in Parquet table format, with fields including messages (full multi-round chat history), uuid (sample ID), entry_point (entry class/function name), final_speedup (final speedup ratio), and more. This dataset is collected via KernelGYM feedback, where each conversation round includes server feedback and improvement requests, enabling multi-turn SFT training. The dataset is derived from ByteDance-Seed/cudaLLM-data, and multi-turn trajectories (iterative feedback + improvement) are constructed based on this source. When using this dataset, please cite DR.Kernel and the original query source.

搜集汇总

数据集介绍

构建方式

在深度学习编译器优化领域，构建高质量的监督微调数据集对于提升模型生成高性能Triton内核的能力至关重要。drkernel-coldstart-8k数据集通过从强大的专有教师模型中蒸馏知识，基于KernelGYM平台收集了多轮交互轨迹。每条轨迹严格遵循五轮对话结构，其中用户提供初始的PyTorch模型代码与优化任务，助手生成Triton内核实现，随后系统反馈编译、正确性及加速比等执行信息，驱动助手进行迭代优化。该过程共积累了8920条完整轨迹，每条轨迹包含十轮固定角色的消息序列，确保了数据在反映真实优化迭代过程上的丰富性与一致性。

特点

该数据集的核心特征在于其结构化多轮对话格式，专门为训练模型掌握Triton代码生成与迭代优化能力而设计。每条数据记录一个完整的五轮优化轨迹，包含用户指令、助手代码实现以及基于KernelGYM反馈的持续改进过程。数据集提供了丰富的元数据，如最终加速比、最佳轮次索引、原始Python代码等，为模型性能评估与训练过程分析提供了细致支持。其固定的角色轮转顺序与增强的对话转换模式，确保了训练输入格式的标准化，同时所有轨迹均未启用思维链标记，保持了训练目标的纯粹性。

使用方法

使用该数据集时，可通过Hugging Face Datasets库直接加载，获取包含消息列表及各类元数据的结构化表格。在监督微调阶段，建议启用多轮对话配置，指定消息键与思维标记键，并设置适当的序列最大长度与截断方向。数据集主要用于DR.Kernel框架中强化学习前的冷启动训练，通过执行提供的脚本即可启动基于8B或14B模型的微调流程。研究人员可依据轨迹中的反馈与代码演变，深入分析模型在编译器优化任务中的学习动态与性能提升路径。

背景与挑战

背景概述

在深度学习与高性能计算领域，Triton作为一种高效的GPU编程语言，其内核代码的自动生成与优化已成为提升计算性能的关键研究方向。DR.Kernel Cold-Start数据集由香港科技大学自然语言处理实验室于2026年发布，作为DR.Kernel研究框架的重要组成部分，该数据集旨在通过监督微调阶段初始化模型在Triton内核生成与迭代优化方面的能力。其核心研究问题聚焦于如何利用多轮交互轨迹，模拟真实开发环境中基于反馈的代码迭代过程，从而为后续强化学习阶段奠定坚实基础，对推动代码生成与程序优化领域的自动化进程具有显著影响力。

当前挑战

该数据集致力于解决Triton内核代码生成与优化这一复杂领域问题，其挑战在于如何精准建模从初始PyTorch架构到高效Triton实现的转换过程，并确保生成代码在编译正确性、运行效率与硬件适配性上达到最优。在构建过程中，研究团队面临多轮轨迹数据收集与对齐的复杂性，需整合KernelGYM服务器提供的即时编译、正确性验证与性能剖析反馈，以构建连贯的迭代优化对话历史。此外，数据蒸馏自强教师模型，确保轨迹质量与多样性之间的平衡，以及处理大规模代码与反馈信息的结构化表示，均是构建过程中需要克服的技术难点。

常用场景

经典使用场景

在深度学习编译器优化领域，DR.Kernel Cold-Start数据集为Triton内核代码生成任务提供了标准化的监督微调基础。该数据集通过包含五轮多轮对话轨迹，模拟了从初始PyTorch架构到优化Triton内核的完整迭代过程，每一轮都整合了KernelGYM反馈机制。这种结构使得模型能够在强化学习之前，预先学习如何分析性能瓶颈、生成正确代码并依据执行反馈进行逐步优化，从而建立起稳健的内核生成能力。

解决学术问题

该数据集主要解决了代码生成模型中冷启动初始化与迭代优化能力不足的学术挑战。传统方法往往依赖单轮生成，缺乏对运行时反馈的适应性，而DR.Kernel Cold-Start通过多轮交互轨迹，为模型提供了学习如何根据编译正确性、速度提升等动态指标进行自我修正的范例。这不仅提升了模型在复杂优化任务中的泛化性能，也为将强化学习有效应用于代码生成领域奠定了数据基础，推动了程序合成与编译器优化的交叉研究。

衍生相关工作

围绕该数据集，衍生出了一系列专注于代码生成与强化学习结合的经典研究工作。例如，DR.Kernel论文本身提出了基于多轮反馈的强化学习框架TRLOO/MRS/PR/PRS，该数据集作为其监督微调阶段的核心数据支撑。同时，它也与KernelGYM评估平台紧密关联，促进了迭代优化基准的建立。此外，基于原始查询源ByteDance-Seed/cudaLLM-data的扩展，进一步推动了面向CUDA与Triton的代码生成数据集的丰富与标准化，为后续研究提供了可复现的基线。