Support piecewise cuda graph for Qwen3-next (#13081)

python/sglang/srt/compilation/backend.py

@@ -28,6 +28,7 @@ logger = logging.getLogger(__name__)
 SPLIT_OPS = [
     "sglang.unified_attention_with_output",
     "sglang.inplace_all_reduce",
+    "sglang.gdn_with_output",
 ]
 
 

python/sglang/srt/layers/attention/fla/chunk_o.py

@@ -149,7 +149,7 @@ def chunk_fwd_o(
     if scale is None:
         scale = k.shape[-1] ** -0.5
 
-    o = torch.empty_like(v)
+    o = torch.zeros_like(v)
 
     def grid(meta):
         return (triton.cdiv(V, meta["BV"]), NT, B * H)

python/sglang/srt/mem_cache/memory_pool.py

@@ -15,6 +15,7 @@ limitations under the License.
 
 from __future__ import annotations
 
+import dataclasses
 from dataclasses import dataclass
 
 from sglang.srt.configs.mamba_utils import BaseLinearStateParams
@@ -137,7 +138,10 @@ class MambaPool:
             return type(self)(**kwargs)
 
         def mem_usage_bytes(self):
-            return sum(get_tensor_size_bytes(t) for t in vars(self).values())
+            return sum(
+                get_tensor_size_bytes(getattr(self, f.name))
+                for f in dataclasses.fields(self)
+            )
 
     @dataclass(frozen=True, kw_only=True)
     class SpeculativeState(State):

python/sglang/srt/model_executor/model_runner.py

@@ -363,6 +363,11 @@ class ModelRunner:
                     elif hasattr(layer.self_attn, "attn_mqa"):
                         # For DeepSeek model
                         self.attention_layers.append(layer.self_attn.attn_mqa)
+                # For hybrid model
+                elif hasattr(layer, "attn"):
+                    self.attention_layers.append(layer.attn)
+                elif hasattr(layer, "linear_attn"):
+                    self.attention_layers.append(layer.linear_attn)
                 # For InternVL model
                 elif hasattr(layer, "attention"):
                     if hasattr(layer.attention, "attn"):

python/sglang/srt/models/qwen3_next.py

@@ -5,6 +5,7 @@ from typing import Any, Iterable, Optional, Set, Tuple
 import torch
 from torch import nn
 
+from sglang.srt.compilation.piecewise_context_manager import get_forward_context
 from sglang.srt.configs.qwen3_next import Qwen3NextConfig
 from sglang.srt.distributed import divide, get_pp_group
 from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder
@@ -53,9 +54,13 @@ logger = logging.getLogger(__name__)
 _is_cuda = is_cuda()
 _is_npu = is_npu()
 
+
 import triton
 import triton.language as tl
 
+from sglang.srt.compilation.piecewise_context_manager import get_forward_context
+from sglang.srt.utils import direct_register_custom_op
+
 
 @triton.jit
 def fused_qkvzba_split_reshape_cat_kernel(
@@ -349,7 +354,11 @@ class Qwen3GatedDeltaNet(nn.Module):
         return query, key, value, z, b, a
 
     def _forward_input_proj(self, hidden_states: torch.Tensor):
-        DUAL_STREAM_TOKEN_THRESHOLD = 1024 if not _is_npu else 0
+        if _is_npu or get_global_server_args().enable_piecewise_cuda_graph:
+            DUAL_STREAM_TOKEN_THRESHOLD = 0
+        else:
+            DUAL_STREAM_TOKEN_THRESHOLD = 1024
+
         seq_len, _ = hidden_states.shape
         if seq_len < DUAL_STREAM_TOKEN_THRESHOLD:
             current_stream = torch.cuda.current_stream()
@@ -367,6 +376,22 @@ class Qwen3GatedDeltaNet(nn.Module):
         self,
         hidden_states: torch.Tensor,
         forward_batch: ForwardBatch,
+    ):
+        output = torch.empty_like(hidden_states)
+        if forward_batch.forward_mode.is_extend() and get_forward_context() is not None:
+            torch.ops.sglang.gdn_with_output(
+                hidden_states,
+                output,
+                self.layer_id,
+            )
+            return output
+        else:
+            return self._forward(hidden_states, forward_batch)
+
+    def _forward(
+        self,
+        hidden_states: torch.Tensor,
+        forward_batch: ForwardBatch,
     ):
         seq_len, _ = hidden_states.shape
         is_cuda_graph = forward_batch.forward_mode.is_cuda_graph()
@@ -1023,3 +1048,39 @@ class Qwen3NextForCausalLM(nn.Module):
 
 
 EntryClass = Qwen3NextForCausalLM
+
+
+def gdn_with_output(
+    hidden_states: torch.Tensor,
+    output: torch.Tensor,
+    layer_id: int,
+) -> None:
+    context = get_forward_context()
+    forward_batch = context.forward_batch
+    attention_layers = context.attention_layers
+    attention_layer = attention_layers[layer_id]
+
+    ret = attention_layer._forward(hidden_states, forward_batch)
+
+    assert (
+        output.numel() == ret.numel()
+    ), f"Output tensor element mismatch: {output.numel()} != {ret.numel()}"
+
+    output.view(ret.shape).copy_(ret)
+    return
+
+
+def gdn_with_output_fake(
+    hidden_states: torch.Tensor,
+    output: torch.Tensor,
+    layer_id: int,
+) -> None:
+    return
+
+
+direct_register_custom_op(
+    op_name="gdn_with_output",
+    op_func=gdn_with_output,
+    mutates_args=["output"],
+    fake_impl=gdn_with_output_fake,
+)