[Diffusion][Attention] Support per-role attention backend via CLI (#2681)

Signed-off-by: gcanlin <canlinguosdu@gmail.com>
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

tests/diffusion/attention/test_attention_config.py

@@ -0,0 +1,427 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Tests for per-role attention backend configuration (RFC: per-role-attention-backend).
+
+Tests cover:
+- AttentionSpec and AttentionConfig normalization
+- Role-aware backend resolution with category fallback
+- OmniDiffusionConfig attention shorthand handling
+- AttentionMetadata.extra field
+"""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+import vllm_omni.diffusion.attention.layer as layer_mod
+from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
+from vllm_omni.diffusion.attention.layer import Attention
+from vllm_omni.diffusion.config import (
+    get_current_diffusion_config,
+    get_current_diffusion_config_or_none,
+    set_current_diffusion_config,
+)
+from vllm_omni.diffusion.data import (
+    AttentionConfig,
+    AttentionSpec,
+    OmniDiffusionConfig,
+    build_attention_config,
+    parse_attention_config,
+)
+
+
+class TestAttentionSpec:
+    def test_construct_no_extra(self):
+        spec = AttentionSpec(backend="FLASH_ATTN")
+        assert spec.extra == {}
+
+    def test_mapping_extra_normalized(self):
+        spec = AttentionSpec(backend="SAGE_ATTN", extra={"quant": "int8"})
+        assert spec.backend == "SAGE_ATTN"
+        assert spec.extra == {"quant": "int8"}
+
+    def test_invalid_backend_type(self):
+        with pytest.raises(TypeError):
+            AttentionSpec(backend=123)  # type: ignore[arg-type]
+
+
+class TestAttentionConfig:
+    def test_empty_config(self):
+        config = AttentionConfig()
+        assert config.default is None
+        assert config.per_role == {}
+
+    def test_constructor_normalizes_mappings(self):
+        config = AttentionConfig(
+            default={"backend": "FLASH_ATTN"},
+            per_role={
+                "self": {"backend": "SPARSE_BLOCK", "extra": {"block_size": 128}},
+                "cross": "SAGE_ATTN",
+            },
+        )
+        assert config.default.backend == "FLASH_ATTN"
+        assert config.per_role["self"].backend == "SPARSE_BLOCK"
+        assert config.per_role["self"].extra == {"block_size": 128}
+        assert config.per_role["cross"].backend == "SAGE_ATTN"
+
+    def test_constructor_flattens_nested_per_role_tree(self):
+        config = AttentionConfig(
+            per_role={
+                "ltx2": {
+                    "audio_self": {"backend": "FLASH_ATTN"},
+                    "audio_to_video": {"backend": "SAGE_ATTN"},
+                }
+            }
+        )
+        assert config.per_role["ltx2.audio_self"].backend == "FLASH_ATTN"
+        assert config.per_role["ltx2.audio_to_video"].backend == "SAGE_ATTN"
+
+    def test_constructor_normalizes_auto_to_unset(self):
+        config = AttentionConfig(
+            default={"backend": "auto"},
+            per_role={
+                "self": "auto",
+                "cross": {"backend": "SAGE_ATTN"},
+            },
+        )
+        assert config.default is None
+        assert "self" not in config.per_role
+        assert config.per_role["cross"].backend == "SAGE_ATTN"
+
+    def test_resolve_exact_match(self):
+        config = AttentionConfig(
+            default=AttentionSpec(backend="FLASH_ATTN"),
+            per_role={
+                "self": AttentionSpec(backend="SPARSE_BLOCK"),
+                "cross": AttentionSpec(backend="SAGE_ATTN"),
+            },
+        )
+        spec, _ = config.resolve_with_source(role="self")
+        assert spec.backend == "SPARSE_BLOCK"
+
+        spec, _ = config.resolve_with_source(role="cross")
+        assert spec.backend == "SAGE_ATTN"
+
+    def test_resolve_with_source_reports_match_origin(self):
+        config = AttentionConfig(
+            default=AttentionSpec(backend="FLASH_ATTN"),
+            per_role={
+                "cross": AttentionSpec(backend="SAGE_ATTN"),
+                "ltx2.audio_to_video": AttentionSpec(backend="SPARSE_BLOCK"),
+            },
+        )
+
+        spec, source = config.resolve_with_source(role="ltx2.audio_to_video", role_category="cross")
+        assert spec is not None
+        assert spec.backend == "SPARSE_BLOCK"
+        assert source == "attention_config.per_role['ltx2.audio_to_video']"
+
+        spec, source = config.resolve_with_source(role="ltx2.video_to_audio", role_category="cross")
+        assert spec is not None
+        assert spec.backend == "SAGE_ATTN"
+        assert source == "attention_config.per_role['cross'] (role_category fallback)"
+
+        spec, source = config.resolve_with_source(role="self")
+        assert spec is not None
+        assert spec.backend == "FLASH_ATTN"
+        assert source == "attention_config.default"
+
+    def test_resolve_category_fallback(self):
+        config = AttentionConfig(
+            default=AttentionSpec(backend="FLASH_ATTN"),
+            per_role={
+                "cross": AttentionSpec(backend="SAGE_ATTN"),
+            },
+        )
+        # "ltx2.audio_to_video" falls back to category "cross"
+        spec, _ = config.resolve_with_source(role="ltx2.audio_to_video", role_category="cross")
+        assert spec.backend == "SAGE_ATTN"
+
+    def test_resolve_exact_overrides_category(self):
+        config = AttentionConfig(
+            per_role={
+                "cross": AttentionSpec(backend="SAGE_ATTN"),
+                "ltx2.audio_to_video": AttentionSpec(backend="FLASH_ATTN"),
+            },
+        )
+        # Exact match wins over category
+        spec, _ = config.resolve_with_source(role="ltx2.audio_to_video", role_category="cross")
+        assert spec.backend == "FLASH_ATTN"
+
+    def test_resolve_default_fallback(self):
+        config = AttentionConfig(
+            default=AttentionSpec(backend="FLASH_ATTN"),
+        )
+        spec, _ = config.resolve_with_source(role="self")
+        assert spec.backend == "FLASH_ATTN"
+
+        spec, _ = config.resolve_with_source(role="joint")
+        assert spec.backend == "FLASH_ATTN"
+
+    def test_resolve_returns_none_when_empty(self):
+        config = AttentionConfig()
+        spec, _ = config.resolve_with_source(role="self")
+        assert spec is None
+
+    def test_resolve_no_category_no_default(self):
+        config = AttentionConfig(
+            per_role={"self": AttentionSpec(backend="SPARSE_BLOCK")},
+        )
+        # Unknown role with no category and no default
+        spec, _ = config.resolve_with_source(role="joint")
+        assert spec is None
+
+    def test_full_ltx2_scenario(self):
+        """Test the LTX2 6-role stress test from the RFC."""
+        config = AttentionConfig(
+            default=AttentionSpec(backend="FLASH_ATTN"),
+            per_role={
+                "self": AttentionSpec(backend="SPARSE_BLOCK", extra={"block_size": 128}),
+                "cross": AttentionSpec(backend="SAGE_ATTN"),
+                "ltx2.audio_self": AttentionSpec(backend="FLASH_ATTN"),
+                "ltx2.audio_to_video": AttentionSpec(backend="FLASH_ATTN", extra={"causal_window": 64}),
+            },
+        )
+
+        # video self → exact match "self"
+        assert config.resolve_with_source("self")[0].backend == "SPARSE_BLOCK"
+
+        # audio self → exact match "ltx2.audio_self"
+        assert config.resolve_with_source("ltx2.audio_self", "self")[0].backend == "FLASH_ATTN"
+
+        # video-text cross → exact match "cross"
+        assert config.resolve_with_source("cross")[0].backend == "SAGE_ATTN"
+
+        # audio-text cross → category fallback to "cross"
+        assert config.resolve_with_source("ltx2.audio_text_cross", "cross")[0].backend == "SAGE_ATTN"
+
+        # audio-to-video → exact match
+        spec, _ = config.resolve_with_source("ltx2.audio_to_video", "cross")
+        assert spec.backend == "FLASH_ATTN"
+        assert spec.extra == {"causal_window": 64}
+
+        # video-to-audio → category fallback to "cross"
+        assert config.resolve_with_source("ltx2.video_to_audio", "cross")[0].backend == "SAGE_ATTN"
+
+
+class TestAttentionMetadataExtra:
+    def test_default_extra_is_empty(self):
+        meta = AttentionMetadata()
+        assert meta.extra == {}
+
+    def test_extra_passthrough(self):
+        block_mask = torch.ones(4, 4)
+        meta = AttentionMetadata(extra={"block_mask": block_mask, "kv_indices": [0, 1, 2]})
+        assert torch.equal(meta.extra["block_mask"], block_mask)
+        assert meta.extra["kv_indices"] == [0, 1, 2]
+
+    def test_extra_does_not_affect_existing_fields(self):
+        mask = torch.ones(2, 8)
+        meta = AttentionMetadata(attn_mask=mask, extra={"foo": "bar"})
+        assert meta.attn_mask is mask
+        assert meta.extra == {"foo": "bar"}
+
+
+class TestBuildAttentionConfig:
+    def test_env_sets_default_when_no_higher_priority_input(self, monkeypatch):
+        monkeypatch.setenv("DIFFUSION_ATTENTION_BACKEND", "TORCH_SDPA")
+
+        config = build_attention_config()
+
+        assert config.default is not None
+        assert config.default.backend == "TORCH_SDPA"
+
+    def test_attention_backend_overrides_env(self, monkeypatch):
+        monkeypatch.setenv("DIFFUSION_ATTENTION_BACKEND", "TORCH_SDPA")
+
+        config = parse_attention_config(attention_backend="SAGE_ATTN")
+
+        assert config.default is not None
+        assert config.default.backend == "SAGE_ATTN"
+
+    def test_parse_attention_config_does_not_read_env(self, monkeypatch):
+        monkeypatch.setenv("DIFFUSION_ATTENTION_BACKEND", "TORCH_SDPA")
+
+        config = parse_attention_config()
+
+        assert config.default is None
+
+    def test_attention_backend_auto_disables_env_fallback(self, monkeypatch):
+        monkeypatch.setenv("DIFFUSION_ATTENTION_BACKEND", "TORCH_SDPA")
+
+        config = parse_attention_config(attention_backend="auto")
+
+        assert config.default is None
+
+    def test_explicit_default_ignores_env(self, monkeypatch):
+        monkeypatch.setenv("DIFFUSION_ATTENTION_BACKEND", "self=FLASH_ATTN,cross=TORCH_SDPA")
+
+        config = build_attention_config(
+            AttentionConfig(default=AttentionSpec(backend="FLASH_ATTN")),
+        )
+
+        assert config.default is not None
+        assert config.default.backend == "FLASH_ATTN"
+
+    def test_env_auto_does_not_set_default(self, monkeypatch):
+        monkeypatch.setenv("DIFFUSION_ATTENTION_BACKEND", "auto")
+
+        config = build_attention_config()
+
+        assert config.default is None
+
+    def test_attention_backend_conflicts_with_explicit_default(self):
+        with pytest.raises(ValueError):
+            parse_attention_config(
+                AttentionConfig(default=AttentionSpec(backend="FLASH_ATTN")),
+                attention_backend="SAGE_ATTN",
+            )
+
+
+class TestOmniDiffusionConfigAttentionParsing:
+    """Test OmniDiffusionConfig attention shorthand and structured config."""
+
+    def test_diffusion_attention_backend_sets_default(self):
+        config = OmniDiffusionConfig.from_kwargs(diffusion_attention_backend="SAGE_ATTN")
+        assert isinstance(config.attention_config, AttentionConfig)
+        assert config.attention_config.default is not None
+        assert config.attention_config.default.backend == "SAGE_ATTN"
+
+    def test_diffusion_attention_backend_auto_means_platform_default(self):
+        config = OmniDiffusionConfig.from_kwargs(diffusion_attention_backend="auto")
+        assert isinstance(config.attention_config, AttentionConfig)
+        assert config.attention_config.default is None
+
+    def test_diffusion_attention_backend_and_default_are_mutually_exclusive(self):
+        with pytest.raises(ValueError):
+            OmniDiffusionConfig.from_kwargs(
+                diffusion_attention_backend="SAGE_ATTN",
+                diffusion_attention_config=AttentionConfig(default=AttentionSpec(backend="FLASH_ATTN")),
+            )
+
+    def test_dict_diffusion_attention_config(self):
+        config = OmniDiffusionConfig(
+            attention_config={
+                "default": {"backend": "FLASH_ATTN"},
+                "per_role": {"self": "SPARSE_BLOCK"},
+            }
+        )
+        assert config.attention_config.default.backend == "FLASH_ATTN"
+        assert config.attention_config.per_role["self"].backend == "SPARSE_BLOCK"
+
+    def test_legacy_attention_config_name_maps_in_from_kwargs(self):
+        config = OmniDiffusionConfig.from_kwargs(
+            attention_config={
+                "default": {"backend": "FLASH_ATTN"},
+            }
+        )
+        assert config.attention_config.default.backend == "FLASH_ATTN"
+
+    def test_no_attention_config_defaults_to_empty(self):
+        config = OmniDiffusionConfig()
+        assert isinstance(config.attention_config, AttentionConfig)
+        assert config.attention_config.default is None
+        assert config.attention_config.per_role == {}
+
+
+class TestCurrentDiffusionConfig:
+    def test_get_current_diffusion_config_or_none_defaults_to_none(self):
+        assert get_current_diffusion_config_or_none() is None
+
+    def test_get_current_diffusion_config_raises_when_unset(self):
+        with pytest.raises(AssertionError, match="Diffusion config is not set"):
+            get_current_diffusion_config()
+
+    def test_set_current_diffusion_config_restores_previous_value(self):
+        outer = SimpleNamespace(name="outer")
+        inner = SimpleNamespace(name="inner")
+
+        with set_current_diffusion_config(outer):
+            assert get_current_diffusion_config() is outer
+            with set_current_diffusion_config(inner):
+                assert get_current_diffusion_config() is inner
+            assert get_current_diffusion_config() is outer
+
+        assert get_current_diffusion_config_or_none() is None
+
+
+class TestAttentionInitUsesCurrentDiffusionConfig:
+    def test_attention_init_uses_current_diffusion_config_without_forward_context(self, monkeypatch):
+        class _FakeAttentionImpl:
+            def __init__(self, **kwargs):
+                self.kwargs = kwargs
+
+            def forward(self, query, key, value, attn_metadata=None):
+                return query
+
+        class _FakeBackend:
+            @staticmethod
+            def get_name() -> str:
+                return "FAKE_BACKEND"
+
+            @staticmethod
+            def get_impl_cls():
+                return _FakeAttentionImpl
+
+        captured = {}
+
+        def _fake_get_attn_backend_for_role(
+            role,
+            head_size,
+            attention_config=None,
+            role_category=None,
+        ):
+            captured["role"] = role
+            captured["head_size"] = head_size
+            captured["role_category"] = role_category
+            captured["attention_config"] = attention_config
+            return _FakeBackend, AttentionSpec(backend="TORCH_SDPA", extra={"block_size": 128})
+
+        class _FakeRingParallelAttention:
+            def __init__(self, sp_group, attn_backend_pref=None):
+                self.sp_group = sp_group
+                self.attn_backend_pref = attn_backend_pref
+
+        monkeypatch.setattr(layer_mod, "get_attn_backend_for_role", _fake_get_attn_backend_for_role)
+        monkeypatch.setattr(layer_mod.SDPABackend, "get_impl_cls", staticmethod(lambda: _FakeAttentionImpl))
+        monkeypatch.setattr(layer_mod, "build_parallel_attention_strategy", lambda **kwargs: object())
+        monkeypatch.setattr(layer_mod, "get_sp_group", lambda: SimpleNamespace(ring_group="ring-group"))
+        monkeypatch.setattr(layer_mod, "RingParallelAttention", _FakeRingParallelAttention)
+        monkeypatch.setattr(layer_mod, "is_forward_context_available", lambda: False)
+        monkeypatch.setattr(
+            layer_mod,
+            "get_forward_context",
+            lambda: (_ for _ in ()).throw(AssertionError("Attention init should not read ForwardContext")),
+        )
+
+        od_config = SimpleNamespace(
+            attention_config=AttentionConfig(
+                default=AttentionSpec(backend="FLASH_ATTN"),
+                per_role={"cross": AttentionSpec(backend="TORCH_SDPA", extra={"block_size": 128})},
+            ),
+            parallel_config=SimpleNamespace(ring_degree=2),
+        )
+
+        with set_current_diffusion_config(od_config):
+            attn = Attention(
+                num_heads=4,
+                head_size=64,
+                causal=False,
+                softmax_scale=1.0,
+                role="cross",
+                role_category="cross",
+                qkv_layout="BSND",
+            )
+
+        assert captured["role"] == "cross"
+        assert captured["role_category"] == "cross"
+        assert captured["head_size"] == 64
+        assert captured["attention_config"] is od_config.attention_config
+        assert attn.backend_pref == "TORCH_SDPA"
+        assert attn.attention.kwargs["backend_kwargs"] == {"block_size": 128}
+        assert attn.attention.kwargs["qkv_layout"] == "BSND"
+        assert attn.use_ring is True
+        assert attn.ring_runner is not None
+        assert attn.ring_runner.attn_backend_pref == "TORCH_SDPA"

tests/diffusion/attention/test_attention_sp.py

@@ -32,6 +32,7 @@ import pytest
 import torch
 
 from vllm_omni.diffusion.attention.layer import Attention
+from vllm_omni.diffusion.config import set_current_diffusion_config
 from vllm_omni.diffusion.data import (
     DiffusionParallelConfig,
     OmniDiffusionConfig,
@@ -403,8 +404,8 @@ def ulysses_attention_on_test_model(
         pipeline_parallel_size=1,
     )
 
-    # Set the config so Attention can access it
-    with set_forward_context(omni_diffusion_config=od_config):
+    # Set the config so Attention can access it during init and forward
+    with set_forward_context(omni_diffusion_config=od_config), set_current_diffusion_config(od_config):
         # Create model
         hidden_size = num_heads * head_size
 

tests/diffusion/attention/test_ulysses_uaa.py

@@ -12,6 +12,7 @@ import pytest
 import torch
 
 from vllm_omni.diffusion.attention.layer import Attention
+from vllm_omni.diffusion.config import set_current_diffusion_config
 from vllm_omni.diffusion.data import DiffusionParallelConfig, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.parallel_state import (
     destroy_distributed_env,
@@ -79,7 +80,7 @@ def _run_attention_case(
     )
     od_config = OmniDiffusionConfig(model="test", dtype=torch.float32, parallel_config=parallel_config)
 
-    with set_forward_context(omni_diffusion_config=od_config):
+    with set_forward_context(omni_diffusion_config=od_config), set_current_diffusion_config(od_config):
         attn = Attention(
             num_heads=num_heads,
             head_size=head_size,

tests/diffusion/cache/test_teacache_extractors.py

@@ -303,7 +303,10 @@ class TestFluxExtractor(BaseExtractorTest):
         """Use the SDPA backend so FLUX can be instantiated in CPU tests."""
         from vllm_omni.diffusion.attention.backends.sdpa import SDPABackend
 
-        with patch("vllm_omni.diffusion.attention.layer.get_attn_backend", return_value=SDPABackend):
+        with patch(
+            "vllm_omni.diffusion.attention.layer.get_attn_backend_for_role",
+            return_value=(SDPABackend, None),
+        ):
             yield
 
     def get_extractor(self):

tests/diffusion/model_loader/test_diffusers_loader.py

@@ -1,12 +1,16 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+from types import SimpleNamespace
+
 import pytest
 import torch
 import torch.nn as nn
 
+from vllm_omni.diffusion.config import get_current_diffusion_config, get_current_diffusion_config_or_none
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
 from vllm_omni.diffusion.model_loader.gguf_adapters import get_gguf_adapter
+from vllm_omni.diffusion.registry import initialize_model
 
 pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
 
@@ -93,3 +97,74 @@ def test_qwen_model_class_selects_qwen_gguf_adapter():
     adapter = get_gguf_adapter("dummy.gguf", object(), source, od_config)
 
     assert adapter.__class__.__name__ == "QwenImageGGUFAdapter"
+
+
+class _ConfigAwareModel(nn.Module):
+    def __init__(self, *, od_config):
+        super().__init__()
+        self.captured_config = get_current_diffusion_config()
+        self.seen_config_during_init = get_current_diffusion_config_or_none()
+        self.od_config = od_config
+
+
+def test_initialize_model_sets_current_diffusion_config_during_model_construction(monkeypatch):
+    import vllm_omni.diffusion.registry as registry_mod
+
+    od_config = SimpleNamespace(
+        model_class_name="DummyPipeline",
+        parallel_config=SimpleNamespace(vae_patch_parallel_size=1, sequence_parallel_size=1),
+        vae_use_slicing=False,
+        vae_use_tiling=False,
+    )
+
+    monkeypatch.setattr(
+        registry_mod.DiffusionModelRegistry,
+        "_try_load_model_cls",
+        staticmethod(lambda _name: _ConfigAwareModel),
+    )
+    monkeypatch.setattr(registry_mod, "_apply_sequence_parallel_if_enabled", lambda *_args, **_kwargs: None)
+
+    model = initialize_model(od_config)
+
+    assert model.captured_config is od_config
+    assert model.seen_config_during_init is od_config
+    assert get_current_diffusion_config_or_none() is None
+
+
+def test_load_model_custom_pipeline_sets_current_diffusion_config(monkeypatch):
+    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
+
+    class _DeviceContext:
+        def __init__(self, device_type: str):
+            self.type = device_type
+
+        def __enter__(self):
+            return self
+
+        def __exit__(self, exc_type, exc, tb):
+            return False
+
+    od_config = SimpleNamespace(
+        dtype=torch.float32,
+        parallel_config=SimpleNamespace(use_hsdp=False),
+        quantization_config=None,
+    )
+
+    loader = object.__new__(DiffusersPipelineLoader)
+    loader.load_weights = lambda model: None  # type: ignore[assignment]
+    loader._process_weights_after_loading = lambda model, target_device: None  # type: ignore[assignment]
+    loader._is_gguf_quantization = lambda _od_config: False  # type: ignore[assignment]
+
+    monkeypatch.setattr(loader_mod, "resolve_obj_by_qualname", lambda _name: _ConfigAwareModel)
+    monkeypatch.setattr(loader_mod.torch, "device", lambda _name: _DeviceContext("cpu"))
+
+    model = loader.load_model(
+        od_config,
+        load_device="cpu",
+        load_format="custom_pipeline",
+        custom_pipeline_name="tests.dummy.ConfigAwarePipeline",
+    )
+
+    assert model.captured_config is od_config
+    assert model.seen_config_during_init is od_config
+    assert get_current_diffusion_config_or_none() is None

tests/engine/test_async_omni_engine_stage_init.py

@@ -6,6 +6,7 @@ import types
 
 import pytest
 
+from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec
 from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
 from vllm_omni.engine.stage_init_utils import (
     LogicalStageInitPlan,
@@ -725,3 +726,158 @@ def test_inject_kv_stage_info_infers_receiver_tp_topology():
     assert stage1.engine_args["omni_kv_config"]["stage_id"] == 1
     assert stage1.engine_args["omni_kv_config"]["engine_input_source"] == [0]
     assert stage1.engine_args["omni_kv_config"]["rank_mapping"] == {"from_tp": 4, "to_tp": 2}
+
+
+def test_resolve_stage_configs_injects_global_diffusion_attention_when_missing(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_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]),
+    )
+
+    _config_path, stage_configs = engine._resolve_stage_configs(
+        model="dummy-model",
+        kwargs={"diffusion_attention_backend": "FLASH_ATTN"},
+    )
+
+    diffusion_attention_config = stage_configs[0].engine_args.diffusion_attention_config
+    assert isinstance(diffusion_attention_config, AttentionConfig)
+    assert diffusion_attention_config.default is not None
+    assert diffusion_attention_config.default.backend == "FLASH_ATTN"
+
+
+def test_resolve_stage_configs_preserves_stage_diffusion_attention(monkeypatch):
+    import vllm_omni.engine.async_omni_engine as engine_mod
+
+    engine = object.__new__(AsyncOmniEngine)
+    existing_attention = AttentionConfig(default=AttentionSpec(backend="TORCH_SDPA"))
+    stage_cfg = types.SimpleNamespace(
+        stage_type="diffusion",
+        engine_args=types.SimpleNamespace(
+            diffusion_attention_config=existing_attention,
+            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]),
+    )
+
+    _config_path, stage_configs = engine._resolve_stage_configs(
+        model="dummy-model",
+        kwargs={"diffusion_attention_backend": "FLASH_ATTN"},
+    )
+
+    assert stage_configs[0].engine_args.diffusion_attention_config is existing_attention
+
+
+def test_resolve_stage_configs_does_not_inject_diffusion_attention_into_llm_stage(monkeypatch):
+    import vllm_omni.engine.async_omni_engine as engine_mod
+
+    engine = object.__new__(AsyncOmniEngine)
+    stage_cfg = types.SimpleNamespace(
+        stage_type="llm",
+        engine_args=types.SimpleNamespace(
+            attention_config={"backend": "FLASH_ATTN"},
+            enable_sleep_mode=None,
+        ),
+    )
+
+    monkeypatch.setattr(
+        engine_mod,
+        "load_and_resolve_stage_configs",
+        lambda *args, **kwargs: ("dummy-config", [stage_cfg]),
+    )
+
+    _config_path, stage_configs = engine._resolve_stage_configs(
+        model="dummy-model",
+        kwargs={"diffusion_attention_backend": "TORCH_SDPA"},
+    )
+
+    assert stage_configs[0].engine_args.attention_config == {"backend": "FLASH_ATTN"}
+    assert not hasattr(stage_configs[0].engine_args, "diffusion_attention_config")
+
+
+def test_extract_stage_metadata_rocm_does_not_inject_diffusion_attention(monkeypatch):
+    """ROCm default attention logic only applies to LLM stages, not diffusion."""
+    from vllm_omni.engine.stage_init_utils import extract_stage_metadata
+
+    monkeypatch.setattr("vllm_omni.engine.stage_init_utils.current_omni_platform.is_rocm", lambda: True)
+
+    stage_cfg = types.SimpleNamespace(
+        stage_id=0,
+        stage_type="diffusion",
+        engine_args={},
+        runtime={},
+        engine_input_source=[],
+        final_output=False,
+        final_output_type=None,
+    )
+
+    metadata = extract_stage_metadata(stage_cfg)
+
+    assert metadata.stage_type == "diffusion"
+    assert "diffusion_attention_config" not in stage_cfg.engine_args
+
+
+def test_build_engine_args_dict_normalizes_diffusion_attention_config():
+    from vllm_omni.engine.stage_init_utils import build_engine_args_dict
+
+    stage_cfg = types.SimpleNamespace(
+        stage_id=0,
+        stage_type="diffusion",
+        engine_args={
+            "diffusion_attention_config": {
+                "default": {"backend": "FLASH_ATTN"},
+                "per_role": {"cross": {"backend": "TORCH_SDPA"}},
+            }
+        },
+        runtime={},
+    )
+
+    engine_args_dict = build_engine_args_dict(stage_cfg, model="dummy-model")
+
+    diffusion_attention_config = engine_args_dict["diffusion_attention_config"]
+    assert isinstance(diffusion_attention_config, AttentionConfig)
+    assert diffusion_attention_config.default is not None
+    assert diffusion_attention_config.default.backend == "FLASH_ATTN"
+    assert diffusion_attention_config.per_role["cross"].backend == "TORCH_SDPA"
+
+
+def test_build_engine_args_dict_uses_diffusion_attention_config_key():
+    from vllm_omni.engine.stage_init_utils import build_engine_args_dict
+
+    stage_cfg = types.SimpleNamespace(
+        stage_id=0,
+        stage_type="diffusion",
+        engine_args={
+            "diffusion_attention_config": {
+                "default": {"backend": "FLASH_ATTN"},
+            }
+        },
+        runtime={},
+    )
+
+    engine_args_dict = build_engine_args_dict(stage_cfg, model="dummy-model")
+
+    assert "attention_config" not in engine_args_dict
+    assert engine_args_dict["diffusion_attention_config"].default.backend == "FLASH_ATTN"

tests/entrypoints/test_async_omni_diffusion_config.py

@@ -5,6 +5,7 @@ import pytest
 from vllm.utils.argparse_utils import FlexibleArgumentParser
 
 from vllm_omni.config.stage_config import deploy_override_field_names
+from vllm_omni.diffusion.data import AttentionConfig
 from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
 from vllm_omni.entrypoints.cli.serve import OmniServeCommand, _create_default_diffusion_stage_cfg
 
@@ -131,6 +132,51 @@ def test_default_stage_config_includes_default_sampling_params():
     }
 
 
+def test_default_stage_config_includes_diffusion_attention_backend():
+    """Ensure diffusion attention shorthand lands in engine_args.diffusion_attention_config."""
+    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
+        {
+            "diffusion_attention_backend": "FLASH_ATTN",
+        }
+    )[0]
+
+    diffusion_attention_config = stage_cfg["engine_args"]["diffusion_attention_config"]
+    assert isinstance(diffusion_attention_config, AttentionConfig)
+    assert diffusion_attention_config.default is not None
+    assert diffusion_attention_config.default.backend == "FLASH_ATTN"
+
+
+def test_default_stage_config_includes_diffusion_attention_config():
+    """Ensure structured diffusion attention config survives default stage creation."""
+    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
+        {
+            "diffusion_attention_config": {
+                "default": {"backend": "FLASH_ATTN"},
+                "per_role": {"cross": {"backend": "TORCH_SDPA"}},
+            },
+        }
+    )[0]
+
+    diffusion_attention_config = stage_cfg["engine_args"]["diffusion_attention_config"]
+    assert isinstance(diffusion_attention_config, AttentionConfig)
+    assert diffusion_attention_config.default is not None
+    assert diffusion_attention_config.default.backend == "FLASH_ATTN"
+    assert diffusion_attention_config.per_role["cross"].backend == "TORCH_SDPA"
+
+
+def test_default_stage_config_rejects_conflicting_diffusion_attention_inputs():
+    """Ensure shorthand and default.backend stay mutually exclusive."""
+    with pytest.raises(ValueError, match="mutually exclusive"):
+        AsyncOmniEngine._create_default_diffusion_stage_cfg(
+            {
+                "diffusion_attention_backend": "FLASH_ATTN",
+                "diffusion_attention_config": {
+                    "default": {"backend": "TORCH_SDPA"},
+                },
+            }
+        )
+
+
 def test_default_stage_config_engine_args():
     """Ensure default diffusion-stage builder sets and propagates engine_args."""
     stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
@@ -216,3 +262,28 @@ def test_serve_cli_accepts_diffusion_pipeline_profiler_flag():
 
     assert args.enable_diffusion_pipeline_profiler is True
     assert stage_cfg["engine_args"]["enable_diffusion_pipeline_profiler"] is True
+
+
+def test_serve_cli_accepts_diffusion_attention_backend():
+    """Ensure diffusion serve CLI exposes the shorthand backend flag."""
+    parser = FlexibleArgumentParser()
+    subparsers = parser.add_subparsers(dest="command")
+    OmniServeCommand().subparser_init(subparsers)
+
+    args = parser.parse_args(
+        [
+            "serve",
+            "Qwen/Qwen-Image",
+            "--omni",
+            "--diffusion-attention-backend",
+            "FLASH_ATTN",
+        ]
+    )
+
+    stage_cfg = _create_default_diffusion_stage_cfg(args)[0]
+    diffusion_attention_config = stage_cfg["engine_args"]["diffusion_attention_config"]
+
+    assert args.diffusion_attention_backend == "FLASH_ATTN"
+    assert isinstance(diffusion_attention_config, AttentionConfig)
+    assert diffusion_attention_config.default is not None
+    assert diffusion_attention_config.default.backend == "FLASH_ATTN"

vllm_omni/diffusion/attention/backends/abstract.py

@@ -2,8 +2,8 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 from abc import ABC, abstractmethod
-from dataclasses import dataclass
-from typing import Generic, TypeVar
+from dataclasses import dataclass, field
+from typing import Any, Generic, TypeVar
 
 import torch
 
@@ -64,6 +64,9 @@ class AttentionMetadata:
     # a replicated tensor among processes appended to the front or rear of value, depends the joint_strategy
     joint_strategy: str = "front"
     # the strategy to joint the query, key, and value, can be "front" or "rear"
+    extra: dict[str, Any] = field(default_factory=dict)
+    # Opaque backend-specific per-forward parameters (e.g. block masks, KV indices).
+    # Backends MUST silently ignore unknown keys.
 
 
 T = TypeVar("T", bound=AttentionMetadata)
@@ -79,6 +82,8 @@ class AttentionImpl(ABC, Generic[T]):
         causal: bool = False,
         num_kv_heads: int | None = None,
         prefix: str = "",
+        qkv_layout: str | None = None,
+        backend_kwargs: dict[str, Any] | None = None,
         **extra_impl_args,
     ) -> None:
         raise NotImplementedError

vllm_omni/diffusion/attention/backends/flash_attn.py

@@ -43,11 +43,16 @@ class FlashAttentionImpl(AttentionImpl):
         causal: bool = False,
         num_kv_heads: int | None = None,
         prefix: str = "",
+        qkv_layout: str | None = None,
+        backend_kwargs: dict | None = None,
         **extra_impl_args,
     ) -> None:
         self.num_heads = num_heads
         self.causal = causal
         self.softmax_scale = softmax_scale
+        self.qkv_layout = qkv_layout
+        if backend_kwargs:
+            logger.warning("FlashAttentionImpl ignoring backend_kwargs: %s", list(backend_kwargs.keys()))
 
     @staticmethod
     def _unwrap_flash_output(out: torch.Tensor | tuple[torch.Tensor, ...]) -> torch.Tensor:
@@ -230,6 +235,7 @@ class FlashAttentionImpl(AttentionImpl):
             )
 
         attention_mask = attn_metadata.attn_mask if attn_metadata else None
+        layout = self.qkv_layout or "BNSD"
         output = attention_forward(
             query,
             key,
@@ -237,6 +243,6 @@ class FlashAttentionImpl(AttentionImpl):
             attn_mask=attention_mask,
             opt_mode="manual",
             op_type="fused_attn_score",
-            layout="BNSD",
+            layout=layout,
         )
         return output

vllm_omni/diffusion/attention/backends/sage_attn.py

@@ -49,10 +49,13 @@ class SageAttentionImpl(AttentionImpl):
         causal: bool = False,
         num_kv_heads: int | None = None,
         prefix: str = "",
+        backend_kwargs: dict | None = None,
         **extra_impl_args,
     ) -> None:
         self.causal = causal
         self.softmax_scale = softmax_scale
+        if backend_kwargs:
+            logger.warning("SageAttentionImpl ignoring backend_kwargs: %s", list(backend_kwargs.keys()))
 
     def forward_cuda(
         self,

vllm_omni/diffusion/attention/backends/sdpa.py

@@ -84,10 +84,13 @@ class SDPAImpl(AttentionImpl):
         causal: bool = False,
         num_kv_heads: int | None = None,
         prefix: str = "",
+        backend_kwargs: dict | None = None,
         **extra_impl_args,
     ) -> None:
         self.causal = causal
         self.softmax_scale = softmax_scale
+        if backend_kwargs:
+            logger.warning("SDPAImpl ignoring backend_kwargs: %s", list(backend_kwargs.keys()))
 
     def _forward_impl(
         self,

vllm_omni/diffusion/attention/layer.py

@@ -16,7 +16,8 @@ from vllm_omni.diffusion.attention.backends.sdpa import SDPABackend
 from vllm_omni.diffusion.attention.parallel import build_parallel_attention_strategy
 from vllm_omni.diffusion.attention.parallel.base import NoParallelAttention
 from vllm_omni.diffusion.attention.parallel.ring import RingParallelAttention
-from vllm_omni.diffusion.attention.selector import get_attn_backend
+from vllm_omni.diffusion.attention.selector import get_attn_backend_for_role
+from vllm_omni.diffusion.config import get_current_diffusion_config_or_none
 from vllm_omni.diffusion.distributed.parallel_state import get_sp_group
 from vllm_omni.diffusion.forward_context import get_forward_context, is_forward_context_available
 
@@ -32,6 +33,11 @@ class Attention(nn.Module):
         softmax_scale: float,
         num_kv_heads: int | None = None,
         prefix: str = "",
+        # Per-role backend selection (RFC: per-role attention backend)
+        role: str = "self",
+        role_category: str | None = None,
+        # Model-defined Q/K/V tensor layout hint for backend execution.
+        qkv_layout: str | None = None,
         # ulysses attention
         scatter_idx: int = 2,
         gather_idx: int = 1,
@@ -39,7 +45,34 @@ class Attention(nn.Module):
         skip_sequence_parallel: bool = False,
     ):
         super().__init__()
-        self.attn_backend = get_attn_backend(-1)
+
+        self.role = role
+        self.role_category = role_category
+        self.qkv_layout = qkv_layout
+
+        # Resolve backend via role-aware config.
+        # The global diffusion config is set during model init via
+        # set_current_diffusion_config(); no env-var re-parsing needed here.
+        backend_kwargs: dict | None = None
+        self.backend_pref = None
+
+        config = get_current_diffusion_config_or_none()
+        attention_config = config.attention_config if config is not None else None
+
+        attn_backend_cls, spec = get_attn_backend_for_role(
+            role=role,
+            head_size=head_size,
+            attention_config=attention_config,
+            role_category=role_category,
+        )
+        if spec is not None:
+            backend_kwargs = spec.extra or None
+            self.backend_pref = spec.backend
+            logger.debug("Attention(role=%s) → backend=%s", role, spec.backend)
+        else:
+            logger.debug("Attention(role=%s) → platform default", role)
+
+        self.attn_backend = attn_backend_cls
         self.attn_impl_cls = self.attn_backend.get_impl_cls()
         self.attention = self.attn_impl_cls(
             num_heads=num_heads,
@@ -47,6 +80,8 @@ class Attention(nn.Module):
             softmax_scale=softmax_scale,
             causal=causal,
             num_kv_heads=num_kv_heads,
+            qkv_layout=qkv_layout,
+            backend_kwargs=backend_kwargs,
         )
         # Instantiate fallback backend for float32 support
         self.sdpa_fallback = SDPABackend.get_impl_cls()(
@@ -55,8 +90,8 @@ class Attention(nn.Module):
             softmax_scale=softmax_scale,
             causal=causal,
             num_kv_heads=num_kv_heads,
+            qkv_layout=qkv_layout,
         )
-        self.backend_pref = None
 
         self.softmax_scale = softmax_scale
         self.scatter_idx = scatter_idx
@@ -69,21 +104,19 @@ class Attention(nn.Module):
         self.ring_pg = None
         self.ring_runner = None
 
-        try:
-            config = get_forward_context().omni_diffusion_config
-            self.backend_pref = config.attention_backend
+        if config is not None:
             if config.parallel_config.ring_degree > 1:
                 self.use_ring = True
                 try:
                     sp_group = get_sp_group()
                     self.ring_pg = sp_group.ring_group
-                    self.ring_runner = RingParallelAttention(sp_group)
+                    self.ring_runner = RingParallelAttention(
+                        sp_group,
+                        attn_backend_pref=self.backend_pref,
+                    )
                 except Exception:
                     self.use_ring = False
                     self.ring_runner = None
-        except Exception:
-            self.use_ring = False
-            self.ring_runner = None
 
         self.parallel_strategy = build_parallel_attention_strategy(
             scatter_idx=scatter_idx,

vllm_omni/diffusion/attention/parallel/ring.py

@@ -18,12 +18,11 @@ from vllm_omni.diffusion.attention.parallel.base import (
 )
 from vllm_omni.diffusion.distributed.group_coordinator import SequenceParallelGroupCoordinator
 
-# from vllm_omni.diffusion.attention.backends.ring_selector import AttnType # Already imported above
-from vllm_omni.diffusion.forward_context import get_forward_context
-
 if TYPE_CHECKING:
     from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
 
+logger = init_logger(__name__)
+
 
 @dataclass(frozen=True, slots=True)
 class _RingCtx(ParallelAttentionContext):
@@ -108,13 +107,8 @@ class RingParallelAttention:
             softmax_scale = query.shape[-1] ** -0.5
 
         backend_pref = self.attn_backend_pref
-        if backend_pref is None:
-            try:
-                config = get_forward_context().omni_diffusion_config
-                # config might not have attention_backend attribute if not updated
-                backend_pref = getattr(config, "attention_backend", None)
-            except Exception:
-                backend_pref = None
+        if backend_pref is not None:
+            backend_pref = backend_pref.lower()
 
         # Determine attention type with fallback chain: FA3 -> FA2 -> SDPA
         # FP32 is not supported by Flash Attention, force SDPA
@@ -122,7 +116,6 @@ class RingParallelAttention:
             backend_pref = "sdpa"
         elif not HAS_FA3 and not HAS_FLASH_ATTN:
             if backend_pref != "sdpa":
-                logger = init_logger(__name__)
                 logger.warning_once("Flash Attention (FA2/FA3) is not available! Force enabling SDPA.")
             backend_pref = "sdpa"
 
@@ -135,7 +128,7 @@ class RingParallelAttention:
             if attn_metadata.joint_strategy is not None:
                 joint_strategy = attn_metadata.joint_strategy
 
-        if backend_pref == "sdpa" or backend_pref == "torch":
+        if backend_pref in {"sdpa", "torch", "torch_sdpa"}:
             from vllm_omni.diffusion.attention.backends.ring_pytorch_attn import ring_pytorch_attn_func
 
             return ring_pytorch_attn_func(

vllm_omni/diffusion/attention/selector.py

@@ -3,23 +3,16 @@
 """
 Diffusion attention backend selector.
 
-This module provides the interface for selecting diffusion attention backends.
-The actual backend selection logic is delegated to the platform layer
-(vllm_omni.platforms), similar to how vLLM handles attention backend selection.
-
-Usage:
-    from vllm_omni.diffusion.attention.selector import get_attn_backend
-
-    # Get the appropriate backend for current platform
-    backend_cls = get_attn_backend(head_size=64)
-
-    # Or override via environment variable
-    # export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
+This module resolves diffusion attention backends from:
+1. per-role AttentionConfig
+2. platform default
 """
 
+from __future__ import annotations
+
 import importlib
-import os
 from functools import cache
+from typing import TYPE_CHECKING
 
 from vllm.logger import init_logger
 
@@ -27,6 +20,9 @@ from vllm_omni.diffusion.attention.backends.abstract import (
     AttentionBackend,
 )
 
+if TYPE_CHECKING:
+    from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec
+
 logger = init_logger(__name__)
 
 
@@ -52,34 +48,103 @@ def _load_backend_cls(cls_path: str) -> type[AttentionBackend]:
 
 
 @cache
-def get_attn_backend(head_size: int) -> type[AttentionBackend]:
-    """
-    Get attention backend for diffusion models.
+def _cached_get_backend_cls(
+    backend_name: str | None,
+    head_size: int,
+) -> type[AttentionBackend]:
+    """Cache backend class resolution by (backend_name, head_size).
 
-    The backend selection is delegated to the current platform
-    (vllm_omni.platforms.current_omni_platform), which selects the
-    appropriate backend based on:
-    1. User override via DIFFUSION_ATTENTION_BACKEND environment variable
-    2. Platform-specific defaults and capabilities
-
-    This is similar to how vLLM's get_attn_backend_cls works, where the
-    platform layer decides which backend to use based on hardware capabilities.
-
-    Args:
-        head_size: Head size for attention computation (may affect backend selection)
-
-    Returns:
-        The selected attention backend class
+    This ensures platform validation (compute capability checks, package
+    availability, etc.) runs only once per unique (backend_name, head_size)
+    combination, avoiding repeated log messages.
     """
     from vllm_omni.platforms import current_omni_platform
 
-    # Check environment variable for user override
-    selected_backend = os.environ.get("DIFFUSION_ATTENTION_BACKEND")
-
-    # Delegate to platform for backend selection
     backend_cls_path = current_omni_platform.get_diffusion_attn_backend_cls(
-        selected_backend=selected_backend,
+        selected_backend=backend_name,
         head_size=head_size,
     )
-
     return _load_backend_cls(backend_cls_path)
+
+
+@cache
+def _log_backend_resolution(
+    role: str,
+    role_category: str | None,
+    backend_name: str,
+    source: str,
+) -> None:
+    if role_category is not None:
+        logger.info(
+            "Resolved diffusion attention backend '%s' for role=%r (role_category=%r) via %s",
+            backend_name,
+            role,
+            role_category,
+            source,
+        )
+        return
+
+    logger.info(
+        "Resolved diffusion attention backend '%s' for role=%r via %s",
+        backend_name,
+        role,
+        source,
+    )
+
+
+def get_attn_backend_for_role(
+    role: str,
+    head_size: int,
+    attention_config: AttentionConfig | None = None,
+    role_category: str | None = None,
+) -> tuple[type[AttentionBackend], AttentionSpec | None]:
+    """
+    Get attention backend for a specific attention role.
+
+    Lookup precedence:
+      1. attention_config.per_role[role]           — exact match
+      2. attention_config.per_role[role_category]   — category fallback
+      3. attention_config.default                   — global default
+      4. Platform default                           — hardware-specific
+
+    Args:
+        role: Attention role string (e.g. "self", "cross", "joint",
+              "ltx2.audio_to_video")
+        head_size: Head size for attention computation
+        attention_config: The AttentionConfig from OmniDiffusionConfig.
+            If None, falls back to platform default behavior.
+        role_category: Optional category for fallback (e.g. "cross" for
+            "ltx2.audio_to_video")
+
+    Returns:
+        Tuple of (backend_class, AttentionSpec or None).
+        AttentionSpec is None when using platform default without explicit config.
+    """
+    # 1. Try config from OmniDiffusionConfig
+    spec = None
+    source = None
+    if attention_config is not None:
+        spec, source = attention_config.resolve_with_source(
+            role=role,
+            role_category=role_category,
+        )
+
+    if spec is not None:
+        backend_cls = _cached_get_backend_cls(spec.backend, head_size)
+        _log_backend_resolution(
+            role=role,
+            role_category=role_category,
+            backend_name=backend_cls.get_name(),
+            source=source or "attention_config",
+        )
+        return backend_cls, spec
+
+    # 2. Platform default
+    backend_cls = _cached_get_backend_cls(None, head_size)
+    _log_backend_resolution(
+        role=role,
+        role_category=role_category,
+        backend_name=backend_cls.get_name(),
+        source="platform default",
+    )
+    return backend_cls, None

vllm_omni/diffusion/config.py

@@ -0,0 +1,43 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Global diffusion config — mirrors vLLM's set_current_vllm_config pattern.
+
+Set once during model initialisation so that every ``Attention`` layer can
+read the ``OmniDiffusionConfig`` without coupling ``__init__`` to the
+runtime ``ForwardContext``.
+"""
+
+from __future__ import annotations
+
+from contextlib import contextmanager
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from vllm_omni.diffusion.data import OmniDiffusionConfig
+
+_current_diffusion_config: OmniDiffusionConfig | None = None
+
+
+@contextmanager
+def set_current_diffusion_config(config: OmniDiffusionConfig):
+    """Context manager that sets the global diffusion config for model init."""
+    global _current_diffusion_config
+    old = _current_diffusion_config
+    _current_diffusion_config = config
+    try:
+        yield
+    finally:
+        _current_diffusion_config = old
+
+
+def get_current_diffusion_config() -> OmniDiffusionConfig:
+    """Return the current diffusion config or raise."""
+    assert _current_diffusion_config is not None, (
+        "Diffusion config is not set. Wrap model construction with set_current_diffusion_config()."
+    )
+    return _current_diffusion_config
+
+
+def get_current_diffusion_config_or_none() -> OmniDiffusionConfig | None:
+    """Return the current diffusion config, or ``None`` if not set."""
+    return _current_diffusion_config

vllm_omni/diffusion/data.py

@@ -1,7 +1,7 @@
 # adapted from sglang and fastvideo
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-import enum
+import copy
 import os
 import random
 from collections.abc import Callable, Mapping
@@ -367,7 +367,7 @@ class OmniDiffusionConfig:
     tf_model_config: TransformerConfig = field(default_factory=TransformerConfig)
 
     # Attention
-    attention_backend: str | None = None
+    attention_config: "AttentionConfig" = field(default_factory=lambda: AttentionConfig())
 
     # Running mode
     # mode: ExecutionMode = ExecutionMode.INFERENCE
@@ -653,6 +653,10 @@ class OmniDiffusionConfig:
                     f"got {type(self.quantization_config)!r}"
                 )
 
+        # Match vLLM's config flow: parse entrypoint shorthands before the
+        # config object is built, and keep a single runtime truth source.
+        self.attention_config = build_attention_config(self.attention_config)
+
         if self.max_cpu_loras is None:
             self.max_cpu_loras = 1
         elif self.max_cpu_loras < 1:
@@ -794,6 +798,24 @@ class OmniDiffusionConfig:
         else:
             kwargs.pop("quantization", None)
 
+        # Map "diffusion_attention_config" to "attention_config" so the
+        # engine_args key (namespaced to avoid collision with vLLM's own
+        # attention_config) is mapped to the dataclass field.
+        if "diffusion_attention_config" in kwargs and "attention_config" not in kwargs:
+            kwargs["attention_config"] = kwargs.pop("diffusion_attention_config")
+        else:
+            kwargs.pop("diffusion_attention_config", None)
+
+        # Handle "diffusion_attention_backend" shorthand: merge into
+        # attention_config before field filtering.
+        diffusion_attn_backend = kwargs.pop("diffusion_attention_backend", None)
+        if diffusion_attn_backend is not None:
+            existing = kwargs.get("attention_config")
+            kwargs["attention_config"] = parse_attention_config(
+                existing,
+                attention_backend=diffusion_attn_backend,
+            )
+
         # Check environment variable as fallback for cache_backend
         # Support both old DIFFUSION_CACHE_ADAPTER and new DIFFUSION_CACHE_BACKEND for backwards compatibility
         if "cache_backend" not in kwargs:
@@ -848,19 +870,189 @@ class DiffusionRequestAbortedError(RuntimeError):
     """Raised when a diffusion request ends via user-visible abort."""
 
 
-class AttentionBackendEnum(enum.Enum):
-    FA = enum.auto()
-    SLIDING_TILE_ATTN = enum.auto()
-    TORCH_SDPA = enum.auto()
-    SAGE_ATTN = enum.auto()
-    SAGE_ATTN_THREE = enum.auto()
-    VIDEO_SPARSE_ATTN = enum.auto()
-    VMOBA_ATTN = enum.auto()
-    AITER = enum.auto()
-    NO_ATTENTION = enum.auto()
+@dataclass
+class AttentionSpec:
+    """Specifies a backend and its backend-specific parameters for one attention role."""
 
-    def __str__(self):
-        return self.name.lower()
+    backend: str  # registry name, e.g. "FLASH_ATTN"
+    extra: dict[str, Any] = field(default_factory=dict)  # backend-specific kwargs
+
+    def __post_init__(self) -> None:
+        if not isinstance(self.backend, str):
+            raise TypeError(f"Expected str for AttentionSpec.backend, got {type(self.backend)!r}")
+
+        if self.extra is None:
+            self.extra = {}
+        elif isinstance(self.extra, Mapping):
+            self.extra = dict(self.extra)
+        else:
+            raise TypeError(f"Expected dict for AttentionSpec.extra, got {type(self.extra)!r}")
+
+
+@dataclass
+class AttentionConfig:
+    """Per-role attention backend configuration.
+
+    Lookup precedence for a given (role, role_category):
+      1. per_role[role]         — exact match
+      2. per_role[role_category] — category fallback (e.g. "ltx2.audio_to_video" → "cross")
+      3. default                — global default
+      4. platform default       — unchanged platform logic
+    """
+
+    default: AttentionSpec | None = None
+    per_role: dict[str, AttentionSpec] = field(default_factory=dict)
+
+    def __post_init__(self) -> None:
+        if self.default is not None:
+            self.default = self._coerce_spec_or_none(self.default, "default")
+
+        normalized_per_role: dict[str, AttentionSpec] = {}
+        for role_key, spec_data in self._normalize_per_role_mapping(self.per_role).items():
+            spec = self._coerce_spec_or_none(spec_data, f"per_role[{role_key!r}]")
+            if spec is not None:
+                normalized_per_role[role_key] = spec
+        self.per_role = normalized_per_role
+
+    @staticmethod
+    def _coerce_spec(spec_data: Any, field_name: str) -> AttentionSpec:
+        if isinstance(spec_data, AttentionSpec):
+            return spec_data
+        if isinstance(spec_data, str):
+            return AttentionSpec(backend=spec_data)
+        if isinstance(spec_data, Mapping):
+            return AttentionSpec(**dict(spec_data))
+        raise TypeError(f"Expected str, dict, or AttentionSpec for {field_name}, got {type(spec_data)!r}")
+
+    @classmethod
+    def _coerce_spec_or_none(cls, spec_data: Any, field_name: str) -> AttentionSpec | None:
+        spec = cls._coerce_spec(spec_data, field_name)
+        if spec.backend.lower() == "auto":
+            return None
+        return spec
+
+    @classmethod
+    def _normalize_per_role_mapping(cls, raw_per_role: Any) -> dict[str, Any]:
+        if raw_per_role is None:
+            return {}
+        if not isinstance(raw_per_role, Mapping):
+            raise TypeError(f"Expected dict for AttentionConfig.per_role, got {type(raw_per_role)!r}")
+
+        normalized: dict[str, Any] = {}
+        for role_key, spec_data in raw_per_role.items():
+            cls._flatten_per_role_entry([role_key], spec_data, normalized)
+        return normalized
+
+    @classmethod
+    def _flatten_per_role_entry(
+        cls,
+        path: list[str],
+        node: Any,
+        normalized: dict[str, Any],
+    ) -> None:
+        role = ".".join(path)
+        if not isinstance(node, Mapping):
+            normalized[role] = node
+            return
+
+        spec_keys = {"backend", "extra"}
+        node_dict = dict(node)
+        node_keys = set(node_dict)
+        if node_keys & spec_keys:
+            if not node_keys <= spec_keys:
+                raise ValueError(
+                    f"Invalid per_role entry for role {role!r}: cannot mix backend/extra with nested role keys."
+                )
+            normalized[role] = node_dict
+            return
+
+        if not node_dict:
+            raise ValueError(f"Empty per_role entry for role {role!r}")
+
+        for child_key, child_value in node_dict.items():
+            cls._flatten_per_role_entry([*path, child_key], child_value, normalized)
+
+    def resolve_with_source(
+        self,
+        role: str = "self",
+        role_category: str | None = None,
+    ) -> tuple[AttentionSpec | None, str | None]:
+        """Resolve the AttentionSpec and report which config entry matched."""
+        spec = self.per_role.get(role)
+        if spec is not None:
+            return spec, f"attention_config.per_role[{role!r}]"
+        if role_category is not None:
+            spec = self.per_role.get(role_category)
+            if spec is not None:
+                return spec, f"attention_config.per_role[{role_category!r}] (role_category fallback)"
+        if self.default is not None:
+            return self.default, "attention_config.default"
+        return None, None
+
+
+def parse_attention_config(
+    attention_config: AttentionConfig | Mapping[str, Any] | None = None,
+    *,
+    attention_backend: str | None = None,
+) -> AttentionConfig:
+    """Pure type-conversion: coerce *attention_config* to an AttentionConfig.
+
+    Optionally merges an ``attention_backend`` shorthand into the config's
+    ``default`` field.  This does **not** read environment variables —
+    use :func:`build_attention_config` for the full normalisation that
+    should happen exactly once in ``OmniDiffusionConfig.__post_init__``.
+    """
+    if attention_config is None:
+        normalized = AttentionConfig()
+    elif isinstance(attention_config, AttentionConfig):
+        normalized = copy.deepcopy(attention_config)
+    elif isinstance(attention_config, Mapping):
+        normalized = AttentionConfig(**dict(attention_config))
+    else:
+        raise TypeError(
+            f"attention_config must be an AttentionConfig, mapping, or None; got {type(attention_config)!r}"
+        )
+
+    if attention_backend is not None:
+        if normalized.default is not None:
+            raise ValueError(
+                "--diffusion-attention-backend is mutually exclusive with --diffusion-attention-config.default.backend."
+            )
+        if attention_backend.lower() != "auto":
+            normalized.default = AttentionSpec(backend=attention_backend)
+
+    return normalized
+
+
+def build_attention_config(
+    attention_config: AttentionConfig | Mapping[str, Any] | None = None,
+) -> AttentionConfig:
+    """Normalize diffusion attention config — the single authoritative entry point.
+
+    Called exactly once in ``OmniDiffusionConfig.__post_init__``.
+    Handles type-conversion **and** env-var fallback
+    (``DIFFUSION_ATTENTION_BACKEND``).
+    """
+    normalized = parse_attention_config(attention_config)
+
+    if normalized.default is not None:
+        return normalized
+
+    env_attention_backend = os.environ.get("DIFFUSION_ATTENTION_BACKEND")
+    if env_attention_backend is None:
+        return normalized
+
+    if env_attention_backend.lower() == "auto":
+        return normalized
+
+    normalized.default = AttentionSpec(backend=env_attention_backend)
+    logger.info(
+        "Parsed attention config from DIFFUSION_ATTENTION_BACKEND '%s': default=%s, per_role=%s",
+        env_attention_backend,
+        normalized.default,
+        {k: v.backend for k, v in normalized.per_role.items()},
+    )
+    return normalized
 
 
 @dataclass

vllm_omni/diffusion/hooks/sequence_parallel.py

@@ -375,7 +375,7 @@ class SequenceParallelSplitHook(ModelHook):
         2. Creates an attention mask indicating valid vs padding positions
         3. Stores the mask and padding info in ForwardContext
         """
-        from vllm_omni.diffusion.attention.selector import get_attn_backend
+        from vllm_omni.diffusion.attention.selector import get_attn_backend_for_role
         from vllm_omni.diffusion.distributed.parallel_state import (
             get_ring_parallel_world_size,
             get_sequence_parallel_rank,
@@ -395,7 +395,17 @@ class SequenceParallelSplitHook(ModelHook):
             return sp_shard(x, dim, validate=False)
 
         # Check backend compatibility
-        attn_backend = get_attn_backend(-1)
+        attention_config = None
+        if is_forward_context_available():
+            od_config = get_forward_context().omni_diffusion_config
+            if od_config is not None:
+                attention_config = od_config.attention_config
+
+        attn_backend, _ = get_attn_backend_for_role(
+            role="self",
+            head_size=-1,
+            attention_config=attention_config,
+        )
         if not attn_backend.supports_attention_mask:
             raise ValueError(
                 f"Sequence length ({seq_len}) is not divisible by SP world size ({world_size}). "

vllm_omni/diffusion/model_loader/diffusers_loader.py

@@ -319,8 +319,11 @@ class DiffusersPipelineLoader:
                     elif load_format == "diffusers":
                         model = DiffusersAdapterPipeline(od_config=od_config, device=target_device)
                     elif load_format == "custom_pipeline":
+                        from vllm_omni.diffusion.config import set_current_diffusion_config
+
                         model_cls = resolve_obj_by_qualname(custom_pipeline_name)
-                        model = model_cls(od_config=od_config)
+                        with set_current_diffusion_config(od_config):
+                            model = model_cls(od_config=od_config)
                     else:
                         raise ValueError(f"Unknown load_format: {load_format}")
                 logger.debug("Loading weights on %s ...", load_device)
@@ -588,8 +591,11 @@ class DiffusersPipelineLoader:
         if load_format == "default":
             model = initialize_model(od_config)
         elif load_format == "custom_pipeline":
+            from vllm_omni.diffusion.config import set_current_diffusion_config
+
             model_cls = resolve_obj_by_qualname(custom_pipeline_name)
-            model = model_cls(od_config=od_config)
+            with set_current_diffusion_config(od_config):
+                model = model_cls(od_config=od_config)
         self.load_weights(model)
 
         # Collect all transformers to shard (some models have transformer_2 for MoE)

vllm_omni/diffusion/models/diffusers_adapter/pipeline_diffusers_adapter.py

@@ -15,7 +15,6 @@ It does NOT support:
 
 import inspect
 import logging
-import os
 import re
 from typing import Any, cast
 
@@ -211,7 +210,8 @@ class DiffusersAdapterPipeline(nn.Module, DiffusionPipelineProfilerMixin):
             logging.info("No transformer found in diffusers pipeline. Skipping attention backend setting.")
             return
 
-        attention_backend_config = self.od_config.attention_backend or os.environ.get("DIFFUSION_ATTENTION_BACKEND")
+        default_spec = self.od_config.attention_config.default
+        attention_backend_config = default_spec.backend if default_spec is not None else None
         attention_backend_attempts: list[str] = []
         match attention_backend_config:
             case "FLASH_ATTN" | None:

vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py

@@ -373,6 +373,7 @@ class WanSelfAttention(nn.Module):
             num_kv_heads=self.num_kv_heads,
             softmax_scale=1.0 / (head_dim**0.5),
             causal=False,
+            role="self",
         )
 
     def forward(
@@ -525,6 +526,8 @@ class WanCrossAttention(nn.Module):
             num_kv_heads=self.num_heads,
             softmax_scale=1.0 / (head_dim**0.5),
             causal=False,
+            role="cross",
+            qkv_layout="BSND",
             skip_sequence_parallel=True,
         )
 

vllm_omni/diffusion/registry.py

@@ -8,6 +8,7 @@ from vllm.logger import init_logger
 from vllm.model_executor.model_loader.utils import configure_quant_config
 from vllm.model_executor.models.registry import _LazyRegisteredModel, _ModelRegistry
 
+from vllm_omni.diffusion.config import set_current_diffusion_config
 from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import DistributedVaeMixin
 from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelConfig, get_sp_plan_from_model
@@ -280,7 +281,7 @@ def _prepare_diffusion_quant_config(
     model_class: type[nn.Module],
 ) -> None:
     """Prepare diffusion quant config using vLLM-style model bindings."""
-    quant_config = od_config.quantization_config
+    quant_config = getattr(od_config, "quantization_config", None)
     if quant_config is None:
         return
     if hasattr(quant_config, "maybe_update_config"):
@@ -314,7 +315,8 @@ def initialize_model(
     model_class = DiffusionModelRegistry._try_load_model_cls(od_config.model_class_name)
     if model_class is not None:
         _prepare_diffusion_quant_config(od_config, model_class)
-        model = model_class(od_config=od_config)
+        with set_current_diffusion_config(od_config):
+            model = model_class(od_config=od_config)
 
         vae_pp_size = od_config.parallel_config.vae_patch_parallel_size
         is_distributed_vae = hasattr(model, "vae") and isinstance(model.vae, DistributedVaeMixin)

vllm_omni/engine/async_omni_engine.py

@@ -34,7 +34,7 @@ from vllm.v1.engine import EngineCoreRequest
 from vllm.v1.engine.input_processor import InputProcessor
 
 from vllm_omni.config.stage_config import strip_parent_engine_args
-from vllm_omni.diffusion.data import DiffusionParallelConfig
+from vllm_omni.diffusion.data import DiffusionParallelConfig, parse_attention_config
 from vllm_omni.diffusion.diffusion_engine import supports_audio_output
 from vllm_omni.diffusion.stage_diffusion_client import StageDiffusionClient
 from vllm_omni.diffusion.stage_diffusion_proc import (
@@ -1441,6 +1441,16 @@ class AsyncOmniEngine:
         model_class_name = kwargs.get("model_class_name", None)
         final_output_type = "audio" if model_class_name and supports_audio_output(model_class_name) else "image"
 
+        attention_config = None
+        if (
+            kwargs.get("diffusion_attention_config") is not None
+            or kwargs.get("diffusion_attention_backend") is not None
+        ):
+            attention_config = parse_attention_config(
+                kwargs.get("diffusion_attention_config"),
+                attention_backend=kwargs.get("diffusion_attention_backend"),
+            )
+
         stage_engine_args = {
             "max_num_seqs": kwargs.get("max_num_seqs") or 1,
             "parallel_config": parallel_config,
@@ -1467,6 +1477,7 @@ class AsyncOmniEngine:
             "enable_multithread_weight_load": kwargs.get("enable_multithread_weight_load", True),
             "num_weight_load_threads": kwargs.get("num_weight_load_threads", 4),
             "quantization": kwargs.get("quantization", None),
+            **({"diffusion_attention_config": attention_config} if attention_config is not None else {}),
             "force_cutlass_fp8": bool(kwargs.get("force_cutlass_fp8", False)),
             "enable_diffusion_pipeline_profiler": kwargs.get("enable_diffusion_pipeline_profiler", False),
             "enable_ar_profiler": kwargs.get("enable_ar_profiler", False),
@@ -1605,6 +1616,19 @@ class AsyncOmniEngine:
                 if lora_scale is not None:
                     if not hasattr(cfg.engine_args, "lora_scale") or cfg.engine_args.lora_scale is None:
                         cfg.engine_args.lora_scale = lora_scale
+                if (
+                    kwargs.get("diffusion_attention_config") is not None
+                    or kwargs.get("diffusion_attention_backend") is not None
+                ):
+                    has_stage_attention = (
+                        hasattr(cfg.engine_args, "diffusion_attention_config")
+                        and cfg.engine_args.diffusion_attention_config is not None
+                    )
+                    if not has_stage_attention:
+                        cfg.engine_args.diffusion_attention_config = parse_attention_config(
+                            kwargs.get("diffusion_attention_config"),
+                            attention_backend=kwargs.get("diffusion_attention_backend"),
+                        )
                 quantization_config = kwargs.get("quantization_config")
                 if quantization_config is not None:
                     if (

vllm_omni/engine/stage_init_utils.py

@@ -318,7 +318,7 @@ def extract_stage_metadata(stage_config: Any) -> StageMetadata:
     engine_args = stage_config.engine_args
 
     if current_omni_platform.is_rocm():
-        if engine_args.get("attention_backend") is None:
+        if stage_type != "diffusion" and engine_args.get("attention_backend") is None:
             from vllm._aiter_ops import rocm_aiter_ops
 
             if rocm_aiter_ops.is_enabled():
@@ -574,6 +574,14 @@ def build_engine_args_dict(
     if engine_args_dict.get("async_chunk", False):
         engine_args_dict["stage_connector_spec"] = dict(stage_connector_spec or {})
 
+    if stage_type == "diffusion":
+        from vllm_omni.diffusion.data import parse_attention_config
+
+        if engine_args_dict.get("diffusion_attention_config") is not None:
+            engine_args_dict["diffusion_attention_config"] = parse_attention_config(
+                engine_args_dict.get("diffusion_attention_config"),
+            )
+
     if stage_type != "diffusion":
         resolve_worker_cls(engine_args_dict)
 

vllm_omni/entrypoints/cli/serve.py

@@ -370,6 +370,32 @@ class OmniServeCommand(CLISubcommand):
             help="Number of replica groups for HSDP. Each group holds a full sharded copy.",
         )
 
+        # Attention backend configuration
+        omni_config_group.add_argument(
+            "--diffusion-attention-backend",
+            dest="diffusion_attention_backend",
+            type=str,
+            default=None,
+            help="Diffusion attention backend (shorthand). "
+            "Sets the default backend for all diffusion attention roles, e.g. 'FLASH_ATTN'. "
+            "May be combined with --diffusion-attention-config.per_role.* overrides, "
+            "but mutually exclusive with --diffusion-attention-config.default.backend.",
+        )
+        omni_config_group.add_argument(
+            "--diffusion-attention-config",
+            "-dac",
+            dest="diffusion_attention_config",
+            type=json.loads,
+            default=None,
+            help="Diffusion attention config. Accepts JSON or vLLM-style dotted flags. "
+            "Examples: "
+            "--diffusion-attention-config.default.backend FLASH_ATTN, "
+            "--diffusion-attention-config.per_role.self.backend SPARSE_BLOCK, "
+            "--diffusion-attention-config.per_role.cross.backend SAGE_ATTN, "
+            '--diffusion-attention-config \'{"default": {"backend": "FLASH_ATTN"}, '
+            '"per_role": {"cross": {"backend": "SAGE_ATTN"}}}\'.',
+        )
+
         # Cache optimization parameters
         omni_config_group.add_argument(
             "--cache-backend",

vllm_omni/platforms/cuda/platform.py

@@ -84,17 +84,17 @@ class CudaOmniPlatform(OmniPlatform, CudaPlatformBase):
                     )
                 elif not packages_available:
                     logger.warning("Flash Attention packages not available. Falling back to TORCH_SDPA backend.")
-                logger.info("Defaulting to diffusion attention backend SDPA")
+                logger.debug("Defaulting to diffusion attention backend SDPA")
                 return DiffusionAttentionBackendEnum.TORCH_SDPA.get_path()
             backend = DiffusionAttentionBackendEnum[backend_upper]
-            logger.info("Using diffusion attention backend '%s'", backend_upper)
+            logger.debug("Using diffusion attention backend '%s'", backend_upper)
             return backend.get_path()
 
         if flash_attn_supported:
-            logger.info("Defaulting to diffusion attention backend FLASH_ATTN")
+            logger.debug("Defaulting to diffusion attention backend FLASH_ATTN")
             return DiffusionAttentionBackendEnum.FLASH_ATTN.get_path()
 
-        logger.info("Defaulting to diffusion attention backend SDPA")
+        logger.debug("Defaulting to diffusion attention backend SDPA")
         return DiffusionAttentionBackendEnum.TORCH_SDPA.get_path()
 
     @classmethod

vllm_omni/platforms/musa/platform.py

@@ -86,17 +86,17 @@ class MUSAOmniPlatform(OmniPlatform, MUSAPlatformBase):
                     )
                 elif not packages_available:
                     logger.warning("Flash Attention (mate package) not available. Falling back to TORCH_SDPA backend.")
-                logger.info("Defaulting to diffusion attention backend SDPA")
+                logger.debug("Defaulting to diffusion attention backend SDPA")
                 return DiffusionAttentionBackendEnum.TORCH_SDPA.get_path()
             backend = DiffusionAttentionBackendEnum[backend_upper]
-            logger.info("Using diffusion attention backend '%s'", backend_upper)
+            logger.debug("Using diffusion attention backend '%s'", backend_upper)
             return backend.get_path()
 
         if flash_attn_supported:
-            logger.info("Defaulting to diffusion attention backend FLASH_ATTN")
+            logger.debug("Defaulting to diffusion attention backend FLASH_ATTN")
             return DiffusionAttentionBackendEnum.FLASH_ATTN.get_path()
 
-        logger.info("Defaulting to diffusion attention backend SDPA")
+        logger.debug("Defaulting to diffusion attention backend SDPA")
         return DiffusionAttentionBackendEnum.TORCH_SDPA.get_path()
 
     @classmethod

vllm_omni/platforms/npu/platform.py

@@ -78,7 +78,7 @@ class NPUOmniPlatform(OmniPlatform, NPUPlatform):
         if selected_backend is not None:
             backend_upper = selected_backend.upper()
             backend = DiffusionAttentionBackendEnum[backend_upper]
-            logger.info("Using diffusion attention backend '%s'", backend_upper)
+            logger.debug("Using diffusion attention backend '%s'", backend_upper)
             return backend.get_path()
 
         # Try FLASH_ATTN if mindiesd is available, otherwise fall back to SDPA
@@ -86,10 +86,10 @@ class NPUOmniPlatform(OmniPlatform, NPUPlatform):
             # Configure ASCEND_CUSTOM_OPP_PATH for mindiesd custom ops upon import
             import mindiesd  # noqa: F401
 
-            logger.info("Defaulting to diffusion attention backend FLASH_ATTN")
+            logger.debug("Defaulting to diffusion attention backend FLASH_ATTN")
             return DiffusionAttentionBackendEnum.FLASH_ATTN.get_path()
 
-        logger.info("Falling back to diffusion attention backend SDPA")
+        logger.debug("Falling back to diffusion attention backend SDPA")
         return DiffusionAttentionBackendEnum.TORCH_SDPA.get_path()
 
     @classmethod

vllm_omni/platforms/rocm/platform.py

@@ -88,19 +88,19 @@ class RocmOmniPlatform(OmniPlatform, RocmPlatform):
                     "Flash Attention requires `aiter` library which is only supported "
                     "on gfx942 and gfx950. Falling back to TORCH_SDPA backend."
                 )
-                logger.info("Defaulting to diffusion attention backend SDPA")
+                logger.debug("Defaulting to diffusion attention backend SDPA")
                 return DiffusionAttentionBackendEnum.TORCH_SDPA.get_path()
             backend = DiffusionAttentionBackendEnum[backend_upper]
-            logger.info("Using diffusion attention backend '%s'", backend_upper)
+            logger.debug("Using diffusion attention backend '%s'", backend_upper)
             return backend.get_path()
 
         # Choose to enable Flash Attention by default on ROCm
         # whenever possible as it is the fastest backend
         if aiter_supported:
-            logger.info("Defaulting to diffusion attention backend FLASH_ATTN")
+            logger.debug("Defaulting to diffusion attention backend FLASH_ATTN")
             return DiffusionAttentionBackendEnum.FLASH_ATTN.get_path()
 
-        logger.info("Defaulting to diffusion attention backend SDPA")
+        logger.debug("Defaulting to diffusion attention backend SDPA")
         return DiffusionAttentionBackendEnum.TORCH_SDPA.get_path()
 
     @classmethod

vllm_omni/platforms/xpu/platform.py

@@ -39,10 +39,10 @@ class XPUOmniPlatform(OmniPlatform, XPUPlatform):
         if selected_backend is not None:
             backend_upper = selected_backend.upper()
             backend = DiffusionAttentionBackendEnum[backend_upper]
-            logger.info("Using diffusion attention backend '%s'", backend_upper)
+            logger.debug("Using diffusion attention backend '%s'", backend_upper)
             return backend.get_path()
 
-        logger.info("Defaulting to diffusion attention backend SDPA")
+        logger.debug("Defaulting to diffusion attention backend SDPA")
         return DiffusionAttentionBackendEnum.TORCH_SDPA.get_path()
 
     @classmethod