Technical Summary: PTStore (Prefix Tensor Store)

Problem Statement

Large Language Model (LLM) inference workloads have become the dominant load in high-performance computing (HPC) data centers, surpassing training in energy consumption and resource demand. LLM inference consists of two phases: prefill (processing the input prompt in parallel) and decode (generating tokens sequentially). To avoid redundant computation of attention mechanisms, systems utilize a Key-Value (KV) cache to store intermediate results.

While state-of-the-art runtimes (e.g., vLLM) optimize KV caching within a single GPU or node, they face significant limitations at scale:

1. Lack of Cross-Node Reuse: Existing systems often fail to aggregate memory across distributed compute nodes. If a request on one node shares a prefix with a request on another node, the second node typically recomputes the prefix rather than reusing the cached tensors.
2. Metadata and Latency Bottlenecks: Approaches that attempt distributed caching (e.g., LMCache, EvoStore) often suffer from high I/O overheads due to remote memory access or complex metadata synchronization (e.g., scaling Radix-Attention beyond a single node).
3. Memory Constraints: Individual GPU memory is insufficient for large context windows, and offloading to host memory or SSDs introduces latency that negates the benefits of caching.

The core challenge is enabling scalable, low-latency reuse of KV cache prefixes across a large number of GPUs distributed over many compute nodes without incurring prohibitive I/O or metadata overheads.

Methodology: PTStore Architecture

PTStore (Prefix Tensor Store) is a distributed, replicated tensor store designed to address these limitations by distributing and replicating popular KV cache prefixes. The system employs a client-server model where each compute node runs a server that aggregates local host memory and SSDs to serve both local and remote GPU clients.

Key Design Principles

1. Incremental Tensor Storage (Trie-like Structure):

   • Instead of storing full KV blocks, PTStore stores incremental differences (tensors) between a new object and the longest common prefix (LCP) of previously stored objects.
   • This allows prefixes to grow redundancy-free over time in divergent directions, similar to a trie, but implemented via tensor-level granularity.
   • Consolidated Metadata: To avoid expensive distributed trie traversal, PTStore uses a flat metadata structure. Each object's metadata contains a list of unique tensor IDs. A load operation iterates through these IDs to check for local existence in the replication cache; if missing, it fetches them remotely from the "owner" server.

2. Distributed Hierarchical Caching with Replication:

   • Owned Cache: Stores the incremental tensors that a specific server is responsible for.
   • Replication Cache: Stores copies of "hot" (popular) prefixes locally on the server to improve access locality.
   • Trade-off Management: The system manages a configurable threshold between owned and replicated caches. It prioritizes discarding replicated tensors (which can be re-fetched) over evicting owned tensors (which require flushing to slower storage) to balance retrieval speed against storage capacity.

3. Access Pattern-Aware Eviction:

   • PTStore utilizes a frequency-based eviction policy (adapted from GDSF) rather than Least Recently Used (LRU), as prefix structures mean earlier tensors are accessed more frequently.
   • It accounts for the size vs. frequency trade-off, ensuring small, frequent tensors do not displace larger, expensive-to-fetch tensors.

4. RDMA-Aware Consolidation:

   • To minimize scattering, increments appended to an LCP are consolidated into a single contiguous region on the owner server.
   • Load operations use bulk RDMA (Remote Direct Memory Access) to fetch scattered segments in parallel via a single RPC, avoiding the overhead of copying data into a contiguous region before transfer.

Key Contributions

1. Design Principles: A set of high-level principles for a distributed repository that integrates incremental tensor storage, consolidated metadata, and prefix replication.
2. PTStore Prototype: A research prototype implementing these principles, featuring a C++ low-level API and a Python interface for seamless integration with LLM runtimes like vLLM.
3. Performance Validation: Extensive experiments demonstrating significant reductions in I/O overhead and end-to-end runtime compared to state-of-the-art baselines.

Experimental Results

The authors evaluated PTStore on the ALCF Polaris HPC testbed (560 nodes, A100 GPUs) using two extractive QA workloads: WikiQA (long context) and SQUAD (high volume of questions). The LLM used was Mistral-7B-instruct-V2.

Baselines

• vLLM Vanilla: Standard vLLM with no cross-request prefix sharing.
• vLLM Prefix: vLLM with local prefix sharing (within a node).
• EvoStore: A distributed tensor store using incremental storage and RDMA but lacking local prefix replication.
• PTStore: The proposed system with distributed awareness and local replication.

Findings

• Weak Scalability (8–32 GPUs): PTStore significantly outperformed EvoStore and vLLM Prefix. While EvoStore suffered from high RDMA I/O overheads when fetching remote prefixes, PTStore's local replication mitigated this, resulting in a "detached advantage" in Time to First Token (TTFT).
• Sequence Length Scalability (1k–8k tokens):
  • For short sequences (1k), vLLM's local caching was competitive.
  • As sequence length increased, PTStore's advantage grew. At 8k tokens, PTStore was nearly 2x faster than vLLM's prefix caching and 20% faster than EvoStore.
  • The performance gap widened with longer contexts because the cost of recomputation or remote I/O outweighed the benefits of local-only caching.
• Efficiency Gains: On long passage Q&A datasets, PTStore executed inferences 5–6 times more efficiently than baselines that do not aggregate memory across nodes and require regenerating KV caches.

Significance and Claims

The paper claims that PTStore addresses a critical gap in scalable LLM inference serving: the inability of current systems to efficiently reuse KV cache prefixes across distributed nodes. By combining incremental storage to minimize redundancy, consolidated metadata for rapid queries, and a replication strategy to optimize locality, PTStore enables:

• Orders of magnitude expansion of the effective KV cache size by aggregating memory across the cluster.
• Significant reduction in TTFT, particularly for long-context workloads where recomputation is costly.
• Scalability that avoids the communication bottlenecks and metadata synchronization issues plaguing previous distributed approaches.

The authors position PTStore as a foundational step toward scalable AI inference, noting that future work will focus on dynamic memory balancing, ML-based eviction policies, and broader benchmarking against systems like LMCache and Mooncake on real-world conversational and code-completion traces.