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.cache=torch.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

问题：

1. 显存浪费：实际序列可能只有 100，但预留了 4096
2. 碎片化：多个请求并发时，大块连续内存难分配
3. 扩展性差：batch size 受限于显存，不能动态调整

二、PagedAttention 的核心思想

把 KV Cache 切成固定大小的 block（page），按需分配。逻辑上连续，物理上分散。

# PagedAttention 的内存管理BLOCK_SIZE=16# 每个 block 存 16 个 token 的 KVclassPagedKVCache:def__init__(self,num_blocks,block_size,num_heads,head_dim):# 预分配所有 blockself.kv_blocks=torch.zeros(num_blocks,# 总 block 数2,# K 和 Vblock_size,# 每个 block 的序列长度num_heads,head_dim).npu()# 空闲 block 池self.free_blocks=list(range(num_blocks))# 每个请求的 block 映射self.request_blocks={}# request_id -> [block_ids]

Block 映射示意图

请求 1：token 0-31（需要 2 个 block） request_blocks[1] = [0, 1] 请求 2：token 0-15（需要 1 个 block） request_blocks[2] = [2] 请求 3：token 0-47（需要 3 个 block） request_blocks[3] = [3, 4, 5] Block 池状态： 已使用：[0, 1, 2, 3, 4, 5] 空闲：[6, 7, 8, ...]

三、Block 分配与释放

分配 Block

defallocate_block(self,request_id):"""为请求分配一个新的 block"""ifnotself.free_blocks:raiseRuntimeError("No free blocks available")block_id=self.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):"""根据序列长度分配足够的 block"""num_blocks=(seq_len+BLOCK_SIZE-1)//BLOCK_SIZEfor_inrange(num_blocks):self.allocate_block(request_id)returnself.request_blocks[request_id]

释放 Block

deffree_request_blocks(self,request_id):"""请求结束后释放所有 block"""ifrequest_idnotinself.request_blocks:returnblock_ids=self.request_blocks[request_id]self.free_blocks.extend(block_ids)delself.request_blocks[request_id]print(f"Freed{len(block_ids)}blocks for request{request_id}")

四、Attention 计算的分块实现

PagedAttention 的核心是分块计算 Attention，不需要把整个 KV Cache 搬到一起。

标准 Attention

defstandard_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] """# 算整个序列的注意力scores=torch.matmul(query,key_cache.transpose(-1,-2))scores=scores/math.sqrt(head_dim)probs=torch.softmax(scores,dim=-1)output=torch.matmul(probs,value_cache)returnoutput

问题：key_cache和value_cache是整个序列，显存占用大。

PagedAttention

defpaged_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_dim=query.shape output=torch.zeros_like(query)forbinrange(batch_size):seq_len=context_lens[b]num_blocks=(seq_len+block_size-1)//block_size# 逐 block 计算注意力forblock_idxinrange(num_blocks):physical_block=block_tables[b,block_idx]# 获取当前 block 的 K 和 Vk_block=kv_blocks[physical_block,0]# [block_size, num_heads, head_dim]v_block=kv_blocks[physical_block,1]# 计算当前 block 的注意力分数block_scores=torch.matmul(query[b].unsqueeze(0),# [1, num_heads, head_dim]k_block.transpose(-1,-2)# [num_heads, head_dim, block_size])# 处理最后一个 block 的 paddingifblock_idx==num_blocks-1:valid_len=seq_len-block_idx*block_size block_scores[:,:,valid_len:]=float('-inf')# 累加到输出block_probs=torch.softmax(block_scores,dim=-1)output[b]+=torch.matmul(block_probs,v_block).squeeze(0)returnoutput

昇腾优化版本

实际实现中，用 Ascend C 写 kernel 更高效：

// PagedAttention kernel（简化版）template<typenameT>__aicore__voidPagedAttentionKernel(LocalTensor<T>query,// 当前 token 的 queryGlobalTensor<T>kv_blocks,// 所有 KV blockLocalTensor<int32_t>block_ids,// 当前请求的 block 映射int32_tnum_blocks,// 当前请求的 block 数int32_tblock_size,LocalTensor<T>output// 输出){// 1. 初始化累加器LocalTensor<T>acc=GetBuffer<T>(output_size);LocalTensor<float>exp_sum=GetBuffer<float>(1);exp_sum[0]=0.0f;// 2. 遍历每个 blockfor(inti=0;i<num_blocks;i++){intblock_id=block_ids[i];// 3. 从 GM 搬运当前 block 的 K/V 到 UBLocalTensor<T>k_block=GetBuffer<T>(block_size*head_dim);LocalTensor<T>v_block=GetBuffer<T>(block_size*head_dim);CopyIn(kv_blocks[block_id][0],k_block);CopyIn(kv_blocks[block_id][1],v_block);// 4. 计算注意力分数LocalTensor<T>scores=MatMul(query,k_block.T());// 5. Softmax（需要跨 block 累加）LocalTensor<float>exp_scores=Exp(scores);exp_sum[0]+=ReduceSum(exp_scores);// 6. 加权求和LocalTensor<T>weighted=MatMul(exp_scores,v_block);acc+=weighted;}// 7. 归一化output=acc/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_blocks=total_memory/(bs*kv_size_per_token)print(f"Block size{bs}: avg waste={avg_waste}, max blocks={num_blocks}")

实测数据

Block size 越小，显存利用率越高，但计算效率越低（更多的 kernel 启动开销）。通常选择 16 或 32 是平衡点。

六、与 vLLM 的对比

vLLM 是最早实现 PagedAttention 的开源项目，昇腾的实现参考了它的设计：

# vLLM 风格的 APIfromascend_transformer_boostimportPagedAttention# 创建 PagedAttention 实例paged_attn=PagedAttention(num_heads=32,head_dim=128,block_size=16,num_blocks=1024)# 分配 blockblock_ids=paged_attn.allocate(request_id=1,seq_len=128)# 执行 Attentionoutput=paged_attn.forward(query,key,value,block_ids)# 释放 blockpaged_attn.free(request_id=1)

七、实际性能对比

LLaMA-7B，A100 vs 昇腾 910，batch=16，序列=2048：

昇腾的 PagedAttention 实现与 vLLM 性能接近，显存利用率都能到 90% 以上。

参考资源

• vLLM 论文：https://arxiv.org/abs/2309.06180
• PagedAttention 源码：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 的性能接近。