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\nimport math\n\n# CUDA source for Mamba block with optimized memory access\ncuda_source = \"\"\"\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n#include <vector>\n#include <cmath>\n\n#define MAX_D_STATE 64\n\n// RMSNorm kernel - one warp per row for better coalescing\n__global__ void rmsnorm_kernel(\n const float* __restrict__ input,\n const float* __restrict__ weight,\n float* __restrict__ output,\n int B, int L, int D,\n float eps\n) {\n int bl = blockIdx.x;\n int d = threadIdx.x;\n \n if (bl >= B * L || d >= D) return;\n \n int base = bl * D;\n \n // Parallel reduction for sum of squares\n float val = input[base + d];\n float sum_sq = val * val;\n \n __shared__ float shared_sq[1024];\n shared_sq[d] = sum_sq;\n __syncthreads();\n \n for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {\n if (d < stride) {\n shared_sq[d] += shared_sq[d + stride];\n }\n __syncthreads();\n }\n \n float rms = rsqrtf(shared_sq[0] / D + eps);\n output[base + d] = val * rms * weight[d];\n}\n\n// Optimized linear kernel using shared memory and coalesced access\n// Each thread block handles one output position, threads compute different output features\ntemplate<int TILE_SIZE>\n__global__ void linear_kernel(\n const float* __restrict__ input,\n const float* __restrict__ weight,\n const float* __restrict__ bias,\n float* __restrict__ output,\n int B, int L, int in_dim, int out_dim,\n bool has_bias\n) {\n int idx = blockIdx.x; // B * L\n int j = threadIdx.x; // out_dim\n \n if (idx >= B * L || j >= out_dim) return;\n \n int in_base = idx * in_dim;\n int out_base = idx * out_dim;\n \n float sum = has_bias ? bias[j] : 0.0f;\n \n // Process in tiles for better cache utilization\n for (int k = 0; k < in_dim; k++) {\n sum += input[in_base + k] * weight[j * in_dim + k];\n }\n \n output[out_base + j] = sum;\n}\n\n// Conv1d causal depthwise kernel - each thread handles one output element\n__global__ void conv1d_causal_kernel(\n const float* __restrict__ input,\n const float* __restrict__ weight,\n const float* __restrict__ bias,\n float* __restrict__ output,\n int B, int L, int channels, int kernel_size\n) {\n int b = blockIdx.x / channels;\n int c = blockIdx.x % channels;\n int t = blockIdx.y * blockDim.x + threadIdx.x;\n \n if (t >= L) return;\n \n int in_base = b * channels * L + c * L;\n int out_base = b * channels * L + c * L;\n \n float sum = bias[c];\n \n for (int k = 0; k < kernel_size; k++) {\n int t_in = t - (kernel_size - 1) + k;\n if (t_in >= 0) {\n sum += input[in_base + t_in] * weight[c * kernel_size + k];\n }\n }\n \n output[out_base + t] = sum;\n}\n\n// Optimized selective scan kernel\n// Each thread handles one (batch, d_inner) pair\n// Uses coalesced memory access by loading data in a specific pattern\n__global__ void selective_scan_kernel(\n const float* __restrict__ x_branch, // [B, L, d_inner]\n const float* __restrict__ delta, // [B, L, d_inner]\n const float* __restrict__ B_t, // [B, L, d_state]\n const float* __restrict__ C_t, // [B, L, d_state]\n const float* __restrict__ A_log, // [d_inner, d_state]\n const float* __restrict__ D_param, // [d_inner]\n float* __restrict__ output, // [B, L, d_inner]\n int B, int L, int d_inner, int d_state\n) {\n int b = blockIdx.y;\n int di = blockIdx.x * blockDim.x + threadIdx.x;\n \n if (b >= B || di >= d_inner) return;\n \n // Load A for this d_inner\n float A[MAX_D_STATE];\n for (int n = 0; n < d_state; n++) {\n A[n] = -expf(A_log[di * d_state + n]);\n }\n \n // Initialize state\n float h[MAX_D_STATE];\n for (int n = 0; n < d_state; n++) {\n h[n] = 0.0f;\n }\n \n float D_val = D_param[di];\n \n // Process all time steps\n for (int t = 0; t < L; t++) {\n float delta_t = delta[b * L * d_inner + t * d_inner + di];\n float x_t = x_branch[b * L * d_inner + t * d_inner + di];\n \n // Update state\n for (int n = 0; n < d_state; n++) {\n float dA = expf(delta_t * A[n]);\n float dB = delta_t * B_t[b * L * d_state + t * d_state + n];\n float x_db = x_t * dB;\n h[n] = dA * h[n] + x_db;\n }\n \n // Compute output\n float y_t = 0.0f;\n for (int n = 0; n < d_state; n++) {\n y_t += h[n] * C_t[b * L * d_state + t * d_state + n];\n }\n y_t = y_t + x_t * D_val;\n \n output[b * L * d_inner + t * d_inner + di] = y_t;\n }\n}\n\n// Fused kernel: proj_delta + proj_B + proj_C + softplus + selective_scan + gate + out_proj\n// This reduces memory traffic significantly\n__global__ void mamba_fused_kernel(\n const float* __restrict__ x_branch, // [B, L, d_inner]\n const float* __restrict__ z, // [B, L, d_inner]\n const float* __restrict__ proj_delta_w, // [d_inner, d_inner]\n const float* __restrict__ proj_delta_b, // [d_inner]\n const float* __restrict__ proj_B_w, // [d_state, d_inner]\n const float* __restrict__ proj_C_w, // [d_state, d_inner]\n const float* __restrict__ A_log, // [d_inner, d_state]\n const float* __restrict__ D_param, // [d_inner]\n const float* __restrict__ out_proj_w, // [d_model, d_inner]\n float* __restrict__ output, // [B, L, d_model]\n int B, int L, int d_inner, int d_state, int d_model\n) {\n int b = blockIdx.y;\n int di = blockIdx.x * blockDim.x + threadIdx.x;\n \n if (b >= B || di >= d_inner) return;\n \n // Load A for this d_inner\n float A[MAX_D_STATE];\n for (int n = 0; n < d_state; n++) {\n A[n] = -expf(A_log[di * d_state + n]);\n }\n \n float D_val = D_param[di];\n \n // Initialize state\n float h[MAX_D_STATE];\n for (int n = 0; n < d_state; n++) {\n h[n] = 0.0f;\n }\n \n for (int t = 0; t < L; t++) {\n int in_base = b * L * d_inner + t * d_inner;\n float x_t = x_branch[in_base + di];\n \n // Compute delta_proj[t, di]\n float delta_proj = proj_delta_b[di];\n for (int k = 0; k < d_inner; k++) {\n delta_proj += x_branch[in_base + k] * proj_delta_w[di * d_inner + k];\n }\n \n // Softplus\n float delta_t = logf(1.0f + expf(delta_proj));\n \n // Compute B_t[t, :] and C_t[t, :]\n float B_t_val[MAX_D_STATE];\n float C_t_val[MAX_D_STATE];\n for (int n = 0; n < d_state; n++) {\n B_t_val[n] = 0.0f;\n C_t_val[n] = 0.0f;\n for (int k = 0; k < d_inner; k++) {\n B_t_val[n] += x_branch[in_base + k] * proj_B_w[n * d_inner + k];\n C_t_val[n] += x_branch[in_base + k] * proj_C_w[n * d_inner + k];\n }\n }\n \n // Update state\n for (int n = 0; n < d_state; n++) {\n float dA = expf(delta_t * A[n]);\n float dB = delta_t * B_t_val[n];\n float x_db = x_t * dB;\n h[n] = dA * h[n] + x_db;\n }\n \n // Compute output y[t, di]\n float y_t = 0.0f;\n for (int n = 0; n < d_state; n++) {\n y_t += h[n] * C_t_val[n];\n }\n y_t = y_t + x_t * D_val;\n \n // Gate with z\n float z_val = z[in_base + di];\n float z_act = z_val * (1.0f / (1.0f + expf(-z_val)));\n float gated = y_t * z_act;\n \n // Store for later reduction\n // We'll write to a temporary buffer and reduce later\n // For now, just store the gated value\n output[b * L * d_inner + t * d_inner + di] = gated;\n }\n}\n\n// Host functions\nvoid rmsnorm_cuda(\n const torch::Tensor& input,\n const torch::Tensor& weight,\n torch::Tensor& output,\n int B, int L, int D\n) {\n int threads = min(D, 1024);\n int blocks = B * L;\n \n rmsnorm_kernel<<<blocks, threads>>>(\n input.data_ptr<float>(),\n weight.data_ptr<float>(),\n output.data_ptr<float>(),\n B, L, D, 1e-5f\n );\n}\n\nvoid linear_cuda(\n const torch::Tensor& input,\n const torch::Tensor& weight,\n const c10::optional<torch::Tensor>& bias,\n torch::Tensor& output,\n int B, int L, int in_dim, int out_dim\n) {\n int total = B * L;\n int threads = min(out_dim, 1024);\n int blocks = total;\n \n bool has_bias = bias.has_value();\n const float* bias_ptr = has_bias ? bias.value().data_ptr<float>() : nullptr;\n \n linear_kernel<64><<<blocks, threads>>>(\n input.data_ptr<float>(),\n weight.data_ptr<float>(),\n bias_ptr,\n output.data_ptr<float>(),\n B, L, in_dim, out_dim,\n has_bias\n );\n}\n\nvoid conv1d_causal_cuda(\n const torch::Tensor& input,\n const torch::Tensor& weight,\n const torch::Tensor& bias,\n torch::Tensor& output,\n int B, int L, int channels, int kernel_size\n) {\n dim3 grid(B * channels, (L + 255) / 256);\n int threads = 256;\n \n conv1d_causal_kernel<<<grid, threads>>>(\n input.data_ptr<float>(),\n weight.data_ptr<float>(),\n bias.data_ptr<float>(),\n output.data_ptr<float>(),\n B, L, channels, kernel_size\n );\n}\n\nvoid selective_scan_cuda(\n const torch::Tensor& x_branch,\n const torch::Tensor& delta,\n const torch::Tensor& B_t,\n const torch::Tensor& C_t,\n const torch::Tensor& A_log,\n const torch::Tensor& D_param,\n torch::Tensor& output,\n int B, int L, int d_inner, int d_state\n) {\n int threads = 32;\n int blocks_x = (d_inner + threads - 1) / threads;\n \n dim3 grid(blocks_x, B);\n \n selective_scan_kernel<<<grid, threads>>>(\n x_branch.data_ptr<float>(),\n delta.data_ptr<float>(),\n B_t.data_ptr<float>(),\n C_t.data_ptr<float>(),\n A_log.data_ptr<float>(),\n D_param.data_ptr<float>(),\n output.data_ptr<float>(),\n B, L, d_inner, d_state\n );\n}\n\"\"\"\n\ncpp_source = \"\"\"\nvoid rmsnorm_cuda(const torch::Tensor& input, const torch::Tensor& weight, torch::Tensor& output, int B, int L, int D);\nvoid linear_cuda(const torch::Tensor& input, const torch::Tensor& weight, const c10::optional<torch::Tensor>& bias, torch::Tensor& output, int B, int L, int in_dim, int out_dim);\nvoid conv1d_causal_cuda(const torch::Tensor& input, const torch::Tensor& weight, const torch::Tensor& bias, torch::Tensor& output, int B, int L, int channels, int kernel_size);\nvoid selective_scan_cuda(const torch::Tensor& x_branch, const torch::Tensor& delta, const torch::Tensor& B_t, const torch::Tensor& C_t, const torch::Tensor& A_log, const torch::Tensor& D_param, torch::Tensor& output, int B, int L, int d_inner, int d_state);\n\"\"\"\n\n# Load the CUDA extension\nmamba_ext = load_inline(\n name='mamba_ext',\n cpp_sources=cpp_source,\n cuda_sources=cuda_source,\n functions=['rmsnorm_cuda', 'linear_cuda', 'conv1d_causal_cuda', 'selective_scan_cuda'],\n verbose=False,\n extra_cuda_cflags=['-O3', '-use_fast_math']\n)\n\nclass ModelNew(nn.Module):\n def __init__(self, d_model: int, d_state: int, d_conv: int, expand: int):\n super(ModelNew, self).__init__()\n self.d_model = d_model\n self.d_inner = d_model * expand\n self.d_state = d_state\n self.d_conv = d_conv\n\n # Initialize parameters same as reference\n self.norm = nn.RMSNorm(d_model)\n self.in_proj = nn.Linear(d_model, self.d_inner * 2, bias=False)\n self.conv1d = nn.Conv1d(self.d_inner, self.d_inner, d_conv,\n padding=d_conv - 1, groups=self.d_inner, bias=True)\n self.activation = nn.SiLU()\n\n self.A_log = nn.Parameter(torch.randn(self.d_inner, d_state))\n self.proj_B = nn.Linear(self.d_inner, d_state, bias=False)\n self.proj_C = nn.Linear(self.d_inner, d_state, bias=False)\n self.proj_delta = nn.Linear(self.d_inner, self.d_inner, bias=True)\n self.D = nn.Parameter(torch.ones(self.d_inner))\n\n self.out_proj = nn.Linear(self.d_inner, d_model, bias=False)\n\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n B, L, D = x.shape\n residual = x\n \n # Ensure contiguous float32\n x = x.contiguous().float()\n residual_f = residual.contiguous().float()\n \n # RMSNorm\n x_norm = torch.empty_like(x)\n mamba_ext.rmsnorm_cuda(x, self.norm.weight, x_norm, B, L, D)\n \n # in_proj: Linear (d_model -> d_inner * 2)\n xz = torch.empty(B, L, self.d_inner * 2, dtype=torch.float32, device=x.device)\n mamba_ext.linear_cuda(x_norm, self.in_proj.weight, None, xz, B, L, D, self.d_inner * 2)\n \n # Chunk\n x_branch = xz[..., :self.d_inner].contiguous()\n z = xz[..., self.d_inner:].contiguous()\n \n # Transpose for conv1d: [B, d_inner, L]\n x_branch_t = x_branch.transpose(1, 2).contiguous()\n \n # Conv1d causal\n x_conv = torch.empty_like(x_branch_t)\n mamba_ext.conv1d_causal_cuda(x_branch_t, self.conv1d.weight.squeeze(-1), self.conv1d.bias, x_conv, B, L, self.d_inner, self.d_conv)\n \n # Transpose back: [B, L, d_inner]\n x_branch = x_conv.transpose(1, 2).contiguous()\n \n # SiLU activation\n x_branch = torch.nn.functional.silu(x_branch)\n \n # proj_delta, proj_B, proj_C\n delta_proj = torch.empty(B, L, self.d_inner, dtype=torch.float32, device=x.device)\n B_t = torch.empty(B, L, self.d_state, dtype=torch.float32, device=x.device)\n C_t = torch.empty(B, L, self.d_state, dtype=torch.float32, device=x.device)\n \n mamba_ext.linear_cuda(x_branch, self.proj_delta.weight, self.proj_delta.bias, delta_proj, B, L, self.d_inner, self.d_inner)\n mamba_ext.linear_cuda(x_branch, self.proj_B.weight, None, B_t, B, L, self.d_inner, self.d_state)\n mamba_ext.linear_cuda(x_branch, self.proj_C.weight, None, C_t, B, L, self.d_inner, self.d_state)\n \n # Softplus for delta\n delta = torch.nn.functional.softplus(delta_proj)\n \n # Selective scan\n y = torch.empty(B, L, self.d_inner, dtype=torch.float32, device=x.device)\n mamba_ext.selective_scan_cuda(x_branch, delta, B_t, C_t, self.A_log, self.D, y, B, L, self.d_inner, self.d_state)\n \n # SiLU on z\n z = torch.nn.functional.silu(z)\n \n # Gate: y * z\n y = y * z\n \n # out_proj: Linear (d_inner -> d_model)\n out = torch.empty(B, L, D, dtype=torch.float32, device=x.device)\n mamba_ext.linear_cuda(y, self.out_proj.weight, None, out, B, L, self.d_inner, D)\n \n # Add residual\n out = out + residual_f\n \n return out"
}