Part 3: Speculative Decoding via Sampling-Mode Accept/Reject
Part 3 of a 4-part series on system-level LLM inference internals. Part 1 tracked entropy during decode. Part 2 measured attention sinks during prefill. This one implements sampling-mode speculative decoding: a small draft model proposes tokens, a large verifier checks them in one batched pass, and a probability-ratio test decides what survives.
The Core Idea
Text generation is fundamentally sequential. Each token's distribution depends on all the tokens before it. A large model doing this one token at a time is wasteful, since you're paying for a full forward pass just to get one token out.
Speculative decoding flips this around. A small, fast draft model proposes several tokens at once, then a large verifier model checks all of them in a single batched forward pass. The verifier's one pass sees every draft position at once, and that's where the efficiency win comes from. The hard part is making this mathematically sound, so the output distribution still matches sampling from the verifier alone, not some hybrid of draft and verifier.
| Part | What We Build |
|---|---|
| 1 | Per-token entropy tracker, visualized in real time |
| 2 | Attention sink detector, context health scoring |
| 3 (this post) | Speculative decoding: sampling-mode accept/reject |
| 4 | Empirical study: correlation plots across 50 prompts |
The Accept/Reject Mechanism
Drafting tokens is the easy part. The hard part is accepting or rejecting them in a way that provably preserves the verifier's distribution. For each draft token, the core test is a probability ratio:
accept_prob = min(1.0, p_verify / p_draft) if p_draft > 0 else 0.0
if torch.rand(1).item() < accept_prob:
accepted.append(token_id)
continue
1) If p_verify > p_draft, the verifier is more confident than the draft in this token, so accept_prob clamps to 1 and the token is always accepted.
2) If p_verify < p_draft, the verifier is less confident, so the token is accepted only with probability p_verify / p_draft. That's the exact rate needed to cancel out the extra mass the draft added.
On rejection, you don't just fall back to the verifier's argmax. You resample from what's left over:
draft_dist = torch.zeros_like(verify_probs)
draft_dist[token_id] = p_draft
residual = torch.clamp(verify_probs - draft_dist, min=0.0)
residual = residual / residual.sum()
correction = int(torch.multinomial(residual, num_samples=1).item())
residual is max(0, p_verify - p_draft), the slice of the verifier's probability mass that the draft's guess didn't already account for. Sampling from that residual, rather than from the verifier's raw distribution, is what makes the whole scheme add up correctly. Accept plus resample composes back to exactly p_verify, so the output is provably indistinguishable from sampling the verifier alone, even though the draft moved first.
A Worked Example
Say the draft samples "cat" at some position, with p_draft(cat) = 0.6. It's worth noting this is a sample, not an argmax. The draft draws randomly from its own filtered distribution, so a lower-probability token could just as easily have come out. This time it happened to be "cat."
The verifier then looks at that same position and checks what probability it would have assigned to "cat." It doesn't sample its own token here, it just reads p_verify(cat) off its distribution. Say that comes out to p_verify(cat) = 0.3.
accept_prob = min(1, 0.3 / 0.6) = 0.5
The verifier is less confident in "cat" than the draft was, so "cat" only survives half the time, decided by a single coin flip: torch.rand(1).item() < 0.5. Say it fails this time, so "cat" is rejected.
Now build the residual. Suppose the vocabulary is just {cat, dog, fox} for simplicity, and the verifier's distribution at this position is cat = 0.3, dog = 0.5, fox = 0.2:
draft_dist = { cat: 0.6, dog: 0, fox: 0 } # only "cat" has known draft mass
residual = { cat: 0, dog: 0.5, fox: 0.2 } # max(0, verify - draft), cat zeroed out
normalized = { cat: 0, dog: 0.71, fox: 0.29 }
"cat" is zeroed out of the residual entirely. It already had its shot during the accept/reject coin flip, so it can't come back as its own replacement, that would double count its mass. The correction token is sampled from {dog: 0.71, fox: 0.29}. Say "dog" comes out.
One more detail worth spelling out: rejection stops the round right there. If this was draft token 3 of 4, token 4 is discarded unchecked. It was never verified, and it was conditioned on the now-rejected "cat," so it no longer causally follows the corrected sequence. The next round's draft phase starts fresh from right after "dog," not from the discarded token 4.
If every draft token in the round survives, there's one more free token to collect. The verifier already computed logits one position past the last draft token, so sampling from those is essentially free:
bonus_logits = logits[prefix_len - 1 + gamma]
bonus_probs = top_p_filter(F.softmax(bonus_logits.float() / temperature, dim=-1), top_p)
bonus_token = int(torch.multinomial(bonus_probs, num_samples=1).item())
That's the payoff case: gamma + 1 tokens for the cost of one verifier forward pass.
Sampling vs Greedy
Two variants show up in the literature. Greedy decoding has the draft argmax its token and accept only if the verifier's argmax matches, which is deterministic and tends toward higher acceptance rates. Sampling mode is stochastic and uses the probability-ratio test above.
Since Part-1 used stochastic sampling at temperature = 0.7, top_p = 0.9, this implementation sticks with sampling mode to keep acceptance rates comparable across the series. The accept/reject formula min(1, p_verify(x) / p_draft(x)) is what the literature proves correct. It guarantees the output distribution equals the verifier's, token for token, regardless of what the draft proposed.
The Flow
Each round has three phases. The draft phase proposes gamma tokens one at a time from the small model, reusing a KV cache so each new token costs one incremental forward step rather than a full replay of the prefix:
output = draft_model(input_ids, use_cache=True)
cache = output.past_key_values
for i in range(gamma):
model_input = input_ids[:, -1:] if i == 0 else torch.tensor([[draft_ids_list[-1]]], device=device)
position_ids = torch.tensor([[prefix_len + i - 1]], device=device)
output = draft_model(model_input, past_key_values=cache, position_ids=position_ids, use_cache=True)
cache = output.past_key_values
logits = output.logits[0, -1]
token_id, p = sample_token(logits, temperature, top_p)
draft_ids_list.append(token_id)
draft_probs.append(p)
Each draft token conditions on the ones before it, and both the token ID and its probability under the draft's filtered distribution get carried forward. p_draft is needed later for the accept/reject ratio. The explicit position_ids matters here: HuggingFace models don't infer a token's absolute position from past_key_values alone, so an incremental call with cache but no position IDs would silently assume position 0 and corrupt every rotary embedding downstream. Passing prefix_len + i keeps attention and position encoding correct even though only one new token enters the forward pass.
The verify phase is a single batched call over the whole thing at once:
full_ids = torch.cat([input_ids, draft_ids], dim=1)
logits = verifier_model(full_ids).logits[0] # [seq_len, vocab]
This is the actual source of the speedup. Instead of gamma sequential verifier calls, one forward pass produces logits for every draft position simultaneously. The accept/reject phase then walks through those gamma positions, and the main loop just appends whatever survives before starting the next round.
One easy detail to miss: both models need to sample with the same temperature and top_p. The ratio p_verify(x) / p_draft(x) is only meaningful if both probabilities were computed under identical filtering, which is why sample_token and the verifier's inline softmax plus top_p_filter call share the same TEMPERATURE and TOP_P constants.
KV Caching on Both Sides
Recomputing a full forward pass over the entire sequence on every single-token step is the obvious thing to avoid. The model has already seen every prior token, so there's no reason to make it re-derive their key and value projections each round. Both the draft phase and the verify phase are built around a KV cache. Each side keeps its own past_key_values, and every subsequent call forwards only the newest token or tokens rather than the whole sequence so far.
Without caching, each draft token costs a full forward pass over everything generated so far:
# no cache: full recompute every iteration
for _ in range(gamma):
logits = draft_model(generated).logits[0, -1]
token_id, p = sample_token(logits, temperature, top_p)
generated = torch.cat([generated, torch.tensor([[token_id]], device=device)], dim=1)
With caching, the model only sees the newest token each step. Everything before it is already encoded in past_key_values:
output = draft_model(input_ids, use_cache=True)
cache = output.past_key_values
for i in range(gamma):
model_input = input_ids[:, -1:] if i == 0 else torch.tensor([[draft_ids_list[-1]]], device=device)
position_ids = torch.tensor([[prefix_len + i - 1]], device=device)
output = draft_model(model_input, past_key_values=cache, position_ids=position_ids, use_cache=True)
cache = output.past_key_values
logits = output.logits[0, -1]
token_id, p = sample_token(logits, temperature, top_p)
draft_ids_list.append(token_id)
draft_probs.append(p)
The verifier's cache works the same way across rounds. Instead of re-forwarding the entire generated sequence every round, it only forwards the newly proposed draft tokens on top of an already-cached prefix.
One subtlety worth calling out on its own: position_ids. HuggingFace models don't infer a token's absolute position from past_key_values alone. An incremental call with a cache but no explicit position info defaults to assuming the new token sits at position 0, which corrupts every rotary embedding downstream. Passing position_ids = prefix_len + i explicitly keeps the model oriented correctly even though it only sees one new token per call.
The effect is a straightforward complexity change. Cost per round scales with the number of new tokens (O(1) for the draft's incremental steps), not with how much has already been generated (O(prefix) for a full re-forward). Measuring both versions head to head on the same prompt and gamma:
No cache: 100 tokens in 14.9s (6.7 tok/s), 26 rounds
Cached: 100 tokens in 7.1s (14.2 tok/s), 36 rounds
Speedup from caching: 2.12x
Caching alone more than doubles throughput on this setup, and the gap only grows with sequence length, since the no-cache cost is quadratic in tokens generated while the cached cost is linear.
Results
With caching on both sides, here's a gamma sweep against the verifier-only baseline, asking for a detailed summary of a factual, structured passage (photosynthesis) at 500 tokens, across two repeated runs per gamma:
| Mode | Gamma | Tok/s | Acceptance (range, mean) | Speedup (range, mean) |
|---|---|---|---|---|
| Verifier-only | - | 13.9-14.1 | - | 1.00x |
| Speculative | 4 | 13.8-14.7 | 47.8-53.4% | 0.99x-1.16x |
| Speculative | 6 | 13.7-15.3 | 45.0-52.6% | 0.99x-1.08x |
| Speculative | 8 | 11.6-13.1 | 38.7-43.0% | 0.82x-0.94x |
Acceptance stayed in the 39 to 53 percent range across gamma values and repeated runs, holding up under a separate temperature and top_p sweep too (0.01, treated as greedy, 0.3, and 0.7 with top_p 0.9 all landed in a similar band). Gamma 4 and 6 both averaged just above 1.0x measured speedup, while gamma 8 consistently underperformed, continuing the pattern that lower gamma does better here. That's a bit counter-intuitive at first, since more proposed tokens per round should mean fewer verifier calls, but a single rejection discards every draft token after it in that round. Higher gamma means more wasted draft work per rejection, not just a bigger payoff when everything gets accepted. The accept/reject math behaved correctly throughout, with stable, repeatable acceptance rates round to round and no correctness bugs.
That acceptance rate is notably higher than an earlier gamma sweep on an open-ended, opinion-style prompt, which landed in the 20 to 40 percent range and stayed below the verifier-only baseline at every gamma tested. Temperature and top_p turned out not to be the variable that mattered, swapping in structured, factual content did. Lower local entropy per generation step seems to give the draft model a much easier job predicting what the verifier would say next, while open-ended, creative text doesn't offer that same easy target. That lines up with Part 1's entropy tracker, and it's a first real signal of that connection ahead of Phase 4's dedicated study.
The natural next question is why even the best runs here only reach roughly 1.0 to 1.16x rather than something closer to the 1.5 to 2x speedups reported in the literature. That's its own investigation, including a cross-check against llama.cpp's Metal backend to separate hardware and runtime effects from the algorithm itself, plus an honest accounting of what the published results actually assumed that this setup doesn't: Chasing Speedup: What I Got Wrong Along the Way
For Phase 4
Acceptance rates (40.4 percent overall) and fine-grained per-round patterns (stretches of 0 to 1 accepts, occasional 4 out of 4) are now clean data for correlation with Part 1's entropy measurements. The variation in acceptance across a generation may correlate with local entropy or text position, a signal Part 2's attention-sink work was designed to explain.
Connections
Part 1 (entropy): Both sample at the same temperature and top_p, so the verifier's distribution here equals Part 1's entropy target. The hypothesis is that high-entropy rounds should show lower acceptance, since the draft diverges more from an uncertain verifier.
Part 2 (attention sinks): Once the draft has a real cache, cache-eviction decisions aren't free. Part 2's per-prompt sink positions become a direct input into cache management.
Links
| Resource | Link |
|---|---|
| GitHub repo | https://github.com/cyprus09/llm-inference-lab |
| Speculative Decoding (Leviathan et al., 2023) | https://arxiv.org/abs/2211.17192 |
| Speculative Decoding with Large Language Models (Chen et al., 2023) | https://arxiv.org/abs/2302.01318 |
| StreamingLLM (Xiao et al., 2023) | https://arxiv.org/abs/2309.17453 |
Series:
- Part 1: Entropy Tracker
- Part 2: Attention Sink Detector
- Part 3: Speculative Decoding, Sampling-Mode Accept/Reject (you are here)
- Part 4: The Empirical Study (coming)
Top comments (0)