Pidoku

Attention Backends

Advanced 45 min Difficulty 4/5 Lesson 03 of 05

Prerequisites Building the Batch

The idea in one minute#

Attention is the one operation in the model that must understand vLLM’s unusual memory layout: a flat batch of tokens from many requests, each request’s context scattered across blocks anywhere in a shared cache. Ordinary attention code cannot do that, so vLLM needs kernels that take a block table as an input. It does not have one such kernel. It has a registry of attention backends — FlashAttention, FlashInfer, a Triton implementation and more than thirty others — and at startup picks the best one that supports your GPU generation, your model’s attention variant and your cache data type. Which backend you got explains a good deal of your performance, and it is one line in the startup log.

A picture#

flowchart LR
  M[":pytorch: <b>Model layer</b><br/><small>query, key, value for this step</small>"] --> AL[":i-layers: <b>Attention layer</b><br/><small>backend-independent</small>"]
  AL --> SEL{"backend chosen<br/>at startup"}
  SEL --> FA[":nvidia: <b>FLASH_ATTN</b><br/><small>FlashAttention with block tables</small>"]
  SEL --> FI[":nvidia: <b>FLASHINFER</b><br/><small>serving-oriented kernels</small>"]
  SEL --> TR[":i-code: <b>TRITON_ATTN</b><br/><small>portable, written in Triton</small>"]
  SEL --> MLA[":i-brain: <b>MLA family</b><br/><small>latent attention models</small>"]
  FA --> KV[(":i-database: <b>KV cache tensor</b><br/><small>shape chosen by the backend</small>")]
  FI --> KV
  TR --> KV
  MLA --> KV
  META[":i-list-checks: <b>Attention metadata</b><br/><small>block_table, slot_mapping,<br/>query_start_loc, seq_lens</small>"] --> SEL
  class M,AL compute
  class SEL queue
  class FA,FI,TR,MLA compute
  class KV memory
  class META neutral

How it really works#

What “paged attention” requires of a kernel#

Standard attention takes three dense tensors — queries, keys, values — of matching shapes. A kernel usable by vLLM has a different contract:

inputs    queries for the T tokens of this step            (flat, all requests)
          the new keys and values for those T tokens
          the whole KV cache tensor for this layer
          block_table[N, max_blocks]                        where each request's context lives
          slot_mapping[T]                                   where to store the new K and V
          query_start_loc[N+1], seq_lens[N]                 request boundaries

does      1. write the T new keys and values into the cache at slot_mapping
          2. for each request, attend from its queries over its first seq_len tokens,
             fetching them block by block through its row of block_table

The mathematics is unchanged; only the addressing differs. The program at the end of this lesson shows that the result is bit-for-bit identical to attention over a contiguous array.

note

vLLM’s documentation still contains a long description of the original hand-written CUDA kernel from the 2023 paper. It opens with a warning that it “no longer describes the code used in vLLM today.” Read it as history. The idea survived; that kernel is no longer the main path.

The registry#

AttentionBackendEnum in vllm/v1/attention/backends/registry.py lists every backend. The ones you are most likely to meet:

BackendWhat it isTypical use
FLASH_ATTNFlashAttention, in a build maintained by vLLM that accepts block tablesDefault on most NVIDIA GPUs
FLASHINFERKernels from the FlashInfer library, written for LLM servingDefault on the newest NVIDIA data-centre generation
TRITON_ATTNAn implementation in the Triton GPU languageFallback on NVIDIA; the portable option
ROCM_ATTN, ROCM_AITER_FAAMD implementationsAMD GPUs
CPU_ATTNA CPU implementationCPU-only installs
FLEX_ATTENTIONBuilt on PyTorch’s FlexAttentionUnusual attention patterns
TRITON_FLASH_ATTN, TRITON_FLASHINFERComposite backends: two kernels sharing one cacheModels whose image tokens need non-causal attention
FLASHMLA, FLASHINFER_MLA, CUTLASS_MLA, TRITON_MLA, …Kernels for multi-head latent attentionDeepSeek-style models
MAMBA1, MAMBA2, LINEAR, …Not attention at all: state-space and linear layers that use the same block-table plumbingHybrid models
NO_ATTENTIONA stubModels with no KV cache

