import math
from dataclasses import dataclass
import torch
from torch import nn  # KVCache inherits from nn.Module
import torch.nn.functional as F
from task import input_t, output_t

# Global cache for RoPE matrices
_rope_cache = {}


def get_rope_cache(d_model: int, max_seq_len: int, device: torch.device, dtype: torch.dtype):
    """Get or create cached RoPE rotation matrices, ensuring they are stored in the specified dtype."""
    # Add dtype to the cache key to ensure dtype correctness
    cache_key = (d_model, max_seq_len, device, dtype)
    if cache_key not in _rope_cache:
        # Create theta and positions. For RoPE, calculations often use float32 for precision
        # then convert to the target dtype. Here, reference used bf16 for theta.
        # Problem states Q,K,V are bfloat16.
        theta_values = 10000.0 ** (
            -torch.arange(0, d_model // 2, dtype=torch.bfloat16, device=device) / (d_model // 2)
        )
        positions = torch.arange(max_seq_len, dtype=torch.bfloat16, device=device)

        angles = torch.einsum("i,j->ij", positions, theta_values)  # bf16 * bf16 = bf16
        angles = torch.cat([angles, angles], dim=-1)  # Shape (max_seq_len, d_model)

        # Standard PyTorch trig functions may upcast bf16 to fp32 for internal computation.
        # Explicitly convert back to the target dtype (bf16) for storage and subsequent use.
        cos_cached = angles.cos().to(dtype)
        sin_cached = angles.sin().to(dtype)
        _rope_cache[cache_key] = (cos_cached, sin_cached)
    return _rope_cache[cache_key]


def apply_rope(x: torch.Tensor, cos_cached: torch.Tensor, sin_cached: torch.Tensor, start_pos: int):
    """Apply RoPE using precomputed cos/sin matrices."""
    # x is expected to be (..., seq_len, dim)
    # cos_cached/sin_cached are (max_seq_len, dim)
    seq_len = x.size(-2)

    cos = cos_cached[start_pos : start_pos + seq_len]  # Shape (seq_len, dim)
    sin = sin_cached[start_pos : start_pos + seq_len]  # Shape (seq_len, dim)

    x1, x2 = x.chunk(2, dim=-1)
    rotated_x = torch.cat((-x2, x1), dim=-1)

    # x, cos, sin are all bf16. Output is bf16.
    return x * cos + rotated_x * sin


class KVCache(nn.Module):
    def __init__(self, kv_cache_shape: tuple) -> None:
        super().__init__()
        self.register_buffer(
            "data", torch.zeros(kv_cache_shape, dtype=torch.bfloat16, device=torch.device("cuda"))
        )
        self.seq_len = 0

    def zero(self) -> None:  # Provided for completeness, matches reference
        self.data.zero_()
        self.seq_len = 0

    def get_data(self) -> torch.Tensor:
        return self.data

    def forward(self, c_kv: torch.Tensor) -> tuple[torch.Tensor, int]:
        num_tokens_to_add = c_kv.size(1)
        if self.seq_len + num_tokens_to_add > self.data.size(1):
            # Match reference error message format
            raise AssertionError(f"KV Cache Exceeded")

        self.data = self.data.to(c_kv.dtype)  # Ensure dtype consistency
        self.data[:, self.seq_len : self.seq_len + num_tokens_to_add, :] = c_kv
        self.seq_len += num_tokens_to_add

        return self.data[:, : self.seq_len], self.seq_len


@dataclass
class Config:  # Matches reference.py structure
    batch_size: int
    dim: int
    n_heads: int
    q_lora_rank: int
    kv_lora_rank: int
    qk_nope_head_dim: int
    qk_rope_head_dim: int
    v_head_dim: int
    seq_len: int  # Input sequence length for x (decode: 1)
    max_seq_len: int  # Max length for RoPE cache, KV cache buffer
    kv_cache_shape: tuple
    Q_proj_down_weight: torch.Tensor
    Q_proj_up_weight: torch.Tensor
    KV_proj_down_weight: torch.Tensor
    KV_proj_up_weight: torch.Tensor
    wo_weight: torch.Tensor


def mla_decode_forward(
    x: torch.Tensor, kv_cache: KVCache, config: Config
) -> tuple[torch.Tensor, torch.Tensor]:
    batch_size, current_sq_len, _ = x.size()  # current_sq_len is config.seq_len, typically 1

    cos_cached, sin_cached = get_rope_cache(
        config.qk_rope_head_dim, config.max_seq_len, device=x.device, dtype=x.dtype
    )

    # Step 1: Down-projection + KV cache update
    q_lora = F.linear(x, config.Q_proj_down_weight)
    kv_lora_current_token = F.linear(x, config.KV_proj_down_weight)

    kv_lora_full_cache_view, total_kv_len = kv_cache(kv_lora_current_token)
    query_pos = total_kv_len - current_sq_len

    # Step 2: Up-project Q, K, V and prepare NoPE + RoPE parts
    # Queries Q (for current token)
    q_nope_and_rope = F.linear(q_lora, config.Q_proj_up_weight).view(
        batch_size,
        current_sq_len,
        config.n_heads,
        config.qk_nope_head_dim + config.qk_rope_head_dim,
    )
    q_nope, q_rope_tensor = q_nope_and_rope.split(
        [config.qk_nope_head_dim, config.qk_rope_head_dim], dim=-1
    )

    # Keys K and Values V (from the entire KV cache)
    kv_nope_full_cache_data, k_rope_tensor_full_cache_data = kv_lora_full_cache_view.split(
        [config.kv_lora_rank, config.qk_rope_head_dim], dim=-1
    )

    # Up-project KV_nope part from cache; K_rope part is taken directly for RoPE
    kv_projected_up_full_cache = F.linear(kv_nope_full_cache_data, config.KV_proj_up_weight).view(
        batch_size, total_kv_len, config.n_heads, config.qk_nope_head_dim + config.v_head_dim
    )
    k_nope_full_cache, v_full_cache = kv_projected_up_full_cache.split(
        [config.qk_nope_head_dim, config.v_head_dim], dim=-1
    )

    # Step 3: Apply RoPE
    # Q RoPE (current token)
    q_rope_tensor = q_rope_tensor.permute(0, 2, 1, 3)  # (bs, n_heads, current_sq_len, rope_dim)
    q_rope_tensor = apply_rope(q_rope_tensor, cos_cached, sin_cached, query_pos)
    q_nope = q_nope.permute(0, 2, 1, 3)
    q = torch.cat([q_nope, q_rope_tensor], dim=-1)

    # K RoPE (entire cache history)
    # Apply RoPE on (bs, 1, total_kv_len, rope_dim) then expand, as in reference and prev. submission.
    k_rope_tensor_full_for_rope = k_rope_tensor_full_cache_data.unsqueeze(
        1
    )  # (bs, 1, total_kv_len, rope_dim)
    k_rope_applied_full = apply_rope(
        k_rope_tensor_full_for_rope, cos_cached, sin_cached, 0
    )  # start_pos=0 for keys
    k_rope_applied_full = k_rope_applied_full.expand(
        -1, config.n_heads, -1, -1
    )  # (bs, n_heads, total_kv_len, rope_dim)

    k_nope_full_cache = k_nope_full_cache.permute(
        0, 2, 1, 3
    )  # (bs, n_heads, total_kv_len, nope_dim)
    k = torch.cat([k_nope_full_cache, k_rope_applied_full], dim=-1)

    v = v_full_cache.permute(0, 2, 1, 3)  # (bs, n_heads, total_kv_len, v_dim)

    # Step 4: Multi-head Attention using scaled_dot_product_attention
    # q: (bs, n_heads, current_sq_len=1, qk_dim)
    # k: (bs, n_heads, total_kv_len, qk_dim)
    # v: (bs, n_heads, total_kv_len, v_dim)
    attn_output = F.scaled_dot_product_attention(
        q, k, v, attn_mask=None, dropout_p=0.0, is_causal=False
    )
    # attn_output shape: (bs, n_heads, current_sq_len, v_dim)

    # Reshape for output projection: (bs, current_sq_len, n_heads * v_dim)
    # Since current_sq_len is 1, this becomes (bs, 1, n_heads * v_dim)
    # We need to bring n_heads next to v_dim.
    # (bs, n_heads, 1, v_dim) -> permute to (bs, 1, n_heads, v_dim) -> view
    y = (
        attn_output.permute(0, 2, 1, 3)
        .contiguous()
        .view(batch_size, current_sq_len, config.n_heads * config.v_head_dim)
    )

    y = F.linear(y, config.wo_weight)

    return y, kv_cache.get_data()


def custom_kernel(data: input_t) -> output_t:
    config, x, kv_cache_instance = data

    attn_out, updated_kv_cache_data = mla_decode_forward(x, kv_cache_instance, config)

    return attn_out, updated_kv_cache_data
python eval.py test /tmp/tmp22sdjzka
Buff name: data, shape: torch.Size([128, 8192, 576])
Buff name: data, shape: torch.Size([128, 8192, 576])
Buff name: data, shape: torch.Size([128, 8192, 576])
Buff name: data, shape: torch.Size([128, 8192, 576])
