call → compile_kernel
callcompile_kernel
{
"kernel_code": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.cpp_extension import load_inline\n\ncuda_src = r'''\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n#include <vector>\n\n#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x \" must be a CUDA tensor\")\n#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x \" must be contiguous\")\n\n__inline__ __device__ float warp_reduce_sum(float val) {\n for (int offset = 16; offset > 0; offset >>= 1)\n val += __shfl_down_sync(0xffffffff, val, offset);\n return val;\n}\n\n__inline__ __device__ float block_reduce_sum(float val) {\n static __shared__ float shared[32];\n int lane = threadIdx.x & 31;\n int wid = threadIdx.x >> 5;\n val = warp_reduce_sum(val);\n if (lane == 0) shared[wid] = val;\n __syncthreads();\n val = (threadIdx.x < (blockDim.x >> 5)) ? shared[lane] : 0.0f;\n if (wid == 0) val = warp_reduce_sum(val);\n return val;\n}\n\n__global__ void layernorm_kernel(const float* __restrict__ x, const float* __restrict__ gamma,\n const float* __restrict__ beta, float* __restrict__ out,\n int rows, int C) {\n int row = blockIdx.x;\n if (row >= rows) return;\n const float* xr = x + (size_t)row * C;\n float sum = 0.0f;\n float ss = 0.0f;\n for (int i = threadIdx.x; i < C; i += blockDim.x) {\n float v = xr[i];\n sum += v;\n ss += v * v;\n }\n sum = block_reduce_sum(sum);\n ss = block_reduce_sum(ss);\n float mean = sum / (float)C;\n float var = ss / (float)C - mean * mean;\n float inv = rsqrtf(var + 1.0e-5f);\n float* yr = out + (size_t)row * C;\n for (int i = threadIdx.x; i < C; i += blockDim.x) {\n yr[i] = (xr[i] - mean) * inv * gamma[i] + beta[i];\n }\n}\n\n__global__ void add_layernorm_kernel(const float* __restrict__ a, const float* __restrict__ b,\n const float* __restrict__ gamma, const float* __restrict__ beta,\n float* __restrict__ add_out, float* __restrict__ ln_out,\n int rows, int C) {\n int row = blockIdx.x;\n if (row >= rows) return;\n const float* ar = a + (size_t)row * C;\n const float* br = b + (size_t)row * C;\n float sum = 0.0f;\n float ss = 0.0f;\n for (int i = threadIdx.x; i < C; i += blockDim.x) {\n float v = ar[i] + br[i];\n sum += v;\n ss += v * v;\n }\n sum = block_reduce_sum(sum);\n ss = block_reduce_sum(ss);\n float mean = sum / (float)C;\n float var = ss / (float)C - mean * mean;\n float inv = rsqrtf(var + 1.0e-5f);\n float* addr = add_out + (size_t)row * C;\n float* lnr = ln_out + (size_t)row * C;\n for (int i = threadIdx.x; i < C; i += blockDim.x) {\n float v = ar[i] + br[i];\n addr[i] = v;\n lnr[i] = (v - mean) * inv * gamma[i] + beta[i];\n }\n}\n\n__global__ void gelu_kernel(const float* __restrict__ x, float* __restrict__ out, size_t n) {\n size_t idx = (size_t)blockIdx.x * blockDim.x + threadIdx.x;\n size_t stride = (size_t)blockDim.x * gridDim.x;\n const float kAlpha = 0.7978845608028654f;\n const float kBeta = 0.044715f;\n for (size_t i = idx; i < n; i += stride) {\n float v = x[i];\n float u = kAlpha * (v + kBeta * v * v * v);\n out[i] = 0.5f * v * (1.0f + tanhf(u));\n }\n}\n\n__global__ void add_kernel(const float* __restrict__ a, const float* __restrict__ b,\n float* __restrict__ out, size_t n) {\n size_t idx = (size_t)blockIdx.x * blockDim.x + threadIdx.x;\n size_t stride = (size_t)blockDim.x * gridDim.x;\n for (size_t i = idx; i < n; i += stride) out[i] = a[i] + b[i];\n}\n\ntorch::Tensor layernorm_forward(torch::Tensor x, torch::Tensor gamma, torch::Tensor beta) {\n CHECK_CUDA(x); CHECK_CUDA(gamma); CHECK_CUDA(beta);\n auto out = torch::empty_like(x);\n int C = x.size(-1);\n int rows = x.numel() / C;\n layernorm_kernel<<<rows, 256>>>(x.data_ptr<float>(), gamma.data_ptr<float>(), beta.data_ptr<float>(), out.data_ptr<float>(), rows, C);\n return out;\n}\n\nstd::vector<torch::Tensor> add_layernorm_forward(torch::Tensor a, torch::Tensor b, torch::Tensor gamma, torch::Tensor beta) {\n CHECK_CUDA(a); CHECK_CUDA(b); CHECK_CUDA(gamma); CHECK_CUDA(beta);\n auto add_out = torch::empty_like(a);\n auto ln_out = torch::empty_like(a);\n int C = a.size(-1);\n int rows = a.numel() / C;\n add_layernorm_kernel<<<rows, 256>>>(a.data_ptr<float>(), b.data_ptr<float>(), gamma.data_ptr<float>(), beta.data_ptr<float>(), add_out.data_ptr<float>(), ln_out.data_ptr<float>(), rows, C);\n return {add_out, ln_out};\n}\n\ntorch::Tensor gelu_forward(torch::Tensor x) {\n CHECK_CUDA(x);\n auto out = torch::empty_like(x);\n size_t n = x.numel();\n int threads = 256;\n int blocks = (int)min((size_t)65535, (n + threads - 1) / threads);\n gelu_kernel<<<blocks, threads>>>(x.data_ptr<float>(), out.data_ptr<float>(), n);\n return out;\n}\n\ntorch::Tensor add_forward(torch::Tensor a, torch::Tensor b) {\n CHECK_CUDA(a); CHECK_CUDA(b);\n auto out = torch::empty_like(a);\n size_t n = a.numel();\n int threads = 256;\n int blocks = (int)min((size_t)65535, (n + threads - 1) / threads);\n add_kernel<<<blocks, threads>>>(a.data_ptr<float>(), b.data_ptr<float>(), out.data_ptr<float>(), n);\n return out;\n}\n'''\n\n_ext = load_inline(name='transformer_block_ext_v1', cpp_sources='', cuda_sources=cuda_src,\n functions=['layernorm_forward','add_layernorm_forward','gelu_forward','add_forward'],\n extra_cuda_cflags=['-O3'], extra_cflags=['-O3'], verbose=False)\n\nclass NewGELU(nn.Module):\n def __init__(self):\n super().__init__()\n def forward(self, x):\n return _ext.gelu_forward(x.contiguous())\n\nclass CausalSelfAttention(nn.Module):\n def __init__(self, n_embd, n_head, attn_pdrop, resid_pdrop, max_seqlen):\n super().__init__()\n assert n_embd % n_head == 0\n self.c_attn = nn.Linear(n_embd, 3 * n_embd)\n self.c_proj = nn.Linear(n_embd, n_embd)\n self.attn_dropout = nn.Dropout(attn_pdrop)\n self.resid_dropout = nn.Dropout(resid_pdrop)\n self.register_buffer('bias', torch.tril(torch.ones(max_seqlen, max_seqlen)).view(1,1,max_seqlen,max_seqlen))\n self.n_head = n_head\n self.n_embd = n_embd\n def forward(self, x):\n B,T,C = x.shape\n qkv = F.linear(x, self.c_attn.weight, self.c_attn.bias)\n hs = C // self.n_head\n qkv = qkv.view(B, T, 3, self.n_head, hs)\n q = qkv[:,:,0].transpose(1,2)\n k = qkv[:,:,1].transpose(1,2)\n v = qkv[:,:,2].transpose(1,2)\n y = F.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=0.0, is_causal=True)\n y = y.transpose(1,2).contiguous().view(B,T,C)\n y = F.linear(y, self.c_proj.weight, self.c_proj.bias)\n return y\n\nclass ModelNew(nn.Module):\n def __init__(self, n_embd, n_head, attn_pdrop, resid_pdrop, max_seqlen):\n super().__init__()\n self.ln_1 = nn.LayerNorm(n_embd)\n self.attn = CausalSelfAttention(n_embd, n_head, attn_pdrop, resid_pdrop, max_seqlen)\n self.ln_2 = nn.LayerNorm(n_embd)\n self.mlp = nn.ModuleDict(dict(\n c_fc=nn.Linear(n_embd, 4*n_embd),\n c_proj=nn.Linear(4*n_embd, n_embd),\n act=NewGELU(),\n dropout=nn.Dropout(resid_pdrop),\n ))\n self.n_embd = n_embd\n self.n_head = n_head\n def forward(self, x):\n x = x.contiguous()\n ln1 = _ext.layernorm_forward(x, self.ln_1.weight, self.ln_1.bias)\n attn_out = self.attn(ln1)\n pair = _ext.add_layernorm_forward(x, attn_out, self.ln_2.weight, self.ln_2.bias)\n x1 = pair[0]\n ln2 = pair[1]\n h = F.linear(ln2, self.mlp['c_fc'].weight, self.mlp['c_fc'].bias)\n h = _ext.gelu_forward(h)\n h = F.linear(h, self.mlp['c_proj'].weight, self.mlp['c_proj'].bias)\n out = _ext.add_forward(x1, h)\n return out\n"
}