docs: define module ownership contracts

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

docs/design/module/archive/README.md

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

docs/design/module/archive/ar_module.md

@@ -0,0 +1,387 @@
+# AutoRegressive (AR) Module
+
+## 1. Overview
+
+The AutoRegressive (AR) module in vLLM-Omni handles autoregressive generation stages, primarily used for text, chain-of-thought(COT), and audio latent tokens generation stages in multi-stage models like Qwen2.5-Omni, Qwen3-Omni, BAGEL, .etc. Unlike some representative non-autoregressive generation stages (e.g., Diffusion), AR stages generate tokens sequentially, one at a time, following the standard transformer decoder pattern.
+
+The AR module of vLLM-Omni extends vLLM's core components to support:
+
+- **Multimodal inputs/outputs**: Processing images, videos, and audio alongside text
+- **Direct embedding transfer**: Passing pre-computed prompt embeddings between pipeline stages via serialized payloads
+- **Additional information flow**: Carrying per-request metadata (tensors, lists) through the pipeline
+- **Hidden state exposure**: Exposing per-request hidden representations for downstream stages
+- **Basic generator support**: Support some basic heterogeneous architecture such as Convolution, LSTM, etc.
+
+As shown in the [end2end example](../../../user_guide/examples/offline_inference/qwen3_omni.md), AR module can be widely applied across multiple stages, generating text tokens in thinker(AR), audio latent tokens in talker(AR) and audio wave in code2wav(Convolution).
+
+## 2. Relationship with vLLM
+
+The AR module builds upon vLLM main framework through inheritance, extending core classes while preserving compatibility with vLLM's scheduling, batching, KV cache management, and execution mechanisms.
+
+### Inheritance Hierarchy
+- Scheduler
+
+```mermaid
+classDiagram
+    class VLLMScheduler {
+        +schedule() SchedulerOutput
+        +update_from_output() EngineCoreOutputs
+    }
+    class OmniARScheduler {
+        +schedule() SchedulerOutput
+    }
+    class OmniGenerationScheduler {
+        +schedule() SchedulerOutput
+        +update_from_output() EngineCoreOutputs
+    }
+    VLLMScheduler <|-- OmniARScheduler
+    VLLMScheduler <|-- OmniGenerationScheduler
+```
+- Worker
+
+```mermaid
+classDiagram
+    class GPUWorker {
+        +init_device()
+        +model_runner
+    }
+    class GPUARWorker {
+        +init_device()
+    }
+    class GPUGenerationWorker {
+        +init_device()
+    }
+    GPUWorker <|-- GPUARWorker
+    GPUWorker <|-- GPUGenerationWorker
+```
+- ModelRunner
+
+```mermaid
+classDiagram
+    class GPUModelRunner {
+        +execute_model()
+        +sample_tokens()
+    }
+    class OmniGPUModelRunner {
+        +_update_states()
+        +_preprocess()
+        +_model_forward()
+    }
+    class GPUARModelRunner {
+        +execute_model()
+        +sample_tokens()
+    }
+    class GPUGenerationModelRunner {
+        +execute_model()
+    }
+    GPUModelRunner <|-- OmniGPUModelRunner
+    OmniGPUModelRunner <|-- GPUARModelRunner
+    OmniGPUModelRunner <|-- GPUGenerationModelRunner
+```
+- InputProcessor/OutputProcessor
+
+```mermaid
+classDiagram
+    class InputProcessor {
+        +process_inputs() EngineCoreRequest
+    }
+
+    class VLLMOutputProcessor {
+        +process_outputs() OutputProcessorOutput
+    }
+    class MultimodalOutputProcessor {
+        +process_outputs() OutputProcessorOutput
+        +_route_and_normalize()
+    }
+    VLLMOutputProcessor <|-- MultimodalOutputProcessor
+```
+
+### Key Extensions
+
+- **Scheduler**: `OmniARScheduler` extends `vllm.v1.core.sched.scheduler.Scheduler` to enrich scheduled requests with omni-specific payloads
+- **Worker**: `GPUARWorker` extends `vllm.v1.worker.gpu_worker.Worker` to initialize AR-specific model runners
+- **ModelRunner**: `GPUARModelRunner` extends `OmniGPUModelRunner` → `vllm.v1.worker.gpu_model_runner.GPUModelRunner` to expose hidden states and handle multimodal outputs
+- **InputProcessor**: Stage-0 uses upstream `vllm.v1.engine.input_processor.InputProcessor`; `AsyncOmniEngine` then restores omni-specific payloads (for example `additional_information` and `prompt_embeds`) when building `OmniEngineCoreRequest`
+- **OutputProcessor**: `MultimodalOutputProcessor` extends `vllm.v1.engine.output_processor.OutputProcessor` to route and accumulate multimodal outputs
+
+## 3. Scheduler Design
+
+The AR module provides two scheduler implementations: one for standard autoregressive generation and one for basic heterogeneous architectures.
+
+### Request Flow
+
+The following diagram illustrates the request flow through the AR module components:
+
+```mermaid
+flowchart TD
+    A[InputProcessor stage-0 in AsyncOmniEngine] -->|EngineCoreRequest then upgraded to OmniEngineCoreRequest| B[OmniARScheduler]
+    B -->|schedule: OmniNewRequestData| C[GPUARWorker]
+    C -->|SchedulerOutput| D[GPUARModelRunner]
+    D -->|execute_model: None| E[Model Forward Pass]
+    E -->|hidden_states, logits| D
+    D -->|sample_tokens: OmniModelRunnerOutput| F[OmniARScheduler]
+    F -->|update_from_output| G[MultimodalOutputProcessor]
+    G -->|RequestOutput| H[Client/Downstream Stage]
+
+    style A fill:#e1f5ff
+    style B fill:#fff4e1
+    style C fill:#e8f5e9
+    style D fill:#f3e5f5
+    style G fill:#fce4ec
+```
+
+The flow follows vLLM's standard pattern: input processing → scheduling → worker execution → output processing, with omni-specific enrichments at each stage.
+
+### OmniARScheduler
+
+`OmniARScheduler` extends the base vLLM scheduler with minimal modifications, focusing on enriching scheduled requests with omni-specific payloads.
+
+#### Modified API: `schedule()`
+
+The scheduler wraps base `NewRequestData` entries with `OmniNewRequestData` to include prompt embeddings and additional information:
+
+```python
+def schedule(self) -> SchedulerOutput:
+    scheduler_output = super().schedule()
+    # Rewrap base NewRequestData entries with OmniNewRequestData
+    new_list = []
+    for nr in scheduler_output.scheduled_new_reqs:
+        request = self.requests.get(nr.req_id)
+        omni_nr = OmniNewRequestData(
+            req_id=nr.req_id,
+            prompt_token_ids=nr.prompt_token_ids,
+            # ... other base fields ...
+            prompt_embeds=getattr(request, "prompt_embeds", None),
+            additional_information=getattr(request, "additional_information", None),
+        )
+        new_list.append(omni_nr)
+    scheduler_output.scheduled_new_reqs = new_list
+    return scheduler_output
+```
+
+The `update_from_output()` method remains unchanged, inheriting standard request lifecycle management from the base scheduler.
+
+### OmniGenerationScheduler
+
+`OmniGenerationScheduler` implements a fast-path scheduling strategy for basic heterogeneous architectures that process all input tokens in a single step.
+
+#### Modified API: `schedule()`
+
+Allocates all input tokens for a request at once (or 1 placeholder if zero), falling back to default scheduling if budget is insufficient:
+
+```python
+def schedule(self) -> SchedulerOutput:
+    # Fast path: allocate all input tokens at once
+    while self.waiting and token_budget > 0:
+        request = self.waiting.peek_request()
+        required_tokens = max(getattr(request, "num_prompt_tokens", 0), 1)
+        if required_tokens > token_budget:
+            break  # Fall back to default scheduling
+        # Allocate and schedule...
+```
+
+#### Modified API: `update_from_output()`
+
+Marks requests as finished immediately after one step, since generation models complete in a single forward pass:
+
+```python
+def update_from_output(self, ...) -> dict[int, EngineCoreOutputs]:
+    # ...
+    # Diffusion request: completes in one step
+    request.status = RequestStatus.FINISHED_STOPPED
+    kv_transfer_params = self._free_request(request)
+    # ...
+```
+
+## 4. Worker and ModelRunner Design
+
+### GPUARWorker
+
+`GPUARWorker` initializes the AR-specific model runner while maintaining standard device initialization:
+
+```python
+class GPUARWorker(GPUWorker):
+    def init_device(self):
+        # ... standard device initialization ...
+        self.model_runner = GPUARModelRunner(self.vllm_config, self.device)
+```
+
+### GPUARModelRunner
+
+`GPUARModelRunner` follows vLLM's two-phase execute/sample flow while exposing hidden states and multimodal outputs.
+
+#### Two-Phase Execution
+
+**Phase 1: `execute_model()`** - Runs forward pass and stores state:
+- Computes logits from hidden states
+- Stores `ExecuteModelState` with hidden states, logits, and multimodal outputs
+- Returns `None` to defer sampling
+
+**Phase 2: `sample_tokens()`** - Samples tokens and builds output:
+- Retrieves stored state from `execute_model()`
+- Samples tokens using logits
+- Extracts per-request hidden states and multimodal outputs
+- Builds `OmniModelRunnerOutput` with `pooler_output` containing hidden states
+
+```python
+def sample_tokens(self, grammar_output) -> OmniModelRunnerOutput:
+    # Retrieve stored state
+    hidden_states, multimodal_outputs = self.execute_model_state
+
+    # Sample tokens
+    sampler_output = self._sample(logits, spec_decode_metadata)
+
+    # Extract per-request hidden states
+    pooler_output = []
+    for rid in req_ids:
+        hidden_slice = hidden_states_cpu[start:end]
+        payload = {"hidden": hidden_slice}
+        # Add multimodal outputs if present
+        pooler_output.append(payload)
+
+    return OmniModelRunnerOutput(
+        pooler_output=pooler_output,
+        # ... other fields ...
+    )
+```
+
+### GPUGenerationModelRunner
+
+`GPUGenerationModelRunner` implements a simplified single-phase execution for basic heterogeneous architectures:
+
+- No logits computation or token sampling
+- Direct generation from forward pass in model implementation
+- Returns outputs via `pooler_output` immediately after forward pass
+
+### OmniGPUModelRunner
+
+`OmniGPUModelRunner` provides shared functionality for both AR and Generation runners:
+
+#### Prompt Embeddings Overlay
+
+During prefill, overlays custom `prompt_embeds` from request state onto `inputs_embeds`:
+
+```python
+def _collect_additional_information_for_prefill(self, num_scheduled_tokens_np):
+    for req_index, req_id in enumerate(self.input_batch.req_ids):
+        req_state = self.requests[req_id]
+        pe_cpu = getattr(req_state, "prompt_embeds_cpu", None)
+        # Overlay prompt_embeds for prefill portion
+        if pe_cpu is not None:
+            src = pe_cpu[num_computed_tokens:num_computed_tokens + overlay_len]
+            self.inputs_embeds[start_offset:start_offset + overlay_len].copy_(src)
+```
+
+#### Additional Information Processing
+
+Decodes and manages `additional_information` payloads:
+- Decodes serialized payloads → CPU tensors in request state
+- Passes runtime information to model via `runtime_additional_information` kwarg
+- Processes model-provided updates via `postprocess()` hook
+- Merges updates back into request state
+
+#### M-RoPE Position Initialization
+
+For multimodal models using M-RoPE (e.g., Qwen2-VL), computes position encodings from multimodal feature metadata (image grids, video grids, audio features).
+
+## 5. Input/Output Processing
+
+### Processing Pipeline
+
+The input/output processing pipeline handles serialization, routing, and accumulation of multimodal data:
+
+```mermaid
+sequenceDiagram
+    participant Client
+    participant AsyncOmniEngine
+    participant InputProcessor
+    participant Scheduler
+    participant ModelRunner
+    participant MultimodalOutputProcessor
+    participant Client
+
+    Client->>AsyncOmniEngine: prompt + prompt_embeds + additional_info
+    AsyncOmniEngine->>InputProcessor: process_inputs()
+    InputProcessor->>Scheduler: EngineCoreRequest
+    AsyncOmniEngine->>AsyncOmniEngine: _upgrade_to_omni_request() + serialize_additional_information()
+    Scheduler->>ModelRunner: OmniNewRequestData (with payloads)
+    ModelRunner->>ModelRunner: Decode payloads → CPU tensors
+    ModelRunner->>ModelRunner: Overlay prompt_embeds on inputs_embeds
+    ModelRunner->>ModelRunner: Forward pass with runtime_additional_information
+    ModelRunner->>ModelRunner: Extract hidden states + multimodal outputs
+    ModelRunner->>MultimodalOutputProcessor: OmniModelRunnerOutput (pooler_output)
+    MultimodalOutputProcessor->>MultimodalOutputProcessor: Route by output_type
+    MultimodalOutputProcessor->>MultimodalOutputProcessor: Accumulate tensors in OmniRequestState
+    MultimodalOutputProcessor->>MultimodalOutputProcessor: Consolidate tensor lists
+    MultimodalOutputProcessor->>Client: RequestOutput (with multimodal_output)
+```
+
+### Stage-0 Input Processing
+
+Stage-0 now uses upstream `InputProcessor` directly, and `AsyncOmniEngine` upgrades the request to `OmniEngineCoreRequest` while restoring omni-specific payloads.
+
+```python
+request = self.input_processor.process_inputs(
+    request_id=request_id,
+    prompt=prompt,
+    params=params,
+    supported_tasks=self.supported_tasks,
+)
+request = _upgrade_to_omni_request(request, prompt)
+```
+
+### MultimodalOutputProcessor
+
+`MultimodalOutputProcessor` routes outputs by modality type and accumulates multimodal tensors.
+
+#### Output Routing
+
+Routes `EngineCoreOutput` by `output_type` attribute:
+- `"text"`: Standard text generation path
+- `"image"`, `"audio"`, `"latents"`: Extract from `pooling_output` or `multimodal_outputs`
+- Fallback: Heuristic based on presence of `pooling_output`
+
+#### Tensor Accumulation
+
+`OmniRequestState` accumulates multimodal tensors across multiple steps:
+
+```python
+def add_multimodal_tensor(self, payload, mm_type):
+    # Normalize payload to dict
+    incoming = {mm_type or "hidden": payload}
+
+    # Accumulate: convert tensors to lists for deferred concatenation
+    if isinstance(v, torch.Tensor) and isinstance(existing, torch.Tensor):
+        self.mm_accumulated[k] = [existing, v]  # List accumulation
+```
+
+Before final output, consolidates tensor lists via concatenation:
+
+```python
+def _consolidate_multimodal_tensors(self):
+    for k, v in self.mm_accumulated.items():
+        if isinstance(v, list) and isinstance(v[0], torch.Tensor):
+            self.mm_accumulated[k] = torch.cat(v, dim=0)  # Concatenate
+```
+
+The consolidated tensors are attached to `RequestOutput.multimodal_output` for consumption by downstream stages or clients.
+
+## 6. Summary
+
+The AR module of vLLM-Omni extends vLLM through strategic inheritance and minimal API modifications:
+
+### Key Design Patterns
+
+1. **Inheritance over composition**: Extends vLLM classes to preserve compatibility with existing scheduling, batching, and execution mechanisms
+2. **Payload serialization**: Uses serialized `additional_information` payloads together with prompt-embedding handoff for efficient inter-stage transfer
+3. **Two-phase execution**: Maintains vLLM's execute/sample separation for AR models while supporting single-phase execution for generation models
+4. **Multimodal routing**: Routes outputs by `output_type` and accumulates tensors incrementally to support streaming
+
+### Differences from vLLM
+
+- **Payload support**: Serialized additional information and prompt embeddings enable direct transfer between pipeline stages
+- **Multimodal handling**: Extended input/output processors support images, audio, and other modalities alongside text
+- **Hidden state exposure**: AR model runners expose per-request hidden states via `pooler_output` for downstream consumption
+- **Generation scheduler**: Fast-path scheduling for basic heterogeneous architectures that complete in one step
+
+The AR module seamlessly integrates with vLLM's existing infrastructure while adding the necessary extensions for multi-stage, multimodal generation pipelines.

