custom_generate/generate.py

import torch
import types
import inspect
import importlib
import transformers
import torch.nn as nn
from transformers import Cache, GenerationConfig

from typing import Any, Dict, List, Optional, Tuple, Union
from transformers.modeling_utils import  PreTrainedModel
from transformers.processing_utils import Unpack
from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
from transformers import Cache, GenerationConfig


UNSUPPORTED_GENERATION_ARGS = [
    "cache_implementation",  # cache-related arguments, here we always use SepCache
    "cache_config",
    "return_legacy_cache",
    "num_beams",  # beam search (and cousin techniques) are not supported
    "compile_config",  # SepCache doesn't support torch.compile
    "assistant_model",  # it also doesn't support speculative decoding
]

##################################################### Functions to Patch #######################################################
def truncate_input_ids_4_autoregression(input_ids, key_states):
    if input_ids.shape[-1] != key_states.shape[-2]:
        assert input_ids.shape[-1] >= key_states.shape[-2]
        truncated_input_ids = input_ids[..., -key_states.shape[-2]: ]
        return truncated_input_ids
    else:
        return input_ids


def llama_atten_forward(
    self,
    hidden_states: torch.Tensor,
    position_embeddings: tuple[torch.Tensor, torch.Tensor],
    attention_mask: Optional[torch.Tensor],
    past_key_value: Optional[Cache] = None,
    cache_position: Optional[torch.LongTensor] = None,
    **kwargs: Unpack[FlashAttentionKwargs],
) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]:
    input_shape = hidden_states.shape[:-1]

    if hasattr(self, "head_dim"):
        head_dim = self.head_dim
    elif hasattr(self, "head_size"):
        head_dim = self.head_size

    hidden_shape = (*input_shape, -1, head_dim)

    query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
    key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
    value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)


    ###########################SepCache########################
    assert isinstance(past_key_value,  SepCache), f"`past_key_value` must be of the type: `SepCache`."
    APPLY_PE_SHIFT = past_key_value.APPLY_PE_SHIFT
    APPLY_PES_INSIDE = past_key_value.APPLY_PES_INSIDE    
    ###########################################################


    ########################Monkey Patching####################
    module = importlib.import_module(self.__module__)
                
    apply_rotary_pos_emb = module.apply_rotary_pos_emb
    rotate_half = module.rotate_half
    eager_attention_forward = module.eager_attention_forward
    ALL_ATTENTION_FUNCTIONS = module.ALL_ATTENTION_FUNCTIONS
    ###########################################################

    if not APPLY_PE_SHIFT:
        cos, sin = position_embeddings            
        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)

    if past_key_value is not None:
        # ##################################################Default#########################################################
        # sin and cos are specific to RoPE models; cache_position needed for the static cache
        # cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}            
        # key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
        # ##################################################################################################################

        ##################################################SepCache#########################################################
        # sin and cos are specific to RoPE models; position_ids needed for the static cache
        if APPLY_PE_SHIFT and (not APPLY_PES_INSIDE):
            ### At least the shifted `sin` and `cos` should be properly provided (not `None`).
            cache_kwargs = {"sin": sin, "cos": cos, "cos_q": cos_q, "sin_q": sin_q, "cache_position": cache_position, "partial_rotation_size": None }
        else:
            cache_kwargs = {}


        if "kwargs" in locals():
            pass
        elif "flash_attn_kwargs" in locals():
            kwargs = flash_attn_kwargs
        else:
            raise NameError("`kwargs` or `flash_attn_kwargs` should be given and they need to contain `sepllm_kwargs` (which contains `input_ids`) and `position_ids`.")

        if "input_ids" not in locals():
            if "input_ids" in kwargs:
                input_ids = kwargs.get("input_ids", None)
            else:
                sepllm_kwargs = kwargs.get("sepllm_kwargs", None)
                assert sepllm_kwargs is not None, f"`sepllm_kwargs` must be provided when `input_ids` is not given."
                input_ids = sepllm_kwargs.get("input_ids", None)
            
            assert input_ids is not None, f"`input_ids` must be properly provided directly or through `sepllm_kwargs` when calling `update()` in `SepCache`." 

        if "position_ids" not in locals():
            position_ids = kwargs.get("position_ids")            
                    
        assert input_ids is not None, f"`input_ids` must be properly provided when calling `update()` in `SepCache`." 
        bsz, q_len, _ = hidden_states.size()                                    

        input_ids = truncate_input_ids_4_autoregression(input_ids = input_ids, key_states = key_states )

        if APPLY_PE_SHIFT:
            key_states, value_states, query_states = past_key_value.update(                
                key_states = key_states,
                value_states = value_states,
                query_states = query_states,
                input_ids = input_ids,
                layer_idx = self.layer_idx,
                position_ids = position_ids,        
                PREFILLING_FLAG = q_len > 1, 
                cache_kwargs = cache_kwargs )

        else:
            key_states, value_states  =  past_key_value.update(
                key_states = key_states,
                value_states = value_states,
                input_ids = input_ids,
                layer_idx = self.layer_idx,
                position_ids = position_ids,
                PREFILLING_FLAG = q_len > 1, 
                cache_kwargs = cache_kwargs )
                                                        
        seq_len = past_key_value.get_usable_length(self.layer_idx)

        if attention_mask is not None:                
            attention_mask = attention_mask[..., :seq_len]
        ##################################################################################################################


    attention_interface: Callable = eager_attention_forward
    if self.config._attn_implementation != "eager":
        attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]

    attn_output, attn_weights = attention_interface(
        self,
        query_states,
        key_states,
        value_states,
        attention_mask,
        dropout=0.0 if not self.training else self.attention_dropout,
        scaling=self.scaling,
        **kwargs,
    )

    attn_output = attn_output.reshape(*input_shape, -1).contiguous()
    attn_output = self.o_proj(attn_output)
    return attn_output, attn_weights


def _validate_model_kwargs(self, model_kwargs: dict[str, Any]):
    """Validates model kwargs for generation. Generate argument typos will also be caught here."""
    # If a `Cache` instance is passed, checks whether the model is compatible with it
    if isinstance(model_kwargs.get("past_key_values", None), Cache) and not self._supports_cache_class:
        raise ValueError(
            f"{self.__class__.__name__} does not support an instance of `Cache` as `past_key_values`. Please "
            "check the model documentation for supported cache formats."
        )

    # Excludes arguments that are handled before calling any model function
    if self.config.is_encoder_decoder:
        for key in ["decoder_input_ids"]:
            model_kwargs.pop(key, None)

    unused_model_args = []
    model_args = set(inspect.signature(self.prepare_inputs_for_generation).parameters)
    # `kwargs`/`model_kwargs` is often used to handle optional forward pass inputs like `attention_mask`. If
    # `prepare_inputs_for_generation` doesn't accept them, then a stricter check can be made ;)
    if "kwargs" in model_args or "model_kwargs" in model_args:
        model_args |= set(inspect.signature(self.forward).parameters)

    # Encoder-Decoder models may also need Encoder arguments from `model_kwargs`
    if self.config.is_encoder_decoder:
        base_model = getattr(self, self.base_model_prefix, None)

        # allow encoder kwargs
        encoder = getattr(self, "encoder", None)
        # `MusicgenForConditionalGeneration` has `text_encoder` and `audio_encoder`.
        # Also, it has `base_model_prefix = "encoder_decoder"` but there is no `self.encoder_decoder`
        # TODO: A better way to handle this.
        if encoder is None and base_model is not None:
            encoder = getattr(base_model, "encoder", None)

        if encoder is not None:
            encoder_model_args = set(inspect.signature(encoder.forward).parameters)
            model_args |= encoder_model_args

        # allow decoder kwargs
        decoder = getattr(self, "decoder", None)
        if decoder is None and base_model is not None:
            decoder = getattr(base_model, "decoder", None)

        if decoder is not None:
            decoder_model_args = set(inspect.signature(decoder.forward).parameters)
            model_args |= {f"decoder_{x}" for x in decoder_model_args}

    for key, value in model_kwargs.items():
        # #############################Default###########################
        # if value is not None and key not in model_args:
        #     unused_model_args.append(key)
        # ###############################################################

        ###############################SepCache###########################
        if (value is not None) and (key not in model_args) and ("sep" not in str(key).lower()):
            unused_model_args.append(key)
        ###################################################################

    if unused_model_args:
        raise ValueError(
            f"The following `model_kwargs` are not used by the model: {unused_model_args} (note: typos in the"
            " generate arguments will also show up in this list)"
        )

#############################################################End################################################################




########################################################## SepCache ############################################################
class SepCache(Cache):
    """
    A cache as described in the [SepLLM paper - ICML 2025](https://arxiv.org/abs/2412.12094). In the training phase, 
    SepLLM condenses the segment information into the KV of the separator that divides the segment. In the inference phase, the 
    corresponding SepCache only needs to store the KVs of initial tokens, separator tokens, and recent tokens for generation.

    It stores the Key and Value states as lists of tensors, two lists for each layer. The expected shape for each tensor is
    `[batch_size, num_heads, seq_len, head_dim]`.

    Frequently-Used Parameters:

        `init_cache_size: Union[int, List]`:
            The maximum number of KVs to be stored for initial tokens.
            In the paper, the hyperparameter `a` is an abbreviated alias for `self.init_cache_size`.                
                
        `sep_cache_size: Union[int, List]`:
            The maximum number of KVs to be stored for separator tokens.
            In the paper, the hyperparameter `s` is an abbreviated alias for `self.sep_cache_size`.

        `local_size: Union[int, List]`: 
            The maximum number of KVs to be stored for local tokens (i.e., sliding window).
            In the paper, the hyperparameter `w` is an abbreviated alias for `self.local_size`.

        `cache_size: Union[int, List]`:    
            The maximum number of KVs to be stored for all the tokens, i.e., the size for the whole KV cache.  
            In the paper, the hyperparameter `c` is an abbreviated alias for `self.cache_size`.

        Concerning these four parameters above:
            When a list is passed (its length must be `layer_num`), it represents different values for each layer. 
            When an integer is passed, it means the setting is the same for all layers.
        
        
        `USE_MAX_SEP_CACHE: bool`: 
            If True, it means we only keep at most `self.sep_cache_size` seperators' KVs.  
            If the number exceeds this limit, older separator's KVs will be discarded, keeping only the most recent `self.sep_cache_size` KVs. 
            In the paper, the hyperparameter `s` is an abbreviated alias for `self.sep_cache_size`.
          
        `separator_token_ids: List[int]`:
            The token ids of the separator tokens for the current model's tokenizer.            
            We have some examples, such as the Llama-3 series models, where setting `model_type='llama'` allows you 
                to skip setting `separator_token_ids` and `PADDING_ID` (SepCache will auto-fill them).

        `PADDING_ID: int`:
            The token id of the padding token. You can just set `PADDING_ID` to the id of "<|endoftext|>" token of the tokenizer for the pretrained model.

    Important Note:
        When `cache_size` and `local_size` are set to infinity (i.e., sufficiently large positive integers), and `USE_MAX_SEP_CACHE` is `False`, `SepCache` degenerates into a regular Cache. 
        However, you must always ensure that `init_cache_size` + `sep_cache_size` + `local_size` + `left_padding_offset` < `cache_size`. 
        Here, `left_padding_offset` denotes the number of padding tokens in the record with the largest left paddings within a runtime batch. `left_padding_offset` can only be determined at runtime.        
        To guarantee the above inequality always holds during runtime, when setting, you can intentionally create a sufficient margin between both sides of the following inequality:
            `init_cache_size` + `sep_cache_size` + `local_size`  < `cache_size`, i.e., `a`+`s`+`w`<`c` in the [SepLLM paper - ICML 2025]
            to leave room for `left_padding_offset`.

        Please refer to the `__init__` function's comments for more details on the parameters.
            
    Example:

        ```python        
        >>> from transformers import AutoTokenizer, AutoModelForCausalLM, SepCache
        >>> import torch
        >>> from huggingface_hub import login
        >>> login("hf_xxxXXXxxx")


        >>> def to_cuda(a_dict: dict) -> dict:
        >>>    new_dict = {}    
        >>>    for k,v in a_dict.items():
        >>>        if isinstance(v, torch.Tensor):
        >>>            new_dict[k] = v.cuda()
        >>>        else:
        >>>            new_dict[k] = v
        >>>    return new_dict

        >>> model = AutoModelForCausalLM.from_pretrained("meta-llama/Meta-Llama-3-8B-Instruct", attn_implementation="flash_attention_2", device_map="cuda:0")
        >>> model.bfloat16().cuda()
        >>> tokenizer = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3-8B-Instruct")
        >>> inputs = tokenizer(text="My name is Llama 3", return_tensors="pt")
        >>> inputs = to_cuda(inputs)
        >>> # Prepare a cache and pass it to model's forward; `layer_num` is the number of layers for the pretrained model.
        >>> past_key_values = SepCache(init_cache_size=4, sep_cache_size=128, local_size=256, cache_size=512, layer_num=32, USE_MAX_SEP_CACHE=True, model_type='llama')
        >>> # `separator_token_ids` and `PADDING_ID` must also be provided if you are not using `model_type='llama'` like this demo.
        >>> outputs = model(**inputs, past_key_values=past_key_values, use_cache=True)
        >>> outputs.past_key_values # access SepCache filled with keys/values
        SepCache()
        ```

        ```python
        >>> ## When using the `update` function of SepCache to update the keys/values and the past token ids (necessary in SepCache), the current `input_ids` must also be provided.        
        >>> key_states, value_states = past_key_values.update(                
                    key_states = key_states,
                    value_states = value_states,    
                    input_ids = input_ids,
                    layer_idx = layer_idx,     
                    PREFILLING_FLAG = q_len > 1, ## `q_len` is the sequence length of the current `query_states`
                    )

        ```
        For detailed usage instructions, please refer to https://github.com/HKUDS/SepLLM
    """
    # is_sliding = True
    
    @staticmethod
    def slice_on_1d(x, start, end):
        return x[:, start:end, ...]
    @staticmethod
    def slice_on_2d(x, start, end):
        return x[:, :, start:end, ...]
    @staticmethod
    def slice_on_3d(x, start, end):
        return x[:, :, :, start:end, ...]


    @staticmethod
    def sep_1bat_select_on_1d(x, Bid, sep_index, min_sep_num=None, max_sep_num=None, SEP_PADDING_IN_BATCH=True):    
        """
        For the record with index `Bid` in a batch, extract the K/V states corresponding to the separators on dimension 1. 
           If `SEP_PADDING_IN_BATCH=True`, pad to the longest length (i.e. `max_sep_num`); 
           otherwise, truncate to the shortest length (i.e. `min_sep_num`). 
        """
        sep_index = sep_index.to(x.device)

        if SEP_PADDING_IN_BATCH: ## Need padding
            assert max_sep_num is not None, f"if `SEP_PADDING_IN_BATCH=True`, `max_sep_num` should not be None"
            new_x_sep =  x[Bid, sep_index, ...]   # # batch x seqlen x head x dim  -->  sep_num x head x dim  
            padding_num = max_sep_num -  new_x_sep.shape[0]
            if padding_num > 0 :
                assert padding_num <= x.shape[1], f"`padding_num` should be <= `x.shape[1]`, i.e.  x's seqlen"
                new_x_pad = x[Bid, -padding_num: , ...]    #  padding_num x head x dim     
                return torch.cat([new_x_sep, new_x_pad ] , dim=0) # max_sep_num x head x dim 
            else:
                return new_x_sep #  max_sep_num x head x dim 

        if min_sep_num is None:
            return x[Bid, sep_index, ...]  # # batch x seqlen x head x dim -->  sep_num x head x dim    
        else: ## `min_sep_num` is provided. Need truncation
            new_x =  x[Bid, sep_index, ...]   # # batch x seqlen x head x dim -->  sep_num x head x dim               
            return new_x[ :min_sep_num, ...] # #  min_sep_num x head x dim      


    @staticmethod
    def sep_1bat_select_on_2d(x, Bid, sep_index, min_sep_num=None, max_sep_num=None, SEP_PADDING_IN_BATCH=True):    
        """
        For the record with index `Bid` in a batch, extract the K/V states corresponding to the separators on dimension 2. 
           If `SEP_PADDING_IN_BATCH=True`, pad to the longest length (i.e. `max_sep_num`); 
           otherwise, truncate to the shortest length (i.e. `min_sep_num`). 
        """
        sep_index = sep_index.to(x.device)

        if SEP_PADDING_IN_BATCH: ## Need padding
            assert max_sep_num is not None, f"if `SEP_PADDING_IN_BATCH=True`, `max_sep_num` should not be None"
            new_x_sep =  x[Bid, :, sep_index, ...]   # # batch x head x seqlen x dim -->  head x sep_num x dim  
            padding_num = max_sep_num -  new_x_sep.shape[-2]
            if padding_num > 0 :
                assert padding_num<= x.shape[-2], f"`padding_num` should be <= `x.shape[-2]`, i.e.  x's seqlen"
                new_x_pad = x[Bid, :, -padding_num: , ...]    # head x padding_num x dim     
                return torch.cat([new_x_sep, new_x_pad ] , dim=-2) # head x max_sep_num x dim 
            else:
                return new_x_sep # head x max_sep_num x dim 

        if min_sep_num is None:
            return x[Bid, :, sep_index, ...]  # # batch x head x seqlen x dim -->  head x sep_num x dim    
        else: ## `min_sep_num` is provided. Need truncation
            new_x =  x[Bid, :, sep_index, ...]   # # batch x head x seqlen x dim -->  head x sep_num x dim            
            return new_x[:, :min_sep_num, ...] # #  head x min_sep_num x dim      


    @staticmethod
    def sep_1bat_select_on_3d(x, Bid, sep_index, min_sep_num=None, max_sep_num=None, SEP_PADDING_IN_BATCH=True):    
        """
        For the record with index `Bid` in a batch, extract the K/V states corresponding to the separators on dimension 3. 
           If `SEP_PADDING_IN_BATCH=True`, pad to the longest length (i.e. `max_sep_num`); 
           otherwise, truncate to the shortest length (i.e. `min_sep_num`). 
        """        
        sep_index = sep_index.to(x.device)

        if SEP_PADDING_IN_BATCH: ## Need padding
            assert max_sep_num is not None, f"if `SEP_PADDING_IN_BATCH=True`, `max_sep_num` should not be None"
            new_x_sep =  x[Bid, :, :, sep_index, ...]   # # batch x head x dim x seqlen  -->  head x dim x sep_num 
            padding_num = max_sep_num -  new_x_sep.shape[-1]
            if padding_num > 0 :
                assert padding_num <= x.shape[-1], f"`padding_num` should be <= `x.shape[-1]`, i.e.  x's seqlen"
                new_x_pad = x[Bid, :, :, -padding_num:, ...]    # head x dim x padding_num     
                return torch.cat([new_x_sep, new_x_pad] , dim=-1) # head x dim x max_sep_num 
            else:
                return new_x_sep # head x dim x max_sep_num 

        if min_sep_num is None:
            return x[Bid, :, :, sep_index, ...]  # # batch x head x dim x seqlen -->  head x dim x sep_num    
        else: ## `min_sep_num` is provided. Need truncation
            new_x =  x[Bid, :, :, sep_index, ...]   # # batch x head x dim x seqlen -->  head x dim x sep_num          
            return new_x[:, :, :min_sep_num, ...] # #  head x dim x min_sep_num       

    DIM_TO_SLICE = {
        1: slice_on_1d,
        2: slice_on_2d,
        3: slice_on_3d,
    }
    
    BAT_DIM_TO_SELECT = {
        1: sep_1bat_select_on_1d,
        2: sep_1bat_select_on_2d,
        3: sep_1bat_select_on_3d,
    }

    def __init__(self,                                                
                ## For SepLLM                                
                init_cache_size: Union[int, List] = 4,        
                sep_cache_size: Union[int, List] = 64,
                local_size: Union[int, List]=256, 
                cache_size: Union[int, List]=512,    
                SEP_ACCUMULATION: bool = True,
                USE_MAX_SEP_CACHE: bool = False,
                SEP_PADDING_IN_BATCH: bool = False,
                separator_token_ids: List[int] = None, ## required for initialization if `model_type` is not provided.
                PADDING_ID: int = None, ## required for initialization if `model_type` is not provided.

                ## For inheritance & initialization states
                past_tok_ids: List[torch.Tensor] = None,  ## It saves all the token ids corresponding to the saved KVs for all layers in SepCache.                
                key_cache: List[torch.Tensor] = None,          
                value_cache: List[torch.Tensor] = None,

                ## For debugging
                PRINT_KV_RATIO_INSIDE: bool = False,
                print_KV_inside_per_steps: int = 1000,   
                _seen_tokens: int = 0, 
                _kept_kv_ratio: List[Tuple[int]] = None,
                
                ### For positional encoding shifting
                APPLY_PE_SHIFT: bool = False,
                APPLY_PES_INSIDE: bool = True,
                _shifted_position_ids:  List[torch.Tensor] = None,
                _rope_unsqueeze_dim: int = 1, ## The unsqueeze_dim when applying RoPE.
                _rope_seq_dim: int=1, ## The seq_len dimension for the `cos` or `sin` tensors.
                pe_scaling_factor:float = 1.0,
                pe_dim:int=128, ## The number of dims for positional encoding. Typically, just set the `head_dim` to this.
                max_position_embeddings: int = 8192, 
                base: int=10000,  ## The base for RoPE.               
                
                ## For basic transformer architecture
                k_seq_dim: int=2, ## The dimension for seq_len in key tensors
                v_seq_dim: int=2, ## The dimension for seq_len in value tensors
                layer_num: int = None, ## required for initialization

                model_type: str = None,  ## The model type for running the example. choose from ['llama', 'pythia','falcon'].
                device = None          
                 ) -> None:
        """        
        `SEP_ACCUMULATION`: If True, it means we will try to accumulate all the KVs for seperators. If False, only the `new_sep_kv` compressed from the `past_win_kv` will be kept (see function `compress_kv_cache_and_tokids_layer_wise`).
                                                             
        `USE_MAX_SEP_CACHE`: If True, it means we only keep at most `self.sep_cache_size` seperators' KVs.  If the number exceeds this limit, older separator's KVs will be discarded, keeping only the most recent `self.sep_cache_size` KVs. In the paper, the hyperparameter `s` is an abbreviated alias for `self.sep_cache_size`.

        `SEP_PADDING_IN_BATCH`: If True, it means that SepCache will pad separator tokens in other records to be aligned with the record with the most separators in a batch. If False, it means that SepCache will truncate older separator tokens in other records to be aligned with the record with the fewest separators in a batch.
        
        Note: If `SEP_ACCUMULATION=True` and `USE_MAX_SEP_CACHE=False`, as the number of input tokens increases, the number of separators in the KV cache will also accumulate endlessly 
              and `self.cache_size` will also be infinitely expanded (no longer fixed).

              When `SEP_PADDING_IN_BATCH=True` is used in combination with `USE_MAX_SEP_CACHE=False` and `SEP_ACCUMULATION=True`, the KV cache will accumulate indefinitely, 
              and since `SEP_PADDING_IN_BATCH=True`, the KVs of all separators will be retained (rather than being truncated).


        For detailed usage instructions, please refer to https://github.com/HKUDS/SepLLM
        """    

        super().__init__()               
        if (key_cache is not None) or (value_cache is not None) or (past_tok_ids is not None):
            assert isinstance(key_cache, list)
            assert isinstance(value_cache, list)
            assert isinstance(past_tok_ids, list), f"For SepCache, if `key_cache` and `value_cache` are given (e.g., provided from legacy `past_key_values`), `past_tok_ids` corresponding to `key_cache` and `value_cache` must also be provided to initialize SepCache."

            assert len(key_cache) == len(past_tok_ids), f"The length of `key_cache` ({len(key_cache)}) should be equal to that of `past_tok_ids` ({len(past_tok_ids)})."
            assert len(value_cache) == len(past_tok_ids), f"The length of `value_cache` ({len(value_cache)}) should be equal to that of `past_tok_ids` ({len(past_tok_ids)})."
        assert layer_num is not None, f"`layer_num` must be provided according to the pretrained model."

        ## For basic parameters & states    
        self.key_cache: List[torch.Tensor] = key_cache if key_cache is not None else []
        self.value_cache: List[torch.Tensor] = value_cache if value_cache is not None else []    

        self.k_seq_dim = k_seq_dim ## The dimension for the seq_len in key states. Typically, 2.
        self.v_seq_dim = v_seq_dim ## The dimension for the seq_len in value states. Typically, 2.

        self.k_slice = self.DIM_TO_SLICE[k_seq_dim]
        self.v_slice = self.DIM_TO_SLICE[v_seq_dim]
        
        self.k_bat_dim_select = self.BAT_DIM_TO_SELECT[k_seq_dim]
        self.v_bat_dim_select = self.BAT_DIM_TO_SELECT[v_seq_dim]
        self._seen_tokens: int = _seen_tokens  # Used in `generate` to keep tally of how many tokens the cache has seen as well as performing statistics.
        self.layer_num =  layer_num
        self.device = device # Deprecated


        ## For debugging
        self.PRINT_KV_RATIO_INSIDE = PRINT_KV_RATIO_INSIDE
        self.print_KV_inside_per_steps = print_KV_inside_per_steps
        self._print_kv_ratio_count = 0
        self._kept_kv_ratio: List[Tuple[int]] = _kept_kv_ratio if _kept_kv_ratio is not None else []   

        ## For Streaming SepLLM
        self.past_tok_ids: List[torch.Tensor] = past_tok_ids if past_tok_ids is not None else []  ## It saves all the token ids corresponding to the saved KVs for all layers in SepCache      
        self.left_padding_offset = None
        self._set_layer_wise_attribute("init_cache_size", init_cache_size, layer_num)
        self._set_layer_wise_attribute("local_size", local_size, layer_num)
        self._set_layer_wise_attribute("cache_size", cache_size, layer_num)
        self._set_layer_wise_attribute("sep_cache_size", sep_cache_size, layer_num)
        self._set_layer_wise_attribute("sep_exrange", 0, layer_num) # runtime right boundary for separators, excluded
        self._set_layer_wise_attribute("max_sep_exidx", self._list_element_add(self.sep_cache_size, self.init_cache_size), layer_num) # max right boundary for separators, excluded
        self.SEP_ACCUMULATION = SEP_ACCUMULATION
        self.USE_MAX_SEP_CACHE = USE_MAX_SEP_CACHE
        self.SEP_PADDING_IN_BATCH = SEP_PADDING_IN_BATCH
        

        ### For positional encoding shifting
        self.APPLY_PE_SHIFT = APPLY_PE_SHIFT
        self.APPLY_PES_INSIDE = APPLY_PES_INSIDE

        self.cos_sin_rerotation_cache = {}
        self._cos_cache = None
        self._sin_cache = None        
        self._shifted_position_ids: List[torch.Tensor] = _shifted_position_ids if _shifted_position_ids is not None else []        
        self._rope_unsqueeze_dim = _rope_unsqueeze_dim
        self._rope_seq_dim = _rope_seq_dim        

        self.pe_dim = pe_dim
        self.max_position_embeddings = max_position_embeddings
        self.base = base
        inv_freq = 1.0 / (self.base ** (torch.arange(0, self.pe_dim, 2, dtype=torch.int64).float().to(device) / self.pe_dim))
        self.inv_freq = inv_freq
        self.pe_scaling_factor = pe_scaling_factor
        self._sin_cached = None
        self._cos_cached = None

        if model_type is None:
            assert isinstance(separator_token_ids, list), f"`separator_token_ids: List[int]` must be correctly provided for initialization unless `model_type` is properly given, which will auto-fiil `separator_token_ids`."
            assert len(separator_token_ids) > 0, f"`separator_token_ids: List[int]` should NOT be empty."
            for i in range(len(separator_token_ids)):
                assert isinstance(separator_token_ids[i], int), f"The ids in `separator_token_ids` must be of the type `int`."  
            assert isinstance(PADDING_ID, int), f"`PADDING_ID: int` must be correctly provided for initialization unless `model_type` is properly given, which will auto-fiil `PADDING_ID`."
            self.separator_token_ids = separator_token_ids
            self.PADDING_ID = PADDING_ID                               
        else:
            assert isinstance(model_type, str), f"`model_type` should be a `str` or `None`."
            if 'llama' in  model_type.lower():
                # print("Debug: For Llama's default separators")
                self.separator_token_ids = [128000, 13, 11, 30, 0, 26, 25, 198, 220, 662, 1174, 949, 758, 2652, 551, 720, 256,262] # llama3 8b
                self.PADDING_ID = 128009
            elif ( 'pythia' in model_type.lower() ) or ( 'gpt_neox' in model_type.lower() ):
                # print("Debug: For GPTNeox's default separators")
                self.separator_token_ids = [15, 13, 32, 2, 28, 27, 209, 186, 187,    964, 1157, 3736, 2195, 3706, 1163, 2490,  50276,    586, 4928, 50275 ]       # pythia 14b
                self.PADDING_ID = 0
            elif 'falcon' in model_type.lower():
                # print(f"Debug: For Falcon's default separators")
                self.separator_token_ids = [25, 23,  42, 12, 38, 37, 193,  4610,  204, 258, 1212, 23787, 466 ]       # falcon-40b
                self.PADDING_ID = 11
            else:
                raise NotImplementedError(f"NOT implemented for the tokenizer of the backbone model type: `{model_type}`. You must provide `separator_token_ids: List[int]` and `PADDING_ID: int` for initialization in this case! ")
        
        if APPLY_PE_SHIFT:
            print(">>>>>>>>---------#####################################################################################-----------<<<<<<<<")
            print(">>>>>>>>---------                                                                                     -----------<<<<<<<<")
            print(">>>>>>>>---------  Warning: When `APPLY_PE_SHIFT=True`, SepCache must store the key/value states       ----------<<<<<<<<")
            print(">>>>>>>>---------              before applying positional encoding (specifically RoPE)                -----------<<<<<<<<")
            print(">>>>>>>>---------#####################################################################################-----------<<<<<<<<")
                
        if APPLY_PES_INSIDE:
            print(">>>>>>>>---------#####################################################################################-----------<<<<<<<<")
            print(">>>>>>>>---------                                                                                     -----------<<<<<<<<")
            print(">>>>>>>>---------  Warning: When `APPLY_PES_INSIDE=True`, there is no need to apply rotary positional embedding--<<<<<<<<")
            print(">>>>>>>>---------  within the self_attention function, as this operation will be handled inside the `update`  ---<<<<<<<<")
            print(">>>>>>>>---------  function of SepCache. Note that `APPLY_PES_INSIDE=True` is typically used together with     ---<<<<<<<<")
            print(">>>>>>>>---------  `APPLY_PE_SHIFT=True`.                                                                     ---<<<<<<<<")
            print(">>>>>>>>---------#####################################################################################-----------<<<<<<<<")                            
            

    def update(
        self,
        key_states: torch.Tensor,
        value_states: torch.Tensor,        
        layer_idx: int,        
        input_ids: torch.Tensor = None,
        PREFILLING_FLAG: bool = True,
        query_states: Optional[torch.Tensor] = None,        
        position_ids: Optional[torch.Tensor]=None,                
        cache_kwargs: Optional[Dict[str, Any]] = None,
    ) -> Union[Tuple[torch.Tensor, torch.Tensor],Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]:
        """
        Updates the cache with the new `key_states` and `value_states` for the layer `layer_idx`.

        Parameters:        
            `key_states` (`torch.Tensor`):
                The new key states to cache.
            `value_states` (`torch.Tensor`):
                The new value states to cache.
            `input_ids` (`torch.Tensor`)
                The ids of the input tokens (context tokens or autoregressive tokens)                
            `layer_idx` (`int`):
                The index of the layer to cache the states for.
            `PREFILLING_FLAG` (`bool`)
                It should be `True` at pre-filling phase and `False` when decoding

            `query_states` (`Optional[torch.Tensor]`)
                The query states that need positional encoding shifting. Only useful when `self.APPLY_PE_SHIFT=True`
            `position_ids` (`Optional[torch.Tensor]`)
                The original positional ids of the tokens in the input sequence (i.e., indices of positions of each input sequence tokens in the position embeddings)
                Only useful when `self.APPLY_PE_SHIFT=True`, i.e., SepCache will utilize `position_ids` to calculate positional shifting.
            `cache_kwargs` (`Dict[str, Any]`, optional):
                Additional arguments for the cache update. The following arguments can be used in `SepCache`: `sin`,
                `cos`, `sin_q`, `cos_q`, `shifted_pos_ids` and `partial_rotation_size`. These arguments are used with models using RoPE, to recompute the
                rotation as the tokens are shifted. (These are only useful when `self.APPLY_PE_SHIFT=True`)

                Only useful when `self.APPLY_PE_SHIFT=True` and `self.APPLY_PES_INSIDE=False`:
                    `cos` and `sin` are the shifted rotation matrices for key states
                    `cos_q` and `sin_q` are the shifted rotation matrices for query states
                    `shifted_pos_ids` is the shifted positional ids for key states
                    
                When `self.APPLY_PE_SHIFT=True` and `self.APPLY_PES_INSIDE=True`:
                    SepCache will utilize `position_ids` to calculate positional shifting.
                
                `partial_rotation_size` means that `partial_rotation_size` slices along certain dimension need to be shifted (i.e., [0, 1, ..., `partial_rotation_size-1`] slices along certain dimension)

        Return:
            A tuple containing the updated key, value, and query states (query states are optional: only applicable when `self.APPLY_PE_SHIFT=True`).

        For detailed usage instructions, please refer to https://github.com/HKUDS/SepLLM
        """

        APPLY_PE_SHIFT = self.APPLY_PE_SHIFT
        APPLY_PES_INSIDE = self.APPLY_PES_INSIDE
        SEP_ACCUMULATION = self.SEP_ACCUMULATION
        USE_MAX_SEP_CACHE = self.USE_MAX_SEP_CACHE
        SEP_PADDING_IN_BATCH = self.SEP_PADDING_IN_BATCH
        
        if input_ids is None:
            input_ids = cache_kwargs.get("input_ids", None)
        assert input_ids is not None, f"`input_ids` must be properly provided when calling `update()` in `SepCache`."

        assert (self.APPLY_PE_SHIFT and (query_states is not None)) or not APPLY_PE_SHIFT, f"If `APPLY_PE_SHIFT=True`, `query_states` should be provided and it will be updated and returned"
                
        # Update the number of seen tokens
        if layer_idx == 0:
            assert key_states.shape[-2] == input_ids.shape[-1], f"`key_states.shape[-2]` ({key_states.shape[-2]}) should be equal to `input_ids.shape[-1]` ({input_ids.shape[-1]})."
            self._seen_tokens += input_ids.shape[-1]

        # [bsz, num_heads, seq_len, head_dim]
        new_kv_pair = (key_states, value_states)
                
        if (key_states.shape[self.k_seq_dim] + self.get_usable_length(layer_idx) < self.cache_size[layer_idx]) or PREFILLING_FLAG:  ## For prefilling
            assert  (PREFILLING_FLAG and key_states.shape[self.k_seq_dim] >= 1)  or (not PREFILLING_FLAG and key_states.shape[self.k_seq_dim] == 1)

            # Update cache and past token ids                
            self.update_kv_cache_and_past_tok_ids(new_kv_pair, input_ids, layer_idx, COMPRESS_KV=False, SEP_ACCUMULATION=SEP_ACCUMULATION, USE_MAX_SEP_CACHE=USE_MAX_SEP_CACHE, SEP_PADDING_IN_BATCH=SEP_PADDING_IN_BATCH)
            
            if APPLY_PE_SHIFT:                     
                shifted_keys, shifted_queries = self.apply_shifted_pos_emb(layer_idx, APPLY_PES_INSIDE, PREFILLING_FLAG, key_states, query_states, position_ids, cache_kwargs ) 
                query_states  = shifted_queries
                self.set_kv_cache( (shifted_keys, self.value_cache[layer_idx]), layer_idx)
            
            if PREFILLING_FLAG and layer_idx == 0:
                self.left_padding_offset = self.get_initial_pos_offset(layer_idx)

            ## Count KV usage
            kv_len_ori = self.get_seq_length(layer_idx)
            kv_len_cmp = self.get_usable_length(layer_idx)
            self._update_kv_ratio(kv_len_cmp=kv_len_cmp, kv_len_ori=kv_len_ori, layer_idx=layer_idx)

        else:
            ## Update the KV cache, count KV usage, and compress the KV cache if necessary                        
            kv_len_ori = self.get_seq_length(layer_idx)
            offset_init_size_layer = self.update_kv_cache_and_past_tok_ids(new_kv_pair, input_ids, layer_idx, COMPRESS_KV=True, SEP_ACCUMULATION=SEP_ACCUMULATION, USE_MAX_SEP_CACHE=USE_MAX_SEP_CACHE, SEP_PADDING_IN_BATCH=SEP_PADDING_IN_BATCH)
            kv_len_cmp = self.get_usable_length(layer_idx)
            self._update_kv_ratio(kv_len_cmp=kv_len_cmp, kv_len_ori=kv_len_ori, layer_idx=layer_idx)
                        
            if APPLY_PE_SHIFT:                
                shifted_keys, shifted_queries = self.apply_shifted_pos_emb(layer_idx, APPLY_PES_INSIDE, PREFILLING_FLAG, key_states, query_states, position_ids, cache_kwargs )                 
                query_states  = shifted_queries
                self.set_kv_cache( (shifted_keys, self.value_cache[layer_idx]), layer_idx)
            
        if self.PRINT_KV_RATIO_INSIDE:    
            self._print_kv_ratio(layer_idx)

        if query_states is not None:
            return self.key_cache[layer_idx], self.value_cache[layer_idx], query_states
        else:
            return self.key_cache[layer_idx], self.value_cache[layer_idx]
            
    
    def update_kv_cache_and_past_tok_ids(self, new_kv_pair: Tuple[torch.Tensor], input_ids: torch.Tensor, layer_idx: int, COMPRESS_KV=False, SEP_ACCUMULATION:bool=True, USE_MAX_SEP_CACHE:bool=False, SEP_PADDING_IN_BATCH:bool=True) -> None:
        """Update the KV cache and past token ids; compress the KV cache if necessary."""
        assert layer_idx is not None, f"`layer_idx` must be given"
        assert len(new_kv_pair) == 2, f"The length of `new_kv_pair` must be 2."
        assert len(self.key_cache) == len(self.value_cache), f"The layer numbers of stored `self.key_cache` and `self.value_cache` must be the same."

        self.append_past_tok_ids(input_ids, layer_idx)

        key, value = new_kv_pair
                
        if len(self.key_cache) <= layer_idx:
            self.key_cache.append(key)                        
            self.value_cache.append(value)
            assert len(self.key_cache) - 1  == layer_idx, f"The key_cache should be updated sequentially according to the layer numbering."              
            assert len(self.value_cache) - 1  == layer_idx, f"The value_cache should be updated sequentially according to the layer numbering."      
        else:            
            self.key_cache[layer_idx] = torch.cat([self.key_cache[layer_idx] , key], dim=self.k_seq_dim)
            self.value_cache[layer_idx] = torch.cat([self.value_cache[layer_idx] , value], dim=self.v_seq_dim)

        assert len(self.key_cache) == len(self.value_cache), f"The layer numbers of stored key_cache and value_cache must be the same."
        assert self.key_cache[layer_idx].shape[self.k_seq_dim] == self.value_cache[layer_idx].shape[self.v_seq_dim], "The seq length for key_cache and value_cache must be the same."

        if COMPRESS_KV:
            cmp_past_kv_pairs, cmp_past_tok_ids, offset_init_size_layer = self.compress_kv_cache_and_tokids_layer_wise((self.key_cache[layer_idx], self.value_cache[layer_idx]), layer_idx ,SEP_ACCUMULATION=SEP_ACCUMULATION, USE_MAX_SEP_CACHE=USE_MAX_SEP_CACHE, SEP_PADDING_IN_BATCH=SEP_PADDING_IN_BATCH )
            self.set_kv_cache(cmp_past_kv_pairs, layer_idx)
            self.set_past_tok_ids(cmp_past_tok_ids, layer_idx)            
            return offset_init_size_layer
        

    def append_past_tok_ids(self, input_ids: torch.Tensor, layer_idx: int) -> None:
        """Naively append the new `input_ids` to `self.past_tok_ids[layer_idx]`"""    
        assert layer_idx is not None, f"`layer_idx` must be given"
        
        if len(self.past_tok_ids) <= layer_idx:                        
            self.past_tok_ids.append(input_ids)
            assert len(self.past_tok_ids) - 1  == layer_idx, f"The past_tok_ids should be updated sequentially according to the layer numbering."                        
        else:      
            self.past_tok_ids[layer_idx] = torch.cat([self.past_tok_ids[layer_idx] , input_ids], dim=-1)


    def compress_kv_cache_and_tokids_layer_wise(self, past_kv_pairs, layer_idx:int ,SEP_ACCUMULATION=False, USE_MAX_SEP_CACHE=False, SEP_PADDING_IN_BATCH=True ):
        """        
        `SEP_ACCUMULATION`: If True, it means we will try to accumulate all the KVs for seperators. If False, only the `new_sep_kv` compressed from the `past_win_kv` will be kept (see function `compress_kv_cache_and_tokids_layer_wise`).
                                                             
        `USE_MAX_SEP_CACHE`: If True, it means we only keep at most `self.sep_cache_size` seperators' KVs.  If the number exceeds this limit, older separator's KVs will be discarded, keeping only the most recent `self.sep_cache_size` KVs. In the paper, the hyperparameter `s` is an abbreviated alias for `self.sep_cache_size`.

        `SEP_PADDING_IN_BATCH`: If True, it means that SepCache will pad separator tokens in other records to be aligned with the record with the most separators in a batch. If False, it means that SepCache will truncate older separator tokens in other records to be aligned with the record with the fewest separators in a batch.
        
        Note: If `SEP_ACCUMULATION=True` and `USE_MAX_SEP_CACHE=False`, as the number of input tokens increases, the number of separators in the KV cache will also accumulate endlessly 
              and `self.cache_size` will also be infinitely expanded (no longer fixed).

              When `SEP_PADDING_IN_BATCH=True` is used in combination with `USE_MAX_SEP_CACHE=False` and `SEP_ACCUMULATION=True`, the KV cache will accumulate indefinitely, 
              and since `SEP_PADDING_IN_BATCH=True`, the KVs of all separators will be retained (rather than being truncated).


        For detailed usage instructions, please refer to https://github.com/HKUDS/SepLLM
        """    

        key, value = past_kv_pairs
        seq_len = key.size(self.k_seq_dim)
        assert seq_len == self.get_usable_length(layer_idx), f"The seq_len of cached past key and value states should be the same as the return of `get_usable_length()`, which is {self.get_usable_length(layer_idx)}"

        
        left_padding_offset =  self.left_padding_offset        
        assert left_padding_offset is not None
        offset_init_size_layer = self.init_cache_size[layer_idx] + left_padding_offset
        self._set_layer_wise_attribute("max_sep_exidx", self._list_element_add(self.sep_cache_size, self.init_cache_size, bias=left_padding_offset), self.layer_num)
        self._CHECK_PARAMS_VALIDITY(layer_idx, left_padding_offset)

        if self.sep_exrange[layer_idx] <=0:            
            self.sep_exrange[layer_idx] = offset_init_size_layer

        assert seq_len - self.local_size[layer_idx] > self.sep_exrange[layer_idx]
        
        if offset_init_size_layer > 0:                                                       
            initial_kv, initial_tokids =  self.slice_kv_cache_and_tokids( past_kv_pairs, self.past_tok_ids[layer_idx], 0, offset_init_size_layer, seq_len=seq_len, _CHECK_IDX=True )        

        Before_First_Time_Compress_Flag = (self.sep_exrange[layer_idx] == offset_init_size_layer)  ## If true, it means the present timestamp is before t1: the 1st time to compress the past window, in which only seperators' kv are kept.
        
        if SEP_ACCUMULATION and not Before_First_Time_Compress_Flag: ## To get the old sep kv and sep token ids.           
            past_sep_kv, past_sep_tokids =  self.slice_kv_cache_and_tokids( past_kv_pairs, self.past_tok_ids[layer_idx], offset_init_size_layer, self.sep_exrange[layer_idx], seq_len=seq_len, _CHECK_IDX=True )            
        
        past_win_kv, past_win_tokids =  self.slice_kv_cache_and_tokids( past_kv_pairs, self.past_tok_ids[layer_idx], self.sep_exrange[layer_idx], seq_len - self.local_size[layer_idx], seq_len=seq_len, _CHECK_IDX=True )        
        
        
        local_kv, local_tokids  =  self.slice_kv_cache_and_tokids( past_kv_pairs, self.past_tok_ids[layer_idx], seq_len - self.local_size[layer_idx], seq_len, seq_len=seq_len, _CHECK_IDX=True )
        
        new_sep_kv, new_sep_tokids, min_sep_num, max_sep_num = self.compress_past_win_2_seps( past_win_kv, past_win_tokids, SEP_PADDING_IN_BATCH = SEP_PADDING_IN_BATCH ) ## To get the new sep kv and sep token ids that were just compressed from the past window
        
        if SEP_ACCUMULATION and not Before_First_Time_Compress_Flag:  ## Try to accumulate all the seen seps           
            sep_kv, sep_tokids  = self.cat_kv_cache_and_tokids( [ past_sep_kv, new_sep_kv ] ,  [past_sep_tokids, new_sep_tokids ] )                
            new_sep_len = new_sep_tokids.shape[-1]
            sep_len = sep_tokids.shape[-1]  
        else: ## Only keep the newly obtained kv (those just compressed from the past window)
            sep_kv, sep_tokids = new_sep_kv, new_sep_tokids
            # new_sep_len = new_sep_tokids.shape[-1]
            sep_len = sep_tokids.shape[-1]            
            assert (SEP_PADDING_IN_BATCH and max_sep_num==sep_len) or ( (not SEP_PADDING_IN_BATCH) and min_sep_num==sep_len)


        if USE_MAX_SEP_CACHE: ## Fixed sep cache size, i.e., only keep max_sep_len seps' kv in the cache. 
            if offset_init_size_layer + sep_len > self.max_sep_exidx[layer_idx]:
                max_sep_len = self.max_sep_exidx[layer_idx] - offset_init_size_layer
                assert sep_kv[0].shape[-2] == sep_tokids.shape[-1], f"The seq_len for seps' KVs and tok_ids should be the same."

                sep_kv, sep_tokids =  self.slice_kv_cache_and_tokids( sep_kv, sep_tokids, sep_len-max_sep_len, sep_len, seq_len = sep_tokids.shape[-1] ,_CHECK_IDX=True )
                self.sep_exrange[layer_idx] =  self.max_sep_exidx[layer_idx]  
            else:
                self.sep_exrange[layer_idx] =  offset_init_size_layer + sep_len             

        else:    ## Extend the sep cache and the whole cache if USE_MAX_SEP_CACHE is not set                           
            self.sep_exrange[layer_idx] =  offset_init_size_layer + sep_len
            if self.sep_exrange[layer_idx] > self.max_sep_exidx[layer_idx]:                    
                cache_incremental_gap = self.sep_exrange[layer_idx] - self.max_sep_exidx[layer_idx]
                self.max_sep_exidx[layer_idx] = self.sep_exrange[layer_idx] 
                self.sep_cache_size[layer_idx] = self.sep_cache_size[layer_idx] + cache_incremental_gap
                self.cache_size[layer_idx] = self.cache_size[layer_idx] + cache_incremental_gap

        if offset_init_size_layer > 0:                                
            cmp_past_kv_pairs, cmp_past_tok_ids  = self.cat_kv_cache_and_tokids( [initial_kv, sep_kv, local_kv ] ,  [initial_tokids, sep_tokids, local_tokids  ] )
        else:
            cmp_past_kv_pairs, cmp_past_tok_ids  = self.cat_kv_cache_and_tokids( [sep_kv, local_kv ] ,  [sep_tokids, local_tokids  ] )
                
        return cmp_past_kv_pairs, cmp_past_tok_ids, offset_init_size_layer
            

    def compress_past_win_2_seps(self, past_win_kv: Tuple[torch.Tensor], past_win_tokids: torch.Tensor, MIN_SEP_ALERT: bool=False, SEP_PADDING_IN_BATCH: bool=True ) -> Tuple[Union[Tuple[torch.Tensor], torch.Tensor, int ]]:
        """Compress the KVs in the past window into the sep cache where only separators' KVs are kept. Padding or Truncating if necessary."""
        sep_index_tensor = torch.zeros_like(past_win_tokids).bool()  # batch x seq_len

        for sp_id in self.separator_token_ids:            
            sep_index_tensor = sep_index_tensor | ( past_win_tokids == sp_id ) # batch x seq_len

        sep_cnt = sep_index_tensor.int().sum(-1)
        min_sep_num = sep_cnt.min()  # the min sep number for the seqs in a batch
        max_sep_num = sep_cnt.max()  # the max sep number for the seqs in a batch

        
        if MIN_SEP_ALERT and not SEP_PADDING_IN_BATCH:
            assert min_sep_num>0, f"The min sep number for each compressing time in a batch should be at least one if `MIN_SEP_ALERT=True` and `SEP_PADDING_IN_BATCH=False`"
                
        batch1_sep_ids_list = []
        batch_size = past_win_tokids.shape[0]
        for b_id in range(batch_size):            
            batch1_sep_ids = past_win_tokids[b_id, sep_index_tensor[b_id]] # #  sep_num
            if SEP_PADDING_IN_BATCH: ## padding
                sep_num = batch1_sep_ids.shape[-1]
                padding_num =  max_sep_num - sep_num                       
                if padding_num > 0:
                    assert padding_num <= past_win_tokids.shape[-1], f"padding_num: {padding_num} should be <= past_win_tokids.shape[-1]:{past_win_tokids.shape[-1]}"
                    batch1_sep_ids = batch1_sep_ids  # #  sep_num
                    batch1_pad_ids = past_win_tokids[b_id, -padding_num:]  # #  padding_num
                    batch1_sep_ids =  torch.cat([batch1_sep_ids, batch1_pad_ids], dim =-1)   ##  max_sep_num                
            else: ## truncating
                batch1_sep_ids = batch1_sep_ids[..., :min_sep_num ]  # #  min_sep_num
            batch1_sep_ids_list.append(batch1_sep_ids)                                                           
            
        new_sep_tokids = torch.stack(batch1_sep_ids_list, dim=0) # #  B x min_sep_num
        key_cache, value_cache = past_win_kv

        assert batch_size==key_cache.shape[0]
        batch1_sep_k_list = []
        batch1_sep_v_list = []
        batch1_sep_ids_list = []
        for b_id in range(batch_size):
            batch1_sep_k = self.k_bat_dim_select(key_cache, b_id, sep_index_tensor[b_id], min_sep_num, max_sep_num, SEP_PADDING_IN_BATCH)
            batch1_sep_k_list.append(batch1_sep_k)

            batch1_sep_v = self.v_bat_dim_select(value_cache, b_id, sep_index_tensor[b_id], min_sep_num, max_sep_num, SEP_PADDING_IN_BATCH)
            batch1_sep_v_list.append( batch1_sep_v )   
        
        sep_k = torch.stack(batch1_sep_k_list, dim=0)  ## batch x head x min_sep_num x dim
        sep_v = torch.stack(batch1_sep_v_list, dim=0)  ## batch x head x min_sep_num x dim                   
        new_sep_kv = (sep_k, sep_v)

        return new_sep_kv, new_sep_tokids, min_sep_num, max_sep_num      


    def apply_shifted_pos_emb(self, layer_idx: int, APPLY_PES_INSIDE: bool, PREFILLING_FLAG: bool, key_states: torch.Tensor, query_states: torch.Tensor, position_ids: torch.Tensor, cache_kwargs: Optional[Dict[str, Any]] = None ) -> torch.Tensor:        
        """Perform positional encoding shifting if required"""
        seq_len = self.get_usable_length(layer_idx)
        keys_to_shift = self.key_cache[layer_idx]
        queries_to_shift = query_states
        assert keys_to_shift.shape[self.k_seq_dim] == seq_len
        
        if cache_kwargs is None:
            cache_kwargs = {}

        if APPLY_PES_INSIDE:           
            if len(self._shifted_position_ids) <= layer_idx:
                self._shifted_position_ids.append(None)

            if PREFILLING_FLAG: ## for prefilling
                assert position_ids.shape[-1] >= seq_len, f"The length of position_ids should be >= the usable length of kv cache when prefilling."                
                self._shifted_position_ids[layer_idx] = position_ids[:, :seq_len].detach()
                shifted_pos_ids = self._shifted_position_ids[layer_idx]

            elif self._shifted_position_ids[layer_idx].shape[-1] >= seq_len:  ## for generation
                assert position_ids.shape[-1] == 1, f"The length of query and position_ids should be 1 during generation."
                shifted_pos_ids = self._shifted_position_ids[layer_idx][:, :seq_len].detach()

            elif self._shifted_position_ids[layer_idx].shape[-1] < seq_len:   ## for generation
                assert position_ids.shape[-1] == 1, f"The length of query and position_ids should be 1 during generation."
                increased_gap = seq_len - self._shifted_position_ids[layer_idx].shape[-1]
                assert increased_gap < self._shifted_position_ids[layer_idx].shape[-1], f"Normally, for auto-regressive model, the input length for each step should be 1 during generation."

                new_position_ids = self._shifted_position_ids[layer_idx][:, -increased_gap: ] + increased_gap
                self._shifted_position_ids[layer_idx] = torch.cat([self._shifted_position_ids[layer_idx], new_position_ids.detach()], dim=-1)
                shifted_pos_ids = self._shifted_position_ids[layer_idx]
            else:
                raise RuntimeError

            cos, sin = self._get_naive_shifted_cos_sin(
                key_states, shifted_pos_ids, seq_len
            )

            q_rope_idx = torch.arange( seq_len - query_states.shape[self.k_seq_dim],  seq_len).to(cos.device)
            cos_q, sin_q = cos.index_select(self._rope_seq_dim, q_rope_idx), sin.index_select(self._rope_seq_dim, q_rope_idx)

        else:
            sin = cache_kwargs.get("sin")
            cos = cache_kwargs.get("cos")                         
            sin_q = cache_kwargs.get("sin_q")
            cos_q = cache_kwargs.get("cos_q")    
            shifted_pos_ids = cache_kwargs.get("shifted_pos_ids") 
            assert (sin is not None) and (cos is not None), f"sin and cos matrices should be be provided"
            if sin_q is None:
                q_rope_idx = torch.arange( seq_len - query_states.shape[self.k_seq_dim],  seq_len).to(sin.device)
                sin_q = sin.index_select(self._rope_seq_dim, q_rope_idx)
            if cos_q is None:
                q_rope_idx = torch.arange( seq_len - query_states.shape[self.k_seq_dim],  seq_len).to(cos.device)
                cos_q = cos.index_select(self._rope_seq_dim, q_rope_idx)
            
        partial_rotation_size = cache_kwargs.get("partial_rotation_size")
        
        # On RoPE models, we need to recompute the Key rotation as the tokens are shifted
        if partial_rotation_size is not None:
            keys_to_shift, keys_pass = (
                keys_to_shift[..., :partial_rotation_size],
                keys_to_shift[..., partial_rotation_size:]
            )
            queries_to_shift, queries_pass = (
                queries_to_shift[..., :partial_rotation_size],
                queries_to_shift[..., partial_rotation_size:]
            )
                                    
        shifted_keys = self._apply_rotary_pos_emb_single(keys_to_shift, cos, sin, shifted_pos_ids, unsqueeze_dim=self._rope_unsqueeze_dim)
        shifted_queries = self._apply_rotary_pos_emb_single(queries_to_shift, cos_q, sin_q, shifted_pos_ids[:,  -queries_to_shift.shape[self.k_seq_dim] : ], unsqueeze_dim=self._rope_unsqueeze_dim)

        if partial_rotation_size is not None:
            shifted_keys = torch.cat( [shifted_keys, keys_pass], dim=-1)
            shifted_queries = torch.cat( [shifted_queries, queries_pass], dim=-1)


        return shifted_keys, shifted_queries


    def get_seq_length(self, layer_idx: Optional[int] = 0) -> int:
        """Returns the sequence length of the seen tokens. A layer index can be optionally passed."""                
        return self._seen_tokens


    def get_usable_length(self, layer_idx: int = 0) -> int:
        """Returns the sequence length of the actual cached states. A layer index must be passed."""                         
        if len(self.key_cache) <= layer_idx :
            return 0
        assert self.key_cache[layer_idx].shape[self.k_seq_dim] == self.value_cache[layer_idx].shape[self.v_seq_dim], f"`self.key_cache` and `self.value_cache` should have the same length."        
        return self.key_cache[layer_idx].shape[self.k_seq_dim]

    def get_initial_pos_offset(self, layer_idx:int = 0) -> int:      
        """Return the number of padding tokens in the record with the most left padding tokens in a batch."""
        assert isinstance(self.PADDING_ID, int), f"`self.PADDING_ID` should be correctly set."
        assert len(self.past_tok_ids) > layer_idx, f"`self.past_tok_ids` for layer {layer_idx} must have been properly set."
                
        past_tok_ids = self.past_tok_ids[layer_idx]
        assert past_tok_ids is not None, f"`past_tok_ids` for layer {layer_idx} should not be None"

        pad_index_tensor = (past_tok_ids == self.PADDING_ID)  ## batch x seq_len
        pad_toks_cnt = pad_index_tensor.int().sum(-1)  ## [batch]
        offset = pad_toks_cnt.max().item()

        return offset

                             
    def get_batch_size(self) -> int:
        """Return the batch size."""
        assert self.key_cache is not None, f"`self.key_cache` should not be None."
        assert self.value_cache is not None, f"`self.value_cache` should not be None."
        assert len(self.key_cache) > 0, f"`self.key_cache` is empty. No batch size is available."
        assert len(self.value_cache) > 0, f"self.value_cache is empty. No batch size is available."

        assert len(self.value_cache) == len(self.key_cache), f"self.value_cache and self.key_cache should be at the same length."
        assert self.value_cache[0].shape[0] == self.key_cache[0].shape[0], f"self.value_cache and self.key_cache should have the same batch size."

        return self.value_cache[0].shape[0]

    def get_kv_pair(self, layer_idx: int = None) -> Tuple[torch.Tensor]:
        assert layer_idx is not None, f"`layer_idx` must be given."

        if (len(self.key_cache) <= layer_idx) and (len(self.value_cache) <= layer_idx ):
            key = self.key_cache[layer_idx]
            value = self.value_cache[layer_idx]
        else:
            raise RuntimeError(f"The KV for layer:{layer_idx} have not been set.")
        return (key, value)


    def set_kv_cache(self, kv_pair: Tuple , layer_idx: int ) -> None:
        assert len(kv_pair) == 2, f"The length of `kv_pair` must be 2."
        self.key_cache[layer_idx] = kv_pair[0]
        self.value_cache[layer_idx] = kv_pair[1]
    
    def set_past_tok_ids(self, tok_ids: torch.Tensor, layer_idx:int) -> None:
        self.past_tok_ids[layer_idx] = tok_ids


    def cat_kv_cache_and_tokids(self, kv_pairs_list: List[Tuple[torch.Tensor]] , tok_ids_list:List[torch.Tensor]) -> Tuple[Union[Tuple[torch.Tensor],torch.Tensor]]:
        
        return self.cat_kv_cache(kv_pairs_list), self.cat_token_ids(tok_ids_list)


    def slice_kv_cache_and_tokids(self, kv_pair:Tuple[torch.Tensor], tok_ids_list:torch.Tensor, start:int, end:int, seq_len:int=None, _CHECK_IDX:bool=True, ) -> Tuple[Union[Tuple[torch.Tensor], torch.Tensor]]:
                             
        sliced_kv = self._slice_kv(start, end,  kv_pair=kv_pair, seq_len=seq_len, _CHECK_IDX=_CHECK_IDX,)                                    
        sliced_tids = self._slice_tok_ids(start, end, tok_ids_list = tok_ids_list, seq_len=seq_len, _CHECK_IDX=_CHECK_IDX)
        
        return sliced_kv , sliced_tids


    def cat_kv_cache(self, kv_pairs_list: List[Tuple[torch.Tensor]] ) -> Tuple[torch.Tensor]:               
        assert len(kv_pairs_list) >= 1 
        
        if len(kv_pairs_list) == 1 :
            return kv_pairs_list[0]
        else:
            ret = None 
            for i, kv_pair in enumerate(kv_pairs_list): # enumerate all the KVs needed to be cat
                if i == 0:
                    ret = kv_pair
                else:
                    ret = self._cat_kv(ret, kv_pair)
            return ret


    def cat_token_ids(self, tok_ids_list:List[torch.Tensor]  ) -> torch.Tensor :
        assert len(tok_ids_list) >= 1 
        
        return torch.cat(tok_ids_list, dim=-1)     


    def _cat_kv(self, kv_pair_a:Tuple[torch.Tensor],  kv_pair_b:Tuple[torch.Tensor]) -> Tuple[torch.Tensor]:            
        k_a, v_a = kv_pair_a
        k_b, v_b = kv_pair_b
        
        cat_k = torch.cat([k_a, k_b], dim=self.k_seq_dim)
        cat_v = torch.cat([v_a, v_b], dim=self.v_seq_dim)
        return (cat_k, cat_v)


    def _slice_kv(self, start:int, end:int, kv_pair: Tuple[torch.Tensor],   seq_len:int=None, _CHECK_IDX:bool=True)  -> Tuple[torch.Tensor] :
        assert kv_pair is not None, f"kv_pair must NOT be None when slicing it."
        key_cache = kv_pair[0]
        value_cache = kv_pair[1]

        if _CHECK_IDX:                                 
            assert seq_len is not None, f"seq_len must be given for checking the index for slicing"
            start, end = self._CHECK_IDX(start, end, seq_len)   
            
        sliced_key_cache = self.k_slice(key_cache, start, end) 
        sliced_value_cache = self.v_slice(value_cache, start, end)

        return ( sliced_key_cache, sliced_value_cache)


    def _slice_tok_ids(self, start:int, end:int, tok_ids_list:torch.Tensor , seq_len:int=None, _CHECK_IDX:bool=False) -> torch.Tensor:
        assert tok_ids_list is not None, f"tok_ids_list must NOT be None when slicing it."
        
        if _CHECK_IDX:
            assert seq_len is not None, f"seq_len must be given for checking the index for slicing"
            start, end = self._CHECK_IDX(start, end, seq_len)        
          
        sliced_tok_ids = tok_ids_list[:, start:end]
        return sliced_tok_ids

    def _set_layer_wise_attribute(self, name: str, value: Any, layer_num:int ):
        """Set layer-wise attributes"""
        if isinstance(value, int):        
            setattr(self, name, [value] * layer_num)
        elif isinstance(value, (list, tuple)):
            assert len(value) == layer_num, f"The length of {name}: {len(value)} must be equal to `layer_num`: {layer_num}"
            setattr(self, name, list(value))
        else:
            raise TypeError(f"{name} must be of the type `int` or `list` but got `{type(value)}`")

    def _list_element_add(self, list_a: List, list_b: List, bias: int=0, dtype = int, device = 'cpu') -> List:  
        """Element-wise addition between two lists."""      
        assert len(list_a) == len(list_b), f"The length of `list_a` ({len(list_a)}) must be equal to that of `list_b` ({len(list_b)})."
        tensor_c = torch.tensor(list_a, dtype=dtype, device=device) + torch.tensor(list_b, dtype=dtype, device=device) + torch.tensor([bias], dtype=dtype, device=device)
        return tensor_c.int().tolist()
        
    def _CHECK_IDX(self, start: int = 0, end: int = 100, seq_len: int = 1000):
        assert isinstance(start, int) and isinstance(end, int) and isinstance(seq_len, int), f"`start`, `end`, `seq_len` must be `int`."
        assert seq_len>0 , f"`seq_len` must > 0"
        
        if start <0 :
            start = start % seq_len
        if end < 0 :
            end = end % seq_len
        assert (start >=0) and (start < seq_len) , f"start:{start}, end:{end}, seq_len:{seq_len}"
        assert (end >= 0) and (end <= seq_len) , f"start:{start}, end:{end}, seq_len:{seq_len}"
        assert  start < end, f"start:{start}, end:{end}, seq_len:{seq_len}"

        return start,end

    def _CHECK_PARAMS_VALIDITY(self, layer_idx:int, left_padding_offset:int):
        assert len(self.cache_size) > layer_idx
        assert len(self.init_cache_size) > layer_idx
        assert len(self.sep_cache_size) > layer_idx
        assert len(self.max_sep_exidx) > layer_idx
        assert len(self.local_size) > layer_idx

        assert self.cache_size[layer_idx] > 0 , f"`self.cache_size` for layer:{layer_idx} must be greater than 0"
        assert self.init_cache_size[layer_idx] >= 0 , f"`self.init_cache_size` for layer:{layer_idx} must be greater than (equal to) 0"
        assert self.local_size[layer_idx] > 0 , f"`self.local_size` for layer:{layer_idx} must be greater than 0"
                    
        assert self.sep_cache_size[layer_idx] > 0 , f"`self.sep_cache_size` for layer:{layer_idx} must be greater than 0"
        assert self.max_sep_exidx[layer_idx] > 0 , f"`self.max_sep_exidx` for layer:{layer_idx} must be greater than 0"
        assert self.init_cache_size[layer_idx] + self.sep_cache_size[layer_idx] + self.local_size[layer_idx] + left_padding_offset < self.cache_size[layer_idx], f"`init_cache_size` ({self.init_cache_size[layer_idx]}) + `sep_cache_size` ({self.sep_cache_size[layer_idx]}) + `local_size` ({self.local_size[layer_idx]}) + `left_padding_offset` ({left_padding_offset}) for layer {layer_idx} should be less than `cache_size`:({self.cache_size[layer_idx]}) for layer {layer_idx}, i.e., a + s + w + (left_padding_offset) < c. Please increase `cache_size` if applicable."
        


    def _rotate_half(self, x):
        """Rotates half the hidden dims of the input."""
        x1 = x[..., : x.shape[-1] // 2]
        x2 = x[..., x.shape[-1] // 2 :]
        return torch.cat((-x2, x1), dim=-1)

    def _apply_rotary_pos_emb_single(self, k, cos, sin, position_ids=None, unsqueeze_dim=1):
        """Applies Rotary Position Embedding to the query and key tensors.

        Args:
            q (`torch.Tensor`): The query tensor.
            k (`torch.Tensor`): The key tensor.
            cos (`torch.Tensor`): The cosine part of the rotary embedding.
            sin (`torch.Tensor`): The sine part of the rotary embedding.
            position_ids (`torch.Tensor`, *optional*):
                Deprecated and unused.
            unsqueeze_dim (`int`, *optional*, defaults to 1):
                The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
                sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
                that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
                k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
                cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
                the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
        Returns:
            `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
        """        
        cos = cos.unsqueeze(unsqueeze_dim)   # batch x seq_len x dim  --> batch x 1 x seq_len x dim
        sin = sin.unsqueeze(unsqueeze_dim)        
        k_embed = (k * cos) + (self._rotate_half(k) * sin)
        return  k_embed


    def _get_naive_shifted_cos_sin(self, x: torch.Tensor, position_ids: torch.Tensor=None, seq_len=None):
        # x: [batch, num_attention_heads, seq_len, head_size]
        inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1)
        position_ids_expanded = position_ids[:, None, :].float()
        freqs = (inv_freq_expanded @ position_ids_expanded).transpose(1, 2)
        emb = torch.cat((freqs, freqs), dim=-1)
        cos = emb.cos().to(dtype=x.dtype)
        sin = emb.sin().to(dtype=x.dtype)
        # backwards compatibility
        self._cos_cached = cos
        self._sin_cached = sin
        return cos, sin
    

    def _get_scaled_shifted_cos_sin(self, x, position_ids, seq_len=None):
        # difference to the original RoPE: a scaling factor is aplied to the position ids
        position_ids = position_ids.float() / self.scaling_factor
        cos, sin = self._get_naive_shifted_cos_sin(x, position_ids, seq_len)
        return cos, sin


    def _get_dynamicNTK_scaling_shifted_cos_sin(self, x, position_ids, seq_len=None):
        # difference to the original RoPE: inv_freq is recomputed when the sequence length > original length
        seq_len = torch.max(position_ids) + 1
        if seq_len > self.max_position_embeddings:
            base = self.base * (
                (self.scaling_factor * seq_len / self.max_position_embeddings) - (self.scaling_factor - 1)
            ) ** (self.dim / (self.dim - 2))
            inv_freq = 1.0 / (
                base ** (torch.arange(0, self.dim, 2, dtype=torch.int64).float().to(x.device) / self.dim)
            )
            self.register_buffer("inv_freq", inv_freq, persistent=False)  # TODO: this may break with compilation

        cos, sin = self._get_naive_shifted_cos_sin(x, position_ids, seq_len)
        return cos, sin


    def _update_kv_ratio(self, kv_len_cmp:int, kv_len_ori:int, layer_idx: int=0) -> None:
        """Update the KV ratios which are for statistics and debugging."""
        if len(self._kept_kv_ratio) <= layer_idx:
            self._kept_kv_ratio.append( (kv_len_cmp,  kv_len_ori ) )    
        else:
            old_kv_len_cmp = self._kept_kv_ratio[layer_idx][0]
            old_kv_len_ori = self._kept_kv_ratio[layer_idx][1]
            self._kept_kv_ratio[layer_idx] = (old_kv_len_cmp + kv_len_cmp,  old_kv_len_ori + kv_len_ori )
            
    def _print_kv_ratio(self, layer_idx : int, LAYER_WISE: bool = False):
        """Print the KV ratios."""
        self._print_kv_ratio_count += 1 
        if LAYER_WISE:
            if self._print_kv_ratio_count % self.print_KV_inside_per_steps == 0:      
                print(f"######################## [Kept Tokens, Seen Tokens] : {self._kept_kv_ratio[layer_idx]}, Ratio: { (self._kept_kv_ratio[layer_idx][0]+1e-6) / (self._kept_kv_ratio[layer_idx][1]+1e-6) :.4f} ########################")    

        elif self._print_kv_ratio_count % (self.print_KV_inside_per_steps * self.layer_num) == 0:                
            print(f"######################## [Kept Tokens, Seen Tokens] : {self._kept_kv_ratio[layer_idx]}, Ratio: { (self._kept_kv_ratio[layer_idx][0]+1e-6) / (self._kept_kv_ratio[layer_idx][1]+1e-6) :.4f} ########################")    


    @classmethod ## Deprecated
    def from_legacy_cache(cls, 
                past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,

                ## For SepLLM                                
                init_cache_size: Union[int, List] = 4,        
                sep_cache_size: Union[int, List] = 64,
                local_size: Union[int, List]=256, 
                cache_size: Union[int, List]=512,    
                SEP_ACCUMULATION: bool = True,
                USE_MAX_SEP_CACHE: bool = False,
                SEP_PADDING_IN_BATCH: bool = False,
                separator_token_ids: List[int] = None, ## required for initialization if `model_type` is not provided. set it to `[-1]` to degrade SepCache to StreamingLLM's SinkCache
                PADDING_ID: int = None, ## required for initialization if `model_type` is not provided.

                ## For inheritance & initialization states
                past_tok_ids: List[torch.Tensor] = None,  ## It saves all the token ids corresponding to the saved KVs for all layers in SepCache.                
                key_cache: List[torch.Tensor] = None,          
                value_cache: List[torch.Tensor] = None,

                ## For debugging
                PRINT_KV_RATIO_INSIDE: bool = False,
                print_KV_inside_per_steps: int = 1000,   
                _seen_tokens: int = 0, 
                _kept_kv_ratio: List[Tuple[int]] = None,
                
                ### For positional encoding shifting
                APPLY_PE_SHIFT: bool = False,
                APPLY_PES_INSIDE: bool = True,
                _shifted_position_ids:  List[torch.Tensor] = None,
                _rope_unsqueeze_dim: int = 1, ## The unsqueeze_dim when applying RoPE.
                _rope_seq_dim: int=1, ## The seq_len dimension for the `cos` or `sin` tensors.
                pe_scaling_factor:float = 1.0,
                pe_dim:int=128, ## The number of dims for positional encoding. Typically, just set the `head_dim` to this.
                max_position_embeddings: int = 8192, 
                base: int=10000,  ## The base for RoPE.               
                
                ## For basic transformer architecture
                k_seq_dim: int=2, ## The dimension for seq_len in key tensors
                v_seq_dim: int=2, ## The dimension for seq_len in value tensors
                layer_num: int = None, ## required for initialization

                model_type: str = None,  ## The model type for running the example. choose from ['llama', 'pythia','falcon'].
                device = None    
    ) -> "SepCache":
        """Deprecated: Converts a cache in the legacy cache format into `SepCache`."""   

        if past_key_values is not None:
            assert len(past_key_values)==0, f"`from_legacy_cache` function is deprecated. You can only use it when `past_key_values=None` or `past_key_values` is empty, in which case, `from_legacy_cache` is equivalent to the `__init__` function."        
            past_key_values = None

        assert past_key_values is None, f"`from_legacy_cache` function is deprecated. You can only use it when `past_key_values=None` or `past_key_values` is empty, in which case, `from_legacy_cache` is equivalent to the `__init__` function."        
        

        if past_key_values is not None: ## Deprecated
            key_cache = []
            value_cache = []               
            
            for i, kv in enumerate(past_key_values):
                if i == 0:
                    past_tok_ids = [] if len(kv) == 4  else past_tok_ids       

                if len(kv) == 4:
                    k, v, p_tok_ids, _seen_tokens  = kv
                    key_cache.append(k)
                    value_cache.append(v)
                    past_tok_ids.append(p_tok_ids)
                    _seen_tokens = _seen_tokens
                elif len(kv) == 2:
                    k, v = kv
                    key_cache.append(k)
                    value_cache.append(v)
                    
        cache = cls(
                ## For SepLLM                
                init_cache_size = init_cache_size,        
                sep_cache_size = sep_cache_size,
                local_size = local_size, 
                cache_size = cache_size,                    
                SEP_ACCUMULATION = SEP_ACCUMULATION,
                USE_MAX_SEP_CACHE = USE_MAX_SEP_CACHE,
                SEP_PADDING_IN_BATCH = SEP_PADDING_IN_BATCH,
                separator_token_ids = separator_token_ids,
                PADDING_ID = PADDING_ID,

                ## For inheritance & initialization states
                past_tok_ids = past_tok_ids,  ## It saves all the token ids corresponding to the saved KVs for all layers in SepCache        
                key_cache = key_cache,          
                value_cache = value_cache,

                ## For debugging
                PRINT_KV_RATIO_INSIDE = PRINT_KV_RATIO_INSIDE,
                print_KV_inside_per_steps = print_KV_inside_per_steps,   
                _seen_tokens = _seen_tokens, 
                _kept_kv_ratio = _kept_kv_ratio,
                
                ### For positional encoding shifting
                APPLY_PE_SHIFT = APPLY_PE_SHIFT,
                APPLY_PES_INSIDE = APPLY_PES_INSIDE,
                _shifted_position_ids = _shifted_position_ids,
                _rope_unsqueeze_dim = _rope_unsqueeze_dim,
                _rope_seq_dim = _rope_seq_dim, 
                pe_scaling_factor = pe_scaling_factor,
                pe_dim = pe_dim,
                max_position_embeddings = max_position_embeddings, 
                base = base,                 
                
                ## For basic transformer architecture
                k_seq_dim = k_seq_dim,
                v_seq_dim = v_seq_dim,
                layer_num = layer_num,
                
                model_type = model_type,  
                device = device,   
        )

        return cache

    
    def to_legacy_cache(self) -> Tuple[Tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]]: ## Deprecated
        """Deprecated: Converts the `SepCache` instance into the legacy cache format, i.e., tuple."""
        print(">>>>>>>>>>>Warnings: Please try to avoid using this deprecated `to_legacy_cache` function since it will drop many useful parameters or states in SepCache.<<<<<<<<<<<")
        legacy_cache = ()
        for layer_idx in range(len(self.key_cache)):
            legacy_cache += ((self.key_cache[layer_idx], self.value_cache[layer_idx], self.past_tok_ids[layer_idx], self._seen_tokens), )
        return legacy_cache


    def __getitem__(self, layer_idx: int) -> List[Tuple[torch.Tensor]]:
        if layer_idx < len(self):
            return (self.key_cache[layer_idx], self.value_cache[layer_idx])
        else:
            raise KeyError(f"Cache only has {len(self)} layers, attempted to access layer with index {layer_idx}")

    def __iter__(self):
        """
        Support for backwards-compatible `past_key_value` iteration, e.g. `for x in past_key_value:` to iterate over
        keys and values
        """
        for layer_idx in range(len(self)):
            yield (self.key_cache[layer_idx], self.value_cache[layer_idx])

    def __len__(self):
        """
        Support for backwards-compatible `past_key_value` length, e.g. `len(past_key_value)`. This value corresponds
        to the number of layers in the model.
        """
        if self.key_cache is not None:
            return len(self.key_cache)
        else:
            return 0

    @property
    def seen_tokens(self):
        if hasattr(self, "_seen_tokens"):
            return self._seen_tokens
        else:
            return None



class KVUsageCounter:
    def __init__(self, PRINT_KV_per_ITERs:int = 100):
        """
        For detailed usage instructions, please refer to sepllm.github.io
        """
        self._total_kept_kv_ratio = (0, 0)
        self._printing_counter = 0
        self.PRINT_KV_per_ITERs = PRINT_KV_per_ITERs

    def accumulate_historical_kv_stats(self, cache: SepCache = None) -> None:
        assert cache is not None, f"The KV cache object (of the class SepCache) must be given."
        assert hasattr(cache, "_kept_kv_ratio"), f"The cache object must have the attribute _kept_kv_ratio."
        assert hasattr(cache, "layer_num"), f"The cache object must have the attribute layer_num."
        
        
        assert len(cache._kept_kv_ratio) == cache.layer_num, f"The length ({cache._kept_kv_ratio}) of cache object's _kept_kv_ratio attribute must be equal to layer_num ({cache.layer_num})"
        for ly in range(cache.layer_num):
            self._total_kept_kv_ratio = (self._total_kept_kv_ratio[0] + cache._kept_kv_ratio[ly][0],   self._total_kept_kv_ratio[1] + cache._kept_kv_ratio[ly][1] )
        self._printing_counter += 1

    def WHETHER_2_PRINT(self) -> bool:
        return (self._printing_counter % self.PRINT_KV_per_ITERs) == 0 


    def print_KV_ratio(self) -> None:       
        print(f"######################## The KVs for ALL layers: (KV number kept for predicting current token)/(Total seen KV number) ########################")             
        print(f"########################>>>>>>>>>>> [Kept Tokens, Seen Tokens] : {self._total_kept_kv_ratio}, Ratio: { (self._total_kept_kv_ratio[0]+1e-6) / (self._total_kept_kv_ratio[1]+1e-6):.4f} <<<<<<<<<<<<##########################")    
        print(f"######################## -------------------------------------------------------------------------------------------- ########################")

