Have you ever imagined an AI system that improves itself?

—Not exactly the machine overlords from sci-fi, but still, an AI system that can expand its capabilities using AI itself sounds pretty cool.

—And it is something we’re building: FlashInfer-Bench—a benchmark and infrastructure that opens the pathway for AI to accelerate real-world AI deployment.

AI agents have become remarkably powerful, capable of writing complex code and even building sophisticated systems. Such powerful capabilities naturally lead us to wonder: can AI agents optimize the very production systems they run on? At the heart of these AI systems, the most intensive parts are GPU kernels—the low-level programs that perform the core operations for AI models. We have seen amazing progress recently that shows LLMs can generate reasonable GPU kernels (https://scalingintelligence.stanford.edu/blogs/kernelbench/).

That prompts us to ask the next natural question: how can we systematically get AI agents to improve the very AI system they depend on? We know there can still be hurdles towards this ultimate dream, but it is time to get ready to build a clear pathway for the future. We build FlashInfer-Bench, a benchmark of real-world AI-system driven GPU workloads, and more importantly, an infrastructure and workflow to 0-day ship AI-generated kernels into production. While the AI today may not yet perform at the expert engineer level, we would like to get our infrastructure ready and open the pathway to accelerate this ultimate goal. The same workflow would also benefit machine learning system engineering in general, helping us to build and evolve next-generation kernel libraries.

Overview: Systematically Approaching AI for AI Systems

As we approach the goal for “AI for AI systems”, we naturally find that we need more than just benchmarks. Real-world production environments are highly complex: they involve a large number of sophisticated kernels with different API designs and input signatures, while each kernel’s performance can also vary depending on the input it receives. We need to clearly communicate these needs to AI agents and relevant engineers; more importantly, we also need to be able to quickly adopt the solutions from agents back onto production. Fundamentally, we need to consider the following three elements to systematically approach our problem:

1. Clarify the landscape. We need a standardized way to describe the workloads, including their inputs and outputs, along with relevant statistics and characteristics (e.g., shape, raggedness, etc.)
2. Set the right goal. We need benchmarks that capture real-world, production-grade workloads—in other words, the scenarios that occur when LLMs are deployed in practice.
3. Establish a day-0 production path. AI-generated and human-written code can be seamlessly deployed into a real LLM engine right away with as little intervention as possible.

FlashInfer-Bench is designed specifically to address these three challenges. As illustrated in the architecture diagram above, its core components work as follows:

• Clarify the landscape — FlashInfer Trace. To formalize the description of GPU workloads, we developed FlashInfer Trace: a systematic and evolving open schema that captures the kernel definitions, implementation, and evaluation in LLM inference scenarios. This format addresses the critical needs through four key components: kernel definition that defines the kernel signature and computation; workload that describes the kernel’s real inputs; solution to the kernel; and the solution’s evaluation results.

• Set the right goal — FlashInfer-Bench Dataset. We collected and curated the most important kernels used in production, along with their real workloads, to create a dataset. This allows us to measure kernel performance under real traffic conditions and further makes it possible to concretely improve production systems.

• Establish a 0-day production path — First-class FlashInfer Integration. We build first-class integration with FlashInfer — an open LLM kernel library widely used in major LLM inference engines. Our solution can dynamically replace the FlashInfer Kernels with the best-performing kernels, as evaluated using the FlashInfer Trace and the FlashInfer-Bench dataset. This enables activating the best kernels in LLM engines and testing end-to-end performance with minimal effort.

In the rest of the post, we will walk through the details about the three main elements.

FlashInfer Trace: Standardizing and Evolving Schema for LLM Kernels

As discussed in the overview section, we need a standard way for humans, infrastructure workflows, and AI agents to communicate about the context throughout the overall production life cycle. We specifically would need to support different stages in the production lifecycle, including handling the problem to the agent/engineers, recording the answer, evaluating, and inspecting the outputs.

We define FlashInfer Trace, a set of JSON schemas to formally describe a kernel’s definition, benchmark workload, solution, and evaluation results. Below is the figure showing FlashInfer Trace’s components:

Definition provides the complete specification of the kernel, including the kernel metadata, input and output specifications, axes of input and output tensors, and a reference implementation written in Python, defining the kernel’s computational behavior. The kernel definition provides all the information needed for AI to generate kernels. This is an example of a General Matmul (GEMM) kernel definition:

In the kernel definition, we aim to capture as many details of a kernel as possible, especially the constant axes value. This is because kernels with different input shapes often require different optimal implementations. For example, a GEMM with a dimension of 4096 may have a very different implementation from one with a dimension of 128. Some axes in AI systems are dynamic, such as the sequence length. These dimensions will be kept variable so that the workload can provide its concrete values.

We also design the op_type that specifies input and output and computation at a high level, such as GEMM (as gemm) and Paged Grouped-Query Attention (as gqa_paged). By defining op_type, we can group together multiple kernel definitions that have different axes but share similar computation. We define op_types for common LLM operators, and we aim to evolve this standard with the community.

Workload captures the kernel’s inputs observed in real traffic, including both tensor inputs and scalar or flag inputs. If the specific values of a tensor do not affect evaluation, it can be stored as random to save space; otherwise, the original tensor can be dumped to reproduce the most realistic production scenario.

Solution gives the concrete implementation of a kernel. It must perform the same computation as the reference in the kernel definition. We support multiple kernel programming languages, including CUDA, Triton, Python, and TVM (coming soon!), among others. We can also specify the kernel’s author, target hardware, compatible software versions, and other related information.

Evaluation is the benchmark result for a specified solution on a certain definition and workload. The hardware and software used for evaluation, as well as the kernel’s correctness and runtime performance, are fully measured. When a kernel fails to compile or has a runtime error, we assign different evaluation statuses accordingly.

With the standardized design of the FlashInfer Trace, we can easily exchange information throughout the production lifecycle:

• Workload tracer can take traffic from an LLM inference engine and produce workload traces;
• AI agents can then take the definition to further generate a kernel solution;
• Benchmark tools can take the workload and solution and fill in the evaluation field;
• The leaderboard will take a collection of traces and visualize the result.

We built the flashinfer-bench Python package following this philosophy, which contains support to work with FlashInfer trace and tooling across various stages, including schema validation, benchmarking, workload tracing, and integration with LLM engine and FlashInfer. The design of the FlashInfer Trace is independent of specific tools, and we encourage the community to develop a more versatile toolchain.

Real CUDA engineers can also benefit from the FlashInfer Trace. It standardizes kernel interfaces and computations, reducing communication overhead and enabling fair comparison across kernels from different sources. It provides testing tools to measure the performance of kernels and tracing tools to easily capture real workloads. Altogether, these greatly simplify the burden of kernel development.

We provide standardized and clear documents for FlashInfer Trace and Op Types.

Dataset & Benchmark — The Yardstick from the Real World

We curate the dataset by one principle: real-world relevance. This means focusing on the most important kernels in the LLM engine and recording their workloads under realistic production traffic. We selected the most popular models, including Llama 3, DeepSeek V3, and Qwen 3, and recorded their main kernels, including attention, GEMM, MoE, normalization, sampling, and more. We also strive to ensure realism in the LLM engine configuration. For example, for DeepSeek-V3, we use the original FP8 quantization and enable tensor parallelism = 8 and expert parallelism = 8. We feed inputs from real datasets, including ShareGPT and more. These ensure that the workloads we measure are as close to real-world scenarios as possible.

Day-Zero Integration Through FlashInfer

The final step of the virtuous cycle is getting AI-generated kernels into production immediately. We achieve this by first-class integration with FlashInfer, a well-adopted kernel library by major LLM engines, and SGLang. Integration with vLLM is also progressing rapidly and is expected to be completed soon.

We can dynamically replace the kernels in the FlashInfer API with the best-performing ones from our evaluations, all with minimal effort. By simply importing flashinfer_bench in the LLM engine and enabling the environment variable FIB_ENABLE_APPLY, the kernels can be automatically replaced with the best ones from the local database.

Under the hood, the dynamic substitution is supported by the flashinfer_bench.apply() decorator:

@apply(lambda A, B: f"gemm_n_{B.shape[0]}_k_{B.shape[1]}")
def gemm_bf16(
    A: torch.Tensor,
    B: torch.Tensor,
) -> torch.Tensor:
    # Fallback Implementation
    return torch.matmul(A, B)

It takes a lambda function that maps the inputs to a specific kernel definition name, then replaces the original computation function. The replaced function filters the best-performing kernel solution from the database based on the inputs, executes it, and returns its result. If no matching kernel solution is found, the original function runs as a fallback.

By using apply(), we can replace a large number of operators in FlashInfer. If you want to replace operators from your own implementations or other libraries, you can also easily achieve this with apply().

FlashInfer-Bench Leaderboard — The Arena for LLM Kernel Optimization

FlashInfer-Bench provides a standardized benchmark for AI agents to improve AI systems. Therefore, we also evaluated the capabilities of different models and built the FlashInfer-Bench Leaderboard. It directly leverages the FlashInfer Trace generated from the benchmark suite and provides a visualized display of the results.

We adopt the \(\text{fast}_p\) metric initially proposed by KernelBench to compare the performance of different kernels. We use the FlashInfer kernel as the baseline for comparison. \(\text{fast}_p\) represents the proportion of workloads on which a given kernel runs p times faster than FlashInfer. If p > 1, it means we have discovered a kernel that outperforms the state-of-the-art kernel library! In practice, most kernels have p < 1, indicating that current models still have significant room for improvement. For different values of p, \(\text{fast}_p\) forms a curve. The larger the area under the curve, the better the kernel’s performance.

We also found a small number of AI-generated kernels with a speedup greater than 1. We deeply understand the importance of ensuring the correctness of kernel implementations, so we are manually verifying the correctness of each kernel. We will soon release kernels with a speedup ratio greater than 1 after review.

We also provide a separate leaderboard for each kernel to make it easier to select the best implementation for each one. The leaderboard also allows further filtering by workload to select the best-performing kernels for specific workloads.

Furthermore, the model architecture overview shows all the kernels used by the model and indicates which ones are already covered by FlashInfer-Bench:

Summary

FlashInfer-Bench provides a systematic solution for building AI that improves AI systems. With it, we have already built a self-optimizing loop for AI—a virtuous cycle. The revolution of AI improving itself will continue—let’s see where it goes!

Looking ahead, we plan to continually expand our coverage of models and kernels. We will collaborate closely with GPU engineers and LLM engine developers to adapt to the fast-evolving needs of the field. This journey is a community effort, and we welcome your discussions and contributions.

For more information, please visit the following links:

• Leaderboard
• FlashInfer-Bench Dataset
• FlashInfer-Bench Python Package
• Documentation

Acknowledgements

FlashInfer-Bench is a research effort initiated in collaboration with CMU Catalyst, NVIDIA, and Bosch. We are bringing up an open community as part of the FlashInfer community and welcome contributions from the ML systems community.

We thank the entire FlashInfer-Bench team for their contributions to the project:

• Shanli Xing* (UW, CMU): Core components and web development
• Yiyan Zhai* (CMU): FlashInfer-Trace dataset, workload tracing system
• Alexander Jiang* (CMU): Benchmark system, agent design
• Yixin Dong* (CMU): Core idea, overall architecture design
• Yong Wu (NVIDIA): RMSNorm, sampling, fused MOE
• Zihao Ye (NVIDIA): FlashInfer support
• Charlie Ruan (UC Berkeley): Workload tracing system
• Yingyi Huang (CMU): Fused MOE
• Yineng Zhang (Independent researcher): SGLang integration
• Liangsheng Yin (Independent researcher): SGLang integration
• Aksara Bayyapu (CMU): LLM agents
• Luis Ceze (UW, NVIDIA): Project guidance and advice
• Tianqi Chen (CMU, NVIDIA): Project guidance and advice

We would also like to thank Mark Saroufim, Zhuoming Chen, Weihua Du, Bohan Hou, Hongyi Jin, Ruihang Lai, Xinyu Yang, Yilong Zhao, Haizhong Zheng, the FlashInfer community, the GPUMODE community, the HuggingFace community, the SGLang community, the TensorRT-LLM community, the vLLM community, Databricks, and xAI for their insightful feedback.

As vocabulary sizes grow larger in Large Language Models (LLMs), categorical sampling (token selection) has emerged as a significant performance bottleneck in LLM inference serving. The sampling operators in FlashInfer were first introduced in v0.0.5, and since then, the FlashInfer team has continuously improved their robustness and performance. In this blog post, we’ll explore the algorithms and implementation details behind FlashInfer’s sampling operators.

LLM Sampling

Categorical Sampling is the process that picks a specific next token from model output probabilities (over the vocabulary). In practice, filtering is applied before sampling to pass tokens with negligible probability, control generation behaviors, and enforce minimum probabilities, such as Top-P, Top-K, or Min-P thresholds:

Figure 1: The compute time break down highlighting the sampling process. In the vLLM 1xH100 configuration, our kernels reduce the overall sampling time by more than 50% across all three models.

1. Top-K

   Top-K sampling keeps only the $K$ tokens with the highest probabilities at each generation step. For example, if $K=50$, the model will ignore all tokens outside the top 50 likely candidates.

2. Top-P (Nucleus Sampling)

   Top-P rather keeps the smallest set of tokens whose cumulative probability just exceeds a threshold $P$. For example, if $P=0.9$, you accumulate token probabilities in descending order until their sum is at least 0.9.

3. Min-P

   Min-p filters out all tokens below a minimum threashold $p_\text{base} \times p_\text{max}$, where $p_\text{base}$ is parameter and $p_\text{max}$ is the largest probability in the inputs. This helps eliminate extremely unlikely tokens while still respecting relative differences among the top candidates.

In practice, the combination of Top-K and Top-P filtering is popular and used as the standard setting for LLM sampling. This allows for finer-grained control over the generation process. For example if we use the Top-K first filtering, we first limit the token set to the Top-K highest probabilities, and then apply a Top-P cutoff to filter the tail portion within those $K$ tokens. ¹

A PyTorch implementation of these samplers might look like this:

# vllm/vllm/model_executor/layers/sampler.py
def _apply_top_k_top_p(
    logits: torch.Tensor,
    p: torch.Tensor,
    k: torch.Tensor,
) -> torch.Tensor:
    logits_sort, logits_idx = logits.sort(dim=-1, descending=False)

    # Apply top-k.
    top_k_mask = logits_sort.size(1) - k.to(torch.long)
    # Get all the top_k values.
    top_k_mask = logits_sort.gather(1, top_k_mask.unsqueeze(dim=1))
    top_k_mask = logits_sort < top_k_mask
    logits_sort.masked_fill_(top_k_mask, -float("inf"))

    # Apply top-p.
    probs_sort = logits_sort.softmax(dim=-1)
    probs_sum = probs_sort.cumsum(dim=-1)
    top_p_mask = probs_sum <= 1 - p.unsqueeze(dim=1)
    # at least one
    top_p_mask[:, -1] = False
    logits_sort.masked_fill_(top_p_mask, -float("inf"))

    # Re-sort the probabilities.
    src = torch.arange(logits_idx.shape[-1],
                       device=logits_idx.device).expand_as(logits_idx)
    logits_idx_inv = torch.empty_like(logits_idx).scatter_(dim=-1,
                                                           index=logits_idx,
                                                           src=src)
    logits = torch.gather(logits_sort, dim=-1, index=logits_idx_inv)
    return logits

This code uses a combination of sorting, cumulative sums, and masking. While it is straightforward to follow, it induces performance bottleneck especially for large vocab size, because of the huge overhead of sorting.

In FlashInfer, we show that sampling under filtering can be done in sorting-free manner, and we introduce the Dual Pivot Rejection Sampling algorithm and design fused sampling kernel templates to fully leverage GPUs’ parallel computing capabilities, ultimately achieving logarithmic (in worst case) time complexity. In this blog, we’ll walk you through how we developed this algorithm integrating ideas from Inverse Sampling, Rejection Sampling, and final version of the algorithm with theorerical guarantee of convergence.

Algorithm

Inverse Transform Sampling

Figure 2: Inverse Transform Sampling. This animation illustrates the per-block process, and in practice the workload gets executed by blocks.

We begin with implementing a basic sampling kernel that selects tokens purely based on their probabilities, particularly in the GPU parallel computing context.

The method is inverse transform sampling, which draws samples from a probability distribution given its cumulative distribution function (CDF). As for the token samling process, the CDF would be the prefix sum of token probabilities. The algorithm proceeds like this:

1. Draw a random $u$ from $U\sim \text{Unif}(0,1)$.
2. Compute the prefix sums (CDF) for each sampled token $j$ with probability $p_j$: $F_j=\sum^{j}_{i=1}p_i$.
3. Locate the token $k$ such that $F_{k-1} \leq u < F_k$ as the result.

NVIDIA’s CUB library (now part of CCCL) provides efficient primitives for parallel computing, and we leverage the reduce and scan primitives to compute the prefix sums. We use one threadblock for each probability distribution, for batch sampling, we launch multiple threadblocks in parallel. Block-level reduce/scan primitives can be applied to a block of elements (BLOCK_SIZE = NUM_THREADS * NUM_ELEMENTS_PER_THREADS, e.g. 1024 * 4 for float input), for vocabulary size greater than BLOCK_SIZE, we split the vocabulary into multiple blocks and sequentially apply the same procedure on each block:

1. Initialize a running total $\texttt{a}=0$. Compute the probability sum $\texttt{a_local}$ for each block. If $\texttt{a} + \texttt{a_local}> u$, the sampled token lies in this block.
2. If not, we add $\texttt{a_local}$ to $\texttt{a}$ and move on to the next block.
3. Once we know the correct block, we perform a prefix sum over its tokens to pinpoint the exact token index.

We use BlockReduce and BlockScan for the per-block partial sum and prefix sums, and AdjacentDifference to locate the token index. In practice, we use early-stopping to terminate the inverse transform sampling process when the cumulative probability exceeds the random number $u$, so we don’t need to go through the whole vocabulary for each round.

Rejection Sampling

Figure 3: Top-P Rejection Sampling. This animation illustrates the per-block process, and in practice the workload gets executed by blocks.

For more advanced strategies such as Top-P sampling, we use rejection sampling to restrict which tokens can be selected. Rejection sampling draws from a target distribution by comparing random samples against a threshold and discarding those that do not meet it.

Taking the sampling kernel under Top-P filtering as an example, here is a simplified look of what happens:

1. Initialize the pivot to $0$, so initially all tokens are considered.
2. Perform an inverse transform sampling pass but ignoring tokens with probabilities below the current pivot. After sampling a token, update the pivot to that token’s probability.
3. Compute the remaining probability $\texttt{q}$ among tokens that still exceed this pivot:
   1. If $\texttt{q}$ remains greater than or equal to $\texttt{top_p}$, another round is needed to raise the pivot further and reject more tokens.
   2. Otherwise, if it is below $\texttt{top_p}$, we finalize the sampled token and mark success.
4. Repeat until successful.

The whole algorithm can be implemented in a single fused kernel, and it works similar for Top-K and other filtering strategies, other than we’ll be checking the number of tokens exceeding the pivot against $\texttt{top_k}$ or $\texttt{min_p}$.

In practice, we find that the number of rounds for returning a sampled token is usually small. It provides a substantial speedup compared to the naive PyTorch implementation because we avoid the sorting and multiple passes over the vocabulary, as well as multiple kernel launch overheads.

Dual Pivot Rejection Sampling

While this rejection sampling approach is simple and efficient in most cases, it has some limitations. There is no theoretical guarantee on the number of rounds needed to obtain a sampled token. This can lead to varying sampling times across different probability distributions, which in turn causes inconsistent inter-token latency during LLM inference serving. Such variability may impact the predictability and reliability of the serving system.

To address this issue, in FlashInfer v0.2.3, we introduce the a new algorithm called Dual Pivot Rejection Sampling, which uses two pivots for faster convergence in rejection sampling. The algorithm is as follows:

1. Let $f$ be a function that checks if a probability value is valid: $f(x)=1$ if valid, $0$ if not.
2. Initialize $\textrm{low} \leftarrow 0$ and $\textrm{high} \leftarrow \max_i(p_i)$ as the initial range, it’s guaranteed that $f(\textrm{low})=0$ and $f(\textrm{high})=1$.
3. Sample over probability values in the range $(\textrm{low}, \infty)$ using inverse transform sampling.
4. Suppose $j$ is the sampled token, let $\textrm{pivot}_1\leftarrow p_j$, and $\textrm{pivot}_2\leftarrow \frac{\textrm{pivot}_1+\textrm{high}}{2}$.
   1. If $f(\textrm{pivot}_1)=1$, we accept the sampled token and return $j$.
   2. If $f(\textrm{pivot}_1)=0$, $f(\textrm{pivot}_2)=1$, we set $\textrm{pivot}_1$ as new $\textrm{low}$ and $\textrm{pivot}_2$ as new $\textrm{high}$.
   3. If $f(\textrm{pivot}_1)=0$, $f(\textrm{pivot}_2)=0$, we set $\textrm{pivot}_2$ as new $\textrm{low}$.
5. Repeat step 3 and 4 until success.

Figure 4: Transition from round(i) to round(i+1) in Dual Pivot Rejection Sampling, we either accept the sampled token (case 1) or shrinking the range by at least half (case 2 and 3).

Figure 4 shows the transition from round(i) to round(i+1) in Dual Pivot Rejection Sampling, in each round, if the sampled token is accepted, we return the token, otherwise, the new range’s extent is $\frac{\text{high}-\text{pivot}_1}{2} < \frac{\text{high}-\text{low}}{2}$, which is at least half of the previous range. Thus it’s guaranteed that the number of rounds is $O(\log(1/\epsilon))$ where $\epsilon$ is the minimal possible value in floating point representation.

Theoretical Proof of the Correctness of Rejection Sampler

In this section, we provide a theoretical proof of the correctness of the rejection sampler, we choose the top-k sampling as an example, and other samplers can be proved in a similar way.

Nomenclature

Theorem

The algorithm outputs each top‑k index $j \in T$ with probability

i.e. exactly the distribution obtained by first discarding all non‑top‑k items and then sampling categorically inside the top‑k set.

Proof

Fix any pivot $\tau < \min_{j \in T} p_j$ (true at every step because $\tau$ is always taken from a rejected non‑top‑k item). Define

With

where $S(\tau)$ is the sum of all the probabilities of the tokens that are greater than $\tau$, and $W(\tau)$ is the remaining mass of “bad” items still above the threshold.

The next draw obeys

Hence

We show that the following formula is a valid solution to (★):

We can verify the solution by substituting it into (★):

because $S(\tau) = Z + W(\tau)$. Thus the claimed form satisfies the recurrence, so $Q_j(\tau) \equiv p_j/Z$.

Now let’s show that the solution is unique. Suppose there is another solution $Q_j’(\tau)$ satisfies (★), let’s define $\Delta_j(\tau) = Q_j(\tau) - Q_j’(\tau)$, we have:

The sum of the coefficient $\sum_{\substack{r \notin T \ p_r > \tau}} \frac{p_r}{S(\tau)} = \frac{W(\tau)}{S(\tau)}$, which satisfies:

Suppose $\tau^*$ is the pivot where \(|\Delta_j(\tau)|\) reach its maximum, if it’s positive, we have:

which leads to a contradiction, which means $\Delta_j(\tau^*) = 0$, and our solution is unique.

The algorithm starts with $\tau = 0$; therefore

exactly the desired top‑k categorical distribution.

Evaluation

Our evaluation demonstrates that FlashInfer’s sampling kernel delivers substantial improvements in both kernel-level latency and end-to-end throughput compared to traditional sorting-based implementations.

Figure 5: Throughput Comparison of Different Engine Kernel.

Figure 6: Sampling Latency Growth with Batch Size.

Community Adoption and Other Applications

The FlashInfer sampling kernel has gained widespread adoption across major LLM frameworks, including MLC-LLM, sglang, and vLLM. The community’s active engagement through feedback and bug reports has been instrumental in refining and improving our implementation.

Beyond token sampling, the rejection sampling algorithm have proven valuable in other areas of LLM inference optimization. Similar algorithm can also be applied to speculative decoding verification, like chain speculative sampling and tree speculative sampling. Recent innovations like Twilight have further advanced the field by successfully combining top-p sampling with sparse attention in a unified approach.

Implementation Details

While the algorithm is elegant in theory, implementing it efficiently in a GPU kernel requires careful attention to detail, particularly in the token selection logic in inverse transform sampling. One key challenge lies in the parallel prefix-sum operation used to locate sampled tokens. Due to the non-associative and non-commutative nature of floating-point arithmetic, parallel prefix-sum cannot guarantee monotonic outputs even with non-negative inputs. This can lead to invalid token generation if not handled properly. Special care must be taken to ensure numerical stability and correctness in the sampling implementation (and we made a lot of mistakes before got it right)

For a detailed look at our implementation and how we tackle these challenges, you can explore our source code. Additionally, FlashInfer offers a comprehensive set of APIs for probability cutoff and renormalization, such as top_p_renorm_probs and top_k_renorm_probs, enabling flexible composition of multiple sampling filters. These tools allow developers to build sophisticated sampling strategies tailored to their specific needs.

Acknowledgement

This blog is written by Shanli Xing, we thank the flashinfer team for their contributions to the flashinfer.sampling module:

• Zihao Ye: design and implementation of sampling kernels in CUDA.
• Bohan Hou: design and implementation of sampling kernels in TVM.
• Shanli Xing: design and implementation of min-p sampling kernels in CUDA.
• Tianqi Chen: propose the idea of rejection sampling for top-p.

Footnotes

After four months of development, we are thrilled to announce the release of FlashInfer 0.2. This major update introduces performance improvements, enhanced flexibility, and critical bug fixes. Key highlights of this release include:

• Faster sparse (page) attention with FlashAttention-3 template
• JIT compilation for attention variants
• Support for Multi-head Latent Attention (MLA) decoding

FlashAttention-3 Template with Block/Vector-Sparsity

FlashAttention-3 brings a breakthrough optimization for Hopper GPUs by cleverly overlapping softmax and matrix multiplication. FlashInfer 0.2 integrates FA-3 templates, achieving significant improvements in prefill attention performance on Hopper architecture.

Flexible Block-Sparsity and Vector-Sparsity

FlashInfer’s standout feature is its highly flexible block-sparse FlashAttention implementation, supporting any block size configuration. Our PageAttention operators are implemented as block-sparse attention kernels, where page_size specifies the block’s column count. At its finest granularity, FlashInfer supports vector-sparsity¹ (page_size=1), allowing for precise memory management (used in sglang) and efficient KV-Cache token pruning.

By leveraging CuTe’s CustomStride and ComposedLayout abstractions, we have extended vector-sparsity to FlashAttention-3. Inspired by CUTLASS’s gather/scatter convolution, this was achieved through a simple modification to the producer’s memory loading module.

Performance Benchmark

We compared two attention implementations: PageAttention with page_size=1 ² (use vector-sparse attention implementation) and variable-length dense attention ³, benchmarking them under identical problem sizes across both FA-2 (v0.1.*) and FA-3 (v0.2) backends. Benchmarks used head_dim=128, causal=True, varying batch sizes (B) and sequence lengths (L) with Gaussian-initialized input Q/K/V tensors.

Performance comparison between dense/vector-sparse attention on FA-2 and FA-3 templates on H100 SXM5, compiled with CUDA 12.4. y-axis: different settings, x-axis: achieved TFLOPs/s

Results: Vector-sparse attention achieves 90% of dense attention’s throughput under identical conditions. The FA-3 backend consistently outperforms FA-2. Thanks to FlashInfer’s stable API, upgrading from FA-2 to FA-3 requires no code changes—just install FlashInfer 0.2. The reference benchmark script for reproducing these results is available here.

JIT Compilation for Attention Customization

Inspired by FlexAttention, FlashInfer 0.2 introduces customizable programming interface to compile different attention variants. We designed a modularized attention template in CUDA/Cutlass. Users can define custom attention variants by specifying functors such as LogitsTransform/QueryTransform/etc in an attention variant class. The class string will specialize our pre-defined Jinja templates and FlashInfer uses PyTorch’s JIT load function to compile and cache these kernels. New variants like FlashSigmoid can be implemented with minimal code. See our JIT examples for more cases.

Left: JIT workflow in FlashInfer. Right: Compile new attention variants.

In addition to supporting new attention variants, other benefits of supporting JIT in FlashInfer include:

• Reduced wheel size: The binary size of FlashInfer increases exponentially in recently releases because we pre-compile combination of all attention variants. We have to reduce specialization to make wheel size managable which harms kernel performance (as observed in #602, FlashInfer v0.1.6’s prefill performance is even worse than FlashInfer v0.1.1 because we move compile-time parameters to runtime which harms performance). FlashInfer v0.2 address the issue by just pre-compile a subset of core kernels ahead-of-time, while leaving most of the attention variants JIT compiled.
• Light development: No need to reinstall FlashInfer for minor CUDA changes, by installing FlashInfer in JIT Mode.

We have optimized the speed of JIT compilation by minimizing header dependencies and utilizing split compilation. As a result, all kernels for Llama models can be JIT-compiled within 15 seconds on server-grade CPUs. For more details, check out our JIT warmup scripts.

Fused Multi-head Latent Attention (MLA) Decoding Kernel

Multi-head Latent Attention (MLA), introduced in Deepseek v2, compresses the KV-Cache by projecting it into low-rank matrices. Achieving high throughput for MLA is challenging due to a lack of optimized kernels. FlashInfer community recently implemented an fused kernel with the Matrix Absorption trick, improving memory efficiency. See #551 for detailed explanation.

MLA decode kernel workflow in FlashInfer

Future plans include accelerating MLA decoding with Tensor Cores, benefiting speculative decoding.

CUDAGraph Compatibility for Variable-Length Inputs

FlashInfer 0.2 fixes prefill attention’s incompatibility with CUDAGraph when query lengths vary during capture and replay stages, by accurately estimating upper resource bounds. CUDAGraphs now can be used to accelerate speculative decoding and chunked-prefil workloads with FlashInfer kernels.

torch.compile Compatibility

FlashInfer 0.2 adheres to the PyTorch Custom Operators Standard, ensuring compatibility with torch.compile.

Packaging and CI/CD

We now provide nightly builds so users can test the latest features without waiting for stable releases.

Other Notable Improvements

FusedAddRMSNorm Fix

Fixed numerical issues in FusedAddRMSNorm, which may cause bad outputs for some models.

Cutlass’s SM90 Grouped-GEMM Integration

We integrated Cutlass 3.5 SM90 Grouped-GEMM into our SegmentGEMM API, accelerating LoRA and MoE serving.

Non-Contiguous KV-Cache Support

KV-Cache can now utilize non-contiguous storage layouts, improving support for offloading.

Faster plan Functions

plan functions now use non-blocking host-to-device memory transfers, improving performance. After FlashInfer v0.2, it’s encouraged to pass host tensors instead of device tensors to reduce synchronization in the plan function.

KV-Cache Append Optimization

KV-Cache append throughput for small batch sizes was improved by parallelizing per element instead of per request. A new API, get_batch_indices_positions, supports this. Note that we made some breaking changes to this API to accomodate different parallelization mode. See our benchmark for the new API usage.

Standardized RoPE Interface

We standardized the RoPE interface to align with other frameworks. FlashInfer adopted fp32 sin/cos computations to avoid numerical issues.

Roadmap

We appreciate the community’s love and support. To enhance transparency, we’ve published our development roadmap, where you can provide feedback and influence FlashInfer’s future.

Community Contributions

The number of contributors grew from 41 to 52 since v0.1.6. We thank the following developers for their contributions:

• @yzh119: JIT, FA-3 Template, and others
• @abcdabcd987: torch.compile support, packaging
• @nandor: Variable-length CUDAGraph support
• @ur4t: Packaging, CI/CD
• @zhyncs: Nightly builds, CI/CD
• @tsu-bin: MLA decoding
• @xslingcn: Cutlass Grouped-GEMM
• @yuxianq: JIT, bug fixes
• @LinHeLurking: Non-contiguous KV-Cache
• @Abatom: FusedAddRMSNorm fix
• @jeejeelee: Grouped-GEMM bug fixes
• @mvpatel: Faster plan functions
• @Ubospica: Pre-commit setup
• @dc3671: Improved unit tests
• @Pzzzzz5142: JIT compilation fix
• @reyoung: Bug fixes
• @xiezhq-hermann: ARM compilation fixes
• @Bruce-Lee-LY: Performance optimizations
• @francheez: Typo fixes

LLM (Large Language Models) Serving quickly became an important workload. The efficacy of operators within Transformers – namely GEMM, Self-Attention, GEMV, and elementwise computations are critical to the overall performance of LLM serving. While optimization efforts have extensively targeted GEMM and GEMV, there is a lack of performance studies focused on Self-Attention in the context of LLM serving. In this blog post, we break Self-Attention down into three stages: prefill, decode, and append; analyze the performance bottleneck of Self-Attention on both single-request and batching scenarios in these three stages; and propose a solution to tackle these challenges. These ideas have been integrated into FlashInfer, an open-source library for accelerating LLM serving released under Apache 2.0 license.

FlashInfer has been developed by researchers from the University of Washington, Carnegie Mellon University, and OctoAI since summer 2023. FlashInfer provides PyTorch APIs for quick prototyping, and a dependency-free, header-only C++ APIs for integration with LLM serving systems. Compared to existing libraries, FlashInfer has several unique advantages:

1. Comprehensive Attention Kernels: FlashInfer implements attention kernels that cover all the common use cases of LLM serving with state-of-the-art performance, including single-request and batching versions of Prefill, Decode, and Append kernels, on various formats of KV-Cache (Padded Tensor, Ragged Tensor, and Page Table).
2. Optimized Shared-Prefix Batch Decoding: FlashInfer enhances shared-prefix batch decoding performance through cascading, resulting in an impressive up to 31x speedup compared to the baseline vLLM PageAttention implementation (for long prompt of 32768 tokens and large batch size of 256), check another blog post for more details.
3. Accelerate Attention for Compressed/Quantized KV-Cache Modern LLMs are often deployed with quantized/compressed KV-Cache to reduce memory traffic. FlashInfer accelerates these scenarios by optimizing performance for Grouped-Query Attention, Fused-RoPE Attention and Quantized Attention. Notably, FlashInfer achieves up to 2-3x speedup for Grouped-Query Attention on A100 & H100, compared to vLLM implementation.

FlashInfer has been adopted by LLM serving systems such as MLC-LLM (for its CUDA backend), Punica and sglang. We welcome wider adoption and contribution from the community. Please join our discussion forum or creating an issue to leave your feedback and suggestions.

Attentions in LLM Serving

There are three generic stages in LLM serving: prefill, decode and append. During the prefill stage, attention computation occurs between the KV-Cache and all queries. In the decode stage, the model generates tokens one at a time, computing attention only between the KV-Cache and a single query. In the append stage, attention is computed between the KV-Cache and queries of the appended tokens. append attention is also useful in speculative decoding: the draft model suggests a sequence of tokens and the larger model decides whether to accept these suggestions. During the attention stage, proposed tokens are added to the KV-Cache, and the large model calculates attention between the KV-Cache and the proposed tokens.

The crucial factor affecting the efficiency of attention computation is the length of the query ($l_q$), determining whether the operation is compute-bound or IO-bound. The operational intensity (number of operations per byte of memory traffic) for attention computation is expressed as $O\left(\frac{1}{1/l_q + 1/l_{kv}} \right)$, where $l_{kv}$ represents the length of the KV-Cache. During the decode stage, where $l_q$ is consistently 1, the operational intensity is close to $O(1)$, making the operator entirely IO-bound. In the append/prefill stages, the attention operational intensity is approximately $O(l_q)$, leading to compute-bound scenarios when $l_q$ is substantial.

The diagram illustrates the attention computation process in the prefill, append, and decode stages:

Figure 1: Decode attention fills one row of the attention map at a time, prefill attention fills the entire attention map (under the causal mask), and the append attention fills the trapezoid region.

The figure below shows the roofline model of the three stages of attention computations. Decode attention performance is always underneath the peak bandwidth ceiling (bounded by peak memory bandwidth in GPU), and thus is IO-bound. Prefill attention has high operational intensity and is under the peak compute performance ceiling (bounded by peak floating point performance). Append attention is IO-bound when the query length is small, and compute-bound when the query length is large.

Figure 2. Roofline model of attention operators in LLM Serving, data from A100 PCIe 80GB.

Single-Request and Batching

There two common ways to serve LLM models: batching and single request. Batching groups several user requests together and process them in parallel to improve the throughput, however, the operational intensity of attention kernels is irrelevant to batch size ¹, and batch decoding attention still has operational intensity of $O(1)$.

FlashInfer Overview

FlashAttention proposes to fuse multi-head attention into a single kernel by generalizing online softmax trick to self-attention, thus avoiding the overhead of materializing the attention matrix on GPU global memory. FlashAttention2 further improves performance by adopting a more reasonable tiling strategy and reducing the number of non tensor ops to alleviate the issue that A100/H100 has low non-tensor cores performance. vLLM proposes PageAttention where KV-Cache is organized as a page table, to alleviate the memory fragmentation issue in LLM serving.

FlashInfer implements single-request and batch version of FlashAttention for all three stages: prefill, append and decode on versatile KV-Cache formats (e.g. Ragged Tensor, Page Table). For single decode/prefill and batch decoding kernels, FlashInfer achieves state-of-the-art performance for single-request decode/prefill and batch decode kernels. Moreover, FlashInfer implements prefill/append kernels for Paged KV-Cache which none of the existing libraries have done before, and it be used to serve models in speculative decoding setting.

Many recent work proposes KV-Cache compression techniques to reduce memory traffic. In light of this, FlashInfer optimize kernels for Grouped-Query Attention, Fused-RoPE Attention and Quantized Attention for efficient serving with compressed KV-Cache:

• Grouped Query Attention: Grouped Query Attention uses a smaller number of heads for keys and values thus saving memory traffic. The operational intensity of Grouped Query Attention grows from $O(1)$ to $O\left(\frac{H_{qo}}{H_{kv}}\right)$ where $H_{qo}$ is the number of heads for queries and $H_{kv}$ is the number of heads for keys and values. GPUs such as A100/H100 has low non-tensor cores performance, and thus traditional implementation of Grouped Query Attention is compute-bound. FlashInfer proposes to use prefill kernels (which utilizes Tensor Cores) for decode attention in GQA, and achieves up to 2-3x speedup compared to vLLM implementation.
• Fused-RoPE Attention: RoPE (Rotary Positional Embeddings) has become a standard component of Transformers, most existing serving systems stores post-RoPE keys (the keys after applying rotary embeddings) in KV-Cache. However, some recent work such as StreamingLLM will prune tokens in KV-Cache, and the position of tokens will change after pruning, thus the post-RoPE keys in KV-Cache become meaningless. In this case, FlashInfer proposes to save pre-RoPE keys in KV-Cache, and fuses RoPE into attention kernel. Experiments on various platform and settings show that FlashInfer’s Fused-RoPE Attention kernel can apply RoPE on the fly with negligible overhead.
• Quantized Attention: Another way to compress KV-Cache is through pruning, FlexGen and Atom show that it’s possible to prune KV-Cache to 4-bit with negligible accuracy loss. FlashInfer implements low-precision attention kernels so that we can achieve nearly linear speedup to the compression ratio (~4x for 4bit, ~2x for 8bit).

Some recent work such as LightLLM and sglang uses a special form of PageAttention where page size equals one, for easy management of KV-Cache in complicated serving scenarios such as structured generation. FlashInfer optimizes PageAttention kernels by pre-fetching page indices in GPU shared memory, so that kernel performance is not affected by the page size.

In the subsequent sections, we will delve into the detailed optimizations and benchmark results achieved by FlashInfer.

Benchmark Settings

Hardware

We benchmarked on 4 different GPUs: H100 SXM 80GB, A100 PCIe 80GB, RTX 6000 Ada and RTX 4090, the first two is widely used data center GPU in Hopper and Ampere architectures, respectively, and latter two are workstation and gaming GPUs in Ada Lovelace architecture that are much more affordable, the specifications are listed in the following table:

H100 SXM 80GB uses HBM3 and A100 PCIe 80GB use HBM2e, both have larger memory bandwidth than RTX Ada 6000 and RTX 4090 that use GDDR6X. RTX Ada 6000 and RTX 4090 have much larger non-Tensor Cores peak performance (90 and 80 TFLops/s respectively) than A100 (20 TFLops/s). The later three GPUs have similar peak Tensor Cores (fp16 input, without sparsity) performance for f16 accumulation, RTX 4090’s Tensor Cores have 2x throughput with fp16 accumulation compared to fp32 accumulation, while the other GPUs have the same throughput for fp16 and fp32 accumulation.

Below is the roofline curve of the four GPUs for both Tensor Cores and CUDA Cores:

Figure 3: Devices Roofline of 4 GPUs, Tensor Cores Performance and CUDA Cores Performance are indicated separately.

The ridge point is determined by the ratio of peak floating point performance and memory bandwidth.

Software

The baselines being compared are: FlashAttention 2.4.2 which has incorporated FlashAttention 2 and Flash-Decoding, and vLLM v0.2.6 that implements PageAttention 1&2. For vLLM we use prebuilt wheels from pip, we build FlashAttention & FlashInfer from source code with the NVCC compiler in CUDA 12.3.1 release. The kernel profiling is done with nvbench library, we take the “cold” GPU time which flushes the L2 cache before each kernel launch.

Metrics

We report achieved TFLops/s for prefill & append attention kernels, and GPU memory bandwidth utilization (computed by $\frac{\textrm{number of bytes read by the kernel}}{\textrm{kernel latency}} / \textrm{hardware GPU memory bandwidth}$) for decode & append attention kernels.

Prefill Kernels

For prefill (multi-query) attention, we reimplemented the FlashAttention 2 algorithm in pure CUDA with some additional optimizations. Standard FlashAttention implementation uses Tensor Cores with fp16 input and fp32 accumulator, however, RTX 4090 GPUs has lower Tensor Cores performance with fp32 accumulator, we observe that the $\frac{\mathbf{q}\cdot \mathbf{k}^{T}}{\sqrt(d)}$ phase in attention computation have small value range and can be accumulated with fp16 (because the head dimension is always small: e.g. 128), FlashInfer provides an allow_fp16_qk_reduction option to allow this optimization (but still use fp32 accumulation for $\mathbf{score} \cdot \mathbf{v}$), this optimization could bring 50% speedup on RTX 4090. Below is the performance comparison of FlashInfer 0.0.1 and FlashAttention 2.4.2 on different GPUs:

Figure 4: Single request prefill kernel performance, use Llama2-7B setting: num_kv_heads=num_qo_heads=32, head_dim=128. Sequence length varies from 32 to 65535.

In f32 accumulation setting, FlashInfer’s prefill kernel implementation achieves best performance on all 4 GPUs. allow_fp16_qk_reduction option can further improve performance, especially for RTX 4090.

Append & Decode Optimizations

Append and decode attention tend to have larger KV length than query length, which could limit the SM(StreamMultiprocessor) utilization in GPUs when batch size is small, FlashInfer propose to use the Split-K trick in GEMM optimizations which splits the KV-Cache on sequence dimension to increase parallelism. Another work, Flash-Decoding also explored this idea, you can check their great blog post for visualizations and explanations. Below is the decode attention performance comparison of FlashInfer 0.0.1 and FlashAttention 2.4.2 on different GPUs:

Figure 5: Single request decode kernel performance, use Llama2-7B setting: num_kv_heads=num_qo_heads=32, head_dim=128. Sequence length varies from 32 to 65536.

FlashInfer achieves best performance on all 4 GPUs, and the GPU bandwidth utilization is close to 100% for long sequences. An interesting fact is that split-KV do not improve performance for GPUs such as RTX Ada 6000 and RTX 4090 because they have relatively smaller memory bandwidth and stronger CUDA Cores performance (decode attention has low operational intensity and we use CUDA Cores in non-GQA setting). Unlike compute units which is SM local, the global memory traffic on GPUs is shared, thus using 32 (number of heads in Llama2-7B setting) of 108 SMs can still fully utilize the memory bandwidth if the operator is not compute-bound. A100 GPUs has low CUDA Cores performance (20 TFLops/s), using 32 of 108 SMs (5.9 TFLops/s) will make the kernel compute-bound (besides multiply and add, there are also time-consuming computations such as exp in attention computation), and split-KV is helpful in this case.

For batch decoding attention, FlashInfer implements PageAttention with optimizations such as pre-fetching page indices, below is performance comparison of FlashInfer PageAttention kernel and vLLM PageAttention kernel:

Figure 6: Batch decode kernel performance, use Llama2-7B setting: num_kv_heads=num_qo_heads=32, head_dim=128, batch_size=[1,16,64]. Sequence length varies from 32 to 65536 for batch_size = 1, from 32 to 4096 for batch_size = 16, and from 32 to 1024 for batch_size = 64.

FlashInfer PageAttention kernel has consistent speedup over vLLM 0.2.6’s implementation in different batch sizes and different sequence lengths.

We also benchmark the append attention kernels (append attention shares the same set of APIs with prefill attention, with the distinction that query length is smaller than key/value length in append attention):

Figure 7: Append attention kernel performance, use Llama2-7B setting, num_kv_heads=num_qo_heads=32, head_dim=128. The append length is set to 128 or 256, KV sequence length varies from 32 to 65536.

FlashInfer still achieves the best performance on all 4 GPUs, either with fp16 or fp32 qk accumulator. Split-KV significantly improves the performance of append kernels for append length of both 128 and 256, because the operational intensity of the operator becomes large, and using 32/100+ SMs no longer provides enough compute units, thus making the kernel compute-bound. Note that the ridge point of RTX 4090’s Tensor Cores fp32 accumulator roofline is 163 (165 TFLops/s / 1008 GB/s), the kernel will be compute bound when query length (which approximately equals operational intensity) reaches 256, using allow_fp16_qk_reduction can alleviate the issue.

FlashInfer also implemented batch append attention kernel where key/value is stored in Page Tables, this could accelerate speculative decoding in LLM serving, and we will discuss this in another blog post.

Grouped-Query Attention

Grouped-Query Attention uses smaller number of key/value heads than the number of query/output heads, makes the operational intensity higher than ordinary multi-head attention. FlashInfer proposes to use prefill(multi-query) attention kernel, which utilize Tensor Cores, for decode attention in GQA. Below is the GQA kernel performance comparison between FlashInfer (w/ CUDA Cores and w/ Tensor Cores), and FlashAttention 2.4.2 on A100 & H100:

Figure 8: Single request GQA decode performance, use llama2-70b setting: tp=2, num_kv_heads=4, num_qo_heads=32, head_dim=128. Sequence length varies from 32 to 65536.

For single-request GQA decoding attention, FlashInfer (Tensor Cores) achieves better performance than FlashAttention 2.4.2 on both A100 & H100, and FlashInfer (CUDA Cores) can only achieve 40%+ bandwidth utilization because of limited CUDA Cores performance.

Figure 9: Batch GQA decode performance, use llama2-70b setting: tp=2, num_kv_heads=4, num_qo_heads=32, head_dim=128. batch_size is set to 64 and sequence length per request varies from 32 to 8192.

For batch GQA decoding attention, FlashInfer w/ Tensor Cores is 3x faster than vLLM PagaAttention when batch_size=64.

Fused-RoPE Attention

KV-Cache compression techniques such as H2O and Streaming-LLM prunes KV-Cache by removing tokens, and the original relative positions of tokens in KV-Cache will be polluted, storing post-RoPE keys in KV-Cache become meaningless. FlashInfer implements high-performance Fused-RoPE attention kernels which applies RoPE on the fly, below is the performance comparison of FlashInfer decoding attention with and without RoPE:

Figure 10: Fused RoPE attention performance, use Llama2-7B setting: um_kv_heads=num_qo_heads=32, head_dim=128. Sequence length varies from 32 to 65536.

RoPE has negligible overhead on all 4 GPUs, especially for RTX 6000 Ada and RTX 4090 GPU which has strong CUDA Cores performance (RoPE requires sin/cos computation that can not be accelerated with Tensor Cores).

Low-Precision Attention

More and more work show that KV-Cache can be quantized to low bits with negligible accuracy loss. FlashInfer implements high-performance fp8 decode decode kernels, which could accelerate the kernel by up to 2x compared with fp16 kernels:

Figure 11: FP8 decode attention performance, use Llama2-7B setting: num_kv_heads=num_qo_heads=32, head_dim=128. Sequence length varies from 32 to 65536.

There is some gap between bandwidth utilization of fp8 and fp16 kernels, however the gap is getting closer as the query length grows.

Atom implemented high-performance decode attention kernels with int4 quantization on top of FlashInfer.

Effect of Page Size on FlashInfer’s PageAttention

The FlashInfer decode kernels prefetches page indices in GPU shared memory, thus minimizing the impact of the page size on kernel performance. Below is the performance comparison of FlashInfer PageAttention with different page sizes on A100:

Figure 12: Batch decode performance on different page_size. batch_size is set to 1, use Llama2-7B setting: num_kv_heads=num_qo_heads=32, head_dim=128. Sequence lengths varies from 32 to 65536. We also add a reference line for the performance of FlashInfer single-request decode attention without using Page Table.

The memory bandwidth utilization of the 4 different page sizes are nearly identical, and they are close to the single-request decode attention curve, which indicates that page size has little effect on FlashInfer PageAttention’s kernel performance, and page table itself has little overhead.

Some recent work such sglang explores novel KV-Cache management algorithm which requires page_size=1, and the performance could benefit from FlashInfer’s optimization.

Remarks and Future Work

The idea of splitting KV-Cache on sequence dimension to increase parallelism was also explored in Flash-Decoding, FlashInfer implemented this idea concurrently, see our github checkpoint on Sept 1st, 2023 and our public talk at TVM Unity Open Development Meeting on Sept 5th, 2023.

Currently FlashInfer only supports NVIDIA GPUs, the AMD and Apple GPU version of FlashInfer have been initially supported in MLC-LLM project with the help of Apache TVM compiler. Our next release will include the 4-bit fused dequantize+attention operators proposed in Atom and LoRA operators used in Punica. In a longer term, we are interested in performance optimization on post-Hopper NVIDIA GPUs and AMD/Apple GPUs, and new operators from emerging LLM architectures. Please check our roadmap for development plans, and leave your suggestions on what features you want to see in FlashInfer.

Acknowledgement

FlashInfer is inspired by FlashAttention 2, vLLM, cutlass and Stream-K project.

This blog post is written by Zihao Ye. We thank the entire FlashInfer team for their contributions to the project:

• Zihao Ye (UW): design and implementation of FlashInfer
• Lequn Chen (UW): page table data structure design, API design, CI/CD and Punica integration
• Ruihang Lai (CMU): KV-Cache design, API design and integration with MLC-LLM
• Yilong Zhao (UW & SJTU): int4 attention operators
• Size Zheng (UW & PKU): CUDA optimizations and speculative decoding
• Junru Shao and Yaxing Cai (OctoAI): MLC-LLM integration
• Bohan Hou and Hongyi Jin (CMU): porting FlashInfer to AMD and Mac GPUs with Apache TVM
• Liangsheng Yin (SJTU & LMSys): PyTorch bindings and sglang integration.
• Yifei Zuo (UW & USTC): PyTorch bindings
• Tianqi Chen (CMU & OctoAI): recursive form of softmax/attention merge and advices
• Luis Ceze (UW & OctoAI): performance breakdown analysis and advices

We also thank Masahiro Masuda (OctoAI), Yixin Dong (UW & SJTU), Roy Lu (UW), Chien-Yu Lin (UW), Ying Sheng (Stanford & LMSys) and Lianmin Zheng (Berkeley & LMSys) for their valuable feedbacks and discussions.

Citation

@misc{flashinfer,
    title = {Accelerating Self-Attentions for LLM Serving with FlashInfer},
    url = {https://flashinfer.ai/2024/02/02/introduce-flashinfer.html},
    author = {Ye, Zihao and Chen, Lequn and Lai, Ruihang and Zhao, Yilong and Zheng, Size and Shao, Junru and Hou, Bohan and Jin, Hongyi and Zuo, Yifei and Yin, Liangsheng and Chen, Tianqi and Ceze, Luis},
    month = {February},
    year = {2024}
}

Footnotes

In this blog post, we introduce Cascade Inference, which simply decouples attention of shared prefix and unique suffixes, and enables storing shared KV-Cache in GPU shared memory (SMEM for short) for fast access in multiple requests. We show that Cascade Inference can greatly accelerate shared-prefix batch decoding operator, with up to 31x speedup compared to the baseline vLLM PageAttention implementation and 26x speedup compared to FlashInfer batch decoding operator without cascading on a H100 SXM 80GB. The kernels have been supported in FlashInfer as PyTorch and C++ APIs.

Figure 1. An example of serving Document QA for multiple users, all of the requests share the same book as prompt.

Background

GPU’s memory hierarchy is composed of global memory, L2 Cache, SMEM/L1 Cache and registers. While global memory and L2 Cache is shared across all streaming multiprocessors (SMs), SMEM/L1 Cache and registers are private to each SM. The throughput of accessing global memory and L2 Cache is much lower than that of accessing SMEM/L1 Cache and registers. Therefore, it is important to minimize the access to global memory and L2 Cache to achieve high throughput.

In CUDA programs, independent tasks are dispatched to different thread blocks ¹, and each thread block is executed by one SM. For pre-Hopper architectures ², each thread block can only access its local SMEM and registers.

Difference between Multi-Query and Single-Query CUDA Kernels

In multi-query attention (used in prefill/append) kernels, multiple queries access the same region of KV-cache. The usual implementation of multi-query attention kernel process multiple queries in a single thread block and loads KV-Cache to shared memory and computes the attention between multiple queries and the KV-Cache in parallel. This approach is bandwidth-efficient and can maximize TFLOPs/s by utilizing Tensor Cores, but not applicable if the KV-Cache for different queries are different.

The single-query attention kernel (used in decode), on the other hand, assumes that each query has its own KV-Cache, and batching cannot increase the operational intensity of this operator. In this case, there will be no benefit of processing multiple queries in a same thread block because the opportunity of reusing KV-Cache is limited. Most implementations of decode attention kernel process one query in a single thread block, to guarantee parallelism so that all SMs are fully utilized. However, this approach is not memory bandwidth efficient because each thread block needs to load the KV-Cache from global memory (or L2 cache, if the cache line has been hit before).

Divide and Conquer

Neither multi-query attention nor single-query attention kernel is a good fit for shared-prefix batch decoding. However, multi-query attention is perfect for attention between queries and shared prefix, while single-query attention can deal with the attention between queries and unique suffixes. Can we combine the advantages of both approaches?

Recursive Attention

The answer is “yes” if we can find a way to “merge” the attention of the same queries with shared prefix and unique suffixes. Fortunately, FlashAttention has shown it’s possible to combine local softmax/attention results by not only storing the local attention result, but also the normalization scales and renormalizing local attention results on the fly. Here we formulate the idea in concise notations:

Suppose $s_i$ is the pre-softmax attention score between query and the key at index $i$:

we can generalize the definition from single index to an index set:

let’s also generalize the value vector $\mathbf{v}$ from index to index sets (Note that the generalization of both $s$ and $\mathbf{v}$ are self-consistent: when $I$ equals $\{i\}$, we have $s(I) = s_i$ and $\mathbf{v}(I) = \mathbf{v}_i$):

the $\textrm{softmax}$ function is restricted to the index set $I$. Note that $\mathbf{v}(\{1,2,\cdots, n\})$ is the self-attention output of the entire sequence. The attention state between a query with KV of an index set $I$ can be defined as a tuple $\begin{bmatrix}\mathbf{v}(I) \\ s(I)\end{bmatrix}$, then we can define a binary merge operator $\oplus$ to combine two states as (in practice we will minus $s$ with maximum value to guarantee numerical stability and here we omit the trick for simplicity):

the merge operator can be generalized to any number of attention state inputs:

The above n-ary merge operator is consistent with the binary merge operator, and we can prove the operator is communicative and associative. There are different ways to get the attention state of the entire sequence by merging the attention states of index subsets, and the final outcome is mathematically equivalent:

Figure 2. Different order to merge attention states are mathematically equivalent.

Recursive Attention allow us to decompose attention computation into multiple stages, different stages can be dispatched to different compute units/devices. The KV sequence partitioning trick in FlashInfer and Flash-Decoding uses the same idea to merge partial attention states from different thread blocks.

Cascade Inference: The Algorithm

With the merge operator, we can dispatch attention on different KV subsets to different kernel implementations. For shared-prefix batch decoding attention, we propose the following Divide-and-Conquer algorithm:

1. Use multi-query (prefill/append) attention kernel to compute the attention state between queries and KV-Cache of shared prefix.
2. Use batch decode attention kernel to compute the attention state between queries and KV-Cache of unique suffixes.
3. Use merge operator to combine two attention states to get the final attention output.

The overall workflow is explained on the left side of Figure 3, different color of rectangles are processed in different thread blocks in GPU. Note that for multi-query attention kernels, we access KV-Cache through SMEM or registers and for decode kernels we can only access KV-Cache through L2 Cache or Global Memory. Cascade Inference allow us to maximize memory reuse for common prefix, thus making the attention computation much more memory efficient.

Figure 3. Workflow of Cascade Inference, throughput values adapted from blog: TPU vs GPU vs Cerebras vs Graphcore: A Fair Comparison between ML Hardware

We call the divide-and-conquer approach for shared-prefix attention the “Cascade Inference”.

Evaluations

We evaluate Cascade Inference on H100 SXM 80GB and A100 PCIE 80GB GPUs. The input shape are adapted from LLaMA2-7B (32 heads, 128 dimension per head). We varies three parameters: number of requests (batch size), shared prefix length and unique suffix length per request. The baseline implementations is PageAttention kernel implemented in vLLM 0.2.6, we also show the performance of FlashInfer batch decoding operator without cascading. The page size (or block size, equivalently) is fixed to 16 for all implementations (FlashInfer w/ and w/o cascading, vLLM PageAttention).

Figure 4. Speedup over vLLM PageAttention on H100 SXM 80GB

Figure 5. Speedup over vLLM PageAttention on A100 PCIe 80GB

Figure 4 and 5 show the normalized performance on FlashInfer kernels in cascading and non-cascading setting over vLLM implementation. FlashInfer kernels in both settings outperforms vLLM kernels, and cascading kernels significant speedup over non-Cascade Inference kernels in most cases. The benefit of cascade inference increases as shared prefix length and batch size grows (where the prefill kernel dominates execution time) and decreases as we increase unique suffix length (where the batch decode kernel dominates execution time). For very long shared prompt (32768), the decode kernel can get up to 31x speedup on H100 SXM 80GB with large batch size(≥128) and short unique kv-length (≤256).

Remarks and Future Work

The idea of Cascade Inference can be generalized to multiple levels (we only show two levels in this blog post) and multiple shared prefixes, the multi-level, multi shared-prefix Cascade Inference has been integrated to MLC-Serving: the universal serving framework based on MLC-LLM, we will show the end-to-end speedup in future blog posts.

Recently, SGLang (a domain-specific language for programming LLMs) proposes RadixAttention, where the KV-Cache is organized as a radix tree structure and the attention can be further accelerated with multiple-level Cascade Inference. We are collaborating with SGLang team to get this feature landed.

Citation

@misc{cascade-inference,
    title = {Cascade Inference: Memory Bandwidth Efficient Shared Prefix Batch Decoding},
    url = {https://flashinfer.ai/2024/02/02/cascade-inference.html},
    author = {Ye, Zihao and Lai, Ruihang and Lu, Bo-Ru and Lin, Chien-Yu and Zheng, Size and Chen, Lequn and Chen, Tianqi and Ceze, Luis},
    month = {February},
    year = {2024}
}

Footnotes & References