[Fix] Keep mHC context out of non-V4 compiled MoE forwards (#40353)
This commit is contained in:
@@ -608,6 +608,7 @@ class DeepseekV2MoE(nn.Module):
|
|||||||
self.alt_stream = alt_stream
|
self.alt_stream = alt_stream
|
||||||
self.routed_quant_stream = routed_quant_stream
|
self.routed_quant_stream = routed_quant_stream
|
||||||
self.is_nextn = is_nextn
|
self.is_nextn = is_nextn
|
||||||
|
self.is_deepseek_v4 = is_deepseek_v4
|
||||||
self._fuse_finalize_all_reduce = (
|
self._fuse_finalize_all_reduce = (
|
||||||
is_deepseek_v4
|
is_deepseek_v4
|
||||||
and getattr(config, "hc_pre_from_prev_sublayer", False)
|
and getattr(config, "hc_pre_from_prev_sublayer", False)
|
||||||
@@ -1185,8 +1186,10 @@ class DeepseekV2MoE(nn.Module):
|
|||||||
)
|
)
|
||||||
|
|
||||||
if not all_reduce_done:
|
if not all_reduce_done:
|
||||||
if self.tp_size > 1 and not should_skip_post_experts_all_reduce(
|
if (
|
||||||
is_tp_path=True
|
self.is_deepseek_v4
|
||||||
|
and self.tp_size > 1
|
||||||
|
and not should_skip_post_experts_all_reduce(is_tp_path=True)
|
||||||
):
|
):
|
||||||
from sglang.srt.layers.moe.mhc_post_fusion import (
|
from sglang.srt.layers.moe.mhc_post_fusion import (
|
||||||
current_mhc_post_fusion,
|
current_mhc_post_fusion,
|
||||||
@@ -1341,8 +1344,10 @@ class DeepseekV2MoE(nn.Module):
|
|||||||
self.routed_scaling_factor,
|
self.routed_scaling_factor,
|
||||||
)
|
)
|
||||||
|
|
||||||
if self.tp_size > 1 and not should_skip_post_experts_all_reduce(
|
if (
|
||||||
is_tp_path=True,
|
self.is_deepseek_v4
|
||||||
|
and self.tp_size > 1
|
||||||
|
and not should_skip_post_experts_all_reduce(is_tp_path=True)
|
||||||
):
|
):
|
||||||
from sglang.srt.layers.moe.mhc_post_fusion import current_mhc_post_fusion
|
from sglang.srt.layers.moe.mhc_post_fusion import current_mhc_post_fusion
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user