#############################################################End################################################################



##################################################### Monkey Patch Utils #######################################################
def get_full_class_import_path(obj):
    """Get the complete class import path of an object"""
    # Get the class of the object
    cls = obj.__class__
    
    # Get the module name where the class is defined
    module = cls.__module__
    
    # Get the qualified name of the class (including outer classes)
    qualname = cls.__qualname__
    
    # Handle nested classes (e.g., ClassA.ClassB)
    if '.' in qualname:
        # Replace nested class separators
        class_path = f"{module}.{qualname.replace('.', '_')}"
    else:
        class_path = f"{module}.{qualname}"
    
    return class_path


def get_importable_class_path(obj):
    """Get the directly importable class path (handling special cases and dynamic classes)"""
    cls = obj.__class__
    module = cls.__module__
    qualname = cls.__qualname__
    
    # Handle built-in types
    if module == 'builtins':
        return qualname
    
    # Handle dynamically generated classes (e.g., functools.partial)
    if not hasattr(cls, '__module__') or module is None:
        return f"<dynamic class {qualname}>"
    
    # Handle nested classes
    if '.' in qualname:
        # Try to import the parent module to validate the path
        try:
            import importlib
            parent_module = importlib.import_module(module)
            
            # Follow the qualified name path
            parts = qualname.split('.')
            current = parent_module
            for part in parts:
                current = getattr(current, part)
            
            # If successful access, return the original path
            return f"{module}.{qualname}"
        except (ImportError, AttributeError):
            # Fallback: use underscore connection
            return f"{module}.{qualname.replace('.', '_')}"
    
    return f"{module}.{qualname}"


