FlashInfer: Kernel Library for LLM Serving

The flashinfer_benchmark.py framework is a unified tool for benchmarking FlashInfer kernels. It allows developers to:

• Benchmark performance for Attention, GEMM, MOE, Normalization, Quantization, Sampling, RoPE, Mamba, and GDN APIs.
• Compare performance across various backends including FlashAttention2/3, cuDNN, cuBLAS, CUTLASS, CuTe-DSL, TensorRT-LLM, and Triton.
• Compare performance across different configurations.
• Perform batch performance testing for multiple test cases.

FlashInfer uses the concept of attention states to characterize the relationship between a query and a set of key/value pairs. An attention state is represented as a tuple $(s(I), \mathbf{v}(I))$, where:

• $s(I)$ is the generalized score, also known as the log-sum-exp (lse).
• $\mathbf{v}(I)$ is the attention output (the weighted sum of values).

Because the merge operator ($\oplus$) is both commutative and associative, you can compute attention on subsets of KV pairs independently and merge them in any order. This allows for efficient parallelization and offloading.

The CuTeDSL MegaMoE tuning surface includes knobs, per-size default profiles, and online autotuning. Performance is measured using microbenchmarks and real-model geometry sweeps.

Key performance characteristics observed:

• Small-batch regime (< 512 tokens/rank): Performance is typically weight-load bound; nvfp4 variants show parity with deep_gemm (dg).
• Large-batch regime (> 1024 tokens/rank): nvfp4 with +combine_nvfp4 or +combine_mxfp8 provides significant speedups (up to 1.78x) over deep_gemm_mega.
• Accuracy: While nvfp4 variants may have higher relative L2 loss (~21-25%) compared to deep_gemm_mega (20.6%), the mxfp8_cutedsl backend offers significantly better accuracy (6.4% loss).

The moe_ep kernel source is organized into layers to isolate the raw kernel code from the FlashInfer (FI) backends:

• src/: A verbatim copy of the kernel-team drop. Do not edit or add files here.
• shim/: Thin adapters over src/. This is where all adaptation (path bootstrapping, symbol re-exports, API shims) lives.
• __init__.py: The public API for moe_ep. It re-exports symbols from shim/ and is the only layer the FI backends should import.
• kernel_helpers.py: A single re-export point for raw-kernel helpers and constants used by backends and tests.

Layer Isolation Rules:

• shim/ is the only layer that imports packages from src/.
• FI backends must import kernel helpers/constants/launch entry points only from the package __init__, never from src/ directly.

FlashInfer provides two primary architectural patterns for Mixture of Experts (MoE) with Expert Parallelism (EP):

1. Split Architecture (MoEEpSplitLayer): This approach splits the MoE operation into distinct stages. It uses a SplitKernelBackend and a Fleet mechanism to manage dispatching, inner computation, and combining results. This is typically used when combining NCCL-based Expert Parallelism with fused MoE kernels.

2. Mega Architecture (MoEEpMegaLayer): This approach uses a MegaKernelBackend to wrap complex configurations into a MegaConfig. It is designed for high-performance execution by binding EP bootstrap processes (supporting torch_dist and nvshmem) and managing internal workspaces for stages like input staging and computation within a single unified kernel flow.

The flashinfer.sampling module provides high-performance kernels specifically designed for Large Language Model (LLM) sampling tasks. These kernels include operations for sampling from probabilities or logits, as well as specialized sampling strategies like Top-P, Top-K, and Min-P.

If you only require efficient Top-K selection without actual sampling, use the flashinfer.top_k API instead.

The recursive attention mechanism in FlashInfer enables several high-performance optimization patterns:

1. Shared-Prefix Batch Decoding: Decomposes attention into shared prefix attention and unique suffix attention. This allows different kernel implementations to handle different parts of the KV-cache, potentially providing up to 30x acceleration in long-context, large-batch scenarios.
2. KV Sequence Parallelism: Partitions the KV sequence dimension across different thread-blocks (similar to the Split-K trick in GEMM) and merges the results in a second pass. This helps overcome GPU memory limits and improves SM utilization during long-context inference.

Level 5 (Statistics) is compatible with CUDA Graphs. Statistics are computed by a captured CUDA kernel that emits one line per tensor via device-side printf during graph replay.

When a tensor is logged, the host-side log records a correlation marker (e.g., [stats deferred to GPU kernel: id=7; ...]). During cuda_graph.replay(), the kernel prints the statistics for that ID (e.g., [flashinfer stats] id=7 numel=4096 min=-3.42 max=3.59 mean=0.01 nan=0 inf=0).

Supported dtypes: float32, float16, bfloat16, int32, int64, uint8. For other dtypes (like fp8/fp4), a message [statistics skipped: CUDA graph capture in progress] is emitted.

with torch.cuda.graph(cuda_graph):
    result = mm_fp4(a, b, scales, ...)

For MoE-EP (NCCL-EP) to function correctly, especially on B200/Pre-Nyx hardware, ensure the following:

• NCCL Version: Requires NCCL ≥ 2.30.7. Ensure this version is bound first in your LD_LIBRARY_PATH.
• Multi-node High-Throughput: Set the environment variable NCCL_MNNVL_ENABLE=1. (Single-node intra-tray NVLink works without this).
• Verification: You can verify the backend availability with:

python -c "import nccl.ep; from flashinfer.moe_ep import available_backends; print(available_backends())"

Expected output should include ['nccl_ep', ...].

To benchmark a specific routine, use the flashinfer_benchmark.py script with the --routine flag followed by the routine name and any required parameters.

Example patterns:

• bmm_fp8: Benchmarking Batch Matrix Multiplication with FP8.
• BatchPrefillWithRaggedKVCacheWrapper: Benchmarking non-paged (ragged) prefill.
• rmsnorm_quant: Benchmarking RMSNorm with FP8 quantized output.
• mxfp8_quantize: Benchmarking MxFP8 Quantization (requires Blackwell SM10.0+).

Refer to samples/sample_testlist.txt for a list of available routines and example flags.

python3 flashinfer_benchmark.py --routine <routine_name> <flags>