DEV Community: mayankpallai

Is Speculative Decoding's Speedup a Hardware Problem or a Model Problem?

mayankpallai — Sat, 25 Jul 2026 22:15:25 +0000

A follow-up/sub-part to Part 3 of the LLM inference internals series. Part 3 built sampling-mode speculative decoding with KV caching on both the draft and verifier sides, and it worked correctly, but the speedup it delivered didn't match expectations set by the paper. This post is the record of chasing that gap: every hypothesis tested, which ones were wrong, and what the real answer turned out to be.

Starting Point: It's Slower, Not Faster

The first full gamma sweep, on an open-ended, opinion-style prompt, came back like this:

Mode	Gamma	Tokens	Tok/s	Acceptance	Measured Speedup
Verifier-only	-	250	14.4	-	1.00x
Speculative	4	250	12.3	39.2%	0.86x
Speculative	7	250	9.5	33.6%	0.66x

Slower across the board. The accept/reject math had already been verified correct, with stable, repeatable acceptance rates and no correctness bugs, so this wasn't a bug. It was a real result that needed an explanation.

Hypothesis 1: It's MPS Dispatch Overhead

The first instinct was to blame the hardware. Apple Silicon's MPS backend pays a large, roughly fixed dispatch cost per model call that doesn't shrink proportionally with model size. That was confirmed via isolated microbenchmarks: the 6x smaller draft model was only about 2.8x faster per single-token call (21 to 28ms versus 60 to 90ms). Two sequential model calls per round, each paying that fixed tax, looked like the obvious explanation.

To test this without conflating hardware cost with algorithm cost, I computed a FLOPs-based theoretical speedup using Leviathan et al.'s formula, E[speedup] = (1 - α^(γ+1)) / ((1 - α)(γ·c + 1)), with a cost ratio c derived from each model's architecture (layers, hidden size, GQA heads, FLOPs per token via the standard 2 × params approximation) rather than measured milliseconds. That makes the number hardware-agnostic by construction.

Verifier (Qwen2.5-3B-Instruct):  6,171,394,048 FLOPs/token
Draft   (Qwen2.5-0.5B-Instruct):   987,922,432 FLOPs/token
c = 0.1601

Plugging in the real measured acceptance rates:

γ	Acceptance	Theoretical (FLOPs-only)	Measured (MPS)
4	39.2%	0.99x	0.86x
7	33.6%	0.71x	0.66x
8	19.7%	0.55x	0.55x

This was the first wrong turn. Even with zero hardware overhead, the theoretical ceiling never cleared 1.0x. If MPS overhead were the whole story, the theoretical column should have shown a clear, comfortable win being eaten by dispatch cost. Instead it showed that the algorithm itself, on paper, wasn't going to win at this acceptance rate. Hypothesis 1 was only a partial explanation at best.

Hypothesis 2: A Bigger Draft Model Would Fix It

If acceptance rate was the real ceiling, the obvious next move was a closer-sized draft: Qwen2.5-1.5B instead of 0.5B, on the theory that a closer size means a closer distribution match to the verifier.

Checking this on paper first, via AutoConfig with no weights downloaded, gave a worse cost ratio: c = 0.5002, about 2x cheaper than the verifier instead of 6x. Solving the formula for breakeven acceptance rate at each gamma showed the 1.5B candidate would need 69 to 84 percent acceptance just to reach 1.0x, far above anything measured so far. Second wrong turn, ruled out before spending compute on it. A bigger draft looked like trading one problem, low acceptance, for a worse version of the same problem, an unfavorable cost ratio.

Then I actually ran it. Real numbers, not just theory:

Run	γ	Acceptance	Theoretical	Measured
1	6	86.4%	1.18x	1.37x
2	6	19.8%	0.31x	0.48x
3	6	39.5%	0.41x	0.72x
4	6	30.1%	0.36x	0.59x

One run cleared 1.0x by a wide margin, the first of the whole investigation. But three others landed back in the 20 to 40 percent range the FLOPs math had already predicted was a losing zone. This wasn't the bigger-draft theory being confirmed. It was acceptance rate itself being far noisier and more content-dependent than a single-number cost-ratio model could capture. The 1.5B draft wasn't reliably better, it was occasionally much better, which pointed at something other than model size driving the variance.

