@@ -15,12 +15,10 @@
|
||||
|
||||
import math
|
||||
import os
|
||||
from threading import Lock
|
||||
from unittest import mock
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from vllm.config import (CompilationConfig, ModelConfig, ParallelConfig,
|
||||
VllmConfig)
|
||||
|
||||
from tests.ut.base import TestBase
|
||||
from vllm_ascend import utils
|
||||
@@ -28,26 +26,14 @@ from vllm_ascend.utils import REGISTERED_ASCEND_OPS
|
||||
|
||||
|
||||
class TestUtils(TestBase):
|
||||
def setUp(self):
|
||||
import importlib
|
||||
|
||||
def test_is_310p(self):
|
||||
utils._IS_310P = None
|
||||
with mock.patch("vllm_ascend._build_info.__soc_version__",
|
||||
"Ascend310P3"):
|
||||
self.assertTrue(utils.is_310p())
|
||||
utils._IS_310P = None
|
||||
with mock.patch("vllm_ascend._build_info.__soc_version__",
|
||||
"Ascend910P1"):
|
||||
self.assertFalse(utils.is_310p())
|
||||
from vllm_ascend import platform
|
||||
|
||||
def test_sleep_mode_enabled(self):
|
||||
utils._SLEEP_MODE_ENABLED = None
|
||||
with mock.patch("vllm_ascend._build_info.__sleep_mode_enabled__",
|
||||
True):
|
||||
self.assertTrue(utils.sleep_mode_enabled())
|
||||
utils._SLEEP_MODE_ENABLED = None
|
||||
with mock.patch("vllm_ascend._build_info.__sleep_mode_enabled__",
|
||||
False):
|
||||
self.assertFalse(utils.sleep_mode_enabled())
|
||||
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
|
||||
@@ -83,22 +69,29 @@ class TestUtils(TestBase):
|
||||
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]]]])
|
||||
[
|
||||
[
|
||||
[
|
||||
[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):
|
||||
@@ -118,42 +111,20 @@ class TestUtils(TestBase):
|
||||
output_tensor = utils.aligned_16(input_tensor)
|
||||
self.assertEqual(output_tensor.shape[0], 32)
|
||||
|
||||
@mock.patch('importlib.util.find_spec')
|
||||
@mock.patch('importlib.import_module')
|
||||
def test_try_register_lib(self, mock_import_module, mock_find_spec):
|
||||
# import OK
|
||||
mock_find_spec.return_value = mock.MagicMock()
|
||||
mock_import_module.return_value = mock.MagicMock()
|
||||
lib_name = "existing_lib"
|
||||
lib_info = "Library found and imported successfully"
|
||||
utils.try_register_lib(lib_name, lib_info)
|
||||
|
||||
# Can't find lib
|
||||
mock_find_spec.return_value = None
|
||||
lib_name = "non_existing_lib"
|
||||
utils.try_register_lib(lib_name)
|
||||
|
||||
# import error
|
||||
mock_find_spec.return_value = mock.MagicMock()
|
||||
mock_import_module.side_effect = ImportError("import error")
|
||||
lib_name = "error_lib"
|
||||
utils.try_register_lib(lib_name)
|
||||
|
||||
@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:
|
||||
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.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"):
|
||||
@@ -163,6 +134,82 @@ class TestUtils(TestBase):
|
||||
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"):
|
||||
@@ -192,7 +239,6 @@ class TestUtils(TestBase):
|
||||
from transformers import PretrainedConfig
|
||||
|
||||
class SimpleConfig(PretrainedConfig):
|
||||
|
||||
def __init__(self, num_hidden_layers=12):
|
||||
self.num_hidden_layers = num_hidden_layers
|
||||
|
||||
@@ -203,39 +249,24 @@ class TestUtils(TestBase):
|
||||
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
|
||||
"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
|
||||
}
|
||||
}
|
||||
"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}
|
||||
|
||||
@@ -243,124 +274,214 @@ class TestUtils(TestBase):
|
||||
utils.get_max_hidden_layers(NoLayerConfig())
|
||||
self.assertIn("num_hidden_layers", str(context.exception))
|
||||
|
||||
def test_update_aclgraph_sizes(self):
|
||||
# max_num_batch_sizes < len(original_sizes)
|
||||
test_compilation_config = CompilationConfig(
|
||||
cudagraph_capture_sizes=[i for i in range(150)])
|
||||
model_path = os.path.join(os.path.dirname(__file__), "fake_weight")
|
||||
test_model_config = ModelConfig(model=model_path, enforce_eager=True)
|
||||
test_parallel_config = ParallelConfig()
|
||||
test_vllm_config = VllmConfig(
|
||||
model_config=test_model_config,
|
||||
compilation_config=test_compilation_config,
|
||||
parallel_config=test_parallel_config,
|
||||
)
|
||||
utils.update_aclgraph_sizes(test_vllm_config)
|
||||
os.environ['HCCL_OP_EXPANSION_MODE'] = 'AIV'
|
||||
utils.update_aclgraph_sizes(test_vllm_config)
|
||||
del os.environ['HCCL_OP_EXPANSION_MODE']
|
||||
self.assertEqual(
|
||||
137,
|
||||
len(test_vllm_config.compilation_config.cudagraph_capture_sizes))
|
||||
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,
|
||||
}
|
||||
|
||||
test_vllm_config.speculative_config = mock.MagicMock()
|
||||
test_vllm_config.speculative_config.draft_model_config = mock.MagicMock(
|
||||
)
|
||||
test_vllm_config.speculative_config.draft_model_config.hf_config = mock.MagicMock(
|
||||
)
|
||||
test_vllm_config.speculative_config.draft_model_config.hf_config.num_hidden_layers = 2
|
||||
os.environ['HCCL_OP_EXPANSION_MODE'] = 'AIV'
|
||||
utils.update_aclgraph_sizes(test_vllm_config)
|
||||
del os.environ['HCCL_OP_EXPANSION_MODE']
|
||||
self.assertEqual(
|
||||
111,
|
||||
len(test_vllm_config.compilation_config.cudagraph_capture_sizes))
|
||||
with mock.patch("vllm_ascend.utils._IS_DRAFTER_MOE_MODEL", None):
|
||||
self.assertFalse(utils.is_drafter_moe_model(vllm_config))
|
||||
|
||||
# max_num_batch_sizes >= len(original_sizes)
|
||||
test_compilation_config = CompilationConfig(
|
||||
cudagraph_capture_sizes=[1, 2, 3])
|
||||
test_vllm_config = VllmConfig(
|
||||
model_config=test_model_config,
|
||||
compilation_config=test_compilation_config,
|
||||
parallel_config=test_parallel_config,
|
||||
)
|
||||
utils.update_aclgraph_sizes(test_vllm_config)
|
||||
os.environ['HCCL_OP_EXPANSION_MODE'] = 'AIV'
|
||||
utils.update_aclgraph_sizes(test_vllm_config)
|
||||
del os.environ['HCCL_OP_EXPANSION_MODE']
|
||||
self.assertEqual(
|
||||
3,
|
||||
len(test_vllm_config.compilation_config.cudagraph_capture_sizes))
|
||||
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):
|
||||
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.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))
|
||||
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)
|
||||
|
||||
|
||||
class TestProfileExecuteDuration(TestBase):
|
||||
def test_is_pd_decode_recompute_scheduler_enabled_without_config():
|
||||
assert utils.is_pd_decode_recompute_scheduler_enabled() is False
|
||||
|
||||
def setUp(self):
|
||||
utils.ProfileExecuteDuration._instance = None
|
||||
utils.ProfileExecuteDuration._observations = []
|
||||
utils.ProfileExecuteDuration._lock = Lock()
|
||||
|
||||
def test_singleton_creation(self):
|
||||
instance1 = utils.ProfileExecuteDuration()
|
||||
self.assertIsNotNone(instance1)
|
||||
self.assertIs(instance1, utils.ProfileExecuteDuration._instance)
|
||||
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
|
||||
|
||||
instance2 = utils.ProfileExecuteDuration()
|
||||
self.assertIs(instance1, instance2)
|
||||
|
||||
def test_thread_safety(self):
|
||||
from threading import Thread
|
||||
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
|
||||
|
||||
instances = []
|
||||
|
||||
def create_instance():
|
||||
instances.append(utils.ProfileExecuteDuration())
|
||||
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
|
||||
|
||||
threads = [Thread(target=create_instance) for _ in range(10)]
|
||||
for t in threads:
|
||||
t.start()
|
||||
for t in threads:
|
||||
t.join()
|
||||
|
||||
first_instance = instances[0]
|
||||
for instance in instances[1:]:
|
||||
self.assertIs(first_instance, instance)
|
||||
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_atexit_registration(self):
|
||||
with mock.patch('atexit.register') as mock_register:
|
||||
instance = utils.ProfileExecuteDuration()
|
||||
mock_register.assert_called_once_with(instance.destroy)
|
||||
|
||||
def test_lock_usage(self):
|
||||
original_lock = utils.ProfileExecuteDuration._lock
|
||||
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
|
||||
|
||||
with mock.patch.object(utils.ProfileExecuteDuration,
|
||||
'_lock',
|
||||
wraps=original_lock) as mock_lock:
|
||||
utils.ProfileExecuteDuration()
|
||||
mock_lock.__enter__.assert_called()
|
||||
mock_lock.__exit__.assert_called()
|
||||
|
||||
def test_observations_initialization(self):
|
||||
instance = utils.ProfileExecuteDuration()
|
||||
self.assertEqual(instance._observations, [])
|
||||
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
|
||||
|
||||
Reference in New Issue
Block a user