IMPRESS：大模型推理存储优化新突破

全文概览

大模型推理技术正广泛应用于聊天、搜索、代码生成等领域，但其高效运行面临关键挑战：用户提问常共享大量上下文知识，导致系统需频繁加载重复数据。现有方案依赖GPU/CPU内存存储前缀键值对，却因内存容量限制陷入性能瓶颈——从SSD加载数据到GPU的I/O延迟使首次生成令牌时间（TTFT）激增51%-98%。

IMPRESS系统提出一种重要性感知的多级存储架构，通过仅加载关键键值对、优化数据布局和智能缓存策略，将I/O延迟降低1.5-3.8倍，同时保持推理准确率几乎无损。这一创新为大规模LLM推理的实时性与资源效率提供了突破性解决方案。

1. 研究背景与问题

• 大模型推理的挑战：共享上下文导致重复数据加载，SSD-I/O成性能瓶颈。
• 现有方案局限：内存容量不足、缓存策略忽视键值重要性。

2. 关键技术设计

• 重要令牌识别：基于头部相似性提取关键键值，减少加载量。
• 键值重新排序：密集打包重要数据，降低读放大。
• 积分制缓存管理：结合访问频率与重要性比例优化缓存优先级。

3. 实验验证

• 准确率：保留率下准确率下降≤0.2%，优于基线方案。
• 性能提升：TTFT提升1.2-2.8倍，I/O时间减少1.5-3.8倍。
• 鲁棒性：对块大小、数据集规模及模型类型均表现稳定。

阅读收获

1. 理解I/O瓶颈本质：SSD加载对LLM推理性能的影响机制。
2. 掌握重要性感知技术：如何通过头部相似性与动态评分优化存储效率。
3. 获得系统级优化思路：多级存储架构设计与缓存策略的协同作用。
4. 验证方案可行性：IMPRESS在真实场景中的性能与准确率平衡。

研究背景与问题意识

20250314-1317-1.png

大模型推理

大模型推理有海量应用场景，目前已应用于多个领域：

• 聊天（Chat）: ChatGPT
• 搜索（Search）: Perplexity
• 代码（Code）: Cursor
• 问答（Q&A）: ChatPDF

在实际使用过程中，构建大模型问答的请求，需要结合丰富的上下文知识+用户针对性的提问，从应用后端统计来看，大量的用户提问共享相同的上下文知识，这启发了高频访问数据的优化设计。

20250314-1317-2.png

预填充 KV 存储系统

共享的KV缓存数据能重复使用，并最终缩短模型推理过程的首Token响应时延。

图左下，示意不同长度预填充数据在推理过程的时间消耗，大头还是重新计算和CPU数据（DRAM）加载两个过程的时间消耗。

20250314-1317-3.png

讨论当共享前缀键值需要存储到 SSD 中时的情况。

指出，从 SSD 加载前缀键值到 GPU 成为了一个新的性能瓶颈，导致首次生成令牌时间（TTFT）显著增加（51%-98%）。图表展示了随着前缀令牌数量的增加，TTFT 的变化情况，表明 I/O 延迟在关键路径上对系统性能产生了重大影响。

20250314-1317-4.png

相关工作

图片介绍了与前缀键值存储相关的现有工作。

大多数系统仅将前缀键值存储在 GPU 和 CPU 内存中，但由于内存空间有限，这种方法很快会达到极限。另一种方法是根据调度器的预测预先将数据加载到 CPU 内存中，但这在高请求量或抢占式调度下存在局限性。

最后，提出了一个关键问题：是否有可能减少需要加载的键值数据，以进一步优化系统性能。

===

• 大多数现有系统仅将前缀键值存储在 GPU 和/或 CPU 内存中
  • PromptCache-MLSys24, RAGCache-arxiv24, ChunkAttention-arxiv24, SGLang-arxiv23

  GPU 和 CPU 内存的有限空间很快就会耗尽

• 根据调度器的预测，预先将它们加载到 CPU 内存中
  • AttentionStore-ATC24

20250314-1317-5.png

KV缓存优化的机会

图片探讨通过仅保留和加载重要的键值来优化系统性能的机会。在解码阶段，仅保留重要的键值可以保持相同的准确度水平。进一步提出，在预填充阶段也仅加载重要的键值，以减少 I/O 瓶颈和首次生成令牌时间（TTFT），从而提高整体效率。

===

左侧图表（H2O-NeurIPS23）显示了键值对的重要性，其中一些键值被标记为不重要并被移除。

右侧流程图（InfiniGen-OSDI24）展示了在离线倾斜、预填充和解码阶段如何选择和处理键值。GPU 和 CPU 之间的数据流动包括部分权重索引生成、键值选择、注意力机制、前馈网络（FFN）、轻量级注意力和预取等步骤。

Question

在预填充阶段仅加载重要的键值是否可以减少 I/O 瓶颈和首次生成令牌时间（TTFT）？

20250314-1317-6.png

存在的挑战

图片讨论了在识别重要键值对时引入大量 I/O 操作所带来的挑战。

