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

TL;DR

Quest introduces a query-aware method to selectively load only the most critical KV cache pages during long-context LLM inference, significantly accelerating self-attention and reducing latency with minimal accuracy loss.

Contribution

The paper presents a novel query-aware KV cache selection algorithm that improves inference efficiency for long-context LLMs by focusing on critical tokens based on query information.

Findings

Achieves up to 2.23x speedup in self-attention

Reduces inference latency by 7.03x

Maintains accuracy on long-dependency tasks

Abstract

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 .