import math
from typing import Tuple
import torch
from torch import nn
import torch.nn.functional as F


class RoPE(nn.Module):
    def __init__(self, d_model: int) -> None:
        super().__init__()
        theta = 10000 ** (
            -torch.arange(0, d_model // 2, dtype=torch.bfloat16) / (d_model // 2)
        )
        self.register_buffer("theta", theta)

    @staticmethod
    def _rotate_half(x: torch.Tensor) -> torch.Tensor:
        a, b = x.chunk(2, dim=-1)
        return torch.cat((-b, a), dim=-1)

    def forward(self, x: torch.Tensor, pos: int) -> torch.Tensor:
        seq = x.size(-2)
        idx = torch.arange(pos, pos + seq, device=x.device)
        ang = torch.einsum("s,d->sd", idx, self.theta)
        ang = torch.cat([ang, ang], dim=-1)
        return x * ang.cos() + self._rotate_half(x) * ang.sin()


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

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

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


class MLA(nn.Module):
    def __init__(self, cfg) -> None:
        super().__init__()
        self.dim = cfg.dim
        self.n_heads = cfg.n_heads
        self.q_low = cfg.q_lora_rank
        self.kv_low = cfg.kv_lora_rank
        self.nope = cfg.qk_nope_head_dim
        self.rope = cfg.qk_rope_head_dim
        self.v_dim = cfg.v_head_dim

        self.q_down = nn.Linear(self.dim, self.q_low, bias=False, dtype=torch.bfloat16)
        self.kv_down = nn.Linear(
            self.dim, self.kv_low + self.rope, bias=False, dtype=torch.bfloat16
        )

        self.q_up = nn.Linear(
            self.q_low,
            (self.nope + self.rope) * self.n_heads,
            bias=False,
            dtype=torch.bfloat16,
        )
        self.kv_up = nn.Linear(
            self.kv_low,
            (self.nope + self.v_dim) * self.n_heads,
            bias=False,
            dtype=torch.bfloat16,
        )

        self.q_rope = RoPE(self.rope)
        self.k_rope = RoPE(self.rope)

        self.wo = nn.Linear(self.v_dim * self.n_heads, self.dim,
                            bias=False, dtype=torch.bfloat16)

    def forward(self, x: torch.Tensor, cache: KVCache):
        bs = x.size(0)

        q_low = self.q_down(x)
        kv_low = self.kv_down(x)
        kv_low, kv_len = cache(kv_low)

        q_big = self.q_up(q_low).view(
            bs, 1, self.n_heads, self.nope + self.rope
        )
        q_nope, q_rope_raw = torch.split(q_big, [self.nope, self.rope], -1)

        kv_nope_raw, k_rope_raw = torch.split(
            kv_low, [self.kv_low, self.rope], -1
        )
        kv_big = self.kv_up(kv_nope_raw).view(
            bs, kv_len, self.n_heads, self.nope + self.v_dim
        )
        k_nope, v = torch.split(kv_big, [self.nope, self.v_dim], -1)

        q_rope = self.q_rope(q_rope_raw.permute(0, 2, 1, 3), kv_len - 1)
        q = torch.cat([q_nope.permute(0, 2, 1, 3), q_rope], -1)

        k_rope = self.k_rope(k_rope_raw[:, None, :, :], 0).expand(
            -1, self.n_heads, -1, -1
        )
        k = torch.cat([k_nope.permute(0, 2, 1, 3), k_rope], -1)

        v = v.permute(0, 2, 1, 3)
        scale = math.sqrt(self.nope + self.rope)
        attn = F.softmax(torch.matmul(q, k.transpose(-1, -2)) / scale, -1)
        attn = attn.to(torch.bfloat16)
        y = torch.matmul(attn, v).view(bs, 1, -1)
        y = self.wo(y)
        return y, cache.get_data()


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/tmptrl2zp1l
Buff name: data, shape: torch.Size([128, 8192, 576])
Buff name: data, shape: torch.Size([128, 8192, 576])
