M7.2 - Medusa: Multi-Head Candidates, Tree Attention, and Verification

M7.2 - Medusa: Multi-Head Candidates, Tree Attention, and Verification

0x00 Historical Introduction And Owner Contract

Owner question: How can one target backbone propose several future-token alternatives, pack their shared prefixes into one verification tree, and publish only a path whose tokens, positions, and K/V rows agree?

M7.1 established the exact one-sequence draft/verify contract. It deliberately left a harder construction problem unresolved. A separate autoregressive drafter proposes one path because each draft token becomes context for the next. A set of parallel future-token heads does not have that property: every head reads the same target hidden state, so its outputs are offset predictions, not one already-conditioned continuation.

Medusa turns that apparent defect into a candidate-set design:

one expensive target hidden state
    -> ordinary LM head: one root token
    -> Medusa head 0: alternatives for a later offset
    -> Medusa head 1: alternatives for another later offset
    -> ...

rank-index path templates
    -> prefix-closed packed tree
    -> one target verification pass with ancestry visibility
    -> one selected and committed path

The method is not merely “several LM heads.” Its useful unit is a round transaction coupling four contracts:

1. future-head logits have explicit token offsets and training labels;
2. a sparse path family compiles into mutually consistent token, mask, position, and retrieval buffers;
3. target logits are interpreted at the correct next-token coordinates under each root-to-node prefix;
4. output tokens and K/V rows are compacted from the same winning path before the next round.

0x00A Why A Second Full Draft Model Is Not The Only Proposal Source

Autoregressive decode repeatedly invokes the full target to expose one next-token distribution. A conventional speculative decoder adds a cheaper model that performs several serial proposal steps, then asks the target to score those tokens together. This can work well, but it introduces another checkpoint, tokenizer and processor compatibility, placement, cache, and scheduling surface.

Medusa asks whether candidate generation can instead reuse the target state already computed at the end of a round. The final hidden vector is semantically richer than a fresh token embedding and is already resident when the ordinary LM head runs. Small residual heads can therefore predict likely tokens at fixed future offsets without executing another transformer stack.

That reuse changes rather than removes the cost:

The result is attractive precisely when target decode is dominated by moving model weights and the extra tree rows use otherwise underfilled compute. It can be unattractive when large batches already saturate matrix units, when the head vocabulary projections are expensive, or when a wide tree verifies many nodes but accepts only a short prefix. Those runtime conditions are measured in M7.8; this article derives the mechanism that creates them.

0x00B The Lineage Changes One Contract At A Time

The dates below are source-qualified submission or publication coordinates, not a claim that no earlier related idea existed.

Blockwise Parallel Decoding supplies the multi-offset head ancestry. SpecInfer supplies the token-tree verification ancestry. Medusa’s contribution is their specific same-backbone composition plus a practical head-training, sparse-tree, acceptance, and cache-update pipeline. Similar execution shapes do not imply identical output-distribution contracts: greedy validation, exact rejection sampling, and Medusa’s fast typical acceptance must remain separate.

0x00C Future-Offset Heads Produce A Product Space, Not One Sequence

Let the committed prefix be x_{<=t} and its final target hidden state be h_t. The ordinary LM head predicts x_{t+1}. Zero-indexed Medusa head j is trained against x_{t+j+2} while still reading h_t:

$$q_j(v\mid h_t) = \Pr_{H_j}(x_{t+j+2}=v\mid h_t).$$

The notation q_j does not mean that all heads form one autoregressive proposal law. In general,

$$\prod_{j=0}^{K-1}q_j(c_{j+2}\mid h_t) \neq p(c_{t+2:t+K+1}\mid x_{\le t},c_{t+1}).$$

Head j+1 did not consume the candidate selected from head j. Taking one token from each top-k list therefore defines a rank tuple, not a verified language-model path. The tree enumerates a budgeted subset of those tuples; the target pass restores the missing conditional dependence by giving every node access to the prompt and exactly its ancestors.

This is the first conceptual firewall of M7.2:

head score product:
    candidate-ranking heuristic under fixed h_t

target path probability:
    autoregressive conditionals under the actual candidate ancestors

tree verification:
    computes the second object for a packed subset chosen using the first

0x00D One Tree Exists In Four Coordinate Systems

A path template such as (0, 1, 0) means “rank 0 at depth one, rank 1 at depth two, rank 0 at depth three.” It is not a token ID, packed row, logical position, or cache row. The implementation must compile the template family into four aligned buffers:

The buffers are derived from one sorted, prefix-closed path set. A shape-compatible permutation in only one buffer is a semantic corruption: the target may score one path, the posterior selector may read another, and the K/V compactor may publish a third.

Source-local visual: A31 - Medusa. The figure is retained because it places all four implementation coordinates beside one toy tree. Its Chinese annotations are supplemented by the table above and by executable fixtures later in this article.

0x00E Verify-And-Propose Creates A Recurrence

After prefill, a Medusa round can verify the current tree and obtain both ordinary target logits and new Medusa-head logits from the verified rows. The winning terminal row becomes the state from which the next candidate family is generated. Conceptually:

prefill -> root/head logits

repeat:
    materialize tree candidates from previous logits
    target verifies packed tree
    reconstruct candidate paths
    select accepted prefix
    compact matching K/V rows
    carry winning target/head logits into next round

This is why A31 describes the steady state as prefill + verify + verify + .... It does not mean proposal generation is free. Vocabulary projections, top-k selection, tree gathering, mask/position handling, posterior evaluation, and K/V compaction remain in the critical path. Loop 10 will account for each slot instead of collapsing them into one “verification” number.

0x00F Owner Contract: Consumes, Owns, Defers, And Forbids

Consumes

• M3.8 owns hidden-state to logit projection, processor ordering, and ordinary next-token semantics.
• M5.1 owns the general K/V lifecycle and the distinction between emitted tokens and materialized cache rows.
• M7.1 owns exact one-path speculative sampling, correction distributions, and transactional target-distribution preservation.

Owns

• Medusa residual-head architecture, offset labels, parameter/compute state, and independent-head conditionality gap;
• rank-path templates, sparse prefix closure, and compilation into the four tree buffers;
• ancestry visibility, logical positions, path reconstruction, and target-logit alignment;
• greedy and typical Medusa posterior selection as distinct contracts;
• Medusa-1, Medusa-2, self-distillation, checkpoint, and gradient-ownership contracts;
• the single-request winning-path K/V compaction transaction.

Defers

• M7.3 owns Jacobi/n-gram lookahead without added heads.
• M7.5 and M7.6 own feature-level and sequential MTP drafters.
• M7.7 owns Hydra/Clover dependent-head internals and hardware-aware dynamic tree optimization.
• M7.8 owns continuous batching, production measurements, admission policy, release, and rollback.

Forbids

• treating independent future heads as one autoregressive draft distribution;
• calling packed-row order a position order;
• calling fast typical acceptance target-distribution preserving;
• committing K/V by length without gathering the winning path’s exact rows;
• presenting source-reported speedups as local measurements;
• copying the general speculative-sampling proof or production scheduler into this article.

0x01 Claim Status And Scope

0x01A Immutable Medusa-Round Coordinate

A tree example is reproducible only after freezing more than num_heads and a list of tokens. Every later derivation in this article consumes the following coordinate:

The evidence fixture uses a small symbolic vocabulary and the same eight path choices as A31’s buffer example:

source commit:
    FasterDecoding/Medusa
    e2a5d20c048a9b0a4092e6933c34313687422518

committed prefix:
    token ids [101, 102, 103]
    next logical position 3

ordinary LM root:
    token 11 at logical position 3

Medusa head 0 top ranks:
    [21, 22, ..., 30] for logical position 4

Medusa head 1 top ranks:
    [31, 32, ..., 40] for logical position 5

path choices:
    (0), (1),
    (0,0), (0,1), (0,2),
    (1,0), (1,1), (1,2)

The token values are synthetic identities, not sampled English text. Their purpose is to make rank slots, packed rows, and logical positions distinguishable even when shapes agree.

0x01B Coordinate Domains Must Not Alias

For a two-head round, five integer domains coexist:

$$\begin{aligned} r &\in \{0,\dots,TOPK-1\} &&\text{candidate rank},\\ j &\in \{0,\dots,K-1\} &&\text{Medusa head index},\\ u &\in \{0,\dots,N_{tree}-1\} &&\text{packed tree row},\\ depth(u) &\in \{0,\dots,K\} &&\text{tree depth after the LM root},\\ s &\in \{0,\dots,S_{max}-1\} &&\text{cache sequence row}. \end{aligned}$$

The official candidate slot used by a path node at head/depth j and rank r is:

$$slot(j,r)=1+j\cdot TOPK+r,$$

where slot zero is reserved for the ordinary LM-root token. tree_indices[u] stores such a slot. It does not store u, depth(u), a token ID, or an absolute position.

The logical position of packed row u is instead:

$$position(u)=prefix\_length+depth(u),$$

with d(root)=0. Siblings share a position even though they have different packed rows. The cache destination is not known until a path is selected; only then do retrieve_indices[path, :] identify the staged rows that become a compact suffix.

0x01C Source Pin And Legacy/Current Firewall

A31 explains the legacy generate_medusa_buffers, tree_decoding, evaluate_posterior, and update_inference_inputs path. The official repository has continued to evolve. This article pins commit e2a5d20c048a9b0a4092e6933c34313687422518 and stores whole-file copies and hashes for:

medusa/model/utils.py
medusa/model/medusa_model.py
medusa/model/medusa_model_legacy.py
medusa/model/kv_cache.py
medusa/model/modeling_llama_kv.py
medusa/model/medusa_choices.py
medusa/train/train_legacy.py

The pin has two roles. It verifies that A31’s named buffers and legacy training offsets correspond to real implementation paths, and it makes clear when a statement is only about that coordinate. It does not claim that the pinned repository is the latest or preferred production integration in every framework.

0x01D Mutation Failures

Later oracles reject changes that preserve tensor shapes but alter semantics:

The manifest binds semantics, not only shapes. Loop 3 can now derive head tensors and cost from one declared model/head coordinate; Loops 4-8 can compile and transact one declared tree without guessing which integer domain a field inhabits.

0x02 Candidate Heads

Let the final hidden state at the current position be h_t. The ordinary LM head predicts x_{t+1}. With zero-indexed Medusa heads, head j predicts the token j+2 positions ahead:

$$\ell_t^{(j)} = H_j(h_t), \qquad H_j(h_t) \text{ targets } x_{t+j+2}.$$

LM head            token t+1
Medusa head 0      token t+2
Medusa head 1      token t+3
Medusa head j      token t+j+2

This offset convention matters in both training-label slicing and inference tree construction. Calling the first Medusa head the t+1 head would train it to duplicate the ordinary LM head.

0x02A Head Architecture And Tensor Trace

The released Medusa implementation does not use only a bare vocabulary projection. Each head is a small residual MLP followed by its own vocabulary projection:

$$z_t^{(j,0)} = h_t,$$ $$z_t^{(j,r+1)} = z_t^{(j,r)} + \operatorname{SiLU}\left(W_{j,r}z_t^{(j,r)}+b_{j,r}\right),$$ $$\ell_t^{(j)} = U_j z_t^{(j,R)}.$$

For hidden width d, vocabulary size V, K Medusa heads, and R_h residual blocks per head:

backbone hidden             h        [B, S, d]
one head residual state     z_j      [B, S, d]
one head logits             logits_j [B, S, V]
stacked Medusa logits       logits   [K, B, S, V]

A code-shaped version of the official architecture is:

class ResidualHeadBlock(nn.Module):
    def __init__(self, d):
        super().__init__()
        self.proj = nn.Linear(d, d)
        self.activation = nn.SiLU()

    def forward(self, x):
        return x + self.activation(self.proj(x))


medusa_head_j = nn.Sequential(
    *[ResidualHeadBlock(d) for _ in range(num_head_layers)],
    nn.Linear(d, vocab_size, bias=False),
)

The released code zero-initializes the residual block’s linear weight, so the residual path starts near the identity while the vocabulary projection learns the future-token mapping. This is an implementation choice, not a mathematical requirement of all multi-head speculation methods.

0x02B Independent Heads And The Conditionality Gap

Every head receives the same h_t. Head 1 predicts x_{t+3} without consuming the candidate chosen for x_{t+2}:

$$H_1(h_t) \neq p(x_{t+3}\mid x_{\le t},\hat{x}_{t+1},\hat{x}_{t+2}).$$

It is a direct future-offset predictor, not a small autoregressive draft model. Combining head top-k sets into paths therefore creates candidates from marginal future predictions that may be mutually inconsistent. Tree verification is what re-evaluates those combinations under the target model’s true path conditionals.

This explains a central design tradeoff:

Each head may produce top candidates:

head 1: [It, I]
head 2: [is, was, can]

The naive Cartesian product grows quickly. With branching factors b_i:

$$N_{nodes} = \sum_{i=1}^{K}\prod_{j=1}^{i} b_j$$

More candidates can raise the chance of a long accepted path, but every node adds verification work.

0x02C Candidate Budget Example

If each future offset keeps a top-k set, the tree can grow geometrically.

Example:

branching by depth:
    b = [3, 2, 2]

Node count:

depth 1: 3
depth 2: 3 * 2 = 6
depth 3: 3 * 2 * 2 = 12

N_nodes = 3 + 6 + 12 = 21

Those 21 candidate nodes are not 21 committed output tokens. They are verification work. A sparse tree construction policy chooses a budgeted subset of high-value nodes rather than blindly expanding every Cartesian product.

The practical policy question is:

Which candidate nodes maximize expected accepted length under a fixed verification budget?

That question belongs inside M7.2 because it determines the tree shape the target model must verify.

Let r_i be the chosen top-k rank at depth i, and let \pi_i(r_i) be that head’s probability for the ranked token. A simple independent-head path score is:

$$\log score(r_{1:\delta}) = \sum_{i=1}^{\delta}\log \pi_i(r_i).$$

This score is useful for ranking candidate templates, but it is not the target probability of the path because the heads do not condition on earlier chosen candidates. The Medusa paper therefore treats sparse-tree shape as an optimization problem informed by observed head accuracy and runtime cost. Its ablation reports that an optimized sparse 64-node tree can outperform dense 256-node trees in the evaluated setup; that is source-reported evidence for budget allocation, not a universal best tree size.

0x02D Two Output-Projection Contracts

The phrase “residual MLP followed by an LM head” hides a checkpoint-level choice.

Paper/current wrapper shape: each Medusa head owns R_h residual blocks and its own bias-free vocabulary projection U_j in R^(V x d).

$$H_j(h)=U_j\,Res_j^{(R_h)}(h).$$

Pinned legacy wrapper shape: each head owns only its residual blocks and all branches reuse base_model.lm_head.

$$H_j(h)=U_{base}\,Res_j^{(R_h)}(h).$$

Both compute a vocabulary projection for every active head. They do not have the same trainable-state or serialization contract.

