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 05bcd76ad..de9acc8b9 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"]), @@ -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", 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 9ab1af157..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,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() 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()