issue/1395 - six more operators for metax#1396
Open
Lfan-ke wants to merge 1 commit into
Open
Conversation
Author
|
@codex review please. read code at pr and repositories. |
|
You have reached your Codex usage limits for code reviews. You can see your limits in the Codex usage dashboard. |
Author
|
@codex review please. read code at pr and repositories. |
|
You have reached your Codex usage limits for code reviews. You can see your limits in the Codex usage dashboard. |
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>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
关联 issue:#1395
做了什么
为下面六个算子补上
metax/后端,它们此前只有nvidia/实现,在 Metax GPU 上不可用:mul_scalarelementwise/metax)moe_sumprepare_moe_inputrwkv5_wkvmamba_selective_scandequantize_gptq每个算子新增
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_t→hcStream_t、device::nvidia::Handle→device::metax::Handle、CHECK_CUDA→CHECK_METAX。三处平台约束
moe_sum/prepare_moe_input/dequantize_gptq的 nvidia 实现整体包在#ifdef ENABLE_NVIDIA_API之内。移植到.maca后该守卫不再成立,翻译单元被预处理器整体剔除;由于目标文件仍然生成、编译亦无诊断,故缺失直到链接期方才显现:hcStream_t→mcStream_t与hpcc_bfloat16→maca_bfloat16这层兼容宏由devices/metax/metax_ht2mc.h提供,而后者经metax_common.h引入。故metax_common.h须排在metax_kernel_common.h之前,否则在--use-mc=y下这些类型均不可识别:__nv_bfloat16的别名定义于metax_kernel_common.h,用到 bf16 的算子须包含之。测试
在 MetaX C500(MACA 3.5.3.20 / 驱动 3.8.30)上:
.maca是带-Werror编译的,六个算子全部编过;metax/后端从 130 增加到 136。一个需要说明的地方
六个算子里只有
dequantize_gptq在test/infiniop/下有测试文件,所以只有它能跑端到端验证。其余五个算子在仓库里本来就没有测试(161 个算子里目前有 90 个测试),这不是这个 PR 引入的缺口。如果维护者希望,我很乐意在这个 PR 里或者另开一个 PR 补上这五个算子的测试 —— 那样对 CPU / NVIDIA 后端同样有用。CPU 与 NVIDIA 路径未做任何改动。