FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving

Zihao Ye; Lequn Chen; Ruihang Lai; Wuwei Lin; Yineng Zhang; Stephanie Wang; Tianqi Chen; Baris Kasikci; Vinod Grover; Arvind Krishnamurthy; Luis Ceze

Source-linked AI summary

FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving

Zihao Ye, Lequn Chen, Ruihang Lai, Wuwei Lin, Yineng Zhang, Stephanie Wang, Tianqi Chen, Baris Kasikci, Vinod Grover, Arvind Krishnamurthy, Luis Ceze

arXiv:2501.01005v2Jan 2, 2025cs.DCcs.AIcs.LG

TL;DR

Efficient LLM serving requires attention kernels that handle diverse workloads, heterogeneous KV-cache layouts, and varied attention mechanisms. FlashInfer combines unified sparse storage, JIT-customizable attention templates, and dynamic scheduling, achieving reported latency and speed improvements across serving scenarios. Its current scope is limited to forward-pass attention, with backward templates needed for training.

Problem
Diverse workloads, KV-cache configurations, attention variants, and hardware characteristics complicate development of a comprehensive efficient attention solution.
Method
FlashInfer combines block-sparse and composable KV-cache formats, JIT-customizable attention templates, and dynamic load-balanced scheduling for LLM serving.
Results
29-69% inter-token-latency reduction compared to Triton backend for LLM serving benchmark, 28-30% latency reduction for long-context inference, and 13-17% speedup for LLM serving with parallel generation.
Takeaways & Limitations
FlashInfer provides a unified attention engine evaluated across standard and specialized LLM serving scenarios and integrated into vLLM, MLC-Engine, and SGLang.
Takeaways & Limitations
FlashInfer currently supports only the forward pass; applying it to training requires customizable backward attention kernel templates.

Abstract

from arXiv · show

Transformers, driven by attention mechanisms, form the foundation of large language models (LLMs). As these models scale up, efficient GPU attention kernels become essential for high-throughput and low-latency inference. Diverse LLM applications demand flexible and high-performance attention solutions. We present FlashInfer: a customizable and efficient attention engine for LLM serving. FlashInfer tackles KV-cache storage heterogeneity using block-sparse format and composable formats to optimize memory access and reduce redundancy. It also offers a customizable attention template, enabling adaptation to various settings through Just-In-Time (JIT) compilation. Additionally, FlashInfer's load-balanced scheduling algorithm adjusts to dynamism of user requests while maintaining compatibility with CUDAGraph which requires static configuration. FlashInfer have been integrated into leading LLM serving frameworks like SGLang, vLLM and MLC-Engine. Comprehensive kernel-level and end-to-end evaluations demonstrate FlashInfer's ability to significantly boost kernel performance across diverse inference scenarios: compared to state-of-the-art LLM serving solutions, FlashInfer achieve 29-69% inter-token-latency reduction compared to compiler backends for LLM serving benchmark, 28-30% latency reduction for long-context inference, and 13-17% speedup for LLM serving with parallel generation.

1 INTRODUCTION

FlashInfer addresses the combined challenges of diverse LLM workloads, heterogeneous KV-cache layouts, varied attention mechanisms, and GPU-specific optimization. It combines unified block-sparse storage, JIT-customizable attention templates, and dynamic load-balanced scheduling, with reported gains across serving scenarios.

Challenges: Diverse workloads create varying attention patterns, query lengths, and KV-cache lengths, making naive implementations vulnerable to load imbalance.The cited workloads include prefill, batched decoding, prefix reuse, and speculative tree decoding.
Challenges: Hardware-efficient attention must accommodate heterogeneous KV-cache storage, GPU-specific pipelines, and increasingly varied attention mechanisms.Examples include paged attention, radix trees, grouped attention heads, specialized masks, and customized score computations.
Design: FlashInfer uses a unified block-sparse KV-cache format with adjustable block sizes to represent diverse configurations and improve memory access efficiency.The format supports fine-grained sparsity, including vector-level sparsity.
Design: A customizable attention template lets users implement attention variants that FlashInfer converts into optimized block-sparse implementations through JIT compilation.This design supports rapid adaptation to varying attention configurations.
Design: FlashInfer separates compile-time tile-size selection from runtime scheduling, enabling load-balanced adaptation to changing KV-cache lengths while remaining compatible with CUDA Graphs.The runtime scheduler uses lightweight APIs while constant configurations are preserved for CUDA-Graph compatibility.
Evaluation: FlashInfer was integrated into vLLM, MLC-Engine, and SGLang, and evaluated for standard serving, prefix sharing, speculative decoding, long-context inference, and parallel generation.The evaluation assessed end-to-end latency and throughput improvements across these environments and scenarios.

2 BACKGROUND

FlashAttention and attention composition provide the computational and reduction foundations relevant to efficient attention serving, while block-sparse formats improve hardware utilization for structured sparsity. These techniques motivate FlashInfer’s unified storage and composable execution design.

FlashAttention: FlashAttention computes exact attention with reduced memory use by applying the online-softmax trick without materializing the attention matrix in global memory.FlashAttention2 and FlashAttention3 further optimize loop ordering and pipelines for Ampere and Hopper GPUs.
FlashAttention: In LLM serving, attention operational intensity simplifies to O(lqo) because query length is no greater than KV-cache length, and batching does not change this intensity.MQA and GQA increase operational intensity to O(g · lqo) while reducing KV-cache size.
Attention Composition: Attention composition preserves both partial attention outputs and their scales so results over separate key-value sets can be combined.BPT motivates this representation for composing attention states.
Attention Composition: For an index set I, the attention scale is defined by log-sum-exp over attention scores, and the corresponding output combines value vectors weighted by normalized scores.The prose identifies ki as the i-th key vector.
Attention Composition: Attention states for I ∪ J are composed with an associative and commutative operator, allowing partial-attention computations to be reduced in any order.FlashInfer adopts the Attention State as its canonical attention output and ⊕ as its standard reduction operator.
Block/Vector Sparsity: BSR groups non-zero elements into contiguous blocks, improving register reuse and compatibility with GPU matrix-multiplication units while allowing empty blocks to be skipped.Traditional block-sparse kernels often use block sizes that are multiples of (16, 16) or larger, which can be suboptimal for fine-grained sparsity.
Block/Vector Sparsity: Smaller block sizes can still use tensor cores after gathering rows or columns into contiguous shared memory, supporting fine-grained sparsity patterns.FlashInfer extends this approach to blocks with arbitrary column sizes Bc.

3 DESIGN

FlashInfer uses unified and composable block-sparse representations to handle heterogeneous KV-cache layouts, while customizable JIT-compiled templates support diverse attention variants and hardware settings.

3.1 Data Representation: FlashInfer stores KV caches in BSR format with application-defined block sizes (Br, Bc), supporting arbitrary values for heterogeneous layouts.Br corresponds to the query tile size, while Bc is determined by KV-cache management algorithms.
3.1 Data Representation: Shared-prefix blocks let multiple queries reuse KV-cache data in high-bandwidth shared memory, whereas unique blocks may require global-memory or L2-cache accesses.The example uses block sizes (3, 1) for shared prefixes and (1, 1) for unique KV cache.
3.1 Data Representation: Composable formats decompose KV caches into multiple block-sparse matrices, allowing shared-prefix regions to use larger blocks while unique regions use smaller blocks.This improves reuse without moving KV-cache data by computing indices and index-pointer arrays for the submatrices.
3.2 Kernel Design: Customizable attention templates use JIT compilation to insert user-defined attention variants into CUDA templates and generate optimized implementations.The templates support dense and block-sparse matrices across NVIDIA architectures from Turing through Hopper, with FA2 or FA3 selected by architecture.
3.2 Kernel Design: FlashInfer transfers scattered sparse KV-cache tiles into contiguous shared memory so sparse and dense attention implementations can use consistent downstream kernels.The last KV-cache dimension remains contiguous for coalesced access, while TMA is used only for contiguous Hopper layouts.
3.2 Kernel Design: FlashInfer selects tile sizes using workload and hardware heuristics, including average query length and register/shared-memory constraints.Supported FA2 query tile sizes are 1, 16, 32, 64, and 128, with corresponding K/V tile sizes of 32, 64, or 128.

