The signature-dish specialist: vLLM's attention backend (FlashAttention, FlashInfer and the prefill/decode asymmetry)

Contents

Part of the under the engine series. The PagedAttention post explained where the KV lives (in paged blocks). This one explains who reads it and how: the attention kernel. And it connects with FlashAttention v1-v4, which took that kernel apart from the inside; here we look at the level above, how vLLM chooses between several kernels and why it needs more than one.

TL;DR

An LLM forward is, for the most part, standard matrix multiplications that any library does well. The exception that decides performance is attention, and having one good kernel is not enough: you need two, because the two phases of inference are physically opposite problems. Prefill processes the whole prompt: many queries against many keys, dense and compute-bound, the home ground of FlashAttention’s IO-aware tiling. Decode generates one token: a single query against all the accumulated KV, thin and memory-bound, where the only thing that matters is saturating HBM bandwidth while reading the paged KV. That is why vLLM does not have “the attention kernel” but a switchable backend (FLASH_ATTN, FLASHINFER, TRITON_ATTN and others) plus logic that picks according to the GPU: by default FA4 on Blackwell (SM100), FA3 on Hopper (SM90), FA2 elsewhere, with FlashInfer as the alternative that compiles bespoke kernels (JIT) and knows how to do cascade attention for shared prefixes. This post explains why prefill and decode are opposites (with arithmetic intensity), how the backend reads paged KV, how the engine chooses, what FlashInfer brings, the 10 knobs and the trap of pinning a backend blindly. On the generic 4×H100 SXM cluster.

Where you are: the specialist, not the commis

In the kitchen, almost all the work is chopping and sautéing: standard operations any competent commis chef executes, the matrix multiplications of the feed-forward layers and the projections. There is one dish that is not delegated: the signature dish, the one that defines the restaurant. That dish is attention, and it has a quirk: it is cooked in two radically different ways depending on the moment of service.

During prefill, when a new order arrives with its whole prompt, you cook big: a lot of raw material at once, a lot of fire, an intense operation that fills the stoves. During decode, when a table asks for “one more dish”, you cook à la carte: a single dish, but you have to go to the larder and bring back every ingredient that table has accumulated over the whole meal. One is a problem of firepower; the other, of larder speed. The same specialist does not do both well. That is why vLLM has several, and a chef who decides which one steps up depending on the GPU and the phase. That is the attention backend.

Why prefill and decode are opposite problems

This is the central idea, and one sum proves it: arithmetic intensity (FLOPs per byte read). An operation with high intensity is limited by compute; one with low intensity, by memory.

Prefill. We serve $N$ queries (the whole prompt) against $N$ keys. The $QK^\top$ operation and the $\text{softmax}\cdot V$ do on the order of $N^2 d$ FLOPs and read on the order of $N d$ of data. Intensity grows with $N$:

$$I_\text{prefill} \sim \frac{N^2 d}{N d} = N$$

With large $N$ (a prompt of thousands of tokens), intensity is high: compute-bound. This is where FlashAttention’s tiling squeezes the tensor cores and where you get close to the GPU’s peak TFLOPS.

Decode. We serve one query (the new token) against $L$ keys (all the accumulated KV). FLOPs on the order of $L d$; bytes read on the order of $L d s$ (the entire KV has to be read from HBM). Intensity is:

$$I_\text{decode} \sim \frac{L d}{L d s} = \frac{1}{s} \quad (\approx 0.5 \text{ FLOP/byte in FP16})$$

Constant and tiny: memory-bound. The decode kernel is not limited by how much the GPU can compute but by how fast it reads the KV from HBM. It makes no difference that the H100 has 132 idle SMs (see the SM post): the bottleneck is the 3.35 TB/s of bandwidth, and the decode kernel exists so as not to waste a single one of those bytes/s.

performancearithmetic intensity (FLOP/byte)compute roof (peak TFLOPS)decode (I≈0.5): memory-bound, limited by HBM BWprefill (I≈N): compute-bound, limited by computethe ramp = thememory-bound region

