[Fix] Keep mHC context out of non-V4 compiled MoE forwards (#40353)

This commit is contained in:
Xiaoyu Zhang
2026-09-19 21:41:51 +08:00
committed by GitHub
parent 9e2298e913
commit 76f9213a41
+9 -4
View File
@@ -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