[Core] Add Ascend Quant Config to main branch

tests/quantization/test_mindie_turbo.py

@@ -0,0 +1,65 @@
+#
+# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
+# This file is a part of the vllm-ascend project.
+# Adapted from vllm/tests/basic_correctness/test_basic_correctness.py
+# Copyright 2023 The vLLM team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+"""Tests whether ascend quantization based on MindIE-Turbo is enabled correctly.
+
+Run `pytest tests/quantization/test_mindie_turbo.py`.
+"""
+
+import os
+
+import pytest
+
+import vllm  # noqa: F401
+
+import vllm_ascend  # noqa: F401
+
+from tests.conftest import VllmRunner
+from tests.quantization.utils import is_mindie_turbo_supported, example_quantization
+
+MODELS = [
+    "Qwen/Qwen2.5-0.5B-Instruct",
+]
+
+
+@pytest.mark.skipif(not is_mindie_turbo_supported(),
+                    reason="MindIE-Turbo is not installed.")
+@pytest.mark.parametrize("model_name_or_path", MODELS)
+@pytest.mark.parametrize("max_tokens", [5])
+def test_mindie_turbo(
+    model_name_or_path: str,
+    max_tokens: int,
+) -> None:
+    # vLLM must load weights from disk. Hence we need to save the quantized
+    # weights at first, and then load it by vLLM.
+    temp_path = os.path.join(os.path.dirname(__file__), "temp_weight")
+    if not os.path.exists(temp_path):
+        os.makedirs(temp_path)
+    example_quantization(model_name_or_path, temp_path)
+
+    prompt = "What's deep learning?"
+    example_prompts = [prompt]
+
+    with VllmRunner(temp_path,
+                    max_model_len=8192,
+                    dtype="bfloat16",
+                    enforce_eager=False,
+                    gpu_memory_utilization=0.7) as vllm_model:
+
+        output = vllm_model.generate_greedy(example_prompts, max_tokens)
+        assert output

tests/quantization/utils.py

@@ -0,0 +1,88 @@
+#
+# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
+# This file is a part of the vllm-ascend project.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from msmodelslim.pytorch.llm_ptq.anti_outlier import AntiOutlierConfig, AntiOutlier
+from msmodelslim.pytorch.llm_ptq.llm_ptq_tools import Calibrator, QuantConfig
+
+
+def is_mindie_turbo_supported() -> bool:
+    try:
+        import mindie_turbo  # noqa: F401
+    except Exception:
+        return False
+
+    return True
+
+
+def example_quantization(model_name_or_path: str, tmp_path: str) -> None:
+
+    tokenizer = AutoTokenizer.from_pretrained(
+        pretrained_model_name_or_path=model_name_or_path
+    )
+
+    model = AutoModelForCausalLM.from_pretrained(
+        pretrained_model_name_or_path=model_name_or_path,
+        device_map="npu:0",
+        torch_dtype="auto"
+    ).eval()
+
+    data_list = ["What's deep learning?"]
+    dataset_calib = []
+    for calib_data in data_list:
+        inputs = tokenizer(calib_data, return_tensors='pt').to("npu:0")
+    dataset_calib.append([inputs.data['input_ids']])
+
+    anti_config = AntiOutlierConfig(anti_method="m2", dev_type="npu", dev_id=0)
+    anti_outlier = AntiOutlier(model, calib_data=dataset_calib, cfg=anti_config)
+    anti_outlier.process()
+
+    disable_names = ['lm_head']
+    for layer_index in range(24):
+        disable_names.append(f'model.layers.{layer_index}.mlp.down_proj')
+
+    quant_config = QuantConfig(
+        a_bit=8,
+        w_bit=8,
+        disable_names=disable_names,
+        dev_type='npu',
+        dev_id=0,
+        act_method=3,
+        pr=1.0,
+        w_sym=True,
+        mm_tensor=False
+    )
+
+    calibrator = Calibrator(model, quant_config, calib_data=dataset_calib, disable_level='L0')
+    calibrator.run()
+
+    # Currently, we need add config.json manualy for quantized weights generated by msmodelslim.
+    # Following codes will be removed once msmodelslim can generate complete weights
+    # except 'calibrator.save(tmp_path, save_type=["safe_tensor"])'.
+    class EmptyModule(torch.nn.Module):
+        def __init__(self) -> None:
+            super(EmptyModule, self).__init__()
+
+        def forward(self, x):
+            return x
+
+    calibrator.model.config.quantization_config = calibrator.quant_model_json_description.quant_model_description
+
+    calibrator.save(tmp_path, save_type=["safe_tensor"])
+    calibrator.model.save_pretrained(tmp_path, state_dict=EmptyModule().state_dict())
+    tokenizer.save_pretrained(tmp_path)

vllm_ascend/platform.py

@@ -16,7 +16,7 @@
 #
 
 import os
-from typing import Optional, Tuple
+from typing import TYPE_CHECKING, Optional, Tuple
 
 import torch
 
@@ -27,6 +27,10 @@
 
 from vllm.config import VllmConfig
 from vllm.platforms import Platform, PlatformEnum
+if TYPE_CHECKING:
+    from vllm.utils import FlexibleArgumentParser
+else:
+    FlexibleArgumentParser = None
 
 os.environ["RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES"] = "1"
 
@@ -53,6 +57,10 @@ class NPUPlatform(Platform):
     ray_device_key: str = "NPU"
     device_control_env_var: str = "ASCEND_RT_VISIBLE_DEVICES"
 
+    supported_quantization: list[str] = [
+        "ascend"
+    ]
+
     @classmethod
     def get_device_capability(cls, device_id: int = 0):
         return None
@@ -86,6 +94,21 @@ def synchronize(cls):
     def mem_get_info(cls) -> Tuple[int, int]:
         return torch.npu.mem_get_info()
 
+    @classmethod
+    def pre_register_and_update(cls,
+                                parser: Optional[FlexibleArgumentParser] = None
+                                ) -> None:
+        """
+        Do some pre-registeration or update action for the current platform.
+        This function is called before global VllmConfig is initialized or cli
+        arguments are parsed. It's used for out-of-tree platforms to register or
+        update the configuration.
+        For example, the out-of-tree quantization config can be imported and
+        registered here dynamically.
+        """
+
+        from vllm_ascend.quantization.quant_config import AscendQuantConfig  # noqa: F401
+
     @classmethod
     def check_and_update_config(cls, vllm_config: VllmConfig) -> None:
         # Register ops when setup.