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>
#include <cublas_v2.h>

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 kv_heads = 4;
    int head_dim = w_up_k.size(1) / kv_heads;
    int bs = batch_size * seq_len;
    int kv_dim = kv_heads * head_dim;

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

    cublasHandle_t handle;
    cublasCreate(&handle);
    float alpha = 1.0f, beta = 0.0f;

    // hidden @ w_down -> latent : (bs x model) @ (model x rank)
    cublasSgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N,
                rank, bs, model_dim,
                &alpha,
                w_down.data_ptr<float>(), rank,
                hidden.data_ptr<float>(), model_dim,
                &beta,
                latent.data_ptr<float>(), rank);

    // latent @ w_up_k -> k
    cublasSgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N,
                kv_dim, bs, rank,
                &alpha,
                w_up_k.data_ptr<float>(), kv_dim,
                latent.data_ptr<float>(), rank,
                &beta,
                k.data_ptr<float>(), kv_dim);

    // latent @ w_up_v -> v
    cublasSgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N,
                kv_dim, bs, rank,
                &alpha,
                w_up_v.data_ptr<float>(), kv_dim,
                latent.data_ptr<float>(), rank,
                &beta,
                v.data_ptr<float>(), kv_dim);

    // reshape and stack
    auto k_view = k.view({batch_size, seq_len, kv_heads, head_dim});
    auto v_view = v.view({batch_size, seq_len, kv_heads, head_dim});
    output.select(2, 0).copy_(k_view);
    output.select(2, 1).copy_(v_view);

    cublasDestroy(handle);
    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_cublas',
    cpp_sources=[cpp_source],
    cuda_sources=[cuda_source],
    functions=['mla_forward'],
    verbose=False,
    extra_cuda_cflags=['-O3'],
    extra_ldflags=['-lcublas']
)

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)