call → run_correctness
callrun_correctness
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\ncpp_src=r'''\n#include <torch/extension.h>\ntorch::Tensor flash_attn(torch::Tensor Q, torch::Tensor K, torch::Tensor V);\n'''\ncuda_src=r'''\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n#include <mma.h>\n#include <math.h>\nusing namespace nvcuda;\n#define L 512\n#define D 1024\n#define BR 16\n#define BC 16\n#define NW 4\n__global__ __launch_bounds__(256,1) void flash_kernel(const float* __restrict__ Q,const float* __restrict__ K,const float* __restrict__ V,float* __restrict__ O,int NH){\n extern __shared__ float sm[];\n float* sO=sm; float* sV=sO+BR*D; float* sPV=sV+BC*D; float* sPart=sPV+BR*D; float* sScores=sPart+NW*BR*BC; float* sP=sScores+BR*BC; float* sAlpha=sP+BR*BC; float* sL=sAlpha+BR; float* sM=sL+BR;\n int tid=threadIdx.x, warp=tid>>5, qb=blockIdx.x*BR, bh=blockIdx.y; long long base=(long long)bh*L*D;\n for(int idx=tid;idx<BR*D;idx+=blockDim.x) sO[idx]=0.f;\n if(tid<BR){sM[tid]=-INFINITY; sL[tid]=0.f; sAlpha[tid]=0.f;} __syncthreads();\n for(int kb=0;kb<L;kb+=BC){\n for(int idx=tid;idx<BC*D;idx+=blockDim.x){int c=idx>>10,d=idx&1023; sV[idx]=V[base+(kb+c)*(long long)D+d];}\n if(warp<NW){\n wmma::fragment<wmma::matrix_a,16,16,8,wmma::precision::tf32,wmma::row_major> af;\n wmma::fragment<wmma::matrix_b,16,16,8,wmma::precision::tf32,wmma::col_major> bf;\n wmma::fragment<wmma::accumulator,16,16,8,float> cf;\n wmma::fill_fragment(cf,0.0f);\n const float* qptr=Q+base+qb*(long long)D; const float* kptr=K+base+kb*(long long)D;\n for(int k0=warp*256;k0<(warp+1)*256;k0+=8){wmma::load_matrix_sync(af,qptr+k0,D); wmma::load_matrix_sync(bf,kptr+k0,D); wmma::mma_sync(cf,af,bf,cf);} \n wmma::store_matrix_sync(sPart+warp*BR*BC,cf,BC,wmma::mem_row_major);\n }\n __syncthreads();\n for(int idx=tid;idx<BR*BC;idx+=blockDim.x){float s=0.f; for(int w=0;w<NW;w++) s+=sPart[w*BR*BC+idx]; sScores[idx]=s*0.03125f;}\n __syncthreads();\n if(tid<BR){int r=tid; float old=sM[r],mx=old; for(int c=0;c<BC;c++) mx=fmaxf(mx,sScores[r*BC+c]); float a=expf(old-mx),ps=0.f; for(int c=0;c<BC;c++){float p=expf(sScores[r*BC+c]-mx); sP[r*BC+c]=p; ps+=p;} sAlpha[r]=a; sL[r]=sL[r]*a+ps; sM[r]=mx;}\n __syncthreads();\n // tensor-core compute sPV = sP(16x16) * sV(16x1024)\n for(int dt=warp; dt<64; dt+=8){\n wmma::fragment<wmma::matrix_a,16,16,8,wmma::precision::tf32,wmma::row_major> af;\n wmma::fragment<wmma::matrix_b,16,16,8,wmma::precision::tf32,wmma::row_major> bf;\n wmma::fragment<wmma::accumulator,16,16,8,float> cf;\n wmma::fill_fragment(cf,0.0f);\n wmma::load_matrix_sync(af,sP,BC);\n wmma::load_matrix_sync(bf,sV+dt*16,D);\n wmma::mma_sync(cf,af,bf,cf);\n wmma::load_matrix_sync(af,sP+8,BC);\n wmma::load_matrix_sync(bf,sV+8*D+dt*16,D);\n wmma::mma_sync(cf,af,bf,cf);\n wmma::store_matrix_sync(sPV+dt*16,cf,D,wmma::mem_row_major);\n }\n __syncthreads();\n for(int idx=tid;idx<BR*D;idx+=blockDim.x){int r=idx>>10; sO[idx]=sO[idx]*sAlpha[r]+sPV[idx];}\n __syncthreads();\n }\n for(int idx=tid;idx<BR*D;idx+=blockDim.x){int r=idx>>10,d=idx&1023; O[base+(qb+r)*(long long)D+d]=sO[idx]/sL[r];}\n}\ntorch::Tensor flash_attn(torch::Tensor Q,torch::Tensor K,torch::Tensor V){auto O=torch::empty_like(Q); int NH=Q.size(0)*Q.size(1); size_t smem=(size_t)(3*BR*D+(NW+2)*BR*BC+3*BR)*sizeof(float); cudaFuncSetAttribute(flash_kernel,cudaFuncAttributeMaxDynamicSharedMemorySize,(int)smem); dim3 grid(L/BR,NH); flash_kernel<<<grid,256,smem>>>(Q.data_ptr<float>(),K.data_ptr<float>(),V.data_ptr<float>(),O.data_ptr<float>(),NH); return O;}\n'''\next=load_inline(name='flash_wmma_v3', cpp_sources=cpp_src, cuda_sources=cuda_src, functions=['flash_attn'], verbose=False)\nclass ModelNew(nn.Module):\n def __init__(self): super().__init__()\n def forward(self,Q,K,V): return ext.flash_attn(Q,K,V)\n"
}