安装 FlagGems#

1. 环境准备#

• 你必须确保为自己的硬件正确地安装了内核态的驱动程序和用户空间的 SDK 或工具链， 并且均已配置正确，工作正常。无论是 NVIDIA 平台还是其他 AI 加速器硬件， 这一点都适用。

• 如果你想尝试将 FlagGems 与 vLLM 集成， 则需要安装 vLLM，或者厂商定制版本 （如果有的话）。

你不需要手动安装 Python、PyTorch 或 Triton。 setup.sh 脚本会根据你选择的后端自动完成所有这些安装。

2. 从 PyPI 安装#

你可以使用自己常用的 Python 包管理器（例如 pip）从 PyPI 安装 FlagGems：

pip install flag_gems

提示

上述命令仅安装 FlagGems 中纯 Python 实现的算子。

如需使用 C++ 封装的算子（可减少性能关键路径上的 dispatch 开销）， 可以通过 extras 安装预构建的原生扩展 wheel：

pip install "flag-gems[cpp-cuda]"

该命令会从 flagOS PyPI 仓库拉取匹配的 flag-gems-cpp-cuda 包。 请将 cpp-cuda 替换为你所使用的硬件对应的扩展： cpp-musa、cpp-npu、cpp-gcu 或 cpp-ix。

如果你的平台上没有可用的预构建 wheel，请参阅 从源码构建安装。

2.1. 使用 flaggems-setup 安装后端依赖#

pip install flag_gems 只安装纯 Python 实现的算子库，不会为你的加速器 安装 PyTorch 或厂商专属的运行时依赖包。随 FlagGems 一同安装的 flaggems-setup 命令行脚本会从 flagOS PyPI 仓库为你选定的后端安装 PyTorch 及厂商依赖包， 从而补全运行环境。

如果你是从 PyPI 安装 FlagGems（而非用 setup.sh 从源码安装，后者已经替你 完成了这些步骤），推荐在安装完成后执行这一步。

列出可用的后端：

flaggems-setup --list

为你的后端安装依赖（后端名称与 setup.sh 接受的完全一致）：

# NVIDIA CUDA 12.8
flaggems-setup nvidia-cuda128

# Huawei Ascend CANN 9.0.0
flaggems-setup ascend-cann900

# MetaX MACA 3.8.1
flaggems-setup metax-maca3810

在不实际执行的情况下预览将要运行的命令：

flaggems-setup nvidia-cuda128 --dry-run

默认情况下，当 uv 位于 PATH 中时脚本会使用 uv pip，否则回退到 pip。 你也可以用 --pip 显式指定安装器：

flaggems-setup nvidia-cuda128 --pip "pip"

flaggems-setup 还会为你安装 Triton 系编译器。默认情况下，当后端提供 FlagTree 时优先选择 FlagTree，否则回退到 Triton——策略与 setup.sh 一致。你可以用 --compiler 标志或 COMPILER 环境变量覆盖该选择：

# 强制使用原生 Triton，而非自动选择的 FlagTree
flaggems-setup nvidia-cuda128 --compiler triton
COMPILER=triton flaggems-setup nvidia-cuda128

COMPILER 的完整行为参见环境变量。

3. 从源码构建、安装#

3.1 克隆源代码#

git clone https://github.com/flagos-ai/FlagGems
cd FlagGems/

3.2 运行 setup.sh#

setup.sh 脚本是从源码安装 FlagGems 的推荐方式。 它从 src/flag_gems/backends.yaml 读取所有配置信息，并自动完成以下操作：

• 安装 uv（如果尚未安装）
• 安装你的后端所需的正确 Python 版本
• 创建虚拟环境（.venv/）
• 安装构建工具、PyTorch 以及厂商特定的依赖包
• 安装 FlagGems 及其对应的 extras
• 安装编译器（FlagTree 或 Triton）
• 安装测试依赖
• 将后端环境变量写入 .venv/bin/activate

./setup.sh <backend>

例如：

# NVIDIA CUDA 12.8
./setup.sh nvidia-cuda128

# Huawei Ascend CANN 9.0.0
./setup.sh ascend-cann900

# MetaX MACA 3.8.1
./setup.sh metax-maca3810

查看可用的后端列表：

./setup.sh invalid  # 打印可用后端列表

安装完成后，激活环境即可开始使用：

source .venv/bin/activate
pytest tests/test_abs.py -vs

提示

• 后端所需的环境变量已自动包含在 .venv/bin/activate 中， 无需额外的环境配置步骤。
• 默认情况下，如果 FlagTree 可用则安装 FlagTree 作为编译器。 如需使用原生 Triton，可在运行 setup.sh 前设置 COMPILER=triton：

  COMPILER=triton ./setup.sh nvidia-cuda128

3.3 可编辑安装（用于开发）#

如果你在参与 FlagGems 的开发（例如开发新的算子）， 可以执行可编辑模式的安装，使得对 Python 源码的修改无需重新安装即可生效：

source .venv/bin/activate
uv pip install --no-build-isolation -e .

setup.sh 默认以非可编辑模式安装 FlagGems。如需切换为可编辑模式， 请在 setup.sh 完成之后执行上述命令。 --no-build-isolation 参数会复用 venv 中已安装的构建工具。

3.4 C++ 扩展（可选）#

FlagGems 通过 C++ 封装算子在性能关键路径上减少 dispatch 开销。 C++ 扩展是一个独立的、按硬件供应商区分的包（如 flag-gems-cpp-cuda）， 它将编译好的 .so 文件安装到 flag_gems/ 命名空间下，与纯 Python 算子并列。

C++ 扩展的获取方式有两种：

方式 A：通过 setup.sh 从源码构建#

ENABLE_CPP=1 ./setup.sh nvidia-cuda128

setup.sh 会自动完成以下操作：

• 通过 tools/set_cpp_vendor.sh 向 cpp/pyproject.toml 注入正确的 vendor 标识
• 设置对应的 CMAKE_ARGS（-DFLAGGEMS_BACKEND=...）
• 从 cpp/ 子目录构建并安装 C++ 扩展

方式 B：在 cpp/ 子目录中手动构建#

C++ 扩展使用 scikit-build-core 作为构建后端，需要 CMake、C++ 工具链以及 对应硬件厂商的 SDK。从 cpp/ 子目录进行构建：

# 设置 vendor 标识（cuda、musa、npu、gcu 或 ix）
tools/set_cpp_vendor.sh cuda

# 构建并安装
CMAKE_ARGS="-DFLAGGEMS_BUILD_C_EXTENSIONS=ON -DFLAGGEMS_BACKEND=CUDA" \
  uv pip install --no-build-isolation ./cpp

如需手动控制 CMake 选项，请参阅 CMake 选项参考。

运行时：通过 USE_C_EXTENSION 启用#

安装完成后，设置以下环境变量来激活 C++ 路径：

export USE_C_EXTENSION=1

如果不设置该变量，只有 torch.ops.flag_gems.* 及 c_operators pybind 模块生效；ATen 替换和 flag_gems.enable() 的 C++ 分支仍需此变量。 详见 C++ 使用指南。

4. 参考资料#

4.1 可用后端#

所有支持的后端定义在 src/flag_gems/backends.yaml 中。 每个后端指定了：

• Python 版本
• PyTorch 及厂商特定的依赖包
• Triton / FlagTree 编译器包
• 运行时环境变量

4.2 环境变量#

环境变量 COMPILER 控制使用哪个编译器：

环境变量 ENABLE_CPP 用来启用 C++ 扩展：

4.3 CMake 选项#

使用 C++ 扩展构建（ENABLE_CPP=1）时，以下 CMake 选项由 setup.sh 自动设置。 如需手动构建，可通过 CMAKE_ARGS 环境变量传递。

4.4 scikit-build-core 选项#

主包 flag-gems 使用 setuptools 作为构建后端。 scikit-build-core 仅用于从 cpp/ 子目录构建的 C++ 扩展。

scikit-build-core 是一个构建后端，用来桥接 CMake 和 Python 构建系统， 简化使用 CMake 构建 Python 模块的过程。 常用的配置环境变量包括：

1. SKBUILD_CMAKE_BUILD_TYPE：配置项目的构建类型。 合法取值包括 Release、Debug、RelWithDebInfo 和 MinSizeRel。

2. SKBUILD_BUILD_DIR：配置项目的构建目录。 默认值为 pyproject.toml 中定义的 build/{cache_tag}。

需要注意的是，环境变量 SKBUILD_CMAKE_ARGS 中多个选项使用分号（;）分隔， 而 CMAKE_ARGS 中使用空格分隔。

4.5 关于 libtriton_jit 库#

FlagGems 的 C++ 扩展依赖于 TritonJIT， 一个用 C++ 实现 Triton JIT 运行时的库，可以在 C++ 代码中调用 Triton JIT 函数。

如需使用外部的 TritonJIT 库，请先单独构建安装它， 然后通过 -DTritonJIT_ROOT=<安装路径> 选项告知 CMake：

CMAKE_ARGS="-DFLAGGEMS_BUILD_C_EXTENSIONS=ON -DFLAGGEMS_USE_EXTERNAL_TRITON_JIT=ON -DTritonJIT_ROOT=/usr/local/lib/libtriton_jit" \
ENABLE_CPP=1 ./setup.sh nvidia-cuda128