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


def apply_rope_optimized(x: torch.Tensor, start_pos: int, rope_dim: int):
    """Optimized RoPE computation"""
    device = x.device
    seq_len = x.size(-2)

    theta = 10000 ** (
        -torch.arange(0, rope_dim // 2, dtype=torch.bfloat16, device=device) / (rope_dim // 2)
    )
    seq_idx = torch.arange(start_pos, start_pos + seq_len, dtype=torch.bfloat16, device=device)
    angles = torch.einsum("s,d->sd", seq_idx, theta)
    angles = torch.cat([angles, angles], dim=-1)
    cos = angles.cos()
    sin = angles.sin()

    x1, x2 = x.chunk(2, dim=-1)
    return x * cos + torch.cat((-x2, x1), dim=-1) * 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="cuda")
        )
        self.seq_len = 0
        self.zero()

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

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

    def forward(self, c_kv: torch.Tensor) -> torch.Tensor:
        assert self.seq_len + c_kv.size(1) <= self.data.size(1), "KV Cache Exceeded"
        self.data = self.data.to(c_kv.dtype)
        self.data[:, self.seq_len : self.seq_len + c_kv.size(1), :] = c_kv
        self.seq_len += c_kv.size(1)
        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
    max_seq_len: int
    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_forward_optimized(x: torch.Tensor, kv_cache: KVCache, config: Config) -> torch.Tensor:
    batch_size, seq_len, model_dim = x.size()

    ################################################################################
    #                 Step 1: Down-projection + KV cache                          #
    ################################################################################
    q_lora = F.linear(x, config.Q_proj_down_weight)
    kv_lora = F.linear(x, config.KV_proj_down_weight)
    kv_lora, kv_len = kv_cache(kv_lora)
    query_pos = kv_len - 1

    ################################################################################
    #                  Step 2: Up-projection with optimized memory layout         #
    ################################################################################
    # Q projection
    q_proj = F.linear(q_lora, config.Q_proj_up_weight)
    q_proj = q_proj.view(
        batch_size, seq_len, config.n_heads, config.qk_nope_head_dim + config.qk_rope_head_dim
    )
    q_nope, q_rope_tensor = q_proj.split([config.qk_nope_head_dim, config.qk_rope_head_dim], dim=-1)

    # KV projection
    kv_nope, k_rope_tensor = kv_lora.split([config.kv_lora_rank, config.qk_rope_head_dim], dim=-1)
    kv_proj = F.linear(kv_nope, config.KV_proj_up_weight)
    kv_proj = kv_proj.view(
        batch_size, kv_len, config.n_heads, config.qk_nope_head_dim + config.v_head_dim
    )
    k_nope, v = kv_proj.split([config.qk_nope_head_dim, config.v_head_dim], dim=-1)

    ################################################################################
    #                    Step 3: RoPE computation with better layout              #
    ################################################################################
    # Transpose to heads-first layout for better performance
    q_nope = q_nope.transpose(1, 2).contiguous()  # [batch, heads, seq, nope_dim]
    q_rope_tensor = q_rope_tensor.transpose(1, 2).contiguous()  # [batch, heads, seq, rope_dim]
    k_nope = k_nope.transpose(1, 2).contiguous()  # [batch, heads, kv_len, nope_dim]
    v = v.transpose(1, 2).contiguous()  # [batch, heads, kv_len, v_dim]

    # Apply RoPE
    q_rope_tensor = apply_rope_optimized(q_rope_tensor, query_pos, config.qk_rope_head_dim)

    k_rope_tensor = k_rope_tensor.unsqueeze(1)
    k_rope_tensor = apply_rope_optimized(k_rope_tensor, 0, config.qk_rope_head_dim)
    k_rope_tensor = k_rope_tensor.expand(-1, config.n_heads, -1, -1).contiguous()

    ################################################################################
    #                    Step 4: Optimized attention computation                  #
    ################################################################################
    # Concatenate Q and K
    q = torch.cat([q_nope, q_rope_tensor], dim=-1)
    k = torch.cat([k_nope, k_rope_tensor], dim=-1)

    # Efficient attention using reshape instead of view
    scale = math.sqrt(config.qk_rope_head_dim + config.qk_nope_head_dim)

    # Reshape for batch matrix multiplication
    q_reshaped = q.reshape(-1, seq_len, q.size(-1))
    k_reshaped = k.reshape(-1, k.size(-2), k.size(-1))
    v_reshaped = v.reshape(-1, v.size(-2), v.size(-1))

    scores = torch.bmm(q_reshaped, k_reshaped.transpose(-1, -2))
    scores = scores.reshape(batch_size, config.n_heads, seq_len, kv_len) / scale

    attn = F.softmax(scores, dim=-1).to(torch.bfloat16)
    attn_reshaped = attn.reshape(-1, seq_len, kv_len)

    y = torch.bmm(attn_reshaped, v_reshaped)
    y = y.reshape(batch_size, config.n_heads, seq_len, config.v_head_dim)

    # Final projection
    y = y.transpose(1, 2).contiguous().view(batch_size, seq_len, -1)
    y = F.linear(y, config.wo_weight)

    return y, kv_cache.get_data()


# entry point for evaluation
def custom_kernel(data: input_t) -> output_t:
    config, x, kv_cache = data
    output, kv_cache = mla_forward_optimized(x, kv_cache, config)
    return output, kv_cache

python eval.py benchmark /tmp/tmp3lyfz4hm
Buff name: data, shape: torch.Size([128, 8192, 576])
Buff name: data, shape: torch.Size([128, 8192, 576])
