From c0ff772d13d5e5abfc29c048de6676c50025258e 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 4ae7665c0..04b1e6cfa 100644 --- a/python/infinilm/base_config.py +++ b/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", diff --git a/test/bench/backends/infinilm.py b/test/bench/backends/infinilm.py index fbd99379d..8a620aca1 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", "kunlun": "cuda", From 21f34ee03b850c43555d1d69b75f266c8226d544 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 upstream/main\047s external InfiniCore runtime. Preserve MetaX and Mars as distinct platform names while mapping both through the legacy CUDA-compatible Python device entry point, 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, legacy device mapping, and benchmark option forwarding. --- csrc/models/deepseek_v2/deepseek_v2_moe.cpp | 1 + python/infinilm/exception_utils.py | 4 +- test/bench/backends/infinilm.py | 2 + test/bench/test_benchmark.py | 15 ++--- test/models/qwen3_moe/attention_test.py | 17 ++++-- test/models/qwen3_moe/moe_test.py | 19 ++++-- test/test_base_config.py | 68 +++++++++++++++++++++ 7 files changed, 106 insertions(+), 20 deletions(-) create mode 100644 test/test_base_config.py diff --git a/csrc/models/deepseek_v2/deepseek_v2_moe.cpp b/csrc/models/deepseek_v2/deepseek_v2_moe.cpp index a18351ac8..ba4ebcf8d 100644 --- a/csrc/models/deepseek_v2/deepseek_v2_moe.cpp +++ b/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: diff --git a/python/infinilm/exception_utils.py b/python/infinilm/exception_utils.py index fef44d248..b0a8f899c 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 / infinirt allocator + # MetaX/Mars and InfiniRT allocators "hcmalloc", "infinirtmalloc", "out of memory", diff --git a/test/bench/backends/infinilm.py b/test/bench/backends/infinilm.py index 8a620aca1..229a51807 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 @@ -61,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() 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/test_base_config.py b/test/test_base_config.py new file mode 100644 index 000000000..96b6e8d1b --- /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_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()