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

0
tests/ut/eplb/core/__init__.py Normal file
View File

0
tests/ut/eplb/core/a2/__init__.py Normal file
View File

17
tests/ut/eplb/core/a2/expert_map.json Normal file
View File

@@ -0,0 +1,17 @@
{
"moe_layer_count":
1,
"layer_list": [{
"layer_id":
0,
"device_count":
2,
"device_list": [{
"device_id": 0,
"device_expert": [7, 2, 0, 3, 5]
}, {
"device_id": 1,
"device_expert": [6, 1, 4, 7, 2]
}]
}]
}

117
tests/ut/eplb/core/a2/test_eplb_utils.py Normal file
View File

@@ -0,0 +1,117 @@
import os
import unittest
from unittest.mock import MagicMock, patch
# isort: off
import torch
from vllm.config import VllmConfig
from vllm.model_executor.layers.fused_moe.config import FusedMoEConfig, FusedMoEParallelConfig
from vllm_ascend.ascend_config import init_ascend_config
from vllm_ascend.eplb.core.eplb_utils import generate_log2phy_map, init_eplb_config
from vllm_ascend.utils import vllm_version_is
# isort: on
class TestAscendConfig(unittest.TestCase):
@patch("vllm.config.VllmConfig.__post_init__", MagicMock())
@patch("vllm_ascend.platform.NPUPlatform._fix_incompatible_config")
def setUp(self, mock_fix_incompatible_config):
vllm_config = VllmConfig()
vllm_config.model_config = MagicMock()
vllm_config.additional_config = {
"refresh": True,
"eplb_config": {"dynamic_eplb": True, "num_redundant_experts": 2},
}
from vllm.model_executor.layers.fused_moe.config import RoutingMethodType
moe_parallel_config = FusedMoEParallelConfig(2, 0, 1, 2, 1, 1, 1, 1, 1, True, "hccl", enable_eplb=True)
if vllm_version_is("0.23.0"):
moe_config = FusedMoEConfig(
num_experts=8,
experts_per_token=8,
hidden_dim=8192,
intermediate_size_per_partition=5,
num_local_experts=8,
num_logical_experts=8,
activation="silu",
device="npu",
routing_method=RoutingMethodType.Simulated,
moe_parallel_config=moe_parallel_config,
in_dtype=torch.float16,
)
else:
from vllm.model_executor.layers.fused_moe.activation import MoEActivation
moe_config = FusedMoEConfig(
num_experts=8,
experts_per_token=8,
hidden_dim=8192,
intermediate_size=10,
num_local_experts=8,
num_logical_experts=8,
activation=MoEActivation.SILU,
device="npu",
routing_method=RoutingMethodType.Simulated,
moe_parallel_config=moe_parallel_config,
in_dtype=torch.float16,
)
moe_config.supports_eplb = True
self.vllm_config = vllm_config
self.moe_config = moe_config
self.mock_npu_patcher = patch("torch.Tensor.npu", new=lambda self: self)
self.mock_npu_patcher.start()
os.environ["DYNAMIC_EPLB"] = "true"
def tearDown(self):
self.mock_npu_patcher.stop()
os.environ.pop("DYNAMIC_EPLB", None)
def test_init_eplb_config_with_eplb(self):
eplb_config = init_ascend_config(self.vllm_config).eplb_config
_, expert_map, log2phy, redundant_experts = init_eplb_config(eplb_config, 0, self.moe_config)
gt_expert_map = torch.tensor([4, -1, -1, -1, 0, 1, 2, 3])
gt_log2phy = torch.tensor([9, 1, 2, 3, 5, 6, 7, 8])
self.assertTrue(torch.equal(expert_map, gt_expert_map))
self.assertTrue(torch.equal(log2phy, gt_log2phy))
self.assertEqual(redundant_experts, 2)
def test_init_eplb_config_with_eplb_withmap(self):
_TEST_DIR = os.path.dirname(__file__)
self.vllm_config.additional_config["eplb_config"]["expert_map_path"] = _TEST_DIR + "/expert_map.json"
eplb_config = init_ascend_config(self.vllm_config).eplb_config
_, expert_map, log2phy, redundant_experts = init_eplb_config(eplb_config, 0, self.moe_config)
gt_expert_map = torch.tensor([-1, 1, 4, -1, 2, -1, 0, 3])
gt_log2phy = torch.tensor([2, 6, 9, 3, 7, 4, 5, 8])
self.assertTrue(torch.equal(expert_map, gt_expert_map))
self.assertTrue(torch.equal(log2phy, gt_log2phy))
self.assertEqual(redundant_experts, 2)
def test_generate_log2phy_map_rotates_tail_tp_rank_with_tp_size(self):
global_expert_map = [
torch.tensor([0, -1], dtype=torch.int32),
torch.tensor([0, -1], dtype=torch.int32),
torch.tensor([0, -1], dtype=torch.int32),
torch.tensor([0, -1], dtype=torch.int32),
torch.tensor([-1, 0], dtype=torch.int32),
torch.tensor([-1, 0], dtype=torch.int32),
torch.tensor([-1, 0], dtype=torch.int32),
torch.tensor([-1, 0], dtype=torch.int32),
]
fallback_tail_dp1 = generate_log2phy_map(global_expert_map, ep_rank=7)
rotated_tail_dp0 = generate_log2phy_map(global_expert_map, ep_rank=3, tp_size=4)
rotated_tail_dp1 = generate_log2phy_map(global_expert_map, ep_rank=7, tp_size=4)
self.assertTrue(torch.equal(fallback_tail_dp1, torch.tensor([3, 7], dtype=torch.int32)))
self.assertTrue(torch.equal(rotated_tail_dp0, torch.tensor([3, 4], dtype=torch.int32)))
self.assertTrue(torch.equal(rotated_tail_dp1, torch.tensor([0, 5], dtype=torch.int32)))
def test_init_eplb_config_without_eplb(self):
self.vllm_config.additional_config = {"refresh": True}
eplb_config = init_ascend_config(self.vllm_config).eplb_config
_, expert_map, log2phy, redundant_experts = init_eplb_config(eplb_config, 0, self.moe_config)
gt_expert_map = torch.tensor([-1, -1, -1, -1, 0, 1, 2, 3])
self.assertIsNone(log2phy)
self.assertTrue(torch.equal(expert_map, gt_expert_map))
self.assertEqual(redundant_experts, 0)

