[Perf] [MI355X] glm53_flash: Cheaper hand-offs: uncached scratch, fewer vector loads ahead of the polls, tagged selected tokens (preview)

sgl_amd_mono/glm53_flash/gfx950/kernel.py

@@ -528,6 +528,12 @@ def ld_bf16(r, i):
                 fx.BFloat16(bo.buffer_load(r, i, vec_width=1, dtype=T.bf16))
             )
 
+        def ld_bf16s(r, elem, n):
+            """bf16 elements elem .. elem + n - 1 (n of 4 or 8, elem a multiple of n)
+            of a plain tensor -> n f32, in one load."""
+            w = fx.Vector(bo.buffer_load(r, elem // 2, vec_width=n // 2, dtype=T.i32))
+            return [v for e in range(n // 2) for v in bf2_f32(w[e])]
+
         def ld_pair(r, elem):
             """bf16 elements elem, elem + 1 (elem even) of a plain tensor -> 2 f32."""
             return bf2_f32(
@@ -1373,11 +1379,15 @@ def seam_post(item, col0, pre=None):
         # misc words of the mix scales (3) and biases (N_MIX), past the mix sums
         SCW = 32
         HBW = SCW + 4
+        # a row's mix partials (every tile's), staged in LDS words 0 .. PM_WORDS
+        PM_WORDS = N_ROW_TILES * PMIX
+        PM_LOADS = PM_WORDS // 2
+        assert PM_WORDS % 2 == 0 and PM_WORDS + HBW + N_MIX <= X_WORDS
 
         def gate_early():
             """The attention gate's constants, loaded at launch start: this thread's norm
             weights, a mix scale and a mix bias (clamped thread indices)."""
-            gam = [ld_bf16(_rsrc(g_in), tid * CPT + q) for q in range(CPT)]
+            gam = ld_bf16s(_rsrc(g_in), tid * CPT, CPT)
             sc = ld_f32(_rsrc(hc_a_scale), fx.min(tid, 2))
             bs = ld_f32(_rsrc(hc_a_base), fx.min(tid, N_MIX - 1))
             return gam, sc, bs
@@ -1399,7 +1409,7 @@ def gate_row(stage, dep, s, attn, sums, early=None):
                 if tid < N_MIX:
                     lds_st(sums, HBW + tid, bs)
             else:
-                gam = [ld_bf16(r_g, tid * CPT + q) for q in range(CPT)]
+                gam = ld_bf16s(r_g, tid * CPT, CPT)
                 if tid < 3:
                     lds_st(sums, SCW + tid, ld_f32(r_scale, tid))
                 if tid < N_MIX:
@@ -1417,18 +1427,22 @@ def gate_row(stage, dep, s, attn, sums, early=None):
                     for q in range(CPT // 4)
                 ]
             )
+            # every tile's mix partials of the row -> LDS words 0 .., two pairs a load
+            pm = [
+                fx.min(tid + i * THREADS, PM_LOADS - 1)
+                for i in range(-(-PM_LOADS // THREADS))
+            ]
+            got_pm = poll([(pmix, s * N_ROW_TILES * PMIX + 2 * q, 2) for q in pm])
+            for i, q in enumerate(pm):
+                for e in range_constexpr(2):
+                    lds_st(xs, 2 * q + e, got_pm[i][e].bitcast(fx.Float32))
+            gpu.barrier()
             if tid < PMIX * NP:
                 k = tid // NP
                 part = tid % NP
-                vals = getf_many(
-                    [
-                        (pmix, (s * N_ROW_TILES + part * TPP + q) * PMIX + k)
-                        for q in range(TPP)
-                    ]
-                )
-                tot = vals[0]
+                tot = lds_ld(xs, part * TPP * PMIX + k)
                 for q in range_constexpr(1, TPP):
-                    tot = tot + vals[q]
+                    tot = tot + lds_ld(xs, (part * TPP + q) * PMIX + k)
                 lds_st(red, k * NP + part, tot)
             gpu.barrier()
             if tid < PMIX:
@@ -2173,7 +2187,7 @@ def dsa_qnorm(tv, t32):
                 s = fx.Int32(t)
                 live = tid < NQ4
                 e0 = fx.min(tid, NQ4 - 1) * 4
-                gam = [ld_bf16(r_gqa, e0 + j) for j in range(4)]
+                gam = ld_bf16s(r_gqa, e0, 4)
                 v2 = get2_many(
                     [(mb("qa"), s * Q_LORA + e0), (mb("qa"), s * Q_LORA + e0 + 2)]
                 )
@@ -2672,7 +2686,6 @@ def score_half(s, n_pools, loaded):
         def dsa_iselect(tv, t32):
             hist = fx.recast_iter(fx.Int32, xs)
             r_scores = _rsrc(mb("iscores"))
-            r_sel = _rsrc(mb("tok_sel"))
 
             def radix_key(bits):
                 return (bits >= 0).select(bits ^ fx.Int32(-(2**31)), ~bits)
@@ -2856,23 +2869,21 @@ def gather(
                                 )
 
                     def publish_all():
-                        """red[0 .. 512) holds the pools; write the row's tokens."""
+                        """red[0 .. 512) holds the pools; write the row's tokens, tagged."""
                         gpu.barrier()
                         for b in range_constexpr(TOPK_TOKENS // THREADS):
                             j = tid + b * THREADS
                             pool = lds_ld(red, j // KPOOL).bitcast(fx.Int32)
-                            bo.buffer_store(
-                                pool * KPOOL + j % KPOOL,
-                                r_sel,
+                            put(
+                                mb("tok_sel"),
                                 s * SEL_STRIDE + j,
-                                cache_modifier=CM_DEV,
+                                pool * KPOOL + j % KPOOL,
                             )
                         if tid < (pos + 1) - bound * KPOOL:
-                            bo.buffer_store(
-                                bound * KPOOL + tid,
-                                r_sel,
+                            put(
+                                mb("tok_sel"),
                                 s * SEL_STRIDE + TOPK_TOKENS + tid,
-                                cache_modifier=CM_DEV,
+                                bound * KPOOL + tid,
                             )
 
                     def candidates(level, prefix, remain, n_cand):
@@ -3099,12 +3110,6 @@ def hist2_chunk(items, item_keys):
                             candidates(8, p24, r24, c24)
                         if c24 > SEL_CAP:
                             full_radix()
-                    # the tokens are device-scope stores: their completion releases
-                    # them
-                    rocdl.s_waitcnt(vmcnt=0)
-                    gpu.barrier()
-                    if tid == 0:
-                        put(mb("tok_ready"), s, fx.Int32(1))
 
         # ================= sparse MLA split: 64 selected keys x HT heads a task
         # Scored rows attend to the tokens the index selected plus their tail, other
@@ -3158,30 +3163,23 @@ def kv_row(blk, off):
                         )
 
                     kidx = t * SPLIT_KEYS + lane
-                    # A task that reads the selection waits for it: one lane polls the
-                    # selection's and the cache's flags; its divergent poll reconverges
-                    # before wave 0's acquire, and the barrier below orders the other
-                    # waves' loads after it. Any other task knows its tokens and
-                    # gathers them at once; of those only the row's own token (its
-                    # last key) is written by this launch's cache task, and it is read
-                    # again once kv_ready is in.
+                    # A task that reads the selection waits for it: wave 0 polls its 64
+                    # tagged tokens and the cache's flag together, then acquires; the
+                    # barrier below orders the other waves' loads after it. Any other task
+                    # knows its tokens and gathers them at once; of those only the row's
+                    # own token (its last key) is written by this launch's cache task, and
+                    # it is read again once kv_ready is in.
                     if picked_keys:
                         if wave == 0:
-                            if lane == 0:
-                                poll([(mb("tok_ready"), s, 1), (mb("kv_ready"), s, 1)])
+                            picked = poll(
+                                [
+                                    (mb("tok_sel"), s * SEL_STRIDE + kidx, 1),
+                                    (mb("kv_ready"), s, 1),
+                                ]
+                            )[0][0]
                             fx.memory_fence(
                                 ordering=fx.AtomicOrdering.Acquire, syncscope="agent"
                             )
-                        if wave == 0:
-                            picked = fx.Int32(
-                                bo.buffer_load(
-                                    _rsrc(mb("tok_sel")),
-                                    s * SEL_STRIDE + kidx,
-                                    vec_width=1,
-                                    dtype=T.i32,
-                                    cache_modifier=CM_DEV,
-                                )
-                            )
                             resolve((kidx < nk).select(picked, fx.Int32(0)))
                     else:
                         if wave == 0:
@@ -3597,7 +3595,7 @@ def tile_task(k, tile, pre):
                         tile_task(k + 1, tile + G, nxt)
 
             pre = ahead(t0)
-            gam = [ld_bf16(r_g, tid * CPT + q) for q in range(CPT)]
+            gam = ld_bf16s(r_g, tid * CPT, CPT)
             # the rows' rstd: wave s sums row s's tile sums (lanes l, l + 64), then the
             # wave in a fixed order
             if wave < S:
@@ -3823,7 +3821,7 @@ def u_r(c):
             # the gate's
             if const_expr(ROUTER_OWN_ROW):
                 li, rstd, gam, _ = gate_row(
-                    "router", "seam", router_sample, False, xs + HIDDEN // 2
+                    "router", "seam", router_sample, False, xs + PM_WORDS
                 )
                 for q in range_constexpr(CPT // 2):
                     lds_st(

sgl_amd_mono/glm53_flash/gfx950/layout.py

@@ -234,8 +234,7 @@ def layout(
             ("q_lat", samples * H * KV_LORA // 2 * pair_bytes),
             ("kv_ready", samples * pair_bytes),
             ("iscores", samples * max(max_pools, 1) * pair_bytes),
-            ("tok_sel", samples * SEL_STRIDE * 4),
-            ("tok_ready", samples * pair_bytes),
+            ("tok_sel", samples * SEL_STRIDE * pair_bytes),
             ("sp_acc", samples * N_SPLIT * H * KV_LORA // 2 * pair_bytes),
             ("sp_m", samples * N_SPLIT * H * pair_bytes),
             ("sp_l", samples * N_SPLIT * H * pair_bytes),

sgl_amd_mono/glm53_flash/gfx950/runner.py

@@ -202,7 +202,10 @@ def __init__(self, key, rank: int, group, device: torch.device):
         self.launch, self.sc, self.sy = build_flash_kernel(key)
         self.peer = PeerBuffer(self.sy["_bytes"], group, rank, key.tp, device)
         self.peer.bytes.zero_()
-        self.scratch = torch.zeros(self.sc["_bytes"], dtype=torch.uint8, device=device)
+        # the hand-off mailboxes, uncached as the peer buffer is: a line is read at its
+        # home memory, not through the producer's L2
+        self._scratch = PeerBuffer(self.sc["_bytes"], group, rank, 1, device)
+        self.scratch = self._scratch.bytes
         self.epoch = torch.zeros(guard.GUARD_WORDS, dtype=torch.int32, device=device)
         # the tail's gathered logits; a layer's two output sets, by layer parity: a
         # layer never writes what it reads
@@ -244,6 +247,8 @@ def advance_step(self) -> None:
         self.epoch[:1].add_(1)
 
     def close(self) -> None:
+        self.scratch = None
+        self._scratch.close()
         self.peer.close()
 
 

test/kernels/glm53_flash/check_layer.py

@@ -709,10 +709,13 @@ def run(kv):
         out = op.launch(x, res, post, comb, caches=c, rows=table, page_table=pt)
         rt.advance_step()
         torch.cuda.synchronize()
-        sel = rt.scratch[rt.sc["tok_sel"] :].view(torch.int32)[: rows * SEL_STRIDE]
-        return out[1].clone(), sel.view(rows, SEL_STRIDE).clone()
+        sel = rt.scratch[rt.sc["tok_sel"] :].view(torch.int32)[: 2 * rows * SEL_STRIDE]
+        return out[1].clone(), sel.view(rows, SEL_STRIDE, 2).clone()
 
     base, sel = run(caches.kv_cache.clone())
+    # (token, launch tag) pairs; kv_ready holds this launch's tag
+    tags = rt.scratch[rt.sc["kv_ready"] :].view(torch.int32)[1 : 2 * rows : 2].clone()
+    sel, sel_tags = sel[..., 0], sel[..., 1]
     g = _gen(30, device)
     ok = True
     for r, n in enumerate(KEY_LENS):
@@ -731,12 +734,13 @@ def run(kv):
             tail_ok = sel[r, TOPK_TOKENS : TOPK_TOKENS + tail].tolist() == [
                 bound * KPOOL + i for i in range(tail)
             ]
+            tagged = bool((sel_tags[r, : TOPK_TOKENS + tail] == tags[r]).all())
             print(
                 f"rank {rank} row {r} ({n} tokens): 512 whole distinct pools "
-                f"{whole and distinct}, tail {tail_ok}",
+                f"{whole and distinct}, tail {tail_ok}, tagged {tagged}",
                 flush=True,
             )
-            ok = ok and whole and distinct and tail_ok
+            ok = ok and whole and distinct and tail_ok and tagged
             victims = picked[TOPK_TOKENS - tail : TOPK_TOKENS].tolist()
             unpicked = sorted(set(range(bound)) - set(ids.tolist()))
             spare = unpicked[0] * KPOOL