diff --git a/benchmarks/qwen3-tts/vllm_omni/configs/qwen3_tts_bs1.yaml b/benchmarks/qwen3-tts/vllm_omni/configs/qwen3_tts_bs1.yaml index d8e4072c0..5d730f605 100644 --- a/benchmarks/qwen3-tts/vllm_omni/configs/qwen3_tts_bs1.yaml +++ b/benchmarks/qwen3-tts/vllm_omni/configs/qwen3_tts_bs1.yaml @@ -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 diff --git a/benchmarks/qwen3-tts/vllm_omni/configs/qwen3_tts_bs4.yaml b/benchmarks/qwen3-tts/vllm_omni/configs/qwen3_tts_bs4.yaml index 93692d60d..10dc89564 100644 --- a/benchmarks/qwen3-tts/vllm_omni/configs/qwen3_tts_bs4.yaml +++ b/benchmarks/qwen3-tts/vllm_omni/configs/qwen3_tts_bs4.yaml @@ -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 diff --git a/docs/configuration/gpu_memory_utilization.md b/docs/configuration/gpu_memory_utilization.md index 19fc042aa..741066035 100644 --- a/docs/configuration/gpu_memory_utilization.md +++ b/docs/configuration/gpu_memory_utilization.md @@ -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 diff --git a/docs/configuration/pd_disaggregation.md b/docs/configuration/pd_disaggregation.md index 5cd5d0e7f..1cf6189e6 100644 --- a/docs/configuration/pd_disaggregation.md +++ b/docs/configuration/pd_disaggregation.md @@ -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 ``` diff --git a/docs/configuration/stage_configs.md b/docs/configuration/stage_configs.md index d8a26d516..95c42afcc 100644 --- a/docs/configuration/stage_configs.md +++ b/docs/configuration/stage_configs.md @@ -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` diff --git a/docs/configuration/stage_configs/qwen2_5_omni.yaml b/docs/configuration/stage_configs/qwen2_5_omni.yaml index 416ad137d..690577b84 100644 --- a/docs/configuration/stage_configs/qwen2_5_omni.yaml +++ b/docs/configuration/stage_configs/qwen2_5_omni.yaml @@ -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 diff --git a/docs/contributing/model/adding_diffusion_model.md b/docs/contributing/model/adding_diffusion_model.md index 3d04bf942..d044e74d7 100644 --- a/docs/contributing/model/adding_diffusion_model.md +++ b/docs/contributing/model/adding_diffusion_model.md @@ -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:** diff --git a/docs/contributing/model/adding_tts_model.md b/docs/contributing/model/adding_tts_model.md index 6f642c7cd..e48ae5049 100644 --- a/docs/contributing/model/adding_tts_model.md +++ b/docs/contributing/model/adding_tts_model.md @@ -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 diff --git a/docs/design/feature/async_chunk_design.md b/docs/design/feature/async_chunk_design.md index 0416856fb..202ef0e18 100644 --- a/docs/design/feature/async_chunk_design.md +++ b/docs/design/feature/async_chunk_design.md @@ -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 diff --git a/docs/getting_started/quickstart.md b/docs/getting_started/quickstart.md index 9b1182c6b..8ded8f55a 100644 --- a/docs/getting_started/quickstart.md +++ b/docs/getting_started/quickstart.md @@ -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. diff --git a/docs/user_guide/examples/offline_inference/bagel.md b/docs/user_guide/examples/offline_inference/bagel.md index cd245aa8d..5f458750b 100644 --- a/docs/user_guide/examples/offline_inference/bagel.md +++ b/docs/user_guide/examples/offline_inference/bagel.md @@ -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 | diff --git a/docs/user_guide/examples/offline_inference/qwen3_tts.md b/docs/user_guide/examples/offline_inference/qwen3_tts.md index 352c076ad..d04a42014 100644 --- a/docs/user_guide/examples/offline_inference/qwen3_tts.md +++ b/docs/user_guide/examples/offline_inference/qwen3_tts.md @@ -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 diff --git a/docs/user_guide/examples/offline_inference/text_to_image.md b/docs/user_guide/examples/offline_inference/text_to_image.md index 653973fe4..ce000cc8b 100644 --- a/docs/user_guide/examples/offline_inference/text_to_image.md +++ b/docs/user_guide/examples/offline_inference/text_to_image.md @@ -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. diff --git a/examples/offline_inference/bagel/README.md b/examples/offline_inference/bagel/README.md index 9649a54c3..226c009f7 100644 --- a/examples/offline_inference/bagel/README.md +++ b/examples/offline_inference/bagel/README.md @@ -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 | diff --git a/examples/offline_inference/qwen3_tts/README.md b/examples/offline_inference/qwen3_tts/README.md index ce2dd588e..90a2541bc 100644 --- a/examples/offline_inference/qwen3_tts/README.md +++ b/examples/offline_inference/qwen3_tts/README.md @@ -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 diff --git a/examples/offline_inference/text_to_image/README.md b/examples/offline_inference/text_to_image/README.md index 0de89c753..eabd349da 100644 --- a/examples/offline_inference/text_to_image/README.md +++ b/examples/offline_inference/text_to_image/README.md @@ -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. diff --git a/examples/online_serving/qwen3_omni/qwen3_omni_moe_thinking.yaml b/examples/online_serving/qwen3_omni/qwen3_omni_moe_thinking.yaml index 34b63e26e..a32ca7d2d 100644 --- a/examples/online_serving/qwen3_omni/qwen3_omni_moe_thinking.yaml +++ b/examples/online_serving/qwen3_omni/qwen3_omni_moe_thinking.yaml @@ -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 diff --git a/tests/conftest.py b/tests/conftest.py index cac78bdad..1c8d1bef3 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -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. diff --git a/tests/e2e/offline_inference/stage_configs/bagel_mooncake_ci.yaml b/tests/e2e/offline_inference/stage_configs/bagel_mooncake_ci.yaml index 7f3a3a6f4..590244acd 100644 --- a/tests/e2e/offline_inference/stage_configs/bagel_mooncake_ci.yaml +++ b/tests/e2e/offline_inference/stage_configs/bagel_mooncake_ci.yaml @@ -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 diff --git a/tests/e2e/offline_inference/stage_configs/bagel_sharedmemory_ci.yaml b/tests/e2e/offline_inference/stage_configs/bagel_sharedmemory_ci.yaml index aa3cc7718..b7999652e 100644 --- a/tests/e2e/offline_inference/stage_configs/bagel_sharedmemory_ci.yaml +++ b/tests/e2e/offline_inference/stage_configs/bagel_sharedmemory_ci.yaml @@ -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 diff --git a/tests/e2e/offline_inference/stage_configs/npu/qwen2_5_omni_ci.yaml b/tests/e2e/offline_inference/stage_configs/npu/qwen2_5_omni_ci.yaml index acca96e53..f93a6c714 100644 --- a/tests/e2e/offline_inference/stage_configs/npu/qwen2_5_omni_ci.yaml +++ b/tests/e2e/offline_inference/stage_configs/npu/qwen2_5_omni_ci.yaml @@ -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 diff --git a/tests/e2e/stage_configs/qwen2_5_omni_ci.yaml b/tests/e2e/stage_configs/qwen2_5_omni_ci.yaml index f3a8ecf79..2d5690d00 100644 --- a/tests/e2e/stage_configs/qwen2_5_omni_ci.yaml +++ b/tests/e2e/stage_configs/qwen2_5_omni_ci.yaml @@ -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 diff --git a/tests/e2e/stage_configs/qwen2_5_omni_thinker_ci.yaml b/tests/e2e/stage_configs/qwen2_5_omni_thinker_ci.yaml index 23eacafdf..940138284 100644 --- a/tests/e2e/stage_configs/qwen2_5_omni_thinker_ci.yaml +++ b/tests/e2e/stage_configs/qwen2_5_omni_thinker_ci.yaml @@ -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 diff --git a/tests/e2e/stage_configs/qwen3_omni_ci.yaml b/tests/e2e/stage_configs/qwen3_omni_ci.yaml index 6f1293f96..8b08bbb5e 100644 --- a/tests/e2e/stage_configs/qwen3_omni_ci.yaml +++ b/tests/e2e/stage_configs/qwen3_omni_ci.yaml @@ -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 diff --git a/tests/e2e/stage_configs/rocm/qwen2_5_omni_ci.yaml b/tests/e2e/stage_configs/rocm/qwen2_5_omni_ci.yaml index b95afd223..0c756ce56 100644 --- a/tests/e2e/stage_configs/rocm/qwen2_5_omni_ci.yaml +++ b/tests/e2e/stage_configs/rocm/qwen2_5_omni_ci.yaml @@ -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 diff --git a/tests/e2e/stage_configs/rocm/qwen3_omni_ci.yaml b/tests/e2e/stage_configs/rocm/qwen3_omni_ci.yaml index 23e75aae7..ac2b1fbd7 100644 --- a/tests/e2e/stage_configs/rocm/qwen3_omni_ci.yaml +++ b/tests/e2e/stage_configs/rocm/qwen3_omni_ci.yaml @@ -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 diff --git a/tests/e2e/stage_configs/xpu/qwen2_5_omni_ci.yaml b/tests/e2e/stage_configs/xpu/qwen2_5_omni_ci.yaml index 9feaaff67..1af08098a 100644 --- a/tests/e2e/stage_configs/xpu/qwen2_5_omni_ci.yaml +++ b/tests/e2e/stage_configs/xpu/qwen2_5_omni_ci.yaml @@ -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 diff --git a/tests/e2e/stage_configs/xpu/qwen3_omni_ci.yaml b/tests/e2e/stage_configs/xpu/qwen3_omni_ci.yaml index 49de8325a..c4586e066 100644 --- a/tests/e2e/stage_configs/xpu/qwen3_omni_ci.yaml +++ b/tests/e2e/stage_configs/xpu/qwen3_omni_ci.yaml @@ -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 diff --git a/tests/perf/stage_configs/qwen3_omni.yaml b/tests/perf/stage_configs/qwen3_omni.yaml index 7b347c7f1..2add22b87 100644 --- a/tests/perf/stage_configs/qwen3_omni.yaml +++ b/tests/perf/stage_configs/qwen3_omni.yaml @@ -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 diff --git a/tests/perf/stage_configs/qwen3_tts.yaml b/tests/perf/stage_configs/qwen3_tts.yaml index 1ad568f26..dd69b248d 100644 --- a/tests/perf/stage_configs/qwen3_tts.yaml +++ b/tests/perf/stage_configs/qwen3_tts.yaml @@ -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 diff --git a/tests/test_config_factory.py b/tests/test_config_factory.py index 21ef22523..e284de48d 100644 --- a/tests/test_config_factory.py +++ b/tests/test_config_factory.py @@ -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 = """\ diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index d12bd10cc..a4e186c3b 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -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) diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index 56eb3b5cc..5562b84ff 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -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), diff --git a/vllm_omni/entrypoints/utils.py b/vllm_omni/entrypoints/utils.py index 0e31bfa7c..619aac5d4 100644 --- a/vllm_omni/entrypoints/utils.py +++ b/vllm_omni/entrypoints/utils.py @@ -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 diff --git a/vllm_omni/model_executor/models/qwen3_tts/pipeline.yaml b/vllm_omni/model_executor/models/qwen3_tts/pipeline.yaml index 75e9dc1b1..67e95d562 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/pipeline.yaml +++ b/vllm_omni/model_executor/models/qwen3_tts/pipeline.yaml @@ -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] diff --git a/vllm_omni/model_executor/stage_configs/bagel.yaml b/vllm_omni/model_executor/stage_configs/bagel.yaml index 4bce816b4..b0c1b0480 100644 --- a/vllm_omni/model_executor/stage_configs/bagel.yaml +++ b/vllm_omni/model_executor/stage_configs/bagel.yaml @@ -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 diff --git a/vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml b/vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml index 93700b6f3..4919395ca 100644 --- a/vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml +++ b/vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml @@ -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 diff --git a/vllm_omni/model_executor/stage_configs/fish_speech_s2_pro.yaml b/vllm_omni/model_executor/stage_configs/fish_speech_s2_pro.yaml index a847cb181..c00528d10 100644 --- a/vllm_omni/model_executor/stage_configs/fish_speech_s2_pro.yaml +++ b/vllm_omni/model_executor/stage_configs/fish_speech_s2_pro.yaml @@ -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 diff --git a/vllm_omni/model_executor/stage_configs/glm_image.yaml b/vllm_omni/model_executor/stage_configs/glm_image.yaml index 8128847c7..3cc23e1e2 100644 --- a/vllm_omni/model_executor/stage_configs/glm_image.yaml +++ b/vllm_omni/model_executor/stage_configs/glm_image.yaml @@ -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 diff --git a/vllm_omni/model_executor/stage_configs/glm_image_muilticonnector.yaml b/vllm_omni/model_executor/stage_configs/glm_image_muilticonnector.yaml index 91fbf7659..719c73a9f 100644 --- a/vllm_omni/model_executor/stage_configs/glm_image_muilticonnector.yaml +++ b/vllm_omni/model_executor/stage_configs/glm_image_muilticonnector.yaml @@ -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 diff --git a/vllm_omni/model_executor/stage_configs/hunyuan_image3_moe_dit.yaml b/vllm_omni/model_executor/stage_configs/hunyuan_image3_moe_dit.yaml index eac7f2d63..0b812ff37 100644 --- a/vllm_omni/model_executor/stage_configs/hunyuan_image3_moe_dit.yaml +++ b/vllm_omni/model_executor/stage_configs/hunyuan_image3_moe_dit.yaml @@ -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 diff --git a/vllm_omni/model_executor/stage_configs/hunyuan_image_3_moe.yaml b/vllm_omni/model_executor/stage_configs/hunyuan_image_3_moe.yaml index 119047fb1..01d7dcca9 100644 --- a/vllm_omni/model_executor/stage_configs/hunyuan_image_3_moe.yaml +++ b/vllm_omni/model_executor/stage_configs/hunyuan_image_3_moe.yaml @@ -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 diff --git a/vllm_omni/model_executor/stage_configs/mammoth_moda2.yaml b/vllm_omni/model_executor/stage_configs/mammoth_moda2.yaml index 505548ecc..f4f523ac5 100644 --- a/vllm_omni/model_executor/stage_configs/mammoth_moda2.yaml +++ b/vllm_omni/model_executor/stage_configs/mammoth_moda2.yaml @@ -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 diff --git a/vllm_omni/model_executor/stage_configs/mammoth_moda2_ar.yaml b/vllm_omni/model_executor/stage_configs/mammoth_moda2_ar.yaml index f95ed2b3d..d455c7973 100644 --- a/vllm_omni/model_executor/stage_configs/mammoth_moda2_ar.yaml +++ b/vllm_omni/model_executor/stage_configs/mammoth_moda2_ar.yaml @@ -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 diff --git a/vllm_omni/model_executor/stage_configs/mimo_audio.yaml b/vllm_omni/model_executor/stage_configs/mimo_audio.yaml index 0807b68ca..dd6f24cc6 100644 --- a/vllm_omni/model_executor/stage_configs/mimo_audio.yaml +++ b/vllm_omni/model_executor/stage_configs/mimo_audio.yaml @@ -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 diff --git a/vllm_omni/model_executor/stage_configs/mimo_audio_async_chunk.yaml b/vllm_omni/model_executor/stage_configs/mimo_audio_async_chunk.yaml index ef9e1dedb..04f7e656e 100644 --- a/vllm_omni/model_executor/stage_configs/mimo_audio_async_chunk.yaml +++ b/vllm_omni/model_executor/stage_configs/mimo_audio_async_chunk.yaml @@ -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 diff --git a/vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml b/vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml index c66d0fcd4..0a307b447 100644 --- a/vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml +++ b/vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml @@ -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 diff --git a/vllm_omni/model_executor/stage_configs/qwen2_5_omni_multiconnector.yaml b/vllm_omni/model_executor/stage_configs/qwen2_5_omni_multiconnector.yaml index 927d71b49..6e4f871e3 100644 --- a/vllm_omni/model_executor/stage_configs/qwen2_5_omni_multiconnector.yaml +++ b/vllm_omni/model_executor/stage_configs/qwen2_5_omni_multiconnector.yaml @@ -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 diff --git a/vllm_omni/model_executor/stage_configs/qwen3_omni_moe.yaml b/vllm_omni/model_executor/stage_configs/qwen3_omni_moe.yaml index 759590bfc..0ce4f0c94 100644 --- a/vllm_omni/model_executor/stage_configs/qwen3_omni_moe.yaml +++ b/vllm_omni/model_executor/stage_configs/qwen3_omni_moe.yaml @@ -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 diff --git a/vllm_omni/model_executor/stage_configs/qwen3_omni_moe_async_chunk.yaml b/vllm_omni/model_executor/stage_configs/qwen3_omni_moe_async_chunk.yaml index d38302f8e..38626fc08 100644 --- a/vllm_omni/model_executor/stage_configs/qwen3_omni_moe_async_chunk.yaml +++ b/vllm_omni/model_executor/stage_configs/qwen3_omni_moe_async_chunk.yaml @@ -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 diff --git a/vllm_omni/model_executor/stage_configs/qwen3_omni_moe_multiconnector.yaml b/vllm_omni/model_executor/stage_configs/qwen3_omni_moe_multiconnector.yaml index b91937b0c..6c2d2a766 100644 --- a/vllm_omni/model_executor/stage_configs/qwen3_omni_moe_multiconnector.yaml +++ b/vllm_omni/model_executor/stage_configs/qwen3_omni_moe_multiconnector.yaml @@ -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 diff --git a/vllm_omni/model_executor/stage_configs/qwen3_tts.yaml b/vllm_omni/model_executor/stage_configs/qwen3_tts.yaml index f92874bf6..543f57ba7 100644 --- a/vllm_omni/model_executor/stage_configs/qwen3_tts.yaml +++ b/vllm_omni/model_executor/stage_configs/qwen3_tts.yaml @@ -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 diff --git a/vllm_omni/model_executor/stage_configs/qwen3_tts_batch.yaml b/vllm_omni/model_executor/stage_configs/qwen3_tts_batch.yaml index 23557bf23..e8c35fb8d 100644 --- a/vllm_omni/model_executor/stage_configs/qwen3_tts_batch.yaml +++ b/vllm_omni/model_executor/stage_configs/qwen3_tts_batch.yaml @@ -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 diff --git a/vllm_omni/model_executor/stage_configs/qwen3_tts_no_async_chunk.yaml b/vllm_omni/model_executor/stage_configs/qwen3_tts_no_async_chunk.yaml index f559bdebd..678222167 100644 --- a/vllm_omni/model_executor/stage_configs/qwen3_tts_no_async_chunk.yaml +++ b/vllm_omni/model_executor/stage_configs/qwen3_tts_no_async_chunk.yaml @@ -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 diff --git a/vllm_omni/model_executor/stage_configs/voxtral_tts.yaml b/vllm_omni/model_executor/stage_configs/voxtral_tts.yaml index bb791e080..31cccb9cc 100644 --- a/vllm_omni/model_executor/stage_configs/voxtral_tts.yaml +++ b/vllm_omni/model_executor/stage_configs/voxtral_tts.yaml @@ -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 diff --git a/vllm_omni/model_executor/stage_configs/voxtral_tts_1773881731.yaml b/vllm_omni/model_executor/stage_configs/voxtral_tts_1773881731.yaml index d004402b2..287d82895 100644 --- a/vllm_omni/model_executor/stage_configs/voxtral_tts_1773881731.yaml +++ b/vllm_omni/model_executor/stage_configs/voxtral_tts_1773881731.yaml @@ -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, diff --git a/vllm_omni/model_executor/stage_input_processors/chunk_size_utils.py b/vllm_omni/model_executor/stage_input_processors/chunk_size_utils.py index 7bdec0bb3..8b8956839 100644 --- a/vllm_omni/model_executor/stage_input_processors/chunk_size_utils.py +++ b/vllm_omni/model_executor/stage_input_processors/chunk_size_utils.py @@ -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] diff --git a/vllm_omni/platforms/npu/stage_configs/hunyuan_image3_moe_dit.yaml b/vllm_omni/platforms/npu/stage_configs/hunyuan_image3_moe_dit.yaml index 270e0c718..053e8a8cc 100644 --- a/vllm_omni/platforms/npu/stage_configs/hunyuan_image3_moe_dit.yaml +++ b/vllm_omni/platforms/npu/stage_configs/hunyuan_image3_moe_dit.yaml @@ -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 diff --git a/vllm_omni/platforms/npu/stage_configs/qwen2_5_omni.yaml b/vllm_omni/platforms/npu/stage_configs/qwen2_5_omni.yaml index 02a6b0dd2..8f7af161d 100644 --- a/vllm_omni/platforms/npu/stage_configs/qwen2_5_omni.yaml +++ b/vllm_omni/platforms/npu/stage_configs/qwen2_5_omni.yaml @@ -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 diff --git a/vllm_omni/platforms/npu/stage_configs/qwen3_omni_moe.yaml b/vllm_omni/platforms/npu/stage_configs/qwen3_omni_moe.yaml index 706bee5dd..59ed8611e 100644 --- a/vllm_omni/platforms/npu/stage_configs/qwen3_omni_moe.yaml +++ b/vllm_omni/platforms/npu/stage_configs/qwen3_omni_moe.yaml @@ -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 diff --git a/vllm_omni/platforms/npu/stage_configs/qwen3_omni_moe_async_chunk.yaml b/vllm_omni/platforms/npu/stage_configs/qwen3_omni_moe_async_chunk.yaml index 27f0ffcc2..9aa20baec 100644 --- a/vllm_omni/platforms/npu/stage_configs/qwen3_omni_moe_async_chunk.yaml +++ b/vllm_omni/platforms/npu/stage_configs/qwen3_omni_moe_async_chunk.yaml @@ -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 diff --git a/vllm_omni/platforms/npu/stage_configs/qwen3_tts.yaml b/vllm_omni/platforms/npu/stage_configs/qwen3_tts.yaml index 530e957f0..331dd9f98 100644 --- a/vllm_omni/platforms/npu/stage_configs/qwen3_tts.yaml +++ b/vllm_omni/platforms/npu/stage_configs/qwen3_tts.yaml @@ -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 diff --git a/vllm_omni/platforms/rocm/stage_configs/qwen2_5_omni.yaml b/vllm_omni/platforms/rocm/stage_configs/qwen2_5_omni.yaml index 8c4ffaa6f..35e819354 100644 --- a/vllm_omni/platforms/rocm/stage_configs/qwen2_5_omni.yaml +++ b/vllm_omni/platforms/rocm/stage_configs/qwen2_5_omni.yaml @@ -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 diff --git a/vllm_omni/platforms/rocm/stage_configs/qwen3_omni_moe.yaml b/vllm_omni/platforms/rocm/stage_configs/qwen3_omni_moe.yaml index 31312673a..0ca150bee 100644 --- a/vllm_omni/platforms/rocm/stage_configs/qwen3_omni_moe.yaml +++ b/vllm_omni/platforms/rocm/stage_configs/qwen3_omni_moe.yaml @@ -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 diff --git a/vllm_omni/platforms/xpu/stage_configs/bagel.yaml b/vllm_omni/platforms/xpu/stage_configs/bagel.yaml index b1443d884..0fc8a25ea 100644 --- a/vllm_omni/platforms/xpu/stage_configs/bagel.yaml +++ b/vllm_omni/platforms/xpu/stage_configs/bagel.yaml @@ -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 diff --git a/vllm_omni/platforms/xpu/stage_configs/qwen2_5_omni.yaml b/vllm_omni/platforms/xpu/stage_configs/qwen2_5_omni.yaml index 221b65602..7dbedb29a 100644 --- a/vllm_omni/platforms/xpu/stage_configs/qwen2_5_omni.yaml +++ b/vllm_omni/platforms/xpu/stage_configs/qwen2_5_omni.yaml @@ -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 diff --git a/vllm_omni/platforms/xpu/stage_configs/qwen3_omni_moe.yaml b/vllm_omni/platforms/xpu/stage_configs/qwen3_omni_moe.yaml index 59fed7c78..49914bebc 100644 --- a/vllm_omni/platforms/xpu/stage_configs/qwen3_omni_moe.yaml +++ b/vllm_omni/platforms/xpu/stage_configs/qwen3_omni_moe.yaml @@ -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