From ddf7e6b6a47e4052bb861937889c0e12feec5486 Mon Sep 17 00:00:00 2001 From: Vedant Jhaveri Date: Wed, 4 Mar 2026 23:55:32 +0000 Subject: [PATCH 01/11] fix nemotron pcg --- python/sglang/srt/mem_cache/memory_pool.py | 21 +++-- .../sglang/srt/model_executor/model_runner.py | 30 ++++--- python/sglang/srt/models/nemotron_h.py | 79 +++++++++++++++---- 3 files changed, 97 insertions(+), 33 deletions(-) diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 1d917137c68d..056840b3c60a 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -190,13 +190,10 @@ class State: temporal: torch.Tensor def at_layer_idx(self, layer: int): - kwargs = {} - for k, v in vars(self).items(): - if k == "conv" or k == "intermediate_conv_window": - kwargs[k] = [conv[layer] for conv in v] - else: - kwargs[k] = v[layer] - return type(self)(**kwargs) + return MambaPool.State( + conv=[conv[layer] for conv in self.conv], + temporal=self.temporal[layer], + ) def mem_usage_bytes(self): return sum( @@ -209,6 +206,16 @@ class SpeculativeState(State): intermediate_ssm: torch.Tensor intermediate_conv_window: List[torch.Tensor] + def at_layer_idx(self, layer: int): + return MambaPool.SpeculativeState( + conv=[conv[layer] for conv in self.conv], + temporal=self.temporal[layer], + intermediate_ssm=self.intermediate_ssm[layer], + intermediate_conv_window=[ + conv[layer] for conv in self.intermediate_conv_window + ], + ) + def __init__( self, *, diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 058030768cc1..39ff077ede56 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -2178,24 +2178,34 @@ def init_piecewise_cuda_graphs(self): 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 + # attn_layer is None for non-attention layers (e.g. Mamba, MLP-only) + self.attention_layers.append(attn_layer) moe_block = None moe_fusion = None @@ -2213,11 +2223,13 @@ def init_piecewise_cuda_graphs(self): self.moe_layers.append(moe_block) self.moe_fusions.append(moe_fusion) - if len(self.attention_layers) < self.model_config.num_hidden_layers: - # TODO(yuwei): support Non-Standard GQA + num_attn_layers_found = sum( + 1 for layer in self.attention_layers if layer is not None + ) + if num_attn_layers_found == 0: log_info_on_rank0( logger, - "Disable piecewise CUDA graph because some layers do not apply Standard GQA", + "Disable piecewise CUDA graph because no attention layers were found", ) return diff --git a/python/sglang/srt/models/nemotron_h.py b/python/sglang/srt/models/nemotron_h.py index 77c21f0e5504..5f1b632abe7f 100644 --- a/python/sglang/srt/models/nemotron_h.py +++ b/python/sglang/srt/models/nemotron_h.py @@ -21,6 +21,8 @@ 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 from sglang.srt.configs import NemotronHConfig from sglang.srt.configs.nemotron_h import ATTENTION, MAMBA, MLP, MOE from sglang.srt.distributed import ( @@ -63,6 +65,7 @@ ) from sglang.srt.models.utils import WeightsMapper from sglang.srt.server_args import get_global_server_args +from sglang.srt.utils.custom_op import register_custom_op from sglang.srt.utils import ( add_prefix, get_current_device_stream_fast, @@ -391,6 +394,21 @@ def __init__( 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 +422,13 @@ def forward( 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 get_forward_context() is not None: + 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 +539,12 @@ def forward( 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 +874,35 @@ def load_weights( 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 From 16b5e3d4ae30cf0affd6b1aef7c919ccf56fadc0 Mon Sep 17 00:00:00 2001 From: Vedant Jhaveri Date: Thu, 5 Mar 2026 00:31:32 +0000 Subject: [PATCH 02/11] resolver merge conflict --- python/sglang/srt/mem_cache/memory_pool.py | 27 ++++++++++------------ 1 file changed, 12 insertions(+), 15 deletions(-) diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 056840b3c60a..5a064e83f748 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -28,7 +28,7 @@ import dataclasses import logging from contextlib import contextmanager, nullcontext -from dataclasses import dataclass +from dataclasses import dataclass, fields from typing import TYPE_CHECKING, Any, List, Optional, Tuple, Union import numpy as np @@ -190,10 +190,17 @@ class State: temporal: torch.Tensor def at_layer_idx(self, layer: int): - return MambaPool.State( - conv=[conv[layer] for conv in self.conv], - temporal=self.temporal[layer], - ) + kwargs = {} + # Use fields instead of vars to avoid torch.compile graph break + for f in fields(self): + name = f.name + v = getattr(self, name) + if name in ("conv", "intermediate_conv_window"): + kwargs[name] = [conv[layer] for conv in v] + else: + kwargs[name] = v[layer] + + return type(self)(**kwargs) def mem_usage_bytes(self): return sum( @@ -206,16 +213,6 @@ class SpeculativeState(State): intermediate_ssm: torch.Tensor intermediate_conv_window: List[torch.Tensor] - def at_layer_idx(self, layer: int): - return MambaPool.SpeculativeState( - conv=[conv[layer] for conv in self.conv], - temporal=self.temporal[layer], - intermediate_ssm=self.intermediate_ssm[layer], - intermediate_conv_window=[ - conv[layer] for conv in self.intermediate_conv_window - ], - ) - def __init__( self, *, From da7daae3bdf62b676e35f5b1efe629cc0641d301 Mon Sep 17 00:00:00 2001 From: Vedant Jhaveri Date: Thu, 5 Mar 2026 00:39:04 +0000 Subject: [PATCH 03/11] fix rmerge conflict --- python/sglang/srt/mem_cache/memory_pool.py | 13 +++++-------- python/sglang/srt/model_executor/model_runner.py | 7 ++----- 2 files changed, 7 insertions(+), 13 deletions(-) diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 5a064e83f748..282489083b76 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -28,7 +28,7 @@ import dataclasses import logging from contextlib import contextmanager, nullcontext -from dataclasses import dataclass, fields +from dataclasses import dataclass from typing import TYPE_CHECKING, Any, List, Optional, Tuple, Union import numpy as np @@ -191,14 +191,11 @@ class State: def at_layer_idx(self, layer: int): kwargs = {} - # Use fields instead of vars to avoid torch.compile graph break - for f in fields(self): - name = f.name - v = getattr(self, name) - if name in ("conv", "intermediate_conv_window"): - kwargs[name] = [conv[layer] for conv in v] + for k, v in vars(self).items(): + if k == "conv" or k == "intermediate_conv_window": + kwargs[k] = [conv[layer] for conv in v] else: - kwargs[name] = v[layer] + kwargs[k] = v[layer] return type(self)(**kwargs) diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 39ff077ede56..cbbdf504e64b 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -2223,13 +2223,10 @@ def init_piecewise_cuda_graphs(self): self.moe_layers.append(moe_block) self.moe_fusions.append(moe_fusion) - num_attn_layers_found = sum( - 1 for layer in self.attention_layers if layer is not None - ) - if num_attn_layers_found == 0: + if len(self.attention_layers) < self.model_config.num_hidden_layers: log_info_on_rank0( logger, - "Disable piecewise CUDA graph because no attention layers were found", + "Disable piecewise CUDA graph because some layers do not apply Standard GQA", ) return From 02b793e1972108f9bdb101c7b46193e92f52b7c6 Mon Sep 17 00:00:00 2001 From: Vedant Jhaveri Date: Thu, 5 Mar 2026 00:42:30 +0000 Subject: [PATCH 04/11] fix lint --- python/sglang/srt/mem_cache/memory_pool.py | 1 - 1 file changed, 1 deletion(-) diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 282489083b76..1d917137c68d 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -196,7 +196,6 @@ def at_layer_idx(self, layer: int): kwargs[k] = [conv[layer] for conv in v] else: kwargs[k] = v[layer] - return type(self)(**kwargs) def mem_usage_bytes(self): From 2459197c7b907a7236549b5ad5aac01886874050 Mon Sep 17 00:00:00 2001 From: Vedant Jhaveri Date: Thu, 5 Mar 2026 01:20:57 +0000 Subject: [PATCH 05/11] make consistent --- python/sglang/srt/models/nemotron_h.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/sglang/srt/models/nemotron_h.py b/python/sglang/srt/models/nemotron_h.py index 5f1b632abe7f..84de84d29bd2 100644 --- a/python/sglang/srt/models/nemotron_h.py +++ b/python/sglang/srt/models/nemotron_h.py @@ -422,7 +422,7 @@ def forward( else: hidden_states, residual = self.norm(hidden_states, residual) - if get_forward_context() is not None: + if forward_batch.forward_mode.is_extend() and get_forward_context() is not None: output = torch.empty_like(hidden_states) nemotron_mamba2_with_output(hidden_states, output, self.layer_id) return output, residual From 4044f86937e12e54ef548b7ebca12614f1bf6002 Mon Sep 17 00:00:00 2001 From: Vedant Jhaveri Date: Thu, 5 Mar 2026 18:56:06 +0000 Subject: [PATCH 06/11] fix lint --- python/sglang/srt/models/nemotron_h.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/models/nemotron_h.py b/python/sglang/srt/models/nemotron_h.py index 84de84d29bd2..cfe12f78a91c 100644 --- a/python/sglang/srt/models/nemotron_h.py +++ b/python/sglang/srt/models/nemotron_h.py @@ -65,13 +65,13 @@ ) from sglang.srt.models.utils import WeightsMapper from sglang.srt.server_args import get_global_server_args -from sglang.srt.utils.custom_op import register_custom_op from sglang.srt.utils import ( add_prefix, get_current_device_stream_fast, is_cuda, make_layers, ) +from sglang.srt.utils.custom_op import register_custom_op from sglang.utils import logger _is_cuda = is_cuda() @@ -394,7 +394,9 @@ def __init__( self.norm = RMSNorm(config.hidden_size, eps=config.layer_norm_epsilon) - def _forward_mamba(self, hidden_states: torch.Tensor, forward_batch: ForwardBatch) -> torch.Tensor: + 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 From 92e0065076df1805d4915e7495ad1abc2c9a6c67 Mon Sep 17 00:00:00 2001 From: Vedant Jhaveri Date: Thu, 5 Mar 2026 19:07:50 +0000 Subject: [PATCH 07/11] put back comment --- python/sglang/srt/model_executor/model_runner.py | 1 + 1 file changed, 1 insertion(+) diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index f2d62e29fbd1..25a34cde0009 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -2241,6 +2241,7 @@ def init_piecewise_cuda_graphs(self): self.moe_fusions.append(moe_fusion) if len(self.attention_layers) < self.model_config.num_hidden_layers: + # TODO(yuwei): support Non-Standard GQA log_info_on_rank0( logger, "Disable piecewise CUDA graph because some layers do not apply Standard GQA", From 92fd37603759e1a17d02a8bf8807322006dd39c6 Mon Sep 17 00:00:00 2001 From: Vedant Jhaveri Date: Thu, 5 Mar 2026 22:45:16 +0000 Subject: [PATCH 08/11] use is_in_piecewise_cuda_graph instead of if forward_batch.forward_mode.is_extend() and get_forward_context() is not None --- python/sglang/srt/models/nemotron_h.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/models/nemotron_h.py b/python/sglang/srt/models/nemotron_h.py index cfe12f78a91c..b2110b067f7e 100644 --- a/python/sglang/srt/models/nemotron_h.py +++ b/python/sglang/srt/models/nemotron_h.py @@ -22,7 +22,10 @@ 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 +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 ( @@ -424,7 +427,7 @@ def forward( else: hidden_states, residual = self.norm(hidden_states, residual) - if forward_batch.forward_mode.is_extend() and get_forward_context() is not None: + 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 From 644558313e0ee644f4aa261773489a904ef88f76 Mon Sep 17 00:00:00 2001 From: Vedant Jhaveri Date: Sat, 7 Mar 2026 05:30:37 +0000 Subject: [PATCH 09/11] fix nemotron moe --- python/sglang/srt/model_executor/model_runner.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 25a34cde0009..c3bb616e193c 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -2237,6 +2237,10 @@ def init_piecewise_cuda_graphs(self): 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) From 29092e0783740af870fb84225bb46b3e5ebc6ae0 Mon Sep 17 00:00:00 2001 From: Vedant Jhaveri Date: Mon, 9 Mar 2026 18:49:36 +0000 Subject: [PATCH 10/11] fix failing test --- python/sglang/srt/models/nemotron_h.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/models/nemotron_h.py b/python/sglang/srt/models/nemotron_h.py index b2110b067f7e..669eb4fe51aa 100644 --- a/python/sglang/srt/models/nemotron_h.py +++ b/python/sglang/srt/models/nemotron_h.py @@ -220,7 +220,9 @@ def _forward_core( 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) From 6faece1c1defd361b94b5c3e59f4f2035c645574 Mon Sep 17 00:00:00 2001 From: Vedant Jhaveri Date: Tue, 10 Mar 2026 17:57:33 +0000 Subject: [PATCH 11/11] add guard for pcg --- python/sglang/srt/model_executor/model_runner.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 34e7127358cc..dadbd702e3b7 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -2263,8 +2263,9 @@ def init_piecewise_cuda_graphs(self): elif hasattr(layer, "_forward_mamba"): # Mamba layer with split op support - store the layer itself attn_layer = layer - # attn_layer is None for non-attention layers (e.g. Mamba, MLP-only) - self.attention_layers.append(attn_layer) + + if attn_layer is not None: + self.attention_layers.append(attn_layer) moe_block = None moe_fusion = None