Merge pull request #3 from Clorist33/Bugfix_eplb_4490

Bugfix eplb 4490

Author: Clorist33

vllm_ascend/ascend_config.py

@@ -34,6 +34,7 @@ class AscendConfig:
 
     def __init__(self, vllm_config):
         additional_config = vllm_config.additional_config if vllm_config.additional_config is not None else {}
+        self.mix_placement = additional_config.get("mix_placement",False)
         torchair_graph_config = additional_config.get("torchair_graph_config",
                                                       {})
 
@@ -349,4 +350,4 @@ def check_ascend_config(vllm_config, enforce_eager):
                     logger.warning(
                         "ACL Graph is currently experimental. Please "
                         "raise an issue on https://github.com/vllm-project/vllm-ascend/issues"
-                        " if you encourage any Error")
+                        " if you encourage any Error")

vllm_ascend/eplb/adaptor/vllm_adaptor.py

@@ -28,19 +28,21 @@
 
 class VllmEplbAdaptor(EplbAdaptor):
 
-    def __init__(self, model, **args):
+    def __init__(self, model, mtp_instance=None, num_mtp_layers=0, **args):
         super().__init__(**args)
         self.model = model
         self.rank_id = dist.get_rank()
         self.world_size = dist.get_world_size()
         self.param_dict = dict(self.model.named_parameters())
+        self.mtp_instance = mtp_instance
+        self.num_mtp_layers = num_mtp_layers
         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.num_moe_layers = self.model.config.num_hidden_layers - self.num_dense_layers  # MTP not included
         self.init_redundancy_expert = get_ascend_config(
         ).init_redundancy_expert
 
@@ -64,6 +66,18 @@ def __init__(self, model, **args):
         else:
             self.expert_weight_names = ["w13_weight", "w2_weight"]
 
+        if self.mtp_instance is not None:
+            if any("w13_weight_offset" in name
+                   for name, _ in self.mtp_instance.named_parameters()):
+                self.mtp_expert_weight_names = [
+                    "w13_weight", "w2_weight", "w13_weight_scale",
+                    "w13_weight_offset", "w2_weight_scale", "w2_weight_offset"
+                ]
+            else:
+                self.mtp_expert_weight_names = ["w13_weight", "w2_weight"]
+        else:
+            self.mtp_expert_weight_names = []
+
         self.expert_map_per_layer = dict(
         )  # reference to expert map on device for expert map update
         self.expert_map_per_layer_cpu = dict(
@@ -72,6 +86,12 @@ def __init__(self, model, **args):
             self.expert_map_per_layer[self.num_dense_layers + layer_idx] = \
                 self.model.get_expert_map(self.num_dense_layers + layer_idx)
 
+        # Currently, MTP only support one layer.
+        if self.mtp_instance is not None:
+            for mtp_layer_idx in range(self.num_mtp_layers):
+                self.expert_map_per_layer[self.num_dense_layers + self.num_moe_layers + mtp_layer_idx] = \
+                    self.mtp_instance.model.get_expert_map(self.num_dense_layers + self.num_moe_layers + mtp_layer_idx)
+
         # TODO: here we set number of buffer tensor equal to number of expert in each laryer, which can be improved
         num_buffer_tensor = torch.where(
             self.expert_map_per_layer[self.num_dense_layers] != -1)[0].numel()
@@ -88,6 +108,11 @@ def __init__(self, model, **args):
             self.log2phy_map_per_layer[self.num_dense_layers + layer_idx] = \
                 self.model.get_log2phy_map(self.num_dense_layers + layer_idx)
 
+        if self.mtp_instance is not None:
+            for mtp_layer_idx in range(self.num_mtp_layers):
+                self.log2phy_map_per_layer[self.num_dense_layers + self.num_moe_layers + mtp_layer_idx] = \
+                    self.mtp_instance.model.get_log2phy_map(self.num_dense_layers + self.num_moe_layers + mtp_layer_idx)
+
         self.all_topk_ids = []
 
     def init_buffer_tensor(self, num_buffer_tensor):
@@ -131,12 +156,46 @@ def init_expert_param_per_layer(self):
                                             name][0].data[local_expert_id])
                 self.expert_param_per_layer[layer_idx].append(per_expert_param)
 
