diff --git a/.github/scripts/aiter_test.sh b/.github/scripts/aiter_test.sh index a5007955a16..e96b340d94c 100755 --- a/.github/scripts/aiter_test.sh +++ b/.github/scripts/aiter_test.sh @@ -125,6 +125,19 @@ for file in "${sharded_files[@]}"; do } | tee -a latest_test.log test_cmd=(env AITER_MLA_DECODE_PERSISTENT_MAX_BATCH=0 timeout 60m python3 "$file") ;; + op_tests/test_flydsl_pa_decode.py) + # The CLI sweep is separate from the compact parametrized regression. + test_cmd=( + timeout 60m + bash -c ' + set -euo pipefail + test_file=$1 + python3 -m pytest -q "$test_file" + python3 "$test_file" + ' + _ "$file" + ) + ;; op_tests/test_gemm_a6w6.py) { echo "Running tuned dispatch plus every compatible A6W6 ASM kernel" diff --git a/aiter/ops/flydsl/__init__.py b/aiter/ops/flydsl/__init__.py index 160fc5b8787..618a498f944 100644 --- a/aiter/ops/flydsl/__init__.py +++ b/aiter/ops/flydsl/__init__.py @@ -96,6 +96,7 @@ ".gather_kv_b_proj", "gather_kv_b_proj_flydsl_supported", ), + "pa_decode": (".pa_decode", "pa_decode"), "gather_kv_b_proj_flydsl_fp8_supported": ( ".gather_kv_b_proj", "gather_kv_b_proj_flydsl_fp8_supported", @@ -126,6 +127,7 @@ "gather_kv_b_proj_flydsl", "gather_kv_b_proj_flydsl_fp8_supported", "gather_kv_b_proj_flydsl_supported", + "pa_decode", ] _fused_moe_impl_path = "aiter.ops.flydsl.fused_moe_gfx942:run_flydsl_moe_gfx942_impl" diff --git a/aiter/ops/flydsl/kernels/dpp_utils.py b/aiter/ops/flydsl/kernels/dpp_utils.py index 9cb557217f1..d93b3b4e65d 100644 --- a/aiter/ops/flydsl/kernels/dpp_utils.py +++ b/aiter/ops/flydsl/kernels/dpp_utils.py @@ -37,3 +37,25 @@ def update_dpp_i32( [], **kw, ) + + +def dpp_xor_f32(src, offset: int, **kw): + """Return ``src`` from the lane selected by a 16-lane XOR DPP pattern.""" + from flydsl.expr.typing import T + + src_i32 = fx.Float32(src).bitcast(fx.Int32) + if offset == 8: + out_i32 = update_dpp_i32(src_i32, src_i32, 280, 0xF, 0xC, False, **kw) + out_i32 = update_dpp_i32(out_i32, src_i32, 264, 0xF, 0x3, False, **kw) + elif offset == 4: + out_i32 = update_dpp_i32(src_i32, src_i32, 276, 0xF, 0xA, False, **kw) + out_i32 = update_dpp_i32(out_i32, src_i32, 260, 0xF, 0x5, False, **kw) + elif offset == 2: + out_i32 = update_dpp_i32(src_i32, src_i32, 78, 0xF, 0xF, False, **kw) + elif offset == 1: + out_i32 = update_dpp_i32(src_i32, src_i32, 177, 0xF, 0xF, False, **kw) + else: + raise ValueError( + f"dpp_xor_f32 only supports 16-lane offsets 1, 2, 4, 8; got {offset}" + ) + return out_i32.bitcast(T.f32) diff --git a/aiter/ops/flydsl/kernels/pa_decode_kernel.py b/aiter/ops/flydsl/kernels/pa_decode_kernel.py new file mode 100644 index 00000000000..f5c3b5da333 --- /dev/null +++ b/aiter/ops/flydsl/kernels/pa_decode_kernel.py @@ -0,0 +1,2081 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (C) 2025-2026 FlyDSL Project Contributors + +"""FP8 paged-attention tile kernel. + +K/V use e4m3 (FNUZ on gfx942, OCP on gfx950); BF16/FP16 Q and probabilities P +are quantized to FP8. Q/key scales fold into QK, value scale and 1/FP8_MAX into +the epilogue; softmax max/sum stay f32. Tuned gfx950 BF16 per-token MTP3/MTP4 +uses K128 MFMA instead of K32, preserving normalized Q/P and operand layouts. +Tuned gfx950 BF16 scalar decode casts Q/P directly, without normalization or +1/FP8_MAX compensation: Q must fit FP8, and small Q/probabilities may underflow. +Per-token scales retain range normalization. + +Logical layouts (not preshuffled): + +* ``query`` [num_seqs, num_q_heads, head_dim] f16/bf16 (head_dim contiguous) +* ``key_cache`` [num_blocks, num_kv_heads, head_dim//16, block_size, 16] fp8 +* ``value_cache`` [num_blocks, num_kv_heads, block_size//16, head_dim, 16] (trans_v) + or [num_blocks, num_kv_heads, head_dim, block_size] (plain), by rank +* ``block_tables`` [num_seqs, max_blocks_per_seq] int32 +* ``context_lengths`` [num_seqs] int32 +* ``output`` [num_seqs, num_q_heads, head_dim] same dtype as query +* K/V scales [1] per-tensor or [num_blocks, num_kv_heads, block_size] per-token + +Four-wave CTAs process 256-token blocks: QK splits tokens, PV splits head dim, +and P passes through LDS to transpose ownership between the MMAs. +""" + +import flydsl.compiler as flyc +import flydsl.expr as fx +from flydsl.compiler.kernel_function import CompilationContext +from flydsl.compiler.protocol import dsl_size_of +from flydsl.expr import arith, const_expr, gpu, range_constexpr +from flydsl.expr import math as fmath +from flydsl.expr.typing import ReductionOp, T +from flydsl.runtime.device import get_rocm_arch + +from . import dpp_utils +from .tensor_shim import buf_base_i64, buf_copy_store, ptr_buf_tensor +from .utils import rcp_f32 + +MFMA_MNK = 16 # M=N=16; query rows are padded to M-tiles. +# One i64 operand pack/lane covers K32, including in K128 layouts. +FP8_PACK_K = 32 +WAVE = 64 +# Accumulator elements/lane are independent of MFMA K. +MFMA_ACC_ELEMS = MFMA_MNK * MFMA_MNK // WAVE +LOG2E = 1.4426950408889634 +KV_COMPUTE_BLOCK = 256 +# Key by selected specialization, not batch/head/CU scheduling inputs. +_PA_DECODE_TILE_CACHE = {} + + +def compile_pa_decode_tile( + *, + head_dim: int, + query_group_size: int, + block_size: int, + num_seqs: int, + num_kv_heads: int, + num_compute_units: int, + num_partitions: int = 1, + softmax_scale: float | None = None, + query_dtype: str = "f16", + per_token_kv: bool = False, + query_length: int = 1, + trans_v: bool = True, + wide_kv_addressing: bool = False, + query_splits: int | None = None, + use_work_plan: bool = False, + work_capacity: int | None = None, + sliding_window: int = 0, + use_sinks: bool = False, + sink_dtype_str: str = "f32", +): + """Select and cache a PA-decode kernel and launch wrapper. + + ``query_splits=None`` selects from host-known grid bounds; an explicit + count overrides splitting and selects the matching prefetch policy. + CTAs receive equal groups of MTP positions, flattened with GQA into 16-row + M-tiles. Partial output stays unsplit for the shared reducer. + + Single-tile specialization requires a standard plan whose capacity proves + one KV tile per active task; padding skips the task body. Selected plans + map packed slots over (B, C/B), taking sequence IDs from plan records. + Launch B must be positive, divide capacity C, and match schedule selection; + the JIT launch does not recheck these conditions. B and C/B are not cache keys. + Planned gfx950 BF16 D128 output uses slot-rebased buffer stores when the + slot's byte span fits signed i32, retaining the dtype's 2-byte alignment. + + Positive ``sliding_window`` requires a plan and includes the query token. + Plans cover the MTP window union; scores are masked per query row. + ``use_sinks`` adds a zero-value per-head logit only to direct NP=1 output; + partitioned/planned output adds it once in the reducer, not in partials. + + Masked V bytes must remain finite because ``0 * NaN == NaN`` in PV MFMA. + Pages past the sequence are pinned to block 0; callers must leave the + unwritten tail of the last owned page finite. ``wide_kv_addressing`` uses + i64 offsets when a cache reaches 2 GiB and the i32 page product would wrap. + """ + if sliding_window > 0 and not use_work_plan: + raise ValueError("positive sliding_window requires work_plan") + is_gfx950 = "gfx95" in get_rocm_arch() + IS_BF16 = query_dtype == "bf16" + TUNED_SHAPE = is_gfx950 and head_dim == 128 and block_size in (16, 128) + TUNED_PER_TOKEN = TUNED_SHAPE and IS_BF16 and per_token_kv + # Scalar scheduling is broader than the BF16 single-query numerical fast path. + TUNED_SCALAR = TUNED_SHAPE and trans_v and not per_token_kv + + dense_workgroups = num_seqs * num_kv_heads * num_partitions + assert query_length >= 1, f"query_length must be >= 1, got {query_length}" + split_workgroups = dense_workgroups + single_tile_plan = False + if use_work_plan: + if work_capacity is not None: + split_workgroups = work_capacity * num_kv_heads + if sliding_window > 0: + # Bound the unaligned MTP window union, excluding padding CTAs. + window_tiles = ( + sliding_window + query_length - 2 + KV_COMPUTE_BLOCK - 1 + ) // KV_COMPUTE_BLOCK + 1 + split_workgroups = min( + split_workgroups, num_seqs * num_kv_heads * window_tiles + ) + # For n nonempty sequences, T <= n * window_tiles and the budget + # leaves >= n * (window_tiles - 1) extras. Prefix apportionment + # therefore gives each visible tile its own task, even after refresh. + single_tile_plan = ( + work_capacity is not None + and num_partitions >= window_tiles + and work_capacity >= num_seqs * window_tiles + ) + if query_splits is None: + # Single-tile page128 MTP4 tolerates more split-query CTAs per CU. + split_weight = ( + 1 + if single_tile_plan and query_length == 4 and block_size == 128 and trans_v + else query_length + ) + query_splits = ( + query_length + if TUNED_PER_TOKEN + and num_kv_heads == 1 + and query_length in (2, 4) + and query_group_size == 16 + and split_weight * split_workgroups <= 2 * num_compute_units + else 1 + ) + assert query_splits in (1, 2, 4), "query_splits must be one of 1, 2, 4" + assert query_length % query_splits == 0, "query_splits must divide query_length" + QUERIES_PER_CTA = query_length // query_splits + TOTAL_ROWS = query_length * query_group_size + + # Prefetch single-M-tile queries; large multi-tile grids favor fewer registers. + PER_TOKEN_M1 = ( + TUNED_PER_TOKEN + and QUERIES_PER_CTA == 1 + and 8 <= query_group_size <= 16 + and ( + query_splits > 1 + or block_size == 16 + or not trans_v + or dense_workgroups <= num_compute_units + or (single_tile_plan and split_workgroups <= 2 * num_compute_units) + ) + ) + prefetch_v = PER_TOKEN_M1 or ( + TUNED_SCALAR + and block_size == 128 + and TOTAL_ROWS <= MFMA_MNK + and num_compute_units < dense_workgroups <= 2 * num_compute_units + ) + # Require an exact one-tile task budget and the small planned reducer. + batch_first_plan_grid = ( + TUNED_PER_TOKEN + and single_tile_plan + and query_length == query_splits == 4 + and num_kv_heads == 1 + and query_group_size == 16 + and block_size == 128 + and 4096 <= sliding_window <= 8192 + and 4 <= num_seqs <= 24 + and num_partitions <= 64 + and work_capacity == num_seqs * window_tiles + and split_workgroups <= 2 * num_compute_units + ) + cache_key = ( + head_dim, + query_group_size, + block_size, + num_partitions, + softmax_scale, + query_dtype, + per_token_kv, + query_length, + trans_v, + wide_kv_addressing, + prefetch_v, + query_splits, + use_work_plan, + single_tile_plan, + batch_first_plan_grid, + sliding_window, + use_sinks, + sink_dtype_str, + ) + cached = _PA_DECODE_TILE_CACHE.get(cache_key) + if cached is not None: + return cached + + buffer_plan_output = ( + use_work_plan + and is_gfx950 + and head_dim == 128 + and IS_BF16 + and 0 < TOTAL_ROWS * head_dim * 2 <= 0x7FFFFFFF + ) + # Larger windows have identical visibility for int32 context lengths. + sliding_window = min(sliding_window, 2**31 - 1) + FP8 = fx.Float8E4M3FN if is_gfx950 else fx.Float8E4M3FNUZ + FP8_MAX = ( + 448.0 if is_gfx950 else 240.0 + ) # max representable magnitude of the format above + + assert ( + head_dim % MFMA_MNK == 0 + ), f"head_dim {head_dim} must be a multiple of {MFMA_MNK}" + assert block_size in ( + 16, + 64, + 128, + ), f"pa_decode_tile only supports block_size in (16, 64, 128), got {block_size}" + assert query_dtype in ( + "f16", + "bf16", + ), f"pa_decode_tile only supports query_dtype in ('f16', 'bf16'), got {query_dtype}" + Q_DTYPE = fx.BFloat16 if IS_BF16 else fx.Float16 + + assert ( + head_dim % 64 == 0 + ), f"pa_decode_tile only supports head_dim that's a multiple of 64, got {head_dim}" + # Query rows flatten as (MTP, GQA). + CTA_ROWS = QUERIES_PER_CTA * query_group_size + M_TILES = (CTA_ROWS + MFMA_MNK - 1) // MFMA_MNK + ROWS_PADDED = M_TILES * MFMA_MNK + # Avoid repeated BF16 max packing/unpacking. + Q_ABSMAX_F32 = sliding_window > 0 and TUNED_PER_TOKEN + # Stage each token's scales once rather than once per rgroup. + UNIQUE_SCALE_STAGING = ( + sliding_window > 0 + and TUNED_PER_TOKEN + and block_size == 128 + and trans_v + and query_group_size == 16 + and query_length == 1 + and single_tile_plan == prefetch_v + ) + SCALAR_FP8_DECODE = ( + TUNED_SCALAR and IS_BF16 and query_length == 1 and query_group_size in (8, 16) + ) + # K128 consumes four K32 packs without changing cache/LDS layouts. + WIDE_FP8_MFMA = ( + TUNED_PER_TOKEN and query_length in (3, 4) and query_group_size == 16 + ) + MFMA_K = 128 if WIDE_FP8_MFMA else FP8_PACK_K + PACKS_PER_MFMA = MFMA_K // FP8_PACK_K + MTP4_FUSED = WIDE_FP8_MFMA and M_TILES == 4 + # Wide page128 addresses make carrying V across QK too register-heavy. + MTP4_PREFETCH_V = MTP4_FUSED and (block_size == 16 or not wide_kv_addressing) + TUNE_PAGE128 = block_size == 128 and (PER_TOKEN_M1 or TUNED_SCALAR) + PAGE16_VPIPE = prefetch_v and block_size == 16 + REUSE_KV_PAGES = PER_TOKEN_M1 + SCALES_BEFORE_CURRENT_V = ( + WIDE_FP8_MFMA and PER_TOKEN_M1 and (block_size == 16 or not trans_v) + ) + # Scale before masking to avoid -inf * 0 and a second mask. + M1_SCALE_BEFORE_MASK = REUSE_KV_PAGES or SCALAR_FP8_DECODE + P_BUFFERS = M_TILES if MTP4_FUSED else 2 if TUNE_PAGE128 and M_TILES == 3 else 1 + # PV uses V=A, P=B: output [head-dim, query-row=lane16]. + NWARP = 4 # 4 waves / CTA + TILE_TOK = KV_COMPUTE_BLOCK + TOK_PER_WARP = TILE_TOK // NWARP + assert TILE_TOK == NWARP * TOK_PER_WARP, "KV tile must split evenly across warps" + assert ( + TOK_PER_WARP == NWARP * MFMA_MNK + ), "per-warp token ownership must match the MFMA chunk layout" + NCHUNK = TOK_PER_WARP // MFMA_MNK # 4 + # A warp owns 64 tokens: four page-16s, one page-64, or half a page-128. + PAGES_PER_CHUNK = (TOK_PER_WARP + block_size - 1) // block_size + KV_EXTENT = (1 << 42) if wide_kv_addressing else (1 << 30) + assert ( + head_dim % (NWARP * MFMA_MNK) == 0 + ), "head_dim must split across the 4 warps for PV" + + # Four 16-element QK loads form each 64-element fetch group. + RGROUP_QUARTERS = 4 + QK_CHUNK_ELEMS = 16 + QKHE_LOOP = head_dim // (RGROUP_QUARTERS * QK_CHUNK_ELEMS) + assert ( + QKHE_LOOP >= 1 + ), f"head_dim {head_dim} must be at least {RGROUP_QUARTERS * QK_CHUNK_ELEMS}" + # QK operand-pack count, not the number of MFMA instructions. + N_SUBCHUNKS = head_dim // FP8_PACK_K + assert N_SUBCHUNKS % PACKS_PER_MFMA == 0, "QK packs must fill whole MFMA atoms" + assert TILE_TOK % MFMA_K == 0, "PV tokens must fill whole MFMA atoms" + + # Absmax's lane16 butterfly fixes the Q chunk count at 16. + NQCHUNK = 16 + QCHUNK = ( + head_dim // NQCHUNK + ) # f16 elements per lane's load chunk (8 for head_dim=128, 4 for head_dim=64) + # Loads are at most 128 bits; larger chunks must not leave an unloaded tail. + assert QCHUNK <= 8 or QCHUNK % 8 == 0, ( + f"head_dim {head_dim} is unsupported: head_dim//{NQCHUNK} ({QCHUNK}) must " + f"be <= 8 or a multiple of 8" + ) + QLOAD_UNIT = min(8, QCHUNK) + N_QLOADS = QCHUNK // QLOAD_UNIT + + VHE_CHUNKS = head_dim // (NWARP * MFMA_MNK) # 2 for head_dim=128, 1 for head_dim=64 + VHE_SIZE = head_dim // VHE_CHUNKS + OP_ELEMS = MFMA_ACC_ELEMS # PV C-fragment elements/lane/chunk + # Eight i64 packs/lane: eight K32 or two K128 PV instructions. + NVOPS = TILE_TOK // FP8_PACK_K + STEPS_PER_PAGE = block_size // MFMA_MNK + STEPS_PER_CHUNK = min(block_size, TOK_PER_WARP) // MFMA_MNK + + if softmax_scale is None: + softmax_scale = 1.0 / (head_dim**0.5) + NP = int(num_partitions) # context partitions (grid.z); compile-time constant + DIRECT_SINKS = use_sinks and NP == 1 and not use_work_plan + SINK_DTYPE = fx.Float32 + if DIRECT_SINKS: + SINK_DTYPE = { + "f32": fx.Float32, + "f16": fx.Float16, + "bf16": fx.BFloat16, + }[sink_dtype_str] + + BLOCK_THREADS = NWARP * WAVE # 256 + + K_SLOT, V_SLOT = 0, 1 + # Per-M-tile loop-carried state after K/V: output chunks, max, and denom. + STATE_PER_M = VHE_CHUNKS + 2 + V_DATA_SLOT = 2 + STATE_PER_M * M_TILES + + # LDS holds Q/P, cross-wave max/sum, V page IDs and per-token scales. + # PV output and online-softmax state remain in registers. + f32 = 4 + sQ_bytes = ROWS_PADDED * head_dim * 1 # fp8 + # First-QK's max barrier retires all sQ reads before P overwrites it. + # Still-live Q scales stay outside the aliased, 16-byte-aligned Q/P regions. + sP_off = 0 + # Padding avoids 32-bank conflicts while preserving ds_read_b128 alignment. + SP_ROW_BYTES = TILE_TOK + 16 + sP_bytes = P_BUFFERS * MFMA_MNK * SP_ROW_BYTES # fp8, padded rows + sQscale_off = max(sQ_bytes, sP_bytes) + sQscale_bytes = 0 if SCALAR_FP8_DECODE else ROWS_PADDED * f32 + # Tuned rows need 16-byte vector alignment; other paths use bank padding. + NWARP_PAD = NWARP if TUNE_PAGE128 or PER_TOKEN_M1 or MTP4_FUSED else NWARP + 1 + # Phase-split slices sLmax per M-tile so all pass-1 writes share one barrier. + sLmax_off = sQscale_off + sQscale_bytes + sLsum_off = sLmax_off + M_TILES * MFMA_MNK * NWARP_PAD * f32 + sVPage_off = sLsum_off + P_BUFFERS * MFMA_MNK * NWARP_PAD * f32 + sVPage_bytes = NWARP * PAGES_PER_CHUNK * 4 # i32 + # Double buffering prevents next-tile prefetch from clobbering live scales. + KV_BUF_STRIDE = 2 * NWARP * TOK_PER_WARP * f32 # k-region + v-region, one buffer + KV_SCALE_BUFFERS = 1 if single_tile_plan else 2 + sKScale_off = sVPage_off + sVPage_bytes + sVScale_off = sKScale_off + NWARP * TOK_PER_WARP * f32 + sKVScale_bytes = KV_SCALE_BUFFERS * KV_BUF_STRIDE if per_token_kv else 0 + sVScaleMax_off = sKScale_off + sKVScale_bytes + sVScaleMax_bytes = ( + NWARP_PAD * f32 if per_token_kv else 0 + ) # m-independent: one cross-warp slot + total_bytes = sVScaleMax_off + sVScaleMax_bytes + + # All typed LDS regions are 4-byte-aligned within this i32 blob. + @fx.struct + class SharedStorage: + buf: fx.Array[fx.Int32, total_bytes // 4, 16] + + @flyc.jit + def _pa_decode_tile_task( + output_ptr: fx.Pointer, # Direct static NP=1 output. + # Static partials: [B,H,NP,rows]; planned: [H,capacity,rows]. + pmax_ptr: fx.Pointer, # Natural-log row max. + psum_ptr: fx.Pointer, # Row sum. + pout_ptr: fx.Pointer, # Adds head_dim; Q_DTYPE normalized O_p/l_p. + query_ptr: fx.Pointer, + key_cache_ptr: fx.Pointer, + value_cache_ptr: fx.Pointer, + block_tables_ptr: fx.Pointer, + context_lengths_ptr: fx.Pointer, + key_scale_ptr: fx.Pointer, + value_scale_ptr: fx.Pointer, + sinks_ptr: fx.Pointer, + max_blocks_per_seq: fx.Int32, + stride_ks_block: fx.Int32, + stride_ks_head: fx.Int32, + stride_o_row: fx.Int32, + stride_o_head: fx.Int32, + stride_q_row: fx.Int32, + stride_q_head: fx.Int32, + num_sequences: fx.Int32, + planned_seq: fx.Int32, + planned_start: fx.Int32, + planned_end: fx.Int32, + planned_context: fx.Int32, + ): + tid = fx.Int32(gpu.thread_id("x")) + warp = tid // WAVE # 0..NWARP-1 + lane = tid - warp * WAVE # 0..63 + seq = fx.Int32(gpu.block_id("x")) + kv_query = fx.Int32(gpu.block_id("y")) + kv_h = kv_query // query_splits + query_begin = (kv_query % query_splits) * QUERIES_PER_CTA + part = fx.Int32(gpu.block_id("z")) # context partition handled by this CTA + n_kv = fx.Int32(gpu.grid_dim.y) // query_splits + if const_expr(use_work_plan): + if const_expr(batch_first_plan_grid): + part = fx.Int32( + fx.Uint32(gpu.block_id("x")) * fx.Uint32(gpu.grid_dim.z) + + fx.Uint32(gpu.block_id("z")) + ) + # Physical x groups packed slots, not sequences. + seq = planned_seq + capacity = fx.Int32( + fx.Uint32(gpu.grid_dim.x) * fx.Uint32(gpu.grid_dim.z) + ) + partial_slot = kv_h * capacity + part + else: + part = seq # Packed work slot, shared by all KV heads. + seq = planned_seq + partial_slot = kv_h * fx.Int32(gpu.grid_dim.x) + part + else: + partial_slot = (seq * n_kv + kv_h) * NP + part + + output = fx.recast_iter(Q_DTYPE, output_ptr) + pmax = fx.recast_iter(fx.Float32, pmax_ptr) + psum = fx.recast_iter(fx.Float32, psum_ptr) + pout = fx.recast_iter(Q_DTYPE, pout_ptr) + if const_expr(buffer_plan_output): + # Widen before multiplying: packed storage can exceed 2 GiB. + if const_expr(batch_first_plan_grid): + pout_slot = fx.Int64(kv_h) * fx.Int64(gpu.grid_dim.x) * fx.Int64( + gpu.grid_dim.z + ) + fx.Int64(part) + else: + pout_slot = fx.Int64(kv_h) * fx.Int64(gpu.grid_dim.x) + fx.Int64(part) + pout_slot_bytes = TOTAL_ROWS * head_dim * 2 + pout_base = buf_base_i64(pout_ptr) + pout_slot * fx.Int64(pout_slot_bytes) + pout_buffer = ptr_buf_tensor( + pout_base, + Q_DTYPE, + n=TOTAL_ROWS * head_dim, + unit_elems=OP_ELEMS, + # BF16 views may be only 2-byte aligned, even for 8-byte stores. + unit_stride=1, + num_records_bytes=pout_slot_bytes, + ) + if const_expr(DIRECT_SINKS): + sink_token = fx.recast_iter(SINK_DTYPE, sinks_ptr) + + # K/V use raw UniversalCopy so their optional i64 offsets remain intact. + def _make_raw_flat_loader(tensor_ptr, elem_ty, reg_width, extent): + copy_atom = fx.make_copy_atom(fx.UniversalCopy128b(), elem_ty) + reg = fx.make_rmem_tensor(fx.make_layout(reg_width, 1), elem_ty) + flat = fx.Tensor( + fx.make_view( + fx.recast_iter(elem_ty, tensor_ptr), fx.make_layout(extent, 1) + ) + ) + tiled = fx.logical_divide(flat, fx.make_layout(1, 1)) + + def _load(elem_idx): + fx.copy(copy_atom, fx.slice(tiled, (None, elem_idx)), reg) + return fx.Vector(fx.memref_load_vec(reg)) + + return _load + + _k_load_fp8x16 = _make_raw_flat_loader(key_cache_ptr, FP8, 16, KV_EXTENT) + _v_load_fp8x16 = _make_raw_flat_loader(value_cache_ptr, FP8, 16, KV_EXTENT) + + def _kv_addr(phys, page_elems, rest): + # Widen before the page product reaches 2^31 FP8 elements. + if const_expr(wide_kv_addressing): + return fx.Int64(phys) * fx.Int64(page_elems) + fx.Int64(rest) + return phys * page_elems + rest + + _q_copy_op = ( + fx.rocdl.BufferCopy128b() if QLOAD_UNIT == 8 else fx.rocdl.BufferCopy64b() + ) + q_buf = ptr_buf_tensor(query_ptr, Q_DTYPE) + q_tiled = fx.logical_divide(q_buf, fx.make_layout(1, 1)) + q_copy_atom = fx.make_copy_atom(_q_copy_op, Q_DTYPE) + q_reg = fx.make_rmem_tensor(fx.make_layout(QLOAD_UNIT, 1), Q_DTYPE) + + def _q_load_chunk(elem_idx): + fx.copy(q_copy_atom, fx.slice(q_tiled, (None, elem_idx)), q_reg) + return fx.Vector(fx.memref_load_vec(q_reg)) + + rgroup = lane // MFMA_MNK # 0..3: quarter-wave (paired with warp -> query row) + lane16 = lane - rgroup * MFMA_MNK # 0..15: this row's head-dim chunk index + qh_local = warp * 4 + rgroup # 0..15: this thread's query row within an M-tile + + def _load_q_row(qi, gs_head): + qh0 = kv_h * query_group_size + gs_head + row_byte0 = ( + (seq * query_length + query_begin + qi) * stride_q_row + + qh0 * stride_q_head + ) * 2 # 16-bit float = 2B/elem + base_elem = ( + row_byte0 + lane16 * (QCHUNK * 2) + ) // 2 # byte offset -> element index + return [ + _q_load_chunk(base_elem + u * QLOAD_UNIT) + for u in range_constexpr(N_QLOADS) + ] + + # Issue independent Q loads ahead of K/page/scale prefetch. + q_units_prefetched = None + if const_expr(MTP4_FUSED): + q_units_prefetched = [] + for m in range_constexpr(M_TILES): + flat_idx = m * MFMA_MNK + qh_local + qi = flat_idx // query_group_size + gs_head = flat_idx - qi * query_group_size + q_units_prefetched.append(_load_q_row(qi, gs_head)) + + def _k_load16(byte_off): + return _k_load_fp8x16(byte_off).bitcast(fx.Int64) + + def _v_load16(byte_off): + return _v_load_fp8x16(byte_off).bitcast(fx.Int64) + + if const_expr(use_work_plan): + context_len = planned_context + else: + ctx_buf = ptr_buf_tensor(context_lengths_ptr, fx.Int32) + ctx_tiled = fx.logical_divide(ctx_buf, fx.make_layout(1, 1)) + ctx_copy_atom = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), fx.Int32) + ctx_reg = fx.make_rmem_tensor(fx.make_layout(1, 1), fx.Int32) + fx.copy(ctx_copy_atom, fx.slice(ctx_tiled, (None, seq)), ctx_reg) + context_len = fx.Int32(fx.Vector(fx.memref_load_vec(ctx_reg))[0]) + # A partial compute tile may read past block_tables; bounded loads + # return page 0 for those masked tail tokens instead of faulting. + bt_num_records_bytes = ( + fx.Int64(num_sequences) * fx.Int64(max_blocks_per_seq) * 4 + ) + # Wide loads must preserve row starts that are only int32-aligned. + bt_buf = ptr_buf_tensor( + block_tables_ptr, + fx.Int32, + unit_elems=PAGES_PER_CHUNK, + unit_stride=1, + num_records_bytes=bt_num_records_bytes, + ) + if const_expr(not per_token_kv): + key_scale_buf = ptr_buf_tensor(key_scale_ptr, fx.Float32) + value_scale_buf = ptr_buf_tensor(value_scale_ptr, fx.Float32) + key_scale = fx.Float32(key_scale_buf[0]) + value_scale = fx.Float32(value_scale_buf[0]) + + num_pages = ( + context_len + block_size - 1 + ) // block_size # pages this sequence really owns + if const_expr(use_work_plan): + part_start = planned_start + part_end = planned_end + else: + num_tiles = (context_len + TILE_TOK - 1) // TILE_TOK + tiles_per_part = (num_tiles + NP - 1) // NP + part_start = part * tiles_per_part + part_end_raw = part_start + tiles_per_part + part_end = (part_end_raw < num_tiles).select(part_end_raw, num_tiles) + + # Keep an ir.Value pointer for use inside scf control flow; + # the Python SharedStorage handle cannot cross that boundary. + lds = fx.SharedAllocator().allocate(SharedStorage).peek() + lds_base = fx.recast_iter(fx.Uint8, lds.buf.ptr) # byte-addressed base + + def _lds_ptr(byte_off, elem_ty): + # Restore element alignment lost by the byte-base recast. + p = fx.add_offset(lds_base, fx.make_int_tuple(byte_off)) + ptr_ty = fx.PointerType.get( + elem_ty.ir_type, fx.AddressSpace.Shared, dsl_size_of(elem_ty) + ) + return fx.recast_iter(ptr_ty, p) + + def _lds_load(byte_off, elem_ty, n): + return fx.ptr_load( + _lds_ptr(byte_off, elem_ty), result_type=fx.Vector.make_type(n, elem_ty) + ) + + def _lds_store(byte_off, elem_ty, vec): + fx.ptr_store(vec, _lds_ptr(byte_off, elem_ty)) + + if const_expr(per_token_kv): + scale_load_width = ( + NCHUNK if block_size >= 64 and not UNIQUE_SCALE_STAGING else 1 + ) + scale_copy_op = ( + fx.rocdl.BufferCopy32b() + if scale_load_width == 1 + else fx.rocdl.BufferCopy128b() + ) + scale_copy_atom = fx.make_copy_atom(scale_copy_op, fx.Float32) + k_scale_buf = ptr_buf_tensor(key_scale_ptr, fx.Float32) + v_scale_buf = ptr_buf_tensor(value_scale_ptr, fx.Float32) + k_scale_tiled = fx.logical_divide(k_scale_buf, fx.make_layout(1, 1)) + v_scale_tiled = fx.logical_divide(v_scale_buf, fx.make_layout(1, 1)) + k_scale_reg = fx.make_rmem_tensor( + fx.make_layout(scale_load_width, 1), fx.Float32 + ) + v_scale_reg = fx.make_rmem_tensor( + fx.make_layout(scale_load_width, 1), fx.Float32 + ) + + def _k_scale_load(elem_idx): + fx.copy( + scale_copy_atom, + fx.slice(k_scale_tiled, (None, elem_idx)), + k_scale_reg, + ) + return fx.Vector(fx.memref_load_vec(k_scale_reg)) + + def _v_scale_load(elem_idx): + fx.copy( + scale_copy_atom, + fx.slice(v_scale_tiled, (None, elem_idx)), + v_scale_reg, + ) + return fx.Vector(fx.memref_load_vec(v_scale_reg)) + + def _load_phys_scalar(page, vec_width=1): + # Padding may contain stale page IDs. Force out-of-context pages + # to block 0 even when the table read itself is in bounds. + element_offset = seq * max_blocks_per_seq + page + if const_expr(vec_width == 1): + result = bt_buf[element_offset] + return (page < num_pages).select(fx.Int32(result), fx.Int32(0)) + bt_copy_atom = fx.make_copy_atom(fx.rocdl.BufferCopy128b(), fx.Int32) + frag = fx.make_fragment_like(fx.slice(bt_buf, (0, None))) + fx.copy( + bt_copy_atom, + fx.slice(bt_buf, (element_offset, None)), + frag, + ) + loaded = fx.Vector(fx.memref_load_vec(frag)) + return fx.Vector.from_elements( + [ + (page + i < num_pages).select(fx.Int32(loaded[i]), fx.Int32(0)) + for i in range_constexpr(vec_width) + ], + dtype=fx.Int32, + ) + + def _stage_v_page_row(phys_vec): + if lane == 0: + _lds_store( + sVPage_off + warp * (PAGES_PER_CHUNK * 4), fx.Int32, phys_vec + ) + + def _v_page_fetch_and_stage(tt_i32): + # Warp w broadcasts the page row later read by rgroup w. + base_page = tt_i32 * TILE_TOK // block_size # tile start is page-aligned + fetched = _load_phys_scalar( + base_page + (warp * TOK_PER_WARP) // block_size, PAGES_PER_CHUNK + ) + fetched_vec = ( + fx.Vector.from_elements([fx.Int32(fetched)], dtype=fx.Int32) + if const_expr(PAGES_PER_CHUNK == 1) + else fx.Vector(fetched) + ) + _stage_v_page_row(fetched_vec) + return fetched_vec + + def _v_page_read_row(): + off = sVPage_off + rgroup * (PAGES_PER_CHUNK * 4) + return _lds_load(off, fx.Int32, PAGES_PER_CHUNK) + + def _k_page_read_warp(): + off = sVPage_off + warp * (PAGES_PER_CHUNK * 4) + return _lds_load(off, fx.Int32, PAGES_PER_CHUNK) + + def _kv_buf_off(tt_val): + if const_expr(per_token_kv and not single_tile_plan): + return (tt_val & fx.Int32(1)) * KV_BUF_STRIDE + return 0 + + def _stage_kv_scale_to_lds(phys_vec, buf_off=0): + if const_expr(block_size >= 64): + phys = fx.Int32(phys_vec[0]) + chunk_tok = ( + lane if const_expr(UNIQUE_SCALE_STAGING) else lane16 * NCHUNK + ) + page_tok = (warp * TOK_PER_WARP) % block_size + chunk_tok + scale_idx = phys * stride_ks_block + kv_h * stride_ks_head + page_tok + k_scale_vec = _k_scale_load(scale_idx) + v_scale_vec = _v_scale_load(scale_idx) + slot = (warp * TOK_PER_WARP + chunk_tok) * f32 + _lds_store(sKScale_off + buf_off + slot, fx.Float32, k_scale_vec) + _lds_store(sVScale_off + buf_off + slot, fx.Float32, v_scale_vec) + else: + # Each rgroup stages its own page16 sub-block. + phys = fx.Int32(fx.Vector(phys_vec)[rgroup]) + scale_idx = phys * stride_ks_block + kv_h * stride_ks_head + lane16 + k_scale_scalar = fx.Float32(_k_scale_load(scale_idx)[0]) + v_scale_scalar = fx.Float32(_v_scale_load(scale_idx)[0]) + fx.rocdl.sched_barrier(fx.rocdl.mask_vmem_rd) + slot = (warp * TOK_PER_WARP + rgroup * MFMA_MNK + lane16) * f32 + _lds_store( + sKScale_off + buf_off + slot, + fx.Float32, + fx.Vector.from_elements([k_scale_scalar], dtype=fx.Float32), + ) + _lds_store( + sVScale_off + buf_off + slot, + fx.Float32, + fx.Vector.from_elements([v_scale_scalar], dtype=fx.Float32), + ) + + def _load_scale_vec(base_off, a, buf_off=0): + slot = (warp * TOK_PER_WARP + a * MFMA_MNK + rgroup * 4) * f32 + return _lds_load(base_off + buf_off + slot, fx.Float32, 4) + + def _load_kv_scale_vecs(a, buf_off=0): + return _load_scale_vec(sKScale_off, a, buf_off), _load_scale_vec( + sVScale_off, a, buf_off + ) + + def _mfma_fp8(a_ops, b_ops, a_base, b_base, k_packs, acc): + # Counts and offsets are i64 operand packs, not MFMA instructions. + if const_expr(WIDE_FP8_MFMA): + for inst in range_constexpr(k_packs // PACKS_PER_MFMA): + a_pack = fx.Vector.from_elements( + [ + a_ops[a_base + inst * PACKS_PER_MFMA + i] + for i in range_constexpr(PACKS_PER_MFMA) + ], + dtype=fx.Int64, + ).bitcast(fx.Int32) + b_pack = fx.Vector.from_elements( + [ + b_ops[b_base + inst * PACKS_PER_MFMA + i] + for i in range_constexpr(PACKS_PER_MFMA) + ], + dtype=fx.Int64, + ).bitcast(fx.Int32) + acc = fx.rocdl.mfma_scale_f32_16x16x128_f8f6f4( + T.f32x4, + [ + a_pack, + b_pack, + acc, + 0, + 0, + 0, + fx.Int32(0x7F7F7F7F), + 0, + fx.Int32(0x7F7F7F7F), + ], + ) + else: + for pack in range_constexpr(k_packs): + acc = fx.rocdl.mfma_f32_16x16x32_fp8_fp8( + T.f32x4, + [a_ops[a_base + pack], b_ops[b_base + pack], acc, 0, 0, 0], + ) + return acc + + # K token = warp*TOK_PER_WARP + a*MFMA_MNK + lane16; + # softmax masks and P-pack writes must use the same mapping. + def _k_ops(phys, a): + within_page_tok = (warp * TOK_PER_WARP + a * MFMA_MNK + lane16) % block_size + ops = [] + for qkhe in range_constexpr(QKHE_LOOP): + he_idx = qkhe * RGROUP_QUARTERS + rgroup + base = _kv_addr( + phys, + n_kv * (QCHUNK * block_size * QK_CHUNK_ELEMS), + ((kv_h * QCHUNK + he_idx) * block_size + within_page_tok) + * QK_CHUNK_ELEMS, + ) + w = _k_load16(base) # head[he_idx*16 : +16] -> two K32 operand packs + if const_expr(block_size == 16): + # Overlap page16 gathers. + fx.rocdl.sched_barrier(fx.rocdl.mask_vmem_rd) + ops.extend([w[0], w[1]]) + return ops # N_SUBCHUNKS i64 operands + + def _k_ops_from_phys(phys_vec): + flat = [] + for a in range_constexpr(NCHUNK): + phys = fx.Int32(phys_vec[(a * MFMA_MNK) // block_size]) + flat.extend(_k_ops(phys, a)) + if const_expr(head_dim == 64): + fx.rocdl.sched_vmem(len(flat) // 2) + + return fx.Vector.from_elements(flat, dtype=fx.Int64) + + def _k_ops_flat(tt_i32): + base_page = tt_i32 * TILE_TOK // block_size # tile start is page-aligned + fetched = _load_phys_scalar( + base_page + (warp * TOK_PER_WARP) // block_size, PAGES_PER_CHUNK + ) + phys_vec = ( + fx.Vector.from_elements([fx.Int32(fetched)], dtype=fx.Int32) + if const_expr(PAGES_PER_CHUNK == 1) + else fx.Vector(fetched) + ) + return _k_ops_from_phys(phys_vec), phys_vec + + # Empty partitions must not read K/V or block_tables. + k_pf0 = fx.Vector.filled(NCHUNK * N_SUBCHUNKS, 0, fx.Int64) + if part_start < part_end: + k_pf0, phys_vec0 = _k_ops_flat(part_start) + if const_expr(REUSE_KV_PAGES): + # Reuse K page IDs for V's LDS broadcast. + _stage_v_page_row(phys_vec0) + else: + _v_page_fetch_and_stage(part_start) + if const_expr(per_token_kv): + _stage_kv_scale_to_lds(phys_vec0, _kv_buf_off(fx.Int32(part_start))) + elif lane == 0: + _lds_store( + sVPage_off + warp * (PAGES_PER_CHUNK * 4), + fx.Int32, + fx.Vector.filled(PAGES_PER_CHUNK, 0, fx.Int32), + ) + + # Per-token scales fold into logits/P rather than scalar QK/epilogue scales. + if const_expr(per_token_kv): + scale_qk = fx.Float32(softmax_scale * LOG2E) + else: + scale_qk = fx.Float32(softmax_scale * LOG2E) * fx.Float32(key_scale) + v_scale_f = fx.Float32(value_scale) + NEG_INF = fx.Float32(float("-inf")) + ZERO_F = fx.Float32(0.0) + # Finite or -inf scores permit nnan's bare max instructions. + # Do not set ninf: -inf is the mask sentinel. + fm_nnan = arith.FastMathFlags.nnan + + def _row_off(byte_off, m_idx, width, elem_ty): + return byte_off + m_idx * (width * dsl_size_of(elem_ty)) + + def _ld1(byte_off, m_idx): + return _lds_load(_row_off(byte_off, m_idx, 1, fx.Float32), fx.Float32, 1)[0] + + def _st1(byte_off, m_idx, val): + _lds_store( + _row_off(byte_off, m_idx, 1, fx.Float32), + fx.Float32, + fx.Vector.from_elements([val], dtype=fx.Float32), + ) + + # Cross-wave scratch: scalar writes, NWARP-wide row reads. + def _st_lw(base_off, row, w, val): + off = base_off + (row * NWARP_PAD + w) * 4 + _lds_store( + off, fx.Float32, fx.Vector.from_elements([val], dtype=fx.Float32) + ) + + def _ld_lw_row(base_off, row): + off = base_off + row * (NWARP_PAD * 4) + return _lds_load(off, fx.Float32, NWARP) + + def _f32_to_fp8_words(vf32): + # Use hardware FP8 conversion; arith.truncf cannot lower it. + n = vf32.shape[0] + words = [] + for i in range_constexpr(n // 4): + b = i * 4 + lo = fx.rocdl.cvt_pk_fp8_f32(T.i32, vf32[b], vf32[b + 1], 0, False) + words.append( + fx.rocdl.cvt_pk_fp8_f32(T.i32, vf32[b + 2], vf32[b + 3], lo, True) + ) + return fx.Vector.from_elements(words, dtype=fx.Int32) + + def _st_words(byte_off, words): + _lds_store(byte_off, fx.Int32, words) + + def _q_local_absmax(q_unit): + if const_expr(Q_ABSMAX_F32): + return fmath.absf(q_unit.to(fx.Float32)).reduce(ReductionOp.MAX) + else: + return fmath.absf(q_unit).reduce(ReductionOp.MAX).to(fx.Float32) + + # M-tiles quantize disjoint rows, requiring no inter-tile barrier. + def _quant_q_row(m, q_row_off, q_units): + if const_expr(SCALAR_FP8_DECODE): + for u in range_constexpr(N_QLOADS): + _st_words( + q_row_off + + qh_local * head_dim + + lane16 * QCHUNK + + u * QLOAD_UNIT, + _f32_to_fp8_words(q_units[u].to(fx.Float32)), + ) + else: + absmax = _q_local_absmax(q_units[0]) + for u in range_constexpr(1, N_QLOADS): + absmax = fx.maxnumf( + absmax, + _q_local_absmax(q_units[u]), + ) + for sh in (8, 4, 2, 1): + absmax = fx.maxnumf(absmax, dpp_utils.dpp_xor_f32(absmax, sh)) + + q_scale = absmax * fx.Float32(1.0 / FP8_MAX) + inv = fx.Float32(rcp_f32(fx.maxnumf(q_scale, fx.Float32(1e-20)))) + inv_b = fx.Vector.from_elements([inv], dtype=fx.Float32).broadcast_to( + QLOAD_UNIT + ) + + for u in range_constexpr(N_QLOADS): + q_scaled_unit = q_units[u].to(fx.Float32) * inv_b + _st_words( + q_row_off + + qh_local * head_dim + + lane16 * QCHUNK + + u * QLOAD_UNIT, + _f32_to_fp8_words(q_scaled_unit), + ) + if lane16 == 0: + # Transposed [qh][m] enables one vector read across M-tiles. + _st1(sQscale_off, qh_local * M_TILES + m, q_scale) + + for m in range_constexpr(M_TILES): + flat_idx = m * MFMA_MNK + qh_local + qi = flat_idx // query_group_size + gs_head = flat_idx - qi * query_group_size + q_row_off = m * MFMA_MNK * head_dim + # Only the final M-tile can need a runtime row guard. + if const_expr((m + 1) * MFMA_MNK <= CTA_ROWS): + q_units = ( + q_units_prefetched[m] + if const_expr(MTP4_FUSED) + else _load_q_row(qi, gs_head) + ) + _quant_q_row(m, q_row_off, q_units) + elif flat_idx < CTA_ROWS: + _quant_q_row(m, q_row_off, _load_q_row(qi, gs_head)) + else: + _st_words( + q_row_off + qh_local * head_dim + lane16 * QCHUNK, + fx.Vector.filled(QCHUNK // 4, 0, fx.Int32), + ) + if const_expr(not SCALAR_FP8_DECODE) and lane16 == 0: + _st1(sQscale_off, qh_local * M_TILES + m, ZERO_F) + + gpu.barrier() + + v_page_pf0 = _v_page_read_row() + + # Register-resident Q (B operand) must match _k_ops' head-dim permutation. + q_ops_all = [] + for m in range_constexpr(M_TILES): + q_row_off = m * MFMA_MNK * head_dim + for qkhe in range_constexpr(QKHE_LOOP): + he_idx = qkhe * RGROUP_QUARTERS + rgroup + chunk = _lds_load( + q_row_off + lane16 * head_dim + he_idx * QK_CHUNK_ELEMS, fx.Int64, 2 + ) + q_ops_all.extend([chunk[0], chunk[1]]) + _ct = [ + fx.Vector.from_elements( + [float(a * MFMA_MNK + r) for r in range_constexpr(4)] + ) + for a in range_constexpr(NCHUNK) + ] + + def _score_mask(a, upper, lower): + valid = _ct[a] < upper + if const_expr(sliding_window > 0): + valid = valid & (_ct[a] >= lower) + return valid + + # Load one 16-token FP8 vector per head element; trans_v selects offsets. + def _v_ops(phys_row, vh): + head_group = ((vh * VHE_SIZE) // 16) + warp + head_element = head_group * 16 + lane16 + ops = [] + for sub in range_constexpr(PAGES_PER_CHUNK): + for step in range_constexpr(STEPS_PER_CHUNK): + # PV's token chunk is owned by rgroup after the LDS transpose. + page_step = ((rgroup * TOK_PER_WARP) % block_size) // 16 + step + if const_expr(trans_v): + base = _kv_addr( + phys_row[sub], + n_kv * (STEPS_PER_PAGE * head_dim * 16), + ( + (kv_h * STEPS_PER_PAGE + page_step) * head_dim + + head_element + ) + * 16, + ) + else: + base = _kv_addr( + phys_row[sub], + n_kv * (head_dim * block_size), + (kv_h * head_dim + head_element) * block_size + + page_step * 16, + ) + w = _v_load16(base) + if const_expr(block_size == 16): + fx.rocdl.sched_barrier(fx.rocdl.mask_vmem_rd) + ops.extend([w[0], w[1]]) + if const_expr(head_dim == 64): + fx.rocdl.sched_vmem(len(ops) // 2) + return ops # NVOPS i64, the 64-token contiguous run for this head + + # Distance to the newest query keeps causal bounds tile-relative. + if const_expr(QUERIES_PER_CTA == 1): + causal_offset = [ + query_length - 1 - query_begin for _m in range_constexpr(M_TILES) + ] + else: + causal_offset = [ + query_length + - 1 + - query_begin + - (m * MFMA_MNK + lane16) // query_group_size + for m in range_constexpr(M_TILES) + ] + if const_expr(CTA_ROWS % MFMA_MNK != 0): + # Padded rows still need nonnegative causal offsets. + causal_offset = [ + (offset > 0).select(offset, 0) for offset in causal_offset + ] + + def _o_slot(m, vh): + return 2 + STATE_PER_M * m + vh + + def _m_slot(m): + return 2 + STATE_PER_M * m + VHE_CHUNKS + + def _l_slot(m): + return 2 + STATE_PER_M * m + VHE_CHUNKS + 1 + + o_zero = fx.Vector.filled(OP_ELEMS, 0.0, fx.Float32) + init_state = [k_pf0, v_page_pf0] + for _m in range_constexpr(M_TILES): + init_state.extend([o_zero] * VHE_CHUNKS + [NEG_INF, ZERO_F]) + # Reuse carried V registers only after the final query consumes a chunk. + if const_expr(MTP4_PREFETCH_V): + v_pf0 = fx.Vector.filled(VHE_CHUNKS * NVOPS, 0, fx.Int64) + if part_start < part_end: + v_flat0 = [] + for vh in range_constexpr(VHE_CHUNKS): + v_flat0.extend(_v_ops(v_page_pf0, vh)) + v_pf0 = fx.Vector.from_elements(v_flat0, dtype=fx.Int64) + init_state.append(v_pf0) + # Single-tile plans eliminate loop/history but keep absolute tt for + # addressing and masks; the outer guard excludes padded tasks. + loop_start = 0 if const_expr(single_tile_plan) else part_start + loop_end = 1 if const_expr(single_tile_plan) else part_end + for loop_i, ostate in range(loop_start, loop_end, 1, init=init_state): + k_cur = ostate[ + K_SLOT + ] # this tile's prefetched K, as one (NCHUNK*N_SUBCHUNKS,) i64 vector + v_page_cur = ostate[ + V_SLOT + ] # this tile's V pages, as one PAGES_PER_CHUNK-wide i32 vector + tt = ( + fx.Int32(part_start) + if const_expr(single_tile_plan) + else fx.Int32(loop_i) + ) + tok0 = tt * TILE_TOK + # Interleave MFMA with VALU/LDS or page16 V loads. + if const_expr((not per_token_kv and M_TILES > 1) or PAGE16_VPIPE): + fx.rocdl.iglp_opt(0) + + tt1 = tt + 1 + + next_state = [ + None, + None, + ] # slots 0/1 (K_SLOT/V_SLOT) filled in at m==0 below + + cur_kv_buf = _kv_buf_off(tt) + + # Exclude unwritten-tail scales from FP8 normalization using the + # context bound, not a per-query causal bound, for shared MTP scales. + tile_valid = context_len - tok0 + window_left = None + if const_expr(sliding_window > 0): + # Clamp before subtracting MTP offsets to avoid int32 underflow; + # negative left edges exclude no tile-relative token. + window_left = tile_valid - sliding_window + window_left = (window_left > 0).select(window_left, 0) + + if const_expr(per_token_kv): + ctx_thr = fx.Vector.from_elements( + [ + tile_valid.to(fx.Float32) + - fx.Int32(warp * TOK_PER_WARP + rgroup * 4).to(fx.Float32) + ], + dtype=fx.Float32, + ).broadcast_to(4) + window_scale_thr = None + if const_expr(sliding_window > 0): + # Only the MTP window union may set the shared FP8 scale; + # invisible large scales could underflow visible probabilities. + first_visible = (window_left - (query_length - 1)).to(fx.Float32) + window_scale_thr = fx.Vector.from_elements( + [ + first_visible + - fx.Int32(warp * TOK_PER_WARP + rgroup * 4).to(fx.Float32) + ], + dtype=fx.Float32, + ).broadcast_to(4) + zero4_scale = fx.Vector.filled(4, 0.0, fx.Float32) + + def _mask_v_scale( + vec, a, thr=ctx_thr, lower=window_scale_thr, zero=zero4_scale + ): + return _score_mask(a, thr, lower).select(vec, zero) + + # Independent V loads overlap QK/softmax and are shared across M-tiles. + v_vh_shared = None + if const_expr(MTP4_PREFETCH_V): + v_carried = ostate[V_DATA_SLOT] + v_vh_shared = [ + [v_carried[vh * NVOPS + i] for i in range_constexpr(NVOPS)] + for vh in range_constexpr(VHE_CHUNKS) + ] + elif const_expr( + (M_TILES > 1 and not MTP4_FUSED) + or (prefetch_v and not SCALES_BEFORE_CURRENT_V) + ): + v_vh_shared = [ + _v_ops(v_page_cur, vh) for vh in range_constexpr(VHE_CHUNKS) + ] + + q_scale_vec = None + if const_expr(M_TILES > 1): + q_scale_vec = _lds_load( + sQscale_off + lane16 * (M_TILES * f32), fx.Float32, M_TILES + ) + + def _lmax_off_m(m): + return sLmax_off + m * MFMA_MNK * NWARP_PAD * f32 + + # Phase A publishes all M-tiles' QK maxima with one barrier. + if const_expr(M_TILES > 1): + masked_chunks_saved = [None] * M_TILES + + # Reread V scales after Phase A to reduce peak register liveness. + k_scale_shared = None + if const_expr(per_token_kv): + v_scale_A = [ + _load_scale_vec(sVScale_off, a, cur_kv_buf) + for a in range_constexpr(NCHUNK) + ] + pv_max = fx.Float32(0.0) + for a in range_constexpr(NCHUNK): + pv_max = fx.maxnumf( + pv_max, + _mask_v_scale(v_scale_A[a], a).reduce(ReductionOp.MAX), + ) + for sh in (16, 32): + pv_max = fx.maxnumf(pv_max, pv_max.shuffle_xor(sh, WAVE)) + _st_lw(sVScaleMax_off, 0, warp, pv_max) + k_scale_shared = [ + _load_scale_vec(sKScale_off, a, cur_kv_buf) + for a in range_constexpr(NCHUNK) + ] + + for m in range_constexpr(M_TILES): + frag_Ss = [] + for a in range_constexpr(NCHUNK): + acc = fx.Vector.filled(MFMA_ACC_ELEMS, 0.0, fx.Float32) + acc = _mfma_fp8( + k_cur, + q_ops_all, + a * N_SUBCHUNKS, + m * N_SUBCHUNKS, + N_SUBCHUNKS, + acc, + ) + frag_Ss.append(fx.Vector(acc)) + + scale = scale_qk * fx.Float32(q_scale_vec[m]) + n_valid_tile = (tile_valid - causal_offset[m]).to(fx.Float32) + base_tok_f = fx.Int32(warp * TOK_PER_WARP + rgroup * 4).to( + fx.Float32 + ) + thr = fx.Vector.from_elements( + [n_valid_tile - base_tok_f], dtype=fx.Float32 + ).broadcast_to(4) + window_thr = None + if const_expr(sliding_window > 0): + # Subtract before converting to avoid rounding large edges. + first_valid_tile = (window_left - causal_offset[m]).to( + fx.Float32 + ) + window_thr = fx.Vector.from_elements( + [first_valid_tile - base_tok_f], dtype=fx.Float32 + ).broadcast_to(4) + neg4 = fx.Vector.filled(4, float("-inf"), fx.Float32) + + # Scale finite logits before selecting -inf to avoid 0 * -inf + # for zero-Q rows. The stored row max is already scaled. + scale_b = fx.Vector.from_elements( + [scale], dtype=fx.Float32 + ).broadcast_to(4) + if const_expr(per_token_kv): + scaled_frags = [ + frag_Ss[a] * k_scale_shared[a] * scale_b + for a in range_constexpr(NCHUNK) + ] + else: + scaled_frags = [ + frag_Ss[a] * scale_b for a in range_constexpr(NCHUNK) + ] + + if const_expr(MTP4_FUSED): + # Interior tiles need no mask for any of the four queries. + masked_all = fx.Vector.from_elements( + [ + scaled_frags[a][r] + for a in range_constexpr(NCHUNK) + for r in range_constexpr(MFMA_ACC_ELEMS) + ], + dtype=fx.Float32, + ) + needs_mask = tile_valid < TILE_TOK + query_length - 1 + if const_expr(sliding_window > 0): + needs_mask = needs_mask | (window_left > 0) + if needs_mask: + masked_all = fx.Vector.from_elements( + [ + _score_mask(a, thr, window_thr).select( + scaled_frags[a], neg4 + )[r] + for a in range_constexpr(NCHUNK) + for r in range_constexpr(MFMA_ACC_ELEMS) + ], + dtype=fx.Float32, + ) + masked_chunks = [ + fx.Vector.from_elements( + [ + masked_all[a * MFMA_ACC_ELEMS + r] + for r in range_constexpr(MFMA_ACC_ELEMS) + ], + dtype=fx.Float32, + ) + for a in range_constexpr(NCHUNK) + ] + else: + masked_chunks = [ + _score_mask(a, thr, window_thr).select( + scaled_frags[a], neg4 + ) + for a in range_constexpr(NCHUNK) + ] + + pm = fx.Float32(float("-inf")) + for a in range_constexpr(NCHUNK): + pm = fx.maxnumf( + pm, + masked_chunks[a].reduce(ReductionOp.MAX, fastmath=fm_nnan), + fastmath=fm_nnan, + ) + for sh in (16, 32): + pm = fx.maxnumf(pm, pm.shuffle_xor(sh, WAVE), fastmath=fm_nnan) + _st_lw(_lmax_off_m(m), lane16, warp, pm) + + masked_chunks_saved[m] = masked_chunks + + # Publish next pages/scales with the Phase A barrier. + k_next = ( + fx.Vector.filled(NCHUNK * N_SUBCHUNKS, 0, fx.Int64) + if const_expr(MTP4_FUSED) + else k_cur + ) + if const_expr(not single_tile_plan) and tt1 < part_end: + if const_expr(MTP4_FUSED): + # Defer K until P packing releases saved score registers. + phys_vec1 = _v_page_fetch_and_stage(tt1) + _stage_kv_scale_to_lds(phys_vec1, _kv_buf_off(tt1)) + else: + k_next, phys_vec1 = _k_ops_flat(tt1) + _v_page_fetch_and_stage(tt1) + if const_expr(per_token_kv): + _stage_kv_scale_to_lds(phys_vec1, _kv_buf_off(tt1)) + next_state[K_SLOT] = k_next + + gpu.barrier() + + v_page_next = v_page_cur + if const_expr(not single_tile_plan) and tt1 < part_end: + v_page_next = _v_page_read_row() + next_state[V_SLOT] = v_page_next + + # Large M-tile groups reread scale chunks to bound VGPR liveness. + v_scale_shared = None + if const_expr(per_token_kv and M_TILES < 4): + v_scale_shared = [ + _load_scale_vec(sVScale_off, a, cur_kv_buf) + for a in range_constexpr(NCHUNK) + ] + + if const_expr(MTP4_FUSED): + if const_expr(not MTP4_PREFETCH_V and not trans_v): + # Hide plain V's strided loads behind P packing. + v_vh_shared = [ + _v_ops(v_page_cur, vh) for vh in range_constexpr(VHE_CHUNKS) + ] + # V-scale normalization is shared by all query rows. + v_max_global = _ld_lw_row(sVScaleMax_off, 0).reduce(ReductionOp.MAX) + v_max_scaled = v_max_global * fx.Float32(1.0 / FP8_MAX) + v_max_safe = v_max_scaled + fx.Float32(1e-8 / FP8_MAX) + norm_factor = fx.Float32(rcp_f32(v_max_safe)) + norm_factor_b = fx.Vector.from_elements( + [norm_factor], dtype=fx.Float32 + ).broadcast_to(4) + + # Private P/Lsum slots share one barrier without overwrite races. + m_new_saved = [] + safe_max_saved = [] + for m in range_constexpr(M_TILES): + tile_max = _ld_lw_row(_lmax_off_m(m), lane16).reduce( + ReductionOp.MAX, fastmath=fm_nnan + ) + m_new = ( + tile_max + if const_expr(single_tile_plan) + else fx.maxnumf( + ostate[_m_slot(m)], tile_max, fastmath=fm_nnan + ) + ) + m_new_saved.append(m_new) + safe_max = (m_new > NEG_INF).select(m_new, ZERO_F) + safe_max_saved.append( + fx.Vector.from_elements( + [safe_max], dtype=fx.Float32 + ).broadcast_to(4) + ) + ls_saved = [ZERO_F for _ in range_constexpr(M_TILES)] + # Share one scale fragment across queries, not the full scale tile. + for a in range_constexpr(NCHUNK): + v_sc = _load_scale_vec(sVScale_off, a, cur_kv_buf) + for m in range_constexpr(M_TILES): + Pa = fx.Vector( + fx.exp2( + masked_chunks_saved[m][a] - safe_max_saved[m], + fastmath="fast", + ) + ) + ls_saved[m] = ls_saved[m] + Pa.reduce(ReductionOp.ADD) + p_scaled = Pa * v_sc * norm_factor_b + word = _f32_to_fp8_words(p_scaled)[0] + p_off = ( + sP_off + + m * MFMA_MNK * SP_ROW_BYTES + + lane16 * SP_ROW_BYTES + + warp * TOK_PER_WARP + + rgroup * 4 + + a * (MFMA_MNK // 4) * f32 + ) + _lds_store( + p_off, + fx.Int32, + fx.Vector.from_elements([word], dtype=fx.Int32), + ) + for m in range_constexpr(M_TILES): + ls = ls_saved[m] + for sh in (16, 32): + ls = ls + ls.shuffle_xor(sh, WAVE) + if rgroup == 0: + _st_lw( + sLsum_off + m * MFMA_MNK * NWARP_PAD * f32, + lane16, + warp, + ls, + ) + if const_expr(not single_tile_plan) and tt1 < part_end: + # Saved scores are dead; next K can overlap the P barrier/PV. + k_next = _k_ops_from_phys(_k_page_read_warp()) + next_state[K_SLOT] = k_next + + gpu.barrier() + + if const_expr(not MTP4_PREFETCH_V and trans_v): + # Delay transposed V until saved score registers are dead. + v_vh_shared = [ + _v_ops(v_page_cur, vh) for vh in range_constexpr(VHE_CHUNKS) + ] + v_next_chunks = [] + for m in range_constexpr(M_TILES): + p_base = sP_off + m * MFMA_MNK * SP_ROW_BYTES + lsum_base = sLsum_off + m * MFMA_MNK * NWARP_PAD * f32 + o_acc = [ + ostate[_o_slot(m, vh)] for vh in range_constexpr(VHE_CHUNKS) + ] + m_prev = ostate[_m_slot(m)] + l_prev = ostate[_l_slot(m)] + m_new = m_new_saved[m] + safe_max = (m_new > NEG_INF).select(m_new, ZERO_F) + corr_reg = ( + ZERO_F + if const_expr(single_tile_plan) + else fx.Float32(fx.exp2(m_prev - safe_max, fastmath="fast")) + ) + gsum = _ld_lw_row(lsum_base, lane16).reduce(ReductionOp.ADD) + l_new = ( + gsum + if const_expr(single_tile_plan) + else l_prev * corr_reg + gsum + ) + p_ops = _lds_load( + p_base + lane16 * SP_ROW_BYTES + rgroup * 64, + fx.Int64, + NVOPS, + ) + corr_b = fx.Vector.from_elements( + [corr_reg], dtype=fx.Float32 + ).broadcast_to(OP_ELEMS) + for vh in range_constexpr(VHE_CHUNKS): + v_vh = v_vh_shared[vh] + acc = fx.Vector.filled(MFMA_ACC_ELEMS, 0.0, fx.Float32) + acc = _mfma_fp8(v_vh, p_ops, 0, 0, NVOPS, acc) + op = fx.Vector(acc) * fx.Vector.from_elements( + [v_max_scaled], dtype=fx.Float32 + ).broadcast_to(OP_ELEMS) + o_acc[vh] = ( + op + if const_expr(single_tile_plan) + else o_acc[vh] * corr_b + op + ) + if const_expr( + MTP4_PREFETCH_V + and m == M_TILES - 1 + and not single_tile_plan + ): + # Reuse the dead V chunk while the final PV finishes. + v_next_chunk = fx.Vector.filled(NVOPS, 0, fx.Int64) + if const_expr(not single_tile_plan) and tt1 < part_end: + v_next_chunk = fx.Vector.from_elements( + _v_ops(v_page_next, vh), dtype=fx.Int64 + ) + v_next_chunks.append(v_next_chunk) + next_state.extend([*o_acc, m_new, l_new]) + if const_expr(m < M_TILES - 1): + fx.rocdl.sched_barrier(0) + if const_expr(MTP4_PREFETCH_V): + if const_expr(single_tile_plan): + # Preserve the unused next-V slot in the loop state. + v_next = ostate[V_DATA_SLOT] + else: + v_next = fx.Vector.from_elements( + [ + v_next_chunks[vh][i] + for vh in range_constexpr(VHE_CHUNKS) + for i in range_constexpr(NVOPS) + ], + dtype=fx.Int64, + ) + else: + for m in range_constexpr(M_TILES): + p_base = sP_off + (m % P_BUFFERS) * MFMA_MNK * SP_ROW_BYTES + lsum_base = ( + sLsum_off + (m % P_BUFFERS) * MFMA_MNK * NWARP_PAD * f32 + ) + o_acc = [ + ostate[_o_slot(m, vh)] for vh in range_constexpr(VHE_CHUNKS) + ] + m_prev = ostate[ + _m_slot(m) + ] # this thread's own running max, carried from last tile + l_prev = ostate[ + _l_slot(m) + ] # this thread's own running denom, carried from last tile + + masked_chunks = masked_chunks_saved[m] + + v_max_scaled = None + norm_factor_b = None + if const_expr(per_token_kv): + v_max_global = _ld_lw_row(sVScaleMax_off, 0).reduce( + ReductionOp.MAX + ) + v_max_scaled = v_max_global * fx.Float32(1.0 / FP8_MAX) + v_max_safe = v_max_scaled + fx.Float32(1e-8 / FP8_MAX) + norm_factor = fx.Float32(rcp_f32(v_max_safe)) + norm_factor_b = fx.Vector.from_elements( + [norm_factor], dtype=fx.Float32 + ).broadcast_to(4) + + tile_max = _ld_lw_row(_lmax_off_m(m), lane16).reduce( + ReductionOp.MAX, fastmath=fm_nnan + ) + m_new = ( + tile_max + if const_expr(single_tile_plan) + else fx.maxnumf(m_prev, tile_max, fastmath=fm_nnan) + ) + # Empty rows need exponent reference 0 to avoid -inf-(-inf). + safe_max = (m_new > NEG_INF).select(m_new, ZERO_F) + m_new_b = fx.Vector.from_elements( + [safe_max], dtype=fx.Float32 + ).broadcast_to(4) + ls = fx.Float32(0.0) + words = [] + for a in range_constexpr(NCHUNK): + Pa = fx.Vector( + fx.exp2(masked_chunks[a] - m_new_b, fastmath="fast") + ) + ls = ls + Pa.reduce(ReductionOp.ADD) + if const_expr(per_token_kv): + v_sc = ( + _load_scale_vec(sVScale_off, a, cur_kv_buf) + if const_expr(M_TILES >= 4) + else v_scale_shared[a] + ) + p_scaled = Pa * v_sc * norm_factor_b + else: + p_scaled = Pa * fx.Vector.filled(4, FP8_MAX, fx.Float32) + words.append(_f32_to_fp8_words(p_scaled)[0]) + + p_off0 = ( + p_base + + lane16 * SP_ROW_BYTES + + warp * TOK_PER_WARP + + rgroup * 4 + ) + # Scatter P words in the interleave expected by PV reads. + for a in range_constexpr(NCHUNK): + _lds_store( + p_off0 + a * (MFMA_MNK // 4) * f32, + fx.Int32, + fx.Vector.from_elements([words[a]], dtype=fx.Int32), + ) + for sh in (16, 32): + ls = ls + ls.shuffle_xor(sh, WAVE) + # Empty history must contribute exp2(-inf-safe_max)=0; + # replacing m_prev with 0 can overflow for negative logits. + corr_reg = ( + ZERO_F + if const_expr(single_tile_plan) + else fx.Float32(fx.exp2(m_prev - safe_max, fastmath="fast")) + ) + if rgroup == 0: + _st_lw(lsum_base, lane16, warp, ls) + gpu.barrier() + gsum = _ld_lw_row(lsum_base, lane16).reduce(ReductionOp.ADD) + l_new = ( + gsum + if const_expr(single_tile_plan) + else l_prev * corr_reg + gsum + ) + + p_ops = _lds_load( + p_base + lane16 * SP_ROW_BYTES + rgroup * 64, + fx.Int64, + NVOPS, + ) + + corr_b = fx.Vector.from_elements( + [corr_reg], dtype=fx.Float32 + ).broadcast_to(OP_ELEMS) + for vh in range_constexpr(VHE_CHUNKS): + v_vh = v_vh_shared[vh] + acc = fx.Vector.filled(MFMA_ACC_ELEMS, 0.0, fx.Float32) + acc = _mfma_fp8(v_vh, p_ops, 0, 0, NVOPS, acc) + op = fx.Vector(acc) + if const_expr(per_token_kv): + op = op * fx.Vector.from_elements( + [v_max_scaled], dtype=fx.Float32 + ).broadcast_to(OP_ELEMS) + o_acc[vh] = ( + op + if const_expr(single_tile_plan) + else o_acc[vh] * corr_b + op + ) + next_state.extend([*o_acc, m_new, l_new]) + # Single P/Lsum slots need a read-retirement barrier. + # Alternating slots use the next write barrier; Phase A + # protects loop boundaries. The fence bounds live registers. + if const_expr(m < M_TILES - 1): + if const_expr(P_BUFFERS == 1): + gpu.barrier() + fx.rocdl.sched_barrier(0) + else: + o_acc = [ostate[_o_slot(0, vh)] for vh in range_constexpr(VHE_CHUNKS)] + m_prev = ostate[_m_slot(0)] # running max, carried from last tile + l_prev = ostate[_l_slot(0)] # running denom, carried from last tile + frag_Ss = [] + for a in range_constexpr(NCHUNK): + acc = fx.Vector.filled(MFMA_ACC_ELEMS, 0.0, fx.Float32) + acc = _mfma_fp8( + k_cur, q_ops_all, a * N_SUBCHUNKS, 0, N_SUBCHUNKS, acc + ) + frag_Ss.append(fx.Vector(acc)) + k_next = k_cur + if const_expr(not single_tile_plan) and tt1 < part_end: + if const_expr(REUSE_KV_PAGES): + # Stage scales before K so scale waits cannot drain K; + # page16 defers K until the P barrier. + phys_vec1 = _v_page_fetch_and_stage(tt1) + if const_expr(per_token_kv): + _stage_kv_scale_to_lds(phys_vec1, _kv_buf_off(tt1)) + if const_expr(not PAGE16_VPIPE): + k_next = _k_ops_from_phys(phys_vec1) + else: + k_next, phys_vec1 = _k_ops_flat(tt1) + _v_page_fetch_and_stage(tt1) + if const_expr(per_token_kv): + _stage_kv_scale_to_lds(phys_vec1, _kv_buf_off(tt1)) + if const_expr(SCALES_BEFORE_CURRENT_V): + # Issue V after scale staging so scale waits do not drain V. + v_vh_shared = [ + _v_ops(v_page_cur, vh) for vh in range_constexpr(VHE_CHUNKS) + ] + scale = ( + scale_qk + if const_expr(SCALAR_FP8_DECODE) + else scale_qk * _ld1(sQscale_off, lane16) + ) # per-qhead positive score scale + n_valid_tile = (tile_valid - causal_offset[0]).to(fx.Float32) + base_tok_f = fx.Int32(warp * TOK_PER_WARP + rgroup * 4).to(fx.Float32) + thr = fx.Vector.from_elements( + [n_valid_tile - base_tok_f], dtype=fx.Float32 + ).broadcast_to(4) + window_thr = None + if const_expr(sliding_window > 0): + first_valid_tile = (window_left - causal_offset[0]).to(fx.Float32) + window_thr = fx.Vector.from_elements( + [first_valid_tile - base_tok_f], dtype=fx.Float32 + ).broadcast_to(4) + neg4 = fx.Vector.filled( + 4, + float("-inf") if M1_SCALE_BEFORE_MASK else -1e30, + fx.Float32, + ) + # As in Phase A, scale before -inf masking to avoid zero-Q NaNs. + scale_b = None + if const_expr(M1_SCALE_BEFORE_MASK): + scale_b = fx.Vector.from_elements( + [scale], dtype=fx.Float32 + ).broadcast_to(4) + v_scale_vecs = None + if const_expr(per_token_kv): + v_scale_vecs = [] + scaled_frags = [] + masked_chunks = [] + for a in range_constexpr(NCHUNK): + k_scale_vec, v_scale_vec = _load_kv_scale_vecs(a, cur_kv_buf) + v_scale_vecs.append(v_scale_vec) + scaled_frag = frag_Ss[a] * k_scale_vec + if const_expr(M1_SCALE_BEFORE_MASK): + scaled_frag = scaled_frag * scale_b + masked_chunks.append( + _score_mask(a, thr, window_thr).select( + scaled_frag, neg4 + ) + ) + else: + scaled_frags.append(scaled_frag) + else: + if const_expr(M1_SCALE_BEFORE_MASK): + masked_chunks = [ + _score_mask(a, thr, window_thr).select( + frag_Ss[a] * scale_b, neg4 + ) + for a in range_constexpr(NCHUNK) + ] + else: + scaled_frags = frag_Ss + if const_expr(not M1_SCALE_BEFORE_MASK): + masked_chunks = [ + _score_mask(a, thr, window_thr).select(scaled_frags[a], neg4) + for a in range_constexpr(NCHUNK) + ] + # pass 1: per-warp max for this qhead + pm = fx.Float32(float("-inf")) + for a in range_constexpr(NCHUNK): + if const_expr(M1_SCALE_BEFORE_MASK): + pm = fx.maxnumf( + pm, + masked_chunks[a].reduce(ReductionOp.MAX, fastmath=fm_nnan), + fastmath=fm_nnan, + ) + else: + pm = fx.maxnumf(pm, masked_chunks[a].reduce(ReductionOp.MAX)) + for sh in (16, 32): + if const_expr(M1_SCALE_BEFORE_MASK): + pm = fx.maxnumf(pm, pm.shuffle_xor(sh, WAVE), fastmath=fm_nnan) + else: + pm = fx.maxnumf(pm, pm.shuffle_xor(sh, WAVE)) + _st_lw( + sLmax_off, + lane16, + warp, + pm if const_expr(M1_SCALE_BEFORE_MASK) else pm * scale, + ) # redundant across the 4 lanes sharing this qhead + if const_expr(per_token_kv): + pv_max = fx.Float32(0.0) + for a in range_constexpr(NCHUNK): + pv_max = fx.maxnumf( + pv_max, + _mask_v_scale(v_scale_vecs[a], a).reduce(ReductionOp.MAX), + ) + for sh in (16, 32): + pv_max = fx.maxnumf(pv_max, pv_max.shuffle_xor(sh, WAVE)) + _st_lw(sVScaleMax_off, 0, warp, pv_max) + gpu.barrier() + v_page_next = v_page_cur + if const_expr(not single_tile_plan) and tt1 < part_end: + v_page_next = _v_page_read_row() + next_state[V_SLOT] = v_page_next + v_max_scaled = None + norm_factor_b = None + if const_expr(per_token_kv): + v_max_global = _ld_lw_row(sVScaleMax_off, 0).reduce(ReductionOp.MAX) + v_max_scaled = v_max_global * fx.Float32(1.0 / FP8_MAX) + v_max_safe = v_max_scaled + fx.Float32(1e-8 / FP8_MAX) + norm_factor = fx.Float32(rcp_f32(v_max_safe)) + norm_factor_b = fx.Vector.from_elements( + [norm_factor], dtype=fx.Float32 + ).broadcast_to(4) + # pass 2: global max over warps -> exp -> fp8 P pack (-> sP) -> sum + if const_expr(M1_SCALE_BEFORE_MASK): + tile_max = _ld_lw_row(sLmax_off, lane16).reduce( + ReductionOp.MAX, fastmath=fm_nnan + ) + m_new = ( + tile_max + if const_expr(single_tile_plan) + else fx.maxnumf(m_prev, tile_max, fastmath=fm_nnan) + ) + else: + tile_max = _ld_lw_row(sLmax_off, lane16).reduce(ReductionOp.MAX) + m_new = ( + tile_max + if const_expr(single_tile_plan) + else fx.maxnumf(m_prev, tile_max) + ) + # Keep empty partitions' persisted max at -inf, but exponentiate + # relative to 0 to avoid -inf-(-inf) NaNs in P and corrections. + softmax_max = m_new + if const_expr(M1_SCALE_BEFORE_MASK): + softmax_max = (m_new > NEG_INF).select(m_new, ZERO_F) + m_new_b = fx.Vector.from_elements( + [softmax_max], dtype=fx.Float32 + ).broadcast_to(4) + ls = fx.Float32(0.0) + words = [] + if const_expr(not M1_SCALE_BEFORE_MASK): + zero4_p = fx.Vector.filled(4, 0.0, fx.Float32) + for a in range_constexpr(NCHUNK): + if const_expr(M1_SCALE_BEFORE_MASK): + Pa = fx.Vector( + fx.exp2(masked_chunks[a] - m_new_b, fastmath="fast") + ) + else: + # Finite mask sentinels require an explicit zero probability. + valid_a = masked_chunks[a] > fx.Vector.filled( + 4, -1e29, fx.Float32 + ) + Pa = valid_a.select( + fx.Vector( + fx.exp2( + masked_chunks[a] * scale - m_new_b, + fastmath="fast", + ) + ), + zero4_p, + ) + ls = ls + Pa.reduce(ReductionOp.ADD) + if const_expr(per_token_kv): + v_scale_this = ( + _load_scale_vec(sVScale_off, a, cur_kv_buf) + if const_expr(head_dim == 64) + else v_scale_vecs[a] + ) + p_scaled = Pa * v_scale_this * norm_factor_b + elif const_expr(SCALAR_FP8_DECODE): + p_scaled = Pa + else: + p_scaled = Pa * fx.Vector.filled(4, FP8_MAX, fx.Float32) + words.append(_f32_to_fp8_words(p_scaled)[0]) + p_off0 = ( + sP_off + lane16 * SP_ROW_BYTES + warp * TOK_PER_WARP + rgroup * 4 + ) + for a in range_constexpr(NCHUNK): + _lds_store( + p_off0 + a * (MFMA_MNK // 4) * f32, + fx.Int32, + fx.Vector.from_elements([words[a]], dtype=fx.Int32), + ) + if const_expr(head_dim == 64): + fx.rocdl.sched_dswr(NCHUNK) + for sh in (16, 32): + ls = ls + ls.shuffle_xor(sh, WAVE) + corr_reg = ( + ZERO_F + if const_expr(single_tile_plan) + else fx.Float32(fx.exp2(m_prev - softmax_max, fastmath="fast")) + ) + if rgroup == 0: + _st_lw(sLsum_off, lane16, warp, ls) + gpu.barrier() + if const_expr(PAGE16_VPIPE): # noqa: SIM102 + # Next K can now overlap P reads/PV. + if const_expr(not single_tile_plan) and tt1 < part_end: + k_next = _k_ops_from_phys(_k_page_read_warp()) + next_state[K_SLOT] = k_next + gsum = _ld_lw_row(sLsum_off, lane16).reduce(ReductionOp.ADD) + l_new = ( + gsum if const_expr(single_tile_plan) else l_prev * corr_reg + gsum + ) + p_ops = _lds_load( + sP_off + lane16 * SP_ROW_BYTES + rgroup * 64, fx.Int64, NVOPS + ) + corr_b = fx.Vector.from_elements( + [corr_reg], dtype=fx.Float32 + ).broadcast_to(OP_ELEMS) + if const_expr(prefetch_v): + v_vh_batch = v_vh_shared + else: + v_vh_batch = [ + _v_ops(v_page_cur, vh) for vh in range_constexpr(VHE_CHUNKS) + ] + for vh in range_constexpr(VHE_CHUNKS): + v_vh = v_vh_batch[vh] + acc = fx.Vector.filled(MFMA_ACC_ELEMS, 0.0, fx.Float32) + acc = _mfma_fp8(v_vh, p_ops, 0, 0, NVOPS, acc) + op = fx.Vector(acc) + if const_expr(per_token_kv): + op = op * fx.Vector.from_elements( + [v_max_scaled], dtype=fx.Float32 + ).broadcast_to(OP_ELEMS) + o_acc[vh] = ( + op if const_expr(single_tile_plan) else o_acc[vh] * corr_b + op + ) + next_state.extend([*o_acc, m_new, l_new]) + if const_expr(MTP4_PREFETCH_V): + next_state.append(v_next) + results = yield next_state + o_final = results + + # Each lane stores four head-dim values for one query row. + inv_fp8 = fx.Float32(1.0 / FP8_MAX) + for m in range_constexpr(M_TILES): + row = m * MFMA_MNK + lane16 # flat (mtp, gqa) query-row for this lane + global_row = query_begin * query_group_size + row + qi_e = row // query_group_size + gs_head_e = row - qi_e * query_group_size + qh = kv_h * query_group_size + gs_head_e + l_row = o_final[_l_slot(m)] + safe_l = (l_row > ZERO_F).select(l_row, fx.Float32(1.0)) + inv_l = fx.Float32(rcp_f32(safe_l)) + if const_expr(DIRECT_SINKS): + # Sinks affect only the denominator. Compare in natural-logit + # units before LOG2E scaling to avoid overflow for finite sinks. + sink_value = fx.Float32(sink_token[qh]) + row_max = o_final[_m_slot(m)] * fx.Float32(1.0 / LOG2E) + total_max = fx.maxnumf(row_max, sink_value) + safe_max = (total_max > NEG_INF).select(total_max, ZERO_F) + kv_mass = (l_row > ZERO_F).select( + fx.exp2((row_max - safe_max) * fx.Float32(LOG2E), fastmath="fast"), + ZERO_F, + ) + # +inf suppresses all KV output; -inf disables this head's + # sink. Select before exp2 to avoid the +inf - +inf case. + sink_shift = (sink_value == safe_max).select( + ZERO_F, sink_value - safe_max + ) + sink_mass = fx.exp2(sink_shift * fx.Float32(LOG2E), fastmath="fast") + denominator = l_row * kv_mass + sink_mass + safe_denominator = (denominator > ZERO_F).select( + denominator, fx.Float32(1.0) + ) + inv_l = kv_mass * fx.Float32(rcp_f32(safe_denominator)) + if const_expr(per_token_kv): + o_scale = inv_l + elif const_expr(SCALAR_FP8_DECODE): + # Direct P conversion needs no compensating 1/FP8_MAX. + o_scale = inv_l * v_scale_f + else: + o_scale = inv_l * (v_scale_f * inv_fp8) + o_scale_b = fx.Vector.from_elements( + [o_scale], dtype=fx.Float32 + ).broadcast_to(OP_ELEMS) + + def _emit(o_norm, sub, query_idx, query_head): + if const_expr(NP == 1 and not use_work_plan): + out_offset = ( + (seq * query_length + query_begin + query_idx) * stride_o_row + + query_head * stride_o_head + + sub * OP_ELEMS + ) + fx.ptr_store(o_norm, fx.add_offset(output, out_offset)) + elif const_expr(buffer_plan_output): + # Row guards keep the full 8-byte store within this slot. + pout_offset = global_row * head_dim + sub * OP_ELEMS # noqa: B023 + buf_copy_store( + pout_buffer, + pout_offset, + o_norm, + elem=Q_DTYPE, + unit_elems=OP_ELEMS, + cache_modifier=0, + ) + else: + base = partial_slot * TOTAL_ROWS + global_row # noqa: B023 + pout_offset = base * head_dim + sub * OP_ELEMS + fx.ptr_store( + o_norm, + fx.add_offset(pout, pout_offset), + ) + + for vh in range_constexpr(VHE_CHUNKS): + o_slot = _o_slot(m, vh) + o_norm = (o_final[o_slot] * o_scale_b).to(Q_DTYPE) + head_base = ( + vh * (NWARP * MFMA_MNK) + warp * MFMA_MNK + rgroup * OP_ELEMS + ) + sub = head_base // OP_ELEMS + if row < CTA_ROWS: + _emit(o_norm, sub, qi_e, qh) + + if const_expr(NP > 1 or use_work_plan): # noqa: SIM102 + if warp == 0 and rgroup == 0: + base = partial_slot * TOTAL_ROWS + global_row + if row < CTA_ROWS: + # The shared reducer expects maxima in natural-log units. + pmax[base] = o_final[_m_slot(m)] * fx.Float32(1.0 / LOG2E) + psum[base] = l_row + + @flyc.kernel(known_block_size=(BLOCK_THREADS, 1, 1)) + def pa_decode_tile_kernel( + output_ptr: fx.Pointer, + pmax_ptr: fx.Pointer, + psum_ptr: fx.Pointer, + pout_ptr: fx.Pointer, + query_ptr: fx.Pointer, + key_cache_ptr: fx.Pointer, + value_cache_ptr: fx.Pointer, + block_tables_ptr: fx.Pointer, + context_lengths_ptr: fx.Pointer, + key_scale_ptr: fx.Pointer, + value_scale_ptr: fx.Pointer, + sinks_ptr: fx.Pointer, + max_blocks_per_seq: fx.Int32, + stride_ks_block: fx.Int32, + stride_ks_head: fx.Int32, + stride_o_row: fx.Int32, + stride_o_head: fx.Int32, + stride_q_row: fx.Int32, + stride_q_head: fx.Int32, + work_info_ptr: fx.Pointer, + num_sequences: fx.Int32, + ): + def _run_task(seq, start, end, context): + _pa_decode_tile_task( + output_ptr, + pmax_ptr, + psum_ptr, + pout_ptr, + query_ptr, + key_cache_ptr, + value_cache_ptr, + block_tables_ptr, + context_lengths_ptr, + key_scale_ptr, + value_scale_ptr, + sinks_ptr, + max_blocks_per_seq, + stride_ks_block, + stride_ks_head, + stride_o_row, + stride_o_head, + stride_q_row, + stride_q_head, + num_sequences, + seq, + start, + end, + context, + ) + + if const_expr(use_work_plan): + # Skip cleared padding records: no scratch writes or Q loads. + # This CTA-uniform guard keeps every task barrier convergent. + if const_expr(batch_first_plan_grid): + slot = fx.Int32( + fx.Uint32(gpu.block_id("x")) * fx.Uint32(gpu.grid_dim.z) + + fx.Uint32(gpu.block_id("z")) + ) + else: + slot = fx.Int32(gpu.block_id("x")) + work = fx.recast_iter(fx.Int32, work_info_ptr) + task = fx.ptr_load( + fx.add_offset(work, slot * 4), + result_type=fx.Vector.make_type(4, fx.Int32), + ) + start = fx.Int32(task[1]) + end = fx.Int32(task[2]) + if start < end: + _run_task(fx.Int32(task[0]), start, end, fx.Int32(task[3])) + else: + zero = fx.Int32(0) + _run_task(zero, zero, zero, zero) + + @flyc.jit + def pa_decode_tile_launch( + output: fx.Pointer, + pmax: fx.Pointer, + psum: fx.Pointer, + pout: fx.Pointer, + query: fx.Pointer, + key_cache: fx.Pointer, + value_cache: fx.Pointer, + block_tables: fx.Pointer, + context_lengths: fx.Pointer, + key_scale: fx.Pointer, + value_scale: fx.Pointer, + sinks: fx.Pointer, + max_blocks_per_seq: fx.Int32, + num_seqs: fx.Int32, + num_kv_heads: fx.Int32, + stride_ks_block: fx.Int32, + stride_ks_head: fx.Int32, + stride_o_row: fx.Int32, + stride_o_head: fx.Int32, + stride_q_row: fx.Int32, + stride_q_head: fx.Int32, + work_info: fx.Pointer, + work_capacity: fx.Int32, + stream: fx.Stream = fx.Stream(None), # noqa: B008 + ): + # Ambient contract permits FMAs; explicit per-op fastmath still wins. + with CompilationContext.compile_hints({"fastmath": "contract"}): + pa_decode_tile_kernel( + output, + pmax, + psum, + pout, + query, + key_cache, + value_cache, + block_tables, + context_lengths, + key_scale, + value_scale, + sinks, + max_blocks_per_seq, + stride_ks_block, + stride_ks_head, + stride_o_row, + stride_o_head, + stride_q_row, + stride_q_head, + work_info, + num_seqs, + ).launch( + grid=( + (num_seqs, num_kv_heads * query_splits, work_capacity // num_seqs) + if batch_first_plan_grid + else ( + work_capacity if use_work_plan else num_seqs, + num_kv_heads * query_splits, + 1 if use_work_plan else NP, + ) + ), + block=(BLOCK_THREADS, 1, 1), + stream=stream, + ) + + compiled = { + "launch": pa_decode_tile_launch, + "kernel": pa_decode_tile_kernel, + } + return _PA_DECODE_TILE_CACHE.setdefault(cache_key, compiled) diff --git a/aiter/ops/flydsl/kernels/pa_decode_plan.py b/aiter/ops/flydsl/kernels/pa_decode_plan.py new file mode 100644 index 00000000000..d0cf319b862 --- /dev/null +++ b/aiter/ops/flydsl/kernels/pa_decode_plan.py @@ -0,0 +1,206 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (C) 2026 FlyDSL Project Contributors + +"""GPU work planning for FlyDSL paged attention with mixed context lengths.""" + +from dataclasses import dataclass + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _plan_pa_decode( + lengths, + work, + reduce_info, + B: tl.constexpr, + CAPACITY: tl.constexpr, + MAX_PARTS: tl.constexpr, + SLIDING_WINDOW: tl.constexpr, + QUERY_LENGTH: tl.constexpr, + BLOCK_B: tl.constexpr, + BLOCK_P: tl.constexpr, +): + seq = tl.program_id(0) + b = tl.arange(0, BLOCK_B) + ctx = tl.maximum(tl.load(lengths + b, b < B, other=0), 0).to(tl.int64) + first_tile = tl.full((BLOCK_B,), 0, tl.int64) + if SLIDING_WINDOW > 0: + # Absolute tiles cover the MTP window union without rebasing cache offsets. + first_tile = tl.maximum(ctx - (QUERY_LENGTH - 1) - SLIDING_WINDOW, 0) // 256 + tiles = tl.where(b < B, (ctx + 255) // 256 - first_tile, 0) + nonempty = (tiles > 0).to(tl.int32) + total = tl.maximum(tl.sum(tiles, 0), 1) + remaining = CAPACITY - tl.sum(nonempty, 0) + cumulative = tl.cumsum(tiles, 0) + # Reserve one task per nonempty row, then apportion the rest using integer + # prefixes so rounding cannot overfill the budget. + upper = cumulative * remaining // total + lower = (cumulative - tiles) * remaining // total + counts = tl.minimum(tl.minimum(nonempty + upper - lower, tiles), MAX_PARTS) + count = tl.sum(tl.where(b == seq, counts, 0), 0).to(tl.int32) + start = tl.sum(tl.where(b < seq, counts, 0), 0).to(tl.int32) + seq_ctx = tl.load(lengths + seq) + seq_length = tl.maximum(seq_ctx, 0).to(tl.int64) + seq_first_tile = tl.full((), 0, tl.int64) + if SLIDING_WINDOW > 0: + seq_first_tile = ( + tl.maximum(seq_length - (QUERY_LENGTH - 1) - SLIDING_WINDOW, 0) // 256 + ) + seq_tiles = (seq_length + 255) // 256 - seq_first_tile + total_tasks = tl.sum(counts, 0).to(tl.int32) + tl.store(reduce_info + seq * 2, start) + tl.store(reduce_info + seq * 2 + 1, count) + + part = tl.arange(0, BLOCK_P) + active = part < count + begin = seq_first_tile + part.to(tl.int64) * seq_tiles // tl.maximum(count, 1) + end = seq_first_tile + (part.to(tl.int64) + 1) * seq_tiles // tl.maximum(count, 1) + slot = start + part + tl.store(work + slot * 4, seq, active) + tl.store(work + slot * 4 + 1, begin, active) + tl.store(work + slot * 4 + 2, end, active) + tl.store(work + slot * 4 + 3, seq_ctx, active) + + # Zero unused slots for fixed-capacity graph launches. Padding stores must + # stay disjoint from active records, including for all-empty batches. + pad = seq * BLOCK_P + part + padding = (pad >= total_tasks) & (pad < CAPACITY) + for field in tl.static_range(4): + tl.store(work + pad * 4 + field, 0, padding) + + +@dataclass(frozen=True) +class PADecodePlan: + """Reusable GPU metadata; refresh it whenever context lengths change. + + ``work_info`` is [capacity, 4]: sequence, begin/end absolute 256-token tile, + original context length. Tile ranges are half-open and cover the MTP union. + ``reduce_info`` is [batch, 2]: first packed task, actual task count. + """ + + work_info: torch.Tensor + reduce_info: torch.Tensor + num_kv_heads: int + max_partitions: int + sliding_window: int = 0 + query_length: int = 1 + + @property + def capacity(self) -> int: + return self.work_info.shape[0] + + def validate(self, batch_size: int, num_kv_heads: int, device: torch.device): + if self.num_kv_heads != num_kv_heads: + raise ValueError("plan KV head count does not match the cache") + num_compute_units = torch.cuda.get_device_properties( + device + ).multi_processor_count + if not 1 <= self.max_partitions <= num_compute_units: + raise ValueError(f"plan max_partitions must be in [1, {num_compute_units}]") + if self.sliding_window < 0 or self.query_length < 1: + raise ValueError("invalid plan sliding_window or query_length") + if self.reduce_info.shape != (batch_size, 2): + raise ValueError("reduce_info must have shape [batch_size, 2]") + if self.work_info.ndim != 2 or self.work_info.shape[1] != 4: + raise ValueError("work_info must have shape [capacity, 4]") + if not batch_size <= self.capacity <= batch_size * self.max_partitions: + raise ValueError("plan capacity must be in [batch, batch * max_partitions]") + for tensor in (self.work_info, self.reduce_info): + if tensor.device != device or tensor.dtype != torch.int32: + raise ValueError("plan metadata must be int32 on the query device") + if not tensor.is_contiguous(): + raise ValueError("plan metadata must be contiguous") + + +def plan_pa_decode( + context_lengths: torch.Tensor, + num_kv_heads: int, + *, + max_partitions: int | None = None, + workgroup_budget: int | None = None, + sliding_window: int = 0, + query_length: int = 1, + plan: PADecodePlan | None = None, +) -> PADecodePlan: + """Build/refresh GPU work metadata on the current stream without readback. + + Allocate outside graph capture; pass ``plan`` to refresh buffers in place. + The budget counts task slots across KV heads before query splitting. + + ``max_partitions`` is in [1, device CU count], defaulting to CU count; + omitting it on refresh preserves the existing limit. + + Lengths include MTP tokens. Positive windows include each query's position; + 0 and -1 disable them. Plans cover the MTP window union: match the window + and query length when decoding or refreshing. Dense plans ignore query length. + """ + if not isinstance(sliding_window, int): + raise TypeError("sliding_window must be an int") + if sliding_window < -1: + raise ValueError("sliding_window must be -1, 0, or positive") + sliding_window = max(sliding_window, 0) + if not isinstance(query_length, int): + raise TypeError("query_length must be an int") + if query_length < 1: + raise ValueError("query_length must be positive") + if context_lengths.device.type != "cuda" or context_lengths.dtype != torch.int32: + raise ValueError("context_lengths must be a CUDA int32 tensor") + if context_lengths.ndim != 1 or not context_lengths.is_contiguous(): + raise ValueError("context_lengths must be a contiguous vector") + batch = context_lengths.numel() + if batch < 1 or batch > 4096: + raise ValueError("plan supports batches in [1, 4096]") + if num_kv_heads < 1: + raise ValueError("num_kv_heads must be positive") + dev = context_lengths.device + num_compute_units = torch.cuda.get_device_properties(dev).multi_processor_count + if max_partitions is None: + max_partitions = num_compute_units if plan is None else plan.max_partitions + if not 1 <= max_partitions <= num_compute_units: + raise ValueError(f"max_partitions must be in [1, {num_compute_units}]") + if plan is None: + if workgroup_budget is None: + workgroup_budget = 2 * num_compute_units + if workgroup_budget < 1: + raise ValueError("workgroup_budget must be positive") + capacity = min( + batch * max_partitions, + max(batch, (workgroup_budget + num_kv_heads - 1) // num_kv_heads), + ) + plan = PADecodePlan( + torch.empty((capacity, 4), dtype=torch.int32, device=dev), + torch.empty((batch, 2), dtype=torch.int32, device=dev), + num_kv_heads, + max_partitions, + sliding_window, + query_length, + ) + else: + if workgroup_budget is not None: + raise ValueError("workgroup_budget is fixed when reusing a plan") + if max_partitions != plan.max_partitions: + raise ValueError("max_partitions must match the reused plan") + if sliding_window != plan.sliding_window: + raise ValueError("sliding_window must match the reused plan") + if sliding_window > 0 and query_length != plan.query_length: + raise ValueError("query_length must match the reused plan") + plan.validate(batch, num_kv_heads, dev) + with torch.cuda.device(dev): + _plan_pa_decode[(batch,)]( + context_lengths, + plan.work_info, + plan.reduce_info, + batch, + plan.capacity, + plan.max_partitions, + # Clamp windows to int32 lengths; dense plans share one specialization. + min(sliding_window, 2**31 - 1), + query_length if sliding_window > 0 else 1, + triton.next_power_of_2(batch), + triton.next_power_of_2(plan.max_partitions), + num_warps=4, + ) + return plan diff --git a/aiter/ops/flydsl/kernels/pa_decode_reduce.py b/aiter/ops/flydsl/kernels/pa_decode_reduce.py new file mode 100644 index 00000000000..1c1f6738c7c --- /dev/null +++ b/aiter/ops/flydsl/kernels/pa_decode_reduce.py @@ -0,0 +1,869 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (C) 2025-2026 FlyDSL Project Contributors + +"""Partitioned-softmax reduction kernel for FlyDSL paged attention.""" + +from functools import lru_cache + +import flydsl.compiler as flyc +import flydsl.expr as fx +from flydsl.expr.typing import T + +from .tensor_shim import buf_base_i64, buf_copy_load, ptr_buf_tensor +from .utils import rcp_f32, udiv_const, urem_const + +# Static scheduling limit; work plans use their device's CU count instead. +MAX_CONTEXT_PARTITIONS = 256 +_DTYPE_MAP = { + "f32": fx.Float32, + "f16": fx.Float16, + "bf16": fx.BFloat16, +} + + +def _validate_pa_decode_ps_reduce_config( + *, + max_context_partition_num: int, + head_size: int, + output_dtype_str: str, + logits_dtype_str: str, + sink_dtype_str: str, + use_work_plan: bool = False, +) -> None: + # The plan and launch wrappers validate the device-specific CU bound. + if max_context_partition_num < 1 or ( + not use_work_plan and max_context_partition_num > MAX_CONTEXT_PARTITIONS + ): + bound = "positive" if use_work_plan else f"in [1, {MAX_CONTEXT_PARTITIONS}]" + raise ValueError( + f"max_context_partition_num must be {bound}, " + f"got {max_context_partition_num}" + ) + if head_size <= 0 or head_size > 1024 or head_size % 64: + raise ValueError( + f"head_size must be a multiple of 64 in [64, 1024], got {head_size}" + ) + for dtype_str in (output_dtype_str, logits_dtype_str, sink_dtype_str): + if dtype_str not in _DTYPE_MAP: + raise ValueError(f"Unsupported FlyDSL dtype: {dtype_str!r}") + + +def is_pa_decode_ps_reduce_supported( + *, + max_context_partition_num: int, + head_size: int, + output_dtype_str: str, + logits_dtype_str: str, + sink_dtype_str: str, + use_work_plan: bool = False, +) -> bool: + """Return whether the FlyDSL reducer supports a dispatch configuration.""" + try: + _validate_pa_decode_ps_reduce_config( + max_context_partition_num=max_context_partition_num, + head_size=head_size, + output_dtype_str=output_dtype_str, + logits_dtype_str=logits_dtype_str, + sink_dtype_str=sink_dtype_str, + use_work_plan=use_work_plan, + ) + except ValueError: + return False + return True + + +@lru_cache(maxsize=256) +def compile_pa_decode_ps_reduce( + *, + max_context_partition_num: int, + head_size: int, + output_dtype_str: str, + logits_dtype_str: str, + sink_dtype_str: str, + use_sinks: bool, + use_work_plan: bool = False, + query_group_size: int | None = None, + bounded_plan_logits: bool = False, + vectorize_plan_logits: bool = False, +): + """Build PA decode's partitioned-softmax reducer. + + D128/NP>64 shares weights in LDS; other shapes stay register-only. + Sinks are per-head zero-value logits: include them in the shared maximum + for stability and add their denominator mass once after summing KV. + ``query_group_size`` is None for runtime GQA, or a matching positive value. + ``bounded_plan_logits`` requires planned NP<=64, packed contiguous rows + and a positive NP-part byte span <= INT32_MAX. Per-sequence bounds + zero-fill inactive loads without reading adjacent sequences. + ``vectorize_plan_logits`` also requires D128 bf16/f16 logits, a four-byte + aligned base and even head/part/group strides. Output stores retain + scalar alignment. + """ + _validate_pa_decode_ps_reduce_config( + max_context_partition_num=max_context_partition_num, + head_size=head_size, + output_dtype_str=output_dtype_str, + logits_dtype_str=logits_dtype_str, + sink_dtype_str=sink_dtype_str, + use_work_plan=use_work_plan, + ) + if query_group_size is not None and query_group_size <= 0: + raise ValueError("query_group_size must be positive when specified") + if bounded_plan_logits and (not use_work_plan or max_context_partition_num > 64): + raise ValueError("bounded_plan_logits requires a planned NP<=64 reducer") + if vectorize_plan_logits and ( + not bounded_plan_logits + or head_size != 128 + or logits_dtype_str not in ("bf16", "f16") + ): + raise ValueError( + "vectorize_plan_logits requires bounded planned D=128 bf16/f16 logits" + ) + static_query_group_size = query_group_size + planned_elements_per_thread = 2 if vectorize_plan_logits else 1 + + output_dtype = _DTYPE_MAP[output_dtype_str] + logits_dtype = _DTYPE_MAP[logits_dtype_str] + sink_dtype = _DTYPE_MAP[sink_dtype_str] + + warp_size = 64 + log2e = 1.4426950408889634 + reduce_width = ( + 1 + if max_context_partition_num == 1 + else 1 << ((max_context_partition_num - 1).bit_length()) + ) + reduce_shuffle_offsets = [ + offset for offset in (32, 16, 8, 4, 2, 1) if offset < reduce_width + ] + + # D128 splits partition ranges across wave pairs, one wave per 64 outputs. + use_parallel_lds = head_size == 128 and max_context_partition_num > warp_size + parallel_groups = 1 + if use_parallel_lds: + parallel_groups = 2 if max_context_partition_num <= 96 else 8 + head_waves = head_size // warp_size + worker_waves = head_waves * parallel_groups + block_shape = ( + [warp_size, worker_waves, 1] + if use_parallel_lds + else [head_size // planned_elements_per_thread, 1, 1] + ) + parts_per_group = ( + max_context_partition_num + parallel_groups - 1 + ) // parallel_groups + + # SharedStorage is allocated only by the parallel branch. + shared_weight_elems = max_context_partition_num if use_parallel_lds else 1 + shared_partial_elems = (parallel_groups - 1) * head_size if use_parallel_lds else 1 + + @fx.struct + class SharedStorage: + weights: fx.Array[fx.Float32, shared_weight_elems, 16] + partials: fx.Array[fx.Float32, shared_partial_elems, 16] + + @flyc.kernel(known_block_size=block_shape) + def pa_decode_ps_reduce_kernel( + output_ptr: fx.Pointer, + exp_sums_ptr: fx.Pointer, + max_logits_ptr: fx.Pointer, + logits_ptr: fx.Pointer, + sink_token_ptr: fx.Pointer, + 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, + query_group_size: fx.Int32, + reduce_info_ptr: fx.Pointer, + ): + tid = fx.thread_idx.x + worker = fx.thread_idx.y + batch_idx = fx.block_idx.x + kv_head_idx = fx.block_idx.y + eqgs_idx = fx.block_idx.z + + output = fx.recast_iter(output_dtype, output_ptr) + exp_sums = fx.recast_iter(fx.Float32, exp_sums_ptr) + max_logits = fx.recast_iter(fx.Float32, max_logits_ptr) + logits = fx.recast_iter(logits_dtype, logits_ptr) + if fx.const_expr(use_sinks): + sink_token = fx.recast_iter(sink_dtype, sink_token_ptr) + + zero_f = fx.Float32(0.0) + one_f = fx.Float32(1.0) + neg_inf = fx.Float32(float("-inf")) + c_log2e = fx.Float32(log2e) + zero_i = fx.Int32(0) + c_warp_size = fx.Int32(warp_size) + c_wave_mask = fx.Int32(warp_size - 1) + c_part_num = fx.Int32(max_context_partition_num) + if fx.const_expr(use_work_plan): + reduce_info = fx.recast_iter(fx.Int32, reduce_info_ptr) + first_part = fx.Int32(reduce_info[batch_idx * 2]) + c_part_num = fx.Int32(reduce_info[batch_idx * 2 + 1]) + # Empty rows may point past packed storage; clamped loads need slot 0. + # Mask every inactive value before arithmetic. + first_part = (c_part_num > zero_i).select(first_part, zero_i) + stats_seq_offset = first_part * stride_exp_sums_part + logits_seq_offset = first_part * stride_logits_part + else: + stats_seq_offset = batch_idx * stride_exp_sums_seq + logits_seq_offset = batch_idx * stride_logits_seq + if fx.const_expr(bounded_plan_logits): + # Rebase in i64 before narrowing: packed storage can exceed 2 GiB. + logits_item_bytes = logits_dtype.width // 8 + planned_logits_base = ( + buf_base_i64(logits_ptr) + + ( + fx.Int64(first_part) * fx.Int64(stride_logits_part) + + fx.Int64(kv_head_idx) * fx.Int64(stride_logits_head) + ) + * logits_item_bytes + ) + planned_logits_bytes = ( + fx.Int64(c_part_num) * fx.Int64(stride_logits_part) * logits_item_bytes + ) + planned_logits_buffer = ptr_buf_tensor( + planned_logits_base, + logits_dtype, + unit_elems=planned_elements_per_thread, + unit_stride=1, + num_records_bytes=planned_logits_bytes, + ) + c_reduce_width = fx.Int32(reduce_width) + c_four = fx.Int32(4) + lane = tid & c_wave_mask + if fx.const_expr(static_query_group_size is not None): + c_qgs = fx.Int32(static_query_group_size) + group_idx = urem_const(eqgs_idx, static_query_group_size) + else: + c_qgs = query_group_size + group_idx = eqgs_idx % c_qgs + + def _wave_reduce_max(value, offsets=reduce_shuffle_offsets): + reduced = value + for offset in offsets: + reduced = reduced.maximumf( + reduced.shuffle_xor(fx.Int32(offset), c_warp_size) + ) + return reduced + + def _wave_reduce_sum(value, offsets=reduce_shuffle_offsets): + reduced = value + for offset in offsets: + reduced = fx.Float32( + reduced.addf( + reduced.shuffle_xor(fx.Int32(offset), c_warp_size), + fastmath="fast", + ) + ) + return reduced + + def _sink_exp(sink_value, safe_max): + # +inf sinks have unit mass at the shared max and suppress KV; + # -inf contributes zero, including for an entirely empty request. + shift = (sink_value == safe_max).select(zero_f, sink_value - safe_max) + return fx.exp2(shift * c_log2e, fastmath="fast") + + def _reduce_planned_wave(part_count, shuffle_offsets): + # Clamp invalid lanes to a valid slot, then mask loaded NaNs + # before arithmetic; empty rows use global slot 0. + planned_lane_in_range = lane < c_part_num + planned_stats_part_idx = planned_lane_in_range.select(lane, zero_i) + planned_stats_offset = ( + stats_seq_offset + + kv_head_idx * stride_exp_sums_head + + planned_stats_part_idx * stride_exp_sums_part + + eqgs_idx + ) + planned_loaded_sum = fx.Float32(exp_sums[planned_stats_offset]) + planned_loaded_max = fx.Float32(max_logits[planned_stats_offset]) + planned_part_sum = planned_lane_in_range.select(planned_loaded_sum, zero_f) + planned_part_max = planned_lane_in_range.select(planned_loaded_max, neg_inf) + + planned_logits = [] + for planned_part in fx.range_constexpr(part_count): + planned_part_idx = fx.Int32(planned_part) + if fx.const_expr(bounded_plan_logits): + # Bound loads before adjacent sequences/NaN padding; + # count=0 gives a zero-sized descriptor that returns zero. + planned_logits_offset = ( + planned_part_idx * stride_logits_part + + eqgs_idx * stride_logits_group + + tid * fx.Int32(planned_elements_per_thread) + ) + # unit_stride=1 uses element offsets; D128 pairs stay in-row. + # A zero-sized descriptor suppresses the entire dword load. + planned_loaded_logits = buf_copy_load( + planned_logits_buffer, + planned_logits_offset, + logits_dtype, + unit_elems=planned_elements_per_thread, + ) + if fx.const_expr(vectorize_plan_logits): + planned_logits.append(planned_loaded_logits) + else: + planned_logits.append(fx.Float32(planned_loaded_logits)) + else: + planned_part_in_range = planned_part_idx < c_part_num + planned_safe_part_idx = planned_part_in_range.select( + planned_part_idx, zero_i + ) + planned_logits_offset = ( + logits_seq_offset + + kv_head_idx * stride_logits_head + + planned_safe_part_idx * stride_logits_part + + eqgs_idx * stride_logits_group + + tid + ) + planned_loaded_logits = fx.Float32(logits[planned_logits_offset]) + planned_logits.append( + planned_part_in_range.select(planned_loaded_logits, zero_f) + ) + + planned_global_max = _wave_reduce_max(planned_part_max, shuffle_offsets) + if fx.const_expr(use_sinks): + planned_sink_value = fx.Float32( + sink_token[kv_head_idx * c_qgs + group_idx] + ) + planned_global_max = planned_global_max.maximumf(planned_sink_value) + planned_safe_max = (planned_global_max > neg_inf).select( + planned_global_max, zero_f + ) + planned_part_scale = (planned_part_max > neg_inf).select( + fx.exp2( + (planned_part_max - planned_safe_max) * c_log2e, + fastmath="fast", + ), + zero_f, + ) + planned_scaled_sum = planned_part_sum * planned_part_scale + planned_exp_sum = _wave_reduce_sum(planned_scaled_sum, shuffle_offsets) + if fx.const_expr(use_sinks): + planned_exp_sum = planned_exp_sum + _sink_exp( + planned_sink_value, planned_safe_max + ) + planned_safe_exp_sum = (planned_exp_sum > zero_f).select( + planned_exp_sum, one_f + ) + planned_inv_exp_sum = fx.Float32(rcp_f32(planned_safe_exp_sum)) + planned_weight_local_i32 = ( + planned_scaled_sum * planned_inv_exp_sum + ).bitcast(fx.Int32) + + if fx.const_expr(vectorize_plan_logits): + planned_acc = fx.Vector.filled(2, 0.0, fx.Float32) + else: + planned_acc = zero_f + for planned_acc_part in fx.range_constexpr(part_count): + planned_acc_part_idx = fx.Int32(planned_acc_part) + planned_weight_i32 = fx.Int32( + fx.rocdl.ds_bpermute( + T.i32, + planned_acc_part_idx * c_four, + planned_weight_local_i32, + ) + ) + planned_weight = planned_weight_i32.bitcast(fx.Float32) + if fx.const_expr(vectorize_plan_logits): + planned_part_logits = planned_logits[planned_acc_part].to( + fx.Float32 + ) + else: + planned_part_logits = planned_logits[planned_acc_part] + # Preserve increasing-part order and separate multiply/add rounding. + planned_acc = planned_acc + planned_part_logits * planned_weight + return planned_acc + + if fx.const_expr(use_parallel_lds): + # Worker 0 publishes weights to LDS. All threads, including inactive + # groups, must reach both barriers; group 0 reads only active partials. + lds = fx.SharedAllocator().allocate(SharedStorage).peek() + lds_weights = lds.weights + lds_partials = lds.partials + + if worker == zero_i: + partitions_per_lane = ( + max_context_partition_num + warp_size - 1 + ) // warp_size + part_sums = [] + part_maxes = [] + lane_max = neg_inf + for chunk_idx in fx.range_constexpr(partitions_per_lane): + chunk_base = chunk_idx * warp_size + chunk_size = min(warp_size, max_context_partition_num - chunk_base) + part_idx = lane + fx.Int32(chunk_base) + if fx.const_expr(chunk_size == warp_size and not use_work_plan): + stats_offset = ( + stats_seq_offset + + kv_head_idx * stride_exp_sums_head + + part_idx * stride_exp_sums_part + + eqgs_idx + ) + part_sum = fx.Float32(exp_sums[stats_offset]) + part_max = fx.Float32(max_logits[stats_offset]) + else: + lane_in_range = (lane < fx.Int32(chunk_size)) & ( + part_idx < c_part_num + ) + stats_offset = ( + stats_seq_offset + + kv_head_idx * stride_exp_sums_head + + part_idx * stride_exp_sums_part + + eqgs_idx + ) + part_sum = zero_f + part_max = neg_inf + if lane_in_range: + part_sum = fx.Float32(exp_sums[stats_offset]) + part_max = fx.Float32(max_logits[stats_offset]) + part_sums.append(part_sum) + part_maxes.append(part_max) + lane_max = lane_max.maximumf(part_max) + + global_max = _wave_reduce_max(lane_max) + if fx.const_expr(use_sinks): + sink_value = fx.Float32(sink_token[kv_head_idx * c_qgs + group_idx]) + global_max = global_max.maximumf(sink_value) + safe_global_max = (global_max > neg_inf).select(global_max, zero_f) + scaled_sums = [] + lane_exp_sum = zero_f + for chunk_idx in fx.range_constexpr(partitions_per_lane): + part_max = part_maxes[chunk_idx] + if fx.const_expr(use_work_plan): + # select evaluates exp2 eagerly; branch around inactive parts. + part_scale = zero_f + if part_max > neg_inf: + part_scale = fx.exp2( + (part_max - safe_global_max) * c_log2e, + fastmath="fast", + ) + else: + part_scale = (part_max > neg_inf).select( + fx.exp2( + (part_max - safe_global_max) * c_log2e, + fastmath="fast", + ), + zero_f, + ) + scaled_sum = part_sums[chunk_idx] * part_scale + scaled_sums.append(scaled_sum) + lane_exp_sum = lane_exp_sum + scaled_sum + + global_exp_sum = _wave_reduce_sum(lane_exp_sum) + if fx.const_expr(use_sinks): + sink_scale = _sink_exp(sink_value, safe_global_max) + global_exp_sum = global_exp_sum + sink_scale + safe_global_exp_sum = (global_exp_sum > zero_f).select( + global_exp_sum, one_f + ) + parallel_inv_exp_sum = fx.Float32(rcp_f32(safe_global_exp_sum)) + + for chunk_idx in fx.range_constexpr(partitions_per_lane): + chunk_base = chunk_idx * warp_size + chunk_size = min(warp_size, max_context_partition_num - chunk_base) + part_idx = lane + fx.Int32(chunk_base) + if fx.const_expr(use_work_plan): + lane_in_range = (lane < fx.Int32(chunk_size)) & ( + part_idx < c_part_num + ) + if lane_in_range: + weight = scaled_sums[chunk_idx] * parallel_inv_exp_sum + lds_weights[part_idx] = weight + else: + weight = scaled_sums[chunk_idx] * parallel_inv_exp_sum + if fx.const_expr(chunk_size == warp_size): + lds_weights[part_idx] = weight + else: + lane_in_range = lane < fx.Int32(chunk_size) + if lane_in_range: + lds_weights[part_idx] = weight + + fx.gpu.barrier() + + head_wave = worker % fx.Int32(head_waves) + partition_group = worker // fx.Int32(head_waves) + output_element = head_wave * c_warp_size + lane + group_part_begin = partition_group * fx.Int32(parts_per_group) + acc = zero_f + if fx.const_expr(use_work_plan): + group_in_range = group_part_begin < c_part_num + if group_in_range: + for local_part in fx.range_constexpr(parts_per_group): + part_idx = group_part_begin + fx.Int32(local_part) + if part_idx < c_part_num: + weight = fx.Float32(lds_weights[part_idx]) + logits_offset = ( + logits_seq_offset + + kv_head_idx * stride_logits_head + + part_idx * stride_logits_part + + eqgs_idx * stride_logits_group + + output_element + ) + part_logits = fx.Float32(logits[logits_offset]) + acc = acc + part_logits * weight + else: + for local_part in fx.range_constexpr(parts_per_group): + part_idx = group_part_begin + fx.Int32(local_part) + part_in_range = part_idx < c_part_num + if part_in_range: + weight = fx.Float32(lds_weights[part_idx]) + logits_offset = ( + logits_seq_offset + + kv_head_idx * stride_logits_head + + part_idx * stride_logits_part + + eqgs_idx * stride_logits_group + + output_element + ) + part_logits = fx.Float32(logits[logits_offset]) + acc = acc + part_logits * weight + + if partition_group > zero_i: + if fx.const_expr(use_work_plan): + if group_in_range: + partial_offset = (partition_group - fx.Int32(1)) * fx.Int32( + head_size + ) + output_element + lds_partials[partial_offset] = acc + else: + partial_offset = (partition_group - fx.Int32(1)) * fx.Int32( + head_size + ) + output_element + lds_partials[partial_offset] = acc + + fx.gpu.barrier() + + if partition_group == zero_i: + for other_group in fx.range_constexpr(1, parallel_groups): + if fx.const_expr(use_work_plan): + if fx.Int32(other_group * parts_per_group) < c_part_num: + partial_offset = ( + fx.Int32((other_group - 1) * head_size) + output_element + ) + acc = acc + fx.Float32(lds_partials[partial_offset]) + else: + partial_offset = ( + fx.Int32((other_group - 1) * head_size) + output_element + ) + acc = acc + fx.Float32(lds_partials[partial_offset]) + + elif fx.const_expr(max_context_partition_num <= warp_size): + if fx.const_expr(use_work_plan): + # Count branches must carry the same scalar/vector type. + if fx.const_expr(vectorize_plan_logits): + acc = fx.Vector.filled(2, 0.0, fx.Float32) + else: + acc = zero_f + if fx.const_expr(max_context_partition_num > 8): + # Count is CTA-uniform; keep short-path loads and math inside it. + if c_part_num <= c_four: + acc = _reduce_planned_wave(4, (2, 1)) + elif c_part_num <= fx.Int32(8): + acc = _reduce_planned_wave(8, (4, 2, 1)) + else: + acc = _reduce_planned_wave( + max_context_partition_num, reduce_shuffle_offsets + ) + else: + acc = _reduce_planned_wave( + max_context_partition_num, reduce_shuffle_offsets + ) + else: + if fx.const_expr(max_context_partition_num == reduce_width): + # Full subgroups have no inactive lanes; tails use EXEC masks. + lane_in_range = lane < c_part_num + lane_in_reduce = lane < c_reduce_width + part_sum = zero_f + part_max = neg_inf + if lane_in_reduce: + part_idx = lane_in_range.select(lane, zero_i) + stats_offset = ( + stats_seq_offset + + kv_head_idx * stride_exp_sums_head + + part_idx * stride_exp_sums_part + + eqgs_idx + ) + loaded_sum = fx.Float32(exp_sums[stats_offset]) + loaded_max = fx.Float32(max_logits[stats_offset]) + part_sum = lane_in_range.select(loaded_sum, zero_f) + part_max = lane_in_range.select(loaded_max, neg_inf) + else: + lane_in_range = lane < c_part_num + stats_offset = ( + stats_seq_offset + + kv_head_idx * stride_exp_sums_head + + lane * stride_exp_sums_part + + eqgs_idx + ) + part_sum = zero_f + part_max = neg_inf + if lane_in_range: + part_sum = fx.Float32(exp_sums[stats_offset]) + part_max = fx.Float32(max_logits[stats_offset]) + + global_max = _wave_reduce_max(part_max) + if fx.const_expr(use_sinks): + sink_value = fx.Float32(sink_token[kv_head_idx * c_qgs + group_idx]) + global_max = global_max.maximumf(sink_value) + safe_global_max = (global_max > neg_inf).select(global_max, zero_f) + part_scale = (part_max > neg_inf).select( + fx.exp2((part_max - safe_global_max) * c_log2e, fastmath="fast"), + zero_f, + ) + scaled_sum = part_sum * part_scale + global_exp_sum = _wave_reduce_sum(scaled_sum) + if fx.const_expr(use_sinks): + sink_scale = _sink_exp(sink_value, safe_global_max) + global_exp_sum = global_exp_sum + sink_scale + safe_global_exp_sum = (global_exp_sum > zero_f).select( + global_exp_sum, one_f + ) + wave_inv_exp_sum = fx.Float32(rcp_f32(safe_global_exp_sum)) + weight_local = scaled_sum * wave_inv_exp_sum + weight_local_i32 = weight_local.bitcast(fx.Int32) + + acc = zero_f + for part_idx in fx.range_constexpr(max_context_partition_num): + c_part_idx = fx.Int32(part_idx) + weight_i32 = fx.Int32( + fx.rocdl.ds_bpermute( + T.i32, + c_part_idx * c_four, + weight_local_i32, + ) + ) + weight = weight_i32.bitcast(fx.Float32) + logits_offset = ( + logits_seq_offset + + kv_head_idx * stride_logits_head + + c_part_idx * stride_logits_part + + eqgs_idx * stride_logits_group + + tid + ) + part_logits = fx.Float32(logits[logits_offset]) + acc = acc + part_logits * weight + else: + # Lane l owns l + 64*k; reduce locally, then broadcast each weight. + partitions_per_lane = ( + max_context_partition_num + warp_size - 1 + ) // warp_size + part_sums = [] + part_maxes = [] + lane_max = neg_inf + for chunk_idx in fx.range_constexpr(partitions_per_lane): + chunk_base = chunk_idx * warp_size + chunk_size = min(warp_size, max_context_partition_num - chunk_base) + part_idx = lane + fx.Int32(chunk_base) + if fx.const_expr(chunk_size == warp_size and not use_work_plan): + stats_offset = ( + stats_seq_offset + + kv_head_idx * stride_exp_sums_head + + part_idx * stride_exp_sums_part + + eqgs_idx + ) + part_sum = fx.Float32(exp_sums[stats_offset]) + part_max = fx.Float32(max_logits[stats_offset]) + else: + lane_in_range = (lane < fx.Int32(chunk_size)) & ( + part_idx < c_part_num + ) + stats_offset = ( + stats_seq_offset + + kv_head_idx * stride_exp_sums_head + + part_idx * stride_exp_sums_part + + eqgs_idx + ) + part_sum = zero_f + part_max = neg_inf + if lane_in_range: + part_sum = fx.Float32(exp_sums[stats_offset]) + part_max = fx.Float32(max_logits[stats_offset]) + part_sums.append(part_sum) + part_maxes.append(part_max) + lane_max = lane_max.maximumf(part_max) + + global_max = _wave_reduce_max(lane_max) + if fx.const_expr(use_sinks): + sink_value = fx.Float32(sink_token[kv_head_idx * c_qgs + group_idx]) + global_max = global_max.maximumf(sink_value) + safe_global_max = (global_max > neg_inf).select(global_max, zero_f) + scaled_sums = [] + lane_exp_sum = zero_f + for chunk_idx in fx.range_constexpr(partitions_per_lane): + part_max = part_maxes[chunk_idx] + if fx.const_expr(use_work_plan): + part_scale = zero_f + if part_max > neg_inf: + part_scale = fx.exp2( + (part_max - safe_global_max) * c_log2e, + fastmath="fast", + ) + else: + part_scale = (part_max > neg_inf).select( + fx.exp2( + (part_max - safe_global_max) * c_log2e, + fastmath="fast", + ), + zero_f, + ) + scaled_sum = part_sums[chunk_idx] * part_scale + scaled_sums.append(scaled_sum) + lane_exp_sum = lane_exp_sum + scaled_sum + + global_exp_sum = _wave_reduce_sum(lane_exp_sum) + if fx.const_expr(use_sinks): + sink_scale = _sink_exp(sink_value, safe_global_max) + global_exp_sum = global_exp_sum + sink_scale + safe_global_exp_sum = (global_exp_sum > zero_f).select( + global_exp_sum, one_f + ) + striped_inv_exp_sum = fx.Float32(rcp_f32(safe_global_exp_sum)) + + acc = zero_f + for chunk_idx in fx.range_constexpr(partitions_per_lane): + chunk_base = chunk_idx * warp_size + chunk_size = min(warp_size, max_context_partition_num - chunk_base) + if fx.const_expr(use_work_plan): + # Initialize per unrolled chunk so the dynamic-if rewriter + # cannot reuse a stale value. + weight_local_i32 = zero_f.bitcast(fx.Int32) + if fx.Int32(chunk_base) < c_part_num: + weight_local_i32 = ( + scaled_sums[chunk_idx] * striped_inv_exp_sum + ).bitcast(fx.Int32) + for part_lane in fx.range_constexpr(chunk_size): + part_idx = chunk_base + part_lane + c_part_idx = fx.Int32(part_idx) + if c_part_idx < c_part_num: + weight_i32 = fx.Int32( + fx.rocdl.ds_bpermute( + T.i32, + fx.Int32(part_lane) * c_four, + weight_local_i32, + ) + ) + weight = weight_i32.bitcast(fx.Float32) + logits_offset = ( + logits_seq_offset + + kv_head_idx * stride_logits_head + + c_part_idx * stride_logits_part + + eqgs_idx * stride_logits_group + + tid + ) + part_logits = fx.Float32(logits[logits_offset]) + acc = acc + part_logits * weight + else: + weight_local_i32 = ( + scaled_sums[chunk_idx] * striped_inv_exp_sum + ).bitcast(fx.Int32) + for part_lane in fx.range_constexpr(chunk_size): + part_idx = chunk_base + part_lane + c_part_idx = fx.Int32(part_idx) + weight_i32 = fx.Int32( + fx.rocdl.ds_bpermute( + T.i32, + fx.Int32(part_lane) * c_four, + weight_local_i32, + ) + ) + weight = weight_i32.bitcast(fx.Float32) + logits_offset = ( + logits_seq_offset + + kv_head_idx * stride_logits_head + + c_part_idx * stride_logits_part + + eqgs_idx * stride_logits_group + + tid + ) + part_logits = fx.Float32(logits[logits_offset]) + acc = acc + part_logits * weight + + if fx.const_expr(static_query_group_size is not None): + query_idx = udiv_const(eqgs_idx, static_query_group_size) + else: + query_idx = eqgs_idx // c_qgs + if fx.const_expr(use_parallel_lds): + if partition_group == zero_i: + output_offset = ( + batch_idx * stride_output_bs + + query_idx * stride_output_len + + kv_head_idx * stride_output_kv_head + + group_idx * stride_output_group_size + + output_element + ) + output[output_offset] = acc.to(output_dtype) + else: + output_offset = ( + batch_idx * stride_output_bs + + query_idx * stride_output_len + + kv_head_idx * stride_output_kv_head + + group_idx * stride_output_group_size + + tid * fx.Int32(planned_elements_per_thread) + ) + if fx.const_expr(vectorize_plan_logits): + # Do not require a four-byte-aligned output pointer/row base. + output[output_offset] = acc[0].to(output_dtype) + output[output_offset + fx.Int32(1)] = acc[1].to(output_dtype) + else: + output[output_offset] = acc.to(output_dtype) + + @flyc.jit + def launch_pa_decode_ps_reduce_kernel( + output: fx.Pointer, + exp_sums: fx.Pointer, + max_logits: fx.Pointer, + logits: fx.Pointer, + sink_token: fx.Pointer, + 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, + query_seq_len: fx.Int32, + query_group_size: fx.Int32, + batch_size: fx.Int32, + num_kv_heads: fx.Int32, + reduce_info: fx.Pointer, + stream: fx.Stream, + ): + pa_decode_ps_reduce_kernel( + output, + exp_sums, + max_logits, + logits, + sink_token, + stride_output_bs, + stride_output_len, + stride_output_kv_head, + stride_output_group_size, + stride_exp_sums_seq, + stride_exp_sums_head, + stride_exp_sums_part, + stride_logits_seq, + stride_logits_head, + stride_logits_part, + stride_logits_group, + query_group_size, + reduce_info, + ).launch( + grid=(batch_size, num_kv_heads, query_seq_len * query_group_size), + block=tuple(block_shape), + stream=stream, + ) + + return { + "launch": launch_pa_decode_ps_reduce_kernel, + "kernel": pa_decode_ps_reduce_kernel, + } diff --git a/aiter/ops/flydsl/kernels/utils.py b/aiter/ops/flydsl/kernels/utils.py new file mode 100644 index 00000000000..48a3a117bb3 --- /dev/null +++ b/aiter/ops/flydsl/kernels/utils.py @@ -0,0 +1,39 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (C) 2025-2026 FlyDSL Project Contributors + +import flydsl.expr as fx +from flydsl.expr import const_expr, rocdl +from flydsl.expr.typing import T + + +def rcp_f32(value): + return rocdl.rcp(T.f32, value) + + +def pow2_shift(value: int) -> int: + assert value > 0 and (value & (value - 1)) == 0 + return value.bit_length() - 1 + + +def is_pow2(value: int) -> bool: + return value > 0 and (value & (value - 1)) == 0 + + +def udiv_pow2(value, divisor: int): + return value >> fx.Int32(pow2_shift(divisor)) + + +def urem_pow2(value, divisor: int): + return value & fx.Int32(divisor - 1) + + +def udiv_const(value, divisor: int): + if const_expr(is_pow2(divisor)): + return udiv_pow2(value, divisor) + return value // fx.Int32(divisor) + + +def urem_const(value, divisor: int): + if const_expr(is_pow2(divisor)): + return urem_pow2(value, divisor) + return value % fx.Int32(divisor) diff --git a/aiter/ops/flydsl/pa_decode.py b/aiter/ops/flydsl/pa_decode.py new file mode 100644 index 00000000000..f9aaf134d96 --- /dev/null +++ b/aiter/ops/flydsl/pa_decode.py @@ -0,0 +1,717 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (C) 2025-2026 FlyDSL Project Contributors + +"""FP8 paged-attention decode with BF16/FP16 queries on gfx942/gfx950. + +Cache layouts are logical, not preshuffled. K/V use E4M3 FNUZ on gfx942 and +OCP on gfx950. See ``kernels.pa_decode_kernel`` for Q/P quantization and +MFMA specialization details. +""" + +import flydsl.compiler as flyc +import flydsl.expr as fx +import torch + +from aiter.jit.utils.chip_info import get_gfx_runtime + +from .kernels.pa_decode_kernel import KV_COMPUTE_BLOCK, compile_pa_decode_tile +from .kernels.pa_decode_plan import PADecodePlan +from .kernels.pa_decode_plan import plan_pa_decode as plan_pa_decode # noqa: PLC0414 +from .kernels.pa_decode_reduce import ( + MAX_CONTEXT_PARTITIONS, + compile_pa_decode_ps_reduce, +) +from .kernels.tensor_shim import _run_compiled, get_dtype_str, ptr_arg + + +def get_recommended_splits( + num_sequences: int, + num_kv_heads: int, + split_kv_blocks: int = 1, + max_partitions: int | None = None, + *, + max_context_length: int | None = None, +) -> int: + """Recommend a uniform split count for scratch allocation and ``pa_decode``. + + Without ``max_context_length``, the default cap is eight. A host length + hint targets two CTAs per CU, bounded by 256-token tiles and the reducer + limit; short contexts and large grids retain the legacy recommendation. + ``max_partitions`` caps either mode. Allocate scratch and call ``pa_decode`` + with the returned count for every sequence; no GPU lengths are read back. + """ + if max_context_length is not None and max_context_length < 0: + raise ValueError("max_context_length must be non-negative") + if max_partitions is None: + max_partitions = 8 if max_context_length is None else MAX_CONTEXT_PARTITIONS + if not 4 <= max_partitions <= MAX_CONTEXT_PARTITIONS: + raise ValueError( + f"max_partitions must be in [4, {MAX_CONTEXT_PARTITIONS}], " + f"got {max_partitions}" + ) + props = torch.cuda.get_device_properties(torch.device("cuda")) + num_sm = props.multi_processor_count * 2 + denom = max(1, num_sequences * num_kv_heads * split_kv_blocks) + n = ((num_sm + denom - 1) // denom) * split_kv_blocks + if max_context_length is not None: + legacy = max(4, min(n, 8)) + # Count compute tiles, not pages; the floor preserves short-grid tuning. + context_tiles = (max_context_length + KV_COMPUTE_BLOCK - 1) // KV_COMPUTE_BLOCK + work_limit = max(8, context_tiles) + sequence_heads = max(1, num_sequences * num_kv_heads) + occupancy_limit = (num_sm + sequence_heads - 1) // sequence_heads + n = max(legacy, min(work_limit, occupancy_limit)) + return max(4, min(n, max_partitions)) + + +def _flydsl_pointer_dtype(dtype: torch.dtype): + return { + torch.float32: fx.Float32, + torch.float16: fx.Float16, + torch.bfloat16: fx.BFloat16, + torch.float8_e4m3fn: fx.Float8E4M3FN, + torch.float8_e4m3fnuz: fx.Float8E4M3FNUZ, + torch.int32: fx.Int32, + }[dtype] + + +def launch_pa_decode_ps_reduce( + output: torch.Tensor, + exp_sums: torch.Tensor, + max_logits: torch.Tensor, + logits: torch.Tensor, + sink_token: torch.Tensor | None, + stride_output_bs: int, + stride_output_len: int, + stride_output_kv_head: int, + stride_output_group_size: int, + stride_exp_sums_seq: int, + stride_exp_sums_head: int, + stride_exp_sums_part: int, + stride_logits_seq: int, + stride_logits_head: int, + stride_logits_part: int, + stride_logits_group: int, + *, + query_seq_len: int, + query_group_size: int, + head_size: int, + context_partition_num: int, + stream: torch.cuda.Stream, + reduce_info: torch.Tensor | None = None, +) -> None: + use_work_plan = reduce_info is not None + partition_limit = ( + torch.cuda.get_device_properties(output.device).multi_processor_count + if use_work_plan + else MAX_CONTEXT_PARTITIONS + ) + if context_partition_num > partition_limit: + raise ImportError( + f"FlyDSL pa_decode reduce supports at most {partition_limit} partitions" + ) + use_sinks = sink_token is not None + # Bound all attempted i32 byte offsets, including inactive partitions, + # so out-of-bounds reads cannot wrap back into valid data. + bounded_plan_logits = ( + use_work_plan + and context_partition_num <= 64 + and query_seq_len > 0 + and query_group_size > 0 + and 0 <= stride_logits_head <= 2**31 - 1 + and stride_logits_group == head_size + and stride_logits_part == query_seq_len * query_group_size * head_size + and 0 + < context_partition_num * stride_logits_part * logits.element_size() + <= 2**31 - 1 + ) + # Paired loads need dword alignment; scalar output stores do not. + vectorize_plan_logits = ( + bounded_plan_logits + and head_size == 128 + and logits.dtype in (torch.bfloat16, torch.float16) + and logits.data_ptr() % 4 == 0 + and stride_logits_head % 2 == 0 + and stride_logits_part % 2 == 0 + and stride_logits_group % 2 == 0 + ) + compiled = compile_pa_decode_ps_reduce( + max_context_partition_num=context_partition_num, + head_size=head_size, + output_dtype_str=get_dtype_str(output.dtype), + logits_dtype_str=get_dtype_str(logits.dtype), + sink_dtype_str=get_dtype_str( + output.dtype if sink_token is None else sink_token.dtype + ), + use_sinks=use_sinks, + use_work_plan=use_work_plan, + query_group_size=query_group_size, + bounded_plan_logits=bounded_plan_logits, + vectorize_plan_logits=vectorize_plan_logits, + ) + sink_ptr = ( + ptr_arg(sink_token, _flydsl_pointer_dtype(sink_token.dtype)) + if use_sinks + else flyc.from_c_void_p(_flydsl_pointer_dtype(output.dtype), 0) + ) + _run_compiled( + compiled["launch"], + ptr_arg(output, _flydsl_pointer_dtype(output.dtype)), + ptr_arg(exp_sums, fx.Float32), + ptr_arg(max_logits, fx.Float32), + ptr_arg(logits, _flydsl_pointer_dtype(logits.dtype)), + sink_ptr, + stride_output_bs, + stride_output_len, + stride_output_kv_head, + stride_output_group_size, + stride_exp_sums_seq, + stride_exp_sums_head, + stride_exp_sums_part, + stride_logits_seq, + stride_logits_head, + stride_logits_part, + stride_logits_group, + query_seq_len, + query_group_size, + output.shape[0], + output.shape[2], + ( + ptr_arg(reduce_info, fx.Int32) + if reduce_info is not None + else flyc.from_c_void_p(fx.Int32, 0) + ), + stream, + ) + + +def pa_decode( + output: torch.Tensor, # [num_seqs * query_length, num_query_heads, head_size] + query: torch.Tensor, # [num_seqs * query_length, num_query_heads, head_size] + key_cache: torch.Tensor, # [num_blocks, num_kv_heads, head_size // x, kv_block_size, x] + value_cache: torch.Tensor, # [num_blocks, num_kv_heads, head_size, kv_block_size] or [num_blocks, num_kv_heads, kv_block_size // x, head_size, x] + context_lengths: torch.Tensor, # [num_seqs] + block_tables: torch.Tensor, # [num_seqs, max_num_blocks_per_seq] + softmax_scale: float, + query_length: int, + max_context_partition_num: int, + context_partition_size: int = 256, + compute_type: torch.dtype = torch.bfloat16, + query_scale: torch.Tensor = None, + key_scale: torch.Tensor = None, # [num_blocks, num_kv_heads, kv_block_size, 1] + value_scale: torch.Tensor = None, # [num_blocks, num_kv_heads, kv_block_size, 1] + exp_sums: torch.Tensor = None, # [num_seqs, num_kv_heads, max_context_partition_num, query_length * query_group_size] + max_logits: torch.Tensor = None, # [num_seqs, num_kv_heads, max_context_partition_num, query_length * query_group_size] + temporary_output: torch.Tensor = None, # [num_seqs, num_kv_heads, max_context_partition_num, query_length * query_group_size, head_size] + alibi_slopes: torch.Tensor = None, + sinks: torch.Tensor = None, + sliding_window: int = 0, + work_plan: PADecodePlan | None = None, +) -> None: + """Decode FP8 K/V with BF16/FP16 queries using 256-token compute tiles. + + Supports page sizes 16/64/128 and head_dim 64 or multiples of 128 up to 1024. + K/V scales are [1] or [num_blocks, num_kv_heads, block_size, 1]. + ALiBi and externally quantized queries are unsupported. + + MTP lengths include the query tokens and use dense causal masking. + Independently selected sparse queries need separate table rows and + query_length=1. Positive ``sliding_window`` requires ``work_plan`` and + includes each query's own position; 0 and -1 disable it. + + ``sinks`` is a contiguous [num_query_heads] BF16/FP16/FP32 tensor on the + query device: unscaled zero-value logits shared across batch/MTP positions. + Each contributes once, independently of the window; -inf disables it and + +inf suppresses the head's output. + + Refresh ``work_plan`` on the current stream after changing lengths. Its + max_partitions must equal max_context_partition_num, bounded by device CUs + (static limit: 256). Match sliding_window and, for windowed plans, + query_length. Planned scratch is [KV heads, capacity, query rows (, D)]; + static scratch uses the per-sequence layouts shown in the signature. + + GPU lengths and table entries are not checked. Callers must ensure + 0 <= context_lengths[i] <= block_tables.shape[1] * block_size and used + block indices in [0, min(key_cache.shape[0], value_cache.shape[0])). + """ + if context_partition_size != KV_COMPUTE_BLOCK: + raise NotImplementedError( + "pa_decode only supports context_partition_size=256, " + f"got {context_partition_size}" + ) + if query_scale is not None: + raise NotImplementedError( + "pa_decode does not support externally quantized FP8 queries" + ) + if alibi_slopes is not None: + raise NotImplementedError("pa_decode does not support ALiBi") + if not isinstance(sliding_window, int): + raise TypeError("sliding_window must be an int") + if sliding_window < -1: + raise ValueError("sliding_window must be -1, 0, or positive") + sliding_window = max(sliding_window, 0) + if sliding_window > 0 and work_plan is None: + raise ValueError("positive sliding_window requires work_plan") + if not isinstance(query_length, int): + raise TypeError("query_length must be an int") + if query_length < 1: + raise ValueError(f"query_length must be positive, got {query_length}") + if ( + work_plan is None + and not 1 <= max_context_partition_num <= MAX_CONTEXT_PARTITIONS + ): + raise ValueError( + f"max_context_partition_num must be in [1, {MAX_CONTEXT_PARTITIONS}], " + f"got {max_context_partition_num}" + ) + + required_tensors = ( + ("output", output), + ("query", query), + ("key_cache", key_cache), + ("value_cache", value_cache), + ("context_lengths", context_lengths), + ("block_tables", block_tables), + ) + for name, tensor in required_tensors: + if not isinstance(tensor, torch.Tensor): + raise TypeError( + f"{name} must be a torch.Tensor, got {type(tensor).__name__}" + ) + + expected_ranks = ( + ("output", output, (3,)), + ("query", query, (3,)), + ("key_cache", key_cache, (5,)), + ("value_cache", value_cache, (4, 5)), + ("context_lengths", context_lengths, (1,)), + ("block_tables", block_tables, (2,)), + ) + for name, tensor, ranks in expected_ranks: + if tensor.dim() not in ranks: + expected = " or ".join(f"{rank}D" for rank in ranks) + raise ValueError( + f"{name} must be {expected}, got shape {tuple(tensor.shape)}" + ) + + num_seqs = context_lengths.shape[0] + total_q_rows, num_q_heads, head_dim = query.shape + if query.device.type != "cuda": + raise ValueError(f"query must be on a CUDA device, got {query.device}") + if num_q_heads < 1: + raise ValueError(f"query must contain at least one head, got {num_q_heads}") + if total_q_rows != num_seqs * query_length: + raise ValueError( + f"query.shape[0] ({total_q_rows}) must equal " + f"context_lengths.shape[0] * query_length ({num_seqs} * {query_length})" + ) + if output.shape != query.shape: + raise ValueError( + f"output shape {tuple(output.shape)} must match " + f"query shape {tuple(query.shape)}" + ) + if block_tables.shape[0] != num_seqs: + raise ValueError( + f"block_tables.shape[0] ({block_tables.shape[0]}) must match " + f"context_lengths.shape[0] ({num_seqs})" + ) + + num_blocks, num_kv_heads, num_hgroups, block_size, hgroup_width = key_cache.shape + if num_kv_heads < 1: + raise ValueError( + f"key_cache must contain at least one KV head, got {num_kv_heads}" + ) + + # Enforce tile Q-load constraints and the reducer's 1024-thread limit, + # including for direct output. + q_chunk = head_dim // 16 + if not ( + 64 <= head_dim <= 1024 + and head_dim % 64 == 0 + and (q_chunk <= 8 or q_chunk % 8 == 0) + ): + raise NotImplementedError( + f"pa_decode does not support head_dim={head_dim}; supported values " + "are 64 and multiples of 128 in [128, 1024]" + ) + if num_hgroups != head_dim // 16 or hgroup_width != 16: + raise ValueError( + "key_cache shape must be " + "[num_blocks, num_kv_heads, head_dim // 16, block_size, 16], " + f"got {tuple(key_cache.shape)} for head_dim={head_dim}" + ) + if block_size not in (16, 64, 128): + raise ValueError( + f"pa_decode only supports block_size in (16, 64, 128), got {block_size}" + ) + + trans_v = value_cache.dim() == 5 + if trans_v: + v_num_blocks, v_num_kv_heads = value_cache.shape[:2] + expected_v_tail = (block_size // 16, head_dim, 16) + if tuple(value_cache.shape[2:]) != expected_v_tail: + raise ValueError( + "transposed value_cache shape must be " + "[num_blocks, num_kv_heads, block_size // 16, head_dim, 16], " + f"got {tuple(value_cache.shape)} for block_size={block_size}, " + f"head_dim={head_dim}" + ) + else: + v_num_blocks, v_num_kv_heads, v_head_dim, v_block_size = value_cache.shape + if v_head_dim != head_dim or v_block_size != block_size: + raise ValueError( + "value_cache shape must be " + "[num_blocks, num_kv_heads, head_dim, block_size], " + f"got {tuple(value_cache.shape)} for block_size={block_size}, " + f"head_dim={head_dim}" + ) + # Packed V is a shifted view and may span fewer blocks than K. + if v_num_blocks > num_blocks: + raise ValueError( + f"value_cache must not span more blocks than key_cache, " + f"got {num_blocks} and {v_num_blocks}" + ) + if v_num_kv_heads != num_kv_heads: + raise ValueError( + "key_cache and value_cache must have the same number of KV heads, " + f"got {num_kv_heads} and {v_num_kv_heads}" + ) + if num_q_heads % num_kv_heads != 0: + raise ValueError( + f"num_q_heads ({num_q_heads}) must be divisible by " + f"num_kv_heads ({num_kv_heads})" + ) + + arch = get_gfx_runtime() + expected_fp8_dtype = { + "gfx942": torch.float8_e4m3fnuz, + "gfx950": torch.float8_e4m3fn, + }.get(arch) + if expected_fp8_dtype is None: + raise NotImplementedError( + f"pa_decode only supports gfx942 and gfx950, got {arch}" + ) + if compute_type != expected_fp8_dtype: + raise NotImplementedError( + f"pa_decode only supports FP8 compute ({expected_fp8_dtype}) on {arch}, " + f"got {compute_type}" + ) + + if block_tables.dtype != torch.int32: + raise TypeError(f"block_tables must be int32, got {block_tables.dtype}") + if context_lengths.dtype != torch.int32: + raise TypeError(f"context_lengths must be int32, got {context_lengths.dtype}") + query_group_size = num_q_heads // num_kv_heads + max_blocks_per_seq = block_tables.shape[1] + if query.dtype == torch.bfloat16: + query_dtype = "bf16" + elif query.dtype == torch.float16: + query_dtype = "f16" + else: + raise TypeError(f"pa_decode only supports f16/bf16 query, got {query.dtype}") + if output.dtype != query.dtype: + raise TypeError( + "pa_decode requires output.dtype == query.dtype, " + f"got {output.dtype} vs {query.dtype}" + ) + + if key_cache.dtype != expected_fp8_dtype: + raise TypeError( + f"pa_decode requires {expected_fp8_dtype} key cache on {arch}, " + f"got {key_cache.dtype}" + ) + if value_cache.dtype != expected_fp8_dtype: + raise TypeError( + f"pa_decode requires {expected_fp8_dtype} value cache on {arch}, " + f"got {value_cache.dtype}" + ) + + if query.stride(2) != 1: + raise ValueError( + "pa_decode requires a contiguous query head_dim axis, " + f"got strides {query.stride()}" + ) + if output.stride(2) != 1: + raise ValueError( + "pa_decode requires a contiguous output head_dim axis, " + f"got strides {output.stride()}" + ) + + dev = query.device + if sinks is not None: + if not isinstance(sinks, torch.Tensor): + raise TypeError("sinks must be a torch.Tensor") + if sinks.shape != (num_q_heads,): + raise ValueError("sinks must have shape [num_query_heads]") + if sinks.dtype not in (torch.bfloat16, torch.float16, torch.float32): + raise TypeError("sinks must have dtype bfloat16, float16, or float32") + if sinks.device != dev: + raise ValueError("sinks must be on the same device as query") + if not sinks.is_contiguous(): + raise ValueError("sinks must be contiguous") + for name, tensor in ( + ("output", output), + ("key_cache", key_cache), + ("value_cache", value_cache), + ("block_tables", block_tables), + ("context_lengths", context_lengths), + ): + if tensor.device != dev: + raise ValueError( + f"{name} must be on the same device as query ({dev}), " + f"got {tensor.device}" + ) + for name, tensor in ( + ("key_cache", key_cache), + ("value_cache", value_cache), + ("block_tables", block_tables), + ("context_lengths", context_lengths), + ): + if not tensor.is_contiguous(): + raise ValueError(f"{name} must be contiguous") + + def normalize_scale(scale, name): + if scale is None: + return torch.ones(1, dtype=torch.float32, device=dev) + if not isinstance(scale, torch.Tensor): + return torch.tensor([float(scale)], dtype=torch.float32, device=dev) + if scale.numel() == 1: + return scale.reshape(1) + if scale.dim() == 4: + if scale.shape[-1] != 1: + raise ValueError( + f"{name} must have a trailing singleton dimension, " + f"got shape {tuple(scale.shape)}" + ) + return scale.squeeze(-1) + if scale.dim() != 3: + raise ValueError( + f"{name} must be scalar or have shape " + "[num_blocks, num_kv_heads, block_size, 1]" + ) + return scale + + if (key_scale is None) != (value_scale is None): + raise ValueError( + "key_scale and value_scale must either both be provided or both be None" + ) + key_scale_t = normalize_scale(key_scale, "key_scale") + value_scale_t = normalize_scale(value_scale, "value_scale") + per_token_kv = key_scale_t.numel() > 1 + if per_token_kv != (value_scale_t.numel() > 1): + raise ValueError( + "key_scale and value_scale must both be per-tensor or both be per-token" + ) + if per_token_kv: + if key_scale_t.shape != value_scale_t.shape: + raise ValueError( + "key_scale/value_scale shape mismatch: " + f"{tuple(key_scale_t.shape)} vs {tuple(value_scale_t.shape)}" + ) + expected_scale_shape = (num_blocks, num_kv_heads, block_size) + if key_scale_t.shape != expected_scale_shape: + raise ValueError( + "per-token key_scale/value_scale must be " + "[num_blocks, num_kv_heads, block_size] matching the KV cache, " + f"got {tuple(key_scale_t.shape)}" + ) + stride_ks_block = int(key_scale_t.stride(0)) + stride_ks_head = int(key_scale_t.stride(1)) + if key_scale_t.stride(2) != 1: + raise ValueError( + "per-token key_scale token dimension must be contiguous, " + f"got strides {key_scale_t.stride()}" + ) + if value_scale_t.stride() != key_scale_t.stride(): + raise ValueError( + "per-token key_scale and value_scale must have matching strides, " + f"got {key_scale_t.stride()} vs {value_scale_t.stride()}" + ) + else: + stride_ks_block = 0 + stride_ks_head = 0 + for name, scale in (("key_scale", key_scale_t), ("value_scale", value_scale_t)): + if scale.dtype != torch.float32: + raise TypeError(f"{name} tensor must be float32, got {scale.dtype}") + if scale.device != dev: + raise ValueError( + f"{name} tensor must be on the same device as query ({dev}), " + f"got {scale.device}" + ) + + num_partitions = max_context_partition_num + if work_plan is not None: + if not isinstance(work_plan, PADecodePlan): + raise TypeError("work_plan must be a PADecodePlan") + work_plan.validate(num_seqs, num_kv_heads, dev) + if work_plan.max_partitions != num_partitions: + raise ValueError( + "max_context_partition_num must match work_plan.max_partitions" + ) + if work_plan.sliding_window != sliding_window: + raise ValueError("sliding_window must match work_plan.sliding_window") + if sliding_window > 0 and work_plan.query_length != query_length: + raise ValueError("query_length must match work_plan.query_length") + pmax = max_logits + psum = exp_sums + pout = temporary_output + + # Widen before either cache's i32 element offsets can wrap. + wide_kv_addressing = max(key_cache.numel(), value_cache.numel()) >= 2**31 + + # Add sinks once: here for static NP=1, otherwise in reduction. + use_direct_sinks = sinks is not None and num_partitions == 1 and work_plan is None + with torch.cuda.device(dev): + compiled = compile_pa_decode_tile( + head_dim=head_dim, + query_group_size=query_group_size, + block_size=int(block_size), + num_seqs=num_seqs, + num_kv_heads=num_kv_heads, + num_compute_units=torch.cuda.get_device_properties( + dev + ).multi_processor_count, + num_partitions=num_partitions, + softmax_scale=softmax_scale, + query_dtype=query_dtype, + per_token_kv=per_token_kv, + query_length=query_length, + trans_v=trans_v, + wide_kv_addressing=wide_kv_addressing, + use_work_plan=work_plan is not None, + work_capacity=work_plan.capacity if work_plan is not None else None, + sliding_window=sliding_window, + use_sinks=use_direct_sinks, + sink_dtype_str=get_dtype_str(sinks.dtype) if use_direct_sinks else "f32", + ) + + if num_partitions == 1 and work_plan is None: + # Direct output ignores caller scratch; supply unused kernel arguments. + dummy = torch.empty(1, dtype=torch.float32, device=dev) + pmax = psum = pout = dummy + else: + total_rows = query_length * query_group_size + expected_scalar_shape = (num_seqs, num_kv_heads, num_partitions, total_rows) + if work_plan is not None: + expected_scalar_shape = (num_kv_heads, work_plan.capacity, total_rows) + if pmax is None or psum is None or pout is None: + if pmax is None: + pmax = torch.empty( + *expected_scalar_shape, dtype=torch.float32, device=dev + ) + if psum is None: + psum = torch.empty( + *expected_scalar_shape, dtype=torch.float32, device=dev + ) + if pout is None: + pout = torch.empty( + *expected_scalar_shape, head_dim, dtype=output.dtype, device=dev + ) + for name, tensor in ( + ("max_logits", pmax), + ("exp_sums", psum), + ("temporary_output", pout), + ): + if not isinstance(tensor, torch.Tensor): + raise TypeError( + f"{name} must be a torch.Tensor, got {type(tensor).__name__}" + ) + if pmax.shape != expected_scalar_shape: + raise ValueError( + f"max_logits shape {tuple(pmax.shape)} != {expected_scalar_shape}" + ) + if psum.shape != expected_scalar_shape: + raise ValueError( + f"exp_sums shape {tuple(psum.shape)} != {expected_scalar_shape}" + ) + expected_output_shape = (*expected_scalar_shape, head_dim) + if pout.shape != expected_output_shape: + raise ValueError( + f"temporary_output shape {tuple(pout.shape)} != {expected_output_shape}" + ) + if pmax.dtype != torch.float32: + raise TypeError(f"max_logits must be float32, got {pmax.dtype}") + if psum.dtype != torch.float32: + raise TypeError(f"exp_sums must be float32, got {psum.dtype}") + if pout.dtype != output.dtype: + raise TypeError( + f"temporary_output dtype {pout.dtype} must match output dtype " + f"{output.dtype}" + ) + for name, tensor in ( + ("max_logits", pmax), + ("exp_sums", psum), + ("temporary_output", pout), + ): + if tensor.device != dev: + raise ValueError( + f"{name} must be on the same device as query ({dev}), " + f"got {tensor.device}" + ) + if not tensor.is_contiguous(): + raise ValueError(f"{name} must be contiguous") + with torch.cuda.device(dev): + s = torch.cuda.current_stream(dev) + _run_compiled( + compiled["launch"], + ptr_arg(output, _flydsl_pointer_dtype(output.dtype)), + ptr_arg(pmax, fx.Float32), + ptr_arg(psum, fx.Float32), + ptr_arg(pout, _flydsl_pointer_dtype(pout.dtype)), + ptr_arg(query, _flydsl_pointer_dtype(query.dtype)), + ptr_arg(key_cache, _flydsl_pointer_dtype(key_cache.dtype)), + ptr_arg(value_cache, _flydsl_pointer_dtype(value_cache.dtype)), + ptr_arg(block_tables, fx.Int32), + ptr_arg(context_lengths, fx.Int32), + ptr_arg(key_scale_t, fx.Float32), + ptr_arg(value_scale_t, fx.Float32), + ( + ptr_arg(sinks, _flydsl_pointer_dtype(sinks.dtype)) + if use_direct_sinks + else flyc.from_c_void_p(fx.Float32, 0) + ), + int(max_blocks_per_seq), + int(num_seqs), + int(num_kv_heads), + stride_ks_block, + stride_ks_head, + int(output.stride(0)), + int(output.stride(1)), + int(query.stride(0)), + int(query.stride(1)), + ( + ptr_arg(work_plan.work_info, fx.Int32) + if work_plan is not None + else flyc.from_c_void_p(fx.Int32, 0) + ), + work_plan.capacity if work_plan is not None else 0, + s, + ) + if num_partitions > 1 or work_plan is not None: + output_5d = output.reshape( + num_seqs, query_length, num_kv_heads, query_group_size, head_dim + ) + launch_pa_decode_ps_reduce( + output_5d, + psum, + pmax, + pout, + sinks, + output_5d.stride(0), + output_5d.stride(1), + output_5d.stride(2), + output_5d.stride(3), + pmax.stride(0) if work_plan is None else 0, + pmax.stride(1) if work_plan is None else pmax.stride(0), + pmax.stride(2) if work_plan is None else pmax.stride(1), + pout.stride(0) if work_plan is None else 0, + pout.stride(1) if work_plan is None else pout.stride(0), + pout.stride(2) if work_plan is None else pout.stride(1), + pout.stride(3) if work_plan is None else pout.stride(2), + query_seq_len=query_length, + query_group_size=query_group_size, + head_size=head_dim, + context_partition_num=num_partitions, + stream=s, + reduce_info=work_plan.reduce_info if work_plan is not None else None, + ) diff --git a/aiter/ops/triton/gluon/pa_decode_gluon.py b/aiter/ops/triton/gluon/pa_decode_gluon.py index 2f844821904..ecb90d23f00 100644 --- a/aiter/ops/triton/gluon/pa_decode_gluon.py +++ b/aiter/ops/triton/gluon/pa_decode_gluon.py @@ -11,6 +11,12 @@ import aiter from aiter.ops.triton.utils._triton import arch_info +_FLYDSL_REDUCE_DTYPE_NAMES = { + torch.float32: "f32", + torch.float16: "f16", + torch.bfloat16: "bf16", +} + CXX_PS_REDUCE_AVAILABLE = True try: from csrc.cpp_itfs.pa.pa_ps import ( @@ -22,34 +28,16 @@ FLYDSL_PS_REDUCE_AVAILABLE = True try: - import flydsl.compiler as flyc - import flydsl.expr as fx - from flydsl._mlir import ir - from flydsl._mlir.dialects import arith as _mlir_arith - from flydsl.compiler.kernel_function import CompilationContext - from flydsl.expr import arith, gpu, range_constexpr, rocdl - from flydsl.expr.typing import Int32, T - from flydsl.runtime.device import get_rocm_arch as get_hip_arch - from flydsl.utils.smem_allocator import SmemAllocator, SmemPtr - - from aiter.ops.flydsl.kernels import buffer_ops + from aiter.ops.flydsl.kernels.pa_decode_reduce import ( + is_pa_decode_ps_reduce_supported, + ) + from aiter.ops.flydsl.pa_decode import ( + launch_pa_decode_ps_reduce as launch_pa_decode_ps_reduce_flydsl, + ) except Exception: # noqa: BLE001 FLYDSL_PS_REDUCE_AVAILABLE = False - flyc = None - fx = None - arith = None - gpu = None - rocdl = None - buffer_ops = None - range_constexpr = None - T = None - Int32 = None - SmemAllocator = None - SmemPtr = None - get_hip_arch = None - ir = None - CompilationContext = None - _mlir_arith = None + launch_pa_decode_ps_reduce_flydsl = None + is_pa_decode_ps_reduce_supported = None GLUON_JIT_KERNEL_ENABLED = True try: @@ -4045,10 +4033,10 @@ def paged_attention_decode_ps_reduce_kernel( @triton.jit def paged_attention_decode_v2_reduce_kernel( - output_ptr, # [num_seqs, num_kv_heads, query_group_size, head_size] - exp_sums_ptr, # [num_seqs, num_kv_heads, max_parts, query_group_size] - max_logits_ptr, # [num_seqs, num_kv_heads, max_parts, query_group_size] - logits_ptr, # [num_seqs, num_kv_heads, max_parts, query_group_size, head_size] + output_ptr, # [num_seqs, query_length, num_kv_heads, query_group_size, head_size] + exp_sums_ptr, # [num_seqs, num_kv_heads, max_parts, query_length * query_group_size] + max_logits_ptr, # [num_seqs, num_kv_heads, max_parts, query_length * query_group_size] + logits_ptr, # [num_seqs, num_kv_heads, max_parts, query_length * query_group_size, head_size] context_lengths_ptr, # [num_seqs] sink_token_ptr, # [num_query_heads] stride_output_bs, @@ -4494,564 +4482,12 @@ def _paged_attention_decode_v2_with_dot_kernel_reshape_wrapper( ) -def _flydsl_dtype_str(dtype: torch.dtype) -> str: - if dtype == torch.float32: - return "f32" - if dtype == torch.float16: - return "f16" - if dtype == torch.bfloat16: - return "bf16" - raise ValueError(f"Unsupported FlyDSL dtype: {dtype!r}") - - -@lru_cache(maxsize=256) -def compile_pa_decode_ps_reduce_flydsl( - *, - max_context_partition_num: int, - query_seq_len: int, - query_group_size: int, - head_size: int, - output_dtype_str: str, - logits_dtype_str: str, - sink_dtype_str: str, - use_sinks: bool, -): - if not FLYDSL_PS_REDUCE_AVAILABLE: - raise ImportError("FlyDSL is unavailable for pa_decode PS reduce") - - FLYDSL_WARP_SIZE = 64 - FLYDSL_LOG2E = 1.4426950408889634 - - block_threads = head_size - assert block_threads > 0, "head_size must be positive" - assert block_threads <= 1024, "head_size must fit in one workgroup" - reduce_width = ( - 1 - if max_context_partition_num <= 1 - else 1 << ((max_context_partition_num - 1).bit_length()) - ) - reduce_shuffle_offsets = [off for off in [32, 16, 8, 4, 2, 1] if off < reduce_width] - red_slots = max(1, (block_threads + FLYDSL_WARP_SIZE - 1) // FLYDSL_WARP_SIZE) - arch = get_hip_arch() - allocator = SmemAllocator(None, arch=arch, global_sym_name="pa_ps_sw_reduce_smem") - red_off = allocator._align(allocator.ptr, 16) - allocator.ptr = red_off + red_slots * 4 - part_weights_off = allocator._align(allocator.ptr, 16) - allocator.ptr = part_weights_off + max_context_partition_num * 4 - - @flyc.kernel(known_block_size=(block_threads, 1, 1)) - def pa_decode_ps_reduce_flydsl_kernel( - output_ptr: fx.Tensor, - exp_sums_ptr: fx.Tensor, - max_logits_ptr: fx.Tensor, - logits_ptr: fx.Tensor, - sink_token_ptr: fx.Tensor, - 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, - ): - tid = gpu.thread_idx.x - batch_idx = gpu.block_idx.x - kv_head_idx = gpu.block_idx.y - eqgs_idx = gpu.block_idx.z - - smem_base = allocator.get_base() - red_scratch = SmemPtr(smem_base, red_off, T.f32, shape=(red_slots,)) - red_scratch.get() - if max_context_partition_num > FLYDSL_WARP_SIZE: - part_weights_lds = SmemPtr( - smem_base, part_weights_off, T.f32, shape=(max_context_partition_num,) - ) - part_weights_lds.get() - - out_rsrc = buffer_ops.create_buffer_resource(output_ptr, max_size=True) - es_rsrc = buffer_ops.create_buffer_resource(exp_sums_ptr, max_size=True) - ml_rsrc = buffer_ops.create_buffer_resource(max_logits_ptr, max_size=True) - logits_rsrc = buffer_ops.create_buffer_resource(logits_ptr, max_size=True) - if use_sinks: - sink_rsrc = buffer_ops.create_buffer_resource(sink_token_ptr, max_size=True) - - c_zero_f = arith.constant(0.0, type=T.f32) - c_one_f = arith.constant(1.0, type=T.f32) - c_neg_inf = arith.constant(float("-inf"), type=T.f32) - c_log2e = arith.constant(FLYDSL_LOG2E, type=T.f32) - fm_fast = arith.FastMathFlags.fast - c_zero_i = arith.constant(0, type=T.i32) - c_w = arith.constant(FLYDSL_WARP_SIZE, type=T.i32) - c_wave_mask = arith.constant(FLYDSL_WARP_SIZE - 1, type=T.i32) - c_wave_shift = arith.constant(6, type=T.i32) - c_red_slots = arith.constant(red_slots, type=T.i32) - lane = tid & c_wave_mask - wave = tid >> c_wave_shift - c_qgs = arith.constant(query_group_size, type=T.i32) - group_idx = eqgs_idx % c_qgs - - def _wave_reduce_max_full(val): - red = val - for sh in [32, 16, 8, 4, 2, 1]: - red = red.maximumf(red.shuffle_xor(arith.constant(sh, type=T.i32), c_w)) - return red - - def _wave_reduce_sum_full(val): - red = val - for sh in [32, 16, 8, 4, 2, 1]: - red = red.addf( - red.shuffle_xor(arith.constant(sh, type=T.i32), c_w), - fastmath=fm_fast, - ) - return red - - def _block_reduce(val, mode): - if red_slots == 1: - return ( - _wave_reduce_max_full(val) - if mode == "max" - else _wave_reduce_sum_full(val) - ) - - neutral = c_neg_inf if mode == "max" else c_zero_f - w = ( - _wave_reduce_max_full(val) - if mode == "max" - else _wave_reduce_sum_full(val) - ) - - if arith.cmpi(arith.CmpIPredicate.eq, lane, c_zero_i): - wave_idx = arith.index_cast(T.index, wave) - red_scratch.store(w, [wave_idx]) - gpu.barrier() - - if arith.cmpi(arith.CmpIPredicate.eq, wave, c_zero_i): - in_range = arith.cmpi(arith.CmpIPredicate.slt, lane, c_red_slots) - lane_safe = arith.select(in_range, lane, c_zero_i) - lane_safe_idx = arith.index_cast(T.index, lane_safe) - red_val = red_scratch.load([lane_safe_idx]) - red_val = arith.select(in_range, red_val, neutral) - red_val = ( - _wave_reduce_max_full(red_val) - if mode == "max" - else _wave_reduce_sum_full(red_val) - ) - if arith.cmpi(arith.CmpIPredicate.eq, lane, c_zero_i): - red_scratch.store(red_val, [arith.constant(0, index=True)]) - gpu.barrier() - - return red_scratch.load([arith.constant(0, index=True)]) - - if max_context_partition_num <= FLYDSL_WARP_SIZE: - c_part_num = arith.constant(max_context_partition_num, type=T.i32) - c_reduce_width = arith.constant(reduce_width, type=T.i32) - c_four = arith.constant(4, type=T.i32) - - def _wave_reduce_max(val): - red = val - for sh in reduce_shuffle_offsets: - red = red.maximumf( - red.shuffle_xor(arith.constant(sh, type=T.i32), c_w) - ) - return red - - def _wave_reduce_sum(val): - red = val - for sh in reduce_shuffle_offsets: - red = red.addf( - red.shuffle_xor(arith.constant(sh, type=T.i32), c_w), - fastmath=fm_fast, - ) - return red - - lane_in_range = arith.cmpi(arith.CmpIPredicate.slt, lane, c_part_num) - lane_in_reduce = arith.cmpi(arith.CmpIPredicate.slt, lane, c_reduce_width) - part_sum = c_zero_f - part_max = c_neg_inf - if lane_in_reduce: - part_i32 = arith.select(lane_in_range, lane, c_zero_i) - es_off = ( - batch_idx * stride_exp_sums_seq - + kv_head_idx * stride_exp_sums_head - + 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 = arith.select(lane_in_range, part_sum_raw, c_zero_f) - part_max = arith.select(lane_in_range, part_max_raw, c_neg_inf) - - global_max = _wave_reduce_max(part_max) - safe_global_max = arith.select( - global_max > c_neg_inf, - global_max, - c_zero_f, - ) - part_scale = arith.select( - part_max > c_neg_inf, - ((part_max - safe_global_max) * c_log2e).exp2(fastmath=fm_fast), - c_zero_f, - ) - scaled_sum = part_sum * part_scale - global_exp_sum = _wave_reduce_sum(scaled_sum) - if use_sinks: - sink_off = kv_head_idx * c_qgs + group_idx - if sink_dtype_str == "f32": - sink_value = buffer_ops.buffer_load( - sink_rsrc, sink_off, vec_width=1, dtype=T.f32 - ) - elif sink_dtype_str == "f16": - sink_value_raw = buffer_ops.buffer_load( - sink_rsrc, sink_off, vec_width=1, dtype=T.f16 - ) - sink_value = _mlir_arith.ExtFOp(T.f32, sink_value_raw).result - else: - sink_value_raw = buffer_ops.buffer_load( - sink_rsrc, sink_off, vec_width=1, dtype=T.bf16 - ) - sink_value = _mlir_arith.ExtFOp(T.f32, sink_value_raw).result - sink_scale = arith.select( - global_max > c_neg_inf, - ((sink_value - safe_global_max) * c_log2e).exp2(fastmath=fm_fast), - c_zero_f, - ) - global_exp_sum = global_exp_sum + sink_scale - safe_global_exp_sum = arith.select( - global_exp_sum > c_zero_f, - global_exp_sum, - c_one_f, - ) - weight_local = scaled_sum / safe_global_exp_sum - weight_local_i32 = arith.bitcast(T.i32, weight_local) - - acc = c_zero_f - for part_idx in range_constexpr(max_context_partition_num): - part_i32 = arith.constant(part_idx, type=T.i32) - bcast_addr = part_i32 * c_four - weight_i32 = rocdl.ds_bpermute( - T.i32, arith.unwrap(bcast_addr), arith.unwrap(weight_local_i32) - ) - weight = arith.bitcast(T.f32, weight_i32) - logits_off = ( - batch_idx * stride_logits_seq - + kv_head_idx * stride_logits_head - + part_i32 * stride_logits_part - + eqgs_idx * stride_logits_group - + tid - ) - if logits_dtype_str == "f32": - part_logits = buffer_ops.buffer_load( - logits_rsrc, logits_off, vec_width=1, dtype=T.f32 - ) - elif logits_dtype_str == "f16": - part_logits_raw = buffer_ops.buffer_load( - logits_rsrc, logits_off, vec_width=1, dtype=T.f16 - ) - part_logits = _mlir_arith.ExtFOp(T.f32, part_logits_raw).result - else: - part_logits_raw = buffer_ops.buffer_load( - logits_rsrc, logits_off, vec_width=1, dtype=T.bf16 - ) - part_logits = _mlir_arith.ExtFOp(T.f32, part_logits_raw).result - acc = acc + part_logits * weight - else: - global_max = c_neg_inf - for chunk_base in range(0, max_context_partition_num, block_threads): - chunk_size = min(block_threads, max_context_partition_num - chunk_base) - c_chunk_size = arith.constant(chunk_size, type=T.i32) - c_chunk_base = arith.constant(chunk_base, type=T.i32) - in_chunk = arith.cmpi(arith.CmpIPredicate.slt, tid, c_chunk_size) - part_i32 = arith.select(in_chunk, tid + c_chunk_base, c_zero_i) - es_off = ( - batch_idx * stride_exp_sums_seq - + kv_head_idx * stride_exp_sums_head - + part_i32 * stride_exp_sums_part - + eqgs_idx - ) - part_max_raw = buffer_ops.buffer_load( - ml_rsrc, es_off, vec_width=1, dtype=T.f32 - ) - 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) - - safe_global_max = arith.select( - global_max > c_neg_inf, - global_max, - c_zero_f, - ) - global_exp_sum = c_zero_f - for chunk_base in range(0, max_context_partition_num, block_threads): - chunk_size = min(block_threads, max_context_partition_num - chunk_base) - c_chunk_size = arith.constant(chunk_size, type=T.i32) - c_chunk_base = arith.constant(chunk_base, type=T.i32) - in_chunk = arith.cmpi(arith.CmpIPredicate.slt, tid, c_chunk_size) - part_i32 = arith.select(in_chunk, tid + c_chunk_base, c_zero_i) - es_off = ( - batch_idx * stride_exp_sums_seq - + kv_head_idx * stride_exp_sums_head - + 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 = 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( - part_max > c_neg_inf, - ((part_max - safe_global_max) * c_log2e).exp2(fastmath=fm_fast), - c_zero_f, - ) - chunk_sum = _block_reduce(part_sum * part_scale, "sum") - global_exp_sum = global_exp_sum + chunk_sum - - if use_sinks: - sink_off = kv_head_idx * c_qgs + group_idx - if sink_dtype_str == "f32": - sink_value = buffer_ops.buffer_load( - sink_rsrc, sink_off, vec_width=1, dtype=T.f32 - ) - elif sink_dtype_str == "f16": - sink_value_raw = buffer_ops.buffer_load( - sink_rsrc, sink_off, vec_width=1, dtype=T.f16 - ) - sink_value = _mlir_arith.ExtFOp(T.f32, sink_value_raw).result - else: - sink_value_raw = buffer_ops.buffer_load( - sink_rsrc, sink_off, vec_width=1, dtype=T.bf16 - ) - sink_value = _mlir_arith.ExtFOp(T.f32, sink_value_raw).result - sink_scale = arith.select( - global_max > c_neg_inf, - ((sink_value - safe_global_max) * c_log2e).exp2(fastmath=fm_fast), - c_zero_f, - ) - global_exp_sum = global_exp_sum + sink_scale - - safe_global_exp_sum = arith.select( - global_exp_sum > c_zero_f, - global_exp_sum, - c_one_f, - ) - - for chunk_base in range(0, max_context_partition_num, block_threads): - chunk_size = min(block_threads, max_context_partition_num - chunk_base) - c_chunk_size = arith.constant(chunk_size, type=T.i32) - c_chunk_base = arith.constant(chunk_base, type=T.i32) - in_chunk = arith.cmpi(arith.CmpIPredicate.slt, tid, c_chunk_size) - part_i32 = arith.select(in_chunk, tid + c_chunk_base, c_zero_i) - es_off = ( - batch_idx * stride_exp_sums_seq - + kv_head_idx * stride_exp_sums_head - + 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 - ) - if in_chunk: - part_sum = part_sum_raw - part_max = part_max_raw - part_scale = arith.select( - part_max > c_neg_inf, - ((part_max - safe_global_max) * c_log2e).exp2(fastmath=fm_fast), - c_zero_f, - ) - weight = (part_sum * part_scale) / safe_global_exp_sum - part_idx_idx = arith.index_cast(T.index, part_i32) - part_weights_lds.store(weight, [part_idx_idx]) - - gpu.barrier() - - acc = c_zero_f - for part_idx in range_constexpr(max_context_partition_num): - part_i32 = arith.constant(part_idx, type=T.i32) - part_idx_idx = arith.constant(part_idx, index=True) - weight = part_weights_lds.load([part_idx_idx]) - logits_off = ( - batch_idx * stride_logits_seq - + kv_head_idx * stride_logits_head - + part_i32 * stride_logits_part - + eqgs_idx * stride_logits_group - + tid - ) - if logits_dtype_str == "f32": - part_logits = buffer_ops.buffer_load( - logits_rsrc, logits_off, vec_width=1, dtype=T.f32 - ) - elif logits_dtype_str == "f16": - part_logits_raw = buffer_ops.buffer_load( - logits_rsrc, logits_off, vec_width=1, dtype=T.f16 - ) - part_logits = _mlir_arith.ExtFOp(T.f32, part_logits_raw).result - else: - part_logits_raw = buffer_ops.buffer_load( - logits_rsrc, logits_off, vec_width=1, dtype=T.bf16 - ) - part_logits = _mlir_arith.ExtFOp(T.f32, part_logits_raw).result - acc = acc + part_logits * weight - - query_idx = eqgs_idx // c_qgs - group_idx = eqgs_idx % c_qgs - out_off = ( - batch_idx * stride_output_bs - + query_idx * stride_output_len - + kv_head_idx * stride_output_kv_head - + group_idx * stride_output_group_size - + tid - ) - if output_dtype_str == "f32": - out_val = acc - elif output_dtype_str == "f16": - out_val = arith.trunc_f(T.f16, acc) - else: - out_val = arith.trunc_f(T.bf16, acc) - buffer_ops.buffer_store(out_val, out_rsrc, out_off) - - @flyc.jit - def launch_pa_decode_ps_reduce_flydsl( - output, - exp_sums, - max_logits, - logits, - sink_token, - stride_output_bs, - stride_output_len, - stride_output_kv_head, - stride_output_group_size, - stride_exp_sums_seq, - stride_exp_sums_head, - stride_exp_sums_part, - stride_logits_seq, - stride_logits_head, - stride_logits_part, - stride_logits_group, - batch_size, - num_kv_heads, - stream: fx.Stream, - ): - allocator.finalized = False - ctx = CompilationContext.get_current() - with ir.InsertionPoint(ctx.gpu_module_body): - allocator.finalize() - pa_decode_ps_reduce_flydsl_kernel( - output, - exp_sums, - max_logits, - logits, - sink_token, - stride_output_bs, - stride_output_len, - stride_output_kv_head, - stride_output_group_size, - stride_exp_sums_seq, - stride_exp_sums_head, - stride_exp_sums_part, - stride_logits_seq, - stride_logits_head, - stride_logits_part, - stride_logits_group, - ).launch( - grid=(batch_size, num_kv_heads, query_seq_len * query_group_size), - block=(block_threads, 1, 1), - stream=stream, - ) - - return { - "launch": launch_pa_decode_ps_reduce_flydsl, - "kernel": pa_decode_ps_reduce_flydsl_kernel, - "allocator": allocator, - } - - -def launch_pa_decode_ps_reduce_flydsl( - output_ptr, - exp_sums_ptr, - max_logits_ptr, - logits_ptr, - sink_token_ptr, - stride_output_bs, - stride_output_len, - stride_output_kv_head, - stride_output_group_size, - stride_exp_sums_seq, - stride_exp_sums_head, - stride_exp_sums_part, - stride_logits_seq, - stride_logits_head, - stride_logits_part, - stride_logits_group, - query_seq_len, - query_group_size, - head_size, - context_partition_num, -): - if logits_ptr.dtype not in (torch.bfloat16, torch.float16, torch.float32): - raise ImportError( - "FlyDSL PS reduce fallback: only bf16/fp16/fp32 logits supported" - ) - - compiled = compile_pa_decode_ps_reduce_flydsl( - max_context_partition_num=context_partition_num, - query_seq_len=query_seq_len, - query_group_size=query_group_size, - head_size=head_size, - output_dtype_str=_flydsl_dtype_str(output_ptr.dtype), - logits_dtype_str=_flydsl_dtype_str(logits_ptr.dtype), - sink_dtype_str=_flydsl_dtype_str( - output_ptr.dtype if sink_token_ptr is None else sink_token_ptr.dtype - ), - use_sinks=sink_token_ptr is not None, - ) - - if sink_token_ptr is None: - sink_token_ptr = torch.empty( - 0, dtype=output_ptr.dtype, device=output_ptr.device - ) - compiled["launch"]( - output_ptr, - exp_sums_ptr, - max_logits_ptr, - logits_ptr, - sink_token_ptr, - stride_output_bs, - stride_output_len, - stride_output_kv_head, - stride_output_group_size, - stride_exp_sums_seq, - stride_exp_sums_head, - stride_exp_sums_part, - stride_logits_seq, - stride_logits_head, - stride_logits_part, - stride_logits_group, - output_ptr.shape[0], - output_ptr.shape[2], - torch.cuda.current_stream(output_ptr.device), - ) - - def _paged_attention_decode_v2_reduce_kernel_wrapper( grid, - output_ptr, # [num_seqs, num_kv_heads, query_group_size, head_size] - exp_sums_ptr, # [num_seqs, num_kv_heads, max_parts, query_group_size] - max_logits_ptr, # [num_seqs, num_kv_heads, max_parts, query_group_size] - logits_ptr, # [num_seqs, num_kv_heads, max_parts, query_group_size, head_size] + output_ptr, # [num_seqs, query_length, num_kv_heads, query_group_size, head_size] + exp_sums_ptr, # [num_seqs, num_kv_heads, max_parts, query_length * query_group_size] + max_logits_ptr, # [num_seqs, num_kv_heads, max_parts, query_length * query_group_size] + logits_ptr, # [num_seqs, num_kv_heads, max_parts, query_length * query_group_size, head_size] context_lengths_ptr, # [num_seqs] sink_token_ptr, # [num_query_heads] stride_output_bs, @@ -5109,7 +4545,25 @@ def _paged_attention_decode_v2_reduce_kernel_wrapper( return except ImportError: pass - try: + output_dtype_str = _FLYDSL_REDUCE_DTYPE_NAMES.get(output_ptr.dtype) + logits_dtype_str = _FLYDSL_REDUCE_DTYPE_NAMES.get(logits_ptr.dtype) + sink_dtype_str = _FLYDSL_REDUCE_DTYPE_NAMES.get( + output_ptr.dtype if sink_token_ptr is None else sink_token_ptr.dtype + ) + flydsl_supported = ( + FLYDSL_PS_REDUCE_AVAILABLE + and output_dtype_str is not None + and logits_dtype_str is not None + and sink_dtype_str is not None + and is_pa_decode_ps_reduce_supported( + max_context_partition_num=context_partition_num, + head_size=head_size, + output_dtype_str=output_dtype_str, + logits_dtype_str=logits_dtype_str, + sink_dtype_str=sink_dtype_str, + ) + ) + if flydsl_supported: launch_pa_decode_ps_reduce_flydsl( output_ptr, exp_sums_ptr, @@ -5131,37 +4585,34 @@ def _paged_attention_decode_v2_reduce_kernel_wrapper( query_group_size=query_group_size, head_size=head_size, context_partition_num=context_partition_num, - # Was the `fx.Stream(None)` parameter default; passed explicitly - # now that the default is gone. fx.Stream(None) is the default queue. - stream=fx.Stream(None), + stream=torch.cuda.current_stream(output_ptr.device), ) return - except ImportError: - ps_reduce_grid = (grid[0], grid[1], query_seq_len * query_group_size) - paged_attention_decode_ps_reduce_kernel[ps_reduce_grid]( - output_ptr, - exp_sums_ptr, - max_logits_ptr, - logits_ptr, - sink_token_ptr, - stride_output_bs, - stride_output_len, - stride_output_kv_head, - stride_output_group_size, - stride_exp_sums_seq, - stride_exp_sums_head, - stride_exp_sums_part, - stride_logits_seq, - stride_logits_head, - stride_logits_part, - stride_logits_group, - query_group_size=query_group_size, - head_size=head_size, - context_partition_num=context_partition_num, - HEAD_SIZE_POW2=triton.next_power_of_2(head_size), - USE_SINKS=sink_token_ptr is not None, - MAX_CONTEXT_PARTITION_NUM=triton.next_power_of_2(context_partition_num), - ) + ps_reduce_grid = (grid[0], grid[1], query_seq_len * query_group_size) + paged_attention_decode_ps_reduce_kernel[ps_reduce_grid]( + output_ptr, + exp_sums_ptr, + max_logits_ptr, + logits_ptr, + sink_token_ptr, + stride_output_bs, + stride_output_len, + stride_output_kv_head, + stride_output_group_size, + stride_exp_sums_seq, + stride_exp_sums_head, + stride_exp_sums_part, + stride_logits_seq, + stride_logits_head, + stride_logits_part, + stride_logits_group, + query_group_size=query_group_size, + head_size=head_size, + context_partition_num=context_partition_num, + HEAD_SIZE_POW2=triton.next_power_of_2(head_size), + USE_SINKS=sink_token_ptr is not None, + MAX_CONTEXT_PARTITION_NUM=triton.next_power_of_2(context_partition_num), + ) else: paged_attention_decode_v2_reduce_kernel[grid]( output_ptr, @@ -5207,9 +4658,9 @@ def pa_decode_gluon( query_scale: torch.Tensor = None, # [num_seqs * query_length, num_query_heads, 1] or [1] key_scale: torch.Tensor = None, # [num_blocks, num_kv_heads, kv_block_size, 1] value_scale: torch.Tensor = None, # [num_blocks, num_kv_heads, kv_block_size, 1] - exp_sums: torch.Tensor = None, # [num_seqs, num_kv_heads, max_context_partition_num, query_group_size] - max_logits: torch.Tensor = None, # [num_seqs, num_kv_heads, max_context_partition_num, query_group_size] - temporary_output: torch.Tensor = None, # [num_seqs, num_kv_heads, max_context_partition_num, query_group_size, head_size] + exp_sums: torch.Tensor = None, # [num_seqs, num_kv_heads, max_context_partition_num, query_length * query_group_size] + max_logits: torch.Tensor = None, # [num_seqs, num_kv_heads, max_context_partition_num, query_length * query_group_size] + temporary_output: torch.Tensor = None, # [num_seqs, num_kv_heads, max_context_partition_num, query_length * query_group_size, head_size] alibi_slopes: torch.Tensor = None, sinks: torch.Tensor = None, sliding_window: int = 0, @@ -5291,19 +4742,22 @@ def pa_decode_gluon( exp_sums : torch.Tensor Buffer for exponential sums used in online softmax computation. - - Shape: [num_seqs, num_kv_heads, max_context_partition_num, query_group_size] + - Shape: [num_seqs, num_kv_heads, max_context_partition_num, + query_length * query_group_size] where max_context_partition_num = ceil(max_context_length / context_partition_size) - Dtype: torch.float32 max_logits : torch.Tensor Buffer for maximum logits used in online softmax computation. - - Shape: [num_seqs, num_kv_heads, max_context_partition_num, query_group_size] + - Shape: [num_seqs, num_kv_heads, max_context_partition_num, + query_length * query_group_size] - Dtype: torch.float32 temporary_output : torch.Tensor Buffer for partial attention outputs from each context partition. - - Shape: [num_seqs, num_kv_heads, max_context_partition_num, query_group_size, head_size] - - Dtype: torch.float32 + - Shape: [num_seqs, num_kv_heads, max_context_partition_num, + query_length * query_group_size, head_size] + - Dtype: same as query/output alibi_slopes : torch.Tensor, optional ALiBi (Attention with Linear Biases) slopes for positional encoding. diff --git a/op_tests/test_flydsl_pa_decode.py b/op_tests/test_flydsl_pa_decode.py new file mode 100644 index 00000000000..eb41734f223 --- /dev/null +++ b/op_tests/test_flydsl_pa_decode.py @@ -0,0 +1,1196 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Parametrized correctness/contract test and CLI benchmark. + + python -m pytest -q op_tests/test_flydsl_pa_decode.py + python op_tests/test_flydsl_pa_decode.py -d bf16 -b 200 -q 4 \ + -s 16,1,128,200000 --block-size 16 128 --trans-v 0 1 \ + --per-token 1 --num-partitions 3 5 + +Contexts include MTP tokens. Positive windows require a plan; 0/-1 allow both +static and planned decoding. Plan caps default to device CU count; explicit +partitions override static recommendations. CLI timing includes the reducer. +""" + +import argparse +import importlib +import itertools +from dataclasses import dataclass, replace +from functools import partial + +import pandas as pd +import pytest +import torch + +import aiter +from aiter import dtypes, per_tensor_quant, pertoken_quant +from aiter.jit.utils.chip_info import get_gfx_runtime +from aiter.test_common import benchmark, run_perftest + +try: + from aiter.ops.flydsl.pa_decode import ( + MAX_CONTEXT_PARTITIONS, + get_recommended_splits, + pa_decode, + plan_pa_decode, + ) +except (ImportError, AttributeError, RuntimeError, OSError): + MAX_CONTEXT_PARTITIONS = 256 + get_recommended_splits = pa_decode = plan_pa_decode = None + +SUPPORTED_GFX = ("gfx942", "gfx950") +KV_COMPUTE_BLOCK = 256 +ACCURACY_TOLERANCE = 5e-3 +DEFAULT_BATCH_SIZES = [3, 81, 128] +DEFAULT_SHAPES = [ + (8, 1, 128, 257), + (4, 1, 128, 1027), + (8, 1, 128, 1027), + (8, 1, 256, 1027), + (16, 1, 128, 8192), +] +BOUNDARY_LENGTHS = ( + 0, + 1, + 2, + 3, + 4, + 255, + 256, + 257, + 258, + 259, + 511, + 512, + 513, + 514, + 515, + 2303, + 2304, + 2305, + 2306, + 2307, + 4099, +) + + +@dataclass(frozen=True) +class DecodeCase: + lengths: tuple[int, ...] = BOUNDARY_LENGTHS + query_length: int = 4 + num_kv_heads: int = 1 + query_group_size: int = 16 + head_dim: int = 128 + block_size: int = 128 + trans_v: bool = True + dtype: torch.dtype = torch.bfloat16 + num_partitions: int | None = 7 + per_token: bool = True + sliding_window: int = 0 + sink_dtype: torch.dtype | None = None + max_partitions: int | None = None + workgroup_budget: int | None = None + sparse: bool = True + masked_scale: bool = False + query_splits: int | None = None + wide_kv_addressing: bool | None = None + + +def _require_gpu(): + if not torch.cuda.is_available(): + pytest.skip("ROCm is not available") + if pa_decode is None: + pytest.skip("FlyDSL is not available") + if get_gfx_runtime() not in SUPPORTED_GFX: + pytest.skip(f"pa_decode is unsupported on {get_gfx_runtime()}") + + +@pytest.fixture(autouse=True) +def _default_cuda_device(): + # Restore the default device for other tests. + _require_gpu() + previous = torch.get_default_device() + torch.set_default_device("cuda") + try: + yield + finally: + torch.set_default_device(previous) + + +def run_torch( + query, + key_cache, + value_cache, + block_tables, + context_lengths, + key_scale, + value_scale, + query_length=1, + sliding_window=0, + sinks=None, +): + """Dequantized FP32 GQA reference, including empty rows and infinite sinks.""" + batch = context_lengths.numel() + heads, dim = query.shape[1:] + kv_heads, page_size = key_cache.shape[1:3] + group = heads // kv_heads + queries = query.float().reshape(batch, query_length, kv_heads, group, dim) + output = torch.zeros_like(queries) + positions = torch.arange(query_length, device=query.device) + + for seq, length in enumerate(context_lengths.cpu().tolist()): + if length <= 0: + continue + tokens = torch.arange(length, device=query.device) + pages = block_tables[seq, tokens // page_size].long() + offsets = tokens % page_size + keys = key_cache[pages, :, offsets, :].float() + values = value_cache[pages, :, :, offsets].float() + if key_scale.numel() == 1: + keys *= key_scale + values *= value_scale + else: + keys *= key_scale[pages, :, offsets, 0, None] + values *= value_scale[pages, :, offsets, 0, None] + scores = torch.einsum("qhgd,khd->qhgk", queries[seq], keys) * dim**-0.5 + visible = length - query_length + 1 + positions + masked = tokens[None, :] >= visible[:, None] + if sliding_window > 0: + masked |= tokens[None, :] < (visible - sliding_window)[:, None] + scores.masked_fill_(masked[:, None, None, :], float("-inf")) + log_denominator = torch.logsumexp(scores, dim=-1, keepdim=True) + if sinks is not None: + log_denominator = torch.logaddexp( + log_denominator, sinks.float().reshape(1, kv_heads, group, 1) + ) + # +inf sinks suppress finite logits; zero empty rows after -inf - -inf. + probs = torch.exp(scores - log_denominator) + probs.masked_fill_(visible[:, None, None, None] <= 0, 0) + output[seq] = torch.einsum("qhgk,khd->qhgd", probs, values) + return output.reshape_as(query) + + +def _make_inputs(case, planned=False): + torch.manual_seed(37 if case.sparse else 0) + batch, page, dim = len(case.lengths), case.block_size, case.head_dim + kv_heads, ql = case.num_kv_heads, case.query_length + heads = kv_heads * case.query_group_size + counts = [max(1, (length + page - 1) // page) for length in case.lengths] + num_pages = sum(counts) + quant_dtype = ( + torch.float8_e4m3fn if get_gfx_runtime() == "gfx950" else torch.float8_e4m3fnuz + ) + query = torch.empty((batch * ql, heads, dim), dtype=case.dtype).uniform_(-0.5, 0.5) + if case.masked_scale: + query.zero_() # Exercise zero-scale Q quantization. + key = torch.empty((num_pages, kv_heads, page, dim), dtype=case.dtype).uniform_( + -0.5, 0.5 + ) + value = torch.empty_like(key).uniform_(-0.5, 0.5) + quantize = pertoken_quant if case.per_token else per_tensor_quant + key_quant, key_scale = quantize(key, quant_dtype=quant_dtype) + value_quant, value_scale = quantize(value, quant_dtype=quant_dtype) + del key, value + + if case.per_token and case.sparse: + token = torch.arange(num_pages * page).reshape(num_pages, 1, page, 1) + head = torch.arange(kv_heads).reshape(1, kv_heads, 1, 1) + key_scale *= torch.exp2(((2 * token + head) % 4 - 2).float()) + factors = torch.exp2(((token + 2 * head) % 5 - 3).float()) + # Power-of-two scales isolate W=1 masking from FP8 P rounding. + value_scale = ( + factors * 2**-10 if case.sliding_window == 1 else value_scale * factors + ) + + start = 0 + for length, count in zip(case.lengths, counts): + end = start + count + if 0 < length <= 4 or (case.sink_dtype is not None and length in (257, 513)): + # Positive periodic V exposes repeated sink mass without cancellation; + # per-tensor cases retain a shared scale and use larger FP8 values. + token_values = (torch.arange(count * page) % 4 + 1).float() + token_values *= 0.25 if case.per_token else 32 + key_quant[start:end].zero_() + value_quant[start:end] = token_values.reshape(count, 1, page, 1).to( + quant_dtype + ) + if case.per_token: + key_scale[start:end].fill_(1) + value_scale[start:end].fill_(1) + if case.masked_scale: + assert case.per_token and case.sliding_window == 1 + tokens = torch.arange(count * page) + visible = (tokens >= max(0, length - ql)) & (tokens < length) + key_quant[start:end].zero_() + key_scale[start:end].fill_(1) + value_scale[start:end] = torch.where(visible, 1.0, 1e9).reshape( + count, 1, page, 1 + ) + values = torch.where(visible, (tokens % 4 + 1).float() * 0.25, 0) + value_quant[start:end] = values.reshape(count, 1, page, 1).to(quant_dtype) + start = end + + selected = ( + 2 * torch.randperm(num_pages) + 1 if case.sparse else torch.arange(num_pages) + ) + physical_pages = 2 * num_pages + 1 if case.sparse else num_pages + + def scatter(tensor): + if not case.sparse: + return tensor + result = torch.zeros((physical_pages, *tensor.shape[1:]), dtype=tensor.dtype) + result[selected] = tensor + return result + + key_quant, value_quant = scatter(key_quant), scatter(value_quant) + if case.per_token: + key_scale, value_scale = scatter(key_scale), scatter(value_scale) + key_cache = ( + key_quant.reshape(physical_pages, kv_heads, page, dim // 16, 16) + .permute(0, 1, 3, 2, 4) + .contiguous() + ) + value_cache = ( + value_quant.reshape(physical_pages, kv_heads, page // 16, 16, dim) + .permute(0, 1, 2, 4, 3) + .contiguous() + if case.trans_v + else value_quant.permute(0, 1, 3, 2).contiguous() + ) + table = torch.zeros((batch, max(counts)), dtype=torch.int32) + start = 0 + for seq, count in enumerate(counts): + table[seq, :count] = selected[start : start + count] + start += count + context = torch.tensor(case.lengths, dtype=torch.int32) + sinks = None + if case.sink_dtype is not None: + logits = [float("-inf"), float("inf"), -1000, 1000, -2, 0, 0.5, 5] + if case.sink_dtype == torch.float32: + logits += [-3e38, 3e38] + sinks = ( + torch.tensor(logits, dtype=case.sink_dtype) + .repeat((heads + len(logits) - 1) // len(logits))[:heads] + .contiguous() + ) + # Distinguish heads even when GQA is a multiple of the sentinel period. + offset = torch.arange(heads).to(case.sink_dtype) * 0.125 + sinks = torch.where(sinks.abs() < 10, sinks + offset, sinks) + parts = case.num_partitions + if parts is None and not planned: + parts = get_recommended_splits( + batch, + kv_heads, + KV_COMPUTE_BLOCK // page, + case.max_partitions, + max_context_length=max(case.lengths), + ) + plan = ( + plan_pa_decode( + context, + kv_heads, + max_partitions=parts, + workgroup_budget=case.workgroup_budget, + sliding_window=case.sliding_window, + query_length=ql if case.sliding_window > 0 else 1, + ) + if planned + else None + ) + if plan is not None: + parts = plan.max_partitions + rows = ql * case.query_group_size + shape = ( + (kv_heads, plan.capacity, rows) if planned else (batch, kv_heads, parts, rows) + ) + psum = torch.full(shape, float("nan"), dtype=torch.float32) + pmax = torch.full_like(psum, float("nan")) + pout = torch.full((*shape, dim), float("nan"), dtype=case.dtype) + args = ( + torch.full_like(query, float("nan")), + query, + key_cache, + value_cache, + context, + table, + dim**-0.5, + ql, + parts, + KV_COMPUTE_BLOCK, + quant_dtype, + None, + key_scale, + value_scale, + psum, + pmax, + pout, + None, + sinks, + ) + options = {"sliding_window": case.sliding_window, "work_plan": plan} + reference = partial( + run_torch, + query, + key_quant, + value_quant.permute(0, 1, 3, 2), + table, + context, + key_scale, + value_scale, + query_length=ql, + sliding_window=case.sliding_window, + sinks=sinks, + ) + return args, options, reference + + +def _run_flydsl(*args, sliding_window=0, work_plan=None): + if work_plan is not None: + plan_pa_decode( + args[4], + args[2].shape[1], + query_length=args[7], + sliding_window=sliding_window, + plan=work_plan, + ) + pa_decode(*args, sliding_window=sliding_window, work_plan=work_plan) + return args[0] + + +def _assert_close(output, reference): + assert torch.isfinite(output).all() + torch.testing.assert_close( + output.float(), + reference.float(), + atol=ACCURACY_TOLERANCE, + rtol=ACCURACY_TOLERANCE, + ) + + +def _assert_plan(plan, lengths): + """Independent integer oracle for absolute tiles, budgets and packed slots.""" + first = [ + ( + max(0, length - plan.query_length + 1 - plan.sliding_window) // 256 + if plan.sliding_window > 0 + else 0 + ) + for length in lengths + ] + last = [(max(0, length) + 255) // 256 for length in lengths] + tiles = [end - begin for begin, end in zip(first, last)] + remaining = plan.capacity - sum(count > 0 for count in tiles) + total = max(sum(tiles), 1) + prefix, work, reductions = 0, [], [] + for seq, (length, begin, count) in enumerate(zip(lengths, first, tiles)): + extra = (prefix + count) * remaining // total - prefix * remaining // total + parts = min(int(count > 0) + extra, count, plan.max_partitions) + prefix += count + reductions.append([len(work), parts]) + for part in range(parts): + work.append( + [ + seq, + begin + part * count // parts, + begin + (part + 1) * count // parts, + length, + ] + ) + assert len(work) <= plan.capacity + active = len(work) + work += [[0, 0, 0, 0]] * (plan.capacity - active) + assert plan.reduce_info.cpu().tolist() == reductions + assert plan.work_info.cpu().tolist() == work + return active + + +def _assert_contracts(args, options): + """Reuse numerical inputs for API validation.""" + heads = args[1].shape[1] + for invalid, error in [ + ([0.0] * heads, TypeError), + (torch.zeros(1, heads), ValueError), + (torch.zeros(heads - 1), ValueError), + (torch.zeros(heads, dtype=torch.int32), TypeError), + (torch.zeros(heads, dtype=torch.float64), TypeError), + (torch.zeros(heads, device="cpu"), ValueError), + (torch.zeros(heads * 2)[::2], ValueError), + ]: + with pytest.raises(error, match="sinks"): + pa_decode(*args[:-1], sinks=invalid, **options) + for value, error in [(-2, ValueError), (1.5, TypeError)]: + with pytest.raises(error, match="sliding_window"): + pa_decode(*args, **{**options, "sliding_window": value}) + with pytest.raises(error, match="sliding_window"): + plan_pa_decode(args[4], args[2].shape[1], sliding_window=value) + for value, error in [(0, ValueError), (1.5, TypeError)]: + with pytest.raises(error, match="query_length"): + pa_decode(*args[:7], value, *args[8:], **options) + with pytest.raises(error, match="query_length"): + plan_pa_decode(args[4], args[2].shape[1], query_length=value) + context, plan = args[4], options["work_plan"] + num_compute_units = torch.cuda.get_device_properties( + context.device + ).multi_processor_count + for parts in (0, num_compute_units + 1): + with pytest.raises(ValueError, match="max_partitions"): + plan_pa_decode(context, args[2].shape[1], max_partitions=parts) + if plan is not None: + with pytest.raises(ValueError, match="max_partitions"): + replace(plan, max_partitions=parts).validate( + context.numel(), args[2].shape[1], context.device + ) + if plan is None: + for parts in (0, MAX_CONTEXT_PARTITIONS + 1): + with pytest.raises(ValueError, match="max_context_partition_num"): + pa_decode(*args[:8], parts, *args[9:], **options) + if plan is not None: + context, heads = args[4], args[2].shape[1] + reuse = { + "sliding_window": plan.sliding_window, + "query_length": plan.query_length, + "plan": plan, + } + metadata = plan.work_info, plan.reduce_info + assert plan_pa_decode(context, heads, **reuse) is plan + assert plan.work_info is metadata[0] and plan.reduce_info is metadata[1] + changes = [ + ( + {"max_partitions": 1 if plan.max_partitions != 1 else 2}, + "max_partitions", + ), + ({"sliding_window": plan.sliding_window + 1}, "sliding_window"), + ({"workgroup_budget": 1}, "workgroup_budget"), + ] + if plan.sliding_window > 0: + changes.append(({"query_length": plan.query_length + 1}, "query_length")) + for change, message in changes: + with pytest.raises(ValueError, match=message): + plan_pa_decode(context, heads, **{**reuse, **change}) + for lengths, kv_heads in [(context.repeat(2), heads), (context, heads + 1)]: + with pytest.raises(ValueError): + plan_pa_decode(lengths, kv_heads, **reuse) + with pytest.raises(ValueError, match="sliding_window"): + pa_decode(*args, **{**options, "sliding_window": plan.sliding_window + 1}) + if plan.sliding_window > 0: + wrong_plan = plan_pa_decode( + context, + heads, + max_partitions=plan.max_partitions, + sliding_window=plan.sliding_window, + query_length=plan.query_length + 1, + ) + with pytest.raises(ValueError, match="query_length"): + pa_decode(*args, **{**options, "work_plan": wrong_plan}) + + +def _case( + name, + shape=(4, 1, 16, 128), + cache=(128, 1, 1), + parts=7, + window=0, + sink=None, + **kwargs, +): + # shape = (QL, KV heads, GQA, D); cache = (page size, transposed V, per-token). + ql, heads, group, dim = shape + page, trans_v, per_token = cache + return pytest.param( + DecodeCase( + query_length=ql, + num_kv_heads=heads, + query_group_size=group, + head_dim=dim, + block_size=page, + trans_v=bool(trans_v), + per_token=bool(per_token), + num_partitions=parts, + sliding_window=window, + sink_dtype=sink, + **kwargs, + ), + id=name, + ) + + +def _cases(columns, rows, *, prefix="", **shared): + """Expand parameter rows over shared defaults.""" + return [ + _case( + prefix + name, + **{**shared, **dict(zip(columns.split(), values, strict=True))}, + ) + for name, *values in rows + ] + + +BF16, FP16, FP32 = torch.bfloat16, torch.float16, torch.float32 +LENS_1024 = (0, 1, 1025, 1281) +LENS_4096 = (0, 1, 4097, 4353) +LENS_8192 = (0, 1, 8193, 8449) +CASES = [ + *_cases( + "shape cache parts window sink", + [ + ("scalar-direct", (1, 2, 8, 128), (16, 1, 0), 1, 0, None), + ("scalar-window-sinks", (1, 2, 8, 128), (128, 1, 0), 7, 257, FP32), + ("head1024", (1, 2, 4, 1024), (128, 1, 0), 256, 8192, FP32), + ("mtp3-window", (3, 1, 16, 128), (128, 0, 1), 7, 257, FP16), + ("mtp2-odd-parts", (2, 2, 16, 128), (16, 0, 1), 3, 0, None), + ("window256", (3, 2, 8, 128), (128, 0, 1), 7, 256, None), + ("window509", (4, 1, 16, 128), (64, 0, 1), 86, 509, FP16), + ("window8192", (2, 1, 16, 128), (128, 1, 1), 256, 8192, None), + ("head64-mtp-window", (4, 2, 4, 64), (16, 1, 1), 256, 1, FP32), + ("hkv2-direct-sinks", (1, 2, 8, 128), (128, 1, 1), 1, 1, BF16), + ], + ), + _case("head64-fp16", (2, 2, 4, 64), (16, 0, 0), 7, 257, BF16, dtype=FP16), + _case("scalar-fp16-mtp", cache=(128, 1, 0), parts=3, dtype=FP16), + _case( + "register-64-65", + (1, 2, 4, 256), + (64, 1, 0), + 86, + sink=FP16, + dtype=FP16, + lengths=(0, 1, 257, 16384, 16385), + ), + _case( + "np1-fused-sink", cache=(16, 1, 1), parts=1, window=1, sink=FP32, query_splits=1 + ), + _case("np1-split-sink", parts=1, window=1, sink=BF16, query_splits=4), + _case( + "split2-hkv2-full-m1", + (2, 2, 16, 128), + parts=17, + window=1024, + sink=FP32, + lengths=(0, 769, 1025, 2049), + query_splits=2, + wide_kv_addressing=True, + ), + _case( + "split4-hkv2-full-m1", + (4, 2, 16, 128), + (16, 0, 1), + 34, + 4096, + BF16, + lengths=(1, 257, 4353, 0), + query_splits=4, + ), + _case( + "full-m1-fp16-d256", + (2, 2, 8, 256), + (64, 1, 1), + 5, + 257, + FP16, + dtype=FP16, + lengths=(0, 1, 257, 513), + query_splits=1, + ), + _case( + "split4-hkv1-g15-plain-v-window-wide", + (4, 1, 15, 128), + (128, 0, 1), + 8, + 1024, + FP32, + lengths=LENS_1024, + workgroup_budget=24, + query_splits=4, + wide_kv_addressing=True, + ), + _case( + "disabled-window-wide", + cache=(16, 0, 1), + parts=256, + window=-1, + sink=FP32, + query_splits=1, + wide_kv_addressing=True, + ), + _case( + "fused-narrow", + cache=(16, 0, 1), + parts=1, + sink=FP32, + query_splits=1, + wide_kv_addressing=False, + ), + _case("fused-wide", parts=8, query_splits=1, wide_kv_addressing=True), + _case( + "fused-hkv2-window", + (4, 2, 16, 128), + (128, 0, 1), + 8, + 257, + BF16, + query_splits=1, + wide_kv_addressing=False, + ), + _case("mtp2-query-split", (2, 1, 16, 128), parts=3, lengths=(257, 259)), + *_cases( + "cache parts window lengths workgroup_budget", + [ + ("window-query-split", (16, 1, 1), 7, 1024, (1023, 1024, 1025, 4099), None), + ("window-capacity", (128, 1, 1), 256, 1024, (0, 3, 1024, 4099), None), + ("dense-capacity", (128, 0, 1), 256, 0, (257, 259, 1027, 4099), 17), + ("window-large-grid", (128, 1, 1), 7, 1024, (1027,) * 200, None), + ], + prefix="mtp4-", + ), + *_cases( + "parts window lengths workgroup_budget", + [ + ("window4096", 18, 4096, LENS_4096, 72), + ("window8192-batch12", 34, 8192, LENS_8192 * 3, 408), + ("window6000-batch5", 25, 6000, (0, 1, 6001, 6145, 6257), 125), + ("window4095-fallback", 18, 4095, LENS_4096, 72), + ("window8193-fallback", 34, 8193, LENS_8192, 136), + ("window1024-rejected", 6, 1024, LENS_1024 * 2, 48), + ("excess-partitions", 35, 8192, LENS_8192, 140), + ("tight-capacity", 34, 8192, LENS_8192, 135), + ( + "window4096-batch24", + 18, + 4096, + (0, 1, 3, 4, 257, 1025, 2049, 4096, 4097, 4098, 4353, 4354) * 2, + 432, + ), + ( + "window4096-np64", + 64, + 4096, + (0, 1, 769, 1793, 4096, 4097, 4353, 4354), + 144, + ), + ("window8192-np256-fallback", 256, 8192, (0, 1, 8449, 8450), 136), + ], + prefix="batch-first-", + sink=FP32, + ), + *_cases( + "parts lengths workgroup_budget", + [ + ("window8192", 34, (0, 1, 3, 4, 257, 8193, 8449, 0), 272), + ("excess-capacity", 35, LENS_8192, 137), + ], + prefix="batch-first-", + window=8192, + sink=BF16, + ), + *_cases( + "cache parts wide_kv_addressing", + [ + ("window4096-wide", (128, 1, 1), 18, True), + ("window4096-plain-v", (128, 0, 1), 18, None), + ("window4096-plain-v-wide", (128, 0, 1), 18, True), + ("window4096-np64-plain-v-wide", (128, 0, 1), 64, True), + ], + prefix="batch-first-", + window=4096, + sink=FP32, + lengths=LENS_4096, + workgroup_budget=72, + ), + _case( + "batch-first-window8192-plain-v-auto", + cache=(128, 0, 1), + parts=34, + window=8192, + sink=FP32, + lengths=LENS_8192, + workgroup_budget=136, + ), + *_cases( + "workgroup_budget wide_kv_addressing", + [ + ("prefetch", 512, None), + ("prefetch-small-capacity", 263, None), + ("prefetch-wide", 512, True), + ], + prefix="ql1-window-", + shape=(1, 1, 16, 128), + parts=64, + window=8192, + lengths=(8193,) * 8, + ), + *_cases( + "shape parts workgroup_budget", + [ + ("hkv2-g8-window-prefetch", (1, 2, 8, 128), 64, 264), + ("hkv1-g8-large-np-prefetch", (1, 1, 8, 128), 256, 132), + ], + prefix="ql1-", + window=8192, + sink=FP32, + lengths=LENS_8192, + ), + *_cases( + "shape cache workgroup_budget", + [ + ("hkv4-page128-multi", (1, 4, 16, 128), (128, 1, 1), 32), + ("hkv4-page16-multi", (1, 4, 16, 128), (16, 1, 1), 32), + ("hkv4-plain-v-multi", (1, 4, 16, 128), (128, 0, 1), 32), + ("hkv2-g15-plain-v-multi", (1, 2, 15, 128), (128, 0, 1), 16), + ], + prefix="ql1-", + parts=2, + window=1024, + sink=FP32, + lengths=LENS_1024, + ), + *_cases( + "shape workgroup_budget query_splits", + [ + ("ql1-hkv4-page16-direct-sinks", (1, 4, 16, 128), 16, None), + ("ql1-hkv2-g9-page16-direct-sinks", (1, 2, 9, 128), 8, None), + ("split2-hkv2-g9-page16-direct-sinks", (2, 2, 9, 128), 8, 2), + ], + cache=(16, 1, 1), + parts=1, + sink=FP32, + lengths=(0, 1, 257, 769), + ), + *_cases( + "shape", + [ + (f"ql1-hkv2-g{group}-window", (1, 2, group, 128)) + for group in (7, *range(9, 16), 17) + ], + parts=8, + window=1024, + sink=FP32, + lengths=LENS_1024, + workgroup_budget=40, + ), + *_cases( + "shape cache window sink lengths", + [ + ("1wg", (1, 2, 8, 128), (128, 1, 1), 1, FP32, (257,) * 16), + ("2wg", (1, 2, 8, 128), (128, 1, 1), 0, None, (257,) * 32), + ("page16", (1, 2, 8, 128), (16, 1, 1), 257, FP16, (257,) * 32), + ("plain-v", (1, 2, 16, 128), (128, 0, 1), 0, BF16, (257,) * 32), + ], + prefix="hkv2-prefetch-", + parts=8, + ), + _case( + "long-200k", + cache=(128, 0, 1), + parts=256, + sink=FP32, + lengths=(0, 1, 3, 4, 255, 256, 257, 200003), + ), + _case( + "long-np64", + cache=(16, 1, 1), + parts=64, + lengths=(0, 3, 257, 16384, 16385, 65537), + ), + _case("large-batch-auto", cache=(128, 0, 1), parts=None, lengths=(257,) * 200), + _case("small-batch-auto", (1, 1, 8, 128), (16, 0, 0), None, lengths=(257,) * 3), + _case("default-cu-window", parts=None, window=1024, sink=FP32, lengths=LENS_1024), + # CU > 256 gives the long row 257 reducer tasks. + _case( + "default-cu-head256", + (2, 1, 8, 256), + (64, 0, 1), + parts=None, + sink=FP16, + dtype=FP16, + lengths=(0, 1, 257, 65537), + ), + _case( + "exact-parts-override", + (1, 1, 16, 128), + parts=5, + lengths=(200000,), + max_partitions=4, + ), + _case("window-int64", cache=(16, 1, 1), window=2**40, wide_kv_addressing=True), + *_cases( + "window sink", + [("max", 2**31 - 1, None), ("overflow", 2**31, FP32)], + prefix="window-int32-", + cache=(16, 0, 1), + parts=1, + ), + _case( + "window-huge-padded-query", + (3, 2, 8, 128), + parts=1, + window=2**31 - 1, + sink=FP32, + lengths=(0, 1, 2, 257), + ), + _case( + "window255-budget", + (2, 2, 16, 128), + (64, 0, 1), + 64, + 255, + BF16, + workgroup_budget=17, + ), + # Inactive scales and NaN scratch must not contaminate live rows. + *_cases( + "shape sink", + [("decode", (1, 1, 16, 128), FP32), ("mtp", (4, 1, 16, 128), None)], + prefix="masked-scale-", + window=1, + lengths=(511,), + masked_scale=True, + ), + *_cases( + "parts window lengths workgroup_budget", + [ + ("np33-tail-empty", 33, 1, (1, 257, 513, 0, 0), 5), + ("shortcount-boundary", 34, 1024, (0, 769, 1025, 0), 24), + ("eightcount-boundary", 34, 2048, (0, 1793, 2049, 0), 48), + ], + prefix="reduce-", + sink=FP32, + ), + _case( + "reduce-np64-tail-padding", + (4, 2, 8, 128), + (16, 0, 1), + 64, + 255, + BF16, + lengths=(1, 3, 257, 258, 514, 515, 0, 0), + workgroup_budget=32, + ), + _case("empty", window=1, sink=FP16, lengths=(0,) * 4), +] + + +@pytest.mark.parametrize("planned", [False, True], ids=["static", "planned"]) +@pytest.mark.parametrize("case", CASES) +def test_pa_decode(case, planned, monkeypatch): + """Check numerics, contracts and graph replays.""" + if planned and case.num_partitions is not None: + context = torch.tensor(case.lengths, dtype=torch.int32) + num_compute_units = torch.cuda.get_device_properties( + context.device + ).multi_processor_count + if case.num_partitions > num_compute_units: + with pytest.raises(ValueError, match="max_partitions"): + plan_pa_decode( + context, case.num_kv_heads, max_partitions=case.num_partitions + ) + return + args, options, reference = _make_inputs(case, planned) + if not planned and case.sliding_window > 0: + with pytest.raises(ValueError, match="work_plan"): + pa_decode(*args, **options) + with pytest.raises(ValueError, match="work_plan"): + _run_flydsl(*args, **options) + return + output, context = args[0], args[4] + scratch, sinks, plan = args[14:17], args[-1], options["work_plan"] + if plan is not None: + num_compute_units = torch.cuda.get_device_properties( + context.device + ).multi_processor_count + expected_parts = ( + num_compute_units if case.num_partitions is None else case.num_partitions + ) + assert args[8] == plan.max_partitions == expected_parts + budget = ( + case.workgroup_budget + if case.workgroup_budget is not None + else 2 * num_compute_units + ) + budget_slots = (budget + case.num_kv_heads - 1) // case.num_kv_heads + assert plan.capacity == min( + context.numel() * args[8], max(context.numel(), budget_slots) + ) + + module = importlib.import_module("aiter.ops.flydsl.pa_decode") + overrides = { + name: getattr(case, name) + for name in ("query_splits", "wide_kv_addressing") + if getattr(case, name) is not None + } + if overrides: + build_tile = module.compile_pa_decode_tile + + def compile_tile(**kwargs): + return build_tile(**{**kwargs, **overrides}) + + monkeypatch.setattr(module, "compile_pa_decode_tile", compile_tile) + if not planned and args[8] == 1: + + def unexpected_reducer(*_args, **_kwargs): + raise AssertionError("static NP=1 must not launch a reducer") + + monkeypatch.setattr(module, "launch_pa_decode_ps_reduce", unexpected_reducer) + + def check(disabled_sink=False): + _assert_close(output, reference(sinks=None) if disabled_sink else reference()) + visible = ( + context[:, None] + - case.query_length + + 1 + + torch.arange(case.query_length, device=context.device) + ) + assert (output[(visible <= 0).flatten()] == 0).all() + if sinks is not None: + assert (output[:, torch.isposinf(sinks)] == 0).all() + if plan is not None: + active = _assert_plan(plan, context.cpu().tolist()) + for tensor in scratch: + assert torch.isnan(tensor[:, active:]).all() + elif args[8] == 1: + assert all(torch.isnan(tensor).all() for tensor in scratch) + + _run_flydsl(*args, **options) + check() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + _run_flydsl(*args, **options) + original_sinks = sinks.clone() if sinks is not None else None + for step, lengths in enumerate( + ( + tuple(max(0, length - 7) for length in case.lengths), + (0,) * len(case.lengths), + case.lengths, + ) + ): + context.copy_(torch.tensor(lengths, dtype=torch.int32)) + if sinks is not None: + if step == 2: + sinks.fill_(float("-inf")) + else: + sinks.copy_(original_sinks.roll(step + 1)) + for tensor in (output, *scratch): + tensor.fill_(float("nan")) + graph.replay() + check(disabled_sink=step == 2) + + _assert_contracts(args, options) + if plan is not None: + # Check int32 limits without allocating a matching KV cache. + lengths = (-1, 0, 1, 257, 2**31 - 1) + extreme_plan = plan_pa_decode( + torch.tensor(lengths, dtype=torch.int32), + case.num_kv_heads, + max_partitions=args[8], + sliding_window=case.sliding_window, + query_length=case.query_length, + ) + _assert_plan(extreme_plan, lengths) + + +@benchmark() +def run_pa_decode_tile_case( + batch_size, + num_query_heads, + num_kv_heads, + head_dim, + context_length, + block_size, + dtype, + trans_v, + max_partitions=None, + per_token=False, + query_length=1, + num_partitions=None, +): + if min(batch_size, context_length, query_length, num_query_heads, num_kv_heads) < 1: + raise ValueError( + "batch, context, query length and head counts must be positive" + ) + if num_query_heads % num_kv_heads: + raise ValueError("num_query_heads must be divisible by num_kv_heads") + if dtype not in (dtypes.fp16, dtypes.bf16): + raise ValueError("pa_decode only supports fp16/bf16") + if num_partitions is not None and not 1 <= num_partitions <= MAX_CONTEXT_PARTITIONS: + raise ValueError(f"num_partitions must be in [1, {MAX_CONTEXT_PARTITIONS}]") + case = DecodeCase( + lengths=(context_length,) * batch_size, + query_length=query_length, + num_kv_heads=num_kv_heads, + query_group_size=num_query_heads // num_kv_heads, + head_dim=head_dim, + block_size=block_size, + dtype=dtype, + trans_v=trans_v, + num_partitions=num_partitions, + max_partitions=max_partitions, + per_token=per_token, + sparse=False, + ) + args, options, reference_call = _make_inputs(case) + reference = reference_call() + del reference_call + # Keep tensors positional so allocation rotation accounts for their memory. + output, us = run_perftest(_run_flydsl, *args, **options) + _assert_close(output, reference) + query, table = args[1], args[5] + attended = sum( + max(0, context_length - query_length + 1 + p) for p in range(query_length) + ) + flops = 4 * batch_size * num_query_heads * attended * head_dim + scale_elements = batch_size * num_kv_heads * context_length if per_token else 1 + nbytes = ( + 2 * query.numel() * query.element_size() + + 2 + * batch_size + * num_kv_heads + * context_length + * head_dim + * args[2].element_size() + + table.numel() * table.element_size() + + args[4].numel() * args[4].element_size() + + scale_elements * (args[12].element_size() + args[13].element_size()) + ) + return { + "gfx": get_gfx_runtime(), + "partitions": args[8], + "trans_v": trans_v, + "per_token": per_token, + "flydsl us": us, + "flydsl us/token": us / (batch_size * query_length), + "flydsl TFLOPS": flops / us / 1e6, + "flydsl TB/s": nbytes / us / 1e6, + "flydsl err": 0, + } + + +def _positive_int(value): + value = int(value) + if value < 1: + raise argparse.ArgumentTypeError("must be a positive integer") + return value + + +def _parse_args(argv=None): + parser = argparse.ArgumentParser( + description="FlyDSL pa_decode correctness + perf sweep" + ) + parser.add_argument( + "-d", "--dtype", type=dtypes.str2Dtype, nargs="*", default=[dtypes.bf16] + ) + parser.add_argument( + "-b", "--batch", type=_positive_int, nargs="*", default=DEFAULT_BATCH_SIZES + ) + parser.add_argument( + "-q", + "--query-length", + type=_positive_int, + nargs="+", + default=[1], + help="Query tokens per sequence; context lengths include them.", + ) + parser.add_argument( + "-s", + "--shapes", + type=dtypes.str2tuple, + nargs="*", + default=DEFAULT_SHAPES, + help="num_query_heads,num_kv_heads,head_dim,context_length", + ) + parser.add_argument( + "--block-size", + type=int, + nargs="*", + choices=[16, 64, 128], + default=[16, 64, 128], + ) + parser.add_argument( + "--trans-v", + type=int, + nargs="*", + choices=[0, 1], + default=[0, 1], + help="0: plain V cache; 1: transposed V cache.", + ) + parser.add_argument( + "--per-token", + type=int, + nargs="*", + choices=[0, 1], + default=[0], + help="0: per-tensor KV scales; 1: per-token KV scales.", + ) + parser.add_argument( + "--max-partitions", + type=int, + default=None, + help="Upper clamp for automatic splits (4..256); 8 keeps the legacy clamp.", + ) + parser.add_argument( + "--num-partitions", + type=_positive_int, + nargs="+", + default=[None], + help="Exact split counts (1..256), overriding --max-partitions.", + ) + args = parser.parse_args(argv) + if ( + args.max_partitions is not None + and not 4 <= args.max_partitions <= MAX_CONTEXT_PARTITIONS + ): + parser.error(f"--max-partitions must be in [4, {MAX_CONTEXT_PARTITIONS}]") + if any(n is not None and n > MAX_CONTEXT_PARTITIONS for n in args.num_partitions): + parser.error(f"--num-partitions must be in [1, {MAX_CONTEXT_PARTITIONS}]") + for shape in args.shapes: + if not isinstance(shape, tuple) or len(shape) != 4 or any(n < 1 for n in shape): + parser.error("each --shapes value must contain four positive integers") + if shape[0] % shape[1]: + parser.error("num_query_heads must be divisible by num_kv_heads") + return args + + +def main(): + # Keep --help and argument errors usable without a GPU. + args = _parse_args() + if not torch.cuda.is_available(): + aiter.logger.warning("ROCm is not available; skipping pa_decode") + return + if get_gfx_runtime() not in SUPPORTED_GFX or pa_decode is None: + aiter.logger.warning("FlyDSL pa_decode is unavailable or unsupported; skipping") + return + torch.set_default_device("cuda") + rows = [] + for dtype, batch, shape, page, trans_v, per_token, ql, parts in itertools.product( + args.dtype, + args.batch, + args.shapes, + args.block_size, + args.trans_v, + args.per_token, + args.query_length, + args.num_partitions, + ): + heads, kv_heads, dim, context = shape + rows.append( + run_pa_decode_tile_case( + batch, + heads, + kv_heads, + dim, + context, + page, + dtype, + bool(trans_v), + args.max_partitions, + bool(per_token), + query_length=ql, + num_partitions=parts, + ) + ) + aiter.logger.info( + "pa_decode summary (markdown):\n%s", pd.DataFrame(rows).to_markdown(index=False) + ) + + +if __name__ == "__main__": + main()