Tree-Structured Speculative Decoding: How Multi-Candidate Trees and Tree Attention Accelerate LLM Serving

Large language model inference is fundamentally constrained by memory bandwidth during the auto-regressive decoding phase. Because each token generation step requires loading billions of model parameters from high-bandwidth memory (HBM) to compute units for a single token, standard auto-regressive generation operates at low arithmetic intensity.

Speculative decoding addresses this bottleneck by using a fast draft mechanism to propose candidate tokens, which are then verified in parallel by the larger target model in a single forward pass. However, early speculative decoding implementations relied on linear candidate chains. If the target model rejects the second token in a five-token chain, the remaining three tokens are discarded immediately, severely capping throughput gains.

Tree-structured speculative decoding overcomes this limitation by evaluating multi-branch token trees simultaneously. By combining multi-candidate drafting algorithms with a specialized 2D tree attention mask, modern inference engines verify branched candidate trajectories in a single target model execution, achieving 2.2x to 3.5x wall-clock speedups without altering output distributions.

The Linear Speculative Decoding Bottleneck

In standard speculative sampling (Leviathan et al., 2023; Chen et al., 2023), a lightweight draft model predicts a sequence of K consecutive tokens x_1, x_2, ..., x_K. The target model then runs a single forward pass over all K tokens to compute their true conditional probabilities P(x_i | x_<i).

The efficiency of this approach is governed by the token acceptance rate alpha. In a sequential chain, the probability of accepting all K tokens scales as the product of individual token probabilities:

P(accept all K tokens) = product(alpha_i)

If the average per-token acceptance rate is alpha = 0.70, the probability of accepting a 5-token chain is roughly 0.70^5 = 0.168.

When the target model rejects candidate token x_k, every subsequent token x_{k+1}, ..., x_K is immediately invalidated. Even if x_{k+1} would have been accepted under the target model's preferred alternative for x_k, the linear pipeline cannot explore alternative continuations.

Because modern GPU Tensor Cores have ample computational headroom during single-batch or low-batch decoding passes, verifying a batch of candidate tokens consumes virtually the same wall-clock time as verifying a single candidate chain. Linear speculative decoding leaves this hardware capacity underutilized.

Tree Representation of Speculative Candidates

To prevent early token rejections from collapsing entire verification steps, SpecInfer (Miao et al., 2023) introduced token tree verification. Instead of generating a single sequential path, the drafting process constructs a directed tree T = (V, E) of speculative candidate tokens:

• Each node v in V represents a speculated token.
• The root of the tree represents the current verified context.
• Each root-to-node path represents a distinct candidate prefix sequence.
• Branching factors at each depth reflect the draft model's probability distribution over top-k alternatives.

                  [Current Token: x_0]
                       /        \
             [Top-1: x_1a]    [Top-2: x_1b]
               /       \            |
        [x_2a]        [x_2b]      [x_2c]
          |             |
        [x_3a]        [x_3b]

In this structure, if the target model rejects branch [x_1a], it can simultaneously evaluate and potentially accept branch [x_1b, x_2c]. This structural redundancy lifts the effective acceptance length without requiring extra target model forward passes.

Drafting Architectures: Medusa vs. EAGLE

Generating speculative candidate trees efficiently requires draft architectures that do not introduce prohibitive latency overheads. Two dominant paradigms have emerged:

1. Multi-Head Parallel Decoding (Medusa)

Medusa (Cai et al., 2024) avoids running an independent autoregressive draft model entirely. Instead, it attaches K lightweight multi-layer perceptron (MLP) decoding heads to the final hidden state of the target foundation model.

Each Medusa head is trained to predict a specific future token offset:

• Head 1 predicts token t+1.
• Head 2 predicts token t+2.
• Head K predicts token t+K.

During decoding, each head generates top-s predictions. Medusa constructs a fixed candidate tree by taking Cartesian products of the highest-confidence predictions across heads, filtering out low-probability branches using pre-configured tree topologies. Because all Medusa heads run concurrently on the target model's final hidden representation, drafting adds minimal compute latency.

2. Feature-Level Autoregression (EAGLE and EAGLE-2)

While Medusa heads make conditionally independent predictions from the same hidden state, EAGLE (Li et al., 2024) introduces autoregressive modeling at the feature level.

EAGLE identifies that token-level sequences exhibit high entropy, whereas the second-to-last layer hidden features of a Transformer are smoother and more predictable. EAGLE deploys a single Transformer decoder layer (representing under 1% of target model parameter count). At each speculative step, it combines the target model's previous hidden state with draft token embeddings to predict the next feature vector autoregressively.

EAGLE-2 (Li et al., 2024) enhances this by implementing dynamic draft tree construction. Instead of relying on a static tree structure, EAGLE-2 evaluates the confidence scores of draft features in real time, dynamically allocating tree expansion to the most confident semantic paths for any given context.

The Tree Attention Mechanism

Verifying a branched tree of candidate tokens in a standard Transformer requires resolving cross-branch causal isolation. A token in branch A must not attend to tokens in branch B, but all tokens in both branches must attend to their common ancestors and the shared prompt history.

