From d456721d8197655f5fdc427d2d6a32719b2dae34 Mon Sep 17 00:00:00 2001 From: Ziming Huang Date: Fri, 26 Dec 2025 21:51:07 +0800 Subject: [PATCH] [BugFix] Make max_batch_size work in Omni online serving (#487) Signed-off-by: ZeldaHuang Co-authored-by: Hongsheng Liu --- tests/e2e/online_serving/stage_configs/qwen3_omni_ci.yaml | 4 ++-- vllm_omni/entrypoints/omni_stage.py | 2 ++ 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/tests/e2e/online_serving/stage_configs/qwen3_omni_ci.yaml b/tests/e2e/online_serving/stage_configs/qwen3_omni_ci.yaml index fec6169d8..ff889043f 100644 --- a/tests/e2e/online_serving/stage_configs/qwen3_omni_ci.yaml +++ b/tests/e2e/online_serving/stage_configs/qwen3_omni_ci.yaml @@ -8,7 +8,7 @@ stage_args: - stage_id: 0 runtime: devices: "0,1" - max_batch_size: 1 + max_batch_size: 5 engine_args: model_stage: thinker model_arch: Qwen3OmniMoeForConditionalGeneration @@ -38,7 +38,7 @@ stage_args: - stage_id: 1 runtime: devices: "1" - max_batch_size: 1 + max_batch_size: 5 engine_args: model_stage: talker model_arch: Qwen3OmniMoeForConditionalGeneration diff --git a/vllm_omni/entrypoints/omni_stage.py b/vllm_omni/entrypoints/omni_stage.py index 4652a921d..3f45a97f2 100644 --- a/vllm_omni/entrypoints/omni_stage.py +++ b/vllm_omni/entrypoints/omni_stage.py @@ -880,6 +880,8 @@ async def _stage_worker_async( except Exception as e: logger.warning("Device setup failed: %s", e) + max_batch_size = int(runtime_cfg.get("max_batch_size", 1) or 1) + engine_args["max_num_seqs"] = max_batch_size # Initialize OmniConnectors if configured to match sync worker behavior connectors: dict[Any, Any] = {} if connectors_config: