技术摘要：PTStore (Prefix Tensor Store，前缀张量存储)

问题陈述

大语言模型（LLM）推理工作负载已成为高性能计算（HPC）数据中心的主要负载，在能源消耗和资源需求方面已超过了训练任务。LLM 推理包含两个阶段：预填充（prefill）（并行处理输入提示词）和解码（decode）（顺序生成 Token）。为了避免注意力机制的冗余计算，系统利用 键值（KV）缓存 来存储中间结果。

虽然最先进的运行时（如 vLLM）可以优化单个 GPU 或节点内的 KV 缓存，但在规模化扩展时面临显著局限性：

1. 缺乏跨节点复用： 现有系统通常无法聚合分布式计算节点的内存。如果一个节点上的请求与另一个节点上的请求共享前缀，第二个节点通常会重新计算该前فع而非复用已缓存的张量。
2. 元数据与延迟瓶颈： 试图实现分布式缓存的方法（如 LMCache、EvoStore）往往由于远程内存访问或复杂的元数据同步（例如将 Radix-Attention 扩展到单节点之外）而导致高 I/O 开销。
3. 内存约束： 单个 GPU 显存不足以支撑超长上下文窗口，而将数据卸载（offloading）到主机内存或 SSD 会引入延迟，从而抵消缓存带来的收益。

核心挑战在于：如何在不产生过高 I/O 或元数据开销的前提下，实现分布在多个计算节点上的大量 GPU 之间，实现可扩展且低延迟的 KV 缓存前缀复用。

方法论：PTStore 架构

PTStore（Prefix Tensor Store）是一种分布式、复制的张量存储系统，旨在通过分发和复制热门 KV 缓存前缀来解决上述局限性。该系统采用客户端-服务器模型，每个计算节点运行一个服务器，聚合本地主机内存和 SSD，为本地及远程 GPU 客户端提供服务。

核心设计原则

1. 增量张量存储（类 Trie 结构）：

   • PTStore 并不存储完整的 KV 块，而是存储新对象与之前存储对象的最长公共前缀（LCP）之间的增量差异（张量）。
   • 这使得前缀能够随着时间的推移在不同的方向上实现无冗余增长，类似于 Trie 树，但通过张量级的粒度来实现。
   • 整合元数据： 为了避免昂贵的分布式 Trie 遍历，PTStore 使用扁平化的元数据结构。每个对象的元数据包含一个唯一的张量 ID 列表。加载操作通过遍历这些 ID 来检查其在复制缓存中是否存在；如果缺失，则从“所有者”服务器远程获取。

2. 分布式分层缓存与复制：

   • 所有权缓存（Owned Cache）： 存储特定服务器负责的增量张量。
   • 复制缓存（Replication Cache）： 在服务器本地存储“热点”（热门）前缀的副本，以提高访问局部性。
   • 权衡管理： 系统管理着一个可配置的阈值，用于平衡所有权缓存与复制缓存。它优先丢弃复制的张量（这些张量可以重新获取），而不是驱逐所有权张量（后者需要刷新到较慢的存储介质），从而在检索速度与存储容量之间取得平衡。

3. 感知访问模式的驱逐策略：

   • PTStore 使用基于频率的驱逐策略（改编自 GDSF），而非最近最少使用（LRU）算法，因为前缀结构意味着早期的张量会被更频繁地访问。
   • 它考虑了大小与频率之间的权衡，确保小型、高频访问的张量不会置换掉体积较大、获取成本较高的张量。

4. RDMA 感知整合：

   • 为了减少数据分散，附加到 LCP 上的增量会被整合到所有者服务器上的单个连续区域中。
   • 加载操作使用批量 RDMA（远程直接内存访问），通过单次 RPC 并行获取分散的段，从而避免了在传输前将数据拷贝到连续区域所带来的开销。

核心贡献

1. 设计原则： 一套用于构建集成增量张量存储、整合元数据和前缀复制功能的分布式仓库的高层原则。
2. PTStore 原型： 一个实现了上述原则的研究原型，具有 C++ 底层 API 和 Python 接口，可无缝集成到 vLLM 等 LLM 运行时中。
3. 性能验证： 通过广泛的实验证明，与最先进的基准方案相比，PTStore 显著降低了 I/O 开销和端到端运行时间。

实验结果

作者在 ALCF Polaris HPC 测试平台（560 个节点，A100 GPU）上使用两种抽取式问答工作负载对 PTStore 进行了评估：WikiQA（长上下文）和 SQUAD（高问题量）。使用的 LLM 为 Mistral-7B-instruct-V2。

基准方案

• vLLM Vanilla： 标准 vLLM，无跨请求前缀共享。
• vLLM Prefix： vLLM 的本地前缀共享版本（节点内共享）。
• EvoStore： 使用增量存储和 RDMA 的分布式张量存储，但缺乏本地前缀复制功能。
• PTStore： 提出的具有分布式感知和本地复制功能的系统。

研究发现

• 弱扩展性表现（8–32 GPUs）： PTStore 的表现显著优于 EvoStore 和 vLLM Prefix。由于 EvoStore 在获取远程前缀时受到高 RDMA I/O 开销的影响，而 PTStore 的本地复制缓解了这一问题，从而在首 Token 延迟（TTFT）方面展现出“脱离式优势”。
• 序列长度扩展性（1k–8k tokens）：
  • 对于短序列（1k），vLLM 的本地缓存具有竞争力。
  • 随着序列长度增加，PTStore 的优势日益扩大。在 8k tokens 时，PTStore 比 vLLM 的前缀缓存快近 2 倍，比 EvoStore 快 20%。
  • 随着上下文变长，性能差距进一步拉大，因为重新计算或远程 I/O 的成本超过了仅进行本地缓存所带来的收益。
• 效率提升： 在长文本问答数据集上，PTStore 的执行效率比那些无法跨节点聚合内存并需要重新生成 KV 缓存的基准方案高出 5–6 倍。

意义与主张

论文声称 PTStore 解决了当前可扩展 LLM 推理服务中的一个关键空白：即现有系统无法高效地在分布式节点间复用 KV 缓存前缀。通过结合增量存储（最小化冗余）、整合元数据（实现快速查询）以及复制策略（优化局部性），PTStore 实现了：

• 通过聚合集群内存，实现有效 KV 缓存容量的数量级扩张。
• 显著降低 TTFT，特别是在重新计算成本极高的长上下文工作负载中。
• 具备可扩展性，避开了以往分布式方法中存在的通信瓶颈和元数据同步问题。

作者将 PTStore 定位为迈向可扩展 AI 推理的基础性一步，并指出未来的工作将侧重于动态内存平衡、基于机器学习的驱逐策略，以及针对真实对话和代码补全轨迹对 LMCache 和 Mooncake 等系统进行更广泛的基准测试。