CUDA-Agent-Ops-6K

我们发布了训练数据集CUDA-Agent-Ops-6K：包含6000个训练样本，构建流程包括从`torch`和`transformers`中收集参考运算符，使用LLM将多个运算符组合成融合任务，并应用基于规则的过滤以保留可执行、确定性和非平凡的样本。过滤标准包括必须在eager模式和`torch.compile`中正确执行，移除随机运算符和退化输出，控制运行时范围并移除与KernelBench测试高度相似的样本以减少污染风险。

We introduce the training dataset CUDA-Agent-Ops-6K, which contains 6000 training samples. The dataset construction workflow includes collecting reference operators from `torch` and `transformers`, combining multiple operators into fused tasks via LLMs, and applying rule-based filtering to retain executable, deterministic, and non-trivial samples. The filtering criteria include ensuring correct execution in both eager mode and `torch.compile`, removing random operators and degenerate outputs, controlling runtime scope, and eliminating samples highly similar to KernelBench tests to reduce contamination risks.

原始信息汇总

CUDA-Agent数据集概述

数据集基本信息

• 数据集名称: CUDA-Agent-Ops-6K
• 发布方: BytedTsinghua-SIA
• 数据集地址: https://huggingface.co/datasets/BytedTsinghua-SIA/CUDA-Agent-Ops-6K
• 数据集规模: 6,000个训练样本

数据集构建方法

• 数据来源: 从torch和transformers中收集参考算子。
• 合成过程: 使用大语言模型将多个算子组合成融合任务。
• 过滤流程: 应用基于规则的过滤，保留可执行、确定性和非平凡的样本。

数据集过滤标准

• 必须在eager模式和torch.compile中都能正确执行。
• 移除随机算子和退化输出。
• 控制运行时范围，并移除与KernelBench测试高度相似的样本，以降低污染风险。

相关资源

• SKILL.md文件: 包含智能体执行的工作流约束和优化规则。
• agent_workdir目录: 提供了一个标准化的智能体工作空间示例，用于实现完整的循环：实现CUDA内核 -> 编译 -> 验证正确性 -> 分析性能 -> 迭代。
• 关键组件: 包含原始PyTorch基线模型、使用自定义CUDA扩展的优化模型、共享的Python绑定注册基础设施、自定义CUDA/C++内核及其绑定、扩展构建脚本、正确性验证脚本以及性能分析脚本。

搜集汇总

数据集介绍

构建方式

在CUDA高性能计算领域，数据集的构建需兼顾多样性与实用性。CUDA-Agent-Ops-6K数据集通过系统化流程合成，首先从PyTorch和Transformers库中收集基础算子作为参考，随后利用大型语言模型将多个算子融合为复合任务，以模拟真实场景中的计算需求。为确保数据质量，采用规则过滤机制，仅保留在即时执行模式和torch.compile环境下均可正确运行的样本，同时排除随机性算子与退化输出，并控制运行时范围，有效降低了与基准测试集之间的相似性风险。

特点

该数据集在CUDA内核生成研究中展现出显著特点，其规模包含6000个训练样本，覆盖了从基础到复杂的算子组合，为模型训练提供了丰富的学习材料。数据集经过精心筛选，强调可执行性与确定性，避免了随机因素干扰，从而提升了训练过程的稳定性。此外，通过严格控制样本与现有基准测试的相似度，有效减少了数据污染的可能性，确保了评估结果的可靠性，为高性能CUDA代码生成任务奠定了坚实的数据基础。

使用方法

针对CUDA内核优化研究，数据集的使用需结合标准化工作流程。用户可参考提供的agent_workdir示例，其中包含完整的实现循环：从编写CUDA内核代码到编译验证，再到性能分析与迭代优化。关键文件如SKILL.md提供了代理执行的约束与规则，而配套的验证与性能分析脚本则支持自动化测试。通过运行内置命令，研究者能够快速构建自定义扩展，并在统一环境中评估生成内核的正确性与效率，从而加速实验进程并促进成果复现。

背景与挑战

背景概述

在人工智能与高性能计算交叉领域，自动生成高效的CUDA内核代码是提升GPU计算效能的关键挑战。CUDA-Agent-Ops-6K数据集由字节跳动与清华大学智能产业研究院（SIA）于2024年联合发布，旨在支持基于大型语言模型的智能体强化学习研究，以生成超越现有先进模型性能的CUDA内核。该数据集聚焦于解决复杂算子融合任务，通过系统化的数据合成流程，为模型训练提供了高质量、可执行的样本，显著推动了自动化代码生成技术在并行计算领域的应用边界。

当前挑战

该数据集致力于应对高性能CUDA内核自动生成这一核心难题，其挑战在于如何确保生成的代码不仅功能正确，还需在运行效率上超越手工优化及现有编译工具（如torch.compile）的基准。在构建过程中，研究团队面临多重挑战：需从PyTorch和Transformers等框架中筛选代表性算子，并利用大语言模型组合成复杂的融合任务；同时，通过严格的规则过滤，保证样本的可执行性、确定性与非平凡性，并有效控制与现有基准测试的数据污染风险，以维持评估的公正性与泛化能力。

常用场景

经典使用场景

在GPU加速计算领域，CUDA-Agent-Ops-6K数据集为大规模智能体强化学习提供了关键支持。该数据集通过从PyTorch和Transformers库中收集参考算子，并利用大型语言模型将其融合为复杂任务，构建了6000个训练样本。这些样本经过严格的规则过滤，确保其可执行性、确定性和非平凡性，从而为训练高性能CUDA内核生成模型奠定了坚实基础。数据集的核心应用场景在于驱动智能体系统实现从CUDA内核实现到编译、验证、性能剖析的完整迭代循环，显著提升了内核生成的效率与质量。

衍生相关工作

围绕CUDA-Agent-Ops-6K数据集，已衍生出一系列经典研究工作。其中，CUDA-Agent项目作为首个已知的基于强化学习训练并超越先进模型的成果，在KernelBench上实现了最先进的性能。数据集配套的SKILL.md文件提供了智能体执行的工作流约束与优化规则，为后续研究设立了标准。此外，基于该数据集构建的绑定注册基础设施、编译脚本及验证工具，进一步推动了LLM驱动的CUDA生成社区发展，为智能体强化学习在代码优化领域的应用开辟了新方向。

数据集最近研究

最新研究方向

在GPU加速计算领域，高性能CUDA内核的自动生成一直是提升计算效率的核心挑战。CUDA-Agent-Ops-6K数据集的发布，标志着基于大规模智能体强化学习的方法在代码生成任务中取得了突破性进展。该数据集通过从主流深度学习框架中收集算子并利用大语言模型进行融合任务构建，为训练能够超越顶尖商业模型的智能体提供了高质量、多样化的样本。当前研究前沿聚焦于如何将此类数据驱动的智能体进一步应用于复杂异构计算场景的优化，以及探索其在降低硬件编程门槛、推动自动化高性能计算方面的深远影响。这一进展不仅加速了AI编译技术的发展，也为未来自适应计算系统的构建奠定了重要基础。