diff --git a/docs/design/module/ar_module.md b/docs/design/module/ar_module.md deleted file mode 100644 index 5e0aa5b07..000000000 --- a/docs/design/module/ar_module.md +++ /dev/null @@ -1,387 +0,0 @@ -# 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. diff --git a/docs/design/module/ar_runtime.md b/docs/design/module/ar_runtime.md new file mode 100644 index 000000000..194aade5f --- /dev/null +++ b/docs/design/module/ar_runtime.md @@ -0,0 +1,52 @@ +--- +title: Autoregressive Runtime +kind: module +status: draft +owners: + - "@tzhouam" + - "@yinpeiqi" + - "@fake0fan" + - "@Sy0307" +primary_code_paths: + - vllm_omni/core/** + - vllm_omni/worker/** +related_code_paths: + - vllm_omni/model_executor/** +depends_on: + - engine_orchestration.md + - model_integration.md + - input_output_modality_contracts.md +validation_paths: + - tests/core/** + - tests/worker/** +upstream_refs: + - vllm.v1.core + - vllm.v1.worker +last_reviewed: 2026-07-16 +--- + +# Autoregressive runtime + +The AR runtime extends vLLM scheduling and worker execution for omni-stage +inputs and outputs while preserving vLLM scheduling and cache semantics. + +## Candidate invariants + +### AR-INV-001: vLLM owns base scheduling semantics + +**Rule:** Omni schedulers MUST preserve upstream request-state and cache +transitions unless an Omni-specific difference is documented and tested. + +### AR-INV-002: Omni data crosses explicit adapters + +**Rule:** Modality-specific stage data MUST be converted at an input or output +adapter, not injected through unrelated scheduler state. + +### AR-INV-003: Workers execute assigned work + +**Rule:** Workers and model runners MUST NOT implement cross-stage routing. + +## Safe-change guide + +Test request lifecycle, abort, cache state, and every affected worker execution +mode against the supported upstream vLLM contract. diff --git a/docs/design/module/async_omni_architecture.md b/docs/design/module/async_omni_architecture.md deleted file mode 100644 index 92b13a3da..000000000 --- a/docs/design/module/async_omni_architecture.md +++ /dev/null @@ -1,203 +0,0 @@ -# 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 -``` diff --git a/docs/design/module/benchmarking.md b/docs/design/module/benchmarking.md new file mode 100644 index 000000000..93eba2c72 --- /dev/null +++ b/docs/design/module/benchmarking.md @@ -0,0 +1,47 @@ +--- +title: Benchmarking +kind: module +status: draft +owners: + - "@alex-jw-brooks" +primary_code_paths: + - vllm_omni/benchmarks/** + - vllm_omni/entrypoints/cli/benchmark/** +related_code_paths: + - benchmarks/** +depends_on: + - entrypoints.md + - profiling.md +validation_paths: + - tests/benchmarks/** + - tests/dfx/perf/** +upstream_refs: + - vllm.benchmarks +last_reviewed: 2026-07-16 +--- + +# Benchmarking + +Benchmarking defines reproducible workloads, metrics, warm-up, measurement, +and result metadata for comparing vLLM-Omni performance. + +## Candidate invariants + +### BENCH-INV-001: Results are reproducible + +**Rule:** Results MUST identify the model, revision, commit, hardware, +configuration, workload, concurrency, warm-up, and measurement window. + +### BENCH-INV-002: Correctness precedes performance + +**Rule:** A benchmark MUST validate successful and semantically acceptable +outputs before using timings as performance evidence. + +### BENCH-INV-003: Metric definitions are explicit + +**Rule:** Metrics MUST state their unit, population, and aggregation. + +## Safe-change guide + +Test workload parsing and metric calculations with deterministic fixtures. Do +not compare materially different workloads or runtime settings. diff --git a/docs/design/module/cache_management.md b/docs/design/module/cache_management.md new file mode 100644 index 000000000..6f22cf919 --- /dev/null +++ b/docs/design/module/cache_management.md @@ -0,0 +1,47 @@ +--- +title: Cache Management +kind: module +status: draft +owners: + - "@Isotr0py" + - "@princepride" +primary_code_paths: + - vllm_omni/diffusion/cache/** +related_code_paths: + - vllm_omni/experimental/ar_diffusion/kv_cache/** + - vllm_omni/core/** +depends_on: + - ar_runtime.md + - diffusion/diffusion_runtime.md +validation_paths: + - tests/diffusion/cache/** +upstream_refs: + - vllm.v1.core.kv_cache_manager +last_reviewed: 2026-07-16 +--- + +# Cache management + +Cache management defines reusable state, cache identity, validity, lifecycle, +and eviction across AR and diffusion execution. + +## Candidate invariants + +### CACHE-INV-001: Cache identity is complete + +**Rule:** A cache key MUST include every input and execution property that can +change the reused value. + +### CACHE-INV-002: Reuse preserves correctness + +**Rule:** A cache optimization MUST provide a disabled path and MUST NOT reuse +state after its validity conditions change. + +### CACHE-INV-003: Ownership has an end + +**Rule:** Request-scoped and model-scoped cache entries MUST have explicit +reset, eviction, or teardown behavior. + +## Safe-change guide + +Test hit, miss, invalidation, disabled, concurrent-request, and teardown paths. diff --git a/docs/design/module/diffusion/continuous_batching.md b/docs/design/module/diffusion/continuous_batching.md new file mode 100644 index 000000000..bffe24309 --- /dev/null +++ b/docs/design/module/diffusion/continuous_batching.md @@ -0,0 +1,51 @@ +--- +title: Diffusion Continuous Batching +kind: module +status: draft +owners: + - "@Isotr0py" + - "@princepride" +primary_code_paths: + - vllm_omni/diffusion/sched/** +related_code_paths: + - vllm_omni/diffusion/executor/** +depends_on: + - diffusion_runtime.md +validation_paths: + - tests/diffusion/batching/** +upstream_refs: [] +last_reviewed: 2026-07-16 +--- + +# Diffusion continuous batching + +Continuous batching defines when diffusion requests are compatible, how they +share an execution step, and how each request advances independently. + +## Candidate invariants + +### BATCH-INV-001: Compatibility is explicit + +**Rule:** Requests MAY share a batch only when all properties required by the +selected pipeline, scheduler step, shape, precision, and execution features are +compatible. + +### BATCH-INV-002: Per-request state remains isolated + +**Rule:** Batch assembly MUST preserve each request's random generator, progress, +conditioning, outputs, cancellation, and errors. + +### BATCH-INV-003: Admission is bounded + +**Rule:** Admission MUST respect configured capacity and MUST NOT rely on an +unbounded waiting or active-request collection. + +### BATCH-INV-004: Batch order is not completion order + +**Rule:** Output association MUST use stable request identity rather than batch +position across execution steps. + +## Safe-change guide + +Test heterogeneous requests, partial completion, cancellation, deterministic +seeds, capacity limits, and batch sizes of one and many. diff --git a/docs/design/module/diffusion/diffusion_model_integration.md b/docs/design/module/diffusion/diffusion_model_integration.md new file mode 100644 index 000000000..3673db9c1 --- /dev/null +++ b/docs/design/module/diffusion/diffusion_model_integration.md @@ -0,0 +1,58 @@ +--- +title: Diffusion Model Integration +kind: module +status: draft +owners: + - "@Isotr0py" + - "@wtomin" + - "@ZJY0516" +primary_code_paths: + - vllm_omni/diffusion/models/** + - vllm_omni/diffusion/model_loader/** +related_code_paths: + - vllm_omni/diffusion/layers/** + - vllm_omni/diffusion/lora/** + - vllm_omni/diffusion/utils/** +depends_on: + - diffusion_runtime.md + - ../input_output_modality_contracts.md +validation_paths: + - tests/diffusion/models/** + - tests/diffusion/model_loader/** + - tests/diffusion/layers/** + - tests/diffusion/lora/** +upstream_refs: + - diffusers.DiffusionPipeline +last_reviewed: 2026-07-16 +--- + +# Diffusion model integration + +Diffusion model integration owns pipeline contracts, registration, checkpoint +loading, adapters, shared layers, and model-specific processing. + +## Candidate invariants + +### DIFF-MODEL-INV-001: Pipelines implement one runtime contract + +**Rule:** A pipeline MUST declare its supported modalities, configuration, +inputs, outputs, loading path, and runtime capabilities. + +### DIFF-MODEL-INV-002: Registration is the selection boundary + +**Rule:** Runtime code MUST select model implementations through the registry or +loader contract rather than scattered model-name conditionals. + +### DIFF-MODEL-INV-003: Model code does not schedule requests + +**Rule:** Pipeline code MUST NOT own admission, batching, cancellation, or +cross-stage routing. + +### DIFF-MODEL-INV-004: Shared behavior stays shared + +**Rule:** Model directories SHOULD contain only genuine model differences. + +## Safe-change guide + +Test registry selection, checkpoint loading, minimal inference, input and output +contracts, and every declared optional capability. diff --git a/docs/design/module/diffusion/diffusion_runtime.md b/docs/design/module/diffusion/diffusion_runtime.md new file mode 100644 index 000000000..9925ed487 --- /dev/null +++ b/docs/design/module/diffusion/diffusion_runtime.md @@ -0,0 +1,58 @@ +--- +title: Diffusion Runtime +kind: module +status: draft +owners: + - "@Isotr0py" + - "@princepride" + - "@SamitHuang" +primary_code_paths: + - vllm_omni/diffusion/sched/** + - vllm_omni/diffusion/executor/** + - vllm_omni/diffusion/worker/** +related_code_paths: + - vllm_omni/diffusion/hooks/** + - vllm_omni/diffusion/postprocess/** +depends_on: + - ../engine_orchestration.md + - ../input_output_modality_contracts.md + - diffusion_model_integration.md +validation_paths: + - tests/diffusion/diffusion_backend/** + - tests/diffusion/hooks/** +upstream_refs: + - diffusers.DiffusionPipeline +last_reviewed: 2026-07-16 +--- + +# Diffusion runtime + +The diffusion runtime owns request admission, scheduling, execution, progress, +output, cancellation, and cleanup inside a diffusion stage. + +## Candidate invariants + +### DIFF-RUNTIME-INV-001: One lifecycle owner + +**Rule:** Every admitted request MUST have exactly one scheduler-owned lifecycle +until completion, cancellation, or failure. + +### DIFF-RUNTIME-INV-002: Execution follows scheduler output + +**Rule:** Executors and workers MUST execute scheduler decisions without +admitting, reordering, or forwarding requests independently. + +### DIFF-RUNTIME-INV-003: Terminal cleanup is complete + +**Rule:** Every terminal path MUST release request state, temporary tensors, +hooks, and runtime-owned resources. + +### DIFF-RUNTIME-INV-004: Optional features use runtime hooks + +**Rule:** Cache, profiling, offload, and parallel features SHOULD integrate at +defined hooks instead of creating another request lifecycle. + +## Safe-change guide + +Test success, cancellation, failure, shutdown, multiple requests, and resource +cleanup. diff --git a/docs/design/module/diffusion/index.md b/docs/design/module/diffusion/index.md new file mode 100644 index 000000000..51d66eda9 --- /dev/null +++ b/docs/design/module/diffusion/index.md @@ -0,0 +1,42 @@ +--- +title: Diffusion Modules +kind: module-index +status: draft +owners: + - "@Isotr0py" + - "@princepride" + - "@SamitHuang" + - "@wtomin" + - "@ZJY0516" + - "@RuixiangMa" + - "@david6666666" + - "@xuechendi" +primary_code_paths: + - vllm_omni/diffusion/** +related_code_paths: + - vllm_omni/platforms/** +depends_on: + - ../engine_orchestration.md + - ../input_output_modality_contracts.md +validation_paths: + - tests/diffusion/** +upstream_refs: + - diffusers +last_reviewed: 2026-07-16 +--- + +# Diffusion modules + +Diffusion is a module family containing runtime, model integration, batching, +distributed execution, and memory-management subsystems. + +## Module documents + +- [Diffusion runtime](diffusion_runtime.md) +- [Diffusion model integration](diffusion_model_integration.md) +- [Continuous batching](continuous_batching.md) +- [Parallelism](parallelism.md) +- [Offloader](offloader.md) + +Cache, quantization, profiling, and benchmarking follow their top-level module +documents. diff --git a/docs/design/module/diffusion/offloader.md b/docs/design/module/diffusion/offloader.md new file mode 100644 index 000000000..52e634a2f --- /dev/null +++ b/docs/design/module/diffusion/offloader.md @@ -0,0 +1,53 @@ +--- +title: Diffusion Offloader +kind: module +status: draft +owners: + - "@yuanheng-zhao" + - "@Isotr0py" +primary_code_paths: + - vllm_omni/diffusion/offloader/** +related_code_paths: + - vllm_omni/diffusion/models/** + - vllm_omni/diffusion/worker/** +depends_on: + - diffusion_runtime.md + - ../execution_platforms.md +validation_paths: + - tests/diffusion/offloader/** +upstream_refs: + - torch.nn.Module.to +last_reviewed: 2026-07-16 +--- + +# Diffusion offloader + +The offloader owns component residency, transfer scheduling, memory accounting, +prefetch, and teardown for diffusion execution. + +## Candidate invariants + +### OFFLOAD-INV-001: Residency is explicit + +**Rule:** Each managed component MUST have one known residency state and one +owner responsible for transitions. + +### OFFLOAD-INV-002: Use follows readiness + +**Rule:** Runtime execution MUST NOT consume a component until its transfer to +the required device is complete. + +### OFFLOAD-INV-003: Transfers preserve model state + +**Rule:** Offloading MUST preserve parameter and buffer identity, dtype, device, +and correctness required by the selected execution mode. + +### OFFLOAD-INV-004: Memory is bounded and released + +**Rule:** Retained host and device copies MUST respect configured limits and +provide deterministic teardown. + +## Safe-change guide + +Test enabled and disabled paths, repeated execution, asynchronous transfer, +memory limits, failure, teardown, and numerical equivalence. diff --git a/docs/design/module/diffusion/parallelism.md b/docs/design/module/diffusion/parallelism.md new file mode 100644 index 000000000..39034d5d3 --- /dev/null +++ b/docs/design/module/diffusion/parallelism.md @@ -0,0 +1,56 @@ +--- +title: Diffusion Parallelism +kind: module +status: draft +owners: + - "@princepride" + - "@yuanheng-zhao" + - "@xuechendi" +primary_code_paths: + - vllm_omni/diffusion/distributed/** + - vllm_omni/diffusion/attention/parallel/** +related_code_paths: + - vllm_omni/distributed/** + - vllm_omni/config/composable_parallel/** +depends_on: + - diffusion_runtime.md + - ../vllm_omni_config.md +validation_paths: + - tests/diffusion/distributed/** + - tests/diffusion/attention/** +upstream_refs: + - torch.distributed +last_reviewed: 2026-07-16 +--- + +# Diffusion parallelism + +Parallelism owns diffusion rank topology, process groups, tensor and sequence +sharding, collectives, and distributed execution. + +## Candidate invariants + +### PARALLEL-INV-001: Topology has one source of truth + +**Rule:** Rank coordinates, group membership, and parallel dimensions MUST be +derived from validated configuration. + +### PARALLEL-INV-002: Sharding contracts are explicit + +**Rule:** Every distributed boundary MUST define tensor shape, sharded +dimension, placement, dtype, producer, and consumer. + +### PARALLEL-INV-003: Collectives are symmetric + +**Rule:** Process-group members MUST invoke compatible collectives in the same +logical order. + +### PARALLEL-INV-004: Single-rank behavior remains valid + +**Rule:** Distributed implementations SHOULD preserve a supported single-rank +path without distributed initialization. + +## Safe-change guide + +Test topology validation, single-rank execution, affected parallel modes, +combined modes, invalid shapes, and orderly teardown. diff --git a/docs/design/module/dit_module.md b/docs/design/module/dit_module.md deleted file mode 100644 index 02d50181b..000000000 --- a/docs/design/module/dit_module.md +++ /dev/null @@ -1,946 +0,0 @@ ---- -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. - -

- vLLM-Omni Diffusion Module Components -

-

- Main Components of the Diffusion Module -

- - -**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: - -

- vLLM-Omni Diffusion Module Architecture -

-

- Diffusion Architecture Overview -

- - ---- - -## 1. Diffusion Engine - -**Location**: `vllm_omni/diffusion/diffusion_engine.py` - -### Responsibilities - -The `DiffusionEngine` is the **orchestrator** of the diffusion inference system. It manages the lifecycle of worker processes and coordinates the execution flow. - -### Key Components - -#### 1.1 Initialization - -```python -class DiffusionEngine: - def __init__(self, od_config: OmniDiffusionConfig): - self.od_config = od_config - self.post_process_func = get_diffusion_post_process_func(od_config) - self.pre_process_func = get_diffusion_pre_process_func(od_config) - self._processes: list[mp.Process] = [] - self._make_client() -``` - -**Key Features**: - -- **Pre/Post Processing**: Registers model-specific pre-processing and post-processing functions via registry pattern - -- **Worker Management**: Launches and manages multiple worker processes (one per GPU) - -- **Process Isolation**: Uses multiprocessing for true parallelism - -#### 1.2 Worker Launch Process - -The engine launches workers using a **spawn** method: - -```python -def _launch_workers(self, broadcast_handle): - # Creates one process per GPU - for i in range(num_gpus): - process = mp.Process( - target=worker_proc.worker_main, - args=(i, od_config, writer, broadcast_handle), - name=f"DiffusionWorker-{i}", - ) - process.start() -``` - -**Design Decisions**: - -- **Spawn Method**: Ensures clean state for each worker (no shared memory issues) - -- **Pipe Communication**: Uses `mp.Pipe` for initialization handshake - -- **Device Selection**: Each worker is assigned a specific GPU (`cuda:{rank}`) - -#### 1.3 Request Processing Flow - -```python -def step(self, requests: list[OmniDiffusionRequest]): - # 1. Pre-process requests - requests = self.pre_process_func(requests) - - # 2. Send to scheduler and wait for response - output = self.add_req_and_wait_for_response(requests) - - # 3. Post-process results - result = self.post_process_func(output.output) - return result -``` - -**Flow**: - -1. **Pre-processing**: Applies model-specific transformations - -2. **Scheduling**: Delegates to scheduler for distribution - -3. **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 Scheduler Interface - -```python -class SchedulerInterface(ABC): - def add_request(self, request: OmniDiffusionRequest) -> str: ... - def schedule(self) -> DiffusionSchedulerOutput: ... - def update_from_output( - self, - sched_output: DiffusionSchedulerOutput, - output: DiffusionOutput, - ) -> 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. - -#### 2.2 Request State Model - -```python -class DiffusionRequestStatus(enum.IntEnum): - WAITING = ... - RUNNING = ... - PREEMPTED = ... - FINISHED_COMPLETED = ... - FINISHED_ABORTED = ... - FINISHED_ERROR = ... - -@dataclass -class DiffusionRequestState: - request_id: str - req: OmniDiffusionRequest - status: DiffusionRequestStatus = DiffusionRequestStatus.WAITING -``` - -**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.3 Shared Bookkeeping in `_BaseScheduler` - -```python -class _BaseScheduler(SchedulerInterface): - def __init__(self) -> None: - self._request_states = {} - self._waiting = deque() - self._running = [] - self._finished_req_ids = set() - self.max_num_running_reqs = 1 -``` - -**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.4 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 `SamplingParamsKey` 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.5 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 - [Request-Level Batching](../feature/diffusion_request_level_batching.md) and - [Continuous Batching for Step-Wise Diffusion](../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 = GPUWorker(local_rank=gpu_id, rank=gpu_id, od_config=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 GPUWorker: - def init_device_and_model(self): - # Set distributed environment variables - os.environ["RANK"] = str(rank) - os.environ["WORLD_SIZE"] = str(world_size) - - # Initialize PyTorch distributed - 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, - ) - - # Load model - model_loader = DiffusersPipelineLoader(load_config, self.od_config) - self.pipeline = model_loader.load_model(load_device=f"cuda:{rank}") - - # Setup cache backend - from vllm_omni.diffusion.cache.selector import get_cache_backend - self.cache_backend = get_cache_backend(od_config.cache_backend, od_config.cache_config) - - if self.cache_backend is not None: - self.cache_backend.enable(self.pipeline) -``` - -**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.recv_message() - - # 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.extra 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..backend`) — finest control; matched on the layer's exact `role` string. - -2. **Role category fallback** (`--diffusion-attention-config.per_role..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. - -`AttentionSpec.extra` is forwarded to the backend constructor as `backend_kwargs`, so backend-specific parameters (e.g. SparseBlock `block_size`) 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//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/cache_dit_backend.py` - -```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 - -

