自定义CUDA算子融合实现模型推理加速

原创

已于 2025-03-31 00:56:59 修改 · 1k 阅读

8 ·

CC 4.0 BY-SA版权

文章标签：

#CUDA

于 2025-03-31 00:55:10 首次发布

对模型进行推理加速的最常用方法就是算子融合，这里用个简单demo记录下：

总共有如下三个步骤：

导出模型权重

用pytorch定义一个多层DNN模型，然后导出其各层的网络参数。

# export_model.py

import torch
import torch.nn as nn

# 定义PyTorch模型结构
class DNNModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.layer1 = nn.Linear(100, 128)
        self.relu1 = nn.ReLU()
        self.layer2 = nn.Linear(128, 64)
        self.relu2 = nn.ReLU()
        self.layer3 = nn.Linear(64, 10)

    def forward(self, x):
        x = self.relu1(self.layer1(x))
        x = self.relu2(self.layer2(x))
        return self.layer3(x)


# 创建模型并随机初始化
model = DNNModel()
model.eval()

# 导出权重为二进制文件
for name, param in model.named_parameters():
    param.detach().cpu().numpy().tofile(f"{
     
     name}.bin")

print("权重导出完成！")

运行 python export_model.py 。

编写CUDA融合算子

神经网络的每一层前向传播，都先从全局内存中读取tensor到寄存器内存，完成计算后再写回到全局内存，IO次数较多。利用算子融合，将多次计算融合成一次计算，减少IO读写，从而实现模型推理加速。

具体的代码包括如下三个步骤：

加载pytorch导出的模型参数
将多次前向传播融合到一个函数中
将优化后的函数绑定到python模块中

// fused_op.cu

#include <torch/extension.h>
#include <cuda_runtime.h>

__global__ void fused_forward(
    const float* input,
    float* output,
    const float* W1, const float* b1,
    const float* W2, const float* b2,
    const float* W3, const float* b3,
    int batch_size, int in_dim, int hid1, int hid2, int out_dim
) {
   
   
    // 每个线程处理一个样本的完整计算
    int sample_idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (sample_idx >= batch_size) return;

    // 指向当前样本的输入和输出
    const float* x = input + sample_idx