feat(io/rdma): reuse thread-local scratch buffers on batch hot path

[staryxchen]

Summary

RdmaBatchReadWrite() (src/io/rdma/common.cpp) allocates several std::vectors
on every batch transfer:

• indices (size = batchSize)
• mergedWrs plus a per-WR sges.reserve(maxSge) (≈ one heap allocation per WR)
• a freshly-returned vector from ExpandChunkedWorkRequests(std::move(...))
• epWrsSinceSignal, epMergedSinceSignal

For small/medium messages these allocations sit on the critical path before the
WRs are posted, so they directly add to per-batch latency and scale with batch size.

This PR converts those scratch containers to thread_local pools reused across
calls (slot reuse preserves each WR's inner sges capacity), expands chunked WRs
in place into a second pooled buffer (the chunk plan is reused via a new
PlanChunksInto() that fills a caller-owned buffer), and skips the chunk-expansion
pass entirely when no WR exceeds chunkBytes. Single-SGE WR initialization is shared
via one initSingleSgeWr helper.

Safety / hygiene addressed in this PR:

• Re-entrancy guard: the pools belong to the outermost call on a thread. A
  cheap RAII depth sentinel makes any same-thread re-entry (e.g. from a completion
  callback) transparently fall back to local buffers, so the outer call's pools are
  never clobbered. (No such re-entry path exists today; this is defensive.)
• Bounded retained memory: a high-water check releases pool capacity if an
  earlier very large batch grew it far beyond the current need.
• The now-unused ExpandChunkedWorkRequests() is removed.

No public header / ABI / wire-format change — internal to src/io/rdma/common.cpp.

Motivation / Measurements

On 2× MI308X (gfx942, ROCm 7.2.0), 2×200G bnxt RoCE, cross-node, GPU memory, 4 QP,
session + batch transfer, op-type write. A/B = same binary with vs. without this
change (initiator side), identical network config.

A microbench of just the per-batch allocation cost (batchSize=128) shows
~7.5 µs → ~0.9 µs (88 % of the allocation cost removed), scaling ~linearly with batch.

End-to-end benchmark (tests/python/io/benchmark.py):

batch	MsgSize	Avg Lat before	Avg Lat after	Δ latency	Δ BW
128	1 KB	47.37 µs	45.89 µs	-3.1 %	+3.2 %
128	2 KB	48.64 µs	46.44 µs	-4.5 %	+4.6 %
128	4 KB	52.06 µs	50.14 µs	-3.7 %	+3.9 %
512	1 KB	106.54 µs	99.19 µs	-6.9 %	+7.5 %
512	2 KB	114.45 µs	107.00 µs	-6.5 %	+7.0 %
512	4 KB	133.04 µs	124.46 µs	-6.4 %	+6.9 %
128/512	1 MB	unchanged	unchanged	~0 %	~0 %

The absolute saving grows with batch size (≈2 µs @128 → ≈7 µs @512), matching the
microbench's linear prediction. Large messages are bandwidth-bound and unaffected,
as expected.

Risk / Correctness

• Output WR list is byte-for-byte equivalent to before; the benchmark's built-in
  data validation (torch.equal) passes across the full size sweep for both
  write and the chunked path.
• The output WR list is byte-for-byte equivalent before/after; re-verified by the
  benchmark's torch.equal validation over the full size sweep after this revision.

Test plan

# two nodes; initiator side exercises the changed path
torchrun --nnodes=2 --node_rank=0 --nproc_per_node=1 \
  --master_addr=<n0> --master_port=1234 \
  tests/python/io/benchmark.py --host=<n0> --mem-type gpu \
  --enable-batch-transfer --enable-sess --buffer-size 16384 \
  --transfer-batch-size 512 --all --sweep-start-size 8 --sweep-max-size 16384 \
  --num-qp-per-transfer 4 --op-type write --poll_cq_mode polling

[Copilot]

Pull request overview

This PR optimizes the RDMA batch transfer hot path by reusing per-thread scratch containers inside RdmaBatchReadWrite() to avoid repeated heap allocations, and by performing chunk-expansion in-place only when needed. The changes are internal to src/io/rdma/common.cpp and preserve the existing PlanChunks() API while introducing a new in-place planning helper.

Changes:

• Add thread_local scratch pools (with a same-thread re-entrancy fallback) to reuse vectors across RdmaBatchReadWrite() calls.
• Replace ExpandChunkedWorkRequests() with an in-place chunk expansion into a pooled buffer, and skip the expansion pass entirely when no WR exceeds chunkBytes.
• Introduce PlanChunksInto() to fill a caller-owned chunk plan buffer while keeping the existing PlanChunks() API intact.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[staryxchen]

Hi @maning00, CI is green and this is ready for review whenever you have time. Thanks!