• 左侧图表展示了数据从 DRAM 和 SSD 加载到 GPU 的两种方式，所有参数加载（红），部分参数加载（蓝）。
• 右侧图表则显示了在预确定重要键值对时，准确率会有所下降，尤其是在使用静态预确定方法（SPI）时，与原始动态确定方法（ORI）相比，准确率和 F1 分数都有显著降低。这表明在减少 I/O 操作的同时，如何保持模型的准确性和性能是一个需要解决的问题。

20250314-1317-7.png

图片指出现有前缀键值存储和缓存系统的两个主要问题。

• 首先，在存储方面，由于每个块都包含了重要和不重要的键值对的混合，导致了读放大的问题，即读取数据时需要处理大量不必要的信息。
• 其次，在缓存方面，现有的策略仅基于最近使用或频率来决定缓存内容，而忽略了键值对本身的重要性，这可能导致缓存效率低下。

图表 (a) 显示了每个块中重要键值对的比例分布，而图表 (b) 则展示了不同访问频率下块中重要令牌的平均比例，进一步说明了现有系统在考虑键值对重要性方面的不足。

观察与研究方法

20250314-1317-8.png

IMPRESS 架构设计

介绍了 IMPRESS 架构设计。用户请求包含前缀和查询，系统通过三个主要设计来处理这些请求：

1. 重要令牌识别 （1）：从磁盘读取前缀键值对，并在 CPU 和 GPU 内存中进行处理，选择重要的键值对。
2. 键值重新排序 （2.1）：对前缀键值对进行重新排序，以优化后续处理。
3. 缓存替换 （2.2）：管理前缀键值的缓存，确保高效的数据访问和处理。

图表展示了数据流和控制消息的流向，实线箭头表示数据传输，虚线箭头表示控制消息。整体设计旨在提高系统的性能和效率。

===

设计 1：重要令牌识别

• 从磁盘读取前缀键值对。
• CPU 内存中存储元数据和 CPU 缓存。
• GPU 内存中选择重要的键值对，并提供运行时空间和 GPU 缓存。

设计 2.1：键值重新排序

• 对前缀键值对进行重新排序。

设计 2.2：缓存替换

• 前缀键值管理

20250314-1317-9.png

研究数据

图片表明在大型语言模型（LLM）的同一层中，不同头部之间的重要令牌索引集具有很高的相似性。

通过计算两个头部之间的 Jaccard 相似度，可以量化这种相似性。右侧的图表进一步验证了这一观察结果，无论选择前 10% 还是前 40% 的最重要令牌，不同规模的 LLM 都表现出相似的分布模式。这表明重要令牌索引的相似性是跨不同 LLM 规模和重要键值比率存在的普遍现象。

关于 Jaccard 相似度

定义： Jaccard 相似度（Jaccard Similarity）用于衡量两个集合的相似度，计算方式如下：

特点：

• 主要用于集合比较，特别适用于离散数据（如文本单词集合、二元特征向量）。
• 值域在 [0,1][0,1][0,1] 之间，值越大表示相似度越高。
• 适用于衡量两个集合的相似性，如文本、标签分类等。

应用：

• 文本相似性（基于词袋模型）
• 计算机视觉中的图像相似性（如二进制特征匹配）

20250314-1317-11.png

20250314-1317-10.png

相似性引导的重要令牌识别机制

图片介绍了相似性引导的重要令牌识别方法。

关键思想是利用从几个选定头部提取的重要令牌索引集，来近似其他头部的重要令牌索引集。

左侧图表展示了在不使用该方法时，需要加载所有键和部分值的过程；而右侧图表则展示了使用该方法后，只需加载部分键和值，从而减少了数据传输量和处理时间，提高了系统的效率。

20250314-1317-13.png

KV 缓存重排序

图片介绍了键值重新排序的目标和关键思想，即通过将重要的键值重新排序并打包到更密集的块中来减少读放大。

• 左侧图表展示了重新排序前后磁盘上令牌分布的变化
• 右侧图表则指出了重新排序可能破坏基数树结构的问题，并提出了避免跨节点重新排序和添加映射列表的解决方案。这些方法有助于在保持数据结构完整性的前提下，提高系统的读取效率。

20250314-1317-12.png

积分制缓存管理

图片介绍基于分数的缓存管理方法，该方法通过计算块的访问频率和重要键值对的比例来确定数据的优先级。

分数越高，数据被放入更快缓存中的优先级越高。左侧图表展示了两个块的访问频率和重要比例，右侧图表则比较了 LFU 方法和本文提出的方法在不同缓存层次中的表现。结果显示，本文方法能够更有效地利用缓存资源，提高系统的性能。

• 关键思想：基于评分的数据准入和缓存替换。
• 分数 = 块访问频率 * 重要键值对的比例。

分数越高，进入更快介质缓存的优先级越高。

Note

FAST 25 很多论文都在讨论推理场景存储访问优化，其核心又集中在缓存管理机制，与早期大数据推荐系统的缓存管理相比，大模型的热数据管理，呈现出更大范围的随机性，全部数据加载是不切实际的。从效率角度来看，大模型的访问入口可能最终不是集中化的，更可能是地域、分布式。

