diff --git a/python/sglang/srt/models/nemotron_h.py b/python/sglang/srt/models/nemotron_h.py index c31c8a4f6..2037cdbb2 100644 --- a/python/sglang/srt/models/nemotron_h.py +++ b/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) diff --git a/test/registered/models_e2e/test_nvidia_nemotron_3_nano.py b/test/registered/models_e2e/test_nvidia_nemotron_3_nano.py index 0813eaf9f..87e2f708f 100644 --- a/test/registered/models_e2e/test_nvidia_nemotron_3_nano.py +++ b/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