使用 C++ 封装算子获取极致吞吐#

针对延迟极其敏感或吞吐要求极高的生产场景，FlagGems 提供了 C++ 封装算子（C++ Wrapped Operators） 这一深度优化路径。 虽然 Triton 语言编写的 Kernel 具备非常优异的硬件计算效率，但 Triton 本质上是嵌入在 Python 环境中的 DSL：算子的外层包装逻辑、参数校验以及运行时分发均由 Python 解释器驱动。在密集小算子调用、短序列推理或极高 QPS 的服务场景中，Python 运行时的开销往往会成为不容忽视的瓶颈。

为此，FlagGems 构建了专属的 C++ 运行时层，将算子的元数据处理、参数特化逻辑与分发调度完全下沉至 C++ 原生实现，同时计算核心依然无缝复用高性能的 Triton Kernel。这既保留了 Triton 在算法层的高效表达与快速迭代优势，又彻底消除了 Python 层的解释与调度损耗，实现了更极致的端到端吞吐表现。

1. 架构#

FlagGems 中 C++ 封装的算子构建于 libtriton_jit 之上。 libtriton_jit 是一个多后端的 Triton JIT C++ 运行时，它在 C++ 中重新实现了 Triton JIT 运行时逻辑（参数特化、内核缓存和发射）， 而实际的 kernel 编译仍然委托给上游的 Triton 编译器完成。

在这一整套技术栈中：

  • Triton 内核（*.py）仍然是设备端计算的唯一源头；
  • libtriton_jit 负责 JIT 特化、内核缓存以及特定后端的 kernel 发射， 目前已支持 NVIDIA（CUDA）、摩尔线程（MUSA）、华为昇腾（NPU） 与天数智芯（IX） 四种后端；
  • FlagGems 的 C++ 封装算子（位于 cpp/lib/ 目录下，例如 rms_norm.cpp、mm.cpp） 以 C++ 实现张量元数据处理、形状/类型提升以及参数准备，最后通过 libtriton_jit::TritonJITFunction 调用 Triton 内核；
  • 在封装算子之上，FlagGems 还提供两个面向 Python 的扩展模块 （src/flag_gems/csrc/cstub.cpp 与 src/flag_gems/csrc/aten_patch.cpp） 以及一个可安装的 C++ 库目标（FlagGems::operators）， 它们共同把同一份 C++ 封装通过四种不同的调用方式暴露给上层用户 （详见§3 C++ 算子的四种调用方式）。

无论最终走的是哪一种调用方式，包装器（wrapper）本身的逻辑都在 C++ 中执行 （张量元数据处理、参数类型与特化判定、kernel 缓存查找与启动参数准备等）， 这部分 Python 开销被无条件消除，同时底层计算仍然复用相同的 Triton 内核。

而 PyTorch dispatcher 的开销是否也能被省掉，则取决于所选的调用路径：

  • 经过 dispatcher 的路径（torch.ops.flag_gems.*、以及 ATen 直接替换） 仍然要付出常规的 dispatcher 代价；不过由于 dispatcher 背后挂的是 C++ 实现而不是 Python 包装器，整体的 boxed-call 开销仍然比纯 Python 自定义算子小得多。
  • 绕过 dispatcher 的路径（c_operators pybind 模块、原生 C++ API） 则完全没有 dispatcher 开销；其中原生 C++ API 还会在调用路径上进一步消除 Python 解释器本身的参与。

各种方式的具体权衡参见§3 C++ 算子的四种调用方式。

2. 安装与启用#

想让 C++ wrapper 全量生效，需要同时满足以下两步：

  1. 编译安装时：打开 C++ 扩展并以 Release 构建

    推荐方式 — 通过 setup.sh 从源码构建：

    ENABLE_CPP=1 ./setup.sh nvidia-cuda128

    setup.sh 从 backends.yaml 读取后端配置，设置对应的 CMake 参数， 向 cpp/pyproject.toml 注入 vendor 标识，并从 cpp/ 子目录构建 C++ 扩展。

    备选方式 — 安装预构建 wheel（如果您的硬件平台支持）：

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

    该命令从 flagOS PyPI 仓库拉取预构建的 flag-gems-cpp-cuda wheel。 请将 cpp-cuda 替换为您的硬件对应的扩展： cpp-musa、cpp-npu、cpp-gcu 或 cpp-ix。

    手动构建 — 从 cpp/ 子目录：

    C++ 源码和 CMake 构建系统位于 cpp/ 子目录下。 构建时至少需要传入 -DFLAGGEMS_BUILD_C_EXTENSIONS=ON 与 -DCMAKE_BUILD_TYPE=Release（后者保证 FlagGems 与随同构建的 libtriton_jit 都按目标平台开启优化，否则 wrapper 会明显变慢）：

    tools/set_cpp_vendor.sh cuda
