diff --git a/tests/diffusion/attention/test_attention_config.py b/tests/diffusion/attention/test_attention_config.py new file mode 100644 index 000000000..297f760fb --- /dev/null +++ b/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" diff --git a/tests/diffusion/attention/test_attention_sp.py b/tests/diffusion/attention/test_attention_sp.py index c99ba9772..e5fb93d71 100644 --- a/tests/diffusion/attention/test_attention_sp.py +++ b/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 diff --git a/tests/diffusion/attention/test_ulysses_uaa.py b/tests/diffusion/attention/test_ulysses_uaa.py index 67907d7ef..a9e2021c3 100644 --- a/tests/diffusion/attention/test_ulysses_uaa.py +++ b/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, diff --git a/tests/diffusion/cache/test_teacache_extractors.py b/tests/diffusion/cache/test_teacache_extractors.py index b4600dba9..ab71d186a 100644 --- a/tests/diffusion/cache/test_teacache_extractors.py +++ b/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): diff --git a/tests/diffusion/model_loader/test_diffusers_loader.py b/tests/diffusion/model_loader/test_diffusers_loader.py index 1b13c400f..9434702f9 100644 --- a/tests/diffusion/model_loader/test_diffusers_loader.py +++ b/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 diff --git a/tests/engine/test_async_omni_engine_stage_init.py b/tests/engine/test_async_omni_engine_stage_init.py index 9474f8727..fcbf3788f 100644 --- a/tests/engine/test_async_omni_engine_stage_init.py +++ b/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" diff --git a/tests/entrypoints/test_async_omni_diffusion_config.py b/tests/entrypoints/test_async_omni_diffusion_config.py index f06ede401..79a6b376b 100644 --- a/tests/entrypoints/test_async_omni_diffusion_config.py +++ b/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" diff --git a/vllm_omni/diffusion/attention/backends/abstract.py b/vllm_omni/diffusion/attention/backends/abstract.py index 472fde422..6dd785e66 100644 --- a/vllm_omni/diffusion/attention/backends/abstract.py +++ b/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 diff --git a/vllm_omni/diffusion/attention/backends/flash_attn.py b/vllm_omni/diffusion/attention/backends/flash_attn.py index f987fbbbe..341398446 100644 --- a/vllm_omni/diffusion/attention/backends/flash_attn.py +++ b/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 diff --git a/vllm_omni/diffusion/attention/backends/sage_attn.py b/vllm_omni/diffusion/attention/backends/sage_attn.py index af34a0633..080a241c1 100644 --- a/vllm_omni/diffusion/attention/backends/sage_attn.py +++ b/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, diff --git a/vllm_omni/diffusion/attention/backends/sdpa.py b/vllm_omni/diffusion/attention/backends/sdpa.py index 3585689dd..ab71e753b 100644 --- a/vllm_omni/diffusion/attention/backends/sdpa.py +++ b/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, diff --git a/vllm_omni/diffusion/attention/layer.py b/vllm_omni/diffusion/attention/layer.py index 4fdf2ff16..4d7c0d362 100644 --- a/vllm_omni/diffusion/attention/layer.py +++ b/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, diff --git a/vllm_omni/diffusion/attention/parallel/ring.py b/vllm_omni/diffusion/attention/parallel/ring.py index 090d2ef61..94cd0d205 100644 --- a/vllm_omni/diffusion/attention/parallel/ring.py +++ b/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( diff --git a/vllm_omni/diffusion/attention/selector.py b/vllm_omni/diffusion/attention/selector.py index 920396e66..6eede02c2 100644 --- a/vllm_omni/diffusion/attention/selector.py +++ b/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 diff --git a/vllm_omni/diffusion/config.py b/vllm_omni/diffusion/config.py new file mode 100644 index 000000000..4faf482ce --- /dev/null +++ b/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 diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 80171e1d6..e2f91d9aa 100644 --- a/vllm_omni/diffusion/data.py +++ b/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 diff --git a/vllm_omni/diffusion/hooks/sequence_parallel.py b/vllm_omni/diffusion/hooks/sequence_parallel.py index b1c515070..2aaa3c169 100644 --- a/vllm_omni/diffusion/hooks/sequence_parallel.py +++ b/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}). " diff --git a/vllm_omni/diffusion/model_loader/diffusers_loader.py b/vllm_omni/diffusion/model_loader/diffusers_loader.py index 1d2c9e357..f4ebfcb24 100644 --- a/vllm_omni/diffusion/model_loader/diffusers_loader.py +++ b/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) diff --git a/vllm_omni/diffusion/models/diffusers_adapter/pipeline_diffusers_adapter.py b/vllm_omni/diffusion/models/diffusers_adapter/pipeline_diffusers_adapter.py index 81733aa8b..e687cb7b2 100644 --- a/vllm_omni/diffusion/models/diffusers_adapter/pipeline_diffusers_adapter.py +++ b/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: diff --git a/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py b/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py index 98e10b408..e80b116ab 100644 --- a/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py +++ b/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, ) diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py index d69e9acd5..f530115d7 100644 --- a/vllm_omni/diffusion/registry.py +++ b/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) diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index 611f48bdc..b2dc839c9 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/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 ( diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py index b6d601de5..ce68a23da 100644 --- a/vllm_omni/engine/stage_init_utils.py +++ b/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) diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index 6880cf7c5..540d5c0cf 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/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", diff --git a/vllm_omni/platforms/cuda/platform.py b/vllm_omni/platforms/cuda/platform.py index 849e7028a..a3cc32c52 100644 --- a/vllm_omni/platforms/cuda/platform.py +++ b/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 diff --git a/vllm_omni/platforms/musa/platform.py b/vllm_omni/platforms/musa/platform.py index acb27c468..75f20a5a7 100644 --- a/vllm_omni/platforms/musa/platform.py +++ b/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 diff --git a/vllm_omni/platforms/npu/platform.py b/vllm_omni/platforms/npu/platform.py index 1d3e221ff..8a30b4c7a 100644 --- a/vllm_omni/platforms/npu/platform.py +++ b/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 diff --git a/vllm_omni/platforms/rocm/platform.py b/vllm_omni/platforms/rocm/platform.py index 96ee5a498..d7ad2b0f7 100644 --- a/vllm_omni/platforms/rocm/platform.py +++ b/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 diff --git a/vllm_omni/platforms/xpu/platform.py b/vllm_omni/platforms/xpu/platform.py index b2982d263..84707df08 100644 --- a/vllm_omni/platforms/xpu/platform.py +++ b/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