Inference: KV Cache, Batching, Speculative Decoding

You trained a model. Now someone types a prompt and waits. The model produces one token, then another, then another, and each one takes a full trip through seven billion weights. Why is that slow? Not for the reason you probably think.

The obvious guess is "seven billion multiplications is a lot of multiplications". It is not, for a GPU. An A100 does about 312 trillion half-precision multiply-adds per second; a 7B model needs about 14 billion per token. That should be 20,000 tokens per second. Real systems get around 30 to 150 per second for one user. Something else is the bottleneck.

That something else is memory bandwidth. Every token, the GPU has to read all 14 GB of weights from memory, and reading is what it cannot do fast enough. Once you see that, every trick in this chapter (caching, batching, speculation, quantization in the next chapter) becomes a variation on one theme: do more useful work per byte read.

The problem: one token at a time

A transformer trained with the next-token objective is a function from a sequence to a distribution over the next token. To write a sentence, you call it, sample, append, call it again. This is autoregressive decoding. Nothing about the architecture forces you to recompute everything from scratch each time, but the naive loop does exactly that.

Prefill and decode are different animals

Look at what the GPU is doing in the two phases of a single request.

Prefill is the first forward pass, over the whole prompt. If the prompt is 2,048 tokens, every weight matrix multiplies a 2,048-row block of activations. That is a big, fat matrix multiply. GPUs love those: read the weight once, use it 2,048 times.

Decode is every step after that. There is one new token, so every weight matrix multiplies a single row. Read 14 GB of weights, do one row's worth of arithmetic, throw the weights away, repeat for the next token. The GPU spends almost all of its time waiting for memory.

Arithmetic intensity: a worked example

To make "memory-bound" precise we need one number: how many floating-point operations you do per byte you move. That ratio is called arithmetic intensity. Here it is for a decode step of a 7B model in fp16, batch size 1.

First the bytes. Every parameter is 2 bytes in fp16, so the weights are $7 \times 10^9 \times 2 \approx 14$ GB. A decode step reads all of them once (plus the KV cache, which we will get to). So the traffic is about $1.4 \times 10^{10}$ bytes.

Now the FLOPs. A forward pass costs roughly 2 FLOPs per parameter per token (one multiply, one add, in every matrix-vector product). For one token, that is $2 \times 7 \times 10^9 = 1.4 \times 10^{10}$ FLOPs.

Divide the two and you get the intensity of decoding at batch 1:

One operation per byte. Now compare that with what the hardware can sustain. An A100 has about 312 TFLOP/s of fp16 tensor-core throughput and about 2 TB/s of memory bandwidth. Divide again:

This is the "ridge point": below 156 FLOP/byte the memory system is the bottleneck, above it the arithmetic units are. Decode sits at 1, about 150 times below the ridge. The GPU is doing less than 1% of the arithmetic it could.

So the time per decode step is just the time to stream the weights: $14\ \text{GB} / 2\ \text{TB/s} \approx 7$ ms, or about 140 tokens per second, no matter how fast the multipliers are. Prefill of a 2,048-token prompt, by contrast, needs $2 \times 7 \times 10^9 \times 2048 \approx 28.7$ TFLOP for the same 14 GB of weight traffic: intensity around 2,000 FLOP/byte, well above the ridge, so it is limited by compute and takes about 90 ms at full tensor-core speed.

The roofline picture

The cleanest way to hold all this in your head is the roofline plot of Williams, Waterman and Patterson (2009). Put arithmetic intensity on the x-axis and achieved FLOP/s on the y-axis, both logarithmic. Two lines bound what any kernel can do: a slanted line with slope equal to memory bandwidth, and a flat line at peak compute. Where you sit on the x-axis decides which line caps you.

The KV cache

Here is the first and biggest waste in the naive loop. At decode step $t$ the model has $t$ tokens of context. Attention in every layer needs the keys and values of all $t$ tokens. But the keys and values of tokens $1 \dots t-1$ are exactly the same as they were at step $t-1$: they are a function of those tokens' hidden states, which do not change (there is no bidirectional information flow in a causal model). Recomputing them is pure repetition.

So: compute each token's K and V once, store them, and at every later step just read them back. That store is the KV cache.

What is stored

For every layer, for every KV head, for every past token, two vectors of length $d_{head}$: the key and the value. Not the query. The query is only ever needed by the token that owns it, at the step it is produced, and then never again. Not the hidden states either; once K and V are out, the rest of that layer's work on that token is done.

The memory formula, with real numbers

Count the numbers stored and multiply by the bytes each takes. There are $n_{layers}$ layers, each with $n_{kv}$ heads, each holding $L$ rows of length $d_{head}$ for K and again for V (hence the leading 2), each number costing $b$ bytes. So the cache for one sequence is:

The formula says the cache grows linearly with context length and with the number of KV heads, and that the weights of the model do not appear at all: the cache is a separate memory cost, on top of the weights, that you pay per concurrent sequence.

How the attention code changes

Three things change in the forward pass. The model takes only the new tokens as input (the whole prompt during prefill, a single token during decode). Each layer appends its freshly computed K and V to a per-layer buffer and attends against the full buffer. And the causal mask has to know that the new queries sit at absolute positions $n_{past}, n_{past}+1, \dots$, not at 0. Here is the shape of it, close to what code/lumen/kv_cache.py does.

import math, torch

class KVCache:
    """One pre-allocated (K, V) buffer per layer, filled left to right."""
    def __init__(self, n_layers, batch, n_kv_heads, max_len, d_head,
                 dtype=torch.float16, device="cpu"):
        shape = (n_layers, batch, n_kv_heads, max_len, d_head)
        self.k = torch.zeros(shape, dtype=dtype, device=device)
        self.v = torch.zeros(shape, dtype=dtype, device=device)
        self.pos = 0                                  # tokens stored so far

    def update(self, layer, k_new, v_new):
        # k_new, v_new: (batch, n_kv_heads, T_new, d_head)
        T = k_new.size(2)
        self.k[layer, :, :, self.pos:self.pos + T] = k_new
        self.v[layer, :, :, self.pos:self.pos + T] = v_new
        n = self.pos + T
        return self.k[layer, :, :, :n], self.v[layer, :, :, :n]   # views, no copy

def attention_with_cache(q, k, v, n_past):
    # q: (B, h, T_new, d)      k, v: (B, h, n_past + T_new, d)
    scores = q @ k.transpose(-2, -1) / math.sqrt(q.size(-1))
    T_new, T_tot = q.size(2), k.size(2)
    i = torch.arange(T_new, device=q.device)[:, None] + n_past    # absolute query positions
    j = torch.arange(T_tot, device=q.device)[None, :]             # key positions
    scores = scores.masked_fill(j > i, float("-inf"))              # causal: key must not be in the future
    return torch.softmax(scores, dim=-1) @ v

@torch.no_grad()
def generate(model, prompt_ids, max_new, cache):
    logits = model(prompt_ids, cache)            # prefill: T_new = len(prompt)
    cache.pos += prompt_ids.size(1)
    out = []
    for _ in range(max_new):
        nxt = logits[:, -1].argmax(-1, keepdim=True)   # greedy, for clarity
        out.append(nxt)
        logits = model(nxt, cache)               # decode: T_new = 1
        cache.pos += 1
    return torch.cat(out, dim=1)

Notice that during decode q has a single row, so scores is $1 \times (n_{past}+1)$: the 1×6 purple strip in Figure 3. Notice also that the mask is only there for prefill; with one query at the last position nothing is in the future, so every key is visible. And the cache is pre-allocated to max_len: a design decision that PagedAttention will revisit.

Inside the model, each layer calls cache.update(layer_idx, k, v) and attends over what comes back. With GQA the cached K and V have $n_{kv}$ heads while q has $n_{heads}$; you expand the KV heads with repeat_interleave (or better, reshape the queries into groups) before the matmul, as in code/lumen/attention.py.

The latency win

How much does the cache save? Without it, decode step $t$ runs the full model over all $t$ tokens, costing $2 N t$ FLOPs. Summing over a generation of $L$ tokens gives $\sum_t 2Nt \approx N L^2$: quadratic in the output length. With the cache each step costs $2N$, so the total is $2NL$: linear. For a 512-token answer that is a 256-fold difference in arithmetic.

Step through it below with a toy sequence and see exactly which rows get computed and which get read.

Batching

The cache fixed the redundant arithmetic. It did not fix the roofline problem: each decode step for one user still streams 14 GB of weights to produce one token. What if two users are waiting? The weights are already in flight; doing a second row of arithmetic against them is almost free.

Why batching helps decode (and barely helps prefill)

With a batch of $B$ sequences, a decode step reads the weights once and does $B$ rows of arithmetic. Intensity goes from 1 FLOP/byte to about $B$ FLOP/byte. Time per step barely moves until $B$ approaches the ridge (around 150 on an A100), so throughput in tokens per second grows nearly linearly with $B$ over that whole range. That is the blue dot in Figure 2 sliding up the slope.