+        if self.mtp_instance is not None:
+            mtp_param_dict = dict(self.mtp_instance.named_parameters())
+            for mtp_layer_idx in range(self.num_mtp_layers):
+                self.expert_param_per_layer[self.num_dense_layers +
+                                            self.num_moe_layers +
+                                            mtp_layer_idx] = list()
+            for local_expert_id in range(num_local_expert):
+                for mtp_layer_idx in range(self.num_mtp_layers):
+                    self.expert_param_per_layer[
+                        self.num_dense_layers + self.num_moe_layers +
+                        mtp_layer_idx].append([
+                            mtp_param_dict["model.layers." +
+                                           str(self.num_dense_layers +
+                                               self.num_moe_layers +
+                                               mtp_layer_idx) +
+                                           ".mtp_block.mlp.experts." +
+                                           name].data[local_expert_id]
+                            for name in self.mtp_expert_weight_names
+                        ])
+
     def get_rank_expert_workload(self) -> torch.Tensor:
         self.moe_load = self.model.get_all_moe_loads()
+        if self.mtp_instance is not None:
+            self.moe_load = torch.cat([
+                self.moe_load,
+                self.mtp_instance.model.get_all_moe_loads().to(
+                    device=self.moe_load.device)
+            ],
+                                      dim=0)
         return self.moe_load
 
     def get_init_expert_map(self, num_moe_layers):
         expert_map = self.model.get_all_expert_map(num_moe_layers)
+        if self.mtp_instance is not None:
+            expert_map = torch.cat([
+                expert_map,
+                self.mtp_instance.model.get_all_expert_map().to(
+                    device=expert_map.device)
+            ],
+                                   dim=0)
         if dist.is_initialized():
             world_size = dist.get_world_size()
 
@@ -288,7 +347,9 @@ def determine_expert_map_all(self):
         local_num_experts = self.global_expert_num // self.world_size
 
         expert_map_all = torch.full(
-            (self.num_moe_layers, self.world_size, self.global_expert_num),
+            (self.num_moe_layers if self.mtp_instance is None else
+             (self.num_moe_layers + self.num_mtp_layers), self.world_size,
+             self.global_expert_num),
             -1,
             dtype=torch.int32)
 
@@ -311,6 +372,7 @@ def determine_expert_map_all(self):
 
             local_ids = torch.arange(local_count, dtype=torch.int32)
             expert_map_all[:, r, start:end] = local_ids.unsqueeze(0).expand(
-                self.num_moe_layers, -1)
+                self.num_moe_layers if self.mtp_instance is None else
+                (self.num_moe_layers + self.num_mtp_layers), -1)
 
-        return expert_map_all
+        return expert_map_all

vllm_ascend/eplb/core/eplb_device_transfer_loader.py

@@ -50,10 +50,6 @@ def generate_expert_d2d_transfer_task(self, expert_send_info,
             )
             return
 
-        # If neither send nor receive task is needed for this layer on this rank, return
-        if not (expert_send_info or expert_recv_info):
-            return
-
         self.updated_expert_map = updated_expert_map
 
         self.layer_id = layer_id
@@ -135,4 +131,4 @@ def update_expert_map_and_weight(self, reqs):
         self.state = ExpertWeightUpdateState.WAITING
 
     def load_impl(self, old_expert_table, new_expert_table):
-        raise NotImplementedError
+        raise NotImplementedError

vllm_ascend/eplb/eplb_updator.py

