Skip to content

feat(cudf): expose live-allocated GPU bytes via statistics_resource_adaptor - #2357

Open
jurossiar wants to merge 3 commits into
IBM:ibm-research-previewfrom
jurossiar:gpu-rmm-stats
Open

jurossiar wants to merge 3 commits into
IBM:ibm-research-previewfrom
jurossiar:gpu-rmm-stats

Conversation

@jurossiar

Copy link
Copy Markdown

What

Wrap the cuDF/RMM device resource (and the separate output MR, when present) in an
rmm::mr::statistics_resource_adaptor, and add a cudfAllocatedBytes() accessor that
returns the live-allocated device bytes.

Why

cudaMemGetInfo only reflects the retained RMM pool high-water mark — it cannot distinguish a
pool that is free-for-reuse from one with active work. The adaptor's current-bytes counter drops
when queries free their allocations, giving a true "busy vs. idle" GPU-memory signal.

This is needed by a downstream Prestissimo change that reports live GPU pool usage on /v1/status
for GPU-aware autoscaling (scale a GPU worker in only when its GPU memory is actually free, not
merely because the retained pool looks full).

How

statistics_resource_adaptor is a cuda::mr::shared_resource (copyable, shared counter state),
so the copies stored in mr_ / output_mr_ feed the same counters read by cudfAllocatedBytes().

Scope / impact

  • Files: velox/experimental/cudf/exec/GpuResources.{h,cpp}, velox/experimental/cudf/exec/ToCudf.cpp.
  • Only active when cuDF is enabled; no effect on non-cuDF builds.
  • Adds a per-allocation counter on the RMM path (statistics adaptor); overhead expected negligible.

Test plan

  • Built with cuDF enabled; cudfAllocatedBytes() tracks a query's allocations and returns to ~0
    when the query frees, whereas cudaMemGetInfo stays at the retained pool high-water mark.

Blink Bot and others added 3 commits June 29, 2026 16:28
…daptor

Wrap the cuDF/RMM device resource (and separate output MR, if any) in an
rmm::mr::statistics_resource_adaptor so live-allocated device bytes can be
queried. cudaMemGetInfo only sees the retained RMM pool high-water mark and
cannot distinguish a free-for-reuse pool from active work; the adaptor's
current-bytes counter drops when queries free their allocations.

Adds cudfAllocatedBytes() accessor. The adaptor is a cuda::mr::shared_resource
(copyable, shared counter state), so the copies stored in mr_/output_mr_ feed
the same counters read here.
@jurossiar

Copy link
Copy Markdown
Author

@dan13bauer FYI.

@dan13bauer

Copy link
Copy Markdown
Collaborator

That looks good to me. Some minor question/comment:

  • What is the overhead of using the statistics resource adaptor? Should this extension be guarded by a compile-time flag?

@jurossiar

Copy link
Copy Markdown
Author

Thanks!
On overhead: the adaptor just forwards to the upstream resource and bumps a mutex-guarded byte/allocation counter per allocate/deallocate — negligible next to the actual pool allocation, and cuDF's allocations are coarse-grained so there's no hot-path concern. I measured/expect no meaningful difference.

On the flag: the whole experimental/cudf module is already gated by the cuDF build option, so this only exists in GPU builds. Rather than a compile-time flag I'd lean toward [always-on / a runtime CudfConfig toggle], since the cost is so low and a compile-time switch fragments builds. Happy to add a runtime toggle if you'd prefer an opt-out.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants