"""CUDA sparse attention with exact causal block means and warp MMA gather."""
import os
import torch
from torch import nn
from torch.utils.cpp_extension import load_inline

os.environ.setdefault('TORCH_CUDA_ARCH_LIST', '12.0')

_CPP = r'''
#include <torch/extension.h>
torch::Tensor nsa_cuda(torch::Tensor q, torch::Tensor k, torch::Tensor v);
'''
_CUDA = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <cuda_bf16.h>
#include <cuda_runtime.h>
#include <math.h>
using bf16 = __nv_bfloat16;

__device__ __forceinline__ unsigned pack(float x, float y) {
    __nv_bfloat162_raw raw = __floats2bfloat162_rn(x,y); return unsigned(raw.x) | (unsigned(raw.y)<<16);
}
__device__ __forceinline__ void mma(float &c0, float &c1, float &c2, float &c3,
 unsigned a0,unsigned a1,unsigned a2,unsigned a3,unsigned b0,unsigned b1) {
    asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 "
        "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3};"
        : "+f"(c0),"+f"(c1),"+f"(c2),"+f"(c3)
        : "r"(a0),"r"(a1),"r"(a2),"r"(a3),"r"(b0),"r"(b1));
}
__device__ __forceinline__ float warpmax(float x) {
    #pragma unroll
    for(int i=16;i;i/=2) x=fmaxf(x,__shfl_xor_sync(0xffffffff,x,i));
    return x;
}

template<int D> __global__ void means_kernel(const bf16* k, bf16* means, float* prefix,int S,int NB) {
    int d=threadIdx.x, bi=blockIdx.x, bh=blockIdx.y;
    float sum=0;
    int lim=min(64,S-bi*64);
    const bf16* ptr=k+(size_t(bh)*S+bi*64)*D;
    #pragma unroll 1
    for(int j=0;j<lim;j++) {
        sum += __bfloat162float(ptr[j*D+d]);
        prefix[(size_t(bh)*S+bi*64+j)*D+d]=sum/float(j+1);
    }
    float mean=sum/float(lim);
    bf16 hi=__float2bfloat16(mean);
    float rest=mean-__bfloat162float(hi);
    bf16 lo=__float2bfloat16(rest);
    bf16 tail=__float2bfloat16(rest-__bfloat162float(lo));
    means[(size_t(bh)*3*NB+bi)*D+d]=hi;
    means[(size_t(bh)*3*NB+NB+bi)*D+d]=lo;
    means[(size_t(bh)*3*NB+2*NB+bi)*D+d]=tail;
}