The pinned repository contains both paths: medusa_model.py constructs per-head bias-free vocabulary projections, while medusa_model_legacy.py applies base_model.lm_head to each residual branch. A checkpoint saved for one representation cannot be loaded into the other merely because K, d, and V agree.

Source-local visual: A31 - Medusa. The equation shows the independent projection form. The pinned legacy implementation provides the shared-projection alternative described above.

0x02E Parameter, Resident-Memory, And FLOP Derivation

One residual block contains a dense [d,d] weight and a length-d bias:

$$P_{res}=d^2+d.$$

With R_h blocks per head, K heads, and independent bias-free vocabulary projections:

$$P_{heads}^{independent} =K\left[R_h(d^2+d)+Vd\right].$$

If all branches reuse the base LM head, the additional trainable parameters are instead:

$$P_{heads}^{shared}=KR_h(d^2+d).$$

This distinction is about additional resident/trainable state. It does not remove output-projection arithmetic. Under the article’s convention that one multiply-add is two FLOPs, projecting T_h hidden rows through all active heads costs approximately:

$$F_{heads} \approx 2KT_h\left(R_hd^2+Vd\right),$$

excluding SiLU, residual addition, top-k, and memory traffic. T_h=1 for proposal heads applied only to the winning terminal row; it can be the packed tree length if the implementation computes head logits for every verification row before gathering the winner.

For the A31-style example d=4096, V=32000, K=4, R_h=1:

one residual block:
    4096^2 + 4096 = 16,781,312 parameters

one independent vocabulary projection:
    32,000 * 4096 = 131,072,000 parameters

four independent heads:
    4 * (16,781,312 + 131,072,000)
    = 591,413,248 parameters
    ~= 1.102 GiB at two bytes/parameter

four shared-projection residual branches:
    4 * 16,781,312
    = 67,125,248 additional parameters
    ~= 128.03 MiB at two bytes/parameter

The shared projection saves checkpoint residency, not projection FLOPs. For one active row, four [V,d] projections alone perform 4 * 32,000 * 4096 = 524,288,000 multiply-accumulates. This is why “the heads are tiny” must be qualified by vocabulary size, projection sharing, number of active rows, and whether logits are materialized densely.

0x02F Source-To-Tensor Head Trace

For a training batch tokens[B,S], target hidden states are hidden[B,S,d]. Head j consumes the slice ending j+2 tokens before the label slice:

shift_j = j + 2

head input:
    hidden[:, :S-shift_j, :]          [B, S-shift_j, d]

head logits:
    H_j(hidden[:, :S-shift_j, :])     [B, S-shift_j, V]

labels:
    tokens[:, shift_j:S]              [B, S-shift_j]

The token-level alignment follows directly. Hidden row t summarizes tokens[:t+1]; the ordinary LM loss uses label t+1; Medusa head zero skips that ordinary target and uses label t+2; head one uses t+3.

For S=7 and K=3:

Later heads receive fewer supervised positions per fixed-length sequence:

$$N_j=B\max(S-j-2,0).$$

Loss weighting can counter different difficulty or scale, but it cannot create labels beyond the sequence boundary. Packing, padding, ignored labels, and conversation loss masks reduce N_j further and must be applied after the offset alignment. The executable head oracle records both geometric coverage and mask-filtered coverage so a shape-correct slice cannot silently train on padding or user tokens.

0x03 Tree Attention

A normal causal mask is not enough. Candidate paths must not leak information sideways.

The tree mask is:

$$M_{u,v}=1 \quad \text{iff} \quad v \in Ancestors(u) \cup \{u\}$$

That means:

a candidate node can see:
    the prompt prefix
    its own ancestor nodes
    itself

a candidate node cannot see:
    sibling branches
    descendants
    unrelated candidate paths

Source: Medusa, Figure 2. Candidate products from different future heads are flattened into one verification batch, while the tree mask permits each node to see only its ancestors. Position IDs must follow the same tree topology.

Without this mask, the verification pass is not checking the same conditional distributions as independent autoregressive paths.

0x03A Position IDs Are Part Of Tree Correctness

The mask determines which candidate ancestors a node can read. Position IDs determine where that node lives in the logical continuation. Siblings at the same depth need the same logical position even though they occupy different packed rows:

shared LM root:        root
packed path rows:      A, B, A1, A2, B1
path depth after root: 1, 1,  2,  2,  2
logical position id:   t+2, t+2, t+3, t+3, t+3

The ordinary LM candidate is the shared root at logical position t+1. The released implementation constructs medusa_position_ids from depth relative to that root, then adds the committed input length before the target forward. This is required for RoPE or absolute-position semantics. Giving packed row B position t+3 merely because it follows A in memory would make sibling paths represent different logical sequences.

0x03B Six-Node Mask Example

Consider this tree:

root
  A
    A1
    A2
  B
    B1

Flatten candidate nodes as:

[A, B, A1, A2, B1]

This is a six-token candidate tree if the root is counted. The root is the ordinary LM-head token and is unconditionally shared by all continuations; the committed prompt prefix precedes it. The matrix below has five post-root candidate rows because root visibility is handled separately. A valid candidate-to-candidate mask is:

A1 can see A, but not B, A2, or B1. This is why tree verification is not ordinary batching.

0x03C Target Verification Is A Rectangular Forward Contract

The square tree mask is only the lower-right block of the real attention relation. Let P be the number of committed prefix rows and N the number of packed tree rows. The target receives only the packed candidate tokens as new inputs, reuses the prefix K/V cache, and materializes N query rows against P+N key rows:

$$N_q=N,\qquad N_k=P+N,$$ $$A_{u,j}=1\quad (0\le j<P),$$ $$A_{u,P+v}=M_{u,v}.$$

Every candidate can therefore read the entire committed prefix. Inside the tree suffix, row u can read only the rows allowed by the ancestry mask M. In implementation terms the effective attention mask has shape [B,1,N,P+N]; the pinned LLaMA integration first constructs the ordinary causal mask and then replaces the final N x N block with the Medusa visibility contract. This is not equivalent to passing nine independent sequences: their shared prompt keys are reused once and their shared tree prefixes occupy one packed suffix.

The absolute position of row u is:

$$\pi_u=P+depth(u),$$

while its staged physical cache row is:

$$c_u=P+u.$$

These equations answer different questions. Siblings have equal depth(u) and therefore equal RoPE or absolute-position phases, but they need different physical rows because their K/V contents differ. Collapsing either coordinate into the other destroys a valid branch even though all tensor ranks still fit.

Rewrite visual grounded in A31 and the official Medusa tree_decoding and mask integration. Token identities are synthetic; the shapes and coordinate relations are implementation contracts, not a local latency measurement.

0x03D Target Logits Are Shifted Along Reconstructed Paths

A target logit row does not score the token fed into that same row. If packed row u_j contains path token x_j, its hidden state has already consumed x_j, so its ordinary LM logits score the successor:

$$\ell_{u_j}=W_{LM}h_{u_j}, \qquad \operatorname{softmax}(\ell_{u_j}) =p_T(\,\cdot\mid x_{<P},x_0,\ldots,x_j).$$

For retrieval path r_p=[u_0,u_1,\ldots,u_{D_{path}}], verification compares:

$$x_{j+1} \quad\text{against}\quad \ell_{u_j}, \qquad 0\le j<D_{path}.$$

The root token x_0 is different: it was selected from the ordinary target logits at the previous frontier. The tree call consumes that root and scores the first Medusa-proposed continuation. The official implementation makes the shift visible by gathering tree_logits[retrieve_indices], then comparing candidates[:,1:] with logits[:,:-1]. The terminal accepted row’s target and Medusa logits become the proposal frontier for the next round.

