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

[weiziyoung]

Refs: #169

What this changes

Tensor reads registered and unregistered a receive buffer around every transfer.
register_buffer is a kernel operation (page pinning + MR setup) whose throughput is an
order of magnitude below the RDMA transfer it enables, so it dominated read time and left
the RDMA backend slower than the TCP one.

This leases the receive buffer from the local buffer that store.setup() already
registered, then copies into the caller's tensors so the returned tensors keep owning their
memory. split_by_bytes() bounds each lease by one reader thread's share of the pool, so a
read no longer needs local_buffer_size to exceed the batch size.

Three cases still register their own receive buffers: builds without
mooncake.store.BufferPool (mooncake-transfer-engine < v0.3.12), a group whose aligned
total exceeds the per-thread share, and a pool that cannot serve the request.

How the pool is sized

max_bytes=local_buffer_size keeps in-flight leases inside the region setup() registered.
On the default max_bytes=0 mooncake budgets twice the local buffer and serves the excess
through allocate_overflow_unlocked() — posix_memalign plus one register_buffer per
lease — which is exactly the cost this path removes, and it would happen silently.

max_regions=MAX_BATCH_WORKER_THREADS bounds concurrent leases, because each get() starts
its own reader pool and nothing else limits how many run at once.

The per-thread share is local_buffer_size // (2 * MAX_BATCH_WORKER_THREADS) — half of an
equal split. The headroom is on our side rather than the pool's because BufferPoolNative
clamps any max_bytes below local_buffer_capacity_ back up to it, so max_bytes can only
bound the overflow arena, never reserve space inside the registered region. mooncake
allocates from that same region for its own read-side staging: remote DISK/DFS replica reads
in batch_get_into_multi_buffers(), and the host staging execute_ranged_read() needs when
the destination is device memory. Splitting the whole buffer across the reader threads would
make those allocations fail — and inside one batch_get_into call, a failure returns
NO_AVAILABLE_HANDLE while the lease is still held, so the 3 retries cannot help.

Benchmarks

Sections 1-3 were measured with the sizing as first submitted (share = 1/4 of the local
buffer, max_bytes=0). Section 4 re-measures the sizing change on a different machine.

1. Tuning BATCH_SIZE_LIMIT does not substitute for the pool

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 and grows worse as the batch grows: register_buffer cost
is dominated by per-byte page pinning, not per-call overhead, so a larger batch does not
reduce the pinning work — it only reduces how many worker threads share it.

2. The pool removes registration entirely

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

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×)

3. Where the win is big, and where it is 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 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 copy-out scales with
  bytes, so a 4GB read narrows the gain to ~1.08×. The pool also lowers step-to-step jitter
  (1024×1MB threads=1 stdev 19.6 → 7.8 ms).
• Many tiny keys can regress. For 16384×64KB the mandatory copy-out does not pay off
  against that shape's cheap merged registration. Vectorizing the copy-out (uniform groups
  copied in one strided pass instead of a per-tensor loop) cut this case from ~2× slower to
  ~0.8×, but it is still a net loss at extreme fragmentation. A num_keys/per-key-size
  heuristic could fall back to register-per-read here; real trajectory reads (≈1024×1MB) do
  not hit it.

The copy is what keeps returned tensors caller-owned so the buffer returns immediately. A
zero-copy view mode would restore the large-read win but requires a "no cross-read retention"
contract, so the copy path is the safe default.

4. What the reserved headroom costs

protocol=rdma (mlx5), single node, local_buffer_size=1GB, threads=4,
BATCH_SIZE_LIMIT=400, median of 10 reads ± stdev. All three modes are timed round-robin
within each rep so machine drift is shared:

workload	register-per-read	share = 1/4	share = 1/8 (this PR)
64×16MB	323.0 ± 7.3	96.6 ± 13.1 (3.34×)	92.1 ± 3.7 (3.51×)
1024×1MB	209.5 ± 40.9	115.6 ± 5.7 (1.81×)	118.4 ± 5.0 (1.77×)
4096×1MB	598.8 ± 56.4	417.7 ± 36.8 (1.43×)	375.7 ± 70.1 (1.59×)
16384×64KB	384.8 ± 39.0	466.7 ± 69.1 (0.82×)	425.2 ± 48.0 (0.90×)

