[optim] Lease pre-registered receive buffers on the MooncakeStore read path

[weiziyoung]

Refs: #169

Thanks for the pointer! I did benchmark the BATCH_SIZE_LIMIT knob before going with
the buffer pool, and I think the data makes a clear case that batch-size tuning
cannot substitute for the pool。

1. Tuning BATCH_SIZE_LIMIT does not help

Register-per-read, threads=4, both reads swept over the same BATCH_SIZE_LIMIT
grid (read time, ms, median; per-key size fixed at 1MB, only the key count differs):

BATCH_SIZE_LIMIT	50	100	200	400	800	1600	3200
1024×1MB (1GB)	109	138	155	160	253	384	329
4096×1MB (4GB)	528	461	455	496	666	723	1049
pool @400 (ref)				1GB 103 / 4GB 434

The optimum sits at small values (~50 for 1GB, ~200 for 4GB) and grows worse as the
batch grows. The reason is
that register_buffer cost is dominated by per-byte page pinning, not per-call
overhead: the total bytes to pin per read is fixed, so a larger batch does not reduce
the pinning work — it only reduces how many worker threads share it (fewer, larger
chunks → less parallelism), plus a marginal merge-before-register saving. Even at 4GB,
where num_keys/threads = 1024 keeps 4 chunks well past BATCH_SIZE_LIMIT=1024, the
curve still degrades — so this isn't purely a chunk-count artifact.

2. The pool removes registration entirely

Breakdown at threads=1 (segment sums ≈ wall), 1024×1MB (1GB):

stage	register-per-read	pool
register_buffer	279 ms (~79%)	0
batch_get_into (RDMA) + copy/meta	~75 ms	104 ms
total	354 ms	104 ms (3.4×)

Registration is ~79% of the read; leasing pre-registered buffers removes it. Crucially,
the pool beats the best-tuned baseline at every size (e.g. 4GB: pool 387 ms vs the
baseline's best 419 ms), so no BATCH_SIZE_LIMIT value closes the gap.

3. Where the win is big, and where it's honestly marginal

baseline vs pool, threads=4, BATCH_SIZE_LIMIT=400 (read time, ms, median):

workload	baseline	pool	speedup
64×16MB (1GB)	367	94	3.9×
1024×1MB (1GB)	178	109	1.6×
4096×1MB (4GB)	419 (best bsl)	387	1.08×
16384×64KB (1GB)	389	480	0.8× (slower)

Two honest caveats:

• The advantage shrinks with more threads and larger reads. At threads=1 the pool
  is ~3.4×; at threads=4 the register cost is spread across workers, and the pool's
  copy-out (staged buffer → caller tensors) scales with bytes, so a 4GB read narrows
  the gain to ~1.08×. The big wins are single-reader / few-large / mid-size reads;
  pool also lowers step-to-step jitter (e.g. 1024×1MB threads=1 stdev 19.6 → 7.8 ms).
• Many tiny keys can regress. For 16384×64KB the mandatory copy-out doesn't pay off
  vs. that shape's cheap merged registration. I vectorized the copy-out (uniform groups
  are copied in one strided pass instead of a per-tensor Python loop), which cut this
  case from ~2× slower to ~0.8×, but it's still a net loss at extreme fragmentation.
  If desirable, a small num_keys/per-key-size heuristic could fall back to
  register-per-read here; happy to add it, though real trajectory reads (≈1024×1MB)
  don't hit it.

The copy is what keeps returned tensors caller-owned so the lease returns immediately.
A zero-copy view mode (return tensors backed by the lease, release on the next read)
would restore the large-read win but requires a "no cross-read retention" contract, so
I kept the copy path as the safe default.

4. Compatibility

Per your suggestion, the register-per-read path is preserved as a fallback: the client
gates on mooncake.store.BufferPool availability, so builds without it (i.e.
mooncake-transfer-engine < v0.3.12) transparently keep registering their own receive
buffers. Requests the pool cannot serve also fall back.

[ascend-robot]

CLA Signature Guide

@weiziyoung , thanks for your pull request.

The following commit(s) are not associated with a signed Contributor License Agreement (CLA).

Commit	Reason
[a4bbdd2 [optim] Lease pre-registered re...](a4bbdd2)	the email used in the commit is not linked to a signed CLA! please verify that it matches the email you used when signing the CLA.

To sign CLA, click here.

To check if your email is configured correctly, refer to the FAQs.

Once you've signed the CLA or updating your email, please comment /check-cla to revalidate CLA status.

[weiziyoung]

/check-cla

[ascend-robot]

CLA Signature Guide

@weiziyoung , thanks for your pull request.

The following commit(s) are not associated with a signed Contributor License Agreement (CLA).

Commit	Reason
[a4bbdd2 [optim] Lease pre-registered re...](a4bbdd2)	the email used in the commit is not linked to a signed CLA! please verify that it matches the email you used when signing the CLA.

To sign CLA, click here.

To check if your email is configured correctly, refer to the FAQs.

Once you've signed the CLA or updating your email, please comment /check-cla to revalidate CLA status.

[weiziyoung]

/check-cla

[ascend-robot]

CLA Signature Pass

weiziyoung, thanks for your pull request. All authors of the commits have signed the CLA. 👍

[0oshowero0]

How about rename the file as test_mooncake_buffer_pool.py?

[0oshowero0]

the name lease may confuse the readers. Maybe we can call it _copy_buffer_content_into_tensors?

[0oshowero0]

We can simplify the AI comments

[0oshowero0]

It's better to tell the user which version supports this feature.

[0oshowero0]

For GDR and TCP path, we can omit this warnning

[0oshowero0]

We better optimize the function name so that it can let users directly get the idea that this is the counterpart of _read_into_own_buffers. Maybe we can consider _read_into_tensors and _read_into_buffers? Just for inspiration :)

[0oshowero0]

【AI Review】

1. max_bytes=0 does the opposite of what the comment claims

# max_bytes=0: lease from that local buffer only, without an extra arena.
self._buffer_pool = BufferPool(self._store, max_bytes=0) if MOONCAKE_BUFFER_POOL_IMPORTED else None

In mooncake-integration/store/buffer_pool.cpp, max_bytes == 0 is the constructor default and it
sets the budget to twice the local buffer:

if (max_bytes == 0) {
    max_bytes_ = local_buffer_capacity_ * 2;
} else {
    max_bytes_ = std::max(max_bytes, local_buffer_capacity_);
}

The upstream docs describe this as "the default allows one local-buffer-sized overflow burst", and
the overflow path is allocate_overflow_unlocked(): posix_memalign plus one
store.register_buffer() per lease, with unregister_buffer() on release. That is precisely the
per-transfer registration this PR exists to remove, now hidden inside mooncake. Consequences:

• The optimization can silently no-op. When a lease can't be served from the local buffer you get
  an overflow buffer, i.e. pre-PR performance, with no error.
• It is not observable from Python. The pool tracks allocate_count_ / oversize_allocate_count_
  but the pybind bindings only expose acquire, buffer, prewarm and close, so there's no way
  to detect the degradation from TransferQueue.
• Up to another local_buffer_size (1 GiB by default) of unbudgeted host memory.
  This is pitfall 1 from your own issue [optim] MooncakeStore read path registers RDMA buffers on every get(), dominating end-to-end time #169 ("silent fallback on capacity shortfall"). Also note that
  because exhaustion now requires more than 2× capacity in flight, the _acquire_lease() -> None
  fallback you wrote essentially never fires.
  Suggestion: pass max_bytes=self.local_buffer_size so the in-flight budget is bounded by the
  registered region, and fix the comment. (To be precise: this bounds overflow rather than
  eliminating it — allocate_overflow_unlocked is still reachable when the local allocator fails on
  fragmentation while total_bytes_ < max_bytes_ — but it stops a whole extra arena from being
  absorbed silently.)

[0oshowero0]

【AI Review】

2. The lease share targets 100% of a buffer mooncake also uses internally

# One share per reader thread, so all of them can hold a lease at the same time.
self._lease_bytes = self.local_buffer_size // MAX_BATCH_WORKER_THREADS

MAX_BATCH_WORKER_THREADS shares of local_buffer_size / MAX_BATCH_WORKER_THREADS is the entire
buffer. The arithmetic base is right — BufferPool reads local_buffer_capacity_ from
client_buffer_allocator_->size(), which is exactly the local_buffer_size passed to setup() —
but full utilization assumes exclusive ownership and lossless sub-allocation, and neither holds:

• Not exclusive. mooncake-store/src/real_client.cpp registers this region once and also
  records it as local_buffer_region_ for internal Store staging. The design doc calls it a "soft
  isolation policy: internal Store paths and external Python leases share the local registered
  buffer".
• Hard boundary, no slack. The allocator is OffsetAllocator. Upstream's own
  client_buffer_test.cpp shows a 64 KB buffer with 8 KB requests succeeding exactly 8 times and
  failing on the 9th, and getLargestFreeRegion() is documented as best-effort because "the actual
  allocation may still fail due to race conditions or fragmentation".
• Requests round up to a bin. OffsetAllocator bins with a 3-bit mantissa (documented worst
  case +12.5%), and the mooncake wrapper notes it rounds the allocated size up to a bin size.
  Replicating uintToFloatRoundUp, a 250 MiB batch consumes 256 MiB, a 225 MiB batch consumes
  240 MiB. So the usable share is slightly smaller than the nominal one.
  With defaults, four 256 MiB shares sum to exactly 1024 MiB — a perfect fit with zero margin (the
  share is a power of two, so bin rounding doesn't inflate it). Any internal staging use or
  fragmentation makes one thread's local allocation fail, which lands on issue 1's silent overflow.
  That is why these two compound: the 100% split makes local allocation failure the expected case,
  and max_bytes=0 makes that failure invisible.
  Suggestion: leave headroom (e.g. // (2 * MAX_BATCH_WORKER_THREADS)) and say in the comment
  that the share is coupled to the local_buffer_size config.
  While we're here, the coupling is worth documenting because it also changes the number of round
  trips. For a 400 × 1 MiB batch, which used to be a single batch_get_into:
  | local_buffer_size | share | sequential rounds |
  |---|---|---|
  | 64 MB | 16 MB | 25 |
  | 256 MB | 64 MB | 7 |
  | 1024 MB (default) | 256 MB | 2 |
  | 4096 MB | 1024 MB | 1 |
  The default is fine, but anyone who shrinks local_buffer_size gets 25 serialized rounds without
  changing anything else.

[0oshowero0]

Following up on point 2 of my earlier review (the lease share consuming 100% of the local buffer).
I described it as a performance concern; looking at what actually happens when the buffer fills up,
the consequence is worse than that, and I think it moves the sizing from "worth tuning" to "worth
fixing before merge".

Filling the buffer doesn't fail — and it doesn't reach your fallback

The pool's admission check doesn't look at the local buffer at all:

bool BufferPoolNative::has_capacity_for_locked(size_t size_class) const {
    if (size_class > max_bytes_) return false;
    if (total_bytes_ > max_bytes_ - reserved_bytes_) return false;
    ...

max_bytes_ is the pool's own budget (2× capacity with max_bytes=0), not the allocator's
remaining space. So once the local buffer is full, acquire() still admits the request,
client_buffer_allocator_->allocate() returns nullopt, and try_acquire_locked falls through to
allocate_overflow_unlocked(). Nothing is raised, so _acquire_lease() never returns None and
the register-per-read fallback you wrote is never reached. The degradation happens entirely inside
mooncake.

The overflow path is slower than the pre-PR baseline

This is the part I'd flag hardest, because the intuition is "worst case we're back where we
started", and that isn't what happens:

	before this PR	overflow lease
receive memory	the caller's own torch.empty tensors	posix_memalign, 8 MiB aligned, fresh pages
registration	register_buffer on the merged target regions	register_buffer on the overflow buffer
RDMA lands in	the target tensors directly	the overflow buffer
copy	none	one full copy-out
teardown	unregister_buffer	unregister_buffer + free back to the OS

The registration cost is unchanged — by your own analysis it's dominated by per-byte page pinning,
and the total bytes pinned are the same. On top of that you now pay a full copy-out (~11 ms per
GiB in my measurements), a large fresh allocation whose pages have to be zero-filled when they're
pinned, and a free that returns them to the OS so the next lease faults them in again.

So whenever overflow kicks in, this PR is a net regression rather than a no-op.

And it's not observable

No exception means the logger.warning in _acquire_lease never fires. The pool does count
allocate_count_ and oversize_allocate_count_, but the pybind bindings only export acquire,
buffer, prewarm and close, so there's no way to read them from Python. The symptom in
production is "the optimization doesn't seem to help" with nothing to grep for.

With offload.enabled: true this can become a hard read failure

This is the part I missed the first time, and it's the real argument for headroom.

Inside batch_get_into, any key served by a remote DISK replica needs a scratch allocation from
the same client local buffer that the leases come from
(RealClient::batch_get_into_multi_buffers):

for (auto &[key, op] : valid_local_disk_ops) {
    if (op.is_local_disk) continue;
    auto alloc_result = client_buffer_allocator_->allocate(op.total_size);
    if (!alloc_result) {
        LOG(ERROR) << "Failed to allocate temp buffer for DISK read, key: " << key ...
        results[op.original_index] = tl::make_unexpected(ErrorCode::NO_AVAILABLE_HANDLE);

That sets up a self-starvation:

1. MAX_BATCH_WORKER_THREADS reader threads each hold a lease, together covering 100% of the local
   buffer.
2. A key inside one of those same batch_get_into calls lives on SSD, so mooncake tries to
   allocate scratch space from that now-empty buffer.
3. The allocation fails, the key comes back NO_AVAILABLE_HANDLE, i.e. a negative return code.
4. _batch_get_into_with_retry retries 3 times with a 1 s delay — but the lease is held across the
   retries, so the memory the retry needs is held by the retrier. All three attempts fail
   identically.
5. The read raises RuntimeError.

The retry logic can't help here by construction, which is what makes this different from ordinary
capacity pressure.

Preconditions: offload.enabled: true (config.yaml describes it as essential when total CPU DRAM
is smaller than GPU HBM), the object already evicted to SSD, and the reader not running on the node
hosting the offload client — which, given the documented single-node centralized offload pool, is
every other node in the cluster.

What headroom actually buys

Two separate things, which is why I'd change the sizing rather than just document it:

• It keeps overflow an exception instead of the steady state at full concurrency, so the
  optimization keeps working and doesn't silently invert.
• It leaves the store's internal paths the scratch space they need within the same call, so TQ
  doesn't starve its own reads.

Restating the suggestion from the earlier comment, now with the reasoning above behind it:

# Leave headroom: the local buffer is shared with mooncake's internal staging paths,
# including the scratch allocation batch_get_into needs for SSD-resident keys.
self._lease_bytes = self.local_buffer_size // (2 * MAX_BATCH_WORKER_THREADS)

together with max_bytes=self.local_buffer_size on the pool. With both in place, genuine exhaustion
raises instead of silently overflowing, and the explicit fallback in _read_group_via_lease finally
does the job it was written for.

[0oshowero0]

【AI Review】

4. lease.release() in finally can mask the original error and then block close()

try:
    self._batch_get_into_with_retry(keys, [lease.ptr + off for off in offsets], nbytes)
    _copy_lease_into_tensors(lease, [batch_tensors[i] for i in group], offsets)
finally:
    lease.release()

release() throws cannot release buffer while exported views exist when exports_ != 0. The
del src / del staged calls handle the happy path, but if copy-out raises, the torch.frombuffer
view stays alive in the traceback's frame, so release() raises from the finally and that
becomes the surfaced exception (the real one survives only as __context__). I reproduced this
with a lease that models the upstream check.

The follow-on matters more: close() throws cannot close buffer pool with active leases while any
lease is in use, and it's the first thing MooncakeStoreClient.close() does — so
_gdr_staging.close() and _store.close() would both be skipped. Suggest tolerating failure on
both release() and pool.close() (log and continue). Mitigating factor:
~BufferLeaseNative calls release_lease(false) on GC, so this isn't a permanent leak.

[0oshowero0]

【AI Review】

5. The fallback logs one WARNING per group per read

except Exception as e:
    logger.warning(f"Leasing {nbytes} B of pre-registered memory failed ({e}); registering own buffer.")

In a constructed exhaustion scenario a single get() emitted 10 WARNING lines. Multiplied by
MAX_BATCH_WORKER_THREADS and by reads per step, that floods logs at WARNING level from the hot
path. Suggest warn-once, or debug per occurrence plus a single warning.

[0oshowero0]

【AI Review】

6. acquire(block=False) never returns None, so the docstring and one test cover an unreachable branch

The docstring says exhaustion "surfaces as None in some builds and as an exception in others", but
BufferPoolNative::acquire only ever throws — buffer pool is exhausted when block is false, or
requested buffer size exceeds pool capacity when the size class exceeds the budget. Meanwhile
FakePool.acquire returns None, so test_falls_back_to_own_buffers_when_pool_cannot_serve
exercises a branch that can't occur while the branch that actually runs in production is untested.
I verified the exception path does behave correctly (values round-trip, falls back to
register-per-read), so this is a test-coverage gap rather than a bug. Suggest making the fake raise
and dropping the speculative sentence.

[0oshowero0]

And please run precommit script:

# install pre-commit
pip install pre-commit

# run the following command in your repo folder, then fix the check before committing your code
pre-commit install && pre-commit run --all-files --show-diff-on-failure --color=always