Delete custom_generate

custom_generate/__init__.py

@@ -1,3 +0,0 @@
-from . import sep_cache_utils
-from . import monkey_patching_utils
-from . import functions_2_patch

custom_generate/functions_2_patch.py

@@ -1,221 +0,0 @@
-import torch
-import inspect
-import importlib
-
-from typing import Callable, Optional, Union, Any, List
-from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
-from transformers.cache_utils import Cache
-from transformers.processing_utils import Unpack
-
-from .sep_cache_utils import SepCache
-
-
-
-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)"
-        )
-
-

custom_generate/generate.py

@@ -1,271 +0,0 @@
-import sys
-import os
-sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
-
-
-import torch
-
-import types
-from typing import Any, Dict, List, Optional, Tuple, Union
-import transformers
-from transformers import Cache, GenerationConfig
-import torch.nn as nn
-from transformers.modeling_utils import  PreTrainedModel
-
-
-from ..functions_2_patch import _validate_model_kwargs, llama_atten_forward
-from ..monkey_patching_utils import monkey_patching
-from ..sep_cache_utils import SepCache
-
-
-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
-]
-
-
-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 .sep_cache_utils 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 for the `update` function of `SepCache`    
-    model_layers = monkey_patching(model,  model_atten_forward=llama_atten_forward, verbose=monkey_patch_verbose)
-
-    # 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`
-    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

custom_generate/monkey_patching_utils.py

@@ -1,154 +0,0 @@
-import torch
-import inspect
-import importlib
-import transformers
-import types
-
-import torch.nn as nn
-from transformers.modeling_utils import  PreTrainedModel
-from typing import Callable, Optional, Union, Any, List
-
-from .functions_2_patch import _validate_model_kwargs, llama_atten_forward
-
-
-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 , verbose = True):
-    transformers.generation.GenerationMixin._validate_model_kwargs = _validate_model_kwargs    
-
-    ## get inner model
-    possible_inner_model_names = ["model", "transformer", "gpt_neox"]
-    inner_model_type = PreTrainedModel
-    inner_model = find_inner_attribute(model_obj, possible_inner_model_names, inner_model_type)
-
-
-    possible_layers_names = ["layers", "h" ]
-    layers_type = nn.ModuleList
-    model_layers = find_inner_attribute(inner_model, possible_layers_names, layers_type)
-    
-    atten_attr_name_pattern_list = ["attention", "self_attn"]
-    atten_attr_name_pattern_exclude = ["norm", "layer"]
-
-    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

custom_generate/sep_cache_utils.py

@@ -1,1205 +0,0 @@
-from transformers import Cache, GenerationConfig
-from typing import Any, Dict, List, Optional, Tuple, Union
-import torch
-from packaging import version
-from dataclasses import dataclass
-
-
-
-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"######################## -------------------------------------------------------------------------------------------- ########################")
-
-
-
-