Source-linked AI summary
RAGO: Systematic Performance Optimization for Retrieval-Augmented Generation Serving
Wenqi Jiang, Suvinay Subramanian, Cat Graves, Gustavo Alonso, Amir Yazdanbakhsh, Vidushi Dadu
TL;DR
LLM-only systems face data staleness, hallucination, and rigid model knowledge, while RAG’s diverse workloads make efficient serving difficult. The paper introduces RAGSchema to represent workloads and RAGO to optimize schedules; it reports up to 2ˆ improvement in QPS per chip and a 55% reduction in time-to-first-token latency.
Problem
RAG serving is difficult to optimize because RAG components and configurations vary substantially, creating a broad scheduling space and shifting performance bottlenecks.
Source: §1 · p. 1Method
RAGSchema represents RAG pipeline stages and model and retrieval configurations, while RAGO uses an analytical cost model to optimize schedules for workload and resource constraints.
Source: §1 · p. 2Results
Up to 2ˆ improvement in QPS per chip and 55% reduction in time-to-first-token latency compared with RAG serving systems built on LLM-only systems.
Source: §1 · p. 2Takeaways & Limitations
Workload characterization can guide future RAG system and hardware design, including retrieval optimization at hyperscale and support for collocated models.
Source: §9 · p. 13Scaling to extremely complicated RAG and agentic AI systems, broader optimization search spaces, and metrics such as energy and cost efficiency remain future work.
Source: §9 · p. 13
Abstract
from arXiv · showhide
Retrieval-augmented generation (RAG), which combines large language models (LLMs) with retrievals from external knowledge databases, is emerging as a popular approach for reliable LLM serving. However, efficient RAG serving remains an open challenge due to the rapid emergence of many RAG variants and the substantial differences in workload characteristics across them. In this paper, we make three fundamental contributions to advancing RAG serving. First, we introduce RAGSchema, a structured abstraction that captures the wide range of RAG algorithms, serving as a foundation for performance optimization. Second, we analyze several representative RAG workloads with distinct RAGSchema, revealing significant performance variability across these workloads. Third, to address this variability and meet diverse performance requirements, we propose RAGO (Retrieval-Augmented Generation Optimizer), a system optimization framework for efficient RAG serving. Our evaluation shows that RAGO achieves up to a 2x increase in QPS per chip and a 55% reduction in time-to-first-token latency compared to RAG systems built on LLM-system extensions.
1 Introduction
RAG can address the staleness, hallucination, and rigid knowledge of LLM-only systems, but its diverse pipelines and configurations make serving performance difficult to optimize. The paper introduces RAGSchema to represent workloads, characterizes four RAG paradigms, and proposes RAGO to optimize serving schedules.
1 Introduction
- RAG workloads combine external retrieval with LLM generation, addressing LLM-only systems’ data staleness, hallucination, and rigid model knowledge.The introduction frames these issues as reliability challenges, especially for applications demanding high factual accuracy.Source: §1 · p. 1
- RAGSchema represents pipeline stages and model and retrieval configurations, simplifying workload representation while retaining information needed for performance characterization and optimization.Its attributes include pipeline components, model size, database size, query vectors per retrieval, and iterative retrieval frequency.Source: §1 · p. 2
- Across four representative RAG paradigms, bottlenecks shift between retrieval and inference, smaller pipeline models can bottleneck, and iterative retrieval can stall decoding.Hyperscale retrieval can account for over 80% of latency, while retrieval is under 1% in long-context scenarios; a database encoder 100ˆ smaller than the generative LLM can bottleneck long-context processing.Source: §1 · p. 2
- RAGO uses an analytical cost model to explore scheduling choices—including placement, resource allocation, and batching—and identify performance Pareto-frontier schedules for user objectives.The framework takes a RAGSchema workload and system resource constraints as inputs.Source: §1 · p. 2
- Up to 2ˆ improvement in QPS per chip and 55% reduction in time-to-first-token latency versus RAG serving systems built on LLM-only systems.These are the reported headline results for RAGO.Source: §1 · p. 2
2 Background
RAG augments LLM prompts with retrieved external knowledge, while its serving pipeline may include distinct generation, retrieval, and optional processing stages. This section outlines the serving characteristics and vector-search approaches relevant to those stages.
Background
- RAG retrieves external knowledge and appends it to prompts; external databases can be updated independently, and grounding responses in retrievable data helps mitigate hallucination.The passage also reports comparable or better generation quality with models one to two orders of magnitude smaller than LLMs.Source: §2 · pp. 2–3
- LLM-only serving separates prompt computation, which produces the first token and populates the KV cache, from autoregressive decoding, which generates later tokens using that cache.Modern systems may disaggregate these stages because prefix is compute-intensive and decode is memory-bound, with TTFT and TPOT as their respective performance metrics.Source: §2 · p. 3
- Vector search ranks database vectors by similarity to a query vector; approximate nearest-neighbor search trades recall for higher system performance on large datasets.L2 distance and cosine similarity are among the stated similarity metrics.Source: §2 · p. 3
- IVF-PQ combines an inverted-file index with product quantization and is favored for large-scale RAG vector search for its memory efficiency.Faiss uses high-precision quantization and is CPU-bound, whereas ScaNN uses lower precision, achieves higher CPU throughput, and shifts toward memory-bound search.Source: §2 · p. 3
3 Structuring the Complex Terrain of RAG Serving
The section organizes diverse RAG workloads through RAGSchema and uses workload-aware models to reason about serving trade-offs. It distinguishes representative RAG paradigms and relates pipeline performance to component workloads and retrieval settings.
3.1 Representative RAG Paradigms
- Hyperscale retrieval uses large knowledge corpora with smaller LLMs, while long-context RAG retrieves relevant passages instead of placing an entire lengthy document in the prompt.The long-context paradigm also uses a database encoder and a much smaller database; a 100K-token context with 100-token chunks yields 1K vectors, versus tens to hundreds of billions in hyperscale retrieval.Source: §3.1 · pp. 3–4
- Iterative retrieval periodically refreshes retrieved content during generation, which can enhance quality and is especially useful for multi-hop reasoning.Generation pauses while newly retrieved context is processed through the prefix phase, then resumes.Source: §3.1 · p. 4
- Query rewriters can clarify or decompose complex queries, while rerankers score retrieved documents for relevance beyond vector similarity.Source: §3.1 · p. 4
3.2 RAGSchema for Workload Abstraction
- RAG workloads vary in database scale, query-vector count, iterative retrieval, and auxiliary models, motivating a structured abstraction for workload comparison.Source: §3.2 · p. 4
- RAGSchema captures pipeline flow and component configurations, including optional stages and performance-relevant attributes, but does not abstract RAG quality.Its attributes include model parameters, vector dimensionality, database size, queries per vector, and iterative retrieval frequency.Source: §3.2 · p. 4
3.3 Empirical RAG Performance Trade-off Analysis
- The performance model estimates short-sequence inference work from FLOPs and retrieval work from database-vector bytes scanned, rather than using FLOPs for retrieval.Inference FLOPs are approximately 2*M*L for short sequences; retrieval bytes depend on vector count, bytes per vector, and the scanned fraction.Source: §3.3 · p. 4
- Retrieval scan percentage is chosen as the minimum that meets required recall; end-to-end latency sums stage latencies, while pipeline throughput is set by its slowest stage.Source: §3.3 · p. 5
4 Methodology
The methodology evaluates RAG serving in a modeled data-center environment using workload configurations, calibrated inference and retrieval simulators, and established latency and throughput metrics.
4 Methodology
- The evaluated workload uses four Llama-3 sizes and a 64-billion-passage vector database, with default retrieval scanning 0.1% of vectors.The typical prompt is approximated at 512 tokens and generation at 256 tokens, based on the selected question, retrieved-passage, and output lengths.Source: §4 · p. 5
- The modeled data-center setup allocates 64–128 XPUs across 16–32 servers, with CPU resources for distributed retrieval; reported default performance uses XPU-C.Each XPU has 96 GB of high-bandwidth memory, and at least 16 servers are used to provide host memory for the quantized dataset.Source: §4 · p. 5, Table 2 · p. 5
- Inference simulation sums operator execution and communication costs, modeling each operator with a roofline estimate based on compute and memory-access latency.Source: §4 · p. 5
- Retrieval simulation models ScaNN search as vector-scan operations whose latency depends on batch size, CPU cores, throughput, and memory bandwidth.The model is calibrated with ScaNN benchmarks and production datasets; distributed search assumes balanced shards and negligible broadcast and gather overhead.Source: §4 · p. 6
- The evaluation reports TTFT, worst-case TPOT, QPS, and QPS per chip, with TPOT reflecting the slowest-generation sequences in continuous batches.Source: §4 · p. 6
5 RAG Serving Performance Characterization
Four representative RAG paradigms show substantial performance variability, with bottlenecks shifting across retrieval, encoding, and inference. The findings show that workload configuration and batch sizing shape the trade-offs between throughput and latency.
5.1 Case I: Hyperscale Retrieval
- RAG 8B achieves 1.5x the QPS/Chip of LLM-only 70B, despite retrieval overhead and longer prompts.The model size is approximately 10x smaller, but the QPS/Chip gain is not proportional to that reduction.Source: Figure 5 · p. 7, §5.1 · p. 7
- Hyperscale retrieval can become dominant with smaller LLMs and more query vectors, while its time share also varies with accelerator capability, sequence lengths, and scanned-vector fraction.Increasing the percentage of database vectors scanned amplifies retrieval’s share of time across model sizes.Source: §5.1 · pp. 7–8
5.2 Case II: Long-Context Sequence Processing
- At 1M tokens, RAG with a 70B LLM achieves 2852.6x TTFT and 6633.9x QPS/Chip speedups versus a long-context LLM; database encoding, not retrieval, is the bottleneck.Caching 10K 768-d database vectors in FP16 for 1M tokens requires 15 MB of CPU memory or storage.Source: §5.2 · p. 8
5.3 Case III: Iterative Retrievals + Prefix
- At decode batch size 64, iterative retrieval-prefix batch size four gives the lowest TPOT; larger decode batches can favor larger retrieval batches.With smaller decode batches, larger iterative retrieval batches can increase latency because the system waits to accumulate enough requests.Source: §5.3 · pp. 8–9
5.4 Case IV: Query Rewriter and reranker
- Adding a query rewriter increases TTFT latency by 2.4x, while reranking has minimal TTFT impact and the modules leave QPS/Chip largely unaffected.The rewriter’s autoregressive generation accounts for its latency impact.Source: §5.4 · p. 9
6 RAGO: Systematic RAG Serving Optimization
RAGO tailors RAG serving schedules to a workload schema and performance target. It combines placement, resource allocation, and batching choices, then searches candidate schedules for a Pareto frontier.
6 RAGO: Systematic RAG Serving Optimization
- RAGO selects schedules tailored to a specific RAGSchema and performance target because one-size-fits-all systems are inadequate across heterogeneous RAG workloads.Its scheduling decisions cover task placement, resource allocation, and batching policies.Source: §6 · p. 10, §6.1 · p. 10
6.1 RAGO Scheduling Decisions
- RAGO supports hybrid placement: neighboring phases up to prefix may be collocated, while retrieval and the main LLM’s prefix and decode remain disaggregated.Collocation can let compute-intensive components share accelerators and mitigate workload imbalance; the search restricts collocation to consecutive neighbors.Source: §6.1 · p. 10
- RAGO allocates CPU and accelerator resources and lets pipeline stages use different batch sizes; execution order for collocated stages is chosen to reduce final-stage completion time.For iterative retrieval, initial retrieval/prefix batches can differ from later iterations, which primarily affect TPOT; collocated stages are time-multiplexed in the described design.Source: §6.1 · pp. 10–11, Figure 14 · p. 11
6.2 Searching for Optimal Scheduling Policies
- RAGO profiles components, generates schedules across placement, allocation, and batching choices, then evaluates them to identify a Pareto frontier.The search can span millions of schedules, with user-set granularity such as restricting accelerator counts or batch sizes to powers of two.Source: §6.2 · p. 11
7 Evaluation
RAGO evaluates scheduling choices across diverse RAG workloads and finds that performance depends on balancing placement, resource allocation, and batching with each workload’s needs. Across the evaluated cases, RAGO outperforms the LLM-only extension baseline and exposes trade-offs between throughput and latency.
7 Evaluation
- Evaluation profiles each RAG stage by resource, generates placement, allocation, and batching options, then evaluates schedules to construct an end-to-end performance Pareto frontier.Placement and allocation are evaluated in C-II and C-IV; batching is evaluated in C-I, C-II, and C-IV.Source: §7 · p. 11
7.1 Overall Performance
- In C-II, RAGO improves maximum QPS/Chip by 1.7× over the baseline; C-IV results also show throughput losses from collocating smaller RAG stages with the LLM prefix system.The C-II baseline can leave 50% of XPUs idle while encoding is overloaded; RAGO allocates 64 of 96 XPUs to encoding.Source: §7.1 · pp. 11–12, Figure 15 · p. 12, Table 4 · p. 12
- The Pareto frontier combines plans with different latency-throughput trade-offs: the throughput-optimized plan reaches 1.5× QPS/Chip with TTFT about 40% higher than the latency-optimized plan.The throughput-oriented plan assigns one chip to the query rewriter, while the latency-oriented plan assigns it 32 chips.Source: Figure 16 · p. 12, §7.1 · p. 12
7.2 Scheduling Policy Sensitivity Analysis
- Placement sensitivity differs by case: C-II plans differ by only 2% in maximum QPS/Chip, while hybrid placement in C-IV improves QPS/Chip by up to 1.5× over collocation.In C-IV, separating the rewriter from the prefix system mitigates prefix-chip underutilization.Source: Figure 17 · p. 12, §7.2 · p. 12
- Resource allocation strongly affects performance: maximum QPS/Chip varies by up to 52.5× for collocated plans and 64.1× for disaggregated plans.The text attributes the variation to insufficient resources at high-workload stages and imbalanced capacity across stages.Source: §7.2 · pp. 12–13, Figure 18 · p. 13
- Micro-batching reduces TTFT by 22% for C-II at batch size two and by 55% at batch size 32 for 1M tokens; reductions reach 46% for C-I and about 25% for C-IV at batch size 32.C-I benefits only at larger batch sizes, while C-IV’s rewriter decoding stage limits the reduction at smaller sequence batches.Source: §7.2 · p. 13
8 Additional Related Work
Prior RAG optimization work targets particular pipeline configurations, while broader LLM and retrieval optimizations address components rather than the full heterogeneous RAG pipeline. The paper argues that RAG needs system-level scheduling, and notes that adopted techniques can shift workload balance.
Additional Related Work
- RAG-specific work targets distinct bottlenecks: Chameleon [68] accelerates product-quantization vector search, PipeRAG [98] and RaLMSpec [100] address iterative-retrieval decoding stalls, and CacheBlend [101] and RAGCache [102] cache document KV states.The caching approach is especially effective when the prompt-to-generation cost ratio is high.Source: §8 · p. 13
- Adopting these techniques can shift RAG workload balance: retrieval acceleration makes systems more inference-bound, document KV precomputation raises retrieval and decoding importance, and prefetching reduces decoding-engine idleness.These are expected shifts described for techniques evaluated by RAGO.Source: §8 · p. 13
- Although LLM and retrieval optimization are extensive, RAG pipeline complexity and heterogeneity motivate system-level RAG scheduling.The cited retrieval optimizations span algorithmic, system-level, and hardware support for product-quantization and graph-based approaches.Source: §8 · p. 13
9 Conclusion and Future Outlook
The paper presents RAGSchema as a foundation for characterizing RAG workloads and RAGO as a framework for optimizing their serving. It reports up to a 2× QPS-per-chip improvement and a 55% TTFT reduction, while identifying more complex systems and broader efficiency objectives as future challenges.
9 Conclusion and Future Outlook
- RAGO delivers up to a 2× improvement in QPS per chip and a 55% reduction in TTFT compared with a strong baseline.The paper frames RAGSchema as a structured abstraction that supports systematic workload characterization and bridges algorithms with system design.Source: §9 · p. 13
- The characterization identifies retrieval as a possible bottleneck in some RAG paradigms, particularly at hyperscale, and points to efficient accelerator support for collocated models as increasingly critical.The authors say these findings can guide future RAG-system and hardware design.Source: §9 · p. 13
- Extremely complex RAG and agentic AI systems remain future work, with broader optimization search spaces and additional metrics such as energy and cost efficiency.Source: §9 · p. 13