docs: add DLO DP behavior animation

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

.buildkite/cuda/test-nightly.yml

@@ -460,6 +460,21 @@ steps:
             exit $$EXIT2
         mirror_hardwares: h100_2
 
+      - label: ":full_moon: Diffusion X2V · Perf Test · MiniMax-H3"
+        key: nightly-diffusion-x2v-performance-minimax-h3
+        timeout_in_minutes: 180
+        commands:
+          - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
+          - export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
+          - |
+            set +e
+            pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_minimax_h3_vllm_omni.json
+            EXIT=$$?
+            buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
+            buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
+            exit $$EXIT
+        mirror_hardwares: h100_4
+
   - group: ":card_index_dividers: Diffusion Test"
     key: nightly-diffusion-test-group
     depends_on: upload-nightly-pipeline

.buildkite/cuda/test-ready.yml

@@ -47,15 +47,17 @@ steps:
           - timeout 15m pytest -sv tests/diffusion/batching -m 'core_model and cuda' --run-level "core_model"
         mirror_hardwares: l4_1
 
-      - label: "Diffusion · Distributed Test"
+      - label: "Diffusion · Distributed & Attention Test"
         source_file_dependencies:
           - tests/diffusion/distributed/
+          - tests/diffusion/attention/test_attention_sp.py
           - vllm_omni/diffusion/distributed/
-          - vllm_omni/diffusion/attention/parallel/
+          - vllm_omni/diffusion/attention/
           - vllm_omni/diffusion/data.py
           - vllm_omni/diffusion/models/
         commands:
           - timeout 20m pytest -sv tests/diffusion/distributed -m 'core_model and cuda' --run-level "core_model"
+          - timeout 20m pytest -sv tests/diffusion/attention/test_attention_sp.py -m 'core_model and cuda' --run-level "core_model"
         mirror_hardwares: l4_4
 
       - label: "Diffusion · Model Test"

.claude/skills/add-diffusion-model/SKILL.md

@@ -309,7 +309,7 @@ For Omni or custom models, create:
 
 Required updates:
 1. `docs/user_guide/diffusion/parallelism/overview.md` — parallelism support overview/table
-2. `docs/user_guide/diffusion/cpu_offload_diffusion.md` — if CPU offload supported (add to supported models table)
+2. `docs/user_guide/diffusion/cpu_offload.md` — if CPU offload supported (add to supported models table)
 3. `docs/user_guide/diffusion/cache_acceleration/teacache.md` — if TeaCache supported
 4. `docs/user_guide/diffusion/cache_acceleration/cache_dit.md` — if Cache-DiT supported
 5. Offline example docs under `examples/offline_inference/<name>/` (`README.md` or category-specific `.md`)

benchmarks/diffusion/backends.py

@@ -1,5 +1,6 @@
 import asyncio
 import base64
+import json
 import mimetypes
 import os
 import time
@@ -13,6 +14,17 @@ from tqdm import tqdm
 DEFAULT_EDITS_BOT_TASK = "think"
 
 
+def _is_minimax_h3_mixed_reference(model: str) -> bool:
+    """MiniMax-H3 Ref2VA consumes reference videos as file paths.
+
+    Its reference preparer ffprobes / ffmpegs the reference file, so single
+    ``video_reference`` uploads are sent as a base64 data URL: the server
+    persists them to disk and keeps ``source_path``, handing the pipeline a
+    real file path instead of a decoded frame list.
+    """
+    return "MiniMax-H3" in (model or "")
+
+
 @dataclass
 class RequestFuncInput:
     prompt: str
@@ -355,13 +367,21 @@ async def async_request_v1_videos(
         )
     elif input.video_paths and len(input.video_paths) > 0:
         video_path = input.video_paths[0]
-        image_file = open(video_path, "rb")
-        form.add_field(
-            "input_reference",
-            image_file,
-            filename=os.path.basename(video_path),
-            content_type=_guess_mime_type(video_path),
-        )
+        if _is_minimax_h3_mixed_reference(input.model):
+            with open(video_path, "rb") as reference_file:
+                video_b64 = base64.b64encode(reference_file.read()).decode("ascii")
+            form.add_field(
+                "video_reference",
+                json.dumps([{"video_url": f"data:video/mp4;base64,{video_b64}"}]),
+            )
+        else:
+            image_file = open(video_path, "rb")
+            form.add_field(
+                "input_reference",
+                image_file,
+                filename=os.path.basename(video_path),
+                content_type=_guess_mime_type(video_path),
+            )
 
     job_id = None
     job_status = None

benchmarks/diffusion/diffusion_benchmark_serving.py

@@ -858,7 +858,12 @@ class RandomDataset(BaseDataset):
             height = int(self.args.height or profile.get("height") or 720)
             num_frames = int(self.args.num_frames or profile.get("num_frames") or 16)
             fps = float(self.args.fps or profile.get("fps") or 8)
-            writer = cv2.VideoWriter(video_path, cv2.VideoWriter_fourcc(*"mp4v"), fps, (width, height))
+            # cv2 typically writes MPEG-4 Part 2 (``mp4v``), which pipelines such
+            # as MiniMax-H3 Ref2VA reject because they require H.264/H.265
+            # reference videos. Write the synthetic clip first, then re-encode
+            # it to H.264 when ffmpeg is available; otherwise keep ``mp4v``.
+            raw_path = video_path + ".raw.mp4"
+            writer = cv2.VideoWriter(raw_path, cv2.VideoWriter_fourcc(*"mp4v"), fps, (width, height))
             if not writer.isOpened():
                 raise RuntimeError("cv2.VideoWriter failed to open")
             for frame_idx in range(num_frames):
@@ -871,6 +876,10 @@ class RandomDataset(BaseDataset):
                 frame[max(y - 36, 0) : min(y + 36, height), x : min(x + 96, width), :] = (64, 128, 220)
                 writer.write(frame)
             writer.release()
+            if _reencode_video_to_h264(raw_path, video_path):
+                os.unlink(raw_path)
+            else:
+                os.replace(raw_path, video_path)
             return [video_path]
         except Exception as e:
             raise RuntimeError(
@@ -878,6 +887,40 @@ class RandomDataset(BaseDataset):
             ) from e
 
 
+def _reencode_video_to_h264(source: str, destination: str) -> bool:
+    """Re-encode a video to H.264/MP4 using ffmpeg when available.
+
+    Returns False when ffmpeg is missing or the re-encode fails so callers can
+    fall back to the original container/codec.
+    """
+    import shutil
+    import subprocess
+
+    if shutil.which("ffmpeg") is None:
+        return False
+    try:
+        result = subprocess.run(
+            [
+                "ffmpeg",
+                "-y",
+                "-loglevel",
+                "error",
+                "-i",
+                source,
+                "-c:v",
+                "libx264",
+                "-pix_fmt",
+                "yuv420p",
+                destination,
+            ],
+            capture_output=True,
+            text=True,
+        )
+        return result.returncode == 0 and os.path.exists(destination)
+    except Exception:
+        return False
+
+
 def _compute_expected_latency_ms_from_base(req: RequestFuncInput, args, base_time_ms: float | None) -> float | None:
     """Compute expected execution time (ms) based on a base per-step-per-frame unit time.
 

docs/.nav.yml

@@ -68,12 +68,12 @@ nav:
         - Sequence Parallel: user_guide/diffusion/parallelism/sequence_parallel.md
         - Tensor Parallel: user_guide/diffusion/parallelism/tensor_parallel.md
         - VAE Parallelism: user_guide/diffusion/parallelism/vae_parallelism.md
-      - CPU Offloading: user_guide/diffusion/cpu_offload_diffusion.md
+      - CPU Offloading: user_guide/diffusion/cpu_offload.md
       - LoRA: user_guide/diffusion/lora.md
       - Custom Pipeline: features/custom_pipeline.md
       - Execution Modes: user_guide/diffusion/execution_modes.md
     - Quantization:
-      - Overview: user_guide/quantization/overview.md
+      - Quantization: user_guide/quantization/overview.md
       - Online Quantization: user_guide/quantization/online.md
       - Quantized KV Cache: user_guide/quantization/quantized_kvcache.md
       - FP8 W8A8: user_guide/quantization/fp8.md
@@ -106,7 +106,6 @@ nav:
     - Feature Design:
       - Runtime and Stage Execution:
         - Disaggregated Inference: design/feature/disaggregated_inference.md
-        - Ray-based Execution: design/feature/ray_based_execution.md
         - Async Chunk: design/feature/async_chunk.md
         - Async Diffusion Output: design/feature/async_diffusion_output.md
         - Async Omni Output Materialization: design/feature/omni_async_output_materialization.md
@@ -130,20 +129,37 @@ nav:
           - VAE Patch Parallelism: design/feature/vae_parallel.md
         - Attention Backends:
           - Skip-Softmax: design/feature/skip_softmax.md
-        - Quantization: user_guide/quantization/overview.md
+        - Quantization: design/feature/quantization.md
         - Cache-DiT: design/feature/cache_dit.md
         - TeaCache: design/feature/teacache.md
         - Diffusion Continuous Batching: design/feature/diffusion_continuous_batching.md
+        - Distributed Layerwise Offload: design/feature/distributed_layerwise_offload.md
     - Infrastructure and Performance:
       - Prometheus Metrics: design/metrics.md
       - Speech Generation Performance Optimizations: design/qwen3_omni_tts_performance_optimization.md
     - Module Design:
-      - Runtime Modules:
-        - AR Module: design/module/ar_module.md
-        - DIT Module: design/module/dit_module.md
-      - Orchestration:
-        - Entrypoint Module: design/module/entrypoint_module.md
-        - AsyncOmni Architecture: design/module/async_omni_architecture.md
+      - Entrypoints: design/module/entrypoints.md
+      - vLLM-Omni Configuration: design/module/vllm_omni_config.md
+      - Input, Output, and Modality Contracts: design/module/input_output_modality_contracts.md
+      - Error Contracts: design/module/error_contracts.md
+      - Engine Orchestration: design/module/engine_orchestration.md
+      - Stage Runtime: design/module/stage_runtime.md
+      - OmniConnector: design/module/omni_connector.md
+      - Model Integration: design/module/model_integration.md
+      - Autoregressive Runtime: design/module/ar_runtime.md
+      - Diffusion:
+        - Overview: design/module/diffusion/index.md
+        - Runtime: design/module/diffusion/diffusion_runtime.md
+        - Model Integration: design/module/diffusion/diffusion_model_integration.md
+        - Continuous Batching: design/module/diffusion/continuous_batching.md
+        - Parallelism: design/module/diffusion/parallelism.md
+        - Offloader: design/module/diffusion/offloader.md
+      - Execution Platforms: design/module/execution_platforms.md
+      - Cache Management: design/module/cache_management.md
+      - Quantization: design/module/quantization.md
+      - Observability: design/module/observability.md
+      - Profiling: design/module/profiling.md
+      - Benchmarking: design/module/benchmarking.md
   - Docs Guide: contributing/DOCS_GUIDE.md
 - API Reference:
   - api/README.md

docs/contributing/ci/test_examples/l4_functionality_tests.inc.md

@@ -1,6 +1,6 @@
 **Scope**
 
-For diffusion models, the L4 functionality test covers all or common *diffusion features* that are supported by this model, including several [parallelism acceleration methods](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/parallelism_acceleration/), [CPU offloading](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/cpu_offload_diffusion/), [TeaCache](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/teacache/) and [Cache-DiT](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/cache_dit_acceleration/) cache backends, [quantization methods](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/quantization/overview/).
+For diffusion models, the L4 functionality test covers all or common *diffusion features* that are supported by this model, including several [parallelism acceleration methods](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/parallelism_acceleration/), [CPU offloading](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/cpu_offload/), [TeaCache](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/teacache/) and [Cache-DiT](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/cache_dit_acceleration/) cache backends, [quantization methods](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/quantization/overview/).
 
 **Test Case Design**
 

docs/contributing/ci/test_examples/l4_performance_tests.inc.md

@@ -200,4 +200,4 @@ You can add any benchmark running parameters you need here. For all optional par
 | num_prompts     | int / array | Yes      | 10,[10, 20, 30] | Number of requests. Supports single values or arrays. If a single value is used, it will be automatically expanded to match the number of qps or max_concurrency, e.g., [10,10,10]. If an array is used, its length must match the number of qps or max_concurrency. |
 | request_rate    | float / array | No  | 0.5, [0.5, 1, inf] | Queries per second. Supports single values or arrays. If a single value is used, it will be automatically expanded to match the number of num_prompts, e.g., [1,1,1]. If an array is used, its length must match the number of num_prompts.                          |
 | max_concurrency | int / array | No       | 1, [1, 2, 3]    | Maximum concurrent in-flight requests. Same array / expansion rules as `request_rate` (mutually exclusive with QPS mode).                                                                                                                                                                                                             |
-| baseline        | object      | No       | see below       | Optional thresholds. Prefer hardware-nested form: `{"H100": {"throughput_qps": …}, "A3": {…}}` (selected via runtime hardware label). Legacy flat `{metric: threshold}` still works. Per metric: scalar, list (sweep step), or object keyed by concurrency / QPS.
+| baseline        | object      | No       | see below       | Optional thresholds. **Hardware-nested only**: top-level keys must be `[hardware-resource]` pytest markers from `pyproject.toml` (`get_hardware_mark_list()`). Under each hardware key, metric names are free-form; values are scalars or sweep-aligned lists. Example: `{"H100": {"throughput_qps": […], "custom_stage_ms": 1.0}, "A3": {…}}`. Runtime filename aliases (`_RUNTIME_DEVICE_ALIASES`) are separate and do not expand this allowlist. |

docs/contributing/model/adding_diffusion_model.md

@@ -824,7 +824,7 @@ omni = Omni(model="your-model",
 
 ### CPU Offload
 
-See detailed guide: [CPU Offloading for Diffusion Models](../../user_guide/diffusion/cpu_offload_diffusion.md)
+See detailed guide: [CPU Offloading for Diffusion Models](../../user_guide/diffusion/cpu_offload.md)
 
 vLLM-Omni provides two offloading strategies to reduce GPU memory usage:
 

docs/design/feature/distributed_layerwise_offload.md

@@ -0,0 +1,156 @@
+# Distributed Layerwise Offload
+
+This document describes distributed layerwise offload (DLO) for diffusion
+models. DLO keeps only a small number of DiT blocks on the accelerator and
+streams the remaining blocks from host memory. The distributed backend can
+either shard those host-side weights across an existing parallel group or keep
+the standard loader's rank-local weights and avoid an additional collective.
+
+For user-facing commands, see the
+[distributed layerwise offloading guide](../../user_guide/diffusion/cpu_offload.md)
+and the [Cosmos3 recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-DistOffload.md).
+
+## Status
+
+DLO is implemented for multi-device diffusion execution. The default
+AllGather path is the primary path for DP and SP deployments. The
+`--dlo-no-use-allgather` path is a rank-local compatibility mode: it is useful
+for standard-loader sharding, workstation bring-up, and systems where an
+additional DLO collective is undesirable, but it does not reduce host weight
+storage across ranks.
+
+The compatibility matrix below describes the current implementation. The
+unit-level guards are covered, but not every parallelism combination has a
+full model-and-hardware end-to-end test.
+
+## Design
+
+### DLO consumes the existing parallel topology
+
+DLO does not create a new DP, TP, or SP topology. It reads the configured
+`DiffusionParallelConfig` and attaches offload hooks to the DiT blocks after the
+standard distributed groups have been initialized.
+
+The DLO weight-sharding group is selected as follows:
+
+1. Use the existing DP group when `data_parallel_size > 1`.
+2. When DP is one and SP is greater than one, use the SP group.
+3. Otherwise, run rank-locally without a DLO process group.
+
+TP is deliberately not used as DLO's AllGather group. HSDP has its own
+parameter-sharding lifecycle and is not allowed to be sharded a second time by
+DLO's AllGather path.
+
+### AllGather path
+
+With the default `dlo_use_allgather=True`, each rank stores approximately
+`1 / group_size` of each streamable block in pinned host memory. The next
+block's shard is copied to a device buffer and reconstructed with
+`all_gather_into_tensor` on a communication stream while the current block is
+executing.
+
+```text
+Compute:    [Block N]             [Block N+1]          [Block N+2]
+H2D:                      [shard N+1]           [shard N+2]
+AllGather:                [full N+1]             [full N+2]
+Buffers:    [current slot]       [prefetch slot]       [current slot]
+```
+
+![DLO double-buffer prefetch pipeline](../figures/dlo/dlo_pipeline.gif)
+
+The backend uses two shared device buffers, so accelerator weight residency is
+bounded by the largest streamed blocks rather than the complete model.
+
+When DP is greater than one, the engine can process one request per DP rank in
+the same denoising wave. Because AllGather is a collective, all participating
+requests must take the same execution path at every denoising step.
+
+### Rank-local path without DLO AllGather
+
+With `--dlo-no-use-allgather`, DLO forces its internal offload shard size to
+one. The regular model loader remains responsible for preparing each rank's
+weights, including TP-local tensors or HSDP-managed parameters. DLO then
+streams those rank-local tensors block by block using H2D copies only.
+
+This mode means:
+
+- DP still provides independent replicas, but DLO does not shard weights
+  across DP ranks.
+- SP still performs its normal activation/attention collectives, but DLO does
+  not shard weights across SP ranks.
+- TP/HSDP/SP collectives, if configured, are not disabled by this flag; only
+  DLO's additional weight AllGather is disabled.
+- Pure DP deployments keep a full host-side model copy per rank, subject to
+  shared OS page-cache behavior.
+- The scheduler does not require a synchronized DP request wave for DLO.
+
+This path is intentionally not implemented by reusing the AllGather mmap
+loader. It relies on the standard loader so model-specific TP/HSDP transforms
+remain intact.
+
+## Parallelism compatibility
+
+| Parallelism | DLO + AllGather | DLO without AllGather |
+|---|---|---|
+| **DP** | Supported primary path. DLO shards host weights across the DP group and can run DP multi-concurrency. | Supported rank-local path. DP replicas remain independent; no DLO weight sharding or cross-DP DLO collective. |
+| **SP** | Supported in the implementation. With DP=1, DLO uses the SP group for host-weight sharding; SP still shards sequence/activation work. | SP remains active, but DLO keeps standard-loader rank-local weights and adds no SP weight collective. |
+| **TP > 1** | Unsupported in the DLO mmap path. TP-aware loader callbacks are bypassed, so the backend rejects this configuration when it enters that path. | Standard loading is retained and TP-local tensors can be streamed. This is the intended compatibility path, but it still needs broader model and hardware validation. |
+| **HSDP** | Rejected. HSDP has already sharded parameters, so DLO AllGather would double-shard them. | Accepted by configuration. HSDP owns parameter sharding and its own gathers; DLO only stages rank-local parameters. End-to-end coverage is limited. |
+
+### Combined dimensions
+
+- **DP + SP:** DLO uses the DP group for weight sharding when DP is greater
+  than one; SP continues to use its own sequence-parallel group. If DP is one,
+  the SP group becomes DLO's sharding group in AllGather mode.
+- **DP + TP/SP without AllGather:** standard model loading defines the
+  rank-local tensor layout. DLO adds no cross-DP, cross-TP, or cross-SP weight
+  collective.
+- **HSDP + SP:** the general parallel configuration permits HSDP over SP, but
+  DLO must use `--dlo-no-use-allgather`. HSDP remains responsible for weight
+  materialization and synchronization.
+- **HSDP + DP or TP:** rejected independently by the diffusion parallel
+  configuration.
+
+## Request and loading constraints
+
+AllGather DP multi-concurrency requires:
+
+- explicit `num_inference_steps`;
+- the same `num_inference_steps` for all requests in a wave; and
+- identical request arguments that affect the collective execution path.
+
+The no-AllGather path does not impose these DLO-specific synchronized-wave
+requirements.
+
+The mmap loader is used only by the supported DLO+AllGather path when the
+model has a compatible checkpoint layout. Online quantization is incompatible
+with that sharded mmap path; use `--dlo-no-use-allgather` or disable online
+quantization. Models that do not meet the mmap requirements use the regular
+loader path.
+
+## Validation coverage
+
+Current source-level validation includes:
+
+- HSDP + DLO + AllGather rejection;
+- HSDP + DLO without AllGather acceptance at configuration level;
+- TP rejection in the DLO+AllGather mmap path;
+- resident-layer requests requiring no-AllGather;
+- DP request-wave validation for denoising-step compatibility;
+- sharding, double-buffer, AllGather-size, and heterogeneous-block regression
+  tests.
+
+The highest-value missing coverage is end-to-end numerical comparison against
+ordinary layerwise offload for DP+SP, TP+no-AllGather, and HSDP+SP+no-AllGather
+on the target CUDA/NCCL or CANN/HCCL hardware.
+
+## Recommendations
+
+- Use **DP + DLO AllGather** for the supported throughput and host-memory
+  scaling path.
+- Use **SP + DLO AllGather** for long-sequence workloads when DP concurrency is
+  not the goal.
+- Use **no-AllGather** to bring up TP or workstation PCIe configurations, with
+  the expectation of higher host-memory use and lower validation confidence.
+- Prefer **HSDP alone** for production HSDP deployments until the combined
+  HSDP + DLO no-AllGather path has broader end-to-end coverage.

docs/design/feature/quantization.md

@@ -0,0 +1,249 @@
+# Quantization
+
+This document describes the quantization architecture in vLLM-Omni and the
+extension points for contributors. For user-facing configuration, supported
+hardware, and method-specific instructions, see the
+[quantization user guide](../../user_guide/quantization/overview.md).
+
+## Goals and scope
+
+The design has four goals:
+
+1. Provide one `quantization_config` entry point for diffusion-only,
+   multi-stage omni/TTS, and multi-stage diffusion models.
+2. Reuse vLLM's quantization registry and `QuantizationConfig` contract where
+   a vLLM backend already provides the required behavior.
+3. Allow different components of one pipeline to use different quantization
+   configurations.
+4. Preserve the distinction between online quantization of a BF16/FP16
+   checkpoint and loading a pre-quantized checkpoint with serialized scales.
+
+This document covers configuration resolution, component routing, model
+integration, backend extension, and validation. Kernel implementations and
+user-facing support matrices are outside its scope.
+
+## Architecture
+
+The quantization path is split across the following layers:
+
+| Layer | Responsibility | Source |
+| --- | --- | --- |
+| Public API | Exposes `build_quant_config`, component routing, and backend registration | `vllm_omni/quantization/__init__.py` |
+| Factory | Normalizes method names, selects Omni overrides or vLLM registry backends, detects ModelOpt configs, and reconciles checkpoint metadata | `vllm_omni/quantization/factory.py` |
+| Component router | Resolves a quantization config from a layer prefix using longest-prefix matching | `vllm_omni/quantization/component_config.py` |
+| Method configs | Implement `QuantizationConfig` and select the platform-specific linear method | `vllm_omni/quantization/*_config.py` |
+| Diffusion configuration | Stores the canonical runtime config and propagates checkpoint configuration | `vllm_omni/diffusion/data.py` |
+| Worker and model integration | Passes the resolved config into vLLM's model configuration and model layers | `vllm_omni/diffusion/worker/diffusion_worker.py` and `vllm_omni/diffusion/models/` |
+
+The runtime contract remains vLLM's `QuantizationConfig`:
+
+```text
+user/CLI or model config
+          |
+          v
+OmniDiffusionConfig.quantization_config
+          |
+          +--> build_quant_config()
+          |        |
+          |        +--> Omni override
+          |        +--> vLLM registry backend
+          |        +--> ComponentQuantizationConfig
+          |
+          +--> checkpoint metadata reconciliation
+          |
+          +--> diffusion worker / vLLM config
+          |
+          +--> model layers call get_quant_method(layer, prefix)
+```
+
+The model layer receives either a method implementation or `None`. Returning
+`None` leaves that layer unquantized.
+
+## Configuration resolution
+
+`build_quant_config()` accepts the following forms:
+
+| Input | Resolution |
+| --- | --- |
+| `None` or `"none"` | Disable quantization. |
+| Method string, such as `"fp8"` | Resolve an Omni override first, then the vLLM quantization registry. |
+| Flat dictionary with `method` or `quant_method` | Normalize the method and construct its `QuantizationConfig`. |
+| Per-component dictionary | Construct a `ComponentQuantizationConfig`. |
+| Existing `QuantizationConfig` | Pass through without rebuilding it. |
+
+Method aliases are normalized case-insensitively, with `-` and `_` treated as
+equivalent. `auto-round` and `auto_round` use the Omni INC/AutoRound adapter.
+ModelOpt checkpoint metadata is detected from its method, producer metadata, or
+`quant_algo`, then mapped to the corresponding vLLM ModelOpt configuration.
+
+The factory uses Omni-specific overrides for methods whose diffusion behavior
+or platform dispatch is not provided by the vLLM registry. The current override
+set includes INT8, BitsAndBytes, MXFP8, MXFP4, MXFP4 dual-scale, and INC/
+AutoRound. Other methods, such as FP8, GGUF, and ModelOpt, are resolved through
+vLLM's registry when their configuration is compatible.
+
+## Per-component routing
+
+Multi-stage models do not necessarily want one quantization policy for every
+module. A component map expresses that policy without requiring each pipeline
+to construct separate vLLM configurations:
+
+```python
+quantization_config = {
+    "transformer": {"method": "fp8"},
+    "vae": None,
+    "default": None,
+}
+```
+
+`ComponentQuantizationConfig` applies the following rules:
+
+1. Component keys are interpreted as runtime layer-prefixes.
+2. The longest matching prefix wins.
+3. `None` disables quantization for the matching component.
+4. `default`, when present, handles prefixes with no explicit match.
+5. A matching child config delegates to its normal
+   `get_quant_method(layer, prefix)` implementation.
+
+Prefix matching happens after the model's weight-name mapping. New model
+integrations must verify the actual runtime prefixes, especially when a
+`WeightsMapper` changes checkpoint names. A prefix mismatch can silently route
+a layer to the default configuration, so component-routing tests should cover
+both positive matches and fall-through behavior.
+
+## Online and pre-quantized checkpoints
+
+The same method name can represent two different loading paths:
+
+- **Online**: the source checkpoint contains BF16/FP16 weights; the method
+  creates quantized weights or scales while the model is loaded.
+- **Pre-quantized**: the checkpoint contains quantized weights and the metadata
+  required by the method; the loader consumes those serialized tensors.
+
+Method configuration owns the mode flag, for example
+`is_checkpoint_int8_serialized` or `is_checkpoint_mxfp8_serialized`. The
+configuration's `get_quant_method()` selects the corresponding linear method.
+
+For pipelines with transformer-specific `config.json` files, such as cascade
+models, `resolve_quant_config_from_disk()` reconciles the active configuration
+with each transformer's checkpoint metadata. Its invariants are:
+
+- If no active configuration exists, valid checkpoint metadata can be used for
+  auto-detection.
+- A method mismatch is an error rather than an implicit conversion.
+- Serialized-checkpoint flags rebuild an online configuration into the
+  matching offline configuration.
+- Checkpoint-specific `ignored_layers` are preserved by rebuilding the active
+  configuration when they differ.
+- AutoRound MXFP8 metadata (`data_type="mx_fp"`) selects its offline path.
+
+This prevents a model from interpreting serialized weights with an online
+loader or silently applying the wrong method to one transformer in a
+multi-transformer pipeline.
+
+## Model and layer integration
+
+Diffusion model constructors receive the resolved configuration and pass it to
+quantizable vLLM layers. Each method configuration is responsible for:
+
+- identifying supported layers in `get_quant_method()`;
+- creating the correct parameter shapes and loaders;
+- processing weights after loading when conversion is required; and
+- applying the platform-specific quantized operation at runtime.
+
+Model code must keep precision-sensitive or unsupported modules unquantized.
+The shared `safe_quant_config()` helper is used by model integrations for
+normalization and modulation layers where a generic FP8 configuration would be
+unsafe. Text encoders, audio/vision encoders, VAEs, schedulers, and tokenizers
+are not implicitly quantized just because the main transformer is quantized;
+they require explicit component routing and method support.
+
+Weight quantization and quantized KV cache are separate paths. The former is
+resolved through `quantization_config`; runtime attention/KV-cache quantization
+is configured by the diffusion attention configuration and must be validated
+independently.
+
+## Platform and parallelism boundaries
+
+Platform checks belong in the method implementation rather than in the common
+factory. This keeps configuration parsing portable while allowing CUDA, NPU,
+and XPU methods to select different kernels or reject unsupported execution at
+model construction time. The user guide is the source of truth for the current
+hardware support matrix.
+
+Quantization methods must also respect the weight layout produced by the
+configured parallelism strategy:
+
+- Tensor-parallel layers may receive partitioned weight shapes and loaders.
+- Sequence-parallel execution changes activation movement but does not by
+  itself change the quantization configuration.
+- Data-parallel replicas build their own method state; quantization metadata
+  must be available consistently on every rank.
+- `WeightsMapper` transformations must be applied before matching ignored-layer
+  lists or component prefixes.
+
+Distributed layerwise offload has an additional boundary: its sharded mmap
+AllGather path is not compatible with online quantization because the path
+expects already materialized checkpoint weights. Use the rank-local
+`--dlo-no-use-allgather` path or a compatible pre-quantized checkpoint when
+combining these features. See the [DLO design document](distributed_layerwise_offload.md)
+for the full compatibility matrix.
+
+## Adding a quantization backend
+
+Contributors should follow this sequence:
+
+1. Check whether an existing vLLM backend can provide the required weight
+   format and kernel behavior. Prefer reuse over a duplicate Omni backend.
+2. If Omni-specific behavior is required, implement a `QuantizationConfig` and
+   its `QuantizeMethodBase`/linear method in `vllm_omni/quantization/`.
+3. Implement `get_name()`, supported activation dtypes, minimum capability,
+   checkpoint parsing, ignored-layer handling, and `get_quant_method()`.
+4. Add a lazy factory override in `factory.py`. Keep optional platform
+   dependencies out of package import paths and normalize aliases to one
+   canonical method name.
+5. If the method is pre-quantized, define the checkpoint metadata and the
+   serialized loading path before adding an online path.
+6. Add model-specific component routing only when the model contains stages
+   that must use different policies.
+7. Document user-visible configuration in the quantization user guide and keep
+   implementation rationale here.
+
+The public factory registration hook,
+`register_quantization_override(method, builder)`, is intended for integrations
+that need to register an Omni-compatible builder without changing the vLLM
+registry.
+
+## Validation
+
+Validation should cover the complete path, not only whether a config object can
+be instantiated:
+
+| Area | Required checks |
+| --- | --- |
+| Factory | Strings, aliases, flat dictionaries, invalid inputs, passthrough, and component maps |
+| Method config | Online/offline selection, ignored layers, platform guards, and weight loading |
+| Checkpoint metadata | Auto-detection, method mismatch errors, serialized flags, and per-transformer metadata |
+| Model integration | Correct quantization scope, unquantized components, and parallel weight layouts |
+| Numerical quality | Same-seed BF16 versus quantized output comparison at representative resolutions and step counts |
+| Performance | Peak device memory, generation latency, and any backend-specific throughput metric |
+
+The unit tests live under `tests/diffusion/quantization/`. The trajectory
+similarity tool in
+`vllm_omni/quantization/tools/compare_diffusion_trajectory_similarity.py`
+provides a repeatable quality and memory comparison for diffusion models.
+Method support and feature combinations remain model- and hardware-dependent;
+passing factory tests alone does not establish production support.
+
+## Design invariants
+
+- `quantization_config` is the single runtime source of truth after config
+  normalization.
+- A checkpoint method mismatch fails loudly instead of risking weight
+  corruption.
+- Components without an explicit supported quantization path remain in their
+  original precision.
+- Optional backend dependencies are imported lazily.
+- Method implementations own platform-specific capability checks and kernels.
+- User-facing support claims require model, hardware, numerical, and memory
+  validation in addition to configuration-level tests.

docs/design/feature/ray_based_execution.md

@@ -1,59 +0,0 @@
-# Distributed utils
-
-This directory (vllm_omni/distributed/ray_utils) contains utilities for distributed execution in vllm-omni, supporting both **Ray** and **Multiprocessing** backends.
-## 1. Installation
-```bash
-pip install "ray[default]"
-```
-## 2. Ray Utils
-
-The `ray_utils` module provides helper functions for managing Ray clusters and actors, which is essential for:
-*   **Multi-node deployment**: Running pipeline stages across different physical machines.
-*   **Resource management**: Efficient GPU/CPU allocation.
-
-### 2.1 Basic Usage
-
-To use the Ray backend, specify `worker_backend="ray"` when initializing the engine.
-
-**Command Line Example:**
-```bash
-vllm serve Qwen/Qwen2.5-Omni-7B \
-  --omni \
-  --port 8091 \
-  --worker-backend ray \
-  --ray-address auto
-```
-
-### 2.2 Cluster Setup
-
-**Step 1: Start Head Node**
-Run this on your primary machine:
-```bash
-ray start --head --port=6399
-```
-
-**Step 2: Connect Worker Nodes**
-Run this on each worker machine:
-```bash
-ray start --address=<HEAD_NODE_IP>:6399
-```
-
-> **Tip**: For a complete cluster setup script, refer to the vLLM example:
-> [run_cluster.sh](https://github.com/vllm-project/vllm/blob/main/examples/online_serving/run_cluster.sh)
-
-### 2.3 Distributed Connector Support
-
-When running on Ray, the system automatically adapts its communication strategy:
-
-*   **Cross-Node**: Recommended to use `MooncakeTransferEngineConnector` (RDMA, fastest) or `MooncakeStoreConnector` (TCP fallback).
-*   **Same-Node**: Can still use `SharedMemoryConnector` for efficiency, or Ray's native object store (plasma).
-*   **SHM payload routing**: `SharedMemoryConnector` always stores payloads in `/dev/shm`.
-
-### 2.4 Internal Helpers
-
-*   **`initialize_ray_cluster`**: Connects to an existing Ray cluster or starts a local one.
-
-## 3. Troubleshooting
-
-*   **Connection Issues**: Ensure the Ray head node is accessible and ports (default 6399 in this example) are open.
-*   **Version Mismatch**: Ensure all nodes run the same version of Ray and Python.

docs/design/index.md

@@ -1,8 +1,9 @@
 # Design Documents
 
 This section contains design documents and architecture specifications for
-vLLM-Omni. The sidebar groups documents by the system concern they describe,
-while the source paths remain stable so existing links continue to work.
+vLLM-Omni. The sidebar groups documents by the system concern they describe;
+retired module pages are preserved in the legacy archive instead of remaining
+in the active navigation.
 
 ## Architecture Documents
 
@@ -13,7 +14,6 @@ while the source paths remain stable so existing links continue to work.
 ### Runtime and stage execution
 
 - [Disaggregated Inference](feature/disaggregated_inference.md)
-- [Ray-based Execution](feature/ray_based_execution.md)
 - [Async Chunk](feature/async_chunk.md)
 - [Async Diffusion Output](feature/async_diffusion_output.md)
 - [Async Omni Output Materialization](feature/omni_async_output_materialization.md)
@@ -48,11 +48,12 @@ while the source paths remain stable so existing links continue to work.
 
 #### Quantization
 
-- [Quantization Overview](../user_guide/quantization/overview.md)
+- [Quantization](feature/quantization.md)
 
 - [Cache-DiT](feature/cache_dit.md)
 - [TeaCache](feature/teacache.md)
 - [Diffusion Continuous Batching](feature/diffusion_continuous_batching.md)
+- [Distributed Layerwise Offload](feature/distributed_layerwise_offload.md)
 
 ## Infrastructure and Performance
 
@@ -61,12 +62,29 @@ while the source paths remain stable so existing links continue to work.
 
 ## Module Design Documents
 
-### Runtime Modules
+- [Entrypoints and Serving Boundaries](module/entrypoints.md)
+- [vLLM-Omni Configuration](module/vllm_omni_config.md)
+- [Input, Output, and Modality Contracts](module/input_output_modality_contracts.md)
+- [Error Classification, Propagation, and Rendering](module/error_contracts.md)
+- [Engine Orchestration](module/engine_orchestration.md)
+- [Stage Runtime and Replica Lifecycle](module/stage_runtime.md)
+- [OmniConnector](module/omni_connector.md)
+- [Model Integration](module/model_integration.md)
+- [Autoregressive Runtime](module/ar_runtime.md)
+- Diffusion
+  - [Overview](module/diffusion/index.md)
+  - [Runtime](module/diffusion/diffusion_runtime.md)
+  - [Model Integration](module/diffusion/diffusion_model_integration.md)
+  - [Continuous Batching](module/diffusion/continuous_batching.md)
+  - [Parallelism](module/diffusion/parallelism.md)
+  - [Offloader](module/diffusion/offloader.md)
+- [Execution Platforms](module/execution_platforms.md)
+- [Cache Management](module/cache_management.md)
+- [Quantization](module/quantization.md)
+- [Observability](module/observability.md)
+- [Profiling](module/profiling.md)
+- [Benchmarking](module/benchmarking.md)
 
-- [AR Module](module/ar_module.md)
-- [DIT Module](module/dit_module.md)
-
-### Orchestration
-
-- [Entrypoint Module](module/entrypoint_module.md)
-- [AsyncOmni Architecture (Qwen3-Omni Example)](module/async_omni_architecture.md)
+The pre-#5137 pages are preserved in the
+[legacy module archive](module/archive/README.md) for historical reference and
+are not active design contracts.

docs/design/module/ar_runtime.md

@@ -0,0 +1,52 @@
+---
+title: Autoregressive Runtime
+kind: module
+status: draft
+owners:
+  - "@tzhouam"
+  - "@yinpeiqi"
+  - "@fake0fan"
+  - "@Sy0307"
+primary_code_paths:
+  - vllm_omni/core/**
+  - vllm_omni/worker/**
+related_code_paths:
+  - vllm_omni/model_executor/**
+depends_on:
+  - engine_orchestration.md
+  - model_integration.md
+  - input_output_modality_contracts.md
+validation_paths:
+  - tests/core/**
+  - tests/worker/**
+upstream_refs:
+  - vllm.v1.core
+  - vllm.v1.worker
+last_reviewed: 2026-07-16
+---
+
+# Autoregressive runtime
+
+The AR runtime extends vLLM scheduling and worker execution for omni-stage
+inputs and outputs while preserving vLLM scheduling and cache semantics.
+
+## Candidate invariants
+
+### AR-INV-001: vLLM owns base scheduling semantics
+
+**Rule:** Omni schedulers MUST preserve upstream request-state and cache
+transitions unless an Omni-specific difference is documented and tested.
+
+### AR-INV-002: Omni data crosses explicit adapters
+
+**Rule:** Modality-specific stage data MUST be converted at an input or output
+adapter, not injected through unrelated scheduler state.
+
+### AR-INV-003: Workers execute assigned work
+
+**Rule:** Workers and model runners MUST NOT implement cross-stage routing.
+
+## Safe-change guide
+
+Test request lifecycle, abort, cache state, and every affected worker execution
+mode against the supported upstream vLLM contract.

docs/design/module/archive/README.md

@@ -0,0 +1,18 @@
+# Archived module design pages
+
+This directory preserves module-design pages that predate the structure
+proposed in [issue #5137](https://github.com/vllm-project/vllm-omni/issues/5137).
+They are historical snapshots, not active architectural contracts, and are
+intentionally excluded from the documentation navigation.
+
+| Archived page | Active destination |
+| --- | --- |
+| [AR Module](ar_module.md) | [Autoregressive Runtime](../ar_runtime.md) |
+| [AsyncOmni Architecture](async_omni_architecture.md) | [Engine Orchestration](../engine_orchestration.md) and [Stage Runtime](../stage_runtime.md) |
+| [DiT Module](dit_module.md) | [Diffusion module overview](../diffusion/index.md) |
+| [Entrypoint Module](entrypoint_module.md) | [Entrypoints and Serving Boundaries](../entrypoints.md) |
+
+Do not update these pages with current design guidance. Move still-useful
+material into the active destination and review it against current code and
+tests instead. No compatibility redirect pages are provided at the former
+paths.

docs/design/module/archive/ar_module.md

@@ -12,7 +12,7 @@ The AR module of vLLM-Omni extends vLLM's core components to support:
 - **Hidden state exposure**: Exposing per-request hidden representations for downstream stages
 - **Basic generator support**: Support some basic heterogeneous architecture such as Convolution, LSTM, etc.
 
-As shown in the [end2end example](../../user_guide/examples/offline_inference/qwen3_omni.md), AR module can be widely applied across multiple stages, generating text tokens in thinker(AR), audio latent tokens in talker(AR) and audio wave in code2wav(Convolution).
+As shown in the [end2end example](../../../user_guide/examples/offline_inference/qwen3_omni.md), AR module can be widely applied across multiple stages, generating text tokens in thinker(AR), audio latent tokens in talker(AR) and audio wave in code2wav(Convolution).
 
 ## 2. Relationship with vLLM
 

docs/design/module/archive/dit_module.md

@@ -252,7 +252,7 @@ while True:
   batch of compatible independent requests for one full pipeline forward when
   the pipeline supports request-level batching; the step-wise path can also
   admit or remove compatible requests between denoise steps. See
-  [Diffusion Continuous Batching](../feature/diffusion_continuous_batching.md).
+  [Diffusion Continuous Batching](../../feature/diffusion_continuous_batching.md).
 
 ---
 
@@ -539,9 +539,9 @@ def get_attn_backend_for_role(role, head_size, attention_config=None, role_categ
 
 4. **Platform default** — `current_omni_platform.get_diffusion_attn_backend_cls(...)` picks the best-available kernel for the hardware.
 
-A backend that needs configuration exposes it as a typed field on `AttentionSpec`; `AttentionSpec.backend_kwargs()` serializes it into the `backend_kwargs` the impl consumes. Today the only such field is `skip_softmax` (TRTLLM_ATTN), so e.g. `target_sparsity` can be set without changing model code.
+A backend that needs configuration exposes it as a typed field on `AttentionSpec`; `AttentionSpec.backend_kwargs()` serializes it into the `backend_kwargs` the impl consumes. Today those fields are `skip_softmax` (TRTLLM_ATTN) and `block_sparse` (any block-sparse backend, currently RAINFUSION_ATTN), so e.g. `target_sparsity` or `sparsity` can be set without changing model code.
 
-For the user-facing CLI surface, see [Diffusion Attention Backends](../../user_guide/diffusion/attention_backends.md). For the role declaration contract on the model side, see [Adding a Diffusion Model](../../contributing/model/adding_diffusion_model.md).
+For the user-facing CLI surface, see [Diffusion Attention Backends](../../../user_guide/diffusion/attention_backends.md). For the role declaration contract on the model side, see [Adding a Diffusion Model](../../../contributing/model/adding_diffusion_model.md).
 
 **Backend Availability**:
 

docs/design/module/benchmarking.md

@@ -0,0 +1,47 @@
+---
+title: Benchmarking
+kind: module
+status: draft
+owners:
+  - "@alex-jw-brooks"
+primary_code_paths:
+  - vllm_omni/benchmarks/**
+  - vllm_omni/entrypoints/cli/benchmark/**
+related_code_paths:
+  - benchmarks/**
+depends_on:
+  - entrypoints.md
+  - profiling.md
+validation_paths:
+  - tests/benchmarks/**
+  - tests/dfx/perf/**
+upstream_refs:
+  - vllm.benchmarks
+last_reviewed: 2026-07-16
+---
+
+# Benchmarking
+
+Benchmarking defines reproducible workloads, metrics, warm-up, measurement,
+and result metadata for comparing vLLM-Omni performance.
+
+## Candidate invariants
+
+### BENCH-INV-001: Results are reproducible
+
+**Rule:** Results MUST identify the model, revision, commit, hardware,
+configuration, workload, concurrency, warm-up, and measurement window.
+
+### BENCH-INV-002: Correctness precedes performance
+
+**Rule:** A benchmark MUST validate successful and semantically acceptable
+outputs before using timings as performance evidence.
+
+### BENCH-INV-003: Metric definitions are explicit
+
+**Rule:** Metrics MUST state their unit, population, and aggregation.
+
+## Safe-change guide
+
+Test workload parsing and metric calculations with deterministic fixtures. Do
+not compare materially different workloads or runtime settings.

docs/design/module/cache_management.md

@@ -0,0 +1,47 @@
+---
+title: Cache Management
+kind: module
+status: draft
+owners:
+  - "@Isotr0py"
+  - "@princepride"
+primary_code_paths:
+  - vllm_omni/diffusion/cache/**
+related_code_paths:
+  - vllm_omni/experimental/ar_diffusion/kv_cache/**
+  - vllm_omni/core/**
+depends_on:
+  - ar_runtime.md
+  - diffusion/diffusion_runtime.md
+validation_paths:
+  - tests/diffusion/cache/**
+upstream_refs:
+  - vllm.v1.core.kv_cache_manager
+last_reviewed: 2026-07-16
+---
+
+# Cache management
+
+Cache management defines reusable state, cache identity, validity, lifecycle,
+and eviction across AR and diffusion execution.
+
+## Candidate invariants
+
+### CACHE-INV-001: Cache identity is complete
+
+**Rule:** A cache key MUST include every input and execution property that can
+change the reused value.
+
+### CACHE-INV-002: Reuse preserves correctness
+
+**Rule:** A cache optimization MUST provide a disabled path and MUST NOT reuse
+state after its validity conditions change.
+
+### CACHE-INV-003: Ownership has an end
+
+**Rule:** Request-scoped and model-scoped cache entries MUST have explicit
+reset, eviction, or teardown behavior.
+
+## Safe-change guide
+
+Test hit, miss, invalidation, disabled, concurrent-request, and teardown paths.

docs/design/module/diffusion/continuous_batching.md

@@ -0,0 +1,51 @@
+---
+title: Diffusion Continuous Batching
+kind: module
+status: draft
+owners:
+  - "@Isotr0py"
+  - "@princepride"
+primary_code_paths:
+  - vllm_omni/diffusion/sched/**
+related_code_paths:
+  - vllm_omni/diffusion/executor/**
+depends_on:
+  - diffusion_runtime.md
+validation_paths:
+  - tests/diffusion/batching/**
+upstream_refs: []
+last_reviewed: 2026-07-16
+---
+
+# Diffusion continuous batching
+
+Continuous batching defines when diffusion requests are compatible, how they
+share an execution step, and how each request advances independently.
+
+## Candidate invariants
+
+### BATCH-INV-001: Compatibility is explicit
+
+**Rule:** Requests MAY share a batch only when all properties required by the
+selected pipeline, scheduler step, shape, precision, and execution features are
+compatible.
+
+### BATCH-INV-002: Per-request state remains isolated
+
+**Rule:** Batch assembly MUST preserve each request's random generator, progress,
+conditioning, outputs, cancellation, and errors.
+
+### BATCH-INV-003: Admission is bounded
+
+**Rule:** Admission MUST respect configured capacity and MUST NOT rely on an
+unbounded waiting or active-request collection.
+
+### BATCH-INV-004: Batch order is not completion order
+
+**Rule:** Output association MUST use stable request identity rather than batch
+position across execution steps.
+
+## Safe-change guide
+
+Test heterogeneous requests, partial completion, cancellation, deterministic
+seeds, capacity limits, and batch sizes of one and many.

docs/design/module/diffusion/diffusion_model_integration.md

@@ -0,0 +1,58 @@
+---
+title: Diffusion Model Integration
+kind: module
+status: draft
+owners:
+  - "@Isotr0py"
+  - "@wtomin"
+  - "@ZJY0516"
+primary_code_paths:
+  - vllm_omni/diffusion/models/**
+  - vllm_omni/diffusion/model_loader/**
+related_code_paths:
+  - vllm_omni/diffusion/layers/**
+  - vllm_omni/diffusion/lora/**
+  - vllm_omni/diffusion/utils/**
+depends_on:
+  - diffusion_runtime.md
+  - ../input_output_modality_contracts.md
+validation_paths:
+  - tests/diffusion/models/**
+  - tests/diffusion/model_loader/**
+  - tests/diffusion/layers/**
+  - tests/diffusion/lora/**
+upstream_refs:
+  - diffusers.DiffusionPipeline
+last_reviewed: 2026-07-16
+---
+
+# Diffusion model integration
+
+Diffusion model integration owns pipeline contracts, registration, checkpoint
+loading, adapters, shared layers, and model-specific processing.
+
+## Candidate invariants
+
+### DIFF-MODEL-INV-001: Pipelines implement one runtime contract
+
+**Rule:** A pipeline MUST declare its supported modalities, configuration,
+inputs, outputs, loading path, and runtime capabilities.
+
+### DIFF-MODEL-INV-002: Registration is the selection boundary
+
+**Rule:** Runtime code MUST select model implementations through the registry or
+loader contract rather than scattered model-name conditionals.
+
+### DIFF-MODEL-INV-003: Model code does not schedule requests
+
+**Rule:** Pipeline code MUST NOT own admission, batching, cancellation, or
+cross-stage routing.
+
+### DIFF-MODEL-INV-004: Shared behavior stays shared
+
+**Rule:** Model directories SHOULD contain only genuine model differences.
+
+## Safe-change guide
+
+Test registry selection, checkpoint loading, minimal inference, input and output
+contracts, and every declared optional capability.

docs/design/module/diffusion/diffusion_runtime.md

@@ -0,0 +1,58 @@
+---
+title: Diffusion Runtime
+kind: module
+status: draft
+owners:
+  - "@Isotr0py"
+  - "@princepride"
+  - "@SamitHuang"
+primary_code_paths:
+  - vllm_omni/diffusion/sched/**
+  - vllm_omni/diffusion/executor/**
+  - vllm_omni/diffusion/worker/**
+related_code_paths:
+  - vllm_omni/diffusion/hooks/**
+  - vllm_omni/diffusion/postprocess/**
+depends_on:
+  - ../engine_orchestration.md
+  - ../input_output_modality_contracts.md
+  - diffusion_model_integration.md
+validation_paths:
+  - tests/diffusion/diffusion_backend/**
+  - tests/diffusion/hooks/**
+upstream_refs:
+  - diffusers.DiffusionPipeline
+last_reviewed: 2026-07-16
+---
+
+# Diffusion runtime
+
+The diffusion runtime owns request admission, scheduling, execution, progress,
+output, cancellation, and cleanup inside a diffusion stage.
+
+## Candidate invariants
+
+### DIFF-RUNTIME-INV-001: One lifecycle owner
+
+**Rule:** Every admitted request MUST have exactly one scheduler-owned lifecycle
+until completion, cancellation, or failure.
+
+### DIFF-RUNTIME-INV-002: Execution follows scheduler output
+
+**Rule:** Executors and workers MUST execute scheduler decisions without
+admitting, reordering, or forwarding requests independently.
+
+### DIFF-RUNTIME-INV-003: Terminal cleanup is complete
+
+**Rule:** Every terminal path MUST release request state, temporary tensors,
+hooks, and runtime-owned resources.
+
+### DIFF-RUNTIME-INV-004: Optional features use runtime hooks
+
+**Rule:** Cache, profiling, offload, and parallel features SHOULD integrate at
+defined hooks instead of creating another request lifecycle.
+
+## Safe-change guide
+
+Test success, cancellation, failure, shutdown, multiple requests, and resource
+cleanup.

docs/design/module/diffusion/index.md

@@ -0,0 +1,42 @@
+---
+title: Diffusion Modules
+kind: module-index
+status: draft
+owners:
+  - "@Isotr0py"
+  - "@princepride"
+  - "@SamitHuang"
+  - "@wtomin"
+  - "@ZJY0516"
+  - "@RuixiangMa"
+  - "@david6666666"
+  - "@xuechendi"
+primary_code_paths:
+  - vllm_omni/diffusion/**
+related_code_paths:
+  - vllm_omni/platforms/**
+depends_on:
+  - ../engine_orchestration.md
+  - ../input_output_modality_contracts.md
+validation_paths:
+  - tests/diffusion/**
+upstream_refs:
+  - diffusers
+last_reviewed: 2026-07-16
+---
+
+# Diffusion modules
+
+Diffusion is a module family containing runtime, model integration, batching,
+distributed execution, and memory-management subsystems.
+
+## Module documents
+
+- [Diffusion runtime](diffusion_runtime.md)
+- [Diffusion model integration](diffusion_model_integration.md)
+- [Continuous batching](continuous_batching.md)
+- [Parallelism](parallelism.md)
+- [Offloader](offloader.md)
+
+Cache, quantization, profiling, and benchmarking follow their top-level module
+documents.

docs/design/module/diffusion/offloader.md

@@ -0,0 +1,53 @@
+---
+title: Diffusion Offloader
+kind: module
+status: draft
+owners:
+  - "@yuanheng-zhao"
+  - "@Isotr0py"
+primary_code_paths:
+  - vllm_omni/diffusion/offloader/**
+related_code_paths:
+  - vllm_omni/diffusion/models/**
+  - vllm_omni/diffusion/worker/**
+depends_on:
+  - diffusion_runtime.md
+  - ../execution_platforms.md
+validation_paths:
+  - tests/diffusion/offloader/**
+upstream_refs:
+  - torch.nn.Module.to
+last_reviewed: 2026-07-16
+---
+
+# Diffusion offloader
+
+The offloader owns component residency, transfer scheduling, memory accounting,
+prefetch, and teardown for diffusion execution.
+
+## Candidate invariants
+
+### OFFLOAD-INV-001: Residency is explicit
+
+**Rule:** Each managed component MUST have one known residency state and one
+owner responsible for transitions.
+
+### OFFLOAD-INV-002: Use follows readiness
+
+**Rule:** Runtime execution MUST NOT consume a component until its transfer to
+the required device is complete.
+
+### OFFLOAD-INV-003: Transfers preserve model state
+
+**Rule:** Offloading MUST preserve parameter and buffer identity, dtype, device,
+and correctness required by the selected execution mode.
+
+### OFFLOAD-INV-004: Memory is bounded and released
+
+**Rule:** Retained host and device copies MUST respect configured limits and
+provide deterministic teardown.
+
+## Safe-change guide
+
+Test enabled and disabled paths, repeated execution, asynchronous transfer,
+memory limits, failure, teardown, and numerical equivalence.

docs/design/module/diffusion/parallelism.md

@@ -0,0 +1,56 @@
+---
+title: Diffusion Parallelism
+kind: module
+status: draft
+owners:
+  - "@princepride"
+  - "@yuanheng-zhao"
+  - "@xuechendi"
+primary_code_paths:
+  - vllm_omni/diffusion/distributed/**
+  - vllm_omni/diffusion/attention/parallel/**
+related_code_paths:
+  - vllm_omni/distributed/**
+  - vllm_omni/config/composable_parallel/**
+depends_on:
+  - diffusion_runtime.md
+  - ../vllm_omni_config.md
+validation_paths:
+  - tests/diffusion/distributed/**
+  - tests/diffusion/attention/**
+upstream_refs:
+  - torch.distributed
+last_reviewed: 2026-07-16
+---
+
+# Diffusion parallelism
+
+Parallelism owns diffusion rank topology, process groups, tensor and sequence
+sharding, collectives, and distributed execution.
+
+## Candidate invariants
+
+### PARALLEL-INV-001: Topology has one source of truth
+
+**Rule:** Rank coordinates, group membership, and parallel dimensions MUST be
+derived from validated configuration.
+
+### PARALLEL-INV-002: Sharding contracts are explicit
+
+**Rule:** Every distributed boundary MUST define tensor shape, sharded
+dimension, placement, dtype, producer, and consumer.
+
+### PARALLEL-INV-003: Collectives are symmetric
+
+**Rule:** Process-group members MUST invoke compatible collectives in the same
+logical order.
+
+### PARALLEL-INV-004: Single-rank behavior remains valid
+
+**Rule:** Distributed implementations SHOULD preserve a supported single-rank
+path without distributed initialization.
+
+## Safe-change guide
+
+Test topology validation, single-rank execution, affected parallel modes,
+combined modes, invalid shapes, and orderly teardown.

docs/design/module/engine_orchestration.md

@@ -0,0 +1,122 @@
+---
+title: Engine Orchestration
+kind: module
+status: draft
+architecture_state: current-plus-in-flight
+owners:
+  - "@tzhouam"
+  - "@fake0fan"
+document_stewards:
+  - "@hsliuustc0106"
+  - "@Gaohan123"
+  - "@david6666666"
+required_reviewers:
+  - "@yinpeiqi"
+  - "@Sy0307"
+  - "@yenuo26"
+  - "@NickCao"
+primary_code_paths:
+  - vllm_omni/engine/async_omni_engine.py
+  - vllm_omni/engine/async_engine_utils.py
+  - vllm_omni/engine/orchestrator.py
+  - vllm_omni/engine/cfg_companion_tracker.py
+  - vllm_omni/engine/rpc_result_router.py
+related_code_paths:
+  - vllm_omni/engine/messages.py
+  - vllm_omni/engine/orchestrator_monitor.py
+  - vllm_omni/engine/stage_client.py
+  - vllm_omni/engine/stage_pool.py
+  - vllm_omni/outputs/output_processor.py
+  - vllm_omni/distributed/omni_coordinator/**
+depends_on:
+  - input_output_modality_contracts.md
+  - error_contracts.md
+  - stage_runtime.md
+  - omni_connector.md
+validation_paths:
+  - tests/engine/test_async_omni_engine_input.py
+  - tests/engine/test_async_omni_engine_outputs.py
+  - tests/engine/test_async_omni_engine_abort.py
+  - tests/engine/test_async_omni_engine_stage_init.py
+  - tests/engine/test_orchestrator.py
+  - tests/engine/test_orchestrator_error_handling.py
+  - tests/engine/test_orchestrator_stage_input_bridge.py
+  - tests/engine/test_cfg_companion_lifecycle.py
+  - tests/engine/test_rpc_result_router.py
+  - tests/e2e/features/fullduplex/engine/**
+upstream_refs:
+  - vllm.v1.engine.EngineCoreRequest
+  - vllm.v1.engine.EngineCoreOutputs
+  - vllm.v1.engine.exceptions.EngineDeadError
+  - vllm.outputs.RequestOutput
+invariant_namespace: ORCH-INV
+last_reviewed: 2026-08-07
+last_verified_commit: 3d7fc3b9ba3cac88d579d4dc35b78b0b641675fc
+---
+
+# Engine orchestration
+
+Engine orchestration coordinates configured AR and diffusion stages while
+keeping public clients independent from stage processes and transports.
+
+## Contract status
+
+This document is a draft description of current behavior. It also identifies
+the boundary affected by the in-flight stage client/process refactor in
+[#5441](https://github.com/vllm-project/vllm-omni/pull/5441). Names and
+responsibilities proposed only by that PR are not current contracts.
+
+## Ownership boundary
+
+This document owns `AsyncOmniEngine`, `Orchestrator`, request-state creation,
+cross-stage routing, output ordering, companion tracking, control/RPC
+correlation, cancellation propagation, and terminal-state convergence.
+
+It does not own stage placement, replica selection, stage-process startup,
+payload schemas, public protocol rendering, connector transport, or semantic
+error classification. Those responsibilities belong to the stage runtime,
+I/O, entrypoint, connector, and error contracts.
+
+## Candidate invariants
+
+These identifiers are proposals while the document is `draft`.
+
+### ORCH-INV-001: The orchestrator owns cross-stage routing
+
+**Rule:** Entrypoints and stage clients MUST NOT independently forward a
+request to a downstream logical stage.
+
+### ORCH-INV-002: Stage clients do not own routing policy
+
+**Rule:** Stage clients MUST implement communication and lifecycle operations
+without selecting the next logical stage.
+
+### ORCH-INV-100: Terminal state is monotonic
+
+**Rule:** Once a request reaches a terminal state, orchestration MUST NOT
+forward new work for that request.
+
+## Invariant namespace
+
+`ORCH-INV` reserves `001-099` for dependency direction, `100-199` for
+request state and ordering, `200-299` for failure/cancellation/cleanup, and
+`300-399` for extension points, shutdown ordering, and upstream alignment.
+Numbers become append-only after normative promotion.
+
+## Safe-change guide
+
+Test routing, output ordering, queue correlation, cancellation, failure
+propagation, shutdown ordering, and representative multi-stage execution.
+Changes that cross into stage runtime or public error behavior require review
+from that contract's owners.
+
+## Promotion gate
+
+- Reconcile current names and boundaries after #5441 merges, closes, or is
+  replaced.
+- Demonstrate one terminal outcome per request across success, abort,
+  request-scoped failure, and fatal engine failure.
+- Verify queue correlation and orchestrator-before-runtime shutdown ordering
+  in the cited validation paths.
+- Obtain approval from a technical owner and an independent validation
+  reviewer.

docs/design/module/entrypoints.md

@@ -0,0 +1,134 @@
+---
+title: Entrypoints and Serving Boundaries
+kind: module
+status: draft
+architecture_state: current-plus-in-flight
+owners:
+  - "@alex-jw-brooks"
+  - "@linyueqian"
+document_stewards:
+  - "@hsliuustc0106"
+  - "@Gaohan123"
+  - "@david6666666"
+required_reviewers:
+  - "@tzhouam"
+  - "@herotai214"
+  - "@yenuo26"
+  - "@NickCao"
+primary_code_paths:
+  - vllm_omni/entrypoints/omni.py
+  - vllm_omni/entrypoints/async_omni.py
+  - vllm_omni/entrypoints/omni_base.py
+  - vllm_omni/entrypoints/cli/**
+  - vllm_omni/entrypoints/openai/**
+  - vllm_omni/entrypoints/openpi/**
+  - vllm_omni/entrypoints/client_request_state.py
+  - vllm_omni/entrypoints/stage_utils.py
+  - vllm_omni/entrypoints/utils.py
+primary_path_exceptions:
+  - path: vllm_omni/entrypoints/openai/errors.py
+    owner: error_contracts.md
+related_code_paths:
+  - vllm_omni/errors.py
+  - vllm_omni/inputs/**
+  - vllm_omni/outputs/**
+  - vllm_omni/engine/async_omni_engine.py
+  - vllm_omni/config/**
+  - vllm_omni/deploy/**
+depends_on:
+  - input_output_modality_contracts.md
+  - error_contracts.md
+  - engine_orchestration.md
+validation_paths:
+  - tests/entrypoints/test_omni_entrypoints.py
+  - tests/entrypoints/test_async_omni.py
+  - tests/entrypoints/test_async_omni_duplex.py
+  - tests/entrypoints/test_serve.py
+  - tests/entrypoints/test_stream_finish_reason.py
+  - tests/entrypoints/openai/**
+  - tests/entrypoints/openai_api/**
+  - tests/e2e/online_serving/**
+upstream_refs:
+  - vllm.engine.protocol.EngineClient
+  - vllm.entrypoints.openai/**
+  - vllm.entrypoints.serve/**
+  - vllm.renderers/**
+  - vllm.v1.engine.exceptions/**
+invariant_namespace: ENTRY-INV
+last_reviewed: 2026-08-07
+last_verified_commit: 3d7fc3b9ba3cac88d579d4dc35b78b0b641675fc
+---
+
+# Entrypoints and serving boundaries
+
+Entrypoints translate offline, CLI, and serving requests into stable engine
+operations and translate engine outputs into public responses.
+
+## Contract status
+
+This document describes current entrypoint responsibilities plus the boundary
+under review in the open roadmap
+[#5227](https://github.com/vllm-project/vllm-omni/issues/5227) and helper-move
+PR [#5453](https://github.com/vllm-project/vllm-omni/pull/5453). In-flight
+helper locations are not treated as current paths.
+
+## Ownership boundary
+
+This document owns offline API semantics, CLI and serve composition,
+supported OpenAI-compatible routes, request validation and normalization,
+response conversion, streaming/session behavior, and engine handoff.
+
+It does not own configuration precedence, cross-stage routing, stage
+placement, payload implementation, or semantic error classification.
+`entrypoints/openai/errors.py` is an explicit primary-path exception owned by
+`error_contracts.md`.
+
+## Candidate invariants
+
+These identifiers are proposals while the document is `draft`.
+
+### ENTRY-INV-001: Entrypoints adapt but do not orchestrate
+
+**Rule:** Entrypoints MUST NOT implement cross-stage routing or stage lifecycle
+policy.
+
+### ENTRY-INV-100: Public requests are normalized once
+
+**Rule:** Public protocol values MUST be validated and converted to internal
+request contracts before engine submission.
+
+### ENTRY-INV-101: Streaming preserves request identity
+
+**Rule:** Every streamed response MUST remain associated with the request and
+output modality that produced it.
+
+The reviewer-proposed rule that model-specific behavior should stay behind a
+common adapter or processor abstraction remains an unnumbered candidate until
+the entrypoint refactor makes that boundary enforceable.
+
+## Invariant namespace
+
+`ENTRY-INV` reserves `001-099` for boundary and dependency direction,
+`100-199` for normalization, sessions, and streaming, `200-299` for
+disconnect, cancellation, rendering, and cleanup, and `300-399` for upstream
+route and renderer compatibility. Numbers become append-only after normative
+promotion.
+
+## Safe-change guide
+
+Test request validation, protocol conversion, model-adapter routing,
+streaming, session identity, disconnect/cancellation, and error mapping for
+each affected offline or serving entrypoint.
+
+## Promotion gate
+
+- Reconcile the page after #5227's P0 ownership work and #5453's final
+  disposition.
+- Publish a supported route/transport matrix rather than claiming blanket
+  OpenAI parity.
+- Verify one normalization and handoff path per route family and one terminal
+  outcome per streaming transport.
+- Promote the `api_server.py` composition-root rule only after helper moves and
+  import boundaries enforce it.
+- Obtain approval from a technical owner and an independent validation
+  reviewer.

docs/design/module/error_contracts.md

@@ -0,0 +1,109 @@
+---
+title: Error Classification, Propagation, and Rendering
+kind: module
+status: draft
+architecture_state: current-plus-rfc
+owners:
+  - "@alex-jw-brooks"
+  - "@NickCao"
+document_stewards:
+  - "@hsliuustc0106"
+  - "@Gaohan123"
+  - "@david6666666"
+required_reviewers:
+  - "@yinpeiqi"
+  - "@tzhouam"
+  - "@yenuo26"
+primary_code_paths:
+  - vllm_omni/errors.py
+  - vllm_omni/entrypoints/openai/errors.py
+related_code_paths:
+  - vllm_omni/engine/messages.py
+  - vllm_omni/engine/orchestrator.py
+  - vllm_omni/entrypoints/omni_base.py
+  - vllm_omni/entrypoints/async_omni.py
+  - vllm_omni/entrypoints/openai/api_server.py
+  - vllm_omni/entrypoints/openai/serving_chat.py
+  - vllm_omni/entrypoints/openai/serving_speech.py
+  - vllm_omni/entrypoints/openai/serving_audio_generate.py
+  - vllm_omni/entrypoints/openai/protocol/images.py
+  - vllm_omni/entrypoints/openai/protocol/videos.py
+depends_on: []
+validation_paths:
+  - tests/engine/test_orchestrator_error_handling.py
+  - tests/engine/test_async_omni_engine_outputs.py
+  - tests/entrypoints/test_omni_entrypoints.py
+  - tests/entrypoints/test_stream_finish_reason.py
+  - tests/entrypoints/openai_api/test_image_server.py
+  - tests/entrypoints/openai_api/test_video_server.py
+  - tests/entrypoints/openai_api/test_serving_audio_generate.py
+  - tests/entrypoints/openai_api/test_serving_speech.py
+  - tests/metrics/test_prometheus.py
+upstream_refs:
+  - vllm.v1.engine.exceptions.EngineDeadError
+  - vllm.v1.engine.exceptions.EngineGenerateError
+  - vllm.entrypoints.serve.utils.error_response
+  - vLLM RFC #48227
+  - vLLM PR #49665
+invariant_namespace: ERR-INV
+last_reviewed: 2026-08-07
+last_verified_commit: 3d7fc3b9ba3cac88d579d4dc35b78b0b641675fc
+---
+
+# Error classification, propagation, and rendering
+
+This contract separates the meaning of a failure from the transport used to
+carry it and the public protocol used to render it.
+
+## Contract status
+
+Current behavior is distributed across engine, message, and entrypoint code.
+The semantic hierarchy proposed by
+[#5570](https://github.com/vllm-project/vllm-omni/issues/5570) remains an open
+RFC. Proposed types and wire records from that RFC are not current contracts.
+
+## Ownership boundary
+
+This document owns semantic classification, fatal versus request-scoped
+failure, cross-process error context, compatibility policy, sanitization, and
+offline, HTTP, SSE, WebSocket, and asynchronous-job rendering rules.
+
+`engine/messages.py` remains primarily owned by the I/O contract; this page
+owns the semantics of its error fields. Route and serving files remain
+primarily owned by `entrypoints.md`; this page owns their rendering
+constraints.
+
+## Candidate invariants
+
+The following topics require owner review before receiving stable IDs:
+
+- A request-scoped failure terminates only the affected request unless the
+  underlying engine is no longer usable.
+- Every public rendering path preserves a stable classification while
+  removing sensitive internal context.
+- Cancellation remains distinguishable from internal failure across process
+  and protocol boundaries.
+
+## Invariant namespace
+
+`ERR-INV` reserves `001-099` for classification ownership, `100-199` for wire
+context and compatibility, `200-299` for fatality, termination, sanitization,
+and rendering, and `300-399` for upstream alignment and transport extension
+points. Numbers become append-only after normative promotion.
+
+## Safe-change guide
+
+Exercise exception classification, stage wire round trips, fatal and
+request-scoped propagation, cancellation, sanitization, offline exceptions,
+HTTP status/body consistency, streaming terminal events, asynchronous-job
+persistence, and metrics classification.
+
+## Promotion gate
+
+- Resolve #5570's hierarchy, wire record, sanitization, and mixed-version
+  questions.
+- Add table-driven exception/status/type tests and stage wire round trips.
+- Verify all public transports terminate once and do not leak sensitive
+  payload data.
+- Obtain approval from both technical owners and an independent validation
+  reviewer.

docs/design/module/execution_platforms.md

@@ -0,0 +1,48 @@
+---
+title: Execution Platforms
+kind: module
+status: draft
+owners:
+  - "@gcanlin"
+  - "@tjtanaa"
+  - "@xuechendi"
+primary_code_paths:
+  - vllm_omni/platforms/**
+related_code_paths:
+  - vllm_omni/attention/**
+depends_on:
+  - ar_runtime.md
+  - diffusion/index.md
+validation_paths:
+  - tests/platforms/**
+upstream_refs:
+  - vllm.platforms
+last_reviewed: 2026-07-16
+---
+
+# Execution platforms
+
+Execution platforms isolate hardware-specific capability detection, worker
+selection, patches, kernels, and configuration adjustments.
+
+## Candidate invariants
+
+### PLATFORM-INV-001: Capabilities are explicit
+
+**Rule:** Hardware-dependent behavior MUST be guarded by platform selection or
+capability detection and MUST fail clearly when unsupported.
+
+### PLATFORM-INV-002: Portable code remains portable
+
+**Rule:** Platform-neutral modules MUST NOT import a vendor implementation
+directly when a common selection interface exists.
+
+### PLATFORM-INV-003: Overrides are minimal
+
+**Rule:** A platform implementation SHOULD override only behavior that differs
+from the common or upstream implementation.
+
+## Safe-change guide
+
+Validate common import paths without the target accelerator and run focused
+tests in a fresh `uv` environment on every affected device platform.

docs/design/module/input_output_modality_contracts.md

@@ -0,0 +1,133 @@
+---
+title: Input, Output, and Modality Contracts
+kind: module
+status: draft
+architecture_state: current
+owners:
+  - "@Sy0307"
+  - "@amy-why-3459"
+document_stewards:
+  - "@hsliuustc0106"
+  - "@Gaohan123"
+  - "@david6666666"
+required_reviewers:
+  - "@tzhouam"
+  - "@linyueqian"
+  - "@david6666666"
+  - "@yenuo26"
+  - "@NickCao"
+primary_code_paths:
+  - vllm_omni/inputs/**
+  - vllm_omni/outputs/**
+  - vllm_omni/request.py
+  - vllm_omni/data_entry_keys.py
+  - vllm_omni/engine/messages.py
+  - vllm_omni/engine/serialization.py
+  - vllm_omni/engine/mm_outputs.py
+  - vllm_omni/engine/output_modality.py
+  - vllm_omni/engine/output_processor.py
+related_code_paths:
+  - vllm_omni/model_executor/stage_input_processors/**
+  - vllm_omni/core/sched/**
+  - vllm_omni/distributed/omni_connectors/**
+  - vllm_omni/diffusion/request.py
+  - vllm_omni/diffusion/worker/request_batch.py
+  - vllm_omni/errors.py
+depends_on:
+  - error_contracts.md
+validation_paths:
+  - tests/inputs/**
+  - tests/engine/test_data_entry_keys.py
+  - tests/engine/test_omni_request.py
+  - tests/engine/test_omni_request_output.py
+  - tests/engine/test_multimodal_accumulation.py
+  - tests/engine/test_output_modality.py
+  - tests/engine/test_output_processor.py
+  - tests/engine/test_wire_multimodal_output.py
+  - tests/utils/test_mm_outputs.py
+  - tests/utils/test_mm_outputs_partition.py
+  - tests/diffusion/test_diffusion_output_metadata.py
+upstream_refs:
+  - vllm.inputs/**
+  - vllm.outputs/**
+  - vllm.v1.request.Request
+  - vllm.v1.engine.EngineCoreRequest
+  - vllm.v1.engine.EngineCoreOutput
+  - vllm.v1.engine.output_processor.OutputProcessor
+invariant_namespace: IO-INV
+last_reviewed: 2026-08-07
+last_verified_commit: 3d7fc3b9ba3cac88d579d4dc35b78b0b641675fc
+---
+
+# Input, output, and modality contracts
+
+These contracts define the data that may cross entrypoint, orchestration,
+stage, connector, and model boundaries.
+
+## Contract status
+
+This document is a draft description of the current request, message,
+serialization, modality, accumulation, and output types. It does not assign
+semantic error ownership to the wire schema.
+
+## Ownership boundary
+
+This document owns request identity, prompt/input types, modality keys and
+metadata, queue and wire message schemas, serialization, output types,
+accumulation, completion semantics, and compatibility shims.
+
+It does not own route-specific validation or rendering, model-specific
+interpretation, connector transport mechanics, scheduling policy, or semantic
+error classification. The `ErrorMessage` schema belongs here; the meaning and
+public rendering of its error fields belong to `error_contracts.md`.
+
+## Candidate invariants
+
+These identifiers are proposals while the document is `draft`.
+
+### IO-INV-001: Boundary data has an explicit modality
+
+**Rule:** Data crossing a module or stage boundary MUST identify its modality
+and use the corresponding validated contract.
+
+### IO-INV-100: Request identity is stable
+
+**Rule:** Request identity MUST be preserved across conversions, stages,
+streaming updates, cancellation, and errors.
+
+### IO-INV-101: Completion is explicit
+
+**Rule:** Producers MUST distinguish partial updates from the terminal update
+for every output modality.
+
+### IO-INV-300: Internal objects do not leak into public protocols
+
+**Rule:** Entrypoints MUST explicitly translate internal output objects into
+public response types.
+
+## Invariant namespace
+
+`IO-INV` reserves `001-099` for schema ownership and dependency direction,
+`100-199` for identity, modality, accumulation, ordering, and completion,
+`200-299` for serialization failure and payload cleanup, and `300-399` for
+upstream extension, deprecation, and compatibility. Numbers become
+append-only after normative promotion.
+
+## Safe-change guide
+
+Test construction, validation, serialization, round trips, streaming,
+accumulation, completion, optional fields, unknown fields, and compatibility
+at every affected producer-consumer boundary.
+
+## Promotion gate
+
+- Add explicit CODEOWNERS for the primary I/O paths or confirm that fallback
+  ownership is intentional.
+- Document the compatibility window for deprecated `vllm_omni.engine`
+  I/O-related imports.
+- Verify wire round trips and rejection of unknown or incompatible schema
+  fields.
+- Define payload-versus-metadata rules and completion semantics for streaming
+  and batched outputs.
+- Obtain approval from a technical owner and independent AR, diffusion,
+  connector, entrypoint, and validation reviewers as applicable.

docs/design/module/model_integration.md

@@ -0,0 +1,51 @@
+---
+title: Model Integration
+kind: module
+status: draft
+owners:
+  - "@tzhouam"
+  - "@gcanlin"
+primary_code_paths:
+  - vllm_omni/model_executor/**
+  - vllm_omni/model_extras/**
+  - vllm_omni/plugins/**
+related_code_paths:
+  - vllm_omni/transformers_utils/**
+  - vllm_omni/tokenizers/**
+depends_on:
+  - input_output_modality_contracts.md
+  - ar_runtime.md
+validation_paths:
+  - tests/model_executor/**
+  - tests/model_extras/**
+  - tests/plugin/**
+upstream_refs:
+  - vllm.model_executor
+last_reviewed: 2026-07-16
+---
+
+# Model integration
+
+Model integration adapts model-specific configuration, preprocessing, loading,
+and execution behavior to stable runtime contracts.
+
+## Candidate invariants
+
+### MODEL-INV-001: Registration is explicit
+
+**Rule:** A model integration MUST declare how its model class, loader, input
+processor, and stage configuration are selected.
+
+### MODEL-INV-002: Model code does not route stages
+
+**Rule:** Model-specific code MUST NOT select or invoke a downstream omni stage.
+
+### MODEL-INV-003: Upstream behavior is reused deliberately
+
+**Rule:** An upstream vLLM implementation SHOULD be reused when its contract is
+sufficient; an override MUST document the behavioral difference.
+
+## Safe-change guide
+
+Test registration, checkpoint loading, input conversion, and representative
+model execution. Test shared utilities against more than one integration.

docs/design/module/observability.md

@@ -0,0 +1,46 @@
+---
+title: Observability
+kind: module
+status: draft
+owners:
+  - "@lishunyang12"
+primary_code_paths:
+  - vllm_omni/metrics/**
+related_code_paths:
+  - vllm_omni/engine/**
+  - vllm_omni/entrypoints/**
+depends_on:
+  - engine_orchestration.md
+validation_paths:
+  - tests/metrics/**
+upstream_refs:
+  - vllm.v1.metrics
+last_reviewed: 2026-07-16
+---
+
+# Observability
+
+Observability defines stable metrics, logs, and correlation fields for tracing a
+request across entrypoints, orchestration, stages, and workers.
+
+## Candidate invariants
+
+### OBS-INV-001: Request correlation survives stage boundaries
+
+**Rule:** Logs and metrics for the same request SHOULD carry a stable request
+identifier and stage identity.
+
+### OBS-INV-002: Metric meaning is stable
+
+**Rule:** A metric's unit, labels, aggregation, and lifecycle MUST be documented;
+changing any of them is a contract change.
+
+### OBS-INV-003: Labels are bounded
+
+**Rule:** Metrics MUST NOT use request IDs, prompts, outputs, or other unbounded
+values as labels.
+
+## Safe-change guide
+
+Test registration, label sets, units, lifecycle resets, and representative
+multi-stage emission.

docs/design/module/omni_connector.md

@@ -0,0 +1,50 @@
+---
+title: OmniConnector
+kind: module
+status: draft
+owners:
+  - "@princepride"
+  - "@yuanheng-zhao"
+  - "@xuechendi"
+  - "@natureofnature"
+  - "@fake0fan"
+primary_code_paths:
+  - vllm_omni/distributed/omni_connectors/**
+related_code_paths:
+  - vllm_omni/platforms/*/omni_connectors/**
+depends_on:
+  - input_output_modality_contracts.md
+  - vllm_omni_config.md
+validation_paths:
+  - tests/distributed/omni_connectors/**
+upstream_refs:
+  - vllm.distributed
+last_reviewed: 2026-07-16
+---
+
+# OmniConnector
+
+OmniConnector defines model-agnostic transport and synchronization contracts
+for data exchanged across stages, processes, devices, and nodes.
+
+## Candidate invariants
+
+### CONNECTOR-INV-001: Connectors transport but do not route
+
+**Rule:** A connector MUST NOT select the next stage or implement model-specific
+execution policy.
+
+### CONNECTOR-INV-002: Producer and consumer contracts agree
+
+**Rule:** Both ends MUST agree on data identity, shape, dtype, placement,
+ownership, and completion semantics.
+
+### CONNECTOR-INV-003: Resources have deterministic cleanup
+
+**Rule:** Connections, buffers, handles, and background work MUST be released on
+normal completion, cancellation, failure, and shutdown.
+
+## Safe-change guide
+
+Test setup, transfer, synchronization, timeout, cancellation, failure, and
+cleanup across every affected backend.

docs/design/module/profiling.md

@@ -0,0 +1,46 @@
+---
+title: Profiling
+kind: module
+status: draft
+owners:
+  - "@gcanlin"
+primary_code_paths:
+  - vllm_omni/profiler/**
+  - vllm_omni/diffusion/profiler/**
+related_code_paths:
+  - vllm_omni/entrypoints/**
+depends_on:
+  - observability.md
+validation_paths:
+  - tests/profile/**
+upstream_refs:
+  - vllm.profiler
+last_reviewed: 2026-07-16
+---
+
+# Profiling
+
+Profiling owns opt-in instrumentation and trace collection used to locate
+runtime bottlenecks without changing model results or execution policy.
+
+## Candidate invariants
+
+### PROFILE-INV-001: Profiling is opt-in
+
+**Rule:** Expensive trace collection MUST be disabled by default and bounded by
+an explicit start and stop lifecycle.
+
+### PROFILE-INV-002: Instrumentation preserves semantics
+
+**Rule:** Profiling MUST NOT change scheduling decisions, generated outputs, or
+resource ownership beyond documented synchronization overhead.
+
+### PROFILE-INV-003: Trace context is sufficient
+
+**Rule:** Trace events SHOULD identify stage, rank, worker, operation, and
+request where available without embedding user payloads.
+
+## Safe-change guide
+
+Test disabled overhead, repeated start and stop, cleanup, multi-rank trace
+naming, and each supported profiler backend.

docs/design/module/quantization.md

@@ -0,0 +1,50 @@
+---
+title: Quantization
+kind: module
+status: draft
+owners:
+  - "@david6666666"
+  - "@Isotr0py"
+  - "@lishunyang12"
+primary_code_paths:
+  - vllm_omni/quantization/**
+  - vllm_omni/diffusion/quantization/**
+related_code_paths:
+  - vllm_omni/platforms/*/quant/**
+depends_on:
+  - model_integration.md
+  - diffusion/diffusion_model_integration.md
+validation_paths:
+  - tests/quantization/**
+  - tests/diffusion/quantization/**
+upstream_refs:
+  - vllm.model_executor.layers.quantization
+last_reviewed: 2026-07-16
+---
+
+# Quantization
+
+Quantization owns method registration, checkpoint interpretation, layer
+replacement, and runtime constraints for reduced-precision execution.
+
+## Candidate invariants
+
+### QUANT-INV-001: Checkpoint and runtime methods agree
+
+**Rule:** The selected method MUST validate checkpoint metadata, supported
+layers, device capabilities, and required dependencies before use.
+
+### QUANT-INV-002: Unsupported combinations fail clearly
+
+**Rule:** Quantization MUST NOT silently fall back to a different precision or
+method when a requested combination is unsupported.
+
+### QUANT-INV-003: Numerical validation is required
+
+**Rule:** A new method or layer mapping MUST include correctness comparison
+against a defined reference result.
+
+## Safe-change guide
+
+Test loading, representative layers, metadata, platform guards, and numerical
+tolerances separately from performance.

docs/design/module/stage_runtime.md

@@ -0,0 +1,126 @@
+---
+title: Stage Runtime and Replica Lifecycle
+kind: module
+status: draft
+architecture_state: current-plus-in-flight
+owners:
+  - "@tzhouam"
+  - "@fake0fan"
+document_stewards:
+  - "@hsliuustc0106"
+  - "@Gaohan123"
+  - "@david6666666"
+required_reviewers:
+  - "@Sy0307"
+  - "@princepride"
+  - "@chickeyton"
+  - "@yenuo26"
+  - "@NickCao"
+primary_code_paths:
+  - vllm_omni/engine/stage_runtime.py
+  - vllm_omni/engine/stage_pool.py
+  - vllm_omni/engine/stage_client.py
+  - vllm_omni/engine/membership_controller.py
+  - vllm_omni/engine/stage_engine_startup.py
+  - vllm_omni/engine/stage_engine_core_client.py
+  - vllm_omni/engine/stage_engine_core_proc.py
+  - vllm_omni/engine/stage_engine_core_proc_manager.py
+  - vllm_omni/engine/stage_init_utils.py
+related_code_paths:
+  - vllm_omni/distributed/omni_coordinator/**
+  - vllm_omni/distributed/ray_utils/**
+  - vllm_omni/diffusion/inline_stage_diffusion_client.py
+  - vllm_omni/diffusion/stage_diffusion_client.py
+  - vllm_omni/diffusion/stage_diffusion_proc.py
+  - vllm_omni/config/**
+depends_on:
+  - input_output_modality_contracts.md
+  - error_contracts.md
+  - omni_connector.md
+validation_paths:
+  - tests/engine/test_async_omni_engine_stage_init.py
+  - tests/engine/test_membership_controller.py
+  - tests/engine/test_single_stage_mode.py
+  - tests/engine/test_stage_engine_core_client.py
+  - tests/engine/test_stage_engine_core_proc.py
+  - tests/diffusion/test_inline_stage_diffusion_client.py
+  - tests/diffusion/test_stage_diffusion_proc.py
+  - tests/distributed/omni_coordinator/**
+upstream_refs:
+  - vllm.v1.engine.core_client.AsyncMPClient
+  - vllm.v1.engine.core_client.DPLBAsyncMPClient
+  - vllm.v1.engine.core.EngineCoreProc
+  - vllm.v1.engine.utils.CoreEngineProcManager
+  - vllm.v1.engine.coordinator.DPCoordinator
+invariant_namespace: STAGE-INV
+last_reviewed: 2026-08-07
+last_verified_commit: 3d7fc3b9ba3cac88d579d4dc35b78b0b641675fc
+---
+
+# Stage runtime and replica lifecycle
+
+The stage runtime turns a logical stage definition into ready local or
+distributed replicas and provides the lifecycle boundary used by the
+orchestrator.
+
+## Contract status
+
+This document describes the current `StageRuntime`, `StagePool`, stage-client,
+and stage-process layout. The unified LLM/diffusion direction in
+[#5441](https://github.com/vllm-project/vllm-omni/pull/5441) remains in flight
+and is not presented as current behavior.
+
+## Ownership boundary
+
+This document owns local and distributed stage placement, startup and
+readiness, one pool per logical stage, replica identity and membership,
+selection and affinity, stage-client/process lifecycle, liveness, draining,
+and shutdown.
+
+It does not own cross-stage request policy, model scheduler policy, connector
+transfer semantics, configuration precedence, or public error rendering.
+
+## Candidate invariants
+
+These identifiers are proposals while the document is `draft`.
+
+### STAGE-INV-001: The runtime owns replica lifecycle
+
+**Rule:** Orchestration MUST acquire stage capacity through the stage runtime
+and MUST NOT construct or retire backend replicas directly.
+
+### STAGE-INV-100: Affinity uses stable replica identity
+
+**Rule:** Follow-up operations for a request with replica affinity MUST resolve
+the same logical replica until the request terminates or the replica is
+declared lost.
+
+### STAGE-INV-200: Shutdown is idempotent
+
+**Rule:** Repeated shutdown or cleanup signals MUST NOT leak stage clients,
+process managers, or membership registrations.
+
+## Invariant namespace
+
+`STAGE-INV` reserves `001-099` for placement and ownership, `100-199` for
+replica identity/readiness/affinity, `200-299` for liveness loss, draining and
+cleanup, and `300-399` for local/distributed and upstream compatibility.
+Numbers become append-only after normative promotion.
+
+## Safe-change guide
+
+Exercise local and distributed startup, readiness, replica selection,
+affinity, membership changes, abort routing, process failure, draining, and
+repeated shutdown. Cross-stage policy changes belong in
+`engine_orchestration.md`.
+
+## Promotion gate
+
+- Reconcile terminology and ownership after #5441 reaches a final state.
+- Add or identify a dedicated StageRuntime/StagePool lifecycle suite; current
+  evidence is distributed across the cited tests.
+- Demonstrate affinity preservation for update, interaction, and abort paths.
+- Demonstrate membership loss and repeated shutdown without leaked clients or
+  processes.
+- Obtain approval from a technical owner and an independent validation
+  reviewer.

docs/design/module/vllm_omni_config.md

@@ -0,0 +1,56 @@
+---
+title: vLLM-Omni Configuration
+kind: module
+status: draft
+architecture_state: deferred-pending-refactor
+ownership_status: provisional
+owners:
+  - "@lishunyang12"
+  - "@alex-jw-brooks"
+primary_code_paths:
+  - vllm_omni/config/**
+  - vllm_omni/deploy/**
+  - vllm_omni/model_executor/stage_configs/**
+related_code_paths:
+  - vllm_omni/platforms/*/stage_configs/**
+depends_on: []
+validation_paths:
+  - tests/config/**
+upstream_refs:
+  - vllm.config
+last_reviewed: 2026-08-07
+last_verified_commit: 3d7fc3b9ba3cac88d579d4dc35b78b0b641675fc
+---
+
+# vLLM-Omni configuration
+
+This page intentionally records discovery paths only. Substantive
+configuration ownership and contract work is deferred until the active
+configuration refactoring has settled.
+
+## Deferred contract scope
+
+The current code and tests remain authoritative for parsing, defaults,
+overrides, deployment topology, stage construction, and runtime projection.
+This draft does not define a stable precedence model, environment-variable
+capture rule, canonical configuration object, or invariant namespace.
+
+The PR review proposal to capture environment-derived values in configuration
+objects at initialization is an open design question, not a current
+invariant.
+
+## Safe-change guide while deferred
+
+Changes should trace every affected producer to its runtime consumer and test
+the structured, legacy, CLI, and deployment paths that are actually supported.
+Do not infer a stable contract from this placeholder page.
+
+## Promotion gate
+
+- Finish or stabilize the configuration refactor and identify the canonical
+  runtime configuration representation.
+- Confirm technical owners, primary path exceptions, and validation owners.
+- Document parsing and override precedence, including when environment values
+  are captured.
+- Allocate an invariant namespace only after those behaviors are enforced and
+  owner-approved.

docs/user_guide/diffusion/attention_backends.md

@@ -14,15 +14,16 @@ The full set of backends and their platform defaults is in the **Backend Options
 
 | Value | Notes |
 |---|---|
-| `TRTLLM_ATTN` | FlashInfer's trtllm-gen FMHA (TensorRT-LLM's generated kernels, vendored by FlashInfer). BF16, GQA native, `head_dim=128`. Datacenter Blackwell only (sm_100 / sm_103). Supports **Skip-Softmax** sparse attention — see below. Requires `flashinfer`. |
+| `TRTLLM_ATTN` | FlashInfer's trtllm-gen FMHA (TensorRT-LLM's generated kernels, vendored by FlashInfer). Dense BF16, GQA native, `head_dim=128`. Datacenter Blackwell only (sm_100 / sm_103). Packed paths can provide `cu_seqlens` directly. Supports optional **Skip-Softmax** sparse attention — see below. Requires `flashinfer`. |
 | `FLASH_ATTN` | Wraps FlashAttention 4 on Blackwell when `flash-attn-4` is installed, then falls back to FlashAttention 3/2. Default on Hopper / Ada / Ampere when a compatible FlashAttention package is installed. |
 | `CUDNN_ATTN` | Pins `sdpa_kernel([CUDNN_ATTENTION])`. Default on Blackwell (sm_10x / sm_12x) with cuDNN ≥ 9.5. Wins on mask-heavy DiTs (HunyuanVideo-1.5: 2× e2e vs SDPA). |
-| `FLASHINFER_ATTN` | Calls FlashInfer's dense `single_prefill_with_kv_cache` directly with `custom_mask` for non-causal masked attention. Used as Blackwell fallback when cuDNN is unavailable. Requires `flashinfer`. |
+| `FLASHINFER_ATTN` | Uses FlashInfer's batch-prefill wrapper and supports mixed Q/K and V input dtypes through backend-specific configuration. Used as Blackwell fallback when cuDNN is unavailable. Requires `flashinfer` >= 0.6.16rc1 for mixed-dtype configurations. |
 | `TORCH_SDPA` | PyTorch `scaled_dot_product_attention` with the default backend dispatcher. Most conservative; always available. |
 | `SAGE_ATTN` | SageAttention 2.2 — INT8-quantized attention with FP16 accumulation. Lossy but typically visually indistinguishable on diffusion outputs. Requires `sageattention`. |
 | `SAGE_ATTN_3` | Requires `sageattn3` from `SageAttention/sageattention3_blackwell`. CUDA only, intended for Blackwell GPUs, with GQA/MQA requests falling back to PyTorch SDPA. |
 | `FLASH_ATTN_HUB` | FlashAttention 2 from HuggingFace `kernels` library. Useful for train/rollout alignment. |
 | `FLASH_ATTN_3_HUB` | FlashAttention 3 from HuggingFace `kernels` library. CUDA Hopper (sm_90+) only; falls back to `FLASH_ATTN_HUB` on older GPUs. |
+| `RAINFUSION_ATTN` | MindIE-SD **RainFusion** block-sparse video attention — see [below](#rainfusion_attn-backend-and-block-sparse-video-attention). Ascend NPU only; requires `mindiesd`. Delegates to `FLASH_ATTN` for anything that is not a packed video sequence. |
 
 
 ## Configuration
@@ -69,8 +70,9 @@ vllm-omni serve <model> \
     --diffusion-attention-config '{"default":{"backend":"FLASH_ATTN"},"per_role":{"cross":{"backend":"TORCH_SDPA"}}}'
 ```
 
-A backend that needs configuration exposes it as a typed field on the spec. Today the only
-one is `TRTLLM_ATTN`'s Skip-Softmax (see [below](#trtllm_attn-backend-and-skip-softmax)):
+A backend that needs configuration exposes it as a typed field on the spec:
+`TRTLLM_ATTN`'s Skip-Softmax (see [below](#trtllm_attn-backend-and-skip-softmax)) and
+`RAINFUSION_ATTN`'s RainFusion block (see [below](#rainfusion_attn-backend-and-block-sparse-video-attention)):
 
 ```bash
 --diffusion-attention-config.default.backend TRTLLM_ATTN \
@@ -82,7 +84,7 @@ one is `TRTLLM_ATTN`'s Skip-Softmax (see [below](#trtllm_attn-backend-and-skip-s
 When constructing `OmniDiffusionConfig` directly:
 
 ```python
-from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, OmniDiffusionConfig
+from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, AttnQuantSpec, OmniDiffusionConfig
 
 config = OmniDiffusionConfig(
     diffusion_attention_config=AttentionConfig(
@@ -97,13 +99,34 @@ config = OmniDiffusionConfig(
 
 A plain dict is also accepted and normalized to `AttentionConfig`.
 
+#### FlashInfer QK16/V8 quantized attention
+
+`FLASHINFER_ATTN` accepts quantization options through `AttentionSpec.quant`. For QK16/V8:
+
+```python
+config = OmniDiffusionConfig(
+    diffusion_attention_config=AttentionConfig(
+        default=AttentionSpec(
+            backend="FLASHINFER_ATTN",
+            quant=AttnQuantSpec(
+                dtype_qk="bfloat16",
+                dtype_vo="fp8_e4m3",
+            ),
+        ),
+    ),
+    ...,
+)
+```
+
+`dtype_qk` configures the Q and K input dtypes. `dtype_vo` configures the V input dtype.
+
 ## Platform Defaults
 
 ### Blackwell (sm_100 / sm_103 / sm_120 / sm_121)
 
 Auto-route preference, in order:
 
-1. `TRTLLM_ATTN` — on **datacenter** Blackwell (sm_100 / sm_103) when `flashinfer` is installed, the model's `head_dim` is 128, **and the model is mask-free** (see below)
+1. `TRTLLM_ATTN` — on **datacenter** Blackwell (sm_100 / sm_103) when `flashinfer` is installed, the model's `head_dim` is 128, **and the model declares a compatible packed/mask-free path**
 2. `CUDNN_ATTN` — when cuDNN ≥ 9.5 is available (ships in PyTorch 2.5+ wheels)
 3. `FLASHINFER_ATTN` — when `flashinfer` is installed but cuDNN < 9.5
 4. `FLASH_ATTN` — when `flash-attn` is installed with the Blackwell CUTE kernel
@@ -111,7 +134,7 @@ Auto-route preference, in order:
 
 `TRTLLM_ATTN` is skipped on workstation Blackwell (sm_120 / sm_121) and for any `head_dim != 128`, so those GPUs keep the `CUDNN_ATTN` route described below.
 
-`TRTLLM_ATTN` outranks `CUDNN_ATTN` on datacenter Blackwell because it is the only backend that can enable Skip-Softmax, not because its dense kernel is faster — dense, the two are comparable. It cannot honor an attention mask, so the auto-default only applies to pipelines verified mask-free (the Wan family); mask-using pipelines keep `CUDNN_ATTN`. `TRTLLM_ATTN` stays available everywhere via explicit `--diffusion-attention-backend TRTLLM_ATTN` (which raises clearly if a mask is received).
+`TRTLLM_ATTN` outranks `CUDNN_ATTN` on datacenter Blackwell for compatible packed/mask-free pipelines. Workstation Blackwell (sm_120 / sm_121) and pipelines that require attention masks retain their normal fallback.
 
 The startup log line `Defaulting to diffusion attention backend CUDNN_ATTN (Blackwell sm_120, cuDNN 91002)` confirms the route.
 
@@ -204,11 +227,80 @@ AttentionConfig(
 )
 ```
 
-
 **Requirements.** Needs FlashInfer ≥ 0.6.16rc1. Kernel
 availability is arch-dependent: `fp8_e4m3` QK has kernels on both **SM100** (B200) and **SM103**
 (B300); `int8` QK kernels are compiled for **SM100 only**.
 
+## RAINFUSION_ATTN Backend and Block-Sparse Video Attention
+
+`RAINFUSION_ATTN` runs MindIE-SD's RainFusion (`rf_v2`) kernel on Ascend NPU. It pools each
+128-token block of keys, ranks them per query block, and computes attention against only the
+top-scoring ones. Video tokens are rearranged into `(t, h, w)` order first, so the kept blocks are
+spatiotemporal neighbours rather than arbitrary rows of the packed sequence.
+
+Only the video segment is sparsified. The prefix rows (text, visual conditions, audio) and the
+first-frame blocks are always kept dense, which is why the realized sparsity is lower than the
+nominal `sparsity` you configure. Anything the kernel cannot handle — a warmup denoise step, an
+exempt layer, a sequence with no published video geometry, a video segment under 32 blocks —
+delegates to `FLASH_ATTN`, so a model can select this backend unconditionally.
+
+Enable it through the typed `block_sparse` block on the attention spec, shared by every
+block-sparse backend:
+
+| Key | Valid values | Meaning |
+|---|---|---|
+| `sparsity` | finite, `[0, 1]` | Nominal fraction of key blocks dropped per query block. `0` disables sparsity. Defaults to `0.8`. |
+| `start_step` | `≥ 0` | Keep the first N denoise steps dense. Layout is decided early, so these steps dominate structural fidelity. |
+| `skip_layers` | index selector, e.g. `"0-3,38"` | DiT blocks that always stay dense. |
+
+```bash
+vllm-omni serve MiniMaxAI/MiniMax-H3 \
+  --diffusion-attention-config '{"default": {"backend": "RAINFUSION_ATTN",
+      "block_sparse": {"sparsity": 0.8, "start_step": 0}}}'
+```
+
+Programmatically the same block is a typed `BlockSparseSpec` (values validated at construction):
+
+```python
+from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, BlockSparseSpec
+
+AttentionConfig(
+    default=AttentionSpec(
+        backend="RAINFUSION_ATTN",
+        block_sparse=BlockSparseSpec(sparsity=0.8, start_step=0, skip_layers="0-1"),
+    ),
+)
+```
+
+**Tune in the order `start_step` → `sparsity` → `skip_layers`.** Raise `start_step` first: it is the
+cheapest way to recover structure, because the early high-noise steps decide global layout while the
+later steps only refine texture. Then walk `sparsity` up to your quality bar. Reach for
+`skip_layers` last, once an A/B against dense at the same seed points at specific blocks.
+
+Requires Ascend NPU with `mindiesd`; selecting it on any other platform raises. It is also
+incompatible with ring sequence parallelism, since `rf_v2` needs the whole key sequence to rank
+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
+is `latent_t × (height / 32) × (width / 32)`. Otherwise `rf_v2`'s block mask and the kernel's own
+tiling disagree on the block count, and the block straddling the seam mixes video and prefix rows
+that selection may then drop. Any resolution still runs — an unaligned geometry falls back to dense
+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
+spatial tiles (two tiles fill one 128-row block), which is what makes a selected block a compact
+patch of the frame. When the latent `h` or `w` is not a multiple of 8, the leftover rows or columns
+are peeled off and appended as a flat run instead of being tiled, so they get pooled with spatially
+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
+resolutions off that grid still run sparse; they just lose more fidelity at the same `sparsity`,
+which you can buy back with `start_step`.
+
 ## End-to-End Benchmark (BF16, sm_120 RTX Pro 6000 Blackwell)
 
 Same prompt and seed across runs. `Total generation time` from `text_to_video.py` / `text_to_image.py`.

docs/user_guide/diffusion/cache_acceleration/cache_dit.md

@@ -6,6 +6,7 @@
 - [Overview](#overview)
 - [Quick Start](#quick-start)
 - [Example Script](#example-script)
+- [Request-Scoped Quality](#request-scoped-quality-minimax-h3)
 - [Acceleration Methods](#acceleration-methods)
 - [Configuration Parameters](#configuration-parameters)
 - [Best Practices](#best-practices)
@@ -140,6 +141,43 @@ vllm serve Qwen/Qwen-Image --omni --port 8091 \
   --cache-config '{"Fn_compute_blocks": 1, "residual_diff_threshold": 0.12}'
 ```
 
+## Request-Scoped Quality (MiniMax H3)
+
+MiniMax H3 can switch Cache-DiT at the request boundary without enabling a
+cache backend at server startup. Add a quality field to an online request:
+
+```bash
+# Native reference path for this request.
+-F 'quality=lossless'
+
+# H3's conservative Cache-DiT profile for this request.
+-F 'quality=high'
+```
+
+For offline inference, set the same field on the sampling parameters:
+
+```python
+OmniDiffusionSamplingParams(
+    quality="high",
+    num_inference_steps=50,
+    extra_args={"task": "t2va", "duration": 5.0, "audio_flow_shift": 3.0},
+)
+```
+
+| Server startup | Request quality | Behavior |
+|---|---|---|
+| `--cache-backend cache_dit` | omitted | Use or restore the startup Cache-DiT profile |
+| `--cache-backend cache_dit` | `lossless` | Remove Cache-DiT hooks and run the reference path |
+| `--cache-backend cache_dit` | `high` | Use or install H3's conservative Cache-DiT profile |
+| no Cache-DiT | omitted or `lossless` | Run the reference path |
+| no Cache-DiT | `high` | Install H3's conservative Cache-DiT profile |
+
+The startup option is only needed when omitted-quality requests should use the
+server-configured Cache-DiT profile by default.
+
+See the [MiniMax H3 recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/MiniMaxAI/MiniMax-H3.md#request-scoped-quality)
+for a complete request and measured trade-off.
+
 ---
 
 ## Acceleration Methods

docs/user_guide/diffusion/cpu_offload.md

@@ -184,10 +184,15 @@ class Flux2Transformer2DModel(nn.Module):
 ### How It Works
 
 Distributed layerwise offloading extends single-GPU layerwise offloading to
-multi-device deployments.  Each DP rank stores only **1/dp_size** of the model
-weights in host memory; full layer weights are reconstructed at runtime via
-**AllGather** on a dedicated communication stream, overlapped with computation
-via a fixed double-buffer scheme.
+multi-device deployments. In the default mode, each rank stores only
+**1/dp_size** of the model weights in host memory; full layer weights are
+reconstructed at runtime via **AllGather** on a dedicated communication stream,
+overlapped with computation via a fixed double-buffer scheme. The
+`--dlo-no-use-allgather` mode instead streams standard-loader rank-local
+weights independently and does not shard weights across ranks.
+
+For the implementation design and compatibility matrix, see
+[Distributed Layerwise Offload](../../design/feature/distributed_layerwise_offload.md).
 
 **Key features:**
 - **Weight sharding + AllGather**: each rank stores 1/dp_size of weights,
@@ -219,7 +224,7 @@ vllm serve /path/to/model --omni \
   --enable-distributed-layerwise-offload \
   --data-parallel-size 4
 
-# Without AllGather (each rank loads full weights, no sharding)
+# Without DLO AllGather (each rank streams standard-loader rank-local weights)
 vllm serve /path/to/model --omni \
   --enable-distributed-layerwise-offload \
   --data-parallel-size 4 \
@@ -247,9 +252,9 @@ m = Omni(
 | Flag | Description | Default |
 |------|-------------|---------|
 | `--enable-distributed-layerwise-offload` | Enable DLO | `false` |
-| `--data-parallel-size N` | Number of DP ranks for weight sharding | `1` |
+| `--data-parallel-size N` | Number of DP ranks; DLO weight-sharding group when AllGather is enabled | `1` |
 | `--dlo-use-allgather` | Shard + AllGather (saves CPU, requires concurrent requests) | `true` |
-| `--dlo-no-use-allgather` | Full weights per rank (no sharding, no AllGather) | `false` |
+| `--dlo-no-use-allgather` | Use standard-loader rank-local weights; no DLO sharding or AllGather | `false` |
 
 ### How model weights are loaded (mmap path)
 
@@ -303,9 +308,15 @@ request while AllGather synchronizes only weight shards (request-independent).
 
 ### Limitations
 - Online quantization (FP8) is incompatible with mmap loading — falls back to
-  regular `load_weights()` automatically
-- Tensor Parallel is not supported (DLO uses DP-based sharding)
-- HSDP + AllGather is rejected (would double-shard weights)
+  regular loading only when the model is outside the DLO+AllGather mmap path;
+  online quantization with DLO+AllGather is rejected. Use
+  `--dlo-no-use-allgather` or disable online quantization.
+- Tensor Parallel > 1 is rejected in the DLO+AllGather mmap path because TP-aware
+  weight-loader callbacks are bypassed. The no-AllGather path retains the
+  standard TP loader and is experimental.
+- HSDP + AllGather is rejected because it would double-shard parameters. HSDP
+  with no-AllGather is accepted at configuration level, but has limited
+  end-to-end validation.
 - `num_inference_steps=None` is not allowed in DP multi-concurrency mode
 
 **Module Discovery**
@@ -338,7 +349,7 @@ OffloadBackend (base class)
 ├── LayerWiseOffloadBackend → uses LayerwiseOffloadHook
 │                          (single-GPU, full weights on host)
 └── DistributedLayerwiseOffloadBackend → uses DistributedLayerwiseOffloadHook
-                                         (multi-GPU, 1/dp_size sharded weights + AllGather)
+                                         (multi-GPU, sharded AllGather or rank-local mode)
 ```
 
 Factory function `get_offload_backend()` selects the appropriate backend based on
@@ -369,4 +380,10 @@ reasoner/generator components inside the model forward pass.
 **Notes:**
 - Model-Level Offloading is expected to be supported by all common diffusion models (DiT and encoders) naturally
 - Layerwise Offloading requires DiT class to define `_layerwise_offload_blocks_attrs` pointing to transformer blocks
-- Distributed Layerwise Offloading works with any model that supports Layerwise Offloading — no additional model changes required.  See [Cosmos3 DistOffload recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-DistOffload.md) for usage examples.
+- Distributed Layerwise Offloading reuses the layerwise block topology, but the
+  model and parallelism combination still needs validation. AllGather mode
+  additionally requires a compatible checkpoint/mmap layout; no-AllGather mode
+  retains the standard loader's rank-local tensors. See the
+  [DLO design](../../design/feature/distributed_layerwise_offload.md) and
+  [Cosmos3 DistOffload recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-DistOffload.md)
+  for usage examples.

docs/user_guide/diffusion_features.md

@@ -59,7 +59,7 @@ Memory optimization methods help reduce GPU memory usage, enabling inference on
 
 | Method | Description | Best For |
 |--------|-------------|----------|
-| **[CPU Offload](diffusion/cpu_offload_diffusion.md)** | Offloads model components to CPU memory | Limited VRAM, large models on consumer GPUs |
+| **[CPU Offload](diffusion/cpu_offload.md)** | Offloads model components to CPU memory | Limited VRAM, large models on consumer GPUs |
 | **[Quantization](quantization/overview.md)** | Reduces transformer stages from BF16 to FP8/INT8/etc. | Limited VRAM, minimal accuracy loss    |
 | **[VAE Parallelism](diffusion/parallelism/vae_parallelism.md)** | Distributes VAE decode work across GPUs | High-resolution generation with reduced VAE memory peak |
 
@@ -104,8 +104,8 @@ The following tables show which models support each feature:
 
 > Notes:
 
-> 1. CPU Offload has two methods: Module-wise (default for models with DiT + text encoder) and Layerwise. The tables below show **Layerwise support** only. Split models like Cosmos3 (no separate text encoder) swap their reasoner/generator components for module-wise offload; see the [CPU Offload Guide](diffusion/cpu_offload_diffusion.md).
-> 2. The **💾Quantization** column is collapsed for readability. See [Quantization Overview](quantization/overview.md) for per-method and per-model support details.
+> 1. CPU Offload has two methods: Module-wise (default for models with DiT + text encoder) and Layerwise. The tables below show **Layerwise support** only. Split models like Cosmos3 (no separate text encoder) swap their reasoner/generator components for module-wise offload; see the [CPU Offload Guide](diffusion/cpu_offload.md).
+> 2. The **💾Quantization** column is collapsed for readability. See [Quantization](quantization/overview.md) for per-method and per-model support details.
 
 ### ImageGen
 
@@ -274,9 +274,9 @@ Measured on NVIDIA H800:
 
 **Memory Optimization:**
 
-- **[CPU Offload Guide](diffusion/cpu_offload_diffusion.md)** - Offload model components to CPU, reduce GPU memory usage
+- **[CPU Offload Guide](diffusion/cpu_offload.md)** - Offload model components to CPU, reduce GPU memory usage
 - **[VAE Parallelism Guide](diffusion/parallelism/vae_parallelism.md)** - Distribute VAE decode work across GPUs for high-resolution images and videos
-- **[Quantization Overview](quantization/overview.md)** - Overview of quantization methods for diffusion, multi-stage omni/TTS, and multi-stage diffusion models
+- **[Quantization](quantization/overview.md)** - Quantization methods for diffusion, multi-stage omni/TTS, and multi-stage diffusion models
 
 **Extensions:**
 

docs/user_guide/quantization/overview.md

@@ -5,6 +5,9 @@ path. The same configuration entrypoint is used across diffusion-only models,
 multi-stage omni/TTS models, and multi-stage diffusion models, but each model
 type has a different quantization scope.
 
+For the internal architecture and backend extension points, see the
+[quantization design document](../../design/feature/quantization.md).
+
 ## Quantization Modes
 
 | Mode | Guide | Description | Methods |

examples/offline_inference/minimax_h3/dlo_dp2_lifecycle.py

@@ -0,0 +1,240 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Reproduce MiniMax-H3 request-mode and DLO/DP2 lifecycle behavior.
+
+Examples:
+
+    VLLM_WORKER_MULTIPROC_METHOD=spawn \
+    VLLM_OMNI_VIDEO_SYNC_TIMEOUT=14400 \
+    VLLM_OMNI_DLO_DP_WAVE_TIMEOUT=600 \
+    CUDA_VISIBLE_DEVICES=0,1 \
+    python examples/offline_inference/minimax_h3/dlo_dp2_lifecycle.py \
+        --model /path/to/MiniMax-H3/FL2VA --mode dlo-dp2
+
+    VLLM_WORKER_MULTIPROC_METHOD=spawn \
+    VLLM_OMNI_VIDEO_SYNC_TIMEOUT=14400 \
+    CUDA_VISIBLE_DEVICES=0,1 \
+    python examples/offline_inference/minimax_h3/dlo_dp2_lifecycle.py \
+        --model /path/to/MiniMax-H3/FL2VA --mode request
+"""
+
+from __future__ import annotations
+
+import argparse
+import asyncio
+import copy
+import json
+import multiprocessing
+import time
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+
+from vllm_omni.entrypoints.async_omni import AsyncOmni
+
+DEFAULT_PROMPTS = (
+    "At night, three cats march into a bedroom playing tiny brass instruments, "
+    "then abruptly file out, with synchronized room ambience.",
+    "A paper boat crosses a rain-filled street while distant traffic and water sounds remain synchronized.",
+)
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--model", required=True, help="Path to MiniMax-H3/FL2VA")
+    parser.add_argument(
+        "--mode",
+        choices=("dlo-dp2", "request"),
+        required=True,
+        help="DLO AllGather DP2 or ordinary non-DLO TP2 request mode",
+    )
+    parser.add_argument("--steps", type=int, default=2)
+    parser.add_argument("--duration", type=float, default=5.0)
+    parser.add_argument("--width", type=int, default=1344)
+    parser.add_argument("--height", type=int, default=768)
+    parser.add_argument("--batch-wait-ms", type=float, default=500.0)
+    parser.add_argument("--init-timeout", type=float, default=1800.0)
+    parser.add_argument("--output", type=Path)
+    return parser.parse_args()
+
+
+def engine_kwargs(args: argparse.Namespace) -> dict[str, Any]:
+    common: dict[str, Any] = {
+        "model": args.model,
+        "trust_remote_code": True,
+        "num_gpus": 2,
+        "usp": 1,
+        "ring": 1,
+        "vae_parallel_mode": "tile",
+        "vae_use_tiling": True,
+        "diffusion_attention_backend": "CUDNN_ATTN",
+        "request_batch_max_wait_ms": args.batch_wait_ms,
+        "enforce_eager": True,
+        "stage_init_timeout": args.init_timeout,
+        "init_timeout": args.init_timeout,
+    }
+    if args.mode == "dlo-dp2":
+        common.update(
+            tensor_parallel_size=1,
+            data_parallel_size=2,
+            text_encoder_tp_size=1,
+            vae_patch_parallel_size=1,
+            enable_distributed_layerwise_offload=True,
+            dlo_use_allgather=True,
+            dlo_resident_layers=0,
+        )
+    else:
+        common.update(
+            tensor_parallel_size=2,
+            data_parallel_size=1,
+            text_encoder_tp_size=2,
+            vae_patch_parallel_size=2,
+            enable_distributed_layerwise_offload=False,
+        )
+    return common
+
+
+def sampling_params(
+    engine: AsyncOmni,
+    args: argparse.Namespace,
+    seed: int,
+) -> list[Any]:
+    params = copy.deepcopy(engine.default_sampling_params_list)
+    diffusion = params[0]
+    diffusion.width = args.width
+    diffusion.height = args.height
+    diffusion.fps = 24
+    diffusion.num_inference_steps = args.steps
+    diffusion.seed = seed
+    diffusion.extra_args = {
+        "task": "t2va",
+        "duration": args.duration,
+        "aspect_ratio": "16:9",
+        "flow_shift": 12.0,
+        "audio_flow_shift": 3.0,
+    }
+    return params
+
+
+async def generate_one(
+    engine: AsyncOmni,
+    args: argparse.Namespace,
+    *,
+    request_id: str,
+    prompt: str,
+    seed: int,
+) -> Any:
+    final_output = None
+    async for output in engine.generate(
+        prompt=prompt,
+        request_id=request_id,
+        sampling_params_list=sampling_params(engine, args, seed),
+    ):
+        if output.finished:
+            final_output = output
+    if final_output is None:
+        raise RuntimeError(f"{request_id} finished without an output")
+    return final_output
+
+
+def output_summary(output: Any, args: argparse.Namespace) -> dict[str, Any]:
+    if not output.images:
+        raise RuntimeError(f"{output.request_id} returned no video")
+    frames = np.asarray(output.images[0])
+    audio = np.asarray(output.multimodal_output.get("audio"))
+    if frames.ndim != 4 or tuple(frames.shape[1:]) != (
+        args.height,
+        args.width,
+        3,
+    ):
+        raise RuntimeError(f"{output.request_id} returned invalid video shape {tuple(frames.shape)}")
+    if audio.ndim != 3 or tuple(audio.shape[:2]) != (1, 2):
+        raise RuntimeError(f"{output.request_id} returned invalid audio shape {tuple(audio.shape)}")
+    if args.duration == 5.0 and args.height == 768 and args.width == 1344:
+        if tuple(frames.shape) != (124, 768, 1344, 3) or tuple(audio.shape) != (1, 2, 165600):
+            raise RuntimeError(
+                f"{output.request_id} default shape mismatch: video={tuple(frames.shape)}, audio={tuple(audio.shape)}"
+            )
+    return {
+        "request_id": output.request_id,
+        "frames_shape": list(frames.shape),
+        "audio_shape": list(audio.shape),
+        "peak_memory_mb": output.peak_memory_mb,
+        "stage_durations": output.stage_durations,
+    }
+
+
+async def run(args: argparse.Namespace) -> dict[str, Any]:
+    engine = AsyncOmni(**engine_kwargs(args))
+    summary: dict[str, Any] = {
+        "mode": args.mode,
+        "engine_kwargs": engine_kwargs(args),
+    }
+    try:
+        if args.mode == "dlo-dp2":
+            started = time.perf_counter()
+            asymmetric = await asyncio.gather(
+                generate_one(
+                    engine,
+                    args,
+                    request_id="invalid-empty",
+                    prompt="",
+                    seed=1001,
+                ),
+                generate_one(
+                    engine,
+                    args,
+                    request_id="invalid-peer",
+                    prompt=DEFAULT_PROMPTS[0],
+                    seed=1002,
+                ),
+                return_exceptions=True,
+            )
+            summary["asymmetric_wave_s"] = time.perf_counter() - started
+            summary["asymmetric_errors"] = [
+                f"{type(result).__name__}: {result}" for result in asymmetric if isinstance(result, BaseException)
+            ]
+            if len(summary["asymmetric_errors"]) != 2:
+                raise RuntimeError("Both requests in the asymmetric DP wave must fail before dispatch")
+
+        started = time.perf_counter()
+        request_count = 2 if args.mode == "dlo-dp2" else 1
+        outputs = await asyncio.gather(
+            *(
+                generate_one(
+                    engine,
+                    args,
+                    request_id=f"recovery-{index}",
+                    prompt=DEFAULT_PROMPTS[index],
+                    seed=2000 + index,
+                )
+                for index in range(request_count)
+            )
+        )
+        summary["valid_wave_s"] = time.perf_counter() - started
+        summary["outputs"] = [output_summary(output, args) for output in outputs]
+    finally:
+        started = time.perf_counter()
+        engine.close()
+        summary["shutdown_s"] = time.perf_counter() - started
+
+    children = multiprocessing.active_children()
+    summary["active_children"] = [{"name": child.name, "pid": child.pid} for child in children]
+    if children:
+        raise RuntimeError(f"Worker processes remain after shutdown: {children}")
+    return summary
+
+
+def main() -> None:
+    args = parse_args()
+    summary = asyncio.run(run(args))
+    rendered = json.dumps(summary, indent=2, sort_keys=True)
+    print(f"E2E_RESULT {rendered}", flush=True)
+    if args.output is not None:
+        args.output.parent.mkdir(parents=True, exist_ok=True)
+        args.output.write_text(rendered + "\n", encoding="utf-8")
+
+
+if __name__ == "__main__":
+    main()

mkdocs.yml

@@ -75,6 +75,8 @@ plugins:
       redirect_maps:
         user_guide/diffusion/request_batching.md: user_guide/diffusion/execution_modes.md
         user_guide/diffusion/step_execution.md: user_guide/diffusion/execution_modes.md
+        user_guide/diffusion/cpu_offload_diffusion.md: user_guide/diffusion/cpu_offload.md
+        user_guide/quantization/index.md: user_guide/quantization/overview.md
         design/feature/diffusion_request_level_batching.md: design/feature/diffusion_continuous_batching.md
         design/feature/diffusion_step_execution.md: design/feature/diffusion_continuous_batching.md
   - glightbox

pyproject.toml

@@ -250,13 +250,16 @@ markers = [
     "npu: Tests that run on NPU/Ascend (auto-added)",
     "musa: Tests that run on MUSA/Moore Threads (auto-added)",
     # specified computation resources marks (auto-added)
-    "H100: Tests that require H100 GPU",
-    "L4: Tests that require L4 GPU",
-    "MI325: Tests that require MI325 GPU (AMD/ROCm)",
-    "B60: Tests that require Intel Arc Pro B60 XPU",
-    "S5000: Tests that require S5000 GPU (Moore Threads/MUSA)",
-    "A2: Tests that require A2 NPU",
-    "A3: Tests that require A3 NPU",
+    # Tag with [hardware-resource] for tests.helpers.mark.get_hardware_mark_list().
+    "H100: [hardware-resource] Tests that require H100 GPU",
+    "H200: [hardware-resource] Tests that require H200 GPU",
+    "L4: [hardware-resource] Tests that require L4 GPU",
+    "B200: [hardware-resource] Tests that require B200 GPU",
+    "MI325: [hardware-resource] Tests that require MI325 GPU (AMD/ROCm)",
+    "B60: [hardware-resource] Tests that require Intel Arc Pro B60 XPU",
+    "S5000: [hardware-resource] Tests that require S5000 GPU (Moore Threads/MUSA)",
+    "A2: [hardware-resource] Tests that require A2 NPU",
+    "A3: [hardware-resource] Tests that require A3 NPU",
     "distributed_cuda: Tests that require multi cards on CUDA platform",
     "distributed_rocm: Tests that require multi cards on ROCm platform",
     "distributed_xpu: Tests that require multi cards on XPU platform",

recipes/MiniMaxAI/MiniMax-H3-MUSA.md

@@ -12,19 +12,17 @@
 - Maintainer: Community
 
 This recipe adapts [MiniMax-H3.md](MiniMax-H3.md) for MUSA environments.
-MiniMax H3 has two independently served checkpoint partitions:
-
-- `FL2VA` serves text-to-video+audio (`t2va`) and
-  first-frame-to-video+audio (`fl2va`).
-- `Ref2VA` serves image/audio or video-reference generation (`ref2va`).
+The modular H3 pipeline can load both task-specific DiTs, but this MUSA recipe
+uses a **single task partition**: `FL2VA` for `t2va`/`fl2va`, or `Ref2VA` for
+`ref2va`. This avoids loading both DiTs.
 
 ## Prerequisites
 
 ### Checkpoint
 
-Only download the partition needed by the server. Downloading the whole model
-repository also retrieves duplicate shared components and the original
-top-level pipeline layout.
+Download the partition required by the selected startup task. The commands below
+run a single-task server; do not use the modular default (combined) mode on MUSA
+until its two-DiT memory and performance profile is validated.
 
 For T2VA and FL2VA:
 
@@ -66,14 +64,13 @@ soundfile. Formats that libsndfile cannot read are demuxed through ffmpeg.
 
 ## Start a server
 
-One server loads one checkpoint partition. Set `MODEL` to `FL2VA` for T2VA
-and FL2VA, or to `Ref2VA` for Ref2VA.
+Use the partition path and set `--task-type` explicitly: `fl2va` for T2VA/FL2VA
+or `ref2va` for Ref2VA.
 
-The validated Ref2VA configuration uses tensor parallelism across four MTT
-S5000 GPUs and offloads inactive model components to CPU:
+The validated Ref2VA configuration uses tensor parallelism across four MTT S5000 GPUs and offloads inactive model components to CPU:
 
 ```bash
-export MODEL="${MODEL_ROOT}/FL2VA"
+export MODEL="${MODEL_ROOT}/Ref2VA"
 export PORT=8091
 
 MUSA_VISIBLE_DEVICES=0,1,2,3 \
@@ -85,6 +82,7 @@ vllm serve "${MODEL}" \
   --port "${PORT}" \
   --trust-remote-code \
   --num-gpus 4 \
+  --task-type ref2va \
   --tensor-parallel-size 4 \
   --vae-patch-parallel-size 4 \
   --enable-cpu-offload \
@@ -133,7 +131,8 @@ Keep `--vae-use-tiling` enabled for this serving profile.
 - Keep Ring Attention at degree 1. The current Ring path does not preserve
   MiniMax H3 packed padding boundaries.
 - H3 is CFG-distilled, so `--cfg-parallel-size` must remain 1.
-- Each server process loads only one checkpoint partition.
+- This MUSA recipe uses an explicitly selected single task partition; combined
+  serving loads both DiTs.
 - H3 currently executes one generation request per diffusion batch.
 - FP8 quantization has not been enabled for MiniMax H3.
 - MP3, M4A, MP4, and reference-video audio fallback requires `ffmpeg` on

recipes/MiniMaxAI/MiniMax-H3-NPU.md

@@ -26,12 +26,12 @@ environments. Differences from the GPU path:
 
 ### Checkpoint
 
-Same as the GPU recipe — Hugging Face access approval is required:
+Same as the GPU recipe — Hugging Face access approval is required. Authenticate
+once; `vllm serve` downloads the required nested components automatically:
 
 ```bash
 hf auth login
-export MODEL_ROOT=/path/to/MiniMax-H3
-hf download MiniMaxAI/MiniMax-H3 --local-dir "${MODEL_ROOT}"
+export MODEL=MiniMaxAI/MiniMax-H3
 ```
 
 ### Environment
@@ -50,7 +50,8 @@ uv pip install -e .
 ```
 
 Install the **mindie-sd** operator library to enable Ascend-optimized fused
-operators (`adalayernorm`, etc.):
+operators (`adalayernorm`, etc.) and the RainFusion `rf_v2` block-sparse
+attention kernel:
 
 ```bash
 git clone https://gitcode.com/Ascend/MindIE-SD.git && cd MindIE-SD
@@ -72,24 +73,21 @@ pip install mindiesd-*.whl
 
 ## Start a server
 
-One server loads one checkpoint partition. Set `MODEL` to `FL2VA` for T2VA
-and FL2VA requests, or to `Ref2VA` for Ref2VA requests.
+Pass the repository ID directly. The pipeline loads the two nested DiTs while
+sharing the tokenizer, processor, text encoder, and VAEs from `FL2VA`.
 
-### Multi-NPU: 768P validated configuration
+### Multi-NPU: 768P combined-service configuration
 
-Validated on eight NPUs of an Atlas 800I A3 server with Ulysses sequence
-parallelism degree 8, text-encoder tensor parallelism degree 8, native tiled
-VAE patch parallelism degree 8, and layerwise offload:
+Use Ulysses sequence parallelism degree 8, text-encoder tensor parallelism
+degree 8, native tiled VAE patch parallelism degree 8, and layerwise offload:
 
 ```bash
 export ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
 export PORT=9098
-export MODEL="${MODEL_ROOT}/FL2VA"
+export MODEL=MiniMaxAI/MiniMax-H3
 export VLLM_WORKER_MULTIPROC_METHOD=spawn
 export VLLM_OMNI_VIDEO_SYNC_TIMEOUT=1800
 export PYTHONDONTWRITEBYTECODE=1
-export HF_HUB_OFFLINE=1
-export TRANSFORMERS_OFFLINE=1
 
 vllm serve "${MODEL}" \
   --omni \
@@ -111,12 +109,31 @@ Do not add `--enforce-eager`. The first request includes regional
 compilation; warm the server once before measuring steady-state latency.
 H3 is CFG-distilled, so `--cfg-parallel-size` must remain 1.
 
-To serve Ref2VA, stop the FL2VA server and restart with:
+The same endpoint accepts `task=t2va`, `task=fl2va`, and `task=ref2va`; no
+partition restart is required. Layerwise offload applies to both DiTs.
+
+### Optional optimizations
+
+Two independent optimizations may be enabled on top of the configuration
+above. Both are validated for T2VA only.
+
+**RainFusion block-sparse attention** — switch the attention backend, keeping
+every other flag unchanged:
 
 ```bash
-export MODEL="${MODEL_ROOT}/Ref2VA"
+  --diffusion-attention-backend RAINFUSION_ATTN
 ```
 
+**INT8 online quantization** — add one flag to the server command above:
+
+```bash
+  --quantization int8
+```
+
+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.
+
 ## HTTP API examples
 
 The request format is identical to the GPU recipe; see
@@ -145,7 +162,8 @@ throughput guarantee.
 
 ## Known limitations
 
-- Each server process loads only one checkpoint partition.
+- Combined serving requires sibling `FL2VA` and `Ref2VA` directories, loads
+  both task-specific DiTs, and loads one copy of every shared component.
 - H3 currently executes one generation request per diffusion batch.
 - The first regional-compile request is a warmup and should not be included
   in steady-state performance measurements.
@@ -154,9 +172,16 @@ throughput guarantee.
   the same HTTP request contract.
 - VAE patch parallelism requires size 1 or the full DiT group size and
   supports the H3 native `tile` mode only.
+- RainFusion block-sparse attention and INT8 quantization are validated for
+  T2VA only; use the BF16 dense configuration for FL2VA and Ref2VA.
+- Online quantization cannot be combined with distributed layerwise offload
+  while AllGather is enabled; pass `--dlo-no-use-allgather` in that case.
 
 ## Additional resources
 
 - [MiniMax-H3.md](MiniMax-H3.md) — full GPU guide
+- [Attention backends § RAINFUSION_ATTN](../../docs/user_guide/diffusion/attention_backends.md#rainfusion_attn-backend-and-block-sparse-video-attention)
+  — RainFusion knobs and tuning
+- [Int8 quantization](../../docs/user_guide/quantization/int8.md)
 - [Supported models](../../docs/models/supported_models.md)
 - [Video API](../../docs/serving/videos_api.md)

recipes/MiniMaxAI/MiniMax-H3.md

@@ -12,27 +12,34 @@
 - Maintainer: Community
 
 MiniMax H3 is a CFG-distilled joint video/audio diffusion transformer. Its
-checkpoint has two independently served partitions:
+checkpoint has two task-specific DiT partitions:
 
 - `FL2VA`: text-to-video+audio (`t2va`) and first-frame-to-video+audio
   (`fl2va`)
 - `Ref2VA`: up to 9 images, 3 videos, and 3 audio references in supported
   image/video/audio combinations (`ref2va`; audio-only is rejected)
 
+One vLLM-Omni diffusion stage can load both DiTs while instantiating the
+tokenizer, processor, Qwen3-VL text encoder, video VAE, and audio VAE only
+once. Requests select the DiT with `extra_params.task`.
+
 The generated MP4 contains H.264 video and synchronized stereo audio.
 
 ## Prerequisites
 
-The checkpoint requires Hugging Face access approval. Download it and point
-`MODEL_ROOT` at the directory containing the `FL2VA` and `Ref2VA`
-subdirectories:
+The checkpoint requires Hugging Face access approval. Authenticate once;
+`vllm serve` downloads the required components automatically:
 
 ```bash
 hf auth login
-export MODEL_ROOT=/path/to/MiniMax-H3
-hf download MiniMaxAI/MiniMax-H3 --local-dir "${MODEL_ROOT}"
+export MODEL=MiniMaxAI/MiniMax-H3
 ```
 
+The vLLM-Omni pipeline downloads `FL2VA/**`, `Ref2VA/model_index.json`, and
+`Ref2VA/transformer/**`. It does not download or load the diffusers-format
+`transformer`, `transformer_ref`, or `vae` weights at the repository root, nor
+duplicate Ref2VA copies of shared components.
+
 Install vLLM-Omni from the checkout containing MiniMax H3 support. The
 two-GPU RTX 5090/4090 profiles use cuDNN attention and do not need
 FlashAttention-4. Install the optional dependency only for the four-GPU
@@ -41,6 +48,13 @@ B300/GB200 `FLASH_ATTN` profile:
 ```bash
 uv venv
 source .venv/bin/activate
+uv pip install -e .
+```
+
+To keep FA4 available as an explicit option on Blackwell, install the
+FlashAttention-4 extra:
+
+```bash
 uv pip install -e '.[fa4]'
 ```
 
@@ -49,16 +63,17 @@ reference-video preparation and MP4 output.
 
 ## Start a server
 
-One server loads one checkpoint partition. Set `MODEL` to `FL2VA` for T2VA
-and FL2VA requests, or to `Ref2VA` for either Ref2VA request.
+Pass the repository ID directly. The pipeline uses `FL2VA` for model discovery
+and shared components, and loads the second DiT from `Ref2VA/transformer`.
 
 ### Memory and storage requirements
 
 Treat GPU HBM, host RAM, and checkpoint storage as separate requirements. Each
 H3 checkpoint partition (`FL2VA` or `Ref2VA`) contains about **134 GiB** of
-BF16 safetensors (about **135 GiB** on disk). Keeping both partitions locally
-therefore needs roughly **270 GiB** of model storage, although the examples
-serve only one partition at a time.
+BF16 safetensors (about **135 GiB** on disk). Keeping both partitions
+locally therefore needs roughly **270 GiB** of model storage. A combined
+service downloads both; `--task-type fl2va` or `--task-type ref2va` downloads
+only the selected partition.
 
 CPU offload and distributed layerwise offload reduce GPU residency; they do
 not make the model weights disappear. With `--dlo-no-use-allgather`, each
@@ -79,7 +94,7 @@ This matches the accuracy-qualified reference path and prevents the Qwen3-VL
 encoder and DiT from being resident on the GPU at the same time.
 
 ```bash
-export MODEL="${MODEL_ROOT}/FL2VA"
+export MODEL=MiniMaxAI/MiniMax-H3
 export PORT=8091
 
 CUDA_VISIBLE_DEVICES=0 \
@@ -96,8 +111,8 @@ vllm serve "${MODEL}" \
 ```
 
 Use a GPU with enough memory for the active H3 component and enough system RAM
-for the offloaded components. CPU offload reduces GPU memory pressure but adds
-PCIe/NVLink transfer latency.
+for both offloaded DiTs plus the shared components. Model-level offload keeps
+the two DiTs mutually exclusive on GPU, but adds PCIe/NVLink transfer latency.
 
 ### Two 24/32 GB GPUs: TP2 distributed layerwise offload
 
@@ -109,7 +124,7 @@ copied to the GPUs once per denoise stage, reused by every sampling step, and
 released before VAE decode so the decoder can reuse their HBM.
 
 ```bash
-export MODEL="${MODEL_ROOT}/FL2VA"
+export MODEL=MiniMaxAI/MiniMax-H3
 export PORT=8091
 
 CUDA_VISIBLE_DEVICES=0,1 \
@@ -120,6 +135,7 @@ vllm serve "${MODEL}" \
   --host 0.0.0.0 \
   --port "${PORT}" \
   --trust-remote-code \
+  --task-type fl2va \
   --num-gpus 2 \
   --tensor-parallel-size 2 \
   --usp 1 \
@@ -191,23 +207,24 @@ bash examples/offline_inference/minimax_h3/run_h3_2gpu_all_tasks.sh
 The script selects 20 resident layers for `PROFILE=rtx5090` and 12 for
 `PROFILE=rtx4090`; `DLO_RESIDENT_LAYERS=N` overrides either default.
 
-### Four GPUs: measured best-practice throughput
+### Four GPUs: throughput-oriented combined service
 
-The best validated four-GPU configuration on four NVIDIA B300 GPUs is:
+For a combined service on four high-memory GPUs, use:
 
 - no CPU or layerwise offload;
 - Ulysses sequence parallelism degree 4;
 - native tiled VAE patch parallelism degree 4;
 - regional `torch.compile` for the repeated DiT blocks;
-- CuTe FlashAttention-4 through the `FLASH_ATTN` backend, with Ring and TP
-  left at 1.
+- dense BF16 `TRTLLM_ATTN`, with Ring and TP left at 1.
+
+Both DiTs remain resident in this no-offload configuration. If they do not fit,
+use model-level CPU offload.
 
 ```bash
-export MODEL="${MODEL_ROOT}/FL2VA"
+export MODEL=MiniMaxAI/MiniMax-H3
 export PORT=8091
 
 CUDA_VISIBLE_DEVICES=0,1,2,3 \
-FLASHINFER_DISABLE_VERSION_CHECK=1 \
 VLLM_WORKER_MULTIPROC_METHOD=spawn \
 VLLM_OMNI_VIDEO_SYNC_TIMEOUT=1800 \
 vllm serve "${MODEL}" \
@@ -220,19 +237,76 @@ vllm serve "${MODEL}" \
   --ring 1 \
   --vae-patch-parallel-size 4 \
   --vae-parallel-mode tile \
-  --vae-use-tiling \
-  --diffusion-attention-backend FLASH_ATTN
+  --vae-use-tiling
 ```
 
-On Blackwell, `FLASH_ATTN` selects FA4. Confirm the server log contains
-`Using CuTe FlashAttention-4 on Blackwell` before recording measurements.
-
 Do not add `--enforce-eager` to this performance configuration. The first
 request includes regional compilation; warm the server once before measuring
 steady-state latency. H3 is CFG-distilled, so `--cfg-parallel-size` must remain
 1. The H3 VAE supports its native `tile` mode, not
 `spatial_shard_height` or `spatial_shard_width`.
 
+### Attention Backends
+
+On datacenter Blackwell GPUs, MiniMax H3 defaults to dense BF16
+`TRTLLM_ATTN`; no attention backend flag is required. To select it explicitly,
+use:
+
+```bash
+--diffusion-attention-backend TRTLLM_ATTN
+```
+
+Stable measurements with the four-GPU profile above put dense `TRTLLM_ATTN`
+and FA4 within 2% of each other. `TRTLLM_ATTN` remains the datacenter Blackwell
+default and enables the optional optimizations below. Confirm the server log
+contains `Defaulting to diffusion attention backend TRTLLM_ATTN` before
+recording measurements when using the default selection.
+
+FA4 remains available by explicitly selecting the `FLASH_ATTN` backend:
+
+```bash
+--diffusion-attention-backend FLASH_ATTN
+```
+
+On Blackwell, `FLASH_ATTN` selects FA4. Confirm the server log contains
+`Using CuTe FlashAttention-4 on Blackwell` before recording FA4 measurements.
+
+`TRTLLM_ATTN` additionally supports two **lossy** optimizations for the long main
+DiT attention sequence: SAGE attention quantization and Skip-Softmax Sparse
+Attention. SAGE quantizes Q/K to the configured dtype and V to FP8. This example uses
+`fp8_e4m3` for Q/K; B200 also supports `int8` Q/K. The TRTLLM SAGE path fixes V
+to FP8, so vLLM-Omni only exposes the Q/K dtype. The token refiner is a short
+attention path, so the `per_role` override leaves SAGE and Skip-Softmax disabled
+for it. The example enables the calibration-free Skip-Softmax path with
+`threshold=0.05`, after the normalized timestep reaches `0.97`:
+
+```bash
+--diffusion-attention-config '{
+  "default": {
+    "backend": "TRTLLM_ATTN",
+    "quant": {
+      "dtype_qk": "fp8_e4m3",
+      "q_block_size": 1,
+      "k_block_size": 16
+    },
+    "skip_softmax": {
+      "threshold": 0.05,
+      "disabled_until_timestep": 0.97
+    }
+  },
+  "per_role": {
+    "minimax_h3.token_refiner": {
+      "backend": "TRTLLM_ATTN"
+    }
+  }
+}'
+```
+
+For configuration details, see
+[TRTLLM_ATTN Backend and Skip-Softmax](https://github.com/vllm-project/vllm-omni/blob/main/docs/user_guide/diffusion/attention_backends.md#trtllm_attn-backend-and-skip-softmax)
+and
+[TRTLLM_ATTN SAGE Quantization](https://github.com/vllm-project/vllm-omni/blob/main/docs/user_guide/diffusion/attention_backends.md#trtllm_attn-sage-quantization).
+
 ### Text encoder tensor parallelism
 
 The Qwen3-VL text encoder (~51.5 GB in BF16 for the retained 50 layers) is by
@@ -243,11 +317,10 @@ vLLM-style tensor-parallel layers and runs with distributed collectives over
 its own encoder process group):
 
 ```bash
-export MODEL="${MODEL_ROOT}/FL2VA"
+export MODEL=MiniMaxAI/MiniMax-H3
 export PORT=8091
 
 CUDA_VISIBLE_DEVICES=0,1,2,3 \
-FLASHINFER_DISABLE_VERSION_CHECK=1 \
 VLLM_WORKER_MULTIPROC_METHOD=spawn \
 VLLM_OMNI_VIDEO_SYNC_TIMEOUT=1800 \
 vllm serve "${MODEL}" \
@@ -261,8 +334,7 @@ vllm serve "${MODEL}" \
   --text-encoder-tp-size 4 \
   --vae-patch-parallel-size 4 \
   --vae-parallel-mode tile \
-  --vae-use-tiling \
-  --diffusion-attention-backend FLASH_ATTN
+  --vae-use-tiling
 ```
 
 `N` must divide the Qwen3-VL head counts (64 attention heads / 8 KV heads), so
@@ -275,12 +347,9 @@ the reference path within bf16 rounding: every encoder rank all-reduces the
 row-parallel projections, so the full `[seq, 5120]` layer-50 hidden state is
 replicated on every rank.
 
-To serve Ref2VA, stop the FL2VA server and restart the same one- or four-GPU
-command with:
-
-```bash
-export MODEL="${MODEL_ROOT}/Ref2VA"
-```
+No restart is needed: `task=fl2va` routes to `FL2VA/transformer`, while
+`task=ref2va` routes to
+`Ref2VA/transformer`. T2VA uses the FL2VA DiT.
 
 ### Online FP8 quantization
 
@@ -336,7 +405,7 @@ export API_URL="http://127.0.0.1:${PORT}/v1/videos/sync"
 
 ### 1. T2VA: text to video and audio
 
-Run this request against the `FL2VA` partition:
+Run this request against the combined service:
 
 ```bash
 curl -sS -X POST "${API_URL}" \
@@ -354,7 +423,7 @@ curl -sS -X POST "${API_URL}" \
 
 ### 2. FL2VA: first frame to video and audio
 
-Run this request against the `FL2VA` partition. When width and height are
+Run this request against the combined service. When width and height are
 omitted, H3 preserves the first-frame aspect ratio and uses a 768-pixel short
 edge.
 
@@ -393,10 +462,11 @@ curl -sS -X POST "${API_URL}" \
 
 ### 3. Ref2VA: image-only, image/audio, or mixed references
 
-Run these requests against the `Ref2VA` partition. Image-only Ref2VA omits
-`audio_reference`; adding one or more audio references is optional. The typed
-fields accept one object or an ordered JSON list. `audio_reference` accepts an
-HTTP(S) URL or a `data:` URL.
+Run these requests against the combined service or a Ref2VA-only service.
+Image-only Ref2VA omits `audio_reference`; adding one or more audio references
+is optional. The typed fields accept one object or an ordered JSON list.
+`audio_reference` accepts an HTTP(S) URL or a `data:` URL. In one terminal,
+expose the local reference assets to the serving host:
 
 ```bash
 python -m http.server 8092 \
@@ -446,7 +516,7 @@ shorter, the reference soundtrack is truncated to the generated clip.
 
 ### 4. Ref2VA: video, separate audio, and mixed references
 
-Run this request against the `Ref2VA` partition. Repeat the
+Run this request against the combined service. Repeat the
 `input_references` multipart field once per source video. H3 consumes the
 videos in form order and preserves their original soundtracks during
 conditioning.
@@ -504,11 +574,44 @@ default. `short_edge` controls the 768-pixel canvas and must be `768`.
 outputs; the synchronous raw-MP4 endpoint returns the first output when more
 than one is requested.
 
+## Request-scoped quality
+
+Add one of these fields to any HTTP request above. No Cache-DiT startup option
+is required; H3 installs its conservative profile when a `high` request
+arrives and removes it for a `lossless` request.
+
+```bash
+-F 'quality=lossless'  # Native reference path for this request.
+-F 'quality=high'      # H3's conservative Cache-DiT profile.
+```
+
+Omitting `quality` preserves the startup default: it uses the reference path
+normally, or the server-configured profile when the server was started with
+`--cache-backend cache_dit`.
+
+The following result was measured on 4× NVIDIA H200 with SP4, text-encoder
+TP4, 1344×768, 124 frames, 24 FPS, and 50 inference steps. One full
+`lossless` warmup was excluded, followed by three fixed prompt/seed pairs in
+balanced switch order.
+
+| `quality` | Median inference latency | Speedup | SSIM vs `lossless` | PSNR vs `lossless` | Expected trade-off |
+|---|---:|---:|---:|---:|---|
+| `lossless` | 85.49 s | 1.00× | 1.0000 | exact | Native reference path |
+| `high` | 63.36 s | 1.35× | 0.9709 | 34.98 dB | Faster with measured same-seed deviation |
+
+> [!NOTE]
+> `high` selects a fixed Cache-DiT profile, but its cache hit rate and
+> resulting latency/quality trade-off may vary by hardware, topology, and
+> workload. The values above apply to this deployment and are not universal
+> guarantees. `lossless` remains the exact reference path.
+
 ## Key parameters
 
 | Parameter | Recommended value | Notes |
 |-----------|-------------------|-------|
-| `task` | `t2va`, `fl2va`, or `ref2va` | Passed in `extra_params`; must match the served partition |
+| `quality` | omitted or `lossless` | Request-level quality intent; `high` dynamically installs H3's conservative Cache-DiT profile |
+| `extra_params.force_refresh_step_hint` | omitted | Optional positive 1-based denoising-step hint for an active Cache-DiT request; pair with `extra_params.force_refresh_step_policy`=`once` or `repeat` |
+| `task` | `t2va`, `fl2va`, or `ref2va` | Passed in `extra_params`; selects the task-specific DiT |
 | `duration` | Workload-specific | Decimal seconds in `extra_params`; converted to H3-compatible frame count |
 | `fps` | `24` | H3 output FPS is fixed |
 | `num_inference_steps` | `50` | Matches the reference accuracy workloads |
@@ -550,7 +653,8 @@ reduction.
 
 ## Known limitations
 
-- Each server process loads only one checkpoint partition.
+- Combined serving requires sibling `FL2VA` and `Ref2VA` directories, loads
+  both task-specific DiTs, and loads shared components once from `FL2VA`.
 - H3 currently executes one generation request per diffusion batch.
 - The first regional-compile request is a warmup and should not be included in
   steady-state performance measurements.

recipes/cosmos3/Cosmos3-DistOffload.md

@@ -56,10 +56,14 @@ vllm serve /path/to/Cosmos3-Super \
   --data-parallel-size 4
 ```
 
-### Without AllGather (full weights per rank, no sharding)
+### Without DLO AllGather (rank-local weights, no DLO sharding)
 
-Use `--dlo-no-use-allgather` when you want each rank to load full weights independently —
-no AllGather synchronization, no concurrent request requirement, but N× host memory.
+Use `--dlo-no-use-allgather` when you want each rank to stream the standard
+loader's rank-local weights independently — no DLO AllGather synchronization,
+no concurrent request requirement, but no DLO weight sharding and higher host
+memory use. In pure DP, this is equivalent to independent full-weight
+layerwise-offload replicas; with TP or HSDP, each rank retains the standard
+loader's local tensor layout.
 
 ```bash
 vllm serve /path/to/Cosmos3-Nano \
@@ -154,9 +158,9 @@ wait
 | Flag | Description | Default |
 |------|-------------|---------|
 | `--enable-distributed-layerwise-offload` | Enable DLO with H2D + AllGather overlap | `false` |
-| `--data-parallel-size N` | Number of DP ranks (weight sharding + concurrent requests) | `1` |
+| `--data-parallel-size N` | Number of DP ranks; DLO weight sharding + concurrent requests when AllGather is enabled | `1` |
 | `--dlo-use-allgather` | Use shard + AllGather for weight reconstruction (recommended) | `true` |
-| `--dlo-no-use-allgather` | Each rank loads full weights independently (no sharding, no AllGather) | `false` |
+| `--dlo-no-use-allgather` | Stream standard-loader rank-local weights independently (no DLO sharding, no DLO AllGather) | `false` |
 | `--usp N` | Use SP instead of DP for weight sharding (long sequences) | `1` |
 
 ## Mode comparison
@@ -166,7 +170,7 @@ wait
 | DLO + AllGather (DP4) | `--enable-distributed-layerwise-offload --data-parallel-size 4` | 1/4 model | 2 blocks | **3.3× HSDP** | Max throughput, short sequences |
 | DLO + AllGather (DP2) | `--enable-distributed-layerwise-offload --data-parallel-size 2` | 1/2 model | 2 blocks | 2× HSDP | Balanced throughput + memory |
 | DLO + AllGather (SP4) | `--enable-distributed-layerwise-offload --usp 4` | 1/4 model | 2 blocks | 1× (single req) | Long sequences (720p+) |
-| DLO no-AllGather | `--enable-distributed-layerwise-offload --dlo-no-use-allgather` | Full model | 2 blocks | 0.4× HSDP | No AllGather overhead, high CPU |
+| DLO no-AllGather | `--enable-distributed-layerwise-offload --dlo-no-use-allgather` | Standard-loader rank-local weights | 2 blocks | 0.4× HSDP | No DLO AllGather overhead, high CPU |
 | HSDP | `--use-hsdp --hsdp-shard-size 4` | 0 | 1/4 model | 1× (baseline) | Weights fit in HBM |
 
 ## Memory expectations
@@ -185,7 +189,11 @@ wait
 
 - `--dlo-use-allgather` (default) requires all concurrent requests to have the same
   `num_inference_steps` (AllGather is a collective that needs all ranks synchronized)
-- Online quantization (FP8) is incompatible with mmap loading — falls back to regular
-  `load_weights()` automatically
-- Tensor Parallel is not supported (DLO uses DP-based sharding, a different dimension)
-- HSDP + AllGather is rejected (would double-shard weights)
+- Online quantization (FP8) is incompatible with the DLO+AllGather mmap path;
+  use `--dlo-no-use-allgather` or disable online quantization
+- Tensor Parallel > 1 is rejected in the DLO+AllGather mmap path. TP with
+  `--dlo-no-use-allgather` retains the standard TP loader and should be treated
+  as an experimental compatibility path
+- HSDP + AllGather is rejected (would double-shard weights). HSDP with
+  `--dlo-no-use-allgather` is accepted at configuration level, but is not yet a
+  broadly validated production combination

tests/config/test_omni_config.py

@@ -39,7 +39,7 @@ from vllm_omni.config.stage_config import (
     load_deploy_config,
     merge_pipeline_deploy,
 )
-from vllm_omni.engine.stage_init_utils import build_engine_args_dict
+from vllm_omni.engine.stage_init_utils import build_legacy_engine_args_dict
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -124,10 +124,16 @@ def test_vllm_omni_config_from_pipeline_config_matches_merge_pipeline_deploy(mod
         assert omni_stage.model_config.enforce_eager == engine_args.get("enforce_eager", False)
         assert omni_stage.load_config.load_format == engine_args.get("load_format", "auto")
         assert omni_stage.load_config.tokenizer_mode == engine_args.get("tokenizer_mode", "auto")
-        assert omni_stage.cache_config.gpu_memory_utilization == engine_args.get("gpu_memory_utilization", 0.90)
-        assert omni_stage.scheduler_config.max_num_seqs == engine_args.get("max_num_seqs", 128)
+        assert omni_stage.cache_config.gpu_memory_utilization == engine_args.get("gpu_memory_utilization")
+        assert omni_stage.cache_config.enable_prefix_caching == engine_args.get("enable_prefix_caching")
+        expected_disable_hybrid = engine_args.get("disable_hybrid_kv_cache_manager")
+        if omni_stage.stage_pipeline_config.execution_type == StageExecutionType.LLM_GENERATION:
+            expected_disable_hybrid = True if expected_disable_hybrid is None else expected_disable_hybrid
+        assert omni_stage.cache_config.disable_hybrid_kv_cache_manager == expected_disable_hybrid
+        assert omni_stage.scheduler_config.max_num_seqs == engine_args.get("max_num_seqs")
         assert omni_stage.scheduler_config.max_num_batched_tokens == engine_args.get("max_num_batched_tokens")
-        assert omni_stage.scheduler_config.async_scheduling == engine_args.get("async_scheduling", True)
+        assert omni_stage.scheduler_config.enable_chunked_prefill == engine_args.get("enable_chunked_prefill")
+        assert omni_stage.scheduler_config.async_scheduling == engine_args.get("async_scheduling")
         legacy_parallel_config = engine_args.get("parallel_config") or {}
         assert omni_stage.parallel_config.tensor_parallel_size == legacy_parallel_config.get(
             "tensor_parallel_size",
@@ -198,6 +204,64 @@ def test_from_pipeline_config_applies_cli_overrides_without_stage_config_runtime
     assert stage1.runtime_config.num_gpus == stage1.parallel_config.world_size
 
 
+@pytest.mark.parametrize(
+    "cli_overrides",
+    [
+        {"enable_lora": True},
+        {"stage_0_enable_lora": True},
+    ],
+    ids=["global", "stage-scoped"],
+)
+def test_from_pipeline_config_rejects_explicit_unowned_engine_cli_fields(cli_overrides):
+    with pytest.raises(ValueError, match=r"no structured config owner: enable_lora"):
+        _from_pipeline_key("qwen3_tts", cli_overrides=cli_overrides)
+
+
+@pytest.mark.parametrize(
+    ("engine_extras", "unowned_field"),
+    [
+        ({"enable_lora": True}, "enable_lora"),
+        ({"parallel_config": {"cfg_parallel_size": 2}}, "parallel_config.cfg_parallel_size"),
+    ],
+    ids=["top-level", "nested-parallel-config"],
+)
+def test_from_pipeline_config_rejects_unowned_deploy_engine_extras(engine_extras, unowned_field):
+    pipeline = _resolve_pipeline_or_skip("qwen3_tts")
+    deploy = DeployConfig(stages=[StageDeployConfig(stage_id=0, engine_extras=engine_extras)])
+
+    with pytest.raises(ValueError, match=rf"no structured config owner: {unowned_field}"):
+        VllmOmniConfig.from_pipeline_config(pipeline, user_deploy_config=deploy)
+
+
+@pytest.mark.parametrize(
+    ("cli_overrides", "stage_id"),
+    [
+        ({"kv_cache_dtype": "fp8"}, 0),
+        ({"stage_0_kv_cache_dtype": "fp8"}, 0),
+        ({"stage_1_kv_cache_dtype": "fp8"}, 1),
+    ],
+    ids=["global", "ar-stage", "generation-stage"],
+)
+def test_from_pipeline_config_rejects_diffusion_only_cli_fields_for_llm_stages(cli_overrides, stage_id):
+    with pytest.raises(ValueError, match=rf"Stage {stage_id} .*no structured config owner: kv_cache_dtype"):
+        _from_pipeline_key("qwen3_tts", cli_overrides=cli_overrides)
+
+
+def test_from_pipeline_config_accepts_diffusion_only_cli_fields_for_diffusion_stage():
+    stage = _from_pipeline_key(
+        "dreamzero",
+        deploy_config_path="dreamzero_tp1_cfg2",
+        cli_overrides={"kv_cache_dtype": "fp8"},
+    ).stage_by_id(0)
+
+    assert isinstance(stage, VllmOmniDiffusionStageConfig)
+    assert stage.diffusion_config.diffusion_kv_cache_dtype == "fp8"
+
+
+def test_stage_cli_field_selection_defers_ownership_validation_until_sources_are_merged():
+    assert omni_config_module._stage_cli_overrides(0, {"enable_lora": True}) == {"enable_lora": True}
+
+
 def test_runtime_num_gpus_is_derived_from_parallel_world_size():
     omni_config = _from_pipeline_key("hunyuan_image3_dit")
     stage = omni_config.stage_by_id(0)
@@ -380,8 +444,10 @@ def test_vllm_omni_stage_config_public_fields_use_typed_stage_realizations():
     assert {f.name for f in fields(VllmOmniGenerationStageConfig)} == public_fields
 
 
-def test_runtime_config_fields_match_rfc_runtime_scope():
+def test_runtime_config_fields_match_structured_runtime_scope():
     assert {f.name for f in fields(OmniStageRuntimeConfig)} == {
+        "distributed_executor_backend",
+        "worker_cls",
         "devices",
         "num_replicas",
         "env",
@@ -392,8 +458,21 @@ def test_runtime_config_fields_match_rfc_runtime_scope():
     }
 
 
-def test_sub_config_fields_match_rfc_scopes():
+def test_sub_config_fields_match_structured_scopes():
     assert {f.name for f in fields(OmniStageModelConfig)} == {
+        "model",
+        "model_arch",
+        "revision",
+        "tokenizer_revision",
+        "code_revision",
+        "seed",
+        "logits_processors",
+        "trust_remote_code",
+        "dtype",
+        "attention_backend",
+        "moe_backend",
+        "hf_overrides",
+        "limit_mm_per_prompt",
         "active_stream_window",
         "duplex_max_sessions",
         "enable_sleep_mode",
@@ -404,6 +483,7 @@ def test_sub_config_fields_match_rfc_scopes():
         "task_type",
         "codec_frame_rate_hz",
         "enforce_eager",
+        "max_cudagraph_capture_size",
         "enable_flashinfer_autotune",
         "compilation_config",
         "enable_multithread_weight_load",
@@ -416,12 +496,16 @@ def test_sub_config_fields_match_rfc_scopes():
         "tokenizer_subdir",
     }
     assert {f.name for f in fields(OmniStageLoadConfig)} == {
+        "tokenizer",
+        "download_dir",
+        "skip_tokenizer_init",
         "load_format",
         "tokenizer_mode",
         "config_format",
         "skip_mm_profiling",
     }
     assert {f.name for f in fields(OmniStageCacheConfig)} == {
+        "kv_cache_memory_bytes",
         "gpu_memory_utilization",
         "enable_prefix_caching",
         "disable_hybrid_kv_cache_manager",
@@ -435,6 +519,8 @@ def test_sub_config_fields_match_rfc_scopes():
         "async_scheduling",
     }
     assert {f.name for f in fields(OmniStageConnectorConfig)} == {
+        "async_chunk",
+        "omni_kv_config",
         "stage_connector",
         "output_connectors",
         "input_connectors",
@@ -827,7 +913,7 @@ def test_from_pipeline_config_matches_build_engine_args_dict_behavior_for_repres
     pipeline = _resolve_pipeline_or_skip("qwen3_tts")
     legacy_stage = merge_pipeline_deploy(pipeline, _load_default_deploy(pipeline))[0]
     omega_stage = legacy_stage.to_omegaconf()
-    legacy_engine_args = build_engine_args_dict(
+    legacy_engine_args = build_legacy_engine_args_dict(
         omega_stage,
         model="/tmp/qwen3-tts",
         stage_connector_spec={"name": "SharedMemoryConnector", "extra": {}},
@@ -922,6 +1008,8 @@ def test_from_pipeline_config_normalizes_diffusion_config_aliases_from_engine_ar
 
 
 def test_diffusion_config_field_classification_covers_current_fields():
+    from vllm_omni.diffusion.data import OmniDiffusionConfig
+
     classified_fields = (
         omni_config_module._DIFFUSION_SHARED_CONFIG_FIELDS
         | omni_config_module._DIFFUSION_RUNTIME_CONFIG_FIELDS
@@ -929,8 +1017,12 @@ def test_diffusion_config_field_classification_covers_current_fields():
     )
 
     assert classified_fields == {f.name for f in fields(omni_config_module._DiffusionConfigProjection)}
+    assert {f.name for f in fields(OmniDiffusionConfig)} <= (
+        classified_fields | omni_config_module._DIFFUSION_MOVED_SHARED_FIELDS
+    )
     assert {
         "enable_prompt_embed_cache",
+        "enable_session_state_manager",
         "prompt_embed_cache_size",
         "diffusion_kv_cache_dtype",
     } <= omni_config_module._DIFFUSION_ONLY_CONFIG_FIELDS

tests/dfx/conftest.py

@@ -1,4 +1,3 @@
-import copy
 import json
 import os
 import re
@@ -10,7 +9,7 @@ from typing import Any
 
 import pytest
 
-from tests.helpers.mark import hardware_marks
+from tests.helpers.mark import get_hardware_mark_list, hardware_marks
 from tests.helpers.runtime import OmniServerParams
 from tests.helpers.stage_config import modify_stage_config
 from vllm_omni.platforms import current_omni_platform
@@ -447,36 +446,16 @@ def _unique_benchmark_param_id_suffixes(params_list: list[dict[str, Any]]) -> li
     return unique
 
 
-def extract_mark_resource_label(mark_field: Any) -> str:
-    """Return a filename-safe hardware label from ``mark.hardware_marks.res`` values.
-
-    Example: ``{"cuda": "H100"}`` -> ``"H100"``; multiple platforms join with ``-``.
-
-    Prefer :func:`get_runtime_resource_label` for perf result filenames so labels
-    reflect the machine that actually ran the benchmark.
-    """
-    if isinstance(mark_field, list):
-        for item in mark_field:
-            if isinstance(item, dict) and "hardware_marks" in item:
-                return extract_mark_resource_label(item)
-        return "na"
-    if not isinstance(mark_field, dict):
-        return "na"
-    hw = mark_field.get("hardware_marks")
-    if not isinstance(hw, dict):
-        return "na"
-    res = hw.get("res")
-    if not isinstance(res, dict) or not res:
-        return "na"
-    labels = [_safe_filename_token(value) for value in res.values()]
-    return "-".join(labels) if labels else "na"
-
-
-_KNOWN_RUNTIME_RESOURCE_TOKENS: tuple[str, ...] = (
+# Device marketing-name tokens for ``get_runtime_resource_label`` only.
+# Longer / more specific tokens must appear before shorter substrings
+# (e.g. ``H200`` before ``H20``, ``910B4`` before ``910``).
+_RUNTIME_DEVICE_ALIASES: tuple[str, ...] = (
     "H100",
     "H800",
     "H200",
     "H20",
+    "B200",
+    "GB200",
     "L40S",
     "L40",
     "L4",
@@ -505,7 +484,7 @@ def _normalize_runtime_device_label(raw: str) -> str:
     if not raw or not str(raw).strip():
         return "na"
     upper = str(raw).upper()
-    for token in _KNOWN_RUNTIME_RESOURCE_TOKENS:
+    for token in _RUNTIME_DEVICE_ALIASES:
         if token.upper() in upper:
             return _safe_filename_token(token)
     compact = re.sub(r"[^a-zA-Z0-9]+", "", str(raw))
@@ -558,10 +537,108 @@ def resource_label_for_filename(resource_label: str | None) -> str:
     return token
 
 
-def extract_configs_resource_label(configs: list[dict[str, Any]]) -> str:
-    """Return runtime hardware label for perf result filenames."""
-    del configs
-    return get_runtime_resource_label()
+def is_hardware_nested_baseline(baseline: dict[str, Any]) -> bool:
+    """True for ``{"H100": {"metric": ...}, "A3": {...}}`` baselines.
+
+    Detection is structural + hardware-label allowlist (not metric-name hints):
+
+    - every top-level value is a non-empty dict (per-hardware metric map);
+    - every top-level key is in :func:`tests.helpers.mark.get_hardware_mark_list`
+      (``[hardware-resource]`` markers from ``pyproject.toml``).
+
+    Metric names under each hardware bucket are unconstrained.
+    Runtime device aliases (:data:`_RUNTIME_DEVICE_ALIASES`) are unrelated.
+    """
+    if not baseline:
+        return False
+    if not all(isinstance(value, dict) and bool(value) for value in baseline.values()):
+        return False
+    allowed = {label.upper() for label in get_hardware_mark_list()}
+    return all((token := str(key).strip()) and token.upper() in allowed for key in baseline.keys())
+
+
+def resolve_baseline_value(
+    baseline_raw: Any,
+    *,
+    sweep_index: int | None,
+) -> Any:
+    """Pick the baseline threshold for one metric at this sweep step.
+
+    Accepts a sweep-aligned list/tuple (indexed by ``sweep_index``) or a scalar.
+    """
+    if baseline_raw is None:
+        return None
+    if isinstance(baseline_raw, dict):
+        raise TypeError(
+            "per-metric baseline dict keyed by concurrency/request-rate is not supported; "
+            "use a sweep-aligned list or a scalar under each hardware bucket"
+        )
+    if isinstance(baseline_raw, (list, tuple)):
+        if sweep_index is None:
+            raise ValueError("list baseline requires sweep_index")
+        if not (0 <= sweep_index < len(baseline_raw)):
+            raise IndexError(f"baseline list len={len(baseline_raw)} has no index {sweep_index}")
+        return baseline_raw[sweep_index]
+    return baseline_raw
+
+
+def resolve_baseline_for_sweep(
+    baseline_config: dict[str, Any] | None,
+    *,
+    sweep_index: int | None = None,
+) -> dict[str, Any]:
+    """Resolve sweep-dependent baselines while keeping every hardware bucket.
+
+    Expects hardware-nested input only::
+
+        {"H100": {"throughput_qps": [a, b, c]}, "A3": {"throughput_qps": [x, y, z]}}
+
+    with ``sweep_index=2`` (e.g. concurrency 32) becomes::
+
+        {"H100": {"throughput_qps": c}, "A3": {"throughput_qps": z}}
+
+    Per-metric values may be a sweep-aligned list or a scalar.
+    """
+    if not baseline_config:
+        return {}
+    if not is_hardware_nested_baseline(baseline_config):
+        allowed = sorted(get_hardware_mark_list())
+        unknown = [
+            str(key).strip()
+            for key in baseline_config.keys()
+            if str(key).strip().upper() not in {label.upper() for label in allowed}
+        ]
+        nested_shape = all(isinstance(value, dict) and bool(value) for value in baseline_config.values())
+        if nested_shape and unknown:
+            raise ValueError(
+                "baseline top-level keys must be [hardware-resource] pytest markers from "
+                "pyproject.toml (tool.pytest.ini_options.markers). "
+                f"Unknown hardware label(s): {unknown!r}. "
+                f"Known markers: {allowed!r}. "
+                "Add a marker such as '\"A100: [hardware-resource] Tests that require A100 GPU\"' "
+                "to pyproject.toml, then retry."
+            )
+        raise ValueError(
+            "baseline must be hardware-nested with [hardware-resource] marker names as "
+            "top-level keys, e.g. "
+            '{"H100": {"throughput_qps": [...]}, "A3": {"custom_metric": 1.0}}. '
+            f"got top-level keys={list(baseline_config.keys())!r}; "
+            f"known markers={allowed!r}. "
+            "If you need a new hardware bucket, add it to pyproject.toml markers with "
+            "the [hardware-resource] tag."
+        )
+
+    return {
+        hardware: {
+            metric_name: resolve_baseline_value(
+                baseline_raw,
+                sweep_index=sweep_index,
+            )
+            for metric_name, baseline_raw in metrics.items()
+        }
+        for hardware, metrics in baseline_config.items()
+        if isinstance(metrics, dict)
+    }
 
 
 def run_benchmark(
@@ -572,17 +649,18 @@ def run_benchmark(
     num_prompt: int,
     *,
     baseline_config: dict[str, Any] | None = None,
+    sweep_index: int | None = None,
     random_input_len: Any | None = None,
     random_output_len: Any | None = None,
     resource_label: str | None = None,
 ) -> dict[str, Any]:
     """Run one ``vllm bench serve --omni`` iteration and return parsed metrics.
 
-    After ``vllm bench`` writes the JSON, ``result["baseline"]`` stores the raw
-    ``baseline`` object from the test config (when ``baseline_config`` is provided),
-    including any hardware-nested map (e.g. ``{"H100": {...}}``). Downstream report
-    tooling can select a hardware bucket when needed. If the benchmark exits without
-    writing a result file, ``result_omni_template.json`` is used as a fallback.
+    After ``vllm bench`` writes the JSON, ``result["baseline"]`` stores the
+    sweep-resolved baseline (when ``baseline_config`` is provided). Hardware-nested
+    maps keep every hardware bucket, but each metric is reduced to the value for
+    this concurrency / request-rate step. If the benchmark exits without writing a
+    result file, ``result_omni_template.json`` is used as a fallback.
     """
     current_dt = datetime.now().strftime("%Y%m%d-%H%M%S")
     ri = _safe_filename_token(random_input_len)
@@ -649,8 +727,11 @@ def run_benchmark(
             result = json.load(f)
 
     if baseline_config:
-        # Persist the JSON baseline as-is (no hardware / sweep resolution).
-        result["baseline"] = copy.deepcopy(baseline_config)
+        # Keep every hardware bucket; resolve list metrics to this sweep step.
+        result["baseline"] = resolve_baseline_for_sweep(
+            baseline_config,
+            sweep_index=sweep_index,
+        )
     else:
         result["baseline"] = {}
     if random_input_len is not None:

tests/dfx/perf/scripts/diffusion_result_template.json

@@ -21,10 +21,12 @@
       "num-input-images": 0,
       "enable-negative-prompt": false,
       "baseline": {
-        "throughput_qps": 0,
-        "latency_mean": 0,
-        "peak_memory_mb_max": 0,
-        "peak_memory_mb_mean": 0
+        "H100": {
+          "throughput_qps": 0,
+          "latency_mean": 0,
+          "peak_memory_mb_max": 0,
+          "peak_memory_mb_mean": 0
+        }
       }
     },
     "result": {

tests/dfx/perf/scripts/result_omni_template.json

@@ -46,9 +46,11 @@
   "median_audio_duration_s": 0,
   "p99_audio_duration_s": 0,
   "baseline": {
-    "mean_ttft_ms": 0,
-    "mean_audio_ttfp_ms": 0,
-    "mean_audio_rtf": 0
+    "H100": {
+      "mean_ttft_ms": 0,
+      "mean_audio_ttfp_ms": 0,
+      "mean_audio_rtf": 0
+    }
   },
   "random_input_len": 0,
   "random_output_len": 0

tests/dfx/perf/scripts/run_benchmark.py

@@ -237,7 +237,7 @@ def test_performance_benchmark(omni_server, benchmark_params):
         break
 
     # QPS / request-rate sweep
-    for qps, num_prompt in zip(qps_list, num_prompt_list):
+    for sweep_index, (qps, num_prompt) in enumerate(zip(qps_list, num_prompt_list)):
         args = args + ["--request-rate", str(qps), "--num-prompts", str(num_prompt)]
         result = run_benchmark(
             args=args,
@@ -246,6 +246,7 @@ def test_performance_benchmark(omni_server, benchmark_params):
             dataset_name=dataset_name,
             num_prompt=num_prompt,
             baseline_config=params.get("baseline"),
+            sweep_index=sweep_index,
             random_input_len=params.get("random_input_len"),
             random_output_len=params.get("random_output_len"),
             resource_label=resource_label,
@@ -253,7 +254,7 @@ def test_performance_benchmark(omni_server, benchmark_params):
         assert_result(result, num_prompt)
 
     # concurrency test
-    for concurrency, num_prompt in zip(max_concurrency_list, num_prompt_list):
+    for sweep_index, (concurrency, num_prompt) in enumerate(zip(max_concurrency_list, num_prompt_list)):
         args = args + ["--max-concurrency", str(concurrency), "--num-prompts", str(num_prompt), "--request-rate", "inf"]
         result = run_benchmark(
             args=args,
@@ -262,6 +263,7 @@ def test_performance_benchmark(omni_server, benchmark_params):
             dataset_name=dataset_name,
             num_prompt=num_prompt,
             baseline_config=params.get("baseline"),
+            sweep_index=sweep_index,
             random_input_len=params.get("random_input_len"),
             random_output_len=params.get("random_output_len"),
             resource_label=resource_label,

tests/dfx/perf/scripts/run_diffusion_benchmark.py

@@ -23,7 +23,6 @@ of all runs from cases in that file). Bulk load without ``--test-config-file`` u
 the same per-file aggregation; ``-m`` only selects which cases run.
 """
 
-import copy
 import json
 import os
 import socket
@@ -45,6 +44,7 @@ from tests.dfx.conftest import (
     get_runtime_resource_label,
     hardware_json_value,
     is_diffusion_perf_config,
+    resolve_baseline_for_sweep,
     resolve_pytest_marks,
     resource_label_for_filename,
 )
@@ -350,10 +350,49 @@ def _kill_process_tree(pid: int) -> None:
 
 
 def _resolve_offline_model(model: str) -> str:
-    """Under HF_HUB_OFFLINE, resolve a HF repo id to its local snapshot dir."""
+    """Resolve a model string to a local directory.
+
+    Local paths pass through unchanged. HF repo ids are resolved to their local
+    snapshot dir when ``HF_HUB_OFFLINE`` is set. ``repo_id/subfolder`` (e.g.
+    ``MiniMaxAI/MiniMax-H3/FL2VA``) is resolved to the matching partition
+    subfolder of the snapshot, downloading only that subfolder when online.
+    """
     import huggingface_hub
 
-    if not model or os.path.isdir(model) or not huggingface_hub.constants.HF_HUB_OFFLINE:
+    if not model or os.path.isdir(model):
+        return model
+
+    # MiniMax-H3 custom code uses relative imports across files, which breaks
+    # when the HF cache snapshot is symlinked (``get_class_from_dynamic_module``
+    # resolves ``realpath`` into the blobs dir). Mirror the accuracy test's env
+    # overrides so CI / local runs can point at a materialized partition. The
+    # repo-root key (used with ``--task-type fl2va|ref2va``) expects a
+    # materialized root containing the FL2VA/Ref2VA subfolders.
+    model_env_overrides = {
+        "MiniMaxAI/MiniMax-H3": "VLLM_TEST_MINIMAX_H3_MODEL",
+        "MiniMaxAI/MiniMax-H3/FL2VA": "VLLM_TEST_MINIMAX_H3_FL2VA_MODEL",
+        "MiniMaxAI/MiniMax-H3/Ref2VA": "VLLM_TEST_MINIMAX_H3_REF2VA_MODEL",
+    }
+    env_name = model_env_overrides.get(model)
+    if env_name:
+        env_model = os.environ.get(env_name)
+        if env_model:
+            return env_model
+
+    parts = model.split("/")
+    if len(parts) >= 3:
+        repo_id = "/".join(parts[:2])
+        subfolder = "/".join(parts[2:])
+        from huggingface_hub import snapshot_download
+
+        snapshot_root = snapshot_download(
+            repo_id,
+            allow_patterns=[f"{subfolder}/**"],
+            local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE,
+        )
+        return str(Path(snapshot_root) / subfolder)
+
+    if not huggingface_hub.constants.HF_HUB_OFFLINE:
         return model
     from huggingface_hub import snapshot_download
 
@@ -387,6 +426,11 @@ class DiffusionServer:
     def _start_server(self) -> None:
         env = os.environ.copy()
         env["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+        # Give each run its own writable video-storage dir. The server default
+        # (/tmp/storage) may be missing or not writable on shared hosts / CI.
+        storage_path = Path(os.environ.get("VLLM_OMNI_STORAGE_PATH", str(BENCHMARK_RESULT_DIR / "storage")))
+        storage_path.mkdir(parents=True, exist_ok=True)
+        env["VLLM_OMNI_STORAGE_PATH"] = str(storage_path)
 
         cmd = [
             sys.executable,
@@ -766,8 +810,8 @@ def run_benchmark(
     completed = metrics.get("completed_requests", metrics.get("completed", 0))
     failed = metrics.get("failed_requests", metrics.get("failed", 0))
 
-    # Persist the JSON baseline as-is (no hardware / sweep resolution).
-    baseline = copy.deepcopy(params.get("baseline") or {})
+    # Persist sweep-resolved baseline from params (already narrowed in _build_run_params).
+    baseline = params.get("baseline") or {}
     metrics["baseline"] = baseline
 
     record: dict[str, Any] = {
@@ -862,6 +906,7 @@ def _build_run_params(
     params: dict[str, Any],
     *,
     num_prompts: int,
+    sweep_index: int | None = None,
     request_rate: Any | None = None,
     max_concurrency: Any | None = None,
 ) -> dict[str, Any]:
@@ -873,6 +918,12 @@ def _build_run_params(
         run_params["request-rate"] = request_rate
     if max_concurrency is not None:
         run_params["max-concurrency"] = max_concurrency
+    if "baseline" in params:
+        # Keep all hardware buckets; pick the metric value for this sweep step only.
+        run_params["baseline"] = resolve_baseline_for_sweep(
+            params.get("baseline"),
+            sweep_index=sweep_index,
+        )
     return run_params
 
 
@@ -895,19 +946,20 @@ def _iter_sweep_runs(params: dict[str, Any]) -> list[dict[str, Any]]:
 
     sweep_runs: list[dict[str, Any]] = []
 
-    for request_rate, num_prompts in zip(request_rate_list, num_prompt_list):
+    for sweep_index, (request_rate, num_prompts) in enumerate(zip(request_rate_list, num_prompt_list)):
         sweep_runs.append(
             {
                 "params": _build_run_params(
                     params,
                     request_rate=request_rate,
                     num_prompts=num_prompts,
+                    sweep_index=sweep_index,
                 ),
                 "num_prompts": num_prompts,
             }
         )
 
-    for max_concurrency, num_prompts in zip(max_concurrency_list, num_prompt_list):
+    for sweep_index, (max_concurrency, num_prompts) in enumerate(zip(max_concurrency_list, num_prompt_list)):
         sweep_runs.append(
             {
                 "params": _build_run_params(
@@ -915,6 +967,7 @@ def _iter_sweep_runs(params: dict[str, Any]) -> list[dict[str, Any]]:
                     max_concurrency=max_concurrency,
                     num_prompts=num_prompts,
                     request_rate="inf",
+                    sweep_index=sweep_index,
                 ),
                 "num_prompts": num_prompts,
             }
@@ -927,6 +980,7 @@ def _iter_sweep_runs(params: dict[str, Any]) -> list[dict[str, Any]]:
                 "params": _build_run_params(
                     params,
                     num_prompts=default_num_prompts,
+                    sweep_index=0,
                 ),
                 "num_prompts": default_num_prompts,
             }

tests/dfx/perf/tests/test_boogu_image_edit_vllm_omni.json

@@ -25,8 +25,8 @@
                 "enable-negative-prompt": true,
                 "baseline": {
                     "H100": {
-                        "throughput_qps": {"1": 0.0303, "2": 0.0302, "4": 0.0302},
-                        "latency_mean": {"1": 32.6, "2": 62.2, "4": 111.7},
+                        "throughput_qps": [0.0303, 0.0302, 0.0302],
+                        "latency_mean": [32.6, 62.2, 111.7],
                         "peak_memory_mb_max": 40500,
                         "peak_memory_mb_mean": 40500
                     }

tests/dfx/perf/tests/test_boogu_image_vllm_omni.json

@@ -25,8 +25,8 @@
                 "enable-negative-prompt": true,
                 "baseline": {
                     "H100": {
-                        "throughput_qps": {"1": 0.0549, "2": 0.0547, "4": 0.0549},
-                        "latency_mean": {"1": 18.0, "2": 34.3, "4": 61.4},
+                        "throughput_qps": [0.0549, 0.0547, 0.0549],
+                        "latency_mean": [18.0, 34.3, 61.4],
                         "peak_memory_mb_max": 40500,
                         "peak_memory_mb_mean": 40500
                     }

tests/dfx/perf/tests/test_hunyuan_image3_it2i.json

@@ -134,7 +134,7 @@
         "max-concurrency": 16,
         "num-prompts": 96,
         "baseline": {
-          "H100": {
+          "H200": {
             "throughput_qps": 0.1874,
             "latency_p99": 101.852,
             "latency_mean": 79.0385,

tests/dfx/perf/tests/test_hunyuan_image_tp2_cfgp2.json

@@ -38,7 +38,7 @@
         "num-prompts": 10,
         "max-concurrency": 1,
         "baseline": {
-          "H100": {
+          "H200": {
             "throughput_qps": 0.21,
             "latency_p99": 4.7469,
             "peak_memory_mb_max": 101100,

tests/dfx/perf/tests/test_hunyuan_image_tp2_sp2.json

@@ -38,7 +38,7 @@
         "num-prompts": 10,
         "max-concurrency": 1,
         "baseline": {
-          "H100": {
+          "H200": {
             "throughput_qps": 0.20,
             "latency_p99": 5.1025,
             "peak_memory_mb_max": 97402,

tests/dfx/perf/tests/test_hunyuan_image_tp4.json

@@ -37,7 +37,7 @@
         "num-prompts": 10,
         "max-concurrency": 1,
         "baseline": {
-          "H100": {
+          "H200": {
             "throughput_qps": 0.213,
             "latency_mean": 4.6953,
             "latency_p99": 4.8601,

tests/dfx/perf/tests/test_minimax_h3_vllm_omni.json

@@ -0,0 +1,295 @@
+[
+  {
+    "test_name": "test_minimax_h3_t2v_4gpu",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 4
+        }
+      },
+      "full_model",
+      "diffusion"
+    ],
+    "description": "MiniMax-H3 T2VA text-to-video+audio on 4 GPUs (USP4 + HSDP4, encoder-TP4, VAE-PP4 tile)",
+    "server_type": "vllm-omni",
+    "benchmark_endpoint": "/v1/videos",
+    "server_params": {
+      "model": "MiniMaxAI/MiniMax-H3",
+      "serve_args": {
+        "trust-remote-code": true,
+        "task-type": "fl2va",
+        "num-gpus": 4,
+        "usp": 4,
+        "ring": 1,
+        "use-hsdp": true,
+        "hsdp-shard-size": 4,
+        "text-encoder-tp-size": 4,
+        "vae-patch-parallel-size": 4,
+        "vae-parallel-mode": "tile",
+        "vae-use-tiling": true,
+        "diffusion-attention-backend": "FLASH_ATTN",
+        "stage-init-timeout": 1800,
+        "init-timeout": 1800,
+        "enable-diffusion-pipeline-profiler": true
+      }
+    },
+    "benchmark_params": [
+      {
+        "name": "1344x768_frames209_steps8",
+        "dataset": "random",
+        "task": "t2v",
+        "num-prompts": 3,
+        "max-concurrency": 1,
+        "warmup-requests": 1,
+        "warmup-concurrency": 1,
+        "seed": 42,
+        "random-request-config": [
+          {
+            "width": 1344,
+            "height": 768,
+            "prompt": "Cinematic drone shot pushing through misty mountain ridges at sunrise, golden light catching the clouds, slow and smooth.",
+            "num_inference_steps": 8,
+            "num_frames": 209,
+            "fps": 24,
+            "weight": 1
+          }
+        ],
+        "extra-body": {
+          "flow_shift": 12.0,
+          "extra_params": "{\"task\":\"t2va\",\"duration\":8.7,\"aspect_ratio\":\"16:9\",\"audio_flow_shift\":3.0}"
+        },
+        "baseline": {
+          "H100": {
+            "throughput_qps": 0.0211,
+            "latency_mean": 38.3252,
+            "peak_memory_mb_mean": 54219.6
+          }
+        }
+      }
+    ]
+  },
+  {
+    "test_name": "test_minimax_h3_ti2v_4gpu",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 4
+        }
+      },
+      "full_model",
+      "diffusion"
+    ],
+    "description": "MiniMax-H3 FL2VA first-frame image-to-video+audio on 4 GPUs (USP4 + HSDP4, encoder-TP4, VAE-PP4 tile)",
+    "server_type": "vllm-omni",
+    "benchmark_endpoint": "/v1/videos",
+    "server_params": {
+      "model": "MiniMaxAI/MiniMax-H3",
+      "serve_args": {
+        "trust-remote-code": true,
+        "task-type": "fl2va",
+        "num-gpus": 4,
+        "usp": 4,
+        "ring": 1,
+        "use-hsdp": true,
+        "hsdp-shard-size": 4,
+        "text-encoder-tp-size": 4,
+        "vae-patch-parallel-size": 4,
+        "vae-parallel-mode": "tile",
+        "vae-use-tiling": true,
+        "diffusion-attention-backend": "FLASH_ATTN",
+        "stage-init-timeout": 1800,
+        "init-timeout": 1800,
+        "enable-diffusion-pipeline-profiler": true
+      }
+    },
+    "benchmark_params": [
+      {
+        "name": "1344x768_frames209_steps8",
+        "dataset": "random",
+        "task": "ti2v",
+        "num-prompts": 3,
+        "max-concurrency": 1,
+        "warmup-requests": 1,
+        "warmup-concurrency": 1,
+        "seed": 42,
+        "random-request-config": [
+          {
+            "width": 1344,
+            "height": 768,
+            "prompt": "Continue the scene: the camera slowly pushes in on the subject as the background city lights drift and blur, with the actor singing softly and emotionally.",
+            "num_inference_steps": 8,
+            "num_frames": 209,
+            "fps": 24,
+            "weight": 1
+          }
+        ],
+        "extra-body": {
+          "flow_shift": 12.0,
+          "extra_params": "{\"task\":\"fl2va\",\"duration\":8.7,\"audio_flow_shift\":3.0}"
+        },
+        "num-input-images": 1,
+        "baseline": {
+          "H100": {
+            "throughput_qps": 0.0212,
+            "latency_mean": 38.1385,
+            "peak_memory_mb_mean": 54622.8
+          }
+        }
+      }
+    ]
+  },
+  {
+    "test_name": "test_minimax_h3_v2v_4gpu",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 4
+        }
+      },
+      "full_model",
+      "diffusion"
+    ],
+    "description": "MiniMax-H3 Ref2VA video-reference-to-video+audio on 4 GPUs (USP4 + HSDP4, encoder-TP4, VAE-PP4 tile)",
+    "server_type": "vllm-omni",
+    "benchmark_endpoint": "/v1/videos",
+    "server_params": {
+      "model": "MiniMaxAI/MiniMax-H3",
+      "serve_args": {
+        "trust-remote-code": true,
+        "task-type": "ref2va",
+        "num-gpus": 4,
+        "usp": 4,
+        "ring": 1,
+        "use-hsdp": true,
+        "hsdp-shard-size": 4,
+        "text-encoder-tp-size": 4,
+        "vae-patch-parallel-size": 4,
+        "vae-parallel-mode": "tile",
+        "vae-use-tiling": true,
+        "diffusion-attention-backend": "FLASH_ATTN",
+        "stage-init-timeout": 1800,
+        "init-timeout": 1800,
+        "enable-diffusion-pipeline-profiler": true
+      }
+    },
+    "benchmark_params": [
+      {
+        "name": "1344x768_frames209_steps8",
+        "dataset": "random",
+        "task": "v2v",
+        "num-prompts": 3,
+        "max-concurrency": 1,
+        "warmup-requests": 1,
+        "warmup-concurrency": 1,
+        "seed": 42,
+        "random-request-config": [
+          {
+            "width": 1344,
+            "height": 768,
+            "prompt": "Restyle the reference footage into a cinematic night scene with neon reflections, keeping the same camera motion and subject.",
+            "num_inference_steps": 8,
+            "num_frames": 209,
+            "fps": 24,
+            "weight": 1
+          }
+        ],
+        "extra-body": {
+          "flow_shift": 12.0,
+          "extra_params": "{\"task\":\"ref2va\",\"duration\":8.7,\"aspect_ratio\":\"16:9\",\"audio_flow_shift\":3.0}"
+        },
+        "baseline": {
+          "H100": {
+            "throughput_qps": 0.0077,
+            "latency_mean": 105.7508,
+            "peak_memory_mb_mean": 55180.8
+          }
+        }
+      }
+    ]
+  },
+  {
+    "test_name": "test_minimax_h3_t2v_dlo_4gpu",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 4
+        }
+      },
+      "full_model",
+      "diffusion"
+    ],
+    "description": "MiniMax-H3 T2VA with 4-card distributed layerwise offload (DLO), 8 denoise steps (USP4 + HSDP4, encoder-TP4, VAE-PP4 tile)",
+    "server_type": "vllm-omni",
+    "benchmark_endpoint": "/v1/videos",
+    "server_params": {
+      "model": "MiniMaxAI/MiniMax-H3",
+      "serve_args": {
+        "trust-remote-code": true,
+        "task-type": "fl2va",
+        "num-gpus": 4,
+        "usp": 4,
+        "ring": 1,
+        "use-hsdp": true,
+        "hsdp-shard-size": 4,
+        "text-encoder-tp-size": 4,
+        "vae-patch-parallel-size": 4,
+        "vae-parallel-mode": "tile",
+        "vae-use-tiling": true,
+        "diffusion-attention-backend": "FLASH_ATTN",
+        "stage-init-timeout": 1800,
+        "init-timeout": 1800,
+        "enable-diffusion-pipeline-profiler": true,
+        "enable-distributed-layerwise-offload": true,
+        "dlo-no-use-allgather": true,
+        "dlo-resident-layers": 8,
+        "enforce-eager": true
+      }
+    },
+    "benchmark_params": [
+      {
+        "name": "1344x768_frames209_steps8",
+        "dataset": "random",
+        "task": "t2v",
+        "num-prompts": 3,
+        "max-concurrency": 1,
+        "warmup-requests": 1,
+        "warmup-concurrency": 1,
+        "seed": 42,
+        "random-request-config": [
+          {
+            "width": 1344,
+            "height": 768,
+            "prompt": "Cinematic drone shot pushing through misty mountain ridges at sunrise, golden light catching the clouds, slow and smooth.",
+            "num_inference_steps": 8,
+            "num_frames": 209,
+            "fps": 24,
+            "weight": 1
+          }
+        ],
+        "extra-body": {
+          "flow_shift": 12.0,
+          "extra_params": "{\"task\":\"t2va\",\"duration\":8.7,\"aspect_ratio\":\"16:9\",\"audio_flow_shift\":3.0}"
+        },
+        "baseline": {
+          "H100": {
+            "throughput_qps": 0.0212,
+            "latency_mean": 38.2089,
+            "peak_memory_mb_mean": 36218.4
+          }
+        }
+      }
+    ]
+  }
+]

tests/dfx/perf/tests/test_runner_metadata.py

@@ -17,7 +17,7 @@ def test_task_excluded_from_cli_args():
         "backend": "openai-audio-speech",
         "endpoint": "/v1/audio/speech",
         "percentile-metrics": "audio_rtf,audio_ttfp",
-        "baseline": {"mean_audio_rtf": [0.5]},
+        "baseline": {"H100": {"mean_audio_rtf": [0.5]}},
     }
     exclude_keys = {
         "request_rate",
@@ -123,35 +123,6 @@ def test_resolve_pytest_marks_rejects_legacy_object_format():
         )
 
 
-def test_extract_mark_resource_label():
-    from tests.dfx.conftest import extract_mark_resource_label
-
-    assert extract_mark_resource_label(None) == "na"
-    assert (
-        extract_mark_resource_label(
-            [
-                {"hardware_marks": {"res": {"cuda": "H100"}, "num_cards": 1}},
-                "full_model",
-            ]
-        )
-        == "H100"
-    )
-    assert (
-        extract_mark_resource_label(
-            [
-                {
-                    "hardware_marks": {
-                        "res": {"cuda": "H100", "rocm": "MI325"},
-                        "num_cards": 2,
-                    }
-                },
-                "diffusion",
-            ]
-        )
-        == "H100-MI325"
-    )
-
-
 def test_resource_label_for_filename():
     from tests.dfx.conftest import resource_label_for_filename
 
@@ -169,16 +140,14 @@ def test_hardware_json_value():
     assert hardware_json_value(None) == ""
 
 
-def test_extract_configs_resource_label(monkeypatch):
-    from tests.dfx.conftest import extract_configs_resource_label, get_runtime_resource_label
+def test_get_runtime_resource_label(monkeypatch):
+    from tests.dfx.conftest import get_runtime_resource_label
 
     monkeypatch.setattr(
         "tests.dfx.conftest._read_runtime_device_name",
         lambda *, device_id=0: "NVIDIA H100 80GB HBM3",
     )
-    get_runtime_resource_label(refresh=True)
-    assert extract_configs_resource_label([]) == "H100"
-    get_runtime_resource_label(refresh=True)
+    assert get_runtime_resource_label(refresh=True) == "H100"
     monkeypatch.setattr(
         "tests.dfx.conftest._read_runtime_device_name",
         lambda *, device_id=0: "Ascend910B2",
@@ -336,3 +305,183 @@ def test_paired_omni_benchmark_reuses_server_and_preserves_case_metadata(tmp_pat
         ("start", tts_server),
         ("stop", tts_server),
     ]
+
+
+def test_is_hardware_nested_baseline():
+    from tests.dfx.conftest import (
+        _RUNTIME_DEVICE_ALIASES,
+        is_hardware_nested_baseline,
+    )
+    from tests.helpers.mark import get_hardware_mark_list
+
+    hardware_marks = get_hardware_mark_list()
+    assert hardware_marks
+    assert {"H100", "L4", "A3", "MI325", "B200"} <= hardware_marks
+    assert "cuda" not in hardware_marks
+    assert "full_model" not in hardware_marks
+
+    # Runtime aliases are a full static list (independent of baseline allowlist).
+    assert "H100" in _RUNTIME_DEVICE_ALIASES
+    assert "A100" in _RUNTIME_DEVICE_ALIASES
+
+    assert is_hardware_nested_baseline(
+        {
+            "H100": {"mean_ttft_ms": [1.0, 2.0], "mean_e2el_ms": [10.0, 20.0]},
+            "A3": {"mean_ttft_ms": [0.9, 1.8], "mean_e2el_ms": [9.0, 18.0]},
+        }
+    )
+    # Custom metric names under a known hardware label are allowed.
+    assert is_hardware_nested_baseline({"H100": {"custom_stage_ms": 12.0, "foo_bar": [1.0, 2.0]}})
+    assert is_hardware_nested_baseline({"A3": {"request_throughput": 0.5}})
+    # Flat metric maps / concurrency-keyed maps are not hardware-nested.
+    assert not is_hardware_nested_baseline({"mean_ttft_ms": [1.0, 2.0], "mean_e2el_ms": [10.0, 20.0]})
+    assert not is_hardware_nested_baseline({"mean_ttft_ms": {"1": 1.0, "32": 2.0}})
+    assert not is_hardware_nested_baseline({})
+    # Alias-only labels (not [hardware-resource] markers) are rejected for baselines.
+    assert not is_hardware_nested_baseline({"A100": {"throughput_qps": 1.0}})
+    # Unknown top-level hardware label is rejected.
+    assert not is_hardware_nested_baseline({"UnknownGPU": {"throughput_qps": 1.0}})
+    # Empty per-hardware metric map is rejected.
+    assert not is_hardware_nested_baseline({"H100": {}})
+
+
+def test_resolve_baseline_for_sweep_keeps_all_hardware_for_one_concurrency():
+    from tests.dfx.conftest import resolve_baseline_for_sweep
+
+    baseline = {
+        "H100": {
+            "mean_ttft_ms": [96.4, 140.8, 271.9, 362.3, 507.8],
+            "mean_e2el_ms": [18507.0, 28365.0, 31907.0, 48161.0, 72630.0],
+        },
+        "A3": {
+            "mean_ttft_ms": [90.0, 130.0, 250.0, 340.0, 480.0],
+            "mean_e2el_ms": [17000.0, 26000.0, 30000.0, 45000.0, 70000.0],
+        },
+    }
+    # max_concurrency=[1,4,8,16,32] -> index 4 is concurrency 32
+    got = resolve_baseline_for_sweep(baseline, sweep_index=4)
+    assert got == {
+        "H100": {"mean_ttft_ms": 507.8, "mean_e2el_ms": 72630.0},
+        "A3": {"mean_ttft_ms": 480.0, "mean_e2el_ms": 70000.0},
+    }
+    # First sweep step keeps both hardware buckets too.
+    got0 = resolve_baseline_for_sweep(baseline, sweep_index=0)
+    assert set(got0) == {"H100", "A3"}
+    assert got0["H100"]["mean_ttft_ms"] == 96.4
+    assert got0["A3"]["mean_ttft_ms"] == 90.0
+
+
+def test_resolve_baseline_for_sweep_rejects_flat_baseline():
+    from tests.dfx.conftest import resolve_baseline_for_sweep
+
+    with pytest.raises(ValueError, match=r"hardware-nested.*pyproject\.toml"):
+        resolve_baseline_for_sweep(
+            {"throughput_qps": [0.4, 0.6], "latency_mean": [1.0, 2.0]},
+            sweep_index=1,
+        )
+
+
+def test_resolve_baseline_for_sweep_rejects_unknown_hardware_label():
+    from tests.dfx.conftest import resolve_baseline_for_sweep
+
+    with pytest.raises(ValueError, match=r"Unknown hardware label\(s\): \['A100'\].*pyproject\.toml"):
+        resolve_baseline_for_sweep(
+            {"A100": {"throughput_qps": 1.0}},
+            sweep_index=0,
+        )
+
+
+def test_resolve_baseline_for_sweep_supports_list_and_scalar_under_hardware():
+    from tests.dfx.conftest import resolve_baseline_for_sweep
+
+    # Sweep-aligned lists under each hardware bucket (canonical form).
+    listed = {
+        "H100": {"throughput_qps": [0.4, 0.6, 0.8], "latency_mean": [1.0, 2.0, 3.0]},
+        "A3": {"throughput_qps": [0.5, 0.7, 0.9], "latency_mean": [0.9, 1.8, 2.7]},
+    }
+    assert resolve_baseline_for_sweep(listed, sweep_index=1) == {
+        "H100": {"throughput_qps": 0.6, "latency_mean": 2.0},
+        "A3": {"throughput_qps": 0.7, "latency_mean": 1.8},
+    }
+
+    # Scalars under hardware stay as-is (single-concurrency cases).
+    scalar = {"H100": {"throughput_qps": 0.5}}
+    assert resolve_baseline_for_sweep(scalar, sweep_index=0) == {"H100": {"throughput_qps": 0.5}}
+    # Custom metric names are preserved.
+    custom = {"H100": {"custom_stage_ms": [10.0, 20.0]}}
+    assert resolve_baseline_for_sweep(custom, sweep_index=1) == {"H100": {"custom_stage_ms": 20.0}}
+    assert resolve_baseline_for_sweep(None) == {}
+    assert resolve_baseline_for_sweep({}) == {}
+
+
+def test_resolve_baseline_value_errors():
+    from tests.dfx.conftest import resolve_baseline_value
+
+    with pytest.raises(ValueError, match="sweep_index"):
+        resolve_baseline_value([1.0, 2.0], sweep_index=None)
+    with pytest.raises(IndexError):
+        resolve_baseline_value([1.0], sweep_index=1)
+    with pytest.raises(TypeError, match="not supported"):
+        resolve_baseline_value({"1": 0.4}, sweep_index=0)
+
+
+def test_diffusion_build_run_params_resolves_baseline_per_sweep(tmp_path, monkeypatch):
+    """``_build_run_params`` / ``_iter_sweep_runs`` narrow list baselines per concurrency."""
+    import importlib
+    import sys
+
+    cfg = tmp_path / "mini_diffusion_perf.json"
+    cfg.write_text(
+        json.dumps(
+            [
+                {
+                    "test_name": "test_mini",
+                    "server_type": "vllm-omni",
+                    "mark": [
+                        {"hardware_marks": {"res": {"cuda": "H100"}, "num_cards": 1}},
+                        "diffusion",
+                        "full_model",
+                    ],
+                    "server_params": {"model": "m/mini"},
+                    "benchmark_params": [{"name": "p0", "num-prompts": 1, "max-concurrency": 1}],
+                }
+            ]
+        ),
+        encoding="utf-8",
+    )
+    monkeypatch.setattr(sys, "argv", ["pytest", "--test-config-file", str(cfg)])
+    sys.modules.pop("tests.dfx.perf.scripts.run_diffusion_benchmark", None)
+    from tests.dfx.perf.scripts import run_diffusion_benchmark as rdb
+
+    importlib.reload(rdb)
+
+    params = {
+        "name": "c_sweep",
+        "dataset": "random",
+        "task": "t2i",
+        "num-prompts": [8, 16, 32],
+        "max-concurrency": [1, 8, 32],
+        "baseline": {
+            "H100": {"throughput_qps": [0.1, 0.2, 0.3], "latency_mean": [10.0, 20.0, 30.0]},
+            "A3": {"throughput_qps": [0.11, 0.22, 0.33], "latency_mean": [9.0, 19.0, 29.0]},
+        },
+    }
+    run32 = rdb._build_run_params(
+        params,
+        num_prompts=32,
+        max_concurrency=32,
+        request_rate="inf",
+        sweep_index=2,
+    )
+    assert run32["max-concurrency"] == 32
+    assert run32["baseline"] == {
+        "H100": {"throughput_qps": 0.3, "latency_mean": 30.0},
+        "A3": {"throughput_qps": 0.33, "latency_mean": 29.0},
+    }
+
+    sweeps = rdb._iter_sweep_runs(params)
+    assert len(sweeps) == 3
+    assert sweeps[2]["params"]["max-concurrency"] == 32
+    assert sweeps[2]["params"]["baseline"]["H100"]["throughput_qps"] == 0.3
+    assert "A3" in sweeps[2]["params"]["baseline"]
+    assert not isinstance(sweeps[2]["params"]["baseline"]["H100"]["throughput_qps"], list)

tests/diffusion/attention/test_attention_config.py

@@ -90,6 +90,35 @@ class TestAttentionSpec:
         with pytest.raises(ValueError, match=match):
             AttentionSpec(**spec)
 
+    def test_block_sparse_defaults_applied_when_backend_selected(self):
+        spec = AttentionSpec(backend="RAINFUSION_ATTN")
+        assert spec.block_sparse.sparsity == 0.8
+        assert spec.backend_kwargs() == {"sparsity": 0.8, "start_step": 0}
+
+    def test_block_sparse_skip_layers_selector_expanded(self):
+        spec = AttentionSpec(
+            backend="RAINFUSION_ATTN",
+            block_sparse={"sparsity": 0.9, "start_step": 12, "skip_layers": "0-2,38"},
+        )
+        assert spec.block_sparse.skip_layer_indices == {0, 1, 2, 38}
+        assert spec.backend_kwargs() == {
+            "sparsity": 0.9,
+            "start_step": 12,
+            "skip_layers": [0, 1, 2, 38],
+        }
+
+    def test_block_sparse_rejected_on_dense_backend(self):
+        with pytest.raises(ValueError, match="block_sparse is only supported by"):
+            AttentionSpec(backend="FLASH_ATTN", block_sparse={"sparsity": 0.8})
+
+    @pytest.mark.parametrize(
+        "block_sparse",
+        [{"sparsity": 1.5}, {"start_step": -1}],
+    )
+    def test_block_sparse_invalid_values(self, block_sparse):
+        with pytest.raises(ValueError):
+            AttentionSpec(backend="RAINFUSION_ATTN", block_sparse=block_sparse)
+
 
 class TestAttentionConfig:
     def test_empty_config(self):

tests/diffusion/attention/test_attention_sp.py

@@ -31,6 +31,7 @@ import tempfile
 import pytest
 import torch
 
+from tests.helpers.mark import hardware_test
 from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata, QueryRange
 from vllm_omni.diffusion.attention.layer import Attention
 from vllm_omni.diffusion.attention.parallel.allgather_kv import (
@@ -409,6 +410,7 @@ def test_allgather_kv_keeps_gathered_kv_compressed_for_gqa():
     torch.testing.assert_close(v_full, expected_value)
 
 
+@hardware_test(res={"cuda": "L4"}, num_cards=4)
 @pytest.mark.parametrize(
     "test_model_cls",
     [
@@ -416,10 +418,12 @@ def test_allgather_kv_keeps_gathered_kv_compressed_for_gqa():
     ],
 )
 @pytest.mark.parametrize(
-    ("ulysses_degree", "ring_degree", "allgather_degree"),
+    ("ulysses_degree", "ring_degree", "allgather_degree", "num_kv_heads"),
     [
-        pytest.param(2, 2, 1, id="ulysses-ring"),
-        pytest.param(1, 1, 2, id="allgather-kv"),
+        pytest.param(2, 2, 1, None, id="ulysses-ring"),
+        pytest.param(1, 1, 2, None, id="allgather-kv"),
+        pytest.param(1, 2, 1, 2, id="ring-gqa"),
+        pytest.param(1, 2, 1, 1, id="ring-mqa"),
     ],
 )
 @pytest.mark.parametrize("batch_size", [2])
@@ -445,6 +449,7 @@ def test_sequence_parallel(
     seq_len: int,
     num_heads: int,
     head_size: int,
+    num_kv_heads: int | None,
 ):
     """Compare Ulysses/Ring and AllGather-KV SP against a single-rank run."""
     sequence_parallel_size = allgather_degree if allgather_degree > 1 else ulysses_degree * ring_degree
@@ -476,6 +481,7 @@ def test_sequence_parallel(
                 seq_len,
                 num_heads,
                 head_size,
+                num_kv_heads,
                 dtype,
                 causal,
                 use_sync,
@@ -507,6 +513,7 @@ def test_sequence_parallel(
                 seq_len,
                 num_heads,
                 head_size,
+                num_kv_heads,
                 dtype,
                 causal,
                 use_sync,
@@ -603,6 +610,7 @@ def ulysses_attention_on_test_model(
     seq_len: int,
     num_heads: int,
     head_size: int,
+    num_kv_heads: int | None,
     dtype: torch.dtype,
     causal: bool,
     use_sync: bool,
@@ -658,10 +666,13 @@ def ulysses_attention_on_test_model(
         cfg_parallel_size=1,
     )
 
-    od_config = OmniDiffusionConfig(
+    od_config = OmniDiffusionConfig.from_kwargs(
         model="test_model",
         dtype=dtype,
         parallel_config=parallel_config,
+        # This regression targets pytorch_attn_forward(). Do not let an
+        # installed FA/FA3 backend silently bypass the SDPA Ring path.
+        diffusion_attention_backend="TORCH_SDPA",
     )
 
     # Initialize model parallel
@@ -687,7 +698,7 @@ def ulysses_attention_on_test_model(
             "head_size": head_size,
             "hidden_size": hidden_size,
             "causal": causal,
-            "num_kv_heads": None,
+            "num_kv_heads": num_kv_heads,
             "scatter_idx": 2,
             "gather_idx": 1,
             "use_sync": use_sync,

tests/diffusion/attention/test_rainfusion_plan.py

@@ -0,0 +1,157 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Geometry contract RAINFUSION_ATTN enforces before handing a forward to rf_v2.
+
+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
+first can run sparse: rf_v2 tiles the sequence as a whole but pools its block
+mask from the video and prefix segments separately, so the two only agree on the
+block count when the video segment ends on a boundary.
+"""
+
+import dataclasses
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata, VideoTokenLayout
+from vllm_omni.diffusion.attention.backends.rainfusion_attn import (
+    _BLOCK_SIZE,
+    RainFusionAttentionImpl,
+    RainFusionPlan,
+)
+from vllm_omni.platforms import current_omni_platform
+
+PREFIX_ROWS = 710  # 14 text rows + 696 audio rows
+ALIGNED_GRID = (62, 24, 40)  # 1280x768 -> 59520 video rows, 465 blocks
+MISALIGNED_GRID = (62, 24, 42)  # 1344x768 -> 62496 video rows, 488 blocks + 32 rows
+
+
+def make_impl(**backend_kwargs):
+    return RainFusionAttentionImpl(
+        num_heads=8,
+        head_size=128,
+        softmax_scale=128**-0.5,
+        prefix="transformer_blocks.0.attn",
+        qkv_layout="BSND",
+        backend_kwargs={"sparsity": 0.8, **backend_kwargs},
+    )
+
+
+def make_metadata(grid, prefix_len=PREFIX_ROWS, max_seqlen_q=None):
+    video_rows = grid[0] * grid[1] * grid[2]
+    return AttentionMetadata(
+        extra={"max_seqlen_q": prefix_len + video_rows if max_seqlen_q is None else max_seqlen_q},
+        video_layout=VideoTokenLayout(prefix_len=prefix_len, latent_grid=grid),
+    )
+
+
+def test_block_aligned_video_segment_runs_sparse():
+    plan = make_impl()._resolve_plan(make_metadata(ALIGNED_GRID))
+
+    assert plan is not None
+    assert plan.prefix_len == PREFIX_ROWS
+    assert plan.used_len == PREFIX_ROWS + 59520
+    assert plan.latent_shape == list(ALIGNED_GRID)
+
+
+def test_misaligned_video_segment_stays_dense():
+    # Realigning the sequence would mean padding it, and rf_v2 ignores attn_mask:
+    # pad keys would take a share of every real query's softmax denominator.
+    assert make_impl()._resolve_plan(make_metadata(MISALIGNED_GRID)) is None
+
+
+@pytest.mark.parametrize("prefix_len", [1, 127, _BLOCK_SIZE, 710, 900])
+def test_prefix_length_does_not_affect_alignment(prefix_len):
+    # Only the video segment has to land on a block boundary; rf_v2 pools the
+    # prefix separately and keeps every one of its blocks.
+    plan = make_impl()._resolve_plan(make_metadata(ALIGNED_GRID, prefix_len=prefix_len))
+
+    assert plan is not None
+    assert plan.prefix_len == prefix_len
+
+
+def test_sparsity_zero_never_resolves_a_plan():
+    assert make_impl(sparsity=0.0)._resolve_plan(make_metadata(ALIGNED_GRID)) is None
+
+
+def test_missing_video_layout_falls_back_to_dense():
+    assert make_impl()._resolve_plan(AttentionMetadata(extra={"max_seqlen_q": 60230})) is None
+
+
+def test_missing_max_seqlen_falls_back_to_dense():
+    metadata = AttentionMetadata(
+        video_layout=VideoTokenLayout(prefix_len=PREFIX_ROWS, latent_grid=ALIGNED_GRID),
+    )
+
+    assert make_impl()._resolve_plan(metadata) is None
+
+
+def test_video_segment_must_be_the_tail_of_packed_document_zero():
+    metadata = make_metadata(ALIGNED_GRID, max_seqlen_q=PREFIX_ROWS + 59520 + 128)
+
+    assert make_impl()._resolve_plan(metadata) is None
+
+
+@pytest.mark.parametrize("grid", [(4, 24, 40), (1, 24, 40)])
+def test_short_video_stays_dense(grid):
+    assert make_impl()._resolve_plan(make_metadata(grid)) is None
+
+
+@pytest.mark.parametrize("qkv_layout", ["BNSD", "BSH"])
+def test_explicitly_wrong_layout_is_rejected(qkv_layout):
+    with pytest.raises(ValueError, match="BSND"):
+        RainFusionAttentionImpl(
+            num_heads=8,
+            head_size=128,
+            softmax_scale=128**-0.5,
+            qkv_layout=qkv_layout,
+            backend_kwargs={"sparsity": 0.8},
+        )
+
+
+def test_undeclared_layout_stays_dense():
+    # An undeclared layout is not an error -- the layer keeps working, just
+    # densely, and the fallback sees the same absent layout plain FLASH_ATTN would.
+    impl = RainFusionAttentionImpl(
+        num_heads=8,
+        head_size=128,
+        softmax_scale=128**-0.5,
+        prefix="transformer_blocks.0.attn",
+        backend_kwargs={"sparsity": 0.8},
+    )
+
+    assert impl._resolve_plan(make_metadata(ALIGNED_GRID)) is None
+    assert impl.dense_fallback.qkv_layout is None
+
+
+@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="rf_v2 runs on Ascend NPU only.")
+def test_fully_populated_mask_reproduces_dense_attention():
+    """At sparsity=0 every key block is kept, so rf_v2 must match dense attention.
+
+    This is the guard on the geometry the backend hands the kernel: a wrong
+    prefix length or latent grid still produces plausible output under sparse
+    selection, but shows up here as a mismatch against SDPA.
+    """
+    torch.manual_seed(0)
+    prefix_len, grid = 384, (4, 16, 64)  # 4096 video rows = 32 blocks
+    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/attention/test_trtllm_attn.py

@@ -8,6 +8,7 @@ import torch
 import torch.nn.functional as F
 
 from vllm_omni.diffusion.attention.backends import trtllm_attn as tg
+from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
 from vllm_omni.diffusion.attention.backends.trtllm_attn import TrtllmAttentionImpl
 
 pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cuda]
@@ -142,6 +143,9 @@ def test_denoise_progress_mixin(monkeypatch):
     _Pipe().record_denoise_step(4)
     assert ctx.denoise_step_idx == 4 and ctx.denoise_timestep == pytest.approx(0.25)
 
+    _Pipe().record_denoise_step(5, normalized_timestep=0.7)
+    assert ctx.denoise_step_idx == 5 and ctx.denoise_timestep == pytest.approx(0.7)
+
 
 @pytest.mark.parametrize(
     "bk, field",
@@ -187,6 +191,7 @@ def test_mask_free_allowlist_gates_trtllm_default():
 
     assert meta("WanPipeline").attention_mask_free is True
     assert meta("WanVACEPipeline").attention_mask_free is True
+    assert meta("MiniMaxH3Pipeline").attention_mask_free is True
     assert meta("FluxPipeline").attention_mask_free is False
     assert meta(None).attention_mask_free is False
 
@@ -204,6 +209,91 @@ def test_masked_layer_raises():
         _impl().forward_cuda(q, k, v, _Meta())
 
 
+def test_packed_metadata_trims_padding_mask(monkeypatch):
+    b, s, h, d = 1, 8, 8, 128
+    q, k, v = (torch.zeros(b, s, h, d) for _ in range(3))
+    mask = torch.tensor([[True, True, True, True, True, True, False, False]])
+    cu_seqlens = torch.tensor([0, 6, 8], dtype=torch.int32)
+    metadata = AttentionMetadata(
+        attn_mask=mask,
+        extra={
+            "cu_seqlens_q": cu_seqlens,
+            "cu_seqlens_k": cu_seqlens,
+            "max_seqlen_q": 6,
+            "max_seqlen_k": 6,
+        },
+    )
+    captured = {}
+
+    def fake_attention(**kwargs):
+        captured.update(kwargs)
+        return torch.zeros_like(kwargs["query"])
+
+    monkeypatch.setattr(tg, "HAS_FLASHINFER", True)
+    monkeypatch.setattr(tg, "trtllm_ragged_attention_deepseek", fake_attention)
+    monkeypatch.setattr(TrtllmAttentionImpl, "_get_workspace", classmethod(lambda cls, device: torch.empty(0)))
+
+    out = _impl().forward_cuda(q, k, v, metadata)
+
+    assert out.shape == q.shape
+    assert captured["query"].shape[0] == 6
+    assert captured["key"].shape[0] == 6
+    assert captured["value"].shape[0] == 6
+    assert captured["batch_size"] == 1
+    assert captured["seq_lens"].tolist() == [6]
+    assert captured["cum_seq_lens_q"].tolist() == [0, 6]
+    assert captured["cum_seq_lens_kv"].tolist() == [0, 6]
+    assert captured["max_q_len"] == 6
+    assert captured["max_kv_len"] == 6
+    assert torch.count_nonzero(out[:, 6:]) == 0
+
+
+def test_packed_metadata_rejects_non_prefix_mask(monkeypatch):
+    b, s, h, d = 1, 8, 8, 128
+    q, k, v = (torch.zeros(b, s, h, d) for _ in range(3))
+    cu_seqlens = torch.tensor([0, 6, 8], dtype=torch.int32)
+    metadata = AttentionMetadata(
+        attn_mask=torch.tensor([[True, True, False, True, True, True, False, False]]),
+        extra={
+            "cu_seqlens_q": cu_seqlens,
+            "cu_seqlens_k": cu_seqlens,
+            "max_seqlen_q": 6,
+            "max_seqlen_k": 6,
+        },
+    )
+    monkeypatch.setattr(tg, "HAS_FLASHINFER", True)
+
+    with pytest.raises(ValueError, match="prefix-valid"):
+        _impl().forward_cuda(q, k, v, metadata)
+
+
+def test_packed_metadata_must_be_complete():
+    b, s, h, d = 1, 8, 8, 128
+    q, k, v = (torch.zeros(b, s, h, d) for _ in range(3))
+    metadata = AttentionMetadata(extra={"cu_seqlens_q": torch.tensor([0, s], dtype=torch.int32)})
+
+    with pytest.raises(ValueError, match="Incomplete packed TRTLLM attention metadata"):
+        _impl().forward_cuda(q, k, v, metadata)
+
+
+def test_packed_metadata_must_cover_inputs(monkeypatch):
+    b, s, h, d = 1, 8, 8, 128
+    q, k, v = (torch.zeros(b, s, h, d) for _ in range(3))
+    cu_seqlens = torch.tensor([0, 6], dtype=torch.int32)
+    metadata = AttentionMetadata(
+        extra={
+            "cu_seqlens_q": cu_seqlens,
+            "cu_seqlens_k": cu_seqlens,
+            "max_seqlen_q": 6,
+            "max_seqlen_k": 6,
+        }
+    )
+    monkeypatch.setattr(tg, "HAS_FLASHINFER", True)
+
+    with pytest.raises(ValueError, match="must cover all Q/K/V tokens"):
+        _impl().forward_cuda(q, k, v, metadata)
+
+
 def test_propagate_calibration_045_schema(monkeypatch):
     import types
 
@@ -295,6 +385,30 @@ def test_bf16_dense_matches_sdpa():
     assert rel < 0.01, f"BF16 dense rel err {rel:.4f} too high"
 
 
+@requires_trtllm_attn
+def test_bf16_packed_padding_matches_sdpa():
+    torch.manual_seed(0)
+    b, s, used, h, d = 1, 256, 192, 8, 128
+    scale = 1.0 / math.sqrt(d)
+    q, k, v = (torch.randn(b, s, h, d, device="cuda", dtype=torch.bfloat16) for _ in range(3))
+    cu_seqlens = torch.tensor([0, used, s], dtype=torch.int32, device="cuda")
+    metadata = AttentionMetadata(
+        attn_mask=torch.arange(s, device="cuda")[None] < used,
+        extra={
+            "cu_seqlens_q": cu_seqlens,
+            "cu_seqlens_k": cu_seqlens,
+            "max_seqlen_q": used,
+            "max_seqlen_k": used,
+        },
+    )
+
+    out = _impl().forward_cuda(q, k, v, metadata).float()
+    ref_used = _sdpa_ref(q[:, :used], k[:, :used], v[:, :used], scale)
+    used_rel = (out[:, :used] - ref_used).abs().mean() / ref_used.abs().mean()
+    assert used_rel < 0.01, f"BF16 packed valid-token rel err {used_rel:.4f} too high"
+    assert torch.count_nonzero(out[:, used:]) == 0
+
+
 requires_sage = pytest.mark.skipif(
     not (_has_trtllm_attn() and tg._sage_quantize_fn() is not None),
     reason="requires Blackwell SM100+ GPU with flashinfer >= 0.6.16rc1 (trtllm_sage_attention_quantize)",
@@ -321,6 +435,47 @@ def test_sage_quant_matches_sdpa(dtype_qk):
     assert rel < 0.1, f"SAGE({dtype_qk}) rel err {rel:.4f} too high"
 
 
+@requires_sage
+def test_sage_packed_non_aligned_length_matches_sdpa():
+    torch.manual_seed(0)
+    b, s, used, h, d = 1, 208, 194, 8, 128
+    scale = 1.0 / math.sqrt(d)
+    q, k, v = (torch.randn(b, s, h, d, device="cuda", dtype=torch.bfloat16) for _ in range(3))
+    cu_seqlens = torch.tensor([0, used, s], dtype=torch.int32, device="cuda")
+    metadata = AttentionMetadata(
+        attn_mask=torch.arange(s, device="cuda")[None] < used,
+        extra={
+            "cu_seqlens_q": cu_seqlens,
+            "cu_seqlens_k": cu_seqlens,
+            "max_seqlen_q": used,
+            "max_seqlen_k": used,
+        },
+    )
+
+    out = _impl(quant={"dtype_qk": "fp8_e4m3", "q_block_size": 1, "k_block_size": 16}).forward_cuda(q, k, v, metadata)
+    assert torch.isfinite(out).all()
+    ref = _sdpa_ref(q[:, :used], k[:, :used], v[:, :used], scale)
+    rel = (out[:, :used].float() - ref).abs().mean() / ref.abs().mean()
+    assert rel < 0.2, f"SAGE packed valid-token rel err {rel:.4f} too high"
+    assert torch.count_nonzero(out[:, used:]) == 0
+
+
+@requires_sage
+def test_sage_short_sequence_uses_dense_kernel():
+    torch.manual_seed(0)
+    b, s, h, d = 1, 14, 8, 128
+    scale = 1.0 / math.sqrt(d)
+    q, k, v = (torch.randn(b, s, h, d, device="cuda", dtype=torch.bfloat16) for _ in range(3))
+
+    impl = _impl(quant={"dtype_qk": "fp8_e4m3", "q_block_size": 1, "k_block_size": 16})
+    impl._sage_quantize_fn = lambda *args, **kwargs: pytest.fail("short sequence must not use SAGE")
+    out = impl.forward_cuda(q, k, v)
+    ref = _sdpa_ref(q, k, v, scale)
+    rel = (out.float() - ref).abs().mean() / ref.abs().mean()
+    assert torch.isfinite(out).all()
+    assert rel < 0.01, f"SAGE short-sequence dense fallback rel err {rel:.4f} too high"
+
+
 @requires_trtllm_attn
 def test_skip_tiny_threshold_matches_dense():
     torch.manual_seed(0)

tests/diffusion/attention/test_ulysses_uaa.py

@@ -6,12 +6,14 @@ from __future__ import annotations
 import os
 import socket
 import tempfile
+from types import SimpleNamespace
 
 import numpy as np
 import pytest
 import torch
 
 from vllm_omni.diffusion.attention.layer import Attention
+from vllm_omni.diffusion.attention.parallel.ulysses import UlyssesParallelAttention
 from vllm_omni.diffusion.config import set_current_diffusion_config
 from vllm_omni.diffusion.data import DiffusionParallelConfig, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.parallel_state import (
@@ -39,6 +41,33 @@ def _set_dist_env(*, rank: int, world_size: int, master_port: int) -> None:
     os.environ["MASTER_PORT"] = str(master_port)
 
 
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_advanced_uaa_hybrid_rejects_padded_mqa_heads(monkeypatch) -> None:
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.attention.parallel.ulysses.get_ulysses_mode",
+        lambda default: "advanced_uaa",
+    )
+    sp_group = SimpleNamespace(
+        ulysses_group=None,
+        ulysses_world_size=2,
+        ulysses_rank=0,
+        ring_world_size=2,
+    )
+    strategy = UlyssesParallelAttention(
+        sp_group=sp_group,
+        scatter_idx=2,
+        gather_idx=1,
+        use_sync=False,
+    )
+    query = torch.zeros(1, 2, 8, 4)
+    key = torch.zeros(1, 2, 1, 4)
+    value = torch.zeros_like(key)
+
+    with pytest.raises(ValueError, match="does not support K/V head padding"):
+        strategy.pre_attention(query, key, value, None)
+
+
 def _run_attention_case(
     local_rank: int,
     world_size: int,
@@ -241,7 +270,9 @@ def test_ulysses_uaa_hybrid_ring_matches_baseline() -> None:
     batch_size = 2
     head_size = 8
     seq_len = 10
-    num_heads = 3  # head_cnt not divisible by ulysses_degree=2 -> triggers head padding
+    # This is the positive Hybrid case: K/V heads must not require padding,
+    # which is unsupported by advanced_uaa when Ring is also enabled.
+    num_heads = 4
 
     # Ensure ring ranks see equal post-Ulysses seq_len:
     # rank0/1 -> 3+2=5, rank2/3 -> 3+2=5

tests/diffusion/cache/test_cache_backends.py

@@ -11,6 +11,7 @@ This module tests the cache backend implementations:
 - DiffusionCacheConfig: configuration dataclass
 """
 
+from types import SimpleNamespace
 from unittest.mock import Mock, patch
 
 import cache_dit
@@ -21,6 +22,7 @@ from vllm_omni.diffusion.cache.cachedit import (
     CacheDiTAdapterConfig,
     CacheDiTBackend,
     CacheDiTConfig,
+    cache_summary,
 )
 from vllm_omni.diffusion.cache.magcache import MagCacheBackend
 from vllm_omni.diffusion.cache.selector import get_cache_backend
@@ -86,9 +88,84 @@ class TestCacheDiTBackend:
 
         # Verify cache-dit was enabled
         assert backend.enabled is True
-        assert backend._refresh_func is not None
+        assert backend._refresh_funcs
         mock_cache_dit.enable_cache.assert_called_once()
 
+    @patch("vllm_omni.diffusion.cache.cachedit.backend.BlockAdapter")
+    @patch("vllm_omni.diffusion.cache.cachedit.backend.cache_dit")
+    def test_lifecycle_applies_to_every_declared_dit(self, mock_cache_dit, mock_block_adapter):
+        """Combined pipelines manage every exact Cache-DiT target."""
+        pipeline = Mock()
+        pipeline.__class__.__name__ = "MiniMaxH3Pipeline"
+        pipeline._dit_modules = ["transformer", "transformers_ref"]
+        pipeline.transformer = Mock()
+        pipeline.transformers_ref = Mock()
+        for transformer in (pipeline.transformer, pipeline.transformers_ref):
+            transformer._cache_dit_adapter_config = CacheDiTAdapterConfig(
+                block_forward_patterns={"blocks": ForwardPattern.Pattern_3}
+            )
+        installed_adapters = [Mock(), Mock()]
+        mock_block_adapter.side_effect = installed_adapters
+
+        backend = CacheDiTBackend({"Fn_compute_blocks": 2})
+        backend.enable(pipeline)
+        backend.refresh(pipeline, num_inference_steps=20)
+        backend.disable(pipeline)
+
+        assert mock_cache_dit.enable_cache.call_count == 2
+        assert mock_cache_dit.refresh_context.call_count == 2
+        assert {call.args[0] for call in mock_cache_dit.refresh_context.call_args_list} == {
+            pipeline.transformer,
+            pipeline.transformers_ref,
+        }
+        assert [call.args[0] for call in mock_cache_dit.disable_cache.call_args_list] == installed_adapters
+        assert backend._refresh_funcs == []
+        assert backend._cache_targets == []
+        assert not backend.is_enabled()
+
+    @patch("vllm_omni.diffusion.cache.cachedit.backend.cache_dit")
+    def test_summary_skips_uncached_nested_dit(self, mock_cache_dit):
+        """Only DiT modules with an active Cache-DiT context are summarized."""
+        language_model = Mock()
+        language_model._is_cached = False
+        transformer = Mock()
+        transformer._is_cached = True
+        transformer.language_model = language_model
+        pipeline = SimpleNamespace(
+            _dit_modules=["transformer.language_model", "transformer"],
+            transformer=transformer,
+        )
+
+        cache_summary(pipeline)
+
+        mock_cache_dit.summary.assert_called_once_with(transformer, details=True)
+
+    @patch("vllm_omni.diffusion.cache.cachedit.backend.BlockAdapter")
+    @patch("vllm_omni.diffusion.cache.cachedit.backend.cache_dit")
+    def test_enable_and_refresh_nested_declared_dit(self, mock_cache_dit, mock_block_adapter):
+        """Dotted component paths resolve to their nested DiT module."""
+        transformer = Mock()
+        transformer._cache_dit_adapter_config = CacheDiTAdapterConfig(
+            block_forward_patterns={"blocks": ForwardPattern.Pattern_3}
+        )
+        pipeline = SimpleNamespace(
+            _dit_modules=["language_model.model"],
+            language_model=SimpleNamespace(model=transformer),
+        )
+
+        backend = CacheDiTBackend({"Fn_compute_blocks": 2})
+        backend.enable(pipeline)
+        backend.refresh(pipeline, num_inference_steps=20)
+
+        cache_summary(pipeline)
+        mock_cache_dit.enable_cache.assert_called_once()
+        mock_cache_dit.refresh_context.assert_called_once_with(
+            transformer,
+            num_inference_steps=20,
+            verbose=True,
+        )
+        mock_cache_dit.summary.assert_called_once_with(transformer, details=True)
+
     @patch("vllm_omni.diffusion.cache.cachedit.backend.logger")
     @patch("vllm_omni.diffusion.cache.cachedit.backend.cache_dit")
     def test_enable_without_model_adapter_uses_cache_dit_registry(self, mock_cache_dit, mock_logger):
@@ -106,7 +183,7 @@ class TestCacheDiTBackend:
         backend.enable(pipeline)
 
         assert backend.enabled is True
-        assert backend._refresh_func is not None
+        assert backend._refresh_funcs
         assert mock_cache_dit.enable_cache.call_args.args[0] is pipeline.transformer
         mock_logger.info.assert_any_call(
             "Transformer %s does not declare _cache_dit_adapter_config; "
@@ -138,7 +215,7 @@ class TestCacheDiTBackend:
 
         assert isinstance(exc_info.value.__cause__, ValueError)
         assert backend.enabled is False
-        assert backend._refresh_func is None
+        assert backend._refresh_funcs == []
 
     @patch("vllm_omni.diffusion.cache.cachedit.backend.BlockAdapter")
     @patch("vllm_omni.diffusion.cache.cachedit.backend.cache_dit")
@@ -220,6 +297,32 @@ class TestCacheDiTBackend:
 
         assert mock_cache_dit.refresh_context.call_count == 2
 
+    @patch("vllm_omni.diffusion.cache.cachedit.backend.BlockAdapter")
+    @patch("vllm_omni.diffusion.cache.cachedit.backend.cache_dit")
+    def test_refresh_rearms_force_refresh_hint(self, mock_cache_dit, mock_block_adapter):
+        """A once hint must be restored for every repeated request."""
+        pipeline = Mock()
+        pipeline.__class__.__name__ = "DiTPipeline"
+        pipeline.transformer = Mock()
+        pipeline.transformer._cache_dit_adapter_config = CacheDiTAdapterConfig(
+            block_forward_patterns={"layers": ForwardPattern.Pattern_0}
+        )
+
+        config = DiffusionCacheConfig(
+            force_refresh_step_hint=7,
+            force_refresh_step_policy="once",
+        )
+        backend = CacheDiTBackend(config)
+        backend.enable(pipeline)
+        backend.refresh(pipeline, num_inference_steps=20)
+        backend.refresh(pipeline, num_inference_steps=20)
+
+        assert mock_cache_dit.refresh_context.call_count == 2
+        for call in mock_cache_dit.refresh_context.call_args_list:
+            refresh_config = call.kwargs["cache_config"]
+            assert refresh_config.force_refresh_step_hint == 7
+            assert refresh_config.force_refresh_step_policy == "once"
+
     @patch("vllm_omni.diffusion.cache.cachedit.backend.BlockAdapter")
     @patch("vllm_omni.diffusion.cache.cachedit.backend.cache_dit")
     def test_enable_hunyuan_pipeline_uses_model_transformer(self, mock_cache_dit, mock_block_adapter):
@@ -245,7 +348,7 @@ class TestCacheDiTBackend:
         backend.enable(mock_pipeline)
 
         assert backend.enabled is True
-        assert backend._refresh_func is not None
+        assert backend._refresh_funcs
         mock_block_adapter.assert_called_once()
         adapter_kwargs = mock_block_adapter.call_args.kwargs
         assert adapter_kwargs["transformer"] is mock_pipeline.model
@@ -285,7 +388,7 @@ class TestCacheDiTBackend:
         backend.enable(mock_pipeline)
 
         assert backend.enabled is True
-        assert backend._refresh_func is not None
+        assert backend._refresh_funcs
         mock_block_adapter.assert_called_once()
         adapter_kwargs = mock_block_adapter.call_args.kwargs
         assert adapter_kwargs["transformer"] is mock_pipeline.transformer

tests/diffusion/cache/test_cache_dit.py

@@ -81,6 +81,8 @@ def test_cachedit_public_api_is_explicit():
         "CacheDiTAdapterConfig",
         "CacheDiTBackend",
         "CacheDiTConfig",
+        "CacheDiTRequestSpec",
+        "RequestScopedCacheDiTRuntime",
         "SensenovaCachedAdapter",
         "cache_summary",
         "enable_cache_for_dit",

tests/diffusion/cache/test_cache_dit_request_runtime.py

@@ -0,0 +1,101 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+
+from vllm_omni.diffusion.cache.cachedit import runtime as runtime_module
+from vllm_omni.diffusion.cache.cachedit.runtime import (
+    CacheDiTRequestSpec,
+    RequestScopedCacheDiTRuntime,
+)
+from vllm_omni.diffusion.data import DiffusionCacheConfig
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+def _spec(key: str, steps: int = 50) -> CacheDiTRequestSpec:
+    return CacheDiTRequestSpec(
+        installation_key=key,
+        cache_config=DiffusionCacheConfig(),
+        num_inference_steps=steps,
+    )
+
+
+def _recording_backend(events):
+    class FakeBackend:
+        def __init__(self, config):
+            events.append(("create", config))
+            self.enabled = False
+
+        def enable(self, pipeline):
+            events.append(("enable", pipeline))
+            self.enabled = True
+
+        def refresh(self, pipeline, num_inference_steps):
+            events.append(("refresh", pipeline, num_inference_steps))
+
+        def disable(self, pipeline):
+            events.append(("disable", pipeline))
+            self.enabled = False
+
+        def is_enabled(self):
+            return self.enabled
+
+    return FakeBackend
+
+
+def test_startup_cache_follows_omit_lossless_high_omit_state_machine(monkeypatch):
+    events = []
+    FakeBackend = _recording_backend(events)
+    monkeypatch.setattr(runtime_module, "CacheDiTBackend", FakeBackend)
+    pipeline = object()
+    config = DiffusionCacheConfig()
+    runtime = RequestScopedCacheDiTRuntime(pipeline)
+    startup_backend = FakeBackend(config)
+    startup_backend.enabled = True
+
+    runtime.adopt(startup_backend, installation_key="generic")
+    runtime.prepare(_spec("generic", 50))  # omitted: keep startup cache
+    runtime.prepare(None)  # lossless: uninstall
+    runtime.prepare(None)  # repeated lossless: no-op
+    runtime.prepare(_spec("high", 40))  # high: install H3 profile
+    runtime.prepare(_spec("high", 30))  # repeated high: refresh only
+    runtime.prepare(_spec("generic", 20))  # omitted: restore startup profile
+
+    assert [event[0] for event in events] == [
+        "create",
+        "refresh",
+        "disable",
+        "create",
+        "enable",
+        "refresh",
+        "refresh",
+        "disable",
+        "create",
+        "enable",
+        "refresh",
+    ]
+    assert events[1][-1] == 50
+    assert events[5][-1] == 40
+    assert events[6][-1] == 30
+    assert events[10][-1] == 20
+    assert runtime.is_enabled
+
+
+def test_high_installs_without_startup_cache_and_lossless_uninstalls(monkeypatch):
+    events = []
+    monkeypatch.setattr(runtime_module, "CacheDiTBackend", _recording_backend(events))
+    pipeline = object()
+    runtime = RequestScopedCacheDiTRuntime(pipeline)
+
+    runtime.prepare(_spec("high", 40))
+    runtime.prepare(None)
+
+    assert [event[0] for event in events] == [
+        "create",
+        "enable",
+        "refresh",
+        "disable",
+    ]
+    assert events[2][-1] == 40
+    assert not runtime.is_enabled

tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import json
 from multiprocessing.reduction import ForkingPickler
 from types import SimpleNamespace
 from unittest.mock import Mock
@@ -12,6 +13,70 @@ import torch
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
 
 
+def test_h3_prepares_resolved_cache_state_immediately_before_denoise():
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+    from vllm_omni.diffusion.request import OmniDiffusionRequest
+    from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
+    from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline.partition = "fl2va"
+    pipeline.supported_tasks = frozenset({"t2va"})
+    pipeline.default_video_shift = 12.0
+    pipeline.default_audio_shift = 3.0
+    pipeline.device = torch.device("cpu")
+    pipeline.od_config = SimpleNamespace()
+    cache_spec = object()
+    quality_plan = SimpleNamespace(cache_dit=cache_spec)
+    pipeline._quality_policy = Mock()
+    pipeline._quality_policy.resolve.return_value = quality_plan
+    events = []
+    pipeline._cache_dit_runtime = SimpleNamespace(prepare=lambda spec: events.append(("prepare", spec)))
+    pipeline.encode_prompt = Mock(
+        return_value=(
+            torch.ones(1, 2),
+            torch.ones(1, dtype=torch.long),
+        )
+    )
+
+    def diffuse(**kwargs):
+        events.append(("diffuse", kwargs))
+        return torch.zeros(1), torch.zeros(1)
+
+    pipeline.diffuse = diffuse
+    pipeline.decode = Mock(return_value=(torch.zeros(1), torch.zeros(1)))
+    sampling = OmniDiffusionSamplingParams(
+        quality="high",
+        width=1344,
+        height=768,
+        fps=24,
+        num_frames=124,
+        num_inference_steps=50,
+        extra_args={"task": "t2va", "aspect_ratio": "16:9"},
+    )
+    batch = DiffusionRequestBatch(
+        [
+            OmniDiffusionRequest(
+                prompt="quality boundary",
+                sampling_params=sampling,
+                request_id="quality-boundary",
+            )
+        ]
+    )
+
+    output = pipeline.forward(batch)
+
+    assert events[0] == ("prepare", cache_spec)
+    assert events[1][0] == "diffuse"
+    pipeline._quality_policy.resolve.assert_called_once_with(
+        quality="high",
+        num_inference_steps=50,
+        extra_args={"task": "t2va", "aspect_ratio": "16:9"},
+    )
+    assert output.output == pipeline.decode.return_value
+
+
 def test_pipeline_import_registry_and_component_discovery():
     from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
     from vllm_omni.diffusion.registry import (
@@ -24,12 +89,336 @@ def test_pipeline_import_registry_and_component_discovery():
         "pipeline_minimax_h3",
         "MiniMaxH3Pipeline",
     )
+    assert _DIFFUSION_MODELS["MiniMaxH3ModularPipeline"] == _DIFFUSION_MODELS["MiniMaxH3Pipeline"]
     assert _DIFFUSION_POST_PROCESS_FUNCS["MiniMaxH3Pipeline"] == "get_minimax_h3_post_process_func"
-    assert MiniMaxH3Pipeline._dit_modules == ["transformer"]
+    assert (
+        _DIFFUSION_POST_PROCESS_FUNCS["MiniMaxH3ModularPipeline"] == _DIFFUSION_POST_PROCESS_FUNCS["MiniMaxH3Pipeline"]
+    )
+    assert MiniMaxH3Pipeline._dit_modules == ["transformer", "transformers_ref"]
     assert MiniMaxH3Pipeline._encoder_modules == ["text_encoder"]
     assert MiniMaxH3Pipeline._vae_modules == ["video_vae", "audio_vae"]
 
 
+def _write_partition_index(path, *, partition, tasks):
+    path.mkdir(parents=True)
+    (path / "model_index.json").write_text(
+        json.dumps(
+            {
+                "_class_name": "MiniMaxH3Pipeline",
+                "_diffusers_version": "0.32.2",
+                "_minimax_h3": {
+                    "partition": partition,
+                    "tasks": tasks,
+                    "sigma_shift_scales": {"video": 12.0, "audio": 3.0},
+                },
+            }
+        ),
+        encoding="utf-8",
+    )
+
+
+def test_modular_diffusers_index_is_resolved_generically(tmp_path):
+    from vllm_omni.diffusion.data import OmniDiffusionConfig, resolve_model_class_name
+    from vllm_omni.diffusion.utils.hf_utils import is_diffusion_model
+    from vllm_omni.entrypoints.utils import resolve_model_config_path
+
+    (tmp_path / "modular_model_index.json").write_text(
+        json.dumps(
+            {
+                "_class_name": "MiniMaxH3ModularPipeline",
+                "_diffusers_version": "0.36.0.dev0",
+            }
+        ),
+        encoding="utf-8",
+    )
+
+    assert is_diffusion_model(str(tmp_path))
+    assert resolve_model_class_name(str(tmp_path)) == "MiniMaxH3ModularPipeline"
+    assert resolve_model_config_path(str(tmp_path)) is None
+
+    config = OmniDiffusionConfig(model=str(tmp_path))
+    config.enrich_config()
+    assert config.model_class_name == "MiniMaxH3ModularPipeline"
+    assert config.supports_multimodal_inputs
+    assert config.max_multimodal_image_inputs == 9
+    assert config.supports_mixed_reference_inputs
+
+
+@pytest.mark.parametrize(
+    ("task_type", "partition"),
+    [
+        (None, "combined"),
+        ("auto", "combined"),
+        ("t2va", "fl2va"),
+        ("fl2va", "fl2va"),
+        ("ref2va", "ref2va"),
+    ],
+)
+def test_startup_task_selects_weight_partition(task_type, partition):
+    from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import _minimax_h3_partition_for_task
+
+    assert _minimax_h3_partition_for_task(task_type) == partition
+
+
+def test_startup_auto_task_uses_explicit_local_partition(tmp_path):
+    from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import _minimax_h3_partition_for_task
+
+    for directory, partition, tasks in (
+        ("FL2VA", "fl2va", ["t2va", "fl2va"]),
+        ("Ref2VA", "ref2va", ["ref2va"]),
+    ):
+        model_path = tmp_path / directory
+        _write_partition_index(model_path, partition=partition, tasks=tasks)
+        assert _minimax_h3_partition_for_task(None, str(model_path)) == partition
+        assert _minimax_h3_partition_for_task("auto", str(model_path)) == partition
+        assert _minimax_h3_partition_for_task("combined", str(model_path)) == "combined"
+
+
+def test_startup_task_rejects_unsupported_value():
+    from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import _minimax_h3_partition_for_task
+
+    with pytest.raises(ValueError, match="task_type must be one of"):
+        _minimax_h3_partition_for_task("unsupported")
+
+
+def test_combined_task_inference_and_transformer_routing():
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline.partition = "combined"
+    pipeline.supported_tasks = frozenset({"t2va", "fl2va", "ref2va"})
+    pipeline.transformer = torch.nn.Identity()
+    pipeline.transformers_ref = torch.nn.Linear(1, 1)
+    assert pipeline._resolve_task(None, {}) == "t2va"
+    assert pipeline._resolve_task(None, {"image": object()}) == "fl2va"
+    assert pipeline._resolve_task(None, {"audio": object()}) == "ref2va"
+    assert pipeline._resolve_task(None, {"video": object()}) == "ref2va"
+    pipeline.partition = "ref2va"
+    pipeline.supported_tasks = frozenset({"ref2va"})
+    assert pipeline._resolve_task(None, {"image": object()}) == "ref2va"
+
+    pipeline.partition = "combined"
+    pipeline.supported_tasks = frozenset({"t2va", "fl2va", "ref2va"})
+    assert pipeline._transformer_for_task("fl2va") is pipeline.transformer
+    assert pipeline._transformer_for_task("ref2va") is pipeline.transformers_ref
+
+
+@pytest.mark.parametrize(
+    (
+        "task_type",
+        "expected_partition",
+        "expected_dits",
+        "component_partition",
+        "source_partitions",
+        "expected_tasks",
+    ),
+    [
+        (None, "combined", 2, "FL2VA", ["FL2VA", "Ref2VA"], {"t2va", "fl2va", "ref2va"}),
+        ("fl2va", "fl2va", 1, "FL2VA", ["FL2VA"], {"t2va", "fl2va"}),
+        ("ref2va", "ref2va", 1, "Ref2VA", ["Ref2VA"], {"ref2va"}),
+    ],
+)
+def test_pipeline_loads_task_selected_dits_and_shared_components_once(
+    monkeypatch,
+    tmp_path,
+    task_type,
+    expected_partition,
+    expected_dits,
+    component_partition,
+    source_partitions,
+    expected_tasks,
+):
+    from vllm_omni.diffusion.data import (
+        DiffusionParallelConfig,
+        OmniDiffusionConfig,
+    )
+    from vllm_omni.diffusion.models.minimax_h3 import (
+        pipeline_minimax_h3 as pipeline_module,
+    )
+
+    _write_partition_index(
+        tmp_path / "FL2VA",
+        partition="fl2va",
+        tasks=["t2va", "fl2va"],
+    )
+    _write_partition_index(
+        tmp_path / "Ref2VA",
+        partition="ref2va",
+        tasks=["ref2va"],
+    )
+    for partition_name in ("FL2VA", "Ref2VA"):
+        for component in (
+            "transformer",
+            "tokenizer",
+            "processor",
+            "text_encoder",
+            "video_vae",
+            "audio_vae",
+        ):
+            (tmp_path / partition_name / component).mkdir()
+
+    created = {"dit": [], "text_encoder": [], "video_vae": [], "audio_vae": []}
+
+    class FakeModule(torch.nn.Module):
+        def __init__(self, *args, **kwargs):
+            super().__init__()
+
+    class FakeDiT(FakeModule):
+        def __init__(self, *args, **kwargs):
+            super().__init__()
+            created["dit"].append((args, kwargs))
+
+    def component_factory(name):
+        def create(path, *args, **kwargs):
+            created[name].append(str(path))
+            return FakeModule()
+
+        return create
+
+    tokenizer_calls = []
+    processor_calls = []
+    tokenizer_cls = Mock(spec=pipeline_module.Qwen2TokenizerFast)
+    tokenizer_cls.from_pretrained.side_effect = lambda *args, **kwargs: (
+        tokenizer_calls.append((args, kwargs)) or object()
+    )
+    processor_cls = Mock(spec=pipeline_module.Qwen3VLProcessor)
+    processor_cls.from_pretrained.side_effect = lambda *args, **kwargs: (
+        processor_calls.append((args, kwargs)) or object()
+    )
+    monkeypatch.setattr(pipeline_module, "MiniMaxH3DiTModel", FakeDiT)
+    monkeypatch.setattr(
+        pipeline_module,
+        "MiniMaxH3Qwen3VLEncoder",
+        component_factory("text_encoder"),
+    )
+    monkeypatch.setattr(
+        pipeline_module,
+        "MiniMaxH3VideoVAE",
+        component_factory("video_vae"),
+    )
+    monkeypatch.setattr(
+        pipeline_module,
+        "MiniMaxH3AudioVAE",
+        component_factory("audio_vae"),
+    )
+    monkeypatch.setattr(
+        pipeline_module,
+        "Qwen2TokenizerFast",
+        tokenizer_cls,
+    )
+    monkeypatch.setattr(
+        pipeline_module,
+        "Qwen3VLProcessor",
+        processor_cls,
+    )
+    monkeypatch.setattr(
+        pipeline_module,
+        "_dit_rank_world",
+        lambda: (None, 0, 1),
+    )
+    monkeypatch.setattr(
+        pipeline_module,
+        "get_local_device",
+        lambda: torch.device("cpu"),
+    )
+    download_calls = []
+
+    def fake_download(**kwargs):
+        download_calls.append(kwargs)
+        return str(tmp_path)
+
+    monkeypatch.setattr(
+        pipeline_module,
+        "download_weights_from_hf_specific",
+        fake_download,
+    )
+
+    od_config = OmniDiffusionConfig(
+        model="MiniMaxAI/MiniMax-H3",
+        revision=None,
+        task_type=task_type,
+        quantization_config=None,
+        parallel_config=DiffusionParallelConfig(
+            cfg_parallel_size=1,
+            text_encoder_tp_size=1,
+        ),
+    )
+    pipeline = pipeline_module.MiniMaxH3Pipeline(od_config=od_config)
+
+    assert pipeline.partition == expected_partition
+    assert pipeline.supported_tasks == expected_tasks
+    assert len(created["dit"]) == expected_dits
+    component_path = tmp_path / component_partition
+    assert created["text_encoder"] == [str(component_path / "text_encoder")]
+    assert created["video_vae"] == [str(component_path / "video_vae")]
+    assert created["audio_vae"] == [str(component_path / "audio_vae")]
+    assert len(tokenizer_calls) == 1
+    assert len(processor_calls) == 1
+    expected_dit_modules = ["transformer", "transformers_ref"] if expected_dits == 2 else ["transformer"]
+    assert pipeline._dit_modules == expected_dit_modules
+    assert hasattr(pipeline, "transformers_ref") is (expected_dits == 2)
+    if expected_partition == "ref2va":
+        assert pipeline._transformer_for_task("ref2va") is pipeline.transformer
+    assert [source.model_or_path for source in pipeline.weights_sources] == [
+        str(tmp_path / partition_name) for partition_name in source_partitions
+    ]
+    expected_patterns = (
+        pipeline_module.MINIMAX_H3_DOWNLOAD_PATTERNS
+        if expected_partition == "combined"
+        else pipeline_module.MINIMAX_H3_TASK_DOWNLOAD_PATTERNS[expected_partition]
+    )
+    assert download_calls == [
+        {
+            "model_name_or_path": "MiniMaxAI/MiniMax-H3",
+            "cache_dir": None,
+            "allow_patterns": expected_patterns,
+            "revision": None,
+            "require_all": True,
+        }
+    ]
+
+
+def test_combined_weight_loader_routes_each_contiguous_partition():
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+
+    class FakeTransformer(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.loaded = []
+            self.post_load_calls = 0
+
+        def load_weights(self, weights):
+            self.loaded = [name for name, _ in weights]
+            return set(self.loaded)
+
+        def post_load_weights(self):
+            self.post_load_calls += 1
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline.transformer = FakeTransformer()
+    pipeline.transformers_ref = FakeTransformer()
+    pipeline.text_encoder = torch.nn.Identity()
+    pipeline.video_vae = torch.nn.Identity()
+    pipeline.audio_vae = torch.nn.Identity()
+    loaded = pipeline.load_weights(
+        iter(
+            [
+                ("transformer.a", torch.ones(1)),
+                ("transformer.b", torch.ones(1)),
+                ("transformers_ref.a", torch.ones(1)),
+            ]
+        )
+    )
+
+    assert pipeline.transformer.loaded == ["a", "b"]
+    assert pipeline.transformers_ref.loaded == ["a"]
+    assert pipeline.transformer.post_load_calls == 1
+    assert pipeline.transformers_ref.post_load_calls == 1
+    assert loaded == {"transformer.a", "transformer.b", "transformers_ref.a"}
+
+
 def test_dlo_offload_plan_includes_token_refiner():
     from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
 
@@ -451,6 +840,21 @@ def test_distributed_layerwise_resident_blocks_release_on_failure():
     controller.offload_resident_layers.assert_called_once_with()
 
 
+def test_distributed_layerwise_resident_blocks_can_be_skipped():
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    controller = Mock()
+    pipeline._dlo_residency_controller = controller
+
+    with pipeline._resident_dit_layers_on_device(enabled=False):
+        pass
+
+    controller.load_resident_layers.assert_not_called()
+    controller.offload_resident_layers.assert_not_called()
+
+
 def test_encoder_layerwise_offload_keeps_tp_blocks_rank_local(monkeypatch):
     from vllm_omni.diffusion.models.minimax_h3.encoder import (
         MiniMaxH3Qwen3VLEncoder,

tests/diffusion/models/minimax_h3/test_minimax_h3_quality_policy.py

@@ -0,0 +1,212 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+
+import pytest
+
+from vllm_omni.diffusion.data import DiffusionCacheConfig
+from vllm_omni.errors import OmniClientError
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+def _quality_od_config(*, cache_backend: str = "cache_dit", **overrides):
+    values = {
+        "cache_backend": cache_backend,
+        "cache_config": DiffusionCacheConfig(),
+        "diffusion_attention_config": SimpleNamespace(default=None, per_role={}),
+        "enable_cpu_offload": False,
+        "enable_layerwise_offload": False,
+        "enable_distributed_layerwise_offload": False,
+        "enforce_eager": True,
+        "num_gpus": 4,
+        "quantization_config": None,
+        "parallel_config": SimpleNamespace(
+            allgather_degree=1,
+            cfg_parallel_size=1,
+            data_parallel_size=1,
+            pipeline_parallel_size=1,
+            ring_degree=1,
+            sequence_parallel_size=4,
+            tensor_parallel_size=1,
+            ulysses_degree=4,
+            use_hsdp=False,
+        ),
+    }
+    values.update(overrides)
+    return SimpleNamespace(**values)
+
+
+def _resolve_quality(policy, **overrides):
+    values = {
+        "quality": "high",
+        "num_inference_steps": 50,
+        "extra_args": None,
+    }
+    values.update(overrides)
+    return policy.resolve(**values)
+
+
+def test_high_quality_policy_emits_cache_profile_without_deployment_gates():
+    from vllm_omni.diffusion.models.minimax_h3.quality_policy import (
+        MiniMaxH3QualityPolicy,
+    )
+
+    policy = MiniMaxH3QualityPolicy(
+        _quality_od_config(
+            cache_backend="none",
+            enable_layerwise_offload=True,
+            enforce_eager=False,
+            num_gpus=1,
+        )
+    )
+
+    plan = _resolve_quality(policy, num_inference_steps=17)
+    spec = plan.cache_dit
+
+    assert spec is not None
+    assert spec.installation_key == "minimax_h3.high"
+    assert spec.num_inference_steps == 17
+    assert spec.cache_config.Fn_compute_blocks == 1
+    assert spec.cache_config.Bn_compute_blocks == 0
+    assert spec.cache_config.max_warmup_steps == 4
+    assert spec.cache_config.residual_diff_threshold == pytest.approx(0.04)
+    assert spec.cache_config.max_continuous_cached_steps == 1
+    assert spec.cache_config.enable_taylorseer is False
+    assert spec.cache_config.scm_steps_mask_policy is None
+
+
+def test_h3_force_refresh_is_model_owned_and_request_scoped():
+    from vllm_omni.diffusion.models.minimax_h3.quality_policy import MiniMaxH3QualityPolicy
+
+    policy = MiniMaxH3QualityPolicy(_quality_od_config())
+    plan = _resolve_quality(
+        policy,
+        quality="high",
+        num_inference_steps=17,
+        extra_args={"force_refresh_step_hint": 7},
+    )
+
+    assert plan.cache_dit is not None
+    assert plan.cache_dit.installation_key == "minimax_h3.high:force_refresh=7:once"
+    assert plan.cache_dit.cache_config.force_refresh_step_hint == 7
+    assert plan.cache_dit.cache_config.force_refresh_step_policy == "once"
+
+
+def test_h3_force_refresh_policy_supports_repeat_and_does_not_mutate_generic_config():
+    from vllm_omni.diffusion.models.minimax_h3.quality_policy import MiniMaxH3QualityPolicy
+
+    config = _quality_od_config()
+    plan = _resolve_quality(
+        MiniMaxH3QualityPolicy(config),
+        quality=None,
+        num_inference_steps=20,
+        extra_args={"force_refresh_step_hint": 5, "force_refresh_step_policy": "repeat"},
+    )
+
+    assert plan.cache_dit is not None
+    assert plan.cache_dit.installation_key == "minimax_h3.generic:force_refresh=5:repeat"
+    assert plan.cache_dit.cache_config is not config.cache_config
+    assert plan.cache_dit.cache_config.force_refresh_step_hint == 5
+    assert plan.cache_dit.cache_config.force_refresh_step_policy == "repeat"
+    assert config.cache_config.force_refresh_step_hint is None
+    assert config.cache_config.force_refresh_step_policy == "once"
+
+    restored = _resolve_quality(MiniMaxH3QualityPolicy(config), quality=None)
+    assert restored.cache_dit is not None
+    assert restored.cache_dit.installation_key == "minimax_h3.generic"
+    assert restored.cache_dit.cache_config is config.cache_config
+
+
+@pytest.mark.parametrize(
+    ("extra_args", "message"),
+    [
+        ({"force_refresh_step_hint": True}, "positive integer"),
+        ({"force_refresh_step_hint": 0}, "between 1"),
+        ({"force_refresh_step_hint": 51}, "between 1"),
+        ({"force_refresh_step_hint": 5, "force_refresh_step_policy": "always"}, "one of"),
+        ({"force_refresh_step_policy": "repeat"}, "requires force_refresh_step_hint"),
+    ],
+)
+def test_h3_force_refresh_rejects_invalid_values(extra_args, message):
+    from vllm_omni.diffusion.models.minimax_h3.quality_policy import MiniMaxH3QualityPolicy
+
+    with pytest.raises(OmniClientError, match=message):
+        _resolve_quality(
+            MiniMaxH3QualityPolicy(_quality_od_config()),
+            extra_args=extra_args,
+        )
+
+
+@pytest.mark.parametrize(
+    ("cache_backend", "quality"),
+    [
+        ("cache_dit", "lossless"),
+        ("none", None),
+    ],
+)
+def test_h3_force_refresh_requires_active_cache_target(cache_backend, quality):
+    from vllm_omni.diffusion.models.minimax_h3.quality_policy import MiniMaxH3QualityPolicy
+
+    policy = MiniMaxH3QualityPolicy(_quality_od_config(cache_backend=cache_backend))
+
+    with pytest.raises(OmniClientError, match="require an active Cache-DiT"):
+        _resolve_quality(
+            policy,
+            quality=quality,
+            extra_args={"force_refresh_step_hint": 7},
+        )
+
+
+@pytest.mark.parametrize(
+    ("cache_backend", "quality", "expected_key"),
+    [
+        ("cache_dit", None, "minimax_h3.generic"),
+        ("cache_dit", "lossless", None),
+        ("cache_dit", "high", "minimax_h3.high"),
+        ("none", None, None),
+        ("none", "lossless", None),
+        ("none", "high", "minimax_h3.high"),
+    ],
+)
+def test_model_policy_owns_request_cache_target(cache_backend, quality, expected_key):
+    from vllm_omni.diffusion.models.minimax_h3.quality_policy import (
+        MiniMaxH3QualityPolicy,
+    )
+
+    config = _quality_od_config(cache_backend=cache_backend)
+    policy = MiniMaxH3QualityPolicy(config)
+
+    plan = _resolve_quality(
+        policy,
+        quality=quality,
+    )
+
+    if expected_key is None:
+        assert plan.cache_dit is None
+    else:
+        assert plan.cache_dit is not None
+        assert plan.cache_dit.installation_key == expected_key
+        if quality is None:
+            assert plan.cache_dit.cache_config is config.cache_config
+
+
+def test_h3_adopts_runner_installed_cache_dit_backend():
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+
+    backend = object()
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    pipeline._cache_dit_runtime = SimpleNamespace(
+        adopt=lambda adopted, *, installation_key: setattr(
+            pipeline,
+            "adopted",
+            (adopted, installation_key),
+        ),
+        is_enabled=True,
+    )
+
+    pipeline.adopt_cache_dit_backend(backend)
+
+    assert pipeline.adopted == (backend, "minimax_h3.generic")
+    assert pipeline.is_cache_dit_enabled()

tests/diffusion/quantization/test_int8_config.py

@@ -147,6 +147,130 @@ def test_get_npu_quant_method(mocker: MockerFixture, monkeypatch: pytest.MonkeyP
     assert isinstance(method, UnquantizedLinearMethod)
 
 
+def test_fused_layer_can_be_ignored_by_its_prefix(monkeypatch: pytest.MonkeyPatch, mocker: MockerFixture):
+    """A layer opts out of quantization under the prefix its module was built with.
+
+    MiniMax H3 stores qkv as one checkpoint tensor, so the fused name is the
+    only handle a user has on it.
+    """
+    from vllm_omni.quantization.int8_config import NPUInt8OnlineLinearMethod
+
+    monkeypatch.setattr(current_omni_platform, "is_cuda", lambda: False)
+    monkeypatch.setattr(current_omni_platform, "is_npu", lambda: True)
+    config = build_quant_config("int8", ignored_layers=["blocks.0.attn.qkv_proj"])
+    layer = mocker.Mock(spec=LinearBase)
+
+    assert isinstance(config.get_quant_method(layer, "blocks.0.attn.qkv_proj"), UnquantizedLinearMethod)
+    assert isinstance(config.get_quant_method(layer, "blocks.1.attn.qkv_proj"), NPUInt8OnlineLinearMethod)
+    assert isinstance(config.get_quant_method(layer, "blocks.0.mlp.fc1"), NPUInt8OnlineLinearMethod)
+
+
+class TestNPUQuantMatmulShapeFallback:
+    """npu_quant_matmul rejects wide outputs, so those layers stay unquantized.
+
+    The limit applies to the per-rank shard, which is what the kernel sees. The
+    same layer can therefore be quantized at a higher TP degree.
+    """
+
+    @pytest.fixture(autouse=True)
+    def _mock_tp(self, mocker):
+        # The fallback delegates to UnquantizedLinearMethod, which reads the TP
+        # group; stand one in so the shape gating can be tested without a
+        # distributed init.
+        mock_group = mocker.Mock()
+        mock_group.rank_in_group = 0
+        mocker.patch("vllm.model_executor.layers.linear.get_tensor_model_parallel_world_size", return_value=1)
+        mocker.patch("vllm.model_executor.layers.linear.get_tensor_model_parallel_rank", return_value=0)
+        mocker.patch("vllm.distributed.parallel_state.get_tp_group", return_value=mock_group)
+
+    @pytest.fixture
+    def quant_config(self):
+        from vllm_omni.quantization.int8_config import DiffusionInt8Config
+
+        return DiffusionInt8Config(is_checkpoint_int8_serialized=False, activation_scheme="dynamic")
+
+    def create_weights(self, method, layer, output_partition_sizes):
+        # Narrow input dim: the fallback path allocates for real, and only the
+        # output dimension is under test.
+        method.create_weights(
+            layer,
+            input_size_per_partition=8,
+            output_partition_sizes=output_partition_sizes,
+            input_size=8,
+            output_size=sum(output_partition_sizes),
+            params_dtype=torch.bfloat16,
+            weight_loader=lambda param, loaded_weight, *args, **kwargs: None,
+        )
+
+    @pytest.mark.parametrize(
+        "method_name",
+        ["NPUInt8LinearMethod", "NPUInt8OnlineLinearMethod"],
+    )
+    def test_over_wide_layer_falls_back_to_unquantized(self, quant_config, method_name):
+        import vllm_omni.quantization.int8_config as int8_config
+
+        method = getattr(int8_config, method_name)(quant_config)
+        layer = Module()
+        layer.quant_method = method
+
+        self.create_weights(method, layer, [int8_config.NPU_QUANT_MATMUL_MAX_OUT_FEATURES + 1])
+
+        assert isinstance(layer.quant_method, UnquantizedLinearMethod)
+        assert layer.weight.dtype == torch.bfloat16
+
+    def test_layer_within_the_limit_keeps_int8(self, quant_config):
+        from vllm_omni.quantization.int8_config import (
+            NPU_QUANT_MATMUL_MAX_OUT_FEATURES,
+            NPUInt8OnlineLinearMethod,
+        )
+
+        method = NPUInt8OnlineLinearMethod(quant_config)
+        layer = Module()
+        layer.quant_method = method
+
+        self.create_weights(method, layer, [NPU_QUANT_MATMUL_MAX_OUT_FEATURES])
+
+        assert layer.quant_method is method
+
+    def test_tensor_parallel_sharding_brings_a_wide_layer_back(self, quant_config):
+        from vllm_omni.quantization.int8_config import (
+            NPU_QUANT_MATMUL_MAX_OUT_FEATURES,
+            NPUInt8OnlineLinearMethod,
+        )
+
+        global_out_features = 4 * NPU_QUANT_MATMUL_MAX_OUT_FEATURES
+        method = NPUInt8OnlineLinearMethod(quant_config)
+        layer = Module()
+        layer.quant_method = method
+
+        self.create_weights(method, layer, [global_out_features // 4])
+
+        assert layer.quant_method is method
+
+
+class TestOffloadAfterQuant:
+    def test_only_methods_advertising_the_capability_are_asked(self):
+        from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
+        from vllm_omni.quantization.int8_config import DiffusionInt8Config, NPUInt8OnlineLinearMethod
+
+        class MetaDeviceOnlyMethod:
+            """Stands in for upstream online FP8: lazy, but no offload hook."""
+
+            uses_meta_device = True
+
+        model = torch.nn.Sequential(Module(), Module())
+        lazy_int8 = NPUInt8OnlineLinearMethod(
+            DiffusionInt8Config(is_checkpoint_int8_serialized=False, activation_scheme="dynamic")
+        )
+        model[0].quant_method = lazy_int8
+        model[1].quant_method = MetaDeviceOnlyMethod()
+
+        marked = DiffusersPipelineLoader._request_offload_after_quant(model)
+
+        assert marked == 1
+        assert lazy_int8._offload_after_quant
+
+
 class TestInt8LinearMethod:
     @pytest.fixture
     def mock_quant_config(self, mocker):

tests/diffusion/test_async_output_worker.py

@@ -178,10 +178,9 @@ class TestAsyncOutputLoopLogic:
         t = threading.Thread(target=proc._async_output_loop, daemon=True)
         t.start()
 
-        # Give the loop time to process, then unblock and stop.
+        # Give the loop time to process, then send the shutdown sentinel.
         time.sleep(0.5)
-        proc._running = False
-        proc._async_output_queue.put((DiffusionOutput(), "sentinel", None))
+        proc._async_output_queue.put(None)
         t.join(timeout=2.0)
 
         mock_pack.assert_called()
@@ -191,3 +190,43 @@ class TestAsyncOutputLoopLogic:
             for c in proc.result_mq.enqueue.call_args_list
         )
         assert output_ready_enqueued, "Expected OUTPUT_READY with async_output_id='abc123' in enqueue calls"
+
+
+class TestWorkerProcShutdown:
+    def test_shutdown_stops_async_thread_and_releases_queues(self):
+        import threading
+
+        proc = _make_worker_proc(step_execution=False)
+        proc.worker = MagicMock()
+        proc.mq = MagicMock()
+        broadcast_mq = proc.mq
+        result_mq = proc.result_mq
+
+        thread = threading.Thread(target=proc._async_output_queue.get)
+        proc._async_output_thread = thread
+        thread.start()
+
+        proc.shutdown()
+
+        assert not thread.is_alive()
+        proc.worker.shutdown.assert_called_once_with()
+        broadcast_mq.shutdown.assert_called_once_with()
+        result_mq.shutdown.assert_called_once_with()
+        assert proc.mq is None
+        assert proc.result_mq is None
+
+    def test_shutdown_releases_queues_when_worker_shutdown_fails(self):
+        proc = _make_worker_proc(step_execution=True)
+        proc.worker = MagicMock()
+        proc.worker.shutdown.side_effect = RuntimeError("shutdown failed")
+        proc.mq = MagicMock()
+        broadcast_mq = proc.mq
+        result_mq = proc.result_mq
+
+        with pytest.raises(RuntimeError, match="shutdown failed"):
+            proc.shutdown()
+
+        broadcast_mq.shutdown.assert_called_once_with()
+        result_mq.shutdown.assert_called_once_with()
+        assert proc.mq is None
+        assert proc.result_mq is None

tests/diffusion/test_diffusion_config_propagation.py

@@ -15,7 +15,10 @@ from vllm_omni.diffusion.data import (
     DiffusionParallelConfig,
     OmniDiffusionConfig,
 )
-from vllm_omni.diffusion.model_metadata import QWEN_IMAGE_EDIT_PLUS_MAX_INPUT_IMAGES
+from vllm_omni.diffusion.model_metadata import (
+    HUNYUAN_IMAGE3_MAX_INPUT_IMAGES,
+    QWEN_IMAGE_EDIT_PLUS_MAX_INPUT_IMAGES,
+)
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -131,6 +134,35 @@ def test_qwen_image_edit_plus_sets_generic_multimodal_limit():
     assert od_config.max_multimodal_image_inputs == QWEN_IMAGE_EDIT_PLUS_MAX_INPUT_IMAGES
 
 
+def test_task_type_roundtrip():
+    od = _roundtrip_diffusion_config(model="x", task_type="model-defined-task")
+    assert od.task_type == "model-defined-task"
+
+
+def test_architecture_name_resolves_via_pipeline_class_fallback():
+    # Direct hit: the architecture name itself is the metadata table key.
+    qwen_od_config = OmniDiffusionConfig(
+        model="Qwen/Qwen-Image-Edit-2511",
+        model_class_name="QwenImageEditPlusPipeline",
+    )
+    qwen_od_config.update_multimodal_support()
+
+    assert qwen_od_config.supports_multimodal_inputs is True
+    assert qwen_od_config.max_multimodal_image_inputs == QWEN_IMAGE_EDIT_PLUS_MAX_INPUT_IMAGES
+
+    # Fallback: the HF architecture name differs from the internal pipeline class
+    # name, so ``get_diffusion_model_metadata`` must consult ``_DIFFUSION_MODELS``
+    # and resolve the metadata via the pipeline class name.
+    hunyuan_od_config = OmniDiffusionConfig(
+        model="tencent/HunyuanImage-3.0-Instruct",
+        model_class_name="HunyuanImage3ForCausalMM",
+    )
+    hunyuan_od_config.update_multimodal_support()
+
+    assert hunyuan_od_config.supports_multimodal_inputs is True
+    assert hunyuan_od_config.max_multimodal_image_inputs == HUNYUAN_IMAGE3_MAX_INPUT_IMAGES
+
+
 def test_additional_config_roundtrip():
     additional_config = {"torchair_graph_config": {"enabled": True}}
     od = _roundtrip_diffusion_config(model="x", additional_config=additional_config)

tests/diffusion/test_diffusion_engine.py

@@ -331,6 +331,55 @@ class TestRequestBatchCapability:
             DiffusionEngine(od_config)
         fake_executor_cls.assert_not_called()
 
+    def test_engine_allows_independent_dlo_dp_requests_for_single_request_pipeline(
+        self,
+        monkeypatch: pytest.MonkeyPatch,
+        mocker: MockerFixture,
+    ) -> None:
+        od_config = SimpleNamespace(
+            model_class_name="SinglePipeline",
+            custom_pipeline_args=None,
+            streaming_output=False,
+            max_num_seqs=2,
+            request_batch_max_wait_ms=0,
+            parallel_config=SimpleNamespace(data_parallel_size=2),
+            enable_distributed_layerwise_offload=True,
+            dlo_use_allgather=True,
+        )
+        fake_executor = SimpleNamespace(
+            execute_request=mocker.Mock(return_value="per-request"),
+            execute_batch=mocker.Mock(return_value="batch"),
+            execute_step=mocker.Mock(return_value="step"),
+        )
+        fake_executor_cls = mocker.Mock(return_value=fake_executor)
+
+        monkeypatch.setattr(
+            "vllm_omni.diffusion.diffusion_engine.get_diffusion_post_process_func",
+            lambda *args, **kwargs: None,
+        )
+        monkeypatch.setattr(
+            "vllm_omni.diffusion.diffusion_engine.get_diffusion_pre_process_func",
+            lambda *args, **kwargs: None,
+        )
+        monkeypatch.setattr(
+            "vllm_omni.diffusion.diffusion_engine.DiffusionExecutor.get_class",
+            lambda *args, **kwargs: fake_executor_cls,
+        )
+        monkeypatch.setattr(
+            diffusion_engine_module.DiffusionModelRegistry,
+            "_try_load_model_cls",
+            lambda model_class_name: _SingleRequestPipeline,
+        )
+
+        engine = DiffusionEngine(od_config)
+
+        assert engine.execution_mode == DiffusionExecutionMode.REQUEST_BATCH
+        assert engine.supports_request_batch is False
+        assert engine.dp_concurrent is True
+        assert engine.scheduler.max_num_running_reqs == 2
+        assert od_config.request_batch_max_wait_ms == 0
+        fake_executor_cls.assert_called_once_with(od_config)
+
     @pytest.mark.parametrize("request_ids", [("req-a",), ("req-a", "req-b")])
     def test_engine_enables_batch_dispatch_for_request_batch_pipeline(
         self,

tests/diffusion/test_diffusion_ipc.py

@@ -2,11 +2,19 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import contextlib
+import multiprocessing as mp
+import queue
+import threading
+import time
+from multiprocessing import shared_memory
 
+import numpy as np
 import pytest
 import torch
+from vllm.distributed.device_communicators.shm_broadcast import MessageQueue
 
 from vllm_omni.diffusion.data import DiffusionOutput
+from vllm_omni.diffusion.executor.multiproc_executor import MultiprocDiffusionExecutor
 from vllm_omni.diffusion.ipc import (
     _SHM_TENSOR_THRESHOLD,
     DIFFUSION_RPC_RESULT_ENVELOPE,
@@ -30,6 +38,129 @@ def _cleanup_shm_handle(value: object) -> None:
             _unpack_if_shm_handle(value)
 
 
+def _result_queue_worker(connection, rank: int) -> None:
+    """Send one nested NumPy result through a worker-owned MessageQueue."""
+    result_mq = MessageQueue(n_reader=1, n_local_reader=1, local_reader_ranks=[0])
+    try:
+        connection.send(result_mq.export_handle())
+        assert connection.recv() == "reader-ready"
+
+        frames = np.full(300_000, rank, dtype=np.float32)
+        audio = np.arange(300_000, dtype=np.float32) + rank
+        output = DiffusionOutput(
+            output={
+                "rank": rank,
+                "nested": {"frames": frames},
+                "audio": [audio],
+            }
+        )
+        pack_diffusion_output_shm(output)
+        connection.send(
+            [
+                output.output["nested"]["frames"]["name"],
+                output.output["audio"][0]["name"],
+            ]
+        )
+        result_mq.enqueue(output)
+        assert connection.recv() == "shutdown"
+    finally:
+        result_mq.shutdown()
+        connection.close()
+
+
+def test_per_worker_result_queues_release_nested_numpy_shm_and_processes() -> None:
+    """Exercise the production per-worker queue/pump/SHM lifecycle."""
+    ctx = mp.get_context("spawn")
+    parent_connections = []
+    processes = []
+    result_mqs = []
+    shm_names = []
+    executor = None
+
+    try:
+        for rank in range(2):
+            parent_connection, child_connection = ctx.Pipe()
+            process = ctx.Process(
+                target=_result_queue_worker,
+                args=(child_connection, rank),
+                name=f"IPCResultWorker-{rank}",
+            )
+            process.start()
+            child_connection.close()
+            parent_connections.append(parent_connection)
+            processes.append(process)
+
+        for connection in parent_connections:
+            handle = connection.recv()
+            result_mqs.append(MessageQueue.create_from_handle(handle, 0))
+            connection.send("reader-ready")
+            shm_names.extend(connection.recv())
+
+        executor = object.__new__(MultiprocDiffusionExecutor)
+        executor._result_mqs = result_mqs
+        executor._result_mq = result_mqs[0]
+        executor._pump_running = False
+        executor._pump_stop = threading.Event()
+        executor._sync_result_buffer = queue.Queue()
+        executor._is_failed = False
+        executor._futures_lock = threading.RLock()
+        executor._rpc_futures = {}
+        executor._output_futures = {}
+        executor._completed_outputs = {}
+        executor._batch_split_map = {}
+        executor._start_result_pump()
+
+        deadline = time.monotonic() + 10.0
+        while executor._sync_result_buffer.qsize() < 2 and time.monotonic() < deadline:
+            time.sleep(0.01)
+        assert executor._sync_result_buffer.qsize() == 2
+
+        outputs = []
+        for _ in range(2):
+            output = executor._sync_result_buffer.get_nowait()
+            unpack_diffusion_output_shm(output)
+            outputs.append(output)
+
+        assert sorted(output.output["rank"] for output in outputs) == [0, 1]
+        for output in outputs:
+            rank = output.output["rank"]
+            assert output.output["nested"]["frames"][0] == rank
+            assert output.output["audio"][0][0] == rank
+
+        executor._pump_stop.set()
+        for thread in executor._result_pump_threads:
+            thread.join(timeout=3.0)
+        assert all(not thread.is_alive() for thread in executor._result_pump_threads)
+    finally:
+        if executor is not None:
+            executor._pump_stop.set()
+            for thread in getattr(executor, "_result_pump_threads", []):
+                thread.join(timeout=3.0)
+
+        for connection in parent_connections:
+            with contextlib.suppress(BrokenPipeError, EOFError, OSError):
+                connection.send("shutdown")
+        for process in processes:
+            process.join(timeout=10.0)
+            if process.is_alive():
+                process.terminate()
+                process.join(timeout=5.0)
+        for result_mq in result_mqs:
+            with contextlib.suppress(Exception):
+                result_mq.shutdown()
+        for connection in parent_connections:
+            connection.close()
+
+    assert all(process.exitcode == 0 for process in processes)
+    for name in shm_names:
+        try:
+            leaked = shared_memory.SharedMemory(name=name)
+        except FileNotFoundError:
+            continue
+        leaked.close()
+        pytest.fail(f"shared memory segment {name} still exists after unpack")
+
+
 def test_diffusion_output_dict_tensors_round_trip_through_shm() -> None:
     image = torch.arange(300_000, dtype=torch.float32)
     video = torch.arange(300_000, dtype=torch.float32) * 2
@@ -83,6 +214,29 @@ def test_diffusion_output_list_tensors_round_trip_through_shm() -> None:
     torch.testing.assert_close(output.output[1], frames[1])
 
 
+def test_diffusion_output_numpy_array_round_trips_through_shm() -> None:
+    frames = np.arange(300_000, dtype=np.float32)
+    output = DiffusionOutput(output=frames)
+
+    pack_diffusion_output_shm(output)
+
+    assert output.output["__ndarray_shm__"] is True
+
+    unpack_diffusion_output_shm(output)
+
+    assert isinstance(output.output, np.ndarray)
+    np.testing.assert_array_equal(output.output, frames)
+
+
+def test_numpy_object_array_stays_inline() -> None:
+    values = np.empty(200_000, dtype=object)
+    values[:] = "safe-to-pickle"
+
+    packed = _pack_value_if_large(values)
+
+    assert packed is values
+
+
 def test_rpc_result_envelope_diffusion_output_round_trips_through_shm() -> None:
     tensor = torch.arange(300_000, dtype=torch.float32)
     envelope = {
@@ -108,6 +262,25 @@ def test_rpc_result_envelope_diffusion_output_round_trips_through_shm() -> None:
     assert unpacked["rank_statuses"] == [{"rank": 0, "ok": True}]
 
 
+def test_rpc_result_envelope_dp_tagged_output_round_trips_through_shm() -> None:
+    frames = np.arange(300_000, dtype=np.float32)
+    envelope = {
+        "type": DIFFUSION_RPC_RESULT_ENVELOPE,
+        "result": {"dp_rank": 1, "output": DiffusionOutput(output=frames)},
+        "rank_statuses": [{"rank": 1, "ok": True}],
+    }
+
+    packed = pack_diffusion_output_shm(envelope)
+    tagged = packed["result"]
+    assert tagged["dp_rank"] == 1
+    assert tagged["output"].output["__ndarray_shm__"] is True
+
+    unpacked = unpack_diffusion_output_shm(packed)
+    tagged = unpacked["result"]
+    assert tagged["dp_rank"] == 1
+    np.testing.assert_array_equal(tagged["output"].output, frames)
+
+
 def test_batch_runner_output_round_trips_nested_results_through_shm() -> None:
     first = torch.arange(_large_numel(torch.float32), dtype=torch.float32)
     second = torch.arange(_large_numel(torch.float32), dtype=torch.float32) + 1

tests/diffusion/test_diffusion_model_runner.py

@@ -162,6 +162,21 @@ def _fake_platform_for_peak_memory():
     )
 
 
+def test_request_scoped_cache_dit_lifecycle_is_pipeline_opt_in():
+    events = []
+    plain_pipeline = object()
+    request_scoped_pipeline = SimpleNamespace(
+        adopt_cache_dit_backend=lambda backend: events.append(backend),
+        is_cache_dit_enabled=lambda: True,
+    )
+    backend = object()
+
+    assert not model_runner_module.adopt_request_scoped_cache_dit(plain_pipeline, backend)
+    assert model_runner_module.adopt_request_scoped_cache_dit(request_scoped_pipeline, backend)
+    assert model_runner_module.is_request_scoped_cache_dit_enabled(request_scoped_pipeline)
+    assert events == [backend]
+
+
 def _make_runner(cache_backend, cache_backend_name: str, enable_cache_dit_summary: bool = True):
     runner = object.__new__(DiffusionModelRunner)
     runner.vllm_config = object()
@@ -419,6 +434,33 @@ def test_execute_model_emits_cache_summary_with_active_cache_dit_backend(monkeyp
     assert cache_backend.refresh_calls == [(runner.pipeline, 4, True)]
 
 
+@pytest.mark.core_model
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("cache_dit_enabled", [False, True])
+def test_execute_model_cache_summary_follows_pipeline_owned_cache_dit_state(monkeypatch, cache_dit_enabled):
+    runner = _make_runner(cache_backend=None, cache_backend_name="cache_dit")
+    runner.pipeline.adopt_cache_dit_backend = lambda backend: None
+    runner.pipeline.is_cache_dit_enabled = lambda: cache_dit_enabled
+    req = _make_request()
+
+    cache_summary_calls = []
+
+    monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
+    monkeypatch.setattr(model_runner_module.current_omni_platform, "reset_peak_memory_stats", lambda: None)
+    monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_reserved", lambda: 0)
+    monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_allocated", lambda: 0)
+    monkeypatch.setattr(
+        model_runner_module,
+        "cache_summary",
+        lambda pipeline, details: cache_summary_calls.append((pipeline, details)),
+    )
+
+    output = DiffusionModelRunner.execute_model(runner, req)
+
+    assert output.output == "ok"
+    assert cache_summary_calls == ([(runner.pipeline, True)] if cache_dit_enabled else [])
+
+
 @pytest.mark.core_model
 @pytest.mark.cpu
 def test_execute_model_passes_single_request_batch_to_non_admission_pipeline(monkeypatch):

tests/diffusion/test_diffusion_scheduler.py

@@ -380,6 +380,63 @@ class TestRequestScheduler:
         assert first.num_running_reqs == 1
         assert first.num_waiting_reqs == 1
 
+    def test_batches_different_quality_levels_separately(self) -> None:
+        scheduler = RequestScheduler()
+        scheduler.initialize(SimpleNamespace(max_num_seqs=2))
+
+        high = scheduler.add_request(
+            _make_step_request(
+                "high",
+                sampling_params=OmniDiffusionSamplingParams(
+                    num_inference_steps=2,
+                    quality="high",
+                ),
+            )
+        )
+        scheduler.add_request(
+            _make_step_request(
+                "lossless",
+                sampling_params=OmniDiffusionSamplingParams(
+                    num_inference_steps=2,
+                    quality="lossless",
+                ),
+            )
+        )
+
+        first = scheduler.schedule()
+
+        assert _new_ids(first) == [high]
+        assert first.num_running_reqs == 1
+        assert first.num_waiting_reqs == 1
+
+    def test_batches_omitted_and_explicit_lossless_separately(self) -> None:
+        scheduler = RequestScheduler()
+        scheduler.initialize(SimpleNamespace(max_num_seqs=2))
+
+        omitted = scheduler.add_request(
+            _make_step_request(
+                "omitted",
+                sampling_params=OmniDiffusionSamplingParams(
+                    num_inference_steps=2,
+                ),
+            )
+        )
+        scheduler.add_request(
+            _make_step_request(
+                "explicit",
+                sampling_params=OmniDiffusionSamplingParams(
+                    num_inference_steps=2,
+                    quality="lossless",
+                ),
+            )
+        )
+
+        first = scheduler.schedule()
+
+        assert _new_ids(first) == [omitted]
+        assert first.num_running_reqs == 1
+        assert first.num_waiting_reqs == 1
+
     def test_batches_different_request_local_seed_together(self) -> None:
         scheduler = RequestScheduler()
         scheduler.initialize(SimpleNamespace(max_num_seqs=2))
@@ -930,6 +987,35 @@ class TestStepScheduler:
         assert second.num_running_reqs == 1
         assert second.num_waiting_reqs == 1
 
+    def test_step_batch_rejects_different_quality_levels(self) -> None:
+        scheduler = StepScheduler()
+        scheduler.initialize(SimpleNamespace(max_num_seqs=2))
+
+        high = scheduler.add_request(
+            _make_step_request(
+                "high",
+                sampling_params=OmniDiffusionSamplingParams(
+                    num_inference_steps=4,
+                    quality="high",
+                ),
+            )
+        )
+        scheduler.add_request(
+            _make_step_request(
+                "lossless",
+                sampling_params=OmniDiffusionSamplingParams(
+                    num_inference_steps=4,
+                    quality="lossless",
+                ),
+            )
+        )
+
+        first = scheduler.schedule()
+
+        assert _new_ids(first) == [high]
+        assert first.num_running_reqs == 1
+        assert first.num_waiting_reqs == 1
+
     def test_step_batch_co_schedules_requests_sharing_lora(self) -> None:
         """Multiple requests with the same LoRA (id + scale) co-batch."""
         from vllm_omni.lora.request import LoRARequest

tests/diffusion/test_inline_stage_diffusion_client.py

@@ -75,6 +75,32 @@ async def test_inline_dispatch_request_success(client, mock_engine):
     assert output.request_id == "req-1"
 
 
+@pytest.mark.asyncio
+async def test_inline_requests_clone_shared_sampling_params(client, mock_engine):
+    requests = []
+
+    async def _step_streaming(request):
+        requests.append(request)
+        yield [OmniRequestOutput.from_diffusion(request_id=request.request_id, images=[MagicMock()])]
+
+    mock_engine.step_streaming = _step_streaming
+
+    shared = OmniDiffusionSamplingParams()
+    await client.add_request_async("req-1", "First prompt", shared)
+    await client.add_request_async("req-2", "Second prompt", shared)
+
+    for _ in range(10):
+        if len(requests) == 2:
+            break
+        await asyncio.sleep(0.01)
+
+    assert len(requests) == 2
+    assert requests[0].sampling_params is not requests[1].sampling_params
+    assert requests[0].sampling_params.guidance_scale_provided is False
+    assert requests[1].sampling_params.guidance_scale_provided is False
+    assert shared.guidance_scale is None
+
+
 @pytest.mark.asyncio
 async def test_inline_dispatch_request_streaming_success(client, mock_engine):
     """Inline StageDiffusionClient correctly receives streaming chunk output from Diffusion Engine"""
@@ -137,6 +163,20 @@ def test_inline_shutdown(client, mock_engine):
     mock_engine.close.assert_called_once()
 
 
+def test_inline_shutdown_runs_after_orchestrator_premark(client, mock_engine):
+    # The orchestrator pre-marks clients to suppress false worker-death errors
+    # before StagePool invokes the actual teardown.
+    client._shutting_down = True
+
+    client.shutdown()
+
+    mock_engine.close.assert_called_once()
+    assert client._shutdown_complete
+
+    client.shutdown()
+    mock_engine.close.assert_called_once()
+
+
 def test_inline_registers_executor_failure_callback(client, mock_engine):
     mock_engine.executor.register_failure_callback.assert_called_once()
 

tests/diffusion/test_ipc_async.py

@@ -137,6 +137,18 @@ class TestPackDiffusionOutputShmWithD2hStream:
             d2h_stream=d2h,
         )
 
+    @pytest.mark.parametrize("wrapped_in_rpc_envelope", [False, True])
+    def test_dp_tagged_output_path_with_stream(self, mocker, wrapped_in_rpc_envelope):
+        mock_pack_fields = mocker.patch(
+            "vllm_omni.diffusion.ipc._pack_diffusion_fields",
+        )
+        d2h = mocker.MagicMock()
+        result = DiffusionOutput()
+        tagged = {"dp_rank": 1, "output": result}
+        value = {"type": "diffusion_rpc_result", "result": tagged} if wrapped_in_rpc_envelope else tagged
+        pack_diffusion_output_shm(value, d2h_stream=d2h)
+        mock_pack_fields.assert_called_once_with(result, d2h_stream=d2h)
+
     def test_runner_output_path_with_stream(self, mocker):
         mock_pack_fields = mocker.patch(
             "vllm_omni.diffusion.ipc._pack_diffusion_fields",

tests/diffusion/test_multiproc_engine_concurrency.py

@@ -276,6 +276,8 @@ class TestRequestModeDispatch:
         def worker():
             request = req_q.get(timeout=2)
             assert "rpc_id" not in request
+            assert request["output_rank"] is None
+            assert request["collect_rank_status"] is False
             wave_id = request["wave_id"]
             for dp_rank in range(2):
                 res_q.put(
@@ -297,6 +299,181 @@ class TestRequestModeDispatch:
 
         assert [result.error for result in results] == ["0", "1"]
 
+    @pytest.mark.parametrize(
+        ("tp_size", "sp_size", "pp_size", "cfg_size", "expected_primary_ranks"),
+        [
+            (2, 1, 1, 1, [0, 2]),
+            (2, 2, 1, 1, [0, 4]),
+            (1, 2, 2, 1, [0, 4]),
+            (1, 1, 2, 2, [0, 4]),
+        ],
+    )
+    def test_step_execution_reads_each_dp_primary_result_queue(
+        self,
+        tp_size,
+        sp_size,
+        pp_size,
+        cfg_size,
+        expected_primary_ranks,
+    ):
+        dp_size = 2
+        num_gpus = dp_size * tp_size * sp_size * pp_size * cfg_size
+        executor, _, _ = _make_executor(num_gpus=num_gpus)
+        executor.od_config = SimpleNamespace(
+            step_execution=True,
+            parallel_config=SimpleNamespace(
+                data_parallel_size=dp_size,
+                tensor_parallel_size=tp_size,
+                sequence_parallel_size=sp_size,
+                pipeline_parallel_size=pp_size,
+                cfg_parallel_size=cfg_size,
+            ),
+        )
+        request: dict = {}
+        executor._broadcast_mq = SimpleNamespace(enqueue=lambda value: request.update(value))
+
+        def result_queue(global_rank):
+            result_mq = MagicMock()
+            result_mq.dequeue.side_effect = lambda timeout=None: {
+                "dp_rank": global_rank,
+                "output": _tagged_output(str(global_rank)),
+                "wave_id": request["wave_id"],
+            }
+            return result_mq
+
+        # Every queue returns a valid, rank-tagged result. If the executor
+        # selects a non-primary queue, the returned rank exposes the mistake.
+        executor._result_mqs = [result_queue(rank) for rank in range(num_gpus)]
+        executor._result_mq = executor._result_mqs[0]
+
+        results = executor.collective_rpc(
+            "execute_model",
+            unique_reply_rank=None,
+            exec_all_ranks=True,
+        )
+
+        assert [result.error for result in results] == [str(rank) for rank in expected_primary_ranks]
+        for rank, result_mq in enumerate(executor._result_mqs):
+            assert result_mq.dequeue.call_count == (1 if rank in expected_primary_ranks else 0)
+
+    @pytest.mark.parametrize("empty_prompt", ["", {"prompt": ""}])
+    def test_dlo_dp_rejects_empty_prompt_before_worker_dispatch(self, empty_prompt):
+        executor, _, _ = _make_executor(num_gpus=2)
+        executor.od_config = SimpleNamespace(
+            step_execution=False,
+            parallel_config=SimpleNamespace(data_parallel_size=2),
+            enable_distributed_layerwise_offload=True,
+            dlo_use_allgather=True,
+        )
+        executor.collective_rpc = Mock()
+        scheduler_output = _make_sched_output("invalid", "valid")
+        scheduler_output.scheduled_new_reqs[0].req.prompt = empty_prompt
+
+        with pytest.raises(ValueError, match="non-empty prompt"):
+            executor.execute_request(scheduler_output)
+
+        executor.collective_rpc.assert_not_called()
+
+    def test_dlo_dp_valid_wave_still_runs_after_rejected_wave(self):
+        executor, _, _ = _make_executor(num_gpus=2)
+        executor.od_config = SimpleNamespace(
+            step_execution=False,
+            parallel_config=SimpleNamespace(data_parallel_size=2),
+            enable_distributed_layerwise_offload=True,
+            dlo_use_allgather=True,
+        )
+        executor.collective_rpc = Mock(return_value=[_tagged_output("A"), _tagged_output("B")])
+        invalid_wave = _make_sched_output("invalid", "valid")
+        invalid_wave.scheduled_new_reqs[0].req.prompt = ""
+
+        with pytest.raises(ValueError, match="non-empty prompt"):
+            executor.execute_request(invalid_wave)
+
+        result = executor.execute_request(_make_sched_output("A", "B"))
+
+        assert [output.result.error for output in result.runner_outputs] == ["A", "B"]
+        executor.collective_rpc.assert_called_once()
+
+    def test_dlo_dp_allows_shared_default_denoise_steps(self):
+        executor, _, _ = _make_executor(num_gpus=2)
+        executor.od_config = SimpleNamespace(
+            step_execution=False,
+            parallel_config=SimpleNamespace(data_parallel_size=2),
+            enable_distributed_layerwise_offload=True,
+            dlo_use_allgather=True,
+        )
+        executor.collective_rpc = Mock(return_value=[_tagged_output("A"), _tagged_output("B")])
+        scheduler_output = _make_sched_output("A", "B")
+        for new_req in scheduler_output.scheduled_new_reqs:
+            new_req.req.sampling_params.num_inference_steps = None
+
+        result = executor.execute_request(scheduler_output)
+
+        assert [output.result.error for output in result.runner_outputs] == ["A", "B"]
+        executor.collective_rpc.assert_called_once()
+
+    @pytest.mark.parametrize(
+        ("field", "value"),
+        [
+            ("num_inference_steps", 2),
+            ("guidance_scale", 7.5),
+            ("width", 1024),
+        ],
+    )
+    def test_dlo_dp_rejects_incompatible_collective_wave(self, field, value):
+        executor, _, _ = _make_executor(num_gpus=2)
+        executor.od_config = SimpleNamespace(
+            step_execution=False,
+            parallel_config=SimpleNamespace(data_parallel_size=2),
+            enable_distributed_layerwise_offload=True,
+            dlo_use_allgather=True,
+        )
+        executor.collective_rpc = Mock()
+        scheduler_output = _make_sched_output("A", "B")
+        setattr(scheduler_output.scheduled_new_reqs[1].req.sampling_params, field, value)
+
+        with pytest.raises(ValueError, match="compatible shape, CFG"):
+            executor.execute_request(scheduler_output)
+
+        executor.collective_rpc.assert_not_called()
+
+    def test_dlo_dp_partial_reply_times_out_and_fails_closed(self, monkeypatch):
+        from vllm_omni.diffusion.executor import multiproc_executor as executor_module
+
+        executor, req_q, res_q = _make_executor(num_gpus=2)
+        executor.od_config = SimpleNamespace(
+            step_execution=False,
+            parallel_config=SimpleNamespace(data_parallel_size=2),
+            enable_distributed_layerwise_offload=True,
+            dlo_use_allgather=True,
+        )
+        executor._sync_result_buffer = res_q
+        executor._fail_closed_on_dp_wave_timeout = Mock()
+        monkeypatch.setattr(executor_module, "_DLO_DP_WAVE_TIMEOUT_S", 0.05)
+
+        def reply_from_only_one_dp_rank():
+            request = req_q.get(timeout=2.0)
+            res_q.put(
+                {
+                    "dp_rank": 0,
+                    "output": _tagged_output("rank-0"),
+                    "wave_id": request["wave_id"],
+                }
+            )
+
+        worker = threading.Thread(target=reply_from_only_one_dp_rank, daemon=True)
+        worker.start()
+        started = time.monotonic()
+
+        result = executor.execute_request(_make_sched_output("A", "B"))
+
+        elapsed = time.monotonic() - started
+        worker.join(timeout=2.0)
+        assert elapsed < 1.0
+        assert all("timed out" in output.result.error for output in result.runner_outputs)
+        executor._fail_closed_on_dp_wave_timeout.assert_called_once()
+        assert isinstance(executor._fail_closed_on_dp_wave_timeout.call_args.args[0], TimeoutError)
+
 
 # ───────────────── concurrent collective RPC ─────────────────
 
@@ -948,6 +1125,42 @@ class TestStageDiffusionClientErrorPropagation:
         assert rpc_id not in client._pending_rpcs
 
 
+class TestExecutorShutdownCleaner:
+    def test_worker_joins_share_one_global_deadline(self, monkeypatch):
+        from vllm_omni.diffusion.executor import multiproc_executor as executor_module
+
+        class FakeProcess:
+            def __init__(self, name):
+                self.name = name
+                self.alive = True
+                self.terminated = False
+                self.join_timeouts = []
+
+            def is_alive(self):
+                return self.alive
+
+            def join(self, timeout):
+                self.join_timeouts.append(timeout)
+                if self.terminated:
+                    self.alive = False
+
+            def terminate(self):
+                self.terminated = True
+
+        monotonic = Mock(side_effect=[100.0, 100.0, 110.0, 120.0, 120.0, 124.0])
+        monkeypatch.setattr(executor_module, "time", SimpleNamespace(monotonic=monotonic))
+        first = FakeProcess("worker-0")
+        second = FakeProcess("worker-1")
+        cleaner = executor_module._ExecutorShutdownCleaner(processes=[first, second])
+
+        cleaner()
+
+        assert first.join_timeouts == [15.0, 5.0]
+        assert second.join_timeouts == [5.0, 1.0]
+        assert first.terminated and second.terminated
+        assert not first.is_alive() and not second.is_alive()
+
+
 # ───────── monitor thread & death sentinel integration tests ─────────
 
 

tests/diffusion/test_result_pump.py

@@ -35,6 +35,7 @@ def _make_executor(step_execution=False):
     executor._pump_stop = threading.Event()
     executor._sync_result_buffer = queue.Queue()
     executor._result_mq = MagicMock()
+    executor._result_mqs = []
     executor._broadcast_mq = MagicMock()
     executor._closed = False
     executor._is_failed = False
@@ -157,6 +158,41 @@ class TestResultPumpDispatch:
         retrieved = executor._sync_result_buffer.get_nowait()
         assert isinstance(retrieved, DiffusionOutput)
 
+    def test_start_result_pump_reads_every_worker_queue(self):
+        executor = _make_executor()
+        messages = [DiffusionOutput(output="rank0"), DiffusionOutput(output="rank1")]
+        result_mqs = [MagicMock(), MagicMock()]
+
+        def make_dequeue(message):
+            emitted = False
+
+            def dequeue(timeout=None):
+                nonlocal emitted
+                if not emitted:
+                    emitted = True
+                    return message
+                executor._pump_stop.wait(0.01)
+                raise TimeoutError
+
+            return dequeue
+
+        for result_mq, message in zip(result_mqs, messages, strict=True):
+            result_mq.dequeue = make_dequeue(message)
+
+        executor._result_mqs = result_mqs
+        executor._result_mq = result_mqs[0]
+        executor._start_result_pump()
+        deadline = time.monotonic() + 2.0
+        while executor._sync_result_buffer.qsize() < 2 and time.monotonic() < deadline:
+            time.sleep(0.01)
+        executor._pump_stop.set()
+        for thread in executor._result_pump_threads:
+            thread.join(timeout=2.0)
+
+        assert len(executor._result_pump_threads) == 2
+        received = [executor._sync_result_buffer.get_nowait().output for _ in range(2)]
+        assert sorted(received) == ["rank0", "rank1"]
+
     def test_compute_done_routes_to_rpc_future(self):
         executor = _make_executor()
         rpc_id = "42"
@@ -233,6 +269,20 @@ class TestResultPumpDispatch:
 class TestShutdownCleansUpFutures:
     """Test that shutdown cancels pending futures."""
 
+    def test_shutdown_joins_result_pump_threads(self):
+        executor = _make_executor()
+        pump = threading.Thread(
+            target=executor._pump_stop.wait,
+            name="test-result-pump",
+        )
+        executor._result_pump_threads = [pump]
+        pump.start()
+
+        executor.shutdown()
+
+        assert not pump.is_alive()
+        assert executor._result_pump_threads == []
+
     def test_shutdown_sets_exception_on_pending_futures(self):
         executor = _make_executor()
 

tests/e2e/accuracy/test_qwen_image.py

@@ -29,7 +29,7 @@ HEIGHT = 512
 NUM_INFERENCE_STEPS = 20
 TRUE_CFG_SCALE = 4.0
 SEED = 42
-SSIM_THRESHOLD = 0.94
+SSIM_THRESHOLD = 0.97
 PSNR_THRESHOLD = 30.0
 
 MODEL_2512_ID = "Qwen/Qwen-Image-2512"
@@ -69,7 +69,15 @@ def _local_files_only(model: str) -> bool:
 
 
 def _run_vllm_omni_qwen_image(*, model: str, output_path: Path) -> Image.Image:
-    server_args = ["--num-gpus", "1", "--stage-init-timeout", "300", "--init-timeout", "900"]
+    server_args = [
+        "--num-gpus",
+        "1",
+        "--stage-init-timeout",
+        "300",
+        "--init-timeout",
+        "900",
+        "--fa-deterministic",
+    ]
     with OmniServer(model, server_args, use_omni=True) as omni_server:
         response = requests.post(
             f"http://{omni_server.host}:{omni_server.port}/v1/images/generations",

tests/e2e/online_serving/test_flux_2_dev_expansion.py

@@ -136,3 +136,36 @@ def test_flux_2_dev(
     }
 
     openai_client.send_diffusion_request(request_config)
+
+
+@pytest.mark.parametrize(
+    "omni_server",
+    _get_flux_2_dev_feature_cases(MODEL),
+    indirect=True,
+)
+def test_flux_2_dev_batched_chat_completions(
+    omni_server: OmniServer,
+    openai_client: OpenAIClientHandler,
+):
+    """Validate batched chat completions for diffusion models."""
+    messages = [
+        [{"role": "user", "content": PROMPT}],
+        [{"role": "user", "content": "A sunset over the ocean."}],
+    ]
+    responses = openai_client.send_batched_chat_completions_http_request(
+        {
+            "json": {
+                "model": omni_server.model,
+                "messages": messages,
+                "num_inference_steps": 2,
+                "height": 512,
+                "width": 512,
+                "seed": 42,
+            },
+        },
+    )
+    assert responses and len(responses) == 1
+    resp = responses[0]
+    assert resp.success
+    choices = resp.json_body["choices"]
+    assert len(choices) == len(messages)

tests/e2e/online_serving/test_qwen3_omni.py

@@ -250,3 +250,101 @@ def test_completions_rejected_for_thinker_talker(omni_server, openai_client) ->
         err_code=400,
     )
     assert not responses[0].success
+
+
+@pytest.mark.advanced_model
+@pytest.mark.core_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+def test_batched_completions_text(omni_server, openai_client) -> None:
+    """Ensure that we can make a batch chat completions request (text only)."""
+    responses = openai_client.send_batched_chat_completions_http_request(
+        {
+            "json": {
+                "model": omni_server.model,
+                "messages": [
+                    [{"role": "user", "content": "What color is the sky? Answer in one word."}],
+                    [{"role": "user", "content": "What is 2+2? Answer in one word."}],
+                ],
+                "modalities": ["text"],
+                "max_tokens": 20,
+            },
+        },
+    )
+    assert responses
+    resp = responses[0]
+    assert resp.success
+    body = resp.json_body
+    choices = body["choices"]
+    assert len(choices) == 2
+    for choice in choices:
+        assert choice["message"]["content"]
+
+
+@pytest.mark.advanced_model
+@pytest.mark.core_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+def test_batched_completions_audio_out(omni_server, openai_client) -> None:
+    """Ensure that we can make a batch chat completions request (audio + text)."""
+    messages = [
+        [{"role": "user", "content": "Say hello."}],
+        [{"role": "user", "content": "Say goodbye."}],
+    ]
+    num_messages = len(messages)
+    responses = openai_client.send_batched_chat_completions_http_request(
+        {
+            "json": {
+                "model": omni_server.model,
+                "messages": messages,
+                "modalities": ["text", "audio"],
+                "max_tokens": 200,
+            },
+        },
+    )
+    assert responses and len(responses) == 1
+    resp = responses[0]
+    assert resp.success
+    choices = resp.json_body["choices"]
+
+    assert len(choices) == num_messages
+
+    # Every message in the batch should contain both the audio and text content
+    for choice in choices:
+        assert choice["message"]["audio"]["data"]
+        assert choice["message"]["content"]
+
+
+@pytest.mark.advanced_model
+@pytest.mark.core_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+@pytest.mark.parametrize(
+    "sad_opts",
+    [
+        {"stream": "not-a-bool"},  # Fails validation
+        {"n": 2},  # vLLM Batched endpoints doesn't support n > 1
+    ],
+)
+def test_batched_completions_with_bad_values(omni_server, openai_client, sad_opts) -> None:
+    """Ensure that a bad values are correctly handled as 400s."""
+    messages = [
+        [{"role": "user", "content": "Say hello."}],
+        [{"role": "user", "content": "Say goodbye."}],
+    ]
+    responses = openai_client.send_batched_chat_completions_http_request(
+        {
+            "json": {
+                "model": omni_server.model,
+                "messages": messages,
+                **sad_opts,
+            },
+        },
+    )
+    assert responses and len(responses) == 1
+    resp = responses[0]
+    assert not responses[0].success
+    assert resp.status_code == 400

tests/engine/test_arg_utils.py

@@ -128,6 +128,24 @@ def test_qwen3_tts_codec_frame_rate_patching():
     assert omni_config.codec_frame_rate_hz == 12.3
 
 
+def test_qwen3_tts_startup_task_type_is_validated():
+    vllm_config = EngineArgs().create_model_config()
+
+    config = OmniModelConfig.from_vllm_model_config(
+        vllm_config,
+        model_arch="Qwen3TTSTalkerForConditionalGenerationARVLLM",
+        task_type="Base",
+    )
+    assert config.task_type == "Base"
+
+    with pytest.raises(ValueError, match="Qwen3-TTS --task-type must be one of"):
+        OmniModelConfig.from_vllm_model_config(
+            vllm_config,
+            model_arch="Qwen3TTSTalkerForConditionalGenerationARVLLM",
+            task_type="fl2va",
+        )
+
+
 def test_from_cli_args_picks_up_stage_configs_path():
     """from_cli_args should pick up stage_configs_path from namespace."""
     ns = argparse.Namespace(

tests/engine/test_stage_engine_args.py

@@ -0,0 +1,510 @@
+from __future__ import annotations
+
+import copy
+import sys
+import types
+from dataclasses import fields, replace
+from pathlib import Path
+
+import pytest
+from pydantic.fields import FieldInfo
+
+from vllm_omni.config.omni_config import VllmOmniConfig
+from vllm_omni.config.pipeline_registry import OMNI_PIPELINES, resolve_pipeline_config
+from vllm_omni.config.stage_config import (
+    DeployConfig,
+    PipelineConfig,
+    StageDeployConfig,
+    StageExecutionType,
+    StagePipelineConfig,
+    StageType,
+    build_stage_runtime_overrides,
+    load_deploy_config,
+    merge_pipeline_deploy,
+)
+from vllm_omni.diffusion.data import AttentionConfig, OmniDiffusionConfig
+from vllm_omni.engine import stage_init_utils
+from vllm_omni.engine.arg_utils import OmniEngineArgs
+from vllm_omni.engine.stage_init_utils import (
+    build_engine_args_dict,
+    build_engine_args_dict_from_omni_stage_config,
+    build_legacy_engine_args_dict,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_DEPLOY_DIR = Path(__file__).parents[2] / "vllm_omni" / "deploy"
+
+
+def _effective_backend_values(config_cls: type, engine_args: dict) -> dict[str, object]:
+    backend_fields = fields(config_cls)
+    backend_field_names = {backend_field.name for backend_field in backend_fields}
+    backend_config = config_cls(
+        **{name: copy.deepcopy(value) for name, value in engine_args.items() if name in backend_field_names}
+    )
+    effective_values: dict[str, object] = {}
+    for backend_field in backend_fields:
+        value = getattr(backend_config, backend_field.name)
+        if isinstance(value, FieldInfo):
+            value = value.get_default(call_default_factory=True)
+        effective_values[backend_field.name] = value
+    return effective_values
+
+
+_LLM_BACKEND_FIELDS = frozenset(field.name for field in fields(OmniEngineArgs))
+_DIFFUSION_BACKEND_FIELDS = frozenset(field.name for field in fields(OmniDiffusionConfig))
+
+
+@pytest.fixture(autouse=True)
+def _stable_engine_arg_environment(monkeypatch, tmp_path):
+    from vllm_omni import platforms
+
+    monkeypatch.delenv("VLLM_USE_FLASHINFER_MOE_FP16", raising=False)
+    xcodec_model = tmp_path / "xcodec-model"
+    xcodec_model.mkdir()
+    monkeypatch.setenv("XCODEC1_PATH", str(xcodec_model))
+    platform = platforms.current_omni_platform
+    monkeypatch.setattr(platform, "device_name", "cpu", raising=False)
+    monkeypatch.setattr(platform, "device_type", "cpu", raising=False)
+    monkeypatch.setattr(platform, "is_rocm", lambda: False)
+    monkeypatch.setattr(
+        OmniDiffusionConfig,
+        "_resolve_master_port",
+        lambda config: config.master_port or 29500,
+    )
+
+    def _resolve_test_worker(engine_args):
+        worker_type = engine_args.get("worker_type")
+        if worker_type is not None and engine_args.get("worker_cls") in (None, "auto"):
+            engine_args["worker_cls"] = f"test.{worker_type}.Worker"
+
+    monkeypatch.setattr(stage_init_utils, "resolve_worker_cls", _resolve_test_worker)
+
+
+def _engine_arg_inputs(tmp_path: Path) -> tuple[PipelineConfig, DeployConfig, str]:
+    parent_model = tmp_path / "parent-model"
+    stage_model = tmp_path / "stage-model"
+    diffusion_model = tmp_path / "diffusion-model"
+    for root in (parent_model, stage_model, diffusion_model):
+        root.mkdir()
+    (stage_model / "ar-model").mkdir()
+    (stage_model / "ar-tokenizer").mkdir()
+
+    pipeline = PipelineConfig(
+        model_type="typed-engine-args",
+        model_arch="PipelineModel",
+        stages=(
+            StagePipelineConfig(
+                stage_id=0,
+                model_stage="thinker",
+                execution_type=StageExecutionType.LLM_AR,
+                requires_multimodal_data=True,
+                hf_config_name="thinker_config",
+                engine_output_type="hidden_states",
+                model_arch="Qwen3OmniMoeForConditionalGeneration",
+                model_subdir="ar-model",
+                tokenizer_subdir="ar-tokenizer",
+                retains_state_across_chunks=True,
+                async_chunk_process_next_stage_input_func="operator.add",
+            ),
+            StagePipelineConfig(
+                stage_id=1,
+                model_stage="talker",
+                execution_type=StageExecutionType.LLM_GENERATION,
+                input_sources=(0,),
+                engine_output_type="audio_tokens",
+            ),
+            StagePipelineConfig(
+                stage_id=2,
+                model_stage="dit",
+                execution_type=StageExecutionType.DIFFUSION,
+                input_sources=(1,),
+                final_output=True,
+                final_output_type="image",
+                model_arch="TypedDiffusionPipeline",
+            ),
+        ),
+    )
+    deploy = DeployConfig(
+        async_chunk=True,
+        trust_remote_code=True,
+        dtype="float16",
+        distributed_executor_backend="mp",
+        stages=[
+            StageDeployConfig(
+                stage_id=0,
+                tensor_parallel_size=2,
+                max_num_seqs=8,
+                default_sampling_params={"extra_args": {"voice": "test"}},
+                engine_extras={
+                    "model": str(stage_model),
+                    "attention_backend": "FLASHINFER",
+                    "hf_overrides": {"rope_scaling": {"factor": 2}},
+                    "limit_mm_per_prompt": {"audio": 1},
+                    "logits_processors": ["test.custom.LogitsProcessor"],
+                    "kv_cache_memory_bytes": 1024,
+                    "omni_kv_config": {"need_send_cache": True},
+                    "max_cudagraph_capture_size": 0,
+                    "worker_cls": "test.custom.Worker",
+                },
+            ),
+            StageDeployConfig(
+                stage_id=1,
+                max_num_seqs=4,
+            ),
+            StageDeployConfig(
+                stage_id=2,
+                tensor_parallel_size=2,
+                vae_parallel_mode="spatial_shard_height",
+                diffusion_attention_config={
+                    "default": {"backend": "FLASH_ATTN"},
+                    "per_role": {"cross": {"backend": "TORCH_SDPA"}},
+                },
+                engine_extras={
+                    "model": str(diffusion_model),
+                    "engine_backend": "test.diffusion.Engine",
+                    "enable_session_state_manager": True,
+                },
+            ),
+        ],
+    )
+    return pipeline, deploy, str(parent_model)
+
+
+def _legacy_and_typed_stages(
+    pipeline: PipelineConfig,
+    deploy: DeployConfig,
+    model: str,
+    cli_overrides: dict[str, object] | None = None,
+):
+    resolved_cli_overrides = {"model": model, **(cli_overrides or {})}
+    legacy_stage_configs = merge_pipeline_deploy(
+        pipeline,
+        copy.deepcopy(deploy),
+    )
+    for stage in legacy_stage_configs:
+        stage.runtime_overrides = build_stage_runtime_overrides(
+            stage.stage_id,
+            resolved_cli_overrides,
+        )
+    legacy_stages = [stage.to_omegaconf() for stage in legacy_stage_configs]
+    omni_config = VllmOmniConfig.from_pipeline_config(
+        pipeline,
+        user_deploy_config=copy.deepcopy(deploy),
+        cli_overrides=resolved_cli_overrides,
+    )
+    return legacy_stages, omni_config
+
+
+def test_typed_llm_engine_args_preserve_legacy_adapter_behavior(tmp_path):
+    pipeline, deploy, model = _engine_arg_inputs(tmp_path)
+    legacy_stages, omni_config = _legacy_and_typed_stages(pipeline, deploy, model)
+    connector_spec = {"name": "SharedMemoryConnector", "extra": {"mode": "test"}}
+    cli_tokenizer = "/external/tokenizer"
+
+    typed_args_by_stage = {}
+    for stage_id in (0, 1):
+        legacy_args = build_legacy_engine_args_dict(
+            legacy_stages[stage_id],
+            model,
+            stage_connector_spec=connector_spec,
+            cli_tokenizer=cli_tokenizer,
+        )
+        typed_args = build_engine_args_dict_from_omni_stage_config(
+            omni_config.stage_by_id(stage_id),
+            model,
+            stage_connector_spec=connector_spec,
+            cli_tokenizer=cli_tokenizer,
+        )
+        typed_args_by_stage[stage_id] = typed_args
+
+        assert {name: typed_args[name] for name in legacy_args} == legacy_args
+
+    thinker_args = typed_args_by_stage[0]
+    assert thinker_args["model"] == str(tmp_path / "stage-model" / "ar-model")
+    assert thinker_args["tokenizer"] == str(tmp_path / "stage-model" / "ar-tokenizer")
+    assert thinker_args["stage_id"] == 0
+    assert thinker_args["stage_connector_spec"] == connector_spec
+    assert thinker_args["has_sampling_extra_args"] is True
+    assert thinker_args["omni_kv_config"] == {"need_send_cache": True}
+    assert thinker_args["worker_cls"] == "test.custom.Worker"
+    assert thinker_args["attention_backend"] == "FLASHINFER"
+    assert thinker_args["trust_remote_code"] is True
+    assert thinker_args["dtype"] == "float16"
+    assert thinker_args["distributed_executor_backend"] == "mp"
+    assert thinker_args["tensor_parallel_size"] == 2
+    assert thinker_args["hf_overrides"] == {"rope_scaling": {"factor": 2}}
+    assert thinker_args["limit_mm_per_prompt"] == {"audio": 1}
+    assert thinker_args["logits_processors"] == ["test.custom.LogitsProcessor"]
+    assert thinker_args["retains_state_across_chunks"] is True
+    assert thinker_args["kv_cache_memory_bytes"] == 1024
+    assert thinker_args["max_cudagraph_capture_size"] == 0
+
+    talker_args = typed_args_by_stage[1]
+    assert talker_args["model"] == model
+    assert talker_args["tokenizer"] == cli_tokenizer
+    assert talker_args["worker_cls"] == "test.generation.Worker"
+    assert talker_args["disable_hybrid_kv_cache_manager"] is True
+    assert talker_args["enable_prefix_caching"] is False
+
+
+def test_typed_diffusion_engine_args_use_structured_diffusion_config(tmp_path):
+    pipeline, deploy, model = _engine_arg_inputs(tmp_path)
+    legacy_stages, omni_config = _legacy_and_typed_stages(pipeline, deploy, model)
+
+    legacy_args = build_legacy_engine_args_dict(legacy_stages[2], model)
+    typed_args = build_engine_args_dict_from_omni_stage_config(
+        omni_config.stage_by_id(2),
+        model,
+    )
+
+    assert typed_args["stage_id"] == legacy_args["stage_id"] == 2
+    assert typed_args["model"] == legacy_args["model"] == str(tmp_path / "diffusion-model")
+    assert omni_config.stage_by_id(2).diffusion_config.model == str(tmp_path / "diffusion-model")
+    assert typed_args["model_stage"] == legacy_args["model_stage"] == "dit"
+    assert typed_args["model_arch"] == legacy_args["model_arch"] == "TypedDiffusionPipeline"
+    assert typed_args["model_class_name"] == legacy_args["model_class_name"] == "TypedDiffusionPipeline"
+    assert typed_args["engine_backend"] == legacy_args["engine_backend"] == "test.diffusion.Engine"
+    assert typed_args["enable_session_state_manager"] is legacy_args["enable_session_state_manager"] is True
+    assert typed_args["parallel_config"]["tensor_parallel_size"] == 2
+    assert typed_args["parallel_config"]["vae_parallel_mode"] == "spatial_shard_height"
+    assert isinstance(typed_args["diffusion_attention_config"], AttentionConfig)
+    assert typed_args["diffusion_attention_config"].default.backend == "FLASH_ATTN"
+    assert typed_args["diffusion_attention_config"].per_role["cross"].backend == "TORCH_SDPA"
+
+
+def test_typed_engine_args_preserve_explicit_backend_default_overrides(tmp_path):
+    pipeline, deploy, model = _engine_arg_inputs(tmp_path)
+    deploy.enable_prefix_caching = False
+    deploy.enable_chunked_prefill = False
+    thinker_deploy = deploy.stages[0]
+    thinker_deploy.gpu_memory_utilization = 0.75
+    thinker_deploy.disable_hybrid_kv_cache_manager = False
+    thinker_deploy.async_scheduling = False
+    legacy_stages, omni_config = _legacy_and_typed_stages(pipeline, deploy, model)
+
+    legacy_args = build_legacy_engine_args_dict(legacy_stages[0], model)
+    typed_args = build_engine_args_dict_from_omni_stage_config(
+        omni_config.stage_by_id(0),
+        model,
+    )
+    expected = {
+        "gpu_memory_utilization": 0.75,
+        "enable_prefix_caching": False,
+        "disable_hybrid_kv_cache_manager": False,
+        "max_num_seqs": 8,
+        "enable_chunked_prefill": False,
+        "async_scheduling": False,
+    }
+
+    assert {name: legacy_args[name] for name in expected} == expected
+    assert {name: typed_args[name] for name in expected} == expected
+
+
+@pytest.mark.parametrize("stage_scoped", [False, True], ids=["global", "stage-scoped"])
+def test_typed_engine_args_preserve_inherited_model_and_load_cli_fields(tmp_path, stage_scoped):
+    pipeline, deploy, model = _engine_arg_inputs(tmp_path)
+    expected = {
+        "revision": "model-rev",
+        "tokenizer_revision": "tokenizer-rev",
+        "code_revision": "code-rev",
+        "seed": 123,
+        "download_dir": str(tmp_path / "downloads"),
+    }
+    cli_overrides = {f"stage_0_{name}" if stage_scoped else name: value for name, value in expected.items()}
+    legacy_stages, omni_config = _legacy_and_typed_stages(
+        pipeline,
+        deploy,
+        model,
+        cli_overrides,
+    )
+
+    legacy_args = build_legacy_engine_args_dict(legacy_stages[0], model)
+    typed_args = build_engine_args_dict_from_omni_stage_config(
+        omni_config.stage_by_id(0),
+        model,
+    )
+
+    assert {name: legacy_args[name] for name in expected} == expected
+    assert {name: typed_args[name] for name in expected} == expected
+
+
+def test_typed_engine_args_preserve_deploy_subdirectory_precedence(tmp_path):
+    pipeline, deploy, model = _engine_arg_inputs(tmp_path)
+    stage_model = tmp_path / "stage-model"
+    (stage_model / "override-model").mkdir()
+    (stage_model / "override-tokenizer").mkdir()
+    deploy.stages[0].engine_extras.update(
+        {
+            "model_subdir": "override-model",
+            "tokenizer_subdir": "override-tokenizer",
+        }
+    )
+    legacy_stages, omni_config = _legacy_and_typed_stages(pipeline, deploy, model)
+
+    typed_stage = omni_config.stage_by_id(0)
+    legacy_args = build_legacy_engine_args_dict(legacy_stages[0], model)
+    typed_args = build_engine_args_dict_from_omni_stage_config(typed_stage, model)
+
+    assert typed_stage.model_config.model_subdir == "override-model"
+    assert typed_stage.model_config.tokenizer_subdir == "override-tokenizer"
+    assert typed_args["model"] == legacy_args["model"] == str(stage_model / "override-model")
+    assert typed_args["tokenizer"] == legacy_args["tokenizer"] == str(stage_model / "override-tokenizer")
+
+
+@pytest.mark.parametrize(
+    ("cli_overrides", "expected"),
+    [
+        ({}, {"need_recv_cache": True}),
+        (
+            {"stage_0_omni_kv_config": {"need_send_cache": False}},
+            {"need_send_cache": False},
+        ),
+    ],
+    ids=["topology-over-deploy", "cli-over-topology"],
+)
+def test_typed_engine_args_preserve_legacy_omni_kv_precedence(tmp_path, cli_overrides, expected):
+    pipeline, deploy, model = _engine_arg_inputs(tmp_path)
+    thinker = replace(
+        pipeline.stages[0],
+        omni_kv_config={"need_recv_cache": True},
+    )
+    pipeline = replace(pipeline, stages=(thinker, *pipeline.stages[1:]))
+    legacy_stages, omni_config = _legacy_and_typed_stages(
+        pipeline,
+        deploy,
+        model,
+        cli_overrides,
+    )
+
+    legacy_args = build_legacy_engine_args_dict(legacy_stages[0], model)
+    typed_stage = omni_config.stage_by_id(0)
+    typed_args = build_engine_args_dict_from_omni_stage_config(typed_stage, model)
+
+    assert typed_stage.connector_config.omni_kv_config == expected
+    assert typed_args["omni_kv_config"] == legacy_args["omni_kv_config"] == expected
+
+
+def test_build_engine_args_dict_preserves_legacy_api(tmp_path):
+    pipeline, deploy, model = _engine_arg_inputs(tmp_path)
+    legacy_stages, _ = _legacy_and_typed_stages(pipeline, deploy, model)
+    connector_spec = {"name": "SharedMemoryConnector", "extra": {}}
+
+    compatibility_args = build_engine_args_dict(
+        copy.deepcopy(legacy_stages[0]),
+        model,
+        stage_connector_spec=connector_spec,
+        cli_tokenizer="/external/tokenizer",
+    )
+    explicitly_legacy_args = build_legacy_engine_args_dict(
+        copy.deepcopy(legacy_stages[0]),
+        model,
+        stage_connector_spec=connector_spec,
+        cli_tokenizer="/external/tokenizer",
+    )
+
+    assert compatibility_args == explicitly_legacy_args
+
+
+def test_legacy_engine_args_drop_none_tp_without_mutating_stage_config():
+    stage_config = types.SimpleNamespace(
+        stage_id=0,
+        stage_type="llm",
+        engine_args={"tensor_parallel_size": None},
+        default_sampling_params={},
+    )
+
+    engine_args = build_legacy_engine_args_dict(
+        stage_config,
+        model="test-model",
+    )
+
+    assert "tensor_parallel_size" not in engine_args
+    assert stage_config.engine_args == {"tensor_parallel_size": None}
+
+
+def test_typed_engine_args_own_rocm_attention_default(monkeypatch, tmp_path):
+    from vllm_omni import platforms
+
+    pipeline, deploy, model = _engine_arg_inputs(tmp_path)
+    legacy_stages, omni_config = _legacy_and_typed_stages(pipeline, deploy, model)
+    monkeypatch.setattr(platforms.current_omni_platform, "is_rocm", lambda: True)
+    aiter_module = types.ModuleType("vllm._aiter_ops")
+    setattr(
+        aiter_module,
+        "rocm_aiter_ops",
+        types.SimpleNamespace(is_enabled=lambda: False),
+    )
+    monkeypatch.setitem(sys.modules, "vllm._aiter_ops", aiter_module)
+
+    legacy_args = build_legacy_engine_args_dict(legacy_stages[1], model)
+    typed_args = build_engine_args_dict_from_omni_stage_config(
+        omni_config.stage_by_id(1),
+        model,
+    )
+
+    assert "attention_backend" not in legacy_args
+    assert typed_args["attention_backend"] == "TRITON_ATTN"
+
+
+def test_typed_ming_image_engine_args_defer_diffusion_batch_default():
+    pipeline = resolve_pipeline_config("ming_flash_omni_image")
+    assert pipeline is not None
+    deploy = load_deploy_config(_DEPLOY_DIR / "ming_flash_omni_image.yaml")
+    legacy_stages, omni_config = _legacy_and_typed_stages(
+        pipeline,
+        deploy,
+        model="/tmp",
+    )
+
+    legacy_args = build_legacy_engine_args_dict(legacy_stages[1], model="/tmp")
+    typed_args = build_engine_args_dict_from_omni_stage_config(
+        omni_config.stage_by_id(1),
+        model="/tmp",
+    )
+    typed_backend_args = {name: value for name, value in typed_args.items() if name in _DIFFUSION_BACKEND_FIELDS}
+
+    assert "max_num_seqs" not in legacy_args
+    assert "max_num_seqs" not in typed_args
+    assert OmniDiffusionConfig(**typed_backend_args).max_num_seqs == 1
+
+
+@pytest.mark.parametrize("model_type", sorted(OMNI_PIPELINES))
+def test_typed_engine_args_match_current_registry_backend_semantics(model_type):
+    pipeline = resolve_pipeline_config(model_type)
+    if pipeline is None:
+        pytest.skip(f"Pipeline {model_type!r} requires an HF config to resolve")
+
+    deploy = (
+        load_deploy_config(_DEPLOY_DIR / pipeline.default_deploy_config_name)
+        if pipeline.default_deploy_config_name is not None
+        else DeployConfig()
+    )
+    legacy_stages, omni_config = _legacy_and_typed_stages(
+        pipeline,
+        deploy,
+        model="/tmp",
+    )
+
+    for legacy_stage in legacy_stages:
+        stage_id = legacy_stage.stage_id
+        legacy_args = build_legacy_engine_args_dict(legacy_stage, model="/tmp")
+        typed_args = build_engine_args_dict_from_omni_stage_config(
+            omni_config.stage_by_id(stage_id),
+            model="/tmp",
+        )
+        if legacy_stage.stage_type == StageType.DIFFUSION:
+            backend_fields = _DIFFUSION_BACKEND_FIELDS
+            backend_config_cls = OmniDiffusionConfig
+        else:
+            backend_fields = _LLM_BACKEND_FIELDS
+            backend_config_cls = OmniEngineArgs
+
+        missing_fields = (legacy_args.keys() & backend_fields) - typed_args.keys()
+        assert not missing_fields, f"{model_type} stage {stage_id} lost backend fields: {sorted(missing_fields)}"
+
+        legacy_effective_args = _effective_backend_values(backend_config_cls, legacy_args)
+        typed_effective_args = _effective_backend_values(backend_config_cls, typed_args)
+        assert typed_effective_args == legacy_effective_args, (
+            f"{model_type} stage {stage_id} changed effective backend arguments"
+        )

tests/entrypoints/openai/test_batch_chat_completions.py

@@ -0,0 +1,150 @@
+import asyncio
+from unittest.mock import AsyncMock
+
+import pytest
+from openai.types.chat import ChatCompletionUserMessageParam
+from openai.types.chat.chat_completion_audio import ChatCompletionAudio
+from starlette.requests import Request
+from vllm.entrypoints.chat_utils import ChatCompletionMessageParam
+from vllm.entrypoints.openai.chat_completion.protocol import (
+    BatchChatCompletionRequest,
+    ChatCompletionResponse,
+    ChatCompletionResponseChoice,
+    ChatMessage,
+)
+from vllm.entrypoints.openai.engine.protocol import ErrorResponse, UsageInfo
+
+from vllm_omni.entrypoints.openai.batch_serving import OmniOpenAIServingChatBatch
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+MSG_BATCH: list[list[ChatCompletionMessageParam]] = [
+    [ChatCompletionUserMessageParam(role="user", content="What color is the sky?")],
+    [ChatCompletionUserMessageParam(role="user", content="What is 2+2?")],
+]
+collapse = OmniOpenAIServingChatBatch._maybe_collapse_choices
+
+
+# Helpers for testing creating a __new__ serving instance and ensuring IDs are correct
+def _make_raw_request(headers: list[tuple[bytes, bytes]]):
+    scope = {
+        "type": "http",
+        "method": "POST",
+        "path": "/v1/chat/completions/batch",
+        "headers": headers,
+    }
+    return Request(scope)
+
+
+def _make_completion(text):
+    return ChatCompletionResponse(
+        model="test-model",
+        choices=[
+            ChatCompletionResponseChoice(
+                index=0,
+                message=ChatMessage(role="assistant", content=text),
+            )
+        ],
+        usage=UsageInfo(prompt_tokens=0, completion_tokens=0, total_tokens=0),
+    )
+
+
+def _make_handler(responses):
+    handler = OmniOpenAIServingChatBatch.__new__(OmniOpenAIServingChatBatch)
+    handler.create_chat_completion = AsyncMock(side_effect=responses)
+    return handler
+
+
+# Helpers for creating choices of different modality types
+def _text_choice(text="hello"):
+    return ChatCompletionResponseChoice(
+        index=0,
+        message=ChatMessage(role="assistant", content=text),
+    )
+
+
+def _audio_choice():
+    audio = ChatCompletionAudio(id="a1", data="base64audio", expires_at=0, transcript="")
+    return ChatCompletionResponseChoice(
+        index=0,
+        message=ChatMessage(role="assistant", content=None, audio=audio),
+    )
+
+
+def _image_choice():
+    content = [{"type": "image_url", "image_url": {"url": "data:image/png;base64,abc"}}]
+    msg = ChatMessage.model_construct(role="assistant")
+    object.__setattr__(msg, "content", content)
+    return ChatCompletionResponseChoice.model_construct(
+        index=0,
+        message=msg,
+    )
+
+
+def _get_subrequest_ids(handler):
+    return [call.args[0].request_id for call in handler.create_chat_completion.call_args_list]
+
+
+@pytest.mark.parametrize("has_header", [True, False])
+def test_subrequest_ids_are_unique(has_header: bool):
+    """Ensure that a submitted batch request creates unique subrequest IDs"""
+    header = None if not has_header else b"user-123"
+    headers = [] if not header else [(b"x-request-id", header)]
+
+    handler = _make_handler([_make_completion("hello"), _make_completion("world")])
+    request = BatchChatCompletionRequest(
+        model="test-model",
+        messages=MSG_BATCH,
+        seed=42,
+    )
+    result = asyncio.run(handler.create_batch_chat_completion(request, _make_raw_request(headers)))
+    assert not isinstance(result, ErrorResponse)
+    assert result.id.startswith("chatcmpl-batch")
+
+    # Ensure request IDs are unique and that if we had a header, it's in all reqs
+    req_ids = _get_subrequest_ids(handler)
+    assert len(set(req_ids)) == len(req_ids) == 2
+
+    if header is not None:
+        assert all(header.decode("utf-8") in req_id for req_id in req_ids)
+
+
+### Tests for ensuring that the (currently batch specific) choice collapsing is correct
+def test_single_text_passthrough():
+    result = collapse([_text_choice()])
+    assert result.message.content == "hello"
+
+
+def test_single_image_passthrough():
+    result = collapse([_image_choice()])
+    assert isinstance(result.message.content, list)
+    assert result.message.content[0]["type"] == "image_url"
+
+
+def test_text_plus_audio():
+    result = collapse([_text_choice(), _audio_choice()])
+    assert result.message.content == "hello"
+    assert result.message.audio.data == "base64audio"
+
+
+def test_audio_plus_text_order_independent():
+    result = collapse([_audio_choice(), _text_choice()])
+    assert result.message.content == "hello"
+    assert result.message.audio.data == "base64audio"
+
+
+def test_image_plus_audio():
+    result = collapse([_image_choice(), _audio_choice()])
+    assert isinstance(result.message.content, list)
+    assert result.message.content[0]["type"] == "image_url"
+    assert result.message.audio.data == "base64audio"
+
+
+def test_two_content_choices_raises():
+    with pytest.raises(ValueError, match="Multiple content choices cannot be set"):
+        collapse([_text_choice(), _text_choice()])
+
+
+def test_three_choices_raises():
+    with pytest.raises(ValueError, match="got 3"):
+        collapse([_text_choice(), _audio_choice(), _text_choice()])

tests/entrypoints/openai_api/test_serving_chat_sampling_params.py

@@ -95,6 +95,7 @@ def test_serving_boundary_normalizes_declared_root_and_nested_extras(serving_cha
         messages=[],
         modalities=["image"],
         num_inference_steps=7,
+        quality="high",
         size="768x512",
         negative_prompt="avoid blur",
         lora={"name": "adapter"},
@@ -110,6 +111,7 @@ def test_serving_boundary_normalizes_declared_root_and_nested_extras(serving_cha
     }
     assert diffusion_request_args["cfg_text_scale"] == 7.0
     assert diffusion_request_args["num_inference_steps"] == 7
+    assert diffusion_request_args["quality"] == "high"
     assert diffusion_request_args["size"] == "768x512"
     assert diffusion_request_args["negative_prompt"] == "avoid blur"
     assert diffusion_request_args["lora"] == {"name": "adapter"}
@@ -133,6 +135,19 @@ def test_unknown_root_extra_does_not_claim_canonical_extra(serving_chat):
     assert "pipeline_option" not in diffusion_request_args
 
 
+def test_serving_boundary_rejects_invalid_quality(serving_chat):
+    serving_chat._diffusion_extra_body_params = frozenset()
+    request = ChatCompletionRequest(
+        model="test",
+        messages=[],
+        modalities=["image"],
+        quality="medium",
+    )
+
+    with pytest.raises(ValueError, match="quality must be one of"):
+        serving_chat._normalize_diffusion_request_args(request)
+
+
 def test_unregistered_cfg_scale_aliases_common_true_cfg_scale(serving_chat):
     serving_chat._diffusion_extra_body_params = frozenset()
     request = ChatCompletionRequest(
@@ -226,6 +241,7 @@ def test_pure_consumer_preserves_defaults_and_separate_cfg_owners(serving_chat,
         negative_prompt="avoid blur",
         cfg_scale=7.0,
         true_cfg_scale=2.0,
+        quality="high",
         extra_body={"extra_args": {"solver": "ddim"}},
     )
 
@@ -239,6 +255,7 @@ def test_pure_consumer_preserves_defaults_and_separate_cfg_owners(serving_chat,
     }
     assert sampling_params.true_cfg_scale == 2.0
     assert sampling_params.num_inference_steps == 7
+    assert sampling_params.quality == "high"
     assert (sampling_params.height, sampling_params.width) == (512, 768)
     assert captured["prompt"]["negative_prompt"] == "avoid blur"
     assert response.error.message == "No output generated from AsyncOmni"
@@ -310,6 +327,7 @@ def test_mixed_consumer_keeps_root_common_args_with_nested_extras(serving_chat,
     assert sampling_params_list[0].extra_args == {"target_h": 512, "target_w": 768}
     diffusion_params = sampling_params_list[1]
     assert diffusion_params.num_inference_steps == 7
+    assert diffusion_params.quality is None
     assert (diffusion_params.height, diffusion_params.width) == (512, 768)
     assert diffusion_params.extra_args == {"solver": "euler"}
     assert captured["prompt"]["negative_prompt"] == "avoid blur"

tests/entrypoints/openai_api/test_serving_video_output_stream.py

@@ -17,6 +17,7 @@ from pytest_mock import MockerFixture
 from starlette.testclient import TestClient
 
 from tests.helpers.media import generate_synthetic_image
+from vllm_omni.entrypoints.openai.protocol.videos import VideoGenerationRequest
 from vllm_omni.entrypoints.openai.serving_video_output_stream import OmniStreamingVideoOutputHandler
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
@@ -103,6 +104,55 @@ def _build_test_app(
 class TestStreamingVideoOutputWebSocket:
     """WebSocket protocol tests for streaming generated video output."""
 
+    def test_sampling_params_quality_reaches_h3_pipeline_boundary(
+        self,
+        mocker: MockerFixture,
+    ) -> None:
+        async def mock_generate(*_args, **_kwargs):
+            if False:
+                yield None
+
+        _app, handler, _engine_client = _build_test_app(
+            mocker=mocker,
+            streaming_chunks=[],
+            mock_generate=mock_generate,
+        )
+        request_data = handler._coerce_request_data(
+            {
+                "type": "session.start",
+                "prompt": "H3 streaming quality",
+                "sampling_params": {"quality": "high"},
+            }
+        )
+        request = VideoGenerationRequest(**request_data)
+
+        _prompt, sampling_params, _video_params = asyncio.run(handler._build_prompt_and_sampling_params(request))
+
+        assert sampling_params.quality == "high"
+
+    def test_omitted_quality_preserves_engine_default(self, mocker: MockerFixture) -> None:
+        async def mock_generate(*_args, **_kwargs):
+            if False:
+                yield None
+
+        _app, handler, engine_client = _build_test_app(
+            mocker=mocker,
+            streaming_chunks=[],
+            mock_generate=mock_generate,
+        )
+        engine_client.default_sampling_params_list = [OmniDiffusionSamplingParams(quality="high")]
+        request_data = handler._coerce_request_data(
+            {
+                "type": "session.start",
+                "prompt": "H3 streaming default quality",
+            }
+        )
+        request = VideoGenerationRequest(**request_data)
+
+        _prompt, sampling_params, _video_params = asyncio.run(handler._build_prompt_and_sampling_params(request))
+
+        assert sampling_params.quality == "high"
+
     def test_streaming_session_emits_video_start_binary_chunks_and_done(self, mocker: MockerFixture):
         """A full session delivers video.start, binary chunks, then session.done."""
 

tests/entrypoints/openai_api/test_video_api_utils.py

@@ -48,6 +48,52 @@ def test_encode_video_bytes_exports_frames_without_interpolation(monkeypatch):
     assert mux_calls[0]["audio"] is None
 
 
+def test_float_frames_are_converted_without_stacking_full_video(monkeypatch):
+    frame = np.array(
+        [
+            [[0.0, 0.5, 1.0, 0.2], [1.5, -0.5, 0.5, 0.8]],
+        ],
+        dtype=np.float32,
+    )
+    original = frame.copy()
+
+    def fail_stack(*args, **kwargs):
+        raise AssertionError("float video conversion must not stack all frames")
+
+    monkeypatch.setattr(video_api_utils.np, "stack", fail_stack)
+
+    frames = video_api_utils._coerce_video_to_uint8_frames([frame, frame])
+
+    expected = np.array([[[128, 191, 255], [255, 64, 191]]], dtype=np.uint8)
+    assert frames.flags.c_contiguous
+    np.testing.assert_array_equal(frames, np.array([expected, expected]))
+    np.testing.assert_array_equal(frame, original)
+
+
+def test_two_dimensional_width_four_frames_are_not_treated_as_rgba():
+    frame = np.arange(8, dtype=np.float32).reshape(2, 4) / 8
+
+    frames = video_api_utils._coerce_video_to_uint8_frames([frame, frame])
+
+    expected = np.rint(frame * 255).astype(np.uint8)
+    assert frames.shape == (2, 2, 4)
+    np.testing.assert_array_equal(frames, np.stack([expected, expected]))
+
+
+def test_mixed_float_dtypes_preserve_stacked_rounding_semantics():
+    frames = [
+        np.full((1, 1, 3), 0.1, dtype=np.float16),
+        np.full((1, 1, 3), 0.2, dtype=np.float32),
+    ]
+    stacked = np.stack(frames)
+    expected = np.rint(np.clip(stacked, 0.0, 1.0) * 255.0).astype(np.uint8)
+
+    actual = video_api_utils._coerce_video_to_uint8_frames(frames)
+
+    assert expected[0, 0, 0, 0] == 25
+    np.testing.assert_array_equal(actual, expected)
+
+
 @pytest.mark.parametrize("frame_count", [3, 4])
 def test_channel_last_video_tensor_preserves_channel_sized_frame_count(frame_count):
     video = torch.arange(frame_count * 2 * 5 * 3, dtype=torch.uint8).reshape(frame_count, 2, 5, 3)

tests/entrypoints/openai_api/test_video_server.py

@@ -868,6 +868,7 @@ def test_default_sampling_params_apply_to_video_requests(test_client, mocker: Mo
         OmniDiffusionSamplingParams(
             num_inference_steps=4,
             guidance_scale=7.5,
+            quality="high",
             generator_device="cpu",
             enable_frame_interpolation=True,
             frame_interpolation_exp=2,
@@ -890,6 +891,7 @@ def test_default_sampling_params_apply_to_video_requests(test_client, mocker: Mo
     captured = engine.captured_sampling_params_list[0]
     assert captured.num_inference_steps == 4
     assert captured.guidance_scale == 7.5
+    assert captured.quality == "high"
     assert captured.generator_device == "cpu"
     assert captured.enable_frame_interpolation is True
     assert captured.frame_interpolation_exp == 2
@@ -1445,6 +1447,7 @@ def test_invalid_uploaded_input_reference_returns_400(test_client):
 def test_video_request_validation():
     req = VideoGenerationRequest(prompt="test")
     assert req.prompt == "test"
+    assert req.quality is None
     assert req.generate_sound is False
     assert req.sound_duration is None
     assert VideoGenerationRequest(prompt="test", generate_sound=True, sound_duration=1.5).generate_sound is True
@@ -1464,6 +1467,8 @@ def test_video_request_validation():
         VideoGenerationRequest(prompt="test", frame_interpolation_scale=0)
     with pytest.raises(ValueError):
         VideoGenerationRequest(prompt="test", sound_duration=0)
+    with pytest.raises(ValueError):
+        VideoGenerationRequest(prompt="test", quality="medium")
 
 
 def test_list_videos_supports_order_after_and_limit(test_client, mocker: MockerFixture):
@@ -1975,6 +1980,7 @@ def test_sync_sampling_params_pass_through(test_client, mocker: MockerFixture):
             "num_inference_steps": "30",
             "guidance_scale": "6.5",
             "seed": "42",
+            "quality": "high",
         },
     )
     assert response.status_code == 200
@@ -1983,6 +1989,7 @@ def test_sync_sampling_params_pass_through(test_client, mocker: MockerFixture):
     assert captured.num_inference_steps == 30
     assert captured.guidance_scale == 6.5
     assert captured.seed == 42
+    assert captured.quality == "high"
 
 
 def test_sync_frame_interpolation_params_pass_to_sampling_params(test_client, mocker: MockerFixture):
@@ -2018,6 +2025,7 @@ def test_sync_default_sampling_params_apply_to_video_requests(test_client, mocke
         OmniDiffusionSamplingParams(
             num_inference_steps=4,
             guidance_scale=7.5,
+            quality="high",
             enable_frame_interpolation=True,
             frame_interpolation_exp=2,
             frame_interpolation_scale=0.5,
@@ -2038,6 +2046,7 @@ def test_sync_default_sampling_params_apply_to_video_requests(test_client, mocke
     captured = engine.captured_sampling_params_list[0]
     assert captured.num_inference_steps == 4
     assert captured.guidance_scale == 7.5
+    assert captured.quality == "high"
     assert captured.enable_frame_interpolation is True
     assert captured.frame_interpolation_exp == 2
     assert captured.frame_interpolation_scale == 0.5

tests/entrypoints/test_async_omni_diffusion_config.py

@@ -240,6 +240,28 @@ def test_serve_cli_accepts_ulysses_mode():
     assert parallel_config.ulysses_mode == "advanced_uaa"
 
 
+def test_serve_cli_forwards_model_defined_task_type_to_diffusion_stage():
+    parser = TrackingArgumentParser()
+    subparsers = parser.add_subparsers(dest="command")
+    OmniServeCommand().subparser_init(subparsers)
+
+    args = parser.parse_args(
+        [
+            "serve",
+            "MiniMaxAI/MiniMax-H3",
+            "--omni",
+            "--task-type",
+            "fl2va",
+        ]
+    )
+
+    explicit_kwargs = args.get_explicit_kwargs_dict()
+    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(explicit_kwargs)[0]
+
+    assert args.task_type == "fl2va"
+    assert stage_cfg["engine_args"]["task_type"] == "fl2va"
+
+
 def test_serve_cli_accepts_diffusion_pipeline_profiler_flag():
     """Ensure diffusion serve CLI exposes the profiler switch."""
     parser = TrackingArgumentParser()

tests/entrypoints/test_omni_snapshot_download.py

@@ -53,6 +53,11 @@ def download_backend(monkeypatch: pytest.MonkeyPatch):
         "download_weights_from_hf_specific",
         lambda **_kwargs: picked.append("huggingface"),
     )
+    monkeypatch.setattr(
+        omni_base,
+        "file_or_path_exists",
+        lambda *_args, **_kwargs: False,
+    )
 
     def run() -> str:
         picked.clear()
@@ -82,3 +87,14 @@ def test_unset_defaults_to_huggingface(monkeypatch, download_backend):
     monkeypatch.delenv("VLLM_USE_MODELSCOPE", raising=False)
 
     assert download_backend() == "huggingface"
+
+
+def test_modular_diffusers_defers_component_download(monkeypatch, download_backend):
+    monkeypatch.delenv("VLLM_USE_MODELSCOPE", raising=False)
+    monkeypatch.setattr(
+        omni_base,
+        "file_or_path_exists",
+        lambda *_args, **_kwargs: True,
+    )
+
+    assert download_backend() == "none"