diff --git a/examples/offline_inference/text_to_video/wan2_2_sparse.py b/examples/offline_inference/text_to_video/wan2_2_sparse.py new file mode 100644 index 00000000000..c2c030daf69 --- /dev/null +++ b/examples/offline_inference/text_to_video/wan2_2_sparse.py @@ -0,0 +1,195 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Wan 2.2 video generation with sparse attention acceleration. + +Sparse attention reduces computation in DiT self-attention layers by +attending only to a subset of tokens via a plugin function (e.g. SpargeAttn). +The SparseAttentionBackend dispatcher resolves the plugin at init time and +handles cross-attention fallback, tensor layout, and config loading. + +Usage: + # SpargeAttn with default topk (0.5) + python wan2_2_sparse.py --attn-backend spargeattn + + # SpargeAttn with custom topk via JSON config (flat form) + python wan2_2_sparse.py --attn-backend '{"backend":"spargeattn","topk":0.3}' + + # Auto-detect installed plugin + python wan2_2_sparse.py --attn-backend auto + + # Dense baseline (no sparse attention) + python wan2_2_sparse.py --attn-backend dense + + # With sequence parallelism + python wan2_2_sparse.py --attn-backend spargeattn --ring-degree 2 +""" + +import argparse +import time +from pathlib import Path + +import numpy as np +import torch + +from vllm_omni.diffusion.data import DiffusionParallelConfig, DiffusionSparseAttnConfig +from vllm_omni.entrypoints.omni import Omni +from vllm_omni.inputs.data import OmniDiffusionSamplingParams +from vllm_omni.outputs import OmniRequestOutput +from vllm_omni.platforms import current_omni_platform + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Wan 2.2 video generation with sparse attention.") + parser.add_argument( + "--model", + default="Wan-AI/Wan2.2-T2V-A14B-Diffusers", + help="Model ID or local path.", + ) + parser.add_argument( + "--prompt", + default="A serene lakeside sunrise with mist over the water.", + help="Text prompt.", + ) + parser.add_argument("--negative-prompt", default="", help="Negative prompt.") + parser.add_argument("--seed", type=int, default=42, help="Random seed.") + parser.add_argument("--height", type=int, default=720, help="Video height.") + parser.add_argument("--width", type=int, default=1280, help="Video width.") + parser.add_argument("--num-frames", type=int, default=81, help="Number of frames.") + parser.add_argument("--num-inference-steps", type=int, default=40, help="Sampling steps.") + parser.add_argument("--guidance-scale", type=float, default=4.0, help="CFG scale.") + parser.add_argument("--fps", type=int, default=24, help="Output video FPS.") + parser.add_argument("--output", type=str, default="wan22_sparse_output.mp4", help="Output path.") + + # Attention backend: simple name or JSON with config + parser.add_argument( + "--attn-backend", + type=str, + default=None, + help="Sparse attention plugin. " + "Simple name: spargeattn, auto, dense. " + 'JSON: \'{"backend":"spargeattn","topk":0.3}\'. ' + "Import path: 'module.path:func_name'.", + ) + + # Parallelism + parser.add_argument("--ulysses-degree", type=int, default=1, help="Ulysses SP degree.") + parser.add_argument("--ring-degree", type=int, default=1, help="Ring SP degree.") + parser.add_argument("--tensor-parallel-size", type=int, default=1, help="TP size.") + parser.add_argument("--cfg-parallel-size", type=int, default=1, choices=[1, 2], help="CFG parallel size.") + + # Other + parser.add_argument("--flow-shift", type=float, default=None, help="Scheduler flow_shift.") + parser.add_argument("--enforce-eager", action="store_true", help="Disable torch.compile.") + return parser.parse_args() + + +def main(): + args = parse_args() + + generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed) + + # Build sparse attention config from --attn-backend + attn_backend_name = args.attn_backend + sparse_attn = None + if attn_backend_name: + if attn_backend_name.strip().startswith("{"): + import json + + sparse_attn = DiffusionSparseAttnConfig.from_dict(json.loads(attn_backend_name)) + else: + sparse_attn = DiffusionSparseAttnConfig(backend=attn_backend_name) + + parallel_config = DiffusionParallelConfig( + ulysses_degree=args.ulysses_degree, + ring_degree=args.ring_degree, + cfg_parallel_size=args.cfg_parallel_size, + tensor_parallel_size=args.tensor_parallel_size, + ) + + omni_kwargs = dict( + model=args.model, + parallel_config=parallel_config, + enforce_eager=args.enforce_eager, + ) + if sparse_attn is not None: + omni_kwargs["sparse_attn"] = sparse_attn + if args.flow_shift is not None: + omni_kwargs["flow_shift"] = args.flow_shift + + backend_str = sparse_attn.backend if sparse_attn else "default" + params_str = str(sparse_attn.params) if sparse_attn else "N/A" + print(f"\n{'=' * 60}") + print("Wan 2.2 Attention Configuration:") + print(f" Backend: {backend_str}") + print(f" Params: {params_str}") + print(f" Video: {args.width}x{args.height}, {args.num_frames} frames") + print(f" Steps: {args.num_inference_steps}") + print(f"{'=' * 60}\n") + + omni = Omni(**omni_kwargs) + + prompt_dict = {"prompt": args.prompt} + if args.negative_prompt: + prompt_dict["negative_prompt"] = args.negative_prompt + + generation_start = time.perf_counter() + result = omni.generate( + prompt_dict, + OmniDiffusionSamplingParams( + height=args.height, + width=args.width, + generator=generator, + guidance_scale=args.guidance_scale, + num_inference_steps=args.num_inference_steps, + num_frames=args.num_frames, + ), + ) + generation_time = time.perf_counter() - generation_start + + print(f"\nGeneration time: {generation_time:.2f}s ({generation_time * 1000:.0f}ms)") + + # Extract frames from result + # Pipeline returns list[OmniRequestOutput]; each .images is + # [ndarray(batch, frames, H, W, C)] with values in [0, 1]. + frames = result + if isinstance(frames, list): + frames = frames[0] if frames else None + if isinstance(frames, OmniRequestOutput) and frames.images: + frames = frames.images + + if frames is None: + raise ValueError("No video frames in output.") + + # Unwrap: list[ndarray] -> ndarray + if isinstance(frames, list) and len(frames) == 1: + frames = frames[0] + + # Save video + output_path = Path(args.output) + output_path.parent.mkdir(parents=True, exist_ok=True) + try: + from diffusers.utils import export_to_video + except ImportError: + raise ImportError("diffusers is required for export_to_video.") + + if isinstance(frames, torch.Tensor): + frames = frames.detach().cpu().float().numpy() + + if isinstance(frames, np.ndarray): + # Remove batch dim: (B, F, H, W, C) -> (F, H, W, C) + while frames.ndim > 4: + frames = frames[0] + # export_to_video expects float [0,1] arrays — it handles uint8 conversion + if frames.dtype in (np.float32, np.float64): + frames = np.clip(frames, 0, 1) + elif np.issubdtype(frames.dtype, np.integer): + frames = frames.astype(np.float32) / 255.0 + frames = list(frames) # list of (H, W, C) arrays + + export_to_video(frames, str(output_path), fps=args.fps) + print(f"Saved to {output_path}") + + +if __name__ == "__main__": + main() diff --git a/tests/diffusion/sparse_attn/__init__.py b/tests/diffusion/sparse_attn/__init__.py new file mode 100644 index 00000000000..208f01a7cb5 --- /dev/null +++ b/tests/diffusion/sparse_attn/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/tests/diffusion/sparse_attn/test_config_cli.py b/tests/diffusion/sparse_attn/test_config_cli.py new file mode 100644 index 00000000000..ae4abe2144e --- /dev/null +++ b/tests/diffusion/sparse_attn/test_config_cli.py @@ -0,0 +1,64 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for sparse attention config wiring. + +Tests: +- OmniDiffusionConfig.from_kwargs with sparse_attn dict +- DiffusionSparseAttnConfig construction with params dict +- sparse_attn_backend field removed +""" + +from __future__ import annotations + +import pytest + +from vllm_omni.diffusion.data import DiffusionSparseAttnConfig, OmniDiffusionConfig + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class TestFromKwargsWithSparseAttn: + """Test OmniDiffusionConfig.from_kwargs with sparse attn config.""" + + def test_sparse_attn_dict_passthrough(self): + config = OmniDiffusionConfig.from_kwargs( + sparse_attn={"backend": "flashinfer", "topk": 0.4}, + ) + assert isinstance(config.sparse_attn, DiffusionSparseAttnConfig) + assert config.sparse_attn.backend == "flashinfer" + assert config.sparse_attn.params == {"topk": 0.4} + + def test_no_sparse_attn(self): + config = OmniDiffusionConfig.from_kwargs() + assert config.sparse_attn is None + + def test_sparse_attn_backend_field_removed(self): + """sparse_attn_backend field no longer exists on OmniDiffusionConfig.""" + assert not hasattr(OmniDiffusionConfig, "sparse_attn_backend") + config = OmniDiffusionConfig() + assert not hasattr(config, "sparse_attn_backend") + + def test_sparse_attn_config_from_dict(self): + config = OmniDiffusionConfig( + sparse_attn={"backend": "dense", "topk": 0.3}, + ) + assert isinstance(config.sparse_attn, DiffusionSparseAttnConfig) + assert config.sparse_attn.backend == "dense" + assert config.sparse_attn.params == {"topk": 0.3} + + +class TestEndToEndConfigFlow: + """Test the full config flow for sparse attention.""" + + def test_sparse_attn_via_dict(self): + """Simulate: --sparse-attn '{"backend":"spargeattn","topk":0.5}'.""" + sparse_config = {"backend": "spargeattn", "topk": 0.5} + config = OmniDiffusionConfig.from_kwargs(sparse_attn=sparse_config) + + assert isinstance(config.sparse_attn, DiffusionSparseAttnConfig) + assert config.sparse_attn.backend == "spargeattn" + assert config.sparse_attn.params == {"topk": 0.5} + + def test_no_sparse_args_preserves_none(self): + config = OmniDiffusionConfig.from_kwargs() + assert config.sparse_attn is None diff --git a/tests/diffusion/sparse_attn/test_protocol.py b/tests/diffusion/sparse_attn/test_protocol.py new file mode 100644 index 00000000000..3063174c9cb --- /dev/null +++ b/tests/diffusion/sparse_attn/test_protocol.py @@ -0,0 +1,89 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for sparse attention config parsing. + +Tests: +- DiffusionSparseAttnConfig dataclass and from_dict factory +- OmniDiffusionConfig integration with sparse_attn +""" + +import pytest + +from vllm_omni.diffusion.data import DiffusionSparseAttnConfig, OmniDiffusionConfig + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class TestDiffusionSparseAttnConfig: + """Test DiffusionSparseAttnConfig dataclass.""" + + def test_defaults(self): + config = DiffusionSparseAttnConfig() + assert config.backend == "auto" + assert config.params == {} + + def test_from_dict_flat(self): + """Flat JSON: unknown keys auto-routed into params.""" + config = DiffusionSparseAttnConfig.from_dict( + { + "backend": "flashinfer", + "topk": 0.3, + "block_size": 128, + } + ) + assert config.backend == "flashinfer" + assert config.params == {"topk": 0.3, "block_size": 128} + + def test_from_dict_nested(self): + """Nested JSON: explicit params key.""" + config = DiffusionSparseAttnConfig.from_dict( + { + "backend": "spargeattn", + "params": {"topk": 0.5}, + } + ) + assert config.backend == "spargeattn" + assert config.params == {"topk": 0.5} + + def test_from_dict_ignores_backend_in_params(self): + config = DiffusionSparseAttnConfig.from_dict( + { + "backend": "dense", + "extra_key": 42, + } + ) + assert config.backend == "dense" + assert config.params == {"extra_key": 42} + + def test_from_dict_empty(self): + config = DiffusionSparseAttnConfig.from_dict({}) + assert config.backend == "auto" + assert config.params == {} + + +class TestOmniDiffusionConfigSparseAttn: + """Test sparse_attn integration in OmniDiffusionConfig.""" + + def test_sparse_attn_none_by_default(self): + config = OmniDiffusionConfig() + assert config.sparse_attn is None + + def test_sparse_attn_from_dict(self): + config = OmniDiffusionConfig(sparse_attn={"backend": "flashinfer", "topk": 0.4}) + assert isinstance(config.sparse_attn, DiffusionSparseAttnConfig) + assert config.sparse_attn.backend == "flashinfer" + assert config.sparse_attn.params == {"topk": 0.4} + + +class TestOmniDiffusionConfigFromKwargs: + """Test OmniDiffusionConfig.from_kwargs with sparse_attn fields.""" + + def test_from_kwargs_with_sparse_attn_dict(self): + config = OmniDiffusionConfig.from_kwargs(sparse_attn={"backend": "flashinfer", "topk": 0.3}) + assert isinstance(config.sparse_attn, DiffusionSparseAttnConfig) + assert config.sparse_attn.backend == "flashinfer" + assert config.sparse_attn.params == {"topk": 0.3} + + def test_from_kwargs_no_sparse_attn(self): + config = OmniDiffusionConfig.from_kwargs() + assert config.sparse_attn is None diff --git a/tests/diffusion/sparse_attn/test_sparse_attention.py b/tests/diffusion/sparse_attn/test_sparse_attention.py new file mode 100644 index 00000000000..03bd15cc438 --- /dev/null +++ b/tests/diffusion/sparse_attn/test_sparse_attention.py @@ -0,0 +1,254 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for SparseAttentionBackend and the function-based plugin system. + +Tests: +- SparseAttentionBackend inherits AttentionBackend +- _SparseAttentionImpl cross-attention path (is_self_attention=False -> dense) +- _SparseAttentionImpl self-attention with no sparse config -> dense +- Function-based plugin resolution (entry points, import path) +- SPARSE_ATTENTION enum registration +- is_self_attention parameter in Attention layer +""" + +from __future__ import annotations + +from unittest.mock import MagicMock, patch + +import pytest + +from vllm_omni.diffusion.attention.backends.abstract import AttentionBackend +from vllm_omni.diffusion.attention.backends.sparse_attention import ( + SparseAttentionBackend, + _SparseAttentionImpl, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class TestSparseAttentionBackendInheritance: + """Test SparseAttentionBackend implements AttentionBackend.""" + + def test_inherits_attention_backend(self): + assert issubclass(SparseAttentionBackend, AttentionBackend) + + def test_get_name(self): + assert SparseAttentionBackend.get_name() == "sparse_attention" + + def test_get_impl_cls(self): + cls = SparseAttentionBackend.get_impl_cls() + assert cls is _SparseAttentionImpl + + def test_get_metadata_cls(self): + from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata + + assert SparseAttentionBackend.get_metadata_cls() is AttentionMetadata + + def test_get_builder_cls_none(self): + assert SparseAttentionBackend.get_builder_cls() is None + + def test_get_supported_head_sizes(self): + assert SparseAttentionBackend.get_supported_head_sizes() == [] + + +class TestSparseAttentionImplCrossAttention: + """Test _SparseAttentionImpl with is_self_attention=False always uses dense.""" + + def test_cross_attn_uses_dense(self): + impl = _SparseAttentionImpl( + num_heads=8, + head_size=64, + softmax_scale=0.125, + causal=False, + is_self_attention=False, + ) + assert impl._is_self_attention is False + assert impl._sparse_fn is None + assert impl._dense is not None + + def test_self_attn_default_true(self): + """Default is_self_attention=True.""" + impl = _SparseAttentionImpl( + num_heads=8, + head_size=64, + softmax_scale=0.125, + causal=False, + ) + assert impl._is_self_attention is True + + def test_self_attn_no_config_falls_back_to_dense(self): + """When no sparse config is available, self-attention also uses dense.""" + impl = _SparseAttentionImpl( + num_heads=8, + head_size=64, + softmax_scale=0.125, + causal=False, + is_self_attention=True, + ) + # Without forward context, _read_sparse_cfg returns None -> no plugin loaded + assert impl._sparse_fn is None + + +class TestEnumRegistration: + """Test SPARSE_ATTENTION is in the registry enum.""" + + def test_sparse_attention_in_enum(self): + from vllm_omni.diffusion.attention.backends.registry import DiffusionAttentionBackendEnum + + assert hasattr(DiffusionAttentionBackendEnum, "SPARSE_ATTENTION") + + def test_sparse_attention_class_path(self): + from vllm_omni.diffusion.attention.backends.registry import DiffusionAttentionBackendEnum + + backend = DiffusionAttentionBackendEnum.SPARSE_ATTENTION + assert "sparse_attention.SparseAttentionBackend" in backend.value + + +class TestEntryPointsPluginLookup: + """Test _get_plugin_backend_cls in selector.py.""" + + def test_no_plugin_returns_none(self): + from vllm_omni.diffusion.attention.selector import _get_plugin_backend_cls + + result = _get_plugin_backend_cls("nonexistent_backend_xyz") + assert result is None + + def test_plugin_found_via_entry_point(self): + """Mock an entry_point and verify it's loaded.""" + from vllm_omni.diffusion.attention.selector import _get_plugin_backend_cls + + mock_ep = MagicMock() + mock_ep.name = "testplugin" + mock_ep.value = "test.module:TestBackend" + mock_backend_cls = type("TestBackend", (AttentionBackend,), {}) + mock_ep.load.return_value = mock_backend_cls + + with patch("importlib.metadata.entry_points", return_value=[mock_ep]): + result = _get_plugin_backend_cls("testplugin") + + assert result is mock_backend_cls + + def test_plugin_case_insensitive(self): + """Plugin lookup should be case-insensitive.""" + from vllm_omni.diffusion.attention.selector import _get_plugin_backend_cls + + mock_ep = MagicMock() + mock_ep.name = "SpargeAttn" + mock_ep.value = "test.module:SpargeBackend" + mock_backend_cls = type("SpargeBackend", (AttentionBackend,), {}) + mock_ep.load.return_value = mock_backend_cls + + with patch("importlib.metadata.entry_points", return_value=[mock_ep]): + result = _get_plugin_backend_cls("spargeattn") + + assert result is mock_backend_cls + + +class TestIsSelfattentionInAttentionLayer: + """Test that Attention layer passes is_self_attention to impl.""" + + def _mock_distributed(self): + from contextlib import ExitStack + + stack = ExitStack() + + mock_tp = MagicMock() + mock_tp.world_size = 1 + mock_tp.rank_in_group = 0 + mock_tp.local_rank = 0 + import vllm.distributed.parallel_state as ps + + stack.enter_context(patch.object(ps, "_TP", mock_tp)) + + mock_config = MagicMock() + mock_config.attention_backend = None + mock_config.parallel_config.ring_degree = 1 + mock_ctx = stack.enter_context(patch("vllm_omni.diffusion.attention.layer.get_forward_context")) + mock_ctx.return_value.omni_diffusion_config = mock_config + stack.enter_context( + patch( + "vllm_omni.diffusion.attention.layer.is_forward_context_available", + return_value=True, + ) + ) + return stack + + def test_default_is_self_attention_true(self): + """By default, Attention sets is_self_attention=True.""" + with self._mock_distributed(): + from vllm_omni.diffusion.attention.layer import Attention + + attn = Attention( + num_heads=8, + head_size=64, + causal=False, + softmax_scale=0.125, + ) + # The impl was constructed — no crash means it accepted the kwarg + assert attn.attention is not None + + def test_is_self_attention_false_accepted(self): + """Attention(is_self_attention=False) works.""" + with self._mock_distributed(): + from vllm_omni.diffusion.attention.layer import Attention + + attn = Attention( + num_heads=8, + head_size=64, + causal=False, + softmax_scale=0.125, + is_self_attention=False, + ) + assert attn.attention is not None + + +class TestSparseAttentionImplResolvePlugin: + """Test _SparseAttentionImpl._resolve_plugin with function-based plugins.""" + + def test_resolve_auto_no_plugins(self): + """Auto backend returns None when no plugins installed.""" + from vllm_omni.diffusion.data import DiffusionSparseAttnConfig + + cfg = DiffusionSparseAttnConfig(backend="auto") + with patch("importlib.metadata.entry_points", return_value=[]): + result = _SparseAttentionImpl._resolve_plugin(cfg) + assert result is None + + def test_resolve_import_path(self): + """Full import path 'module:func' resolves to a callable.""" + from vllm_omni.diffusion.data import DiffusionSparseAttnConfig + + # Use a known function as test target + cfg = DiffusionSparseAttnConfig(backend="torch.nn.functional:scaled_dot_product_attention") + result = _SparseAttentionImpl._resolve_plugin(cfg) + import torch.nn.functional as F + + assert result is F.scaled_dot_product_attention + + def test_resolve_short_name_via_entry_point(self): + """Short name resolved via vllm_omni.sparse_attn entry point group.""" + from vllm_omni.diffusion.data import DiffusionSparseAttnConfig + + mock_fn = MagicMock() + mock_ep = MagicMock() + mock_ep.name = "testbackend" + mock_ep.load.return_value = mock_fn + + cfg = DiffusionSparseAttnConfig(backend="testbackend") + with patch("importlib.metadata.entry_points", return_value=[mock_ep]): + result = _SparseAttentionImpl._resolve_plugin(cfg) + + assert result is mock_fn + + def test_resolve_auto_picks_first(self): + """Auto backend picks the first available plugin function.""" + mock_fn = MagicMock() + mock_ep = MagicMock() + mock_ep.name = "myplugin" + mock_ep.load.return_value = mock_fn + + result_fn = None + with patch("importlib.metadata.entry_points", return_value=[mock_ep]): + result_fn = _SparseAttentionImpl._first_available_plugin() + + assert result_fn is mock_fn diff --git a/tests/diffusion/sparse_attn/test_wan22_gpu_integration.py b/tests/diffusion/sparse_attn/test_wan22_gpu_integration.py new file mode 100644 index 00000000000..a67b87e5cd7 --- /dev/null +++ b/tests/diffusion/sparse_attn/test_wan22_gpu_integration.py @@ -0,0 +1,87 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""GPU integration tests for SparseAttentionBackend with Wan 2.2. + +Tests the production path: SparseAttentionBackend -> _SparseAttentionImpl -> plugin fn. +Skipped when GPU is not available. +""" + +from __future__ import annotations + +from unittest.mock import patch + +import pytest +import torch + +from vllm_omni.diffusion.data import DiffusionSparseAttnConfig + +_has_cuda = torch.cuda.is_available() + +requires_gpu = pytest.mark.skipif(not _has_cuda, reason="Requires CUDA GPU") + + +@requires_gpu +class TestSparseAttentionImplGPU: + """Test _SparseAttentionImpl forward on GPU with plugin detection.""" + + def _make_impl(self, sparse_cfg=None, is_self_attention=True): + """Create a _SparseAttentionImpl with mocked forward context.""" + from vllm_omni.diffusion.attention.backends.sparse_attention import _SparseAttentionImpl + + with patch( + "vllm_omni.diffusion.attention.backends.sparse_attention._SparseAttentionImpl._read_sparse_cfg", + return_value=sparse_cfg, + ): + return _SparseAttentionImpl( + num_heads=4, + head_size=64, + softmax_scale=64**-0.5, + causal=False, + num_kv_heads=4, + is_self_attention=is_self_attention, + ) + + def test_dense_fallback_forward(self): + """No sparse config -> SDPA dense forward works on GPU.""" + impl = self._make_impl(sparse_cfg=None) + q = torch.randn(1, 128, 4, 64, dtype=torch.bfloat16, device="cuda") + k = torch.randn(1, 128, 4, 64, dtype=torch.bfloat16, device="cuda") + v = torch.randn(1, 128, 4, 64, dtype=torch.bfloat16, device="cuda") + out = impl.forward(q, k, v) + assert out.shape == q.shape + assert not torch.isnan(out).any() + + def test_cross_attention_always_dense(self): + """Cross-attention (is_self_attention=False) ignores sparse config.""" + cfg = DiffusionSparseAttnConfig(backend="auto", params={"topk": 0.3}) + impl = self._make_impl(sparse_cfg=cfg, is_self_attention=False) + q = torch.randn(1, 64, 4, 64, dtype=torch.bfloat16, device="cuda") + k = torch.randn(1, 64, 4, 64, dtype=torch.bfloat16, device="cuda") + v = torch.randn(1, 64, 4, 64, dtype=torch.bfloat16, device="cuda") + out = impl.forward(q, k, v) + assert out.shape == q.shape + + def test_mock_sparse_fn_forward(self): + """Self-attention with a mock sparse function.""" + + # Create impl with no sparse config (dense) + impl = self._make_impl(sparse_cfg=None, is_self_attention=True) + + # Inject a mock sparse function that uses SDPA + def fake_sparse_fn(q, k, v, params, is_causal): + return torch.nn.functional.scaled_dot_product_attention( + q, + k, + v, + is_causal=is_causal, + ) + + impl._sparse_fn = fake_sparse_fn + impl._params = {"topk": 0.5} + + q = torch.randn(1, 128, 4, 64, dtype=torch.bfloat16, device="cuda") + k = torch.randn(1, 128, 4, 64, dtype=torch.bfloat16, device="cuda") + v = torch.randn(1, 128, 4, 64, dtype=torch.bfloat16, device="cuda") + out = impl.forward(q, k, v) + assert out.shape == q.shape + assert not torch.isnan(out).any() diff --git a/tests/diffusion/sparse_attn/test_wan22_sparse.py b/tests/diffusion/sparse_attn/test_wan22_sparse.py new file mode 100644 index 00000000000..b247a34845b --- /dev/null +++ b/tests/diffusion/sparse_attn/test_wan22_sparse.py @@ -0,0 +1,117 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for Wan 2.2 sparse attention integration after refactor. + +Tests: +- WanSelfAttention no longer has sparse_attn attribute +- WanCrossAttention passes is_self_attention=False +- Sparse attention is now handled via the attention backend layer +""" + +from __future__ import annotations + +from unittest.mock import MagicMock, patch + +import pytest + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _mock_distributed(): + """Context manager to mock distributed and forward context for unit tests.""" + from contextlib import ExitStack + + stack = ExitStack() + + mock_tp = MagicMock() + mock_tp.world_size = 1 + mock_tp.rank_in_group = 0 + mock_tp.local_rank = 0 + import vllm.distributed.parallel_state as ps + + stack.enter_context(patch.object(ps, "_TP", mock_tp)) + + mock_config = MagicMock() + mock_config.attention_backend = None + mock_config.parallel_config.ring_degree = 1 + mock_ctx = stack.enter_context(patch("vllm_omni.diffusion.attention.layer.get_forward_context")) + mock_ctx.return_value.omni_diffusion_config = mock_config + stack.enter_context( + patch( + "vllm_omni.diffusion.attention.layer.is_forward_context_available", + return_value=True, + ) + ) + + return stack + + +class TestWanSelfAttentionNoSparseAttr: + """Test that WanSelfAttention no longer has sparse_attn attribute.""" + + def _make_self_attn(self): + with _mock_distributed(): + from vllm_omni.diffusion.models.wan2_2.wan2_2_transformer import ( + WanSelfAttention, + ) + + return WanSelfAttention(dim=512, num_heads=8, head_dim=64) + + def test_no_sparse_attn_attribute(self): + attn = self._make_self_attn() + assert not hasattr(attn, "sparse_attn") + + +class TestWanCrossAttentionIsSelfAttention: + """Test that WanCrossAttention passes is_self_attention=False.""" + + def _make_cross_attn(self): + with _mock_distributed(): + from vllm_omni.diffusion.models.wan2_2.wan2_2_transformer import ( + WanCrossAttention, + ) + + return WanCrossAttention(dim=512, num_heads=8, head_dim=64) + + def test_cross_attn_is_self_attention_false(self): + """Verify the Attention layer in WanCrossAttention was constructed + with is_self_attention=False. The impl should have received it.""" + cross_attn = self._make_cross_attn() + # The underlying attention impl should have received is_self_attention + # through extra_impl_args. We can check if the impl was constructed + # with is_self_attention=False by checking the Attention's construction. + # For sparse_attention backend this means _sparse_impl would be None. + # For other backends this is a no-op (they ignore the kwarg). + assert cross_attn.attn is not None + + +class TestNoEnableSparseAttention: + """Test that WanTransformer3DModel no longer has enable_sparse_attention.""" + + def _make_transformer(self): + from vllm.config import VllmConfig, set_current_vllm_config + + with _mock_distributed(), set_current_vllm_config(VllmConfig()): + from vllm_omni.diffusion.models.wan2_2.wan2_2_transformer import ( + WanTransformer3DModel, + ) + + return WanTransformer3DModel( + num_attention_heads=4, + attention_head_dim=64, + in_channels=16, + out_channels=16, + text_dim=512, + freq_dim=256, + ffn_dim=1024, + num_layers=2, + ) + + def test_no_enable_sparse_attention_method(self): + transformer = self._make_transformer() + assert not hasattr(transformer, "enable_sparse_attention") + + def test_self_attn_has_no_sparse_attn(self): + transformer = self._make_transformer() + for block in transformer.blocks: + assert not hasattr(block.attn1, "sparse_attn") diff --git a/vllm_omni/diffusion/attention/backends/registry.py b/vllm_omni/diffusion/attention/backends/registry.py index a77223d1678..2ea54c7b269 100644 --- a/vllm_omni/diffusion/attention/backends/registry.py +++ b/vllm_omni/diffusion/attention/backends/registry.py @@ -59,6 +59,7 @@ class MyCustomBackend: FLASH_ATTN = "vllm_omni.diffusion.attention.backends.flash_attn.FlashAttentionBackend" TORCH_SDPA = "vllm_omni.diffusion.attention.backends.sdpa.SDPABackend" SAGE_ATTN = "vllm_omni.diffusion.attention.backends.sage_attn.SageAttentionBackend" + SPARSE_ATTENTION = "vllm_omni.diffusion.attention.backends.sparse_attention.SparseAttentionBackend" def get_path(self, include_classname: bool = True) -> str: """Get the class path for this backend (respects overrides). diff --git a/vllm_omni/diffusion/attention/backends/sparse_attention.py b/vllm_omni/diffusion/attention/backends/sparse_attention.py new file mode 100644 index 00000000000..ea8546ba133 --- /dev/null +++ b/vllm_omni/diffusion/attention/backends/sparse_attention.py @@ -0,0 +1,245 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Built-in dispatcher backend for DiT sparse attention. + +Selected via --attn-backend sparse_attention. +Reads is_self_attention (from Attention.__init__ extra_impl_args): + - False -> SDPA fallback, zero overhead on cross-attention + - True -> look up sparse attention function from plugin + +Plugin resolution in self-attention path: + 1. sparse_attn.backend == "auto" -> first available in vllm_omni.sparse_attn EP group + 2. sparse_attn.backend == short name -> match in vllm_omni.sparse_attn EP group + 3. sparse_attn.backend == "module.path:func_name" -> import directly + 4. None found -> dense SDPA fallback with warning + +Plugins register a function with signature: + fn(q, k, v, params, is_causal) -> Tensor + q/k/v: (B, H, S, D) tensors + params: dict of plugin-specific parameters +""" + +from __future__ import annotations + +import importlib +import importlib.metadata +from collections.abc import Callable +from typing import TYPE_CHECKING + +import torch +from vllm.logger import init_logger + +from vllm_omni.diffusion.attention.backends.abstract import ( + AttentionBackend, + AttentionImpl, + AttentionMetadata, +) +from vllm_omni.diffusion.attention.backends.sdpa import SDPABackend + +if TYPE_CHECKING: + from vllm_omni.diffusion.data import DiffusionSparseAttnConfig + +logger = init_logger(__name__) + +# Type alias for sparse attention function plugins +SparseAttnFn = Callable[ + [torch.Tensor, torch.Tensor, torch.Tensor, dict, bool], + torch.Tensor, +] + + +class SparseAttentionBackend(AttentionBackend): + @staticmethod + def get_name() -> str: + return "sparse_attention" + + @staticmethod + def get_impl_cls(): + return _SparseAttentionImpl + + @staticmethod + def get_metadata_cls(): + return AttentionMetadata + + @staticmethod + def get_builder_cls(): + return None + + @staticmethod + def get_supported_head_sizes() -> list[int]: + return [] + + +class _SparseAttentionImpl(AttentionImpl): + def __init__( + self, + num_heads: int, + head_size: int, + softmax_scale: float, + causal: bool = False, + num_kv_heads: int | None = None, + prefix: str = "", + **extra_impl_args, + ) -> None: + self._num_heads = num_heads + self._head_size = head_size + self._softmax_scale = softmax_scale + self._causal = causal + self._is_self_attention: bool = extra_impl_args.get("is_self_attention", True) + + # Dense fallback (used for cross-attention, or when no plugin found) + self._dense = SDPABackend.get_impl_cls()( + num_heads=num_heads, + head_size=head_size, + softmax_scale=softmax_scale, + causal=causal, + num_kv_heads=num_kv_heads, + ) + + # Sparse attention function plugin (None = use dense) + self._sparse_fn: SparseAttnFn | None = None + self._params: dict = {} + + if not self._is_self_attention: + return # cross-attention: always dense, done + + cfg = self._read_sparse_cfg() + if cfg is None or cfg.backend in ("none", "dense"): + return # dense: use SDPA fallback, no plugin needed + + plugin = self._resolve_plugin(cfg) + if plugin is None: + logger.warning( + "No sparse attention plugin found for backend='%s'. Falling back to dense.", + cfg.backend, + ) + return + + self._sparse_fn = plugin + self._params = cfg.params + logger.info( + "Using sparse attention plugin '%s' (params=%s)", + cfg.backend, + cfg.params, + ) + + # ------------------------------------------------------------------ + # AttentionImpl interface + # ------------------------------------------------------------------ + + def _forward_sparse(self, query, key, value, attn_metadata=None): + """Sparse plugin path for self-attention on supported platforms.""" + # query/key/value: (B, S, H, D) -> (B, H, S, D) + q = query.transpose(1, 2) + k = key.transpose(1, 2) + v = value.transpose(1, 2) + + out = self._sparse_fn(q, k, v, self._params, self._causal) + + # (B, H, S, D) -> (B, S, H, D) + return out.transpose(1, 2) + + def forward_cuda(self, query, key, value, attn_metadata=None): + if self._sparse_fn is None: + return self._dense.forward_cuda(query, key, value, attn_metadata) + return self._forward_sparse(query, key, value, attn_metadata) + + def forward_hip(self, query, key, value, attn_metadata=None): + return self._dense.forward_hip(query, key, value, attn_metadata) + + def forward_npu(self, query, key, value, attn_metadata=None): + return self._dense.forward_npu(query, key, value, attn_metadata) + + def forward_xpu(self, query, key, value, attn_metadata=None): + return self._dense.forward_xpu(query, key, value, attn_metadata) + + def forward_musa(self, query, key, value, attn_metadata=None): + return self._dense.forward_musa(query, key, value, attn_metadata) + + # ------------------------------------------------------------------ + # Internal helpers + # ------------------------------------------------------------------ + + @staticmethod + def _read_sparse_cfg() -> DiffusionSparseAttnConfig | None: + try: + from vllm_omni.diffusion.forward_context import get_forward_context + + return get_forward_context().omni_diffusion_config.sparse_attn + except (RuntimeError, LookupError, AssertionError): + # RuntimeError / AssertionError: forward context not yet initialised + # LookupError: config key missing during early init + return None + except AttributeError: + logger.debug( + "sparse_attn config not available on omni_diffusion_config", + exc_info=True, + ) + return None + + @staticmethod + def _resolve_plugin(cfg: DiffusionSparseAttnConfig) -> SparseAttnFn | None: + """Resolve a sparse attention function from config. + + Returns a callable fn(q, k, v, params, is_causal) -> Tensor, + or None if no plugin found. + """ + name = cfg.backend + if not name or name == "auto": + return _SparseAttentionImpl._first_available_plugin() + if ":" in name or ("." in name and not name.startswith(".")): + # Fully qualified "module.path:func_name" or "module.path.func" + try: + return _import_function(name) + except Exception as e: + logger.warning("Failed to load plugin '%s': %s", name, e) + return None + # Short name -> vllm_omni.sparse_attn EP group + try: + for ep in importlib.metadata.entry_points(group="vllm_omni.sparse_attn"): + if ep.name.lower() == name.lower(): + try: + return ep.load() + except Exception: + logger.warning( + "Sparse plugin '%s' (%s) found but failed to load", + ep.name, + ep.value, + exc_info=True, + ) + return None + except Exception: + logger.debug("Failed to query entry_points for sparse_attn", exc_info=True) + return None + + @staticmethod + def _first_available_plugin() -> SparseAttnFn | None: + try: + for ep in importlib.metadata.entry_points(group="vllm_omni.sparse_attn"): + try: + fn = ep.load() + logger.info("Auto-selected sparse plugin: %s", ep.name) + return fn + except ImportError: + logger.debug("Sparse plugin '%s' not importable, skipping", ep.name) + continue + except Exception: + logger.warning( + "Sparse plugin '%s' found but failed to load", + ep.name, + exc_info=True, + ) + continue + except Exception: + logger.debug("Failed to query entry_points for sparse_attn", exc_info=True) + return None + + +def _import_function(path: str) -> SparseAttnFn: + """Import a function from 'module.path:func' or 'module.path.func'.""" + if ":" in path: + module_path, func_name = path.rsplit(":", 1) + else: + module_path, func_name = path.rsplit(".", 1) + module = importlib.import_module(module_path) + return getattr(module, func_name) diff --git a/vllm_omni/diffusion/attention/layer.py b/vllm_omni/diffusion/attention/layer.py index 4fdf2ff1612..b21104ccb10 100644 --- a/vllm_omni/diffusion/attention/layer.py +++ b/vllm_omni/diffusion/attention/layer.py @@ -37,6 +37,7 @@ def __init__( gather_idx: int = 1, use_sync: bool = False, skip_sequence_parallel: bool = False, + is_self_attention: bool = True, ): super().__init__() self.attn_backend = get_attn_backend(-1) @@ -47,6 +48,7 @@ def __init__( softmax_scale=softmax_scale, causal=causal, num_kv_heads=num_kv_heads, + is_self_attention=is_self_attention, ) # Instantiate fallback backend for float32 support self.sdpa_fallback = SDPABackend.get_impl_cls()( diff --git a/vllm_omni/diffusion/attention/selector.py b/vllm_omni/diffusion/attention/selector.py index 920396e66b9..098638274c0 100644 --- a/vllm_omni/diffusion/attention/selector.py +++ b/vllm_omni/diffusion/attention/selector.py @@ -30,6 +30,33 @@ logger = init_logger(__name__) +def _get_plugin_backend_cls(name: str) -> "type[AttentionBackend] | None": + """Check vllm_omni.attn_backend entry_points for a plugin with this name. + + For full AttentionBackend plugins (e.g. custom flash_attn replacement). + Sparse attention plugins should use the vllm_omni.sparse_attn entry + point group instead — they are resolved by the sparse_attention dispatcher. + + Plugin packages declare themselves in pyproject.toml: + [project.entry-points."vllm_omni.attn_backend"] + my_backend = "my_package.backend:MyAttentionBackend" + """ + try: + import importlib.metadata + + for ep in importlib.metadata.entry_points(group="vllm_omni.attn_backend"): + if ep.name.lower() == name.lower(): + logger.info( + "Loading attention backend plugin '%s' from %s", + name, + ep.value, + ) + return ep.load() + except Exception as e: + logger.debug("Plugin backend lookup failed for '%s': %s", name, e) + return None + + def _load_backend_cls(cls_path: str) -> type[AttentionBackend]: """Load a backend class from its fully qualified path. @@ -76,7 +103,13 @@ def get_attn_backend(head_size: int) -> type[AttentionBackend]: # Check environment variable for user override selected_backend = os.environ.get("DIFFUSION_ATTENTION_BACKEND") - # Delegate to platform for backend selection + # Check entry_points plugins before enum lookup + if selected_backend: + plugin_cls = _get_plugin_backend_cls(selected_backend) + if plugin_cls is not None: + return plugin_cls + + # Platform-based selection via enum backend_cls_path = current_omni_platform.get_diffusion_attn_backend_cls( selected_backend=selected_backend, head_size=head_size, diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 56a891aa5cf..a56c240f149 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -349,6 +349,48 @@ def __getattr__(self, item: str) -> Any: raise AttributeError(f"'{self.__class__.__name__}' object has no attribute '{item}'") +@dataclass +class DiffusionSparseAttnConfig: + """Sparse attention configuration for DiT modules. + + Stored in OmniDiffusionConfig.sparse_attn. + vllm-omni owns only `backend`. All kernel-level parameters + (e.g. topk, block_size) go in `params` and are passed + directly to the plugin function — vllm-omni does not + interpret them. + + Examples: + DiffusionSparseAttnConfig(backend="spargeattn", params={"topk": 0.5}) + DiffusionSparseAttnConfig.from_dict({"backend": "spargeattn", "topk": 0.5}) + """ + + backend: str = "auto" + # "auto" | "dense" | plugin short name | "module.path:func_name" + + params: dict[str, Any] = field(default_factory=dict) + # Passed verbatim to: fn(q, k, v, params, is_causal) -> Tensor + # The plugin is responsible for reading what it needs. + + @classmethod + def from_dict(cls, d: dict[str, Any]) -> "DiffusionSparseAttnConfig": + """Accept flat or nested JSON. + + Flat (unknown keys auto-routed into params): + {"backend": "spargeattn", "topk": 0.5} + -> DiffusionSparseAttnConfig(backend="spargeattn", params={"topk": 0.5}) + + Nested (explicit params key): + {"backend": "spargeattn", "params": {"topk": 0.5}} + -> DiffusionSparseAttnConfig(backend="spargeattn", params={"topk": 0.5}) + """ + backend = d.get("backend", "auto") + if "params" in d: + params = dict(d["params"]) + else: + params = {k: v for k, v in d.items() if k != "backend"} + return cls(backend=backend, params=params) + + @dataclass class OmniDiffusionConfig: # Model and path configuration (for convenience) @@ -379,6 +421,9 @@ class OmniDiffusionConfig: cache_config: DiffusionCacheConfig | dict[str, Any] = field(default_factory=dict) enable_cache_dit_summary: bool = False + # Sparse attention configuration for DiT modules + sparse_attn: DiffusionSparseAttnConfig | dict[str, Any] | None = None + # Distributed executor backend distributed_executor_backend: str = "mp" nccl_port: int | None = None @@ -611,6 +656,10 @@ def __post_init__(self): # If it's neither dict nor DiffusionCacheConfig, convert to empty config self.cache_config = DiffusionCacheConfig() + # Convert sparse_attn dict to DiffusionSparseAttnConfig if needed + if isinstance(self.sparse_attn, dict): + self.sparse_attn = DiffusionSparseAttnConfig.from_dict(self.sparse_attn) + # Auto-detect quantization from TransformerConfig if not explicitly set. # This covers the case where tf_model_config is passed at construction # time. For late (post-construction) assignment, callers should use @@ -738,7 +787,6 @@ class AttentionBackendEnum(enum.Enum): 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() 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 3b43f3eaf51..aba50887085 100644 --- a/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py +++ b/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py @@ -426,12 +426,10 @@ def forward( query = apply_rotary_emb_wan(query, freqs_cos, freqs_sin) key = apply_rotary_emb_wan(key, freqs_cos, freqs_sin) - # Create attention metadata if mask is provided + # Compute attention attn_metadata = None if attn_mask is not None: attn_metadata = AttentionMetadata(attn_mask=attn_mask) - - # Compute attention using unified attention layer hidden_states = self.attn(query, key, value, attn_metadata) hidden_states = hidden_states.flatten(2, 3) hidden_states = hidden_states.type_as(query) @@ -541,6 +539,7 @@ def __init__( softmax_scale=1.0 / (head_dim**0.5), causal=False, skip_sequence_parallel=True, + is_self_attention=False, ) def forward( diff --git a/vllm_omni/diffusion/worker/diffusion_model_runner.py b/vllm_omni/diffusion/worker/diffusion_model_runner.py index 535f053c388..b300b81fc6e 100644 --- a/vllm_omni/diffusion/worker/diffusion_model_runner.py +++ b/vllm_omni/diffusion/worker/diffusion_model_runner.py @@ -11,6 +11,7 @@ from __future__ import annotations import copy +import os import time from collections.abc import Iterable from contextlib import nullcontext @@ -123,6 +124,22 @@ def get_memory_context(): return memory_pool_context_fn(tag="weights") return nullcontext() + # Propagate attention backend from config to env var for selector.py + if self.od_config.attention_backend and not os.environ.get("DIFFUSION_ATTENTION_BACKEND"): + os.environ["DIFFUSION_ATTENTION_BACKEND"] = self.od_config.attention_backend + + # Warn if sparse_attn config is set but the backend is a known dense one + selected_backend = os.environ.get("DIFFUSION_ATTENTION_BACKEND", "").upper() + if self.od_config.sparse_attn is not None and selected_backend: + dense_backends = {"FLASH_ATTN", "TORCH_SDPA", "SAGE_ATTN"} + if selected_backend in dense_backends: + logger.warning( + "sparse_attn config is set but --attn-backend '%s' is a dense " + "backend that ignores sparse_attn config. " + "Use --attn-backend spargeattn (or another sparse plugin) instead.", + selected_backend, + ) + # Load model within forward context load_config = LoadConfig() model_loader = DiffusersPipelineLoader(load_config, od_config=self.od_config) diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index 9609cf6e26b..5649a6d5377 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -1137,6 +1137,79 @@ def _create_default_diffusion_stage_cfg(kwargs: dict[str, Any]) -> list: ) normalized_kwargs["dtype"] = str(normalized_kwargs["dtype"]).removeprefix("torch.") + # Handle --attn-backend: accepts simple name or JSON with config. + # --attn-backend spargeattn + # --attn-backend '{"backend":"spargeattn","topk":0.3}' + attn_backend_raw = normalized_kwargs.pop("attn_backend", None) + attn_backend = None + if isinstance(attn_backend_raw, str) and attn_backend_raw.strip().startswith("{"): + try: + attn_config = json.loads(attn_backend_raw) + # "backend" in JSON is the sparse *plugin* name, not the + # AttentionBackend. Store it in sparse_attn config and let + # the auto-select logic below pick "sparse_attention". + sparse_plugin = attn_config.pop("backend", None) + merged = dict(attn_config) + if sparse_plugin: + merged["backend"] = sparse_plugin + existing = normalized_kwargs.get("sparse_attn") + if isinstance(existing, dict): + existing.update(merged) + else: + normalized_kwargs["sparse_attn"] = merged + except json.JSONDecodeError: + logger.warning("Invalid --attn-backend JSON: %s", attn_backend_raw) + elif attn_backend_raw: + # Check if this is a known dense backend name — if so, route + # directly to attention_backend rather than sparse_attn. + from vllm_omni.diffusion.attention.backends.registry import DiffusionAttentionBackendEnum + + _dense_names = { + e.name.lower() + for e in DiffusionAttentionBackendEnum + if e is not DiffusionAttentionBackendEnum.SPARSE_ATTENTION + } + if attn_backend_raw.lower() in _dense_names: + attn_backend = attn_backend_raw.lower() + else: + # Plain plugin name (e.g. "spargeattn") — route to sparse_attn + # config so the sparse_attention dispatcher resolves it. + existing = normalized_kwargs.get("sparse_attn") + if isinstance(existing, dict): + existing.setdefault("backend", attn_backend_raw) + else: + normalized_kwargs["sparse_attn"] = {"backend": attn_backend_raw} + # Don't set attn_backend directly — let auto-select below pick + # "sparse_attention" as the dispatcher. + + # Also handle --sparse-attn JSON string if passed separately + sparse_attn = normalized_kwargs.get("sparse_attn") + if isinstance(sparse_attn, str): + try: + normalized_kwargs["sparse_attn"] = json.loads(sparse_attn) + sparse_attn = normalized_kwargs["sparse_attn"] + except json.JSONDecodeError: + logger.warning("Invalid --sparse-attn JSON: %s", sparse_attn) + normalized_kwargs["sparse_attn"] = None + sparse_attn = None + + # If sparse_attn is configured, auto-select the sparse_attention + # dispatcher backend. The sparse_attn.backend field names the + # *plugin* (e.g. "flashinfer", "spargeattn"), not the attention + # backend — the dispatcher is always "sparse_attention". + if not attn_backend: + _sparse_backend = None + if isinstance(sparse_attn, dict): + _sparse_backend = sparse_attn.get("backend") + elif hasattr(sparse_attn, "backend"): + _sparse_backend = sparse_attn.backend + if _sparse_backend and _sparse_backend not in ("none", "dense"): + attn_backend = "sparse_attention" + + # Set env var for selector.py (CLI overrides env var) + if attn_backend: + os.environ["DIFFUSION_ATTENTION_BACKEND"] = attn_backend + cache_backend = normalized_kwargs.get("cache_backend", "none") cache_config = AsyncOmniEngine._normalize_cache_config( cache_backend, @@ -1188,6 +1261,7 @@ def _create_default_diffusion_stage_cfg(kwargs: dict[str, Any]) -> list: "cache_backend": cache_backend, "cache_config": cache_config, "enable_cache_dit_summary": kwargs.get("enable_cache_dit_summary", False), + "sparse_attn": normalized_kwargs.get("sparse_attn", None), "enable_cpu_offload": kwargs.get("enable_cpu_offload", False), "enable_layerwise_offload": kwargs.get("enable_layerwise_offload", False), "enforce_eager": kwargs.get("enforce_eager", False), diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index 6e9adc24615..5e8c35fbdc4 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -307,6 +307,17 @@ def subparser_init(self, subparsers: argparse._SubParsersAction) -> FlexibleArgu help="Enable per-step diffusion execution so running requests can be aborted between denoise steps.", ) + # Attention backend selection + omni_config_group.add_argument( + "--attn-backend", + type=str, + default=None, + help="Attention backend for DiT modules. " + "Simple name: flash_attn, torch_sdpa, sage_attn, spargeattn. " + 'JSON with config: \'{"backend":"spargeattn","topk":0.3}\'. ' + "Can also be set via DIFFUSION_ATTENTION_BACKEND env var.", + ) + # VAE memory optimization parameters omni_config_group.add_argument( "--vae-use-slicing",