The design consequence: a kernel optimised for prefill (dense tiling, maximum tensor-core occupancy) is not the optimum for decode (coalesced reads of the paged KV, minimum latency). Serious servers have different kernels, or one kernel with two paths. In models with MLA (multi-head latent attention), vLLM goes as far as using separate backends for prefill and decode, selectable independently (attention backends, vLLM).

The scheduler’s trick: prefill and decode in the same forward

Here the circle closes with the scheduler post. Because vLLM V1 mixes requests in prefill and in decode in every step, one single forward has to serve both. The backend receives metadata telling it, for each sequence in the batch, how many queries it brings and how much KV it has to read, and applies the right path to each one. That is why the attention backend and the scheduler are coupled: the first has to digest the heterogeneous batch the second builds.

How the backend reads the paged KV

The kernel does not receive a contiguous KV tensor: it receives the block table from the block manager and does a gather over the physical blocks. This imposes a real constraint on the backend: it has to support the paged layout and vLLM’s block_size. Not every kernel in the world does; the ones vLLM integrates (FlashAttention, FlashInfer, Triton) are adapted to read KV in fixed-size blocks scattered across HBM. That is why you cannot plug in any attention kernel from a paper: it has to speak the language of the pigeonhole larder.

The backends and how the engine chooses

vLLM exposes a backend abstraction with several implementations (vLLM deepwiki):

  • FLASH_ATTN — the FlashAttention family. By default the version is picked by architecture: FA4 on SM100 (Blackwell), FA3 on SM90 (Hopper), FA2 elsewhere, configurable with flash_attn_version.
  • FLASHINFER — an attention engine with JIT compilation and specialisable kernels; strong on heterogeneous KV and shared prefixes.
  • TRITON_ATTN — written in Triton, portable and with no dependency on precompiled CUDA binaries (Triton backend deep dive, vLLM, mar-2026).
  • Specific backends for MLA and for non-NVIDIA hardware.

Selection is automatic unless you force it with VLLM_ATTENTION_BACKEND. The heuristic tries FlashAttention first; on Blackwell (SM100) the fallback order for MLA is TRT-LLM Ragged → FlashInfer → others; on other GPUs only FlashAttention is considered for the main path (attention backends, vLLM). The decision depends on: architecture (SM), dtype (FP16/BF16/FP8), head dimension, and whether the load needs a feature only one backend has (cascade attention, certain soft caps, FP8 in KV).

which architecture (SM)?SM100 Blackwell → FA4SM90 Hopper → FA3rest → FA2special feature? (cascade, FP8 KV, MLA) → FlashInfer / specific

What FlashInfer brings: JIT and cascade attention

FlashInfer does not compete with FlashAttention on “being a bit faster”; it attacks a different problem: the heterogeneity of the KV in real serving (FlashInfer, arXiv 2501.01005). Two ideas:

JIT compilation. Instead of a monolithic kernel, FlashInfer generates bespoke kernels for the attention variant, the problem shape and the KV layout you have, injecting functors (query/key/logit transformations, masks). It specialises rather than generalises.

Cascade attention. Here is the jewel for serving with shared prefixes. If $R$ requests share a prefix of $P$ tokens (a common system prompt), naive attention would read that prefix $R$ times. Cascade attention computes it once against the shared prefix and then combines with each request’s own suffix:

$$\text{reads: } \underbrace{R \cdot (P + s_i)}_{\text{naive}} \;\longrightarrow\; \underbrace{P + \textstyle\sum_i s_i}_{\text{cascade}}$$

With $R=50$ requests and a prefix $P=1000$, that is reading 50,000 prefix tokens versus 1,000. It is the natural complement to prefix caching: the block manager shares the prefix’s memory, and cascade attention shares the compute of attending over it.

The arithmetic that matters: when switching backend gives you something

