深入理解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), # 保留最新窗口 ], dimself.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), # 保留最新窗口 ], dimself.v_seq_dim, ), ] for k, v in past_key_values ]关键参数配置attention_sink_size: 注意力汇点的token数量默认4attention_sink_window_size: 滑动窗口的token数量默认1020cache_size: 总缓存大小 汇点大小 窗口大小缓存驱逐策略除了基本的缓存更新该类还实现了两种驱逐策略空间预分配驱逐在新token加入前提前为其预留空间def evict_for_space(self, past_key_values, num_coming): # 计算新token加入后的序列长度提前调整缓存范围驱逐手动指定需要删除的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, # 其他模型... }注入流程从预训练模型加载通过重写from_pretrained方法在模型加载时注入功能分离注意力汇点参数提取以attention_sink开头的参数注入位置偏移注意力根据模型类型替换注意力计算的前向方法注入KV缓存管理为模型添加AttentionSinkKVCache实例重写生成 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_size4, # 注意力汇点数量 attention_sink_window_size1020 # 窗口大小 )通过这种方式加载的模型会自动应用Attention Sinks机制实现超长文本生成。支持的模型架构目前Attention Sinks支持多种主流LLM架构包括Llama/Llama-2系列Falcon系列GPT-NeoXGPT-JMistralQwenMPTStableLMYi每种模型的适配代码位于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创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考