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,
    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()
}
class Timer:
    def __init__(self, label: str):
        self.label = label
        self.start_time = None
        self.end_time = None
        self.interval = None
    def __enter__(self):
        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
        print(f"{self.label}: {self.interval:.4f} ms")
        return False
    
@triton.jit
def mla_decode_kernel(
    q_nope, kv_nope, 
    q_rope, k_rope,
    batch_size,
    kv_len,
    n_heads: tl.constexpr,
    qk_nope_head_dim: tl.constexpr,
    qk_rope_head_dim: tl.constexpr,
    kv_lora_rank: tl.constexpr,
):
    # q_nope: (batch_size, n_heads, kv_lora_rank)
    # kv_nope: (batch_size, kv_len, kv_lora_rank)
    # q_rope: (batch_size, n_heads, qk_rope_head_dim)
    # k_rope: (batch_size, kv_len, qk_rope_head_dim)
    ...
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)
    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)
        # 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
    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)
    # 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
    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)
        #print(f"q_nope_projed shape: {q_nope_projed.shape} stride: {q_nope_projed.stride()} is_contiguous: {q_nope_projed.is_contiguous()}")
        
        # 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)
    # 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
    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"):
        # TODO: MLA
        print(f"q_nope_projed is contiguous: {q_nope_projed.is_contiguous()} shape: {q_nope_projed.shape} stride: {q_nope_projed.stride()}")
        print(f"kv_nope is contiguous: {kv_nope.is_contiguous()} shape: {kv_nope.shape} stride: {kv_nope.stride()}")
        print(f"q_rope_out is contiguous: {q_rope_out.is_contiguous()}")
        print(f"k_rope_out is contiguous: {k_rope_out.is_contiguous()}")
        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)
    
    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
    

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') / math.sqrt(dim)
    KV_proj_down_weight = torch.randn((512 + 64, dim), dtype=torch.bfloat16, generator=gen, device='cuda') / math.sqrt(dim)
    Q_proj_up_weight = torch.randn(((128 + 64) * 128, dq), dtype=torch.bfloat16, generator=gen, device='cuda') / math.sqrt(dq)
    KV_proj_up_weight = torch.randn(((128 + 128) * 128, 512), dtype=torch.bfloat16, generator=gen, device='cuda') / math.sqrt(512)
    wo_weight = torch.randn((dim, 128 * 128), dtype=torch.bfloat16, generator=gen, device='cuda') / math.sqrt(128 * 128)

    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