template<int D,int MAXB> __device__ __forceinline__ void select_rows(
 const bf16* q,const bf16* compressed,const float* prefix,int* selected,
 bf16* tile,const unsigned (&qfrag)[D/16][4],int S,int NB,int tbase) {
    int lane=threadIdx.x%32,g=lane/4,sub=lane%4;
    int t0=tbase+g,t1=tbase+g+8,cb=tbase/64;
    float partial0=0,partial1=0;
    #pragma unroll
    for(int d=0;d<D/4;d++) {
        int col=4*d+sub;
        if(t0<S) partial0=fmaf(__bfloat162float(q[size_t(t0)*D+col]),prefix[size_t(t0)*D+col],partial0);
        if(t1<S) partial1=fmaf(__bfloat162float(q[size_t(t1)*D+col]),prefix[size_t(t1)*D+col],partial1);
    }
    #pragma unroll
    for(int off=1;off<=2;off*=2) {
        partial0+=__shfl_xor_sync(0xffffffff,partial0,off);
        partial1+=__shfl_xor_sync(0xffffffff,partial1,off);
    }
    float scores[MAXB/8][4];
    #pragma unroll
    for(int i=0;i<MAXB/8;i++) {
        float c0=0,c1=0,c2=0,c3=0;
        if(i*8<cb) {
            #pragma unroll
            for(int j=0;j<3*8*D/8/32;j++) {
                int ix=j*32+lane,part=ix/(8*D/8),row=(ix/(D/8))%8,col=(ix%(D/8))*8;
                unsigned dst=__cvta_generic_to_shared(tile+(part*8+row)*(D+8)+col);
                const bf16* src=compressed+(size_t(part)*NB+i*8+row)*D+col;
                int valid=i*8+row<cb?16:0;
                asm volatile("cp.async.ca.shared.global [%0], [%1], 16, %2;"::"r"(dst),"l"(src),"r"(valid));
            }
            asm volatile("cp.async.commit_group; cp.async.wait_group 0;":::"memory");
            __syncwarp();
            #pragma unroll
            for(int kd=0;kd<D/16;kd++) {
                #pragma unroll
                for(int part=0;part<3;part++) {
                    unsigned b0,b1;
                    unsigned addr=__cvta_generic_to_shared(tile+(part*8+lane%8)*(D+8)+kd*16+((lane/8)%2)*8);
                    asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0,%1}, [%2];":"=r"(b0),"=r"(b1):"r"(addr));
                    mma(c0,c1,c2,c3,qfrag[kd][0],qfrag[kd][1],qfrag[kd][2],qfrag[kd][3],b0,b1);
                }
            }
            __syncwarp();
        }
        int b0=i*8+sub*2,b1=b0+1;
        scores[i][0]=t0>=S || b0>cb?-INFINITY:b0==cb?partial0:c0;
        scores[i][1]=t0>=S || b1>cb?-INFINITY:b1==cb?partial0:c1;
        scores[i][2]=t1>=S || b0>cb?-INFINITY:b0==cb?partial1:c2;
        scores[i][3]=t1>=S || b1>cb?-INFINITY:b1==cb?partial1:c3;
    }
    #pragma unroll
    for(int rank=0;rank<8;rank++) {
        float best0=-INFINITY,best1=-INFINITY;int id0=-1,id1=-1;
        #pragma unroll
        for(int i=0;i<MAXB/8;i++) {
            #pragma unroll
            for(int j=0;j<2;j++) {
                int b=i*8+sub*2+j;
                float x0=scores[i][j],x1=scores[i][j+2];
                if(x0>best0 || (x0==best0 && b>id0)) {best0=x0;id0=b;}
                if(x1>best1 || (x1==best1 && b>id1)) {best1=x1;id1=b;}
            }
        }
        #pragma unroll
        for(int off=1;off<=2;off*=2) {
            float x0=__shfl_xor_sync(0xffffffff,best0,off),x1=__shfl_xor_sync(0xffffffff,best1,off);
            int b0=__shfl_xor_sync(0xffffffff,id0,off),b1=__shfl_xor_sync(0xffffffff,id1,off);
            if(x0>best0 || (x0==best0 && b0>id0)) {best0=x0;id0=b0;}
            if(x1>best1 || (x1==best1 && b1>id1)) {best1=x1;id1=b1;}
        }
        if(sub==0) {selected[g*8+rank]=best0==-INFINITY?-1:id0;selected[(g+8)*8+rank]=best1==-INFINITY?-1:id1;}
        #pragma unroll
        for(int i=0;i<MAXB/8;i++) {
            #pragma unroll
            for(int j=0;j<2;j++) {
                int b=i*8+sub*2+j;
                if(b==id0) scores[i][j]=-INFINITY;
                if(b==id1) scores[i][j+2]=-INFINITY;
            }
        }
    }
}

