论文
AMDKernelVault:用于 AMD GPU 内核优化的大规模数据集和代理训练
AMDKernelVault: Large-Scale Datasets and Agentic Training for AMD GPU Kernel Optimization
摘要
我们介绍 AMDKernelVault,这是一个开放的 HIP 和 Triton 内核语料库以及针对最新 AMD CDNA GPU 的训练框架。现有的基于 LLM 的内核代理主要以 CUDA/NVIDIA 为中心,并且通常依赖于重复的前沿 LLM 调用来进行生成、反射和优化。为了解决这一差距,我们开发了 HIPKernelGen 和 TritonKernelGen,这是代理驱动的管道,可将 PyTorch 引用转换为 HIP 或 Triton 内核,在 ROCm 下编译和验证候选者,并在 AMD 硬件上对它们进行延迟分析。该语料库包含 62,153 个经过执行验证的 HIP 内核样本、2,377 个基于生产的 ROCm 库 QA 条目以及 39,893 个 Triton 内核。我们通过监督微调和执行感知强化学习进一步训练 Qwen3-8B,作为语料库实用性的演示。在固定评估预算下,它在 PyTorch-to-HIP (34.0% Pass@1)、TritonBench-G (33.2% Corr@3) 和 ROCmBench (41.94% Corr@3) 上的比较模型中实现了最高的正确率,但并没有统一领先编译或速度指标。语料库和文档可在 https://huggingface.co/datasets/amd/AIG-Datasets 上获取,相关的训练和内核生成代码可在 https://github.com/AMD-AGI/hip_kernel_llm_lab 上获取。