推理高性能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缓存：动态批处理服务的高效内存管理
• 解码、预填充与追加：所有注意力阶段的优化内核
• 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支持

架构	计算能力	示例GPU
Turing (图灵)	SM 7.5	T4, RTX 20系列
Ampere (安培)	SM 8.0, 8.6	A100, A10, RTX 30系列
Ada Lovelace (洛夫莱斯)	SM 8.9	L4, L40, RTX 40系列
Hopper (霍普)	SM 9.0	H100, H200
Blackwell (布莱克威尔)	SM 10.0, 10.3	B200, B300
Blackwell (布莱克威尔)	SM 12.0, 12.1	RTX 50系列, DGX Spark, Jetson Thor

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

最新动态

最新版本:

重要更新：

• [2025-10-08] 在v0.4.0版本中新增对Blackwell的支持
• [2025-03-10] 博客文章 介绍LLM采样的无排序GPU内核设计

快速开始

安装

快速开始：

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：

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

命令行工具

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

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

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

# 管理二进制文件和缓存
flashinfer download-cubin
flashinfer clear-cache

# 开发者工具：生成IDE集成所需的compile_commands.json
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（文本生成推理）
• 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}
}

FlashInfer 中文快速上手指南

---

环境准备

系统要求

• GPU架构：支持 Turing(SM75) 及更新架构（如 RTX 30/40 系列、A100/H100/B200 等）
• CUDA 版本：12.6、12.8、13.0、13.1（需与 PyTorch 版本匹配）
• 操作系统：Linux（推荐 Ubuntu 20.04+）

前置依赖

# 安装基础依赖
sudo apt-get update && sudo apt-get install -y git python3-pip
# 升级 pip 和 setuptools（源码安装时必需）
python -m pip install --upgrade pip setuptools
# 安装 PyTorch（需与 CUDA 版本匹配）
pip install torch --index-url https://download.pytorch.org/whl/cu129

---

安装步骤

推荐安装（含预编译内核）

# 安装核心包 + 预编译内核（推荐国内用户使用清华源加速）
pip install flashinfer-python flashinfer-cubin -i https://pypi.tuna.tsinghua.edu.cn/simple

# 安装 JIT 缓存（替换 cu129 为实际 CUDA 版本）
pip install flashinfer-jit-cache --index-url https://flashinfer.ai/whl/cu129

从源码安装（需科学上网）

# 克隆仓库
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

验证安装

flashinfer show-config
# 输出应包含 GPU 架构信息和 CUDA 版本

---

基本使用

单个解码注意力示例

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)
print(output.shape)  # 输出: torch.Size([32, 128])

---

提示：

1. 国内用户安装时建议在 pip 命令后添加 -i https://pypi.tuna.tsinghua.edu.cn/simple 使用清华源加速
2. CUDA 版本需与 PyTorch 安装版本严格匹配
3. 更多高级用法请参考官方文档