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 <math.h>

#define BATCH 2
#define SEQ 72
#define HEADS 4
#define DIM 32
#define THREADS 128

__global__ __launch_bounds__(THREADS, 4)
void causal_scaled_int8_attn_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) {
    const int tid = threadIdx.x;
    const int bid = blockIdx.x;
    const int t = bid % SEQ;
    const int tmp = bid / SEQ;
    const int h = tmp % HEADS;
    const int b = tmp / HEADS;

    __shared__ float vals[THREADS];  // first 72 entries hold scores/exp; all entries used for reductions

    const int qbase = ((b * SEQ + t) * HEADS + h) * DIM;
    const float qks = qs[h] * ks[h] * 0.17677669529663687f;

    // One thread computes one (or occasionally zero) history dot product.  DIM is fixed at 32.
    for (int s = tid; s <= t; s += THREADS) {
        const int kbase = ((b * SEQ + s) * HEADS + h) * DIM;
        int dot = 0;
        #pragma unroll
        for (int d = 0; d < DIM; ++d) {
            dot += (int)q[qbase + d] * (int)k[kbase + d];
        }
        vals[s] = ((float)dot) * qks;
    }
    __syncthreads();

    // Max reduction over causal scores.
    float m = -3.4028234663852886e38f;
    for (int s = tid; s <= t; s += THREADS) {
        float x = vals[s];
        m = x > m ? x : m;
    }
    vals[tid] = m;
    __syncthreads();
    if (tid < 64) vals[tid] = vals[tid] > vals[tid + 64] ? vals[tid] : vals[tid + 64];
    __syncthreads();
    if (tid < 32) {
        volatile float* vv = vals;
        vv[tid] = vv[tid] > vv[tid + 32] ? vv[tid] : vv[tid + 32];
        vv[tid] = vv[tid] > vv[tid + 16] ? vv[tid] : vv[tid + 16];
        vv[tid] = vv[tid] > vv[tid + 8]  ? vv[tid] : vv[tid + 8];
        vv[tid] = vv[tid] > vv[tid + 4]  ? vv[tid] : vv[tid + 4];
        vv[tid] = vv[tid] > vv[tid + 2]  ? vv[tid] : vv[tid + 2];
        vv[tid] = vv[tid] > vv[tid + 1]  ? vv[tid] : vv[tid + 1];
    }
    __syncthreads();
    const float maxv = vals[0];

    // Exponentiate, keep unnormalised weights in vals[0..t], and reduce sum.
    float sum = 0.0f;
    for (int s = tid; s <= t; s += THREADS) {
        float e = expf(vals[s] - maxv);
        vals[s] = e;
        sum += e;
    }
    __syncthreads();
    vals[tid] = sum;
    __syncthreads();
    if (tid < 64) vals[tid] += vals[tid + 64];
    __syncthreads();
    if (tid < 32) {
        volatile float* vv = vals;
        vv[tid] += vv[tid + 32];
        vv[tid] += vv[tid + 16];
        vv[tid] += vv[tid + 8];
        vv[tid] += vv[tid + 4];
        vv[tid] += vv[tid + 2];
        vv[tid] += vv[tid + 1];
    }
    __syncthreads();
    const float norm = vs[h] / vals[0];

    // Compute the 32 output channels.  Only first warp is needed here.
    if (tid < DIM) {
        float acc = 0.0f;
        #pragma unroll 1
        for (int s = 0; s <= t; ++s) {
            const int vbase = ((b * SEQ + s) * HEADS + h) * DIM;
            acc += vals[s] * (float)v[vbase + tid];
        }
        out[qbase + tid] = acc * norm;
    }
}

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) {
    auto out = torch::empty({BATCH, SEQ, HEADS, DIM}, q_scale.options());
    const int blocks = BATCH * HEADS * SEQ;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    causal_scaled_int8_attn_kernel<<<blocks, THREADS, 0, stream>>>(
        reinterpret_cast<const int8_t*>(q_q.data_ptr<int8_t>()),
        reinterpret_cast<const int8_t*>(k_q.data_ptr<int8_t>()),
        reinterpret_cast<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>());
    return out;
}
"""

_attn_ext = load_inline(
    name="modelnew_int8_causal_attn_ext",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["attention_forward"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    verbose=False,
)

class ModelNew(nn.Module):
    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:
        return _attn_ext.attention_forward(q_q, k_q, v_q, q_scale, k_scale, v_scale)