There was also a second anomaly in this data. In several rows, measured speedup exceeded the theoretical ceiling, which is supposedly impossible since theoretical was meant to be a zero-overhead upper bound. That turned out to be a real finding, not a bug. FLOPs-based c assumes cost scales with parameter count, but on MPS, dispatch overhead compresses the real cost gap between a 1.5B and 3B model far more than FLOPs alone would predict. The FLOPs-based ceiling is a valid algorithmic upper bound, but it is not a valid hardware upper bound on MPS, a subtlety worth naming rather than glossing over.

Hypothesis 3: It's the Prompt

The high-variance 1.5B result raised the real question directly: what made that one run hit 86 percent acceptance? The recurring guess was entropy. A prompt with structured, predictable, lower-branching continuations should let a small draft model track a large verifier much more closely than an open-ended, creative one.

To test this cleanly, I swept temperature (0.01, 0.3, 0.7) and top_p on the original opinion-style prompt, and acceptance barely moved, staying in the 20 to 45 percent band regardless of sampling settings. Then I swapped the prompt itself: same models, same code, same hyperparameter ranges, but asked for a detailed summary of a structured, factual passage (photosynthesis) instead of an open-ended question, at 500 tokens:

Mode	Gamma	Tok/s	Acceptance	Measured Speedup
Verifier-only	-	13.1-14.3	-	1.00x
Speculative	4	12.8-15.5	43.6-54.0%	0.95x-1.16x
Speculative	6	13.7-15.3	45.0-52.6%	0.99x-1.08x

That's the first time changing a single variable moved acceptance meaningfully, and it wasn't temperature or top_p, it was the content itself. Structured, factual text is where this draft and verifier pair actually performs. Open-ended, creative text is where it doesn't. This became the real answer to what causes the acceptance ceiling here: not the models being poorly matched in general, but poorly matched specifically on high-branching-entropy content, a hypothesis Phase 4's entropy correlation study can now test directly instead of guessing at.

Hypothesis 4: It's Still Just MPS, Even at High Acceptance

With acceptance now reliably in the 45 to 54 percent range and several runs already crossing 1.0x on PyTorch and MPS, the last open question was whether the shortfall from the 1.5 to 2x range reported in the literature was still a hardware story. To check, I ran the same models (GGUF, Q4_K_M quantization) through llama.cpp on the same photosynthesis-style prompt:

n_draft   = 4
n_predict = 504
n_accept  = 391 / 448 drafted  ->  87.277% acceptance
decoded 504 tokens in 10.581s -> 47.6 tok/s (speculative)
verifier-only baseline (same prompt, same runtime): 50.0 tok/s
Speedup: 47.6 / 50.0 ≈ 0.95x

llama.cpp's Metal backend is genuinely faster in absolute terms, roughly 3x the raw tok/s of the PyTorch and MPS setup on both the baseline and the speculative run, confirming the earlier microbenchmark finding that MPS pays real dispatch overhead PyTorch doesn't fully hide. But at 87.3 percent acceptance, well above anything achieved on PyTorch, llama.cpp's own speedup was still around 0.95x, not the 1.5 to 2x hoped for. Fourth wrong turn. If MPS and PyTorch overhead were the primary blocker, a near-zero-overhead runtime at excellent acceptance should have shown a clear win. It didn't.

llama.cpp printed something that pointed at the real remaining factor: 35.7 percent of total wall-clock was unaccounted time, not verifier compute, not draft compute, not sampling, but round-trip and orchestration cost between the draft and verify phases. That's consistent with both implementations, this repo's and llama.cpp's CLI example, running draft and verify strictly sequentially, round after round, rather than pipelining verification of round N with drafting of round N+1.

What Actually Explains the Gap

Putting all four hypotheses together, in order of how much each one turned out to matter:

Per-round orchestration overhead, not per-call dispatch cost, is the main remaining bottleneck at good acceptance rates, confirmed by llama.cpp hitting the same 0.95x ceiling at 87 percent acceptance despite near-zero dispatch overhead.
Content and entropy drive acceptance rate far more than temperature, top_p, or draft model size, the single biggest, most reproducible lever found in this whole investigation.
MPS dispatch overhead is real but secondary. It explains part of the PyTorch versus llama.cpp absolute speed gap, but not why even the fast runtime falls short of the literature's numbers.
A bigger draft model is not a reliable fix. It's occasionally much better, but with high enough variance that it isn't a dependable lever on its own.

