Feat/mars devices split

csrc/models/deepseek_v2/deepseek_v2_moe.cpp

@@ -18,6 +18,7 @@ bool supports_fused_deepseek_moe(infinicore::Device::Type device_type) {
     case infinicore::Device::Type::HYGON:
     case infinicore::Device::Type::ILUVATAR:
     case infinicore::Device::Type::METAX:
+    case infinicore::Device::Type::MARS:
     case infinicore::Device::Type::MOORE:
         return true;
     default:

python/infinilm/base_config.py

@@ -155,27 +155,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,
                 )
 
@@ -199,8 +198,8 @@ def _add_common_args(self):
             type=str,
             default="auto",
             help=(
-                "device platform: auto, cpu, nvidia, qy, metax, moore, iluvatar, "
-                "ali, cambricon, ascend, kunlun, hygon, or backend name "
+                "device platform: auto, cpu, nvidia, qy, metax, mars, moore, "
+                "iluvatar, ali, cambricon, ascend, kunlun, hygon, or backend name "
                 "(cuda/mlu/musa/npu)"
             ),
         )
@@ -519,6 +518,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"]),
         ]
@@ -530,7 +530,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"]),
             ("ali", ["ppu-smi"]),
             ("iluvatar", ["ixsmi"]),
@@ -558,6 +559,7 @@ def get_device_str(self, device):
             "cambricon": "mlu",
             "ascend": "npu",
             "metax": "cuda",
+            "mars": "cuda",
             "moore": "musa",
             "iluvatar": "cuda",
             "kunlun": "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 / infinirt allocator
+        # MetaX/Mars and InfiniRT allocators
         "hcmalloc",
         "infinirtmalloc",
         "out of memory",

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
 
@@ -28,6 +29,7 @@ def __init__(
             "cambricon": "mlu",
             "ascend": "npu",
             "metax": "cuda",
+            "mars": "cuda",
             "moore": "musa",
             "iluvatar": "cuda",
             "kunlun": "cuda",
@@ -60,6 +62,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/test_base_config.py

@@ -0,0 +1,68 @@
+import json
+import os
+import tempfile
+import unittest
+import warnings
+from pathlib import Path
+from unittest import mock
+
+from bench.backends.infinilm import InfiniLMBenchmark
+from infinilm.base_config import BaseConfig
+
+
+class TestMarsBaseConfig(unittest.TestCase):
+    def make_config(self, device="mars", weight_load_mode="async"):
+        config = BaseConfig.__new__(BaseConfig)
+        config.device = device
+        config.weight_load_mode = weight_load_mode
+        return config
+
+    def test_explicit_mars_forces_sync_weight_loading(self):
+        config = self.make_config()
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore")
+            config._force_sync_weight_loading()
+        self.assertEqual(config.weight_load_mode, "sync")
+
+    def test_auto_detects_mars_from_hpcc(self):
+        config = self.make_config(device="auto")
+        config._torch_device_available = lambda _device_type: False
+        with mock.patch.dict(os.environ, {"HPCC_PATH": "/opt/hpcc"}, clear=True):
+            with mock.patch("infinilm.base_config.shutil.which", return_value=None):
+                self.assertEqual(config.detect_device(), "mars")
+
+    def test_auto_detects_metax_from_maca(self):
+        config = self.make_config(device="auto")
+        config._torch_device_available = lambda _device_type: False
+        with mock.patch.dict(os.environ, {"MACA_PATH": "/opt/maca"}, clear=True):
+            with mock.patch("infinilm.base_config.shutil.which", return_value=None):
+                self.assertEqual(config.detect_device(), "metax")
+
+    def test_mars_uses_legacy_cuda_device(self):
+        config = self.make_config()
+        self.assertEqual(config.get_device_str("mars"), "cuda")
+
+    def test_metax_uses_legacy_cuda_device(self):
+        config = self.make_config(device="metax")
+        self.assertEqual(config.get_device_str("metax"), "cuda")
+
+
+class TestInfiniLMBenchmark(unittest.TestCase):
+    @mock.patch("infinilm.LLM")
+    def test_forwards_sync_weight_loading(self, llm):
+        processor = mock.Mock()
+        processor.get_tokenizer.return_value = mock.Mock()
+        llm.return_value.engine.processor = processor
+
+        with tempfile.TemporaryDirectory() as model_dir:
+            Path(model_dir, "config.json").write_text(
+                json.dumps({"max_position_embeddings": 2048}),
+                encoding="utf-8",
+            )
+            InfiniLMBenchmark(model_dir, weight_load_mode="sync")
+
+        self.assertEqual(llm.call_args.kwargs["weight_load_mode"], "sync")
+
+
+if __name__ == "__main__":
+    unittest.main()