diff --git a/docs/reranker-共享前缀批量推理.md b/docs/reranker-共享前缀批量推理.md new file mode 100644 index 0000000..bf51fe9 --- /dev/null +++ b/docs/reranker-共享前缀批量推理.md @@ -0,0 +1,460 @@ + +## 共享前缀+批量后缀 + +**"共享前缀批量推理"**(Shared-Prefix Batch Inference): +- 输入结构:`[Prefix A] + [Suffix B1]`、`[Prefix A] + [Suffix B2]`、`[Prefix A] + [Suffix B3]`... +- 痛点:Prefix A被重复计算成千上万次,浪费算力和时间 +- 理想方案:Prefix A只推理一次,KV Cache复用给所有后缀 + +--- + +## 一、顶级推荐方案 + +### 1. **FlashInfer + Cascade Attention** ⭐ 最强技术 +- **核心创新**:Cascade Inference(级联推理) +- **性能提升**:相比vLLM的PageAttention,在32K token共享前缀、batch size 256场景下可达 **31倍加速** +- **原理**: + - 将Attention计算分解为两个阶段: + 1. **多查询阶段**:对共享前缀使用Multi-Query Kernel,只计算一次,结果存入Shared Memory + 2. **批量解码阶段**:对每个唯一后缀使用Batch Decode Kernel + - 使用结合律算子合并部分Attention状态,类似FlashAttention的分块策略 +- **适用场景**:文档QA、系统提示词复用、RAG批量检索 +- **集成**:已集成到SGLang和vLLM中作为后端 +- **GitHub**: https://github.com/flashinfer-ai/flashinfer + +### 2. **SGLang + RadixAttention** ⭐ 最实用框架 +- **核心创新**:RadixTree(基数树)管理KV Cache +- **自动前缀复用**:无需手动配置,自动识别共享前缀并复用KV Cache +- **性能**:相比vLLM、LMQL等基线系统,结构化工作负载上可达 **6.4倍吞吐量提升** 和 **3.7倍延迟降低** +- **关键特性**: + - **In-Batch Prefix Caching**:同一batch内自动共享前缀(如你的A+B1, A+B2场景) + - **Multi-Item Scoring (MIS)**:LinkedIn用于推荐排序的优化,将多个候选项合并为单次前向传播 + - **Zero-Overhead CPU Scheduler**:GPU计算时CPU并行准备下一batch,利用率接近100% +- **特别适合**:Agent系统、工具链、RAG应用 +- **GitHub**: https://github.com/sgl-project/sglang + +### 3. **vLLM + Automatic Prefix Caching** ⭐ 最成熟稳定 +- **核心机制**:基于哈希表的块级前缀缓存 +- **工作原理**: + - 将KV Cache按块(默认16 tokens)哈希 + - 新请求先查哈希表,命中则直接复用,只计算新tokens + - 使用LRU策略管理缓存 eviction +- **使用方式**: + ```python + from vllm import LLM, SamplingParams + # 启用prefix caching + llm = LLM(model="your-model", enable_prefix_caching=True) + + # 第一次调用缓存前缀 + outputs = llm.generate(long_prefix + prompt_1, sampling_params) + # 第二次调用自动命中缓存,prefix部分零计算 + outputs = llm.generate(long_prefix + prompt_2, sampling_params) + ``` +- **注意事项**:vLLM 0.6.3之前调度器未考虑缓存命中率,高并发长序列场景可能性能下降,建议升级到0.6.5+ + +--- + +## 二、其他重要方案 + +### 4. **TensorRT-LLM + In-Flight Batching** +- **优势**:NVIDIA官方优化,与FlashInfer深度集成 +- **特性**: + - 支持Prefix Caching(具体实现闭源,但概念类似vLLM) + - XQA Kernel(Flash Attention 3变体)优化内存访问 + - 层融合技术减少中间结果存储 +- **性能**:在共享前缀数据集上,吞吐量提升 **~34.7%**,TPOT降低 **~20.9%** +- **适用**:NVIDIA GPU生产环境,追求极致性能 + +### 5. **LMDeploy + TurboMind** +- **定位**:纯C++引擎,消除Python开销 +- **性能**:与SGLang相当,在H100上可达 **~16,200 tok/s**(vLLM为12,553 tok/s) +- **优化**:支持KV Cache量化(8-bit)、Continuous Batching +- **适用**:高吞吐生产部署,对延迟敏感的场景 + +### 6. **Daft + Dynamic Prefix Bucketing**(大数据场景) +- **创新点**:动态前缀分桶 + 流式Continuous Batching +- **解决痛点**: + - 全局排序分组会导致GPU空闲 + - 动态分桶在推理同时进行前缀分组,实现流水线 +- **性能**:128 GPU集群上,20万prompts(1.28亿tokens)处理速度提升 **50.7%** +- **适用**:大规模离线批处理(如数据标注、合成数据生成) + +--- + +## 三、针对你的具体场景建议 + +### 场景1:在线服务(RAG检索、实时重排序) +**推荐栈**:**SGLang** 或 **vLLM + FlashInfer后端** + +```python +# SGLang示例:自动前缀复用 +import sglang as sgl + +@sgl.function +def rerank(s, query, docs): + # query是共享前缀,docs是批量后缀 + s += "Query: " + query + "\n" + s += "Document: " + sgl.arg(docs) + "\nRelevance:" + s += sgl.gen("score", max_tokens=5) + +# 批量执行,自动共享query部分的KV Cache +docs = ["doc1 content", "doc2 content", "doc3 content", ...] # 成千上万个 +state = rerank.run_batch( + [{"query": "user query", "docs": d} for d in docs], + max_new_tokens=5 +) +``` + +### 场景2:离线批量处理(数据标注、索引构建) +**推荐栈**:**Daft** 或 **FlashInfer原生API** + +```python +# Daft示例:动态前缀分桶 +import daft +from daft.functions import prompt + +df = daft.from_pydict({ + "query": ["shared query"] * 10000, + "doc": ["doc1", "doc2", ...] # 不同后缀 +}) + +df = df.with_column("score", + prompt( + df["query"] + "\n" + df["doc"], + provider="vllm-prefix-caching", # 利用前缀缓存 + model="your-model" + ) +) +``` + +### 场景3:Embedding/Reranker模型(Bi-Encoder/Cross-Encoder) +**推荐栈**:**Sentence-Transformers优化** + **ONNX/TensorRT** + +```python +# Cross-Encoder批量重排序优化 +from sentence_transformers import CrossEncoder +import numpy as np + +class OptimizedReranker: + def __init__(self, model_name="cross-encoder/ms-marco-MiniLM-L-6-v2"): + self.model = CrossEncoder(model_name, max_length=512, device="cuda") + + def rerank_batch(self, query, documents, batch_size=32): + # 构建所有pairs:[query, doc1], [query, doc2], ... + pairs = [[query, doc] for doc in documents] + + # 单次批量推理,自动共享query的编码计算 + scores = self.model.predict( + pairs, + batch_size=batch_size, + convert_to_numpy=True + ) + return np.argsort(scores)[::-1] +``` + +--- + +## 四、性能优化关键技巧 + +### 1. **Prompt构造策略** +- **对齐块边界**:确保共享前缀长度是KV Cache块大小(通常是16或32)的整数倍,避免部分块重计算 +- **静态前缀前置**:将不变的指令、系统提示放在最前面 + +### 2. **Batch构造策略** +- **前缀分桶(Prefix Bucketing)**:将相同前缀的请求分到同一batch,最大化缓存命中率 +- **长度排序**:同batch内按后缀长度排序,减少padding浪费 + +### 3. **硬件/精度优化** +- **FP16/BF16**:混合精度推理,吞吐量提升1.5-2倍 +- **KV Cache量化**:8-bit量化减少内存带宽压力,提升20%吞吐 +- **CUDA Graphs**:捕获静态计算图,减少CPU开销(适合固定batch size场景) + +--- + +## 五、方案对比总结 + +| 方案 | 前缀复用方式 | 适用场景 | 性能提升 | 易用性 | 成熟度 | +|------|-------------|---------|---------|--------|--------| +| **FlashInfer** | Cascade Attention | 超长前缀(32K+) | 31x | 中(需集成) | 高 | +| **SGLang** | RadixTree自动缓存 | Agent/RAG/工具链 | 6.4x | 高 | 高 | +| **vLLM** | 哈希块缓存 | 通用在线服务 | 2-5x | 高 | 极高 | +| **TensorRT-LLM** | 内部优化 | NVIDIA生产环境 | 1.3-2x | 中 | 高 | +| **LMDeploy** | TurboMind C++引擎 | 低延迟高吞吐 | 1.3x | 中 | 高 | +| **Daft** | 动态分桶 | 大规模离线批处理 | 50%+ | 高 | 中 | + +--- + +## 六、推荐实施路径 + +1. **快速验证**:先用 **vLLM** 开启`enable_prefix_caching=True`验证效果,零代码改动 +2. **性能突破**:切换到 **SGLang**,利用RadixAttention和自动batch优化,获得显著性能提升 +3. **极致优化**:对超长前缀场景,使用 **FlashInfer** 的Cascade Attention内核,或集成到自定义推理引擎中 + +需要针对具体框架的详细配置代码或部署架构建议吗? + + + +## 天然支持Prefix Caching的Reranker模型清单 + +最新版本 (v2 系列):混合架构 +在 BGE-Reranker-v2 中,BAAI 为了提升模型上限,引入了不同的底座: + +BGE-Reranker-v2-m3: 依然保持了 Encoder-only 架构(基于 XLM-RoBERTa),支持多语言和长文本,具有极高的性价比。 + +BGE-Reranker-v2-Gemma/Llama: 采用了 Decoder-only 架构。 + +这些模型是基于 Gemma-2b 或 Llama-3-8B 等大语言模型(LLM)微调而来的。 + +虽然原生 LLM 是单向(Causal)注意力的,但在作为 Reranker 使用时,模型通常会通过特殊的 Prompt 引导,并取最后一个 Token 的输出经过一个线性层来计算相关性分数。 + + +### 1. **BGE-Reranker-V2/V2.5 系列** ⭐ 强烈推荐 +基于Gemma/MiniCPM等Decoder-only架构,FlagEmbedding官方实现已优化 + +| 模型 | 架构 | 参数量 | 特点 | +|------|------|--------|------| +| `BAAI/bge-reranker-v2-gemma` | Gemma-2B (Decoder-only) | 2B | 多语言强,基础版 | +| `BAAI/bge-reranker-v2-minicpm-layerwise` | MiniCPM-2B (Decoder-only) | 2B | **支持层选择**,可截断到第24层加速 | +| `BAAI/bge-reranker-v2.5-gemma2-lightweight` | Gemma2-9B (Decoder-only) | 9B | **Token压缩+层选择**,极致效率 | + +**Prefix Caching友好特性** : +- 输入格式:`[Query A] + [SEP] + [Document B] + [Prompt]` +- Query部分作为前缀,可被所有Document共享 +- 官方代码中已实现`compute_score_single_gpu`的batch处理,自动对齐长度排序减少padding + +**使用示例**: +```python +from FlagEmbedding import FlagAutoReranker + +# 启用vLLM后端 + Prefix Caching +reranker = FlagAutoReranker.from_finetuned( + 'BAAI/bge-reranker-v2-gemma', + model_class='decoder-only-base', + use_fp16=True, + devices=['cuda:0'] +) + +# 批量推理:Query自动复用KV Cache +pairs = [ + ('what is panda?', 'The giant panda is a bear species...'), # A+B1 + ('what is panda?', 'Pandas are popular zoo animals.'), # A+B2 (Query复用) + ('what is panda?', 'Pandas eat bamboo and live in China.'), # A+B3 (Query复用) +] +scores = reranker.compute_score(pairs, batch_size=32) +``` + +### 2. **Qwen3-Reranker 系列** ⭐ 国产最强 +基于Qwen3 Decoder-only架构,阿里云官方支持 + +| 模型 | 架构 | 参数量 | 特点 | +|------|------|--------|------| +| `Qwen/Qwen3-Reranker-0.6B` | Qwen3-0.6B (Decoder-only) | 0.6B | 超轻量,快速 | +| `Qwen/Qwen3-Reranker-4B` | Qwen3-4B (Decoder-only) | 4B | 性能均衡 | +| `Qwen/Qwen3-Reranker-8B` | Qwen3-8B (Decoder-only) | 8B | 精度最高 | + +**架构细节** : +- **纯Decoder-only架构**:使用因果注意力,天然支持Prefix Caching +- **输入模板**: + ``` + <|im_start|>system + You are a helpful assistant. + <|im_end|> + <|im_start|>user + Query: {query} + Document: {document} + Does the document answer the query? Please answer Yes or No. + <|im_end|> + <|im_start|>assistant + Yes + ``` +- **输出**:只生成"Yes"或"No"的logits,作为相关性分数 + +**vLLM部署优化** : +```bash +# 启动vLLM服务,启用Prefix Caching +python -m vllm.entrypoints.openai.api_server \ + --model Qwen/Qwen3-Reranker-8B \ + --tensor-parallel-size 1 \ + --dtype half \ + --max-model-len 32768 \ + --enable-prefix-caching # 关键参数!提速40% +``` + +### 3. **Jina-Reranker-V3** ⭐ 创新架构 +基于Qwen3-0.6B的Listwise重排序器,支持跨Document注意力 + +| 模型 | 架构 | 参数量 | 特点 | +|------|------|--------|------| +| `jinaai/jina-reranker-v3` | Qwen3-0.6B (Decoder-only) | 0.6B | **Listwise**,单次处理64个docs | + +**独特优势** : +- **"Last but Not Late"交互**:在单个context window中同时处理Query+所有Documents +- **跨Document注意力**:通过因果注意力实现Document间交互,捕捉相对相关性 +- **Prefix Caching优化**:Query放在序列开头,可被所有Document共享 + +**输入格式** : +```xml +<|im_start|>system +You are a search relevance expert... +<|im_end|> +<|im_start|>user +Rank the passages based on their relevance to query: [QUERY] + +[DOC_1]<|doc_emb|> +[DOC_2]<|doc_emb|> +... +[DOC_k]<|doc_emb|> + +[QUERY]<|query_emb|> +<|im_end|> +``` + +**性能**:BEIR nDCG@10达61.94,超过Qwen3-Reranker-4B,体积小6倍 + +### 4. **E5-Mistral / NV-Embed-v2 / SFR-Embedding-Mistral** +基于Mistral-7B Decoder-only架构的Embedding/Reranker + +| 模型 | 架构 | 用途 | 特点 | +|------|------|------|------| +| `intfloat/e5-mistral-7b-instruct` | Mistral-7B (Decoder-only) | Embedding | 指令微调,支持多任务 | +| `nvidia/NV-Embed-v2` | Mistral-7B (Decoder-only) | Embedding | 潜在注意力层优化 | +| `Salesforce/SFR-Embedding-Mistral` | Mistral-7B (Decoder-only) | Embedding | 长上下文优化 | + +**注意**:这些主要是**Embedding模型**(Bi-encoder),但可作为Reranker使用(如通过余弦相似度)。若需Cross-encoder式重排序,需配合其他技术。 + +### 5. **RankGPT / RankZephyr / RankLLaMA** +基于LLM的生成式重排序器 + +| 模型 | 架构 | 特点 | +|------|------|------| +| `RankGPT` (GPT-3.5/4) | Decoder-only API | 通过Prompt让LLM判断相关性 | +| `RankZephyr` | Zephyr-7B (Decoder-only) | 蒸馏自RankGPT,开源可部署 | +| `RankLLaMA` | LLaMA-2/3 (Decoder-only) | 本地部署,隐私友好 | + +**实现方式**: +- 使用LLM的logits或生成"Yes/No"判断相关性 +- 完全基于Decoder-only架构,天然支持Prefix Caching + +--- + +## 不支持Prefix Caching的Reranker(双向架构) + +### ❌ **BGE-Reranker-V1 / BGE-Reranker-Base/Large** +- **架构**:基于XLM-RoBERTa(Encoder-only,双向注意力) +- **问题**:Query和Document拼接后`[CLS]` token的表示依赖于整个序列,无法分离缓存 +- **适用场景**:轻量级、短文本,对延迟不敏感 + +### ❌ **Cross-Encoder (BERT-based)** +- **架构**:BERT/RoBERTa等Encoder-only模型 +- **问题**: + - 每个Query-Document对必须**联合编码** + - 前缀的KV Cache与后续token强耦合,无法复用 + - 计算复杂度O((|Q|+|D|)²),无法分解为O(|Q|²) + O(|D|²) + +### ❌ **ColBERT / ColPali / Late Interaction模型** +- **架构**:基于BERT的双向编码 + 后期MaxSim交互 +- **问题**: + - **独立编码,但双向注意力**:Query和Document分别编码,但各自内部仍是双向 + - **无法Prefix Cache**:虽然Query可独立编码,但Document的编码不依赖于Query,所以不存在"共享前缀"场景 + - **优化点**:Document可**预计算并离线存储**,Query实时编码,然后做MaxSim + - **本质区别**:这是"离线预计算"而非"Prefix Caching",适用于Document固定、Query变化的场景 + +**ColBERT的优化策略** : +```python +# ColBERT流程:Document预计算(离线) + Query实时编码(在线) +document_embeddings = encode_documents(docs) # 离线,一次性 +query_embedding = encode_query(query) # 在线,每次查询 +scores = maxsim(query_embedding, document_embeddings) # 轻量级交互 +``` + +--- + +## 实际部署建议 + +### 场景1:高并发在线服务(Query多变,Document固定) +**推荐**:**Jina-Reranker-V3** 或 **Qwen3-Reranker + vLLM Prefix Caching** + +```python +# vLLM配置优化 +from vllm import LLM, SamplingParams + +llm = LLM( + model="Qwen/Qwen3-Reranker-8B", + enable_prefix_caching=True, # 关键! + max_num_seqs=256, # 高并发 + max_model_len=32768 +) + +# 批量构造Prompts:共享Query前缀 +query = "What is the capital of France?" +docs = ["Paris is the capital...", "France is a country...", "Berlin is the capital of Germany..."] + +prompts = [ + f"Query: {query}\nDocument: {doc}\nRelevant:" + for doc in docs +] + +# vLLM自动识别共享前缀,只计算一次Query的KV Cache +sampling_params = SamplingParams(max_tokens=1, temperature=0) +outputs = llm.generate(prompts, sampling_params) +``` + +### 场景2:离线批量重排序(Query固定,Document多变) +**推荐**:**ColBERT / ColPali**(Document预计算策略) + +```python +from rankify.models.reranking import Reranking +from rankify.dataset.dataset import Document, Question, Context + +# ColBERT:Document预计算,Query实时编码 +model = Reranking(method='colbert_ranker', model_name='Colbert') + +# Documents已预计算并存储 +document = Document( + question=Question("What is RAG?"), + contexts=[Context(text=doc, id=i) for i, doc in enumerate(docs)] +) + +# 只需编码Query,然后MaxSim计算 +model.rank([document]) +``` + +### 场景3:极致性能 + 高精度 +**推荐**:**BGE-Reranker-V2.5-Gemma2-Lightweight**(Token压缩 + 层选择) + +```python +from FlagEmbedding import LightWeightFlagLLMReranker + +reranker = LightWeightFlagLLMReranker( + 'BAAI/bge-reranker-v2.5-gemma2-lightweight', + devices=["cuda:0"], + use_fp16=True +) + +# 综合优化:层截断 + Token压缩 +scores = reranker.compute_score( + pairs, + cutoff_layers=[28], # 只用前28层 + compress_ratio=4, # Token压缩4倍 + compress_layers=[24, 40] # 特定层压缩 +) +``` + +--- + +## 总结对比表 + +| 模型 | 架构 | Prefix Caching | 适用场景 | 性能/效率 | +|------|------|---------------|---------|----------| +| **BGE-Reranker-V2/V2.5** | Decoder-only ✅ | ✅ 原生支持 | 多语言、生产环境 | 高 | +| **Qwen3-Reranker** | Decoder-only ✅ | ✅ vLLM支持 | 中文优先、高精度 | 极高 | +| **Jina-Reranker-V3** | Decoder-only ✅ | ✅ Listwise优化 | 跨Doc交互、Top-K排序 | 极高 | +| **E5-Mistral/NV-Embed** | Decoder-only ✅ | ✅ 需配合框架 | Embedding+轻量重排 | 高 | +| **RankGPT/Zephyr** | Decoder-only ✅ | ✅ API/本地 | 生成式判断 | 中 | +| **BGE-Reranker-V1** | Encoder-only ❌ | ❌ 不支持 | 轻量、短文本 | 低 | +| **ColBERT/ColPali** | 双向+后期交互 | ❌ 不适用(预计算替代) | Document固定场景 | 中 | + +**最终建议**: +- 如果追求**Prefix Caching加速** + **高精度**:选择 **Qwen3-Reranker-8B** 或 **BGE-Reranker-V2.5**,配合 **vLLM + FlashInfer** 部署 +- 如果需要**跨Document比较**(Listwise):选择 **Jina-Reranker-V3** +- 如果**Document固定且量大**:选择 **ColBERT** 预计算策略(虽非Prefix Caching,但效率类似) \ No newline at end of file diff --git a/docs/搜索API对接指南.md b/docs/搜索API对接指南.md index e8511a4..aebbc09 100644 --- a/docs/搜索API对接指南.md +++ b/docs/搜索API对接指南.md @@ -148,6 +148,17 @@ curl -X POST "http://120.76.41.98:6002/search/" \ - **端点**: `POST /search/` - **描述**: 执行文本搜索查询,支持多语言、布尔表达式、过滤器和分面搜索 +- **租户标识**:`tenant_id` 通过 HTTP 请求头 **`X-Tenant-ID`** 传递(推荐);也可通过 URL query 参数 **`tenant_id`** 传递。**不要放在请求体中。** + +**请求示例(推荐)**: +```python +url = f"{base_url.rstrip('/')}/search/" +headers = { + "Content-Type": "application/json", + "X-Tenant-ID": "162", # 租户ID,必填 +} +response = requests.post(url, headers=headers, json={"query": "芭比娃娃"}) +``` ### 3.2 请求参数 @@ -605,6 +616,7 @@ curl "http://localhost:6002/search/instant?q=玩具&size=5" - **端点**: `GET /search/{doc_id}` - **描述**: 根据文档 ID 获取单个商品详情,用于点击结果后的详情页或排查问题。 +- **租户标识**:同 [3.1](#31-接口信息),通过请求头 `X-Tenant-ID` 或 query 参数 `tenant_id` 传递。 #### 路径参数 @@ -630,7 +642,8 @@ curl "http://localhost:6002/search/instant?q=玩具&size=5" #### 请求示例 ```bash -curl "http://localhost:6002/search/12345" +curl "http://localhost:6002/search/12345" -H "X-Tenant-ID: 162" +# 或使用 query 参数:curl "http://localhost:6002/search/12345?tenant_id=162" ``` --- @@ -1287,6 +1300,8 @@ curl -X GET "http://localhost:6004/indexer/health" ## 常见场景示例 +以下示例仅展示**请求体**(body);实际调用时请加上请求头 `X-Tenant-ID: <租户ID>`(或 URL 参数 `tenant_id`),参见 [3.1 接口信息](#31-接口信息)。 + ### 7.1 基础搜索与排序 **按价格从低到高排序**: diff --git a/indexer/document_transformer.py b/indexer/document_transformer.py index ba67f99..271f944 100644 --- a/indexer/document_transformer.py +++ b/indexer/document_transformer.py @@ -348,7 +348,8 @@ class SPUDocumentTransformer: if pd.notna(spu_row.get('image_src')): image_src = str(spu_row['image_src']) if not image_src.startswith('http'): - image_src = f"//{image_src}" if image_src.startswith('//') else image_src + # 仅当尚未是协议相对 URL 时才补 "//",避免 "//host" 变成 "////host" + image_src = f"//{image_src}" if not image_src.startswith('//') else image_src doc['image_url'] = image_src def _process_skus( -- libgit2 0.21.2