call → run_correctness
callrun_correctness
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\n_cpp_src=r'''\n#include <torch/extension.h>\nvoid ln_gelu_out_cuda(torch::Tensor x, torch::Tensor gamma, torch::Tensor beta, torch::Tensor out, torch::Tensor pmean, torch::Tensor pm2, torch::Tensor pcnt);\nvoid ln_gelu_out(torch::Tensor x, torch::Tensor gamma, torch::Tensor beta, torch::Tensor out, torch::Tensor pmean, torch::Tensor pm2, torch::Tensor pcnt){ln_gelu_out_cuda(x,gamma,beta,out,pmean,pm2,pcnt);} \n'''\n_cuda_src=r'''\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n#include <math.h>\n#define CHUNKS 256\nstruct WelfordData { float mean; float m2; int n; };\n__device__ __forceinline__ WelfordData combine(WelfordData a, WelfordData b){ if(b.n==0) return a; if(a.n==0) return b; float delta=b.mean-a.mean; int n=a.n+b.n; float nb_over_n=(float)b.n/(float)n; a.mean += delta*nb_over_n; a.m2 += b.m2 + delta*delta*((float)a.n*nb_over_n); a.n=n; return a; }\n__inline__ __device__ WelfordData warp_reduce(WelfordData v){for(int off=16;off>0;off>>=1){WelfordData o; o.mean=__shfl_down_sync(0xffffffff,v.mean,off); o.m2=__shfl_down_sync(0xffffffff,v.m2,off); o.n=__shfl_down_sync(0xffffffff,v.n,off); v=combine(v,o);} return v;}\n__inline__ __device__ WelfordData block_reduce(WelfordData v){static __shared__ WelfordData shared[32]; int lane=threadIdx.x&31, wid=threadIdx.x>>5; v=warp_reduce(v); if(lane==0) shared[wid]=v; __syncthreads(); WelfordData z; z.mean=0; z.m2=0; z.n=0; v=(threadIdx.x<(blockDim.x>>5))?shared[lane]:z; if(wid==0) v=warp_reduce(v); return v;}\n__global__ void partial_kernel(const float* x,float* pmean,float* pm2,float* pcnt,int B,int N){int bid=blockIdx.x; int b=bid/CHUNKS; int c=bid-b*CHUNKS; int chunk=(N+CHUNKS-1)/CHUNKS; int start=c*chunk; int end=min(start+chunk,N); const float* xb=x+(long long)b*N; WelfordData w{0.0f,0.0f,0}; for(int i=start+threadIdx.x;i<end;i+=blockDim.x){float val=xb[i]; w.n++; float delta=val-w.mean; w.mean += delta/(float)w.n; float delta2=val-w.mean; w.m2 += delta*delta2;} w=block_reduce(w); if(threadIdx.x==0){pmean[bid]=w.mean; pm2[bid]=w.m2; pcnt[bid]=(float)w.n;}}\n__global__ void stats_kernel(const float* pmean,const float* pm2,const float* pcnt,float* stats,int B,int N){int b=blockIdx.x; WelfordData w{0.0f,0.0f,0}; for(int i=threadIdx.x;i<CHUNKS;i+=blockDim.x){int idx=b*CHUNKS+i; WelfordData o{pmean[idx],pm2[idx],(int)pcnt[idx]}; w=combine(w,o);} w=block_reduce(w); if(threadIdx.x==0){float var=w.m2/(float)N; stats[2*b]=w.mean; stats[2*b+1]=rsqrtf(var+1e-5f);}}\n__global__ void norm_gelu(const float* x,const float* gamma,const float* beta,const float* stats,float* out,long long total,int N){long long idx=(long long)blockIdx.x*blockDim.x+threadIdx.x; long long stride=(long long)blockDim.x*gridDim.x; const float k=0.7071067811865475244f; for(;idx<total;idx+=stride){int b=(int)(idx/N); int j=(int)(idx-(long long)b*N); float z=(x[idx]-stats[2*b])*stats[2*b+1]; z=z*gamma[j]+beta[j]; out[idx]=0.5f*z*(1.0f+erff(z*k));}}\nvoid ln_gelu_out_cuda(torch::Tensor x, torch::Tensor gamma, torch::Tensor beta, torch::Tensor out, torch::Tensor pmean, torch::Tensor pm2, torch::Tensor pcnt){int B=x.size(0); int N=x.size(1)*x.size(2); long long total=(long long)B*N; auto stats=torch::empty({B,2},x.options()); int t=256; partial_kernel<<<B*CHUNKS,t>>>(x.data_ptr<float>(),pmean.data_ptr<float>(),pm2.data_ptr<float>(),pcnt.data_ptr<float>(),B,N); stats_kernel<<<B,t>>>(pmean.data_ptr<float>(),pm2.data_ptr<float>(),pcnt.data_ptr<float>(),stats.data_ptr<float>(),B,N); int blocks=(total+t-1)/t; if(blocks>65535)blocks=65535; norm_gelu<<<blocks,t>>>(x.data_ptr<float>(),gamma.data_ptr<float>(),beta.data_ptr<float>(),stats.data_ptr<float>(),out.data_ptr<float>(),total,N);}\n'''\n_ext=load_inline(name='ln_welford_ext1',cpp_sources=_cpp_src,cuda_sources=_cuda_src,functions=['ln_gelu_out'],extra_cuda_cflags=['-O3'],extra_cflags=['-O3'],verbose=False)\nclass ModelNew(nn.Module):\n def __init__(self, channels, seq_len, kernel_size, max_dilation_power):\n super().__init__(); self.norms=nn.ModuleList(); self.convs=nn.ModuleList(); self.activation=nn.GELU()\n for i in range(max_dilation_power+1):\n d=2**i; self.norms.append(nn.LayerNorm([channels,seq_len])); self.convs.append(nn.Conv1d(channels,channels,kernel_size,padding=d,dilation=d))\n def forward(self,x):\n residual=x; B=x.shape[0]; tmp=torch.empty_like(x); pmean=torch.empty((B,256),device=x.device,dtype=x.dtype); pm2=torch.empty((B,256),device=x.device,dtype=x.dtype); pcnt=torch.empty((B,256),device=x.device,dtype=x.dtype)\n for norm,conv in zip(self.norms,self.convs):\n _ext.ln_gelu_out(x,norm.weight,norm.bias,tmp,pmean,pm2,pcnt)\n x=conv(tmp); x.add_(residual)\n return x\n"
}