TL;DR: Flatten the tree into nodes, let node attend to the whole cached prefix plus exactly its ancestors and itself, and give it position id prefix plus depth. A 10-node tree (root plus 9 candidates) raises expected tokens per step from 2.77 to 3.28 over a 4-token chain in the example below, but verifies 10 tokens per request instead of 5, so it stops being free at about batch 29 on an H100 instead of 59.
How to approach it
Say what the tree buys first: a chain fails as soon as one guess is wrong, while a tree keeps a second candidate at the positions where the drafter is unsure. Then define the data structure (parent pointers, depth, node order), derive the mask from it, and only then talk about kernels. Ask the batch size early, since it decides how wide a tree is worth building.
A strong answer
The tree. Write each candidate as a path of ranks: is the drafter's top guess for the next token, its second guess, the second guess for the token after the top guess. A typical static tree spends more nodes under high-ranked parents:
The mask. Within the step, node may attend to node only if is an ancestor of or itself, because the candidates on other branches are alternatives, not context. Every node also attends to the full cached prefix. Position ids are the prefix length plus depth, so siblings share a position: nodes 1 and 2 both sit at depth 1. Build it from parent pointers:
import numpy as np
def build_tree(paths):
"""paths: candidate-rank tuples, e.g. (0,), (1,), (0, 0). Node 0 is the root
(last accepted token). Returns parent [N], depth [N], mask [N, N] (bool)."""
paths = sorted(set(paths), key=lambda p: (len(p), p))
index = {(): 0}
parent, depth = [-1], [0]
for p in paths:
if p[:-1] not in index:
raise ValueError(f"missing parent for {p}")
index[p] = len(parent)
parent.append(index[p[:-1]])
depth.append(len(p))
n = len(parent)
mask = np.zeros((n, n), dtype=bool)
for i in range(n):
j = i
while j != -1: # each node attends to itself and its ancestors
mask[i, j] = True
j = parent[j]
return np.array(parent), np.array(depth), mask, paths
def accept_greedy(parent, tokens, target_argmax):
"""tokens[i]: draft token at node i; target_argmax[i]: target's next token
after node i. Walk from the root while some child matches."""
node, out = 0, []
while True:
kids = [c for c in range(len(parent)) if parent[c] == node]
nxt = next((c for c in kids if tokens[c] == target_argmax[node]), None)
out.append(int(target_argmax[node])) # accepted token, or correction/bonus
if nxt is None:
return out
node = nxt
For this tree, row 9 of the mask is true at columns 0, 1, 3, 6 and 9, and row 5 at 0, 2 and 5. The tests also cover an empty tree (a 1-by-1 mask, plain decoding), a path whose parent is missing (rejected), a walk that takes the second-rank branch, and a walk to the deepest leaf that emits all four drafts plus the bonus. The full step mask is of shape , which a FlashAttention-style kernel takes as a custom mask or, for speed, as the parent array it expands per tile.
What it buys. Under greedy verification, let be the probability that the target's choice is the drafter's top candidate and that it is the second (illustrative). A node is reached with the product of its ranks' probabilities, and expected tokens per step is one plus the sum over nodes:
chain of 4 top guesses: 1 + 0.7 + 0.7^2 + 0.7^3 + 0.7^4 = 2.77 tokens, 5 verified
tree above: 1 + [0.7 + 0.15] + [0.49 + 0.105 + 0.105] + [0.343 + 0.0735 + 0.0735]
+ [0.2401] = 3.28 tokens, 10 verified
gain: 3.28 / 2.77 = 1.18x more tokens per step for 2x the verified tokens
What it costs. At batch the target processes tokens per step. An H100 decode step stays bandwidth-bound up to roughly 295 tokens (989 TFLOP/s dense BF16 divided by 3.35 TB/s, for 2-byte weights):
break-even batch, tree: 295 / 10 = 29.5 -> about 29
break-even batch, chain: 295 / 5 = 59
batch 64: tree 640 tokens per step, chain 320; both past the ridge, tree by 2.2x
Past that point every unused branch is paid-for compute. The tree also leaves a cleanup job: node KV was written in node order, so if nodes 0, 2 and 5 are accepted, their entries must be moved to contiguous positions before the next step, or the engine writes tree KV to scratch and copies the accepted path in.
The decision reverses with load: wide trees at batch 1 to roughly 16, shrinking toward a chain as the batch approaches the ridge, and a dynamic tree (EAGLE-2 expands nodes by draft confidence) when the drafter's confidence is informative enough to spend nodes where they will be accepted.
What interviewers probe next
"Sampling at temperature 0.8 instead of greedy?" Verification becomes recursive rejection sampling over siblings: try each child with the residual distribution updated after every rejection, which keeps the output exact.
"Why not give siblings different positions?" RoPE would then rotate a depth-1 candidate as if it were at depth 2, and the target would score it in the wrong place.
"Does CUDA graph capture survive trees?" Only with a fixed tree shape per batch bucket, which is one reason static trees remain common in production.
How it's scored
Weak: builds a causal mask over the flattened nodes, so candidates attend to their siblings.
Hire: derives the ancestor mask and position ids from parent pointers and computes expected accepted length for a given tree.
Strong Hire: prices the tree as verified tokens against the ridge point, handles the accepted-path KV compaction, and shrinks the tree as batch grows.
Common mistakes
- Using sequential position ids for the flattened tree, which breaks rotary embeddings.
- Forgetting that every node must see the full prefix, not just the tree.
- Quoting tree gains measured at batch 1 for a deployment at batch 64.
- Leaving rejected-branch KV in place and reading it on the next step.