Kv Cache Speculative Decoding Tutorial En

§0 TL;DR Cheat Sheet

1. KV cache formula: per-sample memory $= 2 \cdot L_\text{ctx} \cdot N_\text{layers} \cdot N_\text{kv\_heads} \cdot d_\text{head} \cdot \text{bytes}$; the "2" comes from K+V. LLaMA-3-70B (GQA, $H_\text{kv}=8$) at 4K context fp16 ≈ 1.25 GB/sample — which is why MHA is not used.
2. Prefill vs Decode asymmetry: prefill processes the entire prompt ($O(L^2)$ FLOPs, compute-bound); decode generates 1 token per step ($O(L)$ FLOPs per step, but must read the entire KV cache — memory-bandwidth-bound). This asymmetry explains every modern inference system design.
3. PagedAttention (Kwon et al., SOSP 2023, vLLM): partition KV cache into pages, use a block table to resolve fragmentation; memory utilization improves from ~70% to ~96%.
4. Continuous batching (Orca, Yu et al., OSDI 2022): iteration-level scheduling — completed requests don't wait for the whole batch; combined with PagedAttention, these are vLLM's two pillars.
5. MQA → GQA → MLA: MQA (Shazeer 2019) shares K/V extremely aggressively, slightly hurting quality; GQA (Ainslie et al., EMNLP 2023) groups into $G$ as a middle ground; MLA (DeepSeek-V2, May 2024) uses low-rank latent $c_t^{KV}$ + decoupled RoPE — RoPE cannot be absorbed into latent compression directly, so a small independent dimension $d_\text{head}^R$ must carry position.
6. Speculative Decoding core (Leviathan et al., ICML 2023; Chen et al., 2023): a small draft model $q$ proposes $K$ tokens, and the target model $p$ verifies in parallel in one forward; rejection sampling guarantees the output distribution is exactly equivalent to $p$ (exact, not approximate).
7. Acceptance probability formula: per-token acceptance rate $\alpha = \mathbb{E}_{x \sim q}[\min(1, p(x)/q(x))]$; expected generated tokens $E[\tau] = \dfrac{1-\alpha^{K+1}}{1-\alpha}$ ($K$ is the draft length, plus the final bonus token).
8. Medusa / EAGLE / Lookahead: Medusa (Cai et al., ICML 2024) multi-head + static tree attention; EAGLE/2/3 (Li et al., 2024-2025) feature-level draft + dynamic tree; Lookahead Decoding (Fu et al., ICML 2024) Jacobi iteration — all different drafters under the same acceptance-rate framework.

§1 Intuition

1.1　Why is inference systems work so "counter-intuitive"

In training we worry about FLOPs: model size, batch size, when OOM happens. But deploying a 70B model — the bottleneck is often not compute but HBM bandwidth and memory, both eaten by the KV cache.

A core mental model:

Modern LLM inference is bandwidth-bound during decode and memory-bound during long-context prefill, not compute-bound.

KV cache is the classic "swap recomputation for storage + bandwidth" trade. Once you cache all the K/V of the entire conversation history, generating each new token only requires:

• Computing Q/K/V for the new token (tiny compute)
• Appending the new K/V to the cache
• One attention with the new Q against the full cache

But the cost is: for every new token, the entire KV cache must be read from HBM to SRAM — which is why 8× H100 + LLaMA-3-70B at batch=1 decode runs far below theoretical FLOPs utilization (often 1-5%).

Speculative decoding attacks exactly this asymmetry: since decode is bandwidth-bound while GPU compute is spare, why not compute $K$ candidate tokens in one forward, since weights are read only once?

1.2　Differences from training-time attention

A common interview cross-examination: "can training use KV cache?" — No. Training computes all positions at once; there's no "existing K/V waiting to be appended". Using KV cache during training is a beginner mistake.

§2 KV Cache Memory Accounting

2.1　Exact formula

Per sample (batch=1), fp16:

$$\boxed{\;\text{KV cache}_\text{bytes} = 2 \cdot L_\text{ctx} \cdot N_\text{layers} \cdot N_\text{kv\_heads} \cdot d_\text{head} \cdot \text{bytes\_per\_elem}\;}$$

Each factor:

• 2: one K tensor + one V tensor
• $L_\text{ctx}$: current context length (prompt + tokens generated so far)
• $N_\text{layers}$: Transformer layers (each has its own cache)
• $N_\text{kv\_heads}$: number of K/V heads. In MHA = $H$; MQA = 1; GQA = $G$ ($1 < G < H$)
• $d_\text{head}$: head dimension, usually 64 or 128
• bytes_per_elem: fp16 = 2, fp8/int8 = 1, int4 = 0.5

2.2　Cache size for some concrete models (fp16, $L_\text{ctx}=4096$)

Plug into the §2.1 formula $2 \cdot L_\text{ctx} \cdot N_\text{layers} \cdot N_\text{kv\_heads} \cdot d_\text{head} \cdot 2\text{B}$:

DeepSeek-V2 cache uses $d_c=512$ (latent dim shared by K/V via the same latent vector) + decoupled RoPE component $d_r=64$ (shared across all heads), following $N_\text{layers} \cdot L_\text{ctx} \cdot (d_c + d_r) \cdot \text{bytes}$ (no "$\times 2$", since K/V no longer store separately). Under fp16, $60 \cdot 4096 \cdot 576 \cdot 2 \approx$ 0.27 GB / sample — an order of magnitude smaller than same-scale MHA.

2.3　Batch dimension

In real serving, KV cache must also be multiplied by active batch size. A 70B + 4K context + GQA deployment on 8×A100 80GB:

• weights ≈ 140 GB (fp16, sharded across cards)
• single-sample cache ≈ 1.25 GB
• remaining usable memory ≈ $8 \times 80 - 140 \approx 500$ GB → subtract activations + framework overhead, ~400 GB allocatable to KV cache
• theoretical max batch ≈ $400 / 1.25 \approx$ 320 — but fragmentation prevents reaching it in practice; hence PagedAttention.

§3 Prefill vs Decode Asymmetry

3.1　Per-phase FLOPs / bandwidth differences