A three-token synthetic path makes the alignment concrete:

Comparing row 1 logits with input token 21 would be self-alignment, not next-token verification. It can produce plausible numbers while accepting the wrong event.

0x03E Staged K/V Preserves Packed-Row Identity

The target forward writes one staged K/V row per packed input row at every layer. Before posterior selection, these rows are branch-local speculative state:

committed rows:       [0, ..., P-1]
staged packed rows:   [P+0, P+1, ..., P+N-1]
logical positions:    [P+d(0), ..., P+d(N-1)]

retrieve_indices[p,j] is therefore an inverse coordinate shared by three later operations:

1. reconstruct candidate token path p;
2. gather target/head logits for that path;
3. identify the staged K/V rows eligible for compaction if that path wins.

Loop 6 proves only that those identities align. It does not decide which path is admissible, and it does not commit any staged row. Loop 7 owns posterior policy; Loop 8 owns the winning-path cache transaction.

0x03F Executable Verification Oracle

The article evidence package runs the A31 nine-row tree with P=5. It derives absolute positions [5,6,6,7,7,7,7,7,7], a conceptual rectangular mask of shape [1,1,9,14], six gathered paths, and unique staged rows [5..13]. A synthetic target surface makes path [11,21,31] the only two-edge greedy match. The oracle rejects:

packed-row positions substituted for tree depth
one prompt key hidden from a candidate query
sibling visibility leaked into the rectangular mask
logits compared with their own input row
two packed nodes assigned one staged K/V row
a retrieval path that is not a parent-child chain

This is a deterministic coordinate test, not model-quality or speed evidence.

0x04 Decode Pseudocode

hidden = target_backbone(prefix)
root_candidate = sample_or_argmax(target_lm_head(hidden))

head_candidates = []
for head in medusa_heads:
    head_candidates.append(topk(head(hidden), branch_factor))

tree = build_candidate_tree(root_candidate, head_candidates, path_templates)
mask = build_tree_attention_mask(tree)
position_ids = prefix_length + tree.depths

target_logits, branch_cache = target_verify_tree(
    prefix,
    tree.tokens,
    mask=mask,
    position_ids=position_ids,
)

accepted_path = select_accepted_path(
    tree,
    target_logits,
    rule="exact_or_typical",
)

commit_path_kv(branch_cache, accepted_path)
discard_unaccepted_branches(branch_cache)

The important phrase is exact_or_typical. The article must not imply all Medusa acceptance policies are exact speculative sampling. Medusa’s typical acceptance is designed to improve acceptance while maintaining generation quality, but it is not the same derivation as the correction distribution in M7.1.

Mask builder sketch:

def build_tree_attention_mask(nodes):
    # nodes are in packed tree order; each node stores parent_id.
    mask = zeros([len(nodes), len(nodes)], dtype=bool)

    for query in nodes:
        cur = query
        while cur is not None:
            mask[query.row, cur.row] = True
            cur = cur.parent

    return mask

The committed prompt prefix and shared LM root are handled separately as already-visible context. The candidate-to-candidate mask above is the part that prevents sibling leakage. Verification additionally needs a retrieval table mapping each root-to-leaf path back out of packed tree order before prefix acceptance is evaluated.

0x04A Sparse Tree Construction

The full Cartesian product of head candidates is usually too large. In the released design, a sparse tree is represented by rank-index path templates such as [0, 1, 0]: take rank 0 from the first candidate level, rank 1 from the second, then rank 0 from the third. Prefix closure is mandatory: if [0, 1, 0] exists, [0] and [0, 1] must also exist.

def build_candidate_tree(root_token, medusa_topk, path_templates):
    # root_token comes from the ordinary LM head.
    # medusa_topk[depth, rank] comes from the corresponding Medusa head.
    paths = sorted(path_templates, key=lambda p: (len(p), p))
    nodes = [Node(token=root_token, parent_row=None, depth=0)]
    row_for_path = {}

    for path in paths:
        assert len(path) == 1 or tuple(path[:-1]) in row_for_path

        depth = len(path) - 1
        token = medusa_topk[depth, path[-1]]
        parent = 0 if len(path) == 1 else row_for_path[tuple(path[:-1])]

        row_for_path[tuple(path)] = len(nodes)
        nodes.append(Node(token=token, parent_row=parent, depth=len(path)))

    return PackedTree(nodes=nodes, row_for_path=row_for_path)

The actual Medusa utility code turns these templates into four coupled buffers:

tree_indices          gather tokens from flattened per-head top-k candidates
tree_attention_mask   candidate-to-ancestor visibility
position_ids          logical continuation depth
retrieve_indices      unpack root-to-leaf paths for verification

Canonical row map

Let \mathcal{P} be the non-root path set and let:

$$\sigma(\mathcal{P})=sort(\mathcal{P}, key=(|path|,path))$$

be the canonical order. Define row(empty)=0 for the ordinary LM root and:

$$row(path)=1+index_{\sigma(\mathcal{P})}(path).$$

Every other buffer is a function of row. None may sort the paths again with a different key.

Candidate gather

For non-root path p=(r_1,...,r_\delta), its token comes from depth/head \delta-1 and final rank r_\delta:

$$tree\_indices[row(p)] = 1+(\delta-1)TOPK+r_\delta.$$

The same candidate slot can appear under several parents. In A31’s toy tree, head-1 ranks [0,1,2] are reused under both first-depth ranks, so slots 11,12,13 each occur twice. Duplicate gathered token IDs do not merge the nodes: their ancestors and packed K/V rows differ.

Visibility and position

For paths p_u and p_v, including the empty root path:

$$M_{u,v} = \mathbf{1}[p_v\preceq p_u],$$

where preceq means “is a prefix of.” The logical depth is:

$$depth_u=|p_u|, \qquad position_u=prefix\_length+depth_u.$$

This derivation makes three properties mechanical: the mask is reflexive, ancestor-transitive, and sibling-exclusive; position is constant within a depth; and packed row order never enters the position equation.

Retrieval paths and sentinels

retrieve_indices needs only maximal root-to-node paths: if a selected node is already a prefix of a deeper selected node, it appears in that deeper retrieval row and need not create a duplicate candidate sequence. For maximal path p:

$$retrieve(p) = [0,row(p_{1:1}),...,row(p_{1:|p|})].$$

Rows are padded to one width with a negative sentinel. The sentinel supports rectangular path tensors; it must be masked before token, logit, or K/V gather. Treating it as Python index -1 without an appended neutral element would read the final real tree row.

Executable inverse identities

The four-buffer oracle checks identities stronger than shape equality:

candidate_slots[tree_indices[u]] == packed_tree_token[u]

M[u,v] == 1
    iff path(v) is a prefix of path(u)

position_ids[u] == len(path(u))

packed_tree_token[retrieve_indices[path,d]]
    == token on that root-to-leaf path at depth d

It also verifies that every retrieved adjacent row is a parent-child edge and that every non-sentinel row can see the previous retrieved row. These are the bridge from a combinatorial tree to target conditionals.

A31 includes two implementation annotations that make the first three buffers concrete:

Source-local visual: A31 - Medusa. The left tensor is the candidate-rank structure produced per future head; tree_indices gathers a prefix-closed subset into packed node order. The green node numbers on the right are packed row identities, not token IDs.

Source-local visual: A31 - Medusa. Siblings occupy different packed rows but share a continuation depth. That repeated position ID is required before adding the committed prefix length; assigning positions by packed-row order would silently change RoPE phases and verification logits.

The construction policy can be static, searched offline, or chosen dynamically by a later method. The invariants are prefix closure and consistency among all four buffers. A tree mask generated from one ordering and retrieval indices generated from another is a silent correctness bug.

0x04B Calibration Probabilities And Expected Accepted Length

The regular Cartesian tree keeps the same number of ranks at every parent. That ignores two asymmetries measured by the Medusa paper:

earlier heads are usually more accurate than later heads
rank 0 is usually more accurate than rank 1, rank 1 than rank 2, ...

Let TopAcc_k(i) be top-i accuracy for future head/depth k on a calibration set. The probability mass assigned specifically to rank i is the increment:

$$a_k^{(i)} = TopAcc_k(i+1)-TopAcc_k(i),$$

with TopAcc_k(0)=0 under zero-based ranks. This is not the cumulative top-k accuracy itself. Using cumulative accuracy at every rank would count the same correct events repeatedly.

Under the paper’s simplifying independence assumption, the estimated probability that a rank path r_{1:\delta} matches through depth \delta is:

$$w(r_{1:\delta}) = \prod_{j=1}^{\delta}a_j^{(r_j)}.$$

For a prefix-closed node set T, the accepted-length tail identity gives the paper’s objective:

$$\widehat{\mathbb{E}}[A(T)] = \sum_{u\in T} w(u).$$

Why is the objective a sum over all nodes, not only leaves? A round accepts at least depth δ whenever one selected node at depth δ matches. Here δ is path depth, not hidden width d. The expectation of a nonnegative integer prefix length is the sum of its survival probabilities. Every retained node contributes one estimated survival event at its own depth; a leaf contributes through all of its retained prefixes.

The equality is a calibration model, not a theorem about independent heads. All heads read the same h_t; their errors can be correlated, and the target posterior policy can differ from exact token-match accuracy. The selected tree should therefore be versioned with its model, dataset, prompt distribution, acceptance policy, and node budget.

0x04C Greedy Frontier Search And Its Boundary

Prefix closure constrains selection. A deep node may be admitted only after its parent path is present. Starting from an empty non-root tree:

tree = set()
frontier = all_depth_one_rank_paths()

while len(tree) < node_budget:
    path = argmax(frontier, key=estimated_path_weight)
    frontier.remove(path)
    tree.add(path)

    if len(path) < num_heads:
        for next_rank in allowed_ranks[len(path)]:
            frontier.add(path + (next_rank,))

The marginal contribution of an admitted node is exactly w(path) in the paper’s objective. The resulting tree is intentionally unbalanced: strong early ranks receive more descendants, while weak ranks may remain shallow terminals.

Toy calibration:

exact-rank mass by depth:
    head 0: [0.55, 0.25, 0.10, 0.05]
    head 1: [0.50, 0.20, 0.10, 0.05]
    head 2: [0.45, 0.20, 0.10, 0.05]

budget:
    10 non-root nodes

The executable oracle selects the highest available weighted prefixes and records every frontier decision. It also proves two limits:

1. changing the calibration distribution can change the optimal static tree without changing the model architecture;
2. maximizing estimated accepted length per round is not the same as maximizing tokens per second, because nodes at different depths and widths can have different kernel and cache costs.

The second limit is the handoff to M7.7, which owns measured hardware objectives and dynamic tree policies. M7.2 owns the Medusa paper’s probability-only static search.

Source-local visual: A31 - Medusa. The visibly unbalanced tree is the intended result of spending a finite node budget on calibrated rank paths. The highlighted path is illustrative; the figure does not prove a universal tree shape.

0x04D Acceptance Starts After A Target-Selected Root

The ordinary LM head chooses x_0 from the target logits carried into the round. generate_candidates places that token at tree row zero, and every retrieval path begins with it. Posterior evaluation checks only candidates[:,1:] against logits[:,:-1]. Consequently:

published tokens in a round = 1 target-selected root + accepted suffix edges
released accept_length       = number of accepted suffix edges
minimum progress             = 1 token

This convention matters when comparing algorithms. Calling the root a Medusa-head acceptance overstates head accuracy. Counting it neither in output growth nor cache compaction creates an off-by-one state error.

For greedy decoding, suffix event j passes exactly when:

$$c_{j+1}=\arg\max_x p_T(x\mid prefix,c_0,\ldots,c_j).$$

The pass mask must be converted to a contiguous prefix:

$$a_p=\sum_{j=0}^{D_{path}-1}\prod_{r=0}^{j} \mathbf 1[c_{p,r+1}=\arg\max p_{p,r}].$$

A match after the first mismatch contributes nothing. The released greedy path selects the first path attaining max_p a_p; if no suffix edge matches, it defaults to path zero and still publishes the target-selected root. This is exact for the declared deterministic greedy policy. It is not a statement about preserving a stochastic target distribution.

0x04E Fast Typical Acceptance Derivation

Typical acceptance asks a different question:

Is this candidate token probable enough under the target posterior at this path position?

Let temperature-processed target probabilities be p_i and define entropy:

$$H(p_i)=-\sum_x p_i(x)\log p_i(x).$$

The paper and pinned implementation use:

$$\tau_i=\min\left(\epsilon,\alpha e^{-H(p_i)}\right), \qquad m_i=\mathbf 1[p_i(c_i)>\tau_i].$$

Source-local visual: A31 - Medusa. Its epsilon is the hard probability ceiling and delta corresponds to the scale called alpha in this article and posterior_alpha in the pinned code. The source image explains the criterion; it does not prove target-distribution preservation.

When the distribution is sharp, entropy is low and the entropy-dependent term is larger, so the condition can be stricter until capped by epsilon. When the distribution is diffuse, entropy is high, exp(-H) shrinks, and more plausible alternatives pass. The strict comparison in the pinned code is > rather than >=; reproductions must not silently change that boundary.

For each path, the accepted suffix length is again a cumulative product:

$$a_p^{typ}=\sum_{j=0}^{D_{path}-1}\prod_{r=0}^{j}m_{p,r}.$$

Among paths tied for maximal length, the fast path chooses the largest accepted target log-likelihood:

$$p^*=\arg\max_{p:a_p^{typ}=a^*} \sum_{j=0}^{a^*-1}\log p_{p,j}(c_{p,j+1}).$$

Summing logs is the product-of-probabilities comparison in a numerically stable form. Summing probabilities would define a different tie-break.

probs = softmax(path_logits[:, :-1] / temperature, dim=-1)
entropy = -(probs * log(probs + 1e-5)).sum(dim=-1)
threshold = minimum(epsilon, alpha * exp(-entropy))
candidate_prob = gather(probs, candidates[:, 1:])

passes = candidate_prob > threshold
accepted = cumprod(passes, dim=1).sum(dim=1)
best_length = accepted.max()
best_path = argmax(sum(log(candidate_prob[tied, :best_length]), dim=-1))

The 1e-5 is an implementation numerical guard, not part of the mathematical entropy definition. Fast typical acceptance is a quality/acceptance heuristic. It deliberately relaxes target-argmax equality and carries no M7.1-style correction proof.

0x04F Sampled Modes And The Exactness Firewall

The pinned implementation exposes four operational branches, but they do not share one probability theorem:

The non-fast branches sample a target token separately for each path/position from a truncated posterior, then accept candidate equality. Selecting the longest match across a candidate tree is not, by itself, a proof that the final sequence follows the declared target law.

