Scalable Processing-Near-Memory for 1M-Token LLM Inference: CXL-Enabled KV-Cache Management Beyond GPU Limits

TL;DR

The paper tackles the memory and compute bottlenecks of long-context LLM inference by proposing a CXL-enabled Processing-Near-Memory (PNM) KV-cache management framework that offloads KV-cache handling and attention to PNM devices. It introduces a DP-oriented, multi-PNM KV-cache architecture with a GPU-PNM hybrid execution model and a steady-token selection mechanism to maintain large GPU FC throughput while minimizing data movement. Through server- and rack-scale evaluations on models up to 405B and contexts up to 1M tokens, the approach achieves up to 21.9x throughput, up to 60x lower energy per token, and up to 7.3x higher total cost efficiency relative to GPU-only baselines. The work demonstrates that CXL-enabled multi-PNM systems can serve as scalable backbones for future long-context LLM inference, enabling cost-effective and energy-efficient processing beyond traditional GPU limits.

Abstract

The expansion of context windows in large language models (LLMs) to multi-million tokens introduces severe memory and compute bottlenecks, particularly in managing the growing Key-Value (KV) cache. While Compute Express Link (CXL) enables non-eviction frameworks that offload the full KV-cache to scalable external memory, these frameworks still suffer from costly data transfers when recalling non-resident KV tokens to limited GPU memory as context lengths increase. This work proposes scalable Processing-Near-Memory (PNM) for 1M-Token LLM Inference, a CXL-enabled KV-cache management system that coordinates memory and computation beyond GPU limits. Our design offloads token page selection to a PNM accelerator within CXL memory, eliminating costly recalls and enabling larger GPU batch sizes. We further introduce a hybrid parallelization strategy and a steady-token selection mechanism to enhance compute efficiency and scalability. Implemented atop a state-of-the-art CXL-PNM system, our solution delivers consistent performance gains for LLMs with up to 405B parameters and 1M-token contexts. Our PNM-only offloading scheme (PNM-KV) and GPU-PNM hybrid with steady-token execution (PnG-KV) achieve up to 21.9x throughput improvement, up to 60x lower energy per token, and up to 7.3x better total cost efficiency than the baseline, demonstrating that CXL-enabled multi-PNM architectures can serve as a scalable backbone for future long-context LLM inference.

Figures (14)

• Figure 1: Characteristics of long-context LLM decoding. (a) Per-GPU memory demand increases beyond the GPU capacity limit as context length increases. (b) Accuracy comparison of KV-cache management algorithms: dynamic selection with eviction (StreamingLLM xiao2023efficient) and without eviction (ArkVale chen2024arkvale), evaluated on LongBench bai2024longbench tasks.
• Figure 2: CXL-Enabled Memory Expansion System. (a) A GPU system attached with CXL-Mem/PNM. (b) A rack-scale multi-node GPU system with CXL-enabled memory/PNM nodes.
• Figure 3: (a) As sequence length increases from 64 K to 320 K tokens, the average recall count per layer rises steadily, and the runtime recall overhead ratio compared to FC + attention latency also climbs, indicating that longer sequences incur significantly higher recall-induced overhead. (b) Poor FC computation efficiency on GPU under low batch size conditions. The left axis shows the maximum batch size and the right axis shows GPU utilization, both plotted against increasing sequence length (from 32K to 1M). Experiments were conducted on a single GPU (ArkVale) with Llama 3.1-8B. As sequence length grows, memory constraints reduce the achievable batch size, leading to a significant utilization drop
• Figure 4: CXL-enabled PNM system overview.
• Figure 5: Vector processing unit (VPU) for KV-cache management: (a) Hardware architecture of the VPU, (b) VPU configurations for three KV-cache management patterns.