FlashAttention, speculative decoding, quantisation and parallelism

What it is

Four techniques that make LLM inference faster or cheaper, and they attack different bottlenecks, which is the organising fact:

TechniqueAttacksHelpsQuality cost
FlashAttentionAttention memory trafficPrefill, long contextNone: exact
Speculative decodingDecode's memory-bandwidth floorDecode latencyNone: output-identical
QuantisationWeight memory and bandwidthBoth, plus capacitySmall to moderate
Parallelism (TP/PP)Model too big for one GPUCapacity, latencyNone

Two of the four are free: FlashAttention and speculative decoding produce mathematically identical output to the unoptimised version. That is unusual and it is why both were adopted essentially universally within months.

The framing that makes them make sense, from the transformer in tensor shapes and prefill vs decode: prefill is compute-bound, decode is memory-bandwidth-bound. FlashAttention helps the compute-bound phase's memory behaviour; speculative decoding attacks the bandwidth-bound phase directly; quantisation helps both because it shrinks the thing being moved.

What this is confused with: these are not alternatives, they compose. A production stack runs FlashAttention and continuous batching and paged KV cache and probably quantisation, and adds speculative decoding when latency matters more than throughput. Asking "which one should we use" is the wrong question; asking "which bottleneck are we on" is the right one.

The problem it solves

The decode floor is the number that motivates most of this:

8B model, fp16:            16 GB of weights
H100 memory bandwidth:     3.35 TB/s
Minimum time per token:    16 / 3350 = 4.8 ms
Maximum tokens/sec:        ~209 per sequence

You must read every weight to produce one token. That ceiling is independent of batch size, so a single user's generation speed cannot be improved by adding hardware in the usual way. It is a bandwidth problem, and only three things change it: move fewer bytes (quantisation), get more tokens per weight-read (speculative decoding), or split the weights across more memory systems (tensor parallelism).

The prefill memory problem is the other half. The attention score matrix is [B, H, S, S], which at B=4, H=32, S=8192 in fp16 is 17.2 GB for one intermediate tensor. Materialising it is impossible at long context, and it is also unnecessary.

Mechanics

FlashAttention: exact attention without the S x S matrix

Standard attention writes the full score matrix to HBM (high-bandwidth memory), reads it back for softmax, writes the result, reads it again for the value multiply. Four HBM round trips over an S x S tensor.

FlashAttention tiles the computation so the intermediate never leaves SRAM (the small, very fast on-chip memory, roughly 20 MB on an H100 against 80 GB of HBM):

for each block of queries Q_i:                       # ~128 rows
    running_max   = -inf
    running_sum   = 0
    accumulator   = 0
    for each block of keys/values K_j, V_j:
        S_ij = Q_i @ K_j^T                            # stays in SRAM
        # Online softmax: rescale the accumulator as the max updates.
        new_max = max(running_max, rowmax(S_ij))
        correction = exp(running_max - new_max)
        P_ij = exp(S_ij - new_max)
        accumulator = accumulator * correction + P_ij @ V_j
        running_sum = running_sum * correction + rowsum(P_ij)
        running_max = new_max
    write accumulator / running_sum to HBM            # ONE write, [block, d_h]

The online softmax rescaling is the trick. Softmax normally needs the maximum over the whole row before it can start, which forces materialising the row. Keeping a running maximum and rescaling the accumulator when it changes gives the identical result incrementally.

                        HBM traffic         Peak memory for scores
Standard attention      O(S^2)              O(S^2)   = 17.2 GB at S=8192
FlashAttention          O(S^2 d / M)        O(S)     = a few MB
                        (M = SRAM size)

Exact, not approximate. The output is bit-comparable up to floating-point associativity. That is why adoption was immediate and total: there is no trade to evaluate.

FlashAttention-2 improved the work partitioning across warps and cut non-matmul FLOPs; FlashAttention-3 targets Hopper's asynchrony and FP8. Reported speedups over a standard implementation are roughly 2 to 4x for prefill, and the memory saving is what actually enables long context.

Speculative decoding: more tokens per weight-read

The insight: verifying k tokens costs almost the same as generating one, because both read the full weights once. Decode is bandwidth-bound, so the extra compute of a longer sequence is nearly free.

1. A small DRAFT model generates k tokens autoregressively.   (cheap: small weights)
      draft: "the capital of France is Paris and it"          k = 5

2. The TARGET model scores all k+1 positions in ONE forward pass.
      One weight read. Compute is k+1 positions instead of 1, which is free.

3. Accept the longest prefix consistent with the target's distribution.
      Rejection sampling makes the accepted output EXACTLY the target's
      distribution: this is not an approximation.

4. On the first rejection, sample that token from the corrected distribution
      and discard the rest. Repeat.
def speculative_step(target, draft, prefix, k=5):
    # 1. Draft k tokens.
    draft_tokens, draft_probs = [], []
    ctx = prefix
    for _ in range(k):
        p = draft(ctx)
        t = sample(p)
        draft_tokens.append(t); draft_probs.append(p[t])
        ctx = ctx + [t]

    # 2. ONE target forward pass over all k+1 positions.
    target_probs = target(prefix + draft_tokens)     # [k+1, vocab]

    # 3. Accept/reject: preserves the target's distribution exactly.
    accepted = []
    for i, t in enumerate(draft_tokens):
        r = uniform(0, 1)
        if r < min(1.0, target_probs[i][t] / draft_probs[i]):
            accepted.append(t)
        else:
            # Rejected: sample from the residual (p_target - p_draft)+, normalised.
            accepted.append(sample(normalise(relu(target_probs[i] - draft_dist[i]))))
            return accepted                            # discard the rest
    # All k accepted: the target's own next token is free.
    accepted.append(sample(target_probs[k]))
    return accepted

The acceptance rate is everything. If the draft model agrees with the target 70 percent of the time on each token:

Expected accepted tokens per target pass = (1 - a^(k+1)) / (1 - a)
At a = 0.7, k = 5:  (1 - 0.7^6)/(1 - 0.7) = 2.94 tokens per target forward pass
Speedup ≈ 2.94 / (1 + k * cost_draft/cost_target)
        ≈ 2.94 / (1 + 5 x 0.05) = 2.35x

At a = 0.4 the same arithmetic gives about 1.6 accepted tokens and, after draft overhead, roughly break-even. A poorly matched draft model makes things slower, so the draft must be from the same family and trained on similar data.

Variants worth naming:

  • Medusa: extra decoding heads on the target model predicting several positions ahead. No separate draft model, and it needs training the heads.
  • EAGLE: drafts in feature space rather than token space, achieving higher acceptance rates than a comparable small draft model.
  • Lookahead / n-gram / prompt lookup decoding: no model at all, drafting from n-grams in the prompt. Works remarkably well for tasks with heavy input copying (summarisation, code editing, RAG), where the output repeats input spans.

That last one is worth knowing because it is nearly free to deploy and its acceptance rate on copy-heavy tasks can be very high.

Quantisation: fewer bytes per weight

fp16 / bf16:  2 bytes    baseline
fp8:          1 byte     Hopper+ native, minimal quality loss
int8:         1 byte     mature, ~0.5-1% quality loss with good calibration
int4:         0.5 bytes  4x compression, 1-3% loss, method-dependent

The two families, and the distinction that matters:

Weight-only quantisation (GPTQ, AWQ) stores weights at low precision and dequantises to fp16 for the matmul. Since decode is bandwidth-bound, reading half the bytes is roughly half the time even though the arithmetic is unchanged. This is the common case and it helps decode much more than prefill.

Weight and activation quantisation (SmoothQuant, fp8) quantises both, so the matmul itself runs in low precision on hardware that supports it. This helps compute-bound prefill too.

# AWQ: Activation-aware Weight Quantisation.
# The insight: ~1% of weight channels handle outlier activations and carry
# disproportionate importance. Scale those channels up before quantising so
# they retain precision, and scale the activations down to compensate.
#
# GPTQ: quantise column by column, and after each column update the REMAINING
# weights to compensate for the error just introduced (approximate second-order).