docs/design/module/archive/async_omni_architecture.md

@@ -0,0 +1,203 @@
+# AsyncOmni Architecture (Qwen3-Omni Example)
+
+## 1. System Architecture
+
+```text
+• ┌─────────────────────────────────────────────────────────────────────────────────┐
+  │                                    API Layer                                    │
+  │  ┌─────────────────────────────────────┐  ┌──────────────────────────────────┐  │
+  │  │ AsyncOmni (EngineClient)            │  │ Omni                             │  │
+  │  │ • generate() / abort() / shutdown() │  │ • generate()                     │  │
+  │  │ • _final_output_handler()           │  │                                  |  │
+  │  └─────────────────────────────────────┘  └──────────────────────────────────┘  │
+  ├─────────────────────────────────────────────────────────────────────────────────┤
+  │                              Engine Layer (Proxy)                               │
+  │  ┌───────────────────────────────────────────────────────────────────────────┐  │
+  │  │ AsyncOmniEngine                                                           │  │
+  │  │ • _bootstrap_orchestrator() & _initialize_stages()                        │  │
+  │  │ • add_request() / add_request_async() -> input_processor.process_inputs() │  │
+  │  │ • try_get_output() / try_get_output_async()                               │  │
+  │  └───────────────────┬─────────────────────────────────▲─────────────────────┘  │
+  │         request_queue (janus.Queue)        output_queue (janus.Queue)           │
+  ├──────────────────────┼─────────────────────────────────┼────────────────────────┤
+  │                      ▼        Orchestration Layer      │                        │
+  │  ┌───────────────────────────────────────────────────────────────────────────┐  │
+  │  │ Orchestrator [background thread]                                          │  │
+  │  │ • _request_handler()                                                      │  │
+  │  │     -  stage_client.add_request_async() & _prewarm_async_chunk_stages()   │  │
+  │  │ • _orchestration_output_handler()                                         │  │
+  │  │     -  _process_stage_outputs() -> output_processors[i].process_outputs() │  │
+  │  │     -  _route_output() & _forward_to_next_stage()                         │  │
+  │  └──────────┬─────────────────────────┬────────────────────────┬─────────────┘  │
+  ├─────────────┼─────────────────────────┼────────────────────────┼────────────────┤
+  │             │                 Communication Layer              │                │
+  │  ┌───────────────────────┐ ┌───────────────────────┐ ┌───────────────────────┐  │
+  │  │ StageEngineCoreClient │ │ StageEngineCoreClient │ │ StageDiffusionClient  │  │
+  │  │ • ZMQ ROUTER / PULL   │ │ • ZMQ ROUTER / PULL   │ │ • ZMQ ROUTER / PULL   │  │
+  │  │ • Msgpack codec       │ │ • Msgpack codec       │ │ • Msgpack codec       │  │
+  │  └──────────┬────────────┘ └──────────┬────────────┘ └──────────┬────────────┘  │
+  │             ▼ ZMQ IPC                 ▼ ZMQ IPC                 ▼ ZMQ IPC       │
+  ├─────────────────────────────────────────────────────────────────────────────────┤
+  │                                 Execution Layer                                 │
+  │  ┌───────────────────────┐ ┌───────────────────────┐ ┌───────────────────────┐  │
+  │  │ StageCoreProc         │ │ StageCoreProc         │ │ DiffusionEngine       │  │
+  │  │ [background process]  │ │ [background process]  │ │ [background process]  │  │
+  │  └───────────────────────┘ └───────────────────────┘ └───────────────────────┘  │
+  └─────────────────────────────────────────────────────────────────────────────────┘
+```
+
+## 2. Execution Flow (Arrow Steps, one generate request)
+
+```text
+[1] App
+    -> AsyncOmni.generate(prompt, request_id)
+
+[2] AsyncOmni
+    -> _final_output_handler()   (started on first request)
+    -> AsyncOmniEngine.add_request(stage_id=0, ...)
+
+[3] AsyncOmniEngine.add_request
+    -> (if stage-0 is llm and input is not EngineCoreRequest)
+       InputProcessor.process_inputs()
+       OutputProcessor[0].add_request()
+    -> request_queue.put(add_request_msg)
+
+[4] Orchestrator._request_handler
+    -> _handle_add_request(msg)
+    -> stage_clients[0].add_request_async(...)
+
+[5] Orchestrator._orchestration_loop (loop)
+    -> poll stage output
+       - llm stage: await get_output_async()
+       - diffusion stage: get_diffusion_output_nowait()
+    -> (llm stage) output_processors[i].process_outputs(...)
+    -> _route_output(...)
+    -> if finished and not final_stage and non-async-chunk:
+         _forward_to_next_stage(...)
+         -> next_stage.add_request_async(...)
+    -> output_queue.put(output)
+
+[6] AsyncOmni._final_output_loop (background coroutine)
+    -> AsyncOmniEngine.try_get_output_async()
+    -> route by request_id to ClientRequestState.queue
+
+[7] AsyncOmni._process_orchestrator_results
+    -> read from ClientRequestState.queue
+    -> _process_single_result(...)
+    -> yield OmniRequestOutput
+
+[8] Exit condition
+    -> receive result["finished"] == True
+    -> generate() ends
+```
+
+## 3. Runtime Sequence (one generate request)
+
+```mermaid
+sequenceDiagram
+    participant APP as App
+    participant AO as AsyncOmni
+    participant ENG as AsyncOmniEngine
+    participant ORCH as Orchestrator
+    participant S0 as Stage-0 Client
+    participant SN as Next Stage Client
+
+    APP->>AO: generate
+    AO->>AO: start output_handler once
+    AO->>ENG: add_request(stage_id=0, ...)
+    ENG->>ENG: input_processor.process_inputs()
+    ENG->>ORCH: request_queue.put(add_request)
+
+    ORCH->>ORCH: _handle_add_request
+    ORCH->>S0: add_request_async
+
+    loop poll route forward
+        ORCH->>S0: get_output_async / get_diffusion_output_nowait
+        ORCH->>ORCH: _route_output
+        alt need forward to next stage
+            ORCH->>SN: add_request_async
+        end
+        ORCH-->>ENG: output_queue.put
+    end
+
+    AO->>ENG: try_get_output_async
+    ENG-->>AO: message
+    AO-->>APP: yield OmniRequestOutput
+```
+
+## 4. Comparison
+
+Previous topology (reference):
+
+```text
+┌────────────────────────────────────────────────────────────────────────────┐
+│ Main Process                                                               │
+│  ┌──────────────────────┐   ┌────────────────────────────────────────────┐ │
+│  │ generate()           │   │ final_output_handler()                     │ │
+│  └──────────────────────┘   └────────────────────────────────────────────┘ │
+└──────────┬─────────────────────────┬─────────────────────────┬─────────────┘
+  mp.Queue (in_q/out_q)    mp.Queue (in_q/out_q)    mp.Queue (in_q/out_q)
+           ▼▲                        ▼▲                        ▼▲
+┌───────────────────────┐  ┌───────────────────────┐  ┌──────────────────────┐
+│ Worker Proc-0         │  │ Worker Proc-1         │  │ Worker Proc-2        │
+│ (Thinker LLM)         │  │ (Talker LLM)          │  │ (Vocoder)            │
+│  ┌────────────────┐   │  │  ┌────────────────┐   │  │  ┌────────────────┐  │
+│  │_stage_worker   │   │  │  │_stage_worker   │   │  │  │_stage_worker   │  │
+│  │_async()        │   │  │  │_async()        │   │  │  │_async()        │  │
+│  └────────────────┘   │  │  └────────────────┘   │  │  └────────────────┘  │
+│  ┌────────────────┐   │  │  ┌────────────────┐   │  │  ┌────────────────┐  │
+│  │output_handler()│   │  │  │output_handler()│   │  │  │output_handler()│  │
+│  └────────────────┘   │  │  └────────────────┘   │  │  └────────────────┘  │
+└──────────┬────────────┘  └──────────┬────────────┘  └──────────┬───────────┘
+       ZMQ ▼ ▲ ZMQ               ZMQ ▼ ▲ ZMQ               ZMQ ▼ ▲ ZMQ
+┌──────────────────────┐   ┌──────────────────────┐   ┌──────────────────────┐
+│ EngineCore Proc-0    │   │ EngineCore Proc-1    │   │ EngineCore Proc-2    │
+│ (Thinker)            │   │ (Talker)             │   │ (Vocoder)            │
+└──────────────────────┘   └──────────────────────┘   └──────────────────────┘
+```
+
+Current topology:
+
+```text
+┌────────────────────────────────────────────────────────────────────────────┐
+│ Main Process                                                               │
+│  ┌──────────────────────────────────────────────────────────────────────┐  │
+│  │ Main Thread                                                          │  │
+│  │  ┌──────────────────────┐   ┌─────────────────────────────────────┐  │  │
+│  │  │ generate()           │   │ final_output_handler()              │  │  │
+│  │  └──────────────────────┘   └─────────────────────────────────────┘  │  │
+│  └──────────────────────────────────────────────────────────────────────┘  │
+│         janus.Queue (request_queue) ▼  ▲ janus.Queue (output_queue)        │
+│  ┌──────────────────────────────────────────────────────────────────────┐  │
+│  │ Orchestrator Thread                                                  │  │
+│  │  ┌──────────────────────┐  ┌──────────────────────────────────────┐  │  │
+│  │  │ _request_handler()   │  │ _orchestration_output_handler()      │  │  │
+│  │  └──────────────────────┘  └──────────────────────────────────────┘  │  │
+│  │  ┌────────────────────────────────────────────────────────────────┐  │  │
+│  │  │ _orchestration_loop(): poll/process/route outputs for all stages│  │  │
+│  │  └────────────────────────────────────────────────────────────────┘  │  │
+│  └───────┬─────────────────────────┬─────────────────────────┬──────────┘  │
+└──────────┬─────────────────────────┬─────────────────────────┬─────────────┘
+       ZMQ ▼ ▲ ZMQ               ZMQ ▼ ▲ ZMQ               ZMQ ▼ ▲ ZMQ  
+  ┌──────────────────────┐  ┌──────────────────────┐  ┌──────────────────────┐
+  │ EngineCore Proc-0    │  │ EngineCore Proc-1    │  │ EngineCore Proc-2    │
+  │ (Thinker)            │  │ (Talker)             │  │ (Vocoder)            │
+  └──────────────────────┘  └──────────────────────┘  └──────────────────────┘
+```
+
+
+Test scripts:
+```bash
+# enter offline inference folder.
+cd examples/offline_inference/qwen2_5_omni
+python end2end.py --output-dir output_audio --query-type use_mixed_modalities
+
+cd ../qwen3_omni
+python end2end.py --output-dir output_audio --query-type text --async-chunk --enable-stats
+
+cd ../bagel
+python end2end.py --prompts "A cute cat"
+
+cd ../text_to_image
+python text_to_image.py --prompt "a cup of coffee on the table" --output output.png
+```

