FlashInfer: A Customizable Attention Engine for LLM Serving

Abstract

Fast attention during training is largely a problem of executing one regular, dense operator efficiently. Attention during LLM serving is less uniform. A serving batch may mix prefill and decode, contain requests with different sequence lengths, read a paged or ragged KV cache, share prefixes, use grouped-query attention, and apply model-specific masks or score transformations. The best GPU schedule changes with these properties and with the target architecture.

FlashInfer addresses this variability as an inference-engine problem rather than as a single kernel. Its original MLSys design has three central components: a block-sparse representation that unifies heterogeneous KV-cache layouts, JIT-specialized attention templates, and a runtime scheduler that balances variable-length work while preserving CUDA Graph compatibility. This article explains these mechanisms, the recursive attention state that connects them, and their relationship to FlashAttention and PagedAttention. It also distinguishes the scope of the 2025 attention-engine paper from the broader kernel library that the FlashInfer project has since become.

1. Introduction

The mathematical interface to attention is compact:

The serving interface is not. During prefill, a request contributes many query tokens and exposes substantial parallelism. During decode, it usually contributes one query token but reads its entire KV cache. Continuous batching combines requests at different positions. Prefix caching, speculative decoding, sliding windows, custom score transforms, and multi-query or grouped-query attention further change which data is read and how work should be partitioned.

A fixed kernel can be excellent at one point in this design space and poor at another. A serving framework could maintain separate kernels for every layout, attention variant, GPU, and request shape, but this produces a growing implementation and tuning burden. FlashInfer’s systems thesis is that these configurations share enough structure to be handled by a common engine, provided that storage, computation, and scheduling are separated at the right boundaries.

The following diagram places the three systems at their respective abstraction layers. vLLM owns the request lifecycle and KV-cache allocation. FlashInfer converts vLLM’s batch and cache metadata into a specialized execution plan. FlashAttention supplies the tiled, exact-attention algorithm used within the resulting GPU kernel.

FlashInfer is therefore best understood as the layer between an LLM serving framework and GPU attention implementations. It has been integrated into systems including SGLang, vLLM, and MLC-Engine; it does not replace their admission control, batching policy, or KV-page allocator.

2. Why Serving Attention Is Heterogeneous

2.1 Prefill and Decode Have Different Shapes

For a query length $L_q$ and KV length $L_{kv}$, prefill commonly has $L_q$ comparable to $L_{kv}$, whereas decode has $L_q=1$ per request. Prefill can expose enough rows to occupy the GPU with conventional query tiling. Decode has little query-side parallelism and is usually dominated by reading the cache. Long decode contexts may need to split the KV dimension across thread blocks and merge partial attention results.

The same nominal operator therefore moves between compute-oriented and bandwidth-oriented regimes. One tile shape and one decomposition cannot be optimal for both.

The gap between the two regimes is large enough to be the dominant fact about serving performance. Measured on one GPU with $H_q = H_{kv} = 8$ and $d_h = 128$:

Both are near their respective ceilings, and the ceilings are different resources. Prefill at 78% of compute peak is doing what tensor cores are for; decode at 75% of bandwidth peak is doing what a memory controller is for. No single tile shape serves both, which is why the engine carries separate kernels rather than one parameterized kernel.

2.2 A Batch Is Ragged

Serving requests do not have a common sequence length. Padding each request to the longest member of the batch wastes both storage and computation. FlashInfer represents packed query, key, and value tensors with CSR-style index pointers. If indptr[i+1] - indptr[i] is the length of request $i$, all tokens can remain contiguous without padding while each request retains a logical slice.

Dynamic KV caches add a second form of irregularity. A request’s logical token sequence may map to non-contiguous physical pages. FlashInfer accepts page indices and per-request index pointers, but leaves allocation, eviction, and page-table policy to the serving framework.

Ragged packing is worth quantifying, because the effect is larger than intuition suggests and its size is set by the worst request rather than the average. With prompt lengths drawn from a lognormal distribution (median 300, tail to 3,000), comparing the ragged kernel against the same batch padded to its longest member:

Two observations follow. The gain is not purely a kernel-efficiency gain — at batch 1, where there is no padding at all, ragged packing is still 1.26× faster — but most of it is, and it scales with the tail. And the tail problem does not diminish with batch size: as batches grow, the probability of containing one very long request approaches one, so the padded fraction stays high and the ragged advantage grows with it.

2.3 Attention Semantics Also Vary

Modern models vary the number of query and KV heads, masking rules, positional transformations, score caps, and even the normalization function. Serving workloads also introduce shared prefixes and tree-shaped speculative states. These are not merely API differences: they alter memory reuse, fusion opportunities, and the shape of efficient GPU work.

The engine must consequently answer three independent questions:

1. Which KV entries belong to each query? This is a data-layout question.
2. Which attention function is applied? This is a code-specialization question.
3. How is the current batch distributed across GPU thread blocks? This is a runtime-scheduling question.

3. Design Overview

FlashInfer assigns one mechanism to each question.

This decomposition is important. Compile-time specialization handles facts that remain stable across many invocations, such as head dimensions, data types, attention semantics, and GPU architecture. Runtime planning handles facts that change as requests arrive and sequences grow.

4. KV Cache as a Block-Sparse Matrix

PagedAttention views KV memory through a logical page table. FlashInfer takes the next step for kernel construction: it interprets the query-to-KV access relation as a block-sparse matrix. A nonzero block means that a query tile attends to a particular KV block. CSR-style index arrays identify the blocks that each row group accesses.

This abstraction can describe several layouts:

• contiguous or ragged KV tensors;
• paged KV caches with arbitrary physical page indices;
• sparse attention over selected cache regions;
• tree-shaped masks used by speculative decoding;
• multiple requests that share a cached prefix.

The block dimensions are not fixed globally. The column block size can follow the page granularity selected by the cache manager, while the row block size can follow the query tile selected by the kernel. FlashInfer’s kernels gather scattered KV rows into contiguous shared-memory tiles, after which dense tensor-core computation can proceed using the same basic FlashAttention pipeline.

Page size is the obvious candidate for a tuning parameter, and it is worth knowing that it is not one. Measured across block sizes of 16, 32, and 64 tokens, at every combination of batch size and context length tested, decode latency varied by less than 2% — within run-to-run noise. The same held for the choice of splitting long KV ranges across thread blocks versus processing them whole.

This is not a negative result about the design; it is a statement about where the design’s cost lives. Page size determines allocation granularity — how much of the last block a request wastes — and that shows up in memory capacity, not in kernel time. The gather cost is proportional to the number of KV entries read, not to how they are grouped, so a kernel reading the same 16,384 tokens at page 16 and page 64 moves the same bytes. Page size is an allocator decision, and tuning it against latency measures nothing.

4.1 Composable Formats

A single sparse format forces one block size on the entire access pattern. That is inefficient when part of the cache is shared and part is request-specific. A large row block improves reuse for a common prefix because several queries can consume one shared-memory load. The same block size wastes space or work on unique suffixes.

FlashInfer instead decomposes the access relation into multiple sparse submatrices. A shared prefix can use larger row blocks, while unique suffixes use finer blocks. The underlying KV data need not move; only the index views change. This turns prefix structure into an execution optimization without requiring a separate cache-management subsystem.

5. JIT-Specialized Compute Templates

FlashInfer builds on the non-materializing attention algorithms described in FlashAttention: Exact Attention Without Materialization. The contribution is not a replacement for online softmax. It is a template system that adapts the same tiled skeleton to serving-specific layouts, hardware, and attention semantics.

The paper’s template exposes transformations on queries, keys, values, logits, masks, and outputs. A configuration can therefore specialize operations such as RoPE, sliding-window masking, logit soft capping, or other score transformations. The JIT compiler inserts these operations into CUDA/CUTLASS templates and compiles a kernel for the requested data types, dimensions, tile sizes, and target architecture.

This design avoids two unsatisfactory extremes. A fully generic runtime kernel pays branches and indirect dispatch in its hot loop. A collection of handwritten kernels duplicates an increasingly large amount of scheduling and data-movement code. JIT specialization retains a common implementation skeleton while resolving stable choices before execution.

FlashInfer also selects among tile shapes. Short-query decode may use a query tile of one and CUDA cores, whereas larger query tiles can use tensor cores. Hardware-specific backends apply FlashAttention-2-style kernels through Ada and FlashAttention-3-style kernels on Hopper in the system evaluated by the paper.

6. Dynamism-Aware Scheduling

JIT compilation cannot specialize on exact sequence lengths when those lengths change at every generation step. FlashInfer therefore separates planning from execution using an inspector-executor model.

The plan phase receives the current query and KV lengths. It maps work tiles to cooperative thread arrays, splits sufficiently long KV sequences when useful, and records how partial outputs must be reduced. The run phase launches a persistent attention kernel using that cached metadata. A plan can be reused across all Transformer layers whose requests have the same sequence-length structure, amortizing its CPU cost.

