[BugFix] Qwen-image performance regressed - Avoid mapping diffusion_batch_size onto scheduler max_num_seqs (#6525)

Signed-off-by: NumberWan <wantszkin2003@gmail.com>

tests/engine/test_async_omni_engine_stage_init.py

@@ -474,7 +474,7 @@ def test_initialize_local_llm_replica_scopes_runtime_env(monkeypatch):
     assert os.environ[runtime_env_var] == "parent-value"
 
 
-def test_initialize_diffusion_stage_applies_client_batch_size_to_engine(monkeypatch):
+def test_initialize_diffusion_stage_forwards_batch_size_without_writing_max_num_seqs(monkeypatch):
     import vllm_omni.diffusion.stage_diffusion_client as client_mod
     import vllm_omni.engine.stage_init_utils as init_mod
 
@@ -506,7 +506,7 @@ def test_initialize_diffusion_stage_applies_client_batch_size_to_engine(monkeypa
         use_inline=True,
     )
 
-    assert od_config.max_num_seqs == 4
+    assert od_config.max_num_seqs == 1
     assert captured == {
         "model": "dummy-model",
         "config": od_config,
@@ -517,7 +517,7 @@ def test_initialize_diffusion_stage_applies_client_batch_size_to_engine(monkeypa
     }
 
 
-def test_launch_diffusion_stage_replica_applies_batch_size_to_config(monkeypatch):
+def test_launch_diffusion_stage_replica_does_not_write_max_num_seqs_from_batch_size(monkeypatch):
     import vllm_omni.diffusion.stage_diffusion_client as client_mod
     import vllm_omni.diffusion.stage_diffusion_proc as proc_mod
     import vllm_omni.engine.stage_engine_startup as startup_mod
@@ -565,10 +565,84 @@ def test_launch_diffusion_stage_replica_applies_batch_size_to_config(monkeypatch
     )
 
     assert result is sentinel_client
-    assert od_config.max_num_seqs == 4
+    assert od_config.max_num_seqs == 1
     assert resources.manager is proc_manager
 
 
+def test_initialize_diffusion_stage_does_not_write_max_num_seqs(monkeypatch):
+    import vllm_omni.diffusion.stage_diffusion_client as client_mod
+    import vllm_omni.engine.stage_init_utils as init_mod
+
+    od_config = types.SimpleNamespace(max_num_seqs=8, step_execution=True)
+    monkeypatch.setattr(init_mod, "build_diffusion_config", lambda *args: od_config)
+    monkeypatch.setattr(client_mod, "create_diffusion_client", lambda *args: object())
+    metadata = _make_diffusion_metadata(0)
+
+    init_mod.initialize_diffusion_stage(
+        0,
+        "dummy-model",
+        types.SimpleNamespace(),
+        metadata,
+        stage_init_timeout=12,
+        batch_size=4,
+        use_inline=True,
+    )
+
+    assert od_config.max_num_seqs == 8
+
+
+def test_launch_diffusion_stage_replica_preserves_step_execution_max_num_seqs(monkeypatch):
+    import vllm_omni.diffusion.stage_diffusion_client as client_mod
+    import vllm_omni.diffusion.stage_diffusion_proc as proc_mod
+    import vllm_omni.engine.stage_engine_startup as startup_mod
+
+    od_config = types.SimpleNamespace(
+        max_num_seqs=8,
+        step_execution=True,
+        parallel_config=types.SimpleNamespace(world_size=1),
+    )
+    monkeypatch.setattr(startup_mod, "build_diffusion_config", lambda *args: od_config)
+    monkeypatch.setattr(startup_mod, "acquire_device_locks", lambda *args: [])
+    monkeypatch.setattr(
+        startup_mod,
+        "register_stage_with_omni_master",
+        lambda **kwargs: types.SimpleNamespace(
+            handshake_address="tcp://127.0.0.1:26001",
+            input_address="tcp://127.0.0.1:26002",
+            output_address="tcp://127.0.0.1:26003",
+        ),
+    )
+    omni_master_server = types.SimpleNamespace(
+        address="127.0.0.1",
+        port=25000,
+        release_route_port_reservations=lambda *args, **kwargs: None,
+    )
+    proc_manager = types.SimpleNamespace(
+        addresses=types.SimpleNamespace(
+            inputs=["tcp://127.0.0.1:26002"],
+            outputs=["tcp://127.0.0.1:26003"],
+        )
+    )
+    monkeypatch.setattr(proc_mod, "StageDiffusionProcManager", lambda **kwargs: proc_manager)
+    monkeypatch.setattr(
+        client_mod.StageDiffusionClient,
+        "from_addresses",
+        lambda metadata, **kwargs: object(),
+    )
+
+    startup_mod.launch_diffusion_stage_replica(
+        model="dummy-model",
+        stage_config=types.SimpleNamespace(),
+        metadata=types.SimpleNamespace(stage_id=0),
+        stage_init_timeout=12,
+        batch_size=4,
+        use_inline=False,
+        omni_master_server=omni_master_server,
+    )
+
+    assert od_config.max_num_seqs == 8
+
+
 def test_stage_runtime_initializes_stage_pools(monkeypatch):
     import vllm_omni.engine.stage_runtime as runtime_mod
 

tests/engine/test_single_stage_mode.py

@@ -1175,7 +1175,7 @@ class TestSingleStageReplicaInitialization:
                 os.environ[device_env_var] = prev_device_env
 
         assert result is sentinel_client
-        assert od_config.max_num_seqs == 4
+        assert od_config.max_num_seqs is None
         mock_register.assert_called_once_with(
             omni_master_address="127.0.0.1",
             omni_master_port=25000,