EngineX/xc-llm-ascend
3
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

[EPLB][Bugfix] policy_swift_balancer bugfix and renaming (#5897)

Browse Source 
### What this PR does / why we need it?
1. Rename dynamic_ep to default_eplb.
2. Rename dynamic_ep_v2 to swift_balancer
3. Discard func compose_expert_update_info_bipartite.

- vLLM version: v0.13.0
- vLLM main:
bde38c11df

Signed-off-by: shenchuxiaofugui <1311027364@qq.com>
This commit is contained in:
LI SHENGYONG
2026-01-19 13:47:40 +08:00
committed by GitHub
parent b27774dbd6
commit 83de5385b4
8 changed files with 39 additions and 149 deletions
Download Patch File Download Diff File Expand all files Collapse all files

30
tests/ut/eplb/core/policy/test_policy_dynamic_ep.py → tests/ut/eplb/core/policy/test_policy_default_eplb.py
View File

@@ -3,16 +3,16 @@ from unittest.mock import patch
import numpy as np
import pytest
from vllm_ascend.eplb.core.policy.policy_dynamic_ep import DynamicEplb
from vllm_ascend.eplb.core.policy.policy_default_eplb import DefaultEplb
class TestDynamicEplb:
class TestDefaultEplb:
def test_add_redundant_basic(self):
current_expert_table = np.array([[[0, 1], [1, 0]]])
expert_workload = np.array([[[2, 3], [4, 1]]])
num_original_expert = 2
result = DynamicEplb.add_redundant(current_expert_table,
result = DefaultEplb.add_redundant(current_expert_table,
expert_workload,
num_original_expert)
expected = np.array([[2 + 1, 3 + 4]])
@@ -20,51 +20,51 @@ class TestDynamicEplb:
def test_get_redundant_num(self):
counts = np.array([2, 1, 3])
assert DynamicEplb.get_redundant_num(3, counts) == 3
assert DefaultEplb.get_redundant_num(3, counts) == 3
def test_calculate_max_heat_per_layer(self):
workload_table = np.array([[[1, 2], [3, 4]], [[2, 2], [1, 1]]])
max_heat = DynamicEplb.calculate_max_heat_per_layer(workload_table, 2)
max_heat = DefaultEplb.calculate_max_heat_per_layer(workload_table, 2)
assert max_heat == [7, 4]
def test_constraint_expert_local_exchange(self):
current = [[[0, 1], [2, 3]]]
global_dep = [[[1, 0], [3, 2]]]
new_dep = DynamicEplb.constraint_expert_local_exchange(
new_dep = DefaultEplb.constraint_expert_local_exchange(
current, global_dep)
assert new_dep == [[[0, 1], [2, 3]]]
def test_compute_balanced_pack_redundancy_normal(self):
origin_weights = [(0, 10), (1, 20)]
result, boxes = DynamicEplb.compute_balanced_pack_redundancy(
result, boxes = DefaultEplb.compute_balanced_pack_redundancy(
origin_weights, 2, 1)
assert isinstance(result, list) and len(result) == 2
def test_compute_balanced_pack_redundancy_card0(self):
origin_weights = [(0, 10)]
with pytest.raises(RuntimeError):
DynamicEplb.compute_balanced_pack_redundancy(origin_weights, 0, 0)
DefaultEplb.compute_balanced_pack_redundancy(origin_weights, 0, 0)
def test_compute_balanced_pack_normal(self):
origin_weights = np.array([(0, 10), (1, 20)], dtype=object)
result, boxes = DynamicEplb.compute_balanced_pack(origin_weights, 2)
result, boxes = DefaultEplb.compute_balanced_pack(origin_weights, 2)
assert isinstance(result, list) and len(result) == 2
def test_compute_balanced_pack_card0(self):
origin_weights = np.array([(0, 10)], dtype=object)
with pytest.raises(RuntimeError):
DynamicEplb.compute_balanced_pack(origin_weights, 0)
DefaultEplb.compute_balanced_pack(origin_weights, 0)
def test_original_compute_balanced_pack_redundancy(self):
origin_weights = [(0, 5), (1, 10)]
result, boxes = DynamicEplb.original_compute_balanced_pack_redundancy(
result, boxes = DefaultEplb.original_compute_balanced_pack_redundancy(
origin_weights, 2, 1)
assert isinstance(result, list) and len(result) == 2
def test_rebalance_experts_normal(self):
expert_table = np.array([[[0, 1], [1, 0]]])
workload = np.array([[[2, 3], [4, 1]]])
policy = DynamicEplb(config=None)
policy = DefaultEplb(config=None)
change, priority, new_dep = policy.rebalance_experts(
expert_table, workload)
assert change in [0, 1]
@@ -73,12 +73,12 @@ class TestDynamicEplb:
assert np.array(new_dep).shape == expert_table.shape
def test_rebalance_experts_exceptions(self):
policy = DynamicEplb(config=None)
policy = DefaultEplb(config=None)
# case1: num_original_expert != expert_num
expert_table = np.array([[[0, 1], [1, 0]]])
workload = np.array([[[2, 3], [4, 1]]])
with patch.object(DynamicEplb,
with patch.object(DefaultEplb,
'add_redundant',
return_value=np.array([[1, 2, 3]])):
with pytest.raises(ValueError):
@@ -93,6 +93,6 @@ class TestDynamicEplb:
# case3: num_npus < num_redundancy_expert
expert_table_small = np.array([[[0, 0]]]) # 1 layer, 1 NPU, 2 experts
workload_small = np.array([[[1, 1]]])
with patch.object(DynamicEplb, 'get_redundant_num', return_value=2):
with patch.object(DefaultEplb, 'get_redundant_num', return_value=2):
with pytest.raises(ValueError):
policy.rebalance_experts(expert_table_small, workload_small)

8
tests/ut/eplb/core/policy/test_policy_factor.py
View File

@@ -1,8 +1,8 @@
import pytest
from vllm_ascend.eplb.core.policy.policy_abstract import DynamicConfig
from vllm_ascend.eplb.core.policy.policy_dynamic_ep import DynamicEplb
from vllm_ascend.eplb.core.policy.policy_dynamic_ep_v2 import DynamicEplbV2
from vllm_ascend.eplb.core.policy.policy_default_eplb import DefaultEplb
from vllm_ascend.eplb.core.policy.policy_swift_balancer import SwiftBalanceEplb
from vllm_ascend.eplb.core.policy.policy_factory import PolicyFactory
from vllm_ascend.eplb.core.policy.policy_random import RandomLoadBalance
@@ -14,8 +14,8 @@ def dummy_config():
@pytest.mark.parametrize("policy_type, expected_class", [
(0, RandomLoadBalance),
(1, DynamicEplb),
(2, DynamicEplbV2),
(1, DefaultEplb),
(2, SwiftBalanceEplb),
(999, RandomLoadBalance),
])
def test_generate_policy(policy_type, expected_class, dummy_config):

16
tests/ut/eplb/core/policy/test_policy_dynamic_ep_v2.py → tests/ut/eplb/core/policy/test_policy_swift_balancer.py
View File

@@ -3,8 +3,8 @@ from typing import Dict, Set
import numpy as np
import pytest
from vllm_ascend.eplb.core.policy.policy_dynamic_ep_v2 import (DynamicConfig,
DynamicEplbV2)
from vllm_ascend.eplb.core.policy.policy_swift_balancer import (DynamicConfig,
SwiftBalanceEplb)
@pytest.fixture
@@ -14,7 +14,7 @@ def config():
@pytest.fixture
def policy(config):
return DynamicEplbV2(config)
return SwiftBalanceEplb(config)
def test_safe_operations(policy):
@@ -34,19 +34,19 @@ def test_safe_operations(policy):
def test_add_redundant():
workload = np.array([[[1, 2], [3, 4]]])
placement = np.array([[[0, 1], [0, 1]]])
result = DynamicEplbV2.add_redundant(placement, workload, 2)
result = SwiftBalanceEplb.add_redundant(placement, workload, 2)
assert result.shape == (1, 2)
assert np.all(result[0] == [4, 6]) # 0:1+3, 1:2+4
def test_get_redundant_num():
counts = np.array([1, 2, 1])
assert DynamicEplbV2.get_redundant_num(3, counts) == 1 # sum(counts-1)
assert SwiftBalanceEplb.get_redundant_num(3, counts) == 1 # sum(counts-1)
def test_calculate_max_heat_per_layer():
workload = np.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]])
result = DynamicEplbV2.calculate_max_heat_per_layer(workload, 2)
result = SwiftBalanceEplb.calculate_max_heat_per_layer(workload, 2)
assert result == [7, 15]
@@ -69,7 +69,7 @@ def test_compute_redundant_assignments(policy):
def test_prepare_expert_list():
base_experts = [(0, 10), (1, 5)]
redundant_assignments = [[2], []]
result = DynamicEplbV2.prepare_expert_list(base_experts,
result = SwiftBalanceEplb.prepare_expert_list(base_experts,
redundant_assignments, 1)
assert isinstance(result, list)
assert len(result) == 1
@@ -79,7 +79,7 @@ def test_non_redundant_expert_information():
origin_deployment = np.array([[0, 1]])
updated_weights = [(0, 10), (1, 5)]
rendun_pos: Dict[int, Set[int]] = {0: set()}
assignments, weights, loads, counts = DynamicEplbV2.non_redundant_expert_information(
assignments, weights, loads, counts = SwiftBalanceEplb.non_redundant_expert_information(
origin_deployment, updated_weights, rendun_pos)
assert assignments[0] == [0, 1]
assert loads[0] == 15