Sparse Attention Remapping with Clustering for Efficient LLM Decoding on PIM

要約

変圧器ベースのモデルは、最新の機械学習の基礎ですが、特に大規模な言語モデル（LLM）での自己回帰デコード中の実行は、頻繁なメモリアクセスとキー価値（KV）キャッシュの増加により、メモリシステムに大きな圧力をかけます。
これにより、特にコンテキストの長さが増加するにつれて、メモリ帯域幅にボトルネックが作成されます。
メモリ（PIM）の処理（PIM）アーキテクチャは有望なソリューションであり、高い内部帯域幅を提供し、メモリの近くで並列性を計算します。
ただし、現在のPIMデザインは主に密な注意のために最適化されており、最新のKVキャッシュスパーステクニックによって導入された動的で不規則なアクセスパターンとの闘いです。
その結果、彼らはワークロードの不均衡に苦しみ、スループットとリソースの利用を減らします。
この作業では、PIMアーキテクチャでの効率的なLLMデコードのために特別に調整された新しいスパース最適化データマッピングスキームであるSTARCを提案します。
STARCクラスターKVペアはセマンティックな類似性でペアを付け、PIMバンク構造に沿った隣接するメモリ領域にマッピングします。
デコード中、クエリは、事前計算された重心と一致することにより、クラスター粒度で関連するトークンを取得し、頻繁にリクラスタリングまたはデータ移動オーバーヘッドをせずに選択的注意と並列処理を可能にします。
HBM-PIMシステムでの実験は、一般的なトークンごとのスパースの方法と比較して、STARCが注意レイヤーの遅延を19％から31％、エネルギー消費を19％-27％減らすことを示しています。
1024年のKVキャッシュ予算では、完全なKVキャッシュ検索と比較して、最大54％ – 74％のレイテンシー削減と45％ – 67％のエネルギー削減を達成します。
一方、STARCは、最先端のまばらな注意方法に匹敵するモデルの精度を維持し、PIMアーキテクチャで効率的でハードウェアに優しい長いコンテキストLLM推論を可能にする際のその効果を示しています。

要約(オリジナル)

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.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google