Fix issue with memory manager on MPS devices

extensions_built_in/diffusion_models/krea2/src/mmdit.py

@@ -31,7 +31,11 @@ from torch.utils.checkpoint import checkpoint
 
 
 def rope(pos: Tensor, dim: int, theta: float = 1e4, ntk: float = 1.0) -> Tensor:
-    scale = torch.arange(0, dim, 2, dtype=torch.float64, device=pos.device) / dim
+    # MPS has no float64
+    dt = torch.float32 if pos.device.type == "mps" else torch.float64
+    if dt is torch.float32:
+        pos = pos.to(dt)
+    scale = torch.arange(0, dim, 2, dtype=dt, device=pos.device) / dim
     omega = 1.0 / ((theta * ntk) ** scale)
     out = torch.einsum("...n,d->...nd", pos, omega)
     out = torch.stack(

toolkit/memory_management/manager.py

@@ -457,15 +457,18 @@ class MemoryManager:
         gc.collect()
         torch.cuda.empty_cache()
 
-        # torch's pinned-host cache; private API, name varies by torch version
+        # torch's pinned-host cache; private API, name varies by torch version.
-        for fn_name in ("_accelerator_emptyHostCache", "_host_emptyCache"):
+        # Skipped on MPS: it segfaults there (pytorch/pytorch#197561), not catchable
-            fn = getattr(torch._C, fn_name, None)
+        on_mps = getattr(torch.backends, "mps", None) is not None and torch.backends.mps.is_available()
-            if fn is not None:
+        if not on_mps:
-                try:
+            for fn_name in ("_accelerator_emptyHostCache", "_host_emptyCache"):
-                    fn()
+                fn = getattr(torch._C, fn_name, None)
-                    break
+                if fn is not None:
-                except Exception:
+                    try:
-                    pass
+                        fn()
+                        break
+                    except Exception:
+                        pass
 
         # glibc keeps freed arenas mapped; CDLL(None) resolves malloc_trim in
         # the running process where glibc is present and fails cleanly on