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

cpp_src = r'''
#include <torch/extension.h>
torch::Tensor mla_forward(torch::Tensor q, torch::Tensor kv, torch::Tensor wk, torch::Tensor wv);
'''

cuda_src = r'''
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cmath>

__global__ void mla_kernel(const float* __restrict__ q,
                           const float* __restrict__ kv,
                           const float* __restrict__ wk,
                           const float* __restrict__ wv,
                           float* __restrict__ out) {
    int idx = blockIdx.x;
    int t = idx & 31;        // seq_len=32
    int tmp = idx >> 5;
    int h = tmp & 3;         // num_heads=4
    int b = tmp >> 2;        // batch
    int tid = threadIdx.x;
    const unsigned mask = 0xffffffffu;

    __shared__ float coeff[24];
    __shared__ float probs[32];
    __shared__ float lat[24];

    int qbase = (((b * 32 + t) * 4 + h) * 16);

    if (tid < 24) {
        float acc = 0.0f;
        int wkbase = ((tid * 4 + h) * 16);
        #pragma unroll
        for (int d = 0; d < 16; ++d) {
            acc += q[qbase + d] * wk[wkbase + d];
        }
        coeff[tid] = acc;
    }
    __syncwarp(mask);

    float logit = -3.402823466e38f;
    if (tid <= t) {
        float acc = 0.0f;
        int kvbase = (b * 32 + tid) * 24;
        #pragma unroll
        for (int r = 0; r < 24; ++r) {
            acc += kv[kvbase + r] * coeff[r];
        }
        logit = acc * 0.25f;
    }

    float m = logit;
    #pragma unroll
    for (int off = 16; off > 0; off >>= 1) {
        m = fmaxf(m, __shfl_down_sync(mask, m, off));
    }
    m = __shfl_sync(mask, m, 0);

    float e = (tid <= t) ? expf(logit - m) : 0.0f;
    float sum = e;
    #pragma unroll
    for (int off = 16; off > 0; off >>= 1) {
        sum += __shfl_down_sync(mask, sum, off);
    }
    sum = __shfl_sync(mask, sum, 0);
    if (tid <= t) probs[tid] = e / sum;
    __syncwarp(mask);

    if (tid < 24) {
        float acc = 0.0f;
        #pragma unroll
        for (int s = 0; s < 32; ++s) {
            if (s <= t) acc += probs[s] * kv[(b * 32 + s) * 24 + tid];
        }
        lat[tid] = acc;
    }
    __syncwarp(mask);

    if (tid < 16) {
        float acc = 0.0f;
        #pragma unroll
        for (int r = 0; r < 24; ++r) {
            acc += lat[r] * wv[((r * 4 + h) * 16 + tid)];
        }
        out[qbase + tid] = acc;
    }
}

torch::Tensor mla_forward(torch::Tensor q, torch::Tensor kv, torch::Tensor wk, torch::Tensor wv) {
    auto out = torch::empty_like(q);
    mla_kernel<<<256, 32>>>(q.data_ptr<float>(), kv.data_ptr<float>(), wk.data_ptr<float>(), wv.data_ptr<float>(), out.data_ptr<float>());
    return out;
}
'''

_ext = load_inline(name='mla_fused_ext_v5_final', cpp_sources=cpp_src, cuda_sources=cuda_src, functions=['mla_forward'], verbose=False, extra_cuda_cflags=['-O3', '--use_fast_math'])

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 _ext.mla_forward(q.contiguous(), kv_latent.contiguous(), w_up_k.contiguous(), w_up_v.contiguous())