Feat/mars devices split

README.md

@@ -38,7 +38,7 @@
   - 单次推理测试
     - llama示例
     ```bash
-    python examples/test_infer.py --device [cpu | nvidia | qy | metax | moore | iluvatar | ali | cambricon | hygon] --model=<path/to/model_dir>
+    python examples/test_infer.py --device [cpu | nvidia | qy | metax | mars | moore | iluvatar | ali | cambricon | hygon] --model=<path/to/model_dir>
     ```
     - 例如：
     ```bash
@@ -59,7 +59,7 @@
   - 推理服务测试
     - 启动推理服务
       ```bash
-      python python/infinilm/server/inference_server.py --device [cpu | nvidia | qy | metax | moore | iluvatar | ali | cambricon | hygon] --model=<path/to/model-dir> --max-new-tokens=MAX_TOKENS --max-batch-size=MAX_BATCH --tp=NDEV --temperature=TEMP --top-p=TOP_P --top-k=TOP_K --host=HOST --port=PORT
+      python python/infinilm/server/inference_server.py --device [cpu | nvidia | qy | metax | mars | moore | iluvatar | ali | cambricon | hygon] --model=<path/to/model-dir> --max-new-tokens=MAX_TOKENS --max-batch-size=MAX_BATCH --tp=NDEV --temperature=TEMP --top-p=TOP_P --top-k=TOP_K --host=HOST --port=PORT
       ```
 
     - 单卡示例：
@@ -90,7 +90,7 @@
   - 运行推理基准测试（C-Eval/MMLU）
 
     ```bash
-    python test/bench/test_benchmark.py --device [cpu | nvidia | qy | metax | moore | iluvatar | ali | cambricon | hygon] --model <path/to/model_dir> --bench {ceval|mmlu} [--backend cpp] [--tp N] [--subject SUBJECT] [--num-samples N] [--max-new-tokens N] [--output-csv PATH] [--cache-dir PATH]
+    python test/bench/test_benchmark.py --device [cpu | nvidia | qy | metax | mars | moore | iluvatar | ali | cambricon | hygon] --model <path/to/model_dir> --bench {ceval|mmlu} [--backend cpp] [--tp N] [--subject SUBJECT] [--num-samples N] [--max-new-tokens N] [--output-csv PATH] [--cache-dir PATH]
     ```
 
     - 参数说明：
@@ -160,16 +160,20 @@
 xmake && xmake install
 ```
 
+使用独立安装的 InfiniRT 时设置 `INFINI_RT_ROOT`。该前缀必须包含
+`include/infini/rt.h`，并在 `lib` 或 `lib64` 下包含 `libinfinirt.so`；
+路径无效时配置会直接失败，避免链接到 `INFINI_ROOT` 中残留的旧版本。
+
 - 运行模型推理测试
 
 ```bash
-python scripts/jiuge.py [--cpu | --nvidia | --qy | --cambricon | --ascend | --metax | --moore | --iluvatar | --kunlun | --hygon | --ali] path/to/model_dir [n_device]
+python scripts/jiuge.py [--cpu | --nvidia | --qy | --cambricon | --ascend | --metax | --mars | --moore | --iluvatar | --kunlun | --hygon | --ali] path/to/model_dir [n_device]
 ```
 
 - 部署模型推理服务
 
 ```bash
-python scripts/launch_server.py --model MODEL_PATH [-h] [--dev {cpu,nvidia,qy, cambricon,ascend,metax,moore,iluvatar,kunlun,hygon}] [--ndev NDEV] [--max-batch MAX_BATCH] [--max-new-tokens MAX_TOKENS]
+python scripts/launch_server.py --model MODEL_PATH [-h] [--dev {cpu,nvidia,qy,cambricon,ascend,metax,mars,moore,iluvatar,kunlun,hygon}] [--ndev NDEV] [--max-batch MAX_BATCH] [--max-new-tokens MAX_TOKENS]
 ```
 
 - 测试模型推理服务性能

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

