CUDA-Agent-Ops-6K は、BytedTsinghua-SIA により 2026 年にリリースされた、CUDA Kernel の生成と最適化向けの学習データセットです。エージェント強化学習と CUDA コード生成のために、高品質な演算子レベルの学習タスクを提供することを目的としています。

このデータセットには、PyTorch と Transformers の演算子を中心に、6,000 件の合成演算子集合の学習タスクが含まれており、LLM を用いて複数の演算子をより複雑な融合タスクに組み合わせています。データは Parquet 形式で保存されており、主に CUDA Kernel の生成、最適化、およびエージェント強化学習の訓練に使用されます。

データセットの構成

• ops：タスクに含まれる演算子またはモジュールのリスト
• data_source：データソースの識別子
• code：実行可能な Python/PyTorch タスクコード

Citation

@misc{cuda_agent_2026,
  title        = {CUDA Agent: Large-Scale Agentic RL for High-Performance CUDA Kernel Generation},
  author       = {Dai, Weinan and Wu, Hanlin and Yu, Qiying and Gao, Huan-ang and Li, Jiahao and Jiang, Chengquan and Lou, Weiqiang and Song, Yufan and Yu, Hongli and Chen, Jiaze and Ma, Wei-Ying and Zhang, Ya-Qin and Liu, Jingjing and Wang, Mingxuan and Liu, Xin and Zhou, Hao},
  year         = {2026},
  howpublished = {Project page and technical report}
}

@misc{cuda_agent_ops_6k_2026,
  title        = {CUDA-Agent-Ops-6K: Training Dataset for CUDA-Agent},
  author       = {{BytedTsinghua-SIA}},
  year         = {2026},
  howpublished = {Hugging Face dataset},
  url          = {https://huggingface.co/datasets/BytedTsinghua-SIA/CUDA-Agent-Ops-6K}
}