HedraRAG: Co-Optimizing Generation and Retrieval for Heterogeneous RAG Workflows
Harness 2 extremely different stages (retrieval & generation) by breaking into small uniform unit
Background
Retrieval-Augmented Generation
enables LLMs to access external knowledge without the expensive pretraining.
2 stage workflow
- Retrieval: collect information from external DB / return top-k nearest matches
- Generation: incorporates the retrieved content to the prompt -> produce more accurate & grounded results.
hybrid CPU-GPU design
- CPU : retrival / save vector data in host memory
- GPU : generation / execute compute intensive generation
Primary workflow: one shot
sequential : retrieval -> generation
Problem: Heterogeneity in RAG requests increases
in workload pattern & workflow structures.
2 stage pipeline can’t handle this.
-> dynamic & imbalanced workloads often lead to misaligned execution -> CPU-GPU hybrid system under-utilization
Challenge
Optimization in heterogeneous RAG workflow
1.Across stages
generation & retrieval from different requests can be parallelized -> improve CPU-GPU pipeline utilization
Challenge: Hard to parallelize
generation vs. retrieval have different characteristics
- LLM generation : step-wise & benefit from dynamic batching
- retrieval: one-shot operation & favors large static batching -> mismatches lead to resource imbalance & pipeline stall.
2. Within a request: use locality
semantic similarity across sequential stages: enable reuse & approximate retrieval.
Challenge: should handle high-dimensional embeddings & diverse similarity patterns
3. Across request: skewed key access -> GPU caching opportunity
Challenges
Limited memory, high PCIe latency, shifting access patterns
Root cause
Existing systems handle this as stage unit Optimization in heterogeneous RAG workflow
1.Across stages
Challenge
Hard to parallelize generation & retrieval because their characteristics are so different
Solution: fine-grained sub-stage partitioning & dynamic batching
unify the smallest unit for both of the stages & control
2. Within a request: use locality
Challenge
should handle high-dimensional embeddings & diverse similarity patterns
Solution: semantic-aware reordering with speculative execution
3. Across request: skewed key access -> GPU caching opportunity
Challenges
Limited memory, high PCIe latency, shifting access patterns
Solution: partial GPU index caching with asynchronous updates captures skewed access patterns
Idea: Add granularity
RAGraph : graph-based abstraction/ represents diverse RAG workflows
- enables a unified view of workflow heterogeneity
- facilitates runtime optimization: through a set of graph transformation operations (node splitting, reordering, edge addition, rewiring)
Strength
API compatibility
Evaluation
Experiments across a wide range of workflows demonstrate that
HedraRAG achieves more that 1.5× and up to 5× speedup over existing frameworks,
offering a comprehensive solution for heterogeneous RAG workload serving.
Inputs
breakdown the execution unit
example: chunked prefill, contiguous batching