final submitted kernel (95 lines, 6,709 chars) — full source, untruncated
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
_cpp_src = r'''
#include <torch/extension.h>
torch::Tensor fused_forward(torch::Tensor input, torch::Tensor weight, torch::Tensor bias,
torch::Tensor bn_weight, torch::Tensor bn_bias,
torch::Tensor running_mean, torch::Tensor running_var, torch::Tensor scale,
bool training, double eps, double momentum);
'''
_cuda_src = r'''
#include <torch/extension.h>
#include <ATen/ATen.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <float.h>
#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x " must be CUDA")
#define CHECK_FLOAT(x) TORCH_CHECK(x.scalar_type()==at::kFloat, #x " must be float")
__global__ void partial_stats_kernel(const float* __restrict__ y, float* __restrict__ psum, float* __restrict__ psum2, int M, int N) {
int tx=threadIdx.x, ty=threadIdx.y, c=blockIdx.x*32+tx, g=blockIdx.y;
int r0=g*256, r1=min(r0+256,M); float s=0.f, s2=0.f;
if(c<N){ for(int r=r0+ty;r<r1;r+=8){ float v=y[(long long)r*N+c]; s+=v; s2+=v*v; } }
__shared__ float sh[8][32]; __shared__ float sh2[8][32]; sh[ty][tx]=s; sh2[ty][tx]=s2; __syncthreads();
if(ty==0 && c<N){
float a=sh[0][tx]+sh[1][tx]+sh[2][tx]+sh[3][tx]+sh[4][tx]+sh[5][tx]+sh[6][tx]+sh[7][tx];
float b=sh2[0][tx]+sh2[1][tx]+sh2[2][tx]+sh2[3][tx]+sh2[4][tx]+sh2[5][tx]+sh2[6][tx]+sh2[7][tx];
psum[(long long)g*N+c]=a; psum2[(long long)g*N+c]=b;
}
}
__global__ void finalize_train_affine_kernel(const float* __restrict__ psum, const float* __restrict__ psum2,
const float* __restrict__ bn_w, const float* __restrict__ bn_b, const float* __restrict__ scale,
float* __restrict__ alpha, float* __restrict__ beta,
float* __restrict__ running_mean, float* __restrict__ running_var,
int M, int N, int groups, float eps, float momentum) {
int c=blockIdx.x*blockDim.x+threadIdx.x;
if(c<N){
float s=0.f,s2=0.f; for(int g=0;g<groups;++g){ s+=psum[(long long)g*N+c]; s2+=psum2[(long long)g*N+c]; }
float m=s/(float)M; float var=s2/(float)M-m*m; if(var<0.f)var=0.f; float inv=rsqrtf(var+eps);
float sc=scale[0]; float a=inv*bn_w[c]*sc; alpha[c]=a; beta[c]=(bn_b[c]-m*inv*bn_w[c])*sc;
running_mean[c]=(1.f-momentum)*running_mean[c]+momentum*m;
running_var[c]=(1.f-momentum)*running_var[c]+momentum*(var*((float)M/(float)(M-1)));
}
}
__global__ void eval_affine_kernel(const float* __restrict__ rm, const float* __restrict__ rv,
const float* __restrict__ bn_w, const float* __restrict__ bn_b, const float* __restrict__ scale,
float* __restrict__ alpha, float* __restrict__ beta, int N, float eps){
int c=blockIdx.x*blockDim.x+threadIdx.x; if(c<N){ float inv=rsqrtf(rv[c]+eps); float sc=scale[0]; float a=inv*bn_w[c]*sc; alpha[c]=a; beta[c]=(bn_b[c]-rm[c]*inv*bn_w[c])*sc; }
}
__global__ void softmax_affine_kernel(const float* __restrict__ y, const float* __restrict__ alpha, const float* __restrict__ beta,
float* __restrict__ out, int M, int N) {
int row=blockIdx.x, tid=threadIdx.x; float vmax=-FLT_MAX;
for(int c=tid;c<N;c+=blockDim.x){ float v=y[(long long)row*N+c]*alpha[c]+beta[c]; vmax=fmaxf(vmax,v); }
__shared__ float sh[256]; sh[tid]=vmax; __syncthreads();
for(int s=128;s>0;s>>=1){ if(tid<s) sh[tid]=fmaxf(sh[tid],sh[tid+s]); __syncthreads(); }
float m=sh[0], sum=0.f;
for(int c=tid;c<N;c+=blockDim.x){ float v=y[(long long)row*N+c]*alpha[c]+beta[c]; sum+=__expf(v-m); }
sh[tid]=sum; __syncthreads();
for(int s=128;s>0;s>>=1){ if(tid<s) sh[tid]+=sh[tid+s]; __syncthreads(); }
float invs=1.f/sh[0];
for(int c=tid;c<N;c+=blockDim.x){ float v=y[(long long)row*N+c]*alpha[c]+beta[c]; out[(long long)row*N+c]=__expf(v-m)*invs; }
}
torch::Tensor fused_forward(torch::Tensor input, torch::Tensor weight, torch::Tensor bias,
torch::Tensor bn_weight, torch::Tensor bn_bias,
torch::Tensor running_mean, torch::Tensor running_var, torch::Tensor scale,
bool training, double eps, double momentum) {
CHECK_CUDA(input); CHECK_FLOAT(input); const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
auto y=at::linear(input,weight,bias); int M=(int)y.size(0), N=(int)y.size(1); auto out=torch::empty_like(y); auto opts=y.options();
auto alpha=torch::empty({N},opts); auto beta=torch::empty({N},opts); cudaStream_t stream=at::cuda::getCurrentCUDAStream();
if(training){ int groups=(M+255)/256; auto psum=torch::empty({groups,N},opts); auto psum2=torch::empty({groups,N},opts); dim3 block(32,8); dim3 grid((N+31)/32,groups);
partial_stats_kernel<<<grid,block,0,stream>>>(y.data_ptr<float>(),psum.data_ptr<float>(),psum2.data_ptr<float>(),M,N);
finalize_train_affine_kernel<<<(N+255)/256,256,0,stream>>>(psum.data_ptr<float>(),psum2.data_ptr<float>(),bn_weight.data_ptr<float>(),bn_bias.data_ptr<float>(),scale.data_ptr<float>(),alpha.data_ptr<float>(),beta.data_ptr<float>(),running_mean.data_ptr<float>(),running_var.data_ptr<float>(),M,N,groups,(float)eps,(float)momentum);
} else {
eval_affine_kernel<<<(N+255)/256,256,0,stream>>>(running_mean.data_ptr<float>(),running_var.data_ptr<float>(),bn_weight.data_ptr<float>(),bn_bias.data_ptr<float>(),scale.data_ptr<float>(),alpha.data_ptr<float>(),beta.data_ptr<float>(),N,(float)eps);
}
softmax_affine_kernel<<<M,256,0,stream>>>(y.data_ptr<float>(),alpha.data_ptr<float>(),beta.data_ptr<float>(),out.data_ptr<float>(),M,N);
C10_CUDA_KERNEL_LAUNCH_CHECK(); return out;
}
'''
_ext=load_inline(name='gemm_bn_scale_softmax_ext_v3',cpp_sources=[_cpp_src],cuda_sources=[_cuda_src],functions=['fused_forward'],extra_cuda_cflags=['-O3'],extra_cflags=['-O3'],verbose=False)
class ModelNew(nn.Module):
def __init__(self,in_features,out_features,bn_eps=1e-5,bn_momentum=0.1,scale_shape=(1,)):
super(ModelNew,self).__init__(); self.gemm=nn.Linear(in_features,out_features); self.bn=nn.BatchNorm1d(out_features,eps=bn_eps,momentum=bn_momentum); self.scale=nn.Parameter(torch.ones(scale_shape))
def forward(self,x):
return _ext.fused_forward(x,self.gemm.weight,self.gemm.bias,self.bn.weight,self.bn.bias,self.bn.running_mean,self.bn.running_var,self.scale,self.training,self.bn.eps,self.bn.momentum)