- vLLM-Omni Diffusion Module Components -

-

- End-to-end Data Flow in the vLLM-Omni Diffusion Module -

- - -``` -1. User Request - └─> Omni.generate(prompt) - └─> Prepare OmniDiffusionRequest - └─> DiffusionEngine.step(requests) - -2. Pre-processing - └─> pre_process_func(requests) - └─> Model-specific transformations - -3. Scheduling - └─> scheduler.add_request(request) - └─> scheduler.schedule() - └─> DiffusionEngine submits scheduled request to executor.add_req(req) - -4. Worker Execution - └─> WorkerProc.worker_busy_loop() - └─> GPUWorker.execute_model(reqs) - └─> Pipeline.forward(req) - ├─> encode_prompt() - ├─> prepare_latents() - ├─> diffuse() [loop] - │ ├─> transformer.forward() [with cache backend hooks] - │ └─> scheduler.step() - └─> vae.decode() - -5. Result Collection - └─> Executor returns DiffusionOutput - └─> scheduler.update_from_output(...) - └─> DiffusionEngine pops finished request state - -6. Post-processing - └─> post_process_func(output) - └─> Convert to PIL images / final format -``` - ---- diff --git a/docs/design/module/engine_orchestration.md b/docs/design/module/engine_orchestration.md new file mode 100644 index 000000000..cf8ef1fd0 --- /dev/null +++ b/docs/design/module/engine_orchestration.md @@ -0,0 +1,52 @@ +--- +title: Engine Orchestration +kind: module +status: draft +owners: + - "@tzhouam" + - "@yinpeiqi" + - "@fake0fan" + - "@Sy0307" +primary_code_paths: + - vllm_omni/engine/** +related_code_paths: + - vllm_omni/distributed/omni_coordinator/** + - vllm_omni/distributed/ray_utils/** +depends_on: + - vllm_omni_config.md + - input_output_modality_contracts.md + - omni_connector.md +validation_paths: + - tests/engine/** + - tests/distributed/omni_coordinator/** +upstream_refs: + - vllm.v1.engine +last_reviewed: 2026-07-16 +--- + +# Engine orchestration + +Engine orchestration coordinates configured AR and diffusion stages while +keeping public clients independent from stage processes and transports. + +## Candidate invariants + +### ORCH-INV-001: The orchestrator owns cross-stage routing + +**Rule:** Entrypoints and stage clients MUST NOT independently forward a +request to a downstream 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 + +**Rule:** Once a request reaches a terminal state, orchestration MUST NOT +forward new work for that request. + +## Safe-change guide + +Test routing, output ordering, cancellation, failure propagation, shutdown, and +representative multi-stage execution. diff --git a/docs/design/module/entrypoint_module.md b/docs/design/module/entrypoint_module.md deleted file mode 100644 index 7a26fbb7f..000000000 --- a/docs/design/module/entrypoint_module.md +++ /dev/null @@ -1 +0,0 @@ -Architecture design of the entrypoint (update soon) diff --git a/docs/design/module/entrypoints.md b/docs/design/module/entrypoints.md new file mode 100644 index 000000000..8d4da8434 --- /dev/null +++ b/docs/design/module/entrypoints.md @@ -0,0 +1,51 @@ +--- +title: Entrypoints +kind: module +status: draft +owners: + - "@tzhouam" + - "@yinpeiqi" + - "@fake0fan" + - "@alex-jw-brooks" +primary_code_paths: + - vllm_omni/entrypoints/** +related_code_paths: + - vllm_omni/reasoning/** +depends_on: + - vllm_omni_config.md + - input_output_modality_contracts.md + - engine_orchestration.md +validation_paths: + - tests/entrypoints/** + - tests/reasoning/** +upstream_refs: + - vllm.entrypoints +last_reviewed: 2026-07-16 +--- + +# Entrypoints + +Entrypoints translate offline, CLI, and serving requests into stable engine +operations and translate engine outputs into public responses. + +## Candidate invariants + +### 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 + +**Rule:** Public protocol values MUST be validated and converted to internal +request contracts before engine submission. + +### ENTRY-INV-003: Streaming preserves request identity + +**Rule:** Every streamed response MUST remain associated with the request and +output modality that produced it. + +## Safe-change guide + +Test request validation, protocol conversion, streaming, cancellation, and +error mapping for each affected entrypoint. diff --git a/docs/design/module/execution_platforms.md b/docs/design/module/execution_platforms.md new file mode 100644 index 000000000..beeeef922 --- /dev/null +++ b/docs/design/module/execution_platforms.md @@ -0,0 +1,48 @@ +--- +title: Execution Platforms +kind: module +status: draft +owners: + - "@gcanlin" + - "@tjtanaa" + - "@xuechendi" +primary_code_paths: + - vllm_omni/platforms/** +related_code_paths: + - vllm_omni/attention/** +depends_on: + - ar_runtime.md + - diffusion/index.md +validation_paths: + - tests/platforms/** +upstream_refs: + - vllm.platforms +last_reviewed: 2026-07-16 +--- + +# Execution platforms + +Execution platforms isolate hardware-specific capability detection, worker +selection, patches, kernels, and configuration adjustments. + +## Candidate invariants + +### PLATFORM-INV-001: Capabilities are explicit + +**Rule:** Hardware-dependent behavior MUST be guarded by platform selection or +capability detection and MUST fail clearly when unsupported. + +### PLATFORM-INV-002: Portable code remains portable + +**Rule:** Platform-neutral modules MUST NOT import a vendor implementation +directly when a common selection interface exists. + +### PLATFORM-INV-003: Overrides are minimal + +**Rule:** A platform implementation SHOULD override only behavior that differs +from the common or upstream implementation. + +## Safe-change guide + +Validate common import paths without the target accelerator and run focused +tests in a fresh `uv` environment on every affected device platform. diff --git a/docs/design/module/input_output_modality_contracts.md b/docs/design/module/input_output_modality_contracts.md new file mode 100644 index 000000000..10d803d26 --- /dev/null +++ b/docs/design/module/input_output_modality_contracts.md @@ -0,0 +1,47 @@ +--- +title: Input, Output, and Modality Contracts +kind: module +status: draft +owners: + - "@tzhouam" + - "@gcanlin" +primary_code_paths: + - vllm_omni/inputs/** + - vllm_omni/outputs/** +related_code_paths: + - vllm_omni/tokenizers/** +depends_on: [] +validation_paths: + - tests/inputs/** +upstream_refs: + - vllm.inputs + - vllm.outputs +last_reviewed: 2026-07-16 +--- + +# Input, output, and modality contracts + +These contracts define the data that may cross entrypoint, orchestration, +stage, and model boundaries. + +## Candidate invariants + +### 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 + +**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 + +**Rule:** Entrypoints MUST explicitly translate internal output objects into +public response types. + +## Safe-change guide + +Test construction, validation, serialization, streaming, optional fields, and +compatibility at every affected producer-consumer boundary. diff --git a/docs/design/module/model_integration.md b/docs/design/module/model_integration.md new file mode 100644 index 000000000..8b2afe0a0 --- /dev/null +++ b/docs/design/module/model_integration.md @@ -0,0 +1,51 @@ +--- +title: Model Integration +kind: module +status: draft +owners: + - "@tzhouam" + - "@gcanlin" +primary_code_paths: + - vllm_omni/model_executor/** + - vllm_omni/model_extras/** + - vllm_omni/plugins/** +related_code_paths: + - vllm_omni/transformers_utils/** + - vllm_omni/tokenizers/** +depends_on: + - input_output_modality_contracts.md + - ar_runtime.md +validation_paths: + - tests/model_executor/** + - tests/model_extras/** + - tests/plugin/** +upstream_refs: + - vllm.model_executor +last_reviewed: 2026-07-16 +--- + +# Model integration + +Model integration adapts model-specific configuration, preprocessing, loading, +and execution behavior to stable runtime contracts. + +## Candidate invariants + +### MODEL-INV-001: Registration is explicit + +**Rule:** A model integration MUST declare how its model class, loader, input +processor, and stage configuration are selected. + +### MODEL-INV-002: Model code does not route stages + +**Rule:** Model-specific code MUST NOT select or invoke a downstream omni stage. + +### MODEL-INV-003: Upstream behavior is reused deliberately + +**Rule:** An upstream vLLM implementation SHOULD be reused when its contract is +sufficient; an override MUST document the behavioral difference. + +## Safe-change guide + +Test registration, checkpoint loading, input conversion, and representative +model execution. Test shared utilities against more than one integration. diff --git a/docs/design/module/observability.md b/docs/design/module/observability.md new file mode 100644 index 000000000..83971e2d3 --- /dev/null +++ b/docs/design/module/observability.md @@ -0,0 +1,46 @@ +--- +title: Observability +kind: module +status: draft +owners: + - "@lishunyang12" +primary_code_paths: + - vllm_omni/metrics/** +related_code_paths: + - vllm_omni/engine/** + - vllm_omni/entrypoints/** +depends_on: + - engine_orchestration.md +validation_paths: + - tests/metrics/** +upstream_refs: + - vllm.v1.metrics +last_reviewed: 2026-07-16 +--- + +# Observability + +Observability defines stable metrics, logs, and correlation fields for tracing a +request across entrypoints, orchestration, stages, and workers. + +## Candidate invariants + +### OBS-INV-001: Request correlation survives stage boundaries + +**Rule:** Logs and metrics for the same request SHOULD carry a stable request +identifier and stage identity. + +### OBS-INV-002: Metric meaning is stable + +**Rule:** A metric's unit, labels, aggregation, and lifecycle MUST be documented; +changing any of them is a contract change. + +### OBS-INV-003: Labels are bounded + +**Rule:** Metrics MUST NOT use request IDs, prompts, outputs, or other unbounded +values as labels. + +## Safe-change guide + +Test registration, label sets, units, lifecycle resets, and representative +multi-stage emission. diff --git a/docs/design/module/omni_connector.md b/docs/design/module/omni_connector.md new file mode 100644 index 000000000..875f79e3d --- /dev/null +++ b/docs/design/module/omni_connector.md @@ -0,0 +1,50 @@ +--- +title: OmniConnector +kind: module +status: draft +owners: + - "@princepride" + - "@yuanheng-zhao" + - "@xuechendi" + - "@natureofnature" + - "@fake0fan" +primary_code_paths: + - vllm_omni/distributed/omni_connectors/** +related_code_paths: + - vllm_omni/platforms/*/omni_connectors/** +depends_on: + - input_output_modality_contracts.md + - vllm_omni_config.md +validation_paths: + - tests/distributed/omni_connectors/** +upstream_refs: + - vllm.distributed +last_reviewed: 2026-07-16 +--- + +# OmniConnector + +OmniConnector defines model-agnostic transport and synchronization contracts +for data exchanged across stages, processes, devices, and nodes. + +## Candidate invariants + +### CONNECTOR-INV-001: Connectors transport but do not route + +**Rule:** A connector MUST NOT select the next stage or implement model-specific +execution policy. + +### CONNECTOR-INV-002: Producer and consumer contracts agree + +**Rule:** Both ends MUST agree on data identity, shape, dtype, placement, +ownership, and completion semantics. + +### CONNECTOR-INV-003: Resources have deterministic cleanup + +**Rule:** Connections, buffers, handles, and background work MUST be released on +normal completion, cancellation, failure, and shutdown. + +## Safe-change guide + +Test setup, transfer, synchronization, timeout, cancellation, failure, and +cleanup across every affected backend. diff --git a/docs/design/module/profiling.md b/docs/design/module/profiling.md new file mode 100644 index 000000000..c2f1a9f88 --- /dev/null +++ b/docs/design/module/profiling.md @@ -0,0 +1,46 @@ +--- +title: Profiling +kind: module +status: draft +owners: + - "@gcanlin" +primary_code_paths: + - vllm_omni/profiler/** + - vllm_omni/diffusion/profiler/** +related_code_paths: + - vllm_omni/entrypoints/** +depends_on: + - observability.md +validation_paths: + - tests/profile/** +upstream_refs: + - vllm.profiler +last_reviewed: 2026-07-16 +--- + +# Profiling + +Profiling owns opt-in instrumentation and trace collection used to locate +runtime bottlenecks without changing model results or execution policy. + +## Candidate invariants + +### PROFILE-INV-001: Profiling is opt-in + +**Rule:** Expensive trace collection MUST be disabled by default and bounded by +an explicit start and stop lifecycle. + +### PROFILE-INV-002: Instrumentation preserves semantics + +**Rule:** Profiling MUST NOT change scheduling decisions, generated outputs, or +resource ownership beyond documented synchronization overhead. + +### PROFILE-INV-003: Trace context is sufficient + +**Rule:** Trace events SHOULD identify stage, rank, worker, operation, and +request where available without embedding user payloads. + +## Safe-change guide + +Test disabled overhead, repeated start and stop, cleanup, multi-rank trace +naming, and each supported profiler backend. diff --git a/docs/design/module/quantization.md b/docs/design/module/quantization.md new file mode 100644 index 000000000..3e9001aab --- /dev/null +++ b/docs/design/module/quantization.md @@ -0,0 +1,50 @@ +--- +title: Quantization +kind: module +status: draft +owners: + - "@david6666666" + - "@Isotr0py" + - "@lishunyang12" +primary_code_paths: + - vllm_omni/quantization/** + - vllm_omni/diffusion/quantization/** +related_code_paths: + - vllm_omni/platforms/*/quant/** +depends_on: + - model_integration.md + - diffusion/diffusion_model_integration.md +validation_paths: + - tests/quantization/** + - tests/diffusion/quantization/** +upstream_refs: + - vllm.model_executor.layers.quantization +last_reviewed: 2026-07-16 +--- + +# Quantization + +Quantization owns method registration, checkpoint interpretation, layer +replacement, and runtime constraints for reduced-precision execution. + +## Candidate invariants + +### QUANT-INV-001: Checkpoint and runtime methods agree + +**Rule:** The selected method MUST validate checkpoint metadata, supported +layers, device capabilities, and required dependencies before use. + +### QUANT-INV-002: Unsupported combinations fail clearly + +**Rule:** Quantization MUST NOT silently fall back to a different precision or +method when a requested combination is unsupported. + +### QUANT-INV-003: Numerical validation is required + +**Rule:** A new method or layer mapping MUST include correctness comparison +against a defined reference result. + +## Safe-change guide + +Test loading, representative layers, metadata, platform guards, and numerical +tolerances separately from performance. diff --git a/docs/design/module/vllm_omni_config.md b/docs/design/module/vllm_omni_config.md new file mode 100644 index 000000000..5df0cdcb3 --- /dev/null +++ b/docs/design/module/vllm_omni_config.md @@ -0,0 +1,47 @@ +--- +title: vLLM-Omni Configuration +kind: module +status: draft +owners: + - "@lishunyang12" + - "@alex-jw-brooks" +primary_code_paths: + - vllm_omni/config/** + - vllm_omni/deploy/** + - vllm_omni/model_executor/stage_configs/** +related_code_paths: + - vllm_omni/platforms/*/stage_configs/** +depends_on: [] +validation_paths: + - tests/config/** +upstream_refs: + - vllm.config +last_reviewed: 2026-07-16 +--- + +# vLLM-Omni configuration + +Configuration defines how user input becomes validated engine, stage, +deployment, and platform settings. + +## Candidate invariants + +### CONFIG-INV-001: Configuration is validated before startup + +**Rule:** Invalid stage graphs, resource assignments, and incompatible options +MUST fail before managed runtime processes start. + +### CONFIG-INV-002: Precedence is deterministic + +**Rule:** Defaults, files, environment variables, and CLI overrides MUST have a +documented and deterministic precedence. + +### CONFIG-INV-003: Runtime modules consume validated configuration + +**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.