diff --git a/tests/diffusion/batching/test_diffusion_batching.py b/tests/diffusion/batching/test_diffusion_batching.py index 60136ce4d..cbcbc5e1d 100644 --- a/tests/diffusion/batching/test_diffusion_batching.py +++ b/tests/diffusion/batching/test_diffusion_batching.py @@ -260,7 +260,7 @@ async def compare_single_vs_parallel( ) -> None: """Run the same prompts sequentially then in parallel and print a comparison.""" - omni = AsyncOmni(model=model, diffusion_batch_size=batch_size) + omni = AsyncOmni(model=model, max_num_seqs=batch_size, diffusion_batch_size=batch_size) try: await warmup(omni, WARMUP_PROMPTS) single_time = await run_single(omni, prompts) @@ -288,7 +288,7 @@ async def main(model: str, num_prompts: int, mode: str, batch_size: int = 1) -> await compare_single_vs_parallel(model, prompts, batch_size=batch_size) return - omni = AsyncOmni(model=model, diffusion_batch_size=batch_size) + omni = AsyncOmni(model=model, max_num_seqs=batch_size, diffusion_batch_size=batch_size) try: await warmup(omni, WARMUP_PROMPTS) @@ -524,7 +524,12 @@ if __name__ == "__main__": parser = argparse.ArgumentParser(description="E2E diffusion concurrent benchmark / validation") parser.add_argument("--model", type=str, required=True, help="Model name or path") parser.add_argument("--num-prompts", type=int, default=8, help="Number of prompts to run") - parser.add_argument("--batch-size", type=int, default=1, help="Diffusion batch size (1 = no batching)") + parser.add_argument( + "--batch-size", + type=int, + default=1, + help="Scheduler in-flight width (--max-num-seqs) and client diffusion_batch_size", + ) parser.add_argument( "--mode", choices=["batch", "single", "compare", "validate"], diff --git a/tests/engine/test_async_omni_engine_stage_init.py b/tests/engine/test_async_omni_engine_stage_init.py index a5447a652..1b3c3be40 100644 --- a/tests/engine/test_async_omni_engine_stage_init.py +++ b/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 diff --git a/tests/engine/test_single_stage_mode.py b/tests/engine/test_single_stage_mode.py index 64ae6627c..039b5119a 100644 --- a/tests/engine/test_single_stage_mode.py +++ b/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, diff --git a/vllm_omni/deploy/wan2_2_ti2v.yaml b/vllm_omni/deploy/wan2_2_ti2v.yaml index bdb5cfecb..4aa723eee 100644 --- a/vllm_omni/deploy/wan2_2_ti2v.yaml +++ b/vllm_omni/deploy/wan2_2_ti2v.yaml @@ -1,6 +1,7 @@ # Wan2.2 TI2V-5B deploy. # # Single GPU by default. Increase tensor_parallel_size and devices for A14B MoE variants. +# Request-level batching uses scheduler max_num_seqs, not diffusion_batch_size. async_chunk: false trust_remote_code: true distributed_executor_backend: mp @@ -8,7 +9,7 @@ distributed_executor_backend: mp stages: - stage_id: 0 devices: "0" - max_num_seqs: 1 + max_num_seqs: 4 # in-flight request batch width (CLI: --max-num-seqs) tensor_parallel_size: 1 enforce_eager: true # Keep VAE tiling on by default for higher resolutions / longer clips. diff --git a/vllm_omni/engine/stage_engine_startup.py b/vllm_omni/engine/stage_engine_startup.py index 081b14703..5dacc31df 100644 --- a/vllm_omni/engine/stage_engine_startup.py +++ b/vllm_omni/engine/stage_engine_startup.py @@ -1558,7 +1558,6 @@ def launch_diffusion_stage_replica( from vllm_omni.diffusion.stage_diffusion_client import StageDiffusionClient od_config = build_diffusion_config(model, stage_config, metadata) - od_config.max_num_seqs = batch_size parallel_config = getattr(od_config, "parallel_config", None) world_size = getattr(parallel_config, "world_size", 1) try: diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py index a4c3e1650..263804fd6 100644 --- a/vllm_omni/engine/stage_init_utils.py +++ b/vllm_omni/engine/stage_init_utils.py @@ -1506,15 +1506,14 @@ def initialize_diffusion_stage( stage_cfg: Stage configuration. metadata: Extracted stage metadata. stage_init_timeout: Timeout in seconds for stage initialization handshake - batch_size: Maximum number of requests to batch together in the - diffusion engine. Passed through to ``StageDiffusionClient`` - and ultimately to ``AsyncOmni``. + batch_size: Client-side request batch width. Does not set scheduler + ``max_num_seqs``; pass ``--max-num-seqs`` or stage YAML for that. + Forwarded to ``StageDiffusionClient``. use_inline: If True, uses the inline diffusion client instead of subprocess. """ from vllm_omni.diffusion.stage_diffusion_client import create_diffusion_client od_config = build_diffusion_config(model, stage_cfg, metadata) - od_config.max_num_seqs = batch_size return create_diffusion_client(model, od_config, metadata, stage_init_timeout, batch_size, use_inline)