minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 10 Packages Projects Releases Wiki Activity

Allow piecewise CUDA graph with speculative decoding (#22128)

Browse Source 
Co-authored-by: luhongyu.4869 <luhongyu.4869@bytedance.com>
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
narutolhy
2026-04-17 13:39:30 +08:00
committed by GitHub
co-authored by luhongyu.4869 Claude Opus 4.6
parent 91679d935d
commit 5fa0c6a52e
4 changed files with 272 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/model_executor/model_runner.py
+4
View File
@@ -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
+10
View File
@@ -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
+15 -18
View File
@@ -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
test/registered/piecewise_cuda_graph/test_pcg_with_speculative_decoding.py
+243
View File
@@ -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()