import math
import torch
from torch import nn
import torch.nn.functional as F
from typing import Tuple
# Use the new recommended API for SDPA kernel selection
from torch.nn.attention import SDPBackend, sdpa_kernel as sdpa_kernel_dispatch

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

    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, dtype=self.theta.dtype)
        idx_theta = torch.einsum('s,d->sd', seq_idx, self.theta)
        idx_theta2 = torch.cat([idx_theta, idx_theta], dim=-1)
        cos_val = idx_theta2.cos()
        sin_val = idx_theta2.sin()
        return (x * cos_val + self.rotate_half(x) * sin_val).to(x.dtype)

class KVCache(nn.Module):
    def __init__(self, kv_cache_shape: tuple, **kwargs) -> None:
        super().__init__(**kwargs)
        self.register_buffer('data', torch.zeros(kv_cache_shape, dtype=torch.bfloat16))
        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 forward(self, c_kv: torch.Tensor) -> Tuple[torch.Tensor, int]:
        current_batch_size, current_seq_len, _ = c_kv.shape
        if self.seq_len + current_seq_len > self.data.size(1):
            raise ValueError(f"KV Cache Exceeded: cache_len={self.data.size(1)}, current_total_seq_len={self.seq_len + current_seq_len}")
        self.data[
            :current_batch_size, self.seq_len : self.seq_len + current_seq_len, :
        ] = c_kv
        self.seq_len += current_seq_len
        return self.data[:current_batch_size, :self.seq_len, :], self.seq_len

class OptimizedMLA(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.dim = config.dim
        self.n_heads = config.n_heads
        self.q_lora_rank = config.q_lora_rank
        self.kv_lora_rank = config.kv_lora_rank
        self.nope_head_dim = config.qk_nope_head_dim
        self.rope_head_dim = config.qk_rope_head_dim
        self.v_head_dim = config.v_head_dim

        self.Q_proj_down = nn.Linear(self.dim, self.q_lora_rank, dtype=torch.bfloat16, bias=False)
        self.KV_proj_down = nn.Linear(self.dim, self.kv_lora_rank + self.rope_head_dim, dtype=torch.bfloat16, bias=False)
        self.Q_proj_up = nn.Linear(self.q_lora_rank, (self.nope_head_dim + self.rope_head_dim) * self.n_heads, dtype=torch.bfloat16, bias=False)
        self.KV_proj_up = nn.Linear(self.kv_lora_rank, (self.nope_head_dim + self.v_head_dim) * self.n_heads, dtype=torch.bfloat16, bias=False)
        self.wo = nn.Linear(self.v_head_dim * self.n_heads, self.dim, dtype=torch.bfloat16, bias=False)

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

    def forward(self, x: torch.Tensor, kv_cache_instance: KVCache) -> Tuple[torch.Tensor, torch.Tensor]:
        batch_size, seq_len, _ = x.size()

        # --- Q Path ---
        _q_lora = self.Q_proj_down(x)
        _q_up_projected = self.Q_proj_up(_q_lora); del _q_lora
        _q_up_viewed = _q_up_projected.view(batch_size, seq_len, self.n_heads, self.nope_head_dim + self.rope_head_dim); del _q_up_projected
        _q_nope_part_split, _q_rope_part_split = torch.split(_q_up_viewed, [self.nope_head_dim, self.rope_head_dim], dim=-1); del _q_up_viewed
        
        _kv_lora_current_token = self.KV_proj_down(x)
        kv_lora_full, kv_len = kv_cache_instance(_kv_lora_current_token); del _kv_lora_current_token
        query_pos = kv_len - seq_len

        _q_rope_part_permuted = _q_rope_part_split.permute(0, 2, 1, 3); del _q_rope_part_split
        _q_rope_applied = self.q_rope(_q_rope_part_permuted, start_pos=query_pos); del _q_rope_part_permuted
        _q_nope_part_permuted = _q_nope_part_split.permute(0, 2, 1, 3); del _q_nope_part_split
        q = torch.cat([_q_nope_part_permuted, _q_rope_applied], dim=-1); del _q_nope_part_permuted, _q_rope_applied

        # --- KV Path ---
        kv_lora_rank_part_full, k_rope_embedding_full = torch.split(kv_lora_full, [self.kv_lora_rank, self.rope_head_dim], dim=-1)

        _k_rope_to_rotate = k_rope_embedding_full.unsqueeze(1)
        _k_rope_applied = self.k_rope(_k_rope_to_rotate, start_pos=0); del _k_rope_to_rotate
        k_rope_final = _k_rope_applied.expand(-1, self.n_heads, -1, -1); del _k_rope_applied 

        # --- K Nope and V Path (Critical Reordering for Memory) ---
        _kv_up_projected = self.KV_proj_up(kv_lora_rank_part_full) 
        _kv_up_viewed = _kv_up_projected.view(batch_size, kv_len, self.n_heads, self.nope_head_dim + self.v_head_dim)
        _k_nope_part_split_view, _v_full_split_view = torch.split(_kv_up_viewed, [self.nope_head_dim, self.v_head_dim], dim=-1)

        _k_nope_part_permuted = _k_nope_part_split_view.permute(0, 2, 1, 3) 
        v = _v_full_split_view.permute(0, 2, 1, 3) 
        
        del _k_nope_part_split_view 
        del _v_full_split_view    
        del _kv_up_viewed           
        del _kv_up_projected      

        k = torch.cat([_k_nope_part_permuted, k_rope_final], dim=-1) 
        
        del _k_nope_part_permuted 
        del k_rope_final          
        
        torch.cuda.empty_cache() 

        with sdpa_kernel_dispatch(SDPBackend.MATH): 
            attn_output = F.scaled_dot_product_attention(
                q, k, v, 
                attn_mask=None, 
                dropout_p=0.0, 
                is_causal=False 
            )
        del q, k, v 
        torch.cuda.empty_cache() 

        output = attn_output.transpose(1, 2).contiguous().view(batch_size, seq_len, self.n_heads * self.v_head_dim); del attn_output
        output = self.wo(output)

        return output, kv_cache_instance.get_data()

def custom_kernel(data: Tuple) -> Tuple[torch.Tensor, torch.Tensor]:
    config, x, kv_cache_instance = data
    
    with torch.no_grad():
        model = OptimizedMLA(config).to(x.device)
        
        model.Q_proj_down.weight = nn.Parameter(config.Q_proj_down_weight)
        model.Q_proj_up.weight = nn.Parameter(config.Q_proj_up_weight)
        model.KV_proj_down.weight = nn.Parameter(config.KV_proj_down_weight)
        model.KV_proj_up.weight = nn.Parameter(config.KV_proj_up_weight)
        model.wo.weight = nn.Parameter(config.wo_weight)
        
        with torch.amp.autocast(device_type='cuda', dtype=torch.bfloat16, enabled=False): 
            output, updated_kv_data = model(x, kv_cache_instance)
    
    return output, updated_kv_data
python eval.py test /tmp/tmpmwoo67gw
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])
