feat(diffusion): add regional torch compile (#32696)

python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py

@@ -406,11 +406,22 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin):
             compile_kwargs = build_torch_compile_kwargs(mode=mode)
             logger.info(f"Compiling transformer with mode: {mode}")
 
-        # TODO(triple-mu): support customized fullgraph and dynamic in the future
-        self._torch_compile_registry.compile_once(
-            module,
-            compile_kwargs=compile_kwargs,
-        )
+        if getattr(self.server_args, "regional_compile", False):
+            compiled_count = self._torch_compile_registry.compile_regions_once(
+                module,
+                compile_kwargs=compile_kwargs,
+            )
+            logger.info(
+                "Enabled regional torch.compile for %d submodules in %s",
+                compiled_count,
+                type(module).__name__,
+            )
+        else:
+            # TODO(triple-mu): support customized fullgraph and dynamic in the future
+            self._torch_compile_registry.compile_once(
+                module,
+                compile_kwargs=compile_kwargs,
+            )
 
     def _maybe_enable_cache_dit_and_torch_compile(
         self, num_inference_steps: int | tuple[int, int], batch: Req

python/sglang/multimodal_gen/runtime/server_args/server_args.py

@@ -284,6 +284,7 @@ class ServerArgs(DisaggServerArgsMixin):
 
     # Compilation
     enable_torch_compile: bool = False
+    regional_compile: bool = False
 
     # Breakable CUDA graph (BCG): capture the DiT forward as CUDA-graph
     # segments split at attention modules (SP all-to-all / dynamic attention
@@ -1513,6 +1514,16 @@ class ServerArgs(DisaggServerArgsMixin):
             + "so first real requests do not pay compile latency. "
             + "However, will likely cause precision drifts. See (https://github.com/pytorch/pytorch/issues/145213)",
         )