def monkey_patch_by_class_path(model, new_forward):
    """Perform monkey patching through class path"""
    # Get the complete class path
    class_path = get_importable_class_path(model)
    
    # Dynamically import the class
    try:
        import importlib
        module_path, class_name = class_path.rsplit('.', 1)
        module = importlib.import_module(module_path)
        target_class = getattr(module, class_name)
        
        # Save the original method
        if not hasattr(target_class, '_original_forward'):
            target_class._original_forward = target_class.forward
        
        # Apply the patch
        target_class.forward = new_forward
        
        # Update the method binding of the current instance
        model.forward = types.MethodType(target_class.forward, model)
        
        return f"Successful Monkey Patch: {class_path}.forward"
    
    except (ImportError, AttributeError, ValueError) as e:
        return f"Patch Failed: {str(e)}"


def find_inner_attribute(obj,  attr_name_list: List[str], default_type = PreTrainedModel ):
    # try to find the attribute of the name in `attr_name_list`.
    for target_attr_name in attr_name_list:
        if hasattr(obj, target_attr_name):
            return getattr(obj, target_attr_name)
        
    # else: try to find the attribute of the type `default_type`
    for attr_name in dir(obj): 
        attr_value = getattr(obj, attr_name)   
        if isinstance(attr_value, default_type):
            return attr_value

    raise AttributeError(f"In the {obj} object, there is no attribute whose name matches any name in {attr_name_list} or whose type is {default_type}.")


