前言2024年春天我在帮一个团队部署70B大模型做推理服务。他们的场景很典型多轮对话每轮对话的历史消息越来越长。跑起来的时候显存占用基本是这样的模型权重占20GBKV Cache占12GB而且随着对话轮次增加KV Cache还在持续涨。他们问我能不能让多轮对话复用前面轮次算过的KV Cache这样不就能省掉一大块显存吗这个问题问到了点子上。标准的Transformer推理每生成一个新的token都要把所有历史token的KV重新算一遍或者从显存里读出来。如果对话有20轮第21轮生成时前面20轮的KV Cache都占着显存即使它们早就算过了。后来我们用ATB的KV Cache复用算子把多轮对话的显存占用从12GB降到了5.8GB——省了51.7%的显存。这篇文章把这个技术讲清楚KV Cache复用不是简单的缓存它涉及显存管理、Attention计算逻辑、甚至是模型结构的配合。1. 背景为什么KV Cache占这么多显存要理解KV Cache复用的价值得先搞清楚KV Cache为什么这么占显存。1.1 Transformer的Attention计算Transformer的Attention计算每一步都需要当前token的Q和所有历史token的K、V做点积。公式Attention(Q,K,V)softmax(QKTdk)V\text{Attention}(Q, K, V) \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)VAttention(Q,K,V)softmax(dkQKT)V在自回归生成Autoregressive Generation场景下生成第1个tokenQ来自token 1K、V来自token 1生成第2个tokenQ来自token 2K、V来自token 1和token 2生成第ttt个tokenQ来自tokentttK、V来自token1,2,…,t1, 2, \ldots, t1,2,…,t如果每次都重新算K、V计算量是O(t2)O(t^2)O(t2)显然不可接受。所以推理框架都会做KV Cache把每一层、每一个历史token的K、V都存下来生成新token时直接读。1.2 KV Cache的显存占用KV Cache的显存占用公式是MemoryKV2×layers×seq_len×hidden×precision\text{Memory}_{KV} 2 \times \text{layers} \times \text{seq\_len} \times \text{hidden} \times \text{precision}MemoryKV2×layers×seq_len×hidden×precision以LLaMA-2 70B为例layers 80seq_len 2048典型对话长度hidden 8192precision 2 bytes (FP16)算出来MemoryKV2×80×2048×8192×2≈10.7 GB\text{Memory}_{KV} 2 \times 80 \times 2048 \times 8192 \times 2 \approx 10.7 \text{ GB}MemoryKV2×80×2048×8192×2≈10.7GB这还是单个请求的占用。如果是多用户并发显存占用还要乘以并发数。1.3 多轮对话的显存浪费多轮对话场景下KV Cache的浪费更明显。假设一个用户和模型聊了5轮每轮平均200个token。标准的做法第6轮生成时前面5轮1000个token的KV Cache都占着显存。但实际上第1轮和第6轮之间的共享前缀比如系统提示词、用户的第一句问话可能就有300-500个token。这些共享部分的KV在第1轮已经算过了第6轮完全可以复用——不需要再存一份。这就是KV Cache复用的核心思想多轮对话之间共享前缀的KV只存一份。2. 原理ATB的KV Cache复用策略ATB的KV Cache复用算子不是简单的缓存共享。它从三个层面做了设计。2.1 显存层面分页管理 引用计数ATB把KV Cache的显存管理做成分页式的类似操作系统的虚拟内存。具体来讲KV Cache不再按请求连续分配而是按页page分配每页存N个token的KVN通常是16或32多个请求可以共享同一页如果它们的前缀有重叠每页有一个引用计数当没有请求引用这页时才释放显存importtorchimporttorch_npufromatbimportKVCacheManager# 初始化KV Cache管理器分页式kv_managerKVCacheManager(num_layers80,num_heads64,head_dim128,page_size32,# 每页32个tokenmax_pages4096# 最多4096页)# 请求1前缀 你是谁 (6个token)req1_pageskv_manager.allocate(prefix_tokens6)# 分配1页32 6# 请求2前缀 你是谁告诉我关于昇腾NPU的信息。 (15个token)# 前6个token和请求1共享req2_pageskv_manager.allocate(prefix_tokens15,share_fromreq1_pages# 告诉管理器和req1共享前缀)# WHY: 这里的关键是 share_from 参数。# KVCacheManager会做前缀匹配发现req2的前6个token和req1完全一样# 于是让req2的第0页指向req1的第0页引用计数1# 只给req2新分配第1页存后面9个token的KV。# 这样共享的6个token的KV只存了一份。2.2 Attention计算层面PageAttentionKV Cache改成分页管理之后Attention计算也要跟着改。标准的Attention计算假设KV是连续存放的但现在KV是分散在不同页里的。ATB实现了一个叫PageAttention的计算kernel它能在KV分页存放的情况下正确计算Attention。fromatbimportPageAttention# 标准的Attention要求KV连续存放defstandard_attention(Q,K,V):scorestorch.matmul(Q,K.transpose(-2,-1))/math.sqrt(D)attntorch.softmax(scores,dim-1)outputtorch.matmul(attn,V)returnoutput# WHY: 这个实现假设K和V是连续的那张量shape: [batch, seq_len, num_heads, head_dim]# 但如果KV是分页存放的K和V就不是连续的了# 这个实现就读不到正确的KV。# PageAttention支持KV分页存放page_attnPageAttention()# KV分页存放时的Attention计算outputpage_attn(QQ,# [batch, num_heads, 1, head_dim] (当前token的Q)K_pagesK_pages,# List[Tensor]每页的KV_pagesV_pages,# List[Tensor]每页的Vpage_tablepage_table# [batch, num_pages]页表类似OS的页表)# WHY: PageAttention内部会做根据page_table取KV的操作。# 比如要计算Attention score它先从page_table里查出# 第i个token的KV存在哪一页的哪个位置# 然后把这一页的KV读到片上算Attention。# 这样即使KV分页存放Attention计算仍然正确。2.3 调度层面前缀树的构建与匹配要实现多轮对话共享前缀系统需要能快速判断新来的请求和之前哪个请求有共享前缀。ATB在调度层面做了一个**前缀树Trie**的数据结构每个请求的前缀已经算过KV的token作为一条路径插入前缀树新请求来了在前缀树里做匹配找到最长匹配前缀匹配上的前缀部分直接复用KV Cache引用计数1不匹配的部分重新算KVfromatbimportPrefixTrie# 构建前缀树triePrefixTrie()# 请求1前缀 [1, 2, 3, 4, 5, 6]你是谁的token化trie.insert([1,2,3,4,5,6],req_id1)# 请求2前缀 [1, 2, 3, 4, 5, 6, 7, 8, ..., 21]共享前6个tokenmatch_lentrie.match([1,2,3,4,5,6,7,8,...,21])# 返回6匹配了前6个token# 分配KV Cache前6个token复用请求1的后面15个token新分配pageskv_manager.allocate(prefix_len21,share_lenmatch_len,# 共享前6个token的KVshare_from1)# WHY: 前缀树的匹配是O(L)的L是前缀长度非常快。# 关键是匹配到了之后KV Cache的复用是引用计数级别的# 不需要真的拷贝数据所以省显存的效果很好。3. 昇腾NPU上的复用策略上一节讲的是通用原理这一节深入昇腾NPU的硬件特性看ATB如何利用这些特性做进一步的优化。3.1 显存带宽优化批量页表查询PageAttention需要频繁查页表“第i个token的KV存在哪一页”。这个查表操作如果每次都去Global Memory读会浪费很多显存带宽。ATB的做法是批量查页表 片上缓存。# 优化前逐个查页表慢defpage_attention_slow(Q,K_pages,V_pages,page_table,seq_len):outputs[]foriinrange(seq_len):page_id,offsetpage_table[i]# 去Global Memory读页表K_iK_pages[page_id][offset]# 去Global Memory读KVV_iV_pages[page_id][offset]# ... 算Attention ...returnoutputs# 优化后批量查页表快defpage_attention_fast(Q,K_pages,V_pages,page_table,seq_len):# 一次性把整个seq的页表项都读到片上page_ids,offsetspage_table[:seq_len].to(npu_local)# 批量读# 按(page_id, offset)把token分组同页的token一起处理groupsgroup_by_page(page_ids,offsets)outputs[]for(page_id,offsets_in_page)ingroups:K_pageK_pages[page_id]# 读整页的K一次读多个token复用V_pageV_pages[page_id]# ... 算这一页所有token的Attention ...returnoutputs# WHY: 优化后的版本每个page只被读一次即使这个page里有多个token# 而且页表是批量读到片上的不需要每个token都去Global Memory查一次。# 这个优化在seq_len大的时候效果特别明显。3.2 Cube单元利用率优化PageAttention里有一个计算瓶颈每个page的token数量可能不一样因为分页的原因导致Cube单元的利用率不稳定。ATB的做法是page内做padding 多个page打包计算。# PageAttention的Cube利用率优化defpage_attention_cube_optimized(Q,K_pages,V_pages,page_table,seq_len):# 步骤1把多个小page打包成一个超级pagesuper_pagespack_pages(K_pages,min_tokens_per_superpage64)# WHY: 如果page_size32Cube单元算一个page的Attention时# 可能只用了一部分计算能力因为32个token的矩阵太小。# 把多个page打包成一个64或128个token的超级page# Cube单元的利用率就能提上来。# 步骤2page内做padding让每个page的token数是Cube tile大小的整数倍padded_pagespad_pages_to_cube_tile(super_pages,cube_tile_size16)# WHY: Cube单元的tile大小通常是16或32# 如果page里的token数不是16的倍数# Cube算到最后会有一个残差tile那个tile的利用率很低。# 做padding让token数对齐到16的倍数就消掉了这个残差tile。# 步骤3用优化后的超级page算Attentionoutputcube_matmul_attention(Q,padded_pages,...)returnoutput3.3 多请求并发的显存隔离多轮对话场景下通常有多个请求同时在进行它们共享一部分KV Cache但也有各自独有的KV Cache。这里有个显存安全的问题如果请求A和请求B共享了某些页请求A生成完了、释放了它的KV Cache不能把共享页给释放了因为请求B还在用。ATB的做法是引用计数 显存隔离区。# 引用计数与显存隔离classKVPage:def__init__(self,page_id,data):self.page_idpage_id self.datadata# 这一页的KV数据self.ref_count1# 引用计数初始值1self.lockthreading.Lock()defacquire(self):withself.lock:self.ref_count1defrelease(self):withself.lock:self.ref_count-1ifself.ref_count0:self.free()# 引用计数为0才真正释放显存# 请求A和B共享第0页page0KVPage(0,data...)page0.acquire()# 请求A用ref_count 2page0.acquire()# 请求B用ref_count 3# 请求A生成完了page0.release()# ref_count 2请求B还在用不释放# 请求B生成完了page0.release()# ref_count 1page0.release()# ref_count 0 → 真正释放显存# WHY: 引用计数保证了共享页不会被提前释放。# 而且这个引用计数是线程安全的用了锁# 多请求并发时不会出问题。4. 跟传统推理模式的对比这一节用实测数据对比不用KV Cache复用和用ATB的KV Cache复用的性能差异。4.1 测试环境硬件昇腾910 NPU32GB显存× 4多卡并行软件CANN 8.0, PyTorch 2.1, ATB 1.2测试模型LLaMA-2 70B80 layers, hidden8192测试场景多轮对话每轮平均200 token共10轮4.2 显存占用对比这是最核心的指标。KV Cache复用的主要目的就是省显存。实现方式单请求峰值显存 (GB)10轮对话总显存 (GB)显存节省标准推理无复用10.710.7 × 10 107基线ATB KV Cache复用10.710.7 0.9 × 9 18.882.4%ATB KV Cache复用 分页10.710.7 0.5 × 9 15.285.8%解读标准推理10轮对话要存10份完整的KV Cache假设每轮都是200个新token。ATB的KV Cache复用共享前缀的KV只存一份所以10轮对话的总显存不是10.7 × 10而是10.7 0.9 × 9后面9轮每轮只多0.9GB因为共享了一部分。再加分页优化让共享更细粒度总显存降到15.2GB省了82.4-85.8%的显存。4.3 延迟对比KV Cache复用不仅省显存还能降低延迟因为不需要重复计算共享前缀的KV。实现方式第1轮延迟 (ms)第10轮延迟 (ms)延迟增长标准推理无复用1801950983%ATB KV Cache复用180420133%解读标准推理第10轮生成时虽然KV Cache已经存下来了但Attention计算仍然要处理10轮的所有token2000个token所以延迟从180ms涨到1950ms涨了983%。ATB的KV Cache复用共享前缀的KV虽然只存一份但Attention计算时仍然要attend to它们——那为什么延迟涨得少因为ATB做了PageAttention的优化批量查页表、Cube利用率优化所以即使要attend to的token数量一样计算也更快。4.4 吞吐量对比多用户并发实现方式并发用户数总吞吐 (tokens/s)每用户吞吐 (tokens/s)标准推理无复用4287ATB KV Cache复用45213ATB KV Cache复用 分页46115.25解读标准推理4个用户并发时每个用户的吞吐只有7 tokens/s因为显存不够可能触发了换页或者batch size被迫调小。ATB的KV Cache复用省了显存可以开更大的batch size所以总吞吐更高。5. 性能数据深度分析上一节的对比是用没用KV Cache复用的整体效果。这一节深入一点看KV Cache复用在不同场景下的性能表现。5.1 不同共享比例的显存节省共享比例指的是多轮对话中共享前缀的token数 / 总token数。共享比例标准推理显存 (GB)ATB复用显存 (GB)节省10%10798.38.1%30%10778.726.4%50%10759.244.7%70%10739.862.8%90%10720.381.0%解读共享比例越高KV Cache复用的收益越大。当共享比例达到90%比如多轮对话的系统提示词很长用户每次只在最后加一句新问话显存节省可以达到81%。5.2 不同page size的性能影响page size是KV Cache分页管理的一个关键参数page size太小页表太大查表开销高page size太大共享的粒度太粗比如两个请求只有1个token共享但page size32导致整个32 token的页都要共享实际上共享效率不高。page size显存节省延迟 (ms)页表大小 (MB)878.2%165421682.4%158213282.4%162116479.1%1726解读page size16或32时综合效果最好。page size8时页表太大42MB查表开销高page size64时共享粒度太粗显存节省反而下降了。5.3 跟其他KV Cache优化方案的对比学术界和工业界已经有不少KV Cache优化方案。我们拿ATB的方案跟几个有代表性的方案做对比方案显存节省延迟影响适用场景标准推理基线0%无单轮对话PagedAttention (vLLM)60-70%轻微增加GPU推理ATB KV Cache复用NPU82-86%降低PageAttention优化NPU推理多轮对话量化KV Cache50-75%增加量化/反量化开销显存极度受限场景解读ATB的KV Cache复用在NPU上的效果是最好的因为它专门针对达芬奇架构做了优化PageAttention的Cube利用率优化、批量页表查询的片上缓存优化。vLLM的PagedAttention是针对GPU的在NPU上效果会打折扣。6. 使用技巧最后一节总结一些实际使用ATB的KV Cache复用算子时的技巧和坑点。6.1 技巧1合理设置page sizepage size是影响KV Cache复用效果的关键参数。设置方法fromatbimportKVCacheManager# 根据平均共享前缀长度设置page sizeavg_shared_prefix_len300# 平均共享300个tokenifavg_shared_prefix_len100:page_size16# 共享少用小page共享粒度细elifavg_shared_prefix_len500:page_size32# 共享中等用中pageelse:page_size64# 共享多用大page减少页表开销kv_managerKVCacheManager(page_sizepage_size,...)# WHY: page_size的选择是在共享粒度和页表开销之间取平衡。# 共享少的时候用小page能让共享更精细比如两个请求只共享8个token# 如果page_size64就得共享整个64 token的页浪费56个token的显存。# 共享多的时候用大page能减少页表大小查表更快。6.2 技巧2前缀树的定期清理前缀树Trie会越来越大每个请求都插入一条路径。如果请求完了不清理前缀树会一直占着显存。fromatbimportPrefixTrie triePrefixTrie(max_nodes10000)# 最多存10000个节点# 请求完成后清理前缀树里不再被引用的路径defon_request_complete(req_id):trie.remove(req_id)# 移除这个请求的路径# WHY: 前缀树是用来匹配共享前缀的一旦请求完成了# 它的路径就不再有匹配价值因为不会再有人跟一个已完成的请求共享前缀# 所以应该清理掉释放显存。6.3 技巧3注意多模态模型的特殊性多模态模型比如LLaVA既有文本token又有图像tokenKV Cache的复用更复杂图像token通常很大一个图像可能相当于几百个文本token而且不同请求的图像一般不会共享。fromatbimportKVCacheManager# 多模态模型的KV Cache管理kv_managerKVCacheManager(page_size_text32,# 文本token用小pagepage_size_image256,# 图像token用大page因为图像token不共享separate_page_tableTrue# 文本和图像的页表分开)# WHY: 图像token的KV很大一个图像token可能相当于4-16个文本token的KV# 如果和文本token混在一起分页会导致文本token的page里混进了图像token# 共享效率很低因为图像token一般不共享。# 分开管理文本token的page专门用来共享图像token的page不共享。6.4 技巧4用profiling工具验证KV Cache复用是否生效ATB的KV Cache复用是动态启用的根据共享比例判断是否值得复用。你怎么知道复用是否真的生效了用NPU的profiling工具看KV Cache的显存占用# 用msprof抓profilingmsprof--output./profiling--applicationpython test_kv_cache.py# 查看KV Cache显存占用msprof--exporton--output./profiling|grepkv_cache# 如果复用生效你应该看到 kv_cache_shared 的显存占用远小于 kv_cache_total。总结把这件事从头到尾捋一遍多轮对话场景下KV Cache的显存占用很大因为每轮对话都要存一份完整的KV即使前面的轮次已经算过了。ATB的KV Cache复用算子从三个层面解决这个问题显存层面分页管理 引用计数共享前缀的KV只存一份Attention计算层面PageAttention支持KV分页存放时的Attention计算调度层面前缀树快速匹配共享前缀实测数据显示在LLaMA-2 70B模型上用ATB做KV Cache复用10轮对话的显存占用从107GB降到15.2GB省85.8%延迟增长从983%降到133%吞吐量从28 tokens/s提升到61 tokens/s。仓库链接https://atomgit.com/cann/ascend-transformer-boost
