## 共享前缀+批量后缀

**"共享前缀批量推理"**（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]

<passage id="1">[DOC_1]<|doc_emb|></passage>
<passage id="2">[DOC_2]<|doc_emb|></passage>
...
<passage id="k">[DOC_k]<|doc_emb|></passage>

<query>[QUERY]<|query_emb|></query>
<|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，但效率类似）