FlashInfer — infra.bea.ai

FlashAttention was designed for training — fixed batch, uniform sequence lengths, no KV cache. FlashInfer extends flash attention for serving, where every request has a different prompt length, different generation progress, and KV cache stored in non-contiguous pages.

1. The Serving Problem

Training attention: all sequences same length, contiguous memory, batch processed together.

Serving attention: every request is different:

Property	Training	Serving
Sequence lengths	Uniform (padded)	Variable (each request different)
KV cache	Contiguous tensors	Paged (non-contiguous blocks, via PagedAttention)
Batch	Same operation per sample	Mix of prefill and decode requests
Memory layout	Dense, predictable	Sparse, fragmented

Standard FlashAttention assumes contiguous K/V tensors. In serving with PagedAttention, KV cache is scattered across fixed-size pages — FlashAttention can’t handle this directly.

2. FlashInfer’s Approach

Page-Level Attention

FlashInfer natively understands paged KV cache. Each request has a page table mapping logical sequence positions to physical memory pages:

Request 0: page_table = [page_5, page_12, page_3, page_8]
  Tokens 0-63   → page 5
  Tokens 64-127  → page 12
  Tokens 128-191 → page 3
  Tokens 192-255 → page 8

Request 1: page_table = [page_1, page_7]
  Tokens 0-63   → page 1
  Tokens 64-100  → page 7

FlashInfer’s kernel follows the page table to gather KV cache from non-contiguous pages — similar to how a CPU follows page tables for virtual memory.

Ragged Batching

FlashInfer handles ragged tensors — batches where each element has a different length — natively. No padding needed:

import flashinfer

# Variable-length sequences, no padding
q_lens = [128, 256, 64, 512]  # different length per request
kv_lens = [1024, 2048, 512, 4096]  # different KV cache sizes

# FlashInfer handles the ragged batch directly
output = flashinfer.batch_decode_with_paged_kv_cache(
    q, kv_data, kv_indices, kv_indptr, kv_last_page_len
)

Prefill-Decode Mixing

In serving, the batch may contain both prefill requests (processing full prompts) and decode requests (generating one token). FlashInfer supports mixed batches with different kernels per request type.

3. Key Optimizations

Cascade Attention for Shared Prefixes

When multiple requests share a system prompt:

Request 0: [system prompt (2K tokens)] + [user message A (100 tokens)]
Request 1: [system prompt (2K tokens)] + [user message B (200 tokens)]
Request 2: [system prompt (2K tokens)] + [user message C (50 tokens)]

FlashInfer computes attention over the shared prefix once, then extends per-request:

# Step 1: Shared prefix attention (computed once)
shared_output = flashinfer.single_prefill(q_shared, k_prefix, v_prefix)

# Step 2: Per-request suffix attention (computed per request)
# Merges with shared output using online softmax correction
final_output = flashinfer.merge_state(shared_output, suffix_outputs)

This avoids redundant computation on the shared prefix — significant for chat applications where system prompts are long.

FlashDecoding: Parallel KV Split

For long-context decode (one new token, but 100K+ cached KV):

Standard decode: one thread block processes the entire KV sequence → serialized.

FlashDecoding: split KV across multiple thread blocks, each computes partial attention, then merge:

Thread Block 0: KV tokens 0 - 16K      → partial output + (max, sum)
Thread Block 1: KV tokens 16K - 32K    → partial output + (max, sum)
Thread Block 2: KV tokens 32K - 48K    → partial output + (max, sum)
...
Merge: combine partial outputs using online softmax correction

Benefit: Parallelism across the sequence dimension — critical for long-context decode where a single block would be too slow.

4. FlashInfer in the Serving Stack

FlashInfer is the attention backend for major serving engines:

Engine	Attention backend	How FlashInfer is used
vLLM	FlashInfer or FlashAttention	Paged KV cache attention, batch decode
SGLang	FlashInfer (primary)	RadixAttention for prefix caching
TensorRT-LLM	Custom kernels	Similar concepts, proprietary implementation

The Interface

# Typical serving engine integration:

# 1. Build page table for the batch
kv_indices, kv_indptr, kv_last_page_len = build_page_tables(requests)

# 2. Prefill requests
prefill_output = flashinfer.batch_prefill_with_paged_kv_cache(
    q_prefill, kv_data, kv_indices, kv_indptr, kv_last_page_len,
    causal=True
)

# 3. Decode requests (one new token per request)
decode_output = flashinfer.batch_decode_with_paged_kv_cache(
    q_decode, kv_data, kv_indices, kv_indptr, kv_last_page_len
)

5. FlashInfer vs FlashAttention

Feature	FlashAttention	FlashInfer
Target	Training (and inference)	Serving (production inference)
KV layout	Contiguous tensors	Paged (non-contiguous)
Batch type	Uniform lengths	Ragged (variable lengths)
Prefix sharing	Not supported	✅ Cascade attention
Decode parallelism	Single-block	✅ FlashDecoding (multi-block)
GQA/MQA support	✅	✅
MLA support	Limited	✅ (DeepSeek-V2 compatible)

Use FlashAttention for training. Use FlashInfer for serving.

Self-Check

Explain why standard FlashAttention can’t handle paged KV cache directly
Explain how FlashInfer follows page tables to gather non-contiguous KV cache
Explain cascade attention for shared prefixes and its benefit
Explain FlashDecoding and why it’s needed for long-context decode
Name which serving engines use FlashInfer and how
Explain when to use FlashAttention vs FlashInfer