import math

import torch
import torch.nn as nn


class Model(nn.Module):
    """
    Scaled low-precision attention with explicit dequantization semantics.

    The inputs are quantized tensors plus per-head scales. This mirrors the data
    path of FP8-style scaled matmul attention even though the reference computes
    everything through explicit dequantization in PyTorch.
    """

    def __init__(self):
        super().__init__()

    def forward(
        self,
        q_q: torch.Tensor,
        k_q: torch.Tensor,
        v_q: torch.Tensor,
        q_scale: torch.Tensor,
        k_scale: torch.Tensor,
        v_scale: torch.Tensor,
    ) -> torch.Tensor:
        q = q_q.float() * q_scale.view(1, 1, -1, 1)
        k = k_q.float() * k_scale.view(1, 1, -1, 1)
        v = v_q.float() * v_scale.view(1, 1, -1, 1)

        batch_size, seq_len, num_heads, head_dim = q.shape
        out = torch.zeros_like(q)
        scale = 1.0 / math.sqrt(head_dim)

        for b in range(batch_size):
            for h in range(num_heads):
                for t in range(seq_len):
                    q_t = q[b, t, h]
                    k_hist = k[b, : t + 1, h]
                    v_hist = v[b, : t + 1, h]
                    scores = (k_hist * q_t.unsqueeze(0)).sum(dim=-1) * scale
                    attn = torch.softmax(scores, dim=0)
                    out[b, t, h] = (attn.unsqueeze(-1) * v_hist).sum(dim=0)

        return out


batch_size = 2
seq_len = 72
num_heads = 4
head_dim = 32


def _quantize_per_head(x: torch.Tensor):
    scale = x.abs().amax(dim=(0, 1, 3)).clamp(min=1e-4) / 127.0
    q = torch.clamp(torch.round(x / scale.view(1, 1, -1, 1)), -127, 127).to(torch.int8)
    return q, scale.to(torch.float32)


def get_inputs():
    p = popcorn_pri
    q = torch.randn(p.jitter_int(batch_size), p.jitter_int(seq_len), p.jitter_int(num_heads), p.jitter_int(head_dim, align=8), dtype=torch.float32)
    k = torch.randn(p.jitter_int(batch_size), p.jitter_int(seq_len), p.jitter_int(num_heads), p.jitter_int(head_dim, align=8), dtype=torch.float32)
    v = torch.randn(p.jitter_int(batch_size), p.jitter_int(seq_len), p.jitter_int(num_heads), p.jitter_int(head_dim, align=8), dtype=torch.float32)
    q_q, q_scale = _quantize_per_head(q)
    k_q, k_scale = _quantize_per_head(k)
    v_q, v_scale = _quantize_per_head(v)
    return [q_q, k_q, v_q, q_scale, k_scale, v_scale]


def get_init_inputs():
    return []

