import torch import torch.nn as nn import torch.nn.functional as F from torch.utils.cpp_extension import load_inline import math cuda_source = ''' #include #include #include __device__ inline float new_gelu(float x) { return 0.5f * x * (1.0f + tanhf(0.7978845608f * (x + 0.044715f * x * x * x))); } __device__ inline void block_reduce(float& sum, float& sumsq, float* s_sum, float* s_sumsq) { int tid = threadIdx.x; int lane = tid % 32; int wid = tid / 32; int num_warps = blockDim.x / 32; for (int offset = 16; offset > 0; offset /= 2) { sum += __shfl_down_sync(0xffffffff, sum, offset); sumsq += __shfl_down_sync(0xffffffff, sumsq, offset); } if (lane == 0) { s_sum[wid] = sum; s_sumsq[wid] = sumsq; } __syncthreads(); if (wid == 0) { sum = (tid < num_warps) ? s_sum[tid] : 0.0f; sumsq = (tid < num_warps) ? s_sumsq[tid] : 0.0f; for (int offset = 16; offset > 0; offset /= 2) { sum += __shfl_down_sync(0xffffffff, sum, offset); sumsq += __shfl_down_sync(0xffffffff, sumsq, offset); } if (tid == 0) { s_sum[0] = sum; s_sumsq[0] = sumsq; } } __syncthreads(); sum = s_sum[0]; sumsq = s_sumsq[0]; } __global__ void fused_ln_kernel(const float4* x, const float4* gamma, const float4* beta, float4* out, int rows, int cols_4, float eps) { int row = blockIdx.x; int tid = threadIdx.x; __shared__ float s_sum_mem[32]; __shared__ float s_sumsq_mem[32]; float sum = 0.0f; float sumsq = 0.0f; for (int i = tid; i < cols_4; i += blockDim.x) { float4 vx = x[row * cols_4 + i]; sum += vx.x + vx.y + vx.z + vx.w; sumsq += vx.x*vx.x + vx.y*vx.y + vx.z*vx.z + vx.w*vx.w; } block_reduce(sum, sumsq, s_sum_mem, s_sumsq_mem); float mean = sum / (cols_4 * 4); float var = fmaxf(0.0f, (sumsq / (cols_4 * 4)) - (mean * mean)); float rstd = rsqrtf(var + eps); for (int i = tid; i < cols_4; i += blockDim.x) { float4 vx = x[row * cols_4 + i]; float4 vg = gamma[i]; float4 vb = beta[i]; float4 vo; vo.x = ((vx.x - mean) * rstd) * vg.x + vb.x; vo.y = ((vx.y - mean) * rstd) * vg.y + vb.y; vo.z = ((vx.z - mean) * rstd) * vg.z + vb.z; vo.w = ((vx.w - mean) * rstd) * vg.w + vb.w; out[row * cols_4 + i] = vo; } } __global__ void fused_add_bias_ln_kernel(const float4* x, const float4* y, const float4* bias, const float4* gamma, const float4* beta, float4* out_res, float4* out_norm, int rows, int cols_4, float eps) { int row = blockIdx.x; int tid = threadIdx.x; __shared__ float s_sum_mem[32]; __shared__ float s_sumsq_mem[32]; float sum = 0.0f; float sumsq = 0.0f; for (int i = tid; i < cols_4; i += blockDim.x) { float4 vx = x[row * cols_4 + i]; float4 vy = y[row * cols_4 + i]; float4 vb_lin = bias[i]; float4 vr; vr.x = vx.x + vy.x + vb_lin.x; vr.y = vx.y + vy.y + vb_lin.y; vr.z = vx.z + vy.z + vb_lin.z; vr.w = vx.w + vy.w + vb_lin.w; out_res[row * cols_4 + i] = vr; sum += vr.x + vr.y + vr.z + vr.w; sumsq += vr.x*vr.x + vr.y*vr.y + vr.z*vr.z + vr.w*vr.w; } block_reduce(sum, sumsq, s_sum_mem, s_sumsq_mem); float mean = sum / (cols_4 * 4); float var = fmaxf(0.0f, (sumsq / (cols_4 * 4)) - (mean * mean)); float rstd = rsqrtf(var + eps); for (int i = tid; i < cols_4; i += blockDim.x) { float4 vr = out_res[row * cols_4 + i]; float4 vg = gamma[i]; float4 vb_ln = beta[i]; float4 vo; vo.x = ((vr.x - mean) * rstd) * vg.x + vb_ln.x; vo.y = ((vr.y - mean) * rstd) * vg.y + vb_ln.y; vo.z = ((vr.z - mean) * rstd) * vg.z + vb_ln.z; vo.w = ((vr.w - mean) * rstd) * vg.w + vb_ln.w; out_norm[row * cols_4 + i] = vo; } } __global__ void fused_bias_gelu_kernel(const float4* x, const float4* bias, float4* out, int total_size_4, int cols_4) { int idx = blockIdx.x * blockDim.x + threadIdx.x; if (idx < total_size_4) { float4 vx = x[idx]; float4 vb = bias[idx % cols_4]; float4 vo; vo.x = new_gelu(vx.x + vb.x); vo.y = new_gelu(vx.y + vb.y); vo.z = new_gelu(vx.z + vb.z); vo.w = new_gelu(vx.w + vb.w); out[idx] = vo; } } __global__ void fused_add_bias_kernel(const float4* x, const float4* y, const float4* bias, float4* out, int total_size_4, int cols_4) { int idx = blockIdx.x * blockDim.x + threadIdx.x; if (idx < total_size_4) { float4 vx = x[idx]; float4 vy = y[idx]; float4 vb = bias[idx % cols_4]; float4 vo; vo.x = vx.x + vy.x + vb.x; vo.y = vx.y + vy.y + vb.y; vo.z = vx.z + vy.z + vb.z; vo.w = vx.w + vy.w + vb.w; out[idx] = vo; } } torch::Tensor fused_ln_cuda(torch::Tensor x, torch::Tensor gamma, torch::Tensor beta, float eps) { auto out = torch::empty_like(x); int rows = x.view({-1, x.size(-1)}).size(0); int cols_4 = x.size(-1) / 4; fused_ln_kernel<<>>((const float4*)x.data_ptr(), (const float4*)gamma.data_ptr(), (const float4*)beta.data_ptr(), (float4*)out.data_ptr(), rows, cols_4, eps); return out; } std::vector fused_add_bias_ln_cuda(torch::Tensor x, torch::Tensor y, torch::Tensor bias, torch::Tensor gamma, torch::Tensor beta, float eps) { auto out_res = torch::empty_like(x); auto out_norm = torch::empty_like(x); int rows = x.view({-1, x.size(-1)}).size(0); int cols_4 = x.size(-1) / 4; fused_add_bias_ln_kernel<<>>((const float4*)x.data_ptr(), (const float4*)y.data_ptr(), (const float4*)bias.data_ptr(), (const float4*)gamma.data_ptr(), (const float4*)beta.data_ptr(), (float4*)out_res.data_ptr(), (float4*)out_norm.data_ptr(), rows, cols_4, eps); return {out_res, out_norm}; } torch::Tensor fused_bias_gelu_cuda(torch::Tensor x, torch::Tensor bias) { auto out = torch::empty_like(x); int total_size_4 = x.numel() / 4; int cols_4 = x.size(-1) / 4; fused_bias_gelu_kernel<<<(total_size_4 + 255) / 256, 256>>>((const float4*)x.data_ptr(), (const float4*)bias.data_ptr(), (float4*)out.data_ptr(), total_size_4, cols_4); return out; } torch::Tensor fused_add_bias_cuda(torch::Tensor x, torch::Tensor y, torch::Tensor bias) { auto out = torch::empty_like(x); int total_size_4 = x.numel() / 4; int cols_4 = x.size(-1) / 4; fused_add_bias_kernel<<<(total_size_4 + 255) / 256, 256>>>((const float4*)x.data_ptr(), (const float4*)y.data_ptr(), (const float4*)bias.data_ptr(), (float4*)out.data_ptr(), total_size_4, cols_4); return out; } ''' cpp_source = ''' torch::Tensor fused_ln_cuda(torch::Tensor x, torch::Tensor gamma, torch::Tensor beta, float eps); std::vector fused_add_bias_ln_cuda(torch::Tensor x, torch::Tensor y, torch::Tensor bias, torch::Tensor gamma, torch::Tensor beta, float eps); torch::Tensor fused_bias_gelu_cuda(torch::Tensor x, torch::Tensor bias); torch::Tensor fused_add_bias_cuda(torch::Tensor x, torch::Tensor y, torch::Tensor bias); ''' fused_ops = load_inline( name="fused_ops_v2", cpp_sources=cpp_source, cuda_sources=cuda_source, functions=["fused_ln_cuda", "fused_add_bias_ln_cuda", "fused_bias_gelu_cuda", "fused_add_bias_cuda"], extra_cflags=['-O3'] ) class Model(nn.Module): def __init__(self, n_embd, n_head, attn_pdrop, resid_pdrop, max_seqlen): super().__init__() self.ln_1 = nn.LayerNorm(n_embd) assert n_embd % n_head == 0 self.attn = nn.Module() self.attn.c_attn = nn.Linear(n_embd, 3 * n_embd) self.attn.c_proj = nn.Linear(n_embd, n_embd) self.attn.attn_dropout = nn.Dropout(attn_pdrop) self.attn.resid_dropout = nn.Dropout(resid_pdrop) self.attn.n_head = n_head self.attn.n_embd = n_embd self.ln_2 = nn.LayerNorm(n_embd) self.mlp = nn.ModuleDict(dict( c_fc = nn.Linear(n_embd, 4 * n_embd), c_proj = nn.Linear(4 * n_embd, n_embd), act = nn.GELU(), dropout = nn.Dropout(resid_pdrop), )) self.fused_ops = fused_ops def forward(self, x): B, T, C = x.size() nh = self.attn.n_head ln1_out = self.fused_ops.fused_ln_cuda(x, self.ln_1.weight, self.ln_1.bias, self.ln_1.eps) qkv = self.attn.c_attn(ln1_out) qkv = qkv.view(B, T, 3, nh, C // nh).permute(2, 0, 3, 1, 4) q, k, v = qkv[0], qkv[1], qkv[2] y = F.scaled_dot_product_attention(q, k, v, is_causal=True, dropout_p=self.attn.attn_dropout.p if self.training else 0.0) y = y.transpose(1, 2).contiguous().view(B, T, C) attn_out_no_bias = F.linear(y, self.attn.c_proj.weight, None) x, ln2_out = self.fused_ops.fused_add_bias_ln_cuda(x, attn_out_no_bias, self.attn.c_proj.bias, self.ln_2.weight, self.ln_2.bias, self.ln_2.eps) mlp_fc_no_bias = F.linear(ln2_out, self.mlp.c_fc.weight, None) mlp_act = self.fused_ops.fused_bias_gelu_cuda(mlp_fc_no_bias, self.mlp.c_fc.bias) mlp_proj_no_bias = F.linear(mlp_act, self.mlp.c_proj.weight, None) x = self.fused_ops.fused_add_bias_cuda(x, mlp_proj_no_bias, self.mlp.c_proj.bias) return x

Transformer-block Speedup

Step Code View