Medusa, EAGLE-1/2/3, MTP, and tree attention

§6.2 left us with one expression and one disappointment:

$\alpha$ is the per-token acceptance probability, $\kappa$ the chain length, and $c$ the cost of one draft forward pass as a fraction of one target forward pass. The disappointment is $c$. Put Llama-3.2-1B in front of Llama-3-8B on an H100 and the numbers are, derived from the §0.4 bandwidth floor: the target streams 15.01 GB per step (4.48 ms), the draft streams 2.47 GB (0.74 ms), so $c = 0.74/4.48 = 0.165$. At $\kappa=5$ the denominator is $1 + 5(0.165) = 1.83$. Even at a generous $\alpha = 0.70$, where $E = (1-0.7^{6})/0.3 = 2.94$, the speedup is $2.94/1.83 = 1.61\times$. You threw away 45% of the acceptance you paid for, to the drafter's own weight traffic.

Now suppose the draft cost per step were $c = 0.05$ instead. Same $\alpha$, same $\kappa$: denominator $1.25$, speedup $2.35\times$. Nothing about the draft's quality changed. The entire 46% improvement came from making the guess cheaper.

§6.3 drew the $(\alpha, c)$ plane and pointed at its best corner: high acceptance, near-zero cost. A separate draft model sits in the wrong part of it because it must rediscover, from scratch and from tokens alone, context the target has already computed. This chapter is about the family of methods that refuse to rediscover anything — they bolt a head onto the target and read its hidden state directly — and about the mask that lets you verify a branching guess in a single target pass.

Run the target over a prefix. The final decoder layer emits a hidden state $h_t \in \mathbb{R}^{4096}$; the LM head turns it into a distribution over 128,256 tokens and you sample $x_{t+1}$. Then you throw $h_t$ away and start over.

But $h_t$ contains far more than one token's worth of information. It has just been produced by 32 layers of attention over the whole context; it knows the syntactic frame, the entity being described, the fact that the sentence is halfway through a list. Predicting $x_{t+2}$ from $h_t$ is a much easier problem than predicting $x_{t+2}$ from the token string alone, which is what a standalone 1B draft model is doing.

Two families exploit this, and the difference between them is the whole story. Medusa attaches $K$ independent heads to $h_t$; head $k$ predicts $x_{t+k}$ directly. EAGLE attaches one small autoregressive module that consumes $h_t$ and the embedding of the token it just drafted, producing a new feature vector that the next draft step consumes — autoregression at the feature level rather than the token level.

vLLM's Medusa proposer is eleven lines of substance, and every one of them is diagnostic.

    ) -> torch.Tensor:
        assert num_speculative_tokens == self.num_speculative_tokens
        # Generate blocks and compute logits
        blocks = self.model(target_hidden_states)
        logits = self.model.compute_logits(blocks)

        # Compute argmax for each Medusa head and stack into a single tensor
        # Shape: [batch_size, num_heads]
        draft_tokens = torch.stack([logit.argmax(dim=-1) for logit in logits], dim=1)

        return draft_tokens

