From 6c2e5fcd91372f1d08a1d0516fa3888c48bd5b53 Mon Sep 17 00:00:00 2001 From: Chen Haozhe Date: Wed, 19 Nov 2025 04:17:47 +0300 Subject: [PATCH] [feat][Ascend][Mindspore]: support model-impl of mindspore (#9234) --- docs/references/mindspore_models.md | 162 ++++++++++ python/sglang/srt/configs/model_config.py | 1 + .../srt/model_executor/mindspore_runner.py | 118 +++++++ .../sglang/srt/model_executor/model_runner.py | 20 ++ python/sglang/srt/model_loader/utils.py | 4 +- python/sglang/srt/models/mindspore.py | 305 ++++++++++++++++++ python/sglang/srt/server_args.py | 4 + test/srt/ascend/test_mindspore_models.py | 64 ++++ test/srt/run_suite.py | 3 + 9 files changed, 680 insertions(+), 1 deletion(-) create mode 100644 docs/references/mindspore_models.md create mode 100644 python/sglang/srt/model_executor/mindspore_runner.py create mode 100644 python/sglang/srt/models/mindspore.py create mode 100644 test/srt/ascend/test_mindspore_models.py diff --git a/docs/references/mindspore_models.md b/docs/references/mindspore_models.md new file mode 100644 index 000000000..417cb94ad --- /dev/null +++ b/docs/references/mindspore_models.md @@ -0,0 +1,162 @@ +# MindSpore Models + +## Introduction + +SGLang support run MindSpore framework models, this doc guide users to run mindspore models with SGLang. + +## Requirements + +MindSpore with SGLang current only support Ascend Npu device, users need first install Ascend CANN software packages. +The CANN software packages can download from the [Ascend Official Websites](https://www.hiascend.com). The version depends on the MindSpore version [MindSpore Installation](https://www.mindspore.cn/install) + +## Supported Models + +Currently, the following models are supported: + +- **Qwen3**: Dense models supported. MoE models coming soon. +- *More models coming soon...* + +## Installation + +> **Note**: Currently, MindSpore models are provided by an independent package `sgl-mindspore`, which needs to be installed separately. + +```shell +git clone https://github.com/chz34/sgl-mindspore.git +cd sgl-mindspore +pip install -e . +``` + +You will need to install the following packages, due to the support of tensor conversion through `dlpack` on 3rd devices, the minimum version of `PyTorch` is 2.7.1 + +```shell +pip install mindspore +pip install "torch>=2.7.1" +pip install "torch_npu>=2.7.1" +pip install triton_ascend +``` + +```shell +pip install -e "python[all_npu]" +``` + +## Run Model + +Current SGLang-MindSpore support Qwen3 dense model, this doc uses Qwen3-8B as example. + +### Offline infer + +Use the following script for offline infer: + +```python +import sglang as sgl + +# Initialize the engine with MindSpore backend +llm = sgl.Engine( + model_path="/path/to/your/model", # Local model path + device="npu", # Use NPU device + model_impl="mindspore", # MindSpore implementation + attention_backend="ascend", # Attention backend + tp_size=1, # Tensor parallelism size + dp_size=1 # Data parallelism size +) + +# Generate text +prompts = [ + "Hello, my name is", + "The capital of France is", + "The future of AI is" +] + +sampling_params = {"temperature": 0.01, "top_p": 0.9} +outputs = llm.generate(prompts, sampling_params) + +for prompt, output in zip(prompts, outputs): + print(f"Prompt: {prompt}") + print(f"Generated: {output['text']}") + print("---") +``` + +### Start server + +Launch a server with MindSpore backend: + +```bash +# Basic server startup +python3 -m sglang.launch_server \ + --model-path /path/to/your/model \ + --host 0.0.0.0 \ + --device npu \ + --model-impl mindspore \ + --attention-backend ascend \ + --tp-size 1 \ + --dp-size 1 +``` + +For distributed server with multiple nodes: + +```bash +# Multi-node distributed server +python3 -m sglang.launch_server \ + --model-path /path/to/your/model \ + --host 0.0.0.0 \ + --device npu \ + --model-impl mindspore \ + --attention-backend ascend \ + --dist-init-addr 127.0.0.1:29500 \ + --nnodes 2 \ + --node-rank 0 \ + --tp-size 4 \ + --dp-size 2 +``` + +## Troubleshooting + +#### Debug Mode + +Enable sglang debug logging by log-level argument. + +```bash +python3 -m sglang.launch_server \ + --model-path /path/to/your/model \ + --host 0.0.0.0 \ + --device npu \ + --model-impl mindspore \ + --attention-backend ascend \ + --log-level DEBUG +``` + +Enable mindspore info and debug logging by setting environments. + +```bash +export GLOG_v=1 # INFO +export GLOG_v=0 # DEBUG +``` + +#### Explicitly select devices + +Use the following environment variable to explicitly select the devices to use. + +```shell +export ASCEND_RT_VISIBLE_DEVICES=4,5,6,7 # to set device +``` + +#### Some communication environment issues + +In case of some environment with special communication environment, users need set some environment variables. + +```shell +export MS_ENABLE_LCCL=off # current not support LCCL communication mode in SGLang-MindSpore +``` + +#### Some dependencies of protobuf + +In case of some environment with special protobuf version, users need set some environment variables to avoid binary version mismatch. + +```shell +export PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python # to avoid protobuf binary version mismatch +``` + +## Support +For MindSpore-specific issues: + +- Refer to the [MindSpore documentation](https://www.mindspore.cn/) diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index 77c2d81cd..ba0317bfe 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/python/sglang/srt/configs/model_config.py @@ -47,6 +47,7 @@ class ModelImpl(str, Enum): AUTO = "auto" SGLANG = "sglang" TRANSFORMERS = "transformers" + MINDSPORE = "mindspore" def is_deepseek_nsa(config: PretrainedConfig) -> bool: diff --git a/python/sglang/srt/model_executor/mindspore_runner.py b/python/sglang/srt/model_executor/mindspore_runner.py new file mode 100644 index 000000000..832326b1c --- /dev/null +++ b/python/sglang/srt/model_executor/mindspore_runner.py @@ -0,0 +1,118 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the SGLang project +"""ms_runner launch MindSpore distributed modules.""" + +import multiprocessing as mp +import os +import sys +from pathlib import Path + +import mindspore as ms +import torch +from mindspore._c_expression import GroupOptions +from mindspore.communication import create_group + +from sglang.srt.distributed.parallel_state import _groups + + +class _Tmp: + def __init__(self): + self.sched_p = None + + def set_sched_process(self, p): + self.sched_p = p + + def __del__(self): + if self.sched_p: + self.sched_p.kill() + + +_tmp = _Tmp() + + +def _get_host_and_ip(distributed_init_method): + try: + _, ip_str, port_str = distributed_init_method.split(":") + ip = ip_str.split("/")[-1] + port = int(port_str) + except Exception as e: + raise RuntimeError( + "Cannot get host and port information from %s, error: %s!" + % (distributed_init_method, str(e)) + ) + + return ip, port + + +def run_scheduler_init(rank, local_rank, world_size, master_addr, master_port): + with open(str(Path() / "schedule.log"), "w") as scheduler_f: + # For Python outputs. + sys.stdout = scheduler_f + sys.stderr = scheduler_f + # For C++ outputs. + os.dup2(scheduler_f.fileno(), 1) + os.dup2(scheduler_f.fileno(), 2) + os.environ["DEVICE_ID"] = str(local_rank) + os.environ["MS_WORKER_NUM"] = str(world_size) + os.environ["MS_ROLE"] = "MS_SCHED" + os.environ["MS_NODE_ID"] = str(rank) + os.environ["MS_SCHED_HOST"] = str(master_addr) + os.environ["MS_SCHED_PORT"] = str(master_port) + # This function is blocked until the whole cluster exits. + ms.communication.init() + + +def set_ms_parallel_env(rank, local_rank, world_size, init_method): + master_addr, master_port = _get_host_and_ip(init_method) + # change port avoiding port conflicts with torch + master_port = master_port + 35 if master_port < 65500 else master_port - 35 + if not os.getenv("MS_ROLE"): + if rank == 0: + # Create a subprocess for scheduler of MindSpore, just for internal collaboration, not for collective communication + sched_p = mp.Process( + target=run_scheduler_init, + args=(rank, local_rank, world_size, master_addr, master_port), + ) + sched_p.start() + global _tmp + _tmp.set_sched_process(sched_p) + + os.environ["DEVICE_ID"] = str(local_rank) + os.environ["MS_WORKER_NUM"] = str(world_size) + os.environ["MS_ROLE"] = "MS_WORKER" + os.environ["MS_NODE_ID"] = str(rank) + os.environ["MS_SCHED_HOST"] = str(master_addr) + os.environ["MS_SCHED_PORT"] = str(master_port) + + +def reuse_hccl_comm(): + for group_name, group in _groups.items(): + # Torch ProcessGroupHccl + device_group = group().device_group + hccl_comm_handle = device_group._get_backend(torch.device("npu")).get_hccl_comm( + group().local_rank + ) + print( + f"MindSpore reuse torch group: {device_group}, group_name: {group_name}, local rank: {group().local_rank}," + f"hccl communicator handle: {hex(hccl_comm_handle)}", + flush=True, + ) + # Create MS communication group by hccl comm handle to reuse Torch group. + group_options = GroupOptions() + group_options.hccl_config = {"hccl_comm": hccl_comm_handle} + create_group(group_name, group().ranks, group_options) + + +def init_ms_distributed(world_size, rank, local_rank, server_args, port): + if server_args.dist_init_addr: + dist_init_method = f"tcp://{server_args.dist_init_addr}" + else: + dist_init_method = f"tcp://{server_args.host}:{port}" + set_ms_parallel_env(rank, local_rank, world_size, dist_init_method) + + ms.set_context(infer_boost="on", jit_level="O0") + ms.set_context(mode=ms.context.PYNATIVE_MODE) + ms.set_device("Ascend", local_rank) + ms.communication.init("hccl") + # After distributed job is initialized, reuse hccl comms for MindSpore. + reuse_hccl_comm() diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 950d3acfb..1dfc36e92 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -42,6 +42,7 @@ from sglang.srt.configs.load_config import LoadConfig, LoadFormat from sglang.srt.configs.model_config import ( AttentionArch, ModelConfig, + ModelImpl, get_nsa_index_head_dim, is_deepseek_nsa, ) @@ -317,6 +318,8 @@ class ModelRunner: if get_bool_env_var("SGLANG_DETECT_SLOW_RANK"): slow_rank_detector.execute() + # Init mindspore running environment when model impl is "mindspore" + self.init_mindspore_runner() # Update deep gemm configure if deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM: @@ -364,6 +367,20 @@ class ModelRunner: else: self.piecewise_cuda_graph_runner = None + def init_mindspore_runner(self): + # Init the mindspore runner + # for now, there is only some communication initialization work + if self.server_args.model_impl.lower() == ModelImpl.MINDSPORE and _is_npu: + from sglang.srt.model_executor.mindspore_runner import init_ms_distributed + + init_ms_distributed( + world_size=self.tp_size * self.pp_size, + rank=self.tp_size * self.pp_rank + self.tp_rank, + local_rank=self.gpu_id, + server_args=self.server_args, + port=self.dist_port, + ) + def initialize(self, min_per_gpu_memory: float): server_args = self.server_args @@ -2018,6 +2035,9 @@ class ModelRunner: # TODO: Currently, cuda graph only captures decode steps, which only exists for generation models return + if self.server_args.model_impl.lower() == ModelImpl.MINDSPORE: + return + if self.device != "cpu" and self.server_args.disable_cuda_graph: return diff --git a/python/sglang/srt/model_loader/utils.py b/python/sglang/srt/model_loader/utils.py index b0a45148d..29763bc10 100644 --- a/python/sglang/srt/model_loader/utils.py +++ b/python/sglang/srt/model_loader/utils.py @@ -99,7 +99,9 @@ def get_model_architecture(model_config: ModelConfig) -> Tuple[Type[nn.Module], supported_archs = ModelRegistry.get_supported_archs() is_native_supported = any(arch in supported_archs for arch in architectures) - if not is_native_supported or model_config.model_impl == ModelImpl.TRANSFORMERS: + if model_config.model_impl == ModelImpl.MINDSPORE: + architectures = ["MindSporeForCausalLM"] + elif not is_native_supported or model_config.model_impl == ModelImpl.TRANSFORMERS: architectures = resolve_transformers_arch(model_config, architectures) return ModelRegistry.resolve_model_cls(architectures) diff --git a/python/sglang/srt/models/mindspore.py b/python/sglang/srt/models/mindspore.py new file mode 100644 index 000000000..c4e8d1bc5 --- /dev/null +++ b/python/sglang/srt/models/mindspore.py @@ -0,0 +1,305 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the SGLang project +import logging +from typing import Any, Iterable, Optional, Tuple + +import torch + +from sglang.srt.distributed import ( + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, +) +from sglang.srt.layers.logits_processor import LogitsProcessorOutput +from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.model_executor.forward_batch_info import ForwardBatch +from sglang.srt.models.registry import import_model_classes +from sglang.srt.utils import is_npu + +_is_npu = is_npu() + +if _is_npu: + import mindspore as ms + import numpy as np + import torch_npu + from mindspore import Tensor, mint, mutable + +logger = logging.getLogger(__name__) + + +def tensor_torch2ms(x: torch.Tensor): + if x is None or not isinstance(x, torch.Tensor): + return x + + # torch tensor -> dlpack -> mindspore tensor + pt_dlpack = torch.utils.dlpack.to_dlpack(x) + ms_tensor = ms.utils.dlpack.from_dlpack(pt_dlpack) + return ms_tensor + + +def tensor_ms2torch(x: ms.Tensor): + if x is None or not isinstance(x, ms.Tensor): + return x + + # ms tensor -> dlpack -> torch tensor + ms_dlpack = ms.utils.dlpack.to_dlpack(x) + torch_tensor = torch.utils.dlpack.from_dlpack(ms_dlpack) + torch_npu.npu.synchronize() + return torch_tensor + + +# Adapt from: https://gitee.com/mindspore/vllm-mindspore/blob/master/vllm_mindspore/model_executor/models/attention_mask.py +class LowerTriangularMask: + r""" + Provide Infer model attention mask. + Args: + dtype (ms dtype): The compute type of Infer model. + max_model_len (int): The max model length of Infer model. + """ + + def __init__(self, dtype, max_model_len, decode_mask_coeff=-10000.0): + self.dtype = dtype + self.max_model_len = max_model_len + self.cached_mask_len = 8 * 1024 + self.decode_mask_coeff = decode_mask_coeff + + prefill_mask_coeff = 1.0 if self.dtype == ms.bfloat16 else -10000.0 + self.prefill_mask = Tensor( + np.triu(np.ones(shape=(128, 128), dtype=np.float16), k=1) + * prefill_mask_coeff, + dtype=self.dtype, + ) + + self.hard_mask = mint.zeros((1, 1), dtype=dtype) + self.decode_mask = ( + Tensor( + np.triu( + np.ones( + shape=(self.cached_mask_len, self.cached_mask_len), + dtype=np.int8, + ), + k=1, + ), + dtype=self.dtype, + ) + * self.decode_mask_coeff + ) + + def create_mask(self, query_lens_np, seq_lens_np): + """ + when query_lens_np = [3], seq_lens_np = [6], decode_mask_coeff = 1 + init attention mask + 0 0 0 0 0 0 + 0 0 0 0 0 0 + 0 0 0 0 0 0 + """ + max_seq_len = seq_lens_np.max().item() + total_q_len = query_lens_np.sum().item() + attention_mask = mint.zeros((total_q_len, max_seq_len), dtype=self.dtype) + + req_num = query_lens_np.shape[0] + current_row = 0 + for i in range(req_num): + q_len = query_lens_np[i].item() + current_row += q_len + # skip row when q_len <= 1, to decrease execute time + if q_len <= 1: + continue + seq_len = seq_lens_np[i].item() + context_len = seq_len - q_len + """ + set the right half to 1 + 0 0 0 1 1 1 + 0 0 0 1 1 1 + 0 0 0 1 1 1 + """ + attention_mask[current_row - q_len : current_row, context_len:] = ( + self.decode_mask_coeff + ) + """ + set the lower triangle of the right half to 0 + 0 0 0 0 1 1 + 0 0 0 0 0 1 + 0 0 0 0 0 0 + """ + right_tensor = attention_mask[ + current_row - q_len : current_row, context_len:seq_len + ] + + # use masked_fill_ to inplace modify attention_mask + right_tensor.masked_fill_(right_tensor.tril() == self.decode_mask_coeff, 0) + + return attention_mask + + def gen_attention_mask( + self, + is_prefill: bool, + position_ids: Tensor, + query_lens_np: np.ndarray, + seq_lens_np: np.ndarray, + ): + max_query_len = query_lens_np.max() + max_seq_len = seq_lens_np.max() + if is_prefill: + attention_mask = self.prefill_mask + elif max_query_len > 1: + if max_seq_len <= self.cached_mask_len: + attention_mask = mint.index_select(self.decode_mask, 0, position_ids) + else: + attention_mask = self.create_mask(query_lens_np, seq_lens_np) + else: + attention_mask = self.hard_mask + return attention_mask + + +class MindSporeForCausalLM(torch.nn.Module): + def __init__( + self, + config: Any, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.config = config + + ms.set_context(graph_kernel_flags="--disable_pass=gather_pre_rms_norm_fusion") + ms.set_kernel_launch_capture(False) + + logger.info( + "MindSporeForCausalLM tp size %d tp rank %d", + get_tensor_model_parallel_world_size(), + get_tensor_model_parallel_rank(), + ) + if get_tensor_model_parallel_world_size() not in (1, 2, 4, 8): + # MatMulAllReduce only support tp size in (1, 2, 4, 8) + ms.set_context(graph_kernel_flags="--disable_pass=MatMulAllReduce") + + arch = self.get_arch(self.config) + self.model = arch(config=config, quant_config=quant_config) + + self.casual_mask = LowerTriangularMask( + self.config.param_dtype, self.config.max_position_embeddings + ) + self.key_cache = [] + self.value_cache = [] + + def get_arch(self, config): + # Get all implemented models + mindspore_models = import_model_classes("sgl_mindspore.models") + + # Get arch from config + architectures = config.architectures + if isinstance(architectures, str): + architectures = [architectures] + if not architectures: + logger.warning("No model architectures are specified") + + for arch in architectures: + if arch in mindspore_models: + return mindspore_models[arch] + if arch is None: + raise ValueError(f"Unsupported arch {architectures}") + + @property + def use_mla(self): + return self.config.architectures[0] in ("DeepseekV3ForCausalLM") + + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): + self.model.load_weights(weights) + for _, cell in self.model.cells_and_names(): + quant_method = getattr(cell, "quant_method", None) + if quant_method is not None: + quant_method.process_weights_after_loading(cell) + + def get_kvcache(self, forward_batch: ForwardBatch): + def prepare_cache(cache_list, is_key_cache): + for i in range(self.config.num_hidden_layers): + if is_key_cache: + cache = forward_batch.token_to_kv_pool.get_key_buffer(i) + else: + cache = forward_batch.token_to_kv_pool.get_value_buffer(i) + cache_ms = tensor_torch2ms(cache) + if cache_ms.ndim == 3: + cache_ms = mint.unsqueeze(cache_ms, 2) + cache_list.append(cache_ms) + + if self.use_mla: + if not self.key_cache: + prepare_cache(self.key_cache, is_key_cache=True) + return mutable(self.key_cache) + + if self.key_cache and self.value_cache: + return mutable(self.key_cache), mutable(self.value_cache) + + prepare_cache(self.key_cache, is_key_cache=True) + prepare_cache(self.value_cache, is_key_cache=False) + + return mutable(self.key_cache), mutable(self.value_cache) + + def prepare_inputs(self, input_ids, positions, forward_batch): + if self.use_mla: + key_cache = self.get_kvcache(forward_batch) + else: + key_cache, value_cache = self.get_kvcache(forward_batch) + + # Different processing for the mindspore attention operator + # Without any prefix cache => Use FlashAttentionScore + # With cache => Use PagedAttention, no matter the query length is 1 or not + is_prefill = forward_batch.forward_mode.is_extend() + is_prefill = is_prefill and forward_batch.extend_prefix_lens.sum().item() == 0 + + batch_valid_length = forward_batch.seq_lens.cpu().numpy() + + if forward_batch.extend_seq_lens is not None: + q_seq_lens = forward_batch.extend_seq_lens.cpu().numpy() + else: + q_seq_lens = np.ones([forward_batch.batch_size], dtype=np.int32) + + page_size = forward_batch.token_to_kv_pool.page_size + block_tables = tensor_torch2ms( + ( + forward_batch.req_to_token_pool.req_to_token[ + forward_batch.req_pool_indices, : forward_batch.seq_lens.max() + ][:, ::page_size] + // page_size + ) + ).to(ms.int32) + + model_inputs = {} + model_inputs["input_ids"] = tensor_torch2ms(input_ids).to(ms.int32) + model_inputs["batch_valid_length"] = ms.Tensor( + batch_valid_length, dtype=ms.int32 + ) + model_inputs["position_ids"] = tensor_torch2ms(positions) + model_inputs["q_seq_lens"] = ms.Tensor(q_seq_lens, dtype=ms.int32) + model_inputs["attention_mask"] = self.casual_mask.gen_attention_mask( + is_prefill, model_inputs["position_ids"], q_seq_lens, batch_valid_length + ).contiguous() + model_inputs["out_cache_loc"] = tensor_torch2ms(forward_batch.out_cache_loc).to( + ms.int32 + ) + model_inputs["is_prefill"] = is_prefill + model_inputs["key_cache"] = key_cache + if not self.use_mla: + model_inputs["value_cache"] = value_cache + model_inputs["block_tables"] = block_tables + return model_inputs + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + ) -> Tensor: + # prepare base inputs + model_inputs = self.prepare_inputs(input_ids, positions, forward_batch) + # prepare model inputs + model_inputs = self.model.prepare_inputs(forward_batch, model_inputs) + + logits = self.model(**model_inputs) + + # TODO: npu tensor ms2torch error to be fix, remain issues of torch_npu to get tensor from dlpack + logits_result = LogitsProcessorOutput(next_token_logits=tensor_ms2torch(logits)) + return logits_result + + +EntryClass = [MindSporeForCausalLM] diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 90f0d29b9..04cc868ab 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -2059,6 +2059,7 @@ class ServerArgs: "implementation is available.\n" '* "sglang" will use the SGLang model implementation.\n' '* "transformers" will use the Transformers model ' + '* "mindspore" will use the MindSpore model ' "implementation.\n", ) @@ -3957,6 +3958,9 @@ class ServerArgs: self.schedule_conservativeness >= 0 ), "schedule_conservativeness must be non-negative" + if self.model_impl == "mindspore": + assert is_npu(), "MindSpore model impl is only supported on Ascend npu." + def check_lora_server_args(self): assert self.max_loras_per_batch > 0, "max_loras_per_batch must be positive" diff --git a/test/srt/ascend/test_mindspore_models.py b/test/srt/ascend/test_mindspore_models.py new file mode 100644 index 000000000..8c515bc45 --- /dev/null +++ b/test/srt/ascend/test_mindspore_models.py @@ -0,0 +1,64 @@ +""" +Usage: +python3 -m unittest test_mindspore_models.TestMindSporeQwen3.test_gsm8k +""" + +import unittest +from types import SimpleNamespace + +from sglang.srt.utils import kill_process_tree +from sglang.test.few_shot_gsm8k import run_eval +from sglang.test.test_utils import ( + DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + DEFAULT_URL_FOR_TEST, + CustomTestCase, + popen_launch_server, +) + + +class TestMindSporeQwen3(CustomTestCase): + @classmethod + def setUpClass(cls): + cls.model = "Qwen/Qwen3-8B" + cls.base_url = DEFAULT_URL_FOR_TEST + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=[ + "--device", + "npu", + "--model-impl", + "mindspore", + "--attention-backend", + "ascend", + "--tp-size", + "1", + "--dp-size", + "1", + "--mem-fraction-static", + 0.8, + ], + ) + + @classmethod + def tearDownClass(cls): + kill_process_tree(cls.process.pid) + + def test_gsm8k(self): + args = SimpleNamespace( + num_shots=5, + data_path=None, + num_questions=200, + max_new_tokens=512, + parallel=128, + host="http://127.0.0.1", + port=int(self.base_url.split(":")[-1]), + ) + metrics = run_eval(args) + print(f"{metrics=}") + self.assertGreater(metrics["accuracy"], 0.78) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/srt/run_suite.py b/test/srt/run_suite.py index 0f2842891..1fe6ba97a 100644 --- a/test/srt/run_suite.py +++ b/test/srt/run_suite.py @@ -536,6 +536,9 @@ suite_ascend = { TestFile("ascend/test_ascend_deepep.py", 400), TestFile("ascend/test_ascend_deepseek_mtp.py", 400), ], + "__not_in_ascend_ci__": [ + TestFile("ascend/test_mindspore_models.py"), + ], } suites.update(suite_amd)