"""Persistent, float32-accumulating CUDA decode for the fixed Qwen geometry.

The cooperative grid executes complete steps, including all blocks. Attention
shares KV loads between the two grouped query heads and partitions the context.
Weights and caches stay bf16; all activation math inside a block stays float32.
"""
from __future__ import annotations

import os
import ctypes
import hashlib
from pathlib import Path
import subprocess
import torch
from torch import nn

HIDDEN, INTERMEDIATE, NUM_Q, NUM_KV, HEAD_DIM, NUM_LAYERS = 1024, 3072, 16, 8, 128, 4
EPS = 1e-6
_SPLIT = 64
_ext = None

_CUDA = r'''
#include <algorithm>
#include <cuda_runtime.h>
#include <cuda_bf16.h>
#include <cooperative_groups.h>
namespace cg = cooperative_groups;
using bf = __nv_bfloat16;
constexpr int NT=256, NW=8, SPLIT=64;
struct Weights {
    const bf *ln, *q, *k, *v, *qn, *kn, *o, *post, *gate, *up, *down;
};
__device__ __forceinline__ float ws(float x) {
    #pragma unroll
    for(int d=16;d;d>>=1) x+=__shfl_down_sync(0xffffffff,x,d);
    return __shfl_sync(0xffffffff,x,0);
}
__device__ __forceinline__ float wm(float x) {
    #pragma unroll
    for(int d=16;d;d>>=1) x=fmaxf(x,__shfl_down_sync(0xffffffff,x,d));
    return __shfl_sync(0xffffffff,x,0);
}
__device__ __forceinline__ float bs(float x,float *tmp) {
    x=ws(x);
    __syncthreads();
    if((threadIdx.x&31)==0) tmp[threadIdx.x>>5]=x;
    __syncthreads();
    if(threadIdx.x<32) {
        x=ws(threadIdx.x<NW?tmp[threadIdx.x]:0.f);
        if(threadIdx.x==0) tmp[0]=x;
    }
    __syncthreads();
    return tmp[0];
}
__device__ __forceinline__ float bm(float x,float *tmp) {
    x=wm(x);
    __syncthreads();
    if((threadIdx.x&31)==0) tmp[threadIdx.x>>5]=x;
    __syncthreads();
    if(threadIdx.x<32) {
        x=wm(threadIdx.x<NW?tmp[threadIdx.x]:-INFINITY);
        if(threadIdx.x==0) tmp[0]=x;
    }
    __syncthreads();
    return tmp[0];
}
__device__ __forceinline__ float mixed(const bf *hidden,const bf *noise,int i,bool mix) {
    float x=__bfloat162float(hidden[i]);
    if(mix) {
        float a=__bfloat162float(__float2bfloat16(.5f*x));
        float b=__bfloat162float(__float2bfloat16(.5f*__bfloat162float(noise[i])));
        x=__bfloat162float(__float2bfloat16(a+b));
    }
    return x;
}
template<int K>
__device__ __forceinline__ float dot(const bf *row,const float *x) {
    int lane=threadIdx.x&31;
    float sum=0.f;
    #pragma unroll
    for(int c=lane*8;c<K;c+=256) {
        uint4 u=__ldg((const uint4*)(row+c));
        const bf *w=(const bf*)&u;
        #pragma unroll
        for(int j=0;j<8;j++) sum=fmaf(__bfloat162float(w[j]),x[c+j],sum);
    }
    return ws(sum);
}
__device__ __forceinline__ void gateup(const bf *g,const bf *u,const float *x,float &gs,float &us) {
    int lane=threadIdx.x&31;
    gs=0.f;us=0.f;
    #pragma unroll
    for(int c=lane*8;c<1024;c+=256) {
        uint4 gv=__ldg((const uint4*)(g+c)), uv=__ldg((const uint4*)(u+c));
        const bf *gp=(const bf*)&gv, *up=(const bf*)&uv;
        #pragma unroll
        for(int j=0;j<8;j++) {
            float a=x[c+j];
            gs=fmaf(__bfloat162float(gp[j]),a,gs);
            us=fmaf(__bfloat162float(up[j]),a,us);
        }
    }
    gs=ws(gs);us=ws(us);
}
__device__ void attention(const float *q,const bf *kc,const bf *vc,float *parts,
                          float *s,int pos,int maxseq,int maxparts) {
    int tid=threadIdx.x,lane=tid&31,warp=tid>>5;
    int np=(pos+SPLIT)/SPLIT;
    for(int part=blockIdx.x;part<8*np;part+=gridDim.x) {
        int kh=part/np, sp=part%np;
        int begin=sp*SPLIT, end=min(begin+SPLIT,pos+1);
        int accbase=SPLIT*2, metabase=accbase+2048;
        float qa[4],qb[4];
        #pragma unroll
        for(int j=0;j<4;j++) {
            qa[j]=q[(kh*2)*128+lane*4+j];
            qb[j]=q[(kh*2+1)*128+lane*4+j];
        }
        float ma=-INFINITY,mb=-INFINITY;
        for(int t=begin+warp;t<end;t+=NW) {
            size_t off=((size_t)kh*maxseq+t)*128+lane*4;
            uint2 ku=__ldg((const uint2*)(kc+off));
            const bf *kp=(const bf*)&ku;
            float sa=0,sb=0;
            #pragma unroll
            for(int j=0;j<4;j++) {
                float k=__bfloat162float(kp[j]);
                sa=fmaf(qa[j],k,sa);sb=fmaf(qb[j],k,sb);
            }
            sa=ws(sa)*0.08838834764831844f;
            sb=ws(sb)*0.08838834764831844f;
            ma=fmaxf(ma,sa);mb=fmaxf(mb,sb);
            if(lane==0) {s[(t-begin)*2]=sa;s[(t-begin)*2+1]=sb;}
        }
        if(lane==0) {s[metabase+warp*2]=ma;s[metabase+warp*2+1]=mb;}
        __syncthreads();
        ma=-INFINITY;mb=-INFINITY;
        #pragma unroll
        for(int w=0;w<NW;w++) {
            ma=fmaxf(ma,s[metabase+w*2]);mb=fmaxf(mb,s[metabase+w*2+1]);
        }
        float la=0,lb=0;
        for(int t=tid;t<SPLIT;t+=NT) {
            float pa=t<end-begin?__expf(s[t*2]-ma):0.f;
            float pb=t<end-begin?__expf(s[t*2+1]-mb):0.f;
            s[t*2]=pa;s[t*2+1]=pb;la+=pa;lb+=pb;
        }
        la=bs(la,s+3104);lb=bs(lb,s+3104);
        float oa[4]={0,0,0,0},ob[4]={0,0,0,0};
        for(int t=begin+warp;t<end;t+=NW) {
            size_t off=((size_t)kh*maxseq+t)*128+lane*4;
            uint2 vu=__ldg((const uint2*)(vc+off));
            const bf *vp=(const bf*)&vu;
            float pa=s[(t-begin)*2],pb=s[(t-begin)*2+1];
            #pragma unroll
            for(int j=0;j<4;j++) {
                float v=__bfloat162float(vp[j]);
                oa[j]=fmaf(pa,v,oa[j]);ob[j]=fmaf(pb,v,ob[j]);
            }
        }
        #pragma unroll
        for(int j=0;j<4;j++) {
            s[accbase+warp*256+lane*4+j]=oa[j];
            s[accbase+warp*256+128+lane*4+j]=ob[j];
        }
        __syncthreads();
        int h=tid/128,d=tid%128;
        float o=0;
        #pragma unroll
        for(int w=0;w<NW;w++) o+=s[accbase+w*256+h*128+d];
        float *dst=parts+((kh*maxparts+sp)*2+h)*132;
        dst[d]=o;
        if(d==0) {dst[128]=h==0?ma:mb;dst[129]=h==0?la:lb;}
        __syncthreads();
    }
}
__device__ void merge(const float *parts,float *out,float *s,int pos,int maxparts) {
    int h=blockIdx.x,tid=threadIdx.x,lane=tid&31,warp=tid>>5;
    if(h<16) {
        int kh=h/2,h2=h%2,np=(pos+SPLIT)/SPLIT;
        const float *p=parts+(kh*maxparts*2+h2)*132;
        float m=-INFINITY;
        for(int i=tid;i<np;i+=NT) m=fmaxf(m,p[i*264+128]);
        m=bm(m,s+3104);
        float l=0;
        for(int i=tid;i<np;i+=NT) {
            float r=__expf(p[i*264+128]-m);
            s[i]=r;l=fmaf(r,p[i*264+129],l);
        }
        l=bs(l,s+3104);
        if(np<1024) {
            if(tid<128) {
                float o=0;
                for(int i=0;i<np;i++) o=fmaf(s[i],p[i*264+tid],o);
                out[h*128+tid]=o/l;
            }
        } else {
            float a=0,b=0,c=0,d=0;
            for(int i=warp;i<np;i+=NW) {
                float4 v=*(const float4*)(p+i*264+lane*4);
                float r=s[i];
                a=fmaf(r,v.x,a);b=fmaf(r,v.y,b);c=fmaf(r,v.z,c);d=fmaf(r,v.w,d);
            }
            int base=4096+warp*128+lane*4;
            s[base]=a;s[base+1]=b;s[base+2]=c;s[base+3]=d;
            __syncthreads();
            if(tid<128) {
                float o=0;
                #pragma unroll
                for(int w=0;w<NW;w++) o+=s[4096+w*128+tid];
                out[h*128+tid]=o/l;
            }
        }
    }
}
__global__ void __launch_bounds__(NT,2) mega_decode(
    bf *hidden,const bf *noise,const Weights *weights,bf *const *caches,float *work,
    const float *rope,int maxseq,int maxparts,int start,int steps,int layers) {
    cg::grid_group grid=cg::this_grid();
    __shared__ __align__(16) float s[5120];
    int tid=threadIdx.x,lane=tid&31,warp=tid>>5;
    float *q=work,*k=q+2048,*v=k+1024,*att=v+1024,*res=att+2048,*mlp=res+1024,*parts=mlp+3072;
    for(int step=0;step<steps;step++) {
        int pos=start+step;
        const bf *n=noise+(size_t)step*1024;
        for(int layer=0;layer<layers;layer++) {
            Weights w=weights[layer];
            bf *kc=caches[layer*2],*vc=caches[layer*2+1];
            float sum=0;
            #pragma unroll
            for(int i=tid;i<1024;i+=NT) {
                float x=mixed(hidden,n,i,layer==0);s[i]=x;sum=fmaf(x,x,sum);
            }
            float r=rsqrtf(bs(sum,s+3104)/1024.f+1e-6f);
            #pragma unroll
            for(int i=tid;i<1024;i+=NT) s[i]=s[i]*r*__bfloat162float(w.ln[i]);
            __syncthreads();
            for(int row=blockIdx.x*NW+warp;row<4096;row+=gridDim.x*NW) {
                const bf *wr=row<2048?w.q+(size_t)row*1024:
                             row<3072?w.k+(size_t)(row-2048)*1024:w.v+(size_t)(row-3072)*1024;
                float val=dot<1024>(wr,s);
                if(lane==0) work[row]=val;
            }
            grid.sync();
            // Each head has an independent normalization and rotary transform.
            if(blockIdx.x<24) {
                int head=blockIdx.x;
                float *x=head<16?q+head*128:k+(head-16)*128;
                const bf *nw=head<16?w.qn:w.kn;
                float z=tid<128?x[tid]:0;
                float rstd=rsqrtf(bs(z*z,s+3104)/128.f+1e-6f);
                if(tid<128) s[tid]=z*rstd*__bfloat162float(nw[tid]);
                __syncthreads();
                if(tid<64) {
                    float cs=rope[(size_t)pos*128+tid],sn=rope[(size_t)pos*128+64+tid];
                    float a=s[tid],b=s[tid+64];
                    float lo=a*cs-b*sn,hi=a*sn+b*cs;
                    if(head<16) {x[tid]=lo;x[tid+64]=hi;}
                    else {
                        int kh=head-16;
                        size_t off=((size_t)kh*maxseq+pos)*128+tid;
                        kc[off]=__float2bfloat16(lo);kc[off+64]=__float2bfloat16(hi);
                        vc[off]=__float2bfloat16(v[kh*128+tid]);
                        vc[off+64]=__float2bfloat16(v[kh*128+tid+64]);
                    }
                }
            }
            grid.sync();
            attention(q,kc,vc,parts,s,pos,maxseq,maxparts);
            grid.sync();
            merge(parts,att,s,pos,maxparts);
            grid.sync();
            for(int i=tid;i<2048;i+=NT) s[i]=att[i];
            __syncthreads();
            for(int row=blockIdx.x*NW+warp;row<1024;row+=gridDim.x*NW) {
                float val=dot<2048>(w.o+(size_t)row*2048,s);
                if(lane==0) res[row]=val+mixed(hidden,n,row,layer==0);
            }
            grid.sync();
            sum=0;
            for(int i=tid;i<1024;i+=NT) {float x=res[i];s[i]=x;sum=fmaf(x,x,sum);}
            r=rsqrtf(bs(sum,s+3104)/1024.f+1e-6f);
            for(int i=tid;i<1024;i+=NT) s[i]=s[i]*r*__bfloat162float(w.post[i]);
            __syncthreads();
            for(int row=blockIdx.x*NW+warp;row<3072;row+=gridDim.x*NW) {
                float gs,us;gateup(w.gate+(size_t)row*1024,w.up+(size_t)row*1024,s,gs,us);
                if(lane==0) mlp[row]=(gs/(1.f+__expf(-gs)))*us;
            }
            grid.sync();
            for(int i=tid;i<3072;i+=NT) s[i]=mlp[i];
            __syncthreads();
            for(int row=blockIdx.x*NW+warp;row<1024;row+=gridDim.x*NW) {
                float val=dot<3072>(w.down+(size_t)row*3072,s);
                if(lane==0) hidden[row]=__float2bfloat16(res[row]+val);
            }
            grid.sync();
        }
    }
}
extern "C" int execute(void *hidden,void *noise,void *weights,void *caches,
                       void *work,void *rope,int maxseq,int start,int steps,int layers,
                       void *stream,int requested_blocks) {
    bf *h=(bf*)hidden;const bf *n=(const bf*)noise;
    const Weights *w=(const Weights*)weights;
    bf *const *c=(bf *const*)caches;
    float *b=(float*)work;const float *r=(const float*)rope;
    int ms=maxseq,mp=(ms+SPLIT-1)/SPLIT,st=start,ns=steps,nl=layers;
    int dev=0,sms=0,active=0;
    cudaGetDevice(&dev);
    cudaDeviceGetAttribute(&sms,cudaDevAttrMultiProcessorCount,dev);
    cudaOccupancyMaxActiveBlocksPerMultiprocessor(&active,mega_decode,NT,0);
    int blocks=requested_blocks?requested_blocks:sms*std::min(active,2);
    if(blocks>sms*active) return (int)cudaErrorCooperativeLaunchTooLarge;
    void *args[]={&h,&n,&w,&c,&b,&r,&ms,&mp,&st,&ns,&nl};
    return (int)cudaLaunchCooperativeKernel((void*)mega_decode,blocks,NT,args,0,(cudaStream_t)stream);
}
'''


