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
        # Calculate theta on CUDA in float32 first for better numerical stability
        base_val = torch.tensor(10000.0, dtype=torch.float32, device='cuda')
        exponent = torch.arange(0, d_model // 2, dtype=torch.float32, device='cuda') / (d_model // 2)
        theta = 1.0 / torch.pow(base_val, exponent)
        self.register_buffer("theta", theta.to(torch.bfloat16))

    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)
        # Calculate positions on the same device as input
        positions = torch.arange(start_pos, start_pos + seq_len, dtype=torch.float32, device=x.device)
        positions = positions.to(torch.bfloat16)
        
        idx_theta = torch.einsum('s,d->sd', positions, self.theta)
        idx_theta2 = torch.cat([idx_theta, idx_theta], dim=-1)
        
        # Keep cos/sin in bfloat16 but ensure proper calculation
        cos = torch.cos(idx_theta2).to(torch.bfloat16)
        sin = torch.sin(idx_theta2).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()

    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]:
        assert self.seq_len + c_kv.size(1) <= self.data.size(1), "KV Cache Exceeded"
        
        # This conversion might be a source of mismatch if the reference doesn't do it,
        # but given the problem spec implies bfloat16 throughout, c_kv should be bfloat16.
        # If it's not, converting self.data to c_kv.dtype is the right thing to do.
        if self.data.dtype != c_kv.dtype:
            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 custom_kernel(data: input_t) -> output_t:
    config, x, kv_cache = data
    
    batch_size, seq_len, model_dim = x.size()
    
    # Ensure all weights are bfloat16 from the start
    Q_down_w = config.Q_proj_down_weight.to(x.device).to(torch.bfloat16)
    Q_up_w = config.Q_proj_up_weight.to(x.device).to(torch.bfloat16)
    KV_down_w = config.KV_proj_down_weight.to(x.device).to(torch.bfloat16)
    KV_up_w = config.KV_proj_up_weight.to(x.device).to(torch.bfloat16)
    wo_w = config.wo_weight.to(x.device).to(torch.bfloat16)
    
    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) # x should be bfloat16 from input_t
    
    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)
    
    kv_lora_full, kv_len = kv_cache(kv_lora_new)
    
    query_pos = kv_len - 1 
    
    # --- Step 2: Optimized up-projections ---
    
    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)
    
    kv_nope_full, k_rope_full = torch.split(kv_lora_full, [kv_lora_rank, rope_head_dim], dim=-1)
    
    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 ---
    
    q_rope_module = RoPE(rope_head_dim)
    k_rope_module = RoPE(rope_head_dim)
    
    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)
    q_full = torch.cat([q_nope_perm, q_rope_rotated], dim=-1)
    
    # Apply RoPE to keys (unsqueeze for head dim, apply RoPE, then expand)
    # CRITICAL CHANGE: Ensure k_rope_full.unsqueeze(1) is then passed to .to(torch.bfloat16)
    # before applying RoPE module, just to be absolutely sure of precision if it was somehow promoted.
    k_rope_rotated_shared = k_rope_module(k_rope_full.unsqueeze(1).to(torch.bfloat16), start_pos=0)
    k_rope_rotated = k_rope_rotated_shared.expand(-1, n_heads, -1, -1)

    k_nope_perm = k_nope.permute(0, 2, 1, 3)
    k_full = torch.cat([k_nope_perm, k_rope_rotated], dim=-1)
    
    # --- Step 4: Attention computation ---
    
    v_perm = v.permute(0, 2, 1, 3)
    
    head_dim = rope_head_dim + nope_head_dim
    # CRITICAL CHANGE: Ensure scale is computed in bfloat16
    scale = torch.tensor(1.0, dtype=torch.float32, device=x.device) / math.sqrt(head_dim)
    scale = scale.to(torch.bfloat16)
    
    # Matmul results in bfloat16 if inputs are bfloat16. Multiplication with bfloat16 scale.
    scores = torch.matmul(q_full, k_full.transpose(-1, -2)) * scale
    
    # No explicit masking needed for seq_len=1 decode.
    
    attn_weights = F.softmax(scores, dim=-1) 
    
    attn_output = torch.matmul(attn_weights, v_perm)
    
    y = attn_output.transpose(1, 2).contiguous().view(batch_size, seq_len, n_heads * v_head_dim)
    
    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 benchmark /tmp/tmpk_vwb3a9
Buff name: data, shape: torch.Size([128, 8192, 576])
Buff name: data, shape: torch.Size([128, 8192, 576])
