import torch
import torch.nn as nn


class Model(nn.Module):
    """
    Low-rank expansion of compressed KV activations into multi-head K/V tensors.

    This isolates the MLA-style compressed-cache pattern: a hidden activation is
    projected into a low-rank latent space and then expanded back into separate
    key and value tensors.
    """

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

    def forward(
        self,
        hidden: torch.Tensor,
        w_down: torch.Tensor,
        w_up_k: torch.Tensor,
        w_up_v: torch.Tensor,
    ) -> torch.Tensor:
        batch_size, seq_len, model_dim = hidden.shape
        rank = w_down.shape[1]
        kv_heads = 4
        head_dim = w_up_k.shape[1] // kv_heads

        latent = hidden.reshape(batch_size * seq_len, model_dim) @ w_down
        k = latent @ w_up_k
        v = latent @ w_up_v
        k = k.view(batch_size, seq_len, kv_heads, head_dim)
        v = v.view(batch_size, seq_len, kv_heads, head_dim)
        return torch.stack((k, v), dim=2)


batch_size = 4
seq_len = 96
model_dim = 256
rank = 48
kv_heads = 4
head_dim = 32


def get_inputs():
    p = popcorn_pri
    hidden = torch.randn(p.jitter_int(batch_size), p.jitter_int(seq_len), p.jitter_int(model_dim, align=8), dtype=torch.float32)
    w_down = torch.randn(p.jitter_int(model_dim, align=8), rank, dtype=torch.float32)
    w_up_k = torch.randn(rank, p.jitter_int(kv_heads) * p.jitter_int(head_dim, align=8), dtype=torch.float32)
    w_up_v = torch.randn(rank, p.jitter_int(kv_heads) * p.jitter_int(head_dim, align=8), dtype=torch.float32)
    return [hidden, w_down, 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>

__global__ void mla_kernel(
    const float* __restrict__ hidden,
    const float* __restrict__ w_down,
    const float* __restrict__ w_up_k,
    const float* __restrict__ w_up_v,
    float* __restrict__ output,
    int batch_size, int seq_len, int model_dim, int rank, int kv_heads, int head_dim
) {
    int total_tokens = batch_size * seq_len;
    int out_dim = kv_heads * head_dim;
    int b = blockIdx.x;
    int s = blockIdx.y;
    int token = b * seq_len + s;
    if (token >= total_tokens) return;

    extern __shared__ float smem[];
    float* latent = smem;

    for (int r = threadIdx.x; r < rank; r += blockDim.x) {
        double sum = 0.0;
        for (int d = 0; d < model_dim; d++) {
            sum += (double)hidden[token * model_dim + d] * (double)w_down[d * rank + r];
        }
        latent[r] = (float)sum;
    }
    __syncthreads();

    for (int idx = threadIdx.x; idx < out_dim; idx += blockDim.x) {
        double sum_k = 0.0;
        double sum_v = 0.0;
        for (int r = 0; r < rank; r++) {
            double l = (double)latent[r];
            sum_k += l * (double)w_up_k[r * out_dim + idx];
            sum_v += l * (double)w_up_v[r * out_dim + idx];
        }
        int h = idx / head_dim;
        int dd = idx % head_dim;
        size_t base = ((size_t)b * seq_len + s) * 2 * kv_heads * head_dim;
        output[base + (0 * kv_heads * head_dim) + h * head_dim + dd] = (float)sum_k;
        output[base + (1 * kv_heads * head_dim) + h * head_dim + dd] = (float)sum_v;
    }
}

torch::Tensor mla_forward(
    torch::Tensor hidden,
    torch::Tensor w_down,
    torch::Tensor w_up_k,
    torch::Tensor w_up_v
) {
    int batch_size = hidden.size(0);
    int seq_len = hidden.size(1);
    int model_dim = hidden.size(2);
    int rank = w_down.size(1);
    int out_dim = w_up_k.size(1);
    int kv_heads = 4;
    int head_dim = out_dim / kv_heads;

    auto options = torch::TensorOptions().dtype(torch::kFloat32).device(hidden.device());
    auto output = torch::empty({batch_size, seq_len, 2, kv_heads, head_dim}, options);

    dim3 grid(batch_size, seq_len);
    int threads = 256;
    int smem_size = rank * sizeof(float);
    mla_kernel<<<grid, threads, smem_size>>>(
        hidden.data_ptr<float>(),
        w_down.data_ptr<float>(),
        w_up_k.data_ptr<float>(),
        w_up_v.data_ptr<float>(),
        output.data_ptr<float>(),
        batch_size, seq_len, model_dim, rank, kv_heads, head_dim
    );
    return output;
}
"""

cpp_source = """
torch::Tensor mla_forward(torch::Tensor hidden, torch::Tensor w_down, torch::Tensor w_up_k, torch::Tensor w_up_v);
"""

module = load_inline(
    name="mla_cuda",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["mla_forward"],
    verbose=False,
    extra_cuda_cflags=["-O3"]
)

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

    def forward(self, hidden, w_down, w_up_k, w_up_v):
        return self.mla.mla_forward(hidden, w_down, w_up_k, w_up_v)