def _extension():
    global _ext
    if _ext is None:
        key = hashlib.sha256((_CUDA + "sm120-O3-fast").encode()).hexdigest()[:16]
        root = Path(os.environ.get("TORCH_EXTENSIONS_DIR", str(Path.home() / ".cache/torch_extensions")))
        build = root / ("qwen_cuda_" + key)
        build.mkdir(parents=True, exist_ok=True)
        source, binary = build / "kernel.cu", build / "kernel.so"
        if not binary.exists():
            source.write_text(_CUDA)
            cuda_home = os.environ.get("CUDA_HOME", "/usr/local/cuda")
            result = subprocess.run(
                [str(Path(cuda_home) / "bin/nvcc"), "-O3", "--use_fast_math", "-lineinfo",
                 "-arch=sm_120", "-shared", "-Xcompiler=-fPIC", str(source), "-o", str(binary)],
                capture_output=True, text=True,
            )
            if result.returncode:
                raise RuntimeError("CUDA build failed:\n" + result.stdout + result.stderr)
        _ext = ctypes.CDLL(str(binary))
        _ext.execute.argtypes = [ctypes.c_void_p] * 6 + [ctypes.c_int] * 4 + [ctypes.c_void_p, ctypes.c_int]
        _ext.execute.restype = ctypes.c_int
    return _ext


