diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index d91c51328..58c3b89d8 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -1199,6 +1199,17 @@ class ModelRunner: model_runner=self, init_new_workspace=init_new_workspace ) + def _decode_cuda_graph_runner_cls(self): + """Decode CUDA-graph runner class to construct. + + Subclasses can override this to install specialized decode graph runners. + """ + from sglang.srt.model_executor.runner.decode_cuda_graph_runner import ( + DecodeCudaGraphRunner, + ) + + return DecodeCudaGraphRunner + def init_decode_cuda_graph(self): self.decode_cuda_graph_runner = None self.graph_mem_usage = 0 diff --git a/python/sglang/srt/model_executor/model_runner_components/cuda_graph_setup.py b/python/sglang/srt/model_executor/model_runner_components/cuda_graph_setup.py index f4ded560b..74fa9061f 100644 --- a/python/sglang/srt/model_executor/model_runner_components/cuda_graph_setup.py +++ b/python/sglang/srt/model_executor/model_runner_components/cuda_graph_setup.py @@ -373,12 +373,8 @@ def capture_decode_graph(*, model_runner: ModelRunner) -> DecodeGraphCapture: GraphRunnerCls = current_platform.get_graph_runner_cls() runner = GraphRunnerCls(model_runner) else: - from sglang.srt.model_executor.runner.decode_cuda_graph_runner import ( - DecodeCudaGraphRunner, - ) - graph_runners = defaultdict( - lambda: DecodeCudaGraphRunner, + model_runner._decode_cuda_graph_runner_cls, { "cpu": CPUGraphRunner, "npu": NPUGraphRunner, diff --git a/test/registered/unit/model_executor/model_runner_components/test_cuda_graph_setup.py b/test/registered/unit/model_executor/model_runner_components/test_cuda_graph_setup.py index d81f23b5e..0632fe9c3 100644 --- a/test/registered/unit/model_executor/model_runner_components/test_cuda_graph_setup.py +++ b/test/registered/unit/model_executor/model_runner_components/test_cuda_graph_setup.py @@ -1,9 +1,12 @@ import sys +from types import SimpleNamespace import pytest from sglang.srt.model_executor.cuda_graph_config import Phase +from sglang.srt.model_executor.model_runner_components import cuda_graph_setup from sglang.srt.model_executor.model_runner_components.cuda_graph_setup import ( + capture_decode_graph, should_skip_auto_prefill_cuda_graph_for_memory, ) from sglang.test.ci.ci_register import register_cpu_ci @@ -22,5 +25,42 @@ def test_explicit_prefill_backend_bypasses_memory_gate(): ) +def test_model_runner_can_override_decode_graph_runner(monkeypatch): + class CustomGraphRunner: + def __init__(self, model_runner): + self.model_runner = model_runner + + class TestModelRunner: + is_generation = True + device = "cuda" + gpu_id = 0 + is_draft_worker = False + spec_algorithm = SimpleNamespace(is_speculative=lambda: False) + server_args = SimpleNamespace( + model_impl="auto", + cuda_graph_config=SimpleNamespace( + decode=SimpleNamespace(backend="default") + ), + ) + + def _decode_cuda_graph_runner_cls(self): + return CustomGraphRunner + + model_runner = TestModelRunner() + monkeypatch.setattr(cuda_graph_setup, "check_cuda_graph_backend", lambda *_: False) + monkeypatch.setattr(cuda_graph_setup, "get_available_gpu_memory", lambda *_: 10.0) + monkeypatch.setattr( + cuda_graph_setup, "get_batch_sizes_to_capture", lambda *_: ([1], None) + ) + monkeypatch.setattr( + cuda_graph_setup.current_platform, "is_out_of_tree", lambda: False + ) + + capture = capture_decode_graph(model_runner=model_runner) + + assert isinstance(capture.runner, CustomGraphRunner) + assert capture.runner.model_runner is model_runner + + if __name__ == "__main__": sys.exit(pytest.main([__file__, "-v"]))