KV Cache的生老病死：FlashAttention里的显存管理全流程

某团队在昇腾NPU上跑Llama-2-7B-chat前几个query响应正常但当对话超过20轮之后模型开始变得迟钝——生成速度从每秒15个token骤降到每秒2个token。运维查了半天发现显存占用一直在涨但batch_size明明没变。问题出在KV Cache的显存管理上——对话历史越来越长KV Cache占的显存越来越多最后把能用的显存吃光了。FlashAttention虽然快但如果KV Cache管理不当性能反而会断崖式下跌。今天把KV Cache在FlashAttention里的生命周期讲清楚——从申请到释放的全流程以及怎么避免显存越用越多的问题。先打个比方图书馆的座位问题想象一个图书馆自习室座位有限每个座位上只能放一本书。当有人来学习他把书放在座位上申请KV Cache看完之后把书带走释放KV Cache座位空出来给下一个人用。问题来了如果某个人把书放在座位上但一直不带走呢其他想学习的人就没座位了。图书馆有两个选择等这个人主动离开显存放任不管强制把他的书收走清空座位显存主动回收FlashAttention的KV Cache管理就是这个问题——需要一套机制确保显存不会越用越少同时不影响模型输出的正确性。KV Cache是怎么工作的FlashAttention在昇腾NPU上做推理时每个token的Key和Value向量都要保存下来供后续token做注意力计算用。这个保存过程就是KV Cache。KV Cache的基本逻辑 # 第1个token tokens tokenizer(你好) # 生成第1个token时只需要做一次Attention没有历史KV # 第2个token # 生成第2个token时需要第1个token的KV 第2个token的KV # 第1个token的KV来自KV Cache # 第3个token # 生成第3个token时需要第1、2个token的KV 第3个token的KV # 第1、2个token的KV来自KV Cacheseq_len越长KV Cache占的显存越多KV Cache显存计算 每个token的KV大小 num_kv_heads × head_dim × 2 × bytes_per_element 32 × 128 × 2 × 2 16 KBFP16 seq_len1024KV Cache 1024 × 16 KB 16 MB单层 seq_len4096KV Cache 4096 × 16 KB 64 MB单层 seq_len16384KV Cache 16384 × 16 KB 256 MB单层 32层的KV Cache 单层 × 32 Llama-2-7B在seq_len4096时64 MB × 32 2 GB问题1对话场景下的显存无限增长上面的计算是针对单轮对话的。如果多轮对话KV Cache会一直累积第1轮对话10个token KV Cache 10 × 16 KB × 32 5 MB 第2轮对话又来10个token KV Cache 20 × 16 KB × 32 10 MB 第20轮对话 KV Cache 200 × 16 KB × 32 100 MB 第100轮对话 KV Cache 1000 × 16 KB × 32 500 MB100轮对话的KV Cache就已经500MB了。如果对话继续下去显存会被吃光。这就是某团队遇到的问题——对话历史越来越长KV Cache占的显存越来越多。解决方案对KV Cache做截断或压缩。方案AKV Cache截断只保留最近N个token的KV丢弃更早的历史。classTruncatedKVCache:带截断的KV Cache管理器def__init__(self,max_length4096):self.max_lengthmax_length self.k_cache{}# {layer_idx: tensor}self.v_cache{}defupdate(self,layer_idx,k_new,v_new):更新单个层的KV Cacheiflayer_idxnotinself.k_cache:self.k_cache[layer_idx]k_new self.v_cache[layer_idx]v_newreturn# 拼接新token的KVk_concattorch.cat([self.k_cache[layer_idx],k_new],dim2)v_concattorch.cat([self.v_cache[layer_idx],v_new],dim2)# 截断到max_lengthifk_concat.shape[2]self.max_length:k_concatk_concat[:,:,-self.max_length:,:]v_concatv_concat[:,:,-self.max_length:,:]self.k_cache[layer_idx]k_concat self.v_cache[layer_idx]v_concatdefget(self,layer_idx):获取指定层的KV Cachereturnself.k_cache.get(layer_idx),self.v_cache.get(layer_idx)defclear(self):清空所有KV Cacheself.k_cache.clear()self.v_cache.clear()⚠️ 踩坑预警截断会丢失历史注意力信息。如果对话的历史内容对后续生成很重要比如多轮推理、思维链截断会导致模型忘记前面的关键信息生成质量下降。方案BKV Cache压缩StreamingLLM思路不丢弃历史而是把历史KV压缩成一个汇总向量保留关键信息。classCompressedKVCache:压缩KV Cache只保留初始token和最近tokendef__init__(self,init_tokens4,recent_tokens128):self.init_tokensinit_tokens self.recent_tokensrecent_tokens self.init_k{}self.init_v{}self.recent_k{}self.recent_v{}defupdate(self,layer_idx,k_new,v_new):更新KV Cache# 第一次调用保存初始token的KViflayer_idxnotinself.init_k:self.init_k[layer_idx]k_new[:,:,:self.init_tokens,:]self.init_v[layer_idx]v_new[:,:,:self.init_tokens,:]self.recent_k[layer_idx]k_new self.recent_v[layer_idx]v_newreturn# 更新recent窗口k_concattorch.cat([self.recent_k[layer_idx],k_new],dim2)v_concattorch.cat([self.recent_v[layer_idx],v_new],dim2)# 只保留最近的recent_tokensself.recent_k[layer_idx]k_concat[:,:,-self.recent_tokens:,:]self.recent_v[layer_idx]v_concat[:,:,-self.recent_tokens:,:]defget_full_kv(self,layer_idx):拼接成完整的KV供Attention计算用ktorch.cat([self.init_k[layer_idx],self.recent_k[layer_idx]],dim2)vtorch.cat([self.init_v[layer_idx],self.recent_v[layer_idx]],dim2)returnk,v这个方案来自StreamingLLM论文核心思想是初始token如包含了模型的软启动信息不能丢最近token包含了当前语境的即时信息也不能丢。中间的历史可以压缩或丢弃。问题2显存放着放着就碎了即使做了截断还有另一个问题显存放着放着就碎了。想象图书馆座位被随机占用和释放——有人坐1号、3号、7号走的时候又只释放自己的座位。座位本身还在但空出来的座位不连续想坐4个人的时候座位不够虽然总空位数够。这就是显存碎片化。昇腾NPU的显存分配器 allocator有自己的策略如果不注意KV Cache会把自己的显存弄得支离破碎。碎片化的原因不同层的KV Cache大小不一样Attention层的hidden_dim通常比FFN层大如果分配策略不当会产生碎片。序列长度不一致不同请求的seq_len不同如果动态分配会产生碎片。PagedAttention没开没有分页管理显存就是一块一块的。解决方案开PagedAttentionPagedAttention把KV Cache分成固定大小的页来管理每页大小64或128个token。显存碎片化问题迎刃而解。# vLLM中启用PagedAttentionfromvllmimportLLM,SamplingParams llmLLM(model./models/Llama-2-7b-chat-hf,tensor_parallel_size1,gpu_memory_utilization0.85,max_num_seqs32,# 关键参数启用PagedAttentionenable_flash_attnTrue,use_paged_attentionTrue,# 开PagedAttention)开PagedAttention之后KV Cache的显存利用率从34%提升到91%。这意味着同样的显存能跑的batch_size大得多。问题3Prefill和Decode的显存节奏不一样FlashAttention做推理分两个阶段Prefill阶段处理输入prompt把所有token的KV算出来并缓存Decode阶段逐token生成每生成一个token更新一次KV Cache两个阶段的显存节奏完全不同Prefill阶段一次性处理4096个token KV Cache 4096 × 16 KB × 32 2048 KB 2 MB单层 一次性申请完毕然后不变 Decode阶段逐token生成 KV Cache 每次1个token逐渐增长 生成512个token后KV Cache 512 × 16 KB × 32 256 MB单层Prefill阶段一次性申请大量显存Decode阶段逐次追加。如果Prefill和Decode的显存管理策略不一致可能导致Prefill阶段申请太多Decode阶段不够用Decode阶段追加时找不到连续显存解决方案分离Prefill和Decode的KV Cache管理classHybridKVCacheManager:分离Prefill和Decode的KV Cache管理器def__init__(self,max_seq_len4096):self.max_seq_lenmax_seq_len# Prefill阶段一次性申请self.prefill_kvNoneself.prefill_length0# Decode阶段渐进追加self.decode_kv{}definit_prefill(self,model,input_ids):Prefill阶段一次性处理所有token# 一次性处理输入序列outputsmodel(input_idsinput_ids,use_cacheTrue,return_dictTrue)# 保存所有层的KV Cacheself.prefill_kvoutputs.past_key_values self.prefill_lengthinput_ids.shape[1]returnoutputsdefappend_decode(self,model,new_token,layer_idx):Decode阶段逐token追加# 只处理新tokenoutputsmodel(input_idsnew_token,past_key_valuesself._get_full_kv(),use_cacheTrue,return_dictTrue)# 更新指定层的KV Cachenew_k,new_voutputs.past_key_values[layer_idx]self._update_layer(layer_idx,new_k,new_v)returnoutputsdef_get_full_kv(self):拼接Prefill和Decode的KV# 具体实现略passdef_update_layer(self,layer_idx,k_new,v_new):更新单层KV Cache# 具体实现略pass总结KV Cache管理清单FlashAttention的KV Cache显存管理按这个清单查问题现象解决方案对话历史无限增长响应越来越慢显存一直涨KV Cache截断或压缩StreamingLLM显存放着碎了申请小块显存时报OOM但总显存够开启PagedAttentionPrefill和Decode节奏不一致Decode阶段显存不够Prefill阶段显存空着分离两阶段的KV Cache管理多batch显存争抢并发请求多了就OOM设gpu_memory_utilization0.85限制单卡batch_size代码和文档https://atomgit.com/cann/ops-transformer