+        parser.add_argument(
+            "--regional-compile",
+            action=StoreBoolean,
+            default=ServerArgs.regional_compile,
+            help=(
+                "Compile repeated DiT submodules selected by the model's "
+                "_compile_conditions instead of compiling the whole transformer. "
+                "Requires --enable-torch-compile."
+            ),
+        )
         parser.add_argument(
             "--offload-during-compile",
             action=StoreBoolean,

python/sglang/multimodal_gen/runtime/utils/torch_compile.py

@@ -45,6 +45,28 @@ def resolve_torch_compile_mode(
     return default
 
 
+def compile_matching_submodules(
+    module: nn.Module,
+    *,
+    compile_kwargs: dict[str, object],
+) -> int:
+    conditions = getattr(module, "_compile_conditions", ())
+    matches = [
+        submodule
+        for name, submodule in module.named_modules()
+        if name and any(condition(name, submodule) for condition in conditions)
+    ]
+    if not matches:
+        raise ValueError(
+            "regional compile found no matching submodules; "
+            f"check {type(module).__name__}._compile_conditions"
+        )
+
+    for submodule in matches:
+        submodule.compile(**compile_kwargs)
+    return len(matches)
+
+
 @dataclass
 class CompiledModuleRegistry:
     module_ids: set[int] = field(default_factory=set)
@@ -65,6 +87,22 @@ class CompiledModuleRegistry:
         self.module_ids.add(module_id)
         return True
 
+    def compile_regions_once(
+        self,
+        module: nn.Module,
+        *,
+        compile_kwargs: dict[str, object],
+    ) -> int:
+        module_id = id(module)
+        if module_id in self.module_ids:
+            return 0
+        compiled_count = compile_matching_submodules(
+            module,
+            compile_kwargs=compile_kwargs,
+        )
+        self.module_ids.add(module_id)
+        return compiled_count
+
 
 class CallableModule(nn.Module):
     """Module wrapper for compiling non-forward callables with module.compile"""

python/sglang/multimodal_gen/test/unit/test_regional_torch_compile.py

@@ -0,0 +1,126 @@
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import pytest
+import torch.nn as nn
+
+from sglang.multimodal_gen.configs.models.dits.ltx_2 import LTX2ArchConfig
+from sglang.multimodal_gen.runtime.pipelines_core.stages.denoising import (
+    DenoisingStage,
+)
+from sglang.multimodal_gen.runtime.utils.torch_compile import (
+    CompiledModuleRegistry,
+    compile_matching_submodules,
+)
+
+
+class _CompilableModule(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.compile_calls = []
+
+    def compile(self, **kwargs):
+        self.compile_calls.append(kwargs)
+
+
+class _RegionalModel(_CompilableModule):
+    _compile_conditions = [
+        lambda name, _module: name.startswith("transformer_blocks.")
+        and name.count(".") == 1
+    ]
+
+    def __init__(self):
+        super().__init__()
+        self.transformer_blocks = nn.ModuleList(
+            [_CompilableModule(), _CompilableModule()]
+        )
+        self.transformer_blocks[0].inner = _CompilableModule()
+        self.proj_out = _CompilableModule()
+
+
+def test_ltx2_compile_conditions_match_only_direct_blocks():
+    conditions = LTX2ArchConfig()._compile_conditions
+
+    assert conditions
+    assert any(condition("transformer_blocks.0", object()) for condition in conditions)
+    assert not any(
+        condition("transformer_blocks.0.attn1", object()) for condition in conditions
+    )
+    assert not any(
+        condition("transformer_blocks", object()) for condition in conditions
+    )
+
+
+def test_compile_matching_submodules_matches_only_declared_regions():
+    model = _RegionalModel()
+
+    count = compile_matching_submodules(
+        model,
+        compile_kwargs={"mode": "default", "fullgraph": False},
+    )
+
+    assert count == 2
+    assert [len(block.compile_calls) for block in model.transformer_blocks] == [1, 1]
+    assert not model.transformer_blocks[0].inner.compile_calls
+    assert not model.proj_out.compile_calls
+    assert not model.compile_calls
+
+
+def test_compile_matching_submodules_fails_when_no_region_matches():
+    model = _RegionalModel()
+    model._compile_conditions = [lambda _name, _module: False]
+
+    with pytest.raises(ValueError, match="no matching submodules"):
+        compile_matching_submodules(model, compile_kwargs={"mode": "default"})
+
+
+def test_compiled_module_registry_installs_regions_once():
+    model = _RegionalModel()
+    registry = CompiledModuleRegistry()
+
+    assert (
+        registry.compile_regions_once(
+            model,
+            compile_kwargs={"mode": "default"},
+        )
+        == 2
+    )
+    assert (
+        registry.compile_regions_once(
+            model,
+            compile_kwargs={"mode": "default"},
+        )
+        == 0
+    )
+    assert [len(block.compile_calls) for block in model.transformer_blocks] == [1, 1]
+
+
+def test_denoising_stage_selects_regional_compile():
+    model = _RegionalModel()
+    stage = DenoisingStage.__new__(DenoisingStage)
+    stage.server_args = SimpleNamespace(
+        enable_breakable_cuda_graph=False,
+        enable_torch_compile=True,
+        regional_compile=True,
+        pipeline_config=SimpleNamespace(
+            dit_config=SimpleNamespace(torch_compile_mode="default")
+        ),
+    )
+    stage._cache_dit_enabled = False
+    stage._torch_compile_registry = CompiledModuleRegistry()
+
+    with (
+        patch(
+            "sglang.multimodal_gen.runtime.pipelines_core.stages.denoising."
+            "current_platform.is_npu",
+            return_value=False,
+        ),
+        patch(
+            "sglang.multimodal_gen.runtime.pipelines_core.stages.denoising."
+            "maybe_enable_inductor_compute_comm_overlap"
+        ),
+    ):
+        stage._maybe_torch_compile(model)
+
+    assert [len(block.compile_calls) for block in model.transformer_blocks] == [1, 1]
+    assert not model.compile_calls