[논문 리뷰] QUOKA: Query-Oriented KV Selection For Efficient LLM Prefill

Dalton Jones, Junyoung Park외 4명|arXiv (Cornell University)|2026. 02. 09.

Big Data and Digital Economy인용 수 0

한 줄 요약

QuoKA는 학습 없이 작동하는 하드웨어에 구애받지 않는 희소 주의(attention) 방법으로, 청크 단위 LLM 프리필(prefill)을 위해 대표 쿼리와 KV를 코사인 유사도로 선택하며, 상당한 지연 감소를 근접한 기본 정확도와 함께 달성한다.

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.

연구 동기 및 목표

• Transformer 추론에서 청크드 프리필 하에서 프리필 지연 시간을 줄이는 것을 동기부여한다.
• KV 캐시에서 작동하는 경량의 하드웨어에 구애받지 않는 희소 주의 접근법을 제안한다.
• 저-코사인 유사도 쿼리와 그에 가장 관련성이 높은 KV를 선택하면 계산을 감소시키면서 정확도를 보전한다는 것을 보인다.
• 모델과 하드웨어(GPUs/CPUs)에 걸친 견고성과 일반화를 시연한다.
• 긴 컨텍스트 및 생성 지향 벤치마크에 대한 실증적 증거를 제공한다.

제안 방법

• 코사인 디시멀리티를 평균 쿼리 대비 유사도 기준으로 대표 쿼리의 소규모 세트를 유지한다.
• 도트 곱 대신 코사인 유사도 프록시를 사용하여 쿼리-키의 관련성을 점수화한다.
• 쿼리 및 KV 그룹 간 점수를 집계하여 축소된 KV 부분집합을 선택한다.
• 축소된 KV 세트를 표준 밀집 주의 커널(예: FlashAttention)에 공급한다.
• 청크드 프리필 내에서 동작하여 주의(attention)의 제곱근보다 낮은 차수의 복잡도를 달성한다.
• 커스텀 커널 없이 표준 선형대수 연산을 사용해 이식 가능성을 유지한다.

(a) Chunked Prefill with QuoKA

실험 결과

연구 질문

• RQ1 청크드 프리필 중 KV 주의(attention)를 얼마나 감소시켜도 큰 정확도 손실 없이 가능한가?
• RQ2 코사인 기반 점수화 및 기하학적 인식의 쿼리 선택이 프리필에서 생성 지향이나 고정 패턴 희소성보다 우수한가?
• RQ3 긴 컨텍스트 벤치마크 및 서로 다른 모델 계열에서 QuoKA의 정확도-지연 속도 트레이드오프는 어떠한가?
• RQ4 QuoKA가 GPU/CPU 간, 디코더만 LLM 아키텍처 간 얼마나 잘 일반화하는가?
• RQ5 B_CP 및 B_SA 예산이 달라질 때 QuoKA가 성능을 유지할 수 있는가?

주요 결과

• QuoKA는 프리필 중 Nvidia GPU에서 최대 5배의 주의(attention) 속도 향상을 달성한다.
• QuoKA는 긴 프롬프트에서 TTFT(First Token Time) 약 3배 개선을 제공한다.
• Intel Xeon CPU에서 최대 거의 7배의 속도 향상을 달성하고, 소비자용 GPU에서는 5-6배까지 증가한다.
• QuoKA는 주의 평가당 8.8%의 KV 쌍을 더 적게 사용하면서 근사 기본 정확도를 보전한다.
• 벤치마크 Needle-In-A-Haystack, RULER, LongBench, 그리고 Math500 전반에 걸쳐 QuoKA가 경쟁 희소 주의 방법을 능가한다.
• 정확도는 희소성 증가와 함께 점진적으로 저하되며 조정 가능한 효율-정확도 트레이드오프를 가능하게 한다.

(b) KV selection with QuoKA