面向推理的高性能 GPU 内核

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

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

为什么选择 FlashInfer？

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

核心特性

注意力内核

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

GEMM 与线性运算

• BF16 GEMM：适用于 SM10.0+ GPU 的 BF16 矩阵乘法。
• FP8 GEMM：支持逐张量和分组缩放
• FP4 GEMM：适用于 Blackwell GPU 的 NVFP4 和 MXFP4 矩阵乘法
• 分组 GEMM：为 LoRA 和多专家路由提供高效的批处理矩阵运算

专家混合模型 (MoE)

• 融合 MoE 内核
• 多种路由方法：DeepSeek-V3、Llama-4 和标准 top-k 路由
• 量化 MoE：支持块级缩放的 FP8 和 FP4 专家权重

采样与解码

• 免排序采样：无需排序的高效 Top-K、Top-P 和 Min-P 采样
• 推测解码：支持链式推测采样

通信

• AllReduce：自定义实现
• 多节点 NVLink：支持多节点推理的 MNNVL
• NVSHMEM 集成：用于分布式内存操作

其他算子

• RoPE：LLaMA 风格的旋转位置编码（包括 LLaMA 3.1）
• 归一化：RMSNorm、LayerNorm、Gemma 风格的融合操作
• 激活函数：带融合门控的 SiLU、GELU

GPU 支持

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

最新动态

最新版本：

重要更新：
- [2025-10-08] 在 v0.4.0 版本中添加了对 Blackwell 的支持
- [2025-03-10] 博客文章 介绍了用于 LLM 采样的免排序 GPU 内核，解释了 FlashInfer 中采样内核的设计。

快速开始

安装

快速安装：

pip install flashinfer-python

软件包选项：

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

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

pip install flashinfer-python flashinfer-cubin
# JIT 缓存（将 cu129 替换为您的 CUDA 版本）
pip install flashinfer-jit-cache --index-url https://flashinfer.ai/whl/cu129

验证安装

flashinfer show-config

基本用法

import torch
import flashinfer

# 单次解码注意力
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：
bash 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（为目标 GPU 自定义）
export FLASHINFER_CUDA_ARCH_LIST="7.5 8.0 8.9 9.0a 10.0a 10.3a 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  # 从 PyPI 安装依赖项
pip install -U --pre flashinfer-cubin --index-url https://flashinfer.ai/whl/nightly/
# JIT 缓存（将 cu129 替换为您的 CUDA 版本）
pip install -U --pre flashinfer-jit-cache --index-url https://flashinfer.ai/whl/nightly/cu129

CLI 工具

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

# 验证安装并查看配置
flashinfer show-config

# 列出和检查模块
flashinfer list-modules
flashinfer module-status

# 管理构件和缓存
flashinfer download-cubin
flashinfer clear-cache

# 开发者：生成 compile_commands.json 用于 IDE 集成
flashinfer export-compile-commands [output_path]

完整文档请参阅 CLI 参考。

API 日志记录

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

# 启用日志记录（级别：0=关闭（默认），1=基本，3=详细，5=统计信息）
export FLASHINFER_LOGLEVEL=3

# 设置日志目标（stdout（默认），stderr，或文件路径）
export FLASHINFER_LOGDEST=stdout

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

自定义注意力变体

用户可以使用附加参数自定义自己的注意力变体。更多详情，请参考我们的 JIT 示例。

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 对您的项目或研究有帮助，请考虑引用我们的论文：

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