import json
from collections import defaultdict
from dataclasses import asdict, dataclass, field

import torch
import torch_npu
from vllm.logger import logger
from vllm.model_executor.layers.attention import Attention
from vllm.model_executor.models.utils import extract_layer_index
from vllm.utils.math_utils import cdiv

from vllm_ascend.ascend_config import get_ascend_config


# largely follow vllm.v1.worker.utils.bind_kv_cache
def bind_hashk_cache(
    hashk_caches: dict[str, torch.Tensor],
    forward_context: dict[str, Attention],
    runner_hashk_caches: list[torch.Tensor],
    num_attn_module: int = 1,
) -> None:
    """
    Bind the allocated hashk cache to both ModelRunner and forward context so
    that the hashk cache can be used in the forward pass.

    This function:
      1) Fills the ModelRunner's hashk cache list (`runner_hashk_caches`) with
         hashk_caches.
      2) Associates each attention layer in the `forward_context` with its
         corresponding hashk cache in hashk_caches.

    Args:
        hashk_caches: The allocated hashk_caches with layer names as keys.
        forward_context: The global forward context containing all Attention
            layers with layer names as keys.
        runner_hashk_caches: The hashk cache declared by ModelRunner.
    """
    # Bind hashk_caches to ModelRunner; ensure it is empty before binding
    assert len(runner_hashk_caches) == 0

    # Convert hashk_caches dict to a list of tensors in the order of layer_index.
    index2name = defaultdict(list)
    for layer_name in hashk_caches:
        index2name[extract_layer_index(layer_name, num_attn_module)].append(layer_name)

    for layer_index in sorted(index2name.keys()):
        layer_names = index2name[layer_index]
        # TODO: support multiple hashk caches for the same layer index later, e.g., encoder-decoder models.
        layer_name = layer_names[0]
        runner_hashk_caches.append(hashk_caches[layer_name])

    # Bind kv_caches to forward context
    for layer_name, hashk_cache in hashk_caches.items():
        # NOTE: Use list because of v0 PP virtual engine.
        forward_context[layer_name].hashk_cache = [hashk_cache]


def bind_hashk_cache_nope(
    hashk_caches_nope: dict[str, torch.Tensor],
    forward_context: dict[str, Attention],
    runner_hashk_caches_nope: list[torch.Tensor],
    num_attn_module: int = 1,
) -> None:
    """
    Bind the allocated hashk cache for nope in MLA to both ModelRunner and forward context so
    that the hashk cache for nope can be used in the forward pass.

    This function:
      1) Fills the ModelRunner's hashk cache list (`runner_hashk_caches_nope`) with
         hashk_caches_nope.
      2) Associates each attention layer in the `forward_context` with its
         corresponding hashk cache for nope in MLA in hashk_caches_nope.

    Args:
        hashk_caches_nope: The allocated hashk_caches_nope with layer names as keys.
        forward_context: The global forward context containing all Attention
            layers with layer names as keys.
        runner_hashk_caches_nope: The hashk cache for nope declared by ModelRunner.
    """
    # Bind hashk_caches_nope to ModelRunner; ensure it is empty before binding
    assert len(runner_hashk_caches_nope) == 0

    # Convert hashk_caches_nope dict to a list of tensors in the order of layer_index.
    index2name = defaultdict(list)
    for layer_name in hashk_caches_nope:
        index2name[extract_layer_index(layer_name, num_attn_module)].append(layer_name)

    for layer_index in sorted(index2name.keys()):
        layer_names = index2name[layer_index]
        # TODO: support multiple hashk caches for the same layer index later, e.g., encoder-decoder models.
        layer_name = layer_names[0]
        runner_hashk_caches_nope.append(hashk_caches_nope[layer_name])

    # Bind hashk_caches_nope to forward context
    for layer_name, hashk_cache_nope in hashk_caches_nope.items():
        # NOTE: Use list because of v0 PP virtual engine.
        forward_context[layer_name].hashk_cache_nope = [hashk_cache_nope]


def bind_hashk_cache_rope(
    hashk_caches_rope: dict[str, torch.Tensor],
    forward_context: dict[str, Attention],
    runner_hashk_caches_rope: list[torch.Tensor],
    num_attn_module: int = 1,
) -> None:
    """
    Bind the allocated hashk cache for rope in MLA to both ModelRunner and forward context so
    that the hashk cache for rope can be used in the forward pass.

    This function:
      1) Fills the ModelRunner's hashk cache list (`runner_hashk_caches_rope`) with
         hashk_caches_rope.
      2) Associates each attention layer in the `forward_context` with its
         corresponding hashk cache for rope in MLA in hashk_caches_rope.

    Args:
        hashk_caches_rope: The allocated hashk_caches_rope with layer names as keys.
        forward_context: The global forward context containing all Attention
            layers with layer names as keys.
        runner_hashk_caches_rope: The hashk cache for rope declared by ModelRunner.
    """
    # Bind hashk_caches_rope to ModelRunner; ensure it is empty before binding
    assert len(runner_hashk_caches_rope) == 0

    # Convert hashk_caches_rope dict to a list of tensors in the order of layer_index.
    index2name = defaultdict(list)
    for layer_name in hashk_caches_rope:
        index2name[extract_layer_index(layer_name, num_attn_module)].append(layer_name)

    for layer_index in sorted(index2name.keys()):
        layer_names = index2name[layer_index]
        # TODO: support multiple hashk caches for the same layer index later, e.g., encoder-decoder models.
        layer_name = layer_names[0]
        runner_hashk_caches_rope.append(hashk_caches_rope[layer_name])

    # Bind hashk_caches_rope to forward context
    for layer_name, hashk_cache_rope in hashk_caches_rope.items():
        # NOTE: Use list because of v0 PP virtual engine.
        forward_context[layer_name].hashk_cache_rope = [hashk_cache_rope]


@dataclass
class KVCompConfig:
    """
    Dataclass representing the configuration for KVComp.
    """

    model_name: str = "DummyModel"
    is_mla: bool = False

    # either "random" or "fixed"
    hash_weight_type: str | None = "random"

    num_hidden_layers: int = 36

    # the minimal seq_len to trigger KVComp
    seq_len_threshhold: int = 2048

    # any value divisible by 128
    chunk_size: int = 128

    # either "max", "min" or "sum"
    chunk_repre_method: str = "max"

    head_dim: int = 128
    hash_bits: int = 128

    top_k_ratio_per_layer: list[float] = field(default_factory=lambda: [0.3] * 36)
    top_k_index_reuse: list[int] = field(default_factory=lambda: [-1] * 36)

    # nonnegative means slicing from the start, negative means slicing from the end
    must_select_blocks: list[int] = field(default_factory=lambda: [0, -2, -1])

    # used when is_mla=True and hash_weight_type="fixed"
    hash_weight: list[list[float]] | None = None

    # Conditional fields if is_mla=True
    kv_lora_rank: int = 72  # we need to specify it if is_mla=True
    qk_rope_head_dim: int = 72  # we need to specify it if is_mla=True
    hash_bits_kv_lora: int = 8  # we need to specify it if is_mla=True
    hash_bits_qk_rope: int = 8  # we need to specify it if is_mla=True
    hash_weight_kv_lora: list[list[float]] | None = None
    hash_weight_qk_rope: list[list[float]] | None = None

    vllm_hash_attention_topk: int = 4096
    vllm_hash_attention_reduction_head_num: int | None = None
    vllm_hash_attention_rollback_layers: list[int] = field(
        default_factory=lambda: []
    )  # layers to rollback, empty means no rollback
    vllm_hash_attention_skip_layers: list[int] = field(
        default_factory=lambda: []
    )  # layers to skip, empty means no skip

    # generate non-MLA config data
    def generate_config_data(
        self,
        model_name: str,
        hash_weight_type: str,
        num_hidden_layers: int,
        seq_len_threshhold: int,
        chunk_size: int,
        chunk_repre_method: str,
        head_dim: int,
        hash_bits: int,
        top_k_ratio_per_layer: list[float],
        top_k_index_reuse: list[int],
        must_select_blocks: list[int],
    ) -> None:
        self.is_mla = False
        self.model_name = model_name

        if hash_weight_type not in ["uniform", "fixed"]:
            raise ValueError(f"hash_weight_type should be either 'random' or 'fixed', but got {hash_weight_type}")
        self.hash_weight_type = hash_weight_type

        self.num_hidden_layers = num_hidden_layers
        self.seq_len_threshhold = seq_len_threshhold

        if chunk_size % 128 != 0:
            raise ValueError(f"chunk_size should be divisible by 128, but got {chunk_size}")
        self.chunk_size = chunk_size

        if chunk_repre_method not in ["max", "min", "sum"]:
            raise ValueError(f"chunk_size should be divisible by 128, but got {chunk_size}")
        self.chunk_repre_method = chunk_repre_method

        self.head_dim = head_dim
        self.hash_bits = hash_bits

        if len(top_k_ratio_per_layer) != num_hidden_layers:
            raise ValueError(
                f"top_k_ratio_per_layer length should be equal to num_hidden_layers={num_hidden_layers}, "
                f"but got {len(top_k_ratio_per_layer)}"
            )
        self.top_k_ratio_per_layer = top_k_ratio_per_layer
        if len(top_k_index_reuse) != num_hidden_layers:
            raise ValueError(
                f"top_k_index_reuse length should be equal to num_hidden_layers={num_hidden_layers}, "
                f"but got {len(top_k_index_reuse)}"
            )
        self.top_k_index_reuse = top_k_index_reuse

        self.must_select_blocks = must_select_blocks

        if hash_weight_type == "random":
            logger.info("hash_weight_type is 'random', hash_weight will be generated automatically.")
            self.hash_weight = None
        else:
            logger.info("hash_weight_type is 'fixed', please manually set hash_weight in the config json file.")

    # generate MLA config data
    def generate_mla_config_data(
        self,
        model_name: str,
        hash_weight_type: str,
        num_hidden_layers: int,
        seq_len_threshhold: int,
        chunk_size: int,
        chunk_repre_method: str,
        kv_lora_rank: int,
        qk_rope_head_dim: int,
        hash_bits_kv_lora: int,
        hash_bits_qk_rope: int,
        top_k_ratio_per_layer: list[float],
        top_k_index_reuse: list[int],
        must_select_blocks: list[int],
    ) -> None:
        self.is_mla = True
        self.model_name = model_name
        if hash_weight_type not in ["random", "fixed"]:
            raise ValueError(f"hash_weight_type should be either 'random' or 'fixed', but got {hash_weight_type}")
        self.hash_weight_type = hash_weight_type

        self.num_hidden_layers = num_hidden_layers
        self.seq_len_threshhold = seq_len_threshhold
        if chunk_size % 128 != 0:
            raise ValueError(f"chunk_size should be divisible by 128, but got {chunk_size}")
        self.chunk_size = chunk_size
        if chunk_repre_method not in ["max", "min", "sum"]:
            raise ValueError(f"chunk_repre_method should be either 'max', 'min' or 'sum', but got {chunk_repre_method}")
        self.chunk_repre_method = chunk_repre_method
        self.head_dim = qk_rope_head_dim + kv_lora_rank
        self.hash_bits = hash_bits_qk_rope + hash_bits_kv_lora
        self.kv_lora_rank = kv_lora_rank
        self.qk_rope_head_dim = qk_rope_head_dim
        self.hash_bits_kv_lora = hash_bits_kv_lora
        self.hash_bits_qk_rope = hash_bits_qk_rope

        if len(top_k_ratio_per_layer) != num_hidden_layers:
            raise ValueError(
                f"top_k_ratio_per_layer length should be equal to num_hidden_layers={num_hidden_layers}, "
                f"but got {len(top_k_ratio_per_layer)}"
            )
        self.top_k_ratio_per_layer = top_k_ratio_per_layer
        if len(top_k_index_reuse) != num_hidden_layers:
            raise ValueError(
                f"top_k_index_reuse length should be equal to num_hidden_layers={num_hidden_layers}, "
                f"but got {len(top_k_index_reuse)}"
            )
        self.top_k_index_reuse = top_k_index_reuse

        self.must_select_blocks = must_select_blocks

        if hash_weight_type == "random":
            logger.info(
                "hash_weight_type is 'random', "
                "hash_weight_kv_lora and hash_weight_qk_rope will be generated automatically."
            )
            self.hash_weight = None
            self.hash_weight_kv_lora = None
            self.hash_weight_qk_rope = None
        else:
            logger.info(
                "hash_weight_type is 'fixed', "
                "please manually set hash_weight_kv_lora and hash_weight_qk_rope in the config json file."
            )

    # set hash_weight when hash_weight_type is "fixed" for non-MLA models
    def set_hash_weight(self, hash_weight: list[list[float]]) -> None:
        if self.hash_weight_type != "fixed":
            raise ValueError("hash_weight can only be set when hash_weight_type is 'fixed'")

        if len(hash_weight) != self.head_dim or len(hash_weight[0]) != self.hash_bits:
            raise ValueError(
                f"hash_weight shape should be ({self.head_dim}, {self.hash_bits}), "
                f"but got ({len(hash_weight)}, {len(hash_weight[0])})"
            )

        self.hash_weight = hash_weight

    # set hash_weight when hash_weight_type is "fixed" for MLA models
    def set_mla_hash_weight(
        self,
        hash_weight_kv_lora: list[list[float]],
        hash_weight_qk_rope: list[list[float]],
    ) -> None:
        if self.hash_weight_type != "fixed":
            raise ValueError("hash_weight can only be set when hash_weight_type is 'fixed'")

        if len(hash_weight_kv_lora) != self.kv_lora_rank or len(hash_weight_kv_lora[0]) != self.hash_bits_kv_lora:
            raise ValueError(
                f"hash_weight_kv_lora shape should be ({self.kv_lora_rank}, {self.hash_bits_kv_lora}), "
                f"but got ({len(hash_weight_kv_lora)}, {len(hash_weight_kv_lora[0])})"
            )

        if len(hash_weight_qk_rope) != self.qk_rope_head_dim or len(hash_weight_qk_rope[0]) != self.hash_bits_qk_rope:
            raise ValueError(
                f"hash_weight_qk_rope shape should be ({self.qk_rope_head_dim}, {self.hash_bits_qk_rope}), "
                f"but got ({len(hash_weight_qk_rope)}, {len(hash_weight_qk_rope[0])})"
            )

        self.hash_weight_kv_lora = hash_weight_kv_lora
        self.hash_weight_qk_rope = hash_weight_qk_rope

    def to_json(self, file_path: str) -> None:
        with open(file_path, "w") as f:
            json.dump(asdict(self), f, indent=4)

    @classmethod
    def from_json(cls, file_path: str) -> "KVCompConfig":
        with open(file_path) as f:
            config_dict = json.load(f)
        return cls(**config_dict)


class HashEncoder:
    """
    HashEncoder converts a float tensor to a binary hash code tensor,
    and it packs every 8 bits into a uint8 number.
    """

    def __init__(self, input_dim: int, hash_bits: int, dtype: torch.dtype, device: torch.device) -> None:
        self.input_dim = input_dim

        if hash_bits % 8 != 0:
            raise ValueError("hash_bits must be a multiple of 8")

        self.hash_bits = hash_bits

        # number of uint8 numbers to store hash_bits bits
        self.hash_numbers = self.hash_bits // 8

        self.dtype = dtype
        self.device = device

        assert self.device.type == "npu"

        if dtype not in [torch.float16, torch.float32, torch.float64]:
            logger.warning(
                "NPU only supports float16/32/64 for hash_weights; got %s. Using float16.",
                dtype,
            )
            self.dtype = torch.float16

        self._init_hash_weights()

    def _init_hash_weights(self):
        random_weights = torch.normal(
            mean=0,
            std=2,
            size=(self.input_dim, self.hash_bits),
            dtype=self.dtype,
            device=self.device,
        )
        Q, R = torch.linalg.qr(random_weights)

        d = torch.sign(torch.diag(R))
        self.hash_weights = Q * d

    def set_hash_weight(self, hash_weights: torch.Tensor) -> None:
        if hash_weights.shape != (self.input_dim, self.hash_bits):
            raise ValueError(
                f"hash_weights shape {hash_weights.shape} "
                f"does not match required shape {(self.input_dim, self.hash_bits)}"
            )
        if hash_weights.dtype != self.dtype:
            raise ValueError(f"hash_weights dtype {hash_weights.dtype} does not match required dtype {self.dtype}")
        if hash_weights.device != self.device:
            raise ValueError(f"hash_weights device {hash_weights.device} does not match required device {self.device}")

        self.hash_weights.copy_(hash_weights)

    def compute_hash(self, x: torch.Tensor) -> torch.Tensor:
        """
        Compute the hash code for input tensor x.
        Args:
            x: input tensor of shape (..., input_dim)
        Returns:
            A tensor of shape (..., hash_numbers=hash_bits // 8) representing the hash codes.
            Each element is a uint8 number representing 8 bits of the hash code.
        """
        if x.shape[-1] != self.input_dim:
            raise ValueError(f"x must be of shape (..., {self.input_dim}), but got {x.shape}")
        if x.device != self.device:
            raise ValueError(f"x device {x.device} does not match required device {self.device}")

        # original shape without the last dimension
        # e.g. x.shape=[s1,s2,s3,input_dim], orig_shape=[s1,s2,s3]
        orig_shape = x.shape[:-1]

        # [N, input_dim], e.g., N = s1*s2*s3
        x_flat = x.reshape(-1, self.input_dim)

        if x_flat.dtype != self.dtype:
            x_flat = x_flat.to(self.dtype)

        # [N, hash_bits]
        xW = torch.matmul(x_flat, self.hash_weights)
        # [N * hash_bits]
        xW_flat = xW.view(-1)
        # [N*hash_numbers], where hash_numbers = hash_bits // 8
        packed_codes_flat = torch.ops._C_ascend.npu_sign_bits_pack(xW_flat, size=1)

        # e.g., [s1, s2, s3, hash_numbers]
        out_shape = orig_shape + (self.hash_numbers,)
        packed_codes = packed_codes_flat.view(out_shape)

        return packed_codes

    def _unpack_hash(self, packed_codes: torch.Tensor) -> torch.Tensor:
        """
        Unpack the hash codes to +1 or -1 bits.
        Args:
            packed_codes: input tensor of shape (..., hash_numbers), dtype=torch.uint8
        Returns:
            A tensor of shape (..., hash_bits=hash_numbers*8) representing the unpacked bits.
            Each element is either -1 or 1.
        """
        if packed_codes.shape[-1] != self.hash_numbers:
            raise ValueError(f"packed_codes must be of shape (..., {self.hash_numbers}), but got {packed_codes.shape}")
        if packed_codes.device != self.device:
            raise ValueError(f"packed_codes device {packed_codes.device} does not match required device {self.device}")
        if packed_codes.dtype != torch.uint8:
            raise ValueError(f"packed_codes dtype {packed_codes.dtype} is not torch.uint8")

        # e.g., packed_codes.shape=[s1, s2, s3, hash_numbers]
        # orig_shape = [s1, s2, s3]
        orig_shape = packed_codes.shape[:-1]

        # [N * hash_numbers], e.g., N = s1*s2*s3
        packed_codes_flat = packed_codes.view(-1)

        # [N * hash_bits]
        unpacked_bits_flat = torch_npu.npu_sign_bits_unpack(packed_codes_flat, size=1, dtype=torch.float16)
        out_shape = orig_shape + (self.hash_bits,)
        unpacked_bits = unpacked_bits_flat.view(out_shape)

        return unpacked_bits


@dataclass
class KVCompMetaData:
    # for both GQA and MLA
    kvcomp_config: KVCompConfig
    chunk_sizes_for_hamming_full: torch.Tensor
    topk_for_hamming_full: torch.Tensor
    topk_for_hamming_full_cpu: torch.Tensor
    seq_lens_for_hamming: torch.Tensor
    hamming_output: torch.Tensor
    seq_lens_from_hamming: torch.Tensor
    seq_lens_for_reshape: torch.Tensor
    valid_query_mask: torch.Tensor
    sink: int
    recent: int
    hash_encoder: HashEncoder
    hashk_caches: list[torch.Tensor]

    num_actual_tokens: int = 0
    max_seq_len_for_hamming: int = 0
    slot_mapping: torch.Tensor = None
    seq_lens_gpu: torch.Tensor = None
    actual_query_lens: torch.Tensor = None
    block_tables_for_hamming: torch.Tensor | None = None


def initialize_kvcomp_metadata(
    max_num_reqs: int, block_size: int, device: torch.device, vllm_config, parallel_config, dtype: torch.dtype
) -> KVCompMetaData:
    """
    Wrapper function to build KVCompMetaData object
    """

    # Auto-detect KVComp config file
    kvcomp_config_path = get_ascend_config().sparse_json
    if kvcomp_config_path is not None:
        kvcomp_config = KVCompConfig.from_json(kvcomp_config_path)
    else:
        raise RuntimeError("KVComp config file not found")

    # Initialize various tensors (replace self.xxx with input parameters)
    chunk_sizes_for_hamming_full = torch.full(
        [max_num_reqs],
        fill_value=block_size,
        dtype=torch.int32,
        device=device,
    )
    topk_for_hamming_full = torch.full(
        [max_num_reqs],
        fill_value=kvcomp_config.vllm_hash_attention_topk // block_size,
        dtype=torch.int32,
        device=device,
    )
    topk_for_hamming_full_cpu = torch.full(
        [max_num_reqs], fill_value=kvcomp_config.vllm_hash_attention_topk // block_size, dtype=torch.int32, device="cpu"
    )
    seq_lens_for_hamming = torch.zeros([max_num_reqs], dtype=torch.int32, device=device)
    hamming_output = torch.zeros(
        [
            max_num_reqs,
            vllm_config.model_config.get_num_kv_heads(parallel_config),
            cdiv(vllm_config.model_config.max_model_len, block_size),
        ],
        dtype=torch.int32,
        device=device,
    )
    seq_lens_for_reshape = torch.zeros([max_num_reqs], dtype=torch.int32, device=device)
    valid_query_mask = torch.empty((max_num_reqs,), dtype=torch.bool, device=device)
    seq_lens_from_hamming = torch.zeros([max_num_reqs], dtype=torch.int32, device="cpu")

    hash_encoder = HashEncoder(kvcomp_config.head_dim, kvcomp_config.hash_bits, dtype, device)
    hashk_caches: list[torch.Tensor | None] = []

    # Build and return KVCompMetaData object
    kvcomp_meta_data = KVCompMetaData(
        kvcomp_config=kvcomp_config,
        chunk_sizes_for_hamming_full=chunk_sizes_for_hamming_full,
        topk_for_hamming_full=topk_for_hamming_full,
        topk_for_hamming_full_cpu=topk_for_hamming_full_cpu,
        seq_lens_for_hamming=seq_lens_for_hamming,
        hamming_output=hamming_output,
        hash_encoder=hash_encoder,
        hashk_caches=hashk_caches,
        seq_lens_for_reshape=seq_lens_for_reshape,
        valid_query_mask=valid_query_mask,
        seq_lens_from_hamming=seq_lens_from_hamming,
        sink=1,
        recent=4,
    )

    return kvcomp_meta_data


