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.

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.

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.

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.

sequenceDiagram
    participant S as Serving scheduler
    participant P as FlashInfer plan
    participant G as GPU workspace
    participant R as FlashInfer run
    S->>P: request lengths and page indices
    P->>G: CTA queues and reduction metadata
    loop Transformer layers
        S->>R: Q and layer KV cache
        R->>G: read stable plan
        R-->>S: attention output
    end

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.

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.

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.

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.

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.