arXiv:2505.05772cs.CLcs.LG2025-05被引 2

通过聚类映射稀疏注意力,提升PIM上大模型解码效率

Sparse Attention Remapping with Clustering for Efficient LLM Decoding on PIM

  • 按语义相似性聚类KV对,映射到连续内存区
  • 查询匹配聚类中心,实现选择性并行计算
  • 在1024上下文预算下降低67%能耗,适合长文本推理

基于Transformer的模型是现代机器学习的基础,但大型语言模型(LLM)在自回归解码过程中频繁访问内存,且键值(KV)缓存持续增长,给内存系统带来巨大压力,导致内存带宽瓶颈,尤其在长上下文场景下。存内计算(PIM)架构因其高内部带宽和近内存并行计算能力,被视为潜在解决方案。然而,现有PIM设计主要针对密集注意力优化,难以应对现代KV缓存稀疏技术带来的动态、不规则访问模式,导致负载不均,降低吞吐量与资源利用率。本文提出STARC,一种专为PIM架构高效解码而设计的稀疏注意力映射方案。该方法将KV对按语义相似性聚类,并将其映射至与PIM存储体结构对齐的连续内存区域。解码时,查询通过匹配预计算的聚类中心,以簇为粒度检索相关标记,实现选择性注意力与并行处理,避免频繁重聚类与数据移动开销。在HBM-PIM系统上的实验表明,相比常见逐令牌稀疏方法,STARC将注意力层延迟降低19%–31%,能耗减少19%–27%;在1024的KV缓存预算下,相较全缓存检索,延迟降低54%–74%,能耗减少45%–67%。同时,模型精度与当前最优稀疏注意力方法相当,证明其在支持硬件友好的长上下文大模型推理方面具有显著有效性。

原文摘要 · Abstract (English)

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.

PIM稀疏注意力大模型推理聚类映射

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。