[Store] pin host segments with quota

[zxpdemonio]

Description

This PR adds opt-in CUDA pinned-memory registration for Mooncake Store-managed host segment memory.

The pinned memory is limited to Store segment memory only. It does not pin client local buffers, staging buffers, dummy-client SHM buffers, user-provided buffers, file-backed mountSegment() mappings, CXL/device segments, or tensor-wrapper temporary buffers.

The feature is disabled by default and controlled by a process-wide quota:

• MC_STORE_PIN_MEMORY_MAX_BYTES=<bytes> enables pinned registration up to the configured quota.
• unset, empty, 0, or invalid values disable the feature.
• MC_STORE_PIN_MEMORY=0|false|off|no explicitly disables the feature.

When enabled, Store host segments are registered with cudaHostRegister() after allocation and before the segment is mounted. Registration failures fall back to pageable memory so Store startup and allocation remain usable.

This PR is intentionally scoped to Store segment memory. GPU-destination reads still copy through the client local buffer / staging path before the final CPU -> GPU copy, so get only sees a smaller improvement. Removing or pinning that final staging path is follow-up work and is not part of this PR.

The tensor wrapper path is also not the primary target of this PR. put_tensor(cuda) currently goes through put_parts, which still stages through the client local buffer before writing into Store. Optimizing tensor / put_parts to avoid that staging copy is also follow-up work.

This PR is split from #2598.

Module

• Store

Type of Change

• Performance improvement
• Documentation update

How Has This Been Tested?

Manual performance validation

Manual benchmark on one NVIDIA A10 GPU, local TCP Store setup, 4 GiB Store segment, 2 GiB client local buffer.

Baseline:

• MC_STORE_PIN_MEMORY_MAX_BYTES=0

Pinned Store segment:

• MC_STORE_PIN_MEMORY_MAX_BYTES=4294967296

The benchmark directly exercised the lower-level Store APIs with CUDA source/destination buffers.

Direction	Size	Pageable	Pinned Store segment
put from CUDA buffer	64 MB	8.887 GB/s	24.960 GB/s
put from CUDA buffer	256 MB	9.034 GB/s	25.686 GB/s
put from CUDA buffer	512 MB	9.003 GB/s	25.823 GB/s
get into CUDA buffer	64 MB	5.295 GB/s	5.955 GB/s
get into CUDA buffer	256 MB	5.232 GB/s	6.019 GB/s
get into CUDA buffer	512 MB	5.290 GB/s	6.044 GB/s

The main improvement is on GPU-source writes, where the Store segment is the host destination of the device-to-host copy.

Read-side GPU destinations improve less because the current read path still stages through the client local buffer before the final CPU -> GPU copy. Optimizing that staging path is follow-up work.

Checklist

• I have performed a self-review of my code.
• I have run clang-format / formatting checks.
• I have updated the documentation.
• I have added or updated committed tests.
• I have run pre-commit checks.
• The change is scoped to Store segment memory and does not pin unrelated buffers.

AI Assistance Disclosure

• AI tools were used for implementation, review, benchmarking, and PR description preparation.

[gemini-code-assist]

Code Review

This pull request introduces a mechanism to register Store-managed host segments as pinned memory in CUDA-enabled builds, managed by a new RegisteredPinnedMemoryManager and configured via environment variables. Feedback on the changes suggests leaking the singleton instance of the manager to prevent potential crashes on program exit due to static destruction order issues, and refactoring the environment variable parsing to use std::from_chars instead of std::strtoull for safer, locale-independent parsing.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[codecov-commenter]

⚠️ Please install the to ensure uploads and comments are reliably processed by Codecov.

Codecov Report

❌ Patch coverage is 73.48485% with 70 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
mooncake-store/src/registered_pinned_memory.cpp	60.36%	65 Missing ⚠️
mooncake-store/src/real_client.cpp	79.16%	5 Missing ⚠️

📢 Thoughts on this report? Let us know!

[he-yufeng]

Read through the whole thing — the manager and the real_client wiring. This is careful work and the correctness-sensitive parts hold up well:

• The quota check is overflow-safe (size > limit_bytes_ || pinned_bytes_ > limit_bytes_ - size rather than pinned_bytes_ + size > limit_bytes_), and the reserve-then-commit split — insert the key with a nullptr placeholder under the lock, cudaHostRegister outside the lock, then commit the pointer — keeps a single owner per (addr, size) since the emplace is exclusive. A concurrent try_pin for the same key fails at the placeholder insert, so release's window between "mark nullptr" and "erase" can't be hijacked.
• Failure paths refund correctly: a failed cudaHostRegister calls drop_reservation_locked (placeholder removed, quota returned) and falls back to pageable, and cudaGetLastError() clears the sticky error. The env parse is properly guarded (std::from_chars + explicit negative check + trim, invalid → disabled) rather than a bare stoull.
• Lifetime is right: the region shared_ptr is held in setup_segment_pinned_regions_ for the segment's life and released on teardown, and the pin is dropped if MountSegment fails. Scope is correctly limited to Store host segments.

