Cosmos: A CXL-Based Full In-Memory System for Approximate Nearest Neighbor Search

TL;DR

Cosmos is a novel CXL-based in-memory system designed to perform high-throughput, low-latency approximate nearest neighbor searches for retrieval-augmented generation, overcoming limitations of traditional hardware solutions.

Contribution

It introduces a full ANNS offload architecture with rank-level parallel distance computation and adjacency-aware data placement on CXL memory devices.

Findings

Achieves up to 6.72x higher throughput than baseline CXL systems.

Outperforms state-of-the-art CXL solutions by 2.35x.

Demonstrates scalability for large-scale RAG pipelines.

Abstract

Retrieval-Augmented Generation (RAG) is crucial for improving the quality of large language models by injecting proper contexts extracted from external sources. RAG requires high-throughput, low-latency Approximate Nearest Neighbor Search (ANNS) over billion-scale vector databases. Conventional DRAM/SSD solutions face capacity/latency limits, whereas specialized hardware or RDMA clusters lack flexibility or incur network overhead. We present Cosmos, integrating general-purpose cores within CXL memory devices for full ANNS offload and introducing rank-level parallel distance computation to maximize memory bandwidth. We also propose an adjacency-aware data placement that balances search loads across CXL devices based on inter-cluster proximity. Evaluations on SIFT1B and DEEP1B traces show that Cosmos achieves up to 6.72x higher throughput than the baseline CXL system and 2.35x over a state-of-the-art CXL-based solution, demonstrating scalability for RAG pipelines.