Skip to content

Io prepared transfers v3 - #574

Open
amirakb89 wants to merge 5 commits into
ROCm:mainfrom
amirakb89:io-prepared-transfers-v3
Open

Io prepared transfers v3#574
amirakb89 wants to merge 5 commits into
ROCm:mainfrom
amirakb89:io-prepared-transfers-v3

Conversation

@amirakb89

@amirakb89 amirakb89 commented Aug 17, 2026

Copy link
Copy Markdown
Contributor

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_engine via --prepare-once. This mirrors
the createXferReq / postXferReq split nixl reports separately, so the timed
loop 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

  • Core RDMA path (src/io/rdma/common.cpp, backend_impl.cpp, engine.cpp,
    include/mori/io/{backend,engine}.hpp): new PreparedTransfer handle plus
    PrepareBatch / PostPrepared on IOEngineSession. The build phase is
    hoisted 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.
  • C++ benchmark (tests/cpp/io/bench_engine.cpp): --prepare-once. Needs
    --backend rdma + --enable-sess. Works with both submission modes —
    --enable-batch-transfer prepares one N-descriptor handle,
    --disable-batch-transfer prepares one single-descriptor handle per transfer.
  • Tests (tests/cpp/io/test_engine.cpp): prepared-path unit + end-to-end
    coverage.
  • Docs (docs/MORI-IO-BENCHMARK.md): documents --prepare-once as 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-once vs. the normal path:

Msg size Plain BW (GB/s) Prepare-once BW (GB/s) Plain lat (us) Prepare-once lat (us)
1 KB 1.35 1.41 0.76 0.72
4 KB 5.69 5.70 0.72 0.72
1 MB 38.58 39.03 27.18 26.87

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

  • Two-node MI308X/ionic: prepared batched, prepared singles, and non-prepared
    paths all byte-exact
  • --prepare-once with --disable-batch-transfer (one handle per transfer)
  • pre-commit clean (clang-format, etc.)
  • CI green on the runner

Akbarzadeh and others added 4 commits August 17, 2026 18:47
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
amirakb89 marked this pull request as ready for review August 17, 2026 21:53

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 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".

Comment thread src/io/rdma/common.cpp
Comment thread docs/MORI-IO-BENCHMARK.md Outdated
- 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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant