feat: optional CUDA host-memory pinning of the ring buffers - #5
Merged
Conversation
Page-lock each ring buffer via cudaHostRegister so a downstream jax.device_put of the numpy views is a fast, truly-async H2D DMA instead of a synchronous host->device staging copy. libcudart is resolved at runtime via dlopen (libc only) -- no CUDA toolchain / link-time dependency; the code is always compiled and is a graceful no-op unless ECHO_PIN_HOST_MEMORY=1 and libcudart is present. Pinning happens only at ring-buffer construction (PytreeRingBuf::new) and teardown (Drop), off the tokio runtime. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Keeps CLAUDE.md, docs/agents/ and .scratch/ out of the repo. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Rewrites CUDA host-memory pinning so that it works, and so that it cannot
silently not work.
The previous implementation opened the unversioned `libcudart.so` soname,
which the pip CUDA runtime wheels do not ship. The dlopen returned NULL,
every pin call became a no-op, and nothing reported it — so a workload
could be measured against a code path that never executed and be
concluded not to benefit from pinning.
Pinning is now one keyword argument, off by default:
Server(example, batch_size, ..., pin_host_memory=False)
The contract is that if `Server(...)` returns, every ring buffer is
page-locked. Any inability to deliver that raises `RuntimeError` naming
every path probed and the CUDA error symbolically. That guarantee is what
replaces an observability API: construction succeeding *is* the
assertion, and unlike an explicit check it cannot be forgotten.
- Resolution is a three-rung ladder — already-mapped libraries, then the
installed CUDA wheels (searching the `nvidia` vendor package rather
than a named component, since CUDA 13's consolidated layout differs
from CUDA 12's per-component one), then sonames, versioned first. No
soname symlink or LD_LIBRARY_PATH entry is needed.
- Page-locking moved out of `PytreeRingBuf::new` into an explicit
fallible step on a constructed buffer, so the existing `Drop` owns
rollback for both the failure path and normal teardown rather than a
hand-written unregister loop. `new` stays infallible and the eight
existing ring-buffer tests are untouched.
- CUDA entry points go through an injectable struct of function
pointers, so registration, rollback and drop are all exercised with
stubs on a machine with no GPU.
- Registration is portable-flag only. Page-aligned allocation and
read-only registration were both measured and rejected: alignment
gained 1.6% of copy time (bar was 5%), and read-only is unsupported on
the test GPU and is documented as permission rather than optimisation.
- `ECHO_PIN_HOST_MEMORY` is deleted; the argument is the only control.
GPU tests carry a `gpu` marker and skip when no device is present. The
new guide documents the mechanism, the measured numbers, the unswappable
footprint arithmetic, and how to verify pinning engaged — including that
`VmLck` stays at zero even when it is working, which is what made the
original false negative believable.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Every other module keeps its tests under tests/; host_pinning was the only one with an in-file #[cfg(test)] module, and at 1012 lines was three times the next-largest source file. The tests move to src/host_pinning/tests.rs — in-crate rather than under tests/ because the module is private and the stub-injection tests need the internal CudaApi — leaving 488 lines of source. Also cuts the comments back towards the density of the surrounding code. The load-bearing "why" stays (the pageable staging mechanism, why the unversioned soname is not enough, why the vendor package is searched rather than a named component, why pinning is a post-construction step, what forcing runtime init actually costs); the restatements of the guide and the narrative asides go. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The module did two jobs: find the CUDA runtime, and page-lock memory with
it. Resolution was ~75% of the source and the part with no CUDA calls in
it at all, so the two barely interact — resolve hands back a CudaApi and
register uses it.
host_pinning/
├── mod.rs shared types: CudaApi, Region, PinError
├── resolve.rs the three-rung ladder
├── resolve/tests.rs
├── register.rs pin_all / unpin_all / rollback
└── register/tests.rs
Tests nest under the module they cover rather than sitting in one file, so
each keeps its own privacy: the ladder's internals stay private to
resolve, and the stub-injection tests stay next to the unsafe code they
exist to cover. Only scan_mapped_runtimes widens, to pub(super), because
register's GPU test needs the mapped-runtime path to look up
cudaHostGetFlags independently of the pointers under test.
No behaviour change; the same 20 tests pass, 10 per module.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The tests were inside src/ because host_pinning was the crate's only private module, and a private module's internals are invisible to tests/, which are separate crates. That inverted the repo's convention: every other module is pub with pub internals for precisely this reason — PytreeRingBuf::slot_mut is a pub unsafe fn so tests/ring_buf.rs can drive it. So host_pinning becomes a pub mod and tests/host_pinning.rs joins the other six, one per module. This drops ticket 01's "module stays private" checklist item, deliberately and noted there. What that constraint protected is unaffected: the crate is publish = false and ships only as a Python extension, whose entire surface for this feature is pin_host_memory=True. Two things keep the widening honest — CudaApi's fields stay private behind `unsafe fn CudaApi::new`, so no caller can fabricate one with arbitrary function pointers, and pin_all/unpin_all stay unsafe fn with their Safety contracts. Same 20 tests, still passing, GPU ones included. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The Rust-internals section had no entry for host_pinning at all — the module map in overview.md predated the feature and was never updated — and the resolution ladder's design was only described in the user-facing guide. Adds docs/src/design/host-pinning.md covering the internals: the resolve/register split and why it's clean, why resolution is a ladder and what each rung is for, why registration is a post-construction step, the function-pointer injection seam and what it exists to test, why the module is pub, and the absence of any build-time CUDA dependency. Trims ring-buffer.md's pinning section to what is actually about the ring buffer — that its buffers are never reallocated, so a registration stays valid for their life, and that Drop reverses it — with the module-level rationale moved to the new page. Adds the module map row, the nav entry, and a pointer from the guide. development.md now says where the Rust tests live, and distinguishes the Rust skip-and-pass convention from the Python `gpu` marker. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
`just docs-serve` binds to loopback, so it is invisible when the checkout lives on a remote box. Adds a recipe that takes a bind address, and documents both routes in development.md — SSH port forwarding (preferred, no server-side change) and binding to a reachable interface. Also records the /echo/ path prefix that site_url imposes on the dev server, since http://127.0.0.1:8000/ on its own 404s. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The two new pages carried a lot of scaffolding that wasn't doing work:
roadmap sentences ("this page covers..."), self-referential asides
("this section is the reason this page exists"), and a bolded
mini-headline on every paragraph of the Measured section. Technical
content is unchanged; the prose is trimmed to the voice of the
surrounding docs.
Substantive fixes found while editing:
- The sample RuntimeError in the guide omitted the two "(none)" lines
that resolve() emits for a rung that ran and found nothing, which are
the whole point of reporting empty rungs. Replaced with real output.
- "Two jobs, one file each" introduced a three-row table.
- Dropped the CHANGELOG's Changed and Removed sections. Relative to
0.1.1 the resolution ladder is new, not changed, and
ECHO_PIN_HOST_MEMORY was added and removed inside this branch, so it
never shipped and users can't have depended on it.
- The Added entry showed pin_host_memory=False, the default, rather than
the value that turns the feature on.
- Reported the primary-context cost as ~128 MB in register.rs to match
the measured figure in the guide, not ~100 MB.
- Replaced "~1800x" occupancy with "three orders of magnitude"; the
ratio was derived from a 0.002 ms figure at timer resolution.
- ring-buffer.md now links to the constructor rationale rather than
restating it.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…e pinning The failure message is the entire interface to a resolution failure, and nothing tested it against a real failure. The Rust test hand-builds a PinError and checks Display; the Python test that would have exercised the actual accumulation was skipif(HAS_GPU), so it skipped on every machine with a GPU — i.e. everywhere this gets developed. That gap is why the guide's example output was wrong: it omitted the two "(none)" lines that resolve() emits for a rung that ran and found nothing, which are the whole reason empty rungs are reported at all. Replaced it with a subprocess test that forces every rung to fail (nothing mapped, nvidia unimportable, no LD_LIBRARY_PATH) and asserts each rung appears, in order. It runs on GPU hosts too, since it needs resolution to fail rather than a device to be absent. Verified by mutation: dropping the empty-rung note fails it, and reordering the reporting loop fails it on ordering. The soname probe runs inside the subprocess deliberately. Probing in the parent reports every soname as resolvable, because the module loads the runtime by absolute path at import and dlopen then matches the already-loaded object by soname — which would have skipped this test on any host with the wheel installed. Also release the GIL around pinning. Measured, a cold cudaFree(NULL) creating the primary context takes ~210 ms, all of it with the GIL held; detach lets other Python threads run at 82% of their unblocked rate during construction instead of stalling. detach is synchronous, so the constructor still returns only once every buffer is registered and the page-locked guarantee is unchanged. It also removes a deadlock class: holding the GIL across a CUDA call lets echo block against a thread that holds a CUDA lock and wants the GIL. Minor: note why empty_rung_note's Soname arm is unreachable, and why the wheel walk does not follow directory symlinks. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Trimming the Measured and Verifying sections left four things pointing at content that no longer exists: - The footprint section linked to #verifying-that-pinning-engaged. Folded the part that mattered into one sentence instead: VmLck stays at zero because the driver's page-locking bypasses mlock accounting, so it is neither a binding limit nor a way to check pinning engaged. - design/host-pinning.md and the CHANGELOG both advertised the guide as covering measured numbers and how to verify pinning; they now describe what it actually covers. - server.py lost the blank line before `Raises:`, so griffe stopped parsing it as a section and rendered it as loose prose after the parameters table. Confirmed against the built HTML: the page had only a Parameters: section, and now has Parameters: and Raises:. Also: the README said pinning means "even H2D doesn't require copies". Pinning doesn't remove the copy, it makes it a DMA transfer instead of a staged one — and next to echo's zero-copy claims that reads as though H2D were free. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
sash-a
force-pushed
the
feat/pin-host-memory
branch
2 times, most recently
from
August 3, 2026 13:47
c63c481 to
efbf94b
Compare
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Adds
Server(..., pin_host_memory=True), which CUDA-page-locks echo's ringbuffers so a host-to-device copy of a sampled batch is a real DMA transfer
instead of a chunked staging copy the driver walks through on the calling
thread. Off by default;
0.1.1→0.2.0.Why
Copying out of pageable memory isn't a DMA transfer. The driver stages it
through a small internal pinned buffer in CPU-executed chunks, holding the
calling thread and the driver lock throughout — so
cudaMemcpyAsyncstopsbeing asynchronous and blocks the same thread's kernel launches. On a learner
issuing tens of thousands of small launches per step, that contention, not
bandwidth, is what shows up as host-side dispatch dominating the step.
Measured locally (RTX 5060 Ti, driver 595.71.05, CUDA 13.3),
cudaMemcpyAsyncout of pageable memory holds the calling thread for 97% of the copy
(3.645 ms of 3.751 ms); out of page-locked memory it returns in 0.002 ms.
Bandwidth barely moves on this machine (~14 GB/s either way) because the host
link is already saturated — the occupancy figure is the reason to turn this on.
Full numbers and caveats are in the guide.
The guarantee
If the constructor returns, every ring buffer is page-locked. Anything that
would prevent it — no CUDA runtime found, no usable device, a registration
rejected — raises
RuntimeError, naming every path probed and the CUDA errorsymbolically. There is no status object and no silent fallback, because a
silent no-op is indistinguishable from pinning not helping.
That failure mode is not hypothetical: the first cut of this branch resolved the
runtime with a single
dlopen("libcudart.so"). The pip CUDA runtime wheels shiponly the versioned soname, with no unversioned symlink and no ldconfig entry,
so every pin was quietly a no-op. Resolution is now a three-rung ladder —
already-mapped libraries, then the installed CUDA wheels, then sonames
(versioned first) — built by a pure
candidates()function so its ordering isunit-tested, since reordering the rungs is exactly how it broke.
Shape
src/host_pinning/splits in two, meeting only atCudaApi:mod.rsCudaApi,Region,PinErrorresolve.rsregister.rsdlsym'd at pintime,
RTLD_LOCAL. One wheel installs on GPU and CPU-only hosts alike, andwith pinning off nothing is loaded.
Errnever runs
Drop, so registering there forces a hand-written unregister loopon the error path.
pin_host_memoryruns on a fully-built buffer andDropowns rollback for both the failure path and teardown.
CudaApiis injected, not global, so rollback — whose only visibleconsequence is the absence of leaked registrations — is testable. A stub
that fails on the third buffer must produce exactly two unregister calls.
Testing
cargo test73 passed,pytest73 passed / 3 skipped,cargo clippy --all-targetsclean,mkdocs build --strictclean. Verified on a GPU host.Most of the coverage needs no GPU: the resolution ladder is pure, and
registration and rollback run against injected stubs. The GPU tests carry a
gpumarker and skip automatically, so CI needs no per-runner configuration.The Python tests deliberately do not ask echo what it did — they ask the
CUDA runtime through
ctypesfor the registration flags on the address behind asampled batch. A test that trusted echo's own account would have passed against
the silent no-op described above.
Docs
New guide (mechanism, measured
numbers, footprint arithmetic, and how to confirm pinning engaged — including
why
VmLckis not a valid check, since it stays at zero even when pinningworks) and a design page covering the module
split.
🤖 Generated with Claude Code