import math
import torch
import torch.nn.functional as F
from torch import nn
from dataclasses import dataclass
from typing import Tuple
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 in float32 for numerical stability
        base_val = torch.tensor(10000.0, dtype=torch.float32)
        exponent = torch.arange(0, d_model // 2, dtype=torch.float32) / (d_model // 2)
        theta = 1.0 / torch.pow(base_val, exponent)
        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)
        device = x.device
        
        # Compute positions in float32 then convert to bfloat16
        positions = torch.arange(start_pos, start_pos + seq_len, dtype=torch.float32, device=device)
        positions = positions.to(torch.bfloat16)
        
        # Compute rotation angles with higher precision
        idx_theta = torch.einsum('s,d->sd', positions, self.theta.to(device).to(torch.bfloat16))
        idx_theta2 = torch.cat([idx_theta, idx_theta], dim=-1)
        
        # Compute trig functions in float32 then convert
        cos = torch.cos(idx_theta2.to(torch.float32)).to(torch.bfloat16)
        sin = torch.sin(idx_theta2.to(torch.float32)).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

    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"
        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

def custom_kernel(data: input_t) -> output_t:
    config, x, kv_cache = data
    device = x.device
    
    # Unpack config with explicit casting
    batch_size, seq_len, model_dim = x.size()
    Q_down_w = config.Q_proj_down_weight.to(device).to(torch.bfloat16)
    Q_up_w = config.Q_proj_up_weight.to(device).to(torch.bfloat16)
    KV_down_w = config.KV_proj_down_weight.to(device).to(torch.bfloat16)
    KV_up_w = config.KV_proj_up_weight.to(device).to(torch.bfloat16)
    wo_w = config.wo_weight.to(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: Down-projection ---
    x_reshaped = x.reshape(-1, model_dim)
    
    # Use torch.matmul instead of F.linear for precision
    q_lora = torch.matmul(x_reshaped, Q_down_w.T).view(batch_size, seq_len, q_lora_rank)
    kv_lora_new = torch.matmul(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: Up-projections ---
    q_lora_flat = q_lora.reshape(-1, q_lora_rank)
    q_nope_and_rope = torch.matmul(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 with explicit contiguous
    kv_nope_full, k_rope_full = torch.split(
        kv_lora_full, [kv_lora_rank, rope_head_dim], dim=-1
    )
    kv_nope_full = kv_nope_full.contiguous()
    
    kv_nope_flat = kv_nope_full.view(-1, kv_lora_rank)
    kv_nope_up = torch.matmul(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).to(device)
    k_rope_module = RoPE(rope_head_dim).to(device)
    
    # Apply RoPE to query (single token)
    q_rope_perm = q_rope.permute(0, 2, 1, 3)
    q_rope_rotated = q_rope_module(q_rope_perm, start_pos=query_pos)
    
    # Apply RoPE to keys (entire cache)
    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)
    
    # --- Step 4: Attention computation ---
    q_nope_perm = q_nope.permute(0, 2, 1, 3)
    q_full = torch.cat([q_nope_perm, q_rope_rotated], dim=-1)
    
    k_nope_perm = k_nope.permute(0, 2, 1, 3)
    k_full = torch.cat([k_nope_perm, k_rope_rotated], dim=-1)
    
    v_perm = v.permute(0, 2, 1, 3)
    
    head_dim = rope_head_dim + nope_head_dim
    # Compute scale in float32 then convert
    scale = torch.tensor(1.0, dtype=torch.float32, device=device) / math.sqrt(head_dim)
    scale = scale.to(torch.bfloat16)
    
    # Use explicit attention computation
    scores = torch.matmul(q_full, k_full.transpose(-1, -2)) * scale
    attn_weights = F.softmax(scores, dim=-1)
    attn_output = torch.matmul(attn_weights, v_perm)
    
    # --- Step 5: Output projection ---
    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.matmul(y_flat, wo_w.T).view(batch_size, seq_len, model_dim)
    
    return output, kv_cache.get_data()

# Optimize with torch.compile for MI300
custom_kernel_compiled = torch.compile(custom_kernel, mode="max-autotune")
python eval.py test /tmp/tmpnyy6j4bb
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])