The exact stochastic construction in M7.1 requires a normalized proposal distribution q over the same events, stored proposal probabilities, acceptance RNG, and the positive-part correction after rejection. Medusa’s product of independent future-head ranks is a tree-ranking heuristic; it does not automatically define the autoregressive q assumed by that proof. Exact multi-candidate selection is therefore deferred to M7.7.

One version-pinned caveat is concrete: at official commit e2a5d20c048a9b0a4092e6933c34313687422518, the non-fast typical branch calls get_typical_posterior_mask with one more positional argument than the helper signature accepts. That is an inspected implementation mismatch, not a claim that every Medusa release has the bug. A deployment must pin and exercise the actual branch it enables.

0x04G Executable Posterior Oracle

The evidence fixture reuses Loop 6’s six paths and supplies explicit normalized target distributions at the root and two child rows. It checks:

greedy winner:       [11,21,31], two suffix matches
fast-typical winner: [11,22,31], two threshold passes and best tied likelihood
root progress:       one token even when every suffix fails

It rejects a max in place of the threshold min, the wrong entropy sign, non-prefix pass counting, raw-probability tie-breaking, an exact label on fast typical acceptance, and an exact rejection contract without correction state. The fixture tests policy algebra; it does not claim that the synthetic probabilities reproduce a model or benchmark.

0x05 Medusa-1 And Medusa-2

Medusa-1 is the safer conceptual baseline: the original model is unchanged and the heads learn to propose candidates. Medusa-2 can improve candidate quality but requires careful training so the backbone’s original capability is not damaged. The names identify complete training contracts, not merely two optimizer flags.

0x05A Shifted-Label Objective

For zero-indexed Medusa head j, align hidden states and labels as:

head j input hidden:     h[:, :S-j-2, :]
head j target labels:    tokens[:, j+2:S]

Let m_{b,t}^{(j)} be one only where the shifted target is owned by the training objective: it is not padding, not beyond the right boundary, and not masked by a conversation-loss policy. Then:

$$N_j=\sum_{b,t}m_{b,t}^{(j)},$$ $$L_j=-\frac{1}{N_j}\sum_{b,t}m_{b,t}^{(j)} \log p_j(x_{b,t+j+2}\mid h_{b,t}),$$ $$L_{Medusa}=\sum_{j=0}^{K-1}\lambda_j L_j.$$

An empty N_j must skip or reject that head contribution; dividing by zero is not a valid batch loss. Later offsets are harder and have fewer labels. The paper reports using lambda_k=0.8^k for one-indexed head k; that is an experiment setting, not a universal weighting law or part of inference exactness.

hidden = backbone(tokens).last_hidden_state
loss_medusa = 0.0

for j, head in enumerate(medusa_heads):
    shift = j + 2
    logits_j = head(hidden[:, :-shift])
    labels_j = tokens[:, shift:]
    owned = shifted_loss_mask[:, shift:]
    if owned.any():
        loss_j = masked_cross_entropy(logits_j, labels_j, owned)
        loss_medusa += head_weight[j] * loss_j

Gradient ownership differs by variant:

Medusa-1:
    dL/d(head parameters)       yes
    dL/d(backbone parameters)   blocked/frozen

Medusa-2:
    dL/d(head parameters)       yes
    dL/d(backbone parameters)   yes, controlled by combined loss and optimizer recipe

Training ledger:

Medusa-1:
    train head losses at future offsets
    freeze target backbone and ordinary LM head
    optimizer owns only head parameters
    a quantized frozen backbone may reduce activation-generation memory

Medusa-2:
    train backbone and heads together
    include ordinary next-token objective to preserve base behavior
    use separate head and backbone/adapter optimizer groups
    release changed backbone or adapter with the heads

This keeps the deployment contract clear. Medusa-1 is an added candidate generator. Medusa-2 changes the model training process.

0x05B Direct Medusa-2: Loss, Optimizer Groups, And Warmup

Jointly minimizing only future-head loss would let the backbone change in ways that improve proposal heads while degrading the ordinary LM head. Medusa-2 therefore combines objectives:

$$L_{Medusa-2}=L_{LM}+\lambda_0L_{Medusa}.$$

The paper’s recipe adds three controls:

1. Combined loss: retain ordinary next-token cross entropy alongside future-head losses.
2. Differential learning rates: update new heads more aggressively than the already-trained backbone.
3. Head warmup/two-stage training: first train heads with the backbone frozen, then jointly train; alternatively ramp the relevant loss weight gradually.

The experiment appendix reports five one-layer heads, lambda_k=0.8^k, a cosine schedule with warmup and 8-bit AdamW, LoRA/QLoRA backbone updates, and a head learning rate four times the backbone learning rate. One direct Vicuna setting uses 2e-3 for heads, 5e-4 for the backbone adapter, and lambda_0=0.2. These are source-reported coordinates, not defaults that transfer automatically to another model, data distribution, or batch size.

The two-stage form first fits heads with the backbone frozen, then opens the controlled backbone/adapter update. A smooth alternative ramps the relevant joint-training weight. Optimizer ownership should be auditable:

optimizer = AdamW([
    {"params": medusa_head.parameters(), "lr": lr_head},
    {"params": backbone_adapter.parameters(), "lr": lr_backbone},
])

assert lr_head > lr_backbone
assert ordinary_lm_loss_is_enabled

These are not editorial details. They determine whether “Medusa-2” means a useful native multi-head model or a backbone whose original distribution has drifted.

0x05C Self-Distillation Has Two Different Targets

Future-token heads should be trained on sequences representative of the target model’s output distribution. When the original SFT/RLHF data is unavailable, the Medusa paper proposes self-distillation:

seed prompts
    -> target model generates responses
    -> generated conversations become shifted-label head data

For Medusa-1, generated tokens can train the frozen-backbone heads. For Medusa-2, changing the backbone on only hard generated labels can reduce quality. The paper therefore gives the two parameter groups different targets:

$$L_{backbone-distill} =\sum_t KL\left(p_{teacher,t}^{(0)}\,\|\,p_{student,t}^{(0)}\right),$$ $$L_{self-distill} =L_{backbone-distill}+\lambda_0L_{Medusa}.$$

The future heads still learn shifted hard tokens from generated conversations. The backbone/adapter learns the original model’s probability distribution, not only those hard generated labels. With a LoRA backbone update, the teacher can be recovered by disabling the adapter and the student by enabling it, avoiding a second full resident model.

The paper explicitly prefers LoRA without quantization for this self-distillation path: otherwise the adapter-off teacher is itself quantized and may reduce generated quality. That warning must not be inverted into a claim that QLoRA is always wrong; it is scoped to the paper’s teacher-recovery contract.

This creates three distinct training variants:

The generated dataset also has an identity: seed domain, prompt extraction, chat template, decoding policy, target revision, and random seed determine the distribution the heads see. “Self-generated” is not a sufficient data manifest.

0x05D Checkpoint ABI And Pinned Loader Drift

A Medusa checkpoint is not only a head tensor file. A loader needs enough identity to reconstruct the graph and its offset meaning:

The pinned official source exposes two version-specific compatibility hazards. Its legacy trainer saves medusa_lm_head.safetensors, while the checked loader paths request medusa_lm_head.pt. A fallback path in the current model also contains a hard-coded five-head assignment with a TODO about uploaded config. These findings do not generalize to later releases, but they prove why artifact format and head count must be validated before strict=False loading. A loose state-dict load can silently leave a head random or absent.

A head-only artifact is sufficient for Medusa-1 only when it names the exact unchanged base. Medusa-2 must additionally release the changed backbone or its adapter; otherwise the head checkpoint is paired with the wrong hidden-state distribution.

