import os
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:256"
import math
import torch
import torch.nn.functional as F
from torch import nn
from dataclasses import dataclass
import functools
from typing import Tuple
from task import input_t, output_t


@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

# Use torch.compile for kernel fusion
@torch.compile(mode='max-autotune')
def fused_attention(q_full, k_full, v_perm):
    # Use optimized scaled dot product attention
    return F.scaled_dot_product_attention(q_full, k_full, v_perm, is_causal=False)

# optim_fused_attnention = torch.compile(fused_attention, mode='max-autotune', fullgraph=True, dynamic=False)

class RoPE(nn.Module):
    def __init__(self, d_model: int):
        super().__init__()
        self.d_model = d_model
        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.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)
        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)

        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

    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.to(self.data.device)
        self.seq_len += c_kv.size(1)
        return self.data[:, :self.seq_len], self.seq_len


# Pre-instantiate RoPE modules if rope_head_dim is fixed (or move inside forward if dynamic)
qk_rope_head_dim = 128  # adjust based on your config
q_rope_module = RoPE(qk_rope_head_dim).to('cuda')
k_rope_module = RoPE(qk_rope_head_dim).to('cuda')

torch.cuda.empty_cache()


def custom_kernel(data: input_t) -> output_t:
    config, x, kv_cache = data
    device = x.device

    batch_size, seq_len, model_dim = x.size()

    # Extract weights once and keep in bfloat16
    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
    qk_nope_head_dim = config.qk_nope_head_dim
    qk_rope_head_dim = config.qk_rope_head_dim
    v_head_dim = config.v_head_dim

    # --- Step 1: Down-projection ---
    x_reshaped = x.view(-1, model_dim)
    q_lora = torch.matmul(x_reshaped, Q_down_w.T).view(batch_size, seq_len, config.q_lora_rank)
    kv_lora_new = torch.matmul(x_reshaped, KV_down_w.T).view(
        batch_size, seq_len, config.kv_lora_rank + qk_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.view(-1, config.q_lora_rank)
    q_nope_and_rope = torch.matmul(q_lora_flat, Q_up_w.T).view(
        batch_size, seq_len, n_heads, qk_nope_head_dim + qk_rope_head_dim
    )
    q_nope, q_rope = torch.split(q_nope_and_rope, [qk_nope_head_dim, qk_rope_head_dim], dim=-1)

    kv_nope_full, k_rope_full = torch.split(kv_lora_full, [config.kv_lora_rank, qk_rope_head_dim], dim=-1)

    kv_nope_flat = kv_nope_full.contiguous().view(-1, config.kv_lora_rank)
    kv_nope_up = torch.matmul(kv_nope_flat, KV_up_w.T).view(
        batch_size, kv_len, n_heads, qk_nope_head_dim + v_head_dim
    )
    k_nope, v = torch.split(kv_nope_up, [qk_nope_head_dim, v_head_dim], dim=-1)

    # --- Step 3: Apply RoPE (create RoPE modules dynamically) ---
    q_rope_perm = q_rope.permute(0, 2, 1, 3)  # (B, H, S, D)

    # Create new RoPE instance with correct head dim
    rope_head_dim = q_rope_perm.size(-1)

    # @functools.lru_cache(maxsize=None)
    # def get_rope(head_dim):
    #     return RoPE(head_dim).to('cuda')
    
    # q_rope_module = get_rope(rope_head_dim)
    q_rope_module = RoPE(rope_head_dim).to('cuda')
    q_rope_rotated = q_rope_module(q_rope_perm, start_pos=query_pos)

    k_rope_rotated_shared = RoPE(k_rope_full.size(-1)).to(device)(k_rope_full.unsqueeze(1), start_pos=0)
    k_rope_rotated = k_rope_rotated_shared.expand(-1, n_heads, -1, -1)

    # --- Step 4: Combine components ---
    q_full = torch.cat([q_nope.permute(0, 2, 1, 3), q_rope_rotated], dim=-1)
    k_full = torch.cat([k_nope.permute(0, 2, 1, 3), k_rope_rotated], dim=-1)
    v_perm = v.permute(0, 2, 1, 3)

    # --- Step 5: Compute attention efficiently ---
    attn_output = fused_attention(q_full, k_full, v_perm)

    # --- Step 6: 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()
python eval.py benchmark /tmp/tmp48mtfkk3
Buff name: data, shape: torch.Size([128, 8192, 576])
Buff name: data, shape: torch.Size([128, 8192, 576])
