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

@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 KVCache(nn.Module):
    def __init__(self, kv_cache_shape: tuple, **kwargs) -> None:
        super().__init__(**kwargs)
        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

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 q_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) # (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)


@triton.jit
def q_rope_cos_sin_kernel(
    x,
    y,
    cos,
    sin,
    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)
    # cos: (max_seq_len, rope_head_dim // 2)
    # sin: (max_seq_len, rope_head_dim // 2)
    # start_position: (1, )
    # 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),
    )
    cos_block_ptr = tl.make_block_ptr(
        cos + (start_position + seq_idx) * (rope_head_dim//2),
        shape=(1, rope_head_dim//2),
        strides=(0, 1),
        offsets=(0, 0),
        block_shape=(1, rope_head_dim//2),
        order=(1, 0),
    )
    sin_block_ptr = tl.make_block_ptr(
        sin + (start_position + seq_idx) * (rope_head_dim//2),
        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)
    cos = tl.load(cos_block_ptr) # (1, rope_head_dim//2)
    sin = tl.load(sin_block_ptr) # (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)

@triton.jit
def k_rope_cos_sin_kernel(
    x,
    y,
    cos,
    sin,
    batchsize,
    seq_len,
    rope_head_dim: tl.constexpr,
    x_batch_stride,
    x_seq_stride,
    y_batch_stride,
    y_seq_stride,
    seq_block_size: tl.constexpr,
):
    # x:  (batchsize, seq_len, rope_head_dim)
    # y:  (batchsize, seq_len, rope_head_dim)
    # cos: (max_seq_len, rope_head_dim // 2)
    # sin: (max_seq_len, rope_head_dim // 2)
    # grid: (cdiv(seq_len, seq_block_size), batchsize)
    batch_idx = tl.program_id(1)
    seq_begin = tl.program_id(0) * seq_block_size
    
    x0_block_ptr = tl.make_block_ptr(
        x + batch_idx * x_batch_stride + seq_begin * x_seq_stride,
        shape=(tl.minimum(seq_len - seq_begin, seq_block_size), rope_head_dim//2),
        strides=(x_seq_stride, 1),
        offsets=(0, 0),
        block_shape=(seq_block_size, rope_head_dim//2),
        order=(1, 0),
    )
    x1_block_ptr = tl.make_block_ptr(
        x + batch_idx * x_batch_stride + seq_begin * x_seq_stride + rope_head_dim//2,
        shape=(tl.minimum(seq_len - seq_begin, seq_block_size), rope_head_dim//2),
        strides=(x_seq_stride, 1),
        offsets=(0, 0),
        block_shape=(seq_block_size, rope_head_dim//2),
        order=(1, 0),
    )
    y0_block_ptr = tl.make_block_ptr(
        y + batch_idx * y_batch_stride + seq_begin * y_seq_stride,
        shape=(tl.minimum(seq_len - seq_begin, seq_block_size), rope_head_dim//2),
        strides=(y_seq_stride, 1),
        offsets=(0, 0),
        block_shape=(seq_block_size, rope_head_dim//2),
        order=(1, 0),
    )
    y1_block_ptr = tl.make_block_ptr(
        y + batch_idx * y_batch_stride + seq_begin * y_seq_stride + rope_head_dim//2,
        shape=(tl.minimum(seq_len - seq_begin, seq_block_size), rope_head_dim//2),
        strides=(y_seq_stride, 1),
        offsets=(0, 0),
        block_shape=(seq_block_size, rope_head_dim//2),
        order=(1, 0),
    )
    cos_block_ptr = tl.make_block_ptr(
        cos + seq_begin * (rope_head_dim//2),
        shape=(seq_block_size, rope_head_dim//2),
        strides=(rope_head_dim//2, 1),
        offsets=(0, 0),
        block_shape=(seq_block_size, rope_head_dim//2),
        order=(1, 0),
    )
    sin_block_ptr = tl.make_block_ptr(
        sin + seq_begin * (rope_head_dim//2),
        shape=(seq_block_size, rope_head_dim//2),
        strides=(rope_head_dim//2, 1),
        offsets=(0, 0),
        block_shape=(seq_block_size, rope_head_dim//2),
        order=(1, 0),
    )
    x0 = tl.load(x0_block_ptr, boundary_check=(0, )) # (seq_block_size, rope_head_dim//2)
    x1 = tl.load(x1_block_ptr, boundary_check=(0, )) # (seq_block_size, rope_head_dim//2)
    cos = tl.load(cos_block_ptr) # (seq_block_size, rope_head_dim//2)
    sin = tl.load(sin_block_ptr) # (seq_block_size, rope_head_dim//2)
    y0 = (x0 * cos) - (x1 * sin) # (seq_block_size, rope_head_dim//2)
    y1 = (x0 * sin) + (x1 * cos) # (seq_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, )) # (seq_block_size, rope_head_dim//2)
    tl.store(y1_block_ptr, y1, boundary_check=(0, )) # (seq_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()
}
max_seqlen = 8192
cos_tensors = {
    d_model: 
        torch.einsum('s,d->sd', torch.arange(0, max_seqlen,).cuda(), theta).cos()
    for d_model, theta in thetas.items()
}
sin_tensors = {
    d_model: 
        torch.einsum('s,d->sd', torch.arange(0, max_seqlen,).cuda(), theta).sin()
    for d_model, theta in thetas.items()
}
torch.cuda.synchronize()
def mla_decode_v0(
    config: Config,
    x: torch.Tensor,
    kv_cache: 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_out = torch.empty(size=q_rope.shape, dtype=q_rope.dtype, device=q_rope.device) # (batchsize, 1, n_heads, qk_rope_head_dim)
    head_block_size = 32
    q_rope_cos_sin_kernel[triton.cdiv(config.n_heads, head_block_size), 1, config.batch_size](
        x = q_rope,
        y = q_rope_out,
        cos = cos_tensors[config.qk_rope_head_dim],
        sin = sin_tensors[config.qk_rope_head_dim],
        start_position=query_pos,
        batchsize=config.batch_size,
        seq_len=1,
        n_heads=config.n_heads,
        rope_head_dim=config.qk_rope_head_dim,
        x_batch_stride=q_rope.stride(0),
        x_seq_stride=q_rope.stride(1),
        x_head_stride=q_rope.stride(2),
        y_batch_stride=q_rope_out.stride(0),
        y_seq_stride=q_rope_out.stride(1),
        y_head_stride=q_rope_out.stride(2),
        head_block_size=head_block_size,
    )
    # 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_out = torch.empty(size=k_rope.shape, dtype=k_rope.dtype, device=k_rope.device) # (batchsize, kv_len, qk_rope_head_dim)
    seq_block_size = 16
    k_rope_cos_sin_kernel[triton.cdiv(kv_len, seq_block_size), config.batch_size](
        x = k_rope,
        y = k_rope_out,
        cos = cos_tensors[config.qk_rope_head_dim],
        sin = sin_tensors[config.qk_rope_head_dim],
        batchsize=config.batch_size,
        seq_len=kv_len,
        rope_head_dim=config.qk_rope_head_dim,
        x_batch_stride=k_rope.stride(0),
        x_seq_stride=k_rope.stride(1),
        y_batch_stride=k_rope_out.stride(0),
        y_seq_stride=k_rope_out.stride(1),
        seq_block_size=seq_block_size,
        num_warps=8, # TODO: check if this is needed
        num_stages=1, # TODO: check if this is needed
    )
    # k_rope_out = torch_rope(thetas[config.qk_rope_head_dim],k_rope,0)
        
    # 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
    out = mla_decode_v0(
        config,
        x,
        kv_cache,
    )
    torch.cuda.synchronize()
    time.sleep(0.008)
    return out
    


python eval.py test /tmp/tmp4bjzy1id
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])
