| """无需下载模型即可检查后端选择，模拟不同 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): |
| |
| 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() |
|
|
