RAG（Retrieval-Augmented Generation）是 LLM + 知识库的组合：先检索相关文档，再让 LLM 基于文档回答。昇腾NPU 上部署 RAG 需要两个组件：Embedding 模型（做向量检索）和 LLM（做生成）。CANN 的 ATB 都能跑。
RAG 架构
用户提问 ↓ Embedding 模型 → 问题向量 → 向量数据库检索 → Top-K 文档 ↓ LLM（问题 + 检索到的文档 → 生成回答）两个模型都在昇腾NPU 上跑，延迟最低。
Embedding 模型部署
fromatbimportLLM# Embedding 模型（BGE-base-zh，109M 参数）embed_model=LLM("BAAI/bge-base-zh-v1.5",device="npu:0",task="embedding",# 指定是 Embedding 任务)# 生成向量vectors=embed_model.encode(["什么是昇腾NPU","CANN是什么"])# vectors.shape = (2, 768)Embedding 模型很小（100M 级别），单卡跑绰绰有余。
向量数据库
用 FAISS 做内存向量检索（最快），或用 Milvus 做持久化向量数据库：
importfaissimportnumpyasnp# 构建 FAISS 索引dimension=768# BGE-base-zh 的向量维度index=faiss.IndexFlatIP(dimension)# 内积相似度# 文档库docs=["昇腾NPU是华为研发的AI处理器，基于达芬奇架构","CANN是昇腾计算架构，包含算子库、图引擎、编译器等组件","ATB是昇腾NPU上的大模型推理加速库",]# 文档向量化doc_vectors=embed_model.encode(docs)index.add(np.array(doc_vectors))# 检索query_vector=embed_model.encode(["什么是CANN"])scores,indices=index.search(np.array(query_vector),k=2)# Top-2 文档retrieved_docs=[docs[i]foriinindices[0]]LLM 部署
fromatbimportLLM,SamplingParams llm=LLM("meta-llama/Llama-2-7b-hf",device="npu:0",prefix_caching=PrefixCachingConfig(enable=True),# 系统提示缓存)defrag_generate(query,retrieved_docs):# 拼接 promptsystem_prompt="请根据以下参考资料回答问题。如果资料中没有相关信息，请回答'我不确定'。\n\n参考资料：\n"context="\n".join(f"{i+1}.{doc}"fori,docinenumerate(retrieved_docs))prompt=f"{system_prompt}{context}\n\n问题：{query}\n回答："params=SamplingParams(max_tokens=512,temperature=0.3)result=llm.generate(prompt,params)returnresult多卡部署：Embedding 和 LLM 分卡
两张 NPU 各跑一个模型，避免显存竞争：
# NPU:0 跑 Embeddingembed_model=LLM("BAAI/bge-base-zh-v1.5",device="npu:0",task="embedding")# NPU:1 跑 LLMllm=LLM("meta-llama/Llama-2-7b-hf",device="npu:1")检索和生成可以并行：LLM 生成回答时，Embedding 模型同时处理下一个问题的向量。
Chunk 策略
文档切 Chunk 的粒度影响检索质量：
defchunk_text(text,chunk_size=512,overlap=50):"""滑动窗口切分，overlap 保证上下文连续"""chunks=[]foriinrange(0,len(text),chunk_size-overlap):chunk=text[i:i+chunk_size]chunks.append(chunk)returnchunks# 每 512 token 一段，重叠 50 tokenchunks=chunk_text(long_document,chunk_size=512,overlap=50)Chunk 太大（>1024 token）→ 检索不精确，噪声多。
Chunk 太小（<128 token）→ 语义不完整，回答缺上下文。
推荐 256-512 token，overlap 10%。
Reranker
先检索 Top-20，再用 Reranker 精排到 Top-5：
fromatbimportLLM# Cross-Encoder Rerankerreranker=LLM("BAAI/bge-reranker-base",device="npu:0",task="rerank",)# 粗检索 Top-20scores,indices=index.search(query_vector,k=20)retrieved=[docs[i]foriinindices[0]]# 精排 Top-5ranked=reranker.rerank(query,retrieved,top_k=5)final_docs=[r.textforrinranked]Reranker 的 Cross-Encoder 比双塔 Embedding 更准，但慢 10×。所以先粗后精。
端到端延迟
Atlas 800I A2，BGE-base-zh + Llama2-7B：
| 阶段 | 延迟 |
|---|---|
| Embedding（1 条查询） | 5ms |
| FAISS 检索（100K 文档） | 2ms |
| Rerank（20 条） | 50ms |
| LLM 生成（500 token） | 7.5s |
| 端到端 | 7.6s |
Embedding + 检索只占 1%，瓶颈在 LLM 生成。Speculative Decoding 或量化可以加速 LLM。
RAG 在昇腾NPU 上的部署关键是两个模型分卡跑、Chunk 粒度 256-512 token、Reranker 先粗后精。ATB 同时支持 Embedding 和 LLM 推理，一个框架搞定。仓库在这里：
https://atomgit.com/cann/ATB
https://atomgit.com/cann/torch_npu
