CUDA-Agent-Ops-6K est un ensemble de données d'entraînement pour la génération et l'optimisation de noyaux CUDA, publié par BytedTsinghua-SIA en 2026, conçu pour fournir des tâches d'entraînement de haute qualité au niveau des opérateurs destinés à l'apprentissage par renforcement agentique et à la génération de code CUDA.

Cet ensemble de données comprend 6 000 tâches d'entraînement sur des ensembles d'opérateurs synthétiques, principalement issus des opérateurs de PyTorch et de Transformers, combinés en tâches de fusion plus complexes via des LLM. Les données sont stockées au format Parquet et sont principalement utilisées pour la génération de noyaux CUDA, l'optimisation et l'entraînement par apprentissage par renforcement agentique.

Composition de l'ensemble de données

• ops : liste des opérateurs ou modules impliqués dans la tâche ;
• data_source : identifiant de la source des données ;
• code : code Python/PyTorch exécutable de la tâche.

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}
}