Problem
GPU Kudo shuffle serialization performs excessive CUDA pinned-host allocation when spark.rapids.memory.pinnedPool.size=0, which is the default.
GpuDeviceManager disables the cuDF default pinned pool. With no plugin pinned pool to replace it, small cuDF scalar-staging buffers used during GPU Kudo assembly fall through to an independent cudaHostAlloc and cudaFreeHost pair.
In an Nsight Systems profile of a pure PySpark DataFrame shuffle workload, each shuffle_assemble invocation caused exactly 73 allocation/free pairs. Six actions containing 384 assembles produced:
- 28,032
cudaHostAlloc calls
- 28,032
cudaFreeHost calls
- 3.469 seconds of aggregate
cudaHostAlloc API time
A larger profile produced 56,064 pairs across six actions.
End-to-end impact
Configuration: Spark 3.5.0, local[8], 64 source and shuffle partitions, 18 numeric columns, GPU Kudo read/write, 20 threaded shuffle writers, 1 MiB GPU batches, AQE and shuffle compression disabled.
| Rows |
Current median |
With 1 MiB cuDF fallback |
Improvement |
| 1,000,000 |
0.4460 s |
0.2755 s |
38.2% |
| 10,000,000 |
1.4401 s |
1.0330 s |
28.3% |
With the targeted fallback, Nsight reports one 1 MiB cudaHostAlloc for pool initialization instead of 28,032 per-operation allocations.
Proposed fix
When GPU Kudo serialization is enabled and the configured plugin pinned pool is zero, configure a 1 MiB cuDF fallback pinned pool for scalar staging.
The existing behavior remains unchanged when:
- The plugin shared pinned pool is configured.
- GPU Kudo write serialization is disabled.
- CPU Kudo is used.
This avoids restoring the former 100 MiB cuDF default reservation that was deliberately disabled by #10868. Historical issue #10814 also noted that plugin/cuDF pinned-pool alignment should be revisited after 24.06.
Problem
GPU Kudo shuffle serialization performs excessive CUDA pinned-host allocation when
spark.rapids.memory.pinnedPool.size=0, which is the default.GpuDeviceManagerdisables the cuDF default pinned pool. With no plugin pinned pool to replace it, small cuDF scalar-staging buffers used during GPU Kudo assembly fall through to an independentcudaHostAllocandcudaFreeHostpair.In an Nsight Systems profile of a pure PySpark DataFrame shuffle workload, each
shuffle_assembleinvocation caused exactly 73 allocation/free pairs. Six actions containing 384 assembles produced:cudaHostAlloccallscudaFreeHostcallscudaHostAllocAPI timeA larger profile produced 56,064 pairs across six actions.
End-to-end impact
Configuration: Spark 3.5.0,
local[8], 64 source and shuffle partitions, 18 numeric columns, GPU Kudo read/write, 20 threaded shuffle writers, 1 MiB GPU batches, AQE and shuffle compression disabled.With the targeted fallback, Nsight reports one 1 MiB
cudaHostAllocfor pool initialization instead of 28,032 per-operation allocations.Proposed fix
When GPU Kudo serialization is enabled and the configured plugin pinned pool is zero, configure a 1 MiB cuDF fallback pinned pool for scalar staging.
The existing behavior remains unchanged when:
This avoids restoring the former 100 MiB cuDF default reservation that was deliberately disabled by #10868. Historical issue #10814 also noted that plugin/cuDF pinned-pool alignment should be revisited after 24.06.