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

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
Signed-off-by: Zhou Taichang <tzhouam@connect.ust.hk>

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
 

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
 ```

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`
 

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

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:**

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

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

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.
 

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 |

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
 

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.