4 EVALUATION

FlashInfer is evaluated at kernel and end-to-end levels across standard serving, long-context inference, and parallel generation. The results show consistent gains, with benefits varying by workload and parallelism.

4.1 End-to-end LLM serving performance: 29-69% inter-token-latency reduction compared to Triton backend is reported for the LLM serving benchmark.Experiments use SGLang with FlashInfer or Triton under latency-sensitive online serving settings.
4.1 End-to-end LLM serving performance: FlashInfer shows consistent speedup in ITL and TTFT across evaluated SGLang settings for Llama 3.1 8B and 70B models.The comparison uses one H100 for 8B and four H100s for 70B.
4.2 Kernel Performance for Input Dynamism: FlashInfer kernels significantly outperform FlashAttention kernels for uniform and skewed sequence-length distributions.The evaluation measures achieved bandwidth and FLOPs utilization for decode and causal prefill kernels with batch size 16.
4.3 Customizability for Long-Context Inference: 28-30% latency reduction is achieved for Streaming-LLM with FlashInfer fused kernels across different recent-window settings.The fused RoPE-attention kernel also achieves 1.6-3.7x higher bandwidth utilization than the unfused combination.
4.4 Parallel-Generation Performance: 13-17% speedup is achieved for parallel generation with composable formats at moderate parallelism.At n = 4, ITL decreases by 13.73% for 8B and 17.42% for 70B; benefits plateau at larger n.

5 RELATED WORK

Related systems optimize attention through dense, sparse, fused, compiler-generated, or serving-specific techniques. FlashInfer combines these directions into a unified, customizable CUDA-based attention engine for heterogeneous LLM serving workloads.

Dense and sparse attention kernels: FlashAttention reduces shared-memory requirements for long contexts, while FlashAttention2 and FlashAttention3 further optimize loop structure and overlap computation.FlashInfer extends these templates to support sparse attention kernels.
Prefix sharing and serving: Several systems address shared-prefix decoding but require separate KV-cache management, whereas FlashInfer uses unified page-table management for composable formats.The cited comparison includes RelayAttention, Hydragen, ChunkAttention, and Parrot.
Sparse computation: Prior sparse-fusion and graph-attention methods generally omit softmax, serialize aggregation, or do not explore Tensor Cores for attention.FlashInfer targets sparse attention kernels with hardware-oriented execution.
Programmable attention variants: FlexAttention provides a user-facing interface compiled into Triton block-sparse FlashAttention, while FlashInfer expands customization to query/key transformations and load balancing.FlashInfer generates CUDA code and can serve as a forward-pass backend for FlexAttention.
LLM serving systems: FlashInfer provides a unified block-sparse solution across continuous batching, paged KV-cache management, chunked prefill, and radix-tree prefix caching.The related systems include Orca, PagedAttention, Sarathiserve, and SGLang.

6 DISCUSSIONS

FlashInfer currently targets forward-pass attention, while its decoupled templates and runtime scheduler support broad attention variants and multiple GPU-specific scheduling strategies.

Scope and limitation: FlashInfer currently supports only the forward pass for attention computation.Extending it to training would require customizable backward attention kernel templates.
Generalization: FlashInfer’s computation-scheduling separation allows diverse tiling strategies and runtime scheduling for load balancing and wave-quantization reduction.The design targets architectures from Turing through Hopper with architecture-specific FlashAttention variants.
Attention coverage: The template design space covers most attention functions, including Multi-head Latent Attention and the intra-attention component of Linear Attention.The supported form is expressed through composable query, key, logits, value, scan, and epilogue functions.

7 CONCLUSION AND FUTURE WORK

FlashInfer concludes with a unified attention engine for efficient LLM serving and identifies deployment and extension directions. Its design also targets short-query GQA through head-group fusion.

Conclusion: FlashInfer combines unified block-sparse storage, composable formats, JIT customization, and load-balanced scheduling for LLM serving.The project is open source and has been deployed at scale in production-level systems.
Future work: Future work includes compiling higher-level DSLs to attention specifications and generating code for additional backends.The conclusion also identifies customizable backward templates as necessary for extending FlashInfer to training.
GQA optimization: Head-group fusion maps KV heads to threadblocks while fusing query heads with the query-length dimension in GQA.The strategy is intended primarily for short query lengths, where straightforward mapping underuses KV-cache reuse.

