call → run_correctness
callrun_correctness
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\n\n_cpp = r'''\n#include <torch/extension.h>\n#include <vector>\ntorch::Tensor msa_attn_forward(torch::Tensor msa, torch::Tensor ng, torch::Tensor nb, torch::Tensor qw, torch::Tensor kw, torch::Tensor vw, torch::Tensor ow, torch::Tensor ob);\nstd::vector<torch::Tensor> opm_ab_forward(torch::Tensor msa, torch::Tensor ng, torch::Tensor nb, torch::Tensor aw, torch::Tensor ab, torch::Tensor bw, torch::Tensor bb);\ntorch::Tensor opm_outer_forward(torch::Tensor a, torch::Tensor b);\nstd::vector<torch::Tensor> tri_ab_forward(torch::Tensor pair, torch::Tensor ng, torch::Tensor nb, torch::Tensor paw, torch::Tensor pab, torch::Tensor pbw, torch::Tensor pbb, torch::Tensor gaw, torch::Tensor gab, torch::Tensor gbw, torch::Tensor gbb);\ntorch::Tensor tri_einsum_forward(torch::Tensor a, torch::Tensor b);\ntorch::Tensor norm_linear32_64_forward(torch::Tensor inp, torch::Tensor ng, torch::Tensor nb, torch::Tensor w, torch::Tensor bias);\ntorch::Tensor gate_residual_forward(torch::Tensor pair, torch::Tensor o, torch::Tensor w, torch::Tensor bias);\n'''\n\n_cuda = r'''\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n#include <vector>\n#include <cmath>\n#define EPS 1.0e-5f\n__device__ __forceinline__ float sigmoidf_fast(float x){ return 1.0f/(1.0f+expf(-x)); }\n__global__ void msa_qkv_kernel(const float* __restrict__ msa,const float* __restrict__ gamma,const float* __restrict__ beta,const float* __restrict__ qw,const float* __restrict__ kw,const float* __restrict__ vw,float* __restrict__ q,float* __restrict__ k,float* __restrict__ v,int total_tokens,int N){int tok=blockIdx.x,tid=threadIdx.x;__shared__ float x[64];__shared__ float red[128];float val=(tid<64)?msa[tok*64+tid]:0.0f;red[tid]=(tid<64)?val:0.0f;__syncthreads();for(int stride=64;stride>0;stride>>=1){if(tid<stride)red[tid]+=red[tid+stride];__syncthreads();}float mean=red[0]*(1.0f/64.0f);float diff=(tid<64)?(val-mean):0.0f;red[tid]=(tid<64)?diff*diff:0.0f;__syncthreads();for(int stride=64;stride>0;stride>>=1){if(tid<stride)red[tid]+=red[tid+stride];__syncthreads();}float inv=rsqrtf(red[0]*(1.0f/64.0f)+EPS);if(tid<64)x[tid]=diff*inv*gamma[tid]+beta[tid];__syncthreads();int r=tok/N;int n=tok-r*N;if(tid<64){float sq=0,sk=0,sv=0;for(int d=0;d<64;++d){float xv=x[d];sq+=xv*qw[tid*64+d];sk+=xv*kw[tid*64+d];sv+=xv*vw[tid*64+d];}int h=tid>>4;int dd=tid&15;int idx=(((r*4+h)*N+n)*16+dd);q[idx]=sq;k[idx]=sk;v[idx]=sv;}}\n__global__ void msa_attn_out_kernel(const float* __restrict__ msa,const float* __restrict__ q,const float* __restrict__ k,const float* __restrict__ v,const float* __restrict__ ow,const float* __restrict__ ob,float* __restrict__ out,int R,int N){int r=blockIdx.x,tid=threadIdx.x;__shared__ float ctx[32*64];for(int idx=tid;idx<N*64;idx+=blockDim.x){int n=idx/64;int od=idx-n*64;int h=od>>4;int dd=od&15;float maxs=-3.402823e38f;float scores[32];for(int j=0;j<32;++j){float s=0.0f;for(int t=0;t<16;++t)s+=q[(((r*4+h)*N+n)*16+t)]*k[(((r*4+h)*N+j)*16+t)];s*=0.25f;scores[j]=s;if(s>maxs)maxs=s;}float sum=0.0f;for(int j=0;j<32;++j){float e=expf(scores[j]-maxs);scores[j]=e;sum+=e;}float acc=0.0f;float inv=1.0f/sum;for(int j=0;j<32;++j)acc+=(scores[j]*inv)*v[(((r*4+h)*N+j)*16+dd)];ctx[idx]=acc;}__syncthreads();for(int idx=tid;idx<N*64;idx+=blockDim.x){int n=idx/64;int d=idx-n*64;float acc=ob[d];for(int e=0;e<64;++e)acc+=ctx[n*64+e]*ow[d*64+e];out[(r*N+n)*64+d]=msa[(r*N+n)*64+d]+acc;}}\ntorch::Tensor msa_attn_forward(torch::Tensor msa,torch::Tensor ng,torch::Tensor nb,torch::Tensor qw,torch::Tensor kw,torch::Tensor vw,torch::Tensor ow,torch::Tensor ob){auto mc=msa.contiguous();int B=mc.size(0),S=mc.size(1),N=mc.size(2);int R=B*S,total=R*N;auto opts=mc.options();auto q=torch::empty({R,4,N,16},opts),k=torch::empty({R,4,N,16},opts),v=torch::empty({R,4,N,16},opts);auto out=torch::empty_like(mc);msa_qkv_kernel<<<total,128>>>(mc.data_ptr<float>(),ng.data_ptr<float>(),nb.data_ptr<float>(),qw.data_ptr<float>(),kw.data_ptr<float>(),vw.data_ptr<float>(),q.data_ptr<float>(),k.data_ptr<float>(),v.data_ptr<float>(),total,N);msa_attn_out_kernel<<<R,128>>>(mc.data_ptr<float>(),q.data_ptr<float>(),k.data_ptr<float>(),v.data_ptr<float>(),ow.data_ptr<float>(),ob.data_ptr<float>(),out.data_ptr<float>(),R,N);return out;}\n__global__ void opm_ab_kernel(const float* __restrict__ msa,const float* __restrict__ gamma,const float* __restrict__ beta,const float* __restrict__ aw,const float* __restrict__ ab,const float* __restrict__ bw,const float* __restrict__ bb,float* __restrict__ a,float* __restrict__ b,int total){int tok=blockIdx.x,tid=threadIdx.x;__shared__ float x[64];__shared__ float red[128];float val=(tid<64)?msa[tok*64+tid]:0;red[tid]=(tid<64)?val:0;__syncthreads();for(int stride=64;stride>0;stride>>=1){if(tid<stride)red[tid]+=red[tid+stride];__syncthreads();}float mean=red[0]/64.0f;float diff=(tid<64)?val-mean:0;red[tid]=(tid<64)?diff*diff:0;__syncthreads();for(int stride=64;stride>0;stride>>=1){if(tid<stride)red[tid]+=red[tid+stride];__syncthreads();}float inv=rsqrtf(red[0]/64.0f+EPS);if(tid<64)x[tid]=diff*inv*gamma[tid]+beta[tid];__syncthreads();if(tid<32){float sa=ab[tid],sb=bb[tid];for(int d=0;d<64;++d){float xv=x[d];sa+=xv*aw[tid*64+d];sb+=xv*bw[tid*64+d];}a[tok*32+tid]=sa;b[tok*32+tid]=sb;}}\nstd::vector<torch::Tensor> opm_ab_forward(torch::Tensor msa,torch::Tensor ng,torch::Tensor nb,torch::Tensor aw,torch::Tensor ab,torch::Tensor bw,torch::Tensor bb){auto mc=msa.contiguous();int total=mc.numel()/64;auto opts=mc.options();auto a=torch::empty({total,32},opts);auto b=torch::empty({total,32},opts);opm_ab_kernel<<<total,128>>>(mc.data_ptr<float>(),ng.data_ptr<float>(),nb.data_ptr<float>(),aw.data_ptr<float>(),ab.data_ptr<float>(),bw.data_ptr<float>(),bb.data_ptr<float>(),a.data_ptr<float>(),b.data_ptr<float>(),total);return {a,b};}\n__global__ void opm_outer_kernel(const float* __restrict__ a,const float* __restrict__ b,float* __restrict__ out,int B,int S,int N){int p=blockIdx.x,tid=threadIdx.x;int j=p%N;int tmp=p/N;int i=tmp%N;int batch=tmp/N;for(int idx=tid;idx<1024;idx+=blockDim.x){int h=idx/32;int k2=idx&31;float acc=0;for(int s=0;s<S;++s){int base=((batch*S+s)*N);acc+=a[(base+i)*32+h]*b[(base+j)*32+k2];}out[((batch*N*N+i*N+j)*1024)+idx]=acc/(float)S;}}\ntorch::Tensor opm_outer_forward(torch::Tensor a,torch::Tensor b){int total=a.size(0),S=8,N=32,B=total/(S*N);auto out=torch::empty({B,N,N,1024},a.options());opm_outer_kernel<<<B*N*N,256>>>(a.data_ptr<float>(),b.data_ptr<float>(),out.data_ptr<float>(),B,S,N);return out;}\n__global__ void tri_ab_kernel(const float* __restrict__ pair,const float* __restrict__ gamma,const float* __restrict__ beta,const float* __restrict__ paw,const float* __restrict__ pab,const float* __restrict__ pbw,const float* __restrict__ pbb,const float* __restrict__ gaw,const float* __restrict__ gab,const float* __restrict__ gbw,const float* __restrict__ gbb,float* __restrict__ a,float* __restrict__ b,int total){int tok=blockIdx.x,tid=threadIdx.x;__shared__ float x[64];__shared__ float red[128];float val=(tid<64)?pair[tok*64+tid]:0;red[tid]=(tid<64)?val:0;__syncthreads();for(int stride=64;stride>0;stride>>=1){if(tid<stride)red[tid]+=red[tid+stride];__syncthreads();}float mean=red[0]/64.0f;float diff=(tid<64)?val-mean:0;red[tid]=(tid<64)?diff*diff:0;__syncthreads();for(int stride=64;stride>0;stride>>=1){if(tid<stride)red[tid]+=red[tid+stride];__syncthreads();}float inv=rsqrtf(red[0]/64.0f+EPS);if(tid<64)x[tid]=diff*inv*gamma[tid]+beta[tid];__syncthreads();if(tid<32){float pa=pab[tid],pb=pbb[tid],ga=gab[tid],gb=gbb[tid];for(int d=0;d<64;++d){float xv=x[d];pa+=xv*paw[tid*64+d];pb+=xv*pbw[tid*64+d];ga+=xv*gaw[tid*64+d];gb+=xv*gbw[tid*64+d];}a[tok*32+tid]=pa*sigmoidf_fast(ga);b[tok*32+tid]=pb*sigmoidf_fast(gb);}}\nstd::vector<torch::Tensor> tri_ab_forward(torch::Tensor pair,torch::Tensor ng,torch::Tensor nb,torch::Tensor paw,torch::Tensor pab,torch::Tensor pbw,torch::Tensor pbb,torch::Tensor gaw,torch::Tensor gab,torch::Tensor gbw,torch::Tensor gbb){auto pc=pair.contiguous();int total=pc.numel()/64;auto a=torch::empty({total,32},pc.options());auto b=torch::empty({total,32},pc.options());tri_ab_kernel<<<total,128>>>(pc.data_ptr<float>(),ng.data_ptr<float>(),nb.data_ptr<float>(),paw.data_ptr<float>(),pab.data_ptr<float>(),pbw.data_ptr<float>(),pbb.data_ptr<float>(),gaw.data_ptr<float>(),gab.data_ptr<float>(),gbw.data_ptr<float>(),gbb.data_ptr<float>(),a.data_ptr<float>(),b.data_ptr<float>(),total);return {a,b};}\n__global__ void tri_einsum_kernel(const float* __restrict__ a,const float* __restrict__ b,float* __restrict__ out,int B,int N){int p=blockIdx.x,h=threadIdx.x;int j=p%N;int tmp=p/N;int i=tmp%N;int batch=tmp/N;if(h<32){float acc=0;int base=batch*N*N;for(int kk=0;kk<N;++kk)acc+=a[(base+i*N+kk)*32+h]*b[(base+j*N+kk)*32+h];out[p*32+h]=acc;}}\ntorch::Tensor tri_einsum_forward(torch::Tensor a,torch::Tensor b){int total=a.size(0),N=32,B=total/(N*N);auto out=torch::empty({B,N,N,32},a.options());tri_einsum_kernel<<<B*N*N,32>>>(a.data_ptr<float>(),b.data_ptr<float>(),out.data_ptr<float>(),B,N);return out;}\n__global__ void norm_linear32_64_kernel(const float* __restrict__ inp,const float* __restrict__ gamma,const float* __restrict__ beta,const float* __restrict__ w,const float* __restrict__ bias,float* __restrict__ out,int total){int tok=blockIdx.x,tid=threadIdx.x;__shared__ float x[32];__shared__ float red[64];float val=(tid<32)?inp[tok*32+tid]:0;red[tid]=(tid<32)?val:0;__syncthreads();for(int stride=32;stride>0;stride>>=1){if(tid<stride)red[tid]+=red[tid+stride];__syncthreads();}float mean=red[0]/32.0f;float diff=(tid<32)?val-mean:0;red[tid]=(tid<32)?diff*diff:0;__syncthreads();for(int stride=32;stride>0;stride>>=1){if(tid<stride)red[tid]+=red[tid+stride];__syncthreads();}float inv=rsqrtf(red[0]/32.0f+EPS);if(tid<32)x[tid]=diff*inv*gamma[tid]+beta[tid];__syncthreads();if(tid<64){float acc=bias[tid];for(int d=0;d<32;++d)acc+=x[d]*w[tid*32+d];out[tok*64+tid]=acc;}}\ntorch::Tensor norm_linear32_64_forward(torch::Tensor inp,torch::Tensor ng,torch::Tensor nb,torch::Tensor w,torch::Tensor bias){auto ic=inp.contiguous();int total=ic.numel()/32;auto out=torch::empty({total,64},ic.options());norm_linear32_64_kernel<<<total,64>>>(ic.data_ptr<float>(),ng.data_ptr<float>(),nb.data_ptr<float>(),w.data_ptr<float>(),bias.data_ptr<float>(),out.data_ptr<float>(),total);return out;}\n__global__ void gate_residual_kernel(const float* __restrict__ pair,const float* __restrict__ o,const float* __restrict__ w,const float* __restrict__ bias,float* __restrict__ out,int total){int tok=blockIdx.x,tid=threadIdx.x;__shared__ float x[64];if(tid<64)x[tid]=pair[tok*64+tid];__syncthreads();if(tid<64){float g=bias[tid];for(int d=0;d<64;++d)g+=x[d]*w[tid*64+d];out[tok*64+tid]=x[tid]+sigmoidf_fast(g)*o[tok*64+tid];}}\ntorch::Tensor gate_residual_forward(torch::Tensor pair,torch::Tensor o,torch::Tensor w,torch::Tensor bias){auto pc=pair.contiguous();auto oc=o.contiguous();int total=pc.numel()/64;auto out=torch::empty_like(pc);gate_residual_kernel<<<total,64>>>(pc.data_ptr<float>(),oc.data_ptr<float>(),w.data_ptr<float>(),bias.data_ptr<float>(),out.data_ptr<float>(),total);return out;}\n'''\n_ext = load_inline(name='evo_kernel_v1', cpp_sources=_cpp, cuda_sources=_cuda, functions=['msa_attn_forward','opm_ab_forward','opm_outer_forward','tri_ab_forward','tri_einsum_forward','norm_linear32_64_forward','gate_residual_forward'], extra_cuda_cflags=['-O3','--use_fast_math'], verbose=False)\nclass MSARowSelfAttention(nn.Module):\n def __init__(self, msa_dim, num_heads):\n super().__init__(); self.num_heads=num_heads; self.head_dim=msa_dim//num_heads; self.norm=nn.LayerNorm(msa_dim); self.q=nn.Linear(msa_dim,msa_dim,bias=False); self.k=nn.Linear(msa_dim,msa_dim,bias=False); self.v=nn.Linear(msa_dim,msa_dim,bias=False); self.out=nn.Linear(msa_dim,msa_dim)\nclass OuterProductMeanSmall(nn.Module):\n def __init__(self, msa_dim, pair_dim, hidden=32):\n super().__init__(); self.norm=nn.LayerNorm(msa_dim); self.a=nn.Linear(msa_dim,hidden); self.b=nn.Linear(msa_dim,hidden); self.out=nn.Linear(hidden*hidden,pair_dim)\nclass TriMulOut(nn.Module):\n def __init__(self, pair_dim, hidden):\n super().__init__(); self.norm=nn.LayerNorm(pair_dim); self.pa=nn.Linear(pair_dim,hidden); self.pb=nn.Linear(pair_dim,hidden); self.ga=nn.Linear(pair_dim,hidden); self.gb=nn.Linear(pair_dim,hidden); self.norm_out=nn.LayerNorm(hidden); self.out=nn.Linear(hidden,pair_dim); self.gate_out=nn.Linear(pair_dim,pair_dim)\nclass ModelNew(nn.Module):\n def __init__(self, msa_dim, pair_dim, num_heads, tri_hidden):\n super().__init__(); self.msa_attn=MSARowSelfAttention(msa_dim,num_heads); self.opm=OuterProductMeanSmall(msa_dim,pair_dim); self.tri_mul=TriMulOut(pair_dim,tri_hidden)\n def forward(self, msa_repr: torch.Tensor, pair_repr: torch.Tensor):\n msa_repr=_ext.msa_attn_forward(msa_repr,self.msa_attn.norm.weight,self.msa_attn.norm.bias,self.msa_attn.q.weight,self.msa_attn.k.weight,self.msa_attn.v.weight,self.msa_attn.out.weight,self.msa_attn.out.bias)\n ab=_ext.opm_ab_forward(msa_repr,self.opm.norm.weight,self.opm.norm.bias,self.opm.a.weight,self.opm.a.bias,self.opm.b.weight,self.opm.b.bias)\n outer=_ext.opm_outer_forward(ab[0],ab[1])\n pair_repr=pair_repr+self.opm.out(outer)\n tb=_ext.tri_ab_forward(pair_repr,self.tri_mul.norm.weight,self.tri_mul.norm.bias,self.tri_mul.pa.weight,self.tri_mul.pa.bias,self.tri_mul.pb.weight,self.tri_mul.pb.bias,self.tri_mul.ga.weight,self.tri_mul.ga.bias,self.tri_mul.gb.weight,self.tri_mul.gb.bias)\n o=_ext.tri_einsum_forward(tb[0],tb[1])\n o2=_ext.norm_linear32_64_forward(o,self.tri_mul.norm_out.weight,self.tri_mul.norm_out.bias,self.tri_mul.out.weight,self.tri_mul.out.bias)\n pair_repr=_ext.gate_residual_forward(pair_repr,o2,self.tri_mul.gate_out.weight,self.tri_mul.gate_out.bias)\n return msa_repr,pair_repr\n"
}