diff --git a/sgl_amd_mono/glm53/gfx950/kernel.py b/sgl_amd_mono/glm53/gfx950/kernel.py index f09d3660..67fc917e 100644 --- a/sgl_amd_mono/glm53/gfx950/kernel.py +++ b/sgl_amd_mono/glm53/gfx950/kernel.py @@ -118,8 +118,7 @@ INDEX_HEADS, INDEX_KEYS_PER_TASK, INDEX_Q_ROWS, - INDEX_SEL_CHUNK, - INDEX_SEL_PAR_CHUNKS, + INDEX_SEL_DIGITS, INDEX_SEL_SLICES, INDEX_TILE, N_QKV_A, @@ -216,7 +215,8 @@ def build_glm5_monokernel( ``index_block_bytes`` apart. ``req_ids`` picks each row's ``block_table`` row (rows ``bt_stride`` entries apart, a launch argument; the rows of one verify request may share it), keys cover positions ``0 .. positions[s]``, a key that - a row of this launch writes is scored from that row's ``index_k_new`` pair, and + a row of this launch writes is scored from that row's ``index_k8`` pairs (its + cached FP8 values and scale), and the top-k is published as global cache slots to ``out_indices`` at the row's ``sparse_kv_indptr`` offset. ``index_k_bf16`` reads the index K projection as packed bf16. @@ -261,6 +261,9 @@ def build_glm5_monokernel( assert not (attention_bf16 and with_indexer), ( "BF16 attention uses the external indexer" ) + assert not with_indexer or index_paged, ( + "the fused indexer reads SGLang's paged keys" + ) assert not index_paged or ( with_indexer and use_atom_kv_cache @@ -352,14 +355,9 @@ def build_glm5_monokernel( SEL_BINS = 4096 SEL_CAP = (X_WORDS - SEL_BINS) // 2 assert not index_paged or SEL_CAP >= 2048 - # rows of more keys select on up to SEL_SLICES CTAs a row, in the grid's upper half + # a row selects on up to SEL_SLICES CTAs, whole rows a round of the grid SEL_SLICES = INDEX_SEL_SLICES - SEL_PAR_KEYS = INDEX_SEL_PAR_CHUNKS * SEL_CHUNK - assert not index_paged or ( - SEL_CHUNK == INDEX_SEL_CHUNK - and SEL_CAP <= INDEX_CAND_STAGE_CAP - and S * SEL_SLICES <= G // 2 - ) + assert not index_paged or (SEL_CAP <= INDEX_CAND_STAGE_CAP and G % SEL_SLICES == 0) MISC_OFF = X_WORDS MISC_WORDS = max( 8 + S * (XQ_GROUPS if native_fp4_mfma else XQ_BLOCKS), @@ -2226,6 +2224,172 @@ def ld_qa_one(sks): put_bf(mb("q_an"), s_norm * Q_LORA + k, [a0, a1, a2, a3]) stamp("q_norm", s_norm, 4) + if const_expr(with_indexer): + # The 128 CTAs without q_b work produce the other 128 index-query tiles + # concurrently, the q_norm CTAs theirs right after q_norm (ahead of the + # row's index key and KV rows). They reload q_a, but remove one full GEMV + # from the q_b CTAs' serialized critical path. + N_INDEX_Q_EXTRA = INDEX_Q_ROWS // INDEX_TILE - N_QB + IQ_NKC = Q_LORA // 64 + IQ_UNITS = IQ_NKC // (2 * WAVES) + r_wiq, r_siq = _rsrc(index_arg(3)), _rsrc(index_arg(4)) + + def index_q_extra(tt): + iq_t = N_QB + tt + stamp("index_q", iq_t, 0) + + def u_index_q_extra(c): + kc = (wave * IQ_UNITS + c) * 2 + return unit_fp8x2( + r_wiq, + r_siq, + iq_t, + kc, + IQ_NKC, + Q_LORA, + (n_sel() * Q_LORA + kc * 64) // 2, + shuffled=fp8_shuffled, + ) + + pre = [u_index_q_extra(c) for c in range(IQ_UNITS)] + hint_wait( + S, + lambda s: (mb("q_an"), (s * Q_LORA + Q_LORA - 2) // 2), + mark=("index_q", iq_t), + ) + stage_x_pairs("q_an", S * Q_LORA, lambda k: k) + stamp("index_q", iq_t, 2) + gpu.barrier() + iq_acc = run_units(u_index_q_extra, IQ_UNITS, IQ_UNITS, pre) + reduce_rows(1, iq_acc, emit_out(INDEX_TILE)) + stamp("index_q", iq_t, 3) + gpu.barrier() + if tid < S * INDEX_TILE // 4: + s = tid // (INDEX_TILE // 4) + r = (tid % (INDEX_TILE // 4)) * 4 + put_bf( + mb("index_q"), + s * INDEX_Q_ROWS + iq_t * INDEX_TILE + r, + [lds_ld(outs, s * INDEX_TILE + r + j) for j in range(4)], + ) + stamp("index_q", iq_t, 4) + + for tt in range(start("index_q"), N_INDEX_Q_EXTRA, G): + index_q_extra(fx.Int32(tt)) + + # ========= each row's index query, rotated and Hadamard-transformed once: + # 8 heads a task, a wave a head, lane l holding values 2l and 2l + 1 + def index_qh_tasks(): + for t in range(start("index_qh"), S * INDEX_HEADS // 8, G): + t = fx.Int32(t) + s = t // (INDEX_HEADS // 8) + h = (t % (INDEX_HEADS // 8)) * 8 + fx.Int32(wave) + stamp("index_qh", t, 0) + kq = fx.Int32(lane) * 2 + pos_q = fx.max(index_bound(s) - 1, fx.Int32(0)) + at = pos_q * (PE_DIM // 2) + fx.min(kq, PE_DIM - 2) // 2 + c = ld_bf16(_rsrc(rope_cos), at) + sn = ld_bf16(_rsrc(rope_sin), at) + q_pair = ( + s * (INDEX_Q_ROWS // 2) + h * (INDEX_DIM // 2) + fx.Int32(lane) + ) + q0, q1 = bf2_f32(poll([(mb("index_q"), q_pair, 1)])[0][0]) + stamp("index_qh", t, 2) + if kq < PE_DIM: + q0, q1 = q0 * c - q1 * sn, q0 * sn + q1 * c + q0, q1 = fwht128(q0, q1) + put(mb("index_qh"), q_pair, bf16_pair(q0, q1)) + stamp("index_qh", t, 4) + + # At TP4 q_b runs on every CTA and its tasks, later in a CTA's order, + # produce every index-query tile, so these tasks run after q_b there + assert N_INDEX_Q_EXTRA >= 0 + if const_expr(N_INDEX_Q_EXTRA > 0): + index_qh_tasks() + + if const_expr(with_indexer): + # ========= 3a. each row's index key: LayerNorm, RoPE, Hadamard -> index cache + # one task a row, on the row's q_norm CTA after its index-query tile + for s in range(start("index_key"), S, G): + s = fx.Int32(s) + stamp("index_key", s, 0) + # the LayerNorm weight and bias and the RoPE factors ahead of the wait + r_gik, r_bik = _rsrc(index_arg(5)), _rsrc(index_arg(6)) + r_index_cache = _rsrc(indices) + i0 = fx.min(tid, INDEX_DIM // 2 - 1) * 2 + ik_g = [ld_f32(r_gik, i0 + j) for j in range(2)] + ik_b = [ld_f32(r_bik, i0 + j) for j in range(2)] + tpe = tid % (PE_DIM // 2) + c = ld_bf16(_rsrc(rope_cos), row_position(s) * (PE_DIM // 2) + tpe) + sn = ld_bf16(_rsrc(rope_sin), row_position(s) * (PE_DIM // 2) + tpe) + slot = row_slot(s) + active = row_writes_cache(s) + stamp("index_key", s, 1) + # Index keys use LayerNorm (not RMSNorm), interleaved RoPE on the + # first 64 dimensions, then the 128-point Hadamard that SGLang's + # gfx950 indexer applies to the keys it caches and to its queries. + # the row's key in one poll + live = tid < INDEX_DIM + ik_one = fx.min(tid, INDEX_DIM - 1) + ik_two = fx.min(tid, INDEX_DIM // 2 - 1) * 2 + got = poll( + [ + (mb("index_k"), s * INDEX_DIM + ik_one, 1), + (mb("index_k"), s * INDEX_DIM + ik_two, 2), + ] + ) + stamp("index_key", s, 2) + ik = got[0][0].bitcast(fx.Float32) + k0, k1 = got[1][0].bitcast(fx.Float32), got[1][1].bitcast(fx.Float32) + mean = block_sums([live.select(ik, fx.Float32(0.0))])[0] * ( + 1.0 / INDEX_DIM + ) + centered = live.select(ik - mean, fx.Float32(0.0)) + rstd = _rsq( + block_sums([centered * centered])[0] * (1.0 / INDEX_DIM) + 1.0e-6 + ) + if tid < INDEX_DIM // 2: + v0 = (k0 - mean) * rstd * ik_g[0] + ik_b[0] + v1 = (k1 - mean) * rstd * ik_g[1] + ik_b[1] + if tid < PE_DIM // 2: + v0, v1 = v0 * c - v1 * sn, v0 * sn + v1 * c + v0, v1 = fwht128(v0, v1) + # the index cache's ue8m0 quant: scale = 2^ceil(log2(max(amax, 1e-4) / 448)) + amax = wave_max(fx.max(fmath.absf(v0), fmath.absf(v1))) + bits = div_rn( + fx.max(amax, fx.Float32(1.0e-4)), + fx.Float32(FP8_MAX), + fx.Float32(1.0 / FP8_MAX), + ).bitcast(fx.Int32) + exponent = (bits >> 23) & fx.Int32(0xFF) + exponent = exponent + ((bits & fx.Int32(0x7FFFFF)) != 0).select( + fx.Int32(1), fx.Int32(0) + ) + k_scale = (exponent << 23).bitcast(fx.Float32) + k_inv = ((fx.Int32(254) - exponent) << 23).bitcast(fx.Float32) + q0, q1 = v0 * k_inv, v1 * k_inv + word = fp8_pack4(q0, q1, _xshfl(q0, 1), _xshfl(q1, 1)) + block = slot // index_block_size + off = slot % index_block_size + if active & (tid % 2 == 0): + bo.buffer_store( + word, r_index_cache, index_value_word(block, off, i0) + ) + if tid % 2 == 0: + put(mb("index_k8"), s * (INDEX_DIM // 4 + 1) + tid // 2, word) + if tid == 0: + put( + mb("index_k8"), + s * (INDEX_DIM // 4 + 1) + INDEX_DIM // 4, + k_scale, + ) + if active & (tid == 0): + bo.buffer_store( + k_scale, r_index_cache, index_scale_word(block, off) + ) + gpu.barrier() + stamp("index_key", s, 4) + # ================ 3. KV RMSNorm + k_pe RoPE -> cache (+ this launch's rows) # one task a row for s in range(start("cache"), S, G): @@ -2323,85 +2487,6 @@ def ld_qa_one(sks): bo.buffer_store(p0.to(fx.BFloat16), r_pe, pe_offset) bo.buffer_store(p1.to(fx.BFloat16), r_pe, pe_offset + 1) put2(mb("penew"), s * PE_DIM + tid * 2, p0, p1) - - if const_expr(with_indexer): - # Index keys use LayerNorm (not RMSNorm), interleaved RoPE on the - # first 64 dimensions, then the 128-point Hadamard that SGLang's - # gfx950 indexer applies to the keys it caches and to its queries. - r_gik, r_bik = _rsrc(index_arg(5)), _rsrc(index_arg(6)) - r_index_cache = _rsrc(indices) - # the row's key in one poll - live = tid < INDEX_DIM - ik_one = fx.min(tid, INDEX_DIM - 1) - ik_two = fx.min(tid, INDEX_DIM // 2 - 1) * 2 - got = poll( - [ - (mb("index_k"), s * INDEX_DIM + ik_one, 1), - (mb("index_k"), s * INDEX_DIM + ik_two, 2), - ] - ) - ik = got[0][0].bitcast(fx.Float32) - k0, k1 = got[1][0].bitcast(fx.Float32), got[1][1].bitcast(fx.Float32) - mean = block_sums([live.select(ik, fx.Float32(0.0))])[0] * ( - 1.0 / INDEX_DIM - ) - centered = live.select(ik - mean, fx.Float32(0.0)) - rstd = _rsq( - block_sums([centered * centered])[0] * (1.0 / INDEX_DIM) + 1.0e-6 - ) - if tid < INDEX_DIM // 2: - i0 = tid * 2 - v0 = (k0 - mean) * rstd * ld_f32(r_gik, i0) + ld_f32(r_bik, i0) - v1 = (k1 - mean) * rstd * ld_f32(r_gik, i0 + 1) + ld_f32( - r_bik, i0 + 1 - ) - if tid < PE_DIM // 2: - v0, v1 = v0 * c - v1 * sn, v0 * sn + v1 * c - v0, v1 = fwht128(v0, v1) - if const_expr(index_paged): - # the index cache's ue8m0 quant: scale = 2^ceil(log2(max(amax, 1e-4) / 448)) - amax = wave_max(fx.max(fmath.absf(v0), fmath.absf(v1))) - bits = div_rn( - fx.max(amax, fx.Float32(1.0e-4)), - fx.Float32(FP8_MAX), - fx.Float32(1.0 / FP8_MAX), - ).bitcast(fx.Int32) - exponent = (bits >> 23) & fx.Int32(0xFF) - exponent = exponent + ((bits & fx.Int32(0x7FFFFF)) != 0).select( - fx.Int32(1), fx.Int32(0) - ) - k_scale = (exponent << 23).bitcast(fx.Float32) - k_inv = ((fx.Int32(254) - exponent) << 23).bitcast(fx.Float32) - q0, q1 = v0 * k_inv, v1 * k_inv - word = fp8_pack4(q0, q1, _xshfl(q0, 1), _xshfl(q1, 1)) - block = slot // index_block_size - off = slot % index_block_size - if active & (tid % 2 == 0): - bo.buffer_store( - word, r_index_cache, index_value_word(block, off, i0) - ) - if active & (tid == 0): - bo.buffer_store( - k_scale, r_index_cache, index_scale_word(block, off) - ) - d0, d1 = _fp8_roundtrip(q0, q1) - put_bf( - mb("index_k_new"), - s * INDEX_DIM + i0, - [d0 * k_scale, d1 * k_scale], - ) - else: - bo.buffer_store( - fx.Vector.from_elements([v0, v1], fx.Float32).to( - fx.BFloat16 - ), - r_index_cache, - (pos0 + s) * INDEX_DIM + i0, - ) - put_bf(mb("index_k_new"), s * INDEX_DIM + i0, [v0, v1]) - gpu.barrier() - if tid == 0: - put(mb("index_ready"), s, fx.Int32(1)) stamp("cache", s, 4) # =============================================== 4. normalized q_a -> q_b (+RoPE) @@ -2430,6 +2515,25 @@ def u_qb(c): ) pre = [u_qb(c) for c in range(QB_UNITS)] + if const_expr(with_indexer): + # this CTA's index-query tile reuses the normalized q_lora tile: its + # weights are issued before the wait too + iq_t = t + + def u_index_q(c): + kc = (wave * QB_UNITS + c) * 2 + return unit_fp8x2( + r_wiq, + r_siq, + iq_t, + kc, + QB_NKC, + Q_LORA, + (n_sel() * Q_LORA + kc * 64) // 2, + shuffled=fp8_shuffled, + ) + + iq_pre = [u_index_q(c) for c in range(QB_UNITS)] hint_wait( S, lambda s: (mb("q_an"), (s * Q_LORA + Q_LORA - 2) // 2), @@ -2441,6 +2545,24 @@ def u_qb(c): stage_x_pairs("q_an", S * Q_LORA, lambda k: k) stamp("q_b", t, 2) gpu.barrier() + if const_expr(with_indexer): + # This CTA's index-query row tile first, on the normalized q_lora tile + # (the index chain waits on it, split on the selection after it). + # Complementary CTAs compute the remaining half below. + stamp("index_q", iq_t, 0) + iq_acc = run_units(u_index_q, QB_UNITS, QB_UNITS, iq_pre) + reduce_rows(1, iq_acc, emit_out(INDEX_TILE)) + stamp("index_q", iq_t, 3) + gpu.barrier() + if tid < S * INDEX_TILE // 4: + s = tid // (INDEX_TILE // 4) + r = (tid % (INDEX_TILE // 4)) * 4 + put_bf( + mb("index_q"), + s * INDEX_Q_ROWS + iq_t * INDEX_TILE + r, + [lds_ld(outs, s * INDEX_TILE + r + j) for j in range(4)], + ) + stamp("index_q", iq_t, 4) acc = run_units(u_qb, QB_UNITS, QB_UNITS, pre) if const_expr(attention_ptpc): reduce_rows( @@ -2483,98 +2605,8 @@ def u_qb(c): ) stamp("q_b", t, 4) - if const_expr(with_indexer): - # Reuse this CTA's normalized q_lora tile for one index-query - # row tile. Complementary CTAs compute the remaining half below. - iq_t = t - stamp("index_q", iq_t, 0) - - def u_index_q(c): - kc = (wave * QB_UNITS + c) * 2 - return unit_fp8x2( - r_wiq, - r_siq, - iq_t, - kc, - QB_NKC, - Q_LORA, - (n_sel() * Q_LORA + kc * 64) // 2, - shuffled=fp8_shuffled, - ) - - iq_acc = run_units(u_index_q, QB_UNITS, QB_UNITS) - reduce_rows(1, iq_acc, emit_out(INDEX_TILE)) - stamp("index_q", iq_t, 3) - gpu.barrier() - if tid < S * INDEX_TILE // 4: - s = tid // (INDEX_TILE // 4) - r = (tid % (INDEX_TILE // 4)) * 4 - put_bf( - mb("index_q"), - s * INDEX_Q_ROWS + iq_t * INDEX_TILE + r, - [lds_ld(outs, s * INDEX_TILE + r + j) for j in range(4)], - ) - stamp("index_q", iq_t, 4) - - if const_expr(with_indexer): - # The 128 CTAs without q_b work produce the other 128 index-query - # tiles concurrently. They reload q_a, but remove one full GEMV - # from the q_b CTAs' serialized critical path. S of them hold the q_norm - # and cache tasks: their tiles run second on the CTAs of the first S extra - # tiles instead (qkv_a tiles' CTAs), as behind the cache tasks they would - # end last. - N_INDEX_Q_EXTRA = INDEX_Q_ROWS // INDEX_TILE - N_QB - IQ_NORM = (base["q_norm"] - base["index_q"]) % G - assert base["cache"] == base["q_norm"] - # at TP4 the q_b CTAs run every index-query tile: there is no extra one - assert N_INDEX_Q_EXTRA == 0 or S <= IQ_NORM <= N_INDEX_Q_EXTRA - S - - def index_q_extra(tt): - iq_t = N_QB + tt - stamp("index_q", iq_t, 0) - - def u_index_q_extra(c): - kc = (wave * QB_UNITS + c) * 2 - return unit_fp8x2( - r_wiq, - r_siq, - iq_t, - kc, - QB_NKC, - Q_LORA, - (n_sel() * Q_LORA + kc * 64) // 2, - shuffled=fp8_shuffled, - ) - - pre = [u_index_q_extra(c) for c in range(QB_UNITS)] - hint_wait( - S, - lambda s: (mb("q_an"), (s * Q_LORA + Q_LORA - 2) // 2), - mark=("index_q", iq_t), - ) - stage_x_pairs("q_an", S * Q_LORA, lambda k: k) - stamp("index_q", iq_t, 2) - gpu.barrier() - iq_acc = run_units(u_index_q_extra, QB_UNITS, QB_UNITS, pre) - reduce_rows(1, iq_acc, emit_out(INDEX_TILE)) - stamp("index_q", iq_t, 3) - gpu.barrier() - if tid < S * INDEX_TILE // 4: - s = tid // (INDEX_TILE // 4) - r = (tid % (INDEX_TILE // 4)) * 4 - put_bf( - mb("index_q"), - s * INDEX_Q_ROWS + iq_t * INDEX_TILE + r, - [lds_ld(outs, s * INDEX_TILE + r + j) for j in range(4)], - ) - stamp("index_q", iq_t, 4) - - for tt in range(start("index_q"), N_INDEX_Q_EXTRA, G): - tt = fx.Int32(tt) - if (tt < IQ_NORM) | (tt >= IQ_NORM + S): - index_q_extra(tt) - if tt < S: - index_q_extra(tt + IQ_NORM) + if const_expr(with_indexer and N_INDEX_Q_EXTRA == 0): + index_qh_tasks() # ==================================== 4. absorbed query: q_lat = W_UK^T q_nope # 8 row groups (128 latent rows of one head) x 3 chunks: one row group per wave @@ -2666,10 +2698,11 @@ def u_uk(c): ) stamp("uk", t, 4) - # ====================== 4b. fused sparse index score + exact top-2048 - if const_expr(with_indexer and not index_paged): + # ====================== 4b. index scores and the exact top-2048 (paged keys) + if const_expr(index_paged): r_index_cache = _rsrc(indices) - N_INDEX_SPLIT = index_max_seq // INDEX_KEYS_PER_TASK + r_out_indices = _rsrc(out_indices) + hist = fx.recast_iter(fx.Int32, xs) def load_index_q8(head, k): words = fx.Vector.from_elements( @@ -2678,366 +2711,115 @@ def load_index_q8(head, k): ) return words.bitcast(fx.BFloat16) - for tt in range(start("index_score"), S * N_INDEX_SPLIT, G): - tt = fx.Int32(tt) - s = tt // N_INDEX_SPLIT - tile = tt % N_INDEX_SPLIT - stamp("index_score", tt, 0) - bound = pos0 + s + 1 - if tile * INDEX_KEYS_PER_TASK < bound: - get(mb("index_ready"), s) - if tid < INDEX_HEADS: - lds_st(keys, tid, get(mb("index_w"), s * INDEX_HEADS + tid)) - # Every key group reuses the same 32x128 query. Stage and RoPE - # it once per scoring CTA instead of polling and rotating it in - # each of the four key-group wave pairs. - for b in range_constexpr((INDEX_Q_ROWS // 2) // THREADS): - q_pair = tid + b * THREADS - q_elem = q_pair * 2 - kq = q_elem % INDEX_DIM - q0, q1 = get_bf2_many( - [(mb("index_q"), s * INDEX_Q_ROWS + q_elem)] - )[0] - if kq < PE_DIM: - c = ld_bf16( - _rsrc(rope_cos), (pos0 + s) * (PE_DIM // 2) + kq // 2 - ) - sn = ld_bf16( - _rsrc(rope_sin), (pos0 + s) * (PE_DIM // 2) + kq // 2 - ) - q0, q1 = q0 * c - q1 * sn, q0 * sn + q1 * c - q0, q1 = fwht128(q0, q1) - lds_st(xs, q_pair, bf16_pair(q0, q1)) - gpu.barrier() - key_group = wave // 2 - head_group = wave % 2 - key_pos = tile * INDEX_KEYS_PER_TASK + key_group * 16 + lane % 16 - safe_key = fx.min(key_pos, bound - 1) - head = head_group * 16 + lane % 16 - score_frag = fx.Vector.filled(4, 0.0, fx.Float32) - for k32 in range_constexpr(INDEX_DIM // 32): - k = k32 * 32 + (lane // 16) * 8 - qv = load_index_q8(head, k) - kv = fx.Vector( + def radix_key(bits): + return (bits >= 0).select(bits ^ fx.Int32(-(2**31)), ~bits) + + # Scoring. A request's consecutive rows (up to SCORE_ROWS) score its keys + # together, so each key is read once: the request's pages are split into + # runs of whole pages, one a CTA (counted from a CTA its first row offsets), + # and a CTA copies its run into LDS a batch at a time, the next batch's + # loads in flight while it scores one. Each row's ROW_WAVES waves score the + # batch, the row's rotated query in registers. A key this launch writes is + # patched into the batch from its row's index_k8 pairs (its FP8 values and + # scale as cached; the cache store need not be visible yet). + SCORE_ROWS = 1 if S == 1 else 2 if S == 2 else 4 + ROW_WAVES = WAVES // SCORE_ROWS + PAGE_WORDS = index_block_bytes // 4 + BATCH_PAGES = 8 + BATCH_WORDS = BATCH_PAGES * PAGE_WORDS + BATCH_LOADS = (BATCH_WORDS + 4 * THREADS - 1) // (4 * THREADS) + # a CTA scores at least SCORE_PAGES pages of a request, or all of them + SCORE_PAGES = 2 + MAX_RUN_PAGES = max( + SCORE_PAGES, (index_max_seq // index_block_size + G - 1) // G + ) + BT_OFF = BATCH_WORDS # the run's block ids, after the batch + K8 = INDEX_DIM // 4 + 1 # a new key's pairs: its value words, its scale + assert INDEX_Q_ROWS == THREADS * 8 + assert index_block_size == INDEX_KEYS_PER_TASK + assert BT_OFF + MAX_RUN_PAGES <= X_WORDS + assert SCORE_ROWS * INDEX_Q_ROWS // 2 <= BATCH_WORDS + assert SCORE_ROWS * INDEX_HEADS <= LDS_KEYS and S * K8 <= THREADS + kw = fx.recast_iter(fx.Int32, xs) + # a request's runs from this CTA on (past the q_norm and uk CTAs' work) + SCORE_BASE = 101 + + def score_groups(): + """Bit s of the first value: row s leads a group of rows scored together, + up to SCORE_ROWS consecutive rows of one request; 3 bits a row of the + second: the rows of the group it leads, and of the third: its request's + first row in the launch. A row is of the previous row's request when its + page table maps that row's position to the slot that row writes (a page + being written is its request's alone).""" + bounds = [index_bound(r) for r in range(S)] + same = [fx.Int32(0)] + for r in range_constexpr(1, S): + prev = fx.max(bounds[r - 1] - 1, fx.Int32(0)) + at = index_block(index_request(r), prev) * index_block_size + ( + prev % index_block_size + ) + same.append( + ((bounds[r - 1] > 0) & (at == row_slot(r - 1))).select( + fx.Int32(1), fx.Int32(0) + ) + ) + leads, sizes, firsts = fx.Int32(0), fx.Int64(0), fx.Int64(0) + lead, run, first = fx.Int32(0), fx.Int32(0), fx.Int32(0) + for r in range_constexpr(S): + live = bounds[r] > 0 + first = (live & (same[r] != 0)).select(first, fx.Int32(r)) + joins = live & (same[r] != 0) & (run > 0) & (run < SCORE_ROWS) + starts = live & (joins.select(fx.Int32(1), fx.Int32(0)) == 0) + leads = leads | starts.select(fx.Int32(1 << r), fx.Int32(0)) + firsts = firsts | starts.select( + fx.Int64(first) << (3 * r), fx.Int64(0) + ) + lead = starts.select(fx.Int32(r), lead) + run = live.select(joins.select(run + 1, fx.Int32(1)), fx.Int32(0)) + shift = fx.Int64(lead) * 3 + sizes = live.select( + (sizes & ~(fx.Int64(7) << shift)) | (fx.Int64(run) << shift), + sizes, + ) + return _uniform(leads), sizes, firsts + + def fetch_batch(page0, n_pages): + """This thread's words of the run's pages page0 .. page0 + n_pages - 1 + (at most a batch; block ids in LDS), as 16-byte loads.""" + out = [] + for i in range_constexpr(BATCH_LOADS): + w = (tid + i * THREADS) * 4 + p = fx.Int32(w // PAGE_WORDS) + block = lds_ld( + kw, BT_OFF + page0 + fx.max(fx.min(p, n_pages - 1), 0) + ) + v = fx.Vector.filled(4, 0, fx.Int32) + if (p < n_pages) & (fx.Int32(w) < BATCH_WORDS): + v = fx.Vector( bo.buffer_load( r_index_cache, - (safe_key * INDEX_DIM + k) // 2, + block * PAGE_WORDS + w % PAGE_WORDS, vec_width=4, dtype=T.i32, ) - ).bitcast(fx.BFloat16) - if safe_key >= pos0: - sn = safe_key - pos0 - kv_pairs = get_bf2_many( - [ - (mb("index_k_new"), sn * INDEX_DIM + k + j * 2) - for j in range(4) - ] - ) - kv_values = [] - for pair in kv_pairs: - kv_values += list(pair) - kv = fx.Vector.from_elements(kv_values, fx.Float32).to( - fx.BFloat16 - ) - score_frag = fx.Vector( - rocdl.mfma_f32_16x16x32_bf16( - T.vec(4, T.f32), [qv, kv, score_frag] - ) ) - partial = fx.Float32(0.0) - for e in range_constexpr(4): - index_head = head_group * 16 + (lane // 16) * 4 + e - weight = lds_ld(keys, index_head).bitcast(fx.Float32) - partial = ( - partial + fx.max(score_frag[e], fx.Float32(0.0)) * weight - ) - partial = _xred(partial, 16, lambda a, b: a + b) - partial = _xred(partial, 32, lambda a, b: a + b) - if lane < 16: - lds_st(red, wave * 16 + lane, partial) - gpu.barrier() - if (wave % 2 == 0) & (lane < 16) & (key_pos < bound): - score = lds_ld(red, wave * 16 + lane) + lds_ld( - red, (wave + 1) * 16 + lane - ) - put(mb("index_scores"), s * index_max_seq + key_pos, score) - stamp("index_score", tt, 4) - - for s in range(start("index_select"), S, G): - s = fx.Int32(s) - stamp("index_select", s, 0) - bound = pos0 + s + 1 - - def select_digit(shift, prefix, remain): - digit = 255 - fx.min(tid, 255) - count = (tid < 256).select(lds_ld(keys, digit), fx.Int32(0)) - inclusive = fx.coop.warp_inclusive_scan( - count, fx.ReductionOp.ADD, width=64 - ) - wave_total = read_lane_i32(inclusive, 63) - if (wave < 4) & (lane == 63): - lds_st(keys, 256 + wave, wave_total) - gpu.barrier() - before_wave = fx.Int32(0) - for w in range_constexpr(4): - before_wave = before_wave + (wave > w).select( - lds_ld(keys, 256 + w), fx.Int32(0) - ) - above = before_wave + inclusive - count - hit = (tid < 256) & (above < remain) & (above + count >= remain) - gpu.barrier() - if hit: - lds_st( - keys, 256, fx.Int32(prefix | (fx.Uint32(digit) << shift)) - ) - lds_st(keys, 257, remain - above) - gpu.barrier() - return fx.Uint32(lds_ld(keys, 256)), lds_ld(keys, 257) - - # Transform scores to monotonic integer keys and build the high - # byte histogram in the same pass, avoiding one extra 4-K LDS - # scan before the remaining radix digits. - if tid < 256: - lds_st(keys, tid, fx.Int32(0)) - gpu.barrier() - for batch in range_constexpr((index_max_seq + THREADS - 1) // THREADS): - i = tid + batch * THREADS - if i < bound: - bits = getf(mb("index_scores"), s * index_max_seq + i).bitcast( - fx.Int32 - ) - key = (bits >= 0).select(bits ^ fx.Int32(-(2**31)), ~bits) - # The selector CTA no longer needs the large GEMV staging - # region, so reuse it for the 4-K radix keys instead of - # increasing the monokernel's LDS allocation. - lds_st(xs, i, key.bitcast(fx.Float32)) - digit = fx.Int32((fx.Uint32(key) >> 24) & fx.Uint32(255)) - fx.atomic_add(keys + digit, fx.Int32(1), syncscope="workgroup") - gpu.barrier() - prefix, remain = select_digit( - 24, fx.Uint32(0), fx.min(fx.Int32(topk), bound) - ) - prefix_mask = 255 << 24 - for shift in (16, 8, 0): - if tid < 256: - lds_st(keys, tid, fx.Int32(0)) - gpu.barrier() - for batch in range_constexpr( - (index_max_seq + THREADS - 1) // THREADS - ): - i = tid + batch * THREADS - if i < bound: - key = fx.Uint32(lds_ld(xs, i).bitcast(fx.Int32)) - if (key & fx.Uint32(prefix_mask)) == prefix: - digit = fx.Int32((key >> shift) & fx.Uint32(255)) - fx.atomic_add( - keys + digit, fx.Int32(1), syncscope="workgroup" - ) - gpu.barrier() - prefix, remain = select_digit(shift, prefix, remain) - prefix_mask |= 255 << shift - - threshold = prefix - items = index_max_seq // THREADS - item_indices = [tid * items + j for j in range_constexpr(items)] - item_keys = [ - fx.Uint32(lds_ld(xs, fx.min(i, bound - 1)).bitcast(fx.Int32)) - for i in item_indices - ] - gt = [ - (i < bound) & (key > threshold) - for i, key in zip(item_indices, item_keys) - ] - eq = [ - (i < bound) & (key == threshold) - for i, key in zip(item_indices, item_keys) - ] - - def scan_flags(flags): - """Thread-major exclusive offsets for eight flags without extra LDS. - - The radix histogram already reserves keys[256:264] for wave - totals, so the scan can reuse those words and keep S=8 at the - 96-KiB static-LDS boundary. - """ - local = fx.Int32(0) - local_offsets = [] - for flag in flags: - local_offsets.append(local) - local = local + flag.select(fx.Int32(1), fx.Int32(0)) - inclusive = fx.coop.warp_inclusive_scan( - local, fx.ReductionOp.ADD, width=64 - ) - wave_total = read_lane_i32(inclusive, 63) - if lane == 63: - lds_st(keys, 256 + wave, wave_total) - gpu.barrier() - before_wave = fx.Int32(0) - total = fx.Int32(0) - for w in range_constexpr(WAVES): - wave_count = lds_ld(keys, 256 + w) - before_wave = before_wave + (wave > w).select( - wave_count, fx.Int32(0) - ) - total = total + wave_count - thread_base = before_wave + inclusive - local - return [thread_base + off for off in local_offsets], total - - gt_offsets, out_gt = scan_flags(gt) - gpu.barrier() - eq_offsets, _ = scan_flags(eq) - gpu.barrier() - for j in range_constexpr(items): - if gt[j]: - put(mb("indices"), s * topk + gt_offsets[j], item_indices[j]) - if eq[j] & (eq_offsets[j] < topk): - lds_st(red, eq_offsets[j], item_indices[j].bitcast(fx.Float32)) - gpu.barrier() - need_eq = fx.min(fx.Int32(topk), bound) - out_gt - for batch in range_constexpr((topk + THREADS - 1) // THREADS): - j = tid + batch * THREADS - if j >= bound: - put(mb("indices"), s * topk + j, fx.Int32(0)) - if j < need_eq: - i = lds_ld(red, j).bitcast(fx.Int32) - put(mb("indices"), s * topk + out_gt + j, i) - gpu.barrier() - stamp("index_select", s, 4) - - if const_expr(index_paged): - r_index_cache = _rsrc(indices) - r_out_indices = _rsrc(out_indices) - r_index_scores = _rsrc(mb("index_scores")) - hist = fx.recast_iter(fx.Int32, xs) - N_PAGED_SPLIT = index_max_seq // INDEX_KEYS_PER_TASK - - def load_index_q8(head, k): - words = fx.Vector.from_elements( - [lds_ld(xs, (head * INDEX_DIM + k) // 2 + j) for j in range(4)], - fx.Float32, - ) - return words.bitcast(fx.BFloat16) - - def radix_key(bits): - return (bits >= 0).select(bits ^ fx.Int32(-(2**31)), ~bits) - - def stored_key(s, i): - bits = fx.Int32( - bo.buffer_load( - r_index_scores, - (s * index_max_seq + i) * 2, - vec_width=1, - dtype=T.i32, - cache_modifier=CM_DEV, - ) - ) - return fx.Uint32(radix_key(bits)) - - # A key this launch writes (a verify row's earlier drafts, its own) is - # read from its row's index_k_new pair; the cache store is not yet visible. - new_active = [row_writes_cache(n) for n in range(S)] - new_slots = [row_slot(n) for n in range(S)] - - # Each wave scores groups of 16 keys over all 32 heads (two MFMA chains - # off one read of the keys), every 8th group of its CTA's tiles, with - # SCORE_DEPTH groups' reads in flight and their blocks a trip ahead. - SCORE_DEPTH = 2 - SCORE_STRIDE = 2 * G * INDEX_KEYS_PER_TASK - SCORE_BUF = 6 # a group in flight: 4 word pairs, its scale and block - # a wave's groups are SCORE_STRIDE keys apart: a row's last S keys reach one - # group of a wave at most - assert S + 15 < SCORE_STRIDE - - def lds_barrier(): - """A barrier for an LDS hand-off only: gpu.barrier would also wait on - the key reads in flight.""" - fx.memory_fence( - ordering=fx.AtomicOrdering.Release, syncscope="workgroup" - ) - rocdl.s_barrier() - fx.memory_fence( - ordering=fx.AtomicOrdering.Acquire, syncscope="workgroup" - ) - - def score_block(req, pos, key0, live): - """This lane's cache block in the group from key0 (0 when not live).""" - block = fx.Int32(0) - if live: - block = index_block(req, fx.min(key0 + lane % 16, pos)) - return block - - def score_fetch(block, pos, key0, live, kofs): - """A group's buffer: this lane's cache words for its 4 MFMA steps - (``kofs``: their offsets past the token's first word), its key's - scale and block (zeros when not live).""" - block = fx.Int32(block) - off = fx.min(key0 + lane % 16, pos) % index_block_size - w0 = fx.Vector.filled(2, 0, fx.Int32) - w1 = w0 - w2 = w0 - w3 = w0 - scale = fx.Float32(0.0) - if live: - word = index_value_word(block, off, 0) - w0, w1, w2, w3 = [ - fx.Vector( - bo.buffer_load( - r_index_cache, word + d, vec_width=2, dtype=T.i32 - ) - ) - for d in kofs - ] - scale = ld_f32(r_index_cache, index_scale_word(block, off)) - return [w0, w1, w2, w3, scale, block] - - def score_group(s, bound, pos, key0, buf, qf, wf, wave_u, may_new): - """The group's scores: lane l < 16 publishes key key0 + l's. - ``may_new``: the group holds one of the row's last S keys, the - only ones this launch writes.""" - key_pos = key0 + lane % 16 - tile = key0 // INDEX_KEYS_PER_TASK - group = (key0 % INDEX_KEYS_PER_TASK) // 16 - last = fx.min( - fx.Int32(INDEX_KEYS_PER_TASK // 16 - 1), - (pos - tile * INDEX_KEYS_PER_TASK) // 16, - ) - if group == 0: - stamp("index_score", s * N_PAGED_SPLIT + tile, 0, lead=wave_u * 64) - k_scale = fx.Float32(buf[4]) - if const_expr(may_new): - block = fx.Int32(buf[5]) - off = fx.min(key_pos, pos) % index_block_size - key_slot = block * index_block_size + off - new_row = fx.Int32(-1) - for n in range_constexpr(S): - hit = new_active[n] & (key_slot == new_slots[n]) - new_row = hit.select(fx.Int32(n), new_row) - is_new = new_row >= 0 - k_scale = is_new.select(fx.Float32(1.0), k_scale) + out.append(v) + return out + + def score_keys(g, key0, s_w, bound_w, qf, wf, kofs): + """The wave's group g (16 keys) of the batch, from key key0, for row s_w.""" + k = g * 16 + lane % 16 + p = k // index_block_size + off = k % index_block_size + base_w = index_value_word(p, off, 0) + k_scale = lds_ld(kw, index_scale_word(p, off)).bitcast(fx.Float32) acc0 = fx.Vector.filled(4, 0.0, fx.Float32) acc1 = acc0 for k32 in range_constexpr(INDEX_DIM // 32): - k = k32 * 32 + (lane // 16) * 8 - kw = fx.Vector(buf[k32]) - kv = _fp8_to_bf16x8(kw[0], kw[1]) - if const_expr(may_new): - if is_new: - kv_pairs = get_bf2_many( - [ - ( - mb("index_k_new"), - new_row * INDEX_DIM + k + j * 2, - ) - for j in range(4) - ] - ) - kv_values = [] - for pair in kv_pairs: - kv_values += list(pair) - kv = fx.Vector.from_elements(kv_values, fx.Float32).to( - fx.BFloat16 - ) + kv = _fp8_to_bf16x8( + lds_ld(kw, base_w + kofs[k32]), + lds_ld(kw, base_w + kofs[k32] + 1), + ) acc0 = fx.Vector( rocdl.mfma_f32_16x16x32_bf16( T.vec(4, T.f32), [qf[0][k32], kv, acc0] @@ -3055,153 +2837,92 @@ def score_group(s, bound, pos, key0, buf, qf, wf, wave_u, may_new): for e in range_constexpr(4): partial = partial + fx.max(acc[e], fx.Float32(0.0)) * wf[hg][e] partial = partial * k_scale - partial = _xred(partial, 16, lambda a, b: a + b) - partial = _xred(partial, 32, lambda a, b: a + b) + partial = _xred(partial, 16, lambda a, c: a + c) + partial = _xred(partial, 32, lambda a, c: a + c) halves.append(partial) - if (lane < 16) & (key_pos < bound): + key_pos = key0 + k + if (lane < 16) & (key_pos < bound_w): put( mb("index_scores"), - s * index_max_seq + key_pos, + s_w * index_max_seq + key_pos, halves[0] + halves[1], ) - if group == last: - stamp("index_score", s * N_PAGED_SPLIT + tile, 4, lead=wave_u * 64) - - def score_walk(s, row, init, qf, wf, kofs): - """This wave's ``count`` groups from key ``wave_key0``, SCORE_STRIDE - apart, off ``init``: the first SCORE_DEPTH groups' buffers, then the - next SCORE_DEPTH groups' blocks. A buffer is refilled right after its - group, with the group SCORE_DEPTH on. The first ``n_plain`` groups - hold none of the keys this launch writes; the rest (one at most) is - scored after them.""" - req, bound, pos, wave_u, wave_key0, count, n_plain = row - - def key0_of(k): - return wave_key0 + k * SCORE_STRIDE - - trips = n_plain // SCORE_DEPTH - for j, state in range(0, trips, 1, init=init): - k_trip = fx.Int32(j) * SCORE_DEPTH - refills = [] - lookups = [] - for i in range_constexpr(SCORE_DEPTH): - k_i = k_trip + i - score_group( - s, - bound, - pos, - key0_of(k_i), - state[i * SCORE_BUF : (i + 1) * SCORE_BUF], - qf, - wf, - wave_u, - False, - ) - k_fill = k_i + SCORE_DEPTH - refills += score_fetch( - state[SCORE_DEPTH * SCORE_BUF + i], - pos, - key0_of(k_fill), - k_fill < count, - kofs, - ) - k_look = k_i + 2 * SCORE_DEPTH - lookups.append( - score_block(req, pos, key0_of(k_look), k_look < count) - ) - carried = yield refills + lookups - for i in range_constexpr(SCORE_DEPTH): - k_tail = trips * SCORE_DEPTH + i - buf = carried[i * SCORE_BUF : (i + 1) * SCORE_BUF] - if k_tail < n_plain: - score_group( - s, bound, pos, key0_of(k_tail), buf, qf, wf, wave_u, False - ) - else: - if k_tail < count: - score_group( - s, - bound, - pos, - key0_of(k_tail), - buf, - qf, - wf, - wave_u, - True, - ) - def score_row(s): - # wave-uniform (scalar) values: the walk's branches and trip counts - bound = _uniform(index_bound(s)) - req = _uniform(index_request(s)) - pos = bound - 1 - n_tiles = (bound + INDEX_KEYS_PER_TASK - 1) // INDEX_KEYS_PER_TASK - first_cta = (101 + s * (G // S)) % G - first_tile = (bid + (G - first_cta)) & (G - 1) - if first_tile < n_tiles: - wave_u = _uniform(wave) - wave_key0 = ( - first_tile + (wave_u // 4) * G - ) * INDEX_KEYS_PER_TASK + (wave_u % 4) * 16 - count = (bound > wave_key0).select( - (bound - wave_key0 + SCORE_STRIDE - 1) // SCORE_STRIDE, - fx.Int32(0), - ) - plain_end = pos - (S - 1) - 15 - wave_key0 - n_plain = fx.min( - (plain_end > 0).select( - (plain_end + SCORE_STRIDE - 1) // SCORE_STRIDE, fx.Int32(0) - ), - count, + def score_group(s0, n, req, c0): + """This CTA's run of the keys of rows s0 .. s0 + n - 1 of a request whose + first row in the launch is c0.""" + rows = [fx.min(s0 + r, fx.Int32(S - 1)) for r in range(SCORE_ROWS)] + bounds = [ + (fx.Int32(r) < n).select(index_bound(rows[r]), fx.Int32(0)) + for r in range(SCORE_ROWS) + ] + bound_g = bounds[0] + for r in range_constexpr(1, SCORE_ROWS): + bound_g = fx.max(bound_g, bounds[r]) + bound_g = _uniform(bound_g) + n_pages = (bound_g + index_block_size - 1) // index_block_size + runs = fx.min(fx.Int32(G), (n_pages + SCORE_PAGES - 1) // SCORE_PAGES) + per = (n_pages + runs - 1) // runs + first_cta = (SCORE_BASE + s0 * (G // S)) % G + run = (fx.Int32(bid) + (G - first_cta)) % G + p_lo = run * per + p_hi = fx.min(p_lo + per, n_pages) + if (run < runs) & (p_lo < p_hi): + stamp("index_score", bid, 0) + # the run's block ids: loaded with the query's poll, stored after it + bt_word = index_block( + req, + (p_lo + fx.min(fx.Int32(tid), p_hi - p_lo - 1)) + * index_block_size, ) - row = (req, bound, pos, wave_u, wave_key0, count, n_plain) - # a token's word of dims d .. d + 3 is its first word plus that of - # dims d in token 0 of block 0 - kofs = [ - index_value_word(fx.Int32(0), 0, k32 * 32 + (lane // 16) * 8) - for k32 in range(INDEX_DIM // 32) - ] - # the first blocks' lookups overlap the query's staging - blocks = [ - score_block( - req, pos, wave_key0 + k * SCORE_STRIDE, fx.Int32(k) < count - ) - for k in range(2 * SCORE_DEPTH) - ] - get(mb("index_ready"), s) - if tid < INDEX_HEADS: - lds_st(keys, tid, get(mb("index_w"), s * INDEX_HEADS + tid)) - for b in range_constexpr((INDEX_Q_ROWS // 2) // THREADS): - q_pair = tid + b * THREADS - q_elem = q_pair * 2 - kq = q_elem % INDEX_DIM - q0, q1 = get_bf2_many( - [(mb("index_q"), s * INDEX_Q_ROWS + q_elem)] - )[0] - if kq < PE_DIM: - c = ld_bf16(_rsrc(rope_cos), pos * (PE_DIM // 2) + kq // 2) - sn = ld_bf16(_rsrc(rope_sin), pos * (PE_DIM // 2) + kq // 2) - q0, q1 = q0 * c - q1 * sn, q0 * sn + q1 * c - q0, q1 = fwht128(q0, q1) - lds_st(xs, q_pair, bf16_pair(q0, q1)) - init = [] - for i in range_constexpr(SCORE_DEPTH): - init += score_fetch( - blocks[i], - pos, - wave_key0 + i * SCORE_STRIDE, - fx.Int32(i) < count, - kofs, + # the rows' queries, rotated and transformed by index_qh, in one poll + specs = [] + for r in range_constexpr(SCORE_ROWS): + for i in range_constexpr(2): + specs.append( + ( + mb("index_qh"), + rows[r] * (INDEX_Q_ROWS // 2) + + fx.Int32(tid) * 4 + + 2 * i, + 2, + ) + ) + specs.append( + ( + mb("index_w"), + s0 * INDEX_HEADS + + fx.min(fx.Int32(tid), n * INDEX_HEADS - 1), + 1, ) - init += blocks[SCORE_DEPTH:] - lds_barrier() - # the row's rotated query (A operands of both head halves) and - # head weights, held for all of this wave's groups + ) + got = poll(specs, batch=len(specs)) + if tid < p_hi - p_lo: + lds_st(kw, BT_OFF + tid, bt_word) + gpu.barrier() + first = fetch_batch(fx.Int32(0), p_hi - p_lo) + for r in range_constexpr(SCORE_ROWS): + for i in range_constexpr(4): + lds_st( + xs, + r * (INDEX_Q_ROWS // 2) + fx.Int32(tid) * 4 + i, + got[r * 2 + i // 2][i % 2].bitcast(fx.Float32), + ) + if tid < SCORE_ROWS * INDEX_HEADS: + lds_st(keys, tid, got[SCORE_ROWS * 2][0]) + gpu.barrier() + stamp("index_score", bid, 2) + # this wave's row: its query (A operands of both head halves) and + # head weights, held for the whole run + wave_u = _uniform(wave) + r_w = wave_u // ROW_WAVES + s_w = s0 + fx.min(r_w, n - 1) + bound_w = _uniform(index_bound(s_w)) qf = [ [ load_index_q8( - hg * 16 + lane % 16, k32 * 32 + (lane // 16) * 8 + r_w * INDEX_HEADS + hg * 16 + lane % 16, + k32 * 32 + (lane // 16) * 8, ) for k32 in range(INDEX_DIM // 32) ] @@ -3209,113 +2930,90 @@ def score_row(s): ] wf = [ [ - lds_ld(keys, hg * 16 + (lane // 16) * 4 + e).bitcast( - fx.Float32 - ) + lds_ld( + keys, r_w * INDEX_HEADS + hg * 16 + (lane // 16) * 4 + e + ).bitcast(fx.Float32) for e in range(4) ] for hg in range(2) ] - if count > 0: - score_walk(s, row, init, qf, wf, kofs) - lds_barrier() - - for s in range_constexpr(S): - score_row(s) - - def select_digit(shift, prefix, remain): - digit = 255 - fx.min(tid, 255) - count = (tid < 256).select(lds_ld(keys, digit), fx.Int32(0)) - inclusive = fx.coop.warp_inclusive_scan( - count, fx.ReductionOp.ADD, width=64 - ) - wave_total = read_lane_i32(inclusive, 63) - if (wave < 4) & (lane == 63): - lds_st(keys, 256 + wave, wave_total) - gpu.barrier() - before_wave = fx.Int32(0) - for w in range_constexpr(4): - before_wave = before_wave + (wave > w).select( - lds_ld(keys, 256 + w), fx.Int32(0) - ) - above = before_wave + inclusive - count - hit = (tid < 256) & (above < remain) & (above + count >= remain) - gpu.barrier() - if hit: - lds_st(keys, 256, fx.Int32(prefix | (fx.Uint32(digit) << shift))) - lds_st(keys, 257, remain - above) - gpu.barrier() - return fx.Uint32(lds_ld(keys, 256)), lds_ld(keys, 257) - - def scan_flags(flags): - local = fx.Int32(0) - local_offsets = [] - for flag in flags: - local_offsets.append(local) - local = local + flag.select(fx.Int32(1), fx.Int32(0)) - inclusive = fx.coop.warp_inclusive_scan( - local, fx.ReductionOp.ADD, width=64 - ) - wave_total = read_lane_i32(inclusive, 63) - if lane == 63: - lds_st(keys, 256 + wave, wave_total) - gpu.barrier() - before_wave = fx.Int32(0) - total = fx.Int32(0) - for w in range_constexpr(WAVES): - wave_count = lds_ld(keys, 256 + w) - before_wave = before_wave + (wave > w).select( - wave_count, fx.Int32(0) - ) - total = total + wave_count - thread_base = before_wave + inclusive - local - return [thread_base + off for off in local_offsets], total - - def clear_hist(): - for b in range_constexpr(SEL_BINS // THREADS): - lds_st(hist, tid + b * THREADS, fx.Int32(0)) - gpu.barrier() - - def select_wide(shift, prefix, remain): - """Pick the 12-bit digit at ``shift`` holding the remain-th largest key. - - Returns (prefix with that digit, keys still needed, keys in that bin). - """ - per = SEL_BINS // THREADS - cs = [lds_ld(hist, SEL_BINS - 1 - (tid * per + j)) for j in range(per)] - local = cs[0] - for c in cs[1:]: - local = local + c - inclusive = fx.coop.warp_inclusive_scan( - local, fx.ReductionOp.ADD, width=64 - ) - wave_total = read_lane_i32(inclusive, 63) - if lane == 63: - lds_st(keys, 256 + wave, wave_total) - gpu.barrier() - running = inclusive - local - for w in range_constexpr(WAVES): - running = running + (wave > w).select( - lds_ld(keys, 256 + w), fx.Int32(0) - ) - for j in range_constexpr(per): - digit = SEL_BINS - 1 - (tid * per + j) - hit = (running < remain) & (running + cs[j] >= remain) - if hit: - lds_st( - keys, - 280, - fx.Int32(prefix | (fx.Uint32(digit) << shift)), + kofs = [ + index_value_word(fx.Int32(0), 0, k32 * 32 + (lane // 16) * 8) + for k32 in range(INDEX_DIM // 32) + ] + # the keys this launch writes among the rows': rows c0 .. s0 + n - 1 + n_new = s0 + n - c0 + new_bounds = [ + ((fx.Int32(r) >= c0) & (fx.Int32(r) < s0 + n)).select( + index_bound(r), fx.Int32(0) ) - lds_st(keys, 281, remain - running) - lds_st(keys, 282, cs[j]) - running = running + cs[j] - gpu.barrier() - return ( - fx.Uint32(lds_ld(keys, 280)), - lds_ld(keys, 281), - lds_ld(keys, 282), - ) + for r in range(S) + ] + n_run = p_hi - p_lo + n_batches = (n_run + BATCH_PAGES - 1) // BATCH_PAGES + for bi, words in range(0, n_batches, 1, init=first): + bi = fx.Int32(bi) + page0 = bi * BATCH_PAGES + n_pages_b = fx.min(fx.Int32(BATCH_PAGES), n_run - page0) + key0 = (p_lo + page0) * index_block_size + key_end = key0 + n_pages_b * index_block_size + gpu.barrier() + for i in range_constexpr(BATCH_LOADS): + w = (tid + i * THREADS) * 4 + if fx.Int32(w) < BATCH_WORDS: + for e in range_constexpr(4): + lds_st(kw, w + e, fx.Vector(words[i])[e]) + nxt = fetch_batch( + page0 + BATCH_PAGES, n_run - page0 - BATCH_PAGES + ) + gpu.barrier() + # the keys this launch writes, from their rows' pairs + hit = fx.Int32(0) + for r in range_constexpr(S): + hit = hit | ( + (new_bounds[r] > key0) & (new_bounds[r] <= key_end) + ).select(fx.Int32(1), fx.Int32(0)) + if _uniform(hit) != 0: + r_p = fx.min(c0 + tid // K8, fx.Int32(S - 1)) + pos_p = new_bounds[0] - 1 + for r in range_constexpr(1, S): + pos_p = (r_p == r).select(new_bounds[r] - 1, pos_p) + if ( + (fx.Int32(tid) < n_new * K8) + & (pos_p >= key0) + & (pos_p < key_end) + ): + item = tid % K8 + v = get(mb("index_k8"), r_p * K8 + item) + kk = pos_p - key0 + p = kk // index_block_size + off = kk % index_block_size + at = (item < K8 - 1).select( + index_value_word(p, off, item * 4), + index_scale_word(p, off), + ) + lds_st(kw, at, v) + gpu.barrier() + if r_w < n: + for g in range( + wave_u % ROW_WAVES, + n_pages_b * (index_block_size // 16), + fx.Int32(ROW_WAVES), + ): + g = fx.Int32(g) + if key0 + g * 16 < bound_w: + score_keys(g, key0, s_w, bound_w, qf, wf, kofs) + _ = yield nxt + gpu.barrier() + stamp("index_score", bid, 4) + + leads, sizes, firsts = score_groups() + for s in range(0, S, 1): + s = fx.Int32(s) + if ((leads >> s) & 1) != 0: + n = fx.Int32((sizes >> (fx.Int64(s) * 3)) & fx.Int64(7)) + c0 = fx.Int32((firsts >> (fx.Int64(s) * 3)) & fx.Int64(7)) + score_group(s, n, _uniform(index_request(s)), _uniform(c0)) def scan_pair(first, second, parity): """Exclusive block offsets of two flag lists in one scan. @@ -3369,246 +3067,173 @@ def set_bases(a, b): lds_st(keys, 273, b) gpu.barrier() - # Parallel paged top-k for rows of more than SEL_PAR_KEYS keys: up to - # SEL_SLICES CTAs a row (G // 2 ..) each histogram a slice of contiguous - # chunks, the row's index_select CTA merges them and publishes the - # threshold bin, the slice CTAs then collect their winners and candidates in - # position order, and index_select concatenates them in slice order: the - # same lists as the one-CTA path, so the same selection. - def row_items(chunk): - return [ - chunk * SEL_CHUNK + tid * SEL_ITEMS + j - for j in range_constexpr(SEL_ITEMS) + # Selection: each row's exact top-k over its slices, SEL_SLICES CTAs at + # most, each a run of the row's keys in position order. A radix select by + # 6-bit digits from the top: per digit every slice publishes the histogram + # of its keys under the prefix found so far, and every slice reads all of + # them, finds the digit holding the remain-th key and counts the keys above + # it in the slices before its own. Once the prefix has 12 or 24 bits and its + # keys fit the candidate stage, or has all 32, the prefix is the threshold's: + # each slice writes its keys above it (their places known), and its keys + # with the prefix go to the row's first slice, which resolves them (or, at + # 32 bits, ties, take their places directly). The lists are the one-CTA + # radix select's: keys above the threshold's 12-bit (24-, 32-bit) prefix in + # position order, then the prefix's keys above the threshold, then its ties. + # the keys a slice holds at most, as radix keys in LDS, in chunks of SEL_CHUNK + SLICE_KEYS = ( + ((index_max_seq + SEL_SLICES - 1) // SEL_SLICES + SEL_CHUNK - 1) + // SEL_CHUNK + * SEL_CHUNK + ) + DIGITS = ((26, 6), (20, 6), (14, 6), (8, 6), (2, 6), (0, 2)) + assert len(DIGITS) == INDEX_SEL_DIGITS + HIST_PAIRS = 32 # a slice's 64 counts, two 16-bit counts a pair + # LDS words: the slice's keys, its block ids, its bins (and the digit + # found), the row's bins + SBT = SLICE_KEYS + SBINS = SBT + SLICE_KEYS // index_block_size + 2 + STOT = SBINS + 72 + SALIST = STOT + 64 # the slice's keys above the threshold's prefix + assert SALIST + topk <= X_WORDS and SEL_SLICES * HIST_PAIRS <= 4 * THREADS + assert SLICE_KEYS < 1 << 16 + skeys = fx.recast_iter(fx.Int32, xs) + sbins = fx.recast_iter(fx.Int32, xs + SBINS) + stot = fx.recast_iter(fx.Int32, xs + STOT) + + red_i = fx.recast_iter(fx.Int32, red) + + def block_isums(vs): + """Block-wide sums of several per-thread Int32 values.""" + ws = [ + read_lane_i32( + fx.coop.warp_inclusive_scan(v, fx.ReductionOp.ADD, width=64), 63 + ) + for v in vs ] + if lane == 0: + for i in range_constexpr(len(vs)): + lds_st(red_i, i * WAVES + wave, ws[i]) + gpu.barrier() + tots = [] + for i in range_constexpr(len(vs)): + t = lds_ld(red_i, i * WAVES) + for w in range_constexpr(1, WAVES): + t = t + lds_ld(red_i, i * WAVES + w) + tots.append(t) + gpu.barrier() + return tots - def poll_row_keys(s, bound, items): - vals = poll( - [ - ( - mb("index_scores"), - s * index_max_seq + fx.min(i, bound - 1), - 1, - ) - for i in items - ] - ) - return [fx.Uint32(radix_key(v[0])) for v in vals] - - def poll_thresh(s): - th = poll([(mb("index_thresh"), s * 4 + e, 1) for e in range(3)]) - return fx.Uint32(th[0][0]), th[1][0], th[2][0] + def chunk_slots(chunk): + """This thread's places in a chunk of the slice (position order).""" + first_ = chunk * SEL_CHUNK + fx.Int32(tid) * SEL_ITEMS + return [first_ + j for j in range(SEL_ITEMS)] - r_part = _rsrc(mb("index_part_hist")) - r_win = _rsrc(mb("index_win_stage")) - r_cand = _rsrc(mb("index_cand_stage")) + # the selection's stores take resources made once, outside its branches + r_indices = _rsrc(mb("indices")) - for t in range(start("index_hist"), S * SEL_SLICES, G): - t = fx.Int32(t) - s, k = t // SEL_SLICES, t % SEL_SLICES - bound = index_bound(s) - if bound > SEL_PAR_KEYS: - n_sl, per = select_slices(bound) - if k < n_sl: - stamp("index_hist", t, 0) - clear_hist() - c0 = k * per - c1 = fx.min(c0 + per, (bound + SEL_CHUNK - 1) // SEL_CHUNK) - for chunk in range(c0, c1, fx.Int32(1)): - items = row_items(fx.Int32(chunk)) - item_keys = poll_row_keys(s, bound, items) - for j in range_constexpr(SEL_ITEMS): - if items[j] < bound: - fx.atomic_add( - hist + fx.Int32(item_keys[j] >> 20), - fx.Int32(1), - syncscope="workgroup", - ) - gpu.barrier() - for b in range_constexpr(SEL_BINS // THREADS): - bin_ = tid + b * THREADS - bo.buffer_store( - lds_ld(hist, bin_), - r_part, - t * SEL_BINS + bin_, - cache_modifier=CM_DEV, - ) - fx.memory_fence( - ordering=fx.AtomicOrdering.Release, syncscope="agent" - ) - gpu.barrier() - if tid == 0: - put(mb("index_part_ready"), t, fx.Int32(1)) - stamp("index_hist", t, 4) + def publish_slot(s, out_begin, p, slot): + """Place p of row s's list := cache slot ``slot``.""" + bo.buffer_store( + fx.Vector.from_elements([slot, tag], fx.Int32), + r_indices, + (s * topk + p) * 2, + cache_modifier=CM_DEV, + ) + bo.buffer_store( + slot, r_out_indices, out_begin + p, cache_modifier=CM_DEV + ) - for s in range(start("index_thresh"), S, G): - s = fx.Int32(s) - bound = index_bound(s) - if bound > SEL_PAR_KEYS: - stamp("index_thresh", s, 0) - n_sl, _ = select_slices(bound) - if wave == 0: - poll( - [ - ( - mb("index_part_ready"), - s * SEL_SLICES + fx.min(lane, n_sl - 1), - 1, - ) - ] - ) + def resolve(s, req, out_begin, count, level, prefix, remain, n_cand): + """The row's candidates (the keys under ``prefix``'s top 32 - ``level`` + bits with their cache slots, in position order, staged in LDS): the + remain-th largest, then + places count - remain_at_entry .. count - 1 of the list: the + candidates above it, then its first ties.""" + n_win = count - remain + + def cand_pass(shift, bits, prefix, remain): + """The next ``bits`` (≤ 10) of the threshold below ``prefix``: a + histogram of the candidates under it (2^bits bins in LDS, the slice's + dead keys), then the bin holding the remain-th.""" + nb = 1 << bits + mask = (0xFFFFFFFF << (shift + bits)) & 0xFFFFFFFF + per = (nb + THREADS - 1) // THREADS + for b in range_constexpr(per): + if fx.Int32(tid) + b * THREADS < nb: + lds_st(skeys, fx.Int32(tid) + b * THREADS, fx.Int32(0)) gpu.barrier() - # the histograms are plain stores released before their ready pairs - fx.memory_fence( - ordering=fx.AtomicOrdering.Acquire, syncscope="agent" - ) - for b in range_constexpr(SEL_BINS // THREADS): - bin_ = tid + b * THREADS - total = fx.Int32(0) - for kk in range_constexpr(SEL_SLICES): - part = fx.Int32( - bo.buffer_load( - r_part, - (s * SEL_SLICES + kk) * SEL_BINS + bin_, - vec_width=1, - dtype=T.i32, - cache_modifier=CM_DEV, - ) - ) - total = total + (fx.Int32(kk) < n_sl).select( - part, fx.Int32(0) + for c in range(fx.Int32(tid), n_cand, fx.Int32(THREADS)): + key = fx.Uint32( + lds_ld(xs, SEL_BINS + fx.Int32(c) * 2).bitcast(fx.Int32) + ) + if (key & fx.Uint32(mask)) == prefix: + digit = fx.Int32((key >> shift) & fx.Uint32(nb - 1)) + fx.atomic_add( + skeys + digit, fx.Int32(1), syncscope="workgroup" ) - lds_st(hist, bin_, total) gpu.barrier() - p12, r12, c12 = select_wide( - 20, fx.Uint32(0), fx.min(fx.Int32(topk), bound) - ) - if tid == 0: - put(mb("index_thresh"), s * 4, fx.Int32(p12)) - put(mb("index_thresh"), s * 4 + 1, r12) - put(mb("index_thresh"), s * 4 + 2, c12) - stamp("index_thresh", s, 4) - - for t in range(start("index_gather"), S * SEL_SLICES, G): - t = fx.Int32(t) - s, k = t // SEL_SLICES, t % SEL_SLICES - bound = index_bound(s) - if bound > SEL_PAR_KEYS: - n_sl, per = select_slices(bound) - if k < n_sl: - stamp("index_gather", t, 0) - p12, _, c12 = poll_thresh(s) - set_bases(fx.Int32(0), fx.Int32(0)) - c0 = k * per - c1 = fx.min(c0 + per, (bound + SEL_CHUNK - 1) // SEL_CHUNK) - # an overflowing bin is the one-CTA overflow path's: collect nothing - c1 = (c12 <= SEL_CAP).select(c1, c0) - for chunk in range(c0, c1, fx.Int32(1)): - chunk = fx.Int32(chunk) - items = row_items(chunk) - item_keys = poll_row_keys(s, bound, items) - high = [ - (i < bound) & ((key >> 20) > (p12 >> 20)) - for i, key in zip(items, item_keys) - ] - same = [ - (i < bound) & ((key >> 20) == (p12 >> 20)) - for i, key in zip(items, item_keys) - ] - win_off, cand_off, _, _ = scan_pair( - high, same, (chunk - c0) & 1 - ) - for j in range_constexpr(SEL_ITEMS): - if high[j]: - bo.buffer_store( - items[j], - r_win, - t * topk + win_off[j], - cache_modifier=CM_DEV, - ) - if same[j]: - bo.buffer_store( - fx.Vector.from_elements( - [item_keys[j].bitcast(fx.Int32), items[j]], - fx.Int32, - ), - r_cand, - (t * INDEX_CAND_STAGE_CAP + cand_off[j]) * 2, - cache_modifier=CM_DEV, - ) - gpu.barrier() - slot = (c1 - c0) & 1 - n_win = lds_ld(keys, 272 + slot * 2) - n_cand = lds_ld(keys, 273 + slot * 2) - fx.memory_fence( - ordering=fx.AtomicOrdering.Release, syncscope="agent" + # the bins from the top: per bins a thread, then waves, then the block + cs = [ + (fx.Int32(tid) * per + q < nb).select( + lds_ld( + skeys, + fx.max(nb - 1 - (fx.Int32(tid) * per + q), fx.Int32(0)), + ), + fx.Int32(0), ) - gpu.barrier() - if tid == 0: - put(mb("index_gather_count"), t, n_win | (n_cand << 16)) - stamp("index_gather", t, 4) - - for s in range(start("index_select"), S, G): - s = fx.Int32(s) - stamp("index_select", s, 0) - bound = index_bound(s) - req = index_request(s) - count = fx.min(fx.Int32(topk), bound) - out_begin, _ = row_index_bounds(s) - n_chunks = (bound + SEL_CHUNK - 1) // SEL_CHUNK - - def chunk_items(chunk): - return [ - chunk * SEL_CHUNK + tid * SEL_ITEMS + j - for j in range_constexpr(SEL_ITEMS) + for q in range(per) ] - - def publish(p, i): - slot = index_block(req, i) * index_block_size + i % index_block_size - put(mb("indices"), s * topk + p, slot) - bo.buffer_store( - slot, r_out_indices, out_begin + p, cache_modifier=CM_DEV + local = cs[0] + for q in range_constexpr(1, per): + local = local + cs[q] + inc = fx.coop.warp_inclusive_scan( + local, fx.ReductionOp.ADD, width=64 ) - - def poll_keys(items): - vals = poll( - [ - ( - mb("index_scores"), - s * index_max_seq + fx.min(i, bound - 1), - 1, + if lane == 63: + lds_st(keys, 256 + wave, inc) + gpu.barrier() + running = inc - local + for w in range_constexpr(WAVES): + running = running + (wave > w).select( + lds_ld(keys, 256 + w), fx.Int32(0) + ) + for q in range_constexpr(per): + digit = nb - 1 - (fx.Int32(tid) * per + q) + if (running < remain) & (running + cs[q] >= remain): + lds_st( + keys, + 280, + fx.Int32(prefix | (fx.Uint32(digit) << shift)), ) - for i in items - ] - ) - return [fx.Uint32(radix_key(v[0])) for v in vals] + lds_st(keys, 281, remain - running) + running = running + cs[q] + gpu.barrier() + return fx.Uint32(lds_ld(keys, 280)), lds_ld(keys, 281) - def load_keys(items): - """Keys of consecutive item pairs, two (value, tag) pairs per load.""" - out = [] - for j in range_constexpr(0, len(items), 2): - w = fx.Vector( - bo.buffer_load( - r_index_scores, - ( - s * index_max_seq - + fx.min(items[j], index_max_seq - 2) - ) - * 2, - vec_width=4, - dtype=T.i32, - cache_modifier=CM_DEV, - ) - ) - out += [fx.Uint32(radix_key(w[0])), fx.Uint32(radix_key(w[2]))] - return out - - def gather(chunk, valid, item_keys, values, threshold, out_gt, remain): - """Keys above ``threshold`` append to red from the running base; the - first ``remain`` ties fill red[out_gt:].""" - gt = [v & (k > threshold) for v, k in zip(valid, item_keys)] - eq = [v & (k == threshold) for v, k in zip(valid, item_keys)] + if const_expr(level == 20): + prefix, remain = cand_pass(10, 10, prefix, remain) + prefix, remain = cand_pass(0, 10, prefix, remain) + else: + prefix, remain = cand_pass(0, 8, prefix, remain) + set_bases(n_win, fx.Int32(0)) + out_gt = count - remain + n_cand_chunks = (n_cand + SEL_CHUNK - 1) // SEL_CHUNK + for chunk in range(fx.Int32(0), n_cand_chunks, fx.Int32(1)): + chunk = fx.Int32(chunk) + idx = chunk_slots(chunk) + safe = [fx.min(c, n_cand - 1) for c in idx] + valid = [c < n_cand for c in idx] + item_keys = [ + fx.Uint32(lds_ld(xs, SEL_BINS + c * 2).bitcast(fx.Int32)) + for c in safe + ] + values = [ + lds_ld(xs, SEL_BINS + c * 2 + 1).bitcast(fx.Int32) for c in safe + ] + gt = [v & (k > prefix) for v, k in zip(valid, item_keys)] + eq = [v & (k == prefix) for v, k in zip(valid, item_keys)] gt_offsets, eq_offsets, _, _ = scan_pair(gt, eq, chunk & 1) - for j in range_constexpr(len(values)): + for j in range_constexpr(SEL_ITEMS): if gt[j]: lds_st(red, gt_offsets[j], values[j].bitcast(fx.Float32)) if eq[j] & (eq_offsets[j] < remain): @@ -3617,247 +3242,322 @@ def gather(chunk, valid, item_keys, values, threshold, out_gt, remain): out_gt + eq_offsets[j], values[j].bitcast(fx.Float32), ) - - def publish_all(): - gpu.barrier() - for batch in range_constexpr((topk + THREADS - 1) // THREADS): - j = tid + batch * THREADS - if j < count: - publish(j, lds_ld(red, j).bitcast(fx.Int32)) - - def candidates(level, prefix, remain, n_cand, collect=True): - """Keys above the ``level`` bin go to red, the bin's keys to LDS; - for a parallel row, index_select copied the slices' lists there - (``collect=False``).""" - if const_expr(collect): - set_bases(fx.Int32(0), fx.Int32(0)) - for chunk in range(fx.Int32(0), n_chunks, fx.Int32(1)): - chunk = fx.Int32(chunk) - items = chunk_items(chunk) - item_keys = load_keys(items) - high = [ - (i < bound) & ((k >> level) > (prefix >> level)) - for i, k in zip(items, item_keys) - ] - same = [ - (i < bound) & ((k >> level) == (prefix >> level)) - for i, k in zip(items, item_keys) - ] - win_offsets, cand_offsets, _, _ = scan_pair( - high, same, chunk & 1 + gpu.barrier() + for batch in range_constexpr((topk + THREADS - 1) // THREADS): + j = tid + batch * THREADS + if (j >= n_win) & (j < count): + publish_slot(s, out_begin, j, lds_ld(red, j).bitcast(fx.Int32)) + + def sel_digit( + s, c, n_sl, level, shift, bits, lo, hi, n_chunks, prefix, remain + ): + """One radix digit of row s on its slice c (of n_sl; keys lo .. hi - 1 in + LDS): this slice's histogram of its keys under ``prefix``, every + slice's, the digit holding the remain-th key. Returns (digit, keys still + needed in it, keys in it, keys above it, those in the slices before c, + keys in it in the slices before c).""" + top = shift + bits + if tid < 64: + lds_st(sbins, tid, fx.Int32(0)) + lds_st(stot, tid, fx.Int32(0)) + gpu.barrier() + for chunk in range(fx.Int32(0), n_chunks, fx.Int32(1)): + slots = chunk_slots(fx.Int32(chunk)) + for j in range_constexpr(SEL_ITEMS): + key = fx.Uint32(lds_ld(skeys, slots[j])) + under = lo + slots[j] < hi + if const_expr(top < 32): + under = under & ((key >> top) == (prefix >> top)) + if under: + fx.atomic_add( + sbins + + fx.Int32((key >> shift) & fx.Uint32((1 << bits) - 1)), + fx.Int32(1), + syncscope="workgroup", ) - for j in range_constexpr(SEL_ITEMS): - if high[j]: - lds_st( - red, - win_offsets[j], - items[j].bitcast(fx.Float32), - ) - c = SEL_BINS + cand_offsets[j] * 2 - if same[j]: - lds_st(xs, c, item_keys[j].bitcast(fx.Float32)) - lds_st(xs, c + 1, items[j].bitcast(fx.Float32)) - gpu.barrier() - n_win = count - remain - stamp("index_select", s, 2) + gpu.barrier() + hbase = ((s * len(DIGITS) + level) * SEL_SLICES) * HIST_PAIRS + if tid < HIST_PAIRS: + put( + mb("sel_hist"), + hbase + c * HIST_PAIRS + tid, + lds_ld(sbins, 2 * tid) | (lds_ld(sbins, 2 * tid + 1) << 16), + ) + n_pairs = n_sl * HIST_PAIRS + qs = [fx.Int32(tid) + q * THREADS for q in range(4)] + got = poll( + [(mb("sel_hist"), hbase + fx.min(q, n_pairs - 1), 1) for q in qs], + batch=4, + ) + words = [ + (qs[q] < n_pairs).select(got[q][0], fx.Int32(0)) for q in range(4) + ] + for q in range_constexpr(4): + bin2 = (qs[q] % HIST_PAIRS) * 2 + fx.atomic_add( + stot + bin2, words[q] & fx.Int32(0xFFFF), syncscope="workgroup" + ) + fx.atomic_add( + stot + bin2 + 1, + (words[q] >> 16) & fx.Int32(0xFFFF), + syncscope="workgroup", + ) + gpu.barrier() + # the digit holding the remain-th key: wave 0 scans from the top + if wave == 0: + cnt = lds_ld(stot, 63 - fx.Int32(lane)) + inc = fx.coop.warp_inclusive_scan(cnt, fx.ReductionOp.ADD, width=64) + above = inc - cnt + if (above < remain) & (inc >= remain): + lds_st(sbins, 64, 63 - fx.Int32(lane)) + lds_st(sbins, 65, remain - above) + lds_st(sbins, 66, cnt) + lds_st(sbins, 67, above) + gpu.barrier() + digit = lds_ld(sbins, 64) + rem = lds_ld(sbins, 65) + in_bin = lds_ld(sbins, 66) + a_all = lds_ld(sbins, 67) + # the slices before c: their keys above the digit and in it + mine_above = fx.Int32(0) + mine_in = fx.Int32(0) + for q in range_constexpr(4): + before = (qs[q] < n_pairs) & (qs[q] // HIST_PAIRS < c) + for h in range_constexpr(2): + bn = (qs[q] % HIST_PAIRS) * 2 + h + n_b = before.select( + (words[q] >> (16 * h)) & fx.Int32(0xFFFF), fx.Int32(0) + ) + mine_above = mine_above + (bn > digit).select(n_b, fx.Int32(0)) + mine_in = mine_in + (bn == digit).select(n_b, fx.Int32(0)) + sums = block_isums([mine_above, mine_in]) + return digit, rem, in_bin, a_all, sums[0], sums[1] + + def sel_place(st, p, i): + """Place p of the slice's row := key i (its block id from LDS).""" + s, req, out_begin, lo, hi, first_page = st[:6] + blk = lds_ld(hist, SBT + i // index_block_size - first_page) + publish_slot( + s, out_begin, p, blk * index_block_size + i % index_block_size + ) - def cand_pass(shift, mask, prefix, remain): - if tid < 256: - lds_st(keys, tid, fx.Int32(0)) - gpu.barrier() - for c in range(tid, n_cand, fx.Int32(THREADS)): - key = fx.Uint32( - lds_ld(xs, SEL_BINS + fx.Int32(c) * 2).bitcast(fx.Int32) - ) - if (key & fx.Uint32(mask)) == prefix: - digit = fx.Int32((key >> shift) & fx.Uint32(255)) - fx.atomic_add( - keys + digit, fx.Int32(1), syncscope="workgroup" - ) - gpu.barrier() - return select_digit(shift, prefix, remain) + def sel_stage(st, chunk, in_base, e): + """Entry e of a chunk's compacted keys under the prefix: a tie placed + (at 32 bits) or a candidate staged for the row's first slice.""" + s, req, out_begin, lo, hi, first_page, low, lp, low_mask = st[:9] + prefix_i, decided, remain, above_all, in_before = st[9:] + word = lds_ld(skeys, chunk * SEL_CHUNK + e) + pos = lo + chunk * SEL_CHUNK + (word & fx.Int32(0xFFF)) + rk = in_before + in_base + e + if (decided == 32) & (rk < remain): + sel_place(st, above_all + rk, pos) + if decided != 32: + blk = lds_ld(hist, SBT + pos // index_block_size - first_page) + put2( + mb("sel_cand"), + (s * SEL_CAP + rk) * 2, + prefix_i | fx.Int32(fx.Uint32(word) >> 12), + blk * index_block_size + pos % index_block_size, + ) - if const_expr(level == 20): - prefix, remain = cand_pass(12, 0xFFF00000, prefix, remain) - prefix, remain = cand_pass(4, 0xFFFFF000, prefix, remain) - prefix, remain = cand_pass(0, 0xFFFFFFF0, prefix, remain) - else: - prefix, remain = cand_pass(0, 0xFFFFFF00, prefix, remain) - set_bases(n_win, fx.Int32(0)) - stamp("index_select", s, 3) - out_gt = count - remain - n_cand_chunks = (n_cand + SEL_CHUNK - 1) // SEL_CHUNK - for chunk in range(fx.Int32(0), n_cand_chunks, fx.Int32(1)): - chunk = fx.Int32(chunk) - idx = chunk_items(chunk) - safe = [fx.min(c, n_cand - 1) for c in idx] - gather( - chunk, - [c < n_cand for c in idx], + def sel_place_chunk(st, chunk): + """A chunk's keys under the prefix, compacted in place in its own LDS + words (a word each: the key's bits below the prefix, its place in the + chunk), then staged or placed a thread an entry; its keys above the + prefix listed after the slice's keys.""" + s, req, out_begin, lo, hi, first_page, low, lp, low_mask = st[:9] + slots = chunk_slots(chunk) + ks = [fx.Uint32(lds_ld(skeys, q)) for q in slots] + valid = [lo + q < hi for q in slots] + above_f = [v & ((k >> low) > lp) for v, k in zip(valid, ks)] + in_f = [v & ((k >> low) == lp) for v, k in zip(valid, ks)] + a_offs, i_offs, _, in_end = scan_pair(above_f, in_f, chunk & 1) + in_base = lds_ld(keys, 273 + (chunk & 1) * 2) + gpu.barrier() + for j in range_constexpr(SEL_ITEMS): + if above_f[j]: + lds_st(skeys, SALIST + a_offs[j], lo + slots[j]) + if in_f[j]: + lds_st( + skeys, + chunk * SEL_CHUNK + i_offs[j] - in_base, + (fx.Int32(ks[j] & low_mask) << 12) + | (slots[j] - chunk * SEL_CHUNK), + ) + gpu.barrier() + for e in range(fx.Int32(tid), in_end - in_base, fx.Int32(THREADS)): + sel_stage(st, chunk, in_base, fx.Int32(e)) + + # Rows' slices interleave (slice c of row s is task c * S + s) when every + # row has at most G / S slices, so short rows' slices share one compact + # range of CTAs; else each row's take a run of SEL_SLICES tasks, whole rows + # a round of the grid, so a row never waits on a slice queued behind another + # row's. + max_slices = fx.Int32(0) + for r in range_constexpr(S): + max_slices = fx.max( + max_slices, select_slices(index_bound(r), fx.min, fx.max)[0] + ) + interleave = _uniform(max_slices) * S <= G + for t in range(start("index_slice"), S * SEL_SLICES, G): + t = fx.Int32(t) + s = interleave.select(t % S, t // SEL_SLICES) + c = interleave.select(t // S, t % SEL_SLICES) + bound = _uniform(index_bound(s)) + n_sl, per = select_slices(bound, fx.min, fx.max) + # the timeline's task: the row's slice, whatever the CTA + t = s * SEL_SLICES + c + if (bound > 0) & (c < n_sl): + stamp("index_slice", t, 0) + req = _uniform(index_request(s)) + count = fx.min(fx.Int32(topk), bound) + out_begin, _ = row_index_bounds(s) + lo = c * per + hi = fx.min(lo + per, bound) + n_chunks = (per + SEL_CHUNK - 1) // SEL_CHUNK + first_page = lo // index_block_size + n_pages = ( + hi + index_block_size - 1 + ) // index_block_size - first_page + for pg in range(fx.Int32(tid), n_pages, fx.Int32(THREADS)): + pg = fx.Int32(pg) + lds_st( + hist, + SBT + pg, + index_block(req, (first_page + pg) * index_block_size), + ) + # the slice's scores, as radix keys in LDS + for chunk in range(fx.Int32(0), n_chunks, fx.Int32(1)): + slots = chunk_slots(fx.Int32(chunk)) + vals = poll( [ - fx.Uint32( - lds_ld(xs, SEL_BINS + c * 2).bitcast(fx.Int32) + ( + mb("index_scores"), + s * index_max_seq + + fx.max(fx.min(lo + q, hi - 1), lo), + 1, ) - for c in safe - ], - [ - lds_ld(xs, SEL_BINS + c * 2 + 1).bitcast(fx.Int32) - for c in safe + for q in slots ], - prefix, - out_gt, - remain, - ) - stamp("index_select", s, 5) - publish_all() - - def full_radix(): - """Byte-wise radix over every key; used when the threshold bin overflows LDS.""" - - def radix_pass(shift, prefix_mask, prefix, remain): - if tid < 256: - lds_st(keys, tid, fx.Int32(0)) - gpu.barrier() - for chunk in range(fx.Int32(0), n_chunks, fx.Int32(1)): - items = chunk_items(fx.Int32(chunk)) - item_keys = load_keys(items) - for j in range_constexpr(SEL_ITEMS): - key = item_keys[j] - if (items[j] < bound) & ( - (key & fx.Uint32(prefix_mask)) == prefix - ): - digit = fx.Int32((key >> shift) & fx.Uint32(255)) - fx.atomic_add( - keys + digit, fx.Int32(1), syncscope="workgroup" - ) - gpu.barrier() - return select_digit(shift, prefix, remain) - - prefix, remain = radix_pass(24, 0, fx.Uint32(0), count) - prefix, remain = radix_pass(16, 0xFF000000, prefix, remain) - prefix, remain = radix_pass(8, 0xFFFF0000, prefix, remain) - prefix, remain = radix_pass(0, 0xFFFFFF00, prefix, remain) - set_bases(fx.Int32(0), fx.Int32(0)) - out_gt = count - remain - for chunk in range(fx.Int32(0), n_chunks, fx.Int32(1)): - chunk = fx.Int32(chunk) - items = chunk_items(chunk) - gather( - chunk, - [i < bound for i in items], - load_keys(items), - items, - prefix, - out_gt, - remain, + batch=SEL_ITEMS, ) - publish_all() - - def overflow(p12, r12, polled=True): - """The threshold bin overflows LDS: its keys' next 12 bits, then - candidates or a byte-wise radix over every key (this CTA). A - parallel row's CTA has not polled the row's scores - (``polled=False``).""" - clear_hist() - for chunk in range(fx.Int32(0), n_chunks, fx.Int32(1)): - items = chunk_items(fx.Int32(chunk)) - if const_expr(polled): - item_keys = load_keys(items) - else: - item_keys = poll_keys(items) for j in range_constexpr(SEL_ITEMS): - k = item_keys[j] - if (items[j] < bound) & ((k >> 20) == (p12 >> 20)): - fx.atomic_add( - hist + fx.Int32((k >> 8) & fx.Uint32(0xFFF)), - fx.Int32(1), - syncscope="workgroup", - ) + lds_st(skeys, slots[j], radix_key(vals[j][0])) gpu.barrier() - p24, r24, c24 = select_wide(8, p12, r12) - if c24 <= SEL_CAP: - candidates(8, p24, r24, c24) - if c24 > SEL_CAP: - full_radix() - - if (bound > 0) & (bound <= SEL_PAR_KEYS): - clear_hist() - for chunk in range(fx.Int32(0), n_chunks, fx.Int32(1)): - items = chunk_items(fx.Int32(chunk)) - item_keys = poll_keys(items) - for j in range_constexpr(SEL_ITEMS): - if items[j] < bound: - fx.atomic_add( - hist + fx.Int32(item_keys[j] >> 20), - fx.Int32(1), - syncscope="workgroup", + stamp("index_slice", t, 1) + prefix, remain = fx.Uint32(0), count + above_before, above_all = fx.Int32(0), fx.Int32(0) + in_before, in_all = fx.Int32(0), fx.Int32(0) + done, decided = fx.Int32(0), fx.Int32(0) + for level, (shift, bits) in enumerate(DIGITS): + if done == 0: + digit, remain, in_all, a_all, a_before, in_before = ( + sel_digit( + s, + c, + n_sl, + level, + shift, + bits, + lo, + hi, + n_chunks, + prefix, + remain, ) - gpu.barrier() - p12, r12, c12 = select_wide(20, fx.Uint32(0), count) - stamp("index_select", s, 1) - if c12 <= SEL_CAP: - candidates(20, p12, r12, c12) - if c12 > SEL_CAP: - overflow(p12, r12) - if bound > SEL_PAR_KEYS: - p12, r12, c12 = poll_thresh(s) - stamp("index_select", s, 1) - if c12 <= SEL_CAP: - # the slices' winners to red and candidates to LDS, slice order - n_sl, _ = select_slices(bound) - counts = fx.recast_iter(fx.Int32, misc) - if (wave == 0) & (lane < n_sl): - lds_st( - counts, - lane, - get(mb("index_gather_count"), s * SEL_SLICES + lane), ) - gpu.barrier() - # the staged lists are plain stores released before their counts - fx.memory_fence( - ordering=fx.AtomicOrdering.Acquire, syncscope="agent" + above_all = above_all + a_all + above_before = above_before + a_before + prefix = prefix | (fx.Uint32(digit) << shift) + if const_expr(shift == 0): + done = fx.Int32(1) + decided = fx.Int32(32) + elif const_expr(32 - shift in (12, 24)): + fits = in_all <= SEL_CAP + done = fits.select(fx.Int32(1), fx.Int32(0)) + decided = fits.select(fx.Int32(32 - shift), fx.Int32(0)) + stamp("index_slice", t, 2) + # this slice's keys above the threshold's prefix take their places; + # its keys under it are ties placed, or candidates staged + low = fx.Uint32(32 - decided) + lp = prefix >> low + # the key bits below the prefix (none at 32 bits) + low_mask = (decided < 32).select( + fx.Uint32(0xFFFFFFFF) + >> fx.Uint32(fx.min(decided, fx.Int32(31))), + fx.Uint32(0), + ) + + st = ( + s, + req, + out_begin, + lo, + hi, + first_page, + low, + lp, + low_mask, + fx.Int32(prefix.bitcast(fx.Int32)), + decided, + remain, + above_all, + in_before, + ) + + set_bases(fx.Int32(0), fx.Int32(0)) + for pchunk in range(fx.Int32(0), n_chunks, fx.Int32(1)): + sel_place_chunk(st, fx.Int32(pchunk)) + gpu.barrier() + n_above = lds_ld(keys, 272 + (n_chunks & 1) * 2) + for pe in range(fx.Int32(tid), n_above, fx.Int32(THREADS)): + sel_place( + st, + above_before + fx.Int32(pe), + lds_ld(skeys, SALIST + fx.Int32(pe)), ) - n_win, n_cand = fx.Int32(0), fx.Int32(0) - for kk in range_constexpr(SEL_SLICES): - v = (fx.Int32(kk) < n_sl).select( - lds_ld(counts, kk), fx.Int32(0) + stamp("index_slice", t, 3) + # the row's first slice resolves the candidates + if (c == 0) & (decided < 32): + gpu.barrier() + for c0 in range(fx.Int32(0), in_all, fx.Int32(4 * THREADS)): + cs = [ + fx.Int32(c0) + fx.Int32(tid) + q * THREADS + for q in range(4) + ] + kv = poll( + [ + ( + mb("sel_cand"), + (s * SEL_CAP + fx.min(ci, in_all - 1)) * 2, + 2, + ) + for ci in cs + ], + batch=4, ) - w, c = v & fx.Int32(0xFFFF), v >> 16 - t = s * SEL_SLICES + kk - for wj in range(tid, w, fx.Int32(THREADS)): - staged_win = fx.Int32( - bo.buffer_load( - r_win, - t * topk + fx.Int32(wj), - vec_width=1, - dtype=T.i32, - cache_modifier=CM_DEV, + for q in range_constexpr(4): + if cs[q] < in_all: + lds_st( + xs, + SEL_BINS + cs[q] * 2, + kv[q][0].bitcast(fx.Float32), ) - ) - lds_st( - red, - n_win + fx.Int32(wj), - staged_win.bitcast(fx.Float32), - ) - for cj in range(tid, c, fx.Int32(THREADS)): - staged_cand = fx.Vector( - bo.buffer_load( - r_cand, - (t * INDEX_CAND_STAGE_CAP + fx.Int32(cj)) * 2, - vec_width=2, - dtype=T.i32, - cache_modifier=CM_DEV, + lds_st( + xs, + SEL_BINS + cs[q] * 2 + 1, + kv[q][1].bitcast(fx.Float32), ) - ) - at = SEL_BINS + (n_cand + fx.Int32(cj)) * 2 - lds_st(xs, at, staged_cand[0].bitcast(fx.Float32)) - lds_st(xs, at + 1, staged_cand[1].bitcast(fx.Float32)) - n_win, n_cand = n_win + w, n_cand + c - candidates(20, p12, r12, c12, collect=False) - if c12 > SEL_CAP: - overflow(p12, r12, polled=False) - gpu.barrier() - stamp("index_select", s, 4) + gpu.barrier() + stamp("index_slice", t, 5) + if decided == 12: + resolve( + s, req, out_begin, count, 20, prefix, remain, in_all + ) + else: + resolve(s, req, out_begin, count, 8, prefix, remain, in_all) + gpu.barrier() + stamp("index_slice", t, 4) # ================================== 5. sparse MLA split: 32 keys x 8 heads r_kv = _rsrc(kv_cache) diff --git a/sgl_amd_mono/glm53/gfx950/layout.py b/sgl_amd_mono/glm53/gfx950/layout.py index 96e82e4d..576fc21a 100644 --- a/sgl_amd_mono/glm53/gfx950/layout.py +++ b/sgl_amd_mono/glm53/gfx950/layout.py @@ -37,23 +37,25 @@ INDEX_Q_ROWS = INDEX_HEADS * INDEX_DIM INDEX_TILE = 16 INDEX_KEYS_PER_TASK = 64 -# a long row's paged top-k runs on up to this many slices (CTAs) per row, each with -# staging room for this many threshold-bin candidates -INDEX_SEL_SLICES = 16 +# a row's top-k runs on up to this many slices (CTAs), with staging room for this many +# candidates of the threshold's prefix +INDEX_SEL_SLICES = 64 INDEX_CAND_STAGE_CAP = 16384 -# keys a select chunk (a CTA's 512 threads x 8), and the chunks past which a row selects -# on up to INDEX_SEL_SLICES CTAs: up to 8 chunks the hand-offs cost what the slices save -INDEX_SEL_CHUNK = 4096 -INDEX_SEL_PAR_CHUNKS = 8 - - -def select_slices(bound): - """(slices, chunks a slice) of a row of ``bound`` > 0 keys whose top-k runs on - several CTAs: runs of whole chunks in position order, at most INDEX_SEL_SLICES, none - empty. The kernel calls it on traced Int32 and the CPU test on ints.""" - chunks = (bound + INDEX_SEL_CHUNK - 1) // INDEX_SEL_CHUNK - per = (chunks + INDEX_SEL_SLICES - 1) // INDEX_SEL_SLICES - return (chunks + per - 1) // per, per +# radix digits of the sliced top-k; a row has a slice per this many keys (rounded up), +# at most INDEX_SEL_SLICES +INDEX_SEL_DIGITS = 6 +INDEX_SEL_SLICE_KEYS = 1024 + + +def select_slices(bound, minimum=min, maximum=max): + """(slices, keys a slice) of a row of ``bound`` keys: runs of the row's keys in + position order, none empty (one for an empty row). The kernel calls it on traced + Int32 with ``fx.min`` and ``fx.max``, the CPU test on ints.""" + n = minimum( + (bound + INDEX_SEL_SLICE_KEYS - 1) // INDEX_SEL_SLICE_KEYS, INDEX_SEL_SLICES + ) + n = maximum(n, 1) + return n, (bound + n - 1) // n DCP_SUMMARY_PAIRS = KV_LORA // 2 + 2 @@ -213,22 +215,20 @@ def layout( if with_indexer: items += [ ("index_k", samples * INDEX_DIM * pair_bytes), - ("index_k_new", samples * INDEX_DIM // 2 * pair_bytes), - ("index_ready", samples * pair_bytes), + # each row's new key as cached: its FP8 value words, then its scale + ("index_k8", samples * (INDEX_DIM // 4 + 1) * pair_bytes), ("index_w", samples * INDEX_HEADS * pair_bytes), ("index_q", samples * INDEX_Q_ROWS // 2 * pair_bytes), + ("index_qh", samples * INDEX_Q_ROWS // 2 * pair_bytes), ("index_scores", samples * index_max_seq * pair_bytes), ("indices", samples * sparse_attention_topk * pair_bytes), - # parallel top-k: per slice a 4096-bin histogram and its ready pair, per - # row the threshold (bin prefix, keys still needed, keys in the bin), per - # slice its winners' positions, its (key, position) candidates and their - # counts - ("index_part_hist", samples * INDEX_SEL_SLICES * 4096 * 4), - ("index_part_ready", samples * INDEX_SEL_SLICES * pair_bytes), - ("index_thresh", samples * 4 * pair_bytes), - ("index_win_stage", samples * INDEX_SEL_SLICES * sparse_attention_topk * 4), - ("index_cand_stage", samples * INDEX_SEL_SLICES * INDEX_CAND_STAGE_CAP * 8), - ("index_gather_count", samples * INDEX_SEL_SLICES * pair_bytes), + # sliced top-k: per row, digit and slice the slice's 64 counts (two a + # pair); per row the candidates (key, cache slot) under the threshold's prefix + ( + "sel_hist", + samples * INDEX_SEL_DIGITS * INDEX_SEL_SLICES * 32 * pair_bytes, + ), + ("sel_cand", samples * INDEX_CAND_STAGE_CAP * 2 * pair_bytes), ] offset, scratch = 0, {} @@ -287,22 +287,24 @@ def stage_tasks( ), ("q_norm", samples), ("cache", samples), + ] + if with_indexer: + tasks += [("index_key", samples)] + tasks += [ ("q_b", heads * (NOPE_DIM + PE_DIM) // Q_B_TILE), ] if with_indexer: - tasks += [("index_q", INDEX_Q_ROWS // INDEX_TILE)] + # index_qh: each row's query, rotated and Hadamard-transformed, 8 heads a task + tasks += [ + ("index_q", INDEX_Q_ROWS // INDEX_TILE), + ("index_qh", samples * INDEX_HEADS // 8), + ] tasks += [("uk", heads * KV_LORA // UK_TILE)] if with_indexer: tasks += [ - ( - "index_score", - samples - * ((index_max_seq + INDEX_KEYS_PER_TASK - 1) // INDEX_KEYS_PER_TASK), - ), - ("index_hist", samples * INDEX_SEL_SLICES), - ("index_thresh", samples), - ("index_gather", samples * INDEX_SEL_SLICES), - ("index_select", samples), + # one task a CTA: its runs of the requests' keys + ("index_score", BLOCKS), + ("index_slice", samples * INDEX_SEL_SLICES), ] tasks += [ ( @@ -385,9 +387,12 @@ def stage_bases( # tile for the first half, instead of serializing two index-Q tiles on # every q_b CTA. base["index_q"] = (base["q_b"] + tasks["q_b"]) % BLOCKS - base["index_score"] = 101 - base["index_select"] = 100 - base["index_thresh"] = 100 - base["index_hist"] = BLOCKS // 2 - base["index_gather"] = BLOCKS // 2 + # on the first CTAs of index_q's own tiles, after their tile + base["index_qh"] = base["index_q"] + # each row's index key on its q_norm CTA, after the CTA's index-query tile + base["index_key"] = base["q_norm"] + # one scoring task a CTA, numbered by its CTA + base["index_score"] = 0 + # the selection's slices past uk's CTAs + base["index_slice"] = (base["uk"] + tasks["uk"]) % BLOCKS return base diff --git a/test/e2e/test_glm53_server.py b/test/e2e/test_glm53_server.py index 84c17fac..22e0ac14 100644 --- a/test/e2e/test_glm53_server.py +++ b/test/e2e/test_glm53_server.py @@ -157,11 +157,12 @@ def test_each_covered_verify_layer_within_its_numerics_with_fp8_kv(self): ) def test_each_covered_layer_past_the_parallel_top_k(self): - # rows of about 108k keys (27 chunks: 14 slices, 13 of two chunks) at widths 2 - # and 4, and of about 36k keys (9 one-chunk slices) at width 8. Past 98304 - # keys each wave's score walk refills its buffers and looks blocks ahead. + # rows of about 108k keys at widths 2 and 4 and of about 36k keys at width 8: + # past 98304 keys a row's top-k runs on every slice past = 98304 - self.assertGreaterEqual(past, layout.INDEX_SEL_SLICES * layout.INDEX_SEL_CHUNK) + self.assertGreaterEqual( + past, layout.INDEX_SEL_SLICES * layout.INDEX_SEL_SLICE_KEYS + ) self._check_layers( [], ((1, 9000), (3, 9000), (8, 3000)), diff --git a/test/kernels/glm53/check_layer.py b/test/kernels/glm53/check_layer.py index ed31e051..73ef7892 100644 --- a/test/kernels/glm53/check_layer.py +++ b/test/kernels/glm53/check_layer.py @@ -274,8 +274,9 @@ def make_step(positions: list, requests: list, pages: Pages, device): positions=torch.tensor([p or 0 for p in positions], dtype=i64, device=device), slots=torch.tensor(slots, dtype=i64, device=device), indptr=torch.tensor(indptr, dtype=i32, device=device), - req_ids=torch.tensor(requests, dtype=i32, device=device), - # SGLang's DSA backend gives each row its request's page-table row + # SGLang's DSA backend gives each row its request's page-table row, and the + # spec each row that row + req_ids=torch.arange(W, dtype=i32, device=device), page_table=pages.table[torch.tensor(requests, device=device)].contiguous(), topk=topk, ) diff --git a/test/kernels/glm53/check_topk.py b/test/kernels/glm53/check_topk.py new file mode 100644 index 00000000..4268c3c0 --- /dev/null +++ b/test/kernels/glm53/check_topk.py @@ -0,0 +1,315 @@ +"""Exact top-k check of the GLM-5.3 indexer launch at chosen contexts, under torchrun +(TP 4 or 8): + + torchrun --nproc-per-node 8 test/kernels/glm53/check_topk.py --pos 8447,131071 \\ + [--rows 2,4,8] [--caches random,ties,zeros,recent] [--rowsets serving,mixed,chain] + +Per index cache, position, row set and width, ``--layers`` launches of ``--kind`` +(``index``; ``dense``: the dense layers' launch; ``embed``: layer 0 from the token +ids, then dense launches) run ``--steps`` steps eagerly, with the cached scales of each step's own keys poisoned +first (``check.poison_new_index_scales``): a launch must score the keys its rows write +from those rows, not from the cache. After each launch, for every live row: + +- **scores** (``check.launch_scores``) against an f64 reference: the launch's index + query (its ``index_q`` mailbox, rotated and transformed in f32 as the kernel does, + then bf16), its head weights (``index_w``) and the row's keys in the index cache as + the launch leaves it. ``score_err`` is the largest error over the row's largest + reference score; +- **selection**: the row's list (``step.topk`` at its ``indptr``) against the exact + top-k of the launch's own scores, k = min(2048, keys). ``missed`` counts the keys of + the top-k (score descending, then position) it lacks; ``order`` holds when the list + is, entry by entry, the radix select's: the keys above the k-th key's 12-, 24- or + 32-bit prefix in position order, then the prefix's keys above the k-th, then its + ties in position order (the prefix is the first whose keys fit the build's + candidate room). + +Negative controls on copies of each row's result: a selected slot replaced by the +best unselected key (``missed``), two entries swapped (``order``) and one score's bits +flipped (``score_err``) must each fail. Index caches: ``random`` (check_layer's), +``ties`` (four keys repeated over every position, so the k-th score falls in a tie), +``zeros`` (every cached key 0) and ``recent`` (each live row's last 4,096 keys scaled +8x over check_layer's scale). Row sets: ``serving`` (time_layers.py's), ``mixed`` (one request a row, at +``--pos`` and below, some under 2,048 keys) and ``chain`` (one request's rows, as a +verify step of a draft chain as long as the width). ``--dump DIR``: every row's +scores and list as SHA-256 digests, in check_layer.py's format for its ``--compare``. +Prints a ``TOPK`` line per step and ``TOPK_RESULT exact=True`` when all pass.""" + +import argparse +import json +import os +import sys + +import torch +import torch.distributed as dist + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +import check_layer as cl # noqa: E402 +import time_layers as tl # noqa: E402 + +from sgl_amd_mono.glm53 import check # noqa: E402 +from sgl_amd_mono.glm53.gfx950.layout import ( # noqa: E402 + INDEX_DIM, + INDEX_HEADS, + INDEX_Q_ROWS, +) + +SCORE_BOUND = 1e-4 +CACHES = ("random", "ties", "zeros", "recent") + + +def _pairs(rt, name: str, n: int) -> torch.Tensor: + """The value words of a mailbox's first n (value, tag) pairs.""" + at = rt.sc[name] + return rt.scratch[at : at + n * 8].view(torch.int32)[0::2] + + +def _bf16_words(words: torch.Tensor) -> torch.Tensor: + """Words of two bf16 values (the first in the low half) -> f32 [.., 2 * n].""" + lo = (words << 16).view(torch.float32) + hi = (words & -65536).view(torch.float32) + return torch.stack([lo, hi], -1).flatten(-2) + + +def _hadamard(x: torch.Tensor) -> torch.Tensor: + """The kernel's orthonormal 128-point Hadamard transform of f32 rows, in f32: + butterflies of strides 1 to 64, then the scale.""" + x = x.clone() + n = x.shape[-1] + h = 1 + while h < n: + y = x.view(*x.shape[:-1], n // (2 * h), 2, h) + a, b = y[..., 0, :].clone(), y[..., 1, :].clone() + y[..., 0, :], y[..., 1, :] = a + b, a - b + h *= 2 + return x * torch.tensor(128**-0.5, dtype=torch.float32) + + +def query(rt, model, s: int, W: int) -> torch.Tensor: + """Row s's rotated, transformed index query as the launch scores it: f64 [32, 128].""" + q = _bf16_words(_pairs(rt, "index_q", W * INDEX_Q_ROWS // 2)) + q = q.view(W, INDEX_HEADS, INDEX_DIM)[s].clone() + pos = int(model.step_tables.positions[s]) + cos, sin = model.cos[pos].float(), model.sin[pos].float() + q0, q1 = q[:, 0:64:2].clone(), q[:, 1:64:2].clone() + q[:, 0:64:2] = q0 * cos - q1 * sin + q[:, 1:64:2] = q0 * sin + q1 * cos + return _hadamard(q).bfloat16().double() + + +def row_slots(st, s: int) -> torch.Tensor: + """The cache slots of row s's keys 0 .. its position, in key order.""" + bound = int(st.positions[s]) + 1 + pages = st.page_table[int(st.req_ids[s]), : -(-bound // cl.PAGE)].long() + keys = torch.arange(bound, device=pages.device) + return pages[keys // cl.PAGE] * cl.PAGE + keys % cl.PAGE + + +def reference_scores(rt, model, layer: int, s: int, W: int, slots) -> torch.Tensor: + entries = cl._index_entries(model.caches[layer].idx, slots) + keys = entries[:, :INDEX_DIM].contiguous().view(torch.float8_e4m3fn).double() + scales = entries[:, INDEX_DIM:].contiguous().view(torch.float32).double().flatten() + w = _pairs(rt, "index_w", W * INDEX_HEADS).view(torch.float32) + w = w.view(W, INDEX_HEADS)[s].double() + return (torch.relu(keys @ query(rt, model, s, W).T) @ w) * scales + + +def radix_lists(bits: torch.Tensor, k: int) -> list: + """The radix select's lists (key indices) of the top k of f32 score bits, at a + 12-, 24- and 32-bit prefix of the k-th key.""" + keys = check._radix_keys(bits) + kth = torch.sort(keys, descending=True).values[k - 1] + lists = [] + for bits_kept in (12, 24, 32): + shift = 32 - bits_kept + high = keys >> shift + above = (high > (kth >> shift)).nonzero().flatten() + under = ((high == (kth >> shift)) & (keys > kth)).nonzero().flatten() + ties = (keys == kth).nonzero().flatten()[: k - len(above) - len(under)] + lists.append(torch.cat([above, under, ties])) + return lists + + +def check_row(bits, ref, picked, slots) -> dict: + k = min(cl.TOPK, len(slots)) + keys = check._radix_keys(bits) + exact = slots[torch.sort(keys, descending=True, stable=True).indices[:k]] + out = { + "missed": len(set(exact.tolist()) - set(picked.tolist())), + "order": any( + torch.equal(slots[lst], picked.long()) for lst in radix_lists(bits, k) + ), + } + if ref is not None: + got = bits.view(torch.float32).double() + out["score_err"] = float( + (got - ref).abs().max() / ref.abs().max().clamp_min(1e-30) + ) + return out + + +def controls(bits, ref, picked, slots) -> bool: + """A wrong slot, two entries swapped and a flipped score bit each fail.""" + seen = True + if len(slots) > cl.TOPK: + order = torch.sort(check._radix_keys(bits), descending=True, stable=True) + bad = picked.clone() + bad[0] = slots[order.indices[cl.TOPK]] + seen &= check_row(bits, None, bad, slots)["missed"] >= 1 + if len(picked) > 1: + swapped = picked.clone() + swapped[0], swapped[1] = picked[1], picked[0] + seen &= not check_row(bits, None, swapped, slots)["order"] + flipped = bits.clone() + flipped[int(ref.abs().argmax())] ^= 1 << 22 + seen &= check_row(flipped, ref, picked, slots)["score_err"] > SCORE_BOUND + return seen + + +def fill_index_cache(model, mode: str) -> None: + """Index keys of ``mode`` (CACHES) in every layer's cache.""" + st = model.step_tables + for c in model.caches: + n = c.idx.shape[0] + values = c.idx.view(n, -1)[:, : cl.PAGE * INDEX_DIM] + if mode == "ties": + # [page, 16-token group, 16-dim tile, token, dim]: page 0's first four + # tokens repeated over every token + tiles = values.view(n, 4, 8, 16, 16) + tiles.copy_(tiles[:1, :1, :, :4, :].repeat(n, 4, 1, 4, 1)) + elif mode == "zeros": + values.zero_() + elif mode == "recent": + scales = c.idx.view(n, -1)[:, cl.PAGE * INDEX_DIM :].view(torch.float32) + for s in range(st.positions.shape[0]): + if int(st.slots[s]) >= 0: + slots = row_slots(st, s)[-4096:] + scales[slots // cl.PAGE, slots % cl.PAGE] = 2.0**-3 + else: + assert mode == "random", mode + + +def rows_of(rowset: str, width: int, pos: int): + if rowset == "serving": + return tl.serving_rows(width, pos) + if rowset == "chain": + return [pos + s for s in range(width)], [0] * width + assert rowset == "mixed", rowset + mixed = [pos, 1500, pos // 2, 300, pos - 7, 2047, 2048, 9000] + return [min(p, pos) for p in mixed[:width]], list(range(width)) + + +def run(args, model, name: str, digests: dict) -> list: + st = model.step_tables + W = st.positions.shape[0] + live = st.slots[st.slots >= 0] + records = [] + for step in range(args.steps): + for c in model.caches: + check.poison_new_index_scales(c.idx, live) + results = [] + + def after(i, _, step=step): + rt = model.ops[i].rt + bits = check.launch_scores(rt, W) + indptr = st.indptr.tolist() + d = {} + for s in range(W): + if int(st.slots[s]) < 0: + continue + slots = row_slots(st, s) + row_bits = bits[s, : len(slots)].clone() + picked = st.topk[indptr[s] : indptr[s + 1]].clone() + ref = reference_scores(rt, model, i, s, W, slots) + res = check_row(row_bits, ref, picked, slots) + res["controls"] = controls(row_bits, ref, picked, slots) + results.append(res) + d[f"scores.r{s}"] = cl._digest(row_bits) + d[f"list.r{s}"] = cl._digest(picked) + digests[f"{name}/s{step}/l{i}"] = d + + model.step(after=after) + torch.cuda.synchronize() + rec = { + "name": name, + "step": step, + "rows_checked": len(results), + "missed": max(r["missed"] for r in results), + "order": all(r["order"] for r in results), + "score_err": max(r["score_err"] for r in results), + "controls": all(r["controls"] for r in results), + "clean": model.errors() == 0, + } + rec["ok"] = ( + rec["missed"] == 0 + and rec["order"] + and rec["score_err"] <= SCORE_BOUND + and rec["controls"] + and rec["clean"] + ) + records.append(rec) + return records + + +def main() -> int: + ap = argparse.ArgumentParser() + ap.add_argument("--pos", default="8447", help="comma-separated first positions") + ap.add_argument("--rows", default="2,4,8", help="comma-separated widths") + ap.add_argument("--caches", default=",".join(CACHES)) + ap.add_argument("--rowsets", default="serving,mixed,chain") + ap.add_argument("--kind", choices=("index", "dense", "embed"), default="index") + ap.add_argument("--layers", type=int, default=2, help="launches a step") + ap.add_argument("--steps", type=int, default=2) + ap.add_argument("--kv", choices=("bf16", "fp8"), default="fp8") + ap.add_argument("--deadline", type=float, default=60.0) + ap.add_argument("--dump", default="", help="write the digests into DIR") + args = ap.parse_args() + rank, tp = int(os.environ["RANK"]), int(os.environ["WORLD_SIZE"]) + torch.cuda.set_device(int(os.environ["LOCAL_RANK"])) + device = torch.device("cuda", torch.cuda.current_device()) + dist.init_process_group("gloo") + group = dist.new_group(backend="gloo") + ok, digests = True, {} + for mode in args.caches.split(","): + for pos in map(int, args.pos.split(",")): + for rowset in args.rowsets.split(","): + for W in map(int, args.rows.split(",")): + positions, requests = rows_of(rowset, W, pos) + model = cl.Model( + cl.kinds_of(args.kind, args.layers), + positions, + requests, + rank, + tp, + group, + device, + args, + ) + fill_index_cache(model, mode) + name = f"{args.kv}-{args.kind}-{mode}/p{pos}/{rowset}/w{W}" + for rec in run(args, model, name, digests): + ok &= rec["ok"] + if rank == 0: + print( + "TOPK " + " ".join(f"{k}={v}" for k, v in rec.items()), + flush=True, + ) + model.close() + del model + torch.cuda.empty_cache() + if args.dump: + os.makedirs(args.dump, exist_ok=True) + name = f"digests_topk_{args.kind}_{args.kv}_r{rank}.json" + path = os.path.join(args.dump, name) + with open(path, "w") as f: + json.dump(digests, f, indent=0, sort_keys=True) + oks = [None] * tp + dist.all_gather_object(oks, ok, group=group) + if rank == 0: + print(f"TOPK_RESULT tp={tp} exact={all(oks)}", flush=True) + dist.barrier(group=group) + dist.destroy_process_group() + return 0 if all(oks) else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/test/kernels/test_glm53_kernels.py b/test/kernels/test_glm53_kernels.py index 9ebe7861..19d1b9c8 100644 --- a/test/kernels/test_glm53_kernels.py +++ b/test/kernels/test_glm53_kernels.py @@ -8,8 +8,11 @@ selected token of a row past 2048 keys is a pair tagged with its indexer launch's tag that holds the slot the launch publishes; a dense layer's MLP and the tail's logits are within the spec's bounds against f32 references, and fail them with a -scale or the final norm off; and layer 0 from the token ids equals the dense launch -fed the gathered embedding rows.""" +scale or the final norm off; layer 0 from the token ids equals the dense launch fed +the gathered embedding rows; and every indexer row's list is the exact top-k of the +launch's scores in the selection's order (``glm53/check_topk.py``), on random, tied, +zero and recency-heavy index caches, at about 8K and 300K keys (and for layer 0 and +a dense layer at 8K and 128K).""" import subprocess import sys @@ -29,13 +32,14 @@ def _gfx950_gpus() -> int: CHECK = Path(__file__).with_name("glm53") / "check_layer.py" +TOPK = Path(__file__).with_name("glm53") / "check_topk.py" # the widths the spec builds at each TP size (TP4 runs at most 4 rows) WIDTHS = {4: "2,4", 8: "2,4,8"} @unittest.skipIf(_gfx950_gpus() < 8, "needs 8 gfx950 GPUs") class TestMonoGlm53Kernels(CustomTestCase): - def _check(self, *args: str) -> None: + def _check(self, *args: str, script: Path = CHECK) -> None: for tp, rows in WIDTHS.items(): cmd = [ sys.executable, @@ -43,7 +47,7 @@ def _check(self, *args: str) -> None: "torch.distributed.run", f"--nproc-per-node={tp}", f"--master-port={29630 + tp}", - str(CHECK), + str(script), "--rows", rows, *args, @@ -94,6 +98,11 @@ def test_tail_logits_within_their_bounds(self): def test_selected_tokens_carry_the_launch_tag(self): self._check("--kind", "index", "--check-tags", "--rowsets", "p0,v4") + def test_indexer_selects_the_exact_top_k(self): + self._check("--pos", "8447,300000", script=TOPK) + # layer 0 from the tokens and a dense layer: the same indexer stages + self._check("--kind", "embed", "--pos", "8447,131071", script=TOPK) + if __name__ == "__main__": unittest.main() diff --git a/test/unit/test_glm53_select.py b/test/unit/test_glm53_select.py index 4f358aa1..5efcbbd6 100644 --- a/test/unit/test_glm53_select.py +++ b/test/unit/test_glm53_select.py @@ -1,9 +1,9 @@ -"""The GLM-5.3 gfx950 parallel top-k plan on CPU (RFC #43949): a row past -INDEX_SEL_PAR_CHUNKS chunks selects on INDEX_SEL_SLICES CTAs, each a run of whole -chunks, and the kernel calls the same ``select_slices`` on traced values. On Python -ints: the slices tile a row's chunks once in position order, and each width's slice -tasks fit the half of the grid they run on, after the score stage and before the -stages that read the selection.""" +"""The GLM-5.3 gfx950 sliced top-k plan on CPU (RFC #43949): a row's top-k runs on up +to INDEX_SEL_SLICES slices, runs of its keys in position order, and the kernel calls +the same ``select_slices`` on traced values. On Python ints: the slices tile a row's +keys once, none empty, each within a slice's LDS keys and its 16-bit counts; the slice +tasks of every width are whole rows of the grid, and the stages run in the order the +selection's hand-offs need.""" import pytest @@ -11,36 +11,40 @@ from sgl_amd_mono.glm53.gfx950 import layout as L from sgl_amd_mono.glm53.spec import WIDTHS -PAR = L.INDEX_SEL_PAR_CHUNKS * L.INDEX_SEL_CHUNK -BOUNDS = [PAR + 1, 2 * PAR, 2 * PAR + 1, 65536, 65537, 131072, 262144 + 48, - (1 << 20) - 1, 1 << 20] # fmt: skip +MAX_SEQ = 1 << 20 +KEYS = L.INDEX_SEL_SLICE_KEYS +BOUNDS = [1, 2, 2047, 2048, 2049, KEYS, KEYS + 1, 2 * KEYS, 2 * KEYS + 1, 65535, 65536, + 65537, 131072, 262144 + 48, MAX_SEQ - 1, MAX_SEQ] # fmt: skip -@pytest.mark.parametrize("bound", BOUNDS + list(range(PAR + 1, 1 << 20, 7919))) -def test_slices_tile_a_rows_chunks_once_in_position_order(bound): +@pytest.mark.parametrize("bound", BOUNDS + list(range(3, MAX_SEQ, 7919))) +def test_slices_tile_a_rows_keys_once_in_position_order(bound): n, per = L.select_slices(bound) - chunks = -(-bound // L.INDEX_SEL_CHUNK) - assert 1 <= n <= L.INDEX_SEL_SLICES and per >= 1 - # slice k takes chunks [k * per, min((k + 1) * per, chunks)): each nonempty, all taken - assert (n - 1) * per < chunks <= n * per + assert 1 <= n <= L.INDEX_SEL_SLICES + # slice c takes keys [c * per, min((c + 1) * per, bound)): each nonempty, all taken + assert (n - 1) * per < bound <= n * per + # a slice per INDEX_SEL_SLICE_KEYS keys until the most slices + assert n == min(-(-bound // KEYS), L.INDEX_SEL_SLICES) + assert n == L.INDEX_SEL_SLICES or per <= KEYS + # a slice's keys fit its share of the longest row, and its counts 16 bits + assert per <= -(-MAX_SEQ // L.INDEX_SEL_SLICES) < 1 << 16 + + +def test_an_empty_row_has_one_slice(): + assert L.select_slices(0) == (1, 0) @pytest.mark.parametrize("rows", WIDTHS) -def test_slice_tasks_fit_half_the_grid_between_score_and_split(rows): - stages = L.stage_tasks(rows, 8, 2048, with_indexer=True, index_max_seq=1 << 20) +def test_slice_tasks_are_whole_rows_of_the_grid_between_score_and_split(rows): + stages = L.stage_tasks(rows, 8, 2048, with_indexer=True, index_max_seq=MAX_SEQ) tasks = dict(stages) - assert tasks["index_hist"] == tasks["index_gather"] == rows * L.INDEX_SEL_SLICES - assert rows * L.INDEX_SEL_SLICES <= BLOCKS // 2 - assert tasks["index_thresh"] == tasks["index_select"] == rows + assert tasks["index_slice"] == rows * L.INDEX_SEL_SLICES + assert BLOCKS % L.INDEX_SEL_SLICES == 0 + assert tasks["index_score"] == BLOCKS + assert tasks["index_key"] == rows order = [name for name, _ in stages] - chain = [ - "index_score", - "index_hist", - "index_thresh", - "index_gather", - "index_select", - "split", - ] + chain = ["q_norm", "index_key", "index_q", "index_qh", "index_score", + "index_slice", "split"] # fmt: skip assert [order.index(name) for name in chain] == sorted( order.index(name) for name in chain )