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

# Define the CUDA kernel source
cuda_source = """
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cmath>

#define BLOCK_SIZE 256

__device__ float warp_reduce_sum(float val, int warp_size) {
    for (int offset = warp_size / 2; offset > 0; offset /= 2) {
        val += __shfl_down_sync(0xffffffff, val, offset);
    }
    return val;
}

__global__ void grouped_query_attention_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
) {
    // Each thread block handles one (batch, head) pair
    int batch_idx = blockIdx.x / num_q_heads;
    int q_head_idx = blockIdx.x % num_q_heads;
    
    // Determine which KV head this query head uses
    int kv_head_idx = q_head_idx / group_size;
    
    // Thread index within the block
    int tid = threadIdx.x;
    
    // Pointers to the relevant slices
    const float* q_ptr = q + (batch_idx * seq_len * num_q_heads + q_head_idx) * head_dim;
    const float* k_ptr = k + (batch_idx * seq_len * num_kv_heads + kv_head_idx) * head_dim;
    const float* v_ptr = v + (batch_idx * seq_len * num_kv_heads + kv_head_idx) * head_dim;
    float* out_ptr = out + (batch_idx * seq_len * num_q_heads + q_head_idx) * head_dim;
    
    // Scale factor
    float scale = 1.0f / sqrtf((float)head_dim);
    
    // Shared memory for scores and attention weights (one per position in sequence)
    extern __shared__ float shared_mem[];
    float* scores = shared_mem;           // seq_len elements
    float* attn_weights = shared_mem + seq_len;  // seq_len elements
    
    // Iterate over each query position
    for (int t = 0; t < seq_len; t++) {
        // Get query vector at position t
        const float* q_t = q_ptr + t * num_q_heads * head_dim;
        
        // Compute scores for all positions 0..t using all threads in block
        float max_score = -INFINITY;
        
        // Each thread computes scores for some positions
        for (int pos = tid; pos <= t; pos += BLOCK_SIZE) {
            const float* k_pos = k_ptr + pos * num_kv_heads * head_dim;
            
            float score = 0.0f;
            for (int d = 0; d < head_dim; d++) {
                score += q_t[d] * k_pos[d];
            }
            score *= scale;
            scores[pos] = score;
        }
        
        // Also set scores beyond t to -inf for safety
        for (int pos = tid + t + 1; pos < seq_len; pos += BLOCK_SIZE) {
            scores[pos] = -INFINITY;
        }
        
        __syncthreads();
        
        // Find max score (for numerical stability)
        if (tid == 0) {
            max_score = scores[0];
            for (int i = 1; i <= t; i++) {
                if (scores[i] > max_score) max_score = scores[i];
            }
        }
        
        // Broadcast max_score to all threads
        __shared__ float shared_max;
        if (tid == 0) shared_max = max_score;
        __syncthreads();
        max_score = shared_max;
        
        // Compute exp(scores - max) and sum
        float sum_exp = 0.0f;
        for (int pos = tid; pos <= t; pos += BLOCK_SIZE) {
            float exp_val = expf(scores[pos] - max_score);
            attn_weights[pos] = exp_val;
            sum_exp += exp_val;
        }
        
        // Reduce sum_exp across threads
        sum_exp = warp_reduce_sum(sum_exp, min(BLOCK_SIZE, 32));
        
        __shared__ float total_sum[BLOCK_SIZE / 32];
        int warp_id = tid / 32;
        int lane_id = tid % 32;
        if (lane_id == 0) {
            total_sum[warp_id] = sum_exp;
        }
        __syncthreads();
        
        if (tid == 0) {
            float final_sum = 0.0f;
            for (int i = 0; i < (BLOCK_SIZE + 31) / 32; i++) {
                final_sum += total_sum[i];
            }
            shared_max = final_sum;  // Reuse shared_max for final_sum
        }
        __syncthreads();
        float inv_sum = 1.0f / shared_max;
        
        // Normalize attention weights
        for (int pos = tid; pos <= t; pos += BLOCK_SIZE) {
            attn_weights[pos] *= inv_sum;
        }
        __syncthreads();
        
        // Compute output: weighted sum of V
        // Each thread computes part of the output dimension
        for (int d = tid; d < head_dim; d += BLOCK_SIZE) {
            float out_val = 0.0f;
            for (int pos = 0; pos <= t; pos++) {
                const float* v_pos = v_ptr + pos * num_kv_heads * head_dim;
                out_val += attn_weights[pos] * v_pos[d];
            }
            out_ptr[t * num_q_heads * head_dim + d] = out_val;
        }
        __syncthreads();
    }
}

torch::Tensor grouped_query_attention_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::zeros_like(q);
    
    int group_size = num_q_heads / num_kv_heads;
    
    int num_blocks = batch_size * num_q_heads;
    size_t shared_mem_size = seq_len * 2 * sizeof(float);
    
    grouped_query_attention_kernel<<<num_blocks, BLOCK_SIZE, shared_mem_size>>>(
        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, group_size
    );
    
    return out;
}
"""

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

# Load the CUDA extension
grouped_query_attention = load_inline(
    name='grouped_query_attention',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['grouped_query_attention_cuda'],
    extra_cuda_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 grouped_query_attention.grouped_query_attention_cuda(q, k, v)