anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

[Config] Remove diffusion_batch_size and use max_num_seqs instead. (#6484)

Browse Source 
Signed-off-by: LyxWxj <yuxuanlou621@foxmail.com>
Signed-off-by: Yuxuan Rho <104440472+LyxWxj@users.noreply.github.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
This commit is contained in:
Yuxuan RhoandHongsheng Liu authored and GitHub committed 2026-09-01 01:02:54 +00:00
1 parent 8e26dce1f4
commit 7992c2213e
16 files changed
+36 -95
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

tests/diffusion/batching/test_diffusion_batching.py
+6 -6
View File
@@ -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",
tests/diffusion/test_diffusion_streaming_output.py
+1 -1
View File
@@ -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,
tests/diffusion/test_inline_stage_diffusion_client.py
+1 -2
View File
@@ -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,
)
tests/diffusion/test_multiproc_engine_concurrency.py
-1
View File
@@ -1171,7 +1171,6 @@ class TestStageDiffusionClientErrorPropagation:
metadata,
"tcp://req",
"tcp://resp",
batch_size=1,
)
@pytest.mark.asyncio
tests/diffusion/test_prompt_update.py
-1
View File
@@ -387,7 +387,6 @@ class TestPromptUpdateIntegration:
model="test_model",
od_config=od_config,
metadata=metadata,
batch_size=1,
)
@staticmethod
tests/engine/test_async_omni_engine_stage_init.py
+7 -29
View File
@@ -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
tests/engine/test_single_stage_mode.py
+2 -13
View File
@@ -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",
tests/entrypoints/test_omni_base_profiler.py
+10
View File
@@ -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")
vllm_omni/deploy/wan2_2_ti2v.yaml
-1
View File
@@ -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
vllm_omni/diffusion/inline_stage_diffusion_client.py
+2 -4
View File
@@ -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:
vllm_omni/diffusion/stage_diffusion_client.py
+2 -10
View File
@@ -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:
vllm_omni/engine/async_omni_engine.py
-3
View File
@@ -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,
vllm_omni/engine/stage_engine_startup.py
-3
View File
@@ -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,
vllm_omni/engine/stage_init_utils.py
+1 -5
View File
@@ -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:
vllm_omni/engine/stage_runtime.py
+1 -14
View File
@@ -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,
vllm_omni/entrypoints/omni_base.py
+3 -2
View File
@@ -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,