Skip to content

[Bugfix]Fix eplb enable when using mtp float weights. #4571

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

Merged
MengqingCao merged 5 commits into from
Dec 2, 2025
Merged

[Bugfix]Fix eplb enable when using mtp float weights. #4571

Changes from 1 commit
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
fc03cd9
f eplb enable when using mtp float weights.
offline0806 Nov 29, 2025
36bb49c
[ci]Fix pre lint.
offline0806 Nov 29, 2025
18b48c5
Update vllm_ascend/ops/fused_moe/fused_moe.py
offline893 Nov 29, 2025
bf1d3da
Merge remote-tracking branch 'upstream/main' into main-1129
offline0806 Dec 1, 2025
996dcaf
Merge remote-tracking branch 'upstream_gitee/main-1129' into main-1129
offline0806 Dec 1, 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
12 changes: 8 additions & 4 deletions vllm_ascend/ops/fused_moe/fused_moe.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -183,6 +183,9 @@ def __init__(self, *args, **kwargs):
# init moe.
self.local_num_experts, self.expert_map, _ = determine_expert_map(
self.ep_size, self.ep_rank, self.global_num_experts)
#Todo init_eplb_enable is a flag to judge whether enable eplb in drafter or not enable,
# this flag will be remove when eplb supporting mtp and fload weights.
init_eplb_enable = False
# static eplb initializing with expert_map_path
if self.expert_map_path and os.path.exists(
self.expert_map_path) and os.access(self.expert_map_path,
Expand All @@ -199,6 +202,7 @@ def __init__(self, *args, **kwargs):
self.moe_instance_id, self.ep_rank))
self.log2phy = self.expert_load_balancer.get_rank_log2phy_map(
self.moe_instance_id, self.ep_rank).npu()
init_eplb_enable = True
except Exception as e:
logger.warning(
f"Init expert map of mtp/eagle when using sample.{e}")
Expand All @@ -224,10 +228,10 @@ def __init__(self, *args, **kwargs):
self.moe_load = torch.zeros(local_num_experts,
dtype=torch.int64).npu()

eplb_enable = self.dynamic_eplb or (self.expert_map_path is not None)
if eplb_enable and (not hasattr(self.quant_method, "quant_method") or
not isinstance(self.quant_method.quant_method,
AscendW8A8DynamicFusedMoEMethod)):
if init_eplb_enable and (
not hasattr(self.quant_method, "quant_method") or
not isinstance(self.quant_method.quant_method,
AscendW8A8DynamicFusedMoEMethod)):
raise ValueError("Eplb supports only w8a8_dynamic quantization.")

self.moe_config.num_experts = self.global_num_experts
Expand Down
Loading