CUDA-Agent-Ops-6K ist ein 2026 von BytedTsinghua-SIA veröffentlichter Trainingsdatensatz für die Generierung und Optimierung von CUDA-Kernels, der darauf abzielt, hochwertige Trainingsaufgaben auf Operator-Ebene für das Reinforcement Learning von Agenten und die CUDA-Codegenerierung bereitzustellen.

Der Datensatz umfasst 6.000 Trainingsaufgaben für synthetische Operator-Sets, die hauptsächlich aus Operatoren in PyTorch und Transformers stammen und durch die Kombination mehrerer Operatoren zu komplexeren Fusionsaufgaben mittels LLMs erstellt wurden. Die Daten werden im Parquet-Format gespeichert und dienen hauptsächlich der Generierung und Optimierung von CUDA-Kernels sowie dem Reinforcement-Learning-Training von Agenten.

Zusammensetzung des Datensatzes

• ops: Liste der Operatoren oder Module, die in der Aufgabe beteiligt sind;
• data_source: Kennung der Datenquelle;
• code: ausführbarer Python/PyTorch-Aufgabencode.

Zitation

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