# PromptDistill: Query-based Selective Token Retention in Intermediate Layers for Efficient Large Language Model Inference

arXiv: https://arxiv.org/abs/2503.23274 · HTML: https://mlsanity.com/paper/2503.23274/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Placed in KV cache optimization, one of 12 directions in Efficient LLM inference and serving.
- Builds on 25 of its 29 references that are in the map.
- First posted on arXiv 30 Mar 2025.

## Placed in

- Efficient LLM inference and serving / [KV cache optimization](https://mlsanity.com/inference/kv-cache-optimization/)

## Builds on (most cited first, 25 of 25)

- [Efficient Streaming Language Models with Attention Sinks](https://mlsanity.com/paper/2309.17453/) (2023)
- [LongBench: A Bilingual, Multitask Benchmark for Long Context Understanding](https://mlsanity.com/paper/2308.14508/) (2023)
- [SnapKV: LLM Knows What You are Looking for Before Generation](https://mlsanity.com/paper/2404.14469/) (2024)
- [KIVI: A Tuning-Free Asymmetric 2bit Quantization for KV Cache](https://mlsanity.com/paper/2402.02750/) (2024)
- [Model Tells You What to Discard: Adaptive KV Cache Compression for LLMs](https://mlsanity.com/paper/2310.01801/) (2023)
- [Scissorhands: Exploiting the Persistence of Importance Hypothesis for LLM KV Cache Compression at Test Time](https://mlsanity.com/paper/2305.17118/) (2023)
- [MInference 1.0: Accelerating Pre-filling for Long-Context LLMs via Dynamic Sparse Attention](https://mlsanity.com/paper/2407.02490/) (2024)
- [∞Bench: Extending Long Context Evaluation Beyond 100K Tokens](https://mlsanity.com/paper/2402.13718/) (2024)
- [RECOMP: Improving Retrieval-Augmented LMs with Compression and Selective Augmentation](https://mlsanity.com/paper/2310.04408/) (2023)
- [Deja Vu: Contextual Sparsity for Efficient LLMs at Inference Time](https://mlsanity.com/paper/2310.17157/) (2023)
- [Compressing Context to Enhance Inference Efficiency of Large Language Models](https://mlsanity.com/paper/2310.06201/) (2023)
- [How to Train Long-Context Language Models (Effectively)](https://mlsanity.com/paper/2410.02660/) (2024)
- [The Bottom-up Evolution of Representations in the Transformer: A Study with Machine Translation and Language Modeling Objectives](https://mlsanity.com/paper/1909.01380/) (2019)
- [RetrievalAttention: Accelerating Long-Context LLM Inference via Vector Retrieval](https://mlsanity.com/paper/2409.10516/) (2024)
- [Beyond the Limits: A Survey of Techniques to Extend the Context Length in Large Language Models](https://mlsanity.com/paper/2402.02244/) (2024)
- [ThinK: Thinner Key Cache by Query-Driven Pruning](https://mlsanity.com/paper/2407.21018/) (2024)
- [Discovering the Gems in Early Layers: Accelerating Long-Context LLMs with 1000x Input Token Reduction](https://mlsanity.com/paper/2409.17422/) (2024)
- [Needle in the Haystack for Memory Based Large Language Models](https://mlsanity.com/paper/2407.01437/) (2024)
- [Selective Attention: Enhancing Transformer through Principled Context Control](https://mlsanity.com/paper/2411.12892/) (2024)
- [RAPID: Long-Context Inference with Retrieval-Augmented Speculative Decoding](https://mlsanity.com/paper/2502.20330/) (2025)
- [RefreshKV: Updating Small KV Cache During Long-form Generation](https://mlsanity.com/paper/2411.05787/) (2024)
- [A Little Goes a Long Way: Efficient Long Context Training and Inference with Partial Contexts](https://mlsanity.com/paper/2410.01485/) (2024)
- [Attention Reveals More Than Tokens: Training-Free Long-Context Reasoning with Attention-guided Retrieval](https://mlsanity.com/paper/2503.09819/) (2025)
- [QuickLLaMA: Query-aware Inference Acceleration for Large Language Models](https://mlsanity.com/paper/2406.07528/) (2024)
- [MCSD: An Efficient Language Model with Diverse Fusion](https://mlsanity.com/paper/2406.12230/) (2024)

## Built on by (most cited first, 1 of 1)

- [Adaptive Layer Selection for Layer-Wise Token Pruning in LLM Inference](https://mlsanity.com/paper/2601.07667/) (2026)

## Abstract (arXiv)

As large language models (LLMs) tackle increasingly complex tasks and longer documents, their computational and memory costs during inference become a major bottleneck. To address this, we propose PromptDistill, a novel, training-free method that improves inference efficiency while preserving generation quality. PromptDistill identifies and retains the most informative tokens by leveraging attention interactions in early layers, preserving their hidden states while reducing the computational burden in later layers. This allows the model to focus on essential contextual information without fully processing all tokens. Unlike previous methods such as H2O and SnapKV, which perform compression only after processing the entire input, or GemFilter, which selects a fixed portion of the initial prompt without considering contextual dependencies, PromptDistill dynamically allocates computational resources to the most relevant tokens while maintaining a global awareness of the input. Experiments using our method and baseline approaches with base models such as LLaMA 3.1 8B Instruct, Phi 3.5 Mini Instruct, and Qwen2 7B Instruct on benchmarks including LongBench, InfBench, and Needle in a Haystack demonstrate that PromptDistill significantly improves efficiency while having minimal impact on output quality compared to the original models. With a single-stage selection strategy, PromptDistill effectively balances performance and efficiency, outperforming prior methods like GemFilter, H2O, and SnapKV due to its superior ability to retain essential information. Specifically, compared to GemFilter, PromptDistill achieves an overall 1% to 5% performance improvement while also offering better time efficiency. Additionally, we explore multi-stage selection, which further improves efficiency while maintaining strong generation performance.