CMAKE_ARGS="-DFLAGGEMS_BUILD_C_EXTENSIONS=ON -DCMAKE_BUILD_TYPE=Release" \
uv pip install --no-build-isolation ./cpp

    C++ 扩展必须指定目标硬件 (-DFLAGGEMS_BACKEND=<...>)。 tools/set_cpp_vendor.sh 会将 vendor 标识写入 cpp/pyproject.toml， 确保构建出的 wheel 具有正确的包名。

    其他可选参数：

    • -DFLAGGEMS_BACKEND=<CUDA|IX|MUSA|NPU>：选择目标后端（默认 CUDA）；
    • -DFLAGGEMS_BUILD_POINTWISE_DYNAMIC_CPP=ON：编译 add/div/fill 这几个 pointwise dynamic 算子；
    • -DFLAGGEMS_BUILD_CTESTS=ON：编译 cpp/ctests/ 下的 GTest 用例 （验证 §3.4 原生 C++ API 的唯一手段）；
    • -DFLAGGEMS_USE_EXTERNAL_TRITON_JIT=ON -DTritonJIT_ROOT=<path>： 使用外部预装的 libtriton_jit。

    完整的各后端示例与 libtriton_jit 细节，参见 安装指南。

  2. 运行时：export USE_C_EXTENSION=1

    只装上 C++ 扩展还不够。src/flag_gems/config.py 把一部分上层行为挂在这个环境变量上——不设它，下面这些路径会悄悄走 Python 回退：

    路径 / 行为编译好就可用还需 USE_C_EXTENSION=1
    §3.1 torch.ops.flag_gems.*✅—
    §3.3 c_operators pybind✅—
    §3.2 ATen 直接替换（aten_patch）❌✅
    flag_gems.enable() 的 C++ 分支❌✅
    GemsRMSNorm 等 nn.Module 的 C++ 分支❌✅

    所以正常使用请：

    export USE_C_EXTENSION=1

  3. 快速验证

    下面这个片段可以一次性检查 Python 侧能看到的三条路径：

    import torch
import flag_gems
from flag_gems import c_operators, aten_patch
from flag_gems.config import has_c_extension, use_c_extension

assert has_c_extension, "C++ extension was not built"
assert use_c_extension, "please `export USE_C_EXTENSION=1`"
assert hasattr(torch.ops.flag_gems, "mm"), "§3.1 torch.ops.flag_gems.* not registered"
assert aten_patch.get_registered_ops(), "§3.2 no ATen op has been replaced"
_ = c_operators.mm                                                 # §3.3

    §3.4 原生 C++ API 无法从 Python 观察；要验证它，需要编译时加 -DFLAGGEMS_BUILD_CTESTS=ON，然后运行 ctest：

    BUILD_DIR=$(ls -d cpp/build/*/ | head -n 1)
ctest --test-dir "${BUILD_DIR}" --output-on-failure

    手动单独跑某个 test 二进制时（例如 "${BUILD_DIR}/ctests/test_triton_mm"），需要 export FLAGGEMS_SOURCE_DIR=$(pwd)/src/flag_gems， C++ 运行时才能找到 Triton 内核的 .py 源文件； 通过 ctest 跑会自动设置。

  4. 典型使用场景

    在前两步都完成的前提下，下面两种写法都会自动优先使用 C++ 封装算子， 无需修改调用点：

    • 补丁模式 flag_gems.enable()：monkey-patch torch.* 接口， 当 use_c_extension == True 时 patch 后的函数走 torch.ops.flag_gems.*（§3.1），否则回退到 Python 实现。
    • 用 FlagGems 的 nn.Module 类搭模型：例如 flag_gems.modules.GemsRMSNorm， 内部已经内置「C++ 可用就走 torch.ops.flag_gems.*，否则走 Python」 的分支判断。参考 gems_rms_forward 的实现。

3. C++ 算子的四种调用方式#

C++ 扩展构建完成之后，同一份底层 C++ 封装算子实际上可以通过四种 不同的方式被调用。这几种方式面向不同的使用场景，并且分别有着不同的 dispatcher 开销。

3.1 通过 torch.ops.flag_gems.*（自定义算子名字空间）#

所有 C++ 封装算子都会在 src/flag_gems/csrc/cstub.cpp 中通过 TORCH_LIBRARY(flag_gems, m) 注册为 PyTorch 的自定义算子（custom op）， 归入 flag_gems 名字空间。您可以在 Python 中显式地调用它们， 从而绕过所有 patch 逻辑与 Python 侧的回退路径：

output = torch.ops.flag_gems.fused_add_rms_norm(...)
out    = torch.ops.flag_gems.mm(a, b)

3.2 通过 ATen 直接替换（在 dispatcher 层面透明接管 torch.*）#

对于其中的一部分算子，FlagGems 还会在 src/flag_gems/csrc/aten_patch.cpp 中使用 TORCH_LIBRARY_IMPL(aten, <dispatch_key>, m) 把 C++ 实现 直接注册到 aten:: 名字空间。具体的 dispatch key 按后端选择：

  • 对于 NVIDIA CUDA 与天数智芯（IX）使用 CUDA；
  • 对于华为昇腾（NPU）与摩尔线程（MUSA）使用 PrivateUse1。

由于注册是直接发生在 PyTorch 的 dispatcher 中的， 在受支持的设备上调用标准的 PyTorch API（例如 torch.nonzero(x)、 x.copy_(y)）时，会透明地派发到 FlagGems 的 C++ 实现， 完全不需要在 Python 层做 monkey patch。 这是对一个已有模型做加速时门槛最低的一条路径。

由于 TORCH_LIBRARY_IMPL 在模块导入时就会执行， 这种方式所替换的算子集合在构建期就已经确定， 目前还不支持在运行时对单个算子作启用/禁用控制。

3.3 通过 c_operators pybind 模块（直连、不经 dispatcher）#

同样的一组 C++ 封装算子在 src/flag_gems/csrc/cstub.cpp 中还会通过 PYBIND11_MODULE(c_operators, …) 导出为一个 Python 扩展模块：

from flag_gems import c_operators

out = c_operators.mm(a, b)
c_operators.fused_add_rms_norm(input, residual, weight, eps)

这一路径完全绕开了 PyTorch 的 dispatcher，因此是从 Python 调用 FlagGems C++ 算子开销最低的方式。它最适合用在延迟非常敏感的 microbenchmark，或者内循环中 dispatcher 开销已经能被测量出来的场景。

3.4 通过原生 C++ API（flag_gems:: 函数与 GTest）#

每一个封装算子本身都是一个普通的 C++ 函数，位于 flag_gems:: 名字空间， 在 include/flag_gems/operators.h 中声明，并随安装出的 CMake 目标 FlagGems::operators 一起发布。下游的 C++ 代码可以直接链接该目标并调用：

#include "flag_gems/operators.h"

at::Tensor c = flag_gems::mm_tensor(a, b);
at::Tensor y = flag_gems::rms_norm(x, weight, eps);

仓库中 cpp/ctests/ 下的 GTest 用例（例如 cpp/ctests/test_triton_mm.cpp）正是以这种方式调用 FlagGems 的算子的。 当您希望把 FlagGems 嵌入到一个非 Python 的 C++ 应用中，或者需要写 C++ 单元测试时，这就是合适的路径。

小结#

调用方式入口经 dispatcher
torch.ops.flag_gems.*TORCH_LIBRARY(flag_gems, …)是
ATen 直接替换TORCH_LIBRARY_IMPL(aten, …)是
flag_gems.c_operators pybindPYBIND11_MODULE(c_operators, …)否
原生 C++ APIoperators.h 中的 flag_gems::*否

参考：目前支持的 C++ 封装的算子#

以下算子目前在 FlagGems 中提供了 C++ 封装实现。

  • add（pointwise dynamic C++）
  • div（pointwise dynamic C++）
  • fill（pointwise dynamic C++）
  • addmm
  • mm
  • bmm
  • cat
  • contiguous
  • copy
  • embedding
  • exponential_
  • zeros
  • argmax
  • max
  • sum
  • softmax
  • sort
  • topk
  • nonzero
  • rms_norm
  • fused_add_rms_norm
  • rotary_embedding
  • flash_attn_varlen_func
  • reshape_and_cache_flash
  • rwkv_mm_sparsity
  • rwkv_ka_fusion

标注为 pointwise dynamic C++ 的算子仅在启用 CMake 选项 -DFLAGGEMS_BUILD_POINTWISE_DYNAMIC_CPP=ON 时才会被编译； 详情参阅安装指南。

作为持续性能优化工作的一部分，我们一直在努力扩大这一列表。

Calendar October 28, 2026
Backward 性能预调优 后端深度优化 Forward