[Bugfix] Fix DeepSeek ForwardFlags across custom op boundary (#30987)

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)