高性能 GPU 推理内核

| 文档 | 最新版本 | 博客 | Slack | 讨论区 |

FlashInfer 是一个用于推理的库和内核生成器，在多种 GPU 架构上提供最先进的性能。它为注意力、GEMM 和 MoE 操作提供统一的 API，并包含多种后端实现，包括 FlashAttention-2/3、cuDNN、CUTLASS 和 TensorRT-LLM。

为什么选择 FlashInfer？

• 最先进的性能：针对预填充、解码和混合批处理场景优化的内核
• 多种后端：自动为您的硬件和工作负载选择最佳后端
• 现代架构支持：支持 SM75 (Turing) 及更高版本（直至 Blackwell）
• 低精度计算：针对注意力、GEMM 和 MoE 操作的 FP8 和 FP4 量化
• 生产就绪：兼容 CUDAGraph 和 torch.compile，适用于低延迟服务

核心功能

注意力内核

• 分页和非规则 KV-Cache：用于动态批处理服务的高效内存管理
• 解码、预填充和追加：针对所有注意力阶段优化的内核
• MLA 注意力：原生支持 DeepSeek 的多潜变量注意力
• 级联注意力：用于共享前缀的内存高效分层 KV-Cache
• 稀疏注意力：块稀疏和可变块稀疏模式
• POD-Attention：用于混合批处理的融合预填充+解码

GEMM & Linear Operations

• BF16 GEMM: 面向 SM10.0+ GPU 的 BF16 矩阵乘法。
• FP8 GEMM: 张量级和分组缩放
• FP4 GEMM: 面向 Blackwell GPU 的 NVFP4 和 MXFP4 矩阵乘法
• Grouped GEMM: 用于 LoRA 和多专家路由的高效批量矩阵运算

Mixture of Experts (MoE)

• Fused MoE Kernels
• Multiple Routing Methods: DeepSeek-V3、Llama-4 和标准 top-k 路由
• Quantized MoE: 采用块级缩放的 FP8 和 FP4 专家权重

Sampling & Decoding

• Sorting-Free Sampling: 无需排序的高效 Top-K、Top-P 和 Min-P
• Speculative Decoding: 支持链式推测采样

Communication

• AllReduce: 自定义实现
• Multi-Node NVLink: 支持多节点推理的 MNNVL
• NVSHMEM Integration: 用于分布式内存操作

Other Operators

• RoPE: LLaMA 风格的旋转位置嵌入（包括 LLaMA 3.1）
• Normalization: RMSNorm、LayerNorm、Gemma 风格的融合操作
• Activations: 带融合门控的 SiLU、GELU

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

注意： 并非所有功能都支持所有计算能力。

新闻

最新：

重要更新：

• [2025-10-08] 在 v0.4.0 中添加了 Blackwell 支持
• [2025-03-10] 博客文章 Sorting-Free GPU Kernels for LLM Sampling，解释了 FlashInfer 中采样内核的设计。

入门

安装

快速开始：

pip install flashinfer-python

包选项：

• flashinfer-python：核心包，在首次使用时编译/下载内核
• flashinfer-cubin：针对所有受支持 GPU 架构的预编译内核二进制文件
• flashinfer-jit-cache：针对特定 CUDA 版本的预构建内核缓存

为了更快的初始化和离线使用，请安装可选包以预编译大多数内核：

pip install flashinfer-python
flashinfer install-cubin-wheel
flashinfer install-jit-cache-wheel

对于 Blackwell (SM100+) CuTe DSL 内核，请安装 CUDA 13 extra 以启用针对 Blackwell 优化的内核：

pip install flashinfer-python[cu13]

验证安装

flashinfer show-config

基本用法

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)

请参阅文档以获取全面的 API 参考和教程。

从源码安装

git clone https://github.com/flashinfer-ai/flashinfer.git --recursive
cd flashinfer
python -m pip install -v .

用于开发，请以可编辑模式安装：

python -m pip install --no-build-isolation -e . -v

注意： 使用 --no-build-isolation 时，pip 不会自动安装构建依赖项。FlashInfer 需要 setuptools>=77。如果遇到类似 AttributeError: module 'setuptools.build_meta' has no attribute 'prepare_metadata_for_build_editable' 的错误，请先升级 pip 和 setuptools：

python -m pip install --upgrade pip setuptools

构建可选包：

# flashinfer-cubin
cd flashinfer-cubin
python -m build --no-isolation --wheel
python -m pip install dist/*.whl

# 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

有关更多详细信息，请参阅 从源代码安装文档。

每日构建

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 工具

FlashInfer 提供了若干 CLI 命令，用于配置、模块管理和开发：

# 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]

有关完整文档，请参阅 CLI 参考。

API 日志记录

FlashInfer 提供全面的 API 日志记录功能，用于调试。使用环境变量启用它：

# 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

有关日志级别、配置和高级功能的详细信息，请参阅我们文档中的 Logging。

自定义注意力变体

用户可以使用额外参数自定义自己的注意力变体。更多详情，请参阅我们的 JIT examples。

CUDA 支持

支持的 CUDA 版本： 12.6, 12.8, 13.0, 13.1

注意： FlashInfer 致力于遵循 PyTorch 支持的 CUDA 版本以及最新的 CUDA 发布版本。

采用情况

FlashInfer 为以下系统的推理提供支持：

• SGLang
• vLLM
• TensorRT-LLM
• TGI (Text Generation Inference)
• MLC-LLM
• LightLLM
• lorax
• ScaleLLM

致谢

FlashInfer 受到 FlashAttention、vLLM、stream-K、CUTLASS 和 AITemplate 的启发。

引用

如果您发现 FlashInfer 对您的项目或研究有帮助，请考虑引用我们的 paper：

@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}
}