B OVERHEAD OF SPARSE GATHERING

FlashInfer evaluates sparse-gathering overhead across prefill and decode attention kernels, while combining adaptable data formats, scheduling, precision, and backend choices for diverse serving conditions.

Sparse-gathering overhead: Figure 12 reports causal prefill TFLOPs/s and decode bandwidth utilization across batch sizes and sequence lengths for dense and sparse KV-caches.Sparse KV-cache uses PageAttention with page size 1, representing vector-level sparsity.
Sparse-gathering overhead: Sparse and dense KV-cache decode performance differs by at most 1%, while prefill shows approximately a 10% performance gap.The evaluation fixes 32 query heads, 32 KV heads, and head dimension 128 while varying batch size and sequence length.
Sparse-gathering overhead: Sparse gathering on FA3 uses asynchronous copies and manual pointer arithmetic because Hopper TMA supports fixed-stride rather than arbitrary-row accesses.The workaround consumes more registers and requires smaller KV-tile sizes.
Additional optimizations: Mixed-precision kernels keep queries and outputs in fp16 while storing KV-cache in fp8 to reduce memory footprint and increase bandwidth utilization.FlashInfer evaluates performance and robustness across additional experimental conditions and attention variants.

G.2 Evaluation of Shared-Prefix Attention Kernels

Shared-prefix attention kernels benefit from composable formats under long-prefix and large-batch conditions, but kernel speedups do not always translate proportionally to end-to-end gains.

Shared-prefix attention: Composable formats benefit shared-prefix attention with long prefixes, such as 32k, and large batch sizes, such as 64.Table 5 reports kernel latency in microseconds across prefix lengths, scenarios, and batch sizes; “composable” and “single” denote the compared formats.
Shared-prefix attention: End-to-end gains can be smaller than kernel speedups because real-world shared-prefix sizes tend to be smaller.This bounds how directly benchmarked shared-prefix kernel improvements transfer to deployed workloads.

G.3 Ablation Study on Variable Sequence Length and load-balancing scheduler

The ablation evaluates FlashInfer’s load-balancing scheduler under variable sequence lengths using SGLang and Llama 3.1-8B-Instruct on an NVIDIA H100 SXM5 GPU.

Scheduler ablation: The ablation compares SGLang with FlashInfer with and without load balancing using inter-token latency and time-to-first-token.It includes ShareGPT and variable-sequence-length datasets, with inputs sampled from U(512, 2048) and outputs fixed at 256.

G.4 vLLM Integration Evaluation

The vLLM integration evaluation compares FlashInfer’s backend with vLLM’s default backend at a fixed request rate, exposing both performance gains and host-side integration overhead.

vLLM integration: FlashInfer reduces vLLM inter-token latency by around 13% with fp8 KV-cache.The comparison reports throughput, inter-token latency, and time-to-first-token at a fixed request rate of 16.
vLLM integration: bf16 shows minor regressions because array operations in vLLM integration add Python-side host overhead.The authors identify moving scheduler work to C++ and the device as future optimization directions.

G.5 Fine-Grained Block-Sparsity Evaluation

FlashInfer’s fine-grained block-sparse attention is evaluated for long-context KV-cache sparsity, with performance compared against PyTorch SDPA and FlexAttention. The design combines small-block sparsity with sparse-row gathering to retain efficient tensor-core use.

FlashInfer supports fine-grained block-sparse matrices used by KV-cache pruning algorithms.
The evaluation measures Quest batch-decoding attention on an NVIDIA H100 SXM5 GPU against PyTorch SDPA and FlexAttention.The configuration uses block size 16, 32 query heads, 32 key-value heads, and head dimension 128.
Up to 20x speedup is achieved for long sequence lengths in the fine-grained sparsity evaluation.FlashInfer is compared with the latency reported for PyTorch SDPA and FlexAttention across Tables 9–11.
FlashInfer uses sparse-row gathering to leverage dense tensor cores with small block sizes, unlike FlexAttention’s large-block templates.This design choice supports fine-grained KV-cache pruning.