Implement the lossless speculative sampling verifier for a batch with ragged draft lengths, and prove it matches the target.

TL;DR: Accept draft token x at position i with probability min(1, p(x)/q(x)), stop at the first rejection, and sample the replacement from the normalized residual max(0, p − q); if every draft survives, sample a bonus token from the target at the next position. Vectorize with a cumulative product over the accept mask. With a context-free acceptance rate of 0.70 and 3 drafts, the expected output is 2.533 tokens per verification, and 400,000 rows of this code averaged 2.532.

How to approach it

Fix the interface before the math: the draft tokens and the draft probabilities q they were sampled from, the target probabilities p at every draft position plus one more, and the number of valid drafts per row, because a batched engine lets each row propose a different count. Then state what "lossless" means in a testable way: the tokens that come out must be distributed exactly as if the target had sampled them one at a time. The test has to measure that, which rules out tests that only check shapes or acceptance rates.

A strong answer

The rule for one position, written out, is where the residual comes from. The output token at a position has probability

$$P(x) = q(x)\,\min\!\left(1, \frac{p(x)}{q(x)}\right) + \left(1 - \sum_y \min(p(y), q(y))\right) \frac{\max(0, p(x) - q(x))}{\sum_y \max(0, p(y) - q(y))}$$

The first term is min(p, q). Both sums in the second term equal the total variation distance between p and q, so they cancel and leave max(0, p − q). Adding the two gives p(x). Getting the residual wrong (resampling from p, or from p − q without the clamp) breaks this identity.

import numpy as np

def sample_rows(dist, rng):
    """One inverse-CDF sample per row of dist [B, V] (rows sum to 1)."""
    u = rng.random((dist.shape[0], 1))
    return np.minimum((np.cumsum(dist, -1) < u * dist.sum(-1, keepdims=True)).sum(-1),
                      dist.shape[1] - 1)

def verify(draft, q, p, num_draft, rng):
    """draft [B, K] tokens sampled from q [B, K, V]; p [B, K+1, V] target probs at
    each draft position plus the bonus position; num_draft [B] <= K valid drafts.
    Returns out [B, K+1] (padded with -1) and n_out [B] = accepted + 1."""
    B, K, V = q.shape
    out = np.full((B, K + 1), -1, dtype=np.int64)
    if B == 0:
        return out, np.zeros(0, dtype=np.int64)
    idx = draft[..., None]
    p_tok = np.take_along_axis(p[:, :K], idx, -1)[..., 0]           # [B, K]
    q_tok = np.take_along_axis(q, idx, -1)[..., 0]
    u = rng.random((B, K))
    ok = (u * q_tok < p_tok) & (np.arange(K)[None] < num_draft[:, None])
    n_acc = np.cumprod(ok, axis=1).sum(1) if K else np.zeros(B, np.int64)
    rows = np.arange(B)
    p_j = p[rows, n_acc]                                            # [B, V]
    rejected = n_acc < num_draft
    q_j = np.where(rejected[:, None], q[rows, np.minimum(n_acc, K - 1)] if K else 0, 0)
    resid = np.maximum(p_j - q_j, 0.0)
    mass = resid.sum(-1, keepdims=True)
    dist = np.where(rejected[:, None] & (mass > 0), resid, p_j)     # p == q: no reject
    final = sample_rows(dist, rng)
    keep = np.arange(K + 1)[None] < n_acc[:, None]
    out[:, :K][keep[:, :K]] = draft[keep[:, :K]]
    out[rows, n_acc] = final
    return out, n_acc + 1

Details that matter in review:

1. Compare u * q < p, not u < p / q. Same rule, no division, and no special case when a greedy drafter puts probability 1 on its token.
2. The cumulative product finds the first rejection. cumprod(ok) is 1 up to the first false and 0 after, so its sum is the accepted count. Later positions that would have been accepted must not count, because their context was built on a rejected token.
3. Ragged rows. A row with num_draft = j treats position j as its bonus position. The caller fills p[:, j] with the target distribution after the row's j drafts, which it computed anyway in the verification forward pass.
4. The p = q corner. If the residual mass is zero then the rejection probability was zero, so the branch is unreachable in exact arithmetic; the guard falls back to p so float noise cannot produce a division by zero.

The test that proves losslessness generates sequences end to end. The target and the drafter are random Markov chains over a 4-token vocabulary with a start state, so p and q change with context. 300,000 sequences of 3 tokens each were produced by repeated draft and verify rounds with 3 drafts per round, and the joint distribution of the 3 tokens (64 cells) was compared with the exact product of the target's transition probabilities. The chi-square statistic has to stay below the upper 0.1% point for 63 degrees of freedom, about 103.5:

The negative control matters as much as the passes. It shows the test has the power to catch the most common bug, which shifts the output distribution by 18 percentage points of total variation and still produces plausible text. Smaller checks: with p equal to q every draft is accepted; with zero drafts the function samples one token from the target; an empty batch returns shape (0, K + 1); and with a one-hot target the accepted count equals the length of the run of drafts matching the target's argmax, with the replacement token being that argmax.

acceptance check, context-free p = [0.5, 0.3, 0.15, 0.05], q uniform over 4
alpha = sum min(p, q) = 0.25 + 0.25 + 0.15 + 0.05 = 0.70
E[tokens] = (1 - alpha^(K+1)) / (1 - alpha) = (1 - 0.7^4) / 0.3 = 2.533
measured over 400,000 rows: 2.532

The reversal condition: if the serving contract is greedy at temperature 0, none of the randomness is needed; accept while the draft equals the target argmax and emit the argmax at the first mismatch. Any sampling-based check then becomes an exact equality test against plain greedy decoding.

What interviewers probe next

"Now the drafter proposes a tree, not a chain." Walk the tree from the root with the same rule. When a child sampled from q is rejected, replace the target at that node with the normalized residual max(0, p − q) and test the next sibling against it, so the multi-candidate version stays lossless.

"The target applies top-p and temperature." Apply the same processing to p before verification, and to q as the drafter actually sampled. Losslessness is relative to the processed target distribution, so mismatched processing shows up as bias, not as lower acceptance.

"Why is a uniform random per position required?" Reusing one uniform across positions correlates the accept decisions. Swapping that single line into the code above and rerunning the joint test gave a chi-square of 42,096 and a TV distance of 0.068, a clear failure.

How it's scored

Weak: accepts when the draft equals a target sample, or resamples from p after a rejection, and tests by eyeballing outputs.

Hire: correct accept rule, residual, bonus token and first-rejection logic, vectorized over the batch.

Strong Hire: derives why the residual makes the output exactly p, handles ragged draft counts and the p = q corner, and proves losslessness with a joint-distribution test that includes a negative control showing the test can fail.

Common mistakes

• Counting accepted drafts after the first rejection, which keeps tokens conditioned on a context that no longer exists.
• Forgetting the bonus token when every draft is accepted, which wastes the target's extra position and lowers tokens per step from 2.533 to 2.19 in the example above.
• Verifying with probabilities from the unprocessed logits while the drafter sampled from processed ones.
• Testing only the first output token, which cannot detect correlated or position-dependent errors.