Deep Dive - FlashInfer

Three requests reach one GPU step together. R0 processes 3 new tokens after a 3-token prefix, R1 processes 1 after a 1-token prefix, and R2 processes 4 after a 5-token prefix. Their 8 query rows fit in one compact buffer. After appending the new keys and values, their KV lengths are 6, 2, and 9, scattered across physical pages. Attention still has to return one correct state per row without attending to a future token.

Before naming an API, predict the contract. The kernel needs query boundaries, a page list plus tail length for each history, and a work assignment that can split a long history without losing softmax normalization. Those pieces explain FlashInfer: a library and kernel generator for attention, matrix multiplication, mixture-of-experts operations, sampling, and cache updates.^{[1]Reference 1FlashInferhttps://github.com/flashinfer-ai/flashinfer} The serving engine still owns allocation and admission.

FlashAttention supplies the tiled softmax, and paged KV caching supplies the logical-to-physical mapping. Keep those two mechanisms separate: avoiding a score matrix and avoiding cache relocation save different memory traffic. The Python exercises here run on CPU and check indexing and attention arithmetic. They don't execute FlashInfer kernels or measure GPU speed.

Who ships it, who calls it

Start with ownership, because a kernel benchmark and a serving guarantee belong to different layers.

The API walkthrough is pinned to FlashInfer 0.6.18, released August 29, 2026, at commit 69ff11fc4954396d98326656dc85debd2223f637; its relevant source paths were rechecked September 23, 2026.^{[2]Reference 2FlashInfer v0.6.18 Releasehttps://github.com/flashinfer-ai/flashinfer/releases/tag/v0.6.18[3]Reference 3FlashInfer v0.6.18 Sourcehttps://github.com/flashinfer-ai/flashinfer/tree/69ff11fc4954396d98326656dc85debd2223f637} The latest release on that date is v0.7.0, published September 22. Its cuDNN normalization fix matters to the state-merge contract below.^{[4]Reference 4FlashInfer v0.7.0 Releasehttps://github.com/flashinfer-ai/flashinfer/releases/tag/v0.7.0} The paper benchmarks later in the lesson used v0.2. They aren't measurements of either release.

As checked October 10, the latest release is v0.7.1, published October 9. The older reproduction stays pinned. The new release adds planned MLA backend surfaces and a named trtllm_prefill_with_kv_cache_mla entry point, but that name doesn't remove decode constraints: its backend="auto" can choose XQA, which supports one query token per request. Multi-token prefill callers must explicitly choose a compatible backend.^{[5]Reference 5FlashInfer v0.7.1 Releasehttps://github.com/flashinfer-ai/flashinfer/releases/tag/v0.7.1}

FlashInfer began in summer 2023 with researchers from the University of Washington, Carnegie Mellon University, and OctoAI.^{[6]Reference 6FlashInfer: Kernel Library for LLM Servinghttps://flashinfer.ai/2024/02/02/introduce-flashinfer.html} The MLSys 2025 paper lists Zihao Ye, Lequn Chen, Ruihang Lai, Wuwei Lin, Yineng Zhang, Stephanie Wang, Tianqi Chen, Baris Kasikci, Vinod Grover, Arvind Krishnamurthy, and Luis Ceze, with affiliations spanning UW, NVIDIA, Perplexity, CMU, and independent research.^{[7]Reference 7FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving.https://arxiv.org/abs/2501.01005} The cited contribution snapshot documents NVIDIA participation alongside the FlashInfer community, with public GitHub Actions and an NVIDIA-internal GPU matrix on GitLab.^{[8]Reference 8Contributing to FlashInferhttps://github.com/flashinfer-ai/flashinfer/blob/d020372b068f335e2fe427372e134977a2235c49/CONTRIBUTING.md}

The repository names SGLang, vLLM, TensorRT-LLM, Text Generation Inference, MLC-LLM, LightLLM, lorax, and ScaleLLM as adopters or integrations.^{[1]Reference 1FlashInferhttps://github.com/flashinfer-ai/flashinfer} Adoption doesn't mean every framework enables every operator. Check the framework's attention backend, FlashInfer release, CUDA version, and model feature matrix before treating a project-level capability as a service guarantee.

Ownership says who is responsible for a number. It doesn't explain why three requests in one step stop looking like one attention matrix.

Why inference attention becomes irregular

Each request keeps its own length, so the batch is no longer one rectangle. FlashAttention gives exact attention a streaming recipe: load tiles, update online softmax state, and avoid writing a full score matrix to global memory. In serving, the pressure moves to the batch around that kernel. The transformer equation stays regular, but requests arrive and finish independently, sequence lengths change every step, and the key-value (KV) cache comes from a shared page pool.

Start by counting two axes. R0 and R2 are doing chunked prefill, also called append attention. R1 is in the decode loop. Here, every reported KV length includes the current query tokens' keys and values, which the caller has already appended. A ragged query buffer stores the 8 query rows back-to-back. What tells a kernel where one request ends and the next begins? An index-pointer array qo_indptr:

The query tensor is compact. Now predict what the cache needs: a way to name physical pages and a way to mark the occupied part of the last page. Page size $P = 4$ makes allocation independent of exact sequence length, much like virtual memory maps a process address to a physical frame.^{[10]Reference 10Efficient Memory Management for Large Language Model Serving with PagedAttention.https://arxiv.org/abs/2309.06180} FlashInfer's length contract is page_size * (len(page_indices) - 1) + last_page_len, with last_page_len in 1 .. page_size.^{[11]Reference 11KV-Cache Layoutshttps://docs.flashinfer.ai/tutorials/kv_layout.html}

Packed metadata for the whole batch is kv_indptr = [0, 2, 3, 6], kv_indices = [7, 2, 11, 3, 14, 6], and kv_last_page_len = [2, 2, 1]. Read the two pointer arrays in different coordinate systems: qo_indptr counts query rows, while kv_indptr counts page IDs. Confusing those units can still produce in-bounds pointers with the wrong sequence.

A cooperative thread array (CTA) is CUDA's name for a thread block: threads that share on-chip memory and run one tile of work. These tiny lengths illustrate indexing. Benchmark CTA splitting on long histories and account for query tiles, heads, backend, and hardware before choosing a launch strategy.

Can you rebuild both pointer arrays before invoking attention? The script below does so from the same page lists. If a later figure or kernel dump disagrees with qo_indptr = [0, 3, 4, 8] or KV lengths [6, 2, 9], stop there: metadata is wrong before any attention math runs.

PAGE_SIZE = 4
pages = {0: [7, 2], 1: [11], 2: [3, 14, 6]}
last_page_len = {0: 2, 1: 2, 2: 1}

def kv_len(page_ids: list[int], tail: int) -> int:
    if not page_ids or any(type(p) is not int or p < 0 for p in page_ids):
        raise ValueError("this example requires nonempty, nonnegative page IDs")
    if type(tail) is not int or not 1 <= tail <= PAGE_SIZE:
        raise ValueError("tail must occupy 1..PAGE_SIZE slots")
    return PAGE_SIZE * (len(page_ids) - 1) + tail

qo_indptr = [0]
for query_len in (3, 1, 4):
    qo_indptr.append(qo_indptr[-1] + query_len)

kv_indptr = [0]
for request in (0, 1, 2):
    kv_indptr.append(kv_indptr[-1] + len(pages[request]))

lengths = [kv_len(pages[request], last_page_len[request]) for request in (0, 1, 2)]
assert qo_indptr == [0, 3, 4, 8]
assert kv_indptr == [0, 2, 3, 6]
assert lengths == [6, 2, 9]
for bad_pages, bad_tail in (([], 1), ([3], 0), ([3], 5), ([-1], 1)):
    try:
        kv_len(bad_pages, bad_tail)
    except ValueError:
        pass
    else:
        raise AssertionError("invalid page metadata was accepted")
print("qo_indptr", qo_indptr)
print("kv_indptr", kv_indptr)
print("kv_len", lengths)

qo_indptr [0, 3, 4, 8]
kv_indptr [0, 2, 3, 6]
kv_len [6, 2, 9]

View larger

Layouts are an API, not an implementation detail

FlashInfer documents several logical layouts because no single physical arrangement wins every phase. Ask what each phase wants to keep contiguous. Ragged storage packs tokens without padding; paged storage breaks a sequence into fixed-size pages; block-sparse row (BSR) metadata views that page table as a sparse matrix. Its CSR-style indptr / indices arrays tell the kernel which blocks to load. Multi-head latent attention (MLA) can store compressed latent states instead of ordinary K/V heads.^{[11]Reference 11KV-Cache Layoutshttps://docs.flashinfer.ai/tutorials/kv_layout.html}

The same choice extends to precision. Selected backends accept 8-bit floating point (FP8) or FP4 representations, with format-specific scales and packing. Those aren't drop-in replacements for FP16: scale granularity, cache layout, supported GPU, and dequantization path belong to the selected operator's contract. The example below stays in FP16 rather than assuming every wrapper supports every quantized cache.

The kernel needs three kinds of information:

1. Payload: query, key, value, or latent tensors (in FP16, BF16, FP8, or sub-byte quantized formats).
2. Shape metadata: head counts, head dimensions, page size, and request lengths.
3. Indirection: row pointers, page indices, last-page lengths, and optional sparse masks.

Payload buffers can remain stable while metadata changes every scheduling step. A serving engine can recycle pages without asking a kernel to understand allocator policy. A kernel can specialize memory access around a known page size without owning the allocator.

Ragged prefill and paged decode

Prefill consumes many new tokens per request. A ragged layout keeps those tokens contiguous, so a batch prefill wrapper can traverse row ranges with one set of offsets. Decode usually consumes one token per live request and reads a long history. A paged layout avoids copying each history into a newly padded tensor.

Our batch is mixed: R0 and R2 have several query rows, R1 has one. A paged-prefill wrapper can represent all three through qo_indptr; a dedicated decode wrapper assumes one query per request. FlashInfer also has other mixed-phase paths, but those aren't interchangeable APIs. Choose a named wrapper and verify its signature before passing scheduler metadata.^{[1]Reference 1FlashInferhttps://github.com/flashinfer-ai/flashinfer}

Causality uses logical positions, not page numbers

R2 has five earlier tokens and four new ones. Its first query is at logical position 5, so it can attend to positions 0 through 5, not the whole nine-token cache. For local query index $i$, query length $Q$, and total KV length $K$, the allowed key positions satisfy $j\le K-Q+i$. This is a bottom-right-aligned causal mask. A top-left triangle would expose too little prefix; no mask would expose future tokens.

Run the indexing check before any GPU call. It also shows why sorting the page IDs would change the sequence:

page_ids, page_size, q_len, kv_len = [3, 14, 6], 4, 4, 9
addresses = [(page_ids[j // page_size], j % page_size) for j in range(kv_len)]
mask = [[j <= kv_len - q_len + i for j in range(kv_len)] for i in range(q_len)]
assert addresses[5] == (14, 1)
assert addresses[-1] == (6, 0)
assert [sum(row) for row in mask] == [6, 7, 8, 9]
assert [p for p, _ in addresses] != sorted(p for p, _ in addresses)
print("logical token 5 -> page, slot", addresses[5])
for row in mask:
    print("".join("1" if visible else "." for visible in row))

logical token 5 -> page, slot (14, 1)
111111...
1111111..
11111111.
111111111

BSR as a memory-access contract

BSR metadata doesn't mean the data are mathematically sparse in every model. It says the kernel may load fixed-size blocks through an index table. If R0 references pages [7, 2], the GPU can gather those pages without a defragmentation copy. If a prefix is shared, several request rows can point at the same physical pages while their suffix pages diverge.

That indirection adds pointer arithmetic and less predictable memory access. It pays off when avoiding copies and padding saves more work than extra gathers cost. Tiny sequences can lose to a simpler contiguous kernel. Use the layout only when its saved movement outweighs its indexing cost; whether it helps depends on the workload.

A long history may still need several thread blocks. The layout says where the keys live. It doesn't say how partial attention states combine.

Core mechanism: schedule work around state

Suppose one query's KV history is split across two CTAs. What can each worker discard while still letting a parent recover exact attention? Normalization has to survive. An attention kernel normally looks like a matrix multiplication followed by a softmax and a value multiplication, but FlashAttention already showed that the full score matrix is unnecessary. FlashInfer makes the output vector plus a log-sum-exp statistic the composable unit across workers, devices, and shared prefixes.^{[7]Reference 7FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving.https://arxiv.org/abs/2501.01005[12]Reference 12Recursive Attentionhttps://docs.flashinfer.ai/tutorials/recursive_attention.html}

If two workers process disjoint key ranges for the same query row and head, each returns an output vector $o$ and a base-2 log-sum-exp $b$. For ordinary scaled logits $s_j=qk_j^T/\sqrt{d}$, that statistic is $b=\log_2\sum_j e^{s_j}$, not $\log_2\sum_j 2^{s_j}$. The conversion changes the stored normalization base, not the attention distribution. FlashInfer's public merge_state expects base-2 statistics; don't pass a natural-log LSE without converting it by division by $\ln 2$.^{[13]Reference 13FlashInfer v0.6.18: Cascade Attentionhttps://github.com/flashinfer-ai/flashinfer/blob/69ff11fc4954396d98326656dc85debd2223f637/include/flashinfer/attention/cascade.cuh}

This is a real version boundary: v0.7.0 includes a fix converting the cuDNN frontend graph path's natural-log LSE to the expected base-2 value. The pinned v0.6.18 path returns that LSE without conversion.^{[14]Reference 14FlashInfer: Fix cuDNN Prefill LSE Base and Head-Dimension Offsetshttps://github.com/flashinfer-ai/flashinfer/pull/4663[4]Reference 4FlashInfer v0.7.0 Releasehttps://github.com/flashinfer-ai/flashinfer/releases/tag/v0.7.0} Identify the producer's backend and version before adding a conversion; dividing an already corrected base-2 statistic by $\ln 2$ again changes the merge weights.

The v0.7.1 low-level cuDNN prefill interface makes the choice explicit: lse_base="log2" is the default, while lse_base="ln" keeps the frontend's native natural-log statistic. Convert the latter once before a base-2 merge. The cubin cuDNN path only returns base-2 LSE and rejects the "ln" request. Version alone therefore doesn't identify the statistic's units; record the backend and argument too.^{[15]Reference 15FlashInfer v0.7.1 cuDNN Prefill Interfacehttps://github.com/flashinfer-ai/flashinfer/blob/0b46115d238d29ea2e92c6b9b2cb077d40fbc352/flashinfer/cudnn/prefill.py}

Before reading the recurrence, keep one invariant in mind: every partial output is normalized over its own key range, and its statistic records that range's total weight. The merge must restore the weighting between ranges.

The max subtraction keeps the powers of two in a safe range. Normalize the weights before combining vectors: adding two weighted values near 1e308 and dividing afterward can overflow even when the final convex combination is representable. The pinned CUDA merge also computes normalized scales first.^{[13]Reference 13FlashInfer v0.6.18: Cascade Attentionhttps://github.com/flashinfer-ai/flashinfer/blob/69ff11fc4954396d98326656dc85debd2223f637/include/flashinfer/attention/cascade.cuh} The operator is associative in exact arithmetic, so a tree can merge many chunks without replaying their key-value tiles.^{[12]Reference 12Recursive Attentionhttps://docs.flashinfer.ai/tutorials/recursive_attention.html} Floating-point reduction order can still change the result. Overlapping chunks double-count keys; mismatched masks, positions, or score scales combine different attention problems.

View larger

A worked split-KV example

Take one query row and two key chunks. Chunk A has $b_A = 2$ and output $o_A = [1, 0]$. Chunk B has $b_B = 1$ and output $o_B = [0, 1]$.

1. Pick $m = \max(2, 1) = 2$.
2. Compute weights $w_A = 2^{0} = 1$ and $w_B = 2^{-1} = 0.5$.
3. Sum weights $z = 1.5$.
4. Combine outputs: $o = ([1, 0] + 0.5 \cdot [0, 1]) / 1.5 = [2/3, 1/3]$.
5. Keep $b = 2 + \log_2(1.5) \approx 2.585$ for a parent merge.

The merge doesn't need the original scores. Split-KV workers can run on separate pages or sequence ranges, then publish a small state record. Matching state precision, base, and masking rules makes the merged value mathematically equivalent to one larger softmax.

import math

def merge_state(
    o_a: list[float], b_a: float, o_b: list[float], b_b: float
) -> tuple[list[float], float]:
    """CPU reference; -inf with a zero vector represents an empty key set."""
    if not o_a or len(o_a) != len(o_b):
        raise ValueError("state vectors must have the same positive dimension")
    if not all(math.isfinite(x) for x in o_a + o_b):
        raise ValueError("state vectors must be finite")
    if any(math.isnan(b) or b == math.inf for b in (b_a, b_b)):
        raise ValueError("LSE must be finite or negative infinity")
    if any(b == -math.inf and any(o) for o, b in ((o_a, b_a), (o_b, b_b))):
        raise ValueError("empty states must use a zero vector")
    if b_a == -math.inf:
        return list(o_b), b_b
    if b_b == -math.inf:
        return list(o_a), b_a
    shift = max(b_a, b_b)
    w_a = 2 ** (b_a - shift)
    w_b = 2 ** (b_b - shift)
    z = w_a + w_b
    a_scale, b_scale = w_a / z, w_b / z
    merged = [a_scale * x + b_scale * y for x, y in zip(o_a, o_b, strict=True)]
    return merged, shift + math.log2(z)

output, log_sum_exp = merge_state([1.0, 0.0], 2.0, [0.0, 1.0], 1.0)
assert abs(output[0] - 2 / 3) < 1e-12
assert abs(output[1] - 1 / 3) < 1e-12
assert abs(log_sum_exp - (2 + math.log2(1.5))) < 1e-12
assert merge_state([0., 0.], -math.inf, output, log_sum_exp) == (output, log_sum_exp)
assert merge_state([0., 0.], -math.inf, [0., 0.], -math.inf) == ([0., 0.], -math.inf)
assert merge_state([1e308], 0.0, [1e308], 0.0)[0] == [1e308]
print([round(value, 4) for value in output], round(log_sum_exp, 4))

[0.6667, 0.3333] 2.585

An empty or fully masked chunk has no normalized output. This reference defines it as a zero vector with $b=-\infty$ and handles it explicitly. Blindly evaluating $-\infty-(-\infty)$ produces NaN. The pinned public CUDA merge kernel uses the finite-state recurrence directly.^{[13]Reference 13FlashInfer v0.6.18: Cascade Attentionhttps://github.com/flashinfer-ai/flashinfer/blob/69ff11fc4954396d98326656dc85debd2223f637/include/flashinfer/attention/cascade.cuh} Check the producing kernel's masked-state convention before passing its output to a merge; arbitrary empty sentinels may be unsupported.

Now compare split attention with a direct stable softmax, including large logits and a chunk with no visible keys. This continues the preceding cell. The random numbers are synthetic arithmetic inputs, not model outputs or a performance benchmark.

import random

def direct_state(scores, values):
    if not scores:
        return [0.0, 0.0], -math.inf
    shift = max(scores)
    weights = [math.exp(s - shift) for s in scores]
    denominator = sum(weights)
    output = [sum(w * v[d] for w, v in zip(weights, values, strict=True))
              / denominator for d in range(2)]
    return output, (shift + math.log(denominator)) / math.log(2)

rng = random.Random(7)
cases = 0
for length in (1, 2, 6, 9, 33):
    for offset in (-1000.0, 0.0, 1000.0):
        scores = [offset + rng.uniform(-4, 4) for _ in range(length)]
        values = [[rng.uniform(-1, 1), rng.uniform(-1, 1)] for _ in scores]
        expected_o, expected_b = direct_state(scores, values)
        for cut in range(length + 1):
            left = direct_state(scores[:cut], values[:cut])
            right = direct_state(scores[cut:], values[cut:])
            actual_o, actual_b = merge_state(*left, *right)
            assert max(abs(a - b) for a, b in zip(actual_o, expected_o)) < 1e-10
            assert abs(actual_b - expected_b) < 1e-10
            cases += 1

# Overlap counts a key twice. It is not a valid partition.
scores, values = [0.0, 0.0], [[1.0, 0.0], [0.0, 1.0]]
whole = direct_state(scores, values)
duplicate = merge_state(*whole, *direct_state(scores[:1], values[:1]))
assert abs(whole[0][0] - 0.5) < 1e-12
assert abs(duplicate[0][0] - 2 / 3) < 1e-12
# A natural-log producer has the right masses but the wrong merge units.
correct = merge_state([1., 0.], 2., [0., 1.], 1.)
wrong_base = merge_state([1., 0.], 2 * math.log(2), [0., 1.], math.log(2))
assert abs(wrong_base[0][0] - correct[0][0]) > 0.04
print("split/direct checks:", cases)
print("overlapping chunks change 0.5000 to", f"{duplicate[0][0]:.4f}")
print("wrong LSE base changes 0.6667 to", f"{wrong_base[0][0]:.4f}")

split/direct checks: 168
overlapping chunks change 0.5000 to 0.6667
wrong LSE base changes 0.6667 to 0.6179

Plan then run: inspector-executor for GPU kernels

The state merge solves how to combine work. A second question is when to decide where that work goes. FlashInfer wrappers separate a metadata pass (plan) from the hot data pass (run). The plan computes offsets, request-to-tile assignments, temporary-buffer sizes, and choices such as whether to split KV work. The run consumes that plan with query and cache tensors.

Ask which work should happen once per batch and which should happen on every layer launch. The paper describes this split as an inspector-executor pattern: inspect irregular metadata on the CPU, then execute a stable GPU kernel, including under CUDA Graphs, which want the same launch shape and pointer layout each replay.^{[7]Reference 7FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving.https://arxiv.org/abs/2501.01005}

View larger

Metadata values can change while allocated tensor shapes stay fixed. Reuse a plan across compatible layer calls for the same batch specification, not blindly across decode steps whose sequence lengths have grown. The pinned prefill plan performs host-side work and device copies; it can't run inside CUDA Graph capture or torch.compile.^{[16]Reference 16FlashInfer v0.6.18: Prefill APIhttps://github.com/flashinfer-ai/flashinfer/blob/69ff11fc4954396d98326656dc85debd2223f637/flashinfer/prefill.py} Compatible run calls can be captured. Graph-mode wrappers reserve metadata buffers, require a fixed batch size, and bound total query rows and page-index capacity. Stable buffer addresses don't make stale buffer contents correct.

It also creates an explicit lifetime contract:

The ordinary batch-decode wrapper doesn't take qo_indptr, because each request contributes one query row. Our paged-prefill wrapper does. In v0.6.18 the keyword is head_dim_qk, and paged metadata names start with paged_kv_. The earlier shorthand head_dim / kv_indptr would raise a keyword error for this wrapper.^{[16]Reference 16FlashInfer v0.6.18: Prefill APIhttps://github.com/flashinfer-ai/flashinfer/blob/69ff11fc4954396d98326656dc85debd2223f637/flashinfer/prefill.py}

The following is a GPU-only smoke check source-checked against v0.6.18. Run it on a compatible NVIDIA GPU with CUDA-enabled PyTorch and flashinfer-python==0.6.18; JIT dependencies and artifacts must match the chosen CUDA build. It allocates synthetic Q/K/V directly, with no model weights or checkpoint downloads. The 128 MiB scratch allocation follows the wrapper's documented recommendation; profile the target path to measure its actual requirement.

This example explicitly selects fa2 and counts query offsets in tokens. The pinned cuDNN prefill contract instead documents query offsets in elements, multiplying cumulative query lengths by the number of query heads and head_dim_qk. The same v0.7.0 fix corrects output offsets when head_dim_vo differs from head_dim_qk.^{[14]Reference 14FlashInfer: Fix cuDNN Prefill LSE Base and Head-Dimension Offsetshttps://github.com/flashinfer-ai/flashinfer/pull/4663} Changing the backend label requires checking that backend's metadata and argument contract.

In v0.7.1, the low-level cuDNN call defaults to batch_offsets_units="elements" but also accepts explicit "tokens". In element mode, Q offsets scale by num_qo_heads * head_dim_qk, while O offsets scale by num_qo_heads * head_dim_vo. Copying Q offsets to O is wrong when those dimensions differ. In token mode, the same token boundaries can describe both tensors. Choose units explicitly when migrating and verify output rows; an in-bounds integer array can still address the wrong rows.^{[15]Reference 15FlashInfer v0.7.1 cuDNN Prefill Interfacehttps://github.com/flashinfer-ai/flashinfer/blob/0b46115d238d29ea2e92c6b9b2cb077d40fbc352/flashinfer/cudnn/prefill.py}

import torch
import flashinfer
from importlib.metadata import version

assert version("flashinfer-python") == "0.6.18"
assert torch.cuda.is_available(), "requires NVIDIA CUDA; CPU reference is above"
torch.manual_seed(7)
device = "cuda:0"
workspace = torch.empty(128 * 1024 * 1024, dtype=torch.uint8, device=device)
wrapper = flashinfer.BatchPrefillWithPagedKVCacheWrapper(
    workspace, kv_layout="NHD", backend="fa2"
)
index = lambda values: torch.tensor(values, dtype=torch.int32, device=device)
q = torch.randn(8, 32, 128, dtype=torch.float16, device=device)
paged_kv_cache = torch.randn(15, 2, 4, 8, 128, dtype=torch.float16, device=device)
wrapper.plan(
    qo_indptr=index([0, 3, 4, 8]),
    paged_kv_indptr=index([0, 2, 3, 6]),
    paged_kv_indices=index([7, 2, 11, 3, 14, 6]),
    paged_kv_last_page_len=index([2, 2, 1]),
    num_qo_heads=32,
    num_kv_heads=8,
    head_dim_qk=128,
    page_size=4,
    causal=True,
    q_data_type=torch.float16,
    kv_data_type=torch.float16,
)
output = wrapper.run(q, paged_kv_cache)
assert output.shape == q.shape and torch.isfinite(output).all()

# Independent dense reference: gather logical pages and repeat GQA KV heads.
for row_start, row_end, ids, length in (
    (0, 3, [7, 2], 6), (3, 4, [11], 2), (4, 8, [3, 14, 6], 9)
):
    k = paged_kv_cache[ids, 0].reshape(-1, 8, 128)[:length].float()
    v = paged_kv_cache[ids, 1].reshape(-1, 8, 128)[:length].float()
    k, v = k.repeat_interleave(4, dim=1), v.repeat_interleave(4, dim=1)
    query = q[row_start:row_end].float()
    scores = torch.einsum("qhd,khd->hqk", query, k) / (128 ** 0.5)
    q_len = row_end - row_start
    visible = torch.arange(length, device=device)[None, :] <= (
        length - q_len + torch.arange(q_len, device=device)[:, None]
    )
    probabilities = scores.masked_fill(~visible[None], -torch.inf).softmax(-1)
    expected = torch.einsum("hqk,khd->qhd", probabilities, v)
    torch.testing.assert_close(output[row_start:row_end].float(), expected,
                               atol=2e-3, rtol=2e-3)
torch.cuda.synchronize()
print("paged causal GQA reference passed:", tuple(output.shape))

Here NHD means each K or V page has shape [page_size, kv_heads, head_dim]. The 32 query heads share 8 KV heads in groups of four. Those are grouped-query attention (GQA) semantics, not four copies of each physical cache allocation. The dense comparison repeats heads only in the reference. Its tolerance is a starting check for these random FP16 inputs, not an accuracy budget for every model.

Refresh page mappings and lengths through the public planning path when the batch changes. Some integrations can update compatible fixed-address buffers without rebuilding all work assignments; that support is backend-specific. Replanning metadata doesn't normally require recompiling a kernel or recapturing a compatible graph. Conversely, fixed_split_size can change the number of launched CTAs as KV lengths grow. The v0.6.18 docs warn that this can break CUDA Graph compatibility.^{[16]Reference 16FlashInfer v0.6.18: Prefill APIhttps://github.com/flashinfer-ai/flashinfer/blob/69ff11fc4954396d98326656dc85debd2223f637/flashinfer/prefill.py}

Don't transfer that FA2 walkthrough's lifecycle to every current wrapper. In v0.7.1, BatchMLAPagedAttentionWrapper(backend="auto") resolves SM100 backend selection during plan(), using request facts before preparing an executable. Existing non-FA graph plans allow one plan per wrapper and can reject replanning to preserve captured metadata pointers. Prepare a new wrapper for a different backend and capture. Some monolithic CuTe plans permit bounded query redistribution in place; other backends keep different graph contracts. The selected backend's replay contract must authorize the update, even when allocations stay fixed.^{[17]Reference 17FlashInfer v0.7.1 Planned MLA Wrapperhttps://github.com/flashinfer-ai/flashinfer/blob/0b46115d238d29ea2e92c6b9b2cb077d40fbc352/flashinfer/mla/_batch_mla/_wrapper.py}

Load balancing and split-KV decisions

If one request has 128K cached tokens and seven requests have 128 tokens, assigning only one CTA per request can leave the GPU underfilled or dominated by the long request. Head and query tiling also contribute parallelism, so request count alone doesn't determine utilization. FlashInfer can partition long KV ranges across CTAs, compute partial states, and merge them. The [6, 2, 9] example teaches indexing; it doesn't justify a particular split threshold.

The scheduler must choose a split policy. A tiny chunk raises merge overhead and metadata traffic. A large chunk raises tail latency because one CTA owns too much work. Optimize for enough resident warps to hide memory stalls while preserving tile reuse and register headroom; maximum occupancy can sacrifice both. The current mix of decode lengths, page locality, head dimension, and concurrent streams needs balanced work.

What split-KV changes

The paper describes a load-balanced scheduler designed to cope with dynamic user requests while remaining compatible with CUDA Graphs.^{[7]Reference 7FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving.https://arxiv.org/abs/2501.01005} That claim is paper-era and workload-specific. Measure inter-token latency (ITL), tail latency, and memory traffic on the exact model and GPU rather than copying a chunk threshold from a benchmark.

Backends and JIT compilation

FlashInfer isn't one monolithic CUDA kernel. Its Python and C++ APIs select among implementations such as FlashAttention-2/3, cuDNN, CUTLASS, TensorRT-LLM, and generated kernels. The choice depends on architecture, operation, dtype, layout, and optional features.^{[1]Reference 1FlashInferhttps://github.com/flashinfer-ai/flashinfer}

For one attention call, think in a composition chain: a layout reader follows indirection, a backend microkernel computes tiles and online state, and a merge or output transform finishes the result when work was split. Backend label alone leaves that chain unspecified. Record layout, head grouping, dtype, split policy, and selected implementation together when comparing runs.

The just-in-time (JIT) path matters when a supported specialization isn't covered by a precompiled binary, or when a developer implements a custom attention variant. A template specializes code, compiles it, and caches the resulting module. JIT doesn't make an arbitrary head dimension, cache layout, or unsupported GPU valid by itself. Precompiled flashinfer-cubin and JIT-cache packages reduce first-use latency when their CUDA and architecture match.

As rechecked October 10, the current flashinfer-jit-cache package is a CUDA-specific shim that installs architecture-specific provider wheels. Verify that the installed providers cover your fleet before expecting compilation-free startup.^{[1]Reference 1FlashInferhttps://github.com/flashinfer-ai/flashinfer} The current README distinguishes CUDA 12.9 and 13.0 runtime CI from CUDA 13.4 preview coverage with PyTorch nightly. Treat CUDA 13.4 as preview coverage; validate that toolkit on your GPU before deployment.

The boundary is practical:

JIT can make a feature possible, but it moves compiler compatibility into your release. Pin CUDA, compiler, driver, and cache location. Warm kernels before exposing production traffic. Log which backend was selected, so a silent fallback doesn't look like a model regression.

The repository exposes module status, cache management, and API logging commands. API logging can capture calls and system information for diagnosis, but tensor-dumping modes may write sensitive prompts and outputs to disk.^{[18]Reference 18FlashInfer API Logginghttps://docs.flashinfer.ai/logging.html} Level 10 also wraps CUDA Graph replay: it synchronizes the stream and writes replay snapshots. A latency change with that mode enabled can come from diagnostics rather than attention. Use it for a short correctness investigation, then disable it and repeat the timing run under the recorded logging configuration.

Where FlashInfer sits in a serving stack

FlashInfer owns kernel-level execution. It doesn't decide admission, tenant quotas, request cancellation, model weights, or page allocation. A stack such as vLLM or SGLang owns those policies, then passes shape and indirection metadata into FlashInfer. TensorRT-LLM can supply one backend implementation, but backend availability isn't the same as serving-engine integration.

The allocator boundary is a source of bugs. Reusing pages while a CUDA stream still reads them can corrupt another request. A disagreement between the page table and kv_last_page_len can expose uninitialized values in the last tile. Reusing a plan after page indices change may produce numerically plausible output for the wrong sequence.

A release checklist at the boundary

1. Keep page-index, length, and workspace lifetimes tied to one stream or explicit event.
2. Assert page size, head dimensions, dtype, and device at wrapper entry.
3. Compare incremental decode against full-prefix attention on short randomized cases.
4. Record selected backend, split-KV flag, and workspace size in a trace.
5. Stress request finish, cancellation, and page reuse while CUDA Graph capture is enabled.

Strengths

One vocabulary for many inference phases

The same project covers prefill, decode, append, mixed prefill/decode, MLA, sparse patterns, sampling, quantization, and communication. A serving engine can keep a stable integration surface while the selected kernel changes.

Explicit composability

Ragged and paged layouts, plan/run wrappers, split-KV state, and cascade attention expose the metadata that a scheduler already has. The pieces compose instead of requiring one opaque end-to-end graph.

Customization without abandoning fast paths

JIT templates let an unusual attention variant stay close to the established wrapper contract. Mature backend implementations can handle common shapes while generated code covers new combinations.

Hardware breadth with caveats

The release spans several NVIDIA generations, while support remains operator-specific. v0.6.18 adds Rubin paths and removes SM75 and single-request FA2 kernels from prebuilt JIT-cache wheels: those APIs may still require first-use compilation even with a cache package installed.^{[2]Reference 2FlashInfer v0.6.18 Releasehttps://github.com/flashinfer-ai/flashinfer/releases/tag/v0.6.18} Check the chosen attention backend, dtype, and graph mode on the target GPU.

Weaknesses and failure modes

Metadata complexity moves upward

The engine must now maintain indptr, page indices, last-page lengths, workspaces, events, and plan validity. PagedAttention makes memory efficient, but correctness depends on allocator and kernel agreeing on the mapping.^{[10]Reference 10Efficient Memory Management for Large Language Model Serving with PagedAttention.https://arxiv.org/abs/2309.06180}

JIT is a release dependency

First-use compilation can create startup spikes or fail because a driver, CUDA toolkit, compiler, or architecture is missing. A warm cache can go stale after a version change. Precompiled wheels reduce risk but narrow supported combinations.

Kernel breadth raises test burden

Every backend, dtype, head grouping, page layout, causal mode, and architecture multiplies the test matrix. Run compatibility checks across the target GPU fleet, not only on the development GPU.

GPU specificity limits portability

FlashInfer targets NVIDIA CUDA architectures rather than CPU, AMD, or Apple inference backends. Porting the API idea is possible, but the CUDA kernels, backend dispatch, and cache packages are hardware-specific.

A kernel library isn't a training stack

FlashInfer's main focus is inference kernels and serving operators.^{[3]Reference 3FlashInfer v0.6.18 Sourcehttps://github.com/flashinfer-ai/flashinfer/tree/69ff11fc4954396d98326656dc85debd2223f637} The attention walkthrough here checks forward computation only. Verify any backward path separately before using an operator during training; FlashInfer also requires a separate distributed training framework, optimizer, and checkpoint system.

What the paper actually measured

The FlashInfer paper is an MLSys 2025 paper, posted as arXiv:2501.01005. The evaluation used FlashInfer v0.2, CUDA 12.4, and PyTorch 2.4.0 on A100 40GB SXM and H100 80GB SXM GPUs, with FP16 storage and compute.^{[7]Reference 7FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving.https://arxiv.org/abs/2501.01005} Use the abstract's ranges as evidence that layout-aware kernels can matter; set service-level SLOs from measurements on your workload.

The end-to-end SGLang experiment used ShareGPT plus a synthetic Variable workload with sequence lengths uniformly distributed from 512 to 2,048 tokens. It measured time-to-first-token (TTFT) and ITL under online serving, adjusting request rate to keep P99 TTFT below 200 ms.^{[7]Reference 7FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving.https://arxiv.org/abs/2501.01005} Those workload and admission details belong beside the percentage, because changing them changes the question the benchmark answers.

Read every number with its denominator:

For a production decision, replay your trace with the same prompt length distribution, output length, batch policy, quantization, CUDA graph mode, and error budget. Compare p50 and p99 time-to-first-token, ITL, tokens per second, GPU memory, compile time, and correctness against a reference implementation. Freeze a warm steady-state window, change one variable, rerun the same trace, then verify output and tail metrics. That measure-change-verify loop makes a kernel comparison an experiment rather than a before-and-after anecdote.

Correctness is part of benchmark coverage

A fast path is useful only across shapes it computes correctly. Your harness should compare each selected backend with a trusted reference on the same ragged and paged inputs. Cover prefill, decode, and mixed batches; short histories, non-full last pages, shared prefixes, split and unsplit KV ranges; the MHA, GQA, and MLA head layouts your model uses; and every dtype, mask, device, and fallback path you intend to ship. Include metadata dtypes in that matrix: current FlashInfer docs require indptr, page-index, and tail-length arrays to use int32; an int64 index can fail before numeric comparison.^{[11]Reference 11KV-Cache Layoutshttps://docs.flashinfer.ai/tutorials/kv_layout.html} Compare outputs and attention-state statistics under an explicit tolerance, fail on NaN or Inf, and report unsupported shapes separately. One passing dense case is a smoke test, not coverage.

Diagnose the path before tuning the kernel

When a result regresses, first locate the time. Use Nsight Systems to inspect CPU planning, metadata copies, stream dependencies, kernel launch gaps, and overlap; use Nsight Compute to inspect the selected kernel's memory workload, arithmetic intensity, occupancy limits, and warp stalls.^{[19]Reference 19Nsight Systems User Guidehttps://docs.nvidia.com/nsight-systems/UserGuide/[20]Reference 20Nsight Compute Profiling Guidehttps://docs.nvidia.com/nsight-compute/ProfilingGuide/} Correlate that trace with the recorded backend, layout, split-KV choice, merge work, and p50/p99 TTFT or ITL. A memory-bound trace points first to page access, layout, or precision; a compute-bound trace points to tile or backend choice; a launch gap points to planning or scheduler work. High GPU utilization can still hide a bad tail, a synchronizing plan, or a silent fallback. Tune the boundary that the trace identifies, then rerun the same correctness matrix.

A small code-reading route

Use the local repository as a map rather than trying to read every kernel first:

1. Start with flashinfer/decode.py and find BatchDecodeWithPagedKVCacheWrapper.plan and .run.
2. Follow plan metadata into csrc/batch_decode.cu, where offsets and optional split-KV buffers become kernel parameters.
3. Read flashinfer/cascade.py and include/flashinfer/attention/cascade.cuh to see how attention states compose, including the base-2 merge.
4. Read flashinfer/page.py and csrc/page.cu for page append and slot mapping.
5. Compare a Python wrapper test with its reference-correctness test before reading a specialized backend.

At each layer, write down the shape and lifetime the function assumes, and who owns the next buffer. That check finds more bugs than memorizing kernel names.

What an integration review must defend

Start from the three-request page map, the merge of two partial attention states, and the plan that went stale. Use that batch to explain the boundary before proposing a faster kernel:

• Indexing: Reconstruct both pointer arrays, locate R2 token 5 at page 14 slot 1, and derive its four causal row lengths as 6, 7, 8, and 9.
• Numerics: Match split attention to direct softmax with base-2 LSE, disjoint key ranges, empty-state handling, and an explicit floating-point tolerance.
• Integration: Distinguish refreshing metadata, planning work, compiling a specialization, and capturing a graph. Record which changed before claiming a cache or graph is reusable.
• Measurement: Keep the v0.2 paper results separate from v0.6.18, and report a correctness matrix and serving trace before generalizing a speedup.

Follow-up questions