Hacktoberfest 2026: le issue che i maintainer hanno segnato per ottobre, aperte e adatte ai principianti. Sfoglia le issue Hacktoberfest

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

Aperta
  • Qualcosa non va?
#169 3 commenti 1 reazione 2 assegnatari Vedi su GitHub

I maintainer di solito rispondono entro 1 giorno

@weiziyoung ci sta già lavorando.

Dal 15/9/2026.

Valutazione

Questa issue non è ancora stata valutata.

Descrizione

enhancement
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.

Lingua principale
Python
Stelle
166
Fork
56
Merge medio
4g 1h
PR unite (30g)
15

Preparare l'ambiente

Questo progetto non fornisce container di sviluppo, Dockerfile né guida per i contributori, quindi l'ambiente è a tuo carico: parti dal suo README e consulta la nostra guida al primo contributo per i passaggi generali.

Come iniziare

  1. Leggi tutta la issue e poi la guida ai contributi del progetto.
  2. Commenta sulla issue per dire che te ne occupi tu — evita che due persone facciano lo stesso lavoro.
  3. Fai un fork del repository e lavora su un branch.
  4. Apri una pull request che faccia riferimento al numero della issue.

Altre issue di Ascend/TransferQueue

Tutte le issue di Ascend/TransferQueue

Issue simili

Altre issue su Python

Ricevi le nuove issue nella tua casella

Un breve riepilogo di issue GitHub adatte ai principianti.