[Perf][Diffusion] Compile and prewarm VAE decode

Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>

.pre-commit-config.yaml

@@ -66,7 +66,10 @@ repos:
       entry: python tools/pre_commit/check_tts_adapter.py
       language: python
       pass_filenames: false
-      files: ^vllm_omni/entrypoints/openai/(serving_speech\.py|tts_adapters/__init__\.py)$
+      files: ^(vllm_omni/entrypoints/openai/(serving_speech\.py|tts_adapters/__init__\.py)|tools/pre_commit/check_tts_adapter\.py)$
+      # The budget-is-now-slack reminder is printed on a PASSING run, and
+      # pre-commit hides output from passing hooks unless the hook is verbose.
+      verbose: true
 
     # Keep `suggestion` last
     - id: suggestion

docs/user_guide/diffusion/attention_backends.md

@@ -283,23 +283,20 @@ blocks — use Ulysses SP (`ring_degree=1`).
 
 ### Which geometries run sparse
 
-Sparsity is only applied when the video segment is a **multiple of 128 rows**, where the row count
+MindIE-SD rf_v2 handles arbitrary video grids by spatially rearranging video first and, when
-is `latent_t × (height / 32) × (width / 32)`. Otherwise `rf_v2`'s block mask and the kernel's own
+necessary, promoting a real-video suffix to its always-kept prefix segment. It never pads: rf_v2
-tiling disagree on the block count, and the block straddling the seam mixes video and prefix rows
+does not consume an attention mask for padded keys. The remaining video rows stay sparse, while the
-that selection may then drop. Any resolution still runs — an unaligned geometry falls back to dense
+promoted rows remain visible to every query and key.
-attention and logs `RAINFUSION_ATTN staying dense` with the row count it computed — but it gets no
-speedup, so pick an aligned geometry when you want one.
 
-Alignment is necessary, not sufficient, for good quality. `rf_v2` groups video positions into 8x8
+For example, MiniMax-H3 at 1344x768 has grid `(62, 24, 42)`: its 62496 video rows leave a 32-row
-spatial tiles (two tiles fill one 128-row block), which is what makes a selected block a compact
+128-block residue, and its two non-8-wide latent columns form a 2928-row rearranged suffix.
-patch of the frame. When the latent `h` or `w` is not a multiple of 8, the leftover rows or columns
+MindIE-SD promotes 2976 rows, leaving 59520 sparse rows (`465 × 128`). This keeps the block mask
-are peeled off and appended as a flat run instead of being tiled, so they get pooled with spatially
+and kernel tiling aligned, at the cost of retaining more blocks.
-distant positions and selection can no longer rank them meaningfully. This is silent, and invisible
-to a `sparsity=0` check, because a fully populated mask does not care how blocks are grouped.
 
-"Multiple of 8" on the latent grid means **width and height that are multiples of 256**. Aligned
+8x8-aligned spatial grids remain preferable for performance. A latent `h` or `w` not divisible by 8
-resolutions off that grid still run sparse; they just lose more fidelity at the same `sparsity`,
+creates an always-kept residual suffix and lowers realized sparsity. "Multiple of 8" on the latent
-which you can buy back with `start_step`.
+grid means **width and height that are multiples of 256**. This path requires a MindIE-SD release
+that implements protected video tails for rf_v2.
 
 ## End-to-End Benchmark (BF16, sm_120 RTX Pro 6000 Blackwell)
 

docs/user_guide/examples/online_serving/diffusers_pipeline_adapter.md

@@ -99,7 +99,7 @@ a native vLLM-Omni pipeline for those cases.
 ### Attention Backends
 
 The diffusers backend converts
-[vLLM-Omni standard of attention backend setting](../../../docs/user_guide/diffusion/attention_backends.md)
+[vLLM-Omni standard of attention backend setting](../../diffusion/attention_backends.md)
 to [diffusers standard](https://huggingface.co/docs/diffusers/optimization/attention_backends#available-backends).
 
 Specifically for `FLASH_ATTN`, it will first attempt to use FlashAttention-3 and then FlashAttention-2.

tests/config/test_omni_config.py

@@ -751,6 +751,7 @@ def test_from_pipeline_config_routes_regional_compile_dynamic(tmp_path):
                 "  - stage_id: 0",
                 "    diffusion_compile_granularity: regional",
                 "    diffusion_compile_dynamic: false",
+                "    diffusion_compile_vae: false",
             ]
         )
     )
@@ -762,13 +763,16 @@ def test_from_pipeline_config_routes_regional_compile_dynamic(tmp_path):
         cli_overrides={
             "diffusion_compile_granularity": "full",
             "diffusion_compile_dynamic": True,
+            "diffusion_compile_vae": 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 configured_stage.diffusion_config.diffusion_compile_vae is False
     assert overridden_stage.diffusion_config.diffusion_compile_granularity == "full"
     assert overridden_stage.diffusion_config.diffusion_compile_dynamic is True
+    assert overridden_stage.diffusion_config.diffusion_compile_vae is True
 
 
 def test_structured_diffusion_config_rejects_non_boolean_compile_dynamic():
@@ -776,6 +780,11 @@ def test_structured_diffusion_config_rejects_non_boolean_compile_dynamic():
         omni_config_module._DiffusionConfigProjection(diffusion_compile_dynamic="false")
 
 
+def test_structured_diffusion_config_rejects_non_boolean_compile_vae():
+    with pytest.raises(ValidationError, match="diffusion_compile_vae"):
+        omni_config_module._DiffusionConfigProjection(diffusion_compile_vae="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")

tests/diffusion/attention/test_rainfusion_plan.py

@@ -5,10 +5,9 @@
 
 The grids here are the ones MiniMax-H3 FL2VA actually produces for an 8.7s clip:
 1280x768 gives a 62x24x40 latent grid whose 59520 video rows land on the 128-row
-kernel block, and 1344x768 gives 62x24x42 whose 62496 rows straddle it. Only the
+kernel block, and 1344x768 gives 62x24x42 whose 62496 rows do not. Both are
-first can run sparse: rf_v2 tiles the sequence as a whole but pools its block
+handed to MindIE-SD: its rf_v2 preprocessing promotes an irregular video tail
-mask from the video and prefix segments separately, so the two only agree on the
+to the always-kept prefix segment before generating the block mask.
-block count when the video segment ends on a boundary.
 """
 
 import dataclasses
@@ -57,10 +56,13 @@ def test_block_aligned_video_segment_runs_sparse():
     assert plan.latent_shape == list(ALIGNED_GRID)
 
 
-def test_misaligned_video_segment_stays_dense():
+def test_misaligned_video_segment_runs_sparse_with_updated_mindiesd():
-    # Realigning the sequence would mean padding it, and rf_v2 ignores attn_mask:
+    plan = make_impl()._resolve_plan(make_metadata(MISALIGNED_GRID))
-    # pad keys would take a share of every real query's softmax denominator.
+
-    assert make_impl()._resolve_plan(make_metadata(MISALIGNED_GRID)) is None
+    assert plan is not None
+    assert plan.prefix_len == PREFIX_ROWS
+    assert plan.used_len == PREFIX_ROWS + 62496
+    assert plan.latent_shape == list(MISALIGNED_GRID)
 
 
 @pytest.mark.parametrize("prefix_len", [1, 127, _BLOCK_SIZE, 710, 900])
@@ -155,3 +157,28 @@ def test_fully_populated_mask_reproduces_dense_attention():
     ).transpose(1, 2)
     error = (out.float() - reference.float()).abs().mean() / reference.float().abs().mean()
     assert error < 2e-3, f"mean relative error {error:.4%} against dense attention"
+
+
+@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="rf_v2 runs on Ascend NPU only.")
+def test_irregular_video_tail_reproduces_dense_attention_with_updated_mindiesd():
+    """vLLM forwards the full irregular grid; MindIE-SD owns its tail handling."""
+    torch.manual_seed(1)
+    prefix_len, grid = 384, (4, 16, 10)
+    video_len = grid[0] * grid[1] * grid[2]
+    used = prefix_len + video_len
+    heads, head_dim = 4, 128
+    q, k, v = (torch.randn(1, used, heads, head_dim, dtype=torch.bfloat16, device="npu") for _ in range(3))
+
+    impl = make_impl()
+    impl.rainfusion = dataclasses.replace(impl.rainfusion, sparsity=0.0)
+    plan = RainFusionPlan(prefix_len=prefix_len, used_len=used, latent_shape=list(grid))
+    out = impl._forward_sparse_npu(q, k, v, plan)
+
+    reference = torch.nn.functional.scaled_dot_product_attention(
+        q.transpose(1, 2),
+        k.transpose(1, 2),
+        v.transpose(1, 2),
+        scale=impl.softmax_scale,
+    ).transpose(1, 2)
+    error = (out.float() - reference.float()).abs().mean() / reference.float().abs().mean()
+    assert error < 2e-3, f"mean relative error {error:.4%} against dense attention"

tests/diffusion/cache/test_teacache_extractors.py

@@ -488,6 +488,14 @@ class _MiniMaxH3FakeAttention(nn.Module):
     class FakeBackend:
         supports_prefix_kv_slicing = False
 
+        # Mirror AttentionBackend.supports_packed_mask_free: H3's packed
+        # attention path calls attn_backend.supports_packed_mask_free() to decide
+        # whether to build the padding mask. Return False so the fake exercises
+        # the masked path (the abstract backend default).
+        @classmethod
+        def supports_packed_mask_free(cls) -> bool:
+            return False
+
     def __init__(self, **kwargs):
         del kwargs
         super().__init__()

tests/diffusion/test_compile.py

@@ -5,7 +5,7 @@ import pytest
 import torch.nn as nn
 
 import vllm_omni.diffusion.compile as compile_module
-from vllm_omni.diffusion.compile import regionally_compile
+from vllm_omni.diffusion.compile import compile_callable_with_fallback, regionally_compile
 
 pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
 
@@ -74,3 +74,50 @@ def test_regionally_compile_does_not_partially_mutate_on_setup_failure(monkeypat
         regionally_compile(model, dynamic=True)
 
     assert [block.forward.__func__ for block in model.transformer_blocks] == original_forwards
+
+
+def test_compile_callable_falls_back_once_after_lazy_failure(monkeypatch, caplog):
+    compiled_calls = 0
+    eager_calls = 0
+
+    def eager(value):
+        nonlocal eager_calls
+        eager_calls += 1
+        return value + 1
+
+    def compile_fn(fn, **kwargs):
+        assert fn is eager
+        assert kwargs == {"mode": "default"}
+
+        def compiled(value):
+            nonlocal compiled_calls
+            compiled_calls += 1
+            raise RuntimeError("lazy failure")
+
+        return compiled
+
+    monkeypatch.setattr(compile_module.torch, "compile", compile_fn)
+    callable_fn, activated = compile_callable_with_fallback(eager, label="test callable", mode="default")
+
+    assert activated
+    assert callable_fn(1) == 2
+    assert callable_fn(2) == 3
+    assert compiled_calls == 1
+    assert eager_calls == 2
+    assert "disabling compilation and retrying eagerly" in caplog.text
+
+
+def test_compile_callable_keeps_eager_on_setup_failure(monkeypatch, caplog):
+    def eager(value):
+        return value + 1
+
+    def compile_fn(*args, **kwargs):
+        raise RuntimeError("setup failure")
+
+    monkeypatch.setattr(compile_module.torch, "compile", compile_fn)
+    callable_fn, activated = compile_callable_with_fallback(eager, label="test callable")
+
+    assert not activated
+    assert callable_fn is eager
+    assert callable_fn(1) == 2
+    assert "continuing eagerly" in caplog.text

tests/diffusion/test_diffusion_engine.py

@@ -515,16 +515,19 @@ class TestDiffusionCompileConfig:
 
         assert config.diffusion_compile_granularity == "regional"
         assert config.diffusion_compile_dynamic is True
+        assert config.diffusion_compile_vae 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,
+            diffusion_compile_vae=False,
         )
 
         assert config.diffusion_compile_granularity == "full"
         assert config.diffusion_compile_dynamic is False
+        assert config.diffusion_compile_vae is False
 
     def test_config_rejects_invalid_compile_granularity(self) -> None:
         with pytest.raises(ValueError, match="diffusion_compile_granularity"):
@@ -534,6 +537,10 @@ class TestDiffusionCompileConfig:
         with pytest.raises(TypeError, match="diffusion_compile_dynamic"):
             OmniDiffusionConfig(model="test", diffusion_compile_dynamic="false")
 
+    def test_config_rejects_non_boolean_compile_vae(self) -> None:
+        with pytest.raises(TypeError, match="diffusion_compile_vae"):
+            OmniDiffusionConfig(model="test", diffusion_compile_vae="false")
+
     @pytest.mark.parametrize(
         "kwargs, feature",
         [

tests/diffusion/test_diffusion_model_runner.py

@@ -130,6 +130,23 @@ class _CompileTrackingModel:
         return self
 
 
+class _DecodeVAE(torch.nn.Module):
+    def decode(self, value):
+        return value + 1
+
+
+class _PipelineWithVAE(torch.nn.Module):
+    _dit_modules = []
+    _encoder_modules = []
+    _vae_modules = ["vae", "vae_alias"]
+    _resident_modules = []
+
+    def __init__(self):
+        super().__init__()
+        self.vae = _DecodeVAE()
+        self.vae_alias = self.vae
+
+
 def _make_request():
     sampling_params = SimpleNamespace(
         generator=None,
@@ -320,6 +337,31 @@ def test_compile_transformer_falls_back_after_synchronous_setup_failure(monkeypa
     assert "lazy compilation errors" in caplog.text
 
 
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_compile_vae_decodes_uses_discovery_and_deduplicates_aliases(monkeypatch):
+    runner = object.__new__(DiffusionModelRunner)
+    runner.pipeline = _PipelineWithVAE()
+    compile_calls = []
+
+    def compile_fn(fn, **kwargs):
+        compile_calls.append((fn, kwargs))
+        return lambda *args, **call_kwargs: fn(*args, **call_kwargs) + 10
+
+    monkeypatch.setattr(model_runner_module.torch, "compile", compile_fn)
+
+    DiffusionModelRunner._compile_vae_decodes(runner)
+
+    assert len(compile_calls) == 1
+    assert compile_calls[0][1] == {
+        "mode": "default",
+        "fullgraph": False,
+        "dynamic": None,
+    }
+    assert runner.pipeline.vae.decode(1) == 12
+    assert runner.pipeline.vae_alias.decode is runner.pipeline.vae.decode
+
+
 @pytest.mark.core_model
 @pytest.mark.cpu
 def test_execute_stepwise_streaming_returns_chunks_at_boundaries(monkeypatch):

tests/diffusion/test_diffusion_scheduler.py

@@ -971,6 +971,44 @@ class TestDiffusionEngine:
         admitted_request = engine.add_req_and_wait_for_response.call_args.args[0]
         assert admitted_request.request_id == "dummy_req_id"
 
+    def test_compile_warmup_repeats_through_normal_request_path(self, mocker: MockerFixture) -> None:
+        engine = DiffusionEngine.__new__(DiffusionEngine)
+        engine.od_config = SimpleNamespace(
+            model_class_name="mock_model",
+            diffusion_load_format="default",
+            diffusion_compile_vae=True,
+            enforce_eager=False,
+        )
+        engine.add_req_and_wait_for_response = mocker.Mock(
+            side_effect=[DiffusionOutput(output=None), DiffusionOutput(output=None)]
+        )
+
+        engine._dummy_run()
+
+        requests = [call.args[0] for call in engine.add_req_and_wait_for_response.call_args_list]
+        assert len(requests) == 2
+        assert requests[0].is_dummy_run()
+        assert not requests[1].is_dummy_run()
+        assert requests[1].prompt == {"prompt": "A detailed image."}
+        assert requests[1].sampling_params.need_kv_receive is False
+        assert requests[1].sampling_params.num_inference_steps == 1
+
+    def test_compile_real_path_warmup_failure_is_non_fatal(self, mocker: MockerFixture, caplog) -> None:
+        engine = DiffusionEngine.__new__(DiffusionEngine)
+        engine.od_config = SimpleNamespace(
+            model_class_name="mock_model",
+            diffusion_load_format="default",
+            diffusion_compile_vae=True,
+            enforce_eager=False,
+        )
+        engine.add_req_and_wait_for_response = mocker.Mock(
+            side_effect=[DiffusionOutput(output=None), RuntimeError("prewarm failure")]
+        )
+
+        engine._dummy_run()
+
+        assert "serving will continue" in caplog.text
+
 
 class TestStepScheduler:
     def setup_method(self) -> None:

tests/entrypoints/test_async_omni_diffusion_config.py

@@ -87,6 +87,7 @@ def test_default_stage_config_defaults_nullified_parallel_size_kwargs():
             "enforce_eager": None,
             "diffusion_compile_granularity": None,
             "diffusion_compile_dynamic": None,
+            "diffusion_compile_vae": None,
         }
     )[0]
 
@@ -98,6 +99,7 @@ def test_default_stage_config_defaults_nullified_parallel_size_kwargs():
     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
+    assert stage_cfg["engine_args"]["diffusion_compile_vae"] is True
 
 
 def test_default_stage_config_propagates_ulysses_mode():
@@ -315,7 +317,7 @@ def test_serve_cli_forwards_distributed_offload_residency():
 
 
 def test_serve_cli_accepts_diffusion_compile_controls():
-    """Ensure both compile controls reach the diffusion stage."""
+    """Ensure all compile controls reach the diffusion stage."""
     parser = TrackingArgumentParser()
     subparsers = parser.add_subparsers(dest="command")
     OmniServeCommand().subparser_init(subparsers)
@@ -328,6 +330,7 @@ def test_serve_cli_accepts_diffusion_compile_controls():
             "--diffusion-compile-granularity",
             "full",
             "--no-diffusion-compile-dynamic",
+            "--no-diffusion-compile-vae",
         ]
     )
 
@@ -336,8 +339,10 @@ def test_serve_cli_accepts_diffusion_compile_controls():
 
     assert args.diffusion_compile_granularity == "full"
     assert args.diffusion_compile_dynamic is False
+    assert args.diffusion_compile_vae is False
     assert stage_cfg["engine_args"]["diffusion_compile_granularity"] == "full"
     assert stage_cfg["engine_args"]["diffusion_compile_dynamic"] is False
+    assert stage_cfg["engine_args"]["diffusion_compile_vae"] is False
 
 
 def test_serve_cli_accepts_diffusion_attention_backend():

tools/pre_commit/check_tts_adapter.py

@@ -123,18 +123,29 @@ def _legacy_detectors(tree: ast.AST) -> list[tuple[int, str]]:
     return found
 
 
-def _check(label: str, path: Path, actual: int, budget: int, sample: list[str], budget_name: str) -> list[str]:
+def _check(
+    label: str,
+    path: Path,
+    actual: int,
+    budget: int,
+    sample: list[str],
+    budget_name: str,
+    notices: list[str],
+) -> list[str]:
     if actual > budget:
         listing = "\n".join(f"      {line}" for line in sample)
         return [
             f"{path}: {actual} {label} exceeds the budget of {budget}.\n{listing}\n\n{_GUIDANCE}\n",
         ]
     if actual < budget:
-        return [
+        # Progress is never a build failure. Removing branches is the goal, so a
-            f"{path}: {actual} {label}, down from a budget of {budget}. "
+        # PR that does it must not hand-edit this file to stay green -- nor break
-            f"Lower {budget_name} to {actual} in {Path(__file__).relative_to(REPO_ROOT)} "
+        # `main` for everyone else when it merges from a stale base.
-            f"so the ratchet holds the new ground.",
+        notices.append(
-        ]
+            f"{path}: {actual} {label}, below the budget of {budget}. "
+            f"Lower {budget_name} to {actual} in "
+            f"{Path(__file__).relative_to(REPO_ROOT)} to hold the new ground."
+        )
     return []
 
 
@@ -152,6 +163,7 @@ def main(argv: list[str] | None = None) -> int:
             print(f"check_tts_adapter: expected file not found: {path}", file=sys.stderr)
             return 1
 
+    notices: list[str] = []
     branches = _model_type_branches(ast.parse(serving_path.read_text()))
     detectors = _legacy_detectors(ast.parse(adapters_path.read_text()))
 
@@ -162,6 +174,7 @@ def main(argv: list[str] | None = None) -> int:
         MAX_MODEL_TYPE_BRANCHES,
         [f"line {lineno}: {src}" for lineno, src in branches],
         "MAX_MODEL_TYPE_BRANCHES",
+        notices,
     )
     errors += _check(
         "unmigrated model types in LEGACY_TTS_DETECTORS",
@@ -170,8 +183,11 @@ def main(argv: list[str] | None = None) -> int:
         MAX_LEGACY_DETECTORS,
         [f"line {lineno}: {name}" for lineno, name in detectors],
         "MAX_LEGACY_DETECTORS",
+        notices,
     )
 
+    for notice in notices:
+        print(f"check_tts_adapter: {notice}", file=sys.stderr)
     for error in errors:
         print(f"check_tts_adapter: {error}", file=sys.stderr)
     return 1 if errors else 0

vllm_omni/config/omni_config.py

@@ -570,6 +570,7 @@ class _DiffusionConfigProjection:
     pin_cpu_memory: bool = True
     diffusion_compile_granularity: Literal["regional", "full"] = "regional"
     diffusion_compile_dynamic: bool = Field(default=True, strict=True)
+    diffusion_compile_vae: bool = Field(default=True, strict=True)
     fa_deterministic: bool = False
     vae_use_slicing: bool = False
     vae_use_tiling: bool = False

vllm_omni/config/stage_config.py

@@ -394,6 +394,7 @@ class StageDeployConfig:
     # Diffusion execution, cache, and VAE behavior.
     diffusion_compile_granularity: str | None = None
     diffusion_compile_dynamic: bool | None = None
+    diffusion_compile_vae: bool | None = None
     fa_deterministic: bool | None = None
     cache_backend: str | None = None
     cache_config: dict[str, Any] | None = None

vllm_omni/diffusion/attention/backends/rainfusion_attn.py

@@ -128,9 +128,9 @@ class RainFusionAttentionImpl(AttentionImpl):
     whose extent the model publishes as ``AttentionMetadata.video_layout``. Every
     other case — warmup denoise steps, exempt layers, a layer that does not declare
     ``qkv_layout="BSND"``, sequences without a published video segment, video
-    segments too short to pay for block selection or not aligned to the kernel's
+    segments too short to pay for block selection — delegates to FlashAttention,
-    block size — delegates to FlashAttention, so a model can select this backend
+    so a model can select this backend unconditionally. MindIE-SD handles an
-    unconditionally.
+    irregular video tail internally, retaining it outside the sparse blocks.
     """
 
     def __init__(
@@ -292,25 +292,6 @@ class RainFusionAttentionImpl(AttentionImpl):
                 _MIN_VIDEO_BLOCKS,
             )
             return None
-        if video_len % _BLOCK_SIZE != 0:
-            # rf_v2 tiles the sequence into ceil((video + prefix) / 128) blocks but
-            # builds its block mask from the two segments pooled apart. The counts
-            # only agree when the video segment ends on a block boundary; otherwise
-            # a prefix block is forced on out of range, and the block straddling the
-            # seam holds real prefix rows that selection may drop. Padding the
-            # sequence to realign it is not an option because rf_v2 ignores
-            # attn_mask, so pad keys would take a share of every softmax denominator.
-            logger.warning_once(
-                "RAINFUSION_ATTN staying dense: latent grid %s gives %d video rows, which is not a "
-                "multiple of the %d-row kernel block. Choose a geometry whose latent_t * h * w is a "
-                "multiple of %d to enable sparsity.",
-                tuple(latent_shape),
-                video_len,
-                _BLOCK_SIZE,
-                _BLOCK_SIZE,
-            )
-            return None
-
         logger.info_once(
             "RAINFUSION_ATTN active: sparsity=%.2f, start_step=%d, exempt_layers=%d, "
             "latent_grid=%s, prefix_rows=%d, video_rows=%d. Realized sparsity is lower than nominal "

vllm_omni/diffusion/compile.py

@@ -1,5 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from collections.abc import Callable
 from typing import Any
 
 import torch
@@ -9,6 +10,59 @@ from vllm.logger import init_logger
 logger = init_logger(__name__)
 
 
+class _FallbackCompiledCallable:
+    """Run a compiled callable until it fails, then permanently use eager."""
+
+    def __init__(
+        self,
+        eager: Callable[..., Any],
+        compiled: Callable[..., Any],
+        *,
+        label: str,
+    ) -> None:
+        self._eager = eager
+        self._compiled: Callable[..., Any] | None = compiled
+        self._label = label
+
+    def __call__(self, *args: Any, **kwargs: Any) -> Any:
+        compiled = self._compiled
+        if compiled is None:
+            return self._eager(*args, **kwargs)
+        try:
+            return compiled(*args, **kwargs)
+        except Exception as exc:
+            self._compiled = None
+            logger.warning(
+                "%s failed during lazy torch.compile execution (%s); disabling compilation and retrying eagerly.",
+                self._label,
+                exc,
+            )
+            return self._eager(*args, **kwargs)
+
+
+def compile_callable_with_fallback(
+    fn: Callable[..., Any],
+    *,
+    label: str,
+    **compile_kwargs: Any,
+) -> tuple[Callable[..., Any], bool]:
+    """Compile ``fn`` with a fail-closed eager fallback.
+
+    The boolean return value reports whether lazy compilation was activated.
+    Synchronous setup failures return the original callable unchanged.
+    """
+    try:
+        compiled = torch.compile(fn, **compile_kwargs)
+    except Exception as exc:
+        logger.warning(
+            "%s torch.compile setup failed (%s); continuing eagerly.",
+            label,
+            exc,
+        )
+        return fn, False
+    return _FallbackCompiledCallable(fn, compiled, label=label), True
+
+
 def _matches_repeated_block(
     name: str,
     module: nn.Module,

vllm_omni/diffusion/data.py

@@ -726,6 +726,9 @@ class OmniDiffusionConfig:
     # provide its own setup_compile() implementation.
     diffusion_compile_granularity: str = "regional"
     diffusion_compile_dynamic: bool = True
+    # Compile discovered VAE ``decode`` callables whenever generic diffusion
+    # torch.compile is enabled. Runtime failures permanently fall back to eager.
+    diffusion_compile_vae: bool = True
 
     # Parallel weight loading (for faster diffusion model startup)
     enable_multithread_weight_load: bool = True
@@ -930,6 +933,8 @@ class OmniDiffusionConfig:
             )
         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}")
+        if not isinstance(self.diffusion_compile_vae, bool):
+            raise TypeError(f"diffusion_compile_vae must be a bool, got {type(self.diffusion_compile_vae)!r}")
         self.diffusion_kv_mode = parse_diffusion_kv_cache_mode(self.diffusion_kv_mode)
 
         if self.omni_kv_config is None:

vllm_omni/diffusion/diffusion_engine.py

@@ -56,6 +56,7 @@ if TYPE_CHECKING:
 logger = init_logger(__name__)
 
 _ASYNC_OUTPUT_TIMEOUT = 30.0  # seconds
+_VAE_COMPILE_WARMUP_REQUEST_ID = "vae_compile_warmup_req_id"
 
 __all__ = [
     "DiffusionEngine",
@@ -851,28 +852,26 @@ class DiffusionEngine:
         num_inference_steps = 1
         height = 512
         width = 512
-        prompt: OmniTextPrompt = {"prompt": "dummy run"}
 
         supports_image_input, supports_audio_input = supports_multimodal_input(self.od_config)
-        if supports_image_input:
-            # Provide a dummy image input if the model supports it
-            color_format = image_color_format(self.od_config.model_class_name)
-            dummy_image = PIL.Image.new(color_format, (width, height))
-            prompt.setdefault("multi_modal_data", {})["image"] = dummy_image
-
-        if supports_audio_input:
-            audio_sr = 16000
-            dummy_audio = np.random.randn(audio_sr * 2).astype(np.float32)
-            prompt.setdefault("multi_modal_data", {})["audio"] = dummy_audio
-
         num_frames = get_dummy_run_num_frames(self.od_config.model_class_name, supports_audio_input)
         if num_frames <= 0:
             logger.info("Skipping dummy warmup run (num_frames=0)")
             return
-        req = OmniDiffusionRequest(
+
-            prompt=prompt,
+        def build_request(request_id: str, prompt_text: str, *, real_path: bool) -> OmniDiffusionRequest:
-            request_id=DUMMY_DIFFUSION_REQUEST_ID,
+            prompt: OmniTextPrompt = {"prompt": prompt_text}
-            sampling_params=OmniDiffusionSamplingParams(
+            if supports_image_input:
+                color_format = image_color_format(self.od_config.model_class_name)
+                prompt.setdefault("multi_modal_data", {})["image"] = PIL.Image.new(
+                    color_format,
+                    (width, height),
+                )
+            if supports_audio_input:
+                audio_sr = 16000
+                prompt.setdefault("multi_modal_data", {})["audio"] = np.random.randn(audio_sr * 2).astype(np.float32)
+
+            sampling_params = OmniDiffusionSamplingParams(
                 height=height,
                 width=width,
                 num_inference_steps=num_inference_steps,
@@ -885,13 +884,47 @@ class DiffusionEngine:
                 # Disable CFG for warmup to avoid triggering CFG parallel
                 # validation when cfg_parallel_size > 1.
                 extra_args={"cfg_text_scale": 1.0, "cfg_img_scale": 1.0},
-            ),
+            )
-        )
+            if real_path:
+                # Internal pretraffic must exercise ordinary pipeline branches
+                # without waiting for KV produced by an upstream stage.
+                sampling_params.need_kv_receive = False
+            return OmniDiffusionRequest(
+                prompt=prompt,
+                request_id=request_id,
+                sampling_params=sampling_params,
+            )
+
+        req = build_request(DUMMY_DIFFUSION_REQUEST_ID, "dummy run", real_path=False)
         logger.info("dummy run to warm up the model")
         output = self.add_req_and_wait_for_response(req)
         if output.error:
             raise RuntimeError(f"Dummy run failed: {output.error}")
 
+        compile_real_path = bool(
+            getattr(self.od_config, "diffusion_compile_vae", False)
+            and not getattr(self.od_config, "enforce_eager", True)
+        )
+        if not compile_real_path:
+            return
+
+        logger.info("real request-path warmup for compiled VAE decode")
+        real_path_req = build_request(
+            _VAE_COMPILE_WARMUP_REQUEST_ID,
+            "A detailed image.",
+            real_path=True,
+        )
+        try:
+            real_path_output = self.add_req_and_wait_for_response(real_path_req)
+        except Exception as exc:
+            logger.warning("Real request-path compile warmup failed (%s); serving will continue.", exc)
+            return
+        if real_path_output.error:
+            logger.warning(
+                "Real request-path compile warmup returned an error (%s); serving will continue.",
+                real_path_output.error,
+            )
+
     def _submit_rpc(
         self,
         method: str,

vllm_omni/diffusion/worker/diffusion_model_runner.py

@@ -28,7 +28,7 @@ from vllm_omni.diffusion.cache.prompt_embed_cache import (
     resolve_prompt_embed_cache_config,
 )
 from vllm_omni.diffusion.cache.selector import get_cache_backend
-from vllm_omni.diffusion.compile import regionally_compile
+from vllm_omni.diffusion.compile import compile_callable_with_fallback, regionally_compile
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.forward_context import set_forward_context
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
@@ -40,6 +40,7 @@ from vllm_omni.diffusion.models.interface import (
     supports_step_execution,
 )
 from vllm_omni.diffusion.offloader import get_offload_backend
+from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery
 from vllm_omni.diffusion.registry import _NO_CACHE_ACCELERATION
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.sched.interface import DiffusionSchedulerOutput, KVPrefetchJob
@@ -188,6 +189,39 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
             compile_dynamic,
         )
 
+    def _compile_vae_decodes(self) -> None:
+        """Compile declared VAE decode callables with eager fallback."""
+        assert self.pipeline is not None
+        modules = ModuleDiscovery.discover(self.pipeline)
+        compiled_count = 0
+        for name, vae in zip(modules.vae_names, modules.vaes, strict=True):
+            decode = getattr(vae, "decode", None)
+            if not callable(decode):
+                logger.debug("Model runner: VAE component %s has no decode callable; skipping compile.", name)
+                continue
+
+            label = f"VAE component {name}.decode"
+            compiled_decode, activated = compile_callable_with_fallback(
+                decode,
+                label=label,
+                mode="default",
+                fullgraph=False,
+                dynamic=None,
+            )
+            if not activated:
+                continue
+            setattr(vae, "decode", compiled_decode)
+            compiled_count += 1
+
+        if compiled_count:
+            logger.info(
+                "Model runner: configured lazy torch.compile for %d VAE decode callable(s); "
+                "startup warmup will compile the real request path.",
+                compiled_count,
+            )
+        else:
+            logger.info("Model runner: no compile-compatible VAE decode callable was discovered.")
+
     def load_model(
         self,
         memory_pool_context_fn: Callable[[str], AbstractContextManager[Any]] | None = None,
@@ -277,7 +311,8 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
         # Apply torch.compile if not in eager mode
         if not self.od_config.enforce_eager:
             if current_omni_platform.supports_torch_inductor():
-                if hasattr(self.pipeline, "setup_compile"):
+                pipeline_owns_compile = hasattr(self.pipeline, "setup_compile")
+                if pipeline_owns_compile:
                     try:
                         self.pipeline.setup_compile()
                     except Exception as exc:
@@ -291,6 +326,9 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
                         transformer_attrs = ("transformer", "transformer_2")
                     for attr_name in transformer_attrs:
                         self._compile_transformer(attr_name)
+
+                if self.od_config.diffusion_compile_vae and not pipeline_owns_compile:
+                    self._compile_vae_decodes()
             else:
                 logger.warning(
                     "Model runner: Platform %s does not support torch inductor, skipping torch.compile.",

vllm_omni/engine/arg_utils.py

@@ -485,6 +485,7 @@ class OrchestratorArgs:
     diffusion_attention_config: str | None = None
     diffusion_compile_granularity: str | None = None
     diffusion_compile_dynamic: bool | None = None
+    diffusion_compile_vae: 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

@@ -1043,6 +1043,9 @@ class AsyncOmniEngine:
             "diffusion_compile_dynamic": (
                 True if kwargs.get("diffusion_compile_dynamic") is None else kwargs["diffusion_compile_dynamic"]
             ),
+            "diffusion_compile_vae": (
+                True if kwargs.get("diffusion_compile_vae") is None else kwargs["diffusion_compile_vae"]
+            ),
             "fa_deterministic": bool(kwargs.get("fa_deterministic", False)),
             "boundary_ratio": kwargs.get("boundary_ratio", None),
             "flow_shift": kwargs.get("flow_shift", None),

vllm_omni/entrypoints/cli/serve.py

@@ -446,6 +446,15 @@ class OmniServeCommand(CLISubcommand):
                 "Disable for fixed-shape workloads with --no-diffusion-compile-dynamic."
             ),
         )
+        omni_config_group.add_argument(
+            "--diffusion-compile-vae",
+            action=argparse.BooleanOptionalAction,
+            default=None,
+            help=(
+                "Compile discovered VAE decode callables when diffusion torch.compile is enabled. "
+                "Use --no-diffusion-compile-vae to keep VAE decode eager."
+            ),
+        )
         omni_config_group.add_argument(
             "--fa-deterministic",
             dest="fa_deterministic",