Decoding And KV Cache

def top_k_filter(logits, k):
    values, _ = torch.topk(logits, k, dim=-1)
    cutoff = values[..., -1, None]
    return logits.masked_fill(logits < cutoff, float("-inf"))

def temperature_sample(logits, temperature=1.0):
    probs = torch.softmax(logits / temperature, dim=-1)
    return torch.multinomial(probs, num_samples=1).squeeze(-1)

Implementation Drill: Top-p Without Empty Rows

Nucleus sampling keeps the smallest high-probability prefix whose cumulative mass reaches p, then samples from that truncated distribution. The practical bug to avoid is masking every token in a row, because torch.multinomial needs a finite, non-negative distribution with non-zero mass after softmax.

def top_p_filter(logits, p=0.9):
    if not 0.0 < p <= 1.0:
        raise ValueError("p must be in (0, 1]")

    sorted_logits, sorted_idx = torch.sort(logits, descending=True, dim=-1)
    sorted_probs = torch.softmax(sorted_logits, dim=-1)
    cumulative = sorted_probs.cumsum(dim=-1)

    remove = cumulative >= p
    remove[..., 1:] = remove[..., :-1].clone()
    remove[..., 0] = False  # Always leave at least one candidate.

    filtered_sorted = sorted_logits.masked_fill(remove, float("-inf"))
    return torch.empty_like(logits).scatter(-1, sorted_idx, filtered_sorted)

1. Check mass: after filtering and softmax, each row should sum to one and contain at least one finite token.
2. Check boundaries: p=1.0 should leave logits unchanged, while tiny p should reduce each row to greedy sampling.
3. Check ordering: sort for cumulative mass, but scatter back so token ids remain aligned with the original vocabulary.

Implementation Drill: Keep The Sampler Pipeline Sampleable

Generation libraries treat temperature, top-k, top-p, penalties, and constraints as logit processors or warpers that run before the final token draw. The implementation detail that matters in interviews is the invariant at the end of that chain: PyTorch's torch.multinomial can sample from unnormalized weights, but every row must be finite, non-negative, and have non-zero mass. That means each hard mask needs a fallback before probabilities are sampled.

def sample_next_token(logits, generated, *, temperature=1.0, top_k=None, top_p=None):
    scores = ban_invalid_tokens(logits.clone(), generated)

    if temperature <= 0:
        raise ValueError("temperature must be positive")
    scores = scores / temperature

    if top_k is not None:
        scores = top_k_filter(scores, top_k)
    if top_p is not None:
        scores = top_p_filter(scores, top_p)

    finite = torch.isfinite(scores)
    empty_rows = ~finite.any(dim=-1)
    if empty_rows.any():
        scores[empty_rows] = logits[empty_rows]  # Defined fallback, not silent NaNs.

    probs = torch.softmax(scores, dim=-1)
    assert torch.isfinite(probs).all()
    assert (probs >= 0).all()
    assert (probs.sum(dim=-1) > 0).all()
    return torch.multinomial(probs, num_samples=1).squeeze(-1)

1. Process logits, then sample: keep masks and warpers in logit space until the final softmax so token ids stay aligned.
2. Fail deliberately: if constraints remove every candidate, use an explicit fallback such as relaxing a constraint, forcing EOS, or reverting to unmasked logits.
3. Test adversarial rows: include prompts where no-repeat rules, schema constraints, or banned-token lists would otherwise leave an empty vocabulary row.

Implementation Drill: Repetition Controls Are Logit Rules

Repetition penalty and no-repeat n-gram constraints are decoder-side rules layered on top of model logits. Use them to reduce loops, but test the exact token-level behavior: subword tokens, punctuation, and copied source phrases can be penalized even when a human would not call them bad repetition. Keep the sampler contract intact after every processor: probabilities passed to torch.multinomial must be finite, non-negative, and have non-zero row mass.

def ban_repeated_bigram(logits, generated):
    if generated.size(-1) < 2:
        return logits

    prefix = generated[..., -1]
    for row in range(generated.size(0)):
        seen = generated[row].unfold(0, 2, 1)
        banned = seen[seen[:, 0] == prefix[row], 1]
        logits[row, banned] = float("-inf")
    return logits

