[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
@@ -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
|
||||
```
|
||||
Reference in new issue
Block a user