docs/design/module/archive/dit_module.md

@@ -0,0 +1,978 @@
+---
+toc_depth: 4
+---
+
+# Diffusion Module Architecture Design
+
+The vLLM-Omni diffusion module (`vllm_omni/diffusion`) is a high-performance inference engine for diffusion models, designed with a modular architecture that separates concerns across multiple components. It provides efficient execution for non-autoregressive generation tasks such as image and video generation.
+
+This document describes the architecture design of the diffusion module, including the diffusion engine, scheduler, worker, diffusion pipeline, and acceleration components.
+
+<p align="center">
+   <img src="https://github.com/user-attachments/assets/cde1ebea-4006-4d30-8283-15a5721a09d2" alt="vLLM-Omni Diffusion Module Components" width="80%">
+</p>
+<p align="center">
+  <em> Main Components of the Diffusion Module </em>
+</p>
+
+
+**Table of Content:**
+
+- [Architecture Overview](#architecture-overview)
+- [Diffusion Engine](#1-diffusion-engine)
+- [Scheduler](#2-scheduler)
+- [Worker](#3-worker)
+- [Diffusion Pipeline](#4-diffusion-pipeline)
+- [Acceleration Components](#5-acceleration-components)
+    - [Attention Backends](#51-attention-backends)
+    - [Parallel Attention](#52-parallel-attention)
+    - [Cache Backends](#53-cache-backends)
+    - [Parallel Strategies](#54-parallel-strategies)
+- [Data Flow](#6-data-flow)
+
+---
+
+## Architecture Overview
+
+The diffusion module follows a **multi-process, distributed architecture** with clear separation of concerns:
+
+<p align="center">
+  <img src="https://github.com/user-attachments/assets/78c4d446-d238-4406-a057-eb17d2ccc8e0" alt="vLLM-Omni Diffusion Module Architecture" width="100%">
+</p>
+<p align="center">
+  <em> Diffusion Architecture Overview </em>
+</p>
+
+
+---
+
+## 1. Diffusion Engine
+
+**Location**: `vllm_omni/diffusion/diffusion_engine.py`
+
+### Responsibilities
+
+The `DiffusionEngine` is the **orchestrator** of the diffusion inference system. It owns request admission, scheduler/executor coordination, output-stream delivery, cancellation cleanup, and startup warmup.
+
+### Key Components
+
+#### 1.1 Initialization
+
+```python
+class DiffusionEngine:
+    def __init__(
+        self,
+        od_config: OmniDiffusionConfig,
+        scheduler: SchedulerInterface | None = None,
+    ):
+        self.od_config = od_config
+        self._init_process_hooks(od_config)
+        self.execution_mode = self._resolve_execution_mode(od_config)
+        self._init_executor(od_config)
+        self._init_scheduler(od_config, scheduler)
+        self._init_runtime_state()
+        self._init_execute_fn()
+```
+
+**Key Features**:
+
+- **Pre/Post Processing**: Registers model-specific pre-processing and post-processing functions via registry pattern
+
+- **Execution Mode Selection**: Selects one of two execution modes:
+  - `REQUEST_BATCH`: complete requests are scheduled through request-level execution. `max_num_seqs=1` is the serial request path; values above `1` allow compatible requests to use fused `pipeline.forward(batch)` when the pipeline supports it.
+  - `STEP_BATCH`: requests are advanced one denoising step at a time through `prepare_encode()`, `denoise_step()`, `step_scheduler()`, and `post_decode()`. `max_num_seqs` controls the maximum number of compatible active requests in one step wave.
+
+- **Scheduler Selection**: Uses `StepScheduler` for step-batch execution and `RequestScheduler` for request-batch execution. Tests and custom integrations may inject a `SchedulerInterface` instance explicitly.
+
+- **Executor Selection**: Builds the configured `DiffusionExecutor`, then routes scheduled work through `execute_batch` or `execute_step` according to the execution mode. `execute_batch` handles the serial `max_num_seqs=1` request path internally.
+
+- **Factory-Owned Warmup**: `DiffusionEngine.make_engine()` constructs the selected engine backend and then runs startup warmup through `run_startup_warmup()`.
+
+#### 1.2 Execution Modes
+
+```python
+class DiffusionExecutionMode(str, Enum):
+    REQUEST_BATCH = "request_batch"
+    STEP_BATCH = "step_batch"
+```
+
+`streaming_output=True` requires step execution. If streaming is enabled without `step_execution=True`, the engine enables step execution so the output stream can receive intermediate chunks or the final step output.
+
+#### 1.3 Output Stream Lifecycle
+
+```python
+request_id = engine.add_request(request)
+async for output in engine.get_output_stream(request_id):
+    yield output
+```
+
+**Design Decisions**:
+
+- **Unified Stream Semantics**: `step_streaming()` is the public async engine path. Streaming callers forward every yielded output. Non-streaming callers consume the same stream to completion and return only the final output.
+
+- **Queue Delivery**: Each admitted request gets an `asyncio.Queue[DiffusionOutput]`. The busy loop emits chunks/final outputs into that queue.
+
+- **Cancellation Cleanup**: If a client disconnects or cancels stream consumption, `get_output_stream()` removes only the queue. Scheduler terminal state is still finalized when the scheduler reports the request as finished, so request IDs do not remain active indefinitely.
+
+#### 1.4 Request Processing Flow
+
+```python
+async def step_streaming(self, request: OmniDiffusionRequest):
+    # 1. Pre-process requests
+    request = self.pre_process_func(request)
+
+    # 2. Add request and consume the unified output stream
+    async for output in self.async_add_req_and_stream_response(request):
+        # 3. Post-process each delivered output
+        yield self.postprocess_output(request, output, ...)
+```
+
+**Flow**:
+
+1. **Pre-processing**: Applies model-specific transformations
+
+2. **Scheduling**: Delegates request state transitions to the selected scheduler
+
+3. **Execution**: Calls the selected executor function for request, request-batch, or step execution
+
+4. **Output Delivery**: Emits intermediate chunks and final outputs through the request output stream
+
+5. **Post-processing**: Converts raw outputs to final format (e.g., PIL images)
+
+---
+
+## 2. Scheduler
+
+**Location**: `vllm_omni/diffusion/sched/`
+
+### Architecture
+
+The scheduler is a **request-state scheduler**. It owns request lifecycle management and scheduling decisions, while execution stays in `DiffusionEngine` and the executor.
+
+### Key Components
+
+#### 2.1 Request State Model
+
+```python
+class DiffusionRequestStatus(enum.IntEnum):
+    WAITING = ...
+    RUNNING = ...
+    PREEMPTED = ...
+    FINISHED_COMPLETED = ...
+    FINISHED_ABORTED = ...
+    FINISHED_ERROR = ...
+
+@dataclass
+class SchedulerRequestState:
+    request_id: str
+    req: OmniDiffusionRequest
+    sampling_params_key: StepBatchSamplingParamsKey | RequestBatchSamplingParamsKey | None = None
+    status: DiffusionRequestStatus = DiffusionRequestStatus.WAITING
+    error: str | None = None
+```
+
+**Design Features**:
+
+- **Explicit lifecycle**: Requests move through waiting, running, optional preemption, and terminal states.
+
+- **Centralized error handling**: Completion, abort, and error states are all normalized in the scheduler layer.
+
+#### 2.2 Shared Bookkeeping in `BaseScheduler`
+
+```python
+class BaseScheduler:
+    def __init__(self) -> None:
+        self._request_states = {}
+        self._waiting = deque()
+        self._running = []
+        self._finished_req_ids = set()
+        self.max_num_running_reqs = 1
+
+    def add_request(self, request: OmniDiffusionRequest) -> str: ...
+    def schedule(self) -> DiffusionSchedulerOutput: ...
+    def update_from_output(
+        self,
+        sched_output: DiffusionSchedulerOutput,
+        output: BaseRunnerOutput,
+    ) -> set[str]: ...
+```
+
+**Responsibilities**:
+
+- **Lifecycle contract**: Defines how the engine adds requests, triggers one scheduling cycle, and feeds executor results back.
+
+- **Stable boundary**: `DiffusionSchedulerOutput` is the only scheduling result consumed by `DiffusionEngine`.
+
+- **Pluggability**: Different scheduler policies can reuse the same engine integration path.
+
+**Design Features**:
+
+- **Common state storage**: Shared request maps and waiting/running sets live in the base class.
+
+- **Shared cleanup logic**: Duplicate-request checks, finish handling, and state removal are centralized instead of
+  duplicated in each policy.
+
+- **Current constraint boundary**: `BaseScheduler` derives `max_num_running_reqs` from `max_num_seqs`. Request-mode diffusion can use that capacity for compatible independent requests when the configured pipeline declares `supports_request_batch = True`; step-wise diffusion uses the same scheduler capacity for compatible step batches.
+
+#### 2.3 Current `RequestScheduler` Policy
+
+```python
+class RequestScheduler(BaseScheduler):
+    def schedule(self) -> DiffusionSchedulerOutput:
+        # 1. keep existing RUNNING requests in the scheduling result
+        # 2. pull WAITING requests while capacity remains
+        # 3. move newly admitted requests into RUNNING
+```
+
+**Behavior**:
+
+- **FIFO request scheduling**: Waiting requests are promoted in queue order.
+
+- **Compatible request admission**: `RequestScheduler` admits waiting requests while capacity remains and the request's `RequestBatchSamplingParamsKey` is compatible with the active batch. Request-mode execution keeps each logical request independent, while batch-capable pipelines receive the scheduled requests as a runner-side `DiffusionRequestBatch`.
+
+- **Executor result feedback**: `update_from_output()` converts executor output into `FINISHED_COMPLETED` or `FINISHED_ERROR` and returns finished request ids.
+
+#### 2.4 Engine-Driven Execution Loop
+
+```python
+request_id = scheduler.add_request(request)
+while True:
+    sched_output = scheduler.schedule()
+    output = executor.add_req(req)
+    finished_req_ids = scheduler.update_from_output(sched_output, output)
+```
+
+**Design Decisions**:
+
+- **Separation of concerns**: Scheduler manages state and policy; executor handles runtime execution.
+
+- **No scheduler-owned IPC**: Scheduler no longer talks to workers directly.
+
+- **Split concurrency model**: Request-mode diffusion can schedule a static
+  batch of compatible independent requests for one full pipeline forward when
+  the pipeline supports request-level batching; the step-wise path can also
+  admit or remove compatible requests between denoise steps. See
+  [Diffusion Continuous Batching](../../feature/diffusion_continuous_batching.md).
+
+---
+
+## 3. Worker
+
+**Location**: `vllm_omni/diffusion/worker/gpu_worker.py`
+
+### Architecture
+
+Workers are **independent processes** that execute the actual model inference. Each worker runs on a dedicated GPU and participates in distributed inference.
+
+### Key Components
+
+#### 3.1 Worker Process Structure
+
+```python
+class WorkerProc:
+    def __init__(self, od_config, gpu_id, broadcast_handle):
+        # Initialize ZMQ context for IPC
+        self.context = zmq.Context(io_threads=2)
+
+        # Connect to broadcast queue (receive requests)
+        self.mq = MessageQueue.create_from_handle(broadcast_handle, gpu_id)
+
+        # Create result queue (only rank 0)
+        if gpu_id == 0:
+            self.result_mq = MessageQueue(n_reader=1, ...)
+
+        # Initialize GPU worker
+        self.worker = self._create_worker(gpu_id, od_config, ...)
+```
+
+**Initialization Steps**:
+
+1. **IPC Setup**: Creates ZMQ context and message queues
+
+2. **Distributed Environment Setup**: Initializes PyTorch distributed communication
+
+    - For CUDA GPUs: Uses NCCL (fast GPU communication)
+
+    - For NPU: Uses HCCL (Huawei Collective Communications Library)
+
+    - For other devices: Uses appropriate backend (GLOO, MCCL, etc.)
+
+3. **Model Loading**: Loads diffusion pipeline on assigned GPU
+
+4. **Cache Setup**: Enables cache backend if configured.
+
+#### 3.2 GPU Worker
+
+```python
+class DiffusionWorker:
+    def init_device(self):
+        # Set distributed environment variables
+        os.environ["RANK"] = str(rank)
+        os.environ["WORLD_SIZE"] = str(world_size)
+
+        # Select the device and initialize distributed execution
+        self.device = current_omni_platform.get_torch_device(rank)
+        current_omni_platform.set_device(self.device)
+        self.vllm_config = _create_diffusion_worker_vllm_config(
+            self.device, self.od_config
+        )
+        init_distributed_environment(world_size, rank)
+        parallel_config = self.od_config.parallel_config
+        initialize_model_parallel(
+            data_parallel_size=parallel_config.data_parallel_size,
+            cfg_parallel_size=parallel_config.cfg_parallel_size,
+            sequence_parallel_size=parallel_config.sequence_parallel_size,
+            tensor_parallel_size=parallel_config.tensor_parallel_size,
+            pipeline_parallel_size=parallel_config.pipeline_parallel_size,
+        )
+
+    def load_model(self, load_format="default", custom_pipeline_name=None):
+        self.model_runner.load_model(
+            load_format=load_format,
+            custom_pipeline_name=custom_pipeline_name,
+        )
+
+    def init_lora_manager(self):
+        self.lora_manager = DiffusionLoRAManager(
+            pipeline=self.model_runner.pipeline,
+            device=self.device,
+            dtype=self.od_config.dtype,
+        )
+```
+
+**Key Features**:
+
+- **Tensor Parallelism**: Supports multi-GPU tensor parallelism via PyTorch distributed
+
+- **Model Loading**: Uses `DiffusersPipelineLoader` for efficient weight loading
+
+- **Cache Integration**: Enables cache backends (TeaCache, cache-dit, etc.) transparently
+
+#### 3.3 Worker Busy Loop
+
+```python
+def _worker_busy_loop(self):
+    while self._running:
+        # 1. Receive unified message (generation request, RPC request, or shutdown)
+        msg = self.mq.dequeue(timeout=1.0)
+
+        # 2. Route message based on type
+        if isinstance(msg, dict) and msg.get("type") == "rpc":
+            # Handle RPC request
+            result, should_reply = self._execute_rpc(msg)
+            if should_reply:
+                self._return_result(result)
+
+        elif isinstance(msg, dict) and msg.get("type") == "shutdown":
+            # Handle shutdown message
+            self._running = False
+
+        else:
+            # Handle generation request (OmniDiffusionRequest list)
+            output = self.worker.execute_model(msg, self.od_config)
+            self._return_result(output)
+```
+
+**Execution Flow**:
+
+1. **Receive**: Dequeues unified messages from shared memory queue
+
+2. **Route**: Handles different message types (generation, RPC, shutdown)
+
+3. **Execute**: Runs forward pass through pipeline for generation requests
+
+4. **Respond**: Sends results back (rank 0 for generation, specified rank for RPC)
+
+#### 3.4 Model Execution
+
+```python
+@torch.inference_mode()
+def execute_model(self, reqs: list[OmniDiffusionRequest], od_config):
+    req = reqs[0]  # TODO: support batching
+
+    # Refresh cache backend if enabled
+    if self.cache_backend is not None and self.cache_backend.is_enabled():
+        self.cache_backend.refresh(self.pipeline, req.num_inference_steps)
+
+    # Set forward context for parallelism
+    with set_forward_context(
+        vllm_config=self.vllm_config,
+        omni_diffusion_config=self.od_config
+    ):
+        output = self.pipeline.forward(req)
+    return output
+```
+
+The model execution leverages multiple parallelism strategies that are transparently applied during the forward pass. The `set_forward_context()` context manager makes parallel group information available throughout the forward pass:
+
+```python
+# Inside transformer layers, parallel groups are accessed via:
+from vllm_omni.diffusion.distributed.parallel_state import (
+    get_sp_group, get_dp_group, get_cfg_group, get_pp_group
+)
+```
+
+**Optimizations**:
+
+- **Cache Refresh**: Clears cache state before each generation for clean state
+
+- **Context Management**: Forward context ensures parallel groups are available during execution
+
+- **Single Request**: Currently processes one request at a time (batching TODO)
+
+---
+
+## 4. Diffusion Pipeline
+
+**Location**: `vllm_omni/diffusion/models/*/pipeline_*.py`
+
+The pipeline is the **model-specific implementation** that orchestrates the diffusion process. Different models (QwenImage, Wan2.2, Z-Image) have their own pipeline implementations.
+
+Most pipeline implementation are referred from `diffusers`. The multi-step diffusion loop is usually the most time-consuming part during the overall inference process, which is defined by the `diffuse` function in the pipeline class. An example is as follows:
+
+```python
+def diffuse(self, ...):
+    for i, t in enumerate(timesteps):
+        # Forward pass for positive prompt
+        transformer_kwargs = {
+            "hidden_states": latents,
+            "timestep": timestep / 1000,
+            "encoder_hidden_states": prompt_embeds,
+        }
+        noise_pred = self.transformer(**transformer_kwargs)[0]
+
+        # Forward pass for negative prompt (CFG)
+        if do_true_cfg:
+            neg_transformer_kwargs = {...}
+            neg_transformer_kwargs["cache_branch"] = "negative"
+            neg_noise_pred = self.transformer(**neg_transformer_kwargs)[0]
+
+            # Combine predictions
+            comb_pred = neg_noise_pred + true_cfg_scale * (noise_pred - neg_noise_pred)
+            noise_pred = comb_pred * (cond_norm / noise_norm)
+
+        # Scheduler step
+        latents = self.scheduler.step(noise_pred, t, latents)[0]
+
+    return latents
+```
+
+**Key Features**:
+
+- **CFG Support**: Handles classifier-free guidance with separate forward passes
+
+- **Cache Branching**: Uses `cache_branch` parameter for cache-aware execution
+
+- **True CFG**: Implements advanced CFG with norm preservation
+
+To learn more about the diffusion pipeline and how to add a new diffusion pipeline, please view [Adding Diffusion Model](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/model/adding_diffusion_model)
+
+---
+
+## 5. Acceleration Components
+
+### 5.1 Attention Backends
+
+**Location**: `vllm_omni/diffusion/attention/`
+
+#### Architecture
+
+The attention system uses a **role-aware backend selector pattern**. Each `Attention` site declares a semantic `role` (`"self"`, `"cross"`, or a model-specific string). The selector consults the user's `AttentionConfig` to pick a backend per role and falls back to the platform default when nothing is configured.
+
+#### Backend Selection
+
+**Location**: `vllm_omni/diffusion/attention/selector.py`
+
+```python
+class Attention(nn.Module):
+    def __init__(self, num_heads, head_size, causal, softmax_scale, *,
+                 role="self", role_category=None, ...):
+        # Resolve backend for this role from the active OmniDiffusionConfig
+        config = get_current_diffusion_config_or_none()
+        attention_config = config.diffusion_attention_config if config is not None else None
+        attn_backend_cls, spec = get_attn_backend_for_role(
+            role=role,
+            head_size=head_size,
+            attention_config=attention_config,
+            role_category=role_category,
+        )
+        self.attn_impl_cls = attn_backend_cls.get_impl_cls()
+        self.attention = self.attn_impl_cls(..., backend_kwargs=spec.backend_kwargs() if spec else None)
+```
+
+**Available Backends**:
+
+- **FlashAttention** (`FLASH_ATTN`): Memory-efficient kernel that dispatches to FA2/FA3 on CUDA, AITER on ROCm, and `mindiesd` on Ascend NPU.
+
+- **SDPA** (`TORCH_SDPA`): PyTorch's scaled dot-product attention — cross-platform fallback.
+
+- **SageAttention** (`SAGE_ATTN`): Quantized attention implementation from the SageAttention library.
+
+These backends provide the **kernel implementations** for attention computation. For attention-level sequence parallelism strategies (Ring Attention, Ulysses), see [Parallel Attention](#52-parallel-attention).
+
+#### Backend Selection Mechanism
+
+Selection is driven by `OmniDiffusionConfig.diffusion_attention_config`, which carries a global `default` plus a `per_role` map of `AttentionSpec` entries. A role string is resolved in this order:
+
+```python
+def get_attn_backend_for_role(role, head_size, attention_config=None, role_category=None):
+    # 1. attention_config.per_role[role]            — exact match
+    # 2. attention_config.per_role[role_category]   — category fallback
+    # 3. attention_config.default                   — global default
+    # 4. Platform default                           — hardware-specific
+    if attention_config is not None:
+        spec, source = attention_config.resolve_with_source(
+            role=role, role_category=role_category,
+        )
+        if spec is not None:
+            return load_backend(spec.backend), spec
+    return current_omni_platform.get_diffusion_attn_backend_cls(...), None
+```
+
+**Selection Priority**:
+
+1. **Per-role override** (`--diffusion-attention-config.per_role.<role>.backend`) — finest control; matched on the layer's exact `role` string.
+
+2. **Role category fallback** (`--diffusion-attention-config.per_role.<category>.backend`) — used when an exact match is missing and the layer declared `role_category` (e.g. a `"mymodel.audio_to_video"` site can fall back to `"cross"`).
+
+3. **Global default** (`--diffusion-attention-backend FLASH_ATTN`, or `--diffusion-attention-config.default.backend`, or `DIFFUSION_ATTENTION_BACKEND` env var).
+
+4. **Platform default** — `current_omni_platform.get_diffusion_attn_backend_cls(...)` picks the best-available kernel for the hardware.
+
+A backend that needs configuration exposes it as a typed field on `AttentionSpec`; `AttentionSpec.backend_kwargs()` serializes it into the `backend_kwargs` the impl consumes. Today those fields are `skip_softmax` (TRTLLM_ATTN) and `block_sparse` (any block-sparse backend, currently RAINFUSION_ATTN), so e.g. `target_sparsity` or `sparsity` can be set without changing model code.
+
+For the user-facing CLI surface, see [Diffusion Attention Backends](../../../user_guide/diffusion/attention_backends.md). For the role declaration contract on the model side, see [Adding a Diffusion Model](../../../contributing/model/adding_diffusion_model.md).
+
+**Backend Availability**:
+
+- **SDPA**: Always available (PyTorch built-in)
+
+- **FlashAttention**: Requires `flash-attn` on CUDA, `aiter` on ROCm, or `mindiesd` on Ascend NPU
+
+- **SageAttention**: Requires `sage-attention` package (from THU-ML GitHub)
+
+#### Attention Backend Registry
+
+**Location**: `vllm_omni/diffusion/attention/selector.py` + `vllm_omni/platforms/<device>/platform.py`
+
+Backend resolution is delegated to the active platform: `current_omni_platform.get_diffusion_attn_backend_cls(selected_backend, head_size)` returns a fully-qualified class path that `_load_backend_cls` imports lazily, with the result cached by `(backend_name, head_size)`.
+
+```python
+@cache
+def _cached_get_backend_cls(backend_name: str | None, head_size: int) -> type[AttentionBackend]:
+    backend_cls_path = current_omni_platform.get_diffusion_attn_backend_cls(
+        selected_backend=backend_name,
+        head_size=head_size,
+    )
+    return _load_backend_cls(backend_cls_path)
+```
+
+Each platform (`cuda`, `rocm`, `xpu`, `musa`, `npu`) maps backend names like `"FLASH_ATTN"`, `"TORCH_SDPA"`, `"SAGE_ATTN"` to the right backend class path for that hardware, including head-size compatibility checks. Passing `selected_backend=None` lets the platform pick its own default.
+
+#### Attention Backend Integration
+
+The `Attention` layer integrates backends through a unified interface. Here's how **FlashAttentionBackend** is integrated as an example:
+
+```python
+# attention/backends/flash_attn.py
+
+class FlashAttentionBackend(AttentionBackend):
+    @staticmethod
+    def get_name() -> str:
+        return "FLASH_ATTN"
+
+    @staticmethod
+    def get_impl_cls() -> type["FlashAttentionImpl"]:
+        return FlashAttentionImpl
+
+    @staticmethod
+    def get_supported_head_sizes() -> list[int]:
+        return [64, 96, 128, 192, 256]  # FlashAttention supports these head sizes
+
+
+class FlashAttentionImpl(AttentionImpl):
+    def __init__(self, num_heads, head_size, softmax_scale, causal, ...):
+        self.num_heads = num_heads
+        self.causal = causal
+        self.softmax_scale = softmax_scale
+
+    def forward(self, query, key, value, attn_metadata=None):
+        # Call FlashAttention kernel
+        out = flash_attn_func(
+            query, key, value,
+            causal=self.causal,
+            softmax_scale=self.softmax_scale,
+        )
+        return out
+```
+
+---
+
+### 5.2 Parallel Attention
+
+**Location**: `vllm_omni/diffusion/attention/parallel/`
+
+#### Architecture
+
+Parallel attention strategies implement **Sequence Parallelism (SP) at the attention layer level**. These strategies distribute attention computation across multiple GPUs by splitting the sequence dimension, using different communication patterns. They work **on top of** AttentionBackend implementations (FlashAttention, SDPA, etc.), handling the parallelization/communication while the backends handle the actual attention computation.
+
+**Key Distinction**: Unlike AttentionBackend (which provides kernel implementations), ParallelAttentionStrategy provides communication patterns for multi-GPU attention parallelism. These strategies implement the `ParallelAttentionStrategy` interface and use AttentionBackend implementations internally.
+
+Both Ring Attention and Ulysses are forms of Sequence Parallelism (SP) that:
+
+- Split the sequence dimension across GPUs
+
+- Contribute to `sequence_parallel_size` (via `ring_degree` and `ulysses_degree`)
+
+- Work at the attention layer level (not model/pipeline level)
+
+#### Ulysses Sequence Parallelism (USP)
+
+**Location**: `vllm_omni/diffusion/attention/parallel/ulysses.py`
+
+USP is a sequence-parallel attention strategy that splits attention computation across multiple GPUs by distributing both the sequence dimension and attention heads. It uses **all-to-all communication** to efficiently parallelize attention for very long sequences. Specifically, it uses **all-to-all** collective operations to redistribute Q/K/V tensors before attention computation and gather results afterward.
+
+Ulysses splits attention computation in two dimensions:
+
+1. **Sequence Dimension**: Splits the sequence length across GPUs
+
+2. **Head Dimension**: Splits attention heads across GPUs
+
+**Configuration**: `ulysses_degree` contributes to `sequence_parallel_size`
+
+#### Ring Sequence Parallelism
+
+**Location**: `vllm_omni/diffusion/attention/parallel/ring.py`
+
+Ring Attention is a **parallel attention strategy** that implements sequence parallelism using ring-based point-to-point (P2P) communication. Unlike attention backends that provide the attention kernel implementation, Ring Attention is a **communication pattern** that works on top of attention backends (FlashAttention or SDPA).
+
+Ring Attention splits sequence dimension across GPUs in a ring topology, implemented via the `ParallelAttentionStrategy` interface, instead of `AttentionBackend`. P2P ring communication is applied to circulate Key/Value blocks across GPUs. Internally, `ring_flash_attn_func` or `ring_pytorch_attn_func` is used depending on available backends.
+
+**Architecture**:
+```python
+class RingParallelAttention:
+    """Ring sequence-parallel strategy."""
+
+    def run_attention(self, query, key, value, attn_metadata, ...):
+        # Selects underlying attention kernel (FlashAttention or SDPA)
+        if backend_pref == "sdpa":
+            return ring_pytorch_attn_func(...)  # Uses SDPA kernel
+        else:
+            return ring_flash_attn_func(...)    # Uses FlashAttention kernel
+```
+
+**Integration**:
+
+- Ring Attention is activated when `ring_degree > 1` in parallel config
+
+- It's selected by `build_parallel_attention_strategy()` in the attention layer
+
+- The `Attention` layer routes to `_run_ring_attention()` when Ring is enabled
+
+- Works alongside attention backends: Ring handles communication, backends handle computation
+
+**Configuration**: `ring_degree` contributes to `sequence_parallel_size`
+
+#### Relationship with AttentionBackend
+
+Parallel attention strategies (Ring, Ulysses) work **on top of** AttentionBackend implementations:
+
+- They use AttentionBackend for the actual attention computation (FlashAttention, SDPA, etc.)
+
+- They handle the multi-GPU communication/parallelization layer
+
+- They implement `ParallelAttentionStrategy` interface (not `AttentionBackend`)
+
+For general parallelism strategies (Data Parallelism, Tensor Parallelism, Pipeline Parallelism), see [Parallel Strategies](#54-parallel-strategies).
+
+---
+
+### 5.3 Cache Backends
+
+**Location**: `vllm_omni/diffusion/cache/`
+
+#### Architecture
+
+Cache backends provide a **unified interface** for applying different caching strategies to accelerate diffusion inference. The system supports multiple backends (TeaCache, cache-dit) with a consistent API for enabling and refreshing cache state.
+
+#### Cache Backend Interface
+
+```python
+class CacheBackend(ABC):
+    def __init__(self, config: DiffusionCacheConfig):
+        self.config = config
+        self.enabled = False
+
+    @abstractmethod
+    def enable(self, pipeline: Any) -> None:
+        """Enable cache on the pipeline."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def refresh(self, pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None:
+        """Refresh cache state for new generation."""
+        raise NotImplementedError
+
+    def is_enabled(self) -> bool:
+        """Check if cache is enabled."""
+        return self.enabled
+```
+
+**Design Pattern**:
+
+- **Abstract Base Class**: Defines contract for all cache backends
+
+- **Pipeline-based**: Works with pipeline instances (not just transformers)
+
+- **State Management**: Provides refresh mechanism for clean state between generations
+
+#### Available Backends
+
+**1. TeaCache Backend**
+
+**Location**: `vllm_omni/diffusion/cache/teacache/backend.py`
+
+```python
+class TeaCacheBackend(CacheBackend):
+    def enable(self, pipeline: Any):
+        # Extract transformer from pipeline
+        transformer = pipeline.transformer
+        transformer_type = transformer.__class__.__name__
+
+        # Create TeaCacheConfig from DiffusionCacheConfig
+        teacache_config = TeaCacheConfig(
+            transformer_type=transformer_type,
+            rel_l1_thresh=self.config.rel_l1_thresh,
+            coefficients=self.config.coefficients,
+        )
+
+        # Apply hooks to transformer
+        apply_teacache_hook(transformer, teacache_config)
+        self.enabled = True
+
+    def refresh(self, pipeline: Any, num_inference_steps: int, verbose: bool = True):
+        transformer = pipeline.transformer
+        if hasattr(transformer, "_hook_registry"):
+            transformer._hook_registry.reset_hook(TeaCacheHook._HOOK_NAME)
+```
+
+**TeaCache Features**:
+
+- **Timestep-aware**: Caches based on timestep embedding similarity
+
+- **Adaptive**: Dynamically decides when to reuse cached computations
+
+- **CFG-aware**: Handles positive/negative branches separately
+
+- **Custom Hook System**: Uses a custom forward interception mechanism (via `HookRegistry`) that wraps the module's `forward` method, allowing transparent integration without modifying model code
+
+**2. Cache-DiT Backend**
+
+**Location**: `vllm_omni/diffusion/cache/cachedit/`
+
+```python
+class CacheDiTBackend(CacheBackend):
+    def enable(self, pipeline: Any):
+        # Uses cache-dit library for acceleration
+        # Supports DBCache, SCM (Step Computation Masking), TaylorSeer
+        # Works with single and dual-transformer architectures
+        ...
+        self.enabled = True
+
+    def refresh(self, pipeline: Any, num_inference_steps: int, verbose: bool = True):
+        # Updates cache context with new num_inference_steps
+        ...
+```
+
+**Cache-DiT Features**:
+
+- **DBCache**: Dynamic block caching with configurable compute blocks
+
+- **SCM**: Step Computation Masking for additional speedup
+
+- **TaylorSeer**: Advanced calibration for cache accuracy
+
+- **Dual-transformer Support**: Handles models like Wan2.2 with two transformers
+
+#### Cache Backend Selector
+
+**Location**: `vllm_omni/diffusion/cache/selector.py`
+
+```python
+def get_cache_backend(
+    cache_backend: str | None,
+    cache_config: dict | DiffusionCacheConfig
+) -> CacheBackend | None:
+    """Get cache backend instance based on cache_backend string.
+
+    Args:
+        cache_backend: Cache backend name ("cache_dit", "tea_cache", or None)
+        cache_config: Cache configuration (dict or DiffusionCacheConfig)
+
+    Returns:
+        Cache backend instance or None if cache_backend is "none"
+    """
+    if cache_backend is None or cache_backend == "none":
+        return None
+
+    if isinstance(cache_config, dict):
+        cache_config = DiffusionCacheConfig.from_dict(cache_config)
+
+    if cache_backend == "cache_dit":
+        return CacheDiTBackend(cache_config)
+    elif cache_backend == "tea_cache":
+        return TeaCacheBackend(cache_config)
+    else:
+        raise ValueError(f"Unsupported cache backend: {cache_backend}")
+```
+
+**Usage Flow**:
+
+1. **Selection**: `get_cache_backend()` returns appropriate backend instance
+
+2. **Enable**: `backend.enable(pipeline)` called during worker initialization
+
+3. **Refresh**: `backend.refresh(pipeline, num_inference_steps)` called before each generation
+
+4. **Check**: `backend.is_enabled()` verifies cache is active
+
+### 5.4 Parallel Strategies
+
+**Location**: `vllm_omni/diffusion/distributed/parallel_state.py`
+
+#### Parallelism Types
+
+The system supports multiple orthogonal parallelism strategies:
+
+**Sequence Parallelism (SP)**
+
+- **Purpose**: Split sequence dimension across GPUs
+
+- **Attention-level SP**: Ring Attention and Ulysses (USP) implement SP at the attention layer level
+
+    - See [Parallel Attention](#52-parallel-attention) for details
+
+    - Configuration: `ulysses_degree` × `ring_degree` = `sequence_parallel_size`
+
+- **Use Case**: Very long sequences (e.g., high-resolution images)
+
+**Data Parallelism (DP)**
+
+- **Purpose**: Replicate model across GPUs, split batch
+
+- **Use Case**: Batch processing, throughput optimization
+
+**Tensor Parallelism (TP)** (Experimental)
+
+- **Purpose**: Split model weights across GPUs
+
+- **Implementation**: Uses vLLM's tensor parallel groups
+
+- **Use Case**: Large models that don't fit on single GPU
+
+**CFG Parallelism**  (under development)
+
+- **Purpose**: Parallelize Classifier-Free Guidance (positive/negative prompts)
+
+- **Infrastructure**: CFG parallel groups are initialized and available via `get_cfg_group()`
+
+#### Parallel Group Management
+
+```python
+def initialize_model_parallel(
+    data_parallel_size: int = 1,
+    cfg_parallel_size: int = 1,
+    sequence_parallel_size: int | None = None,
+    ulysses_degree: int = 1,
+    ring_degree: int = 1,
+    tensor_parallel_size: int = 1,
+    pipeline_parallel_size: int = 1,
+    vae_parallel_size: int = 0,
+):
+    # Generate orthogonal parallel groups
+    rank_generator = RankGenerator(
+        tensor_parallel_size,
+        sequence_parallel_size,
+        pipeline_parallel_size,
+        cfg_parallel_size,
+        data_parallel_size,
+        "tp-sp-pp-cfg-dp",
+    )
+
+    # Initialize each parallel group
+    _DP = init_model_parallel_group(rank_generator.get_ranks("dp"), ...)
+    _CFG = init_model_parallel_group(rank_generator.get_ranks("cfg"), ...)
+    _SP = init_model_parallel_group(rank_generator.get_ranks("sp"), ...)
+    _PP = init_model_parallel_group(rank_generator.get_ranks("pp"), ...)
+    _TP = init_model_parallel_group(rank_generator.get_ranks("tp"), ...)
+```
+
+**Rank Order**: `tp-sp-pp-cfg-dp` (tensor → sequence → pipeline → cfg → data)
+
+**Note**: For attention-level Sequence Parallelism implementations (Ring Attention and Ulysses), see [Parallel Attention](#52-parallel-attention). This section covers higher-level parallelism strategies.
+
+
+---
+
+## 6. Data Flow
+
+### Complete Request Flow
+
+<p align="center">
+   <img src="https://github.com/user-attachments/assets/6e093a0d-29c0-4efd-85d5-747002bd2fed" alt="vLLM-Omni Diffusion Module Components" width="100%">
+</p>
+<p align="center">
+  <em> End-to-end Data Flow in the vLLM-Omni Diffusion Module </em>
+</p>
+
+
+```
+1. User Request
+   └─> Omni.generate(prompt)
+       └─> Prepare OmniDiffusionRequest
+           └─> DiffusionEngine.step_streaming(request)
+
+2. Pre-processing
+   └─> pre_process_func(request)
+       └─> Model-specific transformations
+
+3. Scheduling
+       └─> scheduler.add_request(request)
+           └─> scheduler.schedule()
+           └─> DiffusionEngine selects execute_batch / execute_step
+
+4. Worker Execution
+   └─> WorkerProc._worker_busy_loop()
+       └─> DiffusionWorker.execute_model(req)
+           └─> DiffusionModelRunner selects execution mode
+               ├─> REQUEST_BATCH
+               │   └─> Pipeline.forward(req_or_batch)
+               │       ├─> encode_prompt()
+               │       ├─> prepare_latents()
+               │       ├─> diffuse() [loop]
+               │       │   ├─> transformer.forward() [with cache backend hooks]
+               │       │   └─> scheduler.step()
+               │       └─> vae.decode()
+               └─> STEP
+                   ├─> pipeline.prepare_encode(state) [new requests]
+                   ├─> pipeline.denoise_step(input_batch)
+                   ├─> pipeline.step_scheduler(state, noise_pred)
+                   └─> pipeline.post_decode(state) [chunk boundary or final step]
+
+5. Result Collection
+   └─> Executor returns RunnerOutput / BatchRunnerOutput
+       └─> scheduler.update_from_output(...)
+           ├─> DiffusionEngine emits chunk/final output to request queue
+           └─> DiffusionEngine finalizes terminal scheduler state
+
+6. Post-processing
+   └─> post_process_func(output)
+       └─> Convert to PIL images / final format
+
+7. Stream Completion / Cancellation
+   ├─> Streaming callers receive every yielded output
+   ├─> Non-streaming callers drain the stream and return the final output
+   └─> Cancelled streams drop their output queue but still allow scheduler state cleanup
+```
+
+---

docs/design/module/archive/entrypoint_module.md

@@ -0,0 +1 @@
+Architecture design of the entrypoint (update soon)

docs/design/module/engine_orchestration.md

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

docs/design/module/entrypoints.md

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

docs/design/module/error_contracts.md

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

docs/design/module/input_output_modality_contracts.md

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

docs/design/module/stage_runtime.md

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

docs/design/module/vllm_omni_config.md

@@ -2,6 +2,8 @@
 title: vLLM-Omni Configuration
 kind: module
 status: draft
+architecture_state: deferred-pending-refactor
+ownership_status: provisional
 owners:
   - "@lishunyang12"
   - "@alex-jw-brooks"
@@ -16,32 +18,39 @@ validation_paths:
   - tests/config/**
 upstream_refs:
   - vllm.config
-last_reviewed: 2026-07-16
+last_reviewed: 2026-08-07
+last_verified_commit: 3d7fc3b9ba3cac88d579d4dc35b78b0b641675fc
 ---
 
 # vLLM-Omni configuration
 
-Configuration defines how user input becomes validated engine, stage,
-deployment, and platform settings.
+This page intentionally records discovery paths only. Substantive
+configuration ownership and contract work is deferred until the active
+configuration refactoring has settled.
 
-## Candidate invariants
+## Deferred contract scope
 
-### CONFIG-INV-001: Configuration is validated before startup
+The current code and tests remain authoritative for parsing, defaults,
+overrides, deployment topology, stage construction, and runtime projection.
+This draft does not define a stable precedence model, environment-variable
+capture rule, canonical configuration object, or invariant namespace.
 
-**Rule:** Invalid stage graphs, resource assignments, and incompatible options
-MUST fail before managed runtime processes start.
+The PR review proposal to capture environment-derived values in configuration
+objects at initialization is an open design question, not a current
+invariant.
 
-### CONFIG-INV-002: Precedence is deterministic
+## Safe-change guide while deferred
 
-**Rule:** Defaults, files, environment variables, and CLI overrides MUST have a
-documented and deterministic precedence.
+Changes should trace every affected producer to its runtime consumer and test
+the structured, legacy, CLI, and deployment paths that are actually supported.
+Do not infer a stable contract from this placeholder page.
 
-### CONFIG-INV-003: Runtime modules consume validated configuration
+## Promotion gate
 
-**Rule:** Runtime modules MUST NOT independently reinterpret raw user
-configuration.
-
-## Safe-change guide
-
-Test defaults, parsing, override precedence, serialization, invalid
-combinations, and representative multi-stage configurations.
+- Finish or stabilize the configuration refactor and identify the canonical
+  runtime configuration representation.
+- Confirm technical owners, primary path exceptions, and validation owners.
+- Document parsing and override precedence, including when environment values
+  are captured.
+- Allocate an invariant namespace only after those behaviors are enforced and
+  owner-approved.