class Block(nn.Module):
    def __init__(self):
        super().__init__()
        self.input_ln = nn.Parameter(torch.ones(1024, dtype=torch.bfloat16))
        self.q_proj = nn.Parameter(torch.empty(2048, 1024, dtype=torch.bfloat16))
        self.k_proj = nn.Parameter(torch.empty(1024, 1024, dtype=torch.bfloat16))
        self.v_proj = nn.Parameter(torch.empty(1024, 1024, dtype=torch.bfloat16))
        self.q_norm = nn.Parameter(torch.ones(128, dtype=torch.bfloat16))
        self.k_norm = nn.Parameter(torch.ones(128, dtype=torch.bfloat16))
        self.o_proj = nn.Parameter(torch.empty(1024, 2048, dtype=torch.bfloat16))
        self.post_ln = nn.Parameter(torch.ones(1024, dtype=torch.bfloat16))
        self.gate_proj = nn.Parameter(torch.empty(3072, 1024, dtype=torch.bfloat16))
        self.up_proj = nn.Parameter(torch.empty(3072, 1024, dtype=torch.bfloat16))
        self.down_proj = nn.Parameter(torch.empty(1024, 3072, dtype=torch.bfloat16))
        for p in self.parameters():
            if p.ndim == 2:
                nn.init.normal_(p, std=0.02)


