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 the attention kernel\ncuda_source = \"\"\"\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n#include <cmath>\n\n// Constants for shared memory and block sizes\nconstexpr int BLOCK_SIZE = 16;\nconstexpr int TILE_SIZE = 32;\n\n__device__ __forceinline__ float warp_reduce_max(float val, int lane_id) {\n #pragma unroll\n for (int offset = 16; offset > 0; offset /= 2) {\n val = fmaxf(val, __shfl_down_sync(0xffffffff, val, offset));\n }\n return val;\n}\n\n__device__ __forceinline__ float warp_reduce_sum(float val, int lane_id) {\n #pragma unroll\n for (int offset = 16; offset > 0; offset /= 2) {\n val += __shfl_down_sync(0xffffffff, val, offset);\n }\n return val;\n}\n\n// Kernel for computing attention scores and applying softmax\n// Computes Q @ K^T + bias, then softmax\n__global__ void attention_softmax_kernel(\n const float* __restrict__ q,\n const float* __restrict__ k,\n const float* __restrict__ bias,\n const bool* __restrict__ mask,\n float* __restrict__ attn_out,\n int B, int S, int H, int N, int D\n) {\n // Block and thread indices\n int b = blockIdx.z; // batch\n int s = blockIdx.y; // sequence depth\n int h = blockIdx.x; // head\n \n int tx = threadIdx.x; // query position\n int ty = threadIdx.y; // key position\n \n // Each thread computes one element of attention matrix\n // attn[b, s, h, i, j] = sum_d(q[b,s,h,i,d] * k[b,s,h,j,d]) / sqrt(d)\n \n float scale = 1.0f / sqrtf((float)D);\n \n // Load Q row for this thread\n extern __shared__ float shared_mem[];\n float* q_tile = shared_mem; // TILE_SIZE * D elements\n float* k_tile = &shared_mem[TILE_SIZE * D]; // TILE_SIZE * D elements\n \n int i = tx; // query index within tile\n int j = ty; // key index within tile\n \n float score = 0.0f;\n \n // Compute dot product in chunks of D\n for (int d_offset = 0; d_offset < D; d_offset += TILE_SIZE) {\n // Cooperative load of Q and K tiles into shared memory\n int d = d_offset + tx;\n if (d < D && i < N) {\n q_tile[i * D + d] = q[b * S * H * N * D + s * H * N * D + h * N * D + i * D + d];\n }\n if (d < D && j < N) {\n k_tile[j * D + d] = k[b * S * H * N * D + s * H * N * D + h * N * D + j * D + d];\n }\n __syncthreads();\n \n // Compute partial dot product\n for (int d_inner = 0; d_inner < TILE_SIZE && (d_offset + d_inner) < D; d_inner++) {\n if (i < N && j < N) {\n score += q_tile[i * D + d_offset + d_inner] * k_tile[j * D + d_offset + d_inner];\n }\n }\n __syncthreads();\n }\n \n // Apply scaling and bias\n if (i < N && j < N) {\n score *= scale;\n // Add bias: bias shape is (B, H, N, N)\n score += bias[b * H * N * N + h * N * N + i * N + j];\n }\n \n // Apply mask - check if position j is masked\n // mask shape: (B, S, N) - 1 means valid, 0 means masked\n if (i < N && j < N) {\n bool mask_val = mask[b * S * N + s * N + j];\n if (!mask_val) {\n score = -INFINITY;\n }\n }\n \n // Shared memory for reduction\n __shared__ float smem[N * N];\n \n // Store score for softmax\n if (i < N && j < N) {\n smem[i * N + j] = score;\n }\n __syncthreads();\n \n // Softmax over j dimension (key positions)\n // Each row i computes its own softmax\n if (i < N) {\n // Find max in row\n float max_val = -INFINITY;\n for (int jj = 0; jj < N; jj++) {\n max_val = fmaxf(max_val, smem[i * N + jj]);\n }\n \n // Compute exp and sum\n float sum_exp = 0.0f;\n for (int jj = 0; jj < N; jj++) {\n float val = expf(smem[i * N + jj] - max_val);\n smem[i * N + jj] = val;\n sum_exp += val;\n }\n \n // Normalize\n float inv_sum = 1.0f / sum_exp;\n for (int jj = 0; jj < N; jj++) {\n smem[i * N + jj] *= inv_sum;\n }\n }\n __syncthreads();\n \n // Write output\n if (i < N && j < N) {\n attn_out[b * S * H * N * N + s * H * N * N + h * N * N + i * N + j] = smem[i * N + j];\n }\n}\n\n// Kernel to compute output from attention and values\n// out[bshid] = sum_j(attn[bshij] * v[bshjd])\n__global__ void attention_value_kernel(\n const float* __restrict__ attn,\n const float* __restrict__ v,\n float* __restrict__ out,\n int B, int S, int H, int N, int D\n) {\n int b = blockIdx.z;\n int s = blockIdx.y;\n int h = blockIdx.x;\n \n int i = threadIdx.x; // query position\n int d = threadIdx.y; // dimension\n \n if (i >= N || d >= D) return;\n \n float val = 0.0f;\n for (int j = 0; j < N; j++) {\n float attn_val = attn[b * S * H * N * N + s * H * N * N + h * N * N + i * N + j];\n float v_val = v[b * S * H * N * D + s * H * N * D + h * N * D + j * D + d];\n val += attn_val * v_val;\n }\n \n out[b * S * H * N * D + s * H * N * D + h * N * D + i * D + d] = val;\n}\n\n// Combined kernel: compute attention, softmax, and multiply with values\n__global__ void fused_attention_kernel(\n const float* __restrict__ q,\n const float* __restrict__ k,\n const float* __restrict__ v,\n const float* __restrict__ bias,\n const bool* __restrict__ mask,\n float* __restrict__ out,\n int B, int S, int H, int N, int D\n) {\n // Grid: (B*S*H) blocks, each block computes N rows of attention\n int bsh_idx = blockIdx.x;\n int b = bsh_idx / (S * H);\n int remainder = bsh_idx % (S * H);\n int s = remainder / H;\n int h = remainder % H;\n \n int i = threadIdx.x; // query index\n \n extern __shared__ float smem[];\n float* k_cache = smem; // N * D\n float* attn_row = &smem[N * D]; // N elements for softmax\n \n float scale = 1.0f / sqrtf((float)D);\n \n // Each thread computes attention for one query position\n if (i >= N) return;\n \n // Load K matrix into shared memory (all threads cooperate)\n for (int j = threadIdx.x; j < N * D; j += blockDim.x) {\n k_cache[j] = k[b * S * H * N * D + s * H * N * D + h * N * D + j];\n }\n __syncthreads();\n \n // Compute attention scores for row i\n float max_score = -INFINITY;\n \n // Load Q row for thread i\n float q_row[64]; // Assume D <= 64\n for (int d = 0; d < D; d++) {\n q_row[d] = q[b * S * H * N * D + s * H * N * D + h * N * D + i * D + d];\n }\n \n // Compute scores\n for (int j = 0; j < N; j++) {\n float score = 0.0f;\n for (int d = 0; d < D; d++) {\n score += q_row[d] * k_cache[j * D + d];\n }\n score *= scale;\n \n // Add bias\n score += bias[b * H * N * N + h * N * N + i * N + j];\n \n // Apply mask\n bool mask_val = mask[b * S * N + s * N + j];\n if (!mask_val) {\n score = -INFINITY;\n }\n \n attn_row[j] = score;\n max_score = fmaxf(max_score, score);\n }\n \n // Warp reduction for max (within thread block)\n __shared__ float max_per_thread[32];\n max_per_thread[threadIdx.x] = max_score;\n __syncthreads();\n \n // Compute global max\n float global_max = -INFINITY;\n for (int t = 0; t < N && t < 32; t++) {\n global_max = fmaxf(global_max, max_per_thread[t]);\n }\n \n // Compute exp and sum\n float sum_exp = 0.0f;\n for (int j = 0; j < N; j++) {\n attn_row[j] = expf(attn_row[j] - global_max);\n sum_exp += attn_row[j];\n }\n \n // Normalize\n float inv_sum = 1.0f / sum_exp;\n for (int j = 0; j < N; j++) {\n attn_row[j] *= inv_sum;\n }\n \n // Multiply with V and accumulate\n for (int d = 0; d < D; d++) {\n float val = 0.0f;\n for (int j = 0; j < N; j++) {\n float v_val = v[b * S * H * N * D + s * H * N * D + h * N * D + j * D + d];\n val += attn_row[j] * v_val;\n }\n out[b * S * H * N * D + s * H * N * D + h * N * D + i * D + d] = val;\n }\n}\n\n// Host function to launch the kernel\ntorch::Tensor attention_forward_cuda(\n torch::Tensor q,\n torch::Tensor k,\n torch::Tensor v,\n torch::Tensor bias,\n torch::Tensor mask\n) {\n // Input shapes:\n // q, k, v: (B, S, H, N, D)\n // bias: (B, H, N, N)\n // mask: (B, S, N)\n \n auto sizes = q.sizes();\n int B = sizes[0];\n int S = sizes[1];\n int H = sizes[2];\n int N = sizes[3];\n int D = sizes[4];\n \n auto out = torch::empty_like(q);\n \n // Use the fused kernel\n int blocks = B * S * H;\n int threads = N; // One thread per query position\n int shared_mem = N * D * sizeof(float) + N * sizeof(float);\n \n dim3 grid(blocks);\n dim3 block(threads);\n \n AT_DISPATCH_FLOATING_TYPES(q.scalar_type(), \"fused_attention_kernel\", ([&] {\n fused_attention_kernel<<<grid, block, shared_mem>>>(\n q.data_ptr<float>(),\n k.data_ptr<float>(),\n v.data_ptr<float>(),\n bias.data_ptr<float>(),\n mask.data_ptr<bool>(),\n out.data_ptr<float>(),\n B, S, H, N, D\n );\n }));\n \n return out;\n}\n\"\"\"\n\n# Load the CUDA extension\nattention_ext = load_inline(\n name='attention_cuda',\n cpp_sources='',\n cuda_sources=cuda_source,\n functions=['attention_forward_cuda'],\n extra_cuda_cflags=['-O3', '--use_fast_math'],\n verbose=True\n)\n\n\nclass ModelNew(nn.Module):\n \"\"\"\n Row-wise gated self-attention over Multiple Sequence Alignment (MSA)\n representations with custom CUDA kernel for attention computation.\n \"\"\"\n\n def __init__(self, msa_dim, pair_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.pair_bias = nn.Linear(pair_dim, num_heads, bias=False)\n self.out_proj = nn.Linear(msa_dim, msa_dim)\n self.dropout = nn.Dropout(dropout)\n\n def forward(\n self,\n msa_repr: torch.Tensor,\n pair_repr: torch.Tensor,\n msa_mask: torch.Tensor,\n ) -> torch.Tensor:\n \"\"\"\n Args:\n msa_repr: (B, S, N, msa_dim)\n pair_repr: (B, N, N, pair_dim)\n msa_mask: (B, S, N)\n Returns:\n updated MSA (B, S, N, msa_dim)\n \"\"\"\n B, S, N, D = msa_repr.shape\n h = self.num_heads\n d = self.head_dim\n\n x = self.layer_norm(msa_repr)\n q = self.query(x).view(B, S, N, h, d).permute(0, 1, 3, 2, 4)\n k = self.key(x).view(B, S, N, h, d).permute(0, 1, 3, 2, 4)\n v = self.value(x).view(B, S, N, h, d).permute(0, 1, 3, 2, 4)\n g = torch.sigmoid(self.gate(x)).view(B, S, N, h, d).permute(0, 1, 3, 2, 4)\n\n # Pair bias: shared across MSA rows\n bias = self.pair_bias(pair_repr).permute(0, 3, 1, 2) # (B, h, N, N)\n\n # Use custom CUDA kernel for attention\n # Convert mask to bool tensor\n mask_bool = msa_mask.bool()\n \n out = attention_ext.attention_forward_cuda(\n q.contiguous(),\n k.contiguous(),\n v.contiguous(),\n bias.contiguous(),\n mask_bool.contiguous()\n )\n \n # Apply gate\n out = g * out\n \n # Reshape and project output\n out = out.permute(0, 1, 3, 2, 4).reshape(B, S, N, D)\n return self.out_proj(out)"
}