A second, smaller enum covers the state-space family. That these share the registry is deliberate: the scheduler and KV cache manager only know about layers that keep per-request state in blocks, and a Mamba layer is one.

How one is chosen#

For each attention layer, the platform produces an ordered list of candidates and the first that accepts the configuration wins. For ordinary attention on NVIDIA (_get_backend_priorities in vllm/platforms/cuda.py):

Python
if device_capability.major == 10 and not use_non_causal:
    return [
        *([AttentionBackendEnum.TRITON_FLASHINFER] if use_mm_prefix else []),
        AttentionBackendEnum.FLASHINFER,
        AttentionBackendEnum.FLASH_ATTN,
        AttentionBackendEnum.TRITON_ATTN,
        AttentionBackendEnum.FLEX_ATTENTION,
        AttentionBackendEnum.TURBOQUANT,
    ]
else:
    return [
        *([AttentionBackendEnum.TRITON_FLASH_ATTN]
          if device_capability.major == 9 and use_mm_prefix else []),
        AttentionBackendEnum.FLASH_ATTN,
        AttentionBackendEnum.FLASHINFER,
        AttentionBackendEnum.TRITON_ATTN,
        AttentionBackendEnum.FLEX_ATTENTION,
        AttentionBackendEnum.TURBOQUANT,
    ]

device_capability.major is the GPU’s CUDA compute-capability generation: 8 for Ampere (A100, A10G), 9 for Hopper (H100, H200), 10 for the Blackwell data-centre parts. So:

GPU generationFirst choiceThen
Blackwell data centre (10.x)FlashInferFlashAttention, Triton
Hopper, Ampere and olderFlashAttentionFlashInfer, Triton

“Accepts the configuration” is decided by each backend’s validate_configuration(). A backend can decline because of the head size, the model data type, the KV cache data type (several do not support an FP8 cache), the block size, or a missing optional dependency. The list for latent-attention models is longer and more conditional: it depends on the number of heads and the cache data type, with comments citing the benchmark issue that justified each ordering.

Forcing a choice#

Shell
vllm serve <model> --attention-backend FLASH_ATTN
# or, equivalently
vllm serve <model> --attention-config.backend FLASH_ATTN

If the backend you force cannot support the configuration, startup fails with the reason. Forcing is useful for two things: bisecting a correctness or performance problem (“is it the kernel?”), and working around a backend bug on new hardware.

What a backend must provide#

Each backend is a class implementing a small interface (vllm/v1/attention/backend.py):

PiecePurpose
validate_configuration()Say whether this model, GPU and cache data type are supported
get_kv_cache_shape()The shape of this layer’s cache tensor. Different kernels want different layouts, so the backend, not the engine, decides.
A metadata builderEach step, convert the generic index tensors into whatever this kernel expects
The implementationThe forward function: store new keys and values, compute attention
CUDA graph support levelDeclares which batch shapes can be recorded (below)

The metadata builder runs once per step, per group of layers, on the CPU. For some backends it is cheap; for others it involves planning work proportional to the batch. It is part of the per-step Python cost that async scheduling hides behind the GPU.

Backends differ in what can be recorded#

CUDA graphs record a fixed sequence of GPU operations for replay. Whether attention can be recorded depends on whether the kernel’s control flow depends on the batch’s contents. Each backend declares a level:

