Trie-Based Parallel Decoding

• The paper introduces trie-based decoding, achieving up to 92% memory reduction by sharing KV caches across beams with common prefixes.
• The method builds a prefix tree to dynamically manage cache sharing and garbage collection, thereby optimizing beam search in Transformer models.
• Empirical evaluations show that the approach maintains throughput and output quality while substantially reducing memory demands.

Trie-based parallel decoding is a method for efficient beam search in Transformer-based sequence generation that addresses the high memory footprint of conventional batch-based approaches by exploiting prefix sharing among candidate sequences. The central innovation is to represent all active beam hypotheses in a prefix tree (trie), allowing key-value (KV) caches for attention to be shared across beams with common prefixes. This results in substantial reductions in GPU memory usage during sequence decoding, especially in large-scale LLM deployments or memory-constrained environments, without significant compromise in throughput or output quality (Chan et al., 31 Jan 2025).

1. Formal Structure and Core Principles

Let $B$ denote the beam size, $L$ the generation length, and $D$ the size (in floating-point elements) of the KV cache per token. The trie structure $\mathcal{T} = (N, E, \ell)$ comprises a set of nodes $N$, directed edges $E \subseteq N\times N$ labeled by tokens via $\ell: E \rightarrow \mathcal{V}$, where $\mathcal{V}$ is the vocabulary. The root node encodes the initial prompt. A directed path in the trie corresponds to a valid partial hypothesis (sequence of tokens), associating each beam’s current prefix with a unique trie node. Formally, at step $i$, a bijection $\phi_i:\{1,\ldots,B\} \to L_i \subseteq N$ maps each beam to its current trie node (which spells out the beam's prefix).

This data structure allows the sharing of decoder state: any beams that traverse the same path up to a given prefix also share the corresponding KV cache, as their forward computations up to that token sequence are identical.

2. Algorithmic Workflow

The decoding process (as formalized in the referenced paper) proceeds as follows:

1. Trie and Cache Initialization: The trie is initialized with the prompt as the root node. A mapping from trie nodes to KV caches is maintained on the GPU.
2. Gather Distinct Prefixes: At each step, collect active leaf nodes (distinct prefixes among all beams), denoted $L$0.
3. Batch Forward Pass: Construct a batch of KV caches associated with each distinct prefix and form an attention mask that constraints each prefix to attend within its branch.
4. Compute Next Token Distributions: For each distinct prefix, the model predicts logits for the next token. The log-probabilities are used to enumerate all expanded candidates.
5. Beam Candidate Selection: For each node, select the top-$L$1 token expansions, then globally prune to the top $L$2 hypotheses across all nodes.
6. Trie and Cache Update: New trie nodes and KV-caches are created only for surviving expansions. Cache storage is thus limited to nodes on active paths.
7. Garbage Collection: At prescribed intervals $L$3, subtrees not reachable from surviving leaves are pruned, and their associated KV caches are deleted to free memory.

Core data structures and routines include dictionaries for trie nodes to caches, functions for constructing attention masks respecting the trie topology, and pruning mechanisms for memory management.

3. Memory Complexity and Sharing Efficiency

For standard batch-based beam search, each of the $L$4 beams maintains its own full KV cache, yielding total memory consumption: $L$5

In the trie-based approach, only unique prefixes are cached: $L$6 where $L$7 is the average number of distinct active prefixes over decoding. Empirical results confirm $L$8 for typical tasks and beams, as many beams share dominating prefixes (especially at early or intermediate stages of generation).

The memory reduction ratio is

$L$9

For instance, if on average only 3 distinct prefixes remain out of 15 beams, memory usage drops by 80%.

4. Computational Complexity and Parallelization

Classic sequential beam search issues one model call per beam per step, for total computational cost: $D$0 where $D$1 is the cost of a single-sequence forward pass.

Batch-based beam search compacts all $D$2 beams into a batch, reducing model calls but retaining high memory usage: $D$3

Trie-based parallel decoding maintains the same parallelism, but the effective batch size is only the number of unique prefixes $D$4: $D$5 Typically, $D$6, so computational efficiency is at least as high as batch decoding, occasionally better due to reduced compute on shared prefixes. Practical results indicate near-ideal speedups proportional to $D$7 (Chan et al., 31 Jan 2025).

5. Empirical Evaluation

Empirical results (reproducing key findings from Table 1 of the original work) show that trie-based decoding achieves drastic GPU memory savings per token, with minimal impact on throughput and output quality.

Dataset	Beam size	Method	Mem/Tok (MB)	Tok/Sec
CNN/DM	3	Batch	4.06	7.50
CNN/DM	3	Trie	1.51	8.01
CNN/DM	9	Batch	11.99	5.25
CNN/DM	9	Trie	1.51	6.73
HumanEval	3	Batch	3.40	7.91
HumanEval	3	Trie	1.51	7.85
HumanEval	15	Batch	16.16	6.13
HumanEval	15	Trie	3.31	6.21

• Peak memory reduction of up to 92% per token.
• Throughput (Tok/Sec) remains within 5–10% of conventional batch-based decoding.
• Output quality (ROUGE-L for CNN/DM and pass@1 for HumanEval) is statistically indistinguishable ($D$8).

Figures 4–6 in the reference demonstrate that for increased output lengths, the memory gap between batch and trie approaches grows almost linearly for batch decoding, but sublinearly for the trie-based approach

(Chan et al., 31 Jan 2025).

6. Implementation and Practical Considerations

Integration with commonly used frameworks such as HuggingFace Transformers or Fairseq requires modification of the standard beam search loop to maintain the trie, per-node KV caches, and trie-aware attention masks. Caches can be stored in Python dictionaries keyed by trie node identifiers; batched concatenation and attention masking can leverage low-level tensor operations (e.g., torch.cat, index_select).

Special considerations include:

• Attention Masking: For deeply branching tries, attention masks may become prohibitively large ($D$9). Sharding or sparse masked attention (FlashAttention-style) may mitigate overhead.
• Dynamic Batching: In workloads with heterogeneous prompts or beam sizes, multiple tries may be maintained and interleaved for optimal GPU utilization.
• Garbage Collection: The interval parameter $\mathcal{T} = (N, E, \ell)$0 balances memory fragmentation against compute overhead; practical values are in the range $\mathcal{T} = (N, E, \ell)$1.
• Positional Embeddings: Correctly preserving positional encodings requires carrying original beam position IDs along the trie path, as detailed in Section 3.3 of the reference.

7. Context and Broader Implications

Trie-based parallel decoding demonstrates that the dominant source of memory inefficiency in Transformer beam search lies in the redundant storage of model state for beams with shared prefixes. By explicitly sharing state along trie branches, the approach delivers scalable decoding suitable for large beam sizes and long generations. This suggests potential for broader application in any autoregressive sequence model where prefix sharing occurs, particularly in scenarios with constrained GPU memory or high-throughput batch serving requirements.

Reference: (Chan et al., 31 Jan 2025) "Efficient Beam Search for LLMs Using Trie-Based Decoding." Chan et al., IJCAI 2025.