Tree attention solves this through custom 2D attention masking and non-linear positional encoding.

1. Tree Flattening

The tree nodes V = {v_1, v_2, ..., v_N} are flattened into a 1D sequence buffer of length N. The target model receives this flattened buffer alongside the existing key-value (KV) cache of the prefix.

2. 2D Causal Tree Mask

The attention mask M in {0, -inf}^(N x N) is defined by tree ancestry:

• If node v_j is an ancestor of node v_i in T (or if i = j), M_{i,j} = 0.
• Otherwise, M_{i,j} = -inf.

For any candidate node v_i, the attention mechanism computes attention scores exclusively over nodes that precede it along its specific branch path. Unrelated sibling branches receive -inf, suppressing cross-branch information leakage in the softmax operation.

3. Tree Positional IDs

Standard auto-regressive decoding assigns sequential positional IDs 1, 2, ..., N. In a candidate tree, tokens at the same depth in different branches share the same semantic distance from the root.

Tree attention assigns positional encodings based on tree depth:

PosID(v_i) = PosID(root) + depth(v_i)

This ensures that Rotary Position Embeddings (RoPE) or learned positional embeddings correctly reflect sequence positions relative to the prompt context.

Verification and Acceptance Algorithms

Once the target model executes a single forward pass over the flattened candidate tree, it produces output logit distributions for all N candidate nodes simultaneously.

Algorithm: Greedy Tree Verification
------------------------------------------------------------
Input: Candidate Tree T = (V, E), Target Logits L(v) for all v in V
Output: Accepted Token Path P, Next Starting Token x_next

1. Initialize Path P = []
2. Current Node curr = root
3. While curr has children in T:
4.    Target Prediction token_target = argmax(L(curr))
5.    Find child c of curr where candidate_token(c) == token_target
6.    If matching child c exists:
7.        Append candidate_token(c) to P
8.        curr = c
9.    Else:
10.       x_next = token_target
11.       Return P, x_next
12. x_next = argmax(L(curr))
13. Return P, x_next

Exact Sampling Equivalence

For stochastic decoding (temperature T > 0), tree speculative decoding applies generalized speculative rejection sampling (Miao et al., 2023; Cai et al., 2024).

At each tree node v, candidate token x proposed with draft probability q(x) is accepted with probability min(1, p(x) / q(x)). If rejected, a replacement token is sampled from the adjusted distribution (p(x) - q(x))^+ / sum(p(x) - q(x))^+. This preserves exact mathematical alignment with the target model's output distribution, ensuring that speculative acceleration introduces zero quality degradation.

KV Cache Rollback

During the forward verification pass, keys and values for all N tree candidates are computed and appended to the GPU memory buffer. Once the verification algorithm determines the accepted branch, unaccepted tree nodes are discarded from the KV cache using index pointer rollbacks or page table updates in PagedAttention systems.

Serving Economics and Framework Benchmarks

Tree-structured speculative decoding transforms the economics of LLM serving by maximizing the utility of every target model memory transfer:

• Draft Structure: Evolves from simple 1D linear chains of K tokens into 2D tree directed acyclic graphs (DAGs) containing N candidate nodes.
• Failure Recovery: Instead of discarding all subsequent tokens upon a single token mismatch, tree verification enables sibling branch fallback where alternative continuations remain viable.
• Mean Accepted Length (MAL): Increases average tokens accepted per verification step from 1.6 to 2.1 tokens in linear drafting up to 2.8 to 4.2 tokens in dynamic tree drafting.
• Wall-Clock Speedup: Delivers 2.2x to 3.5x wall-clock inference speedups over vanilla autoregressive decoding on single-batch and low-batch serving workloads.
• Output Distribution: Preserves exact mathematical equivalence with target model generation without distribution shift or approximation error.

Serving Framework Adoption

Leading inference frameworks have incorporated native tree-attention primitives:

• SGLang: Implements high-performance custom CUDA kernels for RadixTree-based speculative verification, supporting EAGLE and Medusa draft runners.
• vLLM: Supports speculative decoding models with custom tree attention masking, dynamically managing KV cache block allocation across tree candidates.
• TensorRT-LLM: Implements fused tree-attention kernels optimized for NVIDIA Hopper and Blackwell architectures, leveraging asynchronous tensor core execution to overlap draft head compute with memory staging.

By expanding linear speculation into parallel tree verification, tree-structured speculative decoding allows production LLM systems to break free from single-token memory bandwidth limitations without sacrificing model accuracy or output distribution integrity.

Sources

• SpecInfer: Accelerating Generative LLM Serving with Speculative Inference and Token Tree Verification (Miao et al., arXiv:2305.09781)
• Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads (Cai et al., arXiv:2401.10774)
• EAGLE: Speculative Sampling Requires Rethinking Feature Uncertainty (Li et al., arXiv:2401.15077)
• EAGLE-2: Faster Inference of Language Models with Dynamic Draft Trees (Li et al., arXiv:2406.16858)
• Sequoia: Supercharging Speculative Decoding with Speculative Trees (Chen et al., arXiv:2402.12374)
• Fast Inference from Transformers via Speculative Decoding (Leviathan et al., arXiv:2211.17192)