From a6bc0532c9367d4d23b863c8135c9bd3b8feb44e Mon Sep 17 00:00:00 2001 From: elvischenv <219235043+elvischenv@users.noreply.github.com> Date: Thu, 20 Aug 2026 05:13:40 +0800 Subject: [PATCH] [Fix] Fix Nemotron-H Mamba illegal memory access under DP attention with CUDA graph (#34561) Co-authored-by: Brayden Zhong --- python/sglang/srt/models/nemotron_h.py | 11 ++++++++++- .../test_nvidia_nemotron_3_super_nvfp4.py | 2 -- 2 files changed, 10 insertions(+), 3 deletions(-) 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 = [