@@ -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,7 +192,7 @@ def _add_common_args(self):
             type=str,
             default="auto",
             help=(
-                "device platform: auto, cpu, nvidia, qy, metax, moore, iluvatar, "
+                "device platform: auto, cpu, nvidia, qy, metax, mars, moore, iluvatar, "
                 "ali, cambricon, ascend, kunlun, 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"]),
             ("ali", ["ppu-smi"]),
             ("iluvatar", ["ixsmi"]),
@@ -517,6 +518,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

@@ -1,5 +1,5 @@
-import os
 import logging
+import os
 from typing import Iterator
 
 logger = logging.getLogger(__name__)
@@ -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",

scripts/deepseek.py

@@ -740,7 +740,7 @@ def destroy_model_instance(self):
 def test():
     if len(sys.argv) < 3:
         print(
-            "Usage: python deepseek.py [--cpu | --nvidia| --cambricon | --ascend | --metax | --moore] <path/to/model_dir> [n_device]"
+            "Usage: python deepseek.py [--cpu | --nvidia| --cambricon | --ascend | --metax | --mars | --moore] <path/to/model_dir> [n_device]"
         )
         sys.exit(1)
     model_path = sys.argv[2]
@@ -755,13 +755,15 @@ def test():
         device_type = DeviceType.DEVICE_TYPE_ASCEND
     elif sys.argv[1] == "--metax":
         device_type = DeviceType.DEVICE_TYPE_METAX
+    elif sys.argv[1] == "--mars":
+        device_type = DeviceType.DEVICE_TYPE_MARS
     elif sys.argv[1] == "--moore":
         device_type = DeviceType.DEVICE_TYPE_MOORE
     elif sys.argv[1] == "--iluvatar":
         device_type = DeviceType.DEVICE_TYPE_ILUVATAR
     else:
         print(
-            "Usage: python deepseek.py [--cpu | --nvidia| --cambricon | --ascend | --metax | --moore] <path/to/model_dir> [n_device]"
+            "Usage: python deepseek.py [--cpu | --nvidia| --cambricon | --ascend | --metax | --mars | --moore] <path/to/model_dir> [n_device]"
         )
         sys.exit(1)
 

scripts/jiuge.py

@@ -825,7 +825,7 @@ def destroy_model_instance(self):
 def test():
     if len(sys.argv) < 3:
         print(
-            "Usage: python jiuge.py [--cpu | --nvidia| --qy| --cambricon | --ascend | --metax | --moore | --iluvatar | --kunlun | --hygon] <path/to/model_dir> [n_device] [--verbose]"
+            "Usage: python jiuge.py [--cpu | --nvidia| --qy| --cambricon | --ascend | --metax | --mars | --moore | --iluvatar | --kunlun | --hygon] <path/to/model_dir> [n_device] [--verbose]"
         )
         sys.exit(1)
 
@@ -852,6 +852,8 @@ def test():
         device_type = DeviceType.DEVICE_TYPE_ASCEND
     elif sys.argv[1] == "--metax":
         device_type = DeviceType.DEVICE_TYPE_METAX
+    elif sys.argv[1] == "--mars":
+        device_type = DeviceType.DEVICE_TYPE_MARS
     elif sys.argv[1] == "--moore":
         device_type = DeviceType.DEVICE_TYPE_MOORE
     elif sys.argv[1] == "--iluvatar":
@@ -864,7 +866,7 @@ def test():
         device_type = DeviceType.DEVICE_TYPE_ALI
     else:
         print(
-            "Usage: python jiuge.py [--cpu | --nvidia| --qy| --cambricon | --ascend | --metax | --moore | --iluvatar | --kunlun | --hygon | --ali] <path/to/model_dir> [n_device] [--verbose]"
+            "Usage: python jiuge.py [--cpu | --nvidia| --qy| --cambricon | --ascend | --metax | --mars | --moore | --iluvatar | --kunlun | --hygon | --ali] <path/to/model_dir> [n_device] [--verbose]"
         )
         sys.exit(1)
 

scripts/jiuge_awq.py

@@ -334,7 +334,7 @@ def destroy_model_instance(self):
 def test():
     if len(sys.argv) < 3:
         print(
-            "Usage: python jiuge_awq.py [--cpu | --nvidia| --cambricon | --ascend | --metax | --moore] <path/to/model_dir> [n_device]"
+            "Usage: python jiuge_awq.py [--cpu | --nvidia| --cambricon | --ascend | --metax | --mars | --moore] <path/to/model_dir> [n_device]"
         )
         sys.exit(1)
     model_path = sys.argv[2]
