NVIDIA / NVIDIA/cudf

[FEA]AST-Aware Unified Memory Prefetching to Mitigate Page Fault Storms

Open
#21,007 1 comment 0 reactions 0 assignees View on GitHub
feature request libcudf no-oom
Dominant language
C++
Stars
9.8k
Forks
1.1k
Avg merge
3d 6m
Merged PRs (30d)
278

Description

### Problem Description: The Page Fault Bottleneck
In Unified Memory (UM) scenarios, `cudf` operations on datasets larger than VRAM often suffer from severe latency due to **"Page Fault Storms"**. When a kernel accesses host-resident pages, the GPU SMs (Streaming Multiprocessors) stall, waiting for the interconnect (PCIe/NVLink) to serve the page migration.

Currently, `cudf` relies heavily on the driver's heuristic or explicit `rmm` pooling. However, the library typically possesses semantic knowledge of future access patterns (via the AST/computation graph) that the low-level CUDA driver lacks.

### Proposed Architecture: Semantic Prefetching Layer
I propose investigating a **Semantic Prefetching Layer** that analyzes the `cudf` Expression Graph before kernel launch to issue proactive memory hints.

Instead of relying on demand-paging, we can utilize:
1. **`cudaMemPrefetchAsync`**: To move columns involved in the *next* operation to the GPU immediately, overlapping with the current CPU-side query planning.
2. **`cudaMemAdvise` Hints**:
* Mark join keys as `cudaMemAdviseSetPreferredLocation` (GPU).
* Mark scan-only columns as `cudaMemAdviseSetReadMostly`.

Impact & Justification
Latency Masking: By issuing prefetches during the query planning phase (CPU time), we hide the PCIe transfer latency.

TLB Efficiency: Reducing on-the-fly page migrations reduces pressure on the GPU MMU (Memory Management Unit) and TLB misses.

Out-of-Core Scaling: Enables smoother degradation for workloads slightly larger than VRAM, preventing hard stalls.

I believe this aligns with the roadmap for better Managed Memory support in libcudf.

#### Technical Implementation Concept
We can introduce a `MemoryOptimizer` visitor that walks the `libcudf` AST.

```cpp
// Pseudo-code for AST-Aware Prefetcher
void optimize_memory_placement(const aggregation_request& req, cudaStream_t stream) {

// 1. Identify "Hot" Columns (e.g., GroupBy Keys)
for (auto& col : req.keys) {
// We know these pages will be accessed randomly and heavily.
// Force prefetch to VRAM to avoid TLB Thrashing.
cudaMemPrefetchAsync(col.data(), col.size(), current_device, stream);

// Hint the driver that this data should stay on GPU
cudaMemAdvise(col.data(), col.size(), cudaMemAdviseSetPreferredLocation, current_device);
}

// 2. Identify "Cold/Streaming" Columns (e.g., Value columns in simple sum)
for (auto& col : req.values) {
// These might be accessed linearly. Relax residency requirements if VRAM is tight.
if (is_vram_pressure_high()) {
cudaMemAdvise(col.data(), col.size(), cudaMemAdviseSetAccessedBy, current_device);
}
}
}

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.