diff --git a/aiter/ops/flydsl/kernels/qk_norm_rope_quant.py b/aiter/ops/flydsl/kernels/qk_norm_rope_quant.py index 5f13e92c22..c2befff9ec 100644 --- a/aiter/ops/flydsl/kernels/qk_norm_rope_quant.py +++ b/aiter/ops/flydsl/kernels/qk_norm_rope_quant.py @@ -99,6 +99,9 @@ def _cached_from_dlpack(t: torch.Tensor): SMALL_T_THRESHOLD = 96 USE_TDM_PREFILL = True TDM_MIN_ROWS = 32768 +# gfx1250 cache-policy bit 0 reduces the T=512 BF16 output-store drain; larger +# T already reaches steady-state HBM bandwidth and keeps the default policy. +TDM_DECODE_STORE_CACHE_MODIFIER = 1 _SQRT2 = math.sqrt(2.0) @@ -193,7 +196,7 @@ def _store_bf16_tiled(vals_list, p_dst, copy, vec): fx.copy(copy, frag, p_dst) -def _store_bf16_vec(vals, out_rsrc, row_base_bytes, idx, vec): +def _store_bf16_vec(vals, out_rsrc, row_base_bytes, idx, vec, cache_modifier=0): """Convert VEC fp32 → bf16 dwords and raw buffer_store (split for VEC>8).""" if isinstance(vals, (list, tuple)): vals = fx.Vector.from_elements(vals, dtype=fx.Float32) @@ -205,7 +208,11 @@ def _store_bf16_vec(vals, out_rsrc, row_base_bytes, idx, vec): if const_expr(dwords <= 4): buffer_ops.buffer_store( - bf16_as_i32.ir_value(), out_rsrc, off_bytes, offset_is_bytes=True + bf16_as_i32.ir_value(), + out_rsrc, + off_bytes, + cache_modifier=cache_modifier, + offset_is_bytes=True, ) else: lo = fx.Vector.from_elements([bf16_as_i32[i] for i in range(4)], dtype=fx.Int32) @@ -213,10 +220,18 @@ def _store_bf16_vec(vals, out_rsrc, row_base_bytes, idx, vec): [bf16_as_i32[i] for i in range(4, dwords)], dtype=fx.Int32 ) buffer_ops.buffer_store( - lo.ir_value(), out_rsrc, off_bytes, offset_is_bytes=True + lo.ir_value(), + out_rsrc, + off_bytes, + cache_modifier=cache_modifier, + offset_is_bytes=True, ) buffer_ops.buffer_store( - hi.ir_value(), out_rsrc, off_bytes + 16, offset_is_bytes=True + hi.ir_value(), + out_rsrc, + off_bytes + 16, + cache_modifier=cache_modifier, + offset_is_bytes=True, ) @@ -1593,24 +1608,37 @@ def flydsl_qk_norm_rope_quant( bid_arg = q.new_empty(1, dtype=torch.int32) if is_gfx1250: + tdm_quant = quant use_tdm = ( USE_TDM_PREFILL - and not quant - and not kv_write - and not paged + and (not quant or T_tok * H >= 65536) and T_tok * H >= TDM_MIN_ROWS and (D % 256 == 0) and (H & (H - 1) == 0) ) if use_tdm: num_rows = T_tok * H - tiles_per_wg = _tdm_tiles_per_wg(num_rows) + tiles_per_wg, num_buffers, rows_per_tile = _tdm_tiles_per_wg( + num_rows, head_dim=D + ) launcher = _build_kernel_w32_tdm( num_q_heads=H, head_dim=D, rope_head_dim=RD, + num_buffers=num_buffers, tiles_per_wg=tiles_per_wg, + rows_per_tile=rows_per_tile, q_weighted=q_weighted, + store_cache_modifier=( + TDM_DECODE_STORE_CACHE_MODIFIER + if (T_tok, H, D, RD) == (512, 128, 512, 64) + else 0 + ), + quant=tdm_quant, + group_size=G, + scale_dtype=scale_dtype, + kv_write=kv_write, + paged=paged, ) if stream is None: stream = torch.cuda.current_stream() @@ -1624,27 +1652,41 @@ def _t_ptr(t): ) q_2d = q_view.reshape(num_rows, D) - kv_c = kv if kv.is_contiguous() else kv.contiguous() - per_wg = ROWS_PER_TILE * tiles_per_wg + per_wg = rows_per_tile * tiles_per_wg args = ( # q is a per-call activation: uncached, else the adaptor cache # pins its allocation (see _cached_from_dlpack). flyc.from_dlpack(q_2d), - _t_ptr(kv_c), + _t_ptr(kv), _cached_from_dlpack(cos_2d), _cached_from_dlpack(sin_2d), _t_ptr(positions), _t_ptr(q_out.view(num_rows, D)), _t_ptr(kv_out), + _t_ptr(q_scale_arg.view(-1)), + _t_ptr(kv_scale_arg.view(-1)), _cached_from_dlpack(q_weight_arg.reshape(-1)), _cached_from_dlpack(kv_weight.reshape(-1)), + _t_ptr(swa_kv_arg), + _t_ptr(ssm_arg), + _t_ptr(bid_arg), num_rows, T_tok, (num_rows + per_wg - 1) // per_wg, + kv.stride(0), + swa_slot_stride, + swa_pos_stride, + swa_num_rows, + swa_cache_size, stream if has_direct else Stream(stream), ) _run_compiled(launcher, *args) - return q_out, kv_out, None, None + return ( + q_out, + kv_out, + (q_scale if quant else None), + (kv_scale if quant else None), + ) rows_per_wg = ROWS_PER_WG_SMALL if T_tok <= SMALL_T_THRESHOLD else ROWS_PER_WG launcher = _build_kernel_w32( @@ -1736,21 +1778,28 @@ def _stream_arg(): return q_out, kv_out, (q_scale if quant else None), (kv_scale if quant else None) -ROWS_PER_TILE = 32 # = waves per WG (1024-thread WG) -NUM_BUFFERS = 6 # K: rotating LDS buffers (1 WG at K=6: 6*32KB=192KB LDS) -TILES_PER_WG = 40 # CT: tiles streamed per WG (unrolled; CT % (H/RT) == 0) +ROWS_PER_TILE = 8 # = waves per WG (256-thread WG) +ROWS_PER_TILE_SPARSE = 4 # halve the WG so a thin grid still gets 2 WGs/CU +NUM_BUFFERS = 6 # K: rotating LDS buffers (1 WG at K=6: 6*8KB=48KB LDS) +NUM_BUFFERS_DENSE = 2 # K once the grid is deep enough to hide a short prologue +TILES_PER_WG = 16 # CT: one H=128 token/WG; keeps cos/sin hoisting enabled +TDM_DENSE_MIN_ROWS = 131072 # num_rows from which K=2 beats K=6 +TDM_RT8_MIN_ROWS = 65536 # num_rows where RT=8 first reaches 2 WGs per 256 CUs -def _tdm_tiles_per_wg(num_rows): - """Pick CT (tiles per workgroup) for a row count — trades prefetch depth vs occupancy.""" - if num_rows <= 65536: - return 8 - if num_rows <= 131072: - return 16 - return TILES_PER_WG +def _tdm_tiles_per_wg(num_rows, head_dim=512): + """Pick (CT, K, RT). Both knobs ask the same thing: can a workgroup's load + latency be hidden by a neighbour, or must it cover its own? Thin grids need + a smaller WG (more of them) and a deeper prologue; dense grids need neither. + Thresholds assume a 256-CU part. + """ + del head_dim # tile bytes are implied by RT * D in the kernel + k = NUM_BUFFERS_DENSE if num_rows >= TDM_DENSE_MIN_ROWS else NUM_BUFFERS + rt = ROWS_PER_TILE if num_rows >= TDM_RT8_MIN_ROWS else ROWS_PER_TILE_SPARSE + return TILES_PER_WG, k, rt -@lru_cache(maxsize=16) +@lru_cache(maxsize=32) def _build_kernel_w32_tdm( *, num_q_heads, @@ -1758,8 +1807,15 @@ def _build_kernel_w32_tdm( rope_head_dim, num_buffers=NUM_BUFFERS, tiles_per_wg=TILES_PER_WG, + rows_per_tile=ROWS_PER_TILE, hoist_cs=True, q_weighted=False, + store_cache_modifier=0, + quant=False, + group_size=None, + scale_dtype="fp32", + kv_write=False, + paged=False, ): H, D, RD = num_q_heads, head_dim, rope_head_dim NOPE = D - RD @@ -1769,37 +1825,67 @@ def _build_kernel_w32_tdm( PAIRS = VEC // 2 K = num_buffers CT = tiles_per_wg - RT = ROWS_PER_TILE + RT = rows_per_tile + G = D if group_size is None else group_size + NG = D // G + TPG = G // VEC + log2_tpg = int(math.log2(TPG)) + is_e8m0 = scale_dtype == "e8m0" + if CT < K: + raise ValueError( + f"TDM prologue issues K={K} tiles but WG only owns CT={CT}; " + "need CT >= K to avoid overlapping the next WG's rows" + ) eb = 2 log2b = int(math.log2(BLOCK_THREADS)) + amax_start_step = log2b - log2_tpg log2H = int(math.log2(H)) if (H & (H - 1)) == 0 else None + fp8_mx_dtype = _D.FP8_E4M3 EVEN = list(range(0, VEC, 2)) ODD = list(range(1, VEC, 2)) ILV = [(i // 2) if i % 2 == 0 else (PAIRS + i // 2) for i in range(VEC)] tile_bytes = RT * D * eb ARENA = K * tile_bytes # K rotating tile buffers - name = f"qk_norm_rope_tdm_H{H}_D{D}_RD{RD}_k{K}_ct{CT}" + name = f"qk_norm_rope_tdm_H{H}_D{D}_RD{RD}_k{K}_ct{CT}_r{RT}" if hoist_cs: name += "_h" if q_weighted: name += "_qw" + if store_cache_modifier: + name += f"_sc{store_cache_modifier}" + if quant: + name += f"_q8g{G}_{scale_dtype}" + if kv_write: + name += "_kvw" + if paged: + name += "_paged" name += "_fused_w32_flydsl" @flyc.kernel(name=name, known_block_size=[BLOCK_THREADS, RT, 1]) def kernel( q_in: fx.Tensor, # [T*H, D] bf16 - kv_in: fx.Pointer, # [T, D] bf16 (contiguous) + kv_in: fx.Pointer, # [T, D] bf16, may be strided cos_cache: fx.Tensor, sin_cache: fx.Tensor, positions: fx.Pointer, # [T] i64 q_out: fx.Pointer, # [T*H, D] bf16 kv_out: fx.Pointer, # [T, D] bf16 + q_scale: fx.Pointer, # [T*H, NG] fp32/u8 (dummy when not quant) + kv_scale: fx.Pointer, # [T, NG] fp32/u8 (dummy when not quant) q_weight: fx.Tensor, # [D] bf16 (dummy when not q_weighted) kv_weight: fx.Tensor, # [D] bf16 + swa_kv: fx.Pointer, # [num_slots, cache_size, D] bf16 (dummy if not kv_write) + swa_index: fx.Pointer, # paged: block_tables; direct: [T] dest rows + batch_id_per_token: fx.Pointer, # [T] i32, -1 sentinel (dummy if not kv_write) num_rows: Int32, # T*H num_tokens: Int32, # T gx_q: Int32, # workgroups owned by the Q path + kv_in_row_stride: Int32, # KV row stride in bf16 elements + swa_slot_stride: Int32, # bf16 elements (= cache_size * D) + swa_pos_stride: Int32, # bf16 elements (= D) + swa_num_rows: Int32, # rows in the SWA pool; bounds the final row + swa_cache_size: Int32, # ring slot count ): i32 = T.i32 tid = fx.thread_idx.x @@ -1852,13 +1938,23 @@ def _load_tensor_vec(tensor, tid_val): ] ) - def load_cs(tok): - """cos/sin for this token, as PAIRS-wide f32 vectors.""" - pos_i32 = fx.Int32( + def load_pos(tok): + return fx.Int32( buffer_ops.buffer_load(pos_rsrc, tok, vec_width=1, dtype=T.i64).trunci( i32 ) ) + + def load_cs(tok): + """cos/sin for this token, as PAIRS-wide f32 vectors.""" + return _cs_from_pos(load_pos(tok)) + + def issue_pos(tok): + """Fire position buffer_load (returns raw i64, no wait).""" + return buffer_ops.buffer_load(pos_rsrc, tok, vec_width=1, dtype=T.i64) + + def _cs_from_pos(pos_i32): + """Load cos/sin given an already-resolved position value.""" cs = pos_i32 * (RD // 2) + rope_rel * PAIRS if const_expr(PAIRS == 1): return tuple( @@ -1875,27 +1971,122 @@ def load_cs(tok): for r in (cos_rsrc, sin_rsrc) ) - def norm_rope_store(xv, wv, cosv, sinv, out_rsrc, row_base_bytes): - """RMSNorm (+ optional per-channel weight) + GPT-J RoPE + bf16 store. - rstd comes from the unweighted x^2 (stock-kernel semantics); the - weight multiply commutes into `scaled`.""" + def norm_rope_store( + xv, + wv, + cosv, + sinv, + out_rsrc, + row_base_bytes, + scale_rsrc=None, + scale_idx=None, + swa_rsrc=None, + swa_row_base=None, + do_swa=None, + ): + """RMSNorm + RoPE and optional per-row FP8 quantization.""" sq = (xv * xv).reduce(ReductionOp.ADD) + if const_expr(quant): + xw = xv * wv if wv is not None else xv + am = fmath.absf(xw).reduce(ReductionOp.MAX) for sh in range_constexpr(log2b): o = BLOCK_THREADS // (2 << sh) sq = sq + sq.shuffle_xor(o, BLOCK_THREADS) - rstd = fmath.rsqrt(sq * (1.0 / D) + 1e-6) - scaled = (xv * wv * rstd) if wv is not None else (xv * rstd) - ev = scaled.shuffle(scaled, EVEN) - od = scaled.shuffle(scaled, ODD) - rot = (ev * cosv - od * sinv).shuffle(ev * sinv + od * cosv, ILV) - final = is_rope.select(rot, scaled) - _store_bf16_vec(final, out_rsrc, row_base_bytes, tid, VEC) + if const_expr(quant and sh >= amax_start_step): + am = am.maximumf(am.shuffle_xor(o, BLOCK_THREADS)) + if const_expr(quant and is_e8m0): + rstd_q = fmath.rsqrt(sq * (1.0 / D) + 1e-6) + e8m0_biased = fx.Int32( + emit_mx_e8m0_scale( + (am.maximumf(fx.Float32(1e-12)) * rstd_q * _SQRT2).ir_value(), + mode=_DEFAULT_MODE, + dtype=fp8_mx_dtype, + ) + ) + factor = rstd_q * ((fx.Int32(254) - e8m0_biased) << 23).bitcast( + fx.Float32 + ) + scaled = xw * factor + scale_store = e8m0_biased.to(fx.Int8) + elif const_expr(quant): + am_safe = am.maximumf(fx.Float32(1e-12)) + rcp_am = llvm.call_intrinsic( + T.f32, "llvm.amdgcn.rcp.f32", [am_safe.ir_value()], [], [] + ) + fc = _fp8_const() + factor = fx.Float32(fc["max_over_sqrt2"]) * rcp_am + scaled = xw * factor + else: + rstd = fmath.rsqrt(sq * (1.0 / D) + 1e-6) + scaled = (xv * wv * rstd) if wv is not None else (xv * rstd) + + # Only RD / D lanes own the RoPE tail (4 of 32 for D=512, + # RD=64). Keep the passthrough value in register memory so an + # scf.if can update it without an SSA dominance violation, and + # avoid running all RoPE FMAs in the other 28 lanes. + out_rmem = fx.make_rmem_tensor(fx.make_layout(VEC, 1), fx.Float32) + fx.memref_store_vec(scaled, out_rmem) + if is_rope: + cur = fx.memref_load_vec(out_rmem) + ev = cur.shuffle(cur, EVEN) + od = cur.shuffle(cur, ODD) + rot = (ev * cosv - od * sinv).shuffle(ev * sinv + od * cosv, ILV) + fx.memref_store_vec(rot, out_rmem) + final = fx.memref_load_vec(out_rmem) + if const_expr(quant): + _store_fp8_packed_w32( + [final[i] for i in range_constexpr(VEC)], + out_rsrc, + row_base_bytes, + tid, + VEC, + skip_fnuz_clamp=True, + ) + group_idx = tid >> log2_tpg + if (tid & (TPG - 1)) == 0: + if const_expr(is_e8m0): + buffer_ops.buffer_store( + scale_store, scale_rsrc, scale_idx + group_idx + ) + else: + # FP8 values depend only on amax. Defer the RMS scale + # path until after the wide output store so rsqrt + # latency overlaps memory retirement. + buffer_ops.buffer_store( + am_safe + * fmath.rsqrt(sq * (1.0 / D) + 1e-6) + * fx.Float32(fc["inv_max_sqrt2"]), + scale_rsrc, + scale_idx + group_idx, + ) + else: + _store_bf16_vec( + final, + out_rsrc, + row_base_bytes, + tid, + VEC, + cache_modifier=store_cache_modifier, + ) + # Same bytes as kv_out, scattered into the SWA ring. + if const_expr(kv_write) and do_swa: + _store_bf16_vec(final, swa_rsrc, swa_row_base, tid, VEC) def emit_q(): - q_out_rsrc = _ptr_res(q_out) + q_scale_rsrc = _ptr_res(q_scale) if const_expr(quant) else None nr_m1 = num_rows - 1 tile_base = g * CT # first tile index this WG owns - wave_lds_off = wave * (D * eb) # this wave's slot in a tile buffer + wg_row0 = tile_base * RT # first row this WG owns + out_row_bytes = D if const_expr(quant) else D * 2 + # num_records is 32-bit, so one descriptor only reaches 4 GiB; bias + # the base per WG so the offset spans CT*RT rows, not the tensor. + q_out_rsrc = GTensor( + q_out, + dtype=T.bf16, + shape=(CT * RT, D), + static_bytes_offset_i64=fx.Int64(wg_row0) * fx.Int64(out_row_bytes), + ).rsrc + wave_lds_off = wave * (D * eb) row_elem = tid * VEC # LDS read pos within that slot qw = _load_tensor_vec(q_weight, tid) if const_expr(q_weighted) else None @@ -1904,7 +2095,10 @@ def emit_q(): fx.make_view(fx.get_iter(q_in), fx.make_layout((1, D), (D, 1))) ) tdm_atom = fx.rocdl.make_tdm_atom( - q_row0, [num_rows, None], strides=[D, None], num_warps=1 + q_row0, + [num_rows, None], + strides=[D, None], + num_warps=1, ) def row_of(tile_idx): @@ -1913,7 +2107,7 @@ def row_of(tile_idx): def issue(buf, tile_idx): dst = fx.Tensor( fx.make_view( - fx.inttoptr(lds_bf16_ptr_ty, lds_i32 + (wave_lds_off + buf)), + fx.inttoptr(lds_bf16_ptr_ty, lds_i32 + wave_lds_off + buf), fx.make_layout((1, D), (D, 1)), ) ) @@ -1942,36 +2136,58 @@ def compute_store(buf, tile_idx, cosv, sinv): cosv, sinv, q_out_rsrc, - row_of(tile_idx) * (D * 2), + (row_of(tile_idx) - wg_row0) * out_row_bytes, + q_scale_rsrc, + row_of(tile_idx) * NG, ) - def fence(): - tdm_ops.tensor_wait(K - 1) - rocdl.s_wait_dscnt(0) - - def cs_of(tile_idx): + def tok_of(tile_idx): my_row = row_of(tile_idx) if const_expr(log2H is not None): - return load_cs(my_row >> log2H) - return load_cs(_idiv(my_row, fx.Int32(H))) + return my_row >> log2H + return _idiv(my_row, fx.Int32(H)) + + def cs_of(tile_idx): + return load_cs(tok_of(tile_idx)) GROUP = (H // RT) if (log2H is not None and H % RT == 0) else 1 - do_hoist = hoist_cs and GROUP > 1 and (CT % GROUP == 0) + # A WG may cover one or more complete tokens, or an aligned fraction + # of one token. Both cases keep cos/sin constant within the WG. + do_hoist = hoist_cs and GROUP > 1 and (CT % GROUP == 0 or GROUP % CT == 0) for k in range_constexpr(K): issue(bufs[k], tile_base + k) + cs_cache = [None, None] + if const_expr(do_hoist): + pending_pos = [issue_pos(tok_of(tile_base + 0))] + cs_cache[0], cs_cache[1] = _cs_from_pos( + fx.Int32(pending_pos[0].trunci(i32)) + ) for i in range_constexpr(CT): - fence() + # Tile i consumes TDM load #i (issued in tile order: K in the + # prologue, then one per iteration while i + K < CT). In steady + # state K + i are issued, so leaving K-1 outstanding retires + # exactly #0..#i. Once the issues stop the issued count freezes + # at CT and K-1 is too loose: #i is only guaranteed retired with + # at most CT-1-i left, reaching 0 on the last tile. Using K-1 + # throughout lets a drain tile read an LDS buffer whose load has + # not landed; it stays latent only while the per-tile compute + # happens to outlast the load. + tdm_ops.tensor_wait(min(K - 1, CT - 1 - i)) + rocdl.s_wait_dscnt(0) if const_expr(do_hoist): - if const_expr(i % GROUP == 0): - cs_cache[0], cs_cache[1] = cs_of(tile_base + i) cosv, sinv = cs_cache[0], cs_cache[1] else: cosv, sinv = cs_of(tile_base + i) compute_store(bufs[i % K], tile_base + i, cosv, sinv) if const_expr(i + K < CT): issue(bufs[i % K], tile_base + i + K) # reuse after read + if const_expr(do_hoist and (i + 1) % GROUP == 0 and i + 1 < CT): + pending_pos[0] = issue_pos(tok_of(tile_base + i + 1)) + cs_cache[0], cs_cache[1] = _cs_from_pos( + fx.Int32(pending_pos[0].trunci(i32)) + ) def emit_kv(): gk = g - gx_q @@ -1982,7 +2198,7 @@ def emit_kv(): fx.Vector( buffer_ops.buffer_load( kv_rsrc, - tok * D + tid * VEC + (c * 8), + tok * kv_in_row_stride + tid * VEC + (c * 8), vec_width=8, dtype=T.bf16, ) @@ -1991,8 +2207,72 @@ def emit_kv(): ] ) wv = _load_tensor_vec(kv_weight, tid) - cosv, sinv = load_cs(tok) - norm_rope_store(xv, wv, cosv, sinv, _ptr_res(kv_out), tok * (D * 2)) + pos_i32 = load_pos(tok) + cosv, sinv = _cs_from_pos(pos_i32) + + swa_rsrc = None + swa_row_base = None + do_swa = None + if const_expr(kv_write): + # Gates mirror the wave32 path / C++ sibling exactly. + bid_i32 = fx.Int32( + buffer_ops.buffer_load( + _ptr_res(batch_id_per_token), tok, vec_width=1, dtype=i32 + ) + ) + pos_ok = pos_i32 >= 0 + do_swa = (bid_i32 >= 0) & pos_ok + bid_safe = _imax(bid_i32, fx.Int32(0)) + pos_safe = pos_ok.select(pos_i32, fx.Int32(0)) + if const_expr(paged): + blk = _idiv(pos_safe, swa_cache_size) + blk_ok = blk < swa_slot_stride + blk_safe = blk_ok.select(blk, fx.Int32(0)) + phys = fx.Int32( + buffer_ops.buffer_load( + _ptr_res(swa_index), + bid_safe * swa_slot_stride + blk_safe, + vec_width=1, + dtype=i32, + ) + ) + row = phys * swa_cache_size + (pos_safe % swa_cache_size) + row_ok = blk_ok & (phys >= 0) & (row < swa_num_rows) + do_swa = do_swa & row_ok + row_safe = row_ok.select(row, fx.Int32(0)) + else: + row = fx.Int32( + buffer_ops.buffer_load( + _ptr_res(swa_index), tok, vec_width=1, dtype=i32 + ) + ) + dest_ok = (row >= 0) & (row < swa_num_rows) + do_swa = do_swa & dest_ok + row_safe = dest_ok.select(row, fx.Int32(0)) + # row fits 32 bits, row*D*2 does not: widen before the multiply. + swa_rsrc = GTensor( + swa_kv, + dtype=T.bf16, + shape=(D,), + static_bytes_offset_i64=( + fx.Int64(row_safe) * fx.Int64(swa_pos_stride) * fx.Int64(2) + ), + ).rsrc + swa_row_base = fx.Int32(0) + + norm_rope_store( + xv, + wv, + cosv, + sinv, + _ptr_res(kv_out), + tok * (D if const_expr(quant) else D * 2), + _ptr_res(kv_scale) if const_expr(quant) else None, + tok * NG, + swa_rsrc, + swa_row_base, + do_swa, + ) if g < gx_q: emit_q() @@ -2008,11 +2288,21 @@ def launch( positions: fx.Pointer, q_out: fx.Pointer, kv_out: fx.Pointer, + q_scale: fx.Pointer, + kv_scale: fx.Pointer, q_weight: fx.Tensor, kv_weight: fx.Tensor, + swa_kv: fx.Pointer, + swa_index: fx.Pointer, + batch_id_per_token: fx.Pointer, num_rows: fx.Int32, num_tokens: fx.Int32, gx_q: fx.Int32, + kv_in_row_stride: fx.Int32, + swa_slot_stride: fx.Int32, + swa_pos_stride: fx.Int32, + swa_num_rows: fx.Int32, + swa_cache_size: fx.Int32, stream: fx.Stream = fx.Stream(None), # noqa: B008 ): gx_kv = _idiv(num_tokens + (RT - 1), fx.Int32(RT)) @@ -2025,11 +2315,21 @@ def launch( positions, q_out, kv_out, + q_scale, + kv_scale, q_weight, kv_weight, + swa_kv, + swa_index, + batch_id_per_token, num_rows, num_tokens, gx_q, + kv_in_row_stride, + swa_slot_stride, + swa_pos_stride, + swa_num_rows, + swa_cache_size, ) k.launch( grid=(fx.Index(gx), 1, 1), diff --git a/op_tests/test_flydsl_qk_norm_rope_quant.py b/op_tests/test_flydsl_qk_norm_rope_quant.py index a7f7eba0d8..ff354f1cfb 100755 --- a/op_tests/test_flydsl_qk_norm_rope_quant.py +++ b/op_tests/test_flydsl_qk_norm_rope_quant.py @@ -186,8 +186,21 @@ def _dequant(out_fp8, scale, *, D, quant_group_size, scale_dtype): # Main test (per-config) # ============================================================================ -# MI355X HBM3e peak. Used only for the "%peak" perf column. -_PEAK_BW_GBPS = 8000.0 +# Nominal HBM peak per arch, used only for the "%peak" perf column. A single +# constant made that column meaningless on every card but one. +_PEAK_BW_GBPS_BY_GFX = { + "gfx942": 5300.0, # MI300X (MI325X is 6000) + "gfx950": 8000.0, # MI355X HBM3E + "gfx1250": 22000.0, +} + + +def _pct_peak(gbps): + """%peak against this card's nominal HBM peak; None when the arch is unknown.""" + peak = _PEAK_BW_GBPS_BY_GFX.get(get_gfx()) + return round(gbps / peak * 100, 1) if peak else None + + # Pin the arg-rotation count. Left to itself, run_perftest derives it from # `free_memory` at call time, so two rows timed in one process rotate a # different number of times and land in different L2 states -- their `us` @@ -316,7 +329,7 @@ def test_flydsl_qk_norm_rope_quant( return { "us": round(us, 3), "GB/s": round(gbps, 0), - "%peak": round(gbps / _PEAK_BW_GBPS * 100, 1), + "%peak": _pct_peak(gbps), "err_q": err_q, "err_kv": err_kv, }