FlashInfer项目安装指南：从Python到C++的完整部署方案

项目概述

FlashInfer是一个高性能的深度学习推理加速库，基于PyTorch框架构建，同时提供C++原生接口。该项目针对现代GPU架构进行了深度优化，特别适合需要低延迟、高吞吐量的推理场景。

Python包安装

系统要求

在开始安装前，请确保您的环境满足以下基本要求：

• 操作系统：仅支持Linux系统
• Python版本：3.8至3.12
• PyTorch版本：2.2至2.5，需配合CUDA 11.8/12.1/12.4使用
• GPU架构支持：sm75(Turing)、sm80(Ampere A100)、sm86(Ampere RTX 30)、sm89(Ada Lovelace)、sm90(Hopper)

快速安装

FlashInfer提供了预编译的Python包，可以通过pip直接安装。注意包名为flashinfer-python而非flashinfer。

根据您的PyTorch和CUDA版本，选择对应的安装命令：

PyTorch 2.6版本

• CUDA 12.6: pip install flashinfer-python -i 指定源地址
• CUDA 12.4: pip install flashinfer-python -i 指定源地址

PyTorch 2.5版本

• CUDA 12.4/12.1/11.8均有对应版本

PyTorch 2.4版本

• 支持CUDA 12.4/12.1/11.8

PyTorch 2.3版本

• 支持CUDA 12.1/11.8

从源码安装

在某些情况下，您可能需要从源码安装FlashInfer，例如：

• 需要使用最新开发版功能
• 需要自定义库功能
• 目标环境与预编译版本不兼容

FlashInfer提供两种安装模式：

JIT模式(默认)

• 运行时通过PyTorch JIT编译CUDA内核
• 编译后的内核会缓存以供后续使用
• 安装快速，适合开发和测试环境

AOT模式

• 核心CUDA内核预编译并包含在库中
• 减少运行时编译开销
• 推荐用于生产环境

源码安装步骤

1. 获取源码仓库（使用git clone命令）

2. 确认已安装带CUDA支持的PyTorch：

   python -c "import torch; print(torch.__version__, torch.version.cuda)"
   

3. 安装Ninja构建系统：

   pip install ninja
   

4. 选择模式安装：

   • JIT模式：

     cd flashinfer
     pip install --no-build-isolation --verbose .
     

   • AOT模式：

     cd flashinfer
     export TORCH_CUDA_ARCH_LIST="7.5 8.0 8.9 9.0a 10.0a"
     python -m flashinfer.aot  # 在aot-ops/目录生成AOT内核
     python -m pip install --no-build-isolation --verbose .
     

5. (可选)创建分发包：

   • 创建源码包(sdist)
   • 创建JIT模式wheel包
   • 创建AOT模式wheel包

C++ API使用

FlashInfer作为头文件库设计，仅依赖CUDA和C++标准库，可直接集成到您的C++项目中而无需额外安装。

编译测试和基准程序

如需编译C++基准测试和单元测试，请按以下步骤操作：

1. 获取源码仓库（使用git clone命令）

2. 确保已安装conda（或直接安装cmake和ninja）

3. 安装构建工具：

   conda install cmake ninja
   

4. 创建构建目录并复制配置文件：

   mkdir -p build
   cp cmake/config.cmake build/
   

5. 编译项目：

   cd build
   cmake .. -G Ninja -DCMAKE_BUILD_TYPE=Release
   ninja
   

注意事项

1. 项目中的nvbench和googletest依赖仅用于编译测试和基准程序，库本身不需要这些依赖。

2. 在生产环境中推荐使用AOT模式以获得最佳性能。

3. 安装前请仔细检查PyTorch版本与CUDA版本的兼容性。

4. 对于特定GPU架构，可以调整TORCH_CUDA_ARCH_LIST环境变量以优化性能。