当前位置：首页 > news >正文

vLLM 在 CANN 上的推理优化

news 2026/5/23 20:18:17

本文基于昇腾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_size=16,num_blocks=4096):# 物理块池——一次性分配，不碎片化# shape: [num_layers, 2, num_blocks, block_size, num_heads, head_dim]# 2 表示 K 和 Vself.block_pool=torch.empty(num_layers,2,num_blocks,block_size,num_heads,head_dim,dtype=torch.float16,device="npu")# 空闲块链表——哪些块是可以分配的self.free_blocks=list(range(num_blocks))# 每个请求的块表——逻辑块号 → 物理块号self.block_tables={}# request_id → [physical_block_ids]defallocate_slots(self,request_id,num_tokens):""" 为请求分配物理块 比连续显存分配快——块大小固定，只需操作链表 """num_blocks_needed=(num_tokens+block_size-1)//block_size blocks=[]for_inrange(num_blocks_needed):block_id=self.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_table=self.block_tables[request_id]k_out=[]v_out=[]foridxintoken_indices:block_idx=idx//block_size offset=idx%block_size phys_block=block_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(intbi=0;bi<num_blocks;bi++){intphys_block=block_table[bi];// 从块池读取一块AscendC::LocalTensor<float16_t>k_block;AscendC::LocalAlloc(k_block,block_size*num_heads*head_dim);// 块池的物理地址 = base + block_id × block_size × head_dim × num_headsuint64_tk_addr=kv_pool_base+phys_block*block_stride;AscendC::DataCopy(k_block,k_addr,block_size*num_heads*head_dim);// Q @ K_block^T// 这里 q 是当前 Token 的 Query，K_block 是物理块里的完整数据// vLLM 的块大小 16——正好装满 Ascend C 的 L1 Bufferfor(inth=0;h<num_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_op=cann_ops.paged_attention# 用 CANN 的 GE 做图执行——不依赖 CUDA Graphself.use_cann_graph=Truedefexecute_model(self,model_input):""" vLLM 调度层→CANN 执行层 """# vLLM 构造的 Batch 输入input_ids=model_input["input_ids"]# [batch_size]positions=model_input["positions"]# [batch_size]block_tables=model_input["block_tables"]# [batch, num_blocks]# 调用 AscendCL 执行模型# 这里不走 CUDA Graph，而是用 CANN 的 GE 子图执行output=cann_inference(input_ids,positions,block_tables,kv_cache=self.kv_cache_block_pool)returnoutput

CANN 适配后的 vLLM 在 Ascend 910 上跑 LLaMA-7B 的吞吐约 1800 tok/s（连续 Batch 32），比 CUDA 原版的 2200 tok/s 低约 18%。主要差距在 PagedAttention 的 Online-Softmax 实现上——CANN 的 Vector Unit 比 CUDA 的 Tensor Core 做逐块合并时要多搬一次数据。