Why the Paper's Numbers Are Higher

Leviathan et al.'s headline results didn't come from one trick. They stack several advantages this setup doesn't have:

A much larger size gap (T5-XXL to T5-small, roughly 180x parameters) versus this pairing's roughly 6x. A bigger absolute verifier cost makes batching gamma plus 1 tokens into one pass save proportionally more wall-clock time.
Drafts trained to align with the verifier, not independently instruction-tuned same-family siblings. This is very likely the single largest acceptance-rate lever available, and this repo's Qwen2.5-0.5B and 3B pair has none of it.
Pipelined, overlapped execution between draft and verify stages, rather than strict sequential round-trips, directly addressing the orchestration overhead this investigation found via llama.cpp's unaccounted time.
TPU and datacenter-GPU dispatch overhead, lower even than llama.cpp's Metal backend, shrinking the fixed per-round tax further.

None of these is a bug in this implementation. They're conditions the paper's results depended on that this project, by design (local, from scratch, consumer hardware, off-the-shelf same-family models), doesn't have.

Limitations

Draft and verifier size ratio and training. A 6x parameter gap versus the paper's roughly 180x, and no distillation. The draft was never trained to mimic the verifier, which is very likely the dominant reason acceptance tops out where it does even on favorable content.
Sequential execution only. Both this implementation and the llama.cpp comparison run draft and verify strictly in sequence. No round-pipelining was attempted, so the roughly 35.7 percent orchestration overhead measured here isn't necessarily a hard floor, just what this simple execution model costs.
Hardware and runtime. Results span PyTorch and MPS, and llama.cpp and Metal, on Apple Silicon only. No datacenter GPU or TPU data point exists to confirm how much of the remaining gap is Apple Silicon specific versus fundamental to unpipelined two-model decoding.
Sample size. Acceptance-rate variance was large enough (19.8 to 86.4 percent on the same 1.5B and 3B pair, same gamma) that single-prompt, few-run comparisons should be read as suggestive, not conclusive. The entropy and content hypothesis needs Phase 4's larger prompt set to confirm.
FLOPs approximation. The 2 × params estimate doesn't model attention's quadratic term in sequence length, and the theoretical formula assumes no fixed per-round overhead. Both simplifications that this investigation's own data (measured beating theoretical, llama.cpp's unaccounted time) shows don't fully hold in practice.

Series: Part 3, Speculative Decoding: Sampling-Mode Accept/Reject, leads into this post, which leads into Part 4, The Empirical Study (coming)

Building a Terminal Based LLM Inference Internals Explorer - Part 3

mayankpallai — Fri, 24 Jul 2026 16:02:54 +0000

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)

Building a Terminal Based LLM Inference Internals Explorer - Part 2

mayankpallai — Wed, 15 Jul 2026 19:55:50 +0000

Part 2: The Attention Sink Detector

Part 2 of a 4-part series on system-level LLM inference internals. Part 1 built the entropy tracker; this one looks one step earlier in the pipeline — at prefill, before a single token is generated.

Where This Fits

Part 1 tracked entropy during decode: how confident the model is, one generated token at a time. This post goes further upstream. Before decode even starts, the model runs a single forward pass over the entire prompt, and that pass is where attention gets distributed across every token in the context.

The unifying thesis from Part 1 said context quality shapes attention distribution during prefill, and attention distribution shapes generation confidence during decode. This post is where that first link gets measured directly.

Part	What We Build
1	Per-token entropy tracker, visualized in real time
2 — this post	Attention sink detector, context health scoring
3	Entropy-guided adaptive speculative decoding
4	Empirical study: correlation plots proving the causal chain

The Phenomenon

Across a wide range of transformer LLMs, a small number of tokens, very often just the first token in the sequence — receive a hugely disproportionate share of attention from almost every later query, in almost every head, almost regardless of semantic relevance. StreamingLLM (Xiao et al., 2023) is the paper that named this: attention sinks, and the tokens absorbing the attention are sink tokens.

The mechanism comes straight out of the softmax constraint. Attention weights over keys are:

a_i = exp(q · k_i) / Σ_j exp(q · k_j)

This always sums to 1. A query can't express "nothing here is relevant" by producing low weight everywhere — the weights must still add up to 1 across all keys. So the model needs a release valve: somewhere to dump residual probability mass when nothing in the context is a strong match. Position 0 is the natural candidate, since it's visible to every query in a causal model and it's the one position guaranteed to exist as a valid attention target in every training sequence, at every length, regardless of anything else.

Why this matters practically, per the StreamingLLM paper: if you're managing a KV cache for long-context or streaming generation and evict the sink token to save memory, perplexity spikes — even though position 0 seems completely irrelevant by the time you're thousands of tokens deep. Their fix is to pin the first few tokens' KV pairs permanently, alongside a sliding window of recent tokens.

Phase 0: Getting the Numbers to Exist at All

The first obstacle isn't conceptual, it's mechanical. SDPA and FlashAttention kernels — the fast, default attention implementations — never materialize the full seq_len × seq_len attention matrix. That's the entire point of them: they compute attention output through tiled, online-softmax kernels, discarding each block's partial scores as soon as it's consumed. Great for speed and memory at long context. Useless if you want to look at the numbers.

model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen2.5-3B-Instruct",
    attn_implementation="eager",   # non-negotiable for this phase
    torch_dtype=torch.float16,
).to("mps")

attn_implementation="eager" is the only path where output_attentions=True returns real tensors instead of None. It's slower and considerably more memory-hungry — for a seq_len × seq_len matrix per head per layer, that adds up fast — but for a single prefill pass on a 3B model on Apple Silicon, it's a completely reasonable trade.

Phase 1: The Prefill Analyzer

Aggregating Attention

A single forward pass with output_attentions=True returns one tensor per layer, each shaped (batch, num_heads, seq_len, seq_len). Averaging across layers and heads gives one matrix: attn[i, j] = how much query token i attends to key token j.

def aggregate_attention(attentions):
    """Mean attention matrix across layers and heads.
    Returns (seq_len, seq_len); attn[i, j] = query i's attention to key j."""
    stacked = torch.stack([layer_attn[0].float() for layer_attn in attentions], dim=0)
    return stacked.mean(dim=(0, 1))

Causal-Normalized Attention Received

A raw column sum over this matrix favors early tokens for a boring reason: token j is only visible to queries i >= j, so early positions have simply had more chances to be attended to. Dividing by the number of queries that could actually see each position gives a fair, comparable-across-positions score:

def mean_attention_received(attn_matrix):
    seq_len = attn_matrix.shape[0]
    col_sums = attn_matrix.sum(dim=0)
    valid_queries = torch.arange(seq_len, 0, -1, dtype=torch.float32)
    return col_sums / valid_queries

Structural Tokens Aren't Content

Every ChatML-templated prompt carries fixed scaffolding — <|im_start|>, <|im_end|>, role-name tokens, and the newlines the template inserts around them. These tokens are content-independent: their identity is fixed by the template, not by what was actually said. They need to be excluded from the baseline used to judge what "normal" attention looks like, or they distort the comparison for every real content token in the prompt.

def find_structural_positions(input_ids, tokenizer):
    """Positions fixed by the ChatML template rather than by content."""
    ids = input_ids[0].tolist()
    special_ids = set(tokenizer.all_special_ids)
    im_start_id = tokenizer.convert_tokens_to_ids("<|im_start|>")
    im_end_id = tokenizer.convert_tokens_to_ids("<|im_end|>")

    structural = {0}
    for i, tid in enumerate(ids):
        if tid in special_ids:
            structural.add(i)
        if i > 0 and ids[i - 1] in (im_start_id, im_end_id):
            structural.add(i)  # role name (after im_start) or newline (after im_end)
            if ids[i - 1] == im_start_id and i + 1 < len(ids):
                structural.add(i + 1)  # newline after the role name specifically
    return structural

Phase 2: Getting the Statistics Right (Two Times)

The first version of the sink-scoring logic looked reasonable and produced numbers that were, on inspection, badly wrong. Getting this right took two separate corrections.

Correction 1 — z-score, not skew-blind mean/std

Attention weight is strictly non-negative and heavily right-skewed: most tokens sit near zero, a handful sit higher. A plain mean/std z-score assumes something closer to a normal distribution, and against a skewed population that assumption quietly inflates what counts as "anomalous." A first pass at a threshold of 2.0 reliably flagged over a dozen tokens per prompt as sinks — and on inspection they were just ordinary salient content words ("regarding," "potential," "biases"), not sinks in any mechanistic sense.

The fix: log-transform attention values before computing spread, and use median and median absolute deviation (MAD) instead of mean/std — MAD doesn't get dragged upward by the same skew it's measuring.

def detect_sinks(mean_attn, threshold=6.0, exclude_positions=(0,)):
    mask = torch.ones_like(mean_attn, dtype=torch.bool)
    for p in exclude_positions:
        if 0 <= p < len(mask):
            mask[p] = False

    population = mean_attn[mask]
    log_population = torch.log(population.clamp_min(1e-8))
    log_all = torch.log(mean_attn.clamp_min(1e-8))

    median = log_population.median()
    mad = (log_population - median).abs().median()
    scaled_mad = mad * 1.4826  # comparable scale to std under normality

    z = torch.zeros_like(mean_attn) if scaled_mad.item() < 1e-8 else (log_all - median) / scaled_mad
    return z > threshold, z

Correction 2 — structural doesn't mean sink

An earlier version force-set every structural position to "is a sink," on the reasoning that BOS-as-sink is well-established enough to just assume. That reasoning breaks the moment it's generalized to every scaffolding token, most of the non-BOS structural positions turned out to have z-scores near zero or negative once actually measured. The fix was to stop overwriting measured evidence: a position is only a sink if its z-score says so. "Structural" and "sink" are now two independent flags, not one collapsing into the other.

What It Looks Like End to End

Here's an actual run, on a prompt about training-data provenance and self-supervised bias:

╭─ Context health ─╮
│ 6.7/10           │
╰──────────────────╯

    Structural positions (template-fixed, excluded from baseline)
┌──────────┬────────────────┬────────────┬─────────┬───────┐
│ Position │ Token          │ Mean attn. │ Z-score │ Sink? │
├──────────┼────────────────┼────────────┼─────────┼───────┤
│ 0        │ '<|im_start|>' │ 0.3716     │ 7.23    │ yes   │
│ 1        │ 'system'       │ 0.0023     │ -1.06   │ no    │
│ 2        │ '\n'           │ 0.0359     │ 3.42    │ yes   │
│ 19       │ '<|im_end|>'   │ 0.0027     │ -0.82   │ no    │
│ 20       │ '\n'           │ 0.0038     │ -0.26   │ no    │
│ 21       │ '<|im_start|>' │ 0.0018     │ -1.44   │ no    │
│ 22       │ 'user'         │ 0.0018     │ -1.50   │ no    │
│ 23       │ '\n'           │ 0.0062     │ 0.55    │ no    │
│ 274      │ '<|im_end|>'   │ 0.0648     │ 4.38    │ yes   │
│ 275      │ '\n'           │ 0.0684     │ 4.47    │ yes   │
│ 276      │ '<|im_start|>' │ 0.0726     │ 4.57    │ yes   │
│ 277      │ 'assistant'    │ 0.0734     │ 4.59    │ yes   │
│ 278      │ '\n'           │ 0.1230     │ 5.43    │ yes   │
└──────────┴────────────────┴────────────┴─────────┴───────┘

              Anomalous sink tokens
┌──────────┬──────────────┬────────────┬─────────┐
│ Position │ Token        │ Mean attn. │ Z-score │
├──────────┼──────────────┼────────────┼─────────┤
│ 73       │ ' variety'   │ 0.0208     │ 2.53    │
│ 246      │ ' ###'       │ 0.0187     │ 2.36    │
│ 247      │ ' Self'      │ 0.0173     │ 2.23    │
│ 253      │ ' Bias'      │ 0.0217     │ 2.59    │
│ 254      │ '.'          │ 0.0182     │ 2.31    │
│ 255      │ ' When'      │ 0.0213     │ 2.57    │
│ 256      │ ' an'        │ 0.0152     │ 2.02    │
│ 259      │ ' trains'    │ 0.0217     │ 2.60    │
│ 260      │ ' on'        │ 0.0186     │ 2.35    │
│ 264      │ ','          │ 0.0304     │ 3.15    │
│ 265      │ ' several'   │ 0.0249     │ 2.83    │
│ 266      │ ' factors'   │ 0.0288     │ 3.06    │
│ 267      │ ' come'      │ 0.0165     │ 2.15    │
│ 269      │ ' play'      │ 0.0209     │ 2.53    │
│ 270      │ ' regarding' │ 0.0247     │ 2.81    │
│ 271      │ ' potential' │ 0.0202     │ 2.48    │
│ 272      │ ' biases'    │ 0.0215     │ 2.58    │
│ 273      │ ':'          │ 0.0319     │ 3.23    │
└──────────┴──────────────┴────────────┴─────────┘

Reading this correctly, in three parts:

Position 0 is a real, strong sink. z = 7.23 on a prompt where the log-MAD baseline is built entirely from the remaining ~270 positions. This is the textbook case — BOS absorbing residual softmax mass regardless of content, exactly as StreamingLLM describes.

The five-token run right at the generation boundary (<|im_end|> \n <|im_start|> assistant \n, positions 274–278) is a second, milder sink cluster — z climbing from 4.38 to 5.43 as you approach the point where generation actually starts. This wasn't something I expected going in; the working theory is that these tokens function as a second "nothing decided yet" boundary, structurally similar to BOS in that every later query can see them and their identity is fixed by the template rather than by the conversation. Whether this is a general pattern or specific to this template/model is exactly the kind of question the 50-prompt run should answer.

The "anomalous" table is still showing false positives, and that's on me, not the model. Every row here — "regarding," "biases," "trains," even bare punctuation like , and : — is an ordinary content word from a real sentence about training data and bias, not a token with no semantic relationship to its context. The log-MAD correction from Phase 2 fixed how the spread is measured, but this run still uses threshold=2.0, which is a leftover from when the metric was mean/std-based and needs recalibrating against the new scale. With BOS at z=7.23 and the boundary run at z=4.4–5.4, a threshold somewhere in the 5–6 range would separate real structural/positional sink behavior from what's currently just the natural right tail of content-word attention. I'm leaving this table in rather than cropping it out, because the gap between "the statistic is now sound" and "the threshold is correctly tuned" is itself worth showing — they're different bugs, and fixing the first doesn't automatically fix the second.

How This Connects to Parts 1 and 3

Back to entropy (Part 1): The thesis was that degraded attention during prefill should show up as elevated generation entropy during decode. Now that sink detection is trustworthy rather than noisy, this phase produces a single per-prompt "context health" number that Part 4 can actually correlate against Part 1's mean-entropy metric, instead of correlating against a score that was partly measuring content-word variance by accident.

Forward to speculative decoding (Part 3): KV cache eviction under speculative decoding runs into exactly the sink-eviction problem StreamingLLM identified — evict the wrong position and acceptance rates degrade. Knowing precisely which positions in a given prompt are genuine sinks, rather than assuming "it's probably just BOS," is a direct input into cache retention policy once Part 3 starts trading off draft length against verification cost.

Links

Resource	Link
GitHub repo	https://github.com/cyprus09/llm-inference-lab
StreamingLLM	https://arxiv.org/abs/2309.17453
Speculative Decoding	https://arxiv.org/abs/2211.17192
Lost in the Middle	https://arxiv.org/abs/2307.03172
The Illustrated Transformer	https://jalammar.github.io/illustrated-transformer/

Series:

Part 1 — Entropy Tracker
Part 2 — Attention Sink Detector (you are here)
Part 3 — Speculative Decoding with Entropy-Guided Draft Length (coming)
Part 4 — The Empirical Study (coming)

Building a Terminal Based LLM Inference Internals Explorer - Part 1

mayankpallai — Wed, 15 Jul 2026 08:21:39 +0000

Part 1: The Entropy Tracker

Part 1 of a 4-part series on system-level LLM inference internals.

What This Series Builds

Most LLM tooling treats inference as a black box. Hosted APIs make this worse; they strip away logits, attention weights, and intermediate activations entirely. What's left is just surface behavior.

This project goes the other direction. Running a 3B model locally on Apple Silicon means getting everything: raw logit distributions at every decode step, full attention weight tensors during prefill, and direct control over the generation loop.

The unifying thesis:

Context quality shapes attention distribution during prefill. Attention distribution shapes generation confidence during decode. Generation confidence determines how efficiently speculative decoding can run. These three things are causally linked — and this series builds the tools to try and prove it.

📊 Diagram:

Part	What We Build
1 — this post	Per-token entropy tracker, visualized in real time
2	Attention sink detector, context health scoring
3	Entropy-guided adaptive speculative decoding
4	Empirical study: correlation plots proving the causal chain

Hardware & Stack

MacBook Pro M1 Pro, 16GB unified memory
Model: Qwen2.5-3B-Instruct, fp16, MPS backend (~17 tok/s warm)
Python · PyTorch · HuggingFace Transformers · Rich

This project requires local inference. Hosted APIs (Anthropic, OpenAI) don't expose raw logits or attention weights. To see inside the model, you have to run it yourself.

Phase 0: Environment Setup

pyproject.toml

[project]
name = "llm-inference-lab"
version = "0.1.0"
requires-python = ">=3.10,<3.13"
dependencies = [
    "torch>=2.3.0",
    "transformers>=4.42.0",
    "accelerate>=0.31.0",
    "sentencepiece>=0.2.0",
    "protobuf>=4.25.0",
    "rich>=13.7.0",
    "numpy>=1.26.0",
    "matplotlib>=3.8.0",
    "pandas>=2.2.0",
]

MPS sanity check

import torch
print(torch.backends.mps.is_available())   # True on M1/M2/M3
x = torch.rand(1000, 1000, device="mps")
print((x @ x).shape)                       # confirms GPU matmul works

Benchmarks on M1 Pro, 16GB

Run	Tokens	Time	Tok/s
Cold (MPS kernel compilation)	50	5.59s	8.9
Warm	50	2.89s	17.3
Warm, longer	100	5.95s	16.8

The cold-start penalty is a one-time cost per process. All subsequent calls run at ~17 tok/s. All params confirmed on mps:0 in fp16 — no silent CPU fallback.

Phase 1: The Entropy Tracker

The Math

At each decode step, the model produces a logit vector over ~32,000 vocabulary tokens. After softmax this becomes a probability distribution. Shannon entropy measures how uncertain the model is:

H(t) = -Σ p(x) · log2(p(x))    over all vocab tokens x

Low H  → peaked distribution → confident
High H → flat distribution   → uncertain

The Hook: `LogitsProcessor`

from transformers import LogitsProcessor
import torch.nn.functional as F

class EntropyCapture(LogitsProcessor):
    def __init__(self):
        self.entropies = []
        self.top_tokens = []

    def __call__(self, input_ids, scores):
        probs = F.softmax(scores, dim=-1)
        log_probs = torch.clamp(probs.log2(), min=-1e9)
        H = -(probs * log_probs).sum(dim=-1)
        self.entropies.append(H.item())

        top = torch.topk(probs, k=5, dim=-1)
        self.top_tokens.append({
            "values": top.values[0].tolist(),
            "indices": top.indices[0].tolist(),
        })

        return scores  # unchanged -- observing, not modifying

Wiring it in:

processor = EntropyCapture()

outputs = model.generate(
    **inputs,
    max_new_tokens=200,
    do_sample=True,
    temperature=0.7,
    logits_processor=[processor],
)
# processor.entropies now has one float per generated token

Rich Terminal Renderer

from rich.text import Text
from rich.console import Console

def entropy_color(H: float) -> str:
    if H < 1.0: return "green"
    if H < 3.0: return "yellow"
    return "red"

console = Console()
text = Text()
for token_str, H in zip(generated_token_strings, entropies):
    text.append(token_str, style=entropy_color(H))
console.print(text)

Key Finding: Subword Commitment Points

Qwen2.5 uses BPE tokenization — words split into subword units. "civilization" might tokenize as civil + ization. What happens at the entropy level?

civil → high entropy (the model commits to this word here)
ization → low entropy (the continuation is already determined)

📊 Diagram:

Entropy spikes mark commitment points — where the model decides among multiple valid continuations. Once the first subword of a new word is chosen, the rest is nearly deterministic. The real decision happens at the leading edge.

Entropy Profiles by Prompt Type

Prompt Type	Mean Entropy	Notes
Factual ("capital of France")	~0.8	Mostly confident, few spikes
Creative writing("name 10 new colors")	~2.6	Frequent uncertainty
Code generation	~1.1	Surprisingly confident — syntax constrains
Math reasoning	~1.4	Spikes at numeric choices
Ambiguous questions("what makes most sense?")	~3.1	Sustained high entropy

Code being "greener" than prose is counterintuitive but makes sense: the model has strong priors about what syntactically valid Python looks like. In prose, almost any word could plausibly follow.

How This Connects to Parts 2 and 3

To attention sinks (Part 2): If attention during prefill pools into irrelevant sink tokens, the model enters decode with a degraded state — observable as higher mean generation entropy on poisoned contexts. Part 2 measures both sides and plots the correlation.

To speculative decoding (Part 3): The draft model acceptance criterion is min(1, p_verifier / p_draft). The draft gets accepted most when it's confident — low entropy, peaked distribution. High draft entropy signals a likely upcoming rejection. So instead of always drafting a fixed k tokens, we stop early when entropy spikes.

📊 Diagram:

— decision flow: low draft entropy → keep drafting (likely accepted); high draft entropy → call verifier (rejection incoming).

This is entropy-guided adaptive speculative decoding — the thread connecting all three phases.

Links

Resource	Link
GitHub repo	https://github.com/cyprus09/llm-inference-lab
StreamingLLM	https://arxiv.org/abs/2309.17453
Speculative Decoding	https://arxiv.org/abs/2211.17192
Lost in the Middle	https://arxiv.org/abs/2307.03172
The Illustrated Transformer	https://jalammar.github.io/illustrated-transformer/
Making Deep Learning Go Brrrr	https://horace.io/brrr_intro.html
nanoGPT	https://github.com/karpathy/nanoGPT

Series:

Part 1 — Entropy Tracker (you are here)
Part 2 — Attention Sink Detector (coming)
Part 3 — Speculative Decoding with Entropy-Guided Draft Length (coming)
Part 4 — The Empirical Study (coming)

DEV Community: mayankpallai

Is Speculative Decoding's Speedup a Hardware Problem or a Model Problem?

Starting Point: It's Slower, Not Faster

Hypothesis 1: It's MPS Dispatch Overhead

Hypothesis 2: A Bigger Draft Model Would Fix It

Hypothesis 3: It's the Prompt

Hypothesis 4: It's Still Just MPS, Even at High Acceptance

What Actually Explains the Gap

Why the Paper's Numbers Are Higher

Limitations

Building a Terminal Based LLM Inference Internals Explorer - Part 3

Part 3: Speculative Decoding via Sampling-Mode Accept/Reject

The Core Idea

The Accept/Reject Mechanism

A Worked Example

Sampling vs Greedy

The Flow

KV Caching on Both Sides

Entropy-Guided Stopping

Results

Links

Building a Terminal Based LLM Inference Internals Explorer - Part 2

Part 2: The Attention Sink Detector

Where This Fits

The Phenomenon

Phase 0: Getting the Numbers to Exist at All

Phase 1: The Prefill Analyzer

Aggregating Attention

Causal-Normalized Attention Received

Structural Tokens Aren't Content

Phase 2: Getting the Statistics Right (Two Times)

Correction 1 — z-score, not skew-blind mean/std

Correction 2 — structural doesn't mean sink

What It Looks Like End to End

How This Connects to Parts 1 and 3

Links

Building a Terminal Based LLM Inference Internals Explorer - Part 1

Part 1: The Entropy Tracker

What This Series Builds

Hardware & Stack

Phase 0: Environment Setup

Phase 1: The Entropy Tracker

The Math

The Hook: LogitsProcessor

Rich Terminal Renderer

Key Finding: Subword Commitment Points

Entropy Profiles by Prompt Type

How This Connects to Parts 2 and 3

Links

The Hook: `LogitsProcessor`