Leases taken per read confirm the extra rounds (4→8, 5→10, 21→41; 16384×64KB stays at 41
because BATCH_SIZE_LIMIT already caps each batch below both shares), but the extra rounds do
not show up as time: three of four shapes are slightly faster with the smaller share, and
1024×1MB is 0.04× slower — well inside stdev. The headroom is effectively free.

This run also reproduces two results from section 3 on different hardware: the large win on
64×16MB (3.9× there, 3.34-3.51× here) and the regression on 16384×64KB (0.8× there, 0.82×
here, softened to 0.90× by the smaller share).

Caveats: single node, so the transfer itself may be served by a local copy rather than RDMA —
the eliminated registration and the round counts are real, but absolute numbers are not
comparable to a cross-node run. Timing calls MooncakeStoreClient.get directly rather than
scripts/performance_test, and the machine had residual load. Treat the three modes against
each other as the result, not the absolute milliseconds.

Compatibility

The register-per-read path is preserved as a fallback: the client gates on
mooncake.store.BufferPool availability, so builds without it (mooncake-transfer-engine <
v0.3.12) transparently keep registering their own receive buffers. The "upgrade" warning is
emitted once, and only where CPU RDMA reads would pay for registration — TCP and GDR reads
stay quiet. Requests the pool cannot serve also fall back, warning once rather than per group
per read.

The GDR read/write paths and the non-tensor paths are untouched.

Verification

tests/test_mooncake_buffer_pool.py fakes the store and the pool, so it needs neither
mooncake nor RDMA. The fake pool keeps the properties this path depends on: buffers are carved
out of one arena and reused, exhaustion raises (which is what acquire(block=False) actually
does — it never returns None), a buffer cannot be returned while an exported view of it is
alive, and the pool cannot close while a buffer is out on loan.

Covered: pooled reads with mixed dtypes and scalars, the strided copy-out for uniform groups,
multi-round reads when the batch exceeds the share, register fallback for an oversized tensor /
an exhausted pool / a build without BufferPool, buffer reuse not disturbing tensors returned
by an earlier read, a failed transfer and a failed copy-out both returning the buffer and
preserving the original exception, concurrent readers staying within the budget and the region
cap, warn-once, the warning's scope, and close() ordering.

Separately verified against a real BufferPool (TCP, in-process): the pool accepts these
arguments, an over-capacity acquire raises requested buffer size exceeds pool capacity,
the region cap raises buffer pool is exhausted, and close() is refused while a buffer is
leased.

