HA-RAG: Hotness-Aware RAG Acceleration via Mixed Precision and Data Placement
This work addresses efficiency challenges in RAG for LLM applications, offering incremental improvements in inference speed for systems relying on external knowledge bases.
The paper tackles the high memory consumption and inference latency in Retrieval-Augmented Generation (RAG) systems by proposing HA-RAG, a hotness-aware optimization system that uses mixed precision and data placement, achieving an average speedup of 2.10x and up to 10.49x in Time-To-First-Token with minimal accuracy loss.
Retrieval-Augmented Generation (RAG) improves model output accuracy by leveraging external knowledge bases, serving as an effective solution to address hallucination issues and knowledge-update delays in Large Language Models (LLMs). However, the introduction of external knowledge bases presents RAG with challenges in long-context processing, significantly increasing memory consumption and inference latency. Existing research accelerates inference by precomputing Key and Value (KV) of the knowledge base and loading them on-demand during inference. Based on the access frequency of different KV chunks within the external knowledge base, this paper proposes a hotness-aware RAG (HA-RAG) inference optimization system. First, leveraging the numerical distribution of KV chunks, we introduce a hotness-aware mixed-precision compressing and loading method to reduce disk I/O and memory access overhead. Second, we design a hotness-aware data placement strategy that prioritizes storing frequently accessed KV chunks in high-speed memory to improve data access efficiency. Experimental results demonstrate that, compared with TurboRAG, the proposed HA-RAG achieves an average speedup of 2.10x and maximum speedup of 10.49x in Time-To-First-Token (TTFT) with negligible accuracy loss.