Views
No views yet
SepCache is a simple yet effective, native sparse attention Cache class proposed in the SepLLM paper - ICML 2025, 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.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.
meta-llama/Meta-Llama-3-8B-Instruct, for which we have already prepared a targeted monkey patch.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.modeling_xxx.py file or monkey patch file to adapt SepCache to any model.transformers>=4.53.0,<4.54.0, and we recommend using lm_eval>=0.4.9 for running evaluations. We suggest managing your Python environment with conda for better dependency control.1conda create -n sepcache python=3.10
2conda activate sepcache
3pip install transformers==4.53
4pip install lm_eval==0.4.9SepCache 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.1# requires `transformers>=4.53.0,<4.54.0`
2from transformers import AutoModelForCausalLM, AutoTokenizer
3
4# Preparing model, tokenizer, and model inputs
5tokenizer = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3-8B-Instruct")
6model = AutoModelForCausalLM.from_pretrained("meta-llama/Meta-Llama-3-8B-Instruct", device_map="auto")
7
8
9messages = [{"role": "user", "content": "Tell me a story about a cat."}]
10text = tokenizer.apply_chat_template(
11 messages,
12 tokenize=False,
13 add_generation_prompt=True,
14 enable_thinking=False
15)
16model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
17
18
19# Using SepCache for generation
20gen_out = model.generate(
21 # usual `generate` arguments
22 **model_inputs,
23 do_sample=False,
24 max_new_tokens=100,
25 return_dict_in_generate=True,
26 monkey_patch_verbose = True, # To see which functions are actually being monkey patched for `SepCache`.
27
28 # Using SepCache
29 custom_generate="transformers-community/sep_cache", ## Alternatively, you can use `Gausson/sep_cache`
30 trust_remote_code=True,
31
32 # SepCache arguments
33 init_cache_size = 4,
34 sep_cache_size = 128,
35 local_size = 256,
36 cache_size = 512,
37 USE_MAX_SEP_CACHE = True,
38 model_type = 'llama'
39)
40
41print(tokenizer.batch_decode(gen_out.sequences, skip_special_tokens=True))
42assert "sepcache" in str(type(gen_out.past_key_values)).lower()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.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.1# Using SepCache for generation
2gen_out = model.generate(
3 # usual `generate` arguments
4 **model_inputs,
5 do_sample=False,
6 max_new_tokens=100,
7 return_dict_in_generate=True,
8 monkey_patch_verbose = True, # To see which functions are actually being monkey patched for `SepCache`.
9
10 # Using SepCache
11 custom_generate="transformers-community/sep_cache", ## Alternatively, you can use `Gausson/sep_cache`
12 trust_remote_code=True,
13
14 # SepCache arguments
15 init_cache_size = 4,
16 sep_cache_size = 128,
17 local_size = 256,
18 cache_size = 512,
19 USE_MAX_SEP_CACHE = True,
20 separator_token_ids = [128000, 13, 11, 30, 0, 26, 25, 198, 220, 662, 1174, 949, 758, 2652, 551, 720, 256,262],
21 PADDING_ID = 128009
22)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` separators' 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. 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.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.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.APPLY_PE_SHIFT=False (False is also the default setting) and APPLY_PES_INSIDE=False for initialization.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.1key_states, value_states = past_key_values.update(
2 key_states = key_states,
3 value_states = value_states,
4 input_ids = input_ids, ## required
5 layer_idx = layer_idx,
6 PREFILLING_FLAG = q_len > 1, ## `q_len` is the sequence length of the current `query_states`
7 )update function of SepCache mentioned in 2.2.4 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].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).1def monkey_patching(model_obj,
2 model_atten_forward , ## The `forward` function used to patch.
3 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.
4 possible_layers_names: List[str] = ["layers", "h" ], # In `XXXModel` class, the possible name of internal attribute for decoder layers, e.g., "layers", "h", etc.
5 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.
6 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.
7 verbose = True):
8
9 """
10 This `monkey_patching` function is to
11 - find the `forward` function of the `XXXAttention` class.
12 - replace all the related `forward` functions of the instances of `XXXAttention` class with `model_atten_forward`.
13 """
14
15 ## To avoid the argument check failure, i.e., let "sepllm_kwargs" pass the check.
16 transformers.generation.GenerationMixin._validate_model_kwargs = _validate_model_kwargs
17
18 ## Get inner model obj
19 inner_model_type = PreTrainedModel
20 inner_model = find_inner_attribute(model_obj, possible_inner_model_names, inner_model_type)
21
22 ## Get the decoder layers (`nn.ModuleList`) obj
23 layers_type = nn.ModuleList
24 model_layers = find_inner_attribute(inner_model, possible_layers_names, layers_type)
25
26 ## Replace all the related `forward` functions of XXXAttention class's instances.
27 for i, decoder_layer in enumerate(model_layers):
28 self_attn_module = find_attribute_name(decoder_layer, atten_attr_name_pattern_list, atten_attr_name_pattern_exclude, nn.Module)
29 result = monkey_patch_by_class_path(self_attn_module, model_atten_forward)
30 if verbose:
31 decoder_class_name = get_importable_class_path(decoder_layer)
32 print(f"For Layer {i}'s `{decoder_class_name}`: {result}")
33
34 return model_layersmonkey_patching function primarily does three things:forward function of all instances of the XXXAttention class.forward function with the model_atten_forward function you provide.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)).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.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.)1def truncate_input_ids_4_autoregression(input_ids, key_states):
2 if input_ids.shape[-1] != key_states.shape[-2]:
3 assert input_ids.shape[-1] >= key_states.shape[-2]
4 truncated_input_ids = input_ids[..., -key_states.shape[-2]: ]
5 return truncated_input_ids
6 else:
7 return input_idstruncate_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.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).1lm_eval --model hf \
2 --model_args pretrained=meta-llama/Meta-Llama-3-8B-Instruct,attn_implementation=flash_attention_2 \
3 --tasks gsm8k_cot \
4 --gen_kwargs custom_generate=transformers-community/sep_cache,trust_remote_code=True,monkey_patch_verbose=True,init_cache_size=4,sep_cache_size=128,local_size=256,cache_size=512,separator_token_ids="128000;13;11;30;0;26;25;198;220;662;1174;949;758;2652;551;720;256;262",PADDING_ID=128009 \
5 --device cuda:0\
6 --batch_size 80 2>&1 | tee log.txtSepCache is typically used in combination with Flash Attention to maximize generation efficiency.generate Functiongenerate function for SepCache in custom_generate/generate.py file:1def generate(model,
2 ## For SepCache
3 init_cache_size: Union[int, List] = 4,
4 sep_cache_size: Union[int, List] = 128,
5 local_size: Union[int, List]=256,
6 cache_size: Union[int, List]=512,
7 SEP_ACCUMULATION: bool = True,
8 USE_MAX_SEP_CACHE: bool = False,
9 SEP_PADDING_IN_BATCH: bool = False,
10 separator_token_ids: List[int] = None, ## required for initialization if `model_type` is not provided.
11 PADDING_ID: int = None, ## required for initialization if `model_type` is not provided.
12
13 ## For inheritance & initialization states
14 past_tok_ids: List[torch.Tensor] = None, ## It saves all the token ids corresponding to the saved KVs for all layers in SepCache.
15 key_cache: List[torch.Tensor] = None,
16 value_cache: List[torch.Tensor] = None,
17
18 ## For debugging
19 PRINT_KV_RATIO_INSIDE: bool = False,
20 print_KV_inside_per_steps: int = 1000,
21 _seen_tokens: int = 0,
22 _kept_kv_ratio: List[Tuple[int]] = None,
23
24 ### For positional encoding shifting
25 APPLY_PE_SHIFT: bool = False,
26 APPLY_PES_INSIDE: bool = False,
27 _shifted_position_ids: List[torch.Tensor] = None,
28 _rope_unsqueeze_dim: int = 1, ## The unsqueeze_dim when applying RoPE.
29 _rope_seq_dim: int=1, ## The seq_len dimension for the `cos` or `sin` tensors.
30 pe_scaling_factor:float = 1.0,
31 pe_dim:int=128, ## The number of dims for positional encoding. Typically, just set the `head_dim` to this.
32 max_position_embeddings: int = 8192,
33 base: int=10000, ## The base for RoPE.
34
35 ## For basic transformer architecture
36 k_seq_dim: int=2, ## The dimension for seq_len in key tensors
37 v_seq_dim: int=2, ## The dimension for seq_len in value tensors
38 layer_num: int = None, ## required for initialization
39
40 model_type: str = 'llama', ## The model type for running the example. choose from ['llama', 'pythia','falcon'].
41 device = None,
42
43 ## For verbosity of monkey patching
44 monkey_patch_verbose: bool = False,
45
46 **kwargs
47 ):
48 ...SepCache to various models is simple - two approaches:monkey_patching function to correctly locate and target the forward function of your model's XXXAttention class (e.g., LlamaAttention for Llama 3).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.modeling_xxx.py file to implement:past_key_values as a SepCache instance at the appropriate location (e.g., in XXXForCausalLM or XXXModel class' forward function).forward function of the XXXAttention class to pass input_ids to SepCache's update function.input_ids is [batch_size, seq_len] during prefilling, and [batch_size, 1] during generation.@inproceedings{chen2025sepllm,
title={{SepLLM: Accelerate Large Language Models by Compressing One Segment into One Separator}},
author={Chen, Guoxuan and Shi, Han and Li, Jiawei and Gao, Yihang and Ren, Xiaozhe and Chen, Yimeng and Jiang, Xin and Li, Zhenguo and Liu, Weiyang and Huang, Chao},
booktitle={International Conference on Machine Learning},
year={2025},
note={Also available at arXiv:2412.12094}
}