异步 RAG 流水线：embedding、检索与生成分阶段并发的实现细节

一、深度引言与场景痛点

先给你看一组数字：一个典型的 RAG 查询，embedding 耗时约 150ms，向量检索耗时约 80ms，rerank 耗时约 200ms，LLM 生成耗时约 2000ms。如果你把这些步骤串行执行，总延迟大约是 2430ms。看起来还行？

但现实中的 RAG 流程很少有"一个 query → 一个结果"这么简单的。更常见的情况是：用户的 query 需要先做意图分类（需要一次轻量 LLM 调用），然后拆解成 2-3 个子 query（又需要一次 LLM 调用），每个子 query 各自走 embedding + 检索流程，最后把所有检索结果汇聚起来做 rerank，再喂给 LLM 生成最终回答。

当这些步骤全部串行执行时，延迟会放大到令人崩溃的程度。以 3 个子 query 为例：3 × 150ms（embedding）+ 3 × 80ms（检索）+ 200ms（rerank）+ 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 " f"embedding={self.embedding_ms:.0f}ms " f"检索={self.retrieval_ms:.0f}ms " f"rerank={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", weight=1.0), SubQuery(text="子查询示例2", weight=0.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(f"Embedding 失败 [{sq.text[:30]}]: {e}") return sq.text, [] # 降级返回空向量 tasks = [embed_one(sq) for sq in queries] results = await asyncio.gather(*tasks, return_exceptions=False) 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( vector=vector, top_k=top_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_exceptions=False) 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(), key=lambda k: merged[k], reverse=True) 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), key=lambda x: x[1], reverse=True, ) return [item[0] for item in ranked[:top_n]] except Exception as e: logger.error(f"Rerank 失败: {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_client=mock, embedder=mock, vector_store=mock, reranker=mock, ) 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 和向量相似度这两个打分体系，到底怎么调和。