AMDKernelVault: Large-Scale Datasets and Agentic Training for AMD GPU Kernel Optimization

2026-09-11

AI Summary

为解决AMD GPU内核优化问题，通过创建AMDKernelVault数据集和训练框架，并使用HIPKernelGen和TritonKernelGen方法生成、编译和验证内核。

Abstract

We introduce AMDKernelVault, an open HIP and Triton kernel corpus and training framework for recent AMD CDNA GPUs. Existing LLM-based kernel agents are largely CUDA/NVIDIA-centric and often depend on repeated frontier-LLM calls for generation, reflection, and optimization. To address this gap, we develop HIPKernelGen and TritonKernelGen, agent-driven pipelines that transform PyTorch references into HIP or Triton kernels, compile and validate candidates under ROCm, and latency-profile them on AMD hardware. The corpus contains 62,153 execution-verified HIP kernel samples, 2,377 production-grounded ROCm Libraries QA entries, and 39,893 Triton kernels. We further train Qwen3-8B with supervised fine-tuning and execution-aware reinforcement learning as a demonstration of the corpus's utility. Under fixed evaluation budgets, it achieves the highest correctness among the compared models on PyTorch-to-HIP (34.0% Pass@1), TritonBench-G (33.2% Corr@3), and ROCmBench (41.94% Corr@3), but does not uniformly lead compilation or speed metrics. The corpus and documentation are available at https://huggingface.co/datasets/amd/AIG-Datasets, and the associated training and kernel-generation code is available at https://github.com/AMD-AGI/hip_kernel_llm_lab.

Problem

• LLM-based kernel agents
• CUDA/NVIDIA-centric
• AMD CDNA GPUs
• kernel optimization

Innovation

• AMDKernelVault
• HIPKernelGen
• TritonKernelGen
• execution-aware reinforcement learning
• ROCm