Vibe Engines Learn
CODING CHALLENGE · N°54

Speculative Decoding: Accept Step

Medium AIInferenceSampling

Speculative decoding makes LLM inference faster: a small draft model proposes several tokens, and the big target model verifies them in one pass. The accept/reject rule is a clever bit of probability that guarantees the output matches the target model’s own distribution. Implement it. Solve it in Python or TypeScript, with hidden tests.

The problem

A draft model proposed a run of tokens; for each, you have the target-model probability p[i] and the draft-model probability q[i] it was sampled from, plus a uniform random number r[i] in [0, 1). Accept token i if r[i] ≤ min(1, p[i] / q[i]); stop at the first rejection. Implement spec_accept(p, q, r): return how many proposed tokens are accepted before the first rejection.

EXAMPLE 1
Input p = [0.9, 0.2, 0.8], q = [0.5, 0.8, 0.4], r = [0.1, 0.5, 0.2]
Output 1
token 0 accepted (ratio ≥ 1); token 1 ratio 0.25 < r=0.5 → reject, stop
EXAMPLE 2
Input p = [0.5, 0.5], q = [0.5, 0.5], r = [0.3, 0.9]
Output 2
ratios are 1, so any r ≤ 1 accepts — both pass
CONSTRAINTS
  • Accept token i iff r[i] ≤ min(1, p[i] / q[i]). When the target likes the token at least as much as the draft (ratio ≥ 1), it is always accepted.
  • Stop at the first rejection — tokens after a rejected one are discarded (the target model resamples from there).
  • Return the count of accepted tokens (the length of the accepted prefix). This acceptance rule is exactly what makes the sped-up output identical in distribution to plain target-model sampling.
SOLVE IT YOURSELF

Your turn — write it

Edit the stub, hit Run (or ⌘/Ctrl + Enter), and watch the hidden tests. Stuck? the hints are right above and Reveal solution is one click away.

YOUR TASK

Implement spec_accept(p, q, r): iterate the proposed tokens; accept while r[i] ≤ min(1, p[i]/q[i]), and break on the first token that fails. Return the number accepted.

HINTS — 4 IDEAS
  1. For each token compute the acceptance threshold min(1, p[i] / q[i]).
  2. If the target probability is ≥ the draft probability, the ratio is ≥ 1, so the token is always accepted.
  3. Compare the given random draw: accept if r[i] ≤ threshold, otherwise reject.
  4. The moment a token is rejected, stop — return the count so far. Everything after is thrown away.
CPython · WebAssembly

Explore the topic

See this challenge alongside everything else on the same subject — handbooks, system designs, algorithms and tools, in one place.

LLM Engineering