feat(metax): enable standalone InfiniRT graph bridge

README.md

@@ -97,8 +97,8 @@ python scripts/install.py [XMAKE_CONFIG_FLAGS]
 | `--nv-gpu=[y\|n]`        | 是否编译英伟达 GPU 接口实现       | n
 | `--ascend-npu=[y\|n]`    | 是否编译昇腾 NPU 接口实现         | n
 | `--cambricon-mlu=[y\|n]` | 是否编译寒武纪 MLU 接口实现       | n
-| `--metax-gpu=[y\|n]`     | 是否编译沐曦 GPU 接口实现         | n
-| `--use-mc=[y\|n]`        | 是否沐曦 GPU 接口实现使用maca SDK | n
+| `--metax-gpu=[y\|n]`     | 是否编译 MetaX/MACA GPU 接口实现  | n
+| `--mars-gpu=[y\|n]`      | 是否编译 Mars/HPCC GPU 接口实现   | n
 | `--moore-gpu=[y\|n]`     | 是否编译摩尔线程 GPU 接口实现     | n
 | `--iluvatar-gpu=[y\|n]`  | 是否编译天数 GPU 接口实现         | n
 | `--qy-gpu=[y\|n]`        | 是否编译QY GPU 接口实现           | n
@@ -341,9 +341,9 @@ pip install -e .
 
 ```bash
 # 测试单算子
-python test/infinicore/ops/[operator].py [--bench | --debug | --verbose] [--cpu | --nvidia | --cambricon | --ascend | --iluvatar | --metax | --moore | --kunlun | --Hygon | --ali]
+python test/infinicore/ops/[operator].py [--bench | --debug | --verbose] [--cpu | --nvidia | --cambricon | --ascend | --iluvatar | --metax | --mars | --moore | --kunlun | --Hygon | --ali]
 # 测试全部算子
-python test/infinicore/run.py [--bench | --debug | --verbose] [--cpu | --nvidia | --cambricon | --ascend | --iluvatar | --metax | --moore | --kunlun | --ali]
+python test/infinicore/run.py [--bench | --debug | --verbose] [--cpu | --nvidia | --cambricon | --ascend | --iluvatar | --metax | --mars | --moore | --kunlun | --ali]
 ```
 
 使用 -h 查看更多参数。