[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?

[weiziyoung]

Have renamed to tests/test_mooncake_buffer_pool.py.

[0oshowero0]

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

[weiziyoung]

I have renamed to _copy_buffer_content_into_tensors.

[0oshowero0]

We can simplify the AI comments

[weiziyoung]

Trimmed~

[0oshowero0]

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

[0oshowero0]

For GDR and TCP path, we can omit this warnning

[weiziyoung]

The warning now names mooncake-transfer-engine >= 0.3.12.

[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 :)

[weiziyoung]

make sense~

[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

[ascend-robot]

CLA Signature Pass

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

[weiziyoung]

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

I have ran pre-commit run --all-files: ruff, ruff-format and mypy all pass, and no files were modified.

[0oshowero0]

[AI Review]

Overall the direction is right: the change stays inside the storage client, keeps the register-per-read path for older mooncake builds, and the previous round (max_bytes, headroom, preserving the original exception, warn-once) is addressed. The unit tests pass locally (51 passed, 5 skipped).

Before merge I'd address three things (details inline):

1. The pool is enabled on TCP, the default protocol, where by this PR's own reasoning there is no registration cost to save.
2. The fallback registers and transfers group by group, so oversized tensors and pool exhaustion end up slower than before this PR.
3. Zero-numel tensors raise ValueError on the pooled path.

The rest is simplification: the copy-out can drop the uniform/ragged split and get faster on ragged shapes, two helpers can take simpler arguments, the close() wrapper can go, and the tests can go through get() with a much smaller fake.

One residual risk worth stating in the PR description: when mooncake's own staging holds more than half the local buffer, a lease can still be served by mooncake's overflow allocation (fresh posix_memalign + register_buffer), and nothing on the Python side can tell. The logs also don't say whether the pool path is in use at all.

What I checked: issues 1-3 are reproduced with the fakes from this PR's test file, and the copy-out numbers come from a local CPU microbenchmark. I did not run against a real BufferPool or RDMA, so TCP timings are not measured.

[0oshowero0]

[AI Review] The pool is also enabled for TCP, the default protocol.

This only checks MOONCAKE_BUFFER_POOL_IMPORTED, while the comment a few lines below says TCP reads don't benefit, and config.yaml defaults to protocol: tcp. On TCP this path therefore adds a full copy-out per read and splits each batch into ceil(batch_bytes / share) sequential batch_get_into calls per thread. With the 128 MB local_buffer_size used by the e2e config the share is 16 MB, so a 400 x 1 MB batch goes from 1 call to 25. All benchmarks in the description are RDMA.

Suggest gating on the same condition the warning uses, so one condition drives both:

use_pool = MOONCAKE_BUFFER_POOL_IMPORTED and self.protocol == "rdma"

If TCP should use the pool too, a TCP benchmark showing no regression would justify it.

[0oshowero0]

[AI Review] The fallback runs one register + batch_get_into per group, which is slower than before this PR.

Every group that can't get a lease does its own registration and transfer. With the fakes from this PR, 8 float32 tensors that each exceed the share take 8 register_buffer + 8 batch_get_into calls, serially on one thread; before this PR the same batch was 1 + 1 (one merged region, one transfer). That triggers for any key above 128 MB with the default 1 GB buffer (above 16 MB with 128 MB), and on pool exhaustion every group of the batch falls back separately.

Suggest collecting the fallback indices and reading them once after the loop:

fallback: list[int] = []
for group in split_by_bytes(batch_nbytes, self._lease_bytes):
    if not self._read_via_buffer_pool(group, ...):
        fallback.extend(group)
if fallback:
    self._read_into_registered_tensors(
        [batch_keys[i] for i in fallback],
        [batch_buffer_ptrs[i] for i in fallback],
        [batch_nbytes[i] for i in fallback],
    )

That keeps the worst case at the pre-PR cost and leaves _read_via_buffer_pool with a single job (this uses the 3-argument _read_into_registered_tensors suggested on line 493).

[0oshowero0]

[AI Review] Zero-numel tensors raise here.

torch.frombuffer(..., count=0) raises ValueError: both buffer length (256) and count (0) must not be 0. Reading keys with shapes [(0,), (4,)] through the fakes in this PR fails after the transfer on the pooled path, while the register-per-read path returns both tensors. The uniform branch has the same problem when numel == 0.

Whether this is reachable depends on mooncake accepting a 0-byte put, which I haven't verified. A guard (skip zero-byte targets) plus a test would close it; the per-region copy suggested on line 57 doesn't have this problem.

[0oshowero0]

[AI Review] The copy-out could be simpler and faster by mirroring the destination layout.

The uniform fast path only fires when every tensor in the group has the same dtype and shape and sits back to back in storage, and it depends on allocate_empty_tensors internals (storage_offset, untyped_storage().data_ptr()). Ragged shapes, which are common for variable-length RL data, take the per-tensor loop, which creates a fresh lease.buffer export plus frombuffer / view / copy_ for every key.

Alternative: lay the keys out in the lease back to back in destination-pointer order, without the 256-byte alignment (CPU reads don't need it; the register path already reads into tightly packed tensors). Each region from merge_contiguous_memory is then one contiguous span of the lease and needs one copy:

order = sorted(group, key=lambda i: batch_ptrs[i])
nbytes = [batch_nbytes[i] for i in order]
region_ptrs, region_sizes = merge_contiguous_memory([batch_ptrs[i] for i in order], nbytes)
lease_ptrs = list(itertools.accumulate(nbytes[:-1], initial=lease.ptr))
# ... batch_get_into(keys_in_order, lease_ptrs, nbytes) ...
src = torch.frombuffer(lease.buffer, dtype=torch.uint8)
try:
    off = 0
    for dst_ptr, size in zip(region_ptrs, region_sizes, strict=True):
        storage = ...  # untyped_storage() of the target tensor that starts at dst_ptr
        dst = torch.empty(0, dtype=torch.uint8).set_(storage, dst_ptr - storage.data_ptr(), (size,))
        dst.copy_(src[off : off + size])
        off += size
finally:
    del src

Local CPU microbenchmark (macOS, 256 MB per case, best of 5; treat as relative numbers only):

shape	current copy-out	per-region copy
4096 x ~64 KB, ragged	51 ms	8.5 ms
256 x ~1 MB, ragged	25 ms	8 ms
uniform (64 KB / 1 MB / 16 MB keys)	9-11 ms	9-11 ms

This removes the uniform/ragged split, the _aligned_offsets dependency, and the per-tensor del handling (one export per group). Two notes: ctypes.memmove was ~3x slower than torch copy_ for large spans here because it is single-threaded, so keep the torch copy; and split_by_bytes groups by size, which scatters ragged keys across groups in multi-round reads, so cutting groups in destination order would keep each group at one span per dtype region.

It also suggests the 16384 x 64 KB regression in the description, measured with uniform shapes, is the optimistic case for variable-length data.

[0oshowero0]

[AI Review] The two read helpers could have simpler signatures.

• _read_via_buffer_pool spends half its body on the register fallback, so its name doesn't say what it does, and it takes group indices plus whole-batch arrays. With the fallback moved out (see line 523), it can do only the pooled read and return whether it got a lease.
• _read_into_registered_tensors takes region_ptrs / region_sizes, which merge_contiguous_memory(ptrs, nbytes) derives from its other arguments (the second call site already computes them that way). Computing them inside reduces the signature to (keys, ptrs, nbytes).

[0oshowero0]

[AI Review] Naming and log text.

• _buffer_pool_exhausted_logged is set for any acquire failure, including the max_regions cap, and _lease_bytes is the per-reader share (that's what the tests and the description call it). Names like _pool_fallback_warned / _pool_share_bytes would match what they hold.
• This warning always tells the user to raise local_buffer_size, but the cause can also be the region cap; including the first exception text would make it actionable.
• A few comments narrate the review history or the next line (e.g. # warn once and keep the details at debug level., the 7-line docstring on _copy_buffer_content_into_tensors); AGENTS.md asks for 2-4 lines of "why".

[0oshowero0]

[AI Review] This try/except only changes the error message.

If self._buffer_pool.close() raises, the exception propagates and the store isn't closed either way, so the wrapper only replaces mooncake's message. Nothing in the repo calls MooncakeStoreClient.close() today. Suggest just:

if self._buffer_pool is not None:
    self._buffer_pool.close()
    self._buffer_pool = None

Closing the pool before the store is the part worth keeping; test_close_keeps_the_store_open_while_a_buffer_is_out can then go.

[0oshowero0]

[AI Review] The tests could go through the public API with a much smaller fake.

• read() calls the private _get_tensors_thread_worker, and several tests assert private attributes (_lease_bytes, _buffer_pool). Going through client.get(keys, shapes, dtypes) also covers the thread-pool dispatch and keeps the tests valid if the helpers change as suggested above.
• FakePool reimplements first-fit allocation, alignment and weakref export tracking: 458 lines of tests for ~158 lines of code. A fake that hands out slices of one arena and raises when it runs out covers the behavior the client depends on.
• test_concurrent_readers_stay_within_the_local_buffer (line 388) uses 8 concurrent readers, but KVStorageManager.get_data calls get() synchronously, so one client has at most MAX_BATCH_WORKER_THREADS readers. test_failed_copy_out_returns_the_buffer_and_keeps_the_cause (line 369) patches torch.Tensor.copy_ globally for a failure that is hard to hit. Both cover states the supported paths don't reach in practice.
• Missing coverage: zero-numel tensors, protocol gating (if adopted), and ragged shapes on the copy-out.

[weiziyoung]

Thanks — agreed the per-region copy-out is the better design, and the ragged
numbers make the case. I'm going to land it in a follow-up PR rather than this one,
for two reasons:

1. It co-changes the read, not just the copy. To copy once per contiguous
   region the lease has to be laid out in destination-pointer order, so
   _read_via_buffer_pool and the batch_get_into offsets change together with
   the copy. That's a rewrite of the hot read path, not a local edit.

2. The key simplification — dropping 256-byte alignment — is a correctness
   claim the current tests can't cover. The fakes use memmove, which has no
   alignment requirement, so a green test suite wouldn't prove a real RDMA
   batch_get_into accepts tightly-packed lease offsets. I'd rather verify that
   against a real BufferPool (and reproduce the ragged speedup on our hardware)
   before shipping it, instead of bundling an unverifiable change into this
   round's blocker fixes.

So this PR addresses the three merge blockers (TCP gating, batched fallback,
zero-numel guard) plus the naming/close()/comment cleanups. The copy-out
rewrite — which also removes the uniform/ragged split and the _aligned_offsets
dependency, and resolves the zero-numel case structurally — will be its own PR
with the benchmark and the real-backend verification. Tracking it so it doesn't
get lost.

[0oshowero0]

Thanks — agreed the per-region copy-out is the better design, and the ragged numbers make the case. I'm going to land it in a follow-up PR rather than this one, for two reasons:

1. It co-changes the read, not just the copy. To copy once per contiguous
   region the lease has to be laid out in destination-pointer order, so
   _read_via_buffer_pool and the batch_get_into offsets change together with
   the copy. That's a rewrite of the hot read path, not a local edit.
2. The key simplification — dropping 256-byte alignment — is a correctness
   claim the current tests can't cover. The fakes use memmove, which has no
   alignment requirement, so a green test suite wouldn't prove a real RDMA
   batch_get_into accepts tightly-packed lease offsets. I'd rather verify that
   against a real BufferPool (and reproduce the ragged speedup on our hardware)
   before shipping it, instead of bundling an unverifiable change into this
   round's blocker fixes.

So this PR addresses the three merge blockers (TCP gating, batched fallback, zero-numel guard) plus the naming/close()/comment cleanups. The copy-out rewrite — which also removes the uniform/ragged split and the _aligned_offsets dependency, and resolves the zero-numel case structurally — will be its own PR with the benchmark and the real-backend verification. Tracking it so it doesn't get lost.

Thanks a lot for the quick and thorough follow-up. The RDMA-only gating, the batched register fallback and the zero-numel guard all look good to me, and moving the per-region copy-out to a separate PR makes sense. I have no further comments on my side.

Is this PR ready to merge from your end, or is there anything else you'd like to land first?

[weiziyoung]

Thanks! Nothing else I'd like to land here — the per-region copy-out will come
as its own follow-up PR with the real-backend verification and the ragged
benchmarks.

I just merged latest main into the branch to resolve the conflict with #179
(the per-thread ACL/NPU context fix). The two changes are orthogonal: #179's
_npu_device capture and the executor initializers merged cleanly, and the only
hand-resolved spot was keeping both the pool setup in init and #179's
_bind_npu_device method side by side — no logic changed on either side. Local
checks pass (test_mooncake_buffer_pool.py + test_mooncake_client_thread_context.py,
compileall, ruff/mypy).

Ready to merge from my side once CI goes green on the merge result. I can't run
the MooncakeStore e2e in my environment, so I'm relying on the PR's CI for that
path. Happy for you to squash if you prefer a single commit.