Skip to content

[WIP] Load CUDA external data through pinned buffers - #32437

Open
Xavier Dupré (xadupre) wants to merge 5 commits into
mainfrom
perf/cuda-pinned-initializer-staging
Open

[WIP] Load CUDA external data through pinned buffers#32437
Xavier Dupré (xadupre) wants to merge 5 commits into
mainfrom
perf/cuda-pinned-initializer-staging

Conversation

@xadupre

@xadupre Xavier Dupré (xadupre) commented Sep 4, 2026

Copy link
Copy Markdown
Member

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 -> GPU path. 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::LoadTensor is 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 InferenceSession creation with targeted page-cache eviction.

Configuration

session.cuda.external_data_loader_reading_threads controls how many independent
CPU 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 1 uses
one 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

  • One loader-wide mutex protects both pinned buffers and both CUDA streams for an
    entire initializer load. Initializers therefore cannot concurrently reuse the
    same staging resources.
  • Reader tasks do not take a shared mutex: every task writes to a distinct,
    non-overlapping range of the active pinned buffer.
  • Joining all reader futures is the barrier that guarantees the active buffer is
    completely filled before its H2D copy is submitted.
  • Each buffer has its own CUDA stream. The stream is synchronized before that
    buffer is reused, so CPU readers never overwrite memory still consumed by DMA.
  • Both streams are drained before a successful return and on read or copy errors.

Data path

The complete cold-cache path is:

                         External-data file on local NVMe
                                      |
                                      | NVMe DMA / block I/O
                                      v
                         Linux kernel page-cache pages
                                      |
                                      | CPU copies performed by read()
                                      | 4 independent reads per 64M block by default
                                      | measured NVMe -> pinned: ~4.6 GB/s
                                      |
                    +-----------------+-----------------+
                    |                                   |
                    v                                   v
         +-----------------------+           +-----------------------+
         | pinned buffer 0, 64M  |           | pinned buffer 1, 64M  |
         | CPU fills block N     |           | CPU fills block N + 1 |
         +-----------------------+           +-----------------------+
                    |                                   |
                    | cudaMemcpyAsync                   | cudaMemcpyAsync
                    | pinned -> VRAM: ~55.4 GB/s        | pinned -> VRAM: ~55.4 GB/s
                    |                                   |
                    +-----------------+-----------------+
                                      |
                                      v
                      CUDA BFC Arena initializer buffer
                                      |
                                      | optional CUDA prepack,
                                      | transpose and unpack kernels
                                      | device reads + device writes
                                      v
                          Final prepared CUDA weights

  Timeline:

    CPU reads block N + 1 into buffer 1
           || concurrently with
    PCIe DMA transfers block N from buffer 0

    CPU reads block N + 2 into buffer 0
           || concurrently with
    PCIe DMA transfers block N + 1 from buffer 1

There is no intermediate mmap-backed pageable tensor and no
mmap -> pinned user-space copy. Standard buffered file I/O still necessarily
copies 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:

  External-data file on NVMe
              |
              | page faults / block I/O
              v
  Linux page cache + mmap-backed pageable CPU tensor
              |
              | synchronous cudaMemcpy from pageable memory
              | CUDA driver-managed internal pinned staging
              v
  CUDA BFC Arena initializer buffer
              |
              | optional CUDA prepack/transpose/unpack
              v
  Final prepared CUDA weights

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 explicit
asynchronous H2D copies into the same CUDA BFC Arena destination.

Cold-cache benchmark

Model:

/mnt/nvme/xadupre/models/qwen/qwen3.5-35b-cuda-int4/model.onnx

  • 20.895 GB external-data file
  • NVIDIA H200, GPU 4
  • 96 intra-op threads, spinning disabled
  • fresh process for every run
  • targeted POSIX_FADV_DONTNEED for both model files before every run
  • local eight-disk NVMe volume
Version Run 1 Run 2 Run 3 Mean
main (d47fd8824a) 10.05 s 11.61 s 10.22 s 10.63 s
Direct pinned loader 6.64 s 7.30 s 7.27 s 7.07 s

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

  • Built onnxruntime_pybind11_state with CUDA.
  • Loaded the Qwen3.5 35B CUDA INT4 model with the CUDA Execution Provider.
  • Checked C++ and Python formatting.

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.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Copilot AI balanced review requested due to automatic review settings September 4, 2026 13:13

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 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.

Comment on lines +139 to +140
CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(dst_bytes + offset, state->buffers[staging_index],
chunk_size, cudaMemcpyHostToDevice, stream));
Comment on lines +196 to +197
ORT_RETURN_IF_ERROR(
CopyHostToDeviceWithPinnedStaging(src_data, dst_data, bytes, dst_device.Id()));
Comment on lines +75 to +80
elapsed = time.perf_counter() - start

print(
json.dumps(
{
"active_providers": session.get_providers(),
Comment on lines +129 to +131
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_;
@xadupre Xavier Dupré (xadupre) changed the title Stage large CUDA uploads through pinned buffers [WIP] Stage large CUDA uploads through pinned buffers Sep 4, 2026
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
@xadupre Xavier Dupré (xadupre) changed the title [WIP] Stage large CUDA uploads through pinned buffers Load CUDA external data through pinned buffers Sep 4, 2026
@xadupre

Copy link
Copy Markdown
Member Author

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:

  • Qwen3.5 35B CUDA INT4 model with 20.895 GB of external data
  • local eight-disk NVMe volume
  • NVIDIA H200, GPU 4
  • 96 intra-op threads, spinning disabled
  • fresh process for every run
  • targeted POSIX_FADV_DONTNEED before every run
Version Run 1 Run 2 Run 3 Mean
main (d47fd8824a) 10.05 s 11.61 s 10.22 s 10.63 s
Direct pinned loader 6.64 s 7.30 s 7.27 s 7.07 s

This is a 33.5% reduction in complete cold-cache InferenceSession creation time.

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>
@xadupre Xavier Dupré (xadupre) changed the title Load CUDA external data through pinned buffers [WIP] Load CUDA external data through pinned buffers Sep 4, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants