feat(diffusion): add request-gated VAE fast paths

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

examples/offline_inference/text_to_image/text_to_image.py

@@ -8,9 +8,7 @@ import time
 from pathlib import Path
 from typing import Any
 
-import numpy as np
 import torch
-from diffusers.utils import numpy_to_pil
 
 from vllm_omni.diffusion.data import logger
 from vllm_omni.diffusion.utils.image_output import extract_images_from_outputs
@@ -56,17 +54,6 @@ def parse_json_object(value: str, flag_name: str = "argument") -> dict[str, Any]
 parse_profiler_config = functools.partial(parse_json_object, flag_name="--profiler-config")
 
 
-def _normalize_images_for_save(images: list[Any]) -> list[Any]:
-    """Convert NumPy diffusion outputs to PIL images before saving."""
-    normalized = []
-    for image in images:
-        if isinstance(image, np.ndarray):
-            normalized.extend(numpy_to_pil(image))
-        else:
-            normalized.append(image)
-    return normalized
-
-
 def parse_args() -> argparse.Namespace:
     parser = argparse.ArgumentParser(description="Generate an image with supported diffusion models.")
     parser.add_argument(
@@ -632,7 +619,6 @@ def main():
 
     if not images:
         raise ValueError("No images found in request_output")
-    images = _normalize_images_for_save(images)
 
     output_path = Path(args.output)
     output_path.parent.mkdir(parents=True, exist_ok=True)

tests/dfx/reliability/test_reliability_qwen3_omni.py

@@ -315,32 +315,11 @@ def _assert_post_fault_health_terminal(host: str, port: int, *, scenario: str) -
     pytest.fail(f"[{scenario} health] no terminal post-fault health observed: {last_observation}")
 
 
-def _assert_server_alive_after_fault(host: str, port: int, *, scenario: str) -> None:
-    """#4285: a single-stage engine death must NOT take down the API server.
-
-    The process must stay up and keep answering /health (503 is acceptable —
-    errored but listening). A connection-refused means the server process
-    exited, which is the bug this fix closes.
-    """
-    deadline = time.monotonic() + 15.0
-    responded = False
-    while time.monotonic() < deadline:
-        try:
-            status, _ = get_health_raw(host, port, timeout_sec=5)
-            responded = True
-            assert status in (200, 503), f"[{scenario} alive] unexpected /health status={status}"
-        except Exception as exc:  # noqa: BLE001
-            if _looks_like_server_unreachable(exc):
-                pytest.fail(f"[{scenario} alive] server exited after single-stage OOM (#4285): {exc!r}")
-            raise
-        time.sleep(0.5)
-    assert responded, f"[{scenario} alive] /health never responded after fault"
-
-
 QWEN_PARAMS = create_reliability_omni_server_params(RELIABILITY_SCENARIOS, DEPLOY_CONFIGS_DIR)
 
 
 @pytest.mark.slow
+@pytest.mark.skip(reason="issue#4285")
 @hardware_test(res={"cuda": "H100"}, num_cards=2)
 @pytest.mark.parametrize("omni_server_function", QWEN_PARAMS, indirect=True)
 def test_reliability_fault_gpu_oom_error_contract_consistent_chat_speech(
@@ -407,11 +386,9 @@ def test_reliability_fault_gpu_oom_error_contract_consistent_chat_speech(
     assert "code" in chat_error, f"chat error lacks code field: {chat_error!r}"
     assert "code" in speech_error, f"speech error lacks code field: {speech_error!r}"
 
-    # #4285: a single dead stage must not take the whole API server down.
-    _assert_server_alive_after_fault(host, port, scenario="oom_contract")
-
 
 @pytest.mark.slow
+@pytest.mark.skip(reason="issue#4285")
 @hardware_test(res={"cuda": "H100"}, num_cards=2)
 @pytest.mark.parametrize("omni_server_function", QWEN_PARAMS, indirect=True)
 def test_reliability_fault_gpu_oom_chat_large_payload_failure(omni_server_function, openai_client_function) -> None:
@@ -454,6 +431,7 @@ def test_reliability_fault_gpu_oom_chat_large_payload_failure(omni_server_functi
 
 
 @pytest.mark.slow
+@pytest.mark.skip(reason="issue#4285")
 @hardware_test(res={"cuda": "H100"}, num_cards=2)
 @pytest.mark.parametrize("omni_server_function", QWEN_PARAMS, indirect=True)
 def test_reliability_fault_gpu_oom_concurrent_pressure_failure(omni_server_function, openai_client_function) -> None:
@@ -796,11 +774,7 @@ def test_reliability_fault_process_kill_tree_with_load_fast_fail_and_cleanup(
 
 
 @pytest.mark.slow
-@pytest.mark.skip(
-    reason="needs request-scoped OOM recovery: under transient pressure the stage "
-    "EngineCore still dies and is evicted (not respawned), so the engine does not "
-    "recover after the hog stops. Out of scope for the #4285 fault-isolation fix."
-)
+@pytest.mark.skip(reason="issue#4285")
 @hardware_test(res={"cuda": "H100"}, num_cards=2)
 @pytest.mark.parametrize("omni_server_function", QWEN_PARAMS, indirect=True)
 def test_reliability_fault_gpu_oom_state_converges_after_fault_removed(

tests/diffusion/test_diffusion_config_propagation.py

@@ -113,6 +113,14 @@ class TestCreateDefaultDiffusion:
         stages = StageConfigFactory.create_default_diffusion({"model": "x"})
         assert stages[0]["engine_args"]["cache_backend"] == "none"
 
+    def test_vae_decode_precision_roundtrip(self):
+        od = _roundtrip_diffusion_config(model="x", vae_decode_precision="bf16")
+        assert od.vae_decode_precision == "bf16"
+
+    def test_invalid_vae_decode_precision_is_rejected(self):
+        with pytest.raises(ValueError, match="vae_decode_precision"):
+            OmniDiffusionConfig(model="x", vae_decode_precision="int8")
+
     def test_single_gpu_default_devices(self):
         stages = StageConfigFactory.create_default_diffusion({"model": "x"})
         assert stages[0]["runtime"]["devices"] == "0"

tests/diffusion/test_diffusion_model_runner.py

@@ -512,6 +512,35 @@ def test_execute_model_accepts_bare_diffusion_output_from_single_request_pipelin
     assert runner.pipeline.last_req.num_reqs == 1
 
 
+@pytest.mark.core_model
+@pytest.mark.cpu
+@pytest.mark.parametrize(
+    ("quality", "expected_enabled"),
+    [(None, False), ("lossless", False), ("high", True)],
+)
+def test_execute_model_scopes_vae_fast_path_to_request_quality(monkeypatch, quality, expected_enabled):
+    runner = _make_runner(cache_backend=None, cache_backend_name="none")
+    req = _make_request()
+    req.sampling_params.quality = quality
+    gate_calls = []
+
+    @contextmanager
+    def _record_gate(pipeline, enabled):
+        gate_calls.append((pipeline, enabled, "enter"))
+        yield
+        gate_calls.append((pipeline, enabled, "exit"))
+
+    monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
+    monkeypatch.setattr(model_runner_module, "use_pipeline_vae_fast_path", _record_gate)
+
+    DiffusionModelRunner.execute_model(runner, req)
+
+    assert gate_calls == [
+        (runner.pipeline, expected_enabled, "enter"),
+        (runner.pipeline, expected_enabled, "exit"),
+    ]
+
+
 class _BatchPipeline:
     """Pipeline returning a configurable list of outputs from forward()."""
 
@@ -579,6 +608,18 @@ def test_execute_model_batch_rejects_output_count_mismatch(monkeypatch):
         DiffusionModelRunner.execute_model_batch(runner, sched, runner.od_config)
 
 
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_execute_model_batch_rejects_mixed_vae_quality() -> None:
+    runner = _make_batch_runner(_BatchPipeline(outputs=[DiffusionOutput(output="a"), DiffusionOutput(output="b")]))
+    sched = _make_scheduler_output(num_reqs=2)
+    sched.scheduled_new_reqs[0].req.sampling_params.quality = "lossless"
+    sched.scheduled_new_reqs[1].req.sampling_params.quality = "high"
+
+    with pytest.raises(RuntimeError, match="cannot mix request quality levels"):
+        DiffusionModelRunner.execute_model_batch(runner, sched, runner.od_config)
+
+
 @pytest.mark.core_model
 @pytest.mark.cpu
 def test_execute_model_batch_routes_one_output_per_request(monkeypatch):
@@ -728,6 +769,19 @@ def test_execute_model_runs_forward_after_kv_receive(monkeypatch):
 @pytest.mark.core_model
 @pytest.mark.cpu
 def test_load_model_clears_cache_backend_for_unsupported_pipeline(monkeypatch):
+    memory_pool_active = False
+    optimization_pool_states = []
+
+    @contextmanager
+    def _memory_pool_context(tag):
+        nonlocal memory_pool_active
+        assert tag == "weights"
+        memory_pool_active = True
+        try:
+            yield
+        finally:
+            memory_pool_active = False
+
     class _DummyLoader:
         def __init__(self, load_config, od_config=None):
             del load_config, od_config
@@ -775,16 +829,22 @@ def test_load_model_clears_cache_backend_for_unsupported_pipeline(monkeypatch):
     monkeypatch.setattr(model_runner_module, "LoadConfig", lambda: object())
     monkeypatch.setattr(model_runner_module, "DiffusersPipelineLoader", _DummyLoader)
     monkeypatch.setattr(model_runner_module, "DeviceMemoryProfiler", _DummyMemoryProfiler)
+    monkeypatch.setattr(
+        model_runner_module,
+        "optimize_pipeline_vaes",
+        lambda pipeline, od_config: optimization_pool_states.append(memory_pool_active),
+    )
     monkeypatch.setattr(model_runner_module, "get_offload_backend", lambda od_config, device: None)
     monkeypatch.setattr(
         model_runner_module, "get_cache_backend", lambda cache_backend, cache_config: dummy_cache_backend
     )
 
-    DiffusionModelRunner.load_model(runner)
+    DiffusionModelRunner.load_model(runner, memory_pool_context_fn=_memory_pool_context)
 
     assert runner.cache_backend is None
     assert runner.od_config.cache_backend is None
     assert dummy_cache_backend.enabled is False
+    assert optimization_pool_states == [True]
 
 
 @pytest.mark.core_model

tests/diffusion/test_diffusion_worker.py

@@ -102,6 +102,7 @@ class TestDiffusionWorkerSleep:
 
     def test_sleep_level_1(self, mocker: MockerFixture, mock_gpu_worker):
         """Test sleep mode level 1 (offload weights only)."""
+        clear_vae_caches = mocker.patch("vllm_omni.diffusion.vae_optimizations.clear_pipeline_vae_fast_path_caches")
         mock_allocator_class = patch_cumem_allocator(mocker)
         mock_platform = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
         mock_platform.get_free_memory.side_effect = [10 * 1024**3, 12 * 1024**3]
@@ -128,6 +129,7 @@ class TestDiffusionWorkerSleep:
 
         # Verify sleep was called with correct tags
         mock_allocator.sleep.assert_called_once_with(offload_tags=("weights",))
+        clear_vae_caches.assert_called_once_with(mock_gpu_worker.model_runner.pipeline)
         assert bool(result) is True
         # Verify buffers were NOT saved (level 1 doesn't save buffers)
         assert len(mock_gpu_worker._sleep_saved_buffers) == 0

tests/diffusion/vae_optimizations/test_vae_fastpaths.py

@@ -0,0 +1,300 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Correctness contracts for request-gated VAE decoder fast paths."""
+
+from __future__ import annotations
+
+import pytest
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from tests.helpers.mark import hardware_test
+from vllm_omni.diffusion.vae_optimizations.gate import (
+    VaeFastPathGate,
+    register_vae_fast_path_gate,
+    use_pipeline_vae_fast_path,
+    use_vae_fast_path,
+)
+
+pytestmark = [pytest.mark.diffusion]
+
+
+class _PipelineWithVae(nn.Module):
+    def __init__(self, vae: nn.Module) -> None:
+        super().__init__()
+        self.vae = vae
+
+
+class _DistributedExecutorConfig:
+    def __init__(self, parallel_mode: str) -> None:
+        self.parallel_mode = parallel_mode
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_vae_fast_path_gate_restores_nested_state() -> None:
+    vae = nn.Identity()
+    gate = VaeFastPathGate()
+    register_vae_fast_path_gate(vae, gate)
+
+    with use_vae_fast_path(vae, True):
+        assert gate.enabled
+        with use_vae_fast_path(vae, False):
+            assert not gate.enabled
+        assert gate.enabled
+
+    assert not gate.enabled
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_pipeline_vae_fast_path_controls_discovered_vae() -> None:
+    vae = nn.Identity()
+    pipeline = _PipelineWithVae(vae)
+    gate = VaeFastPathGate()
+    register_vae_fast_path_gate(vae, gate)
+
+    with use_pipeline_vae_fast_path(pipeline, True):
+        assert gate.enabled
+
+    assert not gate.enabled
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_image_group_norm_fast_path_uses_compile_safe_fallback(monkeypatch) -> None:
+    from vllm_omni.diffusion.vae_optimizations import image
+
+    gate = VaeFastPathGate()
+    gate.enabled = True
+    norm = image.FusedGroupNormSiLU(nn.GroupNorm(2, 4), gate)
+    inputs = torch.randn(1, 4, 3, 3)
+
+    monkeypatch.setattr(torch.compiler, "is_compiling", lambda: True)
+    monkeypatch.setattr(
+        image,
+        "group_norm_silu_4d",
+        lambda *args, **kwargs: pytest.fail("eager Triton kernel entered during compilation"),
+    )
+
+    actual = norm(inputs)
+    expected = F.silu(F.group_norm(inputs, 2, norm.weight, norm.bias, norm.eps))
+    torch.testing.assert_close(actual, expected)
+
+
+def _small_image_vae() -> nn.Module:
+    from diffusers.models.autoencoders import AutoencoderKL
+
+    return AutoencoderKL(
+        in_channels=3,
+        out_channels=3,
+        down_block_types=("DownEncoderBlock2D", "DownEncoderBlock2D"),
+        up_block_types=("UpDecoderBlock2D", "UpDecoderBlock2D"),
+        block_out_channels=(32, 32),
+        layers_per_block=1,
+        latent_channels=4,
+        norm_num_groups=8,
+        sample_size=16,
+    )
+
+
+def _small_wan_vae(vae_cls: type[nn.Module] | None = None) -> nn.Module:
+    if vae_cls is None:
+        from diffusers.models.autoencoders import AutoencoderKLWan
+
+        vae_cls = AutoencoderKLWan
+
+    return vae_cls(
+        base_dim=8,
+        decoder_base_dim=8,
+        z_dim=4,
+        dim_mult=[1, 2],
+        num_res_blocks=1,
+        attn_scales=[],
+        temperal_downsample=[False],
+        latents_mean=[0.0] * 4,
+        latents_std=[1.0] * 4,
+        scale_factor_temporal=1,
+        scale_factor_spatial=2,
+    )
+
+
+@pytest.mark.core_model
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@torch.no_grad()
+def test_image_vae_install_preserves_lossless_path_and_state_dict() -> None:
+    from vllm_omni.diffusion.vae_optimizations import clear_pipeline_vae_fast_path_caches
+    from vllm_omni.diffusion.vae_optimizations.image import (
+        FusedGroupNormSiLU,
+        FusedUpsample2xConv2d,
+        maybe_optimize_image_vae,
+    )
+
+    torch.manual_seed(0)
+    vae = _small_image_vae().to(device="cuda", dtype=torch.bfloat16).eval()
+    latents = torch.randn(1, 4, 8, 8, device="cuda", dtype=torch.bfloat16)
+    parameter_names = {name for name, _ in vae.named_parameters()}
+    state_dict = {name: value.clone() for name, value in vae.state_dict().items()}
+    reference = vae.decode(latents).sample
+
+    maybe_optimize_image_vae(vae)
+    maybe_optimize_image_vae(vae)
+
+    assert {name for name, _ in vae.named_parameters()} == parameter_names
+    vae.load_state_dict(state_dict, strict=True)
+    assert any(isinstance(module, FusedGroupNormSiLU) for module in vae.modules())
+    assert any(isinstance(module, FusedUpsample2xConv2d) for module in vae.modules())
+    assert torch.equal(vae.decode(latents).sample, reference)
+
+    with use_vae_fast_path(vae, True):
+        fast = vae.decode(latents).sample
+
+    torch.testing.assert_close(fast.float(), reference.float(), atol=0.1, rtol=0)
+    assert any(
+        buffer is not None
+        for module in vae.modules()
+        for name, buffer in module._buffers.items()
+        if name in {"_fused_weight", "_vllm_folded_value_weight", "_vllm_folded_value_bias"}
+    )
+    clear_pipeline_vae_fast_path_caches(_PipelineWithVae(vae))
+    assert all(
+        buffer is None
+        for module in vae.modules()
+        for name, buffer in module._buffers.items()
+        if name in {"_fused_weight", "_vllm_folded_value_weight", "_vllm_folded_value_bias"}
+    )
+    with use_vae_fast_path(vae, True):
+        torch.testing.assert_close(vae.decode(latents).sample.float(), reference.float(), atol=0.1, rtol=0)
+    assert torch.equal(vae.decode(latents).sample, reference)
+    assert vae.decoder.conv_in.weight.is_contiguous()
+
+
+@pytest.mark.core_model
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@torch.no_grad()
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+def test_channels_last_group_norm_silu_matches_reference(dtype: torch.dtype) -> None:
+    from vllm_omni.diffusion.vae_optimizations.triton_group_norm_silu import group_norm_silu_4d
+
+    torch.manual_seed(0)
+    x = torch.randn(1, 64, 16, 16, device="cuda", dtype=dtype).contiguous(memory_format=torch.channels_last)
+    weight = torch.randn(64, device="cuda", dtype=dtype)
+    bias = torch.randn(64, device="cuda", dtype=dtype)
+
+    actual = group_norm_silu_4d(x, weight, bias, num_groups=8, eps=1e-5)
+    expected = F.silu(F.group_norm(x, 8, weight, bias))
+
+    assert actual is not None
+    atol, rtol = (1e-5, 1e-5) if dtype == torch.float32 else (7e-2, 2e-2)
+    torch.testing.assert_close(actual, expected, atol=atol, rtol=rtol)
+
+
+@pytest.mark.core_model
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@torch.no_grad()
+@pytest.mark.parametrize(
+    "activation_dtype,affine_dtype,atol,rtol",
+    [
+        (torch.float32, torch.float32, 1e-5, 1e-5),
+        (torch.bfloat16, torch.float32, 1.5e-1, 3e-2),
+    ],
+)
+def test_wan_rmsnorm_silu_matches_reference(
+    activation_dtype: torch.dtype,
+    affine_dtype: torch.dtype,
+    atol: float,
+    rtol: float,
+) -> None:
+    from vllm_omni.diffusion.vae_optimizations.triton_wan_rmsnorm_silu import wan_rmsnorm_silu
+
+    torch.manual_seed(0)
+    x = torch.randn(1, 96, 3, 10, 14, device="cuda", dtype=activation_dtype).contiguous(
+        memory_format=torch.channels_last_3d
+    )
+    gamma = torch.randn(96, 1, 1, 1, device="cuda", dtype=affine_dtype)
+    bias = torch.randn_like(gamma)
+    normalized = F.normalize(x.float() if activation_dtype != torch.float32 else x, dim=1).to(activation_dtype)
+    expected = F.silu(normalized * 96**0.5 * gamma + bias)
+
+    actual = wan_rmsnorm_silu(x, gamma, bias)
+
+    assert actual is not None
+    assert actual.stride() == x.stride()
+    torch.testing.assert_close(actual, expected, atol=atol, rtol=rtol)
+
+
+@pytest.mark.core_model
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@torch.no_grad()
+def test_wan_vae_install_preserves_lossless_path_and_state_dict() -> None:
+    from vllm_omni.diffusion.vae_optimizations.wan import (
+        FusedWanRMSNormSiLU,
+        maybe_optimize_wan_vae,
+    )
+
+    torch.manual_seed(0)
+    vae = _small_wan_vae().cuda().eval()
+    latents = torch.randn(1, 4, 1, 8, 8, device="cuda")
+    parameter_names = {name for name, _ in vae.named_parameters()}
+    state_dict = {name: value.clone() for name, value in vae.state_dict().items()}
+    reference = vae.decode(latents).sample
+
+    maybe_optimize_wan_vae(vae)
+    maybe_optimize_wan_vae(vae)
+
+    assert {name for name, _ in vae.named_parameters()} == parameter_names
+    vae.load_state_dict(state_dict, strict=True)
+    assert any(isinstance(module, FusedWanRMSNormSiLU) for module in vae.modules())
+    assert torch.equal(vae.decode(latents).sample, reference)
+
+    with use_vae_fast_path(vae, True):
+        fast = vae.decode(latents).sample
+
+    torch.testing.assert_close(fast, reference, atol=5e-4, rtol=1e-4)
+    assert torch.equal(vae.decode(latents).sample, reference)
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+@pytest.mark.parametrize("parallel_mode", ["auto", "spatial_shard_height", "spatial_shard_width"])
+def test_wan_fast_path_skips_spatial_shard_mode(monkeypatch: pytest.MonkeyPatch, parallel_mode: str) -> None:
+    import vllm_omni.diffusion.vae_optimizations.wan as wan_optimizations
+
+    monkeypatch.setattr(wan_optimizations, "_HAS_TRITON", True)
+    vae = _small_wan_vae()
+    vae.distributed_executor = _DistributedExecutorConfig(parallel_mode)
+
+    wan_optimizations.maybe_optimize_wan_vae(vae)
+
+    assert not any(isinstance(module, wan_optimizations.FusedWanRMSNormSiLU) for module in vae.modules())
+
+
+@pytest.mark.core_model
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@torch.no_grad()
+@pytest.mark.parametrize(
+    ("precision", "decode_dtype"),
+    [("fp32", torch.float32), ("fp16", torch.float16), ("bf16", torch.bfloat16)],
+)
+def test_wan_decode_precision_keeps_encode_fp32(precision: str, decode_dtype: torch.dtype) -> None:
+    from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import DistributedAutoencoderKLWan
+    from vllm_omni.diffusion.vae_optimizations.precision import configure_wan_decode_precision
+
+    torch.manual_seed(0)
+    vae = _small_wan_vae(DistributedAutoencoderKLWan).to(device="cuda", dtype=torch.bfloat16).eval()
+    configure_wan_decode_precision(vae, precision)
+
+    assert next(vae.encoder.parameters()).dtype == torch.float32
+    assert next(vae.quant_conv.parameters()).dtype == torch.float32
+    assert next(vae.decoder.parameters()).dtype == decode_dtype
+    assert next(vae.post_quant_conv.parameters()).dtype == decode_dtype
+
+    _, grid_spec = vae.tile_split(torch.randn(1, 4, 1, 8, 8, device="cuda", dtype=decode_dtype))
+    assert grid_spec.output_dtype == decode_dtype
+
+    encoded = vae.encode(torch.randn(1, 3, 1, 16, 16, device="cuda", dtype=torch.bfloat16))
+    decoded = vae.decode(torch.randn(1, 4, 1, 8, 8, device="cuda", dtype=torch.float32))
+
+    assert encoded.latent_dist.parameters.dtype == torch.float32
+    assert decoded.sample.dtype == decode_dtype

tests/engine/test_cfg_companion_lifecycle.py

@@ -83,11 +83,6 @@ class _FakePool:
         self.stage_client = SimpleNamespace(requires_multimodal_data=False, custom_process_input_func=None)
         self.aborted: list[list[str]] = []
 
-    def live_replica_ids(self) -> list[int]:
-        # Single always-live replica: these tests cover CFG bundling, not
-        # replica death, and the forward path fast-fails on an empty pool.
-        return [0]
-
     async def abort_requests(self, request_ids):
         self.aborted.append(list(request_ids))
 

tests/engine/test_orchestrator.py

@@ -74,6 +74,8 @@ async def test_engine_dead_broadcasts_fatal_to_rpc_waiters(monkeypatch: pytest.M
         rpc_async_queue=rpc_queue,
         stage_pools=[],
     )
+    orchestrator._fatal_error = "stage engine died"
+    orchestrator._fatal_error_stage_id = 2
 
     async def wait_for_requests() -> None:
         await asyncio.Event().wait()
@@ -89,7 +91,8 @@ async def test_engine_dead_broadcasts_fatal_to_rpc_waiters(monkeypatch: pytest.M
     fatal = rpc_queue.get_nowait()
     assert isinstance(fatal, ErrorMessage)
     assert fatal.fatal is True
-    assert "stage engine died" in fatal.error
+    assert fatal.error == "stage engine died"
+    assert fatal.stage_id == 2
 
 
 @dataclass

tests/engine/test_orchestrator_error_handling.py

@@ -3,7 +3,7 @@
 """Tests for error propagation paths within the Orchestrator.
 
 Covers:
-- EngineDeadError from an LLM stage poll → fatal error broadcast + replica eviction
+- EngineDeadError from an LLM stage poll → fatal error broadcast + shutdown
 - Diffusion stage error output (OmniRequestOutput.from_error) → routed correctly
 """
 
@@ -14,30 +14,17 @@ import queue
 import time
 from types import SimpleNamespace
 
-import janus
 import pytest
 from vllm.v1.engine.exceptions import EngineDeadError
 
-from vllm_omni.engine.messages import (
-    AddCompanionRequestMessage,
-    EngineQueueMessage,
-    ErrorMessage,
-    ShutdownRequestMessage,
-    StageSubmissionMessage,
-)
-from vllm_omni.engine.orchestrator import Orchestrator, OrchestratorRequestState
-from vllm_omni.engine.stage_pool import StageUnavailableError
+from vllm_omni.engine.messages import EngineQueueMessage, ErrorMessage, ShutdownRequestMessage
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
 
 from .test_orchestrator import (
-    FakeOutputProcessor,
     FakeStageClient,
     OrchestratorFixture,
     _build_harness,
-    _build_request_output,
-    _build_stage_pools,
-    _engine_core_outputs,
     _enqueue_add_request,
     _wait_for,
 )
@@ -63,22 +50,6 @@ async def _get_any_output_message(fixture: OrchestratorFixture, *, timeout: floa
             await asyncio.sleep(0.01)
 
 
-async def _wait_for_error_message(
-    fixture: OrchestratorFixture,
-    *,
-    request_id: str | None = None,
-    max_messages: int = 50,
-) -> ErrorMessage:
-    """Drain output messages until the ErrorMessage arrives, skipping the
-    non-error outputs (e.g. stage metrics, forwarded stage results) that a
-    driven pipeline can interleave ahead of the failure."""
-    for _ in range(max_messages):
-        msg = await _get_any_output_message(fixture)
-        if isinstance(msg, ErrorMessage) and (request_id is None or msg.request_id == request_id):
-            return msg
-    raise AssertionError(f"No ErrorMessage for req={request_id} within {max_messages} messages")
-
-
 @pytest.fixture
 def orchestrator_factory():
     fixtures: list[OrchestratorFixture] = []
@@ -101,14 +72,32 @@ def orchestrator_factory():
 # ───────── EngineDeadError from LLM stage poll ─────────
 
 
-@pytest.mark.asyncio
-async def test_engine_dead_error_evicts_replica_and_keeps_running(orchestrator_factory) -> None:
-    """When a stage replica raises EngineDeadError during poll, the orchestrator must:
-    1. Enqueue a fatal error message for each affected request
-    2. Evict the dead replica and keep running so the API server stays up
-       (per-replica fault isolation, #4285)
+class FakeDeadLLMStageClient(FakeStageClient):
+    """LLM stage client that raises EngineDeadError once a request is in-flight.
+
+    Raising on the first poll races with ``add_request``: under suite load the
+    orchestrator can die and drain before the test enqueues anything, so no
+    fatal ErrorMessage is ever published. Keep returning empty outputs until
+    ``add_request_async`` has been called, then fail on the next poll.
     """
-    stage0 = _DieOnDemandLLMStageClient(stage_type="llm", final_output=True)
+
+    async def get_output_async(self):
+        if not self.add_request_calls:
+            return SimpleNamespace(
+                outputs=[],
+                scheduler_stats=None,
+                finished_requests=None,
+            )
+        raise EngineDeadError("Stage-0 engine core is dead")
+
+
+@pytest.mark.asyncio
+async def test_engine_dead_error_broadcasts_fatal_and_shuts_down(orchestrator_factory) -> None:
+    """When a stage raises EngineDeadError during poll, the orchestrator must:
+    1. Enqueue a fatal error message for each affected request
+    2. Shut itself down (thread exits)
+    """
+    stage0 = FakeDeadLLMStageClient(stage_type="llm", final_output=True)
     orchestrator_fixture = orchestrator_factory([stage0])
     request = SimpleNamespace(request_id="req-dead", prompt_token_ids=[1, 2])
 
@@ -121,12 +110,6 @@ async def test_engine_dead_error_evicts_replica_and_keeps_running(orchestrator_f
             sampling_params_list=[_sampling_params()],
             final_stage_id=0,
         )
-        # Let the request dispatch before the replica dies; a replica dead on
-        # its very first poll can be evicted before dispatch, which routes the
-        # failure through the add-request guard instead (that ordering is
-        # covered by test_add_request_to_dead_stage_fails_request_not_server).
-        await _wait_for(lambda: len(stage0.add_request_calls) == 1)
-        stage0.die = True
 
         # Collect the fatal error message.
         msg = await _get_any_output_message(orchestrator_fixture)
@@ -135,17 +118,13 @@ async def test_engine_dead_error_evicts_replica_and_keeps_running(orchestrator_f
         assert msg.type == "error"
         assert msg.fatal is True
         assert msg.request_id == "req-dead"
-        assert "Stage-0 replica-0 is dead" in msg.error
+        assert "Stage-0 engine core is dead" in msg.error
 
-        # The dead replica is evicted; the orchestrator keeps running.
-        pool = orchestrator_fixture.orchestrator.stage_pools[0]
-        deadline = time.monotonic() + 5.0
-        while pool.live_replica_ids() and time.monotonic() < deadline:
-            await asyncio.sleep(0.02)
-        assert pool.live_replica_ids() == []
-        assert orchestrator_fixture.thread.is_alive()
+        # The orchestrator thread should exit after the fatal error.
+        orchestrator_fixture.thread.join(timeout=5)
+        assert not orchestrator_fixture.thread.is_alive()
 
-        # Affected request state should be cleaned up.
+        # Request state should be cleaned up.
         assert "req-dead" not in orchestrator_fixture.orchestrator.request_states
     finally:
         if orchestrator_fixture.thread.is_alive():
@@ -153,596 +132,6 @@ async def test_engine_dead_error_evicts_replica_and_keeps_running(orchestrator_f
             orchestrator_fixture.thread.join(timeout=5)
 
 
-class _DieOnDemandLLMStageClient(FakeStageClient):
-    """LLM stage replica that raises EngineDeadError on poll once ``die`` is set."""
-
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        self.die = False
-
-    async def get_output_async(self):
-        if self.die:
-            raise EngineDeadError(f"Stage-{self.stage_id} replica-{self.replica_id} is dead")
-        return await super().get_output_async()
-
-
-class _DieOnDemandDiffusionStageClient(FakeStageClient):
-    """Diffusion stage replica that raises EngineDeadError on poll once ``die`` is set."""
-
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        self.die = False
-
-    def get_diffusion_output_nowait(self):
-        if self.die:
-            raise EngineDeadError(f"Stage-{self.stage_id} diffusion replica-{self.replica_id} is dead")
-        return super().get_diffusion_output_nowait()
-
-
-class _UnavailableDiffusionStageClient(FakeStageClient):
-    """Diffusion replica that stays live (so the pool keeps stage_type) but whose
-    dispatch fails with StageUnavailableError, as a downstream prewarm would when
-    its stage has no capacity."""
-
-    async def add_request_async(self, *args, **kwargs) -> None:
-        raise StageUnavailableError(f"stage {self.stage_id} has no live replica")
-
-
-@pytest.mark.asyncio
-async def test_engine_dead_error_fails_only_dead_replica_requests(orchestrator_factory) -> None:
-    """Multi-replica stage: one replica dying must fail only the requests bound
-    to it. Requests bound to a surviving replica keep running, and the stage
-    (still having a live replica) is not torn down (per-replica fault isolation,
-    #4285).
-    """
-    r0 = _DieOnDemandLLMStageClient(stage_type="llm", final_output=True)
-    r1 = FakeStageClient(stage_type="llm", final_output=True)
-    stage_pools = _build_stage_pools([[r0, r1]])
-    orchestrator_fixture = orchestrator_factory([], stage_pools=stage_pools)
-    pool = orchestrator_fixture.orchestrator.stage_pools[0]
-
-    try:
-        # req-0 → replica 0 (round-robin starts at 0)
-        await _enqueue_add_request(
-            orchestrator_fixture,
-            request_id="req-0",
-            prompt=SimpleNamespace(request_id="req-0", prompt_token_ids=[1, 2]),
-            original_prompt={"prompt": "a"},
-            sampling_params_list=[_sampling_params()],
-            final_stage_id=0,
-        )
-        await _wait_for(lambda: len(r0.add_request_calls) == 1)
-
-        # req-1 → replica 1
-        await _enqueue_add_request(
-            orchestrator_fixture,
-            request_id="req-1",
-            prompt=SimpleNamespace(request_id="req-1", prompt_token_ids=[3, 4]),
-            original_prompt={"prompt": "b"},
-            sampling_params_list=[_sampling_params()],
-            final_stage_id=0,
-        )
-        await _wait_for(lambda: len(r1.add_request_calls) == 1)
-        assert pool.get_bound_replica_id("req-0") == 0
-        assert pool.get_bound_replica_id("req-1") == 1
-
-        # Kill replica 0 on its next poll.
-        r0.die = True
-
-        msg = await _get_any_output_message(orchestrator_fixture)
-        assert isinstance(msg, ErrorMessage)
-        assert msg.fatal is True
-        assert msg.request_id == "req-0"
-
-        # Replica 0 evicted, replica 1 still serving; req-1 untouched.
-        await _wait_for(lambda: pool.live_replica_ids() == [1])
-        assert orchestrator_fixture.thread.is_alive()
-        assert "req-0" not in orchestrator_fixture.orchestrator.request_states
-        assert "req-1" in orchestrator_fixture.orchestrator.request_states
-    finally:
-        if orchestrator_fixture.thread.is_alive():
-            orchestrator_fixture.request_sync_q.put_nowait(ShutdownRequestMessage())
-            orchestrator_fixture.thread.join(timeout=5)
-
-
-@pytest.mark.asyncio
-async def test_forward_to_dead_downstream_stage_fails_request_not_server(orchestrator_factory) -> None:
-    """A request mid-pipeline whose downstream stage has lost all replicas must
-    fail that request, not tear down the orchestrator. Forwarding to a stage with
-    no live replica is handled gracefully (per-replica fault isolation, #4285).
-    """
-    stage0 = FakeStageClient(stage_type="llm", final_output=False, next_inputs=[{"prompt_token_ids": [7, 8]}])
-    stage1 = _DieOnDemandLLMStageClient(stage_type="llm", final_output=True)
-    proc0 = FakeOutputProcessor(request_outputs=[_build_request_output("req-x", token_ids=[3], finished=True)])
-    proc1 = FakeOutputProcessor()
-    stage_pools = _build_stage_pools([[stage0], [stage1]], output_processors=[proc0, proc1])
-    orchestrator_fixture = orchestrator_factory([], stage_pools=stage_pools)
-    pool1 = orchestrator_fixture.orchestrator.stage_pools[1]
-
-    try:
-        await _enqueue_add_request(
-            orchestrator_fixture,
-            request_id="req-x",
-            prompt=SimpleNamespace(request_id="req-x", prompt_token_ids=[1, 2]),
-            original_prompt={"prompt": "hello"},
-            sampling_params_list=[_sampling_params(), _sampling_params()],
-            final_stage_id=1,
-        )
-        await _wait_for(lambda: len(stage0.add_request_calls) == 1)
-
-        # Kill stage-1's only replica before stage-0 output is forwarded.
-        stage1.die = True
-        await _wait_for(lambda: pool1.live_replica_ids() == [])
-
-        # Completing stage 0 triggers a forward to the now-empty stage 1.
-        stage0.push_engine_core_outputs(_engine_core_outputs("s0-raw", 1.0))
-
-        # Completing stage 0 also emits its normal outputs; skip past them to
-        # the downstream-forward failure.
-        error_msg = await _wait_for_error_message(orchestrator_fixture, request_id="req-x")
-        assert error_msg.fatal is True
-        assert error_msg.stage_id == 1
-
-        assert orchestrator_fixture.thread.is_alive()
-        assert "req-x" not in orchestrator_fixture.orchestrator.request_states
-    finally:
-        if orchestrator_fixture.thread.is_alive():
-            orchestrator_fixture.request_sync_q.put_nowait(ShutdownRequestMessage())
-            orchestrator_fixture.thread.join(timeout=5)
-
-
-@pytest.mark.asyncio
-async def test_add_request_to_dead_stage_fails_request_not_server(orchestrator_factory) -> None:
-    """A new request entering a stage that already lost all replicas fails that
-    request instead of raising out of the request handler and tearing the server
-    down. The guard runs before request state is registered, so nothing leaks.
-    """
-    r0 = _DieOnDemandLLMStageClient(stage_type="llm", final_output=True)
-    stage_pools = _build_stage_pools([[r0]])
-    orchestrator_fixture = orchestrator_factory([], stage_pools=stage_pools)
-    pool = orchestrator_fixture.orchestrator.stage_pools[0]
-
-    try:
-        # Send a first request so the replica is polled and dies → evicted.
-        await _enqueue_add_request(
-            orchestrator_fixture,
-            request_id="req-0",
-            prompt=SimpleNamespace(request_id="req-0", prompt_token_ids=[1, 2]),
-            original_prompt={"prompt": "a"},
-            sampling_params_list=[_sampling_params()],
-            final_stage_id=0,
-        )
-        r0.die = True
-        await _wait_for(lambda: pool.live_replica_ids() == [])
-
-        # A new request now enters a stage with no live replica.
-        await _enqueue_add_request(
-            orchestrator_fixture,
-            request_id="req-after",
-            prompt=SimpleNamespace(request_id="req-after", prompt_token_ids=[3, 4]),
-            original_prompt={"prompt": "b"},
-            sampling_params_list=[_sampling_params()],
-            final_stage_id=0,
-        )
-
-        error_msg = await _wait_for_error_message(orchestrator_fixture, request_id="req-after")
-        assert error_msg.fatal is True
-        assert error_msg.stage_id == 0
-
-        assert orchestrator_fixture.thread.is_alive()
-        # Guard runs before registration, so the failed request never enters state.
-        assert "req-after" not in orchestrator_fixture.orchestrator.request_states
-    finally:
-        if orchestrator_fixture.thread.is_alive():
-            orchestrator_fixture.request_sync_q.put_nowait(ShutdownRequestMessage())
-            orchestrator_fixture.thread.join(timeout=5)
-
-
-@pytest.mark.asyncio
-async def test_diffusion_replica_death_on_poll_keeps_server(orchestrator_factory) -> None:
-    """A diffusion stage replica raising EngineDeadError during poll must be
-    evicted while the orchestrator keeps serving. The per-replica catch has to
-    cover the diffusion polling path, not just the LLM one (#4285).
-    """
-    stage0 = _DieOnDemandDiffusionStageClient(stage_type="diffusion", final_output=True, final_output_type="image")
-    stage_pools = _build_stage_pools([[stage0]])
-    orchestrator_fixture = orchestrator_factory([], stage_pools=stage_pools)
-    pool = orchestrator_fixture.orchestrator.stage_pools[0]
-
-    try:
-        await _enqueue_add_request(
-            orchestrator_fixture,
-            request_id="req-dit",
-            prompt={"prompt": "a cat"},
-            original_prompt={"prompt": "a cat"},
-            sampling_params_list=[OmniDiffusionSamplingParams()],
-            final_stage_id=0,
-        )
-        await _wait_for(lambda: len(stage0.add_request_calls) == 1)
-
-        stage0.die = True
-
-        error_msg = await _wait_for_error_message(orchestrator_fixture, request_id="req-dit")
-        assert error_msg.fatal is True
-        assert error_msg.stage_id == 0
-
-        await _wait_for(lambda: pool.live_replica_ids() == [])
-        assert orchestrator_fixture.thread.is_alive()
-        assert "req-dit" not in orchestrator_fixture.orchestrator.request_states
-    finally:
-        if orchestrator_fixture.thread.is_alive():
-            orchestrator_fixture.request_sync_q.put_nowait(ShutdownRequestMessage())
-            orchestrator_fixture.thread.join(timeout=5)
-
-
-@pytest.mark.asyncio
-async def test_async_chunk_prewarm_failure_reports_failing_downstream_stage(orchestrator_factory) -> None:
-    """When an async-chunk prewarm dispatch to a downstream stage fails, the
-    failure must be attributed to that stage, not stage 0 (the prewarm's own
-    dispatch stage). Regression for the guard using stage_id=0.
-    """
-    stage0 = FakeStageClient(stage_type="llm", final_output=False)
-    stage1 = FakeStageClient(stage_type="diffusion", final_output=False, final_output_type="image")
-    stage2 = _UnavailableDiffusionStageClient(stage_type="diffusion", final_output=True, final_output_type="image")
-    stage_pools = _build_stage_pools([[stage0], [stage1], [stage2]])
-    orchestrator_fixture = orchestrator_factory([], stage_pools=stage_pools, async_chunk=True)
-
-    try:
-        await _enqueue_add_request(
-            orchestrator_fixture,
-            request_id="req-x",
-            prompt=SimpleNamespace(request_id="req-x", prompt_token_ids=[1, 2, 3]),
-            original_prompt={"prompt": "a cat"},
-            sampling_params_list=[_sampling_params(), OmniDiffusionSamplingParams(), OmniDiffusionSamplingParams()],
-            final_stage_id=2,
-        )
-
-        error_msg = await _wait_for_error_message(orchestrator_fixture, request_id="req-x")
-        assert error_msg.fatal is True
-        assert error_msg.stage_id == 2
-        assert "Stage-2" in error_msg.error
-
-        # Stage 1 prewarm ran before the stage-2 failure; server stays up.
-        assert len(stage1.add_request_calls) == 1
-        assert orchestrator_fixture.thread.is_alive()
-    finally:
-        if orchestrator_fixture.thread.is_alive():
-            orchestrator_fixture.request_sync_q.put_nowait(ShutdownRequestMessage())
-            orchestrator_fixture.thread.join(timeout=5)
-
-
-# ───────── Direct unit tests for the fault-isolation helpers ─────────
-
-
-def _build_bare_orchestrator(stage_pools) -> tuple[Orchestrator, tuple[janus.Queue, ...]]:
-    """Construct an Orchestrator on the current event loop without running it.
-
-    Lets tests await the fault-isolation helpers directly instead of driving
-    them through the orchestration loop.
-    """
-    queues = (janus.Queue(), janus.Queue(), janus.Queue())
-    orchestrator = Orchestrator(
-        request_async_queue=queues[0].async_q,
-        output_async_queue=queues[1].async_q,
-        rpc_async_queue=queues[2].async_q,
-        stage_pools=stage_pools,
-    )
-    return orchestrator, queues
-
-
-def _register_request(orchestrator: Orchestrator, request_id: str, *, submitted_stage: int | None = 0) -> None:
-    state = OrchestratorRequestState(request_id=request_id)
-    if submitted_stage is not None:
-        state.stage_submit_ts[submitted_stage] = time.time()
-    orchestrator.request_states[request_id] = state
-
-
-@pytest.mark.asyncio
-async def test_handle_dead_replica_fails_only_bound_requests() -> None:
-    """Direct helper test: with a surviving replica, only requests bound to the
-    dead replica fail; the survivor's requests and bindings are untouched.
-    """
-    r0 = FakeStageClient(stage_type="llm", final_output=True)
-    r1 = FakeStageClient(stage_type="llm", final_output=True)
-    orchestrator, queues = _build_bare_orchestrator(_build_stage_pools([[r0, r1]]))
-    try:
-        pool = orchestrator.stage_pools[0]
-        _register_request(orchestrator, "req-0")
-        _register_request(orchestrator, "req-1")
-        assert pool.select_replica_id("req-0") == 0
-        assert pool.select_replica_id("req-1") == 1
-
-        await orchestrator._handle_dead_replica(0, 0, EngineDeadError("replica-0 dead"))
-
-        assert pool.live_replica_ids() == [1]
-        assert r0.shutdown_calls == 1
-        msg = queues[1].async_q.get_nowait()
-        assert isinstance(msg, ErrorMessage)
-        assert msg.fatal is True
-        assert msg.request_id == "req-0"
-        assert msg.stage_id == 0
-        assert "replica-0 dead" in msg.error
-        assert queues[1].async_q.empty()
-        assert "req-0" not in orchestrator.request_states
-        assert "req-1" in orchestrator.request_states
-        assert pool.get_bound_replica_id("req-0") is None
-        assert pool.get_bound_replica_id("req-1") == 1
-    finally:
-        for q in queues:
-            q.close()
-
-
-@pytest.mark.asyncio
-async def test_handle_dead_replica_last_replica_fails_all_stage_requests() -> None:
-    """Direct helper test: when the dead replica was the stage's last one, every
-    request routed through the stage fails — bound or not — while requests that
-    never entered the stage survive.
-    """
-    orchestrator, queues = _build_bare_orchestrator(_build_stage_pools([[FakeStageClient(stage_type="llm")]]))
-    try:
-        pool = orchestrator.stage_pools[0]
-        _register_request(orchestrator, "req-bound")
-        _register_request(orchestrator, "req-unbound")
-        _register_request(orchestrator, "req-elsewhere", submitted_stage=None)
-        assert pool.select_replica_id("req-bound") == 0
-
-        await orchestrator._handle_dead_replica(0, 0, EngineDeadError("last replica dead"))
-
-        assert pool.live_replica_ids() == []
-        failed = {queues[1].async_q.get_nowait().request_id for _ in range(2)}
-        assert failed == {"req-bound", "req-unbound"}
-        assert queues[1].async_q.empty()
-        assert set(orchestrator.request_states) == {"req-elsewhere"}
-    finally:
-        for q in queues:
-            q.close()
-
-
-@pytest.mark.asyncio
-async def test_fail_request_dead_stage_cleans_registered_request() -> None:
-    """Direct helper test: a registered request failed against a dead stage gets
-    a fatal ErrorMessage, its state and bindings are released, and its
-    in-flight work on live replicas is aborted.
-    """
-    r0 = FakeStageClient(stage_type="llm")
-    orchestrator, queues = _build_bare_orchestrator(_build_stage_pools([[r0]]))
-    try:
-        pool = orchestrator.stage_pools[0]
-        _register_request(orchestrator, "req-a")
-        assert pool.select_replica_id("req-a") == 0
-
-        await orchestrator._fail_request_dead_stage("req-a", 0)
-
-        msg = queues[1].async_q.get_nowait()
-        assert isinstance(msg, ErrorMessage)
-        assert msg.error == "Stage-0 has no live replica"
-        assert msg.fatal is True
-        assert msg.request_id == "req-a"
-        assert msg.stage_id == 0
-        assert "req-a" not in orchestrator.request_states
-        assert pool.get_bound_replica_id("req-a") is None
-        # abort=True: the failed request's work on the still-live replica is
-        # stopped, not left to run to completion.
-        assert r0.abort_calls == [["req-a"]]
-    finally:
-        for q in queues:
-            q.close()
-
-
-@pytest.mark.asyncio
-async def test_streaming_update_racing_eviction_fails_request_not_server() -> None:
-    """A streaming_update can find its request state still present while the
-    stage has already lost all replicas (_handle_dead_replica awaits between
-    eviction and cleanup). Submitting then raises from replica selection; the
-    handler must fail that request instead of letting the raise escape
-    _request_handler and shut the server down (#4285).
-    """
-    orchestrator, queues = _build_bare_orchestrator(_build_stage_pools([[FakeStageClient(stage_type="llm")]]))
-    try:
-        pool = orchestrator.stage_pools[0]
-        _register_request(orchestrator, "req-s")
-        orchestrator.request_states["req-s"].sampling_params_list = [_sampling_params()]
-        assert pool.select_replica_id("req-s") == 0
-        pool.evict_replica(0)
-
-        msg = StageSubmissionMessage(
-            type="streaming_update",
-            request_id="req-s",
-            prompt=SimpleNamespace(request_id="req-s", prompt_token_ids=[1]),
-            original_prompt={"prompt": "hi"},
-            output_prompt_text=None,
-            sampling_params_list=[],
-            final_stage_id=0,
-            preprocess_ms=0.0,
-            request_timestamp=time.time(),
-            enqueue_ts=time.time(),
-        )
-        await orchestrator._handle_streaming_update(msg)
-
-        err = queues[1].async_q.get_nowait()
-        assert isinstance(err, ErrorMessage)
-        assert err.fatal is True
-        assert err.request_id == "req-s"
-        assert "req-s" not in orchestrator.request_states
-    finally:
-        for q in queues:
-            q.close()
-
-
-@pytest.mark.asyncio
-async def test_streaming_update_unrelated_error_is_not_masked() -> None:
-    """The dispatch guard converts only StageUnavailableError/EngineDeadError;
-    a genuine bug surfacing as a plain RuntimeError must propagate (fail
-    fast), not be misreported as a dead-stage request failure.
-    """
-
-    class _BuggyClient(FakeStageClient):
-        async def add_request_async(self, *args, **kwargs):
-            raise RuntimeError("boom: unrelated bug")
-
-    orchestrator, queues = _build_bare_orchestrator(
-        _build_stage_pools([[_BuggyClient(stage_type="llm", final_output=True)]])
-    )
-    try:
-        _register_request(orchestrator, "req-b")
-        orchestrator.request_states["req-b"].sampling_params_list = [_sampling_params()]
-
-        msg = StageSubmissionMessage(
-            type="streaming_update",
-            request_id="req-b",
-            prompt=SimpleNamespace(request_id="req-b", prompt_token_ids=[1]),
-            original_prompt={"prompt": "hi"},
-            output_prompt_text=None,
-            sampling_params_list=[],
-            final_stage_id=0,
-            preprocess_ms=0.0,
-            request_timestamp=time.time(),
-            enqueue_ts=time.time(),
-        )
-        with pytest.raises(RuntimeError, match="unrelated bug"):
-            await orchestrator._handle_streaming_update(msg)
-    finally:
-        for q in queues:
-            q.close()
-
-
-@pytest.mark.asyncio
-async def test_add_companion_racing_eviction_fails_parent_not_server() -> None:
-    """A CFG companion submission hitting a fully dead stage must fail the
-    parent request (whose cleanup also releases the companion), not raise out
-    of _request_handler (#4285).
-    """
-    orchestrator, queues = _build_bare_orchestrator(_build_stage_pools([[FakeStageClient(stage_type="llm")]]))
-    try:
-        pool = orchestrator.stage_pools[0]
-        _register_request(orchestrator, "req-p")
-        orchestrator.request_states["req-p"].sampling_params_list = [_sampling_params()]
-        pool.evict_replica(0)
-
-        msg = AddCompanionRequestMessage(
-            companion_id="req-p-cfg",
-            parent_id="req-p",
-            role="negative",
-            prompt=SimpleNamespace(request_id="req-p-cfg", prompt_token_ids=[2]),
-            companion_prompt_text=None,
-            sampling_params_list=[_sampling_params()],
-        )
-        await orchestrator._handle_add_companion(msg)
-
-        err = queues[1].async_q.get_nowait()
-        assert isinstance(err, ErrorMessage)
-        assert err.fatal is True
-        assert err.request_id == "req-p"
-        assert "req-p" not in orchestrator.request_states
-        assert "req-p-cfg" not in orchestrator.request_states
-    finally:
-        for q in queues:
-            q.close()
-
-
-@pytest.mark.asyncio
-async def test_fail_request_dead_stage_is_noop_safe_for_unregistered_request() -> None:
-    """Direct helper test: the stage-0 dispatch guard calls the helper before the
-    request is registered; the cleanup must be a safe no-op and still emit the
-    fatal ErrorMessage.
-    """
-    orchestrator, queues = _build_bare_orchestrator(_build_stage_pools([[FakeStageClient(stage_type="llm")]]))
-    try:
-        await orchestrator._fail_request_dead_stage("req-unknown", 0)
-
-        msg = queues[1].async_q.get_nowait()
-        assert isinstance(msg, ErrorMessage)
-        assert msg.error == "Stage-0 has no live replica"
-        assert msg.fatal is True
-        assert msg.request_id == "req-unknown"
-        assert orchestrator.request_states == {}
-    finally:
-        for q in queues:
-            q.close()
-
-
-@pytest.mark.asyncio
-async def test_dispatch_death_evicts_so_subsequent_requests_reroute_to_survivor() -> None:
-    """A replica dying mid-submit (EngineDeadError from dispatch, before the poll
-    loop evicts it) is evicted by the guard so subsequent requests reroute to the
-    survivor instead of round-robining back onto the dead slot until the poll
-    catches up. The survivor and its request are untouched (#4285).
-    """
-    r0 = FakeStageClient(stage_type="llm")
-    r1 = FakeStageClient(stage_type="llm")
-    orchestrator, queues = _build_bare_orchestrator(_build_stage_pools([[r0, r1]]))
-    try:
-        pool = orchestrator.stage_pools[0]
-        _register_request(orchestrator, "req-dead")
-        _register_request(orchestrator, "req-live")
-        assert pool.select_replica_id("req-dead") == 0
-        assert pool.select_replica_id("req-live") == 1
-
-        async def _die() -> None:
-            raise EngineDeadError("stage-0 replica-0 died on submit")
-
-        failed = await orchestrator._dispatch_or_fail_request(
-            _die, req_id="req-dead", stage_id=0, operation="add_request"
-        )
-        assert failed is False
-
-        # The dead replica is evicted; only the request on it fails, the survivor
-        # and its request are untouched.
-        assert pool.live_replica_ids() == [1]
-        assert r0.shutdown_calls == 1
-        assert r1.shutdown_calls == 0
-        msg = queues[1].async_q.get_nowait()
-        assert isinstance(msg, ErrorMessage) and msg.fatal is True and msg.request_id == "req-dead"
-        assert queues[1].async_q.empty()
-        assert "req-dead" not in orchestrator.request_states
-        assert "req-live" in orchestrator.request_states and pool.get_bound_replica_id("req-live") == 1
-
-        # The point of eviction: a burst of subsequent requests all route to the
-        # survivor and never land back on the evicted replica.
-        for rid in ("next-1", "next-2", "next-3"):
-            assert pool.select_replica_id(rid) == 1
-    finally:
-        for q in queues:
-            q.close()
-
-
-@pytest.mark.asyncio
-async def test_dispatch_death_evicts_so_subsequent_requests_fast_fail() -> None:
-    """When the replica dying on submit was the stage's only one, the guard evicts
-    it (leaving the stage empty). A subsequent request then fast-fails at routing —
-    the ``_handle_add_request`` entry guard turns ``StageUnavailableError`` into a
-    fatal 5xx immediately, rather than dispatching onto the dead replica (#4285).
-    """
-    r0 = FakeStageClient(stage_type="llm")
-    orchestrator, queues = _build_bare_orchestrator(_build_stage_pools([[r0]]))
-    try:
-        pool = orchestrator.stage_pools[0]
-        _register_request(orchestrator, "req-a")
-        assert pool.select_replica_id("req-a") == 0
-
-        async def _die() -> None:
-            raise EngineDeadError("stage-0 replica-0 died on submit")
-
-        failed = await orchestrator._dispatch_or_fail_request(_die, req_id="req-a", stage_id=0, operation="add_request")
-        assert failed is False
-
-        assert pool.live_replica_ids() == []
-        assert r0.shutdown_calls == 1
-        msg = queues[1].async_q.get_nowait()
-        assert isinstance(msg, ErrorMessage) and msg.request_id == "req-a"
-        assert queues[1].async_q.empty()
-        assert "req-a" not in orchestrator.request_states
-
-        # A subsequent request cannot be routed at all: the stage has no live
-        # replica, so selection raises fast instead of dispatching onto the dead
-        # slot. _handle_add_request converts this to a fatal 5xx at its entry guard.
-        with pytest.raises(StageUnavailableError):
-            pool.select_replica_id("req-next")
-    finally:
-        for q in queues:
-            q.close()
-
-
 # ───────── Diffusion stage error output routing ─────────
 
 

tests/entrypoints/test_omni_entrypoints.py

@@ -104,9 +104,6 @@ class FakeAsyncOmniEngine:
         ]
         self.supported_tasks = ("generate",)
         self.stage_clients = [SimpleNamespace(is_comprehension=False) for _ in range(self.num_stages)]
-        # One replica per stage, sharing the stage_clients objects so a test that
-        # marks stage_clients[i] dead is reflected in the pool's liveness.
-        self.stage_pools = [_FakeStagePool([client], stage_id=i) for i, client in enumerate(self.stage_clients)]
         self.stage_vllm_configs = [None for _ in range(self.num_stages)]
         self.output_processors = [SimpleNamespace(tokenizer=None) for _ in range(self.num_stages)]
         self.input_processor = None
@@ -972,17 +969,11 @@ async def test_async_omni_propagates_non_400_client_error_status(
     assert str(exc_info.value) == error_text
 
 
-class _FakeStagePool:
-    def __init__(self, clients: list, stage_id: int = 0):
-        self.clients = clients
-        self.stage_id = stage_id
-
-
 def test_async_omni_errored_property_alive():
     omni = object.__new__(AsyncOmni)
     omni.engine = SimpleNamespace(
         is_alive=lambda: True,
-        stage_pools=[_FakeStagePool(clients=[SimpleNamespace()])],
+        stage_clients=[SimpleNamespace(is_comprehension=False)],
     )
 
     assert omni.errored is False
@@ -992,47 +983,21 @@ def test_async_omni_errored_property_dead_engine():
     omni = object.__new__(AsyncOmni)
     omni.engine = SimpleNamespace(
         is_alive=lambda: False,
-        stage_pools=[_FakeStagePool(clients=[SimpleNamespace()])],
+        stage_clients=[SimpleNamespace(is_comprehension=False)],
     )
 
     assert omni.errored is True
 
 
-def test_async_omni_errored_false_when_stage_dead():
-    # errored is process-fatal only: a dead stage must not trip the serving
-    # precheck, or requests that never touch that stage would be rejected too.
-    # Stage liveness surfaces via check_health / per-request dispatch failures.
+def test_async_omni_errored_property_dead_stage():
     omni = object.__new__(AsyncOmni)
+    dead_stage = SimpleNamespace(is_comprehension=False, _engine_dead=True)
     omni.engine = SimpleNamespace(
         is_alive=lambda: True,
-        stage_pools=[_FakeStagePool(clients=[SimpleNamespace(_engine_dead=True)])],
+        stage_clients=[dead_stage],
     )
 
-    assert omni.errored is False
-
-
-def test_async_omni_errored_false_when_a_replica_survives():
-    # Per-replica fault isolation (#4285): one replica dies, another lives →
-    # the stage is still serving, so the engine is not errored.
-    omni = object.__new__(AsyncOmni)
-    omni.engine = SimpleNamespace(
-        is_alive=lambda: True,
-        stage_pools=[_FakeStagePool(clients=[SimpleNamespace(_engine_dead=True), SimpleNamespace()])],
-    )
-
-    assert omni.errored is False
-
-
-def test_async_omni_errored_false_when_evicted_replica_leaves_stage_empty():
-    # Even a fully evicted stage (all slots None/dead) is not process-fatal;
-    # only orchestrator death makes errored True.
-    omni = object.__new__(AsyncOmni)
-    omni.engine = SimpleNamespace(
-        is_alive=lambda: True,
-        stage_pools=[_FakeStagePool(clients=[None, SimpleNamespace(_engine_dead=True)])],
-    )
-
-    assert omni.errored is False
+    assert omni.errored is True
 
 
 def _enqueue_stage_error(
@@ -1102,22 +1067,20 @@ async def test_async_omni_propagates_engine_generate_error(monkeypatch: pytest.M
 
 def test_check_health_passes_when_all_healthy():
     base = _make_base()
+    healthy_stage = MagicMock()
+    healthy_stage.check_health = MagicMock()
     base.engine.is_alive.return_value = True
-    healthy_stage = SimpleNamespace(check_health=lambda: None)
-    base.engine.stage_pools = [_FakeStagePool([healthy_stage])]
+    base.engine.stage_clients = [healthy_stage]
     base.check_health()  # should not raise
 
 
 def test_check_health_raises_when_stage_dead():
     base = _make_base()
+    dead_stage = MagicMock()
+    dead_stage.check_health = MagicMock(side_effect=EngineDeadError("Stage-1 dead"))
     base.engine.is_alive.return_value = True
-
-    def _raise_dead() -> None:
-        raise EngineDeadError("Stage-1 engine core is dead")
-
-    dead_stage = SimpleNamespace(check_health=_raise_dead)
-    base.engine.stage_pools = [_FakeStagePool([dead_stage], stage_id=1)]
-    with pytest.raises(EngineDeadError, match="Stage-1"):
+    base.engine.stage_clients = [dead_stage]
+    with pytest.raises(EngineDeadError, match="Stage-1 dead"):
         base.check_health()
 
 
@@ -1135,25 +1098,22 @@ def test_check_health_raises_when_orchestrator_dead():
 def test_omni_base_errored_false_when_alive():
     base = _make_base()
     base.engine.is_alive.return_value = True
-    base.engine.stage_pools = [_FakeStagePool([SimpleNamespace()])]
+    base.engine.stage_clients = [SimpleNamespace()]
     assert base.errored is False
 
 
-def test_omni_base_is_running_true_when_stage_engine_dead():
-    # is_running tracks the orchestrator only; a dead stage must not make
-    # `errored and not is_running` true, which would let terminate_if_errored
-    # kill the server on a partial failure.
+def test_omni_base_is_running_false_when_stage_engine_dead():
     base = _make_base()
     base.engine.is_alive.return_value = True
-    base.engine.stage_pools = [_FakeStagePool([SimpleNamespace(_engine_dead=True)])]
-    assert base.is_running is True
+    base.engine.stage_clients = [SimpleNamespace(_engine_dead=True)]
+    assert base.is_running is False
 
 
-def test_omni_base_is_running_true_when_stage_resources_engine_dead():
+def test_omni_base_is_running_false_when_stage_resources_engine_dead():
     base = _make_base()
     base.engine.is_alive.return_value = True
-    base.engine.stage_pools = [_FakeStagePool([SimpleNamespace(resources=SimpleNamespace(engine_dead=True))])]
-    assert base.is_running is True
+    base.engine.stage_clients = [SimpleNamespace(resources=SimpleNamespace(engine_dead=True))]
+    assert base.is_running is False
 
 
 def test_omni_base_errored_true_when_orchestrator_dead():
@@ -1163,22 +1123,20 @@ def test_omni_base_errored_true_when_orchestrator_dead():
     assert base.errored is True
 
 
-def test_omni_base_errored_false_when_stage_engine_dead():
-    # Process-fatal only: the serving precheck reads errored before routing,
-    # so a dead stage here would reject requests that never touch it.
+def test_omni_base_errored_true_when_stage_engine_dead():
     base = _make_base()
     base.engine.is_alive.return_value = True
     dead_stage = SimpleNamespace(_engine_dead=True)
-    base.engine.stage_pools = [_FakeStagePool([dead_stage])]
-    assert base.errored is False
+    base.engine.stage_clients = [dead_stage]
+    assert base.errored is True
 
 
-def test_omni_base_errored_false_when_stage_resources_engine_dead():
+def test_omni_base_errored_true_when_stage_resources_engine_dead():
     base = _make_base()
     base.engine.is_alive.return_value = True
     dead_stage = SimpleNamespace(resources=SimpleNamespace(engine_dead=True))
-    base.engine.stage_pools = [_FakeStagePool([dead_stage])]
-    assert base.errored is False
+    base.engine.stage_clients = [dead_stage]
+    assert base.errored is True
 
 
 # ───────── Omni (sync) EngineDeadError / EngineGenerateError ─────────
@@ -1244,14 +1202,13 @@ def test_omni_errored_property_dead_engine(monkeypatch: pytest.MonkeyPatch):
 
 
 def test_omni_errored_property_dead_stage(monkeypatch: pytest.MonkeyPatch):
-    """Omni.errored stays False when only a stage client dies (process-fatal
-    semantics): the request-level failure and check_health carry the signal."""
+    """Omni.errored returns True when a stage client is marked dead."""
     engine = FakeAsyncOmniEngine(stage_metadata=THREE_STAGE_META)
     _patch_engine(monkeypatch, engine)
 
     app = Omni("dummy-model")
     try:
         engine.stage_clients[0]._engine_dead = True
-        assert app.errored is False
+        assert app.errored is True
     finally:
         app.shutdown()

tests/entrypoints/test_serve.py

@@ -76,6 +76,17 @@ def test_serve_parser_accepts_four_way_cfg_parallelism() -> None:
     assert args.cfg_parallel_size == 4
 
 
+def test_serve_parser_accepts_wan_decode_precision() -> None:
+    parser = TrackingArgumentParser()
+    subparsers = parser.add_subparsers(dest="subcommand")
+    OmniServeCommand().subparser_init(subparsers)
+
+    args = parser.parse_args(["serve", "fake-model", "--omni", "--vae-decode-precision", "bf16"])
+
+    assert args.vae_decode_precision == "bf16"
+    assert args.get_explicit_kwargs_dict()["vae_decode_precision"] == "bf16"
+
+
 def _make_headless_args(**kwargs) -> TrackingNamespace:
     defaults = {
         "model": "fake-model",

tests/model_executor/models/voxtral_tts/test_cuda_graph_acoustic_transformer.py

@@ -83,16 +83,11 @@ class SyntheticAcousticTransformer(nn.Module):
         # time_embedding: (B, 1) -> (B, dim)
         self.time_embedding = nn.Linear(1, HIDDEN_DIM, bias=False)
 
-        # Identities: callers pass pre-projected inputs; the synthetic transformer folds
-        # everything into _velocity_proj.
-        self.llm_projection = nn.Identity()
-        self.time_projection = nn.Identity()
-
-        # _predict_velocity: takes x_t (B, C), llm_proj (B, D), t_proj (B, D) -> (B, C)
+        # _predict_velocity: takes x_t (B, C), llm_output (B, D), t_emb (B, D) -> (B, C)
         self._velocity_proj = nn.Linear(N_ACOUSTIC_CODEBOOK + HIDDEN_DIM + HIDDEN_DIM, N_ACOUSTIC_CODEBOOK, bias=False)
 
-    def _predict_velocity(self, x_t: torch.Tensor, llm_proj: torch.Tensor, t_proj: torch.Tensor) -> torch.Tensor:
-        combined = torch.cat([x_t, llm_proj, t_proj], dim=-1)
+    def _predict_velocity(self, x_t: torch.Tensor, llm_output: torch.Tensor, t_emb: torch.Tensor) -> torch.Tensor:
+        combined = torch.cat([x_t, llm_output, t_emb], dim=-1)
         return self._velocity_proj(combined)
 
 
@@ -134,18 +129,14 @@ class SyntheticModel(nn.Module):
         hidden_zero = torch.zeros_like(hidden_states)
         timesteps = torch.linspace(0, 1, 16, device=hidden_states.device, dtype=hidden_states.dtype)
 
-        # llm_projection is constant across steps so hoist it (mirrors decode_one_frame).
-        llm_batched = torch.cat([hidden_states, hidden_zero], dim=0)
-        llm_proj_batched = at.llm_projection(llm_batched)
-
         for i in range(len(timesteps) - 1):
             t = timesteps[i]
             dt = timesteps[i + 1] - timesteps[i]
             t_emb = at.time_embedding(t.view(-1, 1).repeat(B, 1)).to(hidden_states.dtype)
-            t_proj = at.time_projection(t_emb)
             x_batched = torch.cat([x, x], dim=0)
-            t_proj_batched = t_proj.repeat(2, 1)
-            v_all = at._predict_velocity(x_t=x_batched, llm_proj=llm_proj_batched, t_proj=t_proj_batched)
+            llm_batched = torch.cat([hidden_states, hidden_zero], dim=0)
+            t_emb_batched = t_emb.repeat(2, 1)
+            v_all = at._predict_velocity(x_t=x_batched, llm_output=llm_batched, t_emb=t_emb_batched)
             v_t, uncond_v_t = v_all[:B], v_all[B:]
             cfg_alpha = 1.2
             v_t = cfg_alpha * v_t + (1 - cfg_alpha) * uncond_v_t

vllm_omni/config/omni_config.py

@@ -573,6 +573,7 @@ class _DiffusionConfigProjection:
     fa_deterministic: bool = False
     vae_use_slicing: bool = False
     vae_use_tiling: bool = False
+    vae_decode_precision: Literal["fp32", "fp16", "bf16"] | None = None
     mask_strategy_file_path: str | None = None
     skip_time_steps: int = 15
     VSA_sparsity: float = 0.0

vllm_omni/config/stage_config.py

@@ -401,6 +401,7 @@ class StageDeployConfig:
     step_execution: bool | None = None
     vae_use_slicing: bool | None = None
     vae_use_tiling: bool | None = None
+    vae_decode_precision: str | None = None
     boundary_ratio: float | None = None
     flow_shift: float | None = None
     diffusion_kv_cache_dtype: str | None = None

vllm_omni/diffusion/data.py

@@ -711,6 +711,9 @@ class OmniDiffusionConfig:
     # VAE memory optimization parameters
     vae_use_slicing: bool = False
     vae_use_tiling: bool = False
+    # Optional Wan decode-only precision. When set, the VAE encoder remains
+    # FP32 while decoder-side modules use the selected precision.
+    vae_decode_precision: str | None = None
 
     # STA (Sliding Tile Attention) parameters
     mask_strategy_file_path: str | None = None
@@ -923,6 +926,11 @@ class OmniDiffusionConfig:
         )
 
     def __post_init__(self):
+        if self.vae_decode_precision not in (None, "fp32", "fp16", "bf16"):
+            raise ValueError(
+                "vae_decode_precision must be one of {'fp32', 'fp16', 'bf16'} or None, "
+                f"got {self.vae_decode_precision!r}"
+            )
         if self.diffusion_compile_granularity not in {"regional", "full"}:
             raise ValueError(
                 "diffusion_compile_granularity must be 'regional' or 'full', "

vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py

@@ -106,7 +106,7 @@ class DistributedAutoencoderKLWan(OmniAutoencoderKLWan, DistributedVaeMixin):
             split_dims=(3, 4),
             grid_shape=(tiletask_list[-1].grid_coord[0] + 1, tiletask_list[-1].grid_coord[1] + 1),
             tile_spec=tile_spec,
-            output_dtype=self.dtype,
+            output_dtype=getattr(self, "_vllm_decode_dtype", self.dtype),
         )
         return tiletask_list, grid_spec
 

vllm_omni/diffusion/vae_optimizations/__init__.py

@@ -0,0 +1,19 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Optional VAE decoder optimizations."""
+
+from vllm_omni.diffusion.vae_optimizations.gate import (
+    use_pipeline_vae_fast_path,
+    use_vae_fast_path,
+)
+from vllm_omni.diffusion.vae_optimizations.optimize import (
+    clear_pipeline_vae_fast_path_caches,
+    optimize_pipeline_vaes,
+)
+
+__all__ = [
+    "clear_pipeline_vae_fast_path_caches",
+    "optimize_pipeline_vaes",
+    "use_pipeline_vae_fast_path",
+    "use_vae_fast_path",
+]

vllm_omni/diffusion/vae_optimizations/gate.py

@@ -0,0 +1,67 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Decode-scoped gates for optional, non-bit-exact VAE fast paths."""
+
+from __future__ import annotations
+
+from collections.abc import Iterator
+from contextlib import ExitStack, contextmanager
+from weakref import WeakKeyDictionary
+
+from torch import nn
+
+from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery
+
+
+class VaeFastPathGate:
+    """Mutable flag shared by wrappers installed on one VAE instance."""
+
+    __slots__ = ("enabled",)
+
+    def __init__(self) -> None:
+        self.enabled = False
+
+
+_VAE_FAST_PATH_GATES: WeakKeyDictionary[nn.Module, VaeFastPathGate] = WeakKeyDictionary()
+
+
+def register_vae_fast_path_gate(vae: nn.Module, gate: VaeFastPathGate) -> None:
+    """Associate an installed fast path with its owning VAE."""
+
+    _VAE_FAST_PATH_GATES[vae] = gate
+
+
+@contextmanager
+def use_vae_fast_path(vae: nn.Module, enabled: bool) -> Iterator[None]:
+    """Enable an installed fast path for one decode and always restore it."""
+
+    gate = _VAE_FAST_PATH_GATES.get(vae)
+    if gate is None:
+        yield
+        return
+
+    previous_enabled = gate.enabled
+    gate.enabled = enabled
+    try:
+        yield
+    finally:
+        gate.enabled = previous_enabled
+
+
+@contextmanager
+def use_pipeline_vae_fast_path(pipeline: nn.Module, enabled: bool) -> Iterator[None]:
+    """Apply one request's quality choice to every declared VAE component."""
+
+    vaes = ModuleDiscovery.discover(pipeline).vaes
+    with ExitStack() as stack:
+        for vae in vaes:
+            stack.enter_context(use_vae_fast_path(vae, enabled))
+        yield
+
+
+__all__ = [
+    "VaeFastPathGate",
+    "register_vae_fast_path_gate",
+    "use_pipeline_vae_fast_path",
+    "use_vae_fast_path",
+]

vllm_omni/diffusion/vae_optimizations/image.py

@@ -0,0 +1,354 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Quality-gated CUDA fast paths for Diffusers-style 2D KL VAE decoders."""
+
+from __future__ import annotations
+
+from types import MethodType
+from typing import Any
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.vae_optimizations.gate import (
+    VaeFastPathGate,
+    register_vae_fast_path_gate,
+)
+
+logger = init_logger(__name__)
+
+try:
+    from vllm_omni.diffusion.vae_optimizations.triton_group_norm_silu import (
+        group_norm_silu_4d,
+        group_norm_silu_rows,
+    )
+
+    _HAS_TRITON = True
+except ImportError:  # pragma: no cover - exercised on builds without Triton
+    _HAS_TRITON = False
+
+
+class FusedGroupNormSiLU(nn.Module):
+    """Quality-gated GroupNorm + SiLU with an exact eager fallback."""
+
+    def __init__(self, norm: nn.GroupNorm, gate: VaeFastPathGate) -> None:
+        super().__init__()
+        self.num_groups = norm.num_groups
+        self.num_channels = norm.num_channels
+        self.eps = norm.eps
+        self.affine = norm.affine
+        # Register parameters directly so state-dict FQNs do not change.
+        self.weight = norm.weight
+        self.bias = norm.bias
+        self._vllm_gate = gate
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self._vllm_gate.enabled and x.dim() == 4 and not torch.compiler.is_compiling():
+            output = group_norm_silu_4d(
+                x,
+                self.weight,
+                self.bias,
+                self.num_groups,
+                self.eps,
+                apply_silu=True,
+            )
+            if output is not None:
+                return output
+        return F.silu(F.group_norm(x, self.num_groups, self.weight, self.bias, self.eps))
+
+
+def _install_norm_silu(decoder: nn.Module, resnet_cls: type[nn.Module], gate: VaeFastPathGate) -> int:
+    def is_fusable(norm: object) -> bool:
+        return type(norm) is nn.GroupNorm and norm.affine and norm.weight is not None and norm.bias is not None
+
+    count = 0
+    for module in decoder.modules():
+        if (
+            type(module) is resnet_cls
+            and module.time_emb_proj is None
+            and module.time_embedding_norm in ("default", "group")
+            and module.upsample is None
+            and module.downsample is None
+            and isinstance(module.nonlinearity, nn.SiLU)
+            and is_fusable(module.norm1)
+            and is_fusable(module.norm2)
+        ):
+            module.norm1 = FusedGroupNormSiLU(module.norm1, gate)
+            module.norm2 = FusedGroupNormSiLU(module.norm2, gate)
+            module.nonlinearity = nn.Identity()
+            count += 2
+
+    if (
+        type(getattr(decoder, "conv_norm_out", None)) is nn.GroupNorm
+        and isinstance(getattr(decoder, "conv_act", None), nn.SiLU)
+        and is_fusable(decoder.conv_norm_out)
+    ):
+        decoder.conv_norm_out = FusedGroupNormSiLU(decoder.conv_norm_out, gate)
+        decoder.conv_act = nn.Identity()
+        count += 1
+    return count
+
+
+_UPSAMPLE_TAP_MAP = {0: (2,), 1: (1, 2), 2: (0, 1), 3: (0,)}
+
+
+def _fold_upsample2x_conv2d_weight(conv: nn.Conv2d) -> torch.Tensor:
+    weight = conv.weight.detach().float()
+    output_channels, input_channels = weight.shape[:2]
+    folded = weight.new_zeros(input_channels, output_channels, 4, 4)
+    for row in range(4):
+        for column in range(4):
+            value = weight.new_zeros(output_channels, input_channels)
+            for source_row in _UPSAMPLE_TAP_MAP[row]:
+                for source_column in _UPSAMPLE_TAP_MAP[column]:
+                    value += weight[:, :, source_row, source_column]
+            folded[:, :, row, column] = value.t()
+    folded = folded.to(conv.weight.dtype)
+    if conv.weight.is_contiguous(memory_format=torch.channels_last):
+        folded = folded.contiguous(memory_format=torch.channels_last)
+    return folded.to(conv.weight.device)
+
+
+class FusedUpsample2xConv2d(nn.Module):
+    """Quality-gated nearest-2x + Conv2d replacement using ConvTranspose2d."""
+
+    def __init__(self, upsample: nn.Module, gate: VaeFastPathGate) -> None:
+        super().__init__()
+        # Keep only ``conv`` registered so checkpoint parameter names stay
+        # ``...upsamplers.N.conv.*``. The original module is an eager fallback.
+        object.__setattr__(self, "_original", upsample)
+        self.conv = upsample.conv
+        self.channels = upsample.channels
+        self._vllm_gate = gate
+        self.register_buffer("_fused_weight", None, persistent=False)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        output_size: int | None = None,
+        *args: Any,
+        **kwargs: Any,
+    ) -> torch.Tensor:
+        if not self._vllm_gate.enabled or output_size is not None or hidden_states.shape[1] != self.channels:
+            return self._original(hidden_states, output_size=output_size, *args, **kwargs)
+
+        weight = self._fused_weight
+        if weight is None or weight.device != self.conv.weight.device or weight.dtype != self.conv.weight.dtype:
+            weight = _fold_upsample2x_conv2d_weight(self.conv)
+            self._fused_weight = weight
+        return F.conv_transpose2d(hidden_states, weight, self.conv.bias, stride=2, padding=1)
+
+
+def _install_fused_upsample(decoder: nn.Module, upsample_cls: type[nn.Module], gate: VaeFastPathGate) -> int:
+    count = 0
+    for block in decoder.up_blocks:
+        upsamplers = getattr(block, "upsamplers", None)
+        if not upsamplers:
+            continue
+        for index, upsample in enumerate(upsamplers):
+            if type(upsample) is not upsample_cls:
+                continue
+            conv = getattr(upsample, "conv", None)
+            if (
+                upsample.use_conv
+                and not upsample.use_conv_transpose
+                and upsample.interpolate
+                and upsample.norm is None
+                and upsample.name == "conv"
+                and type(conv) is nn.Conv2d
+                and conv.kernel_size == (3, 3)
+                and conv.stride == (1, 1)
+                and conv.padding == (1, 1)
+                and conv.dilation == (1, 1)
+                and conv.groups == 1
+                and conv.padding_mode == "zeros"
+            ):
+                upsamplers[index] = FusedUpsample2xConv2d(upsample, gate)
+                count += 1
+    return count
+
+
+def _fold_attention_value_projection(module: nn.Module) -> tuple[torch.Tensor, torch.Tensor]:
+    value_weight = module.to_v.weight.detach().float()
+    value_bias = module.to_v.bias.detach().float()
+    output_weight = module.to_out[0].weight.detach().float()
+    output_bias = module.to_out[0].bias.detach().float()
+    dtype = module.to_v.weight.dtype
+    return (
+        (output_weight @ value_weight).to(dtype).contiguous(),
+        (output_weight @ value_bias + output_bias).to(dtype),
+    )
+
+
+def _attention_fast_forward(
+    self: nn.Module,
+    hidden_states: torch.Tensor,
+    encoder_hidden_states: torch.Tensor | None = None,
+    attention_mask: torch.Tensor | None = None,
+    temb: torch.Tensor | None = None,
+    **cross_attention_kwargs: Any,
+) -> torch.Tensor:
+    if (
+        not self._vllm_gate.enabled
+        or encoder_hidden_states is not None
+        or attention_mask is not None
+        or temb is not None
+        or hidden_states.ndim != 4
+    ):
+        return type(self).forward(
+            self,
+            hidden_states,
+            encoder_hidden_states=encoder_hidden_states,
+            attention_mask=attention_mask,
+            temb=temb,
+            **cross_attention_kwargs,
+        )
+
+    residual = hidden_states
+    batch_size, channels, height, width = hidden_states.shape
+    hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(batch_size, height * width, channels)
+
+    if self.group_norm is not None:
+        norm = self.group_norm
+        normalized = None
+        if not torch.compiler.is_compiling():
+            normalized = group_norm_silu_rows(
+                hidden_states,
+                norm.weight,
+                norm.bias,
+                norm.num_groups,
+                norm.eps,
+                apply_silu=False,
+            )
+        hidden_states = normalized if normalized is not None else norm(hidden_states.transpose(1, 2)).transpose(1, 2)
+
+    query = self.to_q(hidden_states)
+    key = self.to_k(hidden_states)
+    folded_weight = self._vllm_folded_value_weight
+    folded_bias = self._vllm_folded_value_bias
+    if (
+        folded_weight is None
+        or folded_bias is None
+        or folded_weight.device != hidden_states.device
+        or folded_weight.dtype != self.to_v.weight.dtype
+    ):
+        folded_weight, folded_bias = _fold_attention_value_projection(self)
+        self._vllm_folded_value_weight = folded_weight
+        self._vllm_folded_value_bias = folded_bias
+    value = F.linear(hidden_states, folded_weight, folded_bias)
+
+    output = F.scaled_dot_product_attention(query.unsqueeze(1), key.unsqueeze(1), value.unsqueeze(1))
+    output = output.squeeze(1).to(query.dtype)
+    output = self.to_out[1](output)
+    output = output.reshape(batch_size, height, width, channels).permute(0, 3, 1, 2)
+    if self.residual_connection:
+        output = output + residual
+    return output / self.rescale_output_factor
+
+
+def _attention_fast_compatible(module: nn.Module, attention_cls: type, processor_cls: type) -> bool:
+    return (
+        type(module) is attention_cls
+        and isinstance(module.processor, processor_cls)
+        and module.heads == 1
+        and module.scale_qk
+        and module.spatial_norm is None
+        and not module.norm_cross
+        and getattr(module, "norm_q", None) is None
+        and getattr(module, "norm_k", None) is None
+        and getattr(module, "add_k_proj", None) is None
+        and type(module.to_q) is nn.Linear
+        and type(module.to_k) is nn.Linear
+        and type(module.to_v) is nn.Linear
+        and type(module.to_out[0]) is nn.Linear
+        and isinstance(module.to_out[1], nn.Dropout)
+        and module.to_v.bias is not None
+        and module.to_out[0].bias is not None
+    )
+
+
+def _decoder_layout_forward(self: nn.Module, *args: Any, **kwargs: Any) -> torch.Tensor:
+    use_channels_last = self._vllm_gate.enabled
+    if use_channels_last != self._vllm_channels_last:
+        memory_format = torch.channels_last if use_channels_last else torch.contiguous_format
+        self.to(memory_format=memory_format)
+        self._vllm_channels_last = use_channels_last
+        logger.debug(
+            "%s decoder switched to %s layout",
+            self._vllm_label,
+            "channels_last" if use_channels_last else "contiguous",
+        )
+    return type(self).forward(self, *args, **kwargs)
+
+
+def maybe_optimize_image_vae(vae: nn.Module) -> nn.Module:
+    """Install image VAE decoder fast paths when the structure is fully supported."""
+
+    if not _HAS_TRITON or getattr(vae, "_vllm_vae_fast_path_installed", False):
+        return vae
+
+    from diffusers.models.attention_processor import Attention, AttnProcessor2_0
+    from diffusers.models.autoencoders import AutoencoderKL
+    from diffusers.models.autoencoders.vae import Decoder
+    from diffusers.models.resnet import ResnetBlock2D
+    from diffusers.models.upsampling import Upsample2D
+
+    supported_types: tuple[type[nn.Module], ...] = (AutoencoderKL,)
+    try:
+        from diffusers.models.autoencoders.autoencoder_kl_flux2 import AutoencoderKLFlux2
+
+        supported_types += (AutoencoderKLFlux2,)
+    except ImportError:  # pragma: no cover - older supported Diffusers builds
+        pass
+
+    if not isinstance(vae, supported_types):
+        return vae
+    decoder = getattr(vae, "decoder", None)
+    if type(decoder) is not Decoder:
+        return vae
+
+    attention_modules = [
+        module for module in decoder.modules() if _attention_fast_compatible(module, Attention, AttnProcessor2_0)
+    ]
+    attention_count = sum(isinstance(module, Attention) for module in decoder.modules())
+    if len(attention_modules) != attention_count:
+        logger.info(
+            "%s has %d/%d attention blocks without a layout-safe rewrite; skipping VAE fast paths",
+            type(vae).__name__,
+            attention_count - len(attention_modules),
+            attention_count,
+        )
+        return vae
+
+    gate = VaeFastPathGate()
+    decoder._vllm_gate = gate
+    decoder._vllm_label = type(vae).__name__
+    decoder._vllm_channels_last = False
+    decoder.forward = MethodType(_decoder_layout_forward, decoder)
+    upsample_count = _install_fused_upsample(decoder, Upsample2D, gate)
+    for module in attention_modules:
+        module._vllm_gate = gate
+        module.register_buffer("_vllm_folded_value_weight", None, persistent=False)
+        module.register_buffer("_vllm_folded_value_bias", None, persistent=False)
+        module.forward = MethodType(_attention_fast_forward, module)
+    norm_count = _install_norm_silu(decoder, ResnetBlock2D, gate)
+    register_vae_fast_path_gate(vae, gate)
+    vae._vllm_vae_fast_path_installed = True
+    logger.info(
+        "%s installed quality-gated VAE decoder fast paths "
+        "(%d upsamplers, %d attention blocks, %d GroupNorm+SiLU fusions)",
+        type(vae).__name__,
+        upsample_count,
+        len(attention_modules),
+        norm_count,
+    )
+    return vae
+
+
+__all__ = [
+    "FusedGroupNormSiLU",
+    "FusedUpsample2xConv2d",
+    "maybe_optimize_image_vae",
+]

vllm_omni/diffusion/vae_optimizations/optimize.py

@@ -0,0 +1,70 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Pipeline-level installation for optional VAE optimizations."""
+
+from __future__ import annotations
+
+from torch import nn
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.data import OmniDiffusionConfig
+from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery
+from vllm_omni.diffusion.vae_optimizations.image import maybe_optimize_image_vae
+from vllm_omni.diffusion.vae_optimizations.precision import configure_wan_decode_precision
+from vllm_omni.diffusion.vae_optimizations.wan import maybe_optimize_wan_vae
+from vllm_omni.platforms import current_omni_platform
+
+logger = init_logger(__name__)
+
+_FOLDED_CACHE_NAMES = (
+    "_fused_weight",
+    "_vllm_folded_value_weight",
+    "_vllm_folded_value_bias",
+)
+
+
+def optimize_pipeline_vaes(pipeline: nn.Module, od_config: OmniDiffusionConfig) -> None:
+    """Configure declared VAE components once after checkpoint loading."""
+
+    precision = getattr(od_config, "vae_decode_precision", None)
+    vaes = ModuleDiscovery.discover(pipeline).vaes
+    for vae in vaes:
+        # An explicit precision override is a serving contract, so failures on
+        # a recognized Wan VAE must abort startup instead of being hidden.
+        configure_wan_decode_precision(vae, precision)
+
+    if not current_omni_platform.is_cuda():
+        return
+
+    for vae in vaes:
+        for optimizer in (maybe_optimize_image_vae, maybe_optimize_wan_vae):
+            try:
+                optimizer(vae)
+            except Exception:
+                # A wrapper may already have installed its exact fallback
+                # before a later optional rewrite failed. With no registered
+                # enabled gate, those wrappers remain on the reference path.
+                logger.warning(
+                    "Failed to install optional VAE fast paths on %s; using the reference path",
+                    type(vae).__name__,
+                    exc_info=True,
+                )
+
+
+def clear_pipeline_vae_fast_path_caches(pipeline: nn.Module) -> None:
+    """Release lazily folded buffers before the worker enters sleep mode."""
+
+    cleared = 0
+    for vae in ModuleDiscovery.discover(pipeline).vaes:
+        if not getattr(vae, "_vllm_vae_fast_path_installed", False):
+            continue
+        for module in vae.modules():
+            for name in _FOLDED_CACHE_NAMES:
+                if name in module._buffers and module._buffers[name] is not None:
+                    setattr(module, name, None)
+                    cleared += 1
+    if cleared:
+        logger.debug("Released %d lazy VAE fast-path buffers before sleep", cleared)
+
+
+__all__ = ["clear_pipeline_vae_fast_path_caches", "optimize_pipeline_vaes"]

vllm_omni/diffusion/vae_optimizations/precision.py

@@ -0,0 +1,91 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Decode-only precision configuration for structurally split Wan VAEs."""
+
+from __future__ import annotations
+
+from types import MethodType
+from typing import Any
+
+import torch
+from torch import nn
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+PRECISION_TO_DTYPE = {
+    "fp32": torch.float32,
+    "fp16": torch.float16,
+    "bf16": torch.bfloat16,
+}
+
+
+def _decode_with_precision(
+    self: nn.Module,
+    latents: torch.Tensor,
+    *args: Any,
+    **kwargs: Any,
+) -> Any:
+    return type(self).decode(
+        self,
+        latents.to(dtype=self._vllm_decode_dtype),
+        *args,
+        **kwargs,
+    )
+
+
+def _encode_in_fp32(
+    self: nn.Module,
+    samples: torch.Tensor,
+    *args: Any,
+    **kwargs: Any,
+) -> Any:
+    return type(self).encode(
+        self,
+        samples.to(dtype=torch.float32),
+        *args,
+        **kwargs,
+    )
+
+
+def configure_wan_decode_precision(vae: nn.Module, precision: str | None) -> nn.Module:
+    """Keep Wan encode in FP32 while materializing only decode modules lower."""
+
+    if precision is None:
+        return vae
+
+    from diffusers.models.autoencoders import AutoencoderKLWan
+
+    if not isinstance(vae, AutoencoderKLWan):
+        return vae
+    try:
+        decode_dtype = PRECISION_TO_DTYPE[precision]
+    except KeyError as exc:
+        raise ValueError(
+            f"Unsupported vae_decode_precision={precision!r}; expected one of {sorted(PRECISION_TO_DTYPE)}"
+        ) from exc
+
+    encoder = getattr(vae, "encoder", None)
+    quant_conv = getattr(vae, "quant_conv", None)
+    decoder = getattr(vae, "decoder", None)
+    post_quant_conv = getattr(vae, "post_quant_conv", None)
+    if not all(isinstance(module, nn.Module) for module in (encoder, decoder, post_quant_conv)):
+        logger.info("Wan VAE lacks separate encoder/decoder modules; skipping decode-only precision")
+        return vae
+
+    encoder.to(dtype=torch.float32)
+    if isinstance(quant_conv, nn.Module):
+        quant_conv.to(dtype=torch.float32)
+    decoder.to(dtype=decode_dtype)
+    post_quant_conv.to(dtype=decode_dtype)
+    vae.encode = MethodType(_encode_in_fp32, vae)
+    vae._vllm_decode_dtype = decode_dtype
+    vae.decode = MethodType(_decode_with_precision, vae)
+    logger.info(
+        "Wan VAE configured with FP32 encode and %s decode",
+        precision,
+    )
+    return vae
+
+
+__all__ = ["PRECISION_TO_DTYPE", "configure_wan_decode_precision"]

vllm_omni/diffusion/vae_optimizations/triton_group_norm_silu.py

@@ -0,0 +1,237 @@
+# ruff: noqa: N803
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Channels-last two-pass GroupNorm(+SiLU) Triton kernels.
+
+The kernels keep the channel dimension innermost throughout a 2D VAE
+decoder. Statistics are accumulated in FP32, then a separate apply pass
+performs the affine transform and optional SiLU epilogue. Unsupported inputs
+return ``None`` so callers can fail closed to the reference implementation.
+"""
+
+from __future__ import annotations
+
+import torch
+import triton
+import triton.language as tl
+
+_SUPPORTED_DTYPES = {torch.float16, torch.bfloat16, torch.float32}
+_MAX_CHANNELS = 2048
+
+
+@triton.jit
+def _gn_partial_rows_kernel(
+    x_ptr,
+    psum_ptr,
+    psq_ptr,
+    rows,
+    rows_per_prog,
+    C: tl.constexpr,
+    BLOCK_R: tl.constexpr,
+):
+    chunk = tl.program_id(0).to(tl.int64)
+    batch = tl.program_id(1).to(tl.int64)
+    num_chunks = tl.num_programs(0)
+    row0 = chunk * rows_per_prog
+    cols = tl.arange(0, C)
+    acc_sum = tl.zeros((C,), tl.float32)
+    acc_sq = tl.zeros((C,), tl.float32)
+    x_base = x_ptr + batch * rows * C
+    for row_offset in range(0, rows_per_prog, BLOCK_R):
+        row_indices = row0 + row_offset + tl.arange(0, BLOCK_R)
+        mask = row_indices < rows
+        x = tl.load(
+            x_base + row_indices[:, None] * C + cols[None, :],
+            mask=mask[:, None],
+            other=0.0,
+        ).to(tl.float32)
+        acc_sum += tl.sum(x, 0)
+        acc_sq += tl.sum(x * x, 0)
+    output_offsets = (batch * num_chunks + chunk) * C + cols
+    tl.store(psum_ptr + output_offsets, acc_sum)
+    tl.store(psq_ptr + output_offsets, acc_sq)
+
+
+@triton.jit
+def _gn_finalize_kernel(
+    psum_ptr,
+    psq_ptr,
+    weight_ptr,
+    bias_ptr,
+    scale_shift_ptr,
+    num_chunks,
+    group_numel,
+    eps,
+    channels,
+    CPG: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+):
+    group = tl.program_id(0).to(tl.int64)
+    batch = tl.program_id(1).to(tl.int64)
+    cols = group * CPG + tl.arange(0, CPG)
+    acc_sum = tl.zeros((), tl.float32)
+    acc_sq = tl.zeros((), tl.float32)
+    for chunk0 in range(0, num_chunks, BLOCK_K):
+        chunks = chunk0 + tl.arange(0, BLOCK_K)
+        mask = chunks < num_chunks
+        offsets = (batch * num_chunks + chunks)[:, None] * channels + cols[None, :]
+        acc_sum += tl.sum(tl.load(psum_ptr + offsets, mask=mask[:, None], other=0.0))
+        acc_sq += tl.sum(tl.load(psq_ptr + offsets, mask=mask[:, None], other=0.0))
+    mean = acc_sum / group_numel
+    variance = acc_sq / group_numel - mean * mean
+    variance = tl.maximum(variance, 0.0)
+    reciprocal_std = tl.rsqrt(variance + eps)
+    weight = tl.load(weight_ptr + cols).to(tl.float32)
+    bias = tl.load(bias_ptr + cols).to(tl.float32)
+    scale = weight * reciprocal_std
+    shift = bias - mean * scale
+    tl.store(scale_shift_ptr + batch * 2 * channels + cols, scale)
+    tl.store(scale_shift_ptr + (batch * 2 + 1) * channels + cols, shift)
+
+
+@triton.jit
+def _gn_apply_rows_kernel(
+    x_ptr,
+    scale_shift_ptr,
+    y_ptr,
+    rows,
+    C: tl.constexpr,
+    BLOCK_R: tl.constexpr,
+    SILU: tl.constexpr,
+):
+    program = tl.program_id(0).to(tl.int64)
+    batch = tl.program_id(1).to(tl.int64)
+    cols = tl.arange(0, C)
+    scale = tl.load(scale_shift_ptr + batch * 2 * C + cols)
+    shift = tl.load(scale_shift_ptr + (batch * 2 + 1) * C + cols)
+    row_indices = program * BLOCK_R + tl.arange(0, BLOCK_R)
+    mask = row_indices < rows
+    offsets = batch * rows * C + row_indices[:, None] * C + cols[None, :]
+    x = tl.load(x_ptr + offsets, mask=mask[:, None], other=0.0).to(tl.float32)
+    y = x * scale[None, :] + shift[None, :]
+    if SILU:
+        y = y * tl.sigmoid(y)
+    tl.store(y_ptr + offsets, y.to(y_ptr.dtype.element_ty), mask=mask[:, None])
+
+
+def _group_norm_silu_rows(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    num_groups: int,
+    eps: float,
+    apply_silu: bool,
+) -> torch.Tensor:
+    batch_size, rows, channels = x.shape
+    channels_per_group = channels // num_groups
+    block_rows = max(1, 8192 // channels)
+    rows_per_program = block_rows * 32
+    num_chunks = triton.cdiv(rows, rows_per_program)
+    partial_sum = torch.empty((batch_size, num_chunks, channels), device=x.device, dtype=torch.float32)
+    partial_sq = torch.empty_like(partial_sum)
+    with torch.cuda.device(x.device):
+        _gn_partial_rows_kernel[(num_chunks, batch_size)](
+            x,
+            partial_sum,
+            partial_sq,
+            rows,
+            rows_per_program,
+            C=channels,
+            BLOCK_R=block_rows,
+            num_warps=4,
+        )
+        scale_shift = torch.empty((batch_size, 2, channels), device=x.device, dtype=torch.float32)
+        block_chunks = max(1, min(4096 // channels_per_group, triton.next_power_of_2(num_chunks)))
+        _gn_finalize_kernel[(num_groups, batch_size)](
+            partial_sum,
+            partial_sq,
+            weight,
+            bias,
+            scale_shift,
+            num_chunks,
+            rows * channels_per_group,
+            eps,
+            channels,
+            CPG=channels_per_group,
+            BLOCK_K=block_chunks,
+            num_warps=4,
+        )
+        output = torch.empty_like(x)
+        _gn_apply_rows_kernel[(triton.cdiv(rows, block_rows), batch_size)](
+            x,
+            scale_shift,
+            output,
+            rows,
+            C=channels,
+            BLOCK_R=block_rows,
+            SILU=apply_silu,
+            num_warps=4,
+        )
+    return output
+
+
+def _is_supported(
+    x: torch.Tensor,
+    weight: torch.Tensor | None,
+    bias: torch.Tensor | None,
+    num_groups: int,
+) -> bool:
+    if not (x.is_cuda and x.numel() > 0 and not torch.is_grad_enabled()):
+        return False
+    if x.requires_grad or x.dtype not in _SUPPORTED_DTYPES:
+        return False
+    if weight is None or bias is None:
+        return False
+    channels = x.shape[1] if x.dim() == 4 else x.shape[-1]
+    if weight.shape != (channels,) or bias.shape != (channels,):
+        return False
+    if not (
+        weight.device == bias.device == x.device
+        and weight.dtype == bias.dtype == x.dtype
+        and weight.is_contiguous()
+        and bias.is_contiguous()
+    ):
+        return False
+    if num_groups < 1 or channels % num_groups != 0:
+        return False
+    return triton.next_power_of_2(channels) == channels and channels <= _MAX_CHANNELS
+
+
+def group_norm_silu_4d(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    num_groups: int,
+    eps: float,
+    apply_silu: bool = True,
+) -> torch.Tensor | None:
+    """Run GroupNorm(+SiLU) on a channels-last ``(N, C, H, W)`` tensor."""
+
+    if x.dim() != 4 or not _is_supported(x, weight, bias, num_groups):
+        return None
+    batch_size, channels, height, width = x.shape
+    if not (channels > 1 and (height > 1 or width > 1) and x.is_contiguous(memory_format=torch.channels_last)):
+        return None
+    rows = x.permute(0, 2, 3, 1).reshape(batch_size, height * width, channels)
+    output = _group_norm_silu_rows(rows, weight, bias, num_groups, eps, apply_silu)
+    return output.reshape(batch_size, height, width, channels).permute(0, 3, 1, 2)
+
+
+def group_norm_silu_rows(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    num_groups: int,
+    eps: float,
+    apply_silu: bool = True,
+) -> torch.Tensor | None:
+    """Run GroupNorm(+SiLU) over contiguous ``(N, rows, C)`` input."""
+
+    if x.dim() != 3 or not x.is_contiguous():
+        return None
+    if not _is_supported(x, weight, bias, num_groups):
+        return None
+    return _group_norm_silu_rows(x, weight, bias, num_groups, eps, apply_silu)
+
+
+__all__ = ["group_norm_silu_4d", "group_norm_silu_rows"]

vllm_omni/diffusion/vae_optimizations/triton_wan_rmsnorm_silu.py

@@ -0,0 +1,160 @@
+# ruff: noqa: N803
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Channels-last-3D Wan VAE RMSNorm(+SiLU) Triton kernel."""
+
+from __future__ import annotations
+
+import torch
+import triton
+import triton.language as tl
+
+_SUPPORTED_DTYPES = {torch.float16, torch.bfloat16, torch.float32}
+_MAX_CHANNELS = 1024
+
+
+@triton.jit
+def _wan_rmsnorm_silu_kernel(
+    x_ptr,
+    gamma_ptr,
+    bias_ptr,
+    output_ptr,
+    channels: tl.constexpr,
+    time_size,
+    height_size,
+    width_size,
+    x_stride_batch,
+    x_stride_channel,
+    x_stride_time,
+    x_stride_height,
+    x_stride_width,
+    output_stride_batch,
+    output_stride_channel,
+    output_stride_time,
+    output_stride_height,
+    output_stride_width,
+    rms_scale,
+    eps,
+    has_bias: tl.constexpr,
+    BLOCK_C: tl.constexpr,
+):
+    row = tl.program_id(0).to(tl.int64)
+    channel_offsets = tl.arange(0, BLOCK_C)
+    channel_mask = channel_offsets < channels
+
+    width = row % width_size
+    remainder = row // width_size
+    height = remainder % height_size
+    remainder = remainder // height_size
+    time = remainder % time_size
+    batch = remainder // time_size
+
+    x_base = batch * x_stride_batch + time * x_stride_time + height * x_stride_height + width * x_stride_width
+    output_base = (
+        batch * output_stride_batch
+        + time * output_stride_time
+        + height * output_stride_height
+        + width * output_stride_width
+    )
+
+    x = tl.load(
+        x_ptr + x_base + channel_offsets * x_stride_channel,
+        mask=channel_mask,
+        other=0.0,
+    ).to(tl.float32)
+    norm = tl.sqrt(tl.sum(x * x, axis=0))
+    inverse_norm = 1.0 / tl.maximum(norm, eps)
+
+    # Match the eager dtype boundaries: normalize and scale in activation
+    # dtype, affine in the promoted dtype, then SiLU in FP32.
+    y = (x * inverse_norm).to(x_ptr.dtype.element_ty)
+    gamma = tl.load(gamma_ptr + channel_offsets, mask=channel_mask, other=1.0)
+    y = (y * rms_scale).to(x_ptr.dtype.element_ty)
+    y = (y.to(tl.float32) * gamma.to(tl.float32)).to(output_ptr.dtype.element_ty)
+    if has_bias:
+        bias = tl.load(bias_ptr + channel_offsets, mask=channel_mask, other=0.0)
+        y = (y.to(tl.float32) + bias.to(tl.float32)).to(output_ptr.dtype.element_ty)
+    y = y.to(tl.float32)
+    y = y * tl.sigmoid(y)
+    tl.store(
+        output_ptr + output_base + channel_offsets * output_stride_channel,
+        y,
+        mask=channel_mask,
+    )
+
+
+def _affine_supported(x: torch.Tensor, affine: torch.Tensor) -> bool:
+    return (
+        affine.is_cuda
+        and affine.device == x.device
+        and (affine.dtype == x.dtype or affine.dtype == torch.float32)
+        and affine.numel() == x.shape[1]
+    )
+
+
+def can_use_wan_rmsnorm_silu(
+    x: torch.Tensor,
+    gamma: torch.Tensor,
+    bias: torch.Tensor | None,
+) -> bool:
+    """Return whether ``wan_rmsnorm_silu`` supports this exact input."""
+
+    return (
+        x.is_cuda
+        and not torch.is_grad_enabled()
+        and not x.requires_grad
+        and x.dtype in _SUPPORTED_DTYPES
+        and x.ndim == 5
+        and x.numel() > 0
+        and 0 < x.shape[1] <= _MAX_CHANNELS
+        and x.is_contiguous(memory_format=torch.channels_last_3d)
+        and _affine_supported(x, gamma)
+        and (bias is None or _affine_supported(x, bias))
+    )
+
+
+def wan_rmsnorm_silu(
+    x: torch.Tensor,
+    gamma: torch.Tensor,
+    bias: torch.Tensor | None = None,
+    rms_scale: float | None = None,
+    eps: float = 1e-12,
+) -> torch.Tensor | None:
+    """Fuse ``SiLU(F.normalize(x, dim=1) * scale * gamma + bias)``."""
+
+    if not can_use_wan_rmsnorm_silu(x, gamma, bias):
+        return None
+
+    batch_size, channels, time_size, height_size, width_size = x.shape
+    output_dtype = torch.promote_types(x.dtype, gamma.dtype)
+    output = torch.empty_strided(x.shape, x.stride(), device=x.device, dtype=output_dtype)
+    block_channels = triton.next_power_of_2(channels)
+    num_warps = 1 if block_channels <= 64 else 4 if block_channels <= 512 else 8
+    gamma = gamma.reshape(channels).contiguous()
+    has_bias = bias is not None
+    bias = gamma if bias is None else bias.reshape(channels).contiguous()
+    if rms_scale is None:
+        rms_scale = channels**0.5
+
+    with torch.cuda.device(x.device):
+        _wan_rmsnorm_silu_kernel[(batch_size * time_size * height_size * width_size,)](
+            x,
+            gamma,
+            bias,
+            output,
+            channels,
+            time_size,
+            height_size,
+            width_size,
+            *x.stride(),
+            *output.stride(),
+            float(rms_scale),
+            eps,
+            has_bias,
+            BLOCK_C=block_channels,
+            num_warps=num_warps,
+        )
+    return output
+
+
+__all__ = ["can_use_wan_rmsnorm_silu", "wan_rmsnorm_silu"]

vllm_omni/diffusion/vae_optimizations/wan.py

@@ -0,0 +1,162 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Quality-gated CUDA fast path for Diffusers Wan VAE decoders."""
+
+from __future__ import annotations
+
+from types import MethodType
+from typing import Any
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.vae_optimizations.gate import (
+    VaeFastPathGate,
+    register_vae_fast_path_gate,
+)
+
+logger = init_logger(__name__)
+
+try:
+    from vllm_omni.diffusion.vae_optimizations.triton_wan_rmsnorm_silu import (
+        wan_rmsnorm_silu,
+    )
+
+    _HAS_TRITON = True
+except ImportError:  # pragma: no cover - exercised on builds without Triton
+    _HAS_TRITON = False
+
+
+class FusedWanRMSNormSiLU(nn.Module):
+    """Quality-gated Wan RMSNorm + SiLU with an exact eager fallback."""
+
+    def __init__(self, norm: nn.Module, gate: VaeFastPathGate) -> None:
+        super().__init__()
+        self.gamma = norm.gamma
+        self.bias = norm.bias
+        self.scale = float(norm.scale)
+        self._vllm_gate = gate
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self._vllm_gate.enabled and not torch.compiler.is_compiling():
+            bias = self.bias if isinstance(self.bias, torch.Tensor) else None
+            output = wan_rmsnorm_silu(x, self.gamma, bias, rms_scale=self.scale)
+            if output is not None:
+                return output
+
+        needs_fp32_normalize = x.dtype in (torch.float16, torch.bfloat16) or any(
+            token in str(x.dtype) for token in ("float4_", "float8_")
+        )
+        normalized = F.normalize(x.float() if needs_fp32_normalize else x, dim=1).to(x.dtype)
+        return F.silu(normalized * self.scale * self.gamma + self.bias)
+
+
+def _plain_silu(activation: object) -> bool:
+    return isinstance(activation, nn.SiLU) and not activation.inplace
+
+
+def _fusable_norm(norm: object, norm_cls: type) -> bool:
+    return (
+        type(norm) is norm_cls
+        and getattr(norm, "channel_first", False)
+        and isinstance(getattr(norm, "gamma", None), torch.Tensor)
+    )
+
+
+def _install_norm_silu(decoder: nn.Module, gate: VaeFastPathGate) -> int | None:
+    from diffusers.models.autoencoders.autoencoder_kl_wan import WanResidualBlock, WanRMS_norm
+
+    residual_blocks = [module for module in decoder.modules() if isinstance(module, WanResidualBlock)]
+    eligible = [
+        module
+        for module in residual_blocks
+        if type(module) is WanResidualBlock
+        and _plain_silu(module.nonlinearity)
+        and _fusable_norm(module.norm1, WanRMS_norm)
+        and _fusable_norm(module.norm2, WanRMS_norm)
+    ]
+    if len(eligible) != len(residual_blocks):
+        logger.info(
+            "Wan VAE has %d/%d non-standard residual blocks; skipping fast path",
+            len(residual_blocks) - len(eligible),
+            len(residual_blocks),
+        )
+        return None
+    if not (
+        _fusable_norm(getattr(decoder, "norm_out", None), WanRMS_norm)
+        and _plain_silu(getattr(decoder, "nonlinearity", None))
+    ):
+        logger.info("Wan VAE has a non-standard output head; skipping fast path")
+        return None
+
+    count = 0
+    for module in eligible:
+        module.norm1 = FusedWanRMSNormSiLU(module.norm1, gate)
+        module.norm2 = FusedWanRMSNormSiLU(module.norm2, gate)
+        module.nonlinearity = nn.Identity()
+        count += 2
+    decoder.norm_out = FusedWanRMSNormSiLU(decoder.norm_out, gate)
+    decoder.nonlinearity = nn.Identity()
+    return count + 1
+
+
+def _set_wan_decoder_memory_format(decoder: nn.Module, *, channels_last: bool) -> None:
+    for module in decoder.modules():
+        if isinstance(module, nn.Conv3d):
+            module.to(memory_format=torch.channels_last_3d if channels_last else torch.contiguous_format)
+        elif isinstance(module, (nn.Conv2d, nn.ConvTranspose2d)):
+            module.to(memory_format=torch.channels_last if channels_last else torch.contiguous_format)
+
+
+def _decoder_layout_forward(self: nn.Module, *args: Any, **kwargs: Any) -> torch.Tensor:
+    use_channels_last = self._vllm_gate.enabled
+    if use_channels_last != self._vllm_channels_last:
+        _set_wan_decoder_memory_format(self, channels_last=use_channels_last)
+        self._vllm_channels_last = use_channels_last
+        logger.debug(
+            "Wan VAE decoder switched to %s layout",
+            "channels_last_3d" if use_channels_last else "contiguous",
+        )
+    return type(self).forward(self, *args, **kwargs)
+
+
+def maybe_optimize_wan_vae(vae: nn.Module) -> nn.Module:
+    """Install the Wan decoder fast path when its full structure is supported."""
+
+    if not _HAS_TRITON or getattr(vae, "_vllm_vae_fast_path_installed", False):
+        return vae
+
+    from diffusers.models.autoencoders import AutoencoderKLWan
+    from diffusers.models.autoencoders.autoencoder_kl_wan import WanDecoder3d
+
+    if not isinstance(vae, AutoencoderKLWan):
+        return vae
+    decoder = getattr(vae, "decoder", None)
+    if type(decoder) is not WanDecoder3d:
+        return vae
+
+    distributed_executor = getattr(vae, "distributed_executor", None)
+    parallel_mode = getattr(distributed_executor, "parallel_mode", "tile")
+    if parallel_mode in ("auto", "spatial_shard_height", "spatial_shard_width"):
+        logger.info("Wan spatial-shard decode is configured; skipping quality-gated VAE fast path")
+        return vae
+
+    gate = VaeFastPathGate()
+    norm_count = _install_norm_silu(decoder, gate)
+    if norm_count is None:
+        return vae
+    decoder._vllm_gate = gate
+    decoder._vllm_channels_last = False
+    decoder.forward = MethodType(_decoder_layout_forward, decoder)
+    register_vae_fast_path_gate(vae, gate)
+    vae._vllm_vae_fast_path_installed = True
+    logger.info(
+        "Wan VAE installed quality-gated decoder fast path (%d RMSNorm+SiLU fusions)",
+        norm_count,
+    )
+    return vae
+
+
+__all__ = ["FusedWanRMSNormSiLU", "maybe_optimize_wan_vae"]

vllm_omni/diffusion/worker/diffusion_model_runner.py

@@ -43,6 +43,10 @@ from vllm_omni.diffusion.offloader import get_offload_backend
 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
+from vllm_omni.diffusion.vae_optimizations import (
+    optimize_pipeline_vaes,
+    use_pipeline_vae_fast_path,
+)
 from vllm_omni.diffusion.worker.input_batch import InputBatch, scatter_latents
 from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
 from vllm_omni.diffusion.worker.utils import (
@@ -243,6 +247,9 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
                     custom_pipeline_name=custom_pipeline_name,
                     device=self.device,
                 )
+                # Keep precision rematerialization in the tagged weights pool
+                # so sleep level 1 can offload every resulting parameter.
+                optimize_pipeline_vaes(self.pipeline, self.od_config)
         time_after_load = time.perf_counter()
 
         logger.info(
@@ -499,6 +506,11 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
         if require_request_batch_support and not getattr(self.pipeline, "supports_request_batch", False):
             raise RuntimeError(f"{type(self.pipeline).__name__} does not support request-batch forward.")
 
+        quality = getattr(reqs[0].sampling_params, "quality", None)
+        if any(getattr(req.sampling_params, "quality", None) != quality for req in reqs[1:]):
+            raise RuntimeError("A diffusion worker batch cannot mix request quality levels.")
+        use_fast_vae = quality == "high"
+
         # Use no_grad() for HSDP compatibility, inference_mode() otherwise for
         # better perf. HSDP2's fully_shard pre-forward hooks need tensor version
         # counters, which inference tensors do not track.
@@ -523,7 +535,8 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
 
             with set_forward_context(vllm_config=self.vllm_config, omni_diffusion_config=od_config):
                 with record_function(record_name):
-                    raw_outputs = self.pipeline.forward(batch)
+                    with use_pipeline_vae_fast_path(self.pipeline, use_fast_vae):
+                        raw_outputs = self.pipeline.forward(batch)
                     outputs = _normalize_pipeline_outputs(
                         raw_outputs,
                         expected_count=len(reqs),
@@ -774,7 +787,11 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
 
                             if should_decode:
                                 clear_pipeline_stage_durations(self.pipeline)
-                                result = self.pipeline.post_decode(req)
+                                with use_pipeline_vae_fast_path(
+                                    self.pipeline,
+                                    getattr(req.sampling, "quality", None) == "high",
+                                ):
+                                    result = self.pipeline.post_decode(req)
                                 if result is not None:
                                     self._attach_stepwise_metrics(
                                         req,

vllm_omni/diffusion/worker/diffusion_worker.py

@@ -576,6 +576,14 @@ class DiffusionWorker:
 
         usage_before = allocator.get_current_usage()
 
+        if self.model_runner is not None and self.model_runner.pipeline is not None:
+            from vllm_omni.diffusion.vae_optimizations import clear_pipeline_vae_fast_path_caches
+
+            # Folded weights are allocated lazily on the request path rather
+            # than in the tagged weight pool. Drop them before sleep so they
+            # neither remain resident nor need a separate CPU backup.
+            clear_pipeline_vae_fast_path_caches(self.model_runner.pipeline)
+
         if level == 2 and self.model_runner is not None:
             if hasattr(self.model_runner, "graph_runners"):
                 self.model_runner.graph_runners.clear()

vllm_omni/engine/arg_utils.py

@@ -491,6 +491,7 @@ class OrchestratorArgs:
     step_execution: bool = False
     vae_use_slicing: bool = False
     vae_use_tiling: bool = False
+    vae_decode_precision: str | None = None
     enable_multithread_weight_load: bool = True
     num_weight_load_threads: int = 4
     enable_cpu_offload: bool = False

vllm_omni/engine/async_omni_engine.py

@@ -1026,6 +1026,7 @@ class AsyncOmniEngine:
             "request_batch_max_wait_ms": kwargs.get("request_batch_max_wait_ms", 0.0),
             "vae_use_slicing": kwargs.get("vae_use_slicing", False),
             "vae_use_tiling": kwargs.get("vae_use_tiling", False),
+            "vae_decode_precision": kwargs.get("vae_decode_precision", None),
             "cache_backend": cache_backend,
             "cache_config": cache_config,
             "enable_cache_dit_summary": kwargs.get("enable_cache_dit_summary", False),

vllm_omni/engine/orchestrator.py

@@ -50,7 +50,7 @@ from vllm_omni.engine.messages import (
 )
 from vllm_omni.engine.orchestrator_monitor import create_orch_monitor, replica_key
 from vllm_omni.engine.serialization import serialize_additional_information
-from vllm_omni.engine.stage_pool import StagePool, StageUnavailableError
+from vllm_omni.engine.stage_pool import StagePool
 from vllm_omni.metrics.prometheus import OmniRequestCounter
 from vllm_omni.metrics.stat_logger import OmniPrometheusStatLogger
 from vllm_omni.outputs import OmniRequestOutput
@@ -417,6 +417,8 @@ class Orchestrator:
 
         self._shutdown_event = asyncio.Event()
         self._stages_shutdown = False
+        self._fatal_error: str | None = None
+        self._fatal_error_stage_id: int | None = None
 
         # Distributed membership (optional, injected by DistStageRuntime)
         self._membership = membership_controller
@@ -530,17 +532,19 @@ class Orchestrator:
         except asyncio.CancelledError:
             raise
         except EngineDeadError as e:
-            # The orchestration loop isolates per-replica EngineDeadError
-            # (evict + continue), so this only fires if one escapes outside the
-            # poll handler (e.g. from output post-processing). Treat it as a
-            # teardown signal and let the finally block clean up; do not re-raise.
-            logger.info("[Orchestrator] Engine dead, shutting down: %s", e)
-            # Unblock any collective_rpc callers waiting on rpc_async_queue
-            # instead of leaving them to hit their own timeout.
+            # EngineDeadError from _orchestration_loop means the diffusion
+            # engine died.  All pending requests were already notified and
+            # _shutdown_event was already set by the loop's handler.
+            # During teardown this is expected; the finally block handles
+            # proper cleanup.  Do not re-raise.
+            logger.info("[Orchestrator] Engine dead during shutdown: %s", e)
+            if self._fatal_error is None:
+                self._fatal_error = str(e) or "Stage engine died"
             await self.rpc_async_queue.put(
                 ErrorMessage(
-                    error=str(e) or "Stage engine died",
+                    error=self._fatal_error or str(e),
                     fatal=True,
+                    stage_id=self._fatal_error_stage_id,
                 )
             )
         except Exception:
@@ -558,6 +562,9 @@ class Orchestrator:
             if self.duplex_control_plane is not None:
                 await self.duplex_control_plane.shutdown()
 
+            if self._fatal_error is not None:
+                await self._drain_pending_requests_on_fatal()
+
             if self._membership is not None:
                 await self._membership.drain_tasks(timeout=10.0)
                 self._membership.shutdown()
@@ -644,13 +651,6 @@ class Orchestrator:
         final_stage_id = msg.final_stage_id
         final_output_stage_ids = set(msg.final_output_stage_ids or [final_stage_id])
 
-        if not self.stage_pools[stage_id].live_replica_ids():
-            # Stage 0 lost all replicas between the HTTP-layer errored check and
-            # dispatch. Runs before request state / running counter registration,
-            # so the helper's cleanup is a no-op here.
-            await self._fail_request_dead_stage(request_id, stage_id)
-            return
-
         logger.debug(
             "[Orchestrator] _handle_add_request: stage=%s req=%s "
             "prompt_type=%s original_prompt_type=%s final_stage=%s "
@@ -682,18 +682,12 @@ class Orchestrator:
         preprocess_ms = msg.preprocess_ms
         if preprocess_ms > 0:
             req_state.pipeline_timings["preprocess_ms"] = preprocess_ms
-        if not await self._dispatch_or_fail_request(
-            lambda: self.stage_pools[stage_id].submit_initial(
-                request_id,
-                req_state,
-                prompt,
-                prompt_text=msg.output_prompt_text,
-            ),
-            req_id=request_id,
-            stage_id=stage_id,
-            operation="add_request",
-        ):
-            return
+        await self.stage_pools[stage_id].submit_initial(
+            request_id,
+            req_state,
+            prompt,
+            prompt_text=msg.output_prompt_text,
+        )
 
         if self.async_chunk and stage_id == 0 and final_stage_id > 0:
             await self._prewarm_async_chunk_stages(request_id, prompt, req_state)
@@ -723,18 +717,12 @@ class Orchestrator:
 
         req_state.streaming.enabled = True
         req_state.stage_submit_ts[stage_id] = _time.time()
-        if not await self._dispatch_or_fail_request(
-            lambda: self.stage_pools[stage_id].submit_update(
-                request_id,
-                req_state,
-                request,
-                prompt_text=msg.output_prompt_text,
-            ),
-            req_id=request_id,
-            stage_id=stage_id,
-            operation="streaming_update",
-        ):
-            return
+        await self.stage_pools[stage_id].submit_update(
+            request_id,
+            req_state,
+            request,
+            prompt_text=msg.output_prompt_text,
+        )
 
         if self.async_chunk and stage_id == 0 and final_stage_id > 0:
             await self._prewarm_async_chunk_stages(request_id, request, req_state)
@@ -769,29 +757,20 @@ class Orchestrator:
         )
         self.request_states[companion_id] = companion_state
         companion_state.stage_submit_ts[0] = _time.time()
-        # A dead-stage failure is attributed to the parent request; its cleanup
-        # also releases this companion.
-        if not await self._dispatch_or_fail_request(
-            lambda: self.stage_pools[0].submit_initial(
-                companion_id,
-                companion_state,
-                companion_prompt,
-                prompt_text=msg.companion_prompt_text,
-                affinity_request_id=parent_id,
-            ),
-            req_id=parent_id,
-            stage_id=0,
-            operation=f"companion {companion_id}",
-            dispatch_req_id=companion_id,
-        ):
-            return
+        companion_replica_id = await self.stage_pools[0].submit_initial(
+            companion_id,
+            companion_state,
+            companion_prompt,
+            prompt_text=msg.companion_prompt_text,
+            affinity_request_id=parent_id,
+        )
 
         logger.info(
             "[Orchestrator] CFG companion submitted: %s (role=%s, parent=%s, stage-0 replica-%s)",
             companion_id,
             role,
             parent_id,
-            self.stage_pools[0].get_bound_replica_id(companion_id),
+            companion_replica_id,
         )
 
     async def _handle_abort(self, msg: AbortRequestMessage) -> None:
@@ -924,17 +903,16 @@ class Orchestrator:
                     if self._shutdown_event.is_set():
                         return
 
-                    # Shared catch so a dead replica on either poll path is
-                    # evicted rather than tearing down every stage (#4285).
-                    try:
-                        if pool.stage_type == "diffusion":
-                            diffusion_output = pool.poll_diffusion_output(replica_id)
-                            if diffusion_output is None:
-                                continue
+                    if pool.stage_type == "diffusion":
+                        output = pool.poll_diffusion_output(replica_id)
+                        if output is None:
+                            continue
 
-                            pool.record_output_timestamps([diffusion_output])
-                            processed = [diffusion_output]
-                        else:
+                        pool.record_output_timestamps([output])
+                        await self._handle_processed_outputs(stage_id, replica_id, [output])
+                        idle = False
+                    else:
+                        try:
                             raw_outputs = await pool.poll_llm_raw_output(replica_id, timeout_s=0.001)
                             if raw_outputs is None:
                                 continue
@@ -966,7 +944,7 @@ class Orchestrator:
                             iteration_stats = (
                                 IterationStats() if (self._stat_logger is not None and raw_outputs.outputs) else None
                             )
-                            processed = await pool.process_llm_raw_outputs(
+                            raw_output = await pool.process_llm_raw_outputs(
                                 replica_id,
                                 raw_outputs,
                                 iteration_stats=iteration_stats,
@@ -979,23 +957,80 @@ class Orchestrator:
                                     iteration_stats,
                                     engine_idx=self._stage_replica_to_engine_idx[(stage_id, replica_id)],
                                 )
-                    except asyncio.CancelledError:
-                        raise
-                    except EngineDeadError as e:
-                        await self._handle_dead_replica(stage_id, replica_id, e)
-                        continue
-                    except Exception:
-                        if self._shutdown_event.is_set():
-                            return
-                        logger.exception(
-                            "[Orchestrator] Stage-%s replica-%s processing failed",
-                            stage_id,
-                            replica_id,
-                        )
-                        raise
+                        except asyncio.CancelledError:
+                            raise
+                        except EngineDeadError as e:
+                            logger.error(
+                                "[Orchestrator] Stage-%s replica-%s is dead: %s",
+                                stage_id,
+                                replica_id,
+                                e,
+                            )
+                            affected_request_ids = pool.mark_replica_unavailable(replica_id)
+                            closed_sessions = (
+                                self.duplex_control_plane.close_sessions_for_request_ids(
+                                    affected_request_ids,
+                                    abort=False,
+                                )
+                                if self.duplex_control_plane is not None
+                                else {}
+                            )
+                            for session_id, stale_request_ids in closed_sessions.items():
+                                affected_request_ids.extend(stale_request_ids)
+                                logger.warning(
+                                    "[Orchestrator] closed duplex session %s after stage-%s replica-%s died; "
+                                    "stale_request_ids=%s",
+                                    session_id,
+                                    stage_id,
+                                    replica_id,
+                                    stale_request_ids,
+                                )
+                            affected_request_ids = list(dict.fromkeys(affected_request_ids))
+                            if pool.available_replica_ids():
+                                for req_id in affected_request_ids:
+                                    await self.output_async_queue.put(
+                                        ErrorMessage(
+                                            error=str(e),
+                                            fatal=False,
+                                            request_id=req_id,
+                                            stage_id=stage_id,
+                                        )
+                                    )
+                                await self._cleanup_request_ids(
+                                    affected_request_ids,
+                                    close_duplex_sessions=True,
+                                )
+                                continue
 
-                    await self._handle_processed_outputs(stage_id, replica_id, processed)
-                    idle = False
+                            self._fatal_error = str(e)
+                            self._fatal_error_stage_id = stage_id
+                            for req_id in affected_request_ids:
+                                await self.output_async_queue.put(
+                                    ErrorMessage(
+                                        error=str(e),
+                                        fatal=True,
+                                        request_id=req_id,
+                                        stage_id=stage_id,
+                                    )
+                                )
+                            await self._cleanup_request_ids(
+                                affected_request_ids,
+                                close_duplex_sessions=True,
+                            )
+                            self._shutdown_event.set()
+                            raise
+                        except Exception:
+                            if self._shutdown_event.is_set():
+                                return
+                            logger.exception(
+                                "[Orchestrator] Stage-%s replica-%s processing failed",
+                                stage_id,
+                                replica_id,
+                            )
+                            raise
+
+                        await self._handle_processed_outputs(stage_id, replica_id, raw_output)
+                        idle = False
 
             self._orch_monitor.note_loop(idle=idle)
             if idle:
@@ -1056,143 +1091,6 @@ class Orchestrator:
             close_duplex_sessions=True,
         )
 
-    async def _handle_dead_replica(self, stage_id: int, replica_id: int, error: EngineDeadError) -> None:
-        """Evict a dead stage replica and fail the requests stranded on it (#4285).
-
-        A dead replica must not tear down the whole server: evict it and keep the
-        orchestrator loop alive so remaining live replicas / stages keep serving.
-        If the stage still has a live replica, fail only the requests bound to the
-        dead replica; requests bound to a surviving replica (or not yet bound)
-        keep running and can be (re)dispatched. If the stage has no live replica
-        left, every request routed through it must fail since none can be served.
-        """
-        pool = self.stage_pools[stage_id]
-        logger.error(
-            "[Orchestrator] Stage-%s replica-%s is dead; evicting it and failing its in-flight requests: %s",
-            stage_id,
-            replica_id,
-            error,
-        )
-        pool.evict_replica(replica_id)
-        stage_has_live = bool(pool.live_replica_ids())
-        failed_ids: list[str] = []
-        for req_id, req_state in list(self.request_states.items()):
-            if stage_id not in req_state.stage_submit_ts:
-                continue
-            bound = pool.get_bound_replica_id(req_id)
-            if bound == replica_id or not stage_has_live:
-                await self.output_async_queue.put(
-                    ErrorMessage(
-                        error=str(error),
-                        fatal=True,
-                        request_id=req_id,
-                        stage_id=stage_id,
-                    )
-                )
-                failed_ids.append(req_id)
-        # Use the shared cleanup path so the running counter, PD/CFG state, and
-        # bindings across every stage pool are released (not just this pool).
-        # abort=True stops the failed requests' in-flight work on surviving
-        # stages (abort skips dead/evicted bindings), matching the
-        # distributed-membership eviction path.
-        for req_id in failed_ids:
-            await self._cleanup_request_ids(
-                [req_id, *self._cfg_tracker.cleanup_parent(req_id)],
-                abort=True,
-            )
-
-    async def _fail_request_dead_stage(self, req_id: str, stage_id: int) -> None:
-        """Fail one request whose target stage has no live replica (#4285).
-
-        Used at dispatch sites so a fully dead stage fails just that request
-        instead of raising out of the request/orchestration task and tearing
-        down the server. The cleanup releases the running counter, PD/CFG state,
-        and bindings across every stage pool; it is a no-op for a request that
-        was never registered.
-        """
-        logger.error(
-            "[Orchestrator] req=%s: stage-%d has no live replica; failing the request",
-            req_id,
-            stage_id,
-        )
-        await self.output_async_queue.put(
-            ErrorMessage(
-                error=f"Stage-{stage_id} has no live replica",
-                fatal=True,
-                request_id=req_id,
-                stage_id=stage_id,
-            )
-        )
-        await self._cleanup_request_ids(
-            [req_id, *self._cfg_tracker.cleanup_parent(req_id)],
-            abort=True,
-        )
-
-    async def _dispatch_or_fail_request(
-        self,
-        dispatch: Callable[[], Awaitable[Any]],
-        *,
-        req_id: str,
-        stage_id: int,
-        operation: str,
-        dispatch_req_id: str | None = None,
-    ) -> bool:
-        """Run a dispatch action, converting stage unavailability into a
-        single-request failure.
-
-        ``dispatch`` is a thunk rather than a bare coroutine so that replica
-        selection and other argument setup run inside this guard: a synchronous
-        StageUnavailableError raised while building the dispatch is caught here,
-        not left to propagate out of the caller.
-
-        Replica eviction can interleave with any dispatch, surfacing as
-        StageUnavailableError (no live replica / evicted slot) or
-        EngineDeadError (replica died mid-submit before the poll loop evicted
-        it). Both mean "this request cannot be placed", not "the server is
-        broken": fail the request and keep serving (#4285). Unrelated errors
-        propagate. Returns False when the request was failed.
-
-        ``req_id`` is the request the failure is attributed to; ``dispatch_req_id``
-        is the id actually being dispatched when it differs (e.g. a CFG companion
-        submitted under its parent's attribution) and is used to locate the
-        replica that died.
-        """
-        try:
-            await dispatch()
-            return True
-        except StageUnavailableError as e:
-            # No specific replica to evict: the stage already has no live
-            # replica or the chosen slot was evicted. Fail just this request.
-            logger.error(
-                "[Orchestrator] %s dispatch for req=%s hit unavailable stage-%s: %s",
-                operation,
-                req_id,
-                stage_id,
-                e,
-            )
-            await self._fail_request_dead_stage(req_id, stage_id)
-            return False
-        except EngineDeadError as e:
-            # A live replica died mid-submit before the poll loop evicted it.
-            # Evict it now (and fail every request bound to it) so a burst of
-            # new requests stops round-robining onto the dead slot until the
-            # poll catches up. Capture the replica before failing this request,
-            # since the failure releases its binding.
-            pool = self.stage_pools[stage_id]
-            dead_replica = pool.get_bound_replica_id(dispatch_req_id or req_id)
-            logger.error(
-                "[Orchestrator] %s dispatch for req=%s hit dead stage-%s replica-%s: %s",
-                operation,
-                req_id,
-                stage_id,
-                dead_replica,
-                e,
-            )
-            await self._fail_request_dead_stage(req_id, stage_id)
-            if dead_replica is not None and dead_replica in pool.live_replica_ids():
-                await self._handle_dead_replica(stage_id, dead_replica, e)
-            return False
-
     # ---- Shared helpers ----
 
     async def _cleanup_request_ids(
@@ -1837,42 +1735,10 @@ class Orchestrator:
         src_replica_id: int | None = None,
         is_streaming_session: bool = False,
         is_final_update: bool = False,
-    ) -> None:
-        """Forward output to the next stage; a dead stage fails only this request."""
-        await self._dispatch_or_fail_request(
-            lambda: self._forward_to_next_stage_unguarded(
-                req_id,
-                src_stage_id,
-                output,
-                req_state,
-                src_replica_id=src_replica_id,
-                is_streaming_session=is_streaming_session,
-                is_final_update=is_final_update,
-            ),
-            req_id=req_id,
-            stage_id=src_stage_id + 1,
-            operation="inter-stage forward",
-        )
-
-    async def _forward_to_next_stage_unguarded(
-        self,
-        req_id: str,
-        src_stage_id: int,
-        output: Any,
-        req_state: OrchestratorRequestState,
-        *,
-        src_replica_id: int | None = None,
-        is_streaming_session: bool = False,
-        is_final_update: bool = False,
     ) -> None:
         """Forward output from the current logical stage to the next one."""
         next_logical = src_stage_id + 1
         next_pool = self.stage_pools[next_logical]
-        if not next_pool.live_replica_ids():
-            # The downstream stage lost all of its replicas before this in-flight
-            # request could be forwarded.
-            await self._fail_request_dead_stage(req_id, next_logical)
-            return
         next_client = next_pool.stage_client
         params = req_state.sampling_params_list[next_logical]
         source_outputs = [output]
@@ -2235,22 +2101,16 @@ class Orchestrator:
             _t_submit_start = _time.perf_counter()
 
             if next_pool.stage_type == "diffusion":
-                submit_kwargs = {
-                    "kv_sender_info": self._build_kv_sender_info(
-                        list(getattr(next_pool.stage_client, "engine_input_source", None) or [next_stage_id - 1]),
-                        request_id=request_id,
-                    )
-                }
-                submitted = await self._dispatch_or_fail_request(
-                    lambda: next_pool.submit_initial(
-                        request_id,
-                        req_state,
-                        req_state.prompt,
-                        submit_kwargs=submit_kwargs,
-                    ),
-                    req_id=request_id,
-                    stage_id=next_stage_id,
-                    operation="async-chunk prewarm",
+                replica_id = await next_pool.submit_initial(
+                    request_id,
+                    req_state,
+                    req_state.prompt,
+                    submit_kwargs={
+                        "kv_sender_info": self._build_kv_sender_info(
+                            list(getattr(next_pool.stage_client, "engine_input_source", None) or [next_stage_id - 1]),
+                            request_id=request_id,
+                        )
+                    },
                 )
             else:
                 import copy
@@ -2280,30 +2140,16 @@ class Orchestrator:
                     resumable=downstream_resumable,
                 )
                 request.external_req_id = request.request_id
-                submitted = await self._dispatch_or_fail_request(
-                    lambda: next_pool.submit_initial(
-                        request_id,
-                        req_state,
-                        request,
-                        prompt_text=None,
-                    ),
-                    req_id=request_id,
-                    stage_id=next_stage_id,
-                    operation="async-chunk prewarm",
+                replica_id = await next_pool.submit_initial(
+                    request_id,
+                    req_state,
+                    request,
+                    prompt_text=None,
                 )
-
-            # Attribute the failure to the stage that failed, not stage 0; the
-            # request is already cleaned up, so stop prewarming.
-            if not submitted:
-                return
-
-            # The guard swallows submit_initial's return value; the pool binding
-            # it recorded carries the same replica.
-            bound_replica_id = next_pool.get_bound_replica_id(request_id)
             self._record_duplex_stage_submission(
                 next_stage_id,
                 request_id,
-                bound_replica_id if bound_replica_id is not None else 0,
+                replica_id,
                 req_state,
             )
 
@@ -2355,6 +2201,53 @@ class Orchestrator:
 
     # ---- Shutdown / lifecycle ----
 
+    async def _drain_pending_requests_on_fatal(self) -> None:
+        """Drain the request queue and broadcast fatal errors for any
+        pending add_request messages that were never processed.
+
+        Called from the ``run()`` finally block when a fatal error
+        (e.g. ``EngineDeadError``) caused the orchestrator to shut down
+        before the request handler could process all queued messages.
+        Also broadcasts for any already-tracked requests still in
+        ``request_states`` that were not yet notified.
+        """
+        assert self._fatal_error is not None
+
+        notified: set[str] = set()
+
+        # 1) Drain pending messages from the request queue.
+        while True:
+            try:
+                msg = self.request_async_queue.get_nowait()
+            except Exception:
+                break
+            if msg.type == "add_request":
+                req_id = msg.request_id
+                await self.output_async_queue.put(
+                    ErrorMessage(
+                        error=self._fatal_error,
+                        fatal=True,
+                        request_id=req_id,
+                        stage_id=self._fatal_error_stage_id,
+                    )
+                )
+                notified.add(req_id)
+
+        # 2) Broadcast for any tracked requests not already notified
+        #    (e.g. request was registered but the EngineDeadError handler
+        #    missed it because it wasn't submitted to the dead stage yet).
+        for req_id in list(self.request_states):
+            if req_id not in notified:
+                await self.output_async_queue.put(
+                    ErrorMessage(
+                        error=self._fatal_error,
+                        fatal=True,
+                        request_id=req_id,
+                        stage_id=self._fatal_error_stage_id,
+                    )
+                )
+            self.request_states.pop(req_id, None)
+
     def _shutdown_stages(self) -> None:
         """Shutdown all stage pools."""
         if self._stages_shutdown:

vllm_omni/engine/stage_pool.py

@@ -44,19 +44,6 @@ if TYPE_CHECKING:
 logger = init_logger(__name__)
 
 
-class StageUnavailableError(RuntimeError):
-    """Raised when dispatch cannot reach a live replica of a stage.
-
-    Deliberately a ``RuntimeError`` subclass (not ``EngineDeadError``): no
-    engine died at the raise site — the stage's replica set is empty or the
-    chosen slot was evicted, which is a routing/capacity condition and, in
-    distributed mode, recoverable. Dispatch guards catch exactly this type so
-    unrelated ``RuntimeError``s keep failing fast, and existing
-    ``EngineDeadError`` handlers (teardown-on-dead, poll-path eviction) are
-    not silently enrolled.
-    """
-
-
 @dataclass
 class _ReplicaMetrics:
     """Per-replica metrics accumulators owned by a stage pool."""
@@ -342,9 +329,7 @@ class StagePool:
 
             now = _time.monotonic()
             if now >= deadline:
-                raise StageUnavailableError(
-                    f"no UP replica for stage {self.stage_id} after {self.DISPATCH_WAIT_TIMEOUT_S:.1f}s"
-                )
+                raise RuntimeError(f"no UP replica for stage {self.stage_id} after {self.DISPATCH_WAIT_TIMEOUT_S:.1f}s")
             await asyncio.sleep(min(self.DISPATCH_RETRY_INTERVAL_S, deadline - now))
 
     def preselect_replica_id(
@@ -570,7 +555,7 @@ class StagePool:
             # Prefer replicas that are both live (client up) and available.
             live = [r for r in self.live_replica_ids() if self.is_replica_available(r)]
             if not live:
-                raise StageUnavailableError(f"stage {self.stage_id} has no live replicas")
+                raise RuntimeError(f"stage {self.stage_id} has no live replicas")
             if len(live) == 1:
                 chosen = live[0]
             else:
@@ -585,13 +570,13 @@ class StagePool:
     def _llm_client(self, replica_id: int) -> StagePoolLLMClient:
         client = self.clients[replica_id]
         if client is None:
-            raise StageUnavailableError(f"stage {self.stage_id} replica {replica_id} is not attached")
+            raise RuntimeError(f"stage {self.stage_id} replica {replica_id} is not attached")
         return cast(StagePoolLLMClient, client)
 
     def _diffusion_client(self, replica_id: int) -> StagePoolDiffusionClient:
         client = self.clients[replica_id]
         if client is None:
-            raise StageUnavailableError(f"stage {self.stage_id} replica {replica_id} is not attached")
+            raise RuntimeError(f"stage {self.stage_id} replica {replica_id} is not attached")
         return cast(StagePoolDiffusionClient, client)
 
     # ---- Metrics ----
@@ -976,7 +961,7 @@ class StagePool:
         )
         client = self.clients[replica_id]
         if client is None:
-            raise StageUnavailableError(f"stage {self.stage_id} replica {replica_id} is not attached")
+            raise RuntimeError(f"stage {self.stage_id} replica {replica_id} is not attached")
         try:
             self.output_processor.add_request(
                 request=request,
@@ -1025,7 +1010,7 @@ class StagePool:
 
         client = self.clients[replica_id]
         if client is None:
-            raise StageUnavailableError(f"stage {self.stage_id} replica {replica_id} is not attached")
+            raise RuntimeError(f"stage {self.stage_id} replica {replica_id} is not attached")
 
         if self.stage_type == "diffusion":
             if isinstance(request, list):
@@ -1261,21 +1246,3 @@ class StagePool:
                 replica_id,
                 e,
             )
-
-    def evict_replica(self, replica_id: int) -> None:
-        """Shut down a replica and remove it from the live set.
-
-        After eviction ``live_replica_ids`` / ``live_num_replicas`` no longer
-        include this slot, so the orchestrator stops polling and dispatching to
-        it (per-replica fault isolation). The slot is left as a ``None`` hole,
-        consistent with ``num_replicas``.
-        """
-        # replica_id should always come from live_replica_ids() (a valid, live
-        # index). Fail loud on an out-of-range id rather than letting a negative
-        # index wrap or IndexError on the next line; an ``assert`` would be
-        # stripped under ``python -O``, silently leaving a dead replica in the
-        # live set and breaking isolation.
-        if not 0 <= replica_id < len(self.clients):
-            raise ValueError(f"evict_replica: replica_id {replica_id} out of range (num_replicas={len(self.clients)})")
-        self.shutdown_replica(replica_id)
-        self.clients[replica_id] = None

vllm_omni/entrypoints/async_omni.py

@@ -906,27 +906,17 @@ class AsyncOmni(EngineClient, OmniBase):
                         await self.duplex_lifecycle_events.put(msg)
                         continue
 
-                    if isinstance(msg, ErrorMessage):
-                        # Route request-scoped errors to that request's queue and
-                        # keep the loop alive. A request whose stage replica died
-                        # and was evicted gets a fatal error delivered here; only
-                        # that request fails (its consumer raises), the server
-                        # stays up for other stages/requests (#4285). A fatal
-                        # error without a request_id is a genuine engine-wide
-                        # death and falls through to the except handler below.
-                        if msg.request_id is not None:
-                            req_state = self.request_states.get(msg.request_id)
-                            if req_state is not None:
-                                await req_state.queue.put(msg)
-                            else:
-                                logger.warning(
-                                    "[%s] dropping error for unknown req %s",
-                                    self._name,
-                                    msg.request_id,
-                                )
-                            continue
-                        if not msg.fatal:
-                            continue
+                    if isinstance(msg, ErrorMessage) and not msg.fatal:
+                        req_state = self.request_states.get(msg.request_id)
+                        if req_state is not None:
+                            await req_state.queue.put(msg)
+                        else:
+                            logger.warning(
+                                "[%s] dropping non-fatal error for unknown req %s",
+                                self._name,
+                                msg.request_id,
+                            )
+                        continue
 
                     should_continue, _, stage_id, req_state = self._handle_output_message(msg)
                     if should_continue:
@@ -1327,11 +1317,10 @@ class AsyncOmni(EngineClient, OmniBase):
 
     @property
     def errored(self) -> bool:
-        """Whether the engine is in a process-fatal error state.
+        """Whether the engine is in a non-recoverable error state.
 
-        Delegates to ``OmniBase.errored``, which is true only when the
-        orchestrator thread is dead; per-stage liveness is reported via
-        ``check_health`` instead.  Redeclared here to satisfy the
+        Delegates to ``OmniBase.errored`` which checks the orchestrator
+        thread and all stage clients.  Redeclared here to satisfy the
         ``EngineClient`` abstract-property requirement (Python's ABC
         mechanism does not resolve abstract methods from sibling MRO
         entries).

vllm_omni/entrypoints/cli/serve.py

@@ -629,6 +629,15 @@ class OmniServeCommand(CLISubcommand):
             action="store_true",
             help="Enable VAE tiling for memory optimization (useful for mitigating OOM issues).",
         )
+        omni_config_group.add_argument(
+            "--vae-decode-precision",
+            choices=("fp32", "fp16", "bf16"),
+            default=None,
+            help=(
+                "Optional Wan decode-only VAE precision. The encoder remains FP32; "
+                "unset preserves the pipeline's existing precision behavior."
+            ),
+        )
 
         # Parallel weight loading (faster diffusion startup)
         omni_config_group.add_argument(

vllm_omni/entrypoints/omni_base.py

@@ -4,7 +4,7 @@ import os
 import time
 import weakref
 from collections.abc import Sequence
-from typing import TYPE_CHECKING, Any, Literal
+from typing import Any, Literal
 
 import huggingface_hub
 import vllm.envs as envs
@@ -31,9 +31,6 @@ from vllm_omni.model_executor.model_loader.weight_utils import download_weights_
 from vllm_omni.outputs import OmniRequestOutput
 from vllm_omni.utils.tracking_parser import TrackingNamespace
 
-if TYPE_CHECKING:
-    from vllm_omni.engine.stage_pool import StagePool, StagePoolClient
-
 logger = init_logger(__name__)
 
 
@@ -240,42 +237,6 @@ class OmniBase(PDDisaggregationMixin):
         """Expose engine stage configs for PD disaggregation detection and validation."""
         return self.engine.stage_configs
 
-    @staticmethod
-    def _replica_is_dead(client: StagePoolClient | None) -> bool:
-        """Whether a replica slot is evicted (``None``) or its engine is dead.
-
-        ``check_health()`` is the source of truth: it raises ``EngineDeadError``
-        when dead (inspecting the ``engine_dead`` flags internally), and the
-        diffusion client additionally runs a synchronous ``proc.is_alive()``
-        probe that catches a silent SIGKILL/segfault the monitor thread has not
-        flagged yet — so ``/health`` does not report 200 for a stage whose
-        subprocess is already dead. A client exposing no ``check_health`` cannot
-        confirm liveness, so it is treated as dead (fail closed).
-        """
-        if client is None:
-            return True
-        check_health = getattr(client, "check_health", None)
-        if not callable(check_health):
-            return True
-        try:
-            check_health()
-        except EngineDeadError:
-            return True
-        return False
-
-    def _live_replica_count(self, pool: StagePool) -> int:
-        """Number of replicas in ``pool`` that are neither evicted nor dead.
-
-        A dead replica is evicted from its pool (slot set to ``None`` in
-        ``StagePool.clients``); a replica that died but has not been evicted
-        yet still carries an ``engine_dead`` flag. Both are excluded.
-        """
-        return sum(1 for client in pool.clients if not self._replica_is_dead(client))
-
-    def _stage_has_no_live_replica(self, pool: StagePool) -> bool:
-        """True when a non-empty stage pool has lost all of its replicas."""
-        return len(pool.clients) > 0 and self._live_replica_count(pool) == 0
-
     def _consumed_metric_message_ids(self, request_id: str) -> set[int]:
         consumed_by_request = getattr(self, "_consumed_metric_messages", None)
         if consumed_by_request is None:
@@ -283,33 +244,38 @@ class OmniBase(PDDisaggregationMixin):
             self._consumed_metric_messages = consumed_by_request
         return consumed_by_request.setdefault(request_id, set())
 
+    def _has_dead_stage(self) -> bool:
+        for stage_client in self.engine.stage_clients:
+            if getattr(stage_client, "_engine_dead", False):
+                return True
+            resources = getattr(stage_client, "resources", None)
+            if resources is not None and getattr(resources, "engine_dead", False):
+                return True
+        return False
+
     @property
     def is_running(self) -> bool:
-        return self.engine.is_alive()
+        return self.engine.is_alive() and not self._has_dead_stage()
 
     @property
     def errored(self) -> bool:
-        """Whether the engine is in a process-fatal error state.
+        """Whether the engine is in a non-recoverable error state.
 
-        True only when the orchestrator thread is dead. Per-stage liveness is
-        deliberately excluded: the OpenAI serving paths precheck ``errored``
-        before request routing, so including it would reject requests that do
-        not touch the dead stage (e.g. text-only chat when the talker stage is
-        down). A fully dead stage instead fails only the requests routed
-        through it (dispatch guards) and flips readiness to 503 via
-        :meth:`check_health` (per-replica fault isolation, #4285).
+        True when the orchestrator thread is dead **or** any stage client
+        has been marked dead (e.g. diffusion worker OOM / process death).
+
+        Checks both ``_engine_dead`` (StageDiffusionClient) and
+        ``resources.engine_dead`` (StageEngineCoreClient / AsyncMPClient)
+        since the two client types store the flag differently.
         """
-        return not self.engine.is_alive()
+        return not self.engine.is_alive() or self._has_dead_stage()
 
     def check_health(self) -> None:
         if not self.engine.is_alive():
             raise EngineDeadError("Orchestrator process is not alive")
-        pools = getattr(self.engine, "stage_pools", None)
-        if pools is None:
-            return
-        for pool in pools:
-            if self._stage_has_no_live_replica(pool):
-                raise EngineDeadError(f"Stage-{pool.stage_id} has no live replica")
+        for stage_client in self.engine.stage_clients:
+            if hasattr(stage_client, "check_health"):
+                stage_client.check_health()
 
     def resolve_sampling_params_list(
         self,

vllm_omni/model_executor/models/voxtral_tts/cuda_graph_acoustic_transformer_wrapper.py

@@ -74,12 +74,6 @@ class CUDAGraphAcousticTransformerWrapper:
 
         # Pre-create persistent buffers
         self.timesteps = torch.linspace(0, 1, self.n_steps + 1, device=device, dtype=dtype)
-        # Precompute schedule constants (dt, projected time embeddings) so the captured
-        # graph indexes a table instead of re-running them on every replay.
-        self.dts = self.timesteps[1:] - self.timesteps[:-1]
-        with torch.no_grad():
-            t_emb_table = self.acoustic_transformer.time_embedding(self.timesteps.view(-1, 1)).to(dtype)
-            self.t_proj_table = self.acoustic_transformer.time_projection(t_emb_table)
         self.fake_eos_one = torch.tensor(1.0, dtype=dtype, device=device)
         self.fake_eos_zero = torch.tensor(0.0, dtype=dtype, device=device)
 
@@ -148,21 +142,21 @@ class CUDAGraphAcousticTransformerWrapper:
 
         x = noise
 
-        # Hoist the step-invariant concat + llm_projection out of the loop; inside the
-        # captured graph these would otherwise be baked in and replayed every step.
-        llm_batched = torch.cat([hidden_states, torch.zeros_like(hidden_states)], dim=0)  # (2B, D)
-        llm_proj_batched = at.llm_projection(llm_batched)  # (2B, D)
+        # Pre-compute zero hidden states for unconditional CFG branch
+        hidden_states_zero = torch.zeros_like(hidden_states)
 
         timesteps = self.timesteps
         for i in range(len(timesteps) - 1):
-            dt = self.dts[i]
+            t = timesteps[i]
+            dt = timesteps[i + 1] - timesteps[i]
 
-            # Batch cond + uncond in a single pass; reuse the cached projected-t_emb row.
-            t_proj = self.t_proj_table[i].unsqueeze(0).expand(B, -1)
+            # Batch conditional + unconditional velocity in a single forward pass
+            t_emb = at.time_embedding(t.view(-1, 1).repeat(B, 1)).to(hidden_states.dtype)
             x_batched = torch.cat([x, x], dim=0)  # (2B, C)
-            t_proj_batched = t_proj.repeat(2, 1)  # (2B, D)
+            llm_batched = torch.cat([hidden_states, hidden_states_zero], dim=0)  # (2B, D)
+            t_emb_batched = t_emb.repeat(2, 1)  # (2B, D)
 
-            v_all = at._predict_velocity(x_t=x_batched, llm_proj=llm_proj_batched, t_proj=t_proj_batched)
+            v_all = at._predict_velocity(x_t=x_batched, llm_output=llm_batched, t_emb=t_emb_batched)
             v_t, uncond_v_t = v_all[:B], v_all[B:]
 
             # CFG combination (cfg_alpha is (B, 1), v_t is (B, C))
@@ -245,10 +239,6 @@ class CUDAGraphAcousticTransformerWrapper:
         if not self.enabled or not self._warmed_up:
             return self.model.compute_mm_logits(hidden_states, cfg_alpha=cfg_alpha)
 
-        # Inner graph replay is illegal during an outer stream capture.
-        if torch.cuda.is_current_stream_capturing():
-            return self.model.compute_mm_logits(hidden_states, cfg_alpha=cfg_alpha)
-
         padded_size = self._get_padded_size(actual_size)
         if padded_size is None or padded_size not in self.graphs:
             return self.model.compute_mm_logits(hidden_states, cfg_alpha=cfg_alpha)

vllm_omni/model_executor/models/voxtral_tts/voxtral_tts.py

@@ -45,12 +45,10 @@ def parse_batched_audio_input(input_ids: torch.Tensor, num_codebooks: int) -> tu
     """
     all_audio_tokens: list[torch.Tensor] = []
     all_ctx_frames: list[int] = []
-    # One D2H copy up-front avoids 2 syncs per request inside the loop.
-    header_view = input_ids.cpu() if input_ids.is_cuda else input_ids
     offset = 0
     while offset < input_ids.numel():
-        ctx_frames = int(header_view[offset])
-        context_length = int(header_view[offset + 1])
+        ctx_frames = int(input_ids[offset].item())
+        context_length = int(input_ids[offset + 1].item())
         offset += 2
 
         req_input_ids_length = (ctx_frames + context_length) * num_codebooks
@@ -157,9 +155,6 @@ class VoxtralTTSForConditionalGeneration(
                 architectures=["VoxtralTTSAudioTokenizer"],
             )
             self.model = self.audio_tokenizer
-            # forward() returns a runtime-length per-request list; FULL CUDAGraphWrapper
-            # would freeze that length at capture.
-            self.supports_cudagraph_full = False
         else:
             raise ValueError("Invalid model stage")
 

vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_generation.py

@@ -441,8 +441,6 @@ class FlowMatchingAudioTransformer(nn.Module):
             torch.linspace(0, 1, self._n_steps + 1),
             persistent=False,
         )
-        # Lazy per-dtype cache of schedule constants: (timesteps, t_proj_table, dts) on device.
-        self._timesteps_cache: dict[torch.dtype, tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = {}
 
     def load_weight(self, weight: tuple[str, torch.Tensor]) -> str:
         params_dict = dict(self.named_parameters())
@@ -516,25 +514,12 @@ class FlowMatchingAudioTransformer(nn.Module):
         # Skip decoding if codebook 0 is [END_AUDIO] token.
         should_decode = semantic_code != self._end_audio_token_id
 
-        # acoustic_codes starts from x_0; generate directly on device to skip H2D.
-        x_0 = torch.randn(B, self.model_args.n_acoustic_codebook, dtype=llm_hidden.dtype, device=llm_hidden.device)
+        # acoustic_codes starts from x_0
+        x_0 = torch.randn(B, self.model_args.n_acoustic_codebook).to(dtype=llm_hidden.dtype, device=llm_hidden.device)
         x_0 = self._noise_scale * x_0
 
-        # Build the schedule constants once per dtype and reuse them every frame.
-        cache = self._timesteps_cache.get(llm_hidden.dtype)
-        if cache is None:
-            timesteps_d = self._timesteps.to(dtype=llm_hidden.dtype)
-            t_emb_table = self.time_embedding(timesteps_d.view(-1, 1)).to(llm_hidden.dtype)
-            t_proj_table = self.time_projection(t_emb_table)
-            dts = timesteps_d[1:] - timesteps_d[:-1]
-            self._timesteps_cache[llm_hidden.dtype] = (timesteps_d, t_proj_table, dts)
-            cache = self._timesteps_cache[llm_hidden.dtype]
-        timesteps, t_proj_table, dts = cache
-
-        # Hoist the step-invariant concat + llm_projection out of the loop; inside the
-        # captured graph these would otherwise replay every step.
-        llm_batched = torch.cat([llm_hidden, torch.zeros_like(llm_hidden)], dim=0)
-        llm_proj_batched = self.llm_projection(llm_batched)
+        timesteps = self._timesteps.to(dtype=llm_hidden.dtype)
+        llm_hidden_zero = torch.zeros_like(llm_hidden)
 
         # Reshape cfg_alpha for broadcasting: (B,) -> (B, 1)
         cfg_alpha = cfg_alpha.to(dtype=llm_hidden.dtype, device=llm_hidden.device)
@@ -543,19 +528,20 @@ class FlowMatchingAudioTransformer(nn.Module):
         # Euler integration with batched conditional + unconditional velocity
         sampled = x_0
         for i in range(len(timesteps) - 1):
-            dt = dts[i]  # precomputed constant step size
+            t = timesteps[i]
+            dt = timesteps[i + 1] - timesteps[i]
 
-            # Reuse cached projected-t_emb row, expanded along batch dim.
-            t_proj = t_proj_table[i].unsqueeze(0).expand(B, -1)
+            t_emb = self.time_embedding(t.view(-1, 1).repeat(B, 1)).to(llm_hidden.dtype)
 
             # Batch cond + uncond into a single forward pass (2B batch)
             x_batched = torch.cat([sampled, sampled], dim=0)
-            t_proj_batched = torch.cat([t_proj, t_proj], dim=0)
+            llm_batched = torch.cat([llm_hidden, llm_hidden_zero], dim=0)
+            t_emb_batched = torch.cat([t_emb, t_emb], dim=0)
 
             v_all = self._predict_velocity(
                 x_t=x_batched,
-                llm_proj=llm_proj_batched,
-                t_proj=t_proj_batched,
+                llm_output=llm_batched,
+                t_emb=t_emb_batched,
             )
             v_t, uncond_v_t = v_all[:B], v_all[B:]
             v_t = cfg_alpha * v_t + (1 - cfg_alpha) * uncond_v_t
@@ -572,16 +558,18 @@ class FlowMatchingAudioTransformer(nn.Module):
     def _predict_velocity(
         self,
         x_t: torch.Tensor,  # BxC
-        llm_proj: torch.Tensor,  # BxD, already through llm_projection
-        t_proj: torch.Tensor,  # BxD, already through time_projection
+        llm_output: torch.Tensor,  # BxD
+        t_emb: torch.Tensor,  # BxD
     ) -> torch.Tensor:
-        # llm_proj/t_proj arrive pre-projected (hoisted out of the per-step loop by callers).
-        x_t = x_t.to(llm_proj.dtype)
+        x_t = x_t.to(llm_output.dtype)
+
+        t_emb = self.time_projection(t_emb)
+        llm_output = self.llm_projection(llm_output)
 
         acoustic_and_semantic_embeddings = [
             self.input_projection(x_t.unsqueeze(1)),  # Bx1xD
-            t_proj.unsqueeze(1),
-            llm_proj.unsqueeze(1),
+            t_emb.unsqueeze(1),
+            llm_output.unsqueeze(1),
         ]
         acoustic_transformer_inputs = torch.concatenate(acoustic_and_semantic_embeddings, dim=1)
 
@@ -966,8 +954,6 @@ class VoxtralTTSAudioGenerationForConditionalGeneration(nn.Module, SupportsMulti
         self.audio_tok_id = audio_encoder.audio_token
         self.eos_tok_id = self.tokenizer.instruct.tokenizer.eos_id
         self.vocab_size = config.text_config.vocab_size
-        self._end_audio_token_id = AudioSpecialTokens.id(AudioSpecialTokens.end_audio)
-        self._fake_eos_consts: dict[torch.device, tuple[torch.Tensor, torch.Tensor]] = {}
 
     def get_language_model(self) -> torch.nn.Module:
         return self.language_model
@@ -1066,15 +1052,11 @@ class VoxtralTTSAudioGenerationForConditionalGeneration(nn.Module, SupportsMulti
             llm_hidden=hidden_states,
             cfg_alpha=cfg_alpha,
         )
-        # Cache device-resident 1.0/0.0 scalars to avoid a per-call H2D transfer.
-        consts = self._fake_eos_consts.get(audio_codes.device)
-        if consts is None:
-            consts = (
-                torch.tensor(1.0, dtype=torch.bfloat16, device=audio_codes.device),
-                torch.tensor(0.0, dtype=torch.bfloat16, device=audio_codes.device),
-            )
-            self._fake_eos_consts[audio_codes.device] = consts
-        fake_eos = torch.where(audio_codes[:, 0] == self._end_audio_token_id, consts[0], consts[1])
+        fake_eos = torch.where(
+            audio_codes[:, 0] == AudioSpecialTokens.id(AudioSpecialTokens.end_audio),
+            torch.tensor(1.0, dtype=torch.bfloat16),
+            torch.tensor(0.0, dtype=torch.bfloat16),
+        )
         # BxC -> Bx1xC since this is per-step
         # Make it a list for vllm-omni processing
         audio_list = list(torch.split(audio_codes.unsqueeze(1), 1, dim=0))

vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_tokenizer.py

@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-import logging
 import math
 from copy import deepcopy
 from dataclasses import dataclass
@@ -10,7 +9,6 @@ import numpy as np
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from torch.nn.utils.rnn import pad_sequence
 from vllm.config import VllmConfig
 
 from vllm_omni.model_executor.models.voxtral_tts.voxtral_tts_audio_generation import (
@@ -487,10 +485,9 @@ class Attention(nn.Module):
                 )
             return slopes  # shape [n_heads], values in (2^-8, 1]
 
-        # Keep alibi slopes in fp32 so `forward` can skip the per-call `.to(fp32)` cast.
         self.register_buffer(
             "alibi_slopes",
-            get_alibi_slopes(self.n_local_heads).to(torch.float32).contiguous(),
+            get_alibi_slopes(self.n_local_heads),
             persistent=False,
         )
 
@@ -555,8 +552,8 @@ class Attention(nn.Module):
         # rel_pos[i, j] = j - i
         rel_pos = positions.unsqueeze(0) - positions.unsqueeze(1)  # (S, S)
 
-        # Alibi bias: slope * (j - i), shape (H, S, S). Cast only to xq dtype (already on device).
-        alibi_slopes = self.alibi_slopes.to(dtype=xq.dtype)
+        # Alibi bias: slope * (j - i), shape (H, S, S)
+        alibi_slopes = self.alibi_slopes.to(dtype=xq.dtype, device=xq.device)
         attn_bias = alibi_slopes.view(H, 1, 1) * rel_pos.unsqueeze(0).to(xq.dtype)
 
         # Causal mask: mask out j > i (future positions)
@@ -587,7 +584,7 @@ class Attention(nn.Module):
         xv = xv.view(bsz, seqlen, self.n_local_kv_heads, self.args.head_dim)
 
         if HAS_FLASH_ATTN:
-            # alibi_slopes already fp32 — avoid per-call cast.
+            alibi_slopes = self.alibi_slopes.to(torch.float32)
             output = flash_attn_func(
                 xq,
                 xk,
@@ -597,7 +594,7 @@ class Attention(nn.Module):
                     self.sliding_window,
                     0 if self.args.causal else self.sliding_window,
                 ),
-                alibi_slopes=self.alibi_slopes,
+                alibi_slopes=alibi_slopes,
             )
         else:
             output = self._native_attention(xq, xk, xv)
@@ -1055,15 +1052,14 @@ class VoxtralTTSAudioTokenizer(nn.Module):
         """
         chunk_size = 375  # TODO(chenyos): Hardcode. Fix it
 
-        # Compute per-request EOA cutting points in a single D2H instead of one sync per request.
-        flags = []
-        for c in codes_list:
-            eoa_mask = c[:, 0] == 1
-            flags.append(torch.stack([eoa_mask.any().to(torch.long), eoa_mask.long().argmax()]))
-        resolved = torch.stack(flags).cpu().tolist() if flags else []
-        cutting_points = [int(idx) if has else c.shape[0] for (has, idx), c in zip(resolved, codes_list)]
-
-        processed = [c[:cp] - 2 for c, cp in zip(codes_list, cutting_points)]
+        # Pre-process: find EOA and unshift tokens
+        processed = []
+        for codes in codes_list:
+            eoa_mask = codes[:, 0] == 1
+            eoa_indices = eoa_mask.nonzero(as_tuple=False)
+            cutting_point = eoa_indices[0].item() if len(eoa_indices) > 0 else len(codes)
+            audio_tokens = codes[:cutting_point] - 2
+            processed.append(audio_tokens)
 
         # Separate empty and non-empty
         results: list[torch.Tensor | None] = [None] * len(processed)
@@ -1091,19 +1087,25 @@ class VoxtralTTSAudioTokenizer(nn.Module):
             chunk_map.append((orig_idx, req_chunk_indices))
 
         # Pad chunks to max length and batch decode
-        padded = pad_sequence(all_chunks, batch_first=True, padding_value=0.0)
+        max_chunk_len = max(chunk_lengths)
+        K = all_chunks[0].shape[1]
+        padded = torch.zeros(
+            len(all_chunks),
+            max_chunk_len,
+            K,
+            dtype=all_chunks[0].dtype,
+            device=all_chunks[0].device,
+        )
+        for i, chunk in enumerate(all_chunks):
+            padded[i, : len(chunk)] = chunk
 
         audio_codes = padded.to(device=current_omni_platform.device_type)  # [B, T, K]
         audio_values = self.decode(audio_codes.transpose(1, 2), dtype=torch.bfloat16)  # [B, 1, T_out]
         audio_values = audio_values.detach().cpu().float().squeeze(1)  # [B, T_out]
-        # Gate min/max reductions behind logger level.
-        if logger.isEnabledFor(logging.WARNING):
-            amin = float(audio_values.min())
-            amax = float(audio_values.max())
-            if amin < -1.0:
-                logger.warning("Min value of decoded waveform signal is %s", amin)
-            if amax > 1.0:
-                logger.warning("Max value of decoded waveform signal is %s", amax)
+        if torch.min(audio_values) < -1.0:
+            logger.warning("Min value of decoded waveform signal is %s", torch.min(audio_values))
+        if torch.max(audio_values) > 1.0:
+            logger.warning("Max value of decoded waveform signal is %s", torch.max(audio_values))
 
         # Trim padding and reassemble per request
         for orig_idx, chunk_indices in chunk_map: