M7.5 - EAGLE Feature Drafting: Uncertainty, Dynamic Trees, HASS, and EAGLE-3

M7.5 - EAGLE Feature Drafting: Uncertainty, Dynamic Trees, HASS, and EAGLE-3

0x00 Overview And Background

M7.4 showed why independent future-token heads do not define a coherent autoregressive continuation. EAGLE chooses a different state variable:

do not draft only future token ids;
autoregress over the target model's internal feature space;
condition each feature step on the token that was actually sampled;
then let the target model verify the resulting token tree.

This creates a deeper question than “is a feature easier to predict than a token?”

Which target state is reused, how is sampling uncertainty injected into the drafter, what distribution mismatch appears during training, and how do EAGLE-2, HASS, and EAGLE-3 change those contracts?

Background knowledge: a drafter can imitate target state as well as target tokens

Classical speculative decoding trains or selects a smaller model to approximate the target’s token distribution. EAGLE moves the prediction interface inside the target model: it reuses target features, conditions on the sampled token, and autoregresses with a small draft network before target verification.

This combines ideas from knowledge distillation and speculative execution:

teacher state:
    a feature extracted from the frozen target model

student or drafter:
    predicts the next draft state from feature history and sampled-token input

shared output head:
    maps draft state into token candidates

verifier:
    remains the target model and owns accepted state

Later EAGLE variants change different contracts: EAGLE-2 changes tree allocation, HASS changes training alignment with speculative trajectories, and EAGLE-3 changes the feature target and objective.

0x00A Historical Lineage And Owner Contract

Feature drafting did not begin as a claim that hidden vectors are inherently more meaningful than tokens. It began from a runtime asymmetry. A small token-level draft model must reproduce a target distribution from its own independently learned state. A post-hoc feature drafter can instead start from state the target has already computed, but that shortcut creates a new question: which target state remains predictive after sampling selects one token branch?

The EAGLE family answers that question by changing one contract at a time:

Owner contract. This article owns the exact source-to-feature coordinate, the post-hoc drafter’s frozen-target training path, its provisional tree state, and the EAGLE-family changes above. It does not rederive M7.1’s exact sampling law, M7.2’s generic tree-mask compiler, or M7.4’s generic MTP objective. M7.6 owns the integrated DeepSeek MTP chain, while M7.7 and M7.8 own cross-proposer topology and production serving policy.

A feature drafter can be locally accurate yet still be wired to the wrong target layer, shifted by one token, trained through a supposedly frozen head, or allowed to publish unverified state. Those are different programs even when every tensor shape matches.

0x01 Claim Status And Scope

Benchmark speedups in this article are intentionally not used as architectural proof. They are source-reported measurements under the papers’ model, task, batch, sampling, and hardware settings.

0x01A Immutable Feature-Drafting Manifest

Before extracting one feature row, freeze the semantic address of every object:

The release artifact should therefore carry a target checkpoint fingerprint, a tokenizer fingerprint, the feature-layer address, normalization and LM-head identities, dtype policy, drafter configuration, and tree ABI. A drafter is not portable merely because another model has the same hidden width.

target checkpoint fingerprint:
    model config + weight digest + tokenizer/vocabulary digest

feature address:
    decoder layer id
    pre/post-normalization state
    tensor layout and dtype

drafter ABI:
    [feature_t ; embedding(x_{t+1})] -> predicted_feature_{t+1}
    separate drafter KV
    frozen compatible embedding and frozen LM head

0x02 Why Token Sampling Makes Feature Prediction Ambiguous

Let the frozen target model produce a second-to-top-layer feature:

$$f_t^T \in \mathbb{R}^{d}$$

The target LM head produces logits and a token distribution:

$$q_{t+1} = \operatorname{softmax}(W_U f_t^T)$$

A token is sampled:

$$x_{t+1} \sim q_{t+1}$$

The next target feature depends on the sampled token:

$$f_{t+1}^T = F_T(f_{\le t}^T, x_{\le t+1})$$

If the drafter sees only $f_t^T$, the target for the next feature can be multimodal. The same feature state may lead to different sampled tokens and therefore different next features:

prefix:
    "I"

possible sampled token:
    "am"
    "always"

next target feature:
    follows the "I am ..." branch
    or follows the "I always ..." branch

A deterministic regression model trained only on $f_t^T$ can be pushed toward an average of incompatible feature targets. That average need not lie on either valid continuation manifold.

Source: EAGLE, Figure 4 in arXiv v1. The same feature after I precedes different next features depending on whether sampling chooses am or always. This is a causal ambiguity argument, not evidence that Euclidean feature regression is universally superior to token drafting.

The regression failure can be stated formally. Let $X$ be the sampled token and $F'$ the next target feature. A feature-only squared-error predictor returns the conditional mean

$$g^*(f_t)=\mathbb{E}[F'\mid F_t=f_t] =\sum_x P(X=x\mid f_t)\,\mathbb{E}[F'\mid f_t,X=x].$$

When sampled tokens lead to separated feature modes, this mean lies between the modes. Supplying $X=x_{t+1}$ changes the optimal prediction to

$$g^*(f_t,x_{t+1}) =\mathbb{E}[F'\mid F_t=f_t,X=x_{t+1}],$$

the branch-conditional mean. The advanced token removes uncertainty caused by that sampled branch; it does not remove uncertainty from later sampling, target stochasticity, or an under-capacity drafter.

EAGLE exposes the sampled branch explicitly by feeding the embedding of the advanced token $x_{t+1}$:

$$e_{t+1} = E(x_{t+1})$$ $$\hat f_{t+1}^D = D_\phi(f_{\le t}^D, e_{\le t+1})$$

The shifted token is not redundant with the feature. It carries the discrete outcome of the sampling event.

0x02A The Index Shift

For a sequence of target features and tokens:

target features:
    f_1, f_2, f_3, ..., f_{S-1}

advanced tokens:
    x_2, x_3, x_4, ..., x_S

draft targets:
    f_2, f_3, f_4, ..., f_S

At aligned draft row $t$:

input:
    target/draft feature f_t
    embedding of sampled token x_{t+1}

target:
    next target feature f_{t+1}
    target token distribution for x_{t+2}

An off-by-one error can still produce shape-compatible tensors. The article therefore treats indexes as part of the model contract, not as an implementation detail.

The source convention can be written without overloaded prose:

$$E_{2:i+1}=\operatorname{Embedding}(X_{2:i+1}), \qquad \hat F_{2:i+1}=D_\phi(E_{2:i+1},F_{1:i}),$$

and each predicted feature $\hat f_{t+1}$ is mapped by the frozen target LM head to a distribution for the next token position, $t+2$. That distribution does not require an observed token at that position. With one-based coordinates $x_1,\ldots,x_S$ and teacher features $f_1,\ldots,f_S$, the soft-target rows are:

Soft distillation therefore has up to $S-1$ boundary-valid rows. The final teacher feature supplies $q_{S+1}$ even though $x_{S+1}$ is not in the segment. A hard-label alternative would need the actual $x_{t+2}$ and would have at most $S-2$ rows. These are different objectives, not two names for the same mask.

Separate token existence from the dataset’s choice of supervised rows. Let $v_t$ mark a valid token/teacher-feature address and let $\ell_t$ be a supervision flag indexed by the source row, as in the pinned training pipeline. For the soft objective:

$$b_t=\mathbf 1[\operatorname{seg}(t)=\operatorname{seg}(t+1)]\,v_t v_{t+1}, \qquad m_t^{soft}=b_t\ell_t.$$

If the same row-level supervision policy is used for a hard-label alternative, the extra observed-token requirement is:

$$m_t^{hard}=m_t^{soft}\,v_{t+2} \mathbf 1[\operatorname{seg}(t+1)=\operatorname{seg}(t+2)].$$

An out-of-range address is invalid; it is never gathered. Compile boundary-valid addresses before gathering features or tokens. Multiplying a loss by zero after a cross-segment gather does not undo the invalid read. Preserve boundary-valid rows with $\ell_t=0$ as causal context while excluding their losses: a supervision mask is not an attention mask. Both teacher feature extraction and draft attention must remain causal and segment-isolated, with original batch, segment, and position identities retained after compaction.

The pinned EAGLE training code shifts advanced tokens and feature targets once, pads their last row, and disables that row’s loss. Its soft targets come from those shifted teacher features. The packed-segment safeguards above are a local extension of this coordinate contract, not a claim that this reference loader itself implements arbitrary sequence packing.

0x03 EAGLE-1 Tensor Path

For batch $B$, sequence length $S$, and hidden width $d$:

target feature sequence F:
    [B, S, d]

advanced token ids X+:
    [B, S]

advanced embeddings E(X+):
    [B, S, d]

concatenated draft input:
    [B, S, 2d]

fusion projection output:
    [B, S, d]

draft decoder output:
    [B, S, d]

shared LM-head logits:
    [B, S, V]

Source: EAGLE, Figure 6. The pipeline combines shifted token embeddings with target-model features, autoregresses in feature space, and uses the shared LM head to obtain draft tokens. The target and drafter maintain distinct runtime state.

The fusion step is:

$$c_t = \left[f_t ; e_{t+1}\right] \in \mathbb{R}^{2d}$$ $$u_t = W_F c_t + b_F, \qquad W_F \in \mathbb{R}^{d \times 2d}$$

The draft decoder produces:

$$\hat f_{t+1}^D = \operatorname{DraftBlock}_\phi(u_{\le t})_t$$

The frozen target LM head converts the predicted feature to draft logits:

$$z_{t+2}^D = W_U \hat f_{t+1}^D$$ $$p_{t+2}^D = \operatorname{softmax}(z_{t+2}^D)$$

The draft token $x_{t+2}^D$ is sampled from $p_{t+2}^D$, embedded, and fed into the next draft step.

0x03A Shape-Checked Draft Step

def eagle_step(feature_history, advanced_token_ids, draft, target):
    # Dense, nonempty feature history: [B, L, d]; token ids: [B, L].
    assert len(feature_history.shape) == 3
    assert feature_history.shape[1] > 0
    assert advanced_token_ids.shape == feature_history.shape[:2]
    hidden_size = feature_history.shape[-1]

    token_embed = target.embed(advanced_token_ids)       # [B, L, d]
    assert token_embed.shape == feature_history.shape
    assert target.embed.weight.requires_grad is False
    assert target.lm_head.weight.requires_grad is False

    concatenated = concat([feature_history, token_embed], -1)
    assert concatenated.shape == (*feature_history.shape[:2], 2 * hidden_size)

    fused = draft.fc(concatenated)                      # [B, L, d]
    assert fused.shape == feature_history.shape

    next_feature, draft_kv = draft.decoder(
        fused,
        causal_mask=True,
        use_cache=True,
    )                                                   # [B, L, d]
    assert next_feature.shape == feature_history.shape

    logits = target.lm_head(next_feature[:, -1])         # [B, V]
    next_token = sample(logits)                         # [B]
    return next_feature[:, -1], next_token, draft_kv

The checks belong to different stages: concatenated has width 2d; fused after the projection and next_feature after the decoder both have width d. Their batch and sequence dimensions must stay unchanged. Keeping separate names prevents a post-projection tensor from being checked against the pre-projection width. The checks run inside the function, where these intermediate tensors exist. This is a dense, nonempty inference sketch; the packed training-row contract is developed separately below.

This sketch hides several operations that the pinned reference implementation makes explicit. In eagle/model/cnets1.py, the frozen token embedding produces $e_{t+1}$, a learned 2d -> d projection fuses token and feature state, the draft decoder applies its own causal or tree mask, and the compatible frozen LM head projects draft state to vocabulary logits. The drafter cache is separate from the target cache even when both use the same hidden width.

0x03B Source-To-Tensor Trace

For a padded batch, do not begin from already aligned arrays. Build the object graph and prove each address:

target_mask = segment_causal_mask(token_valid, segment_id, position_id)
with no_grad():
    target_out = frozen_target(
        tokens, attention_mask=target_mask,
        position_ids=position_id, output_hidden_states=True,
    )
    feature = select_feature_tap(
        target_out.hidden_states,
        layer_id=manifest.feature_layer,
        normalization=manifest.feature_normalization,
    )                                                   # [B,S,d]

rows = compile_soft_rows(
    tokens=tokens, features=feature,
    token_valid=token_valid, segment_id=segment_id,
    position_id=position_id, supervise_at_source=loss_mask,
)
# Plan legal adjacent pairs first, then gather f_t, x_{t+1}, f_{t+1}.
# For one unpadded segment these correspond to:
# feature[:, :-1], tokens[:, 1:], feature[:, 1:].
# Keep unsupervised but boundary-valid rows for context.
# Empty plans / no supervised rows skip the training update.

advanced_embedding = frozen_embedding(rows.advanced_token)
draft_input = learned_fusion(concat(rows.source_feature,
                                    advanced_embedding, dim=-1))
predicted_feature = draft_decoder(draft_input,
                                  attention_mask=rows.causal_mask,
                                  position_ids=rows.position_ids)
draft_logits = frozen_lm_head(predicted_feature)
target_logits = frozen_lm_head(rows.feature_target)

The last line is easy to misread. The target feature at $t+1$ produces the distribution for position $t+2$; its feature address advances only once relative to the source feature. No gather of token $x_{t+2}$ is needed for soft distillation. Here the helper retains original batch/segment/position addresses and builds block-causal draft visibility over the compacted rows. It must not connect the end of one packed sample to the start of another.

0x03C The Drafter Has Its Own KV State

EAGLE’s one-layer decoder is autoregressive. Repeated draft steps therefore maintain draft-model K/V:

target KV:
    committed target prefix and temporary verification branches

draft KV:
    EAGLE autoregression state for candidate generation

feature state:
    accepted target feature at the speculation boundary,
    then predicted draft features for unverified steps

The draft K/V is cheaper than target K/V because the drafter is small, but it is not zero. A runtime needs separate cache ownership and rollback rules.

0x04 EAGLE-1 Training Objective

The target model is frozen. Training first obtains target features and distributions, then updates the draft autoregression head.

0x04A Feature Regression

EAGLE uses a Smooth L1 feature objective. For residual $r = \hat f - f^T$, the elementwise Huber-style loss is:

$$\operatorname{smoothL1}_\beta(r) = \begin{cases} \frac{r^2}{2\beta}, & |r| < \beta \\ |r| - \frac{\beta}{2}, & \text{otherwise} \end{cases}$$

The feature loss is:

$$\mathcal{L}_{reg} = \frac{1}{\sum_t m_t} \sum_t m_t\frac{1}{d}\sum_{j=1}^{d} \operatorname{smoothL1}_\beta (\hat f_{t+1,j}^{D}-f_{t+1,j}^{T})$$

Here $t$ indexes the source row and $m_t=m_t^{soft}$; the feature target is at $t+1$. The classification loss below uses the same supervised rows. Both normalized losses require $\sum_t m_t>0$; an empty-supervision batch skips the optimizer update.

Smooth L1 is quadratic near zero and linear for large errors, reducing the influence of feature outliers compared with pure squared error.

0x04B Token Classification

Feature alignment is intermediate. The runtime goal is a draft token distribution that agrees with the target well enough to obtain high acceptance:

$$\mathcal{L}_{cls} = -\frac{1}{\sum_t m_t}\sum_t m_t \sum_{v=1}^{V} q_{t+2}^T(v) \log p_{t+2}^D(v)$$

The EAGLE paper and pinned reference implementation use the frozen target’s soft distribution. Replacing it with a one-hot token label is a possible alternative training program, not an interchangeable reading of the source. Soft-target cross entropy also equals $H(q^T)+D_{KL}(q^T\Vert p^D)$; because the frozen target fixes $H(q^T)$, minimizing it with respect to the drafter minimizes the forward KL divergence.

The combined EAGLE-1 objective is:

$$\mathcal{L}_{EAGLE} = v_w\mathcal{L}_{reg} +p_w\mathcal{L}_{cls}.$$

The paper and pinned implementation use $v_w=1.0$ and $p_w=0.1$. The 0.1 weights classification, not feature regression. For a synthetic example with $\mathcal L_{reg}=0.6$ and $\mathcal L_{cls}=2.4$, the reference weighting gives $0.6+0.1(2.4)=0.84$. These coefficients are source-reported defaults, not a claim that they transfer unchanged to another target, feature tap, or reduction convention.

0x04C Frozen-Target Training Trace

target_mask = segment_causal_mask(token_valid, segment_id, position_id)
with no_grad():
    target_out = target(
        tokens, attention_mask=target_mask, position_ids=position_id,
        output_hidden_states=True, use_cache=False,
    )
    target_feature = select_manifest_feature(target_out) # [B, S, d]

rows = compile_soft_rows(
    tokens=tokens, features=target_feature,
    token_valid=token_valid, segment_id=segment_id,
    position_id=position_id, supervise_at_source=loss_mask,
)
# rows retain legal context even where rows.supervise is false.
if any(rows.supervise):
    optimizer.zero_grad()
    predicted_feature = draft(
        rows.source_feature,
        target.embed(rows.advanced_token),
        attention_mask=rows.causal_mask,
        position_ids=rows.position_ids,
    )
    draft_logits = target.lm_head(predicted_feature)
    with no_grad():
        target_logits = target.lm_head(rows.feature_target)
        target_probability = softmax(target_logits.float(), dim=-1)

    loss_reg = masked_smooth_l1(
        predicted_feature, rows.feature_target, rows.supervise,
    )
    loss_cls = masked_distribution_ce(
        draft_logits.float(), target_probability, rows.supervise,
    )
    loss = 1.0 * loss_reg + 0.1 * loss_cls
    loss.backward()
    optimizer.step()
# Otherwise skip the update, including optimizer weight decay.

Memory warning: materializing target logits over [B,S,V] can dominate training memory. A practical implementation can compute only needed rows, use vocabulary-parallel loss, or use hard labels when that matches the intended objective.

The reference training pipeline can also consume cached target hidden states rather than rerun the target inside every drafter update. That optimization makes the dataset semantic:

cached sample:
    input_ids
    exact target feature rows
    loss_mask / segment ownership
    target checkpoint + feature-layer fingerprint

collation:
    pad tokens and feature rows together
    shift advanced tokens and feature targets together
    invalidate the one synthetic shifted row at each sample end
    apply padding validity and source-row supervision before loss reduction

Cached features from a different target revision are stale labels, not reusable activations. The checkpoint must also restore the learned fusion projection, draft decoder, optimizer state, scheduler state, loss weights, feature address, and compatible frozen embedding/head identities.

0x04D Gradient And Dtype Ownership

The frozen target may participate in the forward graph, but it owns no trainable gradient path:

for parameter in target.parameters():
    assert parameter.requires_grad is False

assert all(parameter.grad is None for parameter in target.parameters())
assert any(parameter.grad is not None for parameter in draft.parameters())

Target and draft matrix multiplications may use BF16, while softmax, log-softmax, loss accumulation, and normalization denominators should use a declared stable accumulation dtype. The pinned training program uses AdamW, betas=(0.9, 0.95), BF16 acceleration, and gradient clipping at 0.5; these are implementation observations, not universal EAGLE requirements.

0x04E Noise Injection

At inference, the drafter consumes its own imperfect predicted features after the first step. Native one-step training sees target features instead. EAGLE adds feature noise during training as a partial robustness mechanism:

$$\tilde f_t^T = f_t^T + \epsilon_t$$

This broadens the training neighborhood, but random noise is not guaranteed to match the structured error distribution of autoregressive draft features. HASS later targets that mismatch directly.

The EAGLE paper reports uniform noise in $[-0.1,0.1]$. At the pinned repository commit, the code samples (rand - 0.5) * 0.2 * 512 / sequence_length; it matches that interval at sequence length 512 and rescales it otherwise. A reproduction must state whether it follows the paper-level distribution or the pinned length-scaled implementation. Quietly calling both “uniform 0.1 noise” hides a real training difference.

0x05 Drafting And Verification Runtime

0x05A Bootstrap

EAGLE cannot draft before it has a target feature and a sampled token. A speculation round begins with target state from prefill or the preceding verification:

target prefix forward:
    produces committed target feature f_t
    produces next-token distribution q_{t+1}

sample:
    x_{t+1}

EAGLE seed:
    [f_t, embed(x_{t+1})]

This first target computation is not removed by EAGLE.

0x05B Autoregressive Tree Drafting

The drafter can expand several candidate nodes at one tree depth in one pass with tree attention. Each node carries:

token id
draft probability
predicted feature
parent index
logical depth / position id
draft-KV row reference

0x05C Target Verification

The selected tree is flattened, assigned ancestry-aware positions and a tree mask, and passed through the target model. M7.2 owns the tree-mask mechanics; M7.1 owns exact accept/reject/correction.

EAGLE’s distribution-preservation claim depends on that target-governed verification. Accurate feature prediction alone does not provide exactness.

0x05D Feature And KV Commit

Verification produces several kinds of state, but acceptance does not grant all of them the same lifetime. Treat publication as one transaction:

The final row is the subtle one. The verifier can emit a target-governed frontier token after accepting a path, but if that token was not one of the verified input rows, its target KV and feature must be computed by a subsequent valid target append before later decoding depends on them. Output commitment and target-cache publication are related transactions, not identical counts.

def publish_verified_path(target_state, draft_state, verified, decision):
    accepted_rows = decision.accepted_flat_rows

    target_state.append_kv(
        gather_rows(verified.target_kv, accepted_rows)
    )
    target_state.append_features(
        gather_rows(verified.target_features, accepted_rows)
    )
    target_state.append_output(decision.accepted_candidate_tokens)

    if decision.has_frontier_token:
        target_state.append_output([decision.frontier_token])
        target_state.mark_append_required(decision.frontier_token)

    target_state.discard_unaccepted_branch_rows()
    draft_state.discard_unaccepted_tree_rows()
    draft_state.reseed(
        target_feature=target_state.last_verified_feature,
        advanced_token=target_state.last_output_token,
    )

After the transaction:

commit:
    accepted output-token sequence
    target KV for accepted candidate rows
    accepted target features needed to seed the next draft round

discard:
    rejected target branch KV
    unaccepted candidate metadata
    draft branches that no longer extend the committed prefix

The target feature used for the next round should come from a valid target forward, not from an unverified draft feature.

The target-KV row map and target-feature row map should be derived from the same flatten/retrieve metadata used by verification. Reconstructing either map from token equality is invalid because repeated token IDs can occupy different tree nodes and positions.

0x06 EAGLE-2: Context-Aware Dynamic Draft Trees

Static trees assign the same branch budget to every context. EAGLE-2 starts from two observations reported by the paper:

1. acceptance varies by tree position;
2. acceptance at the same position also varies substantially by context.

The draft model already produces confidence scores. EAGLE-2 uses them as a low-cost proxy for acceptance.

0x06A Path Value Derivation

A node can be accepted only if every node on its path is accepted. If conditional acceptance probabilities along the path are $a_j$, the path acceptance probability is:

$$P(\text{node }u\text{ accepted}) = \prod_{j \in \operatorname{path}(u)} a_j$$

EAGLE-2 approximates $a_j$ with draft confidence $c_j$:

$$\operatorname{value}(u) \approx \prod_{j \in \operatorname{path}(u)} c_j$$

Because $0 \le c_j \le 1$:

$$\operatorname{value}(u) \le \operatorname{value}(\operatorname{parent}(u))$$

That monotonicity explains a useful property: selecting globally high-value nodes, with shallow-node tie breaking, preserves ancestors before descendants and therefore produces a connected tree.

For numerical stability, accumulate in log space:

$$\log \widehat V(u) =\sum_{j\in\operatorname{path}(u)} \log\max(c_j,\epsilon).$$

Products underflow on long paths; their log ordering does not. This still does not turn $\widehat V$ into a calibrated target acceptance probability.

The connectivity claim needs its assumptions. For a child $u$ of parent $p$, $\widehat V(u)=\widehat V(p)c_u\le\widehat V(p)$. If $c_u<1$, the parent ranks strictly earlier. If $c_u=1$, a shallow-first tie rule ranks the parent first. Therefore every prefix of the deterministic global order is ancestor-closed. Adding ancestors after selecting $m$ nodes is not equivalent: it can exceed the verification budget.

0x06B Expansion And Reranking

Source: EAGLE-2, top portion of Figure 7. Expansion follows high-value frontier nodes, then reranking selects a connected verification budget before flattening. Draft confidence is a proxy for acceptance, not the verifier itself.

The seed is the context anchor: depth zero and log path value zero. It is not counted in the candidate verification budget. The following teaching interfaces keep expansion records distinct from tree nodes:

• draft_expand_in_parallel returns (parent, token, confidence, feature) records for the selected parents. Confidence is the conditional token probability under that parent’s prefix, not a score renormalized over the retained children. Omitted vocabulary mass is allowed.
• Tree.add returns the newly stored node, with its parent, token, feature, confidence, accumulated log value, depth parent.depth + 1, and stable node_id. Only these nodes become the next frontier.

def eagle2_tree(seed, expand_width, verify_budget, max_depth):
    from math import isfinite, log

    if type(expand_width) is not int or expand_width < 1:
        raise ValueError("expansion width must be a positive integer")
    for value in (verify_budget, max_depth):
        if type(value) is not int or value < 0:
            raise ValueError("verification budget and depth must be nonnegative integers")
    epsilon = 1e-12  # numerical ranking floor; never overwrite the stored probability
    tree = Tree(seed)
    frontier = [seed]
    if verify_budget == 0:
        return flatten_with_tree_mask([])

    for depth in range(max_depth):
        if not frontier:
            break
        # EAGLE-2 expands high-value nodes in the latest generated layer.
        expandable = sorted(
            frontier,
            key=lambda n: (-n.log_path_value, n.depth, n.node_id),
        )[:expand_width]
        children = draft_expand_in_parallel(expandable)

        next_frontier = []
        mass_by_parent = {}
        seen = set()
        for parent, token, confidence, feature in children:
            if not any(parent is node for node in expandable):
                raise ValueError("child parent must belong to the selected frontier")
            if not isfinite(confidence) or not 0 <= confidence <= 1:
                raise ValueError("confidence must be a finite probability")
            key = (parent.node_id, token)
            if key in seen:
                raise ValueError("duplicate token alternative under one parent")
            seen.add(key)
            mass = mass_by_parent.get(parent.node_id, 0.0) + confidence
            if mass > 1.0 + 1e-12:
                raise ValueError("sibling probability mass exceeds one")
            mass_by_parent[parent.node_id] = mass
            log_value = parent.log_path_value + log(max(confidence, epsilon))
            node = tree.add(parent, token, confidence, log_value, feature)
            next_frontier.append(node)

        frontier = next_frontier

    selected = sorted(
        tree.non_root_nodes,
        key=lambda n: (-n.log_path_value, n.depth, n.node_id),
    )[:verify_budget]
    assert len(selected) <= verify_budget
    assert ancestor_closed(selected)
    return flatten_with_tree_mask(selected)

EAGLE-2 does not train a separate tree-policy model. It changes expansion and reranking around the existing EAGLE drafter.

0x06C Numeric Tree Example

root
|- A confidence 0.60, path value 0.60
|  |- A1 confidence 0.80, path value 0.48
|  `- A2 confidence 0.20, path value 0.12
`- B confidence 0.40, path value 0.40
   `- B1 confidence 0.95, path value 0.38

Sibling probabilities sum to 1.00 at the root and under A. Under B, the shown child has mass 0.95; the remaining 0.05 belongs to omitted alternatives. These values are synthetic token probabilities, consistent with the confidence definition in EAGLE-2 §3.2, not measured acceptance rates.

With expansion width two and maximum depth two, all five displayed candidates are built. Global reranking gives A (0.60), A1 (0.48), B (0.40), B1 (0.38), then A2 (0.12). A four-node verification budget therefore selects [A, A1, B, B1] without adding ancestors afterward. Breadth-first order [A, B, A1, A2, B1] would instead spend its fourth slot on A2.

0x06D Confidence Is A Proxy

Draft probability is correlated with acceptance in the paper’s experiments. It is not mathematically identical to target acceptance probability. Calibration can shift with:

Source: EAGLE-2, Figure 6 in arXiv v1. On the reported Vicuna-7B/Alpaca coordinate, binned draft confidence is positively associated with observed acceptance. The plot motivates a cheap ranking proxy; it does not establish universal calibration across targets, samplers, domains, or runtimes.

model and finetune
temperature and truncation policy
domain
prompt length
quantization
draft precision

Runtime telemetry should bucket draft confidence against observed acceptance. A stale confidence-to-acceptance relationship can waste the dynamic-tree budget.

The paper’s source-reported deployment settings also separate expansion width, maximum depth, and final verification budget. They are not one tree_size knob. Larger expansion explores more candidate state; larger final budget adds target verification rows; greater depth extends the longest possible accepted path. M7.8 owns the production policy that chooses among those costs.

0x07 HASS: Align The Objective And The Context

HASS identifies two separate train/decode mismatches.

0x07A Objective Mismatch

The draft model does not need to reproduce the target distribution uniformly over every low-probability vocabulary item. Acceptance is driven especially by tokens that the target and decoding policy consider plausible.

Let $q(v)$ be the target distribution, $p(v)$ the draft distribution, and $\Omega_K(q)$ the target’s top-$K$ token set. A teaching sketch of top-$K$ distillation is:

$$\mathcal{L}_{topK} = -\sum_{v \in \Omega_K(q)} q(v)\log p(v).$$

This is the HASS paper’s Top-K objective: target probability mass is not renormalized over the selected set. If $Z_K=\sum_{v\in\Omega_K(q)}q(v)<1$, then

$$\mathcal L_{topK} =Z_K\left[-\sum_{v\in\Omega_K(q)} \frac{q(v)}{Z_K}\log p(v)\right].$$

Omitted mass therefore scales the example’s contribution as well as focusing the support. Renormalizing would erase that scale and define a different loss. HASS combines this objective with feature alignment using a separate weight; paper hyperparameters such as $K$, rollout count, and loss coefficient are source-reported coordinates, not constants in the EAGLE architecture.

0x07B Context Mismatch

Source: HASS, Figure 2. The blue training edge supplies a target feature at a future step, while the red decoding edge must supply a draft-generated feature because the target has not verified that step yet. The architecture is unchanged; the context distribution differs.

Native EAGLE training uses target features as context:

training step 1:
    target feature -> draft prediction

Autoregressive drafting uses its own prior features:

inference step 1:
    target feature -> draft feature 1

inference step 2:
    draft feature 1 + draft KV -> draft feature 2

inference step 3:
    draft feature 2 + accumulated draft KV -> draft feature 3

HASS unrolls multiple draft steps during training and adjusts attention so later training steps see the kind of inaccurate draft context they will receive at inference.

Source: HASS, Figure 3. Each additional training step replaces more target-derived context with draft-derived features while preserving the query/key dependency pattern expected by recursive decoding. The right column shows the corresponding inference states. This is why HASS needs a dependency-aware attention mask rather than simple Gaussian feature noise.

features = target_features
total = 0.0

for rollout_step in range(num_alignment_steps):
    predicted = draft.forward_with_alignment_mask(
        features,
        advanced_tokens,
        rollout_step=rollout_step,
    )
    total += topk_distillation(predicted, target_distributions)
    total += beta * feature_regression(predicted, target_features_shifted)
    features = splice_draft_features_into_next_context(predicted, target_features)

The exact mask determines which queries use draft-produced features and which keys/values remain target- or draft-derived. It must be tested as a dependency graph, not trusted because dimensions match.

The context transition is easier to audit as a row-ownership table:

Sibling or future rollout rows must remain invisible. A dense causal mask over the physically concatenated batch is not sufficient if that order places unrelated branches before one another.

def hass_alignment_mask(rows):
    mask = zeros([len(rows), len(rows)], dtype=bool)
    for query in rows:
        for key in rows:
            mask[query.row, key.row] = (
                key.is_legal_target_prefix_of(query)
                or key.is_draft_ancestor_of(query)
                or key.row == query.row
            )
    assert no_sibling_or_future_edges(mask, rows)
    return mask

0x07C Cost Boundary

HASS adds training work through repeated draft passes or equivalent packed computation. It keeps the inference architecture and draft cost unchanged relative to the underlying EAGLE-style runtime. Therefore:

training cost:
    higher than native one-step EAGLE training

inference cost:
    same draft architecture

intended gain:
    better alignment -> longer accepted prefix

In arXiv v1, the main reported setup uses Top-K $K=10$, loss coefficient 1.0, and three context-alignment steps; ablations explore other values. These are source-reported experiment settings. The portable contract is the unnormalized Top-K target mass and the rollout dependency graph.

0x08 EAGLE-3: Remove The Feature-Prediction Constraint

EAGLE-3 is a genuine architecture/objective change.

0x08A Why Top-Layer Feature Regression Can Limit The Drafter

The top feature is immediately mapped to next-token logits. Constraining a small drafter to reproduce that feature can force its capacity toward matching a target representation optimized for the next token, even though the drafter needs information useful across several speculative steps.

EAGLE-3 removes the feature-regression objective and trains directly for token prediction. The draft network still has hidden states; “direct token prediction” means those states are no longer required to match the target’s top feature.

This changes the optimization geometry. EAGLE-1 asks a small drafter to land near one teacher representation and then checks the induced vocabulary distribution. EAGLE-3 permits an unconstrained draft state $a_t$ as long as its projected token distribution is useful. The removed constraint is $\mathcal L_{reg}$; target verification and sampled-token conditioning remain.

0x08B Multi-Layer Feature Fusion

EAGLE-3 records low-, middle-, and high-level target features:

$$f_t^{low}, f_t^{mid}, f_t^{high} \in \mathbb{R}^{d}$$

It fuses them:

$$r_t = W_R \left[ f_t^{low};f_t^{mid};f_t^{high} \right]$$

with $W_R \in \mathbb{R}^{d \times 3d}$, then combines the fused representation with the sampled-token embedding for the draft decoder.

“Low, middle, and high” are semantic roles, not universal layer numbers. The exact three layer indices, their pre/post-normalization status, and any final normalization before the head belong in the target checkpoint manifest. The pinned repository’s target wrapper returns three configured hidden streams and the drafter applies a learned 3d -> d projection; another checkpoint may need different addresses while preserving the same paper-level contract.

The second fusion is also source-defined, not an unspecified implementation detail. If e(\hat{x}_{t+1}) \in \mathbb{R}^d is the embedding of the sampled token:

$$u_t = W_I \left[ r_t; e(\hat{x}_{t+1}) \right], \qquad W_I \in \mathbb{R}^{d\times 2d}.$$

The single-layer draft decoder consumes u_t and produces draft state \hat{r}_{t+1}; the shared target LM head maps that state to draft-token logits.

Source: EAGLE-3, Figure 5. EAGLE-3 fuses low-, middle-, and high-level target features and directly trains the draft path for token prediction across multiple steps. The diagram is the EAGLE-3 pipeline, not the EAGLE-1 feature-regression objective.

Tensor trace:

three target features:
    [B, S, d] each

concatenate:
    [B, S, 3d]

reduce:
    [B, S, d]

combine with sampled-token embedding:
    concatenate [B, S, 2d]
    project with W_I -> [B, S, d]

one-layer draft decoder:
    [B, S, d]

shared LM head:
    [B, S, V]

0x08C Training-Time Test

At inference, the first draft step can reuse target features. Later speculative steps cannot obtain target features for unverified tokens, so the drafter feeds back its own output state.

EAGLE-3 simulates that test-time recursion during training:

native training pass:
    target-derived fused state -> draft output 1

simulated test pass 2:
    draft output 1 -> draft output 2

simulated test pass 3:
    draft output 2 -> draft output 3

Tensor recurrence:

$$u_t^{(0)} =W_I[r_t;e(\hat{x}_{t+1})],$$ $$\hat{r}_t^{(s+1)} =D_\phi\left(u_t^{(s)},KV_\phi^{(s)}\right), \qquad \hat{p}_t^{(s+1)} =softmax(W_U\hat{r}_t^{(s+1)}),$$ $$u_t^{(s+1)} =W_I[\hat{r}_t^{(s+1)};e(\hat{x}_{t+s+2})].$$

Unlike EAGLE-1, EAGLE-3 does not force \hat{r}^{(s)} to regress to a target-model feature. It applies direct token supervision to the unrolled draft outputs. A code-level training sketch is:

fused_target = reduce_features(low, middle, high)  # [B, S, d]
state = fuse_token(fused_target, shifted_token_embedding(1))

loss = 0.0
for step in range(num_training_time_test_steps):
    state = draft_decoder(
        state,
        attention_mask=training_time_test_mask(step),
    )
    logits = compatible_frozen_or_loaded_head(state)
    target_probability = shifted_target_distribution(step + 2)
    loss += step_weight[step] * distribution_ce(
        logits,
        target_probability,
        rollout_valid_mask(step),
    )

    # Reference training advances the data-token sequence while feeding back
    # the drafter's own state. Sampling belongs to inference, not this loop.
    advanced_token = shifted_training_token(step + 2)
    state = fuse_token(state, frozen_embedding(advanced_token))

The paper’s training mask does not allow every simulated token to attend to every other simulated token. Original training-sequence keys retain their causal context; generated rows follow the parent relationships they will have during recursive drafting. A dense triangular mask over the concatenated simulated rows would train cross-branch dependencies that are absent at inference.

Source: EAGLE-3, Figure 6 in arXiv v1. Gray rows are native training data; blue and yellow rows are successive draft-state rollouts. Arrows encode logical context. Physical concatenation order alone does not define this graph.

Let $A^{(s)}(u,v)$ denote visibility from query row $u$ at simulated step $s$ to key row $v$. The portable rule is:

$$A^{(s)}(u,v)=1 \iff v\in\operatorname{TargetPrefix}(u) \cup\operatorname{DraftAncestors}(u)\cup\{u\}.$$

The direct token objective across rollout steps can be written as

$$\mathcal L_{E3} =\sum_{s=0}^{R-1}\omega_s \left[-\frac{1}{\sum_t m_t^{(s)}} \sum_t m_t^{(s)}\sum_v q_t^{T,(s)}(v) \log p_t^{D,(s)}(v)\right].$$

There is no feature-regression addend. The pinned training code uses seven passes and weights its per-pass token losses by $\omega_s=0.8^s$; it also uses a target soft distribution over a configured draft vocabulary. Those are pinned implementation choices, not requirements that every EAGLE-3 reproduction use seven steps, the same decay, or the same vocabulary filter.

At the pinned commit, the code constructs target-to-draft vocabulary maps (t2d and d2t) and trains a draft-vocabulary head. That optimization changes checkpoint compatibility and must be recorded. The paper-level statement is multi-layer target-state fusion plus direct token prediction; a reduced draft vocabulary is an implementation contract, not part of the mathematical definition.

The attention mask encodes the simulated tree dependencies. This resembles HASS’s concern with train/decode context, but the objectives differ:

0x09 Runtime Integration

An EAGLE-capable serving loop needs more state than ordinary speculative decoding:

while not done:
    seed = get_target_feature_and_sampled_token(target_state)

    draft_tree = eagle_draft(
        seed_feature=seed.feature,
        seed_token=seed.token,
        draft_kv=draft_state.kv,
        tree_policy=dynamic_tree_policy,
    )

    packed = flatten_tree_with_positions_and_mask(draft_tree)
    verified = target_verify_tree(target_state, packed)
    result = exact_tree_accept_reject(verified, draft_tree)

    publish_verified_path(
        target_state=target_state,
        draft_state=draft_state,
        verified=verified,
        decision=result,
    )

That compact loop should be reconciled against concrete source responsibilities rather than treated as one opaque generate call. At pinned SafeAILab/EAGLE commit cb7e0841fe0c206c6ed74a197ad5e2a1f13f5a2b:

The full state transition is:

committed target prefix
  -> target feature tap + next-token distribution
  -> sampled advanced token
  -> EAGLE drafter state and provisional drafter KV
  -> dynamic candidate graph
  -> flatten rows + tree positions + ancestry mask + retrieve map
  -> target logits + target features + provisional target branch KV
  -> target-governed acceptance
  -> gather accepted rows
  -> publish output / target KV / target feature under distinct rules
  -> reseed or append frontier token before the next round

An executable documentation oracle can use tiny arrays instead of a model:

tokens:       [11, 12, 13, 14, 15, 16]
valid:        [1,  1,  1,  1,  1,  0]
coordinates:  zero-based storage indices in this fixture
soft rows:    4 (source indices 0, 1, 2, 3; all supervised here)
hard rows:    3 (an explicitly different, observed-token objective)
first row:    f_0 + token 12 -> f_1 + q_2
last soft:    f_3 + token 15 -> f_4 + q_5
              q_5 comes from f_4; token slot 5 is invalid and is not read

loss fixture:
    L_reg = 0.6, L_cls = 2.4
    L = 0.6 + 0.1 * 2.4 = 0.84

tree fixture:
    root probabilities A=0.60, B=0.40
    path values A=0.60, A1=0.48, B=0.40, B1=0.38, A2=0.12
    budget 4 -> [A, A1, B, B1]
    selected set is ancestor-closed without post-hoc repair

publication fixture:
    accepted flatten rows [0,2]
    target KV rows [0,2]
    frontier token emitted separately and marked target-append-required

The corrected row oracle retains four soft-target rows from these five valid tokens, including the final teacher distribution. Separate tests cover source-row supervision masks, padding, packed segments, and zero eligible rows. The loss, tree, and publication fixtures illustrate separate contracts; this is not a model-training, acceptance, or speed measurement.

Scheduler-visible work includes:

draft rows
draft-tree nodes
draft KV bytes
target verification rows
temporary target branch-KV bytes
accepted output progress
tree construction and mask time

Batching requests only by prompt length ignores major variation in speculative tree work.

0x09A MoE Interaction

For a dense model, target verification processes more token rows but generally reads the same layer weights. For an MoE target, candidate rows may route to a wider union of experts:

$$\mathcal{E}_{verify} = \bigcup_{u \in \text{candidate nodes}} \mathcal{E}_u$$

The verification pass can therefore increase expert-weight traffic and all-to-all fanout even when accepted progress improves. This is a cross-module M6/M7 cost, not an EAGLE correctness issue.

0x0A Paper-Variant Comparison

This table is a design comparison, not a leaderboard. EAGLE-3 was published after the original A33 article and is externally verified here so Transformer System extends the source rather than freezing it in time.

0x0B Failure Lab

0x0C Observability

Training:

feature regression loss
token classification/distillation loss
draft token accuracy and target top-K recall
feature error by rollout depth
loss by training-time-test step
draft-state norm and drift
target feature extraction time

Runtime:

draft confidence versus observed acceptance
tree nodes expanded / retained / verified
accepted length by depth
draft time
tree pack and mask time
target verify time
draft and target temporary KV bytes
union of MoE experts touched during verification
end-to-end latency and throughput by batch size

0x0D Boundary To M7.6

EAGLE trains a post-hoc drafter around a frozen target. DeepSeek MTP integrates sequential future-prediction modules into model pretraining. Both can support self-speculative decoding, but they differ in:

when the draft mechanism is trained
which hidden state is consumed
whether feature alignment is an objective
whether prediction depths share one draft block
which KV caches exist
how weights are packaged and served

M7.6 derives that integrated module chain.

M7.7 then compares feature drafting with model, retrieval, self-draft, and dependent-head proposers and derives hardware-aware candidate topology. M7.8 owns batch-dependent depth, proposal readiness, serving goodput, and release rollback.

0x0E Publish Gate

0xFF References

• Transformer Series A33 - DeepSeek MTP
• M7.1 - Speculative Decoding
• M7.2 - Medusa
• M7.7 - Draft Sources and Hardware-Aware Trees
• M7.8 - Speculative Serving
• EAGLE: Speculative Sampling Requires Rethinking Feature Uncertainty
• EAGLE arXiv v1 used for source and figure pinning
• EAGLE-2: Faster Inference of Language Models with Dynamic Draft Trees
• EAGLE-2 arXiv v1 used for source and figure pinning
• Learning Harmonized Representations for Speculative Sampling
• HASS arXiv v1 used for objective and context pinning
• EAGLE-3: Scaling up Inference Acceleration of Large Language Models via Training-Time Test
• EAGLE-3 arXiv v1 used for source and figure pinning
• EAGLE official implementation
• Pinned EAGLE reference implementation commit