Skip to content

[WIP]feat: add MACA backend support - #1

Open
chen2021673 wants to merge 3 commits into
masterfrom
feat/maca-backend
Open

[WIP]feat: add MACA backend support#1
chen2021673 wants to merge 3 commits into
masterfrom
feat/maca-backend

Conversation

@chen2021673

@chen2021673 chen2021673 commented Aug 20, 2026

Copy link
Copy Markdown

背景

InfiniTrain 已提供基于 PrivateUse1 的后端扩展接口,但厂商相关的编译器、SDK、运行时、算子和通信实现不适合继续放在框架核心仓库中。

本 PR 初始化 InfiniTrain-Backends,建立独立于 InfiniTrain 核心仓库的 accelerator backend 扩展机制,并以 MACA 作为首个 provider 完成验证。

整体目标是形成如下扩展关系:

InfiniTrain Core
        ↓
PrivateUse1 Extension Interface
        ↓
InfiniTrain-Backends
        ↓
Vendor Backend

作为后续国产加速器接入的统一组织方式。

整体架构

仓库按照 provider 进行隔离:

InfiniTrain-Backends
├── backends/
│   └── maca/
├── examples/
├── tests/
└── third_party/
    └── InfiniTrain

每个厂商后端独立维护在:

backends/<provider>

其中包含该后端自身的构建配置、Backend 注册以及具体实现。

InfiniTrain 通过 submodule 固定一个明确版本,Backend 仓库依赖 InfiniTrain 提供的 PrivateUse1 接口,而不是反向修改框架核心代码。

这样可以保证:

  • InfiniTrain 核心保持与具体硬件厂商无关;
  • 不同厂商 Backend 相互隔离;
  • Backend 可以独立开发、测试和演进;
  • InfiniTrain 与 Backend 之间通过统一接口协作。

Backend 接入逻辑

构建时通过:

INFINITRAIN_BACKEND=<provider>

选择当前使用的 Backend。

一个 build tree 只对应一个 provider。

Backend 的接入主要分为两个阶段。

构建阶段

部分厂商 Backend 不只是链接额外的 runtime library,而是需要使用自己的编译器工具链。

CMake 的 compiler 必须在第一次 project() 之前确定,因此顶层工程会先加载:

backends/<provider>/cmake/pre_project.cmake

完成 provider 相关的编译环境配置,再进入 InfiniTrain 的正常构建流程。

因此 Backend 不只是 InfiniTrain 构建完成后附加的一组库,而是整个构建过程的一部分。

运行时注册阶段

每个 provider 对外提供统一的 Backend 注册入口。

MACA 当前通过:

infini_train::maca::RegisterBackend();

将自身能力注册到 InfiniTrain 的 PrivateUse1 Backend。

注册完成后,上层模型和训练逻辑仍然使用 InfiniTrain 原有抽象:

Model / Autograd / Distributed
              ↓
         Dispatcher
              ↓
        PrivateUse1
              ↓
           MACA

上层代码只需要感知 device = maca,而不需要直接依赖 MACA SDK 或 MACA 特有接口。

MACA Backend

本 PR 将 MACA 支持迁移到新的 Backend 架构下,并适配 InfiniTrain 当前的 PrivateUse1 接口。

MACA Backend 负责承接厂商相关能力,并通过统一注册机制暴露给 InfiniTrain。

整体逻辑为:

InfiniTrain 定义 Backend 接口
        ↓
MACA 实现 Backend 接口
        ↓
启动时注册到 PrivateUse1
        ↓
InfiniTrain Dispatcher 统一调度

核心目标是在 InfiniTrain 中不增加 MACA 专用执行分支,而是通过统一 Backend 接口完成硬件能力接入,从而保持框架逻辑与硬件实现解耦。

测试与开发方式

测试基于 InfiniTrain 的 PrivateUse1 公共测试体系。

Backend 仓库主要负责将指定 provider 注册到测试程序中,而 Tensor、Autograd 等通用行为继续由 InfiniTrain 的 shared tests 进行验证。

这样不同 Backend 可以复用同一套接口语义和测试标准,减少厂商后端重复维护测试逻辑。

开发场景下,默认使用仓库中固定版本的 InfiniTrain submodule;同时支持通过 INFINITRAIN_SOURCE_DIR 指向本地 InfiniTrain 工作区,方便框架接口与 Backend 联调。

当前约束

当前设计下:

  • 一个 build tree 只选择一个 provider;
  • 一个进程最多注册一个 PrivateUse1 provider;
  • provider 通过统一接口接入 InfiniTrain;
  • 厂商相关实现保持在 InfiniTrain 核心仓库之外。

本 PR 以 MACA 验证上述 Backend 架构,为后续其他 accelerator provider 接入提供统一的目录结构、构建流程和运行时注册模式。

- align clang-format configuration with InfiniTrain
- format existing MACA backend sources
- add GitHub Actions format checking
- support .maca files in the formatting script
void RegisterTransformKernels();
void RegisterVocabParallelCrossEntropyKernels();

void RegisterMacaKernels() {

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

删掉

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

RegisterMacaKernels() 统一显式调用各 kernel 注册函数,通过普通符号引用保留对应静态库对象,使 MACA kernel 注册不依赖文件级静态初始化,不能删除

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nv 目前也是通过静态对象完成 kernel 注册的,MACA 理论上可以保持相同机制。这里需要保证链接阶段不会把仅包含静态注册对象的目标文件裁掉,例如静态库场景下通过 -Wl,--whole-archive 确保相关 object file 被链接进来。

既然 nv 侧已经通过构建系统保证了这一点,MACA 也更适合补齐对应的链接配置,而不是额外在 RegisterMacaKernels() 中显式调用所有注册函数。否则会形成两套注册机制,且后续新增 kernel 时还需要额外维护集中注册代码。

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

#define REGISTER_MACA_LINEAR_KERNEL(kernel_name) \
    REGISTER_KERNEL(                              \
        infini_train::Device::DeviceType::kPrivateUse1, \
        kernel_name,                              \
        infini_train::kernels::maca::kernel_name)

REGISTER_MACA_LINEAR_KERNEL(LinearForward)
REGISTER_MACA_LINEAR_KERNEL(LinearBackwardInput)
REGISTER_MACA_LINEAR_KERNEL(LinearBackwardWeight)
REGISTER_MACA_LINEAR_KERNEL(LinearBackwardBias)

#undef REGISTER_MACA_LINEAR_KERNEL

我看目前 cmakelist 里已经有 -Wl,--whole-archive 选项了,在 .maca 里改成类似这样也不行吗?

@chen2021673

Copy link
Copy Markdown
Author
img_v3_0214t_f6ab87d7-e0cd-4652-8c6e-dc15de21822g

- run upstream PrivateUse1 suites through BUILD_TEST
- reject incompatible CUDA builds and remove the local allocator test
- register BF16 autocast and fix pthread detection
- correct elementwise launches, embedding indexing, linear reductions, and layer norm gradients
- remove redundant tensor initialization and strengthen dtype and shape checks
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants