CUDA-Agent-Ops-6K is a training dataset for CUDA kernel generation and optimization, released by BytedTsinghua-SIA in 2026. It aims to provide high-quality operator-level training tasks for agentic reinforcement learning and CUDA code generation.

The dataset contains 6,000 synthetic operator set training tasks, primarily derived from operators in PyTorch and Transformers, with multiple operators combined into more complex fused tasks via LLMs. The data is stored in Parquet format and is mainly used for CUDA kernel generation, optimization, and agentic reinforcement learning training.

Dataset Composition

• ops: A list of operators or modules involved in the task;
• data_source: Data source identifier;
• code: Executable Python/PyTorch task code.

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