FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving

Training mostly runs dense attention on fixed-size batches. Serving doesn't. The KV cache lives in pages (vLLM) or a radix tree (SGLang); a batch mixes new prompts and single-token decodes; many requests share a system prompt; speculative decoding verifies trees of guesses; and models keep adding attention variants (sliding windows, logit soft-capping, fused rotary embeddings). Each engine had been writing its own kernels for its own subset.

FlashInfer is an attention library built for that mix, with three ideas.

One format for every KV layout. A page table or a prefix tree is a block-sparse row (BSR) matrix: rows are queries, columns are KV blocks, and a non-zero block means "this query attends to this page". FlashInfer's kernels accept any block size, gathering scattered pages into contiguous shared memory before using tensor cores. Composable formats split one cache into two views: a shared prefix stored with a large row block (many queries together), so its KV is loaded once into fast on-chip memory for the whole group; and each request's own tokens with single-row blocks. No data moves; only the index arrays differ.

Customizable attention, compiled just in time. A variant is specified as small functions (query and key transforms such as RoPE, a logits transform such as soft-capping, a mask), and FlashInfer generates CUDA from FlashAttention-2 or -3 templates, picking a query tile size (1 to 128) to match the batch's average query length.

Load balancing that survives CUDA graphs. Each step, a CPU-side plan splits long KV sequences into equal chunks across the GPU's processors and records where partial results go; persistent kernels with fixed grid sizes then run inside a captured graph and merge partial results with the log-sum-exp composition.

Integrated into SGLang, vLLM and MLC-Engine, it cut SGLang's median inter-token latency 29-69% against Triton kernels, StreamingLLM latency 28-30% with a fused kernel, and parallel-generation inter-token latency 13.7-17.4% with composable formats.

Most engineers never call FlashInfer directly; it sits under the serving engine. It still shows up in deployments in three ways.

It's a setting. SGLang and vLLM let you choose the attention backend. On the same model and GPU, the choice changed inter-token latency by up to 69% in the paper's SGLang tests, so check which backend a slow deployment uses before tuning anything else.

New model features need kernel support. A model with soft-capped logits, an unusual mask or a new positional scheme runs slowly, or not at all, until the attention kernels support it. JIT templates are what let engines support new variants quickly.

Batch shape matters. Mixed lengths, shared prompts and parallel sampling all change which kernel path is fastest. If a customer's traffic is highly skewed (a few very long contexts among many short ones), load-balanced kernels are what keep the GPU busy.