[perf] Add configurable diffusion compile granularity (#4603)

Signed-off-by: TaffyOfficial <2324465096@qq.com>
Co-authored-by: TaffyOfficial <2324465096@qq.com>

docs/.nav.yml

@@ -56,6 +56,7 @@ nav:
     - Diffusion Features:
       - Overview: user_guide/diffusion_features.md
       - Feature Compatibility: user_guide/feature_compatibility.md
+      - Regional Compilation: user_guide/diffusion/regional_compilation.md
       - Cache Acceleration:
         - TeaCache: user_guide/diffusion/cache_acceleration/teacache.md
         - Cache-DiT: user_guide/diffusion/cache_acceleration/cache_dit.md

docs/user_guide/diffusion/regional_compilation.md

@@ -0,0 +1,38 @@
+# Regional Compilation
+
+Regional compilation applies `torch.compile` to the repeated transformer blocks
+declared by a diffusion model. It is the default compilation scope when
+diffusion inference runs without `--enforce-eager`.
+
+## Configuration
+
+Dynamic compilation is enabled by default so the compiled regions can handle
+mixed resolutions. For a fixed-shape workload, disable it explicitly:
+
+```bash
+vllm serve <model> --omni --no-diffusion-compile-dynamic
+```
+
+The equivalent per-stage deploy configuration is:
+
+```yaml
+stages:
+  - stage_id: 0
+    diffusion_compile_granularity: regional
+    diffusion_compile_dynamic: false
+```
+
+For an experimental whole-transformer compile scope, set
+`--diffusion-compile-granularity full` or use
+`diffusion_compile_granularity: full` in the stage configuration. Full scope may
+still contain graph breaks; it does not force one graph. It is rejected when
+HSDP, sequence parallelism, CPU offload, or layerwise offload is enabled. Use
+regional scope with those features.
+
+These settings control the generic model-runner compilation path. Pipelines
+that provide their own `setup_compile()` implementation manage their compilation
+policy independently. Compilation is lazy, so backend or graph errors can first
+surface on the initial request.
+
+Use `--enforce-eager` to disable the model runner's generic compile setup.
+Pipelines that compile internally define their own eager-mode behavior.

tests/config/test_omni_config.py

@@ -628,6 +628,47 @@ def test_from_pipeline_config_preserves_diffusion_parallel_mask_sp_padding(tmp_p
     assert stage.parallel_config.mask_sp_padding is True
 
 
+def test_from_pipeline_config_routes_regional_compile_dynamic(tmp_path):
+    deploy_path = tmp_path / "dreamzero_compile.yaml"
+    deploy_path.write_text(
+        "\n".join(
+            [
+                "pipeline: dreamzero",
+                "async_chunk: false",
+                "stages:",
+                "  - stage_id: 0",
+                "    diffusion_compile_granularity: regional",
+                "    diffusion_compile_dynamic: false",
+            ]
+        )
+    )
+
+    configured_stage = _from_pipeline_key("dreamzero", deploy_config_path=str(deploy_path)).stage_by_id(0)
+    overridden_stage = _from_pipeline_key(
+        "dreamzero",
+        deploy_config_path=str(deploy_path),
+        cli_overrides={
+            "diffusion_compile_granularity": "full",
+            "diffusion_compile_dynamic": True,
+        },
+    ).stage_by_id(0)
+
+    assert configured_stage.diffusion_config.diffusion_compile_granularity == "regional"
+    assert configured_stage.diffusion_config.diffusion_compile_dynamic is False
+    assert overridden_stage.diffusion_config.diffusion_compile_granularity == "full"
+    assert overridden_stage.diffusion_config.diffusion_compile_dynamic is True
+
+
+def test_structured_diffusion_config_rejects_non_boolean_compile_dynamic():
+    with pytest.raises(ValidationError, match="diffusion_compile_dynamic"):
+        omni_config_module._DiffusionConfigProjection(diffusion_compile_dynamic="false")
+
+
+def test_structured_diffusion_config_rejects_invalid_compile_granularity():
+    with pytest.raises(ValidationError, match="diffusion_compile_granularity"):
+        omni_config_module._DiffusionConfigProjection(diffusion_compile_granularity="block")
+
+
 def test_from_pipeline_config_matches_stage_config_to_omegaconf_behavior_for_representative_stage():
     pipeline = _resolve_pipeline_or_skip("qwen3_tts")
     legacy_stage = merge_pipeline_deploy(pipeline, _load_default_deploy(pipeline))[0]

tests/diffusion/test_compile.py

@@ -54,3 +54,23 @@ def test_regionally_compile_matches_wrapped_blocks_by_declared_container_attr(mo
     assert all(not block.compile_called for block in model.transformer_blocks)
     assert not model.other_blocks[0].compile_called
     assert model.transformer_blocks[0].forward("ok") == "compiled:ok"
+
+
+def test_regionally_compile_does_not_partially_mutate_on_setup_failure(monkeypatch):
+    model = _ModelWithWrappedRepeatedBlocks()
+    original_forwards = [block.forward.__func__ for block in model.transformer_blocks]
+    compile_calls = 0
+
+    def _compile(fn, *args, **kwargs):
+        nonlocal compile_calls
+        compile_calls += 1
+        if compile_calls == 2:
+            raise RuntimeError("compile setup failed")
+        return lambda *fn_args, **fn_kwargs: fn(*fn_args, **fn_kwargs)
+
+    monkeypatch.setattr(compile_module.torch, "compile", _compile)
+
+    with pytest.raises(RuntimeError, match="compile setup failed"):
+        regionally_compile(model, dynamic=True)
+
+    assert [block.forward.__func__ for block in model.transformer_blocks] == original_forwards

tests/diffusion/test_diffusion_engine.py

@@ -392,6 +392,57 @@ class TestRequestBatchCapability:
         warmup.assert_called_once_with()
 
 
+class TestDiffusionCompileConfig:
+    pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+    def test_config_defaults_to_dynamic_regional_compile(self) -> None:
+        config = OmniDiffusionConfig(model="test")
+
+        assert config.diffusion_compile_granularity == "regional"
+        assert config.diffusion_compile_dynamic is True
+
+    def test_from_kwargs_preserves_compile_controls(self) -> None:
+        config = OmniDiffusionConfig.from_kwargs(
+            model="test",
+            diffusion_compile_granularity="full",
+            diffusion_compile_dynamic=False,
+        )
+
+        assert config.diffusion_compile_granularity == "full"
+        assert config.diffusion_compile_dynamic is False
+
+    def test_config_rejects_invalid_compile_granularity(self) -> None:
+        with pytest.raises(ValueError, match="diffusion_compile_granularity"):
+            OmniDiffusionConfig(model="test", diffusion_compile_granularity="block")
+
+    def test_config_rejects_non_boolean_compile_dynamic(self) -> None:
+        with pytest.raises(TypeError, match="diffusion_compile_dynamic"):
+            OmniDiffusionConfig(model="test", diffusion_compile_dynamic="false")
+
+    @pytest.mark.parametrize(
+        "kwargs, feature",
+        [
+            ({"parallel_config": {"ulysses_degree": 2}, "num_gpus": 2}, "sequence parallelism"),
+            (
+                {
+                    "parallel_config": {"use_hsdp": True, "hsdp_shard_size": 2},
+                    "num_gpus": 2,
+                },
+                "HSDP",
+            ),
+            ({"enable_cpu_offload": True}, "CPU offload"),
+            ({"enable_layerwise_offload": True}, "layerwise offload"),
+        ],
+    )
+    def test_full_compile_rejects_incompatible_features(self, kwargs, feature) -> None:
+        with pytest.raises(ValueError, match=feature):
+            OmniDiffusionConfig(
+                model="test",
+                diffusion_compile_granularity="full",
+                **kwargs,
+            )
+
+
 class TestRequestBatchAdmission:
     pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
 

tests/diffusion/test_diffusion_model_runner.py

@@ -121,6 +121,15 @@ class _FinalOnlyStepPipeline:
         return DiffusionOutput(output=state.latents.clone())
 
 
+class _CompileTrackingModel:
+    def __init__(self):
+        self.compile_calls = []
+
+    def compile(self, *args, **kwargs):
+        self.compile_calls.append((args, kwargs))
+        return self
+
+
 def _make_request():
     sampling_params = SimpleNamespace(
         generator=None,
@@ -177,10 +186,20 @@ def _make_runner(cache_backend, cache_backend_name: str, enable_cache_dit_summar
     return runner
 
 
-def _make_compile_runner(*, use_hsdp: bool):
+def _make_compile_runner(
+    model=None,
+    *,
+    compile_granularity: str = "regional",
+    compile_dynamic: bool = True,
+    use_hsdp: bool = False,
+):
     runner = object.__new__(DiffusionModelRunner)
-    runner.pipeline = SimpleNamespace(transformer=SimpleNamespace())
-    runner.od_config = SimpleNamespace(parallel_config=SimpleNamespace(use_hsdp=use_hsdp))
+    runner.pipeline = SimpleNamespace(transformer=model or SimpleNamespace())
+    runner.od_config = SimpleNamespace(
+        diffusion_compile_granularity=compile_granularity,
+        diffusion_compile_dynamic=compile_dynamic,
+        parallel_config=SimpleNamespace(use_hsdp=use_hsdp),
+    )
     return runner
 
 
@@ -208,6 +227,65 @@ def test_compile_transformer_regionally_compiles_blocks(monkeypatch, use_hsdp):
     ]
 
 
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_compile_transformer_uses_regional_dynamic_false_config(monkeypatch):
+    model = _CompileTrackingModel()
+    runner = _make_compile_runner(model, compile_dynamic=False)
+    compiled_model = object()
+    regional_calls = []
+
+    def _regionally_compile(target, **kwargs):
+        regional_calls.append((target, kwargs))
+        return compiled_model
+
+    monkeypatch.setattr(model_runner_module, "regionally_compile", _regionally_compile)
+
+    DiffusionModelRunner._compile_transformer(runner, "transformer")
+
+    assert regional_calls == [(model, {"dynamic": False})]
+    assert model.compile_calls == []
+    assert runner.pipeline.transformer is compiled_model
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_compile_transformer_uses_full_granularity(monkeypatch):
+    model = _CompileTrackingModel()
+    runner = _make_compile_runner(model, compile_granularity="full", compile_dynamic=False)
+    regional_calls = []
+
+    monkeypatch.setattr(
+        model_runner_module,
+        "regionally_compile",
+        lambda *args, **kwargs: regional_calls.append((args, kwargs)),
+    )
+
+    DiffusionModelRunner._compile_transformer(runner, "transformer")
+
+    assert model.compile_calls == [((), {"dynamic": False})]
+    assert regional_calls == []
+    assert runner.pipeline.transformer is model
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_compile_transformer_falls_back_after_synchronous_setup_failure(monkeypatch, caplog):
+    model = _CompileTrackingModel()
+    runner = _make_compile_runner(model)
+
+    def _regionally_compile(*args, **kwargs):
+        raise RuntimeError("compile setup failed")
+
+    monkeypatch.setattr(model_runner_module, "regionally_compile", _regionally_compile)
+
+    DiffusionModelRunner._compile_transformer(runner, "transformer")
+
+    assert runner.pipeline.transformer is model
+    assert "failed before activation" in caplog.text
+    assert "lazy compilation errors" in caplog.text
+
+
 @pytest.mark.core_model
 @pytest.mark.cpu
 def test_execute_stepwise_streaming_returns_chunks_at_boundaries(monkeypatch):

tests/entrypoints/test_async_omni_diffusion_config.py

@@ -85,6 +85,8 @@ def test_default_stage_config_defaults_nullified_parallel_size_kwargs():
             "tensor_parallel_size": None,
             "enable_expert_parallel": None,
             "enforce_eager": None,
+            "diffusion_compile_granularity": None,
+            "diffusion_compile_dynamic": None,
         }
     )[0]
 
@@ -94,6 +96,8 @@ def test_default_stage_config_defaults_nullified_parallel_size_kwargs():
     assert parallel_config.tensor_parallel_size == 1
     assert parallel_config.enable_expert_parallel is False
     assert stage_cfg["engine_args"]["enforce_eager"] is False
+    assert stage_cfg["engine_args"]["diffusion_compile_granularity"] == "regional"
+    assert stage_cfg["engine_args"]["diffusion_compile_dynamic"] is True
 
 
 def test_default_stage_config_propagates_ulysses_mode():
@@ -258,6 +262,32 @@ def test_serve_cli_accepts_diffusion_pipeline_profiler_flag():
     assert stage_cfg["engine_args"]["enable_diffusion_pipeline_profiler"] is True
 
 
+def test_serve_cli_accepts_diffusion_compile_controls():
+    """Ensure both compile controls reach the diffusion stage."""
+    parser = TrackingArgumentParser()
+    subparsers = parser.add_subparsers(dest="command")
+    OmniServeCommand().subparser_init(subparsers)
+
+    args = parser.parse_args(
+        [
+            "serve",
+            "Lightricks/LTX-Video-0.9.8-13B-distilled",
+            "--omni",
+            "--diffusion-compile-granularity",
+            "full",
+            "--no-diffusion-compile-dynamic",
+        ]
+    )
+
+    explicit_kwargs = args.get_explicit_kwargs_dict()
+    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(explicit_kwargs)[0]
+
+    assert args.diffusion_compile_granularity == "full"
+    assert args.diffusion_compile_dynamic is False
+    assert stage_cfg["engine_args"]["diffusion_compile_granularity"] == "full"
+    assert stage_cfg["engine_args"]["diffusion_compile_dynamic"] is False
+
+
 def test_serve_cli_accepts_diffusion_attention_backend():
     """Ensure diffusion serve CLI exposes the shorthand backend flag."""
     parser = TrackingArgumentParser()

vllm_omni/config/omni_config.py

@@ -14,7 +14,7 @@ import copy
 from collections.abc import Mapping
 from dataclasses import dataclass, field, fields
 from pathlib import Path
-from typing import Any, TypeAlias, TypedDict, cast
+from typing import Any, Literal, TypeAlias, TypedDict, cast
 
 from pydantic import ConfigDict, Field
 from vllm.config.utils import config
@@ -470,6 +470,8 @@ class _DiffusionConfigProjection:
     enable_cpu_offload: bool = False
     enable_layerwise_offload: bool = False
     pin_cpu_memory: bool = True
+    diffusion_compile_granularity: Literal["regional", "full"] = "regional"
+    diffusion_compile_dynamic: bool = Field(default=True, strict=True)
     vae_use_slicing: bool = False
     vae_use_tiling: bool = False
     mask_strategy_file_path: str | None = None

vllm_omni/config/stage_config.py

@@ -388,6 +388,8 @@ class StageDeployConfig:
     diffusion_attention_config: dict[str, Any] | None = None
 
     # Diffusion execution, cache, and VAE behavior.
+    diffusion_compile_granularity: str | None = None
+    diffusion_compile_dynamic: bool | None = None
     cache_backend: str | None = None
     cache_config: dict[str, Any] | None = None
     enable_cache_dit_summary: bool | None = None

vllm_omni/diffusion/compile.py

@@ -53,23 +53,29 @@ def regionally_compile(
 
     repeated_block_attrs = getattr(model, "_layerwise_offload_blocks_attrs", [])
 
-    # Check if we have modules with the specified class names
-    has_compiled_region = False
-    compiled_region_count = 0
+    # Build all compiled callables before mutating the model. This keeps setup
+    # failures atomic: callers can safely continue with the uncompiled model if
+    # torch.compile raises synchronously for any repeated block.
+    compiled_forwards: list[tuple[nn.Module, Any]] = []
     for name, submod in model.named_modules():
         if _matches_repeated_block(name, submod, repeated_blocks, repeated_block_attrs):
             # Compile the block compute while keeping nn.Module.__call__ hooks
             # outside the compiled graph.
-            submod.forward = torch.compile(submod.forward, *compile_args, **compile_kwargs)
-            has_compiled_region = True
-            compiled_region_count += 1
+            compiled_forwards.append(
+                (
+                    submod,
+                    torch.compile(submod.forward, *compile_args, **compile_kwargs),
+                )
+            )
 
-    if not has_compiled_region:
+    if not compiled_forwards:
         logger.warning(f"Regional compilation skipped because {repeated_blocks} classes are not found in the model.")
     else:
+        for submod, compiled_forward in compiled_forwards:
+            submod.forward = compiled_forward
         logger.info(
             "Regional compilation applied to %d module(s) for repeated blocks %s.",
-            compiled_region_count,
+            len(compiled_forwards),
             repeated_blocks,
         )
 

vllm_omni/diffusion/data.py

@@ -684,6 +684,10 @@ class OmniDiffusionConfig:
 
     # Compilation
     enforce_eager: bool = False
+    # Controls the generic compilation path used when a pipeline does not
+    # provide its own setup_compile() implementation.
+    diffusion_compile_granularity: str = "regional"
+    diffusion_compile_dynamic: bool = True
 
     # Parallel weight loading (for faster diffusion model startup)
     enable_multithread_weight_load: bool = True
@@ -880,6 +884,13 @@ class OmniDiffusionConfig:
         )
 
     def __post_init__(self):
+        if self.diffusion_compile_granularity not in {"regional", "full"}:
+            raise ValueError(
+                "diffusion_compile_granularity must be 'regional' or 'full', "
+                f"got {self.diffusion_compile_granularity!r}"
+            )
+        if not isinstance(self.diffusion_compile_dynamic, bool):
+            raise TypeError(f"diffusion_compile_dynamic must be a bool, got {type(self.diffusion_compile_dynamic)!r}")
         self.master_port = self._resolve_master_port()
         self.request_batch_max_wait_ms = float(self.request_batch_max_wait_ms or 0.0)
         if self.request_batch_max_wait_ms < 0:
@@ -915,6 +926,23 @@ class OmniDiffusionConfig:
                 f"num_gpus ({self.num_gpus}) < parallel_config.world_size ({self.parallel_config.world_size})"
             )
 
+        if self.diffusion_compile_granularity == "full":
+            incompatible_features = []
+            if self.parallel_config.use_hsdp:
+                incompatible_features.append("HSDP")
+            if self.parallel_config.sequence_parallel_size > 1:
+                incompatible_features.append("sequence parallelism")
+            if self.enable_cpu_offload:
+                incompatible_features.append("CPU offload")
+            if self.enable_layerwise_offload:
+                incompatible_features.append("layerwise offload")
+            if incompatible_features:
+                features = ", ".join(incompatible_features)
+                raise ValueError(
+                    "diffusion_compile_granularity='full' is incompatible with "
+                    f"{features}; use 'regional' compilation instead"
+                )
+
         # Convert string dtype to torch.dtype if needed
         if isinstance(self.dtype, str):
             dtype_map = {

vllm_omni/diffusion/worker/diffusion_model_runner.py

@@ -151,15 +151,33 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
         if model is None:
             return
 
+        compile_granularity = self.od_config.diffusion_compile_granularity
+        compile_dynamic = self.od_config.diffusion_compile_dynamic
         try:
-            setattr(self.pipeline, attr_name, regionally_compile(model, dynamic=True))
-            logger.info("Model runner: %s compiled with torch.compile.", attr_name)
+            if compile_granularity == "full":
+                model.compile(dynamic=compile_dynamic)
+                compiled_model = model
+            else:
+                compiled_model = regionally_compile(model, dynamic=compile_dynamic)
+            setattr(self.pipeline, attr_name, compiled_model)
         except Exception as e:
             logger.warning(
-                "Model runner: torch.compile for %s failed: %s. Using eager mode.",
+                "Model runner: %s torch.compile setup for %s failed before activation: %s. "
+                "Continuing with the uncompiled model; lazy compilation errors can still "
+                "surface on the first request.",
+                compile_granularity,
                 attr_name,
                 e,
             )
+            return
+
+        logger.info(
+            "Model runner: %s configured for lazy %s torch.compile with dynamic=%s; "
+            "compilation errors may surface on the first request.",
+            attr_name,
+            compile_granularity,
+            compile_dynamic,
+        )
 
     def load_model(
         self,

vllm_omni/engine/arg_utils.py

@@ -466,6 +466,8 @@ class OrchestratorArgs:
     hsdp_replicate_size: int = 1
     diffusion_attention_backend: str | None = None
     diffusion_attention_config: str | None = None
+    diffusion_compile_granularity: str | None = None
+    diffusion_compile_dynamic: bool | None = None
     cache_backend: str = "none"
     cache_config: str | None = None
     enable_cache_dit_summary: bool = False

vllm_omni/engine/async_omni_engine.py

@@ -999,6 +999,14 @@ class AsyncOmniEngine:
             "enable_cpu_offload": kwargs.get("enable_cpu_offload", False),
             "enable_layerwise_offload": kwargs.get("enable_layerwise_offload", False),
             "enforce_eager": False if kwargs.get("enforce_eager") is None else kwargs.get("enforce_eager"),
+            "diffusion_compile_granularity": (
+                "regional"
+                if kwargs.get("diffusion_compile_granularity") is None
+                else kwargs["diffusion_compile_granularity"]
+            ),
+            "diffusion_compile_dynamic": (
+                True if kwargs.get("diffusion_compile_dynamic") is None else kwargs["diffusion_compile_dynamic"]
+            ),
             "boundary_ratio": kwargs.get("boundary_ratio", None),
             "flow_shift": kwargs.get("flow_shift", None),
             "diffusion_load_format": kwargs.get("diffusion_load_format", "default"),

vllm_omni/entrypoints/cli/serve.py

@@ -423,6 +423,25 @@ class OmniServeCommand(CLISubcommand):
                 "custom_pipeline, dummy, or diffusers for the HF diffusers adapter."
             ),
         )
+        omni_config_group.add_argument(
+            "--diffusion-compile-granularity",
+            choices=["regional", "full"],
+            default=None,
+            help=(
+                "Compilation scope for the generic diffusion model runner. "
+                "'regional' compiles repeated blocks (default); 'full' compiles the whole transformer and is "
+                "incompatible with HSDP, sequence parallelism, CPU offload, and layerwise offload."
+            ),
+        )
+        omni_config_group.add_argument(
+            "--diffusion-compile-dynamic",
+            action=argparse.BooleanOptionalAction,
+            default=None,
+            help=(
+                "Use dynamic shapes for the selected generic diffusion compile scope. "
+                "Disable for fixed-shape workloads with --no-diffusion-compile-dynamic."
+            ),
+        )
         omni_config_group.add_argument(
             "--diffusers-load-kwargs",
             dest="diffusers_load_kwargs",