feat(metax): enable standalone InfiniRT graph bridge - #1400
Open
gongchensu wants to merge 4 commits into
Open
Conversation
gongchensu
force-pushed
the
feat/metax-standalone-infinirt-graph-bridge
branch
2 times, most recently
from
July 24, 2026 01:57
ee2f301 to
c34131e
Compare
Complete the Mars/HPCC split across build configuration, runtime loading, ATen dispatch, operator tests, and environment setup. Validate external InfiniRT prefixes, support lib64 installs, and reject standalone InfiniOps on Mars instead of mapping it back to MetaX. Resolve the active Mars flash_attn_2_cuda extension and detect its C++ ABI from exported symbols because HPCC and flash-attn are versioned independently. Pass the Mars-only attention-mask and auxiliary arguments for MHA, varlen MHA, and KV-cache while leaving the MetaX and NVIDIA signatures unchanged. Add Mars device and preload coverage, and update the build and test documentation for the separate backend.
MetaX flash-attn packages do not expose one fixed C++ signature. The MACA 3.2.1 extension appends s_aux and also passes an attention mask to mha_fwd, while the split backend assumed the standard signature and failed to import _infinicore with an undefined symbol. Reuse symbol-based ABI detection for MetaX and Mars, but map the result to independent MetaX ABI defines so Mars configuration remains isolated. Pass the optional attention-mask and auxiliary arguments only when the selected extension exports them.
gongchensu
force-pushed
the
feat/metax-standalone-infinirt-graph-bridge
branch
from
July 29, 2026 02:38
c34131e to
b1b1436
Compare
voltjia
requested changes
Aug 3, 2026
Comment on lines
162
to
+209
| @@ -180,9 +182,9 @@ def preload_device(device_type: str) -> None: | |||
| Preload runtime libraries for a specific device type if needed. | |||
|
|
|||
| Args: | |||
| device_type: Device type name (e.g., "METAX", "ASCEND", etc.) | |||
| device_type: Device type name (e.g., "MARS" or "HYGON") | |||
| """ | |||
| if device_type == "METAX": | |||
| if device_type == "MARS": | |||
| preload_hpcc() | |||
| elif device_type == "HYGON": | |||
| preload_torch_hip() | |||
| @@ -200,11 +202,11 @@ def preload() -> None: | |||
| This function detects available device types and preloads their runtime libraries | |||
| if the environment indicates they are needed. | |||
| """ | |||
| # Device types that may require preload. Keep Hygon-only preloads gated by | |||
| # Hygon environment markers so other CUDA-compatible platforms do not load | |||
| # unrelated torch/flash-attn libraries during package import. | |||
| # Device types that may require preload. Gate each preload by its own | |||
| # environment markers so CUDA-compatible platforms do not load one another's | |||
| # runtime libraries during package import. | |||
| device_types = [ | |||
| "METAX", # HPCC/METAX | |||
| "MARS", | |||
Collaborator
There was a problem hiding this comment.
这些地方直接把 METAX 改成 MARS 了的话那 METAX 会不会受影响？
Collaborator
Author
There was a problem hiding this comment.
不会，这里之前的METAX实际就是MARS，实际调用的是preload_hpcc，其实只有MARS设备
Collaborator
There was a problem hiding this comment.
这里面怎么有的 API 里是 Mars 有的是 Hccl？
Collaborator
Author
There was a problem hiding this comment.
mars是设备名，hccl是sdk的集合通信API名，黄总说的是getMarsStream这个吗？这个倒是可以改成getHpccStream，表示接收的是HPCC runtime的hcStream_t，但是改成getHcclStream感觉不太对
Collaborator
Author
There was a problem hiding this comment.
嗯嗯，这个我去掉吧，不是特别需要
Collaborator
Author
There was a problem hiding this comment.
这里也主要是环境问题，不改的话加载不到正确的runtime，之前好像是在mars不行，改了之后两边测试过
Comment on lines
+111
to
+116
| | `--mars-flash-attn-abi=[detect\|standard\|extended]` | Mars flash-attn ABI；默认检查实际扩展库符号 | detect | ||
|
|
||
| Mars/HPCC 当前不支持 standalone InfiniOps 后端；配置 Mars 时请使用 | ||
| `--mars-gpu=true --infiniops=false`。 | ||
| 使用独立 InfiniRT 安装时应设置 `INFINI_RT_ROOT` 或 `--infinirt-root`；如果 | ||
| 显式前缀中缺少 `include/infini/rt.h` 或 `lib{,64}/libinfinirt.so`，配置会直接失败。 |
Collaborator
There was a problem hiding this comment.
这部分先去掉吧，有点突兀，包括 --mars-flash-attn-abi 也得想办法去掉。
Collaborator
There was a problem hiding this comment.
这个文件的修改必要嘛？之前为啥不需要来着，我看后面好像也有几个相关的文件跟这个类似，就是跟这个 preload 和 lib64 啥的有关系的，是不是什么环境问题？看看能不能不改动这些。
Collaborator
Author
There was a problem hiding this comment.
这里好像是之前在mars上，只有lib64是符合实际的，lib是不行，环境问题
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
2 participants
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.
No description provided.