diff --git a/tests/config/test_omni_config.py b/tests/config/test_omni_config.py index 3d356ad91..1f6111d64 100644 --- a/tests/config/test_omni_config.py +++ b/tests/config/test_omni_config.py @@ -4,6 +4,7 @@ from __future__ import annotations +import warnings from dataclasses import fields from inspect import Parameter, signature from multiprocessing.reduction import ForkingPickler @@ -24,6 +25,7 @@ from vllm.config import SchedulerConfig as VllmSchedulerConfig from tests.helpers.stage_config import get_deploy_config_path from vllm_omni.config import omni_config as omni_config_module +from vllm_omni.config.config_factory import StageConfigFactory from vllm_omni.config.omni_config import ( BaseVllmOmniStageConfig, OmniStageCacheConfig, @@ -150,6 +152,30 @@ def test_nested_stage_override_deep_merges_structured_model_config() -> None: assert model_config["policy_server_config"]["action_space"] == "joint_position" +def _build_single_diffusion_config( + *, + cli_overrides: dict | None = None, + engine_extras: dict | None = None, +) -> VllmOmniConfig: + pipeline = PipelineConfig( + model_type="diffusion-owner-test", + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="diffusion", + execution_type=StageExecutionType.DIFFUSION, + final_output=True, + ), + ), + ) + deploy_stages = [StageDeployConfig(stage_id=0, engine_extras=engine_extras)] if engine_extras is not None else [] + return VllmOmniConfig.from_pipeline_config( + pipeline, + user_deploy_config=DeployConfig(async_chunk=False, stages=deploy_stages), + cli_overrides=cli_overrides, + ) + + @pytest.mark.parametrize("model_type", sorted(OMNI_PIPELINES)) def test_vllm_omni_config_from_pipeline_config_matches_merge_pipeline_deploy(model_type: str): pipeline = _resolve_pipeline_or_skip(model_type) @@ -475,7 +501,7 @@ def test_from_pipeline_config_normalizes_nested_parallel_config_aliases(field_na assert "data_parallel_master_ip" in stage.parallel_config._omni_explicit_fields -def test_from_pipeline_config_accepts_diffusion_only_cli_fields_for_diffusion_stage(): +def test_from_pipeline_config_keeps_global_kv_cache_dtype_outside_diffusion_stage(): stage = _from_pipeline_key( "dreamzero", deploy_config_path="dreamzero_tp1_cfg2", @@ -483,11 +509,57 @@ def test_from_pipeline_config_accepts_diffusion_only_cli_fields_for_diffusion_st ).stage_by_id(0) assert isinstance(stage, VllmOmniDiffusionStageConfig) - assert stage.diffusion_config.diffusion_kv_cache_dtype == "fp8" + assert stage.diffusion_config.diffusion_kv_cache_dtype is None def test_stage_cli_field_selection_defers_ownership_validation_until_sources_are_merged(): - assert omni_config_module._stage_cli_overrides(0, {"enable_lora": True}) == {"enable_lora": True} + assert omni_config_module._stage_cli_overrides( + 0, + {"enable_lora": True}, + ) == {"enable_lora": True} + + +@pytest.mark.parametrize( + ("cli_overrides", "expected_dtype"), + [ + ({}, torch.float32), + ({"dtype": None}, torch.float32), + ({"dtype": "float16"}, torch.float16), + ({"dtype": "auto"}, torch.bfloat16), + ({"stage_0_dtype": "bfloat16"}, torch.bfloat16), + ], +) +def test_diffusion_deploy_dtype_survives_unset_cli_overrides(cli_overrides, expected_dtype): + stage = _build_single_diffusion_config( + engine_extras={"dtype": "float32", "cache_backend": "tea_cache"}, + cli_overrides=cli_overrides, + ).stage_by_id(0) + + assert stage.diffusion_config.dtype == expected_dtype + assert stage.diffusion_config.cache_backend == "tea_cache" + + +@pytest.mark.parametrize( + "coordination_kwargs", + [ + {"omni_master_address": "127.0.0.1", "omni_master_port": 30000}, + {"omni_dp_size_local": 2, "omni_heartbeat_timeout": 60.0, "omni_lb_policy": "random"}, + ], +) +def test_diffusion_ingress_routes_coordination_fields_away(coordination_kwargs): + normalize = omni_config_module.normalize_and_validate_diffusion_engine_ingress_kwargs + assert normalize({**coordination_kwargs, "enable_sleep_mode": True}, stage_id=0) == {"enable_sleep_mode": True} + with pytest.raises(ValueError, match="omni_master_unknown_field"): + normalize({**coordination_kwargs, "omni_master_unknown_field": "unknown"}, stage_id=0) + + +def test_diffusion_ingress_defers_defaults(monkeypatch): + monkeypatch.setenv("DIFFUSION_CACHE_BACKEND", "tea_cache") + + assert omni_config_module.normalize_and_validate_diffusion_engine_ingress_kwargs({}, stage_id=0) == {} + assert omni_config_module.normalize_and_validate_diffusion_engine_ingress_kwargs( + {"dtype": None, "cache_backend": None}, stage_id=0 + ) == {"dtype": None, "cache_backend": None} @pytest.mark.parametrize( @@ -1415,18 +1487,19 @@ def test_diffusion_config_from_kwargs_reuses_legacy_normalization(monkeypatch): monkeypatch.setenv("DIFFUSION_CACHE_BACKEND", "TEA_CACHE") monkeypatch.setattr(current_omni_platform, "is_cuda", lambda: True) - cfg = omni_config_module._DiffusionConfigProjection.from_kwargs( - diffusion_attention_backend="flash_attn", - fa_deterministic=True, - kv_cache_dtype="fp8", - kv_cache_skip_steps="0-1", - kv_cache_skip_layers=[2], - static_lora_scale=0.25, - diffusion_kv_mode="paged_scheduler", - diffusion_kv_max_rows_per_request=2, - diffusers_load_kwargs=None, - diffusers_call_kwargs=None, - ) + with pytest.warns(FutureWarning): + cfg = omni_config_module._DiffusionConfigProjection.from_kwargs( + diffusion_attention_backend="flash_attn", + fa_deterministic=True, + kv_cache_dtype="fp8", + kv_cache_skip_steps="0-1", + kv_cache_skip_layers=[2], + static_lora_scale=0.25, + diffusion_kv_mode="paged_scheduler", + diffusion_kv_max_rows_per_request=2, + diffusers_load_kwargs=None, + diffusers_call_kwargs=None, + ) assert cfg.diffusion_attention_config.default.backend == "flash_attn" assert cfg.fa_deterministic is True @@ -1463,47 +1536,6 @@ def test_diffusion_config_none_values_preserve_dataclass_defaults(): assert config.enable_sleep_mode is False -@pytest.mark.parametrize( - ("canonical_key", "alias_key", "canonical_value", "alias_value"), - [ - ("lora_scale", "static_lora_scale", 0.75, 0.25), - ( - "quantization_config", - "diffusion_quantization_config", - {"method": "canonical"}, - {"method": "diffusion-alias"}, - ), - ( - "quantization_config", - "quantization", - {"method": "canonical"}, - "legacy-alias", - ), - ("diffusion_kv_cache_dtype", "kv_cache_dtype", "fp8", "fp16"), - ("diffusion_kv_cache_skip_steps", "kv_cache_skip_steps", "0-1", "2-3"), - ("diffusion_kv_cache_skip_layers", "kv_cache_skip_layers", "1-2", "3-4"), - ("streaming_output", "diffusion_streaming_output", False, True), - ], -) -def test_diffusion_alias_conflicts_prefer_canonical_key( - canonical_key, - alias_key, - canonical_value, - alias_value, -): - from vllm_omni.diffusion.data import normalize_omni_diffusion_kwargs - - normalized = normalize_omni_diffusion_kwargs( - { - canonical_key: canonical_value, - alias_key: alias_value, - } - ) - - assert normalized[canonical_key] == canonical_value - assert alias_key not in normalized - - def test_from_pipeline_config_normalizes_diffusion_config_aliases_from_engine_args(tmp_path, monkeypatch): from vllm_omni.platforms import current_omni_platform @@ -1528,10 +1560,11 @@ def test_from_pipeline_config_normalizes_diffusion_config_aliases_from_engine_ar ) ) - stage = _from_pipeline_key( - "dreamzero", - deploy_config_path=str(deploy_path), - ).stage_by_id(0) + with pytest.warns(FutureWarning): + stage = _from_pipeline_key( + "dreamzero", + deploy_config_path=str(deploy_path), + ).stage_by_id(0) assert isinstance(stage, VllmOmniDiffusionStageConfig) assert stage.diffusion_config.diffusion_attention_config.default.backend == "flash_attn" @@ -1544,7 +1577,10 @@ def test_from_pipeline_config_normalizes_diffusion_config_aliases_from_engine_ar from vllm_omni.engine.stage_init_utils import build_engine_args_dict_from_omni_stage_config engine_args = build_engine_args_dict_from_omni_stage_config(stage, model=str(tmp_path)) - od_config = OmniDiffusionConfig.from_kwargs(**engine_args) + diffusion_kwargs = omni_config_module.extract_diffusion_stage_config_kwargs( + engine_args, stage_id=stage.stage_id, include_engine_adapter_metadata=True + ) + od_config = OmniDiffusionConfig.from_kwargs(**diffusion_kwargs) assert od_config.kv_transfer_config.engine_id == "dit-engine-1" @@ -1697,10 +1733,8 @@ def test_video_output_transport_survives_stage_override_filtering() -> None: def test_video_output_transport_reaches_default_diffusion_stage() -> None: - from vllm_omni.engine.async_omni_engine import AsyncOmniEngine - transport = {"enable_device_postprocess": True} - stages = AsyncOmniEngine._create_default_diffusion_stage_cfg( + stages = StageConfigFactory.create_default_diffusion( { "model": "unused", "model_class_name": "UnknownPipeline", @@ -1830,3 +1864,95 @@ def test_diffusion_quantization_origin_survives_projection_and_transport(monkeyp restored = ForkingPickler.loads(ForkingPickler.dumps(cfg)) assert restored.quantization_config_is_auto_detected is not explicit assert restored.quantization_config.w4a8_fallback_steps == ([] if explicit else [37]) + + +@pytest.mark.parametrize( + "factory", + [ + omni_config_module._DiffusionConfigProjection.from_kwargs, + pytest.param( + lambda **kwargs: __import__( + "vllm_omni.diffusion.data", + fromlist=["OmniDiffusionConfig"], + ).OmniDiffusionConfig.from_kwargs(**kwargs), + id="runtime-config", + ), + ], +) +def test_diffusion_config_entrypoints_reject_unknown_none(factory): + with pytest.raises((ValueError, ValidationError), match="enable_sleep_mod"): + factory(enable_sleep_mod=None) + + +def test_diffusion_config_projection_rejects_unknown_none(): + with pytest.raises(ValueError, match="enable_sleep_mod"): + omni_config_module._DiffusionConfigProjection.from_kwargs(enable_sleep_mod=None) + + +def test_runtime_diffusion_config_rejects_unknown_none(): + from vllm_omni.diffusion.data import OmniDiffusionConfig + + with pytest.raises(ValueError, match="enable_sleep_mod"): + OmniDiffusionConfig.from_kwargs(enable_sleep_mod=None) + + +def test_direct_diffusion_aliases_promote_none_and_reject_real_conflicts(monkeypatch): + from vllm_omni.diffusion import data as diffusion_data + from vllm_omni.diffusion.data import OmniDiffusionConfig + + monkeypatch.setattr(diffusion_data, "build_quant_config", lambda config: config) + with pytest.warns(FutureWarning) as warnings: + config = OmniDiffusionConfig.from_kwargs( + quantization={"method": "example"}, + quantization_config=None, + kv_cache_dtype="fp8", + diffusion_kv_cache_dtype=None, + enable_sleep_mode=False, + request_batch_max_wait_ms=0.0, + ) + + warning_messages = [str(warning.message) for warning in warnings] + assert any("quantization" in message for message in warning_messages) + assert any("kv_cache_dtype" in message for message in warning_messages) + assert config.quantization_config == {"method": "example"} + assert config.diffusion_kv_cache_dtype == "fp8" + assert config.enable_sleep_mode is False + assert config.request_batch_max_wait_ms == 0.0 + with pytest.raises(ValueError, match=r"quantization.*quantization_config"): + OmniDiffusionConfig.from_kwargs(quantization="fp8", quantization_config={"method": "example"}) + + +def test_structured_diffusion_stage_rejects_unknown_engine_extra(): + with pytest.raises(ValueError, match=r"stage 0.*enable_sleep_mod"): + _build_single_diffusion_config(engine_extras={"enable_sleep_mod": None}) + + with pytest.raises(ValueError, match=r"stage 0.*enable_sleep_mod"): + _build_single_diffusion_config(cli_overrides={"enable_sleep_mod": None}) + + with pytest.raises(ValueError, match=r"stage 0.*enable_sleep_mod"): + _build_single_diffusion_config(cli_overrides={"stage_0_enable_sleep_mod": None}) + + with pytest.raises(ValueError, match=r"no structured config owner: parallel_config\.unknown_field"): + _build_single_diffusion_config(engine_extras={"parallel_config": {"unknown_field": 1}}) + + +def test_structured_diffusion_stage_keeps_shared_globals_outside_diffusion(): + with warnings.catch_warnings(record=True) as caught: + stage = _build_single_diffusion_config( + cli_overrides={"kv_cache_dtype": "fp8", "seed": 7}, + ).stage_by_id(0) + + assert stage.diffusion_config.diffusion_kv_cache_dtype is None + assert not any(issubclass(warning.category, FutureWarning) for warning in caught) + + +@pytest.mark.parametrize( + ("field_name", "config_kwargs"), + [ + ("seed", {"cli_overrides": {"stage_0_seed": None}}), + ("kv_cache_dtype", {"engine_extras": {"kv_cache_dtype": "fp8"}}), + ], +) +def test_structured_diffusion_stage_rejects_explicit_shared_engine_field(field_name, config_kwargs): + with pytest.raises(ValueError, match=rf"stage 0.*{field_name}"): + _build_single_diffusion_config(**config_kwargs) diff --git a/tests/config/test_sleep_mode_deploy.py b/tests/config/test_sleep_mode_deploy.py new file mode 100644 index 000000000..13040ebd9 --- /dev/null +++ b/tests/config/test_sleep_mode_deploy.py @@ -0,0 +1,51 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Check the sleep/wake E2E topology before loading large BAGEL weights.""" + +import pytest + +from tests.entrypoints.test_omni_sleep_mode import _sleep_deploy_config +from vllm_omni.config.omni_config import ( + VllmOmniConfig, + VllmOmniDiffusionStageConfig, + extract_diffusion_stage_config_kwargs, +) +from vllm_omni.config.pipeline_registry import resolve_pipeline_config +from vllm_omni.config.stage_config import load_deploy_config +from vllm_omni.diffusion.data import OmniDiffusionConfig +from vllm_omni.engine.stage_init_utils import build_engine_args_dict_from_omni_stage_config + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +@pytest.mark.parametrize("diffusion_only,tp_size", [(False, 1), (False, 2), (True, 2)]) +def test_sleep_deploy_resolves_expected_stages(diffusion_only, tp_size): + deploy = load_deploy_config(_sleep_deploy_config(diffusion_only=diffusion_only, tp_size=tp_size)) + pipeline = resolve_pipeline_config(deploy.pipeline or "bagel") + config = VllmOmniConfig.from_pipeline_config(pipeline, user_deploy_config=deploy) + + assert len(config.stage_configs) == (1 if diffusion_only else 2) + for stage in config.stage_configs: + is_diffusion = diffusion_only or stage.stage_id == 1 + assert isinstance(stage, VllmOmniDiffusionStageConfig) == is_diffusion + # Keep BAGEL's batch budget: a single multimodal item needs 8625 tokens. + assert stage.scheduler_config.max_num_batched_tokens == 32768 + assert stage.model_config.enable_sleep_mode is True + assert stage.parallel_config.tensor_parallel_size == tp_size + assert stage.runtime_config.devices == ",".join(str(stage.stage_id * tp_size + rank) for rank in range(tp_size)) + if not is_diffusion: + assert stage.cache_config.gpu_memory_utilization == 0.8 + if is_diffusion: + engine_args = build_engine_args_dict_from_omni_stage_config(stage, "test-bagel") + terminal = OmniDiffusionConfig.from_kwargs( + **extract_diffusion_stage_config_kwargs( + engine_args, stage_id=stage.stage_id, include_engine_adapter_metadata=True + ) + ) + assert terminal.enable_sleep_mode is True + assert terminal.parallel_config.tensor_parallel_size == tp_size + + if not diffusion_only: + assert config.stage_by_id(1).stage_pipeline_config.input_sources == (0,) + assert deploy.stages[1].input_connectors["from_stage_0"] == "shared_memory_connector" diff --git a/tests/diffusion/quantization/test_quantization_fp8.py b/tests/diffusion/quantization/test_quantization_fp8.py index 5aa7faaa6..9f41a6a89 100644 --- a/tests/diffusion/quantization/test_quantization_fp8.py +++ b/tests/diffusion/quantization/test_quantization_fp8.py @@ -414,18 +414,16 @@ def test_quantization_key_maps_to_quantization_config(): @hardware_test(res={"cuda": ["L4", "B200"]}) -def test_quantization_config_key_takes_priority(): - """When both 'quantization' and 'quantization_config' are set, - 'quantization_config' takes priority.""" +def test_quantization_and_quantization_config_conflict(): + """Two explicit quantization sources are rejected instead of silently merged.""" from vllm_omni.diffusion.data import OmniDiffusionConfig - config = OmniDiffusionConfig.from_kwargs( - model="test", - quantization="fp8", - quantization_config={"method": "fp8", "activation_scheme": "static"}, - ) - assert config.quantization_config is not None - assert config.quantization_config.activation_scheme == "static" + with pytest.raises(ValueError, match=r"quantization.*quantization_config"): + OmniDiffusionConfig.from_kwargs( + model="test", + quantization="fp8", + quantization_config={"method": "fp8", "activation_scheme": "static"}, + ) @hardware_test(res={"cuda": ["L4", "B200"]}) diff --git a/tests/diffusion/test_diffusion_config_propagation.py b/tests/diffusion/test_diffusion_config_propagation.py index 6f37f823d..7421dfb79 100644 --- a/tests/diffusion/test_diffusion_config_propagation.py +++ b/tests/diffusion/test_diffusion_config_propagation.py @@ -11,6 +11,7 @@ import pytest import torch from vllm_omni.config.config_factory import StageConfigFactory +from vllm_omni.config.omni_config import extract_diffusion_stage_config_kwargs from vllm_omni.diffusion.data import ( DiffusionParallelConfig, OmniDiffusionConfig, @@ -40,7 +41,8 @@ def _roundtrip_diffusion_config(**kwargs) -> OmniDiffusionConfig: """ stages = StageConfigFactory.create_default_diffusion(kwargs) engine_args = dict(stages[0]["engine_args"]) - return OmniDiffusionConfig.from_kwargs(**engine_args) + diffusion_kwargs = extract_diffusion_stage_config_kwargs(engine_args, stage_id=0) + return OmniDiffusionConfig(**{name: value for name, value in diffusion_kwargs.items() if value is not None}) class TestParallelConfigPropagation: @@ -57,9 +59,7 @@ class TestParallelConfigPropagation: stages = StageConfigFactory.create_default_diffusion({"parallel_config": pc, "model": "x"}) assert stages[0]["runtime"]["devices"] == "0,1,2,3" - # Let __post_init__ reconstruct from dict (real code path) - ea = dict(stages[0]["engine_args"]) - od = OmniDiffusionConfig.from_kwargs(**ea) + od = _roundtrip_diffusion_config(parallel_config=pc, model="x") assert od.parallel_config.tensor_parallel_size == 4 assert od.parallel_config.world_size == 4 @@ -193,6 +193,17 @@ class TestCreateDefaultDiffusion: assert od.max_num_batched_tokens == 2048 assert od.max_model_len == 4096 + @pytest.mark.parametrize( + ("field_name", "value"), + [ + ("enable_sleep_mod", None), + ("enable_lora", True), + ], + ) + def test_unowned_raw_field_is_rejected(self, field_name, value): + with pytest.raises(ValueError, match=field_name): + StageConfigFactory.create_default_diffusion({"model": "x", field_name: value}) + def test_qwen_image_edit_plus_sets_generic_multimodal_limit(): od_config = OmniDiffusionConfig(model="Qwen/Qwen-Image-Edit-2511", model_class_name="QwenImageEditPlusPipeline") diff --git a/tests/distributed/omni_connectors/test_nixl_connector.py b/tests/distributed/omni_connectors/test_nixl_connector.py index 5b2161eee..ba333580d 100644 --- a/tests/distributed/omni_connectors/test_nixl_connector.py +++ b/tests/distributed/omni_connectors/test_nixl_connector.py @@ -317,6 +317,7 @@ def copying_native_agent(consumer, monkeypatch): @pytest.mark.parametrize("case", ["empty", "all_empty", "mixed", "scalar", "structured_empty", "structured_mixed"]) +@pytest.mark.usefixtures("reliable_claim_queries") def test_zero_byte_leaves_roundtrip_without_native_descriptors(producer, consumer, copying_native_agent, case): empty = torch.empty((2, 0, 3), dtype=torch.float64) other = torch.empty((0,), dtype=torch.int64) diff --git a/tests/e2e/features/helpers/verl_omni_server.py b/tests/e2e/features/helpers/verl_omni_server.py index a38a3a12e..9d9815a1e 100644 --- a/tests/e2e/features/helpers/verl_omni_server.py +++ b/tests/e2e/features/helpers/verl_omni_server.py @@ -10,11 +10,9 @@ Ray workers can unpickle ``vLLMOmniHttpServerLocal`` via from __future__ import annotations -import argparse import json import logging import os -from dataclasses import asdict from enum import Enum from typing import Any @@ -22,14 +20,13 @@ import torchvision.transforms as T from omegaconf import DictConfig, OmegaConf from pydantic import BaseModel, ConfigDict from vllm.entrypoints.openai.api_server import build_app -from vllm.utils.argparse_utils import FlexibleArgumentParser import vllm_omni.entrypoints.cli.serve -from vllm_omni.engine.arg_utils import OmniEngineArgs from vllm_omni.entrypoints.async_omni import AsyncOmni from vllm_omni.entrypoints.openai.api_server import omni_init_app_state from vllm_omni.inputs.data import OmniCustomPrompt, OmniDiffusionSamplingParams from vllm_omni.outputs import OmniRequestOutput +from vllm_omni.utils.tracking_parser import TrackingArgumentParser, TrackingNamespace logger = logging.getLogger(__name__) @@ -183,7 +180,6 @@ class vLLMOmniHttpServerLocal: "tensor_parallel_size": _cfg_get(self.config, "tensor_model_parallel_size", 1), "seed": self.replica_rank + int(_cfg_get(self.config, "seed", 0) or 0), "override_generation_config": json.dumps({}), - "scheduling_policy": _cfg_get(self.config, "scheduling_policy", "fcfs"), "compilation_config": compilation_config, **engine_kwargs, } @@ -191,7 +187,7 @@ class vLLMOmniHttpServerLocal: model_path = str(_cfg_get(self.model_config, "local_path", None) or _cfg_get(self.model_config, "path")) server_args_list = ["serve", model_path] + build_cli_args_from_config(args) - parser = FlexibleArgumentParser(description="vLLM-Omni CLI") + parser = TrackingArgumentParser(description="vLLM-Omni CLI") subparsers = parser.add_subparsers(required=False, dest="subparser") cmds: dict[str, Any] = {} for cmd in vllm_omni.entrypoints.cli.serve.cmd_init(): @@ -204,11 +200,10 @@ class vLLMOmniHttpServerLocal: await self.run_server(parsed) - async def run_server(self, args: argparse.Namespace) -> None: - engine_args = OmniEngineArgs.from_cli_args(args) - engine_args = asdict(engine_args) - - engine_args["enable_dummy_pipeline"] = True + async def run_server(self, args: TrackingNamespace) -> None: + # Match the production entrypoint: parser defaults are not user overrides. + engine_args = args.get_explicit_kwargs_dict() + engine_args["model"] = args.model engine_args["custom_pipeline_args"] = {"pipeline_class": CUSTOM_PIPELINE_CLASS} engine_client = AsyncOmni(**engine_args) diff --git a/tests/engine/test_stage_engine_args.py b/tests/engine/test_stage_engine_args.py index 6650c8e45..6d28280e0 100644 --- a/tests/engine/test_stage_engine_args.py +++ b/tests/engine/test_stage_engine_args.py @@ -29,6 +29,7 @@ from vllm_omni.config.omni_config import ( OmniStageSchedulerConfig, VllmOmniARStageConfig, VllmOmniConfig, + extract_diffusion_stage_config_kwargs, ) from vllm_omni.config.pipeline_registry import OMNI_PIPELINES, resolve_pipeline_config from vllm_omni.config.stage_config import ( @@ -567,7 +568,9 @@ def test_engine_args_consume_stage_diffusion_attention_shorthand(tmp_path): assert engine_args.get("diffusion_attention_backend") is None assert isinstance(engine_args["diffusion_attention_config"], AttentionConfig) assert engine_args["diffusion_attention_config"].default.backend == "TORCH_SDPA" - od_config = OmniDiffusionConfig.from_kwargs(**engine_args) + od_config = OmniDiffusionConfig.from_kwargs( + **extract_diffusion_stage_config_kwargs(engine_args, stage_id=2, include_engine_adapter_metadata=True) + ) assert od_config.diffusion_attention_config.default.backend == "TORCH_SDPA" @@ -600,7 +603,9 @@ def test_engine_args_apply_cli_attention_shorthand_over_yaml_config(tmp_path, ya attention_config = engine_args["diffusion_attention_config"] assert attention_config.default.backend == "TORCH_SDPA" assert attention_config.per_role["cross"].backend == "SAGE_ATTN" - od_config = OmniDiffusionConfig.from_kwargs(**engine_args) + od_config = OmniDiffusionConfig.from_kwargs( + **extract_diffusion_stage_config_kwargs(engine_args, stage_id=2, include_engine_adapter_metadata=True) + ) assert od_config.diffusion_attention_config.default.backend == "TORCH_SDPA" diff --git a/tests/entrypoints/test_async_omni_diffusion_config.py b/tests/entrypoints/test_async_omni_diffusion_config.py index 9b99af712..ed886d026 100644 --- a/tests/entrypoints/test_async_omni_diffusion_config.py +++ b/tests/entrypoints/test_async_omni_diffusion_config.py @@ -9,7 +9,7 @@ import torch from pydantic import ValidationError from vllm_omni.config.config_factory import StageConfigFactory -from vllm_omni.config.omni_config import VllmOmniDiffusionStageConfig +from vllm_omni.config.omni_config import VllmOmniDiffusionStageConfig, extract_diffusion_stage_config_kwargs from vllm_omni.config.resolver import OmniConfigResolution, resolve_omni_config from vllm_omni.diffusion.data import AttentionConfig, OmniDiffusionConfig from vllm_omni.engine import stage_init_utils @@ -21,7 +21,8 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu] def _terminal_config(stage_cfg: dict) -> OmniDiffusionConfig: - return OmniDiffusionConfig.from_kwargs(**stage_cfg["engine_args"]) + kwargs = extract_diffusion_stage_config_kwargs(stage_cfg["engine_args"], stage_id=stage_cfg["stage_id"]) + return OmniDiffusionConfig.from_kwargs(**kwargs) def test_default_stage_config_includes_cache_backend(): @@ -32,6 +33,9 @@ def test_default_stage_config_includes_cache_backend(): "cache_config": '{"Fn_compute_blocks": 2}', "vae_use_slicing": True, "ulysses_degree": 2, + "seed": 7, + "kv_cache_dtype": "fp8", + "diffusion_kv_cache_dtype": "fp8_e4m3", } )[0] @@ -42,6 +46,9 @@ def test_default_stage_config_includes_cache_backend(): assert engine_args["vae_use_slicing"] is True assert engine_args["parallel_config"]["ulysses_degree"] == 2 assert engine_args["model_stage"] == "diffusion" + assert "seed" not in engine_args + assert "kv_cache_dtype" not in engine_args + assert engine_args["diffusion_kv_cache_dtype"] == "fp8_e4m3" def test_default_stage_config_preserves_ulysses_a2a_permute() -> None: @@ -130,6 +137,67 @@ def test_default_stage_rejects_unknown_nested_parallel_config_key(): ) +def test_default_stage_routes_ar_profiler_away_before_diffusion_build(mocker): + stage_dict = StageConfigFactory.create_default_diffusion({"enable_ar_profiler": True})[0] + assert "enable_ar_profiler" not in stage_dict["engine_args"] + + +@pytest.mark.parametrize( + ("field_name", "value"), + [ + ("enable_sleep_mod", None), + ("enable_lora", True), + ("kv_cache_dtype", "fp8"), + ("seed", 7), + ], +) +def test_legacy_diffusion_stage_rejects_unowned_field(field_name, value): + from vllm_omni.config.config_factory import StageConfigFactory + from vllm_omni.config.yaml_util import create_config + from vllm_omni.engine.stage_init_utils import build_diffusion_config + + stage_dict = StageConfigFactory.create_default_diffusion({"model": "unused"})[0] + stage_dict["engine_args"][field_name] = value + stage_cfg = create_config(stage_dict) + metadata = SimpleNamespace(stage_id=0, cfg_kv_collect_func=None) + + with pytest.raises(ValueError, match=rf"stage 0.*{field_name}"): + build_diffusion_config("unused", stage_cfg, metadata) + + +@pytest.mark.parametrize( + ("field_name", "value"), + [ + ("enable_sleep_mod", None), + ("enable_lora", True), + ], +) +def test_default_diffusion_factory_rejects_unowned_field(field_name, value): + with pytest.raises(ValueError, match=rf"stage 0.*{field_name}"): + StageConfigFactory.create_default_diffusion({field_name: value}) + + +def test_legacy_default_stage_build_accepts_engine_adapter_metadata(monkeypatch): + from vllm_omni.config.config_factory import StageConfigFactory + from vllm_omni.config.yaml_util import create_config + from vllm_omni.engine import stage_init_utils + + stage_dict = StageConfigFactory.create_default_diffusion( + { + "model": "unused", + "api_key": "frontend-owned", + } + )[0] + assert "api_key" not in stage_dict["engine_args"] + stage_cfg = create_config(stage_dict) + metadata = SimpleNamespace(stage_id=0, cfg_kv_collect_func=None) + monkeypatch.setattr(stage_init_utils.current_omni_platform, "get_device_count", lambda: 1) + + config = stage_init_utils.build_diffusion_config("unused", stage_cfg, metadata) + + assert config.model == "unused" + + def test_default_cache_config_used_when_missing(): """Ensure default cache_config is synthesized when only backend is given.""" stage_cfg = StageConfigFactory.create_default_diffusion( @@ -671,10 +739,11 @@ def test_serve_cli_rejects_invalid_request_batch_max_wait_ms(bad_wait: str): ) -def test_serve_cli_accepts_additional_config(): +@pytest.mark.parametrize("subcommand_dest", ["command", "subparser"]) +def test_serve_cli_accepts_additional_config(subcommand_dest): """Ensure diffusion serve CLI exposes additional_config and forwards it to stage config.""" parser = TrackingArgumentParser() - subparsers = parser.add_subparsers(dest="command") + subparsers = parser.add_subparsers(dest=subcommand_dest) OmniServeCommand().subparser_init(subparsers) args = parser.parse_args( @@ -687,7 +756,7 @@ def test_serve_cli_accepts_additional_config(): ] ) - stage_cfg = StageConfigFactory.create_default_diffusion(vars(args))[0] + stage_cfg = StageConfigFactory.create_default_diffusion(args.get_explicit_kwargs_dict())[0] engine_args = stage_cfg["engine_args"] @@ -757,6 +826,24 @@ def test_default_stage_config_includes_quantization_config(): assert stage_cfg["engine_args"]["quantization_config"] == quantization_config +@pytest.mark.parametrize("typed", [False, True], ids=["legacy", "typed"]) +def test_default_diffusion_factory_preserves_engine_quantization(typed, monkeypatch): + monkeypatch.setattr(OmniDiffusionConfig, "_resolve_master_port", lambda _self: 29500) + monkeypatch.setattr(OmniDiffusionConfig, "enrich_config", lambda _self: None) + kwargs = {"quantization": "fp8"} + + if typed: + stage = StageConfigFactory.create_typed_default_diffusion("generic-diffusion", kwargs).stage_configs[0] + config = stage.diffusion_config + config.enrich_config() + else: + config = _terminal_config(StageConfigFactory.create_default_diffusion(kwargs)[0]) + + assert config.quantization_config is not None + assert config.quantization_config.get_name() == "fp8" + assert config.quantization_config_is_auto_detected is False + + @pytest.mark.parametrize("model_class_name", ["HeliosPipeline", "HunyuanVideo15Pipeline"]) def test_generic_diffusion_uses_canonical_video_output_type(model_class_name): config = StageConfigFactory.create_typed_default_diffusion( diff --git a/tests/entrypoints/test_omni_sleep_mode.py b/tests/entrypoints/test_omni_sleep_mode.py index 9567cbc96..c4900e9ca 100644 --- a/tests/entrypoints/test_omni_sleep_mode.py +++ b/tests/entrypoints/test_omni_sleep_mode.py @@ -18,6 +18,7 @@ import torch from vllm import SamplingParams from tests.helpers.mark import hardware_test +from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config from vllm_omni.entrypoints.async_omni import AsyncOmni from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.platforms import current_omni_platform @@ -98,16 +99,6 @@ MODEL = "ByteDance-Seed/BAGEL-7B-MoT" MODEL_DIFF = "riverclouds/qwen_image_random" -def get_dynamic_devices(stage_idx: int, num_stages: int, tp_size: int) -> str: - total_gpus = torch.accelerator.device_count() - gpus_per_stage = tp_size - start_idx = stage_idx * gpus_per_stage - if start_idx + gpus_per_stage > total_gpus: - start_idx = start_idx % total_gpus - device_ids = [str(start_idx + i) for i in range(gpus_per_stage)] - return ",".join(device_ids) - - async def _ensure_awake(engine: AsyncOmni, stage_ids: list[int]) -> None: """Best-effort full wake + resume so the next shared-engine test starts active.""" try: @@ -120,33 +111,28 @@ async def _ensure_awake(engine: AsyncOmni, stage_ids: list[int]) -> None: logger.warning("ensure_resume failed (stage_ids=%s): %s", stage_ids, e) -def _build_llm_stages() -> tuple[list[dict], list[dict]]: - common_args = { - "worker_type": "ar", - "enable_sleep_mode": True, - "dtype": "bfloat16", - "trust_remote_code": True, - "max_model_len": 2048, - "max_num_batched_tokens": 8192, - "enforce_eager": True, - } - stages = [ - { - "stage_id": 0, - "stage_type": "llm", - "runtime": {"process": True, "devices": "0", "max_batch_size": 1}, - "engine_args": {**common_args, "model_stage": "thinker", "gpu_memory_utilization": 0.1}, +def _sleep_deploy_config(*, diffusion_only: bool = False, tp_size: int = 1) -> str: + """Use registered BAGEL topology and override only deployment settings.""" + filename = "bagel_single_stage.yaml" if diffusion_only else "bagel.yaml" + stage_ids = (0,) if diffusion_only else (0, 1) + return modify_stage_config( + get_deploy_config_path(filename), + updates={ + "stages": { + stage_id: { + "devices": ",".join(str(stage_id * tp_size + rank) for rank in range(tp_size)), + "tensor_parallel_size": tp_size, + "enable_sleep_mode": True, + "enforce_eager": True, + "dtype": "bfloat16", + # Stages use disjoint GPUs. On an 80 GiB H100, 40% cannot + # fit TP=1 Thinker weights, activation peaks and KV cache. + "gpu_memory_utilization": 0.8 if not diffusion_only and stage_id == 0 else 0.4, + } + for stage_id in stage_ids + } }, - { - "stage_id": 1, - "stage_type": "llm", - "engine_input_source": [0], - "runtime": {"process": True, "devices": "1", "max_batch_size": 1, "connector_type": "queue"}, - "engine_args": {**common_args, "model_stage": "talker", "gpu_memory_utilization": 0.1}, - }, - ] - connectors = [{"src_stage_id": 0, "dst_stage_id": 1, "connector_type": "queue"}] - return stages, connectors + ) @pytest.fixture(scope="module", autouse=True) @@ -163,11 +149,10 @@ def _module_device_cleanup(): @pytest_asyncio.fixture(scope="class", loop_scope="class") async def llm_engine(): - """Shared 2-stage BAGEL LLM engine for sleep/wake protocol tests.""" + """Shared BAGEL Thinker + DiT engine for LLM sleep/wake protocol tests.""" if current_omni_platform.is_rocm(): clean_device_envs() - stages, connectors = _build_llm_stages() - engine = AsyncOmni(model=MODEL, stages=stages, connectors=connectors, init_timeout=600, enable_sleep_mode=True) + engine = AsyncOmni(model=MODEL, deploy_config=_sleep_deploy_config(), init_timeout=600, enable_sleep_mode=True) yield engine engine.shutdown() await asyncio.sleep(1.5) @@ -178,27 +163,12 @@ async def diffusion_engine(): """Shared BAGEL diffusion TP=2 engine for sleep/wake + generate checks.""" if current_omni_platform.is_rocm(): clean_device_envs() - stages = [ - { - "stage_id": 0, - "stage_type": "diffusion", - "runtime": {"process": True, "devices": "0,1", "max_batch_size": 1}, - "engine_args": { - "model_stage": "base", - "gpu_memory_utilization": 0.1, - "model_class_name": "BagelPipeline", - "enable_sleep_mode": True, - "enforce_eager": True, - "max_num_batched_tokens": 8192, - "parallel_config": { - "tensor_parallel_size": 2, - }, - }, - "final_output": True, - "final_output_type": "image", - } - ] - engine = AsyncOmni(model=MODEL, stages=stages, init_timeout=600, enable_sleep_mode=True) + engine = AsyncOmni( + model=MODEL, + deploy_config=_sleep_deploy_config(diffusion_only=True, tp_size=2), + init_timeout=600, + enable_sleep_mode=True, + ) yield engine engine.shutdown() await asyncio.sleep(1.5) @@ -208,7 +178,7 @@ class TestOmniLlmSleepMode: """LLM sleep/wake protocol tests sharing one class-scoped engine.""" @pytest.mark.asyncio(loop_scope="class") - @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=1) + @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) async def test_llm_sleep_ack(self, llm_engine: AsyncOmni): """LLM Thinker (GPU0) Signal and Physical Recycling Audit""" device_id = 0 @@ -333,11 +303,10 @@ class TestOmniDiffusionSleepMode: # --- integrity: baseline generate at original resolution/steps --- prompt = "A huge swimming pool, with many people swimming." sp = OmniDiffusionSamplingParams(num_inference_steps=4, height=512, width=512, seed=42) - llm_sp = SamplingParams() logger.info("Running Baseline Generation...") base_output = None - async for output in diffusion_engine.generate(prompt, request_id="base", sampling_params_list=[llm_sp, sp]): + async for output in diffusion_engine.generate(prompt, request_id="base", sampling_params=sp): base_output = output assert base_output is not None and len(base_output.images) > 0 logger.info("Baseline Generation successful.") @@ -384,7 +353,7 @@ class TestOmniDiffusionSleepMode: # --- integrity + lifecycle: post-wake generate --- logger.info("Running Post-Wakeup Generation...") post_output = None - async for output in diffusion_engine.generate(prompt, request_id="post", sampling_params_list=[llm_sp, sp]): + async for output in diffusion_engine.generate(prompt, request_id="post", sampling_params=sp): post_output = output assert post_output is not None assert len(base_output.images) == len(post_output.images) @@ -458,41 +427,22 @@ class TestOmniCoordinatedSleepMode: @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) @pytest.mark.asyncio async def test_multistage_llm_diffusion_sleep_wake(tp_size: int): - """Explicit 2-stage (llm + diffusion) + connectors; sleep/wake both stages.""" + """Registered BAGEL Thinker + DiT topology; sleep/wake both stages.""" if current_omni_platform.is_rocm(): clean_device_envs() num_gpus = torch.accelerator.device_count() if num_gpus < tp_size * 2: pytest.skip("Not enough GPUs") - stages = [] - for i in range(2): - devs = get_dynamic_devices(i, 2, tp_size) - stages.append( - { - "stage_id": i, - "stage_type": "llm" if i == 0 else "diffusion", - "runtime": {"process": True, "devices": devs}, - "engine_args": { - "model": MODEL, - "model_stage": "thinker" if i == 0 else "base", - "tensor_parallel_size": tp_size, - "gpu_memory_utilization": 0.4, - "dtype": "bfloat16", - "enable_sleep_mode": True, - "trust_remote_code": True, - }, - } - ) - - connectors = [{"src_stage_id": 0, "dst_stage_id": 1, "connector_type": "queue"}] - engine = AsyncOmni( - model=MODEL, stages=stages, connectors=connectors, enable_sleep_mode=True, stage_init_timeout=1200 + model=MODEL, + deploy_config=_sleep_deploy_config(tp_size=tp_size), + enable_sleep_mode=True, + stage_init_timeout=1200, ) try: sp = OmniDiffusionSamplingParams(num_inference_steps=2) - async for _ in engine.generate("warmup", sampling_params_list=[SamplingParams(), sp]): + async for _ in engine.generate("warmup", sampling_params_list=[engine.default_sampling_params_list[0], sp]): pass acks = await engine.sleep(stage_ids=[0, 1], level=1) @@ -500,7 +450,7 @@ async def test_multistage_llm_diffusion_sleep_wake(tp_size: int): await engine.wake_up(stage_ids=[0, 1]) await engine.resume_generation(stage_ids=[0, 1]) - async for _ in engine.generate("verify", sampling_params_list=[SamplingParams(), sp]): + async for _ in engine.generate("verify", sampling_params_list=[engine.default_sampling_params_list[0], sp]): pass finally: engine.shutdown() diff --git a/vllm_omni/config/config_factory.py b/vllm_omni/config/config_factory.py index 94decce9a..af0ea3e81 100644 --- a/vllm_omni/config/config_factory.py +++ b/vllm_omni/config/config_factory.py @@ -7,7 +7,7 @@ from __future__ import annotations import functools import json from collections.abc import Mapping -from dataclasses import asdict, dataclass +from dataclasses import asdict, dataclass, fields from pathlib import Path from typing import Any @@ -18,7 +18,10 @@ from vllm.transformers_utils.repo_utils import get_hf_file_to_dict from vllm.transformers_utils.runai_utils import ObjectStorageModel, is_runai_obj_uri from vllm_omni.config.endpoint_policy import EndpointRestriction -from vllm_omni.config.omni_config import VllmOmniConfig +from vllm_omni.config.omni_config import ( + VllmOmniConfig, + normalize_and_validate_diffusion_engine_ingress_kwargs, +) from vllm_omni.config.pipeline_registry import OMNI_PIPELINES, resolve_pipeline_config from vllm_omni.config.stage_config import ( _DEPLOY_DIR, @@ -645,6 +648,14 @@ class StageConfigFactory: kwargs: dict[str, Any], ) -> tuple[dict[str, Any], DiffusionParallelConfig, dict[str, Any], str]: """Normalize inputs shared by typed and compatibility diffusion builders.""" + # Preserve the default builder's explicit top-level override of model + # extras before the strict same-source alias validator runs. + if kwargs.get("auxiliary_text_encoder") is not None: + kwargs = dict(kwargs) + extras = dict(kwargs.get("extras") or {}) + extras.pop("auxiliary_text_encoder", None) + kwargs["extras"] = extras + kwargs = normalize_and_validate_diffusion_engine_ingress_kwargs(kwargs, stage_id=0) raw_sampling_params = kwargs.get("default_sampling_params") if isinstance(raw_sampling_params, str): try: @@ -659,7 +670,15 @@ class StageConfigFactory: parallel_config = DiffusionParallelConfig.from_stage_overrides(kwargs) if kwargs.get("num_gpus") is not None: parallel_config.resolve_data_parallel_size(int(kwargs["num_gpus"])) - engine_args = OmniDiffusionConfig.normalize_init_kwargs(kwargs) + # Ingress keeps the engine spelling; canonicalize it before filtering + # to terminal diffusion fields. Conflicts were validated at ingress. + quantization = kwargs.pop("quantization", None) + if quantization is not None: + kwargs["quantization_config"] = quantization + diffusion_config_fields = frozenset(config_field.name for config_field in fields(OmniDiffusionConfig)) + engine_args = OmniDiffusionConfig.normalize_init_kwargs( + {name: value for name, value in kwargs.items() if name in diffusion_config_fields} + ) extras = dict(engine_args.get("extras") or {}) for key, default in ( diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index f1e5de26a..5e51974a2 100644 --- a/vllm_omni/config/omni_config.py +++ b/vllm_omni/config/omni_config.py @@ -18,6 +18,7 @@ from inspect import Parameter, signature from pathlib import Path from typing import Any, Literal, TypeAlias, TypedDict, cast +import regex as re from pydantic import ConfigDict, Field, field_validator, model_validator from typing_extensions import Self from vllm.config import CacheConfig as VllmCacheConfig @@ -74,10 +75,14 @@ _NON_STAGE_ENGINE_CLI_FIELDS = frozenset( { "async_chunk", "disable_log_stats", + "command", + "headless", "model", + "model_tag", "omni", "output_modalities", "stage_id", + "subparser", "tokenizer", } ) @@ -95,6 +100,8 @@ _LEGACY_STAGE_METADATA_EXTRA_FIELDS = frozenset( _QuantizationConfigType: TypeAlias = QuantizationConfig | str | Mapping[str, Any] | None +_DIFFUSION_SHARED_ONLY_ENGINE_FIELDS = frozenset({"kv_cache_dtype", "seed"}) + class _QuantizationEngineOverrides(TypedDict, total=False): quantization_config: _QuantizationConfigType @@ -290,10 +297,6 @@ class _DiffusionEngineOverrides: _values: dict[str, Any] - @classmethod - def from_engine(cls, engine: Mapping[str, Any]) -> _DiffusionEngineOverrides: - return cls(_select_engine_overrides(engine, _DIFFUSION_STAGE_ENGINE_FIELDS)) - def to_kwargs(self) -> dict[str, Any]: return {name: _copy_value(value) for name, value in self._values.items()} @@ -359,12 +362,32 @@ def _stage_cli_overrides( *, execution_type: StageExecutionType | None = None, ) -> dict[str, Any]: + if execution_type == StageExecutionType.DIFFUSION: + stage_override_pattern = re.compile(r"^stage_\d+_") + stage_scoped = { + name: _copy_value(value) for name, value in cli_overrides.items() if stage_override_pattern.match(name) + } + global_inputs = {name: _copy_value(value) for name, value in cli_overrides.items() if name not in stage_scoped} + cli_overrides = { + **normalize_and_validate_diffusion_engine_ingress_kwargs( + global_inputs, + stage_id=stage_id, + ), + **stage_scoped, + } runtime_overrides = build_stage_runtime_overrides(stage_id, dict(cli_overrides)) global_stage_fields = _global_stage_cli_fields() owned_fields = None if execution_type is None else _STAGE_ENGINE_FIELDS_BY_EXECUTION_TYPE[execution_type] result: dict[str, Any] = {} for key, value in runtime_overrides.items(): - stage_specific = f"stage_{stage_id}_{key}" in cli_overrides + stage_key = f"stage_{stage_id}_{key}" + stage_specific = stage_key in cli_overrides + if ( + execution_type == StageExecutionType.DIFFUSION + and key in _DIFFUSION_SHARED_ONLY_ENGINE_FIELDS + and not stage_specific + ): + continue if stage_specific or (key in global_stage_fields and (owned_fields is None or key in owned_fields)): result[key] = _copy_value(value) @@ -372,6 +395,19 @@ def _stage_cli_overrides( # argument. Keep global and stage-scoped CLI values off AR/generation stages. if execution_type is not None and execution_type is not StageExecutionType.DIFFUSION: result.pop("step_execution", None) + if execution_type == StageExecutionType.DIFFUSION: + prefix = f"stage_{stage_id}_" + for key, value in cli_overrides.items(): + if key.startswith(prefix): + field_name = key.removeprefix(prefix) + if field_name in {"model", "model_arch", "stage_id"}: + raise ValueError(f"Diffusion stage {stage_id} cannot override identity field {field_name!r}.") + if field_name in _DIFFUSION_SHARED_ONLY_ENGINE_FIELDS: + raise ValueError( + f"Diffusion stage {stage_id} cannot override shared engine field {field_name!r}; " + "it has no diffusion stage-config consumer." + ) + result.setdefault(field_name, _copy_value(value)) return result @@ -387,6 +423,10 @@ def _validate_global_stage_cli_ownership( field for stage in pipeline.stages for field in _STAGE_ENGINE_FIELDS_BY_EXECUTION_TYPE[stage.execution_type] } unowned_fields = explicit_global_fields - owned_fields + if any(stage.execution_type == StageExecutionType.DIFFUSION for stage in pipeline.stages): + # Mixed engine ingress accepts these shared globals, but diffusion + # stages deliberately leave them outside their terminal config. + unowned_fields -= _DIFFUSION_SHARED_ONLY_ENGINE_FIELDS if unowned_fields: names = ", ".join(sorted(unowned_fields)) raise ValueError( @@ -736,7 +776,7 @@ class OmniStageDiffusionParallelConfig(OmniStageParallelConfig): self.world_size = other_parallel_world_size -@config(config=ConfigDict(arbitrary_types_allowed=True)) +@config(config=ConfigDict(arbitrary_types_allowed=True, extra="forbid")) class _DiffusionConfigProjection: """Diffusion-specific per-stage settings. @@ -852,15 +892,19 @@ class _DiffusionConfigProjection: @classmethod def from_kwargs(cls, **kwargs: Any) -> _DiffusionConfigProjection: - from vllm_omni.diffusion.data import normalize_omni_diffusion_kwargs + from vllm_omni.diffusion.data import ( + normalize_omni_diffusion_kwargs, + validate_omni_diffusion_kwargs, + ) from vllm_omni.diffusion.offloader.config import parse_diffusion_offload_config - config_kwargs = normalize_omni_diffusion_kwargs(kwargs) + valid_fields = frozenset(f.name for f in fields(cast(Any, cls))) + normalized = normalize_omni_diffusion_kwargs(kwargs) + validate_omni_diffusion_kwargs(normalized, valid_fields) # Validate before stage construction while retaining the raw mapping # needed by dataclass/config serialization across process boundaries. - parse_diffusion_offload_config(config_kwargs.get("diffusion_offload_config")) - valid_fields = {f.name for f in fields(cast(Any, cls))} - return cls(**{k: v for k, v in config_kwargs.items() if k in valid_fields}) + parse_diffusion_offload_config(normalized.get("diffusion_offload_config")) + return cls(**{name: value for name, value in normalized.items() if value is not None}) def __post_init__(self) -> None: # Keep diffusion imports lazy so importing vllm_omni.config does not @@ -1095,20 +1139,7 @@ _DIFFUSION_MOVED_SHARED_FIELDS = frozenset( _STAGE_DEPLOY_ENGINE_FIELDS: tuple[str, ...] = tuple(_STAGE_DEPLOY_FIELDS) -_DIFFUSION_BACKCOMPAT_ENGINE_FIELDS = frozenset( - { - "diffusion_attention_backend", - "fastvideo_vsa_topk", - "kv_cache_dtype", - "kv_cache_skip_layers", - "kv_cache_skip_steps", - "static_lora_scale", - } -) -_DIFFUSION_STAGE_ENGINE_FIELDS = (_DIFFUSION_CONFIG_FIELDS | _DIFFUSION_BACKCOMPAT_ENGINE_FIELDS) - { - "model", - "stage_id", -} +_DIFFUSION_STAGE_ENGINE_FIELDS = _DIFFUSION_CONFIG_FIELDS - {"model", "stage_id"} def _upstream_engine_field_map( @@ -1248,6 +1279,8 @@ def _validate_stage_engine_override_ownership( stage_id: int, execution_type: StageExecutionType, overrides: Mapping[str, Any], + *, + validate_top_level: bool = True, ) -> None: try: owner_fields = _STAGE_ENGINE_FIELDS_BY_EXECUTION_TYPE[execution_type] @@ -1256,7 +1289,7 @@ def _validate_stage_engine_override_ownership( except KeyError as exc: raise ValueError(f"Unsupported stage execution type: {execution_type!r}") from exc - unowned_fields = set(overrides) - owner_fields + unowned_fields = set(overrides) - owner_fields if validate_top_level else set() parallel_config = overrides.get("parallel_config") if isinstance(parallel_config, Mapping): # Nested values traditionally use upstream config names, while flat @@ -1273,6 +1306,144 @@ def _validate_stage_engine_override_ownership( ) +_DIFFUSION_STAGE_METADATA_FIELDS = frozenset( + { + "async_chunk", + "custom_process_next_stage_input_func", + "engine_output_type", + "has_sampling_extra_args", + "hf_config_name", + "model_arch", + "model_stage", + "retains_state_across_chunks", + "scheduler_cls", + "stage_connector_spec", + "worker_type", + } +) + +_DIFFUSION_ENGINE_ADAPTER_METADATA_FIELDS = frozenset( + { + # Serialized shared model settings are consumed outside the terminal + # diffusion config; accepting them here does not widen raw ingress. + "duplex_max_sessions", + "has_sampling_extra_args", + "inline_diffusion", + "requires_full_payload_input", + "sampling_extra_args_keys", + "session_mode", + } +) +_DIFFUSION_DEFAULT_FACTORY_FIELDS = frozenset( + {"default_llama_model_id", "default_sampling_params", "devices", "stage_0_devices"} +) + + +def _frontend_cli_fields() -> frozenset[str]: + """Return vLLM server fields that are consumed before stage startup.""" + from vllm.entrypoints.launchers.cli_args import FrontendArgs + + return frozenset(config_field.name for config_field in fields(FrontendArgs)) + + +def normalize_and_validate_diffusion_engine_ingress_kwargs( + kwargs: Mapping[str, Any], + *, + stage_id: int | str, +) -> dict[str, Any]: + """Normalize and validate raw diffusion input without inserting defaults.""" + from vllm_omni.diffusion.data import ( + OmniDiffusionConfig, + normalize_omni_diffusion_kwargs, + validate_omni_diffusion_kwargs, + ) + from vllm_omni.engine.arg_utils import orchestrator_field_names + + mixed_kwargs = {name: _copy_value(value) for name, value in kwargs.items()} + engine_owned = { + name: mixed_kwargs.pop(name) + for name in _DIFFUSION_SHARED_ONLY_ENGINE_FIELDS | {"quantization"} + if name in mixed_kwargs + } + normalized = normalize_omni_diffusion_kwargs(mixed_kwargs, apply_defaults=False) + if engine_owned.get("quantization") is not None and normalized.get("quantization_config") is not None: + raise ValueError("Diffusion config fields 'quantization' and 'quantization_config' cannot both be provided.") + normalized.update(engine_owned) + + diffusion_fields = frozenset(config_field.name for config_field in fields(OmniDiffusionConfig)) + stage_consumed_fields = ( + diffusion_fields + | _DIFFUSION_OWNED_STAGE_ENGINE_FIELDS + | frozenset(_STAGE_DEPLOY_ENGINE_FIELDS) + | frozenset(_PIPELINE_DEPLOY_CLI_FIELDS) + | _DIFFUSION_STAGE_METADATA_FIELDS + | _DIFFUSION_DEFAULT_FACTORY_FIELDS + ) - _DIFFUSION_SHARED_ONLY_ENGINE_FIELDS + externally_consumed_fields = ( + _DIFFUSION_SHARED_ONLY_ENGINE_FIELDS + | _NON_STAGE_ENGINE_CLI_FIELDS + | _frontend_cli_fields() + | orchestrator_field_names() + # Coordination fields also live on the typed orchestrator config, not + # all of them are present on the CLI-only OrchestratorArgs dataclass. + | frozenset(config_field.name for config_field in fields(VllmOmniOrchestratorConfig)) + ) + allowed_fields = stage_consumed_fields | externally_consumed_fields + validate_omni_diffusion_kwargs(normalized, allowed_fields, stage_id=stage_id) + return {name: value for name, value in normalized.items() if name in stage_consumed_fields} + + +def extract_diffusion_stage_config_kwargs( + kwargs: Mapping[str, Any], + *, + stage_id: int | str, + include_engine_adapter_metadata: bool = False, +) -> dict[str, Any]: + """Take the diffusion-owned payload from resolved mixed stage arguments.""" + from vllm_omni.diffusion.data import ( + OmniDiffusionConfig, + normalize_omni_diffusion_kwargs, + validate_omni_diffusion_kwargs, + ) + + shared_only_fields = sorted(_DIFFUSION_SHARED_ONLY_ENGINE_FIELDS.intersection(kwargs)) + if shared_only_fields: + field_names = ", ".join(repr(name) for name in shared_only_fields) + raise ValueError( + f"Diffusion stage {stage_id} cannot consume shared engine field(s) {field_names}; " + "use diffusion-owned fields or request sampling parameters instead." + ) + + diffusion_fields = frozenset(config_field.name for config_field in fields(OmniDiffusionConfig)) + fields_owned_elsewhere = ( + frozenset(_STAGE_DEPLOY_ENGINE_FIELDS) + | frozenset(_PIPELINE_DEPLOY_CLI_FIELDS) + | _DIFFUSION_OWNED_STAGE_ENGINE_FIELDS + | _DIFFUSION_STAGE_METADATA_FIELDS + ) + if include_engine_adapter_metadata: + fields_owned_elsewhere |= _DIFFUSION_ENGINE_ADAPTER_METADATA_FIELDS + + # ``quantization`` is canonical in the mixed engine namespace. Hand it to + # the compatibility adapter without treating it as a deprecated alias. + mixed_kwargs = {name: _copy_value(value) for name, value in kwargs.items()} + engine_quantization = mixed_kwargs.pop("quantization", None) + normalized = normalize_omni_diffusion_kwargs(mixed_kwargs) + if engine_quantization is not None: + if normalized.get("quantization_config") is not None: + raise ValueError( + "Diffusion config fields 'quantization' and 'quantization_config' cannot both be provided." + ) + normalized["quantization_config"] = engine_quantization + + validate_omni_diffusion_kwargs( + normalized, + diffusion_fields | fields_owned_elsewhere, + stage_id=stage_id, + ) + return {name: _copy_value(value) for name, value in normalized.items() if name in diffusion_fields} + + def _global_stage_cli_fields() -> frozenset[str]: # Lazy import avoids vllm_omni.config -> omni_config -> engine.arg_utils -> # vllm_omni.config during package-level config imports. @@ -1285,6 +1456,7 @@ def _global_stage_cli_fields() -> frozenset[str]: ) externally_consumed = ( _NON_STAGE_ENGINE_CLI_FIELDS + | _frontend_cli_fields() | frozenset(f.name for f in fields(cast(Any, VllmOmniOrchestratorConfig))) | (orchestrator_field_names() - _STAGE_ENGINE_FIELDS) ) @@ -1343,10 +1515,15 @@ def _stage_engine_values( engine[key] = _get_recursively_merged_dict(existing, dict(value)) else: engine[key] = _copy_value(value) + if topology.execution_type == StageExecutionType.DIFFUSION: + diffusion_kwargs = extract_diffusion_stage_config_kwargs(engine, stage_id=topology.stage_id) + else: + diffusion_kwargs = {} _validate_stage_engine_override_ownership( topology.stage_id, topology.execution_type, engine, + validate_top_level=topology.execution_type != StageExecutionType.DIFFUSION, ) if topology.execution_type in { StageExecutionType.LLM_AR, @@ -1383,7 +1560,7 @@ def _stage_engine_values( ), runtime=cast(_RuntimeEngineOverrides, _select_engine_overrides(engine, runtime_engine_fields)), parallel=cast(_ParallelEngineOverrides, _select_engine_overrides(engine, _PARALLEL_ENGINE_FIELDS)), - diffusion=_DiffusionEngineOverrides.from_engine(engine), + diffusion=_DiffusionEngineOverrides(_select_engine_overrides(diffusion_kwargs, _DIFFUSION_STAGE_ENGINE_FIELDS)), compilation_config=_copy_value(engine.get("compilation_config")), profiler_config=_copy_value(engine.get("profiler_config")), ) diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 537148933..5bc5a229b 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -6,6 +6,7 @@ import json import math import os import random +import warnings from collections.abc import Callable, Mapping from dataclasses import dataclass, field, fields from enum import Enum @@ -44,55 +45,63 @@ if TYPE_CHECKING: logger = init_logger(__name__) -def normalize_omni_diffusion_kwargs(raw_kwargs: Mapping[str, Any]) -> dict[str, Any]: - """Normalize legacy diffusion kwargs before config construction.""" +def _move_diffusion_alias( + normalized: dict[str, Any], + legacy_name: str, + canonical_name: str, +) -> None: + legacy_value = normalized.pop(legacy_name, None) + if legacy_value is None: + return + if normalized.get(canonical_name) is not None: + raise ValueError(f"Diffusion config fields {legacy_name!r} and {canonical_name!r} cannot both be provided.") + warnings.warn( + f"Diffusion config field {legacy_name!r} is deprecated; use {canonical_name!r}.", + FutureWarning, + stacklevel=3, + ) + normalized[canonical_name] = legacy_value + + +def normalize_omni_diffusion_kwargs( + raw_kwargs: Mapping[str, Any], + *, + apply_defaults: bool = True, +) -> dict[str, Any]: + """Normalize diffusion kwargs, deferring defaults until sources are merged.""" config_kwargs = dict(raw_kwargs) dtype = config_kwargs.get("dtype") if dtype is None: - config_kwargs["dtype"] = "auto" + if apply_defaults: + config_kwargs["dtype"] = "auto" elif isinstance(dtype, torch.dtype): config_kwargs["dtype"] = str(dtype).removeprefix("torch.") elif not isinstance(dtype, str): raise TypeError(f"Provided dtype must be a string or torch.dtype, got {type(dtype).__name__}") - # Backwards-compatibility: older callers may use a diffusion-specific - # "static_lora_scale" kwarg. Normalize it to the canonical "lora_scale". - if "static_lora_scale" in config_kwargs: - if "lora_scale" not in config_kwargs: - config_kwargs["lora_scale"] = config_kwargs["static_lora_scale"] - config_kwargs.pop("static_lora_scale", None) - - diffusion_quantization = config_kwargs.pop("diffusion_quantization_config", None) - if config_kwargs.get("quantization_config") is None and diffusion_quantization is not None: - config_kwargs["quantization_config"] = diffusion_quantization - - # Backwards-compatibility: map "quantization" to "quantization_config" - # so callers using the old field name still work. - if "quantization" in config_kwargs and config_kwargs.get("quantization_config", None) is None: - config_kwargs["quantization_config"] = config_kwargs.pop("quantization") - else: - config_kwargs.pop("quantization", None) - - # Renamed from kv_cache_* to avoid clashing with vLLM's --kv-cache-dtype. - if config_kwargs.get("diffusion_kv_cache_dtype") is None and "kv_cache_dtype" in config_kwargs: - config_kwargs["diffusion_kv_cache_dtype"] = config_kwargs.pop("kv_cache_dtype") - else: - config_kwargs.pop("kv_cache_dtype", None) - if config_kwargs.get("diffusion_kv_cache_skip_steps") is None and "kv_cache_skip_steps" in config_kwargs: - config_kwargs["diffusion_kv_cache_skip_steps"] = config_kwargs.pop("kv_cache_skip_steps") - else: - config_kwargs.pop("kv_cache_skip_steps", None) - if config_kwargs.get("diffusion_kv_cache_skip_layers") is None and "kv_cache_skip_layers" in config_kwargs: - config_kwargs["diffusion_kv_cache_skip_layers"] = config_kwargs.pop("kv_cache_skip_layers") - else: - config_kwargs.pop("kv_cache_skip_layers", None) + for legacy, canonical in ( + ("static_lora_scale", "lora_scale"), + ("quantization", "quantization_config"), + ("diffusion_quantization_config", "quantization_config"), + ("max_batch_size", "max_num_seqs"), + ("kv_cache_skip_steps", "diffusion_kv_cache_skip_steps"), + ("kv_cache_skip_layers", "diffusion_kv_cache_skip_layers"), + ): + _move_diffusion_alias(config_kwargs, legacy, canonical) + _move_diffusion_alias(config_kwargs, "kv_cache_dtype", "diffusion_kv_cache_dtype") # Handle "diffusion_attention_backend" shorthand: merge into # diffusion_attention_config before field filtering. diffusion_attn_backend = config_kwargs.pop("diffusion_attention_backend", None) fastvideo_vsa_topk = config_kwargs.pop("fastvideo_vsa_topk", None) if diffusion_attn_backend is not None or fastvideo_vsa_topk is not None: + if diffusion_attn_backend is not None: + warnings.warn( + "Diffusion config field 'diffusion_attention_backend' is deprecated; use 'diffusion_attention_config'.", + FutureWarning, + stacklevel=2, + ) existing = config_kwargs.get("diffusion_attention_config") config_kwargs["diffusion_attention_config"] = parse_attention_config( existing, @@ -100,12 +109,23 @@ def normalize_omni_diffusion_kwargs(raw_kwargs: Mapping[str, Any]) -> dict[str, fastvideo_vsa_topk=fastvideo_vsa_topk, ) + auxiliary_text_encoder = config_kwargs.pop("auxiliary_text_encoder", None) + if auxiliary_text_encoder is not None: + extras = dict(config_kwargs.get("extras") or {}) + if extras.get("auxiliary_text_encoder") is not None: + raise ValueError( + "Diffusion engine field 'auxiliary_text_encoder' cannot be provided both at the top level and in " + "'extras'." + ) + extras["auxiliary_text_encoder"] = auxiliary_text_encoder + config_kwargs["extras"] = extras + # Check environment variable as fallback for cache_backend. # Support both old DIFFUSION_CACHE_ADAPTER and new DIFFUSION_CACHE_BACKEND. - if "cache_backend" not in config_kwargs: + if "cache_backend" not in config_kwargs and apply_defaults: cache_backend = os.environ.get("DIFFUSION_CACHE_BACKEND") or os.environ.get("DIFFUSION_CACHE_ADAPTER") config_kwargs["cache_backend"] = cache_backend.lower() if cache_backend else "none" - elif config_kwargs["cache_backend"] is None: + elif "cache_backend" in config_kwargs and config_kwargs["cache_backend"] is None and apply_defaults: # Callers (e.g. example CLIs with `default=None`) pass an explicit # None for "no cache"; canonicalize it so every consumer sees the # declared `str` value instead of relying on per-model None handling. @@ -162,6 +182,20 @@ def validate_dlo_host_registration_options( return value +def validate_omni_diffusion_kwargs( + kwargs: Mapping[str, Any], + allowed_fields: set[str] | frozenset[str], + *, + stage_id: int | str | None = None, +) -> None: + """Reject every field without an owner.""" + unknown = set(kwargs) - allowed_fields + if unknown: + names = ", ".join(repr(name) for name in sorted(unknown)) + suffix = "" if stage_id is None else f" for stage {stage_id}" + raise ValueError(f"Unknown diffusion config field(s){suffix}: {names}") + + def parse_kv_cache_skip_selector( selector: str | list[int] | tuple[int, ...] | set[int] | None, ) -> set[int] | None: @@ -1631,8 +1665,9 @@ class OmniDiffusionConfig: @classmethod def normalize_init_kwargs(cls, raw_kwargs: Mapping[str, Any]) -> dict[str, Any]: + valid_fields = frozenset(f.name for f in fields(cls)) config_kwargs = normalize_omni_diffusion_kwargs(raw_kwargs) - valid_fields = {f.name for f in fields(cls)} + validate_omni_diffusion_kwargs(config_kwargs, valid_fields) # Remaining ``None`` values mean "unset" at the CLI/deploy boundary. # Drop them so non-optional dataclass defaults are not overwritten. # Fields where ``None`` has normalization semantics (for example dtype diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py index a10cf5861..f15e1238f 100644 --- a/vllm_omni/engine/stage_init_utils.py +++ b/vllm_omni/engine/stage_init_utils.py @@ -1965,13 +1965,20 @@ def build_diffusion_config( metadata: StageMetadata, ) -> Any: """Build diffusion config for a stage.""" + from vllm_omni.config.omni_config import extract_diffusion_stage_config_kwargs engine_args_dict = ( build_engine_args_dict_from_omni_stage_config(stage_cfg, model) if isinstance(stage_cfg, BaseVllmOmniStageConfig) else build_engine_args_dict(stage_cfg, model) ) - od_config = OmniDiffusionConfig.from_kwargs(**engine_args_dict) + stage_id = engine_args_dict.get("stage_id", metadata.stage_id) + diffusion_kwargs = extract_diffusion_stage_config_kwargs( + engine_args_dict, + stage_id=stage_id, + include_engine_adapter_metadata=True, + ) + od_config = OmniDiffusionConfig(**{name: value for name, value in diffusion_kwargs.items() if value is not None}) num_devices_per_stage = od_config.parallel_config.world_size device_control_env = current_omni_platform.device_control_env_var