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

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

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