# 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, ) pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] class TestAttentionSpec: def test_construct_no_skip_softmax(self): spec = AttentionSpec(backend="FLASH_ATTN") assert spec.skip_softmax is None def test_skip_softmax_mapping_coerced(self): spec = AttentionSpec(backend="TRTLLM_ATTN", skip_softmax={"target_sparsity": 0.5}) assert spec.backend == "TRTLLM_ATTN" assert spec.skip_softmax.target_sparsity == 0.5 def test_invalid_backend_type(self): with pytest.raises(TypeError): AttentionSpec(backend=123) # type: ignore[arg-type] def test_skip_softmax_rejected_on_non_trtllm(self): with pytest.raises(ValueError, match="only supported by the TRTLLM_ATTN"): AttentionSpec(backend="TORCH_SDPA", skip_softmax={"target_sparsity": 0.5}) def test_quant_serialized_with_defaults_and_overrides(self): assert AttentionSpec(backend="TRTLLM_ATTN", quant={"dtype_qk": "fp8_e4m3"}).backend_kwargs()["quant"] == { "dtype_qk": "fp8_e4m3", "q_block_size": 1, "k_block_size": 16, } assert AttentionSpec( backend="TRTLLM_ATTN", quant={"dtype_qk": "int8", "q_block_size": 4, "k_block_size": 16} ).backend_kwargs()["quant"] == {"dtype_qk": "int8", "q_block_size": 4, "k_block_size": 16} def test_quant_superset_fields_passed_through(self): # FLASHINFER_ATTN-style config (dtype_vo / flashinfer_backend) round-trips on the shared spec. bk = AttentionSpec( backend="FLASHINFER_ATTN", quant={"dtype_qk": "bfloat16", "dtype_vo": "fp8_e4m3", "flashinfer_backend": "trtllm-gen"}, ).backend_kwargs()["quant"] assert bk["dtype_qk"] == "bfloat16" assert bk["dtype_vo"] == "fp8_e4m3" assert bk["flashinfer_backend"] == "trtllm-gen" def test_quant_and_skip_softmax_coexist(self): bk = AttentionSpec( backend="TRTLLM_ATTN", quant={"dtype_qk": "int8"}, skip_softmax={"target_sparsity": 0.5} ).backend_kwargs() assert bk["target_sparsity"] == 0.5 and bk["quant"]["dtype_qk"] == "int8" @pytest.mark.parametrize( "spec, match", [ ({"backend": "TORCH_SDPA", "quant": {"dtype_qk": "int8"}}, "only supported by the TRTLLM_ATTN"), ({"backend": "TRTLLM_ATTN", "quant": {"dtype_qk": "int4"}}, "quant.dtype_qk"), ({"backend": "TRTLLM_ATTN", "quant": {"dtype_qk": "int8", "k_block_size": 8}}, "quant.k_block_size"), ], ) def test_quant_validation_rejects(self, spec, match): with pytest.raises(ValueError, match=match): AttentionSpec(**spec) def test_block_sparse_defaults_applied_when_backend_selected(self): spec = AttentionSpec(backend="RAINFUSION_ATTN") assert spec.block_sparse.sparsity == 0.8 assert spec.backend_kwargs() == {"sparsity": 0.8, "start_step": 0} def test_block_sparse_skip_layers_selector_expanded(self): spec = AttentionSpec( backend="RAINFUSION_ATTN", block_sparse={"sparsity": 0.9, "start_step": 12, "skip_layers": "0-2,38"}, ) assert spec.block_sparse.skip_layer_indices == {0, 1, 2, 38} assert spec.backend_kwargs() == { "sparsity": 0.9, "start_step": 12, "skip_layers": [0, 1, 2, 38], } def test_block_sparse_rejected_on_dense_backend(self): with pytest.raises(ValueError, match="block_sparse is only supported by"): AttentionSpec(backend="FLASH_ATTN", block_sparse={"sparsity": 0.8}) @pytest.mark.parametrize( "block_sparse", [{"sparsity": 1.5}, {"start_step": -1}], ) def test_block_sparse_invalid_values(self, block_sparse): with pytest.raises(ValueError): AttentionSpec(backend="RAINFUSION_ATTN", block_sparse=block_sparse) 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": "TRTLLM_ATTN", "skip_softmax": {"target_sparsity": 0.5}}, "cross": "SAGE_ATTN", }, ) assert config.default.backend == "FLASH_ATTN" assert config.per_role["self"].backend == "TRTLLM_ATTN" assert config.per_role["self"].skip_softmax.target_sparsity == 0.5 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"), "cross": AttentionSpec(backend="SAGE_ATTN"), "ltx2.audio_self": AttentionSpec(backend="FLASH_ATTN"), "ltx2.audio_to_video": AttentionSpec(backend="FLASH_ATTN"), }, ) # 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" # 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.""" @pytest.fixture(autouse=True) def _clear_diffusion_attention_backend_env(self, monkeypatch): # OmniDiffusionConfig.__post_init__ applies DIFFUSION_ATTENTION_BACKEND via # build_attention_config(); clear it so these tests assert config defaults, # not whatever the process inherited from CI / sibling tests. monkeypatch.delenv("DIFFUSION_ATTENTION_BACKEND", raising=False) def test_diffusion_attention_backend_sets_default(self): config = OmniDiffusionConfig.from_kwargs(diffusion_attention_backend="SAGE_ATTN") assert isinstance(config.diffusion_attention_config, AttentionConfig) assert config.diffusion_attention_config.default is not None assert config.diffusion_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.diffusion_attention_config, AttentionConfig) assert config.diffusion_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( diffusion_attention_config={ "default": {"backend": "FLASH_ATTN"}, "per_role": {"self": "SPARSE_BLOCK"}, } ) assert config.diffusion_attention_config.default.backend == "FLASH_ATTN" assert config.diffusion_attention_config.per_role["self"].backend == "SPARSE_BLOCK" def test_no_diffusion_attention_config_defaults_to_empty(self): config = OmniDiffusionConfig() assert isinstance(config.diffusion_attention_config, AttentionConfig) assert config.diffusion_attention_config.default is None assert config.diffusion_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, allow_trtllm_default=False, ): captured["role"] = role captured["head_size"] = head_size captured["role_category"] = role_category captured["attention_config"] = attention_config return _FakeBackend, AttentionSpec(backend="TRTLLM_ATTN", skip_softmax={"target_sparsity": 0.5}) 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( diffusion_attention_config=AttentionConfig( default=AttentionSpec(backend="FLASH_ATTN"), per_role={"cross": AttentionSpec(backend="TRTLLM_ATTN", skip_softmax={"target_sparsity": 0.5})}, ), parallel_config=SimpleNamespace(ring_degree=2), diffusion_kv_cache_dtype=None, diffusion_kv_cache_skip_step_indices=None, diffusion_kv_cache_skip_layer_indices=None, ) 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.diffusion_attention_config assert attn.backend_pref == "TRTLLM_ATTN" assert attn.attention.kwargs["backend_kwargs"] == {"target_sparsity": 0.5} 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 == "TRTLLM_ATTN" class TestDiffusionKvCacheQuantization: @staticmethod def _install_attention_init_stubs(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 @staticmethod def supports_kv_cache_dtype(kv_cache_dtype, platform_key) -> bool: return True 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", lambda role, head_size, attention_config=None, role_category=None, allow_trtllm_default=False: ( _FakeBackend, None, ), ) 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) def test_diffusion_kv_cache_dtype_none_does_not_trigger_ring_quantization_error(self, monkeypatch): self._install_attention_init_stubs(monkeypatch) od_config = SimpleNamespace( diffusion_attention_config=AttentionConfig(), parallel_config=SimpleNamespace(ring_degree=2), diffusion_kv_cache_dtype=None, diffusion_kv_cache_skip_step_indices=None, diffusion_kv_cache_skip_layer_indices=None, ) with set_current_diffusion_config(od_config): attn = Attention( num_heads=4, head_size=64, causal=False, softmax_scale=1.0, ) assert attn._kv_cache_dtype is None def test_diffusion_kv_cache_dtype_auto_does_not_trigger_ring_quantization_error(self, monkeypatch): self._install_attention_init_stubs(monkeypatch) od_config = SimpleNamespace( diffusion_attention_config=AttentionConfig(), parallel_config=SimpleNamespace(ring_degree=2), diffusion_kv_cache_dtype="auto", diffusion_kv_cache_skip_step_indices=None, diffusion_kv_cache_skip_layer_indices=None, ) with set_current_diffusion_config(od_config): attn = Attention( num_heads=4, head_size=64, causal=False, softmax_scale=1.0, ) assert attn._kv_cache_dtype is None def test_diffusion_kv_cache_dtype_fp8_raises_with_ring_attention(self, monkeypatch): self._install_attention_init_stubs(monkeypatch) od_config = SimpleNamespace( diffusion_attention_config=AttentionConfig(), parallel_config=SimpleNamespace(ring_degree=2), diffusion_kv_cache_dtype="fp8", diffusion_kv_cache_skip_step_indices=None, diffusion_kv_cache_skip_layer_indices=None, ) with set_current_diffusion_config(od_config): with pytest.raises(ValueError, match="KV quantization is not compatible with ring attention"): Attention( num_heads=4, head_size=64, causal=False, softmax_scale=1.0, )