[論文レビュー] QUOKA: Query-Oriented KV Selection For Efficient LLM Prefill

Dalton Jones, Junyoung Park+4 more|arXiv (Cornell University)|Feb 9, 2026

Big Data and Digital Economy被引用数 0

ひとこと要約

QuoKA は訓練不要のハードウェア非依存スパースアテンション手法で、コサイン類似度を用いて代表的なクエリと KVs を選択し、ほぼ基準の精度を保ちながら潜在的なレイテンシを大幅に削減します。

ABSTRACT

We present QUOKA: Query-oriented KV selection for efficient attention, a training-free and hardware agnostic sparse attention algorithm for accelerating transformer inference under chunked prefill. While many queries focus on a smaller group of keys in the attention operator, we observe that queries with low cosine similarity with respect to the mean query interact more strongly with more keys and have the greatest contribution to final attention logits. By prioritizing these low cosine similarity queries, the behavior of full attention during the prefill stage can be closely approximated. QUOKA leverages this observation, accelerating attention by (1) first retaining a small set of representative queries and (2) then subselectin the keys most aligned with those queries. Through experiments on Needle-In-A-Haystack, LongBench, RULER, and Math500, we show that, while realizing a 3x reduction in time-to-first-token, 5x speedup in attention on Nvidia GPUs and up to nearly a 7x speedup on Intel Xeon CPUs, QUOKA achieves near-baseline accuracy, utilizing 88% fewer key-value pairs per attention evaluation.

研究の動機と目的

• トランスフォーマー推論のチャンク化プリフェイルにおけるプリフィル遅延の削減を動機づける。
• KV キャッシュ上で動作する軽量のハードウェア非依存スパースアテンション手法を提案する。
• 低コサイン類似性のクエリとそれに最も関連する KVs を選択することで、計算を削減しつつ精度を維持できることを示す。
• モデルとハードウェア（GPU/CPU）を横断した堅牢性と一般化を実証する。
• 長コンテキストおよび生成志向のベンチマークで経験的証拠を提供する。

提案手法

• 平均クエリからのコサイン類似度の不一致性に基づいて代表的なクエリの小さな集合を保持する。
• ドット積の代わりにクエリ–キーの関連性をスコアするコサイン類似度の代理指標を計算する。
• クエリと KV グループ全体のスコアを集約して、縮小された KV のサブセットを選択する。
• 縮小された KV セットを標準的なデンスアテンションカーネル（例：FlashAttention）に入力する。
• チャンク化プリフェイル内で動作し、アテンションの計算量を二次以下に抑える。
• カスタムカーネルを使わず、標準的な線形代数演算を用いてポータビリティを維持する。

(a) Chunked Prefill with QuoKA

実験結果

リサーチクエスチョン

• RQ1 チャンク化プリフェイル中に大幅な精度低下を招かずに KV アテンションをどれだけ削減できるか。
• RQ2 コサインベースのスコアリングと幾何学的に配慮したクエリ選択は、プリフェイルにおける生成志向または固定パターンのスパース性より優れているか。
• RQ3 長いコンテキストのベンチマークと異なるモデルファミリ間での QuoKA の精度とレイテンシのトレードオフはどうなるか。
• RQ4 QuoKA は GPU/CPU およびデコーダ初期 LLM アーキテクチャを横断してどれだけ一般化できるか。
• RQ5 B_CP および B_SA の予算が変化しても QuoKA は性能を維持できるか。

主な発見

• QuoKA はプリフェイル中に Nvidia GPU で最大5x のアテンションスピードアップを達成。
• QuoKA は長いプロンプトに対する TTFT（Time-To-First-Token）を約3x改善。
• Intel Xeon CPU では最大でほぼ7x、消費者向けGPUでも最大で5–6x のスピードアップを達成。
• QuoKA はアテンション評価ごとに 88% 少ないキー-バリュー対を使用しつつほぼ基準の精度を維持。
• Needle-In-A-Haystack、RULER、LongBench、Math500 のベンチマーク全体で、QuoKA は競合するスパースアテンション手法を上回る。
• 精度はスパース性とともに徐々に低下するため、効率性と精度のトレードオフを調整可能。

(b) KV selection with QuoKA