异步 RAG 流水线embedding、检索与生成分阶段并发的实现细节一、深度引言与场景痛点先给你看一组数字一个典型的 RAG 查询embedding 耗时约 150ms向量检索耗时约 80msrerank 耗时约 200msLLM 生成耗时约 2000ms。如果你把这些步骤串行执行总延迟大约是 2430ms。看起来还行但现实中的 RAG 流程很少有一个 query → 一个结果这么简单的。更常见的情况是用户的 query 需要先做意图分类需要一次轻量 LLM 调用然后拆解成 2-3 个子 query又需要一次 LLM 调用每个子 query 各自走 embedding 检索流程最后把所有检索结果汇聚起来做 rerank再喂给 LLM 生成最终回答。当这些步骤全部串行执行时延迟会放大到令人崩溃的程度。以 3 个子 query 为例3 × 150msembedding 3 × 80ms检索 200msrerank 2000ms生成≈ 2690ms但实际上很多东西是可以并行的——比如 3 个子 query 的 embedding 操作互不依赖完全可以同时发起。把 RAG 流水线从串行改为分阶段并行是在不改动任何模型和服务的前提下能直接对延迟动刀的最有效手段。本文就来深入剖析异步 RAG 流水线的工程实现细节。二、底层机制与原理深度剖析异步 RAG 流水线的核心思想很简单找出每一步之间的依赖关系没依赖的就并行有依赖的就等待上游完成后再启动。flowchart LR A[用户 Query] -- B[意图分类] B -- C[Query 拆解] C -- D1[子 Query 1] C -- D2[子 Query 2] C -- D3[子 Query 3] subgraph Embedding_Phase [Embedding 阶段: 并行] D1 -- E1[Embedding 1] D2 -- E2[Embedding 2] D3 -- E3[Embedding 3] end subgraph Retrieval_Phase [检索阶段: 并行] E1 -- F1[检索 1] E2 -- F2[检索 2] E3 -- F3[检索 3] end F1 -- G[结果聚合 去重] F2 -- G F3 -- G G -- H[Rerank] H -- I[LLM 生成] I -- J[最终回答]流水线分四个阶段每个阶段内部的子任务是并行的Phase 1 - 准备阶段意图分类 Query 拆解。这一步通常是串行的因为拆解依赖分类结果而且这两步通常只需要一次轻量 LLM 调用。优化策略是合并为一次调用——在同一个 Prompt 里完成分类和拆解。Phase 2 - Embedding 阶段多个子 query 的 embedding 操作完全独立用asyncio.gather并行发出。这个阶段的加速效果取决于子 query 数量和 embedding 服务的并发处理能力。Phase 3 - 检索阶段每个子 query 的检索操作也独立。但如果向量库本身不支持高并发比如单机 Milvus并行检索的意义就不大——反而会因为连接争抢而变慢。Phase 4 - 生成阶段聚合 Rerank LLM 生成。聚合操作依赖所有检索结果完成Rerank 依赖聚合LLM 依赖 Rerank。这个阶段天然是串行的优化重心应该在 Rerank 的批处理上一次送多个候选而非逐个调用。三、生产级代码实现import asyncio import time import logging from typing import List, Dict, Any, Optional, Tuple from dataclasses import dataclass, field from enum import Enum logger logging.getLogger(__name__) dataclass class PipelineMetrics: 流水线各阶段耗时记录 prepare_ms: float 0 embedding_ms: float 0 retrieval_ms: float 0 rerank_ms: float 0 generation_ms: float 0 property def total_ms(self) - float: return ( self.prepare_ms self.embedding_ms self.retrieval_ms self.rerank_ms self.generation_ms ) def report(self) - str: return ( f总耗时{self.total_ms:.0f}ms | f准备{self.prepare_ms:.0f}ms fembedding{self.embedding_ms:.0f}ms f检索{self.retrieval_ms:.0f}ms frerank{self.rerank_ms:.0f}ms f生成{self.generation_ms:.0f}ms ) dataclass class SubQuery: 子查询 text: str weight: float 1.0 # 权重用于结果融合 class AsyncRAGPipeline: 异步 RAG 流水线 将 RAG 流程拆分为 1. 准备意图分类 子 query 拆解 2. Embedding多 query 并行 3. 检索多 query 并行 4. 聚合 Rerank 5. LLM 生成 def __init__( self, llm_client: Any, # LLM 客户端需支持异步 embedder: Any, # Embedding 服务客户端 vector_store: Any, # 向量库客户端 reranker: Any, # Rerank 服务客户端 max_parallel: int 10, # 最大并行子 query 数 ): self.llm llm_client self.embedder embedder self.vector_store vector_store self.reranker reranker self.max_parallel max_parallel # 并发信号量限制对 embedding 和向量库的并发请求 self._embed_sem asyncio.Semaphore(max_parallel) self._retrieval_sem asyncio.Semaphore(max_parallel) async def _prepare(self, query: str) - List[SubQuery]: Phase 1: 意图分类 Query 拆解 合并为一次 LLM 调用以减少延迟 prompt f分析以下查询的意图并将其拆解为更具体的子查询最多3个 查询{query} 请以 JSON 格式输出 {{intent: 分类结果, sub_queries: [子查询1, 子查询2, ...]}} response await self.llm.generate(prompt) # 解析 LLM 返回的 JSON此处省略解析逻辑 sub_queries [ SubQuery(text子查询示例1, weight1.0), SubQuery(text子查询示例2, weight0.8), ] logger.debug(f拆解为 {len(sub_queries)} 个子查询) return sub_queries async def _embed_batch(self, queries: List[SubQuery]) - Dict[str, List[float]]: Phase 2: 并行 Embedding 使用信号量控制并发避免打爆 embedding 服务 async def embed_one(sq: SubQuery): async with self._embed_sem: try: # 实际调用 embedding API vector await self.embedder.embed(sq.text) return sq.text, vector except Exception as e: logger.error(fEmbedding 失败 [{sq.text[:30]}]: {e}) return sq.text, [] # 降级返回空向量 tasks [embed_one(sq) for sq in queries] results await asyncio.gather(*tasks, return_exceptionsFalse) return {text: vec for text, vec in results if vec} async def _retrieve_batch( self, embeddings: Dict[str, List[float]], top_k: int 10 ) - Dict[str, List[Dict]]: Phase 3: 并行检索 async def retrieve_one(query_text: str, vector: List[float]): async with self._retrieval_sem: try: docs await self.vector_store.search( vectorvector, top_ktop_k, ) return query_text, docs except Exception as e: logger.error(f检索失败 [{query_text[:30]}]: {e}) return query_text, [] tasks [ retrieve_one(text, vec) for text, vec in embeddings.items() ] results await asyncio.gather(*tasks, return_exceptionsFalse) return {text: docs for text, docs in results} def _aggregate_results( self, retrieval_results: Dict[str, List[Dict]], queries: List[SubQuery] ) - List[Dict]: 结果聚合去重 加权融合 # 构建 query_text → weight 映射 weights {sq.text: sq.weight for sq in queries} # 按文档 ID 聚合分数 merged: Dict[str, float] {} doc_map: Dict[str, Dict] {} for query_text, docs in retrieval_results.items(): w weights.get(query_text, 1.0) for doc in docs: doc_id doc.get(id, doc.get(content, )[:50]) raw_score doc.get(score, 0.0) merged[doc_id] merged.get(doc_id, 0) raw_score * w if doc_id not in doc_map: doc_map[doc_id] doc # 按聚合分数排序 sorted_ids sorted(merged.keys(), keylambda k: merged[k], reverseTrue) return [doc_map[doc_id] for doc_id in sorted_ids] async def _rerank( self, query: str, candidates: List[Dict], top_n: int 5 ) - List[Dict]: Phase 4 子步骤: Rerank if not candidates or top_n len(candidates): return candidates try: documents [d.get(content, ) for d in candidates] scores await self.reranker.rerank(query, documents) # 按 rerank 分数重排 ranked sorted( zip(candidates, scores), keylambda x: x[1], reverseTrue, ) return [item[0] for item in ranked[:top_n]] except Exception as e: logger.error(fRerank 失败: {e}) return candidates[:top_n] # 降级直接取 top_n async def _generate(self, query: str, reranked_docs: List[Dict]) - str: Phase 5: LLM 生成 context \n\n.join( d.get(content, )[:500] for d in reranked_docs[:3] ) prompt f基于以下上下文回答问题 上下文 {context} 问题{query} 请给出准确、简洁的回答。如果上下文信息不足请明确说明。 response await self.llm.generate(prompt) return response async def query(self, user_query: str) - Tuple[str, PipelineMetrics]: 完整流水线入口 metrics PipelineMetrics() # Phase 1: 准备 t0 time.monotonic() sub_queries await self._prepare(user_query) metrics.prepare_ms (time.monotonic() - t0) * 1000 # Phase 2: 并行 Embedding t0 time.monotonic() embeddings await self._embed_batch(sub_queries) metrics.embedding_ms (time.monotonic() - t0) * 1000 if not embeddings: return 抱歉查询处理失败。, metrics # Phase 3: 并行检索 t0 time.monotonic() retrieval_results await self._retrieve_batch(embeddings) metrics.retrieval_ms (time.monotonic() - t0) * 1000 # Phase 4: 聚合 Rerank t0 time.monotonic() aggregated self._aggregate_results(retrieval_results, sub_queries) reranked await self._rerank(user_query, aggregated) metrics.rerank_ms (time.monotonic() - t0) * 1000 # Phase 5: LLM 生成 t0 time.monotonic() answer await self._generate(user_query, reranked) metrics.generation_ms (time.monotonic() - t0) * 1000 logger.info(metrics.report()) return answer, metrics # 使用示例 async def main(): # 实际使用中替换为真实的客户端 class MockClient: async def generate(self, prompt): await asyncio.sleep(0.1); return 回答 async def embed(self, text): await asyncio.sleep(0.05); return [0.1]*768 async def search(self, vector, top_k): await asyncio.sleep(0.03); return [] async def rerank(self, query, docs): await asyncio.sleep(0.02); return [0.9]*len(docs) mock MockClient() pipeline AsyncRAGPipeline( llm_clientmock, embeddermock, vector_storemock, rerankermock, ) answer, metrics await pipeline.query(Python asyncio 的最佳实践是什么) print(f回答: {answer}) print(metrics.report()) print(f流水线加速比评估: embedding/检索 可并行部分节省了约 {metrics.embedding_ms metrics.retrieval_ms:.0f}ms 中的并行增益) if __name__ __main__: asyncio.run(main())四、边界分析与架构权衡1. 并行度的甜蜜点不是并行越多越快。当子 query 数量超过 5 个时embedding 和检索的并行加速效应会递减——因为下游服务的并发处理能力有上限。更糟的是过多的子 query 意味着检索结果聚合时需要做更复杂的融合聚合计算本身也会成为瓶颈。建议子 query 控制在 3 个以内超过 5 个就说明你的 query 拆解逻辑需要优化。2. 错误传播在异步流水线中被放大串行执行时一个步骤失败流水线就停了。并行执行时10 个子 query 中有 1 个失败其他 9 个继续执行——这本是好事。但聚合逻辑需要正确处理部分成功的情况1 个失败的子 query 意味着 1/10 的检索路是空的融合结果的质量会下降。需要在监控中追踪embedding_fail_rate和retrieval_fail_rate。3. 阶段间的提前启动当前实现是严格按阶段执行的——Phase 2 必须等 Phase 1 全部完成后才开始。但对于流水线的性能极致优化来说可以做得更激进把 embedding 和检索合并为一个流水线阶段用asyncio.Queue做生产者-消费者模式——embedding 完成一个子 query 就立即放入检索队列而不是等所有 embedding 完成后再批量检索。4. 缓存命中后的短路如果某个子 query 在缓存中已经有了完整的检索结果那它就不需要走 embedding 检索的完整链路。流水线应该支持短路——在 Phase 2 之前检查缓存命中的子 query 直接跳过 embedding 和检索把缓存结果加入聚合阶段即可。五、总结异步 RAG 流水线的核心加速手段是三件事按阶段分层识别独立子任务、层内并行用 asyncio.gather 同时发起、信号量控流保护下游服务不被并发击穿。这种优化的威力取决于你的子 query 数量——数量越多并行收益越大。但不要为了并行而强行拆 query一个好的拆解策略比分阶段并行重要得多。衡量你的流水线是否值得异步化的一个简单标准是如果单次 embedding 耗时 × 子 query 数量 500ms那么异步化带来的延迟改善会让你明显感知到。下一篇聊聊混合检索的权重调优——BM25 和向量相似度这两个打分体系到底怎么调和。