0x05E Executable Training And Artifact Oracle

The evidence fixture uses B=2, S=8, and three future heads. It shifts an assistant-only loss mask with each label offset, derives per-head owned counts, applies source-reported weights [0.8,0.64,0.512] to synthetic mean losses, and evaluates three variant contracts:

Medusa-1:             heads train; base frozen
direct Medusa-2:      LM CE + 0.2 * Medusa loss; separate LR groups
self-distilled M2:    teacher KL + 0.01 * Medusa loss; LoRA teacher off/on

It rejects offset-one training, an unshifted mask, an empty-label divide, backbone gradients in Medusa-1, missing LM/KL preservation in Medusa-2, equalized optimizer groups, a head-only Medusa-2 release, and a filename/config ABI mismatch. Synthetic losses verify arithmetic and ownership only; they are not a training run.

0x06 Winning-Path K/V State

Medusa shares the prompt prefix. Candidate branches are temporary.

prefix KV:
    already committed request state

tree KV:
    temporary verification state for candidate nodes

accepted path:
    committed after verification

rejected branches:
    discarded

For homogeneous full K/V rows at uniform scalar width, the temporary branch code-payload ledger is:

$$M_{tree\_KV,temp} \propto N_{nodes} \cdot L \cdot 2H_{kv}d_h \cdot bytes$$

The exact allocated shape depends on implementation, but the direction is clear: tree size is not free.

Numeric payload-only example:

N_nodes = 21
L = 32
H_kv = 8
d_h = 128
bytes = 2

per node all-layer K/V = 32 * 2 * 8 * 128 * 2
                       = 131,072 bytes
                       = 128 KiB

temporary tree K/V payload ~= 21 * 128 KiB
                            ~= 2.625 MiB

Here bytes=2 is bytes per scalar, not bits; the 2.625 MiB example remains a dense payload illustration. It excludes allocator slack, masks, logits, routing metadata and implementation-specific reuse. It is a pressure ledger, not a profiler result. Generic admission instead consumes the actual cache schema and incremental union of allocations: codes, qparams, residual/sparse state, page rounding and retained segments by layer, counting shared allocations once. M8.6 owns quantized representation costs, M5.8 the MLA representation, and M5 retention/publication/lifetime. Workspace remains a disjoint charge.

0x06A Accepted Suffix Length Is Not Published Row Count

Let p* be the selected candidate path and a the accepted suffix-edge length returned by the released posterior evaluator. The target-selected root is also published, so the number of tokens and staged rows in the transaction is:

$$G=a+1.$$

For 0 <= j < G, the winning packed row, its physical staged source, and its compact destination are:

$$u_j=retrieve\_indices[p^*,j],$$ $$s_j=P+u_j, \qquad d_j=P+j.$$

The token suffix is candidates[p*,:G]. The K/V suffix is the all-layer gather from physical rows [s_0,...,s_(G-1)]. Updating the cache length to P+G without this gather would expose whichever packed branches happen to occupy contiguous rows, not necessarily the winning path.

0x06B Worked Gather-And-Compact Trace

Use the Loop 7 fast-typical winner:

prefix length P:       5
best path p*:          4
accepted suffix a:     2
published count G:     3
candidate tokens:      [11, 22, 31]
retrieve rows:         [0, 2, 6]
physical source rows:  [5, 7, 11]
compact destinations:  [5, 6, 7]
new current length:    8

The pinned implementation first materializes tgt with advanced indexing, then copies that snapshot into the compact destination slice. Snapshotting matters because source and destination ranges can overlap. An element-by-element in-place rewrite could overwrite a source row before it is read.

G = accepted_suffix_length + 1
path_rows = retrieve_indices[best_path, :G]
source_rows = prefix_length + path_rows

# Advanced/index-select gather creates the all-layer K/V snapshot.
accepted_kv = kv_backing[..., source_rows, :]
destination = kv_backing[..., prefix_length:prefix_length + G, :]
destination.copy_(accepted_kv)

input_ids = concat(input_ids, candidates[best_path, :G])
current_length.fill_(prefix_length + G)

Rewrite trace grounded in A31 and the pinned official update_inference_inputs/KVCache.copy paths. Numeric token and row identities come from the synthetic evidence coordinate, not a model execution.

0x06C The Terminal Row Carries The Next Frontier

The transaction must also select the target and Medusa-head logits produced after the final accepted input token. With suffix length a:

$$\ell_{next}=path\_logits[p^*,a:a+1],$$ $$\ell^{Medusa}_{next}=path\_head\_logits[:,p^*,a:a+1].$$

Those distributions are conditioned on exactly the token history whose K/V rows were compacted. The next round chooses a new target root and future-head candidates from this shared frontier. Taking logits from another path, or from a-1, makes proposal state disagree with committed cache state.

Unlike an M7.1 rejection correction that may be emitted before it has been processed as an input, every Medusa token published by this path was already an input row in the tree verification call. Its target K/V therefore exists in the staged cache and can be gathered immediately.

0x06D Publication Invariants And Failure Semantics

After a successful update:

$$|output_{new}|-|output_{old}| =|KV_{new}|-|KV_{old}| =G,$$ $$current\_length=P+G.$$

The following fields must describe one immutable publication plan:

best path identity
accepted suffix length
published token IDs
retrieval and source rows
compact destination rows
target frontier logits
Medusa-head frontier logits
current cache length and emitted-token counter

The pinned reference function mutates token IDs, backing K/V storage, cache length, and counters sequentially. The invariants above define a logical transaction, not an exception-safe ACID transaction. A production runtime must validate indices and capacity before mutation or provide rollback if an operation can fail. It may leave rejected bytes in preallocated backing storage; decreasing the logical current_length makes them unreachable, but does not securely erase them.

0x06E Executable Commit Oracle

The evidence oracle snapshots the synthetic staged rows, commits path 4, and checks the exact source-to-destination map [5,7,11] -> [5,6,7]. It also proves that the terminal frontier is packed row 6 and rejects:

using a instead of a+1 rows
committing a contiguous packed suffix by length
using token path p but K/V path q
carrying frontier logits from the predecessor row
advancing current_length before the gather plan is validated
aliasing source and destination without a snapshot

The runtime commit path is therefore a gather/compact operation, not merely a length update:

# accepted_length counts suffix edges; include the target-selected root.
path_rows = retrieve_indices[best_path, :accepted_length + 1]
accepted_kv = gather(tree_kv, rows=path_rows)
append_compact(prefix_kv, accepted_kv)
discard(tree_kv)

The same chosen path must drive output tokens, target logits for the next iteration, Medusa-head logits for the next candidate set, and K/V compaction. Using best_path for tokens but a different retrieval order for K/V creates a valid-looking output followed by corrupted state.

0x07 Runtime Pipeline And Cost Sensitivity

0x07A Pinned Steady-State Call Graph

After prefill, the checked implementation repeats this source-visible call graph for batch size one:

carried target/head logits
    -> generate_candidates
         target root selection
         K head top-k operations
         tree_indices gather
         retrieve_indices path materialization
    -> tree_decoding
         N-row target backbone forward with staged K/V
         ordinary LM projection on N rows
         K Medusa head paths on N rows
         retrieve_indices logit gather
    -> evaluate_posterior
         greedy or typical/nucleus path checks
    -> update_inference_inputs
         append G tokens
         gather/compact G K/V rows
         carry terminal target/head logits
    -> repeat

The recurrence removes a separate proposal-model forward, not all proposal work. Candidate selection, vocabulary top-k, tree gathers, output projections, posterior softmax/entropy, and cache copies remain.

