import math
from dataclasses import dataclass
import torch
from torch import nn
import numpy as np
import time
import torch.nn.functional as F
from utils import verbose_allclose
from task import input_t, output_t
from torch.utils.cpp_extension import load_inline
import os

if "PYTORCH_ROCM_ARCH" not in os.environ:
    os.environ["PYTORCH_ROCM_ARCH"] = "gfx942:xnack-"

kernel_cpp = r"""
#undef __HIP_NO_HALF_OPERATORS__
#undef __HIP_NO_HALF_CONVERSIONS__
#include <hip/hip_runtime.h>
#include <hip/hip_bfloat16.h>
#include <pybind11/pybind11.h>
#include <pybind11/stl.h>
#include <vector>
#include <string>
#include <iostream>
#include <iomanip>
#include <random>
#include <sstream>
#include <fstream>
#include <chrono>
#include <unistd.h>  

#include <hip/amd_detail/amd_hip_bf16.h>
#include <hip/amd_detail/amd_hip_fp8.h>

#define WARP_SIZE 64
using bf16 = hip_bfloat16;

int divUp(int x,int n){
return (x + n-1)/n;
}

template<int n_split>
__global__ void rope_kernel(bf16* rope_out, const bf16* rope, int start_pos, int seq_len, int batch_size, int stride_bs, int stride_seq) {
  constexpr int d_model = 64;//Must be equal to WARP_SIZE
  constexpr int d_model_2 = d_model/2;
  int tid = blockIdx.x*blockDim.x + threadIdx.x;

  int pos = tid % d_model;
  int seq = (tid / d_model) % seq_len;
  int bs = tid / d_model / seq_len;
  
  if (bs < batch_size){
    auto pos_2 = bf16(tid % d_model_2);
    auto theta = bf16(powf(bf16(10000.0f),bf16(-pos_2)/bf16(d_model_2)));
    auto idx_theta2 = bf16(start_pos+seq)*theta;
    auto c = bf16(cosf(idx_theta2));
    auto s = bf16(sinf(idx_theta2));

    bf16 x_r[n_split];
    #pragma unroll
    for (int split=0;split<n_split;split++){
      int index = ((batch_size/n_split)*split + bs)*stride_bs+
                  seq*stride_seq+
                  pos;
       x_r[split] = rope[index];
    }

    #pragma unroll
    for (int split=0;split<n_split;split++){
     // int index = ((batch_size/n_split)*split + bs)*stride_bs+
       //           seq*stride_seq+
         //         pos;
      
       int index = ((batch_size/n_split)*split + bs)*d_model*seq_len+
                  seq*d_model+
                  pos;
         
      auto x = x_r[split];

      int src_lane = (pos +d_model_2)% d_model;
      auto x_rot = __shfl(x, src_lane, 64);
      x_rot = (pos < d_model_2) ? -x_rot : x_rot;

      auto out  = bf16(x*c)+bf16(x_rot*s);
      rope_out[index] = out;

    }
  }
}

void compute_rope(bf16* rope_out, const bf16* rope, int start_pos, int seq_len, int batch_size,int stride_bs, int stride_seq) {
  constexpr int dim_rope = 64;
  constexpr int n_split = 32;
  dim3 threadsPerBlock(WARP_SIZE*4);
  dim3 blocksPerGrid(divUp(seq_len*dim_rope*batch_size/n_split, threadsPerBlock.x));
  hipLaunchKernelGGL((rope_kernel<n_split>), blocksPerGrid, threadsPerBlock, 0, 0, rope_out,rope, start_pos,seq_len,batch_size,stride_bs,stride_seq);
}

void compute_ropes_wrapper(
  uintptr_t k_rope_out_ptr,
  uintptr_t k_rope_ptr,
  const std::vector<int>& k_rope_stride,
  uintptr_t q_rope_out_ptr,
  uintptr_t q_rope_ptr,
  int batch_size,
  int n_heads,
  int seq_len){

  bf16 *d_k_rope_out = reinterpret_cast<bf16*>(k_rope_out_ptr);
  bf16 *d_q_rope_out = reinterpret_cast<bf16*>(q_rope_out_ptr);
  const bf16 *d_k_rope = reinterpret_cast<const bf16*>(k_rope_ptr);
  const bf16 *d_q_rope = reinterpret_cast<const bf16*>(q_rope_ptr);

  int stride_bs_k = k_rope_stride[0];
  int stride_seq_k = k_rope_stride[1];
  
  int stride_bs_q = 64;
  int stride_seq_q = 64;

  compute_rope(d_k_rope_out,d_k_rope, 0, seq_len, batch_size,stride_bs_k,stride_seq_k);
  compute_rope(d_q_rope_out,d_q_rope, seq_len-1, 1, batch_size*n_heads,stride_bs_q,stride_seq_q);
  
}

__global__ void store_to_kv_cache_kernel(bf16* kv_cache, const bf16 *kv, int batch_size, int seq_id) {
  constexpr int dim_kv = 576;
  int tid = blockIdx.x*blockDim.x + threadIdx.x;
  int bs = tid/dim_kv;
  int col = tid % dim_kv;
  int dst_row = bs*8192 + seq_id;
  int dstIndex = dst_row*dim_kv+col;
  if (bs<batch_size && dst_row<batch_size*8192){
    kv_cache[dstIndex] = kv[bs*dim_kv+col];
  }
}

void store_to_kv_cache(bf16* kv_cache, const bf16 *kv, int batch_size, int seq_id)  {
    constexpr int dim_kv = 576;
    dim3 threadsPerBlock(WARP_SIZE*4);
    dim3 blocksPerGrid(divUp(batch_size*dim_kv, threadsPerBlock.x));
    hipLaunchKernelGGL((store_to_kv_cache_kernel), blocksPerGrid, threadsPerBlock, 0, 0, kv_cache,kv, batch_size, seq_id);
}


void store_to_kv_cache_wrapper(
  uintptr_t kv_lora_ptr,
  uintptr_t kv_cache_ptr,

  int batch_size,
  int seq_id){

  const bf16 *d_kv_lora = reinterpret_cast<const bf16*>(kv_lora_ptr);
  bf16 *d_kv_cache = reinterpret_cast<bf16*>(kv_cache_ptr);

  store_to_kv_cache(d_kv_cache, d_kv_lora, batch_size, seq_id);
  
}

PYBIND11_MODULE(rope, m) {
  m.def("rope", &compute_ropes_wrapper, "Rope kernel");
  m.def("storekv", &store_to_kv_cache_wrapper, "KVCache kernel");
}
"""

        
hip_module = load_inline(
    name="rope",
    cpp_sources="",
    cuda_sources=kernel_cpp.replace('@', chr(92)),
    with_cuda=True,
    verbose=False,
    extra_cuda_cflags=["-std=c++17"],
    no_implicit_headers=True,
)


