L-electron-Rare/mlx
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix failing python tests on Windows (#3076)

Browse Source
This commit is contained in:
Cheng
2026-01-30 17:50:18 +09:00
committed by GitHub
parent 212077f163
commit 8ef539522c
10 changed files with 63 additions and 68 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mlx/CMakeLists.txt
+16 -15
View File
@@ -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
+2 -2
View File
@@ -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
-2
View File
@@ -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
+11 -7
View File
@@ -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
+1 -1
View File
@@ -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
+2 -2
View File
@@ -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
+13 -11
View File
@@ -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
+16 -8
View File
@@ -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,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<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
+1 -19
View File
@@ -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
)
tests/linalg_tests.cpp
+1 -1
View File
@@ -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);