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

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

__global__ void mla_attention_kernel(
    const float* __restrict__ q,
    const float* __restrict__ kv_latent,
    const float* __restrict__ w_up_k,
    const float* __restrict__ w_up_v,
    float* __restrict__ out,
    int batch_size,
    int seq_len,
    int num_heads,
    int head_dim,
    int rank,
    float scale
) {
    int b = blockIdx.x;
    int h = blockIdx.y;
    int t = threadIdx.x;

    if (t >= seq_len) return;

    extern __shared__ float sdata[];

    // Compute for this (b,h,t)
    float q_local[16]; // head_dim <=16
    for (int d=0; d<head_dim; d++) {
        q_local[d] = q[((b * seq_len + t) * num_heads + h) * head_dim + d];
    }

    float logits[32]; // seq_len<=32
    float values[32*16];

    for (int s=0; s<=t; s++) {
        float k[16];
        float v[16];
        for (int d=0; d<head_dim; d++) {
            k[d] = 0.0f;
            v[d] = 0.0f;
        }
        for (int r=0; r<rank; r++) {
            float lat = kv_latent[(b * seq_len + s) * rank + r];
            for (int d=0; d<head_dim; d++) {
                k[d] += lat * w_up_k[(r * num_heads + h) * head_dim + d];
                v[d] += lat * w_up_v[(r * num_heads + h) * head_dim + d];
            }
        }
        float dot = 0.0f;
        for (int d=0; d<head_dim; d++) {
            dot += q_local[d] * k[d];
            values[s * head_dim + d] = v[d];
        }
        logits[s] = dot * scale;
    }

    // softmax
    float max_logit = -1e30f;
    for (int s=0; s<=t; s++) {
        if (logits[s] > max_logit) max_logit = logits[s];
    }
    float sum_exp = 0.0f;
    for (int s=0; s<=t; s++) {
        logits[s] = expf(logits[s] - max_logit);
        sum_exp += logits[s];
    }
    float inv_sum = 1.0f / sum_exp;
    for (int s=0; s<=t; s++) {
        logits[s] *= inv_sum;
    }

    // weighted sum
    for (int d=0; d<head_dim; d++) {
        float res = 0.0f;
        for (int s=0; s<=t; s++) {
            res += logits[s] * values[s * head_dim + d];
        }
        out[((b * seq_len + t) * num_heads + h) * head_dim + d] = res;
    }
}

torch::Tensor mla_attention_cuda(
    torch::Tensor q,
    torch::Tensor kv_latent,
    torch::Tensor w_up_k,
    torch::Tensor w_up_v
) {
    int batch_size = q.size(0);
    int seq_len = q.size(1);
    int num_heads = q.size(2);
    int head_dim = q.size(3);
    int rank = kv_latent.size(2);
    float scale = 1.0f / sqrtf((float)head_dim);

    auto out = torch::zeros_like(q);

    dim3 grid(batch_size, num_heads);
    int threads = seq_len;
    size_t shmem = 0; // not using much

    mla_attention_kernel<<<grid, threads, shmem>>>(
        q.data_ptr<float>(),
        kv_latent.data_ptr<float>(),
        w_up_k.data_ptr<float>(),
        w_up_v.data_ptr<float>(),
        out.data_ptr<float>(),
        batch_size, seq_len, num_heads, head_dim, rank, 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);
'''

module = load_inline(
    name='mla_attention',
    cpp_sources=[cpp_source],
    cuda_sources=[cuda_source],
    functions=['mla_attention_cuda'],
    verbose=False
)

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

    def forward(self, q, kv_latent, w_up_k, w_up_v):
        return self.module.mla_attention_cuda(q, kv_latent, w_up_k, w_up_v)