docs: address architecture review feedback

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

.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"

README.md

@@ -9,7 +9,7 @@ Easy, fast, and cheap omni-modality model serving for everyone
 </h3>
 
 <p align="center">
-| <a href="https://vllm-omni.readthedocs.io/en/latest/"><b>Documentation</b></a> | <a href="https://deepwiki.com/vllm-project/vllm-omni"><b>DeepWiki</b></a> | <a href="https://discuss.vllm.ai"><b>User Forum</b></a> | <a href="https://slack.vllm.ai"><b>Developer Slack</b></a> | <a href="docs/assets/WeChat.jpg"><b>WeChat</b></a> | <a href="https://arxiv.org/abs/2602.02204"><b>Paper</b></a> | <a href="https://docs.google.com/presentation/d/111-L8zF7A1j_YI_cR8JsblofdScdRr2f/edit?usp=sharing&ouid=110473603432222024453&rtpof=true&sd=true"><b>Slides</b></a> |
+| <a href="https://vllm-omni.readthedocs.io/en/latest/"><b>Documentation</b></a> | <a href="https://deepwiki.com/vllm-project/vllm-omni"><b>DeepWiki</b></a> | <a href="https://discuss.vllm.ai"><b>User Forum</b></a> | <a href="https://slack.vllm.ai"><b>Developer Slack</b></a> | <a href="docs/assets/WeChat.jpg"><b>WeChat</b></a> | <a href="https://arxiv.org/abs/2602.02204"><b>Paper</b></a> | <a href="https://docs.google.com/presentation/d/1aPj0OGl_-ZVoib-Qne5dGDAlrRFB-PdHl6E-EE99g8E/edit?usp=sharing"><b>Slides</b></a> |
 </p>
 
 

docs/README.md

@@ -63,5 +63,5 @@ vLLM-Omni seamlessly supports most popular open-source models on HuggingFace, in
 
 For more information, checkout the following:
 
