forked from ccf-ai-infra/GPUCodeForces
122 lines
4.5 KiB
Python
122 lines
4.5 KiB
Python
import torch
|
||
import torch.nn as nn
|
||
from torch.utils.cpp_extension import load_inline
|
||
|
||
|
||
class ModelNew(nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self._compile_cuda_kernel()
|
||
|
||
def _compile_cuda_kernel(self):
|
||
cpp_source = """
|
||
#include <torch/extension.h>
|
||
torch::Tensor reglu_vectorized_parallel(torch::Tensor input);
|
||
"""
|
||
|
||
cuda_source = """
|
||
#include <cuda_runtime.h>
|
||
|
||
// 使用简单的向量化方法
|
||
__global__ void reglu_vectorized_kernel(
|
||
const float* __restrict__ input,
|
||
float* __restrict__ output,
|
||
int feature_dim, int total_elements) {
|
||
|
||
const int tid = threadIdx.x + blockIdx.x * blockDim.x;
|
||
const int stride = blockDim.x * gridDim.x;
|
||
|
||
// 每个线程处理4个元素（向量化）
|
||
const int elements_per_thread = 4;
|
||
const int vectorized_elements = total_elements / elements_per_thread;
|
||
|
||
// 处理向量化部分
|
||
for (int i = tid; i < vectorized_elements; i += stride) {
|
||
int base_idx = i * elements_per_thread;
|
||
int row = base_idx / (feature_dim / 2);
|
||
int base_col = base_idx % (feature_dim / 2);
|
||
|
||
#pragma unroll
|
||
for (int j = 0; j < elements_per_thread; j++) {
|
||
int col = base_col + j;
|
||
if (col < feature_dim / 2) {
|
||
int global_idx = base_idx + j;
|
||
int gate_offset = row * feature_dim + col;
|
||
int act_offset = gate_offset + (feature_dim / 2);
|
||
|
||
float gate_val = input[gate_offset];
|
||
float act_val = input[act_offset];
|
||
output[global_idx] = fmaxf(0.0f, gate_val) * act_val;
|
||
}
|
||
}
|
||
}
|
||
|
||
// 处理剩余元素
|
||
int remaining_start = vectorized_elements * elements_per_thread;
|
||
for (int i = remaining_start + tid; i < total_elements; i += stride) {
|
||
int row = i / (feature_dim / 2);
|
||
int col = i % (feature_dim / 2);
|
||
|
||
float gate_val = input[row * feature_dim + col];
|
||
float act_val = input[row * feature_dim + col + (feature_dim / 2)];
|
||
output[i] = fmaxf(0.0f, gate_val) * act_val;
|
||
}
|
||
}
|
||
|
||
// 更稳定的版本 - 不使用向量化
|
||
__global__ void reglu_simple_kernel(
|
||
const float* __restrict__ input,
|
||
float* __restrict__ output,
|
||
int feature_dim, int total_elements) {
|
||
|
||
const int tid = threadIdx.x + blockIdx.x * blockDim.x;
|
||
const int stride = blockDim.x * gridDim.x;
|
||
|
||
for (int i = tid; i < total_elements; i += stride) {
|
||
int row = i / (feature_dim / 2);
|
||
int col = i % (feature_dim / 2);
|
||
|
||
float gate_val = input[row * feature_dim + col];
|
||
float act_val = input[row * feature_dim + col + (feature_dim / 2)];
|
||
|
||
// 使用fmaxf代替条件判断，性能更好
|
||
output[i] = fmaxf(0.0f, gate_val) * act_val;
|
||
}
|
||
}
|
||
|
||
torch::Tensor reglu_vectorized_parallel(torch::Tensor input) {
|
||
input = input.contiguous();
|
||
auto sizes = input.sizes().vec();
|
||
int feature_dim = sizes.back();
|
||
sizes.back() /= 2;
|
||
auto output = torch::empty(sizes, input.options());
|
||
|
||
int total_elements = output.numel();
|
||
int threads = 256;
|
||
int blocks = min((total_elements + threads - 1) / threads, 128); // 限制最大blocks
|
||
|
||
// 使用简单稳定的内核
|
||
reglu_simple_kernel<<<blocks, threads>>>(
|
||
input.data_ptr<float>(), output.data_ptr<float>(),
|
||
feature_dim, total_elements);
|
||
|
||
cudaError_t err = cudaGetLastError();
|
||
if (err != cudaSuccess) {
|
||
AT_ERROR("CUDA error in reglu_vectorized_parallel: ", cudaGetErrorString(err));
|
||
}
|
||
|
||
return output;
|
||
}
|
||
"""
|
||
|
||
self.op = load_inline(
|
||
name="reglu_vectorized_fixed",
|
||
cpp_sources=cpp_source,
|
||
cuda_sources=cuda_source,
|
||
functions=["reglu_vectorized_parallel"],
|
||
extra_cuda_cflags=["-O3", "--use_fast_math"],
|
||
verbose=True
|
||
)
|
||
|
||
def forward(self, x):
|
||
return self.op.reglu_vectorized_parallel(x) |