Python
class AttentionCGSupport(Enum):
    ALWAYS = 3
    """Cudagraph always supported; supports mixed-prefill-decode"""
    UNIFORM_BATCH = 2
    """Cudagraph supported for batches the only contain query lengths that are
    the same, this can be used for spec-decode"""
    UNIFORM_SINGLE_TOKEN_DECODE = 1
    """Cudagraph supported for batches the only contain query_len==1 decodes"""
    NEVER = 0
    """NO cudagraph support"""

This is why the choice of backend affects more than attention’s own speed. A backend at level ALWAYS lets the whole forward pass be replayed as one graph for any batch; one at UNIFORM_SINGLE_TOKEN_DECODE allows that only for pure decode steps, and mixed steps fall back to a slower mode. vLLM downgrades the configured CUDA graph mode automatically to whatever the backend supports.

Two kernels inside one attention call#

Prefill and decode stress a kernel differently.

  • Prefill has long queries. Attention is a large matrix product; the algorithmic work of FlashAttention — tiling so that nothing bigger than a GPU cache is materialised — is what matters (FlashAttention).
  • Decode has one query token per request and thousands of context tokens. The cost is fetching those keys and values from scattered blocks: a gather, limited by memory bandwidth.

Several backends therefore run different code for the two cases and split a mixed batch internally. This is also the reason the model runner may reorder the batch so that decode requests come first: a backend can then treat the first part and the rest as two homogeneous groups.

Cascade attention#

When many running requests share a long prefix — a common system prompt — their attention over that prefix is partly redundant work. The scheduler computes, every step, how many leading blocks are common to all running requests:

Python
# Get the longest common prefix among all requests in the running queue.
# This can be potentially used for cascade attention.
num_common_prefix_blocks = self.kv_cache_manager.get_num_common_prefix_blocks(any_request_id)

and FlashInfer can use it to process the shared part once and the per-request parts separately. It applies only when every request in the batch shares the prefix, and it is not compatible with full CUDA graphs, so batches using it are dispatched to the piecewise mode.

Latent attention is a different cache, not just a different kernel#

Multi-head latent attention stores one compressed vector per token instead of a key and a value per head. Its backends differ from the ordinary ones in two ways that reach beyond the kernel: the cache page is much smaller, which changes the sizing arithmetic, and prefill and decode use different formulations of the same mathematics. Some of these models default to an FP8 cache. See MQA, GQA and MLA for the technique.

Finding out what you have#

The startup log names the backend selected for the model. To see every candidate and why each was rejected, start with VLLM_LOGGING_LEVEL=DEBUG. The project’s documentation also carries an auto-generated table of every backend against every feature (data types, head sizes, cache types, graph support), built directly from each backend’s validate_configuration().

Code#

Attention through a block table gives exactly the same answer as attention over a contiguous array. This is the whole correctness argument for paging the cache.

Go
package main

import (
	"fmt"
	"math"
)

const (
	blockSize = 4
	dim       = 3 // size of one key or value vector
	numBlocks = 8
)

type vec [dim]float64

func dot(a, b vec) float64 {
	s := 0.0
	for i := range a {
		s += a[i] * b[i]
	}
	return s
}

// attend computes softmax(q.K) . V over the keys and values it is given.
func attend(q vec, keys, values []vec) vec {
	scores := make([]float64, len(keys))
	maxScore := math.Inf(-1)
	for i, k := range keys {
		scores[i] = dot(q, k) / math.Sqrt(dim)
		maxScore = math.Max(maxScore, scores[i])
	}
	sum := 0.0
	for i := range scores {
		scores[i] = math.Exp(scores[i] - maxScore)
		sum += scores[i]
	}
	var out vec
	for i, v := range values {
		for d := range out {
			out[d] += scores[i] / sum * v[d]
		}
	}
	return out
}

// pagedAttend never sees a contiguous sequence. It walks the block table.
func pagedAttend(q vec, keyCache, valueCache []vec, blockTable []int, seqLen int) vec {
	keys := make([]vec, 0, seqLen)
	values := make([]vec, 0, seqLen)
	for pos := 0; pos < seqLen; pos++ {
		slot := blockTable[pos/blockSize]*blockSize + pos%blockSize
		keys = append(keys, keyCache[slot])
		values = append(values, valueCache[slot])
	}
	return attend(q, keys, values)
}