The scheduler is load-aware rather than request-count-aware. A batch containing one very long context and several short ones should not assign one equal unit of work to each request. Long KV ranges can be partitioned across thread blocks; short requests can write directly to final outputs. FlashInfer uses deterministic reduction ordering rather than atomic accumulation so that identical input length information produces a stable aggregation order.

6.1 CUDA Graph Compatibility

CUDA Graph replay reduces CPU launch overhead, but captured graphs require stable kernel launch configurations and memory addresses. Dynamic schedules appear to violate that requirement. FlashInfer resolves the tension by keeping the kernel grid and workspace layout fixed while changing the scheduling metadata stored inside preallocated buffers.

The CPU-side plan operation is outside graph capture. The GPU-side run operation uses persistent kernels, stable pointers, and fixed workspace offsets, so it can be captured and replayed. In effect, dynamic work is represented as data consumed by a static launch structure.

The magnitudes involved are easy to underestimate, and they must never be folded into a throughput number. Measured at a 4,096-token context:

The first call for a new (dtype, head dimension, page size) combination costs 36 milliseconds — roughly a thousand times a steady-state decode step at batch 1 — because it triggers JIT specialization and kernel loading. That is a startup cost, and reporting it as latency would be a category error; reporting throughput that excludes it would be equally misleading if the service is short-lived.

The graph column shows precisely where capture helps: it removes host launch overhead, so it wins exactly when the host is the limit. At batch 1 the host path costs 0.026 ms against 0.008 ms of GPU work, and capture nearly halves the step. At batch 32 the GPU is saturated at 0.096 ms, the host is already hidden behind it, and capture changes nothing. CUDA graphs are not a general speedup; they are a fix for a specific bottleneck, and they are worth exactly the size of the launch path they remove.

7. Attention States and Recursive Composition

When a KV sequence is divided into chunks, each chunk computes only a partial softmax. Partial output vectors cannot be added directly because each was normalized over a different set of keys. FlashInfer makes the required normalization information explicit as an attention state.

For a query $q$ and KV index set $I$, define

and

The state is the pair $(v(I),s(I))$: the normalized output and its log-sum-exp scale. For disjoint sets $I$ and $J$, let

Their exact merged output is

This merge is associative and commutative, subject to ordinary floating-point effects. It allows the runtime to split a long KV range across thread blocks and contract the resulting states. It also enables cascade attention: compute attention over a shared prefix once, compute each request’s unique suffix separately, and merge the states without approximating the result.

This is the same mathematical property that underlies tiled online softmax and Flash-Decoding. FlashInfer promotes it from an internal kernel recurrence to an engine-level interface for scheduling and composition.

8. A Concrete Serving Workflow

A typical paged decode path follows four steps:

workspace = allocate_workspace()
attention = BatchDecodeWithPagedKVCacheWrapper(workspace, kv_layout="NHD")

attention.plan(
    page_indptr,
    page_indices,
    last_page_length,
    num_query_heads,
    num_kv_heads,
    head_dim,
    page_size,
)

for layer in model.layers:
    output = attention.run(layer.query, layer.paged_kv_cache)

First, the serving framework allocates and updates physical KV pages. Second, it supplies the logical page-table metadata and current dimensions to plan. Third, FlashInfer constructs auxiliary scheduling and reduction data in caller-owned workspace. Finally, each layer calls run with its own KV tensors while reusing the same plan.

Prefill follows the same pattern but supplies query index pointers because each request may contribute a different number of query tokens. Ragged-cache wrappers compact contiguous variable-length data; paged-cache wrappers consume page indices. The common plan/run boundary allows the engine to preserve optimized kernels without forcing one cache policy on every serving framework.

9. Relationship to FlashAttention and PagedAttention

These three names solve different problems:

FlashInfer uses FlashAttention algorithms inside its dense and sparse kernel templates. It consumes paged KV layouts introduced by PagedAttention, but it does not decide which requests receive pages or when pages are evicted. Its role is to translate the serving framework’s current layout and workload into efficient GPU computation.

That division leaves one question open, and it is the one that determines whether any of these kernels matter: who chooses the requests that make up the current batch? FlashInfer executes the schedule it is handed. Producing that schedule — admitting, evicting, and interleaving prefill with decode — is the subject of Continuous Batching.

10. Evaluation Evidence

The MLSys 2025 paper evaluates FlashInfer v0.2 on NVIDIA A100 and H100 GPUs using CUDA 12.4, PyTorch 2.4, and FP16 storage and computation. The reported results should be read as measurements of those systems and workloads, not as architecture-independent guarantees.