class RoPE(nn.Module):
    def __init__(self, d_model: int):
        super().__init__()
        self.d_model = d_model
        theta = 10000 ** (-torch.arange(0, d_model//2,dtype=torch.bfloat16) / (d_model//2))
        self.register_buffer("theta", theta)

    def rotate_half(self, x: torch.Tensor) -> torch.Tensor:
        x1, x2 = x.chunk(2, dim=-1)
        return torch.cat((-x2, x1), dim=-1)

    def forward(self, x: torch.Tensor, start_pos: int = 0) -> torch.Tensor:
        seq_len = x.size(-2)
        d_model = x.size(-1)
        assert d_model == self.d_model
        seq_idx = torch.arange(start_pos, start_pos + seq_len, device=x.device)
        idx_theta = torch.einsum('s,d->sd', seq_idx, self.theta)
        idx_theta2 = torch.cat([idx_theta, idx_theta], dim=-1)
        cos = idx_theta2.cos().to(torch.bfloat16)
        sin = idx_theta2.sin().to(torch.bfloat16)
        return x * cos + self.rotate_half(x) * sin


class KVCache(nn.Module):
    def __init__(self, kv_cache_shape: tuple) -> None:
        super().__init__()
        self.register_buffer('data', torch.zeros(kv_cache_shape, dtype=torch.bfloat16, device='cuda'))
        self.seq_len = 0
        self.zero()

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

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

        self.data = self.data.to(c_kv.dtype)
        self.data[
            :, self.seq_len : self.seq_len + c_kv.size(1), :
        ] = c_kv
        self.seq_len += c_kv.size(1)

        return self.data[:, :self.seq_len], self.seq_len
    
@dataclass
class Config:
    batch_size: int
    dim: int
    n_heads: int
    q_lora_rank: int 
    kv_lora_rank: int
    qk_nope_head_dim: int
    qk_rope_head_dim: int
    v_head_dim: int
    seq_len: int
    max_seq_len: int
    kv_cache_shape: tuple
    Q_proj_down_weight: torch.Tensor
    Q_proj_up_weight: torch.Tensor
    KV_proj_down_weight: torch.Tensor
    KV_proj_up_weight: torch.Tensor
    wo_weight: torch.Tensor

def split_alternate_chunks(tensor: torch.Tensor, size_a: int, size_b: int):
    """
    Splits a tensor into two tensors by alternating row chunks:
    - Takes `size_a` rows → goes to A
    - Takes `size_b` rows → goes to B
    - Repeats until the tensor is fully consumed.

    Args:
        tensor (torch.Tensor): Input tensor of shape [N, D]
        size_a (int): Number of rows per chunk for A
        size_b (int): Number of rows per chunk for B

    Returns:
        Tuple[torch.Tensor, torch.Tensor]: Tensors A and B
    """
    A_chunks = []
    B_chunks = []
    i = 0
    N = tensor.size(0)

    while i < N:
        if i + size_a <= N:
            A_chunks.append(tensor[i:i+size_a])
            i += size_a
        if i + size_b <= N:
            B_chunks.append(tensor[i:i+size_b])
            i += size_b

    A = torch.cat(A_chunks, dim=0) if A_chunks else torch.empty(0, tensor.size(1), device=tensor.device)
    B = torch.cat(B_chunks, dim=0) if B_chunks else torch.empty(0, tensor.size(1), device=tensor.device)

    return A, B


def save_tensor_raw(tensor: torch.Tensor, name: str):
    tensor = tensor.detach().cpu()

    if tensor.dtype == torch.bfloat16:
        # Convert to raw uint16 representation
        array = tensor.view(torch.uint16).numpy()
        array.tofile(f"{name}.bin")
    else:
        np_dtype = {
            torch.float16: np.float16,
            torch.float32: np.float32,
            torch.int32: np.int32,
        }.get(tensor.dtype, None)

        if np_dtype is None:
            raise TypeError(f"Unsupported tensor dtype: {tensor.dtype}")

        array = tensor.numpy().astype(np_dtype)
        array.tofile(f"{name}.bin")

    # Save shape
    np.array(tensor.shape, dtype=np.int32).tofile(f"{name}.shape")
     
     
     
class MLA(nn.Module):
    def __init__(self, config: Config):
        super().__init__()
        self.dim = config.dim
        self.n_heads = config.n_heads
        self.q_lora_rank = config.q_lora_rank
        self.kv_lora_rank = config.kv_lora_rank
        self.nope_head_dim = config.qk_nope_head_dim
        self.rope_head_dim = config.qk_rope_head_dim
        self.v_head_dim = config.v_head_dim
        # Down-projection matrices
        self.Q_proj_down = nn.Linear(self.dim, self.q_lora_rank, bias=False, dtype=torch.bfloat16)
        self.KV_proj_down = nn.Linear(self.dim, self.kv_lora_rank + self.rope_head_dim, bias=False, dtype=torch.bfloat16)

        # Up-projection and rope projection matrices
        self.Q_proj_up = nn.Linear(self.q_lora_rank, (self.nope_head_dim + self.rope_head_dim) * self.n_heads, bias=False, dtype=torch.bfloat16)
        self.KV_proj_up = nn.Linear(self.kv_lora_rank, (self.nope_head_dim + self.v_head_dim) * self.n_heads, bias=False, dtype=torch.bfloat16)

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

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

        batch_size, seq_len, model_dim = x.size()

        q_lora = self.Q_proj_down(x)
        kv_lora = self.KV_proj_down(x)

        # print("kv_lora.shape ",kv_lora.shape)        
        # print("kv_lora.stride ",kv_lora.stride())
        # print("kv_cache.shape ",kv_cache.get_data().shape)        
        # print("kv_cache.stride ",kv_cache.get_data().stride())
          
        # save_tensor_raw(kv_lora,"kv_lora")
        # save_tensor_raw(kv_cache.get_data(),"kv_cache")
        # kv_lora, kv_len = kv_cache(kv_lora)
        # print("kv_lora.stride ",kv_lora.stride())
        # save_tensor_raw(kv_cache.get_data(),"kv_cache_out")
        # print("kv_len :",kv_len)
        # print("kv_cache.seq_len=",kv_cache.seq_len)
        hip_module.storekv(
                kv_lora.data_ptr(),
                kv_cache.get_data().data_ptr(),
                batch_size,
                kv_cache.seq_len)
        
        kv_cache.seq_len+=1
        kv_len = kv_cache.seq_len
        # print("kv_len=",kv_cache.seq_len)
        
        query_pos = kv_len - 1
        kv_lora = kv_cache.get_data()[:, :kv_len, :]
        
        
        kvDown_nope, k_rope = torch.split(kv_lora, [self.kv_lora_rank, self.rope_head_dim], dim=-1)

        Wkv_up = self.KV_proj_up.weight.t().view(self.kv_lora_rank,self.n_heads,self.nope_head_dim+self.v_head_dim)
        Wk_up, Wv_up = torch.split(Wkv_up, [self.nope_head_dim, self.v_head_dim], dim=-1)
        Wk_up = Wk_up.permute(1, 2, 0) # (n_head, dim, head), then (n_head, head,dim) for matmul

        Wv_up = Wv_up.permute(1, 0, 2)  # shape: (n_head, dim, head)        
        
        
        
        WqNopeRope = self.Q_proj_up.weight.t().view(self.q_lora_rank,self.n_heads,self.nope_head_dim+self.rope_head_dim)
        Wq_Up_nope, Wq_Up_rope = torch.split(WqNopeRope, [self.nope_head_dim, self.rope_head_dim], dim=-1)
        
        # print("q_lora_rank",self.q_lora_rank)
        
        # Wq_Up_nope, Wq_Up_rope = split_alternate_chunks(self.Q_proj_up.weight, 128,64)
        # print("Wq_Up_nope.shape",Wq_Up_nope.shape)
        # print("Wq_Up_rope.shape",Wq_Up_rope.shape)
        
        
       
        
        # print("QRope.shape",QRope.shape)
              
 
        Wq_Up_nope = Wq_Up_nope.permute(1, 0, 2)  # shape: (128, 1536, 128)
        W_merged = torch.matmul(Wq_Up_nope,Wk_up) #150us
   
        Wq_Up_rope = Wq_Up_rope.reshape(self.q_lora_rank,self.n_heads*self.rope_head_dim)
        # print("q_lora.shape",q_lora.shape)
        # print("Wq_Up_rope.shape",Wq_Up_rope.shape)
        # QRope = torch.einsum("bm,mnl->bnl", q_lora.squeeze(1), Wq_Up_rope)
        QRope = torch.matmul(q_lora,Wq_Up_rope)
        
        q_down = q_lora
        kv_down_permute = kvDown_nope.permute(0,2,1)
        
        q_down = q_lora.unsqueeze(1) # (128,1,1536) -> (128,1,1,1536)
        kv_down_permute = kv_down_permute.unsqueeze(1) #(128,512,513) -> (128,1,512,513)
 
        

        
        ScoreNoRope1 = torch.matmul(q_down.permute(1,0,2,3).squeeze(2),W_merged)
        ScoreNoRope1 = ScoreNoRope1.permute(1,0,2)
        ScoreNoRope_t=ScoreNoRope1.view(128,128,512) #[bs,n_head,512])
        ScoreNoRope = torch.matmul(ScoreNoRope_t,kv_down_permute.squeeze(1))
        
        #ScoreNoRope=ScoreNoRope3.unsqueeze(2)
        #ROPES
         # print("QRope shape", QRope.shape)
        QRope = QRope.view(128, 1, 128, 64)
        q_rope = QRope.permute(0, 2, 1, 3) # bs x n_heads x seq_len x rope_head_dim
 
       
        
        # q_rope = self.q_rope(q_rope, start_pos=query_pos)
        
        # kv_nope, k_rope = torch.split(kv_lora, [self.kv_lora_rank, self.rope_head_dim], dim=-1)
        # k_rope_ref = self.k_rope(k_rope) #bs x kv_len x rope_head_dim

        # k_rope = k_rope.contiguous()
        # q_rope = q_rope.contiguous()
        # print("k_rope.shape",k_rope.shape)
        # print("k_rope.is_contiguous",k_rope.is_contiguous())
        # print("k_rope.stride",k_rope.stride())
        # print("q_rope.shape",q_rope.shape)
        # print("q_rope.is_contiguous",q_rope.is_contiguous())
        # print("q_rope.stride",q_rope.stride())
        k_rope_out = torch.empty(128, kv_len, 64,dtype=torch.bfloat16, device='cuda')
        # print(k_rope_out.stride())
        q_rope_out = torch.empty(128, 128,1, 64,dtype=torch.bfloat16, device='cuda')
        hip_module.rope(k_rope_out.data_ptr(),
                        k_rope.data_ptr(),
                        k_rope.stride(),
                        q_rope_out.data_ptr(),
                        q_rope.data_ptr(),
                        batch_size,
                        self.n_heads,
                        kv_len)
        
        
        k_rope_out = k_rope_out.unsqueeze(1)
        # q_rope_bs = QRope#q_rope
        ScoreOut = torch.baddbmm(
                        ScoreNoRope,
                        q_rope_out.squeeze(2),
                        k_rope_out.permute(0, 1, 3, 2).squeeze(1),
                        beta=1.0 / math.sqrt(self.rope_head_dim + self.nope_head_dim),
                        alpha=1.0 / math.sqrt(self.rope_head_dim + self.nope_head_dim)
                    )

        AttnOutSMax = F.softmax(ScoreOut, dim=-1).to(torch.bfloat16)
        
        kv_down = kvDown_nope.unsqueeze(1)
        y0_c = torch.matmul(AttnOutSMax.squeeze(2),kv_down.squeeze(1))
        y0_c=y0_c.unsqueeze(2)
        #re-order to process per bs
        y0_c = torch.matmul(y0_c.permute(1,0,2,3).squeeze(2),Wv_up)
        y0_c = y0_c.permute(1,0,2)
        y0_c = self.wo(y0_c.flatten(start_dim=1))

        
        return y0_c.view(128,1,7168), kv_cache.get_data()

model = None
def custom_kernel(data: input_t) -> output_t:
    with torch.no_grad():
        config, x, kv_cache = data
        global model
        if model == None:
            model = MLA(config).to('cuda:0')
            model.Q_proj_down.weight = nn.Parameter(config.Q_proj_down_weight)
            model.Q_proj_up.weight = nn.Parameter(config.Q_proj_up_weight)
            model.KV_proj_down.weight = nn.Parameter(config.KV_proj_down_weight)
            model.KV_proj_up.weight = nn.Parameter(config.KV_proj_up_weight)
            model.wo.weight = nn.Parameter(config.wo_weight)
        else:
            model.Q_proj_down.weight.set_(config.Q_proj_down_weight.untyped_storage(), 0, config.Q_proj_down_weight.size(), config.Q_proj_down_weight.stride()) 
            model.Q_proj_up.weight.set_(config.Q_proj_up_weight.untyped_storage(), 0, config.Q_proj_up_weight.size(), config.Q_proj_up_weight.stride()) 
            model.KV_proj_down.weight.set_(config.KV_proj_down_weight.untyped_storage(), 0, config.KV_proj_down_weight.size(), config.KV_proj_down_weight.stride()) 
            model.KV_proj_up.weight.set_(config.KV_proj_up_weight.untyped_storage(), 0, config.KV_proj_up_weight.size(), config.KV_proj_up_weight.stride()) 
            model.wo.weight.set_(config.wo_weight.untyped_storage(), 0, config.wo_weight.size(), config.wo_weight.stride()) 

            # model.Q_proj_down.weight.copy_(config.Q_proj_down_weight)
            # model.Q_proj_up.weight.copy_(config.Q_proj_up_weight)
            # model.KV_proj_down.weight.copy_(config.KV_proj_down_weight)
            # model.KV_proj_up.weight.copy_(config.KV_proj_up_weight)
            # model.wo.weight.copy_(config.wo_weight)

        output, kv_cache = model(x, kv_cache)
        return output, kv_cache
python eval.py test /tmp/tmp5f72gc30
[92mSuccessfully preprocessed all matching files.[0m

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])