Everything the drafter uses is target_hidden_states; the other two parameters of propose() (sampling_metadata and a slot_mappings the signature marks # unused) never reach the model. No input_ids, no positions, no attention metadata, no KV cache. There is no loop. All $K$ draft tokens appear from one call, and the model itself is a list comprehension:

    def forward(self, hidden_states: torch.Tensor) -> list[torch.Tensor]:
        return [block(hidden_states) for block in self.blocks]

Every block sees the same hidden_states. Head 3 has no way to know what head 1 and head 2 produced. Written out, head $k$ is trained to model

a marginal. The draft you actually submit is the tuple $(\arg\max q_1, \arg\max q_2, \dots)$, which is the coordinatewise mode of a product of marginals — not a mode, or even a likely sample, of the joint. When the continuation is genuinely uncertain, independent argmaxes produce incoherent strings: each token is individually the most likely thing to appear at that offset, and together they are a sentence nobody would write. The verification pass then rejects at the first incoherence, and the accepted length collapses toward 1 regardless of how good each individual head is.

The Medusa paper's answer was a tree: take the top-$s_k$ tokens from head $k$ and verify the Cartesian product $s_1 \times s_2 \times \cdots$ as a branching draft, so that at least one branch is coherent. That works — and it is where tree attention entered speculative decoding — but it is a coverage patch, not a fix. The cost of covering a joint distribution with a product of marginals grows multiplicatively in the depth. The field's response was to fix the model instead, and make the draft autoregressive.

What a Medusa head costs

The head body is trivial: ResidualBlock is num_hidden_layers applications of $x \mathrel{+}= \mathrm{SiLU}(Wx)$ with $W \in \mathbb{R}^{4096\times4096}$, so one layer is 16.78 M parameters. The vocabulary projection is not trivial, and the constructor says so:

        self.orig_vocab_size = config.vocab_size
        self.truncated_vocab_size = config.truncated_vocab_size

        if getattr(config, "original_lm_head", False):
            self.lm_head = ParallelLMHead(
                self.truncated_vocab_size,
                config.hidden_size,
                prefix=maybe_prefix(prefix, "lm_head"),
            )
            self.lm_heads = [self.lm_head for _ in range(self.config.num_heads)]
        else:
            self.lm_heads = nn.ModuleList(
                [
                    ParallelLMHead(
                        config.vocab_size,
                        config.hidden_size,
                        prefix=maybe_prefix(prefix, f"lm_heads.{i}"),
                    )
                    for i in range(self.config.num_heads)
                ]
            )

Derived. With four heads and separate full-vocabulary projections, each head carries $4096 \times 128{,}256 = 525.3$ M parameters of LM head against 16.8 M of body: 4 heads = 2.168 B parameters = 4.34 GB in bf16 = 1.29 ms of weight streaming, so $c = 0.289$. That is larger than the 1B draft's streamed-weight estimate. With original_lm_head, unique resident parameters fall to 592 M (1.19 GB); truncating to 32,000 tokens gives 198 M (0.40 GB). The corresponding 0.354 ms and 0.118 ms figures assume each shared weight is streamed only once across heads. Sharing a tensor does not itself guarantee one HBM read: separate projection calls, cache residency and batching determine traffic. These are optimistic cost estimates.

Note also that Medusa is a vLLM-only path at these SHAs. SGLang's builtin algorithm list is EAGLE, EAGLE3, NEXTN, STANDALONE, NGRAM, DFLASH, DSPARK (python/sglang/srt/server_args.py:L2101-L2105) and the string "medusa" does not appear anywhere under python/.

EAGLE keeps the "read the target's hidden state" idea and adds the missing edge. Its draft module is one Llama decoder layer preceded by a projection, and the projection is the interesting part:

        self.fc = ReplicatedLinear(
            input_size=self.config.hidden_size * 2,
            output_size=self.config.hidden_size,
            bias=False,
            params_dtype=vllm_config.model_config.dtype,
            quant_config=self.quant_config,
            prefix=maybe_prefix(prefix, "fc"),
            return_bias=False,
        )

    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.embed_tokens(input_ids)

    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        input_embeds = self.embed_tokens(input_ids)
        hidden_states = self.fc(torch.cat((input_embeds, hidden_states), dim=-1))
        residual = None
        for layer in self.layers:
            hidden_states, residual = layer(
                positions,
                hidden_states,
                residual,
            )
        hidden_states = hidden_states + residual
        return hidden_states, hidden_states

That is the entire idea in one line. The input to draft step $j$ is $\mathrm{fc}\big(\,[\,\mathrm{embed}(\tilde{x}_j)\;;\;f_{j-1}\,]\,\big)$ where $\tilde{x}_j$ is the token drafted at step $j-1$ and $f_{j-1}$ is the feature that step $j-1$ emitted ($f_0 = h_t$, the target's own hidden state). The layer runs with a real KV cache and real positions, so it also attends to every previous position. Its output is returned twice: once as the thing the LM head consumes, once as $f_j$ for the next step.

Two consequences.

One more detail worth reading, because it explains a load-error class you will hit: the first draft layer has its input LayerNorm deleted.

        # Skip the input_layernorm
        # https://github.com/SafeAILab/EAGLE/blob/35c78f6cdc19a73e05cf5c330b4c358dad970c6a/eagle/model/cnets.py#L427
        if disable_input_layernorm:
            del self.input_layernorm
            self.input_layernorm = nn.Identity()

The checkpoint was trained with this normalization placement. Adding another norm changes the learned function and can reduce acceptance even when shapes remain valid. The size and direction of that change require measurement; there is no general factor-of-two penalty.

What one EAGLE step costs

Derived, all Llama-3-8B shapes. One decoder layer is $4\!\cdot\!4096\!\cdot\!4096$-equivalent attention projections ($q$: 16.78 M, $k$: 4.19 M, $v$: 4.19 M, $o$: 16.78 M under GQA with $h_{kv}=8$) plus a 14,336-wide MLP ($3 \times 4096 \times 14336 = 176.2$ M), total 218.1 M. Add fc at $8192 \times 4096 = 33.6$ M. Add the LM head.

Two thirds of a full-vocabulary EAGLE draft's byte traffic is the LM head. The transformer is the cheap part. This is why every serious variant of this family ships a vocabulary reduction: Medusa's truncated_vocab_size/token_map, EAGLE-3's draft_vocab_size with a d2t index map back to target ids. It is also why the head cost is a fixed tax you pay $S$ times per iteration, once per draft forward.

vLLM's draft loop is a chain

vLLM's EagleProposer is a thirteen-line subclass of SpecDecodeBaseProposer — the whole file is 22 lines (vllm/v1/spec_decode/eagle.py:L10-L22) — whose only job is to pass pass_hidden_states_to_model=True. The loop lives in the base class:

        block_size = self.block_size
        assert block_size > 0, "block_size has not been initialized."
        for token_index in range(self.num_speculative_tokens - 1):
            # Update the inputs.
            # cast to int32 is crucial when eagle model is compiled.
            # tensor.argmax() returns int64 by default.
            input_ids = draft_token_ids_list[-1].int()

            draft_token_ids_list.append(draft_token_ids)

        # [batch_size, num_speculative_tokens]
        draft_token_ids = torch.stack(draft_token_ids_list, dim=1)

[batch_size, num_speculative_tokens]. One token per step per request: a chain, not a tree. As of a556f3f, vLLM's V1 speculative path does not implement tree drafting for EAGLE. The only tree_mask identifier anywhere in vllm/ is an optional parameter of the CPU MLA extend op cpu_mla_extend (vllm/_custom_ops.py:L3999-L4020, backed by csrc/cpu/sgl-kernels/extend.cpp); its one caller, the AMX MLA backend at vllm/v1/attention/backends/mla/amx_mla.py:L377-L397, passes None for it, so nothing in the tree ever builds one. Meanwhile the drafter carries a standing note:

        # FIXME: when using tree-based specdec, adjust number of forward-passes
        # according to the depth of the tree.

That is a genuine divergence, and it is not an oversight so much as a different bet. A chain verifies $\kappa+1$ positions per request; a tree of the same depth verifies several times as many. §6.2 showed that the verify pass's token count is what pushes a batched server over the roofline ridge, so trees pay off in the latency-critical, small-batch regime and hurt soonest in the throughput regime. vLLM's V1 spec path is built around a padded, chain-shaped drafter batch (disable_padded_drafter_batch, prepare_inputs_padded) which a tree would have to break. SGLang made the other choice.

EAGLE-2: the tree stops being a constant

EAGLE-1 verified a fixed tree — a shape chosen offline and reused for every step, every prompt. EAGLE-2's observation is that the draft model's own confidence is a usable estimate of how likely each branch is to survive verification, so the tree should be re-shaped every step. SGLang implements this directly. At each draft step it multiplies the surviving paths' cumulative scores by the new conditional probabilities and keeps the best topk paths:

@torch.compile(dynamic=True, disable=_is_npu or _is_xpu)
def _select_top_k_tokens_later(
    i: int,
    topk_p: torch.Tensor,
    topk_index: torch.Tensor,
    hidden_states: torch.Tensor,
    scores: torch.Tensor,
    topk: int,
):
    topk_sq = topk * topk

    expand_scores = scores.unsqueeze(2) * topk_p.view(-1, topk, topk)
    # (b, topk, 1) * (b, topk, topk) -> (b, topk, topk)

    topk_cs_p, topk_cs_index = fast_topk(
        expand_scores.flatten(start_dim=1), topk, dim=-1
    )  # (b, topk)

    topk_index = topk_index.view(-1, topk_sq)
    input_ids = torch.gather(topk_index, 1, topk_cs_index).flatten()

expand_scores is the cumulative path probability $\prod_{j\le i} q(\tilde{x}_j \mid \tilde{x}_{<j})$ under the draft's own distribution. Note what this is not: it is not a fixed per-level branching factor. The $\mathrm{topk}^2$ children of this level's $\mathrm{topk}$ survivors compete against each other on cumulative probability, so a single confident path can consume the whole budget while its uncertain siblings die at depth 1.

Then, once all steps have run, every candidate ever generated competes for the verify budget:

def organize_draft_results(
    score_list: List[torch.Tensor],
    token_list: List[torch.Tensor],
    parents_list: List[torch.Tensor],
    num_draft_token: int,
):
    # b, n, topk; n = 1 + (num_steps-1) * topk
    score_list = torch.cat(score_list, dim=1).flatten(1)
    # b, (topk + (num_steps-1) * topk)
    ss_token_list = torch.cat(token_list, dim=1)
    top_scores = torch.topk(score_list, num_draft_token - 1, dim=-1)
    top_scores_index = top_scores.indices
    top_scores_index = torch.sort(top_scores_index).values
    # ...
    draft_tokens = torch.gather(ss_token_list, index=top_scores_index, dim=1)

With SGLang's auto-chosen Llama parameters — (num_steps, topk, num_draft_tokens) = (5, 4, 8), from python/sglang/srt/arg_groups/speculative_hook.py:L821-L822; those defaults are filled in only when you pass none of the three (speculative_hook.py:L606-L619) — the candidate pool is $\mathrm{topk} + (\text{num\_steps}-1)\cdot\mathrm{topk}^2 = 4 + 4\cdot16 = 68$ tokens, of which the global top 7 are kept, plus the root, giving 8 verified nodes. The selection is global across depths, prioritizing joint draft confidence, not guaranteed target acceptance. It also needs ancestor closure: every selected node's parent must be selected or be the root. Products of probabilities are nonincreasing, but ties (including probability-one edges and floating-point underflow) invalidate a strict-ranking proof. An ancestor-first tie rule or explicit closure repair is required before building the mask.

EAGLE-3: two things the serving code shows, one it does not

EAGLE-3's changes are mostly training-time, but two of them leave unmistakable fingerprints in the model definition. First, the draft consumes several target layers concatenated, not just the last:

            target_hidden_size = getattr(
                self.config, "target_hidden_size", self.config.hidden_size
            )
            self.fc_input_size = target_hidden_size * self.num_aux_hidden_states

    def combine_hidden_states(
        self,
        hidden_states: torch.Tensor,
    ) -> torch.Tensor:
        if not self.model.use_aux_hidden_state:
            return hidden_states
        # combine multiple auxiliary hidden states returned by eagle3

        if self.model.norm_before_fc:
            hidden_states = self.model.input_norm(hidden_states)

        # `norm_before_fc` adds a single RMSNorm before the FC layer, whereas `fc_norm`
        # applies separate RMSNorms to each chunk of the hidden states.
        if self.model.fc_norm is not None:
            chunks = hidden_states.chunk(self.model.num_aux_hidden_states, dim=-1)
            hidden_states = torch.cat(
                [norm(chunk) for norm, chunk in zip(self.model.fc_norm, chunks)],
                dim=-1,
            )

        return self.model.fc(hidden_states)

Which layers get captured is a checkpoint field, eagle_aux_hidden_state_layer_ids, defaulting to three (llama_eagle3.py:L176-L182). For Llama-3-8B that means the fc is $3\times4096 \to 4096 = 50.3$ M parameters instead of 33.6 M — 16.8 M extra parameters. Under one-stream-per-weight accounting this adds 33.6 MB per projection invocation, not per recurrent draft step. The projection provides three target layers' features and is called once per verify iteration, on the whole prefill/verify batch, before the draft loop: vllm/v1/spec_decode/llm_base_proposer.py:L548-L550.

Second, the draft gets its own smaller vocabulary and a map back:

        self.lm_head = ParallelLMHead(
            self.config.draft_vocab_size,
            self.config.hidden_size,
            quant_config=get_draft_quant_config(vllm_config),
            prefix=maybe_prefix(prefix, "lm_head"),
        )
        self.logits_processor = LogitsProcessor(
            self.config.draft_vocab_size, scale=logit_scale
        )
        self.draft_id_to_target_id = nn.Parameter(
            torch.zeros(self.config.draft_vocab_size, dtype=torch.long),
            requires_grad=False,
        )

The d2t tensor in the checkpoint is renamed to draft_id_to_target_id on load (llama_eagle3.py:L387-L390) and t2d is skipped entirely. This is the row of the cost table above that takes the draft from $c=0.104$ to $c=0.051$.

EAGLE modules are trained after the fact, against a specific target, by a third party. Multi-Token Prediction inverts that: the model author trains extra prediction modules jointly with the base model and releases them in the same checkpoint. DeepSeek-V3 is the canonical case, and vLLM's MTP layer is almost line-for-line EAGLE:

    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
        previous_hidden_states: torch.Tensor,
        inputs_embeds: torch.Tensor | None = None,
        spec_step_index: int = 0,
    ) -> torch.Tensor:
        assert inputs_embeds is not None
        # masking inputs at position 0, as not needed by MTP
        inputs_embeds = torch.where(positions.unsqueeze(-1) == 0, 0, inputs_embeds)
        inputs_embeds = self.enorm(inputs_embeds)
        previous_hidden_states = self.hnorm(previous_hidden_states)

        hidden_states = self.eh_proj(
            torch.cat([inputs_embeds, previous_hidden_states], dim=-1)
        )

        hidden_states, residual = self.mtp_block(
            positions=positions,
            hidden_states=hidden_states,
            residual=None,
        )
        hidden_states = residual + hidden_states  # pre-final-norm (logits hidden)
        # ...
        return hidden_states, self.shared_head(hidden_states)

eh_proj(cat(enorm(embed), hnorm(hidden))) resembles EAGLE's feature-and-token fusion, but separate normalization, the underlying decoder block, attention/MoE structure, shared head, and joint-training objective are architectural and training differences. The methods share a speculative interface; their checkpoints are not interchangeable.

SGLang's auto-chosen parameters for DeepSeek architectures are (num_steps, topk, num_draft_tokens) = (3, 1, 4) — topk = 1, i.e. a chain, while Llama with an EAGLE head gets (5, 4, 8), a tree (python/sglang/srt/arg_groups/speculative_hook.py:L819-L843). Worth reading that function to the bottom: only LlamaForCausalLM and the two Grok architectures get (5, 4, 8); the else arm at L842-L843 hands every other architecture (3, 1, 4) too. Trees are live in SGLang but they are not its default shape — a chain is, unless you are serving Llama, Grok, or you set --speculative-eagle-topk yourself. One MTP block of a 671B-parameter MoE is a much more expensive draft step than one dense Llama layer, so the budget goes into depth rather than breadth.

Everything past this — frozen-KV MTP, multi-layer EAGLE drafts, DFlash, DSpark, standalone draft workers, and adaptive speculation that retunes $\kappa$ at runtime — belongs to §6.6.

Here is the problem trees solve. A chain draft of length 5 gives the target exactly one continuation to check. If the draft's first token is wrong, four correct guesses behind it are wasted. A tree gives the target several continuations at once and lets it pick — but a transformer forward pass eats a linear sequence, not a tree.

The trick is that a tree of $T$ nodes can be flattened into $T$ linear positions, and the tree structure recovered entirely inside the attention mask. Node $i$ attends to the shared prefix, to its own ancestors, and to itself — and to nothing else. Its siblings and their subtrees are invisible to it, so the logits at position $i$ are exactly the logits the target would have produced had you fed it that root-to-$i$ path alone. Verify $T$ candidate continuations in one pass, at the cost of $T$ positions instead of one.

Building the mask, for real

SGLang builds mask, positions, and the tree index structures in one kernel. The dispatch in build_tree_kernel_efficient picks one implementation per device: NPU gets a torch.ops.npu op, XPU the in-tree Triton kernel, CPU a C++ one, and everything else — including CUDA — falls through to sgl_build_tree_kernel_efficient, imported from the sgl_kernel package (python/sglang/srt/speculative/eagle_utils.py:L226-L285). That op's CUDA source is nevertheless readable at this SHA, vendored under python/sglang/kernels/aot/csrc/speculative/eagle_utils.cu:L34-L123, and it is a line-for-line translation of the Triton kernel below: same token_tree_idx address arithmetic, same unconditional root bit, same selected_index / topk ancestor walk, same positions = depth + seq_len. So the Triton version really is the reference, and it repays slow reading:

    # Process all draft token indices for tree mask
    for draft_tokenx in range(draft_token_num):
        if tree_mask_mode == 0:  # FULL_MASK
            token_tree_idx = (
                seq_tree_idx + (seq_len + draft_token_num) * draft_tokenx + seq_len + 1
            )
        else:
            token_tree_idx = (
                draft_token_num * draft_token_num * batch_idx
                + draft_token_num * draft_tokenx
                + 1
            )

        tl.store(tree_mask_ptr + token_tree_idx - 1, 1)
        for i in range(draft_token_num - 1):
            tl.store(tree_mask_ptr + token_tree_idx + i, 0)

        if draft_tokenx > 0:
            # Build tree path for draft_tokenx > 0
            cur_position = draft_tokenx - 1
            position = 0
            should_continue = 1

            for _ in range(depth):
                if should_continue:
                    position += 1
                    tl.store(tree_mask_ptr + token_tree_idx + cur_position, 1)

Unpack the address arithmetic. Under FULL_MASK the mask is one flat bool array of $\big(\sum_b \text{seq\_len}_b\big)\cdot T + T^2 \cdot \mathrm{bs}$ entries; each row is $\text{seq\_len} + T$ wide, so row draft_tokenx starts at seq_tree_idx + (seq_len + draft_token_num) * draft_tokenx. token_tree_idx is that plus seq_len + 1, so token_tree_idx - 1 is column seq_len — the tree's node 0. Every row gets that bit set unconditionally: every node descends from the root. The following loop zeroes the remaining $T-1$ tree columns. The prefix columns are never touched by the kernel; they are pre-filled True by the caller, which is why fill_prefix_mask exists as a knob.

Then the ancestor walk. cur_position = draft_tokenx - 1 is an index into selected_index (the $T-1$ chosen candidates, root excluded), and token_tree_idx + cur_position is column seq_len + 1 + cur_position — which on the first iteration is the node's own diagonal. Each iteration sets one bit and climbs one level, up to depth times. position counts the climb, and that count is the node's depth, written out as positions[draft_tokenx] = position + seq_len (spec_tree.py:L170-L173). The RoPE position of a tree node is its depth. Siblings share a position; that is correct and not a bug.

The parent lookup is the one genuinely non-obvious step:

                    parent_tb_idx = (
                        tl.load(
                            selected_index_ptr
                            + batch_idx * selected_index_stride
                            + cur_position
                        )
                        // topk
                    )
                    if parent_tb_idx == 0:
                        should_continue = 0
                    else:
                        parent_token_idx = tl.load(
                            parent_list_ptr
                            + batch_idx * parent_list_stride
                            + parent_tb_idx
                        )

                        # Find cur_position for next iteration
                        found = 0
                        for cp in range(draft_token_num - 1):
                            if found == 0:
                                if (
                                    tl.load(
                                        selected_index_ptr
                                        + batch_idx * selected_index_stride
                                        + cp
                                    )
                                    == parent_token_idx
                                ):
                                    cur_position = cp
                                    found = 1
                        if found == 0:
                            should_continue = 0

selected_index[cur_position] is the node's index into the flat 68-wide candidate pool. Integer-dividing by topk recovers which expansion group produced it: group 0 is the step-0 expansion of the root, group $1+(i-1)\cdot\mathrm{topk}+r$ is the expansion of the rank-$r$ survivor at step $i$. parent_list is the table mapping group id to the flat candidate index of that group's parent — built in _select_top_k_tokens_first as arange(-1, topk) (the $-1$ is the root sentinel) and extended each step by topk_cs_index + (topk_sq * (i - 1) + topk) (python/sglang/srt/speculative/spec_utils.py:L296-L300, L338-L341). So parent_tb_idx == 0 means "my parent is the root, stop climbing". Otherwise a linear scan seeks the selected parent. This assumes a closed selected set; the excerpt alone does not establish tie-safe closure. A missing parent must be rejected/repaired, never treated as a root. Test identical scores and underflow before trusting the downstream mask.

Flattening, and the two link arrays

The same kernel writes three retrieve_* arrays that turn the tree into something a verifier can walk. retrieve_index[b][i] is node $i$'s slot in the flattened batch-wide draft-token buffer. retrieve_next_token[b][p] is $p$'s first child. retrieve_next_sibling[b][i] is $i$'s next sibling. That is the classic first-child/next-sibling encoding: two $T$-wide int arrays represent an arbitrary tree with no pointers and no recursion.

Walking the accepted path

Verification is that walk, done on the GPU with one program per request. Note that target_predict here is the target's argmax at each node's position — computed for all $T$ nodes by the single verify forward.

    # Tree traversal loop
    should_continue = 1
    for j in range(1, num_speculative_tokens):
        if should_continue:  # Early exit guard
            cur_index = tl.load(
                retrieve_next_token_ptr + bx * num_draft_tokens + cur_index
            )

            # Load target token once per level (before sibling search)
            # last_accept_retrieve_idx is constant during sibling traversal
            target_row = last_accept_retrieve_idx // num_draft_tokens
            target_col = last_accept_retrieve_idx % num_draft_tokens
            target_token = tl.load(
                target_predict_ptr + target_row * num_draft_tokens + target_col
            )

            # Traverse siblings
            found_match = 0
            for _ in range(num_draft_tokens):  # Max iterations = num_draft_tokens
                if found_match == 0:  # Early exit guard
                    # Check if we've reached end of sibling list
                    is_valid = cur_index != -1
                    # ...
                    draft_index = tl.load(retrieve_index_ptr + safe_index)
                    draft_token = tl.load(candidates_ptr + safe_index)

                    # Check for token match (only valid when is_valid is True)
                    token_match = is_valid & (draft_token == target_token)

Descend to the first child, then scan siblings for one whose draft token equals the target's argmax at the parent's position. If one matches, accept it and recurse; if the sibling list runs out, stop. Because at most one sibling can match a single argmax, the walk is deterministic and the accepted set is a path. Correctness is inherited from §6.2's argument applied along that path — for sampling the same walk runs under tree_speculative_sampling_target_only, dispatched at python/sglang/srt/speculative/eagle_utils.py:L772-L775.

Note the loop bound: num_speculative_tokens is accept_index.shape[1], which is max_tree_depth = spec_steps + 1 (python/sglang/srt/speculative/eagle_info.py:L43-L48) — not num_draft_tokens. With (5, 4, 8) you verify 8 nodes but can accept at most 6 tokens. A tree buys breadth, not depth.

Discarding the rejected branches' KV

All $T$ nodes wrote KV during the verify pass, into $T$ contiguous slots starting at seq_lens (assign_extend_cache_locs_uniform_func, python/sglang/srt/speculative/eagle_utils.py:L527-L534). Only the accepted path's KV is real. For a chain that is already the front of the block and nothing needs doing; for a tree the accepted nodes are scattered, so SGLang gathers them forward:

def _finalize_accept_tree_path(
    batch: ScheduleBatch,
    accept_index: torch.Tensor,
    accept_lens: torch.Tensor,
    predict: torch.Tensor,
    logits_output: Any,
    bs: int,
    *,
    token_to_kv_pool_allocator: Any,
    num_draft_tokens: int,
) -> torch.Tensor:
    """Tree drafting (topk > 1): move the accepted path -- KV slots, predict,
    hidden_states -- to the contiguous front of each per-req block, which the
    downstream chain-layout code (draft-extend select_index, committed-KV reads)
    assumes. Returns compacted predict; mutates logits_output.hidden_states
    (moved only when present)."""
    move_accept_tokens_to_target_kvcache(
        batch, accept_index, accept_lens - 1, token_to_kv_pool_allocator
    )
    predict = _compact_accept_to_front(
        predict, accept_index, bs, num_draft_tokens=num_draft_tokens
    )

move_accept_tokens_to_target_kvcache computes the destination slots $[\,\text{seq\_lens},\ \text{seq\_lens}+\text{accept\_lens}\,)$, gathers the source slots via out_cache_loc[accept_index], and issues one move_kv_cache(tgt, src) (python/sglang/srt/speculative/spec_utils.py:L723-L756). The rejected branches are never explicitly freed here — their slots become trailing overshoot past the request's new seq_lens, reclaimed by exactly the machinery §6.2 describes for a rejected chain suffix. That is the whole point of allocating the tree contiguously.

Finally, the mask reaches the kernel. FlashInfer's paged prefill wrapper takes an arbitrary boolean mask, and SGLang hands the tree mask straight through:

        wrapper_paged.begin_forward(
            qo_indptr,
            kv_indptr,
            kv_indices,
            self.kv_last_page_len[:bs],
            self.num_qo_heads,
            self.num_kv_heads,
            self.head_dim,
            1,
            q_data_type=self.q_data_type,
            kv_data_type=self.data_type,
            custom_mask=use_custom_mask,

The mask arrives as spec_info.custom_mask from EagleVerifyInput.generate_attn_arg_prefill (python/sglang/srt/speculative/eagle_info.py:L82-L136), which also computes mask_numel = paged_kernel_lens_sum * draft_token_num + draft_token_num**2 * batch_size — the same layout the Triton kernel wrote into.

The mask's memory problem

FULL_MASK's size is $\big(\sum_b \text{seq\_len}_b\big)\cdot T + T^2\,\mathrm{bs}$ bools. Derived: at batch 1, seq 2,000, $T=8$ that is 16,064 bytes — nothing. At batch 256 with 32k contexts it is 67.1 MB, rewritten every decode step. The code says so:

        elif tree_mask_mode == TreeMaskMode.FULL_MASK:
            # Only the [0, seq_len) prefix columns depend on this fill; the
            # kernel below writes every tree cell itself. Skip the (up to
            # 100s of MB) per-step memset when nothing reads the mask.
            if fill_prefix_mask:
                tree_mask.fill_(True)

The information content of a tree mask is $T^2$ bits per request. Everything else is a constant True. So SGLang keeps two cheaper encodings — QLEN_ONLY ($T \times T$ bools per request, 64 bytes at $T=8$) and QLEN_ONLY_BITPACKING, which picks the narrowest unsigned int that holds $T$ bits: packed_dtype_idx = int(math.ceil(math.log2((num_verify_tokens + 7) // 8))) (eagle_utils.py:L193-L199), so $T=8$ becomes one uint8 per row, 8 bytes per request. Which mode is used is a property of the attention kernel, not of the tree: default_tree_mask_mode() returns QLEN_ONLY on CPU because the Intel AMX verify kernel consumes it directly, and FULL_MASK otherwise (eagle_utils.py:L145-L149).

Redo §6.2's accounting with breadth. Let $T$ be the number of verified nodes, $D$ the maximum tree depth, $S$ the number of draft forward passes per iteration, and $c$ the cost of one draft forward as a fraction of one target forward. Let $\beta_i$ be the probability that the target's continuation at depth $i$ appears among the tree's candidates at depth $i$, given that depth $i-1$ was accepted. Then the expected number of tokens emitted per iteration is

Three things to check against §6.2. First, set $\mathrm{topk}=1$: then $T = D+1 = \kappa+1$, $S = \kappa$, every $\beta_j = \alpha$, and $E_{\text{tree}}$ collapses to $\sum_{i=0}^{\kappa}\alpha^i = (1-\alpha^{\kappa+1})/(1-\alpha)$ — §6.2's $E$ exactly, with $1+\kappa c$ in the denominator. The chain is the degenerate tree.

Second, this denominator assumes target verification time is independent of tree width in an ideal weight-bandwidth model. At 8 positions the 1.05 MB of new KV writes is small beside 15.01 GB of weights, but attention, logits, layouts, routing, launches and achieved bandwidth can still change. Width is not free by a hardware theorem. Measure $t_V(T,B)$ and use $E_{\mathrm{tree}}/[t_V(T,B)/t_T+S c]$ when the constant-cost approximation fails.

Third, for a fixed prefix and nested candidate sets, adding candidates cannot decrease greedy target coverage: $\beta(k)=\sum_{r\le k}p_r\ge p_1$. The inequality is strict only if added candidates carry positive target-argmax probability. At a fixed total node budget, widening can remove depth or change surviving prefixes, so neither per-depth acceptance nor expected emitted length is ordered. A perfect chain already has coverage one; breadth cannot improve it. Global reranking is a heuristic whose value depends on target/draft agreement.

One more term the formula hides: at $c=0.051$ a draft forward is 229 µs of weight traffic, and five of them per iteration is 1.15 ms of work split across five separate model launches. At that size, Python dispatch and kernel launch overhead are not rounding errors. Hence SGLang gives the drafter its own CUDA graph — one that captures the entire multi-step loop, not one step:

        def run_once():
            self.draft_attn_backend.init_forward_metadata_in_graph(forward_batch)
            # ...
            output_cache_loc_backup = forward_batch.out_cache_loc
            hidden_states_backup = forward_batch.spec_info.hidden_states
            dsa_topk_indices_backup = forward_batch.spec_info.dsa_topk_indices

            ret = self.eagle_worker.draft_forward(forward_batch)

            forward_batch.out_cache_loc = output_cache_loc_backup
            forward_batch.spec_info.hidden_states = hidden_states_backup
            forward_batch.spec_info.dsa_topk_indices = dsa_topk_indices_backup
            forward_batch.positions.sub_(self.eagle_worker.speculative_num_steps - 1)
            return ret

draft_forward — all num_steps - 1 model calls, all the select_top_k_tokens reranking, the final organize_draft_results — is one graph replay. The positions.sub_() at the end undoes the in-place advance the loop made, so the captured buffers come back to their entry state and the graph is replayable. This also forces the drafter to own a per-step container of attention backends rather than one backend, because each captured step has different metadata: python/sglang/srt/speculative/draft_utils.py:L71-L78 is explicit that create_decode_backend "Returns a per-step CONTAINER, not an AttentionBackend".

Llama-3-8B target, EAGLE-3 draft head, (num_steps, topk, num_draft_tokens) = (5, 4, 8), batch 1, prefix length 2,000, greedy sampling. In order, with real function names:

1. draft() (eagle_worker_v2.py:L501-L512) builds the draft forward batch from EagleDraftInput.bonus_tokens plus the topk_p/topk_index the previous draft-extend left behind.
2. draft_forward() (eagle_worker_v2.py:L621-L633) loops for i in range(5), breaking before the model call at $i=4$, so four draft model calls happen and score_list accumulates $4 + 4\cdot16 = 68$ scored candidates. per_step_draft_out_cache_loc (eagle_utils.py:L64-L86) reshapes the draft's KV slots to [num_steps, bs*topk].
3. organize_draft_results() keeps the global top-7 by cumulative path probability, then build_tree_kernel_efficient() prepends the bonus token as node 0 and writes a $2000\cdot8 + 8\cdot8 = 16{,}064$-bool FULL_MASK plus the three retrieve_* arrays.
4. eagle_prepare_for_verify() allocates 8 contiguous KV slots at position 2,000, then the target verify forward runs 8 query positions against 2,008 KV positions under the custom mask. This is the 4.48 ms.
5. eagle_sample() (eagle_utils.py:L727-L742) takes argmax over all 8 rows and calls verify_tree_greedy_func, which walks the tree. Suppose nodes 0, 2, 5, 6 accept: accept_lens = 4, accept_index = [0, 2, 5, 6, -1, -1].
6. _finalize_accept_tree_path() moves slots $\{0,2,5,6\}$ of the 8-slot block to positions 0..3 and compacts predict and hidden_states the same way. Slots 4..7 become overshoot.
7. _draft_extend_for_decode() runs the draft head over the 4 accepted positions to bring its own KV cache up to date and to produce the next topk_p. That is the fifth and last draft forward of the iteration.

Four emitted tokens cost one target verify, four recurrent draft calls, one draft extension, and one auxiliary-feature projection. If the four calls and extension are each approximated by 0.229 ms and the extra projection traffic by 0.010 ms, the illustration gives 4.48 + 5(0.229) + 0.010 = 5.635 ms, or about 3.18 times against 17.92 ms unspeculated. This optimistic accounting charges the once-per-iteration projection once. Extension batch shape, cache reads, projection invocation boundaries and acceptance require actual profiling; the figure is not a measured speedup.

Confusing num_draft_tokens with max_tree_depth

This is the single most common tree bug and SGLang carries scar tissue for it in comments. The accepted-path arrays are spec_steps + 1 wide; the node arrays are num_draft_tokens wide. For a chain those are equal and everything works. Turn on topk > 1 and they diverge:

        accept_tokens = predict[accept_index]
        bonus_tokens = torch.empty_like(accept_lens, dtype=torch.int32)
        # stride = accept_tokens per-req width = accept_index.shape[1]
        # (spec_steps + 1); NOT num_draft_tokens, wrong for topk > 1 trees.
        fill_bonus_tokens_func(
            accept_tokens,
            accept_lens,
            bonus_tokens,
            accept_index.shape[1],
            bs,
        )

and in move_accept_tokens_to_target_kvcache:

    # accept_index element count, NOT bs * num_draft_tokens: for topk > 1 the
    # tree exceeds the accepted chain, over-reading accept_index (illegal memory).
    size = bs * accept_index.shape[1]

The failure mode when you get it wrong is an illegal memory access inside a kernel launched two steps later, which is about the worst diagnostic in CUDA. If you are writing a new tree-shaped speculative algorithm, the width you want is max_tree_depth everywhere an accepted thing is indexed and num_draft_tokens everywhere a node is indexed.

Draft slot count assertions

A mismatch between the drafter's assumed shape and the allocator's is caught early and loudly:

    expected = batch_size * topk * num_steps
    assert out_cache_loc.shape[0] == expected, (
        f"out_cache_loc.shape[0]={out_cache_loc.shape[0]} != "
        f"batch_size * topk * num_steps = {batch_size}*{topk}*{num_steps}={expected}"
    )

Note what that means for memory: a tree drafter reserves $\mathrm{bs}\cdot\mathrm{topk}\cdot \text{num\_steps}$ draft KV slots per iteration — 20 per request at (5,4,·) — plus $T=8$ verify slots. That is 28 slots per request per iteration held transiently, against 1 for plain decode.

Attention backends that cannot express the tree

Trees are not portable across attention backends, and SGLang refuses at startup rather than silently producing a wrong-but-plausible answer. Three separate gates. trtllm_mha is rejected outright for any tree — "trtllm_mha backend only supports topk = 1 for speculative decoding." (python/sglang/srt/arg_groups/speculative_hook.py:L621-L625). With page_size > 1 a tree needs the two-pass cascade draft-decode, which only flashinfer, fa3 and triton implement; the comment above the check is blunt — "flashmla / trtllm_mla / cutlass_mla can't express the per-branch tree, so reject" (python/sglang/srt/arg_groups/speculative_hook.py:L685-L699). And the multi-step draft loop needs a decode backend from a fixed map at all, or you get "EAGLE is not supported in decode attention backend {backend_type}" (python/sglang/srt/speculative/draft_utils.py:L98-L104). Separately, switching to rejection sampling requires the drafter to expose a target-vocabulary proposal distribution, which a reduced-draft-vocab EAGLE-3 head does not have for free:

                raise ValueError(
                    "Rejection sampling requires a target-vocab draft proposal "
                    "distribution; the current speculative algorithm/draft worker "
                    "does not produce one (draft_probs missing or vocab-mismatched)."
                )

The mask memset you did not budget for

If you run trees at large batch with long contexts and see an unexplained per-step cost that scales with total context length rather than with anything speculative, it is the FULL_MASK fill. Check whether your backend can take QLEN_ONLY, and check verify_mask.is_read — SGLang only pays the memset when something actually reads the prefix columns (eagle_worker_common.py:L350-L356).

The root is already committed and predicts the first draft token. A selected node attends to that root and its own ancestral path, not a sibling. The tiny reference checks closure before constructing the mask. Production global score selection needs a tie rule or repair that supplies this precondition. Target probabilities must be read at the parent prediction position; the final accepted node supplies the bonus distribution.

import numpy as np

parents = [-1, 0, 1, 0]  # root, first branch, its child, sibling
scores = [1.0, 0.5, 0.5, 0.4]
def depth(i):
    n = 0
    while parents[i] != -1:
        i = parents[i]
        n += 1
    return n

ranked = sorted(range(1, 4), key=lambda i: (-scores[i], depth(i), i))
selected = {0, *ranked[:2]}
assert selected == {0, 1, 2}
assert all(parents[i] in selected for i in selected if i != 0)
mask = np.zeros((4, 4), dtype=bool)
for i in selected:
    ancestor = i
    while ancestor != -1:
        mask[i, ancestor] = True
        ancestor = parents[ancestor]
assert mask[2].tolist() == [True, True, True, False]
bad_topk = {0, 2}  # tied child selected without parent
assert any(parents[i] not in bad_topk for i in bad_topk if i)
print("Ancestor closure and sibling exclusion pass.")

The cleanest experiment in this chapter is chain-versus-tree at a fixed depth. Same draft head, same number of draft forwards, only topk changes — so any acceptance difference is the tree's doing.

python -m sglang.launch_server --model-path meta-llama/Llama-3.1-8B-Instruct \
  --speculative-algorithm EAGLE3 \
  --speculative-draft-model-path <eagle3-head> \
  --speculative-num-steps 5 --speculative-eagle-topk 1 --speculative-num-draft-tokens 6

python -m sglang.launch_server --model-path meta-llama/Llama-3.1-8B-Instruct \
  --speculative-algorithm EAGLE3 \
  --speculative-draft-model-path <eagle3-head> \
  --speculative-num-steps 5 --speculative-eagle-topk 4 --speculative-num-draft-tokens 8

Three things to measure. (1) Accepted length per iteration, at batch 1 and again at batch 64 — the tree's advantage should survive at batch 1 and evaporate as $T\cdot B$ approaches $I^{*}=295$. (2) The draft's own share of step time; SGLang's draft loop is one graph replay, so it shows up as a single fat kernel region. (3) Set SGLANG_SIMULATE_ACC_LEN=3.0 (python/sglang/srt/environ.py:L468) to force a synthetic acceptance length: that removes acceptance from the equation entirely and isolates the pure plumbing cost of the tree machinery, which is the number you want when you suspect the mask build rather than the model.

For vLLM, inspect the selected runner before comparing an eagle3 configuration. The quoted V1 path drafts chains, but runner selection is a separate dispatch decision. Cross-engine accepted-length or latency differences also include kernels, numerics, tokenization and scheduling. A within-engine chain/tree ablation with the same checkpoint and sampled workload isolates tree policy more closely.

• Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads — Cai et al., 2024. arXiv:2401.10774. Where the multi-head idea and the Cartesian-product draft tree come from.
• EAGLE: Speculative Sampling Requires Rethinking Feature Uncertainty — Li et al., 2024. arXiv:2401.15077. The feature-level autoregression argument, and the ablation showing why token-level drafting from hidden states is not enough.
• EAGLE-2: Faster Inference of Language Models with Dynamic Draft Trees — Li et al., 2024. arXiv:2406.16858. Context-aware tree shaping; the mechanism SGLang's _select_top_k_tokens_later and organize_draft_results implement.
• EAGLE-3: Scaling up Inference Acceleration of Large Language Models via Training-Time Test — Li et al., 2025. arXiv:2503.01840. Multi-layer feature fusion and the dropped feature-prediction constraint. See the Unverified callout in §5 for what serving code can and cannot confirm.
• DeepSeek-V3 Technical Report — DeepSeek-AI, 2024. arXiv:2412.19437. The MTP training objective, and the reason a pretrained checkpoint ships its own drafter.
• SpecInfer — Miao et al., 2023. arXiv:2305.09781. Tree-based speculative verification stated in general terms, independent of any particular drafter; the ancestry-mask construction in this chapter is its descendant.
• Source, read this session. vllm/v1/spec_decode/{eagle,medusa,llm_base_proposer,extract_hidden_states}.py, vllm/model_executor/models/{medusa,llama_eagle,llama_eagle3,deepseek_mtp}.py; python/sglang/srt/speculative/{eagle_utils,eagle_info,eagle_worker_v2,eagle_worker_common,spec_utils,draft_utils,eagle_draft_cuda_graph_runner,ragged_verify}.py and python/sglang/kernels/ops/speculative/spec_tree.py.
• Neighbours: the algorithm and its proof in §6.2; draft models, n-gram and suffix decoding in §6.3; DFlash, DSpark, frozen-KV MTP, multi-layer EAGLE, standalone workers and adaptive speculation in §6.6. Formula symbols in FORMULAS.