Io prepared transfers v3 - #574
Open
amirakb89 wants to merge 5 commits into
Open
Conversation
Prepared transfers build a batch's work requests once and re-post them, moving the sort/merge/chunk cost out of the hot loop (the NIXL createXferReq / postXferReq split). The pair is deliberately self-contained: RdmaBatchReadWrite keeps its own scratch pools and posts inline, so the existing paths and the executor are untouched. Surfaced as IOEngineSession::PrepareBatch/PostPrepared.
Bind PrepareBatch/PostPrepared into the Python session and add --prepare-once, which prepares a batch per (buffer size, batch size) pair and times only the posts, isolating work-request construction from the transfer itself. The flag requires the rdma backend with session and batch transfers, since no other backend prepares batches today.
The prepared path was only reachable from the Python bench, so its numbers could not be lined up against the nixl-parity ones this tool reports. Build the sweep point's work requests once before warmup and re-post that handle each iteration, leaving sort/merge/chunk outside the timed loop. Rejected at startup without rdma, --enable-sess and --enable-batch-transfer, matching the Python checks.
amirakb89
marked this pull request as ready for review
August 17, 2026 21:53
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: e90aa3d0a9
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
- Revert the Python prepared-transfers path (engine.py, pybind_io.cpp, benchmark.py) to upstream; --prepare-once is now a C++ bench_engine feature only, using the C++ session API directly. - Fix P1: release pending unsignaled WRs when a prepared post batch's SQ-depth reservation fails. Earlier unsignaled batches held reservations with no signaled ledger tail to free them, leaking SQ depth into permanent SQ-full; now orphan them across all EPs, mirroring the ibv_post_send failure path. - Fix P2 docs: --prepare-once is C++-only and works with both --enable-batch-transfer and --disable-batch-transfer, not batch-required.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Title:
bench(io): add opt-in prepared transfers (build-once, post-many) to MORI-IO C++ path
Summary
Adds an opt-in prepared transfers path to MORI-IO's RDMA backend: the RDMA
work requests for a transfer (descriptor sort, contiguous merge, chunk planning)
are built once and re-posted many times, instead of being re-derived every
submission. Exposed in the C++
bench_enginevia--prepare-once. This mirrorsthe
createXferReq/postXferReqsplit nixl reports separately, so the timedloop measures the post alone.
Prepared transfers are a C++-only feature (the Python benchmark is
unchanged); the C++ benchmark uses the C++ session API (
PrepareBatch/PostPrepared) directly.What changed
src/io/rdma/common.cpp,backend_impl.cpp,engine.cpp,include/mori/io/{backend,engine}.hpp): newPreparedTransferhandle plusPrepareBatch/PostPreparedonIOEngineSession. The build phase ishoisted out of the post; the handle is owned by a session. Backends that don't
support it report the lack rather than silently falling back.
tests/cpp/io/bench_engine.cpp):--prepare-once. Needs--backend rdma+--enable-sess. Works with both submission modes —--enable-batch-transferprepares one N-descriptor handle,--disable-batch-transferprepares one single-descriptor handle per transfer.tests/cpp/io/test_engine.cpp): prepared-path unit + end-to-endcoverage.
docs/MORI-IO-BENCHMARK.md): documents--prepare-onceas C++-only,both submission modes.
Benchmark results
Two-node MI308X + ionic (AINIC), GPU memory, write, 1 initiator × 1 target,
8 QP, batch = 64, 100 iters.
--prepare-oncevs. the normal path:Prepared transfers are a latency / CPU-overhead optimization, not a bandwidth
one. The gain is largest at small message sizes (~4-5% at 1 KB), where
work-request construction is a larger share of per-iteration cost, and negligible
at large sizes where the transfer itself dominates. All configs validated
byte-identical.
Test plan
paths all byte-exact
--prepare-oncewith--disable-batch-transfer(one handle per transfer)