import math

import torch
import torch.nn as nn


class Model(nn.Module):
    """
    Fused MLA-style attention from compressed KV latents.

    The compressed latent representation is expanded into keys and values on the
    fly inside the attention computation rather than materializing a KV cache
    first. This isolates the main compute pattern behind fused MLA attention.
    """

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

    def forward(
        self,
        q: torch.Tensor,
        kv_latent: torch.Tensor,
        w_up_k: torch.Tensor,
        w_up_v: torch.Tensor,
    ) -> torch.Tensor:
        batch_size, seq_len, num_heads, head_dim = q.shape
        rank = kv_latent.shape[-1]
        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):
                    logits = []
                    values = []
                    for s in range(t + 1):
                        latent = kv_latent[b, s]
                        k = torch.matmul(latent, w_up_k[:, h])
                        v = torch.matmul(latent, w_up_v[:, h])
                        logits.append((q[b, t, h] * k).sum() * scale)
                        values.append(v)
                    attn = torch.softmax(torch.stack(logits, dim=0), dim=0)
                    value_stack = torch.stack(values, dim=0)
                    out[b, t, h] = (attn.unsqueeze(-1) * value_stack).sum(dim=0)

        return out


batch_size = 2
seq_len = 32
num_heads = 4
head_dim = 16
rank = 24


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)
    kv_latent = torch.randn(p.jitter_int(batch_size), p.jitter_int(seq_len), rank, dtype=torch.float32)
    w_up_k = torch.randn(rank, p.jitter_int(num_heads), p.jitter_int(head_dim, align=8), dtype=torch.float32)
    w_up_v = torch.randn(rank, p.jitter_int(num_heads), p.jitter_int(head_dim, align=8), dtype=torch.float32)
    return [q, kv_latent, w_up_k, w_up_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>

// Kernel for fused MLA attention - optimized version
// Each thread handles one (batch, head, query_position) combination
// Caches logits to avoid recomputation
template <typename scalar_t>
__global__ void mla_attention_kernel(
    const scalar_t* __restrict__ q,
    const scalar_t* __restrict__ kv_latent,
    const scalar_t* __restrict__ w_up_k,
    const scalar_t* __restrict__ w_up_v,
    scalar_t* __restrict__ out,
    const int batch_size,
    const int seq_len,
    const int num_heads,
    const int head_dim,
    const int rank,
    const scalar_t scale) {
    
    // Linear index -> (batch, head, query_t)
    const int tid = blockIdx.x * blockDim.x + threadIdx.x;
    const int total_threads = batch_size * num_heads * seq_len;
    
    if (tid >= total_threads) return;
    
    const int t = tid % seq_len;
    const int h = (tid / seq_len) % num_heads;
    const int b = tid / (seq_len * num_heads);
    
    // Pointers for this batch and head
    // q shape: (batch, seq, heads, head_dim)
    const scalar_t* q_ptr = q + ((b * seq_len + t) * num_heads + h) * head_dim;
    // kv_latent shape: (batch, seq, rank)
    const scalar_t* kv_latent_base = kv_latent + b * seq_len * rank;
    scalar_t* out_ptr = out + ((b * seq_len + t) * num_heads + h) * head_dim;
    
    // Cache logits in local memory (seq_len is small)
    scalar_t logits_cache[32];  // max seq_len expected is 32
    
    // Online softmax for numerical stability
    scalar_t max_logit = -INFINITY;
    scalar_t sum_exp = 0.0f;
    
    // First pass: compute logits and find max_logit, sum_exp
    for (int s = 0; s <= t; s++) {
        const scalar_t* latent_s = kv_latent_base + s * rank;
        
        // Compute k = latent @ w_up_k[:, h]
        // k[d] = sum_r latent[r] * w_up_k[r, h, d]
        scalar_t logits_s = 0.0f;
        for (int d = 0; d < head_dim; d++) {
            scalar_t k_d = 0.0f;
            for (int r = 0; r < rank; r++) {
                // w_up_k[r, h, d] -> r * num_heads * head_dim + h * head_dim + d
                k_d += latent_s[r] * w_up_k[r * num_heads * head_dim + h * head_dim + d];
            }
            // Q·K for this position
            logits_s += q_ptr[d] * k_d;
        }
        logits_s *= scale;
        logits_cache[s] = logits_s;
        
        // Online softmax update
        scalar_t new_max = max(max_logit, logits_s);
        scalar_t exp_old = sum_exp * exp(max_logit - new_max);
        scalar_t exp_new = exp(logits_s - new_max);
        sum_exp = exp_old + exp_new;
        max_logit = new_max;
    }
    
    // Second pass: compute weighted sum of values using cached logits
    for (int d = 0; d < head_dim; d++) {
        out_ptr[d] = 0.0f;
    }
    
    for (int s = 0; s <= t; s++) {
        const scalar_t* latent_s = kv_latent_base + s * rank;
        
        // Use cached logit
        scalar_t logits_s = logits_cache[s];
        
        // Attention weight
        scalar_t attn_s = exp(logits_s - max_logit) / sum_exp;
        
        // Compute v = latent @ w_up_v[:, h] and accumulate
        // v[d] = sum_r latent[r] * w_up_v[r, h, d]
        for (int d = 0; d < head_dim; d++) {
            scalar_t v_d = 0.0f;
            for (int r = 0; r < rank; r++) {
                v_d += latent_s[r] * w_up_v[r * num_heads * head_dim + h * head_dim + d];
            }
            out_ptr[d] += attn_s * v_d;
        }
    }
}

torch::Tensor mla_attention_cuda(
    torch::Tensor q,
    torch::Tensor kv_latent,
    torch::Tensor w_up_k,
    torch::Tensor w_up_v) {
    
    const int batch_size = q.size(0);
    const int seq_len = q.size(1);
    const int num_heads = q.size(2);
    const int head_dim = q.size(3);
    const int rank = kv_latent.size(2);
    
    auto out = torch::zeros_like(q);
    const float scale = 1.0f / sqrt((float)head_dim);
    
    const int total_threads = batch_size * num_heads * seq_len;
    const int threads_per_block = 256;
    const int num_blocks = (total_threads + threads_per_block - 1) / threads_per_block;
    
    AT_DISPATCH_FLOATING_TYPES(q.scalar_type(), "mla_attention_kernel", ([&] {
        mla_attention_kernel<scalar_t><<<num_blocks, threads_per_block>>>(
            q.data_ptr<scalar_t>(),
            kv_latent.data_ptr<scalar_t>(),
            w_up_k.data_ptr<scalar_t>(),
            w_up_v.data_ptr<scalar_t>(),
            out.data_ptr<scalar_t>(),
            batch_size, seq_len, num_heads, head_dim, rank,
            static_cast<scalar_t>(scale)
        );
    }));
    
    return out;
}
"""

cpp_source = """
torch::Tensor mla_attention_cuda(torch::Tensor q, torch::Tensor kv_latent, torch::Tensor w_up_k, torch::Tensor w_up_v);
"""

mla_attention_op = load_inline(
    name='mla_attention_op',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['mla_attention_cuda'],
    extra_cuda_cflags=['-O3'],
    verbose=True
)


class ModelNew(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, q: torch.Tensor, kv_latent: torch.Tensor, w_up_k: torch.Tensor, w_up_v: torch.Tensor) -> torch.Tensor:
        return mla_attention_op.mla_attention_cuda(q, kv_latent, w_up_k, w_up_v)