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 89a711345..63b33bda4 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py +++ b/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 diff --git a/python/sglang/multimodal_gen/runtime/server_args/server_args.py b/python/sglang/multimodal_gen/runtime/server_args/server_args.py index 4c3b4c922..88db80a7d 100644 --- a/python/sglang/multimodal_gen/runtime/server_args/server_args.py +++ b/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, diff --git a/python/sglang/multimodal_gen/runtime/utils/torch_compile.py b/python/sglang/multimodal_gen/runtime/utils/torch_compile.py index 62ca3a178..4bb74f0ea 100644 --- a/python/sglang/multimodal_gen/runtime/utils/torch_compile.py +++ b/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""" 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 new file mode 100644 index 000000000..ae4eaac03 --- /dev/null +++ b/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