[0.9.1]eplb support qwen3-moe (#2000)

### What this PR does / why we need it?
this pr is eplb add support for qwen3-moe

test in qwen3-moe-235b w8a8 tp1dp16ep16 has 3% optimization

Signed-off-by: yangcheng <[email protected]>
Co-authored-by: yangcheng <[email protected]>

Author: shiyuan680

examples/eplb_generate_map.py

@@ -1,9 +1,33 @@
 import argparse
 import json
+import random
 
 import numpy as np
 
 
+def add_unique_number_with_retry(existing_numbers,
+                                 valid_range,
+                                 max_attempts=100):
+    '''
+    generate an unique number not in existing_numbers
+    Args:
+        existing_numbers:
+        valid_range:
+        max_attempts:
+
+    Returns:
+
+    '''
+    existing_set = set(existing_numbers)
+    min_val, max_val = valid_range
+    for _ in range(max_attempts):
+        candidate = random.randint(min_val, max_val)
+        if candidate not in existing_set:
+            return candidate
+
+    raise ValueError('No unique number found')
+
+
 def split_and_insert(n, k, m):
     '''
     n: expert num
@@ -13,13 +37,12 @@ def split_and_insert(n, k, m):
 
     A = np.arange(n)
 
-    B = np.random.choice(n, size=m, replace=False)
-
     groups = np.array_split(A, k)
 
     for j in range(m // k):
         for i in range(k):
-            groups[i] = np.append(groups[i], B[i + j * k])
+            candidate = add_unique_number_with_retry(groups[i], (0, n - 1))
+            groups[i] = np.append(groups[i], candidate)
     return np.concatenate(groups)
 
 

tests/singlecard/ops/test_fused_moe.py

@@ -98,7 +98,7 @@ def test_fused_experts(
     output = fused_experts(a, w1, w2, topk_weights, topk_ids, topk, e_map)
     torch_output = torch_moe(a, w1, w2, topk_weights, topk_ids, topk, e_map)
     # TODO: The native params are: atol=2e-2, rtol=0, maybe related to the nan problem
-    torch.testing.assert_close(output, torch_output, atol=4e-2, rtol=1)
+    torch.testing.assert_close(output[0], torch_output, atol=4e-2, rtol=1)
     torch.npu.empty_cache()
 
 

vllm_ascend/eplb/adaptor/vllm_adaptor.py

@@ -32,15 +32,22 @@ def __init__(self, model, **args):
         self.rank_id = dist.get_rank()
         self.world_size = dist.get_world_size()
         self.param_dict = dict(self.model.named_parameters())
-        self.num_dense_layers = self.model.config.first_k_dense_replace
+        if self.model.config.model_type == "qwen3_moe":
+            self.num_dense_layers = 0
+            self.global_expert_num = self.model.config.num_experts
+        else:
+            self.num_dense_layers = self.model.config.first_k_dense_replace
+            self.global_expert_num = self.model.config.n_routed_experts
         self.num_moe_layers = self.model.config.num_hidden_layers - self.num_dense_layers
-        self.global_expert_num = self.model.config.n_routed_experts
 
-        # TODO: init self.expert_weight_names depending on different model types, only deepseek v3 w8a8 is supported here
-        self.expert_weight_names = [
-            "w13_weight", "w2_weight", "w13_weight_scale", "w13_weight_offset",
-            "w2_weight_scale", "w2_weight_offset"
-        ]
+        # TODO: init self.expert_weight_names depending on different model types, only deepseek v3 w8a8 and qwen3-moe is supported here
+        if self.model.quant_config is not None:
+            self.expert_weight_names = [
+                "w13_weight", "w2_weight", "w13_weight_scale",
+                "w13_weight_offset", "w2_weight_scale", "w2_weight_offset"
+            ]
+        else:
+            self.expert_weight_names = ["w13_weight", "w2_weight"]
 
         self.expert_map_per_layer = dict(
         )  # reference to expert map on device for expert map update
@@ -127,8 +134,12 @@ def get_init_expert_map_from_file(self, num_moe_layers, expert_map_path):
             expert_map_all = self.determine_expert_map_all()
 
         for layer_idx in range(num_moe_layers):
-            self.expert_map_per_layer_cpu[layer_idx + 3] = \
-                expert_map_all[layer_idx][self.rank_id]
+            if self.model.config.model_type == "qwen3_moe":
+                self.expert_map_per_layer_cpu[layer_idx] = \
+                    expert_map_all[layer_idx][self.rank_id]
+            else:
+                self.expert_map_per_layer_cpu[layer_idx + 3] = \
+                    expert_map_all[layer_idx][self.rank_id]
         return expert_map_all
 
     def _expert_file_to_tensor(self, expert_map_path: str):

vllm_ascend/eplb/utils.py

@@ -0,0 +1,68 @@
+#
+# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# This file is a part of the vllm-ascend project.
+#
+import types
+
+import torch
+
+
+def get_expert_map(self, layer_id):
+    return self.model.layers[layer_id].mlp.experts.get_map()
+
+
+def get_log2phy_map(self, layer_id):
+    return self.model.layers[layer_id].mlp.experts.get_log2phy_map()
+
+
+def get_all_expert_map(self, num_moe_layers):
+    all_loads = []
+    for layer_id in range(num_moe_layers):
+        load_tensor = self.get_expert_map(layer_id)  # (num_experts_per_layer,)
+        all_loads.append(load_tensor)
+
+    return torch.stack(all_loads, dim=0)
+
+
+def get_all_moe_loads(self):
+    all_moe_loads = torch.stack(
+        [self.model.layers[layer_id].mlp.experts.moe_load \
+         for layer_id in range(self.num_moe_layers)],
+        dim=0
+    )
+    return all_moe_loads
+
+
+def clear_all_moe_loads(self):
+    for layer_id in range(self.num_moe_layers):
+        self.model.layers[layer_id].mlp.experts.clear_moe_load()
+
+
+def model_register(model, model_config):
+    model.get_expert_map = types.MethodType(get_expert_map, model)
+    model.get_log2phy_map = types.MethodType(get_log2phy_map, model)
+    model.get_all_expert_map = types.MethodType(get_all_expert_map, model)
+    model.get_all_moe_loads = types.MethodType(get_all_moe_loads, model)
+    model.clear_all_moe_loads = types.MethodType(clear_all_moe_loads, model)
+
+    config = model_config.hf_config
+
+    if config.model_type == "qwen3_moe":
+        model.num_moe_layers = config.num_hidden_layers
+    elif config.model_type == "deepseek_v2":
+        num_dense_layers = config.first_k_dense_replace
+        model.num_moe_layers = config.num_hidden_layers - num_dense_layers
+    else:
+        raise NotImplementedError("EPLB is not supported.")

vllm_ascend/models/deepseek_v2.py

@@ -830,8 +830,6 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         quant_config = vllm_config.quant_config
         self.config = config
         self.quant_config = quant_config
-        self.num_dense_layers = self.config.first_k_dense_replace
-        self.num_moe_layers = self.config.num_hidden_layers - self.num_dense_layers
         self.model = CustomDeepseekV2Model(vllm_config=vllm_config,
                                            prefix=maybe_prefix(
                                                prefix, "model"))
@@ -870,34 +868,6 @@ def load_weights(self, weights: Iterable[tuple[str,
 
         return loaded_params
 
-    def get_expert_map(self, layer_id):
-        return self.model.layers[layer_id].mlp.experts.get_map()
-
-    def get_log2phy_map(self, layer_id):
-        return self.model.layers[layer_id].mlp.experts.get_log2phy_map()
-
-    def get_all_expert_map(self, num_moe_layers):
-        all_loads = []
-        for layer_id in range(num_moe_layers):
-            load_tensor = self.get_expert_map(
-                layer_id + self.num_dense_layers)  # (num_experts_per_layer,)
-            all_loads.append(load_tensor)
-
-        return torch.stack(all_loads, dim=0)
-
-    def get_all_moe_loads(self):
-        all_moe_loads = torch.stack(
-            [self.model.layers[layer_id + self.num_dense_layers].mlp.experts.moe_load \
-                for layer_id in range(self.num_moe_layers)],
-            dim=0
-        )
-        return all_moe_loads
-
-    def clear_all_moe_loads(self):
-        for layer_id in range(self.num_moe_layers):
-            self.model.layers[
-                layer_id + self.num_dense_layers].mlp.experts.clear_moe_load()
-
 
 class CustomDeepseekV3ForCausalLM(CustomDeepseekV2ForCausalLM):
     pass

vllm_ascend/ops/fused_moe.py

@@ -141,7 +141,11 @@ def fused_experts_with_mc2(
     is_torchair: bool = False,
     hidden_states_for_share: Optional[Any] = None,
     mc2_mask: Optional[torch.Tensor] = None,
+    log2phy: Optional[torch.Tensor] = None,
+    global_redundant_expert_num: int = 0
 ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+    if log2phy is not None:
+        topk_ids = log2phy[topk_ids]
     quant_mode = 0
     ep_group = get_mc2_group()
     ep_rank_id = ep_group.rank_in_group
@@ -163,7 +167,7 @@ def fused_experts_with_mc2(
 
     enable_dispatch_v2 = hasattr(torch_npu, "npu_moe_distribute_dispatch_v2")
 
-    moe_expert_num = len(expert_map)
+    moe_expert_num = len(expert_map) + global_redundant_expert_num
     kwargs_mc2 = {
         "x": hidden_states,
         "expert_ids": topk_ids,
@@ -349,17 +353,16 @@ def apply_mlp(
 
 # currently expert parallelism implemented with all2all
 # is under-optimized.
-def fused_experts_with_all2all(
-    hidden_states: torch.Tensor,
-    w1: torch.Tensor,
-    w2: torch.Tensor,
-    topk_weights: torch.Tensor,
-    topk_ids: torch.Tensor,
-    top_k: int,
-    expert_map: torch.Tensor = None,
-    ep_group: GroupCoordinator = None,
-    max_num_tokens: Optional[int] = None,
-):
+def fused_experts_with_all2all(hidden_states: torch.Tensor,
+                               w1: torch.Tensor,
+                               w2: torch.Tensor,
+                               topk_weights: torch.Tensor,
+                               topk_ids: torch.Tensor,
+                               top_k: int,
+                               expert_map: torch.Tensor = None,
+                               ep_group: GroupCoordinator = None,
+                               max_num_tokens: Optional[int] = None,
+                               global_redundant_expert_num: int = 0):
     original_shape = hidden_states.shape
     if len(original_shape) == 3:
         hidden_states = hidden_states.view(-1, hidden_states.shape[-1])
@@ -369,7 +372,7 @@ def fused_experts_with_all2all(
     device = hidden_states.device
 
     if expert_map is not None:
-        global_num_experts = len(expert_map)
+        global_num_experts = len(expert_map) + global_redundant_expert_num
         local_num_experts = global_num_experts // ep_group.world_size
         row_idx_len = num_tokens * top_k
         row_idx = (torch.arange(0,
@@ -639,7 +642,10 @@ def fused_experts_with_all2allv(
     hidden_states: torch.Tensor,
     w1: torch.Tensor,
     w2: torch.Tensor,
+    log2phy: Optional[torch.Tensor] = None,
 ):
+    if log2phy is not None:
+        routing_map = log2phy[routing_map]
     # Enable moe alltoallv, it's a balanced policy for precision and efficiency.
     (share_experts_output, dispatched_input,
      tokens_per_expert) = (token_dispatcher.token_permutation(
@@ -824,8 +830,8 @@ def fused_experts(
             expanded_src_to_dst_row=expanded_row_idx,
             export_for_source_row=topk_ids,
         )
-
-    return final_hidden_states
+    group_list_type = 0
+    return final_hidden_states, expert_tokens, group_list_type
 
 
 def native_grouped_topk(
@@ -1015,6 +1021,8 @@ def apply(
         enable_force_load_balance: bool = False,
         hidden_states_for_share: Optional[Any] = None,
         shared_experts: Optional[Any] = None,
+        log2phy: Optional[Any] = None,
+        global_redundant_expert_num: int = 0,
         **kwargs,
     ) -> torch.Tensor:
 
@@ -1071,6 +1079,8 @@ def apply(
                 is_torchair=self.torchair_graph_enabled,
                 hidden_states_for_share=hidden_states_for_share,
                 mc2_mask=mc2_mask,
+                log2phy=log2phy,
+                global_redundant_expert_num=global_redundant_expert_num,
             )
         elif fused_moe_state == FusedMoEState.AllGather:
             max_num_tokens = self.max_num_batched_tokens if self.use_aclgraph else None
@@ -1105,18 +1115,20 @@ def apply(
                 hidden_states=x,
                 w1=layer.w13_weight,
                 w2=layer.w2_weight,
-            )
+                log2phy=log2phy)
         else:
             max_num_tokens = self.max_num_batched_tokens if self.use_aclgraph else None
-            return fused_experts_with_all2all(hidden_states=x,
-                                              w1=layer.w13_weight,
-                                              w2=layer.w2_weight,
-                                              topk_weights=topk_weights,
-                                              topk_ids=topk_ids,
-                                              top_k=top_k,
-                                              expert_map=expert_map,
-                                              ep_group=get_ep_group(),
-                                              max_num_tokens=max_num_tokens)
+            return fused_experts_with_all2all(
+                hidden_states=x,
+                w1=layer.w13_weight,
+                w2=layer.w2_weight,
+                topk_weights=topk_weights,
+                topk_ids=topk_ids,
+                top_k=top_k,
+                expert_map=expert_map,
+                ep_group=get_ep_group(),
+                max_num_tokens=max_num_tokens,
+                global_redundant_expert_num=global_redundant_expert_num)
 
 
 class AscendFusedMoE(FusedMoE):
@@ -1273,6 +1285,10 @@ def __init__(
         if envs_ascend.VLLM_ASCEND_ENABLE_MOE_ALL2ALL_SEQ and isinstance(
                 self.quant_method, AscendUnquantizedFusedMoEMethod):
             self.reduce_results = False
+            if expert_map_path and os.path.exists(expert_map_path):
+                self.global_num_experts = self.global_num_experts + self.global_redundant_expert_num
+                self.local_num_experts = self.global_num_experts // self.ep_size
+
             moe_dispatcher_config = (
                 MoEDispatcherConfig().set_num_moe_experts(
                     self.global_num_experts).set_num_local_experts(

vllm_ascend/worker/model_runner_v1.py

@@ -81,6 +81,7 @@
 from vllm_ascend.distributed.utils import is_lmhead_tp
 from vllm_ascend.eplb.adaptor.vllm_adaptor import VllmEplbAdaptor
 from vllm_ascend.eplb.eplb_updator import EplbUpdator
+from vllm_ascend.eplb.utils import model_register
 from vllm_ascend.multistream.ms_split import compute_split_seq_index
 from vllm_ascend.platform import NPUPlatform
 from vllm_ascend.sample.rejection_sampler import AscendRejectionSampler
@@ -1897,6 +1898,8 @@ def load_model(self) -> None:
 
         with DeviceMemoryProfiler() as m:  # noqa: SIM117
             self.model = get_model(vllm_config=self.vllm_config)
+            if self.dynamic_eplb:
+                model_register(self.model, self.model_config)
             if hasattr(self, "drafter"):
                 logger.info("Loading drafter model...")
                 self.drafter.load_model()