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\n# Define the CUDA kernel source\ncuda_source = r\"\"\"\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n#include <vector>\n#include <cfloat>\n\ntemplate <int BLOCK_SIZE>\n__global__ void forward_kernel(\n const float* __restrict__ pair_repr,\n const float* __restrict__ pair_mask,\n const float* __restrict__ ln_in_weight,\n const float* __restrict__ ln_in_bias,\n const float* __restrict__ proj_a_weight,\n const float* __restrict__ proj_a_bias,\n const float* __restrict__ gate_a_weight,\n const float* __restrict__ gate_a_bias,\n const float* __restrict__ proj_b_weight,\n const float* __restrict__ proj_b_bias,\n const float* __restrict__ gate_b_weight,\n const float* __restrict__ gate_b_bias,\n const float* __restrict__ ln_out_weight,\n const float* __restrict__ ln_out_bias,\n const float* __restrict__ out_proj_weight,\n const float* __restrict__ out_proj_bias,\n const float* __restrict__ out_gate_weight,\n const float* __restrict__ out_gate_bias,\n float* __restrict__ output,\n int B, int N, int pair_dim, int hidden_dim,\n float ln_in_eps, float ln_out_eps\n) {\n // Each thread block handles one (b, i, j) output element\n // We need to reduce over k dimension\n const int b = blockIdx.z;\n const int i = blockIdx.y;\n const int j = blockIdx.x;\n \n const int tid = threadIdx.x;\n const int num_threads = blockDim.x;\n \n extern __shared__ float smem[];\n float* a_shared = smem; // hidden_dim elements\n float* b_shared = smem + hidden_dim * BLOCK_SIZE; // hidden_dim elements\n \n // Each thread handles a subset of k values\n // We'll compute partial sums and then reduce\n \n float local_sum[64]; // Max hidden_dim\n for (int h = 0; h < hidden_dim; h++) {\n local_sum[h] = 0.0f;\n }\n \n // Process k values in chunks\n for (int k = tid; k < N; k += num_threads) {\n // Compute mask for position (b, i, k) and (b, j, k)\n float mask_ik = pair_mask[b * N * N + i * N + k];\n float mask_jk = pair_mask[b * N * N + j * N + k];\n float mask_val = mask_ik * mask_jk;\n \n if (mask_val > 0.0f) {\n // Load pair_repr[b, i, k, :] and apply layer norm\n // Layer norm: (x - mean) / sqrt(var + eps) * weight + bias\n float mean_ik = 0.0f, var_ik = 0.0f;\n for (int d = 0; d < pair_dim; d++) {\n float val = pair_repr[b * N * N * pair_dim + i * N * pair_dim + k * pair_dim + d];\n mean_ik += val;\n }\n mean_ik /= pair_dim;\n for (int d = 0; d < pair_dim; d++) {\n float val = pair_repr[b * N * N * pair_dim + i * N * pair_dim + k * pair_dim + d];\n float diff = val - mean_ik;\n var_ik += diff * diff;\n }\n var_ik /= pair_dim;\n float inv_std_ik = rsqrtf(var_ik + ln_in_eps);\n \n // Compute a[b, i, k, h] = proj_a(x) * sigmoid(gate_a(x)) * mask\n // proj_a and gate_a are linear layers\n for (int h = 0; h < hidden_dim; h++) {\n float proj_a_val = proj_a_bias[h];\n float gate_a_val = gate_a_bias[h];\n for (int d = 0; d < pair_dim; d++) {\n float ln_val = (pair_repr[b * N * N * pair_dim + i * N * pair_dim + k * pair_dim + d] - mean_ik) * inv_std_ik;\n ln_val = ln_val * ln_in_weight[d] + ln_in_bias[d];\n proj_a_val += ln_val * proj_a_weight[d * hidden_dim + h];\n gate_a_val += ln_val * gate_a_weight[d * hidden_dim + h];\n }\n float a_val = proj_a_val * (1.0f / (1.0f + expf(-gate_a_val))) * mask_ik;\n \n // Now for b[b, j, k, h]\n float proj_b_val = proj_b_bias[h];\n float gate_b_val = gate_b_bias[h];\n // Layer norm for (b, j, k)\n float mean_jk = 0.0f, var_jk = 0.0f;\n for (int d = 0; d < pair_dim; d++) {\n float val = pair_repr[b * N * N * pair_dim + j * N * pair_dim + k * pair_dim + d];\n mean_jk += val;\n }\n mean_jk /= pair_dim;\n for (int d = 0; d < pair_dim; d++) {\n float val = pair_repr[b * N * N * pair_dim + j * N * pair_dim + k * pair_dim + d];\n float diff = val - mean_jk;\n var_jk += diff * diff;\n }\n var_jk /= pair_dim;\n float inv_std_jk = rsqrtf(var_jk + ln_in_eps);\n \n for (int d = 0; d < pair_dim; d++) {\n float ln_val = (pair_repr[b * N * N * pair_dim + j * N * pair_dim + k * pair_dim + d] - mean_jk) * inv_std_jk;\n ln_val = ln_val * ln_in_weight[d] + ln_in_bias[d];\n proj_b_val += ln_val * proj_b_weight[d * hidden_dim + h];\n gate_b_val += ln_val * gate_b_weight[d * hidden_dim + h];\n }\n float b_val = proj_b_val * (1.0f / (1.0f + expf(-gate_b_val))) * mask_jk;\n \n local_sum[h] += a_val * b_val;\n }\n }\n }\n \n // Store local sums to shared memory and reduce\n for (int h = tid; h < hidden_dim; h += num_threads) {\n a_shared[h * BLOCK_SIZE + tid] = local_sum[h];\n }\n __syncthreads();\n \n // Reduce within block\n for (int h = 0; h < hidden_dim; h++) {\n float sum = a_shared[h * BLOCK_SIZE + tid];\n for (int stride = num_threads / 2; stride > 0; stride >>= 1) {\n if (tid < stride) {\n sum += a_shared[h * BLOCK_SIZE + tid + stride];\n }\n __syncthreads();\n a_shared[h * BLOCK_SIZE + tid] = sum;\n __syncthreads();\n }\n \n if (tid == 0) {\n // Apply output layer norm, projection, and gating\n float out_val = sum;\n \n // Output layer norm\n float mean_out = 0.0f, var_out = 0.0f;\n mean_out = out_val; // Single value mean is itself\n var_out = 0.0f; // Single value variance is 0\n \n // For single-element layer norm, just scale by weight\n float ln_out = (out_val - mean_out) * rsqrtf(var_out + ln_out_eps);\n ln_out = ln_out * ln_out_weight[h] + ln_out_bias[h];\n \n // Output projection: hidden_dim -> pair_dim\n // We'll store intermediate and handle in separate kernel\n output[b * N * N * hidden_dim + i * N * hidden_dim + j * hidden_dim + h] = ln_out;\n }\n }\n}\n\n__global__ void output_proj_kernel(\n const float* __restrict__ intermediate,\n const float* __restrict__ pair_repr,\n const float* __restrict__ out_proj_weight,\n const float* __restrict__ out_proj_bias,\n const float* __restrict__ out_gate_weight,\n const float* __restrict__ out_gate_bias,\n float* __restrict__ output,\n int B, int N, int pair_dim, int hidden_dim\n) {\n const int b = blockIdx.z;\n const int i = blockIdx.y;\n const int j = blockIdx.x;\n const int tid = threadIdx.x;\n \n // Each block handles one (b, i, j) pair\n // Threads cooperate on computing output projection\n for (int d = tid; d < pair_dim; d += blockDim.x) {\n float proj_val = out_proj_bias[d];\n for (int h = 0; h < hidden_dim; h++) {\n proj_val += intermediate[b * N * N * hidden_dim + i * N * hidden_dim + j * hidden_dim + h] \n * out_proj_weight[h * pair_dim + d];\n }\n \n // Output gate\n float gate_val = out_gate_bias[d];\n for (int h = 0; h < pair_dim; h++) {\n gate_val += pair_repr[b * N * N * pair_dim + i * N * pair_dim + j * pair_dim + h]\n * out_gate_weight[h * pair_dim + d];\n }\n float sigmoid_gate = 1.0f / (1.0f + expf(-gate_val));\n \n output[b * N * N * pair_dim + i * N * pair_dim + j * pair_dim + d] = sigmoid_gate * proj_val;\n }\n}\n\ntorch::Tensor forward_cuda(\n torch::Tensor pair_repr,\n torch::Tensor pair_mask,\n torch::Tensor ln_in_weight,\n torch::Tensor ln_in_bias,\n torch::Tensor proj_a_weight,\n torch::Tensor proj_a_bias,\n torch::Tensor gate_a_weight,\n torch::Tensor gate_a_bias,\n torch::Tensor proj_b_weight,\n torch::Tensor proj_b_bias,\n torch::Tensor gate_b_weight,\n torch::Tensor gate_b_bias,\n torch::Tensor ln_out_weight,\n torch::Tensor ln_out_bias,\n torch::Tensor out_proj_weight,\n torch::Tensor out_proj_bias,\n torch::Tensor out_gate_weight,\n torch::Tensor out_gate_bias,\n int64_t B, int64_t N, int64_t pair_dim, int64_t hidden_dim,\n double ln_in_eps, double ln_out_eps\n) {\n auto output = torch::zeros({B, N, N, pair_dim}, pair_repr.options());\n auto intermediate = torch::zeros({B, N, N, hidden_dim}, pair_repr.options());\n \n const int block_size = 256;\n dim3 grid(N, N, B);\n dim3 block(block_size);\n size_t shared_mem = hidden_dim * block_size * 2 * sizeof(float);\n \n forward_kernel<block_size><<<grid, block, shared_mem>>>(\n pair_repr.data_ptr<float>(),\n pair_mask.data_ptr<float>(),\n ln_in_weight.data_ptr<float>(),\n ln_in_bias.data_ptr<float>(),\n proj_a_weight.data_ptr<float>(),\n proj_a_bias.data_ptr<float>(),\n gate_a_weight.data_ptr<float>(),\n gate_a_bias.data_ptr<float>(),\n proj_b_weight.data_ptr<float>(),\n proj_b_bias.data_ptr<float>(),\n gate_b_weight.data_ptr<float>(),\n gate_b_bias.data_ptr<float>(),\n ln_out_weight.data_ptr<float>(),\n ln_out_bias.data_ptr<float>(),\n out_proj_weight.data_ptr<float>(),\n out_proj_bias.data_ptr<float>(),\n out_gate_weight.data_ptr<float>(),\n out_gate_bias.data_ptr<float>(),\n intermediate.data_ptr<float>(),\n B, N, pair_dim, hidden_dim,\n static_cast<float>(ln_in_eps),\n static_cast<float>(ln_out_eps)\n );\n \n output_proj_kernel<<<grid, block>>>(\n intermediate.data_ptr<float>(),\n pair_repr.data_ptr<float>(),\n out_proj_weight.data_ptr<float>(),\n out_proj_bias.data_ptr<float>(),\n out_gate_weight.data_ptr<float>(),\n out_gate_bias.data_ptr<float>(),\n output.data_ptr<float>(),\n B, N, pair_dim, hidden_dim\n );\n \n return output;\n}\n\"\"\"\n\n# Load the CUDA extension\nmodule = load_inline(\n name='triangular_mult_update_outgoing_cuda',\n cpp_sources='',\n cuda_sources=cuda_source,\n functions=['forward_cuda'],\n extra_cuda_cflags=['-O3'],\n verbose=True\n)\n\n\nclass ModelNew(nn.Module):\n \"\"\"\n Triangular multiplicative update \u2013 outgoing edges.\n Optimized CUDA implementation.\n \"\"\"\n\n def __init__(self, pair_dim, hidden_dim):\n super().__init__()\n self.layer_norm_in = nn.LayerNorm(pair_dim)\n self.proj_a = nn.Linear(pair_dim, hidden_dim)\n self.gate_a = nn.Linear(pair_dim, hidden_dim)\n self.proj_b = nn.Linear(pair_dim, hidden_dim)\n self.gate_b = nn.Linear(pair_dim, hidden_dim)\n self.layer_norm_out = nn.LayerNorm(hidden_dim)\n self.out_proj = nn.Linear(hidden_dim, pair_dim)\n self.out_gate = nn.Linear(pair_dim, pair_dim)\n\n def forward(self, pair_repr: torch.Tensor, pair_mask: torch.Tensor) -> torch.Tensor:\n \"\"\"\n Args:\n pair_repr: (B, N, N, pair_dim)\n pair_mask: (B, N, N)\n Returns:\n updated pair representation (B, N, N, pair_dim)\n \"\"\"\n B, N, _, pair_dim = pair_repr.shape\n hidden_dim = self.proj_a.out_features\n \n # Ensure tensors are contiguous and on GPU\n pair_repr = pair_repr.contiguous()\n pair_mask = pair_mask.contiguous()\n \n return module.forward_cuda(\n pair_repr,\n pair_mask,\n self.layer_norm_in.weight.contiguous(),\n self.layer_norm_in.bias.contiguous(),\n self.proj_a.weight.transpose(0, 1).contiguous(), # Transpose for correct layout\n self.proj_a.bias.contiguous(),\n self.gate_a.weight.transpose(0, 1).contiguous(),\n self.gate_a.bias.contiguous(),\n self.proj_b.weight.transpose(0, 1).contiguous(),\n self.proj_b.bias.contiguous(),\n self.gate_b.weight.transpose(0, 1).contiguous(),\n self.gate_b.bias.contiguous(),\n self.layer_norm_out.weight.contiguous(),\n self.layer_norm_out.bias.contiguous(),\n self.out_proj.weight.transpose(0, 1).contiguous(),\n self.out_proj.bias.contiguous(),\n self.out_gate.weight.transpose(0, 1).contiguous(),\n self.out_gate.bias.contiguous(),\n B, N, pair_dim, hidden_dim,\n self.layer_norm_in.eps,\n self.layer_norm_out.eps\n )"
}