FlashDecoding + FlashInfer — the decode-time attention story

FlashAttention (any version) was designed for prefill — many queries × many keys, compute-bound, matmul-dominated. Decode is different. During decode, each step has:

  • 1 query per sequence (or a handful with speculative decoding)

  • N keys/values (the growing KV cache)

  • Small batch × few queries × huge N = memory-bound

Run FA2 on decode and you’ll leave a large chunk of the GPU idle because the parallelization axis (batch × heads × Q-blocks) has almost no Q-blocks. FlashDecoding fixes this. FlashInfer takes it further and makes decode-attention a first-class serving primitive with paged and block-sparse KV caches.

FlashDecoding — the split-KV trick

Origin: Tri Dao et al. blog, Oct 2023. No dedicated arXiv paper; the idea is documented in the Dao-AILab repo and Tri Dao’s HuggingFace blog. Ships in flash-attention as the default decode kernel.

The trick: parallelize across the KV dimension, not just batch × heads.

At decode, you have Q with shape (B, H, 1, D) and KV with shape (B, H, N, D) where N is the current context length. FA2 launches ~B×H thread blocks — for typical B=1, H=32, that’s 32 blocks. H100 has 132 SMs. Massive underutilization.

FlashDecoding launches B × H × G blocks, where G is a split factor across N. Each block computes its own partial (m, d, o) for its KV chunk. Then a small reduction kernel combines the G partials using the online-softmax merge rule from 01_online_softmax.md:

For two partial results (m_a, d_a, o_a) and (m_b, d_b, o_b):
    m_new = max(m_a, m_b)
    d_new = d_a * exp(m_a - m_new) + d_b * exp(m_b - m_new)
    o_new = o_a * exp(m_a - m_new) + o_b * exp(m_b - m_new)

This is the same recurrence, applied hierarchically — tile-level in the main kernel, split-level in the reduction.

Result on decode: typical ~ 2–8× speedup vs FA2 at small batch + long context.

When to reach for it: always, for decode. The flash-attn Python API auto-picks FlashDecoding on decode-shaped inputs since ~ v2.3.

FlashDecoding++ (a different paper, same idea zone)

arXiv/venue: Hong et al., MLSys 2024. https://proceedings.mlsys.org/paper_files/paper/2024/file/5321b1dabcd2be188d796c21b733e8c7-Paper-Conference.pdf

Additions on top of the FlashDecoding idea:

  1. Async-softmax with unified max. Uses a statically-inferred upper bound on the softmax max to skip the synchronization between partials, saving ~ 20% of the softmax overhead.

  2. Flat-GEMM aware batching. Adapts tile shapes for the tall-and-skinny GEMMs that show up in decode.

  3. Heuristic + dataflow-aware kernel selection.

Result: 1.37× further speedup vs FlashDecoding on Llama-2 at batch=1.

Awareness-level knowledge. Cite it when someone asks “is there anything past FlashDecoding?” Don’t implement it.

FlashInfer — the serving library