Let prompt length be $L$, hidden $D$, FFN intermediate $4D$, and $N$ layers. Per-layer attention + FFN approximately:

$$\text{FLOPs}_\text{layer} \approx \underbrace{6BLD^2}_\text{QKV proj (3 mat)} + \underbrace{4BL^2 D}_\text{attention (QK + AV)} + \underbrace{2BLD^2}_\text{O proj} + \underbrace{16 BLD^2}_\text{FFN (up + down)}$$

Prefill with $L = L_\text{prompt}$: every term is $\Omega(L D^2)$ or $\Omega(L^2 D)$ — compute-bound, GPU saturated.

Decode with $L=1$ (one token per step), but the attention term becomes $4 B L_\text{ctx} D$ (QK and AV each $2 B L_\text{ctx} D$, since $L_q=1, L_k=L_\text{ctx}$):

• per-step FLOPs ≈ $\Theta(B (L_\text{ctx} D + D^2))$
• per-step HBM read: weights ≈ total model bytes (70B fp16 ≈ 140 GB) + KV cache ≈ $2 B L_\text{ctx} N H_\text{kv} d_\text{head} \cdot \text{bytes}$
• Arithmetic intensity (FLOPs / byte) $\to$ far below the GPU roofline (A100 BF16 ≈ 200 FLOPs/byte)

So decode at small-to-medium batch is memory-bandwidth-bound (until batch grows enough to saturate compute) — memorize this single line and you cover 80% of the interview score.

3.2　Chunked Prefill (Sarathi-Serve)

Sarathi-Serve (Agrawal et al., OSDI 2024) splits long prefill into equal-size chunks; each iteration schedules one prefill chunk + a batch of decode tokens together:

       Traditional: prompt 4096 → one prefill saturates GPU → decode stuck 100+ ms
       Sarathi:     prompt split into 4 chunks × 1024 → each iteration coalesces with decode

Stall-free schedule: every iteration contains decode + prefill-chunk coalesced; decode latency jitter disappears. The paper reports Mistral-7B on a single A100 with 2.6× over vLLM, and Yi-34B on 2×A100 with 3.7×.

3.3　Continuous Batching (Orca)

Traditional static batching: wait for the entire batch to finish before admitting new requests — short requests get blocked by long ones. Orca (Yu et al., OSDI 2022) changes the scheduling granularity from request to iteration: every forward checks the batch, kicks out completed sequences (EOS), and admits new requests in their place.

3.4　Prefix Caching

If multiple requests share a prompt prefix (e.g., system prompt, few-shot prompt), the same KV cache can be reused:

• vLLM's prefix caching: index cache pages by hash(prompt prefix); on a hit, skip prefill
• For prompt-heavy services like ChatGPT, prefix cache hit rate can reach 90%+, drastically cutting prefill cost
• Implementation key: page-level sharing + COW (copy-on-write); branched requests only write to their own pages

§4 KV Cache Optimization Routes

4.1　Overview of routes

4.2　MQA / GQA formula recap

MHA: each head has independent $W_k^{(h)}, W_v^{(h)} \in \mathbb{R}^{D \times d_\text{head}}$; total $H \cdot 2 \cdot D \cdot d_\text{head} = 2 D^2$ parameters (K+V).

MQA (Shazeer 2019): $H$ Q-heads share 1 set of K/V. $W_k, W_v \in \mathbb{R}^{D \times d_\text{head}}$, K+V parameters $= 2 D d_\text{head}$, $H$× smaller. At forward, K and V are broadcast to all $H$ heads for attention.

GQA (Ainslie 2023): $H$ Q-heads divided into $G$ groups, each group shares one K/V set. K+V parameters $= 2 G D d_\text{head}$. LLaMA-2-70B uses $H=64, G=8 \Rightarrow$ KV cache shrunk 8×.

4.3　MLA: low-rank latent K/V (DeepSeek-V2's core innovation)

4.3.1　Compression / decompression

Input hidden state $h_t \in \mathbb{R}^D$. MLA introduces:

$$c_t^{KV} = W^{DKV} h_t \in \mathbb{R}^{d_c}, \quad d_c \ll H d_\text{head}$$

Then only $c_t^{KV}$ is cached. To generate the $i$-th head's K and V:

$$k_t^{C, (i)} = W^{UK, (i)} c_t^{KV}, \quad v_t^{(i)} = W^{UV, (i)} c_t^{KV}$$

where $W^{UK, (i)}, W^{UV, (i)} \in \mathbb{R}^{d_\text{head} \times d_c}$ are head-specific up-projections.

Similarly, Q is also low-rank compressed (this step is optional and mainly saves training memory rather than inference):

$$c_t^Q = W^{DQ} h_t, \quad q_t^{C, (i)} = W^{UQ, (i)} c_t^Q$$

4.3.2　Cache size comparison

DeepSeek-V2 takes $d_c = 4 d_\text{head}$; relative to MHA ($2 H d_\text{head}$), compression ratio is about $H/2$ — for $H=128$ that's roughly 64×.

4.3.3　Inference equivalence transformation (the absorb trick)

In a naive implementation, every step would have to un-project $c_t^{KV}$ back to $k_t, v_t$ then compute attention — which defeats the "save cache" benefit (still doing the up-projection). MLA's elegance is in matrix absorption:

Attention score (ignoring RoPE, content only):

$$q_t^{(i)\top} k_s^{(i)} = (W^{UQ, (i)} c_t^Q)^\top (W^{UK, (i)} c_s^{KV}) = c_t^{Q\top} \underbrace{(W^{UQ, (i)\top} W^{UK, (i)})}_\text{constant matrix \tilde W^{QK,(i)}} c_s^{KV}$$

$\tilde W^{QK, (i)} \in \mathbb{R}^{d_c' \times d_c}$ is fixed at inference time, and can be pre-multiplied once when loading the model. So:

• inference only caches $c_s^{KV}$
• attention score computation is just $c_t^{Q\top} \tilde W^{QK, (i)} c_s^{KV}$, no K/V reconstruction at all
• Similarly $W^{UV, (i)}$ can be absorbed into the output projection $W^O$

This is why MLA has tiny cache but inference FLOPs don't blow up: matrix absorption decouples "save cache" from "save compute".

4.4　MLA's RoPE problem — why decoupling is necessary

4.4.1　Problem: RoPE breaks absorb

RoPE injects positional information as a rotation matrix $R_t \in \mathbb{R}^{d_\text{head} \times d_\text{head}}$ into Q and K:

$$q_t^{\text{RoPE}, (i)} = R_t q_t^{(i)}, \quad k_s^{\text{RoPE}, (i)} = R_s k_s^{(i)}$$

Attention score becomes:

$$q_t^{\text{RoPE}, (i)\top} k_s^{\text{RoPE}, (i)} = q_t^{(i)\top} R_t^\top R_s k_s^{(i)} = q_t^{(i)\top} R_{s-t} k_s^{(i)}$$

(using $R_t^\top R_s = R_{s-t}$, the essence of RoPE — relative position depends only on $s-t$.)

Now plug in the latent form:

$$q_t^{\text{RoPE}, (i)\top} k_s^{\text{RoPE}, (i)} = c_t^{Q\top} \underbrace{W^{UQ, (i)\top} R_{s-t} W^{UK, (i)}}_\text{not constant — depends on (s-t)} c_s^{KV}$$

The middle matrix varies with $s-t$ — meaning it can no longer be pre-absorbed into a constant matrix. Each $(t, s)$ pair must compute $R_{s-t}$ on the fly; the absorb trick fails outright, and compute reverts to MHA-equivalent.

4.4.2　Solution: separate the RoPE component into an independent channel

DeepSeek-V2's solution: give RoPE an independent small-dim channel.

• Latent channel (no RoPE): carries content information, cached as $c_t^{KV} \in \mathbb{R}^{d_c}$, with attention scoring via absorb
• RoPE channel (with RoPE): carries positional information, cached as $k_t^R \in \mathbb{R}^{d_r}$, with attention scoring via standard rotation-aware dot product

Specifically, introduce two new projections $W^{KR} \in \mathbb{R}^{D \times d_r}$ and $W^{QR, (i)} \in \mathbb{R}^{d_c' \times d_r}$ (per-head). $k_t^R$ is shared across all heads:

$$k_t^R = \text{RoPE}_t(W^{KR} h_t), \quad q_t^{R, (i)} = \text{RoPE}_t(W^{QR, (i)} c_t^Q)$$

Full K/Q is a concat of two segments:

$$k_t^{(i)} = [k_t^{C, (i)}; k_t^R], \quad q_t^{(i)} = [q_t^{C, (i)}; q_t^{R, (i)}], \quad k_t^{(i)}, q_t^{(i)} \in \mathbb{R}^{d_\text{head} + d_r}$$

Attention score becomes the sum of two parts:

$$q_t^{(i)\top} k_s^{(i)} = \underbrace{q_t^{C, (i)\top} k_s^{C, (i)}}_\text{latent, absorb} + \underbrace{q_t^{R, (i)\top} k_s^{R}}_\text{RoPE, standard dot}$$

4.4.3　Total cache formula

$$\boxed{\;\text{MLA cache}_\text{per token} = \underbrace{d_c}_\text{latent K/V shared} + \underbrace{d_r}_\text{RoPE K (shared across heads)}\;}$$

DeepSeek-V2: $d_c = 512, d_r = 64$, 576 fp16 elements per token. Compared to LLaMA-3-70B (GQA, $H_\text{kv}=8, d_\text{head}=128$) with $2 \cdot 8 \cdot 128 = 2048$ elements per token, MLA is about 1/3.5 of GQA; vs same-scale MHA ($2 \cdot 64 \cdot 128 = 16384$) about 1/28. DeepSeek-V2 paper reports 93.3% KV reduction vs its internal MHA baseline (numbers differ across model scales and head counts; the 1/28 here is an estimate under another parameter set).

4.5　KV quantization (KIVI / KVQuant / FP8)

KV cache quantization routes compress each cache element from fp16 (2 bytes) to fewer:

§5 PagedAttention (vLLM memory management)

5.1　Problem: fragmentation of KV cache

Traditional attention implementations treat each request's KV cache as a contiguous large tensor $[L_\text{max}, n_\text{layers}, 2, H_\text{kv}, d_\text{head}]$. Issues:

• Must pre-allocate length $L_\text{max}$ (actual usage may be only 10%) → internal fragmentation
• Different requests have different lengths, cache block sizes differ → external fragmentation
• After freeing a request, fragmentation means new requests can't find a large enough block → memory utilization ~70%

5.2　Solution: virtual-memory-style paging

PagedAttention (Kwon et al., SOSP 2023) borrows from OS paging:

1. Partition KV cache into equal-size pages (e.g., 16 tokens per page)
2. Each request maintains a block table: logical block idx → physical block idx
3. The physical page pool is global; pages are allocated on demand and reclaimed on free
4. Attention kernel is rewritten as paged attention: indirect lookup via block table (gather)

Effects:

• Memory utilization from ~70% → ~96%
• Active batch size doubled or quadrupled on the same GPU, throughput rises accordingly
• Supports copy-on-write sharing: beam search, parallel sampling, prefix caching naturally fall out

5.3　Block-table data structure (sketch)

Request A:  logical_blocks = [0, 1, 2, 3]   →   physical = [12, 7, 34, 19]
Request B:  logical_blocks = [0, 1, 2]      →   physical = [12, 7, 5]    ← prefix shared!

Request A and B share the first 32 tokens (2 blocks × 16 tokens); vLLM maintains a ref count for each physical block; when A wants to write new content with ref > 1, it triggers COW (copy + change mapping).

§6 KV Cache Implementation Code

6.1　Naive append + autoregressive decode

import math
import torch
import torch.nn.functional as F
from torch import nn

class NaiveCachedAttention(nn.Module):
    """Single-layer attention with KV cache (MHA / for learning, do not deploy)."""
    def __init__(self, d_model, num_heads):
        super().__init__()
        self.H, self.d = num_heads, d_model // num_heads
        self.W_q = nn.Linear(d_model, d_model, bias=False)
        self.W_k = nn.Linear(d_model, d_model, bias=False)
        self.W_v = nn.Linear(d_model, d_model, bias=False)
        self.W_o = nn.Linear(d_model, d_model, bias=False)

    def _split(self, x):  # [B, L, D] → [B, H, L, d]
        B, L, _ = x.shape
        return x.view(B, L, self.H, self.d).transpose(1, 2)

    def forward(self, x, cache=None):
        """
        x:    [B, L_new, D]  new input (L_new=L_prompt during prefill, L_new=1 during decode)
        cache: dict with 'k','v' of shape [B, H, L_past, d] or None
        Returns: out [B, L_new, D], new_cache
        """
        B, L_new, D = x.shape
        q = self._split(self.W_q(x))                   # [B, H, L_new, d]
        k = self._split(self.W_k(x))                   # [B, H, L_new, d]
        v = self._split(self.W_v(x))                   # [B, H, L_new, d]

        # ── KV cache append ────────────────────────────────────────
        if cache is not None:
            k = torch.cat([cache['k'], k], dim=2)      # [B, H, L_total, d]
            v = torch.cat([cache['v'], v], dim=2)
        new_cache = {'k': k, 'v': v}

        # ── causal attention (when L_new=1 in decode, causal is automatic) ──
        L_total = k.size(2)
        scores = (q @ k.transpose(-2, -1)) / math.sqrt(self.d)
        if L_new > 1:                                  # causal mask only needed during prefill
            causal = torch.tril(torch.ones(L_new, L_total, dtype=torch.bool,
                                          device=x.device), diagonal=L_total - L_new)
            scores = scores.masked_fill(~causal, float('-inf'))
        w = F.softmax(scores, dim=-1)
        out = (w @ v).transpose(1, 2).contiguous().view(B, L_new, D)
        return self.W_o(out), new_cache


# ── Autoregressive generation loop (simplified; real models add sampling/stop/multi-layer) ─
@torch.no_grad()
def generate(model, prompt_ids, max_new_tokens, embed, lm_head):
    cache = None
    out, cache = model(embed(prompt_ids), cache=cache)        # prefill
    next_tok = lm_head(out[:, -1:]).argmax(-1)
    generated = [next_tok]
    for _ in range(max_new_tokens - 1):
        out, cache = model(embed(next_tok), cache=cache)      # decode L_new=1
        next_tok = lm_head(out[:, -1:]).argmax(-1)
        generated.append(next_tok)
    return torch.cat(generated, dim=1)

Note: every torch.cat triggers a new memory allocation + memcpy; in production, swap to pre-allocated buffer + index assignment, or PagedAttention.

6.2　PagedAttention data structure sketch

class PagedKVCache:
    """Simplified page table (no CUDA kernel; data structure + COW demo)."""
    def __init__(self, n_layers, n_kv_heads, head_dim, page_size=16,
                 n_pages=1024, dtype=torch.float16, device='cuda'):
        self.page_size = page_size
        # Global page pool: [n_pages, page_size, n_layers, 2 (K,V), n_kv_heads, head_dim]
        self.pool = torch.empty(n_pages, page_size, n_layers, 2,
                                n_kv_heads, head_dim, dtype=dtype, device=device)
        self.free_list = list(range(n_pages))
        self.ref_count = [0] * n_pages
        self.block_table = {}                          # req_id → list of physical page ids

    def allocate(self, req_id, n_tokens):
        n = (n_tokens + self.page_size - 1) // self.page_size
        assert len(self.free_list) >= n, "OOM"
        physical = [self.free_list.pop() for _ in range(n)]
        for pid in physical: self.ref_count[pid] = 1
        self.block_table[req_id] = physical

    def append_token(self, req_id, pos, layer, k_new, v_new):
        """k_new, v_new: [n_kv_heads, head_dim]"""
        page_idx, slot = pos // self.page_size, pos % self.page_size
        if page_idx >= len(self.block_table[req_id]):
            new_pid = self.free_list.pop()
            self.ref_count[new_pid] = 1
            self.block_table[req_id].append(new_pid)
        pid = self.block_table[req_id][page_idx]
        if self.ref_count[pid] > 1:                    # COW
            new_pid = self.free_list.pop()
            self.pool[new_pid] = self.pool[pid]
            self.ref_count[pid] -= 1; self.ref_count[new_pid] = 1
            self.block_table[req_id][page_idx] = new_pid
            pid = new_pid
        self.pool[pid, slot, layer, 0] = k_new
        self.pool[pid, slot, layer, 1] = v_new

    def free(self, req_id):
        for pid in self.block_table[req_id]:
            self.ref_count[pid] -= 1
            if self.ref_count[pid] == 0: self.free_list.append(pid)
        del self.block_table[req_id]

    def share_prefix(self, src, dst, n_tokens):
        """beam search / parallel sampling: reuse the first n_tokens of pages."""
        n = n_tokens // self.page_size
        prefix = self.block_table[src][:n]
        for pid in prefix: self.ref_count[pid] += 1
        self.block_table[dst] = list(prefix)

Production implementations also need a fused paged-attention kernel (per-block gather + FlashAttention-style) and device-side block-table layout.

6.3　Differences in MQA / GQA / MLA forward

class MQA_GQA_Attention(nn.Module):
    """Unified MHA / MQA / GQA version (num_kv_heads ≤ num_heads)."""
    def __init__(self, d_model, num_heads, num_kv_heads):
        super().__init__()
        assert num_heads % num_kv_heads == 0, "H must be divisible by H_kv"
        self.H, self.H_kv = num_heads, num_kv_heads
        self.d = d_model // num_heads
        self.group = num_heads // num_kv_heads        # # of Q heads each KV head serves
        self.W_q = nn.Linear(d_model, num_heads * self.d, bias=False)
        self.W_k = nn.Linear(d_model, num_kv_heads * self.d, bias=False)   # ← smaller
        self.W_v = nn.Linear(d_model, num_kv_heads * self.d, bias=False)   # ← smaller
        self.W_o = nn.Linear(num_heads * self.d, d_model, bias=False)

    def forward(self, x, cache=None):
        B, L_new, _ = x.shape
        q = self.W_q(x).view(B, L_new, self.H, self.d).transpose(1, 2)
        k = self.W_k(x).view(B, L_new, self.H_kv, self.d).transpose(1, 2)
        v = self.W_v(x).view(B, L_new, self.H_kv, self.d).transpose(1, 2)

        if cache is not None:
            k = torch.cat([cache['k'], k], dim=2)
            v = torch.cat([cache['v'], v], dim=2)
        new_cache = {'k': k, 'v': v}

        # ── Key: broadcast K/V to all Q heads ────────────────────────
        k = k.repeat_interleave(self.group, dim=1)    # [B, H, L_total, d]
        v = v.repeat_interleave(self.group, dim=1)
        # repeat_interleave is explicit broadcast; production uses torch implicit broadcast or fused kernel

        scores = (q @ k.transpose(-2, -1)) / math.sqrt(self.d)
        L_total = k.size(2)
        if L_new > 1:
            causal = torch.tril(torch.ones(L_new, L_total, dtype=torch.bool,
                                          device=x.device), diagonal=L_total - L_new)
            scores = scores.masked_fill(~causal, float('-inf'))
        w = F.softmax(scores, dim=-1)
        out = (w @ v).transpose(1, 2).contiguous().view(B, L_new, -1)
        return self.W_o(out), new_cache


class MLAAttention(nn.Module):
    """MLA with decoupled RoPE (simplified DeepSeek-V2)."""
    def __init__(self, d_model, num_heads, d_head, d_c, d_r):
        super().__init__()
        self.H = num_heads
        self.d_head = d_head          # content head dim
        self.d_c = d_c                # latent dim (shared by K/V)
        self.d_r = d_r                # RoPE head dim (per query head)
        # KV side: compress to latent, then up-project to K_C, V
        self.W_DKV = nn.Linear(d_model, d_c, bias=False)
        self.W_UK = nn.Linear(d_c, num_heads * d_head, bias=False)
        self.W_UV = nn.Linear(d_c, num_heads * d_head, bias=False)
        # RoPE channel: K is shared (across heads)
        self.W_KR = nn.Linear(d_model, d_r, bias=False)
        # Q side: compress to latent (saves training memory), then up-project to Q_C and Q_R
        d_c_q = 4 * d_head            # in the paper d_c_q ≠ d_c
        self.W_DQ = nn.Linear(d_model, d_c_q, bias=False)
        self.W_UQ = nn.Linear(d_c_q, num_heads * (d_head + d_r), bias=False)
        self.W_o = nn.Linear(num_heads * d_head, d_model, bias=False)

    def _rope(self, x, positions, l_axis=-2, base=10000.0):
        """Standard RoPE 2D rotation. L is on l_axis, d_r is the last dim; positions: [L]."""
        assert x.size(-1) % 2 == 0
        D, L = x.size(-1), x.size(l_axis)
        assert positions.shape == (L,)
        i = torch.arange(D // 2, device=x.device, dtype=x.dtype)
        theta = base ** (-2.0 * i / D)                              # [D/2]
        angles = positions.to(x.dtype).unsqueeze(-1) * theta        # [L, D/2]
        cos, sin = angles.cos(), angles.sin()
        ndim = x.dim()
        cos_shape = [1] * ndim
        cos_shape[l_axis % ndim] = L
        cos_shape[-1] = D // 2
        cos_b, sin_b = cos.view(cos_shape), sin.view(cos_shape)
        x2 = x.view(*x.shape[:-1], D // 2, 2)
        x_real, x_imag = x2.unbind(-1)
        rot_real = x_real * cos_b - x_imag * sin_b
        rot_imag = x_real * sin_b + x_imag * cos_b
        return torch.stack([rot_real, rot_imag], dim=-1).flatten(-2)

    def forward(self, x, positions, cache=None):
        B, L_new, _ = x.shape
        # ── KV side: compute latent + RoPE-K ─────────────────────────
        c_kv = self.W_DKV(x)                           # [B, L_new, d_c]
        k_r = self._rope(self.W_KR(x), positions)      # [B, L_new, d_r]  (shared across heads)
        if cache is not None:
            c_kv = torch.cat([cache['c_kv'], c_kv], dim=1)
            k_r  = torch.cat([cache['k_r'],  k_r],  dim=1)
        new_cache = {'c_kv': c_kv, 'k_r': k_r}        # ← only cache latent + RoPE, not full K/V

        # Reconstruct K_C, V (this step can be absorbed into Q/O projection at inference; shown naively here)
        k_c = self.W_UK(c_kv).view(B, -1, self.H, self.d_head).transpose(1,2)  # [B,H,L_tot,d_head]
        v   = self.W_UV(c_kv).view(B, -1, self.H, self.d_head).transpose(1,2)

        # ── Q side: split into content and RoPE segments ─────────────
        c_q = self.W_DQ(x)                             # [B, L_new, d_c_q]
        q_full = self.W_UQ(c_q).view(B, L_new, self.H, self.d_head + self.d_r)
        q_c, q_r_raw = q_full.split([self.d_head, self.d_r], dim=-1)
        q_c = q_c.transpose(1, 2)                      # [B, H, L_new, d_head]
        q_r = self._rope(q_r_raw, positions, l_axis=1).transpose(1, 2)  # [B, H, L_new, d_r]

        # ── attention scores: sum the two parts ──────────────────────
        L_tot = c_kv.size(1)
        scale = math.sqrt(self.d_head + self.d_r)
        scores_c = q_c @ k_c.transpose(-2, -1)         # content part [B, H, L_new, L_tot]
        # k_r is shared across heads, broadcast:
        scores_r = q_r @ k_r.transpose(-2, -1).unsqueeze(1)  # [B, H, L_new, L_tot]
        scores = (scores_c + scores_r) / scale

        if L_new > 1:
            causal = torch.tril(torch.ones(L_new, L_tot, dtype=torch.bool,
                                          device=x.device), diagonal=L_tot - L_new)
            scores = scores.masked_fill(~causal, float('-inf'))
        w = F.softmax(scores, dim=-1)
        out = (w @ v).transpose(1, 2).contiguous().view(B, L_new, -1)
        return self.W_o(out), new_cache

§7 Speculative Decoding Core Mechanism

7.1　Setup

• Target model $p$: the large model to accelerate (e.g., LLaMA-70B), $p(x_{t+1} | x_{\le t})$ is its per-step conditional distribution
• Draft model $q$: a much smaller model (e.g., LLaMA-7B or EAGLE's feature head), $q(x_{t+1} | x_{\le t})$ is its conditional distribution
• Goal: make the final output distribution exactly equal $p$ (exact), not approximate

Each spec step:

1. Use $q$ to autoregressively draft $K$ tokens: $\tilde x_1, \tilde x_2, \dots, \tilde x_K$
2. Feed prefix + $K$ drafts to $p$ at once, in parallel computing $p(x_{t+i} | x_{\le t+i-1})$ for $i=1..K$ (plus $i=K+1$ for bonus-token logits)
3. For each draft position, do rejection sampling: accept with probability $\min(1, p(\tilde x_i) / q(\tilde x_i))$
4. At the first rejected position $j$: resample from the corrected residual distribution $p'$; discard positions $j+1, \dots, K$
5. If all accepted ($j = K+1$), also sample a bonus token for free from the target's last logit set

7.2　Acceptance probability $\alpha$ derivation (must-know)

Derivation: at a position, draft gives $\tilde x \sim q(\cdot)$.

Accept rule: accept with probability $r(\tilde x) = \min(1, p(\tilde x)/q(\tilde x))$.

Probability of accepted token:

$$\Pr[\text{accept} \land X = x] = q(x) \cdot r(x) = q(x) \cdot \min\!\left(1, \frac{p(x)}{q(x)}\right) = \min(q(x), p(x))$$

Rejection probability:

$$\beta = 1 - \alpha = \sum_x q(x) - \sum_x \min(q(x), p(x)) = \sum_x \max(0, q(x) - p(x))$$

Overall acceptance:

$$\boxed{\;\alpha = \sum_x \min(q(x), p(x)) = 1 - \tfrac{1}{2}\|p - q\|_1\;}$$

The last step uses the identity $\sum_x \min(p, q) = 1 - \tfrac{1}{2} \sum_x |p - q|$ (since $\sum_x p = \sum_x q = 1$).

Corollary: the closer $p$ and $q$ are (smaller TV distance), the closer $\alpha$ is to 1.

7.3　Residual distribution $p'$ (how to sample on reject)

After rejection, sample a new token from $p$ with the "accepted mass" removed:

$$p'(x) = \frac{\max(0, p(x) - q(x))}{\sum_x \max(0, p(x) - q(x))} = \frac{\max(0, p(x) - q(x))}{1 - \alpha}$$

Equivalence proof (key, must derive in interviews): the total probability that token $x$ is output in one spec step:

$$\Pr[X = x] = \underbrace{q(x) \min(1, p(x)/q(x))}_\text{accept path} + \underbrace{(1-\alpha) \cdot p'(x)}_\text{reject path}$$

• First term $= \min(p(x), q(x))$
• Second term $= (1-\alpha) \cdot \dfrac{\max(0, p(x) - q(x))}{1-\alpha} = \max(0, p(x) - q(x))$

Sum: $\min(p, q) + \max(0, p - q) = p$. ✅

So the output distribution at each position equals $p$ exactly — this is the mathematical basis of spec decoding's "exact" property.

7.4　Expected speedup: $E[\tau]$ formula

Assume each draft position's acceptance is i.i.d. (slightly correlated in practice; the paper uses this approximation). $K$ drafts + 1 bonus:

• If first $j$ accepted, position $j+1$ rejected ($j < K$): output $j$ accepted + 1 resample = $j+1$ tokens
• If all $K$ accepted: output $K$ + 1 bonus = $K+1$ tokens

Expected token count:

$$E[\tau] = \sum_{j=0}^{K-1} \alpha^j (1-\alpha) (j+1) + \alpha^K (K+1)$$

Simplification (standard geometric-series trick):

$$\boxed{\;E[\tau] = \frac{1 - \alpha^{K+1}}{1 - \alpha}\;}$$

Limit analysis:

• $\alpha \to 1$ (perfect draft): $E[\tau] \to K+1$, speedup $K+1$×
• $\alpha \to 0$ (draft always wrong): $E[\tau] \to 1$, no speedup but no regression either
• In practice, LLaMA-7B drafting LLaMA-70B: $\alpha \approx 0.6-0.7$, $K=4$ gives $E[\tau] \approx 2.7$

$$\text{speedup} = \frac{E[\tau]}{1 + Kc}$$

The numerator is mean accepted token count; the denominator's 1 is the single target verify, and $Kc$ is $K$ draft forwards. If $c$ is too large (draft too big), it eats the gain, so picking a small draft matters.

7.5　Sampling equivalence under temperature / top-p

The rejection sampling equivalence requires only two things: (1) at the target side, replace $p$ with the post-sampler distribution $\tilde p$ for acceptance and residual; (2) the draft proposal distribution $\tilde q$ can be any valid probability distribution. Mathematically it does not require draft and target to use the same sampler — pure greedy draft is also legal, just with $\tilde q$ far from $\tilde p$ and $\alpha$ plummeting. In practice the draft typically uses the same temperature/top-p so $\tilde q$ stays close to $\tilde p$.

7.6　Code: speculative decoding loop

Below is a single-batch demo. Convention: both models expose forward(input_ids, cache); the cache object has a length attribute + truncate(L) method (in production, PagedAttention rollbacks via block-table pointer changes in $O(1)$). Core invariant: at the start of each loop iteration, cache.length == seq.size(1) - 1 (i.e., the cache holds all of seq except the last 1 token).

import torch

@torch.no_grad()
def speculative_decode(target, draft, prompt_ids, max_new_tokens, K=4, temperature=1.0):
    """
    target, draft: callable(input_ids, cache) → (logits [1,L_new,V], new_cache).
    Mathematically exact equivalent to directly sampling from target (Leviathan/Chen 2023).
    """
    seq = prompt_ids.clone()                                   # [1, L_prompt]
    L_prompt = seq.size(1)
    # Prefill the first L-1 tokens; keep the last token as the first verify input.
    _, target_cache = target(seq[:, :-1], cache=None)
    _, draft_cache  = draft(seq[:, :-1],  cache=None)

    while seq.size(1) - L_prompt < max_new_tokens:
        last_tok = seq[:, -1:]                                  # [1, 1], not yet in cache
        draft_chk, target_chk = draft_cache.length, target_cache.length

        # ── 1. Draft: feed last_tok, d_1, ..., d_{K-1}; sample d_1..d_K ──
        cur = last_tok
        draft_tokens, draft_probs = [], []
        for _ in range(K):
            logits, draft_cache = draft(cur, cache=draft_cache)
            probs = torch.softmax(logits[:, -1, :] / temperature, dim=-1)
            tok = torch.multinomial(probs, 1)
            draft_tokens.append(tok); draft_probs.append(probs)
            cur = tok
        draft_tokens = torch.cat(draft_tokens, dim=1)            # [1, K]

        # ── 2. Target: one forward over [last_tok, d_1..d_K], yielding K+1 distributions ──
        target_in = torch.cat([last_tok, draft_tokens], dim=1)   # [1, K+1]
        target_logits, target_cache = target(target_in, cache=target_cache)
        # target_logits[:, i, :] verifies d_{i+1} (i<K) or samples bonus (i=K)

        # ── 3. Rejection sampling position by position ──
        accepted = 0; rejected = False
        for i in range(K):
            p_i = torch.softmax(target_logits[:, i, :] / temperature, dim=-1)
            q_i = draft_probs[i]
            x = draft_tokens[:, i:i+1]
            p_x = p_i.gather(-1, x).squeeze(-1)
            q_x = q_i.gather(-1, x).squeeze(-1)
            r = (p_x / (q_x + 1e-9)).clamp(max=1.0)
            if torch.rand_like(r).item() < r.item():             # accept
                accepted += 1
            else:                                                 # reject
                p_prime = (p_i - q_i).clamp(min=0.0)              # residual distribution p'
                p_prime = p_prime / (p_prime.sum(-1, keepdim=True) + 1e-9)
                new_tok = torch.multinomial(p_prime, 1)
                seq = torch.cat([seq, draft_tokens[:, :accepted], new_tok], dim=1)
                # rollback: keep last_tok + accepted drafts; new_tok not yet in cache
                draft_cache.truncate(draft_chk + 1 + accepted)
                target_cache.truncate(target_chk + 1 + accepted)
                rejected = True
                break

        if not rejected:                                          # all accepted → bonus token
            p_bonus = torch.softmax(target_logits[:, K, :] / temperature, dim=-1)
            bonus = torch.multinomial(p_bonus, 1)
            seq = torch.cat([seq, draft_tokens, bonus], dim=1)
            # draft_cache previously saw only d_1..d_{K-1}; feed d_K to maintain invariant
            _, draft_cache = draft(draft_tokens[:, -1:], cache=draft_cache)
            # target_cache already contains last_tok + d_1..d_K, matches invariant

    return seq[:, :L_prompt + max_new_tokens]                     # trim overshoot

§8 Main Speculative Decoding Variants

8.1　Variant overview

8.2　Medusa: multi-head in parallel

Medusa's core (Cai et al., ICML 2024): add $N$ parallel prediction heads on top of the target's final hidden state; the $k$-th head directly predicts the "$(k+1)$-th future token".

• No need for a separate draft model; only fine-tune these heads (small parameter count)
• The top-$K$ candidates from each of $N$ heads compose a static tree (e.g., top-5 per head, $5^N$ paths total, pruned by typical acceptance)
• Tree attention: flatten all tree nodes into one input for the target's forward; each node's attention mask sees only its ancestors (causal on the tree)

• Verification defaults to typical acceptance (proposed in the paper): accept draft tokens based on the typical-set threshold of the target distribution; this rule does not strictly guarantee exact sampling, but in practice quality is essentially unchanged. For exact, switch to standard rejection sampling (Leviathan/Chen formula).
• Medusa-1 vs Medusa-2 differ in training paradigm: Medusa-1 freezes the backbone and only trains heads; Medusa-2 jointly trains backbone + heads for higher quality; both default to typical acceptance.

8.3　EAGLE family: feature-level autoregression

EAGLE's core insight (Li et al., ICML 2024): the target's previous-layer hidden feature $h_{t-1}$ contains more information than tokens — drafting in feature space is more accurate than in token space.

• Draft model: one transformer layer; inputs $h_{t-1}$ (target's feature) + token embedding; predicts $h_t$ and $x_t$
• Training objective: make draft's $h_t$ close to target's $h_t$ (feature loss) + token prediction loss

EAGLE-2 (Li 2024): replaces the static tree with a dynamic tree — at each step, use draft's probability to expand the most promising paths.

EAGLE-3 (Li et al., 2025):

1. Drop feature regression, predict tokens directly (feature error accumulation is the bottleneck)
2. Use multi-layer feature fusion instead of just the top layer
3. Training-Time Test (TTT): simulate draft-chain errors during training, avoiding train-test gap
4. On Vicuna-7B: 4-5× speedup, ~30% improvement over EAGLE-2

8.4　Lookahead Decoding: Jacobi iteration

Jacobi viewpoint (Fu et al., ICML 2024): autoregressive generation is equivalent to solving the nonlinear system $x_i = f(x_{

Step 0:  x = [<random>, <random>, ..., <random>]
Step 1:  x'_i = f(x_{<i})  ∀ i  in parallel
Step 2:  x = x', another round
... until fixed point

Lookahead Decoding:

• Lookahead branch: maintain a 2D window (lookahead size × window depth), Jacobi-style updates over the whole window each step
• Extract apparently-stable n-grams from the trajectory into a pool
• Verification branch: per forward, simultaneously verify promising n-grams in the pool (multi-path tree attention)
• No draft model; no training; pure inference-time trick

Effect: MT-bench 1.8×, code completion multi-card 4× speedup. But for open-ended generation without repetition (e.g., complex reasoning), the gain is limited.

8.5　Long-context specialists: TriForce / MagicDec

TriForce (Sun et al., 2024, arxiv 2404.11912): three-tier hierarchy.

1. First-tier draft: small LM
2. Second-tier draft: target model + sparse KV cache (keep only heavy-hitter / retrieved portions)
3. Third-tier verify: target model + full KV cache

Speedup core: second-tier draft is fast under sparse cache, then the full-cache target verifies a long batch in one go.

MagicDec (Sadhukhan et al., 2024, arxiv 2408.11049, ICLR 2025): observes that the KV cache is the bottleneck in long context, so the draft uses StreamingLLM-style sparse cache (attention sink + sliding window) and the target uses full cache for verify.

8.6　Self-Speculative Decoding

• Take logits from layer $L_d < L$ in target as draft (early exit)
• Use the full target for verify
• Fully backwards-compatible: no new model to train, no fine-tuning
• Typical speedup 1.5-2× (not as fast as EAGLE, but zero extra training)

§9 Complexity / Resource Accounting

9.1　KV cache memory (summary)

Using LLaMA-2/3-70B architecture ($N_\text{layers}=80, d_\text{head}=128$, MHA $H=64$) as baseline, $L_\text{ctx}=4096$, fp16:

9.2　Speculative decoding expected throughput

$$\text{tokens / sec}_\text{SD} = \frac{\text{tokens / sec}_\text{baseline} \cdot E[\tau]}{1 + Kc}$$

Empirical numbers (A100, LLaMA-2-7B target + 68M draft):

• vanilla SD: 1.6-2.2×
• Medusa-2: 2.5-3.0×
• EAGLE-2: 3.0-3.5×
• EAGLE-3: ~4.0× (short context)

Long context (128K+) regime: vanilla SD drops to 1.1-1.3×; TriForce / MagicDec still hold 2-2.5×.

9.3　Rough budget (70B + 4K context + GQA)

For vanilla SD, the draft model (7B fp16) adds +14 GB; EAGLE's draft head is only ~200 MB.

§10 25 Frequently-Asked Interview Questions

By difficulty L1 (must-know) / L2 (advanced) / L3 (top labs). Expand each entry for answers + common pitfalls.

L1 must-know (asked at every inference / serving interview)

L2 advanced (research-oriented / inference systems roles)

L3 top-lab questions (hardest tier)

§A Appendix: Reference Implementation + Sanity Check

A.1　Components summary

From-scratch reference implementations include:

• NaiveCachedAttention — single-layer MHA + KV cache append
• PagedKVCache — page table + COW sharing sketch
• MQA_GQA_Attention — three-in-one unified version
• MLAAttention — with decoupled RoPE channel
• speculative_decode — exact-equivalent spec loop (with rejection + bonus token)

A.2　Sanity-check expected outputs

[a] naive cache append    prefill (1,16,128) → decode 8 token  ✓
[b] MQA/GQA/MHA shape + cache size consistent                   ✓
[c] MLA cache = d_c + d_r elements                              ✓
[d] spec decode rejection: 100k samples estimate α within 1%    ✓
[e] spec decode output vs direct target sampling: TV < 0.01     ✓
[f] paged cache COW: ref_count + share correct                  ✓

A.3　Main references

• KV / Serving systems

  • Kwon et al., "Efficient Memory Management for Large Language Model Serving with PagedAttention", SOSP 2023.
  • Yu et al., "Orca: A Distributed Serving System for Transformer-Based Generative Models", OSDI 2022.
  • Agrawal et al., "Taming Throughput-Latency Tradeoff in LLM Inference with Sarathi-Serve", OSDI 2024.

• Attention variants

  • Shazeer, "Fast Transformer Decoding: One Write-Head is All You Need", arXiv:1911.02150, 2019 (MQA).
  • Ainslie et al., "GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints", EMNLP 2023.
  • DeepSeek-AI, "DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model", arXiv:2405.04434, May 2024 (MLA).

• KV cache quantization

  • Liu et al., "KIVI: A Tuning-Free Asymmetric 2bit Quantization for KV Cache", ICML 2024.
  • Hooper et al., "KVQuant: Towards 10 Million Context Length LLM Inference with KV Cache Quantization", NeurIPS 2024.

• Speculative decoding

  • Leviathan, Kalman, Matias, "Fast Inference from Transformers via Speculative Decoding", ICML 2023.
  • Chen et al., "Accelerating Large Language Model Decoding with Speculative Sampling", arXiv:2302.01318, 2023 (DeepMind).
  • Cai et al., "Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads", ICML 2024.
  • Li et al., "EAGLE: Speculative Sampling Requires Rethinking Feature Uncertainty", ICML 2024.
  • Li et al., "EAGLE-2: Faster Inference of Language Models with Dynamic Draft Trees", EMNLP 2024 (arXiv:2406.16858).
  • Li et al., "EAGLE-3: Scaling up Inference Acceleration of LLMs via Training-Time Test", arXiv:2503.01840, 2025.
  • Fu et al., "Break the Sequential Dependency of LLM Inference Using Lookahead Decoding", ICML 2024.
  • Miao et al., "SpecInfer: Accelerating Large Language Model Serving with Tree-based Speculative Inference and Verification", ASPLOS 2024.
  • Sun et al., "TriForce: Lossless Acceleration of Long Sequence Generation with Hierarchical Speculative Decoding", arXiv:2404.11912, 2024.
  • Sadhukhan et al., "MagicDec: Breaking the Latency-Throughput Tradeoff for Long Context Generation with Speculative Decoding", arXiv:2408.11049, 2024 (ICLR 2025).
  • Zhang et al., "Draft & Verify: Lossless Large Language Model Acceleration via Self-Speculative Decoding", ACL 2024.

Code and formulas have passed independent reviewer static checks (gpt-5.5 xhigh, cross-model); mathematical equivalence arguments verified.