Select attention kernels for the allocated GPU architecture

app/README.md

@@ -100,8 +100,11 @@ resolved and its CUDA modules are registered during startup so the platform can
 restore them efficiently for each generation. Input decoding and resizing run
 before the GPU lease, and MP4 encoding runs after it; only model inference is in
 the decorated generation callback. The app uses Hugging Face's prebuilt
-FlashAttention 3 kernel when compatible and falls back to installed
-FlashAttention or PyTorch SDPA otherwise. Set `WHATMOVES_DISABLE_FA3=1` to skip
+FlashAttention 3 kernel on Hopper GPUs, installed FlashAttention 2 on supported
+Ampere/Hopper GPUs, and PyTorch SDPA on Blackwell (including ZeroGPU's RTX Pro
+6000). Compatibility is checked on the allocated GPU before any kernel launch;
+the selected backend and GPU model appear in the timing logs.
+Set `WHATMOVES_DISABLE_FA3=1` to skip
 FlashAttention 3. A generation uses an immutable snapshot of the selected
 target, source intervals, masks, and settings present when **Generate video**
 was pressed; you may continue editing while it runs. The completed video is

tests/test_attention_backend.py

@@ -0,0 +1,118 @@
+"""CPU-only regression tests for runtime CUDA kernel selection.
+
+Run with: python -m unittest discover -s tests -v
+"""
+
+import importlib.util
+import os
+from pathlib import Path
+import sys
+from types import SimpleNamespace
+import unittest
+from unittest.mock import Mock, patch
+
+import torch
+
+
+class AttentionBackendTests(unittest.TestCase):
+    def setUp(self):
+        self.fa3 = SimpleNamespace(flash_attn_varlen_func=Mock())
+        spec = importlib.util.spec_from_file_location(
+            "attention_under_test",
+            Path(__file__).resolve().parents[1] / "wan/modules/attention.py",
+        )
+        self.module = importlib.util.module_from_spec(spec)
+        # Importing a prebuilt extension must not inspect ZeroGPU's startup
+        # proxy or launch any CUDA work. Hardware is selected at inference.
+        with (
+            patch.dict(os.environ, {"WHATMOVES_DISABLE_FA3": "0"}),
+            patch.dict(sys.modules, {
+                "kernels": SimpleNamespace(get_kernel=Mock(return_value=self.fa3)),
+                "flash_attn": None,
+            }),
+            patch.object(torch.cuda, "get_device_capability",
+                         side_effect=AssertionError("startup CUDA probe")),
+        ):
+            spec.loader.exec_module(self.module)
+        self.fa3.flash_attn_varlen_func.assert_not_called()
+
+    def test_architecture_and_availability_matrix(self):
+        cases = [
+            ((9, 0), True, False, None, 3),
+            ((9, 0), False, True, None, 2),
+            ((9, 0), True, True, 2, 2),
+            ((9, 0), True, False, 2, None),
+            ((8, 0), True, True, None, 2),
+            ((8, 9), True, True, None, 2),
+            ((8, 0), True, False, None, None),
+            ((10, 0), True, True, None, None),
+            ((12, 0), True, True, None, None),
+            ((12, 0), True, False, 3, None),
+            ((7, 5), True, True, None, None),
+        ]
+        for capability, fa3, fa2, requested, expected in cases:
+            with self.subTest(capability=capability, fa3=fa3, fa2=fa2,
+                              requested=requested):
+                self.module._flash_version_for_device.cache_clear()
+                with (
+                    patch.object(self.module, "FLASH_ATTN_3_AVAILABLE", fa3),
+                    patch.object(self.module, "FLASH_ATTN_2_AVAILABLE", fa2),
+                    patch.object(torch.cuda, "get_device_capability",
+                                 return_value=capability),
+                    patch.object(torch.cuda, "get_device_name", return_value="test GPU"),
+                ):
+                    self.assertEqual(
+                        self.module._flash_version_for_device(
+                            torch.device("cuda:0"), requested),
+                        expected,
+                    )
+
+    def test_blackwell_never_launches_imported_fa3(self):
+        torch.manual_seed(0)
+        q = torch.randn(1, 5, 2, 8)
+        k = torch.randn(1, 5, 2, 8)
+        v = torch.randn(1, 5, 2, 8)
+        # The mocked query supplies CUDA metadata while actual SDPA runs on
+        # CPU, so this checks dispatch without needing or charging for a GPU.
+        allocated_q = SimpleNamespace(device=torch.device("cuda:0"),
+                                      transpose=q.transpose)
+        expected = torch.nn.functional.scaled_dot_product_attention(
+            q.transpose(1, 2).to(torch.bfloat16),
+            k.transpose(1, 2).to(torch.bfloat16),
+            v.transpose(1, 2).to(torch.bfloat16),
+        ).transpose(1, 2).contiguous()
+        with (
+            patch.object(torch.cuda, "get_device_capability", return_value=(12, 0)),
+            patch.object(torch.cuda, "get_device_name", return_value="RTX Pro 6000"),
+        ):
+            for entrypoint in (self.module.attention, self.module.flash_attention):
+                torch.testing.assert_close(entrypoint(allocated_q, k, v), expected)
+        self.fa3.flash_attn_varlen_func.assert_not_called()
+
+    def test_hopper_routes_to_flash_attention(self):
+        q = SimpleNamespace(device=torch.device("cuda:0"))
+        with (
+            patch.object(torch.cuda, "get_device_capability", return_value=(9, 0)),
+            patch.object(torch.cuda, "get_device_name", return_value="H100"),
+            patch.object(self.module, "flash_attention", return_value="flash result") as call,
+        ):
+            self.assertEqual(self.module.attention(q, q, q), "flash result")
+            call.assert_called_once()
+
+    def test_cpu_does_not_query_cuda(self):
+        with patch.object(torch.cuda, "get_device_capability",
+                          side_effect=AssertionError("CPU queried CUDA")):
+            self.assertIsNone(self.module._flash_version_for_device(torch.device("cpu")))
+
+    def test_selection_is_cached_per_device(self):
+        with (
+            patch.object(torch.cuda, "get_device_capability", return_value=(12, 0)) as probe,
+            patch.object(torch.cuda, "get_device_name", return_value="RTX Pro 6000"),
+        ):
+            for _ in range(5):
+                self.module._flash_version_for_device(torch.device("cuda:0"))
+            probe.assert_called_once()
+
+
+if __name__ == "__main__":
+    unittest.main()

