[jaccl] Fix race on local_staging in MeshImpl::all_reduce

[kernelpool]

This fixes a bug introduced in #3412. I observed this when running distributed tensor-parallel inference (MiniMax-M2.7-8bit on 2x M3 Ultra over jaccl using mlx 0.31.2/mlx-lm 0.31.3/macOS 26.4) where generation degenerates into sentence loops (or token-level garbage) for prompts above ~170 tokens.

The SEND-completion handler refilled local_staging(buff) the moment all peers ACK'd the previous send, regardless of whether that chunk had been consumed by the own-rank reduction step. RDMA timing made this non-deterministic, producing wrong sums for messages spanning multiple PIPELINE chunks.

Decouple the two: SEND completion only refills send_buffer (free once on the wire) and posts the next send. local_staging(b) is refilled in the reduce loop right after the own-rank reduction reads it, which also bumps recv_end[rank_] to gate the next step.

Checklist

Put an x in the boxes that apply.

• I have read the CONTRIBUTING document
• I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
• I have added tests that prove my fix is effective or that my feature works
• I have updated the necessary documentation (if needed)

[kernelpool]

I'm not sure this completely addresses the issue. I seem to still get repetitive behavior with Kimi-K2.6 over longer context (something i didn't observe when testing this model over several days previously using mlx 0.31.1 / mlx-lm 0.31.2). I'll do some more digging and see if I can more reliably reproduce it.

[angeloskath]

Thanks this is a great catch!

Not sure it was the bug you encountered but it is a race condition indeed. It never showed up in my tests but it certainly can show up especially in non-homogenous cluster setups.