Skip to content

fix(cuda): reach pinned-DMA bandwidth for large host-to-device uploads#1404

Open
matteospanio wants to merge 2 commits into
tracel-ai:mainfrom
matteospanio:fix/create-with-data-size-guard
Open

fix(cuda): reach pinned-DMA bandwidth for large host-to-device uploads#1404
matteospanio wants to merge 2 commits into
tracel-ai:mainfrom
matteospanio:fix/create-with-data-size-guard

Conversation

@matteospanio

Copy link
Copy Markdown

Problem

Large uploads (create_from_slice / create) reach only a fraction of the device's transfer bandwidth. On an RTX 3070 (PCIe 3.0), a 268 MB f32 upload moves at ~0.9 GB/s, while the read path moves 6.5 GB/s on the same box — a 7× H2D/D2H asymmetry.

Cause: in write_to_gpu, contiguous transfers ≥ 100 MB from unpinned host memory go straight to memcpy_htod_async from pageable memory, which the CUDA driver stages internally and serially. Two smaller costs compound it: create_from_slices re-copies its (already owned) Vec into a fresh Vec before wrapping it into Bytes, and create_with_data page-locks a full-transfer-size staging buffer, bypassing the 100 MB size policy that reserve_cpu already encodes.

Fix

  1. write_to_gpu: large contiguous transfers from unpinned memory are pipelined through two pooled 8 MB pinned staging chunks — each chunk's host memcpy overlaps the previous chunk's DMA. Falls back to the direct path when the pinned pool cannot serve the chunks. Behavior for < 100 MB and already-pinned buffers is unchanged.
  2. create_from_slices: drop the redundant full-size host copy (the slice data is already an owned Vec at that point).
  3. create_with_data: route staging through reserve_cpu so the existing size policy applies.

Measurements (RTX 3070, PCIe 3.0, 268 MB f32)

before after
create(Bytes::from_shared(..)) + sync 298 ms (0.9 GB/s) 42.3 ms (6.35 GB/s)
read (unchanged, for reference) 41 ms (6.5 GB/s) 41 ms (6.5 GB/s)

Measured with the same patch backported onto v0.10.0 (identical code in this area). Byte-exact roundtrips verified at 4 KB, 8 MiB, 8 MiB+1, 16 MiB−3, 100 MiB+4 KiB+1, and 268 MB via both create_from_slice and create(Bytes::from_shared(..)); also validated by a downstream GPU DSP test suite (bit-parity of kernel results vs CPU reference).

🤖 Generated with Claude Code

Three related changes on the upload path:

1. write_to_gpu: contiguous transfers >= 100 MB from unpinned host
   memory were handed to memcpy_htod_async from pageable memory, which
   the driver stages internally and serially at a fraction of pinned
   DMA bandwidth. They are now pipelined through two pooled 8 MB pinned
   staging chunks, each chunk's host memcpy overlapping the previous
   chunk's DMA (falls back to the direct path when the pinned pool
   cannot serve the chunks). Behavior for < 100 MB and already-pinned
   buffers is unchanged.

2. create_from_slices: every slice was re-copied into a fresh Vec
   before being wrapped into Bytes, although it is already an owned
   Vec at that point - a full-size host copy per upload.

3. create_with_data: route the staging allocation through reserve_cpu
   so the existing 100 MB pinned-pool size policy applies instead of
   page-locking a full-transfer-size staging buffer.

Measured on an RTX 3070 (PCIe 3.0), 268 MB f32 upload (cubecl 0.10.0
+ this patch, backported):
- create(Bytes::from_shared(..)) + sync: 298 ms -> 42.3 ms
  (0.9 -> 6.35 GB/s, now symmetric with the 6.5 GB/s read path).
- Byte-exact roundtrips verified at 4 KB, 8 MiB, 8 MiB + 1,
  16 MiB - 3, 100 MiB + 4 KiB + 1, and 268 MB, via both
  create_from_slice and create(Bytes::from_shared(..)).
@nathanielsimard

Copy link
Copy Markdown
Member

I'm not sure the current fix is the right one. We already haven a staging utility exposed from the client, I think we should use that instead, since that will improve performance across all backends.

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