call → compile_kernel
callcompile_kernel
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\n\ncpp_src = r'''\n#include <torch/extension.h>\ntorch::Tensor tri_forward(torch::Tensor x, torch::Tensor mask, torch::Tensor ln_w, torch::Tensor ln_b, torch::Tensor qw, torch::Tensor kw, torch::Tensor vw, torch::Tensor gw, torch::Tensor gb, torch::Tensor bw, torch::Tensor ow, torch::Tensor ob, int64_t num_heads, double eps);\n'''\n\ncuda_src = r'''\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n#include <math_constants.h>\n#define D 64\n#define NMAX 32\n__global__ void preprocess_kernel(const float* __restrict__ inp, const float* __restrict__ ln_w, const float* __restrict__ ln_b, const float* __restrict__ qw, const float* __restrict__ kw, const float* __restrict__ vw, const float* __restrict__ gw, const float* __restrict__ gb, const float* __restrict__ bw, float* __restrict__ q, float* __restrict__ k, float* __restrict__ v, float* __restrict__ gate, float* __restrict__ bias, int M, int N, float eps) {\n int t=blockIdx.x, tid=threadIdx.x; __shared__ float xs[D]; __shared__ float red[128]; float val=0.f; if(tid<D) val=inp[t*D+tid]; red[tid]=(tid<D)?val:0.f; __syncthreads();\n for(int s=64;s>0;s>>=1){ if(tid<s) red[tid]+=red[tid+s]; __syncthreads(); } float mean=red[0]*(1.f/64.f); float diff=(tid<D)?(val-mean):0.f; red[tid]=(tid<D)?diff*diff:0.f; __syncthreads();\n for(int s=64;s>0;s>>=1){ if(tid<s) red[tid]+=red[tid+s]; __syncthreads(); } float inv=rsqrtf(red[0]*(1.f/64.f)+eps); if(tid<D){ float xv=diff*inv*ln_w[tid]+ln_b[tid]; xs[tid]=xv; } __syncthreads();\n if(tid<D){ float sq=0.f,sk=0.f,sv=0.f,sg=gb[tid]; const float*qwp=qw+tid*D; const float*kwp=kw+tid*D; const float*vwp=vw+tid*D; const float*gwp=gw+tid*D;\n for(int i=0;i<D;i++){ float xv=xs[i]; sq+=xv*qwp[i]; sk+=xv*kwp[i]; sv+=xv*vwp[i]; sg+=xv*gwp[i]; }\n q[t*D+tid]=sq; k[t*D+tid]=sk; v[t*D+tid]=sv; gate[t*D+tid]=1.f/(1.f+expf(-sg)); }\n if(tid<4){ float sb=0.f; const float*bwp=bw+tid*D; for(int i=0;i<D;i++) sb+=xs[i]*bwp[i]; bias[t*4+tid]=sb; }\n}\n__global__ void attention_kernel(const float* __restrict__ q,const float* __restrict__ k,const float* __restrict__ v,const float* __restrict__ bias,const float* __restrict__ mask,float* __restrict__ attn_out,int B,int N){\n int idx=blockIdx.x; int qcol=idx&31; idx>>=5; int h=idx&3; idx>>=2; int row=idx&31; idx>>=5; int b=idx; int tid=threadIdx.x; __shared__ float scores[NMAX]; __shared__ float red[64];\n if(tid<N){ float dot=0.f; int qbase=(((b*N+row)*N+qcol)*D+h*16); int kbase=(((b*N+row)*N+tid)*D+h*16); for(int r=0;r<16;r++) dot+=q[qbase+r]*k[kbase+r]; float s=dot*0.25f+bias[((b*N+qcol)*N+tid)*4+h]; if(mask[(b*N+qcol)*N+tid]==0.f) s=-CUDART_INF_F; scores[tid]=s; } __syncthreads();\n float x=(tid<N)?scores[tid]:-CUDART_INF_F; red[tid]=x; __syncthreads(); for(int s=32;s>0;s>>=1){ if(tid<s) red[tid]=fmaxf(red[tid],red[tid+s]); __syncthreads(); } float mx=red[0]; float e=0.f; if(tid<N){ e=expf(scores[tid]-mx); scores[tid]=e; } red[tid]=(tid<N)?e:0.f; __syncthreads(); for(int s=32;s>0;s>>=1){ if(tid<s) red[tid]+=red[tid+s]; __syncthreads(); } float invden=1.f/red[0];\n if(tid<16){ float acc=0.f; for(int kk=0;kk<N;kk++){ int vbase=(((b*N+row)*N+kk)*D+h*16); acc+=scores[kk]*invden*v[vbase+tid]; } int obase=(((b*N+row)*N+qcol)*D+h*16); attn_out[obase+tid]=acc; }\n}\n__global__ void post_kernel(const float* __restrict__ attn_out,const float* __restrict__ gate,const float* __restrict__ ow,const float* __restrict__ ob,float* __restrict__ out,int M){ int t=blockIdx.x, tid=threadIdx.x; __shared__ float z[D]; if(tid<D) z[tid]=attn_out[t*D+tid]*gate[t*D+tid]; __syncthreads(); if(tid<D){ float sum=ob[tid]; const float*wp=ow+tid*D; for(int i=0;i<D;i++) sum+=z[i]*wp[i]; out[t*D+tid]=sum; } }\ntorch::Tensor tri_forward(torch::Tensor x, torch::Tensor mask, torch::Tensor ln_w, torch::Tensor ln_b, torch::Tensor qw, torch::Tensor kw, torch::Tensor vw, torch::Tensor gw, torch::Tensor gb, torch::Tensor bw, torch::Tensor ow, torch::Tensor ob, int64_t num_heads, double eps){ int B=x.size(0), N=x.size(1), Dv=x.size(3); TORCH_CHECK(Dv==64 && num_heads==4 && N<=32, \"only D=64\"); int M=B*N*N; auto opts=x.options(); auto q=torch::empty({M,64},opts), k=torch::empty({M,64},opts), v=torch::empty({M,64},opts), gate=torch::empty({M,64},opts), bias=torch::empty({M,4},opts), attn=torch::empty({M,64},opts); auto out=torch::empty_like(x); preprocess_kernel<<<M,128>>>(x.data_ptr<float>(),ln_w.data_ptr<float>(),ln_b.data_ptr<float>(),qw.data_ptr<float>(),kw.data_ptr<float>(),vw.data_ptr<float>(),gw.data_ptr<float>(),gb.data_ptr<float>(),bw.data_ptr<float>(),q.data_ptr<float>(),k.data_ptr<float>(),v.data_ptr<float>(),gate.data_ptr<float>(),bias.data_ptr<float>(),M,N,(float)eps); attention_kernel<<<B*N*4*N,64>>>(q.data_ptr<float>(),k.data_ptr<float>(),v.data_ptr<float>(),bias.data_ptr<float>(),mask.data_ptr<float>(),attn.data_ptr<float>(),B,N); post_kernel<<<M,64>>>(attn.data_ptr<float>(),gate.data_ptr<float>(),ow.data_ptr<float>(),ob.data_ptr<float>(),out.data_ptr<float>(),M); return out; }\n'''\n_ext=load_inline(name='tri_attn_ext_v2c', cpp_sources=[cpp_src], cuda_sources=[cuda_src], functions=['tri_forward'], extra_cuda_cflags=['-O3'], verbose=False)\nclass ModelNew(nn.Module):\n def __init__(self,pair_dim,num_heads,dropout=0.0):\n super().__init__(); self.num_heads=num_heads; self.head_dim=pair_dim//num_heads; assert pair_dim%num_heads==0; self.layer_norm=nn.LayerNorm(pair_dim); self.query=nn.Linear(pair_dim,pair_dim,bias=False); self.key=nn.Linear(pair_dim,pair_dim,bias=False); self.value=nn.Linear(pair_dim,pair_dim,bias=False); self.gate=nn.Linear(pair_dim,pair_dim); self.bias_proj=nn.Linear(pair_dim,num_heads,bias=False); self.out_proj=nn.Linear(pair_dim,pair_dim); self.dropout=nn.Dropout(dropout)\n def forward(self,pair_repr,pair_mask):\n return _ext.tri_forward(pair_repr.contiguous(),pair_mask.contiguous(),self.layer_norm.weight,self.layer_norm.bias,self.query.weight,self.key.weight,self.value.weight,self.gate.weight,self.gate.bias,self.bias_proj.weight,self.out_proj.weight,self.out_proj.bias,self.num_heads,self.layer_norm.eps)\n"
}