深入理解Attention Sinks的KV缓存机制：代码级实现分析

深入理解Attention Sinks的KV缓存机制：代码级实现分析

【免费下载链接】attention_sinksExtend existing LLMs way beyond the original training length with constant memory usage, without retraining项目地址: https://gitcode.com/gh_mirrors/at/attention_sinks

在大型语言模型（LLM）的实际应用中，长文本处理一直面临着内存限制的挑战。Attention Sinks技术通过创新的KV缓存管理机制，实现了在不重新训练模型的情况下突破原始训练长度限制，同时保持恒定的内存占用。本文将从代码实现角度解析这一机制的核心原理与工作流程。

核心原理：Attention Sinks如何突破长度限制？

传统Transformer模型在生成文本时，会将所有先前的token信息存储在KV缓存中，导致内存占用随序列长度线性增长。Attention Sinks提出了一种巧妙的解决方案：

• 固定缓存大小：将缓存划分为两部分——少量固定的"注意力汇点"（Attention Sinks）和可滑动的窗口区域
• 选择性保留：始终保留初始几个token（注意力汇点）和最近的窗口区域token，丢弃中间部分
• 位置偏移补偿：通过位置编码偏移技术，确保模型感知的位置信息与实际缓存内容一致

这种设计使模型能够处理远超原始训练长度的文本，同时内存使用量保持恒定。

KV缓存管理的核心实现

AttentionSinkKVCache类：缓存管理的核心

attention_sinks/attention_sink_kv_cache.py文件实现了缓存管理的核心逻辑。该类通过__call__方法实现了KV缓存的动态更新：

def __call__(self, past_key_values): if past_key_values is None: return None seq_len = past_key_values[0][0].size(self.k_seq_dim) if seq_len <= self.cache_size: return past_key_values return [ [ torch.cat( [ self.k_slice(k, 0, self.attention_sink_size), # 保留注意力汇点 self.k_slice(k, seq_len - self.attention_sink_window_size, seq_len), # 保留最新窗口 ], dim=self.k_seq_dim, ), torch.cat( [ self.v_slice(v, 0, self.attention_sink_size), # 保留注意力汇点 self.v_slice(v, seq_len - self.attention_sink_window_size, seq_len), # 保留最新窗口 ], dim=self.v_seq_dim, ), ] for k, v in past_key_values ]

关键参数配置：

• attention_sink_size: 注意力汇点的token数量（默认4）
• attention_sink_window_size: 滑动窗口的token数量（默认1020）
• cache_size: 总缓存大小 = 汇点大小 + 窗口大小

缓存驱逐策略

除了基本的缓存更新，该类还实现了两种驱逐策略：

1. 空间预分配驱逐：在新token加入前提前为其预留空间

def evict_for_space(self, past_key_values, num_coming): # 计算新token加入后的序列长度，提前调整缓存

2. 范围驱逐：手动指定需要删除的token范围

def evict_range(self, past_key_values, start, end): # 精确删除指定范围内的token缓存

模型注入机制：如何与现有LLM架构集成？

attention_sinks/inject_mixin.py实现了将Attention Sinks功能注入现有模型的关键逻辑。通过InjectAttentionSinksMixin类，实现了对多种模型架构的适配：

模型类型映射

该文件定义了支持的模型类型及其对应的注意力机制类：

MODEL_NAME_MAPPING = { "llama": "LlamaModel", "falcon": "FalconModel", "mpt": "MptModel", "gpt_neox": "GPTNeoXModel", "gptj": "GPTJModel", "mistral": "MistralModel", "qwen": "QWenModel", # 其他模型... }

注入流程

1. 从预训练模型加载：通过重写from_pretrained方法，在模型加载时注入功能
2. 分离注意力汇点参数：提取以attention_sink开头的参数
3. 注入位置偏移注意力：根据模型类型替换注意力计算的前向方法
4. 注入KV缓存管理：为模型添加AttentionSinkKVCache实例
5. 重写生成 kwargs 更新方法：确保生成过程中模型参数正确更新

核心注入代码：

def overwrite_forward(module): # 创建新的缓存实例 module.attention_sink_kv_cache = AttentionSinkKVCache(**attention_sink_kwargs) # 保存旧的前向方法 old_forward = module.forward # 包装前向方法，更新past_key_values def wrapped_forward(self, *args, **kwargs): outputs = old_forward(*args, **kwargs) outputs.past_key_values = self.attention_sink_kv_cache(outputs.past_key_values) return outputs module.forward = types.MethodType(wrapped_forward, module)

实际应用：如何使用Attention Sinks？

在实际应用中，使用Attention Sinks非常简单。以演示脚本demo/endless_generation.py为例：

# 导入注意力汇点版本的AutoModelForCausalLM from attention_sinks import AutoModelForCausalLM # 加载模型时指定注意力汇点参数 model = AutoModelForCausalLM.from_pretrained( model_name_or_path, attention_sink_size=4, # 注意力汇点数量 attention_sink_window_size=1020 # 窗口大小 )

通过这种方式加载的模型会自动应用Attention Sinks机制，实现超长文本生成。

支持的模型架构

目前Attention Sinks支持多种主流LLM架构，包括：

• Llama/Llama-2系列
• Falcon系列
• GPT-NeoX
• GPT-J
• Mistral
• Qwen
• MPT
• StableLM
• Yi

每种模型的适配代码位于attention_sinks/models/目录下，如attention_sinks/models/llama/、attention_sinks/models/mistral/等。

性能评估与基准测试

项目提供了完整的性能评估工具，位于benchmark/perplexity.py。通过该脚本可以对比不同缓存策略的困惑度（perplexity）表现：

python benchmark/perplexity.py --experiment attention_sinks

测试结果会保存在benchmark/outputs_*/目录下，以CSV格式存储，便于分析不同模型在使用Attention Sinks时的性能表现。

总结：Attention Sinks的价值与应用前景

Attention Sinks技术通过创新的KV缓存管理策略，解决了LLM长文本处理中的内存瓶颈问题。其核心优势包括：

• 无需重新训练：直接在预训练模型上注入功能
• 内存占用恒定：不随输入长度增加而增长
• 广泛的模型支持：适配多种主流LLM架构
• 易于使用：通过简单API集成到现有工作流

这一技术为需要处理超长文本的应用场景（如文档摘要、代码生成、日志分析等）提供了强大支持，有望在LLM部署中发挥重要作用。随着研究的深入，我们可以期待更多优化和改进，进一步提升长文本处理能力。

【免费下载链接】attention_sinksExtend existing LLMs way beyond the original training length with constant memory usage, without retraining项目地址: https://gitcode.com/gh_mirrors/at/attention_sinks