Two things I'd like to resolve before approving, neither a correctness bug:

1. release() escalates a non-unloading cudaHostUnregister failure to LOG(FATAL). For a long-running Store process, aborting on a pin-bookkeeping unregister failure is a heavy response — the alternative is LOG(ERROR) and leak the pin so the service survives a transient CUDA hiccup. The cudaErrorCudartUnloading case is already handled gracefully; is aborting on the other errors the intended contract, or would you rather degrade? (The LOG(FATAL) in the try_pin tracking-mismatch path is fine as an assert — that branch is unreachable given the single-owner invariant.)

2. There are no tests, but the second constructor taking std::pair<bool, uint64_t> config is an obvious seam for them. The cudaHostRegister/Unregister calls are the real barrier to full coverage, but the reservation math — over-quota returns nullptr, overlap returns nullptr, exact-key double-pin returns nullptr, quota is refunded after release — is the part most worth locking down, and it's the part most likely to drift later. Worth a small unit test around the manager's bookkeeping even if the CUDA path stays CI-only.

Design and the careful bookkeeping look good; happy to approve once the FATAL question is settled.

[zxpdemonio]

Read through the whole thing — the manager and the real_client wiring. This is careful work and the correctness-sensitive parts hold up well:

• The quota check is overflow-safe (size > limit_bytes_ || pinned_bytes_ > limit_bytes_ - size rather than pinned_bytes_ + size > limit_bytes_), and the reserve-then-commit split — insert the key with a nullptr placeholder under the lock, cudaHostRegister outside the lock, then commit the pointer — keeps a single owner per (addr, size) since the emplace is exclusive. A concurrent try_pin for the same key fails at the placeholder insert, so release's window between "mark nullptr" and "erase" can't be hijacked.
• Failure paths refund correctly: a failed cudaHostRegister calls drop_reservation_locked (placeholder removed, quota returned) and falls back to pageable, and cudaGetLastError() clears the sticky error. The env parse is properly guarded (std::from_chars + explicit negative check + trim, invalid → disabled) rather than a bare stoull.
• Lifetime is right: the region shared_ptr is held in setup_segment_pinned_regions_ for the segment's life and released on teardown, and the pin is dropped if MountSegment fails. Scope is correctly limited to Store host segments.

Two things I'd like to resolve before approving, neither a correctness bug:

1. release() escalates a non-unloading cudaHostUnregister failure to LOG(FATAL). For a long-running Store process, aborting on a pin-bookkeeping unregister failure is a heavy response — the alternative is LOG(ERROR) and leak the pin so the service survives a transient CUDA hiccup. The cudaErrorCudartUnloading case is already handled gracefully; is aborting on the other errors the intended contract, or would you rather degrade? (The LOG(FATAL) in the try_pin tracking-mismatch path is fine as an assert — that branch is unreachable given the single-owner invariant.)

@he-yufeng I've updated this part to be availability-first.
I downgraded non-cudaErrorCudartUnloading cudaHostUnregister failures from LOG(FATAL) to LOG(ERROR). The cleanup is now best-effort: even if cudaHostUnregister fails, the manager still drops the region from its bookkeeping and refunds the pinned quota, so we do not keep stale raw tracking pointers after the Store segment owner is released.
cudaErrorCudartUnloading is still handled as a shutdown-only warning.
I also added a small unit test for this behavior: when the fake unregister backend returns an error, releasing the pinned region no longer aborts, and the same address range can be pinned again afterwards. This verifies that the manager state is cleaned up even on unregister failure.

2. There are no tests, but the second constructor taking std::pair<bool, uint64_t> config is an obvious seam for them. The cudaHostRegister/Unregister calls are the real barrier to full coverage, but the reservation math — over-quota returns nullptr, overlap returns nullptr, exact-key double-pin returns nullptr, quota is refunded after release — is the part most worth locking down, and it's the part most likely to drift later. Worth a small unit test around the manager's bookkeeping even if the CUDA path stays CI-only.

Test cases are added.

Design and the careful bookkeeping look good; happy to approve once the FATAL question is settled.