本文基于昇腾CANN和昇腾NPU围绕 cann-recipes-infer 仓库的相关技术展开。vLLM 是当前最流行的 LLM 推理框架它的核心发明是 PagedAttention——用操作系统分页的思想管理 KV Cache。CANN 上跑 vLLM 需要对接三个层面算子层PagedAttention 的 Ascend C 实现、Runtime 层显存管理、GE 层图执行调度。昇腾 CANN 开源社区有 vLLM 的适配分支核心改动在 Attention 后端。vLLM 的 PagedAttention 数据结构# vLLM 的 KV Cache——物理上分页逻辑上连续classPagedKVCache: 物理块: 每个块存固定数量的 Token通常是 16 或 32 逻辑页表: 跟操作系统一样虚拟地址 → 物理块号 def__init__(self,num_layers,num_heads,head_dim,block_size16,num_blocks4096):# 物理块池——一次性分配不碎片化# shape: [num_layers, 2, num_blocks, block_size, num_heads, head_dim]# 2 表示 K 和 Vself.block_pooltorch.empty(num_layers,2,num_blocks,block_size,num_heads,head_dim,dtypetorch.float16,devicenpu)# 空闲块链表——哪些块是可以分配的self.free_blockslist(range(num_blocks))# 每个请求的块表——逻辑块号 → 物理块号self.block_tables{}# request_id → [physical_block_ids]defallocate_slots(self,request_id,num_tokens): 为请求分配物理块 比连续显存分配快——块大小固定只需操作链表 num_blocks_needed(num_tokensblock_size-1)//block_size blocks[]for_inrange(num_blocks_needed):block_idself.free_blocks.pop(0)# O(1)blocks.append(block_id)self.block_tables[request_id]blocksreturnblocksdefread_kv(self,request_id,token_indices): 按逻辑索引读 KV——通过块表做地址转换 这是 Attention 计算里最频繁的操作 block_tableself.block_tables[request_id]k_out[]v_out[]foridxintoken_indices:block_idxidx//block_size offsetidx%block_size phys_blockblock_table[block_idx]# 读对应物理块的第 offset 个 Tokenk_out.append(self.block_pool[:,0,phys_block,offset])v_out.append(self.block_pool[:,1,phys_block,offset])returntorch.stack(k_out),torch.stack(v_out)PagedAttention 的核心优化在 Attention 计算时直接传物理块号和块表——不拼接连续 K、V避免显存搬运。CANN 上的 PagedAttention 算子// Ascend C 实现 PagedAttention——输入直接是块表classPagedAttentionKernel:publicAscendC::Kernel{__aicore__inlinevoidProcess()override{// 输入// q_heads: [num_heads, head_dim]// block_table: [num_blocks] ——物理块 ID 列表// kv_pool: [num_blocks, block_size, num_heads, head_dim]// actual_seq_len: 当前序列长度// 分块读取 KV——按物理块索引for(intbi0;binum_blocks;bi){intphys_blockblock_table[bi];// 从块池读取一块AscendC::LocalTensorfloat16_tk_block;AscendC::LocalAlloc(k_block,block_size*num_heads*head_dim);// 块池的物理地址 base block_id × block_size × head_dim × num_headsuint64_tk_addrkv_pool_basephys_block*block_stride;AscendC::DataCopy(k_block,k_addr,block_size*num_heads*head_dim);// Q K_block^T// 这里 q 是当前 Token 的 QueryK_block 是物理块里的完整数据// vLLM 的块大小 16——正好装满 Ascend C 的 L1 Bufferfor(inth0;hnum_heads;h){AscendC::MatMul(partial_score[h],q_heads[h],// [1, head_dim]k_block[h],// [block_size, head_dim]CUBE_MATRIX_TYPE::TRANS_B// → score: [1, block_size]);}}// 所有块累加后做 Softmax V// 因为 vLLM 的块是物理不连续的不能一次性做全局 Softmax// 只能用 Online-Softmax 逐块合并}};CANN 上实现 PagedAttention 的关键是把物理块的不连续性藏到 Ascend C 算子的 DataCopy 里——对算子的计算逻辑来说每块数据在 L1 上是连续的。vLLM 的 CANN 适配架构# vLLM 的 CANN 后端——替换 Attention 和显存管理classCannoVLLMBackend: vLLM 抽象了 ModelRunner 和 AttentionBackend。 CANN 适配需要实现 1. CannoAttentionBackend——替换 FlashAttention/PagedAttention 2. CannoModelRunner——替换 CUDA Graph 执行 def__init__(self):# 用 CANN Runtime 做显存分配fromvllm._Cimportcann_ops# PagedAttention 的自定义 opself.paged_attn_opcann_ops.paged_attention# 用 CANN 的 GE 做图执行——不依赖 CUDA Graphself.use_cann_graphTruedefexecute_model(self,model_input): vLLM 调度层→CANN 执行层 # vLLM 构造的 Batch 输入input_idsmodel_input[input_ids]# [batch_size]positionsmodel_input[positions]# [batch_size]block_tablesmodel_input[block_tables]# [batch, num_blocks]# 调用 AscendCL 执行模型# 这里不走 CUDA Graph而是用 CANN 的 GE 子图执行outputcann_inference(input_ids,positions,block_tables,kv_cacheself.kv_cache_block_pool)returnoutputCANN 适配后的 vLLM 在 Ascend 910 上跑 LLaMA-7B 的吞吐约 1800 tok/s连续 Batch 32比 CUDA 原版的 2200 tok/s 低约 18%。主要差距在 PagedAttention 的 Online-Softmax 实现上——CANN 的 Vector Unit 比 CUDA 的 Tensor Core 做逐块合并时要多搬一次数据。参考仓库vLLM CANN 适配PagedAttention 算子Runtime 显存管理