Retrieval-Augmented Generation (RAG) is an increasingly popular technique to enhance LLMs by incorporating external knowledge, mitigating hallucination and enabling access to up-to-date information. A typical RAG pipeline involves a retrieval phase, where relevant documents are fetched from a knowledge source (like a vector database), and a generation phase, where an LLM synthesizes a response based on the input query and retrieved information. This paper profiles the RAG execution pipeline and identifies the retrieval phase, specifically Exact Nearest Neighbor Search (ENNS) over a dense vector database, as a critical performance bottleneck.
The authors demonstrate that while Approximate Nearest Neighbor Search (ANNS) can be faster, it often compromises retrieval quality. To maintain generation accuracy comparable to ENNS, ANNS requires fetching a larger number of documents (higher K). Providing this increased context to the LLM linearly or quadratically increases the computational cost of the generation phase, potentially negating any speedup gained during retrieval and even increasing the end-to-end inference time. Experiments show that ENNS-based RAG often achieves superior accuracy-throughput trade-offs compared to ANNS-based RAG, especially when high generation accuracy is required.
Current hardware platforms face challenges accelerating ENNS. CPUs are memory-bandwidth limited and cannot saturate DRAM bandwidth for ENNS. While GPUs can accelerate ENNS due to their SIMD capabilities, they are expensive, especially for large knowledge bases that exceed the memory capacity of a single GPU, necessitating multi-GPU setups. Furthermore, GPUs provision significant compute resources that are poorly utilized by the memory-bound ENNS workload. Specialized ANNS accelerators exist but are often task-specific due to the complex and varied nature of ANNS algorithms, whereas ENNS is algorithmically simple.
Motivated by the need for high-quality, high-performance, and cost-effective retrieval, the paper proposes the Intelligent Knowledge Store (IKS), a specialized CXL Type-2 memory expander designed to accelerate ENNS. IKS is built with three key requirements: cost-effective and scalable memory capacity, userspace-managed near-memory acceleration, and a shared address space between the host CPU and accelerators.
The IKS architecture features a scale-out near-memory processing design. It comprises multiple Near-Memory Accelerator (NMA) chips, each located near an LPDDR5X memory package. This scale-out approach addresses the physical limitations (like die shoreline) that restrict the number of memory channels accessible by a single large accelerator chip. LPDDR5X memory is chosen for its balance of cost, capacity, and bandwidth, fitting the needs of large vector databases compared to more expensive HBM or less general-purpose DDR.
IKS leverages the CXL.mem and CXL.cache protocols to implement a novel, low-overhead interface with the host CPU. The internal IKS DRAM, NMA scratchpads, and configuration registers are mapped into the host address space, allowing seamless data access and offload management via cache-coherent shared memory. The host CPU manages the vector database and initiates ENNS offloads by writing query vectors and metadata to memory-mapped context buffers and signaling completion via a doorbell register, all synchronized using CXL.cache. This cache-coherent interface is shown to provide higher throughput for offload context communication compared to non-temporal writes typical of PCIe MMIO. The CPU is notified of offload completion via a lightweight mechanism (umwait) and performs the final aggregation of partial top-K results from each NMA.
Each NMA chip contains 64 processing engines, each with a query scratchpad, dot-product unit, and Top-K unit. The architecture is designed to saturate the LPDDR5X memory bandwidth for dot-product calculations by evaluating similarity scores for blocks of embedding vectors read from DRAM in column-major order. Data is reused across processing engines to support batching efficiently. The Top-K unit maintains the best similarity scores found so far, offloading this task from the CPU.
The evaluation uses a cycle-approximate simulator calibrated with RTL synthesis data and real software stack overheads. Experiments on RAG workloads (FiDT5, Llama-8B, Llama-70B) demonstrate that IKS provides 13.4x to 27.9x faster ENNS retrieval for a 512GB vector database compared to an Intel Sapphire Rapids CPU. When integrated into the end-to-end RAG pipeline (with generation on an NVIDIA H100 GPU), IKS reduces inference time by 1.7x to 26.3x depending on the model and configuration. The faster retrieval allows IKS-accelerated RAG to achieve significantly higher throughput than ANNS-based systems while maintaining high generation accuracy.
The authors discuss trade-offs, noting that IKS is particularly effective for applications requiring high recall or datasets not well-suited for existing ANNS algorithms. Potential inefficiencies include full corpus scan and low NMA utilization at small batch sizes, suggesting future work on early termination and power management techniques. Compared to GPUs, IKS is projected to be more cost-effective for large capacities due to the use of LPDDR5X and a smaller silicon area for the NMAs.
In conclusion, the paper identifies the critical role of high-quality retrieval in RAG performance and proposes IKS as a purpose-built, cost-effective CXL device for accelerating ENNS. IKS leverages near-memory processing and a cache-coherent host interface to significantly improve retrieval performance and, consequently, the end-to-end inference time of RAG applications.