diff --git a/kernels/attention/pa_decode_fp8.py b/kernels/attention/pa_decode_fp8.py index b1867529e..e03d6624a 100644 --- a/kernels/attention/pa_decode_fp8.py +++ b/kernels/attention/pa_decode_fp8.py @@ -335,37 +335,27 @@ def pa_decode_ps_launch( query_length = query.shape[0] // context_lengths.shape[0] query_group_size = num_query_heads // num_kv_heads + batch_size = context_lengths.shape[0] + head_size = query.shape[-1] # Strides for key_scale/value_scale - if per_token_kv: - stride_ks_block = key_scale.stride(0) - stride_ks_head = key_scale.stride(1) - else: - stride_ks_block = 0 - stride_ks_head = 0 + stride_ks_block = key_scale.stride(0) if per_token_kv else 0 + stride_ks_head = key_scale.stride(1) if per_token_kv else 0 s = stream or torch.cuda.current_stream() - if sliding_window > 0: - # Launch one CTA per 256-token context partition in the sliding window: - # grid = (batch, kv_heads, max_context_partition_num). - batch_size = context_lengths.shape[0] - head_size = query.shape[-1] + if max_context_partition_num <= 0: + raise ValueError("max_context_partition_num must be positive for sliding-window decode.") + if is_graph_capturing and (exp_sums is None or max_logits is None or temporary_output is None): + raise ValueError( + "CUDA graph capture requires preallocated `exp_sums`, `max_logits`, " + "and `temporary_output` for the sliding-window path." + ) + # ── small-block (block_size 16/64) → tile kernel ── + # Key cache shape is [num_blocks, num_kv_heads, head_size // 16, block_size, 16]. + block_size = key_cache.shape[-2] + if block_size in _PA_DECODE_PS_SMALL_BLOCK_SIZES or sliding_window > 0: eqgs = query_length * query_group_size - context_partition_size = KV_COMPUTE_BLOCK - if max_context_partition_num == 0: - max_context_partition_num = get_recommended_splits( - batch_size, - num_kv_heads, - sliding_window=sliding_window, - context_partition_size=context_partition_size, - query_length=query_length, - ) - if is_graph_capturing and (exp_sums is None or max_logits is None or temporary_output is None): - raise ValueError( - "CUDA graph capture requires preallocated `exp_sums`, `max_logits`, " - "and `temporary_output` for the sliding-window path." - ) if exp_sums is None: exp_sums = torch.zeros( batch_size, num_kv_heads, max_context_partition_num, eqgs, device=dev, dtype=torch.float32 @@ -382,6 +372,9 @@ def pa_decode_ps_launch( batch_size, num_kv_heads, max_context_partition_num, eqgs, head_size, device=dev, dtype=torch.bfloat16 ) + if sliding_window > 0: + # Launch one CTA per 256-token context partition in the sliding window: + # grid = (batch, kv_heads, max_context_partition_num). # The fused SW kernel is useful only when there is no real cross-partition # parallelism to exploit. For the 1023-token window case, one CTA would # serialize six 256-token partitions and regress badly versus the @@ -393,13 +386,13 @@ def pa_decode_ps_launch( compiled_sw = compile_pa_decode_sw( sliding_window=sliding_window, + max_context_partition_num=max_context_partition_num, softmax_scale=softmax_scale, trans_v=trans_v, query_group_size=query_group_size, per_token_kv=per_token_kv, query_length=query_length, query_input_dtype=query_input_dtype, - fuse_partitions=fuse_sw_partitions, head_dim=int(head_size), ) @@ -475,16 +468,12 @@ def pa_decode_ps_launch( ) return "ps_sw_partitioned" - # ── small-block (block_size 16/64) → tile kernel ── - # Key cache shape is [num_blocks, num_kv_heads, head_size // 16, block_size, 16]. - block_size = key_cache.shape[-2] if block_size in _PA_DECODE_PS_SMALL_BLOCK_SIZES: if block_tables is None: raise ValueError( f"pa_decode_ps_launch: block_size={block_size} requires `block_tables` " "(per-sequence physical block index table)." ) - batch_size = context_lengths.shape[0] if is_graph_capturing: # Buffer sizes must be fixed ahead of capture and stay identical # across every replay, so require the caller to have preallocated diff --git a/kernels/attention/pa_decode_swa.py b/kernels/attention/pa_decode_swa.py index 74a092c5e..b9fca6c4b 100644 --- a/kernels/attention/pa_decode_swa.py +++ b/kernels/attention/pa_decode_swa.py @@ -7,17 +7,14 @@ import flydsl.compiler as flyc import flydsl.expr as fx -from flydsl._mlir.dialects import vector -from flydsl.expr import arith, as_ir_value, const_expr, gpu, range_constexpr, rocdl -from flydsl.expr.typing import Int32, T -from kernels.attention.pa_common import _compute_block_base_dw_i64, _prefetch_q_chunks -from kernels.common import buffer_ops, dpp_utils +from flydsl.expr import arith, const_expr, gpu, range_constexpr, rocdl +from flydsl.expr.typing import T +from kernels.attention.pa_common import _compute_block_base_dw_i64 +from kernels.common import dpp_utils +from kernels.common.kernels_common import get_warp_size from kernels.common.utils import ( cdiv, exp2_f32_fast, - global_load_i32, - global_load_i64x2, - global_ptr_from_addr, rcp_f32, udiv_const, unflatten_k, @@ -28,8 +25,8 @@ KV_BLOCK_SIZE = 1024 # physical page size (matches SP3 kBlockSize) KV_COMPUTE_BLOCK = 256 # tile size (matches SP3 kTileKV) NUM_WARPS = 4 -WARP_SIZE = 64 -BLOCK_THREADS = NUM_WARPS * WARP_SIZE # 256 +WARP_SIZE = get_warp_size() +BLOCK_THREADS = NUM_WARPS * WARP_SIZE # 256 on CDNA MFMA_N = 16 MFMA_K = 32 @@ -53,21 +50,45 @@ FP8_MAX = 240.0 LOG2E = 1.4426950408889634 +_FLAT_BUFFER_ELEMENTS = 1 << 30 # Tiles per block (1024 tokens / 256 tokens per tile = 4, matches SP3 kNumBlockTiles) TILES_PER_BLOCK = KV_BLOCK_SIZE // KV_COMPUTE_BLOCK # 4 -def _get_sw_mtp_group_count(query_length: int, query_group_size: int) -> int: - return cdiv(query_length * query_group_size, MFMA_N) +def _global_pointer_from_addr(addr, dtype, *, alignment: int): + ptr_type = fx.PointerType.get( + elem_ty=dtype.ir_type, + address_space=fx.AddressSpace.Global, + alignment=alignment, + ) + return fx.inttoptr(ptr_type, addr) + + +def _copy_load(source, offset, copy_atom, register): + fx.copy(copy_atom, fx.slice(source, (None, fx.Int32(offset))), register) + return fx.memref_load_vec(register) -def _get_sw_mtp_pair_offset(mtp_group_idx: int, mtp_subgroup_idx: int = 0) -> int: - return mtp_group_idx * MFMA_N + mtp_subgroup_idx * MFMA_N +def _copy_store(destination, offset, copy_atom, register, value): + fx.memref_store_vec(value, register) + fx.copy(copy_atom, register, fx.slice(destination, (None, fx.Int32(offset)))) + + +def _load_global_16b(global_ptr, byte_offset, copy_atom, register): + source = fx.make_view(global_ptr + byte_offset, fx.make_layout(16, 1)) + fx.copy(copy_atom, source, register) + return fx.memref_load_vec(register).bitcast(fx.Int64) + + +def _get_sw_mtp_group_count(query_length: int, query_group_size: int) -> int: + return cdiv(query_length * query_group_size, MFMA_N) def _load_k_flat( k_global_ptr, + k_copy_atom, + k_register, k_block_base_dw_i64, tile_token_offset_i32, k_tok_thread_base, @@ -85,7 +106,7 @@ def _load_k_flat( kbo_dw = kbo * c_tok_stride_dw for qkhe in range_constexpr(qkhe_loop): ka_dw = k_block_base_dw_i64 + fx.Int64(kbo_dw + k_he_off_dw[qkhe]) - k2 = global_load_i64x2(k_global_ptr, ka_dw * fx.Int64(4)) + k2 = _load_global_16b(k_global_ptr, ka_dw * fx.Int64(4), k_copy_atom, k_register) if const_expr(sched_vmem_after_load): rocdl.sched_barrier(rocdl.mask_vmem_rd) k2_words = fx.Vector(k2) @@ -95,71 +116,19 @@ def _load_k_flat( return k_flat -def _build_pa_thread_invariants( +def _build_pa_k_thread_invariants( warp_id, lane16id, rowid, *, - trans_v, - per_token_kv, qkhe_loop: int = 2, - vhe_loop: int = 2, ): c_tokens_per_warp = fx.Int32(TOKENS_PER_WARP) - c_mfma_n = fx.Int32(MFMA_N) k_tok_thread_base = warp_id * c_tokens_per_warp + lane16id c_tok_stride_dw = fx.Int32(FP8_ELEMS_16B // 4) c_he_stride_dw = fx.Int32(KV_BLOCK_SIZE * FP8_ELEMS_16B // 4) k_he_off_dw = [rowid * c_he_stride_dw + fx.Int32(qkhe * 4) * c_he_stride_dw for qkhe in range(qkhe_loop)] - - vhead_elems = [fx.Int32(vhe * NUM_WARPS * MFMA_N) + warp_id * c_mfma_n + lane16id for vhe in range(vhe_loop)] - v_tok_thread_off = [fx.Int32(vt * TOKENS_PER_WARP) + rowid * c_mfma_n for vt in range(VTLOOP)] - if const_expr(trans_v): - vhead_elem_dw = [vhead_elems[vhe] * fx.Int32(FP8_ELEMS_16B // 4) for vhe in range(vhe_loop)] - else: - vhead_elem_dw = [vhead_elems[vhe] * fx.Int32(KV_BLOCK_SIZE // 4) for vhe in range(vhe_loop)] - - kv_tok_thread_base = warp_id * c_tokens_per_warp + rowid * 4 - rowid_8x8 = rowid >> fx.Int32(1) - offset_in_slot = rowid & fx.Int32(1) - prob_row_i32 = PROB_ROW_STRIDE_BYTES // 4 - prob_row_i64 = PROB_ROW_STRIDE_BYTES // 8 - prob_wr_thread_base = ( - warp_id * fx.Int32(4 * MFMA_N * prob_row_i32) - + lane16id * fx.Int32(prob_row_i32) - + rowid_8x8 * fx.Int32(2) - + offset_in_slot - ) - pv_prob_read_base = rowid * fx.Int32(MFMA_N * prob_row_i64) + lane16id * fx.Int32(prob_row_i64) - - sm_lane_wave_base = lane16id * fx.Int32(NUM_WARPS) - sm_max_off = sm_lane_wave_base + warp_id - sm_sum_off = fx.Int32(NUM_WARPS * MFMA_N) + sm_lane_wave_base + warp_id - sm_rd_max_offs = [sm_lane_wave_base + fx.Int32(w) for w in range(NUM_WARPS)] - sm_rd_sum_offs = [fx.Int32(NUM_WARPS * MFMA_N) + sm_lane_wave_base + fx.Int32(w) for w in range(NUM_WARPS)] - - sm_vmax_wr_off = None - sm_vmax_rd_offs = None - if const_expr(per_token_kv): - sm_vmax_wr_off = fx.Int32(2 * NUM_WARPS * MFMA_N) + sm_lane_wave_base + warp_id - sm_vmax_rd_offs = [fx.Int32(2 * NUM_WARPS * MFMA_N) + sm_lane_wave_base + fx.Int32(w) for w in range(NUM_WARPS)] - - return ( - k_tok_thread_base, - c_tok_stride_dw, - k_he_off_dw, - v_tok_thread_off, - vhead_elem_dw, - kv_tok_thread_base, - prob_wr_thread_base, - pv_prob_read_base, - sm_max_off, - sm_sum_off, - sm_rd_max_offs, - sm_rd_sum_offs, - sm_vmax_wr_off, - sm_vmax_rd_offs, - ) + return k_tok_thread_base, c_tok_stride_dw, k_he_off_dw def _compute_sw_mtp_group_state( @@ -167,11 +136,10 @@ def _compute_sw_mtp_group_state( local_qhead_idx, *, mtp_group_idx, - mtp_subgroup_idx=0, query_length, query_group_size, ): - g_off = _get_sw_mtp_pair_offset(mtp_group_idx, mtp_subgroup_idx) + g_off = mtp_group_idx * MFMA_N lane_pair_raw = lane16id + fx.Int32(g_off) c_total_pairs = fx.Int32(query_length * query_group_size) c_pair_max = fx.Int32(query_length * query_group_size - 1) @@ -220,7 +188,6 @@ def _finish_q_fragments( abs_mask = fx.Vector.filled(4, 0x7FFFFFFF, fx.Int32) c_zero_f = fx.Float32(0.0) c_one_f = fx.Float32(1.0) - fx.Float32(FP8_MAX) q_f32_chunks = [] local_max = c_zero_f for q_src in q_chunks: @@ -264,9 +231,7 @@ def _finish_q_fragments( q_frags = [] gpu.barrier() - query_scale_lane = fx.ptr_load(softmax_base + (lane16id), result_type=fx.Vector.make_type(1, fx.Float32))[ - 0 - ].ir_value() + query_scale_lane = fx.ptr_load(softmax_base + lane16id, result_type=fx.Vector.make_type(1, fx.Float32))[0] for qkhe in range_constexpr(qkhe_loop): for qkr in range_constexpr(2): lds_rd = lane16id * fx.Int32(head_size // 8) + fx.Int32(qkhe * 8) + rowid * fx.Int32(2) + fx.Int32(qkr) @@ -277,8 +242,10 @@ def _finish_q_fragments( return q_frags, query_scale_lane -def _prefetch_sw_mtp_group_queries( - q_rsrc, +def _prefetch_sw_mtp_group_query( + q_tiles, + q_copy_atom, + q_register, batch_idx, kv_h, stride_q_seq, @@ -287,128 +254,109 @@ def _prefetch_sw_mtp_group_queries( local_qhead_idx, *, mtp_group_idx, - mtp_subgroup_count, query_length, query_group_size, - query_load_is_bf16, q_lanes_per_head, ): - mtp_prefetches = [] - c_query_length = arith.constant(query_length, type=T.i32) - c_query_group_size = arith.constant(query_group_size, type=T.i32) - for mtp_subgroup_idx in range_constexpr(mtp_subgroup_count): - qi_val, qhi_pos, qi_for_q, local_qhead_idx_for_q = _compute_sw_mtp_group_state( - lane16id, - local_qhead_idx, - mtp_group_idx=mtp_group_idx, - mtp_subgroup_idx=mtp_subgroup_idx, - query_length=query_length, - query_group_size=query_group_size, - ) - q_row = batch_idx * c_query_length + qi_for_q - q_base = q_row * stride_q_seq + (kv_h * c_query_group_size + local_qhead_idx_for_q) * stride_q_head - q_chunks = _prefetch_q_chunks( - q_rsrc, - q_base, - lane16id, - query_load_is_bf16=query_load_is_bf16, - q_lanes_per_head=q_lanes_per_head, - ) - mtp_prefetches.append((qi_val, qhi_pos, q_chunks)) - return mtp_prefetches - - -def _finish_sw_mtp_subgroup_q_fragments( - logits_base, - softmax_base, - mtp_prefetches, - lane16id, - rowid, - local_qhead_idx, - *, - mtp_subgroup_idx, - head_size: int, - qkhe_loop: int, - q_lanes_per_head: int, -): - qi_val, qhi_pos, q_chunks = mtp_prefetches[mtp_subgroup_idx] - q_frags, query_scale_lane = _finish_q_fragments( - logits_base, - softmax_base, - q_chunks, + c_query_length = fx.Int32(query_length) + c_query_group_size = fx.Int32(query_group_size) + qi_val, qhi_pos, qi_for_q, local_qhead_idx_for_q = _compute_sw_mtp_group_state( lane16id, - rowid, local_qhead_idx, - head_size=head_size, - qkhe_loop=qkhe_loop, - q_lanes_per_head=q_lanes_per_head, + mtp_group_idx=mtp_group_idx, + query_length=query_length, + query_group_size=query_group_size, ) - return qi_val, qhi_pos, q_frags, query_scale_lane + q_row = batch_idx * c_query_length + qi_for_q + q_base = q_row * stride_q_seq + (kv_h * c_query_group_size + local_qhead_idx_for_q) * stride_q_head + q_load_lane = lane16id + if const_expr(q_lanes_per_head < MFMA_N): + q_load_lane = (lane16id < fx.Int32(q_lanes_per_head)).select(lane16id, fx.Int32(0)) + q_elem = q_base + q_load_lane * fx.Int32(Q_ELEMS_PER_LANE) + q_chunks = [ + _copy_load(q_tiles, q_elem + fx.Int32(qwi * 4), q_copy_atom, q_register) + for qwi in range_constexpr(Q_CHUNKS_PER_LANE) + ] + return qi_val, qhi_pos, q_chunks -def _normalize_pa_output(running_sum, outs, zero_f, vhe_loop: int = 2): - one_f = fx.Float32(1.0).ir_value() - safe_sum = arith.select(running_sum > zero_f, running_sum, one_f) - inv_sum = rcp_f32(safe_sum) - normalized_outs = [] - for vhe in range_constexpr(vhe_loop): - normalized_outs.append(outs[vhe] * vector.broadcast(T.f32x4, inv_sum)) - return normalized_outs +def _normalize_pa_output(running_sum, outs): + safe_sum = arith.select(running_sum > fx.Float32(0.0), running_sum, fx.Float32(1.0)) + inv_sum = fx.Float32(rcp_f32(safe_sum)) + return [out * inv_sum for out in outs] def _make_pa_phase_helpers( *, trans_v, - per_token_q, per_token_kv, - needs_mask, - query_length, kv_h, v_global_ptr, - ks_rsrc, - vs_rsrc, + v_copy_atom, + v_register, + ks_tiles, + vs_tiles, + scale_copy_atom, + scale_register, logits_base, softmax_base, scale_base, stride_ks_block, stride_ks_head, softmax_scale_base, - softmax_q_scale, k_scale_val, - scale, v_scale_val, warp_id, lane16id, rowid, - k_tok_thread_base, - v_tok_thread_off, - vhead_elem_dw, - kv_tok_thread_base, - prob_wr_thread_base, - pv_prob_read_base, - sm_max_off, - sm_sum_off, - sm_rd_max_offs, - sm_rd_sum_offs, - sm_vmax_wr_off, - sm_vmax_rd_offs, - c_w, - neg_inf, - zero_f, head_size: int = 128, - qkhe_loop: int = 2, - vhe_loop: int = 2, ): + qkhe_loop = head_size // QKHE_PER_FETCH + vhe_loop = head_size // MFMA_N // NUM_WARPS + c_mfma_n = fx.Int32(MFMA_N) + + vhead_elems = [fx.Int32(vhe * NUM_WARPS * MFMA_N) + warp_id * c_mfma_n + lane16id for vhe in range(vhe_loop)] + v_tok_thread_off = [fx.Int32(vt * TOKENS_PER_WARP) + rowid * c_mfma_n for vt in range(VTLOOP)] + if const_expr(trans_v): + vhead_elem_dw = [vhead_elems[vhe] * fx.Int32(FP8_ELEMS_16B // 4) for vhe in range(vhe_loop)] + else: + vhead_elem_dw = [vhead_elems[vhe] * fx.Int32(KV_BLOCK_SIZE // 4) for vhe in range(vhe_loop)] + + kv_tok_thread_base = warp_id * fx.Int32(TOKENS_PER_WARP) + rowid * 4 + rowid_8x8 = rowid >> fx.Int32(1) + offset_in_slot = rowid & fx.Int32(1) + prob_row_i32 = PROB_ROW_STRIDE_BYTES // 4 + prob_row_i64 = PROB_ROW_STRIDE_BYTES // 8 + prob_wr_thread_base = ( + warp_id * fx.Int32(4 * MFMA_N * prob_row_i32) + + lane16id * fx.Int32(prob_row_i32) + + rowid_8x8 * fx.Int32(2) + + offset_in_slot + ) + pv_prob_read_base = rowid * fx.Int32(MFMA_N * prob_row_i64) + lane16id * fx.Int32(prob_row_i64) + + sm_lane_wave_base = lane16id * fx.Int32(NUM_WARPS) + sm_max_off = sm_lane_wave_base + warp_id + sm_sum_off = fx.Int32(NUM_WARPS * MFMA_N) + sm_lane_wave_base + warp_id + sm_rd_max_offs = [sm_lane_wave_base + fx.Int32(w) for w in range(NUM_WARPS)] + sm_rd_sum_offs = [fx.Int32(NUM_WARPS * MFMA_N) + sm_lane_wave_base + fx.Int32(w) for w in range(NUM_WARPS)] + + sm_vmax_wr_off = None + sm_vmax_rd_offs = None + if const_expr(per_token_kv): + sm_vmax_wr_off = fx.Int32(2 * NUM_WARPS * MFMA_N) + sm_lane_wave_base + warp_id + sm_vmax_rd_offs = [fx.Int32(2 * NUM_WARPS * MFMA_N) + sm_lane_wave_base + fx.Int32(w) for w in range(NUM_WARPS)] + + c_w = fx.Int32(WARP_SIZE) + neg_inf = fx.Float32(float("-inf")) + zero_f = fx.Float32(0.0) + # Sliding-window decode always needs an upper-bound mask: even for a # single query, the tail block can contain tokens beyond context_len. pv_prob_i64_elem_offs = [] for vt in range_constexpr(VTLOOP): for j in range_constexpr(2): - p_elem = ( - arith.constant(vt * 4 * MFMA_N * (PROB_ROW_STRIDE_BYTES // 8), type=T.i32) - + pv_prob_read_base - + arith.constant(j, type=T.i32) - ) + p_elem = fx.Int32(vt * 4 * MFMA_N * (PROB_ROW_STRIDE_BYTES // 8)) + pv_prob_read_base + fx.Int32(j) pv_prob_i64_elem_offs.append(p_elem) def _load_kv_scale_scalars(tile_token_offset_i32, phys_block): @@ -416,18 +364,18 @@ def _load_kv_scale_scalars(tile_token_offset_i32, phys_block): scale_block_base = phys_block * stride_ks_block + kv_h * stride_ks_head scale_stage_token = warp_id * fx.Int32(WARP_SIZE) + rowid * fx.Int32(MFMA_N) + lane16id scale_global_token = tile_token_offset_i32 + scale_stage_token - k_scale_scalar = buffer_ops.buffer_load( - ks_rsrc, + k_scale_scalar = _copy_load( + ks_tiles, scale_block_base + scale_global_token, - vec_width=1, - dtype=fx.Float32, - ) - v_scale_scalar = buffer_ops.buffer_load( - vs_rsrc, + scale_copy_atom, + scale_register, + )[0] + v_scale_scalar = _copy_load( + vs_tiles, scale_block_base + scale_global_token, - vec_width=1, - dtype=fx.Float32, - ) + scale_copy_atom, + scale_register, + )[0] return k_scale_scalar, v_scale_scalar return None @@ -457,13 +405,11 @@ def _load_v_and_scales( v_token_in_block = tile_token_offset_i32 + v_tok_thread_off[vt] if const_expr(trans_v): vt_group = v_token_in_block >> fx.Int32(4) - va_dw_delta = ( - vt_group * arith.constant(head_size * FP8_ELEMS_16B // 4, type=T.i32) + vhead_elem_dw[vhe] - ) + va_dw_delta = vt_group * fx.Int32(head_size * FP8_ELEMS_16B // 4) + vhead_elem_dw[vhe] else: va_dw_delta = vhead_elem_dw[vhe] + (v_token_in_block >> fx.Int32(2)) va_byte = (v_block_base_dw + fx.Int64(va_dw_delta)) * fx.Int64(4) - v_i64x2 = global_load_i64x2(v_global_ptr, va_byte) + v_i64x2 = _load_global_16b(v_global_ptr, va_byte, v_copy_atom, v_register) rocdl.sched_barrier(rocdl.mask_vmem_rd) vhe_data.append(v_i64x2) v_results.append(vhe_data) @@ -487,37 +433,35 @@ def _store_vmax_warp(partition_start, *, seq_end=None): kv_tok_base = partition_start + kv_tok_thread_base if const_expr(seq_end is not None) else None v_max_warp = zero_f for td in range_constexpr(TLOOP): - vs = _load_v_scale_vec(td) + vs = fx.Vector(_load_v_scale_vec(td)) + masked_values = [] for i in range_constexpr(4): + vs_i = vs[i] if const_expr(kv_tok_base is not None): - kv_tok = kv_tok_base + arith.constant(td * MFMA_N + i, type=T.i32) - vs_i = vector.extract(as_ir_value(vs), static_position=[i], dynamic_position=[]) + kv_tok = kv_tok_base + fx.Int32(td * MFMA_N + i) vs_i = arith.select(kv_tok < seq_end, vs_i, zero_f) - vs = vector.insert(vs_i, vs, static_position=[i], dynamic_position=[]) - v_max_warp = v_max_warp.maximumf(fx.Vector(vs).reduce("max")) + masked_values.append(vs_i) + vs = fx.Vector.from_elements(masked_values, dtype=fx.Float32) + v_max_warp = v_max_warp.maximumf(vs.reduce("max")) for sh in [32, 16]: - v_max_warp = v_max_warp.maximumf(v_max_warp.shuffle_xor(arith.constant(sh, type=T.i32), c_w)) + v_max_warp = v_max_warp.maximumf(v_max_warp.shuffle_xor(fx.Int32(sh), c_w)) fx.ptr_store( fx.Vector.from_elements([v_max_warp], dtype=fx.Float32), softmax_base + sm_vmax_wr_off, ) def _token_vec_i32(kv_tok_base, td: int): - kv_tok_td_base = kv_tok_base + arith.constant(td * MFMA_N, type=T.i32) + kv_tok_td_base = kv_tok_base + fx.Int32(td * MFMA_N) return fx.Vector.from_elements( - [kv_tok_td_base + arith.constant(i, type=T.i32) for i in range_constexpr(4)], + [kv_tok_td_base + fx.Int32(i) for i in range_constexpr(4)], dtype=fx.Int32, ) def _apply_token_mask_vec(logit_vec, td: int, kv_tok_base, causal_bound, seq_start, false_value): tok_vec = _token_vec_i32(kv_tok_base, td) - if const_expr(needs_mask and seq_start is not None): - in_range = (tok_vec < causal_bound) & (tok_vec >= seq_start) - elif const_expr(needs_mask): - in_range = tok_vec < causal_bound - else: - in_range = tok_vec >= seq_start - return arith.select(in_range, logit_vec, vector.broadcast(T.f32x4, arith.unwrap(false_value))) + in_range = (tok_vec < causal_bound) & (tok_vec >= seq_start) + false_vec = fx.Vector.from_elements([false_value], dtype=fx.Float32).broadcast_to(4) + return in_range.select(logit_vec, false_vec) def _qk_and_intra_softmax( k_ops, @@ -526,56 +470,45 @@ def _qk_and_intra_softmax( causal_bound, query_scale_lane=None, *, - seq_start=None, + seq_start, ): - - query_scale_vec = None - if const_expr(per_token_q): - query_scale_vec = vector.broadcast(T.f32x4, query_scale_lane * softmax_scale_base) + query_scale_vec = fx.Vector.from_elements( + [query_scale_lane * softmax_scale_base], + dtype=fx.Float32, + ).broadcast_to(4) d_out = [] for td in range_constexpr(TLOOP): - acc = arith.constant_vector(0.0, T.f32x4) + acc = fx.Vector.filled(4, 0.0, fx.Float32) for k_step in range_constexpr(qkhe_loop * 2): acc = rocdl.mfma_f32_16x16x32_fp8_fp8(T.f32x4, [k_ops[td][k_step], q_frags[k_step], acc, 0, 0, 0]) if const_expr(per_token_kv): k_scale_vec = _load_k_scale_vec(td) - scale_vec = ( - k_scale_vec * query_scale_vec - if const_expr(per_token_q) - else k_scale_vec * vector.broadcast(T.f32x4, softmax_q_scale) - ) - d_out.append(acc * scale_vec) + d_out.append(acc * (k_scale_vec * query_scale_vec)) else: - if const_expr(per_token_q): - d_out.append(acc * (query_scale_vec * vector.broadcast(T.f32x4, k_scale_val))) - else: - d_out.append(acc * vector.broadcast(T.f32x4, scale)) + d_out.append(acc * (query_scale_vec * k_scale_val)) - apply_range_mask = seq_start is not None - kv_tok_base = partition_start + kv_tok_thread_base if const_expr(needs_mask or apply_range_mask) else None + kv_tok_base = partition_start + kv_tok_thread_base qk_max = neg_inf for td in range_constexpr(TLOOP): - logits_vec = d_out[td] - if const_expr(kv_tok_base is not None): - logits_vec = _apply_token_mask_vec(logits_vec, td, kv_tok_base, causal_bound, seq_start, neg_inf) - d_out[td] = logits_vec + logits_vec = _apply_token_mask_vec(d_out[td], td, kv_tok_base, causal_bound, seq_start, neg_inf) + d_out[td] = logits_vec qk_max = qk_max.maximumf(fx.Vector(logits_vec).reduce("max")) for sh in [32, 16]: - qk_max = qk_max.maximumf(qk_max.shuffle_xor(arith.constant(sh, type=T.i32), c_w)) + qk_max = qk_max.maximumf(qk_max.shuffle_xor(fx.Int32(sh), c_w)) fx.ptr_store( fx.Vector.from_elements([qk_max], dtype=fx.Float32), softmax_base + sm_max_off, ) exp_sum = zero_f - safe_qk_max = arith.select(qk_max > neg_inf, qk_max, zero_f) if const_expr(kv_tok_base is not None) else qk_max + safe_qk_max = arith.select(qk_max > neg_inf, qk_max, zero_f) for td in range_constexpr(TLOOP): - diff_vec = fx.Vector(d_out[td]) - vector.broadcast(T.f32x4, arith.unwrap(safe_qk_max)) - p_vec = exp2_f32_fast(diff_vec * vector.broadcast(T.f32x4, arith.unwrap(fx.Float32(LOG2E)))) + diff_vec = fx.Vector(d_out[td]) - safe_qk_max + p_vec = exp2_f32_fast(diff_vec * fx.Float32(LOG2E)) exp_sum = exp_sum + fx.Vector(p_vec).reduce("add") d_out[td] = p_vec for sh in [32, 16]: - exp_sum = exp_sum + exp_sum.shuffle_xor(arith.constant(sh, type=T.i32), c_w) + exp_sum = exp_sum + exp_sum.shuffle_xor(fx.Int32(sh), c_w) fx.ptr_store( fx.Vector.from_elements([exp_sum], dtype=fx.Float32), softmax_base + sm_sum_off, @@ -595,8 +528,7 @@ def _cross_warp_softmax_and_prob_pack(d_out, rmax, rsum, outs): sum_vec = fx.ptr_load(softmax_base + (sm_rd_sum_offs[0]), result_type=fx.Vector.make_type(4, fx.Float32)) for w in range_constexpr(NUM_WARPS): diff_w = warp_rescale_factors[w] - partition_max - if const_expr(needs_mask): - diff_w = arith.select(partition_max > neg_inf, diff_w, zero_f) + diff_w = arith.select(partition_max > neg_inf, diff_w, zero_f) wf = exp2_f32_fast(diff_w * fx.Float32(LOG2E).ir_value()) w_sum = sum_vec[w] wf_sum = arith.mulf(arith.unwrap(w_sum), arith.unwrap(wf), fastmath=arith.FastMathFlags.contract) @@ -606,26 +538,22 @@ def _cross_warp_softmax_and_prob_pack(d_out, rmax, rsum, outs): my_warp_rescale = warp_rescale_factors[0] for w in range_constexpr(1, NUM_WARPS): my_warp_rescale = arith.select( - warp_id == arith.constant(w, type=T.i32), + warp_id == fx.Int32(w), warp_rescale_factors[w], my_warp_rescale, ) new_rmax = rmax.maximumf(partition_max) - if const_expr(needs_mask): - accum_scale = arith.select( - rmax > neg_inf, - exp2_f32_fast((rmax - new_rmax) * fx.Float32(LOG2E).ir_value()), - zero_f, - ) - part_to_new = arith.select( - partition_max > neg_inf, - exp2_f32_fast((partition_max - new_rmax) * fx.Float32(LOG2E).ir_value()), - zero_f, - ) - else: - accum_scale = exp2_f32_fast((rmax - new_rmax) * fx.Float32(LOG2E).ir_value()) - part_to_new = exp2_f32_fast((partition_max - new_rmax) * fx.Float32(LOG2E).ir_value()) + accum_scale = arith.select( + rmax > neg_inf, + exp2_f32_fast((rmax - new_rmax) * fx.Float32(LOG2E).ir_value()), + zero_f, + ) + part_to_new = arith.select( + partition_max > neg_inf, + exp2_f32_fast((partition_max - new_rmax) * fx.Float32(LOG2E).ir_value()), + zero_f, + ) accum_sum = arith.mulf(arith.unwrap(accum_scale), arith.unwrap(rsum), fastmath=arith.FastMathFlags.contract) partition_sum_scaled = arith.mulf( @@ -635,9 +563,8 @@ def _cross_warp_softmax_and_prob_pack(d_out, rmax, rsum, outs): ) rsum = arith.addf(accum_sum, partition_sum_scaled, fastmath=arith.FastMathFlags.contract) rmax = new_rmax - accum_scale_vec = vector.broadcast(T.f32x4, arith.unwrap(accum_scale)) for vhe in range_constexpr(vhe_loop): - outs[vhe] = outs[vhe] * accum_scale_vec + outs[vhe] = outs[vhe] * fx.Float32(accum_scale) if const_expr(per_token_kv): v_max_global = zero_f @@ -651,30 +578,28 @@ def _cross_warp_softmax_and_prob_pack(d_out, rmax, rsum, outs): prob_scale = my_warp_rescale v_correction = v_max_scaled * part_to_new for td in range_constexpr(TLOOP): - d_out[td] = d_out[td] * ( - _load_v_scale_vec(td) * vector.broadcast(T.f32x4, arith.unwrap(prob_scale * norm_factor)) - ) + d_out[td] = d_out[td] * (_load_v_scale_vec(td) * fx.Float32(prob_scale * norm_factor)) else: prob_scale = my_warp_rescale * part_to_new v_correction = v_scale_val for td in range_constexpr(TLOOP): - d_out[td] = d_out[td] * vector.broadcast(T.f32x4, arith.unwrap(prob_scale)) + d_out[td] = d_out[td] * fx.Float32(prob_scale) for td in range_constexpr(TLOOP): pv = fx.Vector(d_out[td]) - lo = rocdl.cvt_pk_fp8_f32(T.i32, pv[0], pv[1], arith.constant(0, type=T.i32), False) + lo = rocdl.cvt_pk_fp8_f32(T.i32, pv[0], pv[1], fx.Int32(0), False) pk = rocdl.cvt_pk_fp8_f32(T.i32, pv[2], pv[3], lo, True) - elem_base = prob_wr_thread_base + arith.constant(td * MFMA_N * (PROB_ROW_STRIDE_BYTES // 4), type=T.i32) + elem_base = prob_wr_thread_base + fx.Int32(td * MFMA_N * (PROB_ROW_STRIDE_BYTES // 4)) pk_vec = fx.Vector.from_elements([pk], dtype=fx.Int32) fx.ptr_store(pk_vec, logits_base + elem_base) return rmax, rsum, outs, v_correction def _pv_mfma(v_ops, outs, v_correction): - v_correction = fx.Float32(v_correction).ir_value() + v_correction = fx.Float32(v_correction) fm_contract = arith.FastMathFlags.contract - v_correction_vec = vector.broadcast(T.f32x4, v_correction) + v_correction_vec = fx.Vector.from_elements([v_correction], dtype=fx.Float32).broadcast_to(4) for vhe in range_constexpr(vhe_loop): - tmp_out = arith.constant_vector(0.0, T.f32x4) + tmp_out = fx.Vector.filled(4, 0.0, fx.Float32) for vt in range_constexpr(VTLOOP): v_i64x2 = fx.Vector(v_ops[vt][vhe]) for j in range_constexpr(2): @@ -711,17 +636,6 @@ def _pv_mfma(v_ops, outs, v_correction): ) -def get_sw_max_context_partition_num( - sliding_window: int, - context_partition_size: int = KV_COMPUTE_BLOCK, - query_length: int = 1, -) -> int: - if sliding_window <= 0: - return 0 - window_token_count = sliding_window + query_length - return cdiv(window_token_count - 1, context_partition_size) + 1 - - @functools.lru_cache(maxsize=256) def compile_pa_decode_sw_reduce( *, @@ -743,6 +657,12 @@ def compile_pa_decode_sw_reduce( LOGITS_DTYPE = fx.Float16 else: LOGITS_DTYPE = fx.BFloat16 + if output_dtype_str == "f32": + OUTPUT_DTYPE = fx.Float32 + elif output_dtype_str == "f16": + OUTPUT_DTYPE = fx.Float16 + else: + OUTPUT_DTYPE = fx.BFloat16 block_threads = head_size assert block_threads > 0, "head_size must be positive" assert block_threads <= 1024, "head_size must fit in one workgroup" @@ -762,17 +682,17 @@ def pa_decode_sw_reduce_kernel( exp_sums_ptr: fx.Int64, max_logits_ptr: fx.Int64, logits_ptr: fx.Int64, - stride_output_bs: Int32, - stride_output_len: Int32, - stride_output_kv_head: Int32, - stride_output_group_size: Int32, - stride_exp_sums_seq: Int32, - stride_exp_sums_head: Int32, - stride_exp_sums_part: Int32, - stride_logits_seq: Int32, - stride_logits_head: Int32, - stride_logits_part: Int32, - stride_logits_group: Int32, + stride_output_bs: fx.Int32, + stride_output_len: fx.Int32, + stride_output_kv_head: fx.Int32, + stride_output_group_size: fx.Int32, + stride_exp_sums_seq: fx.Int32, + stride_exp_sums_head: fx.Int32, + stride_exp_sums_part: fx.Int32, + stride_logits_seq: fx.Int32, + stride_logits_head: fx.Int32, + stride_logits_part: fx.Int32, + stride_logits_group: fx.Int32, ): tid = fx.Int32(gpu.thread_id("x")) batch_idx = fx.Int32(gpu.block_id("x")) @@ -784,10 +704,31 @@ def pa_decode_sw_reduce_kernel( if const_expr(max_context_partition_num > WARP_SIZE): part_weights_lds = lds.part_weights.view(fx.make_layout(max_context_partition_num, 1)) - out_rsrc = buffer_ops.create_buffer_resource_from_addr(output_ptr) - es_rsrc = buffer_ops.create_buffer_resource_from_addr(exp_sums_ptr) - ml_rsrc = buffer_ops.create_buffer_resource_from_addr(max_logits_ptr) - logits_rsrc = buffer_ops.create_buffer_resource_from_addr(logits_ptr) + def _divide_addr(addr, dtype): + pointer = _global_pointer_from_addr(addr, dtype, alignment=dtype.width // 8) + flat = fx.make_view(pointer, fx.make_layout(_FLAT_BUFFER_ELEMENTS, 1)) + return fx.logical_divide( + fx.rocdl.make_buffer_tensor(flat), + fx.make_layout(1, 1), + ) + + output = _divide_addr(output_ptr, OUTPUT_DTYPE) + exp_sums = _divide_addr(exp_sums_ptr, fx.Float32) + max_logits = _divide_addr(max_logits_ptr, fx.Float32) + logits = _divide_addr(logits_ptr, LOGITS_DTYPE) + + copy_f32 = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), fx.Float32) + copy_logits = fx.make_copy_atom( + fx.rocdl.BufferCopy32b() if LOGITS_DTYPE.width == 32 else fx.rocdl.BufferCopy16b(), + LOGITS_DTYPE, + ) + copy_output = fx.make_copy_atom( + fx.rocdl.BufferCopy32b() if OUTPUT_DTYPE.width == 32 else fx.rocdl.BufferCopy16b(), + OUTPUT_DTYPE, + ) + f32_register = fx.make_rmem_tensor(1, fx.Float32) + logits_register = fx.make_rmem_tensor(1, LOGITS_DTYPE) + output_register = fx.make_rmem_tensor(1, OUTPUT_DTYPE) c_zero_f = fx.Float32(0.0) c_one_f = fx.Float32(1.0) @@ -874,8 +815,8 @@ def _wave_reduce_sum(val): + part_i32 * stride_exp_sums_part + eqgs_idx ) - part_sum_raw = buffer_ops.buffer_load(es_rsrc, es_off, vec_width=1, dtype=T.f32) - part_max_raw = buffer_ops.buffer_load(ml_rsrc, es_off, vec_width=1, dtype=T.f32) + part_sum_raw = _copy_load(exp_sums, es_off, copy_f32, f32_register)[0] + part_max_raw = _copy_load(max_logits, es_off, copy_f32, f32_register)[0] part_sum = arith.select(lane_in_range, part_sum_raw, c_zero_f) part_max = arith.select(lane_in_range, part_max_raw, c_neg_inf) @@ -909,7 +850,7 @@ def _wave_reduce_sum(val): + eqgs_idx * stride_logits_group + tid ) - part_logits_raw = buffer_ops.buffer_load(logits_rsrc, logits_off, vec_width=1, dtype=LOGITS_DTYPE) + part_logits_raw = _copy_load(logits, logits_off, copy_logits, logits_register)[0] part_logits = fx.Float32(part_logits_raw) acc = acc + part_logits * weight else: @@ -927,7 +868,7 @@ def _wave_reduce_sum(val): + part_i32 * stride_exp_sums_part + eqgs_idx ) - part_max_raw = buffer_ops.buffer_load(ml_rsrc, es_off, vec_width=1, dtype=fx.Float32) + part_max_raw = _copy_load(max_logits, es_off, copy_f32, f32_register)[0] part_max = arith.select(in_chunk, part_max_raw, c_neg_inf) chunk_max = _block_reduce(part_max, "max") global_max = global_max.maximumf(chunk_max) @@ -945,8 +886,8 @@ def _wave_reduce_sum(val): + part_i32 * stride_exp_sums_part + eqgs_idx ) - part_sum_raw = buffer_ops.buffer_load(es_rsrc, es_off, vec_width=1, dtype=T.f32) - part_max_raw = buffer_ops.buffer_load(ml_rsrc, es_off, vec_width=1, dtype=T.f32) + part_sum_raw = _copy_load(exp_sums, es_off, copy_f32, f32_register)[0] + part_max_raw = _copy_load(max_logits, es_off, copy_f32, f32_register)[0] part_sum = arith.select(in_chunk, part_sum_raw, c_zero_f) part_max = arith.select(in_chunk, part_max_raw, c_neg_inf) part_scale = arith.select( @@ -976,13 +917,13 @@ def _wave_reduce_sum(val): + part_i32 * stride_exp_sums_part + eqgs_idx ) - part_sum_raw = buffer_ops.buffer_load(es_rsrc, es_off, vec_width=1, dtype=T.f32) - part_max_raw = buffer_ops.buffer_load(ml_rsrc, es_off, vec_width=1, dtype=T.f32) + part_sum_raw = _copy_load(exp_sums, es_off, copy_f32, f32_register)[0] + part_max_raw = _copy_load(max_logits, es_off, copy_f32, f32_register)[0] + part_sum = arith.select(in_chunk, part_sum_raw, c_zero_f) + part_max = arith.select(in_chunk, part_max_raw, global_max) + part_scale = exp2_f32_fast((part_max - global_max) * c_log2e) + weight = part_sum * part_scale * inv_global_exp_sum if in_chunk: - part_sum = part_sum_raw - part_max = part_max_raw - part_scale = exp2_f32_fast((part_max - global_max) * c_log2e) - weight = part_sum * part_scale * inv_global_exp_sum part_idx_idx = fx.Int32(part_i32) fx.memref_store(weight, part_weights_lds, part_idx_idx) @@ -1000,7 +941,7 @@ def _wave_reduce_sum(val): + eqgs_idx * stride_logits_group + tid ) - part_logits_raw = buffer_ops.buffer_load(logits_rsrc, logits_off, vec_width=1, dtype=LOGITS_DTYPE) + part_logits_raw = _copy_load(logits, logits_off, copy_logits, logits_register)[0] part_logits = fx.Float32(part_logits_raw) acc = acc + part_logits * weight @@ -1013,13 +954,14 @@ def _wave_reduce_sum(val): + group_idx * stride_output_group_size + tid ) - if const_expr(output_dtype_str == "f32"): - out_val = acc - elif const_expr(output_dtype_str == "f16"): - out_val = acc.to(fx.Float16) - else: - out_val = acc.to(fx.BFloat16) - buffer_ops.buffer_store(out_val, out_rsrc, out_off) + out_val = acc if const_expr(output_dtype_str == "f32") else acc.to(OUTPUT_DTYPE) + _copy_store( + output, + out_off, + copy_output, + output_register, + fx.Vector.from_elements([out_val], dtype=OUTPUT_DTYPE), + ) @flyc.jit def launch_pa_decode_sw_reduce( @@ -1082,13 +1024,13 @@ def launch_pa_decode_sw_reduce( @functools.lru_cache(maxsize=256) def compile_pa_decode_sw( sliding_window: int, # required > 0 -- baked as compile-time constant + max_context_partition_num: int, softmax_scale=None, trans_v=False, query_group_size=16, per_token_kv=False, query_length: int = 1, query_input_dtype: str = "bf16", - fuse_partitions: bool = False, head_dim: int = 128, ): """Compile a Gluon-style partitioned PA decode kernel for sliding window. @@ -1096,27 +1038,23 @@ def compile_pa_decode_sw( Grid = (batch_size, num_kv_heads * mtp_groups, max_context_partition_num). Each GPU block processes one 256-token partition selected from the visible KV region: the sliding tail window. - sliding_window is a compile-time constant. + sliding_window and max_context_partition_num are compile-time constants. """ assert sliding_window > 0, "compile_pa_decode_sw requires sliding_window > 0" if query_input_dtype not in ("bf16", "f16"): raise ValueError("`compile_pa_decode_sw` only supports bf16/f16 query inputs.") - if head_dim % QKHE_PER_FETCH != 0 or head_dim % (MFMA_N * NUM_WARPS) != 0 or head_dim % Q_ELEMS_PER_LANE != 0: - raise ValueError(f"Unsupported head_dim={head_dim}; must be a multiple of {MFMA_N * NUM_WARPS}.") + if head_dim not in (64, 128): + raise ValueError(f"`compile_pa_decode_sw` only supports head_dim 64 or 128, got {head_dim}.") + fuse_partitions = max_context_partition_num <= 1 _HEAD = head_dim _QKHELOOP = head_dim // QKHE_PER_FETCH _VHELOOP = head_dim // MFMA_N // NUM_WARPS _Q_LANES_PER_HEAD = head_dim // Q_ELEMS_PER_LANE - query_load_is_bf16 = query_input_dtype == "bf16" + _QUERY_DTYPE = fx.BFloat16 if query_input_dtype == "bf16" else fx.Float16 if softmax_scale is None: softmax_scale = 1.0 / (head_dim**0.5) _softmax_scale = float(softmax_scale) _bs = KV_BLOCK_SIZE # 1024 - _max_context_partition_num = get_sw_max_context_partition_num( - sliding_window, - KV_COMPUTE_BLOCK, - query_length, - ) _mtp_groups = _get_sw_mtp_group_count(query_length, query_group_size) LDS_VMAX_BYTES = NUM_WARPS * MFMA_N * 4 if const_expr(per_token_kv) else 0 @@ -1151,26 +1089,26 @@ def pa_decode_sw_kernel( context_lengths_ptr: fx.Int64, key_scale_ptr: fx.Int64, value_scale_ptr: fx.Int64, - stride_q_seq: Int32, - stride_q_head: Int32, - stride_k_block: Int32, - stride_k_head: Int32, - stride_v_block: Int32, - stride_v_head: Int32, - stride_es_seq: Int32, - stride_es_head: Int32, - stride_es_part: Int32, - stride_to_seq: Int32, - stride_to_head: Int32, - stride_to_part: Int32, - stride_to_group: Int32, - stride_out_bs: Int32, - stride_out_len: Int32, - stride_out_kv_head: Int32, - stride_out_group_size: Int32, - stride_bt_seq: Int32, - stride_ks_block: Int32, - stride_ks_head: Int32, + stride_q_seq: fx.Int32, + stride_q_head: fx.Int32, + stride_k_block: fx.Int32, + stride_k_head: fx.Int32, + stride_v_block: fx.Int32, + stride_v_head: fx.Int32, + stride_es_seq: fx.Int32, + stride_es_head: fx.Int32, + stride_es_part: fx.Int32, + stride_to_seq: fx.Int32, + stride_to_head: fx.Int32, + stride_to_part: fx.Int32, + stride_to_group: fx.Int32, + stride_out_bs: fx.Int32, + stride_out_len: fx.Int32, + stride_out_kv_head: fx.Int32, + stride_out_group_size: fx.Int32, + stride_bt_seq: fx.Int32, + stride_ks_block: fx.Int32, + stride_ks_head: fx.Int32, ): tid = fx.Int32(gpu.thread_id("x")) batch_idx = fx.Int32(gpu.block_id("x")) @@ -1178,31 +1116,50 @@ def pa_decode_sw_kernel( kv_h = udiv_const(grid_y, _mtp_groups) mtp_group_from_grid = urem_const(grid_y, _mtp_groups) partition_idx = fx.Int32(gpu.block_id("z")) - cl_global_ptr = global_ptr_from_addr(context_lengths_ptr) - context_len = global_load_i32(cl_global_ptr, batch_idx) - lane16id = tid & 15 - rowid = (tid >> 4) & 3 + lane16id = tid & fx.Int32(15) + rowid = (tid >> fx.Int32(4)) & fx.Int32(3) warp_id = fx.Int32(tid >> fx.Int32(6)) - q_rsrc = buffer_ops.create_buffer_resource_from_addr(query_ptr) - k_global_ptr = global_ptr_from_addr(key_cache_ptr) - v_global_ptr = global_ptr_from_addr(value_cache_ptr) - - bt_global_ptr = global_ptr_from_addr(block_tables_ptr) - es_rsrc = buffer_ops.create_buffer_resource_from_addr(exp_sums_ptr) - ml_rsrc = buffer_ops.create_buffer_resource_from_addr(max_logits_ptr) - to_rsrc = buffer_ops.create_buffer_resource_from_addr(tmp_out_ptr) - out_rsrc = buffer_ops.create_buffer_resource_from_addr(out_ptr) - ks_rsrc = buffer_ops.create_buffer_resource_from_addr(key_scale_ptr) - vs_rsrc = buffer_ops.create_buffer_resource_from_addr(value_scale_ptr) + def _divide_addr(addr, dtype): + pointer = _global_pointer_from_addr(addr, dtype, alignment=dtype.width // 8) + flat = fx.make_view(pointer, fx.make_layout(_FLAT_BUFFER_ELEMENTS, 1)) + return fx.logical_divide( + fx.rocdl.make_buffer_tensor(flat), + fx.make_layout(1, 1), + ) - q_scale_val = 1.0 + q_tiles = _divide_addr(query_ptr, _QUERY_DTYPE) + context_lengths = _divide_addr(context_lengths_ptr, fx.Int32) + block_tables = _divide_addr(block_tables_ptr, fx.Int32) + exp_sums = _divide_addr(exp_sums_ptr, fx.Float32) + max_logits = _divide_addr(max_logits_ptr, fx.Float32) + tmp_out_tiles = _divide_addr(tmp_out_ptr, fx.BFloat16) + out_tiles = _divide_addr(out_ptr, fx.BFloat16) + key_scales = _divide_addr(key_scale_ptr, fx.Float32) + value_scales = _divide_addr(value_scale_ptr, fx.Float32) + + copy_i32 = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), fx.Int32) + copy_f32 = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), fx.Float32) + copy_q = fx.make_copy_atom(fx.rocdl.BufferCopy64b(), _QUERY_DTYPE) + copy_bf16x4 = fx.make_copy_atom(fx.rocdl.BufferCopy64b(), fx.BFloat16) + + i32_register = fx.make_rmem_tensor(1, fx.Int32) + f32_register = fx.make_rmem_tensor(1, fx.Float32) + q_register = fx.make_rmem_tensor(4, _QUERY_DTYPE) + bf16x4_register = fx.make_rmem_tensor(4, fx.BFloat16) + + k_global_ptr = _global_pointer_from_addr(key_cache_ptr, fx.Uint8, alignment=16) + v_global_ptr = _global_pointer_from_addr(value_cache_ptr, fx.Uint8, alignment=16) + global_copy_16b = fx.make_copy_atom(fx.UniversalCopy128b(), fx.Uint8) + global_register_16b = fx.make_rmem_tensor(16, fx.Uint8) + + context_len = _copy_load(context_lengths, batch_idx, copy_i32, i32_register)[0] if const_expr(per_token_kv): - k_scale_val = 1.0 - v_scale_val = 1.0 + k_scale_val = fx.Float32(1.0) + v_scale_val = fx.Float32(1.0) else: - k_scale_val = buffer_ops.buffer_load(ks_rsrc, 0, vec_width=1) - v_scale_val = buffer_ops.buffer_load(vs_rsrc, 0, vec_width=1) + k_scale_val = _copy_load(key_scales, 0, copy_f32, f32_register)[0] + v_scale_val = _copy_load(value_scales, 0, copy_f32, f32_register)[0] lds = fx.SharedAllocator().allocate(SharedStorage).peek() logits_base = lds.logits.ptr @@ -1211,48 +1168,25 @@ def pa_decode_sw_kernel( if const_expr(per_token_kv): scale_base = lds.scale.ptr - _softmax_scale_const = arith.constant(_softmax_scale, type=T.f32) - _softmax_q_scale = _softmax_scale_const * q_scale_val - _scale = _softmax_q_scale * k_scale_val # per-tensor only; per-token uses per-token k_scale - c_w = fx.Int32(WARP_SIZE) + _softmax_scale_const = fx.Float32(_softmax_scale) NEG_INF = fx.Float32(float("-inf")) ZERO_F = fx.Float32(0.0) c_cps = fx.Int32(KV_COMPUTE_BLOCK) c_bs = fx.Int32(_bs) local_qhead_idx = warp_id * 4 + rowid - ( - _k_tok_thread_base, - _c_tok_stride_dw, - _k_he_off_dw, - _v_tok_thread_off, - _vhead_elem_dw, - _kv_tok_thread_base, - _prob_wr_thread_base, - _pv_prob_read_base, - _sm_max_off, - _sm_sum_off, - _sm_rd_max_offs, - _sm_rd_sum_offs, - _sm_vmax_wr_off, - _sm_vmax_rd_offs, - ) = _build_pa_thread_invariants( + _k_tok_thread_base, _c_tok_stride_dw, _k_he_off_dw = _build_pa_k_thread_invariants( warp_id, lane16id, rowid, - trans_v=trans_v, - per_token_kv=per_token_kv, qkhe_loop=_QKHELOOP, - vhe_loop=_VHELOOP, ) # ── Context length and partition mapping ── # Visible tiles cover the union of all per-query sliding windows. - _c_sw = fx.Int32(sliding_window) - _c_query_len = fx.Int32(query_length) num_tiles_for_seq = (context_len + c_cps - 1) >> fx.Int32(8) - seq_start_global = context_len - _c_query_len - _c_sw + seq_start_global = context_len - query_length - sliding_window seq_start_global = arith.select(seq_start_global > 0, seq_start_global, 0) tail_start_tile = seq_start_global >> fx.Int32(8) visible_tile_count = num_tiles_for_seq - tail_start_tile @@ -1272,45 +1206,27 @@ def pa_decode_sw_kernel( _pv_mfma, ) = _make_pa_phase_helpers( trans_v=trans_v, - per_token_q=True, per_token_kv=per_token_kv, - needs_mask=True, - query_length=query_length, kv_h=kv_h, v_global_ptr=v_global_ptr, - ks_rsrc=ks_rsrc, - vs_rsrc=vs_rsrc, + v_copy_atom=global_copy_16b, + v_register=global_register_16b, + ks_tiles=key_scales, + vs_tiles=value_scales, + scale_copy_atom=copy_f32, + scale_register=f32_register, logits_base=logits_base, softmax_base=softmax_base, scale_base=scale_base, stride_ks_block=stride_ks_block, stride_ks_head=stride_ks_head, softmax_scale_base=_softmax_scale_const, - softmax_q_scale=_softmax_q_scale, k_scale_val=k_scale_val, - scale=_scale, v_scale_val=v_scale_val, warp_id=warp_id, lane16id=lane16id, rowid=rowid, - k_tok_thread_base=_k_tok_thread_base, - v_tok_thread_off=_v_tok_thread_off, - vhead_elem_dw=_vhead_elem_dw, - kv_tok_thread_base=_kv_tok_thread_base, - prob_wr_thread_base=_prob_wr_thread_base, - pv_prob_read_base=_pv_prob_read_base, - sm_max_off=_sm_max_off, - sm_sum_off=_sm_sum_off, - sm_rd_max_offs=_sm_rd_max_offs, - sm_rd_sum_offs=_sm_rd_sum_offs, - sm_vmax_wr_off=_sm_vmax_wr_off, - sm_vmax_rd_offs=_sm_vmax_rd_offs, - c_w=c_w, - neg_inf=NEG_INF, - zero_f=ZERO_F, head_size=_HEAD, - qkhe_loop=_QKHELOOP, - vhe_loop=_VHELOOP, ) def _process_block_split( @@ -1341,9 +1257,6 @@ def _process_block_split( outs = _pv_mfma(v0_ops, outs, vc0) return rmax, rsum, outs - def _f32_bits_as_i32(value): - return fx.Float32(value).ir_value().bitcast(fx.Int32.ir_type) - def _store_partition_results(eqgs_lane, running_sum, running_max, outelems_norm): for vhe in range_constexpr(_VHELOOP): hs_base = fx.Int32(vhe * NUM_WARPS * MFMA_N) + warp_id * fx.Int32(MFMA_N) + rowid * 4 @@ -1354,22 +1267,38 @@ def _store_partition_results(eqgs_lane, running_sum, running_max, outelems_norm) + eqgs_lane * stride_to_group + hs_base ) - out_i32 = fx.Vector(outelems_norm[vhe]).to(fx.BFloat16).bitcast(fx.Int32) - buffer_ops.buffer_store(out_i32, to_rsrc, to_off * 2, offset_is_bytes=True) + out_bf16 = fx.Vector(outelems_norm[vhe]).to(fx.BFloat16) + _copy_store( + tmp_out_tiles, + to_off, + copy_bf16x4, + bf16x4_register, + out_bf16, + ) es_off = batch_idx * stride_es_seq + kv_h * stride_es_head + partition_idx * stride_es_part + eqgs_lane - es_i32 = _f32_bits_as_i32(running_sum) - ml_i32 = _f32_bits_as_i32(running_max) - buffer_ops.buffer_store(es_i32, es_rsrc, es_off * 4, offset_is_bytes=True) - buffer_ops.buffer_store(ml_i32, ml_rsrc, es_off * 4, offset_is_bytes=True) + _copy_store( + exp_sums, + es_off, + copy_f32, + f32_register, + fx.Vector.from_elements([running_sum], dtype=fx.Float32), + ) + _copy_store( + max_logits, + es_off, + copy_f32, + f32_register, + fx.Vector.from_elements([running_max], dtype=fx.Float32), + ) def _store_group_results(qi_val, qhi_pos, running_sum, running_max, outs): - outelems_norm = _normalize_pa_output(running_sum, outs, ZERO_F, vhe_loop=_VHELOOP) + outelems_norm = _normalize_pa_output(running_sum, outs) eqgs_lane = qi_val * fx.Int32(query_group_size) + qhi_pos _store_partition_results(eqgs_lane, running_sum, running_max, outelems_norm) def _store_fused_group_results(qi_val, qhi_pos, running_sum, outs): - outelems_norm = _normalize_pa_output(running_sum, outs, ZERO_F, vhe_loop=_VHELOOP) + outelems_norm = _normalize_pa_output(running_sum, outs) for vhe in range_constexpr(_VHELOOP): hs_base = fx.Int32(vhe * NUM_WARPS * MFMA_N) + warp_id * fx.Int32(MFMA_N) + rowid * 4 out_off = ( @@ -1379,8 +1308,14 @@ def _store_fused_group_results(qi_val, qhi_pos, running_sum, outs): + qhi_pos * stride_out_group_size + hs_base ) - out_i32 = fx.Vector(outelems_norm[vhe]).to(fx.BFloat16).bitcast(fx.Int32) - buffer_ops.buffer_store(out_i32, out_rsrc, out_off * 2, offset_is_bytes=True) + out_bf16 = fx.Vector(outelems_norm[vhe]).to(fx.BFloat16) + _copy_store( + out_tiles, + out_off, + copy_bf16x4, + bf16x4_register, + out_bf16, + ) def _write_empty_partition(): zero_output = [fx.Vector.filled(4, 0.0, fx.Float32) for _ in range_constexpr(_VHELOOP)] @@ -1388,7 +1323,6 @@ def _write_empty_partition(): lane16id, local_qhead_idx, mtp_group_idx=mtp_group_from_grid, - mtp_subgroup_idx=0, query_length=query_length, query_group_size=query_group_size, ) @@ -1397,18 +1331,20 @@ def _write_empty_partition(): def _run_valid_partition(): def _get_tile_metadata(tile_partition_idx_value, tile_valid): - if const_expr(tile_valid): - safe_tile_partition_idx = tile_partition_idx_value - tile_context_len = context_len - else: - safe_tile_partition_idx = arith.select(tile_valid, tile_partition_idx_value, 0) - tile_context_len = arith.select(tile_valid, context_len, 0) + safe_tile_partition_idx = ( + arith.select(tile_valid, tile_partition_idx_value, fx.Int32(0)) + if const_expr(fuse_partitions) + else tile_partition_idx_value + ) + tile_context_len = ( + arith.select(tile_valid, context_len, fx.Int32(0)) if const_expr(fuse_partitions) else context_len + ) tile_seq_partition_idx = safe_tile_partition_idx >> fx.Int32(2) tile_block_split_idx = safe_tile_partition_idx & fx.Int32(TILES_PER_BLOCK - 1) tile_token_offset_local = tile_block_split_idx * c_cps tile_kv_seq_start = tile_seq_partition_idx * c_bs + tile_token_offset_local tile_bt_off = batch_idx * stride_bt_seq + tile_seq_partition_idx - tile_phys_block = global_load_i32(bt_global_ptr, tile_bt_off) + tile_phys_block = _copy_load(block_tables, tile_bt_off, copy_i32, i32_register)[0] return tile_token_offset_local, tile_kv_seq_start, tile_context_len, tile_phys_block def _load_tile(tile_metadata, tile_scale_scalars): @@ -1417,6 +1353,8 @@ def _load_tile(tile_metadata, tile_scale_scalars): tile_k_flat = _load_k_flat( k_global_ptr, + global_copy_16b, + global_register_16b, tile_k_base, tile_token_offset_local, _k_tok_thread_base, @@ -1439,8 +1377,10 @@ def _load_tile(tile_metadata, tile_scale_scalars): tile_context_len, ) - mtp_prefetches = _prefetch_sw_mtp_group_queries( - q_rsrc, + mtp_prefetch = _prefetch_sw_mtp_group_query( + q_tiles, + copy_q, + q_register, batch_idx, kv_h, stride_q_seq, @@ -1448,80 +1388,54 @@ def _load_tile(tile_metadata, tile_scale_scalars): lane16id, local_qhead_idx, mtp_group_idx=mtp_group_from_grid, - mtp_subgroup_count=1, query_length=query_length, query_group_size=query_group_size, - query_load_is_bf16=query_load_is_bf16, q_lanes_per_head=_Q_LANES_PER_HEAD, ) - if const_expr(fuse_partitions): - tile_valid = fx.Int32(0) < visible_tile_count - prefetched_tile_metadata = _get_tile_metadata(tail_start_tile, tile_valid) - else: - prefetched_tile_metadata = _get_tile_metadata(tile_partition_idx_raw, True) + tile_valid = fx.Int32(0) < visible_tile_count if const_expr(fuse_partitions) else True + tile_partition_idx = tail_start_tile if const_expr(fuse_partitions) else tile_partition_idx_raw + prefetched_tile_metadata = _get_tile_metadata(tile_partition_idx, tile_valid) prefetched_tile_scale_scalars = _load_kv_scale_scalars( prefetched_tile_metadata[0], prefetched_tile_metadata[3], ) - qi_val, qhi_pos, q_frags, query_scale_lane = _finish_sw_mtp_subgroup_q_fragments( + qi_val, qhi_pos, q_chunks = mtp_prefetch + q_frags, query_scale_lane = _finish_q_fragments( logits_base, softmax_base, - mtp_prefetches, + q_chunks, lane16id, rowid, local_qhead_idx, - mtp_subgroup_idx=0, head_size=_HEAD, qkhe_loop=_QKHELOOP, q_lanes_per_head=_Q_LANES_PER_HEAD, ) + ( + tile_k_ops, + tile_v_and_scales, + tile_kv_seq_start, + tile_context_len, + ) = _load_tile(prefetched_tile_metadata, prefetched_tile_scale_scalars) + attention_context_len = tile_context_len if const_expr(fuse_partitions) else context_len + causal_bound = attention_context_len + fx.Int32(1 - query_length) + qi_val + seq_start = attention_context_len - fx.Int32(query_length + sliding_window) + qi_val + outs = [fx.Vector.filled(4, 0.0, fx.Float32) for _ in range_constexpr(_VHELOOP)] + running_max, running_sum, outs = _process_block_split( + NEG_INF, + ZERO_F, + outs, + tile_k_ops, + tile_v_and_scales, + q_frags, + causal_bound, + query_scale_lane, + seq_start, + tile_kv_seq_start, + ) if const_expr(fuse_partitions): - running_max = NEG_INF - running_sum = ZERO_F - outs = [arith.constant_vector(0.0, T.f32x4) for _ in range_constexpr(_VHELOOP)] - ( - tile_k_ops, - tile_v_and_scales, - tile_kv_seq_start, - tile_context_len, - ) = _load_tile(prefetched_tile_metadata, prefetched_tile_scale_scalars) - causal_bound = tile_context_len + fx.Int32(1 - query_length) + qi_val - seq_start = tile_context_len - fx.Int32(query_length + sliding_window) + qi_val - running_max, running_sum, outs = _process_block_split( - running_max, - running_sum, - outs, - tile_k_ops, - tile_v_and_scales, - q_frags, - causal_bound, - query_scale_lane, - seq_start, - tile_kv_seq_start, - ) _store_fused_group_results(qi_val, qhi_pos, running_sum, outs) else: - ( - k_ops, - preloaded_v_and_scales, - tile_kv_seq_start, - _, - ) = _load_tile(prefetched_tile_metadata, prefetched_tile_scale_scalars) - causal_bound = context_len + fx.Int32(1 - query_length) + qi_val - seq_start = context_len - fx.Int32(query_length + sliding_window) + qi_val - outs = [arith.constant_vector(0.0, T.f32x4) for _ in range_constexpr(_VHELOOP)] - running_max, running_sum, outs = _process_block_split( - NEG_INF, - ZERO_F, - outs, - k_ops, - preloaded_v_and_scales, - q_frags, - causal_bound, - query_scale_lane, - seq_start, - tile_kv_seq_start, - ) _store_group_results(qi_val, qhi_pos, running_sum, running_max, outs) if const_expr(fuse_partitions): @@ -1545,29 +1459,29 @@ def launch_pa_decode_sw( cl: fx.Int64, ks: fx.Int64, vs: fx.Int64, - s_q_seq: Int32, - s_q_head: Int32, - s_k_block: Int32, - s_k_head: Int32, - s_v_block: Int32, - s_v_head: Int32, - s_es_seq: Int32, - s_es_head: Int32, - s_es_part: Int32, - s_to_seq: Int32, - s_to_head: Int32, - s_to_part: Int32, - s_to_group: Int32, - s_out_bs: Int32, - s_out_len: Int32, - s_out_kv_head: Int32, - s_out_group_size: Int32, - s_bt_seq: Int32, - s_ks_block: Int32, - s_ks_head: Int32, - gx: Int32, - gy: Int32, - gz: Int32, + s_q_seq: fx.Int32, + s_q_head: fx.Int32, + s_k_block: fx.Int32, + s_k_head: fx.Int32, + s_v_block: fx.Int32, + s_v_head: fx.Int32, + s_es_seq: fx.Int32, + s_es_head: fx.Int32, + s_es_part: fx.Int32, + s_to_seq: fx.Int32, + s_to_head: fx.Int32, + s_to_part: fx.Int32, + s_to_group: fx.Int32, + s_out_bs: fx.Int32, + s_out_len: fx.Int32, + s_out_kv_head: fx.Int32, + s_out_group_size: fx.Int32, + s_bt_seq: fx.Int32, + s_ks_block: fx.Int32, + s_ks_head: fx.Int32, + gx: fx.Int32, + gy: fx.Int32, + gz: fx.Int32, stream: fx.Stream = fx.Stream(None), ): pa_decode_sw_kernel(