Efficient Iterative Retrieval with Heterogeneous Batching
Dohyun Park, Hubertus Franke, Daniel G. Waddington, Swaminathan Sundararaman, Yongjoo Park
EMNLP 2026 (main conference)
cs.AI, cs.CL
2026-09-22
伊利诺伊大学与IBM的服务系统把嵌入推理和文本生成打进同一循环，四张加速卡上相对按卡切分吞吐最高4.52倍，迭代检索增强生成的尾延迟最多低55.8%。
迭代RAG会在嵌入和生成之间来回切：先把查询做成向量去检索，生成模型再根据缺口改写下一跳查询，再嵌入。vLLM这类框架按任务隔离这两类请求。同一张GPU上跑两个进程会抢资源；多卡按任务静态切分又对不齐实时比例，重切还要重新加载模型。嵌入是一次稠密前向，偏算力；生成decode读KV cache，偏显存。齐次batch把轻活卡在重活后面，GPU两边都吃不饱。
论文在A100-40GB上用Mistral-7B量过：纯嵌入平均算力占用85.3%、显存36.5%；纯生成算力75.5%、显存90.5%。两边的峰值都到得了98%，平均值对不上，这就是共调度的空隙。
Orthrus改的是vLLM 0.92的runner，不改显存管理和模型规格，也不额外训练。假设嵌入和生成共享backbone（生成模型加LoRA做嵌入是主设定）。每轮调度把三类东西放进同一batch：嵌入输入的token chunk、生成prefill chunk、生成decode的单token。一次前向走完Attention和FFN，输出路径再分开：生成走采样，嵌入走进增式pooling。
分块嵌入把长序列切成跨iteration的小段。均值/加权均值pooling用流式累加器和权重和，最后做除法，和整段一次pooling在代数上等价，只留下浮点归约顺序的误差。CLS和last-token只保留对应位置的隐状态。任意确定性头可以缓存全部隐状态再在最后算，代价是O(nd)显存。
Intra-Batch Scheduling按队列里剩余token数而不是请求个数分配容量。先给在飞的decode留位置，避免KV搬动；再按嵌入队列与生成队列的剩余token权重交替取请求。同类型保持到达顺序，类型之间按比例交错，避免一种请求把另一种饿死。嵌入请求可以被切到chunk，残差token预算不会因为一条超长嵌入装不进去而整段浪费。
硬件最多四张A100 40GB。主模型Mistral-7B，嵌入用e5-mistral-7b-instruct LoRA。对照包括按任务切GPU、prefill/decode分离、同卡齐次batch、同卡异构但FCFS。
受控比例工作负载上，相对GPU级切分吞吐1.28×–4.52×。Mistral 9:1嵌入:生成时，Dedicated切分嵌入15.64、生成1.49 req/s，Orthrus为58.80和7.00；5:5时从12.54+11.97升到25.32+26.24。Qwen2-7B和LLaMA3.1-8B（LoRA rank到64）上合吞吐同样全面高于Dedicated1:3。例外是Mistral 1:9：嵌入从5.10降到4.20，生成从29.56升到39.10。
Iter-RetGen跑在2WikiMultihopQA上，prompt均长500、decode均长3000、每查询1–10篇文档。四卡合计吞吐：Orthrus 0.624 req/s，Unified-Homo 0.548，Disagg-PD 0.603，Dedicated 0.260。四卡端到端p99：Orthrus 593.4秒，Dedicated 1342.0秒，低55.8%；一卡相对Unified-Homo从673.4秒降到573.9秒（−14.8%）。最坏FCFS阻塞（先1,000条生成再1,000条嵌入）里，生成p99比Dedicated低9%（89秒→81秒），嵌入p99比其他调度低16%（87秒→73秒）。
三阶段工作负载（嵌入占比10%/50%/90%）下，平均GPU利用率79%对Dedicated1:3的38%，每阶段快43%。饱和后（≥64客户端）单卡生成吞吐对齐专用vLLM，嵌入持平或更高。分块与整段嵌入在mean/CLS/weighted-mean、chunk 256/512/1024上最小余弦相似度都超过0.9999。分离模型（OPT-1.3B生成+GTR-T5-XL嵌入，一张A40）相对两个vLLM实例，嵌入吞吐+22.9%、生成+20.2%，生成p95略差3.8%。
这是给迭代RAG的serving层，不碰检索后端，也不要求换模型。嵌入已经从生成模型LoRA出来的团队，可以直接把两类请求丢进同一循环。静态切GPU和prefill-decode分离在混合比例下会留空泡；Orthrus用chunk把嵌入变成和decode同量级的调度单位。
论文把收益写成吞吐和p99，并明确排除了向量检索和网络延迟。生成特别重的时候嵌入尾部仍会变差（90%生成时嵌入p95到1.71秒）。Mistral 1:9那种生成极重的配比，嵌入吞吐甚至低于切分。这是调度器在既定token预算里做的取舍，不是免费午餐。
评估最多四卡、每卡一个副本，没有张量或流水线并行。更大副本数上的负载倾斜和尾延迟没有数据。<10B的开源模型，更大模型的算存比和batch容量可能把异构batch的优势吃掉。主设定依赖共享backbone；分离模型只在一张A40上试了两个能塞进去的网络，且不能把两个模型打进同一次前向。实现绑在vLLM 0.92。