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

[Bugfix] Remove duplicated config keyword max batch size (#1851)

Browse Source 
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
Signed-off-by: Zhou Taichang <tzhouam@connect.ust.hk>
This commit is contained in:
Zhou Taichang authored and GitHub committed 2026-03-20 16:13:24 +08:00
1 parent 4c2293743b
commit 05d8de813e
67 files changed
+226 -159
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

benchmarks/qwen3-tts/vllm_omni/configs/qwen3_tts_bs1.yaml
+2 -2
View File
@@ -6,8 +6,8 @@ stage_args:
stage_type: llm
runtime:
devices: "0"
max_batch_size: 1
engine_args:
max_num_seqs: 1
model_stage: qwen3_tts
model_arch: Qwen3TTSTalkerForConditionalGeneration
worker_type: ar
@@ -37,8 +37,8 @@ stage_args:
stage_type: llm
runtime:
devices: "0"
max_batch_size: 1
engine_args:
max_num_seqs: 1
model_stage: code2wav
model_arch: Qwen3TTSCode2Wav
worker_type: generation
benchmarks/qwen3-tts/vllm_omni/configs/qwen3_tts_bs4.yaml
+2 -2
View File
@@ -7,8 +7,8 @@ stage_args:
stage_type: llm
runtime:
devices: "0"
max_batch_size: 4
engine_args:
max_num_seqs: 4
model_stage: qwen3_tts
model_arch: Qwen3TTSTalkerForConditionalGeneration
worker_type: ar
@@ -38,8 +38,8 @@ stage_args:
stage_type: llm
runtime:
devices: "0"
max_batch_size: 4
engine_args:
max_num_seqs: 4
model_stage: code2wav
model_arch: Qwen3TTSCode2Wav
worker_type: generation
docs/configuration/gpu_memory_utilization.md
+1 -1
View File
@@ -158,7 +158,7 @@ The stage initialized but ran out of memory during processing.
**Solutions:**
1. Reduce `max_num_batched_tokens`
2. Reduce `max_batch_size` in runtime config
2. Reduce `max_num_seqs` in engine_args
3. Lower `gpu_memory_utilization` slightly
4. Enable quantization if supported
docs/configuration/pd_disaggregation.md
+4 -3
View File
@@ -32,8 +32,8 @@ stage_args:
is_prefill_only: true
runtime:
devices: "0"
max_batch_size: 16
engine_args:
max_num_seqs: 16
model_stage: thinker
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_type: ar
@@ -70,8 +70,8 @@ stage_args:
is_decode_only: true
runtime:
devices: "1"
max_batch_size: 64
engine_args:
max_num_seqs: 64
model_stage: thinker
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_type: ar
@@ -128,7 +128,8 @@ After inserting the extra thinker stage, renumber the remaining stages:
- stage_id: 3
runtime:
devices: "2"
max_batch_size: 1
engine_args:
max_num_seqs: 1
engine_input_source: [2]
custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav
```
docs/configuration/stage_configs.md
+5 -5
View File
@@ -36,9 +36,9 @@ stage_args:
runtime: # The disaggregated configuration
process: true # Run this stage in a separate process
devices: "0" # Visible devices for this stage (CUDA_VISIBLE_DEVICES/torch.cuda.set_device)
max_batch_size: 1 # the batch_size for offline inference
engine_args: # Engine arguments for a certain engine
model_stage: thinker
max_num_seqs: 1
model_arch: Qwen2_5OmniForConditionalGeneration # The model implementation registered in model_executor/models/registry.py
worker_type: ar # The specific worker used
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler # The specific scehduler used
@@ -62,9 +62,9 @@ stage_args:
runtime:
process: true
devices: "1"
max_batch_size: 3
engine_args:
model_stage: talker
max_num_seqs: 3
model_arch: Qwen2_5OmniForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
@@ -88,9 +88,9 @@ stage_args:
runtime:
process: true
devices: "0" # Example: use a different GPU than the previous stage; use "0" if single GPU
max_batch_size: 1
engine_args:
model_stage: code2wav
max_num_seqs: 1
model_arch: Qwen2_5OmniForConditionalGeneration
worker_type: generation
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
@@ -159,9 +159,9 @@ Visible devices for this stage, specified as a string. This controls which GPU d
Default: `"0"`
#### `runtime.max_batch_size`
#### `engine_args.max_num_seqs`
The maximum batch size for offline inference in this stage. This limits how many sequences can be processed together in a single batch during offline inference operations.
The maximum number of sequences for concurrent processing in this stage. For LLM stages, this controls the vLLM scheduler's maximum concurrent sequences. For all stage types, this also controls how many tasks can be batched together in the task processing loop.
Default: `1`
docs/configuration/stage_configs/qwen2_5_omni.yaml
+3 -3
View File
@@ -4,9 +4,9 @@ stage_args:
runtime:
process: true # Run this stage in a separate process
devices: "0" # Visible devices for this stage (CUDA_VISIBLE_DEVICES/torch.cuda.set_device)
max_batch_size: 1
engine_args:
model_stage: thinker
max_num_seqs: 1
model_arch: Qwen2_5OmniForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
@@ -30,9 +30,9 @@ stage_args:
runtime:
process: true
devices: "1"
max_batch_size: 1
engine_args:
model_stage: talker
max_num_seqs: 1
model_arch: Qwen2_5OmniForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
@@ -56,9 +56,9 @@ stage_args:
runtime:
process: true
devices: "0" # Example: use a different GPU than the previous stage; use "0" if single GPU
max_batch_size: 1
engine_args:
model_stage: code2wav
max_num_seqs: 1
model_arch: Qwen2_5OmniForConditionalGeneration
worker_type: generation
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
docs/contributing/model/adding_diffusion_model.md
+1 -1
View File
@@ -967,7 +967,7 @@ hidden_states = hidden_states.reshape(batch_size, seq_len, -1)
1. **Reduce batch size:**
```python
omni.generate(prompts=[...], max_batch_size=2)
omni.generate(prompts=[...], max_num_seqs=2)
```
2. **Use smaller image size:**
docs/contributing/model/adding_tts_model.md
+2 -2
View File
@@ -361,9 +361,9 @@ stage_args:
stage_type: llm
runtime:
devices: "0"
max_batch_size: 64
engine_args:
model_stage: ar_stage
max_num_seqs: 64
model_arch: YourTTSModelForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
@@ -403,9 +403,9 @@ stage_args:
stage_type: llm
runtime:
devices: "0"
max_batch_size: 64
engine_args:
model_stage: ar_stage
max_num_seqs: 64
model_arch: YourTTSModelForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
docs/design/feature/async_chunk_design.md
+2 -2
View File
@@ -171,7 +171,7 @@ stage_args:
- `custom_process_next_stage_input_func: str`: Path to custom chunk processing function; receives `(transfer_manager, pooling_output, request)`. For qwen3-omni: `thinker2talker_async_chunk`, `talker2code2wav_async_chunk`
- `stage_connector_config: dict`: Connector configuration
- `worker_type: str`: Model type, e.g. `"ar"` or `"generation"` (used by OmniChunkTransferAdapter for mode-specific payload handling)
- `max_batch_size: int`: Maximum batch size for the stage
- `max_num_seqs: int`: Maximum number of sequences for concurrent processing in the stage
### Connector Configuration
@@ -195,9 +195,9 @@ stage_args:
- stage_id: 2 # code2wav stage
runtime:
devices: "1"
max_batch_size: 64 # Enables batched audio generation
engine_args:
model_stage: code2wav
max_num_seqs: 64 # Enables batched audio generation
```
## Related Files
docs/getting_started/quickstart.md
+1 -1
View File
@@ -82,7 +82,7 @@ if __name__ == "__main__":
!!! info
For diffusion pipelines, the stage config field `stage_args.[].runtime.max_batch_size` is 1 by default, and the input
For diffusion pipelines, the stage config field `stage_args.[].engine_args.max_num_seqs` is 1 by default, and the input
list is sliced into single-item requests before feeding into the diffusion pipeline. For models that do internally support
batched inputs, you can [modify this configuration](../configuration/stage_configs.md) to let the model accept a longer batch of prompts.
docs/user_guide/examples/offline_inference/bagel.md
+2 -2
View File
@@ -130,7 +130,7 @@ The default yaml configuration deploys Thinker and DiT on the same GPU. You can
| :------------------------------- | :------------------------------ | :----------------------- |
| `stage_type` | `llm` | Stage type |
| `devices` | `"0"` | GPU device ID |
| `max_batch_size` | `1` | Maximum batch size |
| `max_num_seqs` | `1` | Maximum batch size |
| `model_stage` | `thinker` | Model stage identifier |
| `model_arch` | `BagelForConditionalGeneration` | Model architecture |
| `gpu_memory_utilization` | `0.4` | GPU memory utilization |
@@ -146,7 +146,7 @@ The default yaml configuration deploys Thinker and DiT on the same GPU. You can
| :------------------------------- | :---------- | :-------------------------- |
| `stage_type` | `diffusion` | Stage type |
| `devices` | `"0"` | GPU device ID |
| `max_batch_size` | `1` | Maximum batch size |
| `max_num_seqs` | `1` | Maximum batch size |
| `model_stage` | `dit` | Model stage identifier |
| `gpu_memory_utilization` | `0.4` | GPU memory utilization |
| `omni_kv_config.need_recv_cache` | `true` | Whether to receive KV cache |
docs/user_guide/examples/offline_inference/qwen3_tts.md
+2 -2
View File
@@ -107,7 +107,7 @@ completes. This demonstrates that audio data is available progressively rather t
## Batched Decoding
The Code2Wav stage (stage 1) supports batched decoding, where multiple requests are decoded in a single forward pass through the SpeechTokenizer. To use it, provide a stage config with `max_batch_size > 1` and pass multiple prompts via `--txt-prompts` with a matching `--batch-size`.
The Code2Wav stage (stage 1) supports batched decoding, where multiple requests are decoded in a single forward pass through the SpeechTokenizer. To use it, provide a stage config with `max_num_seqs > 1` and pass multiple prompts via `--txt-prompts` with a matching `--batch-size`.
```
python end2end.py --query-type CustomVoice \
@@ -116,7 +116,7 @@ python end2end.py --query-type CustomVoice \
--stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_tts_batch.yaml
```
**Important:** `--batch-size` must match a CUDA graph capture size (1, 2, 4, 8, 16...) because the Talker's code predictor KV cache is sized to `max_num_seqs`, and CUDA graphs pad the batch to the next capture size. Both stages need `max_batch_size >= batch_size` in the stage config for batching to take effect. If only stage 1 has a higher `max_batch_size`, it won't help — stage 1 can only batch chunks from requests that are in-flight simultaneously, which requires stage 0 to also process multiple requests concurrently.
**Important:** `--batch-size` must match a CUDA graph capture size (1, 2, 4, 8, 16...) because the Talker's code predictor KV cache is sized to `max_num_seqs`, and CUDA graphs pad the batch to the next capture size. Both stages need `max_num_seqs >= batch_size` in the stage config for batching to take effect. If only stage 1 has a higher `max_num_seqs`, it won't help — stage 1 can only batch chunks from requests that are in-flight simultaneously, which requires stage 0 to also process multiple requests concurrently.
## Notes
docs/user_guide/examples/offline_inference/text_to_image.md
+1 -1
View File
@@ -49,7 +49,7 @@ if __name__ == "__main__":
!!! info
For diffusion pipelines, the stage config field `stage_args.[].runtime.max_batch_size` is 1 by default, and the input
For diffusion pipelines, the stage config field `stage_args.[].engine_args.max_num_seqs` is 1 by default, and the input
list is sliced into single-item requests before feeding into the diffusion pipeline. For models that do internally support
batched inputs, you can [modify this configuration](https://github.com/vllm-project/vllm-omni/tree/main/configuration/stage_configs.md) to let the model accept a longer batch of prompts.
examples/offline_inference/bagel/README.md
+2 -2
View File
@@ -127,7 +127,7 @@ The default yaml configuration deploys Thinker and DiT on the same GPU. You can
| :------------------------------- | :------------------------------ | :----------------------- |
| `stage_type` | `llm` | Stage type |
| `devices` | `"0"` | GPU device ID |
| `max_batch_size` | `1` | Maximum batch size |
| `max_num_seqs` | `1` | Maximum batch size |
| `model_stage` | `thinker` | Model stage identifier |
| `model_arch` | `BagelForConditionalGeneration` | Model architecture |
| `gpu_memory_utilization` | `0.4` | GPU memory utilization |
@@ -143,7 +143,7 @@ The default yaml configuration deploys Thinker and DiT on the same GPU. You can
| :------------------------------- | :---------- | :-------------------------- |
| `stage_type` | `diffusion` | Stage type |
| `devices` | `"0"` | GPU device ID |
| `max_batch_size` | `1` | Maximum batch size |
| `max_num_seqs` | `1` | Maximum batch size |
| `model_stage` | `dit` | Model stage identifier |
| `gpu_memory_utilization` | `0.4` | GPU memory utilization |
| `omni_kv_config.need_recv_cache` | `true` | Whether to receive KV cache |
examples/offline_inference/qwen3_tts/README.md
+2 -2
View File
@@ -104,7 +104,7 @@ completes. This demonstrates that audio data is available progressively rather t
## Batched Decoding
The Code2Wav stage (stage 1) supports batched decoding, where multiple requests are decoded in a single forward pass through the SpeechTokenizer. To use it, provide a stage config with `max_batch_size > 1` and pass multiple prompts via `--txt-prompts` with a matching `--batch-size`.
The Code2Wav stage (stage 1) supports batched decoding, where multiple requests are decoded in a single forward pass through the SpeechTokenizer. To use it, provide a stage config with `max_num_seqs > 1` and pass multiple prompts via `--txt-prompts` with a matching `--batch-size`.
```
python end2end.py --query-type CustomVoice \
@@ -113,7 +113,7 @@ python end2end.py --query-type CustomVoice \
--stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_tts_batch.yaml
```
**Important:** `--batch-size` must match a CUDA graph capture size (1, 2, 4, 8, 16...) because the Talker's code predictor KV cache is sized to `max_num_seqs`, and CUDA graphs pad the batch to the next capture size. Both stages need `max_batch_size >= batch_size` in the stage config for batching to take effect. If only stage 1 has a higher `max_batch_size`, it won't help — stage 1 can only batch chunks from requests that are in-flight simultaneously, which requires stage 0 to also process multiple requests concurrently.
**Important:** `--batch-size` must match a CUDA graph capture size (1, 2, 4, 8, 16...) because the Talker's code predictor KV cache is sized to `max_num_seqs`, and CUDA graphs pad the batch to the next capture size. Both stages need `max_num_seqs >= batch_size` in the stage config for batching to take effect. If only stage 1 has a higher `max_num_seqs`, it won't help — stage 1 can only batch chunks from requests that are in-flight simultaneously, which requires stage 0 to also process multiple requests concurrently.
## Notes
examples/offline_inference/text_to_image/README.md
+1 -1
View File
@@ -46,7 +46,7 @@ if __name__ == "__main__":
!!! info
For diffusion pipelines, the stage config field `stage_args.[].runtime.max_batch_size` is 1 by default, and the input
For diffusion pipelines, the stage config field `stage_args.[].engine_args.max_num_seqs` is 1 by default, and the input
list is sliced into single-item requests before feeding into the diffusion pipeline. For models that do internally support
batched inputs, you can [modify this configuration](../../../configuration/stage_configs.md) to let the model accept a longer batch of prompts.
examples/online_serving/qwen3_omni/qwen3_omni_moe_thinking.yaml
+1 -1
View File
@@ -9,9 +9,9 @@ stage_args:
- stage_id: 0
runtime:
devices: "0,1"
max_batch_size: 1
engine_args:
model_stage: thinker
max_num_seqs: 1
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
tests/conftest.py
+1 -1
View File
@@ -843,7 +843,7 @@ def modify_stage_config(
'async_chunk': True,
'stage_args': {
0: {'engine_args.max_model_len': 5800},
1: {'runtime.max_batch_size': 2}
1: {'engine_args.max_num_seqs': 2}
}
}
deletes: Dictionary containing configurations to delete.
tests/e2e/offline_inference/stage_configs/bagel_mooncake_ci.yaml
+2 -2
View File
@@ -7,9 +7,9 @@ stage_args:
prompt_expand_func: vllm_omni.model_executor.stage_input_processors.bagel.expand_cfg_prompts
runtime:
devices: "0"
max_batch_size: 1
engine_args:
model_stage: thinker
max_num_seqs: 1
model_arch: BagelForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
@@ -44,9 +44,9 @@ stage_args:
cfg_kv_collect_func: vllm_omni.model_executor.stage_input_processors.bagel.collect_cfg_kv_caches
runtime:
devices: "0"
max_batch_size: 1
engine_args:
model_stage: dit
max_num_seqs: 1
gpu_memory_utilization: 0.45
enforce_eager: true
trust_remote_code: true
tests/e2e/offline_inference/stage_configs/bagel_sharedmemory_ci.yaml
+2 -2
View File
@@ -7,9 +7,9 @@ stage_args:
prompt_expand_func: vllm_omni.model_executor.stage_input_processors.bagel.expand_cfg_prompts
runtime:
devices: "0"
max_batch_size: 1
engine_args:
model_stage: thinker
max_num_seqs: 1
model_arch: OmniBagelForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
@@ -43,9 +43,9 @@ stage_args:
cfg_kv_collect_func: vllm_omni.model_executor.stage_input_processors.bagel.collect_cfg_kv_caches
runtime:
devices: "0"
max_batch_size: 1
engine_args:
model_stage: dit
max_num_seqs: 1
gpu_memory_utilization: 0.45
enforce_eager: true
trust_remote_code: true
tests/e2e/offline_inference/stage_configs/npu/qwen2_5_omni_ci.yaml
+1 -3
View File
@@ -6,7 +6,6 @@ stage_args:
runtime:
process: true # Run this stage in a separate process
devices: "0"
max_batch_size: 1
engine_args:
model_stage: thinker
model_arch: Qwen2_5OmniForConditionalGeneration
@@ -37,7 +36,6 @@ stage_args:
runtime:
process: true
devices: "1"
max_batch_size: 1
engine_args:
model_stage: talker
model_arch: Qwen2_5OmniForConditionalGeneration
@@ -67,9 +65,9 @@ stage_args:
runtime:
process: true
devices: "0" # Example: use a different GPU than the previous stage; use "0" if single GPU
max_batch_size: 1
engine_args:
model_stage: code2wav
max_num_seqs: 1
model_arch: Qwen2_5OmniForConditionalGeneration
worker_type: generation
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
tests/e2e/stage_configs/qwen2_5_omni_ci.yaml
+1 -3
View File
@@ -7,7 +7,6 @@ stage_args:
runtime:
process: true # Run this stage in a separate process
devices: "0" # Visible devices for this stage (CUDA_VISIBLE_DEVICES/torch.cuda.set_device)
max_batch_size: 1
engine_args:
model_stage: thinker
model_arch: Qwen2_5OmniForConditionalGeneration
@@ -38,7 +37,6 @@ stage_args:
runtime:
process: true
devices: "1"
max_batch_size: 1
engine_args:
model_stage: talker
model_arch: Qwen2_5OmniForConditionalGeneration
@@ -68,9 +66,9 @@ stage_args:
runtime:
process: true
devices: "2" # Example: use a different GPU than the previous stage; use "0" if single GPU
max_batch_size: 1
engine_args:
model_stage: code2wav
max_num_seqs: 1
model_arch: Qwen2_5OmniForConditionalGeneration
worker_type: generation
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
tests/e2e/stage_configs/qwen2_5_omni_thinker_ci.yaml
-1
View File
@@ -3,7 +3,6 @@ stage_args:
runtime:
process: true # Run this stage in a separate process
devices: "0" # Visible devices for this stage (CUDA_VISIBLE_DEVICES/torch.cuda.set_device)
max_batch_size: 1
engine_args:
model_stage: thinker
model_arch: Qwen2_5OmniForConditionalGeneration
tests/e2e/stage_configs/qwen3_omni_ci.yaml
+3 -3
View File
@@ -8,9 +8,9 @@ stage_args:
- stage_id: 0
runtime:
devices: "0"
max_batch_size: 5
engine_args:
model_stage: thinker
max_num_seqs: 5
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
@@ -42,9 +42,9 @@ stage_args:
- stage_id: 1
runtime:
devices: "1"
max_batch_size: 5
engine_args:
model_stage: talker
max_num_seqs: 5
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
@@ -72,9 +72,9 @@ stage_args:
- stage_id: 2
runtime:
devices: "1"
max_batch_size: 1
engine_args:
model_stage: code2wav
max_num_seqs: 1
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_type: generation
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
tests/e2e/stage_configs/rocm/qwen2_5_omni_ci.yaml
+1 -3
View File
@@ -7,7 +7,6 @@ stage_args:
runtime:
process: true # Run this stage in a separate process
devices: "0" # Visible devices for this stage (CUDA_VISIBLE_DEVICES/torch.cuda.set_device)
max_batch_size: 1
engine_args:
model_stage: thinker
model_arch: Qwen2_5OmniForConditionalGeneration
@@ -38,7 +37,6 @@ stage_args:
runtime:
process: true
devices: "1"
max_batch_size: 1
engine_args:
model_stage: talker
model_arch: Qwen2_5OmniForConditionalGeneration
@@ -68,9 +66,9 @@ stage_args:
runtime:
process: true
devices: "0" # Example: use a different GPU than the previous stage; use "0" if single GPU
max_batch_size: 1
engine_args:
model_stage: code2wav
max_num_seqs: 1
model_arch: Qwen2_5OmniForConditionalGeneration
worker_type: generation
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
tests/e2e/stage_configs/rocm/qwen3_omni_ci.yaml
+3 -3
View File
@@ -8,9 +8,9 @@ stage_args:
- stage_id: 0
runtime:
devices: "0"
max_batch_size: 1
engine_args:
model_stage: thinker
max_num_seqs: 1
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
@@ -39,9 +39,9 @@ stage_args:
- stage_id: 1
runtime:
devices: "1"
max_batch_size: 1
engine_args:
model_stage: talker
max_num_seqs: 1
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
@@ -70,9 +70,9 @@ stage_args:
- stage_id: 2
runtime:
devices: "1"
max_batch_size: 1
engine_args:
model_stage: code2wav
max_num_seqs: 1
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_type: generation
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
tests/e2e/stage_configs/xpu/qwen2_5_omni_ci.yaml
+3 -3
View File
@@ -7,8 +7,8 @@ stage_args:
runtime:
process: true # Run this stage in a separate process
devices: "0" # Visible devices for this stage
max_batch_size: 1
engine_args:
max_num_seqs: 1
model_stage: thinker
model_arch: Qwen2_5OmniForConditionalGeneration
worker_type: ar
@@ -39,8 +39,8 @@ stage_args:
runtime:
process: true
devices: "1"
max_batch_size: 1
engine_args:
max_num_seqs: 1
model_stage: talker
model_arch: Qwen2_5OmniForConditionalGeneration
worker_type: ar
@@ -71,8 +71,8 @@ stage_args:
runtime:
process: true
devices: "2"
max_batch_size: 1
engine_args:
max_num_seqs: 1
model_stage: code2wav
model_arch: Qwen2_5OmniForConditionalGeneration
worker_type: generation
tests/e2e/stage_configs/xpu/qwen3_omni_ci.yaml
+3 -3
View File
@@ -9,8 +9,8 @@ stage_args:
stage_type: llm # Use llm stage type for AR stages
runtime:
devices: "0,1,2,3"
max_batch_size: 1
engine_args:
max_num_seqs: 1
model_stage: thinker
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_type: ar
@@ -44,8 +44,8 @@ stage_args:
stage_type: llm # Use llm stage type for AR stages
runtime:
devices: "4"
max_batch_size: 1
engine_args:
max_num_seqs: 1
model_stage: talker
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_type: ar
@@ -78,8 +78,8 @@ stage_args:
stage_type: llm # Use llm stage type for AR stages
runtime:
devices: "5"
max_batch_size: 1
engine_args:
max_num_seqs: 1
model_stage: code2wav
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_type: generation
tests/perf/stage_configs/qwen3_omni.yaml
+3 -3
View File
@@ -10,9 +10,9 @@ stage_args:
stage_type: llm # Use llm stage type for AR stages
runtime:
devices: "0"
max_batch_size: 64
engine_args:
model_stage: thinker
max_num_seqs: 64
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
@@ -41,9 +41,9 @@ stage_args:
stage_type: llm # Use llm stage type for AR stages
runtime:
devices: "1"
max_batch_size: 64
engine_args:
model_stage: talker
max_num_seqs: 64
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
@@ -72,9 +72,9 @@ stage_args:
stage_type: llm # Use llm stage type for AR stages
runtime:
devices: "1"
max_batch_size: 64
engine_args:
model_stage: code2wav
max_num_seqs: 64
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_type: generation
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
tests/perf/stage_configs/qwen3_tts.yaml
+2 -2
View File
@@ -10,8 +10,8 @@ stage_args:
is_comprehension: true
runtime:
devices: "0"
max_batch_size: 4
engine_args:
max_num_seqs: 4
model_stage: qwen3_tts
model_arch: Qwen3TTSTalkerForConditionalGeneration
worker_type: ar
@@ -41,8 +41,8 @@ stage_args:
stage_type: llm
runtime:
devices: "0"
max_batch_size: 4
engine_args:
max_num_seqs: 4
model_stage: code2wav
model_arch: Qwen3TTSCode2Wav
worker_type: generation
tests/test_config_factory.py
+56 -2
View File
@@ -83,6 +83,8 @@ class TestStageConfig:
def test_to_omegaconf_with_runtime_overrides(self):
"""Test that runtime overrides are applied to OmegaConf output."""
import warnings
config = StageConfig(
stage_id=0,
model_stage="thinker",
@@ -93,12 +95,41 @@ class TestStageConfig:
"max_batch_size": 64,
},
)
omega_config = config.to_omegaconf()
with warnings.catch_warnings():
warnings.simplefilter("ignore", FutureWarning)
omega_config = config.to_omegaconf()
assert omega_config.engine_args.gpu_memory_utilization == 0.9
assert omega_config.engine_args.tensor_parallel_size == 2
assert omega_config.runtime.devices == "0,1"
assert omega_config.runtime.max_batch_size == 64
# max_batch_size is migrated to engine_args.max_num_seqs
assert omega_config.engine_args.max_num_seqs == 64
def test_to_omegaconf_max_batch_size_deprecation(self):
"""Test that runtime.max_batch_size emits a FutureWarning."""
import warnings
config = StageConfig(
stage_id=0,
model_stage="thinker",
runtime_overrides={"max_batch_size": 8},
)
with warnings.catch_warnings(record=True) as w:
warnings.simplefilter("always")
config.to_omegaconf()
deprecation_warnings = [x for x in w if issubclass(x.category, FutureWarning)]
assert len(deprecation_warnings) == 1
assert "max_batch_size" in str(deprecation_warnings[0].message)
def test_to_omegaconf_max_num_seqs_in_engine_args(self):
"""Test that max_num_seqs in yaml_engine_args takes precedence."""
config = StageConfig(
stage_id=0,
model_stage="thinker",
yaml_engine_args={"max_num_seqs": 32},
)
omega_config = config.to_omegaconf()
assert omega_config.engine_args.max_num_seqs == 32
class TestModelPipeline:
@@ -425,6 +456,29 @@ stages:
errors = pipeline.validate_pipeline()
assert errors == [], f"Unexpected validation errors: {errors}"
def test_parse_yaml_migrates_runtime_max_batch_size(self, tmp_path):
"""Test that runtime.max_batch_size is migrated to engine_args.max_num_seqs."""
yaml_content = """\
model_type: legacy_batch
stages:
- stage_id: 0
model_stage: entry
stage_type: llm
input_sources: []
runtime:
devices: "0"
max_batch_size: 16
engine_args:
model_arch: SomeModel
"""
yaml_file = tmp_path / "legacy_batch.yaml"
yaml_file.write_text(yaml_content)
pipeline = StageConfigFactory._parse_pipeline_yaml(yaml_file, "legacy_batch")
s0 = pipeline.stages[0]
assert "max_batch_size" not in s0.yaml_runtime
assert s0.yaml_engine_args.get("max_num_seqs") == 16
def test_parse_diffusion_stage_type(self, tmp_path):
"""Test parsing a diffusion stage type from YAML."""
yaml_content = """\
vllm_omni/config/stage_config.py
+28 -6
View File
@@ -11,6 +11,7 @@ Runtime parameters (gpu_memory_utilization, tp_size, etc.) come from CLI.
from __future__ import annotations
import re
import warnings
from dataclasses import asdict, dataclass, field
from enum import Enum
from pathlib import Path
@@ -74,7 +75,7 @@ class StageConfig:
# Per-stage engine args from pipeline YAML (defaults)
yaml_engine_args: dict[str, Any] = field(default_factory=dict)
# Per-stage runtime config from pipeline YAML (devices, max_batch_size)
# Per-stage runtime config from pipeline YAML (devices, etc.)
yaml_runtime: dict[str, Any] = field(default_factory=dict)
# Pass-through fields from pipeline YAML (default_sampling_params,
# output_connectors, input_connectors, tts_args, etc.)
@@ -109,13 +110,23 @@ class StageConfig:
# Build runtime config from YAML defaults + CLI overrides
runtime: dict[str, Any] = dict(self.yaml_runtime)
runtime.setdefault("process", True)
runtime.setdefault("max_batch_size", self.runtime_overrides.get("max_batch_size", 1))
if "devices" in self.runtime_overrides:
runtime["devices"] = self.runtime_overrides["devices"]
# Inject max_num_seqs from max_batch_size (legacy compat)
max_batch_size = int(runtime.get("max_batch_size", 1) or 1)
engine_args.setdefault("max_num_seqs", max_batch_size)
# Legacy compat: migrate runtime.max_batch_size → engine_args.max_num_seqs
legacy_mbs = runtime.pop("max_batch_size", None)
cli_mbs = self.runtime_overrides.get("max_batch_size")
if legacy_mbs is not None or cli_mbs is not None:
warnings.warn(
"runtime.max_batch_size is deprecated and will be removed in a "
"future release. Use engine_args.max_num_seqs instead.",
FutureWarning,
stacklevel=2,
)
effective_mbs = int(cli_mbs or legacy_mbs or 1)
engine_args.setdefault("max_num_seqs", effective_mbs)
engine_args.setdefault("max_num_seqs", 1)
# Build full config dict
config_dict: dict[str, Any] = {
@@ -322,13 +333,14 @@ class StageConfigFactory:
if "dtype" in engine_args:
engine_args["dtype"] = str(engine_args["dtype"])
engine_args.setdefault("max_num_seqs", 1)
config_dict: dict[str, Any] = {
"stage_id": 0,
"stage_type": StageType.DIFFUSION.value,
"runtime": {
"process": True,
"devices": devices,
"max_batch_size": 1,
},
"engine_args": create_config(engine_args),
"final_output": True,
@@ -426,6 +438,16 @@ class StageConfigFactory:
raw_rt = stage_data.get("runtime", None)
yaml_runtime = to_dict(raw_rt) if raw_rt is not None else {}
# Migrate legacy runtime.max_batch_size → engine_args.max_num_seqs
if "max_batch_size" in yaml_runtime:
mbs = yaml_runtime.pop("max_batch_size")
yaml_engine_args.setdefault("max_num_seqs", int(mbs))
logger.debug(
"Stage %s: migrated runtime.max_batch_size=%s to engine_args.max_num_seqs",
stage_data.get("stage_id", "?"),
mbs,
)
# Topology-level fields that also live inside engine_args in legacy
# YAMLs (worker_type, scheduler_cls, etc.) — read from both places.
worker_type = stage_data.get("worker_type", None) or yaml_engine_args.pop("worker_type", None)
vllm_omni/engine/async_omni_engine.py
+1 -1
View File
@@ -825,9 +825,9 @@ class AsyncOmniEngine:
"runtime": {
"process": True,
"devices": devices,
"max_batch_size": 1,
},
"engine_args": {
"max_num_seqs": 1,
"parallel_config": parallel_config,
"model_class_name": kwargs.get("model_class_name", None),
"vae_use_slicing": kwargs.get("vae_use_slicing", False),
vllm_omni/entrypoints/utils.py
-3
View File
@@ -288,9 +288,6 @@ def load_stage_configs_from_yaml(config_path: str, base_engine_args: dict | None
base_engine_args_tmp = create_config(merge_configs(base_engine_args_tmp, stage_arg.engine_args))
stage_type = getattr(stage_arg, "stage_type", "llm")
if hasattr(stage_arg, "runtime") and stage_arg.runtime is not None and stage_type != "diffusion":
runtime_cfg = stage_arg.runtime
max_batch_size = int(runtime_cfg.get("max_batch_size", 1) or 1)
base_engine_args_tmp["max_num_seqs"] = max_batch_size
base_engine_args_tmp.async_chunk = global_async_chunk
stage_arg.engine_args = base_engine_args_tmp
return stage_args
vllm_omni/model_executor/models/qwen3_tts/pipeline.yaml
+2 -2
View File
@@ -11,8 +11,8 @@ stages:
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
runtime:
devices: "0"
max_batch_size: 10
engine_args:
max_num_seqs: 10
model_arch: Qwen3TTSTalkerForConditionalGeneration
hf_overrides:
architectures: [Qwen3TTSTalkerForConditionalGeneration]
@@ -47,8 +47,8 @@ stages:
final_output_type: audio
runtime:
devices: "0"
max_batch_size: 1
engine_args:
max_num_seqs: 1
model_arch: Qwen3TTSCode2Wav
hf_overrides:
architectures: [Qwen3TTSCode2Wav]
vllm_omni/model_executor/stage_configs/bagel.yaml
+2 -2
View File
@@ -7,9 +7,9 @@ stage_args:
runtime:
devices: "0"
# 3 = 1 user prompt + 2 CFG companions (text-unconditional + image-unconditional).
max_batch_size: 3
engine_args:
model_stage: thinker
max_num_seqs: 3
model_arch: OmniBagelForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
@@ -42,9 +42,9 @@ stage_args:
cfg_kv_collect_func: vllm_omni.model_executor.stage_input_processors.bagel.collect_cfg_kv_caches
runtime:
devices: "0"
max_batch_size: 1
engine_args:
model_stage: dit
max_num_seqs: 1
gpu_memory_utilization: 0.45
enforce_eager: true
trust_remote_code: true
vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml
+2 -2
View File
@@ -5,9 +5,9 @@ stage_args:
stage_type: llm
runtime:
devices: "0"
max_batch_size: 1
engine_args:
model_stage: thinker
max_num_seqs: 1
model_arch: BagelForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
@@ -42,9 +42,9 @@ stage_args:
stage_type: diffusion
runtime:
devices: "0"
max_batch_size: 1
engine_args:
model_stage: dit
max_num_seqs: 1
gpu_memory_utilization: 0.45
enforce_eager: true
trust_remote_code: true
vllm_omni/model_executor/stage_configs/fish_speech_s2_pro.yaml
+2 -2
View File
@@ -4,8 +4,8 @@ stage_args:
stage_type: llm
runtime:
devices: "0"
max_batch_size: 4
engine_args:
max_num_seqs: 4
model_stage: fish_speech_slow_ar
model_arch: FishSpeechSlowARForConditionalGeneration
worker_type: ar
@@ -37,8 +37,8 @@ stage_args:
stage_type: llm
runtime:
devices: "0"
max_batch_size: 1
engine_args:
max_num_seqs: 1
model_stage: dac_decoder
model_arch: FishSpeechDACDecoder
worker_type: generation
vllm_omni/model_executor/stage_configs/glm_image.yaml
+2 -2
View File
@@ -11,10 +11,10 @@ stage_args:
runtime:
process: true
devices: "0"
max_batch_size: 1
requires_multimodal_data: true # Required for i2i mode to receive source images
engine_args:
model_stage: ar
max_num_seqs: 1
model_arch: GlmImageForConditionalGeneration
model_subdir: vision_language_encoder # AR model config.json is in this subdirectory
tokenizer_subdir: processor # Use processor's tokenizer (not ByT5 from tokenizer/)
@@ -45,10 +45,10 @@ stage_args:
runtime:
process: true
devices: "1" # Can use different GPU, or same GPU if memory allows
max_batch_size: 1
requires_multimodal_data: true # Required for i2i mode to pass condition images
engine_args:
model_stage: dit
max_num_seqs: 1
model_arch: GlmImagePipeline # Required for diffusion model class resolution
# Diffusion-specific parameters
num_gpus: 1
vllm_omni/model_executor/stage_configs/glm_image_muilticonnector.yaml
+2 -2
View File
@@ -14,10 +14,10 @@ stage_args:
runtime:
process: true
devices: "0"
max_batch_size: 1
requires_multimodal_data: true # Required for i2i mode to receive source images
engine_args:
model_stage: ar
max_num_seqs: 1
model_arch: GlmImageForConditionalGeneration
worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
@@ -47,10 +47,10 @@ stage_args:
runtime:
process: true
devices: "1" # Use separate GPU for diffusion
max_batch_size: 1
requires_multimodal_data: true # Required for i2i mode to pass condition images
engine_args:
model_stage: dit
max_num_seqs: 1
# Diffusion-specific parameters
num_gpus: 1
enforce_eager: true
vllm_omni/model_executor/stage_configs/hunyuan_image3_moe_dit.yaml
+1 -1
View File
@@ -8,8 +8,8 @@ stage_args:
stage_type: diffusion
runtime:
devices: "0,1,2,3"
max_batch_size: 1
engine_args:
max_num_seqs: 1
model_stage: dit
gpu_memory_utilization: 0.65
enforce_eager: true
vllm_omni/model_executor/stage_configs/hunyuan_image_3_moe.yaml
+1 -1
View File
@@ -8,9 +8,9 @@ stage_args:
runtime:
process: true # Run this stage in a separate process
devices: "0,1,2,3,4,5,6,7" # Visible devices for this stage (CUDA_VISIBLE_DEVICES/torch.cuda.set_device)
max_batch_size: 1
engine_args:
model_stage: AR
max_num_seqs: 1
model_arch: HunyuanImage3ForCausalMM
worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
vllm_omni/model_executor/stage_configs/mammoth_moda2.yaml
+2 -2
View File
@@ -8,9 +8,9 @@ stage_args:
- stage_id: 0
runtime:
devices: "0"
max_batch_size: 100
engine_args:
model_stage: ar
max_num_seqs: 100
model_arch: MammothModa2ForConditionalGeneration
worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
@@ -25,9 +25,9 @@ stage_args:
- stage_id: 1
runtime:
devices: "0"
max_batch_size: 1
engine_args:
model_stage: dit
max_num_seqs: 1
model_arch: MammothModa2ForConditionalGeneration
worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
vllm_omni/model_executor/stage_configs/mammoth_moda2_ar.yaml
+1 -1
View File
@@ -8,9 +8,9 @@ stage_args:
- stage_id: 0
runtime:
devices: "0"
max_batch_size: 16
engine_args:
model_stage: ar
max_num_seqs: 16
model_arch: MammothModa2ForConditionalGeneration
worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
vllm_omni/model_executor/stage_configs/mimo_audio.yaml
+2 -2
View File
@@ -8,9 +8,9 @@ stage_args:
runtime:
process: true # Run this stage in a separate process
devices: "0" # Visible devices for this stage (CUDA_VISIBLE_DEVICES/torch.cuda.set_device)
max_batch_size: 1
engine_args:
dtype: bfloat16
max_num_seqs: 1
model_stage: fused_thinker_talker
model_arch: MiMoAudioForConditionalGeneration
worker_type: ar
@@ -40,9 +40,9 @@ stage_args:
runtime:
process: true # Run this stage in a separate process
devices: "0" # Visible devices for this stage (CUDA_VISIBLE_DEVICES/torch.cuda.set_device)
max_batch_size: 1
engine_args:
model_stage: code2wav
max_num_seqs: 1
model_arch: MiMoAudioForConditionalGeneration
worker_type: generation
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
vllm_omni/model_executor/stage_configs/mimo_audio_async_chunk.yaml
+2 -2
View File
@@ -8,9 +8,9 @@ stage_args:
runtime:
process: true # Run this stage in a separate process
devices: "0" # Visible devices for this stage (CUDA_VISIBLE_DEVICES/torch.cuda.set_device)
max_batch_size: 1
engine_args:
dtype: bfloat16
max_num_seqs: 1
model_stage: fused_thinker_talker
model_arch: MiMoAudioForConditionalGeneration
worker_type: ar
@@ -41,9 +41,9 @@ stage_args:
runtime:
process: true # Run this stage in a separate process
devices: "0" # Visible devices for this stage (CUDA_VISIBLE_DEVICES/torch.cuda.set_device)
max_batch_size: 1
engine_args:
model_stage: code2wav
max_num_seqs: 1
model_arch: MiMoAudioForConditionalGeneration
worker_type: generation
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml
+3 -3
View File
@@ -7,9 +7,9 @@ stage_args:
runtime:
process: true # Run this stage in a separate process
devices: "0" # Visible devices for this stage (CUDA_VISIBLE_DEVICES/torch.cuda.set_device)
max_batch_size: 1
engine_args:
model_stage: thinker
max_num_seqs: 1
model_arch: Qwen2_5OmniForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
@@ -37,9 +37,9 @@ stage_args:
runtime:
process: true
devices: "1"
max_batch_size: 1
engine_args:
model_stage: talker
max_num_seqs: 1
model_arch: Qwen2_5OmniForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
@@ -66,9 +66,9 @@ stage_args:
runtime:
process: true
devices: "0" # Example: use a different GPU than the previous stage; use "0" if single GPU
max_batch_size: 1
engine_args:
model_stage: code2wav
max_num_seqs: 1
model_arch: Qwen2_5OmniForConditionalGeneration
worker_type: generation
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
vllm_omni/model_executor/stage_configs/qwen2_5_omni_multiconnector.yaml
+3 -3
View File
@@ -7,9 +7,9 @@ stage_args:
runtime:
process: true # Run this stage in a separate process
devices: "0" # Visible devices for this stage (CUDA_VISIBLE_DEVICES/torch.cuda.set_device)
max_batch_size: 1
engine_args:
model_stage: thinker
max_num_seqs: 1
model_arch: Qwen2_5OmniForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
@@ -38,9 +38,9 @@ stage_args:
runtime:
process: true
devices: "1"
max_batch_size: 1
engine_args:
model_stage: talker
max_num_seqs: 1
model_arch: Qwen2_5OmniForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
@@ -70,9 +70,9 @@ stage_args:
runtime:
process: true
devices: "2" # Example: use a different GPU than the previous stage; use "0" if single GPU
max_batch_size: 1
engine_args:
model_stage: code2wav
max_num_seqs: 1
model_arch: Qwen2_5OmniForConditionalGeneration
worker_type: generation
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
vllm_omni/model_executor/stage_configs/qwen3_omni_moe.yaml
+3 -3
View File
@@ -10,9 +10,9 @@ stage_args:
stage_type: llm # Use llm stage type for AR stages
runtime:
devices: "0"
max_batch_size: 64
engine_args:
model_stage: thinker
max_num_seqs: 64
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
@@ -41,9 +41,9 @@ stage_args:
stage_type: llm # Use llm stage type for AR stages
runtime:
devices: "1"
max_batch_size: 64
engine_args:
model_stage: talker
max_num_seqs: 64
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
@@ -72,9 +72,9 @@ stage_args:
stage_type: llm # Use llm stage type for AR stages
runtime:
devices: "1"
max_batch_size: 32
engine_args:
model_stage: code2wav
max_num_seqs: 32
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_type: generation
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
vllm_omni/model_executor/stage_configs/qwen3_omni_moe_async_chunk.yaml
+4 -4
View File
@@ -10,9 +10,9 @@ stage_args:
stage_type: llm # Use llm stage type for AR stages
runtime:
devices: "0"
max_batch_size: 64
engine_args:
model_stage: thinker
max_num_seqs: 64
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
@@ -45,9 +45,9 @@ stage_args:
stage_type: llm # Use llm stage type for AR stages
runtime:
devices: "1"
max_batch_size: 64
engine_args:
model_stage: talker
max_num_seqs: 64
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
@@ -79,9 +79,9 @@ stage_args:
stage_type: llm # Use llm stage type for AR stages
runtime:
devices: "1"
max_batch_size: 64
engine_args:
model_stage: code2wav
max_num_seqs: 64
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_type: generation
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
@@ -92,7 +92,7 @@ stage_args:
engine_output_type: audio # Final output: audio waveform
gpu_memory_utilization: 0.1
distributed_executor_backend: "mp"
max_num_batched_tokens: 51200 # [TODO] if max_num_batch_tokens < max_batch_size * 800, there will be precision problem.
max_num_batched_tokens: 51200 # [TODO] if max_num_batch_tokens < max_num_seqs * 800, there will be precision problem.
hf_config_name: thinker_config
engine_input_source: [1]
final_output: true
vllm_omni/model_executor/stage_configs/qwen3_omni_moe_multiconnector.yaml
+3 -3
View File
@@ -9,9 +9,9 @@ stage_args:
stage_type: llm # Use llm stage type for AR stages
runtime:
devices: "0"
max_batch_size: 1
engine_args:
model_stage: thinker
max_num_seqs: 1
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
@@ -42,9 +42,9 @@ stage_args:
stage_type: llm # Use llm stage type for AR stages
runtime:
devices: "1"
max_batch_size: 1
engine_args:
model_stage: talker
max_num_seqs: 1
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
@@ -78,9 +78,9 @@ stage_args:
stage_type: llm # Use llm stage type for AR stages
runtime:
devices: "1"
max_batch_size: 64
engine_args:
model_stage: code2wav
max_num_seqs: 64
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_type: generation
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
vllm_omni/model_executor/stage_configs/qwen3_tts.yaml
+2 -2
View File
@@ -5,9 +5,9 @@ stage_args:
is_comprehension: true
runtime:
devices: "0"
max_batch_size: 10
engine_args:
model_stage: qwen3_tts
max_num_seqs: 10
model_arch: Qwen3TTSTalkerForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
@@ -37,9 +37,9 @@ stage_args:
stage_type: llm
runtime:
devices: "0"
max_batch_size: 1
engine_args:
model_stage: code2wav
max_num_seqs: 1
model_arch: Qwen3TTSCode2Wav
worker_type: generation
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
vllm_omni/model_executor/stage_configs/qwen3_tts_batch.yaml
+4 -4
View File
@@ -1,6 +1,6 @@
# Same as qwen3_tts.yaml with batched talker and code2wav.
# Stage 0: max_batch_size 4, stage 1: max_batch_size 4.
# max_batch_size must be a power of two to align with CUDA graph capture sizes
# Stage 0: max_num_seqs 4, stage 1: max_num_seqs 4.
# max_num_seqs must be a power of two to align with CUDA graph capture sizes
# (stage 0) and must match --batch-size in end2end.py / benchmark scripts.
async_chunk: true
stage_args:
@@ -9,9 +9,9 @@ stage_args:
is_comprehension: true
runtime:
devices: "0"
max_batch_size: 4
engine_args:
model_stage: qwen3_tts
max_num_seqs: 4
model_arch: Qwen3TTSTalkerForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
@@ -41,9 +41,9 @@ stage_args:
stage_type: llm
runtime:
devices: "0"
max_batch_size: 4
engine_args:
model_stage: code2wav
max_num_seqs: 4
model_arch: Qwen3TTSCode2Wav
worker_type: generation
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
vllm_omni/model_executor/stage_configs/qwen3_tts_no_async_chunk.yaml
+2 -2
View File
@@ -5,9 +5,9 @@ stage_args:
is_comprehension: true
runtime:
devices: "0"
max_batch_size: 1
engine_args:
model_stage: qwen3_tts
max_num_seqs: 1
model_arch: Qwen3TTSTalkerForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
@@ -33,9 +33,9 @@ stage_args:
stage_type: llm
runtime:
devices: "0"
max_batch_size: 1
engine_args:
model_stage: code2wav
max_num_seqs: 1
model_arch: Qwen3TTSCode2Wav
worker_type: generation
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
vllm_omni/model_executor/stage_configs/voxtral_tts.yaml
+2 -2
View File
@@ -5,8 +5,8 @@ stage_args:
runtime:
process: true # Run this stage in a separate process
devices: "0" # Visible devices for this stage (CUDA_VISIBLE_DEVICES/torch.cuda.set_device)
max_batch_size: 32
engine_args:
max_num_seqs: 32
model_stage: audio_generation
model_arch: VoxtralTTSForConditionalGeneration
worker_type: ar
@@ -43,8 +43,8 @@ stage_args:
runtime:
process: true
devices: "0"
max_batch_size: 32
engine_args:
max_num_seqs: 32
model_stage: audio_tokenizer
model_arch: VoxtralTTSForConditionalGeneration
worker_type: generation
vllm_omni/model_executor/stage_configs/voxtral_tts_1773881731.yaml
+4 -4
View File
@@ -2,8 +2,8 @@ async_chunk: true
stage_args:
- stage_id: 0
stage_type: llm
runtime: {process: true, devices: '0', max_batch_size: 32}
engine_args: {model_stage: audio_generation, model_arch: VoxtralTTSForConditionalGeneration,
runtime: {process: true, devices: '0'}
engine_args: {max_num_seqs: 32, model_stage: audio_generation, model_arch: VoxtralTTSForConditionalGeneration,
worker_type: ar, 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: false, trust_remote_code: true, async_scheduling: true,
engine_output_type: latent, enable_prefix_caching: false, tokenizer_mode: mistral,
@@ -17,8 +17,8 @@ stage_args:
seed: 42, detokenize: true, repetition_penalty: 1.1}
- stage_id: 1
stage_type: llm
runtime: {process: true, devices: '0', max_batch_size: 32}
engine_args: {model_stage: audio_tokenizer, model_arch: VoxtralTTSForConditionalGeneration,
runtime: {process: true, devices: '0'}
engine_args: {max_num_seqs: 32, model_stage: audio_tokenizer, model_arch: VoxtralTTSForConditionalGeneration,
worker_type: generation, worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker,
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler,
async_scheduling: false, gpu_memory_utilization: 0.1, enforce_eager: true, trust_remote_code: true,
vllm_omni/model_executor/stage_input_processors/chunk_size_utils.py
+3 -3
View File
@@ -11,7 +11,7 @@ def max_ic_for_chunk_size(chunk_size: int) -> int:
def compute_dynamic_initial_chunk_size(
active_requests: int,
max_batch_size: int,
max_num_seqs: int,
max_ic: int,
) -> int:
"""Select IC from power-of-2 steps [2, 4, ..., max_ic] based on load factor.
@@ -26,8 +26,8 @@ def compute_dynamic_initial_chunk_size(
v <<= 1
if not steps:
return max(1, max_ic)
if max_batch_size <= 0:
if max_num_seqs <= 0:
return steps[0]
load_factor = min(active_requests / max_batch_size, 1.0)
load_factor = min(active_requests / max_num_seqs, 1.0)
idx = int(round(load_factor * (len(steps) - 1)))
return steps[idx]
vllm_omni/platforms/npu/stage_configs/hunyuan_image3_moe_dit.yaml
+1 -1
View File
@@ -8,8 +8,8 @@ stage_args:
stage_type: diffusion
runtime:
devices: "0,1,2,3,4,5,6,7"
max_batch_size: 1
engine_args:
max_num_seqs: 1
model_stage: dit
gpu_memory_utilization: 0.65
enforce_eager: true
vllm_omni/platforms/npu/stage_configs/qwen2_5_omni.yaml
+3 -3
View File
@@ -5,9 +5,9 @@ stage_args:
runtime:
process: true # Run this stage in a separate process
devices: "0" # Visible devices for this stage
max_batch_size: 1
engine_args:
model_stage: thinker
max_num_seqs: 1
model_arch: Qwen2_5OmniForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
@@ -32,9 +32,9 @@ stage_args:
runtime:
process: true
devices: "1"
max_batch_size: 1
engine_args:
model_stage: talker
max_num_seqs: 1
model_arch: Qwen2_5OmniForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
@@ -59,9 +59,9 @@ stage_args:
runtime:
process: true
devices: "2" # Example: use a different NPU than the previous stage; use "0" if single NPU
max_batch_size: 1
engine_args:
model_stage: code2wav
max_num_seqs: 1
model_arch: Qwen2_5OmniForConditionalGeneration
worker_type: generation
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
vllm_omni/platforms/npu/stage_configs/qwen3_omni_moe.yaml
+3 -3
View File
@@ -8,9 +8,9 @@ stage_args:
- stage_id: 0
runtime:
devices: "0,1"
max_batch_size: 1
engine_args:
model_stage: thinker
max_num_seqs: 1
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
@@ -37,9 +37,9 @@ stage_args:
- stage_id: 1
runtime:
devices: "2"
max_batch_size: 1
engine_args:
model_stage: talker
max_num_seqs: 1
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
@@ -67,9 +67,9 @@ stage_args:
- stage_id: 2
runtime:
devices: "2"
max_batch_size: 1
engine_args:
model_stage: code2wav
max_num_seqs: 1
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_type: generation
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
vllm_omni/platforms/npu/stage_configs/qwen3_omni_moe_async_chunk.yaml
+4 -4
View File
@@ -10,8 +10,8 @@ stage_args:
stage_type: llm # Use llm stage type for AR stages
runtime:
devices: "0,1"
max_batch_size: 10
engine_args:
max_num_seqs: 10
model_stage: thinker
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_type: ar
@@ -42,8 +42,8 @@ stage_args:
stage_type: llm # Use llm stage type for AR stages
runtime:
devices: "2"
max_batch_size: 10
engine_args:
max_num_seqs: 10
model_stage: talker
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_type: ar
@@ -73,8 +73,8 @@ stage_args:
stage_type: llm # Use llm stage type for AR stages
runtime:
devices: "2"
max_batch_size: 10
engine_args:
max_num_seqs: 10
model_stage: code2wav
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_type: generation
@@ -86,7 +86,7 @@ stage_args:
engine_output_type: audio # Final output: audio waveform
gpu_memory_utilization: 0.3
distributed_executor_backend: "mp"
max_num_batched_tokens: 51200 # [TODO] if max_num_batch_tokens < max_batch_size * 800, there will be precision problem.
max_num_batched_tokens: 51200 # [TODO] if max_num_batched_tokens < max_num_seqs * 800, there will be precision problem.
hf_config_name: thinker_config
engine_input_source: [1]
final_output: true
vllm_omni/platforms/npu/stage_configs/qwen3_tts.yaml
+2 -2
View File
@@ -5,9 +5,9 @@ stage_args:
is_comprehension: true
runtime:
devices: "0"
max_batch_size: 1
engine_args:
model_stage: qwen3_tts
max_num_seqs: 1
model_arch: Qwen3TTSTalkerForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
@@ -37,9 +37,9 @@ stage_args:
stage_type: llm
runtime:
devices: "0"
max_batch_size: 1
engine_args:
model_stage: code2wav
max_num_seqs: 1
model_arch: Qwen3TTSCode2Wav
worker_type: generation
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
vllm_omni/platforms/rocm/stage_configs/qwen2_5_omni.yaml
+3 -3
View File
@@ -6,9 +6,9 @@ stage_args:
runtime:
process: true # Run this stage in a separate process
devices: "0" # Visible devices for this stage (CUDA_VISIBLE_DEVICES/torch.cuda.set_device)
max_batch_size: 1
engine_args:
model_stage: thinker
max_num_seqs: 1
model_arch: Qwen2_5OmniForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
@@ -34,9 +34,9 @@ stage_args:
runtime:
process: true
devices: "1"
max_batch_size: 1
engine_args:
model_stage: talker
max_num_seqs: 1
model_arch: Qwen2_5OmniForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
@@ -62,9 +62,9 @@ stage_args:
runtime:
process: true
devices: "2" # Example: use a different GPU than the previous stage; use "0" if single GPU
max_batch_size: 1
engine_args:
model_stage: code2wav
max_num_seqs: 1
model_arch: Qwen2_5OmniForConditionalGeneration
worker_type: generation
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
vllm_omni/platforms/rocm/stage_configs/qwen3_omni_moe.yaml
+3 -3
View File
@@ -8,9 +8,9 @@ stage_args:
- stage_id: 0
runtime:
devices: "0"
max_batch_size: 1
engine_args:
model_stage: thinker
max_num_seqs: 1
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
@@ -38,9 +38,9 @@ stage_args:
- stage_id: 1
runtime:
devices: "1"
max_batch_size: 1
engine_args:
model_stage: talker
max_num_seqs: 1
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
@@ -69,9 +69,9 @@ stage_args:
- stage_id: 2
runtime:
devices: "1"
max_batch_size: 1
engine_args:
model_stage: code2wav
max_num_seqs: 1
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_type: generation
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
vllm_omni/platforms/xpu/stage_configs/bagel.yaml
+3 -3
View File
@@ -6,9 +6,9 @@ stage_args:
prompt_expand_func: vllm_omni.model_executor.stage_input_processors.bagel.expand_cfg_prompts
runtime:
devices: "0"
# 3 = 1 user prompt + 2 CFG companions (text-unconditional + image-unconditional).
max_batch_size: 3
engine_args:
# 3 = 1 user prompt + 2 CFG companions (text-unconditional + image-unconditional).
max_num_seqs: 3
model_stage: thinker
model_arch: OmniBagelForConditionalGeneration
worker_type: ar
@@ -43,8 +43,8 @@ stage_args:
cfg_kv_collect_func: vllm_omni.model_executor.stage_input_processors.bagel.collect_cfg_kv_caches
runtime:
devices: "1"
max_batch_size: 1
engine_args:
max_num_seqs: 1
model_stage: dit
gpu_memory_utilization: 0.9
enforce_eager: true
vllm_omni/platforms/xpu/stage_configs/qwen2_5_omni.yaml
+3 -3
View File
@@ -7,9 +7,9 @@ stage_args:
runtime:
process: true # Run this stage in a separate process
devices: "0" # Visible devices for this stage
max_batch_size: 1
engine_args:
model_stage: thinker
max_num_seqs: 1
model_arch: Qwen2_5OmniForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
@@ -34,9 +34,9 @@ stage_args:
runtime:
process: true
devices: "1"
max_batch_size: 1
engine_args:
model_stage: talker
max_num_seqs: 1
model_arch: Qwen2_5OmniForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
@@ -62,9 +62,9 @@ stage_args:
runtime:
process: true
devices: "1"
max_batch_size: 1
engine_args:
model_stage: code2wav
max_num_seqs: 1
model_arch: Qwen2_5OmniForConditionalGeneration
worker_type: generation
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
vllm_omni/platforms/xpu/stage_configs/qwen3_omni_moe.yaml
+3 -3
View File
@@ -9,9 +9,9 @@ stage_args:
stage_type: llm # Use llm stage type for AR stages
runtime:
devices: "0,1,2,3"
max_batch_size: 1
engine_args:
model_stage: thinker
max_num_seqs: 1
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
@@ -41,9 +41,9 @@ stage_args:
stage_type: llm # Use llm stage type for AR stages
runtime:
devices: "4"
max_batch_size: 1
engine_args:
model_stage: talker
max_num_seqs: 1
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
@@ -73,9 +73,9 @@ stage_args:
stage_type: llm # Use llm stage type for AR stages
runtime:
devices: "4"
max_batch_size: 1
engine_args:
model_stage: code2wav
max_num_seqs: 1
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_type: generation
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler