callcompile_kernel
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\nimport math\n\n# Define CUDA source\ncuda_source = r'''\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n#include <c10/cuda/CUDAStream.h>\n#include <ATen/cuda/CUDAContext.h>\n#include <cublas_v2.h>\n#include <vector>\n\n#define CUDA_CHECK(call) \\\n do { \\\n cudaError_t err = call; \\\n if (err != cudaSuccess) { \\\n fprintf(stderr, \"CUDA error at %s:%d: %s\\n\", __FILE__, __LINE__, \\\n cudaGetErrorString(err)); \\\n exit(EXIT_FAILURE); \\\n } \\\n } while(0)\n\n#define CUBLAS_CHECK(call) \\\n do { \\\n cublasStatus_t status = call; \\\n if (status != CUBLAS_STATUS_SUCCESS) { \\\n fprintf(stderr, \"cuBLAS error at %s:%d: %d\\n\", __FILE__, __LINE__, status); \\\n exit(EXIT_FAILURE); \\\n } \\\n } while(0)\n\n// Fused kernel: layer norm + projection + gate + sigmoid + mask\n// Processes one position (i,j) across all channels\ntemplate <typename scalar_t, int PROJ_DIM, int HIDDEN_DIM>\n__global__ void fused_proj_gate_kernel(\n const scalar_t* __restrict__ input, // (B, N, N, PROJ_DIM)\n const scalar_t* __restrict__ weight, // (PROJ_DIM, HIDDEN_DIM)\n const scalar_t* __restrict__ bias, // (HIDDEN_DIM,)\n const scalar_t* __restrict__ gate_weight, // (PROJ_DIM, HIDDEN_DIM)\n const scalar_t* __restrict__ gate_bias, // (HIDDEN_DIM,)\n const scalar_t* __restrict__ gamma, // (PROJ_DIM,) for layer norm\n const scalar_t* __restrict__ beta, // (PROJ_DIM,) for layer norm\n scalar_t* __restrict__ output, // (B, N, N, HIDDEN_DIM)\n const scalar_t* __restrict__ mask, // (B, N, N, 1) or nullptr\n int B, int N, int D_in, int D_out,\n float eps\n) {\n const int b = blockIdx.x;\n const int ij = blockIdx.y * blockDim.y + threadIdx.y;\n const int h = threadIdx.x;\n \n if (ij >= N * N) return;\n \n const int i = ij / N;\n const int j = ij % N;\n const int base_in = (b * N * N + ij) * D_in;\n const int base_out = (b * N * N + ij) * D_out;\n \n // Load input and compute mean/var for layer norm\n scalar_t mean = scalar_t(0);\n scalar_t var = scalar_t(0);\n \n for (int d = 0; d < D_in; d++) {\n scalar_t val = input[base_in + d];\n mean += val;\n }\n mean /= scalar_t(D_in);\n \n for (int d = 0; d < D_in; d++) {\n scalar_t val = input[base_in + d] - mean;\n var += val * val;\n }\n var /= scalar_t(D_in);\n scalar_t inv_std = scalar_t(1) / sqrt(var + scalar_t(eps));\n \n // Compute projection\n scalar_t proj_val = bias[h];\n scalar_t gate_val = gate_bias[h];\n \n for (int d = 0; d < D_in; d++) {\n scalar_t normalized = (input[base_in + d] - mean) * inv_std;\n normalized = normalized * gamma[d] + beta[d];\n \n proj_val += normalized * weight[d * D_out + h];\n gate_val += normalized * gate_weight[d * D_out + h];\n }\n \n // Sigmoid for gate: sigmoid(x) = 1 / (1 + exp(-x))\n scalar_t sigmoid_gate = scalar_t(1) / (scalar_t(1) + exp(-gate_val));\n \n // Apply mask if present\n scalar_t mask_val = scalar_t(1);\n if (mask != nullptr) {\n mask_val = mask[b * N * N + ij];\n }\n \n // Final output: proj * sigmoid(gate) * mask\n output[base_out + h] = proj_val * sigmoid_gate * mask_val;\n}\n\n// Fused kernel for output: layer norm + projection + gate\ntemplate <typename scalar_t>\n__global__ void fused_output_kernel(\n const scalar_t* __restrict__ input, // (B, N, N, HIDDEN_DIM)\n const scalar_t* __restrict__ weight, // (HIDDEN_DIM, PAIR_DIM)\n const scalar_t* __restrict__ bias, // (PAIR_DIM,)\n const scalar_t* __restrict__ gate_input, // (B, N, N, PAIR_DIM) original pair_repr\n const scalar_t* __restrict__ gate_weight, // (PAIR_DIM, PAIR_DIM)\n const scalar_t* __restrict__ gate_bias, // (PAIR_DIM,)\n const scalar_t* __restrict__ gamma, // (HIDDEN_DIM,)\n const scalar_t* __restrict__ beta, // (HIDDEN_DIM,)\n scalar_t* __restrict__ output, // (B, N, N, PAIR_DIM)\n int B, int N, int D_in, int D_out,\n float eps\n) {\n const int b = blockIdx.x;\n const int ij = blockIdx.y * blockDim.y + threadIdx.y;\n const int h = threadIdx.x;\n \n if (ij >= N * N || h >= D_out) return;\n \n const int base_in = (b * N * N + ij) * D_in;\n const int base_out = (b * N * N + ij) * D_out;\n const int base_gate_in = (b * N * N + ij) * D_out;\n \n // Compute layer norm for input\n scalar_t mean = scalar_t(0);\n scalar_t var = scalar_t(0);\n \n for (int d = 0; d < D_in; d++) {\n scalar_t val = input[base_in + d];\n mean += val;\n }\n mean /= scalar_t(D_in);\n \n for (int d = 0; d < D_in; d++) {\n scalar_t val = input[base_in + d] - mean;\n var += val * val;\n }\n var /= scalar_t(D_in);\n scalar_t inv_std = scalar_t(1) / sqrt(var + scalar_t(eps));\n \n // Compute projection\n scalar_t proj_val = bias[h];\n for (int d = 0; d < D_in; d++) {\n scalar_t normalized = (input[base_in + d] - mean) * inv_std;\n normalized = normalized * gamma[d] + beta[d];\n proj_val += normalized * weight[d * D_out + h];\n }\n \n // Compute gate from original pair_repr\n scalar_t gate_val = gate_bias[h];\n for (int d = 0; d < D_out; d++) {\n gate_val += gate_input[base_gate_in + d] * gate_weight[d * D_out + h];\n }\n \n scalar_t sigmoid_gate = scalar_t(1) / (scalar_t(1) + exp(-gate_val));\n \n output[base_out + h] = sigmoid_gate * proj_val;\n}\n\n// Optimized einsum kernel: computes out[i,j,h] = sum_k a[i,k,h] * b[j,k,h]\n// Uses shared memory and coalesced accesses\ntemplate <typename scalar_t, int TILE_SIZE>\n__global__ void einsum_outgoing_kernel(\n const scalar_t* __restrict__ a, // (B, N, N, H) - layout is b,i,k,h\n const scalar_t* __restrict__ b, // (B, N, N, H) - layout is b,j,k,h\n scalar_t* __restrict__ out, // (B, N, N, H) - layout is b,i,j,h\n int B, int N, int H\n) {\n const int b = blockIdx.z;\n const int h = blockIdx.x;\n const int i = blockIdx.y / N;\n const int j_base = (blockIdx.y % N) * TILE_SIZE;\n const int tid = threadIdx.x;\n \n // Each block handles one (b, i, j_block, h)\n // Threads in block collaborate on different j values\n \n extern __shared__ char smem[];\n scalar_t* a_cache = reinterpret_cast<scalar_t*>(smem); // size N\n scalar_t* b_cache = a_cache + N; // size TILE_SIZE * N\n \n const int j = j_base + tid / N;\n const int k = tid % N;\n \n // Initialize output accumulator\n scalar_t sum = scalar_t(0);\n \n // Loop over k with tiling\n for (int k_tile_start = 0; k_tile_start < N; k_tile_start += 1) {\n // Load a[b, i, k, h] for all k\n if (tid < N) {\n a_cache[tid] = a[((b * N + i) * N + tid) * H + h];\n }\n \n // Load b[b, j, k, h] for j in our tile and all k\n if (tid < TILE_SIZE * N) {\n int load_j = j_base + tid / N;\n int load_k = tid % N;\n if (load_j < N) {\n b_cache[tid] = b[((b * N + load_j) * N + load_k) * H + h];\n } else {\n b_cache[tid] = scalar_t(0);\n }\n }\n __syncthreads();\n \n // Compute partial sum\n if (j < N) {\n for (int k_idx = 0; k_idx < N; k_idx++) {\n scalar_t a_val = a_cache[k_idx];\n scalar_t b_val = b_cache[(j - j_base) * N + k_idx];\n sum += a_val * b_val;\n }\n }\n }\n \n if (j < N) {\n out[((b * N + i) * N + j) * H + h] = sum;\n }\n}\n\n// Simpler, more efficient einsum kernel using warps\ntemplate <typename scalar_t, int BLOCK_H>\n__global__ void einsum_kernel_v2(\n const scalar_t* __restrict__ a,\n const scalar_t* __restrict__ b,\n scalar_t* __restrict__ out,\n int B, int N, int H\n) {\n // Each block handles (batch, i, j)\n // Each warp handles H_dim values\n \n const int batch = blockIdx.x;\n const int i_idx = blockIdx.y / N;\n const int j_idx = blockIdx.y % N;\n \n const int h_base = (blockIdx.z * BLOCK_H) + threadIdx.x;\n \n if (h_base >= H) return;\n \n // Compute sum over k for this (batch, i, j, h)\n scalar_t sum = scalar_t(0);\n \n for (int k = 0; k < N; k++) {\n scalar_t a_val = a[((batch * N + i_idx) * N + k) * H + h_base];\n scalar_t b_val = b[((batch * N + j_idx) * N + k) * H + h_base];\n sum += a_val * b_val;\n }\n \n out[((batch * N + i_idx) * N + j_idx) * H + h_base] = sum;\n}\n\n// Optimized einsum with vectorized loads\ntemplate <typename scalar_t, int VEC_SIZE>\n__global__ void einsum_kernel_vec(\n const scalar_t* __restrict__ a,\n const scalar_t* __restrict__ b,\n scalar_t* __restrict__ out,\n int B, int N, int H\n) {\n const int batch = blockIdx.x;\n const int i_idx = blockIdx.y / N;\n const int j_idx = blockIdx.y % N;\n const int h_vec = blockIdx.z * blockDim.x + threadIdx.x;\n \n const int h_base = h_vec * VEC_SIZE;\n if (h_base >= H) return;\n \n // Vectorized accumulator\n scalar_t sum[VEC_SIZE];\n #pragma unroll\n for (int v = 0; v < VEC_SIZE; v++) {\n sum[v] = scalar_t(0);\n }\n \n // Sum over k\n for (int k = 0; k < N; k++) {\n const scalar_t* a_ptr = &a[((batch * N + i_idx) * N + k) * H + h_base];\n const scalar_t* b_ptr = &b[((batch * N + j_idx) * N + k) * H + h_base];\n \n #pragma unroll\n for (int v = 0; v < VEC_SIZE; v++) {\n if (h_base + v < H) {\n sum[v] += a_ptr[v] * b_ptr[v];\n }\n }\n }\n \n // Write output\n scalar_t* out_ptr = &out[((batch * N + i_idx) * N + j_idx) * H + h_base];\n #pragma unroll\n for (int v = 0; v < VEC_SIZE; v++) {\n if (h_base + v < H) {\n out_ptr[v] = sum[v];\n }\n }\n}\n\n// Host functions\nstd::vector<torch::Tensor> triangular_update_cuda(\n torch::Tensor pair_repr,\n torch::Tensor pair_mask,\n torch::Tensor ln_in_gamma,\n torch::Tensor ln_in_beta,\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_gamma,\n torch::Tensor ln_out_beta,\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) {\n int B = pair_repr.size(0);\n int N = pair_repr.size(1);\n int pair_dim = pair_repr.size(3);\n int hidden_dim = proj_a_weight.size(1);\n \n auto options = pair_repr.options();\n \n // Allocate intermediate tensors\n auto a = torch::empty({B, N, N, hidden_dim}, options);\n auto b = torch::empty({B, N, N, hidden_dim}, options);\n auto einsum_out = torch::empty({B, N, N, hidden_dim}, options);\n auto output = torch::empty({B, N, N, pair_dim}, options);\n \n // Create mask tensor\n auto mask = pair_mask.unsqueeze(-1);\n \n // Fused projection + gate kernels for a and b\n const int threads_per_pos = hidden_dim;\n const int num_positions = N * N;\n const int block_y = (num_positions + 15) / 16;\n \n dim3 grid_a(B, block_y);\n dim3 block_a(hidden_dim, 16); // threads.x for hidden dim, threads.y for positions\n \n AT_DISPATCH_FLOATING_TYPES(pair_repr.scalar_type(), \"fused_proj_a\", ([&] {\n fused_proj_gate_kernel<scalar_t><<<grid_a, block_a>>>(\n pair_repr.data_ptr<scalar_t>(),\n proj_a_weight.data_ptr<scalar_t>(),\n proj_a_bias.data_ptr<scalar_t>(),\n gate_a_weight.data_ptr<scalar_t>(),\n gate_a_bias.data_ptr<scalar_t>(),\n ln_in_gamma.data_ptr<scalar_t>(),\n ln_in_beta.data_ptr<scalar_t>(),\n a.data_ptr<scalar_t>(),\n mask.data_ptr<scalar_t>(),\n B, N, pair_dim, hidden_dim,\n 1e-5f\n );\n }));\n \n AT_DISPATCH_FLOATING_TYPES(pair_repr.scalar_type(), \"fused_proj_b\", ([&] {\n fused_proj_gate_kernel<scalar_t><<<grid_a, block_a>>>(\n pair_repr.data_ptr<scalar_t>(),\n proj_b_weight.data_ptr<scalar_t>(),\n proj_b_bias.data_ptr<scalar_t>(),\n gate_b_weight.data_ptr<scalar_t>(),\n gate_b_bias.data_ptr<scalar_t>(),\n ln_in_gamma.data_ptr<scalar_t>(),\n ln_in_beta.data_ptr<scalar_t>(),\n b.data_ptr<scalar_t>(),\n mask.data_ptr<scalar_t>(),\n B, N, pair_dim, hidden_dim,\n 1e-5f\n );\n }));\n \n // Einsum kernel\n const int vec_size = 4;\n const int block_dim = 16; // threads per block for h dimension\n const int h_blocks = (hidden_dim + vec_size * block_dim - 1) / (vec_size * block_dim);\n \n dim3 grid_einsum(B, N * N, h_blocks);\n dim3 block_einsum(block_dim);\n \n AT_DISPATCH_FLOATING_TYPES(pair_repr.scalar_type(), \"einsum\", ([&] {\n einsum_kernel_vec<scalar_t, vec_size><<<grid_einsum, block_einsum>>>(\n a.data_ptr<scalar_t>(),\n b.data_ptr<scalar_t>(),\n einsum_out.data_ptr<scalar_t>(),\n B, N, hidden_dim\n );\n }));\n \n // Output kernel\n dim3 grid_out(B, block_y);\n dim3 block_out(pair_dim, 16);\n \n AT_DISPATCH_FLOATING_TYPES(pair_repr.scalar_type(), \"fused_output\", ([&] {\n fused_output_kernel<scalar_t><<<grid_out, block_out>>>(\n einsum_out.data_ptr<scalar_t>(),\n out_proj_weight.data_ptr<scalar_t>(),\n out_proj_bias.data_ptr<scalar_t>(),\n pair_repr.data_ptr<scalar_t>(),\n out_gate_weight.data_ptr<scalar_t>(),\n out_gate_bias.data_ptr<scalar_t>(),\n ln_out_gamma.data_ptr<scalar_t>(),\n ln_out_beta.data_ptr<scalar_t>(),\n output.data_ptr<scalar_t>(),\n B, N, hidden_dim, pair_dim,\n 1e-5f\n );\n }));\n \n return {output, a, b, einsum_out};\n}\n\ntorch::Tensor triangular_update_out(\n torch::Tensor pair_repr,\n torch::Tensor pair_mask,\n torch::Tensor ln_in_gamma,\n torch::Tensor ln_in_beta,\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_gamma,\n torch::Tensor ln_out_beta,\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) {\n auto results = triangular_update_cuda(\n pair_repr, pair_mask,\n ln_in_gamma, ln_in_beta,\n proj_a_weight, proj_a_bias, gate_a_weight, gate_a_bias,\n proj_b_weight, proj_b_bias, gate_b_weight, gate_b_bias,\n ln_out_gamma, ln_out_beta,\n out_proj_weight, out_proj_bias, out_gate_weight, out_gate_bias\n );\n return results[0];\n}\n'''\n\ncpp_source = r'''\nstd::vector<torch::Tensor> triangular_update_cuda(\n torch::Tensor pair_repr,\n torch::Tensor pair_mask,\n torch::Tensor ln_in_gamma,\n torch::Tensor ln_in_beta,\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_gamma,\n torch::Tensor ln_out_beta,\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);\n\ntorch::Tensor triangular_update_out(\n torch::Tensor pair_repr,\n torch::Tensor pair_mask,\n torch::Tensor ln_in_gamma,\n torch::Tensor ln_in_beta,\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_gamma,\n torch::Tensor ln_out_beta,\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);\n'''\n\n# Load the extension\ntriangular_ext = load_inline(\n name='triangular_update_ext',\n cpp_sources=cpp_source,\n cuda_sources=cuda_source,\n functions=['triangular_update_out'],\n extra_cuda_cflags=['-O3', '-use_fast_math'],\n verbose=True\n)\n\n\nclass ModelNew(nn.Module):\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 self.pair_dim = pair_dim\n self.hidden_dim = hidden_dim\n \n def forward(self, pair_repr: torch.Tensor, pair_mask: torch.Tensor) -> torch.Tensor:\n # Use custom CUDA kernel\n output = triangular_ext.triangular_update_out(\n pair_repr,\n pair_mask,\n self.layer_norm_in.weight,\n self.layer_norm_in.bias,\n self.proj_a.weight,\n self.proj_a.bias,\n self.gate_a.weight,\n self.gate_a.bias,\n self.proj_b.weight,\n self.proj_b.bias,\n self.gate_b.weight,\n self.gate_b.bias,\n self.layer_norm_out.weight,\n self.layer_norm_out.bias,\n self.out_proj.weight,\n self.out_proj.bias,\n self.out_gate.weight,\n self.out_gate.bias\n )\n return output"
}