From 777473f4c5dc4ef29e2f77692ac3a217d0a5f978 Mon Sep 17 00:00:00 2001 From: gongchensu Date: Tue, 21 Jul 2026 16:05:41 +0800 Subject: [PATCH 1/2] feat: add Mars device selection --- python/infinilm/base_config.py | 30 ++++++++++++++++-------------- test/bench/backends/infinilm.py | 1 + 2 files changed, 17 insertions(+), 14 deletions(-) diff --git a/python/infinilm/base_config.py b/python/infinilm/base_config.py index 05bcd76ad..ba802d7d6 100644 --- a/python/infinilm/base_config.py +++ b/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"]), @@ -515,6 +516,7 @@ def get_device_str(self, device): "cambricon": "mlu", "ascend": "npu", "metax": "cuda", + "mars": "cuda", "moore": "musa", "iluvatar": "cuda", "hygon": "cuda", diff --git a/test/bench/backends/infinilm.py b/test/bench/backends/infinilm.py index 9ab1af157..a4b010f61 100644 --- a/test/bench/backends/infinilm.py +++ b/test/bench/backends/infinilm.py @@ -28,6 +28,7 @@ def __init__( "cambricon": "mlu", "ascend": "npu", "metax": "cuda", + "mars": "cuda", "moore": "musa", "iluvatar": "cuda", "hygon": "cuda", From a22b2e2938604f9401182a7b523a179df0aa30a3 Mon Sep 17 00:00:00 2001 From: gongchensu Date: Thu, 23 Jul 2026 09:40:17 +0800 Subject: [PATCH 2/2] fix(mars): complete device integration Adapt Mars selection to the modern embedded InfiniCore runtime. Preserve MetaX and Mars as distinct device names, expose and initialize the Mars runtime type, and include Mars in the affected model dispatch paths. Forward synchronous weight loading through the benchmark backend and extend OOM handling and model tests for Mars. Add regression coverage for MetaX/Mars detection, device mapping, embedded runtime registration, and benchmark option forwarding. --- csrc/infinicore/src/context/context_impl.cc | 2 + csrc/infinicore/src/nn/rmsnorm.cc | 1 + csrc/infinicore/src/pybind11/device.hpp | 1 + csrc/models/deepseek_v2/deepseek_v2_moe.cpp | 1 + python/infinicore/device.py | 1 + python/infinilm/base_config.py | 6 +- python/infinilm/exception_utils.py | 4 +- .../infinilm/llm/model_runner/model_runner.py | 3 +- test/bench/backends/infinilm.py | 6 +- test/bench/test_benchmark.py | 15 ++-- test/models/qwen3_moe/attention_test.py | 17 +++-- test/models/qwen3_moe/moe_test.py | 19 ++++-- .../test_infinicore_python_contracts.py | 7 ++ .../test_infinicore_runtime_contracts.py | 5 +- test/test_base_config.py | 68 +++++++++++++++++++ 15 files changed, 127 insertions(+), 29 deletions(-) create mode 100644 test/test_base_config.py diff --git a/csrc/infinicore/src/context/context_impl.cc b/csrc/infinicore/src/context/context_impl.cc index 919ff07b8..fade429c8 100644 --- a/csrc/infinicore/src/context/context_impl.cc +++ b/csrc/infinicore/src/context/context_impl.cc @@ -16,6 +16,7 @@ constexpr std::array(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(); initializeDeviceType(); initializeDeviceType(); + initializeDeviceType(); initializeDeviceType(); initializeDeviceType(); initializeDeviceType(); diff --git a/csrc/infinicore/src/nn/rmsnorm.cc b/csrc/infinicore/src/nn/rmsnorm.cc index a18e183f2..28e77d791 100644 --- a/csrc/infinicore/src/nn/rmsnorm.cc +++ b/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) { diff --git a/csrc/infinicore/src/pybind11/device.hpp b/csrc/infinicore/src/pybind11/device.hpp index ec1d7fd6f..347586dcc 100644 --- a/csrc/infinicore/src/pybind11/device.hpp +++ b/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); diff --git a/csrc/models/deepseek_v2/deepseek_v2_moe.cpp b/csrc/models/deepseek_v2/deepseek_v2_moe.cpp index 7cc8463ce..0f1e91adf 100644 --- a/csrc/models/deepseek_v2/deepseek_v2_moe.cpp +++ b/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: diff --git a/python/infinicore/device.py b/python/infinicore/device.py index 58f8069cf..9b1aa68f9 100644 --- a/python/infinicore/device.py +++ b/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"), diff --git a/python/infinilm/base_config.py b/python/infinilm/base_config.py index ba802d7d6..de9acc8b9 100644 --- a/python/infinilm/base_config.py +++ b/python/infinilm/base_config.py @@ -505,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", @@ -515,8 +515,8 @@ def get_device_str(self, device): "nvidia": "cuda", "cambricon": "mlu", "ascend": "npu", - "metax": "cuda", - "mars": "cuda", + "metax": "metax", + "mars": "mars", "moore": "musa", "iluvatar": "cuda", "hygon": "cuda", diff --git a/python/infinilm/exception_utils.py b/python/infinilm/exception_utils.py index b66ac54b4..beb31438f 100644 --- a/python/infinilm/exception_utils.py +++ b/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 diff --git a/python/infinilm/llm/model_runner/model_runner.py b/python/infinilm/llm/model_runner/model_runner.py index 45f1eaee1..402789fbc 100644 --- a/python/infinilm/llm/model_runner/model_runner.py +++ b/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( diff --git a/test/bench/backends/infinilm.py b/test/bench/backends/infinilm.py index a4b010f61..eab40c5a4 100644 --- a/test/bench/backends/infinilm.py +++ b/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,8 +28,8 @@ def __init__( "nvidia": "cuda", "cambricon": "mlu", "ascend": "npu", - "metax": "cuda", - "mars": "cuda", + "metax": "metax", + "mars": "mars", "moore": "musa", "iluvatar": "cuda", "hygon": "cuda", @@ -66,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() diff --git a/test/bench/test_benchmark.py b/test/bench/test_benchmark.py index b77b11cde..08c2b38b5 100644 --- a/test/bench/test_benchmark.py +++ b/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}") diff --git a/test/models/qwen3_moe/attention_test.py b/test/models/qwen3_moe/attention_test.py index 26f66e406..d8c815d55 100644 --- a/test/models/qwen3_moe/attention_test.py +++ b/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=" + "Usage: python test/models/qwen3_moe/attention_test.py [--cpu | --nvidia | --metax | --mars | --moore | --iluvatar] --model_path=" ) sys.exit(1) diff --git a/test/models/qwen3_moe/moe_test.py b/test/models/qwen3_moe/moe_test.py index 4e0adaf46..1b1804a28 100644 --- a/test/models/qwen3_moe/moe_test.py +++ b/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=" + "Usage: python test/models/qwen3_moe/moe_test.py [--cpu | --nvidia | --metax | --mars | --moore | --iluvatar] --model_path=" ) sys.exit(1) diff --git a/test/static/test_infinicore_python_contracts.py b/test/static/test_infinicore_python_contracts.py index eb55e4ddf..f9a857965 100644 --- a/test/static/test_infinicore_python_contracts.py +++ b/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 = ( diff --git a/test/static/test_infinicore_runtime_contracts.py b/test/static/test_infinicore_runtime_contracts.py index 474b0c526..90d39435b 100644 --- a/test/static/test_infinicore_runtime_contracts.py +++ b/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", diff --git a/test/test_base_config.py b/test/test_base_config.py new file mode 100644 index 000000000..29b3c86d5 --- /dev/null +++ b/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_distinct_infinicore_device(self): + config = self.make_config() + self.assertEqual(config.get_device_str("mars"), "mars") + + def test_metax_uses_distinct_infinicore_device(self): + config = self.make_config(device="metax") + self.assertEqual(config.get_device_str("metax"), "metax") + + +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()