RS + GEMM overlap for BF16, gfx950, w/ HipKittens + UserBuffers by alextmagro · Pull Request #725 · ROCm/TransformerEngine

alextmagro · 2026-08-31T23:10:08Z

New fused reduce-scatter + GEMM forward kernel
New bulk reduce-scatter backward overlap kernel
HK overlap kernel tests moved to test_rocm_fused_overlap.py, replacing old userbuffer tests in CI
Extracted common AG and RS code to a shared header

For MI350:
fused RS + TN GEMM is 1.20x faster geomean, 1.25x faster ms weighted
bulk RS + NT GEMM is 1.18x faster geomean, 1.06x faster ms weighted. (NT gemm is bottleneck, expect a followup PR improving this performance shortly)

github-actions · 2026-09-01T04:01:44Z

Claude review

Reviewed the full diff against the merge base (23 files, ~2.2k insertions): the new fused RS+GEMM TN kernel, the bulk RS NT kernel, the shared overlap_common.cuh extraction, the fused_ag_gemm -> comm_gemm rename, the Python eligibility/setup plumbing, and the test reorganisation.

The structure is good — the AG/RS symmetry in comm_gemm.cpp reads well, the extraction of the shared geometry into overlap_common.cuh is a real simplification, and the graceful !eligible fallback to rocm_split_overlap_rs is the right shape. Three things should be resolved before merge:

The new test file cannot run. tex is used but never imported, so every test in test_rocm_fused_overlap.py raises NameError on gfx950. It stays invisible because the tests skip on non-gfx950 CI.
A tp=4 assertion contradicts the implementation. Python setup, rs_guards_ok, and the kernel's template instantiation all accept 4 ranks, so the "must decline the fused RS" branch should fail.
A question about cross-rank ordering in the fused RS kernel. I can't find what orders a comm workgroup's read of a peer's stage against that peer's GEMM store — the RS_SENT sentinel looks half-wired (nothing writes the poison, no epilogue clamp), and done[] / peers.arrive / peers.ready / peers.recv are allocated and published but never read. Happy to be shown what I'm missing here.

Also flagged, lower severity: an RS path bumping _ag_signal_base while its own fallback uses _rs_signal_base, two pieces of dead state in comm_gemm_overlap.h, a determinism test that never compares its two runs, an unreachable skip in test_comm_gemm_overlap.py, and the CI swap dropping the remaining Userbuffers coverage.

Copyright headers: OK — 21 of 22 in-scope files correct, no NVIDIA years altered; one question on transformer.py where the AMD line was removed (likely correct, see inline).

aris134 · 2026-09-01T21:29:28Z

+    const int bands   = (M / tp_size) / 256;
+
+    // Order matters here
+    return (tp_size == 4 || tp_size == 8) && tp_size <= args.peer_count &&


Just flagging this here because it seems implied that tp_size should be 8

TransformerEngine/tests/pytorch/distributed/test_rocm_fused_overlap.py

Lines 35 to 37 in 392fd78

The reduce-scatter backend is stricter still -- its combine is specialised on 8 ranks -- but

that is asserted per test rather than hidden in collection, so a tp=4 decline stays visible.