def assert_sampleable(probs):
    assert torch.isfinite(probs).all()
    assert (probs >= 0).all()
    assert (probs.sum(dim=-1) > 0).all()

1. Hard constraint: no-repeat n-gram masking makes some next tokens impossible, so always leave at least one legal token or define a fallback.
2. Soft penalty: repetition penalties reshape logits for tokens that already appeared; they are tunable preferences, not correctness guarantees.
3. Task fit: open-ended writing often benefits from truncation sampling, while extractive QA, code, and citation-heavy answers may need weaker repetition controls so copied spans remain possible.

Structured generation is not just a stronger prompt. A decoder can compile a schema, regex, grammar, or required phrase set into a state machine, then mask out invalid next tokens at each step. The model still scores the allowed tokens, but the sampler is prevented from choosing a token that would make the output impossible to parse.

allowed = constraint_state.allowed_token_ids(prefix)
masked_logits = torch.full_like(logits, float("-inf"))
masked_logits[..., allowed] = logits[..., allowed]
next_token = sample(masked_logits, temperature=0.2)
constraint_state = constraint_state.advance(next_token)

1. Use for contracts: JSON schemas, enums, tool arguments, and required phrases belong in decoding constraints or strict API output modes when downstream code depends on them.
2. Still validate: constrained decoding improves syntax, but application rules, safety checks, and semantic validation remain separate gates.
3. Measure overhead: every step now computes an allowed-token mask, so benchmark latency and batching behavior against unconstrained sampling.

for token in prompt:
    run transformer and store K/V for each layer

while not done:
    run transformer on latest token
    read all old K/V from cache
    append new K/V
    sample next token

Cache memory scales with batch, layers, KV heads, sequence length, head dimension, and bytes per element. GQA and MQA reduce the KV-head term.

kv_bytes = 2 * layers * batch * kv_heads * seq_len * head_dim * bytes_per_element

Worked Example: Size The Cache Before Picking A Policy

Do the memory math before reaching for offloading, quantization, or a static cache. The GQA paper motivates sharing K/V heads as an inference efficiency tradeoff, and Hugging Face's cache strategy docs frame cache classes as memory-versus-latency choices. A quick estimate often shows whether the bottleneck is architecture, context length, batch size, or cache policy.

def gib(n_bytes):
    return n_bytes / (1024 ** 3)

def kv_cache_gib(layers, batch, kv_heads, seq_len, head_dim, dtype_bytes):
    return gib(2 * layers * batch * kv_heads * seq_len * head_dim * dtype_bytes)

case = {
    "layers": 32,
    "batch": 8,
    "query_heads": 32,
    "gqa_kv_heads": 8,
    "seq_len": 16_384,
    "head_dim": 128,
    "dtype_bytes": 2,  # fp16 or bf16
}

mha = kv_cache_gib(case["layers"], case["batch"], case["query_heads"],
                   case["seq_len"], case["head_dim"], case["dtype_bytes"])
gqa = kv_cache_gib(case["layers"], case["batch"], case["gqa_kv_heads"],
                   case["seq_len"], case["head_dim"], case["dtype_bytes"])
print(round(mha, 2), round(gqa, 2))  # 32.0 GiB vs 8.0 GiB

1. Architecture lever: if KV heads dominate memory, GQA/MQA reduces the cache footprint before any serving-layer trick is applied.
2. Policy lever: if the model architecture is fixed, compare dynamic, static, offloaded, and quantized cache choices against the measured bottleneck.
3. Batch lever: continuous batching multiplies cache memory by active sequences, so admission control needs a KV budget, not just a request count.

Implementation Drill: GQA Reduces Stored KV Heads

Grouped-query attention keeps many query heads but shares fewer key and value heads across groups. In a decoder cache, this changes the stored tensor count from num_query_heads to num_kv_heads; it does not change the number of query heads used to produce attention outputs. PyTorch's SDPA can expand grouped K/V heads with enable_gqa=True, but the docs mark the feature experimental and list backend constraints, so keep a reference implementation for shape and logit tests.

