Skip to content

[Feat] Enable EPLB to support MTP layers #4598

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
dsxsteven wants to merge 5 commits into
base: main
Choose a base branch
from
+158 −40
Open

[Feat] Enable EPLB to support MTP layers #4598

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
e92b07b
eplb enable MTP layers
BryanChen408 Dec 1, 2025
323dad5
fix pre-commit
dsxsteven Dec 4, 2025
cac8e3d
fix precommit
dsxsteven Dec 4, 2025
944bd71
fix precommit
BryanChen408 Dec 4, 2025
74efc4c
fix precommit
BryanChen408 Dec 4, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
58 changes: 52 additions & 6 deletions vllm_ascend/eplb/adaptor/vllm_adaptor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,26 +21,29 @@
import torch
import torch.distributed as dist
from vllm.logger import logger
from vllm.config import get_current_vllm_config

from vllm_ascend.ascend_config import get_ascend_config
from vllm_ascend.eplb.adaptor.abstract_adaptor import EplbAdaptor


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

Expand All @@ -64,6 +67,19 @@ def __init__(self, model, **args):
else:
self.expert_weight_names = ["w13_weight", "w2_weight"]

# TODO: init self.mtp_expert_weight_names depending on different model types, only deepseek v3 w8a8 and qwen3-moe is supported here
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(
Expand All @@ -72,6 +88,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()
Expand All @@ -87,6 +109,11 @@ def __init__(self, model, **args):
for layer_idx in range(self.num_moe_layers):
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 = []

Expand Down Expand Up @@ -131,12 +158,29 @@ 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
])
Copy link
Contributor

@gemini-code-assist gemini-code-assist bot Dec 1, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

The construction of parameter names for MTP layers is hardcoded and very specific to the current model structure. This makes the code brittle and difficult to maintain or extend to other MTP models. While the suggestion improves readability and removes a redundant calculation, the core issue of the hardcoded path remains. Consider making this more robust, for example by having the MTP model itself expose a method to get expert parameter names for a given layer.

                    layer_idx = self.num_dense_layers + self.num_moe_layers + mtp_layer_idx
                    self.expert_param_per_layer[layer_idx].append([
                        mtp_param_dict[f"model.layers.{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()

Expand Down Expand Up @@ -288,9 +332,11 @@ 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),
-1,
dtype=torch.int32)
(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)

for r in range(self.world_size):
if r < self.world_size - 1:
Expand All @@ -311,6 +357,6 @@ 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
8 changes: 6 additions & 2 deletions vllm_ascend/eplb/eplb_updator.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -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

Expand Down
94 changes: 61 additions & 33 deletions vllm_ascend/eplb/utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,45 +19,72 @@

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()
Copy link

@Sparkheart Sparkheart Dec 1, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It is suggested here to extract the public get_map() function, obtain the experts in the branch, and then add get_map() again in the return statement.



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()

Copy link
Contributor

@gemini-code-assist gemini-code-assist bot Dec 1, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

The functions get_expert_map, get_log2phy_map, get_all_expert_map, get_all_moe_loads, and clear_all_moe_loads all contain repetitive if not isinstance(self, DeepSeekMultiTokenPredictor): ... else: ... logic. The else blocks also use hardcoded attribute access paths (e.g., .mtp_block.mlp.experts), which is brittle and not easily extensible to other MTP models. This duplicated logic and hardcoding makes the code difficult to maintain.

Consider refactoring this by introducing a helper method that abstracts away the difference in accessing the expert-related modules between the main model and the MTP model. This would centralize the model-specific logic and make the main functions cleaner and more generic.

For example:

def _get_experts_for_layer(model, layer_id):
    if not isinstance(model, DeepSeekMultiTokenPredictor):
        return model.model.layers[layer_id].mlp.experts
    else:
        # This is still model-specific, but at least it's centralized.
        return model.layers[str(layer_id)].mtp_block.mlp.experts

def get_expert_map(self, layer_id):
    experts = _get_experts_for_layer(self, layer_id)
    return experts.get_map()

# ... and so on for other functions



def model_register(model, model_config):
model.get_expert_map = types.MethodType(get_expert_map, model)
Expand All @@ -66,12 +93,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.")
24 changes: 22 additions & 2 deletions vllm_ascend/worker/model_runner_v1.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -62,6 +62,7 @@
from vllm.model_executor.models.interfaces import (SupportsMultiModal,
supports_mrope,
supports_transcription)
from vllm.model_executor.models.deepseek_mtp import DeepSeekMTP
from vllm.model_executor.models.interfaces_base import (
VllmModelForPooling, is_pooling_model, is_text_generation_model)
from vllm.multimodal import MULTIMODAL_REGISTRY
Expand Down Expand Up @@ -294,6 +295,7 @@
'decode_max_num_seqs', 0)
self.max_num_reqs = max(self.scheduler_config.max_num_seqs,
decode_max_num_seqs)
self.mtp_instance = None
if self.pcp_size > 1:
self.model_config.max_model_len += 2 * self.pcp_size * self.max_num_reqs
self.max_num_blocks_per_req = cdiv(self.model_config.max_model_len,
Expand Down Expand Up @@ -3042,6 +3044,8 @@
dummy_compute_logits=dummy_drafter_compute_logits)
if self.in_profile_run and self.dynamic_eplb:
self.model.clear_all_moe_loads()
if self.mtp_instance is not None:
self.drafter.model.model.clear_all_moe_loads()
if not self.in_profile_run and self.dynamic_eplb:
self.eplb_updator.take_update_info_from_eplb_process()
self.eplb_updator.forward_end()
Expand Down Expand Up @@ -3143,7 +3147,7 @@
"lowering `max_num_seqs` or `gpu_memory_utilization` when "
"initializing the engine.") from e
else:
raise e

