callcompile_kernel
{
"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>\nstd::vector<torch::Tensor> full_forward(torch::Tensor msa, torch::Tensor pair, std::vector<torch::Tensor> p);\n'''\n\n_cuda = r'''\n#include <torch/extension.h>\n#include <ATen/cuda/CUDAContext.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n#include <vector>\n#include <math.h>\n#define DMSA 64\n#define HHEAD 4\n#define HDIM 16\n#define PH 32\n#define PDIM 64\n#define EPS 1.0e-5f\n__device__ __forceinline__ float sigm(float x){return 1.f/(1.f+expf(-x));}\n\n__global__ void qkv_k(const float* msa,const float* lnw,const float* lnb,const float* qw,const float* kw,const float* vw,float* q,float* k,float* v,int B,int S,int N){\n int row=blockIdx.x; extern __shared__ float x[]; __shared__ float red[128];\n float s=0; for(int i=threadIdx.x;i<DMSA;i+=blockDim.x) s+=msa[row*DMSA+i]; red[threadIdx.x]=s; __syncthreads();\n for(int off=64;off>0;off>>=1){if(threadIdx.x<off) red[threadIdx.x]+=red[threadIdx.x+off]; __syncthreads();}\n float mean=red[0]/DMSA; float ss=0; for(int i=threadIdx.x;i<DMSA;i+=blockDim.x){float t=msa[row*DMSA+i]-mean; ss+=t*t;} red[threadIdx.x]=ss; __syncthreads();\n for(int off=64;off>0;off>>=1){if(threadIdx.x<off) red[threadIdx.x]+=red[threadIdx.x+off]; __syncthreads();}\n float inv=rsqrtf(red[0]/DMSA+EPS); for(int i=threadIdx.x;i<DMSA;i+=blockDim.x)x[i]=(msa[row*DMSA+i]-mean)*inv*lnw[i]+lnb[i]; __syncthreads();\n int bsn=row/N,n=row%N,b=bsn/S,st=bsn%S;\n for(int od=threadIdx.x;od<DMSA;od+=blockDim.x){float a=0,bb=0,c=0; for(int i=0;i<DMSA;i++){float xx=x[i]; a+=xx*qw[od*DMSA+i]; bb+=xx*kw[od*DMSA+i]; c+=xx*vw[od*DMSA+i];} int h=od/HDIM,d=od%HDIM,idx=(((b*S+st)*HHEAD+h)*N+n)*HDIM+d; q[idx]=a;k[idx]=bb;v[idx]=c;}\n}\n\n__global__ void attn_k(const float* q,const float* k,const float* v,float* out,int B,int S,int N){\n int idx=blockIdx.x,nq=idx%N,t=idx/N,h=t%HHEAD; t/=HHEAD; int st=t%S,b=t/S; __shared__ float sc[32]; __shared__ float mv,sv;\n if(threadIdx.x<N){int nk=threadIdx.x; float dot=0; for(int d=0;d<HDIM;d++) dot+=q[(((b*S+st)*HHEAD+h)*N+nq)*HDIM+d]*k[(((b*S+st)*HHEAD+h)*N+nk)*HDIM+d]; sc[nk]=dot*0.25f;} __syncthreads();\n if(threadIdx.x==0){float m=-3.402823e38f; for(int i=0;i<N;i++) if(sc[i]>m)m=sc[i]; mv=m;} __syncthreads();\n if(threadIdx.x<N) sc[threadIdx.x]=expf(sc[threadIdx.x]-mv); __syncthreads();\n if(threadIdx.x==0){float sum=0; for(int i=0;i<N;i++) sum+=sc[i]; sv=sum;} __syncthreads(); float inv=1.f/sv;\n if(threadIdx.x<HDIM){int d=threadIdx.x; float res=0; for(int nk=0;nk<N;nk++) res+=sc[nk]*inv*v[(((b*S+st)*HHEAD+h)*N+nk)*HDIM+d]; out[((b*S+st)*N+nq)*DMSA+h*HDIM+d]=res;}\n}\n\n__global__ void lin64_res_k(const float* x,const float* res,const float* w,const float* bias,float* y,int rows){int row=blockIdx.x,c=threadIdx.x; if(row<rows&&c<DMSA){float s=bias[c]; for(int i=0;i<DMSA;i++) s+=x[row*DMSA+i]*w[c*DMSA+i]; y[row*DMSA+c]=res[row*DMSA+c]+s;}}\n\n__global__ void opm_ab_k(const float* msa,const float* lnw,const float* lnb,const float* aw,const float* ab,const float* bw,const float* bb,float* a,float* bvec,int rows){\n int row=blockIdx.x; extern __shared__ float x[]; __shared__ float red[128]; float s=0; for(int i=threadIdx.x;i<DMSA;i+=blockDim.x)s+=msa[row*DMSA+i]; red[threadIdx.x]=s; __syncthreads(); for(int off=64;off>0;off>>=1){if(threadIdx.x<off)red[threadIdx.x]+=red[threadIdx.x+off];__syncthreads();} float mean=red[0]/DMSA; float ss=0; for(int i=threadIdx.x;i<DMSA;i+=blockDim.x){float t=msa[row*DMSA+i]-mean;ss+=t*t;} red[threadIdx.x]=ss; __syncthreads(); for(int off=64;off>0;off>>=1){if(threadIdx.x<off)red[threadIdx.x]+=red[threadIdx.x+off];__syncthreads();} float inv=rsqrtf(red[0]/DMSA+EPS); for(int i=threadIdx.x;i<DMSA;i+=blockDim.x)x[i]=(msa[row*DMSA+i]-mean)*inv*lnw[i]+lnb[i]; __syncthreads(); int h=threadIdx.x; if(h<PH){float sa=ab[h],sb=bb[h]; for(int i=0;i<DMSA;i++){float xx=x[i]; sa+=xx*aw[h*DMSA+i]; sb+=xx*bw[h*DMSA+i];} a[row*PH+h]=sa; bvec[row*PH+h]=sb;}}\n\n__global__ void outer_k(const float* a,const float* bvec,float* outer,int B,int S,int N){int idx=blockIdx.x*blockDim.x+threadIdx.x,total=B*N*N*1024; if(idx>=total)return; int h2=idx%PH,tmp=idx/PH,h1=tmp%PH; tmp/=PH; int j=tmp%N; tmp/=N; int i=tmp%N; int b=tmp/N; float sm=0; for(int s=0;s<S;s++) sm+=a[((b*S+s)*N+i)*PH+h1]*bvec[((b*S+s)*N+j)*PH+h2]; outer[idx]=sm/(float)S;}\n\n__global__ void opm_lin_k(const float* A,const float* pair,const float* w,const float* bias,float* out,int M){__shared__ float As[16][16],Bs[16][16]; int row=blockIdx.y*16+threadIdx.y,col=blockIdx.x*16+threadIdx.x; float acc=0; for(int kt=0;kt<1024;kt+=16){As[threadIdx.y][threadIdx.x]=(row<M)?A[row*1024+kt+threadIdx.x]:0; Bs[threadIdx.y][threadIdx.x]=(col<64)?w[col*1024+kt+threadIdx.y]:0; __syncthreads(); for(int kk=0;kk<16;kk++) acc+=As[threadIdx.y][kk]*Bs[kk][threadIdx.x]; __syncthreads();} if(row<M&&col<64) out[row*64+col]=pair[row*64+col]+acc+bias[col];}\n\n__global__ void tri_ab_k(const float* pair,const float* lnw,const float* lnb,const float* paw,const float* pab,const float* pbw,const float* pbb,const float* gaw,const float* gab,const float* gbw,const float* gbb,float* a,float* bvec,int M){int row=blockIdx.x; extern __shared__ float x[]; __shared__ float red[128]; float s=0; for(int i=threadIdx.x;i<64;i+=blockDim.x)s+=pair[row*64+i]; red[threadIdx.x]=s; __syncthreads(); for(int off=64;off>0;off>>=1){if(threadIdx.x<off)red[threadIdx.x]+=red[threadIdx.x+off];__syncthreads();} float mean=red[0]/64; float ss=0; for(int i=threadIdx.x;i<64;i+=blockDim.x){float t=pair[row*64+i]-mean;ss+=t*t;} red[threadIdx.x]=ss; __syncthreads(); for(int off=64;off>0;off>>=1){if(threadIdx.x<off)red[threadIdx.x]+=red[threadIdx.x+off];__syncthreads();} float inv=rsqrtf(red[0]/64+EPS); for(int i=threadIdx.x;i<64;i+=blockDim.x)x[i]=(pair[row*64+i]-mean)*inv*lnw[i]+lnb[i]; __syncthreads(); int h=threadIdx.x; if(h<PH){float pa=pab[h],pb=pbb[h],ga=gab[h],gb=gbb[h]; for(int i=0;i<64;i++){float xx=x[i]; pa+=xx*paw[h*64+i]; pb+=xx*pbw[h*64+i]; ga+=xx*gaw[h*64+i]; gb+=xx*gbw[h*64+i];} a[row*PH+h]=pa*sigm(ga); bvec[row*PH+h]=pb*sigm(gb);}}\n__global__ void tri_e_k(const float* a,const float* bvec,float* o,int B,int N){int idx=blockIdx.x*blockDim.x+threadIdx.x,total=B*N*N*PH; if(idx>=total)return; int h=idx%PH,tmp=idx/PH,j=tmp%N; tmp/=N; int i=tmp%N,b=tmp/N; float s=0; for(int k=0;k<N;k++) s+=a[((b*N+i)*N+k)*PH+h]*bvec[((b*N+j)*N+k)*PH+h]; o[idx]=s;}\n__global__ void tri_fin_k(const float* o,const float* pair,const float* nw,const float* nb,const float* ow,const float* ob,const float* gw,const float* gb,float* out,int M){int row=blockIdx.x; __shared__ float norm[32],red[128]; float s=0; for(int i=threadIdx.x;i<PH;i+=blockDim.x)s+=o[row*PH+i]; red[threadIdx.x]=s; __syncthreads(); for(int off=64;off>0;off>>=1){if(threadIdx.x<off)red[threadIdx.x]+=red[threadIdx.x+off];__syncthreads();} float mean=red[0]/PH; float ss=0; for(int i=threadIdx.x;i<PH;i+=blockDim.x){float t=o[row*PH+i]-mean;ss+=t*t;} red[threadIdx.x]=ss; __syncthreads(); for(int off=64;off>0;off>>=1){if(threadIdx.x<off)red[threadIdx.x]+=red[threadIdx.x+off];__syncthreads();} float inv=rsqrtf(red[0]/PH+EPS); for(int i=threadIdx.x;i<PH;i+=blockDim.x)norm[i]=(o[row*PH+i]-mean)*inv*nw[i]+nb[i]; __syncthreads(); int c=threadIdx.x; if(c<64){float so=ob[c]; for(int i=0;i<PH;i++) so+=norm[i]*ow[c*PH+i]; float sg=gb[c]; for(int i=0;i<64;i++) sg+=pair[row*64+i]*gw[c*64+i]; out[row*64+c]=pair[row*64+c]+sigm(sg)*so;}}\n\nstd::vector<torch::Tensor> full_forward(torch::Tensor msa,torch::Tensor pair,std::vector<torch::Tensor> p){int B=msa.size(0),S=msa.size(1),N=msa.size(2),rm=B*S*N,rp=B*N*N; auto opts=msa.options(); cudaStream_t st=at::cuda::getCurrentCUDAStream(); auto q=torch::empty({B*S*HHEAD*N*HDIM},opts),k=torch::empty({B*S*HHEAD*N*HDIM},opts),v=torch::empty({B*S*HHEAD*N*HDIM},opts),att=torch::empty({rm,64},opts),msa2=torch::empty_like(msa); qkv_k<<<rm,128,64*sizeof(float),st>>>(msa.data_ptr<float>(),p[0].data_ptr<float>(),p[1].data_ptr<float>(),p[2].data_ptr<float>(),p[3].data_ptr<float>(),p[4].data_ptr<float>(),q.data_ptr<float>(),k.data_ptr<float>(),v.data_ptr<float>(),B,S,N); attn_k<<<B*S*HHEAD*N,32,0,st>>>(q.data_ptr<float>(),k.data_ptr<float>(),v.data_ptr<float>(),att.data_ptr<float>(),B,S,N); lin64_res_k<<<rm,64,0,st>>>(att.data_ptr<float>(),msa.data_ptr<float>(),p[5].data_ptr<float>(),p[6].data_ptr<float>(),msa2.data_ptr<float>(),rm); auto a=torch::empty({rm,PH},opts),bb=torch::empty({rm,PH},opts); opm_ab_k<<<rm,128,64*sizeof(float),st>>>(msa2.data_ptr<float>(),p[7].data_ptr<float>(),p[8].data_ptr<float>(),p[9].data_ptr<float>(),p[10].data_ptr<float>(),p[11].data_ptr<float>(),p[12].data_ptr<float>(),a.data_ptr<float>(),bb.data_ptr<float>(),rm); auto outer=torch::empty({rp,1024},opts); outer_k<<<(rp*1024+255)/256,256,0,st>>>(a.data_ptr<float>(),bb.data_ptr<float>(),outer.data_ptr<float>(),B,S,N); auto pair1=torch::empty_like(pair); dim3 bl(16,16),gr(4,(rp+15)/16); opm_lin_k<<<gr,bl,0,st>>>(outer.data_ptr<float>(),pair.data_ptr<float>(),p[13].data_ptr<float>(),p[14].data_ptr<float>(),pair1.data_ptr<float>(),rp); auto ta=torch::empty({rp,PH},opts),tb=torch::empty({rp,PH},opts); tri_ab_k<<<rp,128,64*sizeof(float),st>>>(pair1.data_ptr<float>(),p[15].data_ptr<float>(),p[16].data_ptr<float>(),p[17].data_ptr<float>(),p[18].data_ptr<float>(),p[19].data_ptr<float>(),p[20].data_ptr<float>(),p[21].data_ptr<float>(),p[22].data_ptr<float>(),p[23].data_ptr<float>(),p[24].data_ptr<float>(),ta.data_ptr<float>(),tb.data_ptr<float>(),rp); auto to=torch::empty({rp,PH},opts); tri_e_k<<<(rp*PH+255)/256,256,0,st>>>(ta.data_ptr<float>(),tb.data_ptr<float>(),to.data_ptr<float>(),B,N); auto po=torch::empty_like(pair); tri_fin_k<<<rp,128,0,st>>>(to.data_ptr<float>(),pair1.data_ptr<float>(),p[25].data_ptr<float>(),p[26].data_ptr<float>(),p[27].data_ptr<float>(),p[28].data_ptr<float>(),p[29].data_ptr<float>(),p[30].data_ptr<float>(),po.data_ptr<float>(),rp); return {msa2,po};}\n'''\n_ext=load_inline(name='evoformer_custom_v3',cpp_sources=_cpp,cuda_sources=_cuda,functions=['full_forward'],extra_cflags=['-O3'],extra_cuda_cflags=['-O3'],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 p=[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,self.opm.norm.weight,self.opm.norm.bias,self.opm.a.weight,self.opm.a.bias,self.opm.b.weight,self.opm.b.bias,self.opm.out.weight,self.opm.out.bias,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,self.tri_mul.norm_out.weight,self.tri_mul.norm_out.bias,self.tri_mul.out.weight,self.tri_mul.out.bias,self.tri_mul.gate_out.weight,self.tri_mul.gate_out.bias]\n return tuple(_ext.full_forward(msa_repr.contiguous(),pair_repr.contiguous(),p))\n"
}