KV-cache and speculative decoding on a DS interview

Why this comes up on DS interviews

If you are interviewing for a DS / MLE role with an LLM lean at OpenAI, Anthropic, Meta, Google, or any infra-heavy scale-up, expect at least one round to leave model accuracy alone and grill you on inference economics. The reason is simple: training is a one-time capex hit, but serving a chatbot at billions of tokens per day is an ongoing opex bonfire. The team that ships an 8B model with smart inference beats the team that ships a 70B model with a naive loop.

The three building blocks interviewers cycle through are KV-cache, speculative decoding, and FlashAttention, usually wrapped in a question like "your p99 latency is 4 seconds — what would you look at first?" The answer is rarely "train a smaller model." It is almost always one of these systems-level levers.

Read the rest of this post as a checklist of what an L5 ML engineer should be able to whiteboard cold.

Why KV-cache exists

In a Transformer decoder, each new token needs query, key, and value projections to run attention against every previous token in the sequence. Without a cache, the model has no memory between forward passes, so at step N it recomputes K and V for tokens 1 through N-1 just to compute attention for token N.

Without cache:
Token N: compute K, V for tokens [1..N]   -> O(N) work per step
Across N steps: O(N^2) total

This is wasteful because K and V for previous tokens never change — the autoregressive mask guarantees a token at position i cannot see future tokens. So the K/V projections from earlier steps can be reused verbatim.

With KV-cache:
Token N: read cached K, V for [1..N-1] + compute K, V for token N
Across N steps: O(N) total

The result: on long sequences the cache is the difference between a model that streams at 50 tokens/sec and one that crawls at 3 tokens/sec. Every production inference engine — vLLM, TensorRT-LLM, SGLang, llama.cpp — ships KV-caching on by default. If a candidate says "we should add a cache," that is a red flag; the cache is the baseline, not the optimization.

Sizing the KV-cache

The KV-cache is a tensor of shape (2, num_layers, num_heads, head_dim, seq_length) in your chosen precision. The 2 is for K and V. Memory in bytes is straightforward:

size = 2 * num_layers * num_heads * head_dim * seq_length * bytes_per_value

Worked example for Llama-2 70B with grouped query attention (80 layers, 8 KV heads, head_dim 128, FP16 = 2 bytes):

For a 7B model at FP16 weights (~14 GB), the KV-cache for a single 32k-token request can exceed the model weights themselves. This is the punchline interviewers want you to deliver out loud.

Common mitigations:

• GQA (grouped query attention). Use fewer K/V heads than Q heads — Llama-2 70B ships with 8 KV heads vs 64 Q heads, giving an 8x cache reduction with negligible quality loss.
• MQA (multi-query attention). One K/V head shared across all Q heads. Aggressive; sometimes hurts quality on long context.
• Sliding window attention. Keep only the last W tokens (Mistral 7B uses W = 4096). Bounds cache size regardless of sequence length.
• PagedAttention (vLLM). Cache is allocated in fixed-size pages, similar to virtual memory. Eliminates internal fragmentation and enables prefix sharing across requests.
• Quantized cache. Store K/V in INT8 or even INT4. 2-4x memory savings at a small perplexity cost.

Speculative decoding

Speculative decoding is the single most cited inference trick of the last two years. The intuition: a small draft model (say, Llama 3.2 1B) is much faster per token than the target model (Llama 3.1 70B), but its predictions are correlated enough that the big model agrees with the draft most of the time.

Step 1: Draft (1B) generates 5 candidate tokens forward:
        "the cat sat on the"

Step 2: Target (70B) verifies all 5 in a single parallel forward pass:
        - Tokens 1-4 accepted ("the cat sat on")
        - Token 5 rejected (target predicts " a" instead of " the")

Step 3: Continue from "the cat sat on a", redraft 5 more.

The key win is that the target model runs one batched forward pass instead of five sequential ones. Because LLM inference at small batch sizes is bandwidth-bound, that batched pass is barely slower than a single-token pass. Effective speedup is typically 2-3x on chat workloads and can reach 4-5x on highly predictable text like code completion.

Production frameworks that ship it: vLLM, TensorRT-LLM, SGLang, and the EAGLE / Medusa families that train the draft head jointly with the target.

FlashAttention

The standard attention formula is softmax(Q K^T / sqrt(d)) V. The middle term Q K^T is an N x N matrix where N is the sequence length. At N = 32,768, that is a 1-billion-entry matrix per layer per head, which blows past GPU SRAM and forces expensive trips to HBM.

FlashAttention (Dao 2022) rewrites the kernel to never materialize Q K^T in full. It loads small tiles of Q, K, V into SRAM, computes a streaming softmax, accumulates the output, and discards intermediates. Backward pass recomputes attention on the fly instead of storing it.

If you are running v1 on an H100, you are leaving roughly half your throughput on the floor. This is a frequent gotcha interviewers love to surface.

FlashAttention enables long context not because it changes the algorithm — outputs are bit-identical to vanilla attention up to floating-point reordering — but because it makes a previously OOM-prone operation fit in memory.

Batching and continuous batching

Single-request inference wastes the GPU. A modern H100 can run dozens of decode steps in parallel for roughly the same wall-clock as one. The question is how to keep that parallel slot filled.

Static batching groups N requests together at the start, runs them in lockstep, and finishes only when the slowest one is done. On mixed workloads where one user asks for 5 tokens and another asks for 500, you waste 99% of the GPU for the last 495 steps.

Continuous batching (Orca 2022, popularized by vLLM) schedules at the token level instead of the request level. When any request in the batch finishes, its slot is immediately refilled with a new request. Throughput improves 5-10x on production chat traffic versus static batching.

vLLM, TGI, and TensorRT-LLM all use continuous batching by default. If you are rolling your own inference server in 2026, you are reinventing the wheel — and probably a worse one.

Common pitfalls

The first pitfall is not turning on the KV-cache at all. Some research-grade implementations leave it off so the code is simpler, and engineers who copy that code into production silently take an order-of-magnitude latency hit. Always verify use_cache=True (or the framework equivalent) and benchmark with and without it before declaring the model "too slow."

A second trap is storing the cache in FP32. FP32 doubles cache memory versus FP16, which on long-context workloads can be the difference between fitting and OOM. FP16 and BF16 are fine for cache in essentially all production setups, and INT8 or INT4 are viable with a quality check. If the candidate cannot tell you which precision their cache uses, they have not benchmarked their stack.

The third is speculative decoding without rejection sampling. Some naive implementations accept whatever the draft says without verifying against the target's logits, producing subtle quality regressions that pass smoke tests but fail offline evals. The verification step is not optional; it is what makes speculative decoding lossless.

A fourth pitfall is running FlashAttention v1 on an H100. v1 was tuned for A100 and earlier; v2 and v3 add Hopper-specific optimizations like TMA, asymmetric warp scheduling, and FP8 support. Teams that upgraded their GPUs but never upgraded their attention kernel routinely leave 30-50% throughput on the table.

The fifth — and the one most often surfaced in system design rounds — is single-request serving. Spinning up one model instance per request, with no batching, no continuous scheduling, and no paged cache, is the single biggest reason small teams burn through their inference budget. Use vLLM, TGI, or TensorRT-LLM. If you must build your own, at minimum implement continuous batching.

Related reading

• Transformer architecture for data science interviews
• BERT vs GPT for data science interviews
• Attention mechanism for data scientists
• GPT architecture for data science interviews
• Inference optimization for data scientists
• ML latency optimization on DS interviews

If you want to drill LLM systems questions like this every day, NAILDD is launching with hundreds of DS interview problems aimed at exactly this pattern.

FAQ

Does KV-cache apply to encoder models like BERT?

No. Encoders process the full input sequence in one parallel forward pass — there is no concept of "previous tokens" the way there is for an autoregressive decoder, so the optimization does not apply. KV-caching is a feature of generation, not understanding. If you are running a BERT-style retrieval or classification pipeline, latency wins come from quantization, distillation, ONNX/TensorRT export, and batching instead.

How big does speculative decoding's draft model need to be?

Rule of thumb: the draft model should be 10-30x smaller than the target, share the same tokenizer, and be trained on similar data. Llama 3.2 1B as a draft for Llama 3.1 70B is a canonical pair; Qwen 2.5 0.5B for Qwen 2.5 32B is another. If the draft is too small, acceptance rate craters and the verification overhead dominates. If it is too big, you save no compute. Token acceptance rates of 0.6 to 0.8 are the productive range.

Is FlashAttention only for training, or also for inference?

Both. The original FlashAttention paper emphasized training, but inference engines like vLLM and TensorRT-LLM use FlashAttention kernels for the prefill phase (the initial forward pass over the prompt) and a variant called FlashDecoding / FlashAttention-2-Decode for the per-token decode step. The decode-time variant is optimized for the case where the query length is 1 but the K/V length can be very large.

What is PagedAttention and how does it relate to KV-cache?

PagedAttention is vLLM's contribution: instead of allocating a contiguous KV-cache block for each request, it splits the cache into fixed-size pages (typically 16 tokens) and maintains a page table per request. This eliminates internal fragmentation, enables prefix sharing between requests with the same system prompt, and unlocks 2-4x higher throughput on shared workloads. It is orthogonal to FlashAttention — you can and should use both.

When does speculative decoding stop being worth it?

Two regimes. First, very large batch sizes — once the target model is fully compute-bound, the verification batched pass is no longer "free," so the speedup shrinks. Second, very unpredictable text like creative writing with high sampling temperature — token acceptance rates drop below 0.5 and verification overhead exceeds the draft savings. For chat with temperature 0.7 and batch sizes under 32, speculative decoding is almost always a win.

Should I quantize the KV-cache too, or only the weights?

Both, but in that order. Weight quantization (INT8, INT4 via GPTQ or AWQ) gives the biggest one-time memory win and is well-understood. KV-cache quantization comes second; it matters most for long-context workloads where the cache dominates memory. INT8 KV-cache is essentially free in quality, and INT4 is viable if you eval carefully on your downstream tasks. The pitfall is mixing FP16 weights with FP32 cache — a surprisingly common misconfiguration that wastes memory for no benefit.