wan/modules/attention.py

@@ -1,4 +1,5 @@
 # Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+from functools import lru_cache
 import logging
 import os
 import warnings
@@ -28,7 +29,7 @@ def _load_flash_attention_3():
         if not callable(getattr(interface, "flash_attn_varlen_func", None)):
             raise AttributeError("kernel has no flash_attn_varlen_func")
         logger.info(
-            "flash_attention backend=fa3 source=kernels-community/vllm-flash-attn3"
+            "flash_attention fa3_loaded source=kernels-community/vllm-flash-attn3"
         )
         return interface
     except Exception as error:
@@ -49,7 +50,7 @@ try:
 
         flash_attn_interface = _installed_flash_attn_interface
         FLASH_ATTN_3_AVAILABLE = True
-        logger.info("flash_attention backend=fa3 source=installed")
+        logger.info("flash_attention fa3_loaded source=installed")
 except Exception as error:
     logger.info(
         "flash_attention installed_fa3_unavailable error=%s: %s",
@@ -75,6 +76,37 @@ __all__ = [
 ]
 
 
+@lru_cache(maxsize=32)
+def _flash_version_for_device(device, requested_version=None):
+    """Select kernels on the allocated GPU, not ZeroGPU's startup proxy.
+
+    Importing FA3 does not establish hardware compatibility. Its Hopper
+    kernels cannot execute on Blackwell (including ZeroGPU's RTX Pro 6000),
+    and the CUDA extension can terminate the worker instead of raising a
+    catchable Python exception. Reject incompatible architectures before the
+    first kernel launch. PyTorch SDPA dispatches its own supported kernels.
+    """
+    if device.type != "cuda":
+        return None
+    capability = torch.cuda.get_device_capability(device)
+    selected = None
+    if (
+        requested_version in (None, 3)
+        and FLASH_ATTN_3_AVAILABLE
+        and capability == (9, 0)
+    ):
+        selected = 3
+    elif FLASH_ATTN_2_AVAILABLE and capability[0] in (8, 9):
+        selected = 2
+    logger.info(
+        "flash_attention selected_backend=%s gpu=%s compute_capability=%d.%d",
+        f"fa{selected}" if selected else "torch_sdpa",
+        torch.cuda.get_device_name(device),
+        *capability,
+    )
+    return selected
+
+
 def flash_attention(
     q,
     k,
@@ -103,6 +135,14 @@ def flash_attention(
     deterministic:  bool. If True, slightly slower and uses more memory.
     dtype:          torch.dtype. Apply when dtype of q/k/v is not float16/bfloat16.
     """
+    selected_version = _flash_version_for_device(q.device, version)
+    if selected_version is None:
+        return attention(
+            q=q, k=k, v=v, q_lens=q_lens, k_lens=k_lens,
+            dropout_p=dropout_p, softmax_scale=softmax_scale, q_scale=q_scale,
+            causal=causal, window_size=window_size,
+            deterministic=deterministic, dtype=dtype, fa_version=version,
+        )
     half_dtypes = (torch.float16, torch.bfloat16)
     assert dtype in half_dtypes
     assert q.device.type == 'cuda' and q.size(-1) <= 256
@@ -145,7 +185,7 @@ def flash_attention(
         )
 
     # apply attention
-    if (version is None or version == 3) and FLASH_ATTN_3_AVAILABLE:
+    if selected_version == 3:
         # Note: dropout_p, window_size are not supported in FA3 now.
         x = flash_attn_interface.flash_attn_varlen_func(
             q=q,
@@ -199,7 +239,7 @@ def attention(
     dtype=torch.bfloat16,
     fa_version=None,
 ):
-    if FLASH_ATTN_2_AVAILABLE or FLASH_ATTN_3_AVAILABLE:
+    if _flash_version_for_device(q.device, fa_version) is not None:
         return flash_attention(
             q=q,
             k=k,