import math
from dataclasses import dataclass
import torch
from torch import nn
import torch.nn.functional as F
from task import input_t, output_t
from typing import Tuple
import time
import triton
import triton.language as tl
class RoPE(nn.Module):
    def __init__(self, d_model: int):
        super().__init__()
        self.d_model = d_model
        theta = 10000 ** (-torch.arange(0, d_model//2,dtype=torch.bfloat16) / (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
        seq_idx = torch.arange(start_pos, start_pos + seq_len, device=x.device)
        idx_theta = torch.einsum('s,d->sd', seq_idx, self.theta)
        idx_theta2 = torch.cat([idx_theta, idx_theta], dim=-1)
        torch.cos
        cos = idx_theta2.cos().to(torch.bfloat16)
        sin = idx_theta2.sin().to(torch.bfloat16)
        return x * cos + self.rotate_half(x) * sin

class MyRoPE(nn.Module):
    def __init__(self, d_model: int):
        super().__init__()
        self.d_model = d_model
        theta = 10000 ** (-torch.arange(0, d_model//2,dtype=torch.bfloat16) / (d_model//2))
        self.theta = theta.cuda() 
    
    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
        old_shape = x.shape
        x = x.reshape(-1, seq_len, d_model)
        seq_idx = torch.arange(start_pos, start_pos + seq_len, device=x.device) # (seq_len,)
        # print(f"seq_idx: {seq_idx} dtype: {seq_idx.dtype} shape: {seq_idx.shape}")
        idx_theta = torch.einsum('s,d->sd', seq_idx, self.theta) # (seq_len, d_model//2)
        # print(f"idx_theta dtype: {idx_theta.dtype} ")
        idx_theta = idx_theta.view(-1, seq_len, d_model//2) # (-1, seq_len, d_model//2)
        x1 = x[:, :, :d_model//2] # (-1, seq_len, d_model//2)
        x2 = x[:, :, d_model//2:] # (-1, seq_len, d_model//2)
        # print(f"cosine dtype: {idx_theta.cos().dtype} ")
        # print(f"x1 * idx_theta.cos() dtype: {(x1 * idx_theta.cos()).dtype} ")
        new_x1 = x1 * idx_theta.cos() - x2 * idx_theta.sin() # (-1, seq_len, d_model//2)
        new_x2 = x1 * idx_theta.sin() + x2 * idx_theta.cos() # (-1, seq_len, d_model//2)
        # print(f"new_x1 dtype: {new_x1.dtype} new_x2 dtype: {new_x2.dtype}")
        new_x = torch.cat([new_x1, new_x2], dim=-1) # (-1, seq_len, d_model)
        new_x = new_x.view(old_shape) # (batch_size, seq_len, d_model)
        return new_x
        
        
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))
        self.seq_len = 0
        self.zero()

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

    def forward(self, c_kv: torch.Tensor) -> torch.Tensor:
        assert self.seq_len + c_kv.size(1) <= self.data.size(1), "KV Cache Exceeded"

        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
class MLA(nn.Module):
    def __init__(self, config: 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
        # Down-projection matrices
        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)

        # Up-projection and rope projection matrices
        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)

        # RoPE on half embeddings
        self.q_rope = RoPE(self.rope_head_dim)
        self.k_rope = RoPE(self.rope_head_dim)

        # Output projection
        self.wo = nn.Linear(self.v_head_dim * self.n_heads, self.dim, dtype=torch.bfloat16, bias=False)
        self.eps = 1e-6
   
    def forward(self, x: torch.Tensor, kv_cache: KVCache) -> torch.Tensor:
        # seq_len = 1 always here
        batch_size, seq_len, model_dim = x.size()

        ################################################################################
        #                 Step 1: Handle down-projection + KV cache                    #
        ################################################################################
        q_lora = self.Q_proj_down(x)
        kv_lora = self.KV_proj_down(x)
        kv_lora, kv_len = kv_cache(kv_lora)
        query_pos = kv_len - 1

        ################################################################################
        #                  Step 2: Up-project and prepare NoPE + RoPE                  #
        ################################################################################

        # Handle queries Q first
        q_nope_and_rope = self.Q_proj_up(q_lora).view(
            batch_size, seq_len, self.n_heads, self.nope_head_dim + self.rope_head_dim)
        q_nope, q_rope = torch.split(q_nope_and_rope, [self.nope_head_dim, self.rope_head_dim], dim=-1)

        # Handle keys and values K/V. V does not need RoPE
        kv_nope, k_rope = torch.split(kv_lora, [self.kv_lora_rank, self.rope_head_dim], dim=-1)
        kv_nope = self.KV_proj_up(kv_nope).view(
            batch_size, kv_len, self.n_heads, self.nope_head_dim + self.v_head_dim)
        k_nope, v = torch.split(kv_nope, [self.nope_head_dim, self.v_head_dim], dim=-1)

        ################################################################################
        #                    Step 3: Handle RoPE Stream                                #
        ################################################################################

        # Compute RoPE for queries and combine with no-RoPE part
        q_rope = q_rope.permute(0, 2, 1, 3) # bs x n_heads x seq_len x rope_head_dim
        q_rope = self.q_rope(q_rope, start_pos=query_pos)

        q_nope = q_nope.permute(0, 2, 1, 3) # bs x n_heads x seq_len x rope_head_dim
        q = torch.concat([q_nope, q_rope], dim=-1)


        # Compute RoPE for keys and combine with no-RoPE part
        k_rope = k_rope[:, None, :, :]
        k_rope = self.k_rope(k_rope).expand(-1,self.n_heads,-1,-1)
        k_nope = k_nope.permute(0, 2, 1, 3) # bs x kv_len x n_heads x rope_head_dim
        k = torch.concat([k_nope, k_rope], dim=-1)
                
        ################################################################################
        #                        Compute Multi-head Attention                          #
        ################################################################################
        v = v.permute(0, 2, 1, 3) # bs x n_heads x kv_len x v_head_dim
        scores = torch.matmul(q, k.transpose(-1, -2)) / math.sqrt(self.rope_head_dim + self.nope_head_dim)
        attn = F.softmax(scores, dim=-1).to(torch.bfloat16)
        y = torch.matmul(attn, v).view(batch_size, 1, -1)
        y = self.wo(y)

        return y, kv_cache.get_data()

def print_config(config: Config):
    for key in config.__dict__:
        if isinstance(config.__dict__[key], torch.Tensor):
            print(f"{key}: {config.__dict__[key].shape}")
        else:
            print(f"{key}: {config.__dict__[key]}")

def torch_rope(
    theta: torch.Tensor,
    x: torch.Tensor,
    start_pos: int = 0,
):
    # x: (-1, seq_len, d_model)
    seq_len = x.size(-2)
    d_model = x.size(-1)
    old_shape = x.shape
    x = x.reshape(-1, seq_len, d_model)
    seq_idx = torch.arange(start_pos, start_pos + seq_len, device=x.device) # (seq_len,)
    idx_theta = torch.einsum('s,d->sd', seq_idx, theta) # (seq_len, d_model//2)
    
    cos = idx_theta.cos()
    sin = idx_theta.sin()
    x0 = x[:, :, :d_model//2] # (-1, seq_len, d_model//2)
    x1 = x[:, :, d_model//2:] # (-1, seq_len, d_model//2)
    new_x0 = x0 * cos - x1 * sin # (-1, seq_len, d_model//2)
    new_x1 = x0 * sin + x1 * cos # (-1, seq_len, d_model//2)
    new_x = torch.cat([new_x0, new_x1], dim=-1) # (-1, seq_len, d_model)
    new_x = new_x.view(old_shape) # (batch_size, seq_len, d_model)
    return new_x
@triton.jit
def rope_kernel(
    x,
    y,
    theta,
    start_position,
    batchsize,
    seq_len,
    n_heads,
    rope_head_dim: tl.constexpr,
    x_batch_stride,
    x_seq_stride,
    x_head_stride,
    y_batch_stride,
    y_seq_stride,
    y_head_stride,
    head_block_size: tl.constexpr,
):
    # x:  (batchsize, seq_len, n_heads, rope_head_dim) 
    # y:  (batchsize, seq_len, n_heads, rope_head_dim)
    # theta: (rope_head_dim // 2, )
    # grid: (cdiv(n_heads, head_block_size), seq_len, batchsize)
    
    batch_idx = tl.program_id(2)
    seq_idx = tl.program_id(1)
    head_begin = tl.program_id(0) * head_block_size
    
    x0_block_ptr = tl.make_block_ptr(
        x + batch_idx * x_batch_stride + seq_idx * x_seq_stride + head_begin * x_head_stride,
        shape=(tl.minimum(head_block_size, n_heads - head_begin), rope_head_dim//2),
        strides=(x_head_stride, 1),
        offsets=(0, 0),
        block_shape=(head_block_size, rope_head_dim//2),
        order=(1, 0),
    )
    x1_block_ptr = tl.make_block_ptr(
        x + batch_idx * x_batch_stride + seq_idx * x_seq_stride + head_begin * x_head_stride + rope_head_dim//2,
        shape=(tl.minimum(head_block_size, n_heads - head_begin), rope_head_dim//2),
        strides=(x_head_stride, 1),
        offsets=(0, 0),
        block_shape=(head_block_size, rope_head_dim//2),
        order=(1, 0),
    )
    y0_block_ptr = tl.make_block_ptr(
        y + batch_idx * y_batch_stride + seq_idx * y_seq_stride + head_begin * y_head_stride,
        shape=(tl.minimum(head_block_size, n_heads - head_begin), rope_head_dim//2),
        strides=(y_head_stride, 1),
        offsets=(0, 0),
        block_shape=(head_block_size, rope_head_dim//2),
        order=(1, 0),
    )
    y1_block_ptr = tl.make_block_ptr(
        y + batch_idx * y_batch_stride + seq_idx * y_seq_stride + head_begin * y_head_stride + rope_head_dim//2,
        shape=(tl.minimum(head_block_size, n_heads - head_begin), rope_head_dim//2),
        strides=(y_head_stride, 1),
        offsets=(0, 0),
        block_shape=(head_block_size, rope_head_dim//2),
        order=(1, 0),
    )
    theta_block_ptr = tl.make_block_ptr(
        theta,
        shape=(1, rope_head_dim//2, ),
        strides=(0, 1),
        offsets=(0, 0),
        block_shape=(1, rope_head_dim//2),
        order=(1, 0),
    )
    x0 = tl.load(x0_block_ptr, boundary_check=(0, 1)) # (head_block_size, rope_head_dim//2)
    x1 = tl.load(x1_block_ptr, boundary_check=(0, 1)) # (head_block_size, rope_head_dim//2)
    theta = tl.load(theta_block_ptr) # (1, rope_head_dim//2)
    position = (start_position + seq_idx).to(tl.int64) # (1, )
    # precision!!!!
    theta = (theta * position) # (1, rope_head_dim//2)
    cos = tl.cos(theta).to(tl.bfloat16) # (1, rope_head_dim//2)
    sin = tl.sin(theta).to(tl.bfloat16) # (1, rope_head_dim//2)
    
    y0 = (x0 * cos) - (x1 * sin) # (head_block_size, rope_head_dim//2)
    y1 = (x0 * sin) + (x1 * cos) # (head_block_size, rope_head_dim//2)
    # convert to the type of y
    y0 = y0.to(y.dtype.element_ty)
    y1 = y1.to(y.dtype.element_ty)
    tl.store(y0_block_ptr, y0, boundary_check=(0, 1)) # (head_block_size, rope_head_dim//2)
    tl.store(y1_block_ptr, y1, boundary_check=(0, 1)) # (head_block_size, rope_head_dim//2)


thetas = {
    d_model: 10000 ** (-torch.arange(0, d_model//2,dtype=torch.bfloat16) / (d_model//2)) for d_model in [64,]
}
thetas = {
    d_model: theta.cuda() for d_model, theta in thetas.items()
}
def mla_decode_v0(
    config: Config,
    x: torch.Tensor,
    kv_cache: KVCache,
) -> Tuple[torch.Tensor, KVCache]:
    # dim:                7168
    # n_heads:             128
    # q_lora_rank:        1536
    # kv_lora_rank:       512
    # qk_nope_head_dim:   128
    # qk_rope_head_dim:   64
    # v_head_dim:        128
    # Q_proj_down_weight:   (1536, 7168)
    # Q_proj_up_weight:     ((128 + 64) * 128, 1536) = (24576, 1536)
    # KV_proj_down_weight:  (512 + 64, 7168) = (576, 7168)
    # KV_proj_up_weight:    ((128 + 128) * 128, 512) = (32768, 512)
    # wo_weight:            (7168, 128 * 128) = (7168, 16384)
    # kv_cache_shape:     (batchsize, max_seq_len, kv_lora_rank + qk_rope_head_dim)
    # Q_proj_down_weight:   (q_lora_rank, dim)
    # Q_proj_up_weight:     ((qk_nope_head_dim + qk_rope_head_dim) * n_heads, q_lora_rank)
    # KV_proj_down_weight:  (kv_lora_rank + qk_rope_head_dim, dim)
    # KV_proj_up_weight:    ((qk_nope_head_dim + v_head_dim) * n_heads, kv_lora_rank)
    # wo_weight:            (dim, v_head_dim * n_heads)
    # x:                (batchsize, 1, dim)
    q_lora = x @ config.Q_proj_down_weight.T # (batchsize, 1, q_lora_rank)
    cur_kv_lora = x @ config.KV_proj_down_weight.T # (batchsize, 1, kv_lora_rank + qk_rope_head_dim)
    
    # store kv_lora in cache and return the updated cache
    kv_cache_data = kv_cache.get_data() # (batchsize, max_seq_len, kv_lora_rank + qk_rope_head_dim)
    kv_cache.seq_len += 1
    kv_cache_data[:, kv_cache.seq_len - 1, :] = cur_kv_lora.squeeze(1) 
    kv_len = kv_cache.seq_len # (batchsize, kv_len, kv_lora_rank + qk_rope_head_dim)
    kv_cache_data = kv_cache.get_data() # (batchsize, max_seq_len, kv_lora_rank + qk_rope_head_dim)
    kv_lora = kv_cache_data[:, :kv_cache.seq_len, :] # (batchsize, kv_len, kv_lora_rank + qk_rope_head_dim) : (max_seq_len * (kv_lora_rank + qk_rope_head_dim), kv_lora_rank + qk_rope_head_dim ,1)
    query_pos = kv_len - 1 #
    
    q_nope_and_rope = q_lora @ config.Q_proj_up_weight.T # (batchsize, 1, (qk_nope_head_dim + qk_rope_head_dim) * n_heads)
    q_nope_and_rope = q_nope_and_rope.view(config.batch_size, 1, config.n_heads, config.qk_nope_head_dim + config.qk_rope_head_dim) # (batchsize, 1, n_heads, qk_nope_head_dim + qk_rope_head_dim)
    
    q_nope, q_rope = torch.split(q_nope_and_rope, [config.qk_nope_head_dim, config.qk_rope_head_dim], dim=-1) 
    # (batchsize, 1, n_heads, qk_nope_head_dim) : (n_heads * (qk_nope_head_dim + qk_rope_head_dim), n_heads * (qk_nope_head_dim + qk_rope_head_dim), (qk_nope_head_dim + qk_rope_head_dim), 1) 
    # (batchsize, 1, n_heads, qk_rope_head_dim) : (n_heads * (qk_nope_head_dim + qk_rope_head_dim), n_heads * (qk_nope_head_dim + qk_rope_head_dim), (qk_nope_head_dim + qk_rope_head_dim), 1)
    # print(f"q_nope shape: {q_nope.shape} stride: {q_nope.stride()} is_contiguous: {q_nope.is_contiguous()}")
    # print(f"q_rope shape: {q_rope.shape} stride: {q_rope.stride()} is_contiguous: {q_rope.is_contiguous()}")
    # TODO: inplace rope for q_rope
    # rope_module = MyRoPE(config.qk_rope_head_dim)
    
    q_rope_ref = torch_rope(
        thetas[config.qk_rope_head_dim],
        q_rope.permute(0, 2, 1, 3), # (batchsize, n_heads, 1, qk_rope_head_dim)
        start_pos=query_pos,
    ).permute(0, 2, 1, 3) # (batchsize, 1, n_heads, qk_rope_head_dim)
    q_rope_out = q_rope_ref
    KV_proj_up_weight = config.KV_proj_up_weight.view(config.n_heads, config.qk_nope_head_dim + config.v_head_dim, config.kv_lora_rank) # (n_heads, qk_nope_head_dim + v_head_dim, kv_lora_rank)
    k_nope_proj_up_weight, v_proj_up_weight = torch.split(KV_proj_up_weight, [config.qk_nope_head_dim, config.v_head_dim], dim=-2) 
    # (n_heads, qk_nope_head_dim, kv_lora_rank) : ((qk_nope_head_dim + v_head_dim) * kv_lora_rank, kv_lora_rank, 1)
    # (n_heads, v_head_dim, kv_lora_rank) : ((qk_nope_head_dim + v_head_dim) * kv_lora_rank, kv_lora_rank, 1)
    # print(f"k_nope_proj_up_weight shape: {k_nope_proj_up_weight.shape} stride: {k_nope_proj_up_weight.stride()} is_contiguous: {k_nope_proj_up_weight.is_contiguous()}")
    # print(f"v_proj_up_weight shape: {v_proj_up_weight.shape} stride: {v_proj_up_weight.stride()} is_contiguous: {v_proj_up_weight.is_contiguous()}")
    # TODO: check if contiguous is called or not
    q_nope_projed = torch.einsum('bshd,hdc->bshc', q_nope, k_nope_proj_up_weight) # (batchsize, 1, n_heads, kv_lora_rank)
    # print(f"q_nope_projed shape: {q_nope_projed.shape} stride: {q_nope_projed.stride()} is_contiguous: {q_nope_projed.is_contiguous()}")
    
    
    kv_nope, k_rope = torch.split(kv_lora, [config.kv_lora_rank, config.qk_rope_head_dim], dim=-1) 
    # (batchsize, kv_len, kv_lora_rank) : (max_seq_len * (kv_lora_rank + qk_rope_head_dim), kv_lora_rank + qk_rope_head_dim ,1)
    # (batchsize, kv_len, qk_rope_head_dim) : (max_seq_len * (kv_lora_rank + qk_rope_head_dim), kv_lora_rank + qk_rope_head_dim ,1)
    # print(f"kv_nope shape: {kv_nope.shape} stride: {kv_nope.stride()} is_contiguous: {kv_nope.is_contiguous()}")
    # print(f"k_rope shape: {k_rope.shape} stride: {k_rope.stride()} is_contiguous: {k_rope.is_contiguous()}")
    # TODO: out of place rope for k_rope
    k_rope_ref = torch_rope(
        thetas[config.qk_rope_head_dim],
        k_rope,
        start_pos=0,    
    ) # (batchsize, kv_len, qk_rope_head_dim)
    
    k_rope_out = k_rope_ref
    # print(f"k_rope dtype: {k_rope.dtype} shape: {k_rope.shape} stride: {k_rope.stride()} is_contiguous: {k_rope.is_contiguous()}")
    # TODO: MLA
    score = torch.einsum('bshc,bvc->bshv', q_nope_projed, kv_nope) # (batchsize, 1, n_heads, kv_len)
    score += torch.einsum('bshd,bvd->bshv', q_rope_out, k_rope_out) # (batchsize, 1, n_heads, kv_len)
    score = score / math.sqrt(config.qk_nope_head_dim + config.qk_rope_head_dim) # (batchsize, 1, n_heads, kv_len)
    attn = F.softmax(score, dim=-1) # (batchsize, 1, n_heads, kv_len)
    y = torch.einsum('bshv,bvc->bshc', attn, kv_nope) # (batchsize, 1, n_heads, kv_lora_rank)
    
    
    y = torch.einsum('bshc,hdc->bshd', y, v_proj_up_weight) # (batchsize, 1, n_heads, v_head_dim)
    y = torch.einsum('bshd,ohd->bso', y, config.wo_weight.view(config.dim, config.n_heads, config.v_head_dim)) # (batchsize, 1, dim)
    return y, kv_cache_data


def custom_kernel(data: input_t):
    config, x, kv_cache = data
    # print_config(config)
    out = mla_decode_v0(
        config,
        x,
        kv_cache,
    )
    torch.cuda.synchronize()
    time.sleep(0.008)
    return out
    
    # model = MLA(config)
    # 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)

    # output, kv_cache = model(x, kv_cache)
    # return output, kv_cache
    

def generate_input(batchsize, dim, dq, prefill, seed):
    # Sizes derived from: https://github.com/deepseek-ai/DeepSeek-V3/blob/main/inference/model.py
    gen = torch.Generator(device='cuda')
    gen.manual_seed(seed)
    
    # Generate weights for linear layers
    Q_proj_down_weight = torch.randn((dq, dim), dtype=torch.bfloat16, generator=gen, device='cuda')
    KV_proj_down_weight = torch.randn((512 + 64, dim), dtype=torch.bfloat16, generator=gen, device='cuda')
    Q_proj_up_weight = torch.randn(((128 + 64) * 128, dq), dtype=torch.bfloat16, generator=gen, device='cuda')
    KV_proj_up_weight = torch.randn(((128 + 128) * 128, 512), dtype=torch.bfloat16, generator=gen, device='cuda')
    wo_weight = torch.randn((dim, 128 * 128), dtype=torch.bfloat16, generator=gen, device='cuda')

    config = Config(
        batch_size=batchsize,
        dim=dim,
        q_lora_rank=dq,
        n_heads=128,
        kv_lora_rank=512,
        qk_nope_head_dim=128,
        qk_rope_head_dim=64,
        v_head_dim=128,
        seq_len=1,
        max_seq_len=8192,
        kv_cache_shape=(batchsize, 8192, 512 + 64),
        Q_proj_down_weight=Q_proj_down_weight,
        Q_proj_up_weight=Q_proj_up_weight,
        KV_proj_down_weight=KV_proj_down_weight,
        KV_proj_up_weight=KV_proj_up_weight,
        wo_weight=wo_weight,
    )
    x = torch.randn((config.batch_size, 1, config.dim), dtype=torch.bfloat16, generator=gen, device='cuda')
    
    # Pre-fill KV cache
    kv_cache = KVCache((config.batch_size, config.max_seq_len, config.kv_lora_rank + config.qk_rope_head_dim)).to('cuda')
    pre_filled_cache = torch.randn((config.batch_size, prefill, config.kv_lora_rank + config.qk_rope_head_dim), 
                                 dtype=torch.bfloat16, generator=gen, device='cuda')
    kv_cache(pre_filled_cache)

    return config, x, kv_cache


if __name__ == "__main__":
    # Example usage
    batchsize = 128
    dim = 7168
    dq = 1536
    prefill = 4096
    seed = 9817
    with torch.profiler.profile(
        activities=[torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA],
        # schedule=torch.profiler.schedule(wait=1, warmup=1, active=3),  # 控制分析步骤
        record_shapes=True,
        profile_memory=True,
        with_stack=True,
    ) as prof:
        for i in range(2):
            config, x, kv_cache = generate_input(batchsize, dim, dq, prefill, seed)
            custom_kernel((config, x, kv_cache))
    prof.export_chrome_trace("mla_trace.json")
    # print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=5))
python eval.py leaderboard /tmp/tmpi42ajl8w
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])
