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

cuda_source = """
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cmath>

// Optimized: prefetch latents into shared memory, better coalescing
__global__ void __launch_bounds__(256) fused_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 idx = blockIdx.x;
    int t = idx % seq_len;
    idx /= seq_len;
    int h = idx % num_heads;
    idx /= num_heads;
    int b = idx;
    
    int tid = threadIdx.x;
    int nthreads = blockDim.x;
    
    // Shared memory layout:
    // latents: [seq_len, rank]
    // query: [head_dim]
    // keys: [seq_len, head_dim]
    // vals: [seq_len, head_dim]  
    // scores: [seq_len]
    extern __shared__ char smem[];
    float* latents_smem = (float*)smem;
    float* query_smem = latents_smem + seq_len * rank;
    float* keys = query_smem + head_dim;
    float* vals = keys + seq_len * head_dim;
    float* scores = vals + seq_len * head_dim;
    
    // Cooperatively load latents for all positions into shared memory
    // latents are shape [batch, seq_len, rank]
    for (int i = tid; i < seq_len * rank; i += nthreads) {
        int s = i / rank;
        int r = i % rank;
        latents_smem[s * rank + r] = kv_latent[(b * seq_len + s) * rank + r];
    }
    
    // Load query into shared memory
    const float* q_ptr = q + ((b * seq_len + t) * num_heads + h) * head_dim;
    for (int d = tid; d < head_dim; d += nthreads) {
        query_smem[d] = q_ptr[d];
    }
    
    __syncthreads();
    
    // Compute keys and values for all positions s = 0..t
    // Parallelize over output elements (s, d)
    int total_elems = (t + 1) * head_dim;
    for (int elem = tid; elem < total_elems; elem += nthreads) {
        int s = elem / head_dim;
        int d = elem % head_dim;
        
        float k = 0.0f;
        float v = 0.0f;
        // Compute k[d] = latent[s] @ w_up_k[:, h, d]
        for (int r = 0; r < rank; r++) {
            float lat = latents_smem[s * rank + r];
            k += lat * w_up_k[r * num_heads * head_dim + h * head_dim + d];
            v += lat * w_up_v[r * num_heads * head_dim + h * head_dim + d];
        }
        keys[s * head_dim + d] = k;
        vals[s * head_dim + d] = v;
    }
    
    __syncthreads();
    
    // Compute attention scores: score[s] = (q . k[s]) * scale
    for (int s = tid; s <= t; s += nthreads) {
        float score = 0.0f;
        for (int d = 0; d < head_dim; d++) {
            score += query_smem[d] * keys[s * head_dim + d];
        }
        scores[s] = score * scale;
    }
    
    __syncthreads();
    
    // Find max score for numerical stability
    __shared__ float max_score;
    __shared__ float sum_exp;
    
    float local_max = -INFINITY;
    for (int s = tid; s <= t; s += nthreads) {
        local_max = fmaxf(local_max, scores[s]);
    }
    
    // Block-level max reduction using warp shuffles
    __shared__ float shared_max[8];
    int lane = tid % 32;
    int warp_id = tid / 32;
    int num_warps = nthreads / 32;
    
    for (int offset = 16; offset > 0; offset /= 2) {
        local_max = fmaxf(local_max, __shfl_down_sync(0xffffffff, local_max, offset));
    }
    if (lane == 0 && warp_id < 8) {
        shared_max[warp_id] = local_max;
    }
    __syncthreads();
    
    if (warp_id == 0) {
        float final_max = (lane < num_warps && lane < 8) ? shared_max[lane] : -INFINITY;
        for (int offset = 16; offset > 0; offset /= 2) {
            final_max = fmaxf(final_max, __shfl_down_sync(0xffffffff, final_max, offset));
        }
        if (lane == 0) {
            max_score = final_max;
        }
    }
    __syncthreads();
    
    // Compute exp(score - max) and sum
    float local_sum = 0.0f;
    for (int s = tid; s <= t; s += nthreads) {
        float e = expf(scores[s] - max_score);
        scores[s] = e;
        local_sum += e;
    }
    
    // Block-level sum reduction
    __shared__ float shared_sum[8];
    for (int offset = 16; offset > 0; offset /= 2) {
        local_sum += __shfl_down_sync(0xffffffff, local_sum, offset);
    }
    if (lane == 0 && warp_id < 8) {
        shared_sum[warp_id] = local_sum;
    }
    __syncthreads();
    
    if (warp_id == 0) {
        float final_sum = (lane < num_warps && lane < 8) ? shared_sum[lane] : 0.0f;
        for (int offset = 16; offset > 0; offset /= 2) {
            final_sum += __shfl_down_sync(0xffffffff, final_sum, offset);
        }
        if (lane == 0) {
            sum_exp = final_sum;
        }
    }
    __syncthreads();
    
    // Compute weighted sum of values
    float inv_sum = 1.0f / sum_exp;
    for (int d = tid; d < head_dim; d += nthreads) {
        float out_val = 0.0f;
        for (int s = 0; s <= t; s++) {
            out_val += scores[s] * inv_sum * vals[s * head_dim + d];
        }
        out[((b * seq_len + t) * num_heads + h) * head_dim + d] = out_val;
    }
}

torch::Tensor fused_mla_attention_cuda(
    torch::Tensor q,
    torch::Tensor kv_latent,
    torch::Tensor w_up_k,
    torch::Tensor w_up_v
) {
    TORCH_CHECK(q.is_cuda(), "q must be a CUDA tensor");
    TORCH_CHECK(kv_latent.is_cuda(), "kv_latent must be a CUDA tensor");
    TORCH_CHECK(w_up_k.is_cuda(), "w_up_k must be a CUDA tensor");
    TORCH_CHECK(w_up_v.is_cuda(), "w_up_v must be a CUDA tensor");
    
    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);
    
    int num_blocks = batch_size * num_heads * seq_len;
    int threads_per_block = 256;
    
    // Shared memory: latents + query + keys + vals + scores
    size_t smem_size = (seq_len * rank + head_dim + 2 * seq_len * head_dim + seq_len) * sizeof(float);
    
    fused_mla_attention_kernel<<<num_blocks, threads_per_block, smem_size>>>(
        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 fused_mla_attention_cuda(
    torch::Tensor q,
    torch::Tensor kv_latent,
    torch::Tensor w_up_k,
    torch::Tensor w_up_v
);
"""

module = load_inline(
    name="fused_mla_attention",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["fused_mla_attention_cuda"],
    verbose=True,
    extra_cuda_cflags=["-O3"],
)

class ModelNew(nn.Module):
    def __init__(self):
        super().__init__()
        self.fused_mla_attention = module.fused_mla_attention_cuda
    
    def forward(self, q, kv_latent, w_up_k, w_up_v):
        return self.fused_mla_attention(q, kv_latent, w_up_k, w_up_v)