forked from ccf-ai-infra/GPUCodeForces
131 lines
3.7 KiB
Python
131 lines
3.7 KiB
Python
import torch
|
|
import math
|
|
from torch.utils.cpp_extension import load_inline
|
|
|
|
# 简化但高效的Softmax CUDA实现
|
|
softmax_source = """
|
|
#include <torch/extension.h>
|
|
#include <cuda_runtime.h>
|
|
#include <float.h>
|
|
|
|
// 简化但高效的Softmax内核
|
|
__global__ void softmax_kernel(
|
|
const float* __restrict__ input,
|
|
float* __restrict__ output,
|
|
int batch_size,
|
|
int seq_length) {
|
|
|
|
int batch_idx = blockIdx.x;
|
|
int tid = threadIdx.x;
|
|
|
|
// 每个block处理一个batch
|
|
if (batch_idx >= batch_size) return;
|
|
|
|
extern __shared__ float shared_mem[];
|
|
float* shared_max = shared_mem;
|
|
float* shared_sum = &shared_mem[blockDim.x];
|
|
|
|
// 第一步：每个线程计算局部最大值
|
|
float thread_max = -FLT_MAX;
|
|
for (int i = tid; i < seq_length; i += blockDim.x) {
|
|
float val = input[batch_idx * seq_length + i];
|
|
thread_max = fmaxf(thread_max, val);
|
|
}
|
|
|
|
// 块内归约求最大值
|
|
#pragma unroll
|
|
for (int offset = blockDim.x / 2; offset > 0; offset >>= 1) {
|
|
thread_max = fmaxf(thread_max, __shfl_xor_sync(0xFFFFFFFF, thread_max, offset));
|
|
}
|
|
|
|
if (tid == 0) {
|
|
shared_max[0] = thread_max;
|
|
}
|
|
__syncthreads();
|
|
|
|
float row_max = shared_max[0];
|
|
|
|
// 第二步：计算指数和
|
|
float thread_sum = 0.0f;
|
|
for (int i = tid; i < seq_length; i += blockDim.x) {
|
|
float val = input[batch_idx * seq_length + i];
|
|
float exp_val = expf(val - row_max);
|
|
output[batch_idx * seq_length + i] = exp_val;
|
|
thread_sum += exp_val;
|
|
}
|
|
|
|
// 块内归约求和
|
|
#pragma unroll
|
|
for (int offset = blockDim.x / 2; offset > 0; offset >>= 1) {
|
|
thread_sum += __shfl_xor_sync(0xFFFFFFFF, thread_sum, offset);
|
|
}
|
|
|
|
if (tid == 0) {
|
|
shared_sum[0] = thread_sum;
|
|
}
|
|
__syncthreads();
|
|
|
|
float row_sum = shared_sum[0];
|
|
|
|
// 第三步：归一化
|
|
for (int i = tid; i < seq_length; i += blockDim.x) {
|
|
output[batch_idx * seq_length + i] /= row_sum;
|
|
}
|
|
}"
|
|
|
|
torch::Tensor softmax_cuda_forward(torch::Tensor input, int dim) {
|
|
// 检查输入维度
|
|
if (input.dim() != 2) {
|
|
throw std::runtime_error("Softmax currently only supports 2D tensors");
|
|
}
|
|
|
|
int batch_size = input.size(0);
|
|
int seq_length = input.size(1);
|
|
|
|
// 使用简化的内核
|
|
dim3 blocks(batch_size);
|
|
dim3 threads(256); // 优化块大小
|
|
size_t shared_mem_size = threads.x * 2 * sizeof(float);
|
|
|
|
auto output = torch::zeros_like(input);
|
|
|
|
softmax_kernel<<<blocks, threads, shared_mem_size>>>(
|
|
input.data_ptr<float>(),
|
|
output.data_ptr<float>(),
|
|
batch_size,
|
|
seq_length
|
|
);
|
|
|
|
return output;
|
|
}
|
|
"""
|
|
|
|
softmax_cpp_source = """
|
|
torch::Tensor softmax_cuda_forward(torch::Tensor input, int dim);
|
|
"""
|
|
|
|
# 编译内联CUDA代码
|
|
softmax_cuda = None
|
|
if torch.cuda.is_available() and torch.version.cuda is not None:
|
|
try:
|
|
softmax_cuda = load_inline(
|
|
name="softmax_cuda",
|
|
cpp_sources=softmax_cpp_source,
|
|
cuda_sources=softmax_source,
|
|
functions=["softmax_cuda_forward"],
|
|
verbose=True
|
|
)
|
|
except Exception as e:
|
|
print(f"[警告] Softmax CUDA 扩展构建失败，将回退到 PyTorch 实现。错误: {e}")
|
|
|
|
class ModelNew(torch.nn.Module):
|
|
def __init__(self, dim=-1):
|
|
super(ModelNew, self).__init__()
|
|
self.dim = dim
|
|
|
|
def forward(self, x):
|
|
# 若扩展可用则调用 CUDA 内核，否则回退到 PyTorch softmax
|
|
if softmax_cuda is not None:
|
|
return softmax_cuda.softmax_cuda_forward(x, self.dim)
|
|
else:
|
|
return torch.softmax(x, self.dim) |