class Model(nn.Module):
    def __init__(self, num_layers=NUM_LAYERS, max_seq=131072):
        super().__init__()
        self.num_layers, self.max_seq = num_layers, max_seq
        self.blocks = nn.ModuleList([Block() for _ in range(num_layers)])

    def forward(self, x, k_caches, v_caches, position):
        # A single block-stack evaluation without the input-stream mixing.
        # Two identical bf16 operands make the mixing operation recover x.
        return _execute(self, x, k_caches, v_caches, position, x[None], mix_input=False)


def _prepare(model):
    _extension()
    device = next(model.parameters()).device
    ptrs = [p.data_ptr() for b in model.blocks for p in b.parameters()]
    signature = (tuple(ptrs), str(device), model.max_seq)
    if getattr(model, "_cuda_signature", None) != signature:
        model._weights = torch.tensor(ptrs, dtype=torch.int64, device=device)
        inv = 1.0 / (10000 ** (torch.arange(64, device=device, dtype=torch.float32) / 64))
        freq = torch.arange(model.max_seq, device=device, dtype=torch.float32)[:, None] * inv[None, :]
        model._rope = torch.cat((freq.cos(), freq.sin()), dim=1).contiguous()
        maxparts = (model.max_seq + _SPLIT - 1) // _SPLIT
        model._work = torch.empty(10240 + 8 * maxparts * 2 * 132, device=device, dtype=torch.float32)
        model._cuda_signature = signature


