[optim] MooncakeStore read path registers RDMA buffers on every get(), dominating end-to-end time

[weiziyoung]

Problem

MooncakeStoreClient allocates a fresh receive buffer and calls register_buffer /
unregister_buffer around every transfer, on all four hot paths in
transfer_queue/storage/clients/mooncake_client.py:

• _get_tensors_thread_worker
• _get_bytes_thread_worker
• _put_tensors_thread_worker
• _put_bytes_thread_worker

Memory registration is a kernel operation (page pinning + MR setup), and
_register_all_buffers performs it in a serial loop. On our cluster, for a single
multi-GB read split into ~1.6k tensor keys, the breakdown is:

stage	share of read time
register_buffer	~85%
batch_get_into (actual RDMA)	~6%
metadata lookup + others	~9%

Measured separately: registration runs at ~2.8 GB/s while the RDMA transfer itself
reaches ~38 GB/s on the same NICs (ib_write_bw point-to-point ~386 Gb/s). So the
per-call registration is ~14x slower than the transfer it enables.

The practical consequence is counter-intuitive: with the same workload, the
MooncakeStore/RDMA backend came out ~2x slower than SimpleStorage (TCP/ZMQ),
purely because TCP needs no registration at all. This makes the RDMA backend look
worse than it is and is easy to misattribute to the hardware or to Mooncake.

Proposal

Register once at setup and reuse. Mooncake already exposes this:
mooncake.store.BufferPool (a.k.a. RegisteredBufferPool) hands out leases from the
local_buffer_size region that setup() already registered, so acquire() /
release() never enter the kernel.

This is also what the existing TODO in _put_bytes_thread_worker asks for:

# TODO: switch to a pre-registered buffer from MooncakeStore once such an API is available.

The API is available now (verified with mooncake-transfer-engine 0.3.13.post1;
the mooncake extra currently pins >=0.3.10.post2).

Shape-wise this mirrors what GdrStaging in transfer_queue/utils/mooncake_utils.py
already does for the GDR path — a process-level persistent buffer with
acquire/release — so the host-side read path would just follow the same pattern.

Results from a local prototype

Same script, only this implementation changed. Single reader, multi-GB read:

implementation	throughput	vs. current
current (register per read)	1.00x	—
pooled buffer, still copies on return	2.25x	2.25x faster
pooled buffer + zero-copy view	4.83x	4.83x faster

The middle row matters: pooling alone leaves half the win on the table, because the
copy-out of a multi-GB batch is itself ~100ms-scale.

In real async-RL training runs the pooled version was 1.5x-2.2x faster than
SimpleStorage at several read points, and step-to-step standard deviation dropped
by an order of magnitude, which matters for pipelined async training.

Correctness was checked with value-level (not shape-level) equivalence over tens of
thousands of per-column digests: no mismatches.

Two pitfalls worth designing around

1. Silent fallback on capacity shortfall. If local_buffer_size is smaller than
   bytes_per_batch x MAX_BATCH_WORKER_THREADS, BufferPool takes its documented
   overflow path (allocate + register per lease), i.e. it silently degrades back to
   today's slow path with no error. We think the client should check capacity and
   fail loudly (or at least warn) instead.
2. Lease lifetime for the zero-copy variant. Returning torch.frombuffer views
   over a leased region is where the 4.83x comes from, but the lease must not be
   recycled while a caller still holds the view. Recycling per batch produced stale
   data for us under concurrent readers. Tying the lease to the returned tensor's
   lifetime (e.g. weakref.finalize(tensor, lease.release)) looks like the only
   self-consistent option; recycling at the next get() entry relies on an
   assumption callers cannot be expected to honor.

Also note BufferPool.acquire defaults to block_on_exhaustion=True, which can
block forever once the pool is full — the failure policy needs to be explicit.

Suggested split

1. Pooled buffer + copy on return. Return semantics identical to today, no lifetime
   risk, already 2.25x.
2. Zero-copy views with lease lifetime bound to the returned tensors, plus the
   capacity check and the mooncake extra version bump.

Happy to send both PRs if this direction looks right to you. Unit tests can run
without RDMA hardware by mocking the store (in the spirit of
tests/test_yuanrong_client_zero_copy.py), asserting that registration happens once
rather than per read, that values round-trip correctly, and that an undersized
local_buffer_size raises instead of silently degrading.