[Refactor] Cuda Graph Runner/Backend Refactor (#23906)
Co-authored-by: BBuf <1182563586@qq.com> Co-authored-by: Cheng Wan <54331508+ch-wan@users.noreply.github.com> Co-authored-by: Lianmin Zheng <lianminzheng@gmail.com>
This commit is contained in:
Yuwei An
co-authored by
BBuf
Cheng Wan
Lianmin Zheng
BBuf
Cheng Wan
Lianmin Zheng
parent
56f06278c6
commit
2495c02c2c
@@ -0,0 +1,331 @@
|
||||
"""Tests for the breakable CUDA graph (BCG) runner.
|
||||
|
||||
Two test classes:
|
||||
- TestBreakableCUDAGraphBasic / TestCopyOutput / TestBreakGraphHelper:
|
||||
unit tests for the core capture / replay mechanism (simple tensor ops).
|
||||
- TestBreakableCudaGraph: integration test — spin up Qwen3-8B with
|
||||
--enable-breakable-cuda-graph and check mgsm_en accuracy.
|
||||
"""
|
||||
|
||||
import unittest
|
||||
|
||||
import torch
|
||||
|
||||
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,
|
||||
CustomTestCase,
|
||||
SimpleNamespace,
|
||||
popen_launch_server,
|
||||
)
|
||||
|
||||
# CI Registration — large suite to fit the integration test's server startup.
|
||||
register_cuda_ci(est_time=79, stage="base-b", runner_config="1-gpu-large")
|
||||
|
||||
|
||||
def _skip_if_no_cuda(test_func):
|
||||
return unittest.skipUnless(torch.cuda.is_available(), "CUDA not available")(
|
||||
test_func
|
||||
)
|
||||
|
||||
|
||||
def _skip_if_no_cuda_bindings(test_func):
|
||||
try:
|
||||
from cuda.bindings import runtime as rt # noqa: F401
|
||||
|
||||
return test_func
|
||||
except ImportError:
|
||||
return unittest.skip("cuda-python not installed")(test_func)
|
||||
|
||||
|
||||
class TestBreakableCUDAGraphBasic(CustomTestCase):
|
||||
"""Test basic breakable CUDA graph capture and replay."""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
if not torch.cuda.is_available():
|
||||
raise unittest.SkipTest("CUDA not available")
|
||||
try:
|
||||
from cuda.bindings import runtime # noqa: F401
|
||||
except ImportError:
|
||||
raise unittest.SkipTest("cuda-python not installed")
|
||||
|
||||
from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.breakable_cuda_graph import (
|
||||
BreakableCUDAGraph,
|
||||
BreakableCUDAGraphCapture,
|
||||
eager_on_graph,
|
||||
)
|
||||
|
||||
cls.BreakableCUDAGraph = BreakableCUDAGraph
|
||||
cls.BreakableCUDAGraphCapture = BreakableCUDAGraphCapture
|
||||
cls.eager_on_graph = staticmethod(eager_on_graph)
|
||||
cls.device = torch.device("cuda:0")
|
||||
|
||||
def test_no_break_capture_replay(self):
|
||||
"""Capture and replay without any graph breaks should work like normal CUDA graph."""
|
||||
x = torch.zeros(4, device=self.device)
|
||||
y = torch.zeros(4, device=self.device)
|
||||
|
||||
graph = self.BreakableCUDAGraph()
|
||||
stream = torch.cuda.Stream(self.device)
|
||||
with self.BreakableCUDAGraphCapture(graph, stream=stream):
|
||||
y.copy_(x + 1.0)
|
||||
|
||||
# Replay with new input
|
||||
x.fill_(5.0)
|
||||
graph.replay()
|
||||
torch.cuda.synchronize()
|
||||
self.assertTrue(torch.allclose(y, torch.full((4,), 6.0, device=self.device)))
|
||||
|
||||
def test_single_break(self):
|
||||
"""A single graph break should split capture into two segments."""
|
||||
x = torch.zeros(4, device=self.device)
|
||||
intermediate = torch.zeros(4, device=self.device)
|
||||
y = torch.zeros(4, device=self.device)
|
||||
|
||||
@self.eager_on_graph(enable=True)
|
||||
def eager_op(src):
|
||||
return src * 2.0
|
||||
|
||||
graph = self.BreakableCUDAGraph()
|
||||
stream = torch.cuda.Stream(self.device)
|
||||
with self.BreakableCUDAGraphCapture(graph, stream=stream):
|
||||
intermediate.copy_(x + 1.0)
|
||||
broken = eager_op(intermediate)
|
||||
y.copy_(broken + 3.0)
|
||||
|
||||
# Replay with new input
|
||||
x.fill_(10.0)
|
||||
graph.replay()
|
||||
torch.cuda.synchronize()
|
||||
# x=10 -> intermediate=11 -> eager: 11*2=22 -> y=22+3=25
|
||||
self.assertTrue(torch.allclose(y, torch.full((4,), 25.0, device=self.device)))
|
||||
|
||||
def test_multiple_breaks(self):
|
||||
"""Multiple graph breaks should produce correct chained results."""
|
||||
x = torch.zeros(4, device=self.device)
|
||||
y = torch.zeros(4, device=self.device)
|
||||
|
||||
@self.eager_on_graph(enable=True)
|
||||
def add_one(src):
|
||||
return src + 1.0
|
||||
|
||||
@self.eager_on_graph(enable=True)
|
||||
def double(src):
|
||||
return src * 2.0
|
||||
|
||||
graph = self.BreakableCUDAGraph()
|
||||
stream = torch.cuda.Stream(self.device)
|
||||
with self.BreakableCUDAGraphCapture(graph, stream=stream):
|
||||
t1 = x + 1.0 # graph segment 1
|
||||
t2 = add_one(t1) # break 1: eager
|
||||
t3 = t2 + 1.0 # graph segment 2
|
||||
t4 = double(t3) # break 2: eager
|
||||
y.copy_(t4) # graph segment 3
|
||||
|
||||
# Replay: x=5 -> +1=6 -> add_one=7 -> +1=8 -> double=16
|
||||
x.fill_(5.0)
|
||||
graph.replay()
|
||||
torch.cuda.synchronize()
|
||||
self.assertTrue(torch.allclose(y, torch.full((4,), 16.0, device=self.device)))
|
||||
|
||||
def test_eager_on_graph_disabled(self):
|
||||
"""@eager_on_graph(enable=False) should be a no-op passthrough."""
|
||||
|
||||
@self.eager_on_graph(enable=False)
|
||||
def my_fn(x):
|
||||
return x + 1.0
|
||||
|
||||
# Should just be the original function
|
||||
t = torch.tensor([1.0, 2.0], device=self.device)
|
||||
result = my_fn(t)
|
||||
self.assertTrue(
|
||||
torch.allclose(result, torch.tensor([2.0, 3.0], device=self.device))
|
||||
)
|
||||
|
||||
def test_eager_on_graph_outside_capture(self):
|
||||
"""@eager_on_graph called outside capture should run the function directly."""
|
||||
|
||||
@self.eager_on_graph(enable=True)
|
||||
def my_fn(x):
|
||||
return x + 1.0
|
||||
|
||||
t = torch.tensor([1.0, 2.0], device=self.device)
|
||||
result = my_fn(t)
|
||||
self.assertTrue(
|
||||
torch.allclose(result, torch.tensor([2.0, 3.0], device=self.device))
|
||||
)
|
||||
|
||||
def test_replay_updates_output(self):
|
||||
"""Replay should produce different results when input buffers change."""
|
||||
x = torch.zeros(4, device=self.device)
|
||||
y = torch.zeros(4, device=self.device)
|
||||
|
||||
@self.eager_on_graph(enable=True)
|
||||
def scale(src):
|
||||
return src * 3.0
|
||||
|
||||
graph = self.BreakableCUDAGraph()
|
||||
stream = torch.cuda.Stream(self.device)
|
||||
with self.BreakableCUDAGraphCapture(graph, stream=stream):
|
||||
t = x + 1.0
|
||||
t2 = scale(t)
|
||||
y.copy_(t2)
|
||||
|
||||
# First replay: x=0 -> 0+1=1 -> 1*3=3
|
||||
graph.replay()
|
||||
torch.cuda.synchronize()
|
||||
self.assertTrue(torch.allclose(y, torch.full((4,), 3.0, device=self.device)))
|
||||
|
||||
# Second replay: x=10 -> 10+1=11 -> 11*3=33
|
||||
x.fill_(10.0)
|
||||
graph.replay()
|
||||
torch.cuda.synchronize()
|
||||
self.assertTrue(torch.allclose(y, torch.full((4,), 33.0, device=self.device)))
|
||||
|
||||
|
||||
class TestCopyOutput(CustomTestCase):
|
||||
"""Test the _copy_output helper for structured output writeback."""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
if not torch.cuda.is_available():
|
||||
raise unittest.SkipTest("CUDA not available")
|
||||
try:
|
||||
from cuda.bindings import runtime # noqa: F401
|
||||
except ImportError:
|
||||
raise unittest.SkipTest("cuda-python not installed")
|
||||
|
||||
from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.breakable_cuda_graph import (
|
||||
_copy_output,
|
||||
)
|
||||
|
||||
cls._copy_output = staticmethod(_copy_output)
|
||||
cls.device = torch.device("cuda:0")
|
||||
|
||||
def test_tensor_copy(self):
|
||||
dst = torch.zeros(4, device=self.device)
|
||||
src = torch.ones(4, device=self.device) * 5.0
|
||||
result = self._copy_output(dst, src)
|
||||
self.assertIs(result, dst)
|
||||
self.assertTrue(torch.allclose(dst, src))
|
||||
|
||||
def test_dict_copy(self):
|
||||
dst = {
|
||||
"a": torch.zeros(4, device=self.device),
|
||||
"b": torch.zeros(4, device=self.device),
|
||||
}
|
||||
src = {
|
||||
"a": torch.ones(4, device=self.device),
|
||||
"b": torch.ones(4, device=self.device) * 2.0,
|
||||
}
|
||||
result = self._copy_output(dst, src)
|
||||
self.assertIs(result, dst)
|
||||
self.assertTrue(torch.allclose(dst["a"], torch.ones(4, device=self.device)))
|
||||
self.assertTrue(
|
||||
torch.allclose(dst["b"], torch.ones(4, device=self.device) * 2.0)
|
||||
)
|
||||
|
||||
def test_object_copy(self):
|
||||
class FakeOutput:
|
||||
def __init__(self, t, label):
|
||||
self.tensor = t
|
||||
self.label = label
|
||||
|
||||
dst = FakeOutput(torch.zeros(4, device=self.device), "old")
|
||||
src = FakeOutput(torch.ones(4, device=self.device) * 3.0, "new")
|
||||
result = self._copy_output(dst, src)
|
||||
self.assertIs(result, dst)
|
||||
self.assertTrue(
|
||||
torch.allclose(dst.tensor, torch.ones(4, device=self.device) * 3.0)
|
||||
)
|
||||
self.assertEqual(dst.label, "new")
|
||||
|
||||
def test_non_tensor_fallback(self):
|
||||
result = self._copy_output(42, 99)
|
||||
self.assertEqual(result, 99)
|
||||
|
||||
|
||||
class TestBreakGraphHelper(CustomTestCase):
|
||||
"""Test the break_graph() convenience function."""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
if not torch.cuda.is_available():
|
||||
raise unittest.SkipTest("CUDA not available")
|
||||
try:
|
||||
from cuda.bindings import runtime # noqa: F401
|
||||
except ImportError:
|
||||
raise unittest.SkipTest("cuda-python not installed")
|
||||
|
||||
from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.breakable_cuda_graph import (
|
||||
BreakableCUDAGraph,
|
||||
BreakableCUDAGraphCapture,
|
||||
break_graph,
|
||||
)
|
||||
|
||||
cls.BreakableCUDAGraph = BreakableCUDAGraph
|
||||
cls.BreakableCUDAGraphCapture = BreakableCUDAGraphCapture
|
||||
cls.break_graph = staticmethod(break_graph)
|
||||
cls.device = torch.device("cuda:0")
|
||||
|
||||
def test_break_graph_inserts_segment(self):
|
||||
"""break_graph() should insert a graph break even though it does nothing."""
|
||||
x = torch.zeros(4, device=self.device)
|
||||
y = torch.zeros(4, device=self.device)
|
||||
|
||||
graph = self.BreakableCUDAGraph()
|
||||
stream = torch.cuda.Stream(self.device)
|
||||
with self.BreakableCUDAGraphCapture(graph, stream=stream):
|
||||
t = x + 1.0
|
||||
self.break_graph()
|
||||
y.copy_(t + 2.0)
|
||||
|
||||
x.fill_(10.0)
|
||||
graph.replay()
|
||||
torch.cuda.synchronize()
|
||||
# x=10 -> +1=11 -> break -> +2=13
|
||||
self.assertTrue(torch.allclose(y, torch.full((4,), 13.0, device=self.device)))
|
||||
|
||||
|
||||
class TestBreakableCudaGraph(CustomTestCase):
|
||||
"""Integration: Qwen3-8B with --enable-breakable-cuda-graph on mgsm_en."""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.model = "Qwen/Qwen3-8B"
|
||||
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=[
|
||||
"--cuda-graph-backend-prefill=breakable",
|
||||
],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
kill_process_tree(cls.process.pid)
|
||||
|
||||
def test_gsm8k_accuracy(self):
|
||||
args = SimpleNamespace(
|
||||
base_url=self.base_url,
|
||||
model=self.model,
|
||||
eval_name="mgsm_en",
|
||||
num_examples=1319,
|
||||
num_threads=1024,
|
||||
)
|
||||
|
||||
metrics = run_eval(args)
|
||||
score = metrics["score"]
|
||||
print(f"mgsm_en accuracy with breakable CUDA graph: {score:.3f}")
|
||||
|
||||
self.assertGreaterEqual(score, 0.80)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,71 @@
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
|
||||
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,
|
||||
CustomTestCase,
|
||||
popen_launch_server,
|
||||
)
|
||||
|
||||
register_cuda_ci(est_time=900, stage="base-c", runner_config="4-gpu-b200")
|
||||
|
||||
GLM5_FP4_MODEL = "nvidia/GLM-5-NVFP4"
|
||||
|
||||
|
||||
class TestPCGGlm5Fp4(CustomTestCase):
|
||||
"""PCG prefill on GLM-5-NVFP4 (DSA model, TP=4, B200).
|
||||
|
||||
GLM-5 uses GlmMoeDsaForCausalLM (DSA attention). This test verifies that
|
||||
piecewise CUDA graph works correctly after the DSA indexer was updated to
|
||||
cache k_fp8/k_scale for PCG-compatible prefill.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.model = GLM5_FP4_MODEL
|
||||
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-size",
|
||||
"4",
|
||||
"--trust-remote-code",
|
||||
"--reasoning-parser",
|
||||
"glm45",
|
||||
"--tool-call-parser",
|
||||
"glm47",
|
||||
"--quantization",
|
||||
"modelopt_fp4",
|
||||
"--disable-flashinfer-autotune",
|
||||
"--cuda-graph-backend-prefill=tc_piecewise",
|
||||
"--model-loader-extra-config",
|
||||
'{"enable_multithread_load": true, "num_threads": 64}',
|
||||
],
|
||||
)
|
||||
|
||||
@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",
|
||||
num_examples=200,
|
||||
num_threads=200,
|
||||
max_tokens=4096,
|
||||
)
|
||||
metrics = run_eval(args)
|
||||
print(f"{metrics=}")
|
||||
self.assertGreater(metrics["score"], 0.92)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,45 @@
|
||||
"""Test piecewise CUDA graph coexisting with speculative decoding (EAGLE3).
|
||||
|
||||
PCG handles prefill/extend path while speculative decoding (EAGLE3) uses
|
||||
decode CUDA graphs. This test verifies they don't interfere with each
|
||||
other. MTP / STANDALONE / NGRAM variants moved to the sibling file
|
||||
test_pcg_with_speculative_decoding_extra.py.
|
||||
"""
|
||||
|
||||
import unittest
|
||||
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.server_fixtures.pcg_spec_fixture import PCGSpecBase
|
||||
|
||||
register_cuda_ci(est_time=531, stage="base-b", runner_config="2-gpu-large")
|
||||
|
||||
|
||||
class TestPCGWithEAGLE3(PCGSpecBase, unittest.TestCase):
|
||||
"""PCG + EAGLE3 on Qwen3-30B-A3B-Instruct-2507."""
|
||||
|
||||
model = "Qwen/Qwen3-30B-A3B-Instruct-2507"
|
||||
server_args = [
|
||||
"--tp",
|
||||
"2",
|
||||
"--trust-remote-code",
|
||||
"--cuda-graph-backend-prefill=tc_piecewise",
|
||||
"--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",
|
||||
]
|
||||
timeout_mult = 3
|
||||
server_env = {"SGLANG_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN": "1"}
|
||||
accuracy_threshold = 0.75
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,48 @@
|
||||
"""Test piecewise CUDA graph coexisting with speculative decoding (DFLASH).
|
||||
|
||||
PCG handles prefill/extend path while DFlash needs target aux hidden states
|
||||
from prefill to materialize draft KV cache. This verifies PCG captures that
|
||||
path with the DFlash hidden-state variant enabled.
|
||||
"""
|
||||
|
||||
import unittest
|
||||
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.server_fixtures.pcg_spec_fixture import PCGSpecBase
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_DRAFT_MODEL_DFLASH,
|
||||
DEFAULT_TARGET_MODEL_DFLASH,
|
||||
CustomTestCase,
|
||||
)
|
||||
|
||||
register_cuda_ci(est_time=531, stage="base-b", runner_config="1-gpu-small")
|
||||
|
||||
|
||||
class TestPCGWithDFlash(PCGSpecBase, CustomTestCase):
|
||||
"""PCG + DFLASH on Llama-3.1-8B-Instruct."""
|
||||
|
||||
model = DEFAULT_TARGET_MODEL_DFLASH
|
||||
server_args = [
|
||||
"--trust-remote-code",
|
||||
"--attention-backend",
|
||||
"flashinfer",
|
||||
"--cuda-graph-backend-prefill",
|
||||
"tc_piecewise",
|
||||
"--speculative-algorithm",
|
||||
"DFLASH",
|
||||
"--speculative-draft-model-path",
|
||||
DEFAULT_DRAFT_MODEL_DFLASH,
|
||||
"--page-size",
|
||||
"1",
|
||||
"--max-running-requests",
|
||||
"64",
|
||||
"--cuda-graph-bs-decode",
|
||||
*[str(i) for i in range(1, 65)],
|
||||
]
|
||||
server_env = {"SGLANG_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN": "1"}
|
||||
accuracy_threshold = 0.75
|
||||
speedup_threshold = 2.8
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,80 @@
|
||||
"""Extra: PCG coexistence with non-EAGLE3 speculative decoding variants.
|
||||
|
||||
EAGLE3 stays per-commit in the sibling file
|
||||
test_pcg_with_speculative_decoding.py.
|
||||
"""
|
||||
|
||||
import unittest
|
||||
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.server_fixtures.pcg_spec_fixture import PCGSpecBase
|
||||
|
||||
register_cuda_ci(est_time=531, stage="extra-a", runner_config="2-gpu-large")
|
||||
|
||||
|
||||
class TestPCGWithMTP(PCGSpecBase, unittest.TestCase):
|
||||
"""PCG + MTP (NEXTN) on Qwen3.5-35B-A3B with FP8."""
|
||||
|
||||
model = "Qwen/Qwen3.5-35B-A3B"
|
||||
server_args = [
|
||||
"--tp",
|
||||
"2",
|
||||
"--trust-remote-code",
|
||||
"--quantization",
|
||||
"fp8",
|
||||
"--mamba-scheduler-strategy",
|
||||
"extra_buffer",
|
||||
"--speculative-algorithm",
|
||||
"NEXTN",
|
||||
"--reasoning-parser",
|
||||
"qwen3",
|
||||
]
|
||||
timeout_mult = 3
|
||||
max_tokens = 8192
|
||||
thinking_mode = "qwen3"
|
||||
accuracy_threshold = 0.75
|
||||
|
||||
|
||||
class TestPCGWithSTANDALONE(PCGSpecBase, unittest.TestCase):
|
||||
"""PCG + STANDALONE on Llama-3.1-8B-Instruct + Llama-3.2-1B-Instruct."""
|
||||
|
||||
model = "meta-llama/Llama-3.1-8B-Instruct"
|
||||
server_args = [
|
||||
"--trust-remote-code",
|
||||
"--cuda-graph-backend-prefill=tc_piecewise",
|
||||
"--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",
|
||||
]
|
||||
accuracy_threshold = 0.50
|
||||
|
||||
|
||||
class TestPCGWithNGRAM(PCGSpecBase, unittest.TestCase):
|
||||
"""PCG + NGRAM on Qwen2.5-Coder-7B-Instruct."""
|
||||
|
||||
model = "Qwen/Qwen2.5-Coder-7B-Instruct"
|
||||
server_args = [
|
||||
"--trust-remote-code",
|
||||
"--cuda-graph-backend-prefill=tc_piecewise",
|
||||
"--speculative-algorithm",
|
||||
"NGRAM",
|
||||
"--speculative-num-draft-tokens",
|
||||
"16",
|
||||
"--cuda-graph-max-bs",
|
||||
"8",
|
||||
"--mem-fraction-static",
|
||||
"0.8",
|
||||
]
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,113 @@
|
||||
import unittest
|
||||
|
||||
import torch
|
||||
|
||||
from sglang import Engine
|
||||
from sglang.lang.chat_template import get_chat_template_by_model_path
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
|
||||
from sglang.test.run_eval import run_eval
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_IMAGE_URL,
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
CustomTestCase,
|
||||
SimpleNamespace,
|
||||
popen_launch_server,
|
||||
)
|
||||
|
||||
# CI Registration
|
||||
register_cuda_ci(est_time=180, stage="base-b", runner_config="1-gpu-large")
|
||||
register_amd_ci(est_time=180, suite="stage-b-test-1-gpu-large-amd")
|
||||
|
||||
|
||||
class TestPiecewiseCudaGraphQwen25VL(CustomTestCase):
|
||||
"""Test piecewise CUDA graph with Qwen2.5-VL-7B-Instruct model"""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.model = "Qwen/Qwen2.5-VL-7B-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=[
|
||||
"--cuda-graph-backend-prefill=tc_piecewise",
|
||||
"--disable-radix-cache",
|
||||
],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
kill_process_tree(cls.process.pid)
|
||||
|
||||
def test_gsm8k_accuracy(self):
|
||||
args = SimpleNamespace(
|
||||
base_url=self.base_url,
|
||||
model=self.model,
|
||||
eval_name="gsm8k",
|
||||
num_examples=None,
|
||||
num_threads=1024,
|
||||
)
|
||||
|
||||
metrics = run_eval(args)
|
||||
print(f"GSM8K Accuracy: {metrics['score']:.3f}")
|
||||
|
||||
self.assertGreaterEqual(metrics["score"], 0.80)
|
||||
|
||||
|
||||
class TestPiecewiseCudaGraphQwen25VLEmbedding(CustomTestCase):
|
||||
"""Test piecewise CUDA graph with Qwen2.5-VL-3B-Instruct embedding model"""
|
||||
|
||||
def test_embedding(self):
|
||||
model_path = "Qwen/Qwen2.5-VL-3B-Instruct"
|
||||
chat_template = get_chat_template_by_model_path(model_path)
|
||||
text = f"{chat_template.image_token}What is in this picture? Answer: "
|
||||
|
||||
engine = Engine(
|
||||
model_path=model_path,
|
||||
enable_multimodal=True,
|
||||
is_embedding=True,
|
||||
cuda_graph_backend_prefill="tc_piecewise",
|
||||
)
|
||||
out = engine.encode([text], image_data=[DEFAULT_IMAGE_URL])[0]["embedding"]
|
||||
engine.shutdown()
|
||||
self.assertGreater(len(out), 0)
|
||||
|
||||
engine = Engine(
|
||||
model_path=model_path,
|
||||
enable_multimodal=True,
|
||||
is_embedding=True,
|
||||
cuda_graph_backend_prefill="disabled",
|
||||
)
|
||||
out_without_pcg = engine.encode([text], image_data=[DEFAULT_IMAGE_URL])[0][
|
||||
"embedding"
|
||||
]
|
||||
engine.shutdown()
|
||||
self.assertGreater(len(out_without_pcg), 0)
|
||||
|
||||
t_out = torch.tensor(out)
|
||||
t_out_without_pcg = torch.tensor(out_without_pcg)
|
||||
max_abs_diff = (t_out - t_out_without_pcg).abs().max().item()
|
||||
max_rel_diff = (
|
||||
((t_out - t_out_without_pcg).abs() / (t_out_without_pcg.abs() + 1e-8))
|
||||
.max()
|
||||
.item()
|
||||
)
|
||||
print(
|
||||
f"PCG embedding diff: max_abs={max_abs_diff:.6f}, max_rel={max_rel_diff:.6f}"
|
||||
)
|
||||
self.assertTrue(
|
||||
torch.allclose(
|
||||
t_out,
|
||||
t_out_without_pcg,
|
||||
atol=1e-2,
|
||||
rtol=1e-2,
|
||||
),
|
||||
f"Piecewise CUDA graph embedding mismatch: max_abs_diff={max_abs_diff}, max_rel_diff={max_rel_diff}",
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user