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
import aiter


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

@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,
    batch_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: (seq_len, cdiv(batchsize, batch_block_size))
    batch_begin = tl.program_id(1) * batch_block_size
    seq_idx = tl.program_id(0)
    
    x0_block_ptr = tl.make_block_ptr(
        x + batch_begin * x_batch_stride + seq_idx * x_seq_stride,
        shape=(tl.minimum(batch_block_size, batchsize - batch_begin), rope_head_dim // 2),
        strides=(x_batch_stride, 1),
        offsets=(0, 0),
        block_shape=(batch_block_size, rope_head_dim // 2),
        order=(1, 0),
    )
    x1_block_ptr = tl.make_block_ptr(
        x + batch_begin * x_batch_stride + seq_idx * x_seq_stride + rope_head_dim // 2,
        shape=(tl.minimum(batch_block_size, batchsize - batch_begin), rope_head_dim // 2),
        strides=(x_batch_stride, 1),
        offsets=(0, 0),
        block_shape=(batch_block_size, rope_head_dim // 2),
        order=(1, 0),
    )
    y0_block_ptr = tl.make_block_ptr(
        y + batch_begin * y_batch_stride + seq_idx * y_seq_stride,
        shape=(tl.minimum(batch_block_size, batchsize - batch_begin), rope_head_dim // 2),
        strides=(y_batch_stride, 1),
        offsets=(0, 0),
        block_shape=(batch_block_size, rope_head_dim // 2),
        order=(1, 0),
    )
    y1_block_ptr = tl.make_block_ptr(
        y + batch_begin * y_batch_stride + seq_idx * y_seq_stride + rope_head_dim // 2,
        shape=(tl.minimum(batch_block_size, batchsize - batch_begin), rope_head_dim // 2),
        strides=(y_batch_stride, 1),
        offsets=(0, 0),
        block_shape=(batch_block_size, rope_head_dim // 2),
        order=(1, 0),
    )
    cos_block_ptr = tl.make_block_ptr(
        cos + 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 + 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, )) # (batch_block_size, rope_head_dim//2)
    x1 = tl.load(x1_block_ptr, boundary_check=(0, )) # (batch_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) # (batch_block_size, rope_head_dim//2)
    y1 = (x0 * sin) + (x1 * cos) # (batch_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, )) # (batch_block_size, rope_head_dim//2)
    tl.store(y1_block_ptr, y1, boundary_check=(0, )) # (batch_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()
}
o_preallocated = torch.empty(
    size = (128, 128, 512),
    dtype=torch.bfloat16,
    device='cuda'
)
qo_indptr_preallocated = torch.arange(
    start=0, end=129, step=1, dtype=torch.int, device='cuda'
)
kv_buffer = torch.empty(
    size = (128, max_seqlen, 1, 576),
    dtype = torch.bfloat16,
    device = 'cuda'
)
kv_indptr_preallocated = {
    kv_len : torch.arange(
            start=0, end=kv_len*129, step=kv_len, dtype=torch.int, device='cuda'
    )
    for kv_len in range(1, max_seqlen)
}
kv_last_page_lens_preallocated = torch.ones(128, dtype=torch.int, device='cuda')

base = torch.arange(128, device='cuda', dtype=torch.int64) * max_seqlen    # (128,)
offsets = torch.arange(max_seqlen, device='cuda', dtype=torch.int64)              # (max_seqlen,)
full_idx = base.unsqueeze(1) + offsets.unsqueeze(0)

kv_indices_preallocated = {}
for kv_len in range(1, max_seqlen):
    kv_indices_preallocated[kv_len] = full_idx[:, :kv_len].reshape(-1).to(torch.int32)

class Timer:
    def __init__(self, label: str):
        self.label = label
        self.start_time = None
        self.end_time = None
        self.interval = None
    def __enter__(self):
        torch.cuda.synchronize()
        self.start = time.time()
        return self
    def __exit__(self, *args):
        torch.cuda.synchronize()
        self.end = time.time()
        self.interval = (self.end - self.start) * 1000 # in ms
        return False
    
@triton.jit
def mla_decode_kernel(
    q_nope, kv_nope, 
    q_rope, k_rope,
    out,
    scale,
    batch_size,
    kv_len,
    q_nope_batch_stride,
    q_nope_head_stride,
    kv_nope_batch_stride,
    kv_nope_seq_stride,
    n_heads: tl.constexpr,
    qk_rope_head_dim: tl.constexpr,
    kv_lora_rank: tl.constexpr,
    HEAD_BLOCK_SIZE: tl.constexpr = 16,
    SEQ_BLOCK_SIZE: tl.constexpr = 32,
):
    # q_nope: (batch_size, n_heads, kv_lora_rank) : (q_nope_batch_stride, q_nope_head_stride, 1)
    # kv_nope: (batch_size, kv_len, kv_lora_rank) : (kv_nope_batch_stride, kv_nope_seq_stride, 1)
    # q_rope: (batch_size, n_heads, qk_rope_head_dim) : (n_heads * qk_rope_head_dim, qk_rope_head_dim, 1)
    # k_rope: (batch_size, kv_len, qk_rope_head_dim) : (kv_len * qk_rope_head_dim, qk_rope_head_dim, 1)
    # out: (batch_size, n_heads, kv_lora_rank) : (n_heads * kv_lora_rank, kv_lora_rank, 1)
    # grid: (cdiv(n_heads, HEAD_BLOCK_SIZE), batch_size)
    head_begin = tl.program_id(0) * HEAD_BLOCK_SIZE
    batch_idx = tl.program_id(1)
    
    q_rope_head_stride = qk_rope_head_dim
    q_rope_batch_stride = n_heads * qk_rope_head_dim
    q_rope_block_ptr = tl.make_block_ptr(
        q_rope + batch_idx * q_rope_batch_stride + head_begin * q_rope_head_stride,
        shape=(tl.minimum(HEAD_BLOCK_SIZE, n_heads - head_begin), qk_rope_head_dim),
        strides=(q_rope_head_stride, 1),
        offsets=(0, 0),
        block_shape=(HEAD_BLOCK_SIZE, qk_rope_head_dim),
        order=(1, 0),
    ) # (HEAD_BLOCK_SIZE, qk_rope_head_dim)
    
    k_rope_seq_stride = qk_rope_head_dim
    k_rope_batch_stride = kv_len * qk_rope_head_dim
    k_rope_block_ptr = tl.make_block_ptr(
        k_rope + batch_idx * k_rope_batch_stride,
        shape=(kv_len, qk_rope_head_dim),
        strides=(k_rope_seq_stride, 1),
        offsets=(0, 0),
        block_shape=(SEQ_BLOCK_SIZE, qk_rope_head_dim),
        order=(1, 0),
    ) # (SEQ_BLOCK_SIZE, qk_rope_head_dim)
    
    q_nope_block_ptr = tl.make_block_ptr(
        q_nope + batch_idx * q_nope_batch_stride + head_begin * q_nope_head_stride,
        shape=(tl.minimum(HEAD_BLOCK_SIZE, n_heads - head_begin), kv_lora_rank),
        strides=(q_nope_head_stride, 1),
        offsets=(0, 0),
        block_shape=(HEAD_BLOCK_SIZE, kv_lora_rank),
        order=(1, 0),
    ) # (HEAD_BLOCK_SIZE, kv_lora_rank)
    kv_nope_block_ptr = tl.make_block_ptr(
        kv_nope + batch_idx * kv_nope_batch_stride,
        shape=(kv_len, kv_lora_rank),
        strides=(kv_nope_seq_stride, 1),
        offsets=(0, 0),
        block_shape=(SEQ_BLOCK_SIZE, kv_lora_rank),
        order=(1, 0),
    ) # (SEQ_BLOCK_SIZE, kv_lora_rank)
    
    out_block_ptr = tl.make_block_ptr(
        out + batch_idx * (n_heads * kv_lora_rank) + head_begin * kv_lora_rank,
        shape=(tl.minimum(HEAD_BLOCK_SIZE, n_heads - head_begin), kv_lora_rank),
        strides=(kv_lora_rank, 1),
        offsets=(0, 0),
        block_shape=(HEAD_BLOCK_SIZE, kv_lora_rank),
        order=(1, 0),
    ) # (HEAD_BLOCK_SIZE, kv_lora_rank)
    
    # 
    qk_scale = scale * 1.44269504
    # initialize pointer to m and l
    m_i = tl.zeros([HEAD_BLOCK_SIZE], dtype=tl.float32) - float("inf") #(HEAD_BLOCK_SIZE, )
    l_i = tl.zeros([HEAD_BLOCK_SIZE], dtype=tl.float32) + 1.0 #(HEAD_BLOCK_SIZE, )
    acc = tl.zeros([HEAD_BLOCK_SIZE, kv_lora_rank], dtype=tl.float32) # (HEAD_BLOCK_SIZE, kv_lora_rank)
    
    # TODO: remove boundary check
    q_rope = tl.load(q_rope_block_ptr, boundary_check=(0,)) # (HEAD_BLOCK_SIZE, qk_rope_head_dim)
    q_nope = tl.load(q_nope_block_ptr, boundary_check=(0,)) # (HEAD_BLOCK_SIZE, kv_lora_rank)
    for seq_begin in range(0, kv_len, SEQ_BLOCK_SIZE):
        k_rope_block = tl.load(k_rope_block_ptr, boundary_check=(0,)) # (SEQ_BLOCK_SIZE, qk_rope_head_dim)
        qk = tl.dot(q_rope, k_rope_block.T) # (HEAD_BLOCK_SIZE, SEQ_BLOCK_SIZE)

        kv_nope_block = tl.load(kv_nope_block_ptr, boundary_check=(0,)) # (SEQ_BLOCK_SIZE, kv_lora_rank)
        qk += tl.dot(q_nope, kv_nope_block.T) # (HEAD_BLOCK_SIZE, SEQ_BLOCK_SIZE)
        
        m_ij = tl.maximum(m_i, tl.max(qk, 1) * qk_scale)
        qk = qk * qk_scale - m_ij[:, None]
                
        p = tl.math.exp2(qk)
        
        # -- compute correction factor
        alpha = tl.math.exp2(m_i - m_ij)
        l_ij = tl.sum(p, 1)
        # -- update output accumulator --
        acc = acc * alpha[:, None]
        # prepare p and v for the dot
        p = p.to(tl.bfloat16) # convert p to the same type as kv_nope
        
        acc = tl.dot(p, kv_nope_block, acc) # (HEAD_BLOCK_SIZE, kv_lora_rank)
        # update m_i and l_i
        # place this at the end of the loop to reduce register pressure
        l_i = l_i * alpha + l_ij
        m_i = m_ij
        
        k_rope_block_ptr = tl.advance(k_rope_block_ptr, (SEQ_BLOCK_SIZE, 0)) 
        kv_nope_block_ptr = tl.advance(kv_nope_block_ptr, (SEQ_BLOCK_SIZE, 0)) 
    acc = acc / l_i[:, None] # (HEAD_BLOCK_SIZE, kv_lora_rank)
    tl.store(out_block_ptr, acc.to(tl.bfloat16), boundary_check=(0,)) # (HEAD_BLOCK_SIZE, kv_lora_rank)

@triton.jit
def store_kv(
    kv_nope, k_rope_out,
    kv_buffer,
    batch_size,
    kv_len,
    kv_lora_rank: tl.constexpr,
    qk_rope_head_dim: tl.constexpr,
    max_seq_len: tl.constexpr,
    kv_nope_batch_stride,
    kv_nope_kv_len_stride,
    k_rope_out_batch_stride,
    k_rope_out_kv_len_stride,
):
    # kv_nope: (batchsize, kv_len, kv_lora_rank) : (kv_nope_batch_stride, kv_nope_kv_len_stride, 1)
    # k_rope_out: (batchsize, kv_len, qk_rope_head_dim) : (k_rope_out_batch_stride, k_rope_out_kv_len_stride, 1)
    # kv_buffer: (batchsize, max_seq_len, kv_lora_rank + qk_rope_head_dim)  :
    # grid: (kv_len, batch_size)
    batch_idx = tl.program_id(1)
    seq_idx = tl.program_id(0)
    
    kv_nope_block_ptr = tl.make_block_ptr(
        kv_nope + batch_idx * kv_nope_batch_stride + seq_idx * kv_nope_kv_len_stride,
        shape=(1, kv_lora_rank),
        strides=(0, 1),
        offsets=(0, 0),
        block_shape=(1, kv_lora_rank),
        order=(1, 0),
    )
    k_rope_out_block_ptr = tl.make_block_ptr(
        k_rope_out + batch_idx * k_rope_out_batch_stride + seq_idx * k_rope_out_kv_len_stride,
        shape=(1, qk_rope_head_dim),
        strides=(0, 1),
        offsets=(0, 0),
        block_shape=(1, qk_rope_head_dim),
        order=(1, 0),
    )
    kv_buffer_nope_block_ptr = tl.make_block_ptr(
        kv_buffer + batch_idx * max_seq_len * (kv_lora_rank + qk_rope_head_dim) + seq_idx * (kv_lora_rank + qk_rope_head_dim),
        shape=(1, kv_lora_rank),
        strides=(0, 1),
        offsets=(0, 0),
        block_shape=(1, kv_lora_rank),
        order=(1, 0),
    )
    kv_buffer_rope_block_ptr = tl.make_block_ptr(
        kv_buffer + batch_idx * max_seq_len * (kv_lora_rank + qk_rope_head_dim) + seq_idx * (kv_lora_rank + qk_rope_head_dim) + kv_lora_rank,
        shape=(1, qk_rope_head_dim),
        strides=(0, 1),
        offsets=(0, 0),
        block_shape=(1, qk_rope_head_dim),
        order=(1, 0),
    )
    
    kv_nope_block = tl.load(kv_nope_block_ptr, boundary_check=(0,)) # (1, kv_lora_rank)
    k_rope_out_block = tl.load(k_rope_out_block_ptr, boundary_check=(0,)) # (1, qk_rope_head_dim)
    
    tl.store(kv_buffer_nope_block_ptr, kv_nope_block, boundary_check=(0,)) # (1, kv_lora_rank)
    tl.store(kv_buffer_rope_block_ptr, k_rope_out_block, boundary_check=(0,)) # (1, qk_rope_head_dim)
        
def mla_decode_v0(
    config,
    x: torch.Tensor,
    kv_cache: KVCache,
) :
    with Timer(f"kv length: {kv_cache.seq_len + 1} down projection"):
        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)
    
    with Timer(f"kv length: {kv_cache.seq_len + 1} store kv_lora"):
        # 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
        # TODO: use triton to write the kv_lora to the cache
        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 #
    with Timer(f"kv length: {kv_cache.seq_len} up projection"):
        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)
    # TODO: inplace rope for q_rope
    with Timer(f"kv length: {kv_cache.seq_len} q_rope triton kernel"):
        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 = 16
        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,
        )
    
    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)
    
    # TODO: check if contiguous is called or not
    with Timer(f"kv length: {kv_cache.seq_len} q knope projection"):
        # q_nope_projed = torch.einsum('bshd,hdc->bshc', q_nope, k_nope_proj_up_weight) # (batchsize, 1, n_heads, kv_lora_rank)
        
        # b1hd -> bhd
        q_nope = q_nope.squeeze(1) # (batchsize, n_heads, qk_nope_head_dim)
        # bhd -> hbd
        q_nope = q_nope.transpose(0, 1) # (n_heads, batchsize, qk_nope_head_dim)
        
        q_nope_projed = torch.bmm(q_nope, k_nope_proj_up_weight) # (n_heads, batchsize, kv_lora_rank)
        q_nope_projed = q_nope_projed.transpose(0, 1) # (batchsize, n_heads, kv_lora_rank)
        q_nope_projed = q_nope_projed.unsqueeze(1) # (batchsize, 1, n_heads, kv_lora_rank)
    
    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)
    
    # TODO: out of place rope for k_rope
    with Timer(f"kv length: {kv_cache.seq_len} k_rope triton kernel"):
        k_rope_out = torch.empty(size=k_rope.shape, dtype=k_rope.dtype, device=k_rope.device) # (batchsize, kv_len, qk_rope_head_dim)
        batch_block_size = 32
        k_rope_cos_sin_kernel[kv_len, triton.cdiv(config.batch_size, batch_block_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),
            batch_block_size=batch_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)
    
    with Timer(f"kv length: {kv_cache.seq_len} mla decode"):

        # aiter MLA
        q = torch.cat(
            [q_nope_projed, q_rope_out],
            dim=-1
        ).squeeze(1).reshape(-1, config.n_heads, config.kv_lora_rank + config.qk_rope_head_dim)
        with Timer(f"kv length: {kv_cache.seq_len} aiter store kv"):
            store_kv[(kv_len, config.batch_size)](
                kv_nope=kv_nope,
                k_rope_out=k_rope_out,
                kv_buffer=kv_buffer,
                batch_size=config.batch_size,
                kv_len=kv_len,
                kv_lora_rank=config.kv_lora_rank,
                qk_rope_head_dim=config.qk_rope_head_dim,
                max_seq_len=config.max_seq_len,
                kv_nope_batch_stride=kv_nope.stride(0),
                kv_nope_kv_len_stride=kv_nope.stride(1),
                k_rope_out_batch_stride=k_rope_out.stride(0),
                k_rope_out_kv_len_stride=k_rope_out.stride(1)
            )
        o = o_preallocated
        qo_indptr = qo_indptr_preallocated

        kv_indptr = kv_indptr_preallocated[kv_len]
        kv_indices = kv_indices_preallocated[kv_len]
        kv_last_page_lens = kv_last_page_lens_preallocated
        with Timer(f"kv length: {kv_cache.seq_len} aiter mla decode"):
            aiter.mla.mla_decode_fwd(
                q=q, kv_buffer=kv_buffer.view(-1, 1, 1, config.kv_lora_rank + config.qk_rope_head_dim), o=o, qo_indptr=qo_indptr,
                kv_indptr=kv_indptr, kv_indices=kv_indices, 
                kv_last_page_lens=kv_last_page_lens, max_seqlen_q=1,
                sm_scale=1.0 / math.sqrt(config.qk_nope_head_dim + config.qk_rope_head_dim),
            )
        y = o.unsqueeze(1)
        # diff = y.squeeze(1) - o
        
    with Timer(f"kv length: {kv_cache.seq_len} v projection"):
        y = y.squeeze(1) # (batchsize, n_heads, kv_lora_rank)
        y = y.transpose(0, 1) # (n_heads, batchsize, kv_lora_rank)
        v_proj_up_weight = v_proj_up_weight.transpose(1, 2) # (n_heads, kv_lora_rank, v_head_dim)
        y = torch.bmm(y, v_proj_up_weight) # (n_heads, batchsize, v_head_dim)
        y = y.transpose(0, 1) # (batchsize, n_heads, v_head_dim)
        y = y.reshape(config.batch_size, -1) # (batchsize, n_heads * v_head_dim)
        y = torch.matmul(y, config.wo_weight.T) # (batchsize, dim)
        y = y.view(config.batch_size, 1, config.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/tmpwpfvjdvr
[aiter] start build [module_mla_asm] under /home/runner/aiter/aiter/jit/build/module_mla_asm
[92mSuccessfully preprocessed all matching files.[0m
[aiter] finish build [module_mla_asm], cost 21.75835612s

Buff name: data, shape: torch.Size([128, 8192, 576])
[aiter] hipModuleLoad: /home/runner/aiter/hsa//mla/mla_a16w16_dec_subQ128_mqa128.co GetFunction: mla_a16w16_dec_subQ128_mqa128 Success
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])
