Fix failing python tests on Windows (ml-explore#3076)

Author: zcbenz

mlx/CMakeLists.txt

@@ -32,10 +32,11 @@ set_target_properties(
              CXX_VISIBILITY_PRESET hidden
              CUDA_VISIBILITY_PRESET hidden)
 
-# Define MLX_EXPORT for shared libraries.
-set_target_properties(mlx mlx_version PROPERTIES DEFINE_SYMBOL MLX_EXPORT)
-# Define MLX_STATIC for static libraries.
-if(NOT BUILD_SHARED_LIBS)
+# Define MLX_EXPORT for shared libraries, MLX_STATIC for static libraries.
+set_target_properties(mlx PROPERTIES DEFINE_SYMBOL MLX_EXPORT)
+if(BUILD_SHARED_LIBS)
+  target_compile_definitions(mlx_version PUBLIC MLX_EXPORT)
+else()
   target_compile_definitions(mlx PUBLIC MLX_STATIC)
   target_compile_definitions(mlx_version PUBLIC MLX_STATIC)
 endif()
@@ -49,20 +50,20 @@ endif()
 
 if(MSVC)
   # Some of CUDA's headers include windows.h, which defines min/max macros.
-  target_compile_definitions(mlx PRIVATE NOMINMAX)
+  target_compile_definitions(mlx PRIVATE NOMINMAX WIN32_LEAN_AND_MEAN)
+  # Unicode support in fmt does not compile in .cu files.
+  target_compile_definitions(mlx PRIVATE FMT_UNICODE=0)
   # Disable some MSVC warnings to speed up compilation.
   target_compile_options(
     mlx
-    PUBLIC $<$<COMPILE_LANGUAGE:CXX>:/wd4068
-           /wd4244
-           /wd4267
-           /wd4700
-           /wd4804>
-           $<$<COMPILE_LANGUAGE:CUDA>:-Xcompiler=/wd4068
-           -Xcompiler=/wd4244
-           -Xcompiler=/wd4267
-           -Xcompiler=/wd4700
-           -Xcompiler=/wd4804>)
+    PUBLIC $<$<COMPILE_LANGUAGE:CXX>:/wd4244 /wd4267>
+    PRIVATE $<$<COMPILE_LANGUAGE:CXX>:/wd4068
+            /wd4146
+            /wd4700
+            /wd4804
+            /wd4805>
+            $<$<COMPILE_LANGUAGE:CUDA>:-Xcompiler=/wd4244
+            -Xcompiler=/wd4267>)
   # Enable /bigobj for heavily templated code (e.g., binary.cpp) that exceeds
   # the default 65,535 section limit in COFF object files.
   target_compile_options(

mlx/array.h

@@ -489,10 +489,10 @@ class MLX_API array {
     int64_t offset{0};
 
     // The size in elements of the data buffer the array accesses
-    size_t data_size;
+    size_t data_size{0};
 
     // Contains useful meta data about the array
-    Flags flags;
+    Flags flags{true, true, true};
 
     std::vector<array> inputs;
     // An array to keep track of the siblings from a multi-output

mlx/backend/cpu/device_info.cpp

@@ -6,8 +6,6 @@
 #include <sys/sysctl.h>
 #include <sys/utsname.h>
 #elif defined(_WIN32)
-#define WIN32_LEAN_AND_MEAN
-#define NOMINMAX
 #include <windows.h>
 #else
 #include <sys/utsname.h>

mlx/backend/cuda/allocator.cpp

@@ -196,7 +196,7 @@ CudaAllocator::malloc_async(size_t size, int device, cudaStream_t stream) {
       if (device == -1) {
         data = unified_malloc(size);
       } else {
-        if (free_streams_[device]) { // supports memory pools
+        if (mem_pools_[device]) { // supports memory pools
           CHECK_CUDA_ERROR(cudaMallocAsync(&data, size, stream));
         } else {
           CHECK_CUDA_ERROR(cudaMalloc(&data, size));
@@ -283,12 +283,13 @@ void CudaAllocator::move_to_unified_memory(
   void* data = unified_malloc(buf.size);
   cudaMemcpyKind kind =
       supports_managed_memory() ? cudaMemcpyDefault : cudaMemcpyDeviceToHost;
-  if (stream) {
+  if (stream && mem_pools_[buf.device]) {
     CHECK_CUDA_ERROR(cudaMemcpyAsync(data, buf.data, buf.size, kind, stream));
+    free_async(buf, stream);
   } else {
     CHECK_CUDA_ERROR(cudaMemcpy(data, buf.data, buf.size, kind));
+    free_async(buf);
   }
-  cuda_free(buf);
   buf.data = data;
   buf.device = -1;
 }
@@ -298,17 +299,20 @@ void CudaAllocator::free_cuda_buffer(CudaBuffer* buf) {
   if (scalar_pool_.in_pool(buf)) {
     scalar_pool_.free(buf);
   } else {
-    cuda_free(*buf);
+    free_async(*buf);
     delete buf;
   }
 }
 
-void CudaAllocator::cuda_free(CudaBuffer& buf) {
+void CudaAllocator::free_async(CudaBuffer& buf, cudaStream_t stream) {
   if (buf.device == -1) {
     unified_free(buf.data);
   } else {
-    cudaStream_t stream = free_streams_[buf.device];
-    if (stream) {
+    // Free asynchronously when memory pools is supported.
+    if (mem_pools_[buf.device]) {
+      if (!stream) {
+        stream = free_streams_[buf.device];
+      }
       CHECK_CUDA_ERROR(cudaFreeAsync(buf.data, stream));
     } else {
       CHECK_CUDA_ERROR(cudaFree(buf.data));

mlx/backend/cuda/allocator.h

@@ -69,7 +69,7 @@ class CudaAllocator : public allocator::Allocator {
 
  private:
   void free_cuda_buffer(CudaBuffer* buf);
-  void cuda_free(CudaBuffer& buf);
+  void free_async(CudaBuffer& buf, cudaStream_t stream = nullptr);
 
   CudaAllocator();
   friend CudaAllocator& allocator();

mlx/backend/cuda/gemms/grouped_gemm_unaligned.cu

@@ -124,12 +124,12 @@ struct GemmConfiguration : public CommonGemmConfiguration<T, Arch, 1> {
 };
 
 // Specialized GEMM configuration for sm80 and later.
-template <typename T, typename Arch, int kAlignmentC, bool kEnableTF32>
+template <typename T, typename Arch, int kAlignmentC>
 struct GemmConfiguration<
     T,
     Arch,
     kAlignmentC,
-    kEnableTF32,
+    true,
     std::enable_if_t<Arch::kMinComputeCapability >= 80 && sizeof(T) <= 4>>
     : public CommonGemmConfiguration<T, cutlass::arch::Sm80, kAlignmentC> {
   using OpClass = cutlass::arch::OpClassTensorOp;

mlx/backend/cuda/quantized/qmv.cu

@@ -232,8 +232,8 @@ void fp_qmv(
     using T = cuda_type_t<MLX_GET_TYPE(type_tag)>;
     if constexpr (!std::is_same_v<T, double>) {
       dim3 block_dims{WARP_SIZE, rows_per_block};
-      uint B = out.size() / (M * N);
-      uint blocks_y = (N + rows_per_block - 1) / rows_per_block;
+      uint32_t B = out.size() / (M * N);
+      uint32_t blocks_y = (N + rows_per_block - 1) / rows_per_block;
       const uint32_t* mat_ptr = gpu_ptr<uint32_t>(mat);
       const T* vec_ptr = gpu_ptr<T>(vec);
       int n = 1;
@@ -249,16 +249,17 @@ void fp_qmv(
       }
       dispatch_1_2_4(n, [&](auto n) {
         dispatch_bool(B > 1, [&](auto batched) {
-          if (!batched()) {
-            auto kernel = fp_qmv_single<T, rows_per_block, n(), 4, 32, true>;
+          if (!batched.value) {
+            auto kernel =
+                fp_qmv_single<T, rows_per_block, n.value, 4, 32, true>;
             if (bits == 8) {
-              kernel = fp_qmv_single<T, rows_per_block, n(), 8, 32, true>;
+              kernel = fp_qmv_single<T, rows_per_block, n.value, 8, 32, true>;
             } else if (group_size == 16) {
-              kernel = fp_qmv_single<T, rows_per_block, n(), 4, 16, false>;
+              kernel = fp_qmv_single<T, rows_per_block, n.value, 4, 16, false>;
             }
             encoder.add_kernel_node(
                 kernel,
-                {static_cast<uint>(M), blocks_y},
+                {static_cast<uint32_t>(M), blocks_y},
                 block_dims,
                 0,
                 mat_ptr,
@@ -268,15 +269,16 @@ void fp_qmv(
                 N,
                 K);
           } else {
-            auto kernel = fp_qmv_batched<T, rows_per_block, n(), 4, 32, true>;
+            auto kernel =
+                fp_qmv_batched<T, rows_per_block, n.value, 4, 32, true>;
             if (bits == 8) {
-              kernel = fp_qmv_batched<T, rows_per_block, n(), 8, 32, true>;
+              kernel = fp_qmv_batched<T, rows_per_block, n.value, 8, 32, true>;
             } else if (group_size == 16) {
-              kernel = fp_qmv_batched<T, rows_per_block, n(), 4, 16, false>;
+              kernel = fp_qmv_batched<T, rows_per_block, n.value, 4, 16, false>;
             }
             encoder.add_kernel_node(
                 kernel,
-                {static_cast<uint>(M), blocks_y, B},
+                {static_cast<uint32_t>(M), blocks_y, B},
                 block_dims,
                 0,
                 mat_ptr,

mlx/backend/cuda/scaled_dot_product_attention.cpp

@@ -140,7 +140,7 @@ DnnGraph build_sdpa_graph(
     const std::optional<array>& mask_arr,
     bool output_logsumexp,
     const array& o,
-    const array& stats) {
+    const std::optional<array>& stats) {
   DnnGraph graph(handle, q.dtype());
 
   auto q_ = graph.tensor("Q", Q, q);
@@ -161,7 +161,7 @@ DnnGraph build_sdpa_graph(
   auto [o_, stats_] = graph.sdpa(q_, k_, v_, options);
   graph.tensor(o_, O, o)->set_output(true);
   if (output_logsumexp) {
-    graph.tensor(stats_, STATS, stats)->set_output(true);
+    graph.tensor(stats_, STATS, *stats)->set_output(true);
   }
 
   CHECK_CUDNN_FE_ERROR(graph.prepare());
@@ -239,6 +239,11 @@ bool supports_sdpa_cudnn(
     return false;
   }
 
+  // cuDNN does not support bottom right mask when T_q > T_kv.
+  if (do_causal && (q.shape(2) > k.shape(2))) {
+    return false;
+  }
+
   // D_qk and D_v must be a multiple of 8 with maximum value 128.
   if ((q.shape(-1) % 8 != 0) || (q.shape(-1) > 128) || (v.shape(-1) % 8 != 0) ||
       (v.shape(-1) > 128)) {
@@ -255,7 +260,7 @@ void sdpa_cudnn(
     const array& v,
     float scale,
     array& o,
-    array& stats,
+    std::optional<array>& stats,
     bool do_causal,
     const std::optional<array>& mask_arr,
     bool output_logsumexp,
@@ -273,8 +278,8 @@ void sdpa_cudnn(
     encoder.set_input_array(*mask_arr);
   }
   if (output_logsumexp) {
-    stats.set_data(cu::malloc_async(stats.nbytes(), encoder));
-    encoder.set_output_array(stats);
+    stats->set_data(cu::malloc_async(stats->nbytes(), encoder));
+    encoder.set_output_array(*stats);
   }
 
   // Search cache.
@@ -298,7 +303,7 @@ void sdpa_cudnn(
     variant_pack[BIAS] = gpu_ptr<void>(*mask_arr);
   }
   if (output_logsumexp) {
-    variant_pack[STATS] = gpu_ptr<void>(stats);
+    variant_pack[STATS] = gpu_ptr<void>(*stats);
   }
 
   CHECK_CUDNN_FE_ERROR(graph.encode_graph(encoder, std::move(variant_pack)));
@@ -420,15 +425,18 @@ void ScaledDotProductAttention::eval_gpu(
   array q = prepare_sdpa_input(inputs[0], s);
   array k = prepare_sdpa_input(inputs[1], s);
   array v = prepare_sdpa_input(inputs[2], s);
-  auto& out = outputs[0];
-  auto& stats = outputs[1];
+  array& out = outputs[0];
   bool has_mask = inputs.size() - has_sinks_ > 3;
   bool has_arr_mask = has_mask && !do_causal_;
 
   std::optional<array> mask_arr;
   if (has_arr_mask) {
     mask_arr = prepare_sdpa_input(inputs[3], s);
   }
+  std::optional<array> stats;
+  if (output_logsumexp_) {
+    stats = outputs[1];
+  }
 
   if (supports_sdpa_vector(
           q, k, v, has_mask, has_arr_mask, do_causal_, output_logsumexp_)) {

python/tests/test_fast_sdpa.py

@@ -771,20 +771,6 @@ def test_grad(slow, fast, args):
 
             self.assertTrue(mx.allclose(g1, g2, **tolerance))
 
-        sdpa_mask_slow = lambda q, k, v, mask: mlx_ref_attn(
-            q, k, v, scale=scale, mask=mask
-        )
-        sdpa_mask_fast = lambda q, k, v, mask: mx.fast.scaled_dot_product_attention(
-            q, k, v, scale=scale, mask=mask
-        )
-
-        loss_mask_slow = lambda q, k, v, mask: mlx_ref_attn(
-            q, k, v, scale=scale, mask=mask
-        ).sum()
-        loss_mask_fast = lambda q, k, v, mask: (
-            mx.fast.scaled_dot_product_attention(q, k, v, scale=scale, mask=mask)
-        ).sum()
-
         B, N_kv, T, D = (2, 8, 128, 64)
         scale = D**-0.5
 
@@ -796,11 +782,7 @@ def test_grad(slow, fast, args):
             mask_additive = mx.random.normal((B, N_q, T, T), dtype=mx.float16)
             mask_bool = mx.random.uniform(0, 1, (B, N_q, T, T), dtype=mx.float16) < 0.5
 
-            for mask in (mask_additive, mask_bool):
-                test_vjp(sdpa_mask_slow, sdpa_mask_fast, [q, k, v, mask])
-                test_grad(loss_mask_slow, loss_mask_fast, [q, k, v, mask])
-
-            for mask in (None, "causal"):
+            for mask in (None, "causal", mask_additive, mask_bool):
                 sdpa_slow = lambda q, k, v: mlx_ref_attn(
                     q, k, v, scale=scale, mask=mask
                 )

tests/linalg_tests.cpp

@@ -350,7 +350,7 @@ TEST_CASE("test SVD factorization") {
   const auto A_again = matmul(matmul(U_slice, diag(S)), Vt);
 
   CHECK(
-      allclose(A_again, A, /* rtol = */ 1e-4, /* atol = */ 1e-4).item<bool>());
+      allclose(A_again, A, /* rtol = */ 1e-3, /* atol = */ 1e-3).item<bool>());
   CHECK_EQ(U.dtype(), float32);
   CHECK_EQ(S.dtype(), float32);
   CHECK_EQ(Vt.dtype(), float32);