# ---------------------------  submission.py  ---------------------------------
import math
import torch
from torch import nn
import torch.nn.functional as F
from typing import Tuple

# --------------------------------------------------------------------------- #
# Rotary embedding                                                            #
# --------------------------------------------------------------------------- #
class RoPE(nn.Module):
    def __init__(self, d_model: int):
        super().__init__()
        theta = 10_000 ** (
            -torch.arange(0, d_model // 2, dtype=torch.bfloat16) / (d_model // 2)
        )
        self.register_buffer("theta", theta)

    def rotate_half(self, x: torch.Tensor) -> torch.Tensor:
        x1, x2 = x.chunk(2, dim=-1)
        return torch.cat((-x2, x1), dim=-1)

    def forward(self, x: torch.Tensor, start_pos: int = 0) -> torch.Tensor:
        seq_len = x.size(-2)
        seq_idx = torch.arange(start_pos, start_pos + seq_len, device=x.device)
        idx_theta = torch.einsum("s,d->sd", seq_idx, self.theta)
        idx_theta2 = torch.cat([idx_theta, idx_theta], dim=-1)
        cos = idx_theta2.cos().to(torch.bfloat16)
        sin = idx_theta2.sin().to(torch.bfloat16)
        return x * cos + self.rotate_half(x) * sin

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

    def forward(self, kv_chunk: torch.Tensor):
        self.data[:, self.seq_len : self.seq_len + 1] = kv_chunk
        self.seq_len += 1
        return self.data[:, : self.seq_len], self.seq_len

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

# --------------------------------------------------------------------------- #
# MLA (reference math)                                                        #
# --------------------------------------------------------------------------- #
class MLA(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.dim = cfg.dim
        self.n_heads = cfg.n_heads
        self.q_lora_rank = cfg.q_lora_rank
        self.kv_lora_rank = cfg.kv_lora_rank
        self.nope_head_dim = cfg.qk_nope_head_dim
        self.rope_head_dim = cfg.qk_rope_head_dim
        self.v_head_dim = cfg.v_head_dim

        # down projections
        self.Q_down = nn.Linear(self.dim, self.q_lora_rank, bias=False, dtype=torch.bfloat16)
        self.KV_down = nn.Linear(
            self.dim, self.kv_lora_rank + self.rope_head_dim, bias=False, dtype=torch.bfloat16
        )

        # up projections
        self.Q_up = nn.Linear(
            self.q_lora_rank,
            (self.nope_head_dim + self.rope_head_dim) * self.n_heads,
            bias=False,
            dtype=torch.bfloat16,
        )
        self.KV_up = nn.Linear(
            self.kv_lora_rank,
            (self.nope_head_dim + self.v_head_dim) * self.n_heads,
            bias=False,
            dtype=torch.bfloat16,
        )

        # rotary
        self.q_rope = RoPE(self.rope_head_dim)
        self.k_rope = RoPE(self.rope_head_dim)

        # output proj
        self.wo = nn.Linear(self.v_head_dim * self.n_heads, self.dim, bias=False, dtype=torch.bfloat16)

    # ------------------------------------------------------------------ #
    def forward(self, x: torch.Tensor, kv_cache: KVCache):
        bs, _, _ = x.shape  # seq==1

        # 1) down-projection & append
        q_low = self.Q_down(x)            # [bs,1,dq]
        kv_low = self.KV_down(x)          # [bs,1,512+64]
        kv_low, kv_len = kv_cache(kv_low) # append; kv_len = prefill+1

        # 2) up-projection
        q_long = self.Q_up(q_low).view(bs, 1, self.n_heads,
                                       self.nope_head_dim + self.rope_head_dim)
        q_nope, q_rope_raw = torch.split(q_long,
                                         [self.nope_head_dim, self.rope_head_dim], -1)

        kv_nope_raw, k_rope_raw = torch.split(
            kv_low, [self.kv_lora_rank, self.rope_head_dim], -1
        )
        kv_up = self.KV_up(kv_nope_raw).view(
            bs, kv_len, self.n_heads, self.nope_head_dim + self.v_head_dim
        )
        k_nope, v = torch.split(kv_up, [self.nope_head_dim, self.v_head_dim], -1)

        # 3) apply RoPE
        q_rope = self.q_rope(
            q_rope_raw.permute(0, 2, 1, 3), start_pos=kv_len - 1
        )                                     # [bs,H,1,rope]
        q = torch.cat(
            [q_nope.permute(0, 2, 1, 3), q_rope], -1
        )                                     # [bs,H,1,d_qk]

        # -------- FIXED BLOCK: k_rope_rot construction -------- #
        k_rope_rot = self.k_rope(                         # rotate
            k_rope_raw[:, None, :, :],                    # [bs,1,kv_len,64]
            start_pos=0
        ).expand(-1, self.n_heads, -1, -1)                # [bs,H,kv_len,64]
        # ------------------------------------------------------ #

        k = torch.cat(
            [k_nope.permute(0, 2, 1, 3), k_rope_rot], -1
        )                                     # [bs,H,kv_len,d_qk]

        # 4) attention
        v = v.permute(0, 2, 1, 3)             # [bs,H,kv_len,dv]
        scores = torch.matmul(q, k.transpose(-1, -2)) / math.sqrt(
            self.nope_head_dim + self.rope_head_dim
        )
        attn = F.softmax(scores, -1).to(torch.bfloat16)
        y = torch.matmul(attn, v).view(bs, 1, -1)
        y = self.wo(y)

        # return output + full kv tensor (zeros after seq_len)
        return y, kv_cache.get_data()

# --------------------------------------------------------------------------- #
# Public entry point – rebuild model every call                               #
# --------------------------------------------------------------------------- #
def custom_kernel(data: Tuple):
    cfg, x, kv_cache = data

    model = MLA(cfg).to(x.device)
    model.Q_down.weight  = nn.Parameter(cfg.Q_proj_down_weight)
    model.KV_down.weight = nn.Parameter(cfg.KV_proj_down_weight)
    model.Q_up.weight    = nn.Parameter(cfg.Q_proj_up_weight)
    model.KV_up.weight   = nn.Parameter(cfg.KV_proj_up_weight)
    model.wo.weight      = nn.Parameter(cfg.wo_weight)

    with torch.no_grad():
        return model(x, kv_cache)
# ---------------------------------------------------------------------------
python eval.py benchmark /tmp/tmpyymrs23h
Buff name: data, shape: torch.Size([128, 8192, 576])
Buff name: data, shape: torch.Size([128, 8192, 576])
