CANN-昇腾NPU-RAG推理-检索增强生成怎么部署

RAGRetrieval-Augmented Generation是 LLM 知识库的组合先检索相关文档再让 LLM 基于文档回答。昇腾NPU 上部署 RAG 需要两个组件Embedding 模型做向量检索和 LLM做生成。CANN 的 ATB 都能跑。RAG 架构用户提问 ↓ Embedding 模型 → 问题向量 → 向量数据库检索 → Top-K 文档 ↓ LLM问题 检索到的文档 → 生成回答两个模型都在昇腾NPU 上跑延迟最低。Embedding 模型部署fromatbimportLLM# Embedding 模型BGE-base-zh109M 参数embed_modelLLM(BAAI/bge-base-zh-v1.5,devicenpu:0,taskembedding,# 指定是 Embedding 任务)# 生成向量vectorsembed_model.encode([什么是昇腾NPU,CANN是什么])# vectors.shape (2, 768)Embedding 模型很小100M 级别单卡跑绰绰有余。向量数据库用 FAISS 做内存向量检索最快或用 Milvus 做持久化向量数据库importfaissimportnumpyasnp# 构建 FAISS 索引dimension768# BGE-base-zh 的向量维度indexfaiss.IndexFlatIP(dimension)# 内积相似度# 文档库docs[昇腾NPU是华为研发的AI处理器基于达芬奇架构,CANN是昇腾计算架构包含算子库、图引擎、编译器等组件,ATB是昇腾NPU上的大模型推理加速库,]# 文档向量化doc_vectorsembed_model.encode(docs)index.add(np.array(doc_vectors))# 检索query_vectorembed_model.encode([什么是CANN])scores,indicesindex.search(np.array(query_vector),k2)# Top-2 文档retrieved_docs[docs[i]foriinindices[0]]LLM 部署fromatbimportLLM,SamplingParams llmLLM(meta-llama/Llama-2-7b-hf,devicenpu:0,prefix_cachingPrefixCachingConfig(enableTrue),# 系统提示缓存)defrag_generate(query,retrieved_docs):# 拼接 promptsystem_prompt请根据以下参考资料回答问题。如果资料中没有相关信息请回答我不确定。\n\n参考资料\ncontext\n.join(f{i1}.{doc}fori,docinenumerate(retrieved_docs))promptf{system_prompt}{context}\n\n问题{query}\n回答paramsSamplingParams(max_tokens512,temperature0.3)resultllm.generate(prompt,params)returnresult多卡部署Embedding 和 LLM 分卡两张 NPU 各跑一个模型避免显存竞争# NPU:0 跑 Embeddingembed_modelLLM(BAAI/bge-base-zh-v1.5,devicenpu:0,taskembedding)# NPU:1 跑 LLMllmLLM(meta-llama/Llama-2-7b-hf,devicenpu:1)检索和生成可以并行LLM 生成回答时Embedding 模型同时处理下一个问题的向量。Chunk 策略文档切 Chunk 的粒度影响检索质量defchunk_text(text,chunk_size512,overlap50):滑动窗口切分overlap 保证上下文连续chunks[]foriinrange(0,len(text),chunk_size-overlap):chunktext[i:ichunk_size]chunks.append(chunk)returnchunks# 每 512 token 一段重叠 50 tokenchunkschunk_text(long_document,chunk_size512,overlap50)Chunk 太大1024 token→ 检索不精确噪声多。Chunk 太小128 token→ 语义不完整回答缺上下文。推荐 256-512 tokenoverlap 10%。Reranker先检索 Top-20再用 Reranker 精排到 Top-5fromatbimportLLM# Cross-Encoder RerankerrerankerLLM(BAAI/bge-reranker-base,devicenpu:0,taskrerank,)# 粗检索 Top-20scores,indicesindex.search(query_vector,k20)retrieved[docs[i]foriinindices[0]]# 精排 Top-5rankedreranker.rerank(query,retrieved,top_k5)final_docs[r.textforrinranked]Reranker 的 Cross-Encoder 比双塔 Embedding 更准但慢 10×。所以先粗后精。端到端延迟Atlas 800I A2BGE-base-zh Llama2-7B阶段延迟Embedding1 条查询5msFAISS 检索100K 文档2msRerank20 条50msLLM 生成500 token7.5s端到端7.6sEmbedding 检索只占 1%瓶颈在 LLM 生成。Speculative Decoding 或量化可以加速 LLM。RAG 在昇腾NPU 上的部署关键是两个模型分卡跑、Chunk 粒度 256-512 token、Reranker 先粗后精。ATB 同时支持 Embedding 和 LLM 推理一个框架搞定。仓库在这里https://atomgit.com/cann/ATBhttps://atomgit.com/cann/torch_npu