import torch import torch.nn as nn from torch.utils.cpp_extension import load_inline gelu_source = ''' #include #include __global__ void __launch_bounds__(256) gelu_kernel(const float4* __restrict__ input, float4* __restrict__ output, int n_vec) { const float k1 = 0.7978845608f; const float k2 = 0.0356774081f; const float k0 = 0.5f; int i = blockIdx.x * blockDim.x + threadIdx.x; int stride = blockDim.x * gridDim.x; for (; i < n_vec; i += stride) { float4 in = __ldg(input + i); float x1 = in.x; float x2 = in.y; float x3 = in.z; float x4 = in.w; float a1 = fmaf(k2, x1 * x1, k1) * x1; float a2 = fmaf(k2, x2 * x2, k1) * x2; float a3 = fmaf(k2, x3 * x3, k1) * x3; float a4 = fmaf(k2, x4 * x4, k1) * x4; float h1 = k0 * x1; float h2 = k0 * x2; float h3 = k0 * x3; float h4 = k0 * x4; float t1 = tanhf(a1); float t2 = tanhf(a2); float t3 = tanhf(a3); float t4 = tanhf(a4); float4 out_v; out_v.x = fmaf(h1, t1, h1); out_v.y = fmaf(h2, t2, h2); out_v.z = fmaf(h3, t3, h3); out_v.w = fmaf(h4, t4, h4); output[i] = out_v; } } __global__ void gelu_tail_kernel(const float* __restrict__ input, float* __restrict__ output, int start, int end) { const float k1 = 0.7978845608f; const float k2 = 0.0356774081f; const float k0 = 0.5f; int i = start + blockIdx.x * blockDim.x + threadIdx.x; if (i < end) { float x = input[i]; float arg = fmaf(k2, x * x, k1) * x; float t = tanhf(arg); float h = k0 * x; output[i] = fmaf(h, t, h); } } torch::Tensor gelu_cuda(torch::Tensor a) { auto size = a.numel(); auto out = torch::empty_like(a); int n_vec = size / 4; if (n_vec > 0) { const int block_size = 256; const int grid_size = std::min(4096, (n_vec + block_size - 1) / block_size); gelu_kernel<<>>( reinterpret_cast(a.data_ptr()), reinterpret_cast(out.data_ptr()), n_vec ); } int tail_start = n_vec * 4; if (tail_start < size) { gelu_tail_kernel<<<1, 256>>>(a.data_ptr(), out.data_ptr(), tail_start, size); } return out; } ''' gelu_cpp_source = "torch::Tensor gelu_cuda(torch::Tensor a);" gelu_optimized_lib = load_inline( name="gelu_optimized_lib", cpp_sources=gelu_cpp_source, cuda_sources=gelu_source, functions=["gelu_cuda"], verbose=False, extra_cuda_cflags=["--use_fast_math", "-O3", "-Xptxas=-O3"], extra_cflags=["-O3"] ) class Model(nn.Module): def __init__(self): super(Model, self).__init__() self.gelu_lib = gelu_optimized_lib def forward(self, x): if not x.is_contiguous(): x = x.contiguous() return self.gelu_lib.gelu_cuda(x)

GELU-kernel speedup

Step Code View