[JACCL] Fix silent data corruption from unchecked RDMA work completion status

[0xDaizz]

Summary

The JACCL RDMA backend polls ibv_wc work completions to track in-flight RDMA operations but never checks wc[i].status. When an RDMA operation fails (e.g., due to memory pressure), the failure is silently ignored and the receive buffer—still containing stale or uninitialized data—is used as the result of the collective operation. This can cause silent data corruption in JACCL-based distributed workloads when RDMA operations fail. This PR adds wc[i].status validation to all 9 poll loops across ring.cpp and mesh.cpp, converting silent corruption into immediate, descriptive errors.

Problem

All 9 completion-polling loops in the JACCL backend (5 in ring.cpp, 4 in mesh.cpp) follow the same pattern:

ibv_wc wc[WC_NUM];
int n = poll(connections_, WC_NUM, wc);
for (int i = 0; i < n; i++) {
    // Only wr_id is examined to track in-flight count and buffer indices
    int work_type = wc[i].wr_id >> 16;
    int buff = (wc[i].wr_id >> 8) & 0xff;
    // ...
    in_flight--;
    // Proceed to use the receive buffer as if the operation succeeded
}

The wc[i].status field is never checked. Per the IBV specification, a polled completion with status != IBV_WC_SUCCESS indicates that the corresponding RDMA operation failed. The data in the associated receive buffer is undefined.

Additionally, the return value of ibv_poll_cq (wrapped by poll()) is not checked for negative values, which indicate a polling error.

Affected functions:

• RingGroup::all_gather (1 loop)
• RingGroup::send (1 loop)
• RingGroup::recv (1 loop)
• RingGroup::all_reduce_impl (2 loops: reduce-scatter phase + all-gather phase)
• MeshGroup::all_gather (1 loop)
• MeshGroup::send (1 loop)
• MeshGroup::recv (1 loop)
• MeshGroup::all_reduce (1 loop)

Impact

• May affect JACCL-based distributed workloads where RDMA operations can fail
• Particularly relevant with large models that put memory pressure on RDMA buffer allocation
• Observed in practice: a 612GB MoE model (306GB per rank across 2 nodes via Thunderbolt 5 RDMA) produced corrupted output after approximately 22 tokens of autoregressive generation, as RDMA buffer allocation began to fail under memory pressure
• Smaller models that fit comfortably in RDMA buffer limits are unaffected because their RDMA operations always succeed — the bug is latent but real

Fix

1. Centralized helpers in utils.h — instead of duplicating error-handling code in each file:

   • wc_status_name() — maps all ibv_wc_status enum values to human-readable strings (full coverage including UNKNOWN fallback)
   • check_wc_status(const ibv_wc&) — single-line status check that throws a descriptive std::runtime_error

2. poll() helpers in utils.h fixed to check ibv_poll_cq negative returns internally:

   • Connection::poll() — now throws on negative return
   • Free poll(vector, ...) — now throws on negative return from any CQ
   • This prevents negative returns from being masked when multiple CQ results are summed

3. Error messages include qp_num — valid on error per IBV spec, useful for identifying which connection failed in multi-peer topologies

4. Call sites simplified — each of the 9 poll loops in ring.cpp and mesh.cpp now uses just check_wc_status(wc[i]) instead of inline error handling

Example error message:

[jaccl] RDMA work completion error: status=5 (WR_FLUSH_ERR) qp_num=42 vendor_err=0 wr_id=0x10001

Testing

• Tested with a 612GB MoE model (Kimi k2.5) using tensor parallelism across 2 nodes connected via Thunderbolt 5 RDMA (M3 Studio 512GB x2, macOS 26.3)
• Verified with 256, 512, 5478+433, and 8490+512 token generation tests — all produce coherent output with zero RDMA errors
• Previously, the same configuration produced corrupted output after ~22 tokens
• Smaller models (< 10GB per rank) also tested to confirm no regression when RDMA operations succeed normally
• No performance impact: the status check is a single integer comparison on the hot path

[angeloskath]

That looks great thanks! I will run some benchmarks (it shouldn't affect them) and then will merge.

[angeloskath]

@0xDaizz I 'll check after the tests pass.

[0xDaizz]

@angeloskath Just fixed the lint error!

[0xDaizz]

@angeloskath Rebased onto the latest main (post #3174 refactor). The fix now targets the new mesh_impl.h / ring_impl.h structure instead of the old inline code in mesh.cpp / ring.cpp.

Changes (rebased)

• utils.h: add wc_status_name() + check_wc_status() helpers (+52 lines)
• mesh_impl.h: add status checks at 4 poll() sites
• ring_impl.h: add status checks at 5 poll() sites

Total: 3 files, +61 lines. Single commit on top of main.

Contributing checklist

Item	Status
Code formatting (clang-format)	✅ Verified — zero diff
Tests	✅ No new tests needed — this is a defensive error-path check (if (status != SUCCESS) throw). Existing distributed tests (jaccl_test_distributed.py) cover the normal path. Error-injection tests would require real RDMA hardware faults, which is not practical in CI.
Benchmark impact	✅ None — adds a single branch (if != SUCCESS) per poll completion on the hot path
API changes	✅ None — internal-only helpers

[angeloskath]

Ran the benchmarks and it does have a consistent impact for latency sensitive message sizes. ~10% on 16KB 4-way all reduce (or about 1μs extra). The slowdown obviously is not from the condition but could be a code gen issue or similar.

I am not saying we shouldn't merge this but given that we treat the TB as a reliable channel I am inclined to wait for a refactor/optimization pass on all of JACCL to remedy this. I will leave it open for now.

[0xDaizz]

Thanks for the detailed benchmarks. Noted on the ~1μs regression for latency-sensitive sizes. Happy to keep this open and fold it into a broader JACCL optimization pass. Let me know if I can help when the time comes.