func main() {
	const seqLen = 10
	// A sequence's keys and values, as an ordinary contiguous array.
	keys := make([]vec, seqLen)
	values := make([]vec, seqLen)
	for i := range keys {
		for d := 0; d < dim; d++ {
			keys[i][d] = math.Sin(float64(i*3 + d))
			values[i][d] = math.Cos(float64(i*5 + d*2))
		}
	}
	q := vec{0.3, -1.1, 0.7}

	// The same data stored the vLLM way: one big cache, the sequence's three
	// blocks scattered through it in no particular order.
	keyCache := make([]vec, numBlocks*blockSize)
	valueCache := make([]vec, numBlocks*blockSize)
	blockTable := []int{6, 1, 4}
	for pos := 0; pos < seqLen; pos++ {
		slot := blockTable[pos/blockSize]*blockSize + pos%blockSize // slot_mapping
		keyCache[slot] = keys[pos]
		valueCache[slot] = values[pos]
	}

	a := attend(q, keys, values)
	b := pagedAttend(q, keyCache, valueCache, blockTable, seqLen)

	fmt.Printf("contiguous attention: [% .6f % .6f % .6f]\n", a[0], a[1], a[2])
	fmt.Printf("paged attention:      [% .6f % .6f % .6f]\n", b[0], b[1], b[2])
	fmt.Println("identical:", a == b)
	fmt.Printf("\nblock table %v: positions 0-3 are in slots %d-%d, 4-7 in %d-%d, 8-9 in %d-%d\n",
		blockTable, 6*blockSize, 6*blockSize+3, 1*blockSize, 1*blockSize+3, 4*blockSize, 4*blockSize+1)
	fmt.Printf("%d of %d slots hold this sequence; the rest belong to other requests or are free\n",
		seqLen, numBlocks*blockSize)
}

pagedAttend here gathers the keys into a temporary contiguous list and then calls the ordinary function. A real kernel does not copy: it reads each block in place as it iterates. The extra indirection — one table lookup per block — is the entire runtime cost of paging, and at 16 tokens per block it is negligible next to the arithmetic.

Remember this#

  • A vLLM attention kernel takes a block table and a slot mapping in addition to queries, keys and values.
  • Paged attention computes exactly the same result as contiguous attention.
  • Backends are tried in a priority order that depends on GPU generation; the first that validates wins.
  • FlashAttention is first on most NVIDIA GPUs, FlashInfer on the newest data-centre generation, Triton is the portable fallback.
  • --attention-backend NAME forces one; startup fails if it cannot support the configuration.
  • The backend decides the cache tensor’s shape and which batch shapes can be CUDA-graphed.
  • Latent-attention and state-space layers use the same plumbing with their own backends.

Try it#

  1. In the program, change the block table to [0, 1, 2], which makes the paged layout contiguous. Confirm the result is unchanged. Then shuffle it differently.
  2. Start a server with VLLM_LOGGING_LEVEL=DEBUG and find the lines where attention backends are considered. Which were skipped, and what reason was logged?
  3. Force each of FLASH_ATTN, FLASHINFER and TRITON_ATTN in turn on the same model and measure decode throughput at 32 concurrent requests with vllm bench serve.

Check yourself#

  1. What two extra inputs distinguish a paged attention kernel from a standard one, and what is each used for?
  2. On an H100, which backend is tried first for a Llama-style model? On a Blackwell data-centre GPU?
  3. Why does the attention backend influence whether a whole forward pass can be replayed as a CUDA graph?

Sources#

Checked on 5 October 2026 against main at commit 0c16eee.

↑↓ navigate↵ openesc close

drag to pan · scroll to zoom