面向AMD GPU的更好HIP内核生成：合成数据、多智能体搜索与强化学习

# 面向AMD GPU的更优HIP内核生成：合成数据、多智能体搜索与强化学习 来源：https://scalingintelligence.stanford.edu/blogs/hipkernels/ --- 图注插入处 ## 摘要 本文探索如何让语言模型在为**AMD GPU**生成高性能**HIP内核**时表现更好。我们提出以下内容： 1. 一个包含500个新PyTorch参考任务的**合成数据集**，采用**变异**、**组合**和**基于约束的生成**方法，覆盖更广泛的工作负载。 2. 用于HIP内核生成的**多智能体优化流水线**。我们不再依赖单次提示，而是使用专门的智能体进行任务生成、PyTorch到HIP的翻译、硬件评估和进化优化，以搜索更快的内核。 3. 一个基于**小型、低成本开源模型**的框架，使用**SFT**后接**GRPO**强化学习。虽然SFT帮助模型学习正确的HIP模式，但RL通过直接奖励MI350X GPU上的正确性和加速比进一步提升了性能。我们的结果显示，在所有KernelBench级别上，编译率和正确率均有提升，其中RL带来了显著增益。然而，要实现比PyTorch有意义的加速，仍需要更深入的硬件感知和优化推理。接下来，我们计划集成ROCm分析器，以教模型基于硬件分析器的奖励。 --- ## 动机 每个现代AI工作负载的性能瓶颈都在于内核质量。编写高性能内核需要深入熟悉硬件、底层语言以及优化技术，而这些在NVIDIA CUDA生态系统之外严重匮乏。AMD的HIP就是这种不足的典型例子。它是一种经过编译器验证的底层编程语言，可用的开源训练数据相对较少，但其目标加速器在生产AI集群中的部署却日益增多。这种不对称可以通过经验观察到：顶尖LLM通常能生成流畅的CUDA代码，但在生成HIP时，模型可能会虚构API，或者生成看似合理但在编译时或多种子正确性测试中失败的内核。 --- ## 方法 我们研究了三个互补的思路：(1) 通过**合成PyTorch工作负载**扩展任务空间，(2) 通过**多智能体进化搜索**优化内核，(3) 使用**SFT**后接**GRPO**强化学习训练一个小型、低成本的开源模型（**Qwen2.5-Coder-14B-Instruct**）。我们使用扩展到AMD **MI350X** GPU的**KernelBench**（Ouyang等，2025 (https://arxiv.org/abs/2502.10517)）评估所有方法在内核**编译**、**正确性**和**运行时性能**方面的表现。我们的方法如下： ## 1. 合成数据生成 我们使用一个基于Gemini-2.5-Flash的**多智能体流水线**，生成一个经过验证的HIP内核语料库，并配有相应的PyTorch参考。该流水线包含八个协作智能体： 图注插入处 **任务生成器**：将PyTorch参考封装成一个结构化任务，并通过变异*、组合*和基于约束的生成*方法综合出新的参考模块，每个合成模块在进入流水线前都经过合理性检查。 **翻译器**：从PyTorch参考生成第一个可工作的HIP内核，利用验证器的错误信息和之前的尝试结果重新提示，进行重试。对于每个合成数据任务，该智能体在五次尝试内生成了正确内核。 **正确性验证器**：确定性正确性门控，拒绝捷径模式，并在多个种子上将候选内核与PyTorch参考进行对比。 **进化优化器**：根据最相似的先前已验证内核（遵循Lange等，2025 (https://arxiv.org/abs/2509.14279)）、当前最佳内核以及最近失败记录，迭代采样新候选，保留最快的正确内核作为下一代的种子。 **合理性筛选器**：基于LLM的评审器，对每个候选内核的编译和合理性进行评分，只将有希望的内核送到GPU。 **硬件评估器**：在AMD MI350X GPU上编译每个幸存候选，在多个种子上检查与PyTorch参考的正确性，并测量运行时间。 **存档管理器**：将每个候选内核及其标签、分数和运行时间持久化到每个任务的存档中，并向下游的后训练阶段发出SFT和RL训练记录。 **离线审计器**：一对生成器和审计器，运行精心策划的正确、错误和欺骗性测试用例，通过两个验证器，并报告每个验证器相对预期标签的假正和假负情况。 ### *任务生成模式 **变异**：我们选取一部分现有的KernelBench问题，要求模型生成语义相关的变体。这些变异保留了原始工作负载的整体结构，同时修改计算属性，如操作组合、张量形状、批处理结构或融合模式。虽然与源任务密切相关，但生成的内核可能需要截然不同的优化策略。 **组合**：我们通过组合来自自定义14运算符模板库中的运算符，生成全新的工作负载。对于每个候选，我们随机采样一个运算符子集，并指示模型将它们组合成一个单一的形状兼容的PyTorch模块。重复采样可产生具有不同运算符顺序、张量形状和融合结构的多样化工作负载。 **约束**：我们直接通过自然语言约束来指定工作负载，描述所需的计算、张量属性和结构要求。模型必须解读规范，构建有效的模块架构，并生成可执行代码。 ## 2. SFT 我们在合成语料库上对**Qwen2.5-Coder-14B-Instruct**进行微调，训练3个周期，批次大小为2，学习率为2e-5。 ## 3. RL 我们应用基于**组相对策略优化（GRPO）**的强化学习于合成语料库，模型每个提示生成4个候选内核，并从候选之间的相对性能中学习。我们借鉴了Dr. Kernel的**Turn-Reinforce-Leave-One-Out (TRLOO)** 思想用于优势估计和奖励结构（Liu等，2026 (https://arxiv.org/abs/2602.05885)）。TRLOO通过排除组均值奖励估计中的一个候选，解决了自包含偏差问题。奖励信号包括在AMD MI350X硬件上的内核执行；如果内核编译并通过正确性检查，则获得正奖励，奖励幅度由测得的相对于PyTorch的加速比决定，上限为3倍。 图注插入处 我们实现了三个关键修改： 1. **多轮回合**：如果内核失败，模型会收到错误信息和失败尝试，并允许对组内的每个候选内核进行最多3次额外尝试（遵循Baronio等，2025 (https://cognition.ai/blog/kevin-32b#test-time-search)）。 2. **奖励平滑**：我们跟踪最近100个奖励的滚动窗口，并裁剪任何偏离均值超过1.5个标准差的异常值。这可以防止异常的GPU计时测量结果扭曲奖励信号。 3. **总结智能体**：在训练之前，我们在所有智能体流水线失败日志上运行一个总结智能体，以提取经验教训。这些教训直接注入RL提示中，使模型不仅从奖励信号中学习，还从关于过去应避免错误的指导中学习。 --- ## 结果与讨论 ## 编译结果 我们评估生成内核的第一个指标是**编译率**：成功编译并执行的内核百分比。我们比较了三种不同的设置：**基线模型**、**SFT**和**GRPO**。在所有三个KernelBench级别上，我们从基线模型到SFT看到了显著改进，从SFT到RL再次看到了改进。 图注插入处 ### 模型关于HIP编译学到了什么 有趣的是，在对生成的内核进行定性分析后，我们发现基线模型经常产生看似语法正确但编译失败的内核，原因在于对HIP内核更深层的误解，例如访问**无效内存**和**错误的API使用**。这些失败在**KernelBench Level 1**任务中尤其常见。由于许多Level 1任务只包含少量操作，模型常常尝试将整个计算重写为自定义HIP内核，而不是保留一些现有的PyTorch运算符。在某些情况下，这些失败出人意料地基本。例如，Level 1问题1-15的内核都包含多余的markdown标记，使得候选在内核编译开始之前就已经无效。 通过SFT，我们发现模型开始从合成训练数据中学习重复出现的HIP实现模式，包括HIP API的使用、有效的内核启动配置和正确的内存访问模式。更重要的是，通过SFT，模型似乎对**应该优化什么**形成了更好的判断。例如，在**Level 2问题2**中，模型保留了计算量大的ConvTranspose运算符，只将偏置加法、裁剪、缩放和除法操作融合到自定义HIP内核中。类似地，**Level 2问题7**将ReLU、LeakyReLU、GELU、Sigmoid和偏置加法融合到一个pass中。这两个例子都遵循相同的策略：保留核心计算，仅优化开销较低的周边操作。 **Level 2问题2**—— ConvTranspose2d_BiasAdd_Clamp_Scaling_Clamp_Divide `import os import torch import torch.nn as nn import torch.nn.functional as F from torch.utils.cpp_extension import load_inline # Must include this line so PyTorch could use HIP compiler for AMD GPUs os.environ["CXX"] = "hipcc" conv_transpose_bias_add_clamp_scaling_clamp_divide_cpp_source = """ #include #include #include // For fminf, fmaxf __global__ void conv_transpose_bias_add_clamp_scaling_clamp_divide_kernel( const float* conv_out_ptr, const float* bias_ptr, float* out_ptr, int N, int C, int H, int W, int scaling_factor) { int total_elements = N * C * H * W; int idx = blockIdx.x * blockDim.x + threadIdx.x; if (idx < total_elements) { // Calculate the channel index (c) from the flattened index (idx) // For a tensor (N, C, H, W) flattened row-major, the channel index 'c' for element 'idx' is: // idx = n * (C*H*W) + c * (H*W) + h * W + w // So, c = (idx / (H * W)) % C; int c_idx = (idx / (H * W)) % C; float val = conv_out_ptr[idx] + bias_ptr[c_idx]; // First clamp (min=0.0, max=1.0) val = fmaxf(0.0f, fminf(1.0f, val)); // Scale by scaling_factor val = val * scaling_factor; // Second clamp (min=0.0, max=1.0) val = fmaxf(0.0f, fminf(1.0f, val)); // Divide by scaling_factor val = val / scaling_factor; out_ptr[idx] = val; } } torch::Tensor conv_transpose_bias_add_clamp_scaling_clamp_divide_hip( torch::Tensor conv_out, torch::Tensor bias, float scaling_factor) { // Ensure tensors are on the same device and are contiguous // Contiguous memory layout is crucial for correct pointer arithmetic in the kernel conv_out = conv_out.contiguous(); bias = bias.contiguous(); // Get dimensions of the output from conv_transpose int N = conv_out.size(0); int C = conv_out.size(1); int H = conv_out.size(2); int W = conv_out.size(3); // Output tensor will have the same shape as conv_out auto out = torch::empty_like(conv_out); const int block_size = 256; const int total_elements = N * C * H * W; const int num_blocks = (total_elements + block_size - 1) / block_size; hipLaunchKernelGGL(conv_transpose_bias_add_clamp_scaling_clamp_divide_kernel, dim3(num_blocks), dim3(block_size), 0, 0, conv_out.data_ptr<float>(), bias.data_ptr<float>(), out.data_ptr<float>(), N, C, H, W, static_cast<int>(scaling_factor)); return out; } """ conv_transpose_bias_add_clamp_scaling_clamp_divide = load_inline( name="conv_transpose_bias_add_clamp_scaling_clamp_divide", cpp_sources=conv_transpose_bias_add_clamp_scaling_clamp_divide_cpp_source, functions=["conv_transpose_bias_add_clamp_scaling_clamp_divide_hip"], verbose=True, ) class ModelNew(nn.Module): def __init__(self, in_channels, out_channels, kernel_size, stride, padding, output_padding, bias_shape, scaling_factor): super(ModelNew, self).__init__() self.conv_transpose = nn.ConvTranspose2d(in_channels, out_channels, kernel_size, stride=stride, padding=padding, output_padding=output_padding) self.bias = nn.Parameter(torch.randn(bias_shape)) self.scaling_factor = scaling_factor self.fused_op = conv_transpose_bias_add_clamp_scaling_clamp_divide def forward(self, x): x = self.conv_transpose(x) return self.fused_op.conv_transpose_bias_add_clamp_scaling_clamp_divide_hip(x, self.bias, self.scaling_factor) ` **Level 2问题7**—— Conv3d_ReLU_LeakyReLU_GELU_Sigmoid_BiasAdd `import os import torch import torch.nn as nn import torch.nn.functional as F from torch.utils.cpp_extension import load_inline # Must include this line so PyTorch could use HIP compiler for AMD GPUs os.environ["CXX"] = "hipcc" conv3d_relu_leakyrelu_gelu_sigmoid_biasadd_cpp_source = """ #include #include // For expf, erff, sqrtf // Fused kernel for ReLU, LeakyReLU, GELU, Sigmoid, and BiasAdd __global__ void fused_activations_biasadd_kernel( const float* conv_out_ptr, const float* bias_ptr, float* out_ptr, int N, int C, int D, int H, int W, float leaky_relu_slope) { int total_elements = N * C * D * H * W; int idx = blockIdx.x * blockDim.x + threadIdx.x; if (idx < total_elements) { float val = conv_out_ptr[idx]; // 1. ReLU: max(0, val) val = (val > 0) ? val : 0; // 2. LeakyReLU: val if val > 0 else val * negative_slope val = (val > 0) ? val : val * leaky_relu_slope; // 3. GELU (using approximation with erff, matching PyTorch's default 'tanh' approximation if not specified) // PyTorch's default GELU is 0.5 * x * (1 + erf(x / sqrt(2))) val = 0.5f * val * (1.0f + erff(val * 0.70710678118f)); // 0.70710678118f is 1/sqrt(2) // 4. Sigmoid: 1 / (1 + exp(-val)) val = 1.0f / (1.0f + expf(-val)); // 5. BiasAdd: val + bias[c] int c_idx = (idx / (D * H * W)) % C; // Calculate channel index for bias addition val = val + bias_ptr[c_idx]; out_ptr[idx] = val; } } torch::Tensor fused_activations_biasadd_hip( torch::Tensor conv_out, torch::Tensor bias, float leaky_relu_slope) { // Ensure tensors are on the same device and are contiguous TORCH_CHECK(conv_out.is_cuda(), "Input tensor must be on CUDA (HIP) device."); TORCH_CHECK(bias.is_cuda(), "Bias tensor must be on CUDA (HIP) device."); TORCH_CHECK(conv_out.is_contiguous(), "Input tensor must be contiguous."); TORCH_CHECK(bias.is_contiguous(), "Bias tensor must be contiguous."); // Get dimensions from the input tensor int N = conv_out.size(0); int C = conv_out.size(1); int D = conv_out.size(2); int H = conv_out.size(3); int W = conv_out.size(4); // Output tensor will have the same shape as conv_out auto out = torch::empty_like(conv_out); const int block_size = 256; const int total_elements = N * C * D * H * W; const int num_blocks = (total_elements + block_size - 1) / block_size; fused_activations_biasadd_kernel<<<num_blocks, block_size>>>( conv_out.data_ptr<float>(), bias.data_ptr<float>(), out.data_ptr<float>(), N, C, D, H, W, leaky_relu_slope); return out; } """ fused_ops = load_inline( name="fused_activations_biasadd", cpp_sources=conv3d_relu_leakyrelu_gelu_sigmoid_biasadd_cpp_source, functions=["fused_activations_biasadd_hip"], verbose=True, ) class ModelNew(nn.Module): def __init__(self, in_channels, out_channels, kernel_size, bias_shape): super(ModelNew, self).__init__() self.conv = nn.Conv3d(in_channels, out_channels, kernel_size).cuda() self.bias = nn.Parameter(torch.randn(bias_shape)).cuda() self.fused_ops = fused_ops def forward(self, x): x = self.conv(x) # Call the fused kernel for all activation functions and bias addition return self.fused_ops.fused_activations_biasadd_hip(x, self.bias, 0.01) ` SFT生成中最弱的案例则尝试了相反的方法。在**Lev**