struct Mask { unsigned a,b,c,d; };
__device__ __forceinline__ void insert(Mask &m,int b) {
    if(b<0) return;
    unsigned bit=1u<<(b&31);
    if(b<32) m.a|=bit;
    else if(b<64) m.b|=bit;
    else if(b<96) m.c|=bit;
    else m.d|=bit;
}
__device__ __forceinline__ unsigned word(const Mask &x,int b) {
    return b<32?x.a:b<64?x.b:b<96?x.c:x.d;
}
template<int D,int MAXB> __global__ void grouped_attention_kernel(const bf16* q,const bf16* k,const bf16* v,const bf16* compressed,const float* prefix,bf16* out,int S,int NB) {
    int tid=threadIdx.x,lane=tid%32,warp=tid/32,g=lane/4,sub=lane%4;
    int tbase=blockIdx.x*16,bh=blockIdx.y;
    int t0=tbase+g,t1=tbase+g+8;
    int last=min(S-1,tbase+15);
    q+=size_t(bh)*S*D; k+=size_t(bh)*S*D;v+=size_t(bh)*S*D;
    compressed+=size_t(bh)*3*NB*D;prefix+=size_t(bh)*S*D;
    __shared__ int selected[16*8];
    extern __shared__ bf16 staging[];
    bf16* qtile=staging+(D==128?4*16*(D+8):0);
    __shared__ unsigned active[4];
    if(tid<4) active[tid]=0;
    __syncthreads();
    #pragma unroll
    for(int i=0;i<16*D/8/128;i++) {
        int ix=i*128+tid,row=ix/(D/8),col=(ix%(D/8))*8;
        unsigned dst=__cvta_generic_to_shared(qtile+row*(D+8)+col);
        const bf16* src=q+size_t(tbase+row)*D+col;
        int valid=tbase+row<S?16:0;
        asm volatile("cp.async.ca.shared.global [%0], [%1], 16, %2;"::"r"(dst),"l"(src),"r"(valid));
    }
    asm volatile("cp.async.commit_group; cp.async.wait_group 0;":::"memory");
    __syncthreads();
    unsigned qfrag[D/16][4];
    #pragma unroll
    for(int kd=0;kd<D/16;kd++) {
        unsigned addr=__cvta_generic_to_shared(qtile+(lane%16)*(D+8)+kd*16+(lane/16)*8);
        asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];"
            :"=r"(qfrag[kd][0]),"=r"(qfrag[kd][1]),"=r"(qfrag[kd][2]),"=r"(qfrag[kd][3]):"r"(addr));
    }
    if(warp==0) {
        bf16* means_tile=staging+(D==64?16*(D+8):0);
        select_rows<D,MAXB>(q,compressed,prefix,selected,means_tile,qfrag,S,NB,tbase);
    }
    __syncthreads();
    if(tid<16 && tbase+tid<S) {
        int t=tbase+tid;
        #pragma unroll
        for(int i=0;i<8;i++) {
            int b=selected[tid*8+i];
            if(b>=0) atomicOr(active+b/32,1u<<(b%32));
        }
        int b=t/64;
        atomicOr(active+b/32,1u<<(b%32));
        b=max(0,t-63)/64;
        atomicOr(active+b/32,1u<<(b%32));
    }
    __syncthreads();
    Mask bits0={},bits1={};
    #pragma unroll
    for(int i=0;i<8;i++) {
        int id0=t0<S?selected[g*8+i]:-1,id1=t1<S?selected[(g+8)*8+i]:-1;
        insert(bits0,id0);insert(bits1,id1);
    }
    __syncthreads();
    bf16* tile=staging+warp*16*(D+8);
    float accum[D/8][4]={};
    float m0=-INFINITY,m1=-INFINITY,l0=0,l1=0;
    constexpr float scale=D==64?0.125f:0.08838834764831844f;
    for(int b=0;b<=last/64;b++) {
        if(!(active[b/32] & (1u<<(b%32)))) continue;
        int start=b*64+warp*16;
        if(start>last) continue;
        #pragma unroll
        for(int i=0;i<16*D/8/32;i++) {
            int ix=i*32+lane,row=ix/(D/8),col=(ix%(D/8))*8;
            unsigned dst=__cvta_generic_to_shared(tile+row*(D+8)+col);
            const bf16* src=k+size_t(start+row)*D+col;
            int valid=start+row<S?16:0;
            asm volatile("cp.async.ca.shared.global [%0], [%1], 16, %2;"::"r"(dst),"l"(src),"r"(valid));
        }
        asm volatile("cp.async.commit_group; cp.async.wait_group 0;":::"memory");
        __syncwarp();
        float s00=0,s01=0,s10=0,s11=0,s02=0,s03=0,s12=0,s13=0;
        #pragma unroll
        for(int kd=0;kd<D/16;kd++) {
            if constexpr(D==128) {
                unsigned addr=__cvta_generic_to_shared(qtile+(lane%16)*(D+8)+kd*16+(lane/16)*8);
                asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];"
                    :"=r"(qfrag[kd][0]),"=r"(qfrag[kd][1]),"=r"(qfrag[kd][2]),"=r"(qfrag[kd][3]):"r"(addr));
            }
            unsigned bk0,bk1,bk2,bk3;
            unsigned addr0=__cvta_generic_to_shared(tile+(lane%8)*(D+8)+kd*16+((lane/8)%2)*8);
            unsigned addr1=__cvta_generic_to_shared(tile+(lane%8+8)*(D+8)+kd*16+((lane/8)%2)*8);
            asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0,%1}, [%2];": "=r"(bk0),"=r"(bk1):"r"(addr0));
            asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0,%1}, [%2];": "=r"(bk2),"=r"(bk3):"r"(addr1));
            mma(s00,s01,s10,s11,qfrag[kd][0],qfrag[kd][1],qfrag[kd][2],qfrag[kd][3],bk0,bk1);
            mma(s02,s03,s12,s13,qfrag[kd][0],qfrag[kd][1],qfrag[kd][2],qfrag[kd][3],bk2,bk3);
        }
        __syncwarp();
        #pragma unroll
        for(int i=0;i<16*D/8/32;i++) {
            int ix=i*32+lane,row=ix/(D/8),col=(ix%(D/8))*8;
            unsigned dst=__cvta_generic_to_shared(tile+row*(D+8)+col);
            const bf16* src=v+size_t(start+row)*D+col;
            int valid=start+row<S?16:0;
            asm volatile("cp.async.ca.shared.global [%0], [%1], 16, %2;"::"r"(dst),"l"(src),"r"(valid));
        }
        asm volatile("cp.async.commit_group;":::"memory");
        bool sel0=(word(bits0,b)>>(b%32))&1u,sel1=(word(bits1,b)>>(b%32))&1u;
        int j=start+2*sub;
        s00=(t0<S && j<=t0 && (sel0 || j>=t0-63))?s00*scale:-INFINITY;
        s01=(t0<S && j+1<=t0 && (sel0 || j+1>=t0-63))?s01*scale:-INFINITY;
        s02=(t0<S && j+8<=t0 && (sel0 || j+8>=t0-63))?s02*scale:-INFINITY;
        s03=(t0<S && j+9<=t0 && (sel0 || j+9>=t0-63))?s03*scale:-INFINITY;
        s10=(t1<S && j<=t1 && (sel1 || j>=t1-63))?s10*scale:-INFINITY;
        s11=(t1<S && j+1<=t1 && (sel1 || j+1>=t1-63))?s11*scale:-INFINITY;
        s12=(t1<S && j+8<=t1 && (sel1 || j+8>=t1-63))?s12*scale:-INFINITY;
        s13=(t1<S && j+9<=t1 && (sel1 || j+9>=t1-63))?s13*scale:-INFINITY;
        float nm0=fmaxf(fmaxf(s00,s01),fmaxf(s02,s03));
        float nm1=fmaxf(fmaxf(s10,s11),fmaxf(s12,s13));
        #pragma unroll
        for(int off=1;off<=2;off*=2) {
            nm0=fmaxf(nm0,__shfl_xor_sync(0xffffffff,nm0,off));
            nm1=fmaxf(nm1,__shfl_xor_sync(0xffffffff,nm1,off));
        }
        nm0=fmaxf(m0,nm0);nm1=fmaxf(m1,nm1);
        float a0=nm0==-INFINITY?1.f:expf(m0-nm0),a1=nm1==-INFINITY?1.f:expf(m1-nm1);
        s00=s00==-INFINITY?0.f:expf(s00-nm0);s01=s01==-INFINITY?0.f:expf(s01-nm0);
        s02=s02==-INFINITY?0.f:expf(s02-nm0);s03=s03==-INFINITY?0.f:expf(s03-nm0);
        s10=s10==-INFINITY?0.f:expf(s10-nm1);s11=s11==-INFINITY?0.f:expf(s11-nm1);
        s12=s12==-INFINITY?0.f:expf(s12-nm1);s13=s13==-INFINITY?0.f:expf(s13-nm1);
        float sum0=s00+s01+s02+s03,sum1=s10+s11+s12+s13;
        #pragma unroll
        for(int off=1;off<=2;off*=2) {
            sum0+=__shfl_xor_sync(0xffffffff,sum0,off);
            sum1+=__shfl_xor_sync(0xffffffff,sum1,off);
        }
        l0=l0*a0+sum0;l1=l1*a1+sum1;m0=nm0;m1=nm1;
        unsigned ap0=pack(s00,s01),ap1=pack(s10,s11),ap2=pack(s02,s03),ap3=pack(s12,s13);
        asm volatile("cp.async.wait_group 0;":::"memory");
        __syncwarp();
        #pragma unroll
        for(int di=0;di<D/8;di++) {
            unsigned bv0,bv1;
            unsigned addr=__cvta_generic_to_shared(tile+(lane%16)*(D+8)+di*8);
            asm volatile("ldmatrix.sync.aligned.m8n8.x2.trans.shared.b16 {%0,%1}, [%2];":"=r"(bv0),"=r"(bv1):"r"(addr));
            float c0=accum[di][0]*a0,c1=accum[di][1]*a0,c2=accum[di][2]*a1,c3=accum[di][3]*a1;
            mma(c0,c1,c2,c3,ap0,ap1,ap2,ap3,bv0,bv1);
            accum[di][0]=c0;accum[di][1]=c1;accum[di][2]=c2;accum[di][3]=c3;
        }
        __syncwarp();
    }
    __syncthreads();
    float* stats=reinterpret_cast<float*>(staging);
    float* sums=stats+64;
    float* partial=stats+128;
    if(sub==0) {
        stats[warp*16+g]=m0;stats[warp*16+g+8]=m1;
        sums[warp*16+g]=l0;sums[warp*16+g+8]=l1;
    }
    __syncthreads();
    float gm0=-INFINITY,gm1=-INFINITY;
    #pragma unroll
    for(int w=0;w<4;w++) {gm0=fmaxf(gm0,stats[w*16+g]);gm1=fmaxf(gm1,stats[w*16+g+8]);}
    float norm0=0,norm1=0;
    #pragma unroll
    for(int w=0;w<4;w++) {
        norm0+=gm0==-INFINITY?0.f:expf(stats[w*16+g]-gm0)*sums[w*16+g];
        norm1+=gm1==-INFINITY?0.f:expf(stats[w*16+g+8]-gm1)*sums[w*16+g+8];
    }
    float factor0=gm0==-INFINITY?0.f:expf(m0-gm0)/norm0;
    float factor1=gm1==-INFINITY?0.f:expf(m1-gm1)/norm1;
    if(warp==0) {
        #pragma unroll
        for(int di=0;di<D/8;di++) {
            int d=di*8+2*sub;
            partial[g*(D+8)+d]=accum[di][0]*factor0;
            partial[g*(D+8)+d+1]=accum[di][1]*factor0;
            partial[(g+8)*(D+8)+d]=accum[di][2]*factor1;
            partial[(g+8)*(D+8)+d+1]=accum[di][3]*factor1;
        }
    }
    __syncthreads();
    if(warp!=0) {
        #pragma unroll
        for(int di=0;di<D/8;di++) {
            int d=di*8+2*sub;
            atomicAdd(partial+g*(D+8)+d,accum[di][0]*factor0);
            atomicAdd(partial+g*(D+8)+d+1,accum[di][1]*factor0);
            atomicAdd(partial+(g+8)*(D+8)+d,accum[di][2]*factor1);
            atomicAdd(partial+(g+8)*(D+8)+d+1,accum[di][3]*factor1);
        }
    }
    __syncthreads();
    for(int ix=tid;ix<16*D;ix+=128) {
        int row=ix/D,d=ix%D;
        if(tbase+row<S) out[(size_t(bh)*S+tbase+row)*D+d]=__float2bfloat16(partial[row*(D+8)+d]);
    }
}