arXiv: <phone_number_or_numberic_id_or_random_id_65> — FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving. Authors: Zihao Ye, Lequn Chen, Ruihang Lai, Wuwei Lin, Yineng Zhang, Stephanie Wang, Tianqi Chen, Baris Kasikci, Vinod Grover, Arvind Krishnamurthy, Luis Ceze (UW + NVIDIA). Venue: MLSys 2025 — Best Paper Award. (https://news.cs.washington.edu/2025/07/01/allen-school-researchers-receive-best-paper-award-for-speeding-up-llm-performance-with-flashinfer) Repo: https://github.com/flashinfer-ai/flashinfer

What FlashInfer is

A library for serving-time attention. FA-family kernels are algorithms; FlashInfer is the runtime + kernel dispatcher that:

  1. Unifies KV-cache layouts. Handles paged (vLLM-style), ragged (variable-length), block-sparse, and radix-tree KV caches under one API.

  2. JIT-compiles attention kernels at server startup / on-demand based on shape + mask + dtype. No re-compile per request.

  3. Load-balances across CTAs. For decode-heavy mixed batches (some sequences short, some long), balances work across SMs so the longest sequence doesn’t stall the batch.

  4. Composable mask types. Full, causal, sliding-window, custom block-sparse, all as first-class layouts.

  5. First-class integration with vLLM, SGLang, MLC-LLM.

The numbers (from the paper)

  • 29–69% inter-token latency reduction vs compiler-produced attention backends (e.g., torch.compile / cuDNN default paths).

  • 28–30% latency reduction on long-context inference.

  • 13–17% speedup for parallel generation (multi-sample / speculative).

The API you actually use (Python)

import flashinfer

# 1) Build a wrapper. This is expensive; do it once per model instance.
workspace = torch.zeros(128*1024*1024, dtype=torch.uint8, device="cuda")
wrapper = flashinfer.BatchDecodeWithPagedKVCacheWrapper(workspace, "NHD")

# 2) Plan for the shapes/masks you'll use. Amortized.
wrapper.plan(
    indptr,            # page-table indptr (paged KV cache)
    indices,           # page indices
    last_page_len,     # per-request last-page length
    num_qo_heads, num_kv_heads, head_dim,
    page_size,
    q_data_type=torch.bfloat16,
    kv_data_type=torch.bfloat16,
)

# 3) Run. Hot path.
out = wrapper.run(q, paged_kv_cache)

Key concepts:

  • plan() amortizes JIT + workspace setup and mask indexing. Call once per shape configuration.

  • run() is the fast path.

  • Different wrappers exist: BatchDecode, BatchPrefill, single, Ragged, PagedKV, BlockSparse.

Read the docs: https://docs.flashinfer.ai

Where FlashInfer sits in the stack

          Application: vLLM / SGLang / your service
                            │
                     FlashInfer wrappers
                     (plan, run, KV layout)
                            │
                  JIT-compiled attention kernels
                    (FA2-style, FA3-style, decode-split)
                            │
                       CUDA / SASS

Use FlashInfer whenever you’re building a serving path with paged KV, variable context lengths, or custom mask patterns. It saves you re-implementing the plumbing that has been carefully optimized once.

Paged attention (vLLM) — the KV-cache layout FlashInfer targets

Origin: vLLM paper (SOSP 2023). Inspired by OS virtual memory: split the KV cache into fixed-size blocks (“pages”), maintain a per-sequence page table.

Why:

  • Eliminates KV-cache fragmentation. Without paging, allocating for max-context-per-request wastes gigabytes when actual contexts vary.

  • Enables prefix sharing (multiple requests sharing early tokens share pages).

  • Enables copy-on-write for parallel sampling and beam search.

Layout (conceptual):

Sequence 1: page_ids = [7, 2, 15, 88, 3]     # ~ block_size * len tokens
Sequence 2: page_ids = [7, 2, 15, 41]        # shares first 3 pages with seq1 (prefix caching)

pages: uint8/bf16 blocks of shape (num_pages, block_size, num_kv_heads, head_dim)

FlashInfer’s PagedKV wrappers take (page_ids indptr, last_page_len) and gather the right pages inside the kernel.

Read:

Ragged attention — variable-length batching

Rather than pad every sequence to N_max, ragged/varlen attention lays out queries and keys contiguously:

Q (varlen): shape (sum_i(q_len_i), H, D)   with qo_indptr = [0, q_len_0, q_len_0 + q_len_1, ...]
KV (varlen): shape (sum_i(k_len_i), H, D)  with kv_indptr = analogous

Kernels iterate per-sequence using the indptr arrays. Zero wasted compute on padding.

FlashInfer’s BatchPrefillWithRaggedKVCacheWrapper is your handle.

Concrete decisions — which kernel do I use when?

Situation

Use

Training / fixed-length prefill

FA2 (via flash-attn or Triton)

Long-context prefill on H100

FA3 (via flash-attn hopper backend) or SDPA with flash backend

Blackwell prefill

FA4 (via flash-attn blackwell backend when released)

Small-batch / long-context decode

FlashDecoding (auto in flash-attn decode path)

Production serving with paged KV

FlashInfer + vLLM

Sliding-window / custom mask

FlashInfer block-sparse or custom Triton

Sparse / MoE routing on attention

Roll a Triton kernel; check FlashInfer’s sparse wrapper first

What to actually do this week

  1. Read FlashInfer paper §3 (unified KV layout) and §4 (JIT compiler).

  2. Install flashinfer and run BatchDecodeWithPagedKVCacheWrapper on a toy example. Print plan() vs run() timings so you feel the amortization.

  3. Read Tri Dao’s FlashDecoding blog + the relevant section of the flash-attention README.

  4. Sketch the merge rule for two partial (m, d, o)s on paper. Verify against the recurrence from 01_online_softmax.md.

  5. In your notebook: measure decode-attention on Llama-3 8B at (B=1, N=8k) with vanilla SDPA vs FlashDecoding path vs FlashInfer. Report the tokens/sec delta.

References