call → compile_kernel
callcompile_kernel
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\n\n_cpp_src = r\"\"\"\n#include <torch/extension.h>\n\ntorch::Tensor tri_mul_forward(torch::Tensor pair_repr, torch::Tensor pair_mask,\n torch::Tensor ln_in_w, torch::Tensor ln_in_b,\n torch::Tensor proj_a_w, torch::Tensor proj_a_b,\n torch::Tensor gate_a_w, torch::Tensor gate_a_b,\n torch::Tensor proj_b_w, torch::Tensor proj_b_b,\n torch::Tensor gate_b_w, torch::Tensor gate_b_b,\n torch::Tensor ln_out_w, torch::Tensor ln_out_b,\n torch::Tensor out_proj_w, torch::Tensor out_proj_b,\n torch::Tensor out_gate_w, torch::Tensor out_gate_b);\n\"\"\"\n\n_cuda_src = r\"\"\"\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n#include <math.h>\n\n#define PD 64\n#define HD 32\n\n__device__ __forceinline__ float sigmoidf_fast(float x) {\n return 1.0f / (1.0f + expf(-x));\n}\n\n__global__ void compute_ab_kernel(const float* __restrict__ pair,\n const float* __restrict__ mask,\n const float* __restrict__ ln_w,\n const float* __restrict__ ln_b,\n const float* __restrict__ pa_w,\n const float* __restrict__ pa_b,\n const float* __restrict__ ga_w,\n const float* __restrict__ ga_b,\n const float* __restrict__ pb_w,\n const float* __restrict__ pb_b,\n const float* __restrict__ gb_w,\n const float* __restrict__ gb_b,\n float* __restrict__ a,\n float* __restrict__ b,\n int total_edges) {\n int edge = blockIdx.x;\n int t = threadIdx.x;\n if (edge >= total_edges) return;\n __shared__ float sx[PD];\n __shared__ float red[64];\n const float* xptr = pair + edge * PD;\n float v = 0.0f;\n if (t < PD) { v = xptr[t]; red[t] = v; }\n __syncthreads();\n for (int stride = 32; stride > 0; stride >>= 1) {\n if (t < stride) red[t] += red[t + stride];\n __syncthreads();\n }\n float mean = red[0] * (1.0f / 64.0f);\n float diff = 0.0f;\n if (t < PD) { diff = v - mean; red[t] = diff * diff; }\n __syncthreads();\n for (int stride = 32; stride > 0; stride >>= 1) {\n if (t < stride) red[t] += red[t + stride];\n __syncthreads();\n }\n float inv = rsqrtf(red[0] * (1.0f / 64.0f) + 1.0e-5f);\n if (t < PD) sx[t] = diff * inv * ln_w[t] + ln_b[t];\n __syncthreads();\n\n if (t < HD) {\n float spa = pa_b[t], sga = ga_b[t], spb = pb_b[t], sgb = gb_b[t];\n const float* paw = pa_w + t * PD;\n const float* gaw = ga_w + t * PD;\n const float* pbw = pb_w + t * PD;\n const float* gbw = gb_w + t * PD;\n #pragma unroll\n for (int d = 0; d < PD; ++d) {\n float xv = sx[d];\n spa = fmaf(paw[d], xv, spa);\n sga = fmaf(gaw[d], xv, sga);\n spb = fmaf(pbw[d], xv, spb);\n sgb = fmaf(gbw[d], xv, sgb);\n }\n float m = mask[edge];\n int out_idx = edge * HD + t;\n a[out_idx] = spa * sigmoidf_fast(sga) * m;\n b[out_idx] = spb * sigmoidf_fast(sgb) * m;\n }\n}\n\n__global__ void final_kernel(const float* __restrict__ pair,\n const float* __restrict__ a,\n const float* __restrict__ b,\n const float* __restrict__ ln_w,\n const float* __restrict__ ln_b,\n const float* __restrict__ op_w,\n const float* __restrict__ op_b,\n const float* __restrict__ og_w,\n const float* __restrict__ og_b,\n float* __restrict__ out,\n int B, int N, int total_pairs) {\n int edge = blockIdx.x;\n int t = threadIdx.x;\n if (edge >= total_pairs) return;\n int n2 = N * N;\n int bb = edge / n2;\n int rem = edge - bb * n2;\n int i = rem / N;\n int j = rem - i * N;\n __shared__ float sh[64];\n if (t < HD) {\n float sum = 0.0f;\n int abase = ((bb * N + i) * N) * HD + t;\n int bbase = ((bb * N + j) * N) * HD + t;\n #pragma unroll\n for (int k = 0; k < 32; ++k) {\n sum = fmaf(a[abase + k * HD], b[bbase + k * HD], sum);\n }\n sh[t] = sum;\n }\n __syncthreads();\n // Thread 0 computes the small layernorm statistics and normalized hidden vector.\n if (t == 0) {\n float mean = 0.0f;\n #pragma unroll\n for (int h = 0; h < HD; ++h) mean += sh[h];\n mean *= (1.0f / 32.0f);\n float var = 0.0f;\n #pragma unroll\n for (int h = 0; h < HD; ++h) { float d = sh[h] - mean; var += d * d; }\n float inv = rsqrtf(var * (1.0f / 32.0f) + 1.0e-5f);\n #pragma unroll\n for (int h = 0; h < HD; ++h) sh[h] = (sh[h] - mean) * inv * ln_w[h] + ln_b[h];\n }\n __syncthreads();\n if (t < PD) {\n float y = op_b[t];\n const float* opw = op_w + t * HD;\n #pragma unroll\n for (int h = 0; h < HD; ++h) y = fmaf(opw[h], sh[h], y);\n const float* xptr = pair + edge * PD;\n float g = og_b[t];\n const float* ogw = og_w + t * PD;\n #pragma unroll\n for (int d = 0; d < PD; ++d) g = fmaf(ogw[d], xptr[d], g);\n out[edge * PD + t] = sigmoidf_fast(g) * y;\n }\n}\n\ntorch::Tensor tri_mul_forward(torch::Tensor pair_repr, torch::Tensor pair_mask,\n torch::Tensor ln_in_w, torch::Tensor ln_in_b,\n torch::Tensor proj_a_w, torch::Tensor proj_a_b,\n torch::Tensor gate_a_w, torch::Tensor gate_a_b,\n torch::Tensor proj_b_w, torch::Tensor proj_b_b,\n torch::Tensor gate_b_w, torch::Tensor gate_b_b,\n torch::Tensor ln_out_w, torch::Tensor ln_out_b,\n torch::Tensor out_proj_w, torch::Tensor out_proj_b,\n torch::Tensor out_gate_w, torch::Tensor out_gate_b) {\n int B = pair_repr.size(0);\n int N = pair_repr.size(1);\n int total = B * N * N;\n auto opts_ab = pair_repr.options();\n auto a = torch::empty({B, N, N, HD}, opts_ab);\n auto bb = torch::empty({B, N, N, HD}, opts_ab);\n auto out = torch::empty({B, N, N, PD}, opts_ab);\n dim3 block1(64);\n dim3 grid(total);\n compute_ab_kernel<<<grid, block1>>>(pair_repr.data_ptr<float>(), pair_mask.data_ptr<float>(),\n ln_in_w.data_ptr<float>(), ln_in_b.data_ptr<float>(),\n proj_a_w.data_ptr<float>(), proj_a_b.data_ptr<float>(),\n gate_a_w.data_ptr<float>(), gate_a_b.data_ptr<float>(),\n proj_b_w.data_ptr<float>(), proj_b_b.data_ptr<float>(),\n gate_b_w.data_ptr<float>(), gate_b_b.data_ptr<float>(),\n a.data_ptr<float>(), bb.data_ptr<float>(), total);\n dim3 block2(64);\n final_kernel<<<grid, block2>>>(pair_repr.data_ptr<float>(), a.data_ptr<float>(), bb.data_ptr<float>(),\n ln_out_w.data_ptr<float>(), ln_out_b.data_ptr<float>(),\n out_proj_w.data_ptr<float>(), out_proj_b.data_ptr<float>(),\n out_gate_w.data_ptr<float>(), out_gate_b.data_ptr<float>(),\n out.data_ptr<float>(), B, N, total);\n return out;\n}\n\"\"\"\n\n_ext = load_inline(name=\"tri_mul_ext_v1\", cpp_sources=_cpp_src, cuda_sources=_cuda_src,\n functions=[\"tri_mul_forward\"], extra_cuda_cflags=[\"-O3\", \"--use_fast_math\"], verbose=False)\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 def forward(self, pair_repr: torch.Tensor, pair_mask: torch.Tensor) -> torch.Tensor:\n return _ext.tri_mul_forward(pair_repr, pair_mask,\n self.layer_norm_in.weight, self.layer_norm_in.bias,\n self.proj_a.weight, self.proj_a.bias,\n self.gate_a.weight, self.gate_a.bias,\n self.proj_b.weight, self.proj_b.bias,\n self.gate_b.weight, self.gate_b.bias,\n self.layer_norm_out.weight, self.layer_norm_out.bias,\n self.out_proj.weight, self.out_proj.bias,\n self.out_gate.weight, self.out_gate.bias)\n"
}