fix: localize NumPy image conversion

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

docs/design/feature/vae_parallel.md

@@ -449,21 +449,20 @@ Complete examples in the codebase:
 
 ---
 
-## Spatially-Sharded Decode (Wan and QwenImage)
+## Spatially-Sharded Decode (Wan)
 
-The tile-parallel executor above assigns independent spatial **tiles** to ranks. The Wan VAE and distributed QwenImage backend additionally support a **spatially-sharded decode** backend. `DiffusionParallelConfig.vae_parallel_mode="auto"` is the default; explicit `"tile"`, `"spatial_shard_height"`, and `"spatial_shard_width"` modes remain available.
+The tile-parallel executor above assigns independent spatial **tiles** to ranks. The Wan VAE additionally supports a **spatially-sharded decode** backend, selected through `DiffusionParallelConfig.vae_parallel_mode` (`"spatial_shard_height"` or `"spatial_shard_width"`; default is `"tile"`).
 
 ### How it differs from tile parallel
 
-| Aspect | Tile parallel (`"tile"`) | Automatic (`"auto"`) | Spatially-sharded (`"spatial_shard_height"`/`"spatial_shard_width"`) |
-|--------|--------------------------|----------------------|-----------------------------------------------|
-| Unit of work | Independent overlapping tiles | Chooses tile or a spatial shard per request | A single global feature map sharded along H or W |
-| Cross-rank communication | Gather tiles to rank 0, stitch + blend | Depends on the selected path | Per-conv **halo exchange** of boundary rows/cols (P2P) |
-| Output assembly | Blend overlapping tiles | Depends on the selected path | All-gather shards on rank 0 and trim padding; QwenImage then broadcasts |
-| Activation memory | Bounded by tile size | Depends on the selected path | Grows with full output area divided across ranks |
-| Scope | Decode + encode | Wan/QwenImage decode selection; encode remains tiled | Decode only |
+| Aspect | Tile parallel (`"tile"`) | Spatially-sharded (`"spatial_shard_height"`/`"spatial_shard_width"`) |
+|--------|--------------------------|-----------------------------------------------|
+| Unit of work | Independent overlapping tiles | A single global feature map sharded along H or W |
+| Cross-rank communication | Gather tiles to rank 0, stitch + blend | Per-conv **halo exchange** of boundary rows/cols (P2P) |
+| Output assembly | Blend overlapping tiles | All-gather shards on rank 0, trim padding (matches `broadcast_result=False`) |
+| Scope | Decode + encode | Decode only |
 
-Spatial-shard decode swaps the decoder's spatial convolutions/padding for context-aware halo-exchanging variants so each rank only holds a shard of the activations but still sees the correct receptive field at shard boundaries. Wan and QwenImage share the padding, 2D convolution, attention, halo, and trim primitives in `wan_spatial_shard.py`; `qwenimage_spatial_shard.py` supplies a causal-convolution subclass that remains discoverable by Diffusers' QwenImage cache accounting.
+Spatial-shard decode swaps the decoder's spatial convolutions/padding for halo-exchanging variants (`WanDistConv2d`, `WanDistCausalConv3d`, `WanDistZeroPad2d`) so each rank only holds a shard of the activations but still sees the correct receptive field at shard boundaries. Implementation lives in `vllm_omni/diffusion/distributed/autoencoders/wan_spatial_shard.py`.
 
 ### Wiring
 
@@ -474,17 +473,15 @@ serve.py (--vae-parallel-mode) / OmniEngineArgs
   -> DiffusionParallelConfig.vae_parallel_mode
   -> registry.py: model.vae.set_parallel_size(vae_pp_size, mode=...)
   -> DistributedVaeExecutor.parallel_mode
-  -> capability-based pipeline preparation
-  -> DistributedAutoencoderKLWan/DistributedAutoencoderKLQwenImage.tiled_decode dispatch
+  -> DistributedAutoencoderKLWan.tiled_decode dispatch
 ```
 
-The shared spatial policy gates the path: it requires distributed decode to be enabled, a 5D latent, a supported VAE capability, and `vae_patch_parallel_size == DiT group size`. The method is reached only after Diffusers has selected tiled decode for the request. In `auto` mode the longer axis that can cover every rank is selected, with width winning ties. This policy is cached only through normal kernel shape caches; it is reevaluated for every request.
+`DistributedAutoencoderKLWan._spatial_shard_decode_enabled()` gates the path: it requires distributed decode to be enabled, a 5D latent, and `vae_patch_parallel_size == DiT group size`. Otherwise it logs a warning and falls back to tile-parallel decode.
 
 ### Notes
 
-- The decoder is patched **in place** the first time spatial-shard decode runs. Its wrappers consult a context-local split dimension and retain an exact direct fallback, allowing later requests to select tile, height, or width safely.
+- The decoder is patched **in place** the first time spatial-shard decode runs and is bound to a single split dimension for the lifetime of the VAE instance.
 - Numerical correctness vs. single-GPU decode is covered by `tests/diffusion/distributed/test_wan_spatial_shard.py::test_spatial_shard_decode_matches_reference` (multi-GPU, nightly `full_model` + `distributed_cuda`).
-- QwenImage additionally covers cache discovery, compile compatibility, mixed tile/spatial requests, both axes, and all-rank result broadcast in `tests/diffusion/distributed/test_qwenimage_spatial_shard.py`.
 
 ---
 

docs/user_guide/diffusion/parallelism/vae_parallelism.md

@@ -15,7 +15,7 @@
 
 ## Overview
 
-VAE parallelism distributes VAE (Variational AutoEncoder) decode/encode work across multiple GPUs. This guide covers VAE patch/tile parallelism, which splits latent space into spatial tiles or patches, and Wan/QwenImage spatial-shard decode, which shards decoder feature maps along height or width.
+VAE parallelism distributes VAE (Variational AutoEncoder) decode/encode work across multiple GPUs. This guide covers VAE patch/tile parallelism, which splits latent space into spatial tiles or patches, and Wan spatial-shard decode, which shards decoder feature maps along height or width.
 
 This is particularly useful for:
 - **High-resolution image generation** where VAE decode can become a memory bottleneck
@@ -113,7 +113,7 @@ In `DiffusionParallelConfig`:
 | Parameter | Type | Default | Description |
 |-----------|------|---------|-------------|
 | `vae_patch_parallel_size` | int | 1 | Number of GPUs for VAE patch/tile parallelism. Set to 2 or higher to enable. Should typically match `tensor_parallel_size` as they share the same process group. |
-| `vae_parallel_mode` | str | `"auto"` | VAE parallel decode strategy: `"auto"` (shape/model/topology-aware selection), `"tile"`, `"spatial_shard_height"`, or `"spatial_shard_width"`. Spatial sharding supports Wan and the distributed QwenImage backend. See [Spatially-Sharded Decode](#spatially-sharded-decode-wan-and-qwenimage). |
+| `vae_parallel_mode` | str | `"tile"` | VAE parallel decode strategy: `"tile"` (default tile/patch parallel decode), `"spatial_shard_height"`, or `"spatial_shard_width"` (spatially-sharded decode, Wan only). See [Spatially-Sharded Decode](#spatially-sharded-decode-wan). |
 
 Additional requirements:
 
@@ -126,9 +126,9 @@ Additional requirements:
 
 ---
 
-## Spatially-Sharded Decode (Wan and QwenImage)
+## Spatially-Sharded Decode (Wan)
 
-The default `vae_parallel_mode="auto"` selects spatial-shard decode only after a Wan or distributed QwenImage request is already large enough to enter Diffusers' tiled-decode path and the VAE uses the full DiT decode group. It splits the longer eligible spatial axis to reduce halo traffic. Direct-decode requests, partial groups, and other VAE families keep their existing path. Set `vae_parallel_mode="tile"` to force the previous behavior, or `"spatial_shard_height"` / `"spatial_shard_width"` to force an axis.
+The default `vae_parallel_mode="tile"` distributes whole tiles across ranks. For the **Wan** VAE there is an alternative decode strategy, **spatially-sharded decode**, selected via `vae_parallel_mode="spatial_shard_height"` or `vae_parallel_mode="spatial_shard_width"`.
 
 Instead of assigning independent tiles to ranks, spatial-shard decode shards the decoder feature maps along the height (`spatial_shard_height`) or width (`spatial_shard_width`) dimension and exchanges halo rows/columns between neighboring ranks around the spatial convolutions. This keeps the receptive field correct across shard boundaries, so the result matches the single-GPU decode within numerical tolerance.
 
@@ -141,7 +141,7 @@ omni = Omni(
     parallel_config=DiffusionParallelConfig(
         tensor_parallel_size=2,
         vae_patch_parallel_size=2,               # must match the DiT group size
-        vae_parallel_mode="auto",                # default
+        vae_parallel_mode="spatial_shard_width", # or "spatial_shard_height"
     ),
 )
 ```
@@ -152,18 +152,14 @@ Or from the CLI / serving entrypoint:
 vllm serve Wan-AI/Wan2.1-T2V-1.3B-Diffusers --omni \
     --tensor-parallel-size 2 \
     --vae-patch-parallel-size 2 \
