Allow piecewise CUDA graph with speculative decoding (#22128)

Co-authored-by: luhongyu.4869 <luhongyu.4869@bytedance.com>
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

python/sglang/srt/model_executor/model_runner.py

@@ -2600,6 +2600,10 @@ class ModelRunner(ModelRunnerKVCacheMixin):
             )
             return
 
+        # Draft models use decode CUDA graphs, not PCG
+        if self.is_draft_worker:
+            return
+
         # Disable piecewise CUDA graph for non-language models
         if not hasattr(self.model, "model"):
             logger.warning(

python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py

@@ -422,6 +422,16 @@ class PiecewiseCudaGraphRunner:
         # TODO(yuwei): fix it
         if forward_batch.input_embeds is not None:
             return False
+        # PCG graphs are captured with ForwardMode.EXTEND and spec_info=None.
+        # TARGET_VERIFY has different spec_info and capture_hidden_mode,
+        # so it must not use PCG-captured graphs.
+        if forward_batch.forward_mode.is_target_verify():
+            return False
+        # PCG graphs are captured with the runner's capture_hidden_mode.
+        # If the batch needs a different mode (e.g. FULL for speculative
+        # decoding), PCG replay would return wrong/missing hidden_states.
+        if forward_batch.capture_hidden_mode != self.capture_hidden_mode:
+            return False
         # Disable for token embedding overrides (dynamic per-request)
         if forward_batch.replace_embeds is not None:
             return False

python/sglang/srt/server_args.py

@@ -1153,56 +1153,53 @@ class ServerArgs:
         # 1. Disable Model Arch
         if self.get_model_config().is_piecewise_cuda_graph_disabled_model:
             self.disable_piecewise_cuda_graph = True
-        # 2. Speculative decoding
+        # 2. DP attention
-        if self.speculative_algorithm is not None:
-            self.disable_piecewise_cuda_graph = True
-        # 3. DP attention
         if self.enable_dp_attention:
             self.disable_piecewise_cuda_graph = True
-        # 4. Torch compile
+        # 3. Torch compile
         if self.enable_torch_compile:
             self.disable_piecewise_cuda_graph = True
-        # 5. Pipeline parallelism
+        # 4. Pipeline parallelism
         if self.pp_size > 1:
             self.disable_piecewise_cuda_graph = True
-        # 6. Non-CUDA hardware (AMD, NPU, CPU, MPS, XPU, etc.)
+        # 5. Non-CUDA hardware (AMD, NPU, CPU, MPS, XPU, etc.)
         if is_hip() or is_npu() or is_cpu() or is_mps() or is_xpu():
             self.disable_piecewise_cuda_graph = True
-        # 7. MoE A2A backend
+        # 6. MoE A2A backend
         if self.moe_a2a_backend != "none":
             self.disable_piecewise_cuda_graph = True
-        # 8. LoRA
+        # 7. LoRA
         if self.lora_paths or self.enable_lora:
             self.disable_piecewise_cuda_graph = True
-        # 9. Multimodal / VLM models
+        # 8. Multimodal / VLM models
         if self.get_model_config().is_multimodal:
             self.disable_piecewise_cuda_graph = True
-        # 10. GGUF quantized models (custom dequant ops unsupported by torch.compile)
+        # 9. GGUF quantized models (custom dequant ops unsupported by torch.compile)
         if (
             self.load_format == "gguf"
             or self.quantization == "gguf"
             or check_gguf_file(self.model_path)
         ):
             self.disable_piecewise_cuda_graph = True
-        # 11. DLLM (diffusion LLM) models (context manager in forward breaks dynamo)
+        # 10. DLLM (diffusion LLM) models (context manager in forward breaks dynamo)
         if self.dllm_algorithm is not None:
             self.disable_piecewise_cuda_graph = True
-        # 12. CPU offload (breaks dynamo)
+        # 11. CPU offload (breaks dynamo)
         if self.cpu_offload_gb > 0 or self.enable_hierarchical_cache:
             self.disable_piecewise_cuda_graph = True
-        # 13. Deterministic inference
+        # 12. Deterministic inference
         if self.enable_deterministic_inference:
             self.disable_piecewise_cuda_graph = True
-        # 14. PD disaggregation
+        # 13. PD disaggregation
         if self.disaggregation_mode != "null":
             self.disable_piecewise_cuda_graph = True
-        # 15. Symmetric memory (torch.cuda.use_mem_pool is untraceable by dynamo)
+        # 14. Symmetric memory (torch.cuda.use_mem_pool is untraceable by dynamo)
         if self.enable_symm_mem:
             self.disable_piecewise_cuda_graph = True
-        # 16. Expert distribution recorder
+        # 15. Expert distribution recorder
         if self.enable_eplb or self.expert_distribution_recorder_mode is not None:
             self.disable_piecewise_cuda_graph = True
-        # 17. Context parallel
+        # 16. Context parallel
         if self.attn_cp_size > 1:
             self.disable_piecewise_cuda_graph = True
         # 18. CUDA Graph debug mode

test/registered/piecewise_cuda_graph/test_pcg_with_speculative_decoding.py

@@ -0,0 +1,243 @@
+"""Test piecewise CUDA graph coexisting with speculative decoding.
+
+PCG handles prefill/extend path while speculative decoding (MTP/EAGLE3/STANDALONE/NGRAM)
+uses decode CUDA graphs. This test verifies they don't interfere with each other.
+"""
+
+import unittest
+from types import SimpleNamespace
+
+import requests
+
+from sglang.srt.utils import kill_process_tree
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.run_eval import run_eval
+from sglang.test.test_utils import (
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    DEFAULT_URL_FOR_TEST,
+    popen_launch_server,
+)
+
+register_cuda_ci(est_time=600, suite="stage-b-test-2-gpu-large")
+
+
+class TestPCGWithMTP(unittest.TestCase):
+    """Test PCG + MTP (NEXTN) on Qwen3.5-35B-A3B with FP8."""
+
+    @classmethod
+    def setUpClass(cls):
+        cls.model = "Qwen/Qwen3.5-35B-A3B"
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        other_args = [
+            "--tp",
+            "2",
+            "--trust-remote-code",
+            "--quantization",
+            "fp8",
+            "--mamba-scheduler-strategy",
+            "extra_buffer",
+            "--enable-piecewise-cuda-graph",
+            "--speculative-algorithm",
+            "NEXTN",
+            "--reasoning-parser",
+            "qwen3",
+        ]
+        cls.process = popen_launch_server(
+            cls.model,
+            cls.base_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH * 3,
+            other_args=other_args,
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        kill_process_tree(cls.process.pid)
+
+    def test_gsm8k(self):
+        args = SimpleNamespace(
+            base_url=self.base_url,
+            model=self.model,
+            eval_name="gsm8k",
+            max_tokens=8192,
+            num_examples=200,
+            num_threads=200,
+            thinking_mode="qwen3",
+        )
+        metrics = run_eval(args)
+        print(metrics)
+        self.assertGreater(metrics["score"], 0.75)
+
+        server_info = requests.get(self.base_url + "/server_info").json()
+        avg_spec_accept_length = server_info["internal_states"][0][
+            "avg_spec_accept_length"
+        ]
+        print(f"{avg_spec_accept_length=}")
+        self.assertGreater(avg_spec_accept_length, 1.5)
+
+
+class TestPCGWithEAGLE3(unittest.TestCase):
+    """Test PCG + EAGLE3 on Qwen3-30B-A3B-Instruct-2507."""
+
+    @classmethod
+    def setUpClass(cls):
+        cls.model = "Qwen/Qwen3-30B-A3B-Instruct-2507"
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        other_args = [
+            "--tp",
+            "2",
+            "--trust-remote-code",
+            "--enforce-piecewise-cuda-graph",
+            "--mem-fraction-static",
+            "0.6",
+            "--speculative-algorithm",
+            "EAGLE3",
+            "--speculative-draft-model-path",
+            "lmsys/SGLang-EAGLE3-Qwen3-30B-A3B-Instruct-2507-SpecForge-Nex",
+            "--speculative-num-steps",
+            "5",
+            "--speculative-eagle-topk",
+            "4",
+            "--speculative-num-draft-tokens",
+            "8",
+        ]
+        cls.process = popen_launch_server(
+            cls.model,
+            cls.base_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH * 3,
+            other_args=other_args,
+            env={"SGLANG_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN": "1"},
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        kill_process_tree(cls.process.pid)
+
+    def test_gsm8k(self):
+        args = SimpleNamespace(
+            base_url=self.base_url,
+            model=self.model,
+            eval_name="gsm8k",
+            max_tokens=512,
+            num_examples=200,
+            num_threads=200,
+        )
+        metrics = run_eval(args)
+        print(metrics)
+        self.assertGreater(metrics["score"], 0.75)
+
+        server_info = requests.get(self.base_url + "/server_info").json()
+        avg_spec_accept_length = server_info["internal_states"][0][
+            "avg_spec_accept_length"
+        ]
+        print(f"{avg_spec_accept_length=}")
+        self.assertGreater(avg_spec_accept_length, 1.5)
+
+
+class TestPCGWithSTANDALONE(unittest.TestCase):
+    """Test PCG + STANDALONE on Llama-3.1-8B-Instruct + Llama-3.2-1B-Instruct."""
+
+    @classmethod
+    def setUpClass(cls):
+        cls.model = "meta-llama/Llama-3.1-8B-Instruct"
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        other_args = [
+            "--trust-remote-code",
+            "--enforce-piecewise-cuda-graph",
+            "--mem-fraction-static",
+            "0.5",
+            "--speculative-algorithm",
+            "STANDALONE",
+            "--speculative-draft-model-path",
+            "meta-llama/Llama-3.2-1B-Instruct",
+            "--speculative-num-steps",
+            "3",
+            "--speculative-eagle-topk",
+            "1",
+            "--speculative-num-draft-tokens",
+            "4",
+        ]
+        cls.process = popen_launch_server(
+            cls.model,
+            cls.base_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH * 2,
+            other_args=other_args,
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        kill_process_tree(cls.process.pid)
+
+    def test_gsm8k(self):
+        args = SimpleNamespace(
+            base_url=self.base_url,
+            model=self.model,
+            eval_name="gsm8k",
+            max_tokens=512,
+            num_examples=200,
+            num_threads=200,
+        )
+        metrics = run_eval(args)
+        print(metrics)
+        self.assertGreater(metrics["score"], 0.50)
+
+        server_info = requests.get(self.base_url + "/server_info").json()
+        avg_spec_accept_length = server_info["internal_states"][0][
+            "avg_spec_accept_length"
+        ]
+        print(f"{avg_spec_accept_length=}")
+        self.assertGreater(avg_spec_accept_length, 1.5)
+
+
+class TestPCGWithNGRAM(unittest.TestCase):
+    """Test PCG + NGRAM on Qwen2.5-Coder-7B-Instruct."""
+
+    @classmethod
+    def setUpClass(cls):
+        cls.model = "Qwen/Qwen2.5-Coder-7B-Instruct"
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        other_args = [
+            "--trust-remote-code",
+            "--enforce-piecewise-cuda-graph",
+            "--speculative-algorithm",
+            "NGRAM",
+            "--speculative-num-draft-tokens",
+            "16",
+            "--cuda-graph-max-bs",
+            "8",
+            "--mem-fraction-static",
+            "0.8",
+        ]
+        cls.process = popen_launch_server(
+            cls.model,
+            cls.base_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH * 2,
+            other_args=other_args,
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        kill_process_tree(cls.process.pid)
+
+    def test_gsm8k(self):
+        args = SimpleNamespace(
+            base_url=self.base_url,
+            model=self.model,
+            eval_name="gsm8k",
+            max_tokens=512,
+            num_examples=200,
+            num_threads=200,
+        )
+        metrics = run_eval(args)
+        print(metrics)
+        self.assertGreater(metrics["score"], 0.70)
+
+        server_info = requests.get(self.base_url + "/server_info").json()
+        avg_spec_accept_length = server_info["internal_states"][0][
+            "avg_spec_accept_length"
+        ]
+        print(f"{avg_spec_accept_length=}")
+        self.assertGreater(avg_spec_accept_length, 1.5)
+
+
+if __name__ == "__main__":
+    unittest.main()