Attention backends

What attention reads. Llama-3-8B keeps, per token per layer, 8 K/V heads × 128 numbers × (k and v) × 2 bytes = 4 KB, and 32 layers make 128 KB per token (Ch. 3 Step 9). A decode step reads every running request’s whole cache once, so the attention kernels together move $B \times n \times 128$ KB, against a fixed 16 GB of weights. At $B = 256$ the two are equal at $n \approx 480$; at $n = 4096$ the cache read is 137 GB, 41 ms at 3.35 TB/s against the weights’ 4.8.

Why it is a memory pump. For one query against one key in one head: a 128-wide dot product is 256 flops and reads 256 bytes of $k$; scaling and adding a 128-wide $v$ is another 256 flops and 256 bytes. One flop per byte. The H100’s ratio of compute to bandwidth is $989 \cdot 10^{12} / 3.35 \cdot 10^{12} \approx 295$ flops per byte (Ch. 4 Step 8), so a single-query attention sits about 300× below the balance point. Grouped-query attention lifts it to 4 (four query heads per K/V head), and MLA’s decode kernel to 128. Prefill multiplies the 4 by the height of the query tile, 128, to 512: past the balance point, which is why prefill attention is compute-bound and decode attention is not. A tile of 128 without the grouping would sit at 128, still below it; in practice neighbouring tiles also share K/V through the 50 MB L2.

FlashAttention. Written as three kernels, attention materialises the $n \times n$ score matrix per head in HBM: at a 2,000-token prefill that is 8 MB per head for S and again for P, each written once and read once, so the three kernels move 34 MB per head where Q, K, V and the output are 2 MB. FlashAttention (Dao, Fu, Ermon, Rudra, Ré, 2022) computes the softmax online: one kernel walks K and V in chunks, keeps a running max, sum and output per row, and never writes the scores out, so the traffic is those 2 MB. The stage shows the one identity that keeps the running result exact; the rest of FlashAttention’s design (tile sizes, the backward pass used in training) is not covered. FlashAttention-2 (Dao, 2023) reorganised the loop so a thread block owns a tile of queries and sweeps all keys, which suits prefill. Flash-Decoding (Dao, Haziza, Massa, Sizov, 2023) splits the keys across blocks for decode and merges their partial results with the same rescaling. The stage’s timings for it assume HBM bandwidth divides evenly across the 132 SMs, so 8 busy SMs pull 8/132 of it; in practice fewer than all SMs can saturate HBM, so the unsplit kernel’s 5.3 ms is an upper bound. FlashAttention-3 (Shah, Bikshandi, Zhang, Thakkar, Ramani, Dao, 2024) rewrote it for Hopper: producer warps drive TMA (Hopper’s tensor memory accelerator) loads, two consumer warpgroups alternate so one’s softmax overlaps the other’s matmul, and FP8 with per-block scaling; it reaches about 75% of the H100’s peak where FA2 reached about 35%.

Paged, ragged batches. SGLang’s pool is page size 1 (Ch. 7 Step 4), so a request’s K/V is a list of page numbers in token order. The batch is described to the kernel by kv_indptr (where each request’s slice starts) and kv_indices (the concatenated page numbers), a CSR layout; backends with larger pages add kv_last_page_len. The arrays are rebuilt on the CPU each batch in init_forward_metadata; FlashInfer (Ye et al., 2025) calls this plan(), and uses it to choose the split-KV chunking per request and balance long and short contexts across SMs. For decode the arrays live in buffers sized for the largest captured batch, which is how Ch. 11 Step 13’s graph can replay a kernel whose inputs change every step.

The interface. sglang/srt/layers/attention/base_attn_backend.py defines AttentionBackend: init_forward_metadata(forward_batch), forward_extend(q, k, v, layer, forward_batch), forward_decode(...), init_cuda_graph_state(max_bs, ...), and the capture/replay hooks that refresh the captured buffers. For a decode batch that fits a captured graph, the runner writes the plan into those buffers and replays; the eager init_forward_metadata and forward run only when there is no graph. Recent SGLang splits the plan into a part built outside the graph and a part recorded inside it (init_forward_metadata_out_graph / _in_graph). RadixAttention.forward in radix_attention.py does nothing but call attn_backend.forward(q, k, v, self, forward_batch). attention_registry.py maps --attention-backend names to classes; the docs’ feature table (docs.sglang.io, Attention Backend) lists per backend the page sizes, cache dtypes, sliding-window and speculative-decoding support. Defaults, from get_default_attn_backend: fa3 on Hopper, trtllm_mha on Blackwell, aiter on AMD, flashinfer on other NVIDIA (falling back to triton); for MLA models fa3 on Hopper and flashinfer on Blackwell, with DeepSeek-V3 on Blackwell overridden to trtllm_mla. --prefill-attention-backend and --decode-attention-backend mix two.

MLA. DeepSeek-V2/V3 cache one 576-wide vector per token (a 512-wide compressed latent plus 64 dimensions for rotary position) shared by all 128 heads. For decode, each head’s key up-projection is folded into its query and its value up-projection into the output projection, so the kernel scores all 128 query heads directly against the cached vector: in effect one shared K/V head, 576 wide for k and 512 for v, and 128 uses of every byte. Prefill can instead expand the latent into per-head k and v and use an ordinary kernel. The decode widths match no MHA kernel, so MLA has its own family: flashinfer (MLA mode), flashmla, cutlass_mla, trtllm_mla, and fa3 / fa4 for MLA shapes.