[diffusion] feat: support out-of-tree torch.compile backends (#36249)

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):

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()

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:

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