0x07B One-Round Cost Ledger

Separate the target backbone from output and policy work:

$$\begin{aligned} C_{round}={}&C_{topk}+C_{backbone}(P,N) \\ &+C_{LMproj}(N)+C_{heads}(K,T_h) \\ &+C_{path}(N_{path},D_{path}) \\ &+C_{posterior}(N_{path},D_{path})+C_{commit}(G). \end{aligned}$$

Here P is committed prefix length, N packed tree rows, K future heads, T_h rows on which heads execute, N_path reconstructed paths, D_path retrieval width, and G=a+1 published tokens. The pinned model applies every Medusa head to all verification hidden rows, so T_h=N. A terminal-row-only implementation would have T_h=1, but that is an alternative implementation, not the checked call path.

Under 2 FLOPs/MAC, independent vocabulary projections cost approximately:

$$F_{LMproj}=2NdV,$$ $$F_{head-proj}=2KT_hdV,$$

and R_h square residual blocks per head add:

$$F_{head-residual}=2KT_hR_hd^2.$$

The legacy shared-LM-head parameterization saves additional projection weights, but still invokes the [V,d] projection for each head output. It does not make F_head-proj vanish.

For a dense masked attention implementation, target attention score/value work scales with all physical pairs:

$$E_{dense}=N(P+N).$$

The mathematically allowed pairs are:

$$E_{allowed}=NP+\sum_{u,v}M_{u,v}.$$

The difference is exploitable only if the backend has a tree-sparse execution path. Applying -infinity in a dense score matrix enforces correctness but does not prove arithmetic was skipped.

Posterior materialization also has a distinct shape. The pinned path gather creates [N_path,D_path,V] target logits even when many paths share a packed row. Fast typical acceptance then performs probability/entropy work over roughly N_path(D_path-1)V elements. A unique-row posterior kernel could reduce duplication, but it is not equivalent to pretending path reconstruction is free.

Source-local visual: A31 - Medusa, derived from the Medusa paper’s complexity discussion. It distinguishes one-row decode from q-row parallel decode. Symbols and big-O terms are source-reported; the table is not a profile of this site or hardware.

0x07C K/V And Logit Materialization Ledgers

For the same homogeneous full-KV illustration, per packed node payload-only all-layer K/V is, with $b$ in bytes per scalar:

$$B_{KV/node}=L\cdot2H_{kv}d_hb.$$

Therefore:

$$B_{KV,temp}=N B_{KV/node}, \qquad B_{KV,commit}=G B_{KV/node}.$$

The second quantity is logical copy payload, not necessarily interconnect traffic or incremental allocation. Quantized qparams/exception segments, latent caches and aliasing must use the actual representation/allocator receipt above, not the scalar-width product. Allocator layout and fused cache operations can change actual bytes moved. Dense target and head logits have resident/materialized element counts:

$$N V\quad\text{and}\quad K T_h V,$$

while gathered path logits contain RDV elements and may duplicate the same underlying packed-row logits.

0x07D Break-Even Condition And Sensitivity

The Medusa paper separates acceleration rate, per-round overhead, and wall-time speedup. In the same notation:

$$overhead=\frac{C_{round}}{C_{baseline-step}},$$ $$speedup=\frac{\mathbb E[G]}{overhead}.$$

Thus the conditional break-even rule is:

$$\mathbb E[G]>overhead.$$

This equation is bookkeeping, not a universal predictor. C_round varies with prefix length, tree shape, output kernels, dtype, batch composition, cache layout, and backend. E[G] varies with model/head calibration, data domain, temperature, acceptance policy, and tree budget.

A wider tree can fail in two independent ways:

1. marginal nodes do not increase accepted depth enough;
2. they increase target/head/logit/cache work enough to erase the added depth.

The pinned generator asserts batch size one. Continuous batching adds ragged tree lengths, request-specific acceptance, and scheduler coupling; those are owned by M7.8, not inferred from the single-request loop.

0x07E Executable Cost And Sensitivity Oracle

The evidence fixture evaluates the A31 nine-row tree with P=2048, K=4, d=4096, and V=32000. It derives dense versus ancestry-allowed attention pairs, target/head/residual projection FLOPs, dense and gathered logit elements, temporary and committed K/V payload, and terminal-only versus pinned all-row head execution. A separate normalized profile demonstrates the break-even equation under explicitly synthetic costs.

The oracle rejects omitted head projections, treating weight sharing as zero runtime projection work, treating a mathematical mask as sparse-kernel proof, deduplicating gathered paths without an implementation, omitting commit bytes, using rounds rather than published tokens in the speedup numerator, claiming batch support beyond the pinned assertion, and labeling synthetic sensitivity as measured latency.

0x08 Paper Contracts And Failure Closure

0x08A Lineage By Contract

Blockwise Parallel Decoding, SpecInfer, Medusa-1, and Medusa-2 share pieces of an execution shape, but they do not expose the same proposal, topology, verification, training, or publication contract. Reading the methods as a single progression hides the exact state that changes at each step.

The comparison is deliberately by contract rather than by reported speedup. Blockwise Parallel Decoding supplies the future-offset/longest-prefix lineage. SpecInfer demonstrates expanded or merged token trees with parallel target verification. Medusa composes same-backbone heads, a sparse static rank-tree, Medusa-specific posterior policies, and a cache-commit transaction. None of those facts makes the corresponding output-policy proofs interchangeable.

0x08B Dependent-Head Alternatives Belong To M7.7

Medusa’s heads are parallel because every head reads the same base hidden state. That also makes their future-token predictions conditionally independent of the earlier token choices in the candidate continuation. Hydra conditions later draft heads on embeddings of preceding drafted tokens. Clover propagates a regressive draft state with attention to sequential token knowledge. These alternatives can improve proposal consistency, but they add a sequential dependency to the proposal critical path.

That contrast is enough to define the M7.2 boundary. M7.7 owns Hydra/Hydra++, Clover, their training choices, proposal probabilities, dynamic tree selection, and hardware-aware topology optimization. Importing those derivations here would duplicate that owner and obscure Medusa’s independent-head contract.

0x08C Failure Matrix Bound To Executable Guards

The evidence package does not leave these failures as prose warnings. Each row is bound to an oracle or publication-audit guard. The error code names the semantic invariant, not merely a tensor shape.

0x09 Boundary To Later M7 Articles

Medusa accelerates by adding heads.

Lookahead accelerates without adding heads or a separate draft model, using Jacobi/n-gram machinery. MTP and EAGLE use training-time or feature-level mechanisms to improve candidate generation.

All three still need target-governed verification and correct KV commitment.

M7.7 compares Medusa’s independent heads with Hydra/Hydra++ and Clover, derives multi-candidate selectors, and moves tree shape from a fixed template to a measured hardware objective. M7.8 owns ragged tree commits under continuous batching and production rollback.

0x0A Publish Gate

0xFF Reference

• Transformer Series A31 - Medusa
• M7.1 - Speculative Decoding
• M7.7 - Draft Sources and Hardware-Aware Trees
• M7.8 - Speculative Serving
• Blockwise Parallel Decoding for Deep Autoregressive Models
• Medusa, ICML 2024 canonical publication
• Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads
• Official Medusa implementation, pinned commit e2a5d20c
• SpecInfer: Accelerating Generative Large Language Model Serving with Tree-based Speculative Inference and Verification
• Hydra: Sequentially-Dependent Draft Heads for Medusa Decoding
• Clover: Regressive Lightweight Speculative Decoding with Sequential Knowledge