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_src = r'''
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>

#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x " must be a CUDA tensor")
#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
#define CHECK_FLOAT(x) TORCH_CHECK(x.scalar_type() == at::kFloat, #x " must be float32")

// C = A[M,256] @ B[256,48], M=384
// 16x16 output tile, K tile 32
__global__ __launch_bounds__(256)
void down_kernel(const float* __restrict__ A,
                 const float* __restrict__ B,
                 float* __restrict__ C) {
    constexpr int M = 384;
    constexpr int K = 256;
    constexpr int N = 48;
    constexpr int BM = 16;
    constexpr int BN = 16;
    constexpr int BK = 32;
    __shared__ float As[BM][BK + 1];
    __shared__ float Bs[BK][BN + 1];

    int tx = threadIdx.x; // 0..15 col
    int ty = threadIdx.y; // 0..15 row
    int row = blockIdx.y * BM + ty;
    int col = blockIdx.x * BN + tx;
    float acc = 0.0f;

    #pragma unroll
    for (int k0 = 0; k0 < K; k0 += BK) {
        // 256 threads load 512 A elems and 512 B elems in two passes each
        int tid = ty * 16 + tx;
        int aidx0 = tid;
        int aidx1 = tid + 256;
        int ar0 = aidx0 / BK, ak0 = aidx0 - ar0 * BK;
        int ar1 = aidx1 / BK, ak1 = aidx1 - ar1 * BK;
        As[ar0][ak0] = A[(blockIdx.y * BM + ar0) * K + (k0 + ak0)];
        As[ar1][ak1] = A[(blockIdx.y * BM + ar1) * K + (k0 + ak1)];

        int bidx0 = tid;
        int bidx1 = tid + 256;
        int br0 = bidx0 / BN, bc0 = bidx0 - br0 * BN;
        int br1 = bidx1 / BN, bc1 = bidx1 - br1 * BN;
        int gc0 = blockIdx.x * BN + bc0;
        int gc1 = blockIdx.x * BN + bc1;
        Bs[br0][bc0] = (gc0 < N) ? B[(k0 + br0) * N + gc0] : 0.0f;
        Bs[br1][bc1] = (gc1 < N) ? B[(k0 + br1) * N + gc1] : 0.0f;
        __syncthreads();

        #pragma unroll
        for (int kk = 0; kk < BK; ++kk) {
            acc = fmaf(As[ty][kk], Bs[kk][tx], acc);
        }
        __syncthreads();
    }
    if (col < N) C[row * N + col] = acc;
}

// output = stack(latent@w_up_k, latent@w_up_v) as [4,96,2,4,32]
// latent A[M,48], Bk/Bv[48,128]. 16x16 tile, compute K and V together.
__global__ __launch_bounds__(256)
void up_kernel(const float* __restrict__ A,
               const float* __restrict__ Bk,
               const float* __restrict__ Bv,
               float* __restrict__ O) {
    constexpr int M = 384;
    constexpr int K = 48;
    constexpr int N = 128;
    constexpr int BM = 16;
    constexpr int BN = 16;
    constexpr int BK = 16;
    __shared__ float As[BM][BK + 1];
    __shared__ float Bks[BK][BN + 1];
    __shared__ float Bvs[BK][BN + 1];

    int tx = threadIdx.x;
    int ty = threadIdx.y;
    int row = blockIdx.y * BM + ty;
    int col = blockIdx.x * BN + tx;
    float acck = 0.0f;
    float accv = 0.0f;
    int tid = ty * 16 + tx;

    #pragma unroll
    for (int k0 = 0; k0 < K; k0 += BK) {
        // Load A 16x16
        As[ty][tx] = A[row * K + (k0 + tx)];
        // Load Bk/Bv 16x16
        Bks[ty][tx] = Bk[(k0 + ty) * N + col];
        Bvs[ty][tx] = Bv[(k0 + ty) * N + col];
        __syncthreads();
        #pragma unroll
        for (int kk = 0; kk < BK; ++kk) {
            float a = As[ty][kk];
            acck = fmaf(a, Bks[kk][tx], acck);
            accv = fmaf(a, Bvs[kk][tx], accv);
        }
        __syncthreads();
    }
    // contiguous output [M,2,N]
    O[(row * 2) * N + col] = acck;
    O[(row * 2 + 1) * N + col] = accv;
}

torch::Tensor mla_forward(torch::Tensor hidden, torch::Tensor w_down, torch::Tensor w_up_k, torch::Tensor w_up_v) {
    CHECK_CUDA(hidden); CHECK_CUDA(w_down); CHECK_CUDA(w_up_k); CHECK_CUDA(w_up_v);
    CHECK_CONTIGUOUS(hidden); CHECK_CONTIGUOUS(w_down); CHECK_CONTIGUOUS(w_up_k); CHECK_CONTIGUOUS(w_up_v);
    CHECK_FLOAT(hidden); CHECK_FLOAT(w_down); CHECK_FLOAT(w_up_k); CHECK_FLOAT(w_up_v);

    auto opts = hidden.options();
    auto latent = torch::empty({384, 48}, opts);
    auto out = torch::empty({4, 96, 2, 4, 32}, opts);

    dim3 t(16, 16);
    dim3 g1(3, 24);
    down_kernel<<<g1, t, 0, at::cuda::getCurrentCUDAStream()>>>(
        hidden.data_ptr<float>(), w_down.data_ptr<float>(), latent.data_ptr<float>());
    dim3 g2(8, 24);
    up_kernel<<<g2, t, 0, at::cuda::getCurrentCUDAStream()>>>(
        latent.data_ptr<float>(), w_up_k.data_ptr<float>(), w_up_v.data_ptr<float>(), out.data_ptr<float>());
    return out;
}
'''

_ext = load_inline(
    name='mla_lowrank_ext_v1',
    cpp_sources='',
    cuda_sources=_cuda_src,
    functions=['mla_forward'],
    extra_cuda_cflags=['-O3', '--use_fast_math'],
    extra_cflags=['-O3'],
    verbose=False,
)

class ModelNew(nn.Module):
    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:
        return _ext.mla_forward(hidden, w_down, w_up_k, w_up_v)