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>
torch::Tensor attention_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);
"""

_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <stdint.h>
#include <float.h>

__global__ void causal_attn_i8_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 bid = blockIdx.x;
    int t = bid % T;
    int tmp = bid / T;
    int h = tmp % H;
    int b = tmp / H;
    int tid = threadIdx.x;
    __shared__ float smem[128];
    float qk_scale = qs[h] * ks[h] * rsqrtf((float)D);
    if (tid <= t) {
        int qoff = ((b * T + t) * H + h) * D;
        int koff = ((b * T + tid) * H + h) * D;
        int isum = 0;
        #pragma unroll
        for (int d = 0; d < 32; ++d) isum += (int)q[qoff + d] * (int)k[koff + d];
        smem[tid] = (float)isum * qk_scale;
    }
    __syncthreads();
    if (tid == 0) {
        float m = -FLT_MAX;
        for (int s = 0; s <= t; ++s) m = fmaxf(m, smem[s]);
        float denom = 0.0f;
        for (int s = 0; s <= t; ++s) { float e = expf(smem[s] - m); smem[s] = e; denom += e; }
        float inv = 1.0f / denom;
        for (int s = 0; s <= t; ++s) smem[s] *= inv;
    }
    __syncthreads();
    if (tid < 32) {
        float acc = 0.0f;
        for (int s = 0; s <= t; ++s) acc += smem[s] * (float)v[((b * T + s) * H + h) * D + tid];
        out[((b * T + t) * H + h) * D + tid] = acc * vs[h];
    }
}

__global__ void causal_attn_f32_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 bid = blockIdx.x;
    int t = bid % T;
    int tmp = bid / T;
    int h = tmp % H;
    int b = tmp / H;
    int tid = threadIdx.x;
    __shared__ float smem[128];
    float qk_scale = qs[h] * ks[h] * rsqrtf((float)D);
    if (tid <= t) {
        int qoff = ((b * T + t) * H + h) * D;
        int koff = ((b * T + tid) * H + h) * D;
        float sum = 0.0f;
        #pragma unroll
        for (int d = 0; d < 32; ++d) sum += q[qoff + d] * k[koff + d];
        smem[tid] = sum * qk_scale;
    }
    __syncthreads();
    if (tid == 0) {
        float m = -FLT_MAX;
        for (int s = 0; s <= t; ++s) m = fmaxf(m, smem[s]);
        float denom = 0.0f;
        for (int s = 0; s <= t; ++s) { float e = expf(smem[s] - m); smem[s] = e; denom += e; }
        float inv = 1.0f / denom;
        for (int s = 0; s <= t; ++s) smem[s] *= inv;
    }
    __syncthreads();
    if (tid < 32) {
        float acc = 0.0f;
        for (int s = 0; s <= t; ++s) acc += smem[s] * v[((b * T + s) * H + h) * D + tid];
        out[((b * T + t) * H + h) * D + tid] = acc * vs[h];
    }
}

torch::Tensor attention_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));
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    int blocks = B*T*H;
    if (q_q.scalar_type() == torch::kFloat32) {
        causal_attn_f32_kernel<<<blocks, 128, 0, stream>>>(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);
    } else {
        causal_attn_i8_kernel<<<blocks, 128, 0, stream>>>((const int8_t*)q_q.data_ptr<int8_t>(), (const int8_t*)k_q.data_ptr<int8_t>(), (const 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);
    }
    return out;
}
"""

_ext = load_inline(name="causal_attn_i8_ext_v2", cpp_sources=_cpp_src, cuda_sources=_cuda_src,
                   functions=["attention_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.attention_forward(q_q, k_q, v_q, q_scale, k_scale, v_scale)