-- [vllm-omni architecture design and recent roadmaps](https://docs.google.com/presentation/d/1XJWgv79lORl8rbaVvp2d5Sqs6ZEBgAgj/edit?slide=id.p1#slide=id.p1)
+- [vllm-omni architecture design and recent roadmaps](https://docs.google.com/presentation/d/1aPj0OGl_-ZVoib-Qne5dGDAlrRFB-PdHl6E-EE99g8E/edit?usp=sharing)
 - [vllm-omni announcement blogpost](https://blog.vllm.ai/2025/11/30/vllm-omni.html)

docs/design/architecture_overview.md

@@ -1,6 +1,6 @@
 # Architecture Overview
 
-This document outlines the architectural design for vLLM-Omni.
+This document outlines the architecture of vLLM-Omni.
 
 <p align="center">
   <picture>
@@ -9,194 +9,456 @@ This document outlines the architectural design for vLLM-Omni.
   </picture>
 </p>
 
-# Goals
+## Design goals
 
-The primary goal of the vLLM-Omni project is to build the fastest and easiest-to-use open-source Omni-Modality model inference & serving engine. vLLM-Omni extends the original vLLM, which was created to support large language models for text-based autoregressive (AR) generation tasks. vLLM-Omni is designed to support:
+The primary goal of vLLM-Omni is to provide a fast and easy-to-use inference
+and serving engine for omni-modality models. vLLM-Omni extends vLLM's
+text-oriented autoregressive (AR) runtime with stage-based execution for
+non-textual outputs and non-autoregressive model components.
 
-* **Non-textual Output:** Enables the integration, efficient processing and output of various data types, including but not limited to, images, audio, video, and action trajectories, alongside text.
-* **Non-Autoregressive Structure:** Support model structure beyond autoregressive, especially Diffusion Transformer (DiT), which is widely used in visual and audio generation.
-* **Integration with vLLM Core:** Maintain compatibility and leverage existing vLLM key modules and optimizations where applicable.
-* **Extensibility:** Design a modular and flexible architecture that can easily accommodate new modalities, model architectures, and output formats.
+The architecture is designed to:
 
+* support text, image, audio, video, and action inputs and outputs;
+* compose autoregressive, generation, and diffusion stages in one pipeline;
+* reuse vLLM scheduling, cache, distributed-execution, and serving
+  primitives where they fit; and
+* keep model topology, deployment placement, runtime lifecycle, transport,
+  and public API concerns in separate layers.
 
-# Representative omni-modality models
+## Model execution
 
-According to analysis for current popular open-source models, most of them have the combination of AR+DiT. Specifically, they can be further categorized into 3 types below:
+These four examples show how vLLM-Omni maps model components to stage-local
+execution policies. Each summary covers the model-owned path, batching,
+attention, parallelism, quantization, and primary serving targets. The profiles
+are representative; a model may support multiple `PipelineConfig` and
+`DeployConfig` combinations.
 
-**DiT as a main structure, with AR as text encoder (e.g.: Qwen-Image)**
-  A powerful image generation foundation model capable of complex text rendering and precise image editing.
+| Model | Model-owned path | vLLM-Omni stage view | Primary serving target |
+| --- | --- | --- | --- |
+| Qwen3-Omni | Thinker → Talker → Code2Wav | Three stages with asynchronous chunk hand-off | TTFT/TTFP, TPOT, E2EL, audio RTF, and concurrent throughput |
+| HunyuanImage-3.0 | Multimodal AR understanding/reasoning plus image-generation DiT | AR-only, DiT-only, or split AR → DiT deployment | DiT E2EL/denoising latency and image throughput; TTFT/TPOT for AR tasks |
+| MiniMax-H3 | Text encoder → task-specific FL2VA or Ref2VA DiT → video/audio VAE decoder | Three logical component/phase boundaries within one registered diffusion stage | Video/audio E2EL and media throughput; RTF when audio streaming is measured |
+| Cosmos3 | Unified MoT reasoner tower plus diffusion generator tower | One diffusion pipeline materializes both towers inside the stage | Image/video/action E2EL and throughput; RTF for synchronized audio |
+
+### Serving metric vocabulary
+
+The target metric follows the output modality, not just the model name.
+
+| Metric | Meaning in this document | Most useful for |
+| --- | --- | --- |
+| **TTFT** | Time to the first text token | AR reasoning or text output |
+| **TPOT** | Time per output token after the first token | AR decode stages |
+| **TTFP** | Time to the first streamed media packet, such as audio | Interactive speech and multimodal chat |
+| **E2EL** | End-to-end latency from request admission to the final output | Image, video, audio, and action requests |
+| **RTF** | Wall-clock processing time divided by generated audio duration; lower is faster than real time | Audio generation |
+| **Throughput** | Requests, tokens, or media seconds produced per wall-clock second | Offline and concurrent serving |
+
+### Qwen3-Omni: streaming AR pipeline
+
+Qwen3-Omni illustrates a multi-stage streaming AR pipeline: Thinker → Talker →
+Code2Wav. The [technical report](https://arxiv.org/abs/2509.17765) describes
+the model components, and the figure shows the stage-level streaming schedule
+used by vLLM-Omni.
 
 <p align="center">
-  <picture>
-    <source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/dit-main-architecture.png">
-    <img alt="Qwen-Image" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/dit-main-architecture.png" width=30%>
-  </picture>
+  <a href="https://github.com/vllm-project/vllm-omni/blob/main/docs/source/architecture/qwen3-omni-async-chunk.png">
+    <img alt="Qwen3-Omni Thinker, Talker, and Code2Wav streaming stages" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/qwen3-omni-async-chunk.png" width="80%">
+  </a>
 </p>
 
-**AR as a main structure, with DiT as multi-modal generator (e.g. BAGEL)**
-  A unified multimodal comprehension and generation model, with cot text output and visual generation.
+#### Model architecture by stage
+
+| Stage | Model component | Stage output |
+| --- | --- | --- |
+| **Thinker** | AuT audio encoder, SigLIP2 vision encoder, and the 30B-total/3B-active Thinker MoE Transformer | Text tokens and conditioning for the Talker |
+| **Talker** | 3B-total/0.3B-active Talker MoE Transformer with the MTP code-prediction path | Text tokens and audio codec tokens |
+| **Code2Wav** | Approximately 200M-parameter ConvNet audio decoder | Waveform/audio chunks |
+
+#### Representative stage execution
+
+| Stage | Batching | Attention and execution | Parallelism | Quantization |
+| --- | --- | --- | --- | --- |
+| **Thinker** | Continuous batching with stage-local `max_num_seqs` and token budgets | vLLM KV-cached AR attention; CUDA Graph execution when the stage is not eager | Independent stage placement; use the AR runtime's configured tensor parallelism | Supported ModelOpt FP8/NVFP4 or AutoRound checkpoint paths target the Thinker; encoders remain BF16 in the documented paths |
+| **Talker** | Continuous batching over autoregressive decode requests | KV-cached token decode with async-chunk output to Code2Wav | Independent stage placement; typically a separate GPU or replica from Thinker | BF16 baseline; no generic Talker quantization path is assumed |
+| **Code2Wav** | Static/chunk batching for codec-to-waveform decode | Non-AR chunked audio decode; eager execution is commonly retained | Separate stage or colocated placement, depending on the deployment | BF16 baseline; keep codec/vocoder precision model-specific |
+
+For this pipeline, the main serving target is interactive responsiveness:
+TTFT for text, TTFP for the first audio packet, TPOT for autoregressive decode,
+E2EL for the complete answer, and RTF/throughput for sustained audio serving.
+Async chunking and streaming primarily reduce TTFP and improve stage overlap;
+batching and CUDA Graphs primarily improve E2EL and throughput. See the
+[Qwen3-Omni recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Qwen/Qwen3-Omni.md)
+and [async chunk design](feature/async_chunk.md) for the deployment details.
+
+### HunyuanImage-3.0: shared multimodal model, split deployment choices
+
+HunyuanImage-3.0 illustrates a shared multimodal model with flexible AR, DiT,
+and split AR → DiT deployment. The official [technical
+report](https://arxiv.org/abs/2509.23951) provides the model view; the tables
+below map it to vLLM-Omni stages.
 
 <p align="center">
-  <picture>
-    <source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/ar-main-architecture.png">
-    <img alt="Bagel" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/ar-main-architecture.png" width=30%>
-  </picture>
+  <a href="https://github.com/Tencent-Hunyuan/HunyuanImage-3.0/blob/main/assets/framework.png">
+    <img alt="HunyuanImage-3.0 framework for understanding, language modeling, and image generation" src="https://raw.githubusercontent.com/Tencent-Hunyuan/HunyuanImage-3.0/main/assets/framework.png" width="95%">
+  </a>
 </p>
 
-**AR+DiT (e.g. Qwen-Omni)**
-  A natively end-to-end omni-modal LLM for multimodal inputs (text/image/audio/video...) and outputs (text/audio...).
+#### Model architecture by stage
+
+| Stage or component | Model component | Stage output |
+| --- | --- | --- |
+| **AR understanding/reasoning** | Understanding and generation encoders feeding the Hunyuan-A13B decoder-only Transformer and text detokenizer | Text responses or image-generation conditioning |
+| **Image generation** | Generation encoder, diffusion prediction path/Gen. Decoder, and the image VAE decode path | Image latents and final images |
+
+The vLLM-Omni deployment may expose the AR path only, the DiT path only, or
+both as an AR → DiT pipeline with a connector between them. That is a runtime
+decomposition of the model's capabilities; it should not be read as a claim
+that the published framework figure contains two unrelated backbones.
+
+#### Representative stage execution
+
+| Stage | Batching | Attention and execution | Parallelism | Quantization |
+| --- | --- | --- | --- | --- |
+| **AR** | Standard vLLM continuous batching when serving understanding/text tasks | KV-cached causal attention | Tensor/expert parallelism follows the selected AR deployment | Checkpoint-specific; do not infer DiT FP8 settings for the AR checkpoint |
+| **DiT + VAE** | Request batching or step-wise batching; Hunyuan step batching with more than one request requires `TORCH_SDPA` | `TORCH_SDPA` is used because the model mixes causal and full attention | Validated profiles include TP4, TP2 + sequence parallelism, TP2 + CFG parallelism, and expert parallelism for MoE weights | Online FP8 and ModelOpt mixed FP8/NVFP4 are documented for the DiT; VAE precision stays separate |
+
+Image generation is therefore measured mainly with per-image E2EL/denoising
+latency, peak memory, and images-per-second throughput. TTFT and TPOT remain
+meaningful for the optional AR understanding path, but they are not the right
+headline metrics for the DiT path. See the
+[HunyuanImage-3.0 recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Tencent/HunyuanImage-3.0-Instruct.md)
+for the validated batching, attention, and quantization combinations.
+
+### MiniMax-H3: conditioning, denoising, and decode
+
+MiniMax-H3 illustrates logical component and phase boundaries between text
+encoding, task-specific denoising, and media decoding:
+`Text Encoder → DiT → VAE Decoder`. The current registered topology co-locates
+these components in one diffusion stage. These boundaries are useful for
+profiling and offload; independent placement would require a future
+multi-stage `PipelineConfig`. The stage mapping follows the implementation and
+serving recipe; the [release overview](https://minimaxi.com/blog/minimax-h3)
+notes that a detailed technical report is forthcoming.
+
+#### Model architecture by component
+
+| Logical component | Model component | Component output |
+| --- | --- | --- |
+| **Text Encoder** | Tokenizer/processor and Qwen3-VL text encoder; vision/audio references are prepared alongside the text condition | Packed multimodal conditioning for denoising |
+| **DiT** | Task-specific `FL2VA` DiT for text/first-frame generation or `Ref2VA` DiT for mixed references | Video and audio latent tokens |
+| **VAE Decoder** | Video VAE and audio VAE decode the generated latents and assemble the synchronized output | Decoded video frames and stereo waveform, with FPS and audio sample rate metadata |
+
+Both DiT partitions remain behind one task-selection point; `task` chooses
+which partition runs. In the current implementation, the three logical
+components execute inside one diffusion stage. This is different from
+Qwen3-Omni's three independently scheduled stages, even though both models
+produce synchronized audio and video-like outputs.
+
+#### Representative component execution
+
+| Logical component | Batching | Attention and execution | Parallelism | Quantization |
+| --- | --- | --- | --- | --- |
+| **Text Encoder** | The current co-located pipeline follows the diffusion request scheduler; independent prompt/reference batching would require a separate multi-stage topology | Qwen3-VL attention and multimodal preprocessing | `--text-encoder-tp-size N` shards the encoder across the first `N` DiT ranks; independent placement is not part of the current topology | BF16/FP32 baseline; the H3 DiT FP8 path does not quantize the text encoder |
+| **DiT** | The current H3 implementation executes one generation request per diffusion batch; use concurrency for service-level throughput | cuDNN attention for the validated two-GPU consumer profile; TRTLLM attention or FlashAttention-4 on supported Blackwell profiles | TP2 plus text-encoder TP and Ulysses/VAE parallel groups on multi-GPU profiles; DLO/CPU offload trade memory for transfer time | Online FP8 applies to eligible DiT linears and is incompatible with layerwise offload |
+| **VAE Decoder** | Decode follows each generation request; tile/patch work can be distributed even when denoising is not batched | VAE decode kernels rather than DiT attention | VAE patch parallelism and native tiled decode within the diffusion stage | BF16/FP32 baseline; the documented H3 FP8 path leaves both VAEs unchanged |
+
+The primary target is video/audio E2EL and media throughput, with the logical
+component boundaries making prompt encoding, denoising, and VAE decode costs
+separately measurable. RTF is useful when the audio stream is evaluated
+independently. TTFT and TPOT do not describe the main H3 generation path
+because the user-visible output is diffusion media rather than a token stream.
+The
+[MiniMax-H3 recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/MiniMaxAI/MiniMax-H3.md)
+contains the hardware-specific profiles and warmup requirements.
+
+### Cosmos3: unified MoT reasoner and generator
+
+Cosmos3 illustrates a unified Mixture-of-Transformers (MoT) reasoner and
+diffusion generator. The official [technical
+report](https://research.nvidia.com/labs/cosmos-lab/cosmos3/technical-report.pdf)
+and NVIDIA's [technical explanation](https://developer.nvidia.com/blog/develop-physical-ai-reasoning-world-and-action-models-with-nvidia-cosmos-3/)
+provide the two-tower model view used below.
 
 <p align="center">
-  <picture>
-    <source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/ar-dit-main-architecture.png">
-    <img alt="Qwen-Omni" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/ar-dit-main-architecture.png" width=30%>
-  </picture>
+  <a href="https://developer.nvidia.com/blog/develop-physical-ai-reasoning-world-and-action-models-with-nvidia-cosmos-3/">
+    <img alt="Cosmos3 autoregressive reasoner and diffusion generator towers" src="https://developer-blogs.nvidia.com/wp-content/uploads/2026/05/image-11-e1780000686151.webp" width="95%">
+  </a>
 </p>
 
-# vLLM-Omni main architecture
+#### Model architecture by stage
 
-<p align="center">
-  <picture>
-    <source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/vllm-omni-main-architecture.png">
-    <img alt="vLLM-Omni Main Architecture" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/vllm-omni-main-architecture.png" width=55%>
-  </picture>
-</p>
+| Stage or component | Model component | Stage output |
+| --- | --- | --- |
+| **Reasoner tower** | Autoregressive VLM path with causal attention over language and optional vision/action context; runs once per generation request | Discrete reasoning/context tokens |
+| **Generator tower** | Diffusion path with full attention over the generator tokens and the reasoner K/V context; runs for each denoising step | Image, video, audio, or action tokens |
+| **Shared pipeline** | `Cosmos3OmniDiffusersPipeline` plus modality encoders/decoders and VAE paths | Final image/video/audio or action response |
 
-## Key Components
+The reasoner and generator are model-owned components, not two independent
+vLLM-Omni stages by default. The current Cosmos3 pipeline materializes both
+inside one diffusion stage; model-level CPU offload can swap the reasoner and
+generator components between their phases.
 
-| Component         | Description                                                                                                                              |
-| ----------------- | ---------------------------------------------------------------------------------------------------------------------------------------- |
-| **OmniRouter**    | provide an intelligent router for Omni-modality requests dispatch                                                                        |
-| **EntryPoints**   | define the APIs for offline/online serving (APIServer, Omni/AsyncOmni), while `AsyncOmniEngine` and `Orchestrator` coordinate multi-stage AR/DiT execution |
-| **AR**            | adapted for omni-modality models while inheriting efficient features from vLLM, such as cache management                                 |
-| **Diffusion**     | natively implemented and optimized using acceleration components                                                                         |
-| **OmniConnector** | supports fully disaggregation based on E/P/D/G (Encoding/Processing/Decoding/Generation) disaggregation across stages                    |
+#### Representative stage execution
 
-Disaggregated stages are managed through stage configuration. In Qwen3-Omni, Thinker/Talker/Code2wav are declared as separate configured stages, and runtime routing is handled by `Orchestrator` over `StageEngineCoreClient` / `StageDiffusionClient`.
+| Stage | Batching | Attention and execution | Parallelism | Quantization |
+| --- | --- | --- | --- | --- |
+| **Reasoner + generator pipeline** | Use request batching or step execution only when the selected pipeline advertises that capability; otherwise use `max_num_seqs=1` | Causal attention for reasoner tokens and full attention for diffusion tokens; the validated recipe uses a platform-selected aiter FlashAttention path with `TORCH_SDPA` fallback | Ulysses/context parallelism or tensor parallelism for the transformer; VAE tiling and layerwise/model-level offload address activation and weight memory | Online FP8 is supported for the diffusion pipeline; AutoRound/checkpoint-specific paths and quality validation remain separate from the BF16 baseline |
+
+For Cosmos3, generation E2EL and output throughput are the primary targets.
+Action-policy workloads add a control-loop latency target, while synchronized
+audio workloads add RTF. TTFT/TPOT are appropriate for a reasoner-only service,
+but not for the full diffusion generation path. See the
+[Cosmos3 recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-Nano.md)
+and [diffusion execution modes](../user_guide/diffusion/execution_modes.md)
+for the current runtime choices.
+
+## Why the stage-based architecture follows from these examples
+
+The four representative pipelines show that omni-modality serving is not one
+execution problem. Autoregressive decoding, diffusion denoising, multimodal
+encoding, and media decoding have different scheduling, attention, parallelism,
+memory, and quantization requirements. The architecture therefore needs a
+unit that can be optimized and resourced independently while remaining part of
+one request pipeline.
+
+| Observation | Design implication | vLLM-Omni response |
+| --- | --- | --- |
+| AR, diffusion, encoder, and decoder components have different execution loops | One scheduler and one execution policy cannot fit every component | Specialized AR, generation, and diffusion stage runtimes |
+| A model can be deployed as AR-only, DiT-only, or AR → DiT | Model structure must be separate from deployment topology | `PipelineConfig` describes logical stages and relationships; `DeployConfig` describes placement and resources |
+| Text encoders, DiTs, and VAE decoders have different batch and memory behavior | Batching, attention, parallelism, and quantization must be stage-local | Each stage carries its own runtime and engine configuration |
+| MiniMax-H3 can co-locate or logically split its text encoder and VAE decoder | A stage boundary should be configurable rather than forcing a process boundary | Logical stages can be co-located or assigned independent runtime placement |
+| Qwen3-Omni transfers hidden states and codec chunks across stages | Large intermediate data needs typed transport and synchronization | `OmniConnector` transports payloads and KV-cache data without owning model routing |
+| Requests span multiple stages and may be cancelled or produce ordered streaming output | Cross-stage request state needs a separate control plane | `Orchestrator` owns correlation, cancellation, routing state, and output ordering |
+| Qwen3-Omni optimizes first-packet latency while image/video models optimize E2EL and throughput | Performance objectives must be model- and stage-aware | TTFT/TPOT/TTFP, E2EL, RTF, and throughput are measured according to the output modality |
+
+This leads to a precise definition of a stage: it is a logical execution unit
+with its own model component, execution policy, resource ownership, and
+performance objectives. A stage is not necessarily a separate process. Cosmos3
+keeps its reasoner and generator inside one diffusion pipeline, while
+MiniMax-H3 exposes useful logical boundaries between text encoding, denoising,
+and VAE decoding. Both are valid mappings of the same abstraction.
+
+The resulting separation can be summarized as:
+
+```text
+Model structure       What is computed and how components are related
+        ↓
+PipelineConfig        Which logical stages exist and how they are connected
+        ↓
+Stage runtime policy  How each stage batches, attends, parallelizes, and quantizes
+        ↓
+DeployConfig           Where stages run, with which devices, replicas, and connectors
+```
+
+`AsyncOmniEngine` and `Orchestrator` coordinate the pipeline, `OmniConnector`
+transports stage outputs, and `StageRuntime` materializes the selected
+placement. This separation lets one logical model support multiple serving
+topologies without coupling model code to process layout or deployment
+resources.
+
+## System architecture
+
+The figure shows the main components of the stage-based runtime: entrypoints,
+the shared `AsyncOmniEngine`, request orchestration, stage lifecycle
+management, AR and diffusion execution modules, model/layer operations,
+connector transport, and multimodal outputs. `AsyncOmniEngine` is the
+composition root between the entrypoints and the orchestrator.
+
+```mermaid
+flowchart TB
+    subgraph omni["vLLM-Omni"]
+        subgraph entry["Entry points"]
+            api["API server"]
+            client["Omni / AsyncOmni"]
+        end
+        engine["AsyncOmniEngine<br/>composition root and background loop"]
+        orchestrator["Orchestrator<br/>request state and cross-stage routing"]
+        runtime["StageRuntime<br/>placement, replicas, readiness, and lifecycle"]
+        ar["AR module<br/>vLLM engine, scheduler, cache, executor, worker, and model runner"]
+        diffusion["Diffusion module<br/>scheduler, worker, model runner, and pipeline"]
+        ops["Model / layer / ops<br/>attention, parallelism, and quantization"]
+        connector["OmniConnector<br/>payload and KV transport plus synchronization"]
+        outputs["Multimodal outputs<br/>streaming chunks and final artifacts"]
+
+        api --> engine
+        client --> engine
+        engine --> orchestrator
+        engine --> runtime
+        orchestrator --> runtime
+        runtime --> ar
+        runtime --> diffusion
+        ar -. uses .-> ops
+        diffusion -. uses .-> ops
+        ar -. transports through .-> connector
+        diffusion -. transports through .-> connector
+        ar --> outputs
+        diffusion --> outputs
+    end
+```
+
+### Key components
+
+| Component | Responsibility |
+| --- | --- |
+| **Entrypoints** | Translate offline, CLI, OpenAI-compatible, and duplex requests into engine operations and render outputs back to public protocols. |
+| **Configuration resolution** | Combines pipeline topology, deployment settings, model metadata, and user overrides into a validated control-plane configuration. |
+| **AsyncOmniEngine** | Owns engine composition, the background event loop, stage initialization, request submission, and output collection. |
+| **Orchestrator** | Owns cross-stage request state, stage-to-stage routing, request correlation, cancellation, and output ordering. It does not own model selection or deployment placement. |
+| **StageRuntime** | Expands logical stages into local or distributed replicas, starts stage clients and processes, and manages readiness, affinity, failure, and shutdown. |
+| **AR runtime** | Extends vLLM's scheduler, KV-cache, worker, and model-runner path for omni-modality inputs and inter-stage outputs. |
+| **Diffusion runtime** | Schedules and executes denoising workloads through diffusion executors, workers, pipelines, acceleration backends, and output materialization. |
+| **OmniConnector** | Transports stage payloads and KV-cache data and provides synchronization. Connectors transport data; they do not choose the next logical stage. |
+| **Multimodal outputs** | `MultimodalPayload` separates tensor content from metadata, while `OmniRequestOutput` carries pipeline and diffusion results through the common output path. |
+
+## Configuration and runtime resolution
+
+The control plane has five conceptual layers. Authoring inputs are resolved
+once into a complete, transport-safe configuration before runtime processes are
+started. This keeps stage topology and model capabilities distinct from
+deployment placement and from process-local engine objects.
+
+The figure shows only the primary hand-off object for each layer; the detailed
+inputs, fields, and ownership rules are described below.
+
+```mermaid
+flowchart TB
+    layer1["Layer 1 · Authoring inputs<br/>PipelineConfig + DeployConfig"]
+    layer2["Layer 2 · Resolve once<br/>StageConfigFactory.create_from_model()<br/>VllmOmniConfig.from_pipeline_config()"]
+    layer3["Layer 3 · Transport-safe control plane<br/>VllmOmniConfig"]
+    layer4["Layer 4 · Runtime launch planning<br/>StageRuntime"]
+    layer5["Layer 5 · Engine materialization<br/>VllmConfig / OmniDiffusionConfig"]
+
+    layer1 --> layer2 --> layer3 --> layer4 --> layer5
+```
+
+The single resolution boundary is implemented by
+`StageConfigFactory.create_from_model()` and
+`VllmOmniConfig.from_pipeline_config()` in
+[`vllm_omni/config`](https://github.com/vllm-project/vllm-omni/tree/main/vllm_omni/config).
+The legacy `stage_args` YAML path remains only for models that have not yet
+migrated to `PipelineConfig` and `DeployConfig`.
+
+In the typed path, each stage configuration derives from
+`BaseVllmOmniStageConfig` and is specialized as
+`VllmOmniARStageConfig`, `VllmOmniGenerationStageConfig`, or
+`VllmOmniDiffusionStageConfig`. The request and engine-spec fields belong to the
+control-plane boundary; the current implementation stores their equivalent
+projections in the structured stage configuration and materializes
+backend-specific engine objects during stage initialization.
+
+The important ownership rules are:
+
+1. `PipelineConfig` is the source of truth for stage topology, execution type,
+   model capabilities, and stage relationships.
+2. `DeployConfig` describes placement, replicas, devices, connectors, and
+   deploy-time defaults; it does not redefine the model graph.
+3. CLI and Python overrides are applied at the resolution boundary, with
+   per-stage overrides taking precedence over global values where supported.
+4. `StageRuntime` owns launch planning and replica lifecycle. `ReplicaInitPlan`
+   is runtime-private state, not a user configuration object.
+5. `VllmConfig` and the enriched `OmniDiffusionConfig` are materialized in the
+   process that owns the corresponding engine.
 
 ## Main features
 
-vLLM-Omni aims to be fast, flexible, and easy to use with the following features:
+The feature surface is grouped to match the
+[feature design documents](index.md#feature-design-documents). This page
+summarizes each feature's architectural role; the linked design document is
+the source for configuration, compatibility, and implementation details.
 
-### Performance and Acceleration
+### Runtime and stage execution
 
-The framework achieves high performance through several optimization techniques:
+* **Disaggregated inference:** Logical stages can run in separate processes,
+  devices, or nodes while the orchestrator preserves their declared
+  relationships. `OmniConnector` implementations transfer stage data and
+  control-plane metadata. See [Disaggregated Inference](feature/disaggregated_inference.md).
+* **Asynchronous stage and output execution:** [Async Chunk](feature/async_chunk.md)
+  forwards partial stage outputs as they become available. [Async Diffusion
+  Output](feature/async_diffusion_output.md) overlaps device-to-host output
+  packing with the next diffusion request, while [Async Omni Output
+  Materialization](feature/omni_async_output_materialization.md) moves
+  CPU-side payload construction off the AR decode critical path.
+* **Automatic prefix caching:** [Automatic Prefix Caching in Omni Models](feature/prefix_caching.md)
+  reuses KV-cache-aligned stage outputs and multimodal tensors for requests
+  with common prefixes.
 
-* **Efficient AR Support:** Leverages efficient KV cache management inherited from vLLM.
-* **Pipelined Execution:** Uses pipelined stage execution overlapping to ensure high throughput.
-* **Full Disaggregation:** Relies on the OmniConnector and dynamic resource allocation across stages.
-* **Diffusion Acceleration:** Includes integrated support for diffusion acceleration. This is managed by the acceleration layer, which handles:
-    * **Cache:** Includes DBCache, TeaCache and third-party integration(e.g., [cache-dit](https://github.com/vipshop/cache-dit)).
-    * **Parallelism:** Supports TP, CP, USP, and CFG.
-    * **Attention:** Provides an interface for third-party integration (e.g., FA3, SAGE, MindIE-SD).
-    * **Quantization:** Supports various quantization implementations including FP8 and AWQ.
-    * **FusedOps:** Allows for custom and third-party integration.
+### Communication
 
-### Classifier-Free Guidance (CFG) Companion Flow
+* **OmniConnector transport:** The connector contract carries tensors, KV-cache
+  data, and transport metadata across stage boundaries. The available
+  implementations cover shared memory and multi-node Mooncake, Mori, and
+  Yuanrong transports; see [Disaggregated Inference](feature/disaggregated_inference.md)
+  for the connector choices and configuration model.
 
-vLLM-Omni natively models Classifier-Free Guidance (CFG) across disaggregated multi-stage setups via a "companion request" paradigm, eliminating redundant textual/multimodal context computation boundaries:
-1. **Prompt Expansion:** In the initial autoregressive (AR) stage, a customized `prompt_expand_func` hook intercepts incoming generation prompts and pairs them directly with negative companion prompts (e.g., a default negative prompt) on the fly, tagging the secondary prompt with a specific internal role (`cfg_text`).
-2. **Synchronized KV Cache Transfer:** The AR stage evaluates both the primary and companion sequence batches concurrently. The `OmniConnector` captures these specific structural dependencies and reliably passes the positive and negative outcome KV caches seamlessly across stage boundaries via shared memory or network protocols.
-3. **KV Cache Collection & Injection:** Upon reaching the downstream Diffusion (DiT) Engine, an assigned `cfg_kv_collect_func` automatically intercepts the mapped companion caches (`cfg_text_past_key_values`). These auxiliary dependencies are natively gathered and seamlessly bound to the primary generation sequence variables, enabling the DiT Engine to cleanly implement cross-attention CFG guidance over accurate conditioning and unconditioning structures in parallel.
+### Diffusion acceleration
 
-### Flexibility and Usability
+* **Request and step batching:** [Diffusion Continuous Batching](feature/diffusion_continuous_batching.md)
+  defines request-batch and step-batch execution, scheduler admission, and the
+  common streaming output path.
+* **Composable parallelism:** Diffusion stages can combine [CFG-Parallel](feature/cfg_parallel.md),
+  [Expert Parallel](feature/expert_parallel.md), [HSDP](feature/hsdp.md),
+  [Pipeline Parallel](feature/pipeline_parallel.md), [Sequence Parallel](feature/sequence_parallel.md),
+  [Tensor Parallel](feature/tensor_parallel.md), and [VAE Patch Parallelism](feature/vae_parallel.md)
+  according to the pipeline and hardware topology.
+* **Attention and cache acceleration:** [Skip-Softmax](feature/skip_softmax.md),
+  [Cache-DiT](feature/cache_dit.md), and [TeaCache](feature/teacache.md)
+  provide backend and denoising-step optimizations without changing the
+  stage contract.
+* **Quantization and memory efficiency:** [Quantization](feature/quantization.md)
+  resolves per-pipeline or per-component quantization configurations, while
+  [Distributed Layerwise Offload](feature/distributed_layerwise_offload.md)
+  streams diffusion blocks from host memory within the existing parallel
+  topology.
 
-vLLM-Omni is designed to be flexible and straightforward for users:
+## Interfaces
 
-* **Heterogeneous Pipeline Abstraction:** Manages complex model workflows effectively.
-* **Hugging Face Integration:** Offers seamless integration with popular Hugging Face models.
-* **Distributed Inference:** Supports tensor, pipeline, data, and expert parallelism.
-* **Streaming Outputs:** Supports streaming outputs.
-* **Unified API:** Provides a consistent and unified API interface compatible with vLLM.
-* **OpenAI-compatible API Server:** Includes a FastAPI-based server for online serving that is compatible with the OpenAI API.
-
-# Interface design
-
-If you use vLLM, then you know how to use vLLM-Omni from Day 0:
-
-<p align="center">
-  <picture>
-    <source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/vllm-omni-user-interface.png">
-    <img alt="vLLM-Omni interface design" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/vllm-omni-user-interface.png" width=55%>
-  </picture>
-</p>
-
-Taking **Qwen3-Omni** as an example:
-
-## Offline Inference
-The **Omni** class provides a Python interface for offline batched inference. Users initialize the Omni class with a Hugging Face model name and use the generate method, passing inputs that include both text prompts and multi-modal data:
+The public interfaces map onto the same engine and stage boundaries:
 
+```mermaid
+flowchart LR
+    offline["Offline Python<br/>Omni.generate()"] --> engine["AsyncOmniEngine"]
+    online["OpenAI-compatible serving<br/>vllm serve ... --omni"] --> engine
+    duplex["Experimental duplex WebSocket<br/>/v1/duplex or realtime duplex"] --> engine
+    engine --> stages["Configured AR and diffusion stages"]
+    stages --> result["Streaming or final multimodal output"]
 ```
-# Create an omni runtime with HF model name.
+
+### Offline inference
+
+The **Omni** class provides a Python interface for offline batched inference:
+
+```python
 from vllm_omni.entrypoints.omni import Omni
 
 omni = Omni(model="Qwen/Qwen3-Omni-30B-A3B-Instruct")
 
-# Example prompts.
-om_inputs = {"prompt": prompt,
-             "multi_modal_data": {
-                 "video": video_frames,
-                 "audio": audio_signal,
-             }}
+om_inputs = {
+    "prompt": prompt,
+    "multi_modal_data": {
+        "video": video_frames,
+        "audio": audio_signal,
+    },
+}
 
-# Generate texts and audio from the multi-modality inputs.
 outputs = omni.generate(om_inputs, sampling_params_list)
 ```
 
-## Online Serving
-Similar to vLLM, vLLM-Omni also provides a FastAPI-based server for online serving. Users can launch the server using the vllm serve command with the `--omni` flag:
+### Online serving
 
-```
+The OpenAI-compatible server uses the same stage configuration and engine
+boundaries:
+
+```bash
 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
 ```
 
-Users can send requests to the server using curl:
+For example, a Qwen3-Omni chat request can contain text, image, audio, or
+video content and a `sampling_params_list` for its configured stages. See the
+[Qwen3-Omni serving example](../user_guide/examples/online_serving/qwen3_omni.md)
+and the [examples](https://github.com/vllm-project/vllm-omni/tree/main/examples)
+for complete requests.
 
-```
-# prepare user content
-user_content='[
-        {
-          "type": "video_url",
-          "video_url": {
-            "url": "'"$SAMPLE_VIDEO_URL"'"
-          }
-        },
-        {
-          "type": "text",
-          "text": "Why is this video funny?"
-        }
-      ]'
-    sampling_params_list='[
-      '"$thinker_sampling_params"',
-      '"$talker_sampling_params"',
-      '"$code2wav_sampling_params"'
-    ]'
-    mm_processor_kwargs="{}"
-
-# send the request
-curl -sS -X POST http://localhost:8091/v1/chat/completions \
-    -H "Content-Type: application/json" \
-    -d @- <<EOF
-{
-  "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
-  "sampling_params_list": $sampling_params_list,
-  "mm_processor_kwargs": $mm_processor_kwargs,
-  "messages": [
-    {
-      "role": "system",
-      "content": [
-        {
-          "type": "text",
-          "text": "You are Qwen, a virtual human developed by the Qwen Team, Alibaba Group, capable of perceiving auditory and visual inputs, as well as generating text and speech."
-        }
-      ]
-    },
-    {
-      "role": "user",
-      "content": $user_content
-    }
-  ]
-}
-```
-
-For more usages, please refer to [examples](https://github.com/vllm-project/vllm-omni/tree/main/examples).
+Some pipelines expose additional OpenAI-compatible endpoints, such as joint
+video/audio generation. Endpoint support remains model-specific; consult the
+relevant model guide before assuming that every OpenAI route applies to every
+pipeline.

docs/user_guide/diffusion/attention_backends.md

@@ -17,7 +17,7 @@ The full set of backends and their platform defaults is in the **Backend Options
 | `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. |
@@ -84,7 +84,7 @@ A backend that needs configuration exposes it as a typed field on the spec:
 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(
@@ -99,6 +99,27 @@ 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)

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

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()

recipes/MiniMaxAI/MiniMax-H3.md

@@ -574,10 +574,43 @@ 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 |
 |-----------|-------------------|-------|
+| `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 |

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_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
@@ -91,8 +93,8 @@ class TestCacheDiTBackend:
 
     @patch("vllm_omni.diffusion.cache.cachedit.backend.BlockAdapter")
     @patch("vllm_omni.diffusion.cache.cachedit.backend.cache_dit")
-    def test_enable_and_refresh_every_declared_dit(self, mock_cache_dit, mock_block_adapter):
-        """Combined pipelines enable and refresh Cache-DiT for every 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"]
@@ -102,18 +104,67 @@ class TestCacheDiTBackend:
             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 len(backend._refresh_funcs) == 2
         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")
@@ -246,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):

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

@@ -13,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 (

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/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]
 
@@ -136,6 +139,30 @@ def test_task_type_roundtrip():
     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/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_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/helpers/runtime.py

@@ -1252,6 +1252,28 @@ class OpenAIClientHandler:
         )
         return [resp]
 
+    def send_batched_chat_completions_http_request(
+        self,
+        request_config: dict[str, Any],
+        *,
+        err_code: int | tuple[int, ...] | list[int] | None = None,
+        err_message: str | tuple[str, ...] | list[str] | None = None,
+    ) -> list[HttpResponse]:
+        """Post to batched chat completions."""
+        cfg = _merge_http_expectation_kwargs(
+            request_config,
+            err_code=err_code,
+            err_message=err_message,
+        )
+        r = self._post_json_endpoint("/v1/chat/completions/batch", cfg, default_timeout=120.0)
+        resp = self._http_response_from_requests(r)
+        assert_http_error(
+            resp,
+            err_code=cfg.get("err_code"),
+            err_message=cfg.get("err_message"),
+        )
+        return [resp]
+
     def send_completions_http_request(
         self,
         request_config: dict[str, Any],

tests/inputs/test_diffusion_sampling_params.py

@@ -25,6 +25,23 @@ def test_from_params_returns_existing_omni_params() -> None:
     assert converted is params
 
 
+@pytest.mark.parametrize("quality", ["lossless", "high"])
+def test_quality_accepts_supported_request_levels(quality: str) -> None:
+    params = OmniDiffusionSamplingParams(quality=quality)
+
+    assert params.quality == quality
+
+
+def test_quality_defaults_to_model_owned_policy() -> None:
+    params = OmniDiffusionSamplingParams()
+    assert params.quality is None
+
+
+def test_quality_rejects_unsupported_request_level() -> None:
+    with pytest.raises(ValueError, match="quality must be one of"):
+        OmniDiffusionSamplingParams(quality="medium")
+
+
 def test_from_params_converts_sampling_params_seed_and_known_extra_args() -> None:
     params = _sampling_params(
         seed=1234,
@@ -33,6 +50,7 @@ def test_from_params_converts_sampling_params_seed_and_known_extra_args() -> Non
             "height": 1024,
             "width": 768,
             "guidance_scale": 7.5,
+            "quality": "high",
         },
     )
 
@@ -43,6 +61,7 @@ def test_from_params_converts_sampling_params_seed_and_known_extra_args() -> Non
     assert converted.height == 1024
     assert converted.width == 768
     assert converted.guidance_scale == 7.5
+    assert converted.quality == "high"
     assert converted.extra_args == {}
 
 

vllm_omni/config/endpoint_policy.py

@@ -21,7 +21,6 @@ class RouteTarget(NamedTuple):
 class OmniServingCapability(Enum):
     """Serving capabilities that pipelines can shut down."""
 
-    CHAT_COMPLETIONS_BATCH = RouteTarget("/v1/chat/completions/batch", frozenset({"POST"}))
     COMPLETIONS = RouteTarget("/v1/completions", frozenset({"POST"}))
 
     @property
@@ -39,16 +38,6 @@ class EndpointRestriction:
     reason: str
 
 
-# Routes that are not supported for any model, but are supported in vLLM.
-# This is only temporary to avoid 500s for batched chat completions.
-UNSUPPORTED_ROUTES: tuple[EndpointRestriction, ...] = (
-    EndpointRestriction(
-        OmniServingCapability.CHAT_COMPLETIONS_BATCH,
-        "Batched chat completions are not yet supported by vLLM Omni.",
-    ),
-)
-
-
 def build_rejection_handler(reason: str):
     """Build a rejection handler for a given endpoint for the provided reason."""
 
@@ -71,11 +60,7 @@ def shutdown_unsupported_routes(
     """
     from vllm_omni.entrypoints.openai.api_server import _remove_route_from_app
 
-    # Generally these should not overlap since there is no point. If they do,
-    # we use the reason message in UNSUPPORTED_ROUTES, for consistent error messages.
-    restricted_endpoints = (*endpoint_restrictions, *UNSUPPORTED_ROUTES)
-
-    for end_restrict in restricted_endpoints:
+    for end_restrict in endpoint_restrictions:
         capability = end_restrict.capability
         # Remove the route from the app
         _remove_route_from_app(app, capability.path, capability.methods)

vllm_omni/config/omni_config.py

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

vllm_omni/config/stage_config.py

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

vllm_omni/diffusion/attention/backends/flash_attn.py

@@ -15,6 +15,7 @@ from vllm_omni.diffusion.attention.backends.sdpa import _maybe_reshape_attn_mask
 from vllm_omni.diffusion.attention.backends.utils.piecewise_attn import (
     piecewise_attn,
 )
+from vllm_omni.diffusion.config import get_current_diffusion_config_or_none
 
 logger = init_logger(__name__)
 
@@ -71,9 +72,20 @@ class FlashAttentionImpl(AttentionImpl):
         self.causal = causal
         self.softmax_scale = softmax_scale
         self.qkv_layout = qkv_layout
+        cfg = get_current_diffusion_config_or_none()
+        self.fa_deterministic = bool(getattr(cfg, "fa_deterministic", False)) if cfg is not None else False
         if backend_kwargs:
             logger.warning("FlashAttentionImpl ignoring backend_kwargs: %s", list(backend_kwargs.keys()))
 
+    def _warn_fa_deterministic_non_dense(self, path: str) -> None:
+        if not self.fa_deterministic:
+            return
+        logger.warning_once(
+            "fa_deterministic=True is ignored on the %s FlashAttention path; "
+            "only the dense flash_attn_func path passes deterministic=True.",
+            path,
+        )
+
     @staticmethod
     def _unwrap_flash_output(out: torch.Tensor | tuple[torch.Tensor, ...]) -> torch.Tensor:
         # FA3 may return (out, lse), FA2 returns out
@@ -250,6 +262,7 @@ class FlashAttentionImpl(AttentionImpl):
 
         # Try piecewise attention
         if full_attn_spans is not None:
+            self._warn_fa_deterministic_non_dense("piecewise")
             logger.debug("Using piecewise Flash Attention for mixed causal/full mask")
             if flash_attn_func is not None:
                 attn_func = partial(
@@ -280,6 +293,7 @@ class FlashAttentionImpl(AttentionImpl):
             if len(present_packed_keys) != len(packed_keys):
                 missing = sorted(set(packed_keys) - set(present_packed_keys))
                 raise ValueError(f"Incomplete packed FlashAttention metadata; missing {missing}")
+            self._warn_fa_deterministic_non_dense("packed-varlen")
             return self._forward_varlen_packed(
                 query,
                 key,
@@ -291,6 +305,7 @@ class FlashAttentionImpl(AttentionImpl):
             )
 
         if attention_mask is not None and torch.any(~attention_mask):
+            self._warn_fa_deterministic_non_dense("masked-varlen")
             return self._forward_varlen_masked(
                 query,
                 key,
@@ -299,15 +314,16 @@ class FlashAttentionImpl(AttentionImpl):
             )
 
         if flash_attn_func is not None:
-            out = flash_attn_func(
-                query,
-                key,
-                value,
-                causal=self.causal,
-                softmax_scale=self.softmax_scale,
-            )
+            fa_kwargs = {
+                "causal": self.causal,
+                "softmax_scale": self.softmax_scale,
+            }
+            if self.fa_deterministic:
+                fa_kwargs["deterministic"] = True
+            out = flash_attn_func(query, key, value, **fa_kwargs)
             return self._unwrap_flash_output(out)
 
+        self._warn_fa_deterministic_non_dense("dense-varlen-fallback")
         return self._forward_varlen_dense(
             query,
             key,

vllm_omni/diffusion/attention/backends/flashinfer_attn.py

@@ -1,7 +1,10 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+from dataclasses import dataclass
+
 import torch
+from packaging.version import InvalidVersion, Version
 from vllm.logger import init_logger
 
 from vllm_omni.diffusion.attention.backends.abstract import (
@@ -13,7 +16,8 @@ from vllm_omni.diffusion.attention.backends.abstract import (
 logger = init_logger(__name__)
 
 try:
-    from flashinfer.prefill import single_prefill_with_kv_cache
+    import flashinfer
+    from flashinfer.prefill import BatchPrefillWithRaggedKVCacheWrapper
 
     HAS_FLASHINFER = True
 except Exception as e:
@@ -29,10 +33,8 @@ class FlashInferAttentionBackend(AttentionBackend):
 
     @classmethod
     def supports_attention_mask(cls) -> bool:
-        # FlashInfer single_prefill_with_kv_cache accepts a ``custom_mask`` kwarg
-        # (2D boolean, ``True`` = keep) for non-causal attention. See reviewer
-        # comment on #3079 and the flashinfer docs:
-        #   https://docs.flashinfer.ai/generated/flashinfer.prefill.single_prefill_with_kv_cache.html
+        # Non-CuTe batch-prefill backends accept boolean custom masks. CuTe-DSL
+        # falls back to SDPA for nontrivial masks.
         return True
 
     @staticmethod
@@ -52,6 +54,27 @@ class FlashInferAttentionBackend(AttentionBackend):
 
 
 class FlashInferAttentionImpl(AttentionImpl):
+    _QK_DTYPES = {torch.float16, torch.bfloat16}
+    _VO_DTYPES = {torch.float16, torch.bfloat16, torch.float8_e4m3fn}
+
+    @dataclass(frozen=True)
+    class _WrapperPlanKey:
+        batch_size: int
+        qo_len: int
+        kv_len: int
+        num_q_heads: int
+        num_kv_heads: int
+        head_dim_qk: int
+        head_dim_k: int
+        head_dim_vo: int
+        q_dtype: torch.dtype
+        k_dtype: torch.dtype
+        v_dtype: torch.dtype
+        o_dtype: torch.dtype
+        causal: bool
+        softmax_scale: float
+        has_custom_mask: bool
+
     def __init__(
         self,
         num_heads: int,
@@ -60,10 +83,92 @@ class FlashInferAttentionImpl(AttentionImpl):
         causal: bool = False,
         num_kv_heads: int | None = None,
         prefix: str = "",
+        backend_kwargs: dict | None = None,
         **extra_impl_args,
     ) -> None:
         self.causal = causal
         self.softmax_scale = softmax_scale
+        self.device = torch.device("cuda", torch.accelerator.current_device_index())
+        backend_kwargs = backend_kwargs or {}
+        quant = backend_kwargs.get("quant") or {}
+        self.dtype_qk = self._check_dtype(quant.get("dtype_qk"), "dtype_qk", self._QK_DTYPES)
+        self.dtype_vo = self._check_dtype(quant.get("dtype_vo"), "dtype_vo", self._VO_DTYPES)
+        requested_backend = quant.get("flashinfer_backend", "auto")
+
+        if not HAS_FLASHINFER:
+            raise ImportError("FLASHINFER_ATTN backend requires flashinfer")
+
+        self._check_flashinfer_version()
+
+        self.flashinfer_backend = self._select_backend(requested_backend, self.device)
+        workspace_size = 0 if self.flashinfer_backend == "cute-dsl" else 128 * 1024 * 1024
+        self._workspace = torch.empty(
+            workspace_size,
+            device=self.device,
+            dtype=torch.uint8,
+        )
+        self._wrapper = BatchPrefillWithRaggedKVCacheWrapper(
+            self._workspace,
+            kv_layout="NHD",
+            backend=self.flashinfer_backend,
+        )
+        self._qo_indptr: torch.Tensor | None = None
+        self._kv_indptr: torch.Tensor | None = None
+        self._plan_key: FlashInferAttentionImpl._WrapperPlanKey | None = None
+
+        if self.dtype_qk is not None or self.dtype_vo is not None:
+            logger.info_once(
+                "FLASHINFER_ATTN dtype override: Q/K=%s, V=%s.",
+                self.dtype_qk,
+                self.dtype_vo,
+            )
+
+        logger.info_once(
+            "FLASHINFER_ATTN initialized backend=%s on %s.",
+            self.flashinfer_backend,
+            self.device,
+        )
+
+    def _check_flashinfer_version(self) -> None:
+        if self.dtype_qk == self.dtype_vo:
+            return
+        try:
+            flashinfer_version = Version(flashinfer.__version__)
+        except (AttributeError, InvalidVersion):
+            return
+        if flashinfer_version <= Version("0.6.15"):
+            raise RuntimeError(
+                f"FlashInfer {flashinfer_version} is too old for reliable mixed "
+                f"QK/V dtype attention (Q/K={self.dtype_qk}, V={self.dtype_vo}); "
+                "install flashinfer >= 0.6.16rc1."
+            )
+
+    @staticmethod
+    def _select_backend(requested_backend: str, device: torch.device) -> str:
+        if requested_backend != "auto":
+            return requested_backend
+        major, _minor = torch.cuda.get_device_capability(device)
+        if major >= 10:
+            return "cute-dsl"
+        if major >= 9:
+            return "fa3"
+        return "fa2"
+
+    @classmethod
+    def _check_dtype(
+        cls,
+        dtype: torch.dtype | str | None,
+        option_name: str,
+        allowed: set[torch.dtype],
+    ) -> torch.dtype | None:
+        if dtype is None:
+            return None
+        if isinstance(dtype, str):
+            dtype = torch.float8_e4m3fn if dtype == "fp8_e4m3" else getattr(torch, dtype)
+        if dtype not in allowed:
+            choices = ", ".join(sorted(str(item) for item in allowed))
+            raise ValueError(f"Unsupported {option_name}={dtype}; expected one of: {choices}")
+        return dtype
 
     @staticmethod
     def _pack_mask_for_flashinfer(
@@ -128,6 +233,126 @@ class FlashInferAttentionImpl(AttentionImpl):
         )
         return impl.forward_cuda(query, key, value, attn_metadata)
 
+    @torch.compiler.disable
+    def _plan_wrapper(
+        self,
+        key: _WrapperPlanKey,
+        flat_mask: torch.Tensor | None,
+    ) -> None:
+        self._wrapper.plan(
+            self._qo_indptr,
+            self._kv_indptr,
+            key.num_q_heads,
+            key.num_kv_heads,
+            key.head_dim_qk,
+            head_dim_vo=key.head_dim_vo,
+            custom_mask=flat_mask,
+            causal=key.causal,
+            sm_scale=key.softmax_scale,
+            q_data_type=key.q_dtype,
+            # FlashInfer versions newer than 0.6.15 can dispatch K and V with
+            # independent runtime dtypes even though plan() names this K/V.
+            kv_data_type=key.k_dtype,
+            o_data_type=key.o_dtype,
+        )
+
+    @torch.compiler.disable
+    def _make_indptr(self, batch_size: int, sequence_length: int) -> torch.Tensor:
+        return (
+            torch.arange(
+                batch_size + 1,
+                device=self.device,
+                dtype=torch.int32,
+            )
+            * sequence_length
+        )
+
+    def _ensure_plan(
+        self,
+        key: _WrapperPlanKey,
+        flat_mask: torch.Tensor | None,
+    ) -> None:
+        key_changed = key != self._plan_key
+        if not key_changed and flat_mask is None:
+            return
+
+        if key_changed:
+            self._qo_indptr = self._make_indptr(key.batch_size, key.qo_len)
+            self._kv_indptr = self._make_indptr(key.batch_size, key.kv_len)
+
+        # A custom mask's contents may change without its shape changing, so
+        # copy/re-plan it for every masked invocation.
+        self._plan_wrapper(key, flat_mask)
+        self._plan_key = key
+
+    def _run_batch_prefill(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        custom_mask: torch.Tensor | None,
+    ) -> torch.Tensor:
+        if query.device != self.device or key.device != self.device or value.device != self.device:
+            raise ValueError(
+                "FLASHINFER_ATTN inputs must remain on the layer initialization "
+                f"device {self.device}; got Q={query.device}, K={key.device}, "
+                f"V={value.device}"
+            )
+
+        batch_size, qo_len, num_q_heads, head_dim_qk = query.shape
+        kv_len = key.shape[1]
+        num_kv_heads = key.shape[2]
+        head_dim_k = key.shape[3]
+        head_dim_vo = value.shape[3]
+
+        q = query.reshape(batch_size * qo_len, num_q_heads, head_dim_qk)
+        k = key.reshape(batch_size * kv_len, num_kv_heads, head_dim_k)
+        v = value.reshape(batch_size * kv_len, num_kv_heads, head_dim_vo)
+        if self.dtype_qk is not None:
+            q = q.to(self.dtype_qk)
+            k = k.to(self.dtype_qk)
+        if self.dtype_vo is not None:
+            v = v.to(self.dtype_vo)
+
+        flat_mask = None
+        if custom_mask is not None:
+            if custom_mask.dim() == 2:
+                custom_mask = custom_mask.unsqueeze(0).expand(batch_size, -1, -1)
+            flat_mask = custom_mask.contiguous().view(-1)
+
+        self._ensure_plan(
+            FlashInferAttentionImpl._WrapperPlanKey(
+                batch_size=batch_size,
+                qo_len=qo_len,
+                kv_len=kv_len,
+                num_q_heads=num_q_heads,
+                num_kv_heads=num_kv_heads,
+                head_dim_qk=head_dim_qk,
+                head_dim_k=head_dim_k,
+                head_dim_vo=head_dim_vo,
+                q_dtype=q.dtype,
+                k_dtype=k.dtype,
+                v_dtype=v.dtype,
+                o_dtype=query.dtype,
+                causal=self.causal,
+                softmax_scale=self.softmax_scale,
+                has_custom_mask=flat_mask is not None,
+            ),
+            flat_mask,
+        )
+        out = self._run_wrapper(q, k, v)
+        out = out.reshape(batch_size, qo_len, num_q_heads, head_dim_vo)
+        return out.to(query.dtype) if out.dtype != query.dtype else out
+
+    @torch.compiler.disable
+    def _run_wrapper(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+    ) -> torch.Tensor:
+        return self._wrapper.run(query, key, value)
+
     def forward_cuda(
         self,
         query: torch.Tensor,
@@ -135,18 +360,7 @@ class FlashInferAttentionImpl(AttentionImpl):
         value: torch.Tensor,
         attn_metadata: AttentionMetadata | None = None,
     ) -> torch.Tensor:
-        if not HAS_FLASHINFER:
-            raise ImportError(
-                "FLASHINFER_ATTN backend requires flashinfer. "
-                "Install it or set DIFFUSION_ATTENTION_BACKEND to another backend."
-            )
-
-        # Try the custom_mask path; if the mask can't be packed into the
-        # (qo_len, kv_len) layout FlashInfer expects, fall back to SDPA
-        # instead of risking an illegal-memory-access crash in the kernel.
-        # Input layout is (B, S, H, D); FlashInfer dense prefill takes (S, H, D).
         batch_size = query.shape[0]
-
         custom_mask = None
         if attn_metadata is not None and attn_metadata.attn_mask is not None:
             try:
@@ -156,28 +370,14 @@ class FlashInferAttentionImpl(AttentionImpl):
                     qo_len=query.shape[1],
                     kv_len=key.shape[1],
                 )
-            except ValueError as e:
-                logger.debug("Falling back to SDPA for mask path: %s", e)
+            except ValueError as error:
+                logger.debug("Falling back to SDPA for mask path: %s", error)
                 return self._sdpa_fallback(query, key, value, attn_metadata)
-            # FlashInfer cannot combine causal masking with a custom_mask; rather
-            # than silently dropping the explicit mask (diverging from SDPA), let
-            # SDPA handle the causal+mask case correctly.
             if custom_mask is not None and self.causal:
-                logger.debug("causal=True with explicit attn_mask; deferring to SDPA")
                 return self._sdpa_fallback(query, key, value, attn_metadata)
 
-        outputs = []
-        for b in range(batch_size):
-            kwargs: dict = {
-                "sm_scale": self.softmax_scale,
-                "causal": self.causal,
-                "return_lse": False,
-            }
-            if custom_mask is not None:
-                kwargs["custom_mask"] = custom_mask if custom_mask.dim() == 2 else custom_mask[b]
-            out = single_prefill_with_kv_cache(query[b], key[b], value[b], **kwargs)
-            outputs.append(out)
+        if custom_mask is not None and self.flashinfer_backend == "cute-dsl":
+            logger.debug("CuTe DSL does not support custom masks; deferring to SDPA")
+            return self._sdpa_fallback(query, key, value, attn_metadata)
 
-        if batch_size == 1:
-            return outputs[0].unsqueeze(0)
-        return torch.stack(outputs, dim=0)
+        return self._run_batch_prefill(query, key, value, custom_mask)

vllm_omni/diffusion/attention/backends/ring/ring_kernels.py

@@ -83,6 +83,19 @@ def pytorch_attn_forward(
     k = k.transpose(1, 2)
     v = v.transpose(1, 2)
 
+    # GQA/MQA support. Ring attention only ever *communicates* K/V with H_kv
+    # heads -- that is precisely the benefit of GQA -- so expanding here does
+    # not change communication volume. The expansion is needed because the aten
+    # SDPA ops below are required for the LSE output that ring accumulation
+    # consumes, and they do not accept `enable_gqa`.
+    if k.shape[1] != q.shape[1]:
+        num_heads, num_kv_heads = q.shape[1], k.shape[1]
+        if num_heads % num_kv_heads != 0:
+            raise ValueError(f"num_heads ({num_heads}) must be divisible by num_kv_heads ({num_kv_heads}) for GQA/MQA.")
+        n_rep = num_heads // num_kv_heads
+        k = k.repeat_interleave(n_rep, dim=1)
+        v = v.repeat_interleave(n_rep, dim=1)
+
     if op_type == "flash":
         out, lse = _scaled_dot_product_flash_attention(
             q,

vllm_omni/diffusion/attention/parallel/ulysses.py

@@ -204,6 +204,25 @@ class UlyssesParallelAttention:
         attn_metadata: AttentionMetadata | None,
     ):
         mode = get_ulysses_mode(default="strict")
+        ulysses_world_size = self._sp_group.ulysses_world_size
+
+        # advanced_uaa pads non-divisible head counts before the Ulysses
+        # all-to-all. Padding K/V is not valid in hybrid Ulysses+Ring: the Ring
+        # GQA/MQA path would repeat the padded zero heads as real K/V heads.
+        # Reject this layout until K/V replication is performed before Ulysses.
+        if mode == "advanced_uaa" and self._sp_group.ring_world_size > 1:
+            for name, tensor in (("key", key), ("value", value)):
+                kv_head_cnt = int(tensor.shape[2])
+                if kv_head_cnt % ulysses_world_size != 0:
+                    raise ValueError(
+                        "ulysses_mode='advanced_uaa' with hybrid Ulysses+Ring "
+                        "does not support K/V head padding. "
+                        f"{name}_head_cnt={kv_head_cnt}, "
+                        f"ulysses_degree={ulysses_world_size}. "
+                        "Use ring_degree=1, choose a compatible ulysses_degree, "
+                        "or replicate K/V heads before Ulysses."
+                    )
+
         joint_tensor_query = joint_tensor_key = joint_tensor_value = None
         joint_strategy = "front"
         joint_len = 0
@@ -226,7 +245,6 @@ class UlyssesParallelAttention:
 
             # Slice joint_query for this Ulysses rank
             # joint_query is (B, S, H, D). We split H (dim 2).
-            ulysses_world_size = self._sp_group.ulysses_world_size
             ulysses_rank = self._sp_group.ulysses_rank
             joint_head_cnt = int(joint_tensor_query.shape[-2])
             joint_orig_head_cnt = joint_head_cnt

vllm_omni/diffusion/cache/cachedit/__init__.py

@@ -20,6 +20,10 @@ from vllm_omni.diffusion.cache.cachedit.model_specific import (
 from vllm_omni.diffusion.cache.cachedit.model_specific import (
     register_custom_dit_enablers as _register_custom_dit_enablers,
 )
+from vllm_omni.diffusion.cache.cachedit.runtime import (
+    CacheDiTRequestSpec,
+    RequestScopedCacheDiTRuntime,
+)
 
 _register_custom_dit_enablers()
 
@@ -29,6 +33,8 @@ __all__ = [
     "CacheDiTAdapterConfig",
     "CacheDiTBackend",
     "CacheDiTConfig",
+    "CacheDiTRequestSpec",
+    "RequestScopedCacheDiTRuntime",
     "SensenovaCachedAdapter",
     "cache_summary",
     "enable_cache_for_dit",

vllm_omni/diffusion/cache/cachedit/backend.py

@@ -6,6 +6,7 @@
 from __future__ import annotations
 
 from collections.abc import Callable
+from operator import attrgetter
 from typing import TYPE_CHECKING, Any, TypeAlias, cast
 
 import cache_dit
@@ -38,14 +39,27 @@ def _dit_module_names(pipeline: SupportsComponentDiscovery) -> tuple[str, ...]:
     names = getattr(pipeline, "_dit_modules", None)
     if not isinstance(names, (list, tuple)):
         names = ("transformer",)
-    return tuple(name for name in names if isinstance(name, str) and getattr(pipeline, name, None) is not None)
+
+    resolved_names = []
+    for name in names:
+        if not isinstance(name, str):
+            continue
+        try:
+            module = attrgetter(name)(pipeline)
+        except AttributeError:
+            continue
+        if module is not None:
+            resolved_names.append(name)
+    return tuple(resolved_names)
 
 
 def cache_summary(pipeline: SupportsComponentDiscovery, details: bool = True) -> None:
     """Log Cache-DiT statistics for every transformer on the pipeline."""
 
-    transformers = [getattr(pipeline, name) for name in _dit_module_names(pipeline)]
+    transformers = [attrgetter(name)(pipeline) for name in _dit_module_names(pipeline)]
     for transformer in transformers:
+        if not BlockAdapter.is_cached(transformer):
+            continue
         cache_dit.summary(transformer, details=details)
 
     if not transformers:
@@ -64,7 +78,7 @@ def _make_pipeline_transformer_getter(
     def get_pipeline_transformer(
         pipeline: SupportsComponentDiscovery,
     ) -> nn.Module:
-        return cast(nn.Module, getattr(pipeline, name))
+        return cast(nn.Module, attrgetter(name)(pipeline))
 
     return get_pipeline_transformer
 
@@ -85,19 +99,37 @@ def _build_cache_context_refresh(
         # Cache-DiT has no predefined SCM mask for these small step counts.
         scm_supported_steps = num_inference_steps >= 8 or num_inference_steps in (4, 6)
         if projected_config.scm_steps_mask_policy is None or not scm_supported_steps:
+            if projected_config.force_refresh_step_hint is not None:
+                # Cache-DiT's ``once`` policy clears the hint after it fires.
+                # Reapply it at every request boundary so a repeated request
+                # with the same configured policy receives the same behavior.
+                cache_dit.refresh_context(
+                    transformer,
+                    cache_config=DBCacheConfig().reset(
+                        num_inference_steps=num_inference_steps,
+                        force_refresh_step_hint=projected_config.force_refresh_step_hint,
+                        force_refresh_step_policy=projected_config.force_refresh_step_policy,
+                    ),
+                    verbose=verbose,
+                )
+                return
             cache_dit.refresh_context(transformer, num_inference_steps=num_inference_steps, verbose=verbose)
             return
 
+        refresh_config = DBCacheConfig().reset(
+            num_inference_steps=num_inference_steps,
+            steps_computation_mask=cache_dit.steps_mask(
+                mask_policy=projected_config.scm_steps_mask_policy,
+                total_steps=num_inference_steps,
+            ),
+            steps_computation_policy=projected_config.scm_steps_policy,
+        )
+        if projected_config.force_refresh_step_hint is not None:
+            refresh_config.force_refresh_step_hint = projected_config.force_refresh_step_hint
+            refresh_config.force_refresh_step_policy = projected_config.force_refresh_step_policy
         cache_dit.refresh_context(
             transformer,
-            cache_config=DBCacheConfig().reset(
-                num_inference_steps=num_inference_steps,
-                steps_computation_mask=cache_dit.steps_mask(
-                    mask_policy=projected_config.scm_steps_mask_policy,
-                    total_steps=num_inference_steps,
-                ),
-                steps_computation_policy=projected_config.scm_steps_policy,
-            ),
+            cache_config=refresh_config,
             verbose=verbose,
         )
 
@@ -213,19 +245,23 @@ class CacheDiTBackend(CacheBackend):
 
         super().__init__(config)
         self._refresh_funcs: list[RefreshCacheContextFunc] = []
+        self._cache_targets: list[Any] = []
 
     def enable(self, pipeline: SupportsComponentDiscovery) -> None:
         pipeline_name = type(pipeline).__name__
         custom_enabler = CUSTOM_DIT_ENABLERS.get(pipeline_name)
+        self._refresh_funcs = []
+        self._cache_targets = []
         if custom_enabler is not None:
             logger.info("Using custom cache-dit enabler for model: %s", pipeline_name)
             self._refresh_funcs = [custom_enabler(pipeline, self.config)]
+            self._cache_targets = [_default_get_pipeline_transformer(pipeline)]
         else:
-            self._refresh_funcs = []
             for name in _dit_module_names(pipeline):
                 get_transformer = _make_pipeline_transformer_getter(name)
                 block_adapter = _maybe_build_block_adapter(pipeline, get_transformer)
                 adapter_cls = _maybe_get_cached_adapter_cls(pipeline, get_transformer)
+                cache_target = get_transformer(pipeline) if block_adapter is None else block_adapter
                 self._refresh_funcs.append(
                     enable_cache_for_dit(
                         pipeline,
@@ -235,12 +271,32 @@ class CacheDiTBackend(CacheBackend):
                         get_transformer,
                     )
                 )
+                self._cache_targets.append(cache_target)
             if not self._refresh_funcs:
                 raise ValueError(f"Pipeline {pipeline_name} has no declared DiT modules for Cache-DiT")
 
         self.enabled = True
         logger.info("Cache-dit enabled successfully on %s", pipeline_name)
 
+    def disable(self, pipeline: SupportsComponentDiscovery) -> None:
+        """Remove Cache-DiT hooks so later requests use native forwards."""
+
+        if not self.enabled:
+            return
+
+        logger.info(
+            "Disabling cache-dit on %d DiT target(s) for %s",
+            len(self._cache_targets),
+            type(pipeline).__name__,
+        )
+        try:
+            for target in self._cache_targets:
+                cache_dit.disable_cache(target)
+        finally:
+            self._refresh_funcs = []
+            self._cache_targets = []
+            self.enabled = False
+
     def refresh(self, pipeline: SupportsComponentDiscovery, num_inference_steps: int, verbose: bool = True) -> None:
         if not self.enabled or not self._refresh_funcs:
             logger.warning("Cache-dit is not enabled. Cannot refresh cache context.")

vllm_omni/diffusion/cache/cachedit/runtime.py

@@ -0,0 +1,94 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Declarative request-scoped Cache-DiT lifecycle."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+from vllm_omni.diffusion.cache.cachedit.backend import CacheDiTBackend
+from vllm_omni.diffusion.data import DiffusionCacheConfig
+
+
+@dataclass(frozen=True)
+class CacheDiTRequestSpec:
+    """Desired Cache-DiT state for one request.
+
+    ``installation_key`` identifies configurations that require reinstalling
+    hooks when switching between them. Step-count changes only require a
+    context refresh.
+    """
+
+    installation_key: str
+    cache_config: DiffusionCacheConfig
+    num_inference_steps: int
+
+
+class RequestScopedCacheDiTRuntime:
+    """Reconcile declarative request state with installed Cache-DiT hooks."""
+
+    def __init__(self, pipeline: Any) -> None:
+        self._pipeline = pipeline
+        self._backend: CacheDiTBackend | None = None
+        self._installation_key: str | None = None
+
+    @property
+    def is_enabled(self) -> bool:
+        return self._backend is not None and self._backend.is_enabled()
+
+    def adopt(self, backend: CacheDiTBackend, *, installation_key: str) -> None:
+        """Assume ownership of an already-enabled startup backend."""
+
+        if self._backend is not None:
+            if self._backend is backend and self._installation_key == installation_key:
+                return
+            raise RuntimeError(
+                f"Cache-DiT is already installed with key {self._installation_key!r}; cannot adopt {installation_key!r}"
+            )
+        if not backend.is_enabled():
+            raise ValueError("Cannot adopt a Cache-DiT backend before it is enabled")
+
+        self._backend = backend
+        self._installation_key = installation_key
+
+    def prepare(self, spec: CacheDiTRequestSpec | None) -> None:
+        """Apply the minimal transition needed before the request denoise.
+
+        ``None`` disables installed hooks. An unchanged installation key only
+        refreshes request context; a changed key disables the previous hooks
+        before installing the requested profile.
+        """
+
+        desired_key = None if spec is None else spec.installation_key
+        if self._installation_key != desired_key:
+            self.disable()
+
+        if spec is None:
+            return
+
+        if self._backend is None:
+            backend = CacheDiTBackend(spec.cache_config)
+            backend.enable(self._pipeline)
+            self.adopt(backend, installation_key=spec.installation_key)
+
+        assert self._backend is not None
+        self._backend.refresh(
+            self._pipeline,
+            num_inference_steps=spec.num_inference_steps,
+        )
+
+    def disable(self) -> None:
+        try:
+            if self._backend is not None:
+                self._backend.disable(self._pipeline)
+        finally:
+            self._backend = None
+            self._installation_key = None
+
+
+__all__ = [
+    "CacheDiTRequestSpec",
+    "RequestScopedCacheDiTRuntime",
+]

vllm_omni/diffusion/data.py

@@ -613,6 +613,7 @@ class OmniDiffusionConfig:
 
     # Attention
     diffusion_attention_config: "AttentionConfig" = field(default_factory=lambda: AttentionConfig())
+    fa_deterministic: bool = False
 
     # Running mode
     # mode: ExecutionMode = ExecutionMode.INFERENCE

vllm_omni/diffusion/diffusion_engine.py

@@ -116,6 +116,16 @@ def _max_num_seqs(od_config: OmniDiffusionConfig) -> int:
         return 1
 
 
+def _uses_dlo_dp_concurrency(od_config: OmniDiffusionConfig) -> bool:
+    parallel_config = getattr(od_config, "parallel_config", None)
+    dp_size = getattr(parallel_config, "data_parallel_size", 1)
+    return (
+        dp_size > 1
+        and getattr(od_config, "enable_distributed_layerwise_offload", False)
+        and getattr(od_config, "dlo_use_allgather", True)
+    )
+
+
 def _move_tensor_tree_to_cpu(value: object) -> object:
     if isinstance(value, torch.Tensor):
         return value.cpu() if value.device.type != "cpu" else value
@@ -202,7 +212,7 @@ class DiffusionEngine:
             return DiffusionExecutionMode.STEP_BATCH
 
         self.supports_request_batch = supports_request_batch(od_config)
-        if not self.supports_request_batch and _max_num_seqs(od_config) > 1:
+        if not self.supports_request_batch and _max_num_seqs(od_config) > 1 and not _uses_dlo_dp_concurrency(od_config):
             raise ValueError(
                 f"{getattr(od_config, 'model_class_name', None)!r} does not support request-level batching. "
                 "Use max_num_seqs=1 for serial request execution, or choose a pipeline with "
@@ -233,17 +243,10 @@ class DiffusionEngine:
         # for distributed layerwise offload (which shards weights and
         # needs all ranks active simultaneously).  Ordinary DP with a
         # non-batch pipeline should not schedule multiple requests.
-        dp_size = 1
-        if getattr(self.od_config, "parallel_config", None) is not None:
-            dp_size = getattr(self.od_config.parallel_config, "data_parallel_size", 1)
-        dist_offload = getattr(self.od_config, "enable_distributed_layerwise_offload", False)
-        if dp_size > 1 and dist_offload and getattr(self.od_config, "dlo_use_allgather", True):
+        dp_size = getattr(getattr(self.od_config, "parallel_config", None), "data_parallel_size", 1)
+        if _uses_dlo_dp_concurrency(self.od_config):
             self.scheduler.max_num_running_reqs = dp_size
             self.dp_concurrent = True
-            # Set batch admission wait so concurrent requests accumulate
-            # before scheduling.  The scheduler reads this from od_config.
-            if getattr(self.od_config, "request_batch_max_wait_ms", 0) == 0:
-                self.od_config.request_batch_max_wait_ms = 500.0
             logger.info(
                 "dp_concurrent: max_num_running_reqs=%d, batch_wait=%sms",
                 dp_size,

vllm_omni/diffusion/executor/multiproc_executor.py

@@ -4,6 +4,7 @@ import concurrent.futures
 import json
 import multiprocessing as mp
 import multiprocessing.connection
+import os
 import queue
 import threading
 import time
@@ -21,6 +22,7 @@ from vllm.v1.engine.exceptions import EngineDeadError
 from vllm_omni.diffusion.data import SHUTDOWN_MESSAGE, AsyncDiffusionOutput, AsyncOutputKind, DiffusionOutput
 from vllm_omni.diffusion.executor.abstract import DiffusionExecutor
 from vllm_omni.diffusion.ipc import DIFFUSION_RPC_RESULT_ENVELOPE, unpack_diffusion_output_shm
+from vllm_omni.diffusion.sched.request_scheduler import build_request_batch_sampling_params_key
 from vllm_omni.diffusion.worker import WorkerProc
 
 if TYPE_CHECKING:
@@ -30,6 +32,21 @@ if TYPE_CHECKING:
 logger = init_logger(__name__)
 
 _DEQUEUE_TIMEOUT_S = 5.0
+_DLO_DP_WAVE_TIMEOUT_S = float(os.environ.get("VLLM_OMNI_DLO_DP_WAVE_TIMEOUT", 600.0))
+_WORKER_SHUTDOWN_GRACE_S = 15.0
+_WORKER_TERMINATE_GRACE_S = 5.0
+_RESULT_PUMP_JOIN_TIMEOUT_S = 2.0
+
+
+def _is_empty_dp_prompt(prompt: object) -> bool:
+    """Return whether a DP request has no usable text prompt."""
+    if prompt is None:
+        return True
+    if isinstance(prompt, str):
+        return not prompt
+    if isinstance(prompt, dict) and "prompt" in prompt:
+        return not prompt["prompt"]
+    return False
 
 
 @dataclass
@@ -52,14 +69,20 @@ class _ExecutorShutdownCleaner:
                 logger.warning("Failed to send shutdown signal: %s", exc)
 
         if self.processes:
+            join_deadline = time.monotonic() + _WORKER_SHUTDOWN_GRACE_S
             for proc in self.processes:
                 if not proc.is_alive():
                     continue
-                proc.join(5)
-                if proc.is_alive():
-                    logger.warning("Terminating diffusion worker %s after timeout", proc.name)
-                    proc.terminate()
-                    proc.join(5)
+                proc.join(max(0.0, join_deadline - time.monotonic()))
+
+            alive = [proc for proc in self.processes if proc.is_alive()]
+            for proc in alive:
+                logger.warning("Terminating diffusion worker %s after timeout", proc.name)
+                proc.terminate()
+
+            terminate_deadline = time.monotonic() + _WORKER_TERMINATE_GRACE_S
+            for proc in alive:
+                proc.join(max(0.0, terminate_deadline - time.monotonic()))
 
 
 class MultiprocDiffusionExecutor(DiffusionExecutor):
@@ -75,6 +98,7 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
         self._is_failed = False
         self._failure_callbacks: list[Callable[[], None]] = []
         self._result_mq: MessageQueue | None = None
+        self._result_mqs: list[MessageQueue] = []
         self._rpc_wave_id: int = 0
 
         num_workers = cast(int, self.od_config.num_gpus)
@@ -84,7 +108,7 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
         broadcast_handle = self._broadcast_mq.export_handle()
 
         # Launch workers
-        processes, result_handle = self._launch_workers(broadcast_handle, self.wake_events)
+        processes, result_handles = self._launch_workers(broadcast_handle, self.wake_events)
         self._processes = processes
 
         shutdown_cleaner = _ExecutorShutdownCleaner(
@@ -96,7 +120,8 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
         self._finalizer = weakref.finalize(self, shutdown_cleaner)
 
         try:
-            self._result_mq = self._init_result_queue(result_handle)
+            self._result_mqs = [self._init_result_queue(handle) for handle in result_handles]
+            self._result_mq = self._result_mqs[0]
         except Exception:
             self.shutdown()
             raise
@@ -112,8 +137,8 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
         self._futures_lock = threading.RLock()
         self._pump_running = False
         self._pump_stop = threading.Event()
-        # When pump is active it is the sole reader of result_mq; non-async
-        # messages are placed here for collective_rpc() to consume.
+        # When pumps are active they are the sole readers of the worker result
+        # queues; non-async messages are placed here for collective_rpc().
         self._sync_result_buffer: queue.Queue = queue.Queue()
         if not self.od_config.step_execution:
             self._start_result_pump()
@@ -140,12 +165,17 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
         if self._broadcast_mq is None:
             raise RuntimeError("Broadcast queue is closed")
 
-    def _dequeue_one_with_failure_polling(self, deadline: float | None, method: str) -> Any:
+    def _dequeue_one_with_failure_polling(
+        self,
+        deadline: float | None,
+        method: str,
+        result_mq: MessageQueue | None = None,
+    ) -> Any:
         """Block until one result message, polling ``_is_failed`` between chunk timeouts.
 
-        When async output is enabled, pump is the sole reader of result_mq;
-        non-async messages are placed in _sync_result_buffer, so this method
-        reads from there instead of result_mq directly.
+        When async output is enabled, the pumps are the sole readers of the
+        worker result queues; non-async messages are placed in
+        _sync_result_buffer, so this method reads from there instead.
         """
         while True:
             if deadline is None:
@@ -163,8 +193,11 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
                         raise EngineDeadError()
                     continue
             else:
+                queue_to_read = result_mq if result_mq is not None else self._result_mq
+                if queue_to_read is None:
+                    raise RuntimeError("Result queue is closed")
                 try:
-                    return self._result_mq.dequeue(timeout=chunk_timeout)  # pyright: ignore[reportOptionalMemberAccess] MQ is not None before shutdown
+                    return queue_to_read.dequeue(timeout=chunk_timeout)
                 except (TimeoutError, zmq.error.Again):
                     if self._is_failed:
                         raise EngineDeadError()
@@ -223,7 +256,14 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
 
     _MAX_STALE_DISCARDS = 16
 
-    def _validate_wave_id(self, response: Any, expected_wave_id: int, deadline: float | None, method: str) -> Any:
+    def _validate_wave_id(
+        self,
+        response: Any,
+        expected_wave_id: int,
+        deadline: float | None,
+        method: str,
+        result_mq: MessageQueue | None = None,
+    ) -> Any:
         """Discard stale RPC responses from a previous wave."""
         discards = 0
         while discards < self._MAX_STALE_DISCARDS:
@@ -239,7 +279,7 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
                 method,
             )
             discards += 1
-            response = self._dequeue_one_with_failure_polling(deadline, method)
+            response = self._dequeue_one_with_failure_polling(deadline, method, result_mq)
         raise TimeoutError(
             f"Discarded {self._MAX_STALE_DISCARDS} stale RPC responses "
             f"without finding wave_id={expected_wave_id} for method={method}."
@@ -249,7 +289,7 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
         self,
         broadcast_handle: Handle,
         wake_events: list[Event],
-    ) -> tuple[list[mp.Process], Handle | None]:
+    ) -> tuple[list[mp.Process], list[Handle]]:
         od_config = self.od_config
         logger.info("Starting server...")
 
@@ -287,7 +327,7 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
             processes.append(process)
 
         # Wait for all workers to be ready
-        result_handle = None
+        result_handles: list[Handle] = []
         for writer in scheduler_pipe_writers:
             writer.close()
 
@@ -303,14 +343,16 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
             if data["status"] != "ready":
                 raise RuntimeError("Initialization failed. Please see the error messages above.")
 
-            if i == 0:
-                result_handle = data.get("result_handle")
+            result_handle = data.get("result_handle")
+            if result_handle is None:
+                raise RuntimeError(f"Rank {i} did not provide a result queue handle")
+            result_handles.append(result_handle)
 
             reader.close()
 
         logger.debug("All workers are ready")
 
-        return processes, result_handle
+        return processes, result_handles
 
     @property
     def is_dead(self) -> bool:
@@ -369,6 +411,27 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
         t = threading.Thread(target=_monitor, daemon=True, name="diffusion-worker-monitor")
         t.start()
 
+    def _fail_closed_on_dp_wave_timeout(self, exc: TimeoutError) -> None:
+        """Shut down every rank after a partial DP wave times out.
+
+        Once one rank is stuck in a DLO AllGather, that process group cannot be
+        reused safely. Shutting down the executor converts an otherwise
+        permanent request hang into a bounded engine failure and lets the
+        caller restart.
+        """
+        logger.error(
+            "DLO DP collective wave timed out after %.1fs; shutting down the worker group: %s",
+            _DLO_DP_WAVE_TIMEOUT_S,
+            exc,
+        )
+        self._is_failed = True
+        self.shutdown()
+        for callback in self._failure_callbacks:
+            try:
+                callback()
+            except Exception:
+                logger.exception("failure_callback raised")
+
     def register_failure_callback(
         self,
         callback: Callable[[], None],
@@ -398,21 +461,17 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
             and getattr(self.od_config, "enable_distributed_layerwise_offload", False)
             and getattr(self.od_config, "dlo_use_allgather", True)
         ):
-            # Validate: all concurrent requests must share the same
-            # num_inference_steps and identical extra_args, because
-            # AllGather is a collective that requires every rank to
-            # participate at each step.
-            step_counts = {
-                nr.req.sampling_params.num_inference_steps
-                for nr in new_reqs
-                if nr.req.sampling_params.num_inference_steps is not None
-            }
-            has_none = any(nr.req.sampling_params.num_inference_steps is None for nr in new_reqs)
-            if (len(step_counts) > 1) or has_none:
+            # Reuse the request scheduler's complete compatibility key. DLO
+            # AllGather requires every DP rank to execute the same collective
+            # schedule, including shape, CFG, denoise steps, output count,
+            # and LoRA settings. Two default (None) step counts are valid and
+            # resolve identically inside the same pipeline.
+            compatibility_keys = [build_request_batch_sampling_params_key(nr.req) for nr in new_reqs]
+            if any(key != compatibility_keys[0] for key in compatibility_keys[1:]):
                 raise ValueError(
-                    "DP multi-concurrency requires all concurrent requests to have "
-                    "the same explicit num_inference_steps (None is not allowed), got "
-                    f"{[nr.req.sampling_params.num_inference_steps for nr in new_reqs]}."
+                    "DLO DP multi-concurrency requires compatible shape, CFG, "
+                    "denoise schedule, output count, and LoRA settings for all "
+                    "requests in one collective wave."
                 )
             extra_args_signatures: set = set()
             for nr in new_reqs:
@@ -424,11 +483,18 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
                     "share identical extra_args. Different extra_args can change "
                     "the forward schedule and cause AllGather deadlock."
                 )
+            empty_prompt_ids = [nr.request_id for nr in new_reqs if _is_empty_dp_prompt(nr.req.prompt)]
+            if empty_prompt_ids:
+                raise ValueError(
+                    "DP multi-concurrency requires a non-empty prompt for every request; "
+                    f"empty prompt request IDs: {empty_prompt_ids}."
+                )
 
             reqs_list = [nr.req for nr in new_reqs]
             try:
                 results = self.collective_rpc(
                     "execute_model",
+                    timeout=_DLO_DP_WAVE_TIMEOUT_S,
                     args=(reqs_list, self.od_config, scheduler_output.kv_prefetch_job),
                     unique_reply_rank=None,
                     exec_all_ranks=True,
@@ -448,6 +514,8 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
                     else:
                         raise RuntimeError(f"Unexpected response type [{i}]: {type(res)!r}")
             except Exception as exc:
+                if isinstance(exc, TimeoutError):
+                    self._fail_closed_on_dp_wave_timeout(exc)
                 for new_req in new_reqs:
                     runner_outputs.append(
                         RunnerOutput(
@@ -591,14 +659,18 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
         deadline = None if timeout is None else time.monotonic() + timeout
         kwargs = kwargs or {}
 
+        multi_rank_reply = unique_reply_rank is None and exec_all_ranks
         execute_all_ranks = unique_reply_rank is None or exec_all_ranks
-        collect_rank_status = unique_reply_rank is None
+        # Status aggregation is for control-plane RPCs, where rank 0 sends
+        # one envelope representing every rank. DP request concurrency needs
+        # one independent reply per DP primary instead.
+        collect_rank_status = unique_reply_rank is None and not multi_rank_reply
         rpc_request = {
             "type": "rpc",
             "method": method,
             "args": args,
             "kwargs": kwargs,
-            "output_rank": unique_reply_rank if unique_reply_rank is not None else 0,
+            "output_rank": None if multi_rank_reply else (unique_reply_rank if unique_reply_rank is not None else 0),
             "exec_all_ranks": execute_all_ranks,
             "collect_rank_status": collect_rank_status,
         }
@@ -608,7 +680,6 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
         # so the D2H copy runs on a side stream in the worker background
         # thread while the default stream is free for the next forward.
         # pump routes the result back to a Future via rpc_id.
-        multi_rank_reply = unique_reply_rank is None and exec_all_ranks
         if (
             not self.od_config.step_execution
             and method in ("execute_model", "execute_model_batch")
@@ -651,11 +722,28 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
             responses: list = []
             if unique_reply_rank is None and exec_all_ranks and num_responses > 1:
                 # DP multi-concurrency: collect num_responses replies, sort by dp_rank.
+                result_mqs: list[MessageQueue | None] = [None] * num_responses
+                if self.od_config.step_execution:
+                    parallel_config = self.od_config.parallel_config
+                    replica_size = (
+                        getattr(parallel_config, "tensor_parallel_size", 1)
+                        * getattr(parallel_config, "sequence_parallel_size", 1)
+                        * getattr(parallel_config, "pipeline_parallel_size", 1)
+                        * getattr(parallel_config, "cfg_parallel_size", 1)
+                    )
+                    primary_ranks = [dp_rank * replica_size for dp_rank in range(num_responses)]
+                    all_result_mqs = getattr(self, "_result_mqs", [])
+                    if not all_result_mqs or primary_ranks[-1] >= len(all_result_mqs):
+                        raise RuntimeError(
+                            "Missing result queues for DP primary ranks "
+                            f"{primary_ranks}; available queues: {len(all_result_mqs)}"
+                        )
+                    result_mqs = [all_result_mqs[rank] for rank in primary_ranks]
                 tagged: list[tuple[int, Any]] = []
                 collected_errors: list[str] = []
-                for _ in range(num_responses):
-                    response = self._dequeue_one_with_failure_polling(deadline, method)
-                    response = self._validate_wave_id(response, wave_id, deadline, method)
+                for result_mq in result_mqs:
+                    response = self._dequeue_one_with_failure_polling(deadline, method, result_mq)
+                    response = self._validate_wave_id(response, wave_id, deadline, method, result_mq)
                     try:
                         unpack_diffusion_output_shm(response)
                     except Exception as e:
@@ -698,20 +786,35 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
     def _start_result_pump(self) -> None:
         self._pump_running = True
         self._pump_stop.clear()
-        self._result_pump_thread = threading.Thread(target=self._result_pump, daemon=True, name="DiffusionResultPump")
-        self._result_pump_thread.start()
-        logger.info("Async result pump started")
+        result_mqs = self._result_mqs or ([self._result_mq] if self._result_mq is not None else [])
+        self._result_pump_threads = [
+            threading.Thread(
+                target=self._result_pump,
+                args=(result_mq,),
+                daemon=True,
+                name=f"DiffusionResultPump-{rank}",
+            )
+            for rank, result_mq in enumerate(result_mqs)
+        ]
+        for thread in self._result_pump_threads:
+            thread.start()
+        self._result_pump_thread = self._result_pump_threads[0] if self._result_pump_threads else None
+        logger.info("Async result pump started for %d worker queue(s)", len(self._result_pump_threads))
 
-    def _result_pump(self) -> None:
-        """Sole reader of result_mq when async output is enabled.
+    def _result_pump(self, result_mq: MessageQueue | None = None) -> None:
+        """Sole reader of one worker result queue when async output is enabled.
 
         Dispatches AsyncDiffusionOutput messages to the appropriate future:
         * RPC_RESULT / COMPUTE_DONE → _rpc_futures[rpc_id]
         * OUTPUT_READY → _output_futures[async_output_id]
         """
+        result_mq = result_mq or self._result_mq
+        if result_mq is None:
+            return
+
         while not self._pump_stop.is_set():
             try:
-                msg = self._result_mq.dequeue(timeout=1.0)
+                msg = result_mq.dequeue(timeout=1.0)
             except TimeoutError:
                 if self._is_failed:
                     break
@@ -823,8 +926,18 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
         try:
             self._finalizer()
         finally:
+            pump_threads = getattr(self, "_result_pump_threads", [])
+            for thread in pump_threads:
+                if thread is threading.current_thread():
+                    continue
+                thread.join(timeout=_RESULT_PUMP_JOIN_TIMEOUT_S)
+                if thread.is_alive():
+                    logger.warning("Result pump thread %s did not stop before shutdown", thread.name)
+            self._pump_running = False
             self._broadcast_mq = None
             self._result_mq = None
+            self._result_mqs = []
+            self._result_pump_threads = []
             with self._futures_lock:
                 for fut in self._rpc_futures.values():
                     if not fut.done():

vllm_omni/diffusion/inline_stage_diffusion_client.py

@@ -8,6 +8,7 @@ IPC overhead. Used when there is only a single diffusion stage.
 from __future__ import annotations
 
 import asyncio
+import threading
 import time
 from concurrent.futures import ThreadPoolExecutor
 from typing import TYPE_CHECKING, Any
@@ -66,6 +67,8 @@ class InlineStageDiffusionClient(StageClientBase):
         self._tasks: dict[str, asyncio.Task] = {}
         self._engine_dead = False
         self._shutting_down = False
+        self._shutdown_complete = False
+        self._shutdown_lock = threading.Lock()
 
         self._engine.executor.register_failure_callback(self._mark_engine_dead)
 
@@ -101,6 +104,10 @@ class InlineStageDiffusionClient(StageClientBase):
         sampling_params: OmniDiffusionSamplingParams,
         kv_sender_info: dict[int, dict[str, Any]] | None = None,
     ) -> None:
+        # Each request mutates its sampling state while it is normalized and
+        # executed. Callers commonly reuse one params object for concurrent
+        # requests, so take the copy synchronously before either task starts.
+        sampling_params = sampling_params.clone()
         logger.debug(
             "[InlineStageDiffusionClient] stage-%s [rep-%s] add request: %s",
             self.stage_id,
@@ -298,19 +305,25 @@ class InlineStageDiffusionClient(StageClientBase):
             raise
 
     def shutdown(self) -> None:
-        self._shutting_down = True
+        with self._shutdown_lock:
+            if self._shutdown_complete:
+                return
+            self._shutting_down = True
 
-        # Cancel all pending tasks
-        for task in self._tasks.values():
-            task.cancel()
+            # Cancel all pending tasks
+            for task in self._tasks.values():
+                task.cancel()
 
-        try:
-            # Cancel queued futures and wait for the running one to complete deterministically
-            self._executor.shutdown(wait=True, cancel_futures=True)
-        except Exception:
-            pass
+            try:
+                # Stop the engine first so any control RPC running in the thread
+                # pool can observe shutdown instead of keeping stage teardown
+                # blocked while the executor waits for that RPC.
+                self._engine.close()
+            except Exception:
+                pass
 
-        try:
-            self._engine.close()
-        except Exception:
-            pass
+            try:
+                self._executor.shutdown(wait=True, cancel_futures=True)
+            except Exception:
+                pass
+            self._shutdown_complete = True

vllm_omni/diffusion/ipc.py

@@ -13,6 +13,7 @@ from __future__ import annotations
 
 from typing import Any
 
+import numpy as np
 import torch
 
 from vllm_omni.diffusion.data import DiffusionOutput
@@ -21,6 +22,45 @@ _SHM_TENSOR_THRESHOLD = 1_000_000  # 1 MB
 DIFFUSION_RPC_RESULT_ENVELOPE = "diffusion_rpc_result"
 
 
+def _array_to_shm(array: np.ndarray) -> dict[str, Any]:
+    """Copy a contiguous NumPy-compatible array into shared memory."""
+    from multiprocessing import shared_memory
+
+    if array.dtype.hasobject:
+        raise TypeError("NumPy object arrays cannot be transferred through raw shared memory")
+
+    array = np.ascontiguousarray(array)
+    nbytes = array.nbytes
+    shm = shared_memory.SharedMemory(create=True, size=nbytes)
+    shm_array = np.ndarray(array.shape, dtype=array.dtype, buffer=shm.buf[:nbytes])
+    np.copyto(shm_array, array)
+    handle = {
+        "name": shm.name,
+        "shape": list(array.shape),
+        "numpy_dtype": str(array.dtype),
+        "nbytes": nbytes,
+    }
+    shm.close()
+    return handle
+
+
+def _array_from_shm(handle: dict[str, Any]) -> np.ndarray:
+    """Copy an array from shared memory, then close and unlink its segment."""
+    from multiprocessing import shared_memory
+
+    shm = shared_memory.SharedMemory(name=handle["name"])
+    try:
+        array = np.ndarray(
+            handle["shape"],
+            dtype=np.dtype(handle["numpy_dtype"]),
+            buffer=shm.buf[: handle["nbytes"]],
+        ).copy()
+    finally:
+        shm.close()
+        shm.unlink()
+    return array
+
+
 def _tensor_to_shm(
     tensor: torch.Tensor,
     d2h_stream: torch.Stream | None = None,
@@ -35,10 +75,6 @@ def _tensor_to_shm(
     path.  The caller must synchronize *d2h_stream* after all tensors are
     packed.
     """
-    from multiprocessing import shared_memory
-
-    import numpy as np
-
     original_dtype = tensor.dtype
     if d2h_stream is not None:
         # Non-blocking D2H: copy on side stream to pinned CPU memory.
@@ -58,44 +94,26 @@ def _tensor_to_shm(
         tensor = tensor.detach().cpu().contiguous()
         if original_dtype == torch.bfloat16:
             tensor = tensor.to(torch.float32)
-    arr = tensor.numpy()
-    nbytes = arr.nbytes
-    shm = shared_memory.SharedMemory(create=True, size=nbytes)
-    shm_arr = np.ndarray(arr.shape, dtype=arr.dtype, buffer=shm.buf[:nbytes])
-    np.copyto(shm_arr, arr)
-    handle = {
-        "__tensor_shm__": True,
-        "name": shm.name,
-        "shape": list(tensor.shape),
-        "torch_dtype": str(original_dtype),
-        "numpy_dtype": str(arr.dtype),
-        "nbytes": nbytes,
-    }
-    shm.close()
+    handle = _array_to_shm(tensor.numpy())
+    handle.update(
+        {
+            "__tensor_shm__": True,
+            "torch_dtype": str(original_dtype),
+        }
+    )
     return handle
 
 
 def _tensor_from_shm(handle: dict[str, Any]) -> torch.Tensor:
     """Reconstruct a tensor from a shared-memory handle and free the segment."""
-    from multiprocessing import shared_memory
-
-    import numpy as np
-
-    shm = shared_memory.SharedMemory(name=handle["name"])
-    try:
-        np_dtype = np.dtype(handle["numpy_dtype"])
-        arr = np.ndarray(handle["shape"], dtype=np_dtype, buffer=shm.buf[: handle["nbytes"]])
-        tensor = torch.from_numpy(arr.copy())
-        # Restore the original dtype if it differs from the numpy-compatible
-        # dtype used for the SHM transfer (e.g. bfloat16 → float32 → bfloat16).
-        torch_dtype_str = handle.get("torch_dtype", "")
-        if torch_dtype_str:
-            original_dtype = getattr(torch, torch_dtype_str.replace("torch.", ""), None)
-            if original_dtype is not None and tensor.dtype != original_dtype:
-                tensor = tensor.to(original_dtype)
-    finally:
-        shm.close()
-        shm.unlink()
+    tensor = torch.from_numpy(_array_from_shm(handle))
+    # Restore the original dtype if it differs from the numpy-compatible
+    # dtype used for the SHM transfer (e.g. bfloat16 → float32 → bfloat16).
+    torch_dtype_str = handle.get("torch_dtype", "")
+    if torch_dtype_str:
+        original_dtype = getattr(torch, torch_dtype_str.replace("torch.", ""), None)
+        if original_dtype is not None and tensor.dtype != original_dtype:
+            tensor = tensor.to(original_dtype)
     return tensor
 
 
@@ -109,6 +127,18 @@ def _pack_tensor_if_large(
     return val
 
 
+def _ndarray_to_shm(array: np.ndarray) -> dict[str, Any]:
+    """Copy a contiguous NumPy array into POSIX shared memory."""
+    handle = _array_to_shm(array)
+    handle["__ndarray_shm__"] = True
+    return handle
+
+
+def _ndarray_from_shm(handle: dict[str, Any]) -> np.ndarray:
+    """Reconstruct a NumPy array from shared memory and free the segment."""
+    return _array_from_shm(handle)
+
+
 def _pack_value_if_large(
     val: object,
     d2h_stream: torch.Stream | None = None,
@@ -123,6 +153,8 @@ def _pack_value_if_large(
     """
     if isinstance(val, torch.Tensor):
         return _pack_tensor_if_large(val, d2h_stream=d2h_stream)
+    if isinstance(val, np.ndarray):
+        return _ndarray_to_shm(val) if not val.dtype.hasobject and val.nbytes > _SHM_TENSOR_THRESHOLD else val
     if isinstance(val, dict):
         return {key: _pack_value_if_large(value, d2h_stream=d2h_stream) for key, value in val.items()}
     if isinstance(val, list):
@@ -136,6 +168,8 @@ def _unpack_if_shm_handle(val: object) -> object:
     """Reconstruct tensors from SHM handles, mirroring ``_pack_value_if_large``."""
     if isinstance(val, dict) and val.get("__tensor_shm__"):
         return _tensor_from_shm(val)
+    if isinstance(val, dict) and val.get("__ndarray_shm__"):
+        return _ndarray_from_shm(val)
     if isinstance(val, dict):
         return {key: _unpack_if_shm_handle(value) for key, value in val.items()}
     if isinstance(val, list):
@@ -186,13 +220,12 @@ def pack_diffusion_output_shm(
     if isinstance(output, dict) and "dp_rank" in output and "output" in output:
         inner = output["output"]
         if isinstance(inner, DiffusionOutput):
-            output["output"] = _pack_diffusion_fields(inner)
+            output["output"] = _pack_diffusion_fields(inner, d2h_stream=d2h_stream)
         return output
 
     if _is_rpc_result_envelope(output):
         result = output.get("result")
-        if isinstance(result, DiffusionOutput):
-            output["result"] = _pack_diffusion_fields(result, d2h_stream=d2h_stream)
+        output["result"] = pack_diffusion_output_shm(result, d2h_stream=d2h_stream)
         return output
 
     result = getattr(output, "result", None)
@@ -228,8 +261,7 @@ def unpack_diffusion_output_shm(output: object) -> object:
 
     if _is_rpc_result_envelope(output):
         result = output.get("result")
-        if isinstance(result, DiffusionOutput):
-            output["result"] = _unpack_diffusion_fields(result)
+        output["result"] = unpack_diffusion_output_shm(result)
         return output
 
     result = getattr(output, "result", None)

vllm_omni/diffusion/model_metadata.py

@@ -64,4 +64,21 @@ def get_diffusion_model_metadata(model_class_name: str | None) -> DiffusionModel
     # pipelines do not accidentally inherit limits meant for other models.
     if model_class_name is None:
         return DiffusionModelMetadata()
-    return _DIFFUSION_MODEL_METADATA.get(model_class_name, DiffusionModelMetadata())
+    metadata = _DIFFUSION_MODEL_METADATA.get(model_class_name)
+    if metadata is not None:
+        return metadata
+    # Some checkpoints report the HF architecture name diff from internal pipeline class name
+    # (e.g. HunyuanImage3ForCausalMM, WanVACEPipeline, OmniVoice ...).
+    from vllm_omni.diffusion.registry import _DIFFUSION_MODELS
+
+    entry = _DIFFUSION_MODELS.get(model_class_name)
+    if entry is not None:
+        # Unpack instead of indexing so a future change to the registry tuple
+        # shape fails loudly instead of silently reading the wrong element.
+        _, _, pipeline_cls_name = entry
+        # Note: the registry ``cls_name`` and the metadata keys are two separate
+        # key spaces. Aliases whose pipeline class has no metadata entry (e.g.
+        # Wan22VACEPipeline, OmniVoicePipeline) still fall back to the defaults;
+        # that is not a regression, it just means no capability override.
+        return _DIFFUSION_MODEL_METADATA.get(pipeline_cls_name, DiffusionModelMetadata())
+    return DiffusionModelMetadata()

vllm_omni/diffusion/models/interface.py

@@ -17,6 +17,7 @@ if TYPE_CHECKING:
 
     import torch
 
+    from vllm_omni.diffusion.cache.cachedit import CacheDiTBackend
     from vllm_omni.diffusion.data import DiffusionOutput
     from vllm_omni.diffusion.worker.input_batch import InputBatch
     from vllm_omni.diffusion.worker.utils import StepRequestState
@@ -131,3 +132,31 @@ def supports_prompt_update(pipeline: object) -> bool:
     """Return whether ``pipeline`` implements :class:`SupportsPromptUpdate`."""
 
     return isinstance(pipeline, SupportsPromptUpdate)
+
+
+@runtime_checkable
+class SupportsRequestScopedCacheDiT(Protocol):
+    """Optional protocol for pipelines that own Cache-DiT hook transitions."""
+
+    def adopt_cache_dit_backend(self, backend: CacheDiTBackend) -> None:
+        """Assume ownership of an enabled Cache-DiT backend."""
+        ...
+
+    def is_cache_dit_enabled(self) -> bool:
+        """Return whether this pipeline currently has Cache-DiT installed."""
+        ...
+
+
+def adopt_request_scoped_cache_dit(pipeline: object, backend: CacheDiTBackend) -> bool:
+    """Transfer an enabled Cache-DiT backend to an opted-in pipeline."""
+
+    if not isinstance(pipeline, SupportsRequestScopedCacheDiT):
+        return False
+    pipeline.adopt_cache_dit_backend(backend)
+    return True
+
+
+def is_request_scoped_cache_dit_enabled(pipeline: object) -> bool:
+    """Read Cache-DiT state from a pipeline that owns its lifecycle."""
+
+    return isinstance(pipeline, SupportsRequestScopedCacheDiT) and pipeline.is_cache_dit_enabled()

vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py

@@ -22,6 +22,10 @@ from transformers import Qwen2TokenizerFast, Qwen3VLProcessor
 from vllm.logger import init_logger
 
 from vllm_omni.diffusion import envs
+from vllm_omni.diffusion.cache.cachedit import (
+    CacheDiTBackend,
+    RequestScopedCacheDiTRuntime,
+)
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.parallel_state import (
     get_dit_group,
@@ -73,6 +77,7 @@ from .presentation import (
     minimax_h3_ref2va_video_presentation,
     minimax_h3_text_only_ids,
 )
+from .quality_policy import MINIMAX_H3_GENERIC_CACHE_KEY, MiniMaxH3QualityPolicy
 from .reference_video import (
     load_audio_file,
     load_video_audio,
@@ -538,6 +543,19 @@ class MiniMaxH3Pipeline(
     ]
     dummy_run_num_frames: ClassVar[int] = 0
 
+    def adopt_cache_dit_backend(self, backend: CacheDiTBackend) -> None:
+        """Adopt runner-installed generic Cache-DiT for request transitions."""
+
+        self._cache_dit_runtime.adopt(
+            backend,
+            installation_key=MINIMAX_H3_GENERIC_CACHE_KEY,
+        )
+
+    def is_cache_dit_enabled(self) -> bool:
+        """Return the request-scoped Cache-DiT installation state."""
+
+        return self._cache_dit_runtime.is_enabled
+
     def __init__(
         self,
         *,
@@ -676,6 +694,9 @@ class MiniMaxH3Pipeline(
         # Registry-side VAE patch-parallel discovery uses ``pipeline.vae``.
         self.vae = self.video_vae
 
+        self._quality_policy = MiniMaxH3QualityPolicy(od_config)
+        self._cache_dit_runtime = RequestScopedCacheDiTRuntime(self)
+
         self.setup_diffusion_pipeline_profiler(
             enable_diffusion_pipeline_profiler=(od_config.enable_diffusion_pipeline_profiler)
         )
@@ -1551,6 +1572,8 @@ class MiniMaxH3Pipeline(
             raise OmniClientError("MiniMax H3 requires a non-empty prompt")
 
         sampling = request.sampling_params
+        quality = sampling.quality
+        logger.debug("MiniMax H3 request quality=%s", quality)
         extra = sampling.extra_args or {}
         task = self._resolve_task(extra.get("task"), multi_modal_data)
 
@@ -1732,6 +1755,12 @@ class MiniMaxH3Pipeline(
         num_steps = int(sampling.num_inference_steps or 50)
         video_shift = float(extra.get("flow_shift", self.default_video_shift))
         audio_shift = float(extra.get("audio_flow_shift", self.default_audio_shift))
+        quality_plan = self._quality_policy.resolve(
+            quality=quality,
+            num_inference_steps=num_steps,
+            extra_args=extra,
+        )
+        self._cache_dit_runtime.prepare(quality_plan.cache_dit)
         num_outputs = _resolve_minimax_h3_num_outputs(sampling.num_outputs_per_prompt)
         videos = []
         audios = []

vllm_omni/diffusion/models/minimax_h3/quality_policy.py

@@ -0,0 +1,178 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Model-owned request quality policy for MiniMax H3."""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from dataclasses import dataclass, replace
+from numbers import Integral
+from typing import Any
+
+from vllm_omni.diffusion.cache.cachedit import CacheDiTRequestSpec
+from vllm_omni.diffusion.data import DiffusionCacheConfig
+from vllm_omni.errors import OmniClientError
+
+MINIMAX_H3_GENERIC_CACHE_KEY = "minimax_h3.generic"
+MINIMAX_H3_HIGH_CACHE_KEY = "minimax_h3.high"
+MINIMAX_H3_FORCE_REFRESH_STEP_HINT_ARG = "force_refresh_step_hint"
+MINIMAX_H3_FORCE_REFRESH_STEP_POLICY_ARG = "force_refresh_step_policy"
+MINIMAX_H3_FORCE_REFRESH_POLICIES = ("once", "repeat")
+
+
+def _high_quality_cache_config() -> DiffusionCacheConfig:
+    return DiffusionCacheConfig(
+        Fn_compute_blocks=1,
+        Bn_compute_blocks=0,
+        max_warmup_steps=4,
+        residual_diff_threshold=0.04,
+        max_continuous_cached_steps=1,
+        enable_taylorseer=False,
+        scm_steps_mask_policy=None,
+    )
+
+
+def _resolve_force_refresh(
+    extra_args: Mapping[str, Any] | None,
+    *,
+    num_inference_steps: int,
+) -> tuple[int, str] | None:
+    """Resolve H3's request-scoped Cache-DiT refresh hint.
+
+    These are model-specific ``extra_args`` rather than global sampling fields,
+    keeping other diffusion models unchanged.
+    """
+
+    if not extra_args:
+        return None
+
+    raw_hint = extra_args.get(MINIMAX_H3_FORCE_REFRESH_STEP_HINT_ARG)
+    raw_policy = extra_args.get(MINIMAX_H3_FORCE_REFRESH_STEP_POLICY_ARG)
+
+    if raw_hint is None:
+        if raw_policy is not None:
+            raise OmniClientError("MiniMax H3 force_refresh_step_policy requires force_refresh_step_hint")
+        return None
+    if isinstance(raw_hint, bool) or not isinstance(raw_hint, Integral):
+        raise OmniClientError("MiniMax H3 force_refresh_step_hint must be a positive integer")
+
+    hint = int(raw_hint)
+    if not 1 <= hint <= num_inference_steps:
+        raise OmniClientError(
+            "MiniMax H3 force_refresh_step_hint must be between 1 and "
+            f"num_inference_steps ({num_inference_steps}), got {hint}"
+        )
+
+    policy = "once" if raw_policy is None else raw_policy
+    if not isinstance(policy, str) or policy not in MINIMAX_H3_FORCE_REFRESH_POLICIES:
+        raise OmniClientError(
+            "MiniMax H3 force_refresh_step_policy must be one of "
+            f"{list(MINIMAX_H3_FORCE_REFRESH_POLICIES)}, got {policy!r}"
+        )
+    return hint, policy
+
+
+def _has_force_refresh_args(extra_args: Mapping[str, Any] | None) -> bool:
+    if not extra_args:
+        return False
+    return any(
+        extra_args.get(name) is not None
+        for name in (
+            MINIMAX_H3_FORCE_REFRESH_STEP_HINT_ARG,
+            MINIMAX_H3_FORCE_REFRESH_STEP_POLICY_ARG,
+        )
+    )
+
+
+def _with_force_refresh(
+    cache_config: DiffusionCacheConfig,
+    force_refresh: tuple[int, str] | None,
+) -> DiffusionCacheConfig:
+    if force_refresh is None:
+        return cache_config
+    hint, policy = force_refresh
+    return replace(
+        cache_config,
+        force_refresh_step_hint=hint,
+        force_refresh_step_policy=policy,
+    )
+
+
+def _cache_installation_key(base_key: str, force_refresh: tuple[int, str] | None) -> str:
+    """Make refresh-policy transitions explicit to the request runtime.
+
+    Cache-DiT cannot clear an existing ``force_refresh_step_hint`` through its
+    incremental context update API because ``None`` is treated as "keep the
+    old value".  Including the request hint in the installation key therefore
+    makes a hint change perform a safe hook reinstall, while repeated requests
+    with the same hint still only refresh the context.
+    """
+
+    if force_refresh is None:
+        return base_key
+    hint, policy = force_refresh
+    return f"{base_key}:force_refresh={hint}:{policy}"
+
+
+@dataclass(frozen=True)
+class MiniMaxH3QualityPlan:
+    """Resolved execution choices for one MiniMax H3 request."""
+
+    cache_dit: CacheDiTRequestSpec | None
+
+
+class MiniMaxH3QualityPolicy:
+    """Resolve H3 request quality into a declarative Cache-DiT target.
+
+    When the server starts with Cache-DiT, omitted quality selects the
+    server-configured profile. Otherwise, omitted quality selects no cache.
+    In either case, ``lossless`` selects no cache and ``high`` selects H3's
+    high-quality profile. The policy therefore owns whether a request installs
+    Cache-DiT; the startup backend only controls the omitted-quality default.
+    H3-specific Cache-DiT refresh hints are read from request ``extra_args``;
+    they do not change the global diffusion request contract.
+    The pipeline owns applying the resulting target at the request boundary.
+    """
+
+    def __init__(self, od_config: Any) -> None:
+        self._od_config = od_config
+        self._configured_backend = str(getattr(od_config, "cache_backend", "none") or "none").lower()
+
+    def resolve(
+        self,
+        *,
+        quality: str | None,
+        num_inference_steps: int,
+        extra_args: Mapping[str, Any] | None = None,
+    ) -> MiniMaxH3QualityPlan:
+        if quality == "high":
+            base_key = MINIMAX_H3_HIGH_CACHE_KEY
+            base_config = _high_quality_cache_config()
+        elif self._configured_backend == "cache_dit" and quality is None:
+            base_key = MINIMAX_H3_GENERIC_CACHE_KEY
+            base_config = self._od_config.cache_config
+        else:
+            if _has_force_refresh_args(extra_args):
+                raise OmniClientError("MiniMax H3 force-refresh arguments require an active Cache-DiT request target")
+            return MiniMaxH3QualityPlan(cache_dit=None)
+
+        force_refresh = _resolve_force_refresh(
+            extra_args,
+            num_inference_steps=num_inference_steps,
+        )
+        return MiniMaxH3QualityPlan(
+            cache_dit=CacheDiTRequestSpec(
+                installation_key=_cache_installation_key(base_key, force_refresh),
+                cache_config=_with_force_refresh(base_config, force_refresh),
+                num_inference_steps=num_inference_steps,
+            ),
+        )
+
+
+__all__ = [
+    "MINIMAX_H3_GENERIC_CACHE_KEY",
+    "MINIMAX_H3_HIGH_CACHE_KEY",
+    "MiniMaxH3QualityPlan",
+    "MiniMaxH3QualityPolicy",
+]

vllm_omni/diffusion/sched/interface.py

@@ -55,6 +55,10 @@ class StepBatchSamplingParamsKey:
     true_cfg_scale: float | None = None
     cfg_normalize: bool = False
 
+    # Request-scoped execution policy. A batch must use one quality mode
+    # because model acceleration hooks are shared by the whole worker batch.
+    quality: str | None = None
+
     # Output count. Requests with different num_outputs_per_prompt produce
     # differently shaped outputs and cannot share a batch.
     num_outputs_per_prompt: int = 1
@@ -108,6 +112,11 @@ class RequestBatchSamplingParamsKey:
     layers: int = 4
     use_en_prompt: bool = False
 
+    # Request-scoped execution policy. Keep accelerated and reference-path
+    # requests in separate worker batches. ``None`` delegates the default to
+    # the model and must remain distinct from explicit ``lossless``.
+    quality: str | None = None
+
     # LoRA identity.
     lora_int_id: int | None = None
     lora_scale: float = 1.0

vllm_omni/diffusion/sched/request_scheduler.py

@@ -24,17 +24,21 @@ _REQUEST_BATCH_SAMPLING_PARAMS_KEY_FIELD_NAMES = frozenset(
 ) - {"lora_int_id"}
 
 
+def build_request_batch_sampling_params_key(request: OmniDiffusionRequest) -> RequestBatchSamplingParamsKey:
+    """Build the compatibility key shared by scheduling and DP dispatch."""
+    sampling = request.sampling_params
+    # LoRA identity is optional on sampling params (and on test stubs).
+    lora_request = getattr(sampling, "lora_request", None)
+    key_kwargs = {name: getattr(sampling, name) for name in _REQUEST_BATCH_SAMPLING_PARAMS_KEY_FIELD_NAMES}
+    key_kwargs["lora_int_id"] = lora_request.lora_int_id if lora_request is not None else None
+    return RequestBatchSamplingParamsKey(**key_kwargs)
+
+
 class RequestScheduler(BaseScheduler):
     """Diffusion scheduler with vLLM-style waiting/running queues."""
 
     def _build_sampling_params_key(self, request: OmniDiffusionRequest) -> RequestBatchSamplingParamsKey:
-        """Build a request-batch compatibility key from sampling parameters."""
-        sampling = request.sampling_params
-        # LoRA identity is optional on sampling params (and on test stubs).
-        lora_request = getattr(sampling, "lora_request", None)
-        key_kwargs = {name: getattr(sampling, name) for name in _REQUEST_BATCH_SAMPLING_PARAMS_KEY_FIELD_NAMES}
-        key_kwargs["lora_int_id"] = lora_request.lora_int_id if lora_request is not None else None
-        return RequestBatchSamplingParamsKey(**key_kwargs)
+        return build_request_batch_sampling_params_key(request)
 
     def update_from_output(self, sched_output: DiffusionSchedulerOutput, output: RunnerOutput) -> set[str]:
         scheduled_request_ids = sched_output.scheduled_request_ids

vllm_omni/diffusion/worker/diffusion_model_runner.py

@@ -32,7 +32,13 @@ from vllm_omni.diffusion.compile import regionally_compile
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.forward_context import set_forward_context
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
-from vllm_omni.diffusion.models.interface import SupportsPromptUpdate, supports_prompt_update, supports_step_execution
+from vllm_omni.diffusion.models.interface import (
+    SupportsPromptUpdate,
+    adopt_request_scoped_cache_dit,
+    is_request_scoped_cache_dit_enabled,
+    supports_prompt_update,
+    supports_step_execution,
+)
 from vllm_omni.diffusion.offloader import get_offload_backend
 from vllm_omni.diffusion.registry import _NO_CACHE_ACCELERATION
 from vllm_omni.diffusion.request import OmniDiffusionRequest
@@ -304,7 +310,19 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
                 self.cache_backend = None
                 self.od_config.cache_backend = None
             else:
+                # Install configured cache capability once at startup. A model
+                # may explicitly adopt the enabled Cache-DiT backend and then
+                # own all later request-boundary enable/disable transitions.
                 self.cache_backend.enable(self.pipeline)
+                if str(self.od_config.cache_backend).lower() == "cache_dit" and adopt_request_scoped_cache_dit(
+                    self.pipeline,
+                    self.cache_backend,
+                ):
+                    logger.info(
+                        "Pipeline %s owns request-scoped Cache-DiT transitions.",
+                        type(self.pipeline).__name__,
+                    )
+                    self.cache_backend = None
 
         # Install prompt-embedding cache (transparent wrapper around
         # ``pipeline.encode_prompt``). Enabled via config or env var; a no-op
@@ -523,11 +541,11 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
             if is_primary and prompt_embed_cache is not None:
                 logger.debug("prompt-embed cache: %s", prompt_embed_cache.stats())
 
+            runner_cache_dit_enabled = self.cache_backend is not None and self.cache_backend.is_enabled()
             if (
-                self.cache_backend is not None
-                and self.cache_backend.is_enabled()
-                and od_config.cache_backend == "cache_dit"
+                od_config.cache_backend == "cache_dit"
                 and od_config.enable_cache_dit_summary
+                and (runner_cache_dit_enabled or is_request_scoped_cache_dit_enabled(self.pipeline))
             ):
                 cache_summary(self.pipeline, details=True)
 

vllm_omni/diffusion/worker/diffusion_worker.py

@@ -63,6 +63,8 @@ from vllm_omni.worker.gpu_memory_utils import get_process_gpu_memory
 
 logger = init_logger(__name__)
 
+_ASYNC_OUTPUT_THREAD_JOIN_TIMEOUT_S = 10.0
+
 
 @dataclass
 class _DiffusionVllmModelConfig:
@@ -429,8 +431,8 @@ class DiffusionWorker:
         ranks stay synchronised at each AllGather call while computing
         different activations.
 
-        Each rank returns its OWN DiffusionOutput (no gather).  The executor
-        collects N responses via result_mq (all ranks share one queue).
+        Each rank returns its OWN DiffusionOutput (no gather). The executor
+        collects N responses via the per-worker result queues.
         """
         assert self.model_runner is not None, "Model runner not initialized"
 
@@ -895,8 +897,11 @@ class WorkerProc:
         """
         device = torch.device(torch.accelerator.current_accelerator().type, self.gpu_id)
         d2h_stream = torch.Stream(device=device)
-        while self._running:
-            output, async_output_id, gpu_event = self._async_output_queue.get()
+        while True:
+            item = self._async_output_queue.get()
+            if item is None:
+                break
+            output, async_output_id, gpu_event = item
             try:
                 # Cross-stream ordering: wait for default stream to finish
                 # writing the output tensors before the side stream reads.
@@ -925,6 +930,40 @@ class WorkerProc:
                     )
                 )
 
+    def shutdown(self) -> None:
+        """Stop background work and release worker-owned IPC resources."""
+        self._running = False
+
+        if self._async_output_queue is not None:
+            self._async_output_queue.put(None)
+        if self._async_output_thread is not None:
+            self._async_output_thread.join(timeout=_ASYNC_OUTPUT_THREAD_JOIN_TIMEOUT_S)
+            if self._async_output_thread.is_alive():
+                logger.warning(
+                    "Worker %d: Async output thread did not stop before shutdown",
+                    self.gpu_id,
+                )
+
+        try:
+            self.worker.shutdown()
+        finally:
+            if self.mq is not None:
+                self.mq.shutdown()
+                self.mq = None
+
+            # The worker creates this queue's shared-memory ring buffer.
+            # Dropping the final creator reference invokes
+            # ShmRingBuffer.__del__, which closes and unlinks it before
+            # multiprocessing.resource_tracker exits.
+            if self.result_mq is not None and (
+                self._async_output_thread is None or not self._async_output_thread.is_alive()
+            ):
+                result_mq = self.result_mq
+                self.result_mq = None
+                result_mq.shutdown()
+                del result_mq
+                gc.collect()
+
     def _gather_rpc_rank_statuses(self, status: dict[str, Any]) -> list[dict[str, Any]]:
         if not torch.distributed.is_initialized():
             return [status]
@@ -1196,7 +1235,7 @@ class WorkerProc:
         finally:
             if worker_proc is not None:
                 try:
-                    worker_proc.worker.shutdown()
+                    worker_proc.shutdown()
                 except Exception as exc:
                     logger.warning("Worker %d: Shutdown encountered an error: %s", rank, exc)
                 worker_proc.context.term()

vllm_omni/engine/arg_utils.py

@@ -188,6 +188,7 @@ class OmniEngineArgs(EngineArgs):
     omni: bool = False
     # Diffusion request-mode batch admission (forwarded to OmniDiffusionConfig).
     request_batch_max_wait_ms: float = 0.0
+    fa_deterministic: bool = False
 
     @classmethod
     def _add_omni_specific_args(cls, parser: argparse.ArgumentParser) -> argparse.ArgumentParser:

vllm_omni/engine/async_omni_engine.py

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

vllm_omni/entrypoints/cli/serve.py

@@ -434,6 +434,17 @@ class OmniServeCommand(CLISubcommand):
                 "Disable for fixed-shape workloads with --no-diffusion-compile-dynamic."
             ),
         )
+        omni_config_group.add_argument(
+            "--fa-deterministic",
+            dest="fa_deterministic",
+            action="store_true",
+            default=False,
+            help=(
+                "Request FlashAttention deterministic=True on the local FLASH_ATTN dense path. "
+                "Slower than the library default deterministic=False; intended for accuracy CI. "
+                "Serving default remains non-deterministic."
+            ),
+        )
         omni_config_group.add_argument(
             "--diffusers-load-kwargs",
             dest="diffusers_load_kwargs",

vllm_omni/entrypoints/openai/api_server.py

@@ -39,6 +39,7 @@ from vllm.entrypoints.mcp.tool_server import DemoToolServer, MCPToolServer, Tool
 from vllm.entrypoints.openai.api_server import build_app as build_openai_app
 from vllm.entrypoints.openai.api_server import setup_server as setup_openai_server
 from vllm.entrypoints.openai.chat_completion.protocol import (
+    BatchChatCompletionRequest,
     ChatCompletionRequest,
     ChatCompletionResponse,
 )
@@ -95,6 +96,7 @@ from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError
 from vllm_omni.config.endpoint_policy import shutdown_unsupported_routes
 from vllm_omni.diffusion.models.interface import ReferenceVideoDecodeSpec
 from vllm_omni.entrypoints.async_omni import AsyncOmni
+from vllm_omni.entrypoints.openai.batch_serving import OmniOpenAIServingChatBatch
 from vllm_omni.entrypoints.openai.duplex_capability import should_enable_duplex_endpoint
 from vllm_omni.entrypoints.openai.errors import InvalidInputReferenceError
 from vllm_omni.entrypoints.openai.image_api_utils import (
@@ -513,6 +515,7 @@ async def omni_run_server_worker(listen_address, sock, args, client_config=None,
 
         # OMNI: Remove upstream routes that we override with omni-specific handlers
         _remove_route_from_app(app, "/v1/chat/completions", {"POST"})
+        _remove_route_from_app(app, "/v1/chat/completions/batch", {"POST"})
         _remove_route_from_app(app, "/v1/models", {"GET"})  # Remove upstream /v1/models to use omni's handler
         app.include_router(router)
 
@@ -777,6 +780,10 @@ async def omni_init_app_state(
             diffusion_engine=engine_client,  # type: ignore
             model_name=model_name,
         )
+        state.openai_serving_chat_batch = OmniOpenAIServingChatBatch.for_diffusion(
+            diffusion_engine=engine_client,  # type: ignore
+            model_name=model_name,
+        )
 
         # audio related
         state.openai_serving_speech = None
@@ -950,30 +957,33 @@ async def omni_init_app_state(
         if "generate" in supported_tasks
         else None
     )
-    state.openai_serving_chat = (
-        OmniOpenAIServingChat(
-            engine_client,
-            state.openai_serving_models,
-            args.response_role,
-            online_renderer=state.online_renderer,
-            request_logger=request_logger,
-            chat_template=resolved_chat_template,
-            chat_template_content_format=args.chat_template_content_format,
-            default_chat_template_kwargs=args.default_chat_template_kwargs,
-            trust_request_chat_template=args.trust_request_chat_template,
-            return_tokens_as_token_ids=args.return_tokens_as_token_ids,
-            enable_auto_tools=args.enable_auto_tool_choice,
-            exclude_tools_when_tool_choice_none=args.exclude_tools_when_tool_choice_none,
-            tool_parser=args.tool_call_parser,
-            reasoning_parser=args.structured_outputs_config.reasoning_parser,
-            enable_prompt_tokens_details=args.enable_prompt_tokens_details,
-            enable_force_include_usage=args.enable_force_include_usage,
-            enable_log_outputs=args.enable_log_outputs,
-            enable_log_deltas=args.enable_log_deltas,
-        )
-        if "generate" in supported_tasks
-        else None
+
+    _chat_kwargs = dict(
+        engine_client=engine_client,
+        models=state.openai_serving_models,
+        response_role=args.response_role,
+        online_renderer=state.online_renderer,
+        request_logger=request_logger,
+        chat_template=resolved_chat_template,
+        chat_template_content_format=args.chat_template_content_format,
+        default_chat_template_kwargs=args.default_chat_template_kwargs,
+        trust_request_chat_template=args.trust_request_chat_template,
+        return_tokens_as_token_ids=args.return_tokens_as_token_ids,
+        enable_auto_tools=args.enable_auto_tool_choice,
+        exclude_tools_when_tool_choice_none=args.exclude_tools_when_tool_choice_none,
+        tool_parser=args.tool_call_parser,
+        reasoning_parser=args.structured_outputs_config.reasoning_parser,
+        enable_prompt_tokens_details=args.enable_prompt_tokens_details,
+        enable_force_include_usage=args.enable_force_include_usage,
+        enable_log_outputs=args.enable_log_outputs,
+        enable_log_deltas=args.enable_log_deltas,
     )
+
+    state.openai_serving_chat = OmniOpenAIServingChat(**_chat_kwargs) if "generate" in supported_tasks else None
+    state.openai_serving_chat_batch = (
+        OmniOpenAIServingChatBatch(**_chat_kwargs) if "generate" in supported_tasks else None
+    )
+
     # Warm up chat template processing to avoid first-request latency
     if state.openai_serving_chat is not None:
         state.openai_serving_chat.warmup()
@@ -1162,6 +1172,10 @@ def Omnichat(request: Request) -> OmniOpenAIServingChat | None:
     return request.app.state.openai_serving_chat
 
 
+def OmniBatchChat(request: Request) -> OmniOpenAIServingChatBatch | None:
+    return request.app.state.openai_serving_chat_batch
+
+
 def Omnispeech(request: Request) -> OmniOpenAIServingSpeech | None:
     return request.app.state.openai_serving_speech
 
@@ -1245,6 +1259,45 @@ async def create_chat_completion(request: ChatCompletionRequest, raw_request: Re
     return StreamingResponse(content=generator, media_type="text/event-stream")
 
 
+@router.post(
+    "/v1/chat/completions/batch",
+    dependencies=[Depends(validate_json_request)],
+    responses={
+        HTTPStatus.OK.value: {},
+        HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse},
+        HTTPStatus.NOT_FOUND.value: {"model": ErrorResponse},
+        HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse},
+        HTTPStatus.NOT_IMPLEMENTED.value: {"model": ErrorResponse},
+    },
+)
+@with_cancellation
+@load_aware_call
+async def create_batch_chat_completion(request: BatchChatCompletionRequest, raw_request: Request):
+    handler = OmniBatchChat(raw_request)
+    if handler is None:
+        base_server = getattr(raw_request.app.state, "serving_tokenization", None)
+        if base_server is None:
+            raise HTTPException(
+                status_code=HTTPStatus.NOT_FOUND.value,
+                detail="The model does not support Chat Completions API",
+            )
+        return base_server.create_error_response(message="The model does not support Chat Completions API")
+    try:
+        result = await handler.create_batch_chat_completion(request, raw_request)
+    except (EngineGenerateError, EngineDeadError) as exc:
+        return _create_engine_error_json_response(raw_request, exc)
+    except Exception as e:
+        logger.exception("Batched chat completion failed: %s", e)
+        raise HTTPException(status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, detail=str(e)) from e
+
+    if isinstance(result, ErrorResponse):
+        return JSONResponse(
+            content=result.model_dump(),
+            status_code=result.error.code if result.error else 400,
+        )
+    return JSONResponse(content=result.model_dump(mode="json"))
+
+
 _remove_route_from_router(router, "/v1/audio/speech", {"POST"})
 
 
@@ -2811,6 +2864,7 @@ def video_response_from_request(model_name: str, req: VideoGenerationRequest) ->
         status=VideoGenerationStatus.QUEUED,
         size=req.size,
         prompt=req.prompt,
+        quality=req.quality or "default",
     )
     resp.seconds = str(req.seconds or resp.seconds)
     return resp
@@ -3150,6 +3204,7 @@ async def _parse_video_form(
     short_edge: int | None = Form(default=None, ge=1),
     num_outputs_per_prompt: int = Form(default=1, ge=1, le=10),
     start_time_seconds: float | None = Form(default=None, ge=0.0),
+    quality: str | None = Form(default=None),
     num_inference_steps: int | None = Form(default=None),
     guidance_scale: float | None = Form(default=None),
     guidance_scale_2: float | None = Form(default=None),
@@ -3218,6 +3273,7 @@ async def _parse_video_form(
         "short_edge": short_edge,
         "num_outputs_per_prompt": num_outputs_per_prompt,
         "start_time_seconds": start_time_seconds,
+        "quality": quality,
         "num_inference_steps": num_inference_steps,
         "guidance_scale": guidance_scale,
         "guidance_scale_2": guidance_scale_2,

vllm_omni/entrypoints/openai/batch_serving.py

@@ -0,0 +1,151 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import asyncio
+import time
+
+from fastapi import Request
+from pydantic import ValidationError
+from vllm.entrypoints.openai.chat_completion.protocol import (
+    BatchChatCompletionRequest,
+    ChatCompletionRequest,
+    ChatCompletionResponse,
+    ChatCompletionResponseChoice,
+)
+from vllm.entrypoints.openai.engine.protocol import (
+    ErrorResponse,
+    UsageInfo,
+)
+from vllm.logger import init_logger
+
+from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat
+
+logger = init_logger(__name__)
+
+
+class OmniOpenAIServingChatBatch(OmniOpenAIServingChat):
+    @staticmethod
+    def _maybe_collapse_choices(choices: list[ChatCompletionResponseChoice]):
+        """Collapse content + audio choices into one choice.
+
+        NOTE: This is a hack for models that produce separate text and
+        audio choices to ensure we have 1:1 mapping between number of
+        inputs in the batch and responses.
+        """
+        content_choice, audio_choice = None, None
+        if len(choices) == 1:
+            return choices[0]
+        if len(choices) != 2:
+            raise ValueError(f"Expected 1 or 2 choices to collapse, got {len(choices)}. ")
+
+        for choice in choices:
+            if choice.message.audio:
+                if audio_choice is not None:
+                    raise ValueError("Multiple audio choices cannot be set in one completion message")
+                audio_choice = choice
+            else:
+                if content_choice is not None:
+                    raise ValueError("Multiple content choices cannot be set in one completion message")
+                content_choice = choice
+
+        if content_choice is None or audio_choice is None:
+            raise ValueError(
+                "Expected one content and one audio choice, but got: "
+                f"content={content_choice is not None}, audio={audio_choice is not None}"
+            )
+        content_choice.message.audio = audio_choice.message.audio
+        return content_choice
+
+    @staticmethod
+    def _get_subrequest_ids(base_id, request: BatchChatCompletionRequest) -> list[str]:
+        """Get the request ID for each entry in the batch request to be resubmitted to chat completions."""
+        return [f"{base_id}-idx-{idx}" for idx in range(len(request.messages))]
+
+    async def create_batch_chat_completion(
+        self,
+        request: BatchChatCompletionRequest,
+        raw_request: Request,
+    ) -> ChatCompletionResponse | ErrorResponse:
+        """Given a request, submit each request to chat completions & collect the results."""
+        model = ""
+        enabled_streaming = False
+        chat_requests: list[ChatCompletionRequest] = []
+        choices: list[ChatCompletionResponseChoice] = []
+        total_prompt_tokens = 0
+        total_completion_tokens = 0
+
+        # Get the base request from the raw_request before maybe mutating the header
+        base_id = OmniOpenAIServingChatBatch._base_request_id(
+            raw_request,
+            request.request_id,
+        )
+        batch_req_id = f"chatcmpl-batch-{base_id}"
+        sub_req_ids = self._get_subrequest_ids(base_id, request)
+
+        # Remove the raw request header if it exists now that we have subreq IDs,
+        # since this takes priority over the object request_ids
+        if raw_request.headers.get("X-Request-Id"):
+            copy_headers = raw_request.headers.mutablecopy()
+            del copy_headers["X-Request-Id"]
+            raw_request._headers = copy_headers
+
+        for idx, msg in enumerate(request.messages):
+            # Update the current request ID to avoid subrequest collisions
+            request.request_id = sub_req_ids[idx]
+            try:
+                chat_cmp_request = request.to_chat_completion_request(msg)
+            except ValidationError as e:
+                return self._create_error_response(
+                    f"Message {idx} could not be converted to a chat completion request: {e}",
+                )
+            # Streaming isn't supported for batch chat completions,
+            # so always set to False & warn if it was requested.
+            enabled_streaming |= chat_cmp_request.stream
+            chat_cmp_request.stream = False
+            chat_requests.append(chat_cmp_request)
+
+        if enabled_streaming:
+            logger.warning("Streaming is not supported for batched chat completions; ignoring stream=True.")
+
+        # Submit each chat completion request as a task, then gather results.
+        # TODO (Alex): optimize this
+        tasks = [asyncio.create_task(self.create_chat_completion(c, raw_request)) for c in chat_requests]
+        try:
+            results = await asyncio.gather(*tasks)
+        finally:
+            # Ensure we cancel remaining tasks if needed, e.g., early exit due to bad behavior
+            for t in tasks:
+                if not t.done():
+                    t.cancel()
+
+        for i, resp in enumerate(results):
+            if isinstance(resp, ErrorResponse):
+                return resp
+            completion: ChatCompletionResponse = resp
+            model = completion.model
+            # FIXME (Alex): We should probably handle this in chat completions,
+            # not here, but we need to ensure streaming is properly handled
+            try:
+                collapsed_choice = self._maybe_collapse_choices(completion.choices)
+            except ValueError as e:
+                return self._create_error_response(
+                    f"Failed to collapse choices with error: {e}",
+                )
+            collapsed_choice.index = i
+            choices.append(collapsed_choice)
+            if completion.usage:
+                total_prompt_tokens += completion.usage.prompt_tokens
+                total_completion_tokens += completion.usage.completion_tokens
+
+        usage = UsageInfo(
+            prompt_tokens=total_prompt_tokens,
+            completion_tokens=total_completion_tokens,
+            total_tokens=total_prompt_tokens + total_completion_tokens,
+        )
+
+        return ChatCompletionResponse(
+            id=batch_req_id,
+            created=int(time.time()),
+            model=model,
+            choices=choices,
+            usage=usage,
+        )

vllm_omni/entrypoints/openai/diffusion_request_utils.py

@@ -7,7 +7,10 @@ from collections.abc import Collection, Mapping
 
 from vllm.logger import init_logger
 
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.inputs.data import (
+    DIFFUSION_QUALITY_LEVELS,
+    OmniDiffusionSamplingParams,
+)
 
 logger = init_logger(__name__)
 
@@ -90,6 +93,9 @@ def normalize_diffusion_request_args(
     for alias, canonical_name in aliases.items():
         if alias in request_args:
             request_args[canonical_name] = request_args.pop(alias)
+    quality = request_args.get("quality")
+    if quality is not None and quality not in DIFFUSION_QUALITY_LEVELS:
+        raise ValueError(f"quality must be one of {list(DIFFUSION_QUALITY_LEVELS)}, got {quality!r}")
     return normalized_extra_args, request_args
 
 

vllm_omni/entrypoints/openai/protocol/videos.py

@@ -15,9 +15,10 @@ from enum import Enum
 from functools import lru_cache
 from typing import Annotated, Any, Literal
 
-from pydantic import BaseModel, ConfigDict, Field, StringConstraints
+from pydantic import BaseModel, ConfigDict, Field, StringConstraints, field_validator
 
 from vllm_omni.entrypoints.openai.image_api_utils import parse_size
+from vllm_omni.inputs.data import DIFFUSION_QUALITY_LEVELS
 
 
 class VideoGenerationStatus(str, Enum):
@@ -164,6 +165,24 @@ class VideoGenerationRequest(BaseModel):
     )
 
     # vllm-omni extensions for diffusion control
+    quality: str | None = Field(
+        default=None,
+        description=(
+            "Request-level generation quality intent. Supported values are "
+            "'lossless' and 'high'; exact behavior is model-specific. "
+            "When omitted, the model chooses its default policy."
+        ),
+    )
+
+    @field_validator("quality")
+    @classmethod
+    def validate_quality(cls, value: str | None) -> str | None:
+        if value is None:
+            return None
+        if value not in DIFFUSION_QUALITY_LEVELS:
+            raise ValueError(f"quality must be one of {list(DIFFUSION_QUALITY_LEVELS)}, got {value!r}")
+        return value
+
     negative_prompt: str | None = Field(default=None, description="Text describing what to avoid in the video")
     num_inference_steps: int | None = Field(
         default=None,

vllm_omni/entrypoints/openai/serving_chat.py

@@ -170,6 +170,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
             "width",
             "num_outputs_per_prompt",
             "seed",
+            "quality",
             "num_inference_steps",
             "guidance_scale",
             "true_cfg_scale",
@@ -759,6 +760,11 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                     if hasattr(sp, "num_inference_steps") and num_inference_steps is not None:
                         sp.num_inference_steps = num_inference_steps
                     if isinstance(sp, OmniDiffusionSamplingParams):
+                        quality = extra_body.get("quality")
+                        self._set_if_supported(
+                            sp,
+                            quality=quality,
+                        )
                         apply_normalized_diffusion_request_extra_args(sp, normalized_extra_args)
                     else:
                         apply_declared_extra_args(
@@ -2937,6 +2943,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
         generator_device = gen_params.generator_device
         num_outputs_per_prompt = gen_params.num_outputs_per_prompt
         num_inference_steps = extra_body.get("num_inference_steps")
+        quality = extra_body.get("quality")
         guidance_scale = extra_body.get("guidance_scale")
         true_cfg_scale = extra_body.get("true_cfg_scale") or extra_body.get("cfg_scale")
         negative_prompt = extra_body.get("negative_prompt")
@@ -3087,6 +3094,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                     seed=seed,
                     generator_device=generator_device,
                     num_outputs_per_prompt=num_outputs_per_prompt,
+                    quality=quality,
                     num_inference_steps=num_inference_steps,
                     guidance_scale=guidance_scale,
                     true_cfg_scale=true_cfg_scale,
@@ -3151,6 +3159,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
         self._set_if_supported(
             gen_params,
             generator_device=generator_device,
+            quality=extra_body.get("quality"),
             num_inference_steps=extra_body.get("num_inference_steps"),
             guidance_scale=extra_body.get("guidance_scale"),
             true_cfg_scale=extra_body.get("true_cfg_scale") or extra_body.get("cfg_scale"),
@@ -3489,6 +3498,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
             # method apply its own model-specific default when the user does
             # not provide a value.
             num_inference_steps = extra_body.get("num_inference_steps")
+            quality = extra_body.get("quality")
             guidance_scale = extra_body.get("guidance_scale")
             true_cfg_scale = extra_body.get("true_cfg_scale")
             seed = extra_body.get("seed")
@@ -3543,6 +3553,8 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
             # Only override defaults when the user explicitly provides values
             if num_inference_steps is not None:
                 gen_params.num_inference_steps = num_inference_steps
+            if quality is not None:
+                gen_params.quality = quality
             if guidance_scale is not None:
                 gen_params.guidance_scale = guidance_scale
             if true_cfg_scale is not None:

vllm_omni/entrypoints/openai/serving_video.py

@@ -200,6 +200,8 @@ class OmniOpenAIServingVideo:
 
         if "num_inference_steps" in provided_fields and request.num_inference_steps is not None:
             gen_params.num_inference_steps = request.num_inference_steps
+        if "quality" in provided_fields:
+            gen_params.quality = request.quality
         if "guidance_scale" in provided_fields and request.guidance_scale is not None:
             gen_params.guidance_scale = request.guidance_scale
         if "guidance_scale_2" in provided_fields and request.guidance_scale_2 is not None:

vllm_omni/entrypoints/openai/serving_video_output_stream.py

@@ -535,6 +535,8 @@ class OmniStreamingVideoOutputHandler:
         provided_fields = request.model_fields_set
         if "num_inference_steps" in provided_fields and request.num_inference_steps is not None:
             gen_params.num_inference_steps = request.num_inference_steps
+        if "quality" in provided_fields:
+            gen_params.quality = request.quality
         if "guidance_scale" in provided_fields and request.guidance_scale is not None:
             gen_params.guidance_scale = request.guidance_scale
         if "guidance_scale_2" in provided_fields and request.guidance_scale_2 is not None:

vllm_omni/inputs/data.py

@@ -11,6 +11,8 @@ from vllm.sampling_params import SamplingParams
 
 from vllm_omni.lora.request import LoRARequest
 
+DIFFUSION_QUALITY_LEVELS: tuple[str, ...] = ("lossless", "high")
+
 
 class OmniTextPrompt(TextPrompt):
     """Text prompt with optional embeddings and stage payloads.
@@ -208,6 +210,10 @@ class OmniDiffusionSamplingParams:
     do_classifier_free_guidance: bool = False
     output_type: str | None = None
 
+    # Request-scoped quality intent. ``None`` delegates the default behavior
+    # to the model; explicit levels select a model-owned quality policy.
+    quality: str | None = None
+
     # Batch info
     num_outputs_per_prompt: int = 1
     seed: int | None = None
@@ -337,6 +343,10 @@ class OmniDiffusionSamplingParams:
     # results
     output: torch.Tensor | None = None
 
+    def __post_init__(self) -> None:
+        if self.quality is not None and self.quality not in DIFFUSION_QUALITY_LEVELS:
+            raise ValueError(f"quality must be one of {list(DIFFUSION_QUALITY_LEVELS)}, got {self.quality!r}")
+
     @property
     def batch_size(self):
         # This class is changed to only represent a single prompt request