FlagTree Backend Specialization - flagos-ai/FlagTree GitHub Wiki

FlagTree 设计的后端统一特化，目的是整合后端接入范式，清晰化管理后端的特化实现，为后端维护特有变更、复用 FlagTree 公共能力、升级 Triton 版本、避免编辑模式出错等提供了工程基础。

具体实施方案是将各后端对 Triton 的特化，从以往的直接修改主干代码或者是将整套代码放在后端目录，标准化为在复用主干代码的基础上，在后端目录中给出差异化实现。

主干代码在原则上，既不允许直接出现某后端的特化实现，也不允许对后端做选择判断后特化实现，除非该主干代码文件原先已有后端选择实现。

Python Code Specialization

python/triton/ 目录下的主干代码，可通过调用接口，连接到 third_party/${backend_name}/spec/triton/ 目录下的特化实现。 注意各后端需要在 python/triton/spec/ 目录下建立到 third_party/${backend_name}/spec/ 的软链接，并提交到仓库。

cd python/triton/spec/
ln -s ../../../third_party/${backend_name}/spec/ ${backend_name}
git add ${backend_name}

接口定义在 python/triton/_flagtree_spec.py。

1. 特化整个 py 文件

spec_path 接口原理：在 __init__.py 劫持 __path__，优先在后端特化路径查找同名文件（模块），若不存在才使用主干代码

案例一：FlagTree/pull/862（iluvatar）

目标：iluvatar 后端需要在 python/triton/backends/ 另行实现 compile.py，目录下其他文件复用主干代码

特化的整代码文件：在 third_party/iluvatar/spec/triton/backends/ 添加与主干代码不同的 compile.py，不添加与主干代码完全相同或功能完全被主干代码包含的文件，也不必添加 __init__.py

主干目录代码：python/triton/backends/__init__.py 调用接口

from triton._flagtree_spec import spec_path

# flagtree backend path specialization
spec_path(__path__)

...

案例二：FlagTree/pull/687（mthreads）

目标：mthreads 后端需要在 python/triton/tools/ 另行实现 compile.py、link.py、ragged_tma.py，目录下其他文件复用主干代码

特化的整代码文件：在 third_party/mthreads/spec/triton/tools/ 添加与主干代码不同的 compile.py、link.py、ragged_tma.py，不添加与主干代码完全相同或功能完全被主干代码包含的文件，也不必添加 __init__.py

主干目录代码：python/triton/tools/__init__.py 调用接口

from triton._flagtree_spec import spec_path

# flagtree backend path specialization
spec_path(__path__)

...

适用范围：基本原则是不在父子目录同时劫持 __path__，因为这会导致子目录的主干代码被特化的父目录屏蔽。如果遇到这样的特化需求，方案一是父目录的改动直接合入主干代码（例如 python/triton/ 中的文件），方案二是仅对父目录劫持，特化代码需完整实现整套子目录（例如 python/triton/experimental/gluon/ 中的子目录）。

python/triton  # 不允许劫持路径
├── FLAGTREE_BACKEND      # 安装 FlagTree 时自动生成并记录环境变量 FLAGTREE_BACKEND
│                         # 运行时可通过 from triton._flagtree_backend import FLAGTREE_BACKEND 获取
├── _filecheck.py         # 如有特化需求使用 spec_call|spec_func
├── _flagtree_backend.py  # 用于 import FLAGTREE_BACKEND，不允许后端修改
├── _flagtree_spec.py     # 定义特化接口 spec_path|spec_call|spec_func，不允许后端修改
├── _internal_testing.py  # 如有特化需求直接在该主干文件中实现（符合 Triton 惯例）
├── _utils.py             # 如有特化需求使用 spec_call|spec_func
├── backends              # 允许 spec_path （在 __init__.py 调用，下同）
├── compiler              # 允许 spec_path
├── errors.py             # 如有特化需求使用 spec_call|spec_func
├── experimental  # 不允许劫持路径
│   ├── gluon       # 允许 spec_path
│   │   ├── amd       # 不允许劫持路径，如有需求应特化整个 amd 目录
│   │   ├── language  # 不允许劫持路径，如有需求应特化整个 language 目录
│   │   │   ├── amd          # 不允许劫持路径
│   │   │   │   └── xxx        # 不允许劫持路径
│   │   │   ├── extra        # 不允许劫持路径
│   │   │   └── nvidia       # 不允许劫持路径
│   │   │       ├── ampere     # 不允许劫持路径
│   │   │       ├── blackwell  # 不允许劫持路径
│   │   │       └── hopper     # 不允许劫持路径
│   │   └── nvidia    # 不允许劫持路径，如有需求应特化整个 nvidia 目录
│   └── tle         # 不允许特化，必须适配通用接口，或将变更需求合入通用接口
│       ├── language  # 不允许特化
│       │   ├── gpu     # 不允许特化，后端接口直接放在此处的子目录
│       │   │   └── {backend_name}
│       │   └── raw     # 不允许特化
│       └── raw       # 不允许特化
│           └── xxx     # 不允许特化
├── knobs.py          # 定义环境变量开关，如有特化需求直接在该主干文件中实现（符合 Triton 惯例）
├── language          # 允许 spec_path
│   └── extra           # 不允许劫持路径
├── runtime           # 允许 spec_path
├── spec              # 不允许特化，本身就是后端特化软链接集
│   └── {backend_name}  # 不允许特化，这是到 third_party/{backend_name}/spec 的软链接
├── testing.py        # 如有特化需求使用 spec_call|spec_func
└── tools             # 允许 spec_path

