# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni 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 dataclasses import replace from types import SimpleNamespace from unittest.mock import Mock import pytest import torch import vllm_omni.diffusion.attention.layer as layer_mod import vllm_omni.diffusion.attention.selector as selector_mod from vllm_omni.diffusion.attention.backends.abstract import ( AttentionMetadata, PackedPaddingMetadata, QueryRange, ) from vllm_omni.diffusion.attention.backends.flash_attn import FlashAttentionBackend 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_flashinfer_backend_serialized_without_dtype(self): spec = AttentionSpec(backend="FLASHINFER_ATTN", quant={"flashinfer_backend": "fa2"}) assert spec.quant is not None assert spec.quant.enabled is True assert spec.backend_kwargs() == {"quant": {"flashinfer_backend": "fa2"}} 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_fastvideo_vsa_topk_serialized(self): spec = AttentionSpec(backend="FASTVIDEO_VSA", fastvideo_vsa_topk=96) assert spec.backend_kwargs() == {"topk": 96} def test_fastvideo_vsa_topk_rejected_for_other_backend(self): with pytest.raises(ValueError, match="only supported by the FASTVIDEO_VSA"): AttentionSpec(backend="TORCH_SDPA", fastvideo_vsa_topk=96) 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, "end_step": 0, "precision": "bf16", } 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, "end_step": 0, "precision": "bf16", "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}, {"end_step": -1}, {"precision": "bogus"}], ) def test_block_sparse_invalid_values(self, block_sparse): with pytest.raises(ValueError): AttentionSpec(backend="RAINFUSION_ATTN", block_sparse=block_sparse) def test_block_sparse_end_step_precision_forwarded(self): spec = AttentionSpec( backend="RAINFUSION_ATTN", block_sparse={"sparsity": 0.9, "start_step": 5, "end_step": 3, "precision": "fp8"}, ) assert spec.backend_kwargs() == { "sparsity": 0.9, "start_step": 5, "end_step": 3, "precision": "fp8", } def test_block_sparse_precision_enum_input_normalized(self): """RainFusionPrecision enum members must normalize to their string value. Regression: ``str(str, Enum)`` later produced ``RainFusionPrecision.MIX`` instead of ``mix``. ``BlockSparseSpec`` stores ``member.value``. """ from vllm_omni.diffusion.attention.backends.rainfusion_attn import RainFusionConfig from vllm_omni.diffusion.data import RainFusionPrecision spec = AttentionSpec( backend="RAINFUSION_ATTN", block_sparse={"sparsity": 0.8, "precision": RainFusionPrecision.MIX}, ) kwargs = spec.backend_kwargs() assert kwargs["precision"] == "mix" # The config path is where str(value) used to mangle enum members. cfg = RainFusionConfig.from_backend_kwargs(kwargs) assert cfg.precision == "mix" assert type(cfg.precision) is str 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: @pytest.mark.parametrize("mindiesd_available", [False, True]) def test_paged_npu_memory_profile_uses_sdpa_only_without_mindiesd(self, monkeypatch, mindiesd_available): attention = Attention.__new__(Attention) attention.allow_fp32_fallback = False attention._scheduler_paged_kv = True attention.paged_kv_cache_role = "primary" attention.backend_pref = "FLASH_ATTN" attention.attn_backend = SimpleNamespace(supports_piecewise_spans=True, get_name=lambda: "FLASH_ATTN") flash_output = torch.ones(1) sdpa_output = torch.zeros(1) dense_flash = Mock(return_value=flash_output) sdpa = Mock(return_value=sdpa_output) attention.attention = SimpleNamespace(forward=dense_flash) attention.sdpa_fallback = SimpleNamespace(forward=sdpa) monkeypatch.setattr(layer_mod, "is_forward_context_available", lambda: True) monkeypatch.setattr( layer_mod, "get_forward_context", lambda: SimpleNamespace(in_diffusion_kv_memory_profile=True), ) monkeypatch.setattr(layer_mod.current_omni_platform, "is_npu", lambda: True) monkeypatch.setattr( layer_mod.current_omni_platform, "supports_diffusion_dense_flash_attention", lambda: mindiesd_available, ) result = attention._run_local_attention( torch.zeros(1, dtype=torch.bfloat16), torch.zeros(1, dtype=torch.bfloat16), torch.zeros(1, dtype=torch.bfloat16), AttentionMetadata(), ) if mindiesd_available: assert result is flash_output dense_flash.assert_called_once() sdpa.assert_not_called() else: assert result is sdpa_output dense_flash.assert_not_called() sdpa.assert_called_once() @pytest.mark.parametrize("fa4_available", [False, True]) def test_paged_cuda_memory_profile_uses_sdpa_without_dense_fa4(self, monkeypatch, fa4_available): attention = Attention.__new__(Attention) attention.allow_fp32_fallback = False attention._scheduler_paged_kv = True attention.paged_kv_cache_role = "primary" attention.backend_pref = "FLASH_ATTN" attention.attn_backend = SimpleNamespace(supports_piecewise_spans=True, get_name=lambda: "FLASH_ATTN") flash_output = torch.ones(1) sdpa_output = torch.zeros(1) dense_flash = Mock(return_value=flash_output) sdpa = Mock(return_value=sdpa_output) attention.attention = SimpleNamespace(forward=dense_flash) attention.sdpa_fallback = SimpleNamespace(forward=sdpa) monkeypatch.setattr(layer_mod, "is_forward_context_available", lambda: True) monkeypatch.setattr( layer_mod, "get_forward_context", lambda: SimpleNamespace(in_diffusion_kv_memory_profile=True), ) monkeypatch.setattr(layer_mod.current_omni_platform, "is_npu", lambda: False) monkeypatch.setattr(layer_mod.current_omni_platform, "is_cuda", lambda: True) monkeypatch.setattr( layer_mod.current_omni_platform, "supports_diffusion_dense_flash_attention", lambda: fa4_available, ) result = attention._run_local_attention( torch.zeros(1, dtype=torch.bfloat16), torch.zeros(1, dtype=torch.bfloat16), torch.zeros(1, dtype=torch.bfloat16), AttentionMetadata(), ) if fa4_available: assert result is flash_output dense_flash.assert_called_once() sdpa.assert_not_called() else: assert result is sdpa_output dense_flash.assert_not_called() sdpa.assert_called_once() def test_dense_flash_import_error_is_not_reclassified(self, monkeypatch): attention = Attention.__new__(Attention) attention.allow_fp32_fallback = False attention._scheduler_paged_kv = False attention.paged_kv_cache_role = "primary" attention.backend_pref = "FLASH_ATTN" attention.attn_backend = SimpleNamespace(supports_piecewise_spans=True, get_name=lambda: "FLASH_ATTN") attention.attention = SimpleNamespace( forward=Mock(side_effect=ModuleNotFoundError("No module named 'mindiesd'", name="mindiesd")) ) attention.sdpa_fallback = SimpleNamespace(forward=Mock(side_effect=AssertionError("must not fall back"))) monkeypatch.setattr(layer_mod, "is_forward_context_available", lambda: False) with pytest.raises(ModuleNotFoundError, match="mindiesd"): attention._run_local_attention( torch.zeros(1, dtype=torch.bfloat16), torch.zeros(1, dtype=torch.bfloat16), torch.zeros(1, dtype=torch.bfloat16), AttentionMetadata(), ) def test_paged_marker_does_not_change_dense_backend(self, monkeypatch): class _FakeAttentionImpl: def __init__(self, **kwargs): self.kwargs = kwargs def forward(self, query, key, value, attn_metadata=None): return query class _DenseBackend: supports_paged_kv = False @staticmethod def get_name() -> str: return "TORCH_SDPA" @staticmethod def get_impl_cls(): return _FakeAttentionImpl calls = [] def _fake_get_attn_backend_for_role(*args, **kwargs): calls.append((args, kwargs)) return _DenseBackend, None 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, "is_forward_context_available", lambda: False) od_config = SimpleNamespace( diffusion_attention_config=AttentionConfig(), diffusion_kv_mode=layer_mod.DiffusionKVCacheMode.DENSE_LEGACY, parallel_config=SimpleNamespace(ring_degree=1), 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): attention = Attention( num_heads=4, head_size=64, causal=False, softmax_scale=1.0, paged_kv_cache_role="primary", ) assert len(calls) == 1 assert attention.attn_backend is _DenseBackend # Auto path records the resolved backend name for Ring/SP; it is not explicit. assert attention.backend_pref == _DenseBackend.get_name() assert attention.backend_explicit is False def test_internal_paged_selector_does_not_require_dense_fa4(self, monkeypatch): """Blackwell CUDNN default + marked paged layer must not act as explicit FLASH_ATTN.""" class _FakeAttentionImpl: def __init__(self, **kwargs): self.kwargs = kwargs def forward(self, query, key, value, attn_metadata=None): return query class _CuDNNBackend: supports_paged_kv = False @staticmethod def get_name() -> str: return "CUDNN_ATTN" @staticmethod def get_impl_cls(): return _FakeAttentionImpl def _fake_get_attn_backend_for_role(*args, attention_config=None, **kwargs): if attention_config is not None: spec, _ = attention_config.resolve_with_source(role="self", role_category=None) if spec is not None and spec.backend.upper() == "FLASH_ATTN": raise ValueError( "FLASH_ATTN was explicitly selected but is unsupported " "(Blackwell requires CuTe FlashAttention-4 " "(flash_attn.cute / vllm-omni[fa4]); FA2/FA3 kernels are Hopper-only). " "Select a compatible backend." ) return _CuDNNBackend, None 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(FlashAttentionBackend, "get_impl_cls", staticmethod(lambda: _FakeAttentionImpl)) monkeypatch.setattr(layer_mod, "build_parallel_attention_strategy", lambda **kwargs: object()) monkeypatch.setattr(layer_mod, "is_forward_context_available", lambda: False) od_config = SimpleNamespace( diffusion_attention_config=AttentionConfig(), diffusion_kv_mode=layer_mod.DiffusionKVCacheMode.PAGED_SCHEDULER, parallel_config=SimpleNamespace(ring_degree=1), 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): attention = Attention( num_heads=4, head_size=64, causal=False, softmax_scale=1.0, paged_kv_cache_role="primary", ) assert attention.attn_backend is FlashAttentionBackend assert attention.backend_pref == "FLASH_ATTN" assert attention.backend_explicit is False assert attention.attn_spec is None def test_user_explicit_flash_attn_still_requires_dense_fa4(self, monkeypatch): def _fake_get_attn_backend_for_role(*args, attention_config=None, **kwargs): if attention_config is not None: spec, _ = attention_config.resolve_with_source(role="self", role_category=None) if spec is not None and spec.backend.upper() == "FLASH_ATTN": raise ValueError( "FLASH_ATTN was explicitly selected but is unsupported " "(Blackwell requires CuTe FlashAttention-4 " "(flash_attn.cute / vllm-omni[fa4]); FA2/FA3 kernels are Hopper-only). " "Select a compatible backend." ) raise AssertionError("explicit FLASH_ATTN must not fall through to the platform default") monkeypatch.setattr(layer_mod, "get_attn_backend_for_role", _fake_get_attn_backend_for_role) monkeypatch.setattr(layer_mod, "build_parallel_attention_strategy", lambda **kwargs: object()) monkeypatch.setattr(layer_mod, "is_forward_context_available", lambda: False) od_config = SimpleNamespace( diffusion_attention_config=AttentionConfig(default=AttentionSpec(backend="FLASH_ATTN")), diffusion_kv_mode=layer_mod.DiffusionKVCacheMode.PAGED_SCHEDULER, parallel_config=SimpleNamespace(ring_degree=1), 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), pytest.raises(ValueError, match="explicitly selected"): Attention( num_heads=4, head_size=64, causal=False, softmax_scale=1.0, paged_kv_cache_role="primary", ) 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, attn_backend_explicit=False): self.sp_group = sp_group self.attn_backend_pref = attn_backend_pref self.attn_backend_explicit = attn_backend_explicit 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" assert attn.ring_runner.attn_backend_explicit is True def test_attention_init_forwards_flashinfer_backend_without_dtype(self, monkeypatch): class _FakeAttentionImpl: def __init__(self, **kwargs): self.kwargs = kwargs class _FakeBackend: @staticmethod def get_name() -> str: return "FLASHINFER_ATTN" @staticmethod def get_impl_cls(): return _FakeAttentionImpl spec = AttentionSpec(backend="FLASHINFER_ATTN", quant={"flashinfer_backend": "fa2"}) monkeypatch.setattr(layer_mod, "get_attn_backend_for_role", lambda **kwargs: (_FakeBackend, spec)) 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: None) monkeypatch.setattr(layer_mod, "is_forward_context_available", lambda: False) od_config = SimpleNamespace( diffusion_attention_config=AttentionConfig(default=spec), parallel_config=SimpleNamespace(ring_degree=1), 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="self", ) assert attn.backend_pref == "FLASHINFER_ATTN" assert attn.attention.kwargs["backend_kwargs"] == {"quant": {"flashinfer_backend": "fa2"}} assert attn.attention.kwargs["backend_explicit"] is True def test_attention_init_propagates_ring_setup_failure(self, monkeypatch): class _FakeAttentionImpl: def __init__(self, **kwargs): pass class _FakeBackend: @staticmethod def get_name() -> str: return "TORCH_SDPA" @staticmethod def get_impl_cls(): return _FakeAttentionImpl monkeypatch.setattr( layer_mod, "get_attn_backend_for_role", lambda **kwargs: (_FakeBackend, AttentionSpec(backend="TORCH_SDPA")), ) monkeypatch.setattr(layer_mod.SDPABackend, "get_impl_cls", staticmethod(lambda: _FakeAttentionImpl)) monkeypatch.setattr( layer_mod, "get_sp_group", lambda: (_ for _ in ()).throw(RuntimeError("SP group is not initialized")), ) od_config = SimpleNamespace( diffusion_attention_config=AttentionConfig(default="TORCH_SDPA"), parallel_config=SimpleNamespace(ring_degree=2), ) with set_current_diffusion_config(od_config): with pytest.raises(RuntimeError, match="SP group is not initialized"): Attention(num_heads=4, head_size=64, causal=False, softmax_scale=1.0) def test_float32_main_dispatch_does_not_override_selected_backend(self): sentinel = torch.randn(1, 2, 4, 8) selected_calls = [] def _selected_forward(*args): selected_calls.append(args) return sentinel fake_attention = SimpleNamespace( allow_fp32_fallback=False, attention=SimpleNamespace( forward=_selected_forward, ), sdpa_fallback=SimpleNamespace( forward=lambda *args: pytest.fail("unexpected SDPA fallback"), ), _assert_metadata_compatible=lambda metadata: None, _has_custom_attention=False, _scheduler_paged_kv=False, paged_kv_cache_role=None, ) query = torch.randn(1, 2, 4, 8, dtype=torch.float32) output = Attention._run_local_attention(fake_attention, query, query, query, None) assert output is sentinel assert len(selected_calls) == 1 def test_unsupported_attention_mask_is_rejected_before_backend_call(self): fake_backend = SimpleNamespace( get_name=lambda: "NO_MASK_BACKEND", supports_attention_mask=lambda *args, **kwargs: False, ) fake_attention = SimpleNamespace(attn_backend=fake_backend) metadata = AttentionMetadata(attn_mask=torch.ones(1, 2, dtype=torch.bool)) with pytest.raises(ValueError, match="does not support attn_mask"): Attention._assert_metadata_compatible(fake_attention, metadata) def test_ring_attention_rejects_mask_instead_of_ignoring_it(self): fake_attention = SimpleNamespace(attention=SimpleNamespace(skip=None), ring_runner=None) query = torch.randn(1, 2, 4, 8) metadata = AttentionMetadata(attn_mask=torch.ones(1, 2, dtype=torch.bool)) with pytest.raises(ValueError, match="Ring attention does not support attn_mask"): Attention._run_ring_attention(fake_attention, query, query, query, metadata) class TestOptInFloat32Fallback: @staticmethod def _make_attention(monkeypatch, *, backend=None, allow=True, kv_cache_dtype=None): events, state = [], {} context = object() class _SelectedImpl: def __init__(self, **kwargs): self.kwargs = kwargs @staticmethod def supports_kv_cache_dtype(*args): return True def forward(self, query, key, value, metadata): events.append("selected") state["kernel_call"] = (query, key, value, metadata) state["kernel_output"] = query + 4 return state["kernel_output"] class _FallbackImpl(_SelectedImpl): def forward(self, query, key, value, metadata): output = super().forward(query, key, value, metadata) events[-1] = "fallback" return output class _FlashBackend(FlashAttentionBackend): @staticmethod def get_impl_cls(): return _SelectedImpl class _SDPABackend(layer_mod.SDPABackend): @staticmethod def get_impl_cls(): return _SelectedImpl class _Strategy: def pre_attention(self, query, key, value, metadata): events.append("pre") state["prepared"] = ( query + 1, key + 2, value + 3, replace(metadata) if metadata is not None else None, ) return (*state["prepared"], context) def post_attention(self, output, ctx): events.append("post") assert output is state["kernel_output"] assert ctx is context return output + 7 # Keep real role/config resolution and construction; replace hardware # lookup and kernels. These are CPU dispatch tests, not SP collectives. monkeypatch.setattr( selector_mod, "_cached_get_backend_cls", lambda name, *args, **kwargs: _SDPABackend if name == "TORCH_SDPA" else _FlashBackend, ) monkeypatch.setattr(layer_mod.SDPABackend, "get_impl_cls", staticmethod(lambda: _FallbackImpl)) monkeypatch.setattr(layer_mod, "build_parallel_attention_strategy", lambda **kwargs: _Strategy()) monkeypatch.setattr(layer_mod, "is_forward_context_available", lambda: False) config = OmniDiffusionConfig( diffusion_attention_config={"default": backend} if backend else {}, diffusion_kv_cache_dtype=kv_cache_dtype, ) with set_current_diffusion_config(config): attention = Attention( num_heads=4, num_kv_heads=2, head_size=8, causal=False, softmax_scale=0.25, **({"allow_fp32_fallback": True} if allow else {}), ) return attention, events, state @pytest.mark.parametrize( "backend,allow,dtype,is_cuda,expected", [ pytest.param(None, True, torch.float32, True, "fallback", id="auto-fp32"), pytest.param(None, False, torch.float32, True, "selected", id="default-opt-out"), pytest.param("FLASH_ATTN", True, torch.float32, True, "selected", id="explicit-flash"), pytest.param("TORCH_SDPA", True, torch.float32, True, "selected", id="explicit-sdpa"), pytest.param(None, True, torch.bfloat16, True, "selected", id="bf16"), pytest.param(None, True, torch.float32, False, "selected", id="cpu-fp32"), ], ) def test_public_forward_preserves_strategy_and_backend_choice( self, monkeypatch, backend, allow, dtype, is_cuda, expected ): attention, events, state = self._make_attention(monkeypatch, backend=backend, allow=allow) # Only the dispatch predicate is simulated; every tensor stays on CPU. monkeypatch.setattr(torch.Tensor, "is_cuda", property(lambda self: is_cuda)) query = torch.ones(1, 2, 4, 8, dtype=dtype) key = torch.full((1, 2, 2, 8), 2, dtype=dtype) value = torch.full_like(key, 3) metadata = AttentionMetadata(attn_mask=torch.tensor([[True, False]])) output = attention(query, key, value, metadata) assert events == ["pre", expected, "post"] assert all(got is prepared for got, prepared in zip(state["kernel_call"], state["prepared"])) assert state["kernel_call"][3] is not metadata assert state["kernel_call"][3].attn_mask is metadata.attn_mask torch.testing.assert_close(output, query + 12) assert attention.backend_explicit is (backend is not None) impl = attention.sdpa_fallback if expected == "fallback" else attention.attention assert impl.kwargs["num_kv_heads"] == 2 assert impl.kwargs["softmax_scale"] == 0.25 def test_explicit_backend_failure_is_not_retried_with_sdpa(self, monkeypatch): attention, events, _ = self._make_attention(monkeypatch, backend="FLASH_ATTN") monkeypatch.setattr(torch.Tensor, "is_cuda", property(lambda self: True)) selected = Mock(side_effect=RuntimeError("selected backend rejects this dtype")) monkeypatch.setattr(attention.attention, "forward", selected) query = torch.ones(1, 2, 4, 8) with pytest.raises(RuntimeError, match="selected backend rejects this dtype"): attention(query, query, query) selected.assert_called_once() assert events == ["pre"] @pytest.mark.parametrize("contract", ["packed", "packed-padding", "piecewise", "quantized-kv"]) def test_backend_metadata_is_not_silently_sent_to_dense_sdpa(self, monkeypatch, contract): quant_dtype = "fp8" if contract == "quantized-kv" else None attention, events, state = self._make_attention(monkeypatch, kv_cache_dtype=quant_dtype) monkeypatch.setattr(torch.Tensor, "is_cuda", property(lambda self: True)) lengths = torch.tensor([0, 1, 2], dtype=torch.int32) metadata = AttentionMetadata() if contract == "packed": metadata.extra = { "cu_seqlens_q": lengths, "cu_seqlens_k": lengths, "max_seqlen_q": 1, "max_seqlen_k": 1, } elif contract == "packed-padding": metadata.packed_padding = PackedPaddingMetadata(1, 1, lengths[:2], lengths[:2]) elif contract == "piecewise": metadata.full_attn_spans = [[(0, 1)]] metadata.query_ranges = (QueryRange(0, 2, 0),) query = torch.ones(1, 2, 4, 8) attention(query, query, query, metadata) assert events == ["pre", "selected", "post"] forwarded = state["kernel_call"][3] if quant_dtype: # Quantization is injected by the public forward from config. assert forwarded.extra == {"kv_cache_dtype": "fp8"} assert metadata.extra == {} else: assert forwarded is state["prepared"][3] class TestDiffusionKvCacheQuantization: @staticmethod def _install_attention_init_stubs(monkeypatch, *, supports_kv_cache_dtype=True): class _FakeAttentionImpl: def __init__(self, **kwargs): self.kwargs = kwargs def forward(self, query, key, value, attn_metadata=None): return query @staticmethod def supports_kv_cache_dtype(kv_cache_dtype, platform_key) -> bool: return supports_kv_cache_dtype class _FakeBackend: @staticmethod def get_name() -> str: return "FAKE_BACKEND" @staticmethod def get_impl_cls(): return _FakeAttentionImpl class _FakeRingParallelAttention: def __init__(self, sp_group, attn_backend_pref=None, attn_backend_explicit=False): self.sp_group = sp_group self.attn_backend_pref = attn_backend_pref self.attn_backend_explicit = attn_backend_explicit 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, ) def test_explicit_unsupported_kv_cache_dtype_raises(self, monkeypatch): self._install_attention_init_stubs(monkeypatch, supports_kv_cache_dtype=False) od_config = SimpleNamespace( diffusion_attention_config=AttentionConfig(), parallel_config=SimpleNamespace(ring_degree=1), 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="does not support kv_cache_dtype='fp8'"): Attention( num_heads=4, head_size=64, causal=False, softmax_scale=1.0, )