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
Orthrus co-batches embedding chunks with generation decode in vLLM, hitting 1.28–4.52× throughput vs GPU splits and up to 55.8% lower iterative-RAG p99 on four A100s.
Iterative RAG alternates embedding and generation: embed a query, retrieve, let a generator write the next hop, embed again. Serving stacks such as vLLM run those request types in isolation. Two processes on one GPU fight; static GPU splits cannot track a shifting mix and reloading to repartition is costly. Embedding is a dense forward pass, compute-bound. Generation decode reads the KV cache and is memory-bound. Homogeneous batches make the light work wait on the heavy work, so neither resource stays full.
On an A100-40GB with Mistral-7B, pure embedding averaged 85.3% compute and 36.5% memory; pure generation averaged 75.5% compute and 90.5% memory. Both can spike to 98%. The averages do not line up, which is the co-scheduling gap.
Orthrus rewrites the runner in vLLM 0.92. It does not change memory management or model specs, and it trains nothing. The main setting assumes a shared backbone, typically a generator plus a LoRA embedding adapter. Each scheduling step packs three batch units: embedding token chunks, generation prefill chunks, and single-token decode steps. Attention and FFN run once. Outputs then split: sampling for generation, incremental pooling for embedding.
Chunked embedding splits a long input across iterations. Mean and weighted-mean pooling keep a running weighted sum and a weight total, then divide, which is algebraically the same as a single pass aside from floating-point reduction order. CLS and last-token pooling keep the chosen hidden state. An arbitrary deterministic head can buffer all hidden states and run at the end, at O(nd) memory.
Intra-Batch Scheduling allocates capacity by remaining tokens in each queue, not by request count. In-flight decode is reserved first so KV entries do not move. The scheduler then alternates between the embedding and generation queues in proportion to leftover token work. Arrival order is kept within a type; types interleave. Embedding requests can be chunked, so a leftover token budget is not wasted on one oversize request that cannot fit atomically.
Hardware is up to four A100 40GB GPUs. The base model is Mistral-7B with e5-mistral-7b-instruct LoRA for embeddings. Baselines: task-specific GPU pools, prefill/decode disaggregation, same-GPU homogeneous batching, and same-GPU heterogeneous batching with FCFS.
On controlled mix ratios, throughput versus GPU-level splits is 1.28–4.52×. At a 9:1 embedding:generation mix on Mistral, Dedicated scores 15.64 and 1.49 req/s; Orthrus scores 58.80 and 7.00. At 5:5 the jump is 12.54+11.97 to 25.32+26.24. Qwen2-7B and LLaMA3.1-8B (LoRA rank up to 64) also beat Dedicated1:3 on combined throughput. The exception is Mistral 1:9, where embedding drops from 5.10 to 4.20 while generation rises from 29.56 to 39.10.
Iter-RetGen on 2WikiMultihopQA uses 500-token prompts, 3000-token decodes, and 1–10 documents per query. Four-GPU aggregate throughput: Orthrus 0.624 req/s, Unified-Homo 0.548, Disagg-PD 0.603, Dedicated 0.260. Four-GPU end-to-end p99: 593.4 s versus Dedicated 1342.0 s (−55.8%); one GPU versus Unified-Homo is 573.9 s versus 673.4 s (−14.8%). In a worst-case FCFS stall (1,000 generations then 1,000 embeddings at t=0), generation p99 is 9% below Dedicated (89 s to 81 s) and embedding p99 is 16% below the other schedulers (87 s to 73 s).
Across three workload phases (10/50/90% embedding), mean GPU utilization is 79% versus 38% for Dedicated1:3, and each phase finishes 43% faster. At saturation (≥64 clients) single-GPU generation matches dedicated vLLM; embedding matches or exceeds it. Chunked versus single-pass embeddings stay above 0.9999 minimum cosine similarity across mean/CLS/weighted-mean and chunk sizes 256/512/1024. Separate models (OPT-1.3B generation + GTR-T5-XL embeddings on one A40) beat two vLLM instances by +22.9% embedding and +20.2% generation throughput, with generation p95 3.8% worse.
This is a serving change for iterative RAG. It does not touch the vector backend and does not need a new model. Teams that already derive embeddings from a generator via LoRA can drop both request types into one loop. Static GPU splits and prefill-decode disaggregation leave bubbles when the mix moves; chunking turns embedding into a scheduling unit of the same grain as decode.
Gains are throughput and p99, with vector search and network latency excluded by design. Embedding tails still suffer when generation dominates (embedding p95 reaches 1.71 s at 90% generation). At Mistral 1:9, embedding throughput even falls below the split. That is a scheduler trading one queue against the other inside a fixed token budget, not a free lunch.
Evaluation stops at four GPUs with one replica each, no tensor or pipeline parallelism. Skew and tail latency at larger replica counts are unknown. Models are open-weight and under 10B; larger models may shift the compute-memory mix and the batch-capacity constraint enough to erase the heterogeneous-batching edge. The main path needs a shared backbone. The separate-model trial fits two networks on one A40 and cannot fuse them in one forward pass. The prototype is tied to vLLM 0.92.