def kv_cache_bytes(layers, batch, kv_heads, seq_len, head_dim, dtype_bytes):
    return 2 * layers * batch * kv_heads * seq_len * head_dim * dtype_bytes

query_heads = 32
kv_heads = 8
assert query_heads % kv_heads == 0

baseline = kv_cache_bytes(32, 4, query_heads, 8192, 128, 2)
gqa = kv_cache_bytes(32, 4, kv_heads, 8192, 128, 2)
assert gqa == baseline // (query_heads // kv_heads)

q = torch.randn(4, query_heads, 1, 128, device="cuda")
k = torch.randn(4, kv_heads, 8192, 128, device="cuda")
v = torch.randn(4, kv_heads, 8192, 128, device="cuda")
out = F.scaled_dot_product_attention(q, k, v, enable_gqa=True, dropout_p=0.0)
assert out.shape == q.shape

1. Memory check: compute cache size with KV heads, not query heads, or GQA/MQA savings will disappear from the estimate.
2. Divisibility check: require query heads to be an integer multiple of KV heads before repeating or grouping K/V tensors.
3. Quality check: treat MHA, GQA, and MQA as model architecture choices, not sampler settings; changing them after training requires validation or uptraining.
4. Backend check: benchmark the exact SDPA backend and dtype used in production, because grouped attention support is backend-specific.

Implementation Drill: Test GQA Against A Slow Repeat-KV Reference

GQA implementations usually expose compact K/V tensors, then repeat each K/V head across a group of query heads inside the attention path. The GQA paper frames this as the middle point between MHA and MQA, and PyTorch's SDPA docs require query heads to be divisible by K/V heads when enable_gqa=True. Before benchmarking a fused backend, compare it to an explicit repeat_interleave reference on a tiny tensor so shape mistakes cannot hide inside throughput numbers.

def repeat_kv_for_gqa(x, query_heads):
    # x: batch, kv_heads, seq_len, head_dim
    batch, kv_heads, seq_len, head_dim = x.shape
    assert query_heads % kv_heads == 0
    groups = query_heads // kv_heads
    return x.repeat_interleave(groups, dim=1)

torch.manual_seed(0)
q = torch.randn(2, 8, 3, 16, device="cuda", dtype=torch.float16)
k = torch.randn(2, 2, 5, 16, device="cuda", dtype=torch.float16)
v = torch.randn(2, 2, 5, 16, device="cuda", dtype=torch.float16)

slow = F.scaled_dot_product_attention(
    q, repeat_kv_for_gqa(k, q.size(1)), repeat_kv_for_gqa(v, q.size(1)),
    dropout_p=0.0,
)
fast = F.scaled_dot_product_attention(q, k, v, dropout_p=0.0, enable_gqa=True)
torch.testing.assert_close(fast, slow, rtol=2e-3, atol=2e-3)

1. Reference first: write the repeat-KV path in the same tensor layout your cache uses.
2. Disable dropout: SDPA always applies dropout according to dropout_p, so set it to zero for inference and equivalence tests.
3. Then benchmark: after the outputs match, test the production backend, dtype, sequence lengths, and cache layout.

Implementation Checklist: Trust But Pin The GQA Backend

Treat fused GQA as an implementation detail that must match the model's trained attention contract. PyTorch's SDPA docs mark enable_gqa=True as experimental, restrict it to CUDA Flash attention and math backends, and require query heads to divide evenly by K/V heads. FlashAttention and TensorRT-LLM use the same practical contract: pass fewer K/V heads than query heads, keep the divisibility invariant, then benchmark the actual generation kernel and cache layout instead of assuming the fused path is always selected.

from torch.nn.attention import SDPBackend, sdpa_kernel

def gqa_decode_smoke_test(q, k, v):
    assert q.size(1) % k.size(1) == 0
    assert k.size(1) == v.size(1)
    assert q.is_cuda and k.is_cuda and v.is_cuda

    with sdpa_kernel(backends=[SDPBackend.FLASH_ATTENTION]):
        return F.scaled_dot_product_attention(
            q, k, v,
            dropout_p=0.0,
            enable_gqa=True,
        )

# If this falls back, warns, or fails, keep the repeat-KV reference path
# for correctness and investigate dtype, mask, shape, and backend support.

1. Pin the kernel during tests: use sdpa_kernel when you need to prove that Flash attention, not a silent math fallback, is running.
2. Measure both phases: prefill and one-token decode can dispatch different kernels, so record time-to-first-token and inter-token latency separately.
3. Check cache metadata: fused cache-update APIs still need correct sequence lengths, block tables, RoPE positions, and enough pre-allocated space for appended K/V.
4. Keep a fallback: a slow repeat-KV path is useful for CPU tests, unsupported devices, numerical debugging, and explaining the grouping rule in interviews.

Reading Note: MLA Compresses What The Cache Stores

Multi-head latent attention is another answer to the same decode-time bottleneck that motivates MQA and GQA. DeepSeek-V2 describes MLA as low-rank joint compression of keys and values so inference can cache a compact latent state instead of every expanded per-head K/V tensor; DeepSeek-V3 keeps that architecture for efficient inference. When you compare attention variants, separate three questions: what the model computes per head, what gets stored between decode steps, and what extra projection work is needed when a new token attends to the cache.

attention_cache_ledger = [
    {"layout": "MHA", "stored": "K/V per query head", "decode_tradeoff": "largest cache"},
    {"layout": "GQA", "stored": "K/V per KV group", "decode_tradeoff": "smaller cache"},
    {"layout": "MQA", "stored": "one K/V head", "decode_tradeoff": "smallest simple head cache"},
    {"layout": "MLA", "stored": "compressed latent K/V state", "decode_tradeoff": "projection work plus compact cache"},
]

1. Track stored state: cache estimates should count the tensors that persist across decode steps, not only the attention vectors after projection.
2. Keep RoPE explicit: MLA-style designs may treat positional components differently from compressed content components, so position tests still belong in cache correctness.
3. Benchmark the full path: a smaller cache can improve memory pressure, but the projection, kernel support, and model quality need their own measurements.

Implementation Drill: Think In Logical Tokens And Physical Blocks

PagedAttention separates a request's logical token positions from the physical memory blocks that store its K/V tensors. The vLLM paper and design docs describe fixed-size KV blocks plus a block table, which lets the server allocate cache on demand, reuse shared prefixes, and avoid moving a whole request when a generation grows. When debugging a serving engine, inspect the mapping first: wrong block ids or stale prefix references can look like random attention quality bugs.

def logical_to_block(token_pos, block_size):
    block_index = token_pos // block_size
    offset = token_pos % block_size
    return block_index, offset

request = {
    "block_size": 16,
    "block_table": [42, 87, 91],  # logical block -> physical KV block id
}

logical_block, offset = logical_to_block(token_pos=33, block_size=request["block_size"])
physical_block = request["block_table"][logical_block]
assert physical_block == 91 and offset == 1

1. Allocation invariant: a request with n cached tokens should own or reference exactly ceil(n / block_size) logical blocks.
2. Append invariant: decoding one token only changes the final block until it fills, then appends one new block-table entry.
3. Prefix invariant: shared prompt blocks need reference counts or equivalent ownership tracking before multiple requests can reuse them safely.
4. Scheduler invariant: admission control should count allocated and soon-needed KV blocks, not just request count or prompt length.

Implementation Drill: Admit Requests By KV Blocks, Not Prompt Count

A paged cache makes scheduling look like a small memory-accounting problem. vLLM's design docs define each KV block as a fixed number of tokens at one head, and the PagedAttention paper frames the block table as the indirection that keeps logical sequence positions separate from physical GPU memory. For a serving interview, turn that into an admission test: reserve enough blocks for the prompt plus the promised generation budget before adding a request to the active batch.

import math

def required_kv_blocks(prompt_tokens, max_new_tokens, block_size):
    return math.ceil((prompt_tokens + max_new_tokens) / block_size)

