mayankpallai

Posted on Jul 24 • Edited on Jul 25

Building a Terminal Based LLM Inference Internals Explorer - Part 3

#llm #ai #pytorch #tutorial

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 sequential: each token's distribution depends on all previous tokens, so a large model generates one token per full forward pass. Speculative decoding breaks that: a small, fast draft model proposes multiple tokens at once, and a large verifier model checks all of them in a single batched forward pass. The trick 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

The hard part isn't drafting tokens — it's 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

If p_verify > p_draft: the verifier is more confident than the draft in this token, so accept_prob clamps to 1 — always accept.
If p_verify < p_draft: the verifier is less confident, so accept only with probability p_verify / p_draft. This is the rate that exactly cancels out the "extra" mass the draft added.

On rejection, you don't just take the verifier's argmax — you resample from what's left over:

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 verifier's probability mass that isn't already accounted for by the draft's guess. Sampling from it, rather than the verifier's raw distribution, is what makes the whole thing add up: accept + resample composes back to exactly p_verify, so the output is provably indistinguishable from sampling the verifier alone, even though the draft moved first.

draft_dist is the draft's full filtered distribution over the whole vocabulary, not just the probability of the token it sampled — the residual subtraction needs to know what the draft thought about every token, not only the one that got drawn.

A Worked Example

Say the draft samples "cat" at some position, with p_draft(cat) = 0.6. The verifier doesn't sample its own token here — it just reads what probability it would have assigned to "cat" off its own distribution. Say 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 a coin flip at 50%. Say it fails, so "cat" is rejected.

Now build the residual. Suppose the (simplified) vocabulary is just {cat, dog, fox}, the draft's own distribution at this position was cat=0.6, dog=0.3, fox=0.1, and the verifier's distribution is cat=0.3, dog=0.5, fox=0.2:

draft_dist  = { cat: 0.6, dog: 0.3, fox: 0.1 }
residual    = { cat: 0,   dog: 0.2, fox: 0.1 }   # max(0, verify - draft)
normalized  = { cat: 0,   dog: 0.67, fox: 0.33 }

"cat" gets zeroed out of the residual entirely — it already had its shot during the accept/reject coin flip, so it can't be picked again as its own replacement. That would double-count its mass. Note the residual uses the draft's full distribution, not just p_draft(cat) — "dog" and "fox" both had real draft mass too, and that mass has to be subtracted out just like "cat"'s did, or the residual overstates how much the draft actually left on the table for them. The correction token is sampled from {dog: 0.67, fox: 0.33} — say "dog" comes out.

One more detail worth being explicit about: 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 doesn't causally follow the corrected sequence anymore. 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 exist in the literature: greedy (draft argmax, accept iff verifier argmax matches — deterministic, higher acceptance rates) and sampling (stochastic, uses the probability-ratio test above).

Parts 1–2 use stochastic sampling at temperature=0.7, top_p=0.9, so this implementation uses sampling mode to keep acceptance rates comparable across the series. The 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. On the very first round there's no cache yet, so the verifier prefills on the full prompt plus all gamma draft tokens; every round after that, it reuses the cache and only forwards the newest accepted token plus the new draft tokens:

new_ids = input_ids if cache is None else input_ids[:, -1:]
new_ids = torch.cat([new_ids, draft_ids], dim=1)
output = verifier_model(new_ids, past_key_values=cache, use_cache=True)
logits = output.logits[0, -(gamma + 1):]  # last gamma+1 positions: draft scores + bonus

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 (shown above), and the main loop just appends whatever survives and starts the next round.

One easy-to-miss requirement: both models must 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, sample_token and the verifier's inline softmax + top_p_filter call share the same TEMPERATURE/TOP_P constants for exactly this reason.

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/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(s) 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 newest accepted token plus the newly proposed draft tokens on top of an already-cached prefix. The one difference from the draft side is the very first round, where there's no cache yet — that call has to prefill on the full prompt, not just its last token, otherwise the verifier would be scoring draft tokens with no context on what came before them.

Cost per round now 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.

Entropy-Guided Stopping

The draft phase above always proposes a fixed gamma tokens per round. But the draft model knows, at each step, how confident it is in its own guess, that's exactly what Part 1's entropy tracker measures. If the draft's normalized entropy at a position crosses a threshold, it's a signal the draft itself is unsure, and a token it's unsure about is a token likely to get rejected anyway. So instead of always proposing the full gamma, the draft phase checks its own entropy before sampling each token and stops proposing early once it crosses that threshold:

_, normalized_entropy = compute_entropy(logits)
if i > 0 and normalized_entropy >= ENTROPY_STOP_THRESHOLD:
    break

token_id, p, dist = sample_token(logits, temperature, top_p)

The round always proposes at least one token, so a round can never come back empty. Everything downstream — accept/reject, acceptance-rate bookkeeping, the bonus-token logic — is unchanged: a round that stopped early just proposed fewer tokens, exactly as if gamma had been smaller for that one round.

This isn't primarily a speedup feature — cutting a proposal short saves a small draft forward pass, but the verifier's per-round cost is roughly the same regardless. The real value is data: entropy_trace records the draft's own uncertainty at every proposed position, giving Phase 4 a direct per-token signal to line up against acceptance, rather than inferring the entropy/acceptance link from acceptance rate alone.

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 at 500 tokens, across 5 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–50.6%, mean 49.2%	0.99x–1.04x, mean 1.02x
Speculative	6	13.7–15.3	45.0–52.6%, mean 48.8%	0.99x–1.08x, mean 1.04x

Acceptance held in the low-to-mid 40s for both gamma values, and measured speedup tracked close from the acceptance rate — gamma=6 almost exactly, gamma=4 losing a bit more to overhead. That's consistent with a single rejection discarding every draft token after it in that round: higher gamma means more wasted draft work per rejection, not just a bigger payoff when everything's accepted. The accept/reject math held up throughout — stable, repeatable acceptance rates round to round.

That acceptance rate is notably higher than an early sweep on an open-ended, opinion-style prompt, which landed in the 20–40% range regardless of temperature/top_p. Structured, factual content gives the draft model an easier job predicting what the verifier would say next than open-ended, creative text does — a first, real signal ahead of Phase 4's dedicated study.

The natural next question is why even the best runs here land near 1.0x rather than the 1.5–2x speedups reported in the literature. That's its own investigation, including a cross-check against llama.cpp's Metal backend to separate hardware/runtime effects from the algorithm itself: Is Speculative Decoding's Speedup a Hardware Problem or a Model Problem?

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)

DEV Community