docs: reorganize module design documents
Signed-off-by: hsliuustc0106 <liuhongsheng4@huawei.com>
This commit is contained in:
hsliuustc0106 committed
1 parent
238fc0a609
commit
80fd96c1da
23 files changed
+952
-1537
No files matched your search
@@ -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.
|
||||
@@ -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.
|
||||
@@ -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
|
||||
```
|
||||
@@ -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.
|
||||
@@ -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.
|
||||
@@ -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.
|
||||
@@ -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.
|
||||
@@ -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.
|
||||
@@ -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.
|
||||
@@ -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.
|
||||
@@ -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.
|
||||
@@ -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.
|
||||
|
||||
<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 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.<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.
|
||||
|
||||
`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/<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/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
|
||||
|
||||
<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(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
|
||||
```
|
||||
|
||||
---
|
||||
@@ -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.
|
||||
@@ -1 +0,0 @@
|
||||
Architecture design of the entrypoint (update soon)
|
||||
@@ -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.
|
||||
@@ -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.
|
||||
@@ -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.
|
||||
@@ -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.
|
||||
@@ -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.
|
||||
@@ -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.
|
||||
@@ -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.
|
||||
@@ -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.
|
||||
@@ -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.
|
||||
Reference in new issue
Block a user