diff --git a/sgl_amd_mono/glm53_flash/gfx950/kernel.py b/sgl_amd_mono/glm53_flash/gfx950/kernel.py index b337e78e..00554f5a 100644 --- a/sgl_amd_mono/glm53_flash/gfx950/kernel.py +++ b/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( diff --git a/sgl_amd_mono/glm53_flash/gfx950/layout.py b/sgl_amd_mono/glm53_flash/gfx950/layout.py index 4af304fd..921a56c5 100644 --- a/sgl_amd_mono/glm53_flash/gfx950/layout.py +++ b/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), diff --git a/sgl_amd_mono/glm53_flash/gfx950/runner.py b/sgl_amd_mono/glm53_flash/gfx950/runner.py index 53688e0d..1ef4c25e 100644 --- a/sgl_amd_mono/glm53_flash/gfx950/runner.py +++ b/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() diff --git a/test/kernels/glm53_flash/check_layer.py b/test/kernels/glm53_flash/check_layer.py index 49d11792..027145ae 100644 --- a/test/kernels/glm53_flash/check_layer.py +++ b/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