diff --git a/tests/config/test_config_factory.py b/tests/config/test_config_factory.py index e3c6c585f..f6288b219 100644 --- a/tests/config/test_config_factory.py +++ b/tests/config/test_config_factory.py @@ -2860,6 +2860,54 @@ stages: assert omega_config.engine_args.max_model_len == 8192 assert omega_config.engine_args.enforce_eager is True + def test_cli_attention_shorthand_replaces_yaml_structured_default(self): + stage = StageConfig( + stage_id=0, + model_stage="dit", + stage_type=StageType.DIFFUSION, + yaml_engine_args={ + "diffusion_attention_config": { + "default": {"backend": "FLASH_ATTN"}, + "per_role": {"cross": {"backend": "TORCH_SDPA"}}, + }, + }, + runtime_overrides={"diffusion_attention_backend": "SAGE_ATTN"}, + ) + + engine_args = stage.to_omegaconf().engine_args + + assert engine_args.diffusion_attention_backend == "SAGE_ATTN" + assert "default" not in engine_args.diffusion_attention_config + assert engine_args.diffusion_attention_config.per_role.cross.backend == "TORCH_SDPA" + + def test_cli_attention_shorthand_keeps_yaml_per_role_only_config(self): + stage = StageConfig( + stage_id=0, + model_stage="dit", + stage_type=StageType.DIFFUSION, + yaml_engine_args={"diffusion_attention_config": {"per_role": {"cross": {"backend": "TORCH_SDPA"}}}}, + runtime_overrides={"diffusion_attention_backend": "SAGE_ATTN"}, + ) + + engine_args = stage.to_omegaconf().engine_args + + assert engine_args.diffusion_attention_backend == "SAGE_ATTN" + assert engine_args.diffusion_attention_config == {"per_role": {"cross": {"backend": "TORCH_SDPA"}}} + + def test_cli_attention_config_replaces_yaml_shorthand(self): + stage = StageConfig( + stage_id=0, + model_stage="dit", + stage_type=StageType.DIFFUSION, + yaml_engine_args={"diffusion_attention_backend": "TORCH_SDPA"}, + runtime_overrides={"diffusion_attention_config": {"default": {"backend": "SAGE_ATTN"}}}, + ) + + engine_args = stage.to_omegaconf().engine_args + + assert "diffusion_attention_backend" not in engine_args + assert engine_args.diffusion_attention_config.default.backend == "SAGE_ATTN" + class TestSentinelDefaultPrecedence: """Caller-typed (non-None) values win over YAML; None values fall through diff --git a/tests/engine/test_async_omni_engine_stage_init.py b/tests/engine/test_async_omni_engine_stage_init.py index c1e56b7e1..c31009b76 100644 --- a/tests/engine/test_async_omni_engine_stage_init.py +++ b/tests/engine/test_async_omni_engine_stage_init.py @@ -11,7 +11,7 @@ import types import pytest -from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec +from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, normalize_omni_diffusion_kwargs from vllm_omni.engine import async_omni_engine as async_omni_engine_module from vllm_omni.engine.async_omni_engine import AsyncOmniEngine from vllm_omni.engine.stage_init_utils import ( @@ -1795,6 +1795,41 @@ def test_resolve_stage_configs_preserves_stage_diffusion_attention(monkeypatch): assert stage_configs[0].engine_args.diffusion_attention_config is existing_attention +def test_resolve_stage_configs_does_not_inject_over_stage_diffusion_attention_backend(monkeypatch): + import vllm_omni.engine.async_omni_engine as engine_mod + + engine = object.__new__(AsyncOmniEngine) + stage_cfg = types.SimpleNamespace( + stage_type="diffusion", + engine_args=types.SimpleNamespace( + diffusion_attention_backend="TORCH_SDPA", + diffusion_attention_config=None, + lora_path=None, + lora_scale=None, + enable_sleep_mode=None, + quantization_config=None, + ), + ) + + monkeypatch.setattr( + engine_mod, + "load_and_resolve_stage_configs", + lambda *args, **kwargs: ("dummy-config", [stage_cfg], None), + ) + + _config_path, stage_configs = engine._resolve_stage_configs( + model="dummy-model", + kwargs={"diffusion_attention_backend": "FLASH_ATTN"}, + trust_remote_code=False, + ) + + engine_args = stage_configs[0].engine_args + assert engine_args.diffusion_attention_backend == "TORCH_SDPA" + assert engine_args.diffusion_attention_config is None + normalized = normalize_omni_diffusion_kwargs(vars(engine_args)) + assert normalized["diffusion_attention_config"].default.backend == "TORCH_SDPA" + + def test_resolve_stage_configs_does_not_inject_diffusion_attention_into_llm_stage(monkeypatch): import vllm_omni.engine.async_omni_engine as engine_mod diff --git a/tests/engine/test_stage_engine_args.py b/tests/engine/test_stage_engine_args.py index 0a54e2e80..9cba6b19f 100644 --- a/tests/engine/test_stage_engine_args.py +++ b/tests/engine/test_stage_engine_args.py @@ -133,7 +133,7 @@ def _stable_engine_arg_environment(monkeypatch, tmp_path): for worker_type in ("ar", "generation"): worker_module = types.ModuleType(f"test.{worker_type}") - worker_module.Worker = _TestWorker + setattr(worker_module, "Worker", _TestWorker) monkeypatch.setitem(sys.modules, worker_module.__name__, worker_module) @@ -434,6 +434,63 @@ def test_typed_diffusion_engine_args_use_structured_diffusion_config(tmp_path): assert typed_args["diffusion_attention_config"].per_role["cross"].backend == "TORCH_SDPA" +def test_engine_args_consume_stage_diffusion_attention_shorthand(tmp_path): + """A stage-level ``diffusion_attention_backend`` shorthand is folded into the structured config.""" + pipeline, deploy, model = _engine_arg_inputs(tmp_path) + deploy.stages[2] = replace( + deploy.stages[2], + diffusion_attention_config=None, + diffusion_attention_backend="TORCH_SDPA", + ) + legacy_stages, omni_config = _legacy_and_typed_stages(pipeline, deploy, model) + + legacy_args = build_legacy_engine_args_dict(legacy_stages[2], model) + typed_args = build_engine_args_dict_from_omni_stage_config( + omni_config.stage_by_id(2), + model, + ) + + for engine_args in (legacy_args, typed_args): + 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) + assert od_config.diffusion_attention_config.default.backend == "TORCH_SDPA" + + +@pytest.mark.parametrize( + "yaml_attention_config", + [ + {"default": {"backend": "FLASH_ATTN"}, "per_role": {"cross": {"backend": "SAGE_ATTN"}}}, + {"per_role": {"cross": {"backend": "SAGE_ATTN"}}}, + ], + ids=["yaml-default", "yaml-per-role-only"], +) +def test_engine_args_apply_cli_attention_shorthand_over_yaml_config(tmp_path, yaml_attention_config): + pipeline, deploy, model = _engine_arg_inputs(tmp_path) + deploy.stages[2] = replace(deploy.stages[2], diffusion_attention_config=yaml_attention_config) + legacy_stages, omni_config = _legacy_and_typed_stages( + pipeline, + deploy, + model, + cli_overrides={"stage_2_diffusion_attention_backend": "TORCH_SDPA"}, + ) + + legacy_args = build_legacy_engine_args_dict(legacy_stages[2], model) + typed_args = build_engine_args_dict_from_omni_stage_config( + omni_config.stage_by_id(2), + model, + ) + + for engine_args in (legacy_args, typed_args): + assert engine_args.get("diffusion_attention_backend") is None + 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) + assert od_config.diffusion_attention_config.default.backend == "TORCH_SDPA" + + def test_typed_engine_args_preserve_explicit_backend_default_overrides(tmp_path): pipeline, deploy, model = _engine_arg_inputs(tmp_path) deploy.enable_prefix_caching = False diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index 58dc55fb6..2ef6bf276 100644 --- a/vllm_omni/config/omni_config.py +++ b/vllm_omni/config/omni_config.py @@ -47,6 +47,7 @@ from vllm_omni.config.stage_config import ( build_stage_runtime_overrides, load_deploy_config, normalize_pipeline_cli_overrides, + reconcile_diffusion_attention_overrides, ) from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode @@ -1207,6 +1208,9 @@ def _stage_engine_values( if topology.omni_kv_config: engine["omni_kv_config"] = _copy_value(topology.omni_kv_config) if stage_cli_overrides: + if topology.execution_type == StageExecutionType.DIFFUSION: + # Mirror StageConfig.to_omegaconf so both projections resolve alike. + reconcile_diffusion_attention_overrides(engine, stage_cli_overrides) engine.update(_copy_value(stage_cli_overrides)) _validate_stage_engine_override_ownership( topology.stage_id, diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index 2c1dc670b..6bfcd7e25 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -7,7 +7,7 @@ from __future__ import annotations import functools import re import warnings -from collections.abc import Callable +from collections.abc import Callable, Mapping from dataclasses import asdict, dataclass, field, fields from enum import Enum from pathlib import Path @@ -161,6 +161,28 @@ def _apply_diffusion_parallel_runtime_overrides( engine_args["parallel_config"] = parallel_config_dict +def reconcile_diffusion_attention_overrides( + engine_args: dict[str, Any], + runtime_overrides: Mapping[str, Any], +) -> None: + """Apply CLI precedence across the two diffusion attention representations. + + ``diffusion_attention_backend`` and ``diffusion_attention_config.default`` + express the same setting and are rejected downstream when both are set, so + a CLI value in one form replaces the YAML value in the other. + """ + if runtime_overrides.get("diffusion_attention_backend") is not None: + yaml_config = engine_args.get("diffusion_attention_config") + if isinstance(yaml_config, Mapping) and yaml_config.get("default") is not None: + remaining = {k: v for k, v in yaml_config.items() if k != "default"} + if remaining: + engine_args["diffusion_attention_config"] = remaining + else: + engine_args.pop("diffusion_attention_config") + if runtime_overrides.get("diffusion_attention_config") is not None: + engine_args.pop("diffusion_attention_backend", None) + + class StageType(str, Enum): """Type of processing stage in the Omni pipeline.""" @@ -1092,6 +1114,7 @@ class StageConfig: if StageType(self.stage_type) == StageType.DIFFUSION: _apply_diffusion_parallel_runtime_overrides(engine_args, runtime_overrides) + reconcile_diffusion_attention_overrides(engine_args, runtime_overrides) # CLI overrides take precedence over YAML defaults for key, value in runtime_overrides.items(): diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index 1efafbbab..9e894039f 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -1349,8 +1349,8 @@ class AsyncOmniEngine: or kwargs.get("fastvideo_vsa_topk") is not None ): has_stage_attention = ( - hasattr(cfg.engine_args, "diffusion_attention_config") - and cfg.engine_args.diffusion_attention_config is not None + getattr(cfg.engine_args, "diffusion_attention_config", None) is not None + or getattr(cfg.engine_args, "diffusion_attention_backend", None) is not None ) if not has_stage_attention: cfg.engine_args.diffusion_attention_config = parse_attention_config( diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py index 6fdcac81b..b34cdd5da 100644 --- a/vllm_omni/engine/stage_init_utils.py +++ b/vllm_omni/engine/stage_init_utils.py @@ -1253,9 +1253,19 @@ def _finalize_engine_args_dict( if is_diffusion: from vllm_omni.diffusion.data import parse_attention_config - if engine_args_dict.get("diffusion_attention_config") is not None: + # Fold the attention shorthand into the structured config so only one + # representation reaches OmniDiffusionConfig.from_kwargs. + attention_backend = engine_args_dict.pop("diffusion_attention_backend", None) + fastvideo_vsa_topk = engine_args_dict.pop("fastvideo_vsa_topk", None) + if ( + engine_args_dict.get("diffusion_attention_config") is not None + or attention_backend is not None + or fastvideo_vsa_topk is not None + ): engine_args_dict["diffusion_attention_config"] = parse_attention_config( - engine_args_dict["diffusion_attention_config"], + engine_args_dict.get("diffusion_attention_config"), + attention_backend=attention_backend, + fastvideo_vsa_topk=fastvideo_vsa_topk, ) else: resolve_worker_cls(engine_args_dict)