fix: localize NumPy image conversion

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

.github/CODEOWNERS

@@ -1,6 +1,11 @@
 # CODEOWNERS is used to automatically assign PR reviewers
 # Format: path pattern @username or @team
 # More info: https://docs.github.com/en/repositories/managing-your-repositorys-settings-and-features/customizing-your-repository/about-code-owners
+#
+# The vllm_omni/** source-tree rules below are derived from the `owners` and
+# `primary_code_paths` front matter in docs/design/module/*.md. When updating
+# ownership, update the design doc first, then regenerate the matching rule here.
+# Last-match-wins: more specific rules must come AFTER broader ones.
 
 # Default owners (fallback reviewers for all files)
 * @Gaohan123 @hsliuustc0106 @ywang96
@@ -21,25 +26,10 @@
 # CI
 /.buildkite/ @yenuo26 @congw729 @NickCao
 
-# Engine and model_executor
+# ComfyUI / external apps
-/vllm_omni/core/ @tzhouam @yinpeiqi @fake0fan @Sy0307
+/apps/ @fhfuih @alex-jw-brooks
-/vllm_omni/engine/ @tzhouam @yinpeiqi @fake0fan @Sy0307
-/vllm_omni/entrypoints/ @tzhouam @yinpeiqi @fake0fan @alex-jw-brooks @NickCao
-/vllm_omni/model_executor/ @tzhouam @gcanlin
-/vllm_omni/worker/ @tzhouam @gcanlin
 
-# Diffusion
+# Benchmarks (out-of-tree)
-/vllm_omni/diffusion/ @Isotr0py @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @david6666666 @xuechendi
-/vllm_omni/lora/ @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @yuanheng-zhao
-
-# Diffusion CPU offloading
-/vllm_omni/diffusion/offloader/ @Isotr0py @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @yuanheng-zhao @xuechendi
-
-# Omni and TTS support
-/vllm_omni/transformers_utils/ @tzhouam @yuanheng-zhao
-/vllm_omni/model_executor/models/ @linyueqian @ZeldaHuang @princepride @yuanheng-zhao @amy-why-3459 @Bounty-hunter @Sy0307 @gcanlin @NickCao
-
-# Benchmarks
 /benchmarks/accuracy/ @david6666666
 /benchmarks/build_dataset/ @linyueqian @yuanheng-zhao
 /benchmarks/diffusion/ @Isotr0py @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @xuechendi
@@ -47,20 +37,112 @@
 /benchmarks/glm_image/ @princepride @JaredforReal
 /benchmarks/tts/ @linyueqian @ZeldaHuang @princepride @Sy0307 @gcanlin
 
-# Usability and Observability
+# =====================================================================
-/vllm_omni/deploy/ @tzhouam @lishunyang12
+# Module design ownership (docs/design/module/*.md)
-/vllm_omni/metrics/ @lishunyang12
+# Each section header cites the source design doc.
-/vllm_omni/profiler/ @gcanlin
+# =====================================================================
 
-# ComfyUI
+# Autoregressive runtime — docs/design/module/ar_runtime.md
-/apps/ @fhfuih @alex-jw-brooks
+/vllm_omni/core/ @tzhouam @yinpeiqi @fake0fan @Sy0307 @Gaohan123
+/vllm_omni/worker/ @tzhouam @yinpeiqi @fake0fan @Sy0307 @Gaohan123
 
-# Large Scale Serving
+# Engine orchestration & stage runtime — docs/design/module/engine_orchestration.md, stage_runtime.md
+/vllm_omni/engine/ @tzhouam @fake0fan
+# File-level overrides: input/output modality contracts own these engine files
+# — docs/design/module/input_output_modality_contracts.md
+/vllm_omni/engine/messages.py @Sy0307 @amy-why-3459 @Gaohan123 @alex-jw-brooks
+/vllm_omni/engine/serialization.py @Sy0307 @amy-why-3459 @Gaohan123 @alex-jw-brooks
+/vllm_omni/engine/mm_outputs.py @Sy0307 @amy-why-3459 @Gaohan123 @alex-jw-brooks
+/vllm_omni/engine/output_modality.py @Sy0307 @amy-why-3459 @Gaohan123 @alex-jw-brooks
+/vllm_omni/engine/output_processor.py @Sy0307 @amy-why-3459 @Gaohan123 @alex-jw-brooks
+
+# Input/output modality contracts — docs/design/module/input_output_modality_contracts.md
+/vllm_omni/inputs/ @Sy0307 @amy-why-3459 @Gaohan123 @alex-jw-brooks
+/vllm_omni/outputs/ @Sy0307 @amy-why-3459 @Gaohan123 @alex-jw-brooks
+/vllm_omni/request.py @Sy0307 @amy-why-3459 @Gaohan123 @alex-jw-brooks
+/vllm_omni/data_entry_keys.py @Sy0307 @amy-why-3459 @Gaohan123 @alex-jw-brooks
+
+# Entrypoints and serving boundaries — docs/design/module/entrypoints.md
+/vllm_omni/entrypoints/ @alex-jw-brooks @linyueqian @NickCao
+# File-level overrides within entrypoints/
+/vllm_omni/entrypoints/openai/errors.py @alex-jw-brooks @NickCao
+/vllm_omni/entrypoints/cli/benchmark/ @alex-jw-brooks @Bounty-hunter
+
+# Error contracts — docs/design/module/error_contracts.md
+/vllm_omni/errors.py @alex-jw-brooks @NickCao
+
+# OmniConnector (inter-stage transport) — docs/design/module/omni_connector.md
 /vllm_omni/distributed/ @princepride @yuanheng-zhao @xuechendi @natureofnature @fake0fan
 
-# Quantization and configuration
+# Model integration — docs/design/module/model_integration.md
+/vllm_omni/model_executor/ @tzhouam @gcanlin
+/vllm_omni/model_executor/models/ @tzhouam @gcanlin @Sy0307 @amy-why-3459 @linyueqian
+/vllm_omni/model_extras/ @tzhouam @gcanlin @princepride
+/vllm_omni/plugins/ @gcanlin @tjtanaa @xuechendi
+# Config override within model_executor — docs/design/module/vllm_omni_config.md
+/vllm_omni/model_executor/stage_configs/ @lishunyang12 @alex-jw-brooks
+
+# Execution platforms (hardware backends) — docs/design/module/execution_platforms.md
+/vllm_omni/platforms/ @gcanlin @tjtanaa @xuechendi
+/vllm_omni/attention/ @gcanlin @tjtanaa @xuechendi @david6666666 @lishunyang12
+
+# Configuration & deploy — docs/design/module/vllm_omni_config.md
 /vllm_omni/config/ @lishunyang12 @alex-jw-brooks
+/vllm_omni/deploy/ @lishunyang12 @alex-jw-brooks
+
+# Quantization — docs/design/module/quantization.md
 /vllm_omni/quantization/ @david6666666 @Isotr0py @lishunyang12
 
-# Hardware plugins and accelerators
+# Observability — docs/design/module/observability.md
-/vllm_omni/platforms/ @gcanlin @tjtanaa @xuechendi
+/vllm_omni/metrics/ @lishunyang12 @vraiti
+
+# Profiling — docs/design/module/profiling.md
+/vllm_omni/profiler/ @gcanlin
+
+# Benchmarking (in-tree) — docs/design/module/benchmarking.md
+/vllm_omni/benchmarks/ @alex-jw-brooks @Bounty-hunter
+
+# LoRA (no dedicated design doc; ownership intentionally narrowed)
+/vllm_omni/lora/ @SamitHuang @wtomin
+
+# Transformers utils (related path only in design docs; preserves prior ownership)
+/vllm_omni/transformers_utils/ @tzhouam @yuanheng-zhao
+
+# =====================================================================
+# Diffusion module family — docs/design/module/diffusion/
+# Umbrella rule first; sub-module rules after it win via last-match.
+# =====================================================================
+
+# Diffusion umbrella (top-level files: data.py, forward_context.py, etc.)
+# — docs/design/module/diffusion/index.md
+/vllm_omni/diffusion/ @Isotr0py @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @david6666666 @xuechendi
+
+# Cache management — docs/design/module/cache_management.md
+/vllm_omni/diffusion/cache/ @Isotr0py @princepride @SamitHuang
+
+# Diffusion quantization — docs/design/module/quantization.md
+/vllm_omni/diffusion/quantization/ @david6666666 @Isotr0py @lishunyang12
+
+# Diffusion profiling — docs/design/module/profiling.md
+/vllm_omni/diffusion/profiler/ @gcanlin
+
+# Diffusion runtime (scheduler, executor, worker)
+# — docs/design/module/diffusion/diffusion_runtime.md, continuous_batching.md
+/vllm_omni/diffusion/sched/ @Isotr0py @princepride @SamitHuang @fhfuih
+/vllm_omni/diffusion/executor/ @Isotr0py @princepride @SamitHuang @fhfuih
+/vllm_omni/diffusion/worker/ @Isotr0py @princepride @SamitHuang @fhfuih
+
+# Diffusion model integration (pipelines, registration, checkpoint loading)
+# — docs/design/module/diffusion/diffusion_model_integration.md
+/vllm_omni/diffusion/models/ @fhfuih @Bounty-hunter @wtomin
+/vllm_omni/diffusion/model_loader/ @fhfuih @Bounty-hunter @wtomin @RuixiangMa
+/vllm_omni/diffusion/lora/ @fhfuih @Bounty-hunter @wtomin @RuixiangMa
+
+# Diffusion parallelism (rank topology, TP/SP/PP/VAE patch)
+# — docs/design/module/diffusion/parallelism.md
+/vllm_omni/diffusion/distributed/ @princepride @yuanheng-zhao @xuechendi @Bounty-hunter
+/vllm_omni/diffusion/attention/parallel/ @princepride @yuanheng-zhao @xuechendi @Bounty-hunter
+
+# Diffusion offloader (CPU offloading, DLO)
+# — docs/design/module/diffusion/offloader.md
+/vllm_omni/diffusion/offloader/ @yuanheng-zhao @david6666666 @lishunyang12

.pre-commit-config.yaml

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

docs/.nav.yml

@@ -16,31 +16,21 @@ nav:
   - Examples:
     - examples/README.md
     - Offline Inference:
-      - GLM-Image Multistage End-to-End Inference: user_guide/examples/offline_inference/glm_image.md
-      - HunyuanImage-3.0 Image-to-Text Inference: user_guide/examples/offline_inference/hunyuan_image3.md
       - Image-To-Image: user_guide/examples/offline_inference/image_to_image.md
       - Image-To-Video: user_guide/examples/offline_inference/image_to_video.md
-      - InternVLA-A1: user_guide/examples/offline_inference/internvla_a1.md
+      - Speech-To-Video: user_guide/examples/offline_inference/speech_to_video.md
-      - MiMo-Audio: user_guide/examples/offline_inference/mimo_audio.md
-      - Qwen2.5-Omni: user_guide/examples/offline_inference/qwen2_5_omni.md
-      - Qwen3-Omni: user_guide/examples/offline_inference/qwen3_omni.md
       - Text-To-Audio: user_guide/examples/offline_inference/text_to_audio.md
       - Text-To-Image: user_guide/examples/offline_inference/text_to_image.md
       - Text-To-Speech: user_guide/examples/offline_inference/text_to_speech.md
       - Text-To-Video: user_guide/examples/offline_inference/text_to_video.md
+      - X-To-Text: user_guide/examples/offline_inference/x_to_text.md
       - X-To-Video-Audio: user_guide/examples/offline_inference/x_to_video_audio.md
     - Online Serving:
-      - BAGEL-7B-MoT: user_guide/examples/online_serving/bagel.md
-      - vLLM-Omni Helm Chart: user_guide/examples/online_serving/chart-helm.md
-      - Diffusers Backend Adapter: user_guide/examples/online_serving/diffusers_pipeline_adapter.md
-      - GLM-Image Online Serving: user_guide/examples/online_serving/glm_image.md
       - Image-To-Image: user_guide/examples/online_serving/image_to_image.md
       - Image-To-Video: user_guide/examples/online_serving/image_to_video.md
-      - MiMo-Audio: user_guide/examples/online_serving/mimo_audio.md
+      - Speech-To-Video: user_guide/examples/online_serving/speech_to_video.md
-      - Qwen2.5-Omni: user_guide/examples/online_serving/qwen2_5_omni.md
-      - Qwen3-Omni: user_guide/examples/online_serving/qwen3_omni.md
-      - Text-To-Speech: user_guide/examples/online_serving/text_to_speech.md
       - Text-To-Image: user_guide/examples/online_serving/text_to_image.md
+      - Text-To-Speech: user_guide/examples/online_serving/text_to_speech.md
       - Text-To-Video: user_guide/examples/online_serving/text_to_video.md
   - General:
     - usage/*

docs/design/module/ar_runtime.md

@@ -7,6 +7,7 @@ owners:
   - "@yinpeiqi"
   - "@fake0fan"
   - "@Sy0307"
+  - "@Gaohan123"
 primary_code_paths:
   - vllm_omni/core/**
   - vllm_omni/worker/**

docs/design/module/benchmarking.md

@@ -4,6 +4,7 @@ kind: module
 status: draft
 owners:
   - "@alex-jw-brooks"
+  - "@Bounty-hunter"
 primary_code_paths:
   - vllm_omni/benchmarks/**
   - vllm_omni/entrypoints/cli/benchmark/**

docs/design/module/cache_management.md

@@ -5,6 +5,7 @@ status: draft
 owners:
   - "@Isotr0py"
   - "@princepride"
+  - "@SamitHuang"
 primary_code_paths:
   - vllm_omni/diffusion/cache/**
 related_code_paths:

docs/design/module/diffusion/diffusion_model_integration.md

@@ -3,9 +3,10 @@ title: Diffusion Model Integration
 kind: module
 status: draft
 owners:
-  - "@Isotr0py"
+  - "@fhfuih"
+  - "@Bounty-hunter"
   - "@wtomin"
-  - "@ZJY0516"
+  - "@RuixiangMa"
 primary_code_paths:
   - vllm_omni/diffusion/models/**
   - vllm_omni/diffusion/model_loader/**

docs/design/module/diffusion/diffusion_runtime.md

@@ -6,6 +6,7 @@ owners:
   - "@Isotr0py"
   - "@princepride"
   - "@SamitHuang"
+  - "@fhfuih"
 primary_code_paths:
   - vllm_omni/diffusion/sched/**
   - vllm_omni/diffusion/executor/**

docs/design/module/diffusion/offloader.md

@@ -4,7 +4,8 @@ kind: module
 status: draft
 owners:
   - "@yuanheng-zhao"
-  - "@Isotr0py"
+  - "@david6666666"
+  - "@lishunyang12"
 primary_code_paths:
   - vllm_omni/diffusion/offloader/**
 related_code_paths:

docs/design/module/diffusion/parallelism.md

@@ -6,6 +6,7 @@ owners:
   - "@princepride"
   - "@yuanheng-zhao"
   - "@xuechendi"
+  - "@Bounty-hunter"
 primary_code_paths:
   - vllm_omni/diffusion/distributed/**
   - vllm_omni/diffusion/attention/parallel/**

docs/design/module/entrypoints.md

@@ -6,6 +6,7 @@ architecture_state: current-plus-in-flight
 owners:
   - "@alex-jw-brooks"
   - "@linyueqian"
+  - "@NickCao"
 document_stewards:
   - "@hsliuustc0106"
   - "@Gaohan123"

docs/design/module/execution_platforms.md

@@ -8,8 +8,15 @@ owners:
   - "@xuechendi"
 primary_code_paths:
   - vllm_omni/platforms/**
-related_code_paths:
   - vllm_omni/attention/**
+primary_code_path_owners:
+  # Additional owners scoped to attention kernels (quantized/custom attention)
+  - paths:
+      - vllm_omni/attention/**
+    owners:
+      - "@david6666666"
+      - "@lishunyang12"
+related_code_paths: []
 depends_on:
   - ar_runtime.md
   - diffusion/index.md

docs/design/module/input_output_modality_contracts.md

@@ -6,6 +6,8 @@ architecture_state: current
 owners:
   - "@Sy0307"
   - "@amy-why-3459"
+  - "@Gaohan123"
+  - "@alex-jw-brooks"
 document_stewards:
   - "@hsliuustc0106"
   - "@Gaohan123"

docs/design/module/model_integration.md

@@ -9,6 +9,14 @@ primary_code_paths:
   - vllm_omni/model_executor/**
   - vllm_omni/model_extras/**
   - vllm_omni/plugins/**
+primary_code_path_owners:
+  # Additional owners scoped to a sub-path of model_executor/**
+  - paths:
+      - vllm_omni/model_executor/models/**
+    owners:
+      - "@Sy0307"
+      - "@amy-why-3459"
+      - "@linyueqian"
 related_code_paths:
   - vllm_omni/transformers_utils/**
   - vllm_omni/tokenizers/**

docs/design/module/observability.md

@@ -4,6 +4,7 @@ kind: module
 status: draft
 owners:
   - "@lishunyang12"
+  - "@vraiti"
 primary_code_paths:
   - vllm_omni/metrics/**
 related_code_paths:

docs/mkdocs/hooks/generate_examples.py

@@ -24,6 +24,23 @@ SERVING_MODE_TITLES = {
     "online_serving": "Online serving",
 }
 
+# Keep the Examples navigation focused on reusable modality/task entry points.
+# Model-specific scripts remain available under ``examples/`` and are surfaced
+# from these shared pages instead of getting one page per model.
+GENERAL_EXAMPLE_SLUGS = frozenset(
+    {
+        "image_to_image",
+        "image_to_video",
+        "speech_to_video",
+        "text_to_audio",
+        "text_to_image",
+        "text_to_speech",
+        "text_to_video",
+        "x_to_text",
+        "x_to_video_audio",
+    }
+)
+
 
 def load_model_display_names() -> dict[str, str]:
     try:
@@ -46,6 +63,13 @@ def load_model_display_names() -> dict[str, str]:
 MODEL_DISPLAY_NAMES = load_model_display_names()
 
 
+def is_general_example(example: "Example") -> bool:
+    """Return whether an example should get a generated serving page."""
+    if example.category not in SERVING_MODE_TITLES:
+        return True
+    return example.path.stem in GENERAL_EXAMPLE_SLUGS
+
+
 def fix_case(text: str) -> str:
     subs = {
         "api": "API",
@@ -340,22 +364,21 @@ def update_nav_file(examples: list[Example]):
             examples_by_category[category] = []
         examples_by_category[category].append(example)
 
-    # Build the new Examples section - start with preserved items
+    def category_items(category: str) -> list[dict[str, str]]:
-    examples_section = preserved_items.copy()
+        items = []
-
+        for example in sorted(examples_by_category[category], key=lambda e: e.path.stem):
-    # Add examples grouped by category, sorted by category name
-    for category in sorted(examples_by_category.keys()):
-        category_examples = sorted(examples_by_category[category], key=lambda e: e.path.stem)
-        category_items = []
-        for example in category_examples:
             doc_path = EXAMPLE_DOC_DIR / example.category / f"{example.path.stem}.md"
             rel_path = doc_path.relative_to(ROOT_DIR / "docs")
-            category_items.append({example.nav_title: rel_path.as_posix()})
+            items.append({example.nav_title: rel_path.as_posix()})
+        return items
 
-        if category_items:
+    # Build the new Examples section - start with preserved items.
-            # Format category name (e.g., "offline_inference" -> "Offline Inference")
+    examples_section = preserved_items.copy()
+    for category in sorted(examples_by_category.keys()):
+        items = category_items(category)
+        if items:
             category_title = fix_case(category.replace("_", " ").title())
-            examples_section.append({category_title: category_items})
+            examples_section.append({category_title: items})
 
     # Update the nav structure
     nav_list[user_guide_idx]["User Guide"][examples_idx]["Examples"] = examples_section
@@ -390,6 +413,8 @@ def on_startup(command: Literal["build", "gh-deploy", "serve"], dirty: bool):
         for path in category.glob("*/*.md"):
             examples.append(Example(path.parent, category.stem))
 
+    examples = [example for example in examples if is_general_example(example)]
+
     # Generate the example documentation
     for example in sorted(examples, key=lambda e: e.path.stem):
         example_name = f"{example.path.stem}.md"

docs/models/supported_models.md

@@ -17,84 +17,93 @@ th {
 }
 </style>
 
-| Architecture | Models | Example HF Models | NVIDIA GPU | AMD GPU | Ascend NPU | Intel GPU |
+| Architecture | Models | Example HF Models | NVIDIA GPU | AMD GPU | Ascend NPU | Intel GPU | Recipe |
-|--------------|--------|-------------------|------------|---------|-----|-----------|
+|--------------|--------|-------------------|------------|---------|-----|-----------|--------|
-| `Qwen3OmniMoeForConditionalGeneration` | Qwen3-Omni | `Qwen/Qwen3-Omni-30B-A3B-Instruct` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
+| `Qwen3OmniMoeForConditionalGeneration` | Qwen3-Omni | `Qwen/Qwen3-Omni-30B-A3B-Instruct` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Qwen/Qwen3-Omni.md) |
-| `Qwen2_5OmniForConditionalGeneration` | Qwen2.5-Omni | `Qwen/Qwen2.5-Omni-7B`, `Qwen/Qwen2.5-Omni-3B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
+| `Qwen2_5OmniForConditionalGeneration` | Qwen2.5-Omni | `Qwen/Qwen2.5-Omni-7B`, `Qwen/Qwen2.5-Omni-3B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
-| `MingFlashOmniForConditionalGeneration` + `MingImagePipeline` | Ming-flash-omni-2.0 (omni-speech + imagegen<sup>1</sup>) | `Jonathan1909/Ming-flash-omni-2.0` | ✅︎ |   |   |   |
+| `MingFlashOmniForConditionalGeneration` + `MingImagePipeline` | Ming-flash-omni-2.0 (omni-speech + imagegen<sup>1</sup>) | `Jonathan1909/Ming-flash-omni-2.0` | ✅︎ |   |   |   | — |
-| `BagelForConditionalGeneration` | BAGEL (DiT-only) | `ByteDance-Seed/BAGEL-7B-MoT` | ✅︎ | ✅︎ | | ✅︎ |
+| `BagelForConditionalGeneration` | BAGEL (DiT-only) | `ByteDance-Seed/BAGEL-7B-MoT` | ✅︎ | ✅︎ | | ✅︎ | — |
-| `InternVLAA1Pipeline` | InternVLA-A1 | `InternRobotics/InternVLA-A1-3B` | ✅︎ | ✅︎ | | |
+| `InternVLAA1Pipeline` | InternVLA-A1 | `InternRobotics/InternVLA-A1-3B` | ✅︎ | ✅︎ | | | — |
-| `Gr00tN1d7Pipeline` | GR00T N1.7 | `nvidia/GR00T-N1.7-3B` | ✅︎ | | | |
+| `Gr00tN1d7Pipeline` | GR00T N1.7 | `nvidia/GR00T-N1.7-3B` | ✅︎ | | | | — |
-| `HunyuanImage3ForCausalMM` | HunyuanImage3.0 (DiT-only) | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
+| `HunyuanImage3ForCausalMM` | HunyuanImage3.0 (DiT-only) | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
-| `QwenImagePipeline` | Qwen-Image | `Qwen/Qwen-Image` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
+| `QwenImagePipeline` | Qwen-Image | `Qwen/Qwen-Image` | ✅︎ | | | | [Published](https://recipes.vllm.ai/Qwen/Qwen-Image) |
-| `QwenImagePipeline` | Qwen-Image-2512 | `Qwen/Qwen-Image-2512` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
+| `QwenImagePipeline` | Qwen-Image-2512 | `Qwen/Qwen-Image-2512` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Qwen/Qwen-Image-2512.md) |
-| `QwenImageEditPipeline` | Qwen-Image-Edit | `Qwen/Qwen-Image-Edit` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
+| `QwenImageEditPipeline` | Qwen-Image-Edit | `Qwen/Qwen-Image-Edit` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
-| `QwenImageEditPlusPipeline` | Qwen-Image-Edit-2509 | `Qwen/Qwen-Image-Edit-2509` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
+| `QwenImageEditPlusPipeline` | Qwen-Image-Edit-2509 | `Qwen/Qwen-Image-Edit-2509` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
-| `QwenImageLayeredPipeline` | Qwen-Image-Layered | `Qwen/Qwen-Image-Layered` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
+| `QwenImageLayeredPipeline` | Qwen-Image-Layered | `Qwen/Qwen-Image-Layered` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
-| `QwenImageEditPlusPipeline` | Qwen-Image-Edit-2511 | `Qwen/Qwen-Image-Edit-2511` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
+| `QwenImageEditPlusPipeline` | Qwen-Image-Edit-2511 | `Qwen/Qwen-Image-Edit-2511` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
-| `GlmImagePipeline` | GLM-Image | `zai-org/GLM-Image` | ✅︎ | ✅︎ | | |
+| `GlmImagePipeline` | GLM-Image | `zai-org/GLM-Image` | ✅︎ | ✅︎ | | | — |
-| `ZImagePipeline` | Z-Image | `Tongyi-MAI/Z-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
+| `ZImagePipeline` | Z-Image | `Tongyi-MAI/Z-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
-| `Krea2Pipeline` | Krea 2 (Raw + Turbo) | `krea/Krea-2-Raw`, `krea/Krea-2-Turbo` | ✅︎ | | | |
+| `Krea2Pipeline` | Krea 2 (Raw + Turbo) | `krea/Krea-2-Raw`, `krea/Krea-2-Turbo` | ✅︎ | | | | — |
-| `WanPipeline` | Wan2.1-T2V, Wan2.2-T2V, Wan2.2-TI2V | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers`, `Wan-AI/Wan2.1-T2V-14B-Diffusers`, `Wan-AI/Wan2.2-T2V-A14B-Diffusers`, `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
+| `WanPipeline` | Wan2.1-T2V, Wan2.2-T2V, Wan2.2-TI2V | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers`, `Wan-AI/Wan2.1-T2V-14B-Diffusers`, `Wan-AI/Wan2.2-T2V-A14B-Diffusers`, `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
-| `WanImageToVideoPipeline` | Wan2.2-I2V | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
+| `WanImageToVideoPipeline` | Wan2.2-I2V | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
-| `Cosmos3OmniDiffusersPipeline` | Cosmos3 T2I, T2V, I2V, V2V, T2V with sound, action policy | `nvidia/Cosmos3-Nano`, `nvidia/Cosmos3-Super` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
+| `Cosmos3OmniDiffusersPipeline` | Cosmos3 T2I, T2V, I2V, V2V, T2V with sound, action policy | `nvidia/Cosmos3-Nano`, `nvidia/Cosmos3-Super` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
-| `Wan22S2VPipeline` | Wan2.2-S2V | `Wan-AI/Wan2.2-S2V-14B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
+| `Wan22S2VPipeline` | Wan2.2-S2V | `Wan-AI/Wan2.2-S2V-14B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
-| `Wan22VACEPipeline` | Wan2.1-VACE | `Wan-AI/Wan2.1-VACE-1.3B-diffusers`, `Wan-AI/Wan2.1-VACE-14B-diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
+| `Wan22VACEPipeline` | Wan2.1-VACE | `Wan-AI/Wan2.1-VACE-1.3B-diffusers`, `Wan-AI/Wan2.1-VACE-14B-diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
-| `Wan22VACEPipeline` | Wan2.2-VACE | `Pyros13/Wan2.2-VACE-Fun-A14B-Diffusers` | ✅︎ |   |   |   |
+| `Wan22VACEPipeline` | Wan2.2-VACE | `Pyros13/Wan2.2-VACE-Fun-A14B-Diffusers` | ✅︎ |   |   |   | — |
-| `LTX2Pipeline` | LTX-2 / LTX-2.3 one-stage T2V and I2V | `Lightricks/LTX-2`, `diffusers/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | |
+| `LTX2Pipeline` | LTX-2 / LTX-2.3 one-stage T2V and I2V | `Lightricks/LTX-2`, `diffusers/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | | — |
-| `LTX2DistilledPipeline` | LTX-2 distilled two-stage T2V and I2V | `rootonchair/LTX-2-19b-distilled` | ✅︎ | ✅︎ | | |
+| `LTX2DistilledPipeline` | LTX-2 distilled two-stage T2V and I2V | `rootonchair/LTX-2-19b-distilled` | ✅︎ | ✅︎ | | | — |
-| `LingBotVideoPipeline` | LingBot-Video dense and MoE T2I, T2V, TI2V | `robbyant/lingbot-video-dense-1.3b`, `robbyant/lingbot-video-moe-30b-a3b` | ✅︎ | | | |
+| `LingBotVideoPipeline` | LingBot-Video dense and MoE T2I, T2V, TI2V | `robbyant/lingbot-video-dense-1.3b`, `robbyant/lingbot-video-moe-30b-a3b` | ✅︎ | | | | — |
-| `MiniMaxH3Pipeline` | MiniMax H3 T2VA, first/last-keyframe FL2VA, and mixed-reference Ref2VA | `MiniMaxAI/MiniMax-H3` | ✅︎ | | | |
+| `MiniMaxH3Pipeline` | MiniMax H3 T2VA, first/last-keyframe FL2VA, and mixed-reference Ref2VA | `MiniMaxAI/MiniMax-H3` | ✅︎ | ✅︎ | | | [Published](https://recipes.vllm.ai/MiniMaxAI/MiniMax-H3) |
-| `DreamZeroPipeline` | DreamZero-DROID | `GEAR-Dreams/DreamZero-DROID` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
+| `DreamZeroPipeline` | DreamZero-DROID | `GEAR-Dreams/DreamZero-DROID` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
-| `HeliosPipeline`, `HeliosPyramidPipeline` | Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` | ✅︎ | ✅︎ | ✅︎ | |
+| `HeliosPipeline`, `HeliosPyramidPipeline` | Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` | ✅︎ | ✅︎ | ✅︎ | | — |
-| `MagiHumanPipeline` | MagiHuman | `SII-GAIR/daVinci-MagiHuman-Base-1080p` | ✅︎ | ✅︎ | | |
+| `MagiHumanPipeline` | MagiHuman | `SII-GAIR/daVinci-MagiHuman-Base-1080p` | ✅︎ | ✅︎ | | | — |
-| `OvisImagePipeline` | Ovis-Image | `OvisAI/Ovis-Image` | ✅︎ | ✅︎ | | ✅︎ |
+| `OvisImagePipeline` | Ovis-Image | `OvisAI/Ovis-Image` | ✅︎ | ✅︎ | | ✅︎ | — |
-| `LongcatImagePipeline` | LongCat-Image | `meituan-longcat/LongCat-Image` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
+| `LongcatImagePipeline` | LongCat-Image | `meituan-longcat/LongCat-Image` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
-| `LongCatImageEditPipeline` | LongCat-Image-Edit | `meituan-longcat/LongCat-Image-Edit` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
+| `LongCatImageEditPipeline` | LongCat-Image-Edit | `meituan-longcat/LongCat-Image-Edit` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
-| `BooguImagePipeline` | Boogu-Image | `Boogu/Boogu-Image-0.1-Base`, `Boogu/Boogu-Image-0.1-Edit` | ✅︎ | | | |
+| `BooguImagePipeline` | Boogu-Image | `Boogu/Boogu-Image-0.1-Base`, `Boogu/Boogu-Image-0.1-Edit` | ✅︎ | | | | — |
-| `StableDiffusionXLPipeline` | Stable-Diffusion-XL | `stabilityai/stable-diffusion-xl-base-1.0` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
+| `StableDiffusionXLPipeline` | Stable-Diffusion-XL | `stabilityai/stable-diffusion-xl-base-1.0` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
-| `StableDiffusion3Pipeline` | Stable-Diffusion-3 | `stabilityai/stable-diffusion-3.5-medium` | ✅︎ | ✅︎ | | ✅︎ |
+| `StableDiffusion3Pipeline` | Stable-Diffusion-3 | `stabilityai/stable-diffusion-3.5-medium` | ✅︎ | ✅︎ | | ✅︎ | — |
-| `CosyVoice3Model` | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | ✅︎ | ✅︎ | | ✅︎ |
+| `CosyVoice3Model` | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | ✅︎ | ✅︎ | | ✅︎ | — |
-| `OmniVoiceModel` | OmniVoice | `k2-fsa/OmniVoice` | ✅︎ | | | |
+| `OmniVoiceModel` | OmniVoice | `k2-fsa/OmniVoice` | ✅︎ | | | | — |
-| `VoxCPM2TalkerForConditionalGeneration` | VoxCPM2 | `openbmb/VoxCPM2` | ✅︎ | | | |
+| `VoxCPM2TalkerForConditionalGeneration` | VoxCPM2 | `openbmb/VoxCPM2` | ✅︎ | | | | — |
-| `MammothModa2ForConditionalGeneration` | MammothModa2-Preview | `bytedance-research/MammothModa2-Preview` | ✅︎ | ✅︎ | | |
+| `MammothModa2ForConditionalGeneration` | MammothModa2-Preview | `bytedance-research/MammothModa2-Preview` | ✅︎ | ✅︎ | | | — |
-| `MammothModa2ForConditionalGeneration` | MammothModa2-Dev (AR-only image understanding) | `bytedance-research/MammothModa2-Dev` | ✅︎ | | | |
+| `MammothModa2ForConditionalGeneration` | MammothModa2-Dev (AR-only image understanding) | `bytedance-research/MammothModa2-Dev` | ✅︎ | | | | — |
-| `Flux2KleinPipeline` | FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B`, `black-forest-labs/FLUX.2-klein-9B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
+| `Flux2KleinPipeline` | FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B`, `black-forest-labs/FLUX.2-klein-9B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
-| `FluxKontextPipeline` | FLUX.1-Kontext-dev | `black-forest-labs/FLUX.1-Kontext-dev` | ✅︎ | ✅︎ | | |
+| `FluxKontextPipeline` | FLUX.1-Kontext-dev | `black-forest-labs/FLUX.1-Kontext-dev` | ✅︎ | ✅︎ | | | — |
-| `FluxPipeline` | FLUX.1-dev | `black-forest-labs/FLUX.1-dev` | ✅︎ | ✅︎ | | ✅︎ |
+| `FluxPipeline` | FLUX.1-dev | `black-forest-labs/FLUX.1-dev` | ✅︎ | ✅︎ | | ✅︎ | — |
-| `FluxPipeline` | FLUX.1-schnell | `black-forest-labs/FLUX.1-schnell` | ✅︎ | ✅︎ | | ✅︎ |
+| `FluxPipeline` | FLUX.1-schnell | `black-forest-labs/FLUX.1-schnell` | ✅︎ | ✅︎ | | ✅︎ | — |
-| `OmniGen2Pipeline` | OmniGen2 | `OmniGen2/OmniGen2` | ✅︎ | ✅︎ | | ✅︎ |
+| `OmniGen2Pipeline` | OmniGen2 | `OmniGen2/OmniGen2` | ✅︎ | ✅︎ | | ✅︎ | — |
-| `StableAudioPipeline` | Stable-Audio-Open | `stabilityai/stable-audio-open-1.0` | ✅︎ | ✅︎ | | ✅︎ |
+| `StableAudioPipeline` | Stable-Audio-Open | `stabilityai/stable-audio-open-1.0` | ✅︎ | ✅︎ | | ✅︎ | — |
-| `SoulXSingerPipeline` | SoulX-Singer (SVS) | `Soul-AILab/SoulX-Singer` | ✅︎ | | | |
+| `SoulXSingerPipeline` | SoulX-Singer (SVS) | `Soul-AILab/SoulX-Singer` | ✅︎ | | | | — |
-| `SoulXSingerSVCPipeline` | SoulX-Singer-SVC | `Soul-AILab/SoulX-Singer` (`model-svc.pt`) | ✅︎ | | | |
+| `SoulXSingerSVCPipeline` | SoulX-Singer-SVC | `Soul-AILab/SoulX-Singer` (`model-svc.pt`) | ✅︎ | | | | — |
-| `AudioXPipeline` | AudioX | `zhangj1an/AudioX` | ✅︎ | ✅︎ | | |
+| `AudioXPipeline` | AudioX | `zhangj1an/AudioX` | ✅︎ | ✅︎ | | | — |
-| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
+| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | ✅︎ | | | | [Published](https://recipes.vllm.ai/Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice) |
-| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-VoiceDesign | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
+| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-VoiceDesign | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
-| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-Base | `Qwen/Qwen3-TTS-12Hz-1.7B-Base` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
+| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-Base | `Qwen/Qwen3-TTS-12Hz-1.7B-Base` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
-| `MingTTSForConditionalGeneration` | Ming-omni-tts dense 0.5B | `inclusionAI/Ming-omni-tts-0.5B` | ✅︎ | | | |
+| `MingTTSForConditionalGeneration` | Ming-omni-tts dense 0.5B | `inclusionAI/Ming-omni-tts-0.5B` | ✅︎ | | | | — |
-| `GLMTTSForConditionalGeneration` | GLM-TTS | `zai-org/GLM-TTS` | ✅︎ | | | |
+| `GLMTTSForConditionalGeneration` | GLM-TTS | `zai-org/GLM-TTS` | ✅︎ | | | | — |
-| `MossTTSNanoForCausalLM` | MOSS-TTS-Nano | `OpenMOSS-Team/MOSS-TTS-Nano` | ✅︎ | | | |
+| `MossTTSNanoForCausalLM` | MOSS-TTS-Nano | `OpenMOSS-Team/MOSS-TTS-Nano` | ✅︎ | | | | — |
-| `MossTTSDelayModel` | MOSS-TTS, MOSS-TTSD, MOSS-SoundEffect, MOSS-VoiceGenerator | `OpenMOSS-Team/MOSS-VoiceGenerator` | ✅︎ | | | |
+| `MossTTSDelayModel` | MOSS-TTS, MOSS-TTSD, MOSS-SoundEffect, MOSS-VoiceGenerator | `OpenMOSS-Team/MOSS-VoiceGenerator` | ✅︎ | | | | — |
-| `MossTTSRealtime` | MOSS-TTS-Realtime | `OpenMOSS-Team/MOSS-TTS-Realtime` | ✅︎ | | | |
+| `MossTTSRealtime` | MOSS-TTS-Realtime | `OpenMOSS-Team/MOSS-TTS-Realtime` | ✅︎ | | | | — |
-| `HiggsAudioV2ForConditionalGeneration` | Higgs-Audio v2 | `bosonai/higgs-audio-v2-generation-3B-base` | ✅︎ | | | |
+| `HiggsAudioV2ForConditionalGeneration` | Higgs-Audio v2 | `bosonai/higgs-audio-v2-generation-3B-base` | ✅︎ | | | | — |
-| `HiggsMultimodalQwen3ForConditionalGeneration` | Higgs-Audio v3 (TTS) | `bosonai/higgs-audio-v3-tts-4b` | ✅︎ | | | |
+| `HiggsMultimodalQwen3ForConditionalGeneration` | Higgs-Audio v3 (TTS) | `bosonai/higgs-audio-v3-tts-4b` | ✅︎ | | | | [Published](https://recipes.vllm.ai/bosonai/higgs-audio-v3-tts-4b) |
-| `IndexTTS2TalkerForConditionalGeneration` | IndexTTS-2 | `IndexTeam/IndexTTS-2` | ✅︎ | | | |
+| `IndexTTS2TalkerForConditionalGeneration` | IndexTTS-2 | `IndexTeam/IndexTTS-2` | ✅︎ | | | | — |
-| `NextStep11Pipeline` | NextStep-1.1 | `stepfun-ai/NextStep-1.1` | ✅︎ | ✅︎ | | ✅︎ |
+| `NextStep11Pipeline` | NextStep-1.1 | `stepfun-ai/NextStep-1.1` | ✅︎ | ✅︎ | | ✅︎ | — |
-| `MiMoAudioModel` | MiMo-Audio-7B-Instruct | `XiaomiMiMo/MiMo-Audio-7B-Instruct` | ✅︎ | ✅︎ | | |
+| `MiMoAudioModel` | MiMo-Audio-7B-Instruct | `XiaomiMiMo/MiMo-Audio-7B-Instruct` | ✅︎ | ✅︎ | | | — |
-| `MiMoV2ASRForCausalLM` | MiMo-V2.5-ASR | `XiaomiMiMo/MiMo-V2.5-ASR` | ✅︎ | ✅︎ | | |
+| `MiMoV2ASRForCausalLM` | MiMo-V2.5-ASR | `XiaomiMiMo/MiMo-V2.5-ASR` | ✅︎ | ✅︎ | | | — |
-| `Flux2Pipeline` | FLUX.2-dev | `black-forest-labs/FLUX.2-dev` | ✅︎ | ✅︎ | | |
+| `Flux2Pipeline` | FLUX.2-dev | `black-forest-labs/FLUX.2-dev` | ✅︎ | ✅︎ | | | — |
-| `FishSpeechSlowARForConditionalGeneration` | Fish Speech S2 Pro | `fishaudio/s2-pro` | ✅︎ | ✅︎ | | ✅︎ |
+| `FishSpeechSlowARForConditionalGeneration` | Fish Speech S2 Pro | `fishaudio/s2-pro` | ✅︎ | ✅︎ | | ✅︎ | — |
-| `DreamIDOmniPipeline` | DreamID-Omni | `XuGuo699/DreamID-Omni` | ✅︎ | ✅︎ | | |
+| `DreamIDOmniPipeline` | DreamID-Omni | `XuGuo699/DreamID-Omni` | ✅︎ | ✅︎ | | | — |
-| `SenseNovaU1Pipeline` | SenseNova-U1 (DiT-only) | `SenseNova/SenseNova-U1-8B-MoT` | ✅︎ | | | |
+| `SenseNovaU1Pipeline` | SenseNova-U1 (DiT-only) | `SenseNova/SenseNova-U1-8B-MoT` | ✅︎ | | | | — |
-| `LancePipeline` | Lance | `bytedance-research/Lance` | ✅︎ | | | |
+| `LancePipeline` | Lance | `bytedance-research/Lance` | ✅︎ | | | | — |
-| `HunyuanVideo15Pipeline` | HunyuanVideo-1.5-T2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | ✅︎ | ✅︎ | | |
+| `HunyuanVideo15Pipeline` | HunyuanVideo-1.5-T2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | ✅︎ | ✅︎ | | | — |
-| `HunyuanVideo15ImageToVideoPipeline` | HunyuanVideo-1.5-I2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_i2v` | ✅︎ | ✅︎ | | |
+| `HunyuanVideo15ImageToVideoPipeline` | HunyuanVideo-1.5-I2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_i2v` | ✅︎ | ✅︎ | | | — |
-| `VoxtralTTSForConditionalGeneration` | Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | ✅︎ | ✅︎ | | |
+| `VoxtralTTSForConditionalGeneration` | Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | ✅︎ | ✅︎ | | | — |
-| `CovoAudioForConditionalGeneration` | Covo-Audio-Chat | `tencent/Covo-Audio-Chat` | ✅︎ | | | |
+| `CovoAudioForConditionalGeneration` | Covo-Audio-Chat | `tencent/Covo-Audio-Chat` | ✅︎ | | | | — |
-|`DyninOmniForConditionalGeneration` | Dynin-Omni | `snu-aidas/Dynin-Omni` | ✅︎ | | | |
+|`DyninOmniForConditionalGeneration` | Dynin-Omni | `snu-aidas/Dynin-Omni` | ✅︎ | | | | — |
-| `MiniCPMO45OmniForConditionalGeneration` | MiniCPM-o 4.5 | `openbmb/MiniCPM-o-4_5` | ✅︎ | | ✅︎ | |
+| `MiniCPMO45OmniForConditionalGeneration` | MiniCPM-o 4.5 | `openbmb/MiniCPM-o-4_5` | ✅︎ | | ✅︎ | | — |
-| `ErnieImagePipeline` | ERNIE-Image | `baidu/ERNIE-Image`, `baidu/ERNIE-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
+| `ErnieImagePipeline` | ERNIE-Image | `baidu/ERNIE-Image`, `baidu/ERNIE-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
-|`HiDreamImagePipeline` | HiDream-I1-Full | `HiDream-ai/HiDream-I1-Full` | ✅︎ | ✅︎ | | |
+|`HiDreamImagePipeline` | HiDream-I1-Full | `HiDream-ai/HiDream-I1-Full` | ✅︎ | ✅︎ | | | — |
+
+The `Recipe` column links the corresponding deployment recipe. `Published` is
+preferred when a corresponding `recipes.vllm.ai` page exists; otherwise the
+link points to the repository recipe. `—` means that no recipe link has been
+added for that model row yet.
+
+For rows with an audited recipe link, the hardware checkmarks reflect hardware
+explicitly documented in that recipe. Other rows retain the implementation
+support metadata until their recipe is audited.
 
 ✅︎ indicates the model is supported on that backend. Empty cells mean not listed as supported on that backend.
 

docs/user_guide/diffusion/attention_backends.md

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

docs/user_guide/examples/online_serving/diffusers_pipeline_adapter.md

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

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)

examples/quantization/check_modelopt_fp8_export.py

@@ -1,247 +0,0 @@
-#!/usr/bin/env python3
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Verify a ModelOpt FP8 diffusers checkpoint exported by
-quantize_wan2_2_modelopt_fp8.py / quantize_hunyuanvideo_15_modelopt_fp8.py.
-
-Three checks:
-  A. transformer/config.json has a sane quantization_config block.
-  B. transformer/*.safetensors contains FP8 (float8_e4m3fn) quantized tensors.
-  C. transformer disk size is materially smaller than a BF16 baseline.
-
-Example:
-    python examples/quantization/check_modelopt_fp8_export.py \\
-        --output ./hv15-480p-modelopt-fp8
-
-    # Optional: compare disk size against a local or HF BF16 baseline.
-    python examples/quantization/check_modelopt_fp8_export.py \\
-        --output ./hv15-480p-modelopt-fp8 \\
-        --baseline hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v
-"""
-
-from __future__ import annotations
-
-import argparse
-import json
-import sys
-from collections import Counter
-from pathlib import Path
-
-SUPPORTED_ALGOS = ("FP8",)
-
-
-def _check_config(transformer_dir: Path) -> tuple[int, str | None]:
-    """Returns (status, quant_algo). status: 0 pass, 1 fail, 2 warn."""
-    cfg_path = transformer_dir / "config.json"
-    if not cfg_path.exists():
-        print(f"[FAIL] {cfg_path} missing.")
-        return 1, None
-
-    with cfg_path.open(encoding="utf-8") as f:
-        cfg = json.load(f)
-
-    qc = cfg.get("quantization_config")
-    if not isinstance(qc, dict):
-        print(f"[FAIL] No `quantization_config` block in {cfg_path}.")
-        return 1, None
-
-    print(f"[A] quantization_config from {cfg_path}:")
-    print(json.dumps(qc, indent=2))
-
-    quant_algo = qc.get("quant_algo")
-    issues = []
-    if qc.get("quant_method") != "modelopt":
-        issues.append(f"quant_method={qc.get('quant_method')!r} (expected 'modelopt')")
-    if quant_algo not in SUPPORTED_ALGOS:
-        issues.append(
-            f"quant_algo={quant_algo!r} (expected one of {SUPPORTED_ALGOS} — vllm-omni adapter may not auto-detect)"
-        )
-
-    if issues:
-        print("[A] WARN — config looks incomplete:")
-        for issue in issues:
-            print(f"    - {issue}")
-        return 2, quant_algo
-    print(f"[A] PASS — config looks correct (quant_algo={quant_algo}).")
-    return 0, quant_algo
-
-
-def _read_safetensors_header(path: Path) -> dict:
-    """Read the JSON header of a safetensors file. Bypass-safe — doesn't materialize tensors.
-
-    Returns {tensor_name: {'dtype': 'F8_E4M3', 'shape': [...], 'data_offsets': [...]}}.
-    Header dtype strings: F8_E4M3, F8_E5M2, BF16, F16, F32, F64, I8, I16, I32, I64, BOOL, U8, ...
-    """
-    import struct
-
-    with open(path, "rb") as f:
-        header_len = struct.unpack("<Q", f.read(8))[0]
-        header = json.loads(f.read(header_len))
-    header.pop("__metadata__", None)
-    return header
-
-
-def _classify_weight_scale_granularity(weight_scale_shapes: list[list[int]]) -> str:
-    """Infer per-tensor vs per-channel vs per-block from sample weight_scale shapes.
-
-    ModelOpt block-wise produces shapes like `[16, 1, 16, 1]` (broadcasting dims of 1
-    interleaved with block-count dims). We count "meaningful" dims — ones with size > 1 —
-    and classify: 0 meaningful dims = per-tensor (scalar), 1 = per-channel, 2+ = per-block.
-    """
-    if not weight_scale_shapes:
-        return "no weight_scale tensors found"
-
-    def meaningful_dims(shape: list[int]) -> int:
-        return sum(1 for d in shape if d > 1)
-
-    per_tensor = sum(1 for s in weight_scale_shapes if meaningful_dims(s) == 0)
-    per_channel = sum(1 for s in weight_scale_shapes if meaningful_dims(s) == 1)
-    per_block = sum(1 for s in weight_scale_shapes if meaningful_dims(s) >= 2)
-    total = len(weight_scale_shapes)
-    if per_tensor == total:
-        return "per-tensor (all scalar scales)"
-    if per_channel == total:
-        return "per-channel (1 meaningful dim)"
-    if per_block == total:
-        return "per-block (2+ meaningful dims — e.g. [M//bm, 1, N//bn, 1] for tiles)"
-    return f"mixed: per-tensor={per_tensor}, per-channel={per_channel}, per-block={per_block} of {total}"
-
-
-def _check_safetensors(transformer_dir: Path) -> int:
-    """Returns 0 on pass, 1 on fail. Reads on-disk dtype from the safetensors header."""
-    files = sorted(transformer_dir.glob("*.safetensors"))
-    if not files:
-        print(f"[FAIL] No *.safetensors in {transformer_dir}.")
-        return 1
-
-    header_dtype_counts: Counter[str] = Counter()
-    sample_quant_weight_keys: list[str] = []
-    sample_scale_keys: list[str] = []
-    weight_scale_shapes: list[list[int]] = []
-    sample_weight_scale_entries: list[tuple[str, list[int]]] = []
-
-    for f in files:
-        try:
-            header = _read_safetensors_header(f)
-        except Exception as exc:
-            print(f"[B] WARN — could not parse header of {f}: {exc}")
-            continue
-        for k, info in header.items():
-            dtype = info.get("dtype", "?")
-            header_dtype_counts[dtype] += 1
-            # FP8 stores weights as F8_E4M3 directly.
-            if dtype.startswith("F8") and k.endswith(".weight") and len(sample_quant_weight_keys) < 5:
-                sample_quant_weight_keys.append(k)
-            if k.endswith(("_scale", ".weight_scale", ".input_scale", "_scale_inv")) and len(sample_scale_keys) < 5:
-                sample_scale_keys.append(k)
-            if k.endswith(".weight_scale"):
-                weight_scale_shapes.append(info.get("shape", []))
-                if len(sample_weight_scale_entries) < 5:
-                    sample_weight_scale_entries.append((k, info.get("shape", [])))
-
-    print(f"\n[B] On-disk dtype counts across {len(files)} safetensors file(s) (from header, not get_tensor):")
-    for dtype, count in sorted(header_dtype_counts.items(), key=lambda kv: -kv[1]):
-        marker = "  <-- FP8" if dtype.startswith("F8") else ""
-        print(f"    {dtype:10s} {count:>6d}{marker}")
-
-    quant_count = sum(c for d, c in header_dtype_counts.items() if d.startswith("F8"))
-    if quant_count == 0:
-        print("[B] FAIL — no FP8 tensors on disk. Calibration likely did not actually quantize the weights.")
-        return 1
-
-    print(f"[B] PASS — {quant_count} FP8 tensors stored on disk.")
-    if sample_quant_weight_keys:
-        print(f"    sample quantized weight tensors: {sample_quant_weight_keys[:3]}")
-    if sample_scale_keys:
-        print(f"    sample scale tensors:            {sample_scale_keys[:3]}")
-    print("    (Note: torch's get_tensor() may return these as bf16 views on some versions —")
-    print("     irrelevant; vLLM's loader uses native FP8 ops.)")
-
-    # Weight-scale granularity — per-tensor (scalar) vs per-channel (1-D) vs per-block (N-D).
-    print(f"\n    weight_scale granularity: {_classify_weight_scale_granularity(weight_scale_shapes)}")
-    for key, shape in sample_weight_scale_entries[:3]:
-        print(f"      {key}: shape {shape}")
-    return 0
-
-
-def _disk_size_gib(p: Path) -> float:
-    return sum(f.stat().st_size for f in p.rglob("*") if f.is_file()) / (1024**3)
-
-
-def _check_size_vs_baseline(transformer_dir: Path, baseline: str | None) -> int:
-    """Returns 0 always (informational only)."""
-    quant_size = _disk_size_gib(transformer_dir)
-    print(f"\n[C] FP8 transformer disk size: {quant_size:.2f} GiB")
-
-    if baseline is None:
-        print("[C] SKIP — pass --baseline <path or HF id> to compare against BF16.")
-        return 0
-
-    baseline_path = Path(baseline)
-    if not baseline_path.exists():
-        # Try HF download.
-        try:
-            from huggingface_hub import snapshot_download
-        except ImportError:
-            print("[C] SKIP — huggingface_hub not installed and baseline not a local path.")
-            return 0
-        print(f"    Downloading baseline transformer from HF: {baseline}")
-        baseline_path = Path(snapshot_download(baseline, allow_patterns=["transformer/*"]))
-
-    bf16_dir = baseline_path / "transformer" if (baseline_path / "transformer").exists() else baseline_path
-    bf16_size = _disk_size_gib(bf16_dir)
-    if bf16_size == 0:
-        print(f"[C] WARN — baseline transformer dir empty: {bf16_dir}")
-        return 0
-
-    # Expected reduction for FP8: ~50%.
-    min_reduction = 30
-    reduction = (1 - quant_size / bf16_size) * 100
-    print(f"[C] BF16 baseline transformer disk size: {bf16_size:.2f} GiB ({bf16_dir})")
-    print(f"[C] Disk reduction: {reduction:.1f}%  (FP8 transformer is {quant_size / bf16_size:.0%} of BF16)")
-    if reduction < min_reduction:
-        print(
-            f"[C] WARN — FP8 should typically reduce disk by ~40-50%; <{min_reduction}% suggests partial quantization."
-        )
-    return 0
-
-
-def main() -> None:
-    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
-    p.add_argument("--output", required=True, help="Path to the exported ModelOpt FP8 checkpoint root.")
-    p.add_argument(
-        "--baseline",
-        default=None,
-        help="Optional BF16 baseline (local diffusers dir or HF id) for disk-size comparison.",
-    )
-    args = p.parse_args()
-
-    out_root = Path(args.output).expanduser().resolve()
-    transformer_dir = out_root / "transformer"
-    if not transformer_dir.exists():
-        print(f"[FAIL] {transformer_dir} does not exist.")
-        sys.exit(1)
-
-    print(f"Checking: {out_root}\n")
-
-    fail = 0
-    config_status, _quant_algo = _check_config(transformer_dir)
-    fail |= config_status
-    fail |= _check_safetensors(transformer_dir)
-    _check_size_vs_baseline(transformer_dir, args.baseline)
-
-    print()
-    if fail == 0:
-        print("=" * 60)
-        print("ALL CHECKS PASSED — checkpoint looks ready for vllm-omni serving.")
-    elif fail == 1:
-        print("=" * 60)
-        print("FAILURES detected — calibration may need to be re-run.")
-        sys.exit(1)
-    else:
-        print("=" * 60)
-        print("WARNINGS only — checkpoint may serve but with caveats. See [A] above.")
-
-
-if __name__ == "__main__":
-    main()

examples/quantization/quantize_hunyuanvideo_15_modelopt_fp8.py

@@ -1,448 +0,0 @@
-#!/usr/bin/env python3
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Quantize HunyuanVideo-1.5 (480p T2V) to a ModelOpt FP8 Hugging Face checkpoint.
-
-Calibrates the DiT transformer using a small video prompt set and exports a
-diffusers-style directory whose transformer carries ModelOpt FP8 metadata.
-The exported checkpoint is consumable by vllm-omni's ModelOpt FP8 adapter
-(see vllm_omni/diffusion/model_loader/checkpoint_adapters/modelopt.py).
-
-Layers kept full precision match the #2728 / #2795 pattern: modulation,
-AdaLayerNorm, entry/exit projections, embeddings, the token refiner path,
-and final proj_out. MHA quantizers are off by default; HV-1.5 self-attention
-empirically degrades under FP8 (see #2920 ablation).
-
-Example:
-    python examples/quantization/quantize_hunyuanvideo_15_modelopt_fp8.py \\
-        --model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v \\
-        --output ./hv15-480p-modelopt-fp8 \\
-        --overwrite
-"""
-
-from __future__ import annotations
-
-import argparse
-import copy
-import json
-import re
-import shutil
-import sys
-from pathlib import Path
-from typing import Any
-
-import torch
-from diffusers import DiffusionPipeline
-
-DEFAULT_PROMPTS = [
-    "A dog running across a field of golden wheat.",
-    "An astronaut riding a horse across the surface of Mars, red dust swirling, cinematic wide shot.",
-    "A hummingbird hovering in front of a vibrant red flower, slow motion, macro shot.",
-    "A crackling campfire at night under a starry sky, sparks rising into the dark.",
-    "An underwater shot of a coral reef with tropical fish swimming by, sun rays piercing the water.",
-    "A close-up of a blooming rose covered in morning dew, soft natural light.",
-    "A peaceful mountain village at dawn, mist rolling over the rooftops, cinematic establishing shot.",
-    "A skateboarder doing a kickflip in an urban plaza, slow motion, golden hour lighting.",
-]
-
-
-def _build_parser() -> argparse.ArgumentParser:
-    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
-    p.add_argument("--model", required=True, help="Input HV-1.5 diffusers directory or HF id.")
-    p.add_argument("--output", required=True, help="Output directory for the ModelOpt FP8 checkpoint.")
-    p.add_argument("--dtype", choices=("bfloat16", "float16"), default="bfloat16")
-    p.add_argument("--height", type=int, default=480)
-    p.add_argument("--width", type=int, default=832)
-    p.add_argument(
-        "--num-frames",
-        type=int,
-        default=33,
-        help="Frames per calibration sample. 33 matches the typical short benchmark.",
-    )
-    p.add_argument("--guidance-scale", type=float, default=6.0)
-    p.add_argument(
-        "--calib-steps",
-        type=int,
-        default=10,
-        help="Denoising steps per calibration prompt (10 is enough for amax statistics).",
-    )
-    p.add_argument("--calib-size", type=int, default=8, help="How many prompts to use for calibration.")
-    p.add_argument("--seed", type=int, default=42)
-    p.add_argument(
-        "--prompt",
-        action="append",
-        default=[],
-        help="Custom calibration prompt. Repeat to provide multiple.",
-    )
-    p.add_argument(
-        "--quantize-mha",
-        action="store_true",
-        help="Enable FP8 attention K/V/softmax quantizers. Off by default — empirically degrades HV-1.5 video output.",
-    )
-    p.add_argument(
-        "--weight-block-size",
-        type=str,
-        default=None,
-        help="Per-block weight quantization as 'M,N' (e.g. '128,128' for 128x128 tiles). "
-        "Default: per-tensor (one scale per linear). Block-wise typically gives tighter quality at "
-        "negligible memory cost. Static FP8 is exempt from upstream vLLM's online block-wise gate.",
-    )
-    p.add_argument("--overwrite", action="store_true", help="Replace an existing output directory.")
-    return p
-
-
-def _parse_block_size(spec: str | None) -> list[int] | None:
-    if spec is None:
-        return None
-    parts = [int(x) for x in spec.split(",") if x.strip()]
-    if len(parts) != 2:
-        raise SystemExit(f"--weight-block-size must be 'M,N' (2 ints), got {spec!r}")
-    return parts
-
-
-def _require_modelopt() -> Any:
-    try:
-        import modelopt.torch.quantization as mtq
-    except ModuleNotFoundError as exc:
-        raise SystemExit(
-            "NVIDIA ModelOpt is not installed. Install with:\n"
-            "  pip install 'nvidia-modelopt[all]'\n"
-            f"Original error: {exc}"
-        ) from exc
-    return mtq
-
-
-def _ensure_paths(args: argparse.Namespace) -> tuple[str, Path]:
-    model_path = args.model
-    output_dir = Path(args.output).expanduser().resolve()
-    if output_dir.exists():
-        if not args.overwrite:
-            raise SystemExit(f"Output directory already exists: {output_dir}\nPass --overwrite to replace it.")
-        shutil.rmtree(output_dir)
-    return model_path, output_dir
-
-
-def _select_dtype(name: str) -> torch.dtype:
-    return {"bfloat16": torch.bfloat16, "float16": torch.float16}[name]
-
-
-def _build_prompts(args: argparse.Namespace) -> list[str]:
-    prompts = args.prompt or DEFAULT_PROMPTS
-    if args.calib_size <= 0:
-        raise SystemExit("--calib-size must be positive.")
-    if len(prompts) < args.calib_size:
-        repeats = (args.calib_size + len(prompts) - 1) // len(prompts)
-        prompts = (prompts * repeats)[: args.calib_size]
-    return prompts[: args.calib_size]
-
-
-# Layers to KEEP at full precision (mirror of the #2920 wiring + #2728/#2795 skip pattern).
-# - x_embedder, image_embedder, context_embedder*, time_embed*, cond_type_embed: entry/embedding
-# - norm_out, norm1*.linear, norm1_context*.linear, norm2*, norm2_context*: AdaLayerNorm modulation
-# - proj_out: final output projection
-# - token_refiner*: text-encoder refinement uses diffusers raw nn.Linear
-def _filter_func_hv15(name: str) -> bool:
-    pattern = re.compile(
-        r"(proj_out.*|"
-        r".*(x_embedder|image_embedder|context_embedder|context_embedder_2|"
-        r"time_embed|cond_type_embed|"
-        r"norm_out|norm1\.linear|norm1_context\.linear|norm2|norm2_context|"
-        r"token_refiner).*)"
-    )
-    return pattern.match(name) is not None
-
-
-def _mha_filter_func(name: str) -> bool:
-    pattern = re.compile(
-        r".*(q_bmm_quantizer|k_bmm_quantizer|v_bmm_quantizer|softmax_quantizer|bmm2_output_quantizer).*"
-    )
-    return pattern.match(name) is not None
-
-
-def _disable_known_problematic_quantizers(mtq: Any, backbone: torch.nn.Module, *, quantize_mha: bool) -> None:
-    if not hasattr(mtq, "disable_quantizer"):
-        return
-    mtq.disable_quantizer(backbone, _filter_func_hv15)
-    if not quantize_mha:
-        mtq.disable_quantizer(backbone, _mha_filter_func)
-
-
-def _load_pipeline(model_path: str, dtype: torch.dtype) -> DiffusionPipeline:
-    pipe = DiffusionPipeline.from_pretrained(model_path, torch_dtype=dtype)
-    if hasattr(pipe, "set_progress_bar_config"):
-        pipe.set_progress_bar_config(disable=True)
-    pipe.to("cuda")
-    return pipe
-
-
-def _build_forward_loop(pipe: DiffusionPipeline, args: argparse.Namespace, prompts: list[str]):
-    generator = torch.Generator(device="cuda")
-
-    # Try to set guidance on the pipeline's guider object up front (modern
-    # diffusers HV-1.5 uses a Guider abstraction, not a per-call kwarg). Falls
-    # back silently — calibration uses whatever default the pipeline ships with.
-    guider = getattr(pipe, "guider", None)
-    if guider is not None and hasattr(guider, "guidance_scale"):
-        try:
-            guider.guidance_scale = args.guidance_scale
-        except Exception:
-            pass
-
-    base_kwargs = dict(
-        height=args.height,
-        width=args.width,
-        num_frames=args.num_frames,
-        num_inference_steps=args.calib_steps,
-        output_type="latent",
-    )
-
-    def forward_loop(*_unused_args, **_unused_kwargs) -> None:
-        with torch.inference_mode():
-            for idx, prompt in enumerate(prompts):
-                generator.manual_seed(args.seed + idx)
-                # Try with guidance_scale first; fall back without on TypeError
-                # for pipelines (like HV-1.5) that take CFG via guider config.
-                try:
-                    pipe(prompt=prompt, generator=generator, guidance_scale=args.guidance_scale, **base_kwargs)
-                except TypeError as exc:
-                    if "guidance_scale" not in str(exc):
-                        raise
-                    pipe(prompt=prompt, generator=generator, **base_kwargs)
-
-    return forward_loop
-
-
-def _summarize_export(output_dir: Path) -> None:
-    cfg_path = output_dir / "transformer" / "config.json"
-    if not cfg_path.exists():
-        print(f"[warn] {cfg_path} missing.", file=sys.stderr)
-        return
-    with cfg_path.open(encoding="utf-8") as f:
-        cfg = json.load(f)
-    qc = cfg.get("quantization_config")
-    if not isinstance(qc, dict):
-        print("[warn] No quantization_config in transformer/config.json.", file=sys.stderr)
-        return
-    print("Export summary:")
-    print(f"  quant_method: {qc.get('quant_method')}")
-    print(f"  quant_algo:   {qc.get('quant_algo')}")
-    producer = qc.get("producer")
-    if isinstance(producer, dict):
-        print(f"  producer:     {producer.get('name')} {producer.get('version')}")
-    print(f"  config path:  {cfg_path}")
-
-
-def _force_export_quantized_weights(backbone: torch.nn.Module, dtype: torch.dtype) -> int:
-    """Convert in-memory weights of quantized modules to actual FP8 storage.
-
-    `export_hf_checkpoint` skips this step for unknown model types (HV-1.5 isn't
-    in ModelOpt's recognized-model registry), so we must call the per-weight
-    export helper ourselves. Same workaround as the HunyuanImage-3 calibration
-    helper.
-    """
-    from modelopt.torch.export.quant_utils import (
-        QUANTIZATION_NONE,
-        get_quantization_format,
-        quantizer_attr_names,
-        weight_attr_names,
-    )
-    from modelopt.torch.export.unified_export_hf import _export_quantized_weight
-
-    exported = 0
-    for name, module in backbone.named_modules():
-        try:
-            quantization_format = get_quantization_format(module)
-        except Exception as exc:
-            print(f"[warn] Could not inspect quantization format for {name}: {exc}", file=sys.stderr)
-            continue
-        if quantization_format == QUANTIZATION_NONE:
-            continue
-        for weight_name in weight_attr_names(module):
-            quantizer_attrs = quantizer_attr_names(weight_name)
-            weight_quantizer = getattr(module, quantizer_attrs.weight_quantizer, None)
-            if weight_quantizer is None or not getattr(weight_quantizer, "is_enabled", False):
-                continue
-            _export_quantized_weight(module, dtype, weight_name)
-            exported += 1
-    return exported
-
-
-def _hv15_quant_config_block(weight_block_size: list[int] | None = None) -> dict:
-    """Mirror ModelOpt FP8 metadata expected by vllm-omni's adapter (#2913).
-
-    Same shape as the HunyuanImage-3 author's _hunyuan_quant_config(). When
-    `weight_block_size` is given, advertise block-wise weight quantization in
-    the saved metadata (so consumers know to expect multi-element scale tensors).
-    """
-    weights_cfg: dict = {"dynamic": False, "num_bits": 8, "type": "float"}
-    if weight_block_size is not None:
-        weights_cfg["strategy"] = "block"
-        weights_cfg["block_structure"] = f"{weight_block_size[0]}x{weight_block_size[1]}"
-    return {
-        "config_groups": {
-            "group_0": {
-                "input_activations": {"dynamic": False, "num_bits": 8, "type": "float"},
-                "weights": weights_cfg,
-                "targets": ["Linear"],
-            }
-        },
-        "ignore": [
-            "context_embedder*",
-            "context_embedder_2*",
-            "cond_type_embed*",
-            "image_embedder*",
-            "norm1.linear*",
-            "norm1_context.linear*",
-            "norm2*",
-            "norm2_context*",
-            "norm_out*",
-            "proj_out*",
-            "time_embed*",
-            "token_refiner*",
-            "x_embedder*",
-        ],
-        "producer": {"name": "modelopt"},
-        "quant_algo": "FP8",
-        "quant_method": "modelopt",
-    }
-
-
-def _patch_quant_config(output_dir: Path, weight_block_size: list[int] | None = None) -> None:
-    """Inject quant_algo: FP8 + config_groups into transformer/config.json so
-    vllm-omni's adapter (#2913) recognises the checkpoint as ModelOpt FP8."""
-    cfg_path = output_dir / "transformer" / "config.json"
-    with cfg_path.open(encoding="utf-8") as f:
-        cfg = json.load(f)
-
-    new_qc = _hv15_quant_config_block(weight_block_size=weight_block_size)
-    existing = cfg.get("quantization_config")
-    if isinstance(existing, dict):
-        producer = existing.get("producer")
-        if isinstance(producer, dict):
-            new_qc["producer"] = producer
-
-    cfg["quantization_config"] = new_qc
-    with cfg_path.open("w", encoding="utf-8") as f:
-        json.dump(cfg, f, indent=2)
-
-
-def _save_pipeline_with_fp8_transformer(
-    pipe: DiffusionPipeline,
-    model_path: str,
-    output_dir: Path,
-    max_shard_size: str = "5GB",
-) -> None:
-    """Save the pipeline with the (now FP8) transformer.
-
-    Copies the source directory verbatim except for `transformer/`, then
-    saves the transformer with quantizers hidden so the state dict contains
-    only the FP8 weights + scale tensors.
-    """
-    from modelopt.torch.export.diffusers_utils import hide_quantizers_from_state_dict
-
-    src = Path(model_path)
-    if not src.exists():
-        from huggingface_hub import snapshot_download
-
-        src = Path(snapshot_download(model_path))
-
-    if output_dir.exists():
-        shutil.rmtree(output_dir)
-    shutil.copytree(src, output_dir, ignore=shutil.ignore_patterns("transformer"))
-
-    transformer_out = output_dir / "transformer"
-    # `hide_quantizers_from_state_dict` walks named_modules(); pass the actual
-    # nn.Module (transformer), not the diffusers Pipeline wrapper.
-    with hide_quantizers_from_state_dict(pipe.transformer):
-        pipe.transformer.save_pretrained(
-            str(transformer_out),
-            safe_serialization=True,
-            max_shard_size=max_shard_size,
-        )
-
-
-def main() -> None:
-    args = _build_parser().parse_args()
-    if not torch.cuda.is_available():
-        raise SystemExit("CUDA is required for ModelOpt FP8 quantization.")
-
-    mtq = _require_modelopt()
-    model_path, output_dir = _ensure_paths(args)
-    dtype = _select_dtype(args.dtype)
-    prompts = _build_prompts(args)
-
-    print("Quantization plan:")
-    weight_block_size = _parse_block_size(args.weight_block_size)
-
-    print(f"  input:           {args.model}")
-    print(f"  output:          {output_dir}")
-    print(f"  dtype:           {dtype}")
-    print(f"  height/width:    {args.height}x{args.width}")
-    print(f"  num_frames:      {args.num_frames}")
-    print(f"  calib_size:      {len(prompts)}")
-    print(f"  calib_steps:     {args.calib_steps}")
-    print(f"  quantize_mha:    {args.quantize_mha}")
-    print(
-        f"  weight strategy: {'block-wise ' + str(weight_block_size) if weight_block_size else 'per-tensor (default)'}"
-    )
-
-    pipe = _load_pipeline(model_path, dtype)
-    backbone = pipe.transformer
-
-    quant_config = copy.deepcopy(mtq.FP8_DEFAULT_CFG)
-    if weight_block_size is not None:
-        # Switch from per-tensor (default) to block-wise weight quantization.
-        # ModelOpt's wildcard "*weight_quantizer" matches every linear's weight quantizer.
-        quant_config["quant_cfg"]["*weight_quantizer"] = {
-            "num_bits": (4, 3),  # E4M3 (FP8 weights, same as default)
-            "block_sizes": {-1: weight_block_size[1], -2: weight_block_size[0]},
-        }
-        print(
-            f"  -> overriding weight quantizer with block_sizes={weight_block_size} "
-            f"({weight_block_size[0]}x{weight_block_size[1]} tiles)"
-        )
-
-    forward_loop = _build_forward_loop(pipe, args, prompts)
-    quantized = mtq.quantize(backbone, quant_config, forward_loop)
-    if quantized is not None:
-        pipe.transformer = quantized
-        backbone = quantized
-
-    _disable_known_problematic_quantizers(mtq, backbone, quantize_mha=args.quantize_mha)
-
-    print("\nForcing FP8 weight serialization (HV-1.5 isn't in ModelOpt's recognized-model registry,")
-    print("so we have to call the per-weight export helper ourselves)...")
-    exported = _force_export_quantized_weights(backbone, dtype)
-    print(f"  -> {exported} weights converted to FP8 in memory")
-    if exported == 0:
-        raise SystemExit(
-            "No quantized weights were exported. Calibration may have skipped every layer "
-            "(check the disable_quantizer regex) or `mtq.quantize` did not actually wrap any "
-            "weight quantizers."
-        )
-
-    print("\nSaving pipeline with FP8 transformer...")
-    _save_pipeline_with_fp8_transformer(pipe, model_path, output_dir)
-    _patch_quant_config(output_dir, weight_block_size=weight_block_size)
-    print(f"Saved to: {output_dir}")
-    _summarize_export(output_dir)
-
-    print("\nNext: validate the checkpoint with vllm-omni:")
-    print(
-        "  python examples/offline_inference/text_to_video/text_to_video.py \\\n"
-        f"    --model {output_dir} \\\n"
-        "    --quantization fp8 \\\n"
-        "    --prompt 'A dog running across a field of golden wheat.' \\\n"
-        f"    --height {args.height} --width {args.width} --num-frames {args.num_frames} \\\n"
-        "    --num-inference-steps 30 --guidance-scale 6.0 --seed 42 \\\n"
-        "    --output outputs/hv15_modelopt_fp8.mp4 \\\n"
-        "    --enforce-eager"
-    )
-    print(
-        "\n  (--quantization fp8 is auto-upgraded to ModelOpt FP8 at runtime because the "
-        "checkpoint's config.json has modelopt metadata.)"
-    )
-
-
-if __name__ == "__main__":
-    main()

examples/quantization/quantize_wan2_2_modelopt_fp8.py

@@ -1,471 +0,0 @@
-#!/usr/bin/env python3
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Quantize Wan2.2 (TI2V-5B, 704x1280 T2V) to a ModelOpt FP8 Hugging Face checkpoint.
-
-Calibrates the DiT transformer using a small video prompt set and exports a
-diffusers-style directory whose transformer carries ModelOpt FP8 metadata.
-The exported checkpoint is consumable by vllm-omni's ModelOpt FP8 adapter
-(see vllm_omni/diffusion/model_loader/checkpoint_adapters/modelopt.py).
-
-Layers kept full precision match the #2728 / #2795 pattern: condition embedder
-(time/text/image), patch embedding, modulation (scale_shift_table), final
-norm + proj_out, and sequence-parallel helpers. All attention + FFN linears
-are quantized — static calibration handles the numerics that online FP8
-couldn't (see #2920 ablation).
-
-Default target is `Wan-AI/Wan2.2-TI2V-5B-Diffusers`, the dense 5B variant that
-fits 80GB BF16. The A14B MoE variants need 2+ GPUs and are out of scope here.
-
-Example:
-    python examples/quantization/quantize_wan2_2_modelopt_fp8.py \\
-        --model Wan-AI/Wan2.2-TI2V-5B-Diffusers \\
-        --output ./wan22-ti2v-modelopt-fp8 \\
-        --overwrite
-"""
-
-from __future__ import annotations
-
-import argparse
-import copy
-import json
-import re
-import shutil
-import sys
-from pathlib import Path
-from typing import Any
-
-import torch
-from diffusers import DiffusionPipeline
-
-DEFAULT_PROMPTS = [
-    "A dog running across a field of golden wheat.",
-    "An astronaut riding a horse across the surface of Mars, red dust swirling, cinematic wide shot.",
-    "A hummingbird hovering in front of a vibrant red flower, slow motion, macro shot.",
-    "A crackling campfire at night under a starry sky, sparks rising into the dark.",
-    "An underwater shot of a coral reef with tropical fish swimming by, sun rays piercing the water.",
-    "A close-up of a blooming rose covered in morning dew, soft natural light.",
-    "A peaceful mountain village at dawn, mist rolling over the rooftops, cinematic establishing shot.",
-    "A skateboarder doing a kickflip in an urban plaza, slow motion, golden hour lighting.",
-]
-
-
-def _build_parser() -> argparse.ArgumentParser:
-    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
-    p.add_argument("--model", required=True, help="Input Wan2.2 diffusers directory or HF id.")
-    p.add_argument("--output", required=True, help="Output directory for the ModelOpt FP8 checkpoint.")
-    p.add_argument("--dtype", choices=("bfloat16", "float16"), default="bfloat16")
-    p.add_argument("--height", type=int, default=704, help="Calibration video height (Wan2.2 TI2V-5B native: 704).")
-    p.add_argument("--width", type=int, default=1280, help="Calibration video width (Wan2.2 TI2V-5B native: 1280).")
-    p.add_argument(
-        "--num-frames",
-        type=int,
-        default=49,
-        help="Frames per calibration sample. 49 matches the typical short benchmark; "
-        "use 17 to reduce memory pressure during calibration.",
-    )
-    p.add_argument("--guidance-scale", type=float, default=5.0)
-    p.add_argument(
-        "--calib-steps",
-        type=int,
-        default=10,
-        help="Denoising steps per calibration prompt (10 is enough for amax statistics).",
-    )
-    p.add_argument("--calib-size", type=int, default=8, help="How many prompts to use for calibration.")
-    p.add_argument("--seed", type=int, default=42)
-    p.add_argument(
-        "--prompt",
-        action="append",
-        default=[],
-        help="Custom calibration prompt. Repeat to provide multiple.",
-    )
-    p.add_argument(
-        "--quantize-mha",
-        action="store_true",
-        help="Enable FP8 attention K/V/softmax quantizers. Off by default — Wan2.2's long attention "
-        "sequences amplified FP8 drift in the online ablation (see #2920).",
-    )
-    p.add_argument(
-        "--weight-block-size",
-        type=str,
-        default=None,
-        help="Per-block weight quantization as 'M,N' (e.g. '128,128'). Default per-tensor. "
-        "Note: vllm-omni's ModelOpt adapter may not yet dispatch block-wise scales — check #2924 "
-        "for the HV-1.5 investigation status before relying on this.",
-    )
-    p.add_argument("--overwrite", action="store_true", help="Replace an existing output directory.")
-    return p
-
-
-def _parse_block_size(spec: str | None) -> list[int] | None:
-    if spec is None:
-        return None
-    parts = [int(x) for x in spec.split(",") if x.strip()]
-    if len(parts) != 2:
-        raise SystemExit(f"--weight-block-size must be 'M,N' (2 ints), got {spec!r}")
-    return parts
-
-
-def _require_modelopt() -> Any:
-    try:
-        import modelopt.torch.quantization as mtq
-    except ModuleNotFoundError as exc:
-        raise SystemExit(
-            "NVIDIA ModelOpt is not installed. Install with:\n"
-            "  pip install 'nvidia-modelopt[all]'\n"
-            f"Original error: {exc}"
-        ) from exc
-    return mtq
-
-
-def _ensure_paths(args: argparse.Namespace) -> tuple[str, Path]:
-    model_path = args.model
-    output_dir = Path(args.output).expanduser().resolve()
-    if output_dir.exists():
-        if not args.overwrite:
-            raise SystemExit(f"Output directory already exists: {output_dir}\nPass --overwrite to replace it.")
-        shutil.rmtree(output_dir)
-    return model_path, output_dir
-
-
-def _select_dtype(name: str) -> torch.dtype:
-    return {"bfloat16": torch.bfloat16, "float16": torch.float16}[name]
-
-
-def _build_prompts(args: argparse.Namespace) -> list[str]:
-    prompts = args.prompt or DEFAULT_PROMPTS
-    if args.calib_size <= 0:
-        raise SystemExit("--calib-size must be positive.")
-    if len(prompts) < args.calib_size:
-        repeats = (args.calib_size + len(prompts) - 1) // len(prompts)
-        prompts = (prompts * repeats)[: args.calib_size]
-    return prompts[: args.calib_size]
-
-
-# Layers to KEEP at full precision. Wan2.2's module naming:
-# - condition_embedder: time_embedder, time_proj, text_embedder, image_embedder (I2V)
-# - patch_embedding: Conv3dLayer (already not Linear, belt-and-suspenders skip)
-# - scale_shift_table: nn.Parameter modulation (not Linear, but pattern guard)
-# - norm_out: AdaLayerNorm final
-# - proj_out: final nn.Linear
-# - timestep_proj_prepare / output_scale_shift_prepare: SP helpers
-def _filter_func_wan22(name: str) -> bool:
-    pattern = re.compile(
-        r"(proj_out.*|"
-        r".*(condition_embedder|patch_embedding|"
-        r"norm_out|scale_shift_table|"
-        r"timestep_proj_prepare|output_scale_shift_prepare).*)"
-    )
-    return pattern.match(name) is not None
-
-
-def _mha_filter_func(name: str) -> bool:
-    pattern = re.compile(
-        r".*(q_bmm_quantizer|k_bmm_quantizer|v_bmm_quantizer|softmax_quantizer|bmm2_output_quantizer).*"
-    )
-    return pattern.match(name) is not None
-
-
-def _disable_known_problematic_quantizers(mtq: Any, backbone: torch.nn.Module, *, quantize_mha: bool) -> None:
-    if not hasattr(mtq, "disable_quantizer"):
-        return
-    mtq.disable_quantizer(backbone, _filter_func_wan22)
-    if not quantize_mha:
-        mtq.disable_quantizer(backbone, _mha_filter_func)
-
-
-def _load_pipeline(model_path: str, dtype: torch.dtype) -> DiffusionPipeline:
-    pipe = DiffusionPipeline.from_pretrained(model_path, torch_dtype=dtype)
-    if hasattr(pipe, "set_progress_bar_config"):
-        pipe.set_progress_bar_config(disable=True)
-    pipe.to("cuda")
-    return pipe
-
-
-def _build_forward_loop(pipe: DiffusionPipeline, args: argparse.Namespace, prompts: list[str]):
-    generator = torch.Generator(device="cuda")
-
-    # Try setting guidance on the pipeline's guider if present (newer diffusers APIs).
-    guider = getattr(pipe, "guider", None)
-    if guider is not None and hasattr(guider, "guidance_scale"):
-        try:
-            guider.guidance_scale = args.guidance_scale
-        except Exception:
-            pass
-
-    base_kwargs = dict(
-        height=args.height,
-        width=args.width,
-        num_frames=args.num_frames,
-        num_inference_steps=args.calib_steps,
-        output_type="latent",
-    )
-
-    def forward_loop(*_unused_args, **_unused_kwargs) -> None:
-        with torch.inference_mode():
-            for idx, prompt in enumerate(prompts):
-                generator.manual_seed(args.seed + idx)
-                # Try with guidance_scale first; fall back without on TypeError
-                # for pipelines that take CFG via guider config only.
-                try:
-                    pipe(prompt=prompt, generator=generator, guidance_scale=args.guidance_scale, **base_kwargs)
-                except TypeError as exc:
-                    if "guidance_scale" not in str(exc):
-                        raise
-                    pipe(prompt=prompt, generator=generator, **base_kwargs)
-
-    return forward_loop
-
-
-def _summarize_export(output_dir: Path) -> None:
-    cfg_path = output_dir / "transformer" / "config.json"
-    if not cfg_path.exists():
-        print(f"[warn] {cfg_path} missing.", file=sys.stderr)
-        return
-    with cfg_path.open(encoding="utf-8") as f:
-        cfg = json.load(f)
-    qc = cfg.get("quantization_config")
-    if not isinstance(qc, dict):
-        print("[warn] No quantization_config in transformer/config.json.", file=sys.stderr)
-        return
-    print("Export summary:")
-    print(f"  quant_method: {qc.get('quant_method')}")
-    print(f"  quant_algo:   {qc.get('quant_algo')}")
-    producer = qc.get("producer")
-    if isinstance(producer, dict):
-        print(f"  producer:     {producer.get('name')} {producer.get('version')}")
-    print(f"  config path:  {cfg_path}")
-
-
-def _force_export_quantized_weights(backbone: torch.nn.Module, dtype: torch.dtype) -> int:
-    """Convert in-memory weights of quantized modules to actual FP8 storage.
-
-    `export_hf_checkpoint` skips this step for unknown model types (Wan2.2 isn't
-    in ModelOpt's recognized-model registry), so we must call the per-weight
-    export helper ourselves. Same workaround as the HunyuanVideo-1.5 / HunyuanImage-3
-    calibration helpers.
-    """
-    from modelopt.torch.export.quant_utils import (
-        QUANTIZATION_NONE,
-        get_quantization_format,
-        quantizer_attr_names,
-        weight_attr_names,
-    )
-    from modelopt.torch.export.unified_export_hf import _export_quantized_weight
-
-    exported = 0
-    for name, module in backbone.named_modules():
-        try:
-            quantization_format = get_quantization_format(module)
-        except Exception as exc:
-            print(f"[warn] Could not inspect quantization format for {name}: {exc}", file=sys.stderr)
-            continue
-        if quantization_format == QUANTIZATION_NONE:
-            continue
-        for weight_name in weight_attr_names(module):
-            quantizer_attrs = quantizer_attr_names(weight_name)
-            weight_quantizer = getattr(module, quantizer_attrs.weight_quantizer, None)
-            if weight_quantizer is None or not getattr(weight_quantizer, "is_enabled", False):
-                continue
-            _export_quantized_weight(module, dtype, weight_name)
-            exported += 1
-    return exported
-
-
-def _wan22_quant_config_block(weight_block_size: list[int] | None = None) -> dict:
-    """Mirror ModelOpt FP8 metadata expected by vllm-omni's adapter (#2913)."""
-    weights_cfg: dict = {"dynamic": False, "num_bits": 8, "type": "float"}
-    if weight_block_size is not None:
-        weights_cfg["strategy"] = "block"
-        weights_cfg["block_structure"] = f"{weight_block_size[0]}x{weight_block_size[1]}"
-    return {
-        "config_groups": {
-            "group_0": {
-                "input_activations": {"dynamic": False, "num_bits": 8, "type": "float"},
-                "weights": weights_cfg,
-                "targets": ["Linear"],
-            }
-        },
-        "ignore": [
-            "condition_embedder*",
-            "norm_out*",
-            "output_scale_shift_prepare*",
-            "patch_embedding*",
-            "proj_out*",
-            "scale_shift_table*",
-            "timestep_proj_prepare*",
-        ],
-        "producer": {"name": "modelopt"},
-        "quant_algo": "FP8",
-        "quant_method": "modelopt",
-    }
-
-
-def _patch_quant_config(output_dir: Path, weight_block_size: list[int] | None = None) -> None:
-    """Inject quant_algo: FP8 + config_groups into each transformer's config.json
-    so vllm-omni's adapter (#2913) recognises the checkpoint as ModelOpt FP8.
-
-    Patches both ``transformer`` and ``transformer_2`` (A14B) when present."""
-    for sub in ("transformer", "transformer_2"):
-        cfg_path = output_dir / sub / "config.json"
-        if not cfg_path.exists():
-            continue
-        with cfg_path.open(encoding="utf-8") as f:
-            cfg = json.load(f)
-
-        new_qc = _wan22_quant_config_block(weight_block_size=weight_block_size)
-        existing = cfg.get("quantization_config")
-        if isinstance(existing, dict):
-            producer = existing.get("producer")
-            if isinstance(producer, dict):
-                new_qc["producer"] = producer
-
-        cfg["quantization_config"] = new_qc
-        with cfg_path.open("w", encoding="utf-8") as f:
-            json.dump(cfg, f, indent=2)
-
-
-def _list_transformers(pipe: DiffusionPipeline) -> list[tuple[str, torch.nn.Module]]:
-    """Return [(attr_name, module)] for every DiT on the pipeline.
-
-    Wan2.2 A14B has both ``transformer`` (low-noise) and ``transformer_2``
-    (high-noise). TI2V-5B has only ``transformer``.
-    """
-    out = []
-    for attr in ("transformer", "transformer_2"):
-        mod = getattr(pipe, attr, None)
-        if mod is not None:
-            out.append((attr, mod))
-    return out
-
-
-def _save_pipeline_with_fp8_transformer(
-    pipe: DiffusionPipeline,
-    model_path: str,
-    output_dir: Path,
-    max_shard_size: str = "5GB",
-) -> None:
-    """Copy source dir verbatim minus transformer dirs, then save each quantized DiT.
-
-    Handles A14B's dual transformer (``transformer`` + ``transformer_2``) as well
-    as the single-DiT TI2V-5B.
-    """
-    from modelopt.torch.export.diffusers_utils import hide_quantizers_from_state_dict
-
-    src = Path(model_path)
-    if not src.exists():
-        from huggingface_hub import snapshot_download
-
-        src = Path(snapshot_download(model_path))
-
-    if output_dir.exists():
-        shutil.rmtree(output_dir)
-    shutil.copytree(src, output_dir, ignore=shutil.ignore_patterns("transformer", "transformer_2"))
-
-    for attr, backbone in _list_transformers(pipe):
-        transformer_out = output_dir / attr
-        # Pass the nn.Module (transformer), not the Pipeline wrapper.
-        with hide_quantizers_from_state_dict(backbone):
-            backbone.save_pretrained(
-                str(transformer_out),
-                safe_serialization=True,
-                max_shard_size=max_shard_size,
-            )
-
-
-def main() -> None:
-    args = _build_parser().parse_args()
-    if not torch.cuda.is_available():
-        raise SystemExit("CUDA is required for ModelOpt FP8 quantization.")
-
-    mtq = _require_modelopt()
-    model_path, output_dir = _ensure_paths(args)
-    dtype = _select_dtype(args.dtype)
-    prompts = _build_prompts(args)
-    weight_block_size = _parse_block_size(args.weight_block_size)
-
-    print("Quantization plan:")
-    print(f"  input:           {args.model}")
-    print(f"  output:          {output_dir}")
-    print(f"  dtype:           {dtype}")
-    print(f"  height/width:    {args.height}x{args.width}")
-    print(f"  num_frames:      {args.num_frames}")
-    print(f"  calib_size:      {len(prompts)}")
-    print(f"  calib_steps:     {args.calib_steps}")
-    print(f"  quantize_mha:    {args.quantize_mha}")
-    print(
-        f"  weight strategy: {'block-wise ' + str(weight_block_size) if weight_block_size else 'per-tensor (default)'}"
-    )
-
-    pipe = _load_pipeline(model_path, dtype)
-    transformers = _list_transformers(pipe)
-    if not transformers:
-        raise SystemExit("Pipeline has no transformer or transformer_2 attribute.")
-    print(f"  found {len(transformers)} transformer(s): {', '.join(a for a, _ in transformers)}")
-
-    quant_config = copy.deepcopy(mtq.FP8_DEFAULT_CFG)
-    if weight_block_size is not None:
-        quant_config["quant_cfg"]["*weight_quantizer"] = {
-            "num_bits": (4, 3),
-            "block_sizes": {-1: weight_block_size[1], -2: weight_block_size[0]},
-        }
-        print(
-            f"  -> overriding weight quantizer with block_sizes={weight_block_size} "
-            f"({weight_block_size[0]}x{weight_block_size[1]} tiles)"
-        )
-
-    forward_loop = _build_forward_loop(pipe, args, prompts)
-
-    # Quantize each DiT. The forward_loop runs the full pipeline, so quantizers
-    # on every attached transformer (A14B: transformer + transformer_2) get
-    # calibrated during the same passes.
-    for attr, backbone in transformers:
-        print(f"\nQuantizing {attr}...")
-        quantized = mtq.quantize(backbone, quant_config, forward_loop)
-        if quantized is not None:
-            setattr(pipe, attr, quantized)
-        _disable_known_problematic_quantizers(mtq, getattr(pipe, attr), quantize_mha=args.quantize_mha)
-
-    # Export weights to FP8 only AFTER all transformers are calibrated. Exporting
-    # inside the loop would convert transformer's weights to real FP8, and the
-    # next transformer's forward_loop (full pipeline) would then fake-quantize an
-    # already-FP8 weight -> "fake_e4m3fy not implemented for Float8_e4m3fn".
-    print("\nForcing FP8 weight serialization (Wan2.2 isn't in ModelOpt's recognized-model registry,")
-    print("so we have to call the per-weight export helper ourselves)...")
-    total_exported = 0
-    for attr, backbone in _list_transformers(pipe):
-        exported = _force_export_quantized_weights(backbone, dtype)
-        print(f"  -> {exported} weights converted to FP8 in {attr}")
-        total_exported += exported
-
-    if total_exported == 0:
-        raise SystemExit(
-            "No quantized weights were exported. Calibration may have skipped every layer "
-            "(check the disable_quantizer regex) or `mtq.quantize` did not actually wrap any "
-            "weight quantizers."
-        )
-
-    print("\nSaving pipeline with FP8 transformer(s)...")
-    _save_pipeline_with_fp8_transformer(pipe, model_path, output_dir)
-    _patch_quant_config(output_dir, weight_block_size=weight_block_size)
-    print(f"Saved to: {output_dir}")
-    _summarize_export(output_dir)
-
-    print("\nNext: validate the checkpoint with vllm-omni:")
-    print(
-        "  python examples/offline_inference/text_to_video/text_to_video.py \\\n"
-        f"    --model {output_dir} \\\n"
-        "    --quantization fp8 \\\n"
-        "    --prompt 'A dog running across a field of golden wheat.' \\\n"
-        f"    --height {args.height} --width {args.width} --num-frames {args.num_frames} \\\n"
-        "    --num-inference-steps 30 --guidance-scale 5.0 --seed 42 \\\n"
-        "    --output outputs/wan22_modelopt_fp8.mp4"
-    )
-    print(
-        "\n  (--quantization fp8 is auto-upgraded to ModelOpt FP8 at runtime because the "
-        "checkpoint's config.json has modelopt metadata.)"
-    )
-
-
-if __name__ == "__main__":
-    main()

recipes/MiniMaxAI/MiniMax-H3-NPU.md

@@ -79,7 +79,8 @@ sharing the tokenizer, processor, text encoder, and VAEs from `FL2VA`.
 ### Multi-NPU: 768P combined-service configuration
 
 Use Ulysses sequence parallelism degree 8, text-encoder tensor parallelism
-degree 8, native tiled VAE patch parallelism degree 8, and layerwise offload:
+degree 8, native tiled VAE patch parallelism degree 8, and distributed
+layerwise offload:
 
 ```bash
 export ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
@@ -98,7 +99,7 @@ vllm serve "${MODEL}" \
   --usp 8 \
   --ring 1 \
   --text-encoder-tp-size 8 \
-  --enable-layerwise-offload \
+  --enable-distributed-layerwise-offload \
   --vae-parallel-mode tile \
   --vae-use-tiling \
   --vae-patch-parallel-size 8 \
@@ -112,6 +113,11 @@ H3 is CFG-distilled, so `--cfg-parallel-size` must remain 1.
 The same endpoint accepts `task=t2va`, `task=fl2va`, and `task=ref2va`; no
 partition restart is required. Layerwise offload applies to both DiTs.
 
+On Atlas 800I A3 (64 GB HBM per device) the combined service does not fit at
+768P without offloading or sharding: use distributed layerwise offload (as
+above) or HSDP — see
+[§ Memory and attention optimizations](#memory-and-attention-optimizations-a3).
+
 ### Optional optimizations
 
 Two independent optimizations may be enabled on top of the configuration
@@ -134,6 +140,63 @@ Keep `--ring 1` when using RainFusion: the `rf_v2` kernel ranks key blocks
 over the whole sequence, so ring parallelism would split away the keys it
 needs. Scale with `--usp` instead.
 
+## Memory and attention optimizations (A3)
+
+### Fitting 768P into 64 GB HBM: distributed layerwise offload or HSDP
+
+Each Atlas 800I A3 NPU has 64 GB of HBM, which is not enough for the
+combined MiniMax-H3 service at 768P. Enable **one** of the following at
+server startup:
+
+**Distributed layerwise offload** — DiT layers are offloaded with parameters
+gathered across the parallel group instead of replicated per rank:
+
+```bash
+  --enable-distributed-layerwise-offload
+```
+
+Measured peak NPU memory is about 45 GB per device for Ref2VA with a 13.88 s
+input video generating a 15 s 768P (1344x768) output video. Host (CPU) memory
+usage is high with this option.
+
+**HSDP** — hybrid sharded data parallelism for the DiT parameters. The
+multi-stream memory-reuse knob is mandatory with this flag:
+
+```bash
+export MULTI_STREAM_MEMORY_REUSE=2
+```
+
+```bash
+  --use-hsdp
+```
+
+Host memory usage is small, but large shapes may still OOM; HBM usage
+optimizations for this path are ongoing.
+
+### FLASH_ATTN backend with MindIE-SD
+
+Keep `--diffusion-attention-backend FLASH_ATTN` and install MindIE-SD (see
+[§ Environment](#environment)). On NPU this backend carries most of the
+memory-reduction and performance work: a mask-free packed varlen path that
+never materializes the quadratic `full_qk` padding mask, and K/V prefix
+slicing, both driven by the packed `cu_seqlens` metadata emitted by the H3
+transformer.
+
+### Optional: LaserAttention fused kernel
+
+For an additional attention speedup, select the Ascend LaserAttention fused
+kernel before starting the server:
+
+```bash
+export MINDIE_SD_FA_TYPE="ascend_laser_attention"
+```
+
+This requires the FLASH_ATTN backend and MindIE-SD. H3 automatically applies
+exact power-of-two input pre-scaling (`laser_input_scale=256`) so the
+kernel's fp16 workspace cannot overflow on outlier activations. Measured
+speedup numbers will be added here.
+
+
 ## HTTP API examples
 
 The request format is identical to the GPU recipe; see
@@ -154,8 +217,8 @@ configuration above:
 
 | Workload | Configuration |
 |----------|---------------|
-| T2VA, 209 frames, 1344x768 | TE TP8, layerwise offload, Ulysses 8, VPP8 tile, regional compile |
+| T2VA, 209 frames, 1344x768 | TE TP8, distributed layerwise offload, Ulysses 8, VPP8 tile, regional compile |
-| Ref2VA (prompt + video), 124 frames, 1344x768 | TE TP8, layerwise offload, Ulysses 8, VPP8 tile, regional compile |
+| Ref2VA (prompt + video), 124 frames, 1344x768 | TE TP8, distributed layerwise offload, Ulysses 8, VPP8 tile, regional compile |
 
 These measurements describe the validated shapes rather than a general
 throughput guarantee.

recipes/MiniMaxAI/MiniMax-H3-Spark-GB10.md

@@ -4,6 +4,17 @@ This recipe uses online FP8 weight quantization, tiled VAE decode, and a single
 resident partition. GB10 is a unified-memory platform, so unlike the discrete-GPU
 recipes it uses **no offload of any kind** — see the capacity note below.
 
+Validated on:
+
+- Host: DGX Spark (GB10), aarch64
+- GPUs: 1 (unified memory)
+- Driver: <FILL_DRIVER_VERSION>
+- vLLM: 0.26.0
+- vLLM-Omni: `main` at `e1aa6eae75c460cd1893bc320546e81e66973831`
+- Workloads: T2VA and Ref2VA, 960x576, `duration=8.0`, 50 steps, `flow_shift=12`,
+  `seed=1101`
+
+
 ## Capacity requirements
 
 | Resource | DGX Spark (GB10) |
@@ -11,7 +22,8 @@ recipes it uses **no offload of any kind** — see the capacity note below.
 | Unified memory | 128 GB LPDDR5X (~121 GiB usable) |
 | GPUs | 1 |
 | Checkpoint storage | 135 GiB per partition |
-| Measured peak allocator high-water (FP8) | 97.7 GiB |
+| Measured peak allocator high-water, T2VA (FP8) | 97.7 GiB |
+| Measured peak allocator high-water, Ref2VA (FP8) | 102.8 GiB |
 
 `FL2VA` and `Ref2VA` are separate 135 GiB checkpoint partitions. Start one server
 at a time.
@@ -29,9 +41,21 @@ math relative to the RTX and datacenter recipes:
   fit in 121 GiB. Online FP8 quantizes the DiT only (62 GiB to ~31 GiB); the
   Qwen3-VL text encoder and both VAEs stay BF16.
 
-With ~97.7 GiB of the pool held by the allocator at peak, there is roughly 23 GiB
+With ~97.7 GiB of the pool held by the allocator at peak, T2VA leaves roughly
-of headroom for the OS, page cache, and any other process on the box. Treat this
+23 GiB of headroom for the OS, page cache, and any other process on the box.
-platform as single-tenant while a request is in flight.
+Ref2VA is tighter on both counts. The allocator high-water rises to 102.8 GiB,
+and the allocator figure understates true pool usage.
+
+Sampling `/proc/meminfo` at 1 Hz across a 50-step Ref2VA request on a host with
+`MemTotal: 127600524 kB` (121.7 GiB), `MemAvailable` started at 8.1 GiB, fell to
+a trough of **7.4 GiB**, and recovered to 8.9 GiB. At the trough the pool held
+roughly **114.3 GiB** — 11.4 GiB more than the 102.8 GiB the allocator reports,
+which is the CUDA context, non-PyTorch allocations, and page cache the header
+does not count.
+
+Treat this platform as strictly single-tenant for Ref2VA. There is no room for a
+second process, and adding reference images, reference videos, or a larger output
+shape will run into the OOM killer.
 
 ## One GB10: 960x576, 8 seconds
 
@@ -94,6 +118,52 @@ ffprobe -v error -show_entries \
 Both requests return `200 OK` with an MP4 body: H.264 video at the requested
 geometry plus a 32 kHz stereo AAC track.
 
+## One GB10: Ref2VA, 960x576, 8 seconds
+
+`FL2VA` and `Ref2VA` are separate 135 GiB partitions and only one fits at a time.
+Stop the FL2VA server before starting this one.
+
+```bash
+CUDA_VISIBLE_DEVICES=0 \
+FLASHINFER_DISABLE_VERSION_CHECK=1 \
+VLLM_WORKER_MULTIPROC_METHOD=spawn \
+VLLM_OMNI_VIDEO_SYNC_TIMEOUT=14400 \
+vllm serve /path/to/MiniMax-H3/Ref2VA \
+  --omni --trust-remote-code --host 0.0.0.0 --port 8000 \
+  --init-timeout 3600 \
+  --num-gpus 1 --tensor-parallel-size 1 --text-encoder-tp-size 1 \
+  --usp 1 --ring 1 --vae-patch-parallel-size 1 \
+  --vae-parallel-mode tile --vae-use-tiling \
+  --quantization fp8 --enforce-eager \
+  --diffusion-attention-backend CUDNN_ATTN
+```
+
+The 7200 s timeout used for T2VA is too short here: a 50-step Ref2VA request
+takes over an hour on this platform. Use 14400 s.
+
+Unlike `t2va`, `ref2va` defaults to a 16:9 output ratio, so `aspect_ratio` can be
+omitted when `width` and `height` are supplied.
+
+```bash
+curl --fail-with-body -sS -X POST http://127.0.0.1:8000/v1/videos/sync \
+  -D ref2va-headers.txt \
+  -w 'client_e2e_s=%{time_total}\n' \
+  --max-time 14400 \
+  -F 'input_reference=@/path/to/reference.jpg;type=image/jpeg' \
+  -F 'prompt=A hand rests on a wooden table in warm evening light, then slowly opens and turns toward the camera, with quiet room ambience.' \
+  -F 'width=960' \
+  -F 'height=576' \
+  -F 'fps=24' \
+  -F 'num_inference_steps=50' \
+  -F 'flow_shift=12' \
+  -F 'seed=1101' \
+  -F 'extra_params={"task":"ref2va","duration":8.0,"audio_flow_shift":3.0}' \
+  -o h3-gb10-ref2va.mp4
+```
+
+The response is `200 OK` with an MP4 body: H.264 at the requested geometry plus a
+32 kHz stereo AAC track. `ffprobe` reports 8.032 s for `duration=8.0`.
+
 ## Measured latency
 
 Single GB10, FP8, eager, cuDNN attention, tiled VAE, one request at a time:
@@ -102,32 +172,69 @@ Single GB10, FP8, eager, cuDNN attention, tiled VAE, one request at a time:
 | --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: |
 | T2VA, 4 s, smoke shape | 10 | — | — | 8.68 s | — | 0.49 s | 95.9 s | — |
 | T2VA, 8 s, 960x576 | 50 | 0.25 s | 2088 s | 42.61 s | 70 s | 4.84 s | 2169.4 s (36 min 9 s) | 97.7 GiB |
+| Ref2VA, 8 s, 960x576, one image (3 runs) | 10 | 8.43-10.94 s | 665.5-748.8 s | 66.55-74.88 s | 67.9-69.8 s | 3.63 s | 770.3-861.1 s | 102.8 GiB |
+| Ref2VA, 8 s, 960x576, one image | 50 | 9.66 s | 4039.6 s | 80.79 s | 69.4 s | 3.63 s | 4157.0 s (69 min 17 s) | 102.8 GiB |
 
 How these were measured:
 
 - Stage times come from `--enable-diffusion-pipeline-profiler`, read from the
-  `X-Stage-Durations` response header of `/v1/videos/sync`.
+  `X-Stage-Durations` response header of `/v1/videos/sync`. The keys are
+  `MiniMaxH3Pipeline.encode_prompt` (text encode),
+  `MiniMaxH3Pipeline.diffuse` (denoise), and `MiniMaxH3Pipeline.decode`
+  (VAE decode).
+- MP4 mux is not a header field. It is the
+  `Video response encoding (MP4 bytes)` line in the server log.
 - Peak memory is the `X-Peak-Memory-MB` header, which reports the CUDA allocator
   high-water mark (`torch.cuda.max_memory_reserved`) for the request. It excludes
   the CUDA context and non-PyTorch allocations, so true pool usage is somewhat
   higher. `nvidia-smi --query-gpu=memory.used` is not usable on GB10 — unified
-  memory has no discrete VRAM to report and the field returns `[N/A]`.
+  memory has no discrete VRAM to report and the field returns `[N/A]`. See the
-- Text encode and VAE decode were measured on a separate 10-step run at the same
+  capacity section for a `/proc/meminfo` cross-check.
-  960x576 / 8 s geometry. Both are independent of step count, so they transfer to
+- For T2VA, text encode and VAE decode were measured on a separate 10-step run at
-  the 50-step row unchanged.
+  the same 960x576 / 8 s geometry. Both are independent of step count, so they
+  transfer to the 50-step row unchanged. The Ref2VA MP4 mux figure transfers the
+  same way, from the 10-step run at identical output geometry and duration.
+- Ref2VA peak memory is byte-identical between the 10-step and 50-step runs
+  (`X-Peak-Memory-MB: 105318.000` in both), confirming that peak is set by output
+  geometry and reference input rather than step count.
 - Stage times do not sum exactly to end-to-end. Queueing, result transfer, and
   engine bookkeeping sit outside the profiled stages.
 
 Denoise dominates at 96% of end-to-end, so step count and geometry are the only
 knobs that materially change the wall clock on this platform.
 
-Repeated identical 10-step requests varied by up to 23% in denoise time
+Ref2VA costs 1.92x the wall clock of T2VA at the same geometry and step count:
+4157.0 s versus 2169.4 s. Text encoding shows the most dramatic ratio — 9.66 s
+versus 0.25 s, 39x — but at 0.2% of end-to-end it is not what makes Ref2VA slow.
+The cost is in denoise. A single reference image expands the Qwen3-VL
+presentation to 6883 tokens (logged as `MiniMax H3 ref2va Qwen presentation:
+6883 tokens, 1 reference images`), and those tokens sit in the attention context
+of every denoise step. Per-step time rises from 42.61 s to 80.79 s, and that
+1.90x multiplies across all 50 steps. VAE decode is unchanged at ~69 s, as
+expected for an identical output shape.
+
+Repeated identical 10-step T2VA requests varied by up to 23% in denoise time
 (397 s to 490 s) on a machine that had already been under sustained load. Treat
 these as single-run measurements on one machine, not a throughput guarantee, and
 expect the slower end of that range for back-to-back jobs.
 
+Ref2VA per-step time is load-dependent in the same way. On an idle machine the
+first 10-step request ran at 66.55 s per step. The 50-step request that followed
+occupied the box for 69 consecutive minutes and averaged 80.79 s per step. Two
+further 10-step requests taken immediately afterwards settled at 74.88 s and
+74.31 s — 12.5% slower than the idle run, and within 0.8% of each other, so this
+is a stable thermal state rather than measurement noise.
+
+Plan for 81 s per step when sizing a full 50-step Ref2VA job; 66 s per step is
+only reachable on a cold machine and does not hold for the length of one request.
+
+Ref2VA text encode converged downward across the four runs (10.94 s, 9.66 s,
+8.52 s, 8.43 s) as first-touch costs fell away; treat ~8.5 s as steady state.
+VAE decode was stable at 67.9-69.8 s throughout.
+
 Set `VLLM_OMNI_VIDEO_SYNC_TIMEOUT` well above the expected request time — the
-default 1800 s is shorter than a 50-step run on this platform.
+default 1800 s is shorter than a 50-step run on this platform, and 7200 s is
+shorter than a 50-step Ref2VA run.
 
 ## Known limitations
 
@@ -139,11 +246,17 @@ default 1800 s is shorter than a 50-step run on this platform.
 - `t2va` rejects requests without an explicit `aspect_ratio`. When `width` and
   `height` are both supplied they take precedence and `aspect_ratio` only has to
   name one of the supported ratios.
-- Sustained load degrades throughput. See the variance note above.
+- Sustained load degrades throughput. See the variance notes above.
 - Only one partition is resident. For Ref2VA, stop the FL2VA server and restart
-  the same command with `/path/to/MiniMax-H3/Ref2VA`. Ref2VA reference video
+  with `/path/to/MiniMax-H3/Ref2VA`. Ref2VA with one reference image is measured
-  count and prompt length increase activation memory; begin with one request at
+  above; reference videos and additional images push the token count higher and
-  a time and keep the short edge at 576.
+  were not measured. With `MemAvailable` bottoming out at 7.4 GiB during a
+  single-image request, there is very little room to grow — re-measure before
+  adding references or raising the output shape.
+- Ref2VA requires a vLLM-Omni build newer than `v0.26.0`. See the version note
+  under **Validated on**: FP8 weight loading and image-only Ref2VA are both
+  broken on `release/v0.26.0`.suggestion is  `v0.26.1`.
 - Online FP8 is incompatible with layerwise offload — the offload path produces a
   weight stride the Cutlass FP8 kernel rejects. This is not a practical
   restriction here since offload is unusable on GB10 anyway.
+```

recipes/README.md

@@ -19,6 +19,11 @@ add files for this repo. Use one Markdown file per model family by default.
 > Individual pages resolve at `https://recipes.vllm.ai/<provider>/<model-id>`,
 > e.g. [`bosonai/higgs-audio-v3-tts-4b`](https://recipes.vllm.ai/bosonai/higgs-audio-v3-tts-4b).
 
+The `Recipe` column in the Supported Models documentation links to the
+published recipe when available and otherwise to the corresponding recipe in
+this directory. The recipe files remain the source of truth for commands and
+validation details.
+
 Example layout:
 
 ```text

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/attention/test_flash_attn.py

@@ -9,6 +9,10 @@ This script tests two main scenarios:
 2. Case 2: Comparing FlashAttention and SDPA backends for batch_size=2 with padding
 """
 
+import sys
+from types import SimpleNamespace
+from unittest.mock import Mock
+
 import pytest
 import torch
 
@@ -371,6 +375,264 @@ def test_packed_varlen_metadata_must_be_complete(monkeypatch):
         impl.forward_cuda(query, query, query, metadata)
 
 
+# ---------------------------------------------------------------------------
+# NPU mask-free packed attention (PR #5891).
+#
+# These tests cover the backend-side NPU branches that ``forward_fa_npu`` adds:
+#   - boundary resolution in ``_resolve_packed_seq_npu`` (pure Python, no device
+#     sync, no MindIE-SD import);
+#   - env dispatch in ``forward_fa_npu`` (which op a packed forward takes);
+#   - the laser input pre-scaling math in ``_forward_prefix_kv_slice_npu``.
+#
+# They run on CPU without a real NPU by injecting a fake ``mindiesd`` module
+# into ``sys.modules`` (same pattern as ``test_paged_attention`` /
+# ``benchmarks/conftest``). No production code is exercised for real kernels.
+
+
+def _npu_impl(causal: bool = False) -> FlashAttentionImpl:
+    return FlashAttentionImpl(num_heads=2, head_size=4, softmax_scale=0.5, causal=causal)
+
+
+def _fake_mindiesd(monkeypatch, *, attention_forward=None, attention_forward_varlen=None):
+    """Install a stub ``mindiesd`` so local ``from mindiesd import ...`` works.
+
+    Restored automatically by monkeypatch. Defaults are no-op callables so the
+    top-level import in ``forward_fa_npu`` succeeds even when the test mocks the
+    instance methods that would otherwise call them.
+    """
+    fake = SimpleNamespace(
+        attention_forward=attention_forward or Mock(return_value=torch.zeros(1)),
+        attention_forward_varlen=attention_forward_varlen or Mock(return_value=torch.zeros(1)),
+    )
+    monkeypatch.setitem(sys.modules, "mindiesd", fake)
+    return fake
+
+
+# --- Test group A: boundary resolution (_resolve_packed_seq_npu) -------------
+
+
+def test_resolve_packed_seq_accepts_real_plus_pad_contract():
+    impl = _npu_impl()
+    q = torch.randn(1, 8, 2, 4)
+    cu = torch.tensor([0, 5, 8], dtype=torch.int32)
+    extra = {"cu_seqlens_q": cu, "cu_seqlens_k": cu, "max_seqlen_q": 5, "max_seqlen_k": 5}
+
+    assert impl._resolve_packed_seq_npu(q, q, extra) == ([5, 8], [5, 8])
+
+
+def test_resolve_packed_seq_single_document_no_padding():
+    impl = _npu_impl()
+    q = torch.randn(1, 8, 2, 4)
+    cu = torch.tensor([0, 8], dtype=torch.int32)
+    extra = {"cu_seqlens_q": cu, "cu_seqlens_k": cu, "max_seqlen_q": 8, "max_seqlen_k": 8}
+
+    assert impl._resolve_packed_seq_npu(q, q, extra) == ([8], [8])
+
+
+@pytest.mark.parametrize(
+    "case",
+    [
+        "causal",
+        "missing_key",
+        "batch_ne_1",
+        "cu_too_long",
+        "cu_qk_mismatch",
+        "max_seqlen_not_int",
+        "used_gt_total",
+        "used_zero",
+        "pad_longer_than_real",
+    ],
+)
+def test_resolve_packed_seq_rejects_malformed_metadata(case):
+    """Any metadata outside the exact [real, pad] contract returns None so the
+    caller falls back to the masked path."""
+    cu3 = torch.tensor([0, 5, 8], dtype=torch.int32)
+    cu2 = torch.tensor([0, 8], dtype=torch.int32)
+    cu4 = torch.tensor([0, 3, 5, 8], dtype=torch.int32)
+    cu_short = torch.tensor([0, 3, 8], dtype=torch.int32)
+    base = {"cu_seqlens_q": cu3, "cu_seqlens_k": cu3, "max_seqlen_q": 5, "max_seqlen_k": 5}
+    q = torch.randn(1, 8, 2, 4)
+
+    impl = _npu_impl(causal=(case == "causal"))
+    if case == "missing_key":
+        extra = {"cu_seqlens_q": cu3, "max_seqlen_q": 5}  # 3 of 4 keys absent
+    elif case == "batch_ne_1":
+        q = torch.randn(2, 8, 2, 4)
+        extra = base
+    elif case == "cu_too_long":
+        extra = {**base, "cu_seqlens_q": cu4, "cu_seqlens_k": cu4}
+    elif case == "cu_qk_mismatch":
+        extra = {**base, "cu_seqlens_q": cu3, "cu_seqlens_k": cu2}
+    elif case == "max_seqlen_not_int":
+        extra = {**base, "max_seqlen_q": 5.0}
+    elif case == "used_gt_total":
+        extra = {**base, "max_seqlen_q": 9}  # 9 > total length 8
+    elif case == "used_zero":
+        extra = {**base, "max_seqlen_q": 0, "max_seqlen_k": 0}
+    elif case == "pad_longer_than_real":
+        # real=3, pad=5: the real document must be the longer one.
+        extra = {
+            "cu_seqlens_q": cu_short,
+            "cu_seqlens_k": cu_short,
+            "max_seqlen_q": 3,
+            "max_seqlen_k": 3,
+        }
+    else:
+        extra = base
+
+    assert impl._resolve_packed_seq_npu(q, q, extra) is None
+
+
+# --- Test group B: env dispatch in forward_fa_npu (current behavior) --------
+
+
+def _npu_impl_with_mocked_paths() -> tuple[FlashAttentionImpl, torch.Tensor]:
+    """FlashAttentionImpl whose two packed paths are Mocks returning a sentinel.
+
+    Lets a test assert which branch ``forward_fa_npu`` dispatched to without
+    running any real kernel.
+    """
+    impl = _npu_impl()
+    sentinel = torch.tensor([1.0])
+    impl._forward_varlen_packed_npu = Mock(return_value=sentinel)
+    impl._forward_prefix_kv_slice_npu = Mock(return_value=sentinel)
+    return impl, sentinel
+
+
+def test_npu_env_unset_uses_varlen(monkeypatch):
+    monkeypatch.delenv("MINDIE_SD_FA_TYPE", raising=False)
+    _fake_mindiesd(monkeypatch)
+    impl, _ = _npu_impl_with_mocked_paths()
+    metadata = AttentionMetadata(extra={"npu_attn_varlen": True})
+
+    impl.forward_fa_npu(torch.randn(1, 8, 2, 4), *[torch.randn(1, 8, 2, 4)] * 2, metadata)
+
+    impl._forward_varlen_packed_npu.assert_called_once()
+    impl._forward_prefix_kv_slice_npu.assert_not_called()
+
+
+def test_npu_env_laser_uses_slice(monkeypatch):
+    monkeypatch.setenv("MINDIE_SD_FA_TYPE", "ascend_laser_attention")
+    _fake_mindiesd(monkeypatch)
+    impl, _ = _npu_impl_with_mocked_paths()
+    metadata = AttentionMetadata(extra={"npu_attn_varlen": True})
+
+    impl.forward_fa_npu(torch.randn(1, 8, 2, 4), *[torch.randn(1, 8, 2, 4)] * 2, metadata)
+
+    impl._forward_prefix_kv_slice_npu.assert_called_once()
+    impl._forward_varlen_packed_npu.assert_not_called()
+
+
+@pytest.mark.parametrize("fa_type", ["prompt_flash_attn", "fused_attn_score", "ascend_laser_attenton", "garbage"])
+def test_npu_env_non_laser_value_falls_back_to_varlen(monkeypatch, fa_type):
+    """Locks in the CURRENT dispatch behavior (PR #5891).
+
+    Any value other than ``ascend_laser_attention`` — including the valid
+    ``prompt_flash_attn``/``fused_attn_score`` and typos — is silently routed to
+    the TND varlen op, which cannot honor ``MINDIE_SD_FA_TYPE``. This is review
+    comment #1 on the PR and is intentionally captured here as a regression
+    net: update this test when the dispatch starts honoring/rejecting these.
+    """
+    monkeypatch.setenv("MINDIE_SD_FA_TYPE", fa_type)
+    _fake_mindiesd(monkeypatch)
+    impl, _ = _npu_impl_with_mocked_paths()
+    metadata = AttentionMetadata(extra={"npu_attn_varlen": True})
+
+    impl.forward_fa_npu(torch.randn(1, 8, 2, 4), *[torch.randn(1, 8, 2, 4)] * 2, metadata)
+
+    impl._forward_varlen_packed_npu.assert_called_once()
+    impl._forward_prefix_kv_slice_npu.assert_not_called()
+
+
+def test_npu_varlen_opt_in_unset_takes_mask_path(monkeypatch):
+    """Models that do not set ``npu_attn_varlen`` (e.g. Wan/Cosmos) keep using
+    the existing masked attention_forward path; the new branches are skipped."""
+    fake_forward = Mock(return_value=torch.zeros(1, 8, 2, 4))
+    _fake_mindiesd(monkeypatch, attention_forward=fake_forward)
+    impl, _ = _npu_impl_with_mocked_paths()
+    metadata = AttentionMetadata(
+        attn_mask=torch.tensor([[True, True, True, True, True, False, False, False]]),
+        extra={},  # no npu_attn_varlen opt-in
+    )
+    q = torch.randn(1, 8, 2, 4)
+
+    out = impl.forward_fa_npu(q, q, q, metadata)
+
+    impl._forward_varlen_packed_npu.assert_not_called()
+    impl._forward_prefix_kv_slice_npu.assert_not_called()
+    fake_forward.assert_called_once()
+    assert out is fake_forward.return_value
+
+
+# --- Test group C: laser input pre-scaling in _forward_prefix_kv_slice_npu --
+
+
+def test_prefix_kv_slice_applies_laser_input_scaling(monkeypatch):
+    monkeypatch.setenv("MINDIE_SD_FA_TYPE", "ascend_laser_attention")
+    captured: dict = {}
+
+    def fake_attention_forward(query, key, value, **kwargs):
+        captured["query"] = query
+        captured["key"] = key
+        captured["value"] = value
+        captured.update(kwargs)
+        return torch.ones_like(query)  # deterministic stand-in kernel output
+
+    _fake_mindiesd(monkeypatch, attention_forward=fake_attention_forward)
+    impl = _npu_impl()  # softmax_scale == 0.5
+    q = torch.randn(1, 8, 2, 4)
+    cu = torch.tensor([0, 5, 8], dtype=torch.int32)
+    extra = {
+        "cu_seqlens_q": cu,
+        "cu_seqlens_k": cu,
+        "max_seqlen_q": 5,
+        "max_seqlen_k": 5,
+        "laser_input_scale": 256.0,
+    }
+
+    out = impl._forward_prefix_kv_slice_npu(q, q, q, extra)
+
+    # K/V sliced to the valid prefix (used_k == 5); q keeps full length.
+    assert captured["key"].shape == (1, 5, 2, 4)
+    assert captured["value"].shape == (1, 5, 2, 4)
+    assert captured["query"].shape == (1, 8, 2, 4)
+    assert captured["attn_mask"] is None
+    # Pre-divide q/k/v by the factor and compensate the kernel scale by its
+    # square (exact for power-of-two); output scaled back by the factor.
+    torch.testing.assert_close(captured["query"], q / 256.0)
+    torch.testing.assert_close(captured["key"], q[:, :5] / 256.0)
+    assert captured["scale"] == pytest.approx(0.5 * 256.0**2)
+    torch.testing.assert_close(out, torch.ones(1, 8, 2, 4) * 256.0)
+
+
+def test_prefix_kv_slice_no_scaling_without_factor(monkeypatch):
+    monkeypatch.setenv("MINDIE_SD_FA_TYPE", "ascend_laser_attention")
+    captured: dict = {}
+
+    def fake_attention_forward(query, key, value, **kwargs):
+        captured["query"] = query
+        captured.update(kwargs)
+        return torch.ones_like(query)
+
+    _fake_mindiesd(monkeypatch, attention_forward=fake_attention_forward)
+    impl = _npu_impl()
+    q = torch.randn(1, 8, 2, 4)
+    cu = torch.tensor([0, 5, 8], dtype=torch.int32)
+    extra = {
+        "cu_seqlens_q": cu,
+        "cu_seqlens_k": cu,
+        "max_seqlen_q": 5,
+        "max_seqlen_k": 5,
+        # no laser_input_scale → no pre-scaling
+    }
+
+    out = impl._forward_prefix_kv_slice_npu(q, q, q, extra)
+
+    torch.testing.assert_close(captured["query"], q)
+    assert captured["scale"] == pytest.approx(0.5)
+    torch.testing.assert_close(out, torch.ones(1, 8, 2, 4))
+
+
 if __name__ == "__main__":
     print("Running FlashAttention Padding Tests...")
     print("=" * 60)

tests/diffusion/attention/test_rainfusion_plan.py

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

tests/diffusion/cache/test_teacache_extractors.py

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

tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py

@@ -492,6 +492,10 @@ def test_cudnn_packed_attention_uses_python_length_without_padding_mask():
     class FakeBackend:
         supports_prefix_kv_slicing = True
 
+        @classmethod
+        def supports_packed_mask_free(cls) -> bool:
+            return False
+
     class FakeAttention(torch.nn.Module):
         attn_backend = FakeBackend
 
@@ -526,6 +530,48 @@ def test_cudnn_packed_attention_uses_python_length_without_padding_mask():
     assert attention.attention.metadata.extra["valid_kv_length"] == 5
 
 
+def test_packed_attention_skips_mask_for_packed_mask_free_backend():
+    from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
+        MiniMaxH3Attention,
+    )
+
+    class FakeBackend:
+        supports_prefix_kv_slicing = False
+
+        @classmethod
+        def supports_packed_mask_free(cls) -> bool:
+            return True
+
+    class FakeAttention(torch.nn.Module):
+        attn_backend = FakeBackend
+
+        def __init__(self):
+            super().__init__()
+            self.metadata = None
+
+        def forward(self, query, key, value, metadata):
+            self.metadata = metadata
+            return query
+
+    attention = object.__new__(MiniMaxH3Attention)
+    torch.nn.Module.__init__(attention)
+    attention.attention = FakeAttention()
+    q = torch.randn(8, 2, 4)
+
+    attention._run_packed_attention(
+        q,
+        q,
+        q,
+        cu_seqlens=torch.tensor([0, 5, 8], dtype=torch.int32),
+        max_seqlen=5,
+        packed_total=8,
+    )
+
+    assert attention.attention.metadata.attn_mask is None
+    assert attention.attention.metadata.extra["valid_kv_length"] == 5
+    assert attention.attention.metadata.extra["npu_attn_varlen"] is True
+
+
 def test_packed_attention_keeps_padding_mask_for_other_backends():
     from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
         MiniMaxH3Attention,
@@ -534,6 +580,10 @@ def test_packed_attention_keeps_padding_mask_for_other_backends():
     class FakeBackend:
         supports_prefix_kv_slicing = False
 
+        @classmethod
+        def supports_packed_mask_free(cls) -> bool:
+            return False
+
     class FakeAttention(torch.nn.Module):
         attn_backend = FakeBackend
 

tests/diffusion/models/minimax_h3/test_minimax_h3_parallel.py

@@ -69,6 +69,7 @@ def test_h3_fused_rope_matches_reference_and_preserves_unrotated_dims():
 
     attention = object.__new__(MiniMaxH3Attention)
     nn.Module.__init__(attention)
+    attention.rot_dim = 96
     attention.rope = RotaryEmbedding(is_neox_style=True, half_head_dim=False)
     attention.rope._forward_method = attention.rope.forward_native
 

tests/docs/test_generate_examples.py

@@ -3,6 +3,7 @@
 
 import importlib.util
 from pathlib import Path
+from types import SimpleNamespace
 
 import pytest
 
@@ -43,3 +44,11 @@ def test_model_display_names_require_string_mapping(tmp_path, monkeypatch):
 
     with pytest.raises(ValueError, match="MODEL_DISPLAY_NAMES_FILE"):
         generate_examples.load_model_display_names()
+
+
+def test_only_general_serving_examples_get_generated_pages():
+    general_example = SimpleNamespace(category="offline_inference", path=Path("text_to_image"))
+    model_example = SimpleNamespace(category="offline_inference", path=Path("qwen3_omni"))
+
+    assert generate_examples.is_general_example(general_example)
+    assert not generate_examples.is_general_example(model_example)

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 "stage engine died" in fatal.error
-    assert fatal.stage_id == 2
 
 
 @dataclass

tests/engine/test_orchestrator_error_handling.py

@@ -3,7 +3,7 @@
 """Tests for error propagation paths within the Orchestrator.
 
 Covers:
-- EngineDeadError from an LLM stage poll → fatal error broadcast + 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:
+async def test_engine_dead_error_evicts_replica_and_keeps_running(orchestrator_factory) -> None:
-    """When a stage raises EngineDeadError during poll, the orchestrator must:
+    """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.
+        # The dead replica is evicted; the orchestrator keeps running.
-        orchestrator_fixture.thread.join(timeout=5)
+        pool = orchestrator_fixture.orchestrator.stage_pools[0]
-        assert not orchestrator_fixture.thread.is_alive()
+        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)]
+    base.engine.stage_pools = [_FakeStagePool([SimpleNamespace(_engine_dead=True)])]
-    assert base.is_running is False
+    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))]
+    base.engine.stage_pools = [_FakeStagePool([SimpleNamespace(resources=SimpleNamespace(engine_dead=True))])]
-    assert base.is_running is False
+    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]
+    base.engine.stage_pools = [_FakeStagePool([dead_stage])]
-    assert base.errored is True
+    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]
+    base.engine.stage_pools = [_FakeStagePool([dead_stage])]
-    assert base.errored is True
+    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/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:
+    def _predict_velocity(self, x_t: torch.Tensor, llm_proj: torch.Tensor, t_proj: torch.Tensor) -> torch.Tensor:
-        combined = torch.cat([x_t, llm_output, t_emb], dim=-1)
+        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_proj_batched = t_proj.repeat(2, 1)
-            t_emb_batched = t_emb.repeat(2, 1)
+            v_all = at._predict_velocity(x_t=x_batched, llm_proj=llm_proj_batched, t_proj=t_proj_batched)
-            v_all = at._predict_velocity(x_t=x_batched, llm_output=llm_batched, t_emb=t_emb_batched)
             v_t, uncond_v_t = v_all[:B], v_all[B:]
             cfg_alpha = 1.2
             v_t = cfg_alpha * v_t + (1 - cfg_alpha) * uncond_v_t

tools/pre_commit/check_tts_adapter.py

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

vllm_omni/diffusion/attention/backends/abstract.py

@@ -20,6 +20,17 @@ class AttentionBackend(ABC):
     # avoid a slower masked-attention plan when tail padding is not semantic.
     supports_prefix_kv_slicing: bool = False
 
+    @classmethod
+    def supports_packed_mask_free(cls) -> bool:
+        """Whether packed attention never reads attn_mask on this platform.
+
+        When True, models that pack a [real, pad] two-document layout and
+        carry cu_seqlens/max_seqlen in ``AttentionMetadata.extra`` may skip
+        constructing the padding mask entirely. Backends whose mask-free
+        behavior is platform-dependent must check current_omni_platform.
+        """
+        return False
+
     # ``OmniPlatformEnum`` values this backend runs on; None means unrestricted.
     # Platform resolution rejects an explicit selection outside this set, so a
     # hardware-specific backend fails before the model is built.
@@ -121,6 +132,17 @@ class AttentionMetadata:
     #     the packed cu_seqlens tensors.
     #   "valid_kv_length": int — contiguous valid K/V prefix length for a
     #     backend that advertises supports_prefix_kv_slicing.
+    #   "npu_attn_varlen": bool — model opt-in for the NPU packed varlen path
+    #     (TND npu_fusion_attention driven by cu_seqlens, mask never read).
+    #     Requires the [real, pad] two-document packing contract; see
+    #     FlashAttentionImpl._forward_varlen_packed_npu.
+    #   "laser_input_scale": float — model opt-in input pre-scale for the NPU
+    #     ascend_laser_attention path. The kernel stores unscaled QK^T in an
+    #     fp16 workspace, so outlier activations overflow 65504 into NaN rows;
+    #     with this set (>1), q/k/v are divided by the factor before the op,
+    #     the kernel scale_value is multiplied by its square, and the output
+    #     is scaled back (exact for power-of-two factors). Absent means no
+    #     pre-scaling. See FlashAttentionImpl._forward_prefix_kv_slice_npu.
 
     # Piecewise attention metadata (mixed causal/full masks).
     # full_attn_spans: per-sample [start, end) spans in global coordinates using full attention.

vllm_omni/diffusion/attention/backends/flash_attn.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import os
 from functools import partial
 
 import torch
@@ -16,6 +17,7 @@ from vllm_omni.diffusion.attention.backends.utils.piecewise_attn import (
     piecewise_attn,
 )
 from vllm_omni.diffusion.config import get_current_diffusion_config_or_none
+from vllm_omni.platforms import current_omni_platform
 
 logger = init_logger(__name__)
 
@@ -24,6 +26,15 @@ class FlashAttentionBackend(AttentionBackend):
     accept_output_buffer: bool = True
     supports_piecewise_spans: bool = True
 
+    @classmethod
+    def supports_packed_mask_free(cls) -> bool:
+        # CUDA: forward_cuda dispatches the packed cu_seqlens varlen path
+        # before ever reading attn_mask. NPU: forward_fa_npu honors the
+        # npu_attn_varlen opt-in and rebuilds the padding mask itself if the
+        # packed contract fails. XPU reads attn_mask, so models must keep
+        # constructing it there.
+        return current_omni_platform.is_cuda() or current_omni_platform.is_npu()
+
     @classmethod
     def supports_attention_mask(cls) -> bool:
         return True
@@ -415,7 +426,36 @@ class FlashAttentionImpl(AttentionImpl):
                 "For installation details, see https://gitcode.com/Ascend/MindIE-SD"
                 "Otherwise, use SDPA backend by setting DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA"
             )
+        # Opt-in mask-free paths (mirror the CUDA cu_seqlens behavior): the
+        # model marks extra["npu_attn_varlen"] and carries packed metadata, so
+        # the padding document is excluded without reading or materializing
+        # the quadratic full_qk mask.
+        #   - default: packed varlen via mindiesd attention_forward_varlen;
+        #   - when MINDIE_SD_FA_TYPE=ascend_laser_attention: prefix K/V
+        #     slicing via mindiesd attention_forward (so the env-selected op
+        #     type is honored; the TND varlen op cannot be).
+        extra = attn_metadata.extra if attn_metadata else {}
+        if extra.get("npu_attn_varlen", False):
+            if os.environ.get("MINDIE_SD_FA_TYPE") == "ascend_laser_attention":
+                out = self._forward_prefix_kv_slice_npu(query, key, value, extra)
+            else:
+                out = self._forward_varlen_packed_npu(query, key, value, extra)
+            if out is not None:
+                return out
         attention_mask = attn_metadata.attn_mask if attn_metadata else None
+        if attention_mask is None and extra.get("npu_attn_varlen", False):
+            # Models skip mask construction when this opt-in is set
+            # (FlashAttentionBackend.supports_packed_mask_free). The packed
+            # paths above declined (contract mismatch), so rebuild the padding
+            # mask here to keep the masked fallback correct.
+            used = extra.get("valid_kv_length")
+            if not isinstance(used, int) or not 0 < used <= query.shape[1]:
+                raise ValueError(
+                    "npu_attn_varlen packed metadata is unusable and no attn_mask "
+                    f"was constructed (valid_kv_length={used!r}, seq_len={query.shape[1]}); "
+                    "refusing to run unmasked attention over padding rows."
+                )
+            attention_mask = torch.arange(query.shape[1], device=query.device)[None] < used
 
         # NPU aclnnFlashAttentionScore requires mask shape to be one of:
         # [B, N, Sq, Skv], [B, 1, Sq, Skv], [1, 1, Sq, Skv], or [Sq, Skv]
@@ -433,3 +473,161 @@ class FlashAttentionImpl(AttentionImpl):
             op_type="fused_attn_score",
             layout=layout,
         )
+
+    def _resolve_packed_seq_npu(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        extra: dict,
+    ) -> tuple[list[int], list[int]] | None:
+        """Resolve packed document boundaries (cumulative end offsets per doc).
+
+        Returns (seq_q, seq_k), e.g. ([used_q, total_q], [used_k, total_k]),
+        or None unless the metadata matches the exact contract MiniMax-H3
+        produces:
+          - q/k/v are [1, T, N, D] (BSND, single packed batch);
+          - cu_seqlens describe a [real, pad] two-document packing, with the
+            padding document as a strict suffix;
+          - max_seqlen_q/k are Python ints equal to the real document length
+            (so boundaries are derived without any device sync).
+        """
+        if self.causal:
+            return None
+        packed_keys = ("cu_seqlens_q", "cu_seqlens_k", "max_seqlen_q", "max_seqlen_k")
+        if any(k not in extra for k in packed_keys):
+            return None
+        cu_q, cu_k = extra["cu_seqlens_q"], extra["cu_seqlens_k"]
+        used_q, used_k = extra["max_seqlen_q"], extra["max_seqlen_k"]
+        if query.shape[0] != 1 or key.shape[0] != 1:
+            return None
+        if not isinstance(used_q, int) or not isinstance(used_k, int):
+            return None
+        total_q, total_k = query.shape[1], key.shape[1]
+        # .shape is host-side metadata: counting documents never syncs.
+        if cu_q.shape[0] != cu_k.shape[0] or cu_q.shape[0] > 3:
+            return None
+        if not (0 < used_q <= total_q) or not (0 < used_k <= total_k):
+            return None
+        if used_q == total_q and used_k == total_k:
+            # No padding document: one full-length document.
+            return [total_q], [total_k]
+        if cu_q.shape[0] == 3 and used_q >= total_q - used_q and used_k >= total_k - used_k:
+            # [real, pad] packing: the real document must be the longer one
+            # (consistent with the max_seqlen naming).
+            return [used_q, total_q], [used_k, total_k]
+        return None
+
+    def _forward_varlen_packed_npu(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        extra: dict,
+    ) -> torch.Tensor | None:
+        """Packed varlen attention on NPU via mindiesd attention_forward_varlen.
+
+        Returns None (caller falls back to the mask path) when the packed
+        contract does not hold; see _resolve_packed_seq_npu.
+        """
+        resolved = self._resolve_packed_seq_npu(query, key, extra)
+        if resolved is None:
+            return None
+        seq_q, seq_k = resolved
+
+        try:
+            from mindiesd import attention_forward_varlen
+        except ImportError:
+            raise ImportError(
+                "FlashAttentionBackend NPU implementation requires MindIE-SD. "
+                "Please install MindIE-SD to enable NPU attention support. "
+                "For installation details, see https://gitcode.com/Ascend/MindIE-SD"
+                "Otherwise, use SDPA backend by setting DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA"
+            )
+        q = query.squeeze(0)  # [T, N, D] == TND
+        k = key.squeeze(0)
+        v = value.squeeze(0)
+        # attention_forward_varlen wants cu_seqlens as host lists of cumulative
+        # offsets and forwards cu[1:] as actual_seq_qlen/actual_seq_kvlen.
+        out = attention_forward_varlen(
+            q,
+            k,
+            v,
+            [0, *seq_q],
+            [0, *seq_k],
+            softmax_scale=self.softmax_scale,
+        )
+        return out.unsqueeze(0)
+
+    def _forward_prefix_kv_slice_npu(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        extra: dict,
+    ) -> torch.Tensor | None:
+        """Mask-free attention by slicing K/V to the valid prefix (zero-copy
+        views), then mindiesd attention_forward without a mask.
+
+        Same numerical contract as the varlen path: the padding document is a
+        strict suffix, so dropping it from K/V is identical to masking it out.
+        Query keeps full length; outputs on padding rows are never consumed
+        downstream. Returns None when the packed contract does not hold.
+
+        Used when MINDIE_SD_FA_TYPE=ascend_laser_attention: the TND varlen op
+        cannot honor the env-selected op type, while attention_forward can.
+
+        The laser kernel stores unscaled S=QK^T in an fp16 GM workspace, so
+        bf16 activations with large outliers overflow 65504 into ±inf and the
+        affected rows turn NaN. Models may opt into exact power-of-two input
+        pre-scaling via extra["laser_input_scale"] (see abstract.py): q/k/v
+        are divided by the factor, the kernel scale_value is multiplied by its
+        square, and the output is scaled back. Absent or invalid factor means
+        no pre-scaling.
+        """
+        resolved = self._resolve_packed_seq_npu(query, key, extra)
+        if resolved is None:
+            return None
+        _, seq_k = resolved
+        used_k = seq_k[0]  # real document length (first cumulative end)
+
+        try:
+            from mindiesd import attention_forward
+        except ImportError:
+            raise ImportError(
+                "FlashAttentionBackend NPU implementation requires MindIE-SD. "
+                "Please install MindIE-SD to enable NPU attention support. "
+                "For installation details, see https://gitcode.com/Ascend/MindIE-SD"
+                "Otherwise, use SDPA backend by setting DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA"
+            )
+        # mindiesd always takes BSND input regardless of `layout`; the arg
+        # selects the op-internal layout and laser only supports BNSD, so do
+        # NOT forward the model's qkv_layout ("BSND" for MiniMax-H3) here.
+        layout = "BNSD"
+        key = key[:, :used_k]
+        value = value[:, :used_k]
+
+        input_scale = extra.get("laser_input_scale")
+        preserve_input_range = isinstance(input_scale, (int, float)) and input_scale > 1
+        if preserve_input_range:
+            query = query / input_scale
+            key = key / input_scale
+            value = value / input_scale
+        scale = self.softmax_scale
+        if preserve_input_range:
+            scale *= input_scale**2
+
+        def _postprocess(out: torch.Tensor) -> torch.Tensor:
+            return out * input_scale if preserve_input_range else out
+
+        return _postprocess(
+            attention_forward(
+                query,
+                key,
+                value,
+                attn_mask=None,
+                opt_mode="manual",
+                op_type="fused_attn_score",  # MINDIE_SD_FA_TYPE env overrides this
+                layout=layout,
+                scale=scale,
+            )
+        )

vllm_omni/diffusion/attention/backends/rainfusion_attn.py

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

vllm_omni/diffusion/layers/norm.py

@@ -68,6 +68,7 @@ class RMSNorm(CustomOp):
     def __init__(self, hidden_size: int, eps: float = 1e-6, dtype: torch.dtype = torch.float32) -> None:
         super().__init__()
         self.weight = nn.Parameter(torch.ones(hidden_size, dtype=dtype))
+        self.hidden_size = hidden_size
         self.variance_epsilon = eps
 
     def _forward_fused(self, x: torch.Tensor) -> torch.Tensor:
@@ -107,6 +108,19 @@ class RMSNorm(CustomOp):
         except Exception:
             return self.forward_native(x)
 
+    def forward_musa(
+        self,
+        x: torch.Tensor,
+    ) -> torch.Tensor:
+        # Preserve the aten::rms_norm graph so dynamic Inductor can fuse the
+        # H3 Q/K norm with its inline RoPE path on MUSA.
+        return F.rms_norm(
+            x,
+            (self.hidden_size,),
+            self.weight,
+            self.variance_epsilon,
+        )
+
     def forward_npu(
         self,
         x: torch.Tensor,

vllm_omni/diffusion/layers/rope.py

@@ -193,7 +193,16 @@ class RotaryEmbedding(CustomOp):
         cos: torch.Tensor,
         sin: torch.Tensor,
     ) -> torch.Tensor:
-        return self.forward_native(x, cos, sin)
+        # Keep the full-dimension NeoX form inline for MUSA compilation.
+        # Other layouts use the shared native implementation.
+        if self.half_head_dim or self.interleaved:
+            return self.forward_native(x, cos, sin)
+        if cos.dim() == 3:
+            cos, sin = cos[0], sin[0]
+        cos = cos.unsqueeze(-2)
+        sin = sin.unsqueeze(-2)
+        x1, x2 = x.chunk(2, dim=-1)
+        return x * cos + torch.cat((-x2, x1), dim=-1) * sin
 
     def forward_native(
         self,

vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py

@@ -107,6 +107,15 @@ MINIMAX_H3_FP32_BUFFER_NAMES = frozenset({"rope.inv_freq"})
 # lookup and masked out afterwards).
 MINIMAX_H3_ADALN_MODALITY_NUM = 3
 
+# Opt-in fp16-range protection for the NPU ascend_laser_attention kernel
+# (consumed only via the "laser_input_scale" extra key; other backends and
+# platforms ignore it). The kernel stores unscaled QK^T in an fp16 GM
+# workspace, and H3's outlier activations (per-element amax in the hundreds)
+# push dot products past fp16 max 65504, turning whole 128-row blocks NaN.
+# 256 is a power of two, so pre-dividing q/k/v and the compensating
+# kernel-scale/output multiplies are exact in floating point.
+MINIMAX_H3_LASER_INPUT_SCALE = 256.0
+
 
 def _required_kwarg(kwargs: dict[str, Any], key: str) -> Any:
     if key not in kwargs or kwargs[key] is None:
@@ -333,6 +342,7 @@ class MiniMaxH3Attention(nn.Module):
         )
         self.num_heads = self.qkv_proj.num_heads
         self.num_kv_heads = self.qkv_proj.num_kv_heads
+        self.rot_dim = 6 * arch.rope_inv_freq_len
         self.q_norm = _norm(arch.attention_head_dim, eps=arch.qk_norm_eps)
         self.k_norm = _norm(arch.attention_head_dim, eps=arch.qk_norm_eps)
         self.rope = RotaryEmbedding(is_neox_style=True, half_head_dim=False)
@@ -365,7 +375,7 @@ class MiniMaxH3Attention(nn.Module):
         x: [T, heads, head_dim]; freqs: [T, rot_dim]. In the unfused path, cos/sin
         are cast to the activation dtype before the elementwise math.
         """
-        rot_dim = freqs.shape[-1]
+        rot_dim = self.rot_dim
         x_rot, x_pass = x[..., :rot_dim], x[..., rot_dim:]
         cos = torch.cos(freqs).to(x.dtype)  # [T, rot_dim]
         sin = torch.sin(freqs).to(x.dtype)
@@ -401,11 +411,16 @@ class MiniMaxH3Attention(nn.Module):
         used = min(max_seqlen, packed_total)
         attn_mask = None
         # Ring attention can dispatch to a different implementation from the
-        # configured backend, so this no-mask fast path is local-only.
+        # configured backend, so the no-mask fast paths are local-only.
-        prefix_slice = (
+        # supports_prefix_kv_slicing: backend slices K/V itself (cuDNN).
-            not getattr(self.attention, "use_ring", False) and self.attention.attn_backend.supports_prefix_kv_slicing
+        # supports_packed_mask_free: backend consumes the packed metadata
+        # without ever reading attn_mask (CUDA packed varlen, NPU
+        # npu_attn_varlen opt-in with its own fallback rebuild).
+        no_mask = not getattr(self.attention, "use_ring", False) and (
+            self.attention.attn_backend.supports_prefix_kv_slicing
+            or self.attention.attn_backend.supports_packed_mask_free()
         )
-        if used < packed_total and not prefix_slice:
+        if used < packed_total and not no_mask:
             attn_mask = torch.arange(packed_total, device=q.device)[None] < used
         metadata = AttentionMetadata(
             attn_mask=attn_mask,
@@ -415,6 +430,15 @@ class MiniMaxH3Attention(nn.Module):
                 "max_seqlen_q": max_seqlen,
                 "max_seqlen_k": max_seqlen,
                 "valid_kv_length": used,
+                # Opt the NPU flash backend into the packed varlen path so the
+                # quadratic full_qk mask is never materialized. Ring attention
+                # is excluded: it keeps the aligned padding rows for its
+                # fixed-size P2P buffers and still needs the mask.
+                "npu_attn_varlen": not getattr(self.attention, "use_ring", False),
+                # fp16-range protection for the ascend_laser_attention kernel
+                # (see MINIMAX_H3_LASER_INPUT_SCALE). Ignored by every other
+                # backend/path.
+                "laser_input_scale": MINIMAX_H3_LASER_INPUT_SCALE,
             },
             video_layout=video_layout,
         )

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
+            # The orchestration loop isolates per-replica EngineDeadError
-            # engine died.  All pending requests were already notified and
+            # (evict + continue), so this only fires if one escapes outside the
-            # _shutdown_event was already set by the loop's handler.
+            # poll handler (e.g. from output post-processing). Treat it as a
-            # During teardown this is expected; the finally block handles
+            # teardown signal and let the finally block clean up; do not re-raise.
-            # proper cleanup.  Do not re-raise.
+            logger.info("[Orchestrator] Engine dead, shutting down: %s", e)
-            logger.info("[Orchestrator] Engine dead during shutdown: %s", e)
+            # Unblock any collective_rpc callers waiting on rpc_async_queue
-            if self._fatal_error is None:
+            # instead of leaving them to hit their own timeout.
-                self._fatal_error = str(e) or "Stage engine died"
             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(
+        if not await self._dispatch_or_fail_request(
-            request_id,
+            lambda: self.stage_pools[stage_id].submit_initial(
-            req_state,
+                request_id,
-            prompt,
+                req_state,
-            prompt_text=msg.output_prompt_text,
+                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(
+        if not await self._dispatch_or_fail_request(
-            request_id,
+            lambda: self.stage_pools[stage_id].submit_update(
-            req_state,
+                request_id,
-            request,
+                req_state,
-            prompt_text=msg.output_prompt_text,
+                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(
+        # A dead-stage failure is attributed to the parent request; its cleanup
-            companion_id,
+        # also releases this companion.
-            companion_state,
+        if not await self._dispatch_or_fail_request(
-            companion_prompt,
+            lambda: self.stage_pools[0].submit_initial(
-            prompt_text=msg.companion_prompt_text,
+                companion_id,
-            affinity_request_id=parent_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":
+                    # Shared catch so a dead replica on either poll path is
-                        output = pool.poll_diffusion_output(replica_id)
+                    # evicted rather than tearing down every stage (#4285).
-                        if output is None:
+                    try:
-                            continue
+                        if pool.stage_type == "diffusion":
+                            diffusion_output = pool.poll_diffusion_output(replica_id)
+                            if diffusion_output is None:
+                                continue
 
-                        pool.record_output_timestamps([output])
+                            pool.record_output_timestamps([diffusion_output])
-                        await self._handle_processed_outputs(stage_id, replica_id, [output])
+                            processed = [diffusion_output]
-                        idle = False
+                        else:
-                    else:
-                        try:
                             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:
+                    except asyncio.CancelledError:
-                            raise
+                        raise
-                        except EngineDeadError as e:
+                    except EngineDeadError as e:
-                            logger.error(
+                        await self._handle_dead_replica(stage_id, replica_id, e)
-                                "[Orchestrator] Stage-%s replica-%s is dead: %s",
+                        continue
-                                stage_id,
+                    except Exception:
-                                replica_id,
+                        if self._shutdown_event.is_set():
-                                e,
+                            return
-                            )
+                        logger.exception(
-                            affected_request_ids = pool.mark_replica_unavailable(replica_id)
+                            "[Orchestrator] Stage-%s replica-%s processing failed",
-                            closed_sessions = (
+                            stage_id,
-                                self.duplex_control_plane.close_sessions_for_request_ids(
+                            replica_id,
-                                    affected_request_ids,
+                        )
-                                    abort=False,
+                        raise
-                                )
-                                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
 
-                            self._fatal_error = str(e)
+                    await self._handle_processed_outputs(stage_id, replica_id, processed)
-                            self._fatal_error_stage_id = stage_id
+                    idle = False
-                            for req_id in affected_request_ids:
-                                await self.output_async_queue.put(
-                                    ErrorMessage(
-                                        error=str(e),
-                                        fatal=True,
-                                        request_id=req_id,
-                                        stage_id=stage_id,
-                                    )
-                                )
-                            await self._cleanup_request_ids(
-                                affected_request_ids,
-                                close_duplex_sessions=True,
-                            )
-                            self._shutdown_event.set()
-                            raise
-                        except Exception:
-                            if self._shutdown_event.is_set():
-                                return
-                            logger.exception(
-                                "[Orchestrator] Stage-%s replica-%s processing failed",
-                                stage_id,
-                                replica_id,
-                            )
-                            raise
-
-                        await self._handle_processed_outputs(stage_id, replica_id, raw_output)
-                        idle = False
 
             self._orch_monitor.note_loop(idle=idle)
             if idle:
@@ -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(
+                submit_kwargs = {
-                    request_id,
+                    "kv_sender_info": self._build_kv_sender_info(
-                    req_state,
+                        list(getattr(next_pool.stage_client, "engine_input_source", None) or [next_stage_id - 1]),
-                    req_state.prompt,
+                        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]),
+                submitted = await self._dispatch_or_fail_request(
-                            request_id=request_id,
+                    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(
+                submitted = await self._dispatch_or_fail_request(
-                    request_id,
+                    lambda: next_pool.submit_initial(
-                    req_state,
+                        request_id,
-                    request,
+                        req_state,
-                    prompt_text=None,
+                        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:
+                    if isinstance(msg, ErrorMessage):
-                        req_state = self.request_states.get(msg.request_id)
+                        # Route request-scoped errors to that request's queue and
-                        if req_state is not None:
+                        # keep the loop alive. A request whose stage replica died
-                            await req_state.queue.put(msg)
+                        # and was evicted gets a fatal error delivered here; only
-                        else:
+                        # that request fails (its consumer raises), the server
-                            logger.warning(
+                        # stays up for other stages/requests (#4285). A fatal
-                                "[%s] dropping non-fatal error for unknown req %s",
+                        # error without a request_id is a genuine engine-wide
-                                self._name,
+                        # death and falls through to the except handler below.
-                                msg.request_id,
+                        if msg.request_id is not None:
-                            )
+                            req_state = self.request_states.get(msg.request_id)
-                        continue
+                            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
+        Delegates to ``OmniBase.errored``, which is true only when the
-        thread and all stage clients.  Redeclared here to satisfy 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/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
+        True only when the orchestrator thread is dead. Per-stage liveness is
-        has been marked dead (e.g. diffusion worker OOM / process death).
+        deliberately excluded: the OpenAI serving paths precheck ``errored``
-
+        before request routing, so including it would reject requests that do
-        Checks both ``_engine_dead`` (StageDiffusionClient) and
+        not touch the dead stage (e.g. text-only chat when the talker stage is
-        ``resources.engine_dead`` (StageEngineCoreClient / AsyncMPClient)
+        down). A fully dead stage instead fails only the requests routed
-        since the two client types store the flag differently.
+        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:
+        pools = getattr(self.engine, "stage_pools", None)
-            if hasattr(stage_client, "check_health"):
+        if pools is None:
-                stage_client.check_health()
+            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
+        # Hoist the step-invariant concat + llm_projection out of the loop; inside the
-        hidden_states_zero = torch.zeros_like(hidden_states)
+        # 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 = self.dts[i]
-            dt = timesteps[i + 1] - timesteps[i]
 
-            # Batch conditional + unconditional velocity in a single forward pass
+            # Batch cond + uncond in a single pass; reuse the cached projected-t_emb row.
-            t_emb = at.time_embedding(t.view(-1, 1).repeat(B, 1)).to(hidden_states.dtype)
+            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_proj_batched = t_proj.repeat(2, 1)  # (2B, D)
-            t_emb_batched = t_emb.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())
+        ctx_frames = int(header_view[offset])
-        context_length = int(input_ids[offset + 1].item())
+        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
+        # acoustic_codes starts from x_0; generate directly on device to skip H2D.
-        x_0 = torch.randn(B, self.model_args.n_acoustic_codebook).to(dtype=llm_hidden.dtype, device=llm_hidden.device)
+        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)
+        # Build the schedule constants once per dtype and reuse them every frame.
-        llm_hidden_zero = torch.zeros_like(llm_hidden)
+        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 = dts[i]  # precomputed constant step size
-            dt = timesteps[i + 1] - timesteps[i]
 
-            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_proj_batched = torch.cat([t_proj, t_proj], dim=0)
-            t_emb_batched = torch.cat([t_emb, t_emb], dim=0)
 
             v_all = self._predict_velocity(
                 x_t=x_batched,
-                llm_output=llm_batched,
+                llm_proj=llm_proj_batched,
-                t_emb=t_emb_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
+        llm_proj: torch.Tensor,  # BxD, already through llm_projection
-        t_emb: torch.Tensor,  # BxD
+        t_proj: torch.Tensor,  # BxD, already through time_projection
     ) -> torch.Tensor:
-        x_t = x_t.to(llm_output.dtype)
+        # llm_proj/t_proj arrive pre-projected (hoisted out of the per-step loop by callers).
-
+        x_t = x_t.to(llm_proj.dtype)
-        t_emb = self.time_projection(t_emb)
-        llm_output = self.llm_projection(llm_output)
 
         acoustic_and_semantic_embeddings = [
             self.input_projection(x_t.unsqueeze(1)),  # Bx1xD
-            t_emb.unsqueeze(1),
+            t_proj.unsqueeze(1),
-            llm_output.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(
+        # Cache device-resident 1.0/0.0 scalars to avoid a per-call H2D transfer.
-            audio_codes[:, 0] == AudioSpecialTokens.id(AudioSpecialTokens.end_audio),
+        consts = self._fake_eos_consts.get(audio_codes.device)
-            torch.tensor(1.0, dtype=torch.bfloat16),
+        if consts is None:
-            torch.tensor(0.0, dtype=torch.bfloat16),
+            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 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, device=xq.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
+        # Compute per-request EOA cutting points in a single D2H instead of one sync per request.
-        processed = []
+        flags = []
-        for codes in codes_list:
+        for c in codes_list:
-            eoa_mask = codes[:, 0] == 1
+            eoa_mask = c[:, 0] == 1
-            eoa_indices = eoa_mask.nonzero(as_tuple=False)
+            flags.append(torch.stack([eoa_mask.any().to(torch.long), eoa_mask.long().argmax()]))
-            cutting_point = eoa_indices[0].item() if len(eoa_indices) > 0 else len(codes)
+        resolved = torch.stack(flags).cpu().tolist() if flags else []
-            audio_tokens = codes[:cutting_point] - 2
+        cutting_points = [int(idx) if has else c.shape[0] for (has, idx), c in zip(resolved, codes_list)]
-            processed.append(audio_tokens)
+
+        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)
+        padded = pad_sequence(all_chunks, batch_first=True, padding_value=0.0)
-        K = all_chunks[0].shape[1]
-        padded = torch.zeros(
-            len(all_chunks),
-            max_chunk_len,
-            K,
-            dtype=all_chunks[0].dtype,
-            device=all_chunks[0].device,
-        )
-        for i, chunk in enumerate(all_chunks):
-            padded[i, : len(chunk)] = chunk
 
         audio_codes = padded.to(device=current_omni_platform.device_type)  # [B, T, K]
         audio_values = self.decode(audio_codes.transpose(1, 2), dtype=torch.bfloat16)  # [B, 1, T_out]
         audio_values = audio_values.detach().cpu().float().squeeze(1)  # [B, T_out]
-        if torch.min(audio_values) < -1.0:
+        # Gate min/max reductions behind logger level.
-            logger.warning("Min value of decoded waveform signal is %s", torch.min(audio_values))
+        if logger.isEnabledFor(logging.WARNING):
-        if torch.max(audio_values) > 1.0:
+            amin = float(audio_values.min())
-            logger.warning("Max value of decoded waveform signal is %s", torch.max(audio_values))
+            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:

vllm_omni/platforms/npu/platform.py

@@ -139,6 +139,18 @@ class NPUOmniPlatform(OmniPlatform, NPUPlatform):
                 )
                 backend_upper = "FLASH_ATTN"
 
+            if backend_upper == "FLASH_ATTN" and find_spec("mindiesd"):
+                # The NPU FLASH_ATTN backend imports mindiesd lazily at first
+                # forward, but CANN snapshots the custom-op registry at the
+                # first custom-op regInfo lookup in the process (e.g. a
+                # vllm-ascend custom op during model load/warmup). Import
+                # mindiesd here so its env.py prepends the mindiesd vendor
+                # dirs (aie_ascendc etc.) to ASCEND_CUSTOM_OPP_PATH before
+                # that snapshot; otherwise aclnnLaserAttention /
+                # FusedAttentionScore fail with EZ1001 "does not support
+                # opType" for the rest of the process.
+                import mindiesd  # noqa: F401
+
             backend = DiffusionAttentionBackendEnum[backend_upper]
             logger.debug("Using diffusion attention backend '%s'", backend_upper)
             return backend.get_path()