From 8ef539522c089a0465e170d19bd96cd97accdef2 Mon Sep 17 00:00:00 2001 From: Cheng Date: Fri, 30 Jan 2026 17:50:18 +0900 Subject: [PATCH] Fix failing python tests on Windows (#3076) --- mlx/CMakeLists.txt | 31 ++++++++++--------- mlx/array.h | 4 +-- mlx/backend/cpu/device_info.cpp | 2 -- mlx/backend/cuda/allocator.cpp | 18 ++++++----- mlx/backend/cuda/allocator.h | 2 +- .../cuda/gemms/grouped_gemm_unaligned.cu | 4 +-- mlx/backend/cuda/quantized/qmv.cu | 24 +++++++------- .../cuda/scaled_dot_product_attention.cpp | 24 +++++++++----- python/tests/test_fast_sdpa.py | 20 +----------- tests/linalg_tests.cpp | 2 +- 10 files changed, 63 insertions(+), 68 deletions(-) diff --git a/mlx/CMakeLists.txt b/mlx/CMakeLists.txt index 82e72a7e..9c0fd388 100644 --- a/mlx/CMakeLists.txt +++ b/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 $<$:/wd4068 - /wd4244 - /wd4267 - /wd4700 - /wd4804> - $<$:-Xcompiler=/wd4068 - -Xcompiler=/wd4244 - -Xcompiler=/wd4267 - -Xcompiler=/wd4700 - -Xcompiler=/wd4804>) + PUBLIC $<$:/wd4244 /wd4267> + PRIVATE $<$:/wd4068 + /wd4146 + /wd4700 + /wd4804 + /wd4805> + $<$:-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( diff --git a/mlx/array.h b/mlx/array.h index 6ea61c1a..60d5e50b 100644 --- a/mlx/array.h +++ b/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 inputs; // An array to keep track of the siblings from a multi-output diff --git a/mlx/backend/cpu/device_info.cpp b/mlx/backend/cpu/device_info.cpp index 7dd8bcb9..b709b2ac 100644 --- a/mlx/backend/cpu/device_info.cpp +++ b/mlx/backend/cpu/device_info.cpp @@ -6,8 +6,6 @@ #include #include #elif defined(_WIN32) -#define WIN32_LEAN_AND_MEAN -#define NOMINMAX #include #else #include diff --git a/mlx/backend/cuda/allocator.cpp b/mlx/backend/cuda/allocator.cpp index d6daee82..b4ea473e 100644 --- a/mlx/backend/cuda/allocator.cpp +++ b/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)); diff --git a/mlx/backend/cuda/allocator.h b/mlx/backend/cuda/allocator.h index 4bbc8429..af76ad90 100644 --- a/mlx/backend/cuda/allocator.h +++ b/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(); diff --git a/mlx/backend/cuda/gemms/grouped_gemm_unaligned.cu b/mlx/backend/cuda/gemms/grouped_gemm_unaligned.cu index 022f2a9c..0c1bb9fe 100644 --- a/mlx/backend/cuda/gemms/grouped_gemm_unaligned.cu +++ b/mlx/backend/cuda/gemms/grouped_gemm_unaligned.cu @@ -124,12 +124,12 @@ struct GemmConfiguration : public CommonGemmConfiguration { }; // Specialized GEMM configuration for sm80 and later. -template +template struct GemmConfiguration< T, Arch, kAlignmentC, - kEnableTF32, + true, std::enable_if_t= 80 && sizeof(T) <= 4>> : public CommonGemmConfiguration { using OpClass = cutlass::arch::OpClassTensorOp; diff --git a/mlx/backend/cuda/quantized/qmv.cu b/mlx/backend/cuda/quantized/qmv.cu index 71d9e7ad..0cc00464 100644 --- a/mlx/backend/cuda/quantized/qmv.cu +++ b/mlx/backend/cuda/quantized/qmv.cu @@ -232,8 +232,8 @@ void fp_qmv( using T = cuda_type_t; if constexpr (!std::is_same_v) { 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(mat); const T* vec_ptr = gpu_ptr(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; + if (!batched.value) { + auto kernel = + fp_qmv_single; if (bits == 8) { - kernel = fp_qmv_single; + kernel = fp_qmv_single; } else if (group_size == 16) { - kernel = fp_qmv_single; + kernel = fp_qmv_single; } encoder.add_kernel_node( kernel, - {static_cast(M), blocks_y}, + {static_cast(M), blocks_y}, block_dims, 0, mat_ptr, @@ -268,15 +269,16 @@ void fp_qmv( N, K); } else { - auto kernel = fp_qmv_batched; + auto kernel = + fp_qmv_batched; if (bits == 8) { - kernel = fp_qmv_batched; + kernel = fp_qmv_batched; } else if (group_size == 16) { - kernel = fp_qmv_batched; + kernel = fp_qmv_batched; } encoder.add_kernel_node( kernel, - {static_cast(M), blocks_y, B}, + {static_cast(M), blocks_y, B}, block_dims, 0, mat_ptr, diff --git a/mlx/backend/cuda/scaled_dot_product_attention.cpp b/mlx/backend/cuda/scaled_dot_product_attention.cpp index d8ff1218..54700bdc 100644 --- a/mlx/backend/cuda/scaled_dot_product_attention.cpp +++ b/mlx/backend/cuda/scaled_dot_product_attention.cpp @@ -140,7 +140,7 @@ DnnGraph build_sdpa_graph( const std::optional& mask_arr, bool output_logsumexp, const array& o, - const array& stats) { + const std::optional& 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& stats, bool do_causal, const std::optional& 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(*mask_arr); } if (output_logsumexp) { - variant_pack[STATS] = gpu_ptr(stats); + variant_pack[STATS] = gpu_ptr(*stats); } CHECK_CUDNN_FE_ERROR(graph.encode_graph(encoder, std::move(variant_pack))); @@ -420,8 +425,7 @@ 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_; @@ -429,6 +433,10 @@ void ScaledDotProductAttention::eval_gpu( if (has_arr_mask) { mask_arr = prepare_sdpa_input(inputs[3], s); } + std::optional stats; + if (output_logsumexp_) { + stats = outputs[1]; + } if (supports_sdpa_vector( q, k, v, has_mask, has_arr_mask, do_causal_, output_logsumexp_)) { diff --git a/python/tests/test_fast_sdpa.py b/python/tests/test_fast_sdpa.py index 715e4881..fa6d0398 100644 --- a/python/tests/test_fast_sdpa.py +++ b/python/tests/test_fast_sdpa.py @@ -771,20 +771,6 @@ class TestSDPA(mlx_tests.MLXTestCase): 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 @@ class TestSDPA(mlx_tests.MLXTestCase): 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 ) diff --git a/tests/linalg_tests.cpp b/tests/linalg_tests.cpp index 683e7584..591f6910 100644 --- a/tests/linalg_tests.cpp +++ b/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()); + allclose(A_again, A, /* rtol = */ 1e-3, /* atol = */ 1e-3).item()); CHECK_EQ(U.dtype(), float32); CHECK_EQ(S.dtype(), float32); CHECK_EQ(Vt.dtype(), float32);