[Bugfix][Diffusion] Preserve stage attention backend (#6645)

Signed-off-by: AndyZhou952 <jzhoubc@connect.ust.hk>

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

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
 

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

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,

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():

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(

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)