diff --git a/python/sglang/srt/models/nemotron_h.py b/python/sglang/srt/models/nemotron_h.py index c821474cc..7fa2f1d6b 100644 --- a/python/sglang/srt/models/nemotron_h.py +++ b/python/sglang/srt/models/nemotron_h.py @@ -563,7 +563,16 @@ class NemotronHMambaDecoderLayer(NemotronHAttnLikeDecoderLayer): if get_real_num_tokens(hidden_states, forward_batch) == 0: return torch.zeros_like(hidden_states), residual - output = self._forward_mamba(hidden_states, forward_batch) + if is_in_breakable_cuda_graph(): + output = torch.empty_like(hidden_states) + breakable_nemotron_mamba2_with_output( + hidden_states, output, self.layer_id, False + ) + elif is_in_tc_piecewise_cuda_graph(): + output = torch.empty_like(hidden_states) + nemotron_mamba2_with_output(hidden_states, output, self.layer_id, False) + else: + output = self._forward_mamba(hidden_states, forward_batch) return output, residual hidden_states, residual = input_norm_maybe_fuse_allreduce( diff --git a/test/registered/4-gpu-models/test_nvidia_nemotron_3_super_nvfp4.py b/test/registered/4-gpu-models/test_nvidia_nemotron_3_super_nvfp4.py index e67896db3..faca4d38d 100644 --- a/test/registered/4-gpu-models/test_nvidia_nemotron_3_super_nvfp4.py +++ b/test/registered/4-gpu-models/test_nvidia_nemotron_3_super_nvfp4.py @@ -52,8 +52,6 @@ DP_ATTENTION_EP_ARGS = [ "0.93", "--max-prefill-tokens", "8192", - "--cuda-graph-backend-prefill", - "disabled", ] MTP_ARGS = [