diff --git a/.github/workflows/publish.yaml b/.github/workflows/publish.yaml index e0a5a6f..e003fba 100644 --- a/.github/workflows/publish.yaml +++ b/.github/workflows/publish.yaml @@ -36,6 +36,26 @@ jobs: with: python-version: "3.11" + # The wheel is built with the image's newest MSVC toolset, but the + # image's installed VC redistributable can lag behind it (14.40 vs 14.51 + # on windows-2025-vs2026), so delvewheel would vendor an msvcp140.dll + # older than the toolset that built the .pyd. Install the redist that + # ships inside the image's own Visual Studio, which matches the toolset. + - name: Update VC++ redistributable to match the MSVC toolset (Windows) + if: runner.os == 'Windows' + shell: pwsh + run: | + $vswhere = "${env:ProgramFiles(x86)}\Microsoft Visual Studio\Installer\vswhere.exe" + $vs = & $vswhere -latest -products * -property installationPath + if (-not $vs) { throw "Visual Studio installation not found" } + $redist = Get-ChildItem "$vs\VC\Redist\MSVC" -Recurse -Filter vc_redist.x64.exe | + Sort-Object { $_.VersionInfo.FileVersionRaw } | Select-Object -Last 1 + if (-not $redist) { throw "vc_redist.x64.exe not found under $vs" } + Write-Host "Installing $($redist.FullName) ($($redist.VersionInfo.FileVersion))" + $p = Start-Process -FilePath $redist.FullName -ArgumentList '/install','/quiet','/norestart' -Wait -PassThru + # 1638: same or newer version already installed; 3010: success, reboot required + if ($p.ExitCode -notin 0, 1638, 3010) { throw "vc_redist.x64.exe exited with $($p.ExitCode)" } + - name: Install cibuildwheel run: python -m pip install cibuildwheel setuptools pybind11 - name: Build wheels with cibuildwheel @@ -46,6 +66,19 @@ jobs: CIBW_SKIP: "*-musllinux_* *-win32 *-manylinux_i686" CIBW_TEST_COMMAND: "python {project}/tests/smoke_import_cpp.py" CIBW_ARCHS: "${{ matrix.platform.cibw_arch }}" + # Bundle MSVC runtime satellite DLLs into the wheel; the VS2026 runner + # image builds against a toolchain newer than what end-user machines + # (and the test env) may have installed. + CIBW_BEFORE_BUILD_WINDOWS: "pip install delvewheel" + # --add-path makes delvewheel take msvcp140.dll from System32 (kept + # current by the redist step above) instead of a stale copy that some + # other tool ships earlier on PATH. + CIBW_REPAIR_WHEEL_COMMAND_WINDOWS: "delvewheel repair --add-path C:/Windows/System32 -w {dest_dir} {wheel}" + # Diagnostic: print the .pyd import table and where each DLL resolves + # from, so a load failure names the missing DLL instead of Windows' + # generic "module could not be found". + CIBW_TEST_REQUIRES_WINDOWS: "pefile" + CIBW_TEST_COMMAND_WINDOWS: "python {project}/tests/dump_pyd_imports.py && python {project}/tests/smoke_import_cpp.py" - name: Upload wheel artifact uses: actions/upload-artifact@v4 diff --git a/README.md b/README.md index 778d5fc..3bab297 100644 --- a/README.md +++ b/README.md @@ -3,7 +3,7 @@ fastsafetensors fastsafetensors is an efficient safetensors loader. If you develop your own code that loads large safetensors files, you can try fastsafetensors APIs (see [docs](./docs/overview.md)). For example, vLLM and SGLang have `--load-format fastsafetensors` command-line argument to speed up their initialization. -This library supports Linux/CUDA, ROCm without GDS, Windows, [3FS](https://github.com/deepseek-ai/3fs), unified-memory systems such as DGX Spark, and so on. We welcome more platform/storage-specific optimizations like them by adding new [copier backends](fastsafetensors/copier/). Our CI tests Python 3.10-3.14 with PyTorch 2.11.0. +This library supports Linux/CUDA, ROCm, Windows, [3FS](https://github.com/deepseek-ai/3fs), unified-memory systems such as DGX Spark, and so on. We welcome more platform/storage-specific optimizations like them by adding new [copier backends](fastsafetensors/copier/). Our CI tests Python 3.10-3.14 with PyTorch 2.11.0. # Performance Highlights diff --git a/docs/development.md b/docs/development.md index d842a31..99ef74c 100644 --- a/docs/development.md +++ b/docs/development.md @@ -25,7 +25,7 @@ You can also use Makefile on your local environment. ``` make unittest make unittest-parallel -make vllm +make test-vllm ``` # Pre-commit Hooks diff --git a/docs/overview.md b/docs/overview.md index b14d70b..455882c 100644 --- a/docs/overview.md +++ b/docs/overview.md @@ -24,7 +24,7 @@ The technology helps minimize copy overheads from NVMe SSDs to GPU memory by byp # Basic API usage -`SafeTensorsFileLoader` is a low-level entrypoint. To use it, pass either `SingleGroup()` for simple inference or `ProcessGroup()` (from `torch.distributed`) for tensor-parallel inference. The loader supports both CPU and CUDA devices, with optional GPU Direct Storage (GDS) support. You can specify the device and GDS settings using the `device` and `nogds` arguments, respectively. Note that if GDS is not available, the loader will fail to open files when `nogds=False`. For more information on enabling GDS, please refer to the NVIDIA documentation. +`SafeTensorsFileLoader` is a low-level entrypoint. To use it, pass either `SingleGroup()` for simple inference or `ProcessGroup()` (from `torch.distributed`) for tensor-parallel inference. The loader supports both CPU and CUDA devices, with optional GPU Direct Storage (GDS) support. You can specify the device and GDS settings using the `device` and `nogds` arguments, respectively. If GDS turns out to be unavailable at runtime (e.g., file handle registration fails), the loader logs a warning and falls back to the bounce-buffer (`nogds`) path instead of failing; you can also set `nogds=True` explicitly to skip GDS initialization. For more information on enabling GDS, please refer to the NVIDIA documentation. After creating a `SafeTensorsFileLoader` instance, first map target files and a rank using the `.add_filenames()` method. Then, call `.copy_file_to_device()` to trigger the actual file copies on aggregated GPU memory fragments and directly instantiate a group of tensors. Once the files are loaded, you can retrieve a tensor using the `.get_tensor()` method. Additionally, you can obtain sharded tensors by `.get_sharded()`, which internally runs collective operations in `torch.distributed`. @@ -45,8 +45,8 @@ See [Configuration Guide](./configuration.md) for defaults, examples, and all av # ROCm -On ROCm, there is no GDS-equivalent support, so fastsafetensors only supports `nogds=True` mode. -The performance gain example can be found at [amd-perf.md](./amd-perf.md). +On ROCm, direct storage-to-GPU loading is supported through hipFile (ROCm >= 7.2): when `libhipfile.so` is available, the GDS code path uses it transparently. On older ROCm without hipFile, the loader falls back to the bounce-buffer (`nogds`) path. +A performance gain example with the `nogds` path can be found at [amd-perf.md](./amd-perf.md). # Windows diff --git a/fastsafetensors/cpp/dstorage_reader.cpp b/fastsafetensors/cpp/dstorage_reader.cpp index cc70979..80b3465 100644 --- a/fastsafetensors/cpp/dstorage_reader.cpp +++ b/fastsafetensors/cpp/dstorage_reader.cpp @@ -13,12 +13,24 @@ namespace py = pybind11; -#pragma comment(lib, "d3d12.lib") -#pragma comment(lib, "dxgi.lib") -#pragma comment(lib, "dxguid.lib") - static constexpr UINT32 DS_STAGING_BUFFER_BYTES = 256u * 1024u * 1024u; +static const GUID FST_IID_IDXGIFactory1 = { + 0x770aae78, 0xf26f, 0x4dba, + {0xa8, 0x29, 0x25, 0x3c, 0x83, 0xd1, 0xb3, 0x87} +}; +static const GUID FST_IID_ID3D12Device = { + 0x189819f1, 0x1db6, 0x4b57, + {0xbe, 0x54, 0x18, 0x21, 0x33, 0x9b, 0x85, 0xf7} +}; +static const GUID FST_IID_ID3D12Fence = { + 0x0a753dcf, 0xc4d8, 0x4b91, + {0xad, 0xf6, 0xbe, 0x5a, 0x60, 0xd9, 0x5a, 0x76} +}; +static const GUID FST_IID_ID3D12Resource = { + 0x696442be, 0xa72e, 0x4059, + {0xbc, 0x79, 0x5b, 0x5c, 0x98, 0x04, 0x0f, 0xad} +}; static const GUID IID_IDStorageFactory = { 0x6924ea0c, 0xc3cd, 0x4826, {0xb1, 0x0a, 0xf6, 0x4f, 0x4e, 0xd9, 0x27, 0xc1} @@ -39,6 +51,16 @@ static PFN_DStorageGetFactory g_pfnGetFactory = nullptr; typedef HRESULT (__stdcall *PFN_DStorageSetConfiguration1)(DSTORAGE_CONFIGURATION1 const*); static PFN_DStorageSetConfiguration1 g_pfnSetConfig1 = nullptr; +typedef HRESULT (WINAPI *PFN_CreateDXGIFactory1)(REFIID riid, void** ppvFactory); +static PFN_CreateDXGIFactory1 g_pfnCreateDXGIFactory1 = nullptr; + +typedef HRESULT (WINAPI *PFN_D3D12CreateDevice)( + IUnknown* pAdapter, + D3D_FEATURE_LEVEL minimumFeatureLevel, + REFIID riid, + void** ppDevice); +static PFN_D3D12CreateDevice g_pfnD3D12CreateDevice = nullptr; + static std::wstring utf8_to_wstring(const std::string& input) { int wlen = MultiByteToWideChar(CP_UTF8, 0, input.c_str(), -1, nullptr, 0); if (wlen <= 0) { @@ -89,6 +111,42 @@ static bool LoadDirectStorage(const std::string& dll_dir_utf8) { return g_pfnGetFactory != nullptr; } +static bool LoadD3D12AndDXGI(std::string& error) { + static HMODULE hD3D12 = nullptr; + static HMODULE hDXGI = nullptr; + + if (!hD3D12) { + hD3D12 = LoadLibraryExW(L"d3d12.dll", nullptr, LOAD_LIBRARY_SEARCH_SYSTEM32); + if (!hD3D12) { + error = "Failed to load d3d12.dll: " + std::to_string(GetLastError()); + return false; + } + } + if (!hDXGI) { + hDXGI = LoadLibraryExW(L"dxgi.dll", nullptr, LOAD_LIBRARY_SEARCH_SYSTEM32); + if (!hDXGI) { + error = "Failed to load dxgi.dll: " + std::to_string(GetLastError()); + return false; + } + } + + g_pfnD3D12CreateDevice = reinterpret_cast( + GetProcAddress(hD3D12, "D3D12CreateDevice")); + if (!g_pfnD3D12CreateDevice) { + error = "d3d12.dll does not export D3D12CreateDevice"; + return false; + } + + g_pfnCreateDXGIFactory1 = reinterpret_cast( + GetProcAddress(hDXGI, "CreateDXGIFactory1")); + if (!g_pfnCreateDXGIFactory1) { + error = "dxgi.dll does not export CreateDXGIFactory1"; + return false; + } + + return true; +} + // Global state, D3D12 device + DirectStorage factory class GlobalDStorageState { static inline int s_device_id = 0; @@ -121,8 +179,10 @@ class GlobalDStorageState { s_device = reinterpret_cast(provided_device); s_device->AddRef(); } else { + if (!LoadD3D12AndDXGI(last_error_)) return false; + IDXGIFactory1* factory = nullptr; - HRESULT hr = CreateDXGIFactory1(IID_IDXGIFactory1, (void**)&factory); + HRESULT hr = g_pfnCreateDXGIFactory1(FST_IID_IDXGIFactory1, (void**)&factory); if (FAILED(hr)) { last_error_ = "CreateDXGIFactory1 failed"; return false; } IDXGIAdapter1* adapter = nullptr; @@ -134,7 +194,7 @@ class GlobalDStorageState { } factory->Release(); if (!adapter) { last_error_ = "No hardware D3D12 adapter found"; return false; } - HRESULT hr2 = D3D12CreateDevice(adapter, D3D_FEATURE_LEVEL_11_0, IID_ID3D12Device, (void**)&s_device); + HRESULT hr2 = g_pfnD3D12CreateDevice(adapter, D3D_FEATURE_LEVEL_11_0, FST_IID_ID3D12Device, (void**)&s_device); adapter->Release(); if (FAILED(hr2)) { last_error_ = "D3D12CreateDevice failed"; return false; } } @@ -235,7 +295,7 @@ class dstorage_stream_reader { } hr = dev->CreateFence(0, D3D12_FENCE_FLAG_NONE, - IID_ID3D12Fence, (void**)&fence_); + FST_IID_ID3D12Fence, (void**)&fence_); if (FAILED(hr)) { fprintf(stderr, "dstorage_stream_reader: CreateFence failed hr=0x%08X\n", (unsigned)hr); return; @@ -257,7 +317,7 @@ class dstorage_stream_reader { hr = dev->CreateCommittedResource( &hp, D3D12_HEAP_FLAG_SHARED, &desc, D3D12_RESOURCE_STATE_COMMON, nullptr, - IID_ID3D12Resource, (void**)&stage_res_[i]); + FST_IID_ID3D12Resource, (void**)&stage_res_[i]); if (FAILED(hr)) { fprintf(stderr, "dstorage_stream_reader: CreateCommittedResource[%d] failed hr=0x%08X\n", i, (unsigned)hr); return; diff --git a/pyproject.toml b/pyproject.toml index dc0ced3..1d45d1a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "fastsafetensors" -version = "0.3.2" +version = "0.3.3" description = "High-performance safetensors model loader" authors = [{name = "Takeshi Yoshimura", email = "tyos@jp.ibm.com"}] maintainers = [{name = "Takeshi Yoshimura", email = "tyos@jp.ibm.com"}] diff --git a/setup.py b/setup.py index 0b65210..8620f72 100644 --- a/setup.py +++ b/setup.py @@ -26,8 +26,8 @@ def MyExtension(name, sources, mod_name, *args, **kwargs): kwargs["libraries"] = [] # c++20 required for designated initializers at ext.hpp kwargs["extra_compile_args"] = ["/std:c++20"] - # Note: dstorage.dll is loaded at runtime via LoadLibrary, not linked. - kwargs["libraries"].extend(["ole32", "d3d12", "dxgi", "dxguid", "uuid"]) + # DirectStorage, D3D12, and DXGI DLLs are loaded at runtime so importing + # the extension does not require GPU/DirectX runtime DLLs to be present. # CUDA interop headers: if CUDA_HOME/CUDA_PATH is set, add include path # for cudaExternalMemory types used by the interop bridge. diff --git a/tests/dump_pyd_imports.py b/tests/dump_pyd_imports.py new file mode 100644 index 0000000..2ad319b --- /dev/null +++ b/tests/dump_pyd_imports.py @@ -0,0 +1,99 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""Diagnostic run before the Windows wheel smoke test in CI. + +Prints the DLL import table of the installed fastsafetensors.cpp extension and +where each imported DLL resolves from, so an import failure names the missing +DLL instead of Windows' generic "module could not be found". Also reports the +MSVC runtime DLLs present in System32 and their versions. Requires pefile. +""" + +import importlib.machinery +import importlib.metadata +import os +import sys +from pathlib import Path + + +def _find_pyd(): + dist = importlib.metadata.distribution("fastsafetensors") + dist_root = Path(dist.locate_file("")) + for suffix in importlib.machinery.EXTENSION_SUFFIXES: + candidate = dist_root / "fastsafetensors" / f"cpp{suffix}" + if candidate.exists(): + return dist_root, candidate + raise FileNotFoundError( + f"fastsafetensors.cpp extension not found under {dist_root}" + ) + + +def main() -> None: + if os.name != "nt": + print("dump_pyd_imports: not Windows, nothing to do") + return + + import pefile + + dist_root, pyd = _find_pyd() + print(f"extension: {pyd}") + + search_dirs = { + "System32": Path(os.environ.get("SystemRoot", r"C:\Windows")) / "System32", + "beside pyd": pyd.parent, + "python dir": Path(sys.base_exec_prefix), + } + for libs_dir in dist_root.glob("*.libs"): + search_dirs[libs_dir.name] = libs_dir + + pe = pefile.PE(str(pyd), fast_load=True) + pe.parse_data_directories( + directories=[pefile.DIRECTORY_ENTRY["IMAGE_DIRECTORY_ENTRY_IMPORT"]] + ) + for entry in getattr(pe, "DIRECTORY_ENTRY_IMPORT", []): + name = entry.dll.decode() + # API set names (api-ms-*, ext-ms-*) are resolved virtually by the OS + # loader, not by file lookup, so absence on disk is not a failure. + if name.lower().startswith(("api-ms-", "ext-ms-")): + print(f" imports {name}: API set (loader-resolved)") + continue + found_in = [label for label, d in search_dirs.items() if (d / name).exists()] + status = ", ".join(found_in) if found_in else "*** NOT FOUND ***" + print(f" imports {name}: {status}") + + # The failing wheel was never inspected before delvewheel repair, so also + # report which MSVC CRT DLLs the image itself provides and their versions; + # this identifies what the unrepaired .pyd could(n't) have resolved. + print("MSVC runtime DLLs in System32:") + crt_names = [ + "msvcp140.dll", + "msvcp140_1.dll", + "msvcp140_2.dll", + "msvcp140_atomic_wait.dll", + "msvcp140_codecvt_ids.dll", + "vcruntime140.dll", + "vcruntime140_1.dll", + "concrt140.dll", + ] + for crt_name in crt_names: + path = search_dirs["System32"] / crt_name + if path.exists(): + print(f" {crt_name}: {_file_version(pefile, path)}") + else: + print(f" {crt_name}: *** NOT PRESENT ***") + + +def _file_version(pefile, path): + try: + pe = pefile.PE(str(path), fast_load=True) + pe.parse_data_directories( + directories=[pefile.DIRECTORY_ENTRY["IMAGE_DIRECTORY_ENTRY_RESOURCE"]] + ) + info = pe.VS_FIXEDFILEINFO[0] + ms, ls = info.FileVersionMS, info.FileVersionLS + return f"{ms >> 16}.{ms & 0xFFFF}.{ls >> 16}.{ls & 0xFFFF}" + except Exception as e: + return f"present (version unreadable: {e})" + + +if __name__ == "__main__": + main() diff --git a/tests/smoke_import_cpp.py b/tests/smoke_import_cpp.py index 8e8abd5..8c59bee 100644 --- a/tests/smoke_import_cpp.py +++ b/tests/smoke_import_cpp.py @@ -5,13 +5,25 @@ import importlib.machinery import importlib.metadata import importlib.util +import os import sys from pathlib import Path +def _add_vendored_dll_dirs(dist_root: Path) -> None: + # delvewheel vendors dependent DLLs into .libs and registers that + # directory via a patch in the package __init__. This test bypasses + # __init__ on purpose, so register the directory here. + if sys.platform != "win32": + return + for libs_dir in dist_root.glob("*.libs"): + os.add_dll_directory(str(libs_dir)) + + def _load_cpp_extension(): dist = importlib.metadata.distribution("fastsafetensors") dist_root = Path(dist.locate_file("")) + _add_vendored_dll_dirs(dist_root) for suffix in importlib.machinery.EXTENSION_SUFFIXES: candidate = dist_root / "fastsafetensors" / f"cpp{suffix}"