Prefill gains nothing from batching, because it was compute-bound already. That asymmetry is why serving systems think about the two phases separately.

There is a catch, and it is the KV cache again. Weight reads are shared across the batch; cache reads are not. Each sequence in the batch brings its own $L \times$ (bytes per token) of cache that must be read every step. Total bytes per step are $W + B \cdot L \cdot c$ where $c$ is cache bytes per token. Once $B L c$ is comparable to $W$, adding sequences no longer amortizes anything. For Llama-2 7B at 4k context that happens around $B \approx 13.5\,\text{GB} / 2.15\,\text{GB} \approx 6$. This is the second reason GQA matters: it moves that crossover out by the KV-head reduction factor.

Static versus continuous batching

The obvious way to batch is to collect $B$ requests, run them together, and return when they are all finished. This is static batching, and it wastes a lot. Requests do not have the same prompt length (so the short ones are padded to the long one) and do not generate the same number of tokens (so a request that finished after 20 tokens keeps its slot, doing nothing, while another grinds out 500). Utilization can easily drop below half.

Continuous batching (Orca, Yu et al. 2022) schedules at the granularity of a single decode step rather than a whole request. After every step the scheduler looks at the batch: finished sequences leave, and waiting requests take their slots immediately, their prefill mixed into the next step. The batch is a rolling set of sequences at different stages, and the GPU never idles waiting for the slowest request.

PagedAttention: virtual memory for the cache

Continuous batching creates a memory-management problem. Sequences arrive and leave constantly, each with a cache that grows one token at a time to a length you do not know in advance. The pre-allocated max_len buffer from our code sketch reserves the worst case for every sequence, and most of that reservation is never used. Kwon et al. (2023) measured that existing systems wasted 60 to 80 percent of cache memory this way, through reservation and fragmentation.

Their fix, PagedAttention (the core of vLLM), copies the oldest trick in operating systems. Split each sequence's cache into fixed-size blocks (say 16 tokens each). Keep a pool of physical blocks in GPU memory. Give every sequence a block table that maps its logical block $i$ to some physical block anywhere in the pool. The attention kernel follows the table. A sequence's cache no longer has to be contiguous, blocks are allocated only when a sequence actually reaches them, and waste drops to under 4 percent.

Two more things fall out of the block indirection for free. Sampling several continuations from one prompt (beam search, "give me 4 answers") can share the prompt's blocks with copy-on-write, like forked processes share pages. And a scheduler can preempt a sequence under memory pressure by evicting its blocks to CPU memory or simply recomputing them later, rather than refusing new requests.

Prefix caching

In a chat deployment, thousands of requests start with the same system prompt. Their prefill computes the same K and V for those tokens, thousands of times. Prefix caching keeps the blocks of common prefixes in the pool after the request finishes and lets new requests with a matching prefix skip prefill for those tokens, starting the prompt processing at the first token that differs. SGLang's RadixAttention (Zheng et al., 2023) organizes cached prefixes in a radix tree so that partial matches are found quickly; vLLM hashes block contents. A 2,000-token system prompt that costs 90 ms to prefill costs nothing the second time.

Chunked prefill

Mixing prefill and decode in one continuous batch creates a nasty latency effect: when a request with a 10,000-token prompt arrives, its prefill occupies the GPU for hundreds of milliseconds, and every other user in the batch sees their next token stall. Chunked prefill (Sarathi, Agrawal et al. 2023) splits long prompts into chunks of a few hundred tokens and processes one chunk per step, alongside the decode tokens of everyone else. Each step then has enough prefill work to be compute-efficient and enough decode work to keep everybody's stream moving, and the tail latency of the decoders stays bounded.

Speculative decoding

Batching helps when there are many users. What about one user who wants their answer faster? The roofline says a decode step's cost is set by the bytes read, and those bytes buy you one token. But they could buy you several: the same weight read, applied to a batch of a few tokens, costs about the same. The trouble is we do not know the next few tokens yet; each depends on the previous one.

Unless we guess. Speculative decoding (Leviathan et al. 2023; Chen et al. 2023) lets a small, fast draft model guess $k$ tokens ahead, then has the big target model check all $k$ guesses in one forward pass, which is nearly free because that pass is memory-bound anyway. Accept the guesses that check out, fix the first one that does not, and repeat. Crucially, the output has exactly the distribution the target model would have produced on its own. It is a speedup, not an approximation.

