import math

import torch
import torch.nn as nn


class Model(nn.Module):
    """
    Causal grouped-query attention with KV-head expansion.

    Query heads outnumber KV heads. Keys and values are expanded by repeating
    each KV head across its assigned query-head group before causal attention.
    """

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

    def forward(
        self,
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
    ) -> torch.Tensor:
        batch_size, seq_len, num_q_heads, head_dim = q.shape
        num_kv_heads = k.shape[2]
        group_size = num_q_heads // num_kv_heads

        expanded_k = k.repeat_interleave(group_size, dim=2)
        expanded_v = v.repeat_interleave(group_size, dim=2)

        out = torch.zeros_like(q)
        scale = 1.0 / math.sqrt(head_dim)

        for b in range(batch_size):
            for h in range(num_q_heads):
                for t in range(seq_len):
                    q_t = q[b, t, h]
                    k_hist = expanded_k[b, : t + 1, h]
                    v_hist = expanded_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 = 80
num_q_heads = 8
num_kv_heads = 2
head_dim = 32


def get_inputs():
    p = popcorn_pri
    q = torch.randn(p.jitter_int(batch_size), p.jitter_int(seq_len), p.jitter_int(num_q_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_kv_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_kv_heads), p.jitter_int(head_dim, align=8), dtype=torch.float32)
    return [q, k, v]


def get_init_inputs():
    return []

import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

cpp_source = r'''
#include <torch/extension.h>

torch::Tensor gqa_forward_cuda(torch::Tensor q, torch::Tensor k, torch::Tensor v);

torch::Tensor gqa_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v) {
    return gqa_forward_cuda(q, k, v);
}
'''

cuda_source = r'''
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <math.h>

#define TMAX 80

template<int WARPS_PER_BLOCK>
__global__ void gqa_kernel(const float* __restrict__ q,
                           const float* __restrict__ k,
                           const float* __restrict__ v,
                           float* __restrict__ out) {
    __shared__ float scores[WARPS_PER_BLOCK][TMAX];
    int tid = threadIdx.x;
    int warp = tid >> 5;
    int lane = tid & 31;
    int row = blockIdx.x * WARPS_PER_BLOCK + warp;
    if (row >= 2 * 8 * 80) return;

    int t = row % 80;
    int tmp = row / 80;
    int h = tmp % 8;
    int b = tmp / 8;
    int kvh = h >> 2;

    const float scale = 0.17677669529663687f;
    int q_base = ((b * 80 + t) * 8 + h) * 32;
    float qd = q[q_base + lane];

    unsigned mask = 0xffffffffu;
    for (int s = 0; s <= t; ++s) {
        int k_base = ((b * 80 + s) * 2 + kvh) * 32;
        float prod = qd * k[k_base + lane];
        prod += __shfl_down_sync(mask, prod, 16);
        prod += __shfl_down_sync(mask, prod, 8);
        prod += __shfl_down_sync(mask, prod, 4);
        prod += __shfl_down_sync(mask, prod, 2);
        prod += __shfl_down_sync(mask, prod, 1);
        if (lane == 0) scores[warp][s] = prod * scale;
    }
    __syncwarp(mask);

    float m = -INFINITY;
    if (lane == 0) {
        for (int s = 0; s <= t; ++s) {
            float sc = scores[warp][s];
            m = fmaxf(m, sc);
        }
    }
    m = __shfl_sync(mask, m, 0);

    float denom = 0.0f;
    if (lane == 0) {
        for (int s = 0; s <= t; ++s) {
            float e = expf(scores[warp][s] - m);
            scores[warp][s] = e;
            denom += e;
        }
    }
    denom = __shfl_sync(mask, denom, 0);
    __syncwarp(mask);

    float acc = 0.0f;
    for (int s = 0; s <= t; ++s) {
        int v_base = ((b * 80 + s) * 2 + kvh) * 32;
        acc += (scores[warp][s] / denom) * v[v_base + lane];
    }
    out[q_base + lane] = acc;
}

torch::Tensor gqa_forward_cuda(torch::Tensor q, torch::Tensor k, torch::Tensor v) {
    auto out = torch::empty_like(q);
    constexpr int WARPS = 4;
    int rows = q.size(0) * q.size(2) * q.size(1);
    int blocks = (rows + WARPS - 1) / WARPS;
    gqa_kernel<WARPS><<<blocks, WARPS * 32>>>(q.data_ptr<float>(), k.data_ptr<float>(), v.data_ptr<float>(), out.data_ptr<float>());
    return out;
}
'''

_ext = load_inline(
    name='gqa_ext_final',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['gqa_forward'],
    extra_cuda_cflags=['-O3', '--use_fast_math'],
    extra_cflags=['-O3'],
    verbose=False,
)

class ModelNew(nn.Module):
    def __init__(self):
        super().__init__()
    def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
        return _ext.gqa_forward(q, k, v)