[Entrypoint][Refactor] vLLM-Omni Entrypoint Refactoring (#1908)
Signed-off-by: yinpe <11810305@mail.sustech.edu.cn> Signed-off-by: Peiqi Yin <yinpeiqi809@gmail.com> Signed-off-by: yinpeiqi <yinpeiqi809@gmail.com> Signed-off-by: wuhang <whlbx@hotmail.com> Signed-off-by: linyueqian <linyueqian@outlook.com> Signed-off-by: princepride <wangzhipeng628@gmail.com> Signed-off-by: lishunyang <lishunyang12@163.com> Signed-off-by: Chenguang ZHENG <645327136@qq.com> Signed-off-by: Peiqi Yin <60515999+yinpeiqi@users.noreply.github.com> Co-authored-by: yinpe <11810305@mail.sustech.edu.cn> Co-authored-by: Peiqi Yin <yinpeiqi809@gmail.com> Co-authored-by: wuhang <whlbx@hotmail.com> Co-authored-by: Peiqi Yin <60515999+yinpeiqi@users.noreply.github.com> Co-authored-by: linyueqian <linyueqian@outlook.com> Co-authored-by: princepride <wangzhipeng628@gmail.com> Co-authored-by: lishunyang <lishunyang12@163.com>
This commit is contained in:
125 files changed
+5210
-9322
No files matched your search
+1
-9
@@ -6,19 +6,13 @@ Main entry points for vLLM-Omni inference and serving.
|
||||
|
||||
- [vllm_omni.entrypoints.async_omni.AsyncOmni][]
|
||||
- [vllm_omni.entrypoints.async_omni_diffusion.AsyncOmniDiffusion][]
|
||||
- [vllm_omni.entrypoints.async_omni_llm.AsyncOmniLLM][]
|
||||
- [vllm_omni.entrypoints.cli.benchmark.base.OmniBenchmarkSubcommandBase][]
|
||||
- [vllm_omni.entrypoints.cli.benchmark.main.OmniBenchmarkSubcommand][]
|
||||
- [vllm_omni.entrypoints.cli.benchmark.serve.OmniBenchmarkServingSubcommand][]
|
||||
- [vllm_omni.entrypoints.cli.serve.OmniServeCommand][]
|
||||
- [vllm_omni.entrypoints.client_request_state.ClientRequestState][]
|
||||
- [vllm_omni.entrypoints.omni.Omni][]
|
||||
- [vllm_omni.entrypoints.omni.OmniBase][]
|
||||
- [vllm_omni.entrypoints.omni_diffusion.OmniDiffusion][]
|
||||
- [vllm_omni.entrypoints.omni_llm.OmniLLM][]
|
||||
- [vllm_omni.entrypoints.omni_stage.OmniStage][]
|
||||
- [vllm_omni.entrypoints.stage_utils.OmniStageTaskType][]
|
||||
- [vllm_omni.entrypoints.zmq_utils.ZmqQueue][]
|
||||
- [vllm_omni.entrypoints.omni_base.OmniBase][]
|
||||
|
||||
## Inputs
|
||||
|
||||
@@ -48,9 +42,7 @@ Engine classes for offline and online inference.
|
||||
- [vllm_omni.engine.OmniEngineCoreOutputs][]
|
||||
- [vllm_omni.engine.OmniEngineCoreRequest][]
|
||||
- [vllm_omni.engine.PromptEmbedsPayload][]
|
||||
- [vllm_omni.engine.arg_utils.AsyncOmniEngineArgs][]
|
||||
- [vllm_omni.engine.arg_utils.OmniEngineArgs][]
|
||||
- [vllm_omni.engine.input_processor.OmniInputProcessor][]
|
||||
- [vllm_omni.engine.output_processor.MultimodalOutputProcessor][]
|
||||
- [vllm_omni.engine.output_processor.OmniRequestState][]
|
||||
|
||||
|
||||
@@ -18,7 +18,7 @@ If users want to modify some part of it. The custom stage_configs file can be in
|
||||
For offline (Assume necessary dependencies have ben imported):
|
||||
```python
|
||||
model_name = "Qwen/Qwen2.5-Omni-7B"
|
||||
omni_llm = Omni(model=model_name, stage_configs_path="/path/to/custom_stage_configs.yaml")
|
||||
omni = Omni(model=model_name, stage_configs_path="/path/to/custom_stage_configs.yaml")
|
||||
```
|
||||
|
||||
For online serving:
|
||||
@@ -30,7 +30,7 @@ vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091 --stage-configs-path /path/to
|
||||
|
||||
Below is a specific example of stage_configs.yaml in Qwen2.5-omni.
|
||||
```python
|
||||
# stage config for running qwen2.5-omni with architecture of OmniLLM.
|
||||
# stage config for running qwen2.5-omni with AsyncOmniEngine + Orchestrator runtime.
|
||||
stage_args:
|
||||
- stage_id: 0 # mark the unique id for each stage
|
||||
runtime: # The disaggregated configuration
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
# stage config for running qwen2.5-omni with architecture of OmniLLM.
|
||||
# stage config for running qwen2.5-omni with AsyncOmniEngine + Orchestrator runtime.
|
||||
stage_args:
|
||||
- stage_id: 0
|
||||
runtime:
|
||||
|
||||
@@ -107,7 +107,7 @@ Only specific types of PRs will be reviewed. The PR title is prefixed appropriat
|
||||
- `[CI/Build]` for build or continuous integration improvements.
|
||||
- `[Doc]` for documentation fixes and improvements.
|
||||
- `[Model]` for adding a new model or improving an existing model. Model name should appear in the title.
|
||||
- `[Frontend]` For changes on the vLLM-Omni frontend (e.g., OpenAI API server, `OmniLLM` class, etc.)
|
||||
- `[Frontend]` For changes on the vLLM-Omni frontend (e.g., OpenAI API server, `Omni`/`AsyncOmni`, etc.)
|
||||
- `[Kernel]` for changes affecting CUDA kernels or other compute kernels.
|
||||
- `[Core]` for changes in the core vLLM-Omni logic (e.g., `OmniProcessor`, `OmniARScheduler`, etc.)
|
||||
- `[Hardware][Vendor]` for hardware-specific changes. Vendor name should appear in the prefix, such as [Ascend] for Ascend NPUs.
|
||||
|
||||
@@ -168,11 +168,6 @@ vllm_omni/ tests/
|
||||
│ └── arg_utils.py │ └── test_arg_utils.py ⬜
|
||||
│
|
||||
├── entrypoints/ → ├── entrypoints/
|
||||
│ ├── omni.py │ ├── test_omni.py ⬜ (E2E covered by e2e/offline, e2e/online)
|
||||
│ ├── omni_llm.py │ ├── test_omni_llm.py ✅
|
||||
│ ├── omni_stage.py │ ├── test_omni_stage.py ⬜ (partial in test_omni_stage_diffusion_config.py)
|
||||
│ ├── omni_diffusion.py │ ├── test_omni_diffusion.py ✅
|
||||
│ ├── async_omni.py │ ├── test_async_omni.py ✅ actually in e2e/online_serving/test_async_omni.py
|
||||
│ ├── async_omni_diffusion.py │ ├── test_async_omni_diffusion_config.py ✅
|
||||
│ ├── stage_utils.py │ ├── test_stage_utils.py ✅
|
||||
│ ├── cli/ │ ├── cli/ (benchmarks/test_serve_cli.py covers CLI serve)
|
||||
|
||||
@@ -24,8 +24,6 @@ End-to-end tests verify the complete functionality of a system or component. For
|
||||
|
||||
- **`tests/e2e/online_serving/`**: Tests for online serving scenarios (e.g., API server tests)
|
||||
|
||||
**Example:** The test file for `vllm_omni/entrypoints/omni_llm.py` should be located at `tests/entrypoints/test_omni_llm.py`.
|
||||
|
||||
## Test Directory Structure
|
||||
|
||||
The ideal directory structure mirrors the source code organization. Legend: `✅` = test exists, `⬜` = suggested to add.
|
||||
@@ -75,11 +73,6 @@ vllm_omni/ tests/
|
||||
│ └── arg_utils.py │ └── test_arg_utils.py ⬜
|
||||
│
|
||||
├── entrypoints/ → ├── entrypoints/
|
||||
│ ├── omni.py │ ├── test_omni.py ⬜ (E2E covered by e2e/offline, e2e/online)
|
||||
│ ├── omni_llm.py │ ├── test_omni_llm.py ✅
|
||||
│ ├── omni_stage.py │ ├── test_omni_stage.py ⬜ (partial in test_omni_stage_diffusion_config.py)
|
||||
│ ├── omni_diffusion.py │ ├── test_omni_diffusion.py ✅
|
||||
│ ├── async_omni.py │ ├── test_async_omni.py ✅ actually in e2e/online_serving/test_async_omni.py
|
||||
│ ├── async_omni_diffusion.py │ ├── test_async_omni_diffusion_config.py ✅
|
||||
│ ├── stage_utils.py │ ├── test_stage_utils.py ✅
|
||||
│ ├── cli/ │ ├── cli/ (benchmarks/test_serve_cli.py covers CLI serve)
|
||||
@@ -170,7 +163,7 @@ vllm_omni/ tests/
|
||||
|
||||
### Naming Conventions
|
||||
|
||||
- **Unit Tests**: Use `test_<module_name>.py` format. Example: `omni_llm.py` → `test_omni_llm.py`
|
||||
- **Unit Tests**: Use `test_<module_name>.py` format. Example: `stage_utils.py` → `test_stage_utils.py`
|
||||
|
||||
- **E2E Tests**: Place in `tests/e2e/offline_inference/` or `tests/e2e/online_serving/` with descriptive names. Example: `tests/e2e/offline_inference/test_qwen3_omni.py`, `tests/e2e/offline_inference/test_diffusion_model.py`
|
||||
|
||||
|
||||
@@ -330,54 +330,46 @@ Stage transitions are the mechanism by which outputs from one stage are converte
|
||||
|
||||
### Where Stage Transitions Are Called
|
||||
|
||||
Stage transitions happen automatically in the orchestrator (`OmniLLM` class) during the generation loop. Here's the detailed flow:
|
||||
Stage transitions happen automatically in the runtime orchestrator. Here's the detailed flow:
|
||||
|
||||
1. **Location**: `vllm_omni/entrypoints/omni_llm.py` in the `_run_generation()` method
|
||||
1. **Location**: `vllm_omni/engine/orchestrator.py` in `_forward_to_next_stage()`
|
||||
2. **Trigger**: When a stage completes processing and produces outputs
|
||||
3. **Execution Flow**:
|
||||
```python
|
||||
# In omni_llm.py, _run_generation() method (around line 345-460)
|
||||
# In orchestrator.py
|
||||
next_stage_id = stage_id + 1
|
||||
next_client = self.stage_clients[next_stage_id]
|
||||
params = req_state.sampling_params_list[next_stage_id]
|
||||
|
||||
# Main orchestrator loop polls each stage for completed requests
|
||||
for stage_id, stage in enumerate(self.stage_list):
|
||||
result = stage.try_collect() # Get completed request
|
||||
if result is None:
|
||||
continue
|
||||
# Save current stage outputs so stage_input_processors can consume them.
|
||||
self.stage_clients[stage_id].set_engine_outputs([output])
|
||||
|
||||
# Store outputs from this stage
|
||||
engine_outputs = _load(result, obj_key="engine_outputs", shm_key="engine_outputs_shm")
|
||||
stage.set_engine_outputs(engine_outputs)
|
||||
# THIS IS WHERE STAGE TRANSITION HAPPENS
|
||||
next_inputs = next_client.process_engine_inputs(
|
||||
stage_list=self.stage_clients,
|
||||
prompt=req_state.prompt,
|
||||
)
|
||||
|
||||
# Check if there's a next stage to forward to
|
||||
next_stage_id = stage_id + 1
|
||||
if next_stage_id < len(self.stage_list):
|
||||
next_stage: OmniStage = self.stage_list[next_stage_id]
|
||||
|
||||
# THIS IS WHERE STAGE TRANSITION HAPPENS
|
||||
next_inputs = next_stage.process_engine_inputs(
|
||||
self.stage_list,
|
||||
[request_id_to_prompt[req_id]]
|
||||
)
|
||||
|
||||
# Submit to next stage
|
||||
task = {
|
||||
"type": OmniStageTaskType.GENERATE,
|
||||
"request_id": req_id,
|
||||
"engine_inputs": next_inputs[0],
|
||||
"sampling_params": sampling_params_list[next_stage_id],
|
||||
}
|
||||
next_stage.submit(task)
|
||||
# Build and submit request(s) to the next stage.
|
||||
for next_input in next_inputs:
|
||||
request = build_engine_core_request_from_tokens(
|
||||
request_id=req_id,
|
||||
prompt=next_input,
|
||||
params=params,
|
||||
model_config=self.stage_vllm_configs[next_stage_id].model_config,
|
||||
)
|
||||
await next_client.add_request_async(request)
|
||||
```
|
||||
|
||||
### How Stage Transitions Work
|
||||
|
||||
The stage transition process follows these steps:
|
||||
|
||||
1. **Stage Completion**: When a stage finishes processing a request, it stores outputs via `stage.set_engine_outputs(engine_outputs)`
|
||||
1. **Stage Completion**: When a stage finishes processing a request, the orchestrator stores outputs via `stage_client.set_engine_outputs(...)`
|
||||
|
||||
2. **Transition Detection**: The orchestrator checks if there's a next stage and calls `process_engine_inputs()` on it
|
||||
|
||||
3. **Input Processing**: The `process_engine_inputs()` method in `OmniStage` (`omni_stage.py`) handles the transition:
|
||||
3. **Input Processing**: The stage input processor configured in stage YAML (under `vllm_omni/model_executor/stage_input_processors/`) handles the transition:
|
||||
```python
|
||||
def process_engine_inputs(
|
||||
self, stage_list: list[Any], prompt: OmniTokensPrompt | TextPrompt = None
|
||||
|
||||
@@ -23,30 +23,32 @@ export VLLM_PROFILER_MAX_ITERS=1
|
||||
The profiler is default to function across all stages. But It is highly recommended to profile specific stages by passing the stages list, preventing from producing too large trace files:
|
||||
```python
|
||||
# Profile all stages
|
||||
omni_llm.start_profile()
|
||||
omni.start_profile()
|
||||
|
||||
# Only profile Stage 1
|
||||
omni_llm.start_profile(stages=[1])
|
||||
omni.start_profile(stages=[1])
|
||||
```
|
||||
|
||||
```python
|
||||
# Stage 0 (Thinker) and Stage 2 (Audio Decoder) for qwen omni
|
||||
omni_llm.start_profile(stages=[0, 2])
|
||||
omni.start_profile(stages=[0, 2])
|
||||
```
|
||||
|
||||
**Python Usage**: Wrap your generation logic with `start_profile()` and `stop_profile()`.
|
||||
|
||||
```python
|
||||
from vllm_omni import omni_llm
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
|
||||
omni = Omni(model="Qwen/Qwen3-Omni-30B-A3B-Instruct")
|
||||
|
||||
profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR"))
|
||||
|
||||
# 1. Start profiling if enabled
|
||||
if profiler_enabled:
|
||||
omni_llm.start_profile(stages=[0])
|
||||
omni.start_profile(stages=[0])
|
||||
|
||||
# Initialize generator
|
||||
omni_generator = omni_llm.generate(prompts, sampling_params_list, py_generator=args.py_generator)
|
||||
omni_generator = omni.generate(prompts, sampling_params_list, py_generator=args.py_generator)
|
||||
|
||||
total_requests = len(prompts)
|
||||
processed_count = 0
|
||||
@@ -57,21 +59,21 @@ for stage_outputs in omni_generator:
|
||||
# ... [Output processing logic for text/audio would go here] ...
|
||||
|
||||
# Update count to track when to stop profiling
|
||||
processed_count += len(stage_outputs.request_output)
|
||||
processed_count += 1
|
||||
|
||||
# 2. Check if all requests are done to stop the profiler safely
|
||||
if profiler_enabled and processed_count >= total_requests:
|
||||
print(f"[Info] Processed {processed_count}/{total_requests}. Stopping profiler inside active loop...")
|
||||
|
||||
# Stop the profiler while workers are still active
|
||||
omni_llm.stop_profile()
|
||||
omni.stop_profile()
|
||||
|
||||
# Wait for traces to flush to disk
|
||||
print("[Info] Waiting 30s for workers to write trace files to disk...")
|
||||
time.sleep(30)
|
||||
print("[Info] Trace export wait time finished.")
|
||||
|
||||
omni_llm.close()
|
||||
omni.close()
|
||||
```
|
||||
|
||||
|
||||
|
||||
@@ -67,12 +67,12 @@ According to analysis for current popular open-source models, most of them have
|
||||
| Component | Description |
|
||||
| ----------------- | ---------------------------------------------------------------------------------------------------------------------------------------- |
|
||||
| **OmniRouter** | provide an intelligent router for Omni-modality requests dispatch |
|
||||
| **EntryPoints** | define the APIs for offline/online serving (APIServer, Omni/AsyncOmni) and provide the OmniStage abstraction for different AR/DiT stages |
|
||||
| **EntryPoints** | define the APIs for offline/online serving (APIServer, Omni/AsyncOmni), while `AsyncOmniEngine` and `Orchestrator` coordinate multi-stage AR/DiT execution |
|
||||
| **AR** | adapted for omni-modality models while inheriting efficient features from vLLM, such as cache management |
|
||||
| **Diffusion** | natively implemented and optimized using acceleration components |
|
||||
| **OmniConnector** | supports fully disaggregation based on E/P/D/G (Encoding/Processing/Decoding/Generation) disaggregation across stages |
|
||||
|
||||
Disaggregated stages are managed through configuration, such as in the Qwen3-Omni example, where stages like Thinker, Talker, and Code2wav are defined as separate OmniStage instances with specific resources and input/output type.
|
||||
Disaggregated stages are managed through stage configuration. In Qwen3-Omni, Thinker/Talker/Code2wav are declared as separate configured stages, and runtime routing is handled by `Orchestrator` over `StageEngineCoreClient` / `StageDiffusionClient`.
|
||||
|
||||
## Main features
|
||||
|
||||
@@ -127,10 +127,10 @@ Taking **Qwen3-Omni** as an example:
|
||||
The **Omni** class provides a Python interface for offline batched inference. Users initialize the Omni class with a Hugging Face model name and use the generate method, passing inputs that include both text prompts and multi-modal data:
|
||||
|
||||
```
|
||||
# Create an omni_lm with HF model name.
|
||||
# Create an omni runtime with HF model name.
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
|
||||
omni_lm = Omni(model="Qwen/Qwen3-Omni-30B-A3B-Instruct")
|
||||
omni = Omni(model="Qwen/Qwen3-Omni-30B-A3B-Instruct")
|
||||
|
||||
# Example prompts.
|
||||
om_inputs = {"prompt": prompt,
|
||||
@@ -140,7 +140,7 @@ om_inputs = {"prompt": prompt,
|
||||
}}
|
||||
|
||||
# Generate texts and audio from the multi-modality inputs.
|
||||
outputs = omni_lm.generate(om_inputs, sampling_params_list)
|
||||
outputs = omni.generate(om_inputs, sampling_params_list)
|
||||
```
|
||||
|
||||
## Online Serving
|
||||
|
||||
@@ -85,9 +85,6 @@ classDiagram
|
||||
class InputProcessor {
|
||||
+process_inputs() EngineCoreRequest
|
||||
}
|
||||
class OmniInputProcessor {
|
||||
+process_inputs() OmniEngineCoreRequest
|
||||
}
|
||||
|
||||
class VLLMOutputProcessor {
|
||||
+process_outputs() OutputProcessorOutput
|
||||
@@ -96,7 +93,6 @@ classDiagram
|
||||
+process_outputs() OutputProcessorOutput
|
||||
+_route_and_normalize()
|
||||
}
|
||||
InputProcessor <|-- OmniInputProcessor
|
||||
VLLMOutputProcessor <|-- MultimodalOutputProcessor
|
||||
```
|
||||
|
||||
@@ -105,7 +101,7 @@ classDiagram
|
||||
- **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**: `OmniInputProcessor` extends `vllm.v1.engine.input_processor.InputProcessor` to serialize prompt embeddings and additional information
|
||||
- **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
|
||||
@@ -118,7 +114,7 @@ The following diagram illustrates the request flow through the AR module compone
|
||||
|
||||
```mermaid
|
||||
flowchart TD
|
||||
A[OmniInputProcessor] -->|OmniEngineCoreRequest| B[OmniARScheduler]
|
||||
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]
|
||||
@@ -297,15 +293,17 @@ The input/output processing pipeline handles serialization, routing, and accumul
|
||||
```mermaid
|
||||
sequenceDiagram
|
||||
participant Client
|
||||
participant OmniInputProcessor
|
||||
participant AsyncOmniEngine
|
||||
participant InputProcessor
|
||||
participant Scheduler
|
||||
participant ModelRunner
|
||||
participant MultimodalOutputProcessor
|
||||
participant Client
|
||||
|
||||
Client->>OmniInputProcessor: prompt + prompt_embeds + additional_info
|
||||
OmniInputProcessor->>OmniInputProcessor: Serialize tensors to payloads
|
||||
OmniInputProcessor->>Scheduler: OmniEngineCoreRequest (with payloads)
|
||||
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
|
||||
@@ -318,43 +316,18 @@ sequenceDiagram
|
||||
MultimodalOutputProcessor->>Client: RequestOutput (with multimodal_output)
|
||||
```
|
||||
|
||||
### OmniInputProcessor
|
||||
### Stage-0 Input Processing
|
||||
|
||||
`OmniInputProcessor` extends the base input processor to serialize prompt embeddings and additional information for inter-stage transfer.
|
||||
|
||||
#### Payload Serialization
|
||||
|
||||
Converts PyTorch tensors to serialized payloads:
|
||||
Stage-0 now uses upstream `InputProcessor` directly, and `AsyncOmniEngine` upgrades the request to `OmniEngineCoreRequest` while restoring omni-specific payloads.
|
||||
|
||||
```python
|
||||
def process_inputs(self, ...) -> OmniEngineCoreRequest:
|
||||
# Serialize prompt_embeds
|
||||
if "prompt_embeds" in decoder_inputs:
|
||||
pe_cpu = decoder_inputs["prompt_embeds"].detach().to("cpu").contiguous()
|
||||
prompt_embeds_payload = PromptEmbedsPayload(
|
||||
data=pe_cpu.numpy().tobytes(),
|
||||
shape=[seq_len, hidden_size],
|
||||
dtype=dtype_str,
|
||||
)
|
||||
|
||||
# Serialize additional_information
|
||||
if "additional_information" in decoder_inputs:
|
||||
entries = {}
|
||||
for key, value in raw_info.items():
|
||||
if isinstance(value, torch.Tensor):
|
||||
entry = AdditionalInformationEntry(
|
||||
tensor_data=value.numpy().tobytes(),
|
||||
tensor_shape=list(value.shape),
|
||||
tensor_dtype=dtype_str,
|
||||
)
|
||||
entries[key] = entry
|
||||
additional_information_payload = AdditionalInformationPayload(entries=entries)
|
||||
|
||||
return OmniEngineCoreRequest(
|
||||
# ... standard fields ...
|
||||
prompt_embeds=prompt_embeds_payload,
|
||||
additional_information=additional_information_payload,
|
||||
)
|
||||
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
|
||||
@@ -400,13 +373,13 @@ The AR module of vLLM-Omni extends vLLM through strategic inheritance and minima
|
||||
### Key Design Patterns
|
||||
|
||||
1. **Inheritance over composition**: Extends vLLM classes to preserve compatibility with existing scheduling, batching, and execution mechanisms
|
||||
2. **Payload serialization**: Uses `PromptEmbedsPayload` and `AdditionalInformationPayload` for efficient inter-stage data transfer
|
||||
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 prompt embeddings and additional information enable direct transfer between pipeline stages
|
||||
- **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
|
||||
|
||||
@@ -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_async()
|
||||
-> (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_async
|
||||
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
|
||||
```
|
||||
@@ -97,7 +97,7 @@ outputs = omni.generate(
|
||||
)
|
||||
|
||||
# Access custom trajectory data
|
||||
output = outputs[0].request_output[0]
|
||||
output = outputs[0].request_output
|
||||
print(f"Trajectory timesteps shape: {output.metrics['trajectory_timesteps'].shape}")
|
||||
print(f"Trajectory latents shape: {output.latents.shape}")
|
||||
```
|
||||
|
||||
@@ -42,7 +42,7 @@ if __name__ == "__main__":
|
||||
omni = Omni(model="Tongyi-MAI/Z-Image-Turbo")
|
||||
prompt = "a cup of coffee on the table"
|
||||
outputs = omni.generate(prompt)
|
||||
images = outputs[0].request_output[0].images
|
||||
images = outputs[0].request_output.images
|
||||
images[0].save("coffee.png")
|
||||
```
|
||||
|
||||
@@ -70,7 +70,7 @@ if __name__ == "__main__":
|
||||
]
|
||||
omni_outputs = omni.generate(prompts)
|
||||
for i_prompt, prompt_output in enumerate(omni_outputs):
|
||||
this_request_output = prompt_output.request_output[0]
|
||||
this_request_output = prompt_output.request_output
|
||||
this_images = this_request_output.images
|
||||
for i_image, image in enumerate(this_images):
|
||||
image.save(f"p{i_prompt}-img{i_image}.jpg")
|
||||
|
||||
@@ -22,10 +22,16 @@
|
||||
CLI (examples/offline_inference/text_to_image/text_to_image.py)
|
||||
|
|
||||
v
|
||||
Omni (vllm_omni/entrypoints/omni.py)
|
||||
Omni (vllm_omni/entrypoints/__init__.py)
|
||||
|
|
||||
v
|
||||
OmniStage (diffusion)
|
||||
AsyncOmniEngine
|
||||
|
|
||||
v
|
||||
Orchestrator
|
||||
|
|
||||
v
|
||||
StageDiffusionClient
|
||||
|
|
||||
v
|
||||
DiffusionWorker
|
||||
@@ -64,7 +70,13 @@ _generate_with_async_omni
|
||||
AsyncOmni
|
||||
|
|
||||
v
|
||||
DiffusionEngine
|
||||
AsyncOmniEngine
|
||||
|
|
||||
v
|
||||
Orchestrator
|
||||
|
|
||||
v
|
||||
StageDiffusionClient
|
||||
|
|
||||
v
|
||||
OmniRequestOutput
|
||||
|
||||
@@ -19,7 +19,7 @@ if __name__ == "__main__":
|
||||
omni = Omni(model="Qwen/Qwen-Image")
|
||||
prompt = "a cup of coffee on the table"
|
||||
outputs = omni.generate(prompt)
|
||||
images = outputs[0].request_output[0].images
|
||||
images = outputs[0].request_output.images
|
||||
images[0].save("coffee.png")
|
||||
```
|
||||
|
||||
@@ -37,7 +37,7 @@ if __name__ == "__main__":
|
||||
]
|
||||
outputs = omni.generate(prompts)
|
||||
for i, output in enumerate(outputs):
|
||||
image = output.request_output[0].images[0].save(f"{i}.jpg")
|
||||
image = output.request_output.images[0].save(f"{i}.jpg")
|
||||
```
|
||||
|
||||
!!! info
|
||||
@@ -72,7 +72,7 @@ if __name__ == "__main__":
|
||||
}
|
||||
])
|
||||
for i, output in enumerate(outputs):
|
||||
image = output.request_output[0].images[0].save(f"{i}.jpg")
|
||||
image = output.request_output.images[0].save(f"{i}.jpg")
|
||||
```
|
||||
|
||||
## Local CLI Usage
|
||||
|
||||
Reference in new issue
Block a user