Title: Efficient Low Rank Attention for Long-Context Inference in Large Language Models
Venue: NeurIPS 2025
Accepted: yes
Keywords: LLM

KV cache

Low rank decomposition

Long context inference
Abstract: As the length of input text increases, the key-value (KV) cache in LLMs imposes prohibitive GPU memory costs and limits long-context inference on resource constrained devices.
  Existing approaches, such as KV quantization and pruning, reduce memory usage but suffer from numerical precision loss or suboptimal retention of key-value pairs.
  In this work, Low Rank Query and Key attention (LRQK) is introduced, a two-stage framework that jointly decomposes full-precision query and key matrices into compact rank-\(r\) factors during the prefill stage, and then employs these low-dimensional projections to compute proxy attention scores in \(\mathcal{O}(lr)\) time at each decode step.
  By selecting only the top-\(k\) tokens and a small fixed set of recent tokens, LRQK employs a mixed GPU-CPU cache with a hit-and-miss mechanism where only missing full-precision KV pairs are transferred, thereby preserving exact attention outputs while reducing CPU-GPU data movement.
  Extensive experiments on the RULER and LongBench benchmarks with LLaMA-3-8B and Qwen2.5-7B demonstrate that LRQK matches or surpasses leading sparse-attention methods in long context settings, while delivering significant memory savings with minimal accuracy loss. Our code is available at \url{https://github.com/tenghuilee/LRQK}.


Peer reviews:
comments: SUMMARY:
The paper introduces LRQK, a low-rank attention method that reduces KV cache memory in LLMs during long-context inference. It approximates attention using low-rank query/key projections and retrieves only necessary full-precision KV pairs via a GPU–CPU cache. LRQK achieves strong memory savings with minimal accuracy loss, outperforming existing sparse attention methods on long-context benchmarks.

STRENGTHS AND WEAKNESSES:
### Strengths

- This paper is well-written and easy to follow. 
- The proposed LRQK framework introduces a low-rank approximation of query and key matrices during the prefill stage. This reduces the attention computation to a lightweight proxy form during decoding, enabling efficient token selection with significantly lower overhead.
- The approach is empirically validated on strong long-context benchmarks (RULER and LongBench) with language models like LLaMA-3-8B and Qwen2.5-7B.

### Weaknesses
- The observations regarding singular values and key-value decomposition are not novel; similar findings have already been reported in prior work such as Loki [1].

- The proposed method introduces a two-stage approximation involving proxy attention scores and hierarchical cache management, which adds notable complexity to the inference pipeline. However, the paper does not adequately address the implementation or integration overhead.

- Although the authors acknowledge in the Limitations section that the method requires careful hyperparameter tuning, no concrete strategies or guidelines are provided—this significantly limits its practical usability.

[1] Loki: Low-rank Keys for Efficient Sparse Attention

QUESTIONS:
Please refer to the Weakness.

LIMITATIONS:
Yes
TITLE: Official Review
CLARITY: 3
IMPACT: 2
ORIGINALITY: 2
DATE: 2025-06-24

comments: SUMMARY:
This work proposes Low Rank Query and Key attention (LRQK), a two-stage framework decomposing query and key matrices into low-rank factors during prefill and using them for proxy attention scores in decode, with a hierarchical GPU-CPU cache to reduce data movement. It balances memory efficiency, precision, and latency by selecting top-k and recent tokens, preserving exact attention outputs.

STRENGTHS AND WEAKNESSES:
Strengths
1. Joint Low-Rank Approximation: Unlike prior methods that focus solely on key matrix decomposition, LRQK jointly factorizes both query and key matrices into low-rank components. This approach reduces computational complexity per decode step while maintaining representation accuracy, avoiding expensive SVD operations.
2. Precision Preservation: LRQK uses low-rank proxies for attention score estimation but retains original query, key, and value vectors for final attention computation. This ensures mathematical fidelity and minimal accuracy loss compared to full-precision attention.
3. Hierarchical Cache Management: The hybrid GPU-CPU cache with a hit-and-miss mechanism selectively transfers only missing full-precision KV pairs, reducing data movement by ~60% on average. A recency buffer in the GPU cache further optimizes hit rates by leveraging attention’s locality bias.

Weakness:
1. Hyperparameter Sensitivity: The method requires careful tuning of rank and top-k token counts, which vary across model architectures and tasks. Optimal values demand extensive grid searches, increasing implementation overhead.
2. The paper lacks performance analysis. It would be better if the authors could compare their method with baselines in terms of throughput and TTFT.

QUESTIONS:
Mention in weakness

LIMITATIONS:
yes
TITLE: Official Review
CLARITY: 3
IMPACT: 3
ORIGINALITY: 3
DATE: 2025-07-02

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes Low Rank Query and Key attention (LRQK), a two-stage framework for efficient long-context inference in large language models. LRQK jointly factorizes the query and key matrices into low-rank bases during the prefill stage, and then uses these to compute proxy attention scores for top-k token selection during decoding. A hierarchical GPU–CPU caching mechanism with a hit/miss policy transfers only missing key-value pairs, preserving exact attention outputs while reducing memory and data movement costs. Extensive experiments on RULER and LongBench benchmarks with LLaMA-3 and Qwen2.5 models show LRQK achieves comparable or superior accuracy to strong sparse attention baselines while substantially lowering GPU memory consumption.\n\nSTRENGTHS AND WEAKNESSES:\nStrength\n\n1. The paper addresses a real and pressing bottleneck in LLM inference: the GPU memory consumption of KV caches for long contexts. Efficient memory strategies are highly relevant for widespread LLM deployment, especially on resource-constrained devices.\n\n\nWeakness：\n\n1. The empirical comparisons primarily involve sparse/dynamic attention and offloading methods, but the selection of baselines does not cover some of the most recent quantization methods (e.g., PALU [9], which is only referenced).\n\n[1] Efficient Low Rank Attention for Long-Context Inference in Large Language Models\n\n[2] Scalable Efficient Training of Large Language Models with Low-dimensional Projected Attention\n\n2. While sensitivity to rank and top-k is shown (Tables 3, 4), there is limited discussion of the overhead and tuning effort required, particularly when considering different LLM architectures\n\n3. Results are only shown for LLaMA-3-8B and Qwen2.5-7B (Tables 1-2), lacking validation on other popular architectures.\n\n4. Please elaborate further on the differences with other methods focusing on efficient attention for long-context inference.\n\n5. While the paper claims to balance trade-offs between memory efficiency, numerical precision, and computational latency, it lacks experimental validation of these specific trade-offs.\n\nQUESTIONS:\nList in the Weakness part\n\nLIMITATIONS:\n1.Although presented algorithms are clear, actual running times, cache sizes (absolute MB/GB), and detailed hardware utilization are not given beyond high-level statements about GPU/CPU. Code is not referenced as available, and implementation details are not provided.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes Low-Rank Query and Key attention (LRQK), a two-stage framework for reducing the KV-cache memory footprint during long-context inference in decoder-only large language models (LLMs). For a long prompt, the full-precision query (Q) and key (K) matrices are jointly factorised into rank-r bases by minimising a reconstruction loss (Eq. 3). At each generation step, low-rank projections estimate “proxy” attention scores in O(ℓr) time; only the top-k tokens plus a small recency buffer are kept on-GPU, while the rest of the cache resides on CPU. A hit/miss mechanism fetches full-precision KV rows for missed indices, guaranteeing exact attention on the selected subset. Experiments with LLaMA-3-8B-1M and Qwen-2.5-7B on RULER-128 K and LongBench show that LRQK matches or beats shadow-KV, InfiniGen, QUEST and Loki on most tasks, e.g., +2 pp on QA-SQuAD versus the 8B baseline while cutting CPU↔GPU transfers by ≈60 % on WikiText-2 . Hyper-parameter sweeps reveal favourable rank/top-k trade-offs (r = 32 and k = 1024 give near-lossless accuracy)\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n1. This paper provides a novel and innovate approach to solving the memory constraints in long context inference. Unlike many eviction or sparsity methods, the hit/miss fetch ensures identical outputs to full attention on the retained tokens—no approximation error. \n2. The paper is generally well written and each technique introduced comes with proper motivation. This paper is timely and looks at some key memory issues with kv-caching. \n3. Authors provide strong results. With LLaMA-3-8B-1M at 128 K tokens, LRQK matches or beats strong baselines (ShadowKV, Quest, InfiniGen, Loki) on many RULER tasks and improves Passage-Retrieval and LCC on LongBench (Table 1). I also found the memory-transfer savings results interesting like on WikiText-2 the hit/miss cache cuts CPU ↔ GPU KV transfers by ≈60 % across a broad grid of settings (mean miss-rate ≈0.40, Fig. 5) Moreover, authors have done extensive ablations on rank r, active-set size k, miss-rate, and cache policy across two architectures and multiple long-context benchmarks.\n4. Authors promised to open source the code which will be valuable to the community. \n\nWeaknesses:\n1. Experiments list the hardware (A100) but omit tokens/s throughput, wall-clock latency, and peak/resident GPU memory, making it hard to judge real deployment benefits\n2. All results use LLaMA-3-8B-1M and Qwen-2.5-7B; the method’s scalability to 30 B + or MoE LLMs, where KV caches are a bigger bottleneck, remains unverified. The authors themselves flag that rank r and active-set size k require dataset/model-specific tuning, and optimal values can vary across tasks. I am concerned about the ease of adoption of this technique for more general use cases. \n3. Equation 3 formulation is internally inconsistent. The hard constraints $ Q=A_QB_Q, K=A_KB_K $ in Eq. 3 cannot hold for low rank $r<d$ and are immediately relaxed in the Lagrangian; also, the interaction term omits $B_QB_K^T$. This is mostly cosmetic but weakens the mathematical presentation.\n\nQUESTIONS:\n1. Section 5 reports “≈ 60 % reduction in CPU↔GPU transfers” but gives no tokens-per-second or latency figures on the A100 node you mention (lines 210-216). Could you provide end-to-end throughput (tokens/s) and peak GPU memory for LRQK vs. ShadowKV, InfiniGen and vanilla attention at 128 K tokens?\n\n2. All experiments use 7-8 B parameters. Are there bandwidth or PCIe limits that prevent LRQK from scaling to 30 B+ or mixture-of-experts LLMs where KV memory dominates? Even a profiling statement would help assess significance.\n\n3. Since LRQK already uses low-rank factors, can the retained full-precision KV rows be post-quantised to 8-bit or 4-bit (à la KVQuant) without harming accuracy? Discussing complementarity with quantisation would broaden impact.\n\nLIMITATIONS:\nThe authors candidly notes that rank r and active-set size k require task-specific tuning and that optimal values differ across architectures. However, the societal-impact section is missing. Please add a brief discussion on potential misuse (e.g., enabling cheaper deployment of disinformation bots) and mitigation strategies.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-06"}