Sparse Attention Remapping with Clustering for Efficient LLM Decoding on PIM

TL;DR

This work tackles the memory bandwidth bottleneck in autoregressive LLM decoding by introducing STARC, a clustering-based data mapping that remaps KV cache entries to contiguous memory regions aligned with PIM bank structures. STARC performs online K-means clustering to group semantically similar keys, assigning clusters to hardware-friendly memory rows, and retrieves tokens at the cluster level using precomputed centroids, enabling selective attention with minimal data movement. Compared with token-wise and page-wise sparsity baselines, STARC achieves up to 19%–31% lower attention latency and 19%–27% lower energy on PIM systems, and up to 54%–74% latency reduction and 45%–67% energy reduction relative to full KV retrieval, while preserving accuracy close to state-of-the-art sparse methods. The results demonstrate that co-designing memory layout with sparse attention on PIM can realize efficient, high-quality long-context LLM inference, making sparse attention practical for hardware-friendly deployments.

Abstract

Transformer-based models are the foundation of modern machine learning, but their execution, particularly during autoregressive decoding in large language models (LLMs), places significant pressure on memory systems due to frequent memory accesses and growing key-value (KV) caches. This creates a bottleneck in memory bandwidth, especially as context lengths increase. Processing-in-memory (PIM) architectures are a promising solution, offering high internal bandwidth and compute parallelism near memory. However, current PIM designs are primarily optimized for dense attention and struggle with the dynamic, irregular access patterns introduced by modern KV cache sparsity techniques. Consequently, they suffer from workload imbalance, reducing throughput and resource utilization. In this work, we propose STARC, a novel sparsity-optimized data mapping scheme tailored specifically for efficient LLM decoding on PIM architectures. STARC clusters KV pairs by semantic similarity and maps them to contiguous memory regions aligned with PIM bank structures. During decoding, queries retrieve relevant tokens at cluster granularity by matching against precomputed centroids, enabling selective attention and parallel processing without frequent reclustering or data movement overhead. Experiments on the HBM-PIM system show that, compared to common token-wise sparsity methods, STARC reduces attention-layer latency by 19%--31% and energy consumption by 19%--27%. Under a KV cache budget of 1024, it achieves up to 54%--74% latency reduction and 45%--67% energy reduction compared to full KV cache retrieval. Meanwhile, STARC maintains model accuracy comparable to state-of-the-art sparse attention methods, demonstrating its effectiveness in enabling efficient and hardware-friendly long-context LLM inference on PIM architectures.

Figures (12)

• Figure 1: Enhanced execution efficiency through STARC. Due to the coarse granularity of PIM, directly applying sparsity to KV caches often fails to skip computation. STARC addresses this by clustering keys and values such that selected tokens are physically co-located, enabling effective computation skipping and realizing the speedup benefits of sparsity on PIM.
• Figure 2: Targeted system overview: the QKV Generation and the Feed-Forward Networks are executed on xPU such as GPU and NPU, while the Multi-Head Attention is executed on PIM.
• Figure 3: Common attention sparsity patterns.
• Figure 4: Page-wise retrieval with less important tokens.
• Figure 5: Supporting sparsity on PIM with STARC. PIM & sparsity both achieve higher throughput, but enforce vastly different granularities. As a result, the two do not combine readily. STARC applies clustering to improve KV locality in memory, mitigating the granularity required by sparsity while executing with PIM.