The KV cache can be quantised too, and this is under-used. KV cache is often the binding constraint on concurrency:

8B GQA model, 8k context:
  fp16 KV:   1.07 GB per request  ->  54 concurrent on an 80 GB GPU
  int8 KV:   0.54 GB per request  ->  ~100 concurrent

Roughly double the concurrency for a small quality cost, and quality loss from KV quantisation is generally smaller than from weight quantisation because the cache is transient.

The measured picture on an 8B model:

                    memory   decode speed   MMLU     concurrency (80GB, 8k ctx)
fp16                16.0 GB     209 tok/s   68.4        54
fp8                  8.0 GB     ~390 tok/s  68.1        61
int8 (SmoothQuant)   8.0 GB     ~380 tok/s  67.9        61
int4 (AWQ)           4.5 GB     ~610 tok/s  66.8        65
int4 + int8 KV       4.5 GB     ~610 tok/s  66.5       ~118

int4 weights plus int8 KV roughly doubles concurrency and triples decode speed for about 2 points of MMLU. Whether that is acceptable is a product question, and it should be measured on your own task rather than on MMLU, because quantisation damage is uneven: it hits long-tail knowledge and multi-step reasoning harder than it hits common tasks.

Parallelism: when the model does not fit

Tensor parallelism (TP) splits each matrix across GPUs. Every layer requires an all-reduce, so it needs fast interconnect (NVLink) and is normally kept within one node.

W [4096, 14336] on 4 GPUs -> each holds [4096, 3584]
Each GPU computes its shard, then all-reduce to combine.
Communication: 2 all-reduces per layer, every token.

TP reduces per-GPU memory and per-token latency (each GPU reads a quarter of the weights), and adds communication. On NVLink at 900 GB/s the overhead is acceptable; over PCIe or Ethernet it dominates and TP becomes counterproductive.

Pipeline parallelism (PP) splits by layer: GPU 0 holds layers 0 to 7, GPU 1 holds 8 to 15, and so on. Communication is one activation tensor per boundary, which is tiny, so PP works across nodes. The cost is the pipeline bubble: with p stages and m microbatches, utilisation is m / (m + p - 1), so few microbatches means idle GPUs.

Expert parallelism (EP) for mixture-of-experts models: different experts on different GPUs, routed per token.

The standard composition:

Within a node (8 GPUs, NVLink):   tensor parallel
Across nodes (Ethernet/IB):        pipeline parallel
Very large models:                 TP x PP x EP

TP for latency and per-GPU memory, PP for fitting across nodes, and the reason is entirely about interconnect bandwidth per unit of communication.

A worked example: 3.1 seconds to 0.6, and what each step bought

A code-completion product. 34B model, fp16, served on 2x A100 80GB with tensor parallelism. Target: p95 under 800 ms for a ~120-token completion.

Baseline:

model:                     34B fp16 = 68 GB, TP=2
p50 latency:               2,140 ms
p95 latency:               3,120 ms
throughput:                 41 req/s
concurrent requests:        18 (KV cache bound)
GPU cost:                  $6,200/mo per replica, 6 replicas

Step 1: FlashAttention (they were on a naive attention implementation).

p50: 2,140 -> 1,890 ms      (-12%)
prefill time: 410 -> 180 ms  (-56%)
max context: 8k -> 32k possible without OOM

Prefill more than halved; decode was untouched, which is correct, because decode's attention is over a single query position and was never the bottleneck. The 12 percent end-to-end improvement understates it: prefill was only 19 percent of total latency for this workload.

Step 2: int8 weight quantisation (SmoothQuant).

model: 68 GB -> 34 GB
p50: 1,890 -> 1,120 ms      (-41%)
concurrent: 18 -> 44         (freed KV cache space)
throughput: 41 -> 88 req/s
code benchmark (HumanEval): 71.3% -> 70.8%   (-0.5 points)

Decode roughly halved because half the bytes are read per token. Concurrency more than doubled because 34 GB of freed memory became KV cache. This was the largest single win and it cost half a point on the benchmark.

Step 3: int8 KV cache.

concurrent: 44 -> 79
throughput: 88 -> 121 req/s
HumanEval: 70.8% -> 70.6%

Throughput only, no latency change, which is expected: KV quantisation buys capacity.

Step 4: speculative decoding with a 1.3B draft from the same family.

acceptance rate (code):     0.81      <- high, because code is predictable
mean accepted per pass:     3.4 tokens
p50: 1,120 -> 480 ms        (-57%)
p95: 1,640 -> 620 ms
throughput: 121 -> 104 req/s          <- WENT DOWN

Latency more than halved and throughput fell by 14 percent. That is the speculative decoding trade stated exactly: the draft model consumes GPU time and memory that would otherwise serve other requests. Under low load you get the latency win nearly free; under saturation you are spending throughput to buy it.

The 0.81 acceptance rate is unusually high and it is why this worked so well here: code is highly predictable, with long runs of boilerplate, closing brackets and repeated identifiers that a small model gets right. On open-ended prose the same setup would land nearer 0.6 and the speedup would be closer to 1.8x.

Step 5: an adaptive policy, since the throughput cost is only paid when it matters.

# Enable speculation only when the batch is small enough that spare
# capacity exists. Under load, drop it and serve throughput.
use_speculation = current_batch_size < SPEC_BATCH_THRESHOLD   # 24

Final:

                      baseline    final     change
p50 latency            2,140ms     490ms    -77%
p95 latency            3,120ms     640ms    -79%   (target was 800ms)
throughput              41 req/s   118 req/s +188%
concurrent requests      18         79
HumanEval               71.3%      70.6%    -0.7 points
replicas needed           6          2
GPU cost              $37,200/mo  $12,400/mo  -67%

Three quarters off latency, nearly triple the throughput, two thirds off cost, for 0.7 points of benchmark.

The ordering mattered and is the transferable part. They applied the free ones first (FlashAttention, then quantisation which was nearly free), measured after each, and only then added the one with a real trade. Had they started with speculative decoding they would have measured a throughput regression against a baseline that was bandwidth-bound for a different reason, and might have concluded it did not work.

Production evidence

FlashAttention (Dao et al., 2022) and FlashAttention-2 are in vLLM, TensorRT-LLM, PyTorch's scaled_dot_product_attention, and essentially every serving stack. Adoption was universal because it is exact: there is no quality trade to evaluate, only an implementation to adopt.

Speculative decoding (Leviathan et al., Google, and Chen et al., DeepMind, both 2023) independently derived the same rejection-sampling scheme. The distribution-preservation proof is what made it deployable: you are not trading quality for speed. It is in vLLM, TensorRT-LLM and llama.cpp.

Medusa (Cai et al., 2024) and EAGLE (Li et al., 2024) are the higher-acceptance successors; EAGLE's feature-space drafting reports notably better acceptance than token-space drafts of comparable cost.

AWQ (Lin et al., MIT) and GPTQ (Frantar et al.) are the standard weight-only 4-bit methods, both widely deployed. SmoothQuant (Xiao et al.) handles activation outliers to make W8A8 viable. fp8 is native on Hopper and Blackwell, and NVIDIA's published results show near-lossless quality, which is why fp8 is becoming the default rather than int8 on new hardware.

Megatron-LM (NVIDIA) established the tensor-parallel decomposition used almost everywhere, and DeepSpeed (Microsoft) the ZeRO family for training memory. The TP-within-node, PP-across-node composition is standard practice documented by both.

vLLM's PagedAttention is the other half of the serving story: without paged KV cache, naive allocation reserves maximum context per request and wastes most of the cache to internal fragmentation. Their reported 2 to 4x throughput improvement is largely recovered fragmentation, and it composes with everything here.

The debate

What order should you apply these? Free first. FlashAttention and paged KV cache have no quality cost and should be considered table stakes rather than optimisations. Then quantisation, measuring on your own evaluation set. Then speculative decoding, which has a real throughput trade. Applying them in the other order produces measurements you cannot interpret, because each changes which bottleneck you are on.

How aggressive should quantisation be? fp8 or int8 is close to free on modern hardware, typically under a point on most benchmarks, and I would treat it as the default. int4 is a genuine decision: 2 to 3 points on aggregate benchmarks and, more importantly, uneven damage. It disproportionately affects long-tail factual knowledge and multi-step reasoning while leaving common tasks nearly untouched, so an aggregate benchmark understates the risk for a reasoning-heavy product. Measure on your task, with your hardest examples, not on MMLU.

Is speculative decoding worth it? It buys latency and spends throughput, so the question is which you are short of. For an interactive product under moderate load, yes, clearly. For a batch pipeline at saturation, no: you are paying throughput for a latency nobody experiences. The adaptive policy in the worked example is the right general answer: speculate when the batch is small, stop when it is not.

When is tensor parallelism wrong? Across nodes. TP all-reduces twice per layer per token, and on Ethernet or even InfiniBand the communication dominates. The rule is TP within a node over NVLink, PP across nodes. A team running TP=8 across two 4-GPU nodes over PCIe will measure worse performance than TP=4 on one node, which is a surprising and entirely predictable result.

Do these compose without interference? Mostly, with two caveats. Speculative decoding and continuous batching interact: the draft model competes for the same GPU, and a scheduler unaware of speculation will make poor admission decisions. And quantisation plus speculation needs the draft and target quantised compatibly, or acceptance rates drop because the two models' distributions diverge more than they would at full precision.

Follow-up Q&A

"What does FlashAttention actually do?"

It computes exact attention without materialising the S x S score matrix in HBM. Tiling the computation and keeping intermediates in SRAM, with an online softmax that maintains a running maximum and rescales the accumulator as it updates, gives the identical result with O(S) memory instead of O(S^2). It is not an approximation, which is why adoption was immediate: there is no quality trade to weigh, only an implementation to adopt. It helps prefill and long context, and barely touches decode, where attention is over a single query position.

"How does speculative decoding preserve the output distribution?"

Rejection sampling. Accept the draft's token with probability min(1, p_target/p_draft), and on rejection sample from the normalised positive part of p_target - p_draft. The resulting distribution is provably exactly the target's. Combined with the fact that verifying k tokens costs one weight-read, the same as generating one, you get several tokens per pass for free in bandwidth terms. The speedup is governed by the acceptance rate, and a badly matched draft can make things slower.

"When does speculative decoding hurt?"

Two cases. Low acceptance rate: at 40 percent agreement the accepted tokens per pass drop to about 1.6, and after draft overhead you are near break-even or worse. And saturation: the draft model consumes GPU time and memory that would otherwise serve other requests, so at high load you lose throughput to buy latency. In one measurement throughput fell 14 percent while p50 more than halved. The right policy is adaptive: speculate when the batch is small.

"How much quality does quantisation cost?"

fp8 and int8 are typically under a point on aggregate benchmarks and I would call them close to free on modern hardware. int4 is 2 to 3 points and, more importantly, the damage is uneven: long-tail factual recall and multi-step reasoning degrade more than common tasks, so an aggregate number understates the risk for a reasoning-heavy product. Quantise the KV cache too, which is often overlooked and roughly doubles concurrency for less quality cost than weight quantisation, because the cache is transient.

"Tensor parallel or pipeline parallel?"

TP splits each matrix across GPUs, so it reduces per-GPU memory and per-token latency and requires an all-reduce twice per layer. That needs NVLink-class bandwidth, so TP stays within a node. PP splits by layer, communicating one activation tensor per boundary, which is small enough to cross nodes, at the cost of a pipeline bubble of (p-1)/(m+p-1). The standard composition is TP within a node and PP across nodes, and running TP across nodes over PCIe is a common and measurable mistake.

"Your decode is slow. What is the ordered list?"

Confirm it is bandwidth-bound by dividing weight bytes by memory bandwidth and comparing with observed per-token time. Then: quantise weights, which directly reduces bytes read per token, roughly linearly. Quantise the KV cache to free memory for a larger batch. Check that continuous batching and paged KV are on, because a static batch wastes both capacity and bandwidth. Add tensor parallelism if you have NVLink, which splits the weight read across GPUs. Then speculative decoding if latency matters more than throughput. And check the arithmetic first, because if you are not near the bandwidth floor the bottleneck is elsewhere.

What is RadixAttention, and how is it different from ordinary prefix caching? Both exploit the same fact: if two requests share a prompt prefix, the KV cache for that prefix is identical and should be computed once. Ordinary prefix caching handles the simple version, typically a fixed system prompt, by keying a cache on the prefix and reusing its KV blocks. RadixAttention, introduced with SGLang, generalises this by holding the entire KV cache in a radix tree keyed on token sequences, so any shared prefix between any two requests is matched automatically, at whatever depth it diverges, without anyone declaring what the shared part is. Eviction is LRU over the tree, and because a parent's blocks are shared by all its children, evicting is leaf-first.

The reason it matters more than it sounds is the workloads it fits, which are exactly the agentic ones. A multi-turn conversation shares every previous turn with itself. A tree search or self-consistency sample shares the trunk across all branches. A batch of few-shot requests shares the examples. In all three, the shared prefix is dynamic and not known in advance, which is precisely what a static prefix cache cannot handle, and the reported gains on those workloads are multiples rather than percentages. The connection to make in an interview: this is a caching problem, the tree is the index, and the interesting question is the eviction policy, which is the same conversation as any other cache and can be reasoned about with the same tools.

Common misconceptions

"FlashAttention is an approximation." It is exact, up to floating-point associativity. That is precisely why it was adopted universally and immediately.

"Speculative decoding trades quality for speed." The rejection-sampling scheme preserves the target's distribution exactly. What it trades is throughput, because the draft consumes resources.

"Quantisation always degrades quality proportionally." The damage is uneven: it hits long-tail knowledge and multi-step reasoning more than common tasks, so aggregate benchmarks understate the risk for some products and overstate it for others. Measure on your own hardest examples.

"More GPUs means faster generation." Only with tensor parallelism, and only within a fast interconnect. Adding replicas increases throughput and does nothing for single-stream latency, which is bounded by weight bytes over bandwidth.

"These are alternatives." They compose, and they attack different bottlenecks. The real question is always which bottleneck you are on, and the arithmetic answers it in a minute.

Interview delivery note

Say this verbatim: "Decode is memory-bandwidth-bound, so the only three levers are: move fewer bytes, which is quantisation; get more tokens per weight-read, which is speculative decoding; or split the weights across more memory systems, which is tensor parallelism. FlashAttention is not on that list, because it fixes prefill and long-context memory, not decode." Categorising the techniques by bottleneck rather than listing them is what demonstrates understanding.

The senior-versus-staff separator is knowing that speculative decoding costs throughput. A senior engineer explains draft-and-verify and the acceptance rate correctly. A staff engineer adds that the draft model consumes GPU resources that would otherwise serve other requests, so under saturation you are trading throughput for latency, and proposes an adaptive policy: speculate when the batch is small, stop when it is not. In the worked example throughput fell 14 percent while latency halved, and both numbers are the point.

The second signal is applying them in the right order and saying why: free ones first, so that each measurement is interpretable. Each technique changes which bottleneck you are on, so measuring speculative decoding against a baseline that has not been quantised tells you very little.

Further reading

  • Dao et al., "FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness" (2022), and FlashAttention-2.
  • Leviathan et al., "Fast Inference from Transformers via Speculative Decoding" (2023), for the rejection-sampling proof of distribution preservation.
  • Lin et al., "AWQ: Activation-aware Weight Quantization" and Frantar et al., "GPTQ," for the two standard 4-bit weight-only methods.
  • Shoeybi et al., "Megatron-LM" (2019), for the tensor-parallel decomposition used throughout the field.