-    --vae-parallel-mode auto
+    --vae-parallel-mode spatial_shard_width
 ```
 
 **Constraints and behavior:**
 
-- Spatial-shard decode is **decode-only** and implemented for the **Wan** VAE and the distributed **QwenImage** backend used by Qwen-Image and Krea2. Pipelines using a different VAE backend keep their existing behavior.
-- It requires `vae_patch_parallel_size` to **match the DiT process group size**. `auto` quietly falls back to tile decode for a partial group; an explicitly forced spatial mode logs a warning.
-- The decoder's context-aware wrappers retain the direct path, so one process can alternate small/large and portrait/landscape requests safely.
-- Wan preserves its rank-0-only decode result, while QwenImage preserves its existing all-rank result by broadcasting after rank-0 assembly.
-- Spatial sharding trades activation memory for latency as resolution grows. Use `vae_parallel_mode="tile"` when bounded VAE memory is more important than decode speed.
-
-The automatic policy is adapted from [SGLang #28071](https://github.com/sgl-project/sglang/pull/28071) and narrowed to vLLM-Omni's parity-tested Wan and QwenImage backends.
+- Spatial-shard decode is **decode-only** and currently implemented for the **Wan** VAE. Other models ignore `spatial_shard_*` modes.
+- It requires `vae_patch_parallel_size` to **match the DiT process group size**. If it does not, the VAE logs a warning and **falls back to tile-parallel decode** at runtime.
+- `spatial_shard_height` and `spatial_shard_width` are mutually exclusive for a given VAE instance (the decoder is patched in place for a single split dimension).
 
 For end-to-end latency/throughput, launch serving with the desired `vae_parallel_mode` and use the existing diffusion serving benchmark:
 

examples/offline_inference/text_to_image/text_to_image.py

@@ -8,7 +8,9 @@ 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
@@ -54,6 +56,17 @@ 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(
@@ -619,6 +632,7 @@ 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/config/test_omni_config.py

@@ -579,7 +579,6 @@ def test_diffusion_parallel_config_keeps_current_diffusion_parallel_surface():
     assert cfg.data_parallel_size == 3
     assert cfg.cfg_parallel_size == 3
     assert cfg.mask_sp_padding is True
-    assert cfg.vae_parallel_mode == "auto"
     assert cfg.world_size == 72
 
 

tests/dfx/reliability/test_reliability_qwen3_omni.py

@@ -315,11 +315,32 @@ 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(
@@ -386,9 +407,11 @@ 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:
@@ -431,7 +454,6 @@ 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:
@@ -774,7 +796,11 @@ def test_reliability_fault_process_kill_tree_with_load_fast_fail_and_cleanup(
 
 
 @pytest.mark.slow
-@pytest.mark.skip(reason="issue#4285")
+@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."
+)
 @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/distributed/test_qwenimage_spatial_shard.py

@@ -1,233 +0,0 @@
-import os
-from dataclasses import dataclass
-
-import pytest
-import torch
-import torch.distributed as dist
-import torch.multiprocessing as mp
-from diffusers.models.autoencoders.autoencoder_kl_qwenimage import QwenImageCausalConv3d
-
-from tests.helpers.mark import hardware_test
-from tests.helpers.runtime import get_open_port
-from vllm_omni.diffusion.distributed.autoencoders import qwenimage_spatial_shard
-from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_qwenimage import (
-    DistributedAutoencoderKLQwenImage,
-)
-from vllm_omni.diffusion.distributed.autoencoders.spatial_shard import (
-    prepare_pipeline_spatial_shard_decode,
-)
-from vllm_omni.diffusion.offloader.module_collector import PipelineModules
-from vllm_omni.platforms import current_omni_platform
-
-
-@dataclass
-class _FakeExecutor:
-    group: object
-    parallel_size: int
-    parallel_mode: str
-
-
-def _tiny_qwenimage_vae() -> DistributedAutoencoderKLQwenImage:
-    return DistributedAutoencoderKLQwenImage(
-        base_dim=8,
-        z_dim=4,
-        dim_mult=[1, 2, 2, 2],
-        num_res_blocks=1,
-        attn_scales=[],
-        temperal_downsample=[False, True, True],
-        input_channels=3,
-        latents_mean=[0.0] * 4,
-        latents_std=[1.0] * 4,
-    )
-
-
-@pytest.mark.core_model
-@pytest.mark.cpu
-def test_qwenimage_auto_spatial_shard_gate_is_shape_and_topology_aware(monkeypatch: pytest.MonkeyPatch):
-    monkeypatch.setattr(
-        "vllm_omni.diffusion.distributed.autoencoders.spatial_shard.dist.get_world_size",
-        lambda group=None: 4,
-    )
-    vae = DistributedAutoencoderKLQwenImage.__new__(DistributedAutoencoderKLQwenImage)
-    vae.use_tiling = True
-    vae.distributed_executor = _FakeExecutor(group=object(), parallel_size=4, parallel_mode="auto")
-    vae.is_distributed_enabled = lambda: True
-
-    landscape = torch.zeros((1, 16, 1, 60, 104))
-    portrait = torch.zeros((1, 16, 1, 104, 60))
-
-    assert vae._spatial_shard_decode_split_dim(landscape, 4) == "width"
-    assert vae._spatial_shard_decode_split_dim(portrait, 4) == "height"
-    assert vae._spatial_shard_decode_enabled(landscape) is True
-
-    vae.distributed_executor.parallel_size = 2
-    assert vae._spatial_shard_decode_enabled(landscape) is False
-
-
-@pytest.mark.core_model
-@pytest.mark.cpu
-def test_pipeline_dispatches_qwenimage_spatial_preparation(monkeypatch: pytest.MonkeyPatch):
-    vae = DistributedAutoencoderKLQwenImage.__new__(DistributedAutoencoderKLQwenImage)
-    prepared = []
-    vae._prepare_spatial_shard_decode = lambda: prepared.append(vae)
-
-    monkeypatch.setattr(
-        "vllm_omni.diffusion.offloader.module_collector.ModuleDiscovery.discover",
-        lambda pipeline: PipelineModules(
-            dits=[],
-            dit_names=[],
-            encoders=[],
-            encoder_names=[],
-            vaes=[torch.nn.Identity(), vae],
-        ),
-    )
-    monkeypatch.setattr(
-        "vllm_omni.diffusion.distributed.autoencoders.spatial_shard.dist.is_initialized",
-        lambda: True,
-    )
-
-    prepare_pipeline_spatial_shard_decode(torch.nn.Identity())
-
-    assert prepared == [vae]
-
-
-@pytest.mark.core_model
-@pytest.mark.cpu
-@torch.no_grad()
-def test_qwenimage_installed_decoder_preserves_cache_discovery_state_and_direct_path():
-    torch.manual_seed(0)
-    vae = _tiny_qwenimage_vae().eval()
-    latents = torch.randn(1, 4, 2, 4, 5)
-    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, return_dict=False)[0]
-    causal_conv_count = sum(isinstance(module, QwenImageCausalConv3d) for module in vae.decoder.modules())
-
-    qwenimage_spatial_shard.install_qwenimage_spatial_shard_decode(vae, group=object(), split_dim="height")
-    qwenimage_spatial_shard.install_qwenimage_spatial_shard_decode(vae, group=object(), split_dim="width")
-    vae.clear_cache()
-
-    assert causal_conv_count > 0
-    assert len(vae._feat_map) == causal_conv_count
-    assert sum(isinstance(module, QwenImageCausalConv3d) for module in vae.decoder.modules()) == causal_conv_count
-    assert any(
-        isinstance(module, qwenimage_spatial_shard.QwenImageDistCausalConv3d) for module in vae.decoder.modules()
-    )
-    assert {name for name, _ in vae.named_parameters()} == parameter_names
-    vae.load_state_dict(state_dict, strict=True)
-    assert torch.equal(vae.decode(latents, return_dict=False)[0], reference)
-
-
-@pytest.mark.core_model
-@pytest.mark.cpu
-@torch.no_grad()
-def test_qwenimage_installed_decoder_direct_path_can_be_compiled():
-    torch.manual_seed(1)
-    vae = _tiny_qwenimage_vae().eval()
-    latents = torch.randn(1, 4, 1, 3, 5)
-    qwenimage_spatial_shard.install_qwenimage_spatial_shard_decode(vae, group=object())
-    eager = vae.decode(latents, return_dict=False)[0]
-
-    compiled_decode = torch.compile(vae.decode, backend="eager", fullgraph=False)
-    compiled = compiled_decode(latents, return_dict=False)[0]
-
-    torch.testing.assert_close(compiled, eager, rtol=0, atol=0)
-
-
-def _qwenimage_spatial_shard_worker(rank: int, master_port: str, results) -> None:
-    from vllm_omni.diffusion.distributed.parallel_state import (
-        destroy_model_parallel,
-        init_distributed_environment,
-        initialize_model_parallel,
-    )
-
-    os.environ["MASTER_ADDR"] = "127.0.0.1"
-    os.environ["MASTER_PORT"] = master_port
-    device = current_omni_platform.get_torch_device(rank)
-    current_omni_platform.set_device(device)
-    backend = current_omni_platform.dist_backend
-    init_distributed_environment(world_size=2, rank=rank, local_rank=rank, backend=backend)
-    initialize_model_parallel(sequence_parallel_size=2, ulysses_degree=2, backend=backend)
-
-    try:
-        torch.manual_seed(2)
-        vae = _tiny_qwenimage_vae().to(device=device, dtype=torch.float32).eval()
-        vae.init_distributed()
-        generator = torch.Generator(device=device).manual_seed(3)
-        landscape = torch.randn((1, 4, 2, 5, 7), generator=generator, device=device)
-        portrait = torch.randn((1, 4, 2, 7, 5), generator=generator, device=device)
-
-        with torch.inference_mode():
-            vae.use_tiling = False
-            vae.set_parallel_size(1, mode="tile")
-            landscape_reference = vae.decode(landscape, return_dict=False)[0]
-            portrait_reference = vae.decode(portrait, return_dict=False)[0]
-
-            vae.enable_tiling(
-                tile_sample_min_height=32,
-                tile_sample_min_width=32,
-                tile_sample_stride_height=24,
-                tile_sample_stride_width=24,
-            )
-            vae.set_parallel_size(2, mode="tile")
-            tile_before_install = vae.decode(landscape, return_dict=False)[0]
-
-            vae.set_parallel_size(2, mode="auto")
-            assert vae._spatial_shard_decode_split_dim(landscape, 2) == "width"
-            landscape_auto = vae.decode(landscape, return_dict=False)[0]
-            assert vae._spatial_shard_decode_split_dim(portrait, 2) == "height"
-            portrait_auto = vae.decode(portrait).sample
-
-            vae.set_parallel_size(2, mode="spatial_shard_height")
-            landscape_height = vae.decode(landscape, return_dict=False)[0]
-            vae.set_parallel_size(2, mode="spatial_shard_width")
-            portrait_width = vae.decode(portrait, return_dict=False)[0]
-
-            vae.set_parallel_size(2, mode="tile")
-            tile_after_install = vae.decode(landscape, return_dict=False)[0]
-
-        assert tile_after_install.shape == tile_before_install.shape
-        assert torch.equal(tile_after_install, tile_before_install)
-        assert all(
-            buffer is None
-            for module in vae.decoder.modules()
-            for name, buffer in module._buffers.items()
-            if name in {"_halo_recv_top_buf", "_halo_recv_bottom_buf"}
-        )
-
-        diffs = {
-            "landscape_auto": (landscape_auto - landscape_reference).abs().max().item(),
-            "portrait_auto": (portrait_auto - portrait_reference).abs().max().item(),
-            "landscape_height": (landscape_height - landscape_reference).abs().max().item(),
-            "portrait_width": (portrait_width - portrait_reference).abs().max().item(),
-        }
-        results[rank] = {
-            "diffs": diffs,
-            "landscape_shape": tuple(landscape_auto.shape),
-            "portrait_shape": tuple(portrait_auto.shape),
-        }
-    finally:
-        destroy_model_parallel()
-        if dist.is_initialized():
-            dist.destroy_process_group()
-
-
-@pytest.mark.core_model
-@pytest.mark.diffusion
-@pytest.mark.parallel
-@hardware_test(res={"cuda": "L4"}, num_cards=2)
-def test_qwenimage_spatial_shard_matches_reference_broadcasts_and_switches_paths():
-    manager = mp.get_context("spawn").Manager()
-    results = manager.dict()
-    mp.spawn(
-        _qwenimage_spatial_shard_worker,
-        args=(str(get_open_port()), results),
-        nprocs=2,
-        join=True,
-    )
-
-    assert set(results) == {0, 1}
-    assert results[0]["landscape_shape"] == results[1]["landscape_shape"] == (1, 3, 5, 40, 56)
-    assert results[0]["portrait_shape"] == results[1]["portrait_shape"] == (1, 3, 5, 56, 40)
-    for rank_results in results.values():
-        assert max(rank_results["diffs"].values()) <= 2e-5

tests/diffusion/distributed/test_wan_spatial_shard.py

@@ -1,5 +1,4 @@
 import os
-from contextlib import contextmanager
 from types import SimpleNamespace
 
 import pytest
@@ -11,8 +10,6 @@ from tests.helpers.mark import hardware_test
 from vllm_omni.diffusion.distributed.autoencoders import wan_spatial_shard
 from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import (
     DistributedAutoencoderKLWan,
-    prepare_pipeline_wan_spatial_shard_decode,
-    select_auto_spatial_shard_split_dim,
 )
 from vllm_omni.platforms import current_omni_platform
 
@@ -20,23 +17,6 @@ from vllm_omni.platforms import current_omni_platform
 # the end uses hardware_test (H100 x2) plus full_model / diffusion / parallel.
 
 
-@contextmanager
-def _active_spatial_shard(split_dim: str):
-    token = wan_spatial_shard._SPATIAL_SHARD_CONTEXT.set(
-        wan_spatial_shard.SpatialShardContext(
-            input_extent=2,
-            local_input_extent=2,
-            split_dim=split_dim,
-            rank=0,
-            world_size=3,
-        )
-    )
-    try:
-        yield
-    finally:
-        wan_spatial_shard._SPATIAL_SHARD_CONTEXT.reset(token)
-
-
 @pytest.mark.core_model
 @pytest.mark.cpu
 def test_split_for_parallel_decode_pads_uneven_height():
@@ -275,14 +255,12 @@ def test_dist_zero_pad_only_applies_global_height_edges(monkeypatch: pytest.Monk
 
     monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (1, 3))
     mid_rank_pad = wan_spatial_shard.WanDistZeroPad2d((0, 1, 1, 1), group=object())
-    with _active_spatial_shard("height"):
-        mid = mid_rank_pad(x)
+    mid = mid_rank_pad(x)
     assert mid.shape == (1, 1, 2, 3)
 
     monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (2, 3))
     last_rank_pad = wan_spatial_shard.WanDistZeroPad2d((0, 1, 1, 1), group=object())
-    with _active_spatial_shard("height"):
-        last = last_rank_pad(x)
+    last = last_rank_pad(x)
     assert last.shape == (1, 1, 3, 3)
 
 
@@ -292,15 +270,13 @@ def test_dist_zero_pad_only_applies_global_width_edges(monkeypatch: pytest.Monke
     x = torch.ones((1, 1, 2, 2))
 
     monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (1, 3))
-    mid_rank_pad = wan_spatial_shard.WanDistZeroPad2d((1, 1, 0, 0), group=object())
-    with _active_spatial_shard("width"):
-        mid = mid_rank_pad(x)
+    mid_rank_pad = wan_spatial_shard.WanDistZeroPad2d((1, 1, 0, 0), group=object(), split_dim="width")
+    mid = mid_rank_pad(x)
     assert mid.shape == (1, 1, 2, 2)
 
     monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (2, 3))
-    last_rank_pad = wan_spatial_shard.WanDistZeroPad2d((1, 1, 0, 0), group=object())
-    with _active_spatial_shard("width"):
-        last = last_rank_pad(x)
+    last_rank_pad = wan_spatial_shard.WanDistZeroPad2d((1, 1, 0, 0), group=object(), split_dim="width")
+    last = last_rank_pad(x)
     assert last.shape == (1, 1, 2, 3)
 
 
@@ -323,113 +299,6 @@ def test_spatial_shard_height_gate_falls_back_for_partial_group(monkeypatch: pyt
     assert vae._spatial_shard_decode_enabled(z) is False
 
 
-@pytest.mark.core_model
-@pytest.mark.cpu
-@pytest.mark.parametrize(
-    "z_shape,world_size,expected",
-    [
-        ((1, 16, 1, 60, 104), 4, "width"),
-        ((1, 16, 5, 60, 104), 4, "width"),
-        ((1, 16, 5, 104, 60), 4, "height"),
-        ((1, 16, 5, 2, 2), 4, None),
-        ((1, 16, 5, 64, 64), 4, "width"),
-    ],
-)
-def test_auto_spatial_shard_selector_is_shape_aware(z_shape, world_size, expected):
-    assert select_auto_spatial_shard_split_dim(z_shape, world_size) == expected
-
-
-@pytest.mark.core_model
-@pytest.mark.cpu
-def test_auto_spatial_shard_gate_requires_full_group(monkeypatch: pytest.MonkeyPatch):
-    monkeypatch.setattr(
-        "vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan.dist.get_world_size",
-        lambda group=None: 4,
-    )
-
-    vae = DistributedAutoencoderKLWan.__new__(DistributedAutoencoderKLWan)
-    vae.use_tiling = True
-    vae.distributed_executor = SimpleNamespace(group=object(), parallel_size=4, parallel_mode="auto")
-    vae.is_distributed_enabled = lambda: True
-    large = torch.zeros((1, 16, 5, 60, 104))
-    unshardable = torch.zeros((1, 16, 1, 2, 2))
-
-    assert vae._spatial_shard_decode_split_dim(large, 4) == "width"
-    assert vae._spatial_shard_decode_enabled(large) is True
-    assert vae._spatial_shard_decode_enabled(unshardable) is False
-
-    vae.distributed_executor.parallel_size = 2
-    assert vae._spatial_shard_decode_enabled(large) is False
-
-
-@pytest.mark.core_model
-@pytest.mark.cpu
-def test_spatial_wrappers_are_prepared_only_for_a_full_group(monkeypatch: pytest.MonkeyPatch):
-    vae = DistributedAutoencoderKLWan.__new__(DistributedAutoencoderKLWan)
-    object.__setattr__(
-        vae,
-        "distributed_executor",
-        SimpleNamespace(group=object(), parallel_size=2, parallel_mode="auto"),
-    )
-    installed = []
-
-    monkeypatch.setattr(
-        "vllm_omni.diffusion.offloader.module_collector.ModuleDiscovery.discover",
-        lambda pipeline: SimpleNamespace(vaes=[vae]),
-    )
-    monkeypatch.setattr(
-        "vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan.dist.is_initialized",
-        lambda: True,
-    )
-    monkeypatch.setattr(
-        "vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan.dist.get_world_size",
-        lambda group=None: 4,
-    )
-    monkeypatch.setattr(
-        wan_spatial_shard,
-        "install_wan_spatial_shard_decode",
-        lambda candidate, group, split_dim: installed.append((candidate, group, split_dim)),
-    )
-
-    prepare_pipeline_wan_spatial_shard_decode(torch.nn.Identity())
-    assert installed == []
-
-    vae.distributed_executor.parallel_size = 4
-    prepare_pipeline_wan_spatial_shard_decode(torch.nn.Identity())
-    assert installed == [(vae, vae.distributed_executor.group, "height")]
-
-
-@pytest.mark.core_model
-@pytest.mark.cpu
-@torch.no_grad()
-def test_installed_spatial_decoder_keeps_direct_fallback_and_can_change_axis():
-    vae = DistributedAutoencoderKLWan(
-        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,
-    ).eval()
-    latents = torch.randn(1, 4, 1, 4, 5)
-    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, return_dict=False)[0]
-
-    wan_spatial_shard.install_wan_spatial_shard_decode(vae, group=object(), split_dim="height")
-    assert any(getattr(module, "_vllm_omni_dynamic_spatial_shard_conv", False) for module in vae.decoder.modules())
-    wan_spatial_shard.install_wan_spatial_shard_decode(vae, group=object(), split_dim="width")
-
-    assert {name for name, _ in vae.named_parameters()} == parameter_names
-    vae.load_state_dict(state_dict, strict=True)
-    assert torch.equal(vae.decode(latents, return_dict=False)[0], reference)
-
-
 @pytest.mark.core_model
 @pytest.mark.cpu
 def test_spatial_shard_width_gate_selects_width(monkeypatch: pytest.MonkeyPatch):
@@ -575,99 +444,3 @@ def test_spatial_shard_decode_matches_reference(split_dim: str):
     assert mean_abs_diff <= _SPATIAL_SHARD_TOLERANCE, (
         f"spatial_shard_{split_dim} mean_abs_diff {mean_abs_diff} exceeds {_SPATIAL_SHARD_TOLERANCE}"
     )
-
-
-def _auto_spatial_shard_decode_worker(rank: int, return_dict, master_port: str) -> None:
-    from vllm_omni.diffusion.distributed.parallel_state import (
-        destroy_model_parallel,
-        init_distributed_environment,
-        initialize_model_parallel,
-    )
-
-    os.environ["MASTER_ADDR"] = "127.0.0.1"
-    os.environ["MASTER_PORT"] = master_port
-    device = current_omni_platform.get_torch_device(rank)
-    current_omni_platform.set_device(device)
-    backend = current_omni_platform.dist_backend
-    init_distributed_environment(world_size=2, rank=rank, local_rank=rank, backend=backend)
-    initialize_model_parallel(sequence_parallel_size=2, ulysses_degree=2, backend=backend)
-
-    try:
-        vae = DistributedAutoencoderKLWan.from_pretrained(
-            _SPATIAL_SHARD_MODEL,
-            subfolder=_SPATIAL_SHARD_SUBFOLDER,
-            torch_dtype=torch.float32,
-        ).to(device=device, dtype=torch.float32)
-        vae.eval()
-        generator = torch.Generator(device=device).manual_seed(1)
-        small = torch.randn((1, vae.config.z_dim, 1, 60, 104), generator=generator, device=device)
-        landscape = torch.randn((1, vae.config.z_dim, 2, 60, 104), generator=generator, device=device)
-        portrait = torch.randn((1, vae.config.z_dim, 2, 104, 60), generator=generator, device=device)
-
-        with torch.inference_mode():
-            vae.use_tiling = False
-            vae.set_parallel_size(1, mode="tile")
-            small_reference = vae.decode(small, return_dict=False)[0]
-
-            vae.use_tiling = True
-            vae.set_parallel_size(2, mode="tile")
-            tile_before_install = vae.decode(small, return_dict=False)[0]
-
-            vae.set_parallel_size(2, mode="auto")
-            assert vae._spatial_shard_decode_split_dim(small, 2) == "width"
-            small_sharded = vae.decode(small, return_dict=False)[0]
-            assert vae._spatial_shard_decode_split_dim(landscape, 2) == "width"
-            vae.decode(landscape, return_dict=False)
-
-            vae.set_parallel_size(2, mode="tile")
-            tile_after_install = vae.decode(small, return_dict=False)[0]
-
-            vae.use_tiling = False
-            vae.set_parallel_size(1, mode="tile")
-            portrait_reference = vae.decode(portrait, return_dict=False)[0]
-
-            vae.use_tiling = True
-            vae.set_parallel_size(2, mode="auto")
-            assert vae._spatial_shard_decode_split_dim(portrait, 2) == "height"
-            portrait_sharded = vae.decode(portrait, return_dict=False)[0]
-
-        assert all(
-            buffer is None
-            for module in vae.decoder.modules()
-            for name, buffer in module._buffers.items()
-            if name in {"_halo_recv_top_buf", "_halo_recv_bottom_buf"}
-        )
-
-        if rank == 0:
-            small_diff = (small_sharded.float() - small_reference.float()).abs()
-            portrait_diff = (portrait_sharded.float() - portrait_reference.float()).abs()
-            return_dict["tile_fallback_exact"] = torch.equal(tile_after_install, tile_before_install)
-            return_dict["small_max_abs_diff"] = small_diff.max().item()
-            return_dict["small_mean_abs_diff"] = small_diff.mean().item()
-            return_dict["portrait_max_abs_diff"] = portrait_diff.max().item()
-            return_dict["portrait_mean_abs_diff"] = portrait_diff.mean().item()
-    finally:
-        destroy_model_parallel()
-        if dist.is_initialized():
-            dist.destroy_process_group()
-
-
-@pytest.mark.full_model
-@pytest.mark.diffusion
-@pytest.mark.parallel
-@hardware_test(res={"cuda": "H100"}, num_cards=2)
-def test_auto_spatial_shard_switches_shape_paths_safely():
-    manager = mp.get_context("spawn").Manager()
-    return_dict = manager.dict()
-    mp.spawn(
-        _auto_spatial_shard_decode_worker,
-        args=(return_dict, "29502"),
-        nprocs=2,
-        join=True,
-    )
-
-    assert return_dict.get("tile_fallback_exact") is True
-    assert return_dict["small_max_abs_diff"] <= _SPATIAL_SHARD_TOLERANCE
-    assert return_dict["small_mean_abs_diff"] <= _SPATIAL_SHARD_TOLERANCE
-    assert return_dict["portrait_max_abs_diff"] <= _SPATIAL_SHARD_TOLERANCE
-    assert return_dict["portrait_mean_abs_diff"] <= _SPATIAL_SHARD_TOLERANCE

tests/diffusion/models/minimax_h3/test_minimax_h3_parallel.py

@@ -120,34 +120,3 @@ def test_tp_accepts_checkpoint_supported_sizes():
     arch = MiniMaxH3DiTArchConfig()
     for tp_size in (1, 2, 4, 7):
         model._validate_tp_config(arch=arch, tp_size=tp_size)
-
-
-@pytest.mark.parametrize(("parallel_size", "enabled"), [(1, False), (2, True)])
-def test_video_vae_auto_mode_preserves_native_tile_parallelism(monkeypatch, parallel_size, enabled):
-    from vllm_omni.diffusion.models.minimax_h3 import vae as vae_module
-
-    parallel_state = {}
-    video_vae = object.__new__(vae_module.MiniMaxH3VideoVAE)
-    nn.Module.__init__(video_vae)
-    video_vae.remote = nn.Identity()
-    video_vae.model = type("FakeModel", (), {"parallel_tiling": False})()
-
-    monkeypatch.setattr(vae_module, "get_dit_group", lambda: object())
-    monkeypatch.setattr(vae_module.dist, "get_world_size", lambda group: 2)
-    monkeypatch.setattr(vae_module.dist, "get_rank", lambda group: 0)
-    monkeypatch.setattr(
-        vae_module.importlib,
-        "import_module",
-        lambda package: type(
-            "FakeParallelModule",
-            (),
-            {"get_parallel_state": staticmethod(lambda: parallel_state)},
-        )(),
-    )
-
-    video_vae.set_parallel_size(parallel_size, mode="auto")
-
-    assert video_vae.parallel_size == parallel_size
-    assert video_vae.model.parallel_tiling is enabled
-    assert parallel_state["group_size"] == parallel_size
-    assert parallel_state["sp_enabled"] is enabled

tests/diffusion/test_diffusion_model_runner.py

@@ -728,19 +728,6 @@ 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
-    preparation_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
@@ -788,22 +775,16 @@ 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,
-        "prepare_pipeline_spatial_shard_decode",
-        lambda pipeline: preparation_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, memory_pool_context_fn=_memory_pool_context)
+    DiffusionModelRunner.load_model(runner)
 
     assert runner.cache_backend is None
     assert runner.od_config.cache_backend is None
     assert dummy_cache_backend.enabled is False
-    assert preparation_pool_states == [True]
 
 
 @pytest.mark.core_model

tests/engine/test_cfg_companion_lifecycle.py

@@ -83,6 +83,11 @@ 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,8 +74,6 @@ 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()
@@ -91,8 +89,7 @@ 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 fatal.error == "stage engine died"
-    assert fatal.stage_id == 2
+    assert "stage engine died" in fatal.error
 
 
 @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 + shutdown
+- EngineDeadError from an LLM stage poll → fatal error broadcast + replica eviction
 - Diffusion stage error output (OmniRequestOutput.from_error) → routed correctly
 """
 
@@ -14,17 +14,30 @@ 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 EngineQueueMessage, ErrorMessage, ShutdownRequestMessage
+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.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,
 )
@@ -50,6 +63,22 @@ 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] = []
@@ -72,32 +101,14 @@ def orchestrator_factory():
 # ───────── EngineDeadError from LLM stage poll ─────────
 
 
-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.
-    """
-
-    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:
+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. Shut itself down (thread exits)
+    2. Evict the dead replica and keep running so the API server stays up
+       (per-replica fault isolation, #4285)
     """
-    stage0 = FakeDeadLLMStageClient(stage_type="llm", final_output=True)
+    stage0 = _DieOnDemandLLMStageClient(stage_type="llm", final_output=True)
     orchestrator_fixture = orchestrator_factory([stage0])
     request = SimpleNamespace(request_id="req-dead", prompt_token_ids=[1, 2])
 
@@ -110,6 +121,12 @@ async def test_engine_dead_error_broadcasts_fatal_and_shuts_down(orchestrator_fa
             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)
@@ -118,13 +135,17 @@ async def test_engine_dead_error_broadcasts_fatal_and_shuts_down(orchestrator_fa
         assert msg.type == "error"
         assert msg.fatal is True
         assert msg.request_id == "req-dead"
-        assert "Stage-0 engine core is dead" in msg.error
+        assert "Stage-0 replica-0 is dead" in msg.error
 
-        # The orchestrator thread should exit after the fatal error.
-        orchestrator_fixture.thread.join(timeout=5)
-        assert not orchestrator_fixture.thread.is_alive()
+        # 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()
 
-        # Request state should be cleaned up.
+        # Affected request state should be cleaned up.
         assert "req-dead" not in orchestrator_fixture.orchestrator.request_states
     finally:
         if orchestrator_fixture.thread.is_alive():
@@ -132,6 +153,596 @@ async def test_engine_dead_error_broadcasts_fatal_and_shuts_down(orchestrator_fa
             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,6 +104,9 @@ 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
@@ -969,11 +972,17 @@ 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_clients=[SimpleNamespace(is_comprehension=False)],
+        stage_pools=[_FakeStagePool(clients=[SimpleNamespace()])],
     )
 
     assert omni.errored is False
@@ -983,21 +992,47 @@ def test_async_omni_errored_property_dead_engine():
     omni = object.__new__(AsyncOmni)
     omni.engine = SimpleNamespace(
         is_alive=lambda: False,
-        stage_clients=[SimpleNamespace(is_comprehension=False)],
+        stage_pools=[_FakeStagePool(clients=[SimpleNamespace()])],
     )
 
     assert omni.errored is True
 
 
-def test_async_omni_errored_property_dead_stage():
+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.
     omni = object.__new__(AsyncOmni)
-    dead_stage = SimpleNamespace(is_comprehension=False, _engine_dead=True)
     omni.engine = SimpleNamespace(
         is_alive=lambda: True,
-        stage_clients=[dead_stage],
+        stage_pools=[_FakeStagePool(clients=[SimpleNamespace(_engine_dead=True)])],
     )
 
-    assert omni.errored is True
+    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
 
 
 def _enqueue_stage_error(
@@ -1067,20 +1102,22 @@ 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
-    base.engine.stage_clients = [healthy_stage]
+    healthy_stage = SimpleNamespace(check_health=lambda: None)
+    base.engine.stage_pools = [_FakeStagePool([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
-    base.engine.stage_clients = [dead_stage]
-    with pytest.raises(EngineDeadError, match="Stage-1 dead"):
+
+    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.check_health()
 
 
@@ -1098,22 +1135,25 @@ 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_clients = [SimpleNamespace()]
+    base.engine.stage_pools = [_FakeStagePool([SimpleNamespace()])]
     assert base.errored is False
 
 
-def test_omni_base_is_running_false_when_stage_engine_dead():
+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.
     base = _make_base()
     base.engine.is_alive.return_value = True
-    base.engine.stage_clients = [SimpleNamespace(_engine_dead=True)]
-    assert base.is_running is False
+    base.engine.stage_pools = [_FakeStagePool([SimpleNamespace(_engine_dead=True)])]
+    assert base.is_running is True
 
 
-def test_omni_base_is_running_false_when_stage_resources_engine_dead():
+def test_omni_base_is_running_true_when_stage_resources_engine_dead():
     base = _make_base()
     base.engine.is_alive.return_value = True
-    base.engine.stage_clients = [SimpleNamespace(resources=SimpleNamespace(engine_dead=True))]
-    assert base.is_running is False
+    base.engine.stage_pools = [_FakeStagePool([SimpleNamespace(resources=SimpleNamespace(engine_dead=True))])]
+    assert base.is_running is True
 
 
 def test_omni_base_errored_true_when_orchestrator_dead():
@@ -1123,20 +1163,22 @@ def test_omni_base_errored_true_when_orchestrator_dead():
     assert base.errored is True
 
 
-def test_omni_base_errored_true_when_stage_engine_dead():
+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.
     base = _make_base()
     base.engine.is_alive.return_value = True
     dead_stage = SimpleNamespace(_engine_dead=True)
-    base.engine.stage_clients = [dead_stage]
-    assert base.errored is True
+    base.engine.stage_pools = [_FakeStagePool([dead_stage])]
+    assert base.errored is False
 
 
-def test_omni_base_errored_true_when_stage_resources_engine_dead():
+def test_omni_base_errored_false_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_clients = [dead_stage]
-    assert base.errored is True
+    base.engine.stage_pools = [_FakeStagePool([dead_stage])]
+    assert base.errored is False
 
 
 # ───────── Omni (sync) EngineDeadError / EngineGenerateError ─────────
@@ -1202,13 +1244,14 @@ def test_omni_errored_property_dead_engine(monkeypatch: pytest.MonkeyPatch):
 
 
 def test_omni_errored_property_dead_stage(monkeypatch: pytest.MonkeyPatch):
-    """Omni.errored returns True when a stage client is marked dead."""
+    """Omni.errored stays False when only a stage client dies (process-fatal
+    semantics): the request-level failure and check_health carry the signal."""
     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 True
+        assert app.errored is False
     finally:
         app.shutdown()

tests/entrypoints/test_serve.py

@@ -76,18 +76,6 @@ def test_serve_parser_accepts_four_way_cfg_parallelism() -> None:
     assert args.cfg_parallel_size == 4
 
 
-def test_serve_parser_defaults_vae_parallel_mode_to_auto_and_accepts_tile() -> None:
-    parser = TrackingArgumentParser()
-    subparsers = parser.add_subparsers(dest="subcommand")
-    OmniServeCommand().subparser_init(subparsers)
-
-    default_args = parser.parse_args(["serve", "fake-model", "--omni"])
-    tile_args = parser.parse_args(["serve", "fake-model", "--omni", "--vae-parallel-mode", "tile"])
-
-    assert default_args.vae_parallel_mode == "auto"
-    assert tile_args.vae_parallel_mode == "tile"
-
-
 def _make_headless_args(**kwargs) -> TrackingNamespace:
     defaults = {
         "model": "fake-model",

tests/model_executor/models/voxtral_tts/test_cuda_graph_acoustic_transformer.py

@@ -83,11 +83,16 @@ class SyntheticAcousticTransformer(nn.Module):
         # time_embedding: (B, 1) -> (B, dim)
         self.time_embedding = nn.Linear(1, HIDDEN_DIM, bias=False)
 
-        # _predict_velocity: takes x_t (B, C), llm_output (B, D), t_emb (B, D) -> (B, C)
+        # 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)
         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_output: torch.Tensor, t_emb: torch.Tensor) -> torch.Tensor:
-        combined = torch.cat([x_t, llm_output, t_emb], dim=-1)
+    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)
         return self._velocity_proj(combined)
 
 
@@ -129,14 +134,18 @@ 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)
-            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)
+            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)
             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

@@ -452,7 +452,7 @@ class OmniStageDiffusionParallelConfig(OmniStageParallelConfig):
     cfg_parallel_size: int = Field(default=1, ge=1)
     vae_patch_parallel_size: int = Field(default=1, ge=1)
     text_encoder_tp_size: int = Field(default=1, ge=1)
-    vae_parallel_mode: str = "auto"
+    vae_parallel_mode: str = "tile"
     use_hsdp: bool = False
     mask_sp_padding: bool = False
     hsdp_shard_size: int = -1
@@ -466,10 +466,9 @@ class OmniStageDiffusionParallelConfig(OmniStageParallelConfig):
             raise ValueError("allgather_degree > 1 is mutually exclusive with ulysses_degree/ring_degree > 1")
         if self.ulysses_mode not in {"strict", "advanced_uaa"}:
             raise ValueError("ulysses_mode must be 'strict' or 'advanced_uaa'")
-        if self.vae_parallel_mode not in {"auto", "tile", "spatial_shard_height", "spatial_shard_width"}:
+        if self.vae_parallel_mode not in {"tile", "spatial_shard_height", "spatial_shard_width"}:
             raise ValueError(
-                "vae_parallel_mode must be one of "
-                "{'auto', 'tile', 'spatial_shard_height', 'spatial_shard_width'}, "
+                "vae_parallel_mode must be one of {'tile', 'spatial_shard_height', 'spatial_shard_width'}, "
                 f"but got {self.vae_parallel_mode!r}."
             )
 

vllm_omni/diffusion/data.py

@@ -206,13 +206,10 @@ class DiffusionParallelConfig:
     process group.
     """
 
-    vae_parallel_mode: str = "auto"
+    vae_parallel_mode: str = "tile"
     """VAE parallel decode strategy.
 
-    - "auto": Select the longer spatial-shard axis for validated Wan and
-      distributed QwenImage tiled decodes on a full DiT-group topology;
-      otherwise use patch/tile parallel decode.
-    - "tile": Patch/tile parallel decode. Each rank decodes a subset
+    - "tile": Patch/tile parallel decode (default). Each rank decodes a subset
       of spatial tiles and the results are stitched on rank 0.
     - "spatial_shard_height": Spatially-sharded decode that splits decoder
       feature maps along height and exchanges halo rows around spatial
@@ -253,9 +250,8 @@ class DiffusionParallelConfig:
         assert self.allgather_degree > 0, "AllGather degree must be > 0"
         assert self.cfg_parallel_size > 0, "CFG parallel size must be > 0"
         assert self.vae_patch_parallel_size > 0, "VAE patch parallel size must be > 0"
-        assert self.vae_parallel_mode in {"auto", "tile", "spatial_shard_height", "spatial_shard_width"}, (
-            "vae_parallel_mode must be one of "
-            "{'auto', 'tile', 'spatial_shard_height', 'spatial_shard_width'}, "
+        assert self.vae_parallel_mode in {"tile", "spatial_shard_height", "spatial_shard_width"}, (
+            "vae_parallel_mode must be one of {'tile', 'spatial_shard_height', 'spatial_shard_width'}, "
             f"but got {self.vae_parallel_mode!r}."
         )
         if self.allgather_degree > 1:

vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage.py

@@ -4,23 +4,16 @@
 from typing import Any
 
 import torch
-import torch.distributed as dist
 from diffusers.models.autoencoders import AutoencoderKLQwenImage
 from diffusers.models.autoencoders.vae import DecoderOutput
 from vllm.logger import init_logger
 
-from vllm_omni.diffusion.distributed.autoencoders import qwenimage_spatial_shard
 from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import (
     DistributedOperator,
     DistributedVaeMixin,
     GridSpec,
     TileTask,
 )
-from vllm_omni.diffusion.distributed.autoencoders.spatial_shard import (
-    prepare_spatial_shard_decode,
-    resolve_spatial_shard_split_dim,
-    spatial_shard_decode_enabled,
-)
 
 logger = init_logger(__name__)
 
@@ -111,39 +104,10 @@ class DistributedAutoencoderKLQwenImage(AutoencoderKLQwenImage, DistributedVaeMi
         ]
         return dec
 
-    def _spatial_shard_decode_split_dim(
-        self,
-        z: torch.Tensor | None = None,
-        world_size: int | None = None,
-    ) -> str | None:
-        return resolve_spatial_shard_split_dim(self.distributed_executor.parallel_mode, z, world_size)
-
-    def _spatial_shard_decode_enabled(self, z: torch.Tensor) -> bool:
-        return spatial_shard_decode_enabled(self, z, model_name="QwenImage")
-
-    def _prepare_spatial_shard_decode(self) -> None:
-        prepare_spatial_shard_decode(
-            self,
-            install=qwenimage_spatial_shard.install_qwenimage_spatial_shard_decode,
-        )
-
     def tiled_decode(self, z: torch.Tensor, return_dict: bool = True):
         if not self.is_distributed_enabled():
             return super().tiled_decode(z, return_dict=return_dict)
 
-        if self._spatial_shard_decode_enabled(z):
-            world_size = dist.get_world_size(group=self.distributed_executor.group)
-            split_dim = self._spatial_shard_decode_split_dim(z, world_size)
-            assert split_dim is not None
-            logger.debug("Decode running with QwenImage VAE spatial_shard_%s mode", split_dim)
-            return qwenimage_spatial_shard.spatial_shard_decode(
-                self,
-                z,
-                group=self.distributed_executor.group,
-                return_dict=return_dict,
-                split_dim=split_dim,
-            )
-
         logger.debug("Decode running with distributed executor")
         result = self.distributed_executor.execute(
             z,

vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py

@@ -18,25 +18,11 @@ from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor impor
     GridSpec,
     TileTask,
 )
-from vllm_omni.diffusion.distributed.autoencoders.spatial_shard import (
-    prepare_pipeline_spatial_shard_decode,
-    prepare_spatial_shard_decode,
-    resolve_spatial_shard_split_dim,
-    spatial_shard_decode_enabled,
-)
-from vllm_omni.diffusion.distributed.autoencoders.spatial_shard import (
-    select_auto_spatial_shard_split_dim as select_auto_spatial_shard_split_dim,
-)
 from vllm_omni.platforms import current_omni_platform
 
 logger = init_logger(__name__)
 
 
-def prepare_pipeline_wan_spatial_shard_decode(pipeline: torch.nn.Module) -> None:
-    """Backward-compatible alias for the model-agnostic pipeline dispatcher."""
-    prepare_pipeline_spatial_shard_decode(pipeline)
-
-
 class OmniAutoencoderKLWan(AutoencoderKLWan):
     def _execution_context(self):
         try:
@@ -284,30 +270,45 @@ class DistributedAutoencoderKLWan(OmniAutoencoderKLWan, DistributedVaeMixin):
         dec = torch.clamp(dec, min=-1.0, max=1.0)
         return dec
 
-    def _spatial_shard_decode_split_dim(
-        self,
-        z: torch.Tensor | None = None,
-        world_size: int | None = None,
-    ) -> str | None:
-        return resolve_spatial_shard_split_dim(self.distributed_executor.parallel_mode, z, world_size)
+    def _spatial_shard_decode_split_dim(self) -> str | None:
+        mode = self.distributed_executor.parallel_mode
+        if mode == "spatial_shard_width":
+            return "width"
+        if mode == "spatial_shard_height":
+            return "height"
+        return None
 
     def _spatial_shard_decode_enabled(self, z: torch.Tensor) -> bool:
-        return spatial_shard_decode_enabled(self, z, model_name="Wan")
+        split_dim = self._spatial_shard_decode_split_dim()
+        if split_dim is None:
+            return False
+        if z.ndim != 5:
+            logger.warning("Wan VAE spatial sharded decode expects 5D latent input; falling back to tiled decode.")
+            return False
+        if not self.is_distributed_enabled():
+            return False
 
-    def _prepare_spatial_shard_decode(self) -> None:
-        prepare_spatial_shard_decode(
-            self,
-            install=wan_spatial_shard.install_wan_spatial_shard_decode,
-        )
+        group = self.distributed_executor.group
+        world_size = dist.get_world_size(group=group)
+        requested_size = int(self.distributed_executor.parallel_size)
+        if requested_size != world_size:
+            logger.warning(
+                "Wan VAE spatial sharded decode currently requires vae_patch_parallel_size "
+                "to match the DIT group size; falling back to tiled decode. "
+                "requested=%s dit_group=%s split_dim=%s",
+                requested_size,
+                world_size,
+                split_dim,
+            )
+            return False
+        return True
 
     def tiled_decode(self, z: torch.Tensor, return_dict: bool = True):
         if not self.is_distributed_enabled():
             return super().tiled_decode(z, return_dict=return_dict)
 
         if self._spatial_shard_decode_enabled(z):
-            world_size = dist.get_world_size(group=self.distributed_executor.group)
-            split_dim = self._spatial_shard_decode_split_dim(z, world_size)
-            assert split_dim is not None
+            split_dim = self._spatial_shard_decode_split_dim()
             logger.debug("Decode running with Wan VAE spatial_shard_%s mode", split_dim)
             return wan_spatial_shard.spatial_shard_decode(
                 self,

vllm_omni/diffusion/distributed/autoencoders/qwenimage_spatial_shard.py

@@ -1,86 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-"""QwenImage adapter for the shared spatial-shard VAE decoder kernels."""
-
-from __future__ import annotations
-
-import torch
-import torch.distributed as dist
-import torch.nn as nn
-from diffusers.models.autoencoders.autoencoder_kl_qwenimage import QwenImageCausalConv3d
-from diffusers.models.autoencoders.vae import DecoderOutput
-
-from vllm_omni.diffusion.distributed.autoencoders.spatial_shard import SpatialShardVAE
-from vllm_omni.diffusion.distributed.autoencoders.wan_spatial_shard import (
-    SpatialShardCausalConv3dMixin,
-    install_spatial_shard_decode,
-    spatial_shard_decode_impl,
-)
-
-
-class QwenImageDistCausalConv3d(SpatialShardCausalConv3dMixin, QwenImageCausalConv3d):
-    """Context-aware QwenImage causal conv that remains cache-discoverable."""
-
-    def __init__(self, source: nn.Conv3d, group: dist.ProcessGroup) -> None:
-        # Initialize Conv3d directly so the replacement preserves source device,
-        # dtype, groups, and bias while remaining an isinstance() match for
-        # QwenImageCausalConv3d. Diffusers clear_cache() relies on that match.
-        nn.Conv3d.__init__(
-            self,
-            source.in_channels,
-            source.out_channels,
-            source.kernel_size,
-            stride=source.stride,
-            padding=0,
-            dilation=source.dilation,
-            groups=source.groups,
-            bias=source.bias is not None,
-            padding_mode=source.padding_mode,
-            device=source.weight.device,
-            dtype=source.weight.dtype,
-        )
-        self._init_spatial_shard_causal_conv(source, group)
-
-    def forward(self, x: torch.Tensor, cache_x: torch.Tensor | None = None) -> torch.Tensor:
-        return self._spatial_shard_causal_conv_forward(x, cache_x)
-
-
-def install_qwenimage_spatial_shard_decode(
-    vae: SpatialShardVAE,
-    group: dist.ProcessGroup,
-    split_dim: str = "height",
-) -> None:
-    install_spatial_shard_decode(
-        vae,
-        group,
-        split_dim,
-        causal_conv_class_name="QwenImageCausalConv3d",
-        causal_conv_factory=QwenImageDistCausalConv3d,
-        attention_block_class_name="QwenImageAttentionBlock",
-        installed_attr="_vllm_omni_qwenimage_spatial_shard_installed",
-        model_name="QwenImage",
-    )
-
-
-def spatial_shard_decode(
-    vae: SpatialShardVAE,
-    z: torch.Tensor,
-    *,
-    group: dist.ProcessGroup,
-    return_dict: bool = True,
-    split_dim: str = "height",
-) -> DecoderOutput | tuple[torch.Tensor]:
-    """Decode QwenImage latents and preserve its all-rank output contract."""
-    return spatial_shard_decode_impl(
-        vae,
-        z,
-        group=group,
-        install=install_qwenimage_spatial_shard_decode,
-        model_name="QwenImage",
-        pass_first_chunk=False,
-        broadcast_result=True,
-        unpatchify_patch_size=None,
-        return_dict=return_dict,
-        split_dim=split_dim,
-    )

vllm_omni/diffusion/distributed/autoencoders/spatial_shard.py

@@ -1,142 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-from __future__ import annotations
-
-from collections.abc import Callable
-from typing import Protocol
-
-import torch
-import torch.distributed as dist
-from vllm.logger import init_logger
-
-logger = init_logger(__name__)
-
-
-class SpatialShardExecutor(Protocol):
-    group: dist.ProcessGroup
-    parallel_mode: str
-    parallel_size: int
-
-    def _sync_final_result(
-        self,
-        rank0_result: torch.Tensor,
-        output_ndim: int,
-        output_device: torch.device,
-        output_dtype: torch.dtype,
-    ) -> torch.Tensor: ...
-
-
-class SpatialShardVAE(Protocol):
-    distributed_executor: SpatialShardExecutor
-    decoder: torch.nn.Module
-    post_quant_conv: torch.nn.Module
-    dtype: torch.dtype
-    _conv_idx: list[int]
-    _feat_map: list[object]
-
-    def is_distributed_enabled(self) -> bool: ...
-
-    def clear_cache(self) -> None: ...
-
-
-def select_auto_spatial_shard_split_dim(
-    z_shape: tuple[int, ...],
-    world_size: int,
-) -> str | None:
-    """Select the longer spatial axis that can be split across every rank."""
-    if len(z_shape) != 5 or world_size <= 1:
-        return None
-
-    _, _, _, height, width = z_shape
-    candidates = [
-        (extent, split_dim) for extent, split_dim in ((height, "height"), (width, "width")) if extent >= world_size
-    ]
-    if not candidates:
-        return None
-    # Prefer width for square latents. It was modestly faster in the original
-    # four-GPU Wan benchmark and keeps the choice deterministic across models.
-    return max(candidates, key=lambda candidate: (candidate[0], candidate[1] == "width"))[1]
-
-
-def resolve_spatial_shard_split_dim(
-    mode: str,
-    z: torch.Tensor | None = None,
-    world_size: int | None = None,
-) -> str | None:
-    if mode == "spatial_shard_width":
-        return "width"
-    if mode == "spatial_shard_height":
-        return "height"
-    if mode == "auto" and z is not None and world_size is not None:
-        return select_auto_spatial_shard_split_dim(tuple(z.shape), world_size)
-    return None
-
-
-def spatial_shard_decode_enabled(vae: SpatialShardVAE, z: torch.Tensor, *, model_name: str) -> bool:
-    """Apply the shared request- and topology-level spatial decode policy."""
-    executor = vae.distributed_executor
-    if executor.parallel_mode == "tile":
-        return False
-    if z.ndim != 5:
-        logger.warning(
-            "%s VAE spatial sharded decode expects 5D latent input; falling back to tiled decode.", model_name
-        )
-        return False
-    if not vae.is_distributed_enabled():
-        return False
-
-    group = executor.group
-    world_size = dist.get_world_size(group=group)
-    requested_size = int(executor.parallel_size)
-    requested_split_dim = resolve_spatial_shard_split_dim(executor.parallel_mode)
-    if requested_size != world_size:
-        if executor.parallel_mode == "auto":
-            logger.debug(
-                "%s VAE auto decode selected tile mode for a partial DiT group (requested=%s dit_group=%s)",
-                model_name,
-                requested_size,
-                world_size,
-            )
-        elif requested_split_dim is not None:
-            logger.warning(
-                "%s VAE spatial sharded decode currently requires vae_patch_parallel_size "
-                "to match the DIT group size; falling back to tiled decode. "
-                "requested=%s dit_group=%s split_dim=%s",
-                model_name,
-                requested_size,
-                world_size,
-                requested_split_dim,
-            )
-        return False
-    return resolve_spatial_shard_split_dim(executor.parallel_mode, z, world_size) is not None
-
-
-def prepare_spatial_shard_decode(
-    vae: SpatialShardVAE,
-    *,
-    install: Callable[[SpatialShardVAE, dist.ProcessGroup, str], None],
-) -> None:
-    """Install one VAE's dynamic wrappers while allocations are still tagged."""
-    executor = vae.distributed_executor
-    mode = executor.parallel_mode
-    if mode not in {"auto", "spatial_shard_height", "spatial_shard_width"}:
-        return
-    requested_size = int(executor.parallel_size)
-    world_size = dist.get_world_size(group=executor.group)
-    if requested_size <= 1 or requested_size != world_size:
-        return
-    split_dim = "width" if mode == "spatial_shard_width" else "height"
-    install(vae, executor.group, split_dim)
-
-
-def prepare_pipeline_spatial_shard_decode(pipeline: torch.nn.Module) -> None:
-    """Prepare every VAE that advertises spatial-shard decode capability."""
-    from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery
-
-    if not dist.is_initialized():
-        return
-    for vae in ModuleDiscovery.discover(pipeline).vaes:
-        prepare = getattr(vae, "_prepare_spatial_shard_decode", None)
-        if callable(prepare):
-            prepare()

vllm_omni/diffusion/distributed/autoencoders/wan_spatial_shard.py

@@ -7,21 +7,19 @@
 #   (python/sglang/multimodal_gen/runtime/layers/parallel_conv.py)
 # which is in turn adapted from FastVideo (https://github.com/hao-ai-lab/FastVideo).
 # This version generalizes the height-only sharding to shard along height or
-# width and supports the causal-conv ``feat_cache`` handling used by Wan and
-# QwenImage.
-"""Shared spatial-shard kernels plus the Wan VAE adapter.
+# width and adds the Wan causal-conv ``feat_cache`` handling.
+"""Spatially-sharded Wan VAE decode.
 
-The halo, padding, stride-trim, attention-gather, and decode-lifecycle helpers
-are also reused by QwenImage's thin adapter. The backend shards decoder feature
-maps along height or width and exchanges boundary rows/columns before spatial
-convolutions. It is decode-only and keeps checkpoint loading unchanged by
-patching the already-loaded decoder.
+The existing distributed Wan VAE path shards *tiles*.  This module adds an
+opt-in decode backend that shards decoder feature maps along height or width and
+exchanges boundary rows/columns before spatial convolutions.  It is
+intentionally decode-only and keeps checkpoint loading unchanged by patching the
+already-loaded decoder.
 """
 
 from __future__ import annotations
 
 import math
-from collections.abc import Callable
 from contextlib import nullcontext
 from contextvars import ContextVar
 from dataclasses import dataclass
@@ -36,8 +34,6 @@ from diffusers.models.autoencoders.autoencoder_kl_wan import unpatchify
 from diffusers.models.autoencoders.vae import DecoderOutput
 from vllm.logger import init_logger
 
-from vllm_omni.diffusion.distributed.autoencoders.spatial_shard import SpatialShardVAE
-
 logger = init_logger(__name__)
 
 
@@ -56,17 +52,12 @@ _SPATIAL_SHARD_CONTEXT: ContextVar[SpatialShardContext | None] = ContextVar(
 )
 
 
-def _active_spatial_shard_split_dim() -> str | None:
-    context = _SPATIAL_SHARD_CONTEXT.get()
-    return context.split_dim if context is not None else None
-
-
 def _spatial_dim(split_dim: str) -> int:
     if split_dim == "height":
         return -2
     if split_dim == "width":
         return -1
-    raise ValueError(f"Unsupported VAE split_dim={split_dim!r}; expected 'height' or 'width'.")
+    raise ValueError(f"Unsupported Wan VAE split_dim={split_dim!r}; expected 'height' or 'width'.")
 
 
 def _narrow_along_dim(x: torch.Tensor, dim: int, start: int, length: int) -> torch.Tensor:
@@ -164,9 +155,9 @@ def split_for_parallel_decode(
     rank = 0 if rank is None else int(rank)
     world_size = 1 if world_size is None else int(world_size)
     if world_size < 1:
-        raise ValueError(f"VAE world_size must be >= 1, got {world_size}.")
+        raise ValueError(f"Wan VAE world_size must be >= 1, got {world_size}.")
     if not 0 <= rank < world_size:
-        raise ValueError(f"VAE rank must satisfy 0 <= rank < world_size, got rank={rank}, world_size={world_size}.")
+        raise ValueError(f"Wan VAE rank must satisfy 0 <= rank < world_size, got rank={rank}, world_size={world_size}.")
 
     dim = _spatial_dim(split_dim)
     expected_extent = x.shape[dim] * (2**upsample_count)
@@ -332,27 +323,21 @@ class WanDistZeroPad2d(nn.Module):
         padding: tuple[int, int, int, int],
         group: dist.ProcessGroup,
         *,
-        defer_split_after_padding: bool = False,
+        split_dim: str = "height",
+        split_padding: tuple[int, int] | None = None,
     ) -> None:
         super().__init__()
         self.padding = padding
-        self.defer_split_after_padding = defer_split_after_padding
+        self.split_dim = split_dim
+        default_split_padding = (padding[2], padding[3]) if split_dim == "height" else (padding[0], padding[1])
+        self.split_padding = split_padding or default_split_padding
         self.group = group
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        split_dim = _active_spatial_shard_split_dim()
-        if split_dim is None:
-            return F.pad(x, self.padding)
-
         rank, world_size = _rank_world(self.group)
         left, right, top, bottom = self.padding
-        if self.defer_split_after_padding:
-            if split_dim == "height":
-                bottom = 0
-            else:
-                right = 0
         if world_size > 1:
-            if split_dim == "height":
+            if self.split_dim == "height":
                 top = top if rank == 0 else 0
                 bottom = bottom if rank == world_size - 1 else 0
             else:
@@ -366,7 +351,8 @@ class WanDistConv2d(nn.Conv2d):
         self,
         source: nn.Conv2d,
         group: dist.ProcessGroup,
-        deferred_padding: tuple[int, int, int, int] | None = None,
+        split_dim: str = "height",
+        split_padding: tuple[int, int] | None = None,
     ):
         super().__init__(
             source.in_channels,
@@ -383,253 +369,176 @@ class WanDistConv2d(nn.Conv2d):
         )
         self.load_state_dict(source.state_dict())
         self.group = group
-        self._direct_padding = source.padding
-        self._direct_padding_mode = source.padding_mode
-        self._direct_reversed_padding = source._reversed_padding_repeated_twice
-        self._deferred_padding = deferred_padding
-        source_padding = source.padding if isinstance(source.padding, tuple) else (source.padding, source.padding)
-        if len(source_padding) != 2 or not all(isinstance(value, int) for value in source_padding):
-            raise ValueError(f"Spatial-shard Conv2d requires integer padding, got {source.padding!r}.")
-        self._spatial_pad_h = int(source_padding[-2])
-        self._spatial_pad_w = int(source_padding[-1])
-        self.register_buffer("_halo_recv_top_buf", None, persistent=False)
-        self.register_buffer("_halo_recv_bottom_buf", None, persistent=False)
-        self._trim_cache: dict[tuple[str, int], tuple[int, int, int]] = {}
+        self.split_dim = split_dim
+        self.split_tensor_dim = _spatial_dim(split_dim)
+        kernel_extent = self.kernel_size[-2] if split_dim == "height" else self.kernel_size[-1]
+        self.halo_size = (kernel_extent - 1) // 2
+        pad_h = source.padding[-2] if isinstance(source.padding, tuple) else source.padding
+        pad_w = source.padding[-1] if isinstance(source.padding, tuple) else source.padding
+        if split_dim == "height":
+            if split_padding is None:
+                split_padding = (pad_h, pad_h)
+            self._non_split_padding = (pad_w, pad_w, 0, 0)
+            self.kernel_extent = self.kernel_size[-2]
+            self.stride_extent = self.stride[-2]
+        else:
+            if split_padding is None:
+                split_padding = (pad_w, pad_w)
+            self._non_split_padding = (0, 0, pad_h, pad_h)
+            self.kernel_extent = self.kernel_size[-1]
+            self.stride_extent = self.stride[-1]
+        self.split_pad_left, self.split_pad_right = split_padding
+        self._halo_recv_top_buf: torch.Tensor | None = None
+        self._halo_recv_bottom_buf: torch.Tensor | None = None
+        self._trim_cache: dict[int, tuple[int, int, int]] = {}
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        split_dim = _active_spatial_shard_split_dim()
-        if split_dim is None:
-            if self._direct_padding_mode != "zeros":
-                x = F.pad(x, self._direct_reversed_padding, mode=self._direct_padding_mode)
-                padding: str | tuple[int, int] = (0, 0)
-            else:
-                padding = self._direct_padding
-            return F.conv2d(x, self.weight, self.bias, self.stride, padding, self.dilation, self.groups)
-
-        split_tensor_dim = _spatial_dim(split_dim)
-        if split_dim == "height":
-            kernel_extent = self.kernel_size[-2]
-            stride_extent = self.stride[-2]
-            non_split_padding = (self._spatial_pad_w, self._spatial_pad_w, 0, 0)
-            default_split_padding = (self._spatial_pad_h, self._spatial_pad_h)
-            deferred_split_padding = (
-                (self._deferred_padding[2], self._deferred_padding[3]) if self._deferred_padding is not None else None
-            )
-        else:
-            kernel_extent = self.kernel_size[-1]
-            stride_extent = self.stride[-1]
-            non_split_padding = (0, 0, self._spatial_pad_h, self._spatial_pad_h)
-            default_split_padding = (self._spatial_pad_w, self._spatial_pad_w)
-            deferred_split_padding = (
-                (self._deferred_padding[0], self._deferred_padding[1]) if self._deferred_padding is not None else None
-            )
-        split_pad_left, split_pad_right = deferred_split_padding or default_split_padding
-        halo_size = (kernel_extent - 1) // 2
-
-        x = F.pad(x, non_split_padding)
+        x = F.pad(x, self._non_split_padding)
         x_padded, self._halo_recv_top_buf, self._halo_recv_bottom_buf = halo_exchange(
             x,
             group=self.group,
-            halo_size=halo_size,
-            split_dim=split_dim,
+            halo_size=self.halo_size,
+            split_dim=self.split_dim,
             recv_top_buf=self._halo_recv_top_buf,
             recv_bottom_buf=self._halo_recv_bottom_buf,
         )
-        shift, start, upper_bound = self._get_trim_params(
-            x.shape[split_tensor_dim],
-            split_dim=split_dim,
-            halo_size=halo_size,
-            split_pad_left=split_pad_left,
-            split_pad_right=split_pad_right,
-            kernel_extent=kernel_extent,
-            stride_extent=stride_extent,
-        )
+        shift, start, upper_bound = self._get_trim_params(x.shape[self.split_tensor_dim])
         if shift:
             x_padded = _narrow_along_dim(
                 x_padded,
-                split_tensor_dim,
+                self.split_tensor_dim,
                 shift,
-                x_padded.shape[split_tensor_dim] - shift,
+                x_padded.shape[self.split_tensor_dim] - shift,
             )
         out = super().forward(x_padded)
-        out = _trim_local_conv_output(out, halo_size, start, upper_bound, split_dim=split_dim)
-        return _zero_invalid_extent(out, split_dim=split_dim)
+        out = _trim_local_conv_output(out, self.halo_size, start, upper_bound, split_dim=self.split_dim)
+        return _zero_invalid_extent(out, split_dim=self.split_dim)
 
-    def _get_trim_params(
-        self,
-        local_extent: int,
-        *,
-        split_dim: str,
-        halo_size: int,
-        split_pad_left: int,
-        split_pad_right: int,
-        kernel_extent: int,
-        stride_extent: int,
-    ) -> tuple[int, int, int]:
-        cache_key = (split_dim, local_extent)
-        trim_params = self._trim_cache.get(cache_key)
+    def _get_trim_params(self, local_extent: int) -> tuple[int, int, int]:
+        trim_params = self._trim_cache.get(local_extent)
         if trim_params is None:
             rank, world_size = _rank_world(self.group)
             trim_params = _compute_conv_trim_params(
                 local_extent=local_extent,
                 rank=rank,
                 world_size=world_size,
-                halo_size=halo_size,
-                pad_before=split_pad_left,
-                pad_after=split_pad_right,
-                kernel_extent=kernel_extent,
-                stride_extent=stride_extent,
+                halo_size=self.halo_size,
+                pad_before=self.split_pad_left,
+                pad_after=self.split_pad_right,
+                kernel_extent=self.kernel_extent,
+                stride_extent=self.stride_extent,
             )
-            self._trim_cache[cache_key] = trim_params
+            self._trim_cache[local_extent] = trim_params
         return trim_params
 
 
-# Model-neutral aliases let other causal VAEs reuse the implementation while
-# preserving the original Wan class identities for stacked PRs and downstreams.
-SpatialShardZeroPad2d = WanDistZeroPad2d
-SpatialShardConv2d = WanDistConv2d
-
-
-class SpatialShardCausalConv3dMixin:
-    """Shared causal-convolution behavior for supported Diffusers VAE classes."""
-
-    def _init_spatial_shard_causal_conv(
+class WanDistCausalConv3d(nn.Conv3d):
+    def __init__(
         self,
         source: nn.Conv3d,
         group: dist.ProcessGroup,
-    ) -> None:
+        split_dim: str = "height",
+    ):
+        super().__init__(
+            source.in_channels,
+            source.out_channels,
+            source.kernel_size,
+            stride=source.stride,
+            padding=0,
+            dilation=source.dilation,
+            groups=source.groups,
+            bias=source.bias is not None,
+            padding_mode=source.padding_mode,
+            device=source.weight.device,
+            dtype=source.weight.dtype,
+        )
         self.load_state_dict(source.state_dict())
         self.group = group
+        self.split_dim = split_dim
+        self.split_tensor_dim = _spatial_dim(split_dim)
         source_padding = getattr(source, "_padding", None)
         if source_padding is None:
             p_t, p_h, p_w = source.padding
             source_padding = (p_w, p_w, p_h, p_h, 2 * p_t, 0)
         self._source_padding = tuple(source_padding)
-        self._padding = tuple(source_padding)
-        self.register_buffer("_halo_recv_top_buf", None, persistent=False)
-        self.register_buffer("_halo_recv_bottom_buf", None, persistent=False)
-        self._trim_cache: dict[tuple[str, int], tuple[int, int, int]] = {}
+        if split_dim == "height":
+            self.split_pad_left = int(self._source_padding[2])
+            self.split_pad_right = int(self._source_padding[3])
+            self.kernel_extent = self.kernel_size[-2]
+            self.stride_extent = self.stride[-2]
+            padding = (
+                self._source_padding[0],
+                self._source_padding[1],
+                0,
+                0,
+                self._source_padding[4],
+                self._source_padding[5],
+            )
+        else:
+            self.split_pad_left = int(self._source_padding[0])
+            self.split_pad_right = int(self._source_padding[1])
+            self.kernel_extent = self.kernel_size[-1]
+            self.stride_extent = self.stride[-1]
+            padding = (
+                0,
+                0,
+                self._source_padding[2],
+                self._source_padding[3],
+                self._source_padding[4],
+                self._source_padding[5],
+            )
+        self.halo_size = (self.kernel_extent - 1) // 2
+        self._padding = padding if self.halo_size > 0 else self._source_padding
+        self._halo_recv_top_buf: torch.Tensor | None = None
+        self._halo_recv_bottom_buf: torch.Tensor | None = None
+        self._trim_cache: dict[int, tuple[int, int, int]] = {}
 
-    def _spatial_shard_causal_conv_forward(
-        self,
-        x: torch.Tensor,
-        cache_x: torch.Tensor | None = None,
-    ) -> torch.Tensor:
-        split_dim = _active_spatial_shard_split_dim()
-        padding = list(self._source_padding)
+    def forward(self, x: torch.Tensor, cache_x: torch.Tensor | None = None) -> torch.Tensor:
+        padding = list(self._padding)
         if cache_x is not None and padding[4] > 0:
             cache_x = cache_x.to(x.device)
             x = torch.cat([cache_x, x], dim=2)
             padding[4] -= cache_x.shape[2]
 
-        if split_dim is None:
-            return nn.Conv3d.forward(self, F.pad(x, padding))
-
-        split_tensor_dim = _spatial_dim(split_dim)
-        if split_dim == "height":
-            split_pad_left = int(padding[2])
-            split_pad_right = int(padding[3])
-            kernel_extent = self.kernel_size[-2]
-            stride_extent = self.stride[-2]
-            padding[2] = 0
-            padding[3] = 0
-        else:
-            split_pad_left = int(padding[0])
-            split_pad_right = int(padding[1])
-            kernel_extent = self.kernel_size[-1]
-            stride_extent = self.stride[-1]
-            padding[0] = 0
-            padding[1] = 0
-        halo_size = (kernel_extent - 1) // 2
-
         x = F.pad(x, padding)
         x_padded, self._halo_recv_top_buf, self._halo_recv_bottom_buf = halo_exchange(
             x,
             group=self.group,
-            halo_size=halo_size,
-            split_dim=split_dim,
+            halo_size=self.halo_size,
+            split_dim=self.split_dim,
             recv_top_buf=self._halo_recv_top_buf,
             recv_bottom_buf=self._halo_recv_bottom_buf,
         )
-        shift, start, upper_bound = self._get_trim_params(
-            x.shape[split_tensor_dim],
-            split_dim=split_dim,
-            halo_size=halo_size,
-            split_pad_left=split_pad_left,
-            split_pad_right=split_pad_right,
-            kernel_extent=kernel_extent,
-            stride_extent=stride_extent,
-        )
+        shift, start, upper_bound = self._get_trim_params(x.shape[self.split_tensor_dim])
         if shift:
             x_padded = _narrow_along_dim(
                 x_padded,
-                split_tensor_dim,
+                self.split_tensor_dim,
                 shift,
-                x_padded.shape[split_tensor_dim] - shift,
+                x_padded.shape[self.split_tensor_dim] - shift,
             )
-        out = nn.Conv3d.forward(self, x_padded)
-        out = _trim_local_conv_output(out, halo_size, start, upper_bound, split_dim=split_dim)
-        return _zero_invalid_extent(out, split_dim=split_dim)
+        out = super().forward(x_padded)
+        out = _trim_local_conv_output(out, self.halo_size, start, upper_bound, split_dim=self.split_dim)
+        return _zero_invalid_extent(out, split_dim=self.split_dim)
 
-    def _get_trim_params(
-        self,
-        local_extent: int,
-        *,
-        split_dim: str,
-        halo_size: int,
-        split_pad_left: int,
-        split_pad_right: int,
-        kernel_extent: int,
-        stride_extent: int,
-    ) -> tuple[int, int, int]:
-        cache_key = (split_dim, local_extent)
-        trim_params = self._trim_cache.get(cache_key)
+    def _get_trim_params(self, local_extent: int) -> tuple[int, int, int]:
+        trim_params = self._trim_cache.get(local_extent)
         if trim_params is None:
             rank, world_size = _rank_world(self.group)
             trim_params = _compute_conv_trim_params(
                 local_extent=local_extent,
                 rank=rank,
                 world_size=world_size,
-                halo_size=halo_size,
-                pad_before=split_pad_left,
-                pad_after=split_pad_right,
-                kernel_extent=kernel_extent,
-                stride_extent=stride_extent,
+                halo_size=self.halo_size,
+                pad_before=self.split_pad_left,
+                pad_after=self.split_pad_right,
+                kernel_extent=self.kernel_extent,
+                stride_extent=self.stride_extent,
             )
-            self._trim_cache[cache_key] = trim_params
+            self._trim_cache[local_extent] = trim_params
         return trim_params
 
 
-class WanDistCausalConv3d(SpatialShardCausalConv3dMixin, nn.Conv3d):
-    # Cross-PR compatibility contract for the lossless Wan data-movement
-    # installer: it may fuse the no-context direct path, but must leave this
-    # wrapper in control whenever a spatial-shard context is active.
-    _vllm_omni_dynamic_spatial_shard_conv = True
-
-    def __init__(
-        self,
-        source: nn.Conv3d,
-        group: dist.ProcessGroup,
-    ):
-        super().__init__(
-            source.in_channels,
-            source.out_channels,
-            source.kernel_size,
-            stride=source.stride,
-            padding=0,
-            dilation=source.dilation,
-            groups=source.groups,
-            bias=source.bias is not None,
-            padding_mode=source.padding_mode,
-            device=source.weight.device,
-            dtype=source.weight.dtype,
-        )
-        self._init_spatial_shard_causal_conv(source, group)
-
-    def forward(self, x: torch.Tensor, cache_x: torch.Tensor | None = None) -> torch.Tensor:
-        return self._spatial_shard_causal_conv_forward(x, cache_x)
-
-
 def _compute_conv_trim_params(
     *,
     local_extent: int,
@@ -673,15 +582,12 @@ def _trim_local_conv_output(
     return out
 
 
-def _patch_attention_block(module: nn.Module, group: dist.ProcessGroup) -> None:
+def _patch_attention_block(module: nn.Module, group: dist.ProcessGroup, split_dim: str) -> None:
     if getattr(module, "_vllm_omni_spatial_shard_attention", False):
         return
     orig_forward = module.forward
 
     def _forward(self: nn.Module, x: torch.Tensor, *args: Any, **kwargs: Any) -> torch.Tensor:
-        split_dim = _active_spatial_shard_split_dim()
-        if split_dim is None:
-            return orig_forward(x, *args, **kwargs)
         _, world_size = _rank_world(group)
         if world_size <= 1:
             return orig_forward(x, *args, **kwargs)
@@ -703,43 +609,60 @@ def _replace_child(
     name: str,
     child: nn.Module,
     group: dist.ProcessGroup,
-    *,
-    causal_conv_class_name: str,
-    causal_conv_factory: Callable[[nn.Conv3d, dist.ProcessGroup], nn.Module],
+    split_dim: str,
 ) -> None:
-    if child.__class__.__name__ == causal_conv_class_name:
-        setattr(parent, name, causal_conv_factory(child, group))
-        return
-    if isinstance(child, nn.ZeroPad2d):
-        padding = tuple(int(p) for p in child.padding)
+    if child.__class__.__name__ == "WanCausalConv3d":
         setattr(
             parent,
             name,
-            SpatialShardZeroPad2d(
-                padding,
+            WanDistCausalConv3d(
+                child,
                 group,
-                defer_split_after_padding=parent.__class__.__name__ == "Sequential",
+                split_dim=split_dim,
+            ),
+        )
+        return
+    if isinstance(child, nn.ZeroPad2d):
+        padding = tuple(int(p) for p in child.padding)
+        module_padding = padding
+        if parent.__class__.__name__ == "Sequential":
+            # Let the following WanDistConv2d account for global after-edge
+            # padding; this module handles the non-split dimension and before edge.
+            if split_dim == "height":
+                module_padding = (padding[0], padding[1], padding[2], 0)
+            else:
+                module_padding = (padding[0], 0, padding[2], padding[3])
+        setattr(
+            parent,
+            name,
+            WanDistZeroPad2d(
+                module_padding,
+                group,
+                split_dim=split_dim,
+                split_padding=(padding[2], padding[3]) if split_dim == "height" else (padding[0], padding[1]),
             ),
         )
         return
     if isinstance(child, nn.Conv2d):
-        deferred_padding = None
+        split_padding = None
         if name == "1" and parent.__class__.__name__ == "Sequential":
-            # Wan/QwenImage downsampling uses ZeroPad2d((0, 1, 0, 1)) before a
+            # WanResample downsample uses ZeroPad2d((0, 1, 0, 1)) before a
             # stride-2 conv with padding=0.  Only the last rank should see the
             # bottom/right global padding, which is approximated by split_padding.
             prev = getattr(parent, "0", None)
-            if isinstance(prev, SpatialShardZeroPad2d):
-                deferred_padding = prev.padding
+            if isinstance(prev, WanDistZeroPad2d):
+                split_padding = prev.split_padding
             elif isinstance(prev, nn.ZeroPad2d):
-                deferred_padding = tuple(int(value) for value in prev.padding)
+                pad = prev.padding
+                split_padding = (int(pad[2]), int(pad[3])) if split_dim == "height" else (int(pad[0]), int(pad[1]))
         setattr(
             parent,
             name,
-            SpatialShardConv2d(
+            WanDistConv2d(
                 child,
                 group,
-                deferred_padding=deferred_padding,
+                split_dim=split_dim,
+                split_padding=split_padding,
             ),
         )
 
@@ -747,45 +670,21 @@ def _replace_child(
 def _patch_decoder_modules(
     module: nn.Module,
     group: dist.ProcessGroup,
-    *,
-    causal_conv_class_name: str,
-    causal_conv_factory: Callable[[nn.Conv3d, dist.ProcessGroup], nn.Module],
-    attention_block_class_name: str,
+    split_dim: str,
     inside_attention: bool = False,
 ) -> None:
-    if module.__class__.__name__ == attention_block_class_name:
-        _patch_attention_block(module, group)
+    if module.__class__.__name__ == "WanAttentionBlock":
+        _patch_attention_block(module, group, split_dim)
         inside_attention = True
 
     for name, child in list(module.named_children()):
-        if child.__class__.__name__ == causal_conv_class_name:
-            _replace_child(
-                module,
-                name,
-                child,
-                group,
-                causal_conv_class_name=causal_conv_class_name,
-                causal_conv_factory=causal_conv_factory,
-            )
+        if child.__class__.__name__ == "WanCausalConv3d":
+            _replace_child(module, name, child, group, split_dim)
             continue
         if isinstance(child, nn.Conv2d) and not inside_attention:
-            _replace_child(
-                module,
-                name,
-                child,
-                group,
-                causal_conv_class_name=causal_conv_class_name,
-                causal_conv_factory=causal_conv_factory,
-            )
+            _replace_child(module, name, child, group, split_dim)
             continue
-        _patch_decoder_modules(
-            child,
-            group,
-            causal_conv_class_name=causal_conv_class_name,
-            causal_conv_factory=causal_conv_factory,
-            attention_block_class_name=attention_block_class_name,
-            inside_attention=inside_attention,
-        )
+        _patch_decoder_modules(child, group, split_dim, inside_attention=inside_attention)
 
 
 def _decoder_upsample_count(decoder: nn.Module) -> int:
@@ -796,58 +695,47 @@ def _decoder_upsample_count(decoder: nn.Module) -> int:
     return count
 
 
-def clear_spatial_shard_runtime_buffers(vae: SpatialShardVAE) -> None:
-    decoder = getattr(vae, "decoder", None)
-    if decoder is None:
-        return
-    for module in decoder.modules():
-        if "_halo_recv_top_buf" in module._buffers:
-            module._halo_recv_top_buf = None
-        if "_halo_recv_bottom_buf" in module._buffers:
-            module._halo_recv_bottom_buf = None
-
-
-def install_spatial_shard_decode(
-    vae: SpatialShardVAE,
-    group: dist.ProcessGroup,
-    split_dim: str,
-    *,
-    causal_conv_class_name: str,
-    causal_conv_factory: Callable[[nn.Conv3d, dist.ProcessGroup], nn.Module],
-    attention_block_class_name: str,
-    installed_attr: str,
-    model_name: str,
-) -> None:
+def install_wan_spatial_shard_decode(vae: Any, group: dist.ProcessGroup, split_dim: str = "height") -> None:
     """Patch ``vae.decoder`` once for spatially-sharded decode.
 
     This mutates the already-loaded decoder in place by swapping its spatial
-    convolutions/padding for context-aware variants and attaching a dedicated
-    spatial forward. Outside that forward the replacements execute their exact
-    local PyTorch behavior, so a later auto-selected tile request remains safe.
-    The same installed decoder can alternate height and width sharding.
+    convolutions/padding for halo-exchanging variants and wrapping
+    ``decoder.forward``. The patch is permanent for the lifetime of the VAE
+    instance and is applied only once (subsequent calls are no-ops). A given
+    instance is bound to a single ``split_dim``; switching between
+    ``"height"`` and ``"width"`` requires a fresh VAE instance and raises here
+    otherwise.
 
-    Group-relative rank 0 assembles the final decoded frame. The model adapter
-    then preserves its prior contract: Wan keeps the result on rank 0, while
-    QwenImage broadcasts it to every rank.
+    Only group-relative rank 0 assembles the final decoded frame, mirroring the
+    distributed tiled-decode ``broadcast_result=False`` contract; the other ranks
+    take part in the collectives but return an empty placeholder.
     """
     _spatial_dim(split_dim)
-    if getattr(vae, installed_attr, False):
+    if getattr(vae, "_vllm_omni_wan_spatial_shard_installed", False):
+        installed_split_dim = getattr(vae, "_vllm_omni_wan_spatial_shard_split_dim", "height")
+        if installed_split_dim != split_dim:
+            raise ValueError(
+                "Wan spatial-shard VAE decoder was already patched for "
+                f"{installed_split_dim!r} split; create a fresh VAE instance to use {split_dim!r} split."
+            )
         return
     decoder = getattr(vae, "decoder", None)
     if decoder is None:
-        raise ValueError(f"{model_name} spatial-shard VAE decode requires a decoder module.")
+        raise ValueError("Wan spatial-shard VAE decode requires a decoder module.")
 
-    _patch_decoder_modules(
-        decoder,
-        group,
-        causal_conv_class_name=causal_conv_class_name,
-        causal_conv_factory=causal_conv_factory,
-        attention_block_class_name=attention_block_class_name,
-    )
+    _patch_decoder_modules(decoder, group, split_dim)
     upsample_count = _decoder_upsample_count(decoder)
     orig_forward = decoder.forward
 
-    def _forward(self: nn.Module, x: torch.Tensor, *args: object, split_dim: str, **kwargs: object) -> torch.Tensor:
+    def _forward(
+        self: nn.Module,
+        x: torch.Tensor,
+        feat_cache: list[torch.Tensor] | None = None,
+        feat_idx: list[int] | None = None,
+        first_chunk: bool = False,
+    ) -> torch.Tensor:
+        if feat_idx is None:
+            feat_idx = [0]
         tensor_dim = _spatial_dim(split_dim)
         input_extent = x.shape[tensor_dim]
         x, expected_extent = split_for_parallel_decode(
@@ -867,50 +755,29 @@ def install_spatial_shard_decode(
             )
         )
         try:
-            out = orig_forward(x, *args, **kwargs)
+            out = orig_forward(x, feat_cache=feat_cache, feat_idx=feat_idx, first_chunk=first_chunk)
         finally:
             _SPATIAL_SHARD_CONTEXT.reset(token)
         return gather_and_trim_extent(out, expected_extent=expected_extent, split_dim=split_dim, group=group, dst=0)
 
-    decoder._vllm_omni_spatial_shard_forward = MethodType(_forward, decoder)
-    setattr(vae, installed_attr, True)
-    logger.info("Installed %s VAE dynamic-axis spatial-shard decode.", model_name)
+    decoder.forward = MethodType(_forward, decoder)
+    vae._vllm_omni_wan_spatial_shard_installed = True
+    vae._vllm_omni_wan_spatial_shard_split_dim = split_dim
+    logger.info("Installed Wan VAE %s-sharded decode.", split_dim)
 
 
-def install_wan_spatial_shard_decode(
-    vae: SpatialShardVAE,
-    group: dist.ProcessGroup,
-    split_dim: str = "height",
-) -> None:
-    install_spatial_shard_decode(
-        vae,
-        group,
-        split_dim,
-        causal_conv_class_name="WanCausalConv3d",
-        causal_conv_factory=WanDistCausalConv3d,
-        attention_block_class_name="WanAttentionBlock",
-        installed_attr="_vllm_omni_wan_spatial_shard_installed",
-        model_name="Wan",
-    )
-
-
-def spatial_shard_decode_impl(
-    vae: SpatialShardVAE,
+def spatial_shard_decode(
+    vae: Any,
     z: torch.Tensor,
     *,
     group: dist.ProcessGroup,
-    install: Callable[[SpatialShardVAE, dist.ProcessGroup, str], None],
-    model_name: str,
-    pass_first_chunk: bool,
-    broadcast_result: bool,
-    unpatchify_patch_size: int | None,
     return_dict: bool = True,
     split_dim: str = "height",
 ) -> DecoderOutput | tuple[torch.Tensor]:
-    install(vae, group, split_dim)
+    install_wan_spatial_shard_decode(vae, group, split_dim=split_dim)
 
     if z.shape[2] == 0:
-        raise ValueError(f"{model_name} spatial-shard VAE decode expects at least one latent frame.")
+        raise ValueError("Wan spatial-shard VAE decode expects at least one latent frame.")
 
     # Non-rank-0 ranks must still run the decoder every chunk to stay in lockstep with
     # the halo/all-gather collectives; they just skip keeping/assembling the output.
@@ -919,72 +786,31 @@ def spatial_shard_decode_impl(
 
     vae.clear_cache()
     try:
-        context_factory = getattr(vae, "_execution_context", None)
-        context = context_factory() if callable(context_factory) else nullcontext()
+        context = vae._execution_context() if hasattr(vae, "_execution_context") else nullcontext()
         with context:
             x = vae.post_quant_conv(z)
             decoded_chunks = []
-            spatial_forward = getattr(vae.decoder, "_vllm_omni_spatial_shard_forward")
             for i in range(z.shape[2]):
                 vae._conv_idx = [0]
-                if pass_first_chunk:
-                    chunk = spatial_forward(
-                        x[:, :, i : i + 1, :, :],
-                        feat_cache=vae._feat_map,
-                        feat_idx=vae._conv_idx,
-                        first_chunk=i == 0,
-                        split_dim=split_dim,
-                    )
-                else:
-                    chunk = spatial_forward(
-                        x[:, :, i : i + 1, :, :],
-                        feat_cache=vae._feat_map,
-                        feat_idx=vae._conv_idx,
-                        split_dim=split_dim,
-                    )
+                chunk = vae.decoder(
+                    x[:, :, i : i + 1, :, :],
+                    feat_cache=vae._feat_map,
+                    feat_idx=vae._conv_idx,
+                    first_chunk=(i == 0),
+                )
                 if produce_output:
                     decoded_chunks.append(chunk)
 
             if produce_output:
                 out = torch.cat(decoded_chunks, dim=2)
-                if unpatchify_patch_size is not None:
-                    out = unpatchify(out, patch_size=unpatchify_patch_size)
+                if vae.config.patch_size is not None:
+                    out = unpatchify(out, patch_size=vae.config.patch_size)
                 out = torch.clamp(out, min=-1.0, max=1.0)
             else:
                 out = z.new_zeros(0)
     finally:
         vae.clear_cache()
-        # Halo buffers are request-shape scratch space allocated outside the
-        # tagged weight pool. Drop live references so sleep/offload never has
-        # to preserve or discard them as model state.
-        clear_spatial_shard_runtime_buffers(vae)
-
-    if broadcast_result:
-        out = vae.distributed_executor._sync_final_result(out, z.ndim, z.device, vae.dtype)
 
     if not return_dict:
         return (out,)
     return DecoderOutput(sample=out)
-
-
-def spatial_shard_decode(
-    vae: SpatialShardVAE,
-    z: torch.Tensor,
-    *,
-    group: dist.ProcessGroup,
-    return_dict: bool = True,
-    split_dim: str = "height",
-) -> DecoderOutput | tuple[torch.Tensor]:
-    """Decode a Wan latent with spatial activation sharding."""
-    return spatial_shard_decode_impl(
-        vae,
-        z,
-        group=group,
-        install=install_wan_spatial_shard_decode,
-        model_name="Wan",
-        pass_first_chunk=True,
-        broadcast_result=False,
-        unpatchify_patch_size=vae.config.patch_size,
-        return_dict=return_dict,
-        split_dim=split_dim,
-    )

vllm_omni/diffusion/models/minimax_h3/vae.py

@@ -155,11 +155,6 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin):
         parallel_size: int,
         mode: str = "tile",
     ) -> None:
-        # ``auto`` is the cross-model default. MiniMax H3 has no spatial-shard
-        # capability, so preserve its native tile path for that default while
-        # continuing to reject explicitly requested spatial modes.
-        if mode == "auto":
-            mode = "tile"
         if mode != "tile":
             raise ValueError(f"MiniMax H3 VAE supports its native tile parallel mode only, got {mode!r}")
         group = get_dit_group()

vllm_omni/diffusion/worker/diffusion_model_runner.py

@@ -30,9 +30,6 @@ from vllm_omni.diffusion.cache.prompt_embed_cache import (
 from vllm_omni.diffusion.cache.selector import get_cache_backend
 from vllm_omni.diffusion.compile import regionally_compile
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
-from vllm_omni.diffusion.distributed.autoencoders.spatial_shard import (
-    prepare_pipeline_spatial_shard_decode,
-)
 from vllm_omni.diffusion.forward_context import set_forward_context
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
 from vllm_omni.diffusion.models.interface import (
@@ -246,7 +243,6 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
                     custom_pipeline_name=custom_pipeline_name,
                     device=self.device,
                 )
-                prepare_pipeline_spatial_shard_decode(self.pipeline)
         time_after_load = time.perf_counter()
 
         logger.info(

vllm_omni/engine/arg_utils.py

@@ -505,7 +505,7 @@ class OrchestratorArgs:
     diffusion_kv_cache_skip_layers: str | None = None
     cfg_parallel_size: int = 1
     vae_patch_parallel_size: int = 1
-    vae_parallel_mode: str = "auto"
+    vae_parallel_mode: str = "tile"
     text_encoder_tp_size: int = 1
     default_sampling_params: str | None = None
     max_generated_image_size: int | None = None

vllm_omni/engine/async_omni_engine.py

@@ -972,7 +972,7 @@ class AsyncOmniEngine:
             cfg_parallel_size = normalized_kwargs.get("cfg_parallel_size") or 1
             pipeline_parallel_size = normalized_kwargs.get("pipeline_parallel_size") or 1
             vae_patch_parallel_size = normalized_kwargs.get("vae_patch_parallel_size") or 1
-            vae_parallel_mode = normalized_kwargs.get("vae_parallel_mode") or "auto"
+            vae_parallel_mode = normalized_kwargs.get("vae_parallel_mode") or "tile"
             text_encoder_tp_size = normalized_kwargs.get("text_encoder_tp_size") or 1
             enable_expert_parallel = normalized_kwargs.get("enable_expert_parallel") or False
             use_hsdp = normalized_kwargs.get("use_hsdp", 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
+from vllm_omni.engine.stage_pool import StagePool, StageUnavailableError
 from vllm_omni.metrics.prometheus import OmniRequestCounter
 from vllm_omni.metrics.stat_logger import OmniPrometheusStatLogger
 from vllm_omni.outputs import OmniRequestOutput
@@ -417,8 +417,6 @@ 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
@@ -532,19 +530,17 @@ class Orchestrator:
         except asyncio.CancelledError:
             raise
         except EngineDeadError as e:
-            # 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"
+            # 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.
             await self.rpc_async_queue.put(
                 ErrorMessage(
-                    error=self._fatal_error or str(e),
+                    error=str(e) or "Stage engine died",
                     fatal=True,
-                    stage_id=self._fatal_error_stage_id,
                 )
             )
         except Exception:
@@ -562,9 +558,6 @@ 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()
@@ -651,6 +644,13 @@ 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,12 +682,18 @@ class Orchestrator:
         preprocess_ms = msg.preprocess_ms
         if preprocess_ms > 0:
             req_state.pipeline_timings["preprocess_ms"] = preprocess_ms
-        await self.stage_pools[stage_id].submit_initial(
-            request_id,
-            req_state,
-            prompt,
-            prompt_text=msg.output_prompt_text,
-        )
+        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
 
         if self.async_chunk and stage_id == 0 and final_stage_id > 0:
             await self._prewarm_async_chunk_stages(request_id, prompt, req_state)
@@ -717,12 +723,18 @@ class Orchestrator:
 
         req_state.streaming.enabled = True
         req_state.stage_submit_ts[stage_id] = _time.time()
-        await self.stage_pools[stage_id].submit_update(
-            request_id,
-            req_state,
-            request,
-            prompt_text=msg.output_prompt_text,
-        )
+        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
 
         if self.async_chunk and stage_id == 0 and final_stage_id > 0:
             await self._prewarm_async_chunk_stages(request_id, request, req_state)
@@ -757,20 +769,29 @@ class Orchestrator:
         )
         self.request_states[companion_id] = companion_state
         companion_state.stage_submit_ts[0] = _time.time()
-        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,
-        )
+        # 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
 
         logger.info(
             "[Orchestrator] CFG companion submitted: %s (role=%s, parent=%s, stage-0 replica-%s)",
             companion_id,
             role,
             parent_id,
-            companion_replica_id,
+            self.stage_pools[0].get_bound_replica_id(companion_id),
         )
 
     async def _handle_abort(self, msg: AbortRequestMessage) -> None:
@@ -903,16 +924,17 @@ class Orchestrator:
                     if self._shutdown_event.is_set():
                         return
 
-                    if pool.stage_type == "diffusion":
-                        output = pool.poll_diffusion_output(replica_id)
-                        if output is None:
-                            continue
+                    # 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
 
-                        pool.record_output_timestamps([output])
-                        await self._handle_processed_outputs(stage_id, replica_id, [output])
-                        idle = False
-                    else:
-                        try:
+                            pool.record_output_timestamps([diffusion_output])
+                            processed = [diffusion_output]
+                        else:
                             raw_outputs = await pool.poll_llm_raw_output(replica_id, timeout_s=0.001)
                             if raw_outputs is None:
                                 continue
@@ -944,7 +966,7 @@ class Orchestrator:
                             iteration_stats = (
                                 IterationStats() if (self._stat_logger is not None and raw_outputs.outputs) else None
                             )
-                            raw_output = await pool.process_llm_raw_outputs(
+                            processed = await pool.process_llm_raw_outputs(
                                 replica_id,
                                 raw_outputs,
                                 iteration_stats=iteration_stats,
@@ -957,80 +979,23 @@ class Orchestrator:
                                     iteration_stats,
                                     engine_idx=self._stage_replica_to_engine_idx[(stage_id, replica_id)],
                                 )
-                        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
+                    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
 
-                            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
+                    await self._handle_processed_outputs(stage_id, replica_id, processed)
+                    idle = False
 
             self._orch_monitor.note_loop(idle=idle)
             if idle:
@@ -1091,6 +1056,143 @@ 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(
@@ -1735,10 +1837,42 @@ 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]
@@ -2101,16 +2235,22 @@ class Orchestrator:
             _t_submit_start = _time.perf_counter()
 
             if next_pool.stage_type == "diffusion":
-                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,
-                        )
-                    },
+                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",
                 )
             else:
                 import copy
@@ -2140,16 +2280,30 @@ class Orchestrator:
                     resumable=downstream_resumable,
                 )
                 request.external_req_id = request.request_id
-                replica_id = await next_pool.submit_initial(
-                    request_id,
-                    req_state,
-                    request,
-                    prompt_text=None,
+                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",
                 )
+
+            # 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,
-                replica_id,
+                bound_replica_id if bound_replica_id is not None else 0,
                 req_state,
             )
 
@@ -2201,53 +2355,6 @@ 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,6 +44,19 @@ 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."""
@@ -329,7 +342,9 @@ class StagePool:
 
             now = _time.monotonic()
             if now >= deadline:
-                raise RuntimeError(f"no UP replica for stage {self.stage_id} after {self.DISPATCH_WAIT_TIMEOUT_S:.1f}s")
+                raise StageUnavailableError(
+                    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(
@@ -555,7 +570,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 RuntimeError(f"stage {self.stage_id} has no live replicas")
+                raise StageUnavailableError(f"stage {self.stage_id} has no live replicas")
             if len(live) == 1:
                 chosen = live[0]
             else:
@@ -570,13 +585,13 @@ class StagePool:
     def _llm_client(self, replica_id: int) -> StagePoolLLMClient:
         client = self.clients[replica_id]
         if client is None:
-            raise RuntimeError(f"stage {self.stage_id} replica {replica_id} is not attached")
+            raise StageUnavailableError(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 RuntimeError(f"stage {self.stage_id} replica {replica_id} is not attached")
+            raise StageUnavailableError(f"stage {self.stage_id} replica {replica_id} is not attached")
         return cast(StagePoolDiffusionClient, client)
 
     # ---- Metrics ----
@@ -961,7 +976,7 @@ class StagePool:
         )
         client = self.clients[replica_id]
         if client is None:
-            raise RuntimeError(f"stage {self.stage_id} replica {replica_id} is not attached")
+            raise StageUnavailableError(f"stage {self.stage_id} replica {replica_id} is not attached")
         try:
             self.output_processor.add_request(
                 request=request,
@@ -1010,7 +1025,7 @@ class StagePool:
 
         client = self.clients[replica_id]
         if client is None:
-            raise RuntimeError(f"stage {self.stage_id} replica {replica_id} is not attached")
+            raise StageUnavailableError(f"stage {self.stage_id} replica {replica_id} is not attached")
 
         if self.stage_type == "diffusion":
             if isinstance(request, list):
@@ -1246,3 +1261,21 @@ 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,17 +906,27 @@ class AsyncOmni(EngineClient, OmniBase):
                         await self.duplex_lifecycle_events.put(msg)
                         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
+                    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
 
                     should_continue, _, stage_id, req_state = self._handle_output_message(msg)
                     if should_continue:
@@ -1317,10 +1327,11 @@ class AsyncOmni(EngineClient, OmniBase):
 
     @property
     def errored(self) -> bool:
-        """Whether the engine is in a non-recoverable error state.
+        """Whether the engine is in a process-fatal error state.
 
-        Delegates to ``OmniBase.errored`` which checks the orchestrator
-        thread and all stage clients.  Redeclared here to satisfy the
+        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
         ``EngineClient`` abstract-property requirement (Python's ABC
         mechanism does not resolve abstract methods from sibling MRO
         entries).

vllm_omni/entrypoints/cli/serve.py

@@ -750,12 +750,10 @@ class OmniServeCommand(CLISubcommand):
         omni_config_group.add_argument(
             "--vae-parallel-mode",
             type=str,
-            default="auto",
-            choices=["auto", "tile", "spatial_shard_height", "spatial_shard_width"],
+            default="tile",
+            choices=["tile", "spatial_shard_height", "spatial_shard_width"],
             help="VAE parallel decode strategy for diffusion models. "
-            "'auto' (default) selects the longer spatial-shard axis for validated Wan and distributed QwenImage "
-            "tiled decodes "
-            "on a full DiT group, then falls back to patch/tile decode; 'tile' always uses patch/tile decode; "
+            "'tile' (default) uses patch/tile parallel decode; "
             "'spatial_shard_height'/'spatial_shard_width' use spatially-sharded decode that splits "
             "decoder feature maps along height/width and exchanges halo regions. The "
             "'spatial_shard_*' modes require vae_patch_parallel_size to match the DiT group size. "

vllm_omni/entrypoints/omni_base.py

@@ -4,7 +4,7 @@ import os
 import time
 import weakref
 from collections.abc import Sequence
-from typing import Any, Literal
+from typing import TYPE_CHECKING, Any, Literal
 
 import huggingface_hub
 import vllm.envs as envs
@@ -31,6 +31,9 @@ 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__)
 
 
@@ -237,6 +240,42 @@ 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:
@@ -244,38 +283,33 @@ 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() and not self._has_dead_stage()
+        return self.engine.is_alive()
 
     @property
     def errored(self) -> bool:
-        """Whether the engine is in a non-recoverable error state.
+        """Whether the engine is in a process-fatal error state.
 
-        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.
+        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).
         """
-        return not self.engine.is_alive() or self._has_dead_stage()
+        return not self.engine.is_alive()
 
     def check_health(self) -> None:
         if not self.engine.is_alive():
             raise EngineDeadError("Orchestrator process is not alive")
-        for stage_client in self.engine.stage_clients:
-            if hasattr(stage_client, "check_health"):
-                stage_client.check_health()
+        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")
 
     def resolve_sampling_params_list(
         self,

vllm_omni/model_executor/models/voxtral_tts/cuda_graph_acoustic_transformer_wrapper.py

@@ -74,6 +74,12 @@ 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)
 
@@ -142,21 +148,21 @@ class CUDAGraphAcousticTransformerWrapper:
 
         x = noise
 
-        # Pre-compute zero hidden states for unconditional CFG branch
-        hidden_states_zero = torch.zeros_like(hidden_states)
+        # 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)
 
         timesteps = self.timesteps
         for i in range(len(timesteps) - 1):
-            t = timesteps[i]
-            dt = timesteps[i + 1] - timesteps[i]
+            dt = self.dts[i]
 
-            # 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)
+            # 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)
             x_batched = torch.cat([x, x], dim=0)  # (2B, C)
-            llm_batched = torch.cat([hidden_states, hidden_states_zero], dim=0)  # (2B, D)
-            t_emb_batched = t_emb.repeat(2, 1)  # (2B, D)
+            t_proj_batched = t_proj.repeat(2, 1)  # (2B, D)
 
-            v_all = at._predict_velocity(x_t=x_batched, llm_output=llm_batched, t_emb=t_emb_batched)
+            v_all = at._predict_velocity(x_t=x_batched, llm_proj=llm_proj_batched, t_proj=t_proj_batched)
             v_t, uncond_v_t = v_all[:B], v_all[B:]
 
             # CFG combination (cfg_alpha is (B, 1), v_t is (B, C))
@@ -239,6 +245,10 @@ 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,10 +45,12 @@ 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(input_ids[offset].item())
-        context_length = int(input_ids[offset + 1].item())
+        ctx_frames = int(header_view[offset])
+        context_length = int(header_view[offset + 1])
         offset += 2
 
         req_input_ids_length = (ctx_frames + context_length) * num_codebooks
@@ -155,6 +157,9 @@ 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,6 +441,8 @@ 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())
@@ -514,12 +516,25 @@ 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
-        x_0 = torch.randn(B, self.model_args.n_acoustic_codebook).to(dtype=llm_hidden.dtype, device=llm_hidden.device)
+        # 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)
         x_0 = self._noise_scale * x_0
 
-        timesteps = self._timesteps.to(dtype=llm_hidden.dtype)
-        llm_hidden_zero = torch.zeros_like(llm_hidden)
+        # 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)
 
         # Reshape cfg_alpha for broadcasting: (B,) -> (B, 1)
         cfg_alpha = cfg_alpha.to(dtype=llm_hidden.dtype, device=llm_hidden.device)
@@ -528,20 +543,19 @@ class FlowMatchingAudioTransformer(nn.Module):
         # Euler integration with batched conditional + unconditional velocity
         sampled = x_0
         for i in range(len(timesteps) - 1):
-            t = timesteps[i]
-            dt = timesteps[i + 1] - timesteps[i]
+            dt = dts[i]  # precomputed constant step size
 
-            t_emb = self.time_embedding(t.view(-1, 1).repeat(B, 1)).to(llm_hidden.dtype)
+            # Reuse cached projected-t_emb row, expanded along batch dim.
+            t_proj = t_proj_table[i].unsqueeze(0).expand(B, -1)
 
             # Batch cond + uncond into a single forward pass (2B batch)
             x_batched = torch.cat([sampled, sampled], dim=0)
-            llm_batched = torch.cat([llm_hidden, llm_hidden_zero], dim=0)
-            t_emb_batched = torch.cat([t_emb, t_emb], dim=0)
+            t_proj_batched = torch.cat([t_proj, t_proj], dim=0)
 
             v_all = self._predict_velocity(
                 x_t=x_batched,
-                llm_output=llm_batched,
-                t_emb=t_emb_batched,
+                llm_proj=llm_proj_batched,
+                t_proj=t_proj_batched,
             )
             v_t, uncond_v_t = v_all[:B], v_all[B:]
             v_t = cfg_alpha * v_t + (1 - cfg_alpha) * uncond_v_t
@@ -558,18 +572,16 @@ class FlowMatchingAudioTransformer(nn.Module):
     def _predict_velocity(
         self,
         x_t: torch.Tensor,  # BxC
-        llm_output: torch.Tensor,  # BxD
-        t_emb: torch.Tensor,  # BxD
+        llm_proj: torch.Tensor,  # BxD, already through llm_projection
+        t_proj: torch.Tensor,  # BxD, already through time_projection
     ) -> torch.Tensor:
-        x_t = x_t.to(llm_output.dtype)
-
-        t_emb = self.time_projection(t_emb)
-        llm_output = self.llm_projection(llm_output)
+        # llm_proj/t_proj arrive pre-projected (hoisted out of the per-step loop by callers).
+        x_t = x_t.to(llm_proj.dtype)
 
         acoustic_and_semantic_embeddings = [
             self.input_projection(x_t.unsqueeze(1)),  # Bx1xD
-            t_emb.unsqueeze(1),
-            llm_output.unsqueeze(1),
+            t_proj.unsqueeze(1),
+            llm_proj.unsqueeze(1),
         ]
         acoustic_transformer_inputs = torch.concatenate(acoustic_and_semantic_embeddings, dim=1)
 
@@ -954,6 +966,8 @@ 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
@@ -1052,11 +1066,15 @@ class VoxtralTTSAudioGenerationForConditionalGeneration(nn.Module, SupportsMulti
             llm_hidden=hidden_states,
             cfg_alpha=cfg_alpha,
         )
-        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),
-        )
+        # 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])
         # 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,6 +1,7 @@
 # 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
@@ -9,6 +10,7 @@ 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 (
@@ -485,9 +487,10 @@ 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),
+            get_alibi_slopes(self.n_local_heads).to(torch.float32).contiguous(),
             persistent=False,
         )
 
@@ -552,8 +555,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)
-        alibi_slopes = self.alibi_slopes.to(dtype=xq.dtype, device=xq.device)
+        # 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)
         attn_bias = alibi_slopes.view(H, 1, 1) * rel_pos.unsqueeze(0).to(xq.dtype)
 
         # Causal mask: mask out j > i (future positions)
@@ -584,7 +587,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 = self.alibi_slopes.to(torch.float32)
+            # alibi_slopes already fp32 — avoid per-call cast.
             output = flash_attn_func(
                 xq,
                 xk,
@@ -594,7 +597,7 @@ class Attention(nn.Module):
                     self.sliding_window,
                     0 if self.args.causal else self.sliding_window,
                 ),
-                alibi_slopes=alibi_slopes,
+                alibi_slopes=self.alibi_slopes,
             )
         else:
             output = self._native_attention(xq, xk, xv)
@@ -1052,14 +1055,15 @@ class VoxtralTTSAudioTokenizer(nn.Module):
         """
         chunk_size = 375  # TODO(chenyos): Hardcode. Fix it
 
-        # 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)
+        # 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)]
 
         # Separate empty and non-empty
         results: list[torch.Tensor | None] = [None] * len(processed)
@@ -1087,25 +1091,19 @@ class VoxtralTTSAudioTokenizer(nn.Module):
             chunk_map.append((orig_idx, req_chunk_indices))
 
         # Pad chunks to max length and batch decode
-        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
+        padded = pad_sequence(all_chunks, batch_first=True, padding_value=0.0)
 
         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]
-        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))
+        # 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)
 
         # Trim padding and reassemble per request
         for orig_idx, chunk_indices in chunk_map: