fix FA on Hopper

.gitignore

@@ -1,2 +1,4 @@
 AGENTS.md
-app_examples.py
+app_examples.py
+hftoken
+hftoken.txt

README.md

@@ -35,11 +35,11 @@ Gradio Space using the official Z-Image pipeline (`Tongyi-MAI/Z-Image-Turbo`) wi
 - `MODEL_DTYPE` (default `auto`): `bf16` if supported, else `fp16` (override with `bf16`/`fp16`/`fp32`).
 - `ENABLE_COMPILE` (default `true`): Enable `torch.compile` on the transformer.
 - `ENABLE_WARMUP` (default `false`): Run a quick warmup across resolutions after load (adds startup time).
-- `ATTENTION_BACKEND` (default `flash_3`): Backend for transformer attention (falls back to `flash`/`xformers`/`native`).
+- `ATTENTION_BACKEND` (default `native`): PyTorch SDPA by default; other backends fall back to `native` first. FlashAttention 3 (including `_flash_3_hub`) is only enabled on Hopper GPUs (compute capability 9.0); on Blackwell and other GPUs, the app overrides old FA3 settings with `native` before inference.
 - `OFFLOAD_TO_CPU_AFTER_RUN` (default `false`): Move the model back to CPU after each generation (useful on ZeroGPU; slower on normal GPUs).
 - `ENABLE_AOTI` (default `true`): Try to load ZeroGPU AoTI blocks via `spaces.aoti_blocks_load` for faster inference.
 - `AOTI_REPO` (default `zerogpu-aoti/Z-Image`): AoTI blocks repo.
-- `AOTI_VARIANT` (default `fa3`): AoTI variant.
+- `AOTI_VARIANT` (default `fa3`): AoTI variant. FA3 blocks are skipped on non-Hopper GPUs to avoid incompatible CUDA kernels.
 - `AOTI_ALLOW_LORA` (default `false`): Allow AoTI to load even if LoRA adapters are loaded (may crash; AoTI blocks generally don’t support LoRA).
 - `DEBUG` (default `false`): When set to a truthy value (`1`, `true`, `yes`, `on`), hide the Status/Debug floating panel.
 

app.py