@@ -35,9 +35,11 @@ def __init__(self, ascend_config, loader, eplb_process: EplbProcess,
         self.eplb_process = eplb_process
         self.shared_dict = self.eplb_process.shared_dict
 
-    def set_adaptor(self, adaptor):
+    def set_adaptor(self, adaptor, num_mtp_layers):
         self.adaptor = adaptor
-        self.num_moe_layers = self.adaptor.num_moe_layers
+        self.num_moe_layers = (self.adaptor.num_moe_layers
+                               if self.adaptor.mtp_instance is None else
+                               self.adaptor.num_moe_layers + num_mtp_layers)
         self.global_expert_num = self.adaptor.global_expert_num
 
     def init_eplb(self, expert_map_path, process):
@@ -84,6 +86,8 @@ def update_iteration(self):
                     self.expert_map_record_path)
 
             self.adaptor.model.clear_all_moe_loads()
+            if self.adaptor.mtp_instance is not None:
+                self.adaptor.mtp_instance.model.clear_all_moe_loads()
             if not self.gate_eplb:
                 self.cur_iterations = 0
 
@@ -207,4 +211,4 @@ def shutdown(self):
         if self.process.is_alive():
             self.process.terminate()
             self.process.join()
-            logger.info("[ModelRunner] EPLB process terminated")
+            logger.info("[ModelRunner] EPLB process terminated")

vllm_ascend/eplb/utils.py

@@ -18,45 +18,73 @@
 import types
 
 import torch
+from vllm.model_executor.models.deepseek_mtp import DeepSeekMultiTokenPredictor
 
 
 def get_expert_map(self, layer_id):
-    return self.model.layers[layer_id].mlp.experts.get_map()
+    if not isinstance(self, DeepSeekMultiTokenPredictor):
+        return self.model.layers[layer_id].mlp.experts.get_map()
+    else:
+        return self.layers[str(layer_id)].mtp_block.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 = []
-    num_dense_layers = self.num_dense_layers if hasattr(
-        self, "num_dense_layers") else 0
-    for layer_id in range(num_moe_layers):
-        load_tensor = self.get_expert_map(
-            layer_id + num_dense_layers)  # (num_experts_per_layer,)
-        all_loads.append(load_tensor)
+    if not isinstance(self, DeepSeekMultiTokenPredictor):
+        return self.model.layers[layer_id].mlp.experts.get_log2phy_map()
+    else:
+        return self.layers[str(
+            layer_id)].mtp_block.mlp.experts.get_log2phy_map()
+
+
+def get_all_expert_map(self, num_moe_layers=None):
+    if not isinstance(self, DeepSeekMultiTokenPredictor):
+        all_loads = []
+        num_dense_layers = self.num_dense_layers if hasattr(
+            self, "num_dense_layers") else 0
+        for layer_id in range(num_moe_layers):
+            load_tensor = self.get_expert_map(
+                layer_id + num_dense_layers)  # (num_experts_per_layer,)
+            all_loads.append(load_tensor)
+    else:
+        all_loads = []
+        for layer_id in range(self.mtp_start_layer_idx,
+                              self.mtp_start_layer_idx + self.num_mtp_layers):
+            load_tensor = self.get_expert_map(layer_id)
+            all_loads.append(load_tensor)
 
     return torch.stack(all_loads, dim=0)
 
 
 def get_all_moe_loads(self):
-    num_dense_layers = self.num_dense_layers if hasattr(
-        self, "num_dense_layers") else 0
-    all_moe_loads = torch.stack(
-        [self.model.layers[layer_id + num_dense_layers].mlp.experts.moe_load \
-            for layer_id in range(self.num_moe_layers)],
-        dim=0
-    )
+    if not isinstance(self, DeepSeekMultiTokenPredictor):
+        num_dense_layers = self.num_dense_layers if hasattr(
+            self, "num_dense_layers") else 0
+        all_moe_loads = torch.stack(
+            [self.model.layers[layer_id + num_dense_layers].mlp.experts.moe_load \
+                for layer_id in range(self.num_moe_layers)],
+            dim=0
+        )
+    else:
+        all_moe_loads = torch.stack(
+            [self.layers[str(idx)].mtp_block.mlp.experts.moe_load \
+                for idx in range(self.mtp_start_layer_idx,
+                                    self.mtp_start_layer_idx + self.num_mtp_layers)],
+            dim=0
+        )
     return all_moe_loads
 
 
 def clear_all_moe_loads(self):
