## 1. Project Overview & Quickstart (flashinfer-ai/flashinfer) ## File: README.md

High-Performance GPU Kernels for Inference

[](https://ci.tlcpack.ai/job/flashinfer-ci/job/main/) [](https://github.com/flashinfer-ai/flashinfer/actions/workflows/build-doc.yml) **FlashInfer** is a library and kernel generator for inference that delivers state-of-the-art performance across diverse GPU architectures. It provides unified APIs for attention, GEMM, and MoE operations with multiple backend implementations including FlashAttention-2/3, cuDNN, CUTLASS, and TensorRT-LLM. ## Why FlashInfer? - **State-of-the-art Performance**: Optimized kernels for prefill, decode, and mixed batching scenarios - **Multiple Backends**: Automatically selects the best backend for your hardware and workload - **Modern Architecture Support**: Support for SM75 (Turing) and later (through Blackwell) - **Low-Precision Compute**: FP8 and FP4 quantization for attention, GEMM, and MoE operations - **Production-Ready**: CUDAGraph and torch.compile compatible for low-latency serving ## Core Features ### Attention Kernels - **Paged and Ragged KV-Cache**: Efficient memory management for dynamic batch serving - **Decode, Prefill, and Append**: Optimized kernels for all attention phases - **MLA Attention**: Native support for DeepSeek's Multi-Latent Attention - **Cascade Attention**: Memory-efficient hierarchical KV-Cache for shared prefixes - **Sparse Attention**: Block-sparse and variable block-sparse patterns - **POD-Attention**: Fused prefill+decode for mixed batching ### GEMM & Linear Operations - **BF16 GEMM**: BF16 matrix multiplication for SM10.0+ GPUs. - **FP8 GEMM**: Per-tensor and groupwise scaling - **FP4 GEMM**: NVFP4 and MXFP4 matrix multiplication for Blackwell GPUs - **Grouped GEMM**: Efficient batched matrix operations for LoRA and multi-expert routing ### Mixture of Experts (MoE) - **Fused MoE Kernels** - **Multiple Routing Methods**: DeepSeek-V3, Llama-4, and standard top-k routing - **Quantized MoE**: FP8 and FP4 expert weights with block-wise scaling ### Sampling & Decoding - **Sorting-Free Sampling**: Efficient Top-K, Top-P, and Min-P without sorting - **Speculative Decoding**: Chain speculative sampling support ### Communication - **AllReduce**: Custom implementations - **Multi-Node NVLink**: MNNVL support for multi-node inference - **NVSHMEM Integration**: For distributed memory operations ### Other Operators - **RoPE**: LLaMA-style rotary position embeddings (including LLaMA 3.1) - **Normalization**: RMSNorm, LayerNorm, Gemma-style fused operations - **Activations**: SiLU, GELU with fused gating ## GPU Support | Architecture | Compute Capability | Example GPUs | |--------------|-------------------|------| | Turing | SM 7.5 | T4, RTX 20 series | | Ampere | SM 8.0, 8.6 | A100, A10, RTX 30 series | | Ada Lovelace | SM 8.9 | L4, L40, RTX 40 series | | Hopper | SM 9.0 | H100, H200 | | Blackwell | SM 10.0, 10.3 | B200, B300 | | Blackwell | SM 11.0 | Jetson Thor | | Blackwell | SM 12.0, 12.1 | RTX 50 series, DGX Spark | > **Note:** Not all features are supported across all compute capabilities. ## News Latest: [](https://github.com/flashinfer-ai/flashinfer/releases/latest) Notable updates: - [2025-10-08] Blackwell support added in [v0.4.0](https://github.com/flashinfer-ai/flashinfer/releases/tag/v0.4.0) - [2025-03-10] [Blog Post](https://flashinfer.ai/2025/03/10/sampling.html) Sorting-Free GPU Kernels for LLM Sampling, which explains the design of sampling kernels in FlashInfer. ## Getting Started ### Installation **Quickstart:** ```bash pip install flashinfer-python ``` **Package Options:** - **flashinfer-python**: Core package that compiles/downloads kernels on first use - **flashinfer-cubin**: Pre-compiled kernel binaries for all supported GPU architectures - **flashinfer-jit-cache**: Pre-built kernel cache for specific CUDA versions **For faster initialization and offline usage**, install the optional packages to have most kernels pre-compiled: ```bash pip install flashinfer-python flashinfer install-cubin-wheel flashinfer install-jit-cache-wheel ``` **For Blackwell (SM100+) CuTe DSL kernels**, install with the CUDA 13 extra to enable Blackwell-optimized kernels: ```bash pip install flashinfer-python[cu13] ``` ### Verify Installation ```bash flashinfer show-config ``` ### Basic Usage ```python import torch import flashinfer # Single decode attention q = torch.randn(32, 128, device="cuda", dtype=torch.float16) # [num_qo_heads, head_dim] k = torch.randn(2048, 32, 128, device="cuda", dtype=torch.float16) # [kv_len, num_kv_heads, head_dim] v = torch.randn(2048, 32, 128, device="cuda", dtype=torch.float16) output = flashinfer.single_decode_with_kv_cache(q, k, v) ``` See [documentation](https://docs.flashinfer.ai/) for comprehensive API reference and tutorials. ### Install from Source ```bash git clone https://github.com/flashinfer-ai/flashinfer.git --recursive cd flashinfer python -m pip install -v . ``` **For development**, install in editable mode: ```bash python -m pip install --no-build-isolation -e . -v ``` > **Note:** When using `--no-build-isolation`, pip does not automatically install build dependencies. FlashInfer requires `setuptools>=77`. If you encounter an error like `AttributeError: module 'setuptools.build_meta' has no attribute 'prepare_metadata_for_build_editable'`, upgrade pip and setuptools first: > ```bash > python -m pip install --upgrade pip setuptools > ``` Build optional packages: ```bash # flashinfer-cubin cd flashinfer-cubin python -m build --no-isolation --wheel python -m pip install dist/*.whl ``` ```bash # flashinfer-jit-cache (customize for your target GPUs) export FLASHINFER_CUDA_ARCH_LIST="7.5 8.0 8.9 9.0a 10.0a 10.3a 10.7a 11.0a 12.0f" cd flashinfer-jit-cache python -m build --no-isolation --wheel python -m pip install dist/*.whl ``` For more details, see the [Install from Source documentation](https://docs.flashinfer.ai/installation.html#install-from-source). ### Nightly Builds ```bash pip install -U --pre flashinfer-python --index-url https://flashinfer.ai/whl/nightly/ --no-deps pip install flashinfer-python # Install dependencies from PyPI flashinfer install-cubin-wheel --nightly flashinfer install-jit-cache-wheel --nightly ``` ### CLI Tools FlashInfer provides several CLI commands for configuration, module management, and development: ```bash # Verify installation and view configuration flashinfer show-config # List and inspect modules flashinfer list-modules flashinfer module-status # Manage artifacts and cache flashinfer download-cubin flashinfer install-cubin-wheel flashinfer install-jit-cache-wheel flashinfer download-kernels flashinfer clear-cache # For developers: generate compile_commands.json for IDE integration flashinfer export-compile-commands [output_path] ``` For complete documentation, see the [CLI reference](https://docs.flashinfer.ai/cli.html). ## API Logging FlashInfer provides comprehensive API logging for debugging. Enable it using environment variables: ```bash # Enable logging (levels: 0=off (default), 1=basic, 3=detailed, 5=statistics) export FLASHINFER_LOGLEVEL=3 # Set log destination (stdout (default), stderr, or file path) export FLASHINFER_LOGDEST=stdout ``` For detailed information about logging levels, configuration, and advanced features, see [Logging](https://docs.flashinfer.ai/logging.html) in our documentation. ## Custom Attention Variants Users can customize their own attention variants with additional parameters. For more details, refer to our [JIT examples](https://github.com/flashinfer-ai/flashinfer/blob/main/tests/utils/test_jit_example.py). ## CUDA Support **Supported CUDA Versions:** 12.6, 12.8, 13.0, 13.1 > **Note:** FlashInfer strives to follow PyTorch's supported CUDA versions plus the latest CUDA release. ## Adoption FlashInfer powers inference in: - [SGLang](https://github.com/sgl-project/sglang) - [vLLM](https://github.com/vllm-project/vllm) - [TensorRT-LLM](https://github.com/NVIDIA/TensorRT-LLM) - [TGI (Text Generation Inference)](https://github.com/huggingface/text-generation-inference) - [MLC-LLM](https://github.com/mlc-ai/mlc-llm) - [LightLLM](https://github.com/ModelTC/lightllm) - [lorax](https://github.com/predibase/lorax) - [ScaleLLM](https://github.com/vectorch-ai/ScaleLLM) ## Acknowledgement FlashInfer is inspired by [FlashAttention](https://github.com/dao-AILab/flash-attention/), [vLLM](https://github.com/vllm-project/vllm), [stream-K](https://arxiv.org/abs/2301.03598), [CUTLASS](https://github.com/nvidia/cutlass), and [AITemplate](https://github.com/facebookincubator/AITemplate). ## Citation If you find FlashInfer helpful in your project or research, please consider citing our [paper](https://arxiv.org/abs/2501.01005): ```bibtex @article{ye2025flashinfer, title = {FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving}, author = { Ye, Zihao and Chen, Lequn and Lai, Ruihang and Lin, Wuwei and Zhang, Yineng and Wang, Stephanie and Chen, Tianqi and Kasikci, Baris and Grover, Vinod and Krishnamurthy, Arvind and Ceze, Luis }, journal = {arXiv preprint arXiv:2501.01005}, year = {2025}, url = {https://arxiv.org/abs/2501.01005} } ``` --- ## File: benchmarks/README.md # FlashInfer Perf Benchmarking Framework -- `flashinfer_benchmark.py` The aim of `flashinfer_benchmark.py` is to provide a single framework for benchmarking any FlashInfer kernel and replace standalone benchmarking scripts. `bench_recurrent_kda_prefill.py --case-set h12` runs the six Kimi-K3 TP8 H12 public-API cases. `reduce_kda_h12.py` combines successful SM100a and SM103a result files without producing a cross-shape aggregate. The benchmark defaults to the natural device/shape dispatcher and records its resolved module; use `--candidate-route nonpersistent` for a B200 direct-family route A/B. ## Overview This framework provides tools to: - Benchmark FlashInfer's Attention, GEMM, MOE, Norm, Quantization, Sampling, RoPE, Mamba, and GDN API performance from different kernel backends such as FlashAttention2/3, cuDNN, cuBLAS, CUTLASS, PrimTS, CuTe-DSL, TensorRT-LLM, and Triton - Compare performance across different configurations - Batch performance test multiple test cases Currently supports testing attention, gemm, fused MOE, normalization, quantization, sampling, RoPE, Mamba, and GDN (Gated Delta Net) APIs: - Attention: - `BatchDecodeWithPagedKVCacheWrapper` - Decode attention with paged KV cache. - Also supports computationally similar `cudnn_batch_decode_with_kv_cache` and `trtllm_batch_decode_with_kv_cache`. - Speculative decode is supported by setting `--s_qo > 1` (subject to backend limitations noted below). - `BatchPrefillWithPagedKVCacheWrapper` - Prefill attention with paged KV cache. - Also supports computationally similar `cudnn_batch_prefill_with_kv_cache` and `trtllm_batch_context_with_kv_cache`. - `BatchPrefillWithRaggedKVCacheWrapper` - Prefill attention with ragged KV cache. - Also supports computationally similar `cudnn_batch_prefill_with_kv_cache` (cudnn-native) and `trtllm_ragged_attention_deepseek`. - `BatchMLAPagedAttentionWrapper` - MLA attention proposed in DeepSeek series of models. - Also supports computationally similar `trtllm_batch_decode_with_kv_cache_mla` (trtllm-native) and CuTe DSL MLA decode kernel (cute-dsl, SM100+). - All four attention routines accept `--backends prims-ts` on SM100/SM103 to benchmark the experimental task-scheduled attention implementation. `prims_ts` is accepted as an alias. - GEMM: - `gemm_fp8_nt_groupwise` - GEMM with FP8 data types using groupwise scaling. - `group_gemm_fp8_nt_groupwise` - Group GEMM with FP8 data types using groupwise scaling. - `bmm_fp8` - Batched matrix multiplication with FP8 inputs. - `mm_mxfp8` - Dense MXFP8 matrix multiplication. - `mm_fp8` - Matrix multiplication with FP8 inputs using the trtllm-gen low-latency GEMM (Blackwell SM10.0+, small-M optimized, pre-shuffled weights). - `mm_fp4` - Matrix multiplication with NVFP4 inputs. - `mm_bf16` - Matrix multiplication with BF16 inputs (Blackwell SM10.0+). - `bmm_bf16` - Batched matrix multiplication with BF16 inputs (Blackwell SM10.0+). - MOE: - `trtllm_fp4_block_scale_moe` - MOE with FP4 quantized weights and block-wise scaling. - `trtllm_fp8_block_scale_moe` - MOE with FP8 quantized weights and block-wise scaling. - `trtllm_fp8_per_tensor_scale_moe` - MOE with FP8 quantized weights and per-tensor scaling. - `cutlass_fused_moe` - CUTLASS fused MoE (base/fp8/nvfp4 variants with optional TP/EP) - MOE Communication: - `moe_a2a_dispatch_combine` - MoE All-to-All dispatch + combine benchmark for multi-GPU expert-parallel inference. Requires `mpirun` for multi-GPU execution. Supports optional quantization (FP8, NVFP4, FP8 block-scale) and real MoE kernel computation. - AllReduce Communication: - `allreduce_fusion` - AllReduce fusion benchmark for multi-GPU inference. Requires `mpirun` for multi-GPU execution. Supports TRTLLM and TRTLLM MNNVL backends with multiple fusion patterns (plain allreduce, allreduce + residual + RMSNorm). - Norm: - `rmsnorm` - Root Mean Square Layer Normalization. - `fused_add_rmsnorm` - Fused residual add + RMSNorm. - `gemma_rmsnorm` - Gemma-style RMSNorm using `(weight + 1)`. - `gemma_fused_add_rmsnorm` - Gemma-style fused residual add + RMSNorm. - `rmsnorm_quant` - RMSNorm with FP8 quantized output. - `fused_add_rmsnorm_quant` - Fused residual add + RMSNorm with FP8 quantized output. - `rmsnorm_fp4quant` - RMSNorm with FP4 quantized output (CuTe-DSL, Blackwell SM10.0+). - `add_rmsnorm_fp4quant` - Fused residual add + RMSNorm with FP4 quantized output (CuTe-DSL, Blackwell SM10.0+). - Quantization: - `mxfp8_quantize` - Quantize tensor to MxFP8 format (Blackwell SM10.0+). - `mxfp4_quantize` - Quantize tensor to MxFP4 format (Blackwell SM10.0+). - `nvfp4_quantize` - Quantize tensor to NVFP4 format with configurable scale factor layout (Blackwell SM10.0+). - `nvfp4_batched_quantize` - Batched NVFP4 quantization (Blackwell SM10.0+). - Sampling: - `softmax` - Softmax with optional temperature scaling. - `sampling_from_probs` - Sample token indices from probability distributions. - `sampling_from_logits` - Sample token indices from logits (fused softmax + sampling). - `top_k_sampling_from_probs` - Top-K sampling from probabilities. - `top_p_sampling_from_probs` - Top-P (nucleus) sampling from probabilities. - `top_k_top_p_sampling_from_probs` - Combined Top-K and Top-P sampling from probabilities. - `top_k_top_p_sampling_from_logits` - Combined Top-K and Top-P sampling from logits. - `min_p_sampling_from_probs` - Min-P sampling from probabilities. - `top_k_renorm_probs` - Renormalize probabilities after Top-K filtering. - `top_p_renorm_probs` - Renormalize probabilities after Top-P filtering. - `top_k_mask_logits` - Mask logits outside Top-K values. - `chain_speculative_sampling` - Chain speculative sampling for speculative decoding. - `top_k` - Radix-based Top-K selection. - `top_k_page_table_transform` - Fused Top-K with page table lookup. - `top_k_ragged_transform` - Fused Top-K with ragged index transform. - RoPE (Rotary Positional Embeddings): - `apply_rope` - Apply RoPE with indptr/offsets. - `apply_rope_pos_ids` - Apply RoPE with position IDs. - `apply_llama31_rope` - Apply Llama 3.1 style RoPE with indptr/offsets. - `apply_llama31_rope_pos_ids` - Apply Llama 3.1 style RoPE with position IDs. - `apply_rope_with_cos_sin_cache` - Apply RoPE with precomputed cos/sin cache. - `mla_rope_quantize_fp8` - MLA RoPE with FP8 quantization (SM8.9+). - `rope_quantize_fp8` - RoPE with FP8 quantization (SM8.9+). - `rope_quantize_fp8_append_paged_kv_cache` - RoPE with FP8 quantization and paged KV cache append (SM8.9+). - Mamba (Selective State Space Models): - `selective_state_update` - Selective state update for Mamba layers (generation phase). Supports both single-token prediction (STP) and multi-token prediction (MTP) via `--cache_steps`. Backends: `flashinfer` (CUDA, architecture-specific kernels for base/SM90/SM100+) and `triton` (reference). - GDN (Gated Delta Net linear attention, SM90+): - `gated_delta_rule_decode` - Single-token (T=1) gated delta rule decode. `--state_layout` selects between `gated_delta_rule_decode_pretranspose` ([B, HV, V, K] state, default) and `gated_delta_rule_decode` ([B, HV, K, V] state). `--state_dtype bfloat16` selects the BF16 state kernels (head_size=128, pretranspose only). Backends: `flashinfer` (CuTe-DSL) and `triton` (reference). - `gated_delta_rule_mtp` - Multi-token (T>=2) gated delta rule for speculative-decoding verification, with a state pool + indices. `--state_dtype float32` uses `gated_delta_rule_mtp`; `--state_dtype bfloat16` uses the BF16 MTP kernel via `gated_delta_rule_decode_pretranspose`. Backends: `flashinfer`, `triton`. - `chunk_gated_delta_rule` - Chunked GDN prefill over varlen sequences (uniform per-sequence length `--s_qo`). Backends: `flashinfer` (SM90 C++ / SM100 CuTe-DSL) and `fla` (flash-linear-attention Triton baseline, perf-only). ## Quick Start ### Single Test Run A test case is generally invoked as `python3 flashinfer_benchmark.py --routine `. *See samples in samples/sample_testlist.txt for various example test flags.* Example commands and outputs areas follows ``` /* Detailed source-code truncated for AI context efficiency. */ ``` ### Batch Testing Run multiple tests from a file and save results: ```bash python3 flashinfer_benchmark.py --testlist samples/sample_testlist.txt --output_path samples/sample_testlist_output.csv ``` See `samples/sample_testlist.txt` for an example stdout output from the above command; `samples/sample_testlist_output.csv` for csv output from the same run. The output CSV will contain detailed metrics including: - Median execution time - Standard deviation - TFLOPS/sec - Memory throughput (TB/sec) - Input flags - Reproducer commands if `--generate_repro_command` is provided ## Command Line Arguments ### General Flags | Flag | Description | |--------------------------|-------------------------------------------------------------------------------------------------------------| | `--routine` | Test routine to run. See [Overview](#overview) for full list including attention, GEMM, MOE, norm, and quantization routines. | | `--num_iters` | Number of iterations for performance measurement | | `--dry_run_iters` | Number of warmup iterations | | `--no_cuda_graph` | Disable CUDA graph to execute kernels outside of the graph. | | `--use_cupti` | Use CUPTI for timing GPU kernels when available. | | `--refcheck` | Verify outputs match between different backends | | `--allow_output_mismatch`| Continue testing even if outputs don't pass refcheck | | `--random_seed` | Random seed for reproducibility | | `--output_path` | Path to save CSV results | | `--testlist` | Path to a file containing a list of test cases to run in batch mode | | `--verbose`, `-v` | Print additional information (can be used multiple times for more verbosity, e.g. `-vv`) | | `--case_tag` | Optional tag for the test case, useful for annotating or filtering results in the output CSV. | | `--generate_repro_command`| If set, prints a reproducer command for the test case and stores it in the output CSV. | | `--backends` | Space-separated list of backends to test, e.g. fa2, fa2_tc, fa3, auto, cudnn, cudnn-native, cutlass, trtllm, trtllm-gen, trtllm-native, prims-ts, cute-dsl, cublas, trtllm_low_latency. (`prims_ts` aliases `prims-ts`; `auto` support is routine-dependent.)| ### Attention Flags | Flag | Description | |--------------------------|-------------------------------------------------------------------------------------------------------------| | `--page_size` | Page size for paged attention. Required for paged attention tests. | | `--batch_size` | Number of sequences to process in parallel | | `--s_qo` | Query/output sequence length. For decode, `1` is standard decode and `>1` enables speculative decode on supported backends. | | `--s_kv` | Key/value sequence length (context length) | | `--num_qo_heads` | Number of query/output attention heads | | `--num_kv_heads` | Number of key/value attention heads | | `--head_dim_qk` | Head dimension for Q/K. Backend-dependent; PrimTS supports 64/128/256 for FMHA decode and 128/256 for FMHA context. | | `--head_dim_vo` | Head dimension for V/O. Usually equals head_dim_qk. | | `--head_dim_ckv` | Head dimension for C/K/V (MLA attention). | | `--head_dim_kpe` | Head dimension for KPE (MLA attention). | | `--q_dtype` | Data type for the query tensor. Default: bfloat16. Supports float16, bfloat16, fp8_e4m3, and fp8_e5m2 where the selected backend permits them. | | `--kv_dtype` | Data type for the key and value tensors. Default: bfloat16. Supports float16, bfloat16, fp8_e4m3, and fp8_e5m2 where the selected backend permits them. | | `--out_dtype` | Data type for the output tensor. Default: same as q_dtype. Backend-dependent; PrimTS context accepts bfloat16, float16, or fp8_e4m3, while PrimTS FP8 decode accepts float16 or fp8_e4m3. FP8 ragged comparisons with non-PrimTS backends require bfloat16 or float16. | | `--causal` | Use causal attention masking for context/prefill. Multi-query FMHA and MLA decode use bottom-right causal masking automatically. | | `--random_actual_seq_len`| Use random sequence lengths up to max length. If False, use max length. | ### GEMM Flags | Flag | Description | |--------------------------|-------------------------------------------------------------------------------------------------------------| | `--m` | Number of rows of matrix A and output matrix (GEMM M dimension) | | `--n` | Number of columns of matrix B and output matrix (GEMM N dimension) | | `--k` | Number of columns of matrix A / rows of matrix B (GEMM K dimension) | | `--tile_size` | Tile size for the GEMM operation (affects performance and scaling) | | `--group_size` | Number of groups for group GEMM (batching multiple GEMMs together) | | `--scale_major_mode` | Layout for FP8 scaling: `MN` (per output tile) or `K` (per input tile) | | `--out_dtype` | Output data type: `bfloat16` or `float16` | | `--mma_sm` | Number of SMs to use for the MMA operation (1 or 2) | | `--input_dtype` | Data type for input matrix (for FP8 GEMM, e.g. `fp8_e4m3`) | | `--mat2_dtype` | Data type for second matrix (for FP8 GEMM, e.g. `fp8_e4m3`) | | `--use_128x4_sf_layout` | Use 128x4 scale/format layout for FP4 GEMM (for `mm_fp4` routine) | | `--use_nvfp4` | Whether to use nvfp4 quantization or mxfp4 quantization, defaults to False.(for `mm_fp4` routine) | | `--autotune` | Enable autotune for supported operation (`mm_fp4`, `bmm_fp8`, `mm_fp8`, `bmm_mxfp8`, `mm_mxfp8`, `mm_bf16`, `bmm_bf16` routines) | | `--bias` | Use bias for `mm_bf16` (Enabled for TGV backend) | ### MOE Flags | Flag | Description | |--------------------------|-------------------------------------------------------------------------------------------------------------| | `--num_tokens` | Number of input tokens | | `--hidden_size` | Hidden dimension size | | `--intermediate_size` | Intermediate dimension size (FF layer dimension) | | `--num_experts` | Total number of experts | | `--top_k` | Number of experts to route to per token | | `--n_group` | Number of expert groups (for DeepSeek routing). Default: 1 | | `--topk_group` | Number of groups to consider for top-k routing. Default: 1 | | `--routed_scaling_factor`| Scaling factor for routing. Default: 2.5 | | `--local_expert_offset` | Offset of local experts in global expert space. Default: 0 | | `--local_num_experts` | Number of experts handled by this device. Default: equals num_experts | | | `--routing_method` | Routing method: `renormalize`, `deepseek_v3`, `llama4`, `renormalize_naive`. Default: `deepseek_v3`. | | `--use_shuffled_weight` | Whether to use shuffled weight layout | | `--weight_layout` | Weight layout: 0=MajorK, 1=MajorMn, 2=BlockMajorK. Default: 0 | | `--use_routing_bias` | Whether to use routing bias | | `--use_routing_scales_on_input` | Whether to use routing scales on input (for Llama4 routing) | | `--input_dtype` | Data type of the input hidden states. Default: bfloat16 | | `--weight_dtype` | Data type of the weights (before quantization). Default: bfloat16 | | `--cutlass_variant` | CUTLASS MoE variant: `base` (no quant), `fp8` (per-tensor FP8), `nvfp4` (FP4 block-scale) | | `--quantized_input` | For `nvfp4` only: quantize input activations to FP4 | | `--tp_size` | Tensor-parallel world size | | `--tp_rank` | Tensor-parallel rank | | `--ep_size` | Expert-parallel world size | | `--ep_rank` | Expert-parallel rank | | `--activation-type` | Activation function: `Swiglu` (default), `Geglu`, `SwigluStep` (clipped SwiGLU, limit=7.0), `Relu2`, etc. | | `--autotune` | Enable autotune for supported operation | ### MOE Routing Method Compatibility | Routing Method | Requirements | Compatible MOE Types | |------------------------|--------------|---------------------| | **deepseek_v3** | `top_k <= 8`, `topk_group <= 4`, requires `--n_group`, `--topk_group`, `--routed_scaling_factor`, `--use_routing_bias` | FP4, FP8 Block Scale | | **renormalize** | `top_k == 1` for FP8 Block Scale, `top_k <= 8` for FP4. Do NOT use `--n_group` or `--topk_group` | All MOE types | | **llama4** | `top_k == 1`, requires `--routed_scaling_factor`, `--use_routing_bias`, `--use_routing_scales_on_input`. Do NOT use `--n_group` or `--topk_group` | FP8 Per-Tensor | | **renormalize_naive** | `top_k == 1` for FP8 Block Scale, `top_k <= 8` for FP4. Do NOT use `--n_group` or `--topk_group` | FP4 primarily | Notes: - Group parameters (`--n_group`, `--topk_group`) are ONLY used with DeepSeekV3 routing method. Using them with other routing methods will cause the error: "Routing kernel with groups implies DeepSeekV3 routing method." - Different MOE kernel implementations have different `top_k` constraints. FP8 MOE kernels (both Block Scale and Per-Tensor) have stricter limits than FP4 for non-DeepSeekV3 routing methods. - FP8 MOE kernels require integer values for group parameters, while FP4 MOE kernels accept optional values. - CUTLASS fused MoE (`cutlass_fused_moe`) ignores `--routing_method`, `--n_group`, and `--topk_group`; it computes routing via softmax+top-k internally from the provided logits. ### MoE Communication Flags (moe_a2a_dispatch_combine) The `moe_a2a_dispatch_combine` routine benchmarks MoE All-to-All communication for multi-GPU expert-parallel inference. It must be launched with `mpirun`. | Flag | Description | |--------------------------|-------------------------------------------------------------------------------------------------------------| | `--num_tokens` | Number of tokens per rank (local batch size) | | `--hidden_size` | Hidden dimension size | | `--num_experts` | Total number of experts across all ranks | | `--top_k` | Number of experts to route each token to | | `--input_dtype` | Data type for hidden states payload: `bfloat16` (default) or `float16` | | `--quant_dtype` | Quantization format: `fp8` (per-tensor), `nvfp4` (block-scale FP4), `fp8_block_scale` (block-scale FP8) | | `--real_math` | Run actual MoE kernels instead of fake computation. Requires `--intermediate_size` and `--quant_dtype` to be `nvfp4` or `fp8_block_scale` | | `--intermediate_size` | Intermediate FFN size. Required if `--real_math` is set | | `--max_num_tokens` | Max tokens per rank for workspace allocation. Defaults to `--num_tokens` | | `--validate` | Run correctness validation before benchmarking using deterministic fake MoE | | `--per_phase_timing` | Enable per-phase timing (dispatch/combine/moe_kernel). Adds slight overhead from CUDA events | | `--nvtx` | Enable NVTX markers for Nsight Systems profiling | | `--use_lora` | Carry a per-token int32 LoRA adapter ID through dispatch as an extra payload. | **Launch Examples:** ```bash # Basic (no quantization) mpirun -np 8 python benchmarks/flashinfer_benchmark.py \ --routine moe_a2a_dispatch_combine \ --num_tokens 1024 --hidden_size 7168 --num_experts 256 --top_k 8 # With FP8 quantization mpirun -np 8 python benchmarks/flashinfer_benchmark.py \ --routine moe_a2a_dispatch_combine \ --num_tokens 1024 --hidden_size 7168 --num_experts 256 --top_k 8 \ --quant_dtype fp8 # With NVFP4 quantization and real MoE kernel mpirun -np 8 python benchmarks/flashinfer_benchmark.py \ --routine moe_a2a_dispatch_combine \ --num_tokens 1024 --hidden_size 7168 --num_experts 256 --top_k 8 \ --quant_dtype nvfp4 --real_math --intermediate_size 18432 # With validation and per-phase timing mpirun -np 8 python benchmarks/flashinfer_benchmark.py \ --routine moe_a2a_dispatch_combine \ --num_tokens 1024 --hidden_size 7168 --num_experts 256 --top_k 8 \ --validate --per_phase_timing # Multi-tenant LoRA: carry per-token adapter ID through dispatch mpirun -np 8 python benchmarks/flashinfer_benchmark.py \ --routine moe_a2a_dispatch_combine \ --num_tokens 2048 --hidden_size 7168 --num_experts 256 --top_k 8 \ --use_lora --validate ``` ### AllReduce Communication Flags (allreduce_fusion) The `allreduce_fusion` routine benchmarks AllReduce fusion operations for multi-GPU inference. It must be launched with `mpirun`. Both oneshot and twoshot strategies are benchmarked automatically and reported side by side. | Flag | Description | |--------------------------|-------------------------------------------------------------------------------------------------------------| | `--num_tokens` | Number of tokens (rows) in the input tensor. Default: 64 | | `--hidden_size` | Hidden dimension size. Default: 4096 | | `--input_dtype` | Data type for input tensors: `bfloat16` (default) or `float16` | | `--ar_backend` | AllReduce backend: `auto` (default), `trtllm`, or `mnnvl`. `auto` uses heuristic | | `--pattern` | Fusion pattern: `allreduce` (default) or `ar_residual_rmsnorm` (AllReduce + Residual + RMSNorm) | | `--validate` | Run correctness validation before benchmarking | **Launch Examples:** ```bash # Basic allreduce with auto backend mpirun -np 8 python benchmarks/flashinfer_benchmark.py \ --routine allreduce_fusion \ --num_tokens 64 --hidden_size 4096 # With specific backend mpirun -np 8 python benchmarks/flashinfer_benchmark.py \ --routine allreduce_fusion \ --num_tokens 64 --hidden_size 4096 \ --ar_backend mnnvl # AllReduce + Residual + RMSNorm fusion mpirun -np 8 python benchmarks/flashinfer_benchmark.py \ --routine allreduce_fusion \ --num_tokens 64 --hidden_size 4096 \ --pattern ar_residual_rmsnorm # With validation mpirun -np 8 python benchmarks/flashinfer_benchmark.py \ --routine allreduce_fusion \ --num_tokens 64 --hidden_size 4096 \ --validate ``` ### Norm Flags | Flag | Description | |--------------------------|-------------------------------------------------------------------------------------------------------------| | `--batch_size` | Batch size (number of sequences) | | `--hidden_size` | Hidden dimension size | | `--num_heads` | Number of heads for 3D input shape (batch, num_heads, hidden_size). Optional; if not set, uses 2D shape. | | `--input_dtype` | Input data type: `bfloat16` (default) or `float16` | | `--eps` | Epsilon for numerical stability. Default: 1e-6 | | `--enable_pdl` | Enable programmatic dependent launch | | `--scale` | Scale factor for FP8 quantization (used by `rmsnorm_quant`, `fused_add_rmsnorm_quant`). Default: 1.0 | | `--out_dtype` | Output dtype: `fp8_e4m3`, `fp8_e5m2` (for FP8 quant); `nvfp4`, `mxfp4` (for FP4 quant). Default: `fp8_e4m3`| | `--use_global_scale` | Use global scale factor for NVFP4 format (FP4 routines only) | | `--is_sf_swizzled_layout`| Use swizzled scale factor layout for tensor core GEMM (FP4 routines only) | | `--backends` | Backend to test. Defaults to `cute-dsl` for rmsnorm/rmsnorm_quant/fused_add_rmsnorm/fused_add_rmsnorm_quant/gemma_rmsnorm/gemma_fused_add_rmsnorm/rmsnorm_fp4quant/add_rmsnorm_fp4quant (CuTe-DSL kernels) and `cuda` otherwise. Pass `--backends cuda` to force the CUDA JIT fallback (set `FLASHINFER_USE_CUDA_NORM=1` to actually run the CUDA path). | ### Quantization Flags | Flag | Description | |--------------------------|-------------------------------------------------------------------------------------------------------------| | `--m` | Number of rows in input tensor | | `--k` | Number of columns in input tensor (must be divisible by 32) | | `--input_dtype` | Input data type: `bfloat16` (default) or `float16` | | `--is_sf_swizzled_layout`| Use swizzled layout for scale factors. Default: True | | `--no_sf_swizzled_layout`| Disable swizzled layout for scale factors | | `--alignment` | sfVecSize for quantization. Default: 32 | | `--enable_pdl` | Enable programmatic dependent launch | | `--batch_size` | Batch size for batched quantization (`nvfp4_batched_quantize` only) | | `--global_scale` | Global scale factor for NVFP4 quantization. Default: 1.0 | | `--sf_layout` | Scale factor layout for FP4 quantization: `128x4` (default), `8x4`, or `linear` | | `--do_shuffle` | Shuffle scale factors for TRTLLM backend (`nvfp4_quantize` only) | | `--sf_vec_size` | Scale factor vector size for NVFP4 quantization. Default: 16 | | `--backends` | Backend to test. Default: `cuda` | ### Sampling Flags | Flag | Description | |--------------------------|-------------------------------------------------------------------------------------------------------------| | `--batch_size` | Batch size (number of sequences) | | `--vocab_size` | Vocabulary size | | `--input_dtype` | Input data type for logits: `float32` (default), `float16`, or `bfloat16` | | `--top_k` | Top-K value for top-k sampling. Default: 50 | | `--top_p` | Top-P threshold for top-p (nucleus) sampling. Default: 0.9 | | `--min_p` | Min-P threshold for min-p sampling. Default: 0.1 | | `--temperature` | Temperature for softmax. Default: 1.0 | | `--filter_apply_order` | Order of applying top-k and top-p filters: `top_k_first` (default) or `joint` | | `--num_speculate_tokens` | Number of speculative tokens for chain speculative sampling. Default: 5 | | `--max_len` | Max sequence length for `top_k_page_table_transform` and `top_k_ragged_transform`. Default: 4096 | | `--num_rows` | Number of rows for `top_k_page_table_transform` and `top_k_ragged_transform`. Defaults to batch_size | | `--backends` | Backend to test: `cuda` (default) | ### RoPE Flags | Flag | Description | |--------------------------|-------------------------------------------------------------------------------------------------------------| | `--batch_size` | Batch size (number of sequences) | | `--seq_len` | Sequence length (qkv_len or kv_len) | | `--num_qo_heads` | Number of query/output heads | | `--num_kv_heads` | Number of key/value heads | | `--head_dim` | Head dimension | | `--rotary_dim` | Rotary dimension (defaults to head_dim if not specified) | | `--no_rope_dim` | Number of dimensions without RoPE (for MLA). Default: 0 | | `--input_dtype` | Input data type: `float16` (default) or `bfloat16` | | `--quant_dtype` | Quantized data type for FP8 routines: `fp8_e4m3` (default) or `fp8_e5m2` | | `--rope_scale` | RoPE scaling factor. Default: 1.0 | | `--rope_theta` | RoPE theta base frequency. Default: 10000.0 | | `--interleave` | Use interleaved rotary embedding (GPT-J style) | | `--page_size` | Page size for paged KV cache. Default: 16 | | `--kv_layout` | KV cache layout: `NHD` (default) or `HND` | | `--low_freq_factor` | Low frequency factor for Llama 3.1 RoPE. Default: 1.0 | | `--high_freq_factor` | High frequency factor for Llama 3.1 RoPE. Default: 4.0 | | `--old_context_len` | Old context length for Llama 3.1 RoPE. Default: 8192 | | `--backends` | Backend to test: `cuda` (default) | ### Mamba Flags | Flag | Description | |--------------------------|-------------------------------------------------------------------------------------------------------------| | `--batch_size` | Batch size (number of sequences) | | `--nheads` | Number of SSM heads | | `--dim` | Head dimension (headdim) | | `--dstate` | SSM state size | | `--ngroups` | Number of groups for B and C matrices. `nheads` must be divisible by `ngroups`, and `nheads/ngroups` must be 1, 8, or 16. Default: 8 | | `--cache_steps` | Number of steps/tokens for multi-token prediction (MTP). 0 = single-token prediction (STP). Default: 0 | | `--input_dtype` | Data type for input tensors (x, B, C, z): `bfloat16` (default). Only `bfloat16` is supported. | | `--state_dtype` | Data type for the SSM state cache: `bfloat16` (default), `float16`, or `float32` | | `--weight_dtype` | Data type for weight tensors (dt, D, dt_bias): `float32` (default) or `bfloat16` | | `--has_z` | Include z tensor for gating (`z * sigmoid(z)` applied to output) | | `--dt_softplus` | Apply softplus to dt before use | | `--backends` | Backends to test: `flashinfer` (default), `triton` (reference). Refcheck compares against Triton reference | ### GDN Flags Applies to `gated_delta_rule_decode`, `gated_delta_rule_mtp`, and `chunk_gated_delta_rule` (SM90+). | Flag | Description | |-------------------------------|-------------------------------------------------------------------------------------------------------------| | `--batch_size` | Decode/MTP: number of concurrent requests. Prefill: number of sequences | | `--num_q_heads` | Number of query heads. Default: 16 | | `--num_k_heads` | Number of key heads. Default: 16 | | `--num_v_heads` | Number of value heads (GVA when > `num_q_heads`). Default: 32 | | `--head_size` | Head dimension (K = V = head_size). Default: 128 | | `--input_dtype` | Data type for q/k/v/a/b tensors: `bfloat16` (default) or `float16` | | `--state_dtype` | Recurrent state dtype: `float32` (default) or `bfloat16` (BF16 state kernels; decode/MTP, head_size=128, pretranspose) | | `--state_layout` | Decode only: `pretranspose` ([B, HV, V, K], default) or `nontranspose` ([B, HV, K, V]) | | `--pool_mode` | `single` (default, read == write slots) or `split` (pool of 2B; reads slots [0..B), writes [B..2B)) | | `--seq_len` | MTP only: tokens per request (>= 2). Default: 2 | | `--s_qo` | Prefill only: per-sequence length (uniform). Default: 2048 | | `--update_state` | MTP only: write the final state back (`disable_state_update=False`). BF16 state always updates in-place | | `--cache_intermediate_states` | MTP with `float32` state only: cache per-token intermediate states | | `--no_qk_l2norm` | Decode/MTP: disable in-kernel Q/K L2 normalization | | `--backends` | Decode/MTP: `flashinfer` (default), `triton`. Prefill: `flashinfer` (default), `fla` (requires `pip install flash-linear-attention`; perf-only, excluded from refcheck) | Notes: - Refcheck compares against the torch reference in `tests/gdn/reference_delta_rule.py`. - Prefill pre-L2-normalizes k and calls the kernel with `use_qk_l2norm_in_kernel=False` so the kernel and reference see identical inputs. ## `flashinfer_benchmark.py` Routine & Backend Support Matrix The following table summarizes the support surface of each routine & backend's on various [CUDA Compute Capabilities](https://developer.nvidia.com/cuda-gpus). Each column represents a compute capability. Backends inside cells represent supported backends. A blank cell means no backend is supported for that routine at that compute capability. | Routine | 7.5 | 8.0 | 8.6 | 8.9 | 9.0 | 10.0 | 10.3 | 12.0 | |---------|-----|-----|-----|-----|-----|-------|-------|-------| | **BatchDecodeWithPagedKVCacheWrapper** | fa2 | fa2, fa2_tc, cudnn | fa2, fa2_tc, cudnn | fa2, fa2_tc, cudnn | fa2, fa2_tc, cudnn | fa2, fa2_tc, cudnn, trtllm-gen, trtllm-native, prims-ts | fa2, fa2_tc, cudnn, trtllm-gen, trtllm-native, prims-ts | fa2, fa2_tc, cudnn | | **BatchPrefillWithPagedKVCacheWrapper** | | fa2, cudnn, cudnn-native | fa2, cudnn, cudnn-native | fa2, cudnn, cudnn-native | fa2, fa3, cudnn, cudnn-native | fa2, cudnn, cudnn-native, trtllm-gen, trtllm-native, prims-ts | fa2, cudnn, cudnn-native, trtllm-gen, trtllm-native, prims-ts | fa2, cudnn, cudnn-native | | **BatchPrefillWithRaggedKVCacheWrapper** | | fa2, cudnn, cudnn-native | fa2, cudnn, cudnn-native | fa2, cudnn, cudnn-native | fa2, fa3, cudnn, cudnn-native | fa2, cudnn, cudnn-native, cutlass, trtllm-native, prims-ts | fa2, cudnn, cudnn-native, cutlass, trtllm-native, prims-ts | fa2, cudnn, cudnn-native | | **BatchMLAPagedAttentionWrapper** | | fa2 | fa2 | fa2 | fa2, fa3 | fa2, cutlass, trtllm-native, cute-dsl, prims-ts | fa2, cutlass, trtllm-native, prims-ts | fa2 | | **gemm_fp8_nt_groupwise** | | | | | | cutlass | cutlass | | | **group_gemm_fp8_nt_groupwise** | | | | | | cutlass | cutlass | | | **bmm_fp8** | | | | cudnn, cublas | cudnn, cublas | cudnn, cublas, cutlass | cudnn, cublas, cutlass | cudnn, cublas | | **mm_fp8** | | | | | | trtllm_low_latency | trtllm_low_latency | | | **mm_fp4** | | | | | | cudnn, trtllm, cutlass | cudnn, trtllm, cutlass | cudnn | | **mm_bf16** | | | | | | cudnn, cutlass, tgv | cudnn, cutlass, tgv | | | **bmm_bf16** | | | | | | cudnn, cutlass | cudnn, cutlass | | | **trtllm_fp4_block_scale_moe** | | | | | | trtllm | trtllm | | | **trtllm_fp8_block_scale_moe** | | | | | | trtllm | trtllm | | | **trtllm_fp8_per_tensor_scale_moe** | | | | | | trtllm | trtllm | | | **cutlass_fused_moe** | | | | | | cutlass | cutlass | | | **moe_a2a_dispatch_combine** | | | | | | moe_a2a | moe_a2a | | | **allreduce_fusion** | | | | | | allreduce | allreduce | | | **rmsnorm** | cute-dsl | cute-dsl | cute-dsl | cute-dsl | cute-dsl | cute-dsl | cute-dsl | cute-dsl | | **fused_add_rmsnorm** | cute-dsl | cute-dsl | cute-dsl | cute-dsl | cute-dsl | cute-dsl | cute-dsl | cute-dsl | | **gemma_rmsnorm** | cute-dsl | cute-dsl | cute-dsl | cute-dsl | cute-dsl | cute-dsl | cute-dsl | cute-dsl | | **gemma_fused_add_rmsnorm** | cute-dsl | cute-dsl | cute-dsl | cute-dsl | cute-dsl | cute-dsl | cute-dsl | cute-dsl | | **rmsnorm_quant** | cute-dsl | cute-dsl | cute-dsl | cute-dsl | cute-dsl | cute-dsl | cute-dsl | cute-dsl | | **fused_add_rmsnorm_quant** | cute-dsl | cute-dsl | cute-dsl | cute-dsl | cute-dsl | cute-dsl | cute-dsl | cute-dsl | | **rmsnorm_fp4quant** | | | | | | cute-dsl | cute-dsl | | | **add_rmsnorm_fp4quant** | | | | | | cute-dsl | cute-dsl | | | **mxfp8_quantize** | | | | | | cuda | cuda | | | **mxfp4_quantize** | | | | | | cuda | cuda | | | **nvfp4_quantize** | | | | | | cuda | cuda | | | **nvfp4_batched_quantize** | | | | | | cuda | cuda | | | **softmax** | cuda | cuda | cuda | cuda | cuda | cuda | cuda | cuda | | **sampling_from_probs** | cuda | cuda | cuda | cuda | cuda | cuda | cuda | cuda | | **sampling_from_logits** | cuda | cuda | cuda | cuda | cuda | cuda | cuda | cuda | | **top_k_sampling_from_probs** | cuda | cuda | cuda | cuda | cuda | cuda | cuda | cuda | | **top_p_sampling_from_probs** | cuda | cuda | cuda | cuda | cuda | cuda | cuda | cuda | | **top_k_top_p_sampling_from_probs** | cuda | cuda | cuda | cuda | cuda | cuda | cuda | cuda | | **top_k_top_p_sampling_from_logits** | cuda | cuda | cuda | cuda | cuda | cuda | cuda | cuda | | **min_p_sampling_from_probs** | cuda | cuda | cuda | cuda | cuda | cuda | cuda | cuda | | **top_k_renorm_probs** | cuda | cuda | cuda | cuda | cuda | cuda | cuda | cuda | | **top_p_renorm_probs** | cuda | cuda | cuda | cuda | cuda | cuda | cuda | cuda | | **top_k_mask_logits** | cuda | cuda | cuda | cuda | cuda | cuda | cuda | cuda | | **chain_speculative_sampling** | cuda | cuda | cuda | cuda | cuda | cuda | cuda | cuda | | **top_k** | cuda | cuda | cuda | cuda | cuda | cuda | cuda | cuda | | **top_k_page_table_transform** | cuda | cuda | cuda | cuda | cuda | cuda | cuda | cuda | | **top_k_ragged_transform** | cuda | cuda | cuda | cuda | cuda | cuda | cuda | cuda | | **apply_rope** | cuda | cuda | cuda | cuda | cuda | cuda | cuda | cuda | | **apply_rope_pos_ids** | cuda | cuda | cuda | cuda | cuda | cuda | cuda | cuda | | **apply_llama31_rope** | cuda | cuda | cuda | cuda | cuda | cuda | cuda | cuda | | **apply_llama31_rope_pos_ids** | cuda | cuda | cuda | cuda | cuda | cuda | cuda | cuda | | **apply_rope_with_cos_sin_cache** | cuda | cuda | cuda | cuda | cuda | cuda | cuda | cuda | | **mla_rope_quantize_fp8** | | | | cuda | cuda | cuda | cuda | cuda | | **rope_quantize_fp8** | | | | cuda | cuda | cuda | cuda | cuda | | **rope_quantize_fp8_append_paged_kv_cache** | | | | cuda | cuda | cuda | cuda | cuda | | **selective_state_update** | flashinfer, triton | flashinfer, triton | flashinfer, triton | flashinfer, triton | flashinfer, triton | flashinfer, triton | flashinfer, triton | flashinfer, triton | | **gated_delta_rule_decode** | | | | | flashinfer, triton | flashinfer, triton | flashinfer, triton | triton | | **gated_delta_rule_mtp** | | | | | flashinfer, triton | flashinfer, triton | flashinfer, triton | triton | | **chunk_gated_delta_rule** | | | | | flashinfer, fla | flashinfer, fla | flashinfer, fla | | Backend Legend: - fa2: FlashAttention2 - fa2_tc: FlashAttention2 (with Tensor Cores for `BatchDecodeWithPagedKVCacheWrapper`) - fa3: FlashAttention-3 - cublas: cuBLAS - cudnn: cuDNN (via wrapper API) - cudnn-native: cuDNN (direct API call) - cutlass: CUTLASS - tgv: TGV - trtllm: TensorRT-LLM - trtllm-gen: TensorRT-LLM - trtllm-native: TensorRT-LLM (out-of-wrapper) - prims-ts: Experimental task-scheduled attention kernels (Blackwell SM100/SM103) - cuda: FlashInfer CUDA kernels - cute-dsl: FlashInfer CuTe-DSL kernels (Blackwell SM10.0+) - moe_a2a: MoE All-to-All communication (requires mpirun, Blackwell SM10.0+ with MNNVL) - allreduce: AllReduce fusion communication (requires mpirun, Blackwell SM10.0+ with MNNVL) - triton: Triton reference kernels (used for Mamba selective_state_update and GDN decode/MTP) - fla: flash-linear-attention Triton kernels (GDN prefill baseline) --- ## File: docs/api/activation.rst .. _apiactivation: flashinfer.activation ===================== .. currentmodule:: flashinfer.activation This module provides a set of activation operations for up/gate layers in transformer MLPs. Up/Gate output activation ------------------------- .. autosummary:: :toctree: ../generated silu_and_mul gelu_tanh_and_mul gelu_and_mul silu_and_mul_scaled_nvfp4_experts_quantize --- ## File: docs/api/attention.rst .. _apiattention: FlashInfer Attention Kernels ~~~~~~~~~~~~~~~~~~~~~~~~~~~~ Experimental Task-Scheduled Attention ===================================== The experimental Blackwell task-scheduled FMHA context, FMHA decode, and MLA decode APIs are imported from ``flashinfer.attention.prims_ts``. Scheduling, tile selection, and split-KV reduction are automatic implementation details; there are no public tuning knobs. See the `PrimTS guide index `_ for the public entry points, supported contracts, and examples. Current accuracy and performance signoff is on SM100a/B200; SM103a/B300 is architecture-gated but not yet signoff-qualified. .. currentmodule:: flashinfer.attention.prims_ts FMHA Context/Prefill -------------------- .. autosummary:: :toctree: ../generated batch_prefill batch_prefill_with_paged_kv_cache .. autoclass:: BatchPrefillTSWrapper :members: .. automethod:: __init__ .. autoclass:: BatchPrefillPagedTSWrapper :members: .. automethod:: __init__ FMHA Decode ----------- .. autosummary:: :toctree: ../generated batch_decode_with_paged_kv_cache get_prims_ts_batch_decode_workspace_size prims_ts_batch_decode_with_kv_cache .. autoclass:: BatchDecodePagedTSWrapper :members: .. automethod:: __init__ MLA Decode ---------- .. autosummary:: :toctree: ../generated batch_decode_mla_with_paged_kv_cache get_prims_ts_batch_decode_mla_workspace_size prims_ts_batch_decode_with_kv_cache_mla .. autoclass:: BatchMLADecodePagedTSWrapper :members: .. automethod:: __init__ flashinfer.decode ================= .. currentmodule:: flashinfer.decode Single Request Decoding ----------------------- .. autosummary:: :toctree: ../generated single_decode_with_kv_cache single_decode_with_kv_cache_with_jit_module Batch Decoding -------------- .. autosummary:: :toctree: ../generated cudnn_batch_decode_with_kv_cache trtllm_batch_decode_with_kv_cache xqa_batch_decode_with_kv_cache .. autoclass:: BatchDecodeWithPagedKVCacheWrapper :members: :exclude-members: begin_forward, end_forward, forward, forward_return_lse .. automethod:: __init__ .. autoclass:: BatchDecodeMlaWithPagedKVCacheWrapper :members: :exclude-members: begin_forward, end_forward, forward, forward_return_lse .. automethod:: __init__ .. autoclass:: CUDAGraphBatchDecodeWithPagedKVCacheWrapper :members: .. automethod:: __init__ XQA --- .. currentmodule:: flashinfer.xqa .. autosummary:: :toctree: ../generated xqa xqa_mla flashinfer.prefill ================== Attention kernels for prefill & append attention in both single request and batch serving setting. .. currentmodule:: flashinfer.prefill Single Request Prefill/Append Attention --------------------------------------- .. autosummary:: :toctree: ../generated single_prefill_with_kv_cache single_prefill_with_kv_cache_return_lse single_prefill_with_kv_cache_with_jit_module Batch Prefill/Append Attention ------------------------------ .. autosummary:: :toctree: ../generated cudnn_batch_prefill_with_kv_cache trtllm_batch_context_with_kv_cache trtllm_ragged_attention_deepseek fmha_v2_prefill_deepseek trtllm_fmha_v2_prefill fmha_v2_prefill_sm120 .. autoclass:: BatchPrefillWithPagedKVCacheWrapper :members: :exclude-members: begin_forward, end_forward, forward, forward_return_lse .. automethod:: __init__ .. autoclass:: BatchPrefillWithRaggedKVCacheWrapper :members: :exclude-members: begin_forward, end_forward, forward, forward_return_lse .. automethod:: __init__ Unified BatchAttention ---------------------- .. currentmodule:: flashinfer.attention The ``BatchAttention`` class provides a holistic attention wrapper that automatically dispatches between paged-prefill and paged-decode based on per-request sequence lengths. It is the recommended entry point for serving stacks that batch mixed prefill/decode requests in a single kernel launch. .. autoclass:: BatchAttention :members: .. automethod:: __init__ .. autoclass:: BatchAttentionWithAttentionSinkWrapper :members: .. automethod:: __init__ SM120 NVFP4 Attention --------------------- .. currentmodule:: flashinfer.nvfp4_attention_sm120 .. autosummary:: :toctree: ../generated nvfp4_attention_sm120_quantize_qkv nvfp4_attention_sm120_fwd flashinfer.mla ============== MLA (Multi-head Latent Attention) is an attention mechanism proposed in DeepSeek series of models ( `DeepSeek-V2 `_, `DeepSeek-V3 `_, and `DeepSeek-R1 `_). .. currentmodule:: flashinfer.mla PageAttention for MLA --------------------- .. autosummary:: :toctree: ../generated trtllm_batch_decode_with_kv_cache_mla trtllm_batch_decode_sparse_mla_dsv4 convert_compressed_page_aligned_sparse_indices_to_hca_metadata DSV4HCAMetadata xqa_batch_decode_with_kv_cache_mla .. note:: With ``backend="cute-dsl"``, pass ``hca_swa_indices`` as absolute rows into the flattened SWA cache and ``hca_compressed_block_tables`` as physical compressed-cache page IDs. The SWA table has shape ``[B * Q, 128]`` and may express ring rotation or wraparound. Combined tables whose compressed segment is a canonical page expansion can opt into compatibility conversion with ``hca_sparse_indices_format="compressed-page-aligned"``. SWA entries remain arbitrary absolute rows. Precompute that conversion before a CUDA Graph or a latency-sensitive loop. .. autoclass:: BatchMLAPagedAttentionWrapper :members: .. automethod:: __init__ --- ## File: docs/api/cascade.rst .. _apicascade: flashinfer.cascade ================== .. currentmodule:: flashinfer.cascade .. _api-merge-states: Merge Attention States ---------------------- .. autosummary:: :toctree: ../generated merge_state merge_state_in_place merge_states .. _api-cascade-attention: Cascade Attention ----------------- Cascade Attention Wrapper Classes --------------------------------- .. autoclass:: MultiLevelCascadeAttentionWrapper :members: :exclude-members: begin_forward, end_forward, forward, forward_return_lse .. automethod:: __init__ .. autoclass:: BatchDecodeWithSharedPrefixPagedKVCacheWrapper :members: .. automethod:: __init__ .. autoclass:: BatchPrefillWithSharedPrefixPagedKVCacheWrapper :members: .. automethod:: __init__ --- ## File: docs/api/comm.rst .. _apicomm: flashinfer.comm =============== .. currentmodule:: flashinfer.comm This module provides communication primitives and utilities for distributed computing, including CUDA IPC, AllReduce operations, and memory management utilities. CUDA IPC Utilities ------------------ .. autosummary:: :toctree: ../generated CudaRTLibrary create_shared_buffer free_shared_buffer DLPack Utilities ---------------- .. autosummary:: :toctree: ../generated pack_strided_memory Mapping Utilities ----------------- .. autosummary:: :toctree: ../generated Mapping TensorRT-LLM AllReduce ---------------------- Types and Enums ~~~~~~~~~~~~~~~~ .. autosummary:: :toctree: ../generated AllReduceFusionOp AllReduceFusionPattern AllReduceStrategyConfig AllReduceStrategyType QuantizationSFLayout Core Operations ~~~~~~~~~~~~~~~ .. autosummary:: :toctree: ../generated trtllm_allreduce_fusion trtllm_custom_all_reduce trtllm_moe_allreduce_fusion trtllm_moe_finalize_allreduce_fusion Workspace Management ~~~~~~~~~~~~~~~~~~~~ .. autosummary:: :toctree: ../generated trtllm_create_ipc_workspace_for_all_reduce trtllm_create_ipc_workspace_for_all_reduce_fusion trtllm_destroy_ipc_workspace_for_all_reduce trtllm_destroy_ipc_workspace_for_all_reduce_fusion Initialization and Utilities ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ .. autosummary:: :toctree: ../generated trtllm_lamport_initialize trtllm_lamport_initialize_all compute_fp4_swizzled_layout_sf_size Unified AllReduce Fusion API ---------------------------- .. autosummary:: :toctree: ../generated allreduce_fusion create_allreduce_fusion_workspace AllReduceFusionWorkspace All-reduce workspaces backed by ``SymmDeviceMemory`` preserve their CUDA virtual addresses across process checkpoint/restore. After quiescing all work, release the physical handles and restore them with a fresh communication backend before replaying a captured CUDA graph: .. code-block:: python workspace.checkpoint_prepare() workspace.checkpoint_restore(comm_backend) Both methods are collective. Every rank must call them in the same order, and ``comm_backend`` must reproduce the original rank and world size. Repeated calls are no-ops after the workspace reaches the requested state. If an exception occurs after detach or reattach begins, do not retry or reuse the workspace; restart the affected rank. Workspaces backed by torch symmetric memory do not support this lifecycle. .. autoclass:: TRTLLMAllReduceFusionWorkspace :members: :show-inheritance: .. automethod:: __init__ .. autoclass:: MNNVLAllReduceFusionWorkspace :members: :show-inheritance: .. automethod:: __init__ FP8 Quantized AllReduce ----------------------- .. currentmodule:: flashinfer.comm .. autosummary:: :toctree: ../generated quantized_all_reduce vLLM AllReduce -------------- .. autosummary:: :toctree: ../generated vllm_all_reduce vllm_dispose vllm_init_custom_ar vllm_register_buffer vllm_register_graph_buffers vllm_get_graph_buffer_ipc_meta vllm_meta_size Ulysses Context-Parallel All-to-All ----------------------------------- .. currentmodule:: flashinfer.comm Communication for Ulysses context parallelism over the 4-D layout ``[B, S, H, D]``. Two layout transforms are provided; a typical attention layer makes four collective calls (q/k/v through ``scatter_heads``, the output through ``gather_heads``): - ``scatter_heads``: ``[B, S_local, H, D] -> [B, S_global, H_local, D]`` — each rank keeps head slice ``[rank * H_local, (rank+1) * H_local)`` of the *full* sequence; - ``gather_heads``: ``[B, S_global, H_local, D] -> [B, S_local, H, D]`` — the inverse, returning all heads of this rank's sequence shard, with ``H_local = H // world_size`` and ``S_global = S_local * world_size``. Both backends produce bit-identical results. **Backend policy.** :class:`UlyssesCommunicator` selects its backend in the constructor, strictly before any IPC allocation or JIT compilation: ============== ================================================================== ``backend=`` behavior ============== ================================================================== ``"auto"`` fused-transpose NVLink-P2P kernel when the group is a verified single-node all-pairs NVLink mesh with a supported world size (2/4/6/8); NCCL otherwise. The instance exposes ``.backend`` (effective), ``.fallback_reason`` and ``.decision`` / ``.topology_decision``. ``"nvlink"`` force the fused kernel; raises on every rank (before IPC/JIT for topology failures) when it cannot be used. ``"nccl"`` force ``dist.all_to_all_single`` + permute; skips the topology/NVML probe and all IPC/JIT (the constructor still resolves/guards the CUDA device and performs CUDA-backed metadata collectives); any world size. ============== ================================================================== Typical fallback reasons reported by ``.fallback_reason`` (all conservative — anything unknown or unverifiable selects NCCL): unsupported world size (only 2/4/6/8 have fused-kernel instantiations; ``world_size == 1`` is a no-copy passthrough), ranks spanning multiple hosts, missing pair-wise P2P or NVLink between any two concrete GPUs (verified per pair via NVML, not "some active link"), duplicate or unknown physical GPU identity, a topology probe error, inconsistent per-rank decisions, or a runtime NVLink initialization failure after a positive topology decision. **Constraints.** The constructor is always collective (all ranks together); :meth:`UlyssesCommunicator.close` is collective only when the NVLink backend was armed — for the pure NCCL backend, ``world_size == 1``, or an auto fallback whose NVLink cleanup already completed, ``close`` is local and idempotent. Rank-local failures inside the NVLink initialization or a collective ``close`` are exchanged as group outcomes so all ranks jointly clean up and raise (or fall back) instead of deadlocking, and a failed ``close`` may be retried. All ranks must request the same ``backend`` and agree on ``max_elems`` and ``dtype``; each rank may bind a different CUDA device (``device`` accepts ``torch.device``, ``str`` or an ``int`` ordinal, e.g. ``cuda:rank``). With ``world_size > 1`` the NCCL backend (forced or fallen back to) requires ``group`` to support CUDA all-to-all (an NCCL process group), checked at construction. Operands must be contiguous 4-D CUDA tensors of the construction ``dtype`` (float16 / bfloat16 / float32 only) on the construction device, every dim positive, at most ``max_elems`` (≤ 2^31 − 1) elements; ``scatter_heads`` requires ``H % world_size == 0`` and ``gather_heads`` requires ``S_global % world_size == 0``. Collectives run on the current CUDA stream; all ranks must issue the same call sequence with consistent shapes, one collective in flight per communicator at a time. **Known limitations.** - PyTorch builds without ``torch.cuda.get_device_properties(...).uuid`` cannot establish physical GPU identity: ``auto`` conservatively falls back to NCCL (the reason names the missing attribute). - When each process can only see its own GPU (e.g. one ``CUDA_VISIBLE_DEVICES`` entry per rank), peers are invisible to the P2P probe and ``auto`` falls back to NCCL. - Out-of-range CUDA ordinals passed as *strings or ints* are rejected at construction; a pre-built ``torch.device`` object wraps its index into a signed byte before FlashInfer can see it (``torch.device("cuda:256")`` is already ``cuda:0``), so only the surviving index can be range-checked. - Teardown metadata exchanges run bound to the communicator device; an extreme failure in the guard *restore* path after a completed collective can still desynchronize ranks (never observed in tests; tracked as a hardening note). **Example** (Wan2.1-style attention; see the `wan example `_ for the full integration):: with UlyssesCommunicator(group, max_elems=B * S_local * H * D, dtype=torch.bfloat16) as comm: q_ = comm.scatter_heads(q) # [B,S_local,H,D] -> [B,S_global,H_local,D] k_ = comm.scatter_heads(k) v_ = comm.scatter_heads(v) o_ = attention(q_, k_, v_) o = comm.gather_heads(o_) # [B,S_global,H_local,D] -> [B,S_local,H,D] .. autoclass:: UlyssesCommunicator :members: :show-inheritance: .. automethod:: __init__ Topology Probing and Backend Selection ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ .. autosummary:: :toctree: ../generated UlyssesBackendDecision UlyssesRankTopology UlyssesBackendError .. autofunction:: resolve_ulysses_backend .. autofunction:: decide_ulysses_backend .. autofunction:: probe_ulysses_rank_topology Raw Kernel Entry Points (advanced) ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ Prefer :class:`UlyssesCommunicator`; these assume the caller has already verified all-pairs NVLink P2P and owns the IPC workspace lifecycle. .. autosummary:: :toctree: ../generated init_ulysses_a2a dispose_ulysses_a2a ulysses_a2a MNNVL (Multi-Node NVLink) ------------------------- .. currentmodule:: flashinfer.comm.mnnvl Core Classes ~~~~~~~~~~~~ .. autosummary:: :toctree: ../generated MnnvlMemory McastGPUBuffer TensorRT-LLM MNNVL AllReduce ---------------------------- .. currentmodule:: flashinfer.comm.trtllm_mnnvl_ar .. autosummary:: :toctree: ../generated trtllm_mnnvl_all_reduce trtllm_mnnvl_allreduce trtllm_mnnvl_fused_allreduce_add_rmsnorm trtllm_mnnvl_fused_allreduce_add_rmsnorm_quant trtllm_mnnvl_fused_allreduce_rmsnorm mpi_barrier MNNVL A2A (Throughput Backend) ------------------------------- .. currentmodule:: flashinfer.comm .. autosummary:: :toctree: ../generated moe_a2a_initialize moe_a2a_dispatch moe_a2a_combine moe_a2a_sanitize_expert_ids moe_a2a_get_workspace_size_per_rank moe_a2a_wrap_payload_tensor_in_workspace .. autoclass:: MoeAlltoAll :members: :inherited-members: :show-inheritance: .. automethod:: __init__ ``MoeAlltoAll`` preserves its CUDA virtual addresses across process checkpoint/restore. After quiescing all work, call ``checkpoint_prepare`` to release the non-checkpointable physical MNNVL handles. Then call ``checkpoint_restore`` with a fresh communication backend before replaying a captured CUDA graph: .. code-block:: python moe_alltoall.checkpoint_prepare() moe_alltoall.checkpoint_restore(comm_backend) Both methods are collective. Every rank must call them in the same order, and ``comm_backend`` must reproduce the original rank and world size. Repeated calls are no-ops after the workspace reaches the requested state. If an exception occurs after physical handle unmapping or remapping begins, do not retry or reuse the workspace; restart the affected rank. .. autosummary:: :toctree: ../generated MoeAlltoAll.checkpoint_prepare MoeAlltoAll.checkpoint_restore DCP All-to-All (Context-Parallel Attention Reduction) ----------------------------------------------------- .. currentmodule:: flashinfer.comm .. autosummary:: :toctree: ../generated decode_cp_a2a_workspace_size decode_cp_a2a_allocate_mnnvl_workspace decode_cp_a2a_init_workspace decode_cp_a2a_alltoall Mixed Communication ------------------- .. currentmodule:: flashinfer.comm.mixed_comm .. autosummary:: :toctree: ../generated MixedCommOp MixedCommMode MixedCommHandler run_mixed_comm --- ## File: docs/api/cudnn.rst .. _apicudnn: flashinfer.cudnn ================ cuDNN-backed attention kernels. These wrappers call into NVIDIA's cuDNN runtime for batch prefill and batch decode, and are typically used as an alternative backend for ``BatchPrefillWithPagedKVCacheWrapper`` / ``BatchDecodeWithPagedKVCacheWrapper`` when cuDNN is available on the host GPU. .. currentmodule:: flashinfer.cudnn .. autosummary:: :toctree: ../generated cudnn_batch_decode_with_kv_cache cudnn_batch_prefill_with_kv_cache --- ## File: docs/api/cute_dsl.rst .. _apicute_dsl: flashinfer.cute_dsl =================== CuTe-DSL implementations of selected FlashInfer kernels. These symbols are available only when the ``nvidia-cutlass-dsl`` package is installed and the host has a supported NVIDIA GPU; the module guards its imports with ``is_cute_dsl_available()``. .. note:: A handful of GEMM symbols (``grouped_gemm_nt_masked``, ``Sm100BlockScaledPersistentDenseGemmKernel``, ``create_scale_factor_tensor``) used to live in ``flashinfer.cute_dsl`` and are still re-exported for backwards compatibility, but their canonical home is :doc:`gemm`. New code should import from ``flashinfer.gemm``. .. currentmodule:: flashinfer.cute_dsl Availability ------------ .. autosummary:: :toctree: ../generated is_cute_dsl_available RMSNorm + FP4 Quantization -------------------------- .. autosummary:: :toctree: ../generated rmsnorm_fp4quant add_rmsnorm_fp4quant .. autoclass:: RMSNormFP4QuantKernel :members: .. automethod:: __init__ .. autoclass:: AddRMSNormFP4QuantKernel :members: .. automethod:: __init__ Attention Wrappers ------------------ CuTe-DSL implementations of the batch attention wrappers. .. currentmodule:: flashinfer.cute_dsl.attention.wrappers.batch_mla .. autoclass:: BatchMLADecodeCuteDSLWrapper :members: .. automethod:: __init__ .. currentmodule:: flashinfer.cute_dsl.attention.wrappers.batch_prefill .. autoclass:: BatchPrefillCuteDSLWrapper :members: .. automethod:: __init__ .. currentmodule:: flashinfer.cute_dsl.attention.wrappers.batch_decode .. autoclass:: BatchDecodeCuteDSLWrapper :members: .. automethod:: __init__ .. autoclass:: BatchDecodePagedCuteDSLWrapper :members: .. automethod:: __init__ Block Sparse Attention ---------------------- CuTe-DSL block-sparse attention forward kernels. .. currentmodule:: flashinfer.cute_dsl.sparse.bsa_attn_sm100_blk128 .. autosummary:: :toctree: ../generated bsa_attn_sm100_blk128_fwd .. currentmodule:: flashinfer.cute_dsl.sparse.bsa_attn_sm100_blk64 .. autosummary:: :toctree: ../generated bsa_attn_sm100_blk64_fwd .. currentmodule:: flashinfer.cute_dsl.sparse.bsa_attn_sm120 .. autosummary:: :toctree: ../generated bsa_attn_sm120_blk64_fwd HCA Decode ---------- .. currentmodule:: flashinfer.cute_dsl.attention.wrappers.batch_hca .. autosummary:: :toctree: ../generated cute_dsl_hca_decode The recommended public entry point is ``flashinfer.mla.trtllm_batch_decode_sparse_mla_dsv4`` with ``backend="cute-dsl"``. ``cute_dsl_hca_decode`` is the lower-level wrapper for callers that already use the explicit HCA metadata ABI. The sliding-window cache is flattened into token rows and selected by an ``[B * Q, 128]`` INT32 ``window_indices`` tensor of absolute row indices; ring rotation and wraparound are supported. The compressed cache remains paged and uses an ``[B * Q, max_pages]`` INT32 block table. Masked window padding must still contain a legal row index because gather4 reads every coordinate before masking. ``hca_seq_lens`` describes the backing HCA footprint scheduled by TMA, not the effective top-k. Each compressed block-table row must therefore contain legal page IDs for that footprint rounded up to 128-slot tiles, including slots later masked by a shorter ``sparse_topk_lens``. Callers whose existing ``sparse_indices`` have a canonical compressed-page expansion may set ``hca_sparse_indices_format="compressed-page-aligned"`` to generate SWA gather indices, the compressed block table, and HCA lengths. SWA entries remain arbitrary absolute rows; only the compressed segment must be the canonical page expansion. This one-shot compatibility path validates values, allocates metadata, synchronizes the device, immediately launches the decode, and is not CUDA Graph capture safe. It is not a hot-loop path. Latency-sensitive callers must precompute with ``convert_compressed_page_aligned_sparse_indices_to_hca_metadata`` and reuse the returned metadata through the explicit HCA arguments. Arbitrary TRTLLM-GEN token-row selections in the compressed segment cannot be represented by an HCA page table without repacking the compressed KV pool. --- ## File: docs/api/fp4_quantization.rst .. _apifp4_quantization: flashinfer.fp4_quantization =========================== .. note:: Starting in FlashInfer 0.6.12, the canonical home for FP4 quantization APIs is :ref:`apiquantization`. ``flashinfer.fp4_quantization`` remains as a backwards-compatibility shim that re-exports the same symbols, so existing code such as ``from flashinfer.fp4_quantization import fp4_quantize`` keeps working. New code should import from ``flashinfer.quantization.fp4_quantization`` (or its canonical re-export at ``flashinfer.quantization``). This page intentionally does not re-document the FP4 symbols, because each symbol is the same Python object as the one rendered on :ref:`apiquantization` — duplicating the autosummary entries here would make Sphinx emit "duplicate object description" warnings under ``sphinx -W``. See Also -------- * :ref:`apiquantization` — canonical FP4 / FP8 / packbits API reference, including all of the following symbols that ``flashinfer.fp4_quantization`` used to host: - :func:`flashinfer.quantization.fp4_quantize` - :func:`flashinfer.quantization.nvfp4_quantize` - :func:`flashinfer.quantization.nvfp4_batched_quantize` - :func:`flashinfer.quantization.block_scale_interleave` (alias: ``nvfp4_block_scale_interleave``) - :func:`flashinfer.quantization.e2m1_and_ufp8sf_scale_to_float` - :func:`flashinfer.quantization.scaled_fp4_grouped_quantize` - :func:`flashinfer.quantization.silu_and_mul_nvfp4_quantize` - :func:`flashinfer.quantization.shuffle_matrix_a` - :func:`flashinfer.quantization.shuffle_matrix_sf_a` - :func:`flashinfer.quantization.nvfp4_kv_quantize` - :func:`flashinfer.quantization.nvfp4_kv_dequantize` - :func:`flashinfer.quantization.nvfp4_quantize_paged_kv_cache` - :class:`flashinfer.quantization.SfLayout` --- ## File: docs/api/fused_moe.rst .. _apifused_moe: flashinfer.fused_moe ==================== .. currentmodule:: flashinfer.fused_moe This module provides fused Mixture-of-Experts (MoE) operations optimized for different backends and data types. Types and Enums --------------- .. autosummary:: :toctree: ../generated RoutingMethodType WeightLayout Shared activation helpers live in :mod:`flashinfer.tllm_enums` and are used by both the TRT-LLM and CuteDSL MoE paths. .. currentmodule:: flashinfer.tllm_enums .. autosummary:: :toctree: ../generated is_gated_activation .. currentmodule:: flashinfer.fused_moe Utility Functions ----------------- .. autosummary:: :toctree: ../generated convert_to_block_layout reorder_rows_for_gated_act_gemm interleave_moe_weights_for_sm90_mixed_gemm interleave_moe_scales_for_sm90_mixed_gemm preprocess_moe_weights_for_sm90_mixed_gemm_humming fused_topk_deepseek hash_topk The E8M0 range-clamping, residual-scale factorization, and FP4 payload-rewrite scheme used by ``preprocess_moe_weights_for_sm90_mixed_gemm_humming`` is adapted from `Humming `_. Multi-LoRA MoE (BGMV) --------------------- Batched Gather-Matrix-Vector kernels for serving multiple LoRA adapters on top of a Mixture-of-Experts layer (shrink + expand). .. autosummary:: :toctree: ../generated bgmv_moe bgmv_moe_shrink bgmv_moe_expand bgmv_moe_gemm1_lora_delta bgmv_moe_gemm2_lora_delta CUTLASS Fused MoE ----------------- .. autosummary:: :toctree: ../generated cutlass_fused_moe TensorRT-LLM Fused MoE ---------------------- .. autosummary:: :toctree: ../generated trtllm_bf16_moe trtllm_bf16_routed_moe trtllm_fp4_block_scale_moe trtllm_fp4_block_scale_routed_moe trtllm_fp8_block_scale_moe trtllm_fp8_block_scale_routed_moe trtllm_fp8_per_tensor_scale_moe trtllm_fp8_per_tensor_scale_routed_moe trtllm_mxint4_block_scale_moe trtllm_mxint4_block_scale_routed_moe CuteDSL Fused MoE ----------------- The CuteDSL backends are conditionally available when the ``nvidia-cutlass-dsl`` package is installed. .. autosummary:: :toctree: ../generated cute_dsl_fused_moe_nvfp4 cute_dsl_fused_moe_mxfp8_mxfp4 b12x_fused_moe .. autoclass:: CuteDslMoEWrapper :members: :inherited-members: :show-inheritance: .. automethod:: __init__ .. autoclass:: CuteDslMxfp8Mxfp4MoEWrapper :members: :inherited-members: :show-inheritance: .. automethod:: __init__ .. autoclass:: B12xMoEWrapper :members: :inherited-members: :show-inheritance: .. automethod:: __init__ MonoMoE (Single-Kernel Block-FP8, SM90a) ----------------------------------------- Single-kernel top-K Mixture-of-Experts implementation specialized for the Qwen3.5-35B block-FP8 shape on Hopper (SM90a). The full pipeline — routing, up-projection, SiLU, down-projection and reduction — runs inside one kernel launch. Use :func:`has_monomoe` to check availability before calling. .. autosummary:: :toctree: ../generated has_monomoe get_scratchpad_size_bytes alloc_scratchpad interleave_for_tma_wgmma_up mono_moe ## 2. Official Technical Reference & Guides (flashinfer-ai/flashinfer-ai.github.io) # flashinfer.ai Source for the [FlashInfer project website](https://flashinfer.ai), served by GitHub Pages from the `main` branch. ## Local preview ```bash bundle install bundle exec jekyll serve # http://localhost:4000 ``` See `CLAUDE.md` for build constraints and how the site is put together. ## Content ### Blog posts Markdown files in `_posts/`, named `YYYY-MM-DD-slug.md`. They appear on the home page and in the RSS feed. ### Release highlights The [Releases page](https://flashinfer.ai/releases/) renders one entry per release from `_releases/`. After a release is tagged: ```bash ./scripts/import_release.py # writes _releases/.md ``` See `_releases/_README.md` for the entry format. What belongs in the highlights is an editorial question, decided before they are published on the GitHub release — not here, which is also why not every tag has an entry.