fix(custom_all_reduce): use SYSTEM scope + ACQUIRE ordering for cross-device signal loads

[hekhong-png]

Fix cross-device memory ordering in custom_all_reduce on multi-GPU HIP

Summary

The custom_all_reduce kernel's synchronization primitives use __MEMORY_SCOPE_DEVICE for loads of signals that are written by remote GPUs via P2P. On HIP, DEVICE scope only guarantees visibility within the same device — it does not order loads against cross-device stores. This causes the receiving CU to spin on stale data for an extended tail before the store becomes visible, and in pathological cases can lead to a kernel hang on TP8 configurations.

This patch aligns the load scopes with the CUDA reference implementation (which already uses SYSTEM scope) and adds a __threadfence_system() after __builtin_nontemporal_store to guarantee the data payload is visible to remote CUs before the completion signal is sent.

Root cause

start_sync and end_sync each contain a spin-wait loop that polls a flag written by peer GPUs:

// start_sync — original
while(__scoped_atomic_load_n(&self_sg->start[blockIdx.x][threadIdx.x],
                             __ATOMIC_RELAXED,
                             __MEMORY_SCOPE_DEVICE) < flag)   // ← bug
    ;

The flag is written by a remote GPU using __MEMORY_SCOPE_SYSTEM (the store side is already correct), but the load side uses __MEMORY_SCOPE_DEVICE. On HIP:

• DEVICE scope only orders memory operations within the same device. It does not guarantee that a store from another GPU (arriving via P2P / XGMI) is visible to the load.
• SYSTEM scope orders operations across all devices in the system, matching the CUDA reference.

The __ATOMIC_RELAXED ordering on start_sync's load is also too weak — it provides no acquire semantics, so the invalidate queue is not flushed and cross-CU stores may take a long time to become visible. end_sync already used __ATOMIC_ACQUIRE (conditionally), so only the scope was wrong there.

Additionally, the cross_device_reduce_2stage_write_mode path uses __builtin_nontemporal_store to write the reduction result directly to memory (bypassing L2 cache). Without a fence, the data may not have propagated through the memory controller when end_sync signals completion, causing the receiver to spin waiting for data that hasn't become visible yet.

Changes

Three minimal fixes to csrc/include/custom_all_reduce.cuh (4 insertions, 3 deletions):

1. start_sync load: RELAXED → ACQUIRE + DEVICE → SYSTEM

 while(__scoped_atomic_load_n(&self_sg->start[blockIdx.x][threadIdx.x],
-                             __ATOMIC_RELAXED,
-                             __MEMORY_SCOPE_DEVICE) < flag)
+                             __ATOMIC_ACQUIRE,
+                             __MEMORY_SCOPE_SYSTEM) < flag)

ACQUIRE ordering flushes the invalidate queue so cross-CU stores become visible promptly. SYSTEM scope matches the store side and the CUDA reference.

2. end_sync load: DEVICE → SYSTEM

 while(__scoped_atomic_load_n(&self_sg->end[blockIdx.x][threadIdx.x],
                              final_sync ? __ATOMIC_RELAXED : __ATOMIC_ACQUIRE,
-                             __MEMORY_SCOPE_DEVICE) < flag)
+                             __MEMORY_SCOPE_SYSTEM) < flag)

end_sync already had ACQUIRE (conditional on final_sync); only the scope was wrong.

3. __threadfence_system() after __builtin_nontemporal_store

             __builtin_nontemporal_store(*(src_addr + 3), dst_addr + 3);
+            __threadfence_system();
         }

Placed inside the if(is_broadcast_reg_outptr) block so it only executes on the nontemporal_store path (the else branch uses shared memory and doesn't need it). The fence forces all prior stores — including nontemporal stores that bypass L2 — to be visible to all devices before end_sync signals completion.

Why these are safe and correct

• Matches the CUDA reference. The CUDA path (the #else branch guarded by __CUDA_ARCH__) already uses __threadfence_system() and system-scope atomics. The HIP path was inconsistent: stores used SYSTEM scope but loads used DEVICE scope. This patch makes both sides consistent.
• No behavioral change on the happy path. The fixes only affect when the load observes the store — they make the store visible sooner (by flushing the invalidate queue and using the correct scope), not later. The fence adds a small constant cost but eliminates the pathological spin-wait tail.
• No new state, no new allocations, no new threads. The patch is 4 lines of changes to existing synchronization primitives. No instrumentation, no safety nets, no background threads.

Validation

Tested on a TP8 configuration (8× AMD MI308X, gfx942, ROCm 7.2.0) running GLM-5.2-FP8 with EAGLE speculative decoding (speculative-num-steps 5, speculative-eagle-topk 1, speculative-num-draft-tokens 6).

Instrumented profiling (subsequently removed) measured the spin-wait tail length across ~196 million all-reduce invocations:

Version	start_sync max spin	≥16384-spin events	Safety net triggers
v3 (original)	144,279	7,890	n/a
v4 (scope fix only)	136,285	—	n/a
v5 (+ acquire)	85,381	0	0
v6 (+ fence)	76,809	0	0
v6 extended (~55 min)	60,352 (decreasing)	0	0

Key findings:

• The scope + acquire fix (v5) eliminated all ≥16384-spin events (7,890 → 0) and reduced max spin by 41%.
• The fence fix (v6) reduced max spin a further 10% and maintained 0 high-spin events.
• Over a 55-minute extended run, max spin decreased over time (76,809 → 60,352), ruling out divergence or accumulation.
• The safety net (10M spin limit) never triggered across 196M calls — the fixes eliminated the pathological tail entirely.

Production regression test (v7, final patch)

The final patch (3 fixes only, no instrumentation, no safety net) was deployed and load-tested for ~17 minutes with 6 concurrent workers (~680 requests, ~125K generated tokens):

Metric	v6 baseline	v7 final	Verdict
gen_throughput (tok/s)	218.9	203–275 (avg ~228)	No regression
ITL p50 (ms)	25	25	Identical
ITL p90 (ms)	—	35	—
ITL p99 (ms)	40	40	Identical
ITL avg (ms)	25	25.2	Stable
Errors / hangs	0	0	Clean
Correctness (temp=0 "Paris")	PASS	PASS	Identical

Throughput fluctuated 203–275 tok/s across samples (batch composition variance), averaging ~228 tok/s — within noise of the v6 baseline (218.9). Latency percentiles were bit-identical (p50=25ms, p99=40ms). No errors, no hangs, no crashes across the run.

Scope of impact

• Only affects the HIP path (#ifndef __CUDA_ARCH__).
• Only affects configurations using custom_all_reduce (i.e., disable_custom_all_reduce=false, the default on HIP).
• The use_write_mode path (where the fence is added) only triggers when world_size_ == 8 && bytes > 4,194,304 && arch.find("gfx942") != std::string::npos.

[github-actions]

🏷️ CI Guide

Runs automatically on every PR:

• ✅ Pre-checks (submodule verification, code formatting)
• ✅ Aiter op tests (gfx942 + gfx950)
• ✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
ci:gfx1250-ffm-triton	Run the five-shard gfx1250 FFM Triton test suite
ci:triton-300x	Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
ci:sglang	SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
ci:atom	ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
ci:atom_full	ATOM accuracy suite for PR and main models from ATOM models_accuracy.json
ci:vllm	vLLM benchmark: GPT-OSS-120B, DeepSeek-R1-0528, Kimi-K2.5
ci:all	All standard extended tests (excludes ci:atom_full)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 4786 --add-label <label>