QUOKA: 효율적인 LLM 프리필을 위한 쿼리 기반 키-값 선택

QUOKA: Query-Oriented KV Selection For Efficient LLM Prefill

본 논문에서는 QUOKA를 제안합니다. QUOKA는 훈련 과정 없이 하드웨어에 구애받지 않는 희소 어텐션 알고리즘으로, 청킹된 프리필 과정에서 트랜스포머 추론 속도를 향상시키는 데 사용됩니다. 기존 연구에서는 어텐션 연산에서 많은 쿼리가 더 작은 키 그룹에 집중하는 경향을 보이지만, 본 연구에서는 쿼리와 평균 쿼리 간의 코사인 유사도가 낮은 쿼리가 더 많은 키와 강하게 상호 작용하며, 최종 어텐션 로짓에 가장 큰 영향을 미친다는 것을 확인했습니다. 이러한 관찰을 바탕으로, QUOKA는 코사인 유사도가 낮은 쿼리를 우선적으로 처리하여, 프리필 단계에서 전체 어텐션의 동작을 근사화합니다. QUOKA는 (1) 대표적인 쿼리 집합을 먼저 유지하고, (2) 그런 다음 해당 쿼리와 가장 잘 정렬된 키를 부분적으로 선택하여 어텐션 속도를 향상시킵니다. Needle-In-A-Haystack, LongBench, RULER, 그리고 Math500 데이터셋에 대한 실험 결과, QUOKA는 토큰 생성 시간을 3배 단축하고, Nvidia GPU에서 어텐션 속도를 최대 5배, Intel Xeon CPU에서는 최대 7배까지 향상시키면서도, 어텐션 평가당 사용되는 키-값 쌍의 수를 88% 줄여 거의 동일한 정확도를 달성했습니다.

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.