[HIP] [Feat] Single-kernel Lamport fused all-reduce + RMSNorm (opt-in) - #4977
Draft
EricKing626 wants to merge 7 commits into
Draft
[HIP] [Feat] Single-kernel Lamport fused all-reduce + RMSNorm (opt-in)#4977EricKing626 wants to merge 7 commits into
EricKing626 wants to merge 7 commits into
Conversation
…STED)
Motivation
----------
On MI355X the fused AR+RMSNorm path runs as two kernels on one stream:
reduce_scatter_cross_device_store (start_sync ... end_sync)
local_device_load_rmsnorm
Because ROCm has no PDL (Programmatic Dependent Launch) equivalent, two
kernels on the same stream serialise strictly. B200 hides ~27% of this
pair by overlapping rmsNormLamport under twoshotAllreduceKernel via PDL;
we cannot. The only way to recover that time on ROCm is to make it one
kernel -- which requires removing the end_sync barrier between the two
phases, since a device-wide barrier inside a single kernel is what the
grid-level launch boundary used to provide.
Approach
--------
Lamport-style wait-free signalling: the data *is* the flag. The IPC
staging buffer is pre-filled with a sentinel (-0.0, i.e. 0x8000 per bf16
lane / 0x80000000 per fp32). Producers canonicalise -0.0 -> +0.0 before
storing, so the sentinel is unambiguous. A consumer polls its 16-byte
pack and treats it as arrived only when *no* word still equals the
sentinel. There is no separate flag, therefore no flag-vs-data
visibility ordering problem and no end_sync.
New in this commit (all in csrc/include/custom_all_reduce.cuh):
LamportSentinel<T> sentinel word per dtype
lamport_canonicalize<T,pack> producer-side -0.0 -> +0.0
lamport_spin_load<T,pack> volatile poll + s_sleep backoff
lamport_arm<T,pack> re-arm the pack after consuming
start_sync_acqrel<ngpus> RELEASE/ACQUIRE variant of start_sync
lamport_prefill_sentinel<T> first-use arming kernel
fused_allreduce_rmsnorm_lamport<..> the fused kernel
ar_rmsnorm_lamport_enabled() env gate
lamport_ensure_armed<T> host-side lazy arming (mutex + map)
Phase A of the fused kernel is reduce_scatter_cross_device_store's body
verbatim (same warp_id / lane_id / part indexing) plus canonicalisation
and minus end_sync. Phase B is local_device_load_rmsnorm's body with
lamport_spin_load replacing the plain load and lamport_arm immediately
after; each pack is consumed by exactly one thread, so re-arming needs
no extra synchronisation.
Correctness notes
-----------------
* start_sync_acqrel is required, not cosmetic: the stock start_sync uses
RELAXED atomics, which would let the *previous* iteration's re-arming
sentinel stores land after a peer's fresh producer stores and clobber
them.
* No deadlock within a rank: every block finishes all of its producer
work -- which depends on no one -- before it ever spins.
* The kernel is persistent-style, so the grid is clamped to the
co-resident block count via hipOccupancyMaxActiveBlocksPerMultiprocessor
for the exact instantiation. An over-large grid would leave producer
blocks unscheduled behind spinning consumers.
* Phase B adds a trailing __syncthreads() per token iteration that the
original did not have. This is required because the LDS is aliased
between phase A (T* tmp_smem) and phase B (float* smem); it also closes
a pre-existing latent race on smem[0] across bid iterations.
Status
------
DEFAULT OFF. Nothing changes unless AITER_AR_RMSNORM_LAMPORT=1.
*** THIS CODE IS UNCOMPILED AND UNTESTED. *** It was written without
access to a build environment or a GPU. Do not enable it outside of
bring-up.
Known gaps / follow-ups
-----------------------
* CUDA-graph capture: lamport_ensure_armed launches a kernel lazily on
first use and on any growth in m. If that first use happens during
graph capture the arming kernel gets captured (or rejected) rather than
executed once. Needs an explicit warmup/arming call before capture.
* tnum is fixed at 512; no 256-thread variant, so the 256-thread tuning
the two-kernel path can pick is unavailable here.
* Only the local_device_load_rmsnorm semantics are ported. The _naive,
_512n and the quantising variants still take the two-kernel path.
* No unit test. Suggested bring-up: TP=4/8, n=7168, m in {1,64,256},
compare against AITER_AR_RMSNORM_LAMPORT=0 elementwise, then loop it a
few thousand iterations to shake out the re-arming race.
Cheaper thing to try first
--------------------------
allreduce_fusion_kernel_1stage is already a single kernel with only
start_sync, and is disabled for our shapes purely by a byte threshold
(total_bytes <= 128*7168*2/world_size). A/B AITER_AR_1STAGE=1 vs 0 on
the target shape before investing in this path -- if 1-stage wins, most
of the gap is free and Lamport only matters above the crossover.
Contributor
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
PR title tags: |
github-actions
Bot
changed the title
Bring-up on MI355X (TP=4, n=7168, bf16) hit two deadlocks. Both are fixed
here; with them the kernel is numerically correct across
m in {1,8,16,32,64,128,256}, eager and CUDA graph, vs F.rms_norm at
atol=rtol=1e-2.
1. Grid exceeded the signal-slot limit.
Signal::start/end/_flag are sized [kMaxBlocks = 80] and start_sync_acqrel
indexes them by blockIdx.x. The grid was sized min(m, num_cu * 2), which is
512 on MI355X, so any m > 80 wrote past the signal region and the barrier
never completed. m=1 and m=64 passed; m=256 hung with three GPUs pinned at
100% and one idle. Clamped to kMaxBlocks. Both phases are grid-stride
loops, so the smaller grid only costs throughput.
2. Sentinel arming was captured into the graph.
lamport_ensure_armed launched the prefill on the caller's stream. Under
capture that becomes a graph node, so every replay re-armed the whole
staging buffer, overwriting producer stores the peers had already landed
and leaving the consumers spinning. Intermittent; a profiled replay loop
reproduced it 3/3 with the prefill captured and 0/3 when the buffer was
armed before capture. Now detects capture and runs the prefill on a private
non-capturing stream, keeping it out of the graph. Steady state stays armed
because phase B re-arms every pack it consumes. No sync is issued: it is
illegal under global capture mode, and capture is always followed by a
device sync before the first replay.
Performance note (CUDA graph, TP=4, n=7168, bf16, us/iter): Lamport loses at
every comparable shape - m=1 12.1 vs 6.7, m=8 12.4 vs 7.4, m=16 13.3 vs 9.3,
m=32 15.6 vs 13.2. Eager m=256 is 75.5 vs 48.9. Cause: phase B is one block
per token so the grid is min(m, 80) and phase A inherits it; at small m the
reduce-scatter runs on a single block while the two-kernel path launches the
producer with up to 80 blocks and sizes the consumer from occupancy. Removing
one launch and one device-wide barrier is worth a few us at most, less than
the lost parallelism. Decoupling the two phase grids is the prerequisite for
this path to be competitive. Kept opt-in and default-off.
Unrelated: the stock fused AR+RMSNorm path hangs at (64, 7168) TP=4 bf16
under CUDA graph. Reproduces with AITER_AR_RMSNORM_LAMPORT=0 and with the
pre-patch .so, so it predates this branch. To be filed separately.
Phase A's grid was tied to m, so a single-token decode ran the whole
reduce-scatter on one block. Decoupling it exposed a latent hazard and,
once measured, a different bottleneck than expected.
Block-uniform phase A trip count. The loop body contains __syncthreads()
but the bound was per-thread (`idx < part`). With the grid pinned to m
the waves happened to agree; with the grid sized independently they do
not, and the tail block's waves execute different numbers of barriers.
This was already latent before this commit and would have fired at m=2.
The trip count now comes from the block's base index and out-of-range
lanes are masked off. All ngpus warp groups share lane_id, so they agree
on both `idx` and `active`, which is what warp 0's cross-rank reduction
assumes.
Grid policy. A sweep over pinned grids (AITER_AR_RMSNORM_LAMPORT_BLOCKS,
added here as a tuning knob) shows the limiter is not phase A width but
start_sync_acqrel, which costs ~0.1us per block because every block
signals every peer. At (1, 7168) the cost climbs 11.7 -> 18.3us going
from 16 to 80 blocks with no work added. The optimum is one block per
token up to a crossover near 64, with a small floor so tiny m does not
serialise onto a single CU.
MI355X, TP=4, bf16, n=7168, CUDA graph, us/iter:
m grid old -> new old new delta
1 1 -> 4 12.10 11.26 -6.9%
8 8 -> 8 12.40 12.46 +0.5%
16 16 -> 16 13.27 13.24 -0.2%
32 32 -> 32 15.57 15.63 +0.4%
64 64 -> 64 - 23.32
128 80 -> 64 39.90 38.10 -5.6%
256 80 -> 64 68.50 66.62 -2.5%
Correctness passes for m in {1,8,16,32,64,128,256} in both eager and
graph mode. The fused path is still slower than the two-kernel baseline
(+16% at m=32, +68% at m=1); grid tuning alone does not close that, and
the remaining gap is the barrier and the sentinel re-arm write traffic.
Co-authored-by: Cursor <cursoragent@cursor.com>
Chasing the grid-scaling cost from the previous commit found the wrong
culprit. It is not start_sync_acqrel: replacing it with a variant where
only block 0 does the cross-rank handshake and the rest wait on a local
gate left the curve unchanged (at (1, 7168) still 11.2 -> 20.2us from 4
to 80 blocks) while adding ~1.5us of serialisation, so that variant is
not kept.
It is the __threadfence_system() that published phase A's stores. Every
block issues one and they all serialise on the same L2 writeback, at
~0.09us per block. Dropping it flattens (1, 7168) to 10.6 -> 11.6us
across the same grid range -- on a shape where 78 of those 80 blocks
publish nothing at all. Scope is not the issue: __threadfence() costs
the same, so it is the fence itself, not what it makes visible.
The fence was vestigial, inherited from the two-kernel producer where an
end_sync flag followed it. An LL/Lamport protocol needs no publish
fence: the data is its own flag, so there is no separate flag store to
order the payload before, and word order within a pack does not matter
either because the producer canonicalises the sentinel away and a
consumer seeing any word still holding it just keeps spinning. Same as
ll_store_b128 in custom_all_reduce_gfx1250.cuh and the QuickReduce send
path, neither of which fences. Peer stores are now non-temporal, the
quarter that lands in this rank's own buffer stays cached because phase
B reads it back (non-temporal there cost 4% at m=256).
With the grid nearly free (~0.01us per block, all barrier) the phase A
decoupling finally pays: size the grid for the reduce-scatter rather
than for m, floored so a small decode does not serialise onto a couple
of CUs.
MI355X, TP=4, bf16, n=7168, CUDA graph, us/iter:
m grid before after vs before baseline
1 16 11.26 10.31 -8.4% 6.70
8 16 12.46 11.40 -8.5% 9.39
16 28 13.24 12.22 -7.7% 9.37
32 56 15.63 14.50 -7.2% 13.46
64 80 23.32 19.40 -16.8% hangs
128 80 38.10 32.73 -14.1% hangs
256 80 66.62 60.39 -9.4% hangs
Correctness passes for m in {1,8,16,32,64,128,256} in eager and graph,
repeated. Still short of the two-kernel baseline where that baseline
runs (+7.7% at m=32, +54% at m=1); what is left is the sentinel re-arm
write traffic and the spin latency, not the grid.
Co-authored-by: Cursor <cursoragent@cursor.com>
Two findings from re-baselining, one of which invalidates the numbers in
the previous two commits.
The baseline was never the two-kernel path. dispatchFusedAllReduceRMSNorm
checks the Lamport path before MAYBE_DISPATCH_1S_KERNEL, and the plain
fused rmsnorm entry point selects one-stage for total_bytes <= 128*7168*2
/world_size, which at TP=4 and n=7168 is 448 KB -- every shape up to and
including m=32. So all the "baseline" figures reported so far were the
one-shot kernel, and enabling the Lamport path was silently displacing it.
A reduce-scatter cannot beat a one-shot kernel on small messages: the
one-shot reads each peer's input once and stages nothing. Gate on
!use_1stage so this path only ever displaces the two-kernel pair.
Measured at TP=4, bf16, n=7168, CUDA graph, us/iter:
m 1-stage Lamport delta
1 7.97 9.95 +25%
8 7.47 11.37 +52%
16 9.40 12.30 +31%
32 13.40 14.63 +9%
64 22.44 19.47 -13%
128 39.37 33.38 -15%
256 74.62 60.13 -19%
The crossover lands within a factor of two of the existing 448 KB
threshold, so the gate puts the Lamport path exactly where it wins.
The two-kernel path it replaces cannot be compared under a graph at all:
forcing AITER_AR_1STAGE=0 hangs at every shape, not just the m>=64 ones
where it is the default. That is pre-existing (it reproduces on the
pre-patch .so) and is why the earlier runs looked like "baseline hangs at
m=64" -- m<=32 was quietly taking the one-stage path. In eager, where it
does run, it is still ahead: 48.97us vs 60.13us at m=256.
Also adds AITER_AR_RMSNORM_LAMPORT_ARM_REPEAT, which repeats the
consumer's idempotent re-arm store so the extra passes time the re-arm
without invalidating the result. One pass costs ~12us of the 60us at
m=256 (~300 GB/s of uncached stores, so bandwidth-bound and a fair
estimate there); below m=64 the repeat serialises on the same address and
only bounds the cost from above.
Co-authored-by: Cursor <cursoragent@cursor.com>
!use_1stage alone leaves the gate at the mercy of whoever calls in, and the callers disagree. aiter's own communicator picks one-stage below 128*7168*2/world_size, which is 448 KB at TP=4 -- the same place the measured crossover sits, so the previous commit's gate is exactly right there. sglang does not: it uses a flat 128 KB (parallel_state.py, use_1stage_ar = total_bytes <= 128 * 1024), so with that caller every shape from m=10 up would take this path, including the m=16 and m=32 shapes where it is 31% and 9% slower than the one-shot kernel. Gate on the measured crossover directly instead. AITER_AR_RMSNORM_LAMPORT _MIN_BYTES overrides it for sweeps; set to 0 to disable the floor and recover the previous behaviour. The floor is TP-independent by construction: it is a per-rank input size, so the same byte count means the same amount of work per GPU whatever the world size. Whether the crossover itself moves with world size has not been measured -- only TP=4 was swept -- so 448 KB should be re-checked at TP=8 before this is relied on there.
The Lamport path shared the general tmp region (get_tmp_buf) with the 2-stage all-reduce, reduce-scatter and allgather kernels. That is a silent-corruption bug, not a performance detail. Lamport encodes arrival in the data, so the region must hold sentinels at the start of every call, and the only thing that restores them is the Lamport consumer re-arming each pack it reads. Any other kernel staging through the same region overwrites the sentinels with real data. lamport_ensure_armed will not notice: it tracks a high-water mark and still believes the buffer is primed. The next Lamport call reads non-sentinel words, treats stale data as arrived, and returns the previous iteration's values. No hang, wrong results. The microbenchmark cannot see this because m is fixed for the whole run. End-to-end it is unavoidable: decode's m moves with batch size and MTP, so shapes below the Lamport floor take the 2-stage path and poison the buffer for the shapes above it. Re-arming defensively does not fix it under CUDA graphs. The pollution comes from replaying a graph captured at a different shape, and the re-arm cannot be recorded into the Lamport graph -- replaying it would clobber producer stores the peers have already landed, which is the hang fixed in 14a918f. A private region is the only thing that holds across replays. So carve kLamportScratchBytes (32 MiB) out of the meta buffer right after Signal and move get_tmp_buf past it, mirroring how the gfx1250 path already appends its LL staging scratch. Callers size the meta allocation as meta_size() + 2 * max_size, so reporting the scratch from meta_size() is all that is needed to reserve it -- no change on the Python side. 32 MiB covers m=2048 at n=7168 bf16. Larger shapes fail lam_ok and fall back to the two-kernel path.
EricKing626
marked this pull request as draft
August 25, 2026 10:47
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
custom_fused_ar_rmsruns as two kernels — two-shot all-reduce, then RMSNorm — separated by anend_syncgrid barrier. On ROCm there is no Programmatic Dependent Launch, so the two strictlyserialise.
On GLM-5.2 FP4 decode (isl 8192 / osl 1024 / conc 64, TP4) this pair is the largest single
contributor to our remaining MI355X latency gap. Removing the serialisation means collapsing the
two kernels into one, which means removing the barrier between them.
Approach
Lamport wait-free signalling: the data is the flag.
-0.0).-0.0 -> +0.0, so a written value can never be the sentinel.end_syncis no longer needed, so reduce and norm become one kernel — one launch, no gap.Two supporting details worth reviewing:
start_sync_acqrel— RELEASE/ACQUIRE variant of the existing RELAXEDstart_sync, so there-arming sentinel stores of iteration N cannot be reordered past a peer's producer stores of
iteration N+1.
other blocks of the same grid. The grid is clamped to the co-resident block count via
hipOccupancyMaxActiveBlocksPerMultiprocessor. Without the clamp, unscheduled producer blockscan never run behind spinning consumers and the kernel deadlocks.
Changes
Confined to
csrc/include/custom_all_reduce.cuh. Adds a Lamport sentinel/arm/spin helper set,start_sync_acqrel, and the fused kernel plus its dispatch. Nothing existing is modified.Enabling
Off by default. Opt in with
AITER_AR_RMSNORM_LAMPORT=1.The fast path also requires:
n % pack_size == 0,packs_per_row >= 64,1 <= n_loop <= 4,(m * packs_per_row) % world_size == 0, andworld_sizein {2, 4, 8}. Any failure falls throughto the existing 1-stage / 2-stage paths unchanged.
Testing status
Not yet compiled or run on hardware. Opened early for design review of the Lamport protocol
and the occupancy clamp. Correctness and performance numbers will be posted before merge.
Planned: TP=4/8,
n=7168,min {1, 64, 256}, elementwise vsAITER_AR_RMSNORM_LAMPORT=0;several thousand back-to-back iterations for the re-arming race; end-to-end GLM-5.2 decode A/B.
Known gaps
lamport_ensure_armedcall; needs a warm-up before capture.tnumfixed at 512, not swept.local_device_load_rmsnormsemantics ported;_naive,_512nand quantised variants keepthe two-kernel path.