[diffusion] fix: fix using upstream flash_attn on blackwell (#17111)

2026-01-15 22:30:48 +08:00
parent c9a45b7e3c
commit 16831ab6d7
3 changed files with 28 additions and 10 deletions
@@ -115,6 +115,13 @@ class Platform:
            return False
        return torch.cuda.get_device_capability()[0] == 10

+    @classmethod
+    @lru_cache(maxsize=1)
+    def is_hopper(cls):
+        if not cls.is_cuda_static():
+            return False
+        return torch.cuda.get_device_capability() == (9, 0)
+
    @classmethod
    @lru_cache(maxsize=1)
    def is_sm120(cls):