docs: define module ownership contracts
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
This commit is contained in:
Hongsheng Liu committed
commit
01e7d09119
1300 files changed
+167316
-34596
No files matched your search
@@ -0,0 +1,165 @@
|
||||
# AR-Diffusion pipeline capability
|
||||
|
||||
> **Status:** experimental and non-normative. This note documents the current
|
||||
> `experimental/ar_diffusion` contract. It follows the ownership and lifecycle
|
||||
> direction in [#5137](https://github.com/vllm-project/vllm-omni/issues/5137),
|
||||
> but it is not yet part of the stable module-design hierarchy.
|
||||
|
||||
`ARDiffusionEngine` is an opt-in diffusion backend for pipelines that keep
|
||||
autoregressive attention KV across requests. The engine only selects
|
||||
`ARDiffusionModelRunner`; the runner owns paged KV pools, session lookup, LRU
|
||||
eviction, and failure cleanup.
|
||||
|
||||
A pipeline opts in by implementing `SupportsARDiffusionPipeline`. Its immutable
|
||||
`ARDiffusionKVCacheSpec` reports worker-local KV geometry and policy:
|
||||
|
||||
- self-attention layers, TP-local KV heads, head size, and tokens per latent frame;
|
||||
- latent frames committed by one model block, sliding window, and attention sink;
|
||||
- logical KV branches mapped to worker-local storage slots;
|
||||
- named fixed-length cross-attention KV, model-specific scratch geometry, and
|
||||
the pipeline's requested session-capacity upper bound.
|
||||
|
||||
The runner binds an `ARDiffusionKVState` only around one `forward()` call. The
|
||||
pipeline may use it during that context but must not retain it. The session
|
||||
survives subsequent requests with the same `session_id`. Request reset, explicit
|
||||
close, LRU eviction, and forward exceptions all release runner-owned KV and call
|
||||
the pipeline's reset/close lifecycle hook. Cross-attention KV is allocated lazily
|
||||
per session and logical KV branch, and is released on the same path. A pipeline
|
||||
publishes one named cache transactionally by yielding exactly one `(k, v)` pair
|
||||
per layer:
|
||||
|
||||
```python
|
||||
state.populate_cross_attention("main", "text", project_text_kv_by_layer())
|
||||
```
|
||||
|
||||
The cache becomes readable only after the iterable completes successfully.
|
||||
Logical branches remain independent even when they share a worker-local
|
||||
`local_index`; that index only maps lazily executed self-attention and scratch
|
||||
work to worker-local capacity slots. It does not define cross-attention cache
|
||||
ownership.
|
||||
|
||||
Requests select the session with `extra_args["session_id"]`; setting `reset=True`
|
||||
replaces it before the forward, while `close_session=True` releases it after a
|
||||
successful forward. The runner also exposes `reset_session()` and
|
||||
`close_session()` for an owning serving layer.
|
||||
|
||||
## Current execution limits
|
||||
|
||||
The experimental runner supports request-mode execution with
|
||||
`max_num_seqs=1`. It rejects step execution and request-batch execution at load
|
||||
time, and its inherited batch/step entry points also fail explicitly. Supporting
|
||||
those paths requires a batch-aware state-binding contract rather than silently
|
||||
bypassing `bind_ar_diffusion_state()`.
|
||||
|
||||
The current implementation retains one session per model instance. A request
|
||||
for a different `session_id` closes the resident session through the same LRU
|
||||
lifecycle used by explicit close and failure cleanup before allocating the new
|
||||
one. Self-attention and cross-attention memory are therefore sized for one
|
||||
resident session even if the pipeline declares a larger `session_capacity`.
|
||||
This prioritizes rollout throughput and prevents block-pool exhaustion before a
|
||||
count-based eviction can run. Multi-session residency can be added later with a
|
||||
memory-backed capacity calculation or allocation-pressure eviction.
|
||||
|
||||
Scratch storage is also capability-driven. `frames_per_block` reserves space for
|
||||
an uncommitted current video block, while `max_scratch_tokens_per_branch`
|
||||
declares the maximum model-specific tokens, such as action/state registers, that
|
||||
must coexist with it. `AR_DIFFUSION_KV_SCRATCH_BLOCKS_PER_BRANCH` may increase
|
||||
that derived minimum for deployment experiments but cannot reduce it.
|
||||
|
||||
## DreamZero adapter
|
||||
|
||||
DreamZero implements the capability directly on `DreamZeroPipeline`. Its adapter
|
||||
reads DreamZero transformer geometry, describes the positive and negative CFG
|
||||
branches, binds the state with a context manager, and translates the existing
|
||||
boolean CFG calls to the explicit `"positive"` and `"negative"` KV branch names.
|
||||
Its optional `ar_diffusion_warmup_requests()` provider owns the robot observation
|
||||
schema and the one-frame/four-frame warmup convention; none of those details are
|
||||
present in the generic runner.
|
||||
|
||||
Code migrating from the former experimental KV state API should use the
|
||||
mechanical mapping below:
|
||||
|
||||
```python
|
||||
# Before: state.pos, state.neg, state.get_kv_caches(is_negative, ...)
|
||||
positive = state.adapter("positive")
|
||||
negative = state.adapter("negative")
|
||||
contexts = state.get_kv_caches("negative", seq_len=frame_tokens, commit_current=True)
|
||||
```
|
||||
|
||||
Direct experimental `ARDiffusionKVCache` construction also migrates
|
||||
mechanically: replace `local_branches` with explicit `kv_branches`,
|
||||
`num_frame_per_block` with `frames_per_block`, and the fixed text/image length
|
||||
arguments with a named `cross_attention_lengths` mapping plus
|
||||
`session_capacity`. Normal engine users do not construct this object directly.
|
||||
|
||||
Cross-attention writers migrate from incremental publication:
|
||||
|
||||
```python
|
||||
# Before
|
||||
for layer_index, (key, value) in enumerate(project_text_kv_by_layer()):
|
||||
state.write_cross_attention_kv("text", layer_index, branch, key, value)
|
||||
state.mark_cross_attention_populated(branch, "text")
|
||||
|
||||
# After
|
||||
state.populate_cross_attention(branch, "text", project_text_kv_by_layer())
|
||||
```
|
||||
|
||||
The iterable is consumed once and must yield exactly `num_layers` pairs. Failed
|
||||
or incomplete projection leaves the previous complete cache unchanged, or
|
||||
leaves the cache absent on its first population.
|
||||
|
||||
## Single-KV-branch pipeline sketch
|
||||
|
||||
A future single-branch causal DMD pipeline can opt in without changes to the
|
||||
runner or KV core:
|
||||
|
||||
```python
|
||||
class WorldPipeline:
|
||||
def ar_diffusion_kv_cache_spec(self):
|
||||
return ARDiffusionKVCacheSpec(
|
||||
num_layers=self.transformer.num_layers,
|
||||
num_kv_heads=self.transformer.tp_local_kv_heads,
|
||||
head_size=self.transformer.head_size,
|
||||
tokens_per_frame=self.transformer.tokens_per_latent_frame,
|
||||
frames_per_block=3,
|
||||
window_frames=self.transformer.sliding_window_frames,
|
||||
sink_frames=self.transformer.sink_frames,
|
||||
kv_branches=(ARDiffusionKVBranchSpec("main", 0),),
|
||||
cross_attention=(ARDiffusionCrossAttentionKVSpec("text", self.text_length),),
|
||||
max_scratch_tokens_per_branch=self.transformer.max_register_tokens,
|
||||
session_capacity=1,
|
||||
)
|
||||
|
||||
@contextmanager
|
||||
def bind_ar_diffusion_state(self, session_id, state):
|
||||
self.current_ar_state = state
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
self.current_ar_state = None
|
||||
|
||||
def reset_ar_diffusion_session(self, session_id):
|
||||
self.model_states.pop(session_id, None)
|
||||
|
||||
def close_ar_diffusion_session(self, session_id):
|
||||
self.model_states.pop(session_id, None)
|
||||
```
|
||||
|
||||
The model writes/reads named cross-attention KV through the bound state and uses
|
||||
the `"main"` KV branch for paged self-attention. Pipelines without a warmup
|
||||
provider are loaded normally and skip AR rollout warmup.
|
||||
|
||||
## Relationship to the multi-stage KV manager RFC
|
||||
|
||||
[#5244](https://github.com/vllm-project/vllm-omni/issues/5244) describes a
|
||||
future multi-stage KV transfer manager, connector integration, and cache-aware
|
||||
scheduling. This capability has a narrower scope: it owns model-local paged KV
|
||||
and session lifecycle inside one AR-Diffusion runner. It does not define
|
||||
cross-stage transfer, connector behavior, or scheduler-visible cache affinity.
|
||||
|
||||
When the experimental runtime migrates toward #5244, the implementation should
|
||||
reconcile its current full `KVCacheManager` use with the RFC's proposed
|
||||
`BlockPool`-oriented DiT manager. The pipeline capability should remain the
|
||||
model-facing geometry and state-binding boundary where possible, while the
|
||||
experimental `ARDiffusionKVCache` class itself should not be treated as a stable
|
||||
public API.
|
||||
@@ -0,0 +1,174 @@
|
||||
# Async Diffusion Output
|
||||
|
||||
## Table of Contents
|
||||
|
||||
1. [Overview](#overview)
|
||||
2. [Performance](#performance)
|
||||
3. [Architecture](#architecture)
|
||||
4. [Model Coverage](#model-coverage)
|
||||
5. [Configuration](#configuration)
|
||||
6. [Related Files](#related-files)
|
||||
|
||||
## Overview
|
||||
|
||||
The async diffusion output feature moves the D2H (Device→Host) copy and SHM packing from the Worker main thread to a background daemon thread, so the GPU default stream can immediately start the next request's forward pass. This eliminates the GPU bubble caused by synchronous D2H/packing.
|
||||
|
||||
**Core benefit**: the previous request's output D2H/packing overlaps with the next request's forward pass.
|
||||
|
||||
**Automatically enabled** when `step_execution=False` (default, request-mode). No extra configuration needed. When `step_execution=True` (step-mode), the original synchronous path is used — pump thread and Worker background thread are not started.
|
||||
|
||||
## Performance
|
||||
|
||||
### HunyuanImage-3.0 (TP4)
|
||||
|
||||
| Resolution | Async Off (QPS) | Async On (QPS) | Change |
|
||||
|---|---|---|---|
|
||||
| 1024×1024 | 0.4773 | 0.4802 | +0.60% |
|
||||
| 768×768 | 0.8370 | 0.8533 | +1.95% |
|
||||
|
||||
**Analysis**: smaller resolution → larger QPS gain. D2H/packing is relatively fixed overhead; forward is variable (scales with resolution). Smaller images → shorter forward → D2H/packing occupies a larger fraction of total time → more benefit from overlapping it.
|
||||
|
||||
## Architecture
|
||||
|
||||
### Execution Timeline
|
||||
|
||||
```
|
||||
Before (synchronous D2H):
|
||||
[forward req1] [D2H+SHM pack req1] [forward req2] [D2H+SHM pack req2]
|
||||
^^^^^^^^^^^^^^^^ GPU bubble
|
||||
|
||||
After (async D2H):
|
||||
[forward req1] [forward req2]
|
||||
[D2H+SHM pack req1 (side stream)] [D2H+SHM pack req2 (side stream)]
|
||||
^^^^^^^^^^^^^^^^ bubble eliminated
|
||||
```
|
||||
|
||||
### Data Flow
|
||||
|
||||
```
|
||||
┌─────────────────────────────────────────────────────┐
|
||||
│ Worker Process │
|
||||
│ │
|
||||
execute_model ──> │ WorkerBusyLoop ──> forward() ──> DiffusionOutput │
|
||||
│ │ │ (GPU tensor)│
|
||||
│ │ ┌─────────┴──────────┐ │
|
||||
│ │ │ async output path │ │
|
||||
│ │ │ │ │
|
||||
│ │ compute_done ──> result_mq │ │
|
||||
│ │ │ │ │
|
||||
│ │ ▼ │ │
|
||||
│ │ AsyncOutputThread │ │
|
||||
│ │ (side CUDA stream) │ │
|
||||
│ │ wait_event(gpu_event) │ │
|
||||
│ │ pack_diffusion_output_shm() │ │
|
||||
│ │ D2H + SHM write │ │
|
||||
│ │ │ │ │
|
||||
│ │ output_ready ──> result_mq │ │
|
||||
│ │ │ │ │
|
||||
│ ▼ ▼ │ │
|
||||
│ (dequeue next request) │ │
|
||||
└─────────────────────────────────────────────────────┘
|
||||
│
|
||||
▼
|
||||
┌─────────────────────────────────────────────────────┐
|
||||
│ Executor (main process) │
|
||||
│ │
|
||||
│ ResultPumpThread (sole reader of result_mq) │
|
||||
│ │ │
|
||||
│ ├── AsyncDiffusionOutput(COMPUTE_DONE) │
|
||||
│ │ → resolve _rpc_futures[rpc_id] │
|
||||
│ │ → RunnerOutput(async_output_id=...) │
|
||||
│ │ │
|
||||
│ ├── AsyncDiffusionOutput(OUTPUT_READY) │
|
||||
│ │ → resolve _output_futures[async_output_id]│
|
||||
│ │ → or batch split via _batch_split_map │
|
||||
│ │ │
|
||||
│ └── Non-async message │
|
||||
│ → _sync_result_buffer (for other RPCs) │
|
||||
│ │
|
||||
│ wait_output_ready(async_output_id) → Future │
|
||||
└─────────────────────────────────────────────────────┘
|
||||
```
|
||||
|
||||
### Key Components
|
||||
|
||||
1. **AsyncDiffusionOutput** (`data.py`): Protocol envelope for `result_mq`. `kind` field routes messages:
|
||||
- `COMPUTE_DONE` — forward finished, GPU can start next request
|
||||
- `OUTPUT_READY` — D2H/SHM packing finished, final output available
|
||||
- `RPC_RESULT` — ordinary RPC return (including error propagation)
|
||||
|
||||
2. **Worker AsyncOutputThread** (`diffusion_worker.py`): Background daemon thread that:
|
||||
- Waits for `gpu_event` (cross-stream sync ensuring forward wrote the tensor)
|
||||
- Runs `pack_diffusion_output_shm()` on a side CUDA stream
|
||||
- Enqueues `OUTPUT_READY` to `result_mq` when packing completes
|
||||
|
||||
3. **ResultPumpThread** (`multiproc_executor.py`): Sole reader of `result_mq` in request-mode. Dispatches messages:
|
||||
- `COMPUTE_DONE` / `RPC_RESULT` → resolves `_rpc_futures[rpc_id]`
|
||||
- `OUTPUT_READY` → resolves `_output_futures[async_output_id]` (with batch split via `_batch_split_map`)
|
||||
- Non-async messages → `_sync_result_buffer` for other RPCs
|
||||
|
||||
4. **collective_rpc() two-path dispatch** (`multiproc_executor.py`):
|
||||
- **Path 1**: `execute_model` / `execute_model_batch` → generates `rpc_id`, registers Future, waits for pump to deliver `compute_done`
|
||||
- **Path 2**: All other RPCs → `_sync_result_buffer` (pump-fed) or `_result_mq` (step-mode, no pump)
|
||||
|
||||
5. **record_device_event** (`platforms/`): Cross-platform GPU event recording for side-stream synchronization. Implemented for CUDA, NPU, ROCm, XPU, MUSA; base class returns `None` (safe no-op).
|
||||
|
||||
6. **IPC side-stream D2H** (`ipc.py`): `pack_diffusion_output_shm()` and helpers accept `d2h_stream` parameter; use `pin_memory` + `copy_(non_blocking=True)` on the side stream instead of synchronous `.cpu()`.
|
||||
|
||||
### Batch Split
|
||||
|
||||
When `execute_model_batch` returns a `COMPUTE_DONE` with a single `async_output_id` for the entire batch, the executor splits it into per-request `async_output_id`s (format `{batch_id}/{request_id}`) via `_batch_split_map`. When `OUTPUT_READY` arrives, the pump extracts per-request results from the batch output and resolves each request's output future independently.
|
||||
|
||||
## Model Coverage
|
||||
|
||||
### Currently Supported (request-mode, `step_execution=False`)
|
||||
|
||||
All models running in request-mode (`step_execution=False`, the default) automatically use async output.
|
||||
|
||||
**Verified models**:
|
||||
|
||||
| Model | Type | `supports_request_batch` |
|
||||
|---|---|---|
|
||||
| **HunyuanImage-3.0** | Image | `False` |
|
||||
| **Qwen-Image** | Image | `True` |
|
||||
|
||||
Other models with `step_execution=False` are also supported but not yet verified.
|
||||
|
||||
### Not Yet Supported (step-mode, `step_execution=True`)
|
||||
|
||||
When `step_execution=True` (or `streaming_output=True`, which auto-enables step-mode), models use `execute_stepwise` instead of `execute_model`, which is not in the async Path 1 whitelist. Async output is not applicable.
|
||||
|
||||
**Verified**: Helios (`step_execution=True`, does not benefit from this feature).
|
||||
|
||||
Other models with `step_execution=True` are also not applicable.
|
||||
|
||||
Adapting step-mode requires additional design because:
|
||||
|
||||
- Each step's `RunnerOutput` must be available synchronously for the scheduler to advance
|
||||
- Intermediate latent tensors may not need D2H (only the final step needs D2H + postprocess)
|
||||
- Steps are serial (step N+1 depends on step N output), so the overlap benefit is limited to "final step's D2H overlaps with the next request's forward"
|
||||
|
||||
**Planned approach**: first cover only the final step's async D2H, keep intermediate steps synchronous. This matches the request-mode benefit pattern.
|
||||
|
||||
## Configuration
|
||||
|
||||
No configuration needed. Async output is automatically enabled when `step_execution=False` (default).
|
||||
|
||||
| `step_execution` | Mode | Async Output | Pump Thread | Worker BG Thread |
|
||||
|---|---|---|---|---|
|
||||
| `False` (default) | request-mode | ✅ enabled | ✅ started | ✅ started |
|
||||
| `True` | step-mode | ❌ disabled | ❌ not started | ❌ not started |
|
||||
|
||||
## Related Files
|
||||
|
||||
- `vllm_omni/diffusion/data.py`: `AsyncDiffusionOutput`, `AsyncOutputKind`, `DiffusionOutput.async_output_id`
|
||||
- `vllm_omni/diffusion/worker/diffusion_worker.py`: `WorkerProc._return_result()`, `_async_output_loop()`, `_generate_async_output_id()`
|
||||
- `vllm_omni/diffusion/executor/multiproc_executor.py`: `ResultPumpThread`, `collective_rpc()` two-path dispatch, `wait_output_ready()`, `_batch_split_map`
|
||||
- `vllm_omni/diffusion/diffusion_engine.py`: `step()` / `step_streaming()` / `add_req_and_wait_for_response()` async output waiting
|
||||
- `vllm_omni/diffusion/sched/request_scheduler.py`: `update_from_output()` — `async_output_id` → `FINISHED_COMPLETED`
|
||||
- `vllm_omni/diffusion/ipc.py`: `pack_diffusion_output_shm()` side-stream D2H path
|
||||
- `vllm_omni/platforms/interface.py`: `OmniPlatform.record_device_event()` base class
|
||||
- `vllm_omni/platforms/cuda/platform.py`: `CudaOmniPlatform.record_device_event()`
|
||||
- `vllm_omni/platforms/npu/platform.py`: `NPUOmniPlatform.record_device_event()`
|
||||
- `vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py`: `get_hunyuan_image3_post_process_func()`, postprocess removal
|
||||
- `vllm_omni/diffusion/registry.py`: `_DIFFUSION_POST_PROCESS_FUNCS` postprocess registry
|
||||
@@ -32,10 +32,13 @@ The library supports three main caching strategies:
|
||||
|
||||
vLLM-omni integrates cache-dit through the `CacheDiTBackend` class, which provides a unified interface for managing cache-dit acceleration on diffusion models.
|
||||
|
||||
Code outside `vllm_omni/diffusion/cache` must import Cache-DiT symbols from the
|
||||
package-level API: `vllm_omni.diffusion.cache.cachedit`.
|
||||
|
||||
| Method/Class | Purpose | Behavior |
|
||||
|--------------|---------|----------|
|
||||
| [`CacheDiTBackend`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/cache/#vllm_omni.diffusion.cache.CacheBackend) | Unified backend interface | Automatically handles enabler selection and cache refresh |
|
||||
| [`enable_cache_for_dit()`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/cache/cache_dit_backend/#vllm_omni.diffusion.cache.cache_dit_backend.enable_cache_for_dit) | Apply caching to transformer | Configures DBCache on transformer blocks |
|
||||
| `CacheDiTBackend` | Unified backend interface | Automatically handles enabler selection and cache refresh |
|
||||
| `enable_cache_for_dit()` | Apply caching to transformer | Configures DBCache on transformer blocks |
|
||||
|
||||
**Key APIs from Cache-DiT:**
|
||||
|
||||
@@ -47,7 +50,7 @@ vLLM-omni integrates cache-dit through the `CacheDiTBackend` class, which provid
|
||||
| `ForwardPattern` | Defines block forward signature patterns: `Pattern_0`, `Pattern_1`, `Pattern_2` |
|
||||
| `ParamsModifier` | Per-transformer or per-block-list cache configuration customization |
|
||||
| `DBCacheConfig` | Configuration for DBCache parameters (warmup steps, cached steps, thresholds) |
|
||||
| `refresh_context()` | Update cache context | Called when `num_inference_steps` changes |
|
||||
| `refresh_context()` | Update cache context | Called at every generation boundary |
|
||||
|
||||
---
|
||||
|
||||
@@ -179,7 +182,8 @@ cache_dit.enable_cache(
|
||||
|
||||
### Registering Custom Implementations
|
||||
|
||||
After writing your custom enabler, register it in `CUSTOM_DIT_ENABLERS` in `vllm_omni/diffusion/cache/cache_dit_backend.py`:
|
||||
After writing your custom enabler, register it in `CUSTOM_DIT_ENABLERS` in
|
||||
`vllm_omni/diffusion/cache/cachedit/model_specific.py`:
|
||||
|
||||
```python
|
||||
CUSTOM_DIT_ENABLERS = {
|
||||
@@ -266,10 +270,10 @@ Complete examples in the codebase:
|
||||
|
||||
| Model | Path | Pattern | Notes |
|
||||
|-------|------|---------|-------|
|
||||
| **Standard DiT** | [`cache_dit_backend.py::enable_cache_for_dit`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/cache/cache_dit_backend/#vllm_omni.diffusion.cache.cache_dit_backend.enable_cache_for_dit) | Default enabler | Single transformer, automatic |
|
||||
| **Wan2.2** | [`cache_dit_backend.py::enable_cache_for_wan22`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/cache/cache_dit_backend/#vllm_omni.diffusion.cache.cache_dit_backend.enable_cache_for_wan22) | Single or dual-transformer | Auto-detects mode based on transformer_2 presence |
|
||||
| **LongCat** | [`cache_dit_backend.py::enable_cache_for_longcat_image`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/cache/cache_dit_backend/#vllm_omni.diffusion.cache.cache_dit_backend.enable_cache_for_longcat_image) | Multi-block-list | Two block lists in one transformer |
|
||||
| **BAGEL** | [`cache_dit_backend.py::enable_cache_for_bagel`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/cache/cache_dit_backend/#vllm_omni.diffusion.cache.cache_dit_backend.enable_cache_for_bagel) | Omni model | Complex architecture |
|
||||
| **Standard DiT** | `cachedit.backend::enable_cache_for_dit` | Default enabler | Single transformer, automatic |
|
||||
| **Wan2.2** | `cachedit.model_specific::enable_cache_for_wan22` | Single or dual-transformer | Auto-detects mode based on transformer_2 presence |
|
||||
| **LongCat** | `cachedit.config::CacheDiTAdapterConfig` | Declarative block adapter | Two block lists in one transformer |
|
||||
| **BAGEL** | `cachedit.model_specific::BagelCachedAdapter` | Custom cached adapter | Complex architecture |
|
||||
|
||||
---
|
||||
|
||||
|
||||
@@ -1,163 +1,232 @@
|
||||
# Continuous Batching for Step-Wise Diffusion
|
||||
# Diffusion Continuous Batching
|
||||
|
||||
!!! warning "Experimental Feature"
|
||||
This feature is experimental. It currently applies only to native
|
||||
diffusion pipelines running with `step_execution=True`.
|
||||
|
||||
This document describes the batching extension built on top of
|
||||
[Diffusion Step Execution](diffusion_step_execution.md). The base
|
||||
step-execution contract is unchanged. The batching work is mainly in the
|
||||
scheduler and runner layers.
|
||||
|
||||
## Why It Helps
|
||||
|
||||
Step-wise execution breaks a long denoise loop into scheduler-visible units.
|
||||
That gives the runtime a place to admit other compatible requests between
|
||||
steps instead of waiting for an entire request to finish.
|
||||
|
||||
This matters most in low-MFU or bursty serving scenarios:
|
||||
|
||||
- one request's denoise step may not fully saturate the GPU
|
||||
- several compatible requests can share the same denoise forward pass
|
||||
- throughput and device utilization can improve without changing request-local
|
||||
scheduler state
|
||||
|
||||
This is **not** a guaranteed single-request latency win. The main benefit is
|
||||
usually higher utilization and better throughput when the workload contains
|
||||
multiple in-flight compatible requests.
|
||||
This document describes the unified diffusion execution architecture, including
|
||||
request execution, request-level batching, step execution, continuous batching,
|
||||
and output delivery. For user-facing configuration and CLI examples, see
|
||||
[Diffusion Execution Modes](../../user_guide/diffusion/execution_modes.md).
|
||||
|
||||
## Overview
|
||||
|
||||
With continuous batching enabled:
|
||||
`DiffusionEngine` selects an execution mode from `step_execution` and uses
|
||||
`max_num_seqs` as that mode's scheduler capacity:
|
||||
|
||||
- the scheduler may keep multiple compatible requests active at the same time
|
||||
- the runner packs request-local step state into one `InputBatch`
|
||||
- `denoise_step()` runs on that batch
|
||||
- `step_scheduler()` and `post_decode()` still run per request
|
||||
| Configuration | Engine mode | Scheduler | Execution |
|
||||
|---|---|---|---|
|
||||
| `step_execution=False`, `max_num_seqs=1` | `REQUEST_BATCH` | `RequestScheduler` | One complete request-level `forward()` |
|
||||
| `step_execution=False`, `max_num_seqs>1` | `REQUEST_BATCH` | `RequestScheduler` | One fused `forward()` over compatible requests |
|
||||
| `step_execution=True`, `max_num_seqs=1` | `STEP_BATCH` | `StepScheduler` | One request advanced one denoise step per scheduler tick |
|
||||
| `step_execution=True`, `max_num_seqs>1` | `STEP_BATCH` | `StepScheduler` | Compatible requests advanced together in step waves |
|
||||
|
||||
The current implementation is conservative:
|
||||
Serial request execution is handled inside `REQUEST_BATCH`; it is not a
|
||||
separate engine mode. Similarly, continuous batching is the multi-request
|
||||
configuration of `STEP_BATCH`, not an independent execution mode.
|
||||
|
||||
- only compatible requests are batched together
|
||||
- per-request progress and completion remain independent
|
||||
The engine performs the following initialization:
|
||||
|
||||
Here, "continuous batching" means the step-wise path enabled by
|
||||
`step_execution=True`. Request-mode `DiffusionRequestBatch` is static
|
||||
request-level batching for one full pipeline `forward()` call; it does not
|
||||
admit or remove requests between denoise steps.
|
||||
1. Resolve model-specific pre- and post-processing hooks.
|
||||
2. Select `REQUEST_BATCH` or `STEP_BATCH`.
|
||||
3. Construct the configured `DiffusionExecutor`.
|
||||
4. Construct `RequestScheduler` or `StepScheduler`, unless a scheduler was
|
||||
explicitly injected.
|
||||
5. Bind `execute_batch` or `execute_step`.
|
||||
|
||||
## Enablement
|
||||
`DiffusionEngine.make_engine()` owns startup warmup after engine construction.
|
||||
This keeps construction, backend selection, and warmup failure cleanup in one
|
||||
lifecycle.
|
||||
|
||||
Use `--step-execution` as the feature gate, then increase `--max-num-seqs`
|
||||
above `1` if you want batching:
|
||||
## Unified Output Stream
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen-Image --omni \
|
||||
--port 8091 \
|
||||
--step-execution \
|
||||
--max-num-seqs 8
|
||||
All asynchronous diffusion requests use the same per-request output stream,
|
||||
regardless of whether the caller exposes streaming results to the user.
|
||||
|
||||
```python
|
||||
async for output in engine.step_streaming(request):
|
||||
...
|
||||
```
|
||||
|
||||
`--max-num-seqs 1` keeps the step-wise path without enabling batching.
|
||||
When a request is admitted, the engine creates an
|
||||
`asyncio.Queue[DiffusionOutput]` keyed by its request ID. The busy loop sends
|
||||
both intermediate chunks and final results through that queue:
|
||||
|
||||
For a reproducible replay flow using the bundled serving benchmark, see the
|
||||
Qwen-Image replay commands in
|
||||
[`benchmarks/diffusion/README.md`](gh-file:benchmarks/diffusion/README.md)
|
||||
and
|
||||
[`benchmarks/diffusion/performance_dashboard/qwen_image_serving_performance.md`](gh-file:benchmarks/diffusion/performance_dashboard/qwen_image_serving_performance.md).
|
||||
```text
|
||||
request
|
||||
-> scheduler.add_request()
|
||||
-> scheduler.schedule()
|
||||
-> execute_batch() or execute_step()
|
||||
-> scheduler.update_from_output()
|
||||
-> per-request output queue
|
||||
-> step_streaming()
|
||||
```
|
||||
|
||||
## Scheduler
|
||||
Streaming callers forward every yielded output. Non-streaming callers consume
|
||||
the same generator to completion and return only its final output. This avoids
|
||||
maintaining separate future-based and queue-based result paths.
|
||||
|
||||
The scheduler derives its batch capacity from `max_num_seqs` through
|
||||
`max_num_running_reqs`.
|
||||
`DiffusionEngine.step()` and
|
||||
`async_add_req_and_wait_for_response()` remain deprecated compatibility
|
||||
wrappers. New integrations should consume `step_streaming()` or
|
||||
`async_add_req_and_stream_response()`.
|
||||
|
||||
Batch admission is gated by
|
||||
[`SamplingParamsKey`](gh-file:vllm_omni/diffusion/sched/interface.py),
|
||||
which is built from shape-sensitive and CFG-sensitive sampling fields. This is
|
||||
the core correctness rule for batching: requests are only co-batched when they
|
||||
share the same denoise tensor contract.
|
||||
### Completion and Cancellation
|
||||
|
||||
There are three important details:
|
||||
Output delivery and scheduler lifecycle are related but separate:
|
||||
|
||||
- `num_inference_steps` is not part of the key, so requests with different
|
||||
total step counts can still share a batch
|
||||
- requests also do not need to be at the same current denoise progress; active
|
||||
requests can continue batching even when their current step indices diverge
|
||||
- admission is still FIFO, so an incompatible request at the head of the
|
||||
waiting queue blocks later compatible requests
|
||||
- A terminal scheduler result always finalizes request state.
|
||||
- A disconnected or cancelled consumer removes only its delivery queue.
|
||||
- The scheduler may still finish and clean up the request after its queue has
|
||||
been removed.
|
||||
- Engine shutdown sends an error output to all remaining streams.
|
||||
|
||||
Today that compatibility rule is still shape-sensitive. `height`, `width`,
|
||||
`num_frames`, and CFG-related fields remain part of the key, so different
|
||||
resolutions or incompatible guidance settings do **not** co-batch yet. The
|
||||
key also covers LoRA identity (`lora_int_id`, `lora_scale`), so requests
|
||||
targeting different adapters or scales run in separate batches and the
|
||||
worker can activate exactly one adapter per step.
|
||||
This separation prevents cancelled consumers from leaving active request IDs in
|
||||
the scheduler.
|
||||
|
||||
The scheduler batching unit is one logical `OmniDiffusionRequest`. In the
|
||||
step-wise path, runtime tensor batching is represented as `StepInputBatch`. For
|
||||
request-mode prompt semantics, see
|
||||
[Request-Level Batching](../../user_guide/diffusion/request_batching.md).
|
||||
## Request-Batch Execution
|
||||
|
||||
## Runner
|
||||
Request-batch execution runs the complete diffusion pipeline for each scheduler
|
||||
wave. Each `OmniDiffusionRequest` remains an independent logical request with
|
||||
its own prompt, sampling parameters, seed, request ID, output, error, and abort
|
||||
state.
|
||||
|
||||
The runner keeps persistent per-request execution state in
|
||||
[`DiffusionRequestState`](gh-file:vllm_omni/diffusion/worker/utils.py),
|
||||
while the scheduler owns a separate lightweight request state for queueing and
|
||||
lifecycle tracking.
|
||||
With `max_num_seqs=1`, `execute_batch()` handles the serial request path. With a
|
||||
larger value, a pipeline can opt into fused execution by declaring:
|
||||
|
||||
For each step, the runner builds an
|
||||
[`InputBatch`](gh-file:vllm_omni/diffusion/worker/input_batch.py) from the
|
||||
active request states:
|
||||
```python
|
||||
supports_request_batch = True
|
||||
```
|
||||
|
||||
- prompt embeddings and masks are normalized and padded
|
||||
- dynamic tensors such as `latents` and `timesteps` are gathered each step
|
||||
- buffers are reused when batch composition stays the same
|
||||
Its `forward()` method must accept `DiffusionRequestBatch` and return one
|
||||
`DiffusionOutput` per request. The runner validates the result count and maps
|
||||
outputs back to their original request IDs with `BatchRunnerOutput`.
|
||||
|
||||
The step-wise batched path is:
|
||||
Current in-tree request-batch implementations include the Flux, LTX-2, SD3, and
|
||||
Qwen-Image pipelines. Treat the pipeline capability flag in the source as
|
||||
authoritative because support continues to expand.
|
||||
|
||||
1. Run `prepare_encode()` for newly admitted requests.
|
||||
### Request-Batch Data Flow
|
||||
|
||||
- `DiffusionSchedulerOutput` carries scheduled request IDs and payloads.
|
||||
- `DiffusionRequestBatch` provides the pipeline-facing static batch.
|
||||
- `BatchRunnerOutput` carries per-request results back to the engine.
|
||||
|
||||
Request-local setup, such as seed handling and output/error mapping, remains
|
||||
per request. Cache refresh, homogeneous LoRA activation, and
|
||||
`pipeline.forward(req_batch)` can be shared by the batch.
|
||||
|
||||
### Compatibility and Admission
|
||||
|
||||
`RequestBatchSamplingParamsKey` contains fields that affect the tensor contract,
|
||||
including shape, guidance, output count, and LoRA identity. Only compatible
|
||||
requests can share a fused forward call.
|
||||
|
||||
Admission is FIFO and conservative. An incompatible request at the front of the
|
||||
waiting queue can prevent later compatible requests from joining the current
|
||||
wave.
|
||||
|
||||
`request_batch_max_wait_ms` optionally delays the first schedule operation of a
|
||||
new wave so bursty arrivals can accumulate. It applies only when:
|
||||
|
||||
- the pipeline supports request batching;
|
||||
- step execution is disabled;
|
||||
- no requests are currently running; and
|
||||
- the configured wait is greater than zero.
|
||||
|
||||
The wait ends when the queue reaches capacity, stabilizes briefly, reaches its
|
||||
deadline, or the engine stops.
|
||||
|
||||
## Step-Batch Execution
|
||||
|
||||
Step execution exposes denoising progress to the scheduler. A supporting
|
||||
pipeline implements four stateful operations:
|
||||
|
||||
| Operation | Responsibility |
|
||||
|---|---|
|
||||
| `prepare_encode(state)` | Validate input, encode prompts, initialize latents and timesteps, and create request-local scheduler state |
|
||||
| `denoise_step(input_batch, *, states=...)` | Run one denoise forward for the scheduler-provided request states |
|
||||
| `step_scheduler(state, noise_pred)` | Update latents and advance request progress |
|
||||
| `post_decode(state)` | Decode and postprocess a completed request or output boundary |
|
||||
|
||||
Persistent request state lives in `StepRequestState`. Pipeline-specific fields
|
||||
that do not belong in the shared contract should be stored in `state.extra`.
|
||||
Queueing and lifecycle metadata remain in the scheduler's request state.
|
||||
|
||||
Current native pipelines that explicitly enable step execution include
|
||||
Qwen-Image, HunyuanImage3, and Helios. Step execution alone does not imply
|
||||
continuous-batching support: Qwen-Image accepts batched step states.
|
||||
HunyuanImage3 accepts batched step states only when its resolved self-attention
|
||||
backend is `TORCH_SDPA`; otherwise it rejects groups larger than one request.
|
||||
Configure `DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA` or
|
||||
`diffusion_attention_config.default.backend=TORCH_SDPA` when
|
||||
`max_num_seqs>1`. See the
|
||||
[HunyuanImage-3.0 recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Tencent/HunyuanImage-3.0-Instruct.md)
|
||||
for its validated configuration. Helios supports only a single active step
|
||||
request and must use `max_num_seqs=1`.
|
||||
|
||||
### Continuous Batching
|
||||
|
||||
When `max_num_seqs>1`, `StepScheduler` can keep multiple compatible requests
|
||||
active. The runner gathers their state into `InputBatch`, performs one batched
|
||||
denoise forward, then applies scheduler updates per request:
|
||||
|
||||
1. Receive transferred KV payloads, then run `prepare_encode()` for newly
|
||||
admitted requests.
|
||||
2. Build or refresh `InputBatch`.
|
||||
3. Run one batched `denoise_step(input_batch)`.
|
||||
4. Slice the batched `noise_pred` back per request.
|
||||
5. Run per-request `step_scheduler()`.
|
||||
6. Run `post_decode()` only for requests that finished denoising.
|
||||
7. Scatter updated latents back into persistent request state with
|
||||
[`scatter_latents()`](gh-file:vllm_omni/diffusion/worker/input_batch.py).
|
||||
3. Run one batched `denoise_step(input_batch, states=states)`.
|
||||
4. Slice noise predictions back to each request.
|
||||
5. Run each request's `step_scheduler()`.
|
||||
6. Run `post_decode()` at a chunk boundary or request completion.
|
||||
7. Scatter updated latents back into persistent request state.
|
||||
|
||||
This keeps the shared work limited to the denoise forward pass while preserving
|
||||
request-local scheduler state and outputs.
|
||||
`StepBatchSamplingParamsKey` protects the batched tensor contract. Requests can
|
||||
have different total step counts and current step indices, but shape-sensitive,
|
||||
CFG-sensitive, and LoRA fields must be compatible. FIFO head-of-line blocking
|
||||
also applies here.
|
||||
|
||||
## Engine
|
||||
Chunk-capable pipelines may emit intermediate results. Final-only step
|
||||
pipelines emit when the request finishes even if `streaming_output=True`.
|
||||
|
||||
[`DiffusionEngine`](gh-file:vllm_omni/diffusion/diffusion_engine.py) provides
|
||||
the background loop and async add-request path needed for multiple requests to
|
||||
accumulate in the scheduler.
|
||||
## Model Author Guidelines
|
||||
|
||||
When `step_execution=True`, the engine routes work through the step-wise
|
||||
executor path. The continuous batching behavior is defined by scheduler-side
|
||||
compatibility gating and runner-side `StepInputBatch` packing.
|
||||
When converting a request-level pipeline:
|
||||
|
||||
## Current Limitations
|
||||
- Reuse the same helpers as `forward()` to avoid behavior drift.
|
||||
- Copy request-scoped scheduler state into `state.scheduler`.
|
||||
- Advance `state.step_index` only in `step_scheduler()`.
|
||||
- Keep model forward work in `denoise_step()`.
|
||||
- Keep latent mutation in `step_scheduler()`.
|
||||
- Keep final decoding equivalent to the tail of `forward()`.
|
||||
- Store masks, condition latents, and other request-local tensors in the
|
||||
request state rather than on the shared pipeline.
|
||||
|
||||
- Experimental feature; use `max_num_seqs=1` for the older conservative path.
|
||||
- Only native pipelines that already support `step_execution=True`.
|
||||
- Only homogeneous batches keyed by `SamplingParamsKey` are supported.
|
||||
- `cache_backend`, KV transfer, and other request-mode extras are not wired
|
||||
into the batched step-wise path yet.
|
||||
- Future work can relax the current same-shape restriction with richer
|
||||
heterogeneous batching policies such as bucketing or padded execution for
|
||||
different resolutions.
|
||||
Before setting `supports_step_execution = True`, validate output parity for the
|
||||
same seed and sampling parameters, scheduler isolation across concurrent
|
||||
requests, abort cleanup, progress reporting, and all supported CFG paths.
|
||||
|
||||
## Related Files
|
||||
Before setting `supports_request_batch = True`, validate the batch input
|
||||
contract, one-output-per-request result shape, request identity and error
|
||||
mapping, seeded concurrency, LoRA compatibility, and tensor IPC.
|
||||
|
||||
- Scheduler base:
|
||||
[`vllm_omni/diffusion/sched/base_scheduler.py`](gh-file:vllm_omni/diffusion/sched/base_scheduler.py)
|
||||
- Scheduler interface:
|
||||
[`vllm_omni/diffusion/sched/interface.py`](gh-file:vllm_omni/diffusion/sched/interface.py)
|
||||
- Step scheduler:
|
||||
[`vllm_omni/diffusion/sched/step_scheduler.py`](gh-file:vllm_omni/diffusion/sched/step_scheduler.py)
|
||||
- Runner:
|
||||
[`vllm_omni/diffusion/worker/diffusion_model_runner.py`](gh-file:vllm_omni/diffusion/worker/diffusion_model_runner.py)
|
||||
- Input batch:
|
||||
[`vllm_omni/diffusion/worker/input_batch.py`](gh-file:vllm_omni/diffusion/worker/input_batch.py)
|
||||
- Tests:
|
||||
[`tests/diffusion/test_diffusion_scheduler.py`](gh-file:tests/diffusion/test_diffusion_scheduler.py)
|
||||
## Limitations
|
||||
|
||||
- Request batching requires an explicit pipeline capability when
|
||||
`max_num_seqs>1`.
|
||||
- Step execution requires the four-method stateful pipeline contract.
|
||||
- Both batching paths currently require homogeneous compatibility keys.
|
||||
- FIFO scheduling can reduce batching opportunities.
|
||||
- `request_batch_max_wait_ms` trades first-request latency for burst
|
||||
coalescing.
|
||||
- All diffusion cache backends are currently unsupported in step mode.
|
||||
- KV transfer is supported for newly admitted step requests; some other
|
||||
request-mode extras remain unsupported.
|
||||
- Step continuous batching remains experimental; use `max_num_seqs=1` for the
|
||||
conservative step path.
|
||||
|
||||
## Related Implementation
|
||||
|
||||
- Engine: [`vllm_omni/diffusion/diffusion_engine.py`](gh-file:vllm_omni/diffusion/diffusion_engine.py)
|
||||
- Model contract: [`vllm_omni/diffusion/models/interface.py`](gh-file:vllm_omni/diffusion/models/interface.py)
|
||||
- Scheduler interface: [`vllm_omni/diffusion/sched/interface.py`](gh-file:vllm_omni/diffusion/sched/interface.py)
|
||||
- Request scheduler: [`vllm_omni/diffusion/sched/request_scheduler.py`](gh-file:vllm_omni/diffusion/sched/request_scheduler.py)
|
||||
- Step scheduler: [`vllm_omni/diffusion/sched/step_scheduler.py`](gh-file:vllm_omni/diffusion/sched/step_scheduler.py)
|
||||
- Runner: [`vllm_omni/diffusion/worker/diffusion_model_runner.py`](gh-file:vllm_omni/diffusion/worker/diffusion_model_runner.py)
|
||||
- Request batch: [`vllm_omni/diffusion/worker/request_batch.py`](gh-file:vllm_omni/diffusion/worker/request_batch.py)
|
||||
- Step input batch: [`vllm_omni/diffusion/worker/input_batch.py`](gh-file:vllm_omni/diffusion/worker/input_batch.py)
|
||||
@@ -1,175 +0,0 @@
|
||||
# Request-Level Batching for Diffusion
|
||||
|
||||
This document describes the request-mode batching path for diffusion pipelines.
|
||||
For end-user enablement and tuning, see
|
||||
[Request-Level Batching](../../user_guide/diffusion/request_batching.md).
|
||||
|
||||
This is separate from
|
||||
[Continuous Batching for Step-Wise Diffusion](diffusion_continuous_batching.md).
|
||||
Request-level batching runs one full pipeline `forward()` over a static batch of
|
||||
compatible requests. Step-wise continuous batching admits work between denoise
|
||||
steps when `step_execution=True`.
|
||||
|
||||
## Why It Helps
|
||||
|
||||
The request-level design avoids coupling several logical requests to one request
|
||||
object. This keeps request identity, abort/error handling, and per-request
|
||||
metadata unambiguous while still allowing one fused pipeline forward pass for
|
||||
bursty or concurrent traffic.
|
||||
|
||||
## Overview
|
||||
|
||||
With request-level batching enabled:
|
||||
|
||||
- each `OmniDiffusionRequest` contains one `prompt` and one `request_id`
|
||||
- the scheduler groups compatible waiting requests into one scheduler wave
|
||||
- `DiffusionRequestBatch` wraps the scheduled requests for pipeline `forward()`
|
||||
- batch-capable pipelines return `list[DiffusionOutput]`, one output per
|
||||
request
|
||||
- `BatchRunnerOutput` maps each result back to its original `request_id`
|
||||
|
||||
Pipelines opt in with `supports_request_batch = True` and a `forward()` method
|
||||
that accepts `DiffusionRequestBatch` and returns `list[DiffusionOutput]`.
|
||||
Pipelines that do not opt in keep the existing per-request execution path.
|
||||
|
||||
## Enablement
|
||||
|
||||
Request-level batching is the request-mode path, so `step_execution` must remain
|
||||
disabled. Increase `max_num_seqs` above `1` to let the scheduler keep multiple
|
||||
compatible requests active:
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen-Image --omni \
|
||||
--port 8091 \
|
||||
--max-num-seqs 4
|
||||
```
|
||||
|
||||
For bursty online ingress, `request_batch_max_wait_ms` can add a bounded
|
||||
admission wait before the first `schedule()` of a scheduler wave:
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen-Image --omni \
|
||||
--port 8091 \
|
||||
--max-num-seqs 4 \
|
||||
--request-batch-max-wait-ms 20
|
||||
```
|
||||
|
||||
`request_batch_max_wait_ms=0` disables this wait and is the default.
|
||||
|
||||
## Request Contract
|
||||
|
||||
`OmniDiffusionRequest` represents one logical request. It owns one prompt,
|
||||
sampling parameters, request id, and request-local metadata. Runtime batches are
|
||||
formed by the scheduler and represented separately from the request payload.
|
||||
|
||||
Runtime batching is represented by:
|
||||
|
||||
- [`DiffusionSchedulerOutput`](gh-file:vllm_omni/diffusion/sched/interface.py)
|
||||
for scheduled request ids and request payloads
|
||||
- [`DiffusionRequestBatch`](gh-file:vllm_omni/diffusion/worker/request_batch.py)
|
||||
for the pipeline-facing request batch
|
||||
- [`BatchRunnerOutput`](gh-file:vllm_omni/diffusion/worker/utils.py) for
|
||||
per-request results
|
||||
|
||||
`DiffusionRequestBatch` intentionally exposes compatibility properties such as
|
||||
`prompts`, `sampling_params`, `request_id`, and `kv_sender_info` so migrated
|
||||
pipelines can stay close to upstream code while using a batch-aware contract.
|
||||
|
||||
## Scheduler
|
||||
|
||||
The scheduler derives its capacity from `max_num_seqs` through
|
||||
`max_num_running_reqs`. It exposes waiting/running queue counters so the engine
|
||||
can decide whether admission wait is useful before scheduling a new wave.
|
||||
|
||||
Batch compatibility is controlled by
|
||||
[`SamplingParamsKey`](gh-file:vllm_omni/diffusion/sched/interface.py). The key
|
||||
contains shape-sensitive and guidance-sensitive fields, including output count
|
||||
and LoRA identity. Requests with incompatible shapes, CFG settings, output
|
||||
counts, LoRA adapters, or LoRA scales are kept in separate batches.
|
||||
|
||||
Admission is conservative:
|
||||
|
||||
- the scheduler only batches compatible requests
|
||||
- FIFO ordering is preserved
|
||||
- an incompatible request at the head of the waiting queue blocks later
|
||||
compatible requests
|
||||
|
||||
## Engine
|
||||
|
||||
[`DiffusionEngine`](gh-file:vllm_omni/diffusion/diffusion_engine.py) resolves
|
||||
request-batch capability during initialization from the configured pipeline
|
||||
class, including custom pipeline classes.
|
||||
|
||||
The capability check uses the pipeline class attribute
|
||||
`supports_request_batch = True`. Pipelines that set this attribute must implement
|
||||
a request-batch-compatible `forward()` contract and return one
|
||||
`DiffusionOutput` per request; the runner validates that return shape at runtime.
|
||||
|
||||
When the selected pipeline is batch-capable and `step_execution=False`, request
|
||||
mode routes scheduler waves through the batch executor path. Otherwise it keeps
|
||||
the per-request executor path.
|
||||
|
||||
The optional admission wait runs only when:
|
||||
|
||||
- request batching is supported
|
||||
- `step_execution=False`
|
||||
- `request_batch_max_wait_ms > 0`
|
||||
- no requests are currently running
|
||||
|
||||
The wait exits early when the waiting queue reaches capacity, when the queue is
|
||||
stable for a short window, when the deadline expires, or when the engine stops.
|
||||
|
||||
## Executor And Runner
|
||||
|
||||
The executor exposes two request-mode entries:
|
||||
|
||||
- `execute_request`: one worker call per scheduled request
|
||||
- `execute_batch`: one worker call for the whole `DiffusionSchedulerOutput`
|
||||
|
||||
On the batch path, the worker builds a `DiffusionRequestBatch` and runs the
|
||||
pipeline once. Request-local setup remains per request:
|
||||
|
||||
- KV transfer metadata
|
||||
- random generator and seed handling
|
||||
- request output/error/abort mapping
|
||||
|
||||
Shared batch setup happens once per batch when possible:
|
||||
|
||||
- cache refresh
|
||||
- LoRA activation for the homogeneous adapter key
|
||||
- pipeline `forward(req_batch)`
|
||||
|
||||
Large tensor IPC still uses the shared-memory packing path. The packer traverses
|
||||
both normal `RunnerOutput.result` wrappers and nested batch results so batched
|
||||
outputs do not fall back to pickle IPC for tensor payloads.
|
||||
|
||||
## Current Limitations
|
||||
|
||||
- Only pipelines that declare the request-batch contract use fused batch
|
||||
execution.
|
||||
- Batches are homogeneous under `SamplingParamsKey`; heterogeneous resolution or
|
||||
incompatible guidance settings do not co-batch yet.
|
||||
- FIFO scheduling can reduce batching opportunities when an incompatible
|
||||
request is at the front of the queue.
|
||||
- `request_batch_max_wait_ms` improves burst coalescing but can add latency to
|
||||
the first request in a scheduler wave. Keep it small for latency-sensitive
|
||||
serving.
|
||||
- Step-wise continuous batching is documented separately and only applies when
|
||||
`step_execution=True`.
|
||||
|
||||
## Related Files
|
||||
|
||||
- Request object and request batch:
|
||||
[`vllm_omni/diffusion/request.py`](gh-file:vllm_omni/diffusion/request.py)
|
||||
- Scheduler interface:
|
||||
[`vllm_omni/diffusion/sched/interface.py`](gh-file:vllm_omni/diffusion/sched/interface.py)
|
||||
- Scheduler base:
|
||||
[`vllm_omni/diffusion/sched/base_scheduler.py`](gh-file:vllm_omni/diffusion/sched/base_scheduler.py)
|
||||
- Engine:
|
||||
[`vllm_omni/diffusion/diffusion_engine.py`](gh-file:vllm_omni/diffusion/diffusion_engine.py)
|
||||
- Worker runner:
|
||||
[`vllm_omni/diffusion/worker/diffusion_model_runner.py`](gh-file:vllm_omni/diffusion/worker/diffusion_model_runner.py)
|
||||
- Executor interface:
|
||||
[`vllm_omni/diffusion/executor/abstract.py`](gh-file:vllm_omni/diffusion/executor/abstract.py)
|
||||
- Tests:
|
||||
[`tests/diffusion/test_diffusion_engine.py`](gh-file:tests/diffusion/test_diffusion_engine.py)
|
||||
@@ -1,131 +0,0 @@
|
||||
# Adding Step Execution Support for Diffusion Pipelines
|
||||
|
||||
This guide documents vLLM-Omni's stepwise diffusion contract for model authors
|
||||
and contributors implementing `step_execution=True` support for a diffusion
|
||||
pipeline.
|
||||
|
||||
For end-user enablement, supported models, and current limitations, see
|
||||
[Step Execution](../../user_guide/diffusion/step_execution.md).
|
||||
|
||||
This document describes the base step-execution contract only. For the
|
||||
experimental batching policy layered on top of the step-wise path, see
|
||||
[Continuous Batching for Step-Wise Diffusion](diffusion_continuous_batching.md).
|
||||
|
||||
## Current Support Scope
|
||||
|
||||
`step_execution` is **not** a generic diffusion toggle. It only works for
|
||||
pipelines that implement the segmented stateful contract in
|
||||
[`vllm_omni/diffusion/models/interface.py`](gh-file:vllm_omni/diffusion/models/interface.py).
|
||||
|
||||
This page is intentionally author-facing. Treat runtime enablement
|
||||
(`step_execution=True` in Python or `--step-execution` in serving) as an
|
||||
opt-in user knob layered on top of the implementation contract below.
|
||||
|
||||
Current in-tree support:
|
||||
|
||||
| Pipeline | Example models | Step execution |
|
||||
|----------|----------------|----------------|
|
||||
| `QwenImagePipeline` | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512` | Yes |
|
||||
| All other diffusion pipelines | `QwenImageEditPipeline`, `QwenImageEditPlusPipeline`, `QwenImageLayeredPipeline`, GLM-Image, Wan, Flux, etc. | No |
|
||||
|
||||
Current engine/runtime limitations:
|
||||
|
||||
- Continuous batching with `max_num_seqs > 1` is experimental and documented in
|
||||
[Continuous Batching for Step-Wise Diffusion](diffusion_continuous_batching.md).
|
||||
Keep `max_num_seqs=1` if you want the older conservative behavior.
|
||||
- `cache_backend` is not supported in step mode.
|
||||
- Request-mode extras such as KV transfer are not wired into step mode yet.
|
||||
- Unsupported pipelines now fail early during model loading instead of failing on the first request.
|
||||
|
||||
## Execution Contract
|
||||
|
||||
Step mode is driven by four pipeline methods plus the shared mutable request
|
||||
state object:
|
||||
|
||||
- `prepare_encode(state)`: one-time request preparation.
|
||||
- `denoise_step(state)`: compute the noise prediction for the current step.
|
||||
- `step_scheduler(state, noise_pred)`: mutate latents and advance step state.
|
||||
- `post_decode(state)`: decode the final output after denoising is complete.
|
||||
|
||||
The state lives in
|
||||
[`vllm_omni/diffusion/worker/utils.py`](gh-file:vllm_omni/diffusion/worker/utils.py)
|
||||
as `DiffusionRequestState`. Store request-scoped tensors there, or use
|
||||
`state.extra` for model-specific fields that do not justify extending the core
|
||||
dataclass.
|
||||
|
||||
The worker-side step loop lives in
|
||||
[`vllm_omni/diffusion/worker/diffusion_model_runner.py`](gh-file:vllm_omni/diffusion/worker/diffusion_model_runner.py):
|
||||
|
||||
1. `prepare_encode()` runs once for a new request.
|
||||
2. `denoise_step()` runs every scheduler tick.
|
||||
3. `step_scheduler()` mutates `state.latents` and advances `state.step_index`.
|
||||
4. `post_decode()` runs exactly once after `state.denoise_completed` becomes true.
|
||||
|
||||
## Recommended Split
|
||||
|
||||
When converting an existing request-level `forward()` pipeline, keep the split
|
||||
strict and mechanical:
|
||||
|
||||
| Request-level phase | Stepwise method | What belongs there |
|
||||
|---------------------|-----------------|--------------------|
|
||||
| Input validation, prompt encoding, latent init, timestep prep, per-request scheduler creation | `prepare_encode()` | Anything that should happen once per request |
|
||||
| Transformer forward / noise prediction | `denoise_step()` | Pure denoise computation for the current timestep |
|
||||
| `scheduler.step(...)` and `step_index += 1` | `step_scheduler()` | Only latent/state mutation for one step |
|
||||
| VAE decode / postprocess | `post_decode()` | Final decode only |
|
||||
|
||||
Keep the stepwise path reusing the same helpers as the request-level path
|
||||
whenever possible. Reimplementing the denoise loop from scratch is the easiest
|
||||
way to introduce behavioral drift.
|
||||
|
||||
## Qwen-Image Reference
|
||||
|
||||
[`pipeline_qwen_image.py`](gh-file:vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py)
|
||||
is the reference implementation and is split correctly for the current
|
||||
contract:
|
||||
|
||||
- `prepare_encode()` reuses `_prepare_generation_context()` so prompt encoding,
|
||||
latent init, timestep creation, CFG setup, and shape bookkeeping stay aligned
|
||||
with `forward()`.
|
||||
- `prepare_encode()` deep-copies `self.scheduler` **after**
|
||||
`prepare_timesteps()` so request-specific scheduler state is isolated.
|
||||
- `denoise_step()` reuses `_build_denoise_kwargs()` plus
|
||||
`predict_noise_maybe_with_cfg()`, so sequential CFG, CFG-parallel, and
|
||||
non-CFG behavior stay identical to the request-level path.
|
||||
- `step_scheduler()` only calls
|
||||
`scheduler_step_maybe_with_cfg(..., per_request_scheduler=state.scheduler)`
|
||||
and increments `state.step_index`.
|
||||
- `post_decode()` reuses `_decode_latents()`, so the final image decode matches
|
||||
the normal `forward()` path.
|
||||
|
||||
That decomposition is the target pattern for future models.
|
||||
|
||||
## Rules For New Pipelines
|
||||
|
||||
- Do not keep request-scoped scheduler state on `self.scheduler`. Copy it into
|
||||
`state.scheduler` during `prepare_encode()`.
|
||||
- Do not mutate `state.step_index` inside `denoise_step()`. Only
|
||||
`step_scheduler()` should advance the step.
|
||||
- Do not decode partial outputs in `denoise_step()` or `step_scheduler()`.
|
||||
- If the request-level pipeline has condition latents, masks, or edit-specific
|
||||
tensors, store them in `state` or `state.extra`, not in global pipeline
|
||||
attributes.
|
||||
- Preserve CFG behavior by sharing the same helper path used by `forward()`.
|
||||
- Keep `post_decode()` equivalent to the tail of `forward()`.
|
||||
|
||||
## Validation Checklist
|
||||
|
||||
Before marking a pipeline as `supports_step_execution = True`, verify:
|
||||
|
||||
- Stepwise output matches request-level output for the same seed and sampling params.
|
||||
- Per-request scheduler state is isolated across concurrent requests.
|
||||
- Abort during denoise does not leak cached state.
|
||||
- `step_index` reported by `RunnerOutput` matches the scheduler progress.
|
||||
- CFG-parallel and non-CFG paths both work if the request-level pipeline supports them.
|
||||
|
||||
## Related Files
|
||||
|
||||
- Contract: [`vllm_omni/diffusion/models/interface.py`](gh-file:vllm_omni/diffusion/models/interface.py)
|
||||
- State: [`vllm_omni/diffusion/worker/utils.py`](gh-file:vllm_omni/diffusion/worker/utils.py)
|
||||
- Runner loop: [`vllm_omni/diffusion/worker/diffusion_model_runner.py`](gh-file:vllm_omni/diffusion/worker/diffusion_model_runner.py)
|
||||
- Scheduler transport: [`vllm_omni/diffusion/sched/interface.py`](gh-file:vllm_omni/diffusion/sched/interface.py)
|
||||
- Reference pipeline: [`vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py`](gh-file:vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py)
|
||||
@@ -66,8 +66,6 @@ runtime:
|
||||
connectors:
|
||||
connector_of_shared_memory:
|
||||
name: SharedMemoryConnector
|
||||
extra:
|
||||
shm_threshold_bytes: 65536
|
||||
```
|
||||
|
||||
Wire stages to connectors:
|
||||
|
||||
@@ -0,0 +1,156 @@
|
||||
# Distributed Layerwise Offload
|
||||
|
||||
This document describes distributed layerwise offload (DLO) for diffusion
|
||||
models. DLO keeps only a small number of DiT blocks on the accelerator and
|
||||
streams the remaining blocks from host memory. The distributed backend can
|
||||
either shard those host-side weights across an existing parallel group or keep
|
||||
the standard loader's rank-local weights and avoid an additional collective.
|
||||
|
||||
For user-facing commands, see the
|
||||
[distributed layerwise offloading guide](../../user_guide/diffusion/cpu_offload.md)
|
||||
and the [Cosmos3 recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-DistOffload.md).
|
||||
|
||||
## Status
|
||||
|
||||
DLO is implemented for multi-device diffusion execution. The default
|
||||
AllGather path is the primary path for DP and SP deployments. The
|
||||
`--dlo-no-use-allgather` path is a rank-local compatibility mode: it is useful
|
||||
for standard-loader sharding, workstation bring-up, and systems where an
|
||||
additional DLO collective is undesirable, but it does not reduce host weight
|
||||
storage across ranks.
|
||||
|
||||
The compatibility matrix below describes the current implementation. The
|
||||
unit-level guards are covered, but not every parallelism combination has a
|
||||
full model-and-hardware end-to-end test.
|
||||
|
||||
## Design
|
||||
|
||||
### DLO consumes the existing parallel topology
|
||||
|
||||
DLO does not create a new DP, TP, or SP topology. It reads the configured
|
||||
`DiffusionParallelConfig` and attaches offload hooks to the DiT blocks after the
|
||||
standard distributed groups have been initialized.
|
||||
|
||||
The DLO weight-sharding group is selected as follows:
|
||||
|
||||
1. Use the existing DP group when `data_parallel_size > 1`.
|
||||
2. When DP is one and SP is greater than one, use the SP group.
|
||||
3. Otherwise, run rank-locally without a DLO process group.
|
||||
|
||||
TP is deliberately not used as DLO's AllGather group. HSDP has its own
|
||||
parameter-sharding lifecycle and is not allowed to be sharded a second time by
|
||||
DLO's AllGather path.
|
||||
|
||||
### AllGather path
|
||||
|
||||
With the default `dlo_use_allgather=True`, each rank stores approximately
|
||||
`1 / group_size` of each streamable block in pinned host memory. The next
|
||||
block's shard is copied to a device buffer and reconstructed with
|
||||
`all_gather_into_tensor` on a communication stream while the current block is
|
||||
executing.
|
||||
|
||||
```text
|
||||
Compute: [Block N] [Block N+1] [Block N+2]
|
||||
H2D: [shard N+1] [shard N+2]
|
||||
AllGather: [full N+1] [full N+2]
|
||||
Buffers: [current slot] [prefetch slot] [current slot]
|
||||
```
|
||||
|
||||
|
||||
|
||||
The backend uses two shared device buffers, so accelerator weight residency is
|
||||
bounded by the largest streamed blocks rather than the complete model.
|
||||
|
||||
When DP is greater than one, the engine can process one request per DP rank in
|
||||
the same denoising wave. Because AllGather is a collective, all participating
|
||||
requests must take the same execution path at every denoising step.
|
||||
|
||||
### Rank-local path without DLO AllGather
|
||||
|
||||
With `--dlo-no-use-allgather`, DLO forces its internal offload shard size to
|
||||
one. The regular model loader remains responsible for preparing each rank's
|
||||
weights, including TP-local tensors or HSDP-managed parameters. DLO then
|
||||
streams those rank-local tensors block by block using H2D copies only.
|
||||
|
||||
This mode means:
|
||||
|
||||
- DP still provides independent replicas, but DLO does not shard weights
|
||||
across DP ranks.
|
||||
- SP still performs its normal activation/attention collectives, but DLO does
|
||||
not shard weights across SP ranks.
|
||||
- TP/HSDP/SP collectives, if configured, are not disabled by this flag; only
|
||||
DLO's additional weight AllGather is disabled.
|
||||
- Pure DP deployments keep a full host-side model copy per rank, subject to
|
||||
shared OS page-cache behavior.
|
||||
- The scheduler does not require a synchronized DP request wave for DLO.
|
||||
|
||||
This path is intentionally not implemented by reusing the AllGather mmap
|
||||
loader. It relies on the standard loader so model-specific TP/HSDP transforms
|
||||
remain intact.
|
||||
|
||||
## Parallelism compatibility
|
||||
|
||||
| Parallelism | DLO + AllGather | DLO without AllGather |
|
||||
|---|---|---|
|
||||
| **DP** | Supported primary path. DLO shards host weights across the DP group and can run DP multi-concurrency. | Supported rank-local path. DP replicas remain independent; no DLO weight sharding or cross-DP DLO collective. |
|
||||
| **SP** | Supported in the implementation. With DP=1, DLO uses the SP group for host-weight sharding; SP still shards sequence/activation work. | SP remains active, but DLO keeps standard-loader rank-local weights and adds no SP weight collective. |
|
||||
| **TP > 1** | Unsupported in the DLO mmap path. TP-aware loader callbacks are bypassed, so the backend rejects this configuration when it enters that path. | Standard loading is retained and TP-local tensors can be streamed. This is the intended compatibility path, but it still needs broader model and hardware validation. |
|
||||
| **HSDP** | Rejected. HSDP has already sharded parameters, so DLO AllGather would double-shard them. | Accepted by configuration. HSDP owns parameter sharding and its own gathers; DLO only stages rank-local parameters. End-to-end coverage is limited. |
|
||||
|
||||
### Combined dimensions
|
||||
|
||||
- **DP + SP:** DLO uses the DP group for weight sharding when DP is greater
|
||||
than one; SP continues to use its own sequence-parallel group. If DP is one,
|
||||
the SP group becomes DLO's sharding group in AllGather mode.
|
||||
- **DP + TP/SP without AllGather:** standard model loading defines the
|
||||
rank-local tensor layout. DLO adds no cross-DP, cross-TP, or cross-SP weight
|
||||
collective.
|
||||
- **HSDP + SP:** the general parallel configuration permits HSDP over SP, but
|
||||
DLO must use `--dlo-no-use-allgather`. HSDP remains responsible for weight
|
||||
materialization and synchronization.
|
||||
- **HSDP + DP or TP:** rejected independently by the diffusion parallel
|
||||
configuration.
|
||||
|
||||
## Request and loading constraints
|
||||
|
||||
AllGather DP multi-concurrency requires:
|
||||
|
||||
- explicit `num_inference_steps`;
|
||||
- the same `num_inference_steps` for all requests in a wave; and
|
||||
- identical request arguments that affect the collective execution path.
|
||||
|
||||
The no-AllGather path does not impose these DLO-specific synchronized-wave
|
||||
requirements.
|
||||
|
||||
The mmap loader is used only by the supported DLO+AllGather path when the
|
||||
model has a compatible checkpoint layout. Online quantization is incompatible
|
||||
with that sharded mmap path; use `--dlo-no-use-allgather` or disable online
|
||||
quantization. Models that do not meet the mmap requirements use the regular
|
||||
loader path.
|
||||
|
||||
## Validation coverage
|
||||
|
||||
Current source-level validation includes:
|
||||
|
||||
- HSDP + DLO + AllGather rejection;
|
||||
- HSDP + DLO without AllGather acceptance at configuration level;
|
||||
- TP rejection in the DLO+AllGather mmap path;
|
||||
- resident-layer requests requiring no-AllGather;
|
||||
- DP request-wave validation for denoising-step compatibility;
|
||||
- sharding, double-buffer, AllGather-size, and heterogeneous-block regression
|
||||
tests.
|
||||
|
||||
The highest-value missing coverage is end-to-end numerical comparison against
|
||||
ordinary layerwise offload for DP+SP, TP+no-AllGather, and HSDP+SP+no-AllGather
|
||||
on the target CUDA/NCCL or CANN/HCCL hardware.
|
||||
|
||||
## Recommendations
|
||||
|
||||
- Use **DP + DLO AllGather** for the supported throughput and host-memory
|
||||
scaling path.
|
||||
- Use **SP + DLO AllGather** for long-sequence workloads when DP concurrency is
|
||||
not the goal.
|
||||
- Use **no-AllGather** to bring up TP or workstation PCIe configurations, with
|
||||
the expectation of higher host-memory use and lower validation confidence.
|
||||
- Prefer **HSDP alone** for production HSDP deployments until the combined
|
||||
HSDP + DLO no-AllGather path has broader end-to-end coverage.
|
||||
@@ -0,0 +1,386 @@
|
||||
# Async Omni Output Materialization
|
||||
|
||||
## Table of Contents
|
||||
|
||||
1. [Overview](#overview)
|
||||
2. [Performance](#performance)
|
||||
3. [Architecture](#architecture)
|
||||
4. [Enabling the Feature](#enabling-the-feature)
|
||||
5. [Compatibility and Fallbacks](#compatibility-and-fallbacks)
|
||||
6. [Related Files](#related-files)
|
||||
|
||||
## Overview
|
||||
|
||||
Async Omni output materialization moves CPU-side construction of
|
||||
`OmniModelRunnerOutput` out of the autoregressive (AR) decode critical path.
|
||||
It complements [Async Chunk](async_chunk.md): async chunk pipelines partial
|
||||
outputs between stages, while async output materialization prevents the
|
||||
construction of those partial outputs from blocking the next decode step.
|
||||
|
||||
An AR stage must return sampled token IDs quickly because the scheduler needs
|
||||
them for the next decode step. However, an Omni stage can also produce hidden
|
||||
states, multimodal tensors, streaming inter-stage payloads, and connector
|
||||
metadata. Building these payloads can require device-to-host (D2H) copies,
|
||||
tensor slicing, flattening, and Python object construction.
|
||||
|
||||
Without this feature, all of that work runs inline in `sample_tokens()`:
|
||||
|
||||
```text
|
||||
sample tokens
|
||||
-> update decode state
|
||||
-> copy hidden and multimodal outputs to CPU
|
||||
-> build per-request payloads
|
||||
-> build the streaming wire payload
|
||||
-> collect connector signals
|
||||
-> return sampled tokens
|
||||
-> launch the next decode step
|
||||
```
|
||||
|
||||
With async output materialization, only work required by the next decode step
|
||||
stays on the critical path:
|
||||
|
||||
```text
|
||||
sample tokens
|
||||
-> update decode state
|
||||
-> snapshot output state and start asynchronous D2H copies
|
||||
-> register sampled tokens
|
||||
-> launch the next decode step
|
||||
|
||||
background output path
|
||||
-> wait for the payload snapshot
|
||||
-> build per-request payloads
|
||||
-> build the streaming wire payload
|
||||
-> collect connector signals
|
||||
-> construct OmniModelRunnerOutput
|
||||
```
|
||||
|
||||
This allows payload construction for step `N` to overlap with GPU decode work
|
||||
for step `N + 1`. It does not change model computation or generated output.
|
||||
|
||||
!!! note "Platform and validation scope"
|
||||
The performance and correctness validation in this document covers CUDA
|
||||
and ROCm. This is a validation scope, not a runtime platform guard: XPU
|
||||
inherits `GPUARModelRunner`, and MUSA selects `GPUARWorker`, so those
|
||||
backends may enter the asynchronous path when its other guards pass. XPU
|
||||
and MUSA have not been validated for this optimization. Ascend NPU uses a
|
||||
separate `NPUARModelRunner` path that still constructs
|
||||
`OmniModelRunnerOutput` and drains connector output synchronously before
|
||||
wrapping the result as an asynchronous runner output.
|
||||
|
||||
<p align="center">
|
||||
<img
|
||||
alt="Decode step gap before and after async output materialization"
|
||||
src="../../../source/architecture/qwen3-omni-async-output-step-gap.svg"
|
||||
width="100%"
|
||||
>
|
||||
</p>
|
||||
|
||||
<p align="center">
|
||||
<em>
|
||||
Figure 1: Moving payload construction off the decode path reduces the
|
||||
observed gap between consecutive Talker steps from about 2.8 ms to 41 µs.
|
||||
Adapted from the
|
||||
<a href="https://vllm.ai/blog/2026-07-01-qwen3-omni-optimization">
|
||||
Qwen3-Omni optimization blog
|
||||
</a>.
|
||||
</em>
|
||||
</p>
|
||||
|
||||
For Qwen3-Omni, the feature is enabled for both AR stages:
|
||||
|
||||
- **Thinker**: Defers hidden-state and multimodal payload materialization.
|
||||
- **Talker**: Runs its decode-state postprocess eagerly, then defers codec
|
||||
payload materialization. Because Code2Wav consumes codec codes rather than
|
||||
Talker hidden states, the Talker also skips unnecessary hidden-state D2H.
|
||||
- **Code2Wav**: Does not use this path because it is a generation stage rather
|
||||
than an AR stage.
|
||||
|
||||
Qwen3-TTS uses the same mechanism in its AR Talker stage. The Talker updates
|
||||
the decode state it needs for the next step before returning, then constructs
|
||||
the codec payload for Code2Wav in the background.
|
||||
|
||||
## Performance
|
||||
|
||||
In the controlled Qwen3-Omni optimization sweep, enabling async output
|
||||
materialization on top of CUDA Graph and async chunk produced the following
|
||||
results at concurrency 64:
|
||||
|
||||
| Configuration | Request throughput | Mean audio TTFP | Mean audio RTF |
|
||||
|---|---:|---:|---:|
|
||||
| CUDA Graph + async chunk | 9.3 req/s | 655 ms | 0.63 |
|
||||
| + async output materialization | 11.3 req/s | 631 ms | 0.47 |
|
||||
| Change | **+22%** | **-4%** | **-25%** |
|
||||
|
||||
The sweep used `Qwen/Qwen3-Omni-30B-A3B-Instruct`, 640 Seed-TTS English
|
||||
prompts, and one replica each for Thinker, Talker, and Code2Wav. Async output
|
||||
materialization recovers throughput by removing CPU payload construction from
|
||||
the decode critical path without regressing time to first audio packet (TTFP).
|
||||
See the
|
||||
[Qwen3-Omni optimization blog](https://vllm.ai/blog/2026-07-01-qwen3-omni-optimization)
|
||||
for the complete benchmark methodology and results.
|
||||
|
||||
The implementation was also validated against the synchronous safe-copy path.
|
||||
The text and audio output hashes matched, confirming that overlap does not
|
||||
change the generated result. See
|
||||
[PR #4476](https://github.com/vllm-project/vllm-omni/pull/4476) for the
|
||||
correctness and profiling details.
|
||||
|
||||
## Architecture
|
||||
|
||||
### Output Lifecycle
|
||||
|
||||
`OmniAsyncGPUModelRunnerOutput` extends vLLM's
|
||||
`AsyncGPUModelRunnerOutput`. It preserves the existing asynchronous
|
||||
sampled-token feedback while adding a background builder for the complete Omni
|
||||
output.
|
||||
|
||||
The output lifecycle is:
|
||||
|
||||
1. `GPUARModelRunner.sample_tokens()` samples tokens and performs the
|
||||
bookkeeping required by the next decode step.
|
||||
2. The runner snapshots step-local metadata such as request IDs, token spans,
|
||||
scheduler output, and query start locations.
|
||||
3. CUDA payload tensors are cloned before reusable model or CUDA Graph buffers
|
||||
can be overwritten.
|
||||
4. The cloned payload is copied to pinned CPU memory on a dedicated CUDA
|
||||
stream. A CUDA event records when the copy is ready.
|
||||
5. `OmniAsyncGPUModelRunnerOutput` starts a background thread that waits for
|
||||
the payload event and builds `OmniModelRunnerOutput`.
|
||||
6. The AR runner immediately registers the sampled-token CPU copy with the
|
||||
input batch and returns, allowing async scheduling to advance.
|
||||
7. When the engine calls `get_output()`, it joins the background builder,
|
||||
propagates any builder exception, and finalizes sampled tokens and logprobs
|
||||
through the upstream output implementation.
|
||||
|
||||
```text
|
||||
GPU / decode path Background output path
|
||||
|
||||
forward + sample
|
||||
|
|
||||
+-- clone output tensors
|
||||
+-- enqueue D2H copy --------------> wait for D2H event
|
||||
+-- register sampled tokens |
|
||||
+-- return async output +-- slice per request
|
||||
| +-- build multimodal payloads
|
||||
next decode step +-- partition streaming payloads
|
||||
+-- build the wire payload
|
||||
+-- drain connector signals
|
||||
+-- build OmniModelRunnerOutput
|
||||
```
|
||||
|
||||
### Safe State Snapshots
|
||||
|
||||
The next scheduler step can mutate runner state while the background builder is
|
||||
still active. Most step-local builder inputs are therefore detached snapshots.
|
||||
The snapshotted state includes:
|
||||
|
||||
- scheduler token counts and speculative-token metadata
|
||||
- request IDs and request-to-batch-index mappings
|
||||
- sampled token IDs, logprobs, and prompt logprobs
|
||||
- query start locations and scheduled-token spans
|
||||
- hidden states and multimodal output tensors
|
||||
- KV-connector and encoder-cache outputs captured for the step
|
||||
|
||||
CUDA tensor cloning is required because CUDA Graph and model output buffers can
|
||||
be reused by the next forward pass. The dedicated copy stream and pinned host
|
||||
buffers make the D2H transfer asynchronous; the snapshot retains its cloned
|
||||
CUDA sources until the transfer event completes.
|
||||
|
||||
### Live Runner State and Ownership
|
||||
|
||||
The complete builder input is not a detached snapshot. Connector output state
|
||||
remains runner-owned: `get_omni_connector_output()` drains live, per-cycle
|
||||
connector signals after the streaming payload is built. Correctness therefore
|
||||
depends on that builder being the only consumer that drains connector signals
|
||||
for its output cycle. Receive-side connector state written by the background
|
||||
receiver is coordinated by the connector mixin's `_lock`.
|
||||
`OmniAsyncGPUModelRunnerOutput.get_output()` joins the builder and is the
|
||||
completion and exception boundary before the materialized output is consumed.
|
||||
|
||||
The shared output builder also contains a full-payload accumulation branch
|
||||
that reads `self.requests`. That branch is not reachable when async Omni output
|
||||
materialization is enabled: this feature requires `async_chunk`, while
|
||||
`should_accumulate_full_payload_output()` returns `False` whenever
|
||||
`async_chunk` is enabled. The feature path partitions each step's payload into
|
||||
streaming inter-stage and client handoffs instead.
|
||||
|
||||
### Output Builder
|
||||
|
||||
The background builder performs the work that previously ran inline:
|
||||
|
||||
- resolves which requests need downstream payloads
|
||||
- converts hidden and multimodal tensors into per-request CPU payloads
|
||||
- applies model-specific payload processing
|
||||
- partitions per-step payloads into inter-stage and client streams
|
||||
- creates the tensor-only `multimodal_outputs` wire payload
|
||||
- drains live connector readiness signals after building the streaming payload
|
||||
- constructs the final `OmniModelRunnerOutput`
|
||||
|
||||
Input-side connector operations remain synchronous with model execution. In
|
||||
particular, receiving inputs and flushing previously completed connector
|
||||
outputs are not deferred because they affect scheduler-visible request state.
|
||||
|
||||
### Qwen3-Omni Stage Behavior
|
||||
|
||||
| Stage | Async output behavior | Reason |
|
||||
|---|---|---|
|
||||
| Thinker | Snapshots hidden states and multimodal outputs; builds the downstream payload in the background | The Talker needs the payload, but the next Thinker decode step only needs sampled-token feedback |
|
||||
| Talker | Runs lightweight postprocess eagerly; snapshots codec outputs; omits hidden states from the downstream payload | `hidden_states.last` is needed by the next Talker step, while Code2Wav only needs codec codes |
|
||||
| Code2Wav | Uses the normal generation-stage output path | Code2Wav is not executed by `GPUARModelRunner` |
|
||||
|
||||
The Qwen3-TTS Talker follows the same pattern as the Qwen3-Omni Talker:
|
||||
postprocess runs eagerly because the next decode step needs its updated state,
|
||||
while codec payload construction runs asynchronously. Its Code2Wav stage also
|
||||
uses the normal generation-stage path.
|
||||
|
||||
## Enabling the Feature
|
||||
|
||||
There is no separate `async_omni_output` command-line or YAML option. The
|
||||
feature is selected automatically for model stages that opt in when the runtime
|
||||
conditions are safe. On the validated CUDA and ROCm configurations, using the
|
||||
bundled deployment profile is enough to turn it on for the supported
|
||||
Qwen3-Omni and Qwen3-TTS pipelines. The selection logic has no CUDA/ROCm-only
|
||||
guard; see [Compatibility and Fallbacks](#compatibility-and-fallbacks) for the
|
||||
other platform paths. The same profile does not enable background Omni
|
||||
materialization on Ascend NPU.
|
||||
|
||||
### Example 1: Qwen3-Omni
|
||||
|
||||
Start the full Qwen3-Omni pipeline:
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct \
|
||||
--omni \
|
||||
--port 8091
|
||||
```
|
||||
|
||||
The model registry automatically loads
|
||||
`vllm_omni/deploy/qwen3_omni_moe.yaml`. Its relevant settings are:
|
||||
|
||||
```yaml
|
||||
async_chunk: true
|
||||
|
||||
stages:
|
||||
- stage_id: 0 # Thinker
|
||||
enable_prefix_caching: false
|
||||
# async_scheduling defaults to true for AR stages
|
||||
|
||||
- stage_id: 1 # Talker
|
||||
enable_prefix_caching: false
|
||||
# async_scheduling defaults to true for AR stages
|
||||
|
||||
- stage_id: 2 # Code2Wav
|
||||
enable_prefix_caching: false
|
||||
async_scheduling: false
|
||||
```
|
||||
|
||||
With this profile, async output materialization activates automatically for:
|
||||
|
||||
- **Stage 0, Thinker**: Defers hidden-state and multimodal payload
|
||||
construction.
|
||||
- **Stage 1, Talker**: Defers codec payload construction after eagerly
|
||||
updating its decode state.
|
||||
|
||||
Stage 2 does not use this feature because Code2Wav is not an AR stage.
|
||||
|
||||
### Example 2: Qwen3-TTS
|
||||
|
||||
Start Qwen3-TTS with any checkpoint that uses the Qwen3-TTS Talker pipeline.
|
||||
For example:
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
|
||||
--omni \
|
||||
--port 8091
|
||||
```
|
||||
|
||||
The model registry automatically loads `vllm_omni/deploy/qwen3_tts.yaml`. Its
|
||||
relevant settings are:
|
||||
|
||||
```yaml
|
||||
async_chunk: true
|
||||
|
||||
stages:
|
||||
- stage_id: 0 # Talker
|
||||
async_scheduling: true
|
||||
enable_prefix_caching: false
|
||||
|
||||
- stage_id: 1 # Code2Wav
|
||||
enable_prefix_caching: false
|
||||
```
|
||||
|
||||
Async output materialization activates only for Stage 0, the AR Talker. The
|
||||
Talker keeps the state required by its next decode step on the GPU, skips an
|
||||
unnecessary hidden-state D2H copy, and builds the codec payload for Code2Wav in
|
||||
the background. Stage 1 uses the generation-stage output path.
|
||||
|
||||
The same behavior applies to the Base and VoiceDesign Qwen3-TTS checkpoints
|
||||
because they use the same Talker implementation.
|
||||
|
||||
!!! warning
|
||||
Do not pass `--no-async-chunk` or enable prefix caching when you want this
|
||||
optimization. Either change causes the runner to fall back to synchronous
|
||||
output construction. Enabling async chunk alone cannot activate the feature
|
||||
for a model that has not opted into async Omni output.
|
||||
|
||||
!!! note "Prefix cache compatibility"
|
||||
Async Omni output materialization and Omni prefix caching cannot currently
|
||||
run together. `_should_use_async_omni_output()` returns `False` whenever
|
||||
`self.omni_prefix_cache` is present, so prefix caching selects synchronous
|
||||
output materialization. Supporting both features requires snapshotting or
|
||||
otherwise synchronizing prefix-cache merge and update state before the
|
||||
background output path can consume it safely.
|
||||
|
||||
## Compatibility and Fallbacks
|
||||
|
||||
`GPUARModelRunner` uses the async output path only when all of the following
|
||||
runtime conditions hold:
|
||||
|
||||
| Requirement | Reason |
|
||||
|---|---|
|
||||
| AR async scheduling is enabled | The optimization relies on the scheduler advancing while the prior output is materialized |
|
||||
| `async_chunk` is enabled | The feature targets incremental downstream Omni payloads |
|
||||
| The model stage opts in with `use_async_omni_output` | Models must declare that their output lifecycle is safe to defer |
|
||||
| Omni prefix cache is disabled | Prefix-cache merge and update ordering currently requires synchronous materialization |
|
||||
| Speculative decoding is disabled | Speculative output state is not included in this deferred path |
|
||||
| Routed-expert output is disabled | Routed-expert extraction currently requires the synchronous path |
|
||||
| Postprocess is absent or explicitly runs eagerly | State needed by the next decode step must be updated before the runner returns |
|
||||
|
||||
These checks are evaluated per stage on runners based on `GPUARModelRunner`.
|
||||
An unsupported combination does not prevent serving; it only falls back to
|
||||
synchronous output construction for that stage. CUDA and ROCm are the validated
|
||||
platforms. XPU inherits `GPUARModelRunner`, and MUSA selects `GPUARWorker`, so
|
||||
they may select the async path when these checks pass, but this optimization is
|
||||
not yet validated on either backend.
|
||||
|
||||
Ascend NPU does not use this selection path. `NPUARModelRunner.sample_tokens()`
|
||||
fully constructs `OmniModelRunnerOutput`, calls
|
||||
`get_omni_connector_output()`, and only then creates the asynchronous wrapper,
|
||||
so Omni payload materialization remains synchronous.
|
||||
|
||||
!!! note
|
||||
`use_async_omni_output`,
|
||||
`eager_omni_postprocess_before_async_output`, and
|
||||
`omni_pooler_payload_include_hidden` are model implementation contracts,
|
||||
not user-facing configuration fields.
|
||||
|
||||
## Related Files
|
||||
|
||||
- `vllm_omni/worker/gpu_ar_model_runner.py`: Async output object, safe tensor
|
||||
snapshots, compatibility guards, and deferred Omni output builder.
|
||||
- `vllm_omni/platforms/npu/worker/npu_ar_model_runner.py`: Separate Ascend NPU
|
||||
runner, where Omni output materialization remains synchronous.
|
||||
- `vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py`: Thinker and
|
||||
Talker opt-in behavior.
|
||||
- `vllm_omni/deploy/qwen3_omni_moe.yaml`: Default Qwen3-Omni deployment
|
||||
settings.
|
||||
- `vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py`: Qwen3-TTS
|
||||
Talker opt-in and eager postprocess behavior.
|
||||
- `vllm_omni/deploy/qwen3_tts.yaml`: Default Qwen3-TTS deployment settings.
|
||||
- `tests/worker/test_gpu_ar_model_runner.py`: Snapshot, guard, connector
|
||||
ordering, and background error-propagation tests.
|
||||
- [Async Chunk](async_chunk.md): Inter-stage chunking and scheduling design.
|
||||
- [Qwen3-Omni optimization blog](https://vllm.ai/blog/2026-07-01-qwen3-omni-optimization):
|
||||
Optimization context and controlled performance results.
|
||||
- [PR #4476](https://github.com/vllm-project/vllm-omni/pull/4476): Feature
|
||||
implementation, validation, and review history.
|
||||
@@ -8,9 +8,7 @@ auto-configured when no explicit connector is specified for an edge.
|
||||
## How It Works
|
||||
|
||||
All payloads are serialized and stored in shared memory (`/dev/shm`); the SHM
|
||||
segment name is returned in metadata. The configuration exposes a
|
||||
`shm_threshold_bytes` field for a future inline-vs-SHM split, but the current
|
||||
implementation always uses shared memory regardless of payload size.
|
||||
segment name is returned in metadata.
|
||||
|
||||
## Configuration
|
||||
|
||||
@@ -19,8 +17,6 @@ runtime:
|
||||
connectors:
|
||||
connector_of_shared_memory:
|
||||
name: SharedMemoryConnector
|
||||
extra:
|
||||
shm_threshold_bytes: 65536
|
||||
```
|
||||
|
||||
## Notes
|
||||
@@ -166,27 +162,10 @@ This path is mainly for older code paths and is not the preferred mode for the c
|
||||
|
||||
### 6. Key Implementation Characteristics
|
||||
|
||||
#### 6.1 Threshold Exists, but the Current Code Always Uses SHM
|
||||
#### 6.1 All Payloads Use Shared Memory
|
||||
|
||||
The class keeps a `shm_threshold_bytes` field and still exposes metrics for inline writes. However, the current implementation uses:
|
||||
|
||||
```python
|
||||
if True:
|
||||
...
|
||||
```
|
||||
|
||||
inside `put()`, which means the current code path always writes to shared memory.
|
||||
|
||||
So the design still suggests a future split between:
|
||||
|
||||
- small payloads inline
|
||||
- large payloads in shared memory
|
||||
|
||||
but the current behavior is effectively:
|
||||
|
||||
- all payloads go through shared memory
|
||||
|
||||
This should be documented because it affects real runtime behavior.
|
||||
`put()` writes every serialized payload to shared memory. The connector has no
|
||||
inline-payload path or size threshold.
|
||||
|
||||
#### 6.2 Cleanup Is Currently Passive
|
||||
|
||||
|
||||
@@ -0,0 +1,249 @@
|
||||
# Quantization
|
||||
|
||||
This document describes the quantization architecture in vLLM-Omni and the
|
||||
extension points for contributors. For user-facing configuration, supported
|
||||
hardware, and method-specific instructions, see the
|
||||
[quantization user guide](../../user_guide/quantization/overview.md).
|
||||
|
||||
## Goals and scope
|
||||
|
||||
The design has four goals:
|
||||
|
||||
1. Provide one `quantization_config` entry point for diffusion-only,
|
||||
multi-stage omni/TTS, and multi-stage diffusion models.
|
||||
2. Reuse vLLM's quantization registry and `QuantizationConfig` contract where
|
||||
a vLLM backend already provides the required behavior.
|
||||
3. Allow different components of one pipeline to use different quantization
|
||||
configurations.
|
||||
4. Preserve the distinction between online quantization of a BF16/FP16
|
||||
checkpoint and loading a pre-quantized checkpoint with serialized scales.
|
||||
|
||||
This document covers configuration resolution, component routing, model
|
||||
integration, backend extension, and validation. Kernel implementations and
|
||||
user-facing support matrices are outside its scope.
|
||||
|
||||
## Architecture
|
||||
|
||||
The quantization path is split across the following layers:
|
||||
|
||||
| Layer | Responsibility | Source |
|
||||
| --- | --- | --- |
|
||||
| Public API | Exposes `build_quant_config`, component routing, and backend registration | `vllm_omni/quantization/__init__.py` |
|
||||
| Factory | Normalizes method names, selects Omni overrides or vLLM registry backends, detects ModelOpt configs, and reconciles checkpoint metadata | `vllm_omni/quantization/factory.py` |
|
||||
| Component router | Resolves a quantization config from a layer prefix using longest-prefix matching | `vllm_omni/quantization/component_config.py` |
|
||||
| Method configs | Implement `QuantizationConfig` and select the platform-specific linear method | `vllm_omni/quantization/*_config.py` |
|
||||
| Diffusion configuration | Stores the canonical runtime config and propagates checkpoint configuration | `vllm_omni/diffusion/data.py` |
|
||||
| Worker and model integration | Passes the resolved config into vLLM's model configuration and model layers | `vllm_omni/diffusion/worker/diffusion_worker.py` and `vllm_omni/diffusion/models/` |
|
||||
|
||||
The runtime contract remains vLLM's `QuantizationConfig`:
|
||||
|
||||
```text
|
||||
user/CLI or model config
|
||||
|
|
||||
v
|
||||
OmniDiffusionConfig.quantization_config
|
||||
|
|
||||
+--> build_quant_config()
|
||||
| |
|
||||
| +--> Omni override
|
||||
| +--> vLLM registry backend
|
||||
| +--> ComponentQuantizationConfig
|
||||
|
|
||||
+--> checkpoint metadata reconciliation
|
||||
|
|
||||
+--> diffusion worker / vLLM config
|
||||
|
|
||||
+--> model layers call get_quant_method(layer, prefix)
|
||||
```
|
||||
|
||||
The model layer receives either a method implementation or `None`. Returning
|
||||
`None` leaves that layer unquantized.
|
||||
|
||||
## Configuration resolution
|
||||
|
||||
`build_quant_config()` accepts the following forms:
|
||||
|
||||
| Input | Resolution |
|
||||
| --- | --- |
|
||||
| `None` or `"none"` | Disable quantization. |
|
||||
| Method string, such as `"fp8"` | Resolve an Omni override first, then the vLLM quantization registry. |
|
||||
| Flat dictionary with `method` or `quant_method` | Normalize the method and construct its `QuantizationConfig`. |
|
||||
| Per-component dictionary | Construct a `ComponentQuantizationConfig`. |
|
||||
| Existing `QuantizationConfig` | Pass through without rebuilding it. |
|
||||
|
||||
Method aliases are normalized case-insensitively, with `-` and `_` treated as
|
||||
equivalent. `auto-round` and `auto_round` use the Omni INC/AutoRound adapter.
|
||||
ModelOpt checkpoint metadata is detected from its method, producer metadata, or
|
||||
`quant_algo`, then mapped to the corresponding vLLM ModelOpt configuration.
|
||||
|
||||
The factory uses Omni-specific overrides for methods whose diffusion behavior
|
||||
or platform dispatch is not provided by the vLLM registry. The current override
|
||||
set includes INT8, BitsAndBytes, MXFP8, MXFP4, MXFP4 dual-scale, and INC/
|
||||
AutoRound. Other methods, such as FP8, GGUF, and ModelOpt, are resolved through
|
||||
vLLM's registry when their configuration is compatible.
|
||||
|
||||
## Per-component routing
|
||||
|
||||
Multi-stage models do not necessarily want one quantization policy for every
|
||||
module. A component map expresses that policy without requiring each pipeline
|
||||
to construct separate vLLM configurations:
|
||||
|
||||
```python
|
||||
quantization_config = {
|
||||
"transformer": {"method": "fp8"},
|
||||
"vae": None,
|
||||
"default": None,
|
||||
}
|
||||
```
|
||||
|
||||
`ComponentQuantizationConfig` applies the following rules:
|
||||
|
||||
1. Component keys are interpreted as runtime layer-prefixes.
|
||||
2. The longest matching prefix wins.
|
||||
3. `None` disables quantization for the matching component.
|
||||
4. `default`, when present, handles prefixes with no explicit match.
|
||||
5. A matching child config delegates to its normal
|
||||
`get_quant_method(layer, prefix)` implementation.
|
||||
|
||||
Prefix matching happens after the model's weight-name mapping. New model
|
||||
integrations must verify the actual runtime prefixes, especially when a
|
||||
`WeightsMapper` changes checkpoint names. A prefix mismatch can silently route
|
||||
a layer to the default configuration, so component-routing tests should cover
|
||||
both positive matches and fall-through behavior.
|
||||
|
||||
## Online and pre-quantized checkpoints
|
||||
|
||||
The same method name can represent two different loading paths:
|
||||
|
||||
- **Online**: the source checkpoint contains BF16/FP16 weights; the method
|
||||
creates quantized weights or scales while the model is loaded.
|
||||
- **Pre-quantized**: the checkpoint contains quantized weights and the metadata
|
||||
required by the method; the loader consumes those serialized tensors.
|
||||
|
||||
Method configuration owns the mode flag, for example
|
||||
`is_checkpoint_int8_serialized` or `is_checkpoint_mxfp8_serialized`. The
|
||||
configuration's `get_quant_method()` selects the corresponding linear method.
|
||||
|
||||
For pipelines with transformer-specific `config.json` files, such as cascade
|
||||
models, `resolve_quant_config_from_disk()` reconciles the active configuration
|
||||
with each transformer's checkpoint metadata. Its invariants are:
|
||||
|
||||
- If no active configuration exists, valid checkpoint metadata can be used for
|
||||
auto-detection.
|
||||
- A method mismatch is an error rather than an implicit conversion.
|
||||
- Serialized-checkpoint flags rebuild an online configuration into the
|
||||
matching offline configuration.
|
||||
- Checkpoint-specific `ignored_layers` are preserved by rebuilding the active
|
||||
configuration when they differ.
|
||||
- AutoRound MXFP8 metadata (`data_type="mx_fp"`) selects its offline path.
|
||||
|
||||
This prevents a model from interpreting serialized weights with an online
|
||||
loader or silently applying the wrong method to one transformer in a
|
||||
multi-transformer pipeline.
|
||||
|
||||
## Model and layer integration
|
||||
|
||||
Diffusion model constructors receive the resolved configuration and pass it to
|
||||
quantizable vLLM layers. Each method configuration is responsible for:
|
||||
|
||||
- identifying supported layers in `get_quant_method()`;
|
||||
- creating the correct parameter shapes and loaders;
|
||||
- processing weights after loading when conversion is required; and
|
||||
- applying the platform-specific quantized operation at runtime.
|
||||
|
||||
Model code must keep precision-sensitive or unsupported modules unquantized.
|
||||
The shared `safe_quant_config()` helper is used by model integrations for
|
||||
normalization and modulation layers where a generic FP8 configuration would be
|
||||
unsafe. Text encoders, audio/vision encoders, VAEs, schedulers, and tokenizers
|
||||
are not implicitly quantized just because the main transformer is quantized;
|
||||
they require explicit component routing and method support.
|
||||
|
||||
Weight quantization and quantized KV cache are separate paths. The former is
|
||||
resolved through `quantization_config`; runtime attention/KV-cache quantization
|
||||
is configured by the diffusion attention configuration and must be validated
|
||||
independently.
|
||||
|
||||
## Platform and parallelism boundaries
|
||||
|
||||
Platform checks belong in the method implementation rather than in the common
|
||||
factory. This keeps configuration parsing portable while allowing CUDA, NPU,
|
||||
and XPU methods to select different kernels or reject unsupported execution at
|
||||
model construction time. The user guide is the source of truth for the current
|
||||
hardware support matrix.
|
||||
|
||||
Quantization methods must also respect the weight layout produced by the
|
||||
configured parallelism strategy:
|
||||
|
||||
- Tensor-parallel layers may receive partitioned weight shapes and loaders.
|
||||
- Sequence-parallel execution changes activation movement but does not by
|
||||
itself change the quantization configuration.
|
||||
- Data-parallel replicas build their own method state; quantization metadata
|
||||
must be available consistently on every rank.
|
||||
- `WeightsMapper` transformations must be applied before matching ignored-layer
|
||||
lists or component prefixes.
|
||||
|
||||
Distributed layerwise offload has an additional boundary: its sharded mmap
|
||||
AllGather path is not compatible with online quantization because the path
|
||||
expects already materialized checkpoint weights. Use the rank-local
|
||||
`--dlo-no-use-allgather` path or a compatible pre-quantized checkpoint when
|
||||
combining these features. See the [DLO design document](distributed_layerwise_offload.md)
|
||||
for the full compatibility matrix.
|
||||
|
||||
## Adding a quantization backend
|
||||
|
||||
Contributors should follow this sequence:
|
||||
|
||||
1. Check whether an existing vLLM backend can provide the required weight
|
||||
format and kernel behavior. Prefer reuse over a duplicate Omni backend.
|
||||
2. If Omni-specific behavior is required, implement a `QuantizationConfig` and
|
||||
its `QuantizeMethodBase`/linear method in `vllm_omni/quantization/`.
|
||||
3. Implement `get_name()`, supported activation dtypes, minimum capability,
|
||||
checkpoint parsing, ignored-layer handling, and `get_quant_method()`.
|
||||
4. Add a lazy factory override in `factory.py`. Keep optional platform
|
||||
dependencies out of package import paths and normalize aliases to one
|
||||
canonical method name.
|
||||
5. If the method is pre-quantized, define the checkpoint metadata and the
|
||||
serialized loading path before adding an online path.
|
||||
6. Add model-specific component routing only when the model contains stages
|
||||
that must use different policies.
|
||||
7. Document user-visible configuration in the quantization user guide and keep
|
||||
implementation rationale here.
|
||||
|
||||
The public factory registration hook,
|
||||
`register_quantization_override(method, builder)`, is intended for integrations
|
||||
that need to register an Omni-compatible builder without changing the vLLM
|
||||
registry.
|
||||
|
||||
## Validation
|
||||
|
||||
Validation should cover the complete path, not only whether a config object can
|
||||
be instantiated:
|
||||
|
||||
| Area | Required checks |
|
||||
| --- | --- |
|
||||
| Factory | Strings, aliases, flat dictionaries, invalid inputs, passthrough, and component maps |
|
||||
| Method config | Online/offline selection, ignored layers, platform guards, and weight loading |
|
||||
| Checkpoint metadata | Auto-detection, method mismatch errors, serialized flags, and per-transformer metadata |
|
||||
| Model integration | Correct quantization scope, unquantized components, and parallel weight layouts |
|
||||
| Numerical quality | Same-seed BF16 versus quantized output comparison at representative resolutions and step counts |
|
||||
| Performance | Peak device memory, generation latency, and any backend-specific throughput metric |
|
||||
|
||||
The unit tests live under `tests/diffusion/quantization/`. The trajectory
|
||||
similarity tool in
|
||||
`vllm_omni/quantization/tools/compare_diffusion_trajectory_similarity.py`
|
||||
provides a repeatable quality and memory comparison for diffusion models.
|
||||
Method support and feature combinations remain model- and hardware-dependent;
|
||||
passing factory tests alone does not establish production support.
|
||||
|
||||
## Design invariants
|
||||
|
||||
- `quantization_config` is the single runtime source of truth after config
|
||||
normalization.
|
||||
- A checkpoint method mismatch fails loudly instead of risking weight
|
||||
corruption.
|
||||
- Components without an explicit supported quantization path remain in their
|
||||
original precision.
|
||||
- Optional backend dependencies are imported lazily.
|
||||
- Method implementations own platform-specific capability checks and kernels.
|
||||
- User-facing support claims require model, hardware, numerical, and memory
|
||||
validation in addition to configuration-level tests.
|
||||
@@ -1,59 +0,0 @@
|
||||
# Distributed utils
|
||||
|
||||
This directory (vllm_omni/distributed/ray_utils) contains utilities for distributed execution in vllm-omni, supporting both **Ray** and **Multiprocessing** backends.
|
||||
## 1. Installation
|
||||
```bash
|
||||
pip install "ray[default]"
|
||||
```
|
||||
## 2. Ray Utils
|
||||
|
||||
The `ray_utils` module provides helper functions for managing Ray clusters and actors, which is essential for:
|
||||
* **Multi-node deployment**: Running pipeline stages across different physical machines.
|
||||
* **Resource management**: Efficient GPU/CPU allocation.
|
||||
|
||||
### 2.1 Basic Usage
|
||||
|
||||
To use the Ray backend, specify `worker_backend="ray"` when initializing the engine.
|
||||
|
||||
**Command Line Example:**
|
||||
```bash
|
||||
vllm serve Qwen/Qwen2.5-Omni-7B \
|
||||
--omni \
|
||||
--port 8091 \
|
||||
--worker-backend ray \
|
||||
--ray-address auto
|
||||
```
|
||||
|
||||
### 2.2 Cluster Setup
|
||||
|
||||
**Step 1: Start Head Node**
|
||||
Run this on your primary machine:
|
||||
```bash
|
||||
ray start --head --port=6399
|
||||
```
|
||||
|
||||
**Step 2: Connect Worker Nodes**
|
||||
Run this on each worker machine:
|
||||
```bash
|
||||
ray start --address=<HEAD_NODE_IP>:6399
|
||||
```
|
||||
|
||||
> **Tip**: For a complete cluster setup script, refer to the vLLM example:
|
||||
> [run_cluster.sh](https://github.com/vllm-project/vllm/blob/main/examples/online_serving/run_cluster.sh)
|
||||
|
||||
### 2.3 Distributed Connector Support
|
||||
|
||||
When running on Ray, the system automatically adapts its communication strategy:
|
||||
|
||||
* **Cross-Node**: Recommended to use `MooncakeTransferEngineConnector` (RDMA, fastest) or `MooncakeStoreConnector` (TCP fallback).
|
||||
* **Same-Node**: Can still use `SharedMemoryConnector` for efficiency, or Ray's native object store (plasma).
|
||||
* **SHM threshold default differs**: when `worker_backend="ray"`, the SharedMemoryConnector default threshold is set to `sys.maxsize`, which forces payloads to go inline (no SHM). Override `shm_threshold_bytes` in the connector config if you want SHM for Ray runs.
|
||||
|
||||
### 2.4 Internal Helpers
|
||||
|
||||
* **`initialize_ray_cluster`**: Connects to an existing Ray cluster or starts a local one.
|
||||
|
||||
## 3. Troubleshooting
|
||||
|
||||
* **Connection Issues**: Ensure the Ray head node is accessible and ports (default 6399 in this example) are open.
|
||||
* **Version Mismatch**: Ensure all nodes run the same version of Ray and Python.
|
||||
@@ -0,0 +1,88 @@
|
||||
# Skip-Softmax
|
||||
|
||||
Skip-Softmax is the sparse-attention mode of the `TRTLLM_ATTN` backend. Usage — the config keys
|
||||
and how to pick an operating point — is in
|
||||
[Diffusion Attention Backends](../../user_guide/diffusion/attention_backends.md). This page explains
|
||||
the algorithm.
|
||||
|
||||
## Motivation
|
||||
|
||||
In a long attention row (a video DiT can have tens of thousands of keys), the softmax weight
|
||||
concentrates on a small fraction of the keys; the rest receive near-zero weight and barely move the
|
||||
output. Computing softmax and the value-weighted sum over those keys is wasted work. Skip-Softmax
|
||||
detects, per block of keys, when a block cannot matter and skips its softmax and its value
|
||||
multiply.
|
||||
|
||||
It is approximate: a skipped block still carries a small non-zero contribution, so the mode is
|
||||
opt-in and off by default.
|
||||
|
||||
## The online-softmax pass
|
||||
|
||||
Attention is computed in a single streaming pass over the keys, in tiles of 128. Per query row the
|
||||
kernel maintains three running values:
|
||||
|
||||
- `m` — the largest score seen so far,
|
||||
- `l` — the running denominator `Σ exp(sⱼ − m)`,
|
||||
- `O` — the running numerator `Σ exp(sⱼ − m)·vⱼ`,
|
||||
|
||||
and returns `O / l` at the end. For each key tile it computes the tile's scores `Q · K_jᵀ · scale`,
|
||||
updates `m`, and accumulates that tile's contribution into `l` and `O`. Rescaling `l` and `O` when
|
||||
`m` grows keeps the single pass numerically exact.
|
||||
|
||||
## The skip test
|
||||
|
||||
Once a tile's scores are known, its largest score `tile_max` is compared against the running
|
||||
maximum:
|
||||
|
||||
```text
|
||||
if exp(tile_max − running_max) < threshold:
|
||||
skip this tile # do not compute its softmax weights or its P·V accumulation
|
||||
```
|
||||
|
||||
`exp(tile_max − running_max)` is an upper bound on the softmax weight any key in the tile can
|
||||
receive: if even the tile's best key is far below the current maximum, every key in the tile is
|
||||
negligible, and both the softmax (the exponentials) and the `P·V` accumulation for that tile can be
|
||||
dropped. The tile's contribution to `l` and `O` is simply skipped.
|
||||
|
||||
## What this bounds
|
||||
|
||||
Two properties of the test shape the achievable speedup:
|
||||
|
||||
- **`Q · K_jᵀ` always runs.** The test needs `tile_max`, which comes from the tile's scores, so the
|
||||
score matmul is never skipped — only the softmax and the `P·V` accumulation are. The score matmul
|
||||
and the value matmul are comparable in cost, so skipping every eligible tile removes at most
|
||||
roughly half the attention work; the kernel-level speedup is bounded well under 2×, not unbounded.
|
||||
|
||||
- **The decision is per tile, not per key.** A tile is skipped only when *all* of its keys are
|
||||
collectively negligible; a single important key keeps the whole tile. How many tiles actually
|
||||
qualify depends on the data and rounds down to tile granularity, so `target_sparsity` selects an
|
||||
operating point on a calibrated curve — it is not a promise that a fixed fraction of tiles is
|
||||
skipped.
|
||||
|
||||
## The threshold
|
||||
|
||||
The per-tile threshold is normalized by sequence length:
|
||||
|
||||
```text
|
||||
threshold = factor / seqlen
|
||||
```
|
||||
|
||||
`factor` comes from one of two sources:
|
||||
|
||||
- **Calibrated** — `factor = a · exp(b · target_sparsity)`, where `a` and `b` are fit per model (and
|
||||
per expert, for a multi-expert model) so that a given `target_sparsity` lands near that fraction of
|
||||
skipped tiles on the calibration data.
|
||||
- **Direct** — `factor = skip_softmax.threshold · seqlen`, the calibration-free path; the user sets
|
||||
the threshold themselves.
|
||||
|
||||
## Timestep gating
|
||||
|
||||
`disabled_until_timestep = D` keeps the mode off during the early, high-noise denoise steps and
|
||||
turns it on once the normalized timestep `t` drops to `t ≤ D` (`t` runs `1.0` → `0.0` over the
|
||||
schedule). The early steps set the global structure of the output and their errors propagate through
|
||||
every later step, so keeping them dense costs a few skipped-tile opportunities but protects fidelity.
|
||||
|
||||
`t` is the scheduler's own timestep divided by `num_train_timesteps`, published by the pipeline via
|
||||
`DenoiseProgressMixin.record_denoise_step`. It is deliberately not derived from the step index,
|
||||
because schedulers space their steps non-uniformly. A pipeline that does not publish a timestep
|
||||
stays dense when `disabled_until_timestep` is set.
|
||||
@@ -264,7 +264,7 @@ Complete examples in the codebase:
|
||||
| **Z-Image** | `vllm_omni/diffusion/models/z_image/z_image_transformer.py` | Standard TP | Full implementation with validation |
|
||||
| **FLUX** | `vllm_omni/diffusion/models/flux/flux_transformer.py` | Dual-stream | Image + text streams |
|
||||
| **Qwen-Image** | `vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py` | Standard TP | With RoPE |
|
||||
| **TP Tests** | `tests/e2e/offline_inference/test_zimage_parallelism.py` | E2E testing | TP correctness and performance |
|
||||
| **TP Tests** | `tests/diffusion/distributed/test_tensor_parallel.py` | E2E testing | TP correctness and performance |
|
||||
| **Constraint Tests** | `tests/diffusion/models/z_image/test_zimage_tp_constraints.py` | Unit testing | Validation logic |
|
||||
|
||||
---
|
||||
|
||||
@@ -445,6 +445,7 @@ Complete examples in the codebase:
|
||||
| **Z-Image** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl.py` | ✅ | ❌ |
|
||||
| **Wan2.2** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py` | ✅ | ✅ |
|
||||
| **Qwen-Image** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage.py` | ✅ | ❌ |
|
||||
| **FLUX.2-dev** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_flux2.py` | ✅ | ✅ |
|
||||
|
||||
---
|
||||
|
||||
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 195 KiB |
+56
-8
@@ -1,6 +1,9 @@
|
||||
# Design Documents
|
||||
|
||||
This section contains design documents and architecture specifications for vLLM-Omni.
|
||||
This section contains design documents and architecture specifications for
|
||||
vLLM-Omni. The sidebar groups documents by the system concern they describe;
|
||||
retired module pages are preserved in the legacy archive instead of remaining
|
||||
in the active navigation.
|
||||
|
||||
## Architecture Documents
|
||||
|
||||
@@ -8,22 +11,63 @@ This section contains design documents and architecture specifications for vLLM-
|
||||
|
||||
## Feature Design Documents
|
||||
|
||||
- [Disaggregated Inference](feature/disaggregated_inference.md)
|
||||
- [Ray-based Execution](feature/ray_based_execution.md)
|
||||
- [Adding Step Execution Support for Diffusion Pipelines](feature/diffusion_step_execution.md)
|
||||
- [Request-Level Batching for Diffusion](feature/diffusion_request_level_batching.md)
|
||||
- [Continuous Batching for Step-Wise Diffusion](feature/diffusion_continuous_batching.md)
|
||||
### Runtime and stage execution
|
||||
|
||||
## Infrastructure Design Documents
|
||||
- [Disaggregated Inference](feature/disaggregated_inference.md)
|
||||
- [Async Chunk](feature/async_chunk.md)
|
||||
- [Async Diffusion Output](feature/async_diffusion_output.md)
|
||||
- [Async Omni Output Materialization](feature/omni_async_output_materialization.md)
|
||||
- [Automatic Prefix Caching in Omni Models](feature/prefix_caching.md)
|
||||
|
||||
### Communication
|
||||
|
||||
#### OmniConnector implementations
|
||||
|
||||
- [Mooncake Store Connector](feature/omni_connectors/mooncake_store_connector.md)
|
||||
- [Mooncake Transfer Engine Connector](feature/omni_connectors/mooncake_transfer_engine_connector.md)
|
||||
- [Mori Transfer Engine Connector](feature/omni_connectors/mori_transfer_engine_connector.md)
|
||||
- [Shared Memory Connector](feature/omni_connectors/shared_memory_connector.md)
|
||||
- [Yuanrong Store Connector](feature/omni_connectors/yuanrong_connector.md)
|
||||
- [Yuanrong Transfer Engine Connector](feature/omni_connectors/yuanrong_transfer_engine_connector.md)
|
||||
|
||||
### Diffusion acceleration
|
||||
|
||||
#### Parallelism
|
||||
|
||||
- [CFG-Parallel](feature/cfg_parallel.md)
|
||||
- [Expert Parallel](feature/expert_parallel.md)
|
||||
- [Hybrid Sharded Data Parallel (HSDP)](feature/hsdp.md)
|
||||
- [Pipeline Parallel](feature/pipeline_parallel.md)
|
||||
- [Sequence Parallel](feature/sequence_parallel.md)
|
||||
- [Tensor Parallel](feature/tensor_parallel.md)
|
||||
- [VAE Patch Parallelism](feature/vae_parallel.md)
|
||||
|
||||
#### Attention Backends
|
||||
|
||||
- [Skip-Softmax](feature/skip_softmax.md)
|
||||
|
||||
#### Quantization
|
||||
|
||||
- [Quantization](feature/quantization.md)
|
||||
|
||||
- [Cache-DiT](feature/cache_dit.md)
|
||||
- [TeaCache](feature/teacache.md)
|
||||
- [Diffusion Continuous Batching](feature/diffusion_continuous_batching.md)
|
||||
- [Distributed Layerwise Offload](feature/distributed_layerwise_offload.md)
|
||||
|
||||
## Infrastructure and Performance
|
||||
|
||||
- [Prometheus Metrics](metrics.md)
|
||||
- [Speech Generation Performance Optimizations](qwen3_omni_tts_performance_optimization.md)
|
||||
|
||||
## Module Design Documents
|
||||
|
||||
- [Entrypoints](module/entrypoints.md)
|
||||
- [Entrypoints and Serving Boundaries](module/entrypoints.md)
|
||||
- [vLLM-Omni Configuration](module/vllm_omni_config.md)
|
||||
- [Input, Output, and Modality Contracts](module/input_output_modality_contracts.md)
|
||||
- [Error Classification, Propagation, and Rendering](module/error_contracts.md)
|
||||
- [Engine Orchestration](module/engine_orchestration.md)
|
||||
- [Stage Runtime and Replica Lifecycle](module/stage_runtime.md)
|
||||
- [OmniConnector](module/omni_connector.md)
|
||||
- [Model Integration](module/model_integration.md)
|
||||
- [Autoregressive Runtime](module/ar_runtime.md)
|
||||
@@ -40,3 +84,7 @@ This section contains design documents and architecture specifications for vLLM-
|
||||
- [Observability](module/observability.md)
|
||||
- [Profiling](module/profiling.md)
|
||||
- [Benchmarking](module/benchmarking.md)
|
||||
|
||||
The pre-#5137 pages are preserved in the
|
||||
[legacy module archive](module/archive/README.md) for historical reference and
|
||||
are not active design contracts.
|
||||
@@ -148,7 +148,13 @@ Upstream vLLM's `Scheduler.make_stats()` runs on every AR generation step, retur
|
||||
|
||||
`OmniSchedulerMixin.make_stats()` (in `vllm_omni/core/sched/omni_scheduler_mixin.py`) throttles stats emission to at most once per second **per scheduler** — i.e. per `(stage, replica)` since each replica owns its own scheduler instance. Between intervals it returns `None`, which the engine core skips serializing. This keeps gauges fresh enough for Prometheus scrapes (typically 15-30s intervals) while eliminating the per-step overhead.
|
||||
|
||||
The orchestrator side does not add its own throttle on top: the per-replica recording loop gates only on `raw_outputs.scheduler_stats is not None` (i.e. this replica's scheduler passed its own 1 Hz gate). A previous global `_last_stats_ts` on the orchestrator was removed because it starved every replica other than the first to emit within each second.
|
||||
The orchestrator side does not add its own throttle on top. Scheduler gauges are
|
||||
recorded when `raw_outputs.scheduler_stats is not None` (i.e. this replica's
|
||||
scheduler passed its own 1 Hz gate), while output-batch request/token metrics
|
||||
are recorded from `IterationStats` whenever a non-empty output batch is
|
||||
processed. A previous global `_last_stats_ts` on the orchestrator was removed
|
||||
because it starved every replica other than the first to emit within each
|
||||
second.
|
||||
|
||||
## Metric Definitions
|
||||
|
||||
|
||||
@@ -0,0 +1,18 @@
|
||||
# Archived module design pages
|
||||
|
||||
This directory preserves module-design pages that predate the structure
|
||||
proposed in [issue #5137](https://github.com/vllm-project/vllm-omni/issues/5137).
|
||||
They are historical snapshots, not active architectural contracts, and are
|
||||
intentionally excluded from the documentation navigation.
|
||||
|
||||
| Archived page | Active destination |
|
||||
| --- | --- |
|
||||
| [AR Module](ar_module.md) | [Autoregressive Runtime](../ar_runtime.md) |
|
||||
| [AsyncOmni Architecture](async_omni_architecture.md) | [Engine Orchestration](../engine_orchestration.md) and [Stage Runtime](../stage_runtime.md) |
|
||||
| [DiT Module](dit_module.md) | [Diffusion module overview](../diffusion/index.md) |
|
||||
| [Entrypoint Module](entrypoint_module.md) | [Entrypoints and Serving Boundaries](../entrypoints.md) |
|
||||
|
||||
Do not update these pages with current design guidance. Move still-useful
|
||||
material into the active destination and review it against current code and
|
||||
tests instead. No compatibility redirect pages are provided at the former
|
||||
paths.
|
||||
@@ -0,0 +1,387 @@
|
||||
# AutoRegressive (AR) Module
|
||||
|
||||
## 1. Overview
|
||||
|
||||
The AutoRegressive (AR) module in vLLM-Omni handles autoregressive generation stages, primarily used for text, chain-of-thought(COT), and audio latent tokens generation stages in multi-stage models like Qwen2.5-Omni, Qwen3-Omni, BAGEL, .etc. Unlike some representative non-autoregressive generation stages (e.g., Diffusion), AR stages generate tokens sequentially, one at a time, following the standard transformer decoder pattern.
|
||||
|
||||
The AR module of vLLM-Omni extends vLLM's core components to support:
|
||||
|
||||
- **Multimodal inputs/outputs**: Processing images, videos, and audio alongside text
|
||||
- **Direct embedding transfer**: Passing pre-computed prompt embeddings between pipeline stages via serialized payloads
|
||||
- **Additional information flow**: Carrying per-request metadata (tensors, lists) through the pipeline
|
||||
- **Hidden state exposure**: Exposing per-request hidden representations for downstream stages
|
||||
- **Basic generator support**: Support some basic heterogeneous architecture such as Convolution, LSTM, etc.
|
||||
|
||||
As shown in the [end2end example](../../../user_guide/examples/offline_inference/qwen3_omni.md), AR module can be widely applied across multiple stages, generating text tokens in thinker(AR), audio latent tokens in talker(AR) and audio wave in code2wav(Convolution).
|
||||
|
||||
## 2. Relationship with vLLM
|
||||
|
||||
The AR module builds upon vLLM main framework through inheritance, extending core classes while preserving compatibility with vLLM's scheduling, batching, KV cache management, and execution mechanisms.
|
||||
|
||||
### Inheritance Hierarchy
|
||||
- Scheduler
|
||||
|
||||
```mermaid
|
||||
classDiagram
|
||||
class VLLMScheduler {
|
||||
+schedule() SchedulerOutput
|
||||
+update_from_output() EngineCoreOutputs
|
||||
}
|
||||
class OmniARScheduler {
|
||||
+schedule() SchedulerOutput
|
||||
}
|
||||
class OmniGenerationScheduler {
|
||||
+schedule() SchedulerOutput
|
||||
+update_from_output() EngineCoreOutputs
|
||||
}
|
||||
VLLMScheduler <|-- OmniARScheduler
|
||||
VLLMScheduler <|-- OmniGenerationScheduler
|
||||
```
|
||||
- Worker
|
||||
|
||||
```mermaid
|
||||
classDiagram
|
||||
class GPUWorker {
|
||||
+init_device()
|
||||
+model_runner
|
||||
}
|
||||
class GPUARWorker {
|
||||
+init_device()
|
||||
}
|
||||
class GPUGenerationWorker {
|
||||
+init_device()
|
||||
}
|
||||
GPUWorker <|-- GPUARWorker
|
||||
GPUWorker <|-- GPUGenerationWorker
|
||||
```
|
||||
- ModelRunner
|
||||
|
||||
```mermaid
|
||||
classDiagram
|
||||
class GPUModelRunner {
|
||||
+execute_model()
|
||||
+sample_tokens()
|
||||
}
|
||||
class OmniGPUModelRunner {
|
||||
+_update_states()
|
||||
+_preprocess()
|
||||
+_model_forward()
|
||||
}
|
||||
class GPUARModelRunner {
|
||||
+execute_model()
|
||||
+sample_tokens()
|
||||
}
|
||||
class GPUGenerationModelRunner {
|
||||
+execute_model()
|
||||
}
|
||||
GPUModelRunner <|-- OmniGPUModelRunner
|
||||
OmniGPUModelRunner <|-- GPUARModelRunner
|
||||
OmniGPUModelRunner <|-- GPUGenerationModelRunner
|
||||
```
|
||||
- InputProcessor/OutputProcessor
|
||||
|
||||
```mermaid
|
||||
classDiagram
|
||||
class InputProcessor {
|
||||
+process_inputs() EngineCoreRequest
|
||||
}
|
||||
|
||||
class VLLMOutputProcessor {
|
||||
+process_outputs() OutputProcessorOutput
|
||||
}
|
||||
class MultimodalOutputProcessor {
|
||||
+process_outputs() OutputProcessorOutput
|
||||
+_route_and_normalize()
|
||||
}
|
||||
VLLMOutputProcessor <|-- MultimodalOutputProcessor
|
||||
```
|
||||
|
||||
### Key Extensions
|
||||
|
||||
- **Scheduler**: `OmniARScheduler` extends `vllm.v1.core.sched.scheduler.Scheduler` to enrich scheduled requests with omni-specific payloads
|
||||
- **Worker**: `GPUARWorker` extends `vllm.v1.worker.gpu_worker.Worker` to initialize AR-specific model runners
|
||||
- **ModelRunner**: `GPUARModelRunner` extends `OmniGPUModelRunner` → `vllm.v1.worker.gpu_model_runner.GPUModelRunner` to expose hidden states and handle multimodal outputs
|
||||
- **InputProcessor**: Stage-0 uses upstream `vllm.v1.engine.input_processor.InputProcessor`; `AsyncOmniEngine` then restores omni-specific payloads (for example `additional_information` and `prompt_embeds`) when building `OmniEngineCoreRequest`
|
||||
- **OutputProcessor**: `MultimodalOutputProcessor` extends `vllm.v1.engine.output_processor.OutputProcessor` to route and accumulate multimodal outputs
|
||||
|
||||
## 3. Scheduler Design
|
||||
|
||||
The AR module provides two scheduler implementations: one for standard autoregressive generation and one for basic heterogeneous architectures.
|
||||
|
||||
### Request Flow
|
||||
|
||||
The following diagram illustrates the request flow through the AR module components:
|
||||
|
||||
```mermaid
|
||||
flowchart TD
|
||||
A[InputProcessor stage-0 in AsyncOmniEngine] -->|EngineCoreRequest then upgraded to OmniEngineCoreRequest| B[OmniARScheduler]
|
||||
B -->|schedule: OmniNewRequestData| C[GPUARWorker]
|
||||
C -->|SchedulerOutput| D[GPUARModelRunner]
|
||||
D -->|execute_model: None| E[Model Forward Pass]
|
||||
E -->|hidden_states, logits| D
|
||||
D -->|sample_tokens: OmniModelRunnerOutput| F[OmniARScheduler]
|
||||
F -->|update_from_output| G[MultimodalOutputProcessor]
|
||||
G -->|RequestOutput| H[Client/Downstream Stage]
|
||||
|
||||
style A fill:#e1f5ff
|
||||
style B fill:#fff4e1
|
||||
style C fill:#e8f5e9
|
||||
style D fill:#f3e5f5
|
||||
style G fill:#fce4ec
|
||||
```
|
||||
|
||||
The flow follows vLLM's standard pattern: input processing → scheduling → worker execution → output processing, with omni-specific enrichments at each stage.
|
||||
|
||||
### OmniARScheduler
|
||||
|
||||
`OmniARScheduler` extends the base vLLM scheduler with minimal modifications, focusing on enriching scheduled requests with omni-specific payloads.
|
||||
|
||||
#### Modified API: `schedule()`
|
||||
|
||||
The scheduler wraps base `NewRequestData` entries with `OmniNewRequestData` to include prompt embeddings and additional information:
|
||||
|
||||
```python
|
||||
def schedule(self) -> SchedulerOutput:
|
||||
scheduler_output = super().schedule()
|
||||
# Rewrap base NewRequestData entries with OmniNewRequestData
|
||||
new_list = []
|
||||
for nr in scheduler_output.scheduled_new_reqs:
|
||||
request = self.requests.get(nr.req_id)
|
||||
omni_nr = OmniNewRequestData(
|
||||
req_id=nr.req_id,
|
||||
prompt_token_ids=nr.prompt_token_ids,
|
||||
# ... other base fields ...
|
||||
prompt_embeds=getattr(request, "prompt_embeds", None),
|
||||
additional_information=getattr(request, "additional_information", None),
|
||||
)
|
||||
new_list.append(omni_nr)
|
||||
scheduler_output.scheduled_new_reqs = new_list
|
||||
return scheduler_output
|
||||
```
|
||||
|
||||
The `update_from_output()` method remains unchanged, inheriting standard request lifecycle management from the base scheduler.
|
||||
|
||||
### OmniGenerationScheduler
|
||||
|
||||
`OmniGenerationScheduler` implements a fast-path scheduling strategy for basic heterogeneous architectures that process all input tokens in a single step.
|
||||
|
||||
#### Modified API: `schedule()`
|
||||
|
||||
Allocates all input tokens for a request at once (or 1 placeholder if zero), falling back to default scheduling if budget is insufficient:
|
||||
|
||||
```python
|
||||
def schedule(self) -> SchedulerOutput:
|
||||
# Fast path: allocate all input tokens at once
|
||||
while self.waiting and token_budget > 0:
|
||||
request = self.waiting.peek_request()
|
||||
required_tokens = max(getattr(request, "num_prompt_tokens", 0), 1)
|
||||
if required_tokens > token_budget:
|
||||
break # Fall back to default scheduling
|
||||
# Allocate and schedule...
|
||||
```
|
||||
|
||||
#### Modified API: `update_from_output()`
|
||||
|
||||
Marks requests as finished immediately after one step, since generation models complete in a single forward pass:
|
||||
|
||||
```python
|
||||
def update_from_output(self, ...) -> dict[int, EngineCoreOutputs]:
|
||||
# ...
|
||||
# Diffusion request: completes in one step
|
||||
request.status = RequestStatus.FINISHED_STOPPED
|
||||
kv_transfer_params = self._free_request(request)
|
||||
# ...
|
||||
```
|
||||
|
||||
## 4. Worker and ModelRunner Design
|
||||
|
||||
### GPUARWorker
|
||||
|
||||
`GPUARWorker` initializes the AR-specific model runner while maintaining standard device initialization:
|
||||
|
||||
```python
|
||||
class GPUARWorker(GPUWorker):
|
||||
def init_device(self):
|
||||
# ... standard device initialization ...
|
||||
self.model_runner = GPUARModelRunner(self.vllm_config, self.device)
|
||||
```
|
||||
|
||||
### GPUARModelRunner
|
||||
|
||||
`GPUARModelRunner` follows vLLM's two-phase execute/sample flow while exposing hidden states and multimodal outputs.
|
||||
|
||||
#### Two-Phase Execution
|
||||
|
||||
**Phase 1: `execute_model()`** - Runs forward pass and stores state:
|
||||
- Computes logits from hidden states
|
||||
- Stores `ExecuteModelState` with hidden states, logits, and multimodal outputs
|
||||
- Returns `None` to defer sampling
|
||||
|
||||
**Phase 2: `sample_tokens()`** - Samples tokens and builds output:
|
||||
- Retrieves stored state from `execute_model()`
|
||||
- Samples tokens using logits
|
||||
- Extracts per-request hidden states and multimodal outputs
|
||||
- Builds `OmniModelRunnerOutput` with `pooler_output` containing hidden states
|
||||
|
||||
```python
|
||||
def sample_tokens(self, grammar_output) -> OmniModelRunnerOutput:
|
||||
# Retrieve stored state
|
||||
hidden_states, multimodal_outputs = self.execute_model_state
|
||||
|
||||
# Sample tokens
|
||||
sampler_output = self._sample(logits, spec_decode_metadata)
|
||||
|
||||
# Extract per-request hidden states
|
||||
pooler_output = []
|
||||
for rid in req_ids:
|
||||
hidden_slice = hidden_states_cpu[start:end]
|
||||
payload = {"hidden": hidden_slice}
|
||||
# Add multimodal outputs if present
|
||||
pooler_output.append(payload)
|
||||
|
||||
return OmniModelRunnerOutput(
|
||||
pooler_output=pooler_output,
|
||||
# ... other fields ...
|
||||
)
|
||||
```
|
||||
|
||||
### GPUGenerationModelRunner
|
||||
|
||||
`GPUGenerationModelRunner` implements a simplified single-phase execution for basic heterogeneous architectures:
|
||||
|
||||
- No logits computation or token sampling
|
||||
- Direct generation from forward pass in model implementation
|
||||
- Returns outputs via `pooler_output` immediately after forward pass
|
||||
|
||||
### OmniGPUModelRunner
|
||||
|
||||
`OmniGPUModelRunner` provides shared functionality for both AR and Generation runners:
|
||||
|
||||
#### Prompt Embeddings Overlay
|
||||
|
||||
During prefill, overlays custom `prompt_embeds` from request state onto `inputs_embeds`:
|
||||
|
||||
```python
|
||||
def _collect_additional_information_for_prefill(self, num_scheduled_tokens_np):
|
||||
for req_index, req_id in enumerate(self.input_batch.req_ids):
|
||||
req_state = self.requests[req_id]
|
||||
pe_cpu = getattr(req_state, "prompt_embeds_cpu", None)
|
||||
# Overlay prompt_embeds for prefill portion
|
||||
if pe_cpu is not None:
|
||||
src = pe_cpu[num_computed_tokens:num_computed_tokens + overlay_len]
|
||||
self.inputs_embeds[start_offset:start_offset + overlay_len].copy_(src)
|
||||
```
|
||||
|
||||
#### Additional Information Processing
|
||||
|
||||
Decodes and manages `additional_information` payloads:
|
||||
- Decodes serialized payloads → CPU tensors in request state
|
||||
- Passes runtime information to model via `runtime_additional_information` kwarg
|
||||
- Processes model-provided updates via `postprocess()` hook
|
||||
- Merges updates back into request state
|
||||
|
||||
#### M-RoPE Position Initialization
|
||||
|
||||
For multimodal models using M-RoPE (e.g., Qwen2-VL), computes position encodings from multimodal feature metadata (image grids, video grids, audio features).
|
||||
|
||||
## 5. Input/Output Processing
|
||||
|
||||
### Processing Pipeline
|
||||
|
||||
The input/output processing pipeline handles serialization, routing, and accumulation of multimodal data:
|
||||
|
||||
```mermaid
|
||||
sequenceDiagram
|
||||
participant Client
|
||||
participant AsyncOmniEngine
|
||||
participant InputProcessor
|
||||
participant Scheduler
|
||||
participant ModelRunner
|
||||
participant MultimodalOutputProcessor
|
||||
participant Client
|
||||
|
||||
Client->>AsyncOmniEngine: prompt + prompt_embeds + additional_info
|
||||
AsyncOmniEngine->>InputProcessor: process_inputs()
|
||||
InputProcessor->>Scheduler: EngineCoreRequest
|
||||
AsyncOmniEngine->>AsyncOmniEngine: _upgrade_to_omni_request() + serialize_additional_information()
|
||||
Scheduler->>ModelRunner: OmniNewRequestData (with payloads)
|
||||
ModelRunner->>ModelRunner: Decode payloads → CPU tensors
|
||||
ModelRunner->>ModelRunner: Overlay prompt_embeds on inputs_embeds
|
||||
ModelRunner->>ModelRunner: Forward pass with runtime_additional_information
|
||||
ModelRunner->>ModelRunner: Extract hidden states + multimodal outputs
|
||||
ModelRunner->>MultimodalOutputProcessor: OmniModelRunnerOutput (pooler_output)
|
||||
MultimodalOutputProcessor->>MultimodalOutputProcessor: Route by output_type
|
||||
MultimodalOutputProcessor->>MultimodalOutputProcessor: Accumulate tensors in OmniRequestState
|
||||
MultimodalOutputProcessor->>MultimodalOutputProcessor: Consolidate tensor lists
|
||||
MultimodalOutputProcessor->>Client: RequestOutput (with multimodal_output)
|
||||
```
|
||||
|
||||
### Stage-0 Input Processing
|
||||
|
||||
Stage-0 now uses upstream `InputProcessor` directly, and `AsyncOmniEngine` upgrades the request to `OmniEngineCoreRequest` while restoring omni-specific payloads.
|
||||
|
||||
```python
|
||||
request = self.input_processor.process_inputs(
|
||||
request_id=request_id,
|
||||
prompt=prompt,
|
||||
params=params,
|
||||
supported_tasks=self.supported_tasks,
|
||||
)
|
||||
request = _upgrade_to_omni_request(request, prompt)
|
||||
```
|
||||
|
||||
### MultimodalOutputProcessor
|
||||
|
||||
`MultimodalOutputProcessor` routes outputs by modality type and accumulates multimodal tensors.
|
||||
|
||||
#### Output Routing
|
||||
|
||||
Routes `EngineCoreOutput` by `output_type` attribute:
|
||||
- `"text"`: Standard text generation path
|
||||
- `"image"`, `"audio"`, `"latents"`: Extract from `pooling_output` or `multimodal_outputs`
|
||||
- Fallback: Heuristic based on presence of `pooling_output`
|
||||
|
||||
#### Tensor Accumulation
|
||||
|
||||
`OmniRequestState` accumulates multimodal tensors across multiple steps:
|
||||
|
||||
```python
|
||||
def add_multimodal_tensor(self, payload, mm_type):
|
||||
# Normalize payload to dict
|
||||
incoming = {mm_type or "hidden": payload}
|
||||
|
||||
# Accumulate: convert tensors to lists for deferred concatenation
|
||||
if isinstance(v, torch.Tensor) and isinstance(existing, torch.Tensor):
|
||||
self.mm_accumulated[k] = [existing, v] # List accumulation
|
||||
```
|
||||
|
||||
Before final output, consolidates tensor lists via concatenation:
|
||||
|
||||
```python
|
||||
def _consolidate_multimodal_tensors(self):
|
||||
for k, v in self.mm_accumulated.items():
|
||||
if isinstance(v, list) and isinstance(v[0], torch.Tensor):
|
||||
self.mm_accumulated[k] = torch.cat(v, dim=0) # Concatenate
|
||||
```
|
||||
|
||||
The consolidated tensors are attached to `RequestOutput.multimodal_output` for consumption by downstream stages or clients.
|
||||
|
||||
## 6. Summary
|
||||
|
||||
The AR module of vLLM-Omni extends vLLM through strategic inheritance and minimal API modifications:
|
||||
|
||||
### Key Design Patterns
|
||||
|
||||
1. **Inheritance over composition**: Extends vLLM classes to preserve compatibility with existing scheduling, batching, and execution mechanisms
|
||||
2. **Payload serialization**: Uses serialized `additional_information` payloads together with prompt-embedding handoff for efficient inter-stage transfer
|
||||
3. **Two-phase execution**: Maintains vLLM's execute/sample separation for AR models while supporting single-phase execution for generation models
|
||||
4. **Multimodal routing**: Routes outputs by `output_type` and accumulates tensors incrementally to support streaming
|
||||
|
||||
### Differences from vLLM
|
||||
|
||||
- **Payload support**: Serialized additional information and prompt embeddings enable direct transfer between pipeline stages
|
||||
- **Multimodal handling**: Extended input/output processors support images, audio, and other modalities alongside text
|
||||
- **Hidden state exposure**: AR model runners expose per-request hidden states via `pooler_output` for downstream consumption
|
||||
- **Generation scheduler**: Fast-path scheduling for basic heterogeneous architectures that complete in one step
|
||||
|
||||
The AR module seamlessly integrates with vLLM's existing infrastructure while adding the necessary extensions for multi-stage, multimodal generation pipelines.
|
||||
@@ -0,0 +1,203 @@
|
||||
# AsyncOmni Architecture (Qwen3-Omni Example)
|
||||
|
||||
## 1. System Architecture
|
||||
|
||||
```text
|
||||
• ┌─────────────────────────────────────────────────────────────────────────────────┐
|
||||
│ API Layer │
|
||||
│ ┌─────────────────────────────────────┐ ┌──────────────────────────────────┐ │
|
||||
│ │ AsyncOmni (EngineClient) │ │ Omni │ │
|
||||
│ │ • generate() / abort() / shutdown() │ │ • generate() │ │
|
||||
│ │ • _final_output_handler() │ │ | │
|
||||
│ └─────────────────────────────────────┘ └──────────────────────────────────┘ │
|
||||
├─────────────────────────────────────────────────────────────────────────────────┤
|
||||
│ Engine Layer (Proxy) │
|
||||
│ ┌───────────────────────────────────────────────────────────────────────────┐ │
|
||||
│ │ AsyncOmniEngine │ │
|
||||
│ │ • _bootstrap_orchestrator() & _initialize_stages() │ │
|
||||
│ │ • add_request() / add_request_async() -> input_processor.process_inputs() │ │
|
||||
│ │ • try_get_output() / try_get_output_async() │ │
|
||||
│ └───────────────────┬─────────────────────────────────▲─────────────────────┘ │
|
||||
│ request_queue (janus.Queue) output_queue (janus.Queue) │
|
||||
├──────────────────────┼─────────────────────────────────┼────────────────────────┤
|
||||
│ ▼ Orchestration Layer │ │
|
||||
│ ┌───────────────────────────────────────────────────────────────────────────┐ │
|
||||
│ │ Orchestrator [background thread] │ │
|
||||
│ │ • _request_handler() │ │
|
||||
│ │ - stage_client.add_request_async() & _prewarm_async_chunk_stages() │ │
|
||||
│ │ • _orchestration_output_handler() │ │
|
||||
│ │ - _process_stage_outputs() -> output_processors[i].process_outputs() │ │
|
||||
│ │ - _route_output() & _forward_to_next_stage() │ │
|
||||
│ └──────────┬─────────────────────────┬────────────────────────┬─────────────┘ │
|
||||
├─────────────┼─────────────────────────┼────────────────────────┼────────────────┤
|
||||
│ │ Communication Layer │ │
|
||||
│ ┌───────────────────────┐ ┌───────────────────────┐ ┌───────────────────────┐ │
|
||||
│ │ StageEngineCoreClient │ │ StageEngineCoreClient │ │ StageDiffusionClient │ │
|
||||
│ │ • ZMQ ROUTER / PULL │ │ • ZMQ ROUTER / PULL │ │ • ZMQ ROUTER / PULL │ │
|
||||
│ │ • Msgpack codec │ │ • Msgpack codec │ │ • Msgpack codec │ │
|
||||
│ └──────────┬────────────┘ └──────────┬────────────┘ └──────────┬────────────┘ │
|
||||
│ ▼ ZMQ IPC ▼ ZMQ IPC ▼ ZMQ IPC │
|
||||
├─────────────────────────────────────────────────────────────────────────────────┤
|
||||
│ Execution Layer │
|
||||
│ ┌───────────────────────┐ ┌───────────────────────┐ ┌───────────────────────┐ │
|
||||
│ │ StageCoreProc │ │ StageCoreProc │ │ DiffusionEngine │ │
|
||||
│ │ [background process] │ │ [background process] │ │ [background process] │ │
|
||||
│ └───────────────────────┘ └───────────────────────┘ └───────────────────────┘ │
|
||||
└─────────────────────────────────────────────────────────────────────────────────┘
|
||||
```
|
||||
|
||||
## 2. Execution Flow (Arrow Steps, one generate request)
|
||||
|
||||
```text
|
||||
[1] App
|
||||
-> AsyncOmni.generate(prompt, request_id)
|
||||
|
||||
[2] AsyncOmni
|
||||
-> _final_output_handler() (started on first request)
|
||||
-> AsyncOmniEngine.add_request(stage_id=0, ...)
|
||||
|
||||
[3] AsyncOmniEngine.add_request
|
||||
-> (if stage-0 is llm and input is not EngineCoreRequest)
|
||||
InputProcessor.process_inputs()
|
||||
OutputProcessor[0].add_request()
|
||||
-> request_queue.put(add_request_msg)
|
||||
|
||||
[4] Orchestrator._request_handler
|
||||
-> _handle_add_request(msg)
|
||||
-> stage_clients[0].add_request_async(...)
|
||||
|
||||
[5] Orchestrator._orchestration_loop (loop)
|
||||
-> poll stage output
|
||||
- llm stage: await get_output_async()
|
||||
- diffusion stage: get_diffusion_output_nowait()
|
||||
-> (llm stage) output_processors[i].process_outputs(...)
|
||||
-> _route_output(...)
|
||||
-> if finished and not final_stage and non-async-chunk:
|
||||
_forward_to_next_stage(...)
|
||||
-> next_stage.add_request_async(...)
|
||||
-> output_queue.put(output)
|
||||
|
||||
[6] AsyncOmni._final_output_loop (background coroutine)
|
||||
-> AsyncOmniEngine.try_get_output_async()
|
||||
-> route by request_id to ClientRequestState.queue
|
||||
|
||||
[7] AsyncOmni._process_orchestrator_results
|
||||
-> read from ClientRequestState.queue
|
||||
-> _process_single_result(...)
|
||||
-> yield OmniRequestOutput
|
||||
|
||||
[8] Exit condition
|
||||
-> receive result["finished"] == True
|
||||
-> generate() ends
|
||||
```
|
||||
|
||||
## 3. Runtime Sequence (one generate request)
|
||||
|
||||
```mermaid
|
||||
sequenceDiagram
|
||||
participant APP as App
|
||||
participant AO as AsyncOmni
|
||||
participant ENG as AsyncOmniEngine
|
||||
participant ORCH as Orchestrator
|
||||
participant S0 as Stage-0 Client
|
||||
participant SN as Next Stage Client
|
||||
|
||||
APP->>AO: generate
|
||||
AO->>AO: start output_handler once
|
||||
AO->>ENG: add_request(stage_id=0, ...)
|
||||
ENG->>ENG: input_processor.process_inputs()
|
||||
ENG->>ORCH: request_queue.put(add_request)
|
||||
|
||||
ORCH->>ORCH: _handle_add_request
|
||||
ORCH->>S0: add_request_async
|
||||
|
||||
loop poll route forward
|
||||
ORCH->>S0: get_output_async / get_diffusion_output_nowait
|
||||
ORCH->>ORCH: _route_output
|
||||
alt need forward to next stage
|
||||
ORCH->>SN: add_request_async
|
||||
end
|
||||
ORCH-->>ENG: output_queue.put
|
||||
end
|
||||
|
||||
AO->>ENG: try_get_output_async
|
||||
ENG-->>AO: message
|
||||
AO-->>APP: yield OmniRequestOutput
|
||||
```
|
||||
|
||||
## 4. Comparison
|
||||
|
||||
Previous topology (reference):
|
||||
|
||||
```text
|
||||
┌────────────────────────────────────────────────────────────────────────────┐
|
||||
│ Main Process │
|
||||
│ ┌──────────────────────┐ ┌────────────────────────────────────────────┐ │
|
||||
│ │ generate() │ │ final_output_handler() │ │
|
||||
│ └──────────────────────┘ └────────────────────────────────────────────┘ │
|
||||
└──────────┬─────────────────────────┬─────────────────────────┬─────────────┘
|
||||
mp.Queue (in_q/out_q) mp.Queue (in_q/out_q) mp.Queue (in_q/out_q)
|
||||
▼▲ ▼▲ ▼▲
|
||||
┌───────────────────────┐ ┌───────────────────────┐ ┌──────────────────────┐
|
||||
│ Worker Proc-0 │ │ Worker Proc-1 │ │ Worker Proc-2 │
|
||||
│ (Thinker LLM) │ │ (Talker LLM) │ │ (Vocoder) │
|
||||
│ ┌────────────────┐ │ │ ┌────────────────┐ │ │ ┌────────────────┐ │
|
||||
│ │_stage_worker │ │ │ │_stage_worker │ │ │ │_stage_worker │ │
|
||||
│ │_async() │ │ │ │_async() │ │ │ │_async() │ │
|
||||
│ └────────────────┘ │ │ └────────────────┘ │ │ └────────────────┘ │
|
||||
│ ┌────────────────┐ │ │ ┌────────────────┐ │ │ ┌────────────────┐ │
|
||||
│ │output_handler()│ │ │ │output_handler()│ │ │ │output_handler()│ │
|
||||
│ └────────────────┘ │ │ └────────────────┘ │ │ └────────────────┘ │
|
||||
└──────────┬────────────┘ └──────────┬────────────┘ └──────────┬───────────┘
|
||||
ZMQ ▼ ▲ ZMQ ZMQ ▼ ▲ ZMQ ZMQ ▼ ▲ ZMQ
|
||||
┌──────────────────────┐ ┌──────────────────────┐ ┌──────────────────────┐
|
||||
│ EngineCore Proc-0 │ │ EngineCore Proc-1 │ │ EngineCore Proc-2 │
|
||||
│ (Thinker) │ │ (Talker) │ │ (Vocoder) │
|
||||
└──────────────────────┘ └──────────────────────┘ └──────────────────────┘
|
||||
```
|
||||
|
||||
Current topology:
|
||||
|
||||
```text
|
||||
┌────────────────────────────────────────────────────────────────────────────┐
|
||||
│ Main Process │
|
||||
│ ┌──────────────────────────────────────────────────────────────────────┐ │
|
||||
│ │ Main Thread │ │
|
||||
│ │ ┌──────────────────────┐ ┌─────────────────────────────────────┐ │ │
|
||||
│ │ │ generate() │ │ final_output_handler() │ │ │
|
||||
│ │ └──────────────────────┘ └─────────────────────────────────────┘ │ │
|
||||
│ └──────────────────────────────────────────────────────────────────────┘ │
|
||||
│ janus.Queue (request_queue) ▼ ▲ janus.Queue (output_queue) │
|
||||
│ ┌──────────────────────────────────────────────────────────────────────┐ │
|
||||
│ │ Orchestrator Thread │ │
|
||||
│ │ ┌──────────────────────┐ ┌──────────────────────────────────────┐ │ │
|
||||
│ │ │ _request_handler() │ │ _orchestration_output_handler() │ │ │
|
||||
│ │ └──────────────────────┘ └──────────────────────────────────────┘ │ │
|
||||
│ │ ┌────────────────────────────────────────────────────────────────┐ │ │
|
||||
│ │ │ _orchestration_loop(): poll/process/route outputs for all stages│ │ │
|
||||
│ │ └────────────────────────────────────────────────────────────────┘ │ │
|
||||
│ └───────┬─────────────────────────┬─────────────────────────┬──────────┘ │
|
||||
└──────────┬─────────────────────────┬─────────────────────────┬─────────────┘
|
||||
ZMQ ▼ ▲ ZMQ ZMQ ▼ ▲ ZMQ ZMQ ▼ ▲ ZMQ
|
||||
┌──────────────────────┐ ┌──────────────────────┐ ┌──────────────────────┐
|
||||
│ EngineCore Proc-0 │ │ EngineCore Proc-1 │ │ EngineCore Proc-2 │
|
||||
│ (Thinker) │ │ (Talker) │ │ (Vocoder) │
|
||||
└──────────────────────┘ └──────────────────────┘ └──────────────────────┘
|
||||
```
|
||||
|
||||
|
||||
Test scripts:
|
||||
```bash
|
||||
# enter offline inference folder.
|
||||
cd examples/offline_inference/qwen2_5_omni
|
||||
python end2end.py --output-dir output_audio --query-type use_mixed_modalities
|
||||
|
||||
cd ../qwen3_omni
|
||||
python end2end.py --output-dir output_audio --query-type text --async-chunk --enable-stats
|
||||
|
||||
cd ../bagel
|
||||
python end2end.py --prompts "A cute cat"
|
||||
|
||||
cd ../text_to_image
|
||||
python text_to_image.py --prompt "a cup of coffee on the table" --output output.png
|
||||
```
|
||||
@@ -0,0 +1,978 @@
|
||||
---
|
||||
toc_depth: 4
|
||||
---
|
||||
|
||||
# Diffusion Module Architecture Design
|
||||
|
||||
The vLLM-Omni diffusion module (`vllm_omni/diffusion`) is a high-performance inference engine for diffusion models, designed with a modular architecture that separates concerns across multiple components. It provides efficient execution for non-autoregressive generation tasks such as image and video generation.
|
||||
|
||||
This document describes the architecture design of the diffusion module, including the diffusion engine, scheduler, worker, diffusion pipeline, and acceleration components.
|
||||
|
||||
<p align="center">
|
||||
<img src="https://github.com/user-attachments/assets/cde1ebea-4006-4d30-8283-15a5721a09d2" alt="vLLM-Omni Diffusion Module Components" width="80%">
|
||||
</p>
|
||||
<p align="center">
|
||||
<em> Main Components of the Diffusion Module </em>
|
||||
</p>
|
||||
|
||||
|
||||
**Table of Content:**
|
||||
|
||||
- [Architecture Overview](#architecture-overview)
|
||||
- [Diffusion Engine](#1-diffusion-engine)
|
||||
- [Scheduler](#2-scheduler)
|
||||
- [Worker](#3-worker)
|
||||
- [Diffusion Pipeline](#4-diffusion-pipeline)
|
||||
- [Acceleration Components](#5-acceleration-components)
|
||||
- [Attention Backends](#51-attention-backends)
|
||||
- [Parallel Attention](#52-parallel-attention)
|
||||
- [Cache Backends](#53-cache-backends)
|
||||
- [Parallel Strategies](#54-parallel-strategies)
|
||||
- [Data Flow](#6-data-flow)
|
||||
|
||||
---
|
||||
|
||||
## Architecture Overview
|
||||
|
||||
The diffusion module follows a **multi-process, distributed architecture** with clear separation of concerns:
|
||||
|
||||
<p align="center">
|
||||
<img src="https://github.com/user-attachments/assets/78c4d446-d238-4406-a057-eb17d2ccc8e0" alt="vLLM-Omni Diffusion Module Architecture" width="100%">
|
||||
</p>
|
||||
<p align="center">
|
||||
<em> Diffusion Architecture Overview </em>
|
||||
</p>
|
||||
|
||||
|
||||
---
|
||||
|
||||
## 1. Diffusion Engine
|
||||
|
||||
**Location**: `vllm_omni/diffusion/diffusion_engine.py`
|
||||
|
||||
### Responsibilities
|
||||
|
||||
The `DiffusionEngine` is the **orchestrator** of the diffusion inference system. It owns request admission, scheduler/executor coordination, output-stream delivery, cancellation cleanup, and startup warmup.
|
||||
|
||||
### Key Components
|
||||
|
||||
#### 1.1 Initialization
|
||||
|
||||
```python
|
||||
class DiffusionEngine:
|
||||
def __init__(
|
||||
self,
|
||||
od_config: OmniDiffusionConfig,
|
||||
scheduler: SchedulerInterface | None = None,
|
||||
):
|
||||
self.od_config = od_config
|
||||
self._init_process_hooks(od_config)
|
||||
self.execution_mode = self._resolve_execution_mode(od_config)
|
||||
self._init_executor(od_config)
|
||||
self._init_scheduler(od_config, scheduler)
|
||||
self._init_runtime_state()
|
||||
self._init_execute_fn()
|
||||
```
|
||||
|
||||
**Key Features**:
|
||||
|
||||
- **Pre/Post Processing**: Registers model-specific pre-processing and post-processing functions via registry pattern
|
||||
|
||||
- **Execution Mode Selection**: Selects one of two execution modes:
|
||||
- `REQUEST_BATCH`: complete requests are scheduled through request-level execution. `max_num_seqs=1` is the serial request path; values above `1` allow compatible requests to use fused `pipeline.forward(batch)` when the pipeline supports it.
|
||||
- `STEP_BATCH`: requests are advanced one denoising step at a time through `prepare_encode()`, `denoise_step()`, `step_scheduler()`, and `post_decode()`. `max_num_seqs` controls the maximum number of compatible active requests in one step wave.
|
||||
|
||||
- **Scheduler Selection**: Uses `StepScheduler` for step-batch execution and `RequestScheduler` for request-batch execution. Tests and custom integrations may inject a `SchedulerInterface` instance explicitly.
|
||||
|
||||
- **Executor Selection**: Builds the configured `DiffusionExecutor`, then routes scheduled work through `execute_batch` or `execute_step` according to the execution mode. `execute_batch` handles the serial `max_num_seqs=1` request path internally.
|
||||
|
||||
- **Factory-Owned Warmup**: `DiffusionEngine.make_engine()` constructs the selected engine backend and then runs startup warmup through `run_startup_warmup()`.
|
||||
|
||||
#### 1.2 Execution Modes
|
||||
|
||||
```python
|
||||
class DiffusionExecutionMode(str, Enum):
|
||||
REQUEST_BATCH = "request_batch"
|
||||
STEP_BATCH = "step_batch"
|
||||
```
|
||||
|
||||
`streaming_output=True` requires step execution. If streaming is enabled without `step_execution=True`, the engine enables step execution so the output stream can receive intermediate chunks or the final step output.
|
||||
|
||||
#### 1.3 Output Stream Lifecycle
|
||||
|
||||
```python
|
||||
request_id = engine.add_request(request)
|
||||
async for output in engine.get_output_stream(request_id):
|
||||
yield output
|
||||
```
|
||||
|
||||
**Design Decisions**:
|
||||
|
||||
- **Unified Stream Semantics**: `step_streaming()` is the public async engine path. Streaming callers forward every yielded output. Non-streaming callers consume the same stream to completion and return only the final output.
|
||||
|
||||
- **Queue Delivery**: Each admitted request gets an `asyncio.Queue[DiffusionOutput]`. The busy loop emits chunks/final outputs into that queue.
|
||||
|
||||
- **Cancellation Cleanup**: If a client disconnects or cancels stream consumption, `get_output_stream()` removes only the queue. Scheduler terminal state is still finalized when the scheduler reports the request as finished, so request IDs do not remain active indefinitely.
|
||||
|
||||
#### 1.4 Request Processing Flow
|
||||
|
||||
```python
|
||||
async def step_streaming(self, request: OmniDiffusionRequest):
|
||||
# 1. Pre-process requests
|
||||
request = self.pre_process_func(request)
|
||||
|
||||
# 2. Add request and consume the unified output stream
|
||||
async for output in self.async_add_req_and_stream_response(request):
|
||||
# 3. Post-process each delivered output
|
||||
yield self.postprocess_output(request, output, ...)
|
||||
```
|
||||
|
||||
**Flow**:
|
||||
|
||||
1. **Pre-processing**: Applies model-specific transformations
|
||||
|
||||
2. **Scheduling**: Delegates request state transitions to the selected scheduler
|
||||
|
||||
3. **Execution**: Calls the selected executor function for request, request-batch, or step execution
|
||||
|
||||
4. **Output Delivery**: Emits intermediate chunks and final outputs through the request output stream
|
||||
|
||||
5. **Post-processing**: Converts raw outputs to final format (e.g., PIL images)
|
||||
|
||||
---
|
||||
|
||||
## 2. Scheduler
|
||||
|
||||
**Location**: `vllm_omni/diffusion/sched/`
|
||||
|
||||
### Architecture
|
||||
|
||||
The scheduler is a **request-state scheduler**. It owns request lifecycle management and scheduling decisions, while execution stays in `DiffusionEngine` and the executor.
|
||||
|
||||
### Key Components
|
||||
|
||||
#### 2.1 Request State Model
|
||||
|
||||
```python
|
||||
class DiffusionRequestStatus(enum.IntEnum):
|
||||
WAITING = ...
|
||||
RUNNING = ...
|
||||
PREEMPTED = ...
|
||||
FINISHED_COMPLETED = ...
|
||||
FINISHED_ABORTED = ...
|
||||
FINISHED_ERROR = ...
|
||||
|
||||
@dataclass
|
||||
class SchedulerRequestState:
|
||||
request_id: str
|
||||
req: OmniDiffusionRequest
|
||||
sampling_params_key: StepBatchSamplingParamsKey | RequestBatchSamplingParamsKey | None = None
|
||||
status: DiffusionRequestStatus = DiffusionRequestStatus.WAITING
|
||||
error: str | None = None
|
||||
```
|
||||
|
||||
**Design Features**:
|
||||
|
||||
- **Explicit lifecycle**: Requests move through waiting, running, optional preemption, and terminal states.
|
||||
|
||||
- **Centralized error handling**: Completion, abort, and error states are all normalized in the scheduler layer.
|
||||
|
||||
#### 2.2 Shared Bookkeeping in `BaseScheduler`
|
||||
|
||||
```python
|
||||
class BaseScheduler:
|
||||
def __init__(self) -> None:
|
||||
self._request_states = {}
|
||||
self._waiting = deque()
|
||||
self._running = []
|
||||
self._finished_req_ids = set()
|
||||
self.max_num_running_reqs = 1
|
||||
|
||||
def add_request(self, request: OmniDiffusionRequest) -> str: ...
|
||||
def schedule(self) -> DiffusionSchedulerOutput: ...
|
||||
def update_from_output(
|
||||
self,
|
||||
sched_output: DiffusionSchedulerOutput,
|
||||
output: BaseRunnerOutput,
|
||||
) -> set[str]: ...
|
||||
```
|
||||
|
||||
**Responsibilities**:
|
||||
|
||||
- **Lifecycle contract**: Defines how the engine adds requests, triggers one scheduling cycle, and feeds executor results back.
|
||||
|
||||
- **Stable boundary**: `DiffusionSchedulerOutput` is the only scheduling result consumed by `DiffusionEngine`.
|
||||
|
||||
- **Pluggability**: Different scheduler policies can reuse the same engine integration path.
|
||||
|
||||
**Design Features**:
|
||||
|
||||
- **Common state storage**: Shared request maps and waiting/running sets live in the base class.
|
||||
|
||||
- **Shared cleanup logic**: Duplicate-request checks, finish handling, and state removal are centralized instead of
|
||||
duplicated in each policy.
|
||||
|
||||
- **Current constraint boundary**: `BaseScheduler` derives `max_num_running_reqs` from `max_num_seqs`. Request-mode diffusion can use that capacity for compatible independent requests when the configured pipeline declares `supports_request_batch = True`; step-wise diffusion uses the same scheduler capacity for compatible step batches.
|
||||
|
||||
#### 2.3 Current `RequestScheduler` Policy
|
||||
|
||||
```python
|
||||
class RequestScheduler(BaseScheduler):
|
||||
def schedule(self) -> DiffusionSchedulerOutput:
|
||||
# 1. keep existing RUNNING requests in the scheduling result
|
||||
# 2. pull WAITING requests while capacity remains
|
||||
# 3. move newly admitted requests into RUNNING
|
||||
```
|
||||
|
||||
**Behavior**:
|
||||
|
||||
- **FIFO request scheduling**: Waiting requests are promoted in queue order.
|
||||
|
||||
- **Compatible request admission**: `RequestScheduler` admits waiting requests while capacity remains and the request's `RequestBatchSamplingParamsKey` is compatible with the active batch. Request-mode execution keeps each logical request independent, while batch-capable pipelines receive the scheduled requests as a runner-side `DiffusionRequestBatch`.
|
||||
|
||||
- **Executor result feedback**: `update_from_output()` converts executor output into `FINISHED_COMPLETED` or `FINISHED_ERROR` and returns finished request ids.
|
||||
|
||||
#### 2.4 Engine-Driven Execution Loop
|
||||
|
||||
```python
|
||||
request_id = scheduler.add_request(request)
|
||||
while True:
|
||||
sched_output = scheduler.schedule()
|
||||
output = executor.add_req(req)
|
||||
finished_req_ids = scheduler.update_from_output(sched_output, output)
|
||||
```
|
||||
|
||||
**Design Decisions**:
|
||||
|
||||
- **Separation of concerns**: Scheduler manages state and policy; executor handles runtime execution.
|
||||
|
||||
- **No scheduler-owned IPC**: Scheduler no longer talks to workers directly.
|
||||
|
||||
- **Split concurrency model**: Request-mode diffusion can schedule a static
|
||||
batch of compatible independent requests for one full pipeline forward when
|
||||
the pipeline supports request-level batching; the step-wise path can also
|
||||
admit or remove compatible requests between denoise steps. See
|
||||
[Diffusion Continuous Batching](../../feature/diffusion_continuous_batching.md).
|
||||
|
||||
---
|
||||
|
||||
## 3. Worker
|
||||
|
||||
**Location**: `vllm_omni/diffusion/worker/gpu_worker.py`
|
||||
|
||||
### Architecture
|
||||
|
||||
Workers are **independent processes** that execute the actual model inference. Each worker runs on a dedicated GPU and participates in distributed inference.
|
||||
|
||||
### Key Components
|
||||
|
||||
#### 3.1 Worker Process Structure
|
||||
|
||||
```python
|
||||
class WorkerProc:
|
||||
def __init__(self, od_config, gpu_id, broadcast_handle):
|
||||
# Initialize ZMQ context for IPC
|
||||
self.context = zmq.Context(io_threads=2)
|
||||
|
||||
# Connect to broadcast queue (receive requests)
|
||||
self.mq = MessageQueue.create_from_handle(broadcast_handle, gpu_id)
|
||||
|
||||
# Create result queue (only rank 0)
|
||||
if gpu_id == 0:
|
||||
self.result_mq = MessageQueue(n_reader=1, ...)
|
||||
|
||||
# Initialize GPU worker
|
||||
self.worker = self._create_worker(gpu_id, od_config, ...)
|
||||
```
|
||||
|
||||
**Initialization Steps**:
|
||||
|
||||
1. **IPC Setup**: Creates ZMQ context and message queues
|
||||
|
||||
2. **Distributed Environment Setup**: Initializes PyTorch distributed communication
|
||||
|
||||
- For CUDA GPUs: Uses NCCL (fast GPU communication)
|
||||
|
||||
- For NPU: Uses HCCL (Huawei Collective Communications Library)
|
||||
|
||||
- For other devices: Uses appropriate backend (GLOO, MCCL, etc.)
|
||||
|
||||
3. **Model Loading**: Loads diffusion pipeline on assigned GPU
|
||||
|
||||
4. **Cache Setup**: Enables cache backend if configured.
|
||||
|
||||
#### 3.2 GPU Worker
|
||||
|
||||
```python
|
||||
class DiffusionWorker:
|
||||
def init_device(self):
|
||||
# Set distributed environment variables
|
||||
os.environ["RANK"] = str(rank)
|
||||
os.environ["WORLD_SIZE"] = str(world_size)
|
||||
|
||||
# Select the device and initialize distributed execution
|
||||
self.device = current_omni_platform.get_torch_device(rank)
|
||||
current_omni_platform.set_device(self.device)
|
||||
self.vllm_config = _create_diffusion_worker_vllm_config(
|
||||
self.device, self.od_config
|
||||
)
|
||||
init_distributed_environment(world_size, rank)
|
||||
parallel_config = self.od_config.parallel_config
|
||||
initialize_model_parallel(
|
||||
data_parallel_size=parallel_config.data_parallel_size,
|
||||
cfg_parallel_size=parallel_config.cfg_parallel_size,
|
||||
sequence_parallel_size=parallel_config.sequence_parallel_size,
|
||||
tensor_parallel_size=parallel_config.tensor_parallel_size,
|
||||
pipeline_parallel_size=parallel_config.pipeline_parallel_size,
|
||||
)
|
||||
|
||||
def load_model(self, load_format="default", custom_pipeline_name=None):
|
||||
self.model_runner.load_model(
|
||||
load_format=load_format,
|
||||
custom_pipeline_name=custom_pipeline_name,
|
||||
)
|
||||
|
||||
def init_lora_manager(self):
|
||||
self.lora_manager = DiffusionLoRAManager(
|
||||
pipeline=self.model_runner.pipeline,
|
||||
device=self.device,
|
||||
dtype=self.od_config.dtype,
|
||||
)
|
||||
```
|
||||
|
||||
**Key Features**:
|
||||
|
||||
- **Tensor Parallelism**: Supports multi-GPU tensor parallelism via PyTorch distributed
|
||||
|
||||
- **Model Loading**: Uses `DiffusersPipelineLoader` for efficient weight loading
|
||||
|
||||
- **Cache Integration**: Enables cache backends (TeaCache, cache-dit, etc.) transparently
|
||||
|
||||
#### 3.3 Worker Busy Loop
|
||||
|
||||
```python
|
||||
def _worker_busy_loop(self):
|
||||
while self._running:
|
||||
# 1. Receive unified message (generation request, RPC request, or shutdown)
|
||||
msg = self.mq.dequeue(timeout=1.0)
|
||||
|
||||
# 2. Route message based on type
|
||||
if isinstance(msg, dict) and msg.get("type") == "rpc":
|
||||
# Handle RPC request
|
||||
result, should_reply = self._execute_rpc(msg)
|
||||
if should_reply:
|
||||
self._return_result(result)
|
||||
|
||||
elif isinstance(msg, dict) and msg.get("type") == "shutdown":
|
||||
# Handle shutdown message
|
||||
self._running = False
|
||||
|
||||
else:
|
||||
# Handle generation request (OmniDiffusionRequest list)
|
||||
output = self.worker.execute_model(msg, self.od_config)
|
||||
self._return_result(output)
|
||||
```
|
||||
|
||||
**Execution Flow**:
|
||||
|
||||
1. **Receive**: Dequeues unified messages from shared memory queue
|
||||
|
||||
2. **Route**: Handles different message types (generation, RPC, shutdown)
|
||||
|
||||
3. **Execute**: Runs forward pass through pipeline for generation requests
|
||||
|
||||
4. **Respond**: Sends results back (rank 0 for generation, specified rank for RPC)
|
||||
|
||||
#### 3.4 Model Execution
|
||||
|
||||
```python
|
||||
@torch.inference_mode()
|
||||
def execute_model(self, reqs: list[OmniDiffusionRequest], od_config):
|
||||
req = reqs[0] # TODO: support batching
|
||||
|
||||
# Refresh cache backend if enabled
|
||||
if self.cache_backend is not None and self.cache_backend.is_enabled():
|
||||
self.cache_backend.refresh(self.pipeline, req.num_inference_steps)
|
||||
|
||||
# Set forward context for parallelism
|
||||
with set_forward_context(
|
||||
vllm_config=self.vllm_config,
|
||||
omni_diffusion_config=self.od_config
|
||||
):
|
||||
output = self.pipeline.forward(req)
|
||||
return output
|
||||
```
|
||||
|
||||
The model execution leverages multiple parallelism strategies that are transparently applied during the forward pass. The `set_forward_context()` context manager makes parallel group information available throughout the forward pass:
|
||||
|
||||
```python
|
||||
# Inside transformer layers, parallel groups are accessed via:
|
||||
from vllm_omni.diffusion.distributed.parallel_state import (
|
||||
get_sp_group, get_dp_group, get_cfg_group, get_pp_group
|
||||
)
|
||||
```
|
||||
|
||||
**Optimizations**:
|
||||
|
||||
- **Cache Refresh**: Clears cache state before each generation for clean state
|
||||
|
||||
- **Context Management**: Forward context ensures parallel groups are available during execution
|
||||
|
||||
- **Single Request**: Currently processes one request at a time (batching TODO)
|
||||
|
||||
---
|
||||
|
||||
## 4. Diffusion Pipeline
|
||||
|
||||
**Location**: `vllm_omni/diffusion/models/*/pipeline_*.py`
|
||||
|
||||
The pipeline is the **model-specific implementation** that orchestrates the diffusion process. Different models (QwenImage, Wan2.2, Z-Image) have their own pipeline implementations.
|
||||
|
||||
Most pipeline implementation are referred from `diffusers`. The multi-step diffusion loop is usually the most time-consuming part during the overall inference process, which is defined by the `diffuse` function in the pipeline class. An example is as follows:
|
||||
|
||||
```python
|
||||
def diffuse(self, ...):
|
||||
for i, t in enumerate(timesteps):
|
||||
# Forward pass for positive prompt
|
||||
transformer_kwargs = {
|
||||
"hidden_states": latents,
|
||||
"timestep": timestep / 1000,
|
||||
"encoder_hidden_states": prompt_embeds,
|
||||
}
|
||||
noise_pred = self.transformer(**transformer_kwargs)[0]
|
||||
|
||||
# Forward pass for negative prompt (CFG)
|
||||
if do_true_cfg:
|
||||
neg_transformer_kwargs = {...}
|
||||
neg_transformer_kwargs["cache_branch"] = "negative"
|
||||
neg_noise_pred = self.transformer(**neg_transformer_kwargs)[0]
|
||||
|
||||
# Combine predictions
|
||||
comb_pred = neg_noise_pred + true_cfg_scale * (noise_pred - neg_noise_pred)
|
||||
noise_pred = comb_pred * (cond_norm / noise_norm)
|
||||
|
||||
# Scheduler step
|
||||
latents = self.scheduler.step(noise_pred, t, latents)[0]
|
||||
|
||||
return latents
|
||||
```
|
||||
|
||||
**Key Features**:
|
||||
|
||||
- **CFG Support**: Handles classifier-free guidance with separate forward passes
|
||||
|
||||
- **Cache Branching**: Uses `cache_branch` parameter for cache-aware execution
|
||||
|
||||
- **True CFG**: Implements advanced CFG with norm preservation
|
||||
|
||||
To learn more about the diffusion pipeline and how to add a new diffusion pipeline, please view [Adding Diffusion Model](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/model/adding_diffusion_model)
|
||||
|
||||
---
|
||||
|
||||
## 5. Acceleration Components
|
||||
|
||||
### 5.1 Attention Backends
|
||||
|
||||
**Location**: `vllm_omni/diffusion/attention/`
|
||||
|
||||
#### Architecture
|
||||
|
||||
The attention system uses a **role-aware backend selector pattern**. Each `Attention` site declares a semantic `role` (`"self"`, `"cross"`, or a model-specific string). The selector consults the user's `AttentionConfig` to pick a backend per role and falls back to the platform default when nothing is configured.
|
||||
|
||||
#### Backend Selection
|
||||
|
||||
**Location**: `vllm_omni/diffusion/attention/selector.py`
|
||||
|
||||
```python
|
||||
class Attention(nn.Module):
|
||||
def __init__(self, num_heads, head_size, causal, softmax_scale, *,
|
||||
role="self", role_category=None, ...):
|
||||
# Resolve backend for this role from the active OmniDiffusionConfig
|
||||
config = get_current_diffusion_config_or_none()
|
||||
attention_config = config.diffusion_attention_config if config is not None else None
|
||||
attn_backend_cls, spec = get_attn_backend_for_role(
|
||||
role=role,
|
||||
head_size=head_size,
|
||||
attention_config=attention_config,
|
||||
role_category=role_category,
|
||||
)
|
||||
self.attn_impl_cls = attn_backend_cls.get_impl_cls()
|
||||
self.attention = self.attn_impl_cls(..., backend_kwargs=spec.backend_kwargs() if spec else None)
|
||||
```
|
||||
|
||||
**Available Backends**:
|
||||
|
||||
- **FlashAttention** (`FLASH_ATTN`): Memory-efficient kernel that dispatches to FA2/FA3 on CUDA, AITER on ROCm, and `mindiesd` on Ascend NPU.
|
||||
|
||||
- **SDPA** (`TORCH_SDPA`): PyTorch's scaled dot-product attention — cross-platform fallback.
|
||||
|
||||
- **SageAttention** (`SAGE_ATTN`): Quantized attention implementation from the SageAttention library.
|
||||
|
||||
These backends provide the **kernel implementations** for attention computation. For attention-level sequence parallelism strategies (Ring Attention, Ulysses), see [Parallel Attention](#52-parallel-attention).
|
||||
|
||||
#### Backend Selection Mechanism
|
||||
|
||||
Selection is driven by `OmniDiffusionConfig.diffusion_attention_config`, which carries a global `default` plus a `per_role` map of `AttentionSpec` entries. A role string is resolved in this order:
|
||||
|
||||
```python
|
||||
def get_attn_backend_for_role(role, head_size, attention_config=None, role_category=None):
|
||||
# 1. attention_config.per_role[role] — exact match
|
||||
# 2. attention_config.per_role[role_category] — category fallback
|
||||
# 3. attention_config.default — global default
|
||||
# 4. Platform default — hardware-specific
|
||||
if attention_config is not None:
|
||||
spec, source = attention_config.resolve_with_source(
|
||||
role=role, role_category=role_category,
|
||||
)
|
||||
if spec is not None:
|
||||
return load_backend(spec.backend), spec
|
||||
return current_omni_platform.get_diffusion_attn_backend_cls(...), None
|
||||
```
|
||||
|
||||
**Selection Priority**:
|
||||
|
||||
1. **Per-role override** (`--diffusion-attention-config.per_role.<role>.backend`) — finest control; matched on the layer's exact `role` string.
|
||||
|
||||
2. **Role category fallback** (`--diffusion-attention-config.per_role.<category>.backend`) — used when an exact match is missing and the layer declared `role_category` (e.g. a `"mymodel.audio_to_video"` site can fall back to `"cross"`).
|
||||
|
||||
3. **Global default** (`--diffusion-attention-backend FLASH_ATTN`, or `--diffusion-attention-config.default.backend`, or `DIFFUSION_ATTENTION_BACKEND` env var).
|
||||
|
||||
4. **Platform default** — `current_omni_platform.get_diffusion_attn_backend_cls(...)` picks the best-available kernel for the hardware.
|
||||
|
||||
A backend that needs configuration exposes it as a typed field on `AttentionSpec`; `AttentionSpec.backend_kwargs()` serializes it into the `backend_kwargs` the impl consumes. Today those fields are `skip_softmax` (TRTLLM_ATTN) and `block_sparse` (any block-sparse backend, currently RAINFUSION_ATTN), so e.g. `target_sparsity` or `sparsity` can be set without changing model code.
|
||||
|
||||
For the user-facing CLI surface, see [Diffusion Attention Backends](../../../user_guide/diffusion/attention_backends.md). For the role declaration contract on the model side, see [Adding a Diffusion Model](../../../contributing/model/adding_diffusion_model.md).
|
||||
|
||||
**Backend Availability**:
|
||||
|
||||
- **SDPA**: Always available (PyTorch built-in)
|
||||
|
||||
- **FlashAttention**: Requires `flash-attn` on CUDA, `aiter` on ROCm, or `mindiesd` on Ascend NPU
|
||||
|
||||
- **SageAttention**: Requires `sage-attention` package (from THU-ML GitHub)
|
||||
|
||||
#### Attention Backend Registry
|
||||
|
||||
**Location**: `vllm_omni/diffusion/attention/selector.py` + `vllm_omni/platforms/<device>/platform.py`
|
||||
|
||||
Backend resolution is delegated to the active platform: `current_omni_platform.get_diffusion_attn_backend_cls(selected_backend, head_size)` returns a fully-qualified class path that `_load_backend_cls` imports lazily, with the result cached by `(backend_name, head_size)`.
|
||||
|
||||
```python
|
||||
@cache
|
||||
def _cached_get_backend_cls(backend_name: str | None, head_size: int) -> type[AttentionBackend]:
|
||||
backend_cls_path = current_omni_platform.get_diffusion_attn_backend_cls(
|
||||
selected_backend=backend_name,
|
||||
head_size=head_size,
|
||||
)
|
||||
return _load_backend_cls(backend_cls_path)
|
||||
```
|
||||
|
||||
Each platform (`cuda`, `rocm`, `xpu`, `musa`, `npu`) maps backend names like `"FLASH_ATTN"`, `"TORCH_SDPA"`, `"SAGE_ATTN"` to the right backend class path for that hardware, including head-size compatibility checks. Passing `selected_backend=None` lets the platform pick its own default.
|
||||
|
||||
#### Attention Backend Integration
|
||||
|
||||
The `Attention` layer integrates backends through a unified interface. Here's how **FlashAttentionBackend** is integrated as an example:
|
||||
|
||||
```python
|
||||
# attention/backends/flash_attn.py
|
||||
|
||||
class FlashAttentionBackend(AttentionBackend):
|
||||
@staticmethod
|
||||
def get_name() -> str:
|
||||
return "FLASH_ATTN"
|
||||
|
||||
@staticmethod
|
||||
def get_impl_cls() -> type["FlashAttentionImpl"]:
|
||||
return FlashAttentionImpl
|
||||
|
||||
@staticmethod
|
||||
def get_supported_head_sizes() -> list[int]:
|
||||
return [64, 96, 128, 192, 256] # FlashAttention supports these head sizes
|
||||
|
||||
|
||||
class FlashAttentionImpl(AttentionImpl):
|
||||
def __init__(self, num_heads, head_size, softmax_scale, causal, ...):
|
||||
self.num_heads = num_heads
|
||||
self.causal = causal
|
||||
self.softmax_scale = softmax_scale
|
||||
|
||||
def forward(self, query, key, value, attn_metadata=None):
|
||||
# Call FlashAttention kernel
|
||||
out = flash_attn_func(
|
||||
query, key, value,
|
||||
causal=self.causal,
|
||||
softmax_scale=self.softmax_scale,
|
||||
)
|
||||
return out
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### 5.2 Parallel Attention
|
||||
|
||||
**Location**: `vllm_omni/diffusion/attention/parallel/`
|
||||
|
||||
#### Architecture
|
||||
|
||||
Parallel attention strategies implement **Sequence Parallelism (SP) at the attention layer level**. These strategies distribute attention computation across multiple GPUs by splitting the sequence dimension, using different communication patterns. They work **on top of** AttentionBackend implementations (FlashAttention, SDPA, etc.), handling the parallelization/communication while the backends handle the actual attention computation.
|
||||
|
||||
**Key Distinction**: Unlike AttentionBackend (which provides kernel implementations), ParallelAttentionStrategy provides communication patterns for multi-GPU attention parallelism. These strategies implement the `ParallelAttentionStrategy` interface and use AttentionBackend implementations internally.
|
||||
|
||||
Both Ring Attention and Ulysses are forms of Sequence Parallelism (SP) that:
|
||||
|
||||
- Split the sequence dimension across GPUs
|
||||
|
||||
- Contribute to `sequence_parallel_size` (via `ring_degree` and `ulysses_degree`)
|
||||
|
||||
- Work at the attention layer level (not model/pipeline level)
|
||||
|
||||
#### Ulysses Sequence Parallelism (USP)
|
||||
|
||||
**Location**: `vllm_omni/diffusion/attention/parallel/ulysses.py`
|
||||
|
||||
USP is a sequence-parallel attention strategy that splits attention computation across multiple GPUs by distributing both the sequence dimension and attention heads. It uses **all-to-all communication** to efficiently parallelize attention for very long sequences. Specifically, it uses **all-to-all** collective operations to redistribute Q/K/V tensors before attention computation and gather results afterward.
|
||||
|
||||
Ulysses splits attention computation in two dimensions:
|
||||
|
||||
1. **Sequence Dimension**: Splits the sequence length across GPUs
|
||||
|
||||
2. **Head Dimension**: Splits attention heads across GPUs
|
||||
|
||||
**Configuration**: `ulysses_degree` contributes to `sequence_parallel_size`
|
||||
|
||||
#### Ring Sequence Parallelism
|
||||
|
||||
**Location**: `vllm_omni/diffusion/attention/parallel/ring.py`
|
||||
|
||||
Ring Attention is a **parallel attention strategy** that implements sequence parallelism using ring-based point-to-point (P2P) communication. Unlike attention backends that provide the attention kernel implementation, Ring Attention is a **communication pattern** that works on top of attention backends (FlashAttention or SDPA).
|
||||
|
||||
Ring Attention splits sequence dimension across GPUs in a ring topology, implemented via the `ParallelAttentionStrategy` interface, instead of `AttentionBackend`. P2P ring communication is applied to circulate Key/Value blocks across GPUs. Internally, `ring_flash_attn_func` or `ring_pytorch_attn_func` is used depending on available backends.
|
||||
|
||||
**Architecture**:
|
||||
```python
|
||||
class RingParallelAttention:
|
||||
"""Ring sequence-parallel strategy."""
|
||||
|
||||
def run_attention(self, query, key, value, attn_metadata, ...):
|
||||
# Selects underlying attention kernel (FlashAttention or SDPA)
|
||||
if backend_pref == "sdpa":
|
||||
return ring_pytorch_attn_func(...) # Uses SDPA kernel
|
||||
else:
|
||||
return ring_flash_attn_func(...) # Uses FlashAttention kernel
|
||||
```
|
||||
|
||||
**Integration**:
|
||||
|
||||
- Ring Attention is activated when `ring_degree > 1` in parallel config
|
||||
|
||||
- It's selected by `build_parallel_attention_strategy()` in the attention layer
|
||||
|
||||
- The `Attention` layer routes to `_run_ring_attention()` when Ring is enabled
|
||||
|
||||
- Works alongside attention backends: Ring handles communication, backends handle computation
|
||||
|
||||
**Configuration**: `ring_degree` contributes to `sequence_parallel_size`
|
||||
|
||||
#### Relationship with AttentionBackend
|
||||
|
||||
Parallel attention strategies (Ring, Ulysses) work **on top of** AttentionBackend implementations:
|
||||
|
||||
- They use AttentionBackend for the actual attention computation (FlashAttention, SDPA, etc.)
|
||||
|
||||
- They handle the multi-GPU communication/parallelization layer
|
||||
|
||||
- They implement `ParallelAttentionStrategy` interface (not `AttentionBackend`)
|
||||
|
||||
For general parallelism strategies (Data Parallelism, Tensor Parallelism, Pipeline Parallelism), see [Parallel Strategies](#54-parallel-strategies).
|
||||
|
||||
---
|
||||
|
||||
### 5.3 Cache Backends
|
||||
|
||||
**Location**: `vllm_omni/diffusion/cache/`
|
||||
|
||||
#### Architecture
|
||||
|
||||
Cache backends provide a **unified interface** for applying different caching strategies to accelerate diffusion inference. The system supports multiple backends (TeaCache, cache-dit) with a consistent API for enabling and refreshing cache state.
|
||||
|
||||
#### Cache Backend Interface
|
||||
|
||||
```python
|
||||
class CacheBackend(ABC):
|
||||
def __init__(self, config: DiffusionCacheConfig):
|
||||
self.config = config
|
||||
self.enabled = False
|
||||
|
||||
@abstractmethod
|
||||
def enable(self, pipeline: Any) -> None:
|
||||
"""Enable cache on the pipeline."""
|
||||
raise NotImplementedError
|
||||
|
||||
@abstractmethod
|
||||
def refresh(self, pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None:
|
||||
"""Refresh cache state for new generation."""
|
||||
raise NotImplementedError
|
||||
|
||||
def is_enabled(self) -> bool:
|
||||
"""Check if cache is enabled."""
|
||||
return self.enabled
|
||||
```
|
||||
|
||||
**Design Pattern**:
|
||||
|
||||
- **Abstract Base Class**: Defines contract for all cache backends
|
||||
|
||||
- **Pipeline-based**: Works with pipeline instances (not just transformers)
|
||||
|
||||
- **State Management**: Provides refresh mechanism for clean state between generations
|
||||
|
||||
#### Available Backends
|
||||
|
||||
**1. TeaCache Backend**
|
||||
|
||||
**Location**: `vllm_omni/diffusion/cache/teacache/backend.py`
|
||||
|
||||
```python
|
||||
class TeaCacheBackend(CacheBackend):
|
||||
def enable(self, pipeline: Any):
|
||||
# Extract transformer from pipeline
|
||||
transformer = pipeline.transformer
|
||||
transformer_type = transformer.__class__.__name__
|
||||
|
||||
# Create TeaCacheConfig from DiffusionCacheConfig
|
||||
teacache_config = TeaCacheConfig(
|
||||
transformer_type=transformer_type,
|
||||
rel_l1_thresh=self.config.rel_l1_thresh,
|
||||
coefficients=self.config.coefficients,
|
||||
)
|
||||
|
||||
# Apply hooks to transformer
|
||||
apply_teacache_hook(transformer, teacache_config)
|
||||
self.enabled = True
|
||||
|
||||
def refresh(self, pipeline: Any, num_inference_steps: int, verbose: bool = True):
|
||||
transformer = pipeline.transformer
|
||||
if hasattr(transformer, "_hook_registry"):
|
||||
transformer._hook_registry.reset_hook(TeaCacheHook._HOOK_NAME)
|
||||
```
|
||||
|
||||
**TeaCache Features**:
|
||||
|
||||
- **Timestep-aware**: Caches based on timestep embedding similarity
|
||||
|
||||
- **Adaptive**: Dynamically decides when to reuse cached computations
|
||||
|
||||
- **CFG-aware**: Handles positive/negative branches separately
|
||||
|
||||
- **Custom Hook System**: Uses a custom forward interception mechanism (via `HookRegistry`) that wraps the module's `forward` method, allowing transparent integration without modifying model code
|
||||
|
||||
**2. Cache-DiT Backend**
|
||||
|
||||
**Location**: `vllm_omni/diffusion/cache/cachedit/`
|
||||
|
||||
```python
|
||||
class CacheDiTBackend(CacheBackend):
|
||||
def enable(self, pipeline: Any):
|
||||
# Uses cache-dit library for acceleration
|
||||
# Supports DBCache, SCM (Step Computation Masking), TaylorSeer
|
||||
# Works with single and dual-transformer architectures
|
||||
...
|
||||
self.enabled = True
|
||||
|
||||
def refresh(self, pipeline: Any, num_inference_steps: int, verbose: bool = True):
|
||||
# Updates cache context with new num_inference_steps
|
||||
...
|
||||
```
|
||||
|
||||
**Cache-DiT Features**:
|
||||
|
||||
- **DBCache**: Dynamic block caching with configurable compute blocks
|
||||
|
||||
- **SCM**: Step Computation Masking for additional speedup
|
||||
|
||||
- **TaylorSeer**: Advanced calibration for cache accuracy
|
||||
|
||||
- **Dual-transformer Support**: Handles models like Wan2.2 with two transformers
|
||||
|
||||
#### Cache Backend Selector
|
||||
|
||||
**Location**: `vllm_omni/diffusion/cache/selector.py`
|
||||
|
||||
```python
|
||||
def get_cache_backend(
|
||||
cache_backend: str | None,
|
||||
cache_config: dict | DiffusionCacheConfig
|
||||
) -> CacheBackend | None:
|
||||
"""Get cache backend instance based on cache_backend string.
|
||||
|
||||
Args:
|
||||
cache_backend: Cache backend name ("cache_dit", "tea_cache", or None)
|
||||
cache_config: Cache configuration (dict or DiffusionCacheConfig)
|
||||
|
||||
Returns:
|
||||
Cache backend instance or None if cache_backend is "none"
|
||||
"""
|
||||
if cache_backend is None or cache_backend == "none":
|
||||
return None
|
||||
|
||||
if isinstance(cache_config, dict):
|
||||
cache_config = DiffusionCacheConfig.from_dict(cache_config)
|
||||
|
||||
if cache_backend == "cache_dit":
|
||||
return CacheDiTBackend(cache_config)
|
||||
elif cache_backend == "tea_cache":
|
||||
return TeaCacheBackend(cache_config)
|
||||
else:
|
||||
raise ValueError(f"Unsupported cache backend: {cache_backend}")
|
||||
```
|
||||
|
||||
**Usage Flow**:
|
||||
|
||||
1. **Selection**: `get_cache_backend()` returns appropriate backend instance
|
||||
|
||||
2. **Enable**: `backend.enable(pipeline)` called during worker initialization
|
||||
|
||||
3. **Refresh**: `backend.refresh(pipeline, num_inference_steps)` called before each generation
|
||||
|
||||
4. **Check**: `backend.is_enabled()` verifies cache is active
|
||||
|
||||
### 5.4 Parallel Strategies
|
||||
|
||||
**Location**: `vllm_omni/diffusion/distributed/parallel_state.py`
|
||||
|
||||
#### Parallelism Types
|
||||
|
||||
The system supports multiple orthogonal parallelism strategies:
|
||||
|
||||
**Sequence Parallelism (SP)**
|
||||
|
||||
- **Purpose**: Split sequence dimension across GPUs
|
||||
|
||||
- **Attention-level SP**: Ring Attention and Ulysses (USP) implement SP at the attention layer level
|
||||
|
||||
- See [Parallel Attention](#52-parallel-attention) for details
|
||||
|
||||
- Configuration: `ulysses_degree` × `ring_degree` = `sequence_parallel_size`
|
||||
|
||||
- **Use Case**: Very long sequences (e.g., high-resolution images)
|
||||
|
||||
**Data Parallelism (DP)**
|
||||
|
||||
- **Purpose**: Replicate model across GPUs, split batch
|
||||
|
||||
- **Use Case**: Batch processing, throughput optimization
|
||||
|
||||
**Tensor Parallelism (TP)** (Experimental)
|
||||
|
||||
- **Purpose**: Split model weights across GPUs
|
||||
|
||||
- **Implementation**: Uses vLLM's tensor parallel groups
|
||||
|
||||
- **Use Case**: Large models that don't fit on single GPU
|
||||
|
||||
**CFG Parallelism** (under development)
|
||||
|
||||
- **Purpose**: Parallelize Classifier-Free Guidance (positive/negative prompts)
|
||||
|
||||
- **Infrastructure**: CFG parallel groups are initialized and available via `get_cfg_group()`
|
||||
|
||||
#### Parallel Group Management
|
||||
|
||||
```python
|
||||
def initialize_model_parallel(
|
||||
data_parallel_size: int = 1,
|
||||
cfg_parallel_size: int = 1,
|
||||
sequence_parallel_size: int | None = None,
|
||||
ulysses_degree: int = 1,
|
||||
ring_degree: int = 1,
|
||||
tensor_parallel_size: int = 1,
|
||||
pipeline_parallel_size: int = 1,
|
||||
vae_parallel_size: int = 0,
|
||||
):
|
||||
# Generate orthogonal parallel groups
|
||||
rank_generator = RankGenerator(
|
||||
tensor_parallel_size,
|
||||
sequence_parallel_size,
|
||||
pipeline_parallel_size,
|
||||
cfg_parallel_size,
|
||||
data_parallel_size,
|
||||
"tp-sp-pp-cfg-dp",
|
||||
)
|
||||
|
||||
# Initialize each parallel group
|
||||
_DP = init_model_parallel_group(rank_generator.get_ranks("dp"), ...)
|
||||
_CFG = init_model_parallel_group(rank_generator.get_ranks("cfg"), ...)
|
||||
_SP = init_model_parallel_group(rank_generator.get_ranks("sp"), ...)
|
||||
_PP = init_model_parallel_group(rank_generator.get_ranks("pp"), ...)
|
||||
_TP = init_model_parallel_group(rank_generator.get_ranks("tp"), ...)
|
||||
```
|
||||
|
||||
**Rank Order**: `tp-sp-pp-cfg-dp` (tensor → sequence → pipeline → cfg → data)
|
||||
|
||||
**Note**: For attention-level Sequence Parallelism implementations (Ring Attention and Ulysses), see [Parallel Attention](#52-parallel-attention). This section covers higher-level parallelism strategies.
|
||||
|
||||
|
||||
---
|
||||
|
||||
## 6. Data Flow
|
||||
|
||||
### Complete Request Flow
|
||||
|
||||
<p align="center">
|
||||
<img src="https://github.com/user-attachments/assets/6e093a0d-29c0-4efd-85d5-747002bd2fed" alt="vLLM-Omni Diffusion Module Components" width="100%">
|
||||
</p>
|
||||
<p align="center">
|
||||
<em> End-to-end Data Flow in the vLLM-Omni Diffusion Module </em>
|
||||
</p>
|
||||
|
||||
|
||||
```
|
||||
1. User Request
|
||||
└─> Omni.generate(prompt)
|
||||
└─> Prepare OmniDiffusionRequest
|
||||
└─> DiffusionEngine.step_streaming(request)
|
||||
|
||||
2. Pre-processing
|
||||
└─> pre_process_func(request)
|
||||
└─> Model-specific transformations
|
||||
|
||||
3. Scheduling
|
||||
└─> scheduler.add_request(request)
|
||||
└─> scheduler.schedule()
|
||||
└─> DiffusionEngine selects execute_batch / execute_step
|
||||
|
||||
4. Worker Execution
|
||||
└─> WorkerProc._worker_busy_loop()
|
||||
└─> DiffusionWorker.execute_model(req)
|
||||
└─> DiffusionModelRunner selects execution mode
|
||||
├─> REQUEST_BATCH
|
||||
│ └─> Pipeline.forward(req_or_batch)
|
||||
│ ├─> encode_prompt()
|
||||
│ ├─> prepare_latents()
|
||||
│ ├─> diffuse() [loop]
|
||||
│ │ ├─> transformer.forward() [with cache backend hooks]
|
||||
│ │ └─> scheduler.step()
|
||||
│ └─> vae.decode()
|
||||
└─> STEP
|
||||
├─> pipeline.prepare_encode(state) [new requests]
|
||||
├─> pipeline.denoise_step(input_batch)
|
||||
├─> pipeline.step_scheduler(state, noise_pred)
|
||||
└─> pipeline.post_decode(state) [chunk boundary or final step]
|
||||
|
||||
5. Result Collection
|
||||
└─> Executor returns RunnerOutput / BatchRunnerOutput
|
||||
└─> scheduler.update_from_output(...)
|
||||
├─> DiffusionEngine emits chunk/final output to request queue
|
||||
└─> DiffusionEngine finalizes terminal scheduler state
|
||||
|
||||
6. Post-processing
|
||||
└─> post_process_func(output)
|
||||
└─> Convert to PIL images / final format
|
||||
|
||||
7. Stream Completion / Cancellation
|
||||
├─> Streaming callers receive every yielded output
|
||||
├─> Non-streaming callers drain the stream and return the final output
|
||||
└─> Cancelled streams drop their output queue but still allow scheduler state cleanup
|
||||
```
|
||||
|
||||
---
|
||||
@@ -0,0 +1 @@
|
||||
Architecture design of the entrypoint (update soon)
|
||||
@@ -2,26 +2,56 @@
|
||||
title: Engine Orchestration
|
||||
kind: module
|
||||
status: draft
|
||||
architecture_state: current-plus-in-flight
|
||||
owners:
|
||||
- "@tzhouam"
|
||||
- "@yinpeiqi"
|
||||
- "@fake0fan"
|
||||
document_stewards:
|
||||
- "@hsliuustc0106"
|
||||
- "@Gaohan123"
|
||||
- "@david6666666"
|
||||
required_reviewers:
|
||||
- "@yinpeiqi"
|
||||
- "@Sy0307"
|
||||
- "@yenuo26"
|
||||
- "@NickCao"
|
||||
primary_code_paths:
|
||||
- vllm_omni/engine/**
|
||||
- vllm_omni/engine/async_omni_engine.py
|
||||
- vllm_omni/engine/async_engine_utils.py
|
||||
- vllm_omni/engine/orchestrator.py
|
||||
- vllm_omni/engine/cfg_companion_tracker.py
|
||||
- vllm_omni/engine/rpc_result_router.py
|
||||
related_code_paths:
|
||||
- vllm_omni/engine/messages.py
|
||||
- vllm_omni/engine/orchestrator_monitor.py
|
||||
- vllm_omni/engine/stage_client.py
|
||||
- vllm_omni/engine/stage_pool.py
|
||||
- vllm_omni/outputs/output_processor.py
|
||||
- vllm_omni/distributed/omni_coordinator/**
|
||||
- vllm_omni/distributed/ray_utils/**
|
||||
depends_on:
|
||||
- vllm_omni_config.md
|
||||
- input_output_modality_contracts.md
|
||||
- error_contracts.md
|
||||
- stage_runtime.md
|
||||
- omni_connector.md
|
||||
validation_paths:
|
||||
- tests/engine/**
|
||||
- tests/distributed/omni_coordinator/**
|
||||
- tests/engine/test_async_omni_engine_input.py
|
||||
- tests/engine/test_async_omni_engine_outputs.py
|
||||
- tests/engine/test_async_omni_engine_abort.py
|
||||
- tests/engine/test_async_omni_engine_stage_init.py
|
||||
- tests/engine/test_orchestrator.py
|
||||
- tests/engine/test_orchestrator_error_handling.py
|
||||
- tests/engine/test_orchestrator_stage_input_bridge.py
|
||||
- tests/engine/test_cfg_companion_lifecycle.py
|
||||
- tests/engine/test_rpc_result_router.py
|
||||
- tests/e2e/features/fullduplex/engine/**
|
||||
upstream_refs:
|
||||
- vllm.v1.engine
|
||||
last_reviewed: 2026-07-16
|
||||
- vllm.v1.engine.EngineCoreRequest
|
||||
- vllm.v1.engine.EngineCoreOutputs
|
||||
- vllm.v1.engine.exceptions.EngineDeadError
|
||||
- vllm.outputs.RequestOutput
|
||||
invariant_namespace: ORCH-INV
|
||||
last_reviewed: 2026-08-07
|
||||
last_verified_commit: 3d7fc3b9ba3cac88d579d4dc35b78b0b641675fc
|
||||
---
|
||||
|
||||
# Engine orchestration
|
||||
@@ -29,24 +59,64 @@ last_reviewed: 2026-07-16
|
||||
Engine orchestration coordinates configured AR and diffusion stages while
|
||||
keeping public clients independent from stage processes and transports.
|
||||
|
||||
## Contract status
|
||||
|
||||
This document is a draft description of current behavior. It also identifies
|
||||
the boundary affected by the in-flight stage client/process refactor in
|
||||
[#5441](https://github.com/vllm-project/vllm-omni/pull/5441). Names and
|
||||
responsibilities proposed only by that PR are not current contracts.
|
||||
|
||||
## Ownership boundary
|
||||
|
||||
This document owns `AsyncOmniEngine`, `Orchestrator`, request-state creation,
|
||||
cross-stage routing, output ordering, companion tracking, control/RPC
|
||||
correlation, cancellation propagation, and terminal-state convergence.
|
||||
|
||||
It does not own stage placement, replica selection, stage-process startup,
|
||||
payload schemas, public protocol rendering, connector transport, or semantic
|
||||
error classification. Those responsibilities belong to the stage runtime,
|
||||
I/O, entrypoint, connector, and error contracts.
|
||||
|
||||
## Candidate invariants
|
||||
|
||||
These identifiers are proposals while the document is `draft`.
|
||||
|
||||
### ORCH-INV-001: The orchestrator owns cross-stage routing
|
||||
|
||||
**Rule:** Entrypoints and stage clients MUST NOT independently forward a
|
||||
request to a downstream stage.
|
||||
request to a downstream logical stage.
|
||||
|
||||
### ORCH-INV-002: Stage clients do not own routing policy
|
||||
|
||||
**Rule:** Stage clients MUST implement communication and lifecycle operations
|
||||
without selecting the next logical stage.
|
||||
|
||||
### ORCH-INV-003: Terminal state is monotonic
|
||||
### ORCH-INV-100: Terminal state is monotonic
|
||||
|
||||
**Rule:** Once a request reaches a terminal state, orchestration MUST NOT
|
||||
forward new work for that request.
|
||||
|
||||
## Invariant namespace
|
||||
|
||||
`ORCH-INV` reserves `001-099` for dependency direction, `100-199` for
|
||||
request state and ordering, `200-299` for failure/cancellation/cleanup, and
|
||||
`300-399` for extension points, shutdown ordering, and upstream alignment.
|
||||
Numbers become append-only after normative promotion.
|
||||
|
||||
## Safe-change guide
|
||||
|
||||
Test routing, output ordering, cancellation, failure propagation, shutdown, and
|
||||
representative multi-stage execution.
|
||||
Test routing, output ordering, queue correlation, cancellation, failure
|
||||
propagation, shutdown ordering, and representative multi-stage execution.
|
||||
Changes that cross into stage runtime or public error behavior require review
|
||||
from that contract's owners.
|
||||
|
||||
## Promotion gate
|
||||
|
||||
- Reconcile current names and boundaries after #5441 merges, closes, or is
|
||||
replaced.
|
||||
- Demonstrate one terminal outcome per request across success, abort,
|
||||
request-scoped failure, and fatal engine failure.
|
||||
- Verify queue correlation and orchestrator-before-runtime shutdown ordering
|
||||
in the cited validation paths.
|
||||
- Obtain approval from a technical owner and an independent validation
|
||||
reviewer.
|
||||
@@ -1,51 +1,134 @@
|
||||
---
|
||||
title: Entrypoints
|
||||
title: Entrypoints and Serving Boundaries
|
||||
kind: module
|
||||
status: draft
|
||||
architecture_state: current-plus-in-flight
|
||||
owners:
|
||||
- "@tzhouam"
|
||||
- "@yinpeiqi"
|
||||
- "@fake0fan"
|
||||
- "@alex-jw-brooks"
|
||||
- "@linyueqian"
|
||||
document_stewards:
|
||||
- "@hsliuustc0106"
|
||||
- "@Gaohan123"
|
||||
- "@david6666666"
|
||||
required_reviewers:
|
||||
- "@tzhouam"
|
||||
- "@herotai214"
|
||||
- "@yenuo26"
|
||||
- "@NickCao"
|
||||
primary_code_paths:
|
||||
- vllm_omni/entrypoints/**
|
||||
- vllm_omni/entrypoints/omni.py
|
||||
- vllm_omni/entrypoints/async_omni.py
|
||||
- vllm_omni/entrypoints/omni_base.py
|
||||
- vllm_omni/entrypoints/cli/**
|
||||
- vllm_omni/entrypoints/openai/**
|
||||
- vllm_omni/entrypoints/openpi/**
|
||||
- vllm_omni/entrypoints/client_request_state.py
|
||||
- vllm_omni/entrypoints/stage_utils.py
|
||||
- vllm_omni/entrypoints/utils.py
|
||||
primary_path_exceptions:
|
||||
- path: vllm_omni/entrypoints/openai/errors.py
|
||||
owner: error_contracts.md
|
||||
related_code_paths:
|
||||
- vllm_omni/reasoning/**
|
||||
- vllm_omni/errors.py
|
||||
- vllm_omni/inputs/**
|
||||
- vllm_omni/outputs/**
|
||||
- vllm_omni/engine/async_omni_engine.py
|
||||
- vllm_omni/config/**
|
||||
- vllm_omni/deploy/**
|
||||
depends_on:
|
||||
- vllm_omni_config.md
|
||||
- input_output_modality_contracts.md
|
||||
- error_contracts.md
|
||||
- engine_orchestration.md
|
||||
validation_paths:
|
||||
- tests/entrypoints/**
|
||||
- tests/reasoning/**
|
||||
- tests/entrypoints/test_omni_entrypoints.py
|
||||
- tests/entrypoints/test_async_omni.py
|
||||
- tests/entrypoints/test_async_omni_duplex.py
|
||||
- tests/entrypoints/test_serve.py
|
||||
- tests/entrypoints/test_stream_finish_reason.py
|
||||
- tests/entrypoints/openai/**
|
||||
- tests/entrypoints/openai_api/**
|
||||
- tests/e2e/online_serving/**
|
||||
upstream_refs:
|
||||
- vllm.entrypoints
|
||||
last_reviewed: 2026-07-16
|
||||
- vllm.engine.protocol.EngineClient
|
||||
- vllm.entrypoints.openai/**
|
||||
- vllm.entrypoints.serve/**
|
||||
- vllm.renderers/**
|
||||
- vllm.v1.engine.exceptions/**
|
||||
invariant_namespace: ENTRY-INV
|
||||
last_reviewed: 2026-08-07
|
||||
last_verified_commit: 3d7fc3b9ba3cac88d579d4dc35b78b0b641675fc
|
||||
---
|
||||
|
||||
# Entrypoints
|
||||
# Entrypoints and serving boundaries
|
||||
|
||||
Entrypoints translate offline, CLI, and serving requests into stable engine
|
||||
operations and translate engine outputs into public responses.
|
||||
|
||||
## Contract status
|
||||
|
||||
This document describes current entrypoint responsibilities plus the boundary
|
||||
under review in the open roadmap
|
||||
[#5227](https://github.com/vllm-project/vllm-omni/issues/5227) and helper-move
|
||||
PR [#5453](https://github.com/vllm-project/vllm-omni/pull/5453). In-flight
|
||||
helper locations are not treated as current paths.
|
||||
|
||||
## Ownership boundary
|
||||
|
||||
This document owns offline API semantics, CLI and serve composition,
|
||||
supported OpenAI-compatible routes, request validation and normalization,
|
||||
response conversion, streaming/session behavior, and engine handoff.
|
||||
|
||||
It does not own configuration precedence, cross-stage routing, stage
|
||||
placement, payload implementation, or semantic error classification.
|
||||
`entrypoints/openai/errors.py` is an explicit primary-path exception owned by
|
||||
`error_contracts.md`.
|
||||
|
||||
## Candidate invariants
|
||||
|
||||
These identifiers are proposals while the document is `draft`.
|
||||
|
||||
### ENTRY-INV-001: Entrypoints adapt but do not orchestrate
|
||||
|
||||
**Rule:** Entrypoints MUST NOT implement cross-stage routing or stage lifecycle
|
||||
policy.
|
||||
|
||||
### ENTRY-INV-002: Public requests are normalized once
|
||||
### ENTRY-INV-100: Public requests are normalized once
|
||||
|
||||
**Rule:** Public protocol values MUST be validated and converted to internal
|
||||
request contracts before engine submission.
|
||||
|
||||
### ENTRY-INV-003: Streaming preserves request identity
|
||||
### ENTRY-INV-101: Streaming preserves request identity
|
||||
|
||||
**Rule:** Every streamed response MUST remain associated with the request and
|
||||
output modality that produced it.
|
||||
|
||||
The reviewer-proposed rule that model-specific behavior should stay behind a
|
||||
common adapter or processor abstraction remains an unnumbered candidate until
|
||||
the entrypoint refactor makes that boundary enforceable.
|
||||
|
||||
## Invariant namespace
|
||||
|
||||
`ENTRY-INV` reserves `001-099` for boundary and dependency direction,
|
||||
`100-199` for normalization, sessions, and streaming, `200-299` for
|
||||
disconnect, cancellation, rendering, and cleanup, and `300-399` for upstream
|
||||
route and renderer compatibility. Numbers become append-only after normative
|
||||
promotion.
|
||||
|
||||
## Safe-change guide
|
||||
|
||||
Test request validation, protocol conversion, streaming, cancellation, and
|
||||
error mapping for each affected entrypoint.
|
||||
Test request validation, protocol conversion, model-adapter routing,
|
||||
streaming, session identity, disconnect/cancellation, and error mapping for
|
||||
each affected offline or serving entrypoint.
|
||||
|
||||
## Promotion gate
|
||||
|
||||
- Reconcile the page after #5227's P0 ownership work and #5453's final
|
||||
disposition.
|
||||
- Publish a supported route/transport matrix rather than claiming blanket
|
||||
OpenAI parity.
|
||||
- Verify one normalization and handoff path per route family and one terminal
|
||||
outcome per streaming transport.
|
||||
- Promote the `api_server.py` composition-root rule only after helper moves and
|
||||
import boundaries enforce it.
|
||||
- Obtain approval from a technical owner and an independent validation
|
||||
reviewer.
|
||||
@@ -0,0 +1,109 @@
|
||||
---
|
||||
title: Error Classification, Propagation, and Rendering
|
||||
kind: module
|
||||
status: draft
|
||||
architecture_state: current-plus-rfc
|
||||
owners:
|
||||
- "@alex-jw-brooks"
|
||||
- "@NickCao"
|
||||
document_stewards:
|
||||
- "@hsliuustc0106"
|
||||
- "@Gaohan123"
|
||||
- "@david6666666"
|
||||
required_reviewers:
|
||||
- "@yinpeiqi"
|
||||
- "@tzhouam"
|
||||
- "@yenuo26"
|
||||
primary_code_paths:
|
||||
- vllm_omni/errors.py
|
||||
- vllm_omni/entrypoints/openai/errors.py
|
||||
related_code_paths:
|
||||
- vllm_omni/engine/messages.py
|
||||
- vllm_omni/engine/orchestrator.py
|
||||
- vllm_omni/entrypoints/omni_base.py
|
||||
- vllm_omni/entrypoints/async_omni.py
|
||||
- vllm_omni/entrypoints/openai/api_server.py
|
||||
- vllm_omni/entrypoints/openai/serving_chat.py
|
||||
- vllm_omni/entrypoints/openai/serving_speech.py
|
||||
- vllm_omni/entrypoints/openai/serving_audio_generate.py
|
||||
- vllm_omni/entrypoints/openai/protocol/images.py
|
||||
- vllm_omni/entrypoints/openai/protocol/videos.py
|
||||
depends_on: []
|
||||
validation_paths:
|
||||
- tests/engine/test_orchestrator_error_handling.py
|
||||
- tests/engine/test_async_omni_engine_outputs.py
|
||||
- tests/entrypoints/test_omni_entrypoints.py
|
||||
- tests/entrypoints/test_stream_finish_reason.py
|
||||
- tests/entrypoints/openai_api/test_image_server.py
|
||||
- tests/entrypoints/openai_api/test_video_server.py
|
||||
- tests/entrypoints/openai_api/test_serving_audio_generate.py
|
||||
- tests/entrypoints/openai_api/test_serving_speech.py
|
||||
- tests/metrics/test_prometheus.py
|
||||
upstream_refs:
|
||||
- vllm.v1.engine.exceptions.EngineDeadError
|
||||
- vllm.v1.engine.exceptions.EngineGenerateError
|
||||
- vllm.entrypoints.serve.utils.error_response
|
||||
- vLLM RFC #48227
|
||||
- vLLM PR #49665
|
||||
invariant_namespace: ERR-INV
|
||||
last_reviewed: 2026-08-07
|
||||
last_verified_commit: 3d7fc3b9ba3cac88d579d4dc35b78b0b641675fc
|
||||
---
|
||||
|
||||
# Error classification, propagation, and rendering
|
||||
|
||||
This contract separates the meaning of a failure from the transport used to
|
||||
carry it and the public protocol used to render it.
|
||||
|
||||
## Contract status
|
||||
|
||||
Current behavior is distributed across engine, message, and entrypoint code.
|
||||
The semantic hierarchy proposed by
|
||||
[#5570](https://github.com/vllm-project/vllm-omni/issues/5570) remains an open
|
||||
RFC. Proposed types and wire records from that RFC are not current contracts.
|
||||
|
||||
## Ownership boundary
|
||||
|
||||
This document owns semantic classification, fatal versus request-scoped
|
||||
failure, cross-process error context, compatibility policy, sanitization, and
|
||||
offline, HTTP, SSE, WebSocket, and asynchronous-job rendering rules.
|
||||
|
||||
`engine/messages.py` remains primarily owned by the I/O contract; this page
|
||||
owns the semantics of its error fields. Route and serving files remain
|
||||
primarily owned by `entrypoints.md`; this page owns their rendering
|
||||
constraints.
|
||||
|
||||
## Candidate invariants
|
||||
|
||||
The following topics require owner review before receiving stable IDs:
|
||||
|
||||
- A request-scoped failure terminates only the affected request unless the
|
||||
underlying engine is no longer usable.
|
||||
- Every public rendering path preserves a stable classification while
|
||||
removing sensitive internal context.
|
||||
- Cancellation remains distinguishable from internal failure across process
|
||||
and protocol boundaries.
|
||||
|
||||
## Invariant namespace
|
||||
|
||||
`ERR-INV` reserves `001-099` for classification ownership, `100-199` for wire
|
||||
context and compatibility, `200-299` for fatality, termination, sanitization,
|
||||
and rendering, and `300-399` for upstream alignment and transport extension
|
||||
points. Numbers become append-only after normative promotion.
|
||||
|
||||
## Safe-change guide
|
||||
|
||||
Exercise exception classification, stage wire round trips, fatal and
|
||||
request-scoped propagation, cancellation, sanitization, offline exceptions,
|
||||
HTTP status/body consistency, streaming terminal events, asynchronous-job
|
||||
persistence, and metrics classification.
|
||||
|
||||
## Promotion gate
|
||||
|
||||
- Resolve #5570's hierarchy, wire record, sanitization, and mixed-version
|
||||
questions.
|
||||
- Add table-driven exception/status/type tests and stage wire round trips.
|
||||
- Verify all public transports terminate once and do not leak sensitive
|
||||
payload data.
|
||||
- Obtain approval from both technical owners and an independent validation
|
||||
reviewer.
|
||||
@@ -2,46 +2,132 @@
|
||||
title: Input, Output, and Modality Contracts
|
||||
kind: module
|
||||
status: draft
|
||||
architecture_state: current
|
||||
owners:
|
||||
- "@Sy0307"
|
||||
- "@amy-why-3459"
|
||||
document_stewards:
|
||||
- "@hsliuustc0106"
|
||||
- "@Gaohan123"
|
||||
- "@david6666666"
|
||||
required_reviewers:
|
||||
- "@tzhouam"
|
||||
- "@gcanlin"
|
||||
- "@linyueqian"
|
||||
- "@david6666666"
|
||||
- "@yenuo26"
|
||||
- "@NickCao"
|
||||
primary_code_paths:
|
||||
- vllm_omni/inputs/**
|
||||
- vllm_omni/outputs/**
|
||||
- vllm_omni/request.py
|
||||
- vllm_omni/data_entry_keys.py
|
||||
- vllm_omni/engine/messages.py
|
||||
- vllm_omni/engine/serialization.py
|
||||
- vllm_omni/engine/mm_outputs.py
|
||||
- vllm_omni/engine/output_modality.py
|
||||
- vllm_omni/engine/output_processor.py
|
||||
related_code_paths:
|
||||
- vllm_omni/tokenizers/**
|
||||
depends_on: []
|
||||
- vllm_omni/model_executor/stage_input_processors/**
|
||||
- vllm_omni/core/sched/**
|
||||
- vllm_omni/distributed/omni_connectors/**
|
||||
- vllm_omni/diffusion/request.py
|
||||
- vllm_omni/diffusion/worker/request_batch.py
|
||||
- vllm_omni/errors.py
|
||||
depends_on:
|
||||
- error_contracts.md
|
||||
validation_paths:
|
||||
- tests/inputs/**
|
||||
- tests/engine/test_data_entry_keys.py
|
||||
- tests/engine/test_omni_request.py
|
||||
- tests/engine/test_omni_request_output.py
|
||||
- tests/engine/test_multimodal_accumulation.py
|
||||
- tests/engine/test_output_modality.py
|
||||
- tests/engine/test_output_processor.py
|
||||
- tests/engine/test_wire_multimodal_output.py
|
||||
- tests/utils/test_mm_outputs.py
|
||||
- tests/utils/test_mm_outputs_partition.py
|
||||
- tests/diffusion/test_diffusion_output_metadata.py
|
||||
upstream_refs:
|
||||
- vllm.inputs
|
||||
- vllm.outputs
|
||||
last_reviewed: 2026-07-16
|
||||
- vllm.inputs/**
|
||||
- vllm.outputs/**
|
||||
- vllm.v1.request.Request
|
||||
- vllm.v1.engine.EngineCoreRequest
|
||||
- vllm.v1.engine.EngineCoreOutput
|
||||
- vllm.v1.engine.output_processor.OutputProcessor
|
||||
invariant_namespace: IO-INV
|
||||
last_reviewed: 2026-08-07
|
||||
last_verified_commit: 3d7fc3b9ba3cac88d579d4dc35b78b0b641675fc
|
||||
---
|
||||
|
||||
# Input, output, and modality contracts
|
||||
|
||||
These contracts define the data that may cross entrypoint, orchestration,
|
||||
stage, and model boundaries.
|
||||
stage, connector, and model boundaries.
|
||||
|
||||
## Contract status
|
||||
|
||||
This document is a draft description of the current request, message,
|
||||
serialization, modality, accumulation, and output types. It does not assign
|
||||
semantic error ownership to the wire schema.
|
||||
|
||||
## Ownership boundary
|
||||
|
||||
This document owns request identity, prompt/input types, modality keys and
|
||||
metadata, queue and wire message schemas, serialization, output types,
|
||||
accumulation, completion semantics, and compatibility shims.
|
||||
|
||||
It does not own route-specific validation or rendering, model-specific
|
||||
interpretation, connector transport mechanics, scheduling policy, or semantic
|
||||
error classification. The `ErrorMessage` schema belongs here; the meaning and
|
||||
public rendering of its error fields belong to `error_contracts.md`.
|
||||
|
||||
## Candidate invariants
|
||||
|
||||
These identifiers are proposals while the document is `draft`.
|
||||
|
||||
### IO-INV-001: Boundary data has an explicit modality
|
||||
|
||||
**Rule:** Data crossing a module or stage boundary MUST identify its modality
|
||||
and use the corresponding validated contract.
|
||||
|
||||
### IO-INV-002: Request identity is stable
|
||||
### IO-INV-100: Request identity is stable
|
||||
|
||||
**Rule:** Request identity MUST be preserved across conversions, stages,
|
||||
streaming updates, cancellation, and errors.
|
||||
|
||||
### IO-INV-003: Internal objects do not leak into public protocols
|
||||
### IO-INV-101: Completion is explicit
|
||||
|
||||
**Rule:** Producers MUST distinguish partial updates from the terminal update
|
||||
for every output modality.
|
||||
|
||||
### IO-INV-300: Internal objects do not leak into public protocols
|
||||
|
||||
**Rule:** Entrypoints MUST explicitly translate internal output objects into
|
||||
public response types.
|
||||
|
||||
## Invariant namespace
|
||||
|
||||
`IO-INV` reserves `001-099` for schema ownership and dependency direction,
|
||||
`100-199` for identity, modality, accumulation, ordering, and completion,
|
||||
`200-299` for serialization failure and payload cleanup, and `300-399` for
|
||||
upstream extension, deprecation, and compatibility. Numbers become
|
||||
append-only after normative promotion.
|
||||
|
||||
## Safe-change guide
|
||||
|
||||
Test construction, validation, serialization, streaming, optional fields, and
|
||||
compatibility at every affected producer-consumer boundary.
|
||||
Test construction, validation, serialization, round trips, streaming,
|
||||
accumulation, completion, optional fields, unknown fields, and compatibility
|
||||
at every affected producer-consumer boundary.
|
||||
|
||||
## Promotion gate
|
||||
|
||||
- Add explicit CODEOWNERS for the primary I/O paths or confirm that fallback
|
||||
ownership is intentional.
|
||||
- Document the compatibility window for deprecated `vllm_omni.engine`
|
||||
I/O-related imports.
|
||||
- Verify wire round trips and rejection of unknown or incompatible schema
|
||||
fields.
|
||||
- Define payload-versus-metadata rules and completion semantics for streaming
|
||||
and batched outputs.
|
||||
- Obtain approval from a technical owner and independent AR, diffusion,
|
||||
connector, entrypoint, and validation reviewers as applicable.
|
||||
@@ -0,0 +1,126 @@
|
||||
---
|
||||
title: Stage Runtime and Replica Lifecycle
|
||||
kind: module
|
||||
status: draft
|
||||
architecture_state: current-plus-in-flight
|
||||
owners:
|
||||
- "@tzhouam"
|
||||
- "@fake0fan"
|
||||
document_stewards:
|
||||
- "@hsliuustc0106"
|
||||
- "@Gaohan123"
|
||||
- "@david6666666"
|
||||
required_reviewers:
|
||||
- "@Sy0307"
|
||||
- "@princepride"
|
||||
- "@chickeyton"
|
||||
- "@yenuo26"
|
||||
- "@NickCao"
|
||||
primary_code_paths:
|
||||
- vllm_omni/engine/stage_runtime.py
|
||||
- vllm_omni/engine/stage_pool.py
|
||||
- vllm_omni/engine/stage_client.py
|
||||
- vllm_omni/engine/membership_controller.py
|
||||
- vllm_omni/engine/stage_engine_startup.py
|
||||
- vllm_omni/engine/stage_engine_core_client.py
|
||||
- vllm_omni/engine/stage_engine_core_proc.py
|
||||
- vllm_omni/engine/stage_engine_core_proc_manager.py
|
||||
- vllm_omni/engine/stage_init_utils.py
|
||||
related_code_paths:
|
||||
- vllm_omni/distributed/omni_coordinator/**
|
||||
- vllm_omni/distributed/ray_utils/**
|
||||
- vllm_omni/diffusion/inline_stage_diffusion_client.py
|
||||
- vllm_omni/diffusion/stage_diffusion_client.py
|
||||
- vllm_omni/diffusion/stage_diffusion_proc.py
|
||||
- vllm_omni/config/**
|
||||
depends_on:
|
||||
- input_output_modality_contracts.md
|
||||
- error_contracts.md
|
||||
- omni_connector.md
|
||||
validation_paths:
|
||||
- tests/engine/test_async_omni_engine_stage_init.py
|
||||
- tests/engine/test_membership_controller.py
|
||||
- tests/engine/test_single_stage_mode.py
|
||||
- tests/engine/test_stage_engine_core_client.py
|
||||
- tests/engine/test_stage_engine_core_proc.py
|
||||
- tests/diffusion/test_inline_stage_diffusion_client.py
|
||||
- tests/diffusion/test_stage_diffusion_proc.py
|
||||
- tests/distributed/omni_coordinator/**
|
||||
upstream_refs:
|
||||
- vllm.v1.engine.core_client.AsyncMPClient
|
||||
- vllm.v1.engine.core_client.DPLBAsyncMPClient
|
||||
- vllm.v1.engine.core.EngineCoreProc
|
||||
- vllm.v1.engine.utils.CoreEngineProcManager
|
||||
- vllm.v1.engine.coordinator.DPCoordinator
|
||||
invariant_namespace: STAGE-INV
|
||||
last_reviewed: 2026-08-07
|
||||
last_verified_commit: 3d7fc3b9ba3cac88d579d4dc35b78b0b641675fc
|
||||
---
|
||||
|
||||
# Stage runtime and replica lifecycle
|
||||
|
||||
The stage runtime turns a logical stage definition into ready local or
|
||||
distributed replicas and provides the lifecycle boundary used by the
|
||||
orchestrator.
|
||||
|
||||
## Contract status
|
||||
|
||||
This document describes the current `StageRuntime`, `StagePool`, stage-client,
|
||||
and stage-process layout. The unified LLM/diffusion direction in
|
||||
[#5441](https://github.com/vllm-project/vllm-omni/pull/5441) remains in flight
|
||||
and is not presented as current behavior.
|
||||
|
||||
## Ownership boundary
|
||||
|
||||
This document owns local and distributed stage placement, startup and
|
||||
readiness, one pool per logical stage, replica identity and membership,
|
||||
selection and affinity, stage-client/process lifecycle, liveness, draining,
|
||||
and shutdown.
|
||||
|
||||
It does not own cross-stage request policy, model scheduler policy, connector
|
||||
transfer semantics, configuration precedence, or public error rendering.
|
||||
|
||||
## Candidate invariants
|
||||
|
||||
These identifiers are proposals while the document is `draft`.
|
||||
|
||||
### STAGE-INV-001: The runtime owns replica lifecycle
|
||||
|
||||
**Rule:** Orchestration MUST acquire stage capacity through the stage runtime
|
||||
and MUST NOT construct or retire backend replicas directly.
|
||||
|
||||
### STAGE-INV-100: Affinity uses stable replica identity
|
||||
|
||||
**Rule:** Follow-up operations for a request with replica affinity MUST resolve
|
||||
the same logical replica until the request terminates or the replica is
|
||||
declared lost.
|
||||
|
||||
### STAGE-INV-200: Shutdown is idempotent
|
||||
|
||||
**Rule:** Repeated shutdown or cleanup signals MUST NOT leak stage clients,
|
||||
process managers, or membership registrations.
|
||||
|
||||
## Invariant namespace
|
||||
|
||||
`STAGE-INV` reserves `001-099` for placement and ownership, `100-199` for
|
||||
replica identity/readiness/affinity, `200-299` for liveness loss, draining and
|
||||
cleanup, and `300-399` for local/distributed and upstream compatibility.
|
||||
Numbers become append-only after normative promotion.
|
||||
|
||||
## Safe-change guide
|
||||
|
||||
Exercise local and distributed startup, readiness, replica selection,
|
||||
affinity, membership changes, abort routing, process failure, draining, and
|
||||
repeated shutdown. Cross-stage policy changes belong in
|
||||
`engine_orchestration.md`.
|
||||
|
||||
## Promotion gate
|
||||
|
||||
- Reconcile terminology and ownership after #5441 reaches a final state.
|
||||
- Add or identify a dedicated StageRuntime/StagePool lifecycle suite; current
|
||||
evidence is distributed across the cited tests.
|
||||
- Demonstrate affinity preservation for update, interaction, and abort paths.
|
||||
- Demonstrate membership loss and repeated shutdown without leaked clients or
|
||||
processes.
|
||||
- Obtain approval from a technical owner and an independent validation
|
||||
reviewer.
|
||||
@@ -2,6 +2,8 @@
|
||||
title: vLLM-Omni Configuration
|
||||
kind: module
|
||||
status: draft
|
||||
architecture_state: deferred-pending-refactor
|
||||
ownership_status: provisional
|
||||
owners:
|
||||
- "@lishunyang12"
|
||||
- "@alex-jw-brooks"
|
||||
@@ -16,32 +18,39 @@ validation_paths:
|
||||
- tests/config/**
|
||||
upstream_refs:
|
||||
- vllm.config
|
||||
last_reviewed: 2026-07-16
|
||||
last_reviewed: 2026-08-07
|
||||
last_verified_commit: 3d7fc3b9ba3cac88d579d4dc35b78b0b641675fc
|
||||
---
|
||||
|
||||
# vLLM-Omni configuration
|
||||
|
||||
Configuration defines how user input becomes validated engine, stage,
|
||||
deployment, and platform settings.
|
||||
This page intentionally records discovery paths only. Substantive
|
||||
configuration ownership and contract work is deferred until the active
|
||||
configuration refactoring has settled.
|
||||
|
||||
## Candidate invariants
|
||||
## Deferred contract scope
|
||||
|
||||
### CONFIG-INV-001: Configuration is validated before startup
|
||||
The current code and tests remain authoritative for parsing, defaults,
|
||||
overrides, deployment topology, stage construction, and runtime projection.
|
||||
This draft does not define a stable precedence model, environment-variable
|
||||
capture rule, canonical configuration object, or invariant namespace.
|
||||
|
||||
**Rule:** Invalid stage graphs, resource assignments, and incompatible options
|
||||
MUST fail before managed runtime processes start.
|
||||
The PR review proposal to capture environment-derived values in configuration
|
||||
objects at initialization is an open design question, not a current
|
||||
invariant.
|
||||
|
||||
### CONFIG-INV-002: Precedence is deterministic
|
||||
## Safe-change guide while deferred
|
||||
|
||||
**Rule:** Defaults, files, environment variables, and CLI overrides MUST have a
|
||||
documented and deterministic precedence.
|
||||
Changes should trace every affected producer to its runtime consumer and test
|
||||
the structured, legacy, CLI, and deployment paths that are actually supported.
|
||||
Do not infer a stable contract from this placeholder page.
|
||||
|
||||
### CONFIG-INV-003: Runtime modules consume validated configuration
|
||||
## Promotion gate
|
||||
|
||||
**Rule:** Runtime modules MUST NOT independently reinterpret raw user
|
||||
configuration.
|
||||
|
||||
## Safe-change guide
|
||||
|
||||
Test defaults, parsing, override precedence, serialization, invalid
|
||||
combinations, and representative multi-stage configurations.
|
||||
- Finish or stabilize the configuration refactor and identify the canonical
|
||||
runtime configuration representation.
|
||||
- Confirm technical owners, primary path exceptions, and validation owners.
|
||||
- Document parsing and override precedence, including when environment values
|
||||
are captured.
|
||||
- Allocate an invariant namespace only after those behaviors are enforced and
|
||||
owner-approved.
|
||||
@@ -419,8 +419,7 @@ Notes:
|
||||
```bash
|
||||
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct \
|
||||
--omni \
|
||||
--port 8091 \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_omni_moe_async_chunk.yaml
|
||||
--port 8091
|
||||
```
|
||||
|
||||
#### 3) Key config knobs
|
||||
@@ -494,7 +493,7 @@ The default config (`qwen3_tts.yaml`) enables the full optimization stack:
|
||||
vllm-omni serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
|
||||
--omni \
|
||||
--port 8000 \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_tts_no_async_chunk.yaml
|
||||
--no-async-chunk
|
||||
```
|
||||
|
||||
#### 3) Key config knobs
|
||||
|
||||
Reference in new issue
Block a user