From d64bf6c6ce703389cbeaaa44fe5ee3c699397d0d Mon Sep 17 00:00:00 2001 From: Chen1022 Date: Tue, 25 Nov 2025 21:01:27 +0800 Subject: [PATCH] Support piecewise cuda graph for Qwen3-next (#13081) --- python/sglang/srt/compilation/backend.py | 1 + .../srt/layers/attention/fla/chunk_o.py | 2 +- python/sglang/srt/mem_cache/memory_pool.py | 6 +- .../sglang/srt/model_executor/model_runner.py | 5 ++ python/sglang/srt/models/qwen3_next.py | 63 ++++++++++++++++++- test/srt/models/test_qwen3_next_models.py | 38 +++++++++++ 6 files changed, 112 insertions(+), 3 deletions(-) diff --git a/python/sglang/srt/compilation/backend.py b/python/sglang/srt/compilation/backend.py index abeed1c83..c282a23f4 100644 --- a/python/sglang/srt/compilation/backend.py +++ b/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", ] diff --git a/python/sglang/srt/layers/attention/fla/chunk_o.py b/python/sglang/srt/layers/attention/fla/chunk_o.py index b2ae826f7..21d338d5d 100644 --- a/python/sglang/srt/layers/attention/fla/chunk_o.py +++ b/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) diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 2b4d9f65e..de5855741 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/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): diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 324d5f314..a36311c85 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/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"): diff --git a/python/sglang/srt/models/qwen3_next.py b/python/sglang/srt/models/qwen3_next.py index 913ceb53a..f1f965ccd 100644 --- a/python/sglang/srt/models/qwen3_next.py +++ b/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, +) diff --git a/test/srt/models/test_qwen3_next_models.py b/test/srt/models/test_qwen3_next_models.py index 1656b2f75..c03869746 100644 --- a/test/srt/models/test_qwen3_next_models.py +++ b/test/srt/models/test_qwen3_next_models.py @@ -135,5 +135,43 @@ class TestQwen3NextMTPTopk(CustomTestCase): self.assertGreater(metrics["accuracy"], 0.93) +class TestQwen3NextPiecewiseCudaGraph(CustomTestCase): + + @classmethod + def setUpClass(cls): + cls.model = "Qwen/Qwen3-Next-80B-A3B-Instruct" + cls.base_url = DEFAULT_URL_FOR_TEST + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=[ + "--tp", + "4", + "--enable-piecewise-cuda-graph", + "--piecewise-cuda-graph-compiler", + "eager", + ], + ) + + @classmethod + def tearDownClass(cls): + kill_process_tree(cls.process.pid) + + def test_gsm8k(self): + args = SimpleNamespace( + num_shots=5, + data_path=None, + num_questions=200, + max_new_tokens=512, + parallel=128, + host="http://127.0.0.1", + port=int(self.base_url.split(":")[-1]), + ) + metrics = run_eval(args) + print(f"{metrics=}") + self.assertGreater(metrics["accuracy"], 0.93) + + if __name__ == "__main__": unittest.main()