Quest: Query-Aware Sparsity for Efficient Long-Context LLM Inference

As the demand for long-context large language models (LLMs) increases, models with context windows of up to 128K or 1M tokens are becoming increasingly prevalent. However, long-context LLM inference is challenging since the inference speed decreases significantly as the sequence length grows. This slowdown is primarily caused by loading a large KV cache during self-attention. Previous works have shown that a small portion of critical tokens will dominate the attention outcomes. However, we observe the criticality of a token highly depends on the query. To this end, we propose Quest, a query-aware KV cache selection algorithm. Quest keeps track of the minimal and maximal Key values in KV cache pages and estimates the criticality of a given page using Query vectors. By only loading the Top-K critical KV cache pages for attention, Quest significantly speeds up self-attention without sacrificing accuracy. We show that Quest can achieve up to 2.23x self-attention speedup, which reduces inference latency by 7.03x while performing well on tasks with long dependencies with negligible accuracy loss. Code is available at http://github.com/mit-han-lab/Quest .

한 줄 요약

Quest는 쿼리에 따라 KV 캐시의 핵심 토큰을 선택적으로 처리하여, 긴 문맥의 LLM 추론 속도를 7.03× 가속화하는 알고리즘이다.

핵심 기여도

**Query-aware KV 캐시 선택 알고리즘** 제안: KV 페이지의 최소/최대 키 값과 쿼리 벡터를 기반으로 중요도를 추정.
**Top-K 기반 선택적 어텐션 수행**: 전체 KV 캐시 대신 최상위 K개 페이지만 로드하여 추론 속도 향상.
**7.03× self-attention 속도 향상**, **2.23× end-to-end 추론 가속화** 달성.
**LongBench, passkey retrieval 등 다양한 긴 문맥 작업에서 정확도 손실 없이 성능 유지**.

핵심 아이디어

기존 연구는 어텐션 결과에 영향을 주는 토큰이 전체 중 일부뿐임을 밝혔으나, 이 연구는 **쿼리에 따라 토큰의 중요도가 달라진다는 점**을 새로운 통찰로 제시한다. 이에 따라, Quest는 **KV 캐시 페이지 단위로 최소/최대 키 값을 메타데이터로 저장**하고, 이를 쿼리 벡터와 결합해 페이지의 중요도를 실시간으로 추정한다. 이는 **전체 KV 캐시를 로드하지 않고도 핵심 정보만 선택적으로 처리**할 수 있게 하며, 특히 **긴 문맥의 어텐션 연산에서 메모리 이동을 최소화**한다.

기술적 접근법

**KV 캐시 페이지 단위 관리**: 각 페이지의 Key 벡터의 최소/최대 값 저장.
**Query 벡터와 메타데이터 기반 중요도 추정**: 쿼리에 따라 페이지의 중요도를 실시간 계산.
**Top-K 선택적 어텐션 수행**: K는 128, 256 등 사전 설정 값.
**Flash-Attention 활용**: LongBench, passkey retrieval 등 다양한 작업에서 성능 평가.
**4-bit 가중치 양자화 적용**: end-to-end 추론 속도 향상에 기여.

주요 결과

**Self-attention 속도 7.03× 향상** (FlashInfer 대비).
**End-to-end 추론 속도 2.23× 향상** (4-bit 양자화 기반).
**LongBench 데이터셋에서 KV 캐시 스파스비 1/5~1/10 수준에서도 정확도 손실 없이 성능 유지**.
**10k, 100k 토큰 길이의 passkey retrieval 작업에서 100% 정확도 달성** (기존 H2O, TOVA, StreamingLLM 대비).

의의 및 한계

Quest는 **쿼리에 따른 토큰 중요도를 동적으로 반영**함으로써, 기존 KV 캐시 제거 기반 접근법의 한계를 극복한다. 특히, **긴 문맥 추론에서 정확도 유지와 추론 속도 향상을 동시에 달성**하며, 다양한 작업에서 **정확도 손실 없이 스파스비를 최대 1/10까지 적용**할 수 있다. 그러나, **모든 쿼리에서 동일한 Top-K 값이 최적일 수 없으며**, **메타데이터 기반 추정이 일부 정보를 누락할 가능성**이 있다. 또한, **모델 아키텍처에 따라 메타데이터 추출 방식이 달라질 수 있다**는 점도 한계로 지적된다.

실용적 활용

Quest는 **긴 문서 처리, 대화형 추론, 멀티턴 대화 시스템** 등에서 유용하게 활용될 수 있다. 특히, **128K~1M 토큰 길이의 입력을 처리하는 대규모 LLM 서빙 환경**에서 **추론 지연을 최소화하면서도 정확도를 유지**할 수 있어, **실시간 응답이 필요한 산업 분야**에 적합하다.