[Bugfix] Fix Nemotron ForwardFlags across custom op boundary (#30968)

python/sglang/srt/models/nemotron_h.py

@@ -550,11 +550,13 @@ class NemotronHMambaDecoderLayer(NemotronHAttnLikeDecoderLayer):
             if is_in_breakable_cuda_graph():
                 output = torch.empty_like(hidden_states)
                 breakable_nemotron_mamba2_with_output(
-                    hidden_states, output, self.layer_id
+                    hidden_states, output, self.layer_id, fuse_mlp_allreduce
                 )
             elif is_in_tc_piecewise_cuda_graph():
                 output = torch.empty_like(hidden_states)
-                nemotron_mamba2_with_output(hidden_states, output, self.layer_id)
+                nemotron_mamba2_with_output(
+                    hidden_states, output, self.layer_id, fuse_mlp_allreduce
+                )
             else:
                 output = self._forward_mamba(hidden_states, forward_batch)
 
@@ -1189,6 +1191,7 @@ def nemotron_mamba2_with_output(
     hidden_states: torch.Tensor,
     output: torch.Tensor,
     layer_id: int,
+    fuse_mlp_allreduce: bool = False,
 ) -> None:
     """Split op for Mamba2 forward in piecewise CUDA graph mode."""
     context = get_tc_piecewise_forward_context()
@@ -1208,7 +1211,12 @@ def nemotron_mamba2_with_output(
     if hidden_states.shape[0] != num_actual_tokens:
         hidden_states = hidden_states[:num_actual_tokens]
 
-    ret = mamba_layer._forward_mamba(hidden_states, forward_batch)
+    # This function is an opaque custom op under torch.compile. The caller's
+    # ForwardFlags scope is Python control-plane state and is no longer active
+    # when the compiled graph invokes this implementation. Carry the scalar
+    # across the graph boundary and republish it for RowParallelLinear.
+    with get_forward().scoped(fuse_mlp_allreduce=fuse_mlp_allreduce):
+        ret = mamba_layer._forward_mamba(hidden_states, forward_batch)
 
     # Copy result back; output may be larger (padded) so only fill actual tokens
     output[:num_actual_tokens].view(ret.shape).copy_(ret)

test/registered/models_e2e/test_nvidia_nemotron_3_nano.py

@@ -27,11 +27,6 @@ class TestNvidiaNemotron3Nano30BFP8(LMEvalMixin, DefaultServerBase):
     other_args = [
         "--tp-size",
         "2",
-        # FlashInfer trtllm allreduce fusion on flashinfer 0.6.14 degrades
-        # gsm8k for this model (strict-match 0.84 -> 0.78); the ground truth
-        # was calibrated with fusion disabled. Keep it off until the fusion
-        # numerics regression is resolved.
-        "--enforce-disable-flashinfer-allreduce-fusion",
     ] + NEMOTRON_3_NANO_THINKING_ARGS