From 5fa0c6a52e0a7e3927fea99920087d8c7336b6db Mon Sep 17 00:00:00 2001 From: narutolhy <582909902@qq.com> Date: Thu, 16 Apr 2026 22:39:30 -0700 Subject: [PATCH] Allow piecewise CUDA graph with speculative decoding (#22128) Co-authored-by: luhongyu.4869 Co-authored-by: Claude Opus 4.6 (1M context) --- .../sglang/srt/model_executor/model_runner.py | 4 + .../piecewise_cuda_graph_runner.py | 10 + python/sglang/srt/server_args.py | 33 ++- .../test_pcg_with_speculative_decoding.py | 243 ++++++++++++++++++ 4 files changed, 272 insertions(+), 18 deletions(-) create mode 100644 test/registered/piecewise_cuda_graph/test_pcg_with_speculative_decoding.py diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index abc308302..e224d4da2 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/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( diff --git a/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py b/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py index ded66a366..7c2eebf57 100644 --- a/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py +++ b/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 diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index cb58cfd8b..08692e38e 100644 --- a/python/sglang/srt/server_args.py +++ b/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 - if self.speculative_algorithm is not None: - self.disable_piecewise_cuda_graph = True - # 3. DP attention + # 2. 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 diff --git a/test/registered/piecewise_cuda_graph/test_pcg_with_speculative_decoding.py b/test/registered/piecewise_cuda_graph/test_pcg_with_speculative_decoding.py new file mode 100644 index 000000000..e65f16b3f --- /dev/null +++ b/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()