template<int D> void launch(const torch::Tensor& q,const torch::Tensor& k,const torch::Tensor& v,torch::Tensor& out,torch::Tensor& means,torch::Tensor& prefix,int S,int BH,cudaStream_t stream) {
    int nb=(S+63)/64;
    means_kernel<D><<<dim3(nb,BH),D,0,stream>>>((const bf16*)k.data_ptr(),(bf16*)means.data_ptr(),prefix.data_ptr<float>(),S,nb);
    dim3 grid((S+15)/16,BH);
    int smem=(4+(D==128?1:0))*16*(D+8)*sizeof(bf16);
    if(nb<=32) grouped_attention_kernel<D,32><<<grid,128,smem,stream>>>((const bf16*)q.data_ptr(),(const bf16*)k.data_ptr(),(const bf16*)v.data_ptr(),(const bf16*)means.data_ptr(),prefix.data_ptr<float>(),(bf16*)out.data_ptr(),S,nb);
    else if(nb<=64) grouped_attention_kernel<D,64><<<grid,128,smem,stream>>>((const bf16*)q.data_ptr(),(const bf16*)k.data_ptr(),(const bf16*)v.data_ptr(),(const bf16*)means.data_ptr(),prefix.data_ptr<float>(),(bf16*)out.data_ptr(),S,nb);
    else grouped_attention_kernel<D,128><<<grid,128,smem,stream>>>((const bf16*)q.data_ptr(),(const bf16*)k.data_ptr(),(const bf16*)v.data_ptr(),(const bf16*)means.data_ptr(),prefix.data_ptr<float>(),(bf16*)out.data_ptr(),S,nb);
}
torch::Tensor nsa_cuda(torch::Tensor q,torch::Tensor k,torch::Tensor v) {
    c10::cuda::CUDAGuard guard(q.device());
    int S=q.size(2),D=q.size(3),BH=q.size(0)*q.size(1),nb=(S+63)/64;
    TORCH_CHECK(q.is_cuda() && q.scalar_type()==torch::kBFloat16,"bf16 CUDA inputs required");
    TORCH_CHECK(q.is_contiguous() && k.is_contiguous() && v.is_contiguous(),"contiguous inputs required");
    TORCH_CHECK((D==64 || D==128) && nb<=128,"unsupported shape");
    auto out=torch::empty_like(q);
    auto means=torch::empty({BH,3,nb,D},q.options());
    auto prefix=torch::empty({BH,S,D},q.options().dtype(torch::kFloat32));
    auto stream=at::cuda::getCurrentCUDAStream();
    if(D==64) launch<64>(q,k,v,out,means,prefix,S,BH,stream);
    else launch<128>(q,k,v,out,means,prefix,S,BH,stream);
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return out;
}
'''

_ext = load_inline(name='nsa_bf16_cuda', cpp_sources=_CPP, cuda_sources=_CUDA,
                   functions=['nsa_cuda'], extra_cuda_cflags=['-O3', '--use_fast_math', '-lineinfo'],
                   extra_cflags=['-O3'], verbose=False)

class Model(nn.Module):
    def __init__(self, B, H, S, D):
        super().__init__()
        self.B, self.H, self.S, self.D = B, H, S, D
        self.register_buffer('_dummy', torch.zeros(1, dtype=torch.bfloat16))

    def forward(self, q, k, v):
        return _ext.nsa_cuda(q, k, v)