def _noise(n_steps, seed, device):
    g = torch.Generator(device="cpu").manual_seed(seed)
    return torch.randn(n_steps, HIDDEN, generator=g, dtype=torch.bfloat16).to(device)


def _execute(model, hidden, k_caches, v_caches, start_pos, noise, mix_input=True):
    _prepare(model)
    n_steps = noise.shape[0]
    if start_pos < 0 or start_pos + n_steps > model.max_seq:
        raise ValueError("decode positions exceed model.max_seq")
    h = hidden.clone().contiguous()
    if not mix_input:
        # The mixing operation is exact for two identical bf16 operands.
        noise = hidden[None].contiguous()
    cp = torch.tensor([p.data_ptr() for pair in zip(k_caches, v_caches) for p in pair],
                      device=h.device, dtype=torch.int64)
    err = _extension().execute(h.data_ptr(), noise.data_ptr(), model._weights.data_ptr(),
                               cp.data_ptr(), model._work.data_ptr(), model._rope.data_ptr(),
                               model.max_seq, start_pos, n_steps, model.num_layers,
                               torch.cuda.current_stream(h.device).cuda_stream, 0)
    if err:
        raise RuntimeError(f"CUDA cooperative decode failed with error {err}")
    return h, k_caches, v_caches


@torch.no_grad()
def prefill(model, ctx_len, seed, device=None):
    device = device or next(model.parameters()).device
    if torch.device(device).type != "cuda":
        device = torch.device("cuda")
    model = model.to(device).eval()
    if not 0 <= ctx_len <= model.max_seq:
        raise ValueError("invalid context length")
    _prepare(model)
    g = torch.Generator(device="cpu").manual_seed(seed)
    h = torch.randn(HIDDEN, generator=g, dtype=torch.bfloat16).to(device)
    k = [torch.zeros(NUM_KV, model.max_seq, HEAD_DIM, device=device, dtype=torch.bfloat16)
         for _ in range(model.num_layers)]
    v = [torch.zeros_like(t) for t in k]
    if ctx_len:
        h, k, v = _execute(model, h, k, v, 0, _noise(ctx_len, seed + 1, device))
    return h, k, v


@torch.no_grad()
def decode_steps(model, hidden, k_caches, v_caches, start_pos, n_steps, seed):
    if not n_steps:
        return hidden, k_caches, v_caches
    return _execute(model, hidden, k_caches, v_caches, start_pos,
                    _noise(n_steps, seed + 2, hidden.device))


@torch.no_grad()
def run(ctx_len, decode_steps, seed, model=None, max_seq=None):
    if model is None:
        model = Model(NUM_LAYERS, max_seq or max(ctx_len + decode_steps, 512)).cuda().eval()
    if model.max_seq < ctx_len + decode_steps:
        raise ValueError("model.max_seq is too small")
    device = next(model.parameters()).device
    if device.type != "cuda":
        model = model.cuda().eval()
    h, k, v = prefill(model, ctx_len, seed)
    h, k, v = globals()["decode_steps"](model, h, k, v, ctx_len, decode_steps, seed)
    return {"last_hidden": h.detach(), "ctx_len": ctx_len, "decode_steps": decode_steps}


def get_init_inputs():
    return [NUM_LAYERS, 131072]


def get_inputs():
    return []