DReX: Accurate and Scalable Dense Retrieval Acceleration via Algorithmic-Hardware Codesign
Derrick Quinn, E. Ezgi Yücel, Martin Prammer, Zhenxing Fan, Kevin Skadron, Jignesh M. Patel, José F. Martínez, Mohammad Alian · 2025
Retrieval-augmented generation (RAG) supplements large language models (LLM) with information retrieval to ensure up-to-date, accurate, factually grounded, and contextually relevant outputs.RAG implementations often employ dense retrieval methods and approximate k-nearest neighbor search (ANNS).Unfortunately, ANNS is inherently dataset-specific and prone to low recall, potentially leading to inaccuracies when irrelevant or incomplete context is passed to the LLM.Furthermore, sending numerous imprecise documents to the LLM for generation can significantly degrade performance compared to processing a smaller set of accurate documents.We propose DReX, a dataset-agnostic, accurate, and scalable Dense Retrieval Acceleration scheme enabled through a novel algorithmic-hardware co-design.We leverage in-DRAM logic to enable early filtering of embedding vectors far from the query vector.An outside-DRAM near-memory accelerator then performs exact nearest neighbor searches on the remaining filtered embeddings.This resulting design minimizes off-chip data movement and ensures precise and efficient retrieval, laying the foundation for robust and performant RAG systems that are broadly applicable.Our evaluation shows that DReX delivers a 6.2-7× reduction in time-to-first-token for a representative RAG application over a state-of-the-art mechanism while incurring reasonable area and power overheads in the memory subsystem.