实验评估

20250314-1317-14.png

实验设计

详细介绍了实验的设置，包括系统配置、工作负载和数据集、比较的系统以及默认设置。

系统配置包括使用的 CPU、GPU 和内存规格。工作负载和数据集涵盖了多个数据集和模型，以评估不同系统的性能。比较的系统包括 ReComp、AS-like、AS+H2O+LRU、AS+H2O+LFU 和 IMPRESS，每种系统都有其特定的特点和优化方法。默认设置则规定了缓存大小和块大小，以确保实验的一致性和可比性。

===

• 系统配置
  • CPU: 2 × AMD EPYC 7763
  • GPU: 1 × NVIDIA A100 (80GB)
  • 内存 & SSD: 128 GB DRAM, 2TB SSD (5GB/s)
• 工作负载和数据集
  • 数据集: PIQA, RTE, COPA, 和 OpenBookQA
    • 前缀大小: 55GB, 57GB, 64GB, 65GB
  • 模型: OPT-6.7B, OPT-13B, OPT-30B
• 比较的系统
  • ReComp: 不重用前缀键值对的重新计算
  • AS-like: 具有异步键值加载、无调度器的 AttentionStore
  • AS+H2O+LRU: 在带有 LRU 的 AttentionStore 上添加 H2O
  • AS+H2O+LFU: 在带有 LFU 的 AttentionStore 上添加 H2O
  • IMPRESS: 我们的三种优化在 H2O 之上
• 默认设置
  1. 缓存大小: 10GB GPU HBM, 32GB CPU DRAM
  2. 块大小: 64 个令牌的键或值

20250314-1317-15.png

推理准确度评估

通过多个图表展示了不同方法在四个数据集上的推理准确度表现。

结果显示，与 ReComp 方法相比，其他方法（如 AS+H2O+LRU 和 IMPRESS）在不同保留率下仍能保持较高的准确度，且平均准确度下降不超过 0.2%。这表明这些优化方法在提高系统性能的同时，对模型的推理准确度影响很小，具有较好的效果和稳定性。

20250314-1317-16.png

首Token 时延比较

图片通过多个图表展示了 IMPRESS 方法在不同数据集和模型规模下的 TTFT 和 I/O 时间表现。

结果显示，IMPRESS 方法在 TTFT 方面显著优于其他方法，并且通过减少 I/O 时间实现了性能的大幅提升。单个技术的分析进一步验证了这些优化措施的有效性。

IMPRESS 在所有方法中表现出色，与现有最佳解决方案相比，有 1.2 倍到 2.8 倍的提升，这是由于 I/O 时间减少了 1.5 倍到 3.8 倍。

20250314-1317-17.png

通过多个图表展示了 IMPRESS 方法在不同块大小、数据集规模和 Llama 模型上的 TTFT 表现。

结果显示，无论是在不同的块大小、数据集规模还是在不同的模型上，IMPRESS 方法都显著优于 AS+H2O+LFU 方法。这表明 IMPRESS 方法具有良好的鲁棒性和适应性，在各种情况下都能保持优异的性能。

敏感性分析 的实践意义

通过敏感性分析，可以得出以下结论：

(1) 识别关键优化方向

• 如果模型对某个参数特别敏感（如块大小或数据集规模），说明该参数是优化的重点。
• 例如，如果 TTFT 随块大小的增加而显著降低，则应优先考虑增大块大小。

(2) 验证优化方法的有效性

• 对比不同优化方法（如 AS+H2O+LFU 和 IMPRESS）在各种条件下的表现，评估其鲁棒性和适应性。
• 如果某种方法在所有条件下都优于其他方法，则说明该方法具有较强的通用性。

(3) 指导硬件和软件设计

• 敏感性分析的结果可以为硬件选型（如 GPU 内存容量、SSD 带宽）和软件优化（如缓存管理策略）提供依据。

总结&结论

20250314-1317-18.png

• 问题
  • 当从 SSD 加载共享前缀键值对（KVs）用于大型语言模型（LLM）时，I/O 成为瓶颈。
• 关键思想
  • 在预填充阶段仅加载重要的键值对（KVs）。
• 挑战
  • 为了识别重要键值对，引入了大量的 I/O。
  • 存储和缓存系统是次优的。
• iCache 中的技术
  • 基于相似性的关键令牌识别
  • 键值重新排序和基于分数的缓存管理
• 结果
  • IMPRESS 在保持相同推理准确度的情况下，优于其他替代方案。

延伸思考

这次分享的内容就到这里了，或许以下几个问题，能够启发你更多的思考，欢迎留言，说说你的想法~

1. 场景扩展性：IMPRESS是否适用于非文本型数据（如图像、视频）的推理任务？
2. 动态适应性：如何在实时变化的用户请求中，进一步优化重要键值的在线识别效率？
3. 硬件协同：若结合CXL或NVMe等高速接口技术，能否进一步突破存储层级的性能极限？

原文标题：IMPRESS: An Importance-Informed Multi-Tier Prefix KV Storage System for Large Language Model Inference