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

[Refactor] Migrate Legacy Stage Configs to the Pipeline Registry (#5031)

Browse Source 
Signed-off-by: zwhzzz0821 <2831474076@qq.com>
Co-authored-by: zwhzzz0821 <zwhzzz0821@users.noreply.github.com>
Co-authored-by: SYLAR <125541396+lishunyang12@users.noreply.github.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
This commit is contained in:
Yukim1zwhzzz0821SYLARHongsheng Liu
 authored and GitHub committed 2026-07-18 09:34:33 +00:00
1 parent 793d3e9ab7
commit 4fbbdfd250
41 files changed
+520 -516
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

docs/cli/serve.md
+1 -2
View File
@@ -8,8 +8,7 @@ The stage-based CLI is designed for deployments that require launching each pipe
- For **migrated models** that utilize the bundled deployment YAML configurations located in
`vllm_omni/deploy/`, the `--deploy-config` flag is only required to override the default configuration. By default, executing `vllm serve MODEL --omni ...`
automatically loads the bundled deployment configuration.
- For **legacy models** utilizing configuration files located in
`vllm_omni/model_executor/stage_configs/`, the `--stage-configs-path` parameter remains mandatory.
- For custom legacy `stage_args` YAMLs, pass the file with `--stage-configs-path`.
Example: Initializing Stage 0 (Orchestrator and API Server):
The commands below show a common device mapping where Stage 0 uses GPU 0 and
docs/configuration/stage_configs.md
+1 -1
View File
@@ -3,7 +3,7 @@
In vLLM-Omni, a model's `PipelineConfig` defines its fixed stage topology, while a deploy configuration controls how those stages run. Models that have not migrated to this split still use legacy stage configurations with a `stage_args` schema.
!!! note
Default deploy config YAMLs (for example, `vllm_omni/deploy/qwen2_5_omni.yaml`, `vllm_omni/deploy/qwen3_omni_moe.yaml`, and `vllm_omni/deploy/qwen3_tts.yaml`) are bundled and loaded automatically when neither `--stage-configs-path` nor `--deploy-config` is provided — the model registry resolves the right pipeline + deploy YAML by `model_type`. The bundled defaults have been verified on 1xH100 for Qwen2.5-Omni and 2xH100 for Qwen3-Omni. Models that have not yet migrated to the new schema continue to use the legacy `vllm_omni/model_executor/stage_configs/<model>.yaml` files via `--stage-configs-path`.
Default deploy config YAMLs (for example, `vllm_omni/deploy/qwen2_5_omni.yaml`, `vllm_omni/deploy/qwen3_omni_moe.yaml`, and `vllm_omni/deploy/qwen3_tts.yaml`) are bundled and loaded automatically when neither `--stage-configs-path` nor `--deploy-config` is provided — the model registry resolves the right pipeline + deploy YAML by `model_type`. Custom legacy `stage_args` YAMLs are still accepted via `--stage-configs-path`.
## Deploy configuration schema
docs/contributing/model/adding_omni_model.md
+7 -2
View File
@@ -34,6 +34,7 @@ vllm_omni/model_executor/models/
└── your_model_name/ # Model directory (e.g., qwen3_omni)
├── __init__.py # Exports main model class
├── your_model.py # Main unified model class
├── pipeline.py # Pipeline topology registry entry
├── your_model_stage1_implementation.py # Stage 1 implementation (e.g., thinker)
├── your_model_stage2_implementation.py # Stage 2 implementation (e.g., talker)
└── your_model_stage3_implementation.py # Stage 3 implementation (e.g., code2wav)
@@ -42,10 +43,14 @@ vllm_omni/model_executor/models/
vllm_omni/model_executor/stage_input_processors/
└── your_model_name.py # Stage transition processors
vllm_omni/model_executor/stage_configs/
└── your_model_name.yaml # Stage configuration file
vllm_omni/deploy/
└── your_model_name.yaml # Deployment configuration file
```
New in-tree models should register their `PipelineConfig` in
`vllm_omni/config/pipeline_registry.py`; use `--stage-configs-path` only for
custom legacy `stage_args` YAMLs.
## Step-by-Step Implementation
### Step 1: Create the Model Directory
docs/contributing/model/adding_tts_model.md
+13 -4
View File
@@ -175,15 +175,16 @@ vllm_omni/model_executor/models/
your_model_name/
__init__.py
your_model.py # Unified class (stage dispatch)
pipeline.py # Pipeline topology registry entry
your_model_ar_stage.py # Stage 0: AR stage
your_model_decoder.py # Stage 1: audio decoder
vllm_omni/model_executor/stage_input_processors/
your_model_name.py # Stage 0 -> Stage 1 transition
vllm_omni/model_executor/stage_configs/
your_model_name.yaml # Batch mode config
your_model_name_async_chunk.yaml # Streaming mode config
vllm_omni/deploy/
your_model_name.yaml # Batch/deploy config
your_model_name_async_chunk.yaml # Streaming/deploy config
```
### Example placement
@@ -445,6 +446,13 @@ Key configuration fields:
| `final_output` | Whether this stage produces the final user-facing output |
| `final_output_type` | Type of final output (`audio`, `text`, etc.) |
!!! note
New in-tree models should define frozen topology in `models/<model>/pipeline.py`,
register it in `vllm_omni/config/pipeline_registry.py`, and put deployment
knobs in `vllm_omni/deploy/<model>.yaml`. Legacy `stage_args` YAMLs are still
accepted for custom configs via `--stage-configs-path`, but should not be used
for new bundled defaults.
### Batch mode
```yaml
@@ -771,8 +779,9 @@ directly from the model's own generator.
vllm_omni/model_executor/models/your_model_name/
__init__.py
modeling_your_model_name.py # unified class: load_weights + forward + streaming
pipeline.py # single-stage pipeline topology
vllm_omni/model_executor/stage_configs/your_model_name.yaml
vllm_omni/deploy/your_model_name.yaml
```
No stage input processor is needed.
examples/offline_inference/image_to_video/image_to_video.py
+7
View File
@@ -102,6 +102,11 @@ def parse_args() -> argparse.Namespace:
default=None,
help="Override model class name (e.g., LTX2ImageToVideoPipeline or LTX23ImageToVideoPipeline).",
)
parser.add_argument(
"--deploy-config",
default=None,
help="Optional deploy config YAML to use for pipeline-backed runs.",
)
parser.add_argument("--image", help="Path to the first-frame or source image.")
parser.add_argument("--last-image", help="Path to a last-frame condition (used by models such as VACE).")
parser.add_argument("--mask-image", help="Path to an inpainting mask (used by models such as VACE).")
@@ -453,6 +458,8 @@ def main():
enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
profiler_config=args.profiler_config,
)
if args.deploy_config:
omni_kwargs["deploy_config"] = args.deploy_config
if args.quantization is not None:
omni_kwargs["quantization"] = args.quantization
# Cosmos3 loads its (gated) guardrail models at build time, so the guardrails
examples/offline_inference/mimo_audio/end2end.py
+10 -10
View File
@@ -219,31 +219,31 @@ def main(args):
# Build query result based on query type
# Notice: The audio files used in this example are available at: https://github.com/XiaomiMiMo/MiMo-Audio/tree/main/examples
if args.query_type == "tts_sft":
# python3 -u end2end.py --stage-configs-path ${config_file} --model ${MODEL_PATH} --query-type tts_sft
# python3 -u end2end.py --deploy-config ${config_file} --model-name ${MODEL_PATH} --query-type tts_sft
query_result = query_func(text=text, read_text_only=True)
elif args.query_type == "tts_sft_with_instruct":
# python3 -u end2end.py --stage-configs-path ${config_file} --model ${MODEL_PATH} --query-type tts_sft_with_instruct --instruct "Speak happily in a child's voice"
# python3 -u end2end.py --deploy-config ${config_file} --model-name ${MODEL_PATH} --query-type tts_sft_with_instruct --instruct "Speak happily in a child's voice"
query_result = query_func(text=text, instruct=instruct, read_text_only=True)
elif args.query_type == "tts_sft_with_audio":
# python3 -u end2end.py --stage-configs-path ${config_file} --model ${MODEL_PATH} --query-type tts_sft_with_audio --audio_path "./spoken_dialogue_assistant_turn_1.wav"
# python3 -u end2end.py --deploy-config ${config_file} --model-name ${MODEL_PATH} --query-type tts_sft_with_audio --audio-path "./spoken_dialogue_assistant_turn_1.wav"
audio_list = [get_audio_data(audio_path)]
query_result = query_func(text=text, read_text_only=True, prompt_speech=audio_path, audio_list=audio_list)
elif args.query_type == "tts_sft_with_natural_instruction":
# python3 -u end2end.py --stage-configs-path ${config_file} --model ${MODEL_PATH} --query-type tts_sft_with_natural_instruction --text "In a panting young male voice, he said: I can't run anymore, wait for me!"
# python3 -u end2end.py --deploy-config ${config_file} --model-name ${MODEL_PATH} --query-type tts_sft_with_natural_instruction --text "In a panting young male voice, he said: I can't run anymore, wait for me!"
query_result = query_func(text=text, read_text_only=False)
elif args.query_type == "audio_trancribing_sft":
# python3 -u end2end.py --stage-configs-path ${config_file} --model ${MODEL_PATH} --query-type audio_trancribing_sft --audio_path "./spoken_dialogue_assistant_turn_1.wav"
# python3 -u end2end.py --deploy-config ${config_file} --model-name ${MODEL_PATH} --query-type audio_trancribing_sft --audio-path "./spoken_dialogue_assistant_turn_1.wav"
audio_path = "spoken_dialogue_assistant_turn_1.wav"
text = "Please transcribe this audio and repeat it once."
query_result = query_func(text=text, audio_path=audio_path, use_sostm=True)
elif args.query_type == "audio_understanding_sft":
# python3 -u end2end.py --stage-configs-path ${config_file} --model ${MODEL_PATH} --query-type audio_understanding_sft --text "Summarize the audio." --audio_path "./spoken_dialogue_assistant_turn_1.wav"
# python3 -u end2end.py --deploy-config ${config_file} --model-name ${MODEL_PATH} --query-type audio_understanding_sft --text "Summarize the audio." --audio-path "./spoken_dialogue_assistant_turn_1.wav"
query_result = query_func(text=text, audio_path=audio_path)
elif args.query_type == "audio_understanding_sft_with_thinking":
# python3 -u end2end.py --stage-configs-path ${config_file} --model ${MODEL_PATH} --query-type audio_understanding_sft_with_thinking --text "Summarize the audio." --audio_path "./spoken_dialogue_assistant_turn_1.wav"
# python3 -u end2end.py --deploy-config ${config_file} --model-name ${MODEL_PATH} --query-type audio_understanding_sft_with_thinking --text "Summarize the audio." --audio-path "./spoken_dialogue_assistant_turn_1.wav"
query_result = query_func(text=text, audio_path=audio_path, thinking=True)
elif args.query_type == "spoken_dialogue_sft_multiturn":
# python3 -u end2end.py --stage-configs-path ${config_file} --model ${MODEL_PATH} --query-type spoken_dialogue_sft_multiturn --audio_path "./prompt_speech_zh_m.wav"
# python3 -u end2end.py --deploy-config ${config_file} --model-name ${MODEL_PATH} --query-type spoken_dialogue_sft_multiturn --audio-path "./prompt_speech_zh_m.wav"
first_turn_text_response = "我没办法获取实时的天气信息。不过呢，你可以试试几个方法来查看今天的天气。首先，你可以用手机自带的天气功能，比如苹果手机的天气应用，或者直接在系统设置里查看。其次，你也可以用一些专业的天气服务，像是国外的AccuWeather、Weather.com，或者国内的中国天气网、墨迹天气等等。再有就是，你还可以在谷歌或者百度里直接搜索你所在的城市加上天气这两个字。如果你能告诉我你所在的城市，我也可以帮你分析一下历史天气趋势，不过最新的数据还是需要你通过官方渠道去获取哦。"
audio_list = []
s1_audio_path = "weather_of_today.mp3"
@@ -261,7 +261,7 @@ def main(args):
]
query_result = query_func(message_list, system_prompt=None, ref_audio_path=audio_path, audio_list=audio_list)
elif args.query_type == "speech2text_dialogue_sft_multiturn":
# python3 -u end2end.py --stage-configs-path ${config_file} --model ${MODEL_PATH} --query-type speech2text_dialogue_sft_multiturn
# python3 -u end2end.py --deploy-config ${config_file} --model-name ${MODEL_PATH} --query-type speech2text_dialogue_sft_multiturn
s1_audio_path = "weather_of_today.mp3"
s2_audio_path = "beijing.mp3"
audio_list = []
@@ -277,7 +277,7 @@ def main(args):
]
query_result = query_func(message_list, thinking=True, audio_list=audio_list)
elif args.query_type == "text_dialogue_sft_multiturn":
# python3 -u end2end.py --stage-configs-path ${config_file} --model ${MODEL_PATH} --query-type text_dialogue_sft_multiturn
# python3 -u end2end.py --deploy-config ${config_file} --model-name ${MODEL_PATH} --query-type text_dialogue_sft_multiturn
message_list = [
{"role": "user", "content": "Could you recommend some tourist attractions in China?"},
{"role": "assistant", "content": "Hello, which city would you like to travel to?"},
examples/offline_inference/step_audio2/README.md
+34 -39
View File
@@ -51,7 +51,7 @@ Single request latency comparison between vLLM-Omni and official Step-Audio2 imp
**Benchmark environment:**
- GPU: NVIDIA H100 80GB (single card)
- Model: Step-Audio2-mini
- Model: Step-Audio-2-mini
- Warmup: 1 run, Measured: 3 runs (averaged)
### Async Chunk Streaming Performance
@@ -73,7 +73,7 @@ Comparison between sequential (non-async) and async chunk modes via `/v1/audio/s
**Benchmark environment:**
- GPU: 4x NVIDIA RTX 3090 24GB (TP=2 for Thinker, 1 GPU for Token2Wav)
- Model: Step-Audio2-mini
- Model: Step-Audio-2-mini
- Endpoint: `/v1/audio/speech` (10 prompts, concurrency=1)
- Measured via `bench_tts_serve.py`
@@ -91,23 +91,22 @@ pip install step-audio2
## Model Setup
Step-Audio2 ships a custom Transformers configuration, so this model-specific
example enables remote code loading when it initializes the model.
### Option 1: Auto-download from HuggingFace (Recommended)
The script will **automatically download** the model on first run:
The script will **automatically download** the model on first run when a
HuggingFace model id is passed:
```bash
# Just run without specifying --model, it will auto-download stepfun-ai/Step-Audio2-mini
python end2end.py --query-type audio_to_text
# Or explicitly specify the HuggingFace model
python end2end.py --query-type audio_to_text --model stepfun-ai/Step-Audio2-mini
python end2end.py --query-type audio_to_text --model stepfun-ai/Step-Audio-2-mini
```
Models will be cached in `~/.cache/huggingface/hub/` for future use.
**Available models**:
- `stepfun-ai/Step-Audio2-mini` (smaller, faster)
- `stepfun-ai/Step-Audio2-7B` (larger, better quality)
**Supported model**:
- `stepfun-ai/Step-Audio-2-mini`
### Option 2: Manual Download (for offline use)
@@ -115,15 +114,15 @@ Download and use locally:
```bash
# Download from HuggingFace
huggingface-cli download stepfun-ai/Step-Audio2-mini --local-dir ./models/Step-Audio2-mini
hf download stepfun-ai/Step-Audio-2-mini --local-dir ./models/Step-Audio-2-mini
# Then use the local path
python end2end.py --query-type audio_to_text --model ./models/Step-Audio2-mini
python end2end.py --query-type audio_to_text --model ./models/Step-Audio-2-mini
```
Ensure the model directory contains:
```
Step-Audio2-mini/
Step-Audio-2-mini/
├── config.json
├── model.safetensors (or pytorch_model.bin)
├── tokenizer.json
@@ -144,20 +143,18 @@ Transcribe audio to text:
```bash
# Quick start - Using default model and test audio
python end2end.py --query-type audio_to_text
python end2end.py --query-type audio_to_text \
--model stepfun-ai/Step-Audio-2-mini
# Using your own audio file (model will auto-download)
python end2end.py --query-type audio_to_text \
--audio-path /path/to/input.wav
# With specific model
python end2end.py --query-type audio_to_text \
--audio-path /path/to/input.wav \
--model stepfun-ai/Step-Audio2-7B
--model stepfun-ai/Step-Audio-2-mini
# With custom question
python end2end.py --query-type audio_to_text \
--audio-path input.wav \
--model stepfun-ai/Step-Audio-2-mini \
--question "What is the speaker saying?"
```
@@ -170,12 +167,8 @@ Convert text to speech:
```bash
# Basic TTS (model auto-downloads)
python end2end.py --query-type text_to_audio \
--text "Hello, this is a test of Step Audio 2 synthesis."
# With specific model
python end2end.py --query-type text_to_audio \
--text "Hello, this is a test." \
--model stepfun-ai/Step-Audio2-7B
--text "Hello, this is a test of Step Audio 2 synthesis." \
--model stepfun-ai/Step-Audio-2-mini
```
**Note**: Speaker voice is controlled by the `STEP_AUDIO2_DEFAULT_PROMPT_WAV` environment variable or the default prompt wav bundled with the model.
@@ -191,12 +184,8 @@ Process input audio and generate output audio:
```bash
# Basic voice conversion (model auto-downloads)
python end2end.py --query-type audio_to_audio \
--audio-path /path/to/source_audio.wav
# With specific model
python end2end.py --query-type audio_to_audio \
--audio-path source.wav \
--model stepfun-ai/Step-Audio2-7B
--audio-path /path/to/source_audio.wav \
--model stepfun-ai/Step-Audio-2-mini
```
This mode:
@@ -210,8 +199,14 @@ This mode:
```bash
# Use custom stage configuration
python end2end.py --query-type audio_to_text \
--model stepfun-ai/Step-Audio-2-mini \
--stage-configs-path /path/to/custom_config.yaml
# Use custom deploy configuration
python end2end.py --query-type audio_to_text \
--model stepfun-ai/Step-Audio-2-mini \
--deploy-config /path/to/step_audio_2_asr.yaml
# Multiple prompts (for batch testing)
python end2end.py --query-type audio_to_text \
--audio-path input.wav \
@@ -246,14 +241,14 @@ python end2end.py --query-type text_to_audio \
## Configuration
### Stage Configuration
### Deploy Configuration
The default configuration (`step_audio_2.yaml`) uses:
The default configuration (`vllm_omni/deploy/step_audio_2.yaml`) uses:
- **Stage 0 (Thinker)**: GPU 0, 80% memory
- **Stage 1 (Token2Wav)**: GPU 1, 30% memory
- **Stage 0 (Thinker)**: GPUs 0-1 with tensor parallel size 2, 70% memory
- **Stage 1 (Token2Wav)**: GPU 1, 20% memory
For **single GPU** setup, edit the config to use `devices: "0"` for both stages.
For **single GPU** setup, edit a deploy config copy to use `devices: "0"` for both stages.
### Sampling Parameters
@@ -351,7 +346,7 @@ Complete example from audio to final output:
# 1. ASR: Transcribe audio
python end2end.py --query-type audio_to_text \
--audio-path interview.wav \
--model ./models/Step-Audio2-7B \
--model ./models/Step-Audio-2-mini \
--output-dir ./outputs
# 2. Check the transcription
@@ -361,7 +356,7 @@ cat ./outputs/00000_text.txt
STEP_AUDIO2_DEFAULT_PROMPT_WAV=./speaker_samples/female_voice.wav \
python end2end.py --query-type text_to_audio \
--text "The quick brown fox jumps over the lazy dog" \
--model ./models/Step-Audio2-7B \
--model ./models/Step-Audio-2-mini \
--output-dir ./outputs
# 4. Listen to the result
examples/offline_inference/step_audio2/end2end.py
+37 -21
View File
@@ -232,32 +232,43 @@ def main(args):
raise ValueError(f"Unknown query type: {args.query_type}")
# Initialize vLLM-Omni with Step-Audio2
# Resolve stage config path
# Resolve deploy config path. Keep explicit stage_configs_path for legacy
# custom configs, but use bundled deploy configs by default.
deploy_config_path = None
stage_config_path = None
if args.stage_configs_path and args.deploy_config:
raise ValueError("--stage-configs-path and --deploy-config are mutually exclusive")
if args.stage_configs_path:
stage_config_path = args.stage_configs_path
elif args.deploy_config:
deploy_config_path = args.deploy_config
else:
configs_dir = Path(__file__).parent.parent.parent.parent / "vllm_omni/model_executor/stage_configs"
configs_dir = Path(__file__).parent.parent.parent.parent / "vllm_omni/deploy"
if args.query_type == "audio_to_text":
# ASR only needs Thinker (Stage 0), no Token2Wav
stage_config_path = str(configs_dir / "step_audio_2_asr.yaml")
if not os.path.exists(stage_config_path):
stage_config_path = str(configs_dir / "step_audio_2.yaml")
deploy_config_path = str(configs_dir / "step_audio_2_asr.yaml")
else:
# TTS/S2ST need both stages
stage_config_path = str(configs_dir / "step_audio_2.yaml")
deploy_config_path = str(configs_dir / "step_audio_2.yaml")
omni_llm = Omni(
model=model_name,
log_stats=args.enable_stats,
log_file=("step_audio2_pipeline.log" if args.enable_stats else None),
init_sleep_seconds=args.init_sleep_seconds,
batch_timeout=args.batch_timeout,
init_timeout=args.init_timeout,
shm_threshold_bytes=args.shm_threshold_bytes,
stage_configs_path=stage_config_path,
worker_backend=args.worker_backend,
ray_address=args.ray_address,
)
omni_kwargs = {
"model": model_name,
"log_stats": args.enable_stats,
"log_file": ("step_audio2_pipeline.log" if args.enable_stats else None),
"init_sleep_seconds": args.init_sleep_seconds,
"batch_timeout": args.batch_timeout,
"init_timeout": args.init_timeout,
"shm_threshold_bytes": args.shm_threshold_bytes,
"worker_backend": args.worker_backend,
"ray_address": args.ray_address,
"trust_remote_code": True,
}
if stage_config_path is not None:
omni_kwargs["stage_configs_path"] = stage_config_path
else:
omni_kwargs["deploy_config"] = deploy_config_path
omni_llm = Omni(**omni_kwargs)
# Configure sampling parameters for each stage
# Stage 0 (Thinker): Generate text and audio tokens
@@ -405,16 +416,21 @@ def parse_args():
help=(
"Model path (required). Can be:\n"
" - Local path: /path/to/step-audio-2\n"
" - HuggingFace ID: stepfun-ai/Step-Audio2-mini"
" - HuggingFace ID: stepfun-ai/Step-Audio-2-mini"
),
)
parser.add_argument(
"--stage-configs-path",
type=str,
default=None,
help="Path to stage config YAML file (default: use built-in step_audio_2.yaml)",
help="Path to a legacy stage config YAML file (default: use bundled deploy config)",
)
parser.add_argument(
"--deploy-config",
type=str,
default=None,
help="Path to a deploy config YAML file (default: use bundled Step-Audio2 deploy config)",
)
# Query configuration
parser.add_argument(
"--query-type",
examples/offline_inference/text_to_video/text_to_video.py
+7
View File
@@ -124,6 +124,11 @@ def parse_args() -> argparse.Namespace:
default=None,
help="Override model class name (e.g., LTX2TwoStagesVideoPipeline).",
)
parser.add_argument(
"--deploy-config",
default=None,
help="Optional deploy config YAML to use for pipeline-backed runs.",
)
parser.add_argument("--prompt", default="A serene lakeside sunrise with mist over the water.", help="Text prompt.")
parser.add_argument("--negative-prompt", default="", help="Negative prompt.")
parser.add_argument(
@@ -375,6 +380,8 @@ def main():
enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
profiler_config=args.profiler_config,
)
if args.deploy_config:
omni_kwargs["deploy_config"] = args.deploy_config
if args.boundary_ratio is not None:
omni_kwargs["boundary_ratio"] = args.boundary_ratio
if args.flow_shift is not None:
examples/online_serving/step_audio2/README.md
+6 -6
View File
@@ -10,15 +10,15 @@ Please refer to [README.md](../../../README.md)
```bash
# Async chunk mode (recommended — lower first-packet latency for TTS)
vllm serve stepfun-ai/Step-Audio2-mini --omni --port 8092 \
--stage-configs-path vllm_omni/model_executor/stage_configs/step_audio_2_async_chunk.yaml \
vllm serve stepfun-ai/Step-Audio-2-mini --omni --port 8092 \
--deploy-config vllm_omni/deploy/step_audio_2_async_chunk.yaml \
--trust-remote-code --enforce-eager
```
Sequential mode:
```bash
vllm serve stepfun-ai/Step-Audio2-mini --omni --port 8092 \
--stage-configs-path vllm_omni/model_executor/stage_configs/step_audio_2.yaml \
vllm serve stepfun-ai/Step-Audio-2-mini --omni --port 8092 \
--deploy-config vllm_omni/deploy/step_audio_2.yaml \
--trust-remote-code --enforce-eager
```
@@ -51,7 +51,7 @@ Or via curl:
```bash
curl -X POST http://localhost:8092/v1/audio/speech \
-H "Content-Type: application/json" \
-d '{"model":"stepfun-ai/Step-Audio2-mini","input":"你好世界","voice":"default"}' \
-d '{"model":"stepfun-ai/Step-Audio-2-mini","input":"你好世界","voice":"default"}' \
--output output.wav
```
@@ -139,7 +139,7 @@ Step-Audio2 uses the OpenAI-compatible chat completions API:
```json
{
"model": "stepfun-ai/Step-Audio2-mini",
"model": "stepfun-ai/Step-Audio-2-mini",
"messages": [
{
"role": "system",
examples/online_serving/step_audio2/openai_chat_completion_client.py
+2 -2
View File
@@ -9,7 +9,7 @@ Step-Audio2 is a two-stage model:
Usage examples:
# Start server first:
vllm serve stepfun-ai/Step-Audio2-mini --omni --port 8092
vllm serve stepfun-ai/Step-Audio-2-mini --omni --port 8092
# Audio to text (ASR)
python openai_chat_completion_client.py --query-type audio_to_text
@@ -311,7 +311,7 @@ def parse_args():
"--model",
"-m",
type=str,
default="stepfun-ai/Step-Audio2-mini",
default="stepfun-ai/Step-Audio-2-mini",
help="Model name",
)
parser.add_argument(
examples/online_serving/step_audio2/openai_speech_client.py
+1 -1
View File
@@ -87,7 +87,7 @@ def parse_args():
"--model",
"-m",
type=str,
default="stepfun-ai/Step-Audio2-mini",
default="stepfun-ai/Step-Audio-2-mini",
help="Model name/path",
)
parser.add_argument(
examples/online_serving/step_audio2/run_curl.sh
+1 -1
View File
@@ -110,7 +110,7 @@ output=$(curl -sS -X POST "${API_BASE}/v1/chat/completions" \
-H "Content-Type: application/json" \
-d @- <<EOF
{
"model": "stepfun-ai/Step-Audio2-mini",
"model": "stepfun-ai/Step-Audio-2-mini",
"sampling_params_list": $sampling_params_list,
"messages": [
{
pyproject.toml
-1
View File
@@ -151,7 +151,6 @@ include = ["vllm_omni*"]
[tool.setuptools.package-data]
"vllm_omni" = ["_version.py", "py.typed"]
"vllm_omni.model_executor.stage_configs" = ["*.yaml"]
[tool.setuptools_scm]
# no extra settings needed, presence enables setuptools-scm
tests/model_executor/stage_input_processors/test_step_audio2_async_chunk.py
+14 -16
View File
@@ -5,7 +5,6 @@ from collections import defaultdict
from types import SimpleNamespace
import pytest
import torch
from vllm_omni.model_executor.models.step_audio2.step_audio2_constants import (
DEFAULT_STREAM_CONFIG,
@@ -43,14 +42,14 @@ def test_step_audio2_async_chunk_uses_decode_only_tokens_not_prompt_history():
payload = thinker2token2wav_async_chunk(
transfer_manager=transfer_manager,
pooling_output=None,
multimodal_output=None,
request=request,
)
assert payload is not None
assert payload["code_predictor_codes"] == [1, 2]
assert payload["left_context_size"] == 1
assert payload["finished"].item() is True
assert payload.codes.audio.tolist() == [1, 2]
assert payload.meta.left_context_size == 1
assert payload.meta.finished.item() is True
def test_step_audio2_async_chunk_returns_none_when_not_enough_tokens():
@@ -68,7 +67,7 @@ def test_step_audio2_async_chunk_returns_none_when_not_enough_tokens():
payload = thinker2token2wav_async_chunk(
transfer_manager=transfer_manager,
pooling_output=None,
multimodal_output=None,
request=request,
)
@@ -91,14 +90,14 @@ def test_step_audio2_async_chunk_emits_non_last_chunk_and_advances_consumed_by_c
payload = thinker2token2wav_async_chunk(
transfer_manager=transfer_manager,
pooling_output=None,
multimodal_output=None,
request=request,
)
assert payload is not None
assert payload["left_context_size"] == 0
assert payload["finished"].item() is False
assert payload["code_predictor_codes"] == list(range(required))
assert payload.meta.left_context_size == 0
assert payload.meta.finished.item() is False
assert payload.codes.audio.tolist() == list(range(required))
assert transfer_manager.code_prompt_token_ids["rid-ready"] == list(range(chunk_size))
@@ -121,12 +120,11 @@ def test_step_audio2_async_chunk_emits_eof_when_finished_with_no_remaining_audio
payload = thinker2token2wav_async_chunk(
transfer_manager=transfer_manager,
pooling_output=None,
multimodal_output=None,
request=request,
)
assert payload == {
"code_predictor_codes": [],
"left_context_size": 1,
"finished": torch.tensor(True, dtype=torch.bool),
}
assert payload is not None
assert payload.codes.audio.numel() == 0
assert payload.meta.left_context_size == 1
assert payload.meta.finished.item() is True
tests/test_config_factory.py
+76 -5
View File
@@ -1155,11 +1155,17 @@ stages:
assert runtime_config == expected_runtime_config
@pytest.mark.parametrize(
("deploy_name", "pipeline_name", "stage_count", "final_output_type"),
("deploy_name", "pipeline_name", "stage_count", "final_output_type", "declared_pipeline"),
[
("mammoth_moda2.yaml", "mammoth_moda2", 2, "image"),
("mammoth_moda2_ar.yaml", "mammoth_moda2_ar", 1, "text"),
("omnivoice.yaml", "omnivoice", 1, "audio"),
("mammoth_moda2.yaml", "mammoth_moda2", 2, "image", "mammoth_moda2"),
("mammoth_moda2_ar.yaml", "mammoth_moda2_ar", 1, "text", "mammoth_moda2_ar"),
("omnivoice.yaml", "omnivoice", 1, "audio", "omnivoice"),
("mimo_audio.yaml", "mimo_audio", 2, "audio", None),
("step_audio_2.yaml", "step_audio_2", 2, "audio", None),
("step_audio_2_asr.yaml", "step_audio_2_asr", 1, "text", "step_audio_2_asr"),
("step_audio_2_async_chunk.yaml", "step_audio_2", 2, "audio", None),
("hunyuan_video_15.yaml", "hunyuan_video_15", 1, "video", None),
("wan2_2_ti2v.yaml", "wan2_2_ti2v", 1, "video", None),
],
)
def test_load_new_registry_backed_deploy_configs(
@@ -1168,10 +1174,11 @@ stages:
pipeline_name: str,
stage_count: int,
final_output_type: str,
declared_pipeline: str | None,
):
deploy_path = Path(get_deploy_config_path(deploy_name))
deploy = load_deploy_config(deploy_path)
assert deploy.pipeline == pipeline_name
assert deploy.pipeline == declared_pipeline
with patch("vllm_omni.platforms.current_omni_platform") as platform:
platform.device_name = "cuda"
@@ -1179,6 +1186,70 @@ stages:
assert len(stages) == stage_count
assert stages[-1].final_output is True
assert stages[-1].final_output_type == final_output_type
if deploy.trust_remote_code is not None:
assert {s.yaml_engine_args.get("trust_remote_code") for s in stages} == {deploy.trust_remote_code}
@pytest.mark.parametrize(
("config_json", "model_index", "expected_pipeline"),
[
({"model_type": "step_audio_2"}, None, "step_audio_2"),
(None, {"_class_name": "HunyuanVideo15Pipeline"}, "hunyuan_video_15"),
(None, {"_class_name": "WanPipeline"}, "wan2_2_ti2v"),
],
)
def test_migrated_models_are_discovered_without_explicit_deploy(
self,
config_json,
model_index,
expected_pipeline,
):
def get_model_file(filename, _model, revision=None):
del revision
if filename == "config.json":
return config_json
if filename == "model_index.json":
return model_index
return None
with (
patch.object(StageConfigFactory, "get_hf_config", return_value=None),
patch("vllm_omni.config.config_factory.get_hf_file_to_dict", side_effect=get_model_file),
):
pipeline = StageConfigFactory.get_pipeline_config(
model="/models/unrelated-checkpoint-name",
trust_remote_code=False,
)
assert pipeline is not None
assert pipeline.model_type == expected_pipeline
@pytest.mark.parametrize("deploy_name", ["step_audio_2.yaml", "step_audio_2_async_chunk.yaml"])
def test_step_audio2_deploy_configs_fit_two_gpus(self, deploy_name):
deploy = load_deploy_config(Path(get_deploy_config_path(deploy_name)))
assert deploy.stages[0].devices == "0,1"
assert deploy.stages[0].tensor_parallel_size == 2
assert deploy.stages[1].devices == "1"
def test_step_audio2_dispatches_sync_and_async_chunk_processors(self):
pipeline = StageConfigFactory.resolve_pipeline_config("step_audio_2")
assert isinstance(pipeline, PipelineConfig)
sync_stages = merge_pipeline_deploy(pipeline, DeployConfig(async_chunk=False))
assert "custom_process_next_stage_input_func" not in sync_stages[0].yaml_engine_args
assert sync_stages[1].custom_process_input_func.endswith("thinker2token2wav")
async_stages = merge_pipeline_deploy(pipeline, DeployConfig(async_chunk=True))
assert (
async_stages[0]
.yaml_engine_args["custom_process_next_stage_input_func"]
.endswith("thinker2token2wav_async_chunk")
)
assert async_stages[1].custom_process_input_func is None
def test_no_bundled_legacy_stage_config_yamls(self):
stage_config_dir = Path(__file__).parent.parent / "vllm_omni" / "model_executor" / "stage_configs"
assert not list(stage_config_dir.glob("*.yaml"))
def test_merge_pipeline_deploy(self):
deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml"
vllm_omni/config/pipeline_registry.py
+10
View File
@@ -54,6 +54,7 @@ from vllm_omni.model_executor.models.hunyuan_image3.pipeline import (
HUNYUAN_IMAGE3_DIT_PIPELINE,
HUNYUAN_IMAGE3_PIPELINE,
)
from vllm_omni.model_executor.models.hunyuan_video.pipeline import HUNYUAN_VIDEO_15_PIPELINE
from vllm_omni.model_executor.models.indextts2.pipeline import INDEXTTS2_PIPELINE
from vllm_omni.model_executor.models.lance.pipeline import LANCE_PIPELINE
from vllm_omni.model_executor.models.mammoth_moda2.pipeline import (
@@ -85,8 +86,13 @@ from vllm_omni.model_executor.models.qwen2_5_omni.pipeline import (
)
from vllm_omni.model_executor.models.qwen3_omni.pipeline import resolve_qwen3_omni_pipeline
from vllm_omni.model_executor.models.qwen3_tts.pipeline import QWEN3_TTS_PIPELINE
from vllm_omni.model_executor.models.step_audio2.pipeline import (
STEP_AUDIO2_ASR_PIPELINE,
STEP_AUDIO2_PIPELINE,
)
from vllm_omni.model_executor.models.voxcpm2.pipeline import VOXCPM2_PIPELINE
from vllm_omni.model_executor.models.voxtral_tts.pipeline import VOXTRAL_TTS_PIPELINE
from vllm_omni.model_executor.models.wan2_2.pipeline import WAN2_2_TI2V_PIPELINE
logger = init_logger(__name__)
@@ -99,6 +105,8 @@ OMNI_PIPELINES: dict[str, PipelineConfig | PipelineResolverFunc] = {
"qwen2_5_omni_thinker_only": QWEN2_5_OMNI_THINKER_ONLY_PIPELINE,
"qwen3_omni_moe": resolve_qwen3_omni_pipeline,
"qwen3_tts": QWEN3_TTS_PIPELINE,
"step_audio_2": STEP_AUDIO2_PIPELINE,
"step_audio_2_asr": STEP_AUDIO2_ASR_PIPELINE,
"covo_audio": COVO_AUDIO_PIPELINE,
"bagel": BAGEL_PIPELINE,
"bagel_think": BAGEL_THINK_PIPELINE,
@@ -110,6 +118,8 @@ OMNI_PIPELINES: dict[str, PipelineConfig | PipelineResolverFunc] = {
"hunyuan_image_3_moe": HUNYUAN_IMAGE3_PIPELINE,
"hunyuan_image3_ar": HUNYUAN_IMAGE3_AR_PIPELINE,
"hunyuan_image3_dit": HUNYUAN_IMAGE3_DIT_PIPELINE,
"hunyuan_video_15": HUNYUAN_VIDEO_15_PIPELINE,
"wan2_2_ti2v": WAN2_2_TI2V_PIPELINE,
"voxcpm2": VOXCPM2_PIPELINE,
"cosyvoice3": COSYVOICE3_PIPELINE,
"mimo_audio": MIMO_AUDIO_PIPELINE,
vllm_omni/deploy/hunyuan_video_15.yaml
+18
View File
@@ -0,0 +1,18 @@
# HunyuanVideo-1.5 deploy.
#
# Single GPU by default. Increase tensor_parallel_size and devices for multi-GPU TP.
async_chunk: false
trust_remote_code: true
distributed_executor_backend: mp
stages:
- stage_id: 0
devices: "0"
max_num_seqs: 1
tensor_parallel_size: 1
enforce_eager: true
# HunyuanVideo-1.5 peak VRAM is dominated by the single-shot VAE decode of
# the full video, so tiling the decode is the main memory lever.
vae_use_tiling: true
default_sampling_params:
seed: 42
vllm_omni/deploy/step_audio_2.yaml
+39
View File
@@ -0,0 +1,39 @@
# Step-Audio2 deploy: thinker -> token2wav.
#
# Sequential mode waits for the full Thinker output before running Token2Wav.
async_chunk: false
trust_remote_code: true
enable_prefix_caching: false
stages:
- stage_id: 0
devices: "0,1"
tensor_parallel_size: 2
max_num_seqs: 1
gpu_memory_utilization: 0.7
enforce_eager: true
async_scheduling: false
max_num_batched_tokens: 8192
default_sampling_params:
temperature: 0.7
top_p: 0.9
top_k: -1
max_tokens: 1024
seed: 42
repetition_penalty: 1.05
- stage_id: 1
devices: "1"
max_num_seqs: 1
gpu_memory_utilization: 0.2
enforce_eager: true
async_scheduling: false
max_num_batched_tokens: 4096
hf_overrides:
architectures: ["StepAudio2Token2WavForConditionalGeneration"]
default_sampling_params:
temperature: 0.0
top_p: 1.0
top_k: -1
max_tokens: 1
seed: 42
vllm_omni/deploy/step_audio_2_asr.yaml
+22
View File
@@ -0,0 +1,22 @@
# Step-Audio2 ASR deploy: thinker only.
pipeline: step_audio_2_asr
async_chunk: false
trust_remote_code: true
enable_prefix_caching: false
stages:
- stage_id: 0
devices: "0,1"
tensor_parallel_size: 2
max_num_seqs: 1
gpu_memory_utilization: 0.8
enforce_eager: true
async_scheduling: false
max_num_batched_tokens: 2048
default_sampling_params:
temperature: 0.7
top_p: 0.9
top_k: -1
max_tokens: 1024
seed: 42
repetition_penalty: 1.05
vllm_omni/deploy/step_audio_2_async_chunk.yaml
+37
View File
@@ -0,0 +1,37 @@
# Step-Audio2 deploy: thinker -> token2wav with async chunk streaming.
async_chunk: true
trust_remote_code: true
enable_prefix_caching: false
stages:
- stage_id: 0
devices: "0,1"
tensor_parallel_size: 2
max_num_seqs: 1
gpu_memory_utilization: 0.7
enforce_eager: true
async_scheduling: false
max_num_batched_tokens: 1024
default_sampling_params:
temperature: 0.7
top_p: 0.9
top_k: -1
max_tokens: 1024
seed: 42
repetition_penalty: 1.05
- stage_id: 1
devices: "1"
max_num_seqs: 1
gpu_memory_utilization: 0.2
enforce_eager: true
async_scheduling: false
max_num_batched_tokens: 2048
hf_overrides:
architectures: ["StepAudio2Token2WavForConditionalGeneration"]
default_sampling_params:
temperature: 0.0
top_p: 1.0
top_k: -1
max_tokens: 1
seed: 42
vllm_omni/deploy/wan2_2_ti2v.yaml
+17
View File
@@ -0,0 +1,17 @@
# Wan2.2 TI2V-5B deploy.
#
# Single GPU by default. Increase tensor_parallel_size and devices for A14B MoE variants.
async_chunk: false
trust_remote_code: true
distributed_executor_backend: mp
stages:
- stage_id: 0
devices: "0"
max_num_seqs: 1
tensor_parallel_size: 1
enforce_eager: true
# Keep VAE tiling on by default for higher resolutions / longer clips.
vae_use_tiling: true
default_sampling_params:
seed: 42
vllm_omni/model_executor/models/hunyuan_video/__init__.py
+2
View File
@@ -0,0 +1,2 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
vllm_omni/model_executor/models/hunyuan_video/pipeline.py
+26
View File
@@ -0,0 +1,26 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""HunyuanVideo-1.5 pipeline topology."""
from vllm_omni.config.stage_config import (
PipelineConfig,
StageExecutionType,
StagePipelineConfig,
)
HUNYUAN_VIDEO_15_PIPELINE = PipelineConfig(
model_type="hunyuan_video_15",
model_arch="HunyuanVideo15Pipeline",
diffusers_class_name="HunyuanVideo15Pipeline",
stages=(
StagePipelineConfig(
stage_id=0,
model_stage="dit",
execution_type=StageExecutionType.DIFFUSION,
input_sources=(),
final_output=True,
final_output_type="video",
model_arch="HunyuanVideo15Pipeline",
),
),
)
vllm_omni/model_executor/models/step_audio2/__init__.py
-3
View File
@@ -1,3 +0,0 @@
from .step_audio2 import StepAudio2ForConditionalGeneration
__all__ = ["StepAudio2ForConditionalGeneration"]
vllm_omni/model_executor/models/step_audio2/pipeline.py
+71
View File
@@ -0,0 +1,71 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Step-Audio2 pipeline topologies (frozen)."""
from vllm_omni.config.stage_config import (
PipelineConfig,
StageExecutionType,
StagePipelineConfig,
)
_PROC = "vllm_omni.model_executor.stage_input_processors.step_audio2"
_HF_ARCHITECTURES = (
"StepAudio2ForCausalLM",
"StepAudio2ForConditionalGeneration",
)
STEP_AUDIO2_PIPELINE = PipelineConfig(
model_type="step_audio_2",
model_arch="StepAudio2ForConditionalGeneration",
hf_architectures=_HF_ARCHITECTURES,
stages=(
StagePipelineConfig(
stage_id=0,
model_stage="step_audio2_thinker",
execution_type=StageExecutionType.LLM_AR,
input_sources=(),
final_output=True,
final_output_type="text",
owns_tokenizer=True,
requires_multimodal_data=True,
engine_output_type="latent",
model_arch="StepAudio2ThinkerForConditionalGeneration",
async_chunk_process_next_stage_input_func=f"{_PROC}.thinker2token2wav_async_chunk",
sampling_constraints={"detokenize": True},
),
StagePipelineConfig(
stage_id=1,
model_stage="token2wav",
execution_type=StageExecutionType.LLM_GENERATION,
input_sources=(0,),
final_output=True,
final_output_type="audio",
engine_output_type="audio",
model_arch="StepAudio2Token2WavModel",
sync_process_input_func=f"{_PROC}.thinker2token2wav",
sampling_constraints={"detokenize": False},
),
),
)
STEP_AUDIO2_ASR_PIPELINE = PipelineConfig(
model_type="step_audio_2_asr",
model_arch="StepAudio2ForConditionalGeneration",
stages=(
StagePipelineConfig(
stage_id=0,
model_stage="step_audio2_thinker",
execution_type=StageExecutionType.LLM_AR,
input_sources=(),
final_output=True,
final_output_type="text",
owns_tokenizer=True,
requires_multimodal_data=True,
engine_output_type="text",
model_arch="StepAudio2ThinkerForConditionalGeneration",
sampling_constraints={"detokenize": True},
),
),
)
vllm_omni/model_executor/models/wan2_2/__init__.py
+2
View File
@@ -0,0 +1,2 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
vllm_omni/model_executor/models/wan2_2/pipeline.py
+26
View File
@@ -0,0 +1,26 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Wan2.2 TI2V pipeline topology."""
from vllm_omni.config.stage_config import (
PipelineConfig,
StageExecutionType,
StagePipelineConfig,
)
WAN2_2_TI2V_PIPELINE = PipelineConfig(
model_type="wan2_2_ti2v",
model_arch="WanPipeline",
diffusers_class_name="WanPipeline",
stages=(
StagePipelineConfig(
stage_id=0,
model_stage="dit",
execution_type=StageExecutionType.DIFFUSION,
input_sources=(),
final_output=True,
final_output_type="video",
model_arch="WanPipeline",
),
),
)
vllm_omni/model_executor/stage_configs/__init__.py
View File
Whitespace-only changes.
vllm_omni/model_executor/stage_configs/hunyuan_video_15_dit_fp8.yaml
-38
View File
@@ -1,38 +0,0 @@
# Stage config for running HunyuanVideo-1.5 DiT with ModelOpt FP8 auto-detect.
# Single GPU. Bump `tensor_parallel_size` and `devices` for multi-GPU TP.
#
# Use with a ModelOpt FP8 checkpoint (e.g. produced by
# examples/quantization/quantize_hunyuanvideo_15_modelopt_fp8.py).
stage_args:
- stage_id: 0
stage_type: diffusion
runtime:
devices: "0"
max_batch_size: 1
engine_args:
model_stage: dit
model_class_name: HunyuanVideo15Pipeline
max_num_seqs: 1
enforce_eager: true
trust_remote_code: true
distributed_executor_backend: "mp"
# HunyuanVideo-1.5 peak VRAM is dominated by the single-shot VAE decode of
# the full video (the DiT contributes little), so tiling the decode is the
# main memory lever — far more than FP8 (e.g. 480x832x33: ~117 -> ~47 GiB).
# Default it on here; override with --vae-use-tiling/false if you have headroom.
vae_use_tiling: true
parallel_config:
tensor_parallel_size: 1
final_output: true
final_output_type: video
is_comprehension: false
default_sampling_params:
seed: 42
runtime:
enabled: true
defaults:
window_size: -1
max_inflight: 1
vllm_omni/model_executor/stage_configs/mimo_audio.yaml
-61
View File
@@ -1,61 +0,0 @@
# MiMo Audio deploy: 2-stage thinker+talker → code2wav.
#
# Default mode is async-chunk streaming on a single GPU (both stages on
# device 0) through SharedMemoryConnector. For the legacy 2-GPU sync
# pipeline (stage 1 on a second card), pass ``--no-async-chunk
# --stage-overrides '{"1": {"devices": "1", "max_model_len": 18192, "max_num_batched_tokens": 18192}}'``
# to the serve CLI.
async_chunk: true
dtype: bfloat16
connectors:
connector_of_shared_memory:
name: SharedMemoryConnector
extra:
shm_threshold_bytes: 365536
codec_streaming: true
connector_get_sleep_s: 0.001
connector_get_max_wait_first_chunk: 3000
connector_get_max_wait: 300
codec_chunk_frames: 30
codec_left_context_frames: 40
stages:
- stage_id: 0
max_num_seqs: 1
gpu_memory_utilization: 0.5
enforce_eager: true
trust_remote_code: true
enable_prefix_caching: false
max_num_batched_tokens: 8192
max_model_len: 8192
devices: "0"
output_connectors:
to_stage_1: connector_of_shared_memory
default_sampling_params:
temperature: 0.6
top_p: 0.95
top_k: 50
max_tokens: 18192
seed: 42
repetition_penalty: 1.1
- stage_id: 1
max_num_seqs: 1
gpu_memory_utilization: 0.4
enforce_eager: true
trust_remote_code: true
enable_prefix_caching: false
async_scheduling: false
max_num_batched_tokens: 8192
max_model_len: 8192
devices: "0"
input_connectors:
from_stage_0: connector_of_shared_memory
default_sampling_params:
temperature: 0.0
top_p: 1.0
top_k: -1
max_tokens: 18192
seed: 42
vllm_omni/model_executor/stage_configs/step_audio_2.yaml
-122
View File
@@ -1,122 +0,0 @@
# Stage config for Step-Audio2 with vLLM-Omni
#
# Two-stage pipeline:
# Stage 0 (Thinker): Audio understanding → Text + Audio tokens
# Stage 1 (Token2Wav): Audio tokens → Waveform (24kHz)
#
# Hardware requirement: 1-2 GPUs (can run on single GPU with devices: "0" for both)
# Model configuration reference (for documentation)
# Actual values are set in hf_overrides below
#
# Token configuration:
# text_max: 151688 # Maximum text token ID (inclusive)
# audio_start: 151696 # First audio token ID in absolute vocabulary
# audio_vocab_size: 6562 # Total number of audio tokens
# audio_eos: 6561 # Audio EOS token ID (relative to audio_start)
# audio_patch_token_id: 151690 # <audio_patch> placeholder token ID
#
# Encoder configuration:
# n_mels: 128 # Number of mel frequency bins
# n_audio_ctx: 1500 # Audio context length
# n_audio_state: 512 # Audio encoder hidden state dimension
# n_audio_head: 8 # Number of attention heads
# n_audio_layer: 6 # Number of encoder layers
# kernel_size: 3 # Convolution kernel size for adapter
# adapter_stride: 2 # Stride for adapter convolution
stage_args:
# Stage 0: Thinker (Audio Understanding + Token Generation)
- stage_id: 0
runtime:
process: true # Run in separate process
devices: "0,1,2,3" # Use 3 GPUs for tensor parallelism
max_batch_size: 1 # Single sample at a time
engine_args:
model_stage: step_audio2_thinker
model_arch: StepAudio2ThinkerForConditionalGeneration
tensor_parallel_size: 4 # Split model across 3 GPUs
# NOTE: All default values are defined in step_audio2_constants.py (DEFAULT_MODEL_CONFIG)
# Only override here if your model differs from the defaults
#
# Default values (for reference, from constants.py):
# hidden_size: 4096 # Qwen2-7B (change to 1536 for 1.5B, 896 for 0.5B)
# text_max: 151688
# audio_start: 151696
# audio_vocab_size: 6562
# audio_eos: 6561
# audio_patch_token_id: 151690
# n_mels: 128
# n_audio_ctx: 1500
# n_audio_state: 512
# n_audio_head: 8
# n_audio_layer: 6
# kernel_size: 3
# adapter_stride: 2
#
# Uncomment and modify only if needed:
# hidden_size: 4096
worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
gpu_memory_utilization: 0.8
enforce_eager: true # Required for now
trust_remote_code: true
enable_prefix_caching: false
max_num_batched_tokens: 8192
engine_output_type: text # Outputs text + audio tokens
is_comprehension: true # This stage does comprehension
final_output: true # Can be final output (for ASR mode)
final_output_type: text
default_sampling_params:
temperature: 0.7
top_p: 0.9
top_k: -1
max_tokens: 1024
seed: 42
detokenize: true
repetition_penalty: 1.05
# Stage 1: Token2Wav (Audio Synthesis)
- stage_id: 1
runtime:
process: true
devices: "3" # Use GPU 3 (separate from Stage 0's 0-2)
max_batch_size: 1
engine_args:
model_stage: token2wav
model_arch: StepAudio2Token2WavModel
hf_overrides:
architectures: ["StepAudio2Token2WavForConditionalGeneration"]
worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
gpu_memory_utilization: 0.3 # Token2Wav needs less memory
enforce_eager: true
trust_remote_code: true
async_scheduling: false
enable_prefix_caching: false
max_num_batched_tokens: 4096
engine_output_type: audio # Outputs waveform
engine_input_source: [0] # Takes input from Stage 0
custom_process_input_func: vllm_omni.model_executor.stage_input_processors.step_audio2.thinker2token2wav
final_output: true # Final output stage
final_output_type: audio
default_sampling_params:
temperature: 0.0
top_p: 1.0
top_k: -1
max_tokens: 1
seed: 42
detokenize: false
# Runtime configuration
runtime:
enabled: true
defaults:
window_size: -1 # Wait for full input before triggering
max_inflight: 1 # Process one at a time
edges:
- from: 0 # Thinker → Token2Wav
to: 1
window_size: -1 # Wait for complete output from Thinker
vllm_omni/model_executor/stage_configs/step_audio_2_asr.yaml
-35
View File
@@ -1,35 +0,0 @@
# Stage config for Step-Audio2 ASR (speech-to-text only)
#
# Single-stage pipeline: only Thinker, no Token2Wav.
# Use this for audio_to_text tasks where no audio output is needed.
stage_args:
- stage_id: 0
stage_type: llm
runtime:
process: true
devices: "0,1"
max_batch_size: 1
engine_args:
model_stage: step_audio2_thinker
model_arch: StepAudio2ThinkerForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
gpu_memory_utilization: 0.8
enforce_eager: true
trust_remote_code: true
enable_prefix_caching: false
max_num_batched_tokens: 2048
tensor_parallel_size: 2
engine_output_type: text
is_comprehension: true
final_output: true
final_output_type: text
default_sampling_params:
temperature: 0.7
top_p: 0.9
top_k: -1
max_tokens: 1024
seed: 42
detokenize: true
repetition_penalty: 1.05
vllm_omni/model_executor/stage_configs/step_audio_2_async_chunk.yaml
-78
View File
@@ -1,78 +0,0 @@
# Stage config for Step-Audio2 with async_chunk support
#
# Two-stage pipeline with streaming between stages:
# Stage 0 (Thinker): Audio understanding → Text + Audio tokens (streamed in chunks)
# Stage 1 (Token2Wav): Audio tokens → Waveform (24kHz)
#
# Key difference from step_audio_2.yaml:
# - async_chunk: true enables streaming from Thinker to Token2Wav
# - Thinker uses custom_process_next_stage_input_func to send audio token chunks
# - Token2Wav receives chunks incrementally instead of waiting for full output
#
# Hardware requirement: 1-2 GPUs (can run on single GPU with devices: "0" for both)
async_chunk: true
stage_args:
# Stage 0: Thinker (Audio Understanding + Token Generation)
- stage_id: 0
stage_type: llm
runtime:
process: true
devices: "0,1"
max_batch_size: 1
engine_args:
model_stage: step_audio2_thinker
model_arch: StepAudio2ThinkerForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
gpu_memory_utilization: 0.8
enforce_eager: true
trust_remote_code: true
enable_prefix_caching: false
max_num_batched_tokens: 1024
tensor_parallel_size: 2
engine_output_type: latent
custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.step_audio2.thinker2token2wav_async_chunk
is_comprehension: true
final_output: true
final_output_type: text
default_sampling_params:
temperature: 0.7
top_p: 0.9
top_k: -1
max_tokens: 1024
seed: 42
detokenize: true
repetition_penalty: 1.05
# Stage 1: Token2Wav (Audio Synthesis)
- stage_id: 1
stage_type: llm
runtime:
process: true
devices: "3"
max_batch_size: 1
engine_args:
model_stage: token2wav
model_arch: StepAudio2Token2WavModel
worker_type: generation
hf_overrides:
architectures: ["StepAudio2Token2WavForConditionalGeneration"]
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
gpu_memory_utilization: 0.3
enforce_eager: true
trust_remote_code: true
async_scheduling: false
enable_prefix_caching: false
max_num_batched_tokens: 2048
engine_output_type: audio
engine_input_source: [0]
final_output: true
final_output_type: audio
default_sampling_params:
temperature: 0.0
top_p: 1.0
top_k: -1
max_tokens: 1
seed: 42
detokenize: false
vllm_omni/model_executor/stage_configs/wan2_2_ti2v_dit_fp8.yaml
-41
View File
@@ -1,41 +0,0 @@
# Stage config for running Wan2.2 TI2V-5B DiT with ModelOpt FP8 auto-detect.
# Single GPU (TI2V-5B fits 80GB BF16; FP8 drops by ~half).
# For the A14B MoE variants, bump `tensor_parallel_size` and `devices`.
#
# Use with a ModelOpt FP8 checkpoint (e.g. produced by
# examples/quantization/quantize_wan2_2_modelopt_fp8.py).
stage_args:
- stage_id: 0
stage_type: diffusion
runtime:
devices: "0"
max_batch_size: 1
engine_args:
model_stage: dit
# Wan2.2 TI2V-5B is served by the base Wan22Pipeline, which auto-detects
# expand_timesteps (TI2V-5B mode) from model_index.json. (The standalone
# Wan22TI2VPipeline was consolidated into Wan22Pipeline upstream.)
model_class_name: Wan22Pipeline
max_num_seqs: 1
enforce_eager: true
trust_remote_code: true
distributed_executor_backend: "mp"
# Tile the VAE decode by default. Wan2.2's 480p peak is DiT-dominated (so
# FP8 is the main lever here), but tiling adds headroom for higher
# resolutions / longer clips at negligible cost. Override with --vae-use-tiling/false.
vae_use_tiling: true
parallel_config:
tensor_parallel_size: 1
final_output: true
final_output_type: video
is_comprehension: false
default_sampling_params:
seed: 42
runtime:
enabled: true
defaults:
window_size: -1
max_inflight: 1
vllm_omni/model_executor/stage_input_processors/mimo_audio.py
+2 -2
View File
@@ -17,8 +17,8 @@ logger = init_logger(__name__)
# sequence fed to stage-1 must not exceed this, otherwise gpu_input_batch
# add_request will fail with a broadcast error when copying prompt_token_ids
# into token_ids_cpu. Keep in sync with the stage-1 ``max_model_len`` in
# ``vllm_omni/model_executor/stage_configs/mimo_audio.yaml`` and the offline
# example ``examples/offline_inference/mimo_audio/end2end.py``.
# ``vllm_omni/deploy/mimo_audio.yaml`` and the offline example
# ``examples/offline_inference/mimo_audio/end2end.py``.
MAX_CODE2WAV_TOKENS = 18192
# Minimum safe values for codec streaming parameters.
vllm_omni/model_executor/stage_input_processors/step_audio2.py
+27 -21
View File
@@ -7,6 +7,7 @@ import torch
from vllm.inputs import TextPrompt
from vllm.logger import init_logger
from vllm_omni.data_entry_keys import CodesStruct, MetaStruct, OmniPayloadStruct
from vllm_omni.engine import OmniEngineCoreRequest
from vllm_omni.inputs.data import OmniTokensPrompt
from vllm_omni.model_executor.models.step_audio2.step_audio2_constants import (
@@ -35,14 +36,14 @@ def _ensure_list(x):
def thinker2token2wav_async_chunk(
transfer_manager: Any,
pooling_output: dict[str, Any],
multimodal_output: dict[str, Any] | None,
request: OmniEngineCoreRequest,
is_finished: bool = False,
) -> dict[str, Any] | None:
) -> OmniPayloadStruct | None:
"""
Async chunk processor: stream audio tokens from Thinker to Token2Wav.
Unlike Qwen3-Omni which passes hidden states via pooling_output,
Unlike Qwen3-Omni which passes hidden states via multimodal_output,
Step Audio2 generates audio tokens as part of the LLM autoregressive
output (token IDs >= audio_start). This processor extracts those audio
tokens from the running token stream and sends them in chunks.
@@ -62,14 +63,13 @@ def thinker2token2wav_async_chunk(
Args:
transfer_manager: OmniChunkTransferAdapter instance (has
``code_prompt_token_ids`` defaultdict for state tracking).
pooling_output: Pooler output dict (unused — audio tokens are in
multimodal_output: Multimodal output dict (unused — audio tokens are in
the token ID stream, not in hidden states).
request: Current engine request with access to all_token_ids.
is_finished: Whether the upstream request has finished generating.
Returns:
dict with framework-compatible fields, or None if the chunk is
not yet ready.
Structured connector payload, or None if the chunk is not yet ready.
"""
audio_start = DEFAULT_TOKEN_CONFIG.audio_start
audio_eos = DEFAULT_TOKEN_CONFIG.audio_eos
@@ -100,18 +100,22 @@ def thinker2token2wav_async_chunk(
# Last chunk: send all remaining tokens
if available <= 0:
# No audio tokens at all (text-only response) — send EOF marker
return {
"code_predictor_codes": [],
"left_context_size": 1, # 1 = last chunk
"finished": torch.tensor(True, dtype=torch.bool),
}
return OmniPayloadStruct(
codes=CodesStruct(audio=torch.empty(0, dtype=torch.long)),
meta=MetaStruct(
left_context_size=1, # 1 = last chunk
finished=torch.tensor(True, dtype=torch.bool),
),
)
remaining_tokens = audio_tokens[consumed:]
transfer_manager.code_prompt_token_ids[request_id].extend(remaining_tokens)
return {
"code_predictor_codes": remaining_tokens,
"left_context_size": 1, # 1 = last chunk
"finished": torch.tensor(True, dtype=torch.bool),
}
return OmniPayloadStruct(
codes=CodesStruct(audio=torch.tensor(remaining_tokens, dtype=torch.long)),
meta=MetaStruct(
left_context_size=1, # 1 = last chunk
finished=torch.tensor(True, dtype=torch.bool),
),
)
else:
# Non-last chunk: need chunk_size + pre_lookahead_len tokens
required = chunk_size + pre_lookahead_len
@@ -120,11 +124,13 @@ def thinker2token2wav_async_chunk(
chunk_tokens = audio_tokens[consumed : consumed + required]
# Only consume chunk_size tokens; the lookahead portion is re-used
transfer_manager.code_prompt_token_ids[request_id].extend(audio_tokens[consumed : consumed + chunk_size])
return {
"code_predictor_codes": chunk_tokens,
"left_context_size": 0, # 0 = not last chunk
"finished": torch.tensor(False, dtype=torch.bool),
}
return OmniPayloadStruct(
codes=CodesStruct(audio=torch.tensor(chunk_tokens, dtype=torch.long)),
meta=MetaStruct(
left_context_size=0, # 0 = not last chunk
finished=torch.tensor(False, dtype=torch.bool),
),
)
def thinker2token2wav(
vllm_omni/platforms/cuda/platform.py
+1 -1
View File
@@ -36,7 +36,7 @@ class CudaOmniPlatform(OmniPlatform, CudaPlatformBase):
@classmethod
def get_default_stage_config_path(cls) -> str:
return "vllm_omni/model_executor/stage_configs"
return "vllm_omni/deploy"
@classmethod
def has_flash_attn_package(cls) -> bool:
vllm_omni/platforms/musa/platform.py
+1 -1
View File
@@ -32,7 +32,7 @@ class MUSAOmniPlatform(OmniPlatform, MUSAPlatformBase):
@classmethod
def get_default_stage_config_path(cls) -> str:
return "vllm_omni/model_executor/stage_configs"
return "vllm_omni/deploy"
@classmethod
def has_flash_attn_package(cls) -> bool:
vllm_omni/platforms/npu/platform.py
+1 -1
View File
@@ -75,7 +75,7 @@ class NPUOmniPlatform(OmniPlatform, NPUPlatform):
@classmethod
def get_default_stage_config_path(cls) -> str:
return "vllm_omni/platforms/npu/stage_configs"
return "vllm_omni/deploy"
@classmethod
def get_diffusion_model_impl_qualname(cls, op_name: str) -> str:
vllm_omni/platforms/rocm/platform.py
+1 -1
View File
@@ -122,7 +122,7 @@ class RocmOmniPlatform(OmniPlatform, RocmPlatform):
@classmethod
def get_default_stage_config_path(cls) -> str:
return "vllm_omni/platforms/rocm/stage_configs"
return "vllm_omni/deploy"
@classmethod
def get_torch_device(cls, local_rank: int | None = None) -> torch.device: