# code modifications can be made from here
import math
from dataclasses import dataclass
import torch
from torch import nn
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="cuda"):
    """Get or create cached RoPE rotation matrices"""
    cache_key = (d_model, max_seq_len, device)
    if cache_key not in _rope_cache:
        theta = 10000 ** (
            -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)
        angles = torch.cat([angles, angles], dim=-1)
        cos_cached = angles.cos()
        sin_cached = angles.sin()
        _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"""
    seq_len = x.size(-2)
    cos = cos_cached[start_pos : start_pos + seq_len]
    sin = sin_cached[start_pos : start_pos + seq_len]

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

    return x * cos + rotated * sin


class KVCache(nn.Module):  # This is the cache for DOWN-projected K/V values
    def __init__(self, kv_cache_shape: tuple) -> None:
        super().__init__()
        self.register_buffer(
            "data", torch.zeros(kv_cache_shape, dtype=torch.bfloat16, device="cuda")
        )
        self.seq_len = 0
        # self.zero() # Not strictly needed, torch.zeros does it.

    def zero(self) -> None:
        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]:  # c_kv is for the current token(s)
        num_new_tokens = c_kv.size(1)
        assert self.seq_len + num_new_tokens <= self.data.size(1), "KVCache Exceeded"
        self.data = self.data.to(c_kv.dtype)
        self.data[:, self.seq_len : self.seq_len + num_new_tokens, :] = c_kv
        self.seq_len += num_new_tokens
        return self.data[:, : self.seq_len], self.seq_len


# New Cache for UP-projected K/V "nope" components
class KVUpCache(nn.Module):
    def __init__(self, up_cache_shape: tuple) -> None:
        super().__init__()
        self.register_buffer(
            "data", torch.zeros(up_cache_shape, dtype=torch.bfloat16, device="cuda")
        )
        self.seq_len = 0

    def zero(self) -> None:
        self.data.zero_()
        self.seq_len = 0

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

    def append_and_get_history(self, c_kv_up: torch.Tensor) -> tuple[torch.Tensor, int]:
        # c_kv_up is the up-projected "nope" K/V for the current token(s)
        num_new_tokens = c_kv_up.size(1)
        assert self.seq_len + num_new_tokens <= self.data.size(1), "KVUpCache Exceeded"
        self.data = self.data.to(c_kv_up.dtype)
        self.data[:, self.seq_len : self.seq_len + num_new_tokens, :] = c_kv_up
        self.seq_len += num_new_tokens
        return self.data[:, : self.seq_len], self.seq_len


@dataclass
class Config:
    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  # seq_len of current input x, typically 1 for decoding
    max_seq_len: int  # max capacity of caches
    kv_cache_shape: tuple  # For KVCache (down-projected)
    # kv_up_cache_shape: tuple # Will be derived in custom_kernel
    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_forward(
    x: torch.Tensor, kv_cache_down: KVCache, kv_cache_up: KVUpCache, config: Config
) -> tuple[torch.Tensor, torch.Tensor]:
    batch_size, current_input_seq_len, _ = (
        x.size()
    )  # current_input_seq_len is 1 for decoding context

    cos_cached, sin_cached = get_rope_cache(config.qk_rope_head_dim, config.max_seq_len)

    ################################################################################
    #       Step 1: Down-project, (conditionally) Up-project, Cache               #
    ################################################################################
    # Query path (down-projection)
    q_lora = F.linear(x, config.Q_proj_down_weight)  # Shape: (bs, 1, q_lora_rank)

    # Key/Value path (down-projection for the current token)
    # kv_lora_current_token_down: (bs, 1, kv_lora_rank + qk_rope_head_dim)
    kv_lora_current_token_down = F.linear(x, config.KV_proj_down_weight)

    # Append current token's down-projected K/V to kv_cache_down
    # kv_lora_full_history_down: (bs, kv_len, kv_lora_rank + qk_rope_head_dim)
    kv_lora_full_history_down, kv_len = kv_cache_down(kv_lora_current_token_down)
    query_pos = (
        kv_len - current_input_seq_len
    )  # Start position for RoPE for the current query token(s)

    # Split the *current token's* down-projected K/V to get the part for up-projection
    # kv_nope_current_token_down: (bs, 1, kv_lora_rank)
    kv_nope_current_token_down, _ = kv_lora_current_token_down.split(
        [config.kv_lora_rank, config.qk_rope_head_dim], dim=-1
    )

    # Up-project ONLY the "nope" part of the *current token's* K/V
    # This is the F.linear operation that is now significantly smaller.
    # kv_nope_current_token_up: (bs, 1, n_heads * (qk_nope_head_dim + v_head_dim))
    kv_nope_current_token_up = F.linear(kv_nope_current_token_down, config.KV_proj_up_weight)

    # Append current token's up-projected "nope" K/V to kv_cache_up
    # kv_nope_full_history_up: (bs, kv_len, n_heads * (qk_nope_head_dim + v_head_dim))
    kv_nope_full_history_up, _ = kv_cache_up.append_and_get_history(kv_nope_current_token_up)

    ################################################################################
    #                  Step 2: Prepare Q, K, V from caches                       #
    ################################################################################
    # Queries Q (up-project and split)
    q_nope_and_rope = F.linear(q_lora, config.Q_proj_up_weight).view(
        batch_size,
        current_input_seq_len,
        config.n_heads,
        config.qk_nope_head_dim + config.qk_rope_head_dim,
    )
    q_nope, q_rope_tensor_for_q = q_nope_and_rope.split(
        [config.qk_nope_head_dim, config.qk_rope_head_dim], dim=-1
    )

    # Keys K and Values V
    # The RoPE part of K comes from the down-projected cache (kv_lora_full_history_down)
    _, k_rope_tensor_for_k_full = kv_lora_full_history_down.split(
        [config.kv_lora_rank, config.qk_rope_head_dim], dim=-1
    )  # Shape: (bs, kv_len, qk_rope_head_dim)

    # The "nope" part of K and all of V come from the up-projected cache (kv_nope_full_history_up)
    # Reshape to identify heads for k_nope and v
    up_projected_kv_view = kv_nope_full_history_up.view(
        batch_size, kv_len, config.n_heads, config.qk_nope_head_dim + config.v_head_dim
    )
    k_nope_full, v_full = up_projected_kv_view.split(
        [config.qk_nope_head_dim, config.v_head_dim], dim=-1
    )  # k_nope_full: (bs, kv_len, n_heads, qk_nope_head_dim)
    # v_full: (bs, kv_len, n_heads, v_head_dim)

    ################################################################################
    #                    Step 3: Apply RoPE                                      #
    ################################################################################
    # RoPE for Queries
    q_rope_tensor_for_q = q_rope_tensor_for_q.permute(
        0, 2, 1, 3
    )  # (bs, n_heads, 1, qk_rope_head_dim)
    q_rope_tensor_for_q = apply_rope(q_rope_tensor_for_q, cos_cached, sin_cached, query_pos)
    q_nope = q_nope.permute(0, 2, 1, 3)  # (bs, n_heads, 1, qk_nope_head_dim)
    q = torch.cat([q_nope, q_rope_tensor_for_q], dim=-1)  # q: (bs, n_heads, 1, head_dim_qk)

    # RoPE for Keys
    # k_rope_tensor_for_k_full is (bs, kv_len, qk_rope_head_dim)
    k_rope_tensor_for_k_full_roped = k_rope_tensor_for_k_full.unsqueeze(
        1
    )  # (bs, 1, kv_len, qk_rope_head_dim)
    k_rope_tensor_for_k_full_roped = apply_rope(
        k_rope_tensor_for_k_full_roped, cos_cached, sin_cached, 0
    )  # Keys use start_pos=0
    k_rope_tensor_for_k_full_roped = k_rope_tensor_for_k_full_roped.expand(
        -1, config.n_heads, -1, -1
    )  # (bs, n_heads, kv_len, qk_rope_head_dim)

    k_nope_full = k_nope_full.permute(0, 2, 1, 3)  # (bs, n_heads, kv_len, qk_nope_head_dim)
    k = torch.cat(
        [k_nope_full, k_rope_tensor_for_k_full_roped], dim=-1
    )  # k: (bs, n_heads, kv_len, head_dim_qk)

    ################################################################################
    #                        Step 4: Compute Multi-head Attention                  #
    ################################################################################
    v_full = v_full.permute(0, 2, 1, 3)  # (bs, n_heads, kv_len, v_head_dim)

    scale = math.sqrt(config.qk_rope_head_dim + config.qk_nope_head_dim)
    scores = torch.matmul(q, k.transpose(-1, -2)) / scale
    attn = F.softmax(scores, dim=-1).to(torch.bfloat16)  # Ensure bfloat16 for matmul with v

    y = torch.matmul(attn, v_full)  # y: (bs, n_heads, 1, v_head_dim)
    y = y.view(batch_size, current_input_seq_len, -1)  # (bs, 1, n_heads * v_head_dim)
    y = F.linear(y, config.wo_weight)

    return y, kv_cache_down.get_data()  # Return down-projected cache data as per original output_t


# entry point for evaluation


def custom_kernel(data: input_t) -> output_t:
    config, x, ref_kv_cache = data  # ref_kv_cache stores DOWN-projected values

    # 1. Setup KVCache for down-projected values (our_kv_cache_down)
    # This cache stores the direct output of KV_proj_down_weight
    our_kv_cache_down = KVCache(config.kv_cache_shape)
    ref_data_down = ref_kv_cache.get_data()  # Full capacity tensor from ref
    ref_seq_len = ref_kv_cache.seq_len  # Actual length of prefilled data

    # Copy prefilled data to our_kv_cache_down
    if ref_seq_len > 0:
        our_kv_cache_down.data[:, :ref_seq_len, :] = ref_data_down[:, :ref_seq_len, :]
    our_kv_cache_down.seq_len = ref_seq_len

    # 2. Setup KVUpCache for up-projected "nope" K/V components (our_kv_cache_up)
    # This cache stores the output of KV_proj_up_weight (applied to "nope" part)
    up_projected_dim = config.n_heads * (config.qk_nope_head_dim + config.v_head_dim)
    kv_up_cache_shape = (config.batch_size, config.max_seq_len, up_projected_dim)
    our_kv_cache_up = KVUpCache(kv_up_cache_shape)

    # Prefill our_kv_cache_up based on ref_kv_cache's prefilled (down-projected) data
    if ref_seq_len > 0:
        # Extract the "nope" part from the down-projected prefill data
        # ref_data_down_prefill_slice: (bs, ref_seq_len, kv_lora_rank + qk_rope_head_dim)
        ref_data_down_prefill_slice = ref_data_down[:, :ref_seq_len, :]

        pre_filled_nope_down, _ = ref_data_down_prefill_slice.split(
            [config.kv_lora_rank, config.qk_rope_head_dim], dim=-1
        )  # pre_filled_nope_down: (bs, ref_seq_len, kv_lora_rank)

        # Up-project this "nope" part. This is the one-time cost for prefill.
        # Ensure contiguous for potentially better GEMM performance.
        pre_filled_nope_up = F.linear(pre_filled_nope_down.contiguous(), config.KV_proj_up_weight)

        our_kv_cache_up.data[:, :ref_seq_len, :] = pre_filled_nope_up
        our_kv_cache_up.seq_len = ref_seq_len

    # Run the forward pass with both caches
    output, final_kv_cache_down_data = mla_forward(x, our_kv_cache_down, our_kv_cache_up, config)

    # Return structure matches output_t: (output_tensor, kv_cache_data_tensor)
    # The second tensor should be the data from the primary K/V cache, which is kv_cache_down.
    return output, final_kv_cache_down_data
python eval.py test /tmp/tmpxe_6duzs
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])
