From 1fdd6c89214d5c2f8a3529867b29de996ceb3237 Mon Sep 17 00:00:00 2001 From: Rumit Desai Date: Fri, 18 Sep 2026 00:42:46 -0700 Subject: [PATCH] [Runtime] Let out-of-tree platforms provide full graph backends (#37969) Co-authored-by: Xiaoyu Zhang <1182563586@qq.com> --- .../model_executor/runner_backend/utils.py | 18 +- python/sglang/srt/platforms/interface.py | 6 +- .../runner_backend/test_backend_resolution.py | 189 ++++++++++++++++++ 3 files changed, 207 insertions(+), 6 deletions(-) create mode 100644 test/registered/unit/model_executor/runner_backend/test_backend_resolution.py diff --git a/python/sglang/srt/model_executor/runner_backend/utils.py b/python/sglang/srt/model_executor/runner_backend/utils.py index 1c68be8d6..cb9449f41 100644 --- a/python/sglang/srt/model_executor/runner_backend/utils.py +++ b/python/sglang/srt/model_executor/runner_backend/utils.py @@ -37,9 +37,8 @@ from sglang.srt.model_executor.runner_backend.full_cuda_graph_backend import ( from sglang.srt.model_executor.runner_backend.tc_piecewise_cuda_graph_backend import ( TcPiecewiseCudaGraphBackend, ) -from sglang.srt.runtime_context import ( - get_exec, -) +from sglang.srt.platforms import current_platform +from sglang.srt.runtime_context import get_exec if TYPE_CHECKING: from sglang.srt.model_executor.runner.base_cuda_graph_runner import ( @@ -99,8 +98,17 @@ def resolve_decode_backend( "falling back to 'full'." ) _TC_PIECEWISE_DECODE_FALLBACK_LOGGED = True - return FullCudaGraphBackend( - cuda_graph_runner, enable_memory_saver=enable_memory_saver + + full_backend_cls = None + if current_platform.is_out_of_tree(): + full_backend_cls = current_platform.get_full_graph_backend_cls() + + if full_backend_cls is None: + full_backend_cls = FullCudaGraphBackend + + return full_backend_cls( + cuda_graph_runner, + enable_memory_saver=enable_memory_saver, ) diff --git a/python/sglang/srt/platforms/interface.py b/python/sglang/srt/platforms/interface.py index bcf034e69..b7941002a 100644 --- a/python/sglang/srt/platforms/interface.py +++ b/python/sglang/srt/platforms/interface.py @@ -12,7 +12,7 @@ Out-of-tree platforms register via setuptools entry_points under the from __future__ import annotations -from typing import TYPE_CHECKING, Optional, Type +from typing import TYPE_CHECKING, Any, Optional, Type from sglang.srt.platforms.device_mixin import DeviceMixin, PlatformEnum @@ -60,6 +60,10 @@ class SRTPlatform(DeviceMixin): """Return the graph runner class for this platform.""" raise NotImplementedError + def get_full_graph_backend_cls(self) -> type[Any]: + """Return the full device-graph backend class for this platform.""" + return None + def get_mha_kv_pool_cls(self) -> type: """Return the MHA KV pool class for this platform.""" raise NotImplementedError diff --git a/test/registered/unit/model_executor/runner_backend/test_backend_resolution.py b/test/registered/unit/model_executor/runner_backend/test_backend_resolution.py new file mode 100644 index 000000000..a78725a5e --- /dev/null +++ b/test/registered/unit/model_executor/runner_backend/test_backend_resolution.py @@ -0,0 +1,189 @@ +"""CPU-only tests for decode graph backend resolution.""" + +import unittest +from types import SimpleNamespace +from unittest import mock + +from sglang.srt.model_executor.cuda_graph_config import Backend +from sglang.srt.model_executor.runner_backend import utils as backend_utils +from sglang.srt.platforms.device_mixin import PlatformEnum +from sglang.srt.platforms.interface import SRTPlatform +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=1, suite="base-a-test-cpu") + + +def _make_graph_runner(*, device="custom"): + return SimpleNamespace(model_runner=SimpleNamespace(device=device)) + + +def _make_exec_config( + *, + backend=Backend.FULL, + enable_memory_saver=False, + debug_cuda_graph=False, +): + return SimpleNamespace( + graph=SimpleNamespace( + cuda_graph_config=SimpleNamespace(decode=SimpleNamespace(backend=backend)), + debug_cuda_graph=debug_cuda_graph, + ), + features=SimpleNamespace( + enable_memory_saver=enable_memory_saver, + ), + ) + + +def _make_platform(*, is_out_of_tree, full_backend_cls=None): + return SimpleNamespace( + is_out_of_tree=mock.Mock(return_value=is_out_of_tree), + get_full_graph_backend_cls=mock.Mock(return_value=full_backend_cls), + ) + + +class _OutOfTreePlatformWithoutFullGraphBackend(SRTPlatform): + _enum = PlatformEnum.OOT + + +class TestResolveDecodeBackend(CustomTestCase): + def test_out_of_tree_platform_provides_full_backend(self): + runner = _make_graph_runner() + exec_config = _make_exec_config(enable_memory_saver=True) + expected_backend = object() + backend_cls = mock.Mock(return_value=expected_backend) + platform = _make_platform( + is_out_of_tree=True, + full_backend_cls=backend_cls, + ) + + with ( + mock.patch.object( + backend_utils, + "get_exec", + return_value=exec_config, + ), + mock.patch.object( + backend_utils, + "current_platform", + platform, + ), + mock.patch.object( + backend_utils, + "FullCudaGraphBackend", + ) as default_backend_cls, + ): + actual_backend = backend_utils.resolve_decode_backend(runner) + + self.assertIs(actual_backend, expected_backend) + platform.get_full_graph_backend_cls.assert_called_once_with() + backend_cls.assert_called_once_with( + runner, + enable_memory_saver=True, + ) + default_backend_cls.assert_not_called() + + def test_out_of_tree_platform_without_full_backend_uses_default(self): + runner = _make_graph_runner() + exec_config = _make_exec_config() + expected_backend = object() + platform = _OutOfTreePlatformWithoutFullGraphBackend() + + with ( + mock.patch.object( + backend_utils, + "get_exec", + return_value=exec_config, + ), + mock.patch.object( + backend_utils, + "current_platform", + platform, + ), + mock.patch.object( + backend_utils, + "FullCudaGraphBackend", + return_value=expected_backend, + ) as default_backend_cls, + ): + actual_backend = backend_utils.resolve_decode_backend(runner) + + self.assertIs(actual_backend, expected_backend) + default_backend_cls.assert_called_once_with( + runner, + enable_memory_saver=False, + ) + + def test_in_tree_platform_uses_default_full_backend(self): + runner = _make_graph_runner(device="cuda") + exec_config = _make_exec_config() + expected_backend = object() + platform = _make_platform(is_out_of_tree=False) + + with ( + mock.patch.object( + backend_utils, + "get_exec", + return_value=exec_config, + ), + mock.patch.object( + backend_utils, + "current_platform", + platform, + ), + mock.patch.object( + backend_utils, + "FullCudaGraphBackend", + return_value=expected_backend, + ) as default_backend_cls, + ): + actual_backend = backend_utils.resolve_decode_backend(runner) + + self.assertIs(actual_backend, expected_backend) + platform.get_full_graph_backend_cls.assert_not_called() + default_backend_cls.assert_called_once_with( + runner, + enable_memory_saver=False, + ) + + def test_breakable_backend_does_not_consult_full_backend_factory(self): + runner = _make_graph_runner() + exec_config = _make_exec_config( + backend=Backend.BREAKABLE, + enable_memory_saver=True, + debug_cuda_graph=True, + ) + expected_backend = object() + platform = _make_platform(is_out_of_tree=True) + + with ( + mock.patch.object( + backend_utils, + "get_exec", + return_value=exec_config, + ), + mock.patch.object( + backend_utils, + "current_platform", + platform, + ), + mock.patch.object( + backend_utils, + "BreakableCudaGraphBackend", + return_value=expected_backend, + ) as breakable_backend_cls, + ): + actual_backend = backend_utils.resolve_decode_backend(runner) + + self.assertIs(actual_backend, expected_backend) + platform.is_out_of_tree.assert_not_called() + platform.get_full_graph_backend_cls.assert_not_called() + breakable_backend_cls.assert_called_once_with( + runner, + enable_memory_saver=True, + debug_eager=True, + ) + + +if __name__ == "__main__": + unittest.main()