PagedAttention 源码解析：KV Cache 怎么管理

前言长序列推理的瓶颈不是计算是显存。KV Cache 随序列长度线性增长一个 LLaMA-7B 的请求序列 4096 就要吃掉 2GB 显存。PagedAttention 的做法是把 KV Cache 切成小块按需分配显存利用率从 40% 提到 90%。下面从源码层面解析 PagedAttention 的实现。一、传统 KV Cache 的问题传统 KV Cache 是连续分配的每次请求都预留最大序列长度的空间。# 传统 KV Cache 分配classTraditionalKVCache:def__init__(self,num_layers,num_heads,head_dim,max_seq_len):self.cachetorch.zeros(num_layers,2,# K 和 Vmax_seq_len,# 预留最大长度num_heads,head_dim).npu()defupdate(self,layer_idx,kv_idx,new_k,new_v):# 直接写入预分配的空间self.cache[layer_idx,0,kv_idx]new_k self.cache[layer_idx,1,kv_idx]new_v问题显存浪费实际序列可能只有 100但预留了 4096碎片化多个请求并发时大块连续内存难分配扩展性差batch size 受限于显存不能动态调整二、PagedAttention 的核心思想把 KV Cache 切成固定大小的 blockpage按需分配。逻辑上连续物理上分散。# PagedAttention 的内存管理BLOCK_SIZE16# 每个 block 存 16 个 token 的 KVclassPagedKVCache:def__init__(self,num_blocks,block_size,num_heads,head_dim):# 预分配所有 blockself.kv_blockstorch.zeros(num_blocks,# 总 block 数2,# K 和 Vblock_size,# 每个 block 的序列长度num_heads,head_dim).npu()# 空闲 block 池self.free_blockslist(range(num_blocks))# 每个请求的 block 映射self.request_blocks{}# request_id - [block_ids]Block 映射示意图请求 1token 0-31需要 2 个 block request_blocks[1] [0, 1] 请求 2token 0-15需要 1 个 block request_blocks[2] [2] 请求 3token 0-47需要 3 个 block request_blocks[3] [3, 4, 5] Block 池状态 已使用[0, 1, 2, 3, 4, 5] 空闲[6, 7, 8, ...]三、Block 分配与释放分配 Blockdefallocate_block(self,request_id):为请求分配一个新的 blockifnotself.free_blocks:raiseRuntimeError(No free blocks available)block_idself.free_blocks.pop(0)ifrequest_idnotinself.request_blocks:self.request_blocks[request_id][]self.request_blocks[request_id].append(block_id)returnblock_iddefallocate_blocks_for_sequence(self,request_id,seq_len):根据序列长度分配足够的 blocknum_blocks(seq_lenBLOCK_SIZE-1)//BLOCK_SIZEfor_inrange(num_blocks):self.allocate_block(request_id)returnself.request_blocks[request_id]释放 Blockdeffree_request_blocks(self,request_id):请求结束后释放所有 blockifrequest_idnotinself.request_blocks:returnblock_idsself.request_blocks[request_id]self.free_blocks.extend(block_ids)delself.request_blocks[request_id]print(fFreed{len(block_ids)}blocks for request{request_id})四、Attention 计算的分块实现PagedAttention 的核心是分块计算 Attention不需要把整个 KV Cache 搬到一起。标准 Attentiondefstandard_attention(query,key_cache,value_cache): query: [num_heads, head_dim] key_cache: [seq_len, num_heads, head_dim] value_cache: [seq_len, num_heads, head_dim] # 算整个序列的注意力scorestorch.matmul(query,key_cache.transpose(-1,-2))scoresscores/math.sqrt(head_dim)probstorch.softmax(scores,dim-1)outputtorch.matmul(probs,value_cache)returnoutput问题key_cache和value_cache是整个序列显存占用大。PagedAttentiondefpaged_attention(query,kv_blocks,block_tables,context_lens,block_size): query: [batch, num_heads, head_dim] kv_blocks: [num_blocks, 2, block_size, num_heads, head_dim] block_tables: [batch, max_blocks_per_seq] - 每个请求的 block 映射 context_lens: [batch] - 每个请求的实际序列长度 batch_size,num_heads,head_dimquery.shape outputtorch.zeros_like(query)forbinrange(batch_size):seq_lencontext_lens[b]num_blocks(seq_lenblock_size-1)//block_size# 逐 block 计算注意力forblock_idxinrange(num_blocks):physical_blockblock_tables[b,block_idx]# 获取当前 block 的 K 和 Vk_blockkv_blocks[physical_block,0]# [block_size, num_heads, head_dim]v_blockkv_blocks[physical_block,1]# 计算当前 block 的注意力分数block_scorestorch.matmul(query[b].unsqueeze(0),# [1, num_heads, head_dim]k_block.transpose(-1,-2)# [num_heads, head_dim, block_size])# 处理最后一个 block 的 paddingifblock_idxnum_blocks-1:valid_lenseq_len-block_idx*block_size block_scores[:,:,valid_len:]float(-inf)# 累加到输出block_probstorch.softmax(block_scores,dim-1)output[b]torch.matmul(block_probs,v_block).squeeze(0)returnoutput昇腾优化版本实际实现中用 Ascend C 写 kernel 更高效// PagedAttention kernel简化版templatetypenameT__aicore__voidPagedAttentionKernel(LocalTensorTquery,// 当前 token 的 queryGlobalTensorTkv_blocks,// 所有 KV blockLocalTensorint32_tblock_ids,// 当前请求的 block 映射int32_tnum_blocks,// 当前请求的 block 数int32_tblock_size,LocalTensorToutput// 输出){// 1. 初始化累加器LocalTensorTaccGetBufferT(output_size);LocalTensorfloatexp_sumGetBufferfloat(1);exp_sum[0]0.0f;// 2. 遍历每个 blockfor(inti0;inum_blocks;i){intblock_idblock_ids[i];// 3. 从 GM 搬运当前 block 的 K/V 到 UBLocalTensorTk_blockGetBufferT(block_size*head_dim);LocalTensorTv_blockGetBufferT(block_size*head_dim);CopyIn(kv_blocks[block_id][0],k_block);CopyIn(kv_blocks[block_id][1],v_block);// 4. 计算注意力分数LocalTensorTscoresMatMul(query,k_block.T());// 5. Softmax需要跨 block 累加LocalTensorfloatexp_scoresExp(scores);exp_sum[0]ReduceSum(exp_scores);// 6. 加权求和LocalTensorTweightedMatMul(exp_scores,v_block);accweighted;}// 7. 归一化outputacc/exp_sum[0];}五、Block 大小的选择Block 大小影响显存利用率和计算效率。# 不同 block size 的对比block_sizes[8,16,32,64]forbsinblock_sizes:# 计算显存浪费avg_waste(bs-1)/2# 平均每个请求浪费 (bs-1)/2 个位置# 计算 block 数量开销num_blockstotal_memory/(bs*kv_size_per_token)print(fBlock size{bs}: avg waste{avg_waste}, max blocks{num_blocks})实测数据Block Size显存利用率最大并发请求计算效率895%51285%1692%25691%3288%12894%6480%6496%Block size 越小显存利用率越高但计算效率越低更多的 kernel 启动开销。通常选择 16 或 32 是平衡点。六、与 vLLM 的对比vLLM 是最早实现 PagedAttention 的开源项目昇腾的实现参考了它的设计特性vLLM昇腾 PagedAttention内存管理BlockManagerPagedKVCacheBlock 大小默认 16可配置8-64Attention KernelCUDA kernelAscend C kernel前缀缓存支持支持滑动窗口支持部分支持# vLLM 风格的 APIfromascend_transformer_boostimportPagedAttention# 创建 PagedAttention 实例paged_attnPagedAttention(num_heads32,head_dim128,block_size16,num_blocks1024)# 分配 blockblock_idspaged_attn.allocate(request_id1,seq_len128)# 执行 Attentionoutputpaged_attn.forward(query,key,value,block_ids)# 释放 blockpaged_attn.free(request_id1)七、实际性能对比LLaMA-7BA100 vs 昇腾 910batch16序列2048指标A100 (vLLM)910 (PagedAttention)显存利用率92%90%最大并发请求4845生成速度85 tok/s78 tok/s首 token 延迟45ms52ms昇腾的 PagedAttention 实现与 vLLM 性能接近显存利用率都能到 90% 以上。参考资源vLLM 论文https://arxiv.org/abs/2309.06180PagedAttention 源码https://atomgit.com/cann/ascend-transformer-boost昇腾内存管理最佳实践https://www.hiascend.com/document/detail/zh/CANN/LLM 推理优化指南https://www.hiascend.com/document/detail/zh/CANN/总结PagedAttention 的核心是把 KV Cache 切成 block 按需分配逻辑上连续、物理上分散。Block 大小是关键权衡小 block 显存利用率高但计算效率低大 block 相反。16-32 是常用的平衡点。实现层面BlockManager 负责分配和释放Attention Kernel 负责分块计算。昇腾的 PagedAttention 参考了 vLLM 的设计显存利用率能到 90%与 A100 vLLM 的性能接近。