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
@@ -108,6 +108,12 @@ python scripts/install.py [XMAKE_CONFIG_FLAGS]
 | `--ninetoothed=[y\|n]`   | 是否编译九齿实现                 | n
 | `--ccl=[y\|n]`           | 是否编译 InfiniCCL 通信库接口实现 | n
 | `--graph=[y\|n]`         | 是否编译 cuda graph 接口实现      | n
+| `--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`，配置会直接失败。
 
 ##### 手动安装底层库
 
@@ -268,9 +274,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 查看更多参数。
@@ -279,9 +285,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

@@ -8,7 +8,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)
+#elif defined(ENABLE_NVIDIA_API) || defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) || defined(ENABLE_QY_API)
 #include <ATen/cuda/CUDAContext.h>
 #include <c10/cuda/CUDAGuard.h>
 #include <c10/cuda/CUDAStream.h>
@@ -40,8 +40,8 @@ inline at::ScalarType to_at_dtype(DataType dtype) {
 
 inline at::Device to_at_device(const Device &device) {
     // PyTorch ATen only exposes standard device types (e.g. kCPU/kCUDA).
-    // Treat MetaX/QY devices as CUDA devices for ATen tensor interoperability.
-    if (device.getType() == Device::Type::NVIDIA || device.getType() == Device::Type::METAX || device.getType() == Device::Type::QY || device.getType() == Device::Type::HYGON) {
+    // Treat CUDA-compatible devices as CUDA devices for ATen 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::HYGON) {
         return at::Device(at::kCUDA, device.getIndex());
     } else if (device.getType() == Device::Type::CPU) {
         return at::Device(at::kCPU);
@@ -60,7 +60,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)
+#elif defined(ENABLE_NVIDIA_API) || defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) || defined(ENABLE_QY_API)
 c10::cuda::CUDAStream get_cuda_stream();
 #endif
 

include/infinicore/adaptor/flash_attention_adaptor.hpp

@@ -2,10 +2,9 @@
 #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.
-#if !defined(ENABLE_METAX_API)
+// NVIDIA flash-attn-nvidia.so uses namespace flash. MetaX and Mars extension
+// modules export the same entry points at global scope.
+#if !defined(ENABLE_METAX_API) && !defined(ENABLE_MARS_API)
 namespace flash {
 #endif
 std::vector<at::Tensor>
@@ -14,6 +13,10 @@ mha_fwd(at::Tensor &q,                            // batch_size x seqlen_q x num
         const at::Tensor &v,                      // batch_size x seqlen_k x num_heads_k x round_multiple(head_size, 8)
         std::optional<at::Tensor> &out_,          // batch_size x seqlen_q x num_heads x round_multiple(head_size, 8)
         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,
@@ -22,10 +25,10 @@ mha_fwd(at::Tensor &q,                            // batch_size x seqlen_q x num
         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
 );
 
@@ -51,10 +54,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
 );
 
@@ -126,14 +129,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)
+#if !defined(ENABLE_METAX_API) && !defined(ENABLE_MARS_API)
 } // 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:
@@ -30,14 +29,15 @@ def _candidate_prefixes(path: str) -> List[str]:
 def _try_load(paths: Iterable[str], name: str) -> bool:
     """Try to load a shared library from given paths or system search path."""
     for path in paths:
-        full = os.path.join(path, "lib", name)
-        if os.path.exists(full):
-            try:
-                ctypes.CDLL(full, mode=ctypes.RTLD_GLOBAL)
-                return True
-            except OSError:
-                # Try next candidate
-                continue
+        for subdir in ("lib", "lib64"):
+            full = os.path.join(path, subdir, name)
+            if os.path.exists(full):
+                try:
+                    ctypes.CDLL(full, mode=ctypes.RTLD_GLOBAL)
+                    return True
+                except OSError:
+                    # Try next candidate
+                    continue
     # Last resort: rely on loader search path
     try:
         ctypes.CDLL(name, mode=ctypes.RTLD_GLOBAL)
@@ -53,7 +53,9 @@ def preload_hpcc() -> None:
     This mirrors the behavior of torch's HPCC build that loads libtorch_global_deps.so,
     but avoids introducing a hard torch dependency. All failures are swallowed.
     """
-    hpcc_path = os.getenv("HPCC_PATH")
+    hpcc_path = os.getenv("HPCC_PATH") or os.getenv("HPCC_HOME")
+    if not hpcc_path and os.getenv("INFINICORE_PRELOAD_HPCC"):
+        hpcc_path = "/opt/hpcc"
     if not hpcc_path:
         return
 
@@ -157,7 +159,7 @@ 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
+        "MARS": ["HPCC_PATH", "HPCC_HOME", "INFINICORE_PRELOAD_HPCC"],
         "HYGON": ["DTK_ROOT", "INFINICORE_PRELOAD_TORCH_HIP"],
         # Add other device types here as needed:
         # "ASCEND": ["ASCEND_PATH"],
@@ -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",
         "HYGON",
         # Add other device types here as they are implemented:
         # "ASCEND",

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.")