Skip to content

issue/1395 - six more operators for metax#1396

Open
Lfan-ke wants to merge 1 commit into
InfiniTensor:mainfrom
Lfan-ke:issue/1395
Open

issue/1395 - six more operators for metax#1396
Lfan-ke wants to merge 1 commit into
InfiniTensor:mainfrom
Lfan-ke:issue/1395

Conversation

@Lfan-ke

@Lfan-ke Lfan-ke commented Jul 13, 2026

Copy link
Copy Markdown

关联 issue:#1395

做了什么

为下面六个算子补上 metax/ 后端,它们此前只有 nvidia/ 实现,在 Metax GPU 上不可用:

算子 说明
mul_scalar 逐元素乘标量(走 elementwise 框架,复用已合入的 elementwise/metax
moe_sum MoE top-k 专家输出求和
prepare_moe_input MoE 专家偏移 / 置换 / 概率整理
rwkv5_wkv RWKV-5 的 WKV 递推
mamba_selective_scan Mamba 选择性扫描
dequantize_gptq GPTQ 4-bit 权重反量化

每个算子新增 metax/*_metax.h + metax/*_metax.maca,并在 operator.cc 的四处派发(create / get-workspace / calculate / destroy)接入 INFINI_DEVICE_METAX。设备端 kernel 原样复用,operator.cc 每个恰好 +16 行、无删除。xmake/metax.lua 是 glob,不需要改构建。

这六个都不依赖 cuDNN / cuBLAS / CUTLASS / 内联 PTX,所以是同族算子已经确立的那条机械移植路径:运行时头(mc_runtime.h / hc_runtime.h)、devices/metax/*、命名空间、cudaStream_thcStream_tdevice::nvidia::Handledevice::metax::HandleCHECK_CUDACHECK_METAX

三处平台约束

moe_sum / prepare_moe_input / dequantize_gptq 的 nvidia 实现整体包在 #ifdef ENABLE_NVIDIA_API 之内。移植到 .maca 后该守卫不再成立,翻译单元被预处理器整体剔除;由于目标文件仍然生成、编译亦无诊断,故缺失直到链接期方才显现:

undefined symbol: op::moe_sum::metax::Descriptor::~Descriptor()

hcStream_tmcStream_thpcc_bfloat16maca_bfloat16 这层兼容宏由 devices/metax/metax_ht2mc.h 提供,而后者经 metax_common.h 引入。故 metax_common.h 须排在 metax_kernel_common.h 之前,否则在 --use-mc=y 下这些类型均不可识别:

unknown type name 'hcStream_t'; did you mean 'mcStream_t'?
unknown type name 'hpcc_bfloat16'; did you mean 'maca_bfloat16'?

__nv_bfloat16 的别名定义于 metax_kernel_common.h,用到 bf16 的算子须包含之。

测试

在 MetaX C500(MACA 3.5.3.20 / 驱动 3.8.30)上:

XMAKE_ROOT=y xmake f --metax-gpu=y --use-mc=y -y && xmake && xmake install
python test/infiniop/dequantize_gptq.py --metax

.maca 是带 -Werror 编译的,六个算子全部编过;metax/ 后端从 130 增加到 136。

MetaX C500 上的构建与测试

一个需要说明的地方

六个算子里只有 dequantize_gptqtest/infiniop/ 下有测试文件,所以只有它能跑端到端验证。其余五个算子在仓库里本来就没有测试(161 个算子里目前有 90 个测试),这不是这个 PR 引入的缺口。如果维护者希望,我很乐意在这个 PR 里或者另开一个 PR 补上这五个算子的测试 —— 那样对 CPU / NVIDIA 后端同样有用。

CPU 与 NVIDIA 路径未做任何改动。

@Lfan-ke Lfan-ke requested a review from a team July 13, 2026 12:30
@Lfan-ke

Lfan-ke commented Jul 13, 2026

Copy link
Copy Markdown
Author

@codex review please. read code at pr and repositories.

@chatgpt-codex-connector

Copy link
Copy Markdown

You have reached your Codex usage limits for code reviews. You can see your limits in the Codex usage dashboard.

@Lfan-ke

Lfan-ke commented Jul 13, 2026

Copy link
Copy Markdown
Author

@codex review please. read code at pr and repositories.

@chatgpt-codex-connector

Copy link
Copy Markdown

You have reached your Codex usage limits for code reviews. You can see your limits in the Codex usage dashboard.

@Lfan-ke Lfan-ke changed the title issue/1395 为六个算子补齐 Metax 后端 issue/1395 - feat(metax): add Metax backends for six operators Jul 14, 2026
@Lfan-ke Lfan-ke changed the title issue/1395 - feat(metax): add Metax backends for six operators issue/1395 - six more operators for metax Jul 14, 2026
mul_scalar, moe_sum, prepare_moe_input, rwkv5_wkv, mamba_selective_scan and dequantize_gptq only had nvidia/ implementations. Each gains metax/*_metax.h + metax/*_metax.maca and an INFINI_DEVICE_METAX branch in operator.cc; the device kernels are reused as they are. xmake/metax.lua globs, so the build needs no change.

Three points the port has to get right, all of which the same-family operators already established:

- #ifdef ENABLE_NVIDIA_API wraps the whole nvidia translation unit, so a copied .maca compiles and links to nothing until the guard is dropped.
- metax_common.h must precede metax_kernel_common.h: the hcStream_t -> mcStream_t and hpcc_bfloat16 -> maca_bfloat16 shims come from devices/metax/metax_ht2mc.h, which the former pulls in.
- __nv_bfloat16 is aliased in metax_kernel_common.h.

The metax headers carry their own include guards rather than the nvidia ones, since operator.cc includes both under their respective vendor defines and a shared guard would skip the metax header wholesale.

Signed-off-by: 林晨 (Leo Cheng) <chengkelfan@qq.com>
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.

1 participant