def init_and_bind_hashk_cache(
    kv_caches: dict, num_attn_module: int, vllm_config, device: torch.device, compilation_config, kvcomp_meta_data
) -> None:
    """
    Wrapper function to initialize hashk cache and bind to forward context & model runner
    """
    # Initialize hashk cache dict (distinguish MLA/GQA mode)
    if vllm_config.model_config.use_mla:
        hashk_caches_nope: dict[str, torch.Tensor | None] = {}
        hashk_caches_rope: dict[str, torch.Tensor | None] = {}
    else:
        hashk_caches: dict[str, torch.Tensor | None] = {}

    # Iterate all layers' KV cache, initialize corresponding hashk cache
    for layer_name, kv_cache in kv_caches.items():
        # Extract layer index, check if it's rollback/skip layer
        layer_index = extract_layer_index(layer_name, num_attn_module)
        is_rollback_layer = layer_index in kvcomp_meta_data.kvcomp_config.vllm_hash_attention_rollback_layers
        is_skip_layer = layer_index in kvcomp_meta_data.kvcomp_config.vllm_hash_attention_skip_layers

        # Directly assign None for rollback/skip layers
        if is_rollback_layer or is_skip_layer:
            if vllm_config.model_config.use_mla:
                hashk_caches_nope[layer_name] = None
                hashk_caches_rope[layer_name] = None
            else:
                hashk_caches[layer_name] = None
        # Initialize corresponding hashk cache tensor for Hamming calculation layers
        else:
            if vllm_config.model_config.use_mla:
                # MLA mode: handle nope and rope hashk cache separately
                num_blocks_nope, block_size_nope, num_kv_heads_nope, head_size_nope = kv_cache[0].shape
                num_blocks_rope, block_size_rope, num_kv_heads_rope, head_size_rope = kv_cache[1].shape

                hashk_cache_nope = torch.zeros(
                    (num_blocks_nope, num_kv_heads_nope, block_size_nope, head_size_nope // 8),
                    dtype=torch.uint8,
                    device=device,
                )
                hashk_cache_rope = torch.zeros(
                    (num_blocks_rope, num_kv_heads_rope, block_size_rope, head_size_rope // 8 * 2),
                    dtype=torch.uint8,
                    device=device,
                )

                hashk_caches_nope[layer_name] = hashk_cache_nope
                hashk_caches_rope[layer_name] = hashk_cache_rope
            else:
                # GQA mode: initialize normal hashk cache
                num_blocks, block_size, num_kv_heads, head_size = kv_cache[0].shape
                hashk_cache = torch.zeros(
                    (num_blocks, num_kv_heads, block_size, head_size // 8), dtype=torch.uint8, device=device
                )
                hashk_caches[layer_name] = hashk_cache

    # Bind hashk cache to forward context and model runner
    if vllm_config.model_config.use_mla:
        bind_hashk_cache_nope(
            hashk_caches_nope,
            compilation_config.static_forward_context,
            kvcomp_meta_data.hashk_cache_nope,
            num_attn_module,
        )
        bind_hashk_cache_rope(
            hashk_caches_rope,
            compilation_config.static_forward_context,
            kvcomp_meta_data.hashk_cache_rope,
            num_attn_module,
        )
    else:
        bind_hashk_cache(
            hashk_caches, compilation_config.static_forward_context, kvcomp_meta_data.hashk_caches, num_attn_module
        )


def recover_request_lengths(cu_num_tokens: torch.Tensor) -> torch.Tensor:
    """
    Restore the original length of each request from the cumulative sum tensor cu_num_tokens
    """
    if cu_num_tokens.numel() == 0:
        return torch.tensor([], dtype=cu_num_tokens.dtype, device=cu_num_tokens.device)

    request_lengths = cu_num_tokens[1:] - cu_num_tokens[:-1]
    return request_lengths