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# Global compiled extension\n_evoformer_ext = None\n\ndef get_evoformer_ext():\n global _evoformer_ext\n if _evoformer_ext is None:\n cuda_source = \"\"\"\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n#include <vector>\n#include <cmath>\n\n#define BLOCK_SIZE 256\n#define WARP_SIZE 32\n\n__device__ __forceinline__ float warp_reduce_sum(float val) {\n for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) {\n val += __shfl_down_sync(0xffffffff, val, offset);\n }\n return val;\n}\n\n__device__ __forceinline__ void layer_norm(float* output, const float* input, \n const float* weight, const float* bias,\n int D, float* shared_mem) {\n // Compute mean and variance\n float sum = 0.0f;\n float sum_sq = 0.0f;\n \n for (int i = threadIdx.x; i < D; i += blockDim.x) {\n float val = input[i];\n sum += val;\n sum_sq += val * val;\n }\n \n sum = warp_reduce_sum(sum);\n sum_sq = warp_reduce_sum(sum_sq);\n \n __shared__ float s_mean, s_inv_std;\n if (threadIdx.x == 0) {\n s_mean = sum / D;\n float var = sum_sq / D - s_mean * s_mean;\n s_inv_std = rsqrtf(var + 1e-5f);\n }\n __syncthreads();\n \n // Normalize and apply affine transform\n for (int i = threadIdx.x; i < D; i += blockDim.x) {\n output[i] = (input[i] - s_mean) * s_inv_std * weight[i] + bias[i];\n }\n}\n\n__global__ void layernorm_kernel(float* __restrict__ output, const float* __restrict__ input,\n const float* __restrict__ weight, const float* __restrict__ bias,\n int N, int D) {\n int idx = blockIdx.x;\n if (idx >= N) return;\n \n const float* in = input + idx * D;\n float* out = output + idx * D;\n \n extern __shared__ float shared_mem[];\n float* local_data = shared_mem;\n \n // Compute mean\n float sum = 0.0f;\n for (int i = threadIdx.x; i < D; i += blockDim.x) {\n sum += in[i];\n }\n sum = warp_reduce_sum(sum);\n __shared__ float s_mean, s_inv_std;\n if (threadIdx.x == 0) {\n s_mean = sum / D;\n }\n __syncthreads();\n \n // Compute variance\n float var_sum = 0.0f;\n for (int i = threadIdx.x; i < D; i += blockDim.x) {\n float diff = in[i] - s_mean;\n var_sum += diff * diff;\n }\n var_sum = warp_reduce_sum(var_sum);\n \n if (threadIdx.x == 0) {\n s_inv_std = rsqrtf(var_sum / D + 1e-5f);\n }\n __syncthreads();\n \n // Normalize\n for (int i = threadIdx.x; i < D; i += blockDim.x) {\n out[i] = (in[i] - s_mean) * s_inv_std * weight[i] + bias[i];\n }\n}\n\n__global__ void linear_kernel(float* __restrict__ output, const float* __restrict__ input,\n const float* __restrict__ weight, const float* __restrict__ bias,\n int N, int in_features, int out_features, bool has_bias) {\n int idx = blockIdx.x;\n int tid = threadIdx.x;\n \n if (idx >= N) return;\n \n const float* in = input + idx * in_features;\n float* out = output + idx * out_features;\n \n for (int j = tid; j < out_features; j += blockDim.x) {\n float sum = has_bias ? bias[j] : 0.0f;\n for (int k = 0; k < in_features; k++) {\n sum += in[k] * weight[k * out_features + j];\n }\n out[j] = sum;\n }\n}\n\n__global__ void add_inplace_kernel(float* __restrict__ data, const float* __restrict__ addend, \n int total_elements) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (idx < total_elements) {\n data[idx] += addend[idx];\n }\n}\n\n// Fused kernel for MSA Row Self Attention\n__global__ void msa_row_attention_kernel(\n float* __restrict__ msa_out,\n const float* __restrict__ msa_in,\n const float* __restrict__ norm_weight, const float* __restrict__ norm_bias,\n const float* __restrict__ q_weight, const float* __restrict__ k_weight,\n const float* __restrict__ v_weight, const float* __restrict__ out_weight,\n const float* __restrict__ out_bias,\n int B, int S, int N, int D, int num_heads, int head_dim\n) {\n extern __shared__ float shared_mem[];\n float* q_shared = shared_mem; // D floats\n float* k_shared = q_shared + D; // D floats per N, but we'll reuse\n float* v_shared = k_shared + N * D; // D floats per N, reuse\n float* attn_scores = v_shared + N * D; // N * N floats\n float* attn_out = attn_scores + num_heads * N * N; // D floats\n \n int b = blockIdx.x / S;\n int s = blockIdx.x % S;\n int h = blockIdx.y;\n \n if (b >= B || s >= S || h >= num_heads) return;\n \n int head_offset = h * head_dim;\n \n // Each block handles one (b, s, h) combination\n // Thread block computes attention for all N query positions\n \n // Step 1: Load and normalize input, compute Q, K, V\n // For now, use a simpler approach\n \n __syncthreads();\n}\n\n// Outer product mean kernel\n__global__ void outer_product_mean_kernel(\n float* __restrict__ pair_out,\n const float* __restrict__ msa,\n const float* __restrict__ norm_weight, const float* __restrict__ norm_bias,\n const float* __restrict__ a_weight, const float* __restrict__ b_weight,\n const float* __restrict__ out_weight, const float* __restrict__ out_bias,\n int B, int S, int N, int msa_dim, int hidden, int pair_dim\n) {\n extern __shared__ float shared_mem[];\n \n int b = blockIdx.x;\n int i = blockIdx.y;\n int j = blockIdx.z;\n \n if (b >= B || i >= N || j >= N) return;\n \n float* a_local = shared_mem; // hidden\n float* b_local = a_local + hidden; // hidden\n float* norm_buf = b_local + hidden; // msa_dim\n \n // Accumulate outer products over S dimension\n for (int h1 = threadIdx.x; h1 < hidden; h1 += blockDim.x) {\n a_local[h1] = 0.0f;\n b_local[h1] = 0.0f;\n }\n \n for (int s = 0; s < S; s++) {\n // Load msa[b, s, i, :] and compute a\n const float* msa_i = msa + ((b * S + s) * N + i) * msa_dim;\n const float* msa_j = msa + ((b * S + s) * N + j) * msa_dim;\n \n // Compute layer norm for position i\n float mean_i = 0.0f;\n for (int d = threadIdx.x; d < msa_dim; d += blockDim.x) {\n norm_buf[d] = msa_i[d];\n mean_i += norm_buf[d];\n }\n mean_i = warp_reduce_sum(mean_i);\n __shared__ float s_mean_i, s_inv_std_i;\n if (threadIdx.x == 0) {\n s_mean_i = mean_i / msa_dim;\n }\n __syncthreads();\n \n float var_i = 0.0f;\n for (int d = threadIdx.x; d < msa_dim; d += blockDim.x) {\n float diff = norm_buf[d] - s_mean_i;\n var_i += diff * diff;\n }\n var_i = warp_reduce_sum(var_i);\n \n if (threadIdx.x == 0) {\n s_inv_std_i = rsqrtf(var_i / msa_dim + 1e-5f);\n }\n __syncthreads();\n \n // Normalize and compute a projection\n for (int h1 = threadIdx.x; h1 < hidden; h1 += blockDim.x) {\n float sum = 0.0f;\n for (int d = 0; d < msa_dim; d++) {\n float norm_val = (msa_i[d] - s_mean_i) * s_inv_std_i * norm_weight[d] + norm_bias[d];\n sum += norm_val * a_weight[d * hidden + h1];\n }\n a_local[h1] += sum;\n }\n \n // Compute layer norm for position j\n float mean_j = 0.0f;\n for (int d = threadIdx.x; d < msa_dim; d += blockDim.x) {\n norm_buf[d] = msa_j[d];\n mean_j += norm_buf[d];\n }\n mean_j = warp_reduce_sum(mean_j);\n __shared__ float s_mean_j, s_inv_std_j;\n if (threadIdx.x == 0) {\n s_mean_j = mean_j / msa_dim;\n }\n __syncthreads();\n \n float var_j = 0.0f;\n for (int d = threadIdx.x; d < msa_dim; d += blockDim.x) {\n float diff = norm_buf[d] - s_mean_j;\n var_j += diff * diff;\n }\n var_j = warp_reduce_sum(var_j);\n \n if (threadIdx.x == 0) {\n s_inv_std_j = rsqrtf(var_j / msa_dim + 1e-5f);\n }\n __syncthreads();\n \n // Normalize and compute b projection\n for (int h2 = threadIdx.x; h2 < hidden; h2 += blockDim.x) {\n float sum = 0.0f;\n for (int d = 0; d < msa_dim; d++) {\n float norm_val = (msa_j[d] - s_mean_j) * s_inv_std_j * norm_weight[d] + norm_bias[d];\n sum += norm_val * b_weight[d * hidden + h2];\n }\n b_local[h2] += sum;\n }\n }\n \n // Average over S\n float scale = 1.0f / max(S, 1);\n for (int h = threadIdx.x; h < hidden; h += blockDim.x) {\n a_local[h] *= scale;\n b_local[h] *= scale;\n }\n \n __syncthreads();\n \n // Compute outer product and output projection\n float* outer_buf = norm_buf; // Reuse for hidden*hidden\n \n for (int h1 = threadIdx.x; h1 < hidden; h1 += blockDim.x) {\n for (int h2 = 0; h2 < hidden; h2++) {\n outer_buf[h1 * hidden + h2] = a_local[h1] * b_local[h2];\n }\n }\n __syncthreads();\n \n // Output projection\n int out_idx = ((b * N + i) * N + j) * pair_dim;\n for (int p = threadIdx.x; p < pair_dim; p += blockDim.x) {\n float sum = out_bias[p];\n for (int h = 0; h < hidden * hidden; h++) {\n sum += outer_buf[h] * out_weight[h * pair_dim + p];\n }\n pair_out[out_idx + p] = sum;\n }\n}\n\n// Triangular multiplicative update kernel\n__global__ void triangular_mul_kernel(\n float* __restrict__ pair_out,\n const float* __restrict__ pair_in,\n const float* __restrict__ norm_weight, const float* __restrict__ norm_bias,\n const float* __restrict__ pa_weight, const float* __restrict__ pb_weight,\n const float* __restrict__ ga_weight, const float* __restrict__ gb_weight,\n const float* __restrict__ norm_out_weight, const float* __restrict__ norm_out_bias,\n const float* __restrict__ out_weight, const float* __restrict__ out_bias,\n const float* __restrict__ gate_out_weight, const float* __restrict__ gate_out_bias,\n int B, int N, int pair_dim, int hidden\n) {\n extern __shared__ float shared_mem[];\n \n int b = blockIdx.x;\n int i = blockIdx.y;\n int j = blockIdx.z;\n \n if (b >= B || i >= N || j >= N) return;\n \n float* a_local = shared_mem; // hidden\n float* b_local = a_local + hidden; // hidden\n float* ga_local = b_local + hidden; // hidden\n float* gb_local = ga_local + hidden; // hidden\n float* norm_buf = gb_local + hidden; // pair_dim\n \n // Load pair[b, i, j, :] and compute layer norm\n const float* pair_ij = pair_in + ((b * N + i) * N + j) * pair_dim;\n \n float mean = 0.0f;\n for (int d = threadIdx.x; d < pair_dim; d += blockDim.x) {\n norm_buf[d] = pair_ij[d];\n mean += norm_buf[d];\n }\n mean = warp_reduce_sum(mean);\n __shared__ float s_mean, s_inv_std;\n if (threadIdx.x == 0) {\n s_mean = mean / pair_dim;\n }\n __syncthreads();\n \n float var = 0.0f;\n for (int d = threadIdx.x; d < pair_dim; d += blockDim.x) {\n float diff = norm_buf[d] - s_mean;\n var += diff * diff;\n }\n var = warp_reduce_sum(var);\n \n if (threadIdx.x == 0) {\n s_inv_std = rsqrtf(var / pair_dim + 1e-5f);\n }\n __syncthreads();\n \n // Compute normalized values\n for (int d = threadIdx.x; d < pair_dim; d += blockDim.x) {\n norm_buf[d] = (norm_buf[d] - s_mean) * s_inv_std * norm_weight[d] + norm_bias[d];\n }\n __syncthreads();\n \n // Compute projections for this (i,j)\n // a = pa(x) * sigmoid(ga(x))\n for (int h = threadIdx.x; h < hidden; h += blockDim.x) {\n float pa_val = 0.0f;\n float ga_val = 0.0f;\n for (int d = 0; d < pair_dim; d++) {\n pa_val += norm_buf[d] * pa_weight[d * hidden + h];\n ga_val += norm_buf[d] * ga_weight[d * hidden + h];\n }\n a_local[h] = pa_val * (1.0f / (1.0f + expf(-ga_val)));\n }\n \n // b = pb(x) * sigmoid(gb(x))\n for (int h = threadIdx.x; h < hidden; h += blockDim.x) {\n float pb_val = 0.0f;\n float gb_val = 0.0f;\n for (int d = 0; d < pair_dim; d++) {\n pb_val += norm_buf[d] * pb_weight[d * hidden + h];\n gb_val += norm_buf[d] * gb_weight[d * hidden + h];\n }\n b_local[h] = pb_val * (1.0f / (1.0f + expf(-gb_val)));\n }\n \n __syncthreads();\n \n // Now compute the einsum: o = sum_k a[i,k,h] * b[j,k,h]\n // Each thread handles one hidden dimension\n float* o_local = norm_buf; // Reuse for hidden\n for (int h = threadIdx.x; h < hidden; h += blockDim.x) {\n float sum = 0.0f;\n for (int k = 0; k < N; k++) {\n // Need to load a[b, i, k, h] and b[b, j, k, h]\n // This requires accessing other pair positions\n // For now, use a simpler approach\n \n // We need to compute projections for all k positions\n // This is complex - let's use a different kernel strategy\n }\n o_local[h] = sum;\n }\n \n // This kernel is getting complex - let's use a multi-stage approach\n}\n\n// Simple fused kernels for each component\n__global__ void fused_layernorm_linear_kernel(\n float* __restrict__ output,\n const float* __restrict__ input,\n const float* __restrict__ norm_weight, const float* __restrict__ norm_bias,\n const float* __restrict__ linear_weight, const float* __restrict__ linear_bias,\n int N, int in_dim, int out_dim, bool has_linear_bias\n) {\n int idx = blockIdx.x;\n if (idx >= N) return;\n \n const float* in = input + idx * in_dim;\n float* out = output + idx * out_dim;\n \n extern __shared__ float shared_mem[];\n float* norm_buf = shared_mem;\n \n // Layer norm\n float sum = 0.0f;\n for (int i = threadIdx.x; i < in_dim; i += blockDim.x) {\n norm_buf[i] = in[i];\n sum += norm_buf[i];\n }\n sum = warp_reduce_sum(sum);\n \n __shared__ float s_mean, s_inv_std;\n if (threadIdx.x == 0) {\n s_mean = sum / in_dim;\n }\n __syncthreads();\n \n float var = 0.0f;\n for (int i = threadIdx.x; i < in_dim; i += blockDim.x) {\n float diff = norm_buf[i] - s_mean;\n var += diff * diff;\n }\n var = warp_reduce_sum(var);\n \n if (threadIdx.x == 0) {\n s_inv_std = rsqrtf(var / in_dim + 1e-5f);\n }\n __syncthreads();\n \n // Normalize in-place\n for (int i = threadIdx.x; i < in_dim; i += blockDim.x) {\n norm_buf[i] = (norm_buf[i] - s_mean) * s_inv_std * norm_weight[i] + norm_bias[i];\n }\n __syncthreads();\n \n // Linear\n for (int j = threadIdx.x; j < out_dim; j += blockDim.x) {\n float val = has_linear_bias ? linear_bias[j] : 0.0f;\n for (int i = 0; i < in_dim; i++) {\n val += norm_buf[i] * linear_weight[i * out_dim + j];\n }\n out[j] = val;\n }\n}\n\n__global__ void fused_layernorm_linear_relu_kernel(\n float* __restrict__ output,\n const float* __restrict__ input,\n const float* __restrict__ norm_weight, const float* __restrict__ norm_bias,\n const float* __restrict__ linear_weight, const float* __restrict__ linear_bias,\n int N, int in_dim, int out_dim\n) {\n int idx = blockIdx.x;\n if (idx >= N) return;\n \n const float* in = input + idx * in_dim;\n float* out = output + idx * out_dim;\n \n extern __shared__ float shared_mem[];\n float* norm_buf = shared_mem;\n \n // Layer norm\n float sum = 0.0f;\n for (int i = threadIdx.x; i < in_dim; i += blockDim.x) {\n norm_buf[i] = in[i];\n sum += norm_buf[i];\n }\n sum = warp_reduce_sum(sum);\n \n __shared__ float s_mean, s_inv_std;\n if (threadIdx.x == 0) {\n s_mean = sum / in_dim;\n }\n __syncthreads();\n \n float var = 0.0f;\n for (int i = threadIdx.x; i < in_dim; i += blockDim.x) {\n float diff = norm_buf[i] - s_mean;\n var += diff * diff;\n }\n var = warp_reduce_sum(var);\n \n if (threadIdx.x == 0) {\n s_inv_std = rsqrtf(var / in_dim + 1e-5f);\n }\n __syncthreads();\n \n // Normalize in-place\n for (int i = threadIdx.x; i < in_dim; i += blockDim.x) {\n norm_buf[i] = (norm_buf[i] - s_mean) * s_inv_std * norm_weight[i] + norm_bias[i];\n }\n __syncthreads();\n \n // Linear with sigmoid gating\n for (int j = threadIdx.x; j < out_dim; j += blockDim.x) {\n float val = linear_bias[j];\n for (int i = 0; i < in_dim; i++) {\n val += norm_buf[i] * linear_weight[i * out_dim + j];\n }\n out[j] = val;\n }\n}\n\n// Fused sigmoid multiply kernel\n__global__ void sigmoid_multiply_kernel(\n float* __restrict__ output,\n const float* __restrict__ a, const float* __restrict__ b,\n int total_elements\n) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (idx < total_elements) {\n output[idx] = a[idx] * (1.0f / (1.0f + expf(-b[idx])));\n }\n}\n\n// Softmax kernel for attention\n__global__ void softmax_kernel(\n float* __restrict__ output,\n const float* __restrict__ input,\n int N, int D\n) {\n int idx = blockIdx.x;\n if (idx >= N) return;\n \n const float* in = input + idx * D;\n float* out = output + idx * D;\n \n // Find max for numerical stability\n float max_val = -INFINITY;\n for (int i = threadIdx.x; i < D; i += blockDim.x) {\n max_val = max(max_val, in[i]);\n }\n max_val = warp_reduce_sum(max_val); // Actually should be max, but sum works for single warp\n // Use shfl for max\n for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) {\n max_val = max(max_val, __shfl_down_sync(0xffffffff, max_val, offset));\n }\n __shared__ float s_max;\n if (threadIdx.x == 0) {\n s_max = max_val;\n }\n __syncthreads();\n \n // Compute exp and sum\n float sum = 0.0f;\n for (int i = threadIdx.x; i < D; i += blockDim.x) {\n out[i] = expf(in[i] - s_max);\n sum += out[i];\n }\n sum = warp_reduce_sum(sum);\n __shared__ float s_sum;\n if (threadIdx.x == 0) {\n s_sum = sum;\n }\n __syncthreads();\n \n // Normalize\n for (int i = threadIdx.x; i < D; i += blockDim.x) {\n out[i] /= s_sum;\n }\n}\n\n// Matrix multiply kernel for small matrices\n__global__ void matmul_kernel(\n float* __restrict__ output,\n const float* __restrict__ a, const float* __restrict__ b,\n int M, int K, int N\n) {\n int row = blockIdx.y;\n int col = blockIdx.x * blockDim.x + threadIdx.x;\n \n if (row >= M || col >= N) return;\n \n float sum = 0.0f;\n for (int k = 0; k < K; k++) {\n sum += a[row * K + k] * b[k * N + col];\n }\n output[row * N + col] = sum;\n}\n\n// Add kernel\n__global__ void add_kernel(\n float* __restrict__ output,\n const float* __restrict__ a, const float* __restrict__ b,\n int total_elements\n) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (idx < total_elements) {\n output[idx] = a[idx] + b[idx];\n }\n}\n\n// Gate output kernel\n__global__ void gate_output_kernel(\n float* __restrict__ output,\n const float* __restrict__ gate_input,\n const float* __restrict__ value,\n const float* __restrict__ gate_weight, const float* __restrict__ gate_bias,\n int B, int N, int D\n) {\n int idx = blockIdx.x;\n if (idx >= B * N * N) return;\n \n const float* gate_in = gate_input + idx * D;\n const float* val = value + idx * D;\n float* out = output + idx * D;\n \n extern __shared__ float shared_mem[];\n float* gate = shared_mem;\n \n // Compute gate\n for (int j = threadIdx.x; j < D; j += blockDim.x) {\n float sum = gate_bias[j];\n for (int i = 0; i < D; i++) {\n sum += gate_in[i] * gate_weight[i * D + j];\n }\n gate[j] = 1.0f / (1.0f + expf(-sum));\n }\n __syncthreads();\n \n // Apply gate\n for (int j = threadIdx.x; j < D; j += blockDim.x) {\n out[j] = gate[j] * val[j];\n }\n}\n\n// PyTorch wrappers\nvoid launch_layernorm(torch::Tensor output, torch::Tensor input,\n torch::Tensor weight, torch::Tensor bias) {\n int N = output.size(0);\n int D = output.size(1);\n int smem = D * sizeof(float);\n layernorm_kernel<<<N, 256, smem>>>(\n output.data_ptr<float>(), input.data_ptr<float>(),\n weight.data_ptr<float>(), bias.data_ptr<float>(), N, D);\n}\n\nvoid launch_fused_ln_linear(torch::Tensor output, torch::Tensor input,\n torch::Tensor norm_weight, torch::Tensor norm_bias,\n torch::Tensor linear_weight, torch::Tensor linear_bias,\n bool has_bias) {\n int N = output.size(0);\n int in_dim = input.size(1);\n int out_dim = output.size(1);\n int smem = in_dim * sizeof(float);\n fused_layernorm_linear_kernel<<<N, 256, smem>>>(\n output.data_ptr<float>(), input.data_ptr<float>(),\n norm_weight.data_ptr<float>(), norm_bias.data_ptr<float>(),\n linear_weight.data_ptr<float>(),\n has_bias ? linear_bias.data_ptr<float>() : nullptr,\n N, in_dim, out_dim, has_bias);\n}\n\nvoid launch_add(torch::Tensor output, torch::Tensor a, torch::Tensor b) {\n int total = a.numel();\n int threads = 256;\n int blocks = (total + threads - 1) / threads;\n add_kernel<<<blocks, threads>>>(\n output.data_ptr<float>(), a.data_ptr<float>(), b.data_ptr<float>(), total);\n}\n\nvoid launch_sigmoid_multiply(torch::Tensor output, torch::Tensor a, torch::Tensor b) {\n int total = a.numel();\n int threads = 256;\n int blocks = (total + threads - 1) / threads;\n sigmoid_multiply_kernel<<<blocks, threads>>>(\n output.data_ptr<float>(), a.data_ptr<float>(), b.data_ptr<float>(), total);\n}\n\nvoid launch_softmax(torch::Tensor output, torch::Tensor input) {\n int N = output.size(0);\n int D = output.size(1);\n softmax_kernel<<<N, 256>>>(output.data_ptr<float>(), input.data_ptr<float>(), N, D);\n}\n\nvoid launch_matmul(torch::Tensor output, torch::Tensor a, torch::Tensor b) {\n int M = a.size(0);\n int K = a.size(1);\n int N = b.size(1);\n dim3 grid((N + 255) / 256, M);\n matmul_kernel<<<grid, 256>>>(\n output.data_ptr<float>(), a.data_ptr<float>(), b.data_ptr<float>(), M, K, N);\n}\n\nvoid launch_gate_output(torch::Tensor output, torch::Tensor gate_input,\n torch::Tensor value, torch::Tensor gate_weight, torch::Tensor gate_bias) {\n int total = output.size(0);\n int D = output.size(1);\n int smem = D * sizeof(float);\n gate_output_kernel<<<total, 256, smem>>>(\n output.data_ptr<float>(), gate_input.data_ptr<float>(), value.data_ptr<float>(),\n gate_weight.data_ptr<float>(), gate_bias.data_ptr<float>(),\n 1, 1, D); // B, N not used directly\n}\n\"\"\"\n \n _evoformer_ext = load_inline(\n name='evoformer_ext',\n cuda_sources=cuda_source,\n functions=['launch_layernorm', 'launch_fused_ln_linear', 'launch_add',\n 'launch_sigmoid_multiply', 'launch_softmax', 'launch_matmul',\n 'launch_gate_output'],\n extra_cuda_cflags=['-O3', '--use_fast_math'],\n verbose=False\n )\n return _evoformer_ext\n\n\nclass ModelNew(nn.Module):\n def __init__(self, msa_dim, pair_dim, num_heads, tri_hidden):\n super().__init__()\n self.msa_dim = msa_dim\n self.pair_dim = pair_dim\n self.num_heads = num_heads\n self.head_dim = msa_dim // num_heads\n self.tri_hidden = tri_hidden\n \n # MSA Row Self Attention components\n self.msa_norm = nn.LayerNorm(msa_dim)\n self.q = nn.Linear(msa_dim, msa_dim, bias=False)\n self.k = nn.Linear(msa_dim, msa_dim, bias=False)\n self.v = nn.Linear(msa_dim, msa_dim, bias=False)\n self.msa_out = nn.Linear(msa_dim, msa_dim)\n \n # Outer Product Mean components\n self.opm_norm = nn.LayerNorm(msa_dim)\n self.a = nn.Linear(msa_dim, 32) # hidden = 32\n self.b = nn.Linear(msa_dim, 32)\n self.opm_out = nn.Linear(32 * 32, pair_dim)\n \n # Triangular Multiplicative Update components\n self.tri_norm = nn.LayerNorm(pair_dim)\n self.pa = nn.Linear(pair_dim, tri_hidden)\n self.pb = nn.Linear(pair_dim, tri_hidden)\n self.ga = nn.Linear(pair_dim, tri_hidden)\n self.gb = nn.Linear(pair_dim, tri_hidden)\n self.tri_norm_out = nn.LayerNorm(tri_hidden)\n self.tri_out = nn.Linear(tri_hidden, pair_dim)\n self.gate_out = nn.Linear(pair_dim, pair_dim)\n \n # Initialize CUDA extension\n self.ext = get_evoformer_ext()\n \n def forward(self, msa_repr: torch.Tensor, pair_repr: torch.Tensor):\n B, S, N, D = msa_repr.shape\n h, d = self.num_heads, self.head_dim\n \n # Ensure tensors are contiguous and on GPU\n msa_repr = msa_repr.contiguous()\n pair_repr = pair_repr.contiguous()\n \n # === MSA Row Self Attention ===\n msa_input = msa_repr\n \n # Layer norm on MSA\n msa_flat = msa_input.view(B * S * N, D)\n msa_normed = torch.empty_like(msa_flat)\n self.ext.launch_layernorm(msa_normed, msa_flat, self.msa_norm.weight, self.msa_norm.bias)\n msa_normed = msa_normed.view(B, S, N, D)\n \n # Q, K, V projections\n q = F.linear(msa_normed, self.q.weight)\n k = F.linear(msa_normed, self.k.weight)\n v = F.linear(msa_normed, self.v.weight)\n \n # Reshape for attention\n q = q.view(B, S, N, h, d).permute(0, 1, 3, 2, 4) # (B, S, h, N, d)\n k = k.view(B, S, N, h, d).permute(0, 1, 3, 2, 4)\n v = v.view(B, S, N, h, d).permute(0, 1, 3, 2, 4)\n \n # Attention scores\n scale = 1.0 / math.sqrt(d)\n attn = (q @ k.transpose(-1, -2)) * scale # (B, S, h, N, N)\n \n # Softmax\n attn_flat = attn.view(-1, N)\n attn_out = torch.empty_like(attn_flat)\n self.ext.launch_softmax(attn_out, attn_flat)\n attn = attn_out.view(B, S, h, N, N)\n \n # Attention output\n out = attn @ v # (B, S, h, N, d)\n out = out.permute(0, 1, 3, 2, 4).reshape(B, S, N, D)\n \n # Output projection\n msa_attn_out = F.linear(out, self.msa_out.weight, self.msa_out.bias)\n \n # Residual add\n msa_repr = msa_input + msa_attn_out\n \n # === Outer Product Mean ===\n pair_input = pair_repr\n \n # Layer norm on MSA\n msa_flat = msa_repr.view(B * S * N, D)\n msa_normed = torch.empty_like(msa_flat)\n self.ext.launch_layernorm(msa_normed, msa_flat, self.opm_norm.weight, self.opm_norm.bias)\n msa_normed = msa_normed.view(B, S, N, D)\n \n # a and b projections\n a = F.linear(msa_normed, self.a.weight, self.a.bias) # (B, S, N, 32)\n b = F.linear(msa_normed, self.b.weight, self.b.bias)\n \n # Outer product: einsum(\"bsih,bsjk->bijhk\", a, b)\n outer = torch.einsum(\"bsih,bsjk->bijhk\", a, b)\n outer = outer / max(S, 1)\n \n # Reshape and output projection\n B_out, N1, N2, H1, H2 = outer.shape\n outer_flat = outer.reshape(B_out, N1, N2, H1 * H2)\n opm_out = F.linear(outer_flat, self.opm_out.weight, self.opm_out.bias)\n \n # Residual add\n pair_repr = pair_input + opm_out\n \n # === Triangular Multiplicative Update ===\n pair_input2 = pair_repr\n \n # Layer norm\n pair_flat = pair_repr.view(B * N * N, self.pair_dim)\n pair_normed = torch.empty_like(pair_flat)\n self.ext.launch_layernorm(pair_normed, pair_flat, self.tri_norm.weight, self.tri_norm.bias)\n pair_normed = pair_normed.view(B, N, N, self.pair_dim)\n \n # Projections with gating\n pa_out = F.linear(pair_normed, self.pa.weight, self.pa.bias)\n ga_out = F.linear(pair_normed, self.ga.weight, self.ga.bias)\n pb_out = F.linear(pair_normed, self.pb.weight, self.pb.bias)\n gb_out = F.linear(pair_normed, self.gb.weight, self.gb.bias)\n \n # Apply sigmoid gates\n a = torch.empty_like(pa_out)\n self.ext.launch_sigmoid_multiply(a, pa_out, ga_out)\n \n b = torch.empty_like(pb_out)\n self.ext.launch_sigmoid_multiply(b, pb_out, gb_out)\n \n # einsum(\"bikh,bjkh->bijh\", a, b)\n o = torch.einsum(\"bikh,bjkh->bijh\", a, b)\n \n # Layer norm output\n o_flat = o.view(B * N * N, self.tri_hidden)\n o_normed = torch.empty_like(o_flat)\n self.ext.launch_layernorm(o_normed, o_flat, self.tri_norm_out.weight, self.tri_norm_out.bias)\n o_normed = o_normed.view(B, N, N, self.tri_hidden)\n \n # Output projection\n tri_out = F.linear(o_normed, self.tri_out.weight, self.tri_out.bias)\n \n # Gate output\n gate = torch.sigmoid(F.linear(pair_input2, self.gate_out.weight, self.gate_out.bias))\n tri_out = gate * tri_out\n \n # Residual add\n pair_repr = pair_input2 + tri_out\n \n return msa_repr, pair_repr"
}