@@ -19,7 +19,7 @@ LORA_PATH = os.environ.get("LORA_PATH", os.path.join("lora", "zit-mystic-xxx.saf
 HF_TOKEN = os.environ.get("HF_TOKEN")
 ENABLE_COMPILE = os.environ.get("ENABLE_COMPILE", "true").lower() == "true"
 ENABLE_WARMUP = os.environ.get("ENABLE_WARMUP", "false").lower() == "true"
-ATTENTION_BACKEND = os.environ.get("ATTENTION_BACKEND", "flash_3")
+ATTENTION_BACKEND = os.environ.get("ATTENTION_BACKEND", "native")
 OFFLOAD_TO_CPU_AFTER_RUN = os.environ.get("OFFLOAD_TO_CPU_AFTER_RUN", "false").lower() == "true"
 ENABLE_AOTI = os.environ.get("ENABLE_AOTI", "true").lower() == "true"
 AOTI_REPO = os.environ.get("AOTI_REPO", "zerogpu-aoti/Z-Image")
@@ -193,18 +193,33 @@ def parse_resolution(resolution: str) -> Tuple[int, int]:
     return 1024, 1024
 
 
+def supports_flash_attention_3() -> bool:
+    # FA3 的预编译内核面向 Hopper；ZeroGPU 换成 Blackwell 后不能继续使用。
+    try:
+        return torch.cuda.is_available() and torch.cuda.get_device_capability(0) == (9, 0)
+    except Exception:
+        return False
+
+
 def set_attention_backend_safe(transformer, backend: str) -> str:
     global attention_backend_error
+    attention_backend_error = None
     candidates: List[str] = []
-    if backend:
-        candidates.append(backend)
-        if backend.startswith("_"):
-            candidates.append(backend.lstrip("_"))
-        else:
-            candidates.append(f"_{backend}")
-    candidates.extend(["flash", "xformers", "native"])
+    fa3_backends = {"flash_3", "flash_3_hub", "flash_varlen_3", "flash_3_varlen_hub"}
+    if backend and backend.lstrip("_") in fa3_backends and not supports_flash_attention_3():
+        # 设置后端不会运行 CUDA 内核，不能靠 try/except 捕获随后发生的进程退出。
+        attention_backend_error = f"FlashAttention 3 requires Hopper (cc 9.0); current GPU: {get_gpu_summary()}"
+        candidates = ["native"]
+    else:
+        if backend:
+            candidates.append(backend)
+            if backend.startswith("_"):
+                candidates.append(backend.lstrip("_"))
+            else:
+                candidates.append(f"_{backend}")
+        candidates.extend(["native", "flash", "xformers"])
+    candidates = list(dict.fromkeys(candidates))
 
-    attention_backend_error = None
     errors: dict[str, Exception] = {}
     last_exc: Exception | None = None
     for name in candidates:
@@ -607,7 +622,11 @@ def init_app() -> None:
         if ENABLE_COMPILE and pipe is not None:
             ensure_on_gpu()
         if ENABLE_AOTI and not aoti_loaded and pipe is not None and getattr(pipe, "transformer", None) is not None:
-            if lora_loaded and not AOTI_ALLOW_LORA:
+            if AOTI_VARIANT == "fa3" and not supports_flash_attention_3():
+                aoti_loaded = False
+                aoti_error = "disabled: FA3 AoTI blocks require a Hopper GPU (cc 9.0)"
+                print(f"AoTI {aoti_error}")
+            elif lora_loaded and not AOTI_ALLOW_LORA:
                 aoti_loaded = False
                 aoti_error = "disabled: AoTI blocks are incompatible with LoRA adapters"
                 print("AoTI disabled: LoRA adapters are loaded (AoTI blocks are incompatible with LoRA).")

tests/test_attention_backend.py

@@ -0,0 +1,126 @@
+"""无需下载模型即可检查后端选择，模拟不同 GPU 和 CUDA 检测失败。"""
+
+import ast
+import contextlib
+import io
+from pathlib import Path
+from types import SimpleNamespace
+import unittest
+from unittest.mock import Mock
+
+
+APP_PATH = Path(__file__).resolve().parents[1] / "app.py"
+
+
+class AttentionBackendTests(unittest.TestCase):
+    def setUp(self):
+        # 只加载待测函数，避免 app.py 导入时下载模型和启动 Gradio。
+        tree = ast.parse(APP_PATH.read_text(encoding="utf-8"))
+        names = {"get_gpu_summary", "supports_flash_attention_3", "set_attention_backend_safe", "init_app"}
+        functions = [node for node in tree.body if isinstance(node, ast.FunctionDef) and node.name in names]
+        self.cuda = SimpleNamespace(
+            is_available=Mock(return_value=True),
+            get_device_capability=Mock(return_value=(12, 0)),
+            get_device_name=Mock(return_value="RTX PRO 6000 Blackwell"),
+        )
+        self.namespace = {"torch": SimpleNamespace(cuda=self.cuda), "List": list}
+        exec(compile(ast.Module(body=functions, type_ignores=[]), str(APP_PATH), "exec"), self.namespace)
+        self.transformer = SimpleNamespace(set_attention_backend=Mock(), layers=SimpleNamespace())
+
+    def select(self, backend):
+        with contextlib.redirect_stdout(io.StringIO()):
+            return self.namespace["set_attention_backend_safe"](self.transformer, backend)
+
+    def test_blackwell_never_selects_flash_attention_3(self):
+        for capability in [(10, 0), (12, 0)]:
+            for alias in ["flash_3", "_flash_3", "_flash_3_hub", "_flash_varlen_3", "_flash_3_varlen_hub"]:
+                with self.subTest(capability=capability, backend=alias):
+                    self.cuda.get_device_capability.return_value = capability
+                    self.transformer.set_attention_backend.reset_mock()
+                    self.assertEqual(self.select(alias), "native")
+                    self.transformer.set_attention_backend.assert_called_once_with("native")
+                    self.assertIn("Hopper", self.namespace["attention_backend_error"])
+
+    def test_ampere_uses_native_for_flash_attention_3(self):
+        self.cuda.get_device_capability.return_value = (8, 0)
+        self.assertEqual(self.select("_flash_3_hub"), "native")
+        self.transformer.set_attention_backend.assert_called_once_with("native")
+
+    def test_hopper_preserves_requested_flash_attention_3(self):
+        self.cuda.get_device_capability.return_value = (9, 0)
+        self.assertEqual(self.select("_flash_3_hub"), "_flash_3_hub")
+        self.transformer.set_attention_backend.assert_called_once_with("_flash_3_hub")
+        self.assertIsNone(self.namespace["attention_backend_error"])
+
+    def test_unknown_gpu_uses_native(self):
+        self.cuda.get_device_capability.side_effect = RuntimeError("CUDA unavailable")
+        self.assertEqual(self.select("_flash_3_hub"), "native")
+        self.transformer.set_attention_backend.assert_called_once_with("native")
+
+    def test_no_cuda_uses_native_without_querying_device(self):
+        self.cuda.is_available.return_value = False
+        self.assertEqual(self.select("_flash_3_hub"), "native")
+        self.cuda.get_device_capability.assert_not_called()
+
+    def test_native_failure_does_not_try_incompatible_kernels(self):
+        self.transformer.set_attention_backend.side_effect = RuntimeError("native failed")
+        with self.assertRaisesRegex(RuntimeError, "native failed"):
+            self.select("_flash_3_hub")
+        self.transformer.set_attention_backend.assert_called_once_with("native")
+
+    def test_backend_setup_error_falls_back_to_native(self):
+        def configure(name):
+            if name != "native":
+                raise ValueError("backend unavailable")
+
+        self.transformer.set_attention_backend.side_effect = configure
+        self.assertEqual(self.select("missing_backend"), "native")
+        attempts = [call.args[0] for call in self.transformer.set_attention_backend.call_args_list]
+        self.assertEqual(attempts, ["missing_backend", "_missing_backend", "native"])
+        self.assertEqual(self.namespace["attention_backend_error"], "backend unavailable")
+
+    def test_native_selection_clears_previous_error(self):
+        self.select("_flash_3_hub")
+        self.assertEqual(self.select("native"), "native")
+        self.assertIsNone(self.namespace["attention_backend_error"])
+
+    def initialize(self, *, lora=False, allow_lora=False):
+        loader = Mock()
+        self.namespace.update(
+            pipe=SimpleNamespace(transformer=self.transformer),
+            ensure_models_loaded=Mock(), ENABLE_COMPILE=False, ENABLE_AOTI=True,
+            ENABLE_WARMUP=False, aoti_loaded=False, aoti_error=None,
+            lora_loaded=lora, AOTI_ALLOW_LORA=allow_lora, AOTI_VARIANT="fa3",
+            AOTI_REPO="zerogpu-aoti/Z-Image", ATTENTION_BACKEND="_flash_3_hub",
+            spaces=SimpleNamespace(aoti_blocks_load=loader),
+        )
+        with contextlib.redirect_stdout(io.StringIO()):
+            self.namespace["init_app"]()
+        return loader
+
+    def test_blackwell_skips_fa3_aoti(self):
+        loader = self.initialize()
+        loader.assert_not_called()
+        self.assertFalse(self.namespace["aoti_loaded"])
+        self.assertIn("Hopper", self.namespace["aoti_error"])
+
+    def test_allow_lora_cannot_override_hardware_guard(self):
+        loader = self.initialize(lora=True, allow_lora=True)
+        loader.assert_not_called()
+        self.assertFalse(self.namespace["aoti_loaded"])
+
+    def test_hopper_can_load_fa3_aoti(self):
+        self.cuda.get_device_capability.return_value = (9, 0)
+        loader = self.initialize()
+        loader.assert_called_once_with(self.transformer.layers, "zerogpu-aoti/Z-Image", variant="fa3")
+        self.assertTrue(self.namespace["aoti_loaded"])
+
+    def test_hopper_still_skips_aoti_with_lora(self):
+        self.cuda.get_device_capability.return_value = (9, 0)
+        loader = self.initialize(lora=True)
+        loader.assert_not_called()
+        self.assertIn("LoRA", self.namespace["aoti_error"])
+
+
+if __name__ == "__main__":
+    unittest.main()