-    num_dense_layers = self.num_dense_layers if hasattr(
-        self, "num_dense_layers") else 0
-    for layer_id in range(self.num_moe_layers):
-        self.model.layers[layer_id +
-                          num_dense_layers].mlp.experts.clear_moe_load()
+    if not isinstance(self, DeepSeekMultiTokenPredictor):
+        num_dense_layers = self.num_dense_layers if hasattr(
+            self, "num_dense_layers") else 0
+        for layer_id in range(self.num_moe_layers):
+            self.model.layers[layer_id +
+                              num_dense_layers].mlp.experts.clear_moe_load()
+    else:
+        for layer_id in range(self.mtp_start_layer_idx,
+                              self.mtp_start_layer_idx + self.num_mtp_layers):
+            self.layers[str(layer_id)].mtp_block.mlp.experts.clear_moe_load()
 
 
 def model_register(model, model_config):
@@ -66,12 +94,13 @@ def model_register(model, model_config):
     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 not isinstance(model, DeepSeekMultiTokenPredictor):
+        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" or config.model_type == "deepseek_v3":
-        model.num_dense_layers = config.first_k_dense_replace
-        model.num_moe_layers = config.num_hidden_layers - model.num_dense_layers
-    else:
-        raise NotImplementedError("EPLB is not supported.")
+        if config.model_type == "qwen3_moe":
+            model.num_moe_layers = config.num_hidden_layers
+        elif config.model_type == "deepseek_v2" or config.model_type == "deepseek_v3":
+            model.num_dense_layers = config.first_k_dense_replace
+            model.num_moe_layers = config.num_hidden_layers - model.num_dense_layers
+        else:
+            raise NotImplementedError("EPLB is not supported.")

vllm_ascend/ops/fused_moe/experts_selector.py

@@ -33,6 +33,8 @@ def select_experts(hidden_states: torch.Tensor,
                    routed_scaling_factor=1.0,
                    e_score_correction_bias: Optional[torch.Tensor] = None,
                    indices_type: Optional[torch.dtype] = None,
+                   mix_placement: Optional[bool] = False,
+                   num_logical_experts: int = -1,
                    global_num_experts: int = -1):
     """
     Fused experts with select experts.
@@ -95,6 +97,19 @@ def select_experts(hidden_states: torch.Tensor,
             e_score_correction_bias=e_score_correction_bias,
             global_num_experts=global_num_experts,
         )
+    if mix_placement:
+        pad_shared_expert_ids = torch.full((topk_ids.shape[0], 1),
+                                           num_logical_experts,
+                                           dtype=topk_ids.dtype,
+                                           device=topk_ids.device)
+
+        pad_shared_expert_weights = torch.full((topk_weights.shape[0], 1),
+                                               0.4,
+                                               dtype=topk_weights.dtype,
+                                               device=topk_weights.device)
+        topk_ids = torch.cat([topk_ids, pad_shared_expert_ids], dim=1)
+        topk_weights = torch.cat([topk_weights, pad_shared_expert_weights], 
+                                 dim=1)
     return topk_weights, topk_ids
 
 
@@ -302,4 +317,4 @@ def _native_select_experts(
     topk_ids = topk_ids.to(torch.int32)
     topk_weights = _renormalize_topk_weights(topk_weights, renormalize)
 
-    return topk_weights, topk_ids
+    return topk_weights, topk_ids