diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index 32c4c990d..b4b105dc9 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -879,9 +879,12 @@ class DeepseekV2MoE(nn.Module): if not self._enable_a2a_moe: server_args = get_server_args() if self._can_dual_stream_graph(hidden_states, server_args): + fwd = get_forward() return dsv2_flashinfer_moe_dual_stream_graph( hidden_states, self.layer_id, + fwd.fuse_mlp_allreduce, + fwd.mlp_reduce_scatter, ) elif ( self.alt_stream is not None @@ -2984,6 +2987,8 @@ class DeepseekV32ForCausalLM(DeepseekV2ForCausalLM): def dsv2_flashinfer_moe_dual_stream_graph( hidden_states: torch.Tensor, layer_id: int, + fuse_mlp_allreduce: bool, + mlp_reduce_scatter: bool, ) -> torch.Tensor: forward_context = get_tc_piecewise_forward_context() assert forward_context is not None @@ -2991,7 +2996,14 @@ def dsv2_flashinfer_moe_dual_stream_graph( moe_fusion = forward_context.moe_fusions[layer_id] assert moe_fusion is not None - with get_forward().scoped(flashinfer_trtllm_bypass=True): + # Custom-op execution happens outside the caller's Python scope under + # torch.compile. Carry graph-varying control state as scalar operands and + # republish it for the nested MoE/linear consumers. + with get_forward().scoped( + fuse_mlp_allreduce=fuse_mlp_allreduce, + mlp_reduce_scatter=mlp_reduce_scatter, + flashinfer_trtllm_bypass=True, + ): return moe_fusion.forward_normal_dual_stream(hidden_states)