The acceptance rule

Why does this not change the output distribution? Because acceptance is not "did the draft guess the target's argmax". It is a form of rejection sampling. Let $q(x)$ be the draft's probability of the proposed token $x$ and $p(x)$ the target's. The rule is:

What this does: wherever the draft is over-confident ($q > p$) we throw its samples away some of the time; wherever it is under-confident ($q \le p$) we always keep them, and the leftover mass is redistributed by the correction distribution $p'$, which lives exactly on the tokens the draft under-produced. The two effects cancel and the marginal of the token we emit is precisely $p$.

def speculative_step(target, draft, ids, k, cache_t, cache_d):
    # 1. draft proposes k tokens, one at a time (cheap)
    props, q_probs, x = [], [], ids
    for _ in range(k):
        q = torch.softmax(draft(x[:, -1:], cache_d)[:, -1], -1)   # (1, V)
        t = torch.multinomial(q, 1)
        props.append(t); q_probs.append(q); x = torch.cat([x, t], 1)
    # 2. target scores the last known token + all k proposals in ONE pass
    p_all = torch.softmax(target(x[:, -(k + 1):], cache_t), -1)     # (1, k+1, V)
    # 3. accept left to right, fix the first failure
    out = []
    for i in range(k):
        t = props[i].item()
        p, q = p_all[0, i], q_probs[i][0]
        if torch.rand(()) < min(1.0, (p[t] / q[t]).item()):
            out.append(t)
        else:
            resid = torch.clamp(p - q, min=0.0); resid = resid / resid.sum()
            out.append(torch.multinomial(resid, 1).item())
            break
    else:                                   # everything accepted: bonus token
        out.append(torch.multinomial(p_all[0, k], 1).item())
    return out                              # between 1 and k+1 new tokens

Expected speedup versus acceptance rate

How many tokens does one round produce? Suppose each draft token is accepted independently with probability $\alpha$. The round yields $i$ accepted tokens plus one (the correction or the bonus), and the number accepted before the first failure is a truncated geometric variable. Summing it:

At $\alpha = 0.8$ and $k = 4$ this is $(1 - 0.8^5)/0.2 = 3.36$ tokens per big-model pass instead of 1. The formula also shows the diminishing returns: as $k$ grows the sum converges to $1/(1-\alpha) = 5$, because a rejection somewhere in a long draft throws away everything after it.

The speedup is smaller than that, because drafting is not free. If one draft token costs a fraction $c$ of a target step and the verification pass costs about one target step (memory-bound, so $k+1$ tokens cost roughly what one does), a round costs $kc + 1$ target-steps, and

With $c = 0.1$ the example above gives $3.36 / 1.4 = 2.4$. Leviathan et al. report 2 to 3 times on T5-XXL; Chen et al. report 2 to 2.5 times on a 70B Chinchilla with a 4B draft. Two things kill it: a bad draft (low $\alpha$, common on code or unusual domains) and a large batch (the target pass is no longer memory-bound, so verifying $k+1$ tokens really does cost $k+1$ times more).

Self-speculation and multi-token heads

Keeping a separate draft model around is a hassle: it must share the tokenizer, it occupies memory, and it has to be trained or found. Several methods get the draft from the target itself.

Self-speculative decoding (Zhang et al., 2023, "Draft & Verify") drafts by running the target with a subset of its own layers skipped, then verifies with the full stack. No extra weights, and $\alpha$ is decent because the shallow model shares everything with the deep one.

Medusa (Cai et al., 2024) bolts a few small prediction heads onto the target's final hidden state; head $j$ predicts the token $j$ positions ahead. Their top candidates are assembled into a tree of drafts that the target verifies in one pass with a tree-shaped attention mask. EAGLE (Li et al., 2024) instead trains a tiny one-layer draft that autoregresses over the target's feature vectors rather than its tokens, which turns out to be much more predictable, and reports acceptance lengths around 3 to 4 tokens per round.

Multi-token prediction (Gloeckle et al., 2024; used as MTP in DeepSeek-V3, 2024) goes one step earlier: train the model from the start with auxiliary heads that predict several future tokens. At inference those heads are a built-in draft. DeepSeek-V3 reports an acceptance rate of 85 to 90 percent for its second-token head.

