forked from ccf-ai-infra/GPUCodeForces
119 lines
2.8 KiB
Python
119 lines
2.8 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
from torch.utils.cpp_extension import load_inline
|
|
|
|
cuda_source = """
|
|
#include <torch/extension.h>
|
|
#include <cuda_runtime.h>
|
|
|
|
__inline__ __device__ float warpReduceSum(float val) {
|
|
for (int offset = 16; offset > 0; offset /= 2)
|
|
val += __shfl_down_sync(0xffffffff, val, offset);
|
|
return val;
|
|
}
|
|
|
|
__inline__ __device__ float blockReduceSum(float val) {
|
|
static __shared__ float shared[32];
|
|
int lane = threadIdx.x % 32;
|
|
int wid = threadIdx.x / 32;
|
|
|
|
val = warpReduceSum(val);
|
|
|
|
if (lane == 0) shared[wid] = val;
|
|
__syncthreads();
|
|
|
|
val = (threadIdx.x < blockDim.x / 32) ? shared[lane] : 0.0f;
|
|
|
|
if (wid == 0) val = warpReduceSum(val);
|
|
|
|
return val;
|
|
}
|
|
|
|
__global__ void gate_blend_normalize_kernel(
|
|
const float* __restrict__ a,
|
|
const float* __restrict__ b,
|
|
const float* __restrict__ gate,
|
|
float* __restrict__ output,
|
|
int rows,
|
|
int cols,
|
|
float eps
|
|
) {
|
|
int bid = blockIdx.x;
|
|
int tid = threadIdx.x;
|
|
|
|
if (bid >= rows) return;
|
|
|
|
const float* a_row = a + bid * cols;
|
|
const float* b_row = b + bid * cols;
|
|
const float* gate_row = gate + bid * cols;
|
|
float* out_row = output + bid * cols;
|
|
|
|
float sum_sq = 0.0f;
|
|
for (int i = tid; i < cols; i += blockDim.x) {
|
|
float g = gate_row[i];
|
|
float blended = g * a_row[i] + (1.0f - g) * b_row[i];
|
|
sum_sq += blended * blended;
|
|
}
|
|
|
|
sum_sq = blockReduceSum(sum_sq);
|
|
|
|
__shared__ float inv_norm;
|
|
if (tid == 0) {
|
|
inv_norm = rsqrtf(sum_sq + eps);
|
|
}
|
|
__syncthreads();
|
|
|
|
float norm_factor = inv_norm;
|
|
|
|
for (int i = tid; i < cols; i += blockDim.x) {
|
|
float g = gate_row[i];
|
|
float blended = g * a_row[i] + (1.0f - g) * b_row[i];
|
|
out_row[i] = blended * norm_factor;
|
|
}
|
|
}
|
|
|
|
torch::Tensor gate_blend_normalize_cuda(torch::Tensor a, torch::Tensor b, torch::Tensor gate) {
|
|
auto output = torch::empty_like(a);
|
|
|
|
int cols = a.size(a.dim() - 1);
|
|
int rows = a.numel() / cols;
|
|
|
|
int block_size = 256;
|
|
while (block_size < cols && block_size < 1024) {
|
|
block_size *= 2;
|
|
}
|
|
|
|
gate_blend_normalize_kernel<<<rows, block_size>>>(
|
|
a.data_ptr<float>(),
|
|
b.data_ptr<float>(),
|
|
gate.data_ptr<float>(),
|
|
output.data_ptr<float>(),
|
|
rows,
|
|
cols,
|
|
1e-12f
|
|
);
|
|
|
|
return output;
|
|
}
|
|
"""
|
|
|
|
cpp_source = """
|
|
torch::Tensor gate_blend_normalize_cuda(torch::Tensor a, torch::Tensor b, torch::Tensor gate);
|
|
"""
|
|
|
|
module = load_inline(
|
|
name="gate_blend_normalize",
|
|
cpp_sources=cpp_source,
|
|
cuda_sources=cuda_source,
|
|
functions=["gate_blend_normalize_cuda"],
|
|
verbose=True
|
|
)
|
|
|
|
|
|
class ModelNew(nn.Module):
|
|
def __init__(self):
|
|
super(ModelNew, self).__init__()
|
|
self.module = module
|
|
|
|
def forward(self, a, b, gate):
|
|
return self.module.gate_blend_normalize_cuda(a, b, gate) |