EP support for rail isolated topology through Host based RDMA

[itej89]

Motivation

Enable Expert Parallelism (EP) on clusters where GPU-initiated IBGDA doorbells don't work or the vendor DV library is unavailable.

Technical Details

Adds TransportType::PROXY = 3 to the existing dispatch framework. EP kernels route through DISPATCH_TRANSPORT_TYPE → template specializations in shmem_proxy_kernels.hpp, leaving shmem_ibgda_kernels.hpp untouched. Proxy code compiles only when -DMORI_PROXY_ENABLED is set through JIT pipeline.

Per-NIC CPU threads poll host-pinned ring buffers (64K slots × 128B, hipHostRegister mapped). GPU warps write ProxyCmd descriptors; CPU threads convert to ibv_post_send with per-NIC lkey/rkey override. Signal delivery uses RDMA_WRITE for data signals and SEND_WITH_IMM + CPU-side __atomic_fetch_add for barrier atomics (ionic/bnxt). CX7 uses native IBV_WR_ATOMIC_FETCH_AND_ADD for barrier atomics.

Supported NIC backends: Pensando AINIC (ionic), Mellanox CX7 (mlx5), Broadcom Thor2 (bnxt).

When MORI_EP_OVER_RDMA=1, the DV API availability check is skipped for all providers — proxy uses plain ibverbs and does not require vendor DV libraries.

How to Run

export MORI_EP_OVER_RDMA=1
export MORI_RDMA_DEVICES=<nic_list>
export MORI_IB_GID_INDEX=<gid>
export MORI_SHMEM_HEAP_SIZE=4G
export MORI_SOCKET_IFNAME=<ctrl_nic>

torchrun --nnodes=2 --node_rank=<0|1> --nproc_per_node=1 \
  --master_addr=<master_ip> --master_port=1234 \
  examples/ops/dispatch_combine/test_dispatch_combine_internode.py \
  --cmd bench --num-qp 1

Test Results

Commit: 9ae6ebdf — 2-node EP16 benchmark

Proxy mode (MORI_EP_OVER_RDMA=1)

Metric	AINIC (ionic)	CX7 (mlx5)	Broadcom (bnxt)
Dispatch Best (GB/s)	39.78	53.47	63.52
Dispatch Worst (GB/s)	28.57	36.13	43.46
Combine Best (GB/s)	69.93	71.09	82.72
Combine Worst (GB/s)	40.10	44.68	52.42

Native IBGDA (no proxy, CX7 only)

Metric	CX7 (mlx5)
Dispatch Best (GB/s)	70.97
Combine Best (GB/s)	72.42

AINIC native: IBGDA hangs (GPU-written WQEs not visible to NIC on VFIO platform).
Broadcom native: DV library not available on test cluster (RDMA uABI 1).

Clusters tested

Cluster	GPU	NIC
DigitalOcean AINIC	MI350X (gfx950)	Pensando AINIC (ionic)
AMD MI300X	MI300X (gfx942)	Mellanox CX7 (mlx5)
Ruby (rckg)	MI350X (gfx950)	Broadcom Thor2 (bnxt)

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

[jhchouuu]

Hi @itej89, Thanks for the refactor!
Could you fix the conflict before merge?
And could you have a test on crusoe cluster again？

And three of my earlier comments look only partly addressed — following up on each rather than re-opening the threads:

---

1. proxy_thread.cpp:58 — bounds check on the wire pointer

The payload.addr == 0 guard helps, but a null check isn't the property I was asking about. A native IBV_WR_ATOMIC_FETCH_AND_ADD has the responder NIC validate the target address against the MR named by the rkey. A SEND carries no such constraint, so the emulation drops that validation entirely and the CPU does an 8-byte add wherever the payload points. Any wire corruption or a heap-layout mismatch between ranks becomes an arbitrary write into the local process.

GpuStates already tracks heapBaseAddr / heapEndAddr (internal.hpp:128-129); passing those into ProxyThread::Init and rejecting out-of-range payloads would restore the check cheaply. Was there a reason not to?

2. proxy_thread.cpp:140-145 — IBV_SEND_INLINE inline data

This one looks unchanged. PROXY_RDMA_WRITE_INLINE sets IBV_SEND_INLINE but leaves sge.addr as the raw cmd->src_addr assigned at line 122, which is a GPU address — so libibverbs memcpys the payload from a GPU pointer in the proxy thread's host context.

The answer is in the next case down: PROXY_SIGNAL_WRITE (:146-155) correctly redirects sge.addr into the host-side InlineBuf before setting the flag. PROXY_RDMA_WRITE_INLINE needs the same treatment. Its producer is shmem_proxy_kernels.hpp:151.

3. proxy_device_primitives.hpp:158 — fetch atomics returning 0

The fix in DrainCq is gated on qph.use_native_atomics:

if (ring_->cmds[slot].op == PROXY_ATOMIC_FETCH_ADD && qph.use_native_atomics) {
  ring_->cmds[slot].result = *reinterpret_cast<volatile uint64_t*>(ring_->cmds[slot].src_addr);
}

and use_native_atomics is vendorId == Mellanox (init.cpp:784). So ionic and bnxt — the two NICs this PR is aimed at — still take the SEND_WITH_IMM path and still get 0 back. The emulated path would need the responder to send the pre-add value back, or ProxyPostAtomicFetch should fail loudly rather than return a wrong value silently.

---

4. New, and I think it should block: proxy_thread.cpp:96 breaks non-x86 builds even with the feature off

asm volatile("clflush (%0)" :: "r"(target) : "memory");
asm volatile("sfence" ::: "memory");

There's no #if defined(__x86_64__), and CMakeLists.txt:66 adds proxy_thread.cpp to the source list unconditionally — so an aarch64 build fails whether or not MORI_ENABLE_HOST_PROXY is set. Everything else in this PR is cleanly opt-in (git diff -w shows zero removed lines in init.cpp, and symmetric_memory.cpp's native branch is byte-identical to main), so this is the one place the feature reaches past its own switch.

Guarding the asm is enough; __builtin_ia32_clflush plus std::atomic_thread_fence would also work. Worth noting the clflush is a no-op for its stated purpose anyway if the target is uncached device memory.

[itej89]

Hi @itej89, Thanks for the refactor! Could you fix the conflict before merge? And could you have a test on crusoe cluster again？

And three of my earlier comments look only partly addressed — following up on each rather than re-opening the threads:

1. proxy_thread.cpp:58 — bounds check on the wire pointer

The payload.addr == 0 guard helps, but a null check isn't the property I was asking about. A native IBV_WR_ATOMIC_FETCH_AND_ADD has the responder NIC validate the target address against the MR named by the rkey. A SEND carries no such constraint, so the emulation drops that validation entirely and the CPU does an 8-byte add wherever the payload points. Any wire corruption or a heap-layout mismatch between ranks becomes an arbitrary write into the local process.

GpuStates already tracks heapBaseAddr / heapEndAddr (internal.hpp:128-129); passing those into ProxyThread::Init and rejecting out-of-range payloads would restore the check cheaply. Was there a reason not to?

2. proxy_thread.cpp:140-145 — IBV_SEND_INLINE inline data

This one looks unchanged. PROXY_RDMA_WRITE_INLINE sets IBV_SEND_INLINE but leaves sge.addr as the raw cmd->src_addr assigned at line 122, which is a GPU address — so libibverbs memcpys the payload from a GPU pointer in the proxy thread's host context.

The answer is in the next case down: PROXY_SIGNAL_WRITE (:146-155) correctly redirects sge.addr into the host-side InlineBuf before setting the flag. PROXY_RDMA_WRITE_INLINE needs the same treatment. Its producer is shmem_proxy_kernels.hpp:151.

3. proxy_device_primitives.hpp:158 — fetch atomics returning 0

The fix in DrainCq is gated on qph.use_native_atomics:

