AI-CUDA-Engineer-Archive

# The AI CUDA Engineer Archive 👷: Agentic CUDA Kernel Discovery, Optimization & Composition


We release [The AI CUDA Engineer archive](https://sakana.ai/ai-cuda-engineer/), a dataset consisting of approximately 30,000 CUDA kernels generated by [The AI CUDA Engineer](https://pub.sakana.ai/ai-cuda-engineer/paper). It is released under the CC-By-4.0 license and can be accessed via HuggingFace and interactively visualized [here](https://pub.sakana.ai/ai-cuda-engineer). The dataset is based on the Kernel tasks provided in [KernelBench](https://scalingintelligence.stanford.edu/KernelBenchLeaderboard/) and includes a torch reference implementation, torch, NCU and Clang-tidy profiling data, multiple kernels per task, error messages and speedup scores against torch native and compile runtimes.

We envision that this dataset can enable post-training of open-source models to perform better CUDA-enabling modules. This includes offline Reinforcement Learning, preference optimization, and standard supervised fine-tuning.


# Using the AI CUDA Engineer Archive 🕵️

```python
from datasets import load_dataset

dataset = load_dataset("SakanaAI/AI-CUDA-Engineer-Archive")

df_l1 = dataset["level_1"].to_pandas()
print(df_l1.columns)
# Index(['Op_Name', 'Level_ID', 'Task_ID', 'Kernel_Name', 'CUDA_Runtime',
# 'PyTorch_Native_Runtime', 'PyTorch_Compile_Runtime',
# 'CUDA_Speedup_Native', 'CUDA_Speedup_Compile', 'CUDA_Code',
# 'PyTorch_Code_Module', 'PyTorch_Code_Functional', 'Correct', 'Max_Diff',
# 'Error', 'NCU_Profile', 'Torch_Profile', 'Clang_Tidy',
# '__index_level_0__'],
# dtype='object')

df_l2 = dataset["level_2"].to_pandas()
df_l3 = dataset["level_3"].to_pandas()
print(df_l1.shape, df_l2.shape, df_l3.shape)
# (12157, 19) (12938, 19) (5520, 19)

# Group by correct and count
print(df_l1["Correct"].value_counts())
print(df_l2["Correct"].value_counts())
print(df_l3["Correct"].value_counts())
```

# Citing The AI CUDA Engineer Archive 👷

If you use **The AI CUDA Engineer** kernel archive in your research or engineering, please cite it as follows:

```
@article{lange2025aicudaengineer,
title={The AI CUDA Engineer: Agentic CUDA Kernel Discovery, Optimization and Composition},
author={Lange, Robert Tjarko and Prasad, Aaditya and Sun, Qi and Faldor, Maxence and Tang, Yujin and Ha, David},
journal={arXiv preprint},
year={2025}
}
```

# 《AI CUDA工程师档案库》👷：基于智能体的CUDA内核（CUDA Kernel）发现、优化与组合


我们发布了[《AI CUDA工程师档案库》](https://sakana.ai/ai-cuda-engineer/)，该数据集包含约30000个由[AI CUDA工程师（The AI CUDA Engineer）](https://pub.sakana.ai/ai-cuda-engineer/paper)生成的CUDA内核。本数据集采用CC-BY-4.0开源协议，可通过HuggingFace获取，并可在[此处](https://pub.sakana.ai/ai-cuda-engineer)进行交互式可视化。该数据集基于[KernelBench](https://scalingintelligence.stanford.edu/KernelBenchLeaderboard/)提供的内核任务，包含PyTorch参考实现、PyTorch、NCU与Clang-tidy性能分析数据、每个任务对应的多个内核、错误信息，以及相较于PyTorch原生运行时与编译运行时的加速比得分。

我们期望该数据集可支撑开源模型的后训练流程，以提升其CUDA相关模块的性能，涵盖离线强化学习、偏好优化与标准监督微调。


# 《AI CUDA工程师档案库》使用指南 🕵️

python
from datasets import load_dataset

dataset = load_dataset("SakanaAI/AI-CUDA-Engineer-Archive")

df_l1 = dataset["level_1"].to_pandas()
print(df_l1.columns)
# Index(['Op_Name', 'Level_ID', 'Task_ID', 'Kernel_Name', 'CUDA_Runtime',
# 'PyTorch_Native_Runtime', 'PyTorch_Compile_Runtime',
# 'CUDA_Speedup_Native', 'CUDA_Speedup_Compile', 'CUDA_Code',
# 'PyTorch_Code_Module', 'PyTorch_Code_Functional', 'Correct', 'Max_Diff',
# 'Error', 'NCU_Profile', 'Torch_Profile', 'Clang_Tidy',
# '__index_level_0__'],
# dtype='object')

df_l2 = dataset["level_2"].to_pandas()
df_l3 = dataset["level_3"].to_pandas()
print(df_l1.shape, df_l2.shape, df_l3.shape)
# (12157, 19) (12938, 19) (5520, 19)

# Group by correct and count
print(df_l1["Correct"].value_counts())
print(df_l2["Correct"].value_counts())
print(df_l3["Correct"].value_counts())

# 引用《AI CUDA工程师档案库》👷

如果您在研究或工程实践中使用**AI CUDA工程师**内核档案库，请按以下格式引用：

@article{lange2025aicudaengineer,
title={The AI CUDA Engineer: Agentic CUDA Kernel Discovery, Optimization and Composition},
author={Lange, Robert Tjarko and Prasad, Aaditya and Sun, Qi and Faldor, Maxence and Tang, Yujin and Ha, David},
journal={arXiv preprint},
year={2025}
}