Search papers, labs, and topics across Lattice.
This paper introduces a heterogeneous serving system that optimizes the decode phase for retrieval-based sparse attention in large language models (LLMs) by relocating the key-value (KV) cache to processing-near-memory (PNM) nodes. By partitioning operations between GPU nodes and PNM nodes, the proposed KARAT design meets four derived requirements, enabling efficient handling of diverse sparse attention algorithms. The system demonstrates significant performance improvements, achieving 2.09-6.13x throughput enhancements over GPU-only baselines while also improving training-free sparse attention methods by 1.36-3.21x.
Relocating the KV cache to processing-near-memory nodes can boost LLM throughput by over 6x while supporting evolving sparse attention methods.
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.
