Hacktoberfest 2026：メンテナが10月に向けて印を付けた、オープンで初心者向けの issue。 Hacktoberfest の issue を見る

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

オープン
  • 何か問題がありますか？
#169 コメント 3 件 リアクション 1 件 担当者 2 名 GitHub で見る

メンテナーはふだん 1 日以内に返信

@weiziyoung がすでに取り組んでいます。

2026年9月15日 から。

評価

この issue はまだ評価されていません。

説明

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.

主要言語
Python
スター
168
フォーク
56
平均マージ
4日 1時間
マージ済み PR（30日）
15

環境構築

このプロジェクトには開発コンテナ、Dockerfile、コントリビューションガイドがありません。まず README を読み、一般的な手順ははじめてのコントリビューションガイドを参照してください。

はじめの一歩

  1. issue を最後まで読み、次にプロジェクトのコントリビューションガイドを読みます。
  2. 着手することを issue にコメントします — 二人が同じ作業をするのを防げます。
  3. リポジトリをフォークし、ブランチを切って変更します。
  4. issue 番号を参照したプルリクエストを送ります。

Ascend/TransferQueue のほかの issue

Ascend/TransferQueue の issue をすべて見る

似ている issue

Python の issue をもっと見る

新しい issue をメールで受け取る

初心者向けの GitHub issue を短くまとめたダイジェスト。