if (ring_->cmds[slot].op == PROXY_ATOMIC_FETCH_ADD && qph.use_native_atomics) {
  ring_->cmds[slot].result = *reinterpret_cast<volatile uint64_t*>(ring_->cmds[slot].src_addr);
}

and use_native_atomics is vendorId == Mellanox (init.cpp:784). So ionic and bnxt — the two NICs this PR is aimed at — still take the SEND_WITH_IMM path and still get 0 back. The emulated path would need the responder to send the pre-add value back, or ProxyPostAtomicFetch should fail loudly rather than return a wrong value silently.

4. New, and I think it should block: proxy_thread.cpp:96 breaks non-x86 builds even with the feature off

asm volatile("clflush (%0)" :: "r"(target) : "memory");
asm volatile("sfence" ::: "memory");

There's no #if defined(__x86_64__), and CMakeLists.txt:66 adds proxy_thread.cpp to the source list unconditionally — so an aarch64 build fails whether or not MORI_ENABLE_HOST_PROXY is set. Everything else in this PR is cleanly opt-in (git diff -w shows zero removed lines in init.cpp, and symmetric_memory.cpp's native branch is byte-identical to main), so this is the one place the feature reaches past its own switch.

Guarding the asm is enough; __builtin_ia32_clflush plus std::atomic_thread_fence would also work. Worth noting the clflush is a no-op for its stated purpose anyway if the target is uncached device memory.

Hi @jhchouuu, thank you for the feedback. I've made the requested changes as follows:

1. Fixed - heap range passed to ProxyThread::Init, payload.addr validated against [heapBaseAddr, heapEndAddr) before dereference.

2. Fixed - inline_data[48] field added to ProxyCmd. GPU stages data there, CPU reads from it. Tagged with ProxyInlineTag

3. Enabled - Implemented the full-round trip. ender sends PROXY_IMM_ATOMIC_FETCH, receiver does atomic and replies with PROXY_IMM_ATOMIC_REPLY carrying the old value. ~18% combine overhead vs fire-and-forget when forced on all atomics; zero overhead on normal path (ProxyPostAtomicNonFetch stays fire-and-forget)

4. Fixed - replaced with std::atomic_thread_fence(std::memory_order_seq_cst)

[itej89]

cc: @jhchouuu

EP-Over-HOST-PROXY Verification Report - PR(1781fa1) vs main

Branch: feat/ep-rdma-sharing
Date: 2026-08-29
Test: test_dispatch_combine_internode.py --cmd bench (2-node, 8 GPU/node, 16 ranks total)

Results

Dispatch Performance (bfloat16)

Cluster	Commit	Mode	Best (GB/s)	Worst (GB/s)	Avg (GB/s)	Best Latency (us)	Worst Latency (us)	Avg Latency (us)
CX7 MI300X	7c51d18f (main)	Native (IBGDA)	67.87	9.44	48.69	1724	12393	3358
CX7 MI300X	1781fa1a (PR #558)	Native (IBGDA)	71.08	1.36	37.57	1647	86046	17413
CX7 MI300X	7c51d18f (main)	Rail-only	70.77	1.47	45.26	1654	79739	10210
CX7 MI300X	1781fa1a (PR #558)	Rail-only	70.98	1.26	45.92	1649	93179	9555
CX7 MI300X	1781fa1a (PR #558)	Proxy	56.23	36.13	44.46	2082	3238	2659
DO AINIC	1781fa1a (PR #558)	Proxy	50.95	32.27	40.63	2296	3623	2938
Crusoe AINIC	1781fa1a (PR #558)	Proxy	40.52	26.24	31.81	2887	4455	3753

Combine Performance (bfloat16)

Cluster	Commit	Mode	Best (GB/s)	Worst (GB/s)	Avg (GB/s)	Best Latency (us)	Worst Latency (us)	Avg Latency (us)
CX7 MI300X	7c51d18f (main)	Native (IBGDA)	71.86	9.60	55.60	1627	12182	2890
CX7 MI300X	1781fa1a (PR #558)	Native (IBGDA)	72.45	1.37	50.48	1615	85609	12816
CX7 MI300X	7c51d18f (main)	Rail-only	72.38	1.47	52.38	1616	79530	8249
CX7 MI300X	1781fa1a (PR #558)	Rail-only	72.24	1.58	51.69	1619	74113	7128
CX7 MI300X	1781fa1a (PR #558)	Proxy	71.53	38.41	55.88	1636	3046	2129
DO AINIC	1781fa1a (PR #558)	Proxy	75.71	39.66	54.35	1544	2949	2236
Crusoe AINIC	1781fa1a (PR #558)	Proxy	73.80	37.29	54.43	1586	3137	2263

Analysis

• Best-case throughput is identical across main and our branch (67-72 GB/s dispatch, 71-76 GB/s combine) — no regression in steady-state performance.
• Average differences are driven by JIT cold-start outliers (worst-case 12-93 ms) that vary per run. Rail-only averages are within noise between main and ours.
• Proxy mode has the most consistent latency — worst-case is only 3-4.5 ms vs 12-93 ms for Native/Rail-only, because the proxy path doesn't trigger GPU-side JIT recompilation on the IBGDA control path.

Cluster Details

Cluster	GPU	NIC	Firmware	Nodes
DO AINIC	MI350X (gfx950)	AMD Pensando AINIC	1.117.5-a-56	fabric-2, fabric-5
CX7 MI300X	MI300X (gfx942)	Mellanox CX7	mlx5	useocpm2m-097-026, useocpm2m-097-032
Crusoe AINIC	MI355X (gfx950)	AMD Pensando AINIC	1.117.5-a-77	crsuse2-m2m-v2-027, crsuse2-m2m-v2-030

[jhchouuu]

LGTM overall — thanks for working through all of these. Just the one last item above to sort out.

One other thing: async_ll still doesn't seem to work under the proxy — it faults reproducibly right after the first dispatch. v1 and v1_ll are both fine. Is that expected / out of scope for this PR, or worth a look before merging?

[jhchouuu]

@itej89 Would you please take a look at this items?

Blocking: the atomic-fetch reply WR has no wr_id, so its completion corrupts ring slot 0

160:  ibv_send_wr rwr{}, *rbad = nullptr;                    // wr_id zero-initialised
163:  rwr.send_flags = IBV_SEND_SIGNALED | IBV_SEND_INLINE;  // always produces a send CQE
166:  ibv_post_send(qph.qp, &rwr, &rbad);                    // wr_id never assigned

send_cq == recv_cq (ibverbs.cpp:213-214), so this completion lands in the CQ DrainCq polls. Opcode is IBV_WC_SEND, so it skips the RECV branch and hits the tail:

185:  uint32_t slot = static_cast<uint32_t>(wc[i].wr_id) & PROXY_RING_MASK;  // = 0
195:  ring_->cmds[slot].status = PROXY_COMPLETED;

So every atomic-fetch reply marks slot 0 of the responder's own ring as completed. Worst case isn't a lost completion — if slot 0 is PENDING but not yet consumed, MainLoop's if (cmd->status != PROXY_PENDING) break; fires without advancing next_slot_, so that ring stops posting permanently while the GPU's ProxyQuiet sees COMPLETED and moves on. Same route at :106 would mark slot 0 PROXY_ERROR.

Fix — high-bit sentinel, tested before the mask (~0ull & PROXY_RING_MASK is 65535, a valid slot):

static constexpr uint64_t PROXY_WRID_INTERNAL = 1ull << 63;   // proxy_types.hpp
rwr.wr_id = PROXY_WRID_INTERNAL;                              // at :160
if (wc[i].wr_id & PROXY_WRID_INTERNAL) continue;              // before :106 and :185

Testing: no EP operator uses fetch atomics (dispatch/combine use ShmemAtomicTypeNonFetchThread / ShmemPutMemNbiSignalThread), so the benchmarks can't have covered this path.