2. 特化一个函数调用

spec_call 接口原理：通过 third_parth/{backend_name}/spec/triton/__init__.py 调用后端定义的函数，该功能常用于在 __init__.py 修改模块的导出符号，也可用于其他 py 文件

案例一：FlagTree/pull/862（iluvatar）

目标：iluvatar 后端需要在 python/triton/_filecheck.py 修改 run_parser 函数中的部分实现

特化函数定义代码：新增文件 third_party/iluvatar/spec/triton/_filecheck.py 实现 filecheck_make_ir 函数，用于替换 run_parser 函数中的部分实现

def filecheck_make_ir(src, target, options, codegen_fns, module_map, context):
    return src.make_ir(target, options, codegen_fns, module_map, context)

特化顶层初始化代码：third_party/iluvatar/spec/triton/__init__.py 导入该函数

from ._filecheck import filecheck_make_ir

主干目录代码：python/triton/_filecheck.py 的 run_parser 函数中，通过接口调用 filecheck_make_ir 函数

from triton._flagtree_spec import spec_call

def run_parser(...):
    ...
    # module = backend.make_ir(src, options, codegen_fns, module_map, context)  # 原主干代码
    # flagtree backend call specialization
    module = spec_call("filecheck_make_ir", src, target, options, codegen_fns, module_map, context)
    if module is None:
        module = backend.make_ir(src, options, codegen_fns, module_map, context)
    ...

注意：iluvatar 后端同时需要在 python/triton/_filecheck.py 修改 stub_target，由于此前已有后端（mthreads）特化修改 stub_target，因此不用在主干代码重复添加下列代码

# flagtree backend call specialization
stub_target = spec_call("spec_get_stub_target")
if not stub_target:
    stub_target = GPUTarget("cuda", 100, 32)

案例二：FlagTree/pull/687（mthreads）

目标：mthreads 后端需要在 python/triton/language/__init__.py 添加导出 squeeze, unsqueeze 等符号

特化函数定义代码：新增文件 third_party/mthreads/spec/triton/language/__init__.py 实现 language_extend_globals 函数，注意其中的 import 必须使用绝对路径

def language_extend_globals(globals_dict):
    # NOTE: Must use absolute path import.
    from triton.language.standard import squeeze, unsqueeze
    globals_dict["squeeze"] = squeeze
    globals_dict["unsqueeze"] = unsqueeze

特化顶层初始化代码：third_party/mthreads/spec/triton/__init__.py 导入该函数

from .language import language_extend_globals

主干目录代码：python/triton/language/__init__.py 通过接口调用 language_extend_globals 函数，注意用于修改模块的导出符号时，spec 应放在所有 import 之后

from triton._flagtree_spec import spec_call

...  # import

# flagtree backend call specialization
spec_call("language_extend_globals", globals())

__all__ = [ ... ]

3. 新增或修改一个函数定义

spec_func 接口原理：通过 third_parth/{backend_name}/spec/triton/__init__.py 返回后端新增或修改的函数定义，可用于任意的 py 文件

案例一：FlagTree/pull/862（iluvatar）

目标：iluvatar 后端需要在 python/triton/testing.py 修改 get_max_tensorcore_tflops、nvsmi 的函数定义

注意：由于 python/triton/ 是所有目录的父目录，不允许使用 spec_path 劫持 __path__，因此无法整文件特化 _utils.py，转而使用 spec_func 添加函数定义

特化函数定义代码：新增文件 third_party/iluvatar/spec/triton/testing.py

def get_max_tensorcore_tflops(...):
    ...

def nvsmi(...):
    ...

特化顶层初始化代码：third_party/iluvatar/spec/triton/__init__.py 导入该函数

from .testing import get_max_tensorcore_tflops, nvsmi

主干目录代码：python/triton/testing.py 通过接口导入 get_max_tensorcore_tflops、nvsmi 的函数定义

from triton._flagtree_spec import spec_func

# flagtree backend func specialization
nvsmi = spec_func("nvsmi") or nvsmi
get_max_tensorcore_tflops = spec_func("get_max_tensorcore_tflops") or get_max_tensorcore_tflops

案例二：FlagTree/pull/687（mthreads）

目标：mthreads 后端需要在 python/triton/_utils.py 添加 apply_with_path 的函数定义

注意：由于 python/triton/ 是所有目录的父目录，不允许使用 spec_path 劫持 __path__，因此无法整文件特化 _utils.py，转而使用 spec_func 添加函数定义

特化函数定义代码：新增文件 third_party/mthreads/spec/triton/_utils.py

def apply_with_path(...):
    ...

特化顶层初始化代码：third_party/mthreads/spec/triton/__init__.py 导入该函数

from ._utils import apply_with_path

主干目录代码：python/triton/_utils.py 通过接口导入 apply_with_path 的函数定义

from triton._flagtree_spec import spec_func

# flagtree backend func specialization
apply_with_path = spec_func("apply_with_path")

C++ Code Specialization

1. 特化整个.td文件

flagtree_spec_td_set 接口原理：优先在后端特化路径查找同名的.td文件，若不存在才使用主干的.td文件

案例：FlagTree/pull/940（enflame）

目标：enflame 后端需要自定义 include/triton/Dialect/TritonGPU/Transforms/Passes.td

特化的整个.td文件：third_party/enflame/spec_cpp/include/triton/Dialect/TritonGPU/Transforms/Passes.td

主干目录脚本：include/triton/Dialect/TritonGPU/Transforms/CMakeLists.txt 改为调用接口

# set(LLVM_TARGET_DEFINITIONS Passes.td)  # 原主干实现
flagtree_spec_td_set(LLVM_TARGET_DEFINITIONS Passes.td)
...

2. 特化整个.h/.hpp头文件

原理：通过 CMakeLists.txt 调用 cmake/FlagTreeOptions.cmake 中的 flagtree_include_directories，保证 third_party/{backend_name}/spec_cpp/include/ 目录下的头文件优先被包含，早于 include/

案例：FlagTree/pull/940（enflame）

目标：enflame 后端需要自定义 include/triton/Dialect/TritonGPU/Transforms/Passes.h

特化的整个.h文件：third_party/enflame/spec_cpp/include/triton/Dialect/TritonGPU/Transforms/Passes.h

主干目录代码：无需任何修改

3. 特化整个.c/.cc/.cpp/.cxx文件

原理：通过 CMakeLists.txt 调用 cmake/FlagTreeBackendSpec.cmake 中的 flagtree_apply_backend_source_overrides，保证 third_party/{backend_name}/spec_cpp/lib/ 目录下的源码文件优先被使用

案例：FlagTree/pull/940（enflame）

目标：enflame 后端需要自定义 lib/Dialect/TritonGPU/IR/Ops.cpp

特化的整个.cpp文件：third_party/enflame/spec_cpp/lib/Dialect/TritonGPU/IR/Ops.cpp

主干目录代码：无需任何修改

构建日志：

  -- =====================================
  -- flagtree backend cmake specialization
  -- SPEC: lib/Conversion/TritonGPUToLLVM/FuncOpToLLVM.cpp -> TritonGPUToLLVM
  -- SPEC: lib/Dialect/TritonGPU/IR/Ops.cpp -> TritonGPUIR
  -- =====================================

4. 特化python/src/目录下的整个.h/.cc文件

原理：通过 CMakeLists.txt 调用 cmake/FlagTreeOptions.cmake 中的 flagtree_python_src_path_set，保证 third_party/{backend_name}/python/src/ 目录下的头文件优先被包含，源码文件优先被使用

案例：FlagTree/pull/843（ppu）

目标：ppu 后端需要自定义 python/src/gluon_ir.cc，目录下其他文件复用主干代码

特化的整个.cc文件：third_party/ppu/python/src/gluon_ir.cc

主干目录代码：无需任何修改