def find_attribute_name(obj, name_pattern_list: List[str], exclude_pattern_list: List[str], match_type = nn.Module):
    for attr_name in dir(obj): 
        attr_value = getattr(obj, attr_name)   
        for pattern in name_pattern_list:
            for ex_pattern in exclude_pattern_list:
                if isinstance(attr_value, match_type) and (pattern.lower() in attr_value.__class__.__name__.lower()) and ( ex_pattern.lower() not in attr_value.__class__.__name__.lower() ):
                    return attr_value
                elif isinstance(attr_value, match_type) and (pattern.lower() in attr_name.lower()) and (ex_pattern.lower() not in attr_name.lower() ):
                    return attr_value

    raise AttributeError(f"In the {obj} object, there is no attribute whose name matches any pattern in {name_pattern_list} and excludes any pattern in {exclude_pattern_list}, and whose type is {match_type}.")



def monkey_patching(model_obj, 
                    model_atten_forward , ## The `forward` function used to patch.
                    possible_inner_model_names: List[str] = ["model", "transformer", "gpt_neox"] , # In `XXXForCausalLM` class, the possible name of internal attribute for model. e.g.,  "model", "transformer", "gpt_neox", etc.
                    possible_layers_names: List[str] = ["layers", "h" ],  # In `XXXModel` class,  the possible name of internal attribute for decoder layers, e.g.,  "layers", "h", etc.
                    atten_attr_name_pattern_list: List[str] = ["attention", "self_attn"],  # In `XXXDecoderLayer` class, the possible name of internal attribute for self-attention, e.g.,  "attention", "self_attn", etc.
                    atten_attr_name_pattern_exclude: List[str] = ["norm", "layer"], # In `XXXDecoderLayer` class, the impossible name patterns (i.e., the patterns to be excluded) of internal attribute for self-attention module class, e.g., "norm" , etc. Sometimes, there will be some attributes like "post_attention_norm" and we do not want modify the `forward` function of it - we want to modify the `forward` function of `XXXAttention`. So, we need to exclude attribute name patterns like "norm" to accurately find the correct "forward" function to replace.
                    verbose = True):
    
    """
    This `monkey_patching` function is to
        - find the `forward` function of the `XXXAttention` class.
        - replace all the related `forward` functions of the instances of `XXXAttention` class with `model_atten_forward`.
    """
    
    ## To avoid the argument check failure, i.e., let "sepllm_kwargs" pass the check.
    transformers.generation.GenerationMixin._validate_model_kwargs = _validate_model_kwargs    

    ## Get inner model obj
    inner_model_type = PreTrainedModel
    inner_model = find_inner_attribute(model_obj, possible_inner_model_names, inner_model_type)
    
    ## Get the decoder layers (`nn.ModuleList`) obj
    layers_type = nn.ModuleList
    model_layers = find_inner_attribute(inner_model, possible_layers_names, layers_type)
    
    ## Replace all the related `forward` functions of XXXAttention class's instances.
    for i, decoder_layer in enumerate(model_layers):
        self_attn_module = find_attribute_name(decoder_layer, atten_attr_name_pattern_list, atten_attr_name_pattern_exclude, nn.Module)
        result = monkey_patch_by_class_path(self_attn_module, model_atten_forward)
        if verbose:
            decoder_class_name = get_importable_class_path(decoder_layer)
            print(f"For Layer {i}'s `{decoder_class_name}`: {result}")

    return model_layers
#############################################################End################################################################


def str2list(s : str, sep = ';'):
    if isinstance(s, str):
        res =  s.split(sep)
        return [ int(i) for i in res ]
    else:
        return s


def generate(model,                          
            ## For SepCache                              
            init_cache_size: Union[int, List] = 4,        
            sep_cache_size: Union[int, List] = 128,
            local_size: Union[int, List]=256, 
            cache_size: Union[int, List]=512,    
            SEP_ACCUMULATION: bool = True,
            USE_MAX_SEP_CACHE: bool = False,
            SEP_PADDING_IN_BATCH: bool = False,
            separator_token_ids: List[int] = None, ## required for initialization if `model_type` is not provided.
            PADDING_ID: int = None, ## required for initialization if `model_type` is not provided.

            ## For inheritance & initialization states
            past_tok_ids: List[torch.Tensor] = None,  ## It saves all the token ids corresponding to the saved KVs for all layers in SepCache.                
            key_cache: List[torch.Tensor] = None,          
            value_cache: List[torch.Tensor] = None,

            ## For debugging
            PRINT_KV_RATIO_INSIDE: bool = False,
            print_KV_inside_per_steps: int = 1000,   
            _seen_tokens: int = 0, 
            _kept_kv_ratio: List[Tuple[int]] = None,
            
            ### For positional encoding shifting
            APPLY_PE_SHIFT: bool = False,
            APPLY_PES_INSIDE: bool = False,
            _shifted_position_ids:  List[torch.Tensor] = None,
            _rope_unsqueeze_dim: int = 1, ## The unsqueeze_dim when applying RoPE.
            _rope_seq_dim: int=1, ## The seq_len dimension for the `cos` or `sin` tensors.
            pe_scaling_factor:float = 1.0,
            pe_dim:int=128, ## The number of dims for positional encoding. Typically, just set the `head_dim` to this.
            max_position_embeddings: int = 8192, 
            base: int=10000,  ## The base for RoPE.               
            
            ## For basic transformer architecture
            k_seq_dim: int=2, ## The dimension for seq_len in key tensors
            v_seq_dim: int=2, ## The dimension for seq_len in value tensors
            layer_num: int = None, ## required for initialization

            model_type: str = 'llama',  ## The model type for running the example. choose from ['llama', 'pythia','falcon'].
            device = None,   

            ## For verbosity of monkey patching
            monkey_patch_verbose: bool = False,

             **kwargs
             ):
    """Custom generate function for SepCache.

    A cache as described in the [SepLLM paper - ICML 2025](https://arxiv.org/abs/2412.12094). In the training phase, 
    SepLLM condenses the segment information into the KV of the separator that divides the segment. In the inference phase, the 
    corresponding SepCache only needs to store the KVs of initial tokens, separator tokens, and recent tokens for generation.

    It stores the Key and Value states as lists of tensors, two lists for each layer. The expected shape for each tensor is
    `[batch_size, num_heads, seq_len, head_dim]`.

    Frequently-Used Parameters:

        `init_cache_size: Union[int, List]`:
            The maximum number of KVs to be stored for initial tokens.
            In the paper, the hyperparameter `a` is an abbreviated alias for `self.init_cache_size`.                
                
        `sep_cache_size: Union[int, List]`:
            The maximum number of KVs to be stored for separator tokens.
            In the paper, the hyperparameter `s` is an abbreviated alias for `self.sep_cache_size`.

        `local_size: Union[int, List]`: 
            The maximum number of KVs to be stored for local tokens (i.e., sliding window).
            In the paper, the hyperparameter `w` is an abbreviated alias for `self.local_size`.

        `cache_size: Union[int, List]`:    
            The maximum number of KVs to be stored for all the tokens, i.e., the size for the whole KV cache.  
            In the paper, the hyperparameter `c` is an abbreviated alias for `self.cache_size`.

        Concerning these four parameters above:
            When a list is passed (its length must be `layer_num`), it represents different values for each layer. 
            When an integer is passed, it means the setting is the same for all layers.
        
        
        `USE_MAX_SEP_CACHE: bool`: 
            If True, it means we only keep at most `self.sep_cache_size` seperators' KVs.  
            If the number exceeds this limit, older separator's KVs will be discarded, keeping only the most recent `self.sep_cache_size` KVs. 
            In the paper, the hyperparameter `s` is an abbreviated alias for `self.sep_cache_size`.
          
        `separator_token_ids: List[int]`:
            The token ids of the separator tokens for the current model's tokenizer.            
            We have some examples, such as the Llama-3 series models, where setting `model_type='llama'` allows you 
                to skip setting `separator_token_ids` and `PADDING_ID` (SepCache will auto-fill them).

        `PADDING_ID: int`:
            The token id of the padding token. You can just set `PADDING_ID` to the id of "<|endoftext|>" token of the tokenizer for the pretrained model.

    Important Note:
        When `cache_size` and `local_size` are set to infinity (i.e., sufficiently large positive integers), and `USE_MAX_SEP_CACHE` is `False`, `SepCache` degenerates into a regular Cache. 
        However, you must always ensure that `init_cache_size` + `sep_cache_size` + `local_size` + `left_padding_offset` < `cache_size`. 
        Here, `left_padding_offset` denotes the number of padding tokens in the record with the largest left paddings within a runtime batch. `left_padding_offset` can only be determined at runtime.        
        To guarantee the above inequality always holds during runtime, when setting, you can intentionally create a sufficient margin between both sides of the following inequality:
            `init_cache_size` + `sep_cache_size` + `local_size`  < `cache_size`, i.e., `a`+`s`+`w`<`c` in the [SepLLM paper - ICML 2025]
            to leave room for `left_padding_offset`.

        Please refer to the `__init__` function's comments for more details on the parameters.
            
    Example:

        ```python        
        >>> from transformers import AutoTokenizer, AutoModelForCausalLM, 
        >>> from .custom_generate.generate import SepCache
        >>> import torch
        >>> from huggingface_hub import login
        >>> login("hf_xxxXXXxxx")


        >>> def to_cuda(a_dict: dict) -> dict:
        >>>    new_dict = {}    
        >>>    for k,v in a_dict.items():
        >>>        if isinstance(v, torch.Tensor):
        >>>            new_dict[k] = v.cuda()
        >>>        else:
        >>>            new_dict[k] = v
        >>>    return new_dict

        >>> model = AutoModelForCausalLM.from_pretrained("meta-llama/Meta-Llama-3-8B-Instruct", attn_implementation="flash_attention_2", device_map="cuda:0")
        >>> model.bfloat16().cuda()
        >>> tokenizer = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3-8B-Instruct")
        >>> inputs = tokenizer(text="My name is Llama 3", return_tensors="pt")
        >>> inputs = to_cuda(inputs)
        >>> # Prepare a cache and pass it to model's forward; `layer_num` is the number of layers for the pretrained model.
        >>> past_key_values = SepCache(init_cache_size=4, sep_cache_size=128, local_size=256, cache_size=512, layer_num=32, USE_MAX_SEP_CACHE=True, model_type='llama')
        >>> # `separator_token_ids` and `PADDING_ID` must also be provided if you are not using `model_type='llama'` like this demo.
        >>> outputs = model(**inputs, past_key_values=past_key_values, use_cache=True)
        >>> outputs.past_key_values # access SepCache filled with keys/values
        SepCache()
        ```

        ```python
        >>> ## When using the `update` function of SepCache to update the keys/values and the past token ids (necessary in SepCache), the current `input_ids` must also be provided.        
        >>> key_states, value_states = past_key_values.update(                
                    key_states = key_states,
                    value_states = value_states,    
                    input_ids = input_ids,
                    layer_idx = layer_idx,     
                    PREFILLING_FLAG = q_len > 1, ## `q_len` is the sequence length of the current `query_states`
                    )

        ```
        For detailed usage instructions, please refer to https://github.com/HKUDS/SepLLM
    """

    # 0. Monkey Patching towards the "forward" function of `XXXAttention` class in order to pass `input_ids` to the `update` function of `SepCache` when calling it.   
    model_layers = monkey_patching(model,  model_atten_forward=llama_atten_forward, verbose=monkey_patch_verbose)
    separator_token_ids = str2list(separator_token_ids)
    

    # 1. General sanity checks
    # 1.a. A few arguments are not allowed, especially arguments that control caches.
    generation_config = kwargs.get("generation_config")
    default_global_generation_config = GenerationConfig()
    default_model_generation_config = model.generation_config
    for arg in UNSUPPORTED_GENERATION_ARGS:
        has_custom_gen_config_arg = (
            generation_config is not None
            # = and not (match global default or match model-specific default)
            and not (
                getattr(default_model_generation_config, arg) == getattr(generation_config, arg)
                or getattr(default_global_generation_config, arg) == getattr(generation_config, arg)
            )
        )
        kwargs_has_arg = arg in kwargs and kwargs[arg] is not None
        if kwargs_has_arg or has_custom_gen_config_arg:
            raise ValueError(
                f"`{arg}` is set, but it's not supported in this custom generate function. List of "
                f"unsupported arguments: {UNSUPPORTED_GENERATION_ARGS}"
            )

    
    
    # 1.b. The model must be decoder-only
    if model.config.is_encoder_decoder:
        raise ValueError("This custom generate function only works with decoder-only models")

    # 1.c. compatibility with transformers>=4.52: we must pop `custom_generate` from kwargs, otherwise it will result
    # in an infinite loop when we call `model.generate`. This is solved in transformers 4.53.
    kwargs.pop("custom_generate", None)


    sepllm_kwargs = {}
    sepllm_kwargs["input_ids"] = kwargs["input_ids"] ## `input_ids` must be passed to the `update` function of `SepCache` when calling it.
    kwargs["sepllm_kwargs"] = sepllm_kwargs

    # 2. Generate with SepCache
    # 2.a. prepare the cache, if it was not passed.
    past_key_values = kwargs.pop("past_key_values", None)
    if past_key_values is None:
        past_key_values = SepCache(                     
            ## For SepCache                              
            init_cache_size = init_cache_size,        
            sep_cache_size = sep_cache_size,
            local_size = local_size, 
            cache_size = cache_size,    
            SEP_ACCUMULATION = SEP_ACCUMULATION,
            USE_MAX_SEP_CACHE = USE_MAX_SEP_CACHE,
            SEP_PADDING_IN_BATCH = SEP_PADDING_IN_BATCH,
            separator_token_ids = separator_token_ids, ## required for initialization if `model_type` is not provided.
            PADDING_ID = PADDING_ID, ## required for initialization if `model_type` is not provided.

            ## For inheritance & initialization states
            past_tok_ids = past_tok_ids,  ## It saves all the token ids corresponding to the saved KVs for all layers in SepCache.                
            key_cache = key_cache,          
            value_cache = value_cache,

            ## For debugging
            PRINT_KV_RATIO_INSIDE = PRINT_KV_RATIO_INSIDE,
            print_KV_inside_per_steps = print_KV_inside_per_steps,   
            _seen_tokens = _seen_tokens, 
            _kept_kv_ratio = _kept_kv_ratio,
            
            ### For positional encoding shifting
            APPLY_PE_SHIFT = APPLY_PE_SHIFT,
            APPLY_PES_INSIDE = APPLY_PES_INSIDE,
            _shifted_position_ids = _shifted_position_ids,
            _rope_unsqueeze_dim = _rope_unsqueeze_dim, ## The unsqueeze_dim when applying RoPE.
            _rope_seq_dim =_rope_seq_dim, ## The seq_len dimension for the `cos` or `sin` tensors.
            pe_scaling_factor = pe_scaling_factor,
            pe_dim = pe_dim, ## The number of dims for positional encoding. Typically, just set the `head_dim` to this, i.e., model.config.hidden_size // model.config.num_attention_heads
            max_position_embeddings = max_position_embeddings, # i.e.,  model.config.max_position_embeddings                               
            base = base,  ## The base for RoPE.               
            
            ## For basic transformer architecture
            k_seq_dim = k_seq_dim, ## The dimension for seq_len in key tensors
            v_seq_dim = v_seq_dim, ## The dimension for seq_len in value tensors
            layer_num = len(model_layers), ## required for initialization. model.config.num_hidden_layers

            model_type = model_type,  ## The model type for running the example. choose from ['llama', 'pythia','falcon'].
            device = device,    
            )

    elif not isinstance(past_key_values, SepCache):
        raise ValueError(f"`past_key_values` must be a `SepCache` instance, got a {type(past_key_values)} instance")

    # 2.b. generate with the cache
    kwargs["use_cache"] = True
    generation_outputs = model.generate(**kwargs, past_key_values=past_key_values)
    return generation_outputs