The backend only moves the needle where attention is the bottleneck. In memory-bound decode, a kernel that exploits HBM bandwidth better gives a real improvement; in compute-bound prefill with long sequences, FA3/FA4 getting close to the tensor-core peak gives a real improvement. But if your bottleneck is in another layer, launch overhead, a badly sized scheduler, the cold start, then switching backend does not touch that part. The rule, once again: measure the regime before optimising.

The 10 knobs

#KnobWhat it controlsCost / risk
1VLLM_ATTENTION_BACKENDforce a backendmismatch with hardware/feature
2flash_attn_version (2/3/4)FA versionversion unsupported on your SM
3enable FlashInferJIT + cascadeinitial JIT compilation time
4cascade attentionprefix compute reuseonly helps with a heavily shared prefix
5kv_cache_dtype (FP8)FP8 support in the kernelnot every backend/SM supports it
6block_sizelayout the kernel must readconsistency with PagedAttention
7MLA prefill backendkernel for the dense phaseMLA models only
8MLA decode backendkernel for the thin phaseMLA models only
9soft cap / sliding windowfeatures that restrict backendsfewer kernel options
10head_dim / variantwhich kernels are eligibleexotic models with no support

How it connects to the rest of the stack

With FlashAttention. The FA post explains the kernel from the inside (tiling, online softmax, FA1-4); this is the level above, how vLLM chooses between kernels and why it needs more than one.

With PagedAttention. The backend reads the KV that the block manager places in blocks; it has to speak block-table.

With the scheduler. The scheduler builds mixed prefill+decode batches; the backend has to serve both regimes in a single forward.

With CUDA graphs. Attention kernels are captured in the CUDA graphs; a backend that launches many small kernels benefits more from capture.

With prefix caching. Cascade attention is the compute side of what prefix caching does in memory.

With FP8. Attending over KV in FP8 requires the backend to have the FP8 path; not all of them do on every architecture.

Traps and things that are not what they look like

“FlashInfer is always faster than FlashAttention.” No. FlashInfer wins when its specialisation (cascade, heterogeneous KV, a particular attention variant) applies to your load; in classic dense prefill, FA3/FA4 usually does as well or better. It depends on the regime, there is no universal winner.

“A good attention kernel works for everything.” The underlying error of this post. Prefill and decode are compute-bound and memory-bound respectively; a kernel tuned for one wastes effort on the other. That is why there are separate paths (and separate backends in MLA).

“Decode is compute-bound because the GPU is at 100%.” nvidia-smi at 100% is misleading (see the SM post): decode is memory-bound, the GPU is moving KV, not computing. Optimising decode’s compute is polishing what is not the bottleneck.

“I pin VLLM_ATTENTION_BACKEND and forget about it.” Pinning a backend by hand can leave you on a suboptimal one when you change GPU or version, or force a slow fallback if your hardware does not support what you asked for. Autoselection usually gets it right; pin it only with a measurement that justifies it.

“Cascade attention always helps.” Only with a heavily shared prefix across many concurrent requests. If each request has its own context, there is nothing to share and the overhead of organising the cascade is not amortised.

“The attention backend is the bottleneck, that is why I am slow.” Almost always the bottleneck is higher up (launch, scheduling, memory) or lower down (bandwidth). The backend matters where attention dominates; measure it with nsys/DCGM before changing it.

Conclusion

Of everything an LLM does while generating text, almost all of it is matrix multiplications that any library resolves. Performance is decided in a single kernel, attention, and the surprise is that it is not even one kernel: it is two opposite problems wearing the same name. Prefill wants fire, dense compute over thousands of tokens, and decode wants a fast larder, reading a token’s entire KV with the minimum waste of bandwidth. That is why vLLM did not pick a winning kernel but an abstraction that switches: FlashAttention tuned to each architecture for the general case, FlashInfer compiling bespoke kernels when there is heterogeneity or prefixes to share, Triton for portability. The head chef does not cook the signature dish just one way: he looks at who is ordering and at what point of service, and sends out the right specialist. The lesson for whoever is tuning is the usual one in this series: before changing specialist, make sure the signature dish really is what is holding you back.

See also

References