Feat/mars devices split

csrc/infinicore/src/context/context_impl.cc

@@ -16,6 +16,7 @@ constexpr std::array<Device::Type, static_cast<size_t>(Device::Type::kCount)> kD
     Device::Type::kCambricon,
     Device::Type::kAscend,
     Device::Type::kMetax,
+    Device::Type::kMars,
     Device::Type::kMoore,
     Device::Type::kIluvatar,
     Device::Type::kHygon,
@@ -213,6 +214,7 @@ ContextImpl::ContextImpl() {
     initializeDeviceType<Device::Type::kCambricon>();
     initializeDeviceType<Device::Type::kAscend>();
     initializeDeviceType<Device::Type::kMetax>();
+    initializeDeviceType<Device::Type::kMars>();
     initializeDeviceType<Device::Type::kMoore>();
     initializeDeviceType<Device::Type::kIluvatar>();
     initializeDeviceType<Device::Type::kHygon>();

csrc/infinicore/src/nn/rmsnorm.cc

@@ -29,6 +29,7 @@ void RMSNorm::forward_inplace(Tensor &x, Tensor &residual) const {
             || device_.type() == Device::Type::kNvidia
             || device_.type() == Device::Type::kIluvatar
             || device_.type() == Device::Type::kMetax
+            || device_.type() == Device::Type::kMars
             || device_.type() == Device::Type::kMoore
             || device_.type() == Device::Type::kCambricon
             || device_.type() == Device::Type::kHygon) {

csrc/infinicore/src/pybind11/device.hpp

@@ -17,6 +17,7 @@ inline void bind(py::module &m) {
         .value("CAMBRICON", Device::Type::kCambricon)
         .value("ASCEND", Device::Type::kAscend)
         .value("METAX", Device::Type::kMetax)
+        .value("MARS", Device::Type::kMars)
         .value("MOORE", Device::Type::kMoore)
         .value("ILUVATAR", Device::Type::kIluvatar)
         .value("HYGON", Device::Type::kHygon);

csrc/models/deepseek_v2/deepseek_v2_moe.cpp

@@ -17,6 +17,7 @@ bool supports_fused_deepseek_moe(infinicore::Device::Type device_type) {
     case infinicore::Device::Type::kHygon:
     case infinicore::Device::Type::kIluvatar:
     case infinicore::Device::Type::kMetax:
+    case infinicore::Device::Type::kMars:
     case infinicore::Device::Type::kMoore:
         return true;
     default:

python/infinicore/device.py

@@ -11,6 +11,7 @@
     "ascend": (_infinicore.Device.Type.ASCEND, "npu"),
     "npu": (_infinicore.Device.Type.ASCEND, "npu"),
     "metax": (_infinicore.Device.Type.METAX, "metax"),
+    "mars": (_infinicore.Device.Type.MARS, "mars"),
     "moore": (_infinicore.Device.Type.MOORE, "musa"),
     "musa": (_infinicore.Device.Type.MOORE, "musa"),
     "iluvatar": (_infinicore.Device.Type.ILUVATAR, "iluvatar"),

python/infinilm/base_config.py

@@ -149,27 +149,26 @@ def __init__(self):
         if self.enable_paged_attn and self.attn == "default":
             self.attn = "paged-attn"
 
-        # Force sync weight loading for Metax devices
-        self._force_sync_for_metax()
+        self._force_sync_weight_loading()
 
-    def _force_sync_for_metax(self):
-        """Force weight_load_mode to 'sync' for Metax devices."""
-        # Check if device is explicitly set to Metax
-        if self.device.lower() == "metax":
+    def _force_sync_weight_loading(self):
+        """Force synchronous weight loading on MetaX and Mars devices."""
+        device = self.device.lower()
+        if device in ("metax", "mars"):
             self.weight_load_mode = "sync"
             warnings.warn(
-                "Metax device detected: forcing weight_load_mode to 'sync'",
+                f"{device} device detected: forcing weight_load_mode to 'sync'",
                 UserWarning,
             )
             return
 
-        # Check if auto-detected device is Metax
-        if self.device.lower() == "auto":
+        if device == "auto":
             detected_device = self.detect_device()
-            if detected_device.lower() == "metax":
+            if detected_device.lower() in ("metax", "mars"):
                 self.weight_load_mode = "sync"
                 warnings.warn(
-                    "Auto-detected Metax device: forcing weight_load_mode to 'sync'",
+                    f"Auto-detected {detected_device} device: "
+                    "forcing weight_load_mode to 'sync'",
                     UserWarning,
                 )
 
@@ -193,8 +192,8 @@ def _add_common_args(self):
             type=str,
             default="auto",
             help=(
-                "device platform: auto, cpu, nvidia, metax, moore, iluvatar, "
-                "cambricon, ascend, hygon, or backend name "
+                "device platform: auto, cpu, nvidia, metax, mars, moore, "
+                "iluvatar, cambricon, ascend, hygon, or backend name "
                 "(cuda/mlu/musa/npu)"
             ),
         )
@@ -478,6 +477,7 @@ def detect_device(self):
                 return device_name
 
         env_checks = [
+            ("mars", ["HPCC_PATH", "HPCC_HOME"]),
             ("metax", ["MACA_PATH", "MACA_HOME", "MACA_ROOT"]),
             ("hygon", ["DTK_HOME", "DTK_PATH"]),
         ]
@@ -489,7 +489,8 @@ def detect_device(self):
             ("cambricon", ["cnmon"]),
             ("ascend", ["npu-smi"]),
             ("moore", ["mthreads-gmi"]),
-            ("metax", ["mx-smi", "ht-smi"]),
+            ("mars", ["ht-smi"]),
+            ("metax", ["mx-smi"]),
             ("hygon", ["hy-smi"]),
             ("iluvatar", ["ixsmi"]),
             ("nvidia", ["nvidia-smi"]),
@@ -504,7 +505,7 @@ def detect_device(self):
         return "cpu"
 
     def get_device_str(self, device):
-        """Convert device name to backend string (cuda/cpu/musa/mlu)"""
+        """Convert a platform name to the device string used by InfiniLM."""
         DEVICE_STR_MAP = {
             "cpu": "cpu",
             "cuda": "cuda",
@@ -514,7 +515,8 @@ def get_device_str(self, device):
             "nvidia": "cuda",
             "cambricon": "mlu",
             "ascend": "npu",
-            "metax": "cuda",
+            "metax": "metax",
+            "mars": "mars",
             "moore": "musa",
             "iluvatar": "cuda",
             "hygon": "cuda",

python/infinilm/exception_utils.py

@@ -24,7 +24,7 @@ def _iter_exception_chain(
 
 def is_oom_exception(e: BaseException) -> bool:
     """
-    Conservative OOM detector for MetaX allocator failures and CUDA/PyTorch OOMs.
+    Conservative OOM detector for MetaX/Mars allocator failures and CUDA/PyTorch OOMs.
     Checks exception type (when available) and message substrings across chained exceptions.
     """
     # PyTorch OOM exception type (only if torch is present in this environment)
@@ -42,7 +42,7 @@ def is_oom_exception(e: BaseException) -> bool:
     # Common patterns observed for allocator failures.
     # Keep this allowlist small to avoid hard-exiting on unrelated errors.
     patterns = (
-        # MetaX allocator
+        # MetaX/Mars and InfiniRT allocators
         "hcmalloc",
         "out of memory",
         # CUDA / driver / runtime alloc failures

python/infinilm/llm/model_runner/model_runner.py

@@ -4,6 +4,7 @@
 from typing import Any, Generator
 
 import infinicore
+
 from infinilm.cache.cache import PagedKVCacheConfig, StaticKVCacheConfig
 from infinilm.config.engine_config import EngineConfig
 from infinilm.distributed import DistConfig
@@ -138,7 +139,7 @@ def eos_token_id(self):
 
     def _init_device(self):
         """Initialize infinicore device and dtype."""
-        supported_devices = ["cpu", "cuda", "mlu", "musa", "npu"]
+        supported_devices = ["cpu", "cuda", "mlu", "musa", "npu", "metax", "mars"]
         device_str = self.config.device
         if device_str not in supported_devices:
             raise ValueError(

test/bench/backends/infinilm.py

@@ -17,6 +17,7 @@ def __init__(
         enable_paged_attn=False,
         enable_graph=False,
         attn_backend="default",
+        weight_load_mode="async",
     ):
         from infinilm import LLM
 
@@ -27,7 +28,8 @@ def __init__(
             "nvidia": "cuda",
             "cambricon": "mlu",
             "ascend": "npu",
-            "metax": "cuda",
+            "metax": "metax",
+            "mars": "mars",
             "moore": "musa",
             "iluvatar": "cuda",
             "hygon": "cuda",
@@ -65,6 +67,7 @@ def __init__(
             block_size=256,
             enable_graph=enable_graph,
             attn_backend=attn_backend,
+            weight_load_mode=weight_load_mode,
         )
         self.processor = self.model.engine.processor
         self.tokenizer = self.processor.get_tokenizer()

test/bench/test_benchmark.py

@@ -606,13 +606,14 @@ def main():
         model = VLLMBenchmark(cfg.model, device_str, cfg.tp, cfg.bench)
     elif cfg.backend in {"infinilm", "cpp", "python"}:
         model = InfiniLMBenchmark(
-            cfg.model,
-            device_str,
-            cfg.tp,
-            cfg.bench,
-            cfg.enable_paged_attn,
-            cfg.enable_graph,
-            cfg.attn,
+            model_dir_path=cfg.model,
+            device_type_str=device_str,
+            tensor_parallel_size=cfg.tp,
+            benchmark=cfg.bench,
+            enable_paged_attn=cfg.enable_paged_attn,
+            enable_graph=cfg.enable_graph,
+            attn_backend=cfg.attn,
+            weight_load_mode=cfg.weight_load_mode,
         )
     else:
         raise ValueError(f"Unsupported backend: {cfg.backend}")

test/models/qwen3_moe/attention_test.py

@@ -1,10 +1,10 @@
 import os
-import time
 import sys
+import time
+
 import safetensors
 import torch
-from transformers import AutoConfig
-from transformers import DynamicCache
+from transformers import AutoConfig, DynamicCache
 from transformers.models import qwen3_moe
 
 WARMUPS = 10
@@ -47,6 +47,11 @@ def get_args():
         action="store_true",
         help="Run metax test",
     )
+    parser.add_argument(
+        "--mars",
+        action="store_true",
+        help="Run Mars test",
+    )
     parser.add_argument(
         "--moore",
         action="store_true",
@@ -446,14 +451,16 @@ def benchmark_Qwen3attention_decode_torch(
         device = "cuda"
     elif args.metax:
         device = "cuda"
+    elif args.mars:
+        device = "cuda"
     elif args.moore:
         device = "musa"
-        import torch_musa
+        import torch_musa  # noqa: F401 - registers the torch.musa backend
     elif args.iluvatar:
         device = "cuda"
     else:
         print(
-            "Usage:  python test/models/qwen3_moe/attention_test.py [--cpu | --nvidia | --metax | --moore | --iluvatar] --model_path=<path/to/model_path>"
+            "Usage:  python test/models/qwen3_moe/attention_test.py [--cpu | --nvidia | --metax | --mars | --moore | --iluvatar] --model_path=<path/to/model_path>"
         )
         sys.exit(1)
 

test/models/qwen3_moe/moe_test.py

@@ -1,11 +1,11 @@
+import os
+import sys
 import time
-import torch
-import transformers
+
 import safetensors
-import os
+import torch
 from transformers import AutoConfig
 from transformers.models import qwen3_moe
-import sys
 
 WARMUPS = 10
 RUNS = 100
@@ -47,6 +47,11 @@ def get_args():
         action="store_true",
         help="Run metax test",
     )
+    parser.add_argument(
+        "--mars",
+        action="store_true",
+        help="Run Mars test",
+    )
     parser.add_argument(
         "--moore",
         action="store_true",
@@ -139,14 +144,16 @@ def benchmark_moe_torch(moe, testcase, device, dtype):
         device = "cuda"
     elif args.metax:
         device = "cuda"
+    elif args.mars:
+        device = "cuda"
     elif args.moore:
         device = "musa"
-        import torch_musa
+        import torch_musa  # noqa: F401 - registers the torch.musa backend
     elif args.iluvatar:
         device = "cuda"
     else:
         print(
-            "Usage:  python test/models/qwen3_moe/moe_test.py [--cpu | --nvidia | --metax | --moore | --iluvatar] --model_path=<path/to/model_path>"
+            "Usage:  python test/models/qwen3_moe/moe_test.py [--cpu | --nvidia | --metax | --mars | --moore | --iluvatar] --model_path=<path/to/model_path>"
         )
         sys.exit(1)
 

test/static/test_infinicore_python_contracts.py

@@ -136,6 +136,7 @@ def test_device_binding_uses_native_infini_rt_types(self) -> None:
             "kCambricon",
             "kAscend",
             "kMetax",
+            "kMars",
             "kMoore",
             "kIluvatar",
             "kHygon",
@@ -147,6 +148,12 @@ def test_device_binding_uses_native_infini_rt_types(self) -> None:
         self.assertIn("&Device::index", source)
         self.assertIn("&Device::ToString", source)
 
+        python_device = read_source("python/infinicore/device.py")
+        self.assertIn(
+            '"metax": (_infinicore.Device.Type.METAX, "metax")', python_device
+        )
+        self.assertIn('"mars": (_infinicore.Device.Type.MARS, "mars")', python_device)
+
     def test_dtype_binding_matches_native_infini_rt_set(self) -> None:
         source = read_source("csrc/infinicore/src/pybind11/dtype.hpp")
         native_names = (

test/static/test_infinicore_runtime_contracts.py

@@ -466,9 +466,7 @@ def test_static_graph_falls_back_for_unsupported_attention_configs(self) -> None
         source = read_source("csrc/engine/compiler/static_batching_compiler.cpp")
         self.assertIn("bool supports_static_graph_kv_cache(", source)
         self.assertIn("bool supports_static_graph_attention()", source)
-        cache_check = function_body(
-            source, "bool supports_static_graph_kv_cache("
-        )
+        cache_check = function_body(source, "bool supports_static_graph_kv_cache(")
         capability = function_body(source, "bool supports_static_graph_attention()")
         compile_body = function_body(source, "void StaticBatchingCompiler::compile()")
 
@@ -583,6 +581,7 @@ def test_context_owns_one_runtime_per_thread_and_device(self) -> None:
             "kCambricon",
             "kAscend",
             "kMetax",
+            "kMars",
             "kMoore",
             "kIluvatar",
             "kHygon",