昇腾CANN ATB KV Cache 与 PagedAttention：显存碎片消除的完整方案

LLM 推理的最大瓶颈不是计算——是显存。长上下文下KV Cache 的显存占用是二次增长的seq_len128K → KV Cache 128K × 每层 KV 大小 128K × (2 × hidden × head_num) 128K × 2 × 8192 × 32 32GB。加上模型参数70B × 2bytes 140GB→ 总共 172GB → Ascend 910 只有 128GB → OOM。ATB 用 PagedAttention 虚拟内存管理解决这个问题把 KV Cache 分页存储Page Table不连续分配按需申请页面。像操作系统管理虚拟内存一样管理 KV Cache。KV Cache 的显存碎片问题标准 KV Cache 是连续分配的三维张量 [batch, seq_len, hidden, head_num×2]连续分配 KV Cache 的问题 请求 1seq_len128K → 需要 32GB 连续块 请求 2seq_len512 → 需要 16MB 连续块 请求 3seq_len64K → 需要 16GB 连续块 ... 32GB 16MB 16GB ... 180GB 128GB 即使空闲总量够很多小请求释放后但无法分配 32GB 连续块 → OOM对比 PagedAttentionPagedAttention 方式 KV Cache 被分成 16KB 的 pages每 16KB1 page 请求 1 的 32GB → 分配 32GB/16KB 2,097,152 pages 请求 2 的 16MB → 分配 16MB/16KB 1024 pages 请求 3 的 16GB → 分配 16GB/16KB 1,048,576 pages ... page 不需要连续碎片不再是问题——任何空闲 page 都能分配ATB 的 PagedAttention 实现// ascend-transformer-boost/memory/paged_attention.cppclassPagedAttentionMemory{private:// 全局 page 池所有请求共享staticconstexprintPAGE_SIZE16*1024;// 16KBstructPage{intid;// page 编号全局唯一DevicePtr ptr;// page 在 HBM 上的地址boolallocated;// 是否已分配intref_count;// 引用计数多请求共享};std::vectorPageglobal_page_pool_;// 全局 page 池inttotal_pages_;// 总 page 数 HBM 大小 / PAGE_SIZE// 每个请求的 page 表structPageTable{std::vectorintpage_ids;// 虚拟地址 → 物理 page 映射intnum_pages;// 已分配 page 数intseq_len;// 当前序列长度};std::unordered_mapint,PageTablerequest_page_tables_;// request_id → page 表public:// 分配 pages StatusAllocatePages(intrequest_id,intnum_pages_needed){PageTableptrequest_page_tables_[request_id];for(inti0;inum_pages_needed;i){intpage_idFindFreePage();if(page_id-1){returnStatus::OUT_OF_MEMORY;// 没有空闲 page}// 分配 pageglobal_page_pool_[page_id].allocatedtrue;global_page_pool_[page_id].ref_count1;pt.page_ids.push_back(page_id);pt.num_pages;}pt.seq_lennum_pages_needed*(PAGE_SIZE/sizeof(float16)/(hidden*2));returnStatus::OK;}// 逻辑地址 → 物理地址转换 DevicePtrLogicalToPhysical(intrequest_id,intlogical_offset){PageTableptrequest_page_tables_[request_id];// 计算逻辑偏移在哪一页和页内偏移intpage_indexlogical_offset/PAGE_SIZE;intoffset_in_pagelogical_offset%PAGE_SIZE;// 从 page 表查询物理地址intphysical_page_idpt.page_ids[page_index];DevicePtr physical_pageglobal_page_pool_[physical_page_id].ptr;returnphysical_pageoffset_in_page;}// 释放 pages请求完成或溢出voidFreePages(intrequest_id){PageTableptrequest_page_tables_[request_id];for(intpage_id:pt.page_ids){global_page_pool_[page_id].ref_count--;if(global_page_pool_[page_id].ref_count0){global_page_pool_[page_id].allocatedfalse;// 真正释放}}pt.page_ids.clear();pt.num_pages0;}};PagedAttention 的 Attention 计算修改标准注意力计算完整 KV Cache// 标准 Attentionfor(intk0;kseq_len;k){floatscoredot(Q[token],K[k]);// Q 与 K 的一维点积softmax_scores[k]exp(score);}PagedAttention 计算按页计算// ascend-transformer-boost/kernels/paged_attention_kernel.cpp__aicore__voidPagedAttentionKernel(GlobalTensorfloat16Q,// [batch, num_heads, d_head]GlobalTensorfloat16K_pages,// [total_pages, page_size]GlobalTensorfloat16V_pages,// [total_pages, page_size]GlobalTensorfloat16output,// [batch, num_heads, d_head]GlobalTensorintpage_table,// [request_id, max_pages]intnum_pages,inthead_dim){intrequest_idblockIdx.x;// 每个 block 处理一个请求inthead_idthreadIdx.y;// 每个 thread 处理一个注意力头// 初始化累加器LocalTensorfloat16O_local(head_dim);for(intd0;dhead_dim;d)O_local[d]0.0f;floatmax_val-65504.0f;floatsum_exp0.0f;// 逐页计算 Attentionfor(intp0;pnum_pages;p){intphysical_page_idpage_table[request_id*MAX_PAGESp];// 加载一页的 K 和 V连续访问——物理地址LocalTensorfloat16K_page(page_size);LocalTensorfloat16V_page(page_size);DataCopy(K_page,K_pagesphysical_page_id*page_size,page_size);DataCopy(V_page,V_pagesphysical_page_id*page_size,page_size);// 计算 QK^T 在这一页的分数for(inti0;ipage_tokens;i){// K_page[i] 是 K[token_i]计算 Q·K[token_i]floatscore0.0f;for(intd0;dhead_dim;d){scorefloat(Q[head_id*head_dimd])*float(K_page[i*head_dimd]);}// Online softmax逐页更新floatexp_scoreexpf(score-max_val);// 如果这一页有更大的 score → 重新标定累加器if(scoremax_val){floatold_maxmax_val;max_valscore;// 重新标定之前累加的 O_local 和 sum_expfloatcorrectionexpf(old_max-max_val);for(intd0;dhead_dim;d){O_local[d]O_local[d]*correction;}sum_expsum_exp*correction;}// 累加 V * softmax_scorefor(intd0;dhead_dim;d){O_local[d]V_page[i*head_dimd]*exp_score;}sum_expexp_score;}}// 归一化for(intd0;dhead_dim;d){output[request_id*head_dimd]float16(O_local[d]/sum_exp);}}PagedAttention 的关键page 表中的物理地址是离散的但每个 page 内部的访问是连续的。分页解决了碎片不会降低注意力计算的性能因为每个 page 内部依然是连续访问。page 分配的贪心策略// ascend-transformer-boost/memory/page_allocator.cppclassPageAllocator{private:intFindFreePage(){// 贪心找第一个空闲 pagefor(inti0;itotal_pages_;i){if(!global_page_pool_[i].allocated){returni;}}return-1;// 无空闲}// 预取策略预测下一个 page 位置intPrefetchNextPage(intcurrent_page_id){// 如果当前 page 后一个也是该请求的 → 预取减少延迟intnext_pagecurrent_page_id1;if(next_pagetotal_pages_!global_page_pool_[next_page].allocated){PrefetchToCache(next_page);// 预取到 SRAM}}public:// 批量预取所有已分配 pagevoidPrefetchAllPages(intrequest_id){PageTableptrequest_page_tables_[request_id];for(intpage_id:pt.page_ids){PrefetchToCache(page_id);}}};踩坑一page 表查找的延迟PagedAttention 需要频繁查 page 表每次访问 K/V 都要逻辑→物理转换。page 表本身在 HBM 中——每次查表都是 HBM 访问。修复把 page 表拷贝到 L1 缓存// 加速 page 表查找__aicore__voidFastPageLookup(GlobalTensorintpage_table_in_hbm,// page 表在 HBM 中LocalTensorintpage_table_in_l1,// 拷贝到 L1intnum_pages){// 拷贝 page 表到 L1一次性把所有 page 的映射都搬上来DataCopy(page_table_in_l1,page_table_in_hbm,num_pages*sizeof(int));// 之后所有查表都在 L1 中——延迟 1 cycle不是 HBM 的百 cycle}L1 中的 page 表查表延迟1 cycle。HBM 中查表延迟~300 cycles。PagedAttention 每页查一次表——page2MB → 查表延迟节省 2MB × (300-1) ~600M cycles。踩坑二page 引用计数泄漏多个请求可能共享相同的 K/V pages如共享前缀。引用计数减到 0 才真正释放。但如果忘记减引用计数——page 永远不释放 → 内存泄漏。// 引用计数的正确管理classRefCountManager{public:// 分配ref_count 1新请求独占voidAllocPage(intpage_id){global_page_pool_[page_id].ref_count1;}// 共享ref_count其他请求加入voidSharePage(intpage_id,intrequest_id){global_page_pool_[page_id].ref_count;// 记录哪几个请求在共享这个 pageshared_requests_[page_id].push_back(request_id);}// 释放ref_count--只有变成 0 才释放voidReleasePage(intpage_id,intrequest_id){global_page_pool_[page_id].ref_count--;if(global_page_pool_[page_id].ref_count0){// 真正释放标记为可用global_page_pool_[page_id].allocatedfalse;shared_requests_[page_id].clear();}}// 校验多请求释放时的安全检查voidValidateRefCount(intpage_id,intrequest_id){autosharedshared_requests_[page_id];if(std::find(shared.begin(),shared.end(),request_id)shared.end()){// 这个请求没有共享这个 page → 不应该减引用计数throwRefCountError(request not in shared list);}}};踩坑三page 表更新时的时序竞争推理过程中Decoder 生成新 token 时KV Cache 需要扩展添加新的 K, V。如果此时上一个请求的 page 正在被 Attention 计算读 → 数据竞争。方案Copy-on-WriteCoW// Copy-on-Write page 更新StatusExtendKVPage(intrequest_id,intnew_page_id){PageTableptrequest_page_tables_[request_id];intold_page_idpt.page_ids.back();// 如果只有这个请求在用这个 page → 直接更新if(global_page_pool_[old_page_id].ref_count1){// 无竞争直接覆盖旧 pageglobal_page_pool_[old_page_id].allocatedtrue;returnStatus::OK;}// 多个请求在共享这个 page → Copy-on-Write// 分配新 page拷贝旧内容写入新数据intnew_pageFindFreePage();if(new_page-1)returnStatus::OUT_OF_MEMORY;// CoW拷贝旧 page 到新 pagememcpy(global_page_pool_[new_page].ptr,global_page_pool_[old_page_id].ptr,PAGE_SIZE);// 在新 page 上追加 K,V 数据WriteKV(global_page_pool_[new_page].ptr,new_K,new_V);// 更新 page 表pt.page_ids.back()new_page;// 释放旧 page 的引用ReleasePage(old_page_id,request_id);returnStatus::OK;}KV Cache 是 LLM 推理中最大的显存消耗者——128K 上下文下占 32GB。ATB 的 PagedAttention 把连续分配变成分页分配page 池全局共享、page 表做逻辑→物理映射、Copy-on-Write 解决共享页的更新冲突。像操作系统管理虚拟内存一样管理 KV Cache——碎片不再导致 OOM。