HA-RAG: Hotness-Aware RAG Acceleration via Mixed Precision and Data Placement

TL;DR

This paper tackles KV-chunk loading bottlenecks in Retrieval-Augmented Generation by exploiting skewed access patterns and value distributions. It introduces HA-RAG, a framework with hotness-aware mixed-precision compression and hotness-aware data placement that assigns 8-bit representations (INT8, FP8 variants, GSE-8) to KV chunks based on access frequency and memory-tier, reducing disk I/O and memory traffic. Through MS MARCO and TriviaQA experiments with LLaMA-2-7B, HA-RAG achieves an average Time-To-First-Token speedup of $2.10\times$ and a maximum of $10.49\times$ over TurboRAG with negligible accuracy loss, validated by ROUGE-1 F1 metrics. The work demonstrates significant practical gains for deploying RAG systems at scale, balancing inference speed and generation quality by combining mixed-precision encoding with tiered data placement.

Abstract

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.

Figures (16)

• Figure 1: Workflow of RAG.
• Figure 2: Access frequency of different KV chunks.
• Figure 3: Overall architecture of HA-RAG.
• Figure 4: Value range distribution of Keys and Values in KV chunks.
• Figure 5: Exponent distribution of floating-point values in Key (left) and Value (right) chunks. Top-$k$ denotes the proportion of elements covered by the $k$ most frequent exponents.