From 4b2c182d3f29e9bd35ff1271cf7bb9983c4c2c2c Mon Sep 17 00:00:00 2001 From: Joe Date: Tue, 25 Aug 2026 21:53:39 -0400 Subject: [PATCH] [diffusion] feat: support out-of-tree torch.compile backends (#36249) --- .../pipelines_core/stages/denoising.py | 2 +- .../runtime/platforms/interface.py | 8 +++ .../runtime/utils/torch_compile.py | 22 +++++- .../test/unit/test_regional_torch_compile.py | 69 +++++++++++++++++++ 4 files changed, 97 insertions(+), 4 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py index 836ed2a57..c26deaa95 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py @@ -508,7 +508,7 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin): config=dit_config, default="max-autotune-no-cudagraphs", ) - compile_kwargs = build_torch_compile_kwargs(mode=mode) + compile_kwargs = build_torch_compile_kwargs(mode=mode, module=module) logger.info(f"Compiling transformer with mode: {mode}") if getattr(self.server_args, "regional_compile", False): diff --git a/python/sglang/multimodal_gen/runtime/platforms/interface.py b/python/sglang/multimodal_gen/runtime/platforms/interface.py index 2c7866b0d..030dea6d7 100644 --- a/python/sglang/multimodal_gen/runtime/platforms/interface.py +++ b/python/sglang/multimodal_gen/runtime/platforms/interface.py @@ -123,6 +123,14 @@ class Platform: supported_quantization: list[str] = [] + def get_compile_backend(self, mode: str | None = None) -> str: + """Return the backend used to compile diffusion modules.""" + return self.simple_compile_backend + + def get_compile_options(self, module: torch.nn.Module) -> dict[str, object] | None: + """Return backend-specific options for a diffusion module.""" + return None + @lru_cache(maxsize=1) def is_cuda(self) -> bool: return self.is_cuda_static() diff --git a/python/sglang/multimodal_gen/runtime/utils/torch_compile.py b/python/sglang/multimodal_gen/runtime/utils/torch_compile.py index 4bb74f0ea..677ba2871 100644 --- a/python/sglang/multimodal_gen/runtime/utils/torch_compile.py +++ b/python/sglang/multimodal_gen/runtime/utils/torch_compile.py @@ -8,7 +8,6 @@ from typing import Any import torch.nn as nn from sglang.multimodal_gen.runtime.platforms import current_platform -from sglang.srt.utils.common import get_compiler_backend def maybe_enable_inductor_compute_comm_overlap() -> None: @@ -20,9 +19,26 @@ def maybe_enable_inductor_compute_comm_overlap() -> None: pass -def build_torch_compile_kwargs(*, mode: str | None) -> dict[str, object]: +def build_torch_compile_kwargs( + *, mode: str | None, module: nn.Module | None = None +) -> dict[str, object]: compile_kwargs: dict[str, object] = {"fullgraph": False, "dynamic": None} - if current_platform.is_npu(): + if current_platform.is_out_of_tree(): + backend = current_platform.get_compile_backend(mode) + compile_kwargs["backend"] = backend + if module is not None: + options = current_platform.get_compile_options(module) + if options is not None: + compile_kwargs["options"] = options + if ( + "options" not in compile_kwargs + and backend == "inductor" + and mode is not None + ): + compile_kwargs["mode"] = mode + elif current_platform.is_npu(): + from sglang.srt.utils.common import get_compiler_backend + compile_kwargs["backend"] = get_compiler_backend() compile_kwargs["dynamic"] = False elif mode is not None: diff --git a/python/sglang/multimodal_gen/test/unit/test_regional_torch_compile.py b/python/sglang/multimodal_gen/test/unit/test_regional_torch_compile.py index a8d168554..878b74b8f 100644 --- a/python/sglang/multimodal_gen/test/unit/test_regional_torch_compile.py +++ b/python/sglang/multimodal_gen/test/unit/test_regional_torch_compile.py @@ -12,6 +12,7 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.denoising import ( ) from sglang.multimodal_gen.runtime.utils.torch_compile import ( CompiledModuleRegistry, + build_torch_compile_kwargs, compile_matching_submodules, ) @@ -40,6 +41,63 @@ class _RegionalModel(_CompilableModule): self.proj_out = _CompilableModule() +@pytest.mark.parametrize( + ("backend", "options", "expected"), + [ + ( + "custom_backend", + {"pass_manager_config": {"persistent_buffers": ["weight"]}}, + { + "backend": "custom_backend", + "options": {"pass_manager_config": {"persistent_buffers": ["weight"]}}, + }, + ), + ( + "inductor", + {"max_autotune": True}, + {"backend": "inductor", "options": {"max_autotune": True}}, + ), + ( + "inductor", + None, + {"backend": "inductor", "mode": "max-autotune-no-cudagraphs"}, + ), + ], +) +def test_out_of_tree_platform_controls_compile_kwargs(backend, options, expected): + """Out-of-tree hooks select valid backend, mode, and option combinations.""" + module = _CompilableModule() + with ( + patch( + "sglang.multimodal_gen.runtime.utils.torch_compile." + "current_platform.is_out_of_tree", + return_value=True, + ), + patch( + "sglang.multimodal_gen.runtime.utils.torch_compile." + "current_platform.get_compile_backend", + return_value=backend, + ) as get_compile_backend, + patch( + "sglang.multimodal_gen.runtime.utils.torch_compile." + "current_platform.get_compile_options", + return_value=options, + ) as get_compile_options, + ): + compile_kwargs = build_torch_compile_kwargs( + mode="max-autotune-no-cudagraphs", + module=module, + ) + + assert compile_kwargs == { + "dynamic": None, + "fullgraph": False, + **expected, + } + get_compile_backend.assert_called_once_with("max-autotune-no-cudagraphs") + get_compile_options.assert_called_once_with(module) + + def test_ltx2_compile_conditions_match_only_direct_blocks(): conditions = LTX2VideoTransformer3DModel._compile_conditions @@ -99,6 +157,7 @@ def test_compiled_module_registry_installs_regions_once(): def test_denoising_stage_selects_regional_compile(): model = _RegionalModel() + compile_kwargs = {"backend": "custom_backend"} stage = DenoisingStage.__new__(DenoisingStage) stage.server_args = SimpleNamespace( enable_breakable_cuda_graph=False, @@ -121,8 +180,18 @@ def test_denoising_stage_selects_regional_compile(): "sglang.multimodal_gen.runtime.pipelines_core.stages.denoising." "maybe_enable_inductor_compute_comm_overlap" ), + patch( + "sglang.multimodal_gen.runtime.pipelines_core.stages.denoising." + "build_torch_compile_kwargs", + return_value=compile_kwargs, + ) as build_compile_kwargs, ): stage._maybe_torch_compile(model) + build_compile_kwargs.assert_called_once_with(mode="default", module=model) assert [len(block.compile_calls) for block in model.transformer_blocks] == [1, 1] + assert [block.compile_calls for block in model.transformer_blocks] == [ + [compile_kwargs], + [compile_kwargs], + ] assert not model.compile_calls