PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Jiaming Tang, Yilong Zhao, Kan Zhu, Guangxuan Xiao, Baris Kasikci, Song Han

arXiv:2406.10774v2cs.CLcs.LG

TL;DR

Long-context decoding is slowed by repeatedly loading large KV caches, despite evidence that only a small set of tokens often dominates attention. Quest estimates page criticality from the current query and Min/Max Key metadata, then attends only to Top-K pages. It reports up to 7.03× lower self-attention latency and 2.23× end-to-end inference speedup, with accuracy preserved under evaluated settings.

  • Problem

    Long-context decoding requires repeatedly reading large KV caches, while critical tokens vary with the current query.

  • Method

    Quest uses per-page minimum and maximum Key metadata with the current Query to rank KV-cache pages and loads only Top-K pages for attention.

  • Results

    Quest achieves up to 7.03× self-attention latency reduction and 2.23× end-to-end inference speedup, while performing well across evaluated long-context tasks with negligible accuracy loss.

  • Takeaways & Limitations

    Query-aware KV-cache selection can reduce memory movement and accelerate long-context inference while maintaining the stated accuracy targets.

  • Takeaways & Limitations

    Baseline efficiency comparisons are qualitative because the baselines lacked kernel implementations, and Quest is applied only on later layers in one stated setting.

Abstract

from arXiv · show

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 .

1. Introduction

Long-context inference is slowed by repeatedly loading large KV caches, motivating Quest’s query-aware selection of critical pages to reduce memory movement while preserving accuracy. Evaluations report substantial self-attention and end-to-end latency improvements.

  • Motivation: Long-context decoding is costly because generating each token requires reading the entire KV cache.For Llama 7B at 32K context, the KV cache occupies 16GB and reading it requires at least 11 ms, contributing more than 50% of inference latency.
  • Motivation: A small portion of KV-cache tokens can dominate generation accuracy, making critical-token identification a potential route to lower latency.The motivation is to load only critical tokens while maintaining accuracy.
  • Method: Quest dynamically estimates token criticality from the current query and selectively attends to chosen KV-cache tokens.The method uses query-aware criticality estimation rather than a fixed selection policy.
  • Method: Quest represents each KV-cache page with per-feature minimum and maximum Key values, ranks pages using the Query, and loads only Top-K pages for attention.This reduces memory movement from the full KV cache to page metadata and a constant number of selected pages.
  • Results: Up to 7.03× self-attention latency reduction and 2.23× end-to-end latency improvement are reported for Quest.The evaluation includes long-context tasks and compares against FlashInfer under the stated configurations.
  • Contribution: Quest is presented as an efficient and accurate KV-cache acceleration algorithm exploiting query-aware sparsity through dedicated operator designs and implementations.The contribution is framed as both an efficiency and accuracy improvement for KV-cache processing.

2. Related Work

Long-context models have expanded context windows substantially, increasing the importance of efficient inference. Prior KV-cache methods compress or select tokens using policies based on historical attention information.

  • Query dependence: Figure 2 shows that token criticality changes with the query: “B” receives low attention for “D” but high attention for the later query “is”.Rows represent attention scores over previous tokens queried by tokens on the left.
  • KV-cache compression: Prior methods address KV-cache overhead by retaining limited important tokens using historical attention scores or more refined token-selection policies.The related-work passage names H2O, FastGen, and TOVA as examples of KV-cache compression or selection approaches.

3. Methodlogy

Quest targets decode-time KV-cache overhead by estimating page criticality from the current query and page Key extrema, then applying attention only to selected pages. The design exploits substantial layer-wise sparsity while preserving accuracy under stated settings.

  • 3.1. Long-context Inference Is Costly: Over 86% of inference time is spent in decoding for 16K-token prompts with 512-token responses.Because decoding occurs for every generated token, decode-stage performance is crucial for overall latency.
  • 3.1. Long-context Inference Is Costly: A 32K-context Llama-7B KV cache can reach 16GB, and loading it accounts for 53% of decode-stage time.The KV cache must be loaded in every decode stage for self-attention.
  • 3.2. Query-aware sparsity: Less than 10% of tokens are needed for similar accuracy in most layers except the first two, indicating substantial attention sparsity.The layer-wise sparsity is measured while keeping PG19 perplexity increases below 0.01.
  • 3.3. Critical Tokens Depend on the Query: Quest’s two-stage workflow first estimates page criticality using Query and Min/Max Key metadata, then loads only Top-K pages for sparse self-attention.This summarizes the staged design shown in Figure 5.
  • 3.4. Dynamically Estimating Token Criticality: Quest estimates page criticality by taking per-channel upper bounds from Query products with minimum and maximum Key values.For channel i, it computes Ui = max(Q_i m_i, Q_i M_i), which bounds products with Keys in the page.
  • 3.4. Dynamically Estimating Token Criticality: Quest selects Top-K critical pages and performs normal self-attention only on those pages, defining the selected-page token count as the Token Budget.The query-aware sparsity mostly aligns with oracle sparsity while reducing memory movement.
  • 3.4. Dynamically Estimating Token Criticality: Quest is applied only on later layers because the first two layers have low sparsity; skipping them is orthogonal to the KV-cache selection algorithm.This is an explicit design setting used to better preserve model accuracy.
  • Memory movement: With 16 KV pairs per page, a 64K context, and 4K selected pages, Quest reduces memory load by 8×.The passage states that this reduction is universal across models and compatible with existing quantization mechanisms.