# warm up
config, x, kv_cache = generate_input(
    batchsize=128,
    dim=7168,
    dq=1536,
    prefill=4096,
    seed=42,
)
custom_kernel((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(100):
            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 benchmark /tmp/tmprx153vi_
kv length: 4097 down projection: 2000.0336 ms
kv length: 4097 store kv_lora: 0.2103 ms
kv length: 4097 up projection: 0.1719 ms
kv length: 4097 q_rope triton kernel: 371.1481 ms
kv length: 4097 q knope projection: 2656.3666 ms
kv length: 4097 k_rope triton kernel: 103.4415 ms
q_nope_projed is contiguous: False shape: torch.Size([128, 1, 128, 512]) stride: (512, 8388608, 65536, 1)
kv_nope is contiguous: False shape: torch.Size([128, 4097, 512]) stride: (4718592, 576, 1)
q_rope_out is contiguous: True
k_rope_out is contiguous: True
kv length: 4097 mla decode: 39.0062 ms
kv length: 4097 v projection: 0.4017 ms
kv length: 4097 down projection: 0.7653 ms
kv length: 4097 store kv_lora: 0.0789 ms
kv length: 4097 up projection: 0.0923 ms
kv length: 4097 q_rope triton kernel: 0.1445 ms
kv length: 4097 q knope projection: 0.0820 ms
kv length: 4097 k_rope triton kernel: 0.1333 ms
q_nope_projed is contiguous: False shape: torch.Size([128, 1, 128, 512]) stride: (512, 8388608, 65536, 1)
kv_nope is contiguous: False shape: torch.Size([128, 4097, 512]) stride: (4718592, 576, 1)
q_rope_out is contiguous: True
k_rope_out is contiguous: True
kv length: 4097 mla decode: 1.0538 ms
kv length: 4097 v projection: 0.2279 ms
kv length: 4098 down projection: 0.0823 ms
kv length: 4098 store kv_lora: 0.0618 ms
kv length: 4098 up projection: 0.0744 ms
kv length: 4098 q_rope triton kernel: 101.9156 ms
kv length: 4098 q knope projection: 0.1385 ms
kv length: 4098 k_rope triton kernel: 0.1643 ms
q_nope_projed is contiguous: False shape: torch.Size([128, 1, 128, 512]) stride: (512, 8388608, 65536, 1)
kv_nope is contiguous: False shape: torch.Size([128, 4098, 512]) stride: (4718592, 576, 1)
q_rope_out is contiguous: True
k_rope_out is contiguous: True
kv length: 4098 mla decode: 1.1072 ms
kv length: 4098 v projection: 0.2463 ms
kv length: 4099 down projection: 0.0887 ms
kv length: 4099 store kv_lora: 0.0744 ms
kv length: 4099 up projection: 0.0803 ms
kv length: 4099 q_rope triton kernel: 0.1035 ms
kv length: 4099 q knope projection: 0.0620 ms
kv length: 4099 k_rope triton kernel: 0.1204 ms
q_nope_projed is contiguous: False shape: torch.Size([128, 1, 128, 512]) stride: (512, 8388608, 65536, 1)
kv_nope is contiguous: False shape: torch.Size([128, 4099, 512]) stride: (4718592, 576, 1)
q_rope_out is contiguous: True
k_rope_out is contiguous: True
kv length: 4099 mla decode: 1.0273 ms
kv length: 4099 v projection: 0.2112 ms
kv length: 4100 down projection: 0.0818 ms
kv length: 4100 store kv_lora: 0.0684 ms
kv length: 4100 up projection: 0.0739 ms
kv length: 4100 q_rope triton kernel: 0.0954 ms
kv length: 4100 q knope projection: 0.0591 ms
kv length: 4100 k_rope triton kernel: 0.1202 ms
q_nope_projed is contiguous: False shape: torch.Size([128, 1, 128, 512]) stride: (512, 8388608, 65536, 1)
kv_nope is contiguous: False shape: torch.Size([128, 4100, 512]) stride: (4718592, 576, 1)
q_rope_out is contiguous: True
k_rope_out is contiguous: True
kv length: 4100 mla decode: 1.0145 ms
kv length: 4100 v projection: 0.2041 ms
kv length: 4101 down projection: 0.0751 ms
kv length: 4101 store kv_lora: 0.0565 ms
kv length: 4101 up projection: 0.0713 ms
kv length: 4101 q_rope triton kernel: 0.0865 ms
kv length: 4101 q knope projection: 0.0567 ms
kv length: 4101 k_rope triton kernel: 0.1140 ms
q_nope_projed is contiguous: False shape: torch.Size([128, 1, 128, 512]) stride: (512, 8388608, 65536, 1)
kv_nope is contiguous: False shape: torch.Size([128, 4101, 512]) stride: (4718592, 576, 1)
q_rope_out is contiguous: True
k_rope_out is contiguous: True
kv length: 4101 mla decode: 1.0242 ms
kv length: 4101 v projection: 0.2077 ms
kv length: 4102 down projection: 0.0758 ms
kv length: 4102 store kv_lora: 0.0587 ms
kv length: 4102 up projection: 0.0732 ms
kv length: 4102 q_rope triton kernel: 0.0887 ms
kv length: 4102 q knope projection: 0.0589 ms
kv length: 4102 k_rope triton kernel: 0.1252 ms
q_nope_projed is contiguous: False shape: torch.Size([128, 1, 128, 512]) stride: (512, 8388608, 65536, 1)
kv_nope is contiguous: False shape: torch.Size([128, 4102, 512]) stride: (4718592, 576, 1)
q_rope_out is contiguous: True
k_rope_out is contiguous: True
kv length: 4102 mla decode: 1.0128 ms
kv length: 4102 v projection: 0.2053 ms
kv length: 4103 down projection: 0.0753 ms
kv length: 4103 store kv_lora: 0.0589 ms
kv length: 4103 up projection: 0.0722 ms
kv length: 4103 q_rope triton kernel: 0.0885 ms
kv length: 4103 q knope projection: 0.0582 ms
kv length: 4103 k_rope triton kernel: 0.1135 ms
q_nope_projed is contiguous: False shape: torch.Size([128, 1, 128, 512]) stride: (512, 8388608, 65536, 1)
kv_nope is contiguous: False shape: torch.Size([128, 4103, 512]) stride: (4718592, 576, 1)
q_rope_out is contiguous: True
k_rope_out is contiguous: True
kv length: 4103 mla decode: 1.0250 ms
kv length: 4103 v projection: 0.2053 ms
kv length: 4104 down projection: 0.0761 ms
kv length: 4104 store kv_lora: 0.0591 ms
kv length: 4104 up projection: 0.0722 ms
kv length: 4104 q_rope triton kernel: 0.0844 ms
kv length: 4104 q knope projection: 0.0591 ms
kv length: 4104 k_rope triton kernel: 0.1197 ms
q_nope_projed is contiguous: False shape: torch.Size([128, 1, 128, 512]) stride: (512, 8388608, 65536, 1)
kv_nope is contiguous: False shape: torch.Size([128, 4104, 512]) stride: (4718592, 576, 1)
q_rope_out is contiguous: True
k_rope_out is contiguous: True
kv length: 4104 mla decode: 0.9272 ms
kv length: 4104 v projection: 0.2053 ms
kv length: 4105 down projection: 0.0744 ms
kv length: 4105 store kv_lora: 0.0589 ms
kv length: 4105 up projection: 0.0732 ms
kv length: 4105 q_rope triton kernel: 0.0927 ms
kv length: 4105 q knope projection: 0.0606 ms
kv length: 4105 k_rope triton kernel: 0.1163 ms
q_nope_projed is contiguous: False shape: torch.Size([128, 1, 128, 512]) stride: (512, 8388608, 65536, 1)
kv_nope is contiguous: False shape: torch.Size([128, 4105, 512]) stride: (4718592, 576, 1)
q_rope_out is contiguous: True
k_rope_out is contiguous: True
kv length: 4105 mla decode: 1.0197 ms
kv length: 4105 v projection: 0.2086 ms
kv length: 4106 down projection: 0.0751 ms
kv length: 4106 store kv_lora: 0.0560 ms
kv length: 4106 up projection: 0.0722 ms
kv length: 4106 q_rope triton kernel: 0.0868 ms
kv length: 4106 q knope projection: 0.0584 ms
kv length: 4106 k_rope triton kernel: 0.1140 ms
q_nope_projed is contiguous: False shape: torch.Size([128, 1, 128, 512]) stride: (512, 8388608, 65536, 1)
kv_nope is contiguous: False shape: torch.Size([128, 4106, 512]) stride: (4718592, 576, 1)
q_rope_out is contiguous: True
k_rope_out is contiguous: True
kv length: 4106 mla decode: 1.0879 ms
kv length: 4106 v projection: 0.2193 ms
kv length: 4107 down projection: 0.0777 ms
kv length: 4107 store kv_lora: 0.0603 ms
kv length: 4107 up projection: 0.0732 ms
kv length: 4107 q_rope triton kernel: 0.0954 ms
kv length: 4107 q knope projection: 0.0596 ms
kv length: 4107 k_rope triton kernel: 0.1142 ms
q_nope_projed is contiguous: False shape: torch.Size([128, 1, 128, 512]) stride: (512, 8388608, 65536, 1)
kv_nope is contiguous: False shape: torch.Size([128, 4107, 512]) stride: (4718592, 576, 1)
q_rope_out is contiguous: True
k_rope_out is contiguous: True
kv length: 4107 mla decode: 1.0455 ms
kv length: 4107 v projection: 0.2079 ms
kv length: 4108 down projection: 0.0756 ms
kv length: 4108 store kv_lora: 0.0613 ms
kv length: 4108 up projection: 0.0727 ms
kv length: 4108 q_rope triton kernel: 0.0873 ms
kv length: 4108 q knope projection: 0.0584 ms
kv length: 4108 k_rope triton kernel: 0.1206 ms
q_nope_projed is contiguous: False shape: torch.Size([128, 1, 128, 512]) stride: (512, 8388608, 65536, 1)
kv_nope is contiguous: False shape: torch.Size([128, 4108, 512]) stride: (4718592, 576, 1)
q_rope_out is contiguous: True
k_rope_out is contiguous: True
kv length: 4108 mla decode: 0.9918 ms
kv length: 4108 v projection: 0.1991 ms
kv length: 4109 down projection: 0.0753 ms
kv length: 4109 store kv_lora: 0.0548 ms
kv length: 4109 up projection: 0.0727 ms
kv length: 4109 q_rope triton kernel: 0.0861 ms
kv length: 4109 q knope projection: 0.0577 ms
kv length: 4109 k_rope triton kernel: 0.1125 ms
q_nope_projed is contiguous: False shape: torch.Size([128, 1, 128, 512]) stride: (512, 8388608, 65536, 1)
kv_nope is contiguous: False shape: torch.Size([128, 4109, 512]) stride: (4718592, 576, 1)
q_rope_out is contiguous: True
k_rope_out is contiguous: True
kv length: 4109 mla decode: 1.0469 ms
kv length: 4109 v projection: 0.2112 ms
kv length: 4110 down projection: 0.0768 ms
kv length: 4110 store kv_lora: 0.0589 ms
kv length: 4110 up projection: 0.0734 ms
kv length: 4110 q_rope triton kernel: 0.0966 ms
kv length: 4110 q knope projection: 0.0577 ms
kv length: 4110 k_rope triton kernel: 0.1156 ms
q_nope_projed is contiguous: False shape: torch.Size([128, 1, 128, 512]) stride: (512, 8388608, 65536, 1)
kv_nope is contiguous: False shape: torch.Size([128, 4110, 512]) stride: (4718592, 576, 1)
q_rope_out is contiguous: True
k_rope_out is contiguous: True
kv length: 4110 mla decode: 1.0238 ms
kv length: 4110 v projection: 0.2010 ms
kv length: 4111 down projection: 0.0749 ms
kv length: 4111 store kv_lora: 0.0558 ms
kv length: 4111 up projection: 0.0718 ms
kv length: 4111 q_rope triton kernel: 0.0875 ms
kv length: 4111 q knope projection: 0.0577 ms
kv length: 4111 k_rope triton kernel: 0.1125 ms
q_nope_projed is contiguous: False shape: torch.Size([128, 1, 128, 512]) stride: (512, 8388608, 65536, 1)
kv_nope is contiguous: False shape: torch.Size([128, 4111, 512]) stride: (4718592, 576, 1)
q_rope_out is contiguous: True
k_rope_out is contiguous: True
kv length: 4111 mla decode: 1.0259 ms
kv length: 4111 v projection: 0.2050 ms
kv length: 4112 down projection: 0.0756 ms
kv length: 4112 store kv_lora: 0.0548 ms
kv length: 4112 up projection: 0.0722 ms
kv length: 4112 q_rope triton kernel: 0.0863 ms
kv length: 4112 q knope projection: 0.0584 ms
kv length: 4112 k_rope triton kernel: 102.8957 ms
q_nope_projed is contiguous: False shape: torch.Size([128, 1, 128, 512]) stride: (512, 8388608, 65536, 1)
kv_nope is contiguous: False shape: torch.Size([128, 4112, 512]) stride: (4718592, 576, 1)
q_rope_out is contiguous: True
k_rope_out is contiguous: True
kv length: 4112 mla decode: 1.0653 ms
kv length: 4112 v projection: 0.2487 ms
Buff name: data, shape: torch.Size([128, 8192, 576])
kv length: 4097 down projection: 0.2823 ms
kv length: 4097 store kv_lora: 0.0968 ms
kv length: 4097 up projection: 0.0892 ms
kv length: 4097 q_rope triton kernel: 0.1538 ms
kv length: 4097 q knope projection: 0.0854 ms
kv length: 4097 k_rope triton kernel: 0.1304 ms
q_nope_projed is contiguous: False shape: torch.Size([128, 1, 128, 512]) stride: (512, 8388608, 65536, 1)
kv_nope is contiguous: False shape: torch.Size([128, 4097, 512]) stride: (4718592, 576, 1)
q_rope_out is contiguous: True
k_rope_out is contiguous: True
kv length: 4097 mla decode: 1.0366 ms
kv length: 4097 v projection: 0.2251 ms
kv length: 4098 down projection: 0.1903 ms
kv length: 4098 store kv_lora: 0.0994 ms
kv length: 4098 up projection: 0.0861 ms
kv length: 4098 q_rope triton kernel: 0.1640 ms
kv length: 4098 q knope projection: 0.0799 ms
kv length: 4098 k_rope triton kernel: 0.1323 ms
q_nope_projed is contiguous: False shape: torch.Size([128, 1, 128, 512]) stride: (512, 8388608, 65536, 1)
kv_nope is contiguous: False shape: torch.Size([128, 4098, 512]) stride: (4718592, 576, 1)
q_rope_out is contiguous: True
k_rope_out is contiguous: True
kv length: 4098 mla decode: 1.0512 ms
kv length: 4098 v projection: 0.2270 ms
kv length: 4099 down projection: 0.0823 ms
kv length: 4099 store kv_lora: 0.0618 ms
kv length: 4099 up projection: 0.0856 ms
kv length: 4099 q_rope triton kernel: 0.0966 ms
kv length: 4099 q knope projection: 0.0613 ms
kv length: 4099 k_rope triton kernel: 0.1173 ms
q_nope_projed is contiguous: False shape: torch.Size([128, 1, 128, 512]) stride: (512, 8388608, 65536, 1)
kv_nope is contiguous: False shape: torch.Size([128, 4099, 512]) stride: (4718592, 576, 1)
q_rope_out is contiguous: True
k_rope_out is contiguous: True
kv length: 4099 mla decode: 0.9987 ms
kv length: 4099 v projection: 0.2015 ms
kv length: 4100 down projection: 0.0763 ms
kv length: 4100 store kv_lora: 0.0560 ms
kv length: 4100 up projection: 0.0722 ms
kv length: 4100 q_rope triton kernel: 0.0887 ms
kv length: 4100 q knope projection: 0.0577 ms
kv length: 4100 k_rope triton kernel: 0.1173 ms
q_nope_projed is contiguous: False shape: torch.Size([128, 1, 128, 512]) stride: (512, 8388608, 65536, 1)
kv_nope is contiguous: False shape: torch.Size([128, 4100, 512]) stride: (4718592, 576, 1)
q_rope_out is contiguous: True
k_rope_out is contiguous: True
kv length: 4100 mla decode: 0.9823 ms
kv length: 4100 v projection: 0.2024 ms
kv length: 4101 down projection: 0.0765 ms
kv length: 4101 store kv_lora: 0.0565 ms
kv length: 4101 up projection: 0.0732 ms
kv length: 4101 q_rope triton kernel: 0.0949 ms
kv length: 4101 q knope projection: 0.0610 ms
kv length: 4101 k_rope triton kernel: 0.1175 ms
q_nope_projed is contiguous: False shape: torch.Size([128, 1, 128, 512]) stride: (512, 8388608, 65536, 1)
kv_nope is contiguous: False shape: torch.Size([128, 4101, 512]) stride: (4718592, 576, 1)
q_rope_out is contiguous: True
k_rope_out is contiguous: True
kv length: 4101 mla decode: 1.0087 ms
kv length: 4101 v projection: 0.1996 ms
kv length: 4102 down projection: 0.0753 ms
kv length: 4102 store kv_lora: 0.0565 ms
kv length: 4102 up projection: 0.0718 ms
kv length: 4102 q_rope triton kernel: 0.0913 ms
kv length: 4102 q knope projection: 0.0582 ms
kv length: 4102 k_rope triton kernel: 0.1183 ms
q_nope_projed is contiguous: False shape: torch.Size([128, 1, 128, 512]) stride: (512, 8388608, 65536, 1)
kv_nope is contiguous: False shape: torch.Size([128, 4102, 512]) stride: (4718592, 576, 1)
q_rope_out is contiguous: True
k_rope_out is contiguous: True
kv length: 4102 mla decode: 0.9990 ms
kv length: 4102 v projection: 0.2034 ms
kv length: 4103 down projection: 0.0761 ms
kv length: 4103 store kv_lora: 0.0539 ms
kv length: 4103 up projection: 0.0727 ms
kv length: 4103 q_rope triton kernel: 0.0846 ms
kv length: 4103 q knope projection: 0.0598 ms
kv length: 4103 k_rope triton kernel: 0.1214 ms
q_nope_projed is contiguous: False shape: torch.Size([128, 1, 128, 512]) stride: (512, 8388608, 65536, 1)
kv_nope is contiguous: False shape: torch.Size([128, 4103, 512]) stride: (4718592, 576, 1)
q_rope_out is contiguous: True
k_rope_out is contiguous: True
kv length: 4103 mla decode: 0.9825 ms
kv length: 4103 v projection: 0.2103 ms
kv length: 4104 down projection: 0.0780 ms
kv length: 4104 store kv_lora: 0.0579 ms
kv length: 4104 up projection: 0.0730 ms
kv length: 4104 q_rope triton kernel: 0.0899 ms
kv length: 4104 q knope projection: 0.0596 ms
kv length: 4104 k_rope triton kernel: 0.1230 ms
q_nope_projed is contiguous: False shape: torch.Size([128, 1, 128, 512]) stride: (512, 8388608, 65536, 1)
kv_nope is contiguous: False shape: torch.Size([128, 4104, 512]) stride: (4718592, 576, 1)
q_rope_out is contiguous: True
k_rope_out is contiguous: True
kv length: 4104 mla decode: 0.9649 ms
kv length: 4104 v projection: 0.2046 ms
kv length: 4105 down projection: 0.0772 ms
kv length: 4105 store kv_lora: 0.0558 ms
kv length: 4105 up projection: 0.0737 ms
kv length: 4105 q_rope triton kernel: 0.0856 ms
kv length: 4105 q knope projection: 0.0589 ms
kv length: 4105 k_rope triton kernel: 0.1159 ms
q_nope_projed is contiguous: False shape: torch.Size([128, 1, 128, 512]) stride: (512, 8388608, 65536, 1)
kv_nope is contiguous: False shape: torch.Size([128, 4105, 512]) stride: (4718592, 576, 1)
q_rope_out is contiguous: True
k_rope_out is contiguous: True
kv length: 4105 mla decode: 1.0219 ms
kv length: 4105 v projection: 0.2115 ms
kv length: 4106 down projection: 0.0777 ms
kv length: 4106 store kv_lora: 0.0567 ms
kv length: 4106 up projection: 0.0732 ms
kv length: 4106 q_rope triton kernel: 0.0839 ms
kv length: 4106 q knope projection: 0.0587 ms
kv length: 4106 k_rope triton kernel: 0.1154 ms
q_nope_projed is contiguous: False shape: torch.Size([128, 1, 128, 512]) stride: (512, 8388608, 65536, 1)
kv_nope is contiguous: False shape: torch.Size([128, 4106, 512]) stride: (4718592, 576, 1)
q_rope_out is contiguous: True
[...] 353 lines omitted