From 25bd83033d098cee4706449f77bc5120d188003b Mon Sep 17 00:00:00 2001 From: Vedant V Jhaveri Date: Wed, 11 Mar 2026 18:16:38 -0700 Subject: [PATCH] Enable Piecewise CUDA Graph for NemotronH Hybrid (Mamba+Attention) Models (#19903) --- .../sglang/srt/model_executor/model_runner.py | 27 ++++-- python/sglang/srt/models/nemotron_h.py | 88 +++++++++++++++---- 2 files changed, 91 insertions(+), 24 deletions(-) diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 451fc56c6..37d578f2b 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -2238,24 +2238,35 @@ class ModelRunner(ModelRunnerKVCacheMixin): self.moe_layers = [] self.moe_fusions = [] for layer in language_model.model.layers: + attn_layer = None if hasattr(layer, "self_attn"): if hasattr(layer.self_attn, "attn"): - self.attention_layers.append(layer.self_attn.attn) + attn_layer = layer.self_attn.attn elif hasattr(layer.self_attn, "attn_mqa"): # For DeepSeek model - self.attention_layers.append(layer.self_attn.attn_mqa) + attn_layer = layer.self_attn.attn_mqa # For hybrid model elif hasattr(layer, "attn"): - self.attention_layers.append(layer.attn) + attn_layer = layer.attn elif hasattr(layer, "linear_attn"): if hasattr(layer.linear_attn, "attn"): - self.attention_layers.append(layer.linear_attn.attn) + attn_layer = layer.linear_attn.attn else: - self.attention_layers.append(layer.linear_attn) + attn_layer = layer.linear_attn # For InternVL model elif hasattr(layer, "attention"): if hasattr(layer.attention, "attn"): - self.attention_layers.append(layer.attention.attn) + attn_layer = layer.attention.attn + # For NemotronH and similar hybrid models using 'mixer' attribute + elif hasattr(layer, "mixer"): + if hasattr(layer.mixer, "attn"): + attn_layer = layer.mixer.attn + elif hasattr(layer, "_forward_mamba"): + # Mamba layer with split op support - store the layer itself + attn_layer = layer + + if attn_layer is not None: + self.attention_layers.append(attn_layer) moe_block = None moe_fusion = None @@ -2270,6 +2281,10 @@ class ModelRunner(ModelRunnerKVCacheMixin): if hasattr(layer, "moe") and hasattr(layer.moe, "experts"): moe_block = layer.moe.experts moe_fusion = layer.moe + # For NemotronH MoE layers using 'mixer' attribute + if hasattr(layer, "mixer") and hasattr(layer.mixer, "experts"): + moe_block = layer.mixer.experts + moe_fusion = layer.mixer self.moe_layers.append(moe_block) self.moe_fusions.append(moe_fusion) diff --git a/python/sglang/srt/models/nemotron_h.py b/python/sglang/srt/models/nemotron_h.py index 77c21f0e5..669eb4fe5 100644 --- a/python/sglang/srt/models/nemotron_h.py +++ b/python/sglang/srt/models/nemotron_h.py @@ -21,6 +21,11 @@ from typing import Optional, Union import torch from torch import nn +from sglang.srt.compilation.compilation_config import register_split_op +from sglang.srt.compilation.piecewise_context_manager import ( + get_forward_context, + is_in_piecewise_cuda_graph, +) from sglang.srt.configs import NemotronHConfig from sglang.srt.configs.nemotron_h import ATTENTION, MAMBA, MLP, MOE from sglang.srt.distributed import ( @@ -69,6 +74,7 @@ from sglang.srt.utils import ( is_cuda, make_layers, ) +from sglang.srt.utils.custom_op import register_custom_op from sglang.utils import logger _is_cuda = is_cuda() @@ -214,7 +220,9 @@ class NemotronHMoE(nn.Module): self, hidden_states: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor | None]: - if _is_cuda: + # torch.compile cannot trace CUDA streams, so use the non-overlapping + # path when inside piecewise CUDA graph compilation. + if _is_cuda and not is_in_piecewise_cuda_graph(): return self._forward_core_shared_routed_overlap(hidden_states) else: return self._forward_core_normal(hidden_states) @@ -391,6 +399,23 @@ class NemotronHMambaDecoderLayer(nn.Module): self.norm = RMSNorm(config.hidden_size, eps=config.layer_norm_epsilon) + def _forward_mamba( + self, hidden_states: torch.Tensor, forward_batch: ForwardBatch + ) -> torch.Tensor: + """Core Mamba forward logic, called directly or via split op.""" + output = torch.empty_like(hidden_states) + attn_backend = forward_batch.attn_backend + assert isinstance(attn_backend, HybridLinearAttnBackend) + assert isinstance(attn_backend.linear_attn_backend, Mamba2AttnBackend) + attn_backend.linear_attn_backend.forward( + mixer=self.mixer, + layer_id=self.layer_id, + hidden_states=hidden_states, + output=output, + use_triton_causal_conv=True, + ) + return output + def forward( self, *, @@ -404,18 +429,13 @@ class NemotronHMambaDecoderLayer(nn.Module): else: hidden_states, residual = self.norm(hidden_states, residual) - output = torch.empty_like(hidden_states) - attn_backend = forward_batch.attn_backend - assert isinstance(attn_backend, HybridLinearAttnBackend) - assert isinstance(attn_backend.linear_attn_backend, Mamba2AttnBackend) - attn_backend.linear_attn_backend.forward( - mixer=self.mixer, - layer_id=self.layer_id, - hidden_states=hidden_states, - output=output, - use_triton_causal_conv=True, # TODO: investigate need of `use_triton_causal_conv` - ) - return output, residual + if is_in_piecewise_cuda_graph(): + output = torch.empty_like(hidden_states) + nemotron_mamba2_with_output(hidden_states, output, self.layer_id) + return output, residual + else: + output = self._forward_mamba(hidden_states, forward_batch) + return output, residual class NemotronHAttention(nn.Module): @@ -526,12 +546,12 @@ class NemotronHAttentionDecoderLayer(nn.Module): Layers = ( - NemotronHAttentionDecoderLayer - | NemotronHMLPDecoderLayer - | NemotronHMambaDecoderLayer - | NemotronHMoEDecoderLayer + NemotronHAttentionDecoderLayer, + NemotronHMLPDecoderLayer, + NemotronHMambaDecoderLayer, + NemotronHMoEDecoderLayer, ) -ALL_DECODER_LAYER_TYPES: dict[str, type[Layers]] = { +ALL_DECODER_LAYER_TYPES: dict[str, type] = { ATTENTION: NemotronHAttentionDecoderLayer, MLP: NemotronHMLPDecoderLayer, MAMBA: NemotronHMambaDecoderLayer, @@ -861,3 +881,35 @@ class NemotronHForCausalLM(nn.Module): EntryClass = [NemotronHForCausalLM] + + +@register_custom_op(mutates_args=["output"]) +@register_split_op() +def nemotron_mamba2_with_output( + hidden_states: torch.Tensor, + output: torch.Tensor, + layer_id: int, +) -> None: + """Split op for Mamba2 forward in piecewise CUDA graph mode.""" + context = get_forward_context() + forward_batch = context.forward_batch + attention_layers = context.attention_layers + mamba_layer = attention_layers[layer_id] + + # In piecewise CUDA graph mode, hidden_states may be padded to the + # captured graph size. Slice to actual token count for Mamba forward. + attn_backend = forward_batch.attn_backend + metadata = attn_backend.linear_attn_backend.forward_metadata + num_actual_tokens = metadata.num_prefill_tokens + ( + metadata.num_decodes * metadata.draft_token_num + if metadata.is_target_verify + else metadata.num_decodes + ) + if hidden_states.shape[0] != num_actual_tokens: + hidden_states = hidden_states[:num_actual_tokens] + + 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) + return