Heterogeneous LLM Serving with General-Purpose Processing-Near-Memory for Retrieval-Based Sparse Attention

2026-08-04

AI Summary

This work addresses the GPU KV cache memory and access bottlenecks in large language models operating at million-token context lengths, which arise from retrieval-based sparse attention mechanisms. To overcome these limitations, the authors propose KARAT—the first general-purpose near-memory computing (PNM) architecture capable of supporting four distinct classes of sparse attention patterns. By offloading KV cache storage and operations to dedicated PNM nodes while allowing the GPU to focus on model weights and Mixture-of-Experts (MoE) computation, KARAT enables efficient heterogeneous collaboration. This is further enhanced by opportunistic fine-grained micro-batch scheduling (OFMS) and context-aware micro-batch rebalancing (CMR). Evaluated on three state-of-the-art models and real-world agent trajectories, KARAT achieves 2.09–6.13× higher throughput per watt than a pure GPU baseline under service-level objectives, and delivers 1.36–3.21× speedup for training-free sparse attention methods.

Abstract

This paper presents a heterogeneous decode-phase serving system that relocates the KV cache out of GPU memory, motivated by the retrieval-based sparse attention that recent frontier LLMs adopt to serve million-token contexts. It partitions a decode step by operation type: GPU nodes hold the model weights and execute the projections and MoE layers, while processing-near-memory (PNM) nodes hold the KV cache and index keys and execute every operation that reads them. We first show that the assumptions behind prior PIM and PNM designs no longer hold for these operations, and derive four design requirements for such a node. From these requirements, we propose KARAT (KV-cache-resident Accelerator for Retrieval-based ATtention), a general-purpose PNM design that is the design point meeting all four. A KARAT device combines large LPDDR capacity with general-purpose compute sized for the retrieval indexer, serving an operational intensity beyond what PIM/PNM designs built for low-intensity GEMV target while accommodating diverse sparse attention algorithms that fixed-function units cannot support as they evolve. To reduce pipeline bubbles as the two device types alternate between micro-batches, we further propose opportunistic, fine-grained micro-batch scheduling (OFMS), which hides expert all-to-all behind the other micro-batch's GEMMs, and context-length-aware micro-batch rebalancing (CMR), which equalizes their token counts despite the variance in context length. Across three state-of-the-art models and real agentic traces, our proposed system improves throughput per TDP under a service-level objective by 2.09-6.13x over a GPU-only baseline and runs training-free sparse attention methods with 1.36-3.21x improvements.

Problem

• KV cache
• sparse attention
• processing-near-memory
• heterogeneous serving
• large language models

Innovation

• Processing-Near-Memory
• Sparse Attention
• KV Cache
• Heterogeneous Serving
• Micro-batch Scheduling