import math
from dataclasses import dataclass
from typing import Tuple
import torch
from torch import nn
import torch.nn.functional as F
from task import input_t, output_t

class RoPE(nn.Module):
    def __init__(self, d_model: int):
        super().__init__()
        self.d_model = d_model
        # Ensure theta is created directly on 'cuda'
        theta = 10000 ** (-torch.arange(0, d_model//2, dtype=torch.bfloat16, device='cuda') / (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)
        d_model = x.size(-1)
        assert d_model == self.d_model, f"Dimension mismatch: expected {self.d_model}, got {d_model}"
        # Ensure seq_idx is on the same device as x
        seq_idx = torch.arange(start_pos, start_pos + seq_len, device=x.device)
        
        # self.theta is already on cuda, ensure einsum works correctly.
        # This will operate on x.device (cuda) and self.theta.device (cuda)
        idx_theta = torch.einsum('s,d->sd', seq_idx, self.theta)
        idx_theta2 = torch.cat([idx_theta, idx_theta], dim=-1)
        
        # Ensure cos and sin are converted to bfloat16 explicitly, matching MLA.
        cos = idx_theta2.cos().to(torch.bfloat16)
        sin = idx_theta2.sin().to(torch.bfloat16)
        
        return x * cos + self.rotate_half(x) * 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() # This will reset seq_len to 0

    def zero(self) -> None:
        self.data.zero_()
        self.seq_len = 0 # Ensure seq_len is reset here too, as per MLA's KVCache

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

    def forward(self, c_kv: torch.Tensor) -> Tuple[torch.Tensor, int]: # Specify Tuple return type
        assert self.seq_len + c_kv.size(1) <= self.data.size(1), "KV Cache Exceeded"

        # The self.data = self.data.to(c_kv.dtype) line is important for matching dtypes.
        if self.data.dtype != c_kv.dtype: # Defensive, though inputs should be bfloat16
            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 # Return kv_full and kv_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


# The MLA class itself is not needed for the custom_kernel
# but it serves as a reference for operations.
# No changes to MLA unless it's the custom_kernel we're modifying.

def custom_kernel(data: input_t) -> output_t:
    config, x, kv_cache = data
    
    # Extract dimensions
    batch_size, seq_len, model_dim = x.size()
    
    # Get weights - use direct tensor operations for speed
    # Ensure weights are on the correct device (cuda)
    Q_down_w = config.Q_proj_down_weight.to(x.device)
    Q_up_w = config.Q_proj_up_weight.to(x.device)
    KV_down_w = config.KV_proj_down_weight.to(x.device)
    KV_up_w = config.KV_proj_up_weight.to(x.device)
    wo_w = config.wo_weight.to(x.device)
    
    n_heads = config.n_heads
    q_lora_rank = config.q_lora_rank
    kv_lora_rank = config.kv_lora_rank
    nope_head_dim = config.qk_nope_head_dim
    rope_head_dim = config.qk_rope_head_dim
    v_head_dim = config.v_head_dim
    
    # --- Step 1: Optimized down-projection + KV cache ---
    
    x_reshaped = x.view(-1, model_dim)
    
    q_lora = torch.mm(x_reshaped, Q_down_w.T).view(batch_size, seq_len, q_lora_rank)
    kv_lora_new = torch.mm(x_reshaped, KV_down_w.T).view(batch_size, seq_len, kv_lora_rank + rope_head_dim)
    
    # Update KV cache
    kv_lora_full, kv_len = kv_cache(kv_lora_new)
    query_pos = kv_len - 1
    
    # --- Step 2: Optimized up-projections ---
    
    # Q up-projection
    q_lora_flat = q_lora.view(-1, q_lora_rank)
    q_nope_and_rope = torch.mm(q_lora_flat, Q_up_w.T).view(
        batch_size, seq_len, n_heads, nope_head_dim + rope_head_dim)
    q_nope, q_rope = torch.split(q_nope_and_rope, [nope_head_dim, rope_head_dim], dim=-1)
    
    # Split KV cache
    kv_nope_full, k_rope_full = torch.split(kv_lora_full, [kv_lora_rank, rope_head_dim], dim=-1)
    
    # KV up-projection
    kv_nope_flat = kv_nope_full.contiguous().view(-1, kv_lora_rank)
    kv_nope_up = torch.mm(kv_nope_flat, KV_up_w.T).view(
        batch_size, kv_len, n_heads, nope_head_dim + v_head_dim)
    k_nope, v = torch.split(kv_nope_up, [nope_head_dim, v_head_dim], dim=-1)
    
    # --- Step 3: RoPE computation ---
    
    # Instantiate RoPE modules. Since their __init__ now ensures 'theta' is on 'cuda',
    # we just need to ensure the modules themselves are processed correctly.
    # No need to explicitly set theta here anymore.
    q_rope_module = RoPE(rope_head_dim)
    k_rope_module = RoPE(rope_head_dim)
    
    # Apply RoPE to queries 
    q_rope_perm = q_rope.permute(0, 2, 1, 3)  # bs x n_heads x seq_len x rope_head_dim
    q_rope_rotated = q_rope_module(q_rope_perm, start_pos=query_pos)
    
    q_nope_perm = q_nope.permute(0, 2, 1, 3)  # bs x n_heads x seq_len x nope_head_dim
    q_full = torch.concat([q_nope_perm, q_rope_rotated], dim=-1)
    
    # Apply RoPE to keys 
    k_rope_expanded = k_rope_full[:, None, :, :]  # Add head dimension
    k_rope_rotated = k_rope_module(k_rope_expanded).expand(-1, n_heads, -1, -1)
    k_nope_perm = k_nope.permute(0, 2, 1, 3)  # bs x n_heads x kv_len x nope_head_dim
    k_full = torch.concat([k_nope_perm, k_rope_rotated], dim=-1)
    
    # --- Step 4: Attention computation ---
    
    v_perm = v.permute(0, 2, 1, 3)  # bs x n_heads x kv_len x v_head_dim
    
    # Compute scores with exact same scaling as reference
    head_dim = rope_head_dim + nope_head_dim
    scores = torch.matmul(q_full, k_full.transpose(-1, -2)) / math.sqrt(head_dim)
    attn = F.softmax(scores, dim=-1).to(torch.bfloat16) # Ensure bfloat16 for attention weights
    
    # Apply attention
    y = torch.matmul(attn, v_perm).view(batch_size, seq_len, -1)
    
    # Output projection
    y_flat = y.view(-1, n_heads * v_head_dim)
    output = torch.mm(y_flat, wo_w.T).view(batch_size, seq_len, model_dim)
    
    return output, kv_cache.get_data()
python eval.py leaderboard /tmp/tmp0sh1xpe_
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])
