call → compile_kernel
callcompile_kernel
{
"kernel_code": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport math\nfrom torch.utils.cpp_extension import load_inline\n\n# Define CUDA kernel for attention\ncuda_source = \"\"\"\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n#include <cmath>\n\ntemplate <typename scalar_t>\n__device__ inline scalar_t warpReduceSum(scalar_t val, int warpSize = 32) {\n for (int offset = warpSize / 2; offset > 0; offset /= 2)\n val += __shfl_down_sync(0xffffffff, val, offset);\n return val;\n}\n\ntemplate <typename scalar_t>\n__device__ inline void warpReduceMax(scalar_t* val, int warpSize = 32) {\n for (int offset = warpSize / 2; offset > 0; offset /= 2) {\n scalar_t other = __shfl_down_sync(0xffffffff, *val, offset);\n *val = max(*val, other);\n }\n}\n\n// Softmax over S dimension for each (B, N, h, i)\n// attn: (B, N, h, S, S) stored as (B*N*h, S, S)\ntemplate <typename scalar_t>\n__global__ void softmax_kernel(\n scalar_t* __restrict__ attn,\n const bool* __restrict__ mask,\n int B, int N, int h, int S) {\n \n int bnhi = blockIdx.x; // B*N*h*S indices\n int j = threadIdx.x;\n \n int b = bnhi / (N * h * S);\n int rem = bnhi % (N * h * S);\n int n = rem / (h * S);\n rem = rem % (h * S);\n int head = rem / S;\n int i = rem % S;\n \n int idx_base = bnhi * S; // Base index for this row\n \n // Find max for numerical stability\n scalar_t max_val = -INFINITY;\n bool all_masked = true;\n \n for (int jj = j; jj < S; jj += blockDim.x) {\n int mask_idx = b * N * S + n * S + jj;\n if (mask[mask_idx]) {\n all_masked = false;\n scalar_t val = attn[idx_base + jj];\n if (val > max_val) max_val = val;\n }\n }\n \n max_val = warpReduceMax(&max_val);\n __shared__ scalar_t shared_max;\n if (threadIdx.x % 32 == 0) {\n shared_max = max_val;\n }\n __syncthreads();\n max_val = shared_max;\n \n if (all_masked) {\n for (int jj = j; jj < S; jj += blockDim.x) {\n attn[idx_base + jj] = 0.0;\n }\n return;\n }\n \n // Compute exp and sum\n scalar_t sum = 0.0;\n for (int jj = j; jj < S; jj += blockDim.x) {\n int mask_idx = b * N * S + n * S + jj;\n if (mask[mask_idx]) {\n sum += expf(attn[idx_base + jj] - max_val);\n }\n }\n \n sum = warpReduceSum(sum);\n __shared__ scalar_t shared_sum;\n if (threadIdx.x % 32 == 0) {\n shared_sum = sum;\n }\n __syncthreads();\n sum = shared_sum;\n \n // Normalize\n for (int jj = j; jj < S; jj += blockDim.x) {\n int mask_idx = b * N * S + n * S + jj;\n if (mask[mask_idx]) {\n attn[idx_base + jj] = expf(attn[idx_base + jj] - max_val) / sum;\n } else {\n attn[idx_base + jj] = 0.0;\n }\n }\n}\n\n// Attention computation: Q @ K^T / sqrt(d)\n// q: (B, N, h, S, d), k: (B, N, h, S, d) -> attn: (B, N, h, S, S)\ntemplate <typename scalar_t>\n__global__ void attention_score_kernel(\n const scalar_t* __restrict__ q,\n const scalar_t* __restrict__ k,\n scalar_t* __restrict__ attn,\n int B, int N, int h, int S, int d) {\n \n int bnhs = blockIdx.x; // B*N*h*S total blocks\n int tid = threadIdx.x;\n \n int b = bnhs / (N * h * S);\n int rem = bnhs % (N * h * S);\n int n = rem / (h * S);\n rem = rem % (h * S);\n int head = rem / S;\n int i = rem % S;\n \n int q_idx = b * N * h * S * d + n * h * S * d + head * S * d + i * d;\n \n // Compute dot product with all K positions\n for (int j = tid; j < S; j += blockDim.x) {\n int k_idx = b * N * h * S * d + n * h * S * d + head * S * d + j * d;\n scalar_t sum = 0.0;\n for (int dd = 0; dd < d; dd++) {\n sum += q[q_idx + dd] * k[k_idx + dd];\n }\n int attn_idx = b * N * h * S * S + n * h * S * S + head * S * S + i * S + j;\n attn[attn_idx] = sum / sqrtf((float)d);\n }\n}\n\n// Output computation: attn @ V\n// attn: (B, N, h, S, S), v: (B, N, h, S, d) -> out: (B, N, h, S, d)\ntemplate <typename scalar_t>\n__global__ void attention_output_kernel(\n const scalar_t* __restrict__ attn,\n const scalar_t* __restrict__ v,\n const scalar_t* __restrict__ g,\n scalar_t* __restrict__ out,\n int B, int N, int h, int S, int d) {\n \n int bnhs = blockIdx.x; // B*N*h*S total blocks\n int tid = threadIdx.x;\n \n int b = bnhs / (N * h * S);\n int rem = bnhs % (N * h * S);\n int n = rem / (h * S);\n rem = rem % (h * S);\n int head = rem / S;\n int i = rem % S;\n \n int attn_base = b * N * h * S * S + n * h * S * S + head * S * S + i * S;\n int out_idx = b * N * h * S * d + n * h * S * d + head * S * d + i * d;\n \n // Compute weighted sum over v\n for (int dd = tid; dd < d; dd += blockDim.x) {\n scalar_t sum = 0.0;\n for (int j = 0; j < S; j++) {\n int v_idx = b * N * h * S * d + n * h * S * d + head * S * d + j * d + dd;\n sum += attn[attn_base + j] * v[v_idx];\n }\n int g_idx = b * N * h * S * d + n * h * S * d + head * S * d + i * d + dd;\n out[out_idx + dd] = g[g_idx] * sum; // Apply gate\n }\n}\n\ntorch::Tensor attention_cuda(\n torch::Tensor q,\n torch::Tensor k,\n torch::Tensor v,\n torch::Tensor g,\n torch::Tensor mask) {\n \n int B = q.size(0);\n int N = q.size(1);\n int h = q.size(2);\n int S = q.size(3);\n int d = q.size(4);\n \n auto attn = torch::empty({B, N, h, S, S}, q.options());\n auto out = torch::empty({B, N, h, S, d}, q.options());\n \n int blocks = B * N * h * S;\n int threads = min(256, S);\n \n // Compute attention scores\n AT_DISPATCH_FLOATING_TYPES(q.scalar_type(), \"attention_score\", ([&] {\n attention_score_kernel<scalar_t><<<blocks, threads>>>(\n q.data_ptr<scalar_t>(),\n k.data_ptr<scalar_t>(),\n attn.data_ptr<scalar_t>(),\n B, N, h, S, d);\n }));\n \n // Softmax\n AT_DISPATCH_FLOATING_TYPES(q.scalar_type(), \"softmax\", ([&] {\n softmax_kernel<scalar_t><<<blocks, threads>>>(\n attn.data_ptr<scalar_t>(),\n mask.data_ptr<bool>(),\n B, N, h, S);\n }));\n \n // Output\n int threads_out = min(256, d);\n AT_DISPATCH_FLOATING_TYPES(q.scalar_type(), \"attention_output\", ([&] {\n attention_output_kernel<scalar_t><<<blocks, threads_out>>>(\n attn.data_ptr<scalar_t>(),\n v.data_ptr<scalar_t>(),\n g.data_ptr<scalar_t>(),\n out.data_ptr<scalar_t>(),\n B, N, h, S, d);\n }));\n \n return out;\n}\n\"\"\"\n\ncpp_source = \"\"\"\ntorch::Tensor attention_cuda(torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor g, torch::Tensor mask);\n\"\"\"\n\nattention_module = load_inline(\n name='attention_cuda',\n cpp_sources=cpp_source,\n cuda_sources=cuda_source,\n functions=['attention_cuda'],\n verbose=True,\n extra_cuda_cflags=['-O3']\n)\n\n\nclass ModelNew(nn.Module):\n def __init__(self, msa_dim, num_heads, dropout=0.0):\n super().__init__()\n self.num_heads = num_heads\n self.head_dim = msa_dim // num_heads\n assert msa_dim % num_heads == 0\n\n self.layer_norm = nn.LayerNorm(msa_dim)\n self.query = nn.Linear(msa_dim, msa_dim, bias=False)\n self.key = nn.Linear(msa_dim, msa_dim, bias=False)\n self.value = nn.Linear(msa_dim, msa_dim, bias=False)\n self.gate = nn.Linear(msa_dim, msa_dim)\n self.out_proj = nn.Linear(msa_dim, msa_dim)\n self.dropout = nn.Dropout(dropout)\n\n def forward(self, msa_repr: torch.Tensor, msa_mask: torch.Tensor) -> torch.Tensor:\n B, S, N, D = msa_repr.shape\n h = self.num_heads\n d = self.head_dim\n\n # Transpose so column (residue position) is the batch dim for attention\n x = msa_repr.permute(0, 2, 1, 3) # (B, N, S, D)\n x = self.layer_norm(x)\n\n q = self.query(x).view(B, N, S, h, d).permute(0, 1, 3, 2, 4)\n k = self.key(x).view(B, N, S, h, d).permute(0, 1, 3, 2, 4)\n v = self.value(x).view(B, N, S, h, d).permute(0, 1, 3, 2, 4)\n g = torch.sigmoid(self.gate(x)).view(B, N, S, h, d).permute(0, 1, 3, 2, 4)\n\n # Custom CUDA attention\n col_mask = msa_mask.permute(0, 2, 1) # (B, N, S)\n \n out = attention_module.attention_cuda(q, v, k, g, col_mask)\n \n out = out.permute(0, 1, 3, 2, 4).reshape(B, N, S, D)\n out = self.out_proj(out)\n return out.permute(0, 2, 1, 3) # back to (B, S, N, D)"
}