anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

[Entrypoint][Refactor] vLLM-Omni Entrypoint Refactoring (#1908)

Browse Source 
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:
Chenguang ZhengyinpePeiqi YinwuhangPeiqi Yinlinyueqianprincepridelishunyang
 authored and GitHub committed 2026-03-18 17:33:53 +08:00
1 parent 6a66a1fe36
commit 3158e9d28f
125 files changed
+5210 -9322
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

docs/contributing/README.md
+1 -1
View File
@@ -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.
docs/contributing/ci/CI_5levels.md
-5
View File
@@ -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)
docs/contributing/ci/tests_style.md
+1 -8
View File
@@ -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`
docs/contributing/model/adding_omni_model.md
+24 -32
View File
@@ -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
docs/contributing/profiling.md
+11 -9
View File
@@ -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()
```