transformers-community
/

sep_cache

@@ -1,324 +0,0 @@
-## 1. Abstract
-`SepCache` is a simple yet effective, native sparse attention `Cache` class proposed in the [`SepLLM paper - ICML 2025`](https://icml.cc/virtual/2025/poster/45536) , which most closely aligns with the semantic distribution of natural language. 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.
-Notably, `SepCache` also delivers strong performance across many tasks in training-free scenarios. Moreover, `SepLLM` (or simply `SepCache`) is the **most suitable baseline method for sparse attention mechanisms and KV compression/management**, as it is the natively sparse attention mechanism that best aligns with the natural semantic distribution of language.
-See more details and advanced usage in https://github.com/HKUDS/SepLLM
-![image](https://hackmd.io/_uploads/r1POJoR4yg.png)
-## 2. Usage
-### 2.1 Sample Base Model
-We recommend using models from the **Llama 3 series**. Our example model is based on `meta-llama/Meta-Llama-3-8B-Instruct`, for which we have already prepared a targeted `monkey patch`.
-For other models, using `SepCache` requires minor modifications to the corresponding `modeling_xxx.py` file or writing a **custom monkey patch**. These changes are **very simple** -- you only need to pass arguments like `input_ids` to the `update` function of `SepCache` when calling it.
-We will provide a detailed guide later on how to modify your `modeling_xxx.py` file or `monkey patch` file to adapt `SepCache` to any model.
-### 2.2 Quick Start
-#### 2.2.1 Environment Setup
-You need to install `transformers>=4.53`, and we recommend using `lm_eval>=0.4.9` for running evaluations. We suggest managing your Python environment with `conda` for better dependency control.
-```bash
-conda create -n sepcache python=3.10
-conda activate sepcache
-pip install transformers==4.53
-pip install lm_eval==0.4.9
-```
-#### 2.2.2 A Simple Example
-You can use `SepCache` by specifying `custom_generate="transformers-community/sep_cache"` or `custom_generate="Gausson/sep_cache"` when calling the `generate` function. In our demo, we have already prepared sample monkey patching for the `Llama 3 series` models and provided some common parameters for initializing `SepCache`.
-```python
-# requires `transformers>=4.53.0`
-from transformers import AutoModelForCausalLM, AutoTokenizer
-# Preparing model, tokenizer, and model inputs
-tokenizer = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3-8B-Instruct")
-model = AutoModelForCausalLM.from_pretrained("meta-llama/Meta-Llama-3-8B-Instruct", device_map="auto")
-messages = [{"role": "user", "content": "Tell me a story about a cat."}]
-text = tokenizer.apply_chat_template(
-    messages,
-    tokenize=False,
-    add_generation_prompt=True,
-    enable_thinking=False
-)
-model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
-# Using SepCache for generation
-gen_out = model.generate(
-    # usual `generate` arguments
-    **model_inputs,
-    do_sample=False,
-    max_new_tokens=100,
-    return_dict_in_generate=True,
-    monkey_patch_verbose = True, # To see which functions are actually being monkey patched for `SepCache`.
-    # Using SepCache
-    custom_generate="transformers-community/sep_cache", ## Alternatively, you can use `Gausson/sep_cache`
-    trust_remote_code=True,
-    # SepCache arguments
-    init_cache_size = 4,
-    sep_cache_size = 128,
-    local_size = 256,
-    cache_size = 512,
-    USE_MAX_SEP_CACHE = True,
-    model_type = 'llama'
-)
-print(tokenizer.batch_decode(gen_out.sequences, skip_special_tokens=True))
-assert "sepcache" in str(type(gen_out.past_key_values)).lower()
-```
-It is worth noting that you must specify the `separator_token_ids: List[int]` and `PADDING_ID: int` parameters for initializing `SepCache`. In the example above, we did not do this because, for convenience, in the demo above, we specified `model_type = "llama"`, in which case `separator_token_ids` and `PADDING_ID` will be automatically filled.
-However, when you use a tokenizer for a non-Llama 3 series model, you need to specify the specific values of `separator_token_ids` and `PADDING_ID` based on the tokenizer you are using. For example, the following example is based on the values obtained from a Llama 3 series tokenizer.
-```python
-# Using SepCache for generation
-gen_out = model.generate(
-    # usual `generate` arguments
-    **model_inputs,
-    do_sample=False,
-    max_new_tokens=100,
-    return_dict_in_generate=True,
-    monkey_patch_verbose = True, # To see which functions are actually being monkey patched for `SepCache`.
-    # Using SepCache
-    custom_generate="transformers-community/sep_cache", ## Alternatively, you can use `Gausson/sep_cache`
-    trust_remote_code=True,
-    # SepCache arguments
-    init_cache_size = 4,
-    sep_cache_size = 128,
-    local_size = 256,
-    cache_size = 512,
-    USE_MAX_SEP_CACHE = True,
-    separator_token_ids = [128000, 13, 11, 30, 0, 26, 25, 198, 220, 662, 1174, 949, 758, 2652, 551, 720, 256,262],
-    PADDING_ID = 128009
-)
-```
-#### 2.2.3 Frequently-Used Parameters
-Below, we provide explanations and examples for the most commonly used parameters when initializing `SepCache`. These parameters can be passed through the `generate` function.
-```
-`SepCache` 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 `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 `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 `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 `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 `sep_cache_size` seperators' KVs.
-        If the number exceeds this limit, older separators' KVs will be discarded, keeping only the most recent `sep_cache_size` KVs.
-        In the paper, the hyperparameter `s` is an abbreviated alias for `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.
-- 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](https://arxiv.org/abs/2412.12094) to leave room for `left_padding_offset`.
-**More Important Note: In practice, no need to do positional encoding (PE) shifting like [StreamingLLM](https://github.com/mit-han-lab/streaming-llm/) if the actual length does not exceed the pretrained max PE length (which applies to most downstream tasks.) . So, for most basic usages, just set `APPLY_PE_SHIFT=False` (`False` is also the default setting) and `APPLY_PES_INSIDE=False` for initialization.**
-#### 2.2.4 Update Function
-After initialization, another key point to note is that when using the `update` function of `SepCache` to update the **keys/values** and the **past token IDs** (which is necessary in SepCache), the current `input_ids` must also be provided.
-```python
-key_states, value_states = past_key_values.update(
-          key_states = key_states,
-          value_states = value_states,
-          input_ids = input_ids,  ## required
-          layer_idx = layer_idx,
-          PREFILLING_FLAG = q_len > 1, ## `q_len` is the sequence length of the current `query_states`
-          )
-```
-#### 2.2.5 Monkey Patch Demo
-To adapt the `update` function of `SepCache` mentioned in [`2.2.4 Update Function`](#224-update-function), i.e., passing the current `input_ids` as a parameter to the `update` function. It is worth noting that during the prefilling stage, the shape of the input_ids tensor is `[batch_size, seq_len]`, while during the decoding stage of auto-regressive models, the shape of the `input_ids` tensor should be `[batch_size, 1]`.
-In our `custom_generate/generate.py` file, we provide the `monkey_patching` function, which works by replacing the `forward` function in all the related instances of the `XXXAttention` class (for example, in the Llama 3 series model, it would be `LlamaAttention`) with our customized forward function (specified by the `model_atten_forward` parameter of the `monkey_patching` function).
-```python
-def monkey_patching(model_obj,
-                    model_atten_forward , ## The `forward` function used to patch.
-                    possible_inner_model_names: List[str] = ["model", "transformer", "gpt_neox"] , # In `XXXForCausalLM` class, the possible name of internal attribute for model. e.g.,  "model", "transformer", "gpt_neox", etc.
-                    possible_layers_names: List[str] = ["layers", "h" ],  # In `XXXModel` class,  the possible name of internal attribute for decoder layers, e.g.,  "layers", "h", etc.
-                    atten_attr_name_pattern_list: List[str] = ["attention", "self_attn"],  # In `XXXDecoderLayer` class, the possible name of internal attribute for self-attention, e.g.,  "attention", "self_attn", etc.
-                    atten_attr_name_pattern_exclude: List[str] = ["norm", "layer"], # In `XXXDecoderLayer` class, the impossible name patterns (i.e., the patterns to be excluded) of internal attribute for self-attention module class, e.g., "norm" , etc. Sometimes, there will be some attributes like "post_attention_norm" and we do not want modify the `forward` function of it - we want to modify the `forward` function of `XXXAttention`. So, we need to exclude attribute name patterns like "norm" to accurately find the correct "forward" function to replace.
-                    verbose = True):
-    """
-    This `monkey_patching` function is to
-        - find the `forward` function of the `XXXAttention` class.
-        - replace all the related `forward` functions of the instances of `XXXAttention` class with `model_atten_forward`.
-    """
-    ## To avoid the argument check failure, i.e., let "sepllm_kwargs" pass the check.
-    transformers.generation.GenerationMixin._validate_model_kwargs = _validate_model_kwargs
-    ## Get inner model obj
-    inner_model_type = PreTrainedModel
-    inner_model = find_inner_attribute(model_obj, possible_inner_model_names, inner_model_type)
-    ## Get the decoder layers (`nn.ModuleList`) obj
-    layers_type = nn.ModuleList
-    model_layers = find_inner_attribute(inner_model, possible_layers_names, layers_type)
-    ## Replace all the related `forward` functions of XXXAttention class's instances.
-    for i, decoder_layer in enumerate(model_layers):
-        self_attn_module = find_attribute_name(decoder_layer, atten_attr_name_pattern_list, atten_attr_name_pattern_exclude, nn.Module)
-        result = monkey_patch_by_class_path(self_attn_module, model_atten_forward)
-        if verbose:
-            decoder_class_name = get_importable_class_path(decoder_layer)
-            print(f"For Layer {i}'s `{decoder_class_name}`: {result}")
-    return model_layers
-```
-The `monkey_patching` function primarily does three things:
-- Precisely locate the `forward` function of all instances of the `XXXAttention` class.
-- Replace the `forward` function with the `model_atten_forward` function you provide.
-- Return the corresponding properties of the decoder layers found during the process, typically of type `nn.ModuleList`. This return value (`model_layers`) is only used to determine the number of layers in the current model later on (obtained by `len(model_layers)`).
-In addition, the `monkey_patching` function replaces `transformers.generation.GenerationMixin._validate_model_kwargs` with our `_validate_model_kwargs` to bypass some parameter checks, as we will provide an additional `sepllm_kwargs` parameter to wrap the `input_ids` for eventual transmission to the `SepCache` `update` function.
-**Please ensure that the `monkey_patching` function accurately locates and replaces the `forward` function of the `XXXAttention` class. The current `monkey_patching` is designed for the `Llama 3 series` models. For other models, you need to appropriately modify `monkey_patching` to ensure its correctness of targeting and replacement !** You can monitor the monkey patching process by setting `verbose=True` in the `monkey_patching` function (or, `monkey_patch_verbose = True` for the `generate` function.)
-```python
-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
-```
-The `truncate_input_ids_4_autoregression` function in the `custom_generate/generate.py` file is used to shape the `input_ids` tensor to `[batch_size, 1]` during decoding.
-#### 2.2.5 Downstream Task Evaluation
-We recommend using `lm_eval==0.4.9` for downstream task evaluation. You can pass model-related parameters via `--model_args` and generation-related parameters (including those required for initializing `SepCache`) via `--gen_kwargs`. Notably, you typically need to pass a `list` to `separator_token_ids` using a string format like `"id1;id2;id3"` (as shown in the example below).
-```bash
-lm_eval --model hf \
-	--model_args pretrained=meta-llama/Meta-Llama-3-8B-Instruct,attn_implementation=flash_attention_2 \
-	--tasks    gsm8k_cot  \
-	--gen_kwargs custom_generate=transformers-community/sep_cache,trust_remote_code=True,monkey_patch_verbose=True,separator_token_ids="128000;13;11;30;0;26;25;198;220;662;1174;949;758;2652;551;720;256;262",PADDING_ID=128009\
-	--device cuda:0\
-	--batch_size 80 2>&1 | tee log.txt
-```
-Note: `SepCache` is typically used in combination with `Flash Attention` to maximize generation efficiency.
-#### 2.2.5 The Detailed Signature of `generate` Function
-Here is the detailed signature of our customized `generate` function for `SepCache` in `custom_generate/generate.py` file:
-```python
-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
-             ):
-             ...
-```
-### 3. Adaptation for Other Models
-Adapting `SepCache` to various models is simple - two approaches:
-#### 3.1 Method 1 - Monkey Patching
-- Modify the `monkey_patching` function to correctly locate and target the `forward` function of your model's `XXXAttention` class (e.g., `LlamaAttention` for Llame 3).
-- Write your custom `model_atten_forward` function and use `monkey_patching` to replace the `forward` function of all `XXXAttention` class instances. The key modification is passing `input_ids` to `SepCache`'s `update` function.
-#### 3.2 Method 2 - Direct Code Modification (Recommended for Simplicity)
-Simply edit your `modeling_xxx.py` file to implement:
-- Initialize `past_key_values` as a `SepCache` instance at the appropriate location (e.g., in `XXXForCausalLM` or `XXXModel` class' `forward` function).
-- Modify the `forward` function of the `XXXAttention` class to pass `input_ids` to `SepCache`'s `update` function.
- #### 3.3 Important Note
- The shape of `input_ids` is `[batch_size, seq_len]` during prefilling, and `[batch_size, 1]` during generation.
- ##  Other Advanced Usage
- Please refer to https://github.com/HKUDS/SepLLM, in which there are detailed explanations and examples.