From 07caa6a678b3588e5129d48d153764f1fefc3ff2 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 | 28 +++++++++++++++------------- test/bench/backends/infinilm.py | 1 + 2 files changed, 16 insertions(+), 13 deletions(-) diff --git a/python/infinilm/base_config.py b/python/infinilm/base_config.py index 3b8f6533d..fad038b83 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,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", 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 f3e5689a27d7e864644f36006e9f9b885c32c064 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 Propagate Mars through legacy launch scripts, server and device maps, OOM handling, model tests, DeepSeek MoE dispatch, and benchmark configuration. Forward weight_load_mode explicitly so Mars benchmarks can request synchronous loading. Validate INFINI_RT_ROOT and add include, link, and runtime search paths for lib and lib64 so InfiniLM cannot silently bind to a stale runtime under INFINI_ROOT. Add regression tests for Mars and MetaX detection and benchmark option forwarding, and restore the missing GPTQ server import uncovered during verification. --- README.md | 14 +++-- csrc/models/deepseek_v2/deepseek_v2_moe.cpp | 1 + python/infinilm/exception_utils.py | 6 +- scripts/deepseek.py | 6 +- scripts/jiuge.py | 6 +- scripts/jiuge_awq.py | 6 +- scripts/jiuge_gptq.py | 6 +- scripts/jiuge_ppl.py | 2 + scripts/launch_server.py | 2 + scripts/libinfinicore_infer/base.py | 1 + scripts/qwen3vl.py | 6 +- scripts/test_ceval.py | 6 +- test/bench/backends/infinilm.py | 2 + test/bench/test_benchmark.py | 15 ++--- test/models/qwen3_moe/attention_test.py | 9 ++- test/models/qwen3_moe/moe_test.py | 9 ++- test/test_base_config.py | 64 +++++++++++++++++++++ xmake.lua | 41 +++++++++++++ 18 files changed, 173 insertions(+), 29 deletions(-) create mode 100644 test/test_base_config.py diff --git a/README.md b/README.md index 4c1eaf1e1..a509ba57b 100644 --- a/README.md +++ b/README.md @@ -38,7 +38,7 @@ - 单次推理测试 - llama示例 ```bash - python examples/test_infer.py --device [cpu | nvidia | qy | metax | moore | iluvatar | ali | cambricon | hygon] --model= + python examples/test_infer.py --device [cpu | nvidia | qy | metax | mars | moore | iluvatar | ali | cambricon | hygon] --model= ``` - 例如： ```bash @@ -59,7 +59,7 @@ - 推理服务测试 - 启动推理服务 ```bash - python python/infinilm/server/inference_server.py --device [cpu | nvidia | qy | metax | moore | iluvatar | ali | cambricon | hygon] --model= --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= --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 --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 --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] ``` - 测试模型推理服务性能 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..69ea4589e 100644 --- a/python/infinilm/exception_utils.py +++ b/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", diff --git a/scripts/deepseek.py b/scripts/deepseek.py index 29d81ce15..f0dba4b13 100644 --- a/scripts/deepseek.py +++ b/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] [n_device]" + "Usage: python deepseek.py [--cpu | --nvidia| --cambricon | --ascend | --metax | --mars | --moore] [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] [n_device]" + "Usage: python deepseek.py [--cpu | --nvidia| --cambricon | --ascend | --metax | --mars | --moore] [n_device]" ) sys.exit(1) diff --git a/scripts/jiuge.py b/scripts/jiuge.py index 35b2c8ca3..e3d5f8a87 100644 --- a/scripts/jiuge.py +++ b/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] [n_device] [--verbose]" + "Usage: python jiuge.py [--cpu | --nvidia| --qy| --cambricon | --ascend | --metax | --mars | --moore | --iluvatar | --kunlun | --hygon] [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] [n_device] [--verbose]" + "Usage: python jiuge.py [--cpu | --nvidia| --qy| --cambricon | --ascend | --metax | --mars | --moore | --iluvatar | --kunlun | --hygon | --ali] [n_device] [--verbose]" ) sys.exit(1) diff --git a/scripts/jiuge_awq.py b/scripts/jiuge_awq.py index 5191efeba..161bb16b0 100644 --- a/scripts/jiuge_awq.py +++ b/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] [n_device]" + "Usage: python jiuge_awq.py [--cpu | --nvidia| --cambricon | --ascend | --metax | --mars | --moore] [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] [n_device]" + "Usage: python main_jiuge_awq.py [--cpu | --nvidia| --cambricon | --ascend | --metax | --mars | --moore] [n_device]" ) sys.exit(1) diff --git a/scripts/jiuge_gptq.py b/scripts/jiuge_gptq.py index e510847cc..ab171a454 100644 --- a/scripts/jiuge_gptq.py +++ b/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] [n_device]" + "Usage: python jiuge_gptq.py [--cpu | --nvidia | --cambricon | --ascend | --metax | --mars | --moore | --iluvatar] [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] [n_device]" + "Usage: python jiuge_gptq.py [--cpu | --nvidia | --cambricon | --ascend | --metax | --mars | --moore | --iluvatar] [n_device]" ) sys.exit(1) diff --git a/scripts/jiuge_ppl.py b/scripts/jiuge_ppl.py index 923d209c6..8bb7d1bb3 100644 --- a/scripts/jiuge_ppl.py +++ b/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", diff --git a/scripts/launch_server.py b/scripts/launch_server.py index d04d4f69d..a1dccad7f 100644 --- a/scripts/launch_server.py +++ b/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, diff --git a/scripts/libinfinicore_infer/base.py b/scripts/libinfinicore_infer/base.py index 2af0c238a..6e65e1431 100644 --- a/scripts/libinfinicore_infer/base.py +++ b/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): diff --git a/scripts/qwen3vl.py b/scripts/qwen3vl.py index 917bce24e..c7694e6b2 100644 --- a/scripts/qwen3vl.py +++ b/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] [n_device]" + "Usage: python qwen3vl.py [--cpu | --nvidia| --cambricon | --ascend | --metax | --mars | --moore] [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] [n_device]" + "Usage: python qwen3vl.py [--cpu | --nvidia| --cambricon | --ascend | --metax | --mars | --moore] [n_device]" ) sys.exit(1) ndev = int(sys.argv[3]) if len(sys.argv) > 3 else 1 diff --git a/scripts/test_ceval.py b/scripts/test_ceval.py index 0b05dacf9..d0090264c 100644 --- a/scripts/test_ceval.py +++ b/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] [n_device]" + "Usage: python test_ceval.py [--cpu | --nvidia| --cambricon | --ascend | --metax | --mars | --moore | --iluvatar | --kunlun | --hygon] [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] [n_device]" + "Usage: python test_ceval.py [--cpu | --nvidia| --cambricon | --ascend | --metax | --mars | --moore | --iluvatar | --kunlun | --hygon] [n_device]" ) sys.exit(1) 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..88cb5df45 100644 --- a/test/models/qwen3_moe/attention_test.py +++ b/test/models/qwen3_moe/attention_test.py @@ -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,6 +451,8 @@ 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 @@ -453,7 +460,7 @@ def benchmark_Qwen3attention_decode_torch( 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..9e3fafd2b 100644 --- a/test/models/qwen3_moe/moe_test.py +++ b/test/models/qwen3_moe/moe_test.py @@ -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,6 +144,8 @@ 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 @@ -146,7 +153,7 @@ def benchmark_moe_torch(moe, testcase, device, dtype): 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..39120ccf6 --- /dev/null +++ b/test/test_base_config.py @@ -0,0 +1,64 @@ +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_cuda_torch_device(self): + config = self.make_config() + self.assertEqual(config.get_device_str("mars"), "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() diff --git a/xmake.lua b/xmake.lua index aab1a0c70..353faa71e 100644 --- a/xmake.lua +++ b/xmake.lua @@ -1,6 +1,39 @@ add_requires("pybind11") local INFINI_ROOT = os.getenv("INFINI_ROOT") or (os.getenv(is_host("windows") and "HOMEPATH" or "HOME") .. "/.infini") +local INFINI_RT_ROOT = os.getenv("INFINI_RT_ROOT") + +local function validate_external_infinirt(xmake_os) + if not INFINI_RT_ROOT or INFINI_RT_ROOT == "" then + return + end + + local public_header = path.join(INFINI_RT_ROOT, "include", "infini", "rt.h") + if not xmake_os.isfile(public_header) then + xmake_os.raise("InfiniRT public header not found: " .. public_header) + end + + local library = path.join(INFINI_RT_ROOT, "lib", "libinfinirt.so") + if not xmake_os.isfile(library) then + library = path.join(INFINI_RT_ROOT, "lib64", "libinfinirt.so") + end + if not xmake_os.isfile(library) then + xmake_os.raise("InfiniRT library not found under: " .. INFINI_RT_ROOT) + end +end + +local function add_external_infinirt() + if not INFINI_RT_ROOT or INFINI_RT_ROOT == "" then + return + end + add_includedirs(INFINI_RT_ROOT.."/include") + for _, libdir in ipairs({INFINI_RT_ROOT.."/lib", INFINI_RT_ROOT.."/lib64"}) do + if os.isdir(libdir) then + add_linkdirs(libdir) + add_rpathdirs(libdir) + end + end +end set_toolchains("gcc") @@ -10,8 +43,12 @@ add_includedirs("third_party/json/single_include/") target("infinicore_infer") set_kind("shared") + on_load(function (target) + validate_external_infinirt(os) + end) add_includedirs("include", { public = false }) + add_external_infinirt() add_includedirs(INFINI_ROOT.."/include", { public = true }) add_linkdirs(INFINI_ROOT.."/lib") @@ -39,11 +76,15 @@ target("_infinilm") add_rules("python.module", {soabi = true}) set_languages("cxx17") set_kind("shared") + on_load(function (target) + validate_external_infinirt(os) + end) local INFINI_ROOT = os.getenv("INFINI_ROOT") or (os.getenv(is_host("windows") and "HOMEPATH" or "HOME") .. "/.infini") -- add_includedirs("csrc", { public = false }) -- add_includedirs("csrc/pybind11", { public = false }) + add_external_infinirt() add_includedirs(INFINI_ROOT.."/include", { public = true }) add_includedirs("include", { public = false }) -- spdlog is already included globally via add_includedirs at the top