Files
enginex-ascend-910-vllm/tests/ut/test_utils.py
Sun Ruoxi 7f8a1b1f7a init v0.23.0
Signed-off-by: Sun Ruoxi <sunruoxi@4paradigm.com>
2026-08-27 15:11:51 +08:00

488 lines
21 KiB
Python

#
# 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.
# This file is a part of the vllm-ascend project.
#
import math
import os
from unittest import mock
import pytest
import torch
from tests.ut.base import TestBase
from vllm_ascend import utils
from vllm_ascend.utils import REGISTERED_ASCEND_OPS
class TestUtils(TestBase):
def setUp(self):
import importlib
from vllm_ascend import platform
importlib.reload(platform)
utils.enable_dsa_cp_with_layer_shard.cache_clear()
utils.enable_dsa_cp_with_o_proj_tp.cache_clear()
def test_nd_to_nz_2d(self):
# can be divided by 16
input_tensor = torch.randn(32, 64)
output = utils.nd_to_nz_2d(input_tensor)
self.assertEqual(output.shape[0], 1)
self.assertEqual(output.shape[1], 64 // 16)
self.assertEqual(output.shape[2], 32)
self.assertEqual(output.shape[3], 16)
# cannot be divided by 16
input_tensor = torch.randn(30, 62)
output = utils.nd_to_nz_2d(input_tensor)
self.assertEqual(output.shape[0], 1)
self.assertEqual(output.shape[1], math.ceil(62 / 16))
self.assertEqual(output.shape[2], 32)
self.assertEqual(output.shape[3], 16)
# pad to 16
input_tensor = torch.randn(8, 12)
output = utils.nd_to_nz_2d(input_tensor)
self.assertEqual(output.shape[0], 1)
self.assertEqual(output.shape[1], 1) # 12->16, 16//16=1
self.assertEqual(output.shape[2], 16) # 8->16
self.assertEqual(output.shape[3], 16)
# check if the output is contiguous
input_tensor = torch.randn(32, 64)
output = utils.nd_to_nz_2d(input_tensor)
self.assertTrue(output.is_contiguous())
# check if the output values are preserved
input_tensor = torch.tensor([[1, 2, 3, 4], [5, 6, 7, 8]])
output = utils.nd_to_nz_2d(input_tensor)
expected = torch.tensor(
[
[
[
[1, 2, 3, 4, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
[5, 6, 7, 8, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
]
]
]
)
self.assertTrue(torch.allclose(output, expected))
def test_aligned_16(self):
# align to 16
input_tensor = torch.randn(15, 64)
output_tensor = utils.aligned_16(input_tensor)
self.assertEqual(output_tensor.shape[0], 16)
# align to 16
input_tensor = torch.randn(16, 64)
output_tensor = utils.aligned_16(input_tensor)
self.assertEqual(output_tensor.shape[0], 16)
self.assertTrue(torch.equal(input_tensor, output_tensor))
# align to 32
input_tensor = torch.randn(17, 64)
output_tensor = utils.aligned_16(input_tensor)
self.assertEqual(output_tensor.shape[0], 32)
@pytest.mark.skip("Skip as register_kernels has NPU SocName checking in CANN 8.5.0.")
def test_enable_custom_op(self):
result = utils.enable_custom_op()
self.assertTrue(result)
utils._CUSTOM_OP_ENABLED = None
with mock.patch("builtins.__import__") as mock_import_module:
mock_import_module.side_effect = ImportError("import error")
self.assertFalse(utils.enable_custom_op())
def test_find_hccl_library(self):
with mock.patch.dict(os.environ, {"HCCL_SO_PATH": "/path/to/hccl/libhccl.so"}):
self.assertEqual(utils.find_hccl_library(), "/path/to/hccl/libhccl.so")
with mock.patch("torch.version.cann", None):
self.assertRaises(ValueError, utils.find_hccl_library)
with mock.patch("torch.version.cann", "Ascend910"):
self.assertEqual(utils.find_hccl_library(), "libhccl.so")
def test_current_stream(self):
with mock.patch("torch.npu.current_stream") as mock_current_stream:
self.assertEqual(utils.current_stream(), mock_current_stream())
def test_enable_dsa_cp_with_layer_shard_accepts_kv_producer(self):
mock_vllm_config = mock.MagicMock()
mock_vllm_config.kv_transfer_config = mock.MagicMock(
kv_role="kv_producer", is_kv_producer=True, is_kv_consumer=False
)
with (
mock.patch("vllm.config.get_current_vllm_config", return_value=mock_vllm_config),
mock.patch("vllm_ascend.utils.enable_dsa_cp", return_value=True),
):
self.assertTrue(utils.enable_dsa_cp_with_layer_shard())
def test_enable_dsa_cp_with_layer_shard_rejects_kv_both(self):
mock_vllm_config = mock.MagicMock()
mock_vllm_config.kv_transfer_config = mock.MagicMock(
kv_role="kv_both", is_kv_producer=True, is_kv_consumer=True
)
with (
mock.patch("vllm.config.get_current_vllm_config", return_value=mock_vllm_config),
mock.patch("vllm_ascend.utils.enable_dsa_cp", return_value=True),
):
self.assertFalse(utils.enable_dsa_cp_with_layer_shard())
def test_enable_dsa_cp_with_layer_shard_rejects_missing_kv_transfer(self):
mock_vllm_config = mock.MagicMock()
mock_vllm_config.kv_transfer_config = None
with (
mock.patch("vllm.config.get_current_vllm_config", return_value=mock_vllm_config),
mock.patch("vllm_ascend.utils.enable_dsa_cp", return_value=True),
):
self.assertFalse(utils.enable_dsa_cp_with_layer_shard())
def test_enable_dsa_cp_with_layer_shard_rejects_when_dsa_cp_disabled(self):
with mock.patch("vllm_ascend.utils.enable_dsa_cp", return_value=False):
self.assertFalse(utils.enable_dsa_cp_with_layer_shard())
def test_enable_dsa_cp_with_o_proj_tp_accepts_missing_kv_transfer(self):
mock_vllm_config = mock.MagicMock()
mock_vllm_config.kv_transfer_config = None
with (
mock.patch("vllm.config.get_current_vllm_config", return_value=mock_vllm_config),
mock.patch("vllm_ascend.utils.enable_dsa_cp", return_value=True),
):
self.assertTrue(utils.enable_dsa_cp_with_o_proj_tp())
def test_enable_dsa_cp_with_o_proj_tp_accepts_kv_both(self):
mock_vllm_config = mock.MagicMock()
mock_vllm_config.kv_transfer_config = mock.MagicMock(
kv_role="kv_both", is_kv_producer=True, is_kv_consumer=True
)
with (
mock.patch("vllm.config.get_current_vllm_config", return_value=mock_vllm_config),
mock.patch("vllm_ascend.utils.enable_dsa_cp", return_value=True),
):
self.assertTrue(utils.enable_dsa_cp_with_o_proj_tp())
def test_enable_dsa_cp_with_o_proj_tp_rejects_single_role_pd(self):
mock_vllm_config = mock.MagicMock()
mock_vllm_config.kv_transfer_config = mock.MagicMock(
kv_role="kv_producer", is_kv_producer=True, is_kv_consumer=False
)
with (
mock.patch("vllm.config.get_current_vllm_config", return_value=mock_vllm_config),
mock.patch("vllm_ascend.utils.enable_dsa_cp", return_value=True),
):
self.assertFalse(utils.enable_dsa_cp_with_o_proj_tp())
def test_enable_dsa_cp_with_o_proj_tp_rejects_when_dsa_cp_disabled(self):
with mock.patch("vllm_ascend.utils.enable_dsa_cp", return_value=False):
self.assertFalse(utils.enable_dsa_cp_with_o_proj_tp())
def test_vllm_version_is(self):
with mock.patch.dict(os.environ, {"VLLM_VERSION": "1.0.0"}):
with mock.patch("vllm.__version__", "1.0.0"):
self.assertTrue(utils.vllm_version_is.__wrapped__("1.0.0"))
self.assertFalse(utils.vllm_version_is.__wrapped__("2.0.0"))
with mock.patch("vllm.__version__", "2.0.0"):
self.assertTrue(utils.vllm_version_is.__wrapped__("1.0.0"))
self.assertFalse(utils.vllm_version_is.__wrapped__("2.0.0"))
with mock.patch("vllm.__version__", "1.0.0"):
self.assertTrue(utils.vllm_version_is.__wrapped__("1.0.0"))
self.assertFalse(utils.vllm_version_is.__wrapped__("2.0.0"))
with mock.patch("vllm.__version__", "2.0.0"):
self.assertTrue(utils.vllm_version_is.__wrapped__("2.0.0"))
self.assertFalse(utils.vllm_version_is.__wrapped__("1.0.0"))
# Test caching takes effect
utils.vllm_version_is.cache_clear()
utils.vllm_version_is("1.0.0")
misses = utils.vllm_version_is.cache_info().misses
hits = utils.vllm_version_is.cache_info().hits
self.assertEqual(misses, 1)
self.assertEqual(hits, 0)
utils.vllm_version_is("1.0.0")
hits = utils.vllm_version_is.cache_info().hits
self.assertEqual(hits, 1)
def test_get_max_hidden_layers(self):
from transformers import PretrainedConfig
class SimpleConfig(PretrainedConfig):
def __init__(self, num_hidden_layers=12):
self.num_hidden_layers = num_hidden_layers
def to_dict(self):
return {"num_hidden_layers": self.num_hidden_layers}
self.assertEqual(utils.get_max_hidden_layers(SimpleConfig()), 12)
self.assertEqual(utils.get_max_hidden_layers(SimpleConfig(24)), 24)
class NestedConfig(PretrainedConfig):
def to_dict(self):
return {
"model": {"encoder": {"num_hidden_layers": 8}, "decoder": {"num_hidden_layers": 12}},
"other_setting": True,
}
self.assertEqual(utils.get_max_hidden_layers(NestedConfig()), 12)
class MultiValueConfig(PretrainedConfig):
def to_dict(self):
return {
"num_hidden_layers": 6,
"submodule": {"num_hidden_layers": 18, "subsub": {"num_hidden_layers": 9}},
}
self.assertEqual(utils.get_max_hidden_layers(MultiValueConfig()), 18)
class NoLayerConfig(PretrainedConfig):
def to_dict(self):
return {"attention_heads": 8}
with self.assertRaises(ValueError) as context:
utils.get_max_hidden_layers(NoLayerConfig())
self.assertIn("num_hidden_layers", str(context.exception))
def test_is_drafter_moe_model_extract_hidden_states_is_never_moe(self):
"""The extract_hidden_states drafter is a cache-only attention layer
with no MoE layers, but its hf_config copies the (possibly MoE) target
hf_config. The expert-key scan must not misclassify it as MoE,
otherwise _sync_metadata_across_dp(is_draft_model=True) performs a DP
all_reduce that idle DP ranks never match (DP deadlock)."""
vllm_config = mock.MagicMock()
vllm_config.speculative_config.method = "extract_hidden_states"
# Inherited MoE keys from the target model (e.g. MiniMax-M2)
vllm_config.speculative_config.draft_model_config.hf_text_config.to_dict.return_value = {
"num_local_experts": 256,
"num_experts_per_tok": 8,
}
with mock.patch("vllm_ascend.utils._IS_DRAFTER_MOE_MODEL", None):
self.assertFalse(utils.is_drafter_moe_model(vllm_config))
def test_is_drafter_moe_model_eagle_moe_drafter_detected(self):
"""Non-extract_hidden_states drafters keep the expert-key detection."""
vllm_config = mock.MagicMock()
vllm_config.speculative_config.method = "eagle3"
vllm_config.speculative_config.draft_model_config.hf_text_config.to_dict.return_value = {
"num_experts_per_tok": 8,
}
with mock.patch("vllm_ascend.utils._IS_DRAFTER_MOE_MODEL", None):
self.assertTrue(utils.is_drafter_moe_model(vllm_config))
@mock.patch("vllm.model_executor.custom_op.CustomOp")
@mock.patch("vllm_ascend.ops.activation.AscendQuickGELU")
@mock.patch("vllm_ascend.ops.activation.AscendSiluAndMul")
@mock.patch("vllm_ascend.ops.layernorm.AscendRMSNorm")
def test_register_ascend_customop(
self, mock_ascend_rmsnorm, mock_ascend_silu_and_mul, mock_ascend_quick_gelu, mock_customop
):
utils._ASCEND_CUSTOMOP_IS_REIGISTERED = False
# ascend custom op is not registered
utils.register_ascend_customop()
self.assertEqual(mock_customop.register_oot.call_count, len(REGISTERED_ASCEND_OPS))
self.assertTrue(utils._ASCEND_CUSTOMOP_IS_REIGISTERED)
# ascend custom op is already registered
utils.register_ascend_customop()
self.assertEqual(mock_customop.register_oot.call_count, len(REGISTERED_ASCEND_OPS))
@mock.patch("torch_npu.npu_format_cast")
def test_maybe_trans_nz(self, mock_npu_format_cast):
from vllm_ascend.utils import ACL_FORMAT_FRACTAL_NZ
mock_npu_format_cast.side_effect = lambda weight, fmt: weight
def assert_nz_cast(weight):
mock_npu_format_cast.assert_called_once()
args, kwargs = mock_npu_format_cast.call_args
self.assertIs(args[0], weight)
self.assertEqual(args[1], ACL_FORMAT_FRACTAL_NZ)
self.assertEqual(kwargs, {})
# Test case 1: non-310P, NZ is disabled
mock_config = mock.MagicMock()
mock_config.weight_nz_mode = 0
with (
mock.patch("vllm_ascend.utils.get_ascend_config", return_value=mock_config),
mock.patch("vllm_ascend.utils.is_310p", return_value=False),
):
weight = torch.randn(32, 64, dtype=torch.float16)
result = utils.maybe_trans_nz(weight)
self.assertIs(result, weight)
mock_npu_format_cast.assert_not_called()
# Test case 2: 310P always converts non-fp32 weights, even when NZ=0
mock_npu_format_cast.reset_mock()
mock_config.weight_nz_mode = 0
with (
mock.patch("vllm_ascend.utils.get_ascend_config", return_value=mock_config),
mock.patch("vllm_ascend.utils.is_310p", return_value=True),
):
weight = torch.randn(32, 64, dtype=torch.float16)
result = utils.maybe_trans_nz(weight)
self.assertIs(result, weight)
assert_nz_cast(weight)
# Test case 3: fp32 never converts, including on 310P
mock_npu_format_cast.reset_mock()
mock_config.weight_nz_mode = 1
with (
mock.patch("vllm_ascend.utils.get_ascend_config", return_value=mock_config),
mock.patch("vllm_ascend.utils.is_310p", return_value=True),
):
weight = torch.randn(32, 64, dtype=torch.float32)
result = utils.maybe_trans_nz(weight)
self.assertIs(result, weight)
mock_npu_format_cast.assert_not_called()
# Test case 4: non-310P fp16 converts only when NZ=2
mock_npu_format_cast.reset_mock()
mock_config.weight_nz_mode = 1
with (
mock.patch("vllm_ascend.utils.get_ascend_config", return_value=mock_config),
mock.patch("vllm_ascend.utils.is_310p", return_value=False),
):
weight = torch.randn(32, 64, dtype=torch.float16)
result = utils.maybe_trans_nz(weight)
self.assertIs(result, weight)
mock_npu_format_cast.assert_not_called()
# Test case 5: non-310P fp16 converts when NZ=2
mock_npu_format_cast.reset_mock()
mock_config.weight_nz_mode = 2
with (
mock.patch("vllm_ascend.utils.get_ascend_config", return_value=mock_config),
mock.patch("vllm_ascend.utils.is_310p", return_value=False),
):
weight = torch.randn(32, 64, dtype=torch.float16)
result = utils.maybe_trans_nz(weight)
self.assertIs(result, weight)
assert_nz_cast(weight)
# Test case 6: non-310P bf16 converts when NZ=2
mock_npu_format_cast.reset_mock()
mock_config.weight_nz_mode = 2
with (
mock.patch("vllm_ascend.utils.get_ascend_config", return_value=mock_config),
mock.patch("vllm_ascend.utils.is_310p", return_value=False),
):
weight = torch.randn(32, 64, dtype=torch.bfloat16)
result = utils.maybe_trans_nz(weight)
self.assertIs(result, weight)
assert_nz_cast(weight)
# Test case 7: non-310P quantized weights still convert by default
mock_npu_format_cast.reset_mock()
mock_config.weight_nz_mode = 1
with (
mock.patch("vllm_ascend.utils.get_ascend_config", return_value=mock_config),
mock.patch("vllm_ascend.utils.is_310p", return_value=False),
):
weight = torch.zeros(32, 64, dtype=torch.int8)
result = utils.maybe_trans_nz(weight)
self.assertIs(result, weight)
assert_nz_cast(weight)
def test_is_pd_decode_recompute_scheduler_enabled_without_config():
assert utils.is_pd_decode_recompute_scheduler_enabled() is False
def test_is_pd_decode_recompute_scheduler_enabled_kv_producer():
vllm_config = mock.MagicMock()
vllm_config.kv_transfer_config = mock.MagicMock()
vllm_config.kv_transfer_config.is_kv_consumer = False
vllm_config.kv_transfer_config.is_kv_producer = True
assert utils.is_pd_decode_recompute_scheduler_enabled(vllm_config) is False
def test_is_pd_decode_recompute_scheduler_enabled_decode_consumer():
vllm_config = mock.MagicMock()
vllm_config.kv_transfer_config = mock.MagicMock()
vllm_config.kv_transfer_config.is_kv_consumer = True
vllm_config.kv_transfer_config.is_kv_producer = False
ascend_config = mock.MagicMock()
ascend_config.recompute_scheduler_enable = True
with mock.patch("vllm_ascend.utils.get_ascend_config", return_value=ascend_config):
assert utils.is_pd_decode_recompute_scheduler_enabled(vllm_config) is True
def test_is_rc_device_returns_false_on_non_310p():
utils._IS_RC_DEVICE = None
with mock.patch("vllm_ascend.utils.is_310p", return_value=False):
assert utils.is_rc_device() is False
def test_is_rc_device_detects_ep_from_lspci():
utils._IS_RC_DEVICE = None
with (
mock.patch("vllm_ascend.utils.is_310p", return_value=True),
mock.patch("subprocess.run") as mock_run,
):
mock_run.return_value.stdout = "00:00.0 accelerators: Huawei Technologies Co., Ltd."
assert utils.is_rc_device() is False
def test_is_rc_device_detects_rc_from_lspci():
utils._IS_RC_DEVICE = None
with (
mock.patch("vllm_ascend.utils.is_310p", return_value=True),
mock.patch("subprocess.run") as mock_run,
):
mock_run.return_value.stdout = "00:00.0 PCI bridge: Huawei Technologies Co., Ltd."
assert utils.is_rc_device() is True
def test_is_rc_device_defaults_to_ep_when_lspci_unavailable():
utils._IS_RC_DEVICE = None
with (
mock.patch("vllm_ascend.utils.is_310p", return_value=True),
mock.patch("subprocess.run", side_effect=FileNotFoundError),
):
assert utils.is_rc_device() is False
def test_is_pd_decode_recompute_scheduler_enabled_decode_consumer_disabled():
vllm_config = mock.MagicMock()
vllm_config.kv_transfer_config = mock.MagicMock()
vllm_config.kv_transfer_config.is_kv_consumer = True
vllm_config.kv_transfer_config.is_kv_producer = False
ascend_config = mock.MagicMock()
ascend_config.recompute_scheduler_enable = False
with mock.patch("vllm_ascend.utils.get_ascend_config", return_value=ascend_config):
assert utils.is_pd_decode_recompute_scheduler_enabled(vllm_config) is False