In SGLang serving experiments with Llama 3.1 8B and 70B, the paper reports a 29–69% reduction in inter-token latency relative to its Triton backend under the evaluated workloads. For a Streaming-LLM implementation, fusing RoPE transformations into generated attention kernels reduced latency by 28–30%. In MLC-Engine parallel-generation experiments, composable shared-prefix formats reduced inter-token latency by 13–17% at the reported peak configurations.

The evaluation also exposes boundaries. Sparse gathering was within 1% of dense-cache performance for the tested decode kernels, but the paper observed roughly a 10% gap for prefill because Hopper’s TMA could not directly express the fine-grained non-affine accesses. Shared-prefix composition helped most with long prefixes and moderate-to-large parallel batches; its advantage diminished when prefixes were short or attention no longer dominated execution.

10.1 The decode scaling law, measured

One property is worth isolating because it explains most of the variance in decode measurements, and it is not a kernel property at all. Paged decode against a dense contiguous baseline, same KV bytes read, $H_q = H_{kv} = 8$, $d_h = 128$, BF16:

At a 16K context, per-request cost falls from 584 µs (dense, batch 1) to 11 µs (paged, batch 32) — a factor of 51. The dense column is nearly flat across batch because a single query token per request cannot occupy the machine regardless of how many requests are in flight, and its per-request cost therefore never amortizes.

The mechanism is the one from the cache sizing article: decode’s arithmetic intensity is fixed at 1 FLOP/byte, so the only available lever is how much of HBM the batch can keep busy. Measured achieved bandwidth rises from 1,838 GB/s at batch 1 to 6,013 GB/s at batch 32 — 23% to 75% of peak — while intensity never changes. A decode benchmark at batch 1 measures a mostly idle memory system, which is why single-request numbers are so often pessimistic and why they improve so much under production batching.

11. Scope and Evolution

The 2025 paper presents FlashInfer primarily as a customizable attention engine. The current project is broader. Its documentation includes GEMM and grouped GEMM, fused MoE, sampling, communication, normalization, RoPE, quantization, page operations, MLA, and additional model-specific kernels. That expansion is consistent with the project’s role as an inference-kernel library, but those later modules are not evidence for the original paper’s three-part attention design.

This distinction matters when evaluating the system. The enduring idea is not a particular Python wrapper or list of supported operators. It is the separation of stable kernel specialization from dynamic serving schedules, connected by explicit representations for irregular KV access and mergeable attention state.

12. Limitations

FlashInfer does not remove serving complexity; it places a kernel-engine boundary around part of it. The serving framework still owns request scheduling, page allocation, cache eviction, and distributed placement. Planning has CPU and metadata-copy overhead, although plans can be reused across layers. CUDA Graph compatibility requires preallocated workspace and stable upper bounds. Fine-grained sparse gathering may prevent use of hardware data-movement engines designed for affine access.

The paper’s implementation supported only the attention forward pass, making its original design an inference solution rather than a training replacement. JIT specialization also trades runtime generality for compilation and caching complexity. Finally, no backend is uniformly optimal: performance depends on sequence distribution, data type, GPU generation, cache layout, and integration overhead in the host serving framework.

Two of those boundaries are sharper in practice than the list implies. JIT specialization means the set of supported shapes is finite and version-specific: on the build measured here, both the prefill and decode kernels failed to dispatch at a head dimension of 32, and no fused path accepted FP32 storage. Neither is a defect — specializing for the dimensions production models use is the point — but it means an engine’s shape coverage has to be verified rather than assumed, and a model with an unusual head dimension may fall back to a slower path or to no path at all.

13. Conclusion

FlashInfer’s central contribution is to treat attention as a cross-layer serving problem. Block-sparse and composable formats describe where KV data resides and which regions are shared. JIT templates specialize what attention means and how the target GPU executes it. Runtime planning decides how a changing batch occupies the machine while a static execution structure preserves CUDA Graph replay.

The resulting system complements rather than supersedes FlashAttention and PagedAttention. FlashAttention supplies the non-materializing algorithm, PagedAttention supplies a memory-management abstraction, and FlashInfer turns those ideas into a customizable execution engine for heterogeneous inference workloads.

References

1. Zihao Ye et al. FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving. Proceedings of Machine Learning and Systems, 2025.
2. Tri Dao et al. FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. NeurIPS, 2022.
3. Woosuk Kwon et al. Efficient Memory Management for Large Language Model Serving with PagedAttention. SOSP, 2023.
4. FlashInfer contributors. FlashInfer Documentation.