import logging
import os, sys
import time
from typing import Optional

os.environ["BLAS_NUM_THREADS"] = "1"
sys.path.insert(0, os.path.dirname(__file__) + "/../build")
from kt_kernel import kt_kernel_ext
from kt_kernel_ext.kvcache import ggml_type
import torch
from torch import inf, nn
from torch.nn import init
from torch_attention import apply_rotary_pos_emb, DeepseekV2RMSNorm, KDeepSeekV3Cache, DeepseekV3YarnRotaryEmbedding

logger = logging.getLogger("reader")

from gguf.gguf_reader import GGUFReader


def read_gguf_file(gguf_file_path):
    """
    Reads and prints key-value pairs and tensor information from a GGUF file in an improved format.

    Parameters:
    - gguf_file_path: Path to the GGUF file.
    """

    reader = GGUFReader(gguf_file_path)

    # List all key-value pairs in a columnized format
    # print("Key-Value Pairs:") # noqa: NP100
    # max_key_length = max(len(key) for key in reader.fields.keys())
    for key, field in reader.fields.items():
        value = field.parts[field.data[0]]
        # print(f"{key:{max_key_length}} : {value}") # noqa: NP100
    # print("----") # noqa: NP100

    # List all tensors
    # print("Tensors:") # noqa: NP100
    # tensor_info_format = "{:<30} | Shape: {:<15} | Size: {:<12} | Quantization: {}"
    # print(tensor_info_format.format("Tensor Name", "Shape", "Size", "Quantization")) # noqa: NP100
    # print("-" * 80) # noqa: NP100
    re = []
    for tensor in reader.tensors:
        shape_str = "x".join(map(str, tensor.shape))
        size_str = str(tensor.n_elements)
        quantization_str = tensor.tensor_type.name
        # print(tensor_info_format.format(tensor.name, shape_str, size_str, quantization_str)) # noqa: NP100
        re.append(tensor)
    return re


def get_torch_tensor_from_gguf(gguf_weights, name):
    return torch.from_numpy(gguf_weights[name].data).contiguous()


def get_torch_tensor_and_type_from_gguf(gguf_weights, name):
    return torch.from_numpy(gguf_weights[name].data).contiguous(), gguf_weights[name].tensor_type.name


def type_to_ggml_type(type):
    if type == "F32":
        return ggml_type.FP32
    elif type == "F16":
        return ggml_type.FP16
    elif type == "BF16":
        return ggml_type.BF16
    else:
        raise ValueError(f"Unsupported data type: {type}")


use_real_weights = True
gguf_path = "/home/bd/models/DeepSeek-R1-BF16"

seed = 42  # 你可以选择任何整数作为种子
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)

qlen = 3212
kvlen = 0


page_table = range(20)
bsz_tensors = torch.tensor([1])


page_size = 256
pages_count = 200
tp_count = 4


hidden_size = 7168
q_lora_rank = 1536
kv_lora_rank = 512
num_heads = 128
nope_size = 128
rope_size = 64

rope_theta = 10000
max_qlen = 4096
max_kvlen = 4096

max_position_embeddings = 163840


rope_scaling = {
    "beta_fast": 32,
    "beta_slow": 1,
    "factor": 40,
    "mscale": 1.0,
    "mscale_all_dim": 1.0,
    "original_max_position_embeddings": 4096,
    "type": "yarn",
}


CPUInfer = kt_kernel_ext.CPUInfer(30)
validation_iter = 100


# data_type = torch.float32
weight_type = torch.bfloat16
# weight_type = torch.float16


input_type = {
    torch.float32: torch.float32,
    torch.float16: torch.float16,
    torch.bfloat16: torch.float32,
}[weight_type]

q_a_proj = nn.Linear(hidden_size, q_lora_rank, bias=False, dtype=weight_type)
q_b_proj = nn.Linear(q_lora_rank, num_heads * (nope_size + rope_size), bias=False, dtype=weight_type)
kv_a_proj_with_mqa = nn.Linear(hidden_size, kv_lora_rank + rope_size, bias=False, dtype=weight_type)
kv_b_proj = nn.Linear(num_heads * (nope_size + nope_size), kv_lora_rank, bias=False, dtype=weight_type)
o_proj = nn.Linear(num_heads * nope_size, hidden_size, bias=False, dtype=weight_type)
q_a_norm = torch.ones(hidden_size, dtype=torch.float32)
kv_a_norm = torch.ones(hidden_size, dtype=torch.float32)


