TL;DR: Run the target once over the last accepted token plus the drafts, take its argmax at all positions, accept the longest prefix where draft equals argmax, and emit one more token from the target: the correction at the first mismatch or the bonus after a full match. Every step emits between 1 and tokens, 3.36 on average at , and the KV written for rejected drafts must be discarded.
How to approach it
Draw the verify pass before writing code: which tokens go in, which logit rows come out, and which row decides what. Then ask two questions that change the implementation: is the batch ragged (different draft counts per request, which happens with n-gram drafters that sometimes find nothing), and must the engine stop on EOS inside an accepted run? Both turn a one-liner into a batched kernel with masks.
A strong answer
The verify forward pass takes input tokens per request: the last committed token and drafts . Causal attention means the logit row at input predicts the token after it, so the target produces its own argmax in one pass. Draft is correct exactly when , and a mismatch at position invalidates everything after it, because was conditioned on the wrong .
| Position | 1 | 2 | 3 | 4 | 5 |
|---|---|---|---|---|---|
| draft | a | b | c | d | |
| target argmax | a | b | x | y | z |
| outcome | accept | accept | reject, emit x | discard | discard |
That step emits three tokens: a, b and the correction x. If all four drafts had matched, the row at position 5 would still be a valid target prediction conditioned on a fully correct prefix, which is the bonus token, for tokens. One observation makes the code short: the accepted drafts equal the target's argmax, so the emitted tokens are always where is the length of the matching prefix.
import numpy as np
def verify_greedy(draft, target_argmax, num_draft, eos_id=None):
"""draft: [B, G] proposed ids (padded), target_argmax: [B, G+1] target's
argmax at each verified position, num_draft: [B] real drafts per row.
Returns (tokens [B, G+1], n_emit [B]); row b emits tokens[b, :n_emit[b]]."""
B, G = draft.shape
valid = np.arange(G)[None, :] < num_draft[:, None]
match = (draft == target_argmax[:, :G]) & valid
n_acc = np.cumprod(match, axis=1).sum(axis=1) # leading run of matches
n_emit = n_acc + 1 # correction or bonus
tokens = target_argmax.copy() # accepted drafts equal the target's argmax
if eos_id is not None: # stop at the first EOS
is_eos = (tokens == eos_id) & (np.arange(G + 1)[None, :] < n_emit[:, None])
first = np.where(is_eos.any(axis=1), is_eos.argmax(axis=1) + 1, G + 1)
n_emit = np.minimum(n_emit, first)
return tokens, n_emit
Tested cases: a mismatch at position 3 emits 3 tokens; a full match emits 5; a mismatch at position 1 emits only the correction; a row padded to 2 real drafts that both match emits 3, with the bonus read from position 3 rather than 5; an EOS inside the accepted run truncates there; and a row with zero drafts degrades to ordinary decoding, one token. A 200,000-row simulation at per-position match probability 0.8 averaged 3.360 tokens per step, against the closed form:
E[tokens] = (1 - alpha^(gamma+1)) / (1 - alpha)
alpha 0.8, gamma 4: (1 - 0.8^5) / 0.2 = (1 - 0.32768) / 0.2 = 3.36
alpha 0.6, gamma 4: (1 - 0.6^5) / 0.4 = 2.31
alpha 0.9, gamma 4: (1 - 0.9^5) / 0.1 = 4.10
bounds: 1 token (first draft wrong) to gamma + 1 = 5 tokens (all right)
Now the KV cache. The verify pass wrote KV for all five inputs at positions to . With two drafts accepted, the entries for , a and b are correct; the entries for c and d were computed from a wrong prefix and must not be read again. Nothing needs zeroing: the engine rolls the request's computed-token count back to , and the next step's first input (the correction x) overwrites the slot c occupied. Blocks allocated for the speculative tail are kept for the next step rather than freed and reallocated.
When does the "lossless" claim fail at temperature zero? The guarantee is that the output is the target's greedy continuation as computed by the verify pass. That pass runs a different shape (five query tokens per request, a different batch composition) than plain decode, so kernels can pick different reduction orders. On a near-tie between two logits, the argmax can flip. Spec-on and spec-off greedy outputs therefore agree on almost every token but not on every one unless the engine is batch-invariant, and a regression test that diffs them bit for bit will flake.
What interviewers probe next
"Now temperature is 0.7." Replace the equality test with acceptance probability and resample rejections from the normalized ; with a greedy draft, is one-hot and the acceptance probability is just .
"Why not stop the verify pass at the first mismatch?" All positions are computed in one parallel pass, so there is nothing to stop; the cost is the same whether one draft or all four survive.
"Where does this run, CPU or GPU?" On the GPU, as a fused kernel over the batch, because copying logit rows to the host would add a sync and megabytes of traffic per step.
How it's scored
Weak: emits only accepted drafts and loses the extra target token, or checks each draft against the target's prediction from the previous step.
Hire: gets the off-by-one between input positions and logit rows right, emits the correction or the bonus, and handles ragged draft counts.
Strong Hire: explains the KV rollback without zeroing, handles EOS inside the accepted run, and notes that greedy spec-on and spec-off outputs can differ on near-ties without batch-invariant kernels.
Common mistakes
- Accepting every draft that matches, not just the leading run, which commits tokens conditioned on a rejected one.
- Reading the bonus token from row for a request that only proposed two drafts.
- Freeing and reallocating KV blocks after each rejection instead of moving the computed-token pointer.
- Counting acceptance as "tokens per step" without the bonus, which understates mean accepted length by exactly one.