# Efficient Beam Search for Large Language Models Using Trie-Based Decoding

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

## Key facts

- Cited by 10 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 19 of its 28 references that are in the map.
- First posted on arXiv 31 Jan 2025.

## Placed in

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

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks](https://mlsanity.com/paper/2005.11401/) (2020)
- [Evaluating Large Language Models Trained on Code](https://mlsanity.com/paper/2107.03374/) (2021)
- [Measuring Mathematical Problem Solving With the MATH Dataset](https://mlsanity.com/paper/2103.03874/) (2021)
- [Neural Machine Translation by Jointly Learning to Align and Translate](https://mlsanity.com/paper/1409.0473/) (2014)
- [Let's Verify Step by Step](https://mlsanity.com/paper/2305.20050/) (2023)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning](https://mlsanity.com/paper/2307.08691/) (2023)
- [Generating Long Sequences with Sparse Transformers](https://mlsanity.com/paper/1904.10509/) (2019)
- [Google's Neural Machine Translation System: Bridging the Gap between Human and Machine Translation](https://mlsanity.com/paper/1609.08144/) (2016)
- [Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads](https://mlsanity.com/paper/2401.10774/) (2024)
- [Abstractive Text Summarization Using Sequence-to-Sequence RNNs and Beyond](https://mlsanity.com/paper/1602.06023/) (2016)
- [SpecInfer: Accelerating Generative Large Language Model Serving with Tree-based Speculative Inference and Verification](https://mlsanity.com/paper/2305.09781/) (2023)
- [Accelerating LLM Inference with Staged Speculative Decoding](https://mlsanity.com/paper/2308.04623/) (2023)
- [Neural Path Hunter: Reducing Hallucination in Dialogue Systems via Path Grounding](https://mlsanity.com/paper/2104.08455/) (2021)
- [GPT4Rec: A Generative Framework for Personalized Recommendation and User Interests Interpretation](https://mlsanity.com/paper/2304.03879/) (2023)
- [Comparison of Diverse Decoding Methods from Conditional Language Models](https://mlsanity.com/paper/1906.06362/) (2019)
- [How Decoding Strategies Affect the Verifiability of Generated Text](https://mlsanity.com/paper/1911.03587/) (2019)
- [Towards Reliable Medical Question Answering: Techniques and Challenges in Mitigating Hallucinations in Language Models](https://mlsanity.com/paper/2408.13808/) (2024)

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

- [xLLM Technical Report](https://mlsanity.com/paper/2510.14686/) (2025)
- [Saffron-1: Safety Inference Scaling](https://mlsanity.com/paper/2506.06444/) (2025)
- [Efficient Test-Time Inference via Deterministic Exploration of Truncated Decoding Trees](https://mlsanity.com/paper/2604.20500/) (2026)
- [CoVeR: Conformal Calibration for Versatile and Reliable Autoregressive Next-Token Prediction](https://mlsanity.com/paper/2509.04733/) (2025)
- [Democratizing LLM Efficiency: From Hyperscale Optimizations to Universal Deployability](https://mlsanity.com/paper/2511.20662/) (2025)
- [PRISM: Purified Representation and Integrated Semantic Modeling for Generative Sequential Recommendation](https://mlsanity.com/paper/2601.16556/) (2026)
- [Suffix-Constrained Greedy Search Algorithms for Causal Language Models](https://mlsanity.com/paper/2603.01243/) (2026)
- [DocQAC: Adaptive Trie-Guided Decoding for Effective In-Document Query Auto-Completion](https://mlsanity.com/paper/2604.18257/) (2026)
- [Regulating Branch Parallelism in LLM Serving](https://mlsanity.com/paper/2605.06914/) (2026)
- [Trie Automata for Constrained Decoding over Large Finite Sets](https://mlsanity.com/paper/2608.12574/) (2026)

## Abstract (arXiv)

This work presents a novel trie (prefix-tree)-based parallel decoding method that addresses the memory inefficiency of batch-based beam search. By sharing a single KV cache across beams with common prefixes, our approach dramatically reduces memory usage and enables efficient decoding. We evaluated our method across three attention architectures, Multi-Head Attention (Phi-3.5-mini-instruct), Grouped Query Attention (Llama-3.1-8B-Instruct), and Sliding Window Attention (Mistral-Small-24B-Instruct-2501), using CNN/DailyMail for abstractive summarization and HumanEval for code generation. Our experiments demonstrate substantial memory savings (4--8×) and up to 2.4× faster decoding, without compromising generation quality. These results highlight our method's suitability for memory-constrained environments and large-scale deployments.