0
tests/ut/eplb/core/policy/__init__.py Normal file
View File

36
tests/ut/eplb/core/policy/test_policy_factory.py Normal file
View File

@@ -0,0 +1,36 @@
import unittest
import torch
from vllm_ascend.eplb.core.eplb_worker import EplbWorker
from vllm_ascend.eplb.core.policy.policy_factory import PolicyFactory
from vllm_ascend.eplb.core.policy.policy_flashlb import generate_layered_experts
class TestEplbRebalancePolicies(unittest.TestCase):
def setUp(self):
torch.manual_seed(42)
self.current_expert_table = generate_layered_experts()
x = torch.rand(100, 58, 32, 9)
x = x**10
self.expert_workload = (x * 999 + 1).long()
self.hotness = EplbWorker._calculate_hotness(self.current_expert_table, self.expert_workload.sum(0))
@unittest.mock.patch("torch.npu.device_count", return_value=16)
def test_swift_balance_rebalance_experts(self, mock_count):
swift_policy = PolicyFactory.generate_policy(2)
_, _, new_placement = swift_policy.rebalance_experts(self.current_expert_table, self.expert_workload.sum(0))
update_mean, _ = EplbWorker._compute_imbalance(new_placement, self.hotness)
self.assertLessEqual(update_mean, 1.08)
def test_flashlb_rebalance_experts(self):
flashlb_policy = PolicyFactory.generate_policy(3)
_, _, new_placement = flashlb_policy.rebalance_experts(self.current_expert_table, self.expert_workload)
update_mean, _ = EplbWorker._compute_imbalance(new_placement, self.hotness)
self.assertLessEqual(update_mean, 1.1)
if __name__ == "__main__":
unittest.main(verbosity=2)

81
tests/ut/eplb/core/test_eplb_device_transfer_loader.py
View File

