Counting Nodes in a Draft Tree

Medusa builds its draft as the Cartesian product of each head's top- $s_k$ predictions. With the first head contributing its top 2 and the second head its top 3, the tree carries $2 \times 3 = 6$ candidate continuations, and the target model verifies the whole tree in one forward pass using an attention mask built from the tree's shape.

How many draft tokens enter that forward pass, and what does the mask enforce?

5 tokens, 2 from the first head plus 3 from the second, since the heads' predictions are concatenated into one draft sequence.

8 tokens, but each candidate needs its own forward pass, so this tree costs 6 verification passes.

8 nodes: 2 at depth one plus 6 at depth two. The mask lets each node attend only to its own ancestors and the accepted prefix, with position indices set by depth, so all 6 continuations are scored in one pass without seeing each other.

6 tokens, one per candidate continuation, since only the leaves are what get verified.

8 tokens under a standard causal mask over the flattened tree, which is why deeper trees lose accuracy.