@@ -349,13 +349,15 @@ def test():
         device_type = DeviceType.DEVICE_TYPE_ASCEND
     elif sys.argv[1] == "--metax":
         device_type = DeviceType.DEVICE_TYPE_METAX
+    elif sys.argv[1] == "--mars":
+        device_type = DeviceType.DEVICE_TYPE_MARS
     elif sys.argv[1] == "--moore":
         device_type = DeviceType.DEVICE_TYPE_MOORE
     elif sys.argv[1] == "--iluvatar":
         device_type = DeviceType.DEVICE_TYPE_ILUVATAR
     else:
         print(
-            "Usage: python main_jiuge_awq.py [--cpu | --nvidia| --cambricon | --ascend | --metax | --moore] <path/to/model_dir> [n_device]"
+            "Usage: python main_jiuge_awq.py [--cpu | --nvidia| --cambricon | --ascend | --metax | --mars | --moore] <path/to/model_dir> [n_device]"
         )
         sys.exit(1)
 

scripts/jiuge_gptq.py

@@ -334,7 +334,7 @@ def destroy_model_instance(self):
 def test():
     if len(sys.argv) < 3:
         print(
-            "Usage: python jiuge_gptq.py [--cpu | --nvidia | --cambricon | --ascend | --metax | --moore | --iluvatar] <path/to/model_dir> [n_device]"
+            "Usage: python jiuge_gptq.py [--cpu | --nvidia | --cambricon | --ascend | --metax | --mars | --moore | --iluvatar] <path/to/model_dir> [n_device]"
         )
         sys.exit(1)
     model_path = sys.argv[2]
@@ -349,13 +349,15 @@ def test():
         device_type = DeviceType.DEVICE_TYPE_ASCEND
     elif sys.argv[1] == "--metax":
         device_type = DeviceType.DEVICE_TYPE_METAX
+    elif sys.argv[1] == "--mars":
+        device_type = DeviceType.DEVICE_TYPE_MARS
     elif sys.argv[1] == "--moore":
         device_type = DeviceType.DEVICE_TYPE_MOORE
     elif sys.argv[1] == "--iluvatar":
         device_type = DeviceType.DEVICE_TYPE_ILUVATAR
     else:
         print(
-            "Usage: python jiuge_gptq.py [--cpu | --nvidia | --cambricon | --ascend | --metax | --moore | --iluvatar] <path/to/model_dir> [n_device]"
+            "Usage: python jiuge_gptq.py [--cpu | --nvidia | --cambricon | --ascend | --metax | --mars | --moore | --iluvatar] <path/to/model_dir> [n_device]"
         )
         sys.exit(1)
 

scripts/jiuge_ppl.py

@@ -11,6 +11,7 @@
     "cambricon": DeviceType.DEVICE_TYPE_CAMBRICON,
     "ascend": DeviceType.DEVICE_TYPE_ASCEND,
     "metax": DeviceType.DEVICE_TYPE_METAX,
+    "mars": DeviceType.DEVICE_TYPE_MARS,
     "moore": DeviceType.DEVICE_TYPE_MOORE,
     "iluvatar": DeviceType.DEVICE_TYPE_ILUVATAR,
     "kunlun": DeviceType.DEVICE_TYPE_KUNLUN,
@@ -24,6 +25,7 @@
     "cambricon": "mlu",
     "ascend": "npu",
     "metax": "cuda",
+    "mars": "cuda",
     "moore": "cuda",
     "iluvatar": "cuda",
     "kunlun": "cuda",

scripts/launch_server.py

@@ -1,5 +1,6 @@
 from jiuge import JiugeForCauslLM
 from jiuge_awq import JiugeAWQForCausalLM
+from jiuge_gptq import JiugeGPTQForCausalLM
 from libinfinicore_infer import DeviceType
 from infer_task import InferTask
 from kvcache_pool import KVCachePool
@@ -24,6 +25,7 @@
     "cambricon": DeviceType.DEVICE_TYPE_CAMBRICON,
     "ascend": DeviceType.DEVICE_TYPE_ASCEND,
     "metax": DeviceType.DEVICE_TYPE_METAX,
+    "mars": DeviceType.DEVICE_TYPE_MARS,
     "moore": DeviceType.DEVICE_TYPE_MOORE,
     "iluvatar": DeviceType.DEVICE_TYPE_ILUVATAR,
     "kunlun": DeviceType.DEVICE_TYPE_KUNLUN,

scripts/libinfinicore_infer/base.py

@@ -38,6 +38,7 @@ class DeviceType(ctypes.c_int):
     DEVICE_TYPE_HYGON = 8
     DEVICE_TYPE_QY = 9
     DEVICE_TYPE_ALI = 10
+    DEVICE_TYPE_MARS = 11
 
 
 class KVCacheCStruct(ctypes.Structure):

scripts/qwen3vl.py

@@ -888,7 +888,7 @@ def destroy_model_instance(self):
 def test():
     if len(sys.argv) < 3:
         print(
-            "Usage: python qwen3vl.py [--cpu | --nvidia| --cambricon | --ascend | --metax | --moore] <path/to/model_dir> [n_device]"
+            "Usage: python qwen3vl.py [--cpu | --nvidia| --cambricon | --ascend | --metax | --mars | --moore] <path/to/model_dir> [n_device]"
         )
         sys.exit(1)
     model_path = sys.argv[2]
@@ -903,13 +903,15 @@ def test():
         device_type = DeviceType.DEVICE_TYPE_ASCEND
     elif sys.argv[1] == "--metax":
         device_type = DeviceType.DEVICE_TYPE_METAX
+    elif sys.argv[1] == "--mars":
+        device_type = DeviceType.DEVICE_TYPE_MARS
     elif sys.argv[1] == "--moore":
         device_type = DeviceType.DEVICE_TYPE_MOORE
     elif sys.argv[1] == "--iluvatar":
         device_type = DeviceType.DEVICE_TYPE_ILUVATAR
     else:
         print(
-            "Usage: python qwen3vl.py [--cpu | --nvidia| --cambricon | --ascend | --metax | --moore] <path/to/model_dir> [n_device]"
+            "Usage: python qwen3vl.py [--cpu | --nvidia| --cambricon | --ascend | --metax | --mars | --moore] <path/to/model_dir> [n_device]"
         )
         sys.exit(1)
     ndev = int(sys.argv[3]) if len(sys.argv) > 3 else 1

scripts/test_ceval.py

@@ -76,7 +76,7 @@ def generate(self, conversation, max_steps, topp_=1.0, topk_=1, temperature_=1.0
 def test():
     if len(sys.argv) < 3:
         print(
-            "Usage: python test_ceval.py [--cpu | --nvidia| --cambricon | --ascend | --metax | --moore | --iluvatar | --kunlun | --hygon] <path/to/model_dir> [n_device]"
+            "Usage: python test_ceval.py [--cpu | --nvidia| --cambricon | --ascend | --metax | --mars | --moore | --iluvatar | --kunlun | --hygon] <path/to/model_dir> [n_device]"
         )
         sys.exit(1)
 
@@ -117,6 +117,8 @@ def test():
         device_type = DeviceType.DEVICE_TYPE_ASCEND
     elif device_arg == "--metax":
         device_type = DeviceType.DEVICE_TYPE_METAX
+    elif device_arg == "--mars":
+        device_type = DeviceType.DEVICE_TYPE_MARS
     elif device_arg == "--moore":
         device_type = DeviceType.DEVICE_TYPE_MOORE
     elif device_arg == "--iluvatar":
@@ -127,7 +129,7 @@ def test():
         device_type = DeviceType.DEVICE_TYPE_HYGON
     else:
         print(
-            "Usage: python test_ceval.py [--cpu | --nvidia| --cambricon | --ascend | --metax | --moore | --iluvatar | --kunlun | --hygon] <path/to/model_dir> [n_device]"
+            "Usage: python test_ceval.py [--cpu | --nvidia| --cambricon | --ascend | --metax | --mars | --moore | --iluvatar | --kunlun | --hygon] <path/to/model_dir> [n_device]"
         )
         sys.exit(1)
 

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