@@ -11,54 +11,61 @@ import vllm_ascend.eplb.core.eplb_device_transfer_loader as loader
def mock_adaptor():
adaptor = MagicMock()
adaptor.expert_map_per_layer_cpu = {
0: {
10: torch.tensor(1),
20: torch.tensor(0)
}
}
adaptor.expert_map_per_layer_cpu = {0: {10: torch.tensor(1), 20: torch.tensor(0)}}
adaptor.expert_param_per_layer = {
0: {
0: [[torch.tensor([1.0])]],
1: [[torch.tensor([2.0])]]
}
}
adaptor.expert_param_per_layer = {0: {0: [[torch.tensor([1.0])]], 1: [[torch.tensor([2.0])]]}}
adaptor.buffer_tensor_list = [[[torch.tensor([3.0])],
[torch.tensor([4.0])]]]
adaptor.expert_weight_key_per_layer = {0: "weight_key"}
adaptor.buffer_tensor_list = {
"weight_key": [[torch.tensor([3.0]), torch.tensor([4.0])], [torch.tensor([5.0]), torch.tensor([6.0])]]
}
return adaptor
def test_generate_task_and_state_flow(mock_adaptor):
loader_obj = loader.D2DExpertWeightLoader()
with patch("vllm_ascend.eplb.core.eplb_device_transfer_loader.get_dynamic_eplb_group", return_value=None):
loader_obj = loader.D2DExpertWeightLoader()
loader_obj.set_adator(mock_adaptor)
with patch("torch.distributed.P2POp") as mock_p2p, \
patch("torch.distributed.isend", return_value="isend_op"), \
patch("torch.distributed.irecv", return_value="irecv_op"):
with (
patch("torch.distributed.P2POp") as mock_p2p,
patch("torch.distributed.isend", return_value="isend_op"),
patch("torch.distributed.irecv", return_value="irecv_op"),
):
mock_p2p.side_effect = lambda op, tensor, rank: (op, tensor, rank)
loader_obj.state = loader.ExpertWeightUpdateState.READY
loader_obj.generate_expert_d2d_transfer_task([(1, 10)], [(2, 20)],
{20: torch.tensor(0)}, 0)
loader_obj.generate_expert_d2d_transfer_task([(1, 10)], [(2, 20)], {20: torch.tensor(0)}, 0)
assert loader_obj.comm_op_list is None
loader_obj.state = loader.ExpertWeightUpdateState.WAITING
loader_obj.generate_expert_d2d_transfer_task([], [], {}, 0)
assert loader_obj.comm_op_list is None
updated_map = {20: torch.tensor(0)}
loader_obj.generate_expert_d2d_transfer_task([(1, 10)], [(2, 20)],
updated_map, 0)
assert not loader_obj.comm_op_list
assert loader_obj.state == loader.ExpertWeightUpdateState.READY
assert loader_obj.comm_op_list
assert loader_obj.recv_expert_list
def test_generate_task_uses_layer_weight_key_buffer(mock_adaptor):
comm_group = MagicMock()
comm_group.ranks = {2: 20}
comm_group.device_group = object()
with patch("vllm_ascend.eplb.core.eplb_device_transfer_loader.get_dynamic_eplb_group", return_value=comm_group):
loader_obj = loader.D2DExpertWeightLoader()
loader_obj.set_adator(mock_adaptor)
with (
patch("torch.distributed.P2POp") as mock_p2p,
patch("torch.distributed.irecv", return_value="irecv_op"),
):
mock_p2p.side_effect = lambda op, tensor, rank, group=None: (op, tensor, rank, group)
loader_obj.generate_expert_d2d_transfer_task([], [(2, 20)], {20: torch.tensor(0)}, 0)
assert mock_p2p.call_args_list[0].args[1] is mock_adaptor.buffer_tensor_list["weight_key"][0][0]
assert mock_p2p.call_args_list[1].args[1] is mock_adaptor.buffer_tensor_list["weight_key"][0][1]
def test_asyn_transfer_and_update(mock_adaptor):
loader_obj = loader.D2DExpertWeightLoader()
with patch("vllm_ascend.eplb.core.eplb_device_transfer_loader.get_dynamic_eplb_group", return_value=None):
loader_obj = loader.D2DExpertWeightLoader()
loader_obj.set_adator(mock_adaptor)
loader_obj.comm_op_list = ["fake_op"]
@@ -66,8 +73,7 @@ def test_asyn_transfer_and_update(mock_adaptor):
reqs: list[MagicMock] = []
with patch("torch.distributed.batch_isend_irecv",
return_value=[MagicMock(), MagicMock()]):
with patch("torch.distributed.batch_isend_irecv", return_value=[MagicMock(), MagicMock()]):
loader_obj.asyn_expert_weight_transfer(reqs)
assert loader_obj.state == loader.ExpertWeightUpdateState.TRANSFERRING
@@ -94,14 +100,16 @@ def test_asyn_transfer_and_update(mock_adaptor):
def test_set_log2phy_map(mock_adaptor):
loader_obj = loader.D2DExpertWeightLoader()
with patch("vllm_ascend.eplb.core.eplb_device_transfer_loader.get_dynamic_eplb_group", return_value=None):
loader_obj = loader.D2DExpertWeightLoader()
loader_obj.set_adator(mock_adaptor)
loader_obj.set_log2phy_map({"a": 1})
assert loader_obj.updated_log2phy_map == {"a": 1}
def test_invalid_state_asyn_update(mock_adaptor):
loader_obj = loader.D2DExpertWeightLoader()
with patch("vllm_ascend.eplb.core.eplb_device_transfer_loader.get_dynamic_eplb_group", return_value=None):
loader_obj = loader.D2DExpertWeightLoader()
loader_obj.set_adator(mock_adaptor)
loader_obj.state = loader.ExpertWeightUpdateState.WAITING
@@ -113,10 +121,3 @@ def test_invalid_state_asyn_update(mock_adaptor):
loader_obj.update_expert_map_and_weight([])
assert not mock_adaptor.do_update_expert_map.called
def test_load_impl_not_implemented(mock_adaptor):
loader_obj = loader.D2DExpertWeightLoader()
loader_obj.set_adator(mock_adaptor)
with pytest.raises(NotImplementedError):
loader_obj.load_impl({}, {})