Check failure on line 3150 in vllm_ascend/worker/model_runner_v1.py

View workflow job for this annotation

GitHub Actions / lint / pre-commit 

"None" has no attribute "model"  [attr-defined]

Check failure on line 3150 in vllm_ascend/worker/model_runner_v1.py

View workflow job for this annotation

GitHub Actions / lint / pre-commit 

"None" has no attribute "model"  [attr-defined]

Check failure on line 3150 in vllm_ascend/worker/model_runner_v1.py

View workflow job for this annotation

GitHub Actions / lint / pre-commit 

"None" has no attribute "model"  [attr-defined]

Check failure on line 3150 in vllm_ascend/worker/model_runner_v1.py

View workflow job for this annotation

GitHub Actions / lint / pre-commit 

"None" has no attribute "model"  [attr-defined]

@torch.inference_mode()
def _dummy_pooler_run(
Expand All @@ -3164,9 +3168,14 @@
def eplb_warmup(self):
if self.dynamic_eplb and not self.is_eplb_warmuped:
self.is_eplb_warmuped = True
self.eplb_adaptor = VllmEplbAdaptor(model=self.model)
num_mtp_layers = self.mtp_instance.model.num_mtp_layers if self.mtp_instance is not None else 0
self.eplb_adaptor = VllmEplbAdaptor(
model=self.model,
mtp_instance=self.mtp_instance,
num_mtp_layers=num_mtp_layers
)
self.eplb_loader.set_adator(self.eplb_adaptor)
self.eplb_updator.set_adaptor(self.eplb_adaptor)
self.eplb_updator.set_adaptor(self.eplb_adaptor, num_mtp_layers)
self.eplb_updator.warm_up_eplb()

def load_model(self) -> None:
Expand All @@ -3174,8 +3183,8 @@

with DeviceMemoryProfiler() as m: # noqa: SIM117
self.model = get_model(vllm_config=self.vllm_config)
if self.dynamic_eplb:

Check failure on line 3186 in vllm_ascend/worker/model_runner_v1.py

View workflow job for this annotation

GitHub Actions / lint / pre-commit 

Incompatible types in assignment (expression has type "Callable[..., Any]", variable has type "None")  [assignment]

Check failure on line 3186 in vllm_ascend/worker/model_runner_v1.py

View workflow job for this annotation

GitHub Actions / lint / pre-commit 

Incompatible types in assignment (expression has type "Callable[..., Any]", variable has type "None")  [assignment]

Check failure on line 3186 in vllm_ascend/worker/model_runner_v1.py

View workflow job for this annotation

GitHub Actions / lint / pre-commit 

Incompatible types in assignment (expression has type "Callable[..., Any]", variable has type "None")  [assignment]
model_register(self.model, self.model_config)

Check failure on line 3187 in vllm_ascend/worker/model_runner_v1.py

View workflow job for this annotation

GitHub Actions / lint / pre-commit 

"None" has no attribute "model"  [attr-defined]

Check failure on line 3187 in vllm_ascend/worker/model_runner_v1.py

View workflow job for this annotation

GitHub Actions / lint / pre-commit 

"None" has no attribute "model"  [attr-defined]

Check failure on line 3187 in vllm_ascend/worker/model_runner_v1.py

View workflow job for this annotation

GitHub Actions / lint / pre-commit 

"None" has no attribute "model"  [attr-defined]
if get_ascend_device_type() == AscendDeviceType._310P:
from vllm.model_executor.layers.linear import (
MergedColumnParallelLinear, QKVParallelLinear,
Expand All @@ -3189,6 +3198,17 @@
if self.drafter:
logger.info("Loading drafter model...")
self.drafter.load_model(self.model)
if self.speculative_config and self.speculative_config.method == 'deepseek_mtp':
assert isinstance(self.drafter, MtpProposer), \
f"drafter type wrong: {type(self.drafter)}"
assert isinstance(self.drafter.model, (DeepSeekMTP, ACLGraphWrapper)), \
f"drafter type wrong: {type(self.drafter)}, only suport DeepSeekMTP or ACLGraphWrapper"
if isinstance(self.drafter.model, DeepSeekMTP):
self.mtp_instance = self.drafter.model
elif isinstance(self.drafter.model, ACLGraphWrapper):
self.mtp_instance = self.drafter.model.unwrap()
model_register(self.mtp_instance.model, self.vllm_config)

if self.drafter.name == SpecDcodeType.EAGLE3:
self.model.set_aux_hidden_state_layers(
self.model.get_eagle3_aux_hidden_state_layers())
Expand Down
Loading