EngineX-Ascend/enginex-ascend-910-vllm
11
0
Fork 1
Code Issues Pull Requests Actions Projects Releases Wiki Activity

init v0.23.0

Browse Source 
Signed-off-by: Sun Ruoxi <sunruoxi@4paradigm.com>
This commit is contained in:
Sun Ruoxi
2026-08-27 15:11:51 +08:00
parent b582a8e7d1
commit 7f8a1b1f7a
2849 changed files with 712887 additions and 22001 deletions
Download Patch File Download Diff File Expand all files Collapse all files

120
tests/ut/ops/test_linear.py
View File

@@ -1,18 +1,21 @@
import os
import unittest
from unittest import mock
from unittest.mock import MagicMock, patch
import torch
from tests.ut.base import TestBase
from vllm_ascend import ascend_config
from vllm_ascend.distributed import parallel_state
from vllm_ascend.ops.linear import (AscendMergedColumnParallelLinear,
AscendRowParallelLinear)
from vllm_ascend.ops.linear import (
AscendMergedColumnParallelLinear,
AscendReplicatedLinear,
AscendRowParallelLinear,
AscendUnquantizedLinearMethod,
)
class BaseLinearTest(unittest.TestCase):
def setUp(self):
self.mock_group = mock.MagicMock()
self.mock_group.world_size = 2
@@ -22,23 +25,22 @@ class BaseLinearTest(unittest.TestCase):
parallel_state._OTP = self.mock_group
self.mock_ascend_config = MagicMock()
self.mock_ascend_config.oproj_tensor_parallel_size = 2
self.mock_ascend_config.finegrained_tp_config.oproj_tensor_parallel_size = 2
self.mock_ascend_config.finegrained_tp_config.mlp_tensor_parallel_size = 2
self.patches = [
patch("vllm_ascend.ascend_config.get_ascend_config",
return_value=self.mock_ascend_config),
patch("vllm_ascend.distributed.parallel_state.get_otp_group",
return_value=self.mock_group),
patch("vllm_ascend.distributed.parallel_state.get_mlp_tp_group",
return_value=self.mock_group),
patch("vllm_ascend.ops.linear_op.get_tp_group",
return_value=self.mock_group),
patch("vllm_ascend.ascend_config.get_ascend_config", return_value=self.mock_ascend_config),
patch("vllm_ascend.distributed.parallel_state.get_otp_group", return_value=self.mock_group),
patch("vllm_ascend.distributed.parallel_state.get_mlp_tp_group", return_value=self.mock_group),
patch("vllm_ascend.ops.linear_op.get_tp_group", return_value=self.mock_group),
patch(
"vllm.distributed.parallel_state.get_tp_group",
return_value=self.mock_group,
),
patch("vllm_ascend.utils.mlp_tp_enable", return_value=True),
patch("vllm_ascend.utils.oproj_tp_enable", return_value=True)
patch("vllm_ascend.utils.oproj_tp_enable", return_value=True),
patch("vllm_ascend.ops.linear_op.enable_dsa_cp", return_value=False),
patch("vllm_ascend.ops.linear_op.enable_dsa_cp_with_layer_shard", return_value=False),
]
for p in self.patches:
@@ -49,10 +51,57 @@ class BaseLinearTest(unittest.TestCase):
p.stop()
class TestAscendRowParallelLinear(BaseLinearTest):
class TestAscendUnquantizedLinearMethod(TestBase):
def setUp(self):
self.method = AscendUnquantizedLinearMethod()
self.layer = mock.MagicMock()
mock_dtype = mock.PropertyMock(return_value=torch.float16)
type(self.layer.weight.data).dtype = mock_dtype
mock_is_meta = mock.PropertyMock(return_value=False)
type(self.layer.weight.data).is_meta = mock_is_meta
self.layer.precast_fp32_weight = False
def test_mlp_optimize(self):
os.environ["VLLM_ASCEND_ENABLE_MLP_OPTIMIZE"] = "1"
@patch("vllm_ascend.utils.get_ascend_config")
@mock.patch("torch_npu.npu_format_cast")
def test_process_weights_after_loading_with_nz0(self, mock_format_cast, mock_get_config):
mock_config = MagicMock()
mock_config.weight_nz_mode = 0
mock_get_config.return_value = mock_config
self.method.process_weights_after_loading(self.layer)
mock_format_cast.assert_not_called()
@patch("vllm_ascend.utils.get_ascend_config")
@mock.patch("torch_npu.npu_format_cast")
def test_process_weights_after_loading_with_nz1(self, mock_format_cast, mock_get_config):
mock_config = MagicMock()
mock_config.weight_nz_mode = 1
mock_get_config.return_value = mock_config
self.method.process_weights_after_loading(self.layer)
mock_format_cast.assert_not_called()
@patch("vllm_ascend.utils.get_ascend_config")
@mock.patch("torch_npu.npu_format_cast")
def test_process_weights_after_loading_with_nz2(self, mock_format_cast, mock_get_config):
mock_config = MagicMock()
mock_config.weight_nz_mode = 2
mock_get_config.return_value = mock_config
self.method.process_weights_after_loading(self.layer)
mock_format_cast.assert_called_once()
class TestAscendRowParallelLinear(BaseLinearTest):
@patch("vllm_ascend.ops.linear_op.get_weight_prefetch_method", return_value=MagicMock())
@patch("vllm_ascend.ops.linear.get_current_vllm_config", return_value=MagicMock())
@patch("vllm_ascend.ops.linear.enable_sp", return_value=False)
@patch(
"vllm_ascend.ops.linear.AscendUnquantizedLinearMethod.apply",
new=lambda self, layer, x, bias=None: torch.nn.functional.linear(x, layer.weight, bias),
)
def test_mlp_optimize(self, mock_enable_sp, mock_get_current_vllm_config, mock_get_weight_prefetch_method):
ascend_config._ASCEND_CONFIG = MagicMock()
ascend_config._ASCEND_CONFIG.recompute_scheduler_enable = False
ascend_config._ASCEND_CONFIG.finegrained_tp_config.mlp_tensor_parallel_size = 2
ascend_config._ASCEND_CONFIG.ascend_scheduler_config.enabled = False
linear = AscendRowParallelLinear(
input_size=16,
@@ -64,9 +113,18 @@ class TestAscendRowParallelLinear(BaseLinearTest):
input_tensor = torch.randn(16, 8)
linear(input_tensor)
def test_oproj_tp(self):
@patch("vllm_ascend.ops.linear_op.get_weight_prefetch_method", return_value=MagicMock())
@patch("vllm_ascend.ops.linear.get_current_vllm_config", return_value=MagicMock())
@patch("vllm_ascend.ops.linear.enable_sp", return_value=False)
@patch(
"vllm_ascend.ops.linear.AscendUnquantizedLinearMethod.apply",
new=lambda self, layer, x, bias=None: torch.nn.functional.linear(x, layer.weight, bias),
)
def test_oproj_tp(self, mock_enable_sp, mock_get_current_vllm_config, mock_get_weight_prefetch_method):
ascend_config._ASCEND_CONFIG = MagicMock()
ascend_config._ASCEND_CONFIG.oproj_tensor_parallel_size = 2
ascend_config._ASCEND_CONFIG.recompute_scheduler_enable = False
ascend_config._ASCEND_CONFIG.finegrained_tp_config.oproj_tensor_parallel_size = 2
ascend_config._ASCEND_CONFIG.ascend_scheduler_config.enabled = False
linear = AscendRowParallelLinear(
input_size=16,
@@ -80,9 +138,11 @@ class TestAscendRowParallelLinear(BaseLinearTest):
class TestAscendMergedColumnParallelLinear(BaseLinearTest):
def test_merged_mlp_tp_init(self):
os.environ["VLLM_ASCEND_ENABLE_MLP_OPTIMIZE"] = "1"
ascend_config._ASCEND_CONFIG = MagicMock()
ascend_config._ASCEND_CONFIG.recompute_scheduler_enable = False
ascend_config._ASCEND_CONFIG.finegrained_tp_config.mlp_tensor_parallel_size = 2
ascend_config._ASCEND_CONFIG.ascend_scheduler_config.enabled = False
linear = AscendMergedColumnParallelLinear(
input_size=16,
@@ -92,5 +152,21 @@ class TestAscendMergedColumnParallelLinear(BaseLinearTest):
self.assertEqual(linear.custom_op.comm_group, parallel_state._MLP_TP)
if __name__ == '__main__':
class TestAscendReplicatedLinear(BaseLinearTest):
def test_init_disable_tp(self):
linear = AscendReplicatedLinear(
input_size=16,
output_size=8,
)
self.assertTrue(isinstance(linear.quant_method, AscendUnquantizedLinearMethod))
def test_init_without_disable_tp(self):
linear = AscendReplicatedLinear(
input_size=16,
output_size=8,
)
self.assertTrue(isinstance(linear.quant_method, AscendUnquantizedLinearMethod))
if __name__ == "__main__":
unittest.main()