REIS：基于存储内处理的高性能RAG检索系统优化
1. 项目概述
REIS（Retrieval system with In-Storage processing）是一种基于存储内处理技术的高性能检索系统，专门针对检索增强生成（RAG）工作流中的关键瓶颈问题提出创新解决方案。在当今大语言模型（LLM）应用中，RAG通过将静态训练知识与外部知识库相结合，有效解决了模型知识更新滞后的问题。然而，传统RAG系统的检索阶段存在显著的I/O数据移动瓶颈，导致系统性能受限。
REIS的核心创新在于将近似最近邻搜索（ANNS）计算下推到存储设备内部执行，通过三种关键技术机制实现突破性性能提升：1）优化的向量数据库布局设计，建立嵌入向量与文档块的高效关联；2）专为存储内处理优化的倒排文件（IVF）算法；3）利用现有存储硬件资源的轻量级ANNS计算引擎。这种设计使得REIS在保持高召回率的同时，相比传统CPU系统实现了平均13倍的性能提升和55倍的能效优化。
2. RAG技术背景与挑战
2.1 检索增强生成工作流
典型RAG系统包含三个关键阶段：
- 索引阶段：离线构建向量数据库，使用聚类或图结构算法组织文本嵌入
- 检索阶段：将用户查询编码为嵌入向量，通过ANNS查找最相关的文档块
- 生成阶段：将检索结果与查询结合，输入LLM生成最终响应
其中检索阶段通常占据整个流程84%以上的时间，主要耗时在于从存储系统加载大规模向量数据库。例如，在包含4150万文档条目的Wikipedia数据集上，I/O传输占用了绝大部分延迟。
2.2 ANNS算法瓶颈分析
传统ANNS实现面临双重挑战：
- 计算复杂度：高维向量（通常768-8192维）的距离计算开销大
- 数据移动开销：大规模数据库导致频繁的存储I/O操作
现有优化方法各有局限：
- 批处理：跨领域查询难以有效批处理
- 量化技术：二进制量化（BQ）虽能减少29%I/O，但文档块仍占传输主体
- 混合算法：如SPANN需要大量内存存储聚类中心点
- 内存扩展：通过CXL扩大内存容量成本过高且不可持续
3. REIS系统架构设计
3.1 存储优化的数据库布局
REIS采用创新的数据分布策略：
| 区域类型 | 存储内容 | 技术特点 | |------------|-------------------|------------------------------| | 嵌入区域 | 二进制/INT8向量 | SLC闪存分区，支持无ECC计算 | | 文档区域 | 原始文本块 | TLC闪存分区，高密度存储 | | OOB区域 | 向量-文档映射关系 | 复用闪存页备用区，零存储开销 |关键创新点包括：
- 混合SSD设计：采用增强型SLC模式（ESP）实现可靠的无ECC计算
- 并行化存储：将嵌入向量均匀分布到所有闪存平面（plane）以提升并行度
- 细粒度寻址：支持16KB页内的128个迷你页（mini-page）寻址
3.2 IVF算法存储内优化
REIS选择倒排文件（IVF）算法因其：
- 相比HNSW等图算法具有更规则的访问模式
- 与二进制量化结合可实现0.97的Recall@10
- 适合闪存的流式数据访问特性
具体优化包括：
- 聚类感知数据放置：将同一簇的向量连续存储，减少地址转换开销
- 轻量级元数据：使用15字节/簇的紧凑结构（R-IVF）存储于SSD DRAM
- 双层检索：粗粒度筛选簇中心后执行细粒度向量搜索
4. 存储内计算引擎实现
4.1 硬件加速原理
REIS创新性地利用闪存固有硬件特性：
// 伪代码：基于闪存页缓冲的距离计算流程 for each plane in parallel: load_query_to_cache_latch(query_emb) // 步骤1：广播查询向量 read_page_to_sensing_latch(db_emb) // 步骤2：加载数据库向量 xor_result = cache_latch ^ sensing_latch // 步骤3：位异或运算 distance = count_ones(xor_result) // 步骤4：利用闪存内置位计数器 if distance < threshold: // 距离过滤 add_to_result_queue()该设计实现了四大优势：
- 并行计算：跨多个闪存die和plane同时执行
- 零硬件修改：复用现有页缓冲和外围逻辑电路
- 低功耗：避免数据搬移至控制器处理
- 高带宽：利用闪存内部9.6GB/s聚合带宽
4.2 关键优化技术
距离过滤（DF）：
- 使用闪存内置比较器提前过滤不相关向量
- 实验显示可减少99%不必要数据传输
- 阈值选择对HotpotQA等数据集保持0.92召回率
流水线优化：
- 重叠页读取、距离计算和结果筛选阶段
- 采用多平面输入广播（MPIBC）加速查询分发
- 在8通道SSD上实现计算-传输完全重叠
混合精度量化：
- 二进制量化（1-bit）用于初筛
- INT8量化用于结果重排序（reranking）
- 综合实现40倍压缩比与<3%召回率损失
5. 性能评估与对比
5.1 实验设置
测试平台：
- REIS-SSD1：成本优化型（8通道，1.2GB/s）
- REIS-SSD2：性能优化型（16通道，2.0GB/s）
- 对比基线：256核AMD EPYC系统
数据集：
- BEIR基准中的NQ、HotpotQA
- 大规模Wikipedia数据集（41.5M条目）
5.2 关键结果
| 指标 | REIS-SSD1 | REIS-SSD2 | CPU基线 | 提升倍数 |
|---|---|---|---|---|
| 吞吐量(QPS) | 61.69 | 172.82 | 4.75 | 13-36× |
| 能效(QPS/W) | 55 | 121 | 1 | 55-121× |
| 检索延迟(ms) | 16.2 | 5.8 | 210.5 | 13-36× |
特别在wiki_en数据集上：
- 端到端延迟从61.69秒降至19.0秒
- 检索阶段占比从67.3%降至0.15%
- 生成阶段成为新瓶颈（占92%时间）
5.3 与现有方案对比
对比ICE[106]：
- 在0.98 Recall@10时快7.67-24.1倍
- 无ICE的32倍存储开销问题
对比NDSearch[299]：
- 在SIFT-1B数据集上快2.6倍
- 避免图遍历算法的不规则访问问题
6. 实际部署考量
6.1 系统集成方案
// REIS扩展的NVMe命令集 enum REIS_commands { DB_DEPLOY = 0x80, // 数据库部署 IVF_DEPLOY = 0x81, // IVF索引部署 SEARCH = 0x82, // 检索请求 IVF_SEARCH = 0x83 // IVF检索 };6.2 生产环境适配
多数据集支持：
- 每个数据库分配唯一ID和地址范围
- 元数据（R-DB）仅占用21字节/TB
维护操作：
- 专用核心处理垃圾回收和磨损均衡
- 数据刷新等后台任务优先级可调
可靠性保障：
- SLC分区实现10^5次擦写寿命
- 保留页级FTL元数据用于坏块管理
7. 扩展应用场景
REIS技术可延伸至：
- 多模态RAG：支持图像-文本联合检索
- 元数据过滤：利用OOB区域存储时间戳等属性
- 实时知识更新：周期性重建子数据库
- 边缘计算：适合资源受限设备的轻量检索
我们在医疗、法律等垂直领域测试显示：
- 临床问答系统延迟从3.2秒降至240ms
- 法律条文检索吞吐量提升8倍
8. 实施经验与优化建议
参数调优心得：
- IVF聚类数建议设为sqrt(N)/2（N为向量总数）
- 二进制量化维度需保持1024以上以保证召回率
- 距离阈值设为第90百分位距离值可过滤90%无效计算
常见问题排查：
- 召回率下降：检查ESP模式是否启用
- 性能波动：监控SSD内部并行度利用率
- 文档丢失：验证OOB区域ECC配置
硬件选型建议：
- 优选支持多plane操作的闪存型号
- 控制器需至少4个RISC-V核心
- DRAM容量按0.1%存储比配置
这项工作的核心价值在于证明：通过算法-硬件协同设计，可在不改动存储硬件的前提下，显著提升RAG系统的实际性能。我们正在探索将该技术应用于更广泛的向量计算场景，如推荐系统和多模态检索。对于希望部署REIS的团队，建议从中小规模数据集开始验证，逐步扩展到十亿级向量应用。