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 math
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

cuda_source = """
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cmath>

// Optimized: One thread per (batch, q_head, seq_pos)
// Does KV expansion internally and coalesces memory access
__global__ void causal_gqa_kernel(
    const float* __restrict__ q,
    const float* __restrict__ k,
    const float* __restrict__ v,
    float* __restrict__ out,
    int batch_size, int seq_len, int num_q_heads, int num_kv_heads, int head_dim
) {
    int group_size = num_q_heads / num_kv_heads;
    float scale = 1.0f / sqrtf((float)head_dim);
    
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total_threads = batch_size * num_q_heads * seq_len;
    
    if (idx >= total_threads) return;
    
    int t = idx % seq_len;
    int h = (idx / seq_len) % num_q_heads;
    int b = idx / (num_q_heads * seq_len);
    
    // Map query head to KV head
    int kv_h = h / group_size;
    
    // Local array for attention scores
    float scores[80];
    
    // Load query vector into registers
    float q_vec[32];
    for (int d = 0; d < head_dim; d++) {
        q_vec[d] = q[((b * seq_len + t) * num_q_heads + h) * head_dim + d];
    }
    
    // Compute attention scores for positions 0..t
    // Access K from KV head kv_h directly (no expansion needed)
    for (int s = 0; s <= t; s++) {
        float dot = 0.0f;
        for (int d = 0; d < head_dim; d++) {
            dot += q_vec[d] * k[((b * seq_len + s) * num_kv_heads + kv_h) * head_dim + d];
        }
        scores[s] = dot * scale;
    }
    
    // Softmax - find max
    float max_score = scores[0];
    for (int s = 1; s <= t; s++) {
        if (scores[s] > max_score) max_score = scores[s];
    }
    
    // Softmax - compute exp and sum
    float sum_exp = 0.0f;
    for (int s = 0; s <= t; s++) {
        scores[s] = expf(scores[s] - max_score);
        sum_exp += scores[s];
    }
    
    // Softmax - normalize
    float inv_sum = 1.0f / sum_exp;
    for (int s = 0; s <= t; s++) {
        scores[s] *= inv_sum;
    }
    
    // Weighted sum of values
    // Access V from KV head kv_h directly (no expansion needed)
    for (int d = 0; d < head_dim; d++) {
        float val = 0.0f;
        for (int s = 0; s <= t; s++) {
            val += scores[s] * v[((b * seq_len + s) * num_kv_heads + kv_h) * head_dim + d];
        }
        out[((b * seq_len + t) * num_q_heads + h) * head_dim + d] = val;
    }
}

torch::Tensor forward_cuda(torch::Tensor q, torch::Tensor k, torch::Tensor v) {
    int batch_size = q.size(0);
    int seq_len = q.size(1);
    int num_q_heads = q.size(2);
    int head_dim = q.size(3);
    int num_kv_heads = k.size(2);
    
    auto out = torch::empty_like(q);
    
    int total_threads = batch_size * num_q_heads * seq_len;
    int threads_per_block = 256;
    int num_blocks = (total_threads + threads_per_block - 1) / threads_per_block;
    
    causal_gqa_kernel<<<num_blocks, threads_per_block>>>(
        q.data_ptr<float>(),
        k.data_ptr<float>(),
        v.data_ptr<float>(),
        out.data_ptr<float>(),
        batch_size, seq_len, num_q_heads, num_kv_heads, head_dim
    );
    
    return out;
}
"""

cpp_source = """
torch::Tensor forward_cuda(torch::Tensor q, torch::Tensor k, torch::Tensor v);
"""

module = load_inline(
    name="causal_gqa_ext",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["forward_cuda"],
    verbose=True,
)

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