From cb47e8d777c552e23d45a335b428a83f954ead5f Mon Sep 17 00:00:00 2001 From: Daniel Hiltgen Date: Mon, 1 Jun 2026 09:23:26 -0700 Subject: [PATCH 1/2] Use 64-bit file seeks on Windows --- mlx/io/load.h | 34 +++++++++--- tests/load_tests.cpp | 120 +++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 147 insertions(+), 7 deletions(-) diff --git a/mlx/io/load.h b/mlx/io/load.h index 70842cac9d..6b9c9f477b 100644 --- a/mlx/io/load.h +++ b/mlx/io/load.h @@ -8,7 +8,7 @@ #include #include -#ifdef _MSC_VER +#ifdef _WIN32 #include #else #include @@ -79,20 +79,30 @@ class ParallelFileReader : public Reader { } size_t tell() override { +#ifdef _WIN32 + return _lseeki64(fd_, 0, SEEK_CUR); +#else return lseek(fd_, 0, SEEK_CUR); +#endif } // Warning: do not use this function from multiple threads as // it advances the file descriptor void seek(int64_t off, std::ios_base::seekdir way = std::ios_base::beg) override { + int origin; if (way == std::ios_base::beg) { - lseek(fd_, off, SEEK_SET); + origin = SEEK_SET; } else if (way == std::ios_base::end) { - lseek(fd_, off, SEEK_END); + origin = SEEK_END; } else { - lseek(fd_, off, SEEK_CUR); + origin = SEEK_CUR; } +#ifdef _WIN32 + _lseeki64(fd_, off, origin); +#else + lseek(fd_, off, origin); +#endif } // Warning: do not use this function from multiple threads as @@ -159,19 +169,29 @@ class FileWriter : public Writer { size_t tell() override { check_open(); +#ifdef _WIN32 + return _lseeki64(fd_, 0, SEEK_CUR); +#else return lseek(fd_, 0, SEEK_CUR); +#endif } void seek(int64_t off, std::ios_base::seekdir way = std::ios_base::beg) override { check_open(); + int origin; if (way == std::ios_base::beg) { - lseek(fd_, off, SEEK_SET); + origin = SEEK_SET; } else if (way == std::ios_base::end) { - lseek(fd_, off, SEEK_END); + origin = SEEK_END; } else { - lseek(fd_, off, SEEK_CUR); + origin = SEEK_CUR; } +#ifdef _WIN32 + _lseeki64(fd_, off, origin); +#else + lseek(fd_, off, origin); +#endif } void write(const char* data, size_t n) override { diff --git a/tests/load_tests.cpp b/tests/load_tests.cpp index b4b54046a3..8b1d28c858 100644 --- a/tests/load_tests.cpp +++ b/tests/load_tests.cpp @@ -1,12 +1,21 @@ // Copyright © 2023 Apple Inc. +#include #include #include +#include #include #include +#ifdef _WIN32 +#define NOMINMAX +#include +#include +#endif + #include "doctest/doctest.h" +#include "mlx/io/load.h" #include "mlx/mlx.h" using namespace mlx::core; @@ -654,3 +663,114 @@ TEST_CASE("test single array serialization") { CHECK(array_equal(a, b).item()); } } + +namespace { +struct SparseSafetensorsFile { + std::filesystem::path path = std::filesystem::temp_directory_path() / + ("mlx-large-" + + std::to_string( + std::chrono::steady_clock::now().time_since_epoch().count()) + + ".safetensors"); + + ~SparseSafetensorsFile() { + std::error_code error; + std::filesystem::remove(path, error); + } + + void write(size_t rows, const std::vector& payload) { + const size_t padding = rows * 65536; + std::ofstream out(path, std::ios::binary); + REQUIRE(out.good()); +#ifdef _WIN32 + // Windows needs this flag to leave the zero-filled tensor unallocated. + auto handle = CreateFileW( + path.c_str(), + GENERIC_WRITE, + FILE_SHARE_READ | FILE_SHARE_WRITE, + nullptr, + OPEN_EXISTING, + 0, + nullptr); + REQUIRE(handle != INVALID_HANDLE_VALUE); + DWORD returned; + auto sparse = DeviceIoControl( + handle, FSCTL_SET_SPARSE, nullptr, 0, nullptr, 0, &returned, nullptr); + CloseHandle(handle); + REQUIRE(sparse); +#endif + std::ostringstream json; + json << R"({"first":{"dtype":"U8","shape":[4],"data_offsets":[0,4]},)" + << R"("padding":{"dtype":"U8","shape":[)" << rows + << R"(,65536],"data_offsets":[4,)" << padding + 4 << R"(]},)" + << R"("last":{"dtype":"U8","shape":[)" << payload.size() + << R"(],"data_offsets":[)" << padding + 4 << ',' + << padding + 4 + payload.size() << "]}}"; + auto header = json.str(); + header.append((8 - header.size() % 8) % 8, ' '); + uint64_t length = header.size(); + out.write(reinterpret_cast(&length), 8); + out.write(header.data(), header.size()); + out.write("abcd", 4); + out.seekp(padding, std::ios::cur); + out.write(reinterpret_cast(payload.data()), payload.size()); + out.close(); + REQUIRE(out.good()); + REQUIRE( + std::filesystem::file_size(path) == + 8 + header.size() + 4 + padding + payload.size()); + } +}; +} // namespace + +TEST_CASE("test large safetensors load and evaluate") { + size_t rows = 1; + SUBCASE("small control") { + rows = 1; + } + SUBCASE("beyond 2 GiB") { + rows = 49152; + } + SUBCASE("beyond 4 GiB") { + rows = 81920; + } + + // Exceed the parallel reader's 32 MiB batch size. + std::vector payload((33 << 20) + 17); + for (size_t i = 0; i < payload.size(); ++i) { + payload[i] = static_cast((i * 17 + i / (1 << 20)) % 251); + } + SparseSafetensorsFile file; + file.write(rows, payload); + auto [weights, metadata] = load_safetensors(file.path.string()); + REQUIRE(weights.size() == 3); + CHECK(weights.at("padding").nbytes() == rows * 65536); + CHECK(array_equal(weights.at("first"), array({97, 98, 99, 100}, uint8)) + .item()); + CHECK(array_equal( + weights.at("last"), + array(payload.data(), {static_cast(payload.size())}, uint8)) + .item()); +} + +TEST_CASE("test large file writer seeks") { + const size_t offset = size_t{5} << 30; + SparseSafetensorsFile file; + file.write(81920, {11, 22, 33, 44}); + char bytes[4]{}; + io::FileWriter writer(file.path.string()); + writer.open(); + writer.seek(offset); + REQUIRE(writer.tell() == offset); + writer.write("ABCD", 4); + CHECK(writer.tell() == offset + 4); + writer.seek(-2, std::ios::end); + CHECK(writer.tell() == offset + 2); + writer.seek(-1, std::ios::cur); + CHECK(writer.tell() == offset + 1); + writer.write("Z", 1); + std::ifstream check(file.path, std::ios::binary); + check.seekg(offset); + check.read(bytes, 4); + REQUIRE(check.good()); + CHECK(std::string(bytes, 4) == "AZCD"); +} From 5a141534278ab4d0ce24f130dd159672644fa039 Mon Sep 17 00:00:00 2001 From: Cheng Date: Tue, 8 Sep 2026 12:24:34 +0900 Subject: [PATCH 2/2] Remove tests since it is too slow and change is trival --- tests/load_tests.cpp | 120 ------------------------------------------- 1 file changed, 120 deletions(-) diff --git a/tests/load_tests.cpp b/tests/load_tests.cpp index 8b1d28c858..b4b54046a3 100644 --- a/tests/load_tests.cpp +++ b/tests/load_tests.cpp @@ -1,21 +1,12 @@ // Copyright © 2023 Apple Inc. -#include #include #include -#include #include #include -#ifdef _WIN32 -#define NOMINMAX -#include -#include -#endif - #include "doctest/doctest.h" -#include "mlx/io/load.h" #include "mlx/mlx.h" using namespace mlx::core; @@ -663,114 +654,3 @@ TEST_CASE("test single array serialization") { CHECK(array_equal(a, b).item()); } } - -namespace { -struct SparseSafetensorsFile { - std::filesystem::path path = std::filesystem::temp_directory_path() / - ("mlx-large-" + - std::to_string( - std::chrono::steady_clock::now().time_since_epoch().count()) + - ".safetensors"); - - ~SparseSafetensorsFile() { - std::error_code error; - std::filesystem::remove(path, error); - } - - void write(size_t rows, const std::vector& payload) { - const size_t padding = rows * 65536; - std::ofstream out(path, std::ios::binary); - REQUIRE(out.good()); -#ifdef _WIN32 - // Windows needs this flag to leave the zero-filled tensor unallocated. - auto handle = CreateFileW( - path.c_str(), - GENERIC_WRITE, - FILE_SHARE_READ | FILE_SHARE_WRITE, - nullptr, - OPEN_EXISTING, - 0, - nullptr); - REQUIRE(handle != INVALID_HANDLE_VALUE); - DWORD returned; - auto sparse = DeviceIoControl( - handle, FSCTL_SET_SPARSE, nullptr, 0, nullptr, 0, &returned, nullptr); - CloseHandle(handle); - REQUIRE(sparse); -#endif - std::ostringstream json; - json << R"({"first":{"dtype":"U8","shape":[4],"data_offsets":[0,4]},)" - << R"("padding":{"dtype":"U8","shape":[)" << rows - << R"(,65536],"data_offsets":[4,)" << padding + 4 << R"(]},)" - << R"("last":{"dtype":"U8","shape":[)" << payload.size() - << R"(],"data_offsets":[)" << padding + 4 << ',' - << padding + 4 + payload.size() << "]}}"; - auto header = json.str(); - header.append((8 - header.size() % 8) % 8, ' '); - uint64_t length = header.size(); - out.write(reinterpret_cast(&length), 8); - out.write(header.data(), header.size()); - out.write("abcd", 4); - out.seekp(padding, std::ios::cur); - out.write(reinterpret_cast(payload.data()), payload.size()); - out.close(); - REQUIRE(out.good()); - REQUIRE( - std::filesystem::file_size(path) == - 8 + header.size() + 4 + padding + payload.size()); - } -}; -} // namespace - -TEST_CASE("test large safetensors load and evaluate") { - size_t rows = 1; - SUBCASE("small control") { - rows = 1; - } - SUBCASE("beyond 2 GiB") { - rows = 49152; - } - SUBCASE("beyond 4 GiB") { - rows = 81920; - } - - // Exceed the parallel reader's 32 MiB batch size. - std::vector payload((33 << 20) + 17); - for (size_t i = 0; i < payload.size(); ++i) { - payload[i] = static_cast((i * 17 + i / (1 << 20)) % 251); - } - SparseSafetensorsFile file; - file.write(rows, payload); - auto [weights, metadata] = load_safetensors(file.path.string()); - REQUIRE(weights.size() == 3); - CHECK(weights.at("padding").nbytes() == rows * 65536); - CHECK(array_equal(weights.at("first"), array({97, 98, 99, 100}, uint8)) - .item()); - CHECK(array_equal( - weights.at("last"), - array(payload.data(), {static_cast(payload.size())}, uint8)) - .item()); -} - -TEST_CASE("test large file writer seeks") { - const size_t offset = size_t{5} << 30; - SparseSafetensorsFile file; - file.write(81920, {11, 22, 33, 44}); - char bytes[4]{}; - io::FileWriter writer(file.path.string()); - writer.open(); - writer.seek(offset); - REQUIRE(writer.tell() == offset); - writer.write("ABCD", 4); - CHECK(writer.tell() == offset + 4); - writer.seek(-2, std::ios::end); - CHECK(writer.tell() == offset + 2); - writer.seek(-1, std::ios::cur); - CHECK(writer.tell() == offset + 1); - writer.write("Z", 1); - std::ifstream check(file.path, std::ios::binary); - check.seekg(offset); - check.read(bytes, 4); - REQUIRE(check.good()); - CHECK(std::string(bytes, 4) == "AZCD"); -}