def read_gguf_directory(directory):
    """
    Reads all GGUF files in a directory and prints their contents.

    Parameters:
    - directory: Path to the directory containing GGUF files.
    """
    if not os.path.isdir(directory):
        logger.error(f"Directory {directory} does not exist.")
        return

    # List all GGUF files in the directory
    files = [f for f in os.listdir(directory) if f.endswith(".gguf")]
    if not files:
        logger.info(f"No GGUF files found in {directory}.")
        return

    re = []
    for file in files:
        file_path = os.path.join(directory, file)
        # print(f"Reading {file_path}:") # noqa: NP100
        # print("\n") # noqa: NP100
        re.extend(read_gguf_file(file_path))
    re = {r.name: r for r in re}
    return re


if use_real_weights := True:
    gguf_weights = read_gguf_directory(gguf_path)
    layer_idx = 0
    q_a_proj_weight, type = get_torch_tensor_and_type_from_gguf(gguf_weights, f"blk.{layer_idx}.attn_q_a.weight")
    q_a_proj.weight = nn.Parameter(q_a_proj_weight.view(torch.bfloat16), requires_grad=False)
    q_a_type = type

    q_a_norm_weight, type = get_torch_tensor_and_type_from_gguf(gguf_weights, f"blk.{layer_idx}.attn_q_a_norm.weight")
    q_a_norm = q_a_norm_weight.view(torch.float32)
    # config.q_a_norm = q_a_norm_weight.data_ptr()
    # config.q_a_norm_type = type_to_ggml_type(type)

    q_b_proj_weight, type = get_torch_tensor_and_type_from_gguf(gguf_weights, f"blk.{layer_idx}.attn_q_b.weight")
    q_b_proj.weight = nn.Parameter(q_b_proj_weight.view(torch.bfloat16), requires_grad=False)

    kv_a_proj_with_mqa_weight, type = get_torch_tensor_and_type_from_gguf(
        gguf_weights, f"blk.{layer_idx}.attn_kv_a_mqa.weight"
    )
    kv_a_proj_with_mqa.weight = nn.Parameter(kv_a_proj_with_mqa_weight.view(torch.bfloat16), requires_grad=False)

    kv_a_norm_weight, type = get_torch_tensor_and_type_from_gguf(gguf_weights, f"blk.{layer_idx}.attn_kv_a_norm.weight")
    kv_a_norm = kv_a_norm_weight.view(torch.float32)
    # config.kv_a_norm = kv_a_norm_weight.data_ptr()
    # config.kv_a_norm_type = type_to_ggml_type(type)

    kv_b_proj_weight, type = get_torch_tensor_and_type_from_gguf(gguf_weights, f"blk.{layer_idx}.attn_kv_b.weight")
    kv_b_proj.weight = nn.Parameter(kv_b_proj_weight.view(torch.bfloat16), requires_grad=False)

    o_proj_weight, type = get_torch_tensor_and_type_from_gguf(gguf_weights, f"blk.{layer_idx}.attn_output.weight")
    o_proj.weight = nn.Parameter(o_proj_weight.view(torch.bfloat16), requires_grad=False)

else:
    init.normal_(q_a_proj.weight, mean=0.0, std=0.02)
    init.normal_(q_b_proj.weight, mean=0.0, std=0.02)
    init.normal_(kv_a_proj_with_mqa.weight, mean=0.0, std=0.02)
    init.normal_(kv_b_proj.weight, mean=0.0, std=0.02)
    init.normal_(o_proj.weight, mean=0.0, std=0.02)

x_reshaped = kv_b_proj.weight.view(num_heads, 2, nope_size, kv_lora_rank)
q_absorb = x_reshaped[:, 0]
out_absorb = x_reshaped[:, 1]


hidden_states = torch.randn((qlen, hidden_size), dtype=input_type).to("cpu").contiguous()


def test_cpu_mla():
    os.environ["BLAS_NUM_THREADS"] = "1"
    q_a_proj_weight = q_a_proj.weight.to(weight_type).to("cpu").contiguous()
    q_b_proj_weight = q_b_proj.weight.to(weight_type).to("cpu").contiguous()
    kv_a_proj_with_mqa_weight = kv_a_proj_with_mqa.weight.to("cpu").to(weight_type).contiguous()
    kv_b_proj_weight = kv_b_proj.weight.to(weight_type).to("cpu").contiguous()
    o_proj_weight = o_proj.weight.to(weight_type).to("cpu").contiguous()

    config = kt_kernel_ext.mla.MLAConfig(
        hidden_size,
        q_lora_rank,
        kv_lora_rank,
        num_heads,
        nope_size,
        rope_size,
    )
    config.max_qlen = max_qlen
    config.max_kvlen = max_kvlen
    config.max_position_embeddings = max_position_embeddings
    config.rope_scaling_factor = rope_scaling["factor"]
    config.rope_theta = rope_theta
    config.rope_scaling_beta_fast = rope_scaling["beta_fast"]
    config.rope_scaling_beta_slow = rope_scaling["beta_slow"]
    config.rope_scaling_mscale = rope_scaling["mscale"]
    config.rope_scaling_mscale_all_dim = rope_scaling["mscale_all_dim"]
    config.rope_scaling_original_max_position_embeddings = rope_scaling["original_max_position_embeddings"]

    config.q_a_proj = q_a_proj_weight.data_ptr()
    config.q_b_proj = q_b_proj_weight.data_ptr()
    config.kv_a_proj_with_mqa = kv_a_proj_with_mqa_weight.data_ptr()
    config.kv_b_proj = kv_b_proj_weight.data_ptr()
    config.o_proj = o_proj_weight.data_ptr()

    config.q_a_norm = q_a_norm.data_ptr()
    config.q_a_norm_type = ggml_type.FP32
    config.kv_a_norm = kv_a_norm.data_ptr()
    config.kv_a_norm_type = ggml_type.FP32
    config.page_count = pages_count

    if weight_type == torch.float32:
        config.q_a_proj_type = ggml_type.FP32
        config.q_b_proj_type = ggml_type.FP32
        config.kv_a_proj_with_mqa_type = ggml_type.FP32
        config.kv_b_proj_type = ggml_type.FP32
        config.w_o_type = ggml_type.FP32
    elif weight_type == torch.float16:
        config.q_a_proj_type = ggml_type.FP16
        config.q_b_proj_type = ggml_type.FP16
        config.kv_a_proj_with_mqa_type = ggml_type.FP16
        config.kv_b_proj_type = ggml_type.FP16
        config.w_o_type = ggml_type.FP16
    elif weight_type == torch.bfloat16:
        config.q_a_proj_type = ggml_type.BF16
        config.q_b_proj_type = ggml_type.BF16
        config.kv_a_proj_with_mqa_type = ggml_type.BF16
        config.kv_b_proj_type = ggml_type.BF16
        config.w_o_type = ggml_type.BF16
    else:
        raise ValueError(f"Unsupported data type: {weight_type}")

    config.pool = CPUInfer.backend_

    if weight_type == torch.float32:
        mla = kt_kernel_ext.mla.MLA_F32(config)
    elif weight_type == torch.float16:
        mla = kt_kernel_ext.mla.MLA_F16(config)
    elif weight_type == torch.bfloat16:
        # mla = kt_kernel_ext.mla.MLA_F32(config)
        mla = kt_kernel_ext.mla.MLA_QUAN_F32(config)
    else:
        raise ValueError(f"Unsupported data type: {weight_type}")

    mla.load_weights()
    mla.set_local_pages(pages_count)

    output = torch.zeros((qlen, hidden_size), dtype=input_type).to("cpu").contiguous()
    mla.forward([qlen], [page_table], [kvlen], hidden_states.data_ptr(), output.data_ptr())
    print("CPU MLA Output: ", output)
    return output


def load_fp16_tensor(file_path, shape):
    # return load_fp32_tensor(file_path, shape)
    return torch.zeros(shape)
    with open(file_path, "rb") as f:
        raw_data = f.read()
    tensor = torch.frombuffer(raw_data, dtype=weight_type)
    tensor = tensor.view(shape)  # 根据你的 shape reshape
    return tensor


def load_fp32_tensor(file_path, shape):
    return torch.zeros(shape)
    with open(file_path, "rb") as f:
        raw_data = f.read()
    tensor = torch.frombuffer(raw_data, dtype=torch.float32)
    tensor = tensor.view(shape)  # 根据你的 shape reshape
    return tensor


def test_torch():
    torch.set_grad_enabled(False)

    softmax_scale = (nope_size + rope_size) ** -0.5
    # 1代表的是压缩的kv的头数
    k_caches = torch.randn(1, pages_count, page_size, 1, kv_lora_rank + rope_size).to(weight_type)
    kv_cache = KDeepSeekV3Cache(page_size=page_size, kv_lora_rank=kv_lora_rank, k_caches=k_caches)

    q_a_layernorm = DeepseekV2RMSNorm(q_lora_rank)
    q_a_layernorm.weight = nn.Parameter(q_a_norm, requires_grad=False)

    x = torch.randn(q_lora_rank, dtype=weight_type) * 100
    print(x)
    print(q_a_layernorm(x))

    kv_a_layernorm = DeepseekV2RMSNorm(kv_lora_rank)
    kv_a_layernorm.weight = nn.Parameter(kv_a_norm, requires_grad=False)

    # 第三步：拆分成两个 tensor
    # q_absorb, out_absorb = x_permuted[:, 0], x_permuted[:, 1]  # 都是 (num_heads, nope_size, kv_lora_rank
    # q_absorb = kv_b_proj[:, ] # torch.randn(num_heads, nope_size, kv_lora_rank, dtype=data_type)
    # out_absorb = kv_b_proj # torch.randn(num_heads, nope_size, kv_lora_rank, dtype=data_type)

    rotary_emb = DeepseekV3YarnRotaryEmbedding(
        rope_size,
        max_position_embeddings=max_position_embeddings,
        scaling_factor=rope_scaling["factor"],
        base=rope_theta,
        beta_fast=rope_scaling["beta_fast"],
        beta_slow=rope_scaling["beta_slow"],
        mscale=rope_scaling["mscale"],
        mscale_all_dim=rope_scaling["mscale_all_dim"],
        original_max_position_embeddings=rope_scaling["original_max_position_embeddings"],
    )
    # 构造一个qlen 长度的输入 hidden_states, 对应的历史 kv_indptr 是[0:bsz]
    # kv_indices 是[0:bsz]，page_idx=[0:bsz], page_offset=[kvlen:qlen+kvlen]
    # last_page_len = [qlen+kvlen,...] layer_idx = 1
    # position_ids = [kvlen:qlen+kvlen]
    q_indptr = torch.tensor([0, qlen]).to(torch.int32)

    kv_indptr = torch.tensor([0, (qlen + kvlen + page_size - 1) // page_size]).to(torch.int32)
    kv_indices = torch.tensor(range(pages_count)).to(torch.int32)

    page_idx = torch.tensor([i // page_size for i in range(kvlen, kvlen + qlen)]).to(torch.int32)
    page_offset = torch.tensor([i % page_size for i in range(kvlen, kvlen + qlen)]).to(torch.int32)

    last_page_len = torch.tensor([256], device=hidden_states.device)
    position_ids = torch.tensor(range(kvlen, kvlen + qlen)).to(torch.int32)

    # 按照行创建 mask [qlen,kvlen+qlen]
    attention_masks = torch.zeros((max_qlen, max_kvlen), dtype=weight_type)
    for i in range(max_qlen):
        attention_masks[i, i + kvlen + 1 :] = -inf

    def torch_attn(
        hidden_states_i: torch.Tensor,
        kv_cache: KDeepSeekV3Cache,
        position_ids: torch.Tensor,
        page_idx: torch.Tensor,
        page_offset: torch.Tensor,
        attention_masks: Optional[list[torch.Tensor]] = None,
        q_indptr: Optional[torch.Tensor] = None,
        kv_indices: Optional[torch.Tensor] = None,
        kv_indptr: Optional[torch.Tensor] = None,
        bsz_tensors: Optional[torch.Tensor] = None,
        last_page_len: Optional[torch.Tensor] = None,
        layer_idx: Optional[int] = None,
    ):
        global out_absorb
        global q_absorb
        hidden_states = hidden_states_i.to(weight_type)
        # range bsz_tensors
        final_attention_output = torch.tensor([], device=hidden_states.device)
        for i in range(bsz_tensors[0]):
            batch_num_tokens_tensors = q_indptr[i + 1] - q_indptr[i]
            batch_last_page_len = last_page_len[i]
            # kv_total_len is kv_len, batch_compressed_kv is compressed_kv, batch_k_pe is k_pe
            batch_page_idx = page_idx[q_indptr[i] : q_indptr[i + 1]]
            batch_page_offset = page_offset[q_indptr[i] : q_indptr[i + 1]]
            # kv_page_nums is the number of pages for the current batch
            kv_page_nums = kv_indptr[i + 1] - kv_indptr[i]
            # kv_total_len is the total length of the kv cache for the current batch (kv_len for algorithm)
            kv_total_len = kv_page_nums * page_size
            if batch_last_page_len is not None:
                kv_total_len = kv_total_len - (page_size - batch_last_page_len)
            # print(f"kv_total_len's shape {kv_total_len.shape}")
            # kv_index is the index of the kv cache pages for the current batch
            kv_index = kv_indices[kv_indptr[i] : kv_indptr[i + 1]]
            # we can index [kv_index, page_offset_indices] to get the kv cache for the current batch
            # from q_indptr[i] to q_indptr[i+1] is the range of the current batch
            batch_hidden_states = hidden_states[q_indptr[i] : q_indptr[i + 1]]
            batch_position_ids = position_ids[q_indptr[i] : q_indptr[i + 1]]
            qlen, _ = batch_hidden_states.size()
            # print("qlen -> ", qlen)

            hidden_states_to_check = load_fp16_tensor("./debug/query_0_tp_0_input.bin", batch_hidden_states.shape)
            diff = torch.abs(batch_hidden_states - hidden_states_to_check).max()
            print("hidden_states diff -> ", diff)

            q_lora = q_a_proj(batch_hidden_states)
            # q_lora_to_check = load_fp16_tensor('./debug/query_0_tp_0_qlora.bin', q_lora.shape)
            # q_lora_to_check_test = load_fp16_tensor('./debug/query_0_tp_0_qlora_test.bin', q_lora.shape)
            # diff = torch.abs(q_lora - q_lora_to_check).max()
            # diff_test = torch.abs(q_lora - q_lora_to_check_test).max()
            # print("q_lora max diff -> ", diff)
            # print("q_lora max diff test -> ", diff_test)
            # mae =  torch.mean(torch.abs(q_lora - q_lora_to_check))
            # mae_test =  torch.mean(torch.abs(q_lora - q_lora_to_check_test))
            # print("q_lora mae -> ", mae)
            # print("q_lora mae test -> ", mae_test)

            q_lora_norm = q_a_layernorm(q_lora)
            # q_lora_norm_to_check = load_fp16_tensor('./debug/query_0_tp_0_qlora_norm.bin', q_lora_norm.shape)
            # q_lora_norm_to_check_test = load_fp16_tensor('./debug/query_0_tp_0_qlora_norm_test.bin', q_lora_norm.shape)
            # diff = torch.abs(q_lora_norm - q_lora_norm_to_check).max()
            # mae =  torch.mean(torch.abs(q_lora_norm - q_lora_norm_to_check))
            # diff_test = torch.abs(q_lora_norm - q_lora_norm_to_check_test).max()
            # mae_test =  torch.mean(torch.abs(q_lora_norm - q_lora_norm_to_check_test))
            # print("q_lora_norm diff -> ", diff)
            # print("q_lora_norm mae -> ", mae)
            # print("q_lora_norm diff test -> ", diff_test)
            # print("q_lora_norm mae test -> ", mae_test)

            q = q_b_proj(q_lora_norm)
            # for v3, bsz, qlen, num_heads(128), qk_head_dim(192=128(nope)+64(rope))
            q = q.view(qlen, num_heads, nope_size + rope_size)
            # q_nope is [qlen, num_heads(128), qk_nope_head_dim(128)]
            # q_pe is [qlen, num_heads(128), qk_rope_head_dim(64)]
            q_nope, q_pe = torch.split(q, [nope_size, rope_size], dim=-1)

            # compressed_kv is [qlen, kv_lora_rank(512) + rope(64)]
            compressed_kv = kv_a_proj_with_mqa(batch_hidden_states)
            # compressed_kv is [qlen, kv_lora_rank(512)], k_pe is [qlen, rope(64)]
            compressed_kv, k_pe = torch.split(compressed_kv, [kv_lora_rank, rope_size], dim=-1)
            compressed_kv = compressed_kv.contiguous()

            # compressed_kv_page_0 = compressed_kv[0:page_size, :]
            # compressed_kv_to_check = load_fp16_tensor('./debug/query_0_tp_0_page_0_kv_lora_rank',
            #                                           compressed_kv_page_0.shape)
            # diff = torch.abs(compressed_kv_page_0 - compressed_kv_to_check).max()
            # mae =  torch.mean(torch.abs(compressed_kv_page_0 - compressed_kv_to_check))
            # print("compressed_kv diff -> ", diff)
            # print("compressed_kv mae -> ", mae)

            compressed_kv = kv_a_layernorm(compressed_kv)
            # k_pe is [qlen, 1, qk_rope_head_dim(64)]

            # compressed_kv_page_0 = compressed_kv[0:page_size, :]
            # compressed_kv_to_check = load_fp16_tensor('./debug/query_0_tp_0_page_0_kv_lora_rank_norm',
            #                                           compressed_kv_page_0.shape)
            # diff = torch.abs(compressed_kv_page_0 - compressed_kv_to_check).max()
            # mae =  torch.mean(torch.abs(compressed_kv_page_0 - compressed_kv_to_check))
            # print("compressed_kv diff norm -> ", diff)
            # print("compressed_kv mae norm -> ", mae)

            k_pe = k_pe.view(qlen, 1, rope_size)
            # compressed_kv is [qlen, 1, kv_lora_rank(512)]
            compressed_kv = compressed_kv.view(qlen, 1, kv_lora_rank)

            cos, sin = rotary_emb(q_pe, batch_position_ids)

            # q_nope_check = q_nope.transpose(0, 1) # qlen is 1, no GPU overhead, same below

            # q_nope_0_to_check = load_fp16_tensor('./debug/query_0_tp_0_q_nope', q_nope_check[0].shape)
            # q_nope_0_to_check_test = load_fp16_tensor('./debug/query_0_tp_0_q_nope_test', q_nope_check[0].shape)
            # diff = torch.abs(q_nope_check[0] - q_nope_0_to_check).max()
            # mae =  torch.mean(torch.abs(q_nope_check[0] - q_nope_0_to_check))
            # diff_test = torch.abs(q_nope_check[0] - q_nope_0_to_check_test).max()
            # mae_test =  torch.mean(torch.abs(q_nope_check[0] - q_nope_0_to_check_test))
            # print("q_nope[0] diff -> ", diff)
            # print("q_nope[0] mae -> ", mae)
            # print("q_nope[0] diff test -> ", diff_test)
            # print("q_nope[0] mae test -> ", mae_test)

            q_pe_nope = q_pe.transpose(0, 1)
            # q_pe_0_to_check = load_fp16_tensor('./debug/query_0_tp_0_q_rope', q_pe_nope[0].shape)
            # q_pe_0_to_check = load_fp16_tensor('./debug/query_0_tp_0_q_rope_no_rope', q_pe_nope[0].shape)
            # q_pe_0_to_check_test = load_fp16_tensor('./debug/query_0_tp_0_q_rope_no_rope_test', q_pe_nope[0].shape)
            # diff = torch.abs(q_pe_nope[0] - q_pe_0_to_check).max()
            # mae =  torch.mean(torch.abs(q_pe_nope[0] - q_pe_0_to_check))
            # diff_test = torch.abs(q_pe_nope[0] - q_pe_0_to_check_test).max()
            # mae_test =  torch.mean(torch.abs(q_pe_nope[0] - q_pe_0_to_check_test))
            # print("q_pe nope[0] diff -> ", diff)
            # print("q_pe nope[0] mae -> ", mae)
            # print("q_pe nope[0] diff test -> ", diff_test)
            # print("q_pe nope[0] mae test -> ", mae_test)

            # cos_to_check = load_fp32_tensor('./debug/query_0_tp_0_rope_cos', (qlen,32))
            # diff = torch.abs(cos[:,:32]-cos_to_check).max()
            # mae =  torch.mean(torch.abs(cos[:,:32]-cos_to_check))
            # print("cos diff -> ", diff)
            # print("cos mae -> ", mae)
            # sin_to_check = load_fp32_tensor('./debug/query_0_tp_0_rope_sin', (qlen,32))
            # diff = torch.abs(sin[:,:32]-sin_to_check).max()
            # mae =  torch.mean(torch.abs(sin[:,:32]-sin_to_check))
            # print("sin diff -> ", diff)
            # print("sin mae -> ", mae)

            # new_q_pe = q_pe.transpose(0, 1)
            # qa = new_q_pe[:,:,range(0,64,2)]
            # qb = new_q_pe[:,:,range(1,65,2)]
            # # q1 = (qa * cos[:,:32] - qb * sin[:,:32])
            # # q2 = (qb*cos[:,:32] + qa*sin[:,:32])
            # q1 = (qa * cos_to_check - qb * sin_to_check)
            # q2 = (qb*cos_to_check + qa*sin_to_check)
            # q_new = torch.cat((q1,q2), dim=-1)
            # print(f"q_pe shape{q_pe.shape}, k_pe shape {k_pe.shape}")
            # new_q_pe = torch.zeros_like(q_pe)
            # new_q_pe[:,:,range(0,64,2)] = 1
            # new_q_pe[:,:,range(1,65,2)] = 10
            q_pe, k_pe = apply_rotary_pos_emb(q_pe.unsqueeze(0), k_pe.unsqueeze(0), cos, sin, unsqueeze_dim=1)
            q_pe = q_pe.squeeze(0)
            # q_pe is [num_heads(128), qlen, qk_rope_head_dim(64)]
            q_pe.transpose_(0, 1)

            # diff = torch.abs(q_pe - q_new).max()
            # print("q_pe diff -> ", diff)

            # q_pe_0_to_check = load_fp16_tensor('./debug/query_0_tp_0_q_rope', q_pe[0].shape)
            # diff = torch.abs(q_pe[0] - q_pe_0_to_check).max()
            # mae =  torch.mean(torch.abs(q_pe[0] - q_pe_0_to_check))
            # print("q_pe[0] diff -> ", diff)
            # print("q_pe[0] mae -> ", mae)

            # diff = torch.abs(q_pe_0_to_check - q_new[0]).max()
            # mae =  torch.mean(torch.abs(q_pe_0_to_check - q_new[0]))
            # print("q_pe[0] 2  diff -> ", diff)
            # print("q_pe[0] 2 mae -> ", mae)

            if kv_cache is not None:
                cache_kwargs = {
                    "sin": sin,
                    "cos": cos,
                    "page_idx": batch_page_idx,
                    "page_offset": batch_page_offset,
                }  # Specific to RoPE models
                compressed_kv_with_k_pe = kv_cache.update(
                    compressed_kv.unsqueeze(0), k_pe, layer_idx, batch_page_idx, batch_page_offset, cache_kwargs
                )
                compressed_kv = compressed_kv_with_k_pe[:, :, :, :kv_lora_rank].view(-1, page_size, kv_lora_rank)
                k_pe = compressed_kv_with_k_pe[:, :, :, kv_lora_rank:].view(-1, page_size, rope_size)
            # q_absorb is [num_heads(128), qk_nope_head_dim(128), kv_lora_rank(512)]
            # out_absorb is [num_heads(128), kv_lora_rank(512), v_head_dim(128)] v_head_dim is also the nope dim
            # q_absorb, out_absorb = get_absorbed()
            # q_nope is [num_heads(128), qlen, qk_nope_head_dim(128)]
            q_nope = q_nope.transpose(0, 1)  # qlen is 1, no GPU overhead, same below

            # q_nope_0_to_check = load_fp16_tensor('./debug/query_0_tp_0_q_nope', q_nope[0].shape)
            # diff = torch.abs(q_nope[0] - q_nope_0_to_check).max()
            # mae =  torch.mean(torch.abs(q_nope[0] - q_nope_0_to_check))
            # print("q_nope[0] diff -> ", diff)

            # q_nope is [num_heads(128), qlen, kv_lora_rank(512)]
            q_nope = torch.matmul(q_nope, q_absorb)  # batched MM

            # k_b_proj_check = load_fp16_tensor('./debug/query_0_tp_0_k_b_lora', (nope_size,kv_lora_rank))
            # diff = torch.abs(q_absorb[0] - k_b_proj_check).max()
            # print("kv b lora weight[0] diff -> ", diff)

            # q_absorb_check = load_fp16_tensor('./debug/query_0_tp_0_q_absorb', (kv_lora_rank,1024))
            # q_absorb_check = q_absorb_check[:,0:qlen].transpose(0,1)
            # diff = torch.abs(q_nope[0] - q_absorb_check).max()
            # mae =  torch.mean(torch.abs(q_nope[0] - q_absorb_check))
            # print("q_nope absorb diff -> ", diff)
            # print("q_nope absorb mae -> ", mae)

            # # q_nope is [qlen, num_heads(128), kv_lora_rank(512)]
            # q_nope = q_nope.transpose(0, 1)

            # we need to index out the compressed_kv and k_pe for the current batch
            batch_compressed_kv = None
            batch_k_pe = None
            for page_index in kv_index:
                if kv_total_len > page_size:
                    tmp_compressed_kv = compressed_kv[page_index, 0:page_size, :]
                    tmp_k_pe = k_pe[page_index, 0:page_size, :]
                    if batch_compressed_kv is None or batch_k_pe is None:
                        batch_compressed_kv = tmp_compressed_kv
                        batch_k_pe = tmp_k_pe
                    else:
                        batch_compressed_kv = torch.cat((batch_compressed_kv, tmp_compressed_kv), dim=0)
                        batch_k_pe = torch.cat((batch_k_pe, tmp_k_pe), dim=0)
                    kv_total_len -= page_size
                else:
                    tmp_compressed_kv = compressed_kv[page_index, 0:kv_total_len, :]
                    tmp_k_pe = k_pe[page_index, 0:kv_total_len, :]
                    if batch_compressed_kv is None or batch_k_pe is None:
                        batch_compressed_kv = tmp_compressed_kv
                        batch_k_pe = tmp_k_pe
                    else:
                        batch_compressed_kv = torch.cat((batch_compressed_kv, tmp_compressed_kv), dim=0)
                        batch_k_pe = torch.cat((batch_k_pe, tmp_k_pe), dim=0)
                    break
            # batch_compressed_kv is [kv_total_len(k_len), kv_lora_rank(512)]
            # batch_k_pe is [kv_total_len(k_len), qk_rope_head_dim(64)]

            # k_pe_to_check = load_fp16_tensor('./debug/query_0_tp_0_page_0_k_rope', (256,64))
            # diff = torch.abs(batch_k_pe[:256] - k_pe_to_check).max()
            # mae =  torch.mean(torch.abs(batch_k_pe[:256] - k_pe_to_check))
            # print("k_pe diff -> ", diff)
            # print("k_pe mae -> ", mae)

            pe_weights = torch.matmul(q_pe, batch_k_pe.mT)
            kv_total_len = kv_page_nums * page_size
            # pe_weights_0 = load_fp16_tensor('./debug/query_0_tp_0_pe_attention_weights', (1024,4096))
            # pe_weights_0 = pe_weights_0[0:qlen, 0:kv_total_len]
            # diff = torch.abs(pe_weights[0] - pe_weights_0).max()
            # print("pe_weights[0] diff -> ", diff)

            attention_weights = pe_weights + torch.matmul(q_nope, batch_compressed_kv.mT)

            # raw_weights = load_fp16_tensor('./debug/query_0_tp_0_raw_attention_weights', (1024, 4096))
            # raw_weights = raw_weights[0:qlen, 0:kv_total_len]
            # diff = torch.abs(attention_weights[0] - raw_weights).max()
            # print("raw attention_weights[0] diff -> ", diff)

            attention_weights = attention_weights * softmax_scale
            # attention_weights is [num_heads(128), qlen, k_len]

            # attention_weights = attention_weights.transpose(0,1).unsqueeze(0).squeeze(-1).expand(qlen,-1,-1).transpose(0,1)

            # attention_masks[i] is [qlen, k_len]

            print(attention_weights.shape)
            print(attention_masks.shape)
            attention_weights = (
                attention_weights + attention_masks[: attention_weights.shape[1], : attention_weights.shape[2]]
            )
            # attention_weights shape is [num_heads(128), qlen, k_len]

            attention_weights = nn.functional.softmax(attention_weights, dim=-1, dtype=weight_type).to(q_pe.dtype)

            # attention_weights_0 = load_fp16_tensor('./debug/query_0_tp_0_attention_weights', (1024, 4096))
            # attention_weights_0 = attention_weights_0[0:qlen, 0:kv_total_len]
            # diff = torch.abs(attention_weights[0] - attention_weights_0).max()
            # print("attention_weights[0] diff -> ", diff)

            attn_output = torch.matmul(attention_weights, batch_compressed_kv)  # [num_heads(128),qlen, lora_rank(512)]
            # out_absorb shape is [num_heads(128), kv_lora_rank(512), v_head_dim(128)]

            # o_absorb_check = load_fp16_tensor('./debug/query_0_tp_0_o_absorb', (qlen,kv_lora_rank))
            # diff = torch.abs(attn_output[0] - o_absorb_check).max()
            # print("o absorb[0] diff -> ", diff)

            out_absorb = out_absorb.transpose(1, 2)  # [qlen, num_heads(128), v_head_dim(128)]
            # q for qlen, n for num_heads, h for v_head_dim, v for kv_lora_rank
            attn_output = torch.matmul(attn_output, out_absorb)  # [num_heads(128), qlen, v_head_dim(128)]

            # attn_output_check_0 = load_fp16_tensor('./debug/query_0_tp_0_attention_output', (qlen, nope_size))
            # diff = torch.abs(attn_output[0] - attn_output_check_0).max()
            # print("attn_output[0] diff -> ", diff)

            attn_output = attn_output.transpose(0, 1)  # [qlen, num_heads(128), v_head_dim(128)]
            attn_output = attn_output.reshape(qlen, num_heads * nope_size)

            w_o = o_proj.weight.view([hidden_size, num_heads * nope_size])
            output = torch.matmul(attn_output, w_o.transpose(0, 1))
            output = output.view(qlen, hidden_size)

            # output_0_check = load_fp16_tensor('./debug/query_0_tp_0_qlen_output', (qlen, hidden_size))
            # h1_o = w_o[:,:128]
            # local_o_check = load_fp16_tensor('./debug/query_0_tp_0_local_w_o', (hidden_size, 128))
            # diff = torch.abs(local_o_check - h1_o).max()
            # print("local w_o diff -> ", diff)

            # h1_output = torch.matmul(attn_output[:,:128],h1_o.transpose(0,1))
            # diff = torch.abs(h1_output - output_0_check).max()
            # print("h1_output diff -> ", diff)

            # output_check = load_fp16_tensor('./debug/output.bin', output.shape)
            # diff = torch.abs(output - output_check).max()
            # mae =   torch.mean(torch.abs(output - output_check))
            # print("output diff -> ", diff)

            final_attention_output = torch.cat((final_attention_output, output), dim=0)
        return final_attention_output

    torch_output = torch_attn(
        hidden_states,
        kv_cache,
        position_ids,
        page_idx,
        page_offset,
        attention_masks=attention_masks,
        q_indptr=q_indptr,
        kv_indices=kv_indices,
        kv_indptr=kv_indptr,
        bsz_tensors=bsz_tensors,
        last_page_len=last_page_len,
        layer_idx=0,
    )
    print("Torch Output: ", torch_output)
    return torch_output


torch.set_printoptions(sci_mode=False, precision=5)
output_cpu = test_cpu_mla()
output_torch = test_torch()
print("Output CPU: ", output_cpu)
print("Output Torch: ", output_torch)
diff = (output_cpu - output_torch).abs()
# 计算相对误差
diff_relative = diff / (output_cpu.abs())
# 把 diff_relative 中的 NaN 替换为 0
diff_relative = torch.where(torch.isnan(diff_relative), torch.zeros_like(diff_relative), diff_relative)
diff_relative_mean = torch.mean(torch.abs(output_cpu - output_torch)) / torch.mean(torch.abs(output_torch))

print(
    f"Diff: ave:{diff.mean()}, max:{diff.max()}, min:{diff.min()},  relative_mean:{diff_relative_mean}, relative_max:{diff_relative.max()}, relative_min:{diff_relative.min()}"
)
assert diff_relative_mean < 2e-1, "CPU and Torch outputs are not close enough!"
