Quick clarification on the post above: The benchmarks and watermark bypass (PYTORCH_MPS_HIGH_WATERMARK_RATIO) reflect the current limitations we hit when running heavy diffusion loops under PyTorch’s MPS allocator on 16GB M4 nodes.
The reason I brought this to the MLX discussion is to understand how MLX’s native Metal allocator (mlx/backend/metal/allocator.cpp) compares when pushed past physical memory limits:
-
Does MLX’s allocator allow contiguous array allocations to gracefully page into virtual memory swap without throwing Metal buffer allocation failures, or does it enforce an internal ceiling?
-
For quantization types that lack dedicated GPU Metal kernels (like fp8 variants), does MLX support seamless in-place CPU stream evaluation on the same underlying buffer without data duplication?
Would love to hear how the PyTorch MPS/MLX architecture handles these edge cases differently from MPS.
Execution Demo:https://youtube.com/shorts/JXFWwQ1Vhu0?si=KurgQ8Ii6IRA49Hd