File size: 6,166 Bytes
f29c645 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 | """无需下载模型即可检查后端选择，模拟不同 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()
|