Skip to content

Commit 5e19a5a

Browse files
freyfwtfreyfwt
authored
[Refactor][EPLB] Support mixed expert weight layouts (vllm-project#11120)
### What this PR does / why we need it? This PR refactors EPLB expert weight handling so layers can use different expert weight layouts in the same model. - Adds per-layer expert weight keys based on quantization type and fused MC2 state. - Builds EPLB buffer tensors per expert weight key, with shape validation before reusing buffers across layers. - Uses the layer's expert weight key when generating D2D expert transfer tasks. - Propagates EPLB heat collection status into draft model forward contexts so registered draft MoE layers do not stay at zero load. - Fixes the non-last PP rank `forward_end()` path to use the current heat collection status. This also avoids `Expert hotness` mean becoming `nan` when registered draft or PP-stage MoE layers did not collect heat. ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? - Local formatting/preflight: - `bash format.sh ci` - Local UT: - `tests/ut/eplb/adaptor/test_vllm_adaptor.py` - Result: `6 tests OK` - NPU focused UT: - `python -m pytest -q --confcutdir=tests/ut/eplb tests/ut/eplb --tb=short` - Result: `19 passed, 17 warnings in 24.77s` - vLLM version: v0.23.0 - vLLM main: vllm-project/vllm@dc68bd8 --------- Signed-off-by: freyfwt <freytian1996@gmail.com>
1 parent 23bdf08 commit 5e19a5a

7 files changed

Lines changed: 193 additions & 56 deletions

File tree

‎.github/workflows/scripts/run_selected_tests.sh‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,15 @@ pytest_log_dir="${RUNNER_TEMP:-/tmp}/selected-tests-${npu_type}-${num_npus}card"
3232

3333
mkdir -p "${pytest_log_dir}"
3434

35+
setup_vllm_cache_root() {
36+
if [ "${CI:-}" != "true" ]; then
37+
return
38+
fi
39+
export VLLM_CACHE_ROOT
40+
VLLM_CACHE_ROOT="$(mktemp -d "${RUNNER_TEMP:-/tmp}/vllm-cache-${npu_type}-${num_npus}card.XXXXXX")"
41+
echo "Using vLLM cache root: ${VLLM_CACHE_ROOT}"
42+
}
43+
3544
print_test_info() {
3645
echo -e "\033[1;34m=== TEST INFO ===\033[0m"
3746
echo -e " \033[33mDevice:\033[0m ${npu_type}"
@@ -153,6 +162,7 @@ print_timing_json() {
153162
}
154163

155164
print_test_info
165+
setup_vllm_cache_root
156166

157167
if [ "${npu_type}" = "cpu" ]; then
158168
run_pytest_batch "cpu-ut (${#targets[@]} targets)" "${targets[@]}"

‎tests/ut/eplb/adaptor/test_vllm_adaptor.py‎

Lines changed: 91 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -4,8 +4,8 @@
44
import torch
55
from transformers import DeepseekV2Config
66

7-
from vllm_ascend.eplb.adaptor.vllm_adaptor import VllmEplbAdaptor
8-
from vllm_ascend.quantization.methods.base import QuantType
7+
from vllm_ascend.eplb.adaptor.vllm_adaptor import EPLB_EXPERT_WEIGHT_NAMES, VllmEplbAdaptor
8+
from vllm_ascend.quantization.quant_type import QuantType
99

1010

1111
class TestVllmAdaptor(unittest.TestCase):
@@ -42,9 +42,14 @@ def setUp(self):
4242
self.mock_size = patch("vllm_ascend.eplb.adaptor.vllm_adaptor.dist.get_world_size", return_value=4).start()
4343

4444
@patch("torch.empty_like", return_value=torch.zeros(16, 32))
45-
def test_init_fp16(self, mock_func):
45+
@patch("vllm_ascend.eplb.adaptor.vllm_adaptor.get_ascend_config")
46+
def test_init_fp16(self, mock_get_config, mock_func):
47+
mock_config = MagicMock()
48+
mock_config.enable_fused_mc2 = 1
49+
mock_get_config.return_value = mock_config
4650
self.model.quant_config = None
47-
VllmEplbAdaptor(self.model)
51+
adaptor = VllmEplbAdaptor(self.model)
52+
self.assertEqual(adaptor.expert_weight_key_per_layer[0], (QuantType.NONE, True))
4853

4954
@patch("torch.empty_like", return_value=torch.zeros(16, 32))
5055
@patch("vllm_ascend.eplb.adaptor.vllm_adaptor.get_ascend_config")
@@ -95,6 +100,88 @@ def test_pp_eplb_adaptor_init_with_registered_layer(self):
95100
self.assertEqual(adaptor.num_local_experts, 4)
96101
self.assertEqual(adaptor.ep_rank, 0)
97102

103+
@patch("vllm_ascend.eplb.adaptor.vllm_adaptor.get_ascend_config")
104+
def test_init_mixed_quant_type_per_layer(self, mock_get_config):
105+
mock_config = MagicMock()
106+
mock_config.enable_fused_mc2 = 1
107+
mock_get_config.return_value = mock_config
108+
109+
VllmEplbAdaptor._registered_moe_layers = []
110+
num_local_experts = 2
111+
w8a8_layer = MagicMock()
112+
w8a8_layer.local_num_experts = num_local_experts
113+
w8a8_layer.ep_rank = 0
114+
w8a8_layer.quant_type = QuantType.W8A8
115+
w8a8_layer.w13_weight_list = [torch.randn(2, 2) for _ in range(num_local_experts)]
116+
w8a8_layer.w2_weight_list = [torch.randn(2, 2) for _ in range(num_local_experts)]
117+
w8a8_layer.w13_weight_scale_fp32_list = [torch.randn(1) for _ in range(num_local_experts)]
118+
w8a8_layer.w2_weight_scale_list = [torch.randn(1) for _ in range(num_local_experts)]
119+
w8a8_layer.fused_w1_scale_list = [torch.randn(1) for _ in range(num_local_experts)]
120+
w8a8_layer.fused_w2_scale_list = [torch.randn(1) for _ in range(num_local_experts)]
121+
w8a8_layer.moe_load = torch.zeros(num_local_experts)
122+
w8a8_layer.global_expert_map = torch.arange(num_local_experts * 4).reshape(num_local_experts, 4)
123+
w8a8_layer.get_log2phy_map.return_value = torch.arange(4)
124+
125+
mxfp8_layer = MagicMock()
126+
mxfp8_layer.local_num_experts = num_local_experts
127+
mxfp8_layer.ep_rank = 0
128+
mxfp8_layer.quant_type = QuantType.MXFP8
129+
mxfp8_layer.w13_weight = torch.randn(num_local_experts, 2, 2)
130+
mxfp8_layer.w2_weight = torch.randn(num_local_experts, 2, 2)
131+
mxfp8_layer.w13_weight_scale = torch.randn(num_local_experts, 1)
132+
mxfp8_layer.w2_weight_scale = torch.randn(num_local_experts, 1)
133+
mxfp8_layer.moe_load = torch.zeros(num_local_experts)
134+
mxfp8_layer.global_expert_map = torch.arange(num_local_experts * 4).reshape(num_local_experts, 4)
135+
mxfp8_layer.get_log2phy_map.return_value = torch.arange(4)
136+
137+
VllmEplbAdaptor.register_layer(w8a8_layer)
138+
VllmEplbAdaptor.register_layer(mxfp8_layer)
139+
140+
model = MagicMock()
141+
model.quant_config = MagicMock()
142+
model.config.first_k_dense_replace = 0
143+
del model.language_model
144+
adaptor = VllmEplbAdaptor(model)
145+
146+
w8a8_key = (QuantType.W8A8, True)
147+
mxfp8_key = (QuantType.MXFP8, True)
148+
self.assertEqual(adaptor.expert_weight_key_per_layer[0], w8a8_key)
149+
self.assertEqual(adaptor.expert_weight_key_per_layer[1], mxfp8_key)
150+
self.assertEqual(len(adaptor.buffer_tensor_list[w8a8_key][0]), len(EPLB_EXPERT_WEIGHT_NAMES[w8a8_key]))
151+
self.assertEqual(len(adaptor.buffer_tensor_list[mxfp8_key][0]), len(EPLB_EXPERT_WEIGHT_NAMES[mxfp8_key]))
152+
self.assertEqual(len(adaptor.expert_param_per_layer[0][0]), len(EPLB_EXPERT_WEIGHT_NAMES[w8a8_key]))
153+
self.assertEqual(len(adaptor.expert_param_per_layer[1][0]), len(EPLB_EXPERT_WEIGHT_NAMES[mxfp8_key]))
154+
155+
@patch("vllm_ascend.eplb.adaptor.vllm_adaptor.get_ascend_config")
156+
def test_reused_buffer_requires_same_expert_weight_shape(self, mock_get_config):
157+
mock_config = MagicMock()
158+
mock_config.enable_fused_mc2 = 0
159+
mock_get_config.return_value = mock_config
160+
161+
VllmEplbAdaptor._registered_moe_layers = []
162+
num_local_experts = 2
163+
for weight_shape in [(2, 2), (3, 2)]:
164+
layer = MagicMock()
165+
layer.local_num_experts = num_local_experts
166+
layer.ep_rank = 0
167+
layer.quant_type = QuantType.W8A8
168+
layer.w13_weight_list = [torch.randn(*weight_shape) for _ in range(num_local_experts)]
169+
layer.w2_weight_list = [torch.randn(2, 2) for _ in range(num_local_experts)]
170+
layer.w13_weight_scale_fp32_list = [torch.randn(1) for _ in range(num_local_experts)]
171+
layer.w2_weight_scale_list = [torch.randn(1) for _ in range(num_local_experts)]
172+
layer.moe_load = torch.zeros(num_local_experts)
173+
layer.global_expert_map = torch.arange(num_local_experts * 4).reshape(num_local_experts, 4)
174+
layer.get_log2phy_map.return_value = torch.arange(4)
175+
VllmEplbAdaptor.register_layer(layer)
176+
177+
model = MagicMock()
178+
model.quant_config = MagicMock()
179+
model.config.first_k_dense_replace = 0
180+
del model.language_model
181+
182+
with self.assertRaisesRegex(AssertionError, "EPLB expert weight shapes mismatch"):
183+
VllmEplbAdaptor(model)
184+
98185
def tearDown(self):
99186
self.mock_rank.stop()
100187
self.mock_size.stop()

‎tests/ut/eplb/core/test_eplb_device_transfer_loader.py‎

Lines changed: 23 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,10 @@ def mock_adaptor():
1515

1616
adaptor.expert_param_per_layer = {0: {0: [[torch.tensor([1.0])]], 1: [[torch.tensor([2.0])]]}}
1717

18-
adaptor.buffer_tensor_list = [[[torch.tensor([3.0])], [torch.tensor([4.0])]]]
18+
adaptor.expert_weight_key_per_layer = {0: "weight_key"}
19+
adaptor.buffer_tensor_list = {
20+
"weight_key": [[torch.tensor([3.0]), torch.tensor([4.0])], [torch.tensor([5.0]), torch.tensor([6.0])]]
21+
}
1922
return adaptor
2023

2124

@@ -41,6 +44,25 @@ def test_generate_task_and_state_flow(mock_adaptor):
4144
assert loader_obj.state == loader.ExpertWeightUpdateState.READY
4245

4346

47+
def test_generate_task_uses_layer_weight_key_buffer(mock_adaptor):
48+
comm_group = MagicMock()
49+
comm_group.ranks = {2: 20}
50+
comm_group.device_group = object()
51+
with patch("vllm_ascend.eplb.core.eplb_device_transfer_loader.get_dynamic_eplb_group", return_value=comm_group):
52+
loader_obj = loader.D2DExpertWeightLoader()
53+
loader_obj.set_adator(mock_adaptor)
54+
55+
with (
56+
patch("torch.distributed.P2POp") as mock_p2p,
57+
patch("torch.distributed.irecv", return_value="irecv_op"),
58+
):
59+
mock_p2p.side_effect = lambda op, tensor, rank, group=None: (op, tensor, rank, group)
60+
loader_obj.generate_expert_d2d_transfer_task([], [(2, 20)], {20: torch.tensor(0)}, 0)
61+
62+
assert mock_p2p.call_args_list[0].args[1] is mock_adaptor.buffer_tensor_list["weight_key"][0][0]
63+
assert mock_p2p.call_args_list[1].args[1] is mock_adaptor.buffer_tensor_list["weight_key"][0][1]
64+
65+
4466
def test_asyn_transfer_and_update(mock_adaptor):
4567
with patch("vllm_ascend.eplb.core.eplb_device_transfer_loader.get_dynamic_eplb_group", return_value=None):
4668
loader_obj = loader.D2DExpertWeightLoader()

‎vllm_ascend/eplb/adaptor/vllm_adaptor.py‎

Lines changed: 63 additions & 49 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,38 @@
2323
from vllm.logger import logger
2424

2525
from vllm_ascend.ascend_config import get_ascend_config
26-
from vllm_ascend.quantization.methods.base import QuantType
26+
from vllm_ascend.quantization.quant_type import QuantType
27+
28+
EPLB_EXPERT_WEIGHT_NAMES = {
29+
(QuantType.NONE, False): ("w13_weight", "w2_weight"),
30+
(QuantType.NONE, True): ("w13_weight", "w2_weight"),
31+
(QuantType.W8A8, False): (
32+
"w13_weight_list",
33+
"w2_weight_list",
34+
"w13_weight_scale_fp32_list",
35+
"w2_weight_scale_list",
36+
),
37+
(QuantType.W8A8, True): (
38+
"w13_weight_list",
39+
"w2_weight_list",
40+
"w13_weight_scale_fp32_list",
41+
"w2_weight_scale_list",
42+
"fused_w1_scale_list",
43+
"fused_w2_scale_list",
44+
),
45+
(QuantType.W4A8, True): (
46+
"w13_weight_list",
47+
"w2_weight_list",
48+
"w13_weight_scale_list",
49+
"w2_weight_scale_list",
50+
"w13_scale_bias_list",
51+
"w2_scale_bias_list",
52+
),
53+
(QuantType.MXFP4, False): ("w13_weight", "w2_weight", "w13_weight_scale", "w2_weight_scale"),
54+
(QuantType.MXFP4, True): ("w13_weight", "w2_weight", "w13_weight_scale", "w2_weight_scale"),
55+
(QuantType.MXFP8, False): ("w13_weight", "w2_weight", "w13_weight_scale", "w2_weight_scale"),
56+
(QuantType.MXFP8, True): ("w13_weight", "w2_weight", "w13_weight_scale", "w2_weight_scale"),
57+
}
2758

2859

2960
class VllmEplbAdaptor:
@@ -61,73 +92,54 @@ def __init__(self, model, **args):
6192
self.ep_rank = first_layer.ep_rank
6293

6394
self.expert_param_per_layer = dict()
95+
self.expert_weight_key_per_layer = dict()
6496
self.init_expert_param_per_layer()
6597

6698
num_buffer_tensor = self.num_local_experts
67-
self.buffer_tensor_list: list[list[Any]] = [[] for _ in range(num_buffer_tensor)]
99+
self.buffer_tensor_list: dict[Any, list[list[Any]]] = dict()
68100
self.init_buffer_tensor(num_buffer_tensor)
69101

70102
self.log2phy_map_per_layer = dict()
71103
for local_idx, layer in enumerate(self.moe_layers):
72104
self.log2phy_map_per_layer[local_idx] = layer.get_log2phy_map()
73105

74106
def init_buffer_tensor(self, num_buffer_tensor):
75-
for buffer_id in range(num_buffer_tensor):
76-
for name in self.expert_weight_names:
77-
expert_tensor = self.param_dict[f"0.{name}"][0]
78-
buffer_tensor = torch.empty_like(expert_tensor)
79-
self.buffer_tensor_list[buffer_id].append(buffer_tensor)
107+
buffer_tensor_shapes: dict[Any, list[torch.Size]] = dict()
108+
for local_idx, _ in enumerate(self.moe_layers):
109+
expert_weight_key = self.expert_weight_key_per_layer[local_idx]
110+
expert_weight_names = EPLB_EXPERT_WEIGHT_NAMES[expert_weight_key]
111+
expert_tensors = [self.param_dict[f"{local_idx}.{name}"][0] for name in expert_weight_names]
112+
expert_tensor_shapes = [tensor.shape for tensor in expert_tensors]
113+
if expert_weight_key in self.buffer_tensor_list:
114+
assert expert_tensor_shapes == buffer_tensor_shapes[expert_weight_key], (
115+
f"EPLB expert weight shapes mismatch for {expert_weight_key}: "
116+
f"expected {buffer_tensor_shapes[expert_weight_key]}, got {expert_tensor_shapes}"
117+
)
118+
continue
119+
buffer_tensor_shapes[expert_weight_key] = expert_tensor_shapes
120+
self.buffer_tensor_list[expert_weight_key] = [[] for _ in range(num_buffer_tensor)]
121+
for buffer_id in range(num_buffer_tensor):
122+
for expert_tensor in expert_tensors:
123+
buffer_tensor = torch.empty_like(expert_tensor)
124+
self.buffer_tensor_list[expert_weight_key][buffer_id].append(buffer_tensor)
80125

81126
def init_expert_param_per_layer(self):
82127
self.param_dict = dict()
83128

84-
first_layer = self.moe_layers[0]
85-
86-
if self.model.quant_config is not None:
87-
quant_type = first_layer.quant_type
88-
if quant_type == QuantType.W8A8:
89-
self.expert_weight_names = [
90-
"w13_weight_list",
91-
"w2_weight_list",
92-
"w13_weight_scale_fp32_list",
93-
"w2_weight_scale_list",
94-
]
95-
if get_ascend_config().enable_fused_mc2 == 1:
96-
self.expert_weight_names.append("fused_w1_scale_list")
97-
self.expert_weight_names.append("fused_w2_scale_list")
98-
99-
elif quant_type == QuantType.W4A8:
100-
if get_ascend_config().enable_fused_mc2 != 1:
101-
raise ValueError("EPLB not support W4A8 with fused MC2 disabled")
102-
self.expert_weight_names = [
103-
"w13_weight_list",
104-
"w2_weight_list",
105-
"w13_weight_scale_list",
106-
"w2_weight_scale_list",
107-
"w13_scale_bias_list",
108-
"w2_scale_bias_list",
109-
]
110-
111-
elif quant_type in (QuantType.MXFP4, QuantType.MXFP8):
112-
self.expert_weight_names = [
113-
"w13_weight",
114-
"w2_weight",
115-
"w13_weight_scale",
116-
"w2_weight_scale",
117-
]
118-
else:
119-
raise ValueError(f"EPLB not support {quant_type}")
120-
else:
121-
self.expert_weight_names = ["w13_weight", "w2_weight"]
122-
123129
for local_idx, layer in enumerate(self.moe_layers):
130+
quant_type = QuantType.NONE if self.model.quant_config is None else layer.quant_type
131+
expert_weight_key = (quant_type, get_ascend_config().enable_fused_mc2 == 1)
132+
if expert_weight_key not in EPLB_EXPERT_WEIGHT_NAMES:
133+
raise ValueError(f"EPLB not support {quant_type} with fused MC2 {expert_weight_key[1]}")
134+
expert_weight_names = EPLB_EXPERT_WEIGHT_NAMES[expert_weight_key]
135+
self.expert_weight_key_per_layer[local_idx] = expert_weight_key
124136
self.expert_param_per_layer[local_idx] = list()
125-
for name in self.expert_weight_names:
137+
for name in expert_weight_names:
126138
param_key = f"{local_idx}.{name}"
127139
self.param_dict[param_key] = getattr(layer, name)
128140
for local_expert_id in range(self.num_local_experts):
129141
per_expert_param = list()
130-
for name in self.expert_weight_names:
142+
for name in expert_weight_names:
131143
per_expert_param.append(self.param_dict[f"{local_idx}.{name}"][local_expert_id])
132144
self.expert_param_per_layer[local_idx].append(per_expert_param)
133145

@@ -168,8 +180,10 @@ def do_update_expert_map(self, layer_id, updated_expert_map):
168180
self.expert_map_per_layer_cpu[layer_id].copy_(updated_expert_map)
169181

170182
def do_update_expert_weight(self, layer_id, local_expert_to_replace, buffer_tensor_id):
183+
expert_weight_key = self.expert_weight_key_per_layer[layer_id]
171184
for expert_tensor, buffer_tensor in zip(
172-
self.expert_param_per_layer[layer_id][local_expert_to_replace], self.buffer_tensor_list[buffer_tensor_id]
185+
self.expert_param_per_layer[layer_id][local_expert_to_replace],
186+
self.buffer_tensor_list[expert_weight_key][buffer_tensor_id],
173187
):
174188
expert_tensor.copy_(buffer_tensor)
175189
logger.debug("Expert tensor shape is :%s", expert_tensor.shape)

‎vllm_ascend/eplb/core/eplb_device_transfer_loader.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -67,7 +67,8 @@ def generate_expert_d2d_transfer_task(self, expert_send_info, expert_recv_info,
6767

6868
for buffer_tensor_id, recv_info in enumerate(expert_recv_info):
6969
recv_rank, global_expert_id_to_recv = recv_info
70-
for buffer_tensor in self.eplb_adaptor.buffer_tensor_list[buffer_tensor_id]:
70+
expert_weight_key = self.eplb_adaptor.expert_weight_key_per_layer[layer_id]
71+
for buffer_tensor in self.eplb_adaptor.buffer_tensor_list[expert_weight_key][buffer_tensor_id]:
7172
self.comm_op_list.append(
7273
dist.P2POp(
7374
dist.irecv, buffer_tensor, self.comm_group.ranks[recv_rank], group=self.comm_group.device_group

‎vllm_ascend/spec_decode/llm_base_proposer.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -966,6 +966,9 @@ def _propose(
966966
aclgraph_runtime_mode=aclgraph_runtime_mode,
967967
is_draft_model=True,
968968
draft_attn_metadatas=multi_steps_attn_metadata,
969+
eplb_heat_collection_status=(
970+
self.runner.eplb_heat_collection_status if self.runner.dynamic_eplb else False
971+
),
969972
):
970973
# Reset MOE layer index for forward pass
971974
forward_context = get_forward_context()

‎vllm_ascend/worker/model_runner_v1.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2326,7 +2326,7 @@ def execute_model(
23262326
self.kv_connector_output = kv_connector_output
23272327
self._finalize_dump_data()
23282328
if self.dynamic_eplb:
2329-
self.eplb_updator.forward_end()
2329+
self.eplb_updator.forward_end(self.eplb_heat_collection_status)
23302330
return hidden_states
23312331
if self.is_pooling_model:
23322332
# Return the pooling output.

0 commit comments

Comments
 (0)