4.1. Setting

Quest is evaluated on long-context language modeling, retrieval, and LongBench settings using two models, multiple cache budgets, and established KV-cache baselines.

  • Evaluation setup: Quest is evaluated on PG19, passkey retrieval, and six LongBench datasets using LongChat-v1.5-7b-32k and Yarn-Llama-2-7b-128k.Baselines include H2O, TOVA, and StreamingLLM.
  • Efficiency accounting: The top-K operator takes 5–10 us and is excluded from the efficiency analysis.The operator’s memory loading and execution time are described as negligible.
  • Passkey retrieval: Table 1 covers 10K passkey retrieval with LongChat-7b-v1.5-32k and 100K retrieval with Yarn-Llama-2-7b-128k.Quest uses 64 and 1024 token budgets in the two tests, respectively.
  • Layer treatment: The evaluation applies Quest and baselines only after the first two model layers because those layers have low sparsity ratios.The first two layers retain a full cache to better preserve accuracy.

4.2. Accuracy Evaluation

Quest maintains accuracy across language modeling, long-dependency retrieval, and diverse long-context benchmarks while using sparse KV-cache budgets.

  • Language modeling: Quest’s PG19 perplexity closely matches the oracle baseline with a full KV cache at a 4096-token budget.PG19 contains 100 books averaging 70K tokens, and the test uses 32K-token contexts.
  • Long-dependency retrieval: Quest achieves perfect accuracy on both 10K and 100K passkey retrieval tests with minimal budgets.The query-aware method avoids discarding tokens that become critical for future queries.
  • Evaluation protocol: Long-dependency evaluation simulates decoding by feeding the question and instruction token by token.This setup tests whether methods retain information needed by later query tokens.
  • LongBench: Quest consistently outperforms all baselines across six LongBench datasets and varied KV-cache budgets.With a 1K-token budget, Quest reaches performance comparable to the full-cache model for most datasets.
  • LongBench: Quest achieves lossless performance with KV-cache sparsity of 1/6 on Qasper, HotpotQA, and MultifieldQA; 1/5 on GovReport and NarrativeQA; and 1/10 on TriviaQA.These results include the full cache used in the first two layers.

4.3. Efficiency evaluation

Quest’s efficiency comes from sparse, query-aware attention and reduced memory movement, producing substantial kernel, self-attention, end-to-end, and accuracy-constrained speedups.

  • Kernel evaluation: Quest’s criticality-estimation overhead approaches 1/Page Size of FlashInfer as sequence length grows.Estimation consumes one token per page, improving its relative performance at longer contexts.
  • Kernel evaluation: Top-K filtering adds only 5–10 us for sequence lengths below 128K.Criticality estimation reduces each token page to one criticality score, limiting memory movement.
  • Kernel evaluation: At a fixed token budget B, Quest’s approximate-attention latency remains constant as sequence length increases.It has latency similar to FlashInfer at sequence length B and is compatible with PageAttention.
  • Kernel evaluation: 7.03×: Quest reduces self-attention time versus FlashInfer at 32K sequence length with a 2048-token budget.The reduction is attributed to decreased memory movement.
  • End-to-end evaluation: 2.23×: Quest accelerates end-to-end inference with 4-bit quantized weights at 32K sequence length and a 2048-token budget.With FP16 weights, the corresponding speedup is 1.74×.
  • Comparison with baselines: Under lossless LongBench accuracy, Quest requires 5K tokens for NarrativeQA versus 14K for TOVA.The lower budget yields higher sparsity for Quest.
  • Comparison with baselines: Baseline efficiency is estimated from FlashInfer latency without including other runtime overheads, while Quest includes all operators.The baselines lack kernel implementations of their proposed methods.
  • Comparison with baselines: Quest improves inference speed by 3.82× on GovReport and 4.54× on TriviaQA under comparable accuracy.The comparison uses lossless accuracy targets across six LongBench tasks.

5. Conclusion

Quest is an efficient KV cache selection algorithm that exploits query-aware sparsity to reduce self-attention latency while maintaining accuracy. Its evaluations report substantial self-attention and end-to-end latency reductions under long-context settings.

  • Quest dynamically estimates KV cache token criticality from per-page metadata and the current query, then attends only to critical tokens.This reduces memory movement while targeting high sparsity with negligible accuracy loss.
  • 7.03× self-attention speedup contributes to a 2.23× end-to-end latency reduction in the decode phase.
  • 4.5× lower self-attention latency is achieved than prior baselines at the same accuracy target under long-context benchmarks.
Loading 2406.10774v2…