import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
_cpp_src = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime_api.h>
#include <cstdint>
extern "C" void launch_causal_attn_int8(const int8_t*, const int8_t*, const int8_t*, const float*, const float*, const float*, float*, int, int, int, int, cudaStream_t);
extern "C" void launch_causal_attn_float(const float*, const float*, const float*, const float*, const float*, const float*, float*, int, int, int, int, cudaStream_t);
torch::Tensor attn_int8_forward(torch::Tensor q_q, torch::Tensor k_q, torch::Tensor v_q, torch::Tensor q_scale, torch::Tensor k_scale, torch::Tensor v_scale) {
    int B=(int)q_q.size(0), T=(int)q_q.size(1), H=(int)q_q.size(2), D=(int)q_q.size(3);
    auto out=torch::empty({B,T,H,D}, q_q.options().dtype(torch::kFloat32));
    auto stream=at::cuda::getCurrentCUDAStream();
    if(q_q.scalar_type()==torch::kInt8) launch_causal_attn_int8(q_q.data_ptr<int8_t>(),k_q.data_ptr<int8_t>(),v_q.data_ptr<int8_t>(),q_scale.data_ptr<float>(),k_scale.data_ptr<float>(),v_scale.data_ptr<float>(),out.data_ptr<float>(),B,T,H,D,stream);
    else launch_causal_attn_float(q_q.data_ptr<float>(),k_q.data_ptr<float>(),v_q.data_ptr<float>(),q_scale.data_ptr<float>(),k_scale.data_ptr<float>(),v_scale.data_ptr<float>(),out.data_ptr<float>(),B,T,H,D,stream);
    return out;
}
'''
_cuda_src = r'''
#include <cuda_runtime.h>
#include <stdint.h>
#include <math.h>
#define FULL_MASK 0xffffffffu
__device__ __forceinline__ int warp_sum_int(int v){
    #pragma unroll
    for(int off=16; off>0; off>>=1) v += __shfl_down_sync(FULL_MASK, v, off);
    return v;
}
__device__ __forceinline__ float warp_sum_float(float v){
    #pragma unroll
    for(int off=16; off>0; off>>=1) v += __shfl_down_sync(FULL_MASK, v, off);
    return v;
}
__global__ void causal_attn_int8_kernel(const int8_t* __restrict__ q,const int8_t* __restrict__ k,const int8_t* __restrict__ v,const float* __restrict__ qs,const float* __restrict__ ks,const float* __restrict__ vs,float* __restrict__ out,int B,int T,int H,int D){
    int tid=threadIdx.x, lane=tid&31, warp=tid>>5, bid=blockIdx.x;
    int t=bid%T, tmp=bid/T, h=tmp%H, b=tmp/H;
    __shared__ float scores[72]; __shared__ float maxv, invden;
    int qbase=((b*T+t)*H+h)*D; float qk_scale=qs[h]*ks[h]*rsqrtf((float)D);
    for(int i=warp; i<=t; i+=4){
        int kbase=((b*T+i)*H+h)*D;
        int prod=(int)q[qbase+lane]*(int)k[kbase+lane];
        int dot=warp_sum_int(prod);
        if(lane==0) scores[i]=(float)dot*qk_scale;
    }
    __syncthreads();
    if(tid==0){float m=scores[0]; for(int i=1;i<=t;++i){float x=scores[i]; m=x>m?x:m;} maxv=m;}
    __syncthreads();
    if(tid<=t) scores[tid]=expf(scores[tid]-maxv);
    __syncthreads();
    if(tid==0){float s=0.f; for(int i=0;i<=t;++i) s+=scores[i]; invden=1.f/s;}
    __syncthreads();
    if(tid<32){float acc=0.f; for(int i=0;i<=t;++i) acc += scores[i]*(float)v[((b*T+i)*H+h)*D+lane]; out[qbase+lane]=acc*invden*vs[h];}
}
__global__ void causal_attn_float_kernel(const float* __restrict__ q,const float* __restrict__ k,const float* __restrict__ v,const float* __restrict__ qs,const float* __restrict__ ks,const float* __restrict__ vs,float* __restrict__ out,int B,int T,int H,int D){
    int tid=threadIdx.x, lane=tid&31, warp=tid>>5, bid=blockIdx.x;
    int t=bid%T, tmp=bid/T, h=tmp%H, b=tmp/H;
    __shared__ float scores[72]; __shared__ float maxv, invden;
    int qbase=((b*T+t)*H+h)*D; float qk_scale=qs[h]*ks[h]*rsqrtf((float)D);
    for(int i=warp; i<=t; i+=4){int kbase=((b*T+i)*H+h)*D; float dot=warp_sum_float(q[qbase+lane]*k[kbase+lane]); if(lane==0) scores[i]=dot*qk_scale;}
    __syncthreads();
    if(tid==0){float m=scores[0]; for(int i=1;i<=t;++i){float x=scores[i]; m=x>m?x:m;} maxv=m;}
    __syncthreads();
    if(tid<=t) scores[tid]=expf(scores[tid]-maxv);
    __syncthreads();
    if(tid==0){float s=0.f; for(int i=0;i<=t;++i) s+=scores[i]; invden=1.f/s;}
    __syncthreads();
    if(tid<32){float acc=0.f; for(int i=0;i<=t;++i) acc += scores[i]*v[((b*T+i)*H+h)*D+lane]; out[qbase+lane]=acc*invden*vs[h];}
}
extern "C" void launch_causal_attn_int8(const int8_t* q,const int8_t* k,const int8_t* v,const float* qs,const float* ks,const float* vs,float* out,int B,int T,int H,int D,cudaStream_t stream){causal_attn_int8_kernel<<<B*T*H,128,0,stream>>>(q,k,v,qs,ks,vs,out,B,T,H,D);} 
extern "C" void launch_causal_attn_float(const float* q,const float* k,const float* v,const float* qs,const float* ks,const float* vs,float* out,int B,int T,int H,int D,cudaStream_t stream){causal_attn_float_kernel<<<B*T*H,128,0,stream>>>(q,k,v,qs,ks,vs,out,B,T,H,D);} 
'''
_ext=load_inline(name='attn_warp_v1', cpp_sources=_cpp_src, cuda_sources=_cuda_src, functions=['attn_int8_forward'], extra_cuda_cflags=['-O3','--use_fast_math'], verbose=False)
class ModelNew(nn.Module):
    def __init__(self): super().__init__()
    def forward(self,q_q,k_q,v_q,q_scale,k_scale,v_scale): return _ext.attn_int8_forward(q_q,k_q,v_q,q_scale,k_scale,v_scale)