[WIP] Load CUDA external data through pinned buffers - #32437
[WIP] Load CUDA external data through pinned buffers#32437Xavier Dupré (xadupre) wants to merge 5 commits into
Conversation
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
There was a problem hiding this comment.
🟡 Changes recommended
Error paths can leave DMA in flight, and setup failures regress previously valid transfers.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Pull request overview
Adds pinned-buffer staging to accelerate large synchronous pageable host-to-CUDA transfers.
Changes:
- Alternates 64 MiB chunks across two CUDA streams.
- Retains per-device staging resources.
- Adds a cold model-loading benchmark.
File summaries
| File | Description |
|---|---|
gpu_data_transfer.cc |
Implements staged CUDA transfers. |
gpu_data_transfer.h |
Declares staging state and synchronization. |
benchmark_cuda_model_loading.py |
Benchmarks CUDA session creation. |
Review details
- Files reviewed: 3/3 changed files
- Comments generated: 5
- Review effort level: Balanced
💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(dst_bytes + offset, state->buffers[staging_index], | ||
| chunk_size, cudaMemcpyHostToDevice, stream)); |
| ORT_RETURN_IF_ERROR( | ||
| CopyHostToDeviceWithPinnedStaging(src_data, dst_data, bytes, dst_device.Id())); |
| elapsed = time.perf_counter() - start | ||
|
|
||
| print( | ||
| json.dumps( | ||
| { | ||
| "active_providers": session.get_providers(), |
| for (size_t offset = 0, chunk_index = 0; offset < bytes; ++chunk_index) { | ||
| const size_t staging_index = chunk_index % state->buffers.size(); | ||
| const size_t chunk_size = std::min(kPinnedStagingBufferSize, bytes - offset); |
| void ReleaseAllPinnedStaging() const noexcept; | ||
|
|
||
| mutable std::mutex pinned_staging_mutex_; | ||
| mutable std::unordered_map<int, std::unique_ptr<PinnedStagingState>> pinned_staging_by_device_; |
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
|
Cold-cache CUDA loading now reads external initializers directly into two reusable pinned buffers instead of copying from an mmap-backed pageable tensor. Benchmark configuration:
This is a 33.5% reduction in complete cold-cache The external-data file reads at approximately 4.6 GB/s on this volume. Reading 20.895 GB therefore has an incompressible lower bound of about 4.5 seconds, before CUDA weight preparation and the remaining session initialization. |
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Description
Load large CUDA external initializers directly from their files into two reusable 64 MiB pinned host buffers.
The CUDA Execution Provider now supplies an
IExternalDataLoader. For each block, independent reads fill disjoint ranges of the next pinned buffer while the preceding buffer is transferred asynchronously to the GPU. The two buffers and CUDA streams are retained for the lifetime of the loader and synchronized before reuse and before returning.This avoids the previous
mmap -> pageable CPU memory -> pinned memory -> GPUpath. CPU and other execution providers keep their existing external-data behavior, and ordinary CUDA data transfers keep the existing CUDA-managed pageable-memory staging.IExternalDataLoader::LoadTensoris made pure virtual so its type information is emitted in each provider shared library. All existing external data loaders already implement this method.The PR also includes a standalone benchmark for CUDA
InferenceSessioncreation with targeted page-cache eviction.Configuration
session.cuda.external_data_loader_reading_threadscontrols how many independentCPU read tasks fill each 64 MiB pinned staging buffer. The default is 4, which
was the fastest setting on the benchmarked eight-disk NVMe volume. Each task reads
a disjoint range of the active buffer; once every range is complete, the whole
buffer is submitted to CUDA while the next buffer is filled. A value of
1usesone sequential read per block. Values from 1 through 64 are accepted because the
best value depends on the storage device, filesystem, and host.
The benchmark script exposes the same setting as
--reading-threads.Synchronization and locking
entire initializer load. Initializers therefore cannot concurrently reuse the
same staging resources.
non-overlapping range of the active pinned buffer.
completely filled before its H2D copy is submitted.
buffer is reused, so CPU readers never overwrite memory still consumed by DMA.
Data path
The complete cold-cache path is:
There is no intermediate mmap-backed pageable tensor and no
mmap -> pinneduser-space copy. Standard buffered file I/O still necessarilycopies data from Linux page-cache pages into the pinned user-space buffer.
That CPU copy was not measured independently; the 4.6 GB/s figure covers
the complete cold buffered-read path from the NVMe file into pinned memory.
CUDA then performs one H2D DMA from that pinned buffer into the initializer's
device allocation at approximately 55.4 GB/s in the pinned-memory H2D
microbenchmark. Operators that prepack weights may subsequently read that CUDA
allocation and write a transformed CUDA allocation.
In the benchmarked default configuration, the H2D destination is the
initializer buffer planned and allocated from the CUDA BFC Arena. The new
loader writes each block directly into its final offset in that arena buffer.
This replaces the previous path:
The new path removes the full mmap-backed CPU tensor and the CUDA driver's
implicit pageable-memory staging. It replaces them with controlled parallel
read()calls directly into two persistent pinned buffers followed by explicitasynchronous H2D copies into the same CUDA BFC Arena destination.
Cold-cache benchmark
Model:
/mnt/nvme/xadupre/models/qwen/qwen3.5-35b-cuda-int4/model.onnxPOSIX_FADV_DONTNEEDfor both model files before every runmain(d47fd8824a)The measured cold-load improvement is 33.5%.
The external-data file reads at approximately 4.6 GB/s from this NVMe volume. Reading 20.895 GB therefore has a lower bound of approximately 4.5 seconds before CUDA weight preparation and the remaining session initialization are considered.
Validation
onnxruntime_pybind11_statewith CUDA.TODO: with this design, I still need to find a way to tell the data loader the number of threads it should use and this is defined in the session options.