All of these preserve the target distribution exactly when they use the acceptance rule above. Some deployments relax it (accept if the token is in the target's top-$k$, say) for extra speed at the cost of no longer sampling from $p$; know which one you are running.

Attention variants that exist for inference

Everything above treated the model architecture as fixed. But the KV cache is so dominant a cost that several architectural choices are made for inference rather than for modeling quality. Each of these trades a little quality for a lot of cache.

MQA and GQA, recapped

Multi-query attention (Shazeer, 2019) gives every query head the same single K and V head: $n_{kv} = 1$. For Llama-2 7B that would cut the cache from 512 KB to 16 KB per token. It hurts quality a little, and it turned out to be unstable to train at scale. Grouped-query attention (Ainslie et al., 2023) is the compromise: $n_{kv}$ somewhere between 1 and $n_{heads}$, with each KV head shared by a group of $n_{heads}/n_{kv}$ query heads. Llama-3 uses 8 KV heads for 32 (or 64) query heads; Mistral 7B also uses 8. Ainslie et al. showed you can convert a trained MHA model to GQA by mean-pooling the K and V heads of each group and fine-tuning briefly, which is how most GQA models before 2023 were made. The implementation lives in code/lumen/attention.py.

Sliding-window attention

If every token only attends to the previous $w$ tokens (a sliding window; Beltagy et al. 2020 for Longformer, Jiang et al. 2023 for Mistral 7B with $w = 4096$), the cache never needs more than $w$ rows per layer: it becomes a ring buffer of fixed size, and cost per step stops growing with context. Information from further back still arrives, indirectly, because layer $\ell$'s window sees layer $\ell-1$'s outputs, which themselves saw $w$ tokens further back; after $n_{layers}$ layers the receptive field is $n_{layers} \times w$. Many recent models interleave a few full-attention layers with mostly sliding-window ones (Gemma 2 and 3 do this) to keep exact long-range recall on a fraction of the cache budget.

KV-cache compression and quantization

Alternatively keep the architecture and shrink the stored numbers. Quantizing the cache to 8 bits halves it with essentially no loss; 4-bit and even 2-bit schemes work if you are careful about the per-channel outliers in keys (KIVI, Liu et al. 2024, quantizes keys per channel and values per token for exactly this reason). Eviction methods keep only the tokens that matter: H2O (Zhang et al., 2023) observed that a small set of "heavy hitter" tokens receive most attention mass and drops the rest. And latent compression changes the projection: DeepSeek-V2's multi-head latent attention (2024) projects K and V down to a shared low-rank latent vector per token, caches only that, and re-expands it inside the attention kernel, cutting the cache by an order of magnitude while, they report, matching MHA quality. The quantization chapter works through the arithmetic of the first of these.

Latency metrics

"Fast" means three different things to three different people. A chat user cares how long until the first word appears and whether the words then arrive at reading speed. An operator cares how many tokens per second the whole GPU produces. Serving is the art of trading these against each other.

TTFT, TPOT and throughput

Time to first token (TTFT) is the latency of prefill plus queueing: how long from hitting enter until the first token streams back. It is dominated by prompt length and by how busy the scheduler is; chunked prefill and prefix caching are TTFT tools.

Time per output token (TPOT), also called inter-token latency (ITL), is the average gap between consecutive tokens once streaming has started, which is the decode step time. Human reading speed is around 5 to 10 tokens per second, so a TPOT of 50 to 100 ms feels instant; 20 ms is only better for agents and code generation, where nobody is reading in real time. Speculative decoding and quantization are TPOT tools.

End-to-end latency is roughly $\text{TTFT} + \text{TPOT} \times (\text{output tokens} - 1)$. Throughput is total tokens per second across all users on the GPU, and it is what determines cost per token. Batching is the throughput tool. Some teams report goodput: the throughput counting only requests that met their latency target, which is what a serving team is actually paid for.

The throughput–latency trade-off

Look back at the batching interactive. Doubling the batch nearly doubles throughput while leaving step time almost unchanged, right up to the point where the step becomes compute-bound. Past that point, step time grows linearly with batch, so TPOT for every user gets worse and throughput stops improving. The knee of that curve is where a well-tuned server sits: as much batching as the memory line allows, and no more. With a latency budget, you run at the largest batch whose step time is still within it; with a cost budget, you run at the knee and accept the TPOT that comes with it.

The serving stack in one section

What actually runs when a request hits an inference server? Roughly five layers, and you have now met the ideas behind each.

Tokenizer and API. The front door speaks an HTTP API (most engines copy the OpenAI chat format), applies the model's chat template to turn messages into a single token sequence, and tokenizes. It is the same BPE you built in code/lumen/tokenizer.py, run once per request on the way in and incrementally on the way out. Servers usually run it in a separate process so the GPU loop never waits on string handling.

Scheduler. Every step, it decides which sequences run: admitting new requests if the cache pool has room (or preempting old ones if it does not), choosing how much prefill to chunk in alongside the decoders, and grouping sequences for the kernels. This is where continuous batching lives and where all the latency-versus-throughput policy is set.

Engine. The forward pass itself: fused attention kernels that read the block tables (PagedAttention, FlashAttention-style tiling from Dao et al. 2022 for prefill), quantized matmul kernels, and the memory manager for the block pool. Speculative decoding, when enabled, is a change to the engine loop: draft, verify, roll back.

Sampler. Logits become tokens with temperature, top-k and top-p (code/lumen/sampling.py), plus repetition penalties and stop sequences. Structured output is implemented here: given a JSON schema or a grammar, a finite-state machine or pushdown automaton tracks which tokens are legal next (Willard & Louf, 2023, "Efficient Guided Generation"), and the sampler masks every illegal logit to $-\infty$ before sampling. The model literally cannot produce invalid JSON, at zero cost per token once the masks are precomputed.

Streaming. Each sampled token is detokenized and sent immediately, typically as server-sent events, so the client renders text as it appears. Detokenization has to be incremental and careful: a single token can be an incomplete UTF-8 byte sequence, so the streamer holds bytes back until they form a valid character.

What actually works in practice

If you remember one thing from this chapter, make it the roofline. Then, in rough order of how much each buys you for a single GPU serving a chat model:

• A KV cache. Non-negotiable; without it decode is quadratic. Every framework does this for you, but you should have written one once.
• Weight quantization to 8 or 4 bits. Nearly linear decode speedup in the memory-bound regime, and the difference between fitting and not fitting on a consumer card. Next chapter.
• Continuous batching with a paged cache. The difference between a demo and a service: 5 to 20 times the throughput of static batching at the same latency in typical mixed workloads.
• GQA or MLA in the architecture. Decided at training time; the most important inference decision that is not made by the inference team.
• Prefix caching and chunked prefill. Cheap TTFT wins for chat and agent workloads with big shared prompts.
• Speculative decoding. 1.5 to 3 times on TPOT for low-batch, latency-sensitive traffic; near zero benefit at high batch. Free if the model shipped with MTP heads.

The one number to sanity-check any of it: bytes moved per step divided by bandwidth. If a reported tokens-per-second figure is faster than that, something is being batched, quantized or speculated, and you should be able to say which.

Practice

Further reading

• Pope et al. (2022). Efficiently Scaling Transformer Inference. The clearest treatment of the memory-bound decode analysis, batching and partitioning at scale.
• Kwon et al. (2023). Efficient Memory Management for Large Language Model Serving with PagedAttention. The vLLM paper; block tables, copy-on-write, and the fragmentation measurements.
• Yu et al. (2022). Orca: A Distributed Serving System for Transformer-Based Generative Models. Iteration-level (continuous) batching.
• Leviathan, Kalman & Matias (2023). Fast Inference from Transformers via Speculative Decoding. The acceptance rule and the expected-tokens formula.
• Chen et al. (2023). Accelerating Large Language Model Decoding with Speculative Sampling. Independent discovery; results on a 70B model.
• Cai et al. (2024). Medusa and Li et al. (2024). EAGLE. Draft heads on the target model; tree verification.
• Gloeckle et al. (2024). Better & Faster Large Language Models via Multi-token Prediction, and DeepSeek-AI (2024). DeepSeek-V3 Technical Report for MTP and MLA in a production model.
• Ainslie et al. (2023). GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints, and Shazeer (2019). Fast Transformer Decoding: One Write-Head is All You Need.
• Agrawal et al. (2023). SARATHI: Efficient LLM Inference by Piggybacking Decodes with Chunked Prefills. Chunked prefill.
• Zheng et al. (2023). SGLang / RadixAttention. Prefix caching with a radix tree.
• Liu et al. (2024). KIVI: A Tuning-Free Asymmetric 2bit Quantization for KV Cache, and Zhang et al. (2023). H2O: Heavy-Hitter Oracle. Shrinking the cache by quantizing or evicting.
• Willard & Louf (2023). Efficient Guided Generation for Large Language Models. Grammar-constrained sampling.
• Dao et al. (2022). FlashAttention. The tiled, IO-aware attention kernel used for prefill everywhere.