def can_admit(active_block_counts, request, free_blocks, block_size):
    already_reserved = sum(active_block_counts)
    needed = required_kv_blocks(
        request["prompt_tokens"],
        request["max_new_tokens"],
        block_size,
    )
    return already_reserved + needed <= free_blocks

request = {"prompt_tokens": 900, "max_new_tokens": 300}
assert required_kv_blocks(900, 300, block_size=16) == 75

1. Reserve decode room: a request that fits at prefill can still OOM later if the scheduler ignores its maximum generation budget.
2. Update on completion: release blocks when a sequence stops, and decrement prefix-reference counts only when every sharer is done.
3. Track internal waste: each active request can waste up to block_size - 1 token slots in its final block, so smaller blocks reduce waste but increase metadata and scheduling overhead.
4. Report the unit: dashboards should show allocated KV blocks and token slots, not only number of active requests.

Implementation Drill: Cache Correctness Before Cache Speed

Start with a tiny decode loop that compares cached and uncached logits on the same prompt. Only optimize after the last-token logits match within tolerance; most cache bugs are position, mask, dtype, or append order bugs rather than math bugs.

# Pseudocode: one test prompt, deterministic mode, no dropout.
full_logits = model(input_ids, use_cache=False).logits[:, -1]

prefill = model(input_ids[:, :-1], use_cache=True)
step = model(
    input_ids[:, -1:],
    past_key_values=prefill.past_key_values,
    use_cache=True,
)
cached_logits = step.logits[:, -1]

torch.testing.assert_close(cached_logits, full_logits, rtol=1e-3, atol=1e-3)

1. Mask alignment: with PyTorch SDPA, make sure the query length and key/value length imply the causal mask you intend, especially when decoding one token against a longer cache.
2. Cache shape: inspect every layer's K/V tensors as batch, kv_heads, cache_len, head_dim or the local equivalent, then verify that cache length increases by exactly one per decode step.
3. Policy choice: dynamic caches are convenient for correctness tests; static caches can support compiled generation but require careful max-length budgeting; prefix caching is a serving-layer reuse policy, not a conversation memory feature.

Implementation Drill: Grow The Attention Mask With The Cache

In a custom generation loop, the model sees current tokens plus cached past K/V tensors. Hugging Face's cache docs call out the shape contract: the attention mask must cover past_kv_length + new_tokens_length, even when the next model call only passes one new token id. If cached and uncached logits diverge after the first decode step, check this mask length before blaming sampling or numerical precision.

def extend_attention_mask(attention_mask, new_tokens=1):
    batch = attention_mask.size(0)
    extra = attention_mask.new_ones((batch, new_tokens))
    return torch.cat([attention_mask, extra], dim=-1)

prefill = model(input_ids, attention_mask=attention_mask, use_cache=True)
next_ids = prefill.logits[:, -1:].argmax(dim=-1)
next_mask = extend_attention_mask(attention_mask, new_tokens=1)

step = model(
    next_ids,
    attention_mask=next_mask,
    past_key_values=prefill.past_key_values,
    use_cache=True,
)
assert next_mask.size(-1) == input_ids.size(-1) + 1

1. Mask invariant: at every decode call, the mask length should equal cached token count plus the number of new input tokens.
2. Batch invariant: append one visible position for every active sequence; finished sequences need an explicit stopping policy, not a shorter mask.
3. Debug invariant: log input_ids.shape, mask shape, and per-layer cache length together so off-by-one bugs are visible in one line.

Implementation Drill: Keep KV Cache Out Of Training Graphs

A decode cache is an inference-time shortcut, not a replacement for teacher-forced training. Hugging Face's cache explanation warns that caching should be used for inference, because training needs gradients through every token position and usually processes the full shifted sequence in one pass. In practice, make use_cache=False part of the training/eval contract, then test cached generation separately under torch.inference_mode().

def training_step(model, input_ids, labels):
    model.train()
    out = model(input_ids, labels=labels, use_cache=False)
    out.loss.backward()
    return out.loss

@torch.inference_mode()
def generate_step(model, input_ids):
    model.eval()
    return model(input_ids, use_cache=True)

1. Separate contracts: train with full-sequence logits and no cache; generate with cached prefill and one-token decode steps.
2. Disable gradients: cache tensors should not keep autograd history or they can leak memory across decode steps.
3. Test both paths: loss tests catch training regressions, while cached-versus-uncached logit tests catch inference state bugs.

Implementation Drill: Non-Square Causal Masks Need Position Tests

Cached decoding often calls attention with one query token and many cached key/value tokens. PyTorch documents is_causal=True as using an upper-left causal alignment for non-square masks, so do not assume that a one-token query automatically means "attend to every cached key." Test the exact mask, position ids, and cache length that your model wrapper sends to the attention kernel.

def make_decode_attention_mask(cache_len, device):
    # A decode query at absolute position cache_len should see all cached keys
    # plus its own current key after the new K/V is appended.
    return torch.ones(1, cache_len + 1, dtype=torch.bool, device=device)

q = torch.randn(batch, heads, 1, head_dim, device=device)
k = torch.randn(batch, heads, cache_len + 1, head_dim, device=device)
v = torch.randn(batch, heads, cache_len + 1, head_dim, device=device)

out = F.scaled_dot_product_attention(
    q, k, v,
    attn_mask=make_decode_attention_mask(cache_len, q.device),
    dropout_p=0.0,
)

1. Square prefill path: use the normal causal path for full prompts and compare it against a slow reference.
2. Non-square decode path: separately test one-token and multi-token decode queries against the intended absolute positions.
3. Boolean polarity: in PyTorch SDPA, True means the element participates in attention, which is the opposite of some padding-mask APIs.
4. Failure signal: if cached and uncached logits diverge only after the first generated token, inspect mask alignment before tuning numerical tolerances.

Implementation Drill: Choose Cache Classes From The Constraint

Cache policy is not one universal speed switch. Hugging Face's cache strategy docs describe DynamicCache as the default growing cache, StaticCache as the compile-friendly fixed-shape option, offloaded cache as a GPU-memory relief valve, and quantized cache as a lower-memory path that can add latency on short contexts. A practical implementation should pick the cache from the bottleneck the benchmark actually shows.

def select_cache_profile(workload):
    if workload.cached_logits_do_not_match_uncached:
        return {"cache_implementation": "dynamic", "reason": "debug correctness"}
    if workload.decode_compile_needed and workload.length_bucket_is_tight:
        return {"cache_implementation": "static", "reason": "stable decode shapes"}
    if workload.gpu_oom_on_long_context:
        return {"cache_implementation": "offloaded", "reason": "fit before speed"}
    if workload.kv_memory_is_limit and workload.long_context_dominates:
        return {"cache_implementation": "quantized", "reason": "reduce KV footprint"}
    return {"cache_implementation": "dynamic", "reason": "simple baseline"}

1. Dynamic first: use the default growing cache while validating positions, masks, and append order.
2. Static for compile: pre-allocate only when request lengths are bucketed tightly enough that unused slots do not dominate memory or attention work.
3. Offload for fit: move cache layers through CPU memory when OOM is the failure, then charge the copy overhead in latency reports.
4. Quantize for long contexts: test quality and latency before launching, because smaller KV tensors can still cost extra dequantization work.

Benchmark Drill: Quantized KV Cache Needs A Quality Gate

KV-cache quantization targets the memory term that grows with active batch and context length. Hugging Face documents quantized cache as the low-memory cache option and cautions that it can hurt latency on short contexts. KIVI shows why keys and values may want different grouping rules, while vLLM's FP8 cache docs make scale calibration an explicit serving choice. Treat cache quantization like a deployment profile: prove it fits more useful work, then prove it keeps answers stable.

def kv_cache_gib(layers, batch, kv_heads, seq_len, head_dim, bytes_per_value):
    return 2 * layers * batch * kv_heads * seq_len * head_dim * bytes_per_value / 2**30

baseline = kv_cache_gib(32, 8, 8, 32768, 128, 2)   # bf16 cache
fp8 = kv_cache_gib(32, 8, 8, 32768, 128, 1)        # FP8 cache
int4 = kv_cache_gib(32, 8, 8, 32768, 128, 0.5)     # idealized 4-bit cache

print(round(baseline, 1), round(fp8, 1), round(int4, 1))
# 16.0 GiB, 8.0 GiB, 4.0 GiB before metadata, scales, and residual buffers

1. Fit test: report max admitted batch, prompt length, and generated tokens under the same GPU memory cap.
2. Latency test: measure time to first token and inter-token latency separately; dequantization or scale handling can erase wins on short contexts.
3. Quality test: compare exact-match retrieval prompts, long-context QA, code completion, and regression golden cases against the unquantized cache.
4. Calibration test: record whether scales are default, random-token, or dataset-calibrated, because that choice changes the accuracy and portability story.

Implementation Drill: Static Cache Is A Compile Contract

A dynamic KV cache is convenient because it grows with generation, but changing cache shapes can block just-in-time optimizations. A static cache pre-allocates a maximum length so the decode step sees steadier shapes and can be paired with torch.compile. The tradeoff is explicit: unused cache slots still consume memory and may add masked attention work, so static cache is best tested on batches with similar prompt and output lengths.

def choose_cache(request_lengths, max_context, compile_decode):
    spread = max(request_lengths) - min(request_lengths)
    similar_lengths = spread <= 0.25 * max_context

    if compile_decode and similar_lengths:
        return {"cache_implementation": "static", "max_cache_len": max_context}
    return {"cache_implementation": "dynamic"}

1. Prove correctness first: compare cached and uncached logits before compiling the decode step.
2. Bucket by shape: group requests by nearby context budgets so a static cache does not waste most of its slots on short prompts.
3. Watch recompile signals: PyTorch documents guard failures and dynamic-shape tracing as the clues that shape variation is defeating compile reuse.
4. Measure memory separately: reduce-overhead and CUDA graph paths can lower Python overhead, but PyTorch notes that this may cache extra workspace memory.

Autoregressive decoding is serial because the target model normally runs once per accepted token. Speculative decoding asks a cheaper draft model to propose several candidate tokens, then uses one target-model forward pass to verify the proposed block. If the acceptance rate is high, the large model commits multiple tokens for nearly the cost of one verification pass; if the draft model is poorly matched, rejected tokens erase the speedup.

while not done:
    draft_tokens = assistant.generate(prefix, max_draft_tokens)
    target_logits = target.verify(prefix + draft_tokens)
    accepted = accept_prefix_by_target_distribution(draft_tokens, target_logits)
    prefix.extend(accepted or sample_one_from_target(target_logits))

1. Measure acceptance: track accepted draft tokens per target verification step, not just tokens per second.
2. Keep distributions honest: the verifier must preserve the target model's intended sampling distribution.
3. Choose assistants carefully: fast is useful only when the assistant's guesses are often accepted.

Implementation Drill: Test Acceptance On A Toy Distribution

The easy mistake is to treat speculative decoding as "use the draft when it agrees with the target." The original speculative decoding and speculative sampling papers are stricter: acceptance and fallback sampling must leave the target model's distribution unchanged. Before optimizing kernels, write a tiny categorical test where the draft distribution q and target distribution p are known, then verify that accepted tokens use min(1, p_i / q_i) and rejected tokens are resampled from the positive residual mass.

def speculative_acceptance_prob(target_p, draft_q, draft_token):
    p_i = float(target_p[draft_token])
    q_i = float(draft_q[draft_token])
    return min(1.0, p_i / max(q_i, 1e-12))

def residual_after_rejection(target_p, draft_q):
    residual = (target_p - draft_q).clamp_min(0)
    return residual / residual.sum()

target_p = torch.tensor([0.50, 0.30, 0.15, 0.05])
draft_q = torch.tensor([0.40, 0.40, 0.10, 0.10])

assert speculative_acceptance_prob(target_p, draft_q, 0) == 1.0
assert round(speculative_acceptance_prob(target_p, draft_q, 1), 2) == 0.75
fallback = residual_after_rejection(target_p, draft_q)
torch.testing.assert_close(fallback, torch.tensor([2 / 3, 0.0, 1 / 3, 0.0]))

1. Acceptance check: never accept with probability above one; tokens underpredicted by the draft can be accepted whenever proposed.
2. Residual check: after rejection, sample only from tokens where the target assigns more mass than the draft.
3. Monte Carlo check: run many one-step trials and compare empirical token frequencies to the target distribution before testing long generations.

Implementation Drill: Match The Assisted-Decoding Contract

Assisted decoding is easy to call but easy to benchmark poorly. Hugging Face documents several proposer paths: a smaller assistant_model, prompt lookup decoding, self-speculation with early exits, and universal assisted decoding when the assistant and target use different tokenizers. vLLM exposes a broader production menu, including draft models, EAGLE, MTP, n-gram, suffix decoding, and custom proposers. Treat the proposer choice as part of the decoding contract, then test whether it reduces target-model forward passes for the actual prompt slice.

spec_decode_case = {
    "target": "large-model@revision",
    "proposer": "assistant_model | prompt_lookup | early_exit | ngram | suffix",
    "same_tokenizer": True,
    "draft_tokens": 4,
    "metrics": [
        "accepted_tokens_per_verify",
        "target_forwards_per_output_token",
        "p95_inter_token_latency",
        "fallback_rate",
    ],
}

1. Start same-tokenizer: paired model families are the simplest baseline because candidate tokens do not need retokenization.
2. Use UAD deliberately: cross-tokenizer assistants need explicit target and assistant tokenizers, plus alignment checks before claiming speedup.
3. Try lookup for copied text: prompt lookup or n-gram speculation can help summarization, translation, and retrieval-heavy answers without loading a draft model.
4. Compare to no speculation: keep the same target model, prompt, sampling settings, and max tokens so latency changes come from the proposer path.

Benchmark Drill: Tune Lookahead From Acceptance, Not Hope

The original speculative decoding algorithm uses a fixed number of draft tokens, but current assisted-generation implementations can vary the speculation length during generation. Hugging Face describes dynamic speculation as the default assisted-decoding mode in recent Transformers releases, with knobs such as assistant_confidence_threshold, num_assistant_tokens_schedule, and num_assistant_tokens. Treat those as benchmark variables: more draft tokens help only when they are cheap and likely to be accepted by the target verifier.

def report_speculation_run(run):
    accepted = run.accepted_draft_tokens
    verified = run.target_verify_steps
    produced = run.output_tokens
    return {
        "accepted_per_verify": accepted / max(verified, 1),
        "target_forwards_per_token": verified / max(produced, 1),
        "assistant_tokens_wasted": run.rejected_draft_tokens,
        "p95_itl_ms": percentile(run.inter_token_latency_ms, 95),
    }

1. Sweep lookahead: compare constant, heuristic, and dynamic schedules on the same prompts before changing the assistant model.
2. Cap bad guesses: lower confidence or shorter draft limits can win when the assistant is uncertain and rejections are frequent.
3. Separate costs: report assistant forward time, target verification time, retokenization time, and queueing time separately.
4. Keep a rollback profile: disable speculation when traffic, prompts, or sampling settings push acceptance below the break-even point.

Interactive LLM serving mixes requests with different prompt lengths and decode lengths. A scheduler tries to keep the GPU busy while protecting tail latency, memory limits, and fairness.

Sources: vLLM / PagedAttention, vLLM paged attention design, FlashAttention, The Curious Case of Neural Text Degeneration, Hugging Face generation strategies, Hugging Face generation configuration, Fast Inference from Transformers via Speculative Decoding, Hugging Face assisted decoding, Hugging Face dynamic speculation lookahead, Hugging Face universal assisted generation, vLLM speculative decoding, Hugging Face KV cache strategies, Hugging Face cache explanation, KIVI KV-cache quantization, vLLM quantized KV cache, PyTorch torch.compile, PyTorch torch.multinomial, PyTorch SDPA, PyTorch SDPA causal mask notes, Grouped-query attention, DeepSeek-V2 / Multi-head Latent Attention, DeepSeek-V3 technical report, vLLM automatic prefix caching, OpenAI structured outputs, Hugging Face constrained beam search, Outlines structured generation docs.