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

Recent advances in LLMs like GPT-4, GLaM, and PaLM have demonstrated unprecedented capabilities in NLP tasks. However, scaling these models to handle long contexts (millions of tokens) introduces hardware bottlenecks, especially in memory capacity and bandwidth. Traditional GPU-based systems struggle with storing and accessing large KV caches, which are essential for efficient attention mechanisms. Sparse attention algorithms, such as retrieval-based methods like GQA, DeepSeek, and MLA, have been proposed to reduce computational load by selectively attending to relevant tokens. Despite these innovations, existing hardware solutions—such as offloading KV caches to host memory or employing fixed-function PIM/PNM units—fail to meet the combined demands of capacity, bandwidth, and algorithmic flexibility. This gap motivates the development of a heterogeneous system that can adapt to diverse algorithms and long contexts, ensuring scalable and efficient deployment of future large models.

The core challenge lies in balancing the conflicting demands of large KV cache storage, high bandwidth for retrieval operations, and flexible computation to support evolving sparse attention algorithms. GPU memory limitations restrict batch sizes and increase latency, while fixed-function near-memory units cannot accommodate algorithmic diversity. Existing solutions either sacrifice capacity or flexibility, leading to performance bottlenecks. The problem becomes more acute with multi-turn agentic workloads, where context length can reach hundreds of thousands of tokens, demanding hardware that can sustain high operational intensity without compromising energy efficiency. Designing a system that integrates large-capacity memory, high-bandwidth access, and general-purpose compute remains an open and pressing issue.

This work introduces KARAT, a novel heterogeneous hardware architecture that combines large LPDDR memory modules with general-purpose GPNM cores, supporting diverse retrieval-based sparse attention algorithms. Unlike prior fixed-function PIM/PNM units, KARAT's cores are programmable, enabling adaptation to algorithmic variations. The system employs micro-batch scheduling (OFMS) to overlap communication and computation, reducing pipeline stalls, and context-aware rebalancing (CMR) to equalize token counts across micro-batches, minimizing idle time. The architecture facilitates seamless collaboration between GPU and KARAT devices, with data placement strategies ensuring minimal communication overhead. These innovations collectively address the core bottlenecks of capacity, bandwidth, compute, and flexibility, enabling efficient long-context processing at scale.

• �� System architecture: Decouples dense and sparse operations, assigning dense projections and MoE layers to GPU nodes, while KARAT nodes handle KV cache storage and retrieval operations.
• �� Hardware design: Each KARAT device integrates large LPDDR memory with GPNM cores capable of executing diverse sparse attention algorithms.
• �� Scheduling strategies: OFMS overlaps micro-batch communication with computation by interleaving tasks; CMR balances token counts across micro-batches based on context length, reducing pipeline bubbles.
• �� Data flow: During prefill, KV caches are transferred to KARAT nodes; during decode, GPU constructs index queries, KARAT performs index scan, gather, and attention, then results are sent back to GPU.
• �� Optimization: Data placement ensures requests are localized, minimizing cross-device communication; pipeline is optimized by dynamic scheduling to adapt to workload variance.

Using three models—GLM-5.2, MiniMax-M3, DeepSeek-V3.2—the authors evaluate throughput, energy efficiency, and latency under real agentic traces with context lengths up to 1 million tokens. Baseline comparisons involve GPU-only systems running the same models with native sparse attention methods. The experiments test the impact of OFMS and CMR, measuring micro-batch throughput, pipeline bubbles, and resource utilization. Hyperparameters include micro-batch sizes from 1 to 16, with service-level objectives set at 100ms per token. Results demonstrate consistent performance gains, validating the effectiveness of the heterogeneous approach and scheduling techniques.

The proposed system achieves 2.09× to 6.13× throughput improvements across models under the P99 time-between-tokens constraint, with energy efficiency gains of 1.36× to 3.21×. Micro-batch scheduling (OFMS) effectively overlaps communication with computation, reducing idle time. Context-aware rebalancing (CMR) maintains balanced token counts despite variance in request length, further boosting pipeline efficiency. The hardware supports multiple sparse attention variants, demonstrating flexibility and robustness. Overall, the results confirm that heterogeneous hardware collaboration can overcome traditional bottlenecks in long-context LLM deployment, enabling scalable and energy-efficient inference.

This architecture is ideal for deploying large-scale, multi-turn conversational AI, long-text summarization, and document understanding systems. It enables efficient inference with millions of tokens, reducing latency and energy consumption. The system's flexibility supports evolving sparse attention algorithms, making it adaptable for future AI models. It can be integrated into data centers and cloud services, facilitating real-time, long-context AI applications across industries such as customer service, content creation, and legal analysis.

Dependence on high-speed network infrastructure may limit performance in distributed environments with high latency. Hardware complexity and costs are significant, requiring advanced manufacturing and integration. Supporting a broad spectrum of algorithms demands ongoing hardware/software co-design, which could hinder rapid adaptation. Future work must address these challenges to ensure broader adoption and cost-effectiveness.