Embedding模型是将文本转化为高维向量的核心组件广泛应用于语义搜索、推荐系统和 RAG 场景。在昇腾NPU上加速向量检索需要从模型推理和向量索引两个层面协同优化。一、昇腾NPU的Embedding推理优势昇腾CANN为Embedding推理提供了专用的算子库。ops-math 库中的矩阵乘法算子对高维向量计算进行了深度优化配合 HCCL 的多核并行能在昇腾NPU上实现远超 CPU 的向量化吞吐。对于 Embedding 模型而言推理瓶颈通常在两部分模型前向传播的计算耗时以及批量请求下的调度开销。ops-nn 库提供了融合算子将多个小操作合并执行减少中间结果的访存次数。二、批量推理与图优化向量检索场景通常有高并发的特点批量推理是提升吞吐的关键手段。关键在于动态 batching——将多个请求打包成同一批次执行importtorchimporttorch_npu# 动态 batching将等待队列中的请求汇聚成批次# 好处是昇腾NPU的计算资源得到充分利用单个请求的平均耗时降低classDynamicBatcher:def__init__(self,max_batch32,max_wait_ms10):self.queue[]self.max_batchmax_batch self.max_waitmax_wait_ms/1000defadd(self,texts):self.queue.extend(texts)iflen(self.queue)self.max_batch:returnself._run_batch()# 等待一定时间后强制执行避免小批次长时间阻塞# 这里用超时机制而非立即返回保证延迟和吞吐的平衡time.sleep(self.max_wait)returnself._run_batch()def_run_batch(self):tensors[self._tokenize(t)fortinself.queue]# 在昇腾NPU上批量处理避免逐条处理的开销batchtorch_npu.npu_pad_and_concatenate(tensors)embeddingsself.model(batch).cpu().numpy()self.queue.clear()returnembeddings三、向量索引加速推理得到向量后需要通过向量索引实现高效检索。昇腾CANN的 ATB 库对接了主流向量索引库支持 HNSW、IVF 等算法在昇腾NPU上执行frompyatb.retrievalimportVectorIndex# 构建 HNSW 索引适用于对延迟敏感的高召回场景# 相比暴力搜索HNSW 在精度损失可接受的前提下大幅降低检索耗时indexVectorIndex(dimensionembedding_dim,index_typeHNSW,metric_typeCOSINE,ef_construction200,# 控制构建时的搜索广度m16# 控制图的结构稀疏度)# 批量插入向量利用昇腾NPU的并行能力加速索引构建forbatch_embeddingsinbatches(embeddings,batch_size1024):index.add(batch_embeddings)四、混合检索策略生产环境中向量检索通常与关键词检索配合使用形成混合检索策略。ops-transformer 库支持将稠密向量和稀疏向量联合检索# 混合检索融合语义相似度和关键词匹配# 稠密向量负责语义理解稀疏向量负责精确词匹配classHybridRetriever:def__init__(self,vector_index,bm25_index):self.vector_indexvector_index self.bm25_indexbm25_indexdefsearch(self,query,top_k10,alpha0.7):# alpha 控制语义和关键词的权重配比dense_scoresself.vector_index.search(query,top_ktop_k*2)sparse_scoresself.bm25_index.search(query,top_ktop_k*2)# RRF 融合比简单加权更稳定避免不同量纲的分数相互干扰fused_scores{}fori,(doc_id,score)inenumerate(dense_scores):fused_scores[doc_id]fused_scores.get(doc_id,0)1/(alpha*(i1))fori,(doc_id,score)inenumerate(sparse_scores):fused_scores[doc_id]fused_scores.get(doc_id,0)1/((1-alpha)*(i1))returnsorted(fused_scores.items(),keylambdax:-x[1])[:top_k]五、端到端性能调优完整的向量检索加速需要端到端优化从前处理、模型推理到后检索全流程协同调优。使用 torch_npu 的 Profiler 可以定位瓶颈withtorch.npu.profile(profile_dir./profile_results)asprof:forbatchindataloader:embeddingsembedding_model(batch)# 模型推理耗时resultsindex.search(embeddings)# 索引检索耗时# 通过 Timeline 分析各阶段耗时比例针对性优化热点print(prof.export_chrome_trace(trace.json))六、总结在昇腾NPU上加速向量检索需要从模型推理和向量索引双管齐下利用 ops-math 的融合算子和动态 batching 提升推理吞吐借助 ATB 库的 HNSW 索引加速检索结合混合检索策略提升召回质量。ops-transformer 库的检索模块提供了开箱即用的端到端方案。项目地址https://github.com/ascend/open-aimodel-res