"""

aris134 · 2026-09-02T13:04:54Z

+};
+
+using namespace kittens;
+using G_group = kittens::group<NUM_WARPS>;


this line appears to be duplicated across fused_ag_gemm_nn.cuh, fused_ag_gemm_tn.cuh, and fused_rs_gemm_tn.cuh. Since NUM_WARPS is already specified in overlap_common.cuh, perhaps factor this line out into the common header as well

aris134 · 2026-09-02T13:12:11Z

+#define RS_PUBLISH(p) __hip_atomic_fetch_add((p), 1u, __ATOMIC_RELEASE, __HIP_MEMORY_SCOPE_SYSTEM)
+#define RS_SPIN(p)    __hip_atomic_load((p), __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_SYSTEM)
+#define RS_ACQUIRE(p) ((void)__hip_atomic_load((p), __ATOMIC_ACQUIRE, __HIP_MEMORY_SCOPE_SYSTEM))


I don't think these macros are used anywhere

aris134 · 2026-09-02T13:52:18Z

+            if (xcd_bucket) {
+                // Steal order: own bucket, then the local chunk, then the other XCDs.
+                int found = -1;
+                const int b0 = (int)blockIdx.x % NUM_XCDS_AFF;
+                for (int s = 0; s <= NUM_XCDS_AFF; s++) {
+                    int bb;
+                    if (s == 0) {
+                        bb = b0;
+                    } else if (s == 1) {
+                        bb = my_pe;
+                    } else {
+                        bb = (b0 + s - 1) & (NUM_XCDS_AFF - 1);
+                    }
+                    if (buckets.cnt[bb] == 0) continue;
+                    if (__hip_atomic_load(&bucket_ctr[bb], __ATOMIC_RELAXED,
+                                          __HIP_MEMORY_SCOPE_AGENT) >= buckets.cnt[bb]) continue;
+                    const int idx = atomicAdd(&bucket_ctr[bb], 1);
+                    if (idx < buckets.cnt[bb]) {
+                        found = buckets.off[bb] + idx;
+                        break;
+                    }
+                }
+                s_tile_idx = (found < 0) ? num_tiles : found;


The work stealing code block appears identical across all headers, and could likely be factored out as a device helper in overlap_common.cuh

aris134 · 2026-09-02T14:49:10Z

+    typedef int v4i __attribute__((ext_vector_type(4)));
+
+    {
+        const int b0 = 0;


nit: b0 is a frozen 0. Do we need it in this block?

aris134 · 2026-09-02T15:12:35Z

+                          size_t band_elems) {
+    const int w = (int)blockIdx.x;
+    const size_t lines = band_elems / 8;
+    const int peer = (w < TP - 1) ? (w + (w >= my_pe ? 1 : 0)) : -1;


defined here but not used in this kernel

wangye805 · 2026-09-02T15:12:10Z

+    if (hidden == 16384 && k_local == 6656 &&
+        (tokens == 16384 || tokens == 32768 || tokens == 65536)) {
+        return 4;
+    }
+    if (hidden == 8192 && k_local == 1024 && (tokens == 32768 || tokens == 65536)) {
+        return 12;
+    }
+    if (hidden == 4096 && k_local == 512) {
+        if (tokens == 65536) return 16;
+        if (tokens == 32768) return 12;
+    }


nit: this hardcoding seemed to be applied onto specific customer requested configs. Okay for now. Is there a systematic way to tune comm_wg? Like with gemm bounded cases, decrease the comm_wg and for rs bounded cases, increase the comm_wg?

wangye805 · 2026-09-02T15:20:07Z

+#define RS_SENT_BF16 0xFFAAu
+#define RS_SENT_DW ((unsigned int)RS_SENT_BF16 * 0x00010001u)


Do you need a hipmemset write this 0xffaa into output buffer before each gemm call? I didn't see the writing function call

Also, this poison value usage is based on the assumption that user will not use produce -nan. What if user use this 0xffaa for the padding tokens?

wangye805 · 2026-09-02T15:36:02Z

+}
+
+typedef int rs_v4i __attribute__((ext_vector_type(4)));
+typedef const volatile __attribute__((address_space(1))) rs_v4i *rs_gvol4;


This is for poison value 0xffaa reading, right? It was written by other ranks and we are trying to read it in local rank. Without atomic writing, can you check if this read lower to cache-bypassing load (sc0/sc1/nt) on gfx950 in asm?

If not, consumer may keep reading the stale l2 cache data...

wangye805 · 2026-09-02T15:38:42Z

+#pragma unroll 1
+        for (int tile = 0; tile < k_tiles - 2; tile += 2) {
+
+            auto st_b = subtile_inplace<REG_N, K_STEP>(Bs[0][0], {warp_n, 0});
+            load(b_tile_0, st_b);
+            auto st_a = subtile_inplace<REG_M, K_STEP>(As[0][0], {warp_m, 0});
+            load(a_tile, st_a);
+            G_group::load(As[1][1], A, {0, 0, block_row * 2 + 1, tile + 1}, sw_A, a_srd, a_base, a_lds_11);
+            asm volatile("s_waitcnt lgkmcnt(8)");
+            __builtin_amdgcn_s_barrier();
+
+            asm volatile("s_waitcnt lgkmcnt(0)");
+            __builtin_amdgcn_s_setprio(1);
+            mma_ABt(c00, a_tile, b_tile_0, c00);
+            __builtin_amdgcn_s_setprio(0);
+            __builtin_amdgcn_s_barrier();
+            __builtin_amdgcn_sched_barrier(0);
+
+            st_b = subtile_inplace<REG_N, K_STEP>(Bs[0][1], {warp_n, 0});
+            load(b_tile_1, st_b);
+            G_group::load(Bs[0][0], B, {0, 0, block_col * 2, tile + 2}, sw_B, b_srd, b_base, b_lds_00);
+            __builtin_amdgcn_s_barrier();
+
+            asm volatile("s_waitcnt lgkmcnt(0)");
+            __builtin_amdgcn_s_setprio(1);
+            mma_ABt(c01, a_tile, b_tile_1, c01);
+            __builtin_amdgcn_s_setprio(0);
+            __builtin_amdgcn_s_barrier();
+
+            st_a = subtile_inplace<REG_M, K_STEP>(As[0][1], {warp_m, 0});
+            load(a_tile, st_a);
+            G_group::load(As[0][0], A, {0, 0, block_row * 2, tile + 2}, sw_A, a_srd, a_base, a_lds_00);
+            __builtin_amdgcn_s_barrier();
+
+            asm volatile("s_waitcnt lgkmcnt(0)");
+            __builtin_amdgcn_s_setprio(1);
+            mma_ABt(c10, a_tile, b_tile_0, c10);
+            __builtin_amdgcn_s_setprio(0);
+            __builtin_amdgcn_s_barrier();
+            __builtin_amdgcn_sched_barrier(0);
+
+            st_b = subtile_inplace<REG_N, K_STEP>(Bs[1][0], {warp_n, 0});
+            load(b_tile_0, st_b);
+            G_group::load(Bs[0][1], B, {0, 0, block_col * 2 + 1, tile + 2}, sw_B, b_srd, b_base, b_lds_01);
+            asm volatile("s_waitcnt vmcnt(6)");
+            __builtin_amdgcn_s_barrier();
+
+            __builtin_amdgcn_s_setprio(1);
+            mma_ABt(c11, a_tile, b_tile_1, c11);
+            __builtin_amdgcn_s_setprio(0);
+            __builtin_amdgcn_s_barrier();
+
+            st_a = subtile_inplace<REG_M, K_STEP>(As[1][0], {warp_m, 0});
+            load(a_tile, st_a);
+            G_group::load(As[0][1], A, {0, 0, block_row * 2 + 1, tile + 2}, sw_A, a_srd, a_base, a_lds_01);
+            asm volatile("s_waitcnt lgkmcnt(8)");
+            __builtin_amdgcn_s_barrier();
+
+            asm volatile("s_waitcnt lgkmcnt(0)");
+            __builtin_amdgcn_s_setprio(1);
+            mma_ABt(c00, a_tile, b_tile_0, c00);
+            __builtin_amdgcn_s_setprio(0);
+            __builtin_amdgcn_s_barrier();
+            __builtin_amdgcn_sched_barrier(0);
+
+            st_b = subtile_inplace<REG_N, K_STEP>(Bs[1][1], {warp_n, 0});
+            load(b_tile_1, st_b);
+            G_group::load(Bs[1][0], B, {0, 0, block_col * 2, tile + 3}, sw_B, b_srd, b_base, b_lds_10);
+            __builtin_amdgcn_s_barrier();
+
+            asm volatile("s_waitcnt lgkmcnt(0)");
+            __builtin_amdgcn_s_setprio(1);
+            mma_ABt(c01, a_tile, b_tile_1, c01);
+            __builtin_amdgcn_s_setprio(0);
+            __builtin_amdgcn_s_barrier();
+
+            st_a = subtile_inplace<REG_M, K_STEP>(As[1][1], {warp_m, 0});
+            load(a_tile, st_a);
+            G_group::load(As[1][0], A, {0, 0, block_row * 2, tile + 3}, sw_A, a_srd, a_base, a_lds_10);
+            __builtin_amdgcn_s_barrier();
+
+            asm volatile("s_waitcnt lgkmcnt(0)");
+            __builtin_amdgcn_s_setprio(1);
+            mma_ABt(c10, a_tile, b_tile_0, c10);
+            __builtin_amdgcn_s_setprio(0);
+            __builtin_amdgcn_s_barrier();
+            __builtin_amdgcn_sched_barrier(0);
+
+            G_group::load(Bs[1][1], B, {0, 0, block_col * 2 + 1, tile + 3}, sw_B, b_srd, b_base, b_lds_11);
+            asm volatile("s_waitcnt vmcnt(4)");
+            __builtin_amdgcn_s_barrier();
+
+            __builtin_amdgcn_s_setprio(1);
+            mma_ABt(c11, a_tile, b_tile_1, c11);
+            __builtin_amdgcn_s_setprio(0);
+            __builtin_amdgcn_s_barrier();
+        }


nit: is this tn gemm tile same as that in ag+tn gemm? If so, can we consolidate them with a common macro?

alextmagro requested review from aris134 and matthiasdiener August 31, 2026 23:10

alextmagro requested review from ipanfilo, wangye805 and wenchenvincent as code owners August 31, 2026 23:10

Reduce scatter overlap kernels for gfx950 + bf16

392fd78

alextmagro force-pushed the hipkittens_overlap_rs branch from 62f89ba to 392fd78 Compare September 1, 2026 03:49

alextmagro changed the title ~~Hipkittens overlap rs~~ RS + GEMM overlap for BF16, gfx950, w/ HipKittens + UserBuffers Sep 1, 2026

alextmagro added the claude-review label Sep 1, 2026