@@ -352,9 +352,9 @@ python test/infinicore/run.py [--bench | --debug | --verbose] [--cpu | --nvidia
 
 ```shell
 # 测试单算子
-python test/infiniop/[operator].py [--cpu | --nvidia | --cambricon | --ascend | --iluvatar | --metax | --moore | --kunlun | --Hygon | --ali]
+python test/infiniop/[operator].py [--cpu | --nvidia | --cambricon | --ascend | --iluvatar | --metax | --mars | --moore | --kunlun | --Hygon | --ali]
 # 测试全部算子
-python scripts/python_test.py [--cpu | --nvidia | --cambricon | --ascend | --iluvatar | --metax | --moore | --kunlun | --Hygon | --ali]
+python scripts/python_test.py [--cpu | --nvidia | --cambricon | --ascend | --iluvatar | --metax | --mars | --moore | --kunlun | --Hygon | --ali]
 ```
 
 #### 通信库（InfiniCCL）测试

include/infinicore.h

@@ -48,6 +48,7 @@ typedef enum {
     INFINI_DEVICE_HYGON = 8,
     INFINI_DEVICE_QY = 9,
     INFINI_DEVICE_ALI = 10,
+    INFINI_DEVICE_MARS = 11,
     INFINI_DEVICE_TYPE_COUNT
 } infiniDevice_t;
 

include/infinicore/adaptor/aten_adaptor.hpp

@@ -9,7 +9,7 @@
 #if defined(ENABLE_HYGON_API)
 #include <c10/hip/HIPGuard.h>
 #include <c10/hip/HIPStream.h>
-#elif defined(ENABLE_NVIDIA_API) || defined(ENABLE_METAX_API) || defined(ENABLE_QY_API) || defined(ENABLE_ALI_API)
+#elif defined(ENABLE_NVIDIA_API) || defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) || defined(ENABLE_QY_API) || defined(ENABLE_ALI_API)
 #include <ATen/cuda/CUDAContext.h>
 #include <c10/cuda/CUDAGuard.h>
 #include <c10/cuda/CUDAStream.h>
@@ -53,6 +53,7 @@ inline at::Device to_at_device(const Device &device) {
     // Treat CUDA-compatible vendor devices as CUDA devices for ATen tensor interoperability.
     if (device.getType() == Device::Type::NVIDIA
         || device.getType() == Device::Type::METAX
+        || device.getType() == Device::Type::MARS
         || device.getType() == Device::Type::QY
         || device.getType() == Device::Type::ILUVATAR
         || device.getType() == Device::Type::ALI
@@ -85,7 +86,7 @@ at::Tensor to_aten_tensor(const infinicore::Tensor &t);
 
 #if defined(ENABLE_HYGON_API)
 c10::hip::HIPStream get_hip_stream();
-#elif defined(ENABLE_NVIDIA_API) || defined(ENABLE_METAX_API) || defined(ENABLE_QY_API) || defined(ENABLE_ALI_API)
+#elif defined(ENABLE_NVIDIA_API) || defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) || defined(ENABLE_QY_API) || defined(ENABLE_ALI_API)
 c10::cuda::CUDAStream get_cuda_stream();
 #endif
 

include/infinicore/adaptor/flash_attention_adaptor.hpp

@@ -2,15 +2,14 @@
 #pragma once
 #include "aten_adaptor.hpp"
 
-// NVIDIA flash-attn-nvidia.so uses namespace flash. The pip/MetaX flash_attn_2_cuda extension
-// exports the same entry points at global scope (no namespace), matching FLASH_NAMESPACE builds
-// where the namespace is empty.
+// NVIDIA flash-attn-nvidia.so uses namespace flash. The MetaX and Mars
+// extensions export the same entry points at global scope.
 //
 // Ascend (aclnn C API path) does NOT use the flash:: namespace at all — the aclnn kernels are
 // called directly from the dedicated *._ascend.cc implementation files, so this header is only
-// included by the NVIDIA/MetaX/QY code paths.  We still guard the namespace below so that
+// included by the NVIDIA/MetaX/Mars/QY code paths. We still guard the namespace below so that
 // existing code compiles unchanged when ENABLE_ASCEND_FLASH_ATTN is defined.
-#if !defined(ENABLE_METAX_API) && !defined(ENABLE_ASCEND_FLASH_ATTN)
+#if !defined(ENABLE_METAX_API) && !defined(ENABLE_MARS_API) && !defined(ENABLE_ASCEND_FLASH_ATTN)
 namespace flash {
 #endif
 std::vector<at::Tensor>
@@ -22,6 +21,10 @@ mha_fwd(at::Tensor &q,                   // batch_size x seqlen_q x num_heads x
         std::optional<at::Tensor> &softmax_lse_, // MetaX flash-attn dense fwd ABI includes an optional preallocated LSE tensor
 #endif
         std::optional<at::Tensor> &alibi_slopes_, // num_heads or batch_size x num_heads
+#if defined(INFINICORE_FLASH_ATTN_MARS_EXT) || defined(INFINICORE_FLASH_ATTN_METAX_EXT)
+        // Extended device ABIs accept an attention mask before the scalar options.
+        std::optional<at::Tensor> &attn_mask_,
+#endif
         const float p_dropout,
         const float softmax_scale,
         bool is_causal,
@@ -30,10 +33,10 @@ mha_fwd(at::Tensor &q,                   // batch_size x seqlen_q x num_heads x
         const float softcap,
         const bool return_softmax,
         std::optional<at::Generator> gen_
-#if defined(ENABLE_METAX_API) && defined(INFINICORE_HPCC_VERSION_MAJOR) && (INFINICORE_HPCC_VERSION_MAJOR >= 3)
-        // MetaX/Mars `flash_attn_2_cuda` (e.g. 2.6.x+mars) appends this argument vs upstream Dao-AILab flash-attn.
+#if defined(INFINICORE_FLASH_ATTN_MARS_EXT) || defined(INFINICORE_FLASH_ATTN_METAX_EXT) || defined(INFINICORE_FLASH_ATTN_METAX_S_AUX)
+        // Some device extensions append an auxiliary tensor after the generator.
         ,
-        std::optional<at::Tensor> &flash_attn_mars_ext_
+        std::optional<at::Tensor> &s_aux_
 #endif
 );
 
@@ -59,10 +62,10 @@ mha_varlen_fwd(at::Tensor &q,                               // total_q x num_hea
                const float softcap,
                const bool return_softmax,
                std::optional<at::Generator> gen_
-#if defined(ENABLE_METAX_API) && defined(INFINICORE_HPCC_VERSION_MAJOR) && (INFINICORE_HPCC_VERSION_MAJOR >= 3)
-               // MetaX/Mars `flash_attn_2_cuda` (e.g. 2.6.x+mars) appends this argument vs upstream Dao-AILab flash-attn.
+#if defined(INFINICORE_FLASH_ATTN_MARS_EXT) || defined(INFINICORE_FLASH_ATTN_METAX_EXT) || defined(INFINICORE_FLASH_ATTN_METAX_S_AUX)
+               // Some device extensions append an auxiliary tensor.
                ,
-               std::optional<at::Tensor> &flash_attn_mars_ext_
+               std::optional<at::Tensor> &s_aux_
 #endif
 );
 
@@ -134,14 +137,14 @@ mha_fwd_kvcache(at::Tensor &q,                                     // batch_size
                 const float softcap,
                 bool is_rotary_interleaved, // if true, rotary combines indices 0 & 1, else indices 0 & rotary_dim / 2
                 int num_splits
-#if defined(ENABLE_METAX_API) && defined(INFINICORE_HPCC_VERSION_MAJOR) && (INFINICORE_HPCC_VERSION_MAJOR >= 3)
-                // MetaX/Mars `flash_attn_2_cuda` (e.g. 2.6.x+mars) appends this argument vs upstream Dao-AILab flash-attn.
+#if defined(INFINICORE_FLASH_ATTN_MARS_EXT) || defined(INFINICORE_FLASH_ATTN_METAX_EXT) || defined(INFINICORE_FLASH_ATTN_METAX_S_AUX)
+                // Some device extensions append an auxiliary tensor.
                 ,
-                std::optional<at::Tensor> &flash_attn_mars_ext_
+                std::optional<at::Tensor> &s_aux_
 #endif
 );
 
-#if !defined(ENABLE_METAX_API) && !defined(ENABLE_ASCEND_FLASH_ATTN)
+#if !defined(ENABLE_METAX_API) && !defined(ENABLE_MARS_API) && !defined(ENABLE_ASCEND_FLASH_ATTN)
 } // namespace flash
 #endif
 #endif // ENABLE_FLASH_ATTN

include/infinicore/device.hpp

@@ -23,6 +23,7 @@ class Device {
         HYGON = INFINI_DEVICE_HYGON,
         QY = INFINI_DEVICE_QY,
         ALI = INFINI_DEVICE_ALI,
+        MARS = INFINI_DEVICE_MARS,
         COUNT = INFINI_DEVICE_TYPE_COUNT,
     };
 

python/infinicore/_preload.py

@@ -11,8 +11,7 @@
 
 def _candidate_prefixes(path: str) -> List[str]:
     """
-    Return HPCC install prefixes to search for libs.
-    Prefer HPCC_PATH; if absent and explicitly opted-in, fall back to /opt/hpcc.
+    Return unique runtime install prefixes to search for libraries.
     """
     prefixes: List[str] = []
     if path:
@@ -175,7 +174,9 @@ def _should_preload_device(device_type: str) -> bool:
     Check if preload is needed for a specific device type.
     """
     device_env_map = {
-        "METAX": ["HPCC_PATH", "INFINICORE_PRELOAD_HPCC"],  # HPCC/METAX
+        # The old combined MetaX backend used this hook for HPCC. After the
+        # split, HPCC belongs to Mars; MetaX/MACA does not need this preload.
+        "MARS": ["HPCC_PATH", "INFINICORE_PRELOAD_HPCC"],
         "HYGON": ["DTK_ROOT", "INFINICORE_PRELOAD_TORCH_HIP"],
         "ASCEND": ["ASCEND_HOME", "ASCEND_TOOLKIT_HOME"],
         "CAMBRICON": ["NEUWARE_HOME", "INFINICORE_PRELOAD_CAMBRICON"],
@@ -200,9 +201,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()
         preload_torch()
     elif device_type == "HYGON":
@@ -224,11 +225,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",
         "HYGON",
         "ASCEND",
         "CAMBRICON",

python/infinicore/device.py

@@ -83,6 +83,7 @@ def _from_infinicore_device(infinicore_device: _infinicore.Device):
     _infinicore.Device.Type.HYGON: "cuda",
     _infinicore.Device.Type.QY: "cuda",
     _infinicore.Device.Type.ALI: "cuda",
+    _infinicore.Device.Type.MARS: "cuda",
 }
 
 

scripts/metax_env.py

@@ -8,15 +8,12 @@ def _first_existing_dir(paths: list[str]) -> str:
     return ""
 
 
-def _metax_toolkit_root(use_mc: bool) -> str:
-    """Return toolkit root for MetaX builds (MACA when use-mc; otherwise HPCC)."""
-    if use_mc:
-        for key in ("MACA_PATH", "MACA_HOME", "MACA_ROOT"):
-            v = os.environ.get(key, "").strip()
-            if v:
-                return v
-        return _first_existing_dir(["/opt/maca"])
-    return _first_existing_dir(["/opt/hpcc"])
+def _toolkit_root(env_names: tuple[str, ...], fallback: str) -> str:
+    for key in env_names:
+        value = os.environ.get(key, "").strip()
+        if value:
+            return value
+    return _first_existing_dir([fallback])
 
 
 def _prepend_path_var(name: str, prefixes: list[str]) -> None:
@@ -37,16 +34,14 @@ def set_env_for_metax_gpu(
     """
     Prepend compiler include paths needed when building ATen-enabled C++ against torch headers.
 
-    This chooses paths based on xmake backend flags (e.g. --metax-gpu) and toolkit selection
-    (e.g. MetaX HPCC vs MACA when --use-mc=y).
+    MetaX always uses the MACA SDK. Mars/HPCC is configured separately.
     """
     d = parse_xmake_cli_flag_values(flags)
     if not truthy_flag_value(d.get("aten", "n")):
         return
 
     if truthy_flag_value(d.get("metax-gpu", "n")):
-        use_mc = truthy_flag_value(d.get("use-mc", "n"))
-        root = _metax_toolkit_root(use_mc=use_mc)
+        root = _toolkit_root(("MACA_PATH", "MACA_HOME", "MACA_ROOT"), "/opt/maca")
         if not root:
             return
         dirs = [
@@ -64,4 +59,32 @@ def set_env_for_metax_gpu(
         ]
         for var in ("CPATH", "CPLUS_INCLUDE_PATH", "C_INCLUDE_PATH"):
             _prepend_path_var(var, dirs)
+
+
+def set_env_for_mars_gpu(
+    flags: str,
+    *,
+    parse_xmake_cli_flag_values,
+    truthy_flag_value,
+) -> None:
+    """Prepend HPCC compatibility headers for ATen-enabled Mars builds."""
+    d = parse_xmake_cli_flag_values(flags)
+    if not truthy_flag_value(d.get("aten", "n")):
+        return
+    if not truthy_flag_value(d.get("mars-gpu", "n")):
+        return
+
+    root = _toolkit_root(("HPCC_PATH", "HPCC_HOME"), "/opt/hpcc")
+    if not root:
         return
+    dirs = [
+        os.path.join(root, "tools", "cu-bridge", "include"),
+        os.path.join(root, "include", "hcr"),
+        os.path.join(root, "include", "common"),
+        os.path.join(root, "include", "hcsparse"),
+        os.path.join(root, "include", "hcblas"),
+        os.path.join(root, "include", "hcsolver"),
+        os.path.join(root, "include"),
+    ]
+    for var in ("CPATH", "CPLUS_INCLUDE_PATH", "C_INCLUDE_PATH"):
+        _prepend_path_var(var, dirs)

scripts/set_env.py

@@ -1,7 +1,7 @@
 import os
 import platform
 
-from metax_env import set_env_for_metax_gpu
+from metax_env import set_env_for_mars_gpu, set_env_for_metax_gpu
 
 
 def _parse_xmake_cli_flag_values(flags: str):
@@ -38,12 +38,16 @@ def set_env_by_config(flags: str) -> None:
             parse_xmake_cli_flag_values=_parse_xmake_cli_flag_values,
             truthy_flag_value=_truthy_flag_value,
         )
-    else:
-        pass
+    if _truthy_flag_value(d.get("mars-gpu", "n")):
+        set_env_for_mars_gpu(
+            flags,
+            parse_xmake_cli_flag_values=_parse_xmake_cli_flag_values,
+            truthy_flag_value=_truthy_flag_value,
+        )
 
 
 def set_env():
-    if os.environ.get("INFINI_ROOT") == None:
+    if os.environ.get("INFINI_ROOT") is None:
         os.environ["INFINI_ROOT"] = os.path.expanduser("~/.infini")
 
     if platform.system() == "Windows":
@@ -56,10 +60,23 @@ def set_env():
         if new_path not in os.environ.get("PATH", ""):
             os.environ["PATH"] = f"{new_path}:{os.environ.get('PATH', '')}"
 
-        new_lib_path = os.path.expanduser(os.environ.get("INFINI_ROOT") + "/lib")
-        if new_lib_path not in os.environ.get("LD_LIBRARY_PATH", ""):
-            os.environ["LD_LIBRARY_PATH"] = (
-                f"{new_lib_path}:{os.environ.get('LD_LIBRARY_PATH', '')}"
-            )
+        new_lib_paths = []
+        infinirt_root = os.environ.get("INFINI_RT_ROOT")
+        if infinirt_root:
+            for subdir in ("lib", "lib64"):
+                candidate = os.path.join(infinirt_root, subdir)
+                if os.path.isdir(candidate):
+                    new_lib_paths.append(candidate)
+        new_lib_paths.append(os.path.expanduser(os.environ["INFINI_ROOT"] + "/lib"))
+
+        current_lib_paths = [
+            path
+            for path in os.environ.get("LD_LIBRARY_PATH", "").split(":")
+            if path
+        ]
+        for new_lib_path in reversed(new_lib_paths):
+            if new_lib_path not in current_lib_paths:
+                current_lib_paths.insert(0, new_lib_path)
+        os.environ["LD_LIBRARY_PATH"] = ":".join(current_lib_paths)
     else:
         raise RuntimeError("Unsupported platform.")

src/bridge/infini/rt.hpp

@@ -35,6 +35,8 @@ inline constexpr ::infini::rt::Device::Type translate_to(infiniDevice_t device)
         return ::infini::rt::Device::Type::kHygon;
     case INFINI_DEVICE_ALI:
         return ::infini::rt::Device::Type::kThead;
+    case INFINI_DEVICE_MARS:
+        return ::infini::rt::Device::Type::kMars;
     default:
         return ::infini::rt::Device::Type::kCount;
     }
@@ -60,6 +62,8 @@ inline constexpr infiniDevice_t translate_from(::infini::rt::Device::Type device
         return INFINI_DEVICE_HYGON;
     case ::infini::rt::Device::Type::kThead:
         return INFINI_DEVICE_ALI;
+    case ::infini::rt::Device::Type::kMars:
+        return INFINI_DEVICE_MARS;
     default:
         return INFINI_DEVICE_TYPE_COUNT;
     }

src/infiniccl-test/main.cpp

@@ -29,7 +29,7 @@ void printUsage() {
               << "[--master-addr <ip>] [--master-port <port>]" << std::endl
               << std::endl;
     std::cout << "  --<device>" << std::endl;
-    std::cout << "    Specify the device type --(nvidia|cambricon|ascend|metax|moore|iluvatar|qy|kunlun|hygon|ali)." << std::endl
+    std::cout << "    Specify the device type --(nvidia|cambricon|ascend|metax|mars|moore|iluvatar|qy|kunlun|hygon|ali)." << std::endl
               << std::endl;
     std::cout << "  --func" << std::endl
               << "    allreduce: local all-reduce test across all visible devices." << std::endl
@@ -64,6 +64,7 @@ ParsedArgs parseArgs(int argc, char *argv[]) {
         else PARSE_DEVICE("--cambricon", INFINI_DEVICE_CAMBRICON)
         else PARSE_DEVICE("--ascend", INFINI_DEVICE_ASCEND)
         else PARSE_DEVICE("--metax", INFINI_DEVICE_METAX)
+        else PARSE_DEVICE("--mars", INFINI_DEVICE_MARS)
         else PARSE_DEVICE("--moore", INFINI_DEVICE_MOORE)
         else PARSE_DEVICE("--iluvatar", INFINI_DEVICE_ILUVATAR)
         else PARSE_DEVICE("--qy", INFINI_DEVICE_QY)