diff --git a/tests/diffusion/batching/test_diffusion_batching.py b/tests/diffusion/batching/test_diffusion_batching.py index cbcbc5e1d..c9ef78383 100644 --- a/tests/diffusion/batching/test_diffusion_batching.py +++ b/tests/diffusion/batching/test_diffusion_batching.py @@ -7,7 +7,7 @@ End-to-end test for diffusion batching via AsyncOmni. This test fires multiple concurrent ``AsyncOmni.generate()`` calls for a diffusion model and validates that every caller receives its correct individual result. When the underlying diffusion stage is configured with -``batch_size > 1`` (via stage config or ``StageDiffusionClient``), the +``max_num_seqs > 1``, the requests will be batched internally. Even without explicit batching config this test is useful for verifying @@ -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, max_num_seqs=batch_size, diffusion_batch_size=batch_size) + omni = AsyncOmni(model=model, max_num_seqs=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, max_num_seqs=batch_size, diffusion_batch_size=batch_size) + omni = AsyncOmni(model=model, max_num_seqs=batch_size) try: await warmup(omni, WARMUP_PROMPTS) @@ -389,7 +389,7 @@ def test_diffusion_batching_async_concurrent(model_name: str): fired via asyncio.gather. Each request_id must appear in the results.""" async def _inner(): - omni = AsyncOmni(model=model_name, diffusion_batch_size=1) + omni = AsyncOmni(model=model_name, max_num_seqs=1) try: prompts = TEST_PROMPTS[:4] sp = _default_sampling_params() @@ -430,7 +430,7 @@ def test_diffusion_batching_list_prompt_rejected(model_name: str): """ async def _inner(): - omni = AsyncOmni(model=model_name, diffusion_batch_size=4) + omni = AsyncOmni(model=model_name, max_num_seqs=4) try: prompts = TEST_PROMPTS[:4] sp = _default_sampling_params() @@ -528,7 +528,7 @@ if __name__ == "__main__": "--batch-size", type=int, default=1, - help="Scheduler in-flight width (--max-num-seqs) and client diffusion_batch_size", + help="Scheduler in-flight width (--max-num-seqs)", ) parser.add_argument( "--mode", diff --git a/tests/diffusion/test_diffusion_streaming_output.py b/tests/diffusion/test_diffusion_streaming_output.py index 2e1d24540..c833ae7a1 100644 --- a/tests/diffusion/test_diffusion_streaming_output.py +++ b/tests/diffusion/test_diffusion_streaming_output.py @@ -441,7 +441,6 @@ class TestPipelineStreamingOutputToEntrypoint: model="test_model", od_config=od_config, metadata=metadata, - batch_size=1, ) @staticmethod @@ -526,6 +525,7 @@ class TestSupportedPipelines: runner = object.__new__(DiffusionModelRunner) runner.vllm_config = _make_vllm_config() runner.od_config = SimpleNamespace( + max_num_seqs=1, enable_cpu_offload=False, enable_layerwise_offload=False, enforce_eager=True, diff --git a/tests/diffusion/test_inline_stage_diffusion_client.py b/tests/diffusion/test_inline_stage_diffusion_client.py index 25d4f1feb..e39c18ff4 100644 --- a/tests/diffusion/test_inline_stage_diffusion_client.py +++ b/tests/diffusion/test_inline_stage_diffusion_client.py @@ -47,7 +47,7 @@ def client(mock_engine): with patch.object(InlineStageDiffusionClient, "_enrich_config"): od_config = MagicMock(spec=OmniDiffusionConfig) od_config.streaming_output = False - c = InlineStageDiffusionClient(model="test_model", od_config=od_config, metadata=metadata, batch_size=1) + c = InlineStageDiffusionClient(model="test_model", od_config=od_config, metadata=metadata) yield c c.shutdown() @@ -234,5 +234,4 @@ def test_inline_client_requires_replica_id(mock_engine): model="test_model", od_config=od_config, metadata=metadata, - batch_size=1, ) diff --git a/tests/diffusion/test_multiproc_engine_concurrency.py b/tests/diffusion/test_multiproc_engine_concurrency.py index 67e6f4751..855f38c4a 100644 --- a/tests/diffusion/test_multiproc_engine_concurrency.py +++ b/tests/diffusion/test_multiproc_engine_concurrency.py @@ -1171,7 +1171,6 @@ class TestStageDiffusionClientErrorPropagation: metadata, "tcp://req", "tcp://resp", - batch_size=1, ) @pytest.mark.asyncio diff --git a/tests/diffusion/test_prompt_update.py b/tests/diffusion/test_prompt_update.py index cdee35979..f08825d1a 100644 --- a/tests/diffusion/test_prompt_update.py +++ b/tests/diffusion/test_prompt_update.py @@ -387,7 +387,6 @@ class TestPromptUpdateIntegration: model="test_model", od_config=od_config, metadata=metadata, - batch_size=1, ) @staticmethod diff --git a/tests/engine/test_async_omni_engine_stage_init.py b/tests/engine/test_async_omni_engine_stage_init.py index 859e03289..c1e56b7e1 100644 --- a/tests/engine/test_async_omni_engine_stage_init.py +++ b/tests/engine/test_async_omni_engine_stage_init.py @@ -172,7 +172,6 @@ def _make_stage_runtime() -> StageRuntime: model="dummy-model", config_path="dummy-config", stage_init_timeout=1, - diffusion_batch_size=1, async_chunk=False, ) @@ -196,7 +195,6 @@ def test_async_omni_engine_initialize_stages_passes_log_stats_to_runtime(monkeyp engine.model = "dummy-model" engine.config_path = "dummy-config" engine.single_stage_mode = False - engine.diffusion_batch_size = 1 engine.async_chunk = False engine.tokenizer = None engine._single_stage_id_filter = None @@ -282,7 +280,6 @@ def test_build_logical_stage_init_plans_handles_stage_without_devices(monkeypatc model="dummy-model", config_path="dummy-config", stage_init_timeout=1, - diffusion_batch_size=1, async_chunk=False, ) @@ -371,7 +368,6 @@ def test_initialize_local_diffusion_replica_restores_device_visibility_after_loc model="dummy-model", config_path="dummy-config", stage_init_timeout=1, - diffusion_batch_size=1, async_chunk=False, ) @@ -416,7 +412,6 @@ def test_initialize_local_diffusion_replica_passes_stage_init_timeout_and_inline model="dummy-model", config_path="dummy-config", stage_init_timeout=1, - diffusion_batch_size=4, async_chunk=False, ) @@ -429,7 +424,6 @@ def test_initialize_local_diffusion_replica_passes_stage_init_timeout_and_inline def _capture_launch_diffusion_stage_replica(**kwargs): captured["stage_id"] = kwargs["metadata"].stage_id captured["stage_init_timeout"] = kwargs["stage_init_timeout"] - captured["batch_size"] = kwargs["batch_size"] captured["use_inline"] = kwargs["use_inline"] captured["omni_master_server"] = kwargs["omni_master_server"] return types.SimpleNamespace(), StageReplicaResources() @@ -441,7 +435,6 @@ def test_initialize_local_diffusion_replica_passes_stage_init_timeout_and_inline assert captured == { "stage_id": 0, "stage_init_timeout": 302, - "batch_size": 4, "use_inline": expected_inline, "omni_master_server": None, } @@ -511,21 +504,20 @@ def test_initialize_local_llm_replica_scopes_runtime_env(monkeypatch): assert os.environ[runtime_env_var] == "parent-value" -def test_initialize_diffusion_stage_forwards_batch_size_without_writing_max_num_seqs(monkeypatch): +def test_initialize_diffusion_stage_preserves_configured_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=1) + od_config = types.SimpleNamespace(max_num_seqs=4) captured: dict[str, object] = {} monkeypatch.setattr(init_mod, "build_diffusion_config", lambda *args: od_config) - def _capture_client(model, config, metadata, stage_init_timeout, batch_size, use_inline): + def _capture_client(model, config, metadata, stage_init_timeout, use_inline): captured.update( model=model, config=config, metadata=metadata, stage_init_timeout=stage_init_timeout, - batch_size=batch_size, use_inline=use_inline, ) return object() @@ -539,27 +531,25 @@ def test_initialize_diffusion_stage_forwards_batch_size_without_writing_max_num_ types.SimpleNamespace(), metadata, stage_init_timeout=12, - batch_size=4, use_inline=True, ) - assert od_config.max_num_seqs == 1 + assert od_config.max_num_seqs == 4 assert captured == { "model": "dummy-model", "config": od_config, "metadata": metadata, "stage_init_timeout": 12, - "batch_size": 4, "use_inline": True, } -def test_launch_diffusion_stage_replica_does_not_write_max_num_seqs_from_batch_size(monkeypatch): +def test_launch_diffusion_stage_replica_preserves_configured_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=1, parallel_config=types.SimpleNamespace(world_size=1)) + od_config = types.SimpleNamespace(max_num_seqs=4, 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( @@ -596,13 +586,12 @@ def test_launch_diffusion_stage_replica_does_not_write_max_num_seqs_from_batch_s 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 result is sentinel_client - assert od_config.max_num_seqs == 1 + assert od_config.max_num_seqs == 4 assert resources.manager is proc_manager @@ -621,7 +610,6 @@ def test_initialize_diffusion_stage_does_not_write_max_num_seqs(monkeypatch): types.SimpleNamespace(), metadata, stage_init_timeout=12, - batch_size=4, use_inline=True, ) @@ -672,7 +660,6 @@ def test_launch_diffusion_stage_replica_preserves_step_execution_max_num_seqs(mo 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, ) @@ -688,7 +675,6 @@ def test_stage_runtime_initializes_stage_pools(monkeypatch): model="dummy-model", config_path="dummy-config", stage_init_timeout=1, - diffusion_batch_size=1, async_chunk=False, ) @@ -754,7 +740,6 @@ def test_stage_runtime_passes_log_stats_to_llm_replica_launch(monkeypatch): model="dummy-model", config_path="dummy-config", stage_init_timeout=1, - diffusion_batch_size=1, async_chunk=False, log_stats=True, ) @@ -798,7 +783,6 @@ def test_stage_runtime_passes_log_stats_to_output_processor(monkeypatch): model="dummy-model", config_path="dummy-config", stage_init_timeout=1, - diffusion_batch_size=1, async_chunk=False, log_stats=True, ) @@ -851,7 +835,6 @@ def test_build_logical_stage_init_plans_applies_replica_device_splits(monkeypatc model="dummy-model", config_path="dummy-config", stage_init_timeout=1, - diffusion_batch_size=1, async_chunk=False, ) @@ -892,7 +875,6 @@ def test_initialize_stage_replicas_collects_results_by_stage_and_replica_id(monk model="dummy-model", config_path="dummy-config", stage_init_timeout=123, - diffusion_batch_size=1, async_chunk=False, ) @@ -930,7 +912,6 @@ def test_remote_replicas_use_distinct_init_group_keys(): model="dummy-model", config_path="dummy-config", stage_init_timeout=123, - diffusion_batch_size=1, async_chunk=False, ) plan = _make_llm_plan( @@ -956,7 +937,6 @@ def test_initialize_stages_cleans_up_successful_replicas_after_partial_multi_rep model="dummy-model", config_path="dummy-config", stage_init_timeout=1, - diffusion_batch_size=1, async_chunk=False, ) @@ -993,7 +973,6 @@ def test_initialize_stages_cleans_up_late_successful_replicas_after_early_multi_ model="dummy-model", config_path="dummy-config", stage_init_timeout=1, - diffusion_batch_size=1, async_chunk=False, ) @@ -1032,7 +1011,6 @@ def test_initialize_local_llm_replica_passes_stage_init_timeout_to_complete_stag model="dummy-model", config_path="dummy-config", stage_init_timeout=1, - diffusion_batch_size=1, async_chunk=False, ) stage_init_timeout = 302 diff --git a/tests/engine/test_single_stage_mode.py b/tests/engine/test_single_stage_mode.py index 334001a65..cacda1f60 100644 --- a/tests/engine/test_single_stage_mode.py +++ b/tests/engine/test_single_stage_mode.py @@ -575,7 +575,6 @@ class TestSingleStageInitialization: model="fake-model", config_path="/fake/stages.yaml", stage_init_timeout=60, - diffusion_batch_size=2, async_chunk=False, single_stage_id_filter=stage_id_filter, omni_master_address="127.0.0.1", @@ -807,7 +806,6 @@ class TestSingleStageInitialization: model="fake-model", config_path="/fake/stages.yaml", stage_init_timeout=60, - diffusion_batch_size=2, async_chunk=False, ) mocker.patch.object(runtime_mod, "prepare_engine_environment") @@ -860,7 +858,6 @@ class TestSingleStageReplicaInitialization: model="fake-model", config_path="/fake/stages.yaml", stage_init_timeout=60, - diffusion_batch_size=2, async_chunk=False, single_stage_id_filter=None, omni_master_address="127.0.0.1", @@ -928,7 +925,6 @@ class TestSingleStageReplicaInitialization: model="fake-model", config_path="/fake/stages.yaml", stage_init_timeout=60, - diffusion_batch_size=2, async_chunk=False, single_stage_id_filter=None, omni_master_address="127.0.0.1", @@ -950,7 +946,6 @@ class TestSingleStageReplicaInitialization: model="fake-model", config_path="/fake/stages.yaml", stage_init_timeout=60, - diffusion_batch_size=2, async_chunk=False, single_stage_id_filter=None, omni_master_address="127.0.0.1", @@ -1002,7 +997,6 @@ class TestSingleStageReplicaInitialization: model="fake-model", config_path="/fake/stages.yaml", stage_init_timeout=60, - diffusion_batch_size=2, async_chunk=False, single_stage_id_filter=None, omni_master_address="127.0.0.1", @@ -1068,7 +1062,6 @@ class TestSingleStageReplicaInitialization: model="fake-model", config_path="/fake/stages.yaml", stage_init_timeout=60, - diffusion_batch_size=4, async_chunk=False, single_stage_id_filter=None, omni_master_address="127.0.0.1", @@ -1120,7 +1113,6 @@ class TestSingleStageReplicaInitialization: model="fake-model", config_path="/fake/stages.yaml", stage_init_timeout=60, - diffusion_batch_size=4, async_chunk=False, single_stage_id_filter=None, omni_master_address="127.0.0.1", @@ -1141,7 +1133,7 @@ class TestSingleStageReplicaInitialization: runtime._init_visible_devices_baseline = "0" mocker.patch.object(runtime_mod, "inject_kv_stage_info") - od_config = SimpleNamespace(max_num_seqs=None, parallel_config=SimpleNamespace(world_size=1)) + od_config = SimpleNamespace(max_num_seqs=4, parallel_config=SimpleNamespace(world_size=1)) mocker.patch("vllm_omni.engine.stage_engine_startup.build_diffusion_config", return_value=od_config) mock_register = mocker.patch( "vllm_omni.engine.stage_engine_startup.register_stage_with_omni_master", @@ -1178,7 +1170,7 @@ class TestSingleStageReplicaInitialization: os.environ[device_env_var] = prev_device_env assert result is sentinel_client - assert od_config.max_num_seqs is None + assert od_config.max_num_seqs == 4 mock_register.assert_called_once_with( omni_master_address="127.0.0.1", omni_master_port=25000, @@ -1203,7 +1195,6 @@ class TestSingleStageReplicaInitialization: request_address="tcp://127.0.0.1:26002", response_address="tcp://127.0.0.1:26003", proc_manager=mocker.ANY, - batch_size=4, ) assert mock_from_addresses.call_args.kwargs["proc_manager"].proc is proc @@ -1216,7 +1207,6 @@ class TestSingleStageReplicaInitialization: model="fake-model", config_path="/fake/stages.yaml", stage_init_timeout=60, - diffusion_batch_size=4, async_chunk=False, single_stage_id_filter=None, omni_master_address="127.0.0.1", @@ -1258,7 +1248,6 @@ class TestSingleStageReplicaInitialization: model="fake-model", config_path="/fake/stages.yaml", stage_init_timeout=60, - diffusion_batch_size=4, async_chunk=False, single_stage_id_filter=None, omni_master_address="127.0.0.1", diff --git a/tests/entrypoints/test_omni_base_profiler.py b/tests/entrypoints/test_omni_base_profiler.py index 3f793e89e..827ab67b6 100644 --- a/tests/entrypoints/test_omni_base_profiler.py +++ b/tests/entrypoints/test_omni_base_profiler.py @@ -46,6 +46,16 @@ class TestOmniBaseProfiler: stage_ids=None, ) + def test_removed_diffusion_batch_size_fails_loudly(self, mock_engine, mocker: MockerFixture): + """Legacy diffusion_batch_size callers receive an actionable error.""" + mocker.patch("vllm_omni.entrypoints.omni_base.AsyncOmniEngine", return_value=mock_engine) + mocker.patch("vllm_omni.entrypoints.omni_base.omni_snapshot_download", side_effect=lambda x: x) + mocker.patch("vllm_omni.entrypoints.omni_base.weakref.finalize") + from vllm_omni.entrypoints.omni_base import OmniBase + + with pytest.raises(TypeError, match="diffusion_batch_size.*max_num_seqs"): + OmniBase(model="test-model", diffusion_batch_size=8) + def test_start_profile_with_prefix(self, omni_base_instance, mock_engine): """Test that start_profile passes profile_prefix to collective_rpc.""" omni_base_instance.start_profile(profile_prefix="test_trace") diff --git a/vllm_omni/deploy/wan2_2_ti2v.yaml b/vllm_omni/deploy/wan2_2_ti2v.yaml index 4aa723eee..0d0a9124c 100644 --- a/vllm_omni/deploy/wan2_2_ti2v.yaml +++ b/vllm_omni/deploy/wan2_2_ti2v.yaml @@ -1,7 +1,6 @@ # 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 diff --git a/vllm_omni/diffusion/inline_stage_diffusion_client.py b/vllm_omni/diffusion/inline_stage_diffusion_client.py index b8983ab20..3d3bc4683 100644 --- a/vllm_omni/diffusion/inline_stage_diffusion_client.py +++ b/vllm_omni/diffusion/inline_stage_diffusion_client.py @@ -48,7 +48,6 @@ class InlineStageDiffusionClient(StageClientBase): model: str, od_config: OmniDiffusionConfig, metadata: StageMetadata, - batch_size: int = 1, ) -> None: self.model = model self.od_config = od_config @@ -63,7 +62,6 @@ class InlineStageDiffusionClient(StageClientBase): self.prompt_transform_func = None self.prompt_expand_func = None self.engine_input_source = metadata.engine_input_source - self.batch_size = batch_size self._enrich_config() self._engine = DiffusionEngine.make_engine(self.od_config) @@ -79,10 +77,10 @@ class InlineStageDiffusionClient(StageClientBase): self._engine.executor.register_failure_callback(self._mark_engine_dead) logger.info( - "[InlineStageDiffusionClient] stage-%s [rep-%s] initialized inline (batch_size=%d)", + "[InlineStageDiffusionClient] stage-%s [rep-%s] initialized inline (max_num_seqs=%d)", self.stage_id, self.replica_id, - self.batch_size, + self.od_config.max_num_seqs, ) def _enrich_config(self) -> None: diff --git a/vllm_omni/diffusion/stage_diffusion_client.py b/vllm_omni/diffusion/stage_diffusion_client.py index 2bd4f2888..fe7b6c408 100644 --- a/vllm_omni/diffusion/stage_diffusion_client.py +++ b/vllm_omni/diffusion/stage_diffusion_client.py @@ -49,14 +49,13 @@ def create_diffusion_client( od_config: OmniDiffusionConfig, metadata: StageMetadata, stage_init_timeout: int, - batch_size: int = 1, use_inline: bool = False, ) -> Any: """Factory to create either an inline or out-of-process diffusion client.""" if use_inline: from vllm_omni.diffusion.inline_stage_diffusion_client import InlineStageDiffusionClient - return InlineStageDiffusionClient(model, od_config, metadata, batch_size=batch_size) + return InlineStageDiffusionClient(model, od_config, metadata) proc_manager = StageDiffusionProcManager( model=model, od_config=od_config, @@ -67,7 +66,6 @@ def create_diffusion_client( request_address=proc_manager.addresses.inputs[0], response_address=proc_manager.addresses.outputs[0], proc_manager=proc_manager, - batch_size=batch_size, ) @@ -91,14 +89,12 @@ class StageDiffusionClient(StageClientBase): response_address: str, *, proc_manager: StageDiffusionProcManager | None = None, - batch_size: int = 1, ) -> None: self._initialize_client( metadata, request_address, response_address, proc_manager=proc_manager, - batch_size=batch_size, ) @classmethod @@ -109,7 +105,6 @@ class StageDiffusionClient(StageClientBase): response_address: str, *, proc_manager: StageDiffusionProcManager | None = None, - batch_size: int = 1, ) -> StageDiffusionClient: """Create a client for an already-running diffusion subprocess.""" return cls( @@ -117,7 +112,6 @@ class StageDiffusionClient(StageClientBase): request_address, response_address, proc_manager=proc_manager, - batch_size=batch_size, ) def _initialize_client( @@ -127,7 +121,6 @@ class StageDiffusionClient(StageClientBase): response_address: str, *, proc_manager: StageDiffusionProcManager | None = None, - batch_size: int, ) -> None: self._set_stage_metadata(metadata) self._proc_manager = proc_manager @@ -144,11 +137,10 @@ class StageDiffusionClient(StageClientBase): self._start_proc_monitor() logger.info( - "[StageDiffusionClient] stage-%s [rep-%s] initialized (owns_process=%s, batch_size=%d)", + "[StageDiffusionClient] stage-%s [rep-%s] initialized (owns_process=%s)", self.stage_id, self.replica_id, self._proc_manager is not None, - batch_size, ) def _set_stage_metadata(self, metadata: StageMetadata) -> None: diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index 7e0d9f629..1efafbbab 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -129,7 +129,6 @@ class AsyncOmniEngine: model: str, stage_init_timeout: int = 300, init_timeout: int = 600, - diffusion_batch_size: int = 1, single_stage_mode: bool = False, transfer_emitter: Any = None, prom_metrics: Any = None, @@ -140,7 +139,6 @@ class AsyncOmniEngine: ) -> None: self.model = model self.tokenizer = tokenizer - self.diffusion_batch_size = diffusion_batch_size # Cached by get_diffusion_od_config(). self._diffusion_od_config_view: Any = None startup_timeout = int(init_timeout) @@ -320,7 +318,6 @@ class AsyncOmniEngine: config_path=self.config_path, single_stage_mode=self.single_stage_mode, stage_init_timeout=stage_init_timeout, - diffusion_batch_size=self.diffusion_batch_size, async_chunk=self.async_chunk, tokenizer=self.tokenizer, single_stage_id_filter=self._single_stage_id_filter, diff --git a/vllm_omni/engine/stage_engine_startup.py b/vllm_omni/engine/stage_engine_startup.py index e9622b83a..5abcbd249 100644 --- a/vllm_omni/engine/stage_engine_startup.py +++ b/vllm_omni/engine/stage_engine_startup.py @@ -1546,7 +1546,6 @@ def launch_diffusion_stage_replica( stage_config: Any, metadata: Any, stage_init_timeout: int, - batch_size: int, use_inline: bool, replica_id: int = 0, omni_master_server: OmniMasterServer | None = None, @@ -1565,7 +1564,6 @@ def launch_diffusion_stage_replica( stage_config, metadata, stage_init_timeout=stage_init_timeout, - batch_size=batch_size, use_inline=use_inline, ) return client, StageReplicaResources() @@ -1624,7 +1622,6 @@ def launch_diffusion_stage_replica( request_address=proc_manager.addresses.inputs[0], response_address=proc_manager.addresses.outputs[0], proc_manager=proc_manager, - batch_size=batch_size, ) return client, StageReplicaResources( manager=proc_manager, diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py index b1589aef5..6fdcac81b 100644 --- a/vllm_omni/engine/stage_init_utils.py +++ b/vllm_omni/engine/stage_init_utils.py @@ -1813,7 +1813,6 @@ def initialize_diffusion_stage( stage_cfg: Any, metadata: StageMetadata, stage_init_timeout: int, - batch_size: int = 1, use_inline: bool = False, ) -> Any: """Build a diffusion stage client. @@ -1823,15 +1822,12 @@ def initialize_diffusion_stage( stage_cfg: Stage configuration. metadata: Extracted stage metadata. stage_init_timeout: Timeout in seconds for stage initialization handshake - 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) - return create_diffusion_client(model, od_config, metadata, stage_init_timeout, batch_size, use_inline) + return create_diffusion_client(model, od_config, metadata, stage_init_timeout, use_inline) def _stage_declares_cfg_pairs(model_config: Any) -> bool: diff --git a/vllm_omni/engine/stage_runtime.py b/vllm_omni/engine/stage_runtime.py index 6088814ce..049d09688 100644 --- a/vllm_omni/engine/stage_runtime.py +++ b/vllm_omni/engine/stage_runtime.py @@ -87,7 +87,6 @@ class StageRemoteFactoryContext: base_metadata: Any vllm_config: Any | None = None executor_class: type | None = None - diffusion_batch_size: int = 1 def _build_load_balancer_factory(policy: str) -> Callable[[], LoadBalancer]: @@ -124,7 +123,6 @@ class StageRuntime: config_path: str, *, stage_init_timeout: int, - diffusion_batch_size: int, async_chunk: bool, tokenizer: str | None = None, log_stats: bool = False, @@ -133,7 +131,6 @@ class StageRuntime: self._model = model self._config_path = config_path self._stage_init_timeout = stage_init_timeout - self._diffusion_batch_size = diffusion_batch_size self._async_chunk = async_chunk self._tokenizer = tokenizer self._log_stats = log_stats @@ -672,7 +669,6 @@ class StageRuntime: stage_config=plan.stage_cfg, metadata=plan.metadata, stage_init_timeout=stage_init_timeout, - batch_size=self._diffusion_batch_size, use_inline=plan.num_replicas == 1 and bool(self._num_stages == 1 or inline_diffusion or custom_pipeline_args), replica_id=plan.replica_id, @@ -681,10 +677,9 @@ class StageRuntime: ) logger.info( - "[StageRuntime] Stage %s replica %s initialized (diffusion, batch_size=%d)", + "[StageRuntime] Stage %s replica %s initialized (diffusion)", plan.metadata.stage_id, plan.replica_id, - self._diffusion_batch_size, ) return cast(StagePoolClient, client) except Exception: @@ -768,7 +763,6 @@ class DistStageRuntime(StageRuntime): config_path: str, *, stage_init_timeout: int, - diffusion_batch_size: int, async_chunk: bool, single_stage_id_filter: int | None, omni_master_address: str, @@ -785,7 +779,6 @@ class DistStageRuntime(StageRuntime): model=model, config_path=config_path, stage_init_timeout=stage_init_timeout, - diffusion_batch_size=diffusion_batch_size, async_chunk=async_chunk, tokenizer=tokenizer, log_stats=log_stats, @@ -914,7 +907,6 @@ class DistStageRuntime(StageRuntime): base_metadata=metadata, vllm_config=plan.stage_vllm_config, executor_class=plan.executor_class, - diffusion_batch_size=self._diffusion_batch_size, ) return self._create_remote_replica_client(ctx, plan.replica_id) @@ -977,7 +969,6 @@ class DistStageRuntime(StageRuntime): base_metadata=template.metadata, vllm_config=template.stage_vllm_config, executor_class=template.executor_class, - diffusion_batch_size=self._diffusion_batch_size, ) return contexts @@ -1045,7 +1036,6 @@ class DistStageRuntime(StageRuntime): metadata, request_address=resources.addresses.inputs[0], response_address=resources.addresses.outputs[0], - batch_size=ctx.diffusion_batch_size, ) logger.info( "[DistStageRuntime] Remote diffusion replica attached stage=%d replica=%d", @@ -1109,7 +1099,6 @@ def create_stage_runtime( *, single_stage_mode: bool, stage_init_timeout: int, - diffusion_batch_size: int, async_chunk: bool, tokenizer: str | None = None, # Distributed-only params: @@ -1131,7 +1120,6 @@ def create_stage_runtime( model=model, config_path=config_path, stage_init_timeout=stage_init_timeout, - diffusion_batch_size=diffusion_batch_size, async_chunk=async_chunk, tokenizer=tokenizer, log_stats=log_stats, @@ -1148,7 +1136,6 @@ def create_stage_runtime( model=model, config_path=config_path, stage_init_timeout=stage_init_timeout, - diffusion_batch_size=diffusion_batch_size, async_chunk=async_chunk, tokenizer=tokenizer, log_stats=log_stats, diff --git a/vllm_omni/entrypoints/omni_base.py b/vllm_omni/entrypoints/omni_base.py index af86b43d5..0aa6832d0 100644 --- a/vllm_omni/entrypoints/omni_base.py +++ b/vllm_omni/entrypoints/omni_base.py @@ -187,7 +187,9 @@ class OmniBase(PDDisaggregationMixin): # override the deploy YAML's ``async_chunk: true`` default. async_chunk = kwargs.get("async_chunk") output_modalities = kwargs.pop("output_modalities", None) - diffusion_batch_size: int = kwargs.pop("diffusion_batch_size", 1) + + if "diffusion_batch_size" in kwargs: + raise TypeError("`diffusion_batch_size` has been removed in Omni/AsyncOmni. Use `max_num_seqs`.") if "log_requests" in kwargs: raise TypeError("`log_requests` has been removed in Omni/AsyncOmni. Use `log_stats`.") @@ -213,7 +215,6 @@ class OmniBase(PDDisaggregationMixin): model=model, init_timeout=init_timeout, stage_init_timeout=stage_init_timeout, - diffusion_batch_size=diffusion_batch_size, transfer_emitter=self.transfer_metrics, prom_metrics=self.prom_metrics, log_stats=log_stats,