From c6be02ea02425950c967be45e71deb5ec3f1d6ec Mon Sep 17 00:00:00 2001 From: fsx950223 Date: Wed, 22 Jul 2026 08:17:39 +0000 Subject: [PATCH 01/66] feat(flydsl): Add paged-attention Tile kernel Signed-off-by: fsx950223 Co-authored-by: Cursor --- aiter/ops/flydsl/__init__.py | 2 + aiter/ops/flydsl/kernels/dpp_utils.py | 23 + aiter/ops/flydsl/kernels/mem_ops.py | 202 +++ aiter/ops/flydsl/kernels/utils.py | 97 ++ aiter/ops/flydsl/pa_decode_tile.py | 1477 ++++++++++++++++++ op_tests/flydsl_tests/test_pa_decode_tile.py | 160 ++ 6 files changed, 1961 insertions(+) create mode 100644 aiter/ops/flydsl/kernels/mem_ops.py create mode 100644 aiter/ops/flydsl/kernels/utils.py create mode 100644 aiter/ops/flydsl/pa_decode_tile.py create mode 100644 op_tests/flydsl_tests/test_pa_decode_tile.py diff --git a/aiter/ops/flydsl/__init__.py b/aiter/ops/flydsl/__init__.py index fa12af5cdf2..5b4fbee1170 100644 --- a/aiter/ops/flydsl/__init__.py +++ b/aiter/ops/flydsl/__init__.py @@ -82,6 +82,7 @@ ".gather_kv_b_proj", "gather_kv_b_proj_flydsl_supported", ), + "pa_decode_tile": (".pa_decode_tile", "pa_decode_tile"), } __all__ = [ @@ -104,6 +105,7 @@ "flydsl_qk_norm_rope_quant", "gather_kv_b_proj_flydsl", "gather_kv_b_proj_flydsl_supported", + "pa_decode_tile", ] diff --git a/aiter/ops/flydsl/kernels/dpp_utils.py b/aiter/ops/flydsl/kernels/dpp_utils.py index 8226da6a0f8..2eadaefb239 100644 --- a/aiter/ops/flydsl/kernels/dpp_utils.py +++ b/aiter/ops/flydsl/kernels/dpp_utils.py @@ -55,3 +55,26 @@ 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._mlir.dialects import arith as _arith_dialect + from flydsl.expr.typing import T + + src_i32 = _to_ir(src).bitcast(T.i32) + 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 _arith_dialect.BitcastOp(T.f32, out_i32).result diff --git a/aiter/ops/flydsl/kernels/mem_ops.py b/aiter/ops/flydsl/kernels/mem_ops.py new file mode 100644 index 00000000000..bc0d08e2ca1 --- /dev/null +++ b/aiter/ops/flydsl/kernels/mem_ops.py @@ -0,0 +1,202 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (C) 2025-2026 FlyDSL Project Contributors + +"""Single home for kernel-facing memory & atomic helpers. + +Consolidates the LLVM-pointer construction, global load/store, and atomic +primitives that were previously scattered across ``kernels_common.py`` and +``utils.py`` (and hand-rolled inline in individual kernels). The original +modules re-export these names, so existing imports keep working unchanged. + +Two distinct atomic mechanisms live here and are NOT interchangeable: + * ``atomic_add`` - LLVM ``atomicrmw`` on an ``!llvm.ptr`` (pointer atomics) + * ``buffer_atomic_add`` - AMD ``raw.ptr.buffer.atomic.fadd`` on a buffer resource +""" + +import flydsl.expr as fx +from flydsl._mlir import ir +from flydsl._mlir.dialects import arith as _std_arith +from flydsl._mlir.dialects import fly as _fly +from flydsl._mlir.dialects import llvm as _llvm +from flydsl.expr import arith as _expr_arith +from flydsl.expr import buffer_ops, const_expr, rocdl +from flydsl.expr.typing import T + +# Public API. Prefer the unified names; the legacy aliases at the bottom of this +# list exist only so pre-consolidation imports keep working (re-exported by +# kernels_common / utils) and should not be used in new code: +# element_ptr <- get_llvm_ptr aligned_ptr <- extract_global_ptr +# to_llvm_ptr <- _create_llvm_ptr +__all__ = [ + # pointer construction / extraction (preferred) + "element_ptr", + "to_llvm_ptr", + "aligned_ptr", + "global_ptr_from_addr", + # global load / store + "global_load", + "global_load_i32", + "global_load_i64x2", + "global_store", + # atomics (two distinct mechanisms) + "atomic_add", + "buffer_atomic_add", + # legacy aliases (back-compat only) + "get_llvm_ptr", + "_create_llvm_ptr", + "extract_global_ptr", +] + + +# --------------------------------------------------------------------------- # +# Pointer construction / extraction +# --------------------------------------------------------------------------- # +def get_llvm_ptr(ptr, offset, dtype_bytes, ptr_type=None): + """Build a global (address-space 1) ``!llvm.ptr`` at ``ptr + offset*dtype_bytes``. + + Shared home for the LLVM-ptr arithmetic used by atomic/global accesses + (previously duplicated in hgemm_splitk.py, small_m_hgemm.py, splitk_hgemm.py + and rmsnorm_kernel.py). + """ + if ptr_type is None: + ptr_type = ir.Type.parse("!llvm.ptr<1>") + base_ptr = _fly.extract_aligned_pointer_as_index(ptr_type, ptr) + base_ptr = _llvm.PtrToIntOp(T.i64, base_ptr).result + byte_offset = _expr_arith.index_cast( + T.i64, fx.Index(offset) * fx.Index(dtype_bytes) + ) + llvm_ptr = _llvm.AddOp(base_ptr, byte_offset, _llvm.IntegerOverflowFlags(0)).result + llvm_ptr = _llvm.IntToPtrOp(ptr_type, llvm_ptr).result + return llvm_ptr._value if const_expr(hasattr(llvm_ptr, "_value")) else llvm_ptr + + +# Unified name for the element-pointer helper above. +element_ptr = get_llvm_ptr + + +def _create_llvm_ptr(value, address_space: int = 1): + value = buffer_ops._unwrap_value(value) + if isinstance(value.type, ir.IndexType): + i64_type = T.i64 + value = buffer_ops._unwrap_value(_std_arith.IndexCastOp(i64_type, value).result) + ptr_type = ir.Type.parse(f"!llvm.ptr<{address_space}>") + return _llvm.IntToPtrOp(ptr_type, value).result + + +# Unified name: int/index scalar -> typed llvm.ptr (any address space). +to_llvm_ptr = _create_llvm_ptr + + +def extract_global_ptr(tensor): + raw = ( + tensor.ir_value() + if hasattr(tensor, "ir_value") and not isinstance(tensor, ir.Value) + else tensor + ) + ptr_type = ir.Type.parse("!llvm.ptr<1>") + return _fly.extract_aligned_pointer_as_index(ptr_type, raw) + + +# Unified name: tensor/memref -> aligned global base pointer. +aligned_ptr = extract_global_ptr + + +def global_ptr_from_addr(addr_i64): + raw = ( + addr_i64.ir_value() + if hasattr(addr_i64, "ir_value") and not isinstance(addr_i64, ir.Value) + else addr_i64 + ) + return _llvm.IntToPtrOp(ir.Type.parse("!llvm.ptr<1>"), raw).result + + +# --------------------------------------------------------------------------- # +# Global load / store +# --------------------------------------------------------------------------- # +def global_load(global_ptr, byte_offset, result_type, *, alignment): + ptr = buffer_ops.get_element_ptr( + global_ptr, byte_offset=fx.Int64(byte_offset), elem_type=T.i8 + ) + return _llvm.LoadOp(result_type, ptr, alignment=alignment).result + + +def global_load_i64x2(global_ptr, byte_offset_i64): + return global_load(global_ptr, byte_offset_i64, T.i64x2, alignment=16) + + +def global_load_i32(global_ptr, elem_offset_i32): + return global_load( + global_ptr, fx.Int64(elem_offset_i32) * fx.Int64(4), T.i32, alignment=4 + ) + + +def global_store(global_ptr, byte_offset, value, *, alignment): + """Store ``value`` at ``global_ptr + byte_offset`` (mirror of ``global_load``). + + Unwraps a FlyDSL expression operand the same way ``atomic_add`` does, so + callers may pass either a raw ``ir.Value`` or a FlyDSL expression. New in the + consolidation; the hand-rolled ``llvm.StoreOp`` sites are collected in a later + phase, so this has no call site yet. + """ + val = value.ir_value() if const_expr(hasattr(value, "ir_value")) else value + ptr = buffer_ops.get_element_ptr( + global_ptr, byte_offset=fx.Int64(byte_offset), elem_type=T.i8 + ) + _llvm.StoreOp(val, ptr, alignment=alignment) + + +# --------------------------------------------------------------------------- # +# Atomics -- two distinct mechanisms (see module docstring) +# --------------------------------------------------------------------------- # +def atomic_add( + dst, + offset, + value, + *, + dtype_bytes=4, + syncscope="agent", + ordering=None, + alignment=None, + ptr_type=None, +): + """Atomically add ``value`` into ``dst[offset]`` in global memory. + + inline (e.g. the rmsnorm backward ``dweight`` accumulation). Selects + ``fadd`` for a floating-point operand and integer ``add`` otherwise, from + the operand's IR type, so a single call covers both cases. Returns the + atomicrmw result (the value previously stored at ``dst[offset]``). + + ``dtype_bytes`` sizes the byte offset and, unless ``alignment`` is given, is + reused as the access alignment. + """ + ptr = get_llvm_ptr(dst, offset, dtype_bytes, ptr_type=ptr_type) + val = value.ir_value() if const_expr(hasattr(value, "ir_value")) else value + elem_ty = val.type.element_type if isinstance(val.type, ir.VectorType) else val.type + bin_op = ( + _llvm.AtomicBinOp.fadd + if isinstance(elem_ty, ir.FloatType) + else _llvm.AtomicBinOp.add + ) + if ordering is None: + ordering = _llvm.AtomicOrdering.monotonic + if alignment is None: + alignment = dtype_bytes + return _llvm.AtomicRMWOp( + bin_op, + ptr, + val, + ordering, + syncscope=syncscope, + alignment=alignment, + ).result + + +def buffer_atomic_add(vdata, rsrc, offset, soffset, aux): + """Buffer-resource atomic fadd (AMD ``raw.ptr.buffer.atomic.fadd``). + + Single entry point for the buffer-resource atomic accumulation used by the + MoE gemm/blockscale wrappers (``atomic_add_f16x2 / _f32 / _x2``) and conv3d. + Distinct from ``atomic_add``: operates on a buffer resource + byte offset, + not an ``!llvm.ptr``, so the two are not interchangeable. + """ + return rocdl.raw_ptr_buffer_atomic_fadd(vdata, rsrc, offset, soffset, aux) diff --git a/aiter/ops/flydsl/kernels/utils.py b/aiter/ops/flydsl/kernels/utils.py new file mode 100644 index 00000000000..e8633e08ccb --- /dev/null +++ b/aiter/ops/flydsl/kernels/utils.py @@ -0,0 +1,97 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (C) 2025-2026 FlyDSL Project Contributors + +import flydsl.expr as fx +from flydsl._mlir import ir +from flydsl._mlir.dialects import llvm +from flydsl.expr import arith, const_expr, rocdl +from flydsl.expr.typing import T + +# Pointer/global-load helpers now live in mem_ops; re-exported here for back-compat. +from .mem_ops import extract_global_ptr as extract_global_ptr +from .mem_ops import global_load as global_load +from .mem_ops import global_load_i32 as global_load_i32 +from .mem_ops import global_load_i64x2 as global_load_i64x2 +from .mem_ops import global_ptr_from_addr as global_ptr_from_addr + + +def rcp_f32(value): + return rocdl.rcp(T.f32, value) + + +def exp2_amdgcn_scalar(scalar_value): + raw = ( + arith.unwrap(scalar_value) + if hasattr(scalar_value, "ir_value") or hasattr(scalar_value, "type") + else scalar_value + ) + f32_ty = ir.F32Type.get() + return llvm.call_intrinsic(f32_ty, "llvm.amdgcn.exp2.f32", [raw], [], []) + + +def exp2_f32_fast(value): + from flydsl._mlir.dialects import vector as _vector_dialect + + raw = ( + arith.unwrap(value) + if hasattr(value, "ir_value") or hasattr(value, "type") + else value + ) + ty = raw.type + if isinstance(ty, ir.VectorType): + n = ty.shape[0] + elems = [] + for i in range(n): + scalar = _vector_dialect.extract( + raw, static_position=[i], dynamic_position=[] + ) + elems.append(exp2_amdgcn_scalar(scalar)) + return _vector_dialect.from_elements(ty, elems) + return exp2_amdgcn_scalar(raw) + + +def cdiv(numer: int, denom: int) -> int: + return (numer + denom - 1) // denom + + +# Alias: several kernels historically spelled this ``ceildiv``. +ceildiv = cdiv + + +def align_up(value: int, align: int) -> int: + """Round *value* up to the next multiple of *align* (static ints).""" + return ((int(value) + int(align) - 1) // int(align)) * int(align) + + +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) + + +def unflatten_k(k_flat, qkhe_loop: int = 2): + n = qkhe_loop * 2 + return [[k_flat[td * n + j] for j in range(n)] for td in range(len(k_flat) // n)] diff --git a/aiter/ops/flydsl/pa_decode_tile.py b/aiter/ops/flydsl/pa_decode_tile.py new file mode 100644 index 00000000000..e867720f5f3 --- /dev/null +++ b/aiter/ops/flydsl/pa_decode_tile.py @@ -0,0 +1,1477 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (C) 2025-2026 FlyDSL Project Contributors + +"""Readable tile-programming reference for paged-attention fp8 decode. + +K/V are fp8 e4m3 (FNUZ on gfx942, OCP on gfx950) fed straight into +``mfma_f32_16x16x32_fp8_fp8``; Q (bf16/f16) and the softmax probabilities P are +quantized to fp8 too. Scales fold out of the matmuls (q/key scale into the QK +score, value scale + 1/FP8_MAX into the epilogue); softmax max/sum stay f32. +``key_scale``/``value_scale`` are either a ``[1]`` per-tensor scalar or a +``[num_blocks, num_kv_heads, block_size]`` per-token tensor (chosen by rank). + +``block_size`` (16/64) and ``head_dim`` (multiple of 64) are compile-time +constants. Layouts are logical, not production's preshuffle. + +* ``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 + +One CTA (4 waves) per (seq, kv_head) runs a flash-style online softmax over +256-token blocks; the 4 waves split tokens for Q.KT and head-dim for P.V, with +an LDS round-trip on P transposing ownership between the two MMAs. +""" + +import functools + +import torch + +import flydsl.compiler as flyc +import flydsl.expr as fx +from flydsl.compiler.protocol import dsl_size_of +from flydsl.expr import arith, buffer_ops, const_expr, gpu, range_constexpr, vector +from flydsl.expr import math as fmath +from flydsl.expr.typing import T +from flydsl.expr.vector import ReductionOp +from flydsl.runtime.device import get_rocm_arch + +from .kernels import dpp_utils +from .kernels.tensor_shim import _run_compiled +from .kernels.utils import cdiv, exp2_amdgcn_scalar, exp2_f32_fast, rcp_f32 + +MFMA_MNK = ( + 16 # M = N = 16 for the MMA atom; also query rows handled per CTA (padded to 16) +) +MFMA_K = 32 # fp8 MFMA contracts K = 32 per instruction (mfma_f32_16x16x32_fp8_fp8) +WAVE = 64 +LOG2E = 1.4426950408889634 +KV_COMPUTE_BLOCK = 256 + + +@functools.lru_cache(maxsize=None) +def compile_pa_decode_tile( + *, + head_dim: int, + query_group_size: int, + block_size: 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, +): + """Build the tile-programming PA-decode kernel + launch wrapper. + + ``block_size``, ``head_dim``, and ``query_dtype`` are compile-time + constants (lru_cache keys). ``query_length`` (MTP) and ``query_group_size`` + flatten into ``TOTAL_ROWS = query_length * query_group_size``, tiled into + ``M_TILES = ceil(TOTAL_ROWS / 16)`` independent 16-row MFMA tiles; each + extra M-tile duplicates loop-carried state, so VGPR/LDS/occupancy scale + roughly linearly with ``M_TILES``. + """ + is_gfx950 = "gfx95" in get_rocm_arch() + 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, + ), f"pa_decode_tile only supports block_size in (16, 64), 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 query_dtype == "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}" + assert query_length >= 1, f"query_length must be >= 1, got {query_length}" + # Flattened query-row axis (MTP outer, GQA head inner), tiled into 16-row M-tiles. + TOTAL_ROWS = query_length * query_group_size + M_TILES = cdiv(TOTAL_ROWS, MFMA_MNK) + ROWS_PADDED = M_TILES * MFMA_MNK + # PV layout: V=A, P=B -> output [head-dim (row), query-row (col=lane16)], + # generalized over head_dim via the VHE_CHUNKS loop. + NWARP = 4 # 4 waves / CTA + TOK_PER_WARP = 64 # tokens each warp owns per compute block (matches production KV_COMPUTE_BLOCK) + TILE_TOK = NWARP * TOK_PER_WARP # 256 tokens / compute block + PAGES_PER_CHUNK = ( + TOK_PER_WARP // block_size + ) # pages spanned by one 64-token warp-chunk: 1 (bs=64) or 4 (bs=16) + assert ( + head_dim % (NWARP * MFMA_MNK) == 0 + ), "head_dim must split across the 4 warps for PV" + + # head_dim splits into 16-element chunks (QK_CHUNK_ELEMS, one dwordx4 load); + # RGROUP_QUARTERS of them make a 64-element fetch group, QKHE_LOOP groups total. + 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}" + N_SUBCHUNKS = QKHE_LOOP * (QK_CHUNK_ELEMS // 8) + + # Q-quant chunk width: NQCHUNK stays fixed at 16 (tied to `lane16`'s role + # as the absmax butterfly width); QCHUNK scales with head_dim instead. + NQCHUNK = 16 + QCHUNK = ( + head_dim // NQCHUNK + ) # f16 elements per lane's load chunk (8 for head_dim=128, 4 for head_dim=64) + + VHE_CHUNKS = head_dim // (NWARP * MFMA_MNK) # 2 for head_dim=128, 1 for head_dim=64 + + if softmax_scale is None: + softmax_scale = 1.0 / (head_dim**0.5) + NP = int(num_partitions) # context partitions (grid.z); compile-time constant + + BLOCK_THREADS = NWARP * WAVE # 256 + + # -- LDS layout (shared across the 4 warps) -- + # sQ: fp8[ROWS_PADDED,head_dim] staged+quantized query. sP: fp8[16,TILE_TOK] + # quantized probs. sQscale: f32[ROWS_PADDED]. sLmax/sLsum: cross-warp scratch. + # sVPage: V page-index broadcast. sKScale/sVScale/sVScaleMax: per-token K/V + # scale staging. No sO/sM/sL/sCorr: PV output is register-resident/loop-carried + # and stored straight to global (V=A/P=B swap). + f32 = 4 + sQ_bytes = ROWS_PADDED * head_dim * 1 # fp8 + sP_off = sQ_bytes + # +16B row padding breaks a 32-bank LDS conflict on the P-pack writes while + # keeping the row 16B-aligned for PV's ds_read_b128. + SP_ROW_BYTES = TILE_TOK + 16 + sP_bytes = MFMA_MNK * SP_ROW_BYTES # fp8, padded rows + sQscale_off = sP_off + sP_bytes + NWARP_PAD = ( + NWARP + 1 + ) # coprime with 32 banks -> no LDS bank conflict on the cross-warp scratch + # Phase-split slices sLmax per M-tile so all pass-1 writes share one barrier. + sLmax_off = sQscale_off + ROWS_PADDED * f32 + sLsum_off = sLmax_off + M_TILES * MFMA_MNK * NWARP_PAD * f32 + # V page-index broadcast (V's page depends on `rgroup`, shared across warps). + sVPage_off = sLsum_off + MFMA_MNK * NWARP_PAD * f32 + sVPage_bytes = NWARP * PAGES_PER_CHUNK * 4 # i32 + # Per-token KV-scale staging (per_token only), double-buffered so the tt+1 + # prefetch (into the other buffer) doesn't clobber the current tile's scales. + KV_BUFS = 2 if per_token_kv else 1 + KV_BUF_STRIDE = 2 * NWARP * TOK_PER_WARP * f32 # k-region + v-region, one buffer + sKScale_off = sVPage_off + sVPage_bytes + sVScale_off = sKScale_off + NWARP * TOK_PER_WARP * f32 + sKVScale_bytes = KV_BUFS * 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 + + # LDS blob: one i32 array carved into typed byte-offset views (see the + # `_lds_*` helpers in the kernel). Every region size above is 4-byte + # aligned, so `total_bytes // 4` covers the blob exactly. + @fx.struct + class SharedStorage: + buf: fx.Array[fx.Int32, total_bytes // 4, 16] + + @flyc.kernel(known_block_size=(BLOCK_THREADS, 1, 1)) + def pa_decode_tile_kernel( + output_ptr: fx.Tensor, # [num_seqs*query_length, num_q_heads, head_dim] (written directly when NP==1) + # per-partition partial outputs (combined by the reduce kernel when NP>1): + pmax_ptr: fx.Tensor, # [num_seqs*query_length, num_kv_heads, num_partitions, query_group_size] row max + psum_ptr: fx.Tensor, # [num_seqs*query_length, num_kv_heads, num_partitions, query_group_size] row sum + pout_ptr: fx.Tensor, # [num_seqs*query_length, num_kv_heads, num_partitions, query_group_size, head_dim] Q_DTYPE, normalized O_p/l_p + query_ptr: fx.Tensor, # [num_seqs*query_length, num_q_heads, head_dim] -- row = seq*query_length + qi (MTP position) + key_cache_ptr: fx.Tensor, # [num_blocks, num_kv_heads, head_dim//16, block_size, 16] (blocked, see module docstring) + value_cache_ptr: fx.Tensor, # [num_blocks, num_kv_heads, block_size//16, head_dim, 16] (blocked, see module docstring) + block_tables_ptr: fx.Tensor, # [num_seqs, max_blocks_per_seq] + context_lengths_ptr: fx.Tensor, # [num_seqs] + key_scale_ptr: fx.Tensor, # [1] per-tensor OR [num_blocks, num_kv_heads, block_size] per-token + value_scale_ptr: fx.Tensor, # same shape as key_scale_ptr + max_blocks_per_seq: fx.Int32, + stride_ks_block: fx.Int32, + stride_ks_head: fx.Int32, + stride_q_row: fx.Int32, + stride_q_head: 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_h = fx.Int32(gpu.block_id("y")) + part = fx.Int32(gpu.block_id("z")) # context partition handled by this CTA + n_kv = fx.Int32(gpu.grid_dim.y) # num_kv_heads == gridDim.y + + # fx.copy-based flat loaders: K/V/context_len over a raw pointer, Q over + # a buffer resource. + def _make_flat_loader(tensor_ptr, elem_ty, reg_width, copy_op): + use_buffer_resource = isinstance( + copy_op, fx.rocdl.CopyOpCDNA3BufferCopyType + ) + copy_atom = fx.make_copy_atom(copy_op, elem_ty) + reg = fx.make_rmem_tensor(fx.make_layout(reg_width, 1), elem_ty) + base_iter = ( + fx.get_iter(fx.rocdl.make_buffer_tensor(tensor_ptr, max_size=True)) + if use_buffer_resource + else fx.get_iter(tensor_ptr) + ) + flat = fx.Tensor(fx.make_view(base_iter, fx.make_layout(1 << 30, 1))) + div = fx.logical_divide(flat, fx.make_layout(1, 1)) + + def _load(elem_idx): + fx.copy(copy_atom, fx.slice(div, (None, elem_idx)), reg) + return fx.Vector(fx.memref_load_vec(reg)) + + return _load + + _k_load_fp8x16 = _make_flat_loader( + key_cache_ptr, FP8, 16, fx.UniversalCopy128b() + ) + _v_load_fp8x16 = _make_flat_loader( + value_cache_ptr, FP8, 16, fx.UniversalCopy128b() + ) + # Per-lane Q chunk (QCHUNK 16-bit elems) fetched in QLOAD_UNIT-wide + # pieces (128b max per buffer load): head_dim=256 needs 2 pieces. + QLOAD_UNIT = QCHUNK if QCHUNK < 8 else 8 + N_QLOADS = QCHUNK // QLOAD_UNIT + _q_copy_op = ( + fx.rocdl.BufferCopy128b() if QLOAD_UNIT == 8 else fx.rocdl.BufferCopy64b() + ) + _q_load_chunk = _make_flat_loader(query_ptr, Q_DTYPE, QLOAD_UNIT, _q_copy_op) + _ctxlen_load = _make_flat_loader( + context_lengths_ptr, fx.Int32, 1, fx.rocdl.BufferCopy32b() + ) + + 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) + + context_len = fx.Int32(_ctxlen_load(seq)[0]) + # Bound block_tables to its real extent: the last (partial) 256-token tile + # can index a page past ceil(context/block_size); the bounded resource + # returns page 0 for that out-of-range read instead of faulting (those + # tail tokens are masked out anyway). + bt_num_records_bytes = ( + fx.Index(gpu.grid_dim.x) * fx.Index(max_blocks_per_seq) * 4 + ) # int32 entries + bt_rsrc = buffer_ops.create_buffer_resource( + block_tables_ptr, max_size=False, num_records_bytes=bt_num_records_bytes + ) + ks_rsrc = buffer_ops.create_buffer_resource(key_scale_ptr, max_size=True) + vs_rsrc = buffer_ops.create_buffer_resource(value_scale_ptr, max_size=True) + # Per-tensor: a single global scale, read once. Per-token: read + # per-token instead (see _kv_scale_ops/_stage_kv_scale_to_lds below). + if const_expr(not per_token_kv): + key_scale = fx.Int32( + buffer_ops.buffer_load( + ks_rsrc, arith.constant(0, type=T.i32), vec_width=1, is_scalar=True + ) + ).bitcast(fx.Float32) + value_scale = fx.Int32( + buffer_ops.buffer_load( + vs_rsrc, arith.constant(0, type=T.i32), vec_width=1, is_scalar=True + ) + ).bitcast(fx.Float32) + + num_tiles = cdiv(context_len, TILE_TOK) + tiles_per_part = cdiv(num_tiles, NP) + part_start = part * tiles_per_part + part_end_raw = part_start + tiles_per_part + part_end = arith.select(part_end_raw < num_tiles, part_end_raw, num_tiles) + + # One i8 blob carved into typed byte-offset pointers. `lds_base` is an + # ir.Value pointer (safe inside scf control flow); the Python `lds` + # handle is not. + 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): + # Pin the element-size alignment explicitly (a bare recast off the + # byte base would gcd the alignment down to 1). + 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)) + + c16 = 16 + rgroup = lane // c16 # 0..3: which quarter-wave (paired with warp -> query row) + lane16 = lane - rgroup * c16 # 0..15: this row's head-dim chunk index + + TOK_CHUNK = NWARP * MFMA_MNK # 64 + NCHUNK = TILE_TOK // TOK_CHUNK # 4 + + def _load_phys_scalar(page, vec_width=1): + result = buffer_ops.buffer_load( + bt_rsrc, + seq * max_blocks_per_seq + page, + vec_width=vec_width, + is_scalar=True, + ) + return fx.Int32(result) if vec_width == 1 else result + + def _v_page_fetch_and_stage(tt_i32): + # V's page depends on `rgroup` (shared across warps): warp w fetches + # its rgroup row and broadcasts via LDS (read back by _v_page_read_row). + base_page = tt_i32 * TILE_TOK // block_size # tile start is page-aligned + fetched = _load_phys_scalar( + base_page + warp * PAGES_PER_CHUNK, PAGES_PER_CHUNK + ) + if lane == 0: + fetched_vec = ( + fx.Vector.from_elements([fx.Int32(fetched)], dtype=fx.Int32) + if const_expr(PAGES_PER_CHUNK == 1) + else fx.Vector(fetched) + ) + _lds_store( + sVPage_off + warp * (PAGES_PER_CHUNK * 4), fx.Int32, 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 _kv_buf_off(tt_val): + # ping-pong buffer byte offset (0 when single-buffered). + if const_expr(per_token_kv): + 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]) + base_tok = lane16 * NCHUNK + scale_idx = phys * stride_ks_block + kv_h * stride_ks_head + base_tok + k_scale_vec = fx.Vector( + buffer_ops.buffer_load( + ks_rsrc, scale_idx, vec_width=NCHUNK, dtype=fx.Float32 + ) + ) + v_scale_vec = fx.Vector( + buffer_ops.buffer_load( + vs_rsrc, scale_idx, vec_width=NCHUNK, dtype=fx.Float32 + ) + ) + slot = (warp * TOK_PER_WARP + base_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: + # block_size==16: each lane stages its own rgroup's page/token, + # so the 4 sub-blocks stage in parallel across rgroup-groups. + phys = fx.Int32( + vector.extract( + arith.unwrap(phys_vec), + static_position=[], + dynamic_position=[fx.Index(rgroup)], + ) + ) + scale_idx = phys * stride_ks_block + kv_h * stride_ks_head + lane16 + k_scale_scalar = fx.Float32( + buffer_ops.buffer_load( + ks_rsrc, scale_idx, vec_width=1, dtype=fx.Float32 + ) + ) + v_scale_scalar = fx.Float32( + buffer_ops.buffer_load( + vs_rsrc, scale_idx, vec_width=1, dtype=fx.Float32 + ) + ) + fx.rocdl.sched_barrier(fx.rocdl.mask_vmem_rd) + slot = (warp * TOK_PER_WARP + rgroup * c16 + 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): + # This lane's 4 per-token scales for chunk `a` from an LDS scale region. + slot = (warp * TOK_CHUNK + a * c16 + 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 + ) + + # -- raw dwordx4 K load (A operand) -- + # token = warp*TOK_CHUNK + a*c16 + lane16 (the softmax mask and P-pack + # write position below must encode this same formula). + def _k_ops(phys, a): + within_page_tok = (a * c16 + lane16) % block_size + ops = [] + for qkhe in range_constexpr(QKHE_LOOP): + he_idx = qkhe * RGROUP_QUARTERS + rgroup + base = ( + ((phys * n_kv + kv_h) * QCHUNK + he_idx) * block_size + + within_page_tok + ) * QK_CHUNK_ELEMS + w = _k_load16(base) # head[he_idx*16 : +16] -> k_step 2*qkhe, 2*qkhe+1 + if const_expr(block_size == 16): + # help the scheduler overlap the PAGES_PER_CHUNK gathered loads + fx.rocdl.sched_barrier(fx.rocdl.mask_vmem_rd) + ops.extend([w[0], w[1]]) + return ops # N_SUBCHUNKS i64 operands + + 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 * PAGES_PER_CHUNK, 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) + ) + flat = [] + for a in range_constexpr(NCHUNK): + phys = fx.Int32(phys_vec[(a * c16) // 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), phys_vec + + # -- prologue: prefetch the first tile's K -- + num_tiles_m1 = num_tiles - 1 + start_safe = arith.select(part_start < num_tiles, part_start, num_tiles_m1) + k_pf0, phys_vec0 = _k_ops_flat(start_safe) + # V page-index prefetch, issued here too for the same overlap; the + # LDS write is visible after the barrier below. + _v_page_fetch_and_stage(start_safe) + if const_expr(per_token_kv): + _stage_kv_scale_to_lds(phys_vec0, _kv_buf_off(fx.Int32(start_safe))) + + # per_token_kv has no single global key_scale/value_scale: scale_qk + # drops the key_scale factor (folded in per-token, see masked_chunks + # below) and v_scale_f is unused (replaced by v_max_scaled). + 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) + fm_contract = arith.FastMathFlags.contract + # Softmax scores are finite or the -inf mask sentinel -- never NaN -- so + # nnan lets maxnum lower to a bare v_max (no v_cmp_u NaN check + its s_nop + # hazard) and fuse to v_max3. (ninf must NOT be set: -inf is load-bearing.) + 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), + ) + + # f32[16, NWARP] cross-warp scratch: scalar write at (row, warp), vec read of a row's NWARP valid slots. + 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): + # f32 -> fp8 must use the HW cvt (arith.truncf to fp8 doesn't lower); + # pack 4 f32 -> 1 i32 (4 fp8) via two cvt_pk_fp8_f32 calls. + 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) + + qh_local = warp * 4 + rgroup # 0..15: this thread's query row within an M-tile + + # Each M-tile quantizes 16 rows of the flattened (MTP position, GQA + # head) axis: `flat_idx = m*16 + qh_local`, decomposed as + # `qi = flat_idx // query_group_size`, `gs_head = flat_idx % + # query_group_size` (same convention as `_mtp_groups`). No + # cross-M-tile dependency, so no barriers needed between iterations. + def _quant_q_row(m, qi, gs_head, q_row_off): + qh0 = kv_h * query_group_size + gs_head + row_byte0 = ( + (seq * query_length + 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 + # QCHUNK elements, loaded as N_QLOADS contiguous QLOAD_UNIT pieces + # (a buffer load is 128b max); head_dim=256 splits into 2 pieces. + q_units = [ + _q_load_chunk(base_elem + u * QLOAD_UNIT) + for u in range_constexpr(N_QLOADS) + ] + + absmax = fmath.absf(q_units[0]).reduce(ReductionOp.MAX).to(fx.Float32) + for u in range_constexpr(1, N_QLOADS): + absmax = fx.maxnumf( + absmax, + fmath.absf(q_units[u]).reduce(ReductionOp.MAX).to(fx.Float32), + ) + 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] (not [m][qh]) so the whole M_TILES-wide + # row for a fixed qh is contiguous, letting the KV-loop read + # it back in one wide load instead of M_TILES separate + # narrow ones -- see the read site below. + _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 + # `flat_idx < TOTAL_ROWS` is statically true for every lane except + # possibly on the last M-tile (only it can be a partial tile), so + # skip the runtime branch (and its EXEC-mask overhead) entirely + # for every other M-tile. + if const_expr((m + 1) * MFMA_MNK <= TOTAL_ROWS): + _quant_q_row(m, qi, gs_head, q_row_off) + elif flat_idx < TOTAL_ROWS: + _quant_q_row(m, qi, gs_head, q_row_off) + else: + _st_words( + q_row_off + qh_local * head_dim + lane16 * QCHUNK, + fx.Vector.filled(QCHUNK // 4, 0, fx.Int32), + ) + if lane16 == 0: + _st1(sQscale_off, qh_local * M_TILES + m, ZERO_F) + + gpu.barrier() + + # First tile's V page-index row, now visible after the barrier above + # (the fetch+LDS-store was issued earlier, alongside k_pf0). + v_page_pf0 = _v_page_read_row() + + # Q is the B operand, read raw from sQ once per M-tile and held in + # registers. MUST use the exact same (qkhe, rgroup, qkr) -> head_dim + # permutation as K's `_k_ops`. + 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]]) + # q_ops_all[m*N_SUBCHUNKS+s] for s=0..N_SUBCHUNKS-1, s = qkhe*2+qkr, + # = M-tile m's head[he_idx*16+qkr*8 : +8] of qhead=lane16 + + # QK in NCHUNK chunks of 4 tokens: each chunk yields a f32x4 + # C-fragment, so softmax processes 4 scores at a time (low VGPR peak). + _ct = [ + fx.Vector.from_elements([float(a * c16 + r) for r in range_constexpr(4)]) + for a in range_constexpr(NCHUNK) + ] + # P.V is loop-tiled over head-dim (like production's VHELOOP): each + # step computes O[:, vh*VHE_SIZE:+VHE_SIZE] instead of materializing + # the full [16, head_dim] at once. + VHE_SIZE = head_dim // VHE_CHUNKS + OP_ELEMS = ( + MFMA_MNK * VHE_SIZE // (NWARP * WAVE) + ) # PV C-fragment elements/lane/chunk (probed = 4) + + # -- raw dwordx4 V load (B operand) -- + # One dwordx4 load per (16-token sub-block, head_elem); trans_v only + # changes the offset formula. `sub`/`step` walk pages/16-token sub-blocks. + NVOPS = TILE_TOK // MFMA_K # 8 PV k_steps (256 tokens / K=32) + STEPS_PER_PAGE = block_size // 16 + + 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_PAGE): + if const_expr(trans_v): + base = ( + ((phys_row[sub] * n_kv + kv_h) * STEPS_PER_PAGE + step) + * head_dim + + head_element + ) * 16 + else: + base = ( + (phys_row[sub] * n_kv + kv_h) * head_dim + head_element + ) * block_size + step * 16 + w = _v_load16(base) + if const_expr(block_size == 16): + # help the scheduler overlap the per-page gathered loads (see _k_ops) + 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 + + # query_length==1: plain wave-uniform `context_len` (folding in the + # per-lane `lane16` would cost a live VGPR for no behavioral change). + if const_expr(query_length == 1): + causal_bound = [context_len for _m in range_constexpr(M_TILES)] + else: + causal_bound = [ + context_len + - (query_length - 1) + + (m * MFMA_MNK + lane16) // query_group_size + for m in range_constexpr(M_TILES) + ] + + K_SLOT, V_SLOT = 0, 1 + # Per-M-tile loop-carried state after the K/V slots: one output chunk + # per VHE_CHUNKS (2 for head_dim=128, 4 for head_dim=256, 1 for 64), + # then running-max + running-denom. + STATE_PER_M = VHE_CHUNKS + 2 + + 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]) + for tt, ostate in range(part_start, part_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(tt) + tok0 = tt * TILE_TOK + # per_tensor phase-split: let IGLP interleave MFMA with softmax + # VALU/LDS to hide the MFMA-hazard s_nop (per_token is VGPR-cliffed). + if const_expr(not per_token_kv and M_TILES > 1): + 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 + + # This tile's ping-pong scale buffer; the tt+1 prefetch stages the other. + cur_kv_buf = _kv_buf_off(tt) + + # V data doesn't depend on `m`; hoist once for the phase-split + # instead of reloading per M-tile. + v_vh_shared = None + if const_expr(M_TILES > 1): + v_vh_shared = [ + _v_ops(v_page_cur, vh) for vh in range_constexpr(VHE_CHUNKS) + ] + + # q_scale doesn't depend on `m`; read the whole M_TILES-wide row once + # (contiguous via the transposed [qh][m] sQscale layout). + 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 if const_expr(M_TILES > 1) else 0 + ) + + # Phase-split (M_TILES>1): pass-1 (QK+mask+max) loops all M-tiles into + # per-M-tile LDS slices so they share ONE barrier. M_TILES==1: else below. + if const_expr(M_TILES > 1): + masked_chunks_saved = [None] * M_TILES + scale_saved = [None] * M_TILES + + # per_token: compute the m-independent pv_max once, hold only + # k_scale across Phase A, re-read v_scale for Phase B after the + # barrier (peak scale liveness 16, not 32). + 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, v_scale_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 = arith.constant_vector(0.0, T.f32x4) + for s in range_constexpr(N_SUBCHUNKS): + acc = fx.rocdl.mfma_f32_16x16x32_fp8_fp8( + T.f32x4, + [ + k_cur[a * N_SUBCHUNKS + s], + q_ops_all[m * N_SUBCHUNKS + s], + acc, + 0, + 0, + 0, + ], + ) + frag_Ss.append(fx.Vector(acc)) + + scale = scale_qk * fx.Float32(q_scale_vec[m]) + n_valid_tile = (causal_bound[m] - tok0).to(fx.Float32) + base_tok_f = fx.Int32(warp * TOK_CHUNK + rgroup * 4).to(fx.Float32) + thr = fx.Vector.from_elements( + [n_valid_tile - base_tok_f], dtype=fx.Float32 + ).broadcast_to(4) + neg4 = fx.Vector.filled(4, float("-inf"), fx.Float32) + + if const_expr(per_token_kv): + scaled_frags = [ + frag_Ss[a] * k_scale_shared[a] + for a in range_constexpr(NCHUNK) + ] + else: + scaled_frags = frag_Ss + + masked_chunks = [ + (_ct[a] < 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 * scale) + + masked_chunks_saved[m] = masked_chunks + scale_saved[m] = scale + + # tt+1 K/V/scale prefetch, issued once here so the V-page read + # reuses this barrier. + k_next = k_cur + if tt1 < part_end: + k_next, phys_vec1 = _k_ops_flat(tt1) + _v_page_fetch_and_stage(tt1) + if const_expr(per_token_kv): + # stage tt+1 into the OTHER buffer (current v_scale re-read below) + _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 tt1 < part_end: + v_page_next = _v_page_read_row() + next_state[V_SLOT] = v_page_next + + # Phase B v_scale: M_TILES>=4 re-reads per chunk (VGPR cliff); + # smaller M_TILES hold all NCHUNK. Both read the current buffer. + 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) + ] + + for m in range_constexpr(M_TILES): + 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] + scale = scale_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) + + m_new = fx.maxnumf( + m_prev, + _ld_lw_row(_lmax_off_m(m), lane16).reduce( + ReductionOp.MAX, fastmath=fm_nnan + ), + fastmath=fm_nnan, + ) + # Fully-invalid row: use 0 as the effective max so masked lanes + # give exp2(-inf-0)==0 (avoids the -inf-(-inf) cancellation). + safe_max = arith.select(m_new > NEG_INF, 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( + exp2_f32_fast(masked_chunks[a] * scale - m_new_b) + ) + 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 = ( + sP_off + lane16 * SP_ROW_BYTES + warp * TOK_CHUNK + rgroup * 4 + ) + # The NCHUNK P words scatter across the row at stride c16//4 + # i32 (the token->fp8-lane interleave the PV ds_read_b128 + # expects); one strided store per chunk. + for a in range_constexpr(NCHUNK): + _lds_store( + p_off0 + a * (c16 // 4) * f32, + fx.Int32, + fx.Vector.from_elements([words[a]], dtype=fx.Int32), + ) + for sh in (16, 32): + ls = ls.addf(ls.shuffle_xor(sh, WAVE), fastmath=fm_contract) + # PV output is [head-dim, query-row=lane16] after the operand + # swap, so correction/denominator are per-lane scalars (no sCorr). + safe_prev = arith.select(m_prev > NEG_INF, m_prev, ZERO_F) + corr_reg = fx.Float32(exp2_amdgcn_scalar(safe_prev - safe_max)) + if rgroup == 0: + _st_lw(sLsum_off, lane16, warp, ls) + gpu.barrier() + gsum = _ld_lw_row(sLsum_off, lane16).reduce(ReductionOp.ADD) + l_new = fx.Float32( + arith.mulf( + arith.unwrap(l_prev), + arith.unwrap(corr_reg), + fastmath=fm_contract, + ) + ).addf(gsum, fastmath=fm_contract) + + 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) + for vh in range_constexpr(VHE_CHUNKS): + v_vh = v_vh_shared[vh] + acc = arith.constant_vector(0.0, T.f32x4) + for s in range_constexpr(NVOPS): + # SWAPPED operands (V=A, P=B): output row = + # head-dim, output col = query-row=lane16. + acc = fx.rocdl.mfma_f32_16x16x32_fp8_fp8( + T.f32x4, [v_vh[s], p_ops[s], acc, 0, 0, 0] + ) + 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] = o_acc[vh] * corr_b + op + next_state.extend([*o_acc, m_new, l_new]) + # sP and sLsum are reused by the next M-tile. Synchronize + # all waves after their LDS reads before any wave overwrites + # those regions, then retire this tile's dependency chain. + if const_expr(m < M_TILES - 1): + gpu.barrier() + fx.rocdl.sched_barrier(0) + else: + # M_TILES==1 single tile (m==0 for the _o_slot/_m_slot/_l_slot helpers). + 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 + # QK: each NCHUNK chunk accumulates N_SUBCHUNKS k_steps into an f32x4. + frag_Ss = [] + for a in range_constexpr(NCHUNK): + acc = arith.constant_vector(0.0, T.f32x4) + for s in range_constexpr(N_SUBCHUNKS): + acc = fx.rocdl.mfma_f32_16x16x32_fp8_fp8( + T.f32x4, + [k_cur[a * N_SUBCHUNKS + s], q_ops_all[s], acc, 0, 0, 0], + ) + frag_Ss.append(fx.Vector(acc)) + # tt+1 K/V/scale prefetch, issued here to reuse the pass-1 barrier. + k_next = k_cur + if tt1 < part_end: + 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 + # Softmax: each lane owns one qhead (lane%16); register reduce + shuffle_xor. + scale = scale_qk * _ld1( + sQscale_off, lane16 + ) # per-qhead positive score scale + n_valid_tile = (causal_bound[0] - tok0).to(fx.Float32) + base_tok_f = fx.Int32(warp * TOK_CHUNK + rgroup * 4).to(fx.Float32) + thr = fx.Vector.from_elements( + [n_valid_tile - base_tok_f], dtype=fx.Float32 + ).broadcast_to(4) + neg4 = fx.Vector.filled(4, -1e30, fx.Float32) + # per_token_kv: K-scale varies per token, so fold it in BEFORE the max-reduce. + v_scale_vecs = None + if const_expr(per_token_kv): + v_scale_vecs = [] + scaled_frags = [] + 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_frags.append(frag_Ss[a] * k_scale_vec) + else: + scaled_frags = frag_Ss + # Reused in pass 2 below, halving the mask instruction count. + masked_chunks = [ + (_ct[a] < 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): + pm = fx.maxnumf(pm, masked_chunks[a].reduce(ReductionOp.MAX)) + for sh in (16, 32): + pm = fx.maxnumf(pm, pm.shuffle_xor(sh, WAVE)) + _st_lw( + sLmax_off, lane16, warp, pm * scale + ) # redundant across the 4 lanes sharing this qhead + # per_token_kv: max V-scale for the per-tile fp8 normalization + # (any positive value is correct, so skip the causal mask). + if const_expr(per_token_kv): + pv_max = fx.Float32(0.0) + for a in range_constexpr(NCHUNK): + pv_max = fx.maxnumf( + pv_max, v_scale_vecs[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-index row for next tile, now visible after the barrier. + v_page_next = v_page_cur + if tt1 < part_end: + v_page_next = _v_page_read_row() + next_state[V_SLOT] = v_page_next + # per_token_kv: combine warps' max V-scale into the tile's + # normalization factor (also the PV correction below). + 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 + m_new = fx.maxnumf( + m_prev, _ld_lw_row(sLmax_off, lane16).reduce(ReductionOp.MAX) + ) + m_new_b = fx.Vector.from_elements( + [m_new], dtype=fx.Float32 + ).broadcast_to(4) + ls = fx.Float32(0.0) + words = [] + zero4_p = fx.Vector.filled(4, 0.0, fx.Float32) + for a in range_constexpr(NCHUNK): + # re-mask Pa so a fully-masked chunk contributes exactly 0 + valid_a = masked_chunks[a] > fx.Vector.filled(4, -1e29, fx.Float32) + Pa = valid_a.select( + fx.Vector(exp2_f32_fast(masked_chunks[a] * scale - m_new_b)), + 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 + 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_CHUNK + rgroup * 4 + # NCHUNK P words scatter at stride c16//4 i32 (see phase-split). + for a in range_constexpr(NCHUNK): + _lds_store( + p_off0 + a * (c16 // 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.addf(ls.shuffle_xor(sh, WAVE), fastmath=fm_contract) + # PV (V=A, P=B) -> output [head-dim, query-row=lane16]; same as + # the phase-split path. + corr_reg = fx.Float32(exp2_amdgcn_scalar(m_prev - m_new)) + if rgroup == 0: + _st_lw(sLsum_off, lane16, warp, ls) + gpu.barrier() + gsum = _ld_lw_row(sLsum_off, lane16).reduce(ReductionOp.ADD) + l_new = fx.Float32( + arith.mulf( + arith.unwrap(l_prev), + arith.unwrap(corr_reg), + fastmath=fm_contract, + ) + ).addf(gsum, fastmath=fm_contract) + 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) + # Single tile: batch both vh's V loads upfront (no sibling chain + # to hide the latency behind). + 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 = arith.constant_vector(0.0, T.f32x4) + for s in range_constexpr(NVOPS): + acc = fx.rocdl.mfma_f32_16x16x32_fp8_fp8( + T.f32x4, [v_vh[s], p_ops[s], acc, 0, 0, 0] + ) + 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] = o_acc[vh] * corr_b + op + next_state.extend([*o_acc, m_new, l_new]) + results = yield next_state + o_final = results + + # Direct-store epilogue: after the PV swap each lane holds its 4 head-dim + # values for one query-row and writes them straight to global. + 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 + l_row = o_final[_l_slot(m)] + safe_l = arith.select(l_row > ZERO_F, l_row, fx.Float32(1.0)) + inv_l = fx.Float32(rcp_f32(safe_l)) + if const_expr(per_token_kv): + o_scale = inv_l + else: + o_scale = fx.Float32( + arith.mulf( + arith.unwrap(inv_l), + arith.unwrap(v_scale_f * inv_fp8), + fastmath=fm_contract, + ) + ) + o_scale_b = fx.Vector.from_elements( + [o_scale], dtype=fx.Float32 + ).broadcast_to(OP_ELEMS) + qi_e = row // query_group_size + gs_head_e = row - qi_e * query_group_size + qh = kv_h * query_group_size + gs_head_e + + def _emit(o_norm, sub): + if const_expr(NP == 1): + out_row = output_ptr[seq * query_length + qi_e, qh, None] + out_chunk = fx.slice( + fx.logical_divide(out_row, fx.make_layout(OP_ELEMS, 1)), + (None, sub), + ) + out_chunk.store(o_norm) + else: + base = ((seq * n_kv + kv_h) * NP + part) * TOTAL_ROWS + row + pout_div = fx.logical_divide(pout_ptr, fx.make_layout(OP_ELEMS, 1)) + pout_chunk = fx.slice( + pout_div, (None, base * (head_dim // OP_ELEMS) + sub) + ) + pout_chunk.store(o_norm) + + 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 + # Guard the partial last tile's out-of-range rows (folded away for full tiles). + if row < TOTAL_ROWS: + _emit(o_norm, sub) + + if const_expr(NP > 1): + if warp == 0 and rgroup == 0: + base = ((seq * n_kv + kv_h) * NP + part) * TOTAL_ROWS + row + if row < TOTAL_ROWS: + # Convert the running max from log2 units (scale_qk folds + # in LOG2E) to natural-log units: the shared reduce + # re-applies LOG2E itself when combining partitions. + pmax_ptr[base] = o_final[_m_slot(m)] * fx.Float32(1.0 / LOG2E) + psum_ptr[base] = l_row + + @flyc.jit + def pa_decode_tile_launch( + output: fx.Tensor, + pmax: fx.Tensor, + psum: fx.Tensor, + pout: fx.Tensor, + query: fx.Tensor, + key_cache: fx.Tensor, + value_cache: fx.Tensor, + block_tables: fx.Tensor, + context_lengths: fx.Tensor, + key_scale: fx.Tensor, # [1] per-tensor OR [num_blocks, num_kv_heads, block_size] per-token + value_scale: fx.Tensor, # same shape as key_scale + 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_q_row: fx.Int32, + stride_q_head: fx.Int32, + stream: fx.Stream = fx.Stream(None), + ): + pa_decode_tile_kernel( + output, + pmax, + psum, + pout, + query, + key_cache, + value_cache, + block_tables, + context_lengths, + key_scale, + value_scale, + max_blocks_per_seq, + stride_ks_block, + stride_ks_head, + stride_q_row, + stride_q_head, + ).launch( + grid=(num_seqs, num_kv_heads, NP), + block=(BLOCK_THREADS, 1, 1), + stream=stream, + ) + + return {"launch": pa_decode_tile_launch, "kernel": pa_decode_tile_kernel} + + +def _is_current_stream_capturing() -> bool: + if not torch.cuda.is_available(): + return False + try: + return torch.cuda.is_current_stream_capturing() + except RuntimeError: + return False + + +def get_recommended_splits( + num_sequences: int, + num_kv_heads: int, + split_kv_blocks: int = 1, +) -> int: + 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 = cdiv(num_sm, denom) * split_kv_blocks + return max(4, min(n, 8)) + + +def pa_decode_tile( + output: torch.Tensor, + query: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + block_tables: torch.Tensor, + context_lengths: torch.Tensor, + key_scale: float | torch.Tensor, + value_scale: float | torch.Tensor, + softmax_scale: float | None = None, + stream=None, + *, + num_partitions: int | None = None, + pmax: torch.Tensor | None = None, + psum: torch.Tensor | None = None, + pout: torch.Tensor | None = None, +) -> None: + """Host entry point. See module docstring for tensor layouts. + + ``num_partitions``/``pmax``/``psum``/``pout`` are optional caller overrides + (e.g. for CUDA-graph capture, where nothing may be allocated on-the-fly); + when omitted they are picked/allocated here. + """ + num_seqs = context_lengths.shape[0] + total_q_rows, num_q_heads, head_dim = query.shape + assert ( + total_q_rows % num_seqs == 0 + ), f"query.shape[0] ({total_q_rows}) must be a multiple of context_lengths.shape[0] ({num_seqs})" + query_length = total_q_rows // num_seqs + _, num_kv_heads, num_hgroups, block_size, hgroup_width = key_cache.shape + + assert num_hgroups == head_dim // 16 and hgroup_width == 16 + assert block_size in ( + 16, + 64, + ), f"pa_decode_tile only supports block_size in (16, 64), got {block_size}" + + trans_v = value_cache.dim() == 5 + if trans_v: + _, v_num_kv_heads, v_subblocks, v_head_dim, v_width = value_cache.shape + assert ( + v_head_dim == head_dim and v_width == 16 and v_subblocks == block_size // 16 + ), f"value_cache shape {tuple(value_cache.shape)} doesn't match block_size={block_size}, head_dim={head_dim}" + else: + _, v_num_kv_heads, v_head_dim, v_block_size = value_cache.shape + assert v_head_dim == head_dim and v_block_size == block_size, ( + f"value_cache shape {tuple(value_cache.shape)} doesn't match " + f"block_size={block_size}, head_dim={head_dim}" + ) + assert v_num_kv_heads == num_kv_heads + assert ( + block_tables.dtype == torch.int32 + ), f"block_tables must be int32, got {block_tables.dtype}" + assert ( + context_lengths.dtype == torch.int32 + ), 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 ValueError( + f"pa_decode_tile only supports f16/bf16 query, got {query.dtype}" + ) + assert ( + output.dtype == query.dtype + ), f"pa_decode_tile requires output.dtype == query.dtype, got {output.dtype} vs {query.dtype}" + + assert ( + query.stride(2) == 1 + ), f"pa_decode_tile requires a contiguous head_dim axis, got strides {query.stride()}" + + dev = query.device + key_scale_t = ( + key_scale + if isinstance(key_scale, torch.Tensor) + else torch.tensor([float(key_scale)], device=dev) + ) + value_scale_t = ( + value_scale + if isinstance(value_scale, torch.Tensor) + else torch.tensor([float(value_scale)], device=dev) + ) + per_token_kv = key_scale_t.dim() > 1 + if per_token_kv: + assert ( + value_scale_t.dim() > 1 + ), "value_scale must also be per-token (dim>1) when key_scale is per-token" + assert ( + key_scale_t.shape == value_scale_t.shape + ), f"key_scale/value_scale shape mismatch: {tuple(key_scale_t.shape)} vs {tuple(value_scale_t.shape)}" + assert key_scale_t.shape == (key_cache.shape[0], num_kv_heads, block_size), ( + "per-token key_scale/value_scale must be [num_blocks, num_kv_heads, block_size] " + f"matching the KV cache, got {tuple(key_scale_t.shape)}" + ) + stride_ks_block = int(key_scale_t.stride(0)) + stride_ks_head = int(key_scale_t.stride(1)) + else: + stride_ks_block = 0 + stride_ks_head = 0 + assert ( + key_scale_t.dtype == torch.float32 and key_scale_t.device == dev + ), f"key_scale tensor must be float32 on {dev}, got {key_scale_t.dtype} on {key_scale_t.device}" + assert ( + value_scale_t.dtype == torch.float32 and value_scale_t.device == dev + ), f"value_scale tensor must be float32 on {dev}, got {value_scale_t.dtype} on {value_scale_t.device}" + + if not num_partitions: + blocks_per_partition = KV_COMPUTE_BLOCK // block_size + num_partitions = get_recommended_splits( + num_seqs, + num_kv_heads, + split_kv_blocks=blocks_per_partition, + ) + + compiled = compile_pa_decode_tile( + head_dim=head_dim, + query_group_size=query_group_size, + block_size=int(block_size), + 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, + ) + + is_graph_capturing = _is_current_stream_capturing() + if num_partitions == 1: + # NP==1 writes output directly; partials unused (caller buffers ignored). + if pmax is None: + if is_graph_capturing: + raise ValueError( + "CUDA graph capture requires preallocated `pmax`/`psum`/`pout` " + "even when num_partitions==1 (nothing may be allocated mid-capture)." + ) + pmax = psum = pout = torch.empty(1, dtype=torch.float32, device=dev) + else: + total_rows = query_length * query_group_size + expected_scalar_shape = (num_seqs, num_kv_heads, num_partitions, total_rows) + if pmax is None or psum is None or pout is None: + if is_graph_capturing: + raise ValueError( + "CUDA graph capture requires preallocated `pmax`/`psum`/`pout` " + "for num_partitions>1 (nothing may be allocated mid-capture)." + ) + pmax = torch.empty(*expected_scalar_shape, dtype=torch.float32, device=dev) + psum = torch.empty(*expected_scalar_shape, dtype=torch.float32, device=dev) + pout = torch.empty( + *expected_scalar_shape, head_dim, dtype=output.dtype, device=dev + ) + else: + assert ( + pmax.shape == expected_scalar_shape + ), f"pmax shape {tuple(pmax.shape)} != {expected_scalar_shape}" + assert ( + psum.shape == expected_scalar_shape + ), f"psum shape {tuple(psum.shape)} != {expected_scalar_shape}" + assert pout.shape == ( + *expected_scalar_shape, + head_dim, + ), f"pout shape {tuple(pout.shape)} != {(*expected_scalar_shape, head_dim)}" + s = stream or torch.cuda.current_stream() + + _run_compiled( + compiled["launch"], + output, + pmax.view(-1), + psum.view(-1), + pout.view(-1), + query, + key_cache, + value_cache, + block_tables, + context_lengths, + key_scale_t, + value_scale_t, + int(max_blocks_per_seq), + int(num_seqs), + int(num_kv_heads), + stride_ks_block, + stride_ks_head, + int(query.stride(0)), + int(query.stride(1)), + s, + ) + if num_partitions > 1: + from csrc.cpp_itfs.pa.pa_ps import launch_pa_decode_ps_reduce + + output_5d = output.reshape( + num_seqs, query_length, num_kv_heads, query_group_size, head_dim + ) + with torch.cuda.stream(s): + launch_pa_decode_ps_reduce( + output_5d, + psum, + pmax, + pout, + None, + output_5d.stride(0), + output_5d.stride(1), + output_5d.stride(2), + output_5d.stride(3), + pmax.stride(0), + pmax.stride(1), + pmax.stride(2), + pout.stride(0), + pout.stride(1), + pout.stride(2), + pout.stride(3), + query_seq_len=query_length, + query_group_size=query_group_size, + head_size=head_dim, + context_partition_num=num_partitions, + ) diff --git a/op_tests/flydsl_tests/test_pa_decode_tile.py b/op_tests/flydsl_tests/test_pa_decode_tile.py new file mode 100644 index 00000000000..cb48aae54f1 --- /dev/null +++ b/op_tests/flydsl_tests/test_pa_decode_tile.py @@ -0,0 +1,160 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Correctness tests for the FlyDSL paged-attention Tile kernel.""" + +from __future__ import annotations + +import pytest +import torch + +from aiter import per_tensor_quant +from aiter.ops.flydsl import is_flydsl_available + +pytest.importorskip("flydsl") + +if not torch.cuda.is_available(): + pytest.skip("ROCm is not available", allow_module_level=True) +if not is_flydsl_available(): + pytest.skip("FlyDSL is not available", allow_module_level=True) + +from aiter.ops.flydsl.pa_decode_tile import pa_decode_tile # noqa: E402 + + +def _get_arch() -> str: + return torch.cuda.get_device_properties(0).gcnArchName.lower().split(":")[0] + + +pytestmark = pytest.mark.skipif( + not (_get_arch().startswith("gfx942") or _get_arch().startswith("gfx95")), + reason="PA Tile requires gfx942 or gfx95x", +) + + +def _quant_dtype() -> torch.dtype: + if _get_arch().startswith("gfx95"): + return torch.float8_e4m3fn + return torch.float8_e4m3fnuz + + +def _reference_attention( + query: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + block_tables: torch.Tensor, + context_lengths: torch.Tensor, + key_scale: torch.Tensor, + value_scale: torch.Tensor, +) -> torch.Tensor: + batch_size, num_query_heads, head_dim = query.shape + block_size = key_cache.shape[2] + num_kv_heads = key_cache.shape[1] + query_group_size = num_query_heads // num_kv_heads + softmax_scale = head_dim**-0.5 + output = torch.empty_like(query) + + for seq_idx in range(batch_size): + context_length = int(context_lengths[seq_idx].item()) + token_ids = torch.arange(context_length, device=query.device) + logical_pages = token_ids // block_size + token_offsets = token_ids % block_size + physical_pages = block_tables[seq_idx, logical_pages].long() + + keys = ( + key_cache[physical_pages, :, token_offsets, :].float() * key_scale.float() + ) + values = ( + value_cache[physical_pages, :, :, token_offsets].float() + * value_scale.float() + ) + keys = keys.repeat_interleave(query_group_size, dim=1) + values = values.repeat_interleave(query_group_size, dim=1) + scores = ( + torch.einsum("hd,khd->hk", query[seq_idx].float(), keys) * softmax_scale + ) + probs = torch.softmax(scores, dim=-1) + output[seq_idx] = torch.einsum("hk,khd->hd", probs, values).to(query.dtype) + return output + + +@pytest.mark.parametrize("block_size", [16, 64]) +def test_pa_decode_tile_fp8(block_size: int) -> None: + torch.manual_seed(0) + device = torch.device("cuda") + batch_size = 3 + num_query_heads = 8 + num_kv_heads = 1 + head_dim = 128 + context_length = 257 + blocks_per_sequence = (context_length + block_size - 1) // block_size + num_blocks = batch_size * blocks_per_sequence + + query = torch.empty( + batch_size, + num_query_heads, + head_dim, + dtype=torch.bfloat16, + device=device, + ).uniform_(-0.5, 0.5) + key = torch.empty( + num_blocks, + num_kv_heads, + block_size, + head_dim, + dtype=torch.bfloat16, + device=device, + ).uniform_(-0.5, 0.5) + value = torch.empty( + num_blocks, + num_kv_heads, + head_dim, + block_size, + dtype=torch.bfloat16, + device=device, + ).uniform_(-0.5, 0.5) + + quant_dtype = _quant_dtype() + key_quant, key_scale = per_tensor_quant(key, quant_dtype=quant_dtype) + value_quant, value_scale = per_tensor_quant(value, quant_dtype=quant_dtype) + key_cache = ( + key_quant.view( + num_blocks, + num_kv_heads, + block_size, + head_dim // 16, + 16, + ) + .permute(0, 1, 3, 2, 4) + .contiguous() + ) + value_cache = value_quant.contiguous() + block_tables = torch.arange(num_blocks, dtype=torch.int32, device=device).reshape( + batch_size, blocks_per_sequence + ) + context_lengths = torch.full( + (batch_size,), context_length, dtype=torch.int32, device=device + ) + output = torch.empty_like(query) + + pa_decode_tile( + output, + query, + key_cache, + value_cache, + block_tables, + context_lengths, + key_scale, + value_scale, + ) + torch.cuda.synchronize() + + reference = _reference_attention( + query, + key_quant, + value_quant, + block_tables, + context_lengths, + key_scale, + value_scale, + ) + torch.testing.assert_close(output, reference, atol=3e-2, rtol=3e-2) From 9183f94193e8acd7c34e4523a0f843febe3c43be Mon Sep 17 00:00:00 2001 From: fsx950223 Date: Wed, 22 Jul 2026 09:31:19 +0000 Subject: [PATCH 02/66] refactor(flydsl): Remove unused memory helpers Signed-off-by: fsx950223 Co-authored-by: Cursor --- aiter/ops/flydsl/kernels/mem_ops.py | 202 ---------------------------- aiter/ops/flydsl/kernels/utils.py | 7 - 2 files changed, 209 deletions(-) delete mode 100644 aiter/ops/flydsl/kernels/mem_ops.py diff --git a/aiter/ops/flydsl/kernels/mem_ops.py b/aiter/ops/flydsl/kernels/mem_ops.py deleted file mode 100644 index bc0d08e2ca1..00000000000 --- a/aiter/ops/flydsl/kernels/mem_ops.py +++ /dev/null @@ -1,202 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# Copyright (C) 2025-2026 FlyDSL Project Contributors - -"""Single home for kernel-facing memory & atomic helpers. - -Consolidates the LLVM-pointer construction, global load/store, and atomic -primitives that were previously scattered across ``kernels_common.py`` and -``utils.py`` (and hand-rolled inline in individual kernels). The original -modules re-export these names, so existing imports keep working unchanged. - -Two distinct atomic mechanisms live here and are NOT interchangeable: - * ``atomic_add`` - LLVM ``atomicrmw`` on an ``!llvm.ptr`` (pointer atomics) - * ``buffer_atomic_add`` - AMD ``raw.ptr.buffer.atomic.fadd`` on a buffer resource -""" - -import flydsl.expr as fx -from flydsl._mlir import ir -from flydsl._mlir.dialects import arith as _std_arith -from flydsl._mlir.dialects import fly as _fly -from flydsl._mlir.dialects import llvm as _llvm -from flydsl.expr import arith as _expr_arith -from flydsl.expr import buffer_ops, const_expr, rocdl -from flydsl.expr.typing import T - -# Public API. Prefer the unified names; the legacy aliases at the bottom of this -# list exist only so pre-consolidation imports keep working (re-exported by -# kernels_common / utils) and should not be used in new code: -# element_ptr <- get_llvm_ptr aligned_ptr <- extract_global_ptr -# to_llvm_ptr <- _create_llvm_ptr -__all__ = [ - # pointer construction / extraction (preferred) - "element_ptr", - "to_llvm_ptr", - "aligned_ptr", - "global_ptr_from_addr", - # global load / store - "global_load", - "global_load_i32", - "global_load_i64x2", - "global_store", - # atomics (two distinct mechanisms) - "atomic_add", - "buffer_atomic_add", - # legacy aliases (back-compat only) - "get_llvm_ptr", - "_create_llvm_ptr", - "extract_global_ptr", -] - - -# --------------------------------------------------------------------------- # -# Pointer construction / extraction -# --------------------------------------------------------------------------- # -def get_llvm_ptr(ptr, offset, dtype_bytes, ptr_type=None): - """Build a global (address-space 1) ``!llvm.ptr`` at ``ptr + offset*dtype_bytes``. - - Shared home for the LLVM-ptr arithmetic used by atomic/global accesses - (previously duplicated in hgemm_splitk.py, small_m_hgemm.py, splitk_hgemm.py - and rmsnorm_kernel.py). - """ - if ptr_type is None: - ptr_type = ir.Type.parse("!llvm.ptr<1>") - base_ptr = _fly.extract_aligned_pointer_as_index(ptr_type, ptr) - base_ptr = _llvm.PtrToIntOp(T.i64, base_ptr).result - byte_offset = _expr_arith.index_cast( - T.i64, fx.Index(offset) * fx.Index(dtype_bytes) - ) - llvm_ptr = _llvm.AddOp(base_ptr, byte_offset, _llvm.IntegerOverflowFlags(0)).result - llvm_ptr = _llvm.IntToPtrOp(ptr_type, llvm_ptr).result - return llvm_ptr._value if const_expr(hasattr(llvm_ptr, "_value")) else llvm_ptr - - -# Unified name for the element-pointer helper above. -element_ptr = get_llvm_ptr - - -def _create_llvm_ptr(value, address_space: int = 1): - value = buffer_ops._unwrap_value(value) - if isinstance(value.type, ir.IndexType): - i64_type = T.i64 - value = buffer_ops._unwrap_value(_std_arith.IndexCastOp(i64_type, value).result) - ptr_type = ir.Type.parse(f"!llvm.ptr<{address_space}>") - return _llvm.IntToPtrOp(ptr_type, value).result - - -# Unified name: int/index scalar -> typed llvm.ptr (any address space). -to_llvm_ptr = _create_llvm_ptr - - -def extract_global_ptr(tensor): - raw = ( - tensor.ir_value() - if hasattr(tensor, "ir_value") and not isinstance(tensor, ir.Value) - else tensor - ) - ptr_type = ir.Type.parse("!llvm.ptr<1>") - return _fly.extract_aligned_pointer_as_index(ptr_type, raw) - - -# Unified name: tensor/memref -> aligned global base pointer. -aligned_ptr = extract_global_ptr - - -def global_ptr_from_addr(addr_i64): - raw = ( - addr_i64.ir_value() - if hasattr(addr_i64, "ir_value") and not isinstance(addr_i64, ir.Value) - else addr_i64 - ) - return _llvm.IntToPtrOp(ir.Type.parse("!llvm.ptr<1>"), raw).result - - -# --------------------------------------------------------------------------- # -# Global load / store -# --------------------------------------------------------------------------- # -def global_load(global_ptr, byte_offset, result_type, *, alignment): - ptr = buffer_ops.get_element_ptr( - global_ptr, byte_offset=fx.Int64(byte_offset), elem_type=T.i8 - ) - return _llvm.LoadOp(result_type, ptr, alignment=alignment).result - - -def global_load_i64x2(global_ptr, byte_offset_i64): - return global_load(global_ptr, byte_offset_i64, T.i64x2, alignment=16) - - -def global_load_i32(global_ptr, elem_offset_i32): - return global_load( - global_ptr, fx.Int64(elem_offset_i32) * fx.Int64(4), T.i32, alignment=4 - ) - - -def global_store(global_ptr, byte_offset, value, *, alignment): - """Store ``value`` at ``global_ptr + byte_offset`` (mirror of ``global_load``). - - Unwraps a FlyDSL expression operand the same way ``atomic_add`` does, so - callers may pass either a raw ``ir.Value`` or a FlyDSL expression. New in the - consolidation; the hand-rolled ``llvm.StoreOp`` sites are collected in a later - phase, so this has no call site yet. - """ - val = value.ir_value() if const_expr(hasattr(value, "ir_value")) else value - ptr = buffer_ops.get_element_ptr( - global_ptr, byte_offset=fx.Int64(byte_offset), elem_type=T.i8 - ) - _llvm.StoreOp(val, ptr, alignment=alignment) - - -# --------------------------------------------------------------------------- # -# Atomics -- two distinct mechanisms (see module docstring) -# --------------------------------------------------------------------------- # -def atomic_add( - dst, - offset, - value, - *, - dtype_bytes=4, - syncscope="agent", - ordering=None, - alignment=None, - ptr_type=None, -): - """Atomically add ``value`` into ``dst[offset]`` in global memory. - - inline (e.g. the rmsnorm backward ``dweight`` accumulation). Selects - ``fadd`` for a floating-point operand and integer ``add`` otherwise, from - the operand's IR type, so a single call covers both cases. Returns the - atomicrmw result (the value previously stored at ``dst[offset]``). - - ``dtype_bytes`` sizes the byte offset and, unless ``alignment`` is given, is - reused as the access alignment. - """ - ptr = get_llvm_ptr(dst, offset, dtype_bytes, ptr_type=ptr_type) - val = value.ir_value() if const_expr(hasattr(value, "ir_value")) else value - elem_ty = val.type.element_type if isinstance(val.type, ir.VectorType) else val.type - bin_op = ( - _llvm.AtomicBinOp.fadd - if isinstance(elem_ty, ir.FloatType) - else _llvm.AtomicBinOp.add - ) - if ordering is None: - ordering = _llvm.AtomicOrdering.monotonic - if alignment is None: - alignment = dtype_bytes - return _llvm.AtomicRMWOp( - bin_op, - ptr, - val, - ordering, - syncscope=syncscope, - alignment=alignment, - ).result - - -def buffer_atomic_add(vdata, rsrc, offset, soffset, aux): - """Buffer-resource atomic fadd (AMD ``raw.ptr.buffer.atomic.fadd``). - - Single entry point for the buffer-resource atomic accumulation used by the - MoE gemm/blockscale wrappers (``atomic_add_f16x2 / _f32 / _x2``) and conv3d. - Distinct from ``atomic_add``: operates on a buffer resource + byte offset, - not an ``!llvm.ptr``, so the two are not interchangeable. - """ - return rocdl.raw_ptr_buffer_atomic_fadd(vdata, rsrc, offset, soffset, aux) diff --git a/aiter/ops/flydsl/kernels/utils.py b/aiter/ops/flydsl/kernels/utils.py index e8633e08ccb..5ae043e1e24 100644 --- a/aiter/ops/flydsl/kernels/utils.py +++ b/aiter/ops/flydsl/kernels/utils.py @@ -7,13 +7,6 @@ from flydsl.expr import arith, const_expr, rocdl from flydsl.expr.typing import T -# Pointer/global-load helpers now live in mem_ops; re-exported here for back-compat. -from .mem_ops import extract_global_ptr as extract_global_ptr -from .mem_ops import global_load as global_load -from .mem_ops import global_load_i32 as global_load_i32 -from .mem_ops import global_load_i64x2 as global_load_i64x2 -from .mem_ops import global_ptr_from_addr as global_ptr_from_addr - def rcp_f32(value): return rocdl.rcp(T.f32, value) From 5824387f8b7e19f6346ac43fa96c46a3a7ed089f Mon Sep 17 00:00:00 2001 From: fsx950223 Date: Thu, 23 Jul 2026 05:41:45 +0000 Subject: [PATCH 03/66] refactor(flydsl): Separate PA decode kernel from host Signed-off-by: fsx950223 Co-authored-by: Cursor --- aiter/ops/flydsl/kernels/pa_decode_tile.py | 1234 ++++++++++++++++++ aiter/ops/flydsl/pa_decode_tile.py | 1208 +---------------- op_tests/flydsl_tests/test_pa_decode_tile.py | 290 +++- 3 files changed, 1469 insertions(+), 1263 deletions(-) create mode 100644 aiter/ops/flydsl/kernels/pa_decode_tile.py diff --git a/aiter/ops/flydsl/kernels/pa_decode_tile.py b/aiter/ops/flydsl/kernels/pa_decode_tile.py new file mode 100644 index 00000000000..de8142e99e9 --- /dev/null +++ b/aiter/ops/flydsl/kernels/pa_decode_tile.py @@ -0,0 +1,1234 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (C) 2025-2026 FlyDSL Project Contributors + +"""Readable tile-programming reference for paged-attention fp8 decode. + +K/V are fp8 e4m3 (FNUZ on gfx942, OCP on gfx950) fed straight into +``mfma_f32_16x16x32_fp8_fp8``; Q (bf16/f16) and the softmax probabilities P are +quantized to fp8 too. Scales fold out of the matmuls (q/key scale into the QK +score, value scale + 1/FP8_MAX into the epilogue); softmax max/sum stay f32. +``key_scale``/``value_scale`` are either a ``[1]`` per-tensor scalar or a +``[num_blocks, num_kv_heads, block_size]`` per-token tensor (chosen by rank). + +``block_size`` (16/64) and ``head_dim`` (multiple of 64) are compile-time +constants. Layouts are logical, not production's preshuffle. + +* ``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 + +One CTA (4 waves) per (seq, kv_head) runs a flash-style online softmax over +256-token blocks; the 4 waves split tokens for Q.KT and head-dim for P.V, with +an LDS round-trip on P transposing ownership between the two MMAs. +""" + +import functools + +import flydsl.compiler as flyc +import flydsl.expr as fx +from flydsl.compiler.protocol import dsl_size_of +from flydsl.expr import arith, buffer_ops, const_expr, gpu, range_constexpr, vector +from flydsl.expr import math as fmath +from flydsl.expr.typing import T +from flydsl.expr.vector import ReductionOp +from flydsl.runtime.device import get_rocm_arch + +from . import dpp_utils +from .utils import cdiv, exp2_amdgcn_scalar, exp2_f32_fast, rcp_f32 + +MFMA_MNK = ( + 16 # M = N = 16 for the MMA atom; also query rows handled per CTA (padded to 16) +) +MFMA_K = 32 # fp8 MFMA contracts K = 32 per instruction (mfma_f32_16x16x32_fp8_fp8) +WAVE = 64 +LOG2E = 1.4426950408889634 +KV_COMPUTE_BLOCK = 256 + + +@functools.lru_cache(maxsize=None) +def compile_pa_decode_tile( + *, + head_dim: int, + query_group_size: int, + block_size: 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, +): + """Build the tile-programming PA-decode kernel + launch wrapper. + + ``block_size``, ``head_dim``, and ``query_dtype`` are compile-time + constants (lru_cache keys). ``query_length`` (MTP) and ``query_group_size`` + flatten into ``TOTAL_ROWS = query_length * query_group_size``, tiled into + ``M_TILES = ceil(TOTAL_ROWS / 16)`` independent 16-row MFMA tiles; each + extra M-tile duplicates loop-carried state, so VGPR/LDS/occupancy scale + roughly linearly with ``M_TILES``. + """ + is_gfx950 = "gfx95" in get_rocm_arch() + 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, + ), f"pa_decode_tile only supports block_size in (16, 64), 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 query_dtype == "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}" + assert query_length >= 1, f"query_length must be >= 1, got {query_length}" + # Flattened query-row axis (MTP outer, GQA head inner), tiled into 16-row M-tiles. + TOTAL_ROWS = query_length * query_group_size + M_TILES = cdiv(TOTAL_ROWS, MFMA_MNK) + ROWS_PADDED = M_TILES * MFMA_MNK + # PV layout: V=A, P=B -> output [head-dim (row), query-row (col=lane16)], + # generalized over head_dim via the VHE_CHUNKS loop. + NWARP = 4 # 4 waves / CTA + TOK_PER_WARP = 64 # tokens each warp owns per compute block (matches production KV_COMPUTE_BLOCK) + TILE_TOK = NWARP * TOK_PER_WARP # 256 tokens / compute block + PAGES_PER_CHUNK = ( + TOK_PER_WARP // block_size + ) # pages spanned by one 64-token warp-chunk: 1 (bs=64) or 4 (bs=16) + assert ( + head_dim % (NWARP * MFMA_MNK) == 0 + ), "head_dim must split across the 4 warps for PV" + + # head_dim splits into 16-element chunks (QK_CHUNK_ELEMS, one dwordx4 load); + # RGROUP_QUARTERS of them make a 64-element fetch group, QKHE_LOOP groups total. + 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}" + N_SUBCHUNKS = QKHE_LOOP * (QK_CHUNK_ELEMS // 8) + + # Q-quant chunk width: NQCHUNK stays fixed at 16 (tied to `lane16`'s role + # as the absmax butterfly width); QCHUNK scales with head_dim instead. + NQCHUNK = 16 + QCHUNK = ( + head_dim // NQCHUNK + ) # f16 elements per lane's load chunk (8 for head_dim=128, 4 for head_dim=64) + + VHE_CHUNKS = head_dim // (NWARP * MFMA_MNK) # 2 for head_dim=128, 1 for head_dim=64 + + if softmax_scale is None: + softmax_scale = 1.0 / (head_dim**0.5) + NP = int(num_partitions) # context partitions (grid.z); compile-time constant + + BLOCK_THREADS = NWARP * WAVE # 256 + + # -- LDS layout (shared across the 4 warps) -- + # sQ: fp8[ROWS_PADDED,head_dim] staged+quantized query. sP: fp8[16,TILE_TOK] + # quantized probs. sQscale: f32[ROWS_PADDED]. sLmax/sLsum: cross-warp scratch. + # sVPage: V page-index broadcast. sKScale/sVScale/sVScaleMax: per-token K/V + # scale staging. No sO/sM/sL/sCorr: PV output is register-resident/loop-carried + # and stored straight to global (V=A/P=B swap). + f32 = 4 + sQ_bytes = ROWS_PADDED * head_dim * 1 # fp8 + sP_off = sQ_bytes + # +16B row padding breaks a 32-bank LDS conflict on the P-pack writes while + # keeping the row 16B-aligned for PV's ds_read_b128. + SP_ROW_BYTES = TILE_TOK + 16 + sP_bytes = MFMA_MNK * SP_ROW_BYTES # fp8, padded rows + sQscale_off = sP_off + sP_bytes + NWARP_PAD = ( + NWARP + 1 + ) # coprime with 32 banks -> no LDS bank conflict on the cross-warp scratch + # Phase-split slices sLmax per M-tile so all pass-1 writes share one barrier. + sLmax_off = sQscale_off + ROWS_PADDED * f32 + sLsum_off = sLmax_off + M_TILES * MFMA_MNK * NWARP_PAD * f32 + # V page-index broadcast (V's page depends on `rgroup`, shared across warps). + sVPage_off = sLsum_off + MFMA_MNK * NWARP_PAD * f32 + sVPage_bytes = NWARP * PAGES_PER_CHUNK * 4 # i32 + # Per-token KV-scale staging (per_token only), double-buffered so the tt+1 + # prefetch (into the other buffer) doesn't clobber the current tile's scales. + KV_BUFS = 2 if per_token_kv else 1 + KV_BUF_STRIDE = 2 * NWARP * TOK_PER_WARP * f32 # k-region + v-region, one buffer + sKScale_off = sVPage_off + sVPage_bytes + sVScale_off = sKScale_off + NWARP * TOK_PER_WARP * f32 + sKVScale_bytes = KV_BUFS * 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 + + # LDS blob: one i32 array carved into typed byte-offset views (see the + # `_lds_*` helpers in the kernel). Every region size above is 4-byte + # aligned, so `total_bytes // 4` covers the blob exactly. + @fx.struct + class SharedStorage: + buf: fx.Array[fx.Int32, total_bytes // 4, 16] + + @flyc.kernel(known_block_size=(BLOCK_THREADS, 1, 1)) + def pa_decode_tile_kernel( + output_ptr: fx.Tensor, # [num_seqs*query_length, num_q_heads, head_dim] (written directly when NP==1) + # per-partition partial outputs (combined by the reduce kernel when NP>1): + pmax_ptr: fx.Tensor, # [num_seqs*query_length, num_kv_heads, num_partitions, query_group_size] row max + psum_ptr: fx.Tensor, # [num_seqs*query_length, num_kv_heads, num_partitions, query_group_size] row sum + pout_ptr: fx.Tensor, # [num_seqs*query_length, num_kv_heads, num_partitions, query_group_size, head_dim] Q_DTYPE, normalized O_p/l_p + query_ptr: fx.Tensor, # [num_seqs*query_length, num_q_heads, head_dim] -- row = seq*query_length + qi (MTP position) + key_cache_ptr: fx.Tensor, # [num_blocks, num_kv_heads, head_dim//16, block_size, 16] (blocked, see module docstring) + value_cache_ptr: fx.Tensor, # [num_blocks, num_kv_heads, block_size//16, head_dim, 16] (blocked, see module docstring) + block_tables_ptr: fx.Tensor, # [num_seqs, max_blocks_per_seq] + context_lengths_ptr: fx.Tensor, # [num_seqs] + key_scale_ptr: fx.Tensor, # [1] per-tensor OR [num_blocks, num_kv_heads, block_size] per-token + value_scale_ptr: fx.Tensor, # same shape as key_scale_ptr + max_blocks_per_seq: fx.Int32, + stride_ks_block: fx.Int32, + stride_ks_head: fx.Int32, + stride_q_row: fx.Int32, + stride_q_head: 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_h = fx.Int32(gpu.block_id("y")) + part = fx.Int32(gpu.block_id("z")) # context partition handled by this CTA + n_kv = fx.Int32(gpu.grid_dim.y) # num_kv_heads == gridDim.y + + # fx.copy-based flat loaders: K/V/context_len over a raw pointer, Q over + # a buffer resource. + def _make_flat_loader(tensor_ptr, elem_ty, reg_width, copy_op): + use_buffer_resource = isinstance( + copy_op, fx.rocdl.CopyOpCDNA3BufferCopyType + ) + copy_atom = fx.make_copy_atom(copy_op, elem_ty) + reg = fx.make_rmem_tensor(fx.make_layout(reg_width, 1), elem_ty) + base_iter = ( + fx.get_iter(fx.rocdl.make_buffer_tensor(tensor_ptr, max_size=True)) + if use_buffer_resource + else fx.get_iter(tensor_ptr) + ) + flat = fx.Tensor(fx.make_view(base_iter, fx.make_layout(1 << 30, 1))) + div = fx.logical_divide(flat, fx.make_layout(1, 1)) + + def _load(elem_idx): + fx.copy(copy_atom, fx.slice(div, (None, elem_idx)), reg) + return fx.Vector(fx.memref_load_vec(reg)) + + return _load + + _k_load_fp8x16 = _make_flat_loader( + key_cache_ptr, FP8, 16, fx.UniversalCopy128b() + ) + _v_load_fp8x16 = _make_flat_loader( + value_cache_ptr, FP8, 16, fx.UniversalCopy128b() + ) + # Per-lane Q chunk (QCHUNK 16-bit elems) fetched in QLOAD_UNIT-wide + # pieces (128b max per buffer load): head_dim=256 needs 2 pieces. + QLOAD_UNIT = QCHUNK if QCHUNK < 8 else 8 + N_QLOADS = QCHUNK // QLOAD_UNIT + _q_copy_op = ( + fx.rocdl.BufferCopy128b() if QLOAD_UNIT == 8 else fx.rocdl.BufferCopy64b() + ) + _q_load_chunk = _make_flat_loader(query_ptr, Q_DTYPE, QLOAD_UNIT, _q_copy_op) + _ctxlen_load = _make_flat_loader( + context_lengths_ptr, fx.Int32, 1, fx.rocdl.BufferCopy32b() + ) + + 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) + + context_len = fx.Int32(_ctxlen_load(seq)[0]) + # Bound block_tables to its real extent: the last (partial) 256-token tile + # can index a page past ceil(context/block_size); the bounded resource + # returns page 0 for that out-of-range read instead of faulting (those + # tail tokens are masked out anyway). + bt_num_records_bytes = ( + fx.Index(gpu.grid_dim.x) * fx.Index(max_blocks_per_seq) * 4 + ) # int32 entries + bt_rsrc = buffer_ops.create_buffer_resource( + block_tables_ptr, max_size=False, num_records_bytes=bt_num_records_bytes + ) + ks_rsrc = buffer_ops.create_buffer_resource(key_scale_ptr, max_size=True) + vs_rsrc = buffer_ops.create_buffer_resource(value_scale_ptr, max_size=True) + # Per-tensor: a single global scale, read once. Per-token: read + # per-token instead (see _kv_scale_ops/_stage_kv_scale_to_lds below). + if const_expr(not per_token_kv): + key_scale = fx.Int32( + buffer_ops.buffer_load( + ks_rsrc, arith.constant(0, type=T.i32), vec_width=1, is_scalar=True + ) + ).bitcast(fx.Float32) + value_scale = fx.Int32( + buffer_ops.buffer_load( + vs_rsrc, arith.constant(0, type=T.i32), vec_width=1, is_scalar=True + ) + ).bitcast(fx.Float32) + + num_tiles = cdiv(context_len, TILE_TOK) + tiles_per_part = cdiv(num_tiles, NP) + part_start = part * tiles_per_part + part_end_raw = part_start + tiles_per_part + part_end = arith.select(part_end_raw < num_tiles, part_end_raw, num_tiles) + + # One i8 blob carved into typed byte-offset pointers. `lds_base` is an + # ir.Value pointer (safe inside scf control flow); the Python `lds` + # handle is not. + 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): + # Pin the element-size alignment explicitly (a bare recast off the + # byte base would gcd the alignment down to 1). + 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)) + + c16 = 16 + rgroup = lane // c16 # 0..3: which quarter-wave (paired with warp -> query row) + lane16 = lane - rgroup * c16 # 0..15: this row's head-dim chunk index + + TOK_CHUNK = NWARP * MFMA_MNK # 64 + NCHUNK = TILE_TOK // TOK_CHUNK # 4 + + def _load_phys_scalar(page, vec_width=1): + result = buffer_ops.buffer_load( + bt_rsrc, + seq * max_blocks_per_seq + page, + vec_width=vec_width, + is_scalar=True, + ) + return fx.Int32(result) if vec_width == 1 else result + + def _v_page_fetch_and_stage(tt_i32): + # V's page depends on `rgroup` (shared across warps): warp w fetches + # its rgroup row and broadcasts via LDS (read back by _v_page_read_row). + base_page = tt_i32 * TILE_TOK // block_size # tile start is page-aligned + fetched = _load_phys_scalar( + base_page + warp * PAGES_PER_CHUNK, PAGES_PER_CHUNK + ) + if lane == 0: + fetched_vec = ( + fx.Vector.from_elements([fx.Int32(fetched)], dtype=fx.Int32) + if const_expr(PAGES_PER_CHUNK == 1) + else fx.Vector(fetched) + ) + _lds_store( + sVPage_off + warp * (PAGES_PER_CHUNK * 4), fx.Int32, 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 _kv_buf_off(tt_val): + # ping-pong buffer byte offset (0 when single-buffered). + if const_expr(per_token_kv): + 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]) + base_tok = lane16 * NCHUNK + scale_idx = phys * stride_ks_block + kv_h * stride_ks_head + base_tok + k_scale_vec = fx.Vector( + buffer_ops.buffer_load( + ks_rsrc, scale_idx, vec_width=NCHUNK, dtype=fx.Float32 + ) + ) + v_scale_vec = fx.Vector( + buffer_ops.buffer_load( + vs_rsrc, scale_idx, vec_width=NCHUNK, dtype=fx.Float32 + ) + ) + slot = (warp * TOK_PER_WARP + base_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: + # block_size==16: each lane stages its own rgroup's page/token, + # so the 4 sub-blocks stage in parallel across rgroup-groups. + phys = fx.Int32( + vector.extract( + arith.unwrap(phys_vec), + static_position=[], + dynamic_position=[fx.Index(rgroup)], + ) + ) + scale_idx = phys * stride_ks_block + kv_h * stride_ks_head + lane16 + k_scale_scalar = fx.Float32( + buffer_ops.buffer_load( + ks_rsrc, scale_idx, vec_width=1, dtype=fx.Float32 + ) + ) + v_scale_scalar = fx.Float32( + buffer_ops.buffer_load( + vs_rsrc, scale_idx, vec_width=1, dtype=fx.Float32 + ) + ) + fx.rocdl.sched_barrier(fx.rocdl.mask_vmem_rd) + slot = (warp * TOK_PER_WARP + rgroup * c16 + 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): + # This lane's 4 per-token scales for chunk `a` from an LDS scale region. + slot = (warp * TOK_CHUNK + a * c16 + 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 + ) + + # -- raw dwordx4 K load (A operand) -- + # token = warp*TOK_CHUNK + a*c16 + lane16 (the softmax mask and P-pack + # write position below must encode this same formula). + def _k_ops(phys, a): + within_page_tok = (a * c16 + lane16) % block_size + ops = [] + for qkhe in range_constexpr(QKHE_LOOP): + he_idx = qkhe * RGROUP_QUARTERS + rgroup + base = ( + ((phys * n_kv + kv_h) * QCHUNK + he_idx) * block_size + + within_page_tok + ) * QK_CHUNK_ELEMS + w = _k_load16(base) # head[he_idx*16 : +16] -> k_step 2*qkhe, 2*qkhe+1 + if const_expr(block_size == 16): + # help the scheduler overlap the PAGES_PER_CHUNK gathered loads + fx.rocdl.sched_barrier(fx.rocdl.mask_vmem_rd) + ops.extend([w[0], w[1]]) + return ops # N_SUBCHUNKS i64 operands + + 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 * PAGES_PER_CHUNK, 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) + ) + flat = [] + for a in range_constexpr(NCHUNK): + phys = fx.Int32(phys_vec[(a * c16) // 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), phys_vec + + # -- prologue: prefetch the first tile's K -- + num_tiles_m1 = num_tiles - 1 + start_safe = arith.select(part_start < num_tiles, part_start, num_tiles_m1) + k_pf0, phys_vec0 = _k_ops_flat(start_safe) + # V page-index prefetch, issued here too for the same overlap; the + # LDS write is visible after the barrier below. + _v_page_fetch_and_stage(start_safe) + if const_expr(per_token_kv): + _stage_kv_scale_to_lds(phys_vec0, _kv_buf_off(fx.Int32(start_safe))) + + # per_token_kv has no single global key_scale/value_scale: scale_qk + # drops the key_scale factor (folded in per-token, see masked_chunks + # below) and v_scale_f is unused (replaced by v_max_scaled). + 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) + fm_contract = arith.FastMathFlags.contract + # Softmax scores are finite or the -inf mask sentinel -- never NaN -- so + # nnan lets maxnum lower to a bare v_max (no v_cmp_u NaN check + its s_nop + # hazard) and fuse to v_max3. (ninf must NOT be set: -inf is load-bearing.) + 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), + ) + + # f32[16, NWARP] cross-warp scratch: scalar write at (row, warp), vec read of a row's NWARP valid slots. + 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): + # f32 -> fp8 must use the HW cvt (arith.truncf to fp8 doesn't lower); + # pack 4 f32 -> 1 i32 (4 fp8) via two cvt_pk_fp8_f32 calls. + 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) + + qh_local = warp * 4 + rgroup # 0..15: this thread's query row within an M-tile + + # Each M-tile quantizes 16 rows of the flattened (MTP position, GQA + # head) axis: `flat_idx = m*16 + qh_local`, decomposed as + # `qi = flat_idx // query_group_size`, `gs_head = flat_idx % + # query_group_size` (same convention as `_mtp_groups`). No + # cross-M-tile dependency, so no barriers needed between iterations. + def _quant_q_row(m, qi, gs_head, q_row_off): + qh0 = kv_h * query_group_size + gs_head + row_byte0 = ( + (seq * query_length + 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 + # QCHUNK elements, loaded as N_QLOADS contiguous QLOAD_UNIT pieces + # (a buffer load is 128b max); head_dim=256 splits into 2 pieces. + q_units = [ + _q_load_chunk(base_elem + u * QLOAD_UNIT) + for u in range_constexpr(N_QLOADS) + ] + + absmax = fmath.absf(q_units[0]).reduce(ReductionOp.MAX).to(fx.Float32) + for u in range_constexpr(1, N_QLOADS): + absmax = fx.maxnumf( + absmax, + fmath.absf(q_units[u]).reduce(ReductionOp.MAX).to(fx.Float32), + ) + 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] (not [m][qh]) so the whole M_TILES-wide + # row for a fixed qh is contiguous, letting the KV-loop read + # it back in one wide load instead of M_TILES separate + # narrow ones -- see the read site below. + _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 + # `flat_idx < TOTAL_ROWS` is statically true for every lane except + # possibly on the last M-tile (only it can be a partial tile), so + # skip the runtime branch (and its EXEC-mask overhead) entirely + # for every other M-tile. + if const_expr((m + 1) * MFMA_MNK <= TOTAL_ROWS): + _quant_q_row(m, qi, gs_head, q_row_off) + elif flat_idx < TOTAL_ROWS: + _quant_q_row(m, qi, gs_head, q_row_off) + else: + _st_words( + q_row_off + qh_local * head_dim + lane16 * QCHUNK, + fx.Vector.filled(QCHUNK // 4, 0, fx.Int32), + ) + if lane16 == 0: + _st1(sQscale_off, qh_local * M_TILES + m, ZERO_F) + + gpu.barrier() + + # First tile's V page-index row, now visible after the barrier above + # (the fetch+LDS-store was issued earlier, alongside k_pf0). + v_page_pf0 = _v_page_read_row() + + # Q is the B operand, read raw from sQ once per M-tile and held in + # registers. MUST use the exact same (qkhe, rgroup, qkr) -> head_dim + # permutation as K's `_k_ops`. + 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]]) + # q_ops_all[m*N_SUBCHUNKS+s] for s=0..N_SUBCHUNKS-1, s = qkhe*2+qkr, + # = M-tile m's head[he_idx*16+qkr*8 : +8] of qhead=lane16 + + # QK in NCHUNK chunks of 4 tokens: each chunk yields a f32x4 + # C-fragment, so softmax processes 4 scores at a time (low VGPR peak). + _ct = [ + fx.Vector.from_elements([float(a * c16 + r) for r in range_constexpr(4)]) + for a in range_constexpr(NCHUNK) + ] + # P.V is loop-tiled over head-dim (like production's VHELOOP): each + # step computes O[:, vh*VHE_SIZE:+VHE_SIZE] instead of materializing + # the full [16, head_dim] at once. + VHE_SIZE = head_dim // VHE_CHUNKS + OP_ELEMS = ( + MFMA_MNK * VHE_SIZE // (NWARP * WAVE) + ) # PV C-fragment elements/lane/chunk (probed = 4) + + # -- raw dwordx4 V load (B operand) -- + # One dwordx4 load per (16-token sub-block, head_elem); trans_v only + # changes the offset formula. `sub`/`step` walk pages/16-token sub-blocks. + NVOPS = TILE_TOK // MFMA_K # 8 PV k_steps (256 tokens / K=32) + STEPS_PER_PAGE = block_size // 16 + + 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_PAGE): + if const_expr(trans_v): + base = ( + ((phys_row[sub] * n_kv + kv_h) * STEPS_PER_PAGE + step) + * head_dim + + head_element + ) * 16 + else: + base = ( + (phys_row[sub] * n_kv + kv_h) * head_dim + head_element + ) * block_size + step * 16 + w = _v_load16(base) + if const_expr(block_size == 16): + # help the scheduler overlap the per-page gathered loads (see _k_ops) + 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 + + # query_length==1: plain wave-uniform `context_len` (folding in the + # per-lane `lane16` would cost a live VGPR for no behavioral change). + if const_expr(query_length == 1): + causal_bound = [context_len for _m in range_constexpr(M_TILES)] + else: + causal_bound = [ + context_len + - (query_length - 1) + + (m * MFMA_MNK + lane16) // query_group_size + for m in range_constexpr(M_TILES) + ] + + K_SLOT, V_SLOT = 0, 1 + # Per-M-tile loop-carried state after the K/V slots: one output chunk + # per VHE_CHUNKS (2 for head_dim=128, 4 for head_dim=256, 1 for 64), + # then running-max + running-denom. + STATE_PER_M = VHE_CHUNKS + 2 + + 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]) + for tt, ostate in range(part_start, part_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(tt) + tok0 = tt * TILE_TOK + # per_tensor phase-split: let IGLP interleave MFMA with softmax + # VALU/LDS to hide the MFMA-hazard s_nop (per_token is VGPR-cliffed). + if const_expr(not per_token_kv and M_TILES > 1): + 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 + + # This tile's ping-pong scale buffer; the tt+1 prefetch stages the other. + cur_kv_buf = _kv_buf_off(tt) + + # V data doesn't depend on `m`; hoist once for the phase-split + # instead of reloading per M-tile. + v_vh_shared = None + if const_expr(M_TILES > 1): + v_vh_shared = [ + _v_ops(v_page_cur, vh) for vh in range_constexpr(VHE_CHUNKS) + ] + + # q_scale doesn't depend on `m`; read the whole M_TILES-wide row once + # (contiguous via the transposed [qh][m] sQscale layout). + 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 if const_expr(M_TILES > 1) else 0 + ) + + # Phase-split (M_TILES>1): pass-1 (QK+mask+max) loops all M-tiles into + # per-M-tile LDS slices so they share ONE barrier. M_TILES==1: else below. + if const_expr(M_TILES > 1): + masked_chunks_saved = [None] * M_TILES + scale_saved = [None] * M_TILES + + # per_token: compute the m-independent pv_max once, hold only + # k_scale across Phase A, re-read v_scale for Phase B after the + # barrier (peak scale liveness 16, not 32). + 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, v_scale_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 = arith.constant_vector(0.0, T.f32x4) + for s in range_constexpr(N_SUBCHUNKS): + acc = fx.rocdl.mfma_f32_16x16x32_fp8_fp8( + T.f32x4, + [ + k_cur[a * N_SUBCHUNKS + s], + q_ops_all[m * N_SUBCHUNKS + s], + acc, + 0, + 0, + 0, + ], + ) + frag_Ss.append(fx.Vector(acc)) + + scale = scale_qk * fx.Float32(q_scale_vec[m]) + n_valid_tile = (causal_bound[m] - tok0).to(fx.Float32) + base_tok_f = fx.Int32(warp * TOK_CHUNK + rgroup * 4).to(fx.Float32) + thr = fx.Vector.from_elements( + [n_valid_tile - base_tok_f], dtype=fx.Float32 + ).broadcast_to(4) + neg4 = fx.Vector.filled(4, float("-inf"), fx.Float32) + + if const_expr(per_token_kv): + scaled_frags = [ + frag_Ss[a] * k_scale_shared[a] + for a in range_constexpr(NCHUNK) + ] + else: + scaled_frags = frag_Ss + + masked_chunks = [ + (_ct[a] < 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 * scale) + + masked_chunks_saved[m] = masked_chunks + scale_saved[m] = scale + + # tt+1 K/V/scale prefetch, issued once here so the V-page read + # reuses this barrier. + k_next = k_cur + if tt1 < part_end: + k_next, phys_vec1 = _k_ops_flat(tt1) + _v_page_fetch_and_stage(tt1) + if const_expr(per_token_kv): + # stage tt+1 into the OTHER buffer (current v_scale re-read below) + _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 tt1 < part_end: + v_page_next = _v_page_read_row() + next_state[V_SLOT] = v_page_next + + # Phase B v_scale: M_TILES>=4 re-reads per chunk (VGPR cliff); + # smaller M_TILES hold all NCHUNK. Both read the current buffer. + 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) + ] + + for m in range_constexpr(M_TILES): + 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] + scale = scale_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) + + m_new = fx.maxnumf( + m_prev, + _ld_lw_row(_lmax_off_m(m), lane16).reduce( + ReductionOp.MAX, fastmath=fm_nnan + ), + fastmath=fm_nnan, + ) + # Fully-invalid row: use 0 as the effective max so masked lanes + # give exp2(-inf-0)==0 (avoids the -inf-(-inf) cancellation). + safe_max = arith.select(m_new > NEG_INF, 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( + exp2_f32_fast(masked_chunks[a] * scale - m_new_b) + ) + 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 = ( + sP_off + lane16 * SP_ROW_BYTES + warp * TOK_CHUNK + rgroup * 4 + ) + # The NCHUNK P words scatter across the row at stride c16//4 + # i32 (the token->fp8-lane interleave the PV ds_read_b128 + # expects); one strided store per chunk. + for a in range_constexpr(NCHUNK): + _lds_store( + p_off0 + a * (c16 // 4) * f32, + fx.Int32, + fx.Vector.from_elements([words[a]], dtype=fx.Int32), + ) + for sh in (16, 32): + ls = ls.addf(ls.shuffle_xor(sh, WAVE), fastmath=fm_contract) + # PV output is [head-dim, query-row=lane16] after the operand + # swap, so correction/denominator are per-lane scalars (no sCorr). + safe_prev = arith.select(m_prev > NEG_INF, m_prev, ZERO_F) + corr_reg = fx.Float32(exp2_amdgcn_scalar(safe_prev - safe_max)) + if rgroup == 0: + _st_lw(sLsum_off, lane16, warp, ls) + gpu.barrier() + gsum = _ld_lw_row(sLsum_off, lane16).reduce(ReductionOp.ADD) + l_new = fx.Float32( + arith.mulf( + arith.unwrap(l_prev), + arith.unwrap(corr_reg), + fastmath=fm_contract, + ) + ).addf(gsum, fastmath=fm_contract) + + 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) + for vh in range_constexpr(VHE_CHUNKS): + v_vh = v_vh_shared[vh] + acc = arith.constant_vector(0.0, T.f32x4) + for s in range_constexpr(NVOPS): + # SWAPPED operands (V=A, P=B): output row = + # head-dim, output col = query-row=lane16. + acc = fx.rocdl.mfma_f32_16x16x32_fp8_fp8( + T.f32x4, [v_vh[s], p_ops[s], acc, 0, 0, 0] + ) + 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] = o_acc[vh] * corr_b + op + next_state.extend([*o_acc, m_new, l_new]) + # sP and sLsum are reused by the next M-tile. Synchronize + # all waves after their LDS reads before any wave overwrites + # those regions, then retire this tile's dependency chain. + if const_expr(m < M_TILES - 1): + gpu.barrier() + fx.rocdl.sched_barrier(0) + else: + # M_TILES==1 single tile (m==0 for the _o_slot/_m_slot/_l_slot helpers). + 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 + # QK: each NCHUNK chunk accumulates N_SUBCHUNKS k_steps into an f32x4. + frag_Ss = [] + for a in range_constexpr(NCHUNK): + acc = arith.constant_vector(0.0, T.f32x4) + for s in range_constexpr(N_SUBCHUNKS): + acc = fx.rocdl.mfma_f32_16x16x32_fp8_fp8( + T.f32x4, + [k_cur[a * N_SUBCHUNKS + s], q_ops_all[s], acc, 0, 0, 0], + ) + frag_Ss.append(fx.Vector(acc)) + # tt+1 K/V/scale prefetch, issued here to reuse the pass-1 barrier. + k_next = k_cur + if tt1 < part_end: + 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 + # Softmax: each lane owns one qhead (lane%16); register reduce + shuffle_xor. + scale = scale_qk * _ld1( + sQscale_off, lane16 + ) # per-qhead positive score scale + n_valid_tile = (causal_bound[0] - tok0).to(fx.Float32) + base_tok_f = fx.Int32(warp * TOK_CHUNK + rgroup * 4).to(fx.Float32) + thr = fx.Vector.from_elements( + [n_valid_tile - base_tok_f], dtype=fx.Float32 + ).broadcast_to(4) + neg4 = fx.Vector.filled(4, -1e30, fx.Float32) + # per_token_kv: K-scale varies per token, so fold it in BEFORE the max-reduce. + v_scale_vecs = None + if const_expr(per_token_kv): + v_scale_vecs = [] + scaled_frags = [] + 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_frags.append(frag_Ss[a] * k_scale_vec) + else: + scaled_frags = frag_Ss + # Reused in pass 2 below, halving the mask instruction count. + masked_chunks = [ + (_ct[a] < 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): + pm = fx.maxnumf(pm, masked_chunks[a].reduce(ReductionOp.MAX)) + for sh in (16, 32): + pm = fx.maxnumf(pm, pm.shuffle_xor(sh, WAVE)) + _st_lw( + sLmax_off, lane16, warp, pm * scale + ) # redundant across the 4 lanes sharing this qhead + # per_token_kv: max V-scale for the per-tile fp8 normalization + # (any positive value is correct, so skip the causal mask). + if const_expr(per_token_kv): + pv_max = fx.Float32(0.0) + for a in range_constexpr(NCHUNK): + pv_max = fx.maxnumf( + pv_max, v_scale_vecs[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-index row for next tile, now visible after the barrier. + v_page_next = v_page_cur + if tt1 < part_end: + v_page_next = _v_page_read_row() + next_state[V_SLOT] = v_page_next + # per_token_kv: combine warps' max V-scale into the tile's + # normalization factor (also the PV correction below). + 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 + m_new = fx.maxnumf( + m_prev, _ld_lw_row(sLmax_off, lane16).reduce(ReductionOp.MAX) + ) + m_new_b = fx.Vector.from_elements( + [m_new], dtype=fx.Float32 + ).broadcast_to(4) + ls = fx.Float32(0.0) + words = [] + zero4_p = fx.Vector.filled(4, 0.0, fx.Float32) + for a in range_constexpr(NCHUNK): + # re-mask Pa so a fully-masked chunk contributes exactly 0 + valid_a = masked_chunks[a] > fx.Vector.filled(4, -1e29, fx.Float32) + Pa = valid_a.select( + fx.Vector(exp2_f32_fast(masked_chunks[a] * scale - m_new_b)), + 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 + 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_CHUNK + rgroup * 4 + # NCHUNK P words scatter at stride c16//4 i32 (see phase-split). + for a in range_constexpr(NCHUNK): + _lds_store( + p_off0 + a * (c16 // 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.addf(ls.shuffle_xor(sh, WAVE), fastmath=fm_contract) + # PV (V=A, P=B) -> output [head-dim, query-row=lane16]; same as + # the phase-split path. + corr_reg = fx.Float32(exp2_amdgcn_scalar(m_prev - m_new)) + if rgroup == 0: + _st_lw(sLsum_off, lane16, warp, ls) + gpu.barrier() + gsum = _ld_lw_row(sLsum_off, lane16).reduce(ReductionOp.ADD) + l_new = fx.Float32( + arith.mulf( + arith.unwrap(l_prev), + arith.unwrap(corr_reg), + fastmath=fm_contract, + ) + ).addf(gsum, fastmath=fm_contract) + 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) + # Single tile: batch both vh's V loads upfront (no sibling chain + # to hide the latency behind). + 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 = arith.constant_vector(0.0, T.f32x4) + for s in range_constexpr(NVOPS): + acc = fx.rocdl.mfma_f32_16x16x32_fp8_fp8( + T.f32x4, [v_vh[s], p_ops[s], acc, 0, 0, 0] + ) + 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] = o_acc[vh] * corr_b + op + next_state.extend([*o_acc, m_new, l_new]) + results = yield next_state + o_final = results + + # Direct-store epilogue: after the PV swap each lane holds its 4 head-dim + # values for one query-row and writes them straight to global. + 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 + l_row = o_final[_l_slot(m)] + safe_l = arith.select(l_row > ZERO_F, l_row, fx.Float32(1.0)) + inv_l = fx.Float32(rcp_f32(safe_l)) + if const_expr(per_token_kv): + o_scale = inv_l + else: + o_scale = fx.Float32( + arith.mulf( + arith.unwrap(inv_l), + arith.unwrap(v_scale_f * inv_fp8), + fastmath=fm_contract, + ) + ) + o_scale_b = fx.Vector.from_elements( + [o_scale], dtype=fx.Float32 + ).broadcast_to(OP_ELEMS) + qi_e = row // query_group_size + gs_head_e = row - qi_e * query_group_size + qh = kv_h * query_group_size + gs_head_e + + def _emit(o_norm, sub): + if const_expr(NP == 1): + out_row = output_ptr[seq * query_length + qi_e, qh, None] + out_chunk = fx.slice( + fx.logical_divide(out_row, fx.make_layout(OP_ELEMS, 1)), + (None, sub), + ) + out_chunk.store(o_norm) + else: + base = ((seq * n_kv + kv_h) * NP + part) * TOTAL_ROWS + row + pout_div = fx.logical_divide(pout_ptr, fx.make_layout(OP_ELEMS, 1)) + pout_chunk = fx.slice( + pout_div, (None, base * (head_dim // OP_ELEMS) + sub) + ) + pout_chunk.store(o_norm) + + 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 + # Guard the partial last tile's out-of-range rows (folded away for full tiles). + if row < TOTAL_ROWS: + _emit(o_norm, sub) + + if const_expr(NP > 1): + if warp == 0 and rgroup == 0: + base = ((seq * n_kv + kv_h) * NP + part) * TOTAL_ROWS + row + if row < TOTAL_ROWS: + # Convert the running max from log2 units (scale_qk folds + # in LOG2E) to natural-log units: the shared reduce + # re-applies LOG2E itself when combining partitions. + pmax_ptr[base] = o_final[_m_slot(m)] * fx.Float32(1.0 / LOG2E) + psum_ptr[base] = l_row + + @flyc.jit + def pa_decode_tile_launch( + output: fx.Tensor, + pmax: fx.Tensor, + psum: fx.Tensor, + pout: fx.Tensor, + query: fx.Tensor, + key_cache: fx.Tensor, + value_cache: fx.Tensor, + block_tables: fx.Tensor, + context_lengths: fx.Tensor, + key_scale: fx.Tensor, # [1] per-tensor OR [num_blocks, num_kv_heads, block_size] per-token + value_scale: fx.Tensor, # same shape as key_scale + 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_q_row: fx.Int32, + stride_q_head: fx.Int32, + stream: fx.Stream = fx.Stream(None), + ): + pa_decode_tile_kernel( + output, + pmax, + psum, + pout, + query, + key_cache, + value_cache, + block_tables, + context_lengths, + key_scale, + value_scale, + max_blocks_per_seq, + stride_ks_block, + stride_ks_head, + stride_q_row, + stride_q_head, + ).launch( + grid=(num_seqs, num_kv_heads, NP), + block=(BLOCK_THREADS, 1, 1), + stream=stream, + ) + + return {"launch": pa_decode_tile_launch, "kernel": pa_decode_tile_kernel} diff --git a/aiter/ops/flydsl/pa_decode_tile.py b/aiter/ops/flydsl/pa_decode_tile.py index e867720f5f3..e9d498cb0d5 100644 --- a/aiter/ops/flydsl/pa_decode_tile.py +++ b/aiter/ops/flydsl/pa_decode_tile.py @@ -26,1215 +26,11 @@ an LDS round-trip on P transposing ownership between the two MMAs. """ -import functools - import torch -import flydsl.compiler as flyc -import flydsl.expr as fx -from flydsl.compiler.protocol import dsl_size_of -from flydsl.expr import arith, buffer_ops, const_expr, gpu, range_constexpr, vector -from flydsl.expr import math as fmath -from flydsl.expr.typing import T -from flydsl.expr.vector import ReductionOp -from flydsl.runtime.device import get_rocm_arch - -from .kernels import dpp_utils +from .kernels.pa_decode_tile import KV_COMPUTE_BLOCK, compile_pa_decode_tile from .kernels.tensor_shim import _run_compiled -from .kernels.utils import cdiv, exp2_amdgcn_scalar, exp2_f32_fast, rcp_f32 - -MFMA_MNK = ( - 16 # M = N = 16 for the MMA atom; also query rows handled per CTA (padded to 16) -) -MFMA_K = 32 # fp8 MFMA contracts K = 32 per instruction (mfma_f32_16x16x32_fp8_fp8) -WAVE = 64 -LOG2E = 1.4426950408889634 -KV_COMPUTE_BLOCK = 256 - - -@functools.lru_cache(maxsize=None) -def compile_pa_decode_tile( - *, - head_dim: int, - query_group_size: int, - block_size: 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, -): - """Build the tile-programming PA-decode kernel + launch wrapper. - - ``block_size``, ``head_dim``, and ``query_dtype`` are compile-time - constants (lru_cache keys). ``query_length`` (MTP) and ``query_group_size`` - flatten into ``TOTAL_ROWS = query_length * query_group_size``, tiled into - ``M_TILES = ceil(TOTAL_ROWS / 16)`` independent 16-row MFMA tiles; each - extra M-tile duplicates loop-carried state, so VGPR/LDS/occupancy scale - roughly linearly with ``M_TILES``. - """ - is_gfx950 = "gfx95" in get_rocm_arch() - 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, - ), f"pa_decode_tile only supports block_size in (16, 64), 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 query_dtype == "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}" - assert query_length >= 1, f"query_length must be >= 1, got {query_length}" - # Flattened query-row axis (MTP outer, GQA head inner), tiled into 16-row M-tiles. - TOTAL_ROWS = query_length * query_group_size - M_TILES = cdiv(TOTAL_ROWS, MFMA_MNK) - ROWS_PADDED = M_TILES * MFMA_MNK - # PV layout: V=A, P=B -> output [head-dim (row), query-row (col=lane16)], - # generalized over head_dim via the VHE_CHUNKS loop. - NWARP = 4 # 4 waves / CTA - TOK_PER_WARP = 64 # tokens each warp owns per compute block (matches production KV_COMPUTE_BLOCK) - TILE_TOK = NWARP * TOK_PER_WARP # 256 tokens / compute block - PAGES_PER_CHUNK = ( - TOK_PER_WARP // block_size - ) # pages spanned by one 64-token warp-chunk: 1 (bs=64) or 4 (bs=16) - assert ( - head_dim % (NWARP * MFMA_MNK) == 0 - ), "head_dim must split across the 4 warps for PV" - - # head_dim splits into 16-element chunks (QK_CHUNK_ELEMS, one dwordx4 load); - # RGROUP_QUARTERS of them make a 64-element fetch group, QKHE_LOOP groups total. - 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}" - N_SUBCHUNKS = QKHE_LOOP * (QK_CHUNK_ELEMS // 8) - - # Q-quant chunk width: NQCHUNK stays fixed at 16 (tied to `lane16`'s role - # as the absmax butterfly width); QCHUNK scales with head_dim instead. - NQCHUNK = 16 - QCHUNK = ( - head_dim // NQCHUNK - ) # f16 elements per lane's load chunk (8 for head_dim=128, 4 for head_dim=64) - - VHE_CHUNKS = head_dim // (NWARP * MFMA_MNK) # 2 for head_dim=128, 1 for head_dim=64 - - if softmax_scale is None: - softmax_scale = 1.0 / (head_dim**0.5) - NP = int(num_partitions) # context partitions (grid.z); compile-time constant - - BLOCK_THREADS = NWARP * WAVE # 256 - - # -- LDS layout (shared across the 4 warps) -- - # sQ: fp8[ROWS_PADDED,head_dim] staged+quantized query. sP: fp8[16,TILE_TOK] - # quantized probs. sQscale: f32[ROWS_PADDED]. sLmax/sLsum: cross-warp scratch. - # sVPage: V page-index broadcast. sKScale/sVScale/sVScaleMax: per-token K/V - # scale staging. No sO/sM/sL/sCorr: PV output is register-resident/loop-carried - # and stored straight to global (V=A/P=B swap). - f32 = 4 - sQ_bytes = ROWS_PADDED * head_dim * 1 # fp8 - sP_off = sQ_bytes - # +16B row padding breaks a 32-bank LDS conflict on the P-pack writes while - # keeping the row 16B-aligned for PV's ds_read_b128. - SP_ROW_BYTES = TILE_TOK + 16 - sP_bytes = MFMA_MNK * SP_ROW_BYTES # fp8, padded rows - sQscale_off = sP_off + sP_bytes - NWARP_PAD = ( - NWARP + 1 - ) # coprime with 32 banks -> no LDS bank conflict on the cross-warp scratch - # Phase-split slices sLmax per M-tile so all pass-1 writes share one barrier. - sLmax_off = sQscale_off + ROWS_PADDED * f32 - sLsum_off = sLmax_off + M_TILES * MFMA_MNK * NWARP_PAD * f32 - # V page-index broadcast (V's page depends on `rgroup`, shared across warps). - sVPage_off = sLsum_off + MFMA_MNK * NWARP_PAD * f32 - sVPage_bytes = NWARP * PAGES_PER_CHUNK * 4 # i32 - # Per-token KV-scale staging (per_token only), double-buffered so the tt+1 - # prefetch (into the other buffer) doesn't clobber the current tile's scales. - KV_BUFS = 2 if per_token_kv else 1 - KV_BUF_STRIDE = 2 * NWARP * TOK_PER_WARP * f32 # k-region + v-region, one buffer - sKScale_off = sVPage_off + sVPage_bytes - sVScale_off = sKScale_off + NWARP * TOK_PER_WARP * f32 - sKVScale_bytes = KV_BUFS * 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 - - # LDS blob: one i32 array carved into typed byte-offset views (see the - # `_lds_*` helpers in the kernel). Every region size above is 4-byte - # aligned, so `total_bytes // 4` covers the blob exactly. - @fx.struct - class SharedStorage: - buf: fx.Array[fx.Int32, total_bytes // 4, 16] - - @flyc.kernel(known_block_size=(BLOCK_THREADS, 1, 1)) - def pa_decode_tile_kernel( - output_ptr: fx.Tensor, # [num_seqs*query_length, num_q_heads, head_dim] (written directly when NP==1) - # per-partition partial outputs (combined by the reduce kernel when NP>1): - pmax_ptr: fx.Tensor, # [num_seqs*query_length, num_kv_heads, num_partitions, query_group_size] row max - psum_ptr: fx.Tensor, # [num_seqs*query_length, num_kv_heads, num_partitions, query_group_size] row sum - pout_ptr: fx.Tensor, # [num_seqs*query_length, num_kv_heads, num_partitions, query_group_size, head_dim] Q_DTYPE, normalized O_p/l_p - query_ptr: fx.Tensor, # [num_seqs*query_length, num_q_heads, head_dim] -- row = seq*query_length + qi (MTP position) - key_cache_ptr: fx.Tensor, # [num_blocks, num_kv_heads, head_dim//16, block_size, 16] (blocked, see module docstring) - value_cache_ptr: fx.Tensor, # [num_blocks, num_kv_heads, block_size//16, head_dim, 16] (blocked, see module docstring) - block_tables_ptr: fx.Tensor, # [num_seqs, max_blocks_per_seq] - context_lengths_ptr: fx.Tensor, # [num_seqs] - key_scale_ptr: fx.Tensor, # [1] per-tensor OR [num_blocks, num_kv_heads, block_size] per-token - value_scale_ptr: fx.Tensor, # same shape as key_scale_ptr - max_blocks_per_seq: fx.Int32, - stride_ks_block: fx.Int32, - stride_ks_head: fx.Int32, - stride_q_row: fx.Int32, - stride_q_head: 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_h = fx.Int32(gpu.block_id("y")) - part = fx.Int32(gpu.block_id("z")) # context partition handled by this CTA - n_kv = fx.Int32(gpu.grid_dim.y) # num_kv_heads == gridDim.y - - # fx.copy-based flat loaders: K/V/context_len over a raw pointer, Q over - # a buffer resource. - def _make_flat_loader(tensor_ptr, elem_ty, reg_width, copy_op): - use_buffer_resource = isinstance( - copy_op, fx.rocdl.CopyOpCDNA3BufferCopyType - ) - copy_atom = fx.make_copy_atom(copy_op, elem_ty) - reg = fx.make_rmem_tensor(fx.make_layout(reg_width, 1), elem_ty) - base_iter = ( - fx.get_iter(fx.rocdl.make_buffer_tensor(tensor_ptr, max_size=True)) - if use_buffer_resource - else fx.get_iter(tensor_ptr) - ) - flat = fx.Tensor(fx.make_view(base_iter, fx.make_layout(1 << 30, 1))) - div = fx.logical_divide(flat, fx.make_layout(1, 1)) - - def _load(elem_idx): - fx.copy(copy_atom, fx.slice(div, (None, elem_idx)), reg) - return fx.Vector(fx.memref_load_vec(reg)) - - return _load - - _k_load_fp8x16 = _make_flat_loader( - key_cache_ptr, FP8, 16, fx.UniversalCopy128b() - ) - _v_load_fp8x16 = _make_flat_loader( - value_cache_ptr, FP8, 16, fx.UniversalCopy128b() - ) - # Per-lane Q chunk (QCHUNK 16-bit elems) fetched in QLOAD_UNIT-wide - # pieces (128b max per buffer load): head_dim=256 needs 2 pieces. - QLOAD_UNIT = QCHUNK if QCHUNK < 8 else 8 - N_QLOADS = QCHUNK // QLOAD_UNIT - _q_copy_op = ( - fx.rocdl.BufferCopy128b() if QLOAD_UNIT == 8 else fx.rocdl.BufferCopy64b() - ) - _q_load_chunk = _make_flat_loader(query_ptr, Q_DTYPE, QLOAD_UNIT, _q_copy_op) - _ctxlen_load = _make_flat_loader( - context_lengths_ptr, fx.Int32, 1, fx.rocdl.BufferCopy32b() - ) - - 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) - - context_len = fx.Int32(_ctxlen_load(seq)[0]) - # Bound block_tables to its real extent: the last (partial) 256-token tile - # can index a page past ceil(context/block_size); the bounded resource - # returns page 0 for that out-of-range read instead of faulting (those - # tail tokens are masked out anyway). - bt_num_records_bytes = ( - fx.Index(gpu.grid_dim.x) * fx.Index(max_blocks_per_seq) * 4 - ) # int32 entries - bt_rsrc = buffer_ops.create_buffer_resource( - block_tables_ptr, max_size=False, num_records_bytes=bt_num_records_bytes - ) - ks_rsrc = buffer_ops.create_buffer_resource(key_scale_ptr, max_size=True) - vs_rsrc = buffer_ops.create_buffer_resource(value_scale_ptr, max_size=True) - # Per-tensor: a single global scale, read once. Per-token: read - # per-token instead (see _kv_scale_ops/_stage_kv_scale_to_lds below). - if const_expr(not per_token_kv): - key_scale = fx.Int32( - buffer_ops.buffer_load( - ks_rsrc, arith.constant(0, type=T.i32), vec_width=1, is_scalar=True - ) - ).bitcast(fx.Float32) - value_scale = fx.Int32( - buffer_ops.buffer_load( - vs_rsrc, arith.constant(0, type=T.i32), vec_width=1, is_scalar=True - ) - ).bitcast(fx.Float32) - - num_tiles = cdiv(context_len, TILE_TOK) - tiles_per_part = cdiv(num_tiles, NP) - part_start = part * tiles_per_part - part_end_raw = part_start + tiles_per_part - part_end = arith.select(part_end_raw < num_tiles, part_end_raw, num_tiles) - - # One i8 blob carved into typed byte-offset pointers. `lds_base` is an - # ir.Value pointer (safe inside scf control flow); the Python `lds` - # handle is not. - 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): - # Pin the element-size alignment explicitly (a bare recast off the - # byte base would gcd the alignment down to 1). - 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)) - - c16 = 16 - rgroup = lane // c16 # 0..3: which quarter-wave (paired with warp -> query row) - lane16 = lane - rgroup * c16 # 0..15: this row's head-dim chunk index - - TOK_CHUNK = NWARP * MFMA_MNK # 64 - NCHUNK = TILE_TOK // TOK_CHUNK # 4 - - def _load_phys_scalar(page, vec_width=1): - result = buffer_ops.buffer_load( - bt_rsrc, - seq * max_blocks_per_seq + page, - vec_width=vec_width, - is_scalar=True, - ) - return fx.Int32(result) if vec_width == 1 else result - - def _v_page_fetch_and_stage(tt_i32): - # V's page depends on `rgroup` (shared across warps): warp w fetches - # its rgroup row and broadcasts via LDS (read back by _v_page_read_row). - base_page = tt_i32 * TILE_TOK // block_size # tile start is page-aligned - fetched = _load_phys_scalar( - base_page + warp * PAGES_PER_CHUNK, PAGES_PER_CHUNK - ) - if lane == 0: - fetched_vec = ( - fx.Vector.from_elements([fx.Int32(fetched)], dtype=fx.Int32) - if const_expr(PAGES_PER_CHUNK == 1) - else fx.Vector(fetched) - ) - _lds_store( - sVPage_off + warp * (PAGES_PER_CHUNK * 4), fx.Int32, 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 _kv_buf_off(tt_val): - # ping-pong buffer byte offset (0 when single-buffered). - if const_expr(per_token_kv): - 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]) - base_tok = lane16 * NCHUNK - scale_idx = phys * stride_ks_block + kv_h * stride_ks_head + base_tok - k_scale_vec = fx.Vector( - buffer_ops.buffer_load( - ks_rsrc, scale_idx, vec_width=NCHUNK, dtype=fx.Float32 - ) - ) - v_scale_vec = fx.Vector( - buffer_ops.buffer_load( - vs_rsrc, scale_idx, vec_width=NCHUNK, dtype=fx.Float32 - ) - ) - slot = (warp * TOK_PER_WARP + base_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: - # block_size==16: each lane stages its own rgroup's page/token, - # so the 4 sub-blocks stage in parallel across rgroup-groups. - phys = fx.Int32( - vector.extract( - arith.unwrap(phys_vec), - static_position=[], - dynamic_position=[fx.Index(rgroup)], - ) - ) - scale_idx = phys * stride_ks_block + kv_h * stride_ks_head + lane16 - k_scale_scalar = fx.Float32( - buffer_ops.buffer_load( - ks_rsrc, scale_idx, vec_width=1, dtype=fx.Float32 - ) - ) - v_scale_scalar = fx.Float32( - buffer_ops.buffer_load( - vs_rsrc, scale_idx, vec_width=1, dtype=fx.Float32 - ) - ) - fx.rocdl.sched_barrier(fx.rocdl.mask_vmem_rd) - slot = (warp * TOK_PER_WARP + rgroup * c16 + 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): - # This lane's 4 per-token scales for chunk `a` from an LDS scale region. - slot = (warp * TOK_CHUNK + a * c16 + 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 - ) - - # -- raw dwordx4 K load (A operand) -- - # token = warp*TOK_CHUNK + a*c16 + lane16 (the softmax mask and P-pack - # write position below must encode this same formula). - def _k_ops(phys, a): - within_page_tok = (a * c16 + lane16) % block_size - ops = [] - for qkhe in range_constexpr(QKHE_LOOP): - he_idx = qkhe * RGROUP_QUARTERS + rgroup - base = ( - ((phys * n_kv + kv_h) * QCHUNK + he_idx) * block_size - + within_page_tok - ) * QK_CHUNK_ELEMS - w = _k_load16(base) # head[he_idx*16 : +16] -> k_step 2*qkhe, 2*qkhe+1 - if const_expr(block_size == 16): - # help the scheduler overlap the PAGES_PER_CHUNK gathered loads - fx.rocdl.sched_barrier(fx.rocdl.mask_vmem_rd) - ops.extend([w[0], w[1]]) - return ops # N_SUBCHUNKS i64 operands - - 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 * PAGES_PER_CHUNK, 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) - ) - flat = [] - for a in range_constexpr(NCHUNK): - phys = fx.Int32(phys_vec[(a * c16) // 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), phys_vec - - # -- prologue: prefetch the first tile's K -- - num_tiles_m1 = num_tiles - 1 - start_safe = arith.select(part_start < num_tiles, part_start, num_tiles_m1) - k_pf0, phys_vec0 = _k_ops_flat(start_safe) - # V page-index prefetch, issued here too for the same overlap; the - # LDS write is visible after the barrier below. - _v_page_fetch_and_stage(start_safe) - if const_expr(per_token_kv): - _stage_kv_scale_to_lds(phys_vec0, _kv_buf_off(fx.Int32(start_safe))) - - # per_token_kv has no single global key_scale/value_scale: scale_qk - # drops the key_scale factor (folded in per-token, see masked_chunks - # below) and v_scale_f is unused (replaced by v_max_scaled). - 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) - fm_contract = arith.FastMathFlags.contract - # Softmax scores are finite or the -inf mask sentinel -- never NaN -- so - # nnan lets maxnum lower to a bare v_max (no v_cmp_u NaN check + its s_nop - # hazard) and fuse to v_max3. (ninf must NOT be set: -inf is load-bearing.) - 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), - ) - - # f32[16, NWARP] cross-warp scratch: scalar write at (row, warp), vec read of a row's NWARP valid slots. - 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): - # f32 -> fp8 must use the HW cvt (arith.truncf to fp8 doesn't lower); - # pack 4 f32 -> 1 i32 (4 fp8) via two cvt_pk_fp8_f32 calls. - 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) - - qh_local = warp * 4 + rgroup # 0..15: this thread's query row within an M-tile - - # Each M-tile quantizes 16 rows of the flattened (MTP position, GQA - # head) axis: `flat_idx = m*16 + qh_local`, decomposed as - # `qi = flat_idx // query_group_size`, `gs_head = flat_idx % - # query_group_size` (same convention as `_mtp_groups`). No - # cross-M-tile dependency, so no barriers needed between iterations. - def _quant_q_row(m, qi, gs_head, q_row_off): - qh0 = kv_h * query_group_size + gs_head - row_byte0 = ( - (seq * query_length + 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 - # QCHUNK elements, loaded as N_QLOADS contiguous QLOAD_UNIT pieces - # (a buffer load is 128b max); head_dim=256 splits into 2 pieces. - q_units = [ - _q_load_chunk(base_elem + u * QLOAD_UNIT) - for u in range_constexpr(N_QLOADS) - ] - - absmax = fmath.absf(q_units[0]).reduce(ReductionOp.MAX).to(fx.Float32) - for u in range_constexpr(1, N_QLOADS): - absmax = fx.maxnumf( - absmax, - fmath.absf(q_units[u]).reduce(ReductionOp.MAX).to(fx.Float32), - ) - 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] (not [m][qh]) so the whole M_TILES-wide - # row for a fixed qh is contiguous, letting the KV-loop read - # it back in one wide load instead of M_TILES separate - # narrow ones -- see the read site below. - _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 - # `flat_idx < TOTAL_ROWS` is statically true for every lane except - # possibly on the last M-tile (only it can be a partial tile), so - # skip the runtime branch (and its EXEC-mask overhead) entirely - # for every other M-tile. - if const_expr((m + 1) * MFMA_MNK <= TOTAL_ROWS): - _quant_q_row(m, qi, gs_head, q_row_off) - elif flat_idx < TOTAL_ROWS: - _quant_q_row(m, qi, gs_head, q_row_off) - else: - _st_words( - q_row_off + qh_local * head_dim + lane16 * QCHUNK, - fx.Vector.filled(QCHUNK // 4, 0, fx.Int32), - ) - if lane16 == 0: - _st1(sQscale_off, qh_local * M_TILES + m, ZERO_F) - - gpu.barrier() - - # First tile's V page-index row, now visible after the barrier above - # (the fetch+LDS-store was issued earlier, alongside k_pf0). - v_page_pf0 = _v_page_read_row() - - # Q is the B operand, read raw from sQ once per M-tile and held in - # registers. MUST use the exact same (qkhe, rgroup, qkr) -> head_dim - # permutation as K's `_k_ops`. - 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]]) - # q_ops_all[m*N_SUBCHUNKS+s] for s=0..N_SUBCHUNKS-1, s = qkhe*2+qkr, - # = M-tile m's head[he_idx*16+qkr*8 : +8] of qhead=lane16 - - # QK in NCHUNK chunks of 4 tokens: each chunk yields a f32x4 - # C-fragment, so softmax processes 4 scores at a time (low VGPR peak). - _ct = [ - fx.Vector.from_elements([float(a * c16 + r) for r in range_constexpr(4)]) - for a in range_constexpr(NCHUNK) - ] - # P.V is loop-tiled over head-dim (like production's VHELOOP): each - # step computes O[:, vh*VHE_SIZE:+VHE_SIZE] instead of materializing - # the full [16, head_dim] at once. - VHE_SIZE = head_dim // VHE_CHUNKS - OP_ELEMS = ( - MFMA_MNK * VHE_SIZE // (NWARP * WAVE) - ) # PV C-fragment elements/lane/chunk (probed = 4) - - # -- raw dwordx4 V load (B operand) -- - # One dwordx4 load per (16-token sub-block, head_elem); trans_v only - # changes the offset formula. `sub`/`step` walk pages/16-token sub-blocks. - NVOPS = TILE_TOK // MFMA_K # 8 PV k_steps (256 tokens / K=32) - STEPS_PER_PAGE = block_size // 16 - - 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_PAGE): - if const_expr(trans_v): - base = ( - ((phys_row[sub] * n_kv + kv_h) * STEPS_PER_PAGE + step) - * head_dim - + head_element - ) * 16 - else: - base = ( - (phys_row[sub] * n_kv + kv_h) * head_dim + head_element - ) * block_size + step * 16 - w = _v_load16(base) - if const_expr(block_size == 16): - # help the scheduler overlap the per-page gathered loads (see _k_ops) - 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 - - # query_length==1: plain wave-uniform `context_len` (folding in the - # per-lane `lane16` would cost a live VGPR for no behavioral change). - if const_expr(query_length == 1): - causal_bound = [context_len for _m in range_constexpr(M_TILES)] - else: - causal_bound = [ - context_len - - (query_length - 1) - + (m * MFMA_MNK + lane16) // query_group_size - for m in range_constexpr(M_TILES) - ] - - K_SLOT, V_SLOT = 0, 1 - # Per-M-tile loop-carried state after the K/V slots: one output chunk - # per VHE_CHUNKS (2 for head_dim=128, 4 for head_dim=256, 1 for 64), - # then running-max + running-denom. - STATE_PER_M = VHE_CHUNKS + 2 - - 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]) - for tt, ostate in range(part_start, part_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(tt) - tok0 = tt * TILE_TOK - # per_tensor phase-split: let IGLP interleave MFMA with softmax - # VALU/LDS to hide the MFMA-hazard s_nop (per_token is VGPR-cliffed). - if const_expr(not per_token_kv and M_TILES > 1): - 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 - - # This tile's ping-pong scale buffer; the tt+1 prefetch stages the other. - cur_kv_buf = _kv_buf_off(tt) - - # V data doesn't depend on `m`; hoist once for the phase-split - # instead of reloading per M-tile. - v_vh_shared = None - if const_expr(M_TILES > 1): - v_vh_shared = [ - _v_ops(v_page_cur, vh) for vh in range_constexpr(VHE_CHUNKS) - ] - - # q_scale doesn't depend on `m`; read the whole M_TILES-wide row once - # (contiguous via the transposed [qh][m] sQscale layout). - 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 if const_expr(M_TILES > 1) else 0 - ) - - # Phase-split (M_TILES>1): pass-1 (QK+mask+max) loops all M-tiles into - # per-M-tile LDS slices so they share ONE barrier. M_TILES==1: else below. - if const_expr(M_TILES > 1): - masked_chunks_saved = [None] * M_TILES - scale_saved = [None] * M_TILES - - # per_token: compute the m-independent pv_max once, hold only - # k_scale across Phase A, re-read v_scale for Phase B after the - # barrier (peak scale liveness 16, not 32). - 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, v_scale_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 = arith.constant_vector(0.0, T.f32x4) - for s in range_constexpr(N_SUBCHUNKS): - acc = fx.rocdl.mfma_f32_16x16x32_fp8_fp8( - T.f32x4, - [ - k_cur[a * N_SUBCHUNKS + s], - q_ops_all[m * N_SUBCHUNKS + s], - acc, - 0, - 0, - 0, - ], - ) - frag_Ss.append(fx.Vector(acc)) - - scale = scale_qk * fx.Float32(q_scale_vec[m]) - n_valid_tile = (causal_bound[m] - tok0).to(fx.Float32) - base_tok_f = fx.Int32(warp * TOK_CHUNK + rgroup * 4).to(fx.Float32) - thr = fx.Vector.from_elements( - [n_valid_tile - base_tok_f], dtype=fx.Float32 - ).broadcast_to(4) - neg4 = fx.Vector.filled(4, float("-inf"), fx.Float32) - - if const_expr(per_token_kv): - scaled_frags = [ - frag_Ss[a] * k_scale_shared[a] - for a in range_constexpr(NCHUNK) - ] - else: - scaled_frags = frag_Ss - - masked_chunks = [ - (_ct[a] < 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 * scale) - - masked_chunks_saved[m] = masked_chunks - scale_saved[m] = scale - - # tt+1 K/V/scale prefetch, issued once here so the V-page read - # reuses this barrier. - k_next = k_cur - if tt1 < part_end: - k_next, phys_vec1 = _k_ops_flat(tt1) - _v_page_fetch_and_stage(tt1) - if const_expr(per_token_kv): - # stage tt+1 into the OTHER buffer (current v_scale re-read below) - _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 tt1 < part_end: - v_page_next = _v_page_read_row() - next_state[V_SLOT] = v_page_next - - # Phase B v_scale: M_TILES>=4 re-reads per chunk (VGPR cliff); - # smaller M_TILES hold all NCHUNK. Both read the current buffer. - 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) - ] - - for m in range_constexpr(M_TILES): - 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] - scale = scale_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) - - m_new = fx.maxnumf( - m_prev, - _ld_lw_row(_lmax_off_m(m), lane16).reduce( - ReductionOp.MAX, fastmath=fm_nnan - ), - fastmath=fm_nnan, - ) - # Fully-invalid row: use 0 as the effective max so masked lanes - # give exp2(-inf-0)==0 (avoids the -inf-(-inf) cancellation). - safe_max = arith.select(m_new > NEG_INF, 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( - exp2_f32_fast(masked_chunks[a] * scale - m_new_b) - ) - 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 = ( - sP_off + lane16 * SP_ROW_BYTES + warp * TOK_CHUNK + rgroup * 4 - ) - # The NCHUNK P words scatter across the row at stride c16//4 - # i32 (the token->fp8-lane interleave the PV ds_read_b128 - # expects); one strided store per chunk. - for a in range_constexpr(NCHUNK): - _lds_store( - p_off0 + a * (c16 // 4) * f32, - fx.Int32, - fx.Vector.from_elements([words[a]], dtype=fx.Int32), - ) - for sh in (16, 32): - ls = ls.addf(ls.shuffle_xor(sh, WAVE), fastmath=fm_contract) - # PV output is [head-dim, query-row=lane16] after the operand - # swap, so correction/denominator are per-lane scalars (no sCorr). - safe_prev = arith.select(m_prev > NEG_INF, m_prev, ZERO_F) - corr_reg = fx.Float32(exp2_amdgcn_scalar(safe_prev - safe_max)) - if rgroup == 0: - _st_lw(sLsum_off, lane16, warp, ls) - gpu.barrier() - gsum = _ld_lw_row(sLsum_off, lane16).reduce(ReductionOp.ADD) - l_new = fx.Float32( - arith.mulf( - arith.unwrap(l_prev), - arith.unwrap(corr_reg), - fastmath=fm_contract, - ) - ).addf(gsum, fastmath=fm_contract) - - 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) - for vh in range_constexpr(VHE_CHUNKS): - v_vh = v_vh_shared[vh] - acc = arith.constant_vector(0.0, T.f32x4) - for s in range_constexpr(NVOPS): - # SWAPPED operands (V=A, P=B): output row = - # head-dim, output col = query-row=lane16. - acc = fx.rocdl.mfma_f32_16x16x32_fp8_fp8( - T.f32x4, [v_vh[s], p_ops[s], acc, 0, 0, 0] - ) - 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] = o_acc[vh] * corr_b + op - next_state.extend([*o_acc, m_new, l_new]) - # sP and sLsum are reused by the next M-tile. Synchronize - # all waves after their LDS reads before any wave overwrites - # those regions, then retire this tile's dependency chain. - if const_expr(m < M_TILES - 1): - gpu.barrier() - fx.rocdl.sched_barrier(0) - else: - # M_TILES==1 single tile (m==0 for the _o_slot/_m_slot/_l_slot helpers). - 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 - # QK: each NCHUNK chunk accumulates N_SUBCHUNKS k_steps into an f32x4. - frag_Ss = [] - for a in range_constexpr(NCHUNK): - acc = arith.constant_vector(0.0, T.f32x4) - for s in range_constexpr(N_SUBCHUNKS): - acc = fx.rocdl.mfma_f32_16x16x32_fp8_fp8( - T.f32x4, - [k_cur[a * N_SUBCHUNKS + s], q_ops_all[s], acc, 0, 0, 0], - ) - frag_Ss.append(fx.Vector(acc)) - # tt+1 K/V/scale prefetch, issued here to reuse the pass-1 barrier. - k_next = k_cur - if tt1 < part_end: - 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 - # Softmax: each lane owns one qhead (lane%16); register reduce + shuffle_xor. - scale = scale_qk * _ld1( - sQscale_off, lane16 - ) # per-qhead positive score scale - n_valid_tile = (causal_bound[0] - tok0).to(fx.Float32) - base_tok_f = fx.Int32(warp * TOK_CHUNK + rgroup * 4).to(fx.Float32) - thr = fx.Vector.from_elements( - [n_valid_tile - base_tok_f], dtype=fx.Float32 - ).broadcast_to(4) - neg4 = fx.Vector.filled(4, -1e30, fx.Float32) - # per_token_kv: K-scale varies per token, so fold it in BEFORE the max-reduce. - v_scale_vecs = None - if const_expr(per_token_kv): - v_scale_vecs = [] - scaled_frags = [] - 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_frags.append(frag_Ss[a] * k_scale_vec) - else: - scaled_frags = frag_Ss - # Reused in pass 2 below, halving the mask instruction count. - masked_chunks = [ - (_ct[a] < 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): - pm = fx.maxnumf(pm, masked_chunks[a].reduce(ReductionOp.MAX)) - for sh in (16, 32): - pm = fx.maxnumf(pm, pm.shuffle_xor(sh, WAVE)) - _st_lw( - sLmax_off, lane16, warp, pm * scale - ) # redundant across the 4 lanes sharing this qhead - # per_token_kv: max V-scale for the per-tile fp8 normalization - # (any positive value is correct, so skip the causal mask). - if const_expr(per_token_kv): - pv_max = fx.Float32(0.0) - for a in range_constexpr(NCHUNK): - pv_max = fx.maxnumf( - pv_max, v_scale_vecs[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-index row for next tile, now visible after the barrier. - v_page_next = v_page_cur - if tt1 < part_end: - v_page_next = _v_page_read_row() - next_state[V_SLOT] = v_page_next - # per_token_kv: combine warps' max V-scale into the tile's - # normalization factor (also the PV correction below). - 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 - m_new = fx.maxnumf( - m_prev, _ld_lw_row(sLmax_off, lane16).reduce(ReductionOp.MAX) - ) - m_new_b = fx.Vector.from_elements( - [m_new], dtype=fx.Float32 - ).broadcast_to(4) - ls = fx.Float32(0.0) - words = [] - zero4_p = fx.Vector.filled(4, 0.0, fx.Float32) - for a in range_constexpr(NCHUNK): - # re-mask Pa so a fully-masked chunk contributes exactly 0 - valid_a = masked_chunks[a] > fx.Vector.filled(4, -1e29, fx.Float32) - Pa = valid_a.select( - fx.Vector(exp2_f32_fast(masked_chunks[a] * scale - m_new_b)), - 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 - 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_CHUNK + rgroup * 4 - # NCHUNK P words scatter at stride c16//4 i32 (see phase-split). - for a in range_constexpr(NCHUNK): - _lds_store( - p_off0 + a * (c16 // 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.addf(ls.shuffle_xor(sh, WAVE), fastmath=fm_contract) - # PV (V=A, P=B) -> output [head-dim, query-row=lane16]; same as - # the phase-split path. - corr_reg = fx.Float32(exp2_amdgcn_scalar(m_prev - m_new)) - if rgroup == 0: - _st_lw(sLsum_off, lane16, warp, ls) - gpu.barrier() - gsum = _ld_lw_row(sLsum_off, lane16).reduce(ReductionOp.ADD) - l_new = fx.Float32( - arith.mulf( - arith.unwrap(l_prev), - arith.unwrap(corr_reg), - fastmath=fm_contract, - ) - ).addf(gsum, fastmath=fm_contract) - 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) - # Single tile: batch both vh's V loads upfront (no sibling chain - # to hide the latency behind). - 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 = arith.constant_vector(0.0, T.f32x4) - for s in range_constexpr(NVOPS): - acc = fx.rocdl.mfma_f32_16x16x32_fp8_fp8( - T.f32x4, [v_vh[s], p_ops[s], acc, 0, 0, 0] - ) - 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] = o_acc[vh] * corr_b + op - next_state.extend([*o_acc, m_new, l_new]) - results = yield next_state - o_final = results - - # Direct-store epilogue: after the PV swap each lane holds its 4 head-dim - # values for one query-row and writes them straight to global. - 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 - l_row = o_final[_l_slot(m)] - safe_l = arith.select(l_row > ZERO_F, l_row, fx.Float32(1.0)) - inv_l = fx.Float32(rcp_f32(safe_l)) - if const_expr(per_token_kv): - o_scale = inv_l - else: - o_scale = fx.Float32( - arith.mulf( - arith.unwrap(inv_l), - arith.unwrap(v_scale_f * inv_fp8), - fastmath=fm_contract, - ) - ) - o_scale_b = fx.Vector.from_elements( - [o_scale], dtype=fx.Float32 - ).broadcast_to(OP_ELEMS) - qi_e = row // query_group_size - gs_head_e = row - qi_e * query_group_size - qh = kv_h * query_group_size + gs_head_e - - def _emit(o_norm, sub): - if const_expr(NP == 1): - out_row = output_ptr[seq * query_length + qi_e, qh, None] - out_chunk = fx.slice( - fx.logical_divide(out_row, fx.make_layout(OP_ELEMS, 1)), - (None, sub), - ) - out_chunk.store(o_norm) - else: - base = ((seq * n_kv + kv_h) * NP + part) * TOTAL_ROWS + row - pout_div = fx.logical_divide(pout_ptr, fx.make_layout(OP_ELEMS, 1)) - pout_chunk = fx.slice( - pout_div, (None, base * (head_dim // OP_ELEMS) + sub) - ) - pout_chunk.store(o_norm) - - 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 - # Guard the partial last tile's out-of-range rows (folded away for full tiles). - if row < TOTAL_ROWS: - _emit(o_norm, sub) - - if const_expr(NP > 1): - if warp == 0 and rgroup == 0: - base = ((seq * n_kv + kv_h) * NP + part) * TOTAL_ROWS + row - if row < TOTAL_ROWS: - # Convert the running max from log2 units (scale_qk folds - # in LOG2E) to natural-log units: the shared reduce - # re-applies LOG2E itself when combining partitions. - pmax_ptr[base] = o_final[_m_slot(m)] * fx.Float32(1.0 / LOG2E) - psum_ptr[base] = l_row - - @flyc.jit - def pa_decode_tile_launch( - output: fx.Tensor, - pmax: fx.Tensor, - psum: fx.Tensor, - pout: fx.Tensor, - query: fx.Tensor, - key_cache: fx.Tensor, - value_cache: fx.Tensor, - block_tables: fx.Tensor, - context_lengths: fx.Tensor, - key_scale: fx.Tensor, # [1] per-tensor OR [num_blocks, num_kv_heads, block_size] per-token - value_scale: fx.Tensor, # same shape as key_scale - 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_q_row: fx.Int32, - stride_q_head: fx.Int32, - stream: fx.Stream = fx.Stream(None), - ): - pa_decode_tile_kernel( - output, - pmax, - psum, - pout, - query, - key_cache, - value_cache, - block_tables, - context_lengths, - key_scale, - value_scale, - max_blocks_per_seq, - stride_ks_block, - stride_ks_head, - stride_q_row, - stride_q_head, - ).launch( - grid=(num_seqs, num_kv_heads, NP), - block=(BLOCK_THREADS, 1, 1), - stream=stream, - ) - - return {"launch": pa_decode_tile_launch, "kernel": pa_decode_tile_kernel} +from .kernels.utils import cdiv def _is_current_stream_capturing() -> bool: diff --git a/op_tests/flydsl_tests/test_pa_decode_tile.py b/op_tests/flydsl_tests/test_pa_decode_tile.py index cb48aae54f1..5720a8decea 100644 --- a/op_tests/flydsl_tests/test_pa_decode_tile.py +++ b/op_tests/flydsl_tests/test_pa_decode_tile.py @@ -1,43 +1,51 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. -"""Correctness tests for the FlyDSL paged-attention Tile kernel.""" +"""Correctness and performance sweep for FlyDSL paged-attention Tile.""" -from __future__ import annotations +import argparse +import itertools -import pytest +import aiter +import pandas as pd import torch +from aiter import dtypes, per_tensor_quant +from aiter.jit.utils.chip_info import get_gfx +from aiter.test_common import benchmark, checkAllclose, run_perftest -from aiter import per_tensor_quant -from aiter.ops.flydsl import is_flydsl_available +torch.set_default_device("cuda") -pytest.importorskip("flydsl") +SUPPORTED_GFX = ["gfx942", "gfx950"] +KV_COMPUTE_BLOCK = 256 -if not torch.cuda.is_available(): - pytest.skip("ROCm is not available", allow_module_level=True) -if not is_flydsl_available(): - pytest.skip("FlyDSL is not available", allow_module_level=True) +# Pairwise coverage of the normal-accuracy axes in FlyDSL's PA regression test: +# batches {3, 81, 128}, Q/KV heads {(4,1), (8,1), (16,1)}, head dims +# {128, 256}, and contexts {1027, 8192}. Keep the original 257-token boundary +# case as well. Both supported block sizes are crossed with every case in main(). +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), +] -from aiter.ops.flydsl.pa_decode_tile import pa_decode_tile # noqa: E402 - - -def _get_arch() -> str: - return torch.cuda.get_device_properties(0).gcnArchName.lower().split(":")[0] - - -pytestmark = pytest.mark.skipif( - not (_get_arch().startswith("gfx942") or _get_arch().startswith("gfx95")), - reason="PA Tile requires gfx942 or gfx95x", -) +try: + from aiter.ops.flydsl.pa_decode_tile import ( + get_recommended_splits, + pa_decode_tile, + ) +except ImportError: + get_recommended_splits = None + pa_decode_tile = None def _quant_dtype() -> torch.dtype: - if _get_arch().startswith("gfx95"): - return torch.float8_e4m3fn - return torch.float8_e4m3fnuz + return torch.float8_e4m3fn if get_gfx() == "gfx950" else torch.float8_e4m3fnuz -def _reference_attention( +def run_torch( query: torch.Tensor, key_cache: torch.Tensor, value_cache: torch.Tensor, @@ -77,15 +85,55 @@ def _reference_attention( return output -@pytest.mark.parametrize("block_size", [16, 64]) -def test_pa_decode_tile_fp8(block_size: int) -> None: +def _run_flydsl( + output, + query, + key_cache, + value_cache, + block_tables, + context_lengths, + key_scale, + value_scale, + num_partitions, + pmax, + psum, + pout, +): + pa_decode_tile( + output, + query, + key_cache, + value_cache, + block_tables, + context_lengths, + key_scale, + value_scale, + num_partitions=num_partitions, + pmax=pmax, + psum=psum, + pout=pout, + ) + return output + + +@benchmark() +def test_pa_decode_tile( + batch_size, + num_query_heads, + num_kv_heads, + head_dim, + context_length, + block_size, + dtype, +): + if pa_decode_tile is None or get_recommended_splits is None: + raise RuntimeError("FlyDSL is not available") + if dtype not in (dtypes.fp16, dtypes.bf16): + raise ValueError(f"pa_decode_tile only supports fp16/bf16, got {dtype}") + if num_query_heads % num_kv_heads != 0: + raise ValueError("num_query_heads must be divisible by num_kv_heads") + torch.manual_seed(0) - device = torch.device("cuda") - batch_size = 3 - num_query_heads = 8 - num_kv_heads = 1 - head_dim = 128 - context_length = 257 blocks_per_sequence = (context_length + block_size - 1) // block_size num_blocks = batch_size * blocks_per_sequence @@ -93,24 +141,21 @@ def test_pa_decode_tile_fp8(block_size: int) -> None: batch_size, num_query_heads, head_dim, - dtype=torch.bfloat16, - device=device, + dtype=dtype, ).uniform_(-0.5, 0.5) key = torch.empty( num_blocks, num_kv_heads, block_size, head_dim, - dtype=torch.bfloat16, - device=device, + dtype=dtype, ).uniform_(-0.5, 0.5) value = torch.empty( num_blocks, num_kv_heads, head_dim, block_size, - dtype=torch.bfloat16, - device=device, + dtype=dtype, ).uniform_(-0.5, 0.5) quant_dtype = _quant_dtype() @@ -128,27 +173,13 @@ def test_pa_decode_tile_fp8(block_size: int) -> None: .contiguous() ) value_cache = value_quant.contiguous() - block_tables = torch.arange(num_blocks, dtype=torch.int32, device=device).reshape( + block_tables = torch.arange(num_blocks, dtype=torch.int32).reshape( batch_size, blocks_per_sequence ) - context_lengths = torch.full( - (batch_size,), context_length, dtype=torch.int32, device=device - ) + context_lengths = torch.full((batch_size,), context_length, dtype=torch.int32) output = torch.empty_like(query) - pa_decode_tile( - output, - query, - key_cache, - value_cache, - block_tables, - context_lengths, - key_scale, - value_scale, - ) - torch.cuda.synchronize() - - reference = _reference_attention( + reference = run_torch( query, key_quant, value_quant, @@ -157,4 +188,149 @@ def test_pa_decode_tile_fp8(block_size: int) -> None: key_scale, value_scale, ) - torch.testing.assert_close(output, reference, atol=3e-2, rtol=3e-2) + + query_group_size = num_query_heads // num_kv_heads + num_partitions = get_recommended_splits( + batch_size, + num_kv_heads, + split_kv_blocks=KV_COMPUTE_BLOCK // block_size, + ) + partial_shape = ( + batch_size, + num_kv_heads, + num_partitions, + query_group_size, + ) + pmax = torch.empty(partial_shape, dtype=dtypes.fp32) + psum = torch.empty_like(pmax) + pout = torch.empty(*partial_shape, head_dim, dtype=dtype) + + candidates = { + "flydsl": lambda: _run_flydsl( + output, + query, + key_cache, + value_cache, + block_tables, + context_lengths, + key_scale, + value_scale, + num_partitions, + pmax, + psum, + pout, + ) + } + + # QK and PV each perform one multiply-add per query-head/context pair. + flops = 4 * batch_size * num_query_heads * context_length * head_dim + # Logical tensor traffic: Q + O + the referenced K/V tokens and metadata. + nbytes = ( + 2 * query.numel() * query.element_size() + + 2 + * batch_size + * num_kv_heads + * context_length + * head_dim + * key_cache.element_size() + + block_tables.numel() * block_tables.element_size() + + context_lengths.numel() * context_lengths.element_size() + + key_scale.numel() * key_scale.element_size() + + value_scale.numel() * value_scale.element_size() + ) + + ret = {"gfx": get_gfx(), "partitions": num_partitions} + for name, fn in candidates.items(): + out, us = run_perftest(fn) + err = checkAllclose( + reference.to(dtypes.fp32), + out.to(dtypes.fp32), + rtol=3e-2, + atol=3e-2, + tol_err_ratio=0.0, + msg=f"{name}: pa_decode_tile", + ) + ret[f"{name} us"] = us + ret[f"{name} TFLOPS"] = flops / us / 1e6 + ret[f"{name} TB/s"] = nbytes / us / 1e6 + ret[f"{name} err"] = err + return ret + + +def main(): + if not torch.cuda.is_available(): + aiter.logger.warning("ROCm is not available; skipping pa_decode_tile") + return + if get_gfx() not in SUPPORTED_GFX: + aiter.logger.warning("pa_decode_tile unsupported on %s; skipping", get_gfx()) + return + if pa_decode_tile is None: + aiter.logger.warning("flydsl is unavailable; skipping pa_decode_tile") + return + + parser = argparse.ArgumentParser( + formatter_class=argparse.RawTextHelpFormatter, + description="FlyDSL pa_decode_tile correctness + perf sweep", + ) + parser.add_argument( + "-d", + "--dtype", + type=dtypes.str2Dtype, + nargs="*", + default=[dtypes.bf16], + help="""Query/output data type. + e.g.: -d bf16 fp16""", + ) + parser.add_argument( + "-b", + "--batch", + type=int, + nargs="*", + default=DEFAULT_BATCH_SIZES, + help="""Batch sizes. + e.g.: -b 1 3 16""", + ) + parser.add_argument( + "-s", + "--shapes", + type=dtypes.str2tuple, + nargs="*", + default=DEFAULT_SHAPES, + help="""(num_query_heads,num_kv_heads,head_dim,context_length). + e.g.: -s 8,1,128,257""", + ) + parser.add_argument( + "--block-size", + type=int, + nargs="*", + choices=[16, 64], + default=[16, 64], + help="""KV-cache block sizes.""", + ) + args = parser.parse_args() + + rows = [] + for dtype, batch_size, shape, block_size in itertools.product( + args.dtype, args.batch, args.shapes, args.block_size + ): + num_query_heads, num_kv_heads, head_dim, context_length = shape + rows.append( + test_pa_decode_tile( + batch_size, + num_query_heads, + num_kv_heads, + head_dim, + context_length, + block_size, + dtype, + ) + ) + + df = pd.DataFrame(rows) + aiter.logger.info( + "pa_decode_tile summary (markdown):\n%s", df.to_markdown(index=False) + ) + + +if __name__ == "__main__": + main() From ed3e0d6a62da17f35cc528595593f06b93ca939f Mon Sep 17 00:00:00 2001 From: fsx950223 Date: Thu, 23 Jul 2026 05:50:36 +0000 Subject: [PATCH 04/66] refactor(flydsl): Rename PA decode host API Signed-off-by: fsx950223 Co-authored-by: Cursor --- aiter/ops/flydsl/__init__.py | 4 +-- aiter/ops/flydsl/pa_decode.py | 16 +++++++++++ aiter/ops/flydsl/pa_decode_tile.py | 12 ++++----- op_tests/flydsl_tests/test_pa_decode_tile.py | 28 +++++++++----------- 4 files changed, 36 insertions(+), 24 deletions(-) create mode 100644 aiter/ops/flydsl/pa_decode.py diff --git a/aiter/ops/flydsl/__init__.py b/aiter/ops/flydsl/__init__.py index 5b4fbee1170..8a4b6137f5d 100644 --- a/aiter/ops/flydsl/__init__.py +++ b/aiter/ops/flydsl/__init__.py @@ -82,7 +82,7 @@ ".gather_kv_b_proj", "gather_kv_b_proj_flydsl_supported", ), - "pa_decode_tile": (".pa_decode_tile", "pa_decode_tile"), + "pa_decode": (".pa_decode", "pa_decode"), } __all__ = [ @@ -105,7 +105,7 @@ "flydsl_qk_norm_rope_quant", "gather_kv_b_proj_flydsl", "gather_kv_b_proj_flydsl_supported", - "pa_decode_tile", + "pa_decode", ] diff --git a/aiter/ops/flydsl/pa_decode.py b/aiter/ops/flydsl/pa_decode.py new file mode 100644 index 00000000000..5df9b68546b --- /dev/null +++ b/aiter/ops/flydsl/pa_decode.py @@ -0,0 +1,16 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (C) 2025-2026 FlyDSL Project Contributors + +"""Public host dispatch for FlyDSL paged-attention decode.""" + +from .pa_decode_tile import ( + compile_pa_decode_tile, + get_recommended_splits, + pa_decode, +) + +__all__ = [ + "compile_pa_decode_tile", + "get_recommended_splits", + "pa_decode", +] diff --git a/aiter/ops/flydsl/pa_decode_tile.py b/aiter/ops/flydsl/pa_decode_tile.py index e9d498cb0d5..423e629c12c 100644 --- a/aiter/ops/flydsl/pa_decode_tile.py +++ b/aiter/ops/flydsl/pa_decode_tile.py @@ -54,7 +54,7 @@ def get_recommended_splits( return max(4, min(n, 8)) -def pa_decode_tile( +def pa_decode( output: torch.Tensor, query: torch.Tensor, key_cache: torch.Tensor, @@ -89,7 +89,7 @@ def pa_decode_tile( assert block_size in ( 16, 64, - ), f"pa_decode_tile only supports block_size in (16, 64), got {block_size}" + ), f"pa_decode only supports block_size in (16, 64), got {block_size}" trans_v = value_cache.dim() == 5 if trans_v: @@ -117,16 +117,14 @@ def pa_decode_tile( elif query.dtype == torch.float16: query_dtype = "f16" else: - raise ValueError( - f"pa_decode_tile only supports f16/bf16 query, got {query.dtype}" - ) + raise ValueError(f"pa_decode only supports f16/bf16 query, got {query.dtype}") assert ( output.dtype == query.dtype - ), f"pa_decode_tile requires output.dtype == query.dtype, got {output.dtype} vs {query.dtype}" + ), f"pa_decode requires output.dtype == query.dtype, got {output.dtype} vs {query.dtype}" assert ( query.stride(2) == 1 - ), f"pa_decode_tile requires a contiguous head_dim axis, got strides {query.stride()}" + ), f"pa_decode requires a contiguous head_dim axis, got strides {query.stride()}" dev = query.device key_scale_t = ( diff --git a/op_tests/flydsl_tests/test_pa_decode_tile.py b/op_tests/flydsl_tests/test_pa_decode_tile.py index 5720a8decea..05a8d8bfca3 100644 --- a/op_tests/flydsl_tests/test_pa_decode_tile.py +++ b/op_tests/flydsl_tests/test_pa_decode_tile.py @@ -32,13 +32,13 @@ ] try: - from aiter.ops.flydsl.pa_decode_tile import ( + from aiter.ops.flydsl.pa_decode import ( get_recommended_splits, - pa_decode_tile, + pa_decode, ) except ImportError: get_recommended_splits = None - pa_decode_tile = None + pa_decode = None def _quant_dtype() -> torch.dtype: @@ -99,7 +99,7 @@ def _run_flydsl( psum, pout, ): - pa_decode_tile( + pa_decode( output, query, key_cache, @@ -126,10 +126,10 @@ def test_pa_decode_tile( block_size, dtype, ): - if pa_decode_tile is None or get_recommended_splits is None: + if pa_decode is None or get_recommended_splits is None: raise RuntimeError("FlyDSL is not available") if dtype not in (dtypes.fp16, dtypes.bf16): - raise ValueError(f"pa_decode_tile only supports fp16/bf16, got {dtype}") + raise ValueError(f"pa_decode only supports fp16/bf16, got {dtype}") if num_query_heads % num_kv_heads != 0: raise ValueError("num_query_heads must be divisible by num_kv_heads") @@ -248,7 +248,7 @@ def test_pa_decode_tile( rtol=3e-2, atol=3e-2, tol_err_ratio=0.0, - msg=f"{name}: pa_decode_tile", + msg=f"{name}: pa_decode", ) ret[f"{name} us"] = us ret[f"{name} TFLOPS"] = flops / us / 1e6 @@ -259,18 +259,18 @@ def test_pa_decode_tile( def main(): if not torch.cuda.is_available(): - aiter.logger.warning("ROCm is not available; skipping pa_decode_tile") + aiter.logger.warning("ROCm is not available; skipping pa_decode") return if get_gfx() not in SUPPORTED_GFX: - aiter.logger.warning("pa_decode_tile unsupported on %s; skipping", get_gfx()) + aiter.logger.warning("pa_decode unsupported on %s; skipping", get_gfx()) return - if pa_decode_tile is None: - aiter.logger.warning("flydsl is unavailable; skipping pa_decode_tile") + if pa_decode is None: + aiter.logger.warning("flydsl is unavailable; skipping pa_decode") return parser = argparse.ArgumentParser( formatter_class=argparse.RawTextHelpFormatter, - description="FlyDSL pa_decode_tile correctness + perf sweep", + description="FlyDSL pa_decode correctness + perf sweep", ) parser.add_argument( "-d", @@ -327,9 +327,7 @@ def main(): ) df = pd.DataFrame(rows) - aiter.logger.info( - "pa_decode_tile summary (markdown):\n%s", df.to_markdown(index=False) - ) + aiter.logger.info("pa_decode summary (markdown):\n%s", df.to_markdown(index=False)) if __name__ == "__main__": From da010a7c90d810873c4072b64762491cab86bedc Mon Sep 17 00:00:00 2001 From: fsx950223 Date: Thu, 23 Jul 2026 05:53:59 +0000 Subject: [PATCH 05/66] refactor(flydsl): Merge PA decode host modules Signed-off-by: fsx950223 Co-authored-by: Cursor --- aiter/ops/flydsl/pa_decode.py | 281 +++++++++++++++++++++++++++-- aiter/ops/flydsl/pa_decode_tile.py | 271 ---------------------------- 2 files changed, 268 insertions(+), 284 deletions(-) delete mode 100644 aiter/ops/flydsl/pa_decode_tile.py diff --git a/aiter/ops/flydsl/pa_decode.py b/aiter/ops/flydsl/pa_decode.py index 5df9b68546b..423e629c12c 100644 --- a/aiter/ops/flydsl/pa_decode.py +++ b/aiter/ops/flydsl/pa_decode.py @@ -1,16 +1,271 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (C) 2025-2026 FlyDSL Project Contributors -"""Public host dispatch for FlyDSL paged-attention decode.""" - -from .pa_decode_tile import ( - compile_pa_decode_tile, - get_recommended_splits, - pa_decode, -) - -__all__ = [ - "compile_pa_decode_tile", - "get_recommended_splits", - "pa_decode", -] +"""Readable tile-programming reference for paged-attention fp8 decode. + +K/V are fp8 e4m3 (FNUZ on gfx942, OCP on gfx950) fed straight into +``mfma_f32_16x16x32_fp8_fp8``; Q (bf16/f16) and the softmax probabilities P are +quantized to fp8 too. Scales fold out of the matmuls (q/key scale into the QK +score, value scale + 1/FP8_MAX into the epilogue); softmax max/sum stay f32. +``key_scale``/``value_scale`` are either a ``[1]`` per-tensor scalar or a +``[num_blocks, num_kv_heads, block_size]`` per-token tensor (chosen by rank). + +``block_size`` (16/64) and ``head_dim`` (multiple of 64) are compile-time +constants. Layouts are logical, not production's preshuffle. + +* ``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 + +One CTA (4 waves) per (seq, kv_head) runs a flash-style online softmax over +256-token blocks; the 4 waves split tokens for Q.KT and head-dim for P.V, with +an LDS round-trip on P transposing ownership between the two MMAs. +""" + +import torch + +from .kernels.pa_decode_tile import KV_COMPUTE_BLOCK, compile_pa_decode_tile +from .kernels.tensor_shim import _run_compiled +from .kernels.utils import cdiv + + +def _is_current_stream_capturing() -> bool: + if not torch.cuda.is_available(): + return False + try: + return torch.cuda.is_current_stream_capturing() + except RuntimeError: + return False + + +def get_recommended_splits( + num_sequences: int, + num_kv_heads: int, + split_kv_blocks: int = 1, +) -> int: + 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 = cdiv(num_sm, denom) * split_kv_blocks + return max(4, min(n, 8)) + + +def pa_decode( + output: torch.Tensor, + query: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + block_tables: torch.Tensor, + context_lengths: torch.Tensor, + key_scale: float | torch.Tensor, + value_scale: float | torch.Tensor, + softmax_scale: float | None = None, + stream=None, + *, + num_partitions: int | None = None, + pmax: torch.Tensor | None = None, + psum: torch.Tensor | None = None, + pout: torch.Tensor | None = None, +) -> None: + """Host entry point. See module docstring for tensor layouts. + + ``num_partitions``/``pmax``/``psum``/``pout`` are optional caller overrides + (e.g. for CUDA-graph capture, where nothing may be allocated on-the-fly); + when omitted they are picked/allocated here. + """ + num_seqs = context_lengths.shape[0] + total_q_rows, num_q_heads, head_dim = query.shape + assert ( + total_q_rows % num_seqs == 0 + ), f"query.shape[0] ({total_q_rows}) must be a multiple of context_lengths.shape[0] ({num_seqs})" + query_length = total_q_rows // num_seqs + _, num_kv_heads, num_hgroups, block_size, hgroup_width = key_cache.shape + + assert num_hgroups == head_dim // 16 and hgroup_width == 16 + assert block_size in ( + 16, + 64, + ), f"pa_decode only supports block_size in (16, 64), got {block_size}" + + trans_v = value_cache.dim() == 5 + if trans_v: + _, v_num_kv_heads, v_subblocks, v_head_dim, v_width = value_cache.shape + assert ( + v_head_dim == head_dim and v_width == 16 and v_subblocks == block_size // 16 + ), f"value_cache shape {tuple(value_cache.shape)} doesn't match block_size={block_size}, head_dim={head_dim}" + else: + _, v_num_kv_heads, v_head_dim, v_block_size = value_cache.shape + assert v_head_dim == head_dim and v_block_size == block_size, ( + f"value_cache shape {tuple(value_cache.shape)} doesn't match " + f"block_size={block_size}, head_dim={head_dim}" + ) + assert v_num_kv_heads == num_kv_heads + assert ( + block_tables.dtype == torch.int32 + ), f"block_tables must be int32, got {block_tables.dtype}" + assert ( + context_lengths.dtype == torch.int32 + ), 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 ValueError(f"pa_decode only supports f16/bf16 query, got {query.dtype}") + assert ( + output.dtype == query.dtype + ), f"pa_decode requires output.dtype == query.dtype, got {output.dtype} vs {query.dtype}" + + assert ( + query.stride(2) == 1 + ), f"pa_decode requires a contiguous head_dim axis, got strides {query.stride()}" + + dev = query.device + key_scale_t = ( + key_scale + if isinstance(key_scale, torch.Tensor) + else torch.tensor([float(key_scale)], device=dev) + ) + value_scale_t = ( + value_scale + if isinstance(value_scale, torch.Tensor) + else torch.tensor([float(value_scale)], device=dev) + ) + per_token_kv = key_scale_t.dim() > 1 + if per_token_kv: + assert ( + value_scale_t.dim() > 1 + ), "value_scale must also be per-token (dim>1) when key_scale is per-token" + assert ( + key_scale_t.shape == value_scale_t.shape + ), f"key_scale/value_scale shape mismatch: {tuple(key_scale_t.shape)} vs {tuple(value_scale_t.shape)}" + assert key_scale_t.shape == (key_cache.shape[0], num_kv_heads, block_size), ( + "per-token key_scale/value_scale must be [num_blocks, num_kv_heads, block_size] " + f"matching the KV cache, got {tuple(key_scale_t.shape)}" + ) + stride_ks_block = int(key_scale_t.stride(0)) + stride_ks_head = int(key_scale_t.stride(1)) + else: + stride_ks_block = 0 + stride_ks_head = 0 + assert ( + key_scale_t.dtype == torch.float32 and key_scale_t.device == dev + ), f"key_scale tensor must be float32 on {dev}, got {key_scale_t.dtype} on {key_scale_t.device}" + assert ( + value_scale_t.dtype == torch.float32 and value_scale_t.device == dev + ), f"value_scale tensor must be float32 on {dev}, got {value_scale_t.dtype} on {value_scale_t.device}" + + if not num_partitions: + blocks_per_partition = KV_COMPUTE_BLOCK // block_size + num_partitions = get_recommended_splits( + num_seqs, + num_kv_heads, + split_kv_blocks=blocks_per_partition, + ) + + compiled = compile_pa_decode_tile( + head_dim=head_dim, + query_group_size=query_group_size, + block_size=int(block_size), + 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, + ) + + is_graph_capturing = _is_current_stream_capturing() + if num_partitions == 1: + # NP==1 writes output directly; partials unused (caller buffers ignored). + if pmax is None: + if is_graph_capturing: + raise ValueError( + "CUDA graph capture requires preallocated `pmax`/`psum`/`pout` " + "even when num_partitions==1 (nothing may be allocated mid-capture)." + ) + pmax = psum = pout = torch.empty(1, dtype=torch.float32, device=dev) + else: + total_rows = query_length * query_group_size + expected_scalar_shape = (num_seqs, num_kv_heads, num_partitions, total_rows) + if pmax is None or psum is None or pout is None: + if is_graph_capturing: + raise ValueError( + "CUDA graph capture requires preallocated `pmax`/`psum`/`pout` " + "for num_partitions>1 (nothing may be allocated mid-capture)." + ) + pmax = torch.empty(*expected_scalar_shape, dtype=torch.float32, device=dev) + psum = torch.empty(*expected_scalar_shape, dtype=torch.float32, device=dev) + pout = torch.empty( + *expected_scalar_shape, head_dim, dtype=output.dtype, device=dev + ) + else: + assert ( + pmax.shape == expected_scalar_shape + ), f"pmax shape {tuple(pmax.shape)} != {expected_scalar_shape}" + assert ( + psum.shape == expected_scalar_shape + ), f"psum shape {tuple(psum.shape)} != {expected_scalar_shape}" + assert pout.shape == ( + *expected_scalar_shape, + head_dim, + ), f"pout shape {tuple(pout.shape)} != {(*expected_scalar_shape, head_dim)}" + s = stream or torch.cuda.current_stream() + + _run_compiled( + compiled["launch"], + output, + pmax.view(-1), + psum.view(-1), + pout.view(-1), + query, + key_cache, + value_cache, + block_tables, + context_lengths, + key_scale_t, + value_scale_t, + int(max_blocks_per_seq), + int(num_seqs), + int(num_kv_heads), + stride_ks_block, + stride_ks_head, + int(query.stride(0)), + int(query.stride(1)), + s, + ) + if num_partitions > 1: + from csrc.cpp_itfs.pa.pa_ps import launch_pa_decode_ps_reduce + + output_5d = output.reshape( + num_seqs, query_length, num_kv_heads, query_group_size, head_dim + ) + with torch.cuda.stream(s): + launch_pa_decode_ps_reduce( + output_5d, + psum, + pmax, + pout, + None, + output_5d.stride(0), + output_5d.stride(1), + output_5d.stride(2), + output_5d.stride(3), + pmax.stride(0), + pmax.stride(1), + pmax.stride(2), + pout.stride(0), + pout.stride(1), + pout.stride(2), + pout.stride(3), + query_seq_len=query_length, + query_group_size=query_group_size, + head_size=head_dim, + context_partition_num=num_partitions, + ) diff --git a/aiter/ops/flydsl/pa_decode_tile.py b/aiter/ops/flydsl/pa_decode_tile.py deleted file mode 100644 index 423e629c12c..00000000000 --- a/aiter/ops/flydsl/pa_decode_tile.py +++ /dev/null @@ -1,271 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# Copyright (C) 2025-2026 FlyDSL Project Contributors - -"""Readable tile-programming reference for paged-attention fp8 decode. - -K/V are fp8 e4m3 (FNUZ on gfx942, OCP on gfx950) fed straight into -``mfma_f32_16x16x32_fp8_fp8``; Q (bf16/f16) and the softmax probabilities P are -quantized to fp8 too. Scales fold out of the matmuls (q/key scale into the QK -score, value scale + 1/FP8_MAX into the epilogue); softmax max/sum stay f32. -``key_scale``/``value_scale`` are either a ``[1]`` per-tensor scalar or a -``[num_blocks, num_kv_heads, block_size]`` per-token tensor (chosen by rank). - -``block_size`` (16/64) and ``head_dim`` (multiple of 64) are compile-time -constants. Layouts are logical, not production's preshuffle. - -* ``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 - -One CTA (4 waves) per (seq, kv_head) runs a flash-style online softmax over -256-token blocks; the 4 waves split tokens for Q.KT and head-dim for P.V, with -an LDS round-trip on P transposing ownership between the two MMAs. -""" - -import torch - -from .kernels.pa_decode_tile import KV_COMPUTE_BLOCK, compile_pa_decode_tile -from .kernels.tensor_shim import _run_compiled -from .kernels.utils import cdiv - - -def _is_current_stream_capturing() -> bool: - if not torch.cuda.is_available(): - return False - try: - return torch.cuda.is_current_stream_capturing() - except RuntimeError: - return False - - -def get_recommended_splits( - num_sequences: int, - num_kv_heads: int, - split_kv_blocks: int = 1, -) -> int: - 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 = cdiv(num_sm, denom) * split_kv_blocks - return max(4, min(n, 8)) - - -def pa_decode( - output: torch.Tensor, - query: torch.Tensor, - key_cache: torch.Tensor, - value_cache: torch.Tensor, - block_tables: torch.Tensor, - context_lengths: torch.Tensor, - key_scale: float | torch.Tensor, - value_scale: float | torch.Tensor, - softmax_scale: float | None = None, - stream=None, - *, - num_partitions: int | None = None, - pmax: torch.Tensor | None = None, - psum: torch.Tensor | None = None, - pout: torch.Tensor | None = None, -) -> None: - """Host entry point. See module docstring for tensor layouts. - - ``num_partitions``/``pmax``/``psum``/``pout`` are optional caller overrides - (e.g. for CUDA-graph capture, where nothing may be allocated on-the-fly); - when omitted they are picked/allocated here. - """ - num_seqs = context_lengths.shape[0] - total_q_rows, num_q_heads, head_dim = query.shape - assert ( - total_q_rows % num_seqs == 0 - ), f"query.shape[0] ({total_q_rows}) must be a multiple of context_lengths.shape[0] ({num_seqs})" - query_length = total_q_rows // num_seqs - _, num_kv_heads, num_hgroups, block_size, hgroup_width = key_cache.shape - - assert num_hgroups == head_dim // 16 and hgroup_width == 16 - assert block_size in ( - 16, - 64, - ), f"pa_decode only supports block_size in (16, 64), got {block_size}" - - trans_v = value_cache.dim() == 5 - if trans_v: - _, v_num_kv_heads, v_subblocks, v_head_dim, v_width = value_cache.shape - assert ( - v_head_dim == head_dim and v_width == 16 and v_subblocks == block_size // 16 - ), f"value_cache shape {tuple(value_cache.shape)} doesn't match block_size={block_size}, head_dim={head_dim}" - else: - _, v_num_kv_heads, v_head_dim, v_block_size = value_cache.shape - assert v_head_dim == head_dim and v_block_size == block_size, ( - f"value_cache shape {tuple(value_cache.shape)} doesn't match " - f"block_size={block_size}, head_dim={head_dim}" - ) - assert v_num_kv_heads == num_kv_heads - assert ( - block_tables.dtype == torch.int32 - ), f"block_tables must be int32, got {block_tables.dtype}" - assert ( - context_lengths.dtype == torch.int32 - ), 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 ValueError(f"pa_decode only supports f16/bf16 query, got {query.dtype}") - assert ( - output.dtype == query.dtype - ), f"pa_decode requires output.dtype == query.dtype, got {output.dtype} vs {query.dtype}" - - assert ( - query.stride(2) == 1 - ), f"pa_decode requires a contiguous head_dim axis, got strides {query.stride()}" - - dev = query.device - key_scale_t = ( - key_scale - if isinstance(key_scale, torch.Tensor) - else torch.tensor([float(key_scale)], device=dev) - ) - value_scale_t = ( - value_scale - if isinstance(value_scale, torch.Tensor) - else torch.tensor([float(value_scale)], device=dev) - ) - per_token_kv = key_scale_t.dim() > 1 - if per_token_kv: - assert ( - value_scale_t.dim() > 1 - ), "value_scale must also be per-token (dim>1) when key_scale is per-token" - assert ( - key_scale_t.shape == value_scale_t.shape - ), f"key_scale/value_scale shape mismatch: {tuple(key_scale_t.shape)} vs {tuple(value_scale_t.shape)}" - assert key_scale_t.shape == (key_cache.shape[0], num_kv_heads, block_size), ( - "per-token key_scale/value_scale must be [num_blocks, num_kv_heads, block_size] " - f"matching the KV cache, got {tuple(key_scale_t.shape)}" - ) - stride_ks_block = int(key_scale_t.stride(0)) - stride_ks_head = int(key_scale_t.stride(1)) - else: - stride_ks_block = 0 - stride_ks_head = 0 - assert ( - key_scale_t.dtype == torch.float32 and key_scale_t.device == dev - ), f"key_scale tensor must be float32 on {dev}, got {key_scale_t.dtype} on {key_scale_t.device}" - assert ( - value_scale_t.dtype == torch.float32 and value_scale_t.device == dev - ), f"value_scale tensor must be float32 on {dev}, got {value_scale_t.dtype} on {value_scale_t.device}" - - if not num_partitions: - blocks_per_partition = KV_COMPUTE_BLOCK // block_size - num_partitions = get_recommended_splits( - num_seqs, - num_kv_heads, - split_kv_blocks=blocks_per_partition, - ) - - compiled = compile_pa_decode_tile( - head_dim=head_dim, - query_group_size=query_group_size, - block_size=int(block_size), - 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, - ) - - is_graph_capturing = _is_current_stream_capturing() - if num_partitions == 1: - # NP==1 writes output directly; partials unused (caller buffers ignored). - if pmax is None: - if is_graph_capturing: - raise ValueError( - "CUDA graph capture requires preallocated `pmax`/`psum`/`pout` " - "even when num_partitions==1 (nothing may be allocated mid-capture)." - ) - pmax = psum = pout = torch.empty(1, dtype=torch.float32, device=dev) - else: - total_rows = query_length * query_group_size - expected_scalar_shape = (num_seqs, num_kv_heads, num_partitions, total_rows) - if pmax is None or psum is None or pout is None: - if is_graph_capturing: - raise ValueError( - "CUDA graph capture requires preallocated `pmax`/`psum`/`pout` " - "for num_partitions>1 (nothing may be allocated mid-capture)." - ) - pmax = torch.empty(*expected_scalar_shape, dtype=torch.float32, device=dev) - psum = torch.empty(*expected_scalar_shape, dtype=torch.float32, device=dev) - pout = torch.empty( - *expected_scalar_shape, head_dim, dtype=output.dtype, device=dev - ) - else: - assert ( - pmax.shape == expected_scalar_shape - ), f"pmax shape {tuple(pmax.shape)} != {expected_scalar_shape}" - assert ( - psum.shape == expected_scalar_shape - ), f"psum shape {tuple(psum.shape)} != {expected_scalar_shape}" - assert pout.shape == ( - *expected_scalar_shape, - head_dim, - ), f"pout shape {tuple(pout.shape)} != {(*expected_scalar_shape, head_dim)}" - s = stream or torch.cuda.current_stream() - - _run_compiled( - compiled["launch"], - output, - pmax.view(-1), - psum.view(-1), - pout.view(-1), - query, - key_cache, - value_cache, - block_tables, - context_lengths, - key_scale_t, - value_scale_t, - int(max_blocks_per_seq), - int(num_seqs), - int(num_kv_heads), - stride_ks_block, - stride_ks_head, - int(query.stride(0)), - int(query.stride(1)), - s, - ) - if num_partitions > 1: - from csrc.cpp_itfs.pa.pa_ps import launch_pa_decode_ps_reduce - - output_5d = output.reshape( - num_seqs, query_length, num_kv_heads, query_group_size, head_dim - ) - with torch.cuda.stream(s): - launch_pa_decode_ps_reduce( - output_5d, - psum, - pmax, - pout, - None, - output_5d.stride(0), - output_5d.stride(1), - output_5d.stride(2), - output_5d.stride(3), - pmax.stride(0), - pmax.stride(1), - pmax.stride(2), - pout.stride(0), - pout.stride(1), - pout.stride(2), - pout.stride(3), - query_seq_len=query_length, - query_group_size=query_group_size, - head_size=head_dim, - context_partition_num=num_partitions, - ) From 430b8cda40f39131056d507751031fd4610b5d68 Mon Sep 17 00:00:00 2001 From: fsx950223 Date: Fri, 24 Jul 2026 09:27:44 +0000 Subject: [PATCH 06/66] feat(flydsl): Add Gluon-compatible PA decode dispatch Signed-off-by: fsx950223 Co-authored-by: Cursor --- aiter/ops/attention.py | 179 ++++++++- aiter/ops/flydsl/kernels/pa_decode_tile.py | 7 +- aiter/ops/flydsl/pa_decode.py | 390 ++++++++++++------- op_tests/flydsl_tests/test_pa_decode_tile.py | 235 ++++++++++- 4 files changed, 649 insertions(+), 162 deletions(-) diff --git a/aiter/ops/attention.py b/aiter/ops/attention.py index ce2085fb35f..1e7a976637b 100644 --- a/aiter/ops/attention.py +++ b/aiter/ops/attention.py @@ -9,7 +9,9 @@ from aiter import dtypes from aiter.ops.enum import Enum, MlaVersion, QuantType -from aiter.ops.triton.gluon.pa_decode_gluon import pa_decode_gluon +from aiter.ops.triton.gluon.pa_decode_gluon import ( + pa_decode_gluon as _pa_decode_gluon_fallback, +) from aiter.utility.dtypes import _aiter_dtype_id from csrc.cpp_itfs.pa.pa import paged_attention_rocm as paged_attention_rocm_core from csrc.cpp_itfs.pa.pa_ragged import ( @@ -19,10 +21,183 @@ from csrc.cpp_itfs.torch_utils import direct_register_custom_op from ..jit.core import compile_ops, is_experimental_enabled -from ..jit.utils.chip_info import get_cu_num, get_gfx +from ..jit.utils.chip_info import get_cu_num, get_gfx, get_gfx_runtime + +try: + from aiter.ops.flydsl import pa_decode as _pa_decode_flydsl +except (ImportError, AttributeError, RuntimeError, OSError): + _pa_decode_flydsl = None MD_NAME = "module_attention" + +def _can_use_flydsl_pa_decode( + output, + query, + key_cache, + value_cache, + context_lengths, + block_tables, + query_length, + max_context_partition_num, + context_partition_size, + compute_type, + query_scale, + key_scale, + value_scale, + exp_sums, + max_logits, + temporary_output, + alibi_slopes, + sinks, + sliding_window, +): + runtime_gfx = get_gfx_runtime() + runtime_fp8_dtype = { + "gfx942": torch.float8_e4m3fnuz, + "gfx950": torch.float8_e4m3fn, + }.get(runtime_gfx) + + def scratch_buffers_match(): + if max_context_partition_num == 1: + return True + query_group_size = query.shape[1] // key_cache.shape[1] + total_rows = query_length * query_group_size + scalar_shape = ( + context_lengths.shape[0], + key_cache.shape[1], + max_context_partition_num, + total_rows, + ) + output_shape = (*scalar_shape, query.shape[-1]) + scalar_buffers_match = all( + tensor is None + or ( + tensor.shape == scalar_shape + and tensor.dtype == torch.float32 + and tensor.device == query.device + and tensor.is_contiguous() + ) + for tensor in (exp_sums, max_logits) + ) + return scalar_buffers_match and ( + temporary_output is None + or ( + temporary_output.shape == output_shape + and temporary_output.dtype == output.dtype + and temporary_output.device == query.device + and temporary_output.is_contiguous() + ) + ) + + return ( + _pa_decode_flydsl is not None + and runtime_fp8_dtype is not None + and query.dim() == 3 + and output.shape == query.shape + and output.dtype == query.dtype + and query.dtype in (torch.bfloat16, torch.float16) + and query.stride(-1) == 1 + and key_cache.dim() == 5 + and key_cache.dtype == runtime_fp8_dtype + and key_cache.shape[-2] in (16, 64) + and key_cache.is_contiguous() + and value_cache.dim() in (4, 5) + and value_cache.dtype == runtime_fp8_dtype + and value_cache.is_contiguous() + and query.shape[-1] % 64 == 0 + and key_cache.shape[1] > 0 + and query.shape[1] % key_cache.shape[1] == 0 + and context_lengths.dim() == 1 + and context_lengths.is_contiguous() + and block_tables.dim() == 2 + and block_tables.is_contiguous() + and query.shape[0] == context_lengths.shape[0] * query_length + and 1 <= max_context_partition_num <= 64 + and context_partition_size == 256 + and compute_type == runtime_fp8_dtype + and query_scale is None + and (key_scale is None) == (value_scale is None) + and scratch_buffers_match() + and alibi_slopes is None + and sinks is None + and sliding_window == 0 + ) + + +def pa_decode_gluon( + output: torch.Tensor, + query: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + context_lengths: torch.Tensor, + block_tables: torch.Tensor, + 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, + value_scale: torch.Tensor = None, + exp_sums: torch.Tensor = None, + max_logits: torch.Tensor = None, + temporary_output: torch.Tensor = None, + alibi_slopes: torch.Tensor = None, + sinks: torch.Tensor = None, + sliding_window: int = 0, + ps: bool = True, +) -> None: + implementation = ( + _pa_decode_flydsl + if _can_use_flydsl_pa_decode( + output, + query, + key_cache, + value_cache, + context_lengths, + block_tables, + query_length, + max_context_partition_num, + context_partition_size, + compute_type, + query_scale, + key_scale, + value_scale, + exp_sums, + max_logits, + temporary_output, + alibi_slopes, + sinks, + sliding_window, + ) + else _pa_decode_gluon_fallback + ) + implementation( + output, + query, + key_cache, + value_cache, + context_lengths, + block_tables, + softmax_scale, + query_length, + max_context_partition_num, + context_partition_size, + compute_type, + query_scale, + key_scale, + value_scale, + exp_sums, + max_logits, + temporary_output, + alibi_slopes, + sinks, + sliding_window, + ps, + ) + + direct_register_custom_op( "pa_decode_gluon", pa_decode_gluon, diff --git a/aiter/ops/flydsl/kernels/pa_decode_tile.py b/aiter/ops/flydsl/kernels/pa_decode_tile.py index de8142e99e9..e2cb40a279a 100644 --- a/aiter/ops/flydsl/kernels/pa_decode_tile.py +++ b/aiter/ops/flydsl/kernels/pa_decode_tile.py @@ -61,16 +61,19 @@ def compile_pa_decode_tile( per_token_kv: bool = False, query_length: int = 1, trans_v: bool = True, + device_index: int | None = None, ): """Build the tile-programming PA-decode kernel + launch wrapper. ``block_size``, ``head_dim``, and ``query_dtype`` are compile-time - constants (lru_cache keys). ``query_length`` (MTP) and ``query_group_size`` - flatten into ``TOTAL_ROWS = query_length * query_group_size``, tiled into + constants. ``device_index`` keeps the cached launch/compiled function + device-local. ``query_length`` (MTP) and ``query_group_size`` flatten into + ``TOTAL_ROWS = query_length * query_group_size``, tiled into ``M_TILES = ceil(TOTAL_ROWS / 16)`` independent 16-row MFMA tiles; each extra M-tile duplicates loop-carried state, so VGPR/LDS/occupancy scale roughly linearly with ``M_TILES``. """ + del device_index is_gfx950 = "gfx95" in get_rocm_arch() FP8 = fx.Float8E4M3FN if is_gfx950 else fx.Float8E4M3FNUZ FP8_MAX = ( diff --git a/aiter/ops/flydsl/pa_decode.py b/aiter/ops/flydsl/pa_decode.py index 423e629c12c..e15047acc34 100644 --- a/aiter/ops/flydsl/pa_decode.py +++ b/aiter/ops/flydsl/pa_decode.py @@ -8,7 +8,7 @@ quantized to fp8 too. Scales fold out of the matmuls (q/key scale into the QK score, value scale + 1/FP8_MAX into the epilogue); softmax max/sum stay f32. ``key_scale``/``value_scale`` are either a ``[1]`` per-tensor scalar or a -``[num_blocks, num_kv_heads, block_size]`` per-token tensor (chosen by rank). +``[num_blocks, num_kv_heads, block_size, 1]`` per-token tensor. ``block_size`` (16/64) and ``head_dim`` (multiple of 64) are compile-time constants. Layouts are logical, not production's preshuffle. @@ -28,20 +28,13 @@ import torch +from aiter.jit.utils.chip_info import get_gfx_runtime + from .kernels.pa_decode_tile import KV_COMPUTE_BLOCK, compile_pa_decode_tile from .kernels.tensor_shim import _run_compiled from .kernels.utils import cdiv -def _is_current_stream_capturing() -> bool: - if not torch.cuda.is_available(): - return False - try: - return torch.cuda.is_current_stream_capturing() - except RuntimeError: - return False - - def get_recommended_splits( num_sequences: int, num_kv_heads: int, @@ -55,34 +48,88 @@ def get_recommended_splits( def pa_decode( - output: torch.Tensor, - query: torch.Tensor, - key_cache: torch.Tensor, - value_cache: torch.Tensor, - block_tables: torch.Tensor, - context_lengths: torch.Tensor, - key_scale: float | torch.Tensor, - value_scale: float | torch.Tensor, - softmax_scale: float | None = None, - stream=None, - *, - num_partitions: int | None = None, - pmax: torch.Tensor | None = None, - psum: torch.Tensor | None = None, - pout: torch.Tensor | None = None, + 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, # [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] + alibi_slopes: torch.Tensor = None, + sinks: torch.Tensor = None, + sliding_window: int = 0, + ps: bool = True, ) -> None: - """Host entry point. See module docstring for tensor layouts. + """FlyDSL replacement for ``pa_decode_gluon``. - ``num_partitions``/``pmax``/``psum``/``pout`` are optional caller overrides - (e.g. for CUDA-graph capture, where nothing may be allocated on-the-fly); - when omitted they are picked/allocated here. + The call signature and intermediate-buffer layouts match + ``pa_decode_gluon``. This kernel currently supports FP8 K/V caches, + BF16/FP16 queries, a 256-token context partition, and block sizes 16/64. + ALiBi, attention sinks, sliding-window attention, and externally quantized + FP8 queries are not supported. """ + 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 sinks is not None: + raise NotImplementedError("pa_decode does not support attention sinks") + if sliding_window != 0: + raise NotImplementedError("pa_decode does not support sliding-window attention") + if query_length < 1: + raise ValueError(f"query_length must be positive, got {query_length}") + if not 1 <= max_context_partition_num <= 64: + raise ValueError( + "max_context_partition_num must be in [1, 64], " + f"got {max_context_partition_num}" + ) + + 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}" + ) + + # ``ps`` is retained for drop-in API compatibility. Both partitioning + # policies are represented by the caller-provided partition count. + del ps + num_seqs = context_lengths.shape[0] total_q_rows, num_q_heads, head_dim = query.shape - assert ( - total_q_rows % num_seqs == 0 - ), f"query.shape[0] ({total_q_rows}) must be a multiple of context_lengths.shape[0] ({num_seqs})" - query_length = total_q_rows // num_seqs + assert total_q_rows == num_seqs * query_length, ( + f"query.shape[0] ({total_q_rows}) must equal " + f"context_lengths.shape[0] * query_length ({num_seqs} * {query_length})" + ) + assert output.shape == query.shape, ( + f"output shape {tuple(output.shape)} must match " + f"query shape {tuple(query.shape)}" + ) _, num_kv_heads, num_hgroups, block_size, hgroup_width = key_cache.shape assert num_hgroups == head_dim // 16 and hgroup_width == 16 @@ -104,6 +151,9 @@ def pa_decode( f"block_size={block_size}, head_dim={head_dim}" ) assert v_num_kv_heads == num_kv_heads + assert ( + num_q_heads % num_kv_heads == 0 + ), f"num_q_heads ({num_q_heads}) must be divisible by num_kv_heads ({num_kv_heads})" assert ( block_tables.dtype == torch.int32 ), f"block_tables must be int32, got {block_tables.dtype}" @@ -122,26 +172,67 @@ def pa_decode( output.dtype == query.dtype ), f"pa_decode requires output.dtype == query.dtype, got {output.dtype} vs {query.dtype}" + assert ( + key_cache.dtype == expected_fp8_dtype + ), f"pa_decode requires {expected_fp8_dtype} key cache on {arch}, got {key_cache.dtype}" + assert ( + value_cache.dtype == expected_fp8_dtype + ), f"pa_decode requires {expected_fp8_dtype} value cache on {arch}, got {value_cache.dtype}" + assert ( query.stride(2) == 1 ), f"pa_decode requires a contiguous head_dim axis, got strides {query.stride()}" dev = query.device - key_scale_t = ( - key_scale - if isinstance(key_scale, torch.Tensor) - else torch.tensor([float(key_scale)], device=dev) - ) - value_scale_t = ( - value_scale - if isinstance(value_scale, torch.Tensor) - else torch.tensor([float(value_scale)], device=dev) - ) - per_token_kv = key_scale_t.dim() > 1 + for name, tensor in ( + ("output", output), + ("key_cache", key_cache), + ("value_cache", value_cache), + ("block_tables", block_tables), + ("context_lengths", context_lengths), + ): + assert tensor.device == dev, ( + 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), + ): + assert tensor.is_contiguous(), 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: + assert scale.shape[-1] == 1, ( + f"{name} must have a trailing singleton dimension, " + f"got shape {tuple(scale.shape)}" + ) + return scale.squeeze(-1) + assert scale.dim() == 3, ( + 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 + assert per_token_kv == ( + value_scale_t.numel() > 1 + ), "key_scale and value_scale must both be per-tensor or both be per-token" if per_token_kv: - assert ( - value_scale_t.dim() > 1 - ), "value_scale must also be per-token (dim>1) when key_scale is per-token" assert ( key_scale_t.shape == value_scale_t.shape ), f"key_scale/value_scale shape mismatch: {tuple(key_scale_t.shape)} vs {tuple(value_scale_t.shape)}" @@ -151,6 +242,14 @@ def pa_decode( ) stride_ks_block = int(key_scale_t.stride(0)) stride_ks_head = int(key_scale_t.stride(1)) + assert key_scale_t.stride(2) == 1, ( + f"per-token key_scale token dimension must be contiguous, " + f"got strides {key_scale_t.stride()}" + ) + assert value_scale_t.stride() == key_scale_t.stride(), ( + "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 @@ -161,111 +260,114 @@ def pa_decode( value_scale_t.dtype == torch.float32 and value_scale_t.device == dev ), f"value_scale tensor must be float32 on {dev}, got {value_scale_t.dtype} on {value_scale_t.device}" - if not num_partitions: - blocks_per_partition = KV_COMPUTE_BLOCK // block_size - num_partitions = get_recommended_splits( - num_seqs, - num_kv_heads, - split_kv_blocks=blocks_per_partition, - ) + num_partitions = max_context_partition_num + pmax = max_logits + psum = exp_sums + pout = temporary_output - compiled = compile_pa_decode_tile( - head_dim=head_dim, - query_group_size=query_group_size, - block_size=int(block_size), - 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, - ) + 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_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, + device_index=dev.index, + ) - is_graph_capturing = _is_current_stream_capturing() if num_partitions == 1: # NP==1 writes output directly; partials unused (caller buffers ignored). - if pmax is None: - if is_graph_capturing: - raise ValueError( - "CUDA graph capture requires preallocated `pmax`/`psum`/`pout` " - "even when num_partitions==1 (nothing may be allocated mid-capture)." - ) - pmax = psum = pout = torch.empty(1, dtype=torch.float32, device=dev) + 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 pmax is None or psum is None or pout is None: - if is_graph_capturing: - raise ValueError( - "CUDA graph capture requires preallocated `pmax`/`psum`/`pout` " - "for num_partitions>1 (nothing may be allocated mid-capture)." + if pmax is None: + pmax = torch.empty( + *expected_scalar_shape, dtype=torch.float32, device=dev ) - pmax = torch.empty(*expected_scalar_shape, dtype=torch.float32, device=dev) - psum = torch.empty(*expected_scalar_shape, dtype=torch.float32, device=dev) - pout = torch.empty( - *expected_scalar_shape, head_dim, dtype=output.dtype, device=dev - ) - else: - assert ( - pmax.shape == expected_scalar_shape - ), f"pmax shape {tuple(pmax.shape)} != {expected_scalar_shape}" - assert ( - psum.shape == expected_scalar_shape - ), f"psum shape {tuple(psum.shape)} != {expected_scalar_shape}" - assert pout.shape == ( - *expected_scalar_shape, - head_dim, - ), f"pout shape {tuple(pout.shape)} != {(*expected_scalar_shape, head_dim)}" - s = stream or torch.cuda.current_stream() - - _run_compiled( - compiled["launch"], - output, - pmax.view(-1), - psum.view(-1), - pout.view(-1), - query, - key_cache, - value_cache, - block_tables, - context_lengths, - key_scale_t, - value_scale_t, - int(max_blocks_per_seq), - int(num_seqs), - int(num_kv_heads), - stride_ks_block, - stride_ks_head, - int(query.stride(0)), - int(query.stride(1)), - s, - ) - if num_partitions > 1: - from csrc.cpp_itfs.pa.pa_ps import launch_pa_decode_ps_reduce - - output_5d = output.reshape( - num_seqs, query_length, num_kv_heads, query_group_size, head_dim + 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 + ) + assert ( + pmax.shape == expected_scalar_shape + ), f"max_logits shape {tuple(pmax.shape)} != {expected_scalar_shape}" + assert ( + psum.shape == expected_scalar_shape + ), f"exp_sums shape {tuple(psum.shape)} != {expected_scalar_shape}" + assert pout.shape == ( + *expected_scalar_shape, + head_dim, + ), ( + f"temporary_output shape {tuple(pout.shape)} != " + f"{(*expected_scalar_shape, head_dim)}" + ) + assert pmax.dtype == torch.float32 and pmax.device == dev + assert psum.dtype == torch.float32 and psum.device == dev + assert pout.dtype == output.dtype and pout.device == dev + assert pmax.is_contiguous() + assert psum.is_contiguous() + assert pout.is_contiguous() + with torch.cuda.device(dev): + s = torch.cuda.current_stream(dev) + _run_compiled( + compiled["launch"], + output, + pmax.view(-1), + psum.view(-1), + pout.view(-1), + query, + key_cache, + value_cache, + block_tables, + context_lengths, + key_scale_t, + value_scale_t, + int(max_blocks_per_seq), + int(num_seqs), + int(num_kv_heads), + stride_ks_block, + stride_ks_head, + int(query.stride(0)), + int(query.stride(1)), + s, ) - with torch.cuda.stream(s): - launch_pa_decode_ps_reduce( - output_5d, - psum, - pmax, - pout, - None, - output_5d.stride(0), - output_5d.stride(1), - output_5d.stride(2), - output_5d.stride(3), - pmax.stride(0), - pmax.stride(1), - pmax.stride(2), - pout.stride(0), - pout.stride(1), - pout.stride(2), - pout.stride(3), - query_seq_len=query_length, - query_group_size=query_group_size, - head_size=head_dim, - context_partition_num=num_partitions, + if num_partitions > 1: + from csrc.cpp_itfs.pa.pa_ps import launch_pa_decode_ps_reduce + + output_5d = output.reshape( + num_seqs, query_length, num_kv_heads, query_group_size, head_dim ) + with torch.cuda.stream(s): + launch_pa_decode_ps_reduce( + output_5d, + psum, + pmax, + pout, + None, + output_5d.stride(0), + output_5d.stride(1), + output_5d.stride(2), + output_5d.stride(3), + pmax.stride(0), + pmax.stride(1), + pmax.stride(2), + pout.stride(0), + pout.stride(1), + pout.stride(2), + pout.stride(3), + query_seq_len=query_length, + query_group_size=query_group_size, + head_size=head_dim, + context_partition_num=num_partitions, + ) diff --git a/op_tests/flydsl_tests/test_pa_decode_tile.py b/op_tests/flydsl_tests/test_pa_decode_tile.py index 05a8d8bfca3..270b8ad4446 100644 --- a/op_tests/flydsl_tests/test_pa_decode_tile.py +++ b/op_tests/flydsl_tests/test_pa_decode_tile.py @@ -4,13 +4,18 @@ """Correctness and performance sweep for FlyDSL paged-attention Tile.""" import argparse +import importlib +import inspect import itertools import aiter import pandas as pd +import pytest import torch from aiter import dtypes, per_tensor_quant -from aiter.jit.utils.chip_info import get_gfx +from aiter.jit.utils.chip_info import get_gfx_runtime +from aiter.ops.attention import pa_decode_gluon as public_pa_decode +from aiter.ops.triton.gluon.pa_decode_gluon import pa_decode_gluon from aiter.test_common import benchmark, checkAllclose, run_perftest torch.set_default_device("cuda") @@ -36,13 +41,185 @@ get_recommended_splits, pa_decode, ) -except ImportError: +except (ImportError, AttributeError, RuntimeError, OSError): get_recommended_splits = None pa_decode = None def _quant_dtype() -> torch.dtype: - return torch.float8_e4m3fn if get_gfx() == "gfx950" else torch.float8_e4m3fnuz + return ( + torch.float8_e4m3fn if get_gfx_runtime() == "gfx950" else torch.float8_e4m3fnuz + ) + + +def test_pa_decode_api_matches_gluon(): + if pa_decode is None: + pytest.skip("FlyDSL is not available") + + flydsl_parameters = inspect.signature(pa_decode).parameters + gluon_parameters = inspect.signature(pa_decode_gluon).parameters + public_parameters = inspect.signature(public_pa_decode).parameters + + assert tuple(flydsl_parameters) == tuple(gluon_parameters) + assert tuple(public_parameters) == tuple(gluon_parameters) + for name, parameter in flydsl_parameters.items(): + gluon_parameter = gluon_parameters[name] + assert parameter.kind == gluon_parameter.kind + assert parameter.default == gluon_parameter.default + assert public_parameters[name].kind == gluon_parameter.kind + assert public_parameters[name].default == gluon_parameter.default + + +def test_pa_decode_maps_gluon_buffers_and_scale_layout(monkeypatch): + if not torch.cuda.is_available(): + pytest.skip("ROCm is not available") + if pa_decode is None: + pytest.skip("FlyDSL is not available") + + pa_decode_module = importlib.import_module("aiter.ops.flydsl.pa_decode") + attention_module = importlib.import_module("aiter.ops.attention") + captured = {} + + monkeypatch.setattr( + pa_decode_module, + "compile_pa_decode_tile", + lambda **kwargs: {"launch": object()}, + ) + + def capture_launch( + launch, + output, + max_logits, + exp_sums, + temporary_output, + query, + key_cache, + value_cache, + block_tables, + context_lengths, + key_scale, + value_scale, + *args, + ): + captured.update( + max_logits=max_logits, + exp_sums=exp_sums, + temporary_output=temporary_output, + key_scale=key_scale, + value_scale=value_scale, + ) + + monkeypatch.setattr(pa_decode_module, "_run_compiled", capture_launch) + pa_ps_module = importlib.import_module("csrc.cpp_itfs.pa.pa_ps") + monkeypatch.setattr( + pa_ps_module, + "launch_pa_decode_ps_reduce", + lambda *args, **kwargs: None, + ) + + query = torch.empty(1, 8, 128, dtype=torch.bfloat16) + output = torch.empty_like(query) + key_cache = torch.empty(1, 1, 8, 16, 16, dtype=_quant_dtype()) + value_cache = torch.empty(1, 1, 128, 16, dtype=_quant_dtype()) + context_lengths = torch.tensor([16], dtype=torch.int32) + block_tables = torch.tensor([[0]], dtype=torch.int32) + key_scale = torch.ones(1, 1, 16, 1, dtype=torch.float32) + value_scale = torch.ones_like(key_scale) + max_logits = torch.empty(1, 1, 2, 8, dtype=torch.float32) + exp_sums = torch.empty_like(max_logits) + temporary_output = torch.empty(1, 1, 2, 8, 128, dtype=query.dtype) + + assert attention_module._can_use_flydsl_pa_decode( + output, + query, + key_cache, + value_cache, + context_lengths, + block_tables, + 1, + 2, + 256, + key_cache.dtype, + None, + key_scale, + value_scale, + exp_sums, + max_logits, + temporary_output, + None, + None, + 0, + ) + pa_decode( + output, + query, + key_cache, + value_cache, + context_lengths, + block_tables, + 128**-0.5, + 1, + 2, + compute_type=key_cache.dtype, + key_scale=key_scale, + value_scale=value_scale, + exp_sums=exp_sums, + max_logits=max_logits, + temporary_output=temporary_output, + ) + + assert captured["key_scale"].shape == (1, 1, 16) + assert captured["value_scale"].shape == (1, 1, 16) + assert captured["max_logits"].data_ptr() == max_logits.data_ptr() + assert captured["exp_sums"].data_ptr() == exp_sums.data_ptr() + assert captured["temporary_output"].data_ptr() == temporary_output.data_ptr() + + dispatches = [] + monkeypatch.setattr( + attention_module, + "_pa_decode_flydsl", + lambda *args, **kwargs: dispatches.append("flydsl"), + ) + monkeypatch.setattr( + attention_module, + "_pa_decode_gluon_fallback", + lambda *args, **kwargs: dispatches.append("gluon"), + ) + common_kwargs = { + "key_scale": key_scale, + "value_scale": value_scale, + "exp_sums": exp_sums, + "max_logits": max_logits, + "temporary_output": temporary_output, + } + public_pa_decode( + output, + query, + key_cache, + value_cache, + context_lengths, + block_tables, + 128**-0.5, + 1, + 2, + compute_type=key_cache.dtype, + **common_kwargs, + ) + public_pa_decode( + output, + query, + key_cache, + value_cache, + context_lengths, + block_tables, + 128**-0.5, + 1, + 2, + compute_type=key_cache.dtype, + query_scale=torch.ones(1, dtype=torch.float32), + **common_kwargs, + ) + assert dispatches == ["flydsl", "gluon"] def run_torch( @@ -95,29 +272,36 @@ def _run_flydsl( key_scale, value_scale, num_partitions, + softmax_scale, pmax, psum, pout, ): - pa_decode( + torch.ops.aiter.pa_decode_gluon( output, query, key_cache, value_cache, - block_tables, context_lengths, + block_tables, + softmax_scale, + query.shape[0] // context_lengths.shape[0], + num_partitions, + 256, + key_cache.dtype, + None, key_scale, value_scale, - num_partitions=num_partitions, - pmax=pmax, - psum=psum, - pout=pout, + exp_sums=psum, + max_logits=pmax, + temporary_output=pout, + ps=True, ) return output @benchmark() -def test_pa_decode_tile( +def run_pa_decode_tile_case( batch_size, num_query_heads, num_kv_heads, @@ -204,6 +388,7 @@ def test_pa_decode_tile( pmax = torch.empty(partial_shape, dtype=dtypes.fp32) psum = torch.empty_like(pmax) pout = torch.empty(*partial_shape, head_dim, dtype=dtype) + softmax_scale = head_dim**-0.5 candidates = { "flydsl": lambda: _run_flydsl( @@ -216,6 +401,7 @@ def test_pa_decode_tile( key_scale, value_scale, num_partitions, + softmax_scale, pmax, psum, pout, @@ -239,7 +425,7 @@ def test_pa_decode_tile( + value_scale.numel() * value_scale.element_size() ) - ret = {"gfx": get_gfx(), "partitions": num_partitions} + ret = {"gfx": get_gfx_runtime(), "partitions": num_partitions} for name, fn in candidates.items(): out, us = run_perftest(fn) err = checkAllclose( @@ -257,12 +443,33 @@ def test_pa_decode_tile( return ret +@pytest.mark.parametrize("block_size", [16, 64]) +def test_pa_decode_tile(block_size): + 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()}") + + result = run_pa_decode_tile_case( + batch_size=3, + num_query_heads=8, + num_kv_heads=1, + head_dim=128, + context_length=257, + block_size=block_size, + dtype=dtypes.bf16, + ) + assert result["flydsl err"] == 0 + + def main(): if not torch.cuda.is_available(): aiter.logger.warning("ROCm is not available; skipping pa_decode") return - if get_gfx() not in SUPPORTED_GFX: - aiter.logger.warning("pa_decode unsupported on %s; skipping", get_gfx()) + if get_gfx_runtime() not in SUPPORTED_GFX: + aiter.logger.warning("pa_decode unsupported on %s; skipping", get_gfx_runtime()) return if pa_decode is None: aiter.logger.warning("flydsl is unavailable; skipping pa_decode") @@ -315,7 +522,7 @@ def main(): ): num_query_heads, num_kv_heads, head_dim, context_length = shape rows.append( - test_pa_decode_tile( + run_pa_decode_tile_case( batch_size, num_query_heads, num_kv_heads, From 6f8d4fdf966917b03c35c5982619b6e21e3188ef Mon Sep 17 00:00:00 2001 From: fsx950223 Date: Fri, 24 Jul 2026 09:52:56 +0000 Subject: [PATCH 07/66] fix(ci): Resolve Ruff 0.16 lint failures Signed-off-by: fsx950223 Co-authored-by: Cursor --- aiter/ops/flydsl/kernels/pa_decode_tile.py | 18 +++++++++++------- op_tests/flydsl_tests/test_pa_decode_tile.py | 3 ++- 2 files changed, 13 insertions(+), 8 deletions(-) diff --git a/aiter/ops/flydsl/kernels/pa_decode_tile.py b/aiter/ops/flydsl/kernels/pa_decode_tile.py index e2cb40a279a..c5977cd9133 100644 --- a/aiter/ops/flydsl/kernels/pa_decode_tile.py +++ b/aiter/ops/flydsl/kernels/pa_decode_tile.py @@ -49,7 +49,7 @@ KV_COMPUTE_BLOCK = 256 -@functools.lru_cache(maxsize=None) +@functools.cache def compile_pa_decode_tile( *, head_dim: int, @@ -239,7 +239,7 @@ def _load(elem_idx): ) # Per-lane Q chunk (QCHUNK 16-bit elems) fetched in QLOAD_UNIT-wide # pieces (128b max per buffer load): head_dim=256 needs 2 pieces. - QLOAD_UNIT = QCHUNK if QCHUNK < 8 else 8 + QLOAD_UNIT = min(8, QCHUNK) N_QLOADS = QCHUNK // QLOAD_UNIT _q_copy_op = ( fx.rocdl.BufferCopy128b() if QLOAD_UNIT == 8 else fx.rocdl.BufferCopy64b() @@ -579,7 +579,7 @@ def _quant_q_row(m, qi, gs_head, q_row_off): # possibly on the last M-tile (only it can be a partial tile), so # skip the runtime branch (and its EXEC-mask overhead) entirely # for every other M-tile. - if const_expr((m + 1) * MFMA_MNK <= TOTAL_ROWS): + if const_expr((m + 1) * MFMA_MNK <= TOTAL_ROWS): # noqa: SIM114 _quant_q_row(m, qi, gs_head, q_row_off) elif flat_idx < TOTAL_ROWS: _quant_q_row(m, qi, gs_head, q_row_off) @@ -1154,14 +1154,18 @@ def _lmax_off_m(m): def _emit(o_norm, sub): if const_expr(NP == 1): - out_row = output_ptr[seq * query_length + qi_e, qh, None] + out_row = output_ptr[ + seq * query_length + qi_e, qh, None # noqa: B023 + ] out_chunk = fx.slice( fx.logical_divide(out_row, fx.make_layout(OP_ELEMS, 1)), (None, sub), ) out_chunk.store(o_norm) else: - base = ((seq * n_kv + kv_h) * NP + part) * TOTAL_ROWS + row + base = ( + (seq * n_kv + kv_h) * NP + part + ) * TOTAL_ROWS + row # noqa: B023 pout_div = fx.logical_divide(pout_ptr, fx.make_layout(OP_ELEMS, 1)) pout_chunk = fx.slice( pout_div, (None, base * (head_dim // OP_ELEMS) + sub) @@ -1179,7 +1183,7 @@ def _emit(o_norm, sub): if row < TOTAL_ROWS: _emit(o_norm, sub) - if const_expr(NP > 1): + if const_expr(NP > 1): # noqa: SIM102 if warp == 0 and rgroup == 0: base = ((seq * n_kv + kv_h) * NP + part) * TOTAL_ROWS + row if row < TOTAL_ROWS: @@ -1209,7 +1213,7 @@ def pa_decode_tile_launch( stride_ks_head: fx.Int32, stride_q_row: fx.Int32, stride_q_head: fx.Int32, - stream: fx.Stream = fx.Stream(None), + stream: fx.Stream = fx.Stream(None), # noqa: B008 ): pa_decode_tile_kernel( output, diff --git a/op_tests/flydsl_tests/test_pa_decode_tile.py b/op_tests/flydsl_tests/test_pa_decode_tile.py index 270b8ad4446..38a27905936 100644 --- a/op_tests/flydsl_tests/test_pa_decode_tile.py +++ b/op_tests/flydsl_tests/test_pa_decode_tile.py @@ -8,10 +8,11 @@ import inspect import itertools -import aiter import pandas as pd import pytest import torch + +import aiter from aiter import dtypes, per_tensor_quant from aiter.jit.utils.chip_info import get_gfx_runtime from aiter.ops.attention import pa_decode_gluon as public_pa_decode From 3500f28bc434f4449560f1b55cfaba63aa91d149 Mon Sep 17 00:00:00 2001 From: fsx950223 Date: Fri, 14 Aug 2026 08:12:29 +0000 Subject: [PATCH 08/66] fix(flydsl): use in-tree PA buffer helpers Import the buffer helpers from AITER because current FlyDSL releases no longer export buffer_ops from flydsl.expr. Signed-off-by: fsx950223 Co-authored-by: Cursor --- aiter/ops/flydsl/kernels/pa_decode_tile.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/aiter/ops/flydsl/kernels/pa_decode_tile.py b/aiter/ops/flydsl/kernels/pa_decode_tile.py index c5977cd9133..19b2c973a06 100644 --- a/aiter/ops/flydsl/kernels/pa_decode_tile.py +++ b/aiter/ops/flydsl/kernels/pa_decode_tile.py @@ -31,13 +31,13 @@ import flydsl.compiler as flyc import flydsl.expr as fx from flydsl.compiler.protocol import dsl_size_of -from flydsl.expr import arith, buffer_ops, const_expr, gpu, range_constexpr, vector +from flydsl.expr import arith, const_expr, gpu, range_constexpr, vector from flydsl.expr import math as fmath from flydsl.expr.typing import T from flydsl.expr.vector import ReductionOp from flydsl.runtime.device import get_rocm_arch -from . import dpp_utils +from . import buffer_ops, dpp_utils from .utils import cdiv, exp2_amdgcn_scalar, exp2_f32_fast, rcp_f32 MFMA_MNK = ( From 48e320d32dd74fe577258d378063cbc23eb3dfb4 Mon Sep 17 00:00:00 2001 From: fsx950223 Date: Fri, 14 Aug 2026 08:32:48 +0000 Subject: [PATCH 09/66] fix(flydsl): support FlyDSL 0.3 vector APIs Signed-off-by: fsx950223 Co-authored-by: Cursor --- aiter/ops/flydsl/kernels/pa_decode_tile.py | 13 +++---------- 1 file changed, 3 insertions(+), 10 deletions(-) diff --git a/aiter/ops/flydsl/kernels/pa_decode_tile.py b/aiter/ops/flydsl/kernels/pa_decode_tile.py index 19b2c973a06..4b4aa5665ae 100644 --- a/aiter/ops/flydsl/kernels/pa_decode_tile.py +++ b/aiter/ops/flydsl/kernels/pa_decode_tile.py @@ -31,10 +31,9 @@ import flydsl.compiler as flyc import flydsl.expr as fx from flydsl.compiler.protocol import dsl_size_of -from flydsl.expr import arith, const_expr, gpu, range_constexpr, vector +from flydsl.expr import arith, const_expr, gpu, range_constexpr from flydsl.expr import math as fmath -from flydsl.expr.typing import T -from flydsl.expr.vector import ReductionOp +from flydsl.expr.typing import ReductionOp, T from flydsl.runtime.device import get_rocm_arch from . import buffer_ops, dpp_utils @@ -375,13 +374,7 @@ def _stage_kv_scale_to_lds(phys_vec, buf_off=0): else: # block_size==16: each lane stages its own rgroup's page/token, # so the 4 sub-blocks stage in parallel across rgroup-groups. - phys = fx.Int32( - vector.extract( - arith.unwrap(phys_vec), - static_position=[], - dynamic_position=[fx.Index(rgroup)], - ) - ) + 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( buffer_ops.buffer_load( From 10ad6d34b4d52d9843c4bd40733b68c05a92f47f Mon Sep 17 00:00:00 2001 From: fsx950223 Date: Fri, 21 Aug 2026 08:22:06 +0000 Subject: [PATCH 10/66] [FlyDSL] fix(pa_decode): rename to pa_decode_flydsl and fix review blockers Renames the public entry point added by #4332 and addresses the correctness items raised in review. Dispatch - `pa_decode_gluon` in aiter/ops/attention.py goes back to main's plain import + registration; the FlyDSL kernel is registered separately as `torch.ops.aiter.pa_decode_flydsl`. The eligibility gate (`_can_use_flydsl_pa_decode`) and the Gluon fallback are removed, so nothing is silently rerouted and no input can be captured into FlyDSL and then rejected by an assert. Kernel - Reject head_dim whose head_dim//16 is above 8 and not a multiple of 8 (192/320/448/...). Q staging fetches each lane's chunk in <=8-wide pieces, so those widths silently dropped the tail: measured 7.8% error at 192 and 97.6% at 320 against a torch reference. - Fold the per-row score scale in before the -inf mask in the phase-split path. An all-zero query row quantises to scale 0, and `-inf * 0` reached exp2 and turned the whole output into NaN (4096/4096 at ql=4). max() commutes with a non-negative scale, so the row max is unchanged. - Restrict the per-token pv_max reduction to context-visible tokens, so an extreme scale on an unwritten slot cannot shrink valid probabilities to zero in the fp8 pack. Uses context_len rather than the per-row causal bound, which keeps pv_max independent of the query row. - Pin block-table entries past the sequence's extent to block 0. They hold whatever the caller left there, often a stale id into another sequence's live page, and the PV matmul still multiplies those V bytes because a zero probability is not an annihilator (0 * NaN == NaN). - Widen the K/V element offsets to 64-bit once a cache tensor passes 2 GiB, where the i32 page-stride product wraps; reproduced as a GPU memory-access fault at 2.00 GiB, correct afterwards. Gated on the cache size because the wider math costs ~18% at block_size=64. - Drop the device_index cache key. Move arith.select, arith.constant and arith.constant_vector to the fx surface, and replace the per-op arith.FastMathFlags.contract plumbing with an ambient `CompilationContext.compile_hints({"fastmath": "contract"})` around kernel emission -- byte-identical IR (103 contract / 24 nnan flags on both the single- and multi-M-tile paths). `fm_nnan` stays explicit, since an explicit fastmath arg wins over the ambient hint. Tests - Move op_tests/flydsl_tests/test_pa_decode_tile.py to op_tests/test_flydsl_pa_decode.py so `find op_tests -maxdepth 1` in split_tests.sh collects it; neither the aiter nor the triton shard saw the old path. - Add adversarial cases for each fix above, driven through torch.ops.aiter.pa_decode_flydsl, with explicit isfinite assertions. - Scope the module-level torch.set_default_device("cuda") into a fixture now that this file shares a pytest session. Verified on gfx942 with FlyDSL 0.3.1: 22 passed, 1 skipped. Co-Authored-By: Claude Opus 5 (1M context) --- aiter/ops/attention.py | 137 +-------- aiter/ops/flydsl/kernels/pa_decode_tile.py | 254 ++++++++++------ aiter/ops/flydsl/pa_decode.py | 23 +- ...ecode_tile.py => test_flydsl_pa_decode.py} | 271 ++++++++++++++---- 4 files changed, 413 insertions(+), 272 deletions(-) rename op_tests/{flydsl_tests/test_pa_decode_tile.py => test_flydsl_pa_decode.py} (64%) diff --git a/aiter/ops/attention.py b/aiter/ops/attention.py index 1e7a976637b..4a0a5fe13a8 100644 --- a/aiter/ops/attention.py +++ b/aiter/ops/attention.py @@ -9,9 +9,7 @@ from aiter import dtypes from aiter.ops.enum import Enum, MlaVersion, QuantType -from aiter.ops.triton.gluon.pa_decode_gluon import ( - pa_decode_gluon as _pa_decode_gluon_fallback, -) +from aiter.ops.triton.gluon.pa_decode_gluon import pa_decode_gluon from aiter.utility.dtypes import _aiter_dtype_id from csrc.cpp_itfs.pa.pa import paged_attention_rocm as paged_attention_rocm_core from csrc.cpp_itfs.pa.pa_ragged import ( @@ -21,7 +19,7 @@ from csrc.cpp_itfs.torch_utils import direct_register_custom_op from ..jit.core import compile_ops, is_experimental_enabled -from ..jit.utils.chip_info import get_cu_num, get_gfx, get_gfx_runtime +from ..jit.utils.chip_info import get_cu_num, get_gfx try: from aiter.ops.flydsl import pa_decode as _pa_decode_flydsl @@ -31,101 +29,7 @@ MD_NAME = "module_attention" -def _can_use_flydsl_pa_decode( - output, - query, - key_cache, - value_cache, - context_lengths, - block_tables, - query_length, - max_context_partition_num, - context_partition_size, - compute_type, - query_scale, - key_scale, - value_scale, - exp_sums, - max_logits, - temporary_output, - alibi_slopes, - sinks, - sliding_window, -): - runtime_gfx = get_gfx_runtime() - runtime_fp8_dtype = { - "gfx942": torch.float8_e4m3fnuz, - "gfx950": torch.float8_e4m3fn, - }.get(runtime_gfx) - - def scratch_buffers_match(): - if max_context_partition_num == 1: - return True - query_group_size = query.shape[1] // key_cache.shape[1] - total_rows = query_length * query_group_size - scalar_shape = ( - context_lengths.shape[0], - key_cache.shape[1], - max_context_partition_num, - total_rows, - ) - output_shape = (*scalar_shape, query.shape[-1]) - scalar_buffers_match = all( - tensor is None - or ( - tensor.shape == scalar_shape - and tensor.dtype == torch.float32 - and tensor.device == query.device - and tensor.is_contiguous() - ) - for tensor in (exp_sums, max_logits) - ) - return scalar_buffers_match and ( - temporary_output is None - or ( - temporary_output.shape == output_shape - and temporary_output.dtype == output.dtype - and temporary_output.device == query.device - and temporary_output.is_contiguous() - ) - ) - - return ( - _pa_decode_flydsl is not None - and runtime_fp8_dtype is not None - and query.dim() == 3 - and output.shape == query.shape - and output.dtype == query.dtype - and query.dtype in (torch.bfloat16, torch.float16) - and query.stride(-1) == 1 - and key_cache.dim() == 5 - and key_cache.dtype == runtime_fp8_dtype - and key_cache.shape[-2] in (16, 64) - and key_cache.is_contiguous() - and value_cache.dim() in (4, 5) - and value_cache.dtype == runtime_fp8_dtype - and value_cache.is_contiguous() - and query.shape[-1] % 64 == 0 - and key_cache.shape[1] > 0 - and query.shape[1] % key_cache.shape[1] == 0 - and context_lengths.dim() == 1 - and context_lengths.is_contiguous() - and block_tables.dim() == 2 - and block_tables.is_contiguous() - and query.shape[0] == context_lengths.shape[0] * query_length - and 1 <= max_context_partition_num <= 64 - and context_partition_size == 256 - and compute_type == runtime_fp8_dtype - and query_scale is None - and (key_scale is None) == (value_scale is None) - and scratch_buffers_match() - and alibi_slopes is None - and sinks is None - and sliding_window == 0 - ) - - -def pa_decode_gluon( +def pa_decode_flydsl( output: torch.Tensor, query: torch.Tensor, key_cache: torch.Tensor, @@ -148,32 +52,9 @@ def pa_decode_gluon( sliding_window: int = 0, ps: bool = True, ) -> None: - implementation = ( - _pa_decode_flydsl - if _can_use_flydsl_pa_decode( - output, - query, - key_cache, - value_cache, - context_lengths, - block_tables, - query_length, - max_context_partition_num, - context_partition_size, - compute_type, - query_scale, - key_scale, - value_scale, - exp_sums, - max_logits, - temporary_output, - alibi_slopes, - sinks, - sliding_window, - ) - else _pa_decode_gluon_fallback - ) - implementation( + if _pa_decode_flydsl is None: + raise RuntimeError("pa_decode_flydsl requires the `flydsl` package") + _pa_decode_flydsl( output, query, key_cache, @@ -204,6 +85,12 @@ def pa_decode_gluon( ["output", "exp_sums", "max_logits", "temporary_output"], ) +direct_register_custom_op( + "pa_decode_flydsl", + pa_decode_flydsl, + ["output", "exp_sums", "max_logits", "temporary_output"], +) + def gen_pa_fwd_native_fake( # [num_seqs, num_heads, head_size] diff --git a/aiter/ops/flydsl/kernels/pa_decode_tile.py b/aiter/ops/flydsl/kernels/pa_decode_tile.py index 4b4aa5665ae..8d355b0a864 100644 --- a/aiter/ops/flydsl/kernels/pa_decode_tile.py +++ b/aiter/ops/flydsl/kernels/pa_decode_tile.py @@ -30,6 +30,7 @@ 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 @@ -44,6 +45,8 @@ ) MFMA_K = 32 # fp8 MFMA contracts K = 32 per instruction (mfma_f32_16x16x32_fp8_fp8) WAVE = 64 +# f32 C-fragment elements each lane holds for one 16x16x32 atom (16*16 / 64). +MFMA_ACC_ELEMS = MFMA_MNK * MFMA_MNK // WAVE LOG2E = 1.4426950408889634 KV_COMPUTE_BLOCK = 256 @@ -60,19 +63,31 @@ def compile_pa_decode_tile( per_token_kv: bool = False, query_length: int = 1, trans_v: bool = True, - device_index: int | None = None, + wide_kv_addressing: bool = False, ): """Build the tile-programming PA-decode kernel + launch wrapper. ``block_size``, ``head_dim``, and ``query_dtype`` are compile-time - constants. ``device_index`` keeps the cached launch/compiled function - device-local. ``query_length`` (MTP) and ``query_group_size`` flatten into + constants. ``query_length`` (MTP) and ``query_group_size`` flatten into ``TOTAL_ROWS = query_length * query_group_size``, tiled into ``M_TILES = ceil(TOTAL_ROWS / 16)`` independent 16-row MFMA tiles; each extra M-tile duplicates loop-carried state, so VGPR/LDS/occupancy scale roughly linearly with ``M_TILES``. + + Tokens at/after ``context_len`` still take part in the PV matmul with a zero + probability, so their V bytes must be finite: the MFMA does not treat a zero + operand as an annihilator (``0 * NaN == NaN``). Whole pages past the + sequence's extent are pinned to block 0 (see ``_load_phys_scalar``); the + remaining slots are the unwritten tail of the last page the sequence owns, + which the caller is expected to leave finite. + + ``wide_kv_addressing`` computes K/V element offsets in 64-bit. It is only + correct-critical once a single cache tensor passes 2 GiB, where the i32 + page-stride product wraps, and it is not free: measured +18% at + ``block_size=64`` (one page index feeds four step-loads there, so i32 lets + the page-base multiply be hoisted and the offsets stay 32-bit adds) and + neutral at ``block_size=16``. The wrapper turns it on from the cache size. """ - del device_index is_gfx950 = "gfx95" in get_rocm_arch() FP8 = fx.Float8E4M3FN if is_gfx950 else fx.Float8E4M3FNUZ FP8_MAX = ( @@ -128,6 +143,12 @@ def compile_pa_decode_tile( QCHUNK = ( head_dim // NQCHUNK ) # f16 elements per lane's load chunk (8 for head_dim=128, 4 for head_dim=64) + # The chunk is fetched in min(8, QCHUNK)-wide pieces (128b max per load), so + # a QCHUNK above 8 that isn't a multiple of 8 would leave its tail unloaded. + 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" + ) VHE_CHUNKS = head_dim // (NWARP * MFMA_MNK) # 2 for head_dim=128, 1 for head_dim=64 @@ -210,7 +231,7 @@ def pa_decode_tile_kernel( # fx.copy-based flat loaders: K/V/context_len over a raw pointer, Q over # a buffer resource. - def _make_flat_loader(tensor_ptr, elem_ty, reg_width, copy_op): + def _make_flat_loader(tensor_ptr, elem_ty, reg_width, copy_op, extent=1 << 30): use_buffer_resource = isinstance( copy_op, fx.rocdl.CopyOpCDNA3BufferCopyType ) @@ -221,7 +242,7 @@ def _make_flat_loader(tensor_ptr, elem_ty, reg_width, copy_op): if use_buffer_resource else fx.get_iter(tensor_ptr) ) - flat = fx.Tensor(fx.make_view(base_iter, fx.make_layout(1 << 30, 1))) + flat = fx.Tensor(fx.make_view(base_iter, fx.make_layout(extent, 1))) div = fx.logical_divide(flat, fx.make_layout(1, 1)) def _load(elem_idx): @@ -230,12 +251,23 @@ def _load(elem_idx): return _load + # A cache above 2 GiB overflows an i32 element offset, so the flat view + # and the offsets below widen together. + KV_EXTENT = (1 << 42) if wide_kv_addressing else (1 << 30) _k_load_fp8x16 = _make_flat_loader( - key_cache_ptr, FP8, 16, fx.UniversalCopy128b() + key_cache_ptr, FP8, 16, fx.UniversalCopy128b(), KV_EXTENT ) _v_load_fp8x16 = _make_flat_loader( - value_cache_ptr, FP8, 16, fx.UniversalCopy128b() + value_cache_ptr, FP8, 16, fx.UniversalCopy128b(), KV_EXTENT ) + + def _kv_addr(phys, page_elems, rest): + # `phys * page_elems` is the term that overflows first: it reaches + # 2^31 elements at a 2 GiB cache. `rest` stays inside one page. + if const_expr(wide_kv_addressing): + return fx.Int64(phys) * fx.Int64(page_elems) + fx.Int64(rest) + return phys * page_elems + rest + # Per-lane Q chunk (QCHUNK 16-bit elems) fetched in QLOAD_UNIT-wide # pieces (128b max per buffer load): head_dim=256 needs 2 pieces. QLOAD_UNIT = min(8, QCHUNK) @@ -272,20 +304,21 @@ def _v_load16(byte_off): if const_expr(not per_token_kv): key_scale = fx.Int32( buffer_ops.buffer_load( - ks_rsrc, arith.constant(0, type=T.i32), vec_width=1, is_scalar=True + ks_rsrc, fx.Int32(0), vec_width=1, is_scalar=True ) ).bitcast(fx.Float32) value_scale = fx.Int32( buffer_ops.buffer_load( - vs_rsrc, arith.constant(0, type=T.i32), vec_width=1, is_scalar=True + vs_rsrc, fx.Int32(0), vec_width=1, is_scalar=True ) ).bitcast(fx.Float32) num_tiles = cdiv(context_len, TILE_TOK) + num_pages = cdiv(context_len, block_size) # pages this sequence really owns tiles_per_part = cdiv(num_tiles, NP) part_start = part * tiles_per_part part_end_raw = part_start + tiles_per_part - part_end = arith.select(part_end_raw < num_tiles, part_end_raw, num_tiles) + part_end = (part_end_raw < num_tiles).select(part_end_raw, num_tiles) # One i8 blob carved into typed byte-offset pointers. `lds_base` is an # ir.Value pointer (safe inside scf control flow); the Python `lds` @@ -318,13 +351,27 @@ def _lds_store(byte_off, elem_ty, vec): NCHUNK = TILE_TOK // TOK_CHUNK # 4 def _load_phys_scalar(page, vec_width=1): + # Past `num_pages` the block-table entry is padding: whatever the + # caller left there, often a stale id pointing at another sequence's + # live page. Pin those to block 0 so the tokens the softmax masks + # out always resolve to one known, in-bounds page instead. `page` is + # wave-uniform, so the clamp is scalar. result = buffer_ops.buffer_load( bt_rsrc, seq * max_blocks_per_seq + page, vec_width=vec_width, is_scalar=True, ) - return fx.Int32(result) if vec_width == 1 else result + if const_expr(vec_width == 1): + return (page < num_pages).select(fx.Int32(result), fx.Int32(0)) + loaded = fx.Vector(result) + 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 _v_page_fetch_and_stage(tt_i32): # V's page depends on `rgroup` (shared across warps): warp w fetches @@ -417,10 +464,12 @@ def _k_ops(phys, a): ops = [] for qkhe in range_constexpr(QKHE_LOOP): he_idx = qkhe * RGROUP_QUARTERS + rgroup - base = ( - ((phys * n_kv + kv_h) * QCHUNK + he_idx) * block_size - + within_page_tok - ) * QK_CHUNK_ELEMS + 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] -> k_step 2*qkhe, 2*qkhe+1 if const_expr(block_size == 16): # help the scheduler overlap the PAGES_PER_CHUNK gathered loads @@ -449,7 +498,7 @@ def _k_ops_flat(tt_i32): # -- prologue: prefetch the first tile's K -- num_tiles_m1 = num_tiles - 1 - start_safe = arith.select(part_start < num_tiles, part_start, num_tiles_m1) + start_safe = (part_start < num_tiles).select(part_start, num_tiles_m1) k_pf0, phys_vec0 = _k_ops_flat(start_safe) # V page-index prefetch, issued here too for the same overlap; the # LDS write is visible after the barrier below. @@ -467,7 +516,6 @@ def _k_ops_flat(tt_i32): v_scale_f = fx.Float32(value_scale) NEG_INF = fx.Float32(float("-inf")) ZERO_F = fx.Float32(0.0) - fm_contract = arith.FastMathFlags.contract # Softmax scores are finite or the -inf mask sentinel -- never NaN -- so # nnan lets maxnum lower to a bare v_max (no v_cmp_u NaN check + its s_nop # hazard) and fuse to v_max3. (ninf must NOT be set: -inf is load-bearing.) @@ -632,15 +680,18 @@ def _v_ops(phys_row, vh): for sub in range_constexpr(PAGES_PER_CHUNK): for step in range_constexpr(STEPS_PER_PAGE): if const_expr(trans_v): - base = ( - ((phys_row[sub] * n_kv + kv_h) * STEPS_PER_PAGE + step) - * head_dim - + head_element - ) * 16 + base = _kv_addr( + phys_row[sub], + n_kv * (STEPS_PER_PAGE * head_dim * 16), + ((kv_h * STEPS_PER_PAGE + step) * head_dim + head_element) + * 16, + ) else: - base = ( - (phys_row[sub] * n_kv + kv_h) * head_dim + head_element - ) * block_size + step * 16 + base = _kv_addr( + phys_row[sub], + n_kv * (head_dim * block_size), + (kv_h * head_dim + head_element) * block_size + step * 16, + ) w = _v_load16(base) if const_expr(block_size == 16): # help the scheduler overlap the per-page gathered loads (see _k_ops) @@ -705,6 +756,28 @@ def _l_slot(m): # This tile's ping-pong scale buffer; the tt+1 prefetch stages the other. cur_kv_buf = _kv_buf_off(tt) + # Tokens at/after context_len index cache slots the caller never + # wrote. They are masked out of the scores, but their V bytes and V + # scales still reach the PV MFMA / the fp8 normalization, so both + # have to be neutralised here. `context_len` (not the per-row causal + # bound) is the right cutoff: causally-masked tokens inside the + # context hold real data, and using it keeps both masks independent + # of the query row. + tile_valid = context_len - tok0 + + if const_expr(per_token_kv): + ctx_thr = fx.Vector.from_elements( + [ + tile_valid.to(fx.Float32) + - fx.Int32(warp * TOK_CHUNK + 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, zero=zero4_scale): + return (_ct[a] < thr).select(vec, zero) + # V data doesn't depend on `m`; hoist once for the phase-split # instead of reloading per M-tile. v_vh_shared = None @@ -730,7 +803,6 @@ def _lmax_off_m(m): # per-M-tile LDS slices so they share ONE barrier. M_TILES==1: else below. if const_expr(M_TILES > 1): masked_chunks_saved = [None] * M_TILES - scale_saved = [None] * M_TILES # per_token: compute the m-independent pv_max once, hold only # k_scale across Phase A, re-read v_scale for Phase B after the @@ -744,7 +816,8 @@ def _lmax_off_m(m): pv_max = fx.Float32(0.0) for a in range_constexpr(NCHUNK): pv_max = fx.maxnumf( - pv_max, v_scale_A[a].reduce(ReductionOp.MAX) + 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)) @@ -757,7 +830,7 @@ def _lmax_off_m(m): for m in range_constexpr(M_TILES): frag_Ss = [] for a in range_constexpr(NCHUNK): - acc = arith.constant_vector(0.0, T.f32x4) + acc = fx.Vector.filled(MFMA_ACC_ELEMS, 0.0, fx.Float32) for s in range_constexpr(N_SUBCHUNKS): acc = fx.rocdl.mfma_f32_16x16x32_fp8_fp8( T.f32x4, @@ -780,13 +853,23 @@ def _lmax_off_m(m): ).broadcast_to(4) neg4 = fx.Vector.filled(4, float("-inf"), fx.Float32) + # Fold the per-row score scale in BEFORE the -inf mask. A + # zero q_scale (all-zero Q row) would otherwise reach + # `-inf * 0 == NaN` in pass 2 and poison the whole row. + # max() commutes with a non-negative scale, so `pm` below is + # still the correctly scaled row max. + 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] + frag_Ss[a] * k_scale_shared[a] * scale_b for a in range_constexpr(NCHUNK) ] else: - scaled_frags = frag_Ss + scaled_frags = [ + frag_Ss[a] * scale_b for a in range_constexpr(NCHUNK) + ] masked_chunks = [ (_ct[a] < thr).select(scaled_frags[a], neg4) @@ -802,10 +885,11 @@ def _lmax_off_m(m): ) 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 * scale) + # `pm` is already scaled (see scaled_frags above); a fully + # masked row keeps -inf, which `safe_max` turns into 0. + _st_lw(_lmax_off_m(m), lane16, warp, pm) masked_chunks_saved[m] = masked_chunks - scale_saved[m] = scale # tt+1 K/V/scale prefetch, issued once here so the V-page read # reuses this barrier. @@ -846,7 +930,6 @@ def _lmax_off_m(m): ] # this thread's own running denom, carried from last tile masked_chunks = masked_chunks_saved[m] - scale = scale_saved[m] v_max_scaled = None norm_factor_b = None @@ -870,16 +953,14 @@ def _lmax_off_m(m): ) # Fully-invalid row: use 0 as the effective max so masked lanes # give exp2(-inf-0)==0 (avoids the -inf-(-inf) cancellation). - safe_max = arith.select(m_new > NEG_INF, m_new, ZERO_F) + 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( - exp2_f32_fast(masked_chunks[a] * scale - m_new_b) - ) + Pa = fx.Vector(exp2_f32_fast(masked_chunks[a] - m_new_b)) ls = ls + Pa.reduce(ReductionOp.ADD) if const_expr(per_token_kv): v_sc = ( @@ -905,22 +986,16 @@ def _lmax_off_m(m): fx.Vector.from_elements([words[a]], dtype=fx.Int32), ) for sh in (16, 32): - ls = ls.addf(ls.shuffle_xor(sh, WAVE), fastmath=fm_contract) + ls = ls + ls.shuffle_xor(sh, WAVE) # PV output is [head-dim, query-row=lane16] after the operand # swap, so correction/denominator are per-lane scalars (no sCorr). - safe_prev = arith.select(m_prev > NEG_INF, m_prev, ZERO_F) + safe_prev = (m_prev > NEG_INF).select(m_prev, ZERO_F) corr_reg = fx.Float32(exp2_amdgcn_scalar(safe_prev - safe_max)) if rgroup == 0: _st_lw(sLsum_off, lane16, warp, ls) gpu.barrier() gsum = _ld_lw_row(sLsum_off, lane16).reduce(ReductionOp.ADD) - l_new = fx.Float32( - arith.mulf( - arith.unwrap(l_prev), - arith.unwrap(corr_reg), - fastmath=fm_contract, - ) - ).addf(gsum, fastmath=fm_contract) + l_new = l_prev * corr_reg + gsum p_ops = _lds_load( sP_off + lane16 * SP_ROW_BYTES + rgroup * 64, fx.Int64, NVOPS @@ -931,7 +1006,7 @@ def _lmax_off_m(m): ).broadcast_to(OP_ELEMS) for vh in range_constexpr(VHE_CHUNKS): v_vh = v_vh_shared[vh] - acc = arith.constant_vector(0.0, T.f32x4) + acc = fx.Vector.filled(MFMA_ACC_ELEMS, 0.0, fx.Float32) for s in range_constexpr(NVOPS): # SWAPPED operands (V=A, P=B): output row = # head-dim, output col = query-row=lane16. @@ -959,7 +1034,7 @@ def _lmax_off_m(m): # QK: each NCHUNK chunk accumulates N_SUBCHUNKS k_steps into an f32x4. frag_Ss = [] for a in range_constexpr(NCHUNK): - acc = arith.constant_vector(0.0, T.f32x4) + acc = fx.Vector.filled(MFMA_ACC_ELEMS, 0.0, fx.Float32) for s in range_constexpr(N_SUBCHUNKS): acc = fx.rocdl.mfma_f32_16x16x32_fp8_fp8( T.f32x4, @@ -1009,13 +1084,18 @@ def _lmax_off_m(m): _st_lw( sLmax_off, lane16, warp, pm * scale ) # redundant across the 4 lanes sharing this qhead - # per_token_kv: max V-scale for the per-tile fp8 normalization - # (any positive value is correct, so skip the causal mask). + # per_token_kv: max V-scale for the per-tile fp8 normalization. + # Slots at/after context_len are uninitialised, so an arbitrarily + # large one there would shrink every valid probability to 0 in the + # fp8 conversion below -- restrict the reduction to real tokens. + # (Causally-masked tokens inside the context keep real scales and + # stay in, which also keeps pv_max independent of the query row.) if const_expr(per_token_kv): pv_max = fx.Float32(0.0) for a in range_constexpr(NCHUNK): pv_max = fx.maxnumf( - pv_max, v_scale_vecs[a].reduce(ReductionOp.MAX) + 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)) @@ -1077,7 +1157,7 @@ def _lmax_off_m(m): if const_expr(head_dim == 64): fx.rocdl.sched_dswr(NCHUNK) for sh in (16, 32): - ls = ls.addf(ls.shuffle_xor(sh, WAVE), fastmath=fm_contract) + ls = ls + ls.shuffle_xor(sh, WAVE) # PV (V=A, P=B) -> output [head-dim, query-row=lane16]; same as # the phase-split path. corr_reg = fx.Float32(exp2_amdgcn_scalar(m_prev - m_new)) @@ -1085,13 +1165,7 @@ def _lmax_off_m(m): _st_lw(sLsum_off, lane16, warp, ls) gpu.barrier() gsum = _ld_lw_row(sLsum_off, lane16).reduce(ReductionOp.ADD) - l_new = fx.Float32( - arith.mulf( - arith.unwrap(l_prev), - arith.unwrap(corr_reg), - fastmath=fm_contract, - ) - ).addf(gsum, fastmath=fm_contract) + l_new = l_prev * corr_reg + gsum p_ops = _lds_load( sP_off + lane16 * SP_ROW_BYTES + rgroup * 64, fx.Int64, NVOPS ) @@ -1105,7 +1179,7 @@ def _lmax_off_m(m): ] for vh in range_constexpr(VHE_CHUNKS): v_vh = v_vh_batch[vh] - acc = arith.constant_vector(0.0, T.f32x4) + acc = fx.Vector.filled(MFMA_ACC_ELEMS, 0.0, fx.Float32) for s in range_constexpr(NVOPS): acc = fx.rocdl.mfma_f32_16x16x32_fp8_fp8( T.f32x4, [v_vh[s], p_ops[s], acc, 0, 0, 0] @@ -1126,18 +1200,12 @@ def _lmax_off_m(m): for m in range_constexpr(M_TILES): row = m * MFMA_MNK + lane16 # flat (mtp, gqa) query-row for this lane l_row = o_final[_l_slot(m)] - safe_l = arith.select(l_row > ZERO_F, l_row, fx.Float32(1.0)) + safe_l = (l_row > ZERO_F).select(l_row, fx.Float32(1.0)) inv_l = fx.Float32(rcp_f32(safe_l)) if const_expr(per_token_kv): o_scale = inv_l else: - o_scale = fx.Float32( - arith.mulf( - arith.unwrap(inv_l), - arith.unwrap(v_scale_f * inv_fp8), - fastmath=fm_contract, - ) - ) + 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) @@ -1208,27 +1276,31 @@ def pa_decode_tile_launch( stride_q_head: fx.Int32, stream: fx.Stream = fx.Stream(None), # noqa: B008 ): - pa_decode_tile_kernel( - output, - pmax, - psum, - pout, - query, - key_cache, - value_cache, - block_tables, - context_lengths, - key_scale, - value_scale, - max_blocks_per_seq, - stride_ks_block, - stride_ks_head, - stride_q_row, - stride_q_head, - ).launch( - grid=(num_seqs, num_kv_heads, NP), - block=(BLOCK_THREADS, 1, 1), - stream=stream, - ) + # `contract` lets the mul/add pairs in the softmax rescale and the + # epilogue fuse into FMAs. Set as an ambient hint rather than per op; + # an explicit `fastmath=` on an op still wins (see fm_nnan). + 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, + max_blocks_per_seq, + stride_ks_block, + stride_ks_head, + stride_q_row, + stride_q_head, + ).launch( + grid=(num_seqs, num_kv_heads, NP), + block=(BLOCK_THREADS, 1, 1), + stream=stream, + ) return {"launch": pa_decode_tile_launch, "kernel": pa_decode_tile_kernel} diff --git a/aiter/ops/flydsl/pa_decode.py b/aiter/ops/flydsl/pa_decode.py index e15047acc34..90680dadd8e 100644 --- a/aiter/ops/flydsl/pa_decode.py +++ b/aiter/ops/flydsl/pa_decode.py @@ -70,10 +70,10 @@ def pa_decode( sliding_window: int = 0, ps: bool = True, ) -> None: - """FlyDSL replacement for ``pa_decode_gluon``. + """FlyDSL paged-attention fp8 decode. - The call signature and intermediate-buffer layouts match - ``pa_decode_gluon``. This kernel currently supports FP8 K/V caches, + The call signature and intermediate-buffer layouts follow the shared + aiter paged-attention decode API. This kernel currently supports FP8 K/V caches, BF16/FP16 queries, a 256-token context partition, and block sizes 16/64. ALiBi, attention sinks, sliding-window attention, and externally quantized FP8 queries are not supported. @@ -133,6 +133,15 @@ def pa_decode( _, num_kv_heads, num_hgroups, block_size, hgroup_width = key_cache.shape assert num_hgroups == head_dim // 16 and hgroup_width == 16 + # Q staging loads each lane's head_dim//16 elements in <=8-wide pieces, so a + # chunk wider than 8 that isn't a multiple of 8 would silently drop its tail + # (head_dim 192/320/448/...). Reject those rather than return wrong results. + q_chunk = head_dim // 16 + if q_chunk > 8 and q_chunk % 8 != 0: + raise NotImplementedError( + f"pa_decode does not support head_dim={head_dim}: head_dim//16 " + f"({q_chunk}) must be <=8 or a multiple of 8" + ) assert block_size in ( 16, 64, @@ -265,6 +274,12 @@ def normalize_scale(scale, name): psum = exp_sums pout = temporary_output + # An i32 element offset wraps once a single cache tensor passes 2^31 + # elements (2 GiB at fp8). The wider math costs ~18% at block_size=64, so + # pay it only where it is needed. Both caches share the code path, so widen + # if either does. + wide_kv_addressing = max(key_cache.numel(), value_cache.numel()) >= 2**31 + with torch.cuda.device(dev): compiled = compile_pa_decode_tile( head_dim=head_dim, @@ -276,7 +291,7 @@ def normalize_scale(scale, name): per_token_kv=per_token_kv, query_length=query_length, trans_v=trans_v, - device_index=dev.index, + wide_kv_addressing=wide_kv_addressing, ) if num_partitions == 1: diff --git a/op_tests/flydsl_tests/test_pa_decode_tile.py b/op_tests/test_flydsl_pa_decode.py similarity index 64% rename from op_tests/flydsl_tests/test_pa_decode_tile.py rename to op_tests/test_flydsl_pa_decode.py index 38a27905936..1ed19e97eeb 100644 --- a/op_tests/flydsl_tests/test_pa_decode_tile.py +++ b/op_tests/test_flydsl_pa_decode.py @@ -15,11 +15,21 @@ import aiter from aiter import dtypes, per_tensor_quant from aiter.jit.utils.chip_info import get_gfx_runtime -from aiter.ops.attention import pa_decode_gluon as public_pa_decode +from aiter.ops.attention import pa_decode_flydsl as public_pa_decode from aiter.ops.triton.gluon.pa_decode_gluon import pa_decode_gluon from aiter.test_common import benchmark, checkAllclose, run_perftest -torch.set_default_device("cuda") + +@pytest.fixture(autouse=True) +def _default_cuda_device(): + # Scoped rather than set at import time: this module now runs in the shared + # Standard Tests shard, where a global default-device switch would follow + # every later test in the session. + previous = torch.get_default_device() + torch.set_default_device("cuda") + yield + torch.set_default_device(previous) + SUPPORTED_GFX = ["gfx942", "gfx950"] KV_COMPUTE_BLOCK = 256 @@ -130,27 +140,6 @@ def capture_launch( exp_sums = torch.empty_like(max_logits) temporary_output = torch.empty(1, 1, 2, 8, 128, dtype=query.dtype) - assert attention_module._can_use_flydsl_pa_decode( - output, - query, - key_cache, - value_cache, - context_lengths, - block_tables, - 1, - 2, - 256, - key_cache.dtype, - None, - key_scale, - value_scale, - exp_sums, - max_logits, - temporary_output, - None, - None, - 0, - ) pa_decode( output, query, @@ -181,31 +170,6 @@ def capture_launch( "_pa_decode_flydsl", lambda *args, **kwargs: dispatches.append("flydsl"), ) - monkeypatch.setattr( - attention_module, - "_pa_decode_gluon_fallback", - lambda *args, **kwargs: dispatches.append("gluon"), - ) - common_kwargs = { - "key_scale": key_scale, - "value_scale": value_scale, - "exp_sums": exp_sums, - "max_logits": max_logits, - "temporary_output": temporary_output, - } - public_pa_decode( - output, - query, - key_cache, - value_cache, - context_lengths, - block_tables, - 128**-0.5, - 1, - 2, - compute_type=key_cache.dtype, - **common_kwargs, - ) public_pa_decode( output, query, @@ -217,10 +181,13 @@ def capture_launch( 1, 2, compute_type=key_cache.dtype, - query_scale=torch.ones(1, dtype=torch.float32), - **common_kwargs, + key_scale=key_scale, + value_scale=value_scale, + exp_sums=exp_sums, + max_logits=max_logits, + temporary_output=temporary_output, ) - assert dispatches == ["flydsl", "gluon"] + assert dispatches == ["flydsl"] def run_torch( @@ -278,7 +245,7 @@ def _run_flydsl( psum, pout, ): - torch.ops.aiter.pa_decode_gluon( + torch.ops.aiter.pa_decode_flydsl( output, query, key_cache, @@ -465,7 +432,207 @@ def test_pa_decode_tile(block_size): assert result["flydsl err"] == 0 +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()}") + + +def _adversarial_case( + head_dim=128, + context_length=257, + block_size=16, + query_group_size=8, + query_length=1, + per_token=False, + zero_query=False, + poison_padding_blocks=False, + tail_value_scale=None, + seed=0, +): + """Build one case plus its torch reference over the same dequantised fp8 KV. + + ``poison_padding_blocks`` fills every block past the sequence's real extent + with NaN and pads the block table out to reach them -- the entries a caller + leaves behind, which the kernel must resolve to block 0 rather than follow. + """ + quant_dtype = _quant_dtype() + generator = torch.Generator(device="cuda").manual_seed(seed) + owned = (context_length + block_size - 1) // block_size + # A 256-token tile always walks 256/block_size pages, so pad past `owned` to + # give the block table entries that must not be dereferenced. + blocks = owned + 256 // block_size if poison_padding_blocks else owned + tokens = blocks * block_size + + def rand(*shape): + return torch.rand(*shape, generator=generator, device="cuda") - 0.5 + + query = ( + torch.zeros(query_length, query_group_size, head_dim, dtype=dtypes.bf16) + if zero_query + else rand(query_length, query_group_size, head_dim).to(dtypes.bf16) + ) + key = rand(blocks, 1, block_size, head_dim).to(quant_dtype) + value = rand(blocks, 1, block_size, head_dim).to(quant_dtype) + + invalid = (torch.arange(tokens, device="cuda") >= context_length).view( + blocks, 1, block_size, 1 + ) + if poison_padding_blocks: + padding = (torch.arange(tokens, device="cuda") >= owned * block_size).view( + blocks, 1, block_size, 1 + ) + poison = torch.full_like(value, float("nan"), dtype=dtypes.fp32).to(quant_dtype) + value = torch.where(padding.expand_as(value), poison, value) + key = torch.where(padding.expand_as(key), poison, key) + + if per_token: + key_scale = ( + torch.rand(blocks, 1, block_size, 1, generator=generator, device="cuda") + * 0.5 + + 0.75 + ) + value_scale = ( + torch.rand(blocks, 1, block_size, 1, generator=generator, device="cuda") + * 0.5 + + 0.75 + ) + if tail_value_scale is not None: + value_scale = torch.where( + invalid, torch.full_like(value_scale, tail_value_scale), value_scale + ) + else: + key_scale = torch.ones(1, dtype=dtypes.fp32) + value_scale = torch.ones(1, dtype=dtypes.fp32) + + key_cache = ( + key.view(blocks, 1, block_size, head_dim // 16, 16) + .permute(0, 1, 3, 2, 4) + .contiguous() + ) + value_cache = value.permute(0, 1, 3, 2).contiguous() + block_tables = torch.arange(blocks, dtype=dtypes.i32).reshape(1, blocks) + context_lengths = torch.full((1,), context_length, dtype=dtypes.i32) + + keys = key.float().reshape(tokens, head_dim) + values = value.float().reshape(tokens, head_dim) + if per_token: + keys = keys * key_scale.float().reshape(tokens, 1) + values = values * value_scale.float().reshape(tokens, 1) + reference = torch.empty( + query_length, query_group_size, head_dim, dtype=dtypes.fp32, device="cuda" + ) + for position in range(query_length): + # MTP position `p` sees context_length - (query_length - 1) + p tokens. + visible = context_length - (query_length - 1) + position + scores = query[position].float() @ keys[:visible].T * head_dim**-0.5 + reference[position] = torch.softmax(scores, dim=-1) @ values[:visible] + + output = torch.empty_like(query) + torch.ops.aiter.pa_decode_flydsl( + output, + query, + key_cache, + value_cache, + context_lengths, + block_tables, + head_dim**-0.5, + query_length, + 1, + 256, + quant_dtype, + None, + key_scale, + value_scale, + ) + return output.float(), reference + + +def _assert_matches(output, reference, tolerance=0.06): + assert torch.isfinite(output).all(), ( + f"{int((~torch.isfinite(output)).sum())} of {output.numel()} outputs " + "are NaN or inf" + ) + scale = float(reference.abs().max()) or 1.0 + error = float((output - reference).abs().max()) / scale + assert error < tolerance, f"relative error {error:.3%} exceeds {tolerance:.1%}" + + +@pytest.mark.parametrize("query_length", [1, 4]) +@pytest.mark.parametrize("per_token", [False, True]) +def test_zero_query_stays_finite(query_length, per_token): + """An all-zero query row quantises to scale 0; -inf * 0 must not reach exp2.""" + _require_gpu() + output, reference = _adversarial_case( + query_length=query_length, per_token=per_token, zero_query=True + ) + _assert_matches(output, reference) + + +@pytest.mark.parametrize("block_size", [16, 64]) +@pytest.mark.parametrize("context_length", [1, 257, 1027]) +def test_block_table_padding_is_not_dereferenced(block_size, context_length): + """Entries past the sequence's extent must resolve to block 0. + + A 256-token tile always walks 256/block_size pages, so it reads block-table + entries the sequence does not own. Those hold whatever the caller left -- + often a stale id pointing at another sequence's live page. The tokens behind + them get a zero probability, but the PV matmul still multiplies their V + bytes (``0 * NaN == NaN``), so the page index itself has to be pinned. + + Not covered here, by design: the unwritten tail of the last page the + sequence does own. Those slots are inside a real block and the caller is + expected to leave them finite. + """ + _require_gpu() + output, reference = _adversarial_case( + context_length=context_length, + block_size=block_size, + poison_padding_blocks=True, + ) + _assert_matches(output, reference) + + +@pytest.mark.parametrize("query_length", [1, 4]) +@pytest.mark.parametrize("context_length", [1, 257]) +def test_extreme_tail_value_scale_is_ignored(context_length, query_length): + """A huge scale on an unwritten slot must not underflow valid probabilities. + + Only finite tails are covered: scale tensors are allocated zeroed, so a NaN + scale past ``context_len`` is not a case the kernel has to survive. + """ + _require_gpu() + if context_length < query_length: + pytest.skip("MTP position 0 would see a non-positive causal bound") + output, reference = _adversarial_case( + context_length=context_length, + query_length=query_length, + per_token=True, + tail_value_scale=1e9, + ) + _assert_matches(output, reference) + + +@pytest.mark.parametrize("head_dim", [64, 128, 256]) +def test_supported_head_dims_are_accurate(head_dim): + _require_gpu() + output, reference = _adversarial_case(head_dim=head_dim, context_length=1027) + _assert_matches(output, reference) + + +@pytest.mark.parametrize("head_dim", [192, 320]) +def test_unsupported_head_dim_is_rejected(head_dim): + """head_dim//16 above 8 and not a multiple of 8 leaves the Q tail unloaded.""" + _require_gpu() + with pytest.raises(NotImplementedError, match="head_dim"): + _adversarial_case(head_dim=head_dim, context_length=64) + + def main(): + torch.set_default_device("cuda") if not torch.cuda.is_available(): aiter.logger.warning("ROCm is not available; skipping pa_decode") return From e65d6bae4ce4a709c418b5161fa931e904927159 Mon Sep 17 00:00:00 2001 From: fsx950223 Date: Tue, 8 Sep 2026 07:06:35 +0000 Subject: [PATCH 11/66] refactor(pa_decode): move FlyDSL reduce kernel --- aiter/ops/flydsl/pa_decode.py | 381 +++++++++++++- aiter/ops/triton/gluon/pa_decode_gluon.py | 588 +--------------------- op_tests/test_flydsl_pa_decode.py | 31 +- 3 files changed, 387 insertions(+), 613 deletions(-) diff --git a/aiter/ops/flydsl/pa_decode.py b/aiter/ops/flydsl/pa_decode.py index 90680dadd8e..34f62fcb967 100644 --- a/aiter/ops/flydsl/pa_decode.py +++ b/aiter/ops/flydsl/pa_decode.py @@ -26,12 +26,17 @@ an LDS round-trip on P transposing ownership between the two MMAs. """ +from functools import lru_cache + +import flydsl.compiler as flyc +import flydsl.expr as fx import torch +from flydsl.expr.typing import T from aiter.jit.utils.chip_info import get_gfx_runtime from .kernels.pa_decode_tile import KV_COMPUTE_BLOCK, compile_pa_decode_tile -from .kernels.tensor_shim import _run_compiled +from .kernels.tensor_shim import _run_compiled, ptr_arg from .kernels.utils import cdiv @@ -47,6 +52,332 @@ def get_recommended_splits( return max(4, min(n, 8)) +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}") + + +def _flydsl_pointer_dtype(dtype: torch.dtype): + return { + torch.float32: fx.Float32, + torch.float16: fx.Float16, + torch.bfloat16: fx.BFloat16, + }[dtype] + + +@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, +): + """Build the partitioned-softmax reduction used by ``pa_decode``. + + Each wave reduces the partition statistics independently, then reuses the + per-partition weights to combine one output element per thread. The host + API limits the partition count to one wave, so the LDS path from the more + general Gluon implementation is unnecessary here. + """ + if not 1 <= max_context_partition_num <= 64: + raise ValueError( + "max_context_partition_num must be in [1, 64], " + 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}" + ) + + dtype_map = { + "f32": fx.Float32, + "f16": fx.Float16, + "bf16": fx.BFloat16, + } + try: + output_dtype = dtype_map[output_dtype_str] + logits_dtype = dtype_map[logits_dtype_str] + sink_dtype = dtype_map[sink_dtype_str] + except KeyError as exc: + raise ValueError(f"Unsupported FlyDSL dtype: {exc.args[0]!r}") from exc + + 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 + ] + + @flyc.kernel(known_block_size=[head_size, 1, 1]) + 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, + ): + tid = fx.thread_idx.x + 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) + c_reduce_width = fx.Int32(reduce_width) + c_four = fx.Int32(4) + c_qgs = query_group_size + lane = tid & c_wave_mask + group_idx = eqgs_idx % c_qgs + + def _wave_reduce_max(value): + reduced = value + for offset in reduce_shuffle_offsets: + reduced = reduced.maximumf( + reduced.shuffle_xor(fx.Int32(offset), c_warp_size) + ) + return reduced + + def _wave_reduce_sum(value): + reduced = value + for offset in reduce_shuffle_offsets: + reduced = fx.Float32( + reduced.addf( + reduced.shuffle_xor(fx.Int32(offset), c_warp_size), + fastmath="fast", + ) + ) + return reduced + + 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 = ( + batch_idx * stride_exp_sums_seq + + 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) + + global_max = _wave_reduce_max(part_max) + 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_value = fx.Float32(sink_token[kv_head_idx * c_qgs + group_idx]) + sink_scale = (global_max > neg_inf).select( + fx.exp2( + (sink_value - safe_global_max) * c_log2e, + fastmath="fast", + ), + zero_f, + ) + global_exp_sum = global_exp_sum + sink_scale + safe_global_exp_sum = (global_exp_sum > zero_f).select(global_exp_sum, one_f) + weight_local = scaled_sum / safe_global_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 = ( + batch_idx * stride_logits_seq + + 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 + + query_idx = eqgs_idx // c_qgs + group_idx = eqgs_idx % c_qgs + 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 + ) + 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, + 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, + ).launch( + grid=(batch_size, num_kv_heads, query_seq_len * query_group_size), + block=(head_size, 1, 1), + stream=stream, + ) + + return { + "launch": launch_pa_decode_ps_reduce_kernel, + "kernel": pa_decode_ps_reduce_kernel, + } + + +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, +) -> None: + if context_partition_num > 64: + raise ImportError("FlyDSL pa_decode reduce supports at most 64 partitions") + use_sinks = sink_token is not None + compiled = compile_pa_decode_ps_reduce( + max_context_partition_num=context_partition_num, + head_size=head_size, + output_dtype_str=_flydsl_dtype_str(output.dtype), + logits_dtype_str=_flydsl_dtype_str(logits.dtype), + sink_dtype_str=_flydsl_dtype_str( + output.dtype if sink_token is None else sink_token.dtype + ), + use_sinks=use_sinks, + ) + 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], + 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] @@ -358,31 +689,29 @@ def normalize_scale(scale, name): s, ) if num_partitions > 1: - from csrc.cpp_itfs.pa.pa_ps import launch_pa_decode_ps_reduce - output_5d = output.reshape( num_seqs, query_length, num_kv_heads, query_group_size, head_dim ) - with torch.cuda.stream(s): - launch_pa_decode_ps_reduce( - output_5d, - psum, - pmax, - pout, - None, - output_5d.stride(0), - output_5d.stride(1), - output_5d.stride(2), - output_5d.stride(3), - pmax.stride(0), - pmax.stride(1), - pmax.stride(2), - pout.stride(0), - pout.stride(1), - pout.stride(2), - pout.stride(3), - query_seq_len=query_length, - query_group_size=query_group_size, - head_size=head_dim, - context_partition_num=num_partitions, - ) + launch_pa_decode_ps_reduce( + output_5d, + psum, + pmax, + pout, + None, + output_5d.stride(0), + output_5d.stride(1), + output_5d.stride(2), + output_5d.stride(3), + pmax.stride(0), + pmax.stride(1), + pmax.stride(2), + pout.stride(0), + pout.stride(1), + pout.stride(2), + pout.stride(3), + query_seq_len=query_length, + query_group_size=query_group_size, + head_size=head_dim, + context_partition_num=num_partitions, + stream=s, + ) diff --git a/aiter/ops/triton/gluon/pa_decode_gluon.py b/aiter/ops/triton/gluon/pa_decode_gluon.py index 2f844821904..30acce1fe21 100644 --- a/aiter/ops/triton/gluon/pa_decode_gluon.py +++ b/aiter/ops/triton/gluon/pa_decode_gluon.py @@ -22,34 +22,12 @@ 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.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 GLUON_JIT_KERNEL_ENABLED = True try: @@ -4494,558 +4472,6 @@ 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] @@ -5110,6 +4536,8 @@ def _paged_attention_decode_v2_reduce_kernel_wrapper( except ImportError: pass try: + if not FLYDSL_PS_REDUCE_AVAILABLE: + raise ImportError("FlyDSL pa_decode reduce is unavailable") launch_pa_decode_ps_reduce_flydsl( output_ptr, exp_sums_ptr, @@ -5131,9 +4559,7 @@ 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: diff --git a/op_tests/test_flydsl_pa_decode.py b/op_tests/test_flydsl_pa_decode.py index 1ed19e97eeb..a264976f73d 100644 --- a/op_tests/test_flydsl_pa_decode.py +++ b/op_tests/test_flydsl_pa_decode.py @@ -88,9 +88,15 @@ def test_pa_decode_maps_gluon_buffers_and_scale_layout(monkeypatch): pytest.skip("FlyDSL is not available") pa_decode_module = importlib.import_module("aiter.ops.flydsl.pa_decode") + gluon_module = importlib.import_module("aiter.ops.triton.gluon.pa_decode_gluon") attention_module = importlib.import_module("aiter.ops.attention") captured = {} + assert ( + gluon_module.launch_pa_decode_ps_reduce_flydsl + is pa_decode_module.launch_pa_decode_ps_reduce + ) + monkeypatch.setattr( pa_decode_module, "compile_pa_decode_tile", @@ -121,12 +127,19 @@ def capture_launch( ) monkeypatch.setattr(pa_decode_module, "_run_compiled", capture_launch) - pa_ps_module = importlib.import_module("csrc.cpp_itfs.pa.pa_ps") - monkeypatch.setattr( - pa_ps_module, - "launch_pa_decode_ps_reduce", - lambda *args, **kwargs: None, - ) + + def capture_reduce( + output_5d, reduce_exp_sums, reduce_max_logits, logits, *args, **kwargs + ): + captured.update( + reduce_output=output_5d, + reduce_exp_sums=reduce_exp_sums, + reduce_max_logits=reduce_max_logits, + reduce_logits=logits, + reduce_kwargs=kwargs, + ) + + monkeypatch.setattr(pa_decode_module, "launch_pa_decode_ps_reduce", capture_reduce) query = torch.empty(1, 8, 128, dtype=torch.bfloat16) output = torch.empty_like(query) @@ -163,6 +176,11 @@ def capture_launch( assert captured["max_logits"].data_ptr() == max_logits.data_ptr() assert captured["exp_sums"].data_ptr() == exp_sums.data_ptr() assert captured["temporary_output"].data_ptr() == temporary_output.data_ptr() + assert captured["reduce_output"].data_ptr() == output.data_ptr() + assert captured["reduce_exp_sums"].data_ptr() == exp_sums.data_ptr() + assert captured["reduce_max_logits"].data_ptr() == max_logits.data_ptr() + assert captured["reduce_logits"].data_ptr() == temporary_output.data_ptr() + assert captured["reduce_kwargs"]["context_partition_num"] == 2 dispatches = [] monkeypatch.setattr( @@ -429,6 +447,7 @@ def test_pa_decode_tile(block_size): block_size=block_size, dtype=dtypes.bf16, ) + assert result["partitions"] > 1 assert result["flydsl err"] == 0 From 7bac8689fdcb1d1ff39b72fae6b91be09519e6de Mon Sep 17 00:00:00 2001 From: fsx950223 Date: Tue, 8 Sep 2026 07:42:15 +0000 Subject: [PATCH 12/66] refactor(flydsl): address pa decode review feedback --- aiter/ops/flydsl/kernels/pa_decode_reduce.py | 250 +++++++++++++++++++ aiter/ops/flydsl/kernels/pa_decode_tile.py | 68 ++--- aiter/ops/flydsl/kernels/tensor_shim.py | 34 ++- aiter/ops/flydsl/pa_decode.py | 244 +----------------- 4 files changed, 310 insertions(+), 286 deletions(-) create mode 100644 aiter/ops/flydsl/kernels/pa_decode_reduce.py 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..c358dd8b8bf --- /dev/null +++ b/aiter/ops/flydsl/kernels/pa_decode_reduce.py @@ -0,0 +1,250 @@ +# 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 + + +@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, +): + """Build the partitioned-softmax reduction used by ``pa_decode``. + + Each wave reduces the partition statistics independently, then reuses the + per-partition weights to combine one output element per thread. The host + API limits the partition count to one wave, so the LDS path from the more + general Gluon implementation is unnecessary here. + """ + if not 1 <= max_context_partition_num <= 64: + raise ValueError( + "max_context_partition_num must be in [1, 64], " + 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}" + ) + + dtype_map = { + "f32": fx.Float32, + "f16": fx.Float16, + "bf16": fx.BFloat16, + } + try: + output_dtype = dtype_map[output_dtype_str] + logits_dtype = dtype_map[logits_dtype_str] + sink_dtype = dtype_map[sink_dtype_str] + except KeyError as exc: + raise ValueError(f"Unsupported FlyDSL dtype: {exc.args[0]!r}") from exc + + 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 + ] + + @flyc.kernel(known_block_size=[head_size, 1, 1]) + 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, + ): + tid = fx.thread_idx.x + 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) + c_reduce_width = fx.Int32(reduce_width) + c_four = fx.Int32(4) + c_qgs = query_group_size + lane = tid & c_wave_mask + group_idx = eqgs_idx % c_qgs + + def _wave_reduce_max(value): + reduced = value + for offset in reduce_shuffle_offsets: + reduced = reduced.maximumf( + reduced.shuffle_xor(fx.Int32(offset), c_warp_size) + ) + return reduced + + def _wave_reduce_sum(value): + reduced = value + for offset in reduce_shuffle_offsets: + reduced = fx.Float32( + reduced.addf( + reduced.shuffle_xor(fx.Int32(offset), c_warp_size), + fastmath="fast", + ) + ) + return reduced + + 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 = ( + batch_idx * stride_exp_sums_seq + + 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) + + global_max = _wave_reduce_max(part_max) + 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_value = fx.Float32(sink_token[kv_head_idx * c_qgs + group_idx]) + sink_scale = (global_max > neg_inf).select( + fx.exp2( + (sink_value - safe_global_max) * c_log2e, + fastmath="fast", + ), + zero_f, + ) + global_exp_sum = global_exp_sum + sink_scale + safe_global_exp_sum = (global_exp_sum > zero_f).select(global_exp_sum, one_f) + weight_local = scaled_sum / safe_global_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 = ( + batch_idx * stride_logits_seq + + 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 + + query_idx = eqgs_idx // c_qgs + group_idx = eqgs_idx % c_qgs + 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 + ) + 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, + 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, + ).launch( + grid=(batch_size, num_kv_heads, query_seq_len * query_group_size), + block=(head_size, 1, 1), + stream=stream, + ) + + return { + "launch": launch_pa_decode_ps_reduce_kernel, + "kernel": pa_decode_ps_reduce_kernel, + } diff --git a/aiter/ops/flydsl/kernels/pa_decode_tile.py b/aiter/ops/flydsl/kernels/pa_decode_tile.py index 8d355b0a864..7289a33b723 100644 --- a/aiter/ops/flydsl/kernels/pa_decode_tile.py +++ b/aiter/ops/flydsl/kernels/pa_decode_tile.py @@ -37,7 +37,8 @@ from flydsl.expr.typing import ReductionOp, T from flydsl.runtime.device import get_rocm_arch -from . import buffer_ops, dpp_utils +from . import dpp_utils +from .tensor_shim import ptr_buf_scalar from .utils import cdiv, exp2_amdgcn_scalar, exp2_f32_fast, rcp_f32 MFMA_MNK = ( @@ -294,24 +295,16 @@ def _v_load16(byte_off): bt_num_records_bytes = ( fx.Index(gpu.grid_dim.x) * fx.Index(max_blocks_per_seq) * 4 ) # int32 entries - bt_rsrc = buffer_ops.create_buffer_resource( - block_tables_ptr, max_size=False, num_records_bytes=bt_num_records_bytes + bt_load = ptr_buf_scalar( + block_tables_ptr, num_records_bytes=bt_num_records_bytes ) - ks_rsrc = buffer_ops.create_buffer_resource(key_scale_ptr, max_size=True) - vs_rsrc = buffer_ops.create_buffer_resource(value_scale_ptr, max_size=True) # Per-tensor: a single global scale, read once. Per-token: read # per-token instead (see _kv_scale_ops/_stage_kv_scale_to_lds below). if const_expr(not per_token_kv): - key_scale = fx.Int32( - buffer_ops.buffer_load( - ks_rsrc, fx.Int32(0), vec_width=1, is_scalar=True - ) - ).bitcast(fx.Float32) - value_scale = fx.Int32( - buffer_ops.buffer_load( - vs_rsrc, fx.Int32(0), vec_width=1, is_scalar=True - ) - ).bitcast(fx.Float32) + key_scale_load = ptr_buf_scalar(key_scale_ptr) + value_scale_load = ptr_buf_scalar(value_scale_ptr) + key_scale = fx.Int32(key_scale_load()).bitcast(fx.Float32) + value_scale = fx.Int32(value_scale_load()).bitcast(fx.Float32) num_tiles = cdiv(context_len, TILE_TOK) num_pages = cdiv(context_len, block_size) # pages this sequence really owns @@ -350,18 +343,27 @@ def _lds_store(byte_off, elem_ty, vec): TOK_CHUNK = NWARP * MFMA_MNK # 64 NCHUNK = TILE_TOK // TOK_CHUNK # 4 + if const_expr(per_token_kv): + scale_load_width = NCHUNK if block_size == 64 else 1 + scale_copy_op = ( + fx.rocdl.BufferCopy128b() + if block_size == 64 + else fx.rocdl.BufferCopy32b() + ) + _k_scale_load = _make_flat_loader( + key_scale_ptr, fx.Float32, scale_load_width, scale_copy_op + ) + _v_scale_load = _make_flat_loader( + value_scale_ptr, fx.Float32, scale_load_width, scale_copy_op + ) + def _load_phys_scalar(page, vec_width=1): # Past `num_pages` the block-table entry is padding: whatever the # caller left there, often a stale id pointing at another sequence's # live page. Pin those to block 0 so the tokens the softmax masks # out always resolve to one known, in-bounds page instead. `page` is # wave-uniform, so the clamp is scalar. - result = buffer_ops.buffer_load( - bt_rsrc, - seq * max_blocks_per_seq + page, - vec_width=vec_width, - is_scalar=True, - ) + result = bt_load(seq * max_blocks_per_seq + page, vec_width) if const_expr(vec_width == 1): return (page < num_pages).select(fx.Int32(result), fx.Int32(0)) loaded = fx.Vector(result) @@ -405,16 +407,8 @@ def _stage_kv_scale_to_lds(phys_vec, buf_off=0): phys = fx.Int32(phys_vec[0]) base_tok = lane16 * NCHUNK scale_idx = phys * stride_ks_block + kv_h * stride_ks_head + base_tok - k_scale_vec = fx.Vector( - buffer_ops.buffer_load( - ks_rsrc, scale_idx, vec_width=NCHUNK, dtype=fx.Float32 - ) - ) - v_scale_vec = fx.Vector( - buffer_ops.buffer_load( - vs_rsrc, scale_idx, vec_width=NCHUNK, dtype=fx.Float32 - ) - ) + k_scale_vec = _k_scale_load(scale_idx) + v_scale_vec = _v_scale_load(scale_idx) slot = (warp * TOK_PER_WARP + base_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) @@ -423,16 +417,8 @@ def _stage_kv_scale_to_lds(phys_vec, buf_off=0): # so the 4 sub-blocks stage in parallel across rgroup-groups. 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( - buffer_ops.buffer_load( - ks_rsrc, scale_idx, vec_width=1, dtype=fx.Float32 - ) - ) - v_scale_scalar = fx.Float32( - buffer_ops.buffer_load( - vs_rsrc, scale_idx, vec_width=1, dtype=fx.Float32 - ) - ) + 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 * c16 + lane16) * f32 _lds_store( diff --git a/aiter/ops/flydsl/kernels/tensor_shim.py b/aiter/ops/flydsl/kernels/tensor_shim.py index 991575d870a..fea268bb4ee 100644 --- a/aiter/ops/flydsl/kernels/tensor_shim.py +++ b/aiter/ops/flydsl/kernels/tensor_shim.py @@ -53,7 +53,7 @@ BUF_VIEW_MAX_ELEMS = 0xFFFFFFFF -def ptr_rsrc(ptr): +def ptr_rsrc(ptr, num_records_bytes=None): """Convert an fx.Pointer kernel arg to a buffer resource for buffer_load/store. Kept for kernels still on the raw `buffer_ops` path; migrated kernels should @@ -61,7 +61,37 @@ def ptr_rsrc(ptr): """ from aiter.ops.flydsl.kernels import buffer_ops - return buffer_ops.create_buffer_resource_from_addr(fx.Int64(ptrtoint(ptr))) + return buffer_ops.create_buffer_resource_from_addr( + fx.Int64(ptrtoint(ptr)), num_records_bytes=num_records_bytes + ) + + +def ptr_buf_scalar(ptr, num_records_bytes=None): + """Return a wave-uniform 32-bit ``s.buffer.load`` accessor. + + The returned callable accepts an element offset and a width of one or four + dwords. Values are returned as raw i32 bits so callers can either consume + integer data directly or bitcast floating-point data without routing the + load through per-lane VGPR addressing. ``ptr`` may be either an opaque + ``fx.Pointer`` or a shaped ``fx.Tensor`` kernel argument. + """ + from aiter.ops.flydsl.kernels import buffer_ops + + if isinstance(ptr, fx.Pointer): + rsrc = ptr_rsrc(ptr, num_records_bytes=num_records_bytes) + else: + rsrc = buffer_ops.create_buffer_resource( + ptr, + max_size=num_records_bytes is None, + num_records_bytes=num_records_bytes, + ) + + def load(offset=0, vec_width=1): + return buffer_ops.buffer_load( + rsrc, offset, vec_width=vec_width, is_scalar=True + ) + + return load def buf_base_i64(base): diff --git a/aiter/ops/flydsl/pa_decode.py b/aiter/ops/flydsl/pa_decode.py index 34f62fcb967..b977cc8916d 100644 --- a/aiter/ops/flydsl/pa_decode.py +++ b/aiter/ops/flydsl/pa_decode.py @@ -26,15 +26,13 @@ an LDS round-trip on P transposing ownership between the two MMAs. """ -from functools import lru_cache - import flydsl.compiler as flyc import flydsl.expr as fx import torch -from flydsl.expr.typing import T from aiter.jit.utils.chip_info import get_gfx_runtime +from .kernels.pa_decode_reduce import compile_pa_decode_ps_reduce from .kernels.pa_decode_tile import KV_COMPUTE_BLOCK, compile_pa_decode_tile from .kernels.tensor_shim import _run_compiled, ptr_arg from .kernels.utils import cdiv @@ -70,246 +68,6 @@ def _flydsl_pointer_dtype(dtype: torch.dtype): }[dtype] -@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, -): - """Build the partitioned-softmax reduction used by ``pa_decode``. - - Each wave reduces the partition statistics independently, then reuses the - per-partition weights to combine one output element per thread. The host - API limits the partition count to one wave, so the LDS path from the more - general Gluon implementation is unnecessary here. - """ - if not 1 <= max_context_partition_num <= 64: - raise ValueError( - "max_context_partition_num must be in [1, 64], " - 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}" - ) - - dtype_map = { - "f32": fx.Float32, - "f16": fx.Float16, - "bf16": fx.BFloat16, - } - try: - output_dtype = dtype_map[output_dtype_str] - logits_dtype = dtype_map[logits_dtype_str] - sink_dtype = dtype_map[sink_dtype_str] - except KeyError as exc: - raise ValueError(f"Unsupported FlyDSL dtype: {exc.args[0]!r}") from exc - - 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 - ] - - @flyc.kernel(known_block_size=[head_size, 1, 1]) - 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, - ): - tid = fx.thread_idx.x - 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) - c_reduce_width = fx.Int32(reduce_width) - c_four = fx.Int32(4) - c_qgs = query_group_size - lane = tid & c_wave_mask - group_idx = eqgs_idx % c_qgs - - def _wave_reduce_max(value): - reduced = value - for offset in reduce_shuffle_offsets: - reduced = reduced.maximumf( - reduced.shuffle_xor(fx.Int32(offset), c_warp_size) - ) - return reduced - - def _wave_reduce_sum(value): - reduced = value - for offset in reduce_shuffle_offsets: - reduced = fx.Float32( - reduced.addf( - reduced.shuffle_xor(fx.Int32(offset), c_warp_size), - fastmath="fast", - ) - ) - return reduced - - 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 = ( - batch_idx * stride_exp_sums_seq - + 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) - - global_max = _wave_reduce_max(part_max) - 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_value = fx.Float32(sink_token[kv_head_idx * c_qgs + group_idx]) - sink_scale = (global_max > neg_inf).select( - fx.exp2( - (sink_value - safe_global_max) * c_log2e, - fastmath="fast", - ), - zero_f, - ) - global_exp_sum = global_exp_sum + sink_scale - safe_global_exp_sum = (global_exp_sum > zero_f).select(global_exp_sum, one_f) - weight_local = scaled_sum / safe_global_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 = ( - batch_idx * stride_logits_seq - + 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 - - query_idx = eqgs_idx // c_qgs - group_idx = eqgs_idx % c_qgs - 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 - ) - 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, - 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, - ).launch( - grid=(batch_size, num_kv_heads, query_seq_len * query_group_size), - block=(head_size, 1, 1), - stream=stream, - ) - - return { - "launch": launch_pa_decode_ps_reduce_kernel, - "kernel": pa_decode_ps_reduce_kernel, - } - - def launch_pa_decode_ps_reduce( output: torch.Tensor, exp_sums: torch.Tensor, From 8de2af2011b67f4b71063ab0d20e22fceb564b5a Mon Sep 17 00:00:00 2001 From: fsx950223 Date: Tue, 8 Sep 2026 07:44:46 +0000 Subject: [PATCH 13/66] style(flydsl): format scalar buffer helper --- aiter/ops/flydsl/kernels/tensor_shim.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/aiter/ops/flydsl/kernels/tensor_shim.py b/aiter/ops/flydsl/kernels/tensor_shim.py index fea268bb4ee..252e80a5f64 100644 --- a/aiter/ops/flydsl/kernels/tensor_shim.py +++ b/aiter/ops/flydsl/kernels/tensor_shim.py @@ -87,9 +87,7 @@ def ptr_buf_scalar(ptr, num_records_bytes=None): ) def load(offset=0, vec_width=1): - return buffer_ops.buffer_load( - rsrc, offset, vec_width=vec_width, is_scalar=True - ) + return buffer_ops.buffer_load(rsrc, offset, vec_width=vec_width, is_scalar=True) return load From 5b232b136a6f0d85c39d01cbf616cc9337967a96 Mon Sep 17 00:00:00 2001 From: fsx950223 Date: Tue, 8 Sep 2026 08:13:43 +0000 Subject: [PATCH 14/66] refactor(flydsl): use buffer tensor helper in pa decode --- aiter/ops/flydsl/kernels/pa_decode_tile.py | 44 +++++++++++++--------- aiter/ops/flydsl/kernels/tensor_shim.py | 32 +--------------- 2 files changed, 29 insertions(+), 47 deletions(-) diff --git a/aiter/ops/flydsl/kernels/pa_decode_tile.py b/aiter/ops/flydsl/kernels/pa_decode_tile.py index 7289a33b723..441139b1965 100644 --- a/aiter/ops/flydsl/kernels/pa_decode_tile.py +++ b/aiter/ops/flydsl/kernels/pa_decode_tile.py @@ -38,7 +38,7 @@ from flydsl.runtime.device import get_rocm_arch from . import dpp_utils -from .tensor_shim import ptr_buf_scalar +from .tensor_shim import ptr_buf_tensor from .utils import cdiv, exp2_amdgcn_scalar, exp2_f32_fast, rcp_f32 MFMA_MNK = ( @@ -237,13 +237,13 @@ def _make_flat_loader(tensor_ptr, elem_ty, reg_width, copy_op, extent=1 << 30): copy_op, fx.rocdl.CopyOpCDNA3BufferCopyType ) copy_atom = fx.make_copy_atom(copy_op, elem_ty) + if const_expr(use_buffer_resource): + flat = ptr_buf_tensor(tensor_ptr, elem_ty, n=extent) + else: + flat = fx.Tensor( + fx.make_view(fx.get_iter(tensor_ptr), fx.make_layout(extent, 1)) + ) reg = fx.make_rmem_tensor(fx.make_layout(reg_width, 1), elem_ty) - base_iter = ( - fx.get_iter(fx.rocdl.make_buffer_tensor(tensor_ptr, max_size=True)) - if use_buffer_resource - else fx.get_iter(tensor_ptr) - ) - flat = fx.Tensor(fx.make_view(base_iter, fx.make_layout(extent, 1))) div = fx.logical_divide(flat, fx.make_layout(1, 1)) def _load(elem_idx): @@ -295,16 +295,19 @@ def _v_load16(byte_off): bt_num_records_bytes = ( fx.Index(gpu.grid_dim.x) * fx.Index(max_blocks_per_seq) * 4 ) # int32 entries - bt_load = ptr_buf_scalar( - block_tables_ptr, num_records_bytes=bt_num_records_bytes + bt_buf = ptr_buf_tensor( + block_tables_ptr, + fx.Int32, + unit_elems=PAGES_PER_CHUNK, + num_records_bytes=bt_num_records_bytes, ) # Per-tensor: a single global scale, read once. Per-token: read # per-token instead (see _kv_scale_ops/_stage_kv_scale_to_lds below). if const_expr(not per_token_kv): - key_scale_load = ptr_buf_scalar(key_scale_ptr) - value_scale_load = ptr_buf_scalar(value_scale_ptr) - key_scale = fx.Int32(key_scale_load()).bitcast(fx.Float32) - value_scale = fx.Int32(value_scale_load()).bitcast(fx.Float32) + 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_tiles = cdiv(context_len, TILE_TOK) num_pages = cdiv(context_len, block_size) # pages this sequence really owns @@ -361,12 +364,19 @@ def _load_phys_scalar(page, vec_width=1): # Past `num_pages` the block-table entry is padding: whatever the # caller left there, often a stale id pointing at another sequence's # live page. Pin those to block 0 so the tokens the softmax masks - # out always resolve to one known, in-bounds page instead. `page` is - # wave-uniform, so the clamp is scalar. - result = bt_load(seq * max_blocks_per_seq + page, vec_width) + # out always resolve to one known, in-bounds page instead. + 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)) - loaded = fx.Vector(result) + 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 // fx.Int32(vec_width), 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)) diff --git a/aiter/ops/flydsl/kernels/tensor_shim.py b/aiter/ops/flydsl/kernels/tensor_shim.py index 252e80a5f64..991575d870a 100644 --- a/aiter/ops/flydsl/kernels/tensor_shim.py +++ b/aiter/ops/flydsl/kernels/tensor_shim.py @@ -53,7 +53,7 @@ BUF_VIEW_MAX_ELEMS = 0xFFFFFFFF -def ptr_rsrc(ptr, num_records_bytes=None): +def ptr_rsrc(ptr): """Convert an fx.Pointer kernel arg to a buffer resource for buffer_load/store. Kept for kernels still on the raw `buffer_ops` path; migrated kernels should @@ -61,35 +61,7 @@ def ptr_rsrc(ptr, num_records_bytes=None): """ from aiter.ops.flydsl.kernels import buffer_ops - return buffer_ops.create_buffer_resource_from_addr( - fx.Int64(ptrtoint(ptr)), num_records_bytes=num_records_bytes - ) - - -def ptr_buf_scalar(ptr, num_records_bytes=None): - """Return a wave-uniform 32-bit ``s.buffer.load`` accessor. - - The returned callable accepts an element offset and a width of one or four - dwords. Values are returned as raw i32 bits so callers can either consume - integer data directly or bitcast floating-point data without routing the - load through per-lane VGPR addressing. ``ptr`` may be either an opaque - ``fx.Pointer`` or a shaped ``fx.Tensor`` kernel argument. - """ - from aiter.ops.flydsl.kernels import buffer_ops - - if isinstance(ptr, fx.Pointer): - rsrc = ptr_rsrc(ptr, num_records_bytes=num_records_bytes) - else: - rsrc = buffer_ops.create_buffer_resource( - ptr, - max_size=num_records_bytes is None, - num_records_bytes=num_records_bytes, - ) - - def load(offset=0, vec_width=1): - return buffer_ops.buffer_load(rsrc, offset, vec_width=vec_width, is_scalar=True) - - return load + return buffer_ops.create_buffer_resource_from_addr(fx.Int64(ptrtoint(ptr))) def buf_base_i64(base): From 03d44cb4b2f4f86763119e18769ff40ac030537d Mon Sep 17 00:00:00 2001 From: fsx950223 Date: Tue, 8 Sep 2026 08:24:30 +0000 Subject: [PATCH 15/66] refactor(flydsl): inline pa decode buffer loaders --- aiter/ops/flydsl/kernels/pa_decode_tile.py | 81 ++++++++++++++-------- 1 file changed, 52 insertions(+), 29 deletions(-) diff --git a/aiter/ops/flydsl/kernels/pa_decode_tile.py b/aiter/ops/flydsl/kernels/pa_decode_tile.py index 441139b1965..11a1e8b7da2 100644 --- a/aiter/ops/flydsl/kernels/pa_decode_tile.py +++ b/aiter/ops/flydsl/kernels/pa_decode_tile.py @@ -230,24 +230,17 @@ def pa_decode_tile_kernel( part = fx.Int32(gpu.block_id("z")) # context partition handled by this CTA n_kv = fx.Int32(gpu.grid_dim.y) # num_kv_heads == gridDim.y - # fx.copy-based flat loaders: K/V/context_len over a raw pointer, Q over - # a buffer resource. - def _make_flat_loader(tensor_ptr, elem_ty, reg_width, copy_op, extent=1 << 30): - use_buffer_resource = isinstance( - copy_op, fx.rocdl.CopyOpCDNA3BufferCopyType - ) - copy_atom = fx.make_copy_atom(copy_op, elem_ty) - if const_expr(use_buffer_resource): - flat = ptr_buf_tensor(tensor_ptr, elem_ty, n=extent) - else: - flat = fx.Tensor( - fx.make_view(fx.get_iter(tensor_ptr), fx.make_layout(extent, 1)) - ) + # 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) - div = fx.logical_divide(flat, fx.make_layout(1, 1)) + flat = fx.Tensor( + fx.make_view(fx.get_iter(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(div, (None, elem_idx)), reg) + fx.copy(copy_atom, fx.slice(tiled, (None, elem_idx)), reg) return fx.Vector(fx.memref_load_vec(reg)) return _load @@ -255,12 +248,8 @@ def _load(elem_idx): # A cache above 2 GiB overflows an i32 element offset, so the flat view # and the offsets below widen together. KV_EXTENT = (1 << 42) if wide_kv_addressing else (1 << 30) - _k_load_fp8x16 = _make_flat_loader( - key_cache_ptr, FP8, 16, fx.UniversalCopy128b(), KV_EXTENT - ) - _v_load_fp8x16 = _make_flat_loader( - value_cache_ptr, FP8, 16, fx.UniversalCopy128b(), KV_EXTENT - ) + _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): # `phys * page_elems` is the term that overflows first: it reaches @@ -276,10 +265,23 @@ def _kv_addr(phys, page_elems, rest): _q_copy_op = ( fx.rocdl.BufferCopy128b() if QLOAD_UNIT == 8 else fx.rocdl.BufferCopy64b() ) - _q_load_chunk = _make_flat_loader(query_ptr, Q_DTYPE, QLOAD_UNIT, _q_copy_op) - _ctxlen_load = _make_flat_loader( - context_lengths_ptr, fx.Int32, 1, fx.rocdl.BufferCopy32b() - ) + 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)) + + 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) + + def _ctxlen_load(elem_idx): + fx.copy(ctx_copy_atom, fx.slice(ctx_tiled, (None, elem_idx)), ctx_reg) + return fx.Vector(fx.memref_load_vec(ctx_reg)) def _k_load16(byte_off): return _k_load_fp8x16(byte_off).bitcast(fx.Int64) @@ -353,13 +355,34 @@ def _lds_store(byte_off, elem_ty, vec): if block_size == 64 else fx.rocdl.BufferCopy32b() ) - _k_scale_load = _make_flat_loader( - key_scale_ptr, fx.Float32, scale_load_width, scale_copy_op + 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_load = _make_flat_loader( - value_scale_ptr, fx.Float32, scale_load_width, scale_copy_op + 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): # Past `num_pages` the block-table entry is padding: whatever the # caller left there, often a stale id pointing at another sequence's From e847b801d2f9e0b43aed2c6bb5d29bf75f4d3d15 Mon Sep 17 00:00:00 2001 From: fsx950223 Date: Tue, 8 Sep 2026 09:16:27 +0000 Subject: [PATCH 16/66] refactor(flydsl): use pointer ABI for pa decode tile --- aiter/ops/flydsl/kernels/pa_decode_tile.py | 87 ++++++++++++---------- aiter/ops/flydsl/pa_decode.py | 30 +++++--- op_tests/test_flydsl_pa_decode.py | 1 + 3 files changed, 69 insertions(+), 49 deletions(-) diff --git a/aiter/ops/flydsl/kernels/pa_decode_tile.py b/aiter/ops/flydsl/kernels/pa_decode_tile.py index 11a1e8b7da2..6e5a8598d0a 100644 --- a/aiter/ops/flydsl/kernels/pa_decode_tile.py +++ b/aiter/ops/flydsl/kernels/pa_decode_tile.py @@ -204,21 +204,23 @@ class SharedStorage: @flyc.kernel(known_block_size=(BLOCK_THREADS, 1, 1)) def pa_decode_tile_kernel( - output_ptr: fx.Tensor, # [num_seqs*query_length, num_q_heads, head_dim] (written directly when NP==1) + output_ptr: fx.Pointer, # [num_seqs*query_length, num_q_heads, head_dim] (written directly when NP==1) # per-partition partial outputs (combined by the reduce kernel when NP>1): - pmax_ptr: fx.Tensor, # [num_seqs*query_length, num_kv_heads, num_partitions, query_group_size] row max - psum_ptr: fx.Tensor, # [num_seqs*query_length, num_kv_heads, num_partitions, query_group_size] row sum - pout_ptr: fx.Tensor, # [num_seqs*query_length, num_kv_heads, num_partitions, query_group_size, head_dim] Q_DTYPE, normalized O_p/l_p - query_ptr: fx.Tensor, # [num_seqs*query_length, num_q_heads, head_dim] -- row = seq*query_length + qi (MTP position) - key_cache_ptr: fx.Tensor, # [num_blocks, num_kv_heads, head_dim//16, block_size, 16] (blocked, see module docstring) - value_cache_ptr: fx.Tensor, # [num_blocks, num_kv_heads, block_size//16, head_dim, 16] (blocked, see module docstring) - block_tables_ptr: fx.Tensor, # [num_seqs, max_blocks_per_seq] - context_lengths_ptr: fx.Tensor, # [num_seqs] - key_scale_ptr: fx.Tensor, # [1] per-tensor OR [num_blocks, num_kv_heads, block_size] per-token - value_scale_ptr: fx.Tensor, # same shape as key_scale_ptr + pmax_ptr: fx.Pointer, # [num_seqs*query_length, num_kv_heads, num_partitions, query_group_size] row max + psum_ptr: fx.Pointer, # [num_seqs*query_length, num_kv_heads, num_partitions, query_group_size] row sum + pout_ptr: fx.Pointer, # [num_seqs*query_length, num_kv_heads, num_partitions, query_group_size, head_dim] Q_DTYPE, normalized O_p/l_p + query_ptr: fx.Pointer, # [num_seqs*query_length, num_q_heads, head_dim] -- row = seq*query_length + qi (MTP position) + key_cache_ptr: fx.Pointer, # [num_blocks, num_kv_heads, head_dim//16, block_size, 16] (blocked, see module docstring) + value_cache_ptr: fx.Pointer, # [num_blocks, num_kv_heads, block_size//16, head_dim, 16] (blocked, see module docstring) + block_tables_ptr: fx.Pointer, # [num_seqs, max_blocks_per_seq] + context_lengths_ptr: fx.Pointer, # [num_seqs] + key_scale_ptr: fx.Pointer, # [1] per-tensor OR [num_blocks, num_kv_heads, block_size] per-token + value_scale_ptr: fx.Pointer, # same shape as key_scale_ptr 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, ): @@ -230,12 +232,19 @@ def pa_decode_tile_kernel( part = fx.Int32(gpu.block_id("z")) # context partition handled by this CTA n_kv = fx.Int32(gpu.grid_dim.y) # num_kv_heads == gridDim.y + 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) + # 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.get_iter(tensor_ptr), fx.make_layout(extent, 1)) + 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)) @@ -1232,25 +1241,23 @@ def _lmax_off_m(m): gs_head_e = row - qi_e * query_group_size qh = kv_h * query_group_size + gs_head_e - def _emit(o_norm, sub): + def _emit(o_norm, sub, query_idx, query_head): if const_expr(NP == 1): - out_row = output_ptr[ - seq * query_length + qi_e, qh, None # noqa: B023 - ] - out_chunk = fx.slice( - fx.logical_divide(out_row, fx.make_layout(OP_ELEMS, 1)), - (None, sub), + out_offset = ( + (seq * query_length + query_idx) * stride_o_row + + query_head * stride_o_head + + sub * OP_ELEMS ) - out_chunk.store(o_norm) + fx.ptr_store(o_norm, fx.add_offset(output, out_offset)) else: base = ( (seq * n_kv + kv_h) * NP + part ) * TOTAL_ROWS + row # noqa: B023 - pout_div = fx.logical_divide(pout_ptr, fx.make_layout(OP_ELEMS, 1)) - pout_chunk = fx.slice( - pout_div, (None, base * (head_dim // OP_ELEMS) + sub) + pout_offset = base * head_dim + sub * OP_ELEMS + fx.ptr_store( + o_norm, + fx.add_offset(pout, pout_offset), ) - pout_chunk.store(o_norm) for vh in range_constexpr(VHE_CHUNKS): o_slot = _o_slot(m, vh) @@ -1261,7 +1268,7 @@ def _emit(o_norm, sub): sub = head_base // OP_ELEMS # Guard the partial last tile's out-of-range rows (folded away for full tiles). if row < TOTAL_ROWS: - _emit(o_norm, sub) + _emit(o_norm, sub, qi_e, qh) if const_expr(NP > 1): # noqa: SIM102 if warp == 0 and rgroup == 0: @@ -1270,27 +1277,29 @@ def _emit(o_norm, sub): # Convert the running max from log2 units (scale_qk folds # in LOG2E) to natural-log units: the shared reduce # re-applies LOG2E itself when combining partitions. - pmax_ptr[base] = o_final[_m_slot(m)] * fx.Float32(1.0 / LOG2E) - psum_ptr[base] = l_row + pmax[base] = o_final[_m_slot(m)] * fx.Float32(1.0 / LOG2E) + psum[base] = l_row @flyc.jit def pa_decode_tile_launch( - output: fx.Tensor, - pmax: fx.Tensor, - psum: fx.Tensor, - pout: fx.Tensor, - query: fx.Tensor, - key_cache: fx.Tensor, - value_cache: fx.Tensor, - block_tables: fx.Tensor, - context_lengths: fx.Tensor, - key_scale: fx.Tensor, # [1] per-tensor OR [num_blocks, num_kv_heads, block_size] per-token - value_scale: fx.Tensor, # same shape as key_scale + 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, # [1] per-tensor OR [num_blocks, num_kv_heads, block_size] per-token + value_scale: fx.Pointer, # same shape as key_scale 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, stream: fx.Stream = fx.Stream(None), # noqa: B008 @@ -1314,6 +1323,8 @@ def pa_decode_tile_launch( max_blocks_per_seq, stride_ks_block, stride_ks_head, + stride_o_row, + stride_o_head, stride_q_row, stride_q_head, ).launch( diff --git a/aiter/ops/flydsl/pa_decode.py b/aiter/ops/flydsl/pa_decode.py index b977cc8916d..3b5dc59788d 100644 --- a/aiter/ops/flydsl/pa_decode.py +++ b/aiter/ops/flydsl/pa_decode.py @@ -65,6 +65,9 @@ def _flydsl_pointer_dtype(dtype: torch.dtype): 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] @@ -280,6 +283,9 @@ def pa_decode( assert ( query.stride(2) == 1 ), f"pa_decode requires a contiguous head_dim axis, got strides {query.stride()}" + assert ( + output.stride(2) == 1 + ), f"pa_decode requires a contiguous output head_dim axis, got strides {output.stride()}" dev = query.device for name, tensor in ( @@ -426,22 +432,24 @@ def normalize_scale(scale, name): s = torch.cuda.current_stream(dev) _run_compiled( compiled["launch"], - output, - pmax.view(-1), - psum.view(-1), - pout.view(-1), - query, - key_cache, - value_cache, - block_tables, - context_lengths, - key_scale_t, - value_scale_t, + 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), 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)), s, diff --git a/op_tests/test_flydsl_pa_decode.py b/op_tests/test_flydsl_pa_decode.py index a264976f73d..73cf14f98af 100644 --- a/op_tests/test_flydsl_pa_decode.py +++ b/op_tests/test_flydsl_pa_decode.py @@ -127,6 +127,7 @@ def capture_launch( ) monkeypatch.setattr(pa_decode_module, "_run_compiled", capture_launch) + monkeypatch.setattr(pa_decode_module, "ptr_arg", lambda tensor, dtype=None: tensor) def capture_reduce( output_5d, reduce_exp_sums, reduce_max_logits, logits, *args, **kwargs From 27b6f587fa14b8c0985eb3c0f2a3f5dbcd3c42b4 Mon Sep 17 00:00:00 2001 From: Elton Date: Sat, 12 Sep 2026 19:30:38 +0800 Subject: [PATCH 17/66] [FlyDSL] Fix PA edge cases and add tuned page-128 support (#5474) * fix(flydsl): fix PA page-table edge cases and support page 128 Signed-off-by: Elton * [FlyDSL] Tune page-128 PA load and LDS scheduling Signed-off-by: Elton * [FlyDSL] Align page-128 PA reduction scratch rows Signed-off-by: Elton --------- Signed-off-by: Elton --- aiter/ops/flydsl/kernels/pa_decode_tile.py | 135 +++++--- aiter/ops/flydsl/pa_decode.py | 36 +- op_tests/test_flydsl_pa_decode.py | 366 ++++++++++++++++++++- 3 files changed, 474 insertions(+), 63 deletions(-) diff --git a/aiter/ops/flydsl/kernels/pa_decode_tile.py b/aiter/ops/flydsl/kernels/pa_decode_tile.py index 6e5a8598d0a..847340a3bb8 100644 --- a/aiter/ops/flydsl/kernels/pa_decode_tile.py +++ b/aiter/ops/flydsl/kernels/pa_decode_tile.py @@ -10,7 +10,7 @@ ``key_scale``/``value_scale`` are either a ``[1]`` per-tensor scalar or a ``[num_blocks, num_kv_heads, block_size]`` per-token tensor (chosen by rank). -``block_size`` (16/64) and ``head_dim`` (multiple of 64) are compile-time +``block_size`` (16/64/128) and ``head_dim`` (multiple of 64) are compile-time constants. Layouts are logical, not production's preshuffle. * ``query`` [num_seqs, num_q_heads, head_dim] f16/bf16 (head_dim contiguous) @@ -65,6 +65,7 @@ def compile_pa_decode_tile( query_length: int = 1, trans_v: bool = True, wide_kv_addressing: bool = False, + prefetch_v: bool = False, ): """Build the tile-programming PA-decode kernel + launch wrapper. @@ -101,7 +102,8 @@ def compile_pa_decode_tile( assert block_size in ( 16, 64, - ), f"pa_decode_tile only supports block_size in (16, 64), got {block_size}" + 128, + ), f"pa_decode_tile only supports block_size in (16, 64, 128), got {block_size}" assert query_dtype in ( "f16", "bf16", @@ -116,14 +118,24 @@ def compile_pa_decode_tile( TOTAL_ROWS = query_length * query_group_size M_TILES = cdiv(TOTAL_ROWS, MFMA_MNK) ROWS_PADDED = M_TILES * MFMA_MNK + # The page-128 gfx950 scalar-scale path benefits from earlier V loads + # for decode and alternating P/reduction buffers for three M-tiles. + tune_page128 = ( + is_gfx950 + and head_dim == 128 + and block_size == 128 + and trans_v + and not per_token_kv + ) + EARLY_V = prefetch_v and tune_page128 and M_TILES == 1 + P_BUFFERS = 2 if tune_page128 and M_TILES == 3 else 1 # PV layout: V=A, P=B -> output [head-dim (row), query-row (col=lane16)], # generalized over head_dim via the VHE_CHUNKS loop. NWARP = 4 # 4 waves / CTA TOK_PER_WARP = 64 # tokens each warp owns per compute block (matches production KV_COMPUTE_BLOCK) TILE_TOK = NWARP * TOK_PER_WARP # 256 tokens / compute block - PAGES_PER_CHUNK = ( - TOK_PER_WARP // block_size - ) # pages spanned by one 64-token warp-chunk: 1 (bs=64) or 4 (bs=16) + # A warp owns 64 tokens: four page-16s, one page-64, or half a page-128. + PAGES_PER_CHUNK = cdiv(TOK_PER_WARP, block_size) assert ( head_dim % (NWARP * MFMA_MNK) == 0 ), "head_dim must split across the 4 warps for PV" @@ -171,16 +183,17 @@ def compile_pa_decode_tile( # +16B row padding breaks a 32-bank LDS conflict on the P-pack writes while # keeping the row 16B-aligned for PV's ds_read_b128. SP_ROW_BYTES = TILE_TOK + 16 - sP_bytes = MFMA_MNK * SP_ROW_BYTES # fp8, padded rows + sP_bytes = P_BUFFERS * MFMA_MNK * SP_ROW_BYTES # fp8, padded rows sQscale_off = sP_off + sP_bytes - NWARP_PAD = ( - NWARP + 1 - ) # coprime with 32 banks -> no LDS bank conflict on the cross-warp scratch + # Keep reduction rows 16-byte aligned on the tuned gfx950 page-128 path + # so each cross-wave row can be read as one vector. Other paths retain + # the padding used to avoid LDS bank conflicts. + NWARP_PAD = NWARP if tune_page128 else NWARP + 1 # Phase-split slices sLmax per M-tile so all pass-1 writes share one barrier. sLmax_off = sQscale_off + ROWS_PADDED * f32 sLsum_off = sLmax_off + M_TILES * MFMA_MNK * NWARP_PAD * f32 # V page-index broadcast (V's page depends on `rgroup`, shared across warps). - sVPage_off = sLsum_off + MFMA_MNK * NWARP_PAD * f32 + sVPage_off = sLsum_off + P_BUFFERS * MFMA_MNK * NWARP_PAD * f32 sVPage_bytes = NWARP * PAGES_PER_CHUNK * 4 # i32 # Per-token KV-scale staging (per_token only), double-buffered so the tt+1 # prefetch (into the other buffer) doesn't clobber the current tile's scales. @@ -306,10 +319,12 @@ def _v_load16(byte_off): bt_num_records_bytes = ( fx.Index(gpu.grid_dim.x) * fx.Index(max_blocks_per_seq) * 4 ) # int32 entries + # 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, ) # Per-tensor: a single global scale, read once. Per-token: read @@ -358,10 +373,10 @@ def _lds_store(byte_off, elem_ty, vec): NCHUNK = TILE_TOK // TOK_CHUNK # 4 if const_expr(per_token_kv): - scale_load_width = NCHUNK if block_size == 64 else 1 + scale_load_width = NCHUNK if block_size >= 64 else 1 scale_copy_op = ( fx.rocdl.BufferCopy128b() - if block_size == 64 + if block_size >= 64 else fx.rocdl.BufferCopy32b() ) scale_copy_atom = fx.make_copy_atom(scale_copy_op, fx.Float32) @@ -405,7 +420,7 @@ def _load_phys_scalar(page, vec_width=1): frag = fx.make_fragment_like(fx.slice(bt_buf, (0, None))) fx.copy( bt_copy_atom, - fx.slice(bt_buf, (element_offset // fx.Int32(vec_width), None)), + fx.slice(bt_buf, (element_offset, None)), frag, ) loaded = fx.Vector(fx.memref_load_vec(frag)) @@ -422,7 +437,7 @@ def _v_page_fetch_and_stage(tt_i32): # its rgroup row and broadcasts via LDS (read back by _v_page_read_row). base_page = tt_i32 * TILE_TOK // block_size # tile start is page-aligned fetched = _load_phys_scalar( - base_page + warp * PAGES_PER_CHUNK, PAGES_PER_CHUNK + base_page + (warp * TOK_PER_WARP) // block_size, PAGES_PER_CHUNK ) if lane == 0: fetched_vec = ( @@ -445,13 +460,14 @@ def _kv_buf_off(tt_val): return 0 def _stage_kv_scale_to_lds(phys_vec, buf_off=0): - if const_expr(block_size == 64): + if const_expr(block_size >= 64): phys = fx.Int32(phys_vec[0]) - base_tok = lane16 * NCHUNK - scale_idx = phys * stride_ks_block + kv_h * stride_ks_head + base_tok + chunk_tok = 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 + base_tok) * f32 + 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: @@ -488,7 +504,7 @@ def _load_kv_scale_vecs(a, buf_off=0): # token = warp*TOK_CHUNK + a*c16 + lane16 (the softmax mask and P-pack # write position below must encode this same formula). def _k_ops(phys, a): - within_page_tok = (a * c16 + lane16) % block_size + within_page_tok = (warp * TOK_PER_WARP + a * c16 + lane16) % block_size ops = [] for qkhe in range_constexpr(QKHE_LOOP): he_idx = qkhe * RGROUP_QUARTERS + rgroup @@ -508,7 +524,7 @@ def _k_ops(phys, a): 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 * PAGES_PER_CHUNK, PAGES_PER_CHUNK + base_page + (warp * TOK_PER_WARP) // block_size, PAGES_PER_CHUNK ) phys_vec = ( fx.Vector.from_elements([fx.Int32(fetched)], dtype=fx.Int32) @@ -524,15 +540,22 @@ def _k_ops_flat(tt_i32): return fx.Vector.from_elements(flat, dtype=fx.Int64), phys_vec - # -- prologue: prefetch the first tile's K -- - num_tiles_m1 = num_tiles - 1 - start_safe = (part_start < num_tiles).select(part_start, num_tiles_m1) - k_pf0, phys_vec0 = _k_ops_flat(start_safe) - # V page-index prefetch, issued here too for the same overlap; the - # LDS write is visible after the barrier below. - _v_page_fetch_and_stage(start_safe) - if const_expr(per_token_kv): - _stage_kv_scale_to_lds(phys_vec0, _kv_buf_off(fx.Int32(start_safe))) + # Empty partitions retain neutral state and must not read K/V or the + # block table. In particular, an empty context has no valid first tile. + 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) + # Issue the V page-index prefetch alongside K; the LDS write is + # visible after the barrier below. + _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_kv has no single global key_scale/value_scale: scale_qk # drops the key_scale factor (folded in per-token, see masked_chunks @@ -700,25 +723,32 @@ def _quant_q_row(m, qi, gs_head, q_row_off): # changes the offset formula. `sub`/`step` walk pages/16-token sub-blocks. NVOPS = TILE_TOK // MFMA_K # 8 PV k_steps (256 tokens / K=32) STEPS_PER_PAGE = block_size // 16 + STEPS_PER_CHUNK = min(block_size, TOK_PER_WARP) // 16 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_PAGE): + 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 + step) * head_dim + head_element) + ( + (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 + step * 16, + (kv_h * head_dim + head_element) * block_size + + page_step * 16, ) w = _v_load16(base) if const_expr(block_size == 16): @@ -806,10 +836,10 @@ def _l_slot(m): def _mask_v_scale(vec, a, thr=ctx_thr, zero=zero4_scale): return (_ct[a] < thr).select(vec, zero) - # V data doesn't depend on `m`; hoist once for the phase-split - # instead of reloading per M-tile. + # V is independent of QK/softmax. Load it early for the tuned + # decode path and reuse it across M-tiles in the phase-split path. v_vh_shared = None - if const_expr(M_TILES > 1): + if const_expr(M_TILES > 1 or EARLY_V): v_vh_shared = [ _v_ops(v_page_cur, vh) for vh in range_constexpr(VHE_CHUNKS) ] @@ -947,6 +977,8 @@ def _lmax_off_m(m): ] 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) ] @@ -1002,7 +1034,7 @@ def _lmax_off_m(m): words.append(_f32_to_fp8_words(p_scaled)[0]) p_off0 = ( - sP_off + lane16 * SP_ROW_BYTES + warp * TOK_CHUNK + rgroup * 4 + p_base + lane16 * SP_ROW_BYTES + warp * TOK_CHUNK + rgroup * 4 ) # The NCHUNK P words scatter across the row at stride c16//4 # i32 (the token->fp8-lane interleave the PV ds_read_b128 @@ -1020,13 +1052,13 @@ def _lmax_off_m(m): safe_prev = (m_prev > NEG_INF).select(m_prev, ZERO_F) corr_reg = fx.Float32(exp2_amdgcn_scalar(safe_prev - safe_max)) if rgroup == 0: - _st_lw(sLsum_off, lane16, warp, ls) + _st_lw(lsum_base, lane16, warp, ls) gpu.barrier() - gsum = _ld_lw_row(sLsum_off, lane16).reduce(ReductionOp.ADD) + gsum = _ld_lw_row(lsum_base, lane16).reduce(ReductionOp.ADD) l_new = l_prev * corr_reg + gsum p_ops = _lds_load( - sP_off + lane16 * SP_ROW_BYTES + rgroup * 64, fx.Int64, NVOPS + p_base + lane16 * SP_ROW_BYTES + rgroup * 64, fx.Int64, NVOPS ) corr_b = fx.Vector.from_elements( @@ -1048,11 +1080,15 @@ def _lmax_off_m(m): ).broadcast_to(OP_ELEMS) o_acc[vh] = o_acc[vh] * corr_b + op next_state.extend([*o_acc, m_new, l_new]) - # sP and sLsum are reused by the next M-tile. Synchronize - # all waves after their LDS reads before any wave overwrites - # those regions, then retire this tile's dependency chain. + # A shared P/Lsum slot needs all reads to finish before + # the next M-tile overwrites it. With alternating slots, + # the next tile's write barrier retires those reads before + # that slot is reused. The Phase A barrier protects reuse + # across loop iterations. Keep the scheduler fence to + # bound register liveness across M-tiles. if const_expr(m < M_TILES - 1): - gpu.barrier() + if const_expr(P_BUFFERS == 1): + gpu.barrier() fx.rocdl.sched_barrier(0) else: # M_TILES==1 single tile (m==0 for the _o_slot/_m_slot/_l_slot helpers). @@ -1200,11 +1236,14 @@ def _lmax_off_m(m): corr_b = fx.Vector.from_elements( [corr_reg], dtype=fx.Float32 ).broadcast_to(OP_ELEMS) - # Single tile: batch both vh's V loads upfront (no sibling chain - # to hide the latency behind). - v_vh_batch = [ - _v_ops(v_page_cur, vh) for vh in range_constexpr(VHE_CHUNKS) - ] + # Use the early V loads where enabled; other shapes retain + # the existing batched loads before PV. + if const_expr(EARLY_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) diff --git a/aiter/ops/flydsl/pa_decode.py b/aiter/ops/flydsl/pa_decode.py index 3b5dc59788d..6c7047d1a43 100644 --- a/aiter/ops/flydsl/pa_decode.py +++ b/aiter/ops/flydsl/pa_decode.py @@ -10,7 +10,7 @@ ``key_scale``/``value_scale`` are either a ``[1]`` per-tensor scalar or a ``[num_blocks, num_kv_heads, block_size, 1]`` per-token tensor. -``block_size`` (16/64) and ``head_dim`` (multiple of 64) are compile-time +``block_size`` (16/64/128) and ``head_dim`` (multiple of 64) are compile-time constants. Layouts are logical, not production's preshuffle. * ``query`` [num_seqs, num_q_heads, head_dim] f16/bf16 (head_dim contiguous) @@ -166,7 +166,10 @@ def pa_decode( The call signature and intermediate-buffer layouts follow the shared aiter paged-attention decode API. This kernel currently supports FP8 K/V caches, - BF16/FP16 queries, a 256-token context partition, and block sizes 16/64. + BF16/FP16 queries, a 256-token compute tile, and block sizes 16/64/128. + Sparse attention uses caller-prepared block tables and selected context + lengths. Each independently selected MTP query must have its own table row + and use query_length=1; query_length>1 applies dense causal masking. ALiBi, attention sinks, sliding-window attention, and externally quantized FP8 queries are not supported. """ @@ -183,7 +186,7 @@ def pa_decode( raise NotImplementedError("pa_decode does not support ALiBi") if sinks is not None: raise NotImplementedError("pa_decode does not support attention sinks") - if sliding_window != 0: + if sliding_window not in (0, -1): raise NotImplementedError("pa_decode does not support sliding-window attention") if query_length < 1: raise ValueError(f"query_length must be positive, got {query_length}") @@ -237,7 +240,8 @@ def pa_decode( assert block_size in ( 16, 64, - ), f"pa_decode only supports block_size in (16, 64), got {block_size}" + 128, + ), f"pa_decode only supports block_size in (16, 64, 128), got {block_size}" trans_v = value_cache.dim() == 5 if trans_v: @@ -295,10 +299,9 @@ def pa_decode( ("block_tables", block_tables), ("context_lengths", context_lengths), ): - assert tensor.device == dev, ( - f"{name} must be on the same device as query ({dev}), " - f"got {tensor.device}" - ) + assert ( + tensor.device == dev + ), f"{name} must be on the same device as query ({dev}), got {tensor.device}" for name, tensor in ( ("key_cache", key_cache), ("value_cache", value_cache), @@ -375,6 +378,22 @@ def normalize_scale(scale, name): # if either does. wide_kv_addressing = max(key_cache.numel(), value_cache.numel()) >= 2**31 + # Early V loads help the measured one-to-two-workgroup-per-CU regime. + # Keep other grids on the existing schedule: early loads regress + # short-context decode at larger grid sizes. + prefetch_v = False + if ( + arch == "gfx950" + and head_dim == 128 + and block_size == 128 + and trans_v + and not per_token_kv + and query_length * query_group_size <= 16 + ): + num_cus = torch.cuda.get_device_properties(dev).multi_processor_count + workgroups = num_seqs * num_kv_heads * num_partitions + prefetch_v = num_cus < workgroups <= 2 * num_cus + with torch.cuda.device(dev): compiled = compile_pa_decode_tile( head_dim=head_dim, @@ -387,6 +406,7 @@ def normalize_scale(scale, name): query_length=query_length, trans_v=trans_v, wide_kv_addressing=wide_kv_addressing, + prefetch_v=prefetch_v, ) if num_partitions == 1: diff --git a/op_tests/test_flydsl_pa_decode.py b/op_tests/test_flydsl_pa_decode.py index 73cf14f98af..203c58a032f 100644 --- a/op_tests/test_flydsl_pa_decode.py +++ b/op_tests/test_flydsl_pa_decode.py @@ -37,7 +37,7 @@ def _default_cuda_device(): # Pairwise coverage of the normal-accuracy axes in FlyDSL's PA regression test: # batches {3, 81, 128}, Q/KV heads {(4,1), (8,1), (16,1)}, head dims # {128, 256}, and contexts {1027, 8192}. Keep the original 257-token boundary -# case as well. Both supported block sizes are crossed with every case in main(). +# case as well. All supported block sizes are crossed with every case in main(). DEFAULT_BATCH_SIZES = [3, 81, 128] DEFAULT_SHAPES = [ (8, 1, 128, 257), @@ -430,7 +430,7 @@ def run_pa_decode_tile_case( return ret -@pytest.mark.parametrize("block_size", [16, 64]) +@pytest.mark.parametrize("block_size", [16, 64, 128]) def test_pa_decode_tile(block_size): if not torch.cuda.is_available(): pytest.skip("ROCm is not available") @@ -472,6 +472,8 @@ def _adversarial_case( poison_padding_blocks=False, tail_value_scale=None, seed=0, + trans_v=False, + num_partitions=1, ): """Build one case plus its torch reference over the same dequantised fp8 KV. @@ -533,7 +535,14 @@ def rand(*shape): .permute(0, 1, 3, 2, 4) .contiguous() ) - value_cache = value.permute(0, 1, 3, 2).contiguous() + if trans_v: + value_cache = ( + value.view(blocks, 1, block_size // 16, 16, head_dim) + .permute(0, 1, 2, 4, 3) + .contiguous() + ) + else: + value_cache = value.permute(0, 1, 3, 2).contiguous() block_tables = torch.arange(blocks, dtype=dtypes.i32).reshape(1, blocks) context_lengths = torch.full((1,), context_length, dtype=dtypes.i32) @@ -561,7 +570,7 @@ def rand(*shape): block_tables, head_dim**-0.5, query_length, - 1, + num_partitions, 256, quant_dtype, None, @@ -592,7 +601,7 @@ def test_zero_query_stays_finite(query_length, per_token): _assert_matches(output, reference) -@pytest.mark.parametrize("block_size", [16, 64]) +@pytest.mark.parametrize("block_size", [16, 64, 128]) @pytest.mark.parametrize("context_length", [1, 257, 1027]) def test_block_table_padding_is_not_dereferenced(block_size, context_length): """Entries past the sequence's extent must resolve to block 0. @@ -651,6 +660,349 @@ def test_unsupported_head_dim_is_rejected(head_dim): _adversarial_case(head_dim=head_dim, context_length=64) +def _constant_page_decode_case( + lengths, + block_size, + num_partitions, + block_table_width, + poison_padding=False, + per_token=False, + trans_v=False, +): + """Each sequence's V pages are constant, so its attention output is known.""" + _require_gpu() + batch_size, query_heads, head_dim = len(lengths), 8, 128 + pages_per_sequence = [(length + block_size - 1) // block_size for length in lengths] + page_values = [ + float(seq + 1) + for seq, pages in enumerate(pages_per_sequence) + for _ in range(pages) + ] + num_pages = len(page_values) + quant_dtype = _quant_dtype() + query = torch.ones(batch_size, query_heads, head_dim, dtype=dtypes.bf16) + key_cache = torch.zeros( + num_pages, 1, head_dim // 16, block_size, 16, dtype=quant_dtype + ) + value_cache = ( + torch.tensor(page_values, dtype=dtypes.fp32) + .reshape(num_pages, 1, 1, 1) + .expand(num_pages, 1, head_dim, block_size) + .contiguous() + .to(quant_dtype) + ) + if trans_v: + value_cache = ( + value_cache.view(num_pages, 1, head_dim, block_size // 16, 16) + .permute(0, 1, 3, 2, 4) + .contiguous() + ) + block_tables = torch.full( + (batch_size, block_table_width), + num_pages + 1024 if poison_padding else 0, + dtype=dtypes.i32, + ) + next_page = 0 + for seq, pages in enumerate(pages_per_sequence): + block_tables[seq, :pages] = torch.arange( + next_page, next_page + pages, dtype=dtypes.i32 + ) + next_page += pages + context_lengths = torch.tensor(lengths, dtype=dtypes.i32) + scale = torch.ones( + (num_pages, 1, block_size, 1) if per_token else (1,), dtype=dtypes.fp32 + ) + output = torch.full_like(query, float("nan")) + torch.ops.aiter.pa_decode_flydsl( + output, + query, + key_cache, + value_cache, + context_lengths, + block_tables, + head_dim**-0.5, + 1, + num_partitions, + 256, + quant_dtype, + None, + scale, + scale, + ) + # The reference uses the values actually representable in the FP8 cache. + expected = ( + torch.tensor( + [float(seq + 1) if length else 0.0 for seq, length in enumerate(lengths)], + dtype=dtypes.fp32, + ) + .to(quant_dtype) + .to(dtypes.fp32) + .reshape(batch_size, 1, 1) + .expand_as(output) + ) + _assert_matches(output.float(), expected, tolerance=0.02) + + +@pytest.mark.parametrize("num_partitions", [1, 4]) +@pytest.mark.parametrize("block_table_width", [1, 2, 3, 4, 5]) +def test_block_table_row_alignment(block_table_width, num_partitions): + """A contiguous block table need not have four-entry-aligned row starts.""" + _constant_page_decode_case([16, 16, 16], 16, num_partitions, block_table_width) + + +@pytest.mark.parametrize("num_partitions", [1, 4]) +@pytest.mark.parametrize("block_size", [16, 64, 128]) +@pytest.mark.parametrize( + "lengths,per_token", + [ + pytest.param([64, 0], False, id="mixed"), + pytest.param([64, 0], True, id="mixed-per-token"), + pytest.param([0, 0], False, id="all-empty"), + ], +) +def test_empty_contexts_do_not_read_kv(lengths, per_token, block_size, num_partitions): + """Empty rows must not follow padding or read an empty KV allocation.""" + width = 2 * KV_COMPUTE_BLOCK // block_size if any(lengths) else 0 + _constant_page_decode_case( + lengths, + block_size, + num_partitions, + width, + poison_padding=True, + per_token=per_token, + ) + + +@pytest.mark.parametrize("context_length", [63, 64, 65, 127, 128, 129, 257]) +@pytest.mark.parametrize("trans_v", [False, True]) +@pytest.mark.parametrize("per_token", [False, True]) +def test_page128_chunk_boundaries(context_length, trans_v, per_token): + """Two waves share a page-128, but must load different K/V/scales halves.""" + _require_gpu() + output, reference = _adversarial_case( + block_size=128, + context_length=context_length, + query_group_size=16, + trans_v=trans_v, + per_token=per_token, + num_partitions=4, + ) + _assert_matches(output, reference) + + +@pytest.mark.parametrize("context_length", [2048, 100000, 200000]) +@pytest.mark.parametrize("query_length", [1, 3]) +@pytest.mark.parametrize("per_token", [False, True]) +def test_tp4_page128_contexts(context_length, query_length, per_token): + """64 Q / 4 KV heads under TP4: Hq16, Hkv1, D128, including dense MTP.""" + _require_gpu() + output, reference = _adversarial_case( + block_size=128, + query_group_size=16, + context_length=context_length, + query_length=query_length, + trans_v=True, + per_token=per_token, + num_partitions=8, + ) + _assert_matches(output, reference) + + +@pytest.mark.parametrize("trans_v", [False, True]) +def test_page128_cache_above_2gib(trans_v): + """Large batches at 200k context cross the signed-i32 cache offset limit.""" + _require_gpu() + block_size, head_dim = 128, 128 + boundary_page = 2**31 // (block_size * head_dim) + pages = boundary_page + 2 + quant_dtype = _quant_dtype() + key_cache = torch.empty(pages, 1, 8, block_size, 16, dtype=quant_dtype) + value_shape = ( + (pages, 1, 8, head_dim, 16) if trans_v else (pages, 1, head_dim, block_size) + ) + value_cache = torch.empty(value_shape, dtype=quant_dtype) + key_cache[0].zero_() + value_cache[0].zero_() + page_ids = [boundary_page - 1, boundary_page, boundary_page + 1] + for index, page in enumerate(page_ids): + key_cache[page].zero_() + value_cache[page].fill_(index + 1) + query = torch.ones(1, 16, head_dim, dtype=dtypes.bf16) + output = torch.empty_like(query) + scale = torch.ones(1, dtype=dtypes.fp32) + torch.ops.aiter.pa_decode_flydsl( + output, + query, + key_cache, + value_cache, + torch.tensor([257], dtype=dtypes.i32), + torch.tensor([page_ids], dtype=dtypes.i32), + head_dim**-0.5, + 1, + 8, + 256, + quant_dtype, + None, + scale, + scale, + ) + expected = torch.full_like(output, (128 + 2 * 128 + 3) / 257, dtype=dtypes.fp32) + _assert_matches(output.float(), expected, tolerance=0.02) + + +@pytest.mark.parametrize( + "block_size,per_token,packed", + [ + (16, False, False), + (16, True, False), + (16, False, True), + (128, False, False), + (128, True, False), + ], +) +def test_prepared_sparse_page_tables(block_size, per_token, packed): + """PA consumes selected pages; each MTP query already has its own table. + + Select 2048 tokens from a 200k-token cache. In the page-16 adapter each + selected logical page-128 expands to eight physical pages. Packed K/V + sides share a block and V has a different base pointer and cache extent. + """ + _require_gpu() + torch.manual_seed(7) + rows, query_heads, head_dim = 6, 16, 128 # two requests with three MTP queries + logical_blocks = (200000 + 127) // 128 + pages_per_logical = 128 // block_size + page_stride = pages_per_logical * (2 if packed else 1) + num_pages = logical_blocks * page_stride + quant_dtype = _quant_dtype() + key = (torch.rand(num_pages, 1, block_size, head_dim) - 0.5).to(quant_dtype) + value = (torch.rand_like(key, dtype=dtypes.fp32) - 0.5).to(quant_dtype) + key_cache = ( + key.view(num_pages, 1, block_size, head_dim // 16, 16) + .permute(0, 1, 3, 2, 4) + .contiguous() + ) + value_cache = ( + value.view(num_pages, 1, block_size // 16, 16, head_dim) + .permute(0, 1, 2, 4, 3) + .contiguous() + ) + if packed: + # Mirror the adapter's contiguous, overlapping views into interleaved + # K/V storage. Only K-side page ids are put into the block table. + value_cache = key_cache.flatten()[ + pages_per_logical * block_size * head_dim : + ].view(num_pages - pages_per_logical, 1, block_size // 16, head_dim, 16) + value = value_cache.permute(0, 1, 2, 4, 3).reshape( + num_pages - pages_per_logical, 1, block_size, head_dim + ) + + query = (torch.rand(rows, query_heads, head_dim) - 0.5).to(dtypes.bf16) + lengths = [2046, 2047, 2048] * 2 + # Spare entry deliberately leaves page-16 rows non-vector-aligned. + block_tables = torch.full( + (rows, 16 * pages_per_logical + 1), num_pages + 1024, dtype=dtypes.i32 + ) + if per_token: + key_scale = 0.5 + torch.rand(num_pages, 1, block_size, 1) + value_scale = 0.5 + torch.rand_like(key_scale) + else: + key_scale = torch.ones(1, dtype=dtypes.fp32) + value_scale = torch.ones(1, dtype=dtypes.fp32) + reference = torch.empty_like(query, dtype=dtypes.fp32) + for row, length in enumerate(lengths): + selected = torch.randperm(logical_blocks - 1)[:15] + # Partial current block comes last; earlier selections differ per query. + selected = torch.cat((selected, selected.new_tensor([logical_blocks - 1]))) + pages = ( + selected[:, None] * page_stride + torch.arange(pages_per_logical)[None, :] + ).flatten() + block_tables[row, : pages.numel()] = pages.to(dtypes.i32) + keys = key[pages].float() + values = value[pages].float() + if per_token: + keys = keys * key_scale[pages] + values = values * value_scale[pages] + keys = keys.reshape(-1, head_dim)[:length] + values = values.reshape(-1, head_dim)[:length] + scores = query[row].float() @ keys.T * head_dim**-0.5 + reference[row] = torch.softmax(scores, dim=-1) @ values + output = torch.empty_like(query) + torch.ops.aiter.pa_decode_flydsl( + output, + query, + key_cache, + value_cache, + torch.tensor(lengths, dtype=dtypes.i32), + block_tables, + head_dim**-0.5, + 1, + 8, + 256, + quant_dtype, + None, + key_scale, + value_scale, + sliding_window=-1, + ) + _assert_matches(output.float(), reference) + + +@pytest.mark.parametrize("context_length", [3, 257, 2051]) +@pytest.mark.parametrize("num_partitions", [1, 4]) +@pytest.mark.parametrize("trans_v", [False, True]) +@pytest.mark.parametrize("zero_query", [False, True]) +def test_mtp3_page128_partition_boundaries( + context_length, num_partitions, trans_v, zero_query +): + """MTP rows retain independent softmax state across partial/empty partitions.""" + output, reference = _adversarial_case( + head_dim=128, + context_length=context_length, + block_size=128, + query_group_size=16, + query_length=3, + num_partitions=num_partitions, + trans_v=trans_v, + zero_query=zero_query, + ) + _assert_matches(output, reference) + + +@pytest.mark.parametrize("all_empty", [False, True]) +def test_page128_decode_prefetch_mixed_contexts(all_empty): + """A two-CTA-per-CU decode grid handles empty partitions and poisoned padding.""" + lengths = [0] * 64 if all_empty else [0, 1, 63, 127, 128, 129, 257, 2051] * 8 + _constant_page_decode_case( + lengths, + block_size=128, + num_partitions=8, + block_table_width=0 if all_empty else 17, + poison_padding=True, + trans_v=True, + ) + + +@pytest.mark.parametrize( + "query_group_size,query_length", [(8, 3), (16, 2), (16, 4), (16, 5)] +) +@pytest.mark.parametrize("context_length", [129, 1027]) +def test_page128_reduction_rows_across_m_tiles( + query_group_size, query_length, context_length +): + """Aligned scratch rows stay independent across padded M-tiles and partitions.""" + output, reference = _adversarial_case( + head_dim=128, + context_length=context_length, + block_size=128, + query_group_size=query_group_size, + query_length=query_length, + num_partitions=4, + trans_v=True, + ) + _assert_matches(output, reference) + + def main(): torch.set_default_device("cuda") if not torch.cuda.is_available(): @@ -698,8 +1050,8 @@ def main(): "--block-size", type=int, nargs="*", - choices=[16, 64], - default=[16, 64], + choices=[16, 64, 128], + default=[16, 64, 128], help="""KV-cache block sizes.""", ) args = parser.parse_args() From da579b86e53347492ab6f453a95660a526231621 Mon Sep 17 00:00:00 2001 From: Wulley Date: Sun, 13 Sep 2026 13:21:50 +0000 Subject: [PATCH 18/66] [FlyDSL] Fix phase-split PA decode softmax overflow Preserve the -inf running maximum when correcting empty softmax history to avoid exp2 overflow on large negative logits. Add 16 BF16/FP16 regressions covering scale layouts, MTP, partitions, padded rows, and retained history. Run them from the standard CI entry point. Validated on gfx950 GPU 6: 98 related tests passed, plus the CI entry point with 16 regressions and 45 sweep configurations. The 32 paired benchmarks stayed within the 3% regression threshold; register and LDS usage were unchanged in eight configurations. --- .github/scripts/aiter_test.sh | 14 ++++ aiter/ops/flydsl/kernels/pa_decode_tile.py | 5 +- op_tests/test_flydsl_pa_decode.py | 76 ++++++++++++++++++++++ 3 files changed, 93 insertions(+), 2 deletions(-) diff --git a/.github/scripts/aiter_test.sh b/.github/scripts/aiter_test.sh index 5fae2fa2c7f..d17c4bb530d 100755 --- a/.github/scripts/aiter_test.sh +++ b/.github/scripts/aiter_test.sh @@ -124,6 +124,20 @@ 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 does not collect parametrized pytest regressions. + test_cmd=( + timeout 60m + bash -c ' + set -euo pipefail + test_file=$1 + python3 -m pytest -q \ + "${test_file}::test_large_negative_logits_preserve_online_softmax" + 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/kernels/pa_decode_tile.py b/aiter/ops/flydsl/kernels/pa_decode_tile.py index 847340a3bb8..a4cb85edf26 100644 --- a/aiter/ops/flydsl/kernels/pa_decode_tile.py +++ b/aiter/ops/flydsl/kernels/pa_decode_tile.py @@ -1049,8 +1049,9 @@ def _lmax_off_m(m): ls = ls + ls.shuffle_xor(sh, WAVE) # PV output is [head-dim, query-row=lane16] after the operand # swap, so correction/denominator are per-lane scalars (no sCorr). - safe_prev = (m_prev > NEG_INF).select(m_prev, ZERO_F) - corr_reg = fx.Float32(exp2_amdgcn_scalar(safe_prev - safe_max)) + # Empty history must contribute zero: exp2(-inf - safe_max). + # Replacing m_prev with 0 can overflow for negative logits. + corr_reg = fx.Float32(exp2_amdgcn_scalar(m_prev - safe_max)) if rgroup == 0: _st_lw(lsum_base, lane16, warp, ls) gpu.barrier() diff --git a/op_tests/test_flydsl_pa_decode.py b/op_tests/test_flydsl_pa_decode.py index 203c58a032f..652a4b5b849 100644 --- a/op_tests/test_flydsl_pa_decode.py +++ b/op_tests/test_flydsl_pa_decode.py @@ -601,6 +601,82 @@ def test_zero_query_stays_finite(query_length, per_token): _assert_matches(output, reference) +@pytest.mark.parametrize("query_dtype", [dtypes.bf16, dtypes.fp16]) +@pytest.mark.parametrize("per_token", [False, True]) +@pytest.mark.parametrize( + "query_group_size,query_length,block_size,context_length,num_partitions,trans_v,split_values", + [ + pytest.param(32, 1, 64, 128, 1, False, False, id="decode"), + pytest.param(16, 3, 128, 257, 4, True, False, id="mtp-partitions"), + pytest.param(8, 3, 128, 257, 1, True, False, id="padded-m-tile"), + pytest.param(32, 1, 64, 512, 1, False, True, id="history"), + ], +) +def test_large_negative_logits_preserve_online_softmax( + query_group_size, + query_length, + block_size, + context_length, + num_partitions, + trans_v, + split_values, + per_token, + query_dtype, +): + """Large negative logits must preserve both initial and accumulated state.""" + _require_gpu() + head_dim = 128 + blocks = (context_length + block_size - 1) // block_size + quant_dtype = _quant_dtype() + query = torch.full( + (query_length, query_group_size, head_dim), -3.0, dtype=query_dtype + ) + key = torch.full((blocks, 1, block_size, head_dim), 3.0).to(quant_dtype) + value = torch.ones(blocks, 1, block_size, head_dim) + if split_values: + # With equal logits, both KV tiles contribute equally. Resetting the + # correction to zero on every tile would incorrectly return 1, not 0.5. + value[: KV_COMPUTE_BLOCK // block_size] = 0 + value = value.to(quant_dtype) + key_cache = ( + key.view(blocks, 1, block_size, head_dim // 16, 16) + .permute(0, 1, 3, 2, 4) + .contiguous() + ) + if trans_v: + value_cache = ( + value.view(blocks, 1, block_size // 16, 16, head_dim) + .permute(0, 1, 2, 4, 3) + .contiguous() + ) + else: + value_cache = value.permute(0, 1, 3, 2).contiguous() + block_tables = torch.arange(blocks, dtype=dtypes.i32).reshape(1, blocks) + context_lengths = torch.full((1,), context_length, dtype=dtypes.i32) + scale = torch.ones( + (blocks, 1, block_size, 1) if per_token else (1,), dtype=dtypes.fp32 + ) + output = torch.full_like(query, float("nan")) + torch.ops.aiter.pa_decode_flydsl( + output, + query, + key_cache, + value_cache, + context_lengths, + block_tables, + head_dim**-0.5, + query_length, + num_partitions, + KV_COMPUTE_BLOCK, + quant_dtype, + None, + scale, + scale, + ) + reference = torch.full_like(output, 0.5 if split_values else 1.0, dtype=dtypes.fp32) + _assert_matches(output.float(), reference, tolerance=0.02) + + @pytest.mark.parametrize("block_size", [16, 64, 128]) @pytest.mark.parametrize("context_length", [1, 257, 1027]) def test_block_table_padding_is_not_dereferenced(block_size, context_length): From be0963c14d313781c88ebd2ebcb28a5e12bf6e51 Mon Sep 17 00:00:00 2001 From: Wulley Date: Sun, 13 Sep 2026 14:28:16 +0000 Subject: [PATCH 19/66] [FlyDSL] Remove duplicate group_idx calculation in PA reduce --- aiter/ops/flydsl/kernels/pa_decode_reduce.py | 1 - 1 file changed, 1 deletion(-) diff --git a/aiter/ops/flydsl/kernels/pa_decode_reduce.py b/aiter/ops/flydsl/kernels/pa_decode_reduce.py index c358dd8b8bf..dff93eb040d 100644 --- a/aiter/ops/flydsl/kernels/pa_decode_reduce.py +++ b/aiter/ops/flydsl/kernels/pa_decode_reduce.py @@ -186,7 +186,6 @@ def _wave_reduce_sum(value): acc = acc + part_logits * weight query_idx = eqgs_idx // c_qgs - group_idx = eqgs_idx % c_qgs output_offset = ( batch_idx * stride_output_bs + query_idx * stride_output_len From 77572ba8d989a93f2eeab3ccf9339cc7f9cf6116 Mon Sep 17 00:00:00 2001 From: fsx950223 Date: Mon, 14 Sep 2026 03:19:23 +0000 Subject: [PATCH 20/66] fix(flydsl): import PA decode function explicitly Signed-off-by: fsx950223 --- aiter/ops/attention.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/aiter/ops/attention.py b/aiter/ops/attention.py index 4a0a5fe13a8..eae49925f4d 100644 --- a/aiter/ops/attention.py +++ b/aiter/ops/attention.py @@ -22,7 +22,7 @@ from ..jit.utils.chip_info import get_cu_num, get_gfx try: - from aiter.ops.flydsl import pa_decode as _pa_decode_flydsl + from aiter.ops.flydsl.pa_decode import pa_decode as _pa_decode_flydsl except (ImportError, AttributeError, RuntimeError, OSError): _pa_decode_flydsl = None From 3bcde89bbb4c2bcc3d0efeb90867e3dfe1719d94 Mon Sep 17 00:00:00 2001 From: fsx950223 Date: Mon, 14 Sep 2026 03:56:51 +0000 Subject: [PATCH 21/66] perf(flydsl): prototype page16 PA V/K pipeline --- aiter/ops/flydsl/kernels/pa_decode_tile.py | 89 +++++++++++++++------- aiter/ops/flydsl/pa_decode.py | 32 +++++++- 2 files changed, 94 insertions(+), 27 deletions(-) diff --git a/aiter/ops/flydsl/kernels/pa_decode_tile.py b/aiter/ops/flydsl/kernels/pa_decode_tile.py index a4cb85edf26..4bddea58096 100644 --- a/aiter/ops/flydsl/kernels/pa_decode_tile.py +++ b/aiter/ops/flydsl/kernels/pa_decode_tile.py @@ -118,8 +118,10 @@ def compile_pa_decode_tile( TOTAL_ROWS = query_length * query_group_size M_TILES = cdiv(TOTAL_ROWS, MFMA_MNK) ROWS_PADDED = M_TILES * MFMA_MNK - # The page-128 gfx950 scalar-scale path benefits from earlier V loads - # for decode and alternating P/reduction buffers for three M-tiles. + # The gfx950 scalar-scale decode paths can opt into earlier V loads. The + # page-16 single-M-tile variant additionally delays the next K load until + # immediately before the current PV, shortening its live range across the + # softmax while leaving the existing schedule as the default. tune_page128 = ( is_gfx950 and head_dim == 128 @@ -127,7 +129,15 @@ def compile_pa_decode_tile( and trans_v and not per_token_kv ) - EARLY_V = prefetch_v and tune_page128 and M_TILES == 1 + tune_page16_vpipe = ( + is_gfx950 + and head_dim == 128 + and block_size == 16 + and trans_v + and not per_token_kv + ) + EARLY_V = prefetch_v and (tune_page128 or tune_page16_vpipe) and M_TILES == 1 + PAGE16_VPIPE = prefetch_v and tune_page16_vpipe and M_TILES == 1 P_BUFFERS = 2 if tune_page128 and M_TILES == 3 else 1 # PV layout: V=A, P=B -> output [head-dim (row), query-row (col=lane16)], # generalized over head_dim via the VHE_CHUNKS loop. @@ -432,6 +442,12 @@ def _load_phys_scalar(page, vec_width=1): 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): # V's page depends on `rgroup` (shared across warps): warp w fetches # its rgroup row and broadcasts via LDS (read back by _v_page_read_row). @@ -439,20 +455,21 @@ def _v_page_fetch_and_stage(tt_i32): fetched = _load_phys_scalar( base_page + (warp * TOK_PER_WARP) // block_size, PAGES_PER_CHUNK ) - if lane == 0: - fetched_vec = ( - fx.Vector.from_elements([fx.Int32(fetched)], dtype=fx.Int32) - if const_expr(PAGES_PER_CHUNK == 1) - else fx.Vector(fetched) - ) - _lds_store( - sVPage_off + warp * (PAGES_PER_CHUNK * 4), fx.Int32, fetched_vec - ) + 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) 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): # ping-pong buffer byte offset (0 when single-buffered). if const_expr(per_token_kv): @@ -521,6 +538,16 @@ def _k_ops(phys, a): 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 * c16) // 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( @@ -531,14 +558,7 @@ def _k_ops_flat(tt_i32): if const_expr(PAGES_PER_CHUNK == 1) else fx.Vector(fetched) ) - flat = [] - for a in range_constexpr(NCHUNK): - phys = fx.Int32(phys_vec[(a * c16) // 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), phys_vec + return _k_ops_from_phys(phys_vec), phys_vec # Empty partitions retain neutral state and must not read K/V or the # block table. In particular, an empty context has no valid first tile. @@ -547,7 +567,12 @@ def _k_ops_flat(tt_i32): k_pf0, phys_vec0 = _k_ops_flat(part_start) # Issue the V page-index prefetch alongside K; the LDS write is # visible after the barrier below. - _v_page_fetch_and_stage(part_start) + if const_expr(PAGE16_VPIPE): + # The K page vector is already resident. Reuse it for the V + # page matrix instead of issuing a second block-table load. + _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: @@ -1109,11 +1134,16 @@ def _lmax_off_m(m): # tt+1 K/V/scale prefetch, issued here to reuse the pass-1 barrier. k_next = k_cur if tt1 < part_end: - 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 + if const_expr(PAGE16_VPIPE): + # Publish only the next page ids here. K itself is + # issued after the P barrier so its VMEM can overlap + # the current tile's PV without spanning softmax. + _v_page_fetch_and_stage(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)) # Softmax: each lane owns one qhead (lane%16); register reduce + shuffle_xor. scale = scale_qk * _ld1( sQscale_off, lane16 @@ -1229,6 +1259,13 @@ def _lmax_off_m(m): if rgroup == 0: _st_lw(sLsum_off, lane16, warp, ls) gpu.barrier() + if const_expr(PAGE16_VPIPE): # noqa: SIM102 + # The page matrix is now visible. Issue next K before the + # current P read/PV so those independent operations can + # cover its VMEM latency. + if 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 = l_prev * corr_reg + gsum p_ops = _lds_load( diff --git a/aiter/ops/flydsl/pa_decode.py b/aiter/ops/flydsl/pa_decode.py index 6c7047d1a43..4d9c736c3c4 100644 --- a/aiter/ops/flydsl/pa_decode.py +++ b/aiter/ops/flydsl/pa_decode.py @@ -26,6 +26,8 @@ an LDS round-trip on P transposing ownership between the two MMAs. """ +import os + import flydsl.compiler as flyc import flydsl.expr as fx import torch @@ -37,6 +39,17 @@ from .kernels.tensor_shim import _run_compiled, ptr_arg from .kernels.utils import cdiv +_PAGE16_VPIPE_ENV = "AITER_FLYDSL_PA_PAGE16_VPIPE" + + +def _page16_vpipe_enabled() -> bool: + return os.getenv(_PAGE16_VPIPE_ENV, "0").strip().lower() in { + "1", + "true", + "yes", + "on", + } + def get_recommended_splits( num_sequences: int, @@ -378,10 +391,27 @@ def normalize_scale(scale, name): # if either does. wide_kv_addressing = max(key_cache.numel(), value_cache.numel()) >= 2**31 + # The page-16 long-context pipeline is opt-in until it can be benchmarked + # on gfx950. It loads current V before QK and delays next K until the + # current PV, while every unsupported shape keeps the existing schedule. + page16_vpipe = ( + _page16_vpipe_enabled() + and arch == "gfx950" + and head_dim == 128 + and block_size == 16 + and trans_v + and not per_token_kv + and query.dtype == torch.bfloat16 + and num_kv_heads == 1 + and query_length == 1 + and query_group_size in (8, 16) + and num_partitions == 8 + ) + # Early V loads help the measured one-to-two-workgroup-per-CU regime. # Keep other grids on the existing schedule: early loads regress # short-context decode at larger grid sizes. - prefetch_v = False + prefetch_v = page16_vpipe if ( arch == "gfx950" and head_dim == 128 From 68cf4912d7675ed8e5d08249b9105d9f2dabd0d9 Mon Sep 17 00:00:00 2001 From: fsx950223 Date: Mon, 14 Sep 2026 03:59:45 +0000 Subject: [PATCH 22/66] perf(flydsl): add page16 PA IGLP variant --- aiter/ops/flydsl/kernels/pa_decode_tile.py | 9 +++++++-- aiter/ops/flydsl/pa_decode.py | 10 +++++++--- 2 files changed, 14 insertions(+), 5 deletions(-) diff --git a/aiter/ops/flydsl/kernels/pa_decode_tile.py b/aiter/ops/flydsl/kernels/pa_decode_tile.py index 4bddea58096..a34dfc77ca5 100644 --- a/aiter/ops/flydsl/kernels/pa_decode_tile.py +++ b/aiter/ops/flydsl/kernels/pa_decode_tile.py @@ -66,6 +66,7 @@ def compile_pa_decode_tile( trans_v: bool = True, wide_kv_addressing: bool = False, prefetch_v: bool = False, + prefetch_v_iglp: bool = False, ): """Build the tile-programming PA-decode kernel + launch wrapper. @@ -138,6 +139,7 @@ def compile_pa_decode_tile( ) EARLY_V = prefetch_v and (tune_page128 or tune_page16_vpipe) and M_TILES == 1 PAGE16_VPIPE = prefetch_v and tune_page16_vpipe and M_TILES == 1 + PAGE16_VPIPE_IGLP = PAGE16_VPIPE and prefetch_v_iglp P_BUFFERS = 2 if tune_page128 and M_TILES == 3 else 1 # PV layout: V=A, P=B -> output [head-dim (row), query-row (col=lane16)], # generalized over head_dim via the VHE_CHUNKS loop. @@ -826,7 +828,9 @@ def _l_slot(m): tok0 = tt * TILE_TOK # per_tensor phase-split: let IGLP interleave MFMA with softmax # VALU/LDS to hide the MFMA-hazard s_nop (per_token is VGPR-cliffed). - if const_expr(not per_token_kv and M_TILES > 1): + if const_expr( + (not per_token_kv and M_TILES > 1) or PAGE16_VPIPE_IGLP + ): fx.rocdl.iglp_opt(0) tt1 = tt + 1 @@ -1131,7 +1135,8 @@ def _lmax_off_m(m): [k_cur[a * N_SUBCHUNKS + s], q_ops_all[s], acc, 0, 0, 0], ) frag_Ss.append(fx.Vector(acc)) - # tt+1 K/V/scale prefetch, issued here to reuse the pass-1 barrier. + # Prepare tt+1 before the pass-1 barrier. The page-16 deep + # pipeline publishes page ids only; other paths also load K. k_next = k_cur if tt1 < part_end: if const_expr(PAGE16_VPIPE): diff --git a/aiter/ops/flydsl/pa_decode.py b/aiter/ops/flydsl/pa_decode.py index 4d9c736c3c4..4385953fc0c 100644 --- a/aiter/ops/flydsl/pa_decode.py +++ b/aiter/ops/flydsl/pa_decode.py @@ -40,10 +40,11 @@ from .kernels.utils import cdiv _PAGE16_VPIPE_ENV = "AITER_FLYDSL_PA_PAGE16_VPIPE" +_PAGE16_VPIPE_IGLP_ENV = "AITER_FLYDSL_PA_PAGE16_VPIPE_IGLP" -def _page16_vpipe_enabled() -> bool: - return os.getenv(_PAGE16_VPIPE_ENV, "0").strip().lower() in { +def _env_enabled(name: str) -> bool: + return os.getenv(name, "0").strip().lower() in { "1", "true", "yes", @@ -394,8 +395,9 @@ def normalize_scale(scale, name): # The page-16 long-context pipeline is opt-in until it can be benchmarked # on gfx950. It loads current V before QK and delays next K until the # current PV, while every unsupported shape keeps the existing schedule. + page16_vpipe_iglp = _env_enabled(_PAGE16_VPIPE_IGLP_ENV) page16_vpipe = ( - _page16_vpipe_enabled() + (_env_enabled(_PAGE16_VPIPE_ENV) or page16_vpipe_iglp) and arch == "gfx950" and head_dim == 128 and block_size == 16 @@ -407,6 +409,7 @@ def normalize_scale(scale, name): and query_group_size in (8, 16) and num_partitions == 8 ) + page16_vpipe_iglp = page16_vpipe and page16_vpipe_iglp # Early V loads help the measured one-to-two-workgroup-per-CU regime. # Keep other grids on the existing schedule: early loads regress @@ -437,6 +440,7 @@ def normalize_scale(scale, name): trans_v=trans_v, wide_kv_addressing=wide_kv_addressing, prefetch_v=prefetch_v, + prefetch_v_iglp=page16_vpipe_iglp, ) if num_partitions == 1: From cdc4c74014dc55a852a05d47782716f6b8fd5983 Mon Sep 17 00:00:00 2001 From: fsx950223 Date: Mon, 14 Sep 2026 04:06:41 +0000 Subject: [PATCH 23/66] perf(flydsl): allow page16 PA pipeline for all split counts --- aiter/ops/flydsl/pa_decode.py | 1 - 1 file changed, 1 deletion(-) diff --git a/aiter/ops/flydsl/pa_decode.py b/aiter/ops/flydsl/pa_decode.py index 4385953fc0c..1de3acb225a 100644 --- a/aiter/ops/flydsl/pa_decode.py +++ b/aiter/ops/flydsl/pa_decode.py @@ -407,7 +407,6 @@ def normalize_scale(scale, name): and num_kv_heads == 1 and query_length == 1 and query_group_size in (8, 16) - and num_partitions == 8 ) page16_vpipe_iglp = page16_vpipe and page16_vpipe_iglp From 7bc366057a0200c0bae3e6173489ddf00b66909d Mon Sep 17 00:00:00 2001 From: fsx950223 Date: Mon, 14 Sep 2026 04:22:11 +0000 Subject: [PATCH 24/66] perf(flydsl): add Gluon numeric-policy PA ablation --- aiter/ops/flydsl/kernels/pa_decode_tile.py | 117 +++++++++++++++------ aiter/ops/flydsl/pa_decode.py | 20 ++++ 2 files changed, 106 insertions(+), 31 deletions(-) diff --git a/aiter/ops/flydsl/kernels/pa_decode_tile.py b/aiter/ops/flydsl/kernels/pa_decode_tile.py index a34dfc77ca5..bc596bd392a 100644 --- a/aiter/ops/flydsl/kernels/pa_decode_tile.py +++ b/aiter/ops/flydsl/kernels/pa_decode_tile.py @@ -67,6 +67,7 @@ def compile_pa_decode_tile( wide_kv_addressing: bool = False, prefetch_v: bool = False, prefetch_v_iglp: bool = False, + match_gluon_numerics: bool = False, ): """Build the tile-programming PA-decode kernel + launch wrapper. @@ -111,6 +112,18 @@ def compile_pa_decode_tile( ), f"pa_decode_tile only supports query_dtype in ('f16', 'bf16'), got {query_dtype}" Q_DTYPE = fx.BFloat16 if query_dtype == "bf16" else fx.Float16 + if match_gluon_numerics: + assert ( + is_gfx950 + and head_dim == 128 + and block_size == 16 + and trans_v + and not per_token_kv + and query_dtype == "bf16" + and query_length == 1 + and query_group_size in (8, 16) + ), "matched Gluon numerics is restricted to the MiniMax-M3 page-16 shapes" + assert ( head_dim % 64 == 0 ), f"pa_decode_tile only supports head_dim that's a multiple of 64, got {head_dim}" @@ -185,7 +198,8 @@ def compile_pa_decode_tile( # -- LDS layout (shared across the 4 warps) -- # sQ: fp8[ROWS_PADDED,head_dim] staged+quantized query. sP: fp8[16,TILE_TOK] - # quantized probs. sQscale: f32[ROWS_PADDED]. sLmax/sLsum: cross-warp scratch. + # quantized probs. sQscale: f32[ROWS_PADDED] except in the matched-Gluon + # ablation, where Q has no dynamic scale. sLmax/sLsum: cross-warp scratch. # sVPage: V page-index broadcast. sKScale/sVScale/sVScaleMax: per-token K/V # scale staging. No sO/sM/sL/sCorr: PV output is register-resident/loop-carried # and stored straight to global (V=A/P=B swap). @@ -197,12 +211,13 @@ def compile_pa_decode_tile( SP_ROW_BYTES = TILE_TOK + 16 sP_bytes = P_BUFFERS * MFMA_MNK * SP_ROW_BYTES # fp8, padded rows sQscale_off = sP_off + sP_bytes + sQscale_bytes = 0 if match_gluon_numerics else ROWS_PADDED * f32 # Keep reduction rows 16-byte aligned on the tuned gfx950 page-128 path # so each cross-wave row can be read as one vector. Other paths retain # the padding used to avoid LDS bank conflicts. NWARP_PAD = NWARP if tune_page128 else NWARP + 1 # Phase-split slices sLmax per M-tile so all pass-1 writes share one barrier. - sLmax_off = sQscale_off + ROWS_PADDED * f32 + sLmax_off = sQscale_off + sQscale_bytes sLsum_off = sLmax_off + M_TILES * MFMA_MNK * NWARP_PAD * f32 # V page-index broadcast (V's page depends on `rgroup`, shared across warps). sVPage_off = sLsum_off + P_BUFFERS * MFMA_MNK * NWARP_PAD * f32 @@ -661,33 +676,55 @@ def _quant_q_row(m, qi, gs_head, q_row_off): for u in range_constexpr(N_QLOADS) ] - absmax = fmath.absf(q_units[0]).reduce(ReductionOp.MAX).to(fx.Float32) - for u in range_constexpr(1, N_QLOADS): - absmax = fx.maxnumf( - absmax, - fmath.absf(q_units[u]).reduce(ReductionOp.MAX).to(fx.Float32), + if const_expr(match_gluon_numerics): + # Match Gluon's Hkv1 policy when QUERY_QUANT_MODE==-1: cast + # BF16 Q directly to the fp8 MFMA type, with no per-row absmax + # normalization and therefore no Q scale to fold into QK. + 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 = ( + fmath.absf(q_units[0]) + .reduce(ReductionOp.MAX) + .to(fx.Float32) ) - 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(1, N_QLOADS): + absmax = fx.maxnumf( + absmax, + fmath.absf(q_units[u]) + .reduce(ReductionOp.MAX) + .to(fx.Float32), + ) + for sh in (8, 4, 2, 1): + absmax = fx.maxnumf(absmax, dpp_utils.dpp_xor_f32(absmax, sh)) - 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), + 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 ) - if lane16 == 0: - # Transposed [qh][m] (not [m][qh]) so the whole M_TILES-wide - # row for a fixed qh is contiguous, letting the KV-loop read - # it back in one wide load instead of M_TILES separate - # narrow ones -- see the read site below. - _st1(sQscale_off, qh_local * M_TILES + m, q_scale) + + 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] (not [m][qh]) so the whole M_TILES-wide + # row for a fixed qh is contiguous, letting the KV-loop read + # it back in one wide load instead of M_TILES separate + # narrow ones -- see the read site below. + _st1(sQscale_off, qh_local * M_TILES + m, q_scale) for m in range_constexpr(M_TILES): flat_idx = m * MFMA_MNK + qh_local @@ -707,7 +744,7 @@ def _quant_q_row(m, qi, gs_head, q_row_off): q_row_off + qh_local * head_dim + lane16 * QCHUNK, fx.Vector.filled(QCHUNK // 4, 0, fx.Int32), ) - if lane16 == 0: + if const_expr(not match_gluon_numerics) and lane16 == 0: _st1(sQscale_off, qh_local * M_TILES + m, ZERO_F) gpu.barrier() @@ -876,7 +913,7 @@ def _mask_v_scale(vec, a, thr=ctx_thr, zero=zero4_scale): # q_scale doesn't depend on `m`; read the whole M_TILES-wide row once # (contiguous via the transposed [qh][m] sQscale layout). q_scale_vec = None - if const_expr(M_TILES > 1): + if const_expr(M_TILES > 1 and not match_gluon_numerics): q_scale_vec = _lds_load( sQscale_off + lane16 * (M_TILES * f32), fx.Float32, M_TILES ) @@ -932,7 +969,11 @@ def _lmax_off_m(m): ) frag_Ss.append(fx.Vector(acc)) - scale = scale_qk * fx.Float32(q_scale_vec[m]) + scale = ( + scale_qk + if const_expr(match_gluon_numerics) + else scale_qk * fx.Float32(q_scale_vec[m]) + ) n_valid_tile = (causal_bound[m] - tok0).to(fx.Float32) base_tok_f = fx.Int32(warp * TOK_CHUNK + rgroup * 4).to(fx.Float32) thr = fx.Vector.from_elements( @@ -1058,6 +1099,10 @@ def _lmax_off_m(m): else v_scale_shared[a] ) p_scaled = Pa * v_sc * norm_factor_b + elif const_expr(match_gluon_numerics): + # Match Gluon Hkv1's scalar-scale path: P is cast + # to fp8 in its native [0, 1] range. + p_scaled = Pa else: p_scaled = Pa * fx.Vector.filled(4, FP8_MAX, fx.Float32) words.append(_f32_to_fp8_words(p_scaled)[0]) @@ -1150,8 +1195,10 @@ def _lmax_off_m(m): if const_expr(per_token_kv): _stage_kv_scale_to_lds(phys_vec1, _kv_buf_off(tt1)) # Softmax: each lane owns one qhead (lane%16); register reduce + shuffle_xor. - scale = scale_qk * _ld1( - sQscale_off, lane16 + scale = ( + scale_qk + if const_expr(match_gluon_numerics) + else scale_qk * _ld1(sQscale_off, lane16) ) # per-qhead positive score scale n_valid_tile = (causal_bound[0] - tok0).to(fx.Float32) base_tok_f = fx.Int32(warp * TOK_CHUNK + rgroup * 4).to(fx.Float32) @@ -1243,6 +1290,10 @@ def _lmax_off_m(m): else v_scale_vecs[a] ) p_scaled = Pa * v_scale_this * norm_factor_b + elif const_expr(match_gluon_numerics): + # Gluon Hkv1 does not multiply scalar-scale P by 448 + # before its fp8 conversion. + p_scaled = Pa else: p_scaled = Pa * fx.Vector.filled(4, FP8_MAX, fx.Float32) words.append(_f32_to_fp8_words(p_scaled)[0]) @@ -1314,6 +1365,10 @@ def _lmax_off_m(m): inv_l = fx.Float32(rcp_f32(safe_l)) if const_expr(per_token_kv): o_scale = inv_l + elif const_expr(match_gluon_numerics): + # The matching P path did not scale by FP8_MAX, so preserve + # value_scale but omit FlyDSL's 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( diff --git a/aiter/ops/flydsl/pa_decode.py b/aiter/ops/flydsl/pa_decode.py index 1de3acb225a..a6aecbb7a8b 100644 --- a/aiter/ops/flydsl/pa_decode.py +++ b/aiter/ops/flydsl/pa_decode.py @@ -41,6 +41,7 @@ _PAGE16_VPIPE_ENV = "AITER_FLYDSL_PA_PAGE16_VPIPE" _PAGE16_VPIPE_IGLP_ENV = "AITER_FLYDSL_PA_PAGE16_VPIPE_IGLP" +_MATCH_GLUON_NUMERICS_ENV = "AITER_FLYDSL_PA_MATCH_GLUON_NUMERICS" def _env_enabled(name: str) -> bool: @@ -410,6 +411,24 @@ def normalize_scale(scale, name): ) page16_vpipe_iglp = page16_vpipe and page16_vpipe_iglp + # Numerical-policy ablation against Gluon's Hkv1 kernel. Keep this + # independent from the page-16 scheduling switches so the two effects can + # be measured separately, and constrain it to the exact MiniMax-M3 shapes + # whose Gluon path directly casts both Q and P to fp8. All other shapes, + # including per-token KV scales, retain FlyDSL's range-normalized policy. + match_gluon_numerics = ( + _env_enabled(_MATCH_GLUON_NUMERICS_ENV) + and arch == "gfx950" + and head_dim == 128 + and block_size == 16 + and trans_v + and not per_token_kv + and query.dtype == torch.bfloat16 + and num_kv_heads == 1 + and query_length == 1 + and query_group_size in (8, 16) + ) + # Early V loads help the measured one-to-two-workgroup-per-CU regime. # Keep other grids on the existing schedule: early loads regress # short-context decode at larger grid sizes. @@ -440,6 +459,7 @@ def normalize_scale(scale, name): wide_kv_addressing=wide_kv_addressing, prefetch_v=prefetch_v, prefetch_v_iglp=page16_vpipe_iglp, + match_gluon_numerics=match_gluon_numerics, ) if num_partitions == 1: From f4f84314aa552f3b045c4032dca06f56e4749dd5 Mon Sep 17 00:00:00 2001 From: fsx950223 Date: Mon, 14 Sep 2026 04:39:02 +0000 Subject: [PATCH 25/66] fix(flydsl): address PA decode review follow-ups Signed-off-by: fsx950223 --- .github/scripts/aiter_test.sh | 14 +- aiter/ops/flydsl/kernels/pa_decode_reduce.py | 80 +++- aiter/ops/flydsl/kernels/pa_decode_tile.py | 6 +- aiter/ops/flydsl/pa_decode.py | 421 +++++++++++++------ aiter/ops/triton/gluon/pa_decode_gluon.py | 116 +++-- op_tests/test_flydsl_pa_decode.py | 421 +++++++++++++++++-- 6 files changed, 819 insertions(+), 239 deletions(-) diff --git a/.github/scripts/aiter_test.sh b/.github/scripts/aiter_test.sh index d17c4bb530d..49873b4cd1d 100755 --- a/.github/scripts/aiter_test.sh +++ b/.github/scripts/aiter_test.sh @@ -132,7 +132,19 @@ for file in "${sharded_files[@]}"; do set -euo pipefail test_file=$1 python3 -m pytest -q \ - "${test_file}::test_large_negative_logits_preserve_online_softmax" + "${test_file}::test_pa_decode_api_matches_gluon" \ + "${test_file}::test_pa_decode_maps_gluon_buffers_and_scale_layout" \ + "${test_file}::test_gluon_unsupported_flydsl_reducer_falls_back" \ + "${test_file}::test_gluon_supported_flydsl_reducer_does_not_mask_errors" \ + "${test_file}::test_v_prefetch_workgroup_interval" \ + "${test_file}::test_pa_decode_fixed_length_accuracy" \ + "${test_file}::test_pa_decode_variable_length_accuracy" \ + "${test_file}::test_large_negative_logits_preserve_online_softmax" \ + "${test_file}::test_additional_supported_head_dims_are_accurate" \ + "${test_file}::test_supported_head_dim_boundaries_are_accepted" \ + "${test_file}::test_unsupported_head_dim_is_rejected" \ + "${test_file}::test_ps_false_is_rejected" \ + "${test_file}::test_invalid_tensor_structure_is_rejected" python3 "$test_file" ' _ "$file" diff --git a/aiter/ops/flydsl/kernels/pa_decode_reduce.py b/aiter/ops/flydsl/kernels/pa_decode_reduce.py index dff93eb040d..dc401b8c997 100644 --- a/aiter/ops/flydsl/kernels/pa_decode_reduce.py +++ b/aiter/ops/flydsl/kernels/pa_decode_reduce.py @@ -9,6 +9,56 @@ import flydsl.expr as fx from flydsl.expr.typing import T +_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, +) -> None: + if not 1 <= max_context_partition_num <= 64: + raise ValueError( + "max_context_partition_num must be in [1, 64], " + 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, +) -> 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, + ) + except ValueError: + return False + return True + @lru_cache(maxsize=256) def compile_pa_decode_ps_reduce( @@ -27,27 +77,17 @@ def compile_pa_decode_ps_reduce( API limits the partition count to one wave, so the LDS path from the more general Gluon implementation is unnecessary here. """ - if not 1 <= max_context_partition_num <= 64: - raise ValueError( - "max_context_partition_num must be in [1, 64], " - 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}" - ) + _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, + ) - dtype_map = { - "f32": fx.Float32, - "f16": fx.Float16, - "bf16": fx.BFloat16, - } - try: - output_dtype = dtype_map[output_dtype_str] - logits_dtype = dtype_map[logits_dtype_str] - sink_dtype = dtype_map[sink_dtype_str] - except KeyError as exc: - raise ValueError(f"Unsupported FlyDSL dtype: {exc.args[0]!r}") from exc + 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 diff --git a/aiter/ops/flydsl/kernels/pa_decode_tile.py b/aiter/ops/flydsl/kernels/pa_decode_tile.py index a4cb85edf26..8e2cc8f721e 100644 --- a/aiter/ops/flydsl/kernels/pa_decode_tile.py +++ b/aiter/ops/flydsl/kernels/pa_decode_tile.py @@ -219,9 +219,9 @@ class SharedStorage: def pa_decode_tile_kernel( output_ptr: fx.Pointer, # [num_seqs*query_length, num_q_heads, head_dim] (written directly when NP==1) # per-partition partial outputs (combined by the reduce kernel when NP>1): - pmax_ptr: fx.Pointer, # [num_seqs*query_length, num_kv_heads, num_partitions, query_group_size] row max - psum_ptr: fx.Pointer, # [num_seqs*query_length, num_kv_heads, num_partitions, query_group_size] row sum - pout_ptr: fx.Pointer, # [num_seqs*query_length, num_kv_heads, num_partitions, query_group_size, head_dim] Q_DTYPE, normalized O_p/l_p + pmax_ptr: fx.Pointer, # [num_seqs, num_kv_heads, num_partitions, query_length*query_group_size] row max + psum_ptr: fx.Pointer, # [num_seqs, num_kv_heads, num_partitions, query_length*query_group_size] row sum + pout_ptr: fx.Pointer, # same shape plus head_dim; Q_DTYPE normalized O_p/l_p query_ptr: fx.Pointer, # [num_seqs*query_length, num_q_heads, head_dim] -- row = seq*query_length + qi (MTP position) key_cache_ptr: fx.Pointer, # [num_blocks, num_kv_heads, head_dim//16, block_size, 16] (blocked, see module docstring) value_cache_ptr: fx.Pointer, # [num_blocks, num_kv_heads, block_size//16, head_dim, 16] (blocked, see module docstring) diff --git a/aiter/ops/flydsl/pa_decode.py b/aiter/ops/flydsl/pa_decode.py index 6c7047d1a43..6341da7bfec 100644 --- a/aiter/ops/flydsl/pa_decode.py +++ b/aiter/ops/flydsl/pa_decode.py @@ -10,8 +10,9 @@ ``key_scale``/``value_scale`` are either a ``[1]`` per-tensor scalar or a ``[num_blocks, num_kv_heads, block_size, 1]`` per-token tensor. -``block_size`` (16/64/128) and ``head_dim`` (multiple of 64) are compile-time -constants. Layouts are logical, not production's preshuffle. +``block_size`` (16/64/128) and ``head_dim`` (64, or a multiple of 128 up to +1024) are compile-time constants. Layouts are logical, not production's +preshuffle. * ``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 @@ -50,6 +51,22 @@ def get_recommended_splits( return max(4, min(n, 8)) +def _workgroup_count_enables_v_prefetch( + num_sequences: int, + num_kv_heads: int, + num_partitions: int, + num_compute_units: int, +) -> bool: + """Return whether the launch is in the tuned one-to-two-CTA-per-CU window. + + This result is a kernel specialization and therefore part of the compile + cache key. Applications serving shapes on both sides of this interval must + materialize both variants before graph capture or latency-sensitive use. + """ + workgroups = num_sequences * num_kv_heads * num_partitions + return num_compute_units < workgroups <= 2 * num_compute_units + + def _flydsl_dtype_str(dtype: torch.dtype) -> str: if dtype == torch.float32: return "f32" @@ -154,9 +171,9 @@ def pa_decode( 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, @@ -166,12 +183,20 @@ def pa_decode( The call signature and intermediate-buffer layouts follow the shared aiter paged-attention decode API. This kernel currently supports FP8 K/V caches, - BF16/FP16 queries, a 256-token compute tile, and block sizes 16/64/128. + BF16/FP16 queries, a 256-token compute tile, block sizes 16/64/128, and + ``head_dim=64`` or a multiple of 128 up to 1024. Sparse attention uses caller-prepared block tables and selected context lengths. Each independently selected MTP query must have its own table row and use query_length=1; query_length>1 applies dense causal masking. ALiBi, attention sinks, sliding-window attention, and externally quantized - FP8 queries are not supported. + FP8 queries are not supported. The ``ps=False`` partitioning policy is also + not supported. + + ``context_lengths`` and ``block_tables`` are GPU-resident, so their values + are not inspected here (which would synchronize the device). Callers must + ensure ``0 <= context_lengths[i] <= block_tables.shape[1] * block_size`` and + that every block-table entry used by a sequence is a physical block index in + ``[0, key_cache.shape[0])``. """ if context_partition_size != KV_COMPUTE_BLOCK: raise NotImplementedError( @@ -188,6 +213,10 @@ def pa_decode( raise NotImplementedError("pa_decode does not support attention sinks") if sliding_window not in (0, -1): raise NotImplementedError("pa_decode does not support sliding-window attention") + if not isinstance(ps, bool): + raise TypeError(f"ps must be a bool, got {type(ps).__name__}") + if not ps: + raise NotImplementedError("pa_decode does not support ps=False") if query_length < 1: raise ValueError(f"query_length must be positive, got {query_length}") if not 1 <= max_context_partition_num <= 64: @@ -196,6 +225,123 @@ def pa_decode( 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}" + ) + + # Keep NP==1 and NP>1 on the same domain. The tile requires a multiple of + # 64 and complete <=8-element Q-load pieces; the NP>1 reducer additionally + # caps its thread block (head_dim) at 1024. + 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}" + ) + if v_num_blocks != num_blocks: + raise ValueError( + f"key_cache and value_cache must have the same number of blocks, " + 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, @@ -211,60 +357,10 @@ def pa_decode( f"got {compute_type}" ) - # ``ps`` is retained for drop-in API compatibility. Both partitioning - # policies are represented by the caller-provided partition count. - del ps - - num_seqs = context_lengths.shape[0] - total_q_rows, num_q_heads, head_dim = query.shape - assert total_q_rows == num_seqs * query_length, ( - f"query.shape[0] ({total_q_rows}) must equal " - f"context_lengths.shape[0] * query_length ({num_seqs} * {query_length})" - ) - assert output.shape == query.shape, ( - f"output shape {tuple(output.shape)} must match " - f"query shape {tuple(query.shape)}" - ) - _, num_kv_heads, num_hgroups, block_size, hgroup_width = key_cache.shape - - assert num_hgroups == head_dim // 16 and hgroup_width == 16 - # Q staging loads each lane's head_dim//16 elements in <=8-wide pieces, so a - # chunk wider than 8 that isn't a multiple of 8 would silently drop its tail - # (head_dim 192/320/448/...). Reject those rather than return wrong results. - q_chunk = head_dim // 16 - if q_chunk > 8 and q_chunk % 8 != 0: - raise NotImplementedError( - f"pa_decode does not support head_dim={head_dim}: head_dim//16 " - f"({q_chunk}) must be <=8 or a multiple of 8" - ) - assert block_size in ( - 16, - 64, - 128, - ), 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_kv_heads, v_subblocks, v_head_dim, v_width = value_cache.shape - assert ( - v_head_dim == head_dim and v_width == 16 and v_subblocks == block_size // 16 - ), f"value_cache shape {tuple(value_cache.shape)} doesn't match block_size={block_size}, head_dim={head_dim}" - else: - _, v_num_kv_heads, v_head_dim, v_block_size = value_cache.shape - assert v_head_dim == head_dim and v_block_size == block_size, ( - f"value_cache shape {tuple(value_cache.shape)} doesn't match " - f"block_size={block_size}, head_dim={head_dim}" - ) - assert v_num_kv_heads == num_kv_heads - assert ( - num_q_heads % num_kv_heads == 0 - ), f"num_q_heads ({num_q_heads}) must be divisible by num_kv_heads ({num_kv_heads})" - assert ( - block_tables.dtype == torch.int32 - ), f"block_tables must be int32, got {block_tables.dtype}" - assert ( - context_lengths.dtype == torch.int32 - ), f"context_lengths must be int32, got {context_lengths.dtype}" + 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: @@ -272,24 +368,34 @@ def pa_decode( elif query.dtype == torch.float16: query_dtype = "f16" else: - raise ValueError(f"pa_decode only supports f16/bf16 query, got {query.dtype}") - assert ( - output.dtype == query.dtype - ), f"pa_decode requires output.dtype == query.dtype, got {output.dtype} vs {query.dtype}" - - assert ( - key_cache.dtype == expected_fp8_dtype - ), f"pa_decode requires {expected_fp8_dtype} key cache on {arch}, got {key_cache.dtype}" - assert ( - value_cache.dtype == expected_fp8_dtype - ), f"pa_decode requires {expected_fp8_dtype} value cache on {arch}, got {value_cache.dtype}" - - assert ( - query.stride(2) == 1 - ), f"pa_decode requires a contiguous head_dim axis, got strides {query.stride()}" - assert ( - output.stride(2) == 1 - ), f"pa_decode requires a contiguous output head_dim axis, got strides {output.stride()}" + 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 for name, tensor in ( @@ -299,16 +405,19 @@ def pa_decode( ("block_tables", block_tables), ("context_lengths", context_lengths), ): - assert ( - tensor.device == dev - ), f"{name} must be on the same device as query ({dev}), got {tensor.device}" + 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), ): - assert tensor.is_contiguous(), f"{name} must be contiguous" + if not tensor.is_contiguous(): + raise ValueError(f"{name} must be contiguous") def normalize_scale(scale, name): if scale is None: @@ -318,15 +427,17 @@ def normalize_scale(scale, name): if scale.numel() == 1: return scale.reshape(1) if scale.dim() == 4: - assert scale.shape[-1] == 1, ( - f"{name} must have a trailing singleton dimension, " - f"got shape {tuple(scale.shape)}" - ) + 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) - assert scale.dim() == 3, ( - f"{name} must be scalar or have shape " - "[num_blocks, num_kv_heads, block_size, 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): @@ -336,36 +447,46 @@ def normalize_scale(scale, name): 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 - assert per_token_kv == ( - value_scale_t.numel() > 1 - ), "key_scale and value_scale must both be per-tensor or both be per-token" - if per_token_kv: - assert ( - key_scale_t.shape == value_scale_t.shape - ), f"key_scale/value_scale shape mismatch: {tuple(key_scale_t.shape)} vs {tuple(value_scale_t.shape)}" - assert key_scale_t.shape == (key_cache.shape[0], num_kv_heads, block_size), ( - "per-token key_scale/value_scale must be [num_blocks, num_kv_heads, block_size] " - f"matching the KV cache, got {tuple(key_scale_t.shape)}" + 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)) - assert key_scale_t.stride(2) == 1, ( - f"per-token key_scale token dimension must be contiguous, " - f"got strides {key_scale_t.stride()}" - ) - assert value_scale_t.stride() == key_scale_t.stride(), ( - "per-token key_scale and value_scale must have matching strides, " - f"got {key_scale_t.stride()} vs {value_scale_t.stride()}" - ) + 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 - assert ( - key_scale_t.dtype == torch.float32 and key_scale_t.device == dev - ), f"key_scale tensor must be float32 on {dev}, got {key_scale_t.dtype} on {key_scale_t.device}" - assert ( - value_scale_t.dtype == torch.float32 and value_scale_t.device == dev - ), f"value_scale tensor must be float32 on {dev}, got {value_scale_t.dtype} on {value_scale_t.device}" + 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 pmax = max_logits @@ -391,8 +512,12 @@ def normalize_scale(scale, name): and query_length * query_group_size <= 16 ): num_cus = torch.cuda.get_device_properties(dev).multi_processor_count - workgroups = num_seqs * num_kv_heads * num_partitions - prefetch_v = num_cus < workgroups <= 2 * num_cus + prefetch_v = _workgroup_count_enables_v_prefetch( + num_seqs, + num_kv_heads, + num_partitions, + num_cus, + ) with torch.cuda.device(dev): compiled = compile_pa_decode_tile( @@ -429,25 +554,49 @@ def normalize_scale(scale, name): pout = torch.empty( *expected_scalar_shape, head_dim, dtype=output.dtype, device=dev ) - assert ( - pmax.shape == expected_scalar_shape - ), f"max_logits shape {tuple(pmax.shape)} != {expected_scalar_shape}" - assert ( - psum.shape == expected_scalar_shape - ), f"exp_sums shape {tuple(psum.shape)} != {expected_scalar_shape}" - assert pout.shape == ( - *expected_scalar_shape, - head_dim, - ), ( - f"temporary_output shape {tuple(pout.shape)} != " - f"{(*expected_scalar_shape, head_dim)}" - ) - assert pmax.dtype == torch.float32 and pmax.device == dev - assert psum.dtype == torch.float32 and psum.device == dev - assert pout.dtype == output.dtype and pout.device == dev - assert pmax.is_contiguous() - assert psum.is_contiguous() - assert pout.is_contiguous() + 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( diff --git a/aiter/ops/triton/gluon/pa_decode_gluon.py b/aiter/ops/triton/gluon/pa_decode_gluon.py index 30acce1fe21..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,12 +28,16 @@ FLYDSL_PS_REDUCE_AVAILABLE = True try: + 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 launch_pa_decode_ps_reduce_flydsl = None + is_pa_decode_ps_reduce_supported = None GLUON_JIT_KERNEL_ENABLED = True try: @@ -4023,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, @@ -4474,10 +4484,10 @@ def _paged_attention_decode_v2_with_dot_kernel_reshape_wrapper( 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, @@ -4535,9 +4545,25 @@ def _paged_attention_decode_v2_reduce_kernel_wrapper( return except ImportError: pass - try: - if not FLYDSL_PS_REDUCE_AVAILABLE: - raise ImportError("FlyDSL pa_decode reduce is unavailable") + 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, @@ -4562,32 +4588,31 @@ def _paged_attention_decode_v2_reduce_kernel_wrapper( 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, @@ -4633,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, @@ -4717,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 index 652a4b5b849..e9f3c2133c3 100644 --- a/op_tests/test_flydsl_pa_decode.py +++ b/op_tests/test_flydsl_pa_decode.py @@ -33,6 +33,11 @@ def _default_cuda_device(): SUPPORTED_GFX = ["gfx942", "gfx950"] KV_COMPUTE_BLOCK = 256 +# Match the corresponding FP8 Gluon accuracy contracts: fixed-length outputs +# use the tighter bound, while independently selected variable-length rows use +# the established varlen bound. +FIXED_LENGTH_ACCURACY_TOLERANCE = 5e-3 +VARIABLE_LENGTH_ACCURACY_TOLERANCE = 5e-2 # Pairwise coverage of the normal-accuracy axes in FlyDSL's PA regression test: # batches {3, 81, 128}, Q/KV heads {(4,1), (8,1), (16,1)}, head dims @@ -63,6 +68,15 @@ def _quant_dtype() -> torch.dtype: ) +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()}") + + def test_pa_decode_api_matches_gluon(): if pa_decode is None: pytest.skip("FlyDSL is not available") @@ -82,10 +96,7 @@ def test_pa_decode_api_matches_gluon(): def test_pa_decode_maps_gluon_buffers_and_scale_layout(monkeypatch): - if not torch.cuda.is_available(): - pytest.skip("ROCm is not available") - if pa_decode is None: - pytest.skip("FlyDSL is not available") + _require_gpu() pa_decode_module = importlib.import_module("aiter.ops.flydsl.pa_decode") gluon_module = importlib.import_module("aiter.ops.triton.gluon.pa_decode_gluon") @@ -209,6 +220,128 @@ def capture_reduce( assert dispatches == ["flydsl"] +def _call_gluon_reduce_wrapper(gluon_module, head_size): + output = torch.empty(1, 1, 1, 1, head_size, dtype=dtypes.bf16, device="cpu") + exp_sums = torch.empty(1, 1, 4, 1, dtype=dtypes.fp32, device="cpu") + max_logits = torch.empty_like(exp_sums) + logits = torch.empty(1, 1, 4, 1, head_size, dtype=dtypes.bf16, device="cpu") + context_lengths = torch.tensor([256], dtype=dtypes.i32, device="cpu") + + gluon_module._paged_attention_decode_v2_reduce_kernel_wrapper( + (1, 1, 1), + output, + exp_sums, + max_logits, + logits, + context_lengths, + None, + output.stride(0), + output.stride(1), + output.stride(2), + output.stride(3), + exp_sums.stride(0), + exp_sums.stride(1), + exp_sums.stride(2), + logits.stride(0), + logits.stride(1), + logits.stride(2), + logits.stride(3), + 1, + 1, + head_size, + 256, + PS=True, + context_partition_num=4, + ) + + +def test_gluon_unsupported_flydsl_reducer_falls_back(monkeypatch): + """An unsupported optional reducer must fall back instead of failing dispatch.""" + _require_gpu() + gluon_module = importlib.import_module("aiter.ops.triton.gluon.pa_decode_gluon") + triton_launches = [] + flydsl_launches = [] + + class TritonReduceStub: + def __getitem__(self, grid): + def launch(*args, **kwargs): + triton_launches.append((grid, args, kwargs)) + + return launch + + monkeypatch.setattr(gluon_module, "CXX_PS_REDUCE_AVAILABLE", False) + monkeypatch.setattr(gluon_module, "FLYDSL_PS_REDUCE_AVAILABLE", True) + monkeypatch.setattr( + gluon_module, + "launch_pa_decode_ps_reduce_flydsl", + lambda *args, **kwargs: flydsl_launches.append((args, kwargs)), + ) + monkeypatch.setattr( + gluon_module, + "paged_attention_decode_ps_reduce_kernel", + TritonReduceStub(), + ) + + _call_gluon_reduce_wrapper(gluon_module, head_size=96) + + assert not flydsl_launches + assert len(triton_launches) == 1 + assert triton_launches[0][0] == (1, 1, 1) + + +def test_gluon_supported_flydsl_reducer_does_not_mask_errors(monkeypatch): + """Errors from a selected FlyDSL reducer are not mistaken for unsupported.""" + _require_gpu() + gluon_module = importlib.import_module("aiter.ops.triton.gluon.pa_decode_gluon") + triton_launches = [] + + class TritonReduceStub: + def __getitem__(self, grid): + def launch(*args, **kwargs): + triton_launches.append((grid, args, kwargs)) + + return launch + + def fail_flydsl(*args, **kwargs): + raise ImportError("FlyDSL reducer failed internally") + + monkeypatch.setattr(gluon_module, "CXX_PS_REDUCE_AVAILABLE", False) + monkeypatch.setattr(gluon_module, "FLYDSL_PS_REDUCE_AVAILABLE", True) + monkeypatch.setattr(gluon_module, "launch_pa_decode_ps_reduce_flydsl", fail_flydsl) + monkeypatch.setattr( + gluon_module, + "paged_attention_decode_ps_reduce_kernel", + TritonReduceStub(), + ) + monkeypatch.setattr( + gluon_module.torch.cuda, "current_stream", lambda device: object() + ) + + with pytest.raises(ImportError, match="failed internally"): + _call_gluon_reduce_wrapper(gluon_module, head_size=128) + assert not triton_launches + + +@pytest.mark.parametrize( + "batch_size,expected", + [(32, False), (33, True), (64, True), (65, False)], +) +def test_v_prefetch_workgroup_interval(batch_size, expected): + """Warmup boundaries cover both compile-time prefetch variants.""" + if pa_decode is None: + pytest.skip("FlyDSL is not available") + pa_decode_module = importlib.import_module("aiter.ops.flydsl.pa_decode") + assert ( + pa_decode_module._workgroup_count_enables_v_prefetch( + batch_size, + num_kv_heads=1, + num_partitions=8, + num_compute_units=256, + ) + is expected + ) + + def run_torch( query: torch.Tensor, key_cache: torch.Tensor, @@ -418,8 +551,8 @@ def run_pa_decode_tile_case( err = checkAllclose( reference.to(dtypes.fp32), out.to(dtypes.fp32), - rtol=3e-2, - atol=3e-2, + rtol=FIXED_LENGTH_ACCURACY_TOLERANCE, + atol=FIXED_LENGTH_ACCURACY_TOLERANCE, tol_err_ratio=0.0, msg=f"{name}: pa_decode", ) @@ -430,35 +563,215 @@ def run_pa_decode_tile_case( return ret -@pytest.mark.parametrize("block_size", [16, 64, 128]) -def test_pa_decode_tile(block_size): - 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()}") +def _run_accuracy_case( + lengths, + *, + num_query_heads, + num_kv_heads, + head_dim, + block_size, + query_dtype, + num_partitions, + tolerance, +): + """Run PA against a dequantized-FP8 torch reference without benchmarking.""" + _require_gpu() + if num_query_heads % num_kv_heads: + raise ValueError("num_query_heads must be divisible by num_kv_heads") + + torch.manual_seed(0) + batch_size = len(lengths) + pages_per_sequence = [ + (context_length + block_size - 1) // block_size for context_length in lengths + ] + max_pages = max(pages_per_sequence) + num_blocks = sum(pages_per_sequence) + + query = torch.empty( + batch_size, num_query_heads, head_dim, dtype=query_dtype + ).uniform_(-0.5, 0.5) + key = torch.empty( + num_blocks, num_kv_heads, block_size, head_dim, dtype=query_dtype + ).uniform_(-0.5, 0.5) + value = torch.empty( + num_blocks, num_kv_heads, head_dim, block_size, dtype=query_dtype + ).uniform_(-0.5, 0.5) + + quant_dtype = _quant_dtype() + key_quant, key_scale = per_tensor_quant(key, quant_dtype=quant_dtype) + value_quant, value_scale = per_tensor_quant(value, quant_dtype=quant_dtype) + key_cache = ( + key_quant.view(num_blocks, num_kv_heads, block_size, head_dim // 16, 16) + .permute(0, 1, 3, 2, 4) + .contiguous() + ) + value_cache = value_quant.contiguous() + + block_tables = torch.zeros( + batch_size, max_pages, dtype=dtypes.i32, device=query.device + ) + next_page = 0 + for seq_idx, page_count in enumerate(pages_per_sequence): + block_tables[seq_idx, :page_count] = torch.arange( + next_page, + next_page + page_count, + dtype=dtypes.i32, + device=query.device, + ) + next_page += page_count + context_lengths = torch.tensor(lengths, dtype=dtypes.i32, device=query.device) - result = run_pa_decode_tile_case( - batch_size=3, + reference = run_torch( + query, + key_quant, + value_quant, + block_tables, + context_lengths, + key_scale, + value_scale, + ) + output = torch.empty_like(query) + torch.ops.aiter.pa_decode_flydsl( + output, + query, + key_cache, + value_cache, + context_lengths, + block_tables, + head_dim**-0.5, + 1, + num_partitions, + KV_COMPUTE_BLOCK, + quant_dtype, + None, + key_scale, + value_scale, + ) + + torch.testing.assert_close( + output.float(), + reference.float(), + atol=tolerance, + rtol=tolerance, + ) + + +@pytest.mark.parametrize( + "block_size,query_dtype,num_kv_heads,num_partitions", + [ + pytest.param(16, dtypes.bf16, 1, 1, id="page16-bf16-hkv1-np1"), + pytest.param(64, dtypes.fp16, 2, 1, id="page64-fp16-hkv2-np1"), + pytest.param(128, dtypes.bf16, 2, 4, id="page128-bf16-hkv2-np4"), + pytest.param(128, dtypes.fp16, 1, 4, id="page128-fp16-hkv1-np4"), + ], +) +def test_pa_decode_fixed_length_accuracy( + block_size, query_dtype, num_kv_heads, num_partitions +): + _run_accuracy_case( + [257, 257, 257], num_query_heads=8, - num_kv_heads=1, + num_kv_heads=num_kv_heads, head_dim=128, - context_length=257, block_size=block_size, - dtype=dtypes.bf16, + query_dtype=query_dtype, + num_partitions=num_partitions, + tolerance=FIXED_LENGTH_ACCURACY_TOLERANCE, ) - assert result["partitions"] > 1 - assert result["flydsl err"] == 0 -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.mark.parametrize("query_dtype", [dtypes.bf16, dtypes.fp16]) +def test_pa_decode_variable_length_accuracy(query_dtype): + _run_accuracy_case( + [1, 127, 257, 1027], + num_query_heads=8, + num_kv_heads=2, + head_dim=128, + block_size=128, + query_dtype=query_dtype, + num_partitions=4, + tolerance=VARIABLE_LENGTH_ACCURACY_TOLERANCE, + ) + + +def _valid_pa_decode_arguments(head_dim=128, num_partitions=1): + """Small valid call used by host-contract tests that fail before compile.""" + quant_dtype = _quant_dtype() + query = torch.empty(1, 8, head_dim, dtype=dtypes.bf16, device="cuda") + scale = torch.ones(1, dtype=dtypes.fp32, device=query.device) + return { + "output": torch.empty_like(query), + "query": query, + "key_cache": torch.empty( + 1, 1, head_dim // 16, 16, 16, dtype=quant_dtype, device=query.device + ), + "value_cache": torch.empty( + 1, 1, head_dim, 16, dtype=quant_dtype, device=query.device + ), + "context_lengths": torch.tensor([16], dtype=dtypes.i32, device=query.device), + "block_tables": torch.tensor([[0]], dtype=dtypes.i32, device=query.device), + "softmax_scale": head_dim**-0.5, + "query_length": 1, + "max_context_partition_num": num_partitions, + "compute_type": quant_dtype, + "key_scale": scale, + "value_scale": scale, + } + + +def test_ps_false_is_rejected(): + _require_gpu() + with pytest.raises(NotImplementedError, match="ps=False"): + pa_decode(**_valid_pa_decode_arguments(), ps=False) + + +@pytest.mark.parametrize( + "invalid_input,error_match", + [ + pytest.param("query_rank", "query", id="query-rank"), + pytest.param("output_rank", "output", id="output-rank"), + pytest.param("key_cache_rank", "key_cache", id="key-cache-rank"), + pytest.param("value_cache_rank", "value_cache", id="value-cache-rank"), + pytest.param("context_lengths_rank", "context_lengths", id="context-rank"), + pytest.param("block_tables_rank", "block_tables", id="block-table-rank"), + pytest.param("query_device", "CUDA device", id="query-device"), + pytest.param("zero_query_heads", "at least one head", id="zero-query-heads"), + pytest.param("block_table_rows", "block_tables", id="block-table-rows"), + pytest.param("kv_block_count", "block", id="kv-block-count"), + pytest.param("kv_head_count", "head", id="kv-head-count"), + ], +) +def test_invalid_tensor_structure_is_rejected(invalid_input, error_match): + _require_gpu() + arguments = _valid_pa_decode_arguments() + if invalid_input == "query_rank": + arguments["query"] = arguments["query"].squeeze(0) + elif invalid_input == "output_rank": + arguments["output"] = arguments["output"].squeeze(0) + elif invalid_input == "key_cache_rank": + arguments["key_cache"] = arguments["key_cache"].squeeze(0) + elif invalid_input == "value_cache_rank": + arguments["value_cache"] = arguments["value_cache"].squeeze(0) + elif invalid_input == "context_lengths_rank": + arguments["context_lengths"] = arguments["context_lengths"].reshape(1, 1) + elif invalid_input == "block_tables_rank": + arguments["block_tables"] = arguments["block_tables"].flatten() + elif invalid_input == "query_device": + arguments["query"] = arguments["query"].cpu() + elif invalid_input == "zero_query_heads": + arguments["query"] = arguments["query"][:, :0] + arguments["output"] = arguments["output"][:, :0] + elif invalid_input == "block_table_rows": + arguments["block_tables"] = arguments["block_tables"].repeat(2, 1) + elif invalid_input == "kv_block_count": + arguments["value_cache"] = arguments["value_cache"].repeat(2, 1, 1, 1) + elif invalid_input == "kv_head_count": + arguments["value_cache"] = arguments["value_cache"].repeat(1, 2, 1, 1) + else: + raise AssertionError(f"unknown invalid input: {invalid_input}") + + with pytest.raises(ValueError, match=error_match): + pa_decode(**arguments) def _adversarial_case( @@ -721,19 +1034,54 @@ def test_extreme_tail_value_scale_is_ignored(context_length, query_length): _assert_matches(output, reference) -@pytest.mark.parametrize("head_dim", [64, 128, 256]) -def test_supported_head_dims_are_accurate(head_dim): +@pytest.mark.parametrize( + "head_dim,num_partitions", + [(64, 1), (256, 1), (1024, 1), (1024, 4)], +) +def test_additional_supported_head_dims_are_accurate(head_dim, num_partitions): _require_gpu() - output, reference = _adversarial_case(head_dim=head_dim, context_length=1027) + output, reference = _adversarial_case( + head_dim=head_dim, + context_length=64, + num_partitions=num_partitions, + ) _assert_matches(output, reference) -@pytest.mark.parametrize("head_dim", [192, 320]) -def test_unsupported_head_dim_is_rejected(head_dim): - """head_dim//16 above 8 and not a multiple of 8 leaves the Q tail unloaded.""" +@pytest.mark.parametrize("num_partitions", [1, 4]) +@pytest.mark.parametrize("head_dim", [64, 128, 256, 1024]) +def test_supported_head_dim_boundaries_are_accepted( + monkeypatch, head_dim, num_partitions +): + """The public support check is identical before direct and reduced paths.""" + _require_gpu() + pa_decode_module = importlib.import_module("aiter.ops.flydsl.pa_decode") + compile_arguments = [] + + def capture_compile(**kwargs): + compile_arguments.append(kwargs) + return {"launch": object()} + + monkeypatch.setattr(pa_decode_module, "compile_pa_decode_tile", capture_compile) + monkeypatch.setattr(pa_decode_module, "_run_compiled", lambda *args: None) + monkeypatch.setattr(pa_decode_module, "ptr_arg", lambda tensor, dtype=None: tensor) + monkeypatch.setattr( + pa_decode_module, "launch_pa_decode_ps_reduce", lambda *args, **kwargs: None + ) + + pa_decode(**_valid_pa_decode_arguments(head_dim, num_partitions)) + assert len(compile_arguments) == 1 + assert compile_arguments[0]["head_dim"] == head_dim + assert compile_arguments[0]["num_partitions"] == num_partitions + + +@pytest.mark.parametrize("num_partitions", [1, 4]) +@pytest.mark.parametrize("head_dim", [96, 192, 1152]) +def test_unsupported_head_dim_is_rejected(head_dim, num_partitions): + """NP=1 and NP>1 expose the same public head-dimension support domain.""" _require_gpu() with pytest.raises(NotImplementedError, match="head_dim"): - _adversarial_case(head_dim=head_dim, context_length=64) + pa_decode(**_valid_pa_decode_arguments(head_dim, num_partitions)) def _constant_page_decode_case( @@ -1032,6 +1380,7 @@ def test_mtp3_page128_partition_boundaries( context_length, num_partitions, trans_v, zero_query ): """MTP rows retain independent softmax state across partial/empty partitions.""" + _require_gpu() output, reference = _adversarial_case( head_dim=128, context_length=context_length, @@ -1048,6 +1397,7 @@ def test_mtp3_page128_partition_boundaries( @pytest.mark.parametrize("all_empty", [False, True]) def test_page128_decode_prefetch_mixed_contexts(all_empty): """A two-CTA-per-CU decode grid handles empty partitions and poisoned padding.""" + _require_gpu() lengths = [0] * 64 if all_empty else [0, 1, 63, 127, 128, 129, 257, 2051] * 8 _constant_page_decode_case( lengths, @@ -1067,6 +1417,7 @@ def test_page128_reduction_rows_across_m_tiles( query_group_size, query_length, context_length ): """Aligned scratch rows stay independent across padded M-tiles and partitions.""" + _require_gpu() output, reference = _adversarial_case( head_dim=128, context_length=context_length, From d39ba45279e305e553a8c2069903fd8e5838291c Mon Sep 17 00:00:00 2001 From: fsx950223 Date: Mon, 14 Sep 2026 04:41:41 +0000 Subject: [PATCH 26/66] feat(flydsl): extend PA reduction to 256 partitions --- aiter/ops/flydsl/kernels/pa_decode_reduce.py | 215 ++++++++++++++----- aiter/ops/flydsl/pa_decode.py | 16 +- op_tests/test_flydsl_pa_decode_reduce_cpu.py | 139 ++++++++++++ 3 files changed, 306 insertions(+), 64 deletions(-) create mode 100644 op_tests/test_flydsl_pa_decode_reduce_cpu.py diff --git a/aiter/ops/flydsl/kernels/pa_decode_reduce.py b/aiter/ops/flydsl/kernels/pa_decode_reduce.py index dff93eb040d..a208e8cc8fb 100644 --- a/aiter/ops/flydsl/kernels/pa_decode_reduce.py +++ b/aiter/ops/flydsl/kernels/pa_decode_reduce.py @@ -9,6 +9,8 @@ import flydsl.expr as fx from flydsl.expr.typing import T +MAX_CONTEXT_PARTITIONS = 256 + @lru_cache(maxsize=256) def compile_pa_decode_ps_reduce( @@ -23,13 +25,14 @@ def compile_pa_decode_ps_reduce( """Build the partitioned-softmax reduction used by ``pa_decode``. Each wave reduces the partition statistics independently, then reuses the - per-partition weights to combine one output element per thread. The host - API limits the partition count to one wave, so the LDS path from the more - general Gluon implementation is unnecessary here. + per-partition weights to combine one output element per thread. Counts up + to one wave retain the original one-partition-per-lane path. For larger + counts, each lane owns up to four partitions separated by a wave, keeping + all weights in registers and avoiding an LDS round trip. """ - if not 1 <= max_context_partition_num <= 64: + if not 1 <= max_context_partition_num <= MAX_CONTEXT_PARTITIONS: raise ValueError( - "max_context_partition_num must be in [1, 64], " + f"max_context_partition_num must be in [1, {MAX_CONTEXT_PARTITIONS}], " f"got {max_context_partition_num}" ) if head_size <= 0 or head_size > 1024 or head_size % 64: @@ -125,65 +128,159 @@ def _wave_reduce_sum(value): ) return reduced - 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 = ( - batch_idx * stride_exp_sums_seq - + 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) - - global_max = _wave_reduce_max(part_max) - 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_value = fx.Float32(sink_token[kv_head_idx * c_qgs + group_idx]) - sink_scale = (global_max > neg_inf).select( - fx.exp2( - (sink_value - safe_global_max) * c_log2e, - fastmath="fast", - ), + if fx.const_expr(max_context_partition_num <= warp_size): + # Preserve the original <=64 path verbatim: one partition per lane + # and one bpermute source register per wave. + 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 = ( + batch_idx * stride_exp_sums_seq + + 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) + + global_max = _wave_reduce_max(part_max) + 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, ) - global_exp_sum = global_exp_sum + sink_scale - safe_global_exp_sum = (global_exp_sum > zero_f).select(global_exp_sum, one_f) - weight_local = scaled_sum / safe_global_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, + scaled_sum = part_sum * part_scale + global_exp_sum = _wave_reduce_sum(scaled_sum) + if fx.const_expr(use_sinks): + sink_value = fx.Float32(sink_token[kv_head_idx * c_qgs + group_idx]) + sink_scale = (global_max > neg_inf).select( + fx.exp2( + (sink_value - safe_global_max) * c_log2e, + fastmath="fast", + ), + zero_f, ) + global_exp_sum = global_exp_sum + sink_scale + safe_global_exp_sum = (global_exp_sum > zero_f).select( + global_exp_sum, one_f ) - weight = weight_i32.bitcast(fx.Float32) - logits_offset = ( - batch_idx * stride_logits_seq - + kv_head_idx * stride_logits_head - + c_part_idx * stride_logits_part - + eqgs_idx * stride_logits_group - + tid + weight_local = scaled_sum / safe_global_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 = ( + batch_idx * stride_logits_seq + + 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: + # A wave covers several 64-partition chunks. Lane ``l`` owns + # partitions l, l+64, l+128, and l+192 (as present). Reduce the + # lane-local maxima/sums before the usual wave reduction; later, + # select the corresponding local weight and broadcast from + # ``part_idx % 64``. This stays register-only through NP=256. + 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): + part_idx = lane + fx.Int32(chunk_idx * warp_size) + lane_in_range = part_idx < c_part_num + safe_part_idx = lane_in_range.select(part_idx, zero_i) + stats_offset = ( + batch_idx * stride_exp_sums_seq + + kv_head_idx * stride_exp_sums_head + + safe_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) + part_sums.append(part_sum) + part_maxes.append(part_max) + lane_max = lane_max.maximumf(part_max) + + global_max = _wave_reduce_max(lane_max) + 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] + 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_value = fx.Float32(sink_token[kv_head_idx * c_qgs + group_idx]) + sink_scale = (global_max > neg_inf).select( + fx.exp2( + (sink_value - safe_global_max) * c_log2e, + fastmath="fast", + ), + zero_f, + ) + global_exp_sum = global_exp_sum + sink_scale + safe_global_exp_sum = (global_exp_sum > zero_f).select( + global_exp_sum, one_f ) - part_logits = fx.Float32(logits[logits_offset]) - acc = acc + part_logits * weight + + acc = zero_f + for chunk_idx in fx.range_constexpr(partitions_per_lane): + weight_local_i32 = ( + scaled_sums[chunk_idx] / safe_global_exp_sum + ).bitcast(fx.Int32) + chunk_base = chunk_idx * warp_size + chunk_size = min(warp_size, max_context_partition_num - chunk_base) + 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 = ( + batch_idx * stride_logits_seq + + 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 query_idx = eqgs_idx // c_qgs output_offset = ( diff --git a/aiter/ops/flydsl/pa_decode.py b/aiter/ops/flydsl/pa_decode.py index a6aecbb7a8b..f9287ba5107 100644 --- a/aiter/ops/flydsl/pa_decode.py +++ b/aiter/ops/flydsl/pa_decode.py @@ -34,7 +34,10 @@ from aiter.jit.utils.chip_info import get_gfx_runtime -from .kernels.pa_decode_reduce import compile_pa_decode_ps_reduce +from .kernels.pa_decode_reduce import ( + MAX_CONTEXT_PARTITIONS, + compile_pa_decode_ps_reduce, +) from .kernels.pa_decode_tile import KV_COMPUTE_BLOCK, compile_pa_decode_tile from .kernels.tensor_shim import _run_compiled, ptr_arg from .kernels.utils import cdiv @@ -110,8 +113,11 @@ def launch_pa_decode_ps_reduce( context_partition_num: int, stream: torch.cuda.Stream, ) -> None: - if context_partition_num > 64: - raise ImportError("FlyDSL pa_decode reduce supports at most 64 partitions") + if context_partition_num > MAX_CONTEXT_PARTITIONS: + raise ImportError( + "FlyDSL pa_decode reduce supports at most " + f"{MAX_CONTEXT_PARTITIONS} partitions" + ) use_sinks = sink_token is not None compiled = compile_pa_decode_ps_reduce( max_context_partition_num=context_partition_num, @@ -205,9 +211,9 @@ def pa_decode( raise NotImplementedError("pa_decode does not support sliding-window attention") if query_length < 1: raise ValueError(f"query_length must be positive, got {query_length}") - if not 1 <= max_context_partition_num <= 64: + if not 1 <= max_context_partition_num <= MAX_CONTEXT_PARTITIONS: raise ValueError( - "max_context_partition_num must be in [1, 64], " + f"max_context_partition_num must be in [1, {MAX_CONTEXT_PARTITIONS}], " f"got {max_context_partition_num}" ) diff --git a/op_tests/test_flydsl_pa_decode_reduce_cpu.py b/op_tests/test_flydsl_pa_decode_reduce_cpu.py new file mode 100644 index 00000000000..ed27726026b --- /dev/null +++ b/op_tests/test_flydsl_pa_decode_reduce_cpu.py @@ -0,0 +1,139 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""CPU-only coverage for the FlyDSL PA partition reducer.""" + +import importlib.util +import math +import random +import sys +from pathlib import Path + +import pytest + +pytest.importorskip("flydsl") + + +def _load_reduce_module(): + """Load the self-contained kernel module without importing GPU-facing aiter.""" + path = ( + Path(__file__).parents[1] + / "aiter" + / "ops" + / "flydsl" + / "kernels" + / "pa_decode_reduce.py" + ) + name = "_test_flydsl_pa_decode_reduce" + spec = importlib.util.spec_from_file_location(name, path) + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module + spec.loader.exec_module(module) + return module + + +REDUCE = _load_reduce_module() + + +@pytest.mark.parametrize("num_partitions", [1, 63, 64, 65, 128, 192, 256]) +def test_reduce_kernel_accepts_supported_partition_counts(num_partitions): + """Kernel construction is CPU-only and covers both generated code paths.""" + compiled = REDUCE.compile_pa_decode_ps_reduce( + max_context_partition_num=num_partitions, + head_size=128, + output_dtype_str="bf16", + logits_dtype_str="bf16", + sink_dtype_str="bf16", + use_sinks=False, + ) + assert set(compiled) == {"launch", "kernel"} + + +@pytest.mark.parametrize("num_partitions", [0, 257]) +def test_reduce_kernel_rejects_unsupported_partition_counts(num_partitions): + with pytest.raises(ValueError, match=r"must be in \[1, 256\]"): + REDUCE.compile_pa_decode_ps_reduce( + max_context_partition_num=num_partitions, + head_size=128, + output_dtype_str="bf16", + logits_dtype_str="bf16", + sink_dtype_str="bf16", + use_sinks=False, + ) + + +def _flat_partition_reduce(exp_sums, max_logits, partials): + global_max = max(max_logits) + scaled_sums = [ + part_sum * math.exp(part_max - global_max) if part_max != -math.inf else 0.0 + for part_sum, part_max in zip(exp_sums, max_logits) + ] + denominator = sum(scaled_sums) + if denominator == 0.0: + return [0.0] * len(partials[0]) + return [ + sum( + partials[part][element] * scaled_sums[part] for part in range(len(exp_sums)) + ) + / denominator + for element in range(len(partials[0])) + ] + + +def _lane_striped_partition_reduce(exp_sums, max_logits, partials): + """Model the >64 kernel: lane owns p, p+64, p+128, and p+192.""" + wave_size = 64 + lane_maxes = [-math.inf] * wave_size + for part, part_max in enumerate(max_logits): + lane = part % wave_size + lane_maxes[lane] = max(lane_maxes[lane], part_max) + global_max = max(lane_maxes) + + scaled_sums = [] + lane_sums = [0.0] * wave_size + for part, (part_sum, part_max) in enumerate(zip(exp_sums, max_logits)): + scaled_sum = ( + part_sum * math.exp(part_max - global_max) if part_max != -math.inf else 0.0 + ) + scaled_sums.append(scaled_sum) + lane_sums[part % wave_size] += scaled_sum + + denominator = sum(lane_sums) + if denominator == 0.0: + return [0.0] * len(partials[0]) + accumulators = [0.0] * len(partials[0]) + for chunk_base in range(0, len(exp_sums), wave_size): + chunk_end = min(chunk_base + wave_size, len(exp_sums)) + for part in range(chunk_base, chunk_end): + weight = scaled_sums[part] / denominator + for element, value in enumerate(partials[part]): + accumulators[element] += value * weight + return accumulators + + +@pytest.mark.parametrize("num_partitions", [65, 127, 128, 129, 192, 255, 256]) +def test_lane_striped_reduce_matches_flat_reference(num_partitions): + rng = random.Random(num_partitions) + max_logits = [rng.uniform(-20.0, 5.0) for _ in range(num_partitions)] + exp_sums = [rng.uniform(0.01, 10.0) for _ in range(num_partitions)] + partials = [ + [rng.uniform(-2.0, 2.0) for _ in range(7)] for _ in range(num_partitions) + ] + + # Exercise masked lanes in the final chunk and neutral empty partitions. + for part in range(11, num_partitions, 37): + max_logits[part] = -math.inf + exp_sums[part] = 0.0 + + expected = _flat_partition_reduce(exp_sums, max_logits, partials) + actual = _lane_striped_partition_reduce(exp_sums, max_logits, partials) + assert actual == pytest.approx(expected, rel=2e-15, abs=2e-15) + + +@pytest.mark.parametrize("num_partitions", [128, 192, 256]) +def test_lane_striped_reduce_all_empty_is_zero(num_partitions): + exp_sums = [0.0] * num_partitions + max_logits = [-math.inf] * num_partitions + # The tile kernel writes finite neutral scratch for empty partitions. + partials = [[0.0] * 7 for _ in range(num_partitions)] + assert _lane_striped_partition_reduce(exp_sums, max_logits, partials) == [0.0] * 7 From dbfff5fdb37534bb07690308c28b7551265dbf6a Mon Sep 17 00:00:00 2001 From: fsx950223 Date: Mon, 14 Sep 2026 04:48:22 +0000 Subject: [PATCH 27/66] fix(flydsl): mask partial PA reduction waves --- aiter/ops/flydsl/kernels/pa_decode_reduce.py | 80 +++++++++++++------- op_tests/test_flydsl_pa_decode_reduce_cpu.py | 5 +- 2 files changed, 58 insertions(+), 27 deletions(-) diff --git a/aiter/ops/flydsl/kernels/pa_decode_reduce.py b/aiter/ops/flydsl/kernels/pa_decode_reduce.py index a208e8cc8fb..1fd7a2f091f 100644 --- a/aiter/ops/flydsl/kernels/pa_decode_reduce.py +++ b/aiter/ops/flydsl/kernels/pa_decode_reduce.py @@ -129,24 +129,40 @@ def _wave_reduce_sum(value): return reduced if fx.const_expr(max_context_partition_num <= warp_size): - # Preserve the original <=64 path verbatim: one partition per lane - # and one bpermute source register per wave. - 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) + # Exact powers of two have no inactive lanes inside their reduction + # subgroup. Keep that original path unchanged; only partial + # subgroups need an EXEC-masked load to avoid carrying a predicate + # across the shuffle sequence. + if fx.const_expr(max_context_partition_num == reduce_width): + 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 = ( + batch_idx * stride_exp_sums_seq + + 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 = ( batch_idx * stride_exp_sums_seq + kv_head_idx * stride_exp_sums_head - + part_idx * stride_exp_sums_part + + lane * 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) + 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) safe_global_max = (global_max > neg_inf).select(global_max, zero_f) @@ -205,19 +221,31 @@ def _wave_reduce_sum(value): part_maxes = [] lane_max = neg_inf for chunk_idx in fx.range_constexpr(partitions_per_lane): - part_idx = lane + fx.Int32(chunk_idx * warp_size) - lane_in_range = part_idx < c_part_num - safe_part_idx = lane_in_range.select(part_idx, zero_i) - stats_offset = ( - batch_idx * stride_exp_sums_seq - + kv_head_idx * stride_exp_sums_head - + safe_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) + 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): + stats_offset = ( + batch_idx * stride_exp_sums_seq + + 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) + stats_offset = ( + batch_idx * stride_exp_sums_seq + + 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) diff --git a/op_tests/test_flydsl_pa_decode_reduce_cpu.py b/op_tests/test_flydsl_pa_decode_reduce_cpu.py index ed27726026b..dbc52cc9bc2 100644 --- a/op_tests/test_flydsl_pa_decode_reduce_cpu.py +++ b/op_tests/test_flydsl_pa_decode_reduce_cpu.py @@ -111,7 +111,10 @@ def _lane_striped_partition_reduce(exp_sums, max_logits, partials): return accumulators -@pytest.mark.parametrize("num_partitions", [65, 127, 128, 129, 192, 255, 256]) +@pytest.mark.parametrize( + "num_partitions", + [26, 30, 32, 34, 36, 40, 64, 65, 96, 127, 128, 129, 160, 192, 255, 256], +) def test_lane_striped_reduce_matches_flat_reference(num_partitions): rng = random.Random(num_partitions) max_logits = [rng.uniform(-20.0, 5.0) for _ in range(num_partitions)] From cbec7698afc3d76d30c1618b7db34ac8dae8ce47 Mon Sep 17 00:00:00 2001 From: fsx950223 Date: Mon, 14 Sep 2026 04:44:16 +0000 Subject: [PATCH 28/66] perf(flydsl): streamline single-tile PA softmax --- aiter/ops/flydsl/kernels/pa_decode_tile.py | 119 +++++++++++++++++---- 1 file changed, 96 insertions(+), 23 deletions(-) diff --git a/aiter/ops/flydsl/kernels/pa_decode_tile.py b/aiter/ops/flydsl/kernels/pa_decode_tile.py index bc596bd392a..22749ed141a 100644 --- a/aiter/ops/flydsl/kernels/pa_decode_tile.py +++ b/aiter/ops/flydsl/kernels/pa_decode_tile.py @@ -153,6 +153,10 @@ def compile_pa_decode_tile( EARLY_V = prefetch_v and (tune_page128 or tune_page16_vpipe) and M_TILES == 1 PAGE16_VPIPE = prefetch_v and tune_page16_vpipe and M_TILES == 1 PAGE16_VPIPE_IGLP = PAGE16_VPIPE and prefetch_v_iglp + # Keep the single-M-tile softmax rewrite confined to the two opt-in + # MiniMax-M3 experiments. In particular, every default/non-target kernel + # continues to instantiate the original instruction stream. + M1_SCALE_BEFORE_MASK = PAGE16_VPIPE or match_gluon_numerics P_BUFFERS = 2 if tune_page128 and M_TILES == 3 else 1 # PV layout: V=A, P=B -> output [head-dim (row), query-row (col=lane16)], # generalized over head_dim via the VHE_CHUNKS loop. @@ -1205,31 +1209,73 @@ def _lmax_off_m(m): thr = fx.Vector.from_elements( [n_valid_tile - base_tok_f], dtype=fx.Float32 ).broadcast_to(4) - neg4 = fx.Vector.filled(4, -1e30, fx.Float32) - # per_token_kv: K-scale varies per token, so fold it in BEFORE the max-reduce. + neg4 = fx.Vector.filled( + 4, + float("-inf") if M1_SCALE_BEFORE_MASK else -1e30, + fx.Float32, + ) + # The tuned path follows the M_TILES>1 numerical ordering: + # fold every positive score scale into finite logits before + # selecting -inf for invalid tokens. This makes the reduced + # max directly usable by pass 2 and avoids `-inf * 0` for an + # all-zero Q row. The per-token branch is retained here so + # this ordering stays correct if the gate is widened later. + 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_frags.append(frag_Ss[a] * k_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( + (_ct[a] < thr).select(scaled_frag, neg4) + ) + else: + scaled_frags.append(scaled_frag) else: - scaled_frags = frag_Ss + if const_expr(M1_SCALE_BEFORE_MASK): + masked_chunks = [ + (_ct[a] < thr).select(frag_Ss[a] * scale_b, neg4) + for a in range_constexpr(NCHUNK) + ] + else: + scaled_frags = frag_Ss # Reused in pass 2 below, halving the mask instruction count. - masked_chunks = [ - (_ct[a] < thr).select(scaled_frags[a], neg4) - for a in range_constexpr(NCHUNK) - ] + if const_expr(not M1_SCALE_BEFORE_MASK): + masked_chunks = [ + (_ct[a] < 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): - pm = fx.maxnumf(pm, masked_chunks[a].reduce(ReductionOp.MAX)) + 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): - pm = fx.maxnumf(pm, pm.shuffle_xor(sh, WAVE)) + 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 * scale + sLmax_off, + lane16, + warp, + pm if const_expr(M1_SCALE_BEFORE_MASK) else pm * scale, ) # redundant across the 4 lanes sharing this qhead # per_token_kv: max V-scale for the per-tile fp8 normalization. # Slots at/after context_len are uninitialised, so an arbitrarily @@ -1266,22 +1312,49 @@ def _lmax_off_m(m): [norm_factor], dtype=fx.Float32 ).broadcast_to(4) # pass 2: global max over warps -> exp -> fp8 P pack (-> sP) -> sum - m_new = fx.maxnumf( - m_prev, _ld_lw_row(sLmax_off, lane16).reduce(ReductionOp.MAX) - ) + if const_expr(M1_SCALE_BEFORE_MASK): + m_new = fx.maxnumf( + m_prev, + _ld_lw_row(sLmax_off, lane16).reduce( + ReductionOp.MAX, fastmath=fm_nnan + ), + fastmath=fm_nnan, + ) + else: + m_new = fx.maxnumf( + m_prev, + _ld_lw_row(sLmax_off, lane16).reduce(ReductionOp.MAX), + ) + # Keep -inf in the loop-carried/persisted max for an empty + # partition, but use zero as the exponent reference so + # -inf-(-inf) cannot manufacture NaNs in P or the correction. + 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( - [m_new], dtype=fx.Float32 + [softmax_max], dtype=fx.Float32 ).broadcast_to(4) ls = fx.Float32(0.0) words = [] - zero4_p = fx.Vector.filled(4, 0.0, fx.Float32) + if const_expr(not M1_SCALE_BEFORE_MASK): + zero4_p = fx.Vector.filled(4, 0.0, fx.Float32) for a in range_constexpr(NCHUNK): - # re-mask Pa so a fully-masked chunk contributes exactly 0 - valid_a = masked_chunks[a] > fx.Vector.filled(4, -1e29, fx.Float32) - Pa = valid_a.select( - fx.Vector(exp2_f32_fast(masked_chunks[a] * scale - m_new_b)), - zero4_p, - ) + if const_expr(M1_SCALE_BEFORE_MASK): + # Invalid lanes are -inf, hence exp2(-inf-safe_max)=0 + # without a second validity compare/select. + Pa = fx.Vector(exp2_f32_fast(masked_chunks[a] - m_new_b)) + else: + # Legacy path: re-mask Pa so a fully-masked chunk + # contributes exactly 0. + valid_a = masked_chunks[a] > fx.Vector.filled( + 4, -1e29, fx.Float32 + ) + Pa = valid_a.select( + fx.Vector( + exp2_f32_fast(masked_chunks[a] * scale - m_new_b) + ), + zero4_p, + ) ls = ls + Pa.reduce(ReductionOp.ADD) if const_expr(per_token_kv): v_scale_this = ( @@ -1311,7 +1384,7 @@ def _lmax_off_m(m): ls = ls + ls.shuffle_xor(sh, WAVE) # PV (V=A, P=B) -> output [head-dim, query-row=lane16]; same as # the phase-split path. - corr_reg = fx.Float32(exp2_amdgcn_scalar(m_prev - m_new)) + corr_reg = fx.Float32(exp2_amdgcn_scalar(m_prev - softmax_max)) if rgroup == 0: _st_lw(sLsum_off, lane16, warp, ls) gpu.barrier() From f03c6af852459c029e38e5dbf1fb49ce12436f72 Mon Sep 17 00:00:00 2001 From: fsx950223 Date: Mon, 14 Sep 2026 06:30:31 +0000 Subject: [PATCH 29/66] perf(flydsl): tune native page128 long-context decode --- aiter/ops/flydsl/kernels/pa_decode_tile.py | 18 +++++------------- aiter/ops/flydsl/pa_decode.py | 11 +++++++---- 2 files changed, 12 insertions(+), 17 deletions(-) diff --git a/aiter/ops/flydsl/kernels/pa_decode_tile.py b/aiter/ops/flydsl/kernels/pa_decode_tile.py index 22749ed141a..81c83744a86 100644 --- a/aiter/ops/flydsl/kernels/pa_decode_tile.py +++ b/aiter/ops/flydsl/kernels/pa_decode_tile.py @@ -116,13 +116,13 @@ def compile_pa_decode_tile( assert ( is_gfx950 and head_dim == 128 - and block_size == 16 + and block_size in (16, 128) and trans_v and not per_token_kv and query_dtype == "bf16" and query_length == 1 and query_group_size in (8, 16) - ), "matched Gluon numerics is restricted to the MiniMax-M3 page-16 shapes" + ), "matched Gluon numerics is restricted to the MiniMax-M3 page-16/page-128 shapes" assert ( head_dim % 64 == 0 @@ -693,17 +693,11 @@ def _quant_q_row(m, qi, gs_head, q_row_off): _f32_to_fp8_words(q_units[u].to(fx.Float32)), ) else: - absmax = ( - fmath.absf(q_units[0]) - .reduce(ReductionOp.MAX) - .to(fx.Float32) - ) + absmax = fmath.absf(q_units[0]).reduce(ReductionOp.MAX).to(fx.Float32) for u in range_constexpr(1, N_QLOADS): absmax = fx.maxnumf( absmax, - fmath.absf(q_units[u]) - .reduce(ReductionOp.MAX) - .to(fx.Float32), + fmath.absf(q_units[u]).reduce(ReductionOp.MAX).to(fx.Float32), ) for sh in (8, 4, 2, 1): absmax = fx.maxnumf(absmax, dpp_utils.dpp_xor_f32(absmax, sh)) @@ -869,9 +863,7 @@ def _l_slot(m): tok0 = tt * TILE_TOK # per_tensor phase-split: let IGLP interleave MFMA with softmax # VALU/LDS to hide the MFMA-hazard s_nop (per_token is VGPR-cliffed). - if const_expr( - (not per_token_kv and M_TILES > 1) or PAGE16_VPIPE_IGLP - ): + if const_expr((not per_token_kv and M_TILES > 1) or PAGE16_VPIPE_IGLP): fx.rocdl.iglp_opt(0) tt1 = tt + 1 diff --git a/aiter/ops/flydsl/pa_decode.py b/aiter/ops/flydsl/pa_decode.py index f9287ba5107..e7d2ea981df 100644 --- a/aiter/ops/flydsl/pa_decode.py +++ b/aiter/ops/flydsl/pa_decode.py @@ -45,6 +45,7 @@ _PAGE16_VPIPE_ENV = "AITER_FLYDSL_PA_PAGE16_VPIPE" _PAGE16_VPIPE_IGLP_ENV = "AITER_FLYDSL_PA_PAGE16_VPIPE_IGLP" _MATCH_GLUON_NUMERICS_ENV = "AITER_FLYDSL_PA_MATCH_GLUON_NUMERICS" +_PAGE128_EARLY_V_ENV = "AITER_FLYDSL_PA_PAGE128_EARLY_V" def _env_enabled(name: str) -> bool: @@ -426,7 +427,7 @@ def normalize_scale(scale, name): _env_enabled(_MATCH_GLUON_NUMERICS_ENV) and arch == "gfx950" and head_dim == 128 - and block_size == 16 + and block_size in (16, 128) and trans_v and not per_token_kv and query.dtype == torch.bfloat16 @@ -436,8 +437,8 @@ def normalize_scale(scale, name): ) # Early V loads help the measured one-to-two-workgroup-per-CU regime. - # Keep other grids on the existing schedule: early loads regress - # short-context decode at larger grid sizes. + # Keep other grids on the existing schedule by default; the opt-in override + # lets long-context under-filled grids be measured without changing policy. prefetch_v = page16_vpipe if ( arch == "gfx950" @@ -449,7 +450,9 @@ def normalize_scale(scale, name): ): num_cus = torch.cuda.get_device_properties(dev).multi_processor_count workgroups = num_seqs * num_kv_heads * num_partitions - prefetch_v = num_cus < workgroups <= 2 * num_cus + prefetch_v = _env_enabled(_PAGE128_EARLY_V_ENV) or ( + num_cus < workgroups <= 2 * num_cus + ) with torch.cuda.device(dev): compiled = compile_pa_decode_tile( From c154696e0730bae0172a3865381c88db99a9073d Mon Sep 17 00:00:00 2001 From: root Date: Mon, 14 Sep 2026 06:43:34 +0000 Subject: [PATCH 30/66] [FlyDSL] fix(pa_decode): allow a packed V cache to span fewer blocks f4f84314a added a K/V block-count equality check, which rejects the packed cache layout: V is reached through a view shifted past the first interleaved pages, so it legitimately spans fewer blocks than K. That broke test_prepared_sparse_page_tables[16-False-True] with "key_cache and value_cache must have the same number of blocks, got 25008 and 25000". Block ids live on the device, so V's usable extent cannot be validated on the host without a synchronize; it stays a caller contract, now stated in the docstring as [0, min(key_cache.shape[0], value_cache.shape[0])). Only the inverted-argument direction is still rejected, which keeps the existing kv-block-count negative test passing. Co-authored-by: Cursor --- aiter/ops/flydsl/pa_decode.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/aiter/ops/flydsl/pa_decode.py b/aiter/ops/flydsl/pa_decode.py index 6341da7bfec..30179db2ee9 100644 --- a/aiter/ops/flydsl/pa_decode.py +++ b/aiter/ops/flydsl/pa_decode.py @@ -196,7 +196,8 @@ def pa_decode( are not inspected here (which would synchronize the device). Callers must ensure ``0 <= context_lengths[i] <= block_tables.shape[1] * block_size`` and that every block-table entry used by a sequence is a physical block index in - ``[0, key_cache.shape[0])``. + ``[0, min(key_cache.shape[0], value_cache.shape[0]))`` -- a packed cache + reaches V through a shifted view that spans fewer blocks than K. """ if context_partition_size != KV_COMPUTE_BLOCK: raise NotImplementedError( @@ -326,9 +327,13 @@ def pa_decode( f"got {tuple(value_cache.shape)} for block_size={block_size}, " f"head_dim={head_dim}" ) - if v_num_blocks != num_blocks: + # A packed cache interleaves K and V in one allocation, so the V view starts + # part-way into it and legitimately spans fewer blocks than K. Block ids live + # on the device, so V's usable extent is a caller contract either way; only + # the inverted-argument direction is worth rejecting here. + if v_num_blocks > num_blocks: raise ValueError( - f"key_cache and value_cache must have the same number of blocks, " + 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: From 71241bb16fcc9472da3c7252cd253013c11258a7 Mon Sep 17 00:00:00 2001 From: root Date: Mon, 14 Sep 2026 06:43:47 +0000 Subject: [PATCH 31/66] [FlyDSL] test(pa_decode): sweep served TP4 batches over both page sizes A tp4 shard of a 64-Q-head / 4-KV-head / D128 model serves 16 Q heads over 1 KV head. Nothing covered that grid across batch, so add test_tp4_decode_batches: batch 1-64 x KV 2048/100k/200k x page 16/128, and give test_tp4_contexts the same page-size axis. At NP=8 on a 256-CU part the batch axis straddles the prefetch_v window (256 < batch * 8 <= 512), so page-128 batch 64 compiles the tuned specialization while smaller batches compile the default schedule; page 16 never enters it because the gate also requires block_size == 128. Confirmed by instrumenting compile_pa_decode_tile. _run_accuracy_case and run_pa_decode_tile_case grow a trans_v option, and the perf sweep grows a --trans-v axis: production serves the transposed 5-D V cache, but the harness only ever measured the plain 4-D one. run_perftest now receives the tensors as arguments so it rotates input buffers instead of replaying one hot cache. Co-authored-by: Cursor --- op_tests/test_flydsl_pa_decode.py | 115 +++++++++++++++++++++++------- 1 file changed, 89 insertions(+), 26 deletions(-) diff --git a/op_tests/test_flydsl_pa_decode.py b/op_tests/test_flydsl_pa_decode.py index e9f3c2133c3..4b3ff4763aa 100644 --- a/op_tests/test_flydsl_pa_decode.py +++ b/op_tests/test_flydsl_pa_decode.py @@ -429,6 +429,7 @@ def run_pa_decode_tile_case( context_length, block_size, dtype, + trans_v, ): if pa_decode is None or get_recommended_splits is None: raise RuntimeError("FlyDSL is not available") @@ -476,7 +477,16 @@ def run_pa_decode_tile_case( .permute(0, 1, 3, 2, 4) .contiguous() ) - value_cache = value_quant.contiguous() + if trans_v: + value_cache = ( + value_quant.permute(0, 1, 3, 2) + .contiguous() + .view(num_blocks, num_kv_heads, block_size // 16, 16, head_dim) + .permute(0, 1, 2, 4, 3) + .contiguous() + ) + else: + value_cache = value_quant.contiguous() block_tables = torch.arange(num_blocks, dtype=torch.int32).reshape( batch_size, blocks_per_sequence ) @@ -510,27 +520,12 @@ def run_pa_decode_tile_case( pout = torch.empty(*partial_shape, head_dim, dtype=dtype) softmax_scale = head_dim**-0.5 - candidates = { - "flydsl": lambda: _run_flydsl( - output, - query, - key_cache, - value_cache, - block_tables, - context_lengths, - key_scale, - value_scale, - num_partitions, - softmax_scale, - pmax, - psum, - pout, - ) - } + candidates = {"flydsl": _run_flydsl} # QK and PV each perform one multiply-add per query-head/context pair. flops = 4 * batch_size * num_query_heads * context_length * head_dim - # Logical tensor traffic: Q + O + the referenced K/V tokens and metadata. + # Effective bandwidth counts Q + O + referenced K/V tokens and metadata once. + # It excludes padded tokens, repeated loads, and partition scratch traffic. nbytes = ( 2 * query.numel() * query.element_size() + 2 @@ -547,7 +542,23 @@ def run_pa_decode_tile_case( ret = {"gfx": get_gfx_runtime(), "partitions": num_partitions} for name, fn in candidates.items(): - out, us = run_perftest(fn) + # Pass tensors explicitly so perftest can rotate their allocations. + out, us = run_perftest( + fn, + output, + query, + key_cache, + value_cache, + block_tables, + context_lengths, + key_scale, + value_scale, + num_partitions, + softmax_scale, + pmax, + psum, + pout, + ) err = checkAllclose( reference.to(dtypes.fp32), out.to(dtypes.fp32), @@ -573,6 +584,7 @@ def _run_accuracy_case( query_dtype, num_partitions, tolerance, + trans_v=False, ): """Run PA against a dequantized-FP8 torch reference without benchmarking.""" _require_gpu() @@ -605,7 +617,16 @@ def _run_accuracy_case( .permute(0, 1, 3, 2, 4) .contiguous() ) - value_cache = value_quant.contiguous() + if trans_v: + value_cache = ( + value_quant.permute(0, 1, 3, 2) + .contiguous() + .view(num_blocks, num_kv_heads, block_size // 16, 16, head_dim) + .permute(0, 1, 2, 4, 3) + .contiguous() + ) + else: + value_cache = value_quant.contiguous() block_tables = torch.zeros( batch_size, max_pages, dtype=dtypes.i32, device=query.device @@ -1214,14 +1235,20 @@ def test_page128_chunk_boundaries(context_length, trans_v, per_token): _assert_matches(output, reference) +@pytest.mark.parametrize("block_size", [16, 128]) @pytest.mark.parametrize("context_length", [2048, 100000, 200000]) @pytest.mark.parametrize("query_length", [1, 3]) @pytest.mark.parametrize("per_token", [False, True]) -def test_tp4_page128_contexts(context_length, query_length, per_token): - """64 Q / 4 KV heads under TP4: Hq16, Hkv1, D128, including dense MTP.""" +def test_tp4_contexts(block_size, context_length, query_length, per_token): + """64 Q / 4 KV heads under TP4: Hq16, Hkv1, D128, including dense MTP. + + Page size decides how many block-table entries a 256-token tile walks -- + two at page-128, sixteen at page-16 -- so both served page sizes are run + against every context. + """ _require_gpu() output, reference = _adversarial_case( - block_size=128, + block_size=block_size, query_group_size=16, context_length=context_length, query_length=query_length, @@ -1232,6 +1259,32 @@ def test_tp4_page128_contexts(context_length, query_length, per_token): _assert_matches(output, reference) +@pytest.mark.parametrize("batch_size", [1, 2, 4, 8, 16, 32, 64]) +@pytest.mark.parametrize("block_size", [16, 128]) +@pytest.mark.parametrize("context_length", [2048, 100000, 200000]) +def test_tp4_decode_batches(batch_size, block_size, context_length): + """Served TP4 decode grid: Hq16, Hkv1, D128, transposed V. + + Batch is the only axis that moves the CTA count, and at NP=8 the sweep + straddles the tuned V-prefetch window (256 < batch * 8 <= 512 on a 256-CU + part), so page-128 batch 64 compiles the prefetch specialization while the + smaller batches compile the default schedule. Page 16 never enters that + window -- the gate also requires block_size == 128 -- so it covers the + sixteen-pages-per-tile block-table walk over the same grids instead. + """ + _run_accuracy_case( + [context_length] * batch_size, + num_query_heads=16, + num_kv_heads=1, + head_dim=128, + block_size=block_size, + query_dtype=dtypes.bf16, + num_partitions=8, + tolerance=FIXED_LENGTH_ACCURACY_TOLERANCE, + trans_v=True, + ) + + @pytest.mark.parametrize("trans_v", [False, True]) def test_page128_cache_above_2gib(trans_v): """Large batches at 200k context cross the signed-i32 cache offset limit.""" @@ -1481,11 +1534,20 @@ def main(): default=[16, 64, 128], help="""KV-cache block sizes.""", ) + parser.add_argument( + "--trans-v", + type=int, + nargs="*", + choices=[0, 1], + default=[0, 1], + help="""V-cache layouts: 0 is the plain 4-D cache, 1 the transposed 5-D + cache production serves.""", + ) args = parser.parse_args() rows = [] - for dtype, batch_size, shape, block_size in itertools.product( - args.dtype, args.batch, args.shapes, args.block_size + for dtype, batch_size, shape, block_size, trans_v in itertools.product( + args.dtype, args.batch, args.shapes, args.block_size, args.trans_v ): num_query_heads, num_kv_heads, head_dim, context_length = shape rows.append( @@ -1497,6 +1559,7 @@ def main(): context_length, block_size, dtype, + bool(trans_v), ) ) From 959a148f6cb19c9583d4e008f9fa98c3a9f03ba3 Mon Sep 17 00:00:00 2001 From: fsx950223 Date: Mon, 14 Sep 2026 06:46:25 +0000 Subject: [PATCH 32/66] perf(flydsl): parallelize large PA reductions --- aiter/ops/flydsl/kernels/pa_decode_reduce.py | 202 +++++++++++++++++-- 1 file changed, 186 insertions(+), 16 deletions(-) diff --git a/aiter/ops/flydsl/kernels/pa_decode_reduce.py b/aiter/ops/flydsl/kernels/pa_decode_reduce.py index 1fd7a2f091f..d0d3d1b9ffd 100644 --- a/aiter/ops/flydsl/kernels/pa_decode_reduce.py +++ b/aiter/ops/flydsl/kernels/pa_decode_reduce.py @@ -24,11 +24,10 @@ def compile_pa_decode_ps_reduce( ): """Build the partitioned-softmax reduction used by ``pa_decode``. - Each wave reduces the partition statistics independently, then reuses the - per-partition weights to combine one output element per thread. Counts up - to one wave retain the original one-partition-per-lane path. For larger - counts, each lane owns up to four partitions separated by a wave, keeping - all weights in registers and avoiding an LDS round trip. + Counts up to one wave use one partition per lane and one output element per + thread. For D=128 and larger counts, a 2-D workgroup materializes weights + once in LDS and splits each output element's partition chain over several + waves. Other head sizes retain the register-only lane-striped fallback. """ if not 1 <= max_context_partition_num <= MAX_CONTEXT_PARTITIONS: raise ValueError( @@ -63,7 +62,39 @@ def compile_pa_decode_ps_reduce( offset for offset in (32, 16, 8, 4, 2, 1) if offset < reduce_width ] - @flyc.kernel(known_block_size=[head_size, 1, 1]) + # The original mapping gives each output element to one thread, so every + # thread walks every partition. That is a good fit for <=1 wave of + # partitions, but it leaves NP=160..256 as a long dependent load/FMA + # chain. For the decode shape used by PA (D=128), split that chain over + # two or four independent wave pairs. A pair covers the two 64-element + # halves of the output vector, while its y-coordinate selects a disjoint + # contiguous range of partitions. + use_parallel_lds = head_size == 128 and max_context_partition_num > warp_size + parallel_groups = 1 + if use_parallel_lds: + # Four groups win from NP=128 onward on gfx950; two avoid excessive + # synchronization/thread overhead for the small >64 tail. + parallel_groups = 2 if max_context_partition_num <= 96 else 4 + 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, 1, 1] + ) + parts_per_group = ( + max_context_partition_num + parallel_groups - 1 + ) // parallel_groups + + # Keep the legacy specializations effectively LDS-free. The fields are + # only allocated from the compile-time parallel branch below. + 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, @@ -84,6 +115,7 @@ def pa_decode_ps_reduce_kernel( query_group_size: fx.Int32, ): 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 @@ -128,7 +160,134 @@ def _wave_reduce_sum(value): ) return reduced - if fx.const_expr(max_context_partition_num <= warp_size): + if fx.const_expr(use_parallel_lds): + # One wave materializes the normalized partition weights once. + # All output waves then reuse those weights from LDS and split the + # long partition loop. This avoids both duplicated exp2 work and + # a ds_bpermute for every output FMA. + 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): + stats_offset = ( + batch_idx * stride_exp_sums_seq + + 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) + stats_offset = ( + batch_idx * stride_exp_sums_seq + + 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) + 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] + 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_value = fx.Float32(sink_token[kv_head_idx * c_qgs + group_idx]) + sink_scale = (global_max > neg_inf).select( + fx.exp2( + (sink_value - safe_global_max) * c_log2e, + fastmath="fast", + ), + zero_f, + ) + global_exp_sum = global_exp_sum + sink_scale + safe_global_exp_sum = (global_exp_sum > zero_f).select( + global_exp_sum, one_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) + part_idx = lane + fx.Int32(chunk_base) + weight = scaled_sums[chunk_idx] / safe_global_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 + 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 = ( + batch_idx * stride_logits_seq + + 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: + 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): + 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): # Exact powers of two have no inactive lanes inside their reduction # subgroup. Keep that original path unchanged; only partial # subgroups need an EXEC-masked load to avoid carrying a predicate @@ -311,14 +470,25 @@ def _wave_reduce_sum(value): acc = acc + part_logits * weight query_idx = eqgs_idx // c_qgs - 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 - ) - output[output_offset] = acc.to(output_dtype) + 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 + ) + output[output_offset] = acc.to(output_dtype) @flyc.jit def launch_pa_decode_ps_reduce_kernel( @@ -364,7 +534,7 @@ def launch_pa_decode_ps_reduce_kernel( query_group_size, ).launch( grid=(batch_size, num_kv_heads, query_seq_len * query_group_size), - block=(head_size, 1, 1), + block=tuple(block_shape), stream=stream, ) From 610db54ba520c1e9c89ab4b58d87bf48b4f295ee Mon Sep 17 00:00:00 2001 From: fsx950223 Date: Mon, 14 Sep 2026 06:49:36 +0000 Subject: [PATCH 33/66] test(flydsl): model parallel PA reduction --- op_tests/test_flydsl_pa_decode_reduce_cpu.py | 71 +++++++++++++++++++- 1 file changed, 70 insertions(+), 1 deletion(-) diff --git a/op_tests/test_flydsl_pa_decode_reduce_cpu.py b/op_tests/test_flydsl_pa_decode_reduce_cpu.py index dbc52cc9bc2..28fa20bc687 100644 --- a/op_tests/test_flydsl_pa_decode_reduce_cpu.py +++ b/op_tests/test_flydsl_pa_decode_reduce_cpu.py @@ -62,13 +62,15 @@ def test_reduce_kernel_rejects_unsupported_partition_counts(num_partitions): ) -def _flat_partition_reduce(exp_sums, max_logits, partials): +def _flat_partition_reduce(exp_sums, max_logits, partials, sink=None): global_max = max(max_logits) scaled_sums = [ part_sum * math.exp(part_max - global_max) if part_max != -math.inf else 0.0 for part_sum, part_max in zip(exp_sums, max_logits) ] denominator = sum(scaled_sums) + if sink is not None and global_max != -math.inf: + denominator += math.exp(sink - global_max) if denominator == 0.0: return [0.0] * len(partials[0]) return [ @@ -111,6 +113,42 @@ def _lane_striped_partition_reduce(exp_sums, max_logits, partials): return accumulators +def _parallel_lds_partition_reduce(exp_sums, max_logits, partials, sink=None): + """Model the D=128 large-NP workgroup and its cross-wave LDS merge.""" + num_partitions = len(exp_sums) + parallel_groups = 2 if num_partitions <= 96 else 4 + parts_per_group = (num_partitions + parallel_groups - 1) // parallel_groups + global_max = max(max_logits) + scaled_sums = [ + part_sum * math.exp(part_max - global_max) if part_max != -math.inf else 0.0 + for part_sum, part_max in zip(exp_sums, max_logits) + ] + denominator = sum(scaled_sums) + if sink is not None and global_max != -math.inf: + denominator += math.exp(sink - global_max) + if denominator == 0.0: + return [0.0] * len(partials[0]) + weights = [value / denominator for value in scaled_sums] + + group_accumulators = [] + for group in range(parallel_groups): + begin = group * parts_per_group + end = min(begin + parts_per_group, num_partitions) + group_accumulators.append( + [ + sum( + partials[part][element] * weights[part] + for part in range(begin, end) + ) + for element in range(len(partials[0])) + ] + ) + return [ + sum(group[element] for group in group_accumulators) + for element in range(len(partials[0])) + ] + + @pytest.mark.parametrize( "num_partitions", [26, 30, 32, 34, 36, 40, 64, 65, 96, 127, 128, 129, 160, 192, 255, 256], @@ -140,3 +178,34 @@ def test_lane_striped_reduce_all_empty_is_zero(num_partitions): # The tile kernel writes finite neutral scratch for empty partitions. partials = [[0.0] * 7 for _ in range(num_partitions)] assert _lane_striped_partition_reduce(exp_sums, max_logits, partials) == [0.0] * 7 + + +@pytest.mark.parametrize("num_partitions", [65, 96, 97, 128, 160, 192, 255, 256]) +@pytest.mark.parametrize("use_sink", [False, True]) +def test_parallel_lds_reduce_matches_flat_reference(num_partitions, use_sink): + rng = random.Random(10_000 + num_partitions) + max_logits = [rng.uniform(-20.0, 5.0) for _ in range(num_partitions)] + exp_sums = [rng.uniform(0.01, 10.0) for _ in range(num_partitions)] + partials = [ + [rng.uniform(-2.0, 2.0) for _ in range(128)] for _ in range(num_partitions) + ] + for part in range(7, num_partitions, 37): + max_logits[part] = -math.inf + exp_sums[part] = 0.0 + sink = rng.uniform(-10.0, 4.0) if use_sink else None + + expected = _flat_partition_reduce(exp_sums, max_logits, partials, sink) + actual = _parallel_lds_partition_reduce(exp_sums, max_logits, partials, sink) + assert actual == pytest.approx(expected, rel=2e-15, abs=2e-15) + + +@pytest.mark.parametrize("num_partitions", [65, 96, 128, 160, 192, 256]) +@pytest.mark.parametrize("sink", [None, -2.0]) +def test_parallel_lds_reduce_all_empty_is_zero(num_partitions, sink): + exp_sums = [0.0] * num_partitions + max_logits = [-math.inf] * num_partitions + partials = [[0.0] * 128 for _ in range(num_partitions)] + assert ( + _parallel_lds_partition_reduce(exp_sums, max_logits, partials, sink) + == [0.0] * 128 + ) From 53cac035897019c163c81bf76a681e9cdaab6a4a Mon Sep 17 00:00:00 2001 From: fsx950223 Date: Mon, 14 Sep 2026 06:54:44 +0000 Subject: [PATCH 34/66] perf(flydsl): use eight PA reduction groups --- aiter/ops/flydsl/kernels/pa_decode_reduce.py | 6 +++--- op_tests/test_flydsl_pa_decode_reduce_cpu.py | 2 +- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/aiter/ops/flydsl/kernels/pa_decode_reduce.py b/aiter/ops/flydsl/kernels/pa_decode_reduce.py index d0d3d1b9ffd..5d8c26274b0 100644 --- a/aiter/ops/flydsl/kernels/pa_decode_reduce.py +++ b/aiter/ops/flydsl/kernels/pa_decode_reduce.py @@ -66,15 +66,15 @@ def compile_pa_decode_ps_reduce( # thread walks every partition. That is a good fit for <=1 wave of # partitions, but it leaves NP=160..256 as a long dependent load/FMA # chain. For the decode shape used by PA (D=128), split that chain over - # two or four independent wave pairs. A pair covers the two 64-element + # two or eight independent wave pairs. A pair covers the two 64-element # halves of the output vector, while its y-coordinate selects a disjoint # contiguous range of partitions. use_parallel_lds = head_size == 128 and max_context_partition_num > warp_size parallel_groups = 1 if use_parallel_lds: - # Four groups win from NP=128 onward on gfx950; two avoid excessive + # Eight groups win from NP=128 onward on gfx950; two avoid excessive # synchronization/thread overhead for the small >64 tail. - parallel_groups = 2 if max_context_partition_num <= 96 else 4 + 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 = ( diff --git a/op_tests/test_flydsl_pa_decode_reduce_cpu.py b/op_tests/test_flydsl_pa_decode_reduce_cpu.py index 28fa20bc687..048996703c2 100644 --- a/op_tests/test_flydsl_pa_decode_reduce_cpu.py +++ b/op_tests/test_flydsl_pa_decode_reduce_cpu.py @@ -116,7 +116,7 @@ def _lane_striped_partition_reduce(exp_sums, max_logits, partials): def _parallel_lds_partition_reduce(exp_sums, max_logits, partials, sink=None): """Model the D=128 large-NP workgroup and its cross-wave LDS merge.""" num_partitions = len(exp_sums) - parallel_groups = 2 if num_partitions <= 96 else 4 + parallel_groups = 2 if num_partitions <= 96 else 8 parts_per_group = (num_partitions + parallel_groups - 1) // parallel_groups global_max = max(max_logits) scaled_sums = [ From 5ae39e4cd96576e4b99a53e2f46bea635f20466f Mon Sep 17 00:00:00 2001 From: fsx950223 Date: Mon, 14 Sep 2026 07:12:56 +0000 Subject: [PATCH 35/66] test(flydsl): expose transposed-V PA sweep option --- op_tests/test_flydsl_pa_decode.py | 35 +++++++++++++++++++++++++++---- 1 file changed, 31 insertions(+), 4 deletions(-) diff --git a/op_tests/test_flydsl_pa_decode.py b/op_tests/test_flydsl_pa_decode.py index 652a4b5b849..2daadaf3cfb 100644 --- a/op_tests/test_flydsl_pa_decode.py +++ b/op_tests/test_flydsl_pa_decode.py @@ -296,6 +296,7 @@ def run_pa_decode_tile_case( context_length, block_size, dtype, + trans_v=False, ): if pa_decode is None or get_recommended_splits is None: raise RuntimeError("FlyDSL is not available") @@ -343,7 +344,20 @@ def run_pa_decode_tile_case( .permute(0, 1, 3, 2, 4) .contiguous() ) - value_cache = value_quant.contiguous() + if trans_v: + value_cache = ( + value_quant.view( + num_blocks, + num_kv_heads, + head_dim, + block_size // 16, + 16, + ) + .permute(0, 1, 3, 2, 4) + .contiguous() + ) + else: + value_cache = value_quant.contiguous() block_tables = torch.arange(num_blocks, dtype=torch.int32).reshape( batch_size, blocks_per_sequence ) @@ -412,7 +426,11 @@ def run_pa_decode_tile_case( + value_scale.numel() * value_scale.element_size() ) - ret = {"gfx": get_gfx_runtime(), "partitions": num_partitions} + ret = { + "gfx": get_gfx_runtime(), + "partitions": num_partitions, + "trans_v": trans_v, + } for name, fn in candidates.items(): out, us = run_perftest(fn) err = checkAllclose( @@ -1130,11 +1148,19 @@ def main(): default=[16, 64, 128], help="""KV-cache block sizes.""", ) + parser.add_argument( + "--trans-v", + type=int, + nargs="*", + choices=[0, 1], + default=[0], + help="""Value-cache layout: 0 for rank-4, 1 for rank-5 transposed V.""", + ) args = parser.parse_args() rows = [] - for dtype, batch_size, shape, block_size in itertools.product( - args.dtype, args.batch, args.shapes, args.block_size + for dtype, batch_size, shape, block_size, trans_v in itertools.product( + args.dtype, args.batch, args.shapes, args.block_size, args.trans_v ): num_query_heads, num_kv_heads, head_dim, context_length = shape rows.append( @@ -1146,6 +1172,7 @@ def main(): context_length, block_size, dtype, + bool(trans_v), ) ) From d5edb5c47b6f9ac6a1a79d636d65ecea07b7c4af Mon Sep 17 00:00:00 2001 From: fsx950223 Date: Mon, 14 Sep 2026 07:46:44 +0000 Subject: [PATCH 36/66] feat(flydsl): make PA split clamp configurable --- aiter/ops/flydsl/pa_decode.py | 14 +++++++++- op_tests/test_flydsl_pa_decode.py | 43 +++++++++++++++++++++++++++++++ 2 files changed, 56 insertions(+), 1 deletion(-) diff --git a/aiter/ops/flydsl/pa_decode.py b/aiter/ops/flydsl/pa_decode.py index cec5953b26c..d8045d01249 100644 --- a/aiter/ops/flydsl/pa_decode.py +++ b/aiter/ops/flydsl/pa_decode.py @@ -62,12 +62,24 @@ def get_recommended_splits( num_sequences: int, num_kv_heads: int, split_kv_blocks: int = 1, + max_partitions: int = 8, ) -> int: + """Choose a split count, with a caller-configurable upper clamp. + + The default of eight preserves the existing policy. Long-context callers + can opt into the FlyDSL reducer's larger domain with, for example, + ``max_partitions=256``. + """ + 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 = cdiv(num_sm, denom) * split_kv_blocks - return max(4, min(n, 8)) + return max(4, min(n, max_partitions)) def _workgroup_count_enables_v_prefetch( diff --git a/op_tests/test_flydsl_pa_decode.py b/op_tests/test_flydsl_pa_decode.py index fb217f269d7..0fa8bb57655 100644 --- a/op_tests/test_flydsl_pa_decode.py +++ b/op_tests/test_flydsl_pa_decode.py @@ -342,6 +342,40 @@ def test_v_prefetch_workgroup_interval(batch_size, expected): ) +@pytest.mark.parametrize( + "batch_size,max_partitions,expected", + [(1, 8, 8), (1, 256, 256), (16, 256, 32), (64, 256, 8)], +) +def test_recommended_splits_has_configurable_upper_clamp( + monkeypatch, batch_size, max_partitions, expected +): + """The default stays at eight while long-context callers may opt into 256.""" + if get_recommended_splits is None: + pytest.skip("FlyDSL is not available") + monkeypatch.setattr( + torch.cuda, + "get_device_properties", + lambda *_args, **_kwargs: type("Props", (), {"multi_processor_count": 256})(), + ) + assert ( + get_recommended_splits( + batch_size, + num_kv_heads=1, + split_kv_blocks=2, + max_partitions=max_partitions, + ) + == expected + ) + + +@pytest.mark.parametrize("max_partitions", [0, 3, 257]) +def test_recommended_splits_rejects_invalid_upper_clamp(max_partitions): + if get_recommended_splits is None: + pytest.skip("FlyDSL is not available") + with pytest.raises(ValueError, match="max_partitions"): + get_recommended_splits(1, 1, max_partitions=max_partitions) + + def run_torch( query: torch.Tensor, key_cache: torch.Tensor, @@ -430,6 +464,7 @@ def run_pa_decode_tile_case( block_size, dtype, trans_v, + max_partitions=8, ): if pa_decode is None or get_recommended_splits is None: raise RuntimeError("FlyDSL is not available") @@ -508,6 +543,7 @@ def run_pa_decode_tile_case( batch_size, num_kv_heads, split_kv_blocks=KV_COMPUTE_BLOCK // block_size, + max_partitions=max_partitions, ) partial_shape = ( batch_size, @@ -1547,6 +1583,12 @@ def main(): help="""V-cache layouts: 0 is the plain 4-D cache, 1 the transposed 5-D cache production serves.""", ) + parser.add_argument( + "--max-partitions", + type=int, + default=8, + help="""Upper clamp passed to get_recommended_splits (4..256).""", + ) args = parser.parse_args() rows = [] @@ -1564,6 +1606,7 @@ def main(): block_size, dtype, bool(trans_v), + args.max_partitions, ) ) From 38291a412503db70fc24b44f1f8f2a0b467e9d15 Mon Sep 17 00:00:00 2001 From: fsx950223 Date: Mon, 14 Sep 2026 08:13:19 +0000 Subject: [PATCH 37/66] test(flydsl): remove standalone PA reducer model test --- op_tests/test_flydsl_pa_decode_reduce_cpu.py | 211 ------------------- 1 file changed, 211 deletions(-) delete mode 100644 op_tests/test_flydsl_pa_decode_reduce_cpu.py diff --git a/op_tests/test_flydsl_pa_decode_reduce_cpu.py b/op_tests/test_flydsl_pa_decode_reduce_cpu.py deleted file mode 100644 index 048996703c2..00000000000 --- a/op_tests/test_flydsl_pa_decode_reduce_cpu.py +++ /dev/null @@ -1,211 +0,0 @@ -# SPDX-License-Identifier: MIT -# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. - -"""CPU-only coverage for the FlyDSL PA partition reducer.""" - -import importlib.util -import math -import random -import sys -from pathlib import Path - -import pytest - -pytest.importorskip("flydsl") - - -def _load_reduce_module(): - """Load the self-contained kernel module without importing GPU-facing aiter.""" - path = ( - Path(__file__).parents[1] - / "aiter" - / "ops" - / "flydsl" - / "kernels" - / "pa_decode_reduce.py" - ) - name = "_test_flydsl_pa_decode_reduce" - spec = importlib.util.spec_from_file_location(name, path) - module = importlib.util.module_from_spec(spec) - sys.modules[name] = module - spec.loader.exec_module(module) - return module - - -REDUCE = _load_reduce_module() - - -@pytest.mark.parametrize("num_partitions", [1, 63, 64, 65, 128, 192, 256]) -def test_reduce_kernel_accepts_supported_partition_counts(num_partitions): - """Kernel construction is CPU-only and covers both generated code paths.""" - compiled = REDUCE.compile_pa_decode_ps_reduce( - max_context_partition_num=num_partitions, - head_size=128, - output_dtype_str="bf16", - logits_dtype_str="bf16", - sink_dtype_str="bf16", - use_sinks=False, - ) - assert set(compiled) == {"launch", "kernel"} - - -@pytest.mark.parametrize("num_partitions", [0, 257]) -def test_reduce_kernel_rejects_unsupported_partition_counts(num_partitions): - with pytest.raises(ValueError, match=r"must be in \[1, 256\]"): - REDUCE.compile_pa_decode_ps_reduce( - max_context_partition_num=num_partitions, - head_size=128, - output_dtype_str="bf16", - logits_dtype_str="bf16", - sink_dtype_str="bf16", - use_sinks=False, - ) - - -def _flat_partition_reduce(exp_sums, max_logits, partials, sink=None): - global_max = max(max_logits) - scaled_sums = [ - part_sum * math.exp(part_max - global_max) if part_max != -math.inf else 0.0 - for part_sum, part_max in zip(exp_sums, max_logits) - ] - denominator = sum(scaled_sums) - if sink is not None and global_max != -math.inf: - denominator += math.exp(sink - global_max) - if denominator == 0.0: - return [0.0] * len(partials[0]) - return [ - sum( - partials[part][element] * scaled_sums[part] for part in range(len(exp_sums)) - ) - / denominator - for element in range(len(partials[0])) - ] - - -def _lane_striped_partition_reduce(exp_sums, max_logits, partials): - """Model the >64 kernel: lane owns p, p+64, p+128, and p+192.""" - wave_size = 64 - lane_maxes = [-math.inf] * wave_size - for part, part_max in enumerate(max_logits): - lane = part % wave_size - lane_maxes[lane] = max(lane_maxes[lane], part_max) - global_max = max(lane_maxes) - - scaled_sums = [] - lane_sums = [0.0] * wave_size - for part, (part_sum, part_max) in enumerate(zip(exp_sums, max_logits)): - scaled_sum = ( - part_sum * math.exp(part_max - global_max) if part_max != -math.inf else 0.0 - ) - scaled_sums.append(scaled_sum) - lane_sums[part % wave_size] += scaled_sum - - denominator = sum(lane_sums) - if denominator == 0.0: - return [0.0] * len(partials[0]) - accumulators = [0.0] * len(partials[0]) - for chunk_base in range(0, len(exp_sums), wave_size): - chunk_end = min(chunk_base + wave_size, len(exp_sums)) - for part in range(chunk_base, chunk_end): - weight = scaled_sums[part] / denominator - for element, value in enumerate(partials[part]): - accumulators[element] += value * weight - return accumulators - - -def _parallel_lds_partition_reduce(exp_sums, max_logits, partials, sink=None): - """Model the D=128 large-NP workgroup and its cross-wave LDS merge.""" - num_partitions = len(exp_sums) - parallel_groups = 2 if num_partitions <= 96 else 8 - parts_per_group = (num_partitions + parallel_groups - 1) // parallel_groups - global_max = max(max_logits) - scaled_sums = [ - part_sum * math.exp(part_max - global_max) if part_max != -math.inf else 0.0 - for part_sum, part_max in zip(exp_sums, max_logits) - ] - denominator = sum(scaled_sums) - if sink is not None and global_max != -math.inf: - denominator += math.exp(sink - global_max) - if denominator == 0.0: - return [0.0] * len(partials[0]) - weights = [value / denominator for value in scaled_sums] - - group_accumulators = [] - for group in range(parallel_groups): - begin = group * parts_per_group - end = min(begin + parts_per_group, num_partitions) - group_accumulators.append( - [ - sum( - partials[part][element] * weights[part] - for part in range(begin, end) - ) - for element in range(len(partials[0])) - ] - ) - return [ - sum(group[element] for group in group_accumulators) - for element in range(len(partials[0])) - ] - - -@pytest.mark.parametrize( - "num_partitions", - [26, 30, 32, 34, 36, 40, 64, 65, 96, 127, 128, 129, 160, 192, 255, 256], -) -def test_lane_striped_reduce_matches_flat_reference(num_partitions): - rng = random.Random(num_partitions) - max_logits = [rng.uniform(-20.0, 5.0) for _ in range(num_partitions)] - exp_sums = [rng.uniform(0.01, 10.0) for _ in range(num_partitions)] - partials = [ - [rng.uniform(-2.0, 2.0) for _ in range(7)] for _ in range(num_partitions) - ] - - # Exercise masked lanes in the final chunk and neutral empty partitions. - for part in range(11, num_partitions, 37): - max_logits[part] = -math.inf - exp_sums[part] = 0.0 - - expected = _flat_partition_reduce(exp_sums, max_logits, partials) - actual = _lane_striped_partition_reduce(exp_sums, max_logits, partials) - assert actual == pytest.approx(expected, rel=2e-15, abs=2e-15) - - -@pytest.mark.parametrize("num_partitions", [128, 192, 256]) -def test_lane_striped_reduce_all_empty_is_zero(num_partitions): - exp_sums = [0.0] * num_partitions - max_logits = [-math.inf] * num_partitions - # The tile kernel writes finite neutral scratch for empty partitions. - partials = [[0.0] * 7 for _ in range(num_partitions)] - assert _lane_striped_partition_reduce(exp_sums, max_logits, partials) == [0.0] * 7 - - -@pytest.mark.parametrize("num_partitions", [65, 96, 97, 128, 160, 192, 255, 256]) -@pytest.mark.parametrize("use_sink", [False, True]) -def test_parallel_lds_reduce_matches_flat_reference(num_partitions, use_sink): - rng = random.Random(10_000 + num_partitions) - max_logits = [rng.uniform(-20.0, 5.0) for _ in range(num_partitions)] - exp_sums = [rng.uniform(0.01, 10.0) for _ in range(num_partitions)] - partials = [ - [rng.uniform(-2.0, 2.0) for _ in range(128)] for _ in range(num_partitions) - ] - for part in range(7, num_partitions, 37): - max_logits[part] = -math.inf - exp_sums[part] = 0.0 - sink = rng.uniform(-10.0, 4.0) if use_sink else None - - expected = _flat_partition_reduce(exp_sums, max_logits, partials, sink) - actual = _parallel_lds_partition_reduce(exp_sums, max_logits, partials, sink) - assert actual == pytest.approx(expected, rel=2e-15, abs=2e-15) - - -@pytest.mark.parametrize("num_partitions", [65, 96, 128, 160, 192, 256]) -@pytest.mark.parametrize("sink", [None, -2.0]) -def test_parallel_lds_reduce_all_empty_is_zero(num_partitions, sink): - exp_sums = [0.0] * num_partitions - max_logits = [-math.inf] * num_partitions - partials = [[0.0] * 128 for _ in range(num_partitions)] - assert ( - _parallel_lds_partition_reduce(exp_sums, max_logits, partials, sink) - == [0.0] * 128 - ) From 42edab9476f54bc086516b3dd87cd866dd36dc2b Mon Sep 17 00:00:00 2001 From: fsx950223 Date: Mon, 14 Sep 2026 08:23:31 +0000 Subject: [PATCH 38/66] test(flydsl): remove Gluon-specific PA cases --- op_tests/test_flydsl_pa_decode.py | 135 +----------------------------- 1 file changed, 3 insertions(+), 132 deletions(-) diff --git a/op_tests/test_flydsl_pa_decode.py b/op_tests/test_flydsl_pa_decode.py index 0fa8bb57655..4c105568dd3 100644 --- a/op_tests/test_flydsl_pa_decode.py +++ b/op_tests/test_flydsl_pa_decode.py @@ -5,7 +5,6 @@ import argparse import importlib -import inspect import itertools import pandas as pd @@ -16,7 +15,6 @@ from aiter import dtypes, per_tensor_quant from aiter.jit.utils.chip_info import get_gfx_runtime from aiter.ops.attention import pa_decode_flydsl as public_pa_decode -from aiter.ops.triton.gluon.pa_decode_gluon import pa_decode_gluon from aiter.test_common import benchmark, checkAllclose, run_perftest @@ -33,9 +31,8 @@ def _default_cuda_device(): SUPPORTED_GFX = ["gfx942", "gfx950"] KV_COMPUTE_BLOCK = 256 -# Match the corresponding FP8 Gluon accuracy contracts: fixed-length outputs -# use the tighter bound, while independently selected variable-length rows use -# the established varlen bound. +# Fixed-length outputs use the tighter bound, while independently selected +# variable-length rows use the established varlen bound. FIXED_LENGTH_ACCURACY_TOLERANCE = 5e-3 VARIABLE_LENGTH_ACCURACY_TOLERANCE = 5e-2 @@ -77,37 +74,13 @@ def _require_gpu(): pytest.skip(f"pa_decode is unsupported on {get_gfx_runtime()}") -def test_pa_decode_api_matches_gluon(): - if pa_decode is None: - pytest.skip("FlyDSL is not available") - - flydsl_parameters = inspect.signature(pa_decode).parameters - gluon_parameters = inspect.signature(pa_decode_gluon).parameters - public_parameters = inspect.signature(public_pa_decode).parameters - - assert tuple(flydsl_parameters) == tuple(gluon_parameters) - assert tuple(public_parameters) == tuple(gluon_parameters) - for name, parameter in flydsl_parameters.items(): - gluon_parameter = gluon_parameters[name] - assert parameter.kind == gluon_parameter.kind - assert parameter.default == gluon_parameter.default - assert public_parameters[name].kind == gluon_parameter.kind - assert public_parameters[name].default == gluon_parameter.default - - -def test_pa_decode_maps_gluon_buffers_and_scale_layout(monkeypatch): +def test_pa_decode_maps_buffers_and_scale_layout(monkeypatch): _require_gpu() pa_decode_module = importlib.import_module("aiter.ops.flydsl.pa_decode") - gluon_module = importlib.import_module("aiter.ops.triton.gluon.pa_decode_gluon") attention_module = importlib.import_module("aiter.ops.attention") captured = {} - assert ( - gluon_module.launch_pa_decode_ps_reduce_flydsl - is pa_decode_module.launch_pa_decode_ps_reduce - ) - monkeypatch.setattr( pa_decode_module, "compile_pa_decode_tile", @@ -220,108 +193,6 @@ def capture_reduce( assert dispatches == ["flydsl"] -def _call_gluon_reduce_wrapper(gluon_module, head_size): - output = torch.empty(1, 1, 1, 1, head_size, dtype=dtypes.bf16, device="cpu") - exp_sums = torch.empty(1, 1, 4, 1, dtype=dtypes.fp32, device="cpu") - max_logits = torch.empty_like(exp_sums) - logits = torch.empty(1, 1, 4, 1, head_size, dtype=dtypes.bf16, device="cpu") - context_lengths = torch.tensor([256], dtype=dtypes.i32, device="cpu") - - gluon_module._paged_attention_decode_v2_reduce_kernel_wrapper( - (1, 1, 1), - output, - exp_sums, - max_logits, - logits, - context_lengths, - None, - output.stride(0), - output.stride(1), - output.stride(2), - output.stride(3), - exp_sums.stride(0), - exp_sums.stride(1), - exp_sums.stride(2), - logits.stride(0), - logits.stride(1), - logits.stride(2), - logits.stride(3), - 1, - 1, - head_size, - 256, - PS=True, - context_partition_num=4, - ) - - -def test_gluon_unsupported_flydsl_reducer_falls_back(monkeypatch): - """An unsupported optional reducer must fall back instead of failing dispatch.""" - _require_gpu() - gluon_module = importlib.import_module("aiter.ops.triton.gluon.pa_decode_gluon") - triton_launches = [] - flydsl_launches = [] - - class TritonReduceStub: - def __getitem__(self, grid): - def launch(*args, **kwargs): - triton_launches.append((grid, args, kwargs)) - - return launch - - monkeypatch.setattr(gluon_module, "CXX_PS_REDUCE_AVAILABLE", False) - monkeypatch.setattr(gluon_module, "FLYDSL_PS_REDUCE_AVAILABLE", True) - monkeypatch.setattr( - gluon_module, - "launch_pa_decode_ps_reduce_flydsl", - lambda *args, **kwargs: flydsl_launches.append((args, kwargs)), - ) - monkeypatch.setattr( - gluon_module, - "paged_attention_decode_ps_reduce_kernel", - TritonReduceStub(), - ) - - _call_gluon_reduce_wrapper(gluon_module, head_size=96) - - assert not flydsl_launches - assert len(triton_launches) == 1 - assert triton_launches[0][0] == (1, 1, 1) - - -def test_gluon_supported_flydsl_reducer_does_not_mask_errors(monkeypatch): - """Errors from a selected FlyDSL reducer are not mistaken for unsupported.""" - _require_gpu() - gluon_module = importlib.import_module("aiter.ops.triton.gluon.pa_decode_gluon") - triton_launches = [] - - class TritonReduceStub: - def __getitem__(self, grid): - def launch(*args, **kwargs): - triton_launches.append((grid, args, kwargs)) - - return launch - - def fail_flydsl(*args, **kwargs): - raise ImportError("FlyDSL reducer failed internally") - - monkeypatch.setattr(gluon_module, "CXX_PS_REDUCE_AVAILABLE", False) - monkeypatch.setattr(gluon_module, "FLYDSL_PS_REDUCE_AVAILABLE", True) - monkeypatch.setattr(gluon_module, "launch_pa_decode_ps_reduce_flydsl", fail_flydsl) - monkeypatch.setattr( - gluon_module, - "paged_attention_decode_ps_reduce_kernel", - TritonReduceStub(), - ) - monkeypatch.setattr( - gluon_module.torch.cuda, "current_stream", lambda device: object() - ) - - with pytest.raises(ImportError, match="failed internally"): - _call_gluon_reduce_wrapper(gluon_module, head_size=128) - assert not triton_launches - - @pytest.mark.parametrize( "batch_size,expected", [(32, False), (33, True), (64, True), (65, False)], From 29689c3173d65b36f3d5257208086ccc754d2be6 Mon Sep 17 00:00:00 2001 From: fsx950223 Date: Mon, 14 Sep 2026 08:54:37 +0000 Subject: [PATCH 39/66] test(flydsl): add per-token PA performance sweep --- op_tests/test_flydsl_pa_decode.py | 51 +++++++++++++++++++++++-------- 1 file changed, 39 insertions(+), 12 deletions(-) diff --git a/op_tests/test_flydsl_pa_decode.py b/op_tests/test_flydsl_pa_decode.py index 4c105568dd3..580e9ceebb6 100644 --- a/op_tests/test_flydsl_pa_decode.py +++ b/op_tests/test_flydsl_pa_decode.py @@ -12,7 +12,7 @@ import torch import aiter -from aiter import dtypes, per_tensor_quant +from aiter import dtypes, per_tensor_quant, pertoken_quant from aiter.jit.utils.chip_info import get_gfx_runtime from aiter.ops.attention import pa_decode_flydsl as public_pa_decode from aiter.test_common import benchmark, checkAllclose, run_perftest @@ -270,13 +270,16 @@ def run_torch( token_offsets = token_ids % block_size physical_pages = block_tables[seq_idx, logical_pages].long() - keys = ( - key_cache[physical_pages, :, token_offsets, :].float() * key_scale.float() - ) - values = ( - value_cache[physical_pages, :, :, token_offsets].float() - * value_scale.float() - ) + keys = key_cache[physical_pages, :, token_offsets, :].float() + values = value_cache[physical_pages, :, :, token_offsets].float() + if key_scale.numel() == 1: + keys = keys * key_scale.float() + values = values * value_scale.float() + else: + token_key_scale = key_scale[physical_pages, :, token_offsets, 0].float() + token_value_scale = value_scale[physical_pages, :, token_offsets, 0].float() + keys = keys * token_key_scale.unsqueeze(-1) + values = values * token_value_scale.unsqueeze(-1) keys = keys.repeat_interleave(query_group_size, dim=1) values = values.repeat_interleave(query_group_size, dim=1) scores = ( @@ -336,6 +339,7 @@ def run_pa_decode_tile_case( dtype, trans_v, max_partitions=8, + per_token=False, ): if pa_decode is None or get_recommended_splits is None: raise RuntimeError("FlyDSL is not available") @@ -370,8 +374,16 @@ def run_pa_decode_tile_case( ).uniform_(-0.5, 0.5) quant_dtype = _quant_dtype() - key_quant, key_scale = per_tensor_quant(key, quant_dtype=quant_dtype) - value_quant, value_scale = per_tensor_quant(value, quant_dtype=quant_dtype) + if per_token: + key_quant, key_scale = pertoken_quant(key, quant_dtype=quant_dtype) + value_token_major = value.permute(0, 1, 3, 2).contiguous() + value_token_quant, value_scale = pertoken_quant( + value_token_major, quant_dtype=quant_dtype + ) + value_quant = value_token_quant.permute(0, 1, 3, 2).contiguous() + else: + key_quant, key_scale = per_tensor_quant(key, quant_dtype=quant_dtype) + value_quant, value_scale = per_tensor_quant(value, quant_dtype=quant_dtype) key_cache = ( key_quant.view( num_blocks, @@ -451,6 +463,7 @@ def run_pa_decode_tile_case( "gfx": get_gfx_runtime(), "partitions": num_partitions, "trans_v": trans_v, + "per_token": per_token, } for name, fn in candidates.items(): # Pass tensors explicitly so perftest can rotate their allocations. @@ -1460,11 +1473,24 @@ def main(): default=8, help="""Upper clamp passed to get_recommended_splits (4..256).""", ) + parser.add_argument( + "--per-token", + type=int, + nargs="*", + choices=[0, 1], + default=[0], + help="""KV scale layout: 0 for per-tensor, 1 for per-token.""", + ) args = parser.parse_args() rows = [] - for dtype, batch_size, shape, block_size, trans_v in itertools.product( - args.dtype, args.batch, args.shapes, args.block_size, args.trans_v + for dtype, batch_size, shape, block_size, trans_v, per_token in itertools.product( + args.dtype, + args.batch, + args.shapes, + args.block_size, + args.trans_v, + args.per_token, ): num_query_heads, num_kv_heads, head_dim, context_length = shape rows.append( @@ -1478,6 +1504,7 @@ def main(): dtype, bool(trans_v), args.max_partitions, + bool(per_token), ) ) From 60d4474fa5ce88c82d98bf6aafb93e8d77dbf106 Mon Sep 17 00:00:00 2001 From: root Date: Mon, 14 Sep 2026 09:18:32 +0000 Subject: [PATCH 40/66] [FlyDSL] perf(pa_decode): improve page-16 and page-128 load scheduling Overlap page-16 V loads with QK/softmax for gfx950 D128 scalar-scale decode and relax the corresponding K/V scheduling fences. Extend page-128 single-query prefetch to smaller grids. Expand transposed-V, variable-length, and empty-context regressions, and include the prefetch regression in the standard CI entry point. Signed-off-by: root --- .github/scripts/aiter_test.sh | 8 +-- aiter/ops/flydsl/kernels/pa_decode_tile.py | 48 ++++++++-------- aiter/ops/flydsl/pa_decode.py | 41 +++++++++----- op_tests/test_flydsl_pa_decode.py | 64 ++++++++++++++-------- 4 files changed, 96 insertions(+), 65 deletions(-) diff --git a/.github/scripts/aiter_test.sh b/.github/scripts/aiter_test.sh index 49873b4cd1d..77f40d57e16 100755 --- a/.github/scripts/aiter_test.sh +++ b/.github/scripts/aiter_test.sh @@ -132,11 +132,11 @@ for file in "${sharded_files[@]}"; do set -euo pipefail test_file=$1 python3 -m pytest -q \ - "${test_file}::test_pa_decode_api_matches_gluon" \ - "${test_file}::test_pa_decode_maps_gluon_buffers_and_scale_layout" \ - "${test_file}::test_gluon_unsupported_flydsl_reducer_falls_back" \ - "${test_file}::test_gluon_supported_flydsl_reducer_does_not_mask_errors" \ + "${test_file}::test_pa_decode_maps_buffers_and_scale_layout" \ + "${test_file}::test_recommended_splits_has_configurable_upper_clamp" \ + "${test_file}::test_recommended_splits_rejects_invalid_upper_clamp" \ "${test_file}::test_v_prefetch_workgroup_interval" \ + "${test_file}::test_decode_prefetch_mixed_contexts" \ "${test_file}::test_pa_decode_fixed_length_accuracy" \ "${test_file}::test_pa_decode_variable_length_accuracy" \ "${test_file}::test_large_negative_logits_preserve_online_softmax" \ diff --git a/aiter/ops/flydsl/kernels/pa_decode_tile.py b/aiter/ops/flydsl/kernels/pa_decode_tile.py index 609d86e9a26..5114f40a465 100644 --- a/aiter/ops/flydsl/kernels/pa_decode_tile.py +++ b/aiter/ops/flydsl/kernels/pa_decode_tile.py @@ -68,6 +68,7 @@ def compile_pa_decode_tile( prefetch_v: bool = False, prefetch_v_iglp: bool = False, match_gluon_numerics: bool = False, + early_v_only: bool = False, ): """Build the tile-programming PA-decode kernel + launch wrapper. @@ -132,30 +133,26 @@ def compile_pa_decode_tile( TOTAL_ROWS = query_length * query_group_size M_TILES = cdiv(TOTAL_ROWS, MFMA_MNK) ROWS_PADDED = M_TILES * MFMA_MNK - # The gfx950 scalar-scale decode paths can opt into earlier V loads. The - # page-16 single-M-tile variant additionally delays the next K load until - # immediately before the current PV, shortening its live range across the - # softmax while leaving the existing schedule as the default. - tune_page128 = ( - is_gfx950 - and head_dim == 128 - and block_size == 128 - and trans_v - and not per_token_kv + tune_scalar_v = is_gfx950 and head_dim == 128 and trans_v and not per_token_kv + tune_page128 = tune_scalar_v and block_size == 128 + tune_page16_vpipe = tune_scalar_v and block_size == 16 + # Native page-16 decode overlaps V with QK/softmax and relaxes the K/V + # scheduling fences. The opt-in V/K pipeline also moves next K and changes + # the softmax instruction sequence, so keep its selection independent. + PAGE16_EARLY_V = ( + early_v_only + and not prefetch_v + and tune_page16_vpipe + and query_length == 1 + and M_TILES == 1 ) - tune_page16_vpipe = ( - is_gfx950 - and head_dim == 128 - and block_size == 16 - and trans_v - and not per_token_kv - ) - EARLY_V = prefetch_v and (tune_page128 or tune_page16_vpipe) and M_TILES == 1 + EARLY_V = ( + prefetch_v and (tune_page128 or tune_page16_vpipe) and M_TILES == 1 + ) or PAGE16_EARLY_V PAGE16_VPIPE = prefetch_v and tune_page16_vpipe and M_TILES == 1 PAGE16_VPIPE_IGLP = PAGE16_VPIPE and prefetch_v_iglp - # Keep the single-M-tile softmax rewrite confined to the two opt-in - # MiniMax-M3 experiments. In particular, every default/non-target kernel - # continues to instantiate the original instruction stream. + # Keep the softmax rewrite confined to the opt-in experiments. Native + # early-V scheduling retains the original numerical policy. M1_SCALE_BEFORE_MASK = PAGE16_VPIPE or match_gluon_numerics P_BUFFERS = 2 if tune_page128 and M_TILES == 3 else 1 # PV layout: V=A, P=B -> output [head-dim (row), query-row (col=lane16)], @@ -553,8 +550,9 @@ def _k_ops(phys, a): * QK_CHUNK_ELEMS, ) w = _k_load16(base) # head[he_idx*16 : +16] -> k_step 2*qkhe, 2*qkhe+1 - if const_expr(block_size == 16): - # help the scheduler overlap the PAGES_PER_CHUNK gathered loads + if const_expr(block_size == 16 and not PAGE16_EARLY_V): + # Native early V needs K/V loads to interleave with MFMA + # and softmax; retain the fences for the other schedules. fx.rocdl.sched_barrier(fx.rocdl.mask_vmem_rd) ops.extend([w[0], w[1]]) return ops # N_SUBCHUNKS i64 operands @@ -813,8 +811,8 @@ def _v_ops(phys_row, vh): + page_step * 16, ) w = _v_load16(base) - if const_expr(block_size == 16): - # help the scheduler overlap the per-page gathered loads (see _k_ops) + if const_expr(block_size == 16 and not PAGE16_EARLY_V): + # Match the K-load scheduling policy in _k_ops. fx.rocdl.sched_barrier(fx.rocdl.mask_vmem_rd) ops.extend([w[0], w[1]]) if const_expr(head_dim == 64): diff --git a/aiter/ops/flydsl/pa_decode.py b/aiter/ops/flydsl/pa_decode.py index d8045d01249..dcabaf639af 100644 --- a/aiter/ops/flydsl/pa_decode.py +++ b/aiter/ops/flydsl/pa_decode.py @@ -87,15 +87,20 @@ def _workgroup_count_enables_v_prefetch( num_kv_heads: int, num_partitions: int, num_compute_units: int, + query_length: int = 1, ) -> bool: - """Return whether the launch is in the tuned one-to-two-CTA-per-CU window. + """Return whether the launch fits the tuned page-128 V-prefetch range. + + Single-query decode benefits on grids up to two CTAs per CU, including + grids smaller than one CTA per CU. MTP retains the one-to-two-CTA window. This result is a kernel specialization and therefore part of the compile cache key. Applications serving shapes on both sides of this interval must materialize both variants before graph capture or latency-sensitive use. """ workgroups = num_sequences * num_kv_heads * num_partitions - return num_compute_units < workgroups <= 2 * num_compute_units + min_workgroups = 0 if query_length == 1 else num_compute_units + return min_workgroups < workgroups <= 2 * num_compute_units def _flydsl_dtype_str(dtype: torch.dtype) -> str: @@ -574,27 +579,34 @@ def normalize_scale(scale, name): and query_group_size in (8, 16) ) - # Early V loads help the measured one-to-two-workgroup-per-CU regime. - # Keep other grids on the existing schedule by default; the opt-in override - # lets long-context under-filled grids be measured without changing policy. + # Native page-16 decode overlaps V with QK/softmax without selecting the + # opt-in pipeline or numerical ablation. Page-128 decode also benefits on + # smaller grids, up to two CTAs/CU, and keeps its explicit override. prefetch_v = page16_vpipe + early_v_only = False if ( arch == "gfx950" and head_dim == 128 - and block_size == 128 + and block_size in (16, 128) and trans_v and not per_token_kv and query_length * query_group_size <= 16 ): - num_cus = torch.cuda.get_device_properties(dev).multi_processor_count - prefetch_v = _env_enabled(_PAGE128_EARLY_V_ENV) or ( - _workgroup_count_enables_v_prefetch( - num_seqs, - num_kv_heads, - num_partitions, - num_cus, + if block_size == 16: + early_v_only = ( + query_length == 1 and not page16_vpipe and not match_gluon_numerics + ) + else: + num_cus = torch.cuda.get_device_properties(dev).multi_processor_count + prefetch_v = _env_enabled(_PAGE128_EARLY_V_ENV) or ( + _workgroup_count_enables_v_prefetch( + num_seqs, + num_kv_heads, + num_partitions, + num_cus, + query_length=query_length, + ) ) - ) with torch.cuda.device(dev): compiled = compile_pa_decode_tile( @@ -611,6 +623,7 @@ def normalize_scale(scale, name): prefetch_v=prefetch_v, prefetch_v_iglp=page16_vpipe_iglp, match_gluon_numerics=match_gluon_numerics, + early_v_only=early_v_only, ) if num_partitions == 1: diff --git a/op_tests/test_flydsl_pa_decode.py b/op_tests/test_flydsl_pa_decode.py index 580e9ceebb6..e711e784f68 100644 --- a/op_tests/test_flydsl_pa_decode.py +++ b/op_tests/test_flydsl_pa_decode.py @@ -194,10 +194,20 @@ def capture_reduce( @pytest.mark.parametrize( - "batch_size,expected", - [(32, False), (33, True), (64, True), (65, False)], + "batch_size,query_length,expected", + [ + (1, 1, True), + (32, 1, True), + (33, 1, True), + (64, 1, True), + (65, 1, False), + (32, 2, False), + (33, 2, True), + (64, 2, True), + (65, 2, False), + ], ) -def test_v_prefetch_workgroup_interval(batch_size, expected): +def test_v_prefetch_workgroup_interval(batch_size, query_length, expected): """Warmup boundaries cover both compile-time prefetch variants.""" if pa_decode is None: pytest.skip("FlyDSL is not available") @@ -208,6 +218,7 @@ def test_v_prefetch_workgroup_interval(batch_size, expected): num_kv_heads=1, num_partitions=8, num_compute_units=256, + query_length=query_length, ) is expected ) @@ -602,16 +613,18 @@ def _run_accuracy_case( @pytest.mark.parametrize( - "block_size,query_dtype,num_kv_heads,num_partitions", + "block_size,query_dtype,num_kv_heads,num_partitions,trans_v", [ - pytest.param(16, dtypes.bf16, 1, 1, id="page16-bf16-hkv1-np1"), - pytest.param(64, dtypes.fp16, 2, 1, id="page64-fp16-hkv2-np1"), - pytest.param(128, dtypes.bf16, 2, 4, id="page128-bf16-hkv2-np4"), - pytest.param(128, dtypes.fp16, 1, 4, id="page128-fp16-hkv1-np4"), + pytest.param(16, dtypes.bf16, 1, 1, False, id="page16-bf16-hkv1-np1"), + pytest.param(64, dtypes.fp16, 2, 1, False, id="page64-fp16-hkv2-np1"), + pytest.param(128, dtypes.bf16, 2, 4, False, id="page128-bf16-hkv2-np4"), + pytest.param(128, dtypes.fp16, 1, 4, False, id="page128-fp16-hkv1-np4"), + pytest.param(16, dtypes.bf16, 1, 1, True, id="page16-bf16-hkv1-np1-transv"), + pytest.param(16, dtypes.fp16, 2, 4, True, id="page16-fp16-hkv2-np4-transv"), ], ) def test_pa_decode_fixed_length_accuracy( - block_size, query_dtype, num_kv_heads, num_partitions + block_size, query_dtype, num_kv_heads, num_partitions, trans_v ): _run_accuracy_case( [257, 257, 257], @@ -622,20 +635,23 @@ def test_pa_decode_fixed_length_accuracy( query_dtype=query_dtype, num_partitions=num_partitions, tolerance=FIXED_LENGTH_ACCURACY_TOLERANCE, + trans_v=trans_v, ) @pytest.mark.parametrize("query_dtype", [dtypes.bf16, dtypes.fp16]) -def test_pa_decode_variable_length_accuracy(query_dtype): +@pytest.mark.parametrize("block_size,trans_v", [(128, False), (16, True)]) +def test_pa_decode_variable_length_accuracy(query_dtype, block_size, trans_v): _run_accuracy_case( [1, 127, 257, 1027], num_query_heads=8, num_kv_heads=2, head_dim=128, - block_size=128, + block_size=block_size, query_dtype=query_dtype, num_partitions=4, tolerance=VARIABLE_LENGTH_ACCURACY_TOLERANCE, + trans_v=trans_v, ) @@ -1189,12 +1205,10 @@ def test_tp4_contexts(block_size, context_length, query_length, per_token): def test_tp4_decode_batches(batch_size, block_size, context_length): """Served TP4 decode grid: Hq16, Hkv1, D128, transposed V. - Batch is the only axis that moves the CTA count, and at NP=8 the sweep - straddles the tuned V-prefetch window (256 < batch * 8 <= 512 on a 256-CU - part), so page-128 batch 64 compiles the prefetch specialization while the - smaller batches compile the default schedule. Page 16 never enters that - window -- the gate also requires block_size == 128 -- so it covers the - sixteen-pages-per-tile block-table walk over the same grids instead. + Batch is the only axis that moves the CTA count. At NP=8 the page-128 + sweep stays within its tuned decode range (batch * 8 <= 512 on a 256-CU + part). Page 16 uses early V loads across these grids and walks sixteen + physical pages per compute tile. """ _run_accuracy_case( [context_length] * batch_size, @@ -1372,15 +1386,21 @@ def test_mtp3_page128_partition_boundaries( @pytest.mark.parametrize("all_empty", [False, True]) -def test_page128_decode_prefetch_mixed_contexts(all_empty): - """A two-CTA-per-CU decode grid handles empty partitions and poisoned padding.""" +@pytest.mark.parametrize("batch_size", [8, 64]) +@pytest.mark.parametrize("block_size", [16, 128]) +def test_decode_prefetch_mixed_contexts(all_empty, batch_size, block_size): + """Early V loads handle empty partitions and poisoned page-table padding.""" _require_gpu() - lengths = [0] * 64 if all_empty else [0, 1, 63, 127, 128, 129, 257, 2051] * 8 + lengths = ( + [0] * batch_size + if all_empty + else [0, 1, 63, 127, 128, 129, 257, 2051] * (batch_size // 8) + ) _constant_page_decode_case( lengths, - block_size=128, + block_size=block_size, num_partitions=8, - block_table_width=0 if all_empty else 17, + block_table_width=(max(lengths) + block_size - 1) // block_size, poison_padding=True, trans_v=True, ) From 3d53dca93486a34781c614196000ea1ae5bf73f6 Mon Sep 17 00:00:00 2001 From: root Date: Mon, 14 Sep 2026 09:50:36 +0000 Subject: [PATCH 41/66] Revert PA load-scheduling optimization for separate review Remove the page-16 and page-128 load-scheduling changes from 60d4474fa5ce88c82d98bf6aafb93e8d77dbf106 so they can be reviewed in a separate PR targeting flydsl-pa-decode-tile. Retain the CI test-selector corrections required by preceding test-cleanup and configurable-split commits. --- .github/scripts/aiter_test.sh | 1 - aiter/ops/flydsl/kernels/pa_decode_tile.py | 48 ++++++++-------- aiter/ops/flydsl/pa_decode.py | 41 +++++--------- op_tests/test_flydsl_pa_decode.py | 64 ++++++++-------------- 4 files changed, 61 insertions(+), 93 deletions(-) diff --git a/.github/scripts/aiter_test.sh b/.github/scripts/aiter_test.sh index 77f40d57e16..1ef643c1f54 100755 --- a/.github/scripts/aiter_test.sh +++ b/.github/scripts/aiter_test.sh @@ -136,7 +136,6 @@ for file in "${sharded_files[@]}"; do "${test_file}::test_recommended_splits_has_configurable_upper_clamp" \ "${test_file}::test_recommended_splits_rejects_invalid_upper_clamp" \ "${test_file}::test_v_prefetch_workgroup_interval" \ - "${test_file}::test_decode_prefetch_mixed_contexts" \ "${test_file}::test_pa_decode_fixed_length_accuracy" \ "${test_file}::test_pa_decode_variable_length_accuracy" \ "${test_file}::test_large_negative_logits_preserve_online_softmax" \ diff --git a/aiter/ops/flydsl/kernels/pa_decode_tile.py b/aiter/ops/flydsl/kernels/pa_decode_tile.py index 5114f40a465..609d86e9a26 100644 --- a/aiter/ops/flydsl/kernels/pa_decode_tile.py +++ b/aiter/ops/flydsl/kernels/pa_decode_tile.py @@ -68,7 +68,6 @@ def compile_pa_decode_tile( prefetch_v: bool = False, prefetch_v_iglp: bool = False, match_gluon_numerics: bool = False, - early_v_only: bool = False, ): """Build the tile-programming PA-decode kernel + launch wrapper. @@ -133,26 +132,30 @@ def compile_pa_decode_tile( TOTAL_ROWS = query_length * query_group_size M_TILES = cdiv(TOTAL_ROWS, MFMA_MNK) ROWS_PADDED = M_TILES * MFMA_MNK - tune_scalar_v = is_gfx950 and head_dim == 128 and trans_v and not per_token_kv - tune_page128 = tune_scalar_v and block_size == 128 - tune_page16_vpipe = tune_scalar_v and block_size == 16 - # Native page-16 decode overlaps V with QK/softmax and relaxes the K/V - # scheduling fences. The opt-in V/K pipeline also moves next K and changes - # the softmax instruction sequence, so keep its selection independent. - PAGE16_EARLY_V = ( - early_v_only - and not prefetch_v - and tune_page16_vpipe - and query_length == 1 - and M_TILES == 1 + # The gfx950 scalar-scale decode paths can opt into earlier V loads. The + # page-16 single-M-tile variant additionally delays the next K load until + # immediately before the current PV, shortening its live range across the + # softmax while leaving the existing schedule as the default. + tune_page128 = ( + is_gfx950 + and head_dim == 128 + and block_size == 128 + and trans_v + and not per_token_kv ) - EARLY_V = ( - prefetch_v and (tune_page128 or tune_page16_vpipe) and M_TILES == 1 - ) or PAGE16_EARLY_V + tune_page16_vpipe = ( + is_gfx950 + and head_dim == 128 + and block_size == 16 + and trans_v + and not per_token_kv + ) + EARLY_V = prefetch_v and (tune_page128 or tune_page16_vpipe) and M_TILES == 1 PAGE16_VPIPE = prefetch_v and tune_page16_vpipe and M_TILES == 1 PAGE16_VPIPE_IGLP = PAGE16_VPIPE and prefetch_v_iglp - # Keep the softmax rewrite confined to the opt-in experiments. Native - # early-V scheduling retains the original numerical policy. + # Keep the single-M-tile softmax rewrite confined to the two opt-in + # MiniMax-M3 experiments. In particular, every default/non-target kernel + # continues to instantiate the original instruction stream. M1_SCALE_BEFORE_MASK = PAGE16_VPIPE or match_gluon_numerics P_BUFFERS = 2 if tune_page128 and M_TILES == 3 else 1 # PV layout: V=A, P=B -> output [head-dim (row), query-row (col=lane16)], @@ -550,9 +553,8 @@ def _k_ops(phys, a): * QK_CHUNK_ELEMS, ) w = _k_load16(base) # head[he_idx*16 : +16] -> k_step 2*qkhe, 2*qkhe+1 - if const_expr(block_size == 16 and not PAGE16_EARLY_V): - # Native early V needs K/V loads to interleave with MFMA - # and softmax; retain the fences for the other schedules. + if const_expr(block_size == 16): + # help the scheduler overlap the PAGES_PER_CHUNK gathered loads fx.rocdl.sched_barrier(fx.rocdl.mask_vmem_rd) ops.extend([w[0], w[1]]) return ops # N_SUBCHUNKS i64 operands @@ -811,8 +813,8 @@ def _v_ops(phys_row, vh): + page_step * 16, ) w = _v_load16(base) - if const_expr(block_size == 16 and not PAGE16_EARLY_V): - # Match the K-load scheduling policy in _k_ops. + if const_expr(block_size == 16): + # help the scheduler overlap the per-page gathered loads (see _k_ops) fx.rocdl.sched_barrier(fx.rocdl.mask_vmem_rd) ops.extend([w[0], w[1]]) if const_expr(head_dim == 64): diff --git a/aiter/ops/flydsl/pa_decode.py b/aiter/ops/flydsl/pa_decode.py index dcabaf639af..d8045d01249 100644 --- a/aiter/ops/flydsl/pa_decode.py +++ b/aiter/ops/flydsl/pa_decode.py @@ -87,20 +87,15 @@ def _workgroup_count_enables_v_prefetch( num_kv_heads: int, num_partitions: int, num_compute_units: int, - query_length: int = 1, ) -> bool: - """Return whether the launch fits the tuned page-128 V-prefetch range. - - Single-query decode benefits on grids up to two CTAs per CU, including - grids smaller than one CTA per CU. MTP retains the one-to-two-CTA window. + """Return whether the launch is in the tuned one-to-two-CTA-per-CU window. This result is a kernel specialization and therefore part of the compile cache key. Applications serving shapes on both sides of this interval must materialize both variants before graph capture or latency-sensitive use. """ workgroups = num_sequences * num_kv_heads * num_partitions - min_workgroups = 0 if query_length == 1 else num_compute_units - return min_workgroups < workgroups <= 2 * num_compute_units + return num_compute_units < workgroups <= 2 * num_compute_units def _flydsl_dtype_str(dtype: torch.dtype) -> str: @@ -579,34 +574,27 @@ def normalize_scale(scale, name): and query_group_size in (8, 16) ) - # Native page-16 decode overlaps V with QK/softmax without selecting the - # opt-in pipeline or numerical ablation. Page-128 decode also benefits on - # smaller grids, up to two CTAs/CU, and keeps its explicit override. + # Early V loads help the measured one-to-two-workgroup-per-CU regime. + # Keep other grids on the existing schedule by default; the opt-in override + # lets long-context under-filled grids be measured without changing policy. prefetch_v = page16_vpipe - early_v_only = False if ( arch == "gfx950" and head_dim == 128 - and block_size in (16, 128) + and block_size == 128 and trans_v and not per_token_kv and query_length * query_group_size <= 16 ): - if block_size == 16: - early_v_only = ( - query_length == 1 and not page16_vpipe and not match_gluon_numerics - ) - else: - num_cus = torch.cuda.get_device_properties(dev).multi_processor_count - prefetch_v = _env_enabled(_PAGE128_EARLY_V_ENV) or ( - _workgroup_count_enables_v_prefetch( - num_seqs, - num_kv_heads, - num_partitions, - num_cus, - query_length=query_length, - ) + num_cus = torch.cuda.get_device_properties(dev).multi_processor_count + prefetch_v = _env_enabled(_PAGE128_EARLY_V_ENV) or ( + _workgroup_count_enables_v_prefetch( + num_seqs, + num_kv_heads, + num_partitions, + num_cus, ) + ) with torch.cuda.device(dev): compiled = compile_pa_decode_tile( @@ -623,7 +611,6 @@ def normalize_scale(scale, name): prefetch_v=prefetch_v, prefetch_v_iglp=page16_vpipe_iglp, match_gluon_numerics=match_gluon_numerics, - early_v_only=early_v_only, ) if num_partitions == 1: diff --git a/op_tests/test_flydsl_pa_decode.py b/op_tests/test_flydsl_pa_decode.py index e711e784f68..580e9ceebb6 100644 --- a/op_tests/test_flydsl_pa_decode.py +++ b/op_tests/test_flydsl_pa_decode.py @@ -194,20 +194,10 @@ def capture_reduce( @pytest.mark.parametrize( - "batch_size,query_length,expected", - [ - (1, 1, True), - (32, 1, True), - (33, 1, True), - (64, 1, True), - (65, 1, False), - (32, 2, False), - (33, 2, True), - (64, 2, True), - (65, 2, False), - ], + "batch_size,expected", + [(32, False), (33, True), (64, True), (65, False)], ) -def test_v_prefetch_workgroup_interval(batch_size, query_length, expected): +def test_v_prefetch_workgroup_interval(batch_size, expected): """Warmup boundaries cover both compile-time prefetch variants.""" if pa_decode is None: pytest.skip("FlyDSL is not available") @@ -218,7 +208,6 @@ def test_v_prefetch_workgroup_interval(batch_size, query_length, expected): num_kv_heads=1, num_partitions=8, num_compute_units=256, - query_length=query_length, ) is expected ) @@ -613,18 +602,16 @@ def _run_accuracy_case( @pytest.mark.parametrize( - "block_size,query_dtype,num_kv_heads,num_partitions,trans_v", + "block_size,query_dtype,num_kv_heads,num_partitions", [ - pytest.param(16, dtypes.bf16, 1, 1, False, id="page16-bf16-hkv1-np1"), - pytest.param(64, dtypes.fp16, 2, 1, False, id="page64-fp16-hkv2-np1"), - pytest.param(128, dtypes.bf16, 2, 4, False, id="page128-bf16-hkv2-np4"), - pytest.param(128, dtypes.fp16, 1, 4, False, id="page128-fp16-hkv1-np4"), - pytest.param(16, dtypes.bf16, 1, 1, True, id="page16-bf16-hkv1-np1-transv"), - pytest.param(16, dtypes.fp16, 2, 4, True, id="page16-fp16-hkv2-np4-transv"), + pytest.param(16, dtypes.bf16, 1, 1, id="page16-bf16-hkv1-np1"), + pytest.param(64, dtypes.fp16, 2, 1, id="page64-fp16-hkv2-np1"), + pytest.param(128, dtypes.bf16, 2, 4, id="page128-bf16-hkv2-np4"), + pytest.param(128, dtypes.fp16, 1, 4, id="page128-fp16-hkv1-np4"), ], ) def test_pa_decode_fixed_length_accuracy( - block_size, query_dtype, num_kv_heads, num_partitions, trans_v + block_size, query_dtype, num_kv_heads, num_partitions ): _run_accuracy_case( [257, 257, 257], @@ -635,23 +622,20 @@ def test_pa_decode_fixed_length_accuracy( query_dtype=query_dtype, num_partitions=num_partitions, tolerance=FIXED_LENGTH_ACCURACY_TOLERANCE, - trans_v=trans_v, ) @pytest.mark.parametrize("query_dtype", [dtypes.bf16, dtypes.fp16]) -@pytest.mark.parametrize("block_size,trans_v", [(128, False), (16, True)]) -def test_pa_decode_variable_length_accuracy(query_dtype, block_size, trans_v): +def test_pa_decode_variable_length_accuracy(query_dtype): _run_accuracy_case( [1, 127, 257, 1027], num_query_heads=8, num_kv_heads=2, head_dim=128, - block_size=block_size, + block_size=128, query_dtype=query_dtype, num_partitions=4, tolerance=VARIABLE_LENGTH_ACCURACY_TOLERANCE, - trans_v=trans_v, ) @@ -1205,10 +1189,12 @@ def test_tp4_contexts(block_size, context_length, query_length, per_token): def test_tp4_decode_batches(batch_size, block_size, context_length): """Served TP4 decode grid: Hq16, Hkv1, D128, transposed V. - Batch is the only axis that moves the CTA count. At NP=8 the page-128 - sweep stays within its tuned decode range (batch * 8 <= 512 on a 256-CU - part). Page 16 uses early V loads across these grids and walks sixteen - physical pages per compute tile. + Batch is the only axis that moves the CTA count, and at NP=8 the sweep + straddles the tuned V-prefetch window (256 < batch * 8 <= 512 on a 256-CU + part), so page-128 batch 64 compiles the prefetch specialization while the + smaller batches compile the default schedule. Page 16 never enters that + window -- the gate also requires block_size == 128 -- so it covers the + sixteen-pages-per-tile block-table walk over the same grids instead. """ _run_accuracy_case( [context_length] * batch_size, @@ -1386,21 +1372,15 @@ def test_mtp3_page128_partition_boundaries( @pytest.mark.parametrize("all_empty", [False, True]) -@pytest.mark.parametrize("batch_size", [8, 64]) -@pytest.mark.parametrize("block_size", [16, 128]) -def test_decode_prefetch_mixed_contexts(all_empty, batch_size, block_size): - """Early V loads handle empty partitions and poisoned page-table padding.""" +def test_page128_decode_prefetch_mixed_contexts(all_empty): + """A two-CTA-per-CU decode grid handles empty partitions and poisoned padding.""" _require_gpu() - lengths = ( - [0] * batch_size - if all_empty - else [0, 1, 63, 127, 128, 129, 257, 2051] * (batch_size // 8) - ) + lengths = [0] * 64 if all_empty else [0, 1, 63, 127, 128, 129, 257, 2051] * 8 _constant_page_decode_case( lengths, - block_size=block_size, + block_size=128, num_partitions=8, - block_table_width=(max(lengths) + block_size - 1) // block_size, + block_table_width=0 if all_empty else 17, poison_padding=True, trans_v=True, ) From 64a07d6db43577a110bdbadb53391fc7000314e3 Mon Sep 17 00:00:00 2001 From: fsx950223 Date: Tue, 15 Sep 2026 03:55:32 +0000 Subject: [PATCH 42/66] perf(flydsl): optimize FP8 per-token MTP4 paged attention Simplify per-token dispatch and keep the fast numerical path. Add automatic query splitting for small MTP4 grids, K128 MFMA for fused MTP4, shared V-scale processing, and pipelined Q/K/V loads with safe full-tile causal-mask bypass. Add regression coverage for dispatch, partition and causal-mask boundaries. Keep native FlyDSL reduction and its ABI unchanged; do not modify C++ reduction. Validated the exact source with FlyDSL 0.3.2: 395 tests passed, 1 skipped; the 84-case NP8 MTP4 matrix passed strict 0.005 checks. The separate B200 NP3 variable-length comparison retains a known strict-0.005 numerical limitation in both implementations. --- aiter/ops/flydsl/kernels/pa_decode_tile.py | 773 ++++++++++++++------- aiter/ops/flydsl/pa_decode.py | 84 ++- op_tests/test_flydsl_pa_decode.py | 654 +++++++++++++++++ 3 files changed, 1235 insertions(+), 276 deletions(-) diff --git a/aiter/ops/flydsl/kernels/pa_decode_tile.py b/aiter/ops/flydsl/kernels/pa_decode_tile.py index 609d86e9a26..ed5749eda9e 100644 --- a/aiter/ops/flydsl/kernels/pa_decode_tile.py +++ b/aiter/ops/flydsl/kernels/pa_decode_tile.py @@ -7,6 +7,13 @@ ``mfma_f32_16x16x32_fp8_fp8``; Q (bf16/f16) and the softmax probabilities P are quantized to fp8 too. Scales fold out of the matmuls (q/key scale into the QK score, value scale + 1/FP8_MAX into the epilogue); softmax max/sum stay f32. +The gfx950 BF16 per-token MTP4 specialization uses K=128 FP8 MFMA atoms, +with the same normalized Q/P values and operand layouts. +The tuned gfx950 BF16 scalar-scale decode specialization instead casts Q/P +directly to fp8, without range normalization or the compensating 1/FP8_MAX. +This is a speed/precision tradeoff: small Q values and long probability tails +can round to zero, and Q must fit the fp8 range. Per-token scales retain the +range-normalized Q/P policy. ``key_scale``/``value_scale`` are either a ``[1]`` per-tensor scalar or a ``[num_blocks, num_kv_heads, block_size]`` per-token tensor (chosen by rank). @@ -21,7 +28,7 @@ * ``context_lengths`` [num_seqs] int32 * ``output`` [num_seqs, num_q_heads, head_dim] same dtype as query -One CTA (4 waves) per (seq, kv_head) runs a flash-style online softmax over +One CTA (4 waves) per (seq, kv_head, partition, query group) runs online softmax over 256-token blocks; the 4 waves split tokens for Q.KT and head-dim for P.V, with an LDS round-trip on P transposing ownership between the two MMAs. """ @@ -66,17 +73,17 @@ def compile_pa_decode_tile( trans_v: bool = True, wide_kv_addressing: bool = False, prefetch_v: bool = False, - prefetch_v_iglp: bool = False, - match_gluon_numerics: bool = False, + query_splits: int = 1, ): """Build the tile-programming PA-decode kernel + launch wrapper. ``block_size``, ``head_dim``, and ``query_dtype`` are compile-time constants. ``query_length`` (MTP) and ``query_group_size`` flatten into - ``TOTAL_ROWS = query_length * query_group_size``, tiled into - ``M_TILES = ceil(TOTAL_ROWS / 16)`` independent 16-row MFMA tiles; each - extra M-tile duplicates loop-carried state, so VGPR/LDS/occupancy scale - roughly linearly with ``M_TILES``. + ``TOTAL_ROWS = query_length * query_group_size``. ``query_splits`` assigns + equal groups of query positions to separate CTAs, each processing + ``M_TILES = ceil(TOTAL_ROWS / query_splits / 16)`` MFMA tiles. The partial + output layout remains unsplit so the existing FlyDSL reducer can combine + context partitions without an intermediate transpose. Tokens at/after ``context_len`` still take part in the PV matmul with a zero probability, so their V bytes must be finite: the MFMA does not treat a zero @@ -112,52 +119,84 @@ def compile_pa_decode_tile( ), f"pa_decode_tile only supports query_dtype in ('f16', 'bf16'), got {query_dtype}" Q_DTYPE = fx.BFloat16 if query_dtype == "bf16" else fx.Float16 - if match_gluon_numerics: - assert ( - is_gfx950 - and head_dim == 128 - and block_size in (16, 128) - and trans_v - and not per_token_kv - and query_dtype == "bf16" - and query_length == 1 - and query_group_size in (8, 16) - ), "matched Gluon numerics is restricted to the MiniMax-M3 page-16/page-128 shapes" + # Fix the tuned scalar-scale shape to the direct-fp8 fast path. This + # removes the numerical-policy switch, not the fp8 range/precision tradeoff + # documented above. Per-token and other shapes keep normalized conversion. + SCALAR_FP8_DECODE = ( + is_gfx950 + and head_dim == 128 + and block_size in (16, 128) + and trans_v + and not per_token_kv + and query_dtype == "bf16" + and query_length == 1 + and query_group_size in (8, 16) + ) assert ( head_dim % 64 == 0 ), f"pa_decode_tile only supports head_dim that's a multiple of 64, got {head_dim}" assert query_length >= 1, f"query_length must be >= 1, got {query_length}" + 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 # Flattened query-row axis (MTP outer, GQA head inner), tiled into 16-row M-tiles. TOTAL_ROWS = query_length * query_group_size - M_TILES = cdiv(TOTAL_ROWS, MFMA_MNK) + CTA_ROWS = QUERIES_PER_CTA * query_group_size + M_TILES = cdiv(CTA_ROWS, MFMA_MNK) ROWS_PADDED = M_TILES * MFMA_MNK - # The gfx950 scalar-scale decode paths can opt into earlier V loads. The - # page-16 single-M-tile variant additionally delays the next K load until - # immediately before the current PV, shortening its live range across the - # softmax while leaving the existing schedule as the default. + # CDNA4 contracts four legacy fp8 K=32 groups in one K=128 instruction. + # Concatenating the existing packs preserves cache/LDS layouts for both + # fused and query-split MTP4 CTAs. + WIDE_FP8_MFMA = ( + is_gfx950 + and per_token_kv + and query_dtype == "bf16" + and head_dim == 128 + and block_size in (16, 128) + and query_length == 4 + and query_group_size == 16 + ) + # Small grids split query positions into separate CTAs; otherwise retain + # all four queries here and share KV loads, scales, and P publication. + MTP4_FUSED = WIDE_FP8_MFMA and M_TILES == 4 + PER_TOKEN_M1 = ( + per_token_kv + and is_gfx950 + and head_dim == 128 + and block_size in (16, 128) + and query_dtype == "bf16" + and QUERIES_PER_CTA == 1 + and query_group_size in (8, 16) + and M_TILES == 1 + and prefetch_v + ) + # The wrapper selects earlier V loads for tuned gfx950 decode shapes. + # Page 16 additionally delays next K until PV and uses IGLP to overlap + # groups of V loads with QK MFMA instructions. tune_page128 = ( is_gfx950 and head_dim == 128 and block_size == 128 - and trans_v - and not per_token_kv + and (trans_v or PER_TOKEN_M1) + and (not per_token_kv or PER_TOKEN_M1) ) tune_page16_vpipe = ( is_gfx950 and head_dim == 128 and block_size == 16 - and trans_v - and not per_token_kv + and (trans_v or PER_TOKEN_M1) + and (not per_token_kv or PER_TOKEN_M1) ) EARLY_V = prefetch_v and (tune_page128 or tune_page16_vpipe) and M_TILES == 1 PAGE16_VPIPE = prefetch_v and tune_page16_vpipe and M_TILES == 1 - PAGE16_VPIPE_IGLP = PAGE16_VPIPE and prefetch_v_iglp - # Keep the single-M-tile softmax rewrite confined to the two opt-in - # MiniMax-M3 experiments. In particular, every default/non-target kernel - # continues to instantiate the original instruction stream. - M1_SCALE_BEFORE_MASK = PAGE16_VPIPE or match_gluon_numerics - P_BUFFERS = 2 if tune_page128 and M_TILES == 3 else 1 + SCALES_BEFORE_CURRENT_V = WIDE_FP8_MFMA and ( + PAGE16_VPIPE or (PER_TOKEN_M1 and not trans_v) + ) + # Apply score scales before the -inf mask on the tuned single-M-tile + # paths. This removes the second mask without introducing -inf * 0. + M1_SCALE_BEFORE_MASK = PAGE16_VPIPE or SCALAR_FP8_DECODE or PER_TOKEN_M1 + P_BUFFERS = M_TILES if MTP4_FUSED else 2 if tune_page128 and M_TILES == 3 else 1 # PV layout: V=A, P=B -> output [head-dim (row), query-row (col=lane16)], # generalized over head_dim via the VHE_CHUNKS loop. NWARP = 4 # 4 waves / CTA @@ -202,8 +241,8 @@ def compile_pa_decode_tile( # -- LDS layout (shared across the 4 warps) -- # sQ: fp8[ROWS_PADDED,head_dim] staged+quantized query. sP: fp8[16,TILE_TOK] - # quantized probs. sQscale: f32[ROWS_PADDED] except in the matched-Gluon - # ablation, where Q has no dynamic scale. sLmax/sLsum: cross-warp scratch. + # quantized probs. sQscale: f32[ROWS_PADDED] except when Q is converted + # directly to fp8. sLmax/sLsum: cross-warp scratch. # sVPage: V page-index broadcast. sKScale/sVScale/sVScaleMax: per-token K/V # scale staging. No sO/sM/sL/sCorr: PV output is register-resident/loop-carried # and stored straight to global (V=A/P=B swap). @@ -215,11 +254,10 @@ def compile_pa_decode_tile( SP_ROW_BYTES = TILE_TOK + 16 sP_bytes = P_BUFFERS * MFMA_MNK * SP_ROW_BYTES # fp8, padded rows sQscale_off = sP_off + sP_bytes - sQscale_bytes = 0 if match_gluon_numerics else ROWS_PADDED * f32 - # Keep reduction rows 16-byte aligned on the tuned gfx950 page-128 path - # so each cross-wave row can be read as one vector. Other paths retain - # the padding used to avoid LDS bank conflicts. - NWARP_PAD = NWARP if tune_page128 else NWARP + 1 + sQscale_bytes = 0 if SCALAR_FP8_DECODE else ROWS_PADDED * f32 + # Keep tuned cross-wave reduction rows 16-byte aligned for vector reads. + # Other paths retain the original bank-conflict 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 @@ -272,9 +310,11 @@ def pa_decode_tile_kernel( warp = tid // WAVE # 0..NWARP-1 lane = tid - warp * WAVE # 0..63 seq = fx.Int32(gpu.block_id("x")) - kv_h = fx.Int32(gpu.block_id("y")) + 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) # num_kv_heads == gridDim.y + n_kv = fx.Int32(gpu.grid_dim.y) // query_splits output = fx.recast_iter(Q_DTYPE, output_ptr) pmax = fx.recast_iter(fx.Float32, pmax_ptr) @@ -327,6 +367,38 @@ 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)) + c16 = 16 + rgroup = lane // c16 # 0..3: which quarter-wave (paired with warp -> query row) + lane16 = lane - rgroup * c16 # 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 + # QCHUNK elements, loaded as N_QLOADS contiguous QLOAD_UNIT pieces + # (a buffer load is 128b max); head_dim=256 splits into 2 pieces. + return [ + _q_load_chunk(base_elem + u * QLOAD_UNIT) + for u in range_constexpr(N_QLOADS) + ] + + # Start all four independent Q loads before the initial K/page/scale + # prefetch. Their quantization and LDS publication remain below. + 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)) + 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) @@ -396,10 +468,6 @@ def _lds_load(byte_off, elem_ty, n): def _lds_store(byte_off, elem_ty, vec): fx.ptr_store(vec, _lds_ptr(byte_off, elem_ty)) - c16 = 16 - rgroup = lane // c16 # 0..3: which quarter-wave (paired with warp -> query row) - lane16 = lane - rgroup * c16 # 0..15: this row's head-dim chunk index - TOK_CHUNK = NWARP * MFMA_MNK # 64 NCHUNK = TILE_TOK // TOK_CHUNK # 4 @@ -482,6 +550,7 @@ def _v_page_fetch_and_stage(tt_i32): 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) @@ -538,6 +607,39 @@ def _load_kv_scale_vecs(a, buf_off=0): sVScale_off, a, buf_off ) + def _mfma_fp8(a_ops, b_ops, a_base, b_base, k_steps, acc): + if const_expr(WIDE_FP8_MFMA): + for ks in range_constexpr(k_steps // 4): + a_pack = fx.Vector.from_elements( + [a_ops[a_base + ks * 4 + i] for i in range_constexpr(4)], + dtype=fx.Int64, + ).bitcast(fx.Int32) + b_pack = fx.Vector.from_elements( + [b_ops[b_base + ks * 4 + i] for i in range_constexpr(4)], + 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 ks in range_constexpr(k_steps): + acc = fx.rocdl.mfma_f32_16x16x32_fp8_fp8( + T.f32x4, + [a_ops[a_base + ks], b_ops[b_base + ks], acc, 0, 0, 0], + ) + return acc + # -- raw dwordx4 K load (A operand) -- # token = warp*TOK_CHUNK + a*c16 + lane16 (the softmax mask and P-pack # write position below must encode this same formula). @@ -588,7 +690,7 @@ def _k_ops_flat(tt_i32): k_pf0, phys_vec0 = _k_ops_flat(part_start) # Issue the V page-index prefetch alongside K; the LDS write is # visible after the barrier below. - if const_expr(PAGE16_VPIPE): + if const_expr(PAGE16_VPIPE or PER_TOKEN_M1): # The K page vector is already resident. Reuse it for the V # page matrix instead of issuing a second block-table load. _stage_v_page_row(phys_vec0) @@ -658,32 +760,15 @@ def _f32_to_fp8_words(vf32): def _st_words(byte_off, words): _lds_store(byte_off, fx.Int32, words) - qh_local = warp * 4 + rgroup # 0..15: this thread's query row within an M-tile - # Each M-tile quantizes 16 rows of the flattened (MTP position, GQA # head) axis: `flat_idx = m*16 + qh_local`, decomposed as # `qi = flat_idx // query_group_size`, `gs_head = flat_idx % # query_group_size` (same convention as `_mtp_groups`). No # cross-M-tile dependency, so no barriers needed between iterations. - def _quant_q_row(m, qi, gs_head, q_row_off): - qh0 = kv_h * query_group_size + gs_head - row_byte0 = ( - (seq * query_length + 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 - # QCHUNK elements, loaded as N_QLOADS contiguous QLOAD_UNIT pieces - # (a buffer load is 128b max); head_dim=256 splits into 2 pieces. - q_units = [ - _q_load_chunk(base_elem + u * QLOAD_UNIT) - for u in range_constexpr(N_QLOADS) - ] - - if const_expr(match_gluon_numerics): - # Match Gluon's Hkv1 policy when QUERY_QUANT_MODE==-1: cast - # BF16 Q directly to the fp8 MFMA type, with no per-row absmax - # normalization and therefore no Q scale to fold into QK. + def _quant_q_row(m, q_row_off, q_units): + if const_expr(SCALAR_FP8_DECODE): + # Cast Q directly to the fp8 MFMA type; no per-row absmax + # normalization or Q scale is needed on this specialization. for u in range_constexpr(N_QLOADS): _st_words( q_row_off @@ -729,20 +814,25 @@ def _quant_q_row(m, qi, gs_head, q_row_off): qi = flat_idx // query_group_size gs_head = flat_idx - qi * query_group_size q_row_off = m * MFMA_MNK * head_dim - # `flat_idx < TOTAL_ROWS` is statically true for every lane except + # `flat_idx < CTA_ROWS` is statically true for every lane except # possibly on the last M-tile (only it can be a partial tile), so # skip the runtime branch (and its EXEC-mask overhead) entirely # for every other M-tile. - if const_expr((m + 1) * MFMA_MNK <= TOTAL_ROWS): # noqa: SIM114 - _quant_q_row(m, qi, gs_head, q_row_off) - elif flat_idx < TOTAL_ROWS: - _quant_q_row(m, qi, gs_head, q_row_off) + 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 match_gluon_numerics) and lane16 == 0: + if const_expr(not SCALAR_FP8_DECODE) and lane16 == 0: _st1(sQscale_off, qh_local * M_TILES + m, ZERO_F) gpu.barrier() @@ -821,14 +911,18 @@ def _v_ops(phys_row, vh): fx.rocdl.sched_vmem(len(ops) // 2) return ops # NVOPS i64, the 64-token contiguous run for this head - # query_length==1: plain wave-uniform `context_len` (folding in the - # per-lane `lane16` would cost a live VGPR for no behavioral change). - if const_expr(query_length == 1): - causal_bound = [context_len for _m in range_constexpr(M_TILES)] + # One query per CTA has a wave-uniform causal bound, including split + # MTP. Keep the original query position when selecting visible tokens. + if const_expr(QUERIES_PER_CTA == 1): + causal_bound = [ + context_len - (query_length - 1) + query_begin + for _m in range_constexpr(M_TILES) + ] else: causal_bound = [ context_len - (query_length - 1) + + query_begin + (m * MFMA_MNK + lane16) // query_group_size for m in range_constexpr(M_TILES) ] @@ -852,6 +946,17 @@ def _l_slot(m): 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]) + # Fused pipeline: carry next V in registers, reusing each chunk only + # after its final query has consumed the current tile's values. + V_DATA_SLOT = 2 + STATE_PER_M * M_TILES + if const_expr(MTP4_FUSED): + 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) for tt, ostate in range(part_start, part_end, 1, init=init_state): k_cur = ostate[ K_SLOT @@ -861,9 +966,9 @@ def _l_slot(m): ] # this tile's V pages, as one PAGES_PER_CHUNK-wide i32 vector tt = fx.Int32(tt) tok0 = tt * TILE_TOK - # per_tensor phase-split: let IGLP interleave MFMA with softmax - # VALU/LDS to hide the MFMA-hazard s_nop (per_token is VGPR-cliffed). - if const_expr((not per_token_kv and M_TILES > 1) or PAGE16_VPIPE_IGLP): + # Interleave MFMA with VALU/LDS in the scalar phase-split path, + # and with V loads in the tuned page-16 decode pipeline. + if const_expr((not per_token_kv and M_TILES > 1) or PAGE16_VPIPE): fx.rocdl.iglp_opt(0) tt1 = tt + 1 @@ -901,7 +1006,13 @@ def _mask_v_scale(vec, a, thr=ctx_thr, zero=zero4_scale): # V is independent of QK/softmax. Load it early for the tuned # decode path and reuse it across M-tiles in the phase-split path. v_vh_shared = None - if const_expr(M_TILES > 1 or EARLY_V): + if const_expr(MTP4_FUSED): + 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 or (EARLY_V and not SCALES_BEFORE_CURRENT_V)): v_vh_shared = [ _v_ops(v_page_cur, vh) for vh in range_constexpr(VHE_CHUNKS) ] @@ -909,7 +1020,7 @@ def _mask_v_scale(vec, a, thr=ctx_thr, zero=zero4_scale): # q_scale doesn't depend on `m`; read the whole M_TILES-wide row once # (contiguous via the transposed [qh][m] sQscale layout). q_scale_vec = None - if const_expr(M_TILES > 1 and not match_gluon_numerics): + if const_expr(M_TILES > 1): q_scale_vec = _lds_load( sQscale_off + lane16 * (M_TILES * f32), fx.Float32, M_TILES ) @@ -951,25 +1062,17 @@ def _lmax_off_m(m): frag_Ss = [] for a in range_constexpr(NCHUNK): acc = fx.Vector.filled(MFMA_ACC_ELEMS, 0.0, fx.Float32) - for s in range_constexpr(N_SUBCHUNKS): - acc = fx.rocdl.mfma_f32_16x16x32_fp8_fp8( - T.f32x4, - [ - k_cur[a * N_SUBCHUNKS + s], - q_ops_all[m * N_SUBCHUNKS + s], - acc, - 0, - 0, - 0, - ], - ) + 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 - if const_expr(match_gluon_numerics) - else scale_qk * fx.Float32(q_scale_vec[m]) - ) + scale = scale_qk * fx.Float32(q_scale_vec[m]) n_valid_tile = (causal_bound[m] - tok0).to(fx.Float32) base_tok_f = fx.Int32(warp * TOK_CHUNK + rgroup * 4).to(fx.Float32) thr = fx.Vector.from_elements( @@ -995,10 +1098,41 @@ def _lmax_off_m(m): frag_Ss[a] * scale_b for a in range_constexpr(NCHUNK) ] - masked_chunks = [ - (_ct[a] < thr).select(scaled_frags[a], neg4) - for a in range_constexpr(NCHUNK) - ] + if const_expr(MTP4_FUSED): + # All four causal rows fully cover non-tail tiles. + # Keep one uniform branch around the sixteen masks. + 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, + ) + if tile_valid < TILE_TOK + query_length - 1: + masked_all = fx.Vector.from_elements( + [ + (_ct[a] < 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 = [ + (_ct[a] < thr).select(scaled_frags[a], neg4) + for a in range_constexpr(NCHUNK) + ] pm = fx.Float32(float("-inf")) for a in range_constexpr(NCHUNK): @@ -1017,13 +1151,22 @@ def _lmax_off_m(m): # tt+1 K/V/scale prefetch, issued once here so the V-page read # reuses this barrier. - k_next = k_cur + k_next = ( + fx.Vector.filled(NCHUNK * N_SUBCHUNKS, 0, fx.Int64) + if const_expr(MTP4_FUSED) + else k_cur + ) if tt1 < part_end: - k_next, phys_vec1 = _k_ops_flat(tt1) - _v_page_fetch_and_stage(tt1) - if const_expr(per_token_kv): - # stage tt+1 into the OTHER buffer (current v_scale re-read below) + if const_expr(MTP4_FUSED): + # Publish next pages/scales now; defer next K until P + # packing releases the 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() @@ -1033,8 +1176,9 @@ def _lmax_off_m(m): v_page_next = _v_page_read_row() next_state[V_SLOT] = v_page_next - # Phase B v_scale: M_TILES>=4 re-reads per chunk (VGPR cliff); - # smaller M_TILES hold all NCHUNK. Both read the current buffer. + # Small generic M-tile groups retain the whole V-scale tile. + # Fused MTP4 instead shares one chunk across all four queries; + # other large groups reread per chunk to bound VGPR liveness. v_scale_shared = None if const_expr(per_token_kv and M_TILES < 4): v_scale_shared = [ @@ -1042,125 +1186,253 @@ def _lmax_off_m(m): for a in range_constexpr(NCHUNK) ] - 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] + if const_expr(MTP4_FUSED): + # V-scale normalization is independent of the query row. + # Compute it once instead of once per M tile. + 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) - 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 + # Give every M tile private P/Lsum storage. Publish all four + # rows with one barrier, then consume them without any + # inter-M overwrite barrier. + m_new_saved = [] + safe_max_saved = [] + for m in range_constexpr(M_TILES): + m_new = fx.maxnumf( + ostate[_m_slot(m)], + _ld_lw_row(_lmax_off_m(m), lane16).reduce( + ReductionOp.MAX, fastmath=fm_nnan + ), + fastmath=fm_nnan, ) - 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) - - m_new = fx.maxnumf( - m_prev, - _ld_lw_row(_lmax_off_m(m), lane16).reduce( - ReductionOp.MAX, fastmath=fm_nnan - ), - fastmath=fm_nnan, - ) - # Fully-invalid row: use 0 as the effective max so masked lanes - # give exp2(-inf-0)==0 (avoids the -inf-(-inf) cancellation). - 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 = [] + 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)] + # A four-value V-scale fragment is shared by all four + # query tiles without retaining the full scale tile. for a in range_constexpr(NCHUNK): - Pa = fx.Vector(exp2_f32_fast(masked_chunks[a] - m_new_b)) - 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] + v_sc = _load_scale_vec(sVScale_off, a, cur_kv_buf) + for m in range_constexpr(M_TILES): + Pa = fx.Vector( + exp2_f32_fast( + masked_chunks_saved[m][a] - safe_max_saved[m] + ) ) + ls_saved[m] = ls_saved[m] + Pa.reduce(ReductionOp.ADD) p_scaled = Pa * v_sc * norm_factor_b - elif const_expr(match_gluon_numerics): - # Match Gluon Hkv1's scalar-scale path: P is cast - # to fp8 in its native [0, 1] range. - p_scaled = Pa - else: - p_scaled = Pa * fx.Vector.filled(4, FP8_MAX, fx.Float32) - words.append(_f32_to_fp8_words(p_scaled)[0]) + word = _f32_to_fp8_words(p_scaled)[0] + p_off = ( + sP_off + + m * MFMA_MNK * SP_ROW_BYTES + + lane16 * SP_ROW_BYTES + + warp * TOK_CHUNK + + rgroup * 4 + + a * (c16 // 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 tt1 < part_end: + # Let next K overlap the P-publication barrier and PV, + # after saved scores have been consumed by P packing. + k_next = _k_ops_from_phys(_k_page_read_warp()) + next_state[K_SLOT] = k_next - p_off0 = ( - p_base + lane16 * SP_ROW_BYTES + warp * TOK_CHUNK + rgroup * 4 - ) - # The NCHUNK P words scatter across the row at stride c16//4 - # i32 (the token->fp8-lane interleave the PV ds_read_b128 - # expects); one strided store per chunk. - for a in range_constexpr(NCHUNK): - _lds_store( - p_off0 + a * (c16 // 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) - # PV output is [head-dim, query-row=lane16] after the operand - # swap, so correction/denominator are per-lane scalars (no sCorr). - # Empty history must contribute zero: exp2(-inf - safe_max). - # Replacing m_prev with 0 can overflow for negative logits. - corr_reg = fx.Float32(exp2_amdgcn_scalar(m_prev - safe_max)) - if rgroup == 0: - _st_lw(lsum_base, lane16, warp, ls) gpu.barrier() - gsum = _ld_lw_row(lsum_base, lane16).reduce(ReductionOp.ADD) - l_new = l_prev * corr_reg + gsum - p_ops = _lds_load( - p_base + lane16 * SP_ROW_BYTES + rgroup * 64, fx.Int64, NVOPS + 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 = fx.Float32(exp2_amdgcn_scalar(m_prev - safe_max)) + gsum = _ld_lw_row(lsum_base, lane16).reduce(ReductionOp.ADD) + l_new = 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] = o_acc[vh] * corr_b + op + if const_expr(m == M_TILES - 1): + # The current chunk is dead; reuse its registers + # for the next tile while the final PV finishes. + v_next_chunk = fx.Vector.filled(NVOPS, 0, fx.Int64) + if 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) + 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 - 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) - for s in range_constexpr(NVOPS): + 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) + + m_new = fx.maxnumf( + m_prev, + _ld_lw_row(_lmax_off_m(m), lane16).reduce( + ReductionOp.MAX, fastmath=fm_nnan + ), + fastmath=fm_nnan, + ) + # Fully-invalid row: use 0 as the effective max so masked lanes + # give exp2(-inf-0)==0 (avoids the -inf-(-inf) cancellation). + 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(exp2_f32_fast(masked_chunks[a] - m_new_b)) + 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_CHUNK + + rgroup * 4 + ) + # The NCHUNK P words scatter across the row at stride c16//4 + # i32 (the token->fp8-lane interleave the PV ds_read_b128 + # expects); one strided store per chunk. + for a in range_constexpr(NCHUNK): + _lds_store( + p_off0 + a * (c16 // 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) + # PV output is [head-dim, query-row=lane16] after the operand + # swap, so correction/denominator are per-lane scalars (no sCorr). + # Empty history must contribute zero: exp2(-inf - safe_max). + # Replacing m_prev with 0 can overflow for negative logits. + corr_reg = fx.Float32(exp2_amdgcn_scalar(m_prev - safe_max)) + if rgroup == 0: + _st_lw(lsum_base, lane16, warp, ls) + gpu.barrier() + gsum = _ld_lw_row(lsum_base, lane16).reduce(ReductionOp.ADD) + l_new = 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) # SWAPPED operands (V=A, P=B): output row = # head-dim, output col = query-row=lane16. - acc = fx.rocdl.mfma_f32_16x16x32_fp8_fp8( - T.f32x4, [v_vh[s], p_ops[s], acc, 0, 0, 0] - ) - 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] = o_acc[vh] * corr_b + op - next_state.extend([*o_acc, m_new, l_new]) - # A shared P/Lsum slot needs all reads to finish before - # the next M-tile overwrites it. With alternating slots, - # the next tile's write barrier retires those reads before - # that slot is reused. The Phase A barrier protects reuse - # across loop iterations. Keep the scheduler fence to - # bound register liveness across M-tiles. - if const_expr(m < M_TILES - 1): - if const_expr(P_BUFFERS == 1): - gpu.barrier() - fx.rocdl.sched_barrier(0) + 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] = o_acc[vh] * corr_b + op + next_state.extend([*o_acc, m_new, l_new]) + # A shared P/Lsum slot needs all reads to finish before + # the next M-tile overwrites it. With alternating slots, + # the next tile's write barrier retires those reads before + # that slot is reused. The Phase A barrier protects reuse + # across loop iterations. Keep the scheduler fence to + # bound register liveness across M-tiles. + if const_expr(m < M_TILES - 1): + if const_expr(P_BUFFERS == 1): + gpu.barrier() + fx.rocdl.sched_barrier(0) else: # M_TILES==1 single tile (m==0 for the _o_slot/_m_slot/_l_slot helpers). o_acc = [ostate[_o_slot(0, vh)] for vh in range_constexpr(VHE_CHUNKS)] @@ -1170,30 +1442,42 @@ def _lmax_off_m(m): frag_Ss = [] for a in range_constexpr(NCHUNK): acc = fx.Vector.filled(MFMA_ACC_ELEMS, 0.0, fx.Float32) - for s in range_constexpr(N_SUBCHUNKS): - acc = fx.rocdl.mfma_f32_16x16x32_fp8_fp8( - T.f32x4, - [k_cur[a * N_SUBCHUNKS + s], q_ops_all[s], acc, 0, 0, 0], - ) + acc = _mfma_fp8( + k_cur, q_ops_all, a * N_SUBCHUNKS, 0, N_SUBCHUNKS, acc + ) frag_Ss.append(fx.Vector(acc)) - # Prepare tt+1 before the pass-1 barrier. The page-16 deep - # pipeline publishes page ids only; other paths also load K. + # Publish next page ids and scales before the pass-1 barrier. + # Page 16 defers K until PV; page 128 overlaps K with softmax. k_next = k_cur if tt1 < part_end: if const_expr(PAGE16_VPIPE): - # Publish only the next page ids here. K itself is - # issued after the P barrier so its VMEM can overlap - # the current tile's PV without spanning softmax. - _v_page_fetch_and_stage(tt1) + # Publish next page ids/scales here. K itself is issued + # after the P barrier to overlap the current tile's PV. + 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)) + elif const_expr(PER_TOKEN_M1): + # Finish the scale VMEM-to-LDS dependency before + # issuing K, so its wait cannot drain the K prefetch. + phys_vec1 = _v_page_fetch_and_stage(tt1) + _stage_kv_scale_to_lds(phys_vec1, _kv_buf_off(tt1)) + 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): + # Next scales must reach LDS before their barrier. Issue + # current V afterward so that scale VMEM waits do not + # drain V; softmax can then hide the current V latency. + v_vh_shared = [ + _v_ops(v_page_cur, vh) for vh in range_constexpr(VHE_CHUNKS) + ] # Softmax: each lane owns one qhead (lane%16); register reduce + shuffle_xor. scale = ( scale_qk - if const_expr(match_gluon_numerics) + if const_expr(SCALAR_FP8_DECODE) else scale_qk * _ld1(sQscale_off, lane16) ) # per-qhead positive score scale n_valid_tile = (causal_bound[0] - tok0).to(fx.Float32) @@ -1355,9 +1639,8 @@ def _lmax_off_m(m): else v_scale_vecs[a] ) p_scaled = Pa * v_scale_this * norm_factor_b - elif const_expr(match_gluon_numerics): - # Gluon Hkv1 does not multiply scalar-scale P by 448 - # before its fp8 conversion. + elif const_expr(SCALAR_FP8_DECODE): + # This scalar specialization converts P directly. p_scaled = Pa else: p_scaled = Pa * fx.Vector.filled(4, FP8_MAX, fx.Float32) @@ -1406,10 +1689,7 @@ def _lmax_off_m(m): for vh in range_constexpr(VHE_CHUNKS): v_vh = v_vh_batch[vh] acc = fx.Vector.filled(MFMA_ACC_ELEMS, 0.0, fx.Float32) - for s in range_constexpr(NVOPS): - acc = fx.rocdl.mfma_f32_16x16x32_fp8_fp8( - T.f32x4, [v_vh[s], p_ops[s], acc, 0, 0, 0] - ) + 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( @@ -1417,6 +1697,8 @@ def _lmax_off_m(m): ).broadcast_to(OP_ELEMS) o_acc[vh] = o_acc[vh] * corr_b + op next_state.extend([*o_acc, m_new, l_new]) + if const_expr(MTP4_FUSED): + next_state.append(v_next) results = yield next_state o_final = results @@ -1425,13 +1707,14 @@ def _lmax_off_m(m): 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 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(per_token_kv): o_scale = inv_l - elif const_expr(match_gluon_numerics): - # The matching P path did not scale by FP8_MAX, so preserve + elif const_expr(SCALAR_FP8_DECODE): + # Direct P conversion did not scale by FP8_MAX, so preserve # value_scale but omit FlyDSL's compensating 1/FP8_MAX. o_scale = inv_l * v_scale_f else: @@ -1446,7 +1729,7 @@ def _lmax_off_m(m): def _emit(o_norm, sub, query_idx, query_head): if const_expr(NP == 1): out_offset = ( - (seq * query_length + query_idx) * stride_o_row + (seq * query_length + query_begin + query_idx) * stride_o_row + query_head * stride_o_head + sub * OP_ELEMS ) @@ -1454,7 +1737,7 @@ def _emit(o_norm, sub, query_idx, query_head): else: base = ( (seq * n_kv + kv_h) * NP + part - ) * TOTAL_ROWS + row # noqa: B023 + ) * TOTAL_ROWS + global_row # noqa: B023 pout_offset = base * head_dim + sub * OP_ELEMS fx.ptr_store( o_norm, @@ -1469,13 +1752,13 @@ def _emit(o_norm, sub, query_idx, query_head): ) sub = head_base // OP_ELEMS # Guard the partial last tile's out-of-range rows (folded away for full tiles). - if row < TOTAL_ROWS: + if row < CTA_ROWS: _emit(o_norm, sub, qi_e, qh) if const_expr(NP > 1): # noqa: SIM102 if warp == 0 and rgroup == 0: - base = ((seq * n_kv + kv_h) * NP + part) * TOTAL_ROWS + row - if row < TOTAL_ROWS: + base = ((seq * n_kv + kv_h) * NP + part) * TOTAL_ROWS + global_row + if row < CTA_ROWS: # Convert the running max from log2 units (scale_qk folds # in LOG2E) to natural-log units: the shared reduce # re-applies LOG2E itself when combining partitions. @@ -1530,7 +1813,7 @@ def pa_decode_tile_launch( stride_q_row, stride_q_head, ).launch( - grid=(num_seqs, num_kv_heads, NP), + grid=(num_seqs, num_kv_heads * query_splits, NP), block=(BLOCK_THREADS, 1, 1), stream=stream, ) diff --git a/aiter/ops/flydsl/pa_decode.py b/aiter/ops/flydsl/pa_decode.py index d8045d01249..abf1451df01 100644 --- a/aiter/ops/flydsl/pa_decode.py +++ b/aiter/ops/flydsl/pa_decode.py @@ -7,6 +7,11 @@ ``mfma_f32_16x16x32_fp8_fp8``; Q (bf16/f16) and the softmax probabilities P are quantized to fp8 too. Scales fold out of the matmuls (q/key scale into the QK score, value scale + 1/FP8_MAX into the epilogue); softmax max/sum stay f32. +The gfx950 BF16 per-token MTP4 specialization uses K=128 FP8 MFMA atoms +without changing this normalization policy. +The tuned gfx950 BF16 scalar-scale decode path uses direct fp8 Q/P conversion +instead. It trades dynamic range for speed and can truncate small probabilities; +see ``kernels.pa_decode_tile`` for its scope. Per-token scales keep normalization. ``key_scale``/``value_scale`` are either a ``[1]`` per-tensor scalar or a ``[num_blocks, num_kv_heads, block_size, 1]`` per-token tensor. @@ -22,9 +27,10 @@ * ``context_lengths`` [num_seqs] int32 * ``output`` [num_seqs, num_q_heads, head_dim] same dtype as query -One CTA (4 waves) per (seq, kv_head) runs a flash-style online softmax over -256-token blocks; the 4 waves split tokens for Q.KT and head-dim for P.V, with -an LDS round-trip on P transposing ownership between the two MMAs. +Each CTA (4 waves) processes a KV head and context partition over 256-token +blocks. Small MTP4 grids split query positions among CTAs; larger grids keep +queries fused. The waves split tokens for Q.KT and head-dim for P.V, with an +LDS round-trip on P transposing ownership between the two MMAs. """ import os @@ -45,7 +51,6 @@ _PAGE16_VPIPE_ENV = "AITER_FLYDSL_PA_PAGE16_VPIPE" _PAGE16_VPIPE_IGLP_ENV = "AITER_FLYDSL_PA_PAGE16_VPIPE_IGLP" -_MATCH_GLUON_NUMERICS_ENV = "AITER_FLYDSL_PA_MATCH_GLUON_NUMERICS" _PAGE128_EARLY_V_ENV = "AITER_FLYDSL_PA_PAGE128_EARLY_V" @@ -538,12 +543,10 @@ def normalize_scale(scale, name): # if either does. wide_kv_addressing = max(key_cache.numel(), value_cache.numel()) >= 2**31 - # The page-16 long-context pipeline is opt-in until it can be benchmarked - # on gfx950. It loads current V before QK and delays next K until the - # current PV, while every unsupported shape keeps the existing schedule. - page16_vpipe_iglp = _env_enabled(_PAGE16_VPIPE_IGLP_ENV) + # Keep the scalar page-16 pipeline opt-in. Both environment spellings + # select its tuned schedule; IGLP is now derived from that specialization. page16_vpipe = ( - (_env_enabled(_PAGE16_VPIPE_ENV) or page16_vpipe_iglp) + (_env_enabled(_PAGE16_VPIPE_ENV) or _env_enabled(_PAGE16_VPIPE_IGLP_ENV)) and arch == "gfx950" and head_dim == 128 and block_size == 16 @@ -554,26 +557,6 @@ def normalize_scale(scale, name): and query_length == 1 and query_group_size in (8, 16) ) - page16_vpipe_iglp = page16_vpipe and page16_vpipe_iglp - - # Numerical-policy ablation against Gluon's Hkv1 kernel. Keep this - # independent from the page-16 scheduling switches so the two effects can - # be measured separately, and constrain it to the exact MiniMax-M3 shapes - # whose Gluon path directly casts both Q and P to fp8. All other shapes, - # including per-token KV scales, retain FlyDSL's range-normalized policy. - match_gluon_numerics = ( - _env_enabled(_MATCH_GLUON_NUMERICS_ENV) - and arch == "gfx950" - and head_dim == 128 - and block_size in (16, 128) - and trans_v - and not per_token_kv - and query.dtype == torch.bfloat16 - and num_kv_heads == 1 - and query_length == 1 - and query_group_size in (8, 16) - ) - # Early V loads help the measured one-to-two-workgroup-per-CU regime. # Keep other grids on the existing schedule by default; the opt-in override # lets long-context under-filled grids be measured without changing policy. @@ -596,6 +579,46 @@ def normalize_scale(scale, name): ) ) + # Per-token decode benefits from overlapping V with QK. Page 16 also + # delays next K until PV; page 128 stages next scales before issuing K. + # With transposed page-128 V, more than one workgroup per CU already hides + # latency: retain the lower-register schedule in that regime. + if ( + arch == "gfx950" + and per_token_kv + and head_dim == 128 + and query.dtype == torch.bfloat16 + and num_kv_heads == 1 + and query_length == 1 + and query_group_size in (8, 16) + and block_size in (16, 128) + ): + if block_size == 16 or not trans_v: + prefetch_v = True + else: + num_cus = torch.cuda.get_device_properties(dev).multi_processor_count + prefetch_v = num_seqs * num_kv_heads * num_partitions <= num_cus + + # Split small MTP4 grids into one query per CTA, retaining the original + # partial layout and partition count. Larger grids amortize KV loads by + # processing all four queries together. Keep QL1's selector above intact. + query_splits = 1 + if ( + arch == "gfx950" + and per_token_kv + and head_dim == 128 + and query.dtype == torch.bfloat16 + and num_kv_heads == 1 + and query_length == 4 + and query_group_size == 16 + and block_size in (16, 128) + ): + num_cus = torch.cuda.get_device_properties(dev).multi_processor_count + base_workgroups = num_seqs * num_kv_heads * num_partitions + if 4 * base_workgroups <= 2 * num_cus: + query_splits = 4 + prefetch_v = True + with torch.cuda.device(dev): compiled = compile_pa_decode_tile( head_dim=head_dim, @@ -609,8 +632,7 @@ def normalize_scale(scale, name): trans_v=trans_v, wide_kv_addressing=wide_kv_addressing, prefetch_v=prefetch_v, - prefetch_v_iglp=page16_vpipe_iglp, - match_gluon_numerics=match_gluon_numerics, + query_splits=query_splits, ) if num_partitions == 1: diff --git a/op_tests/test_flydsl_pa_decode.py b/op_tests/test_flydsl_pa_decode.py index 580e9ceebb6..5bb1906e139 100644 --- a/op_tests/test_flydsl_pa_decode.py +++ b/op_tests/test_flydsl_pa_decode.py @@ -979,6 +979,425 @@ def test_extreme_tail_value_scale_is_ignored(context_length, query_length): _assert_matches(output, reference) +@pytest.mark.parametrize("num_partitions", [1, 4, 8]) +@pytest.mark.parametrize("trans_v", [False, True]) +@pytest.mark.parametrize("block_size", [16, 128]) +def test_per_token_scale_extremes_with_mixed_lengths( + block_size, trans_v, num_partitions +): + """Per-token normalization ignores tails and survives odd split starts. + + K is zero, so attention is uniform and the expected output is simply the + mean of the valid per-token V scales. Valid scales vary by 256-token tile + and 64-token wave, the first tile has an exactly-zero maximum, and unused + slots in the final physical page carry a hostile 1e9 scale. Lengths 769 + and 2305 put a one-token tail after three and nine full tiles respectively; + with four/eight partitions they also exercise odd global tile indices in + the ping-pong scale buffer. The empty row checks the neutral result. + """ + _require_gpu() + lengths = [0, 769, 2305] + batch_size, num_query_heads, num_kv_heads, head_dim = 3, 16, 1, 128 + pages_per_sequence = [(length + block_size - 1) // block_size for length in lengths] + max_pages = max(pages_per_sequence) + num_pages = sum(pages_per_sequence) + quant_dtype = _quant_dtype() + + query = torch.ones(batch_size, num_query_heads, head_dim, dtype=dtypes.bf16) + key_cache = torch.zeros( + num_pages, + num_kv_heads, + head_dim // 16, + block_size, + 16, + dtype=quant_dtype, + ) + value_plain = torch.ones( + num_pages, num_kv_heads, head_dim, block_size, dtype=quant_dtype + ) + if trans_v: + value_cache = ( + value_plain.view( + num_pages, + num_kv_heads, + head_dim, + block_size // 16, + 16, + ) + .permute(0, 1, 3, 2, 4) + .contiguous() + ) + else: + value_cache = value_plain + + block_tables = torch.zeros(batch_size, max_pages, dtype=dtypes.i32) + key_scale = torch.full( + (num_pages, num_kv_heads, block_size, 1), 1e9, dtype=dtypes.fp32 + ) + value_scale = torch.full_like(key_scale, 1e9) + expected_values = [] + next_page = 0 + for seq, (length, page_count) in enumerate(zip(lengths, pages_per_sequence)): + if length == 0: + expected_values.append(torch.zeros((), dtype=dtypes.fp32)) + continue + block_tables[seq, :page_count] = torch.arange( + next_page, next_page + page_count, dtype=dtypes.i32 + ) + capacity = page_count * block_size + token = torch.arange(capacity, dtype=dtypes.i32) + tile = token // KV_COMPUTE_BLOCK + wave = (token % KV_COMPUTE_BLOCK) // 64 + exponent = ((tile + wave) % 5 - 4).to(dtypes.fp32) + scales = torch.exp2(exponent) + scales = torch.where(token < KV_COMPUTE_BLOCK, 0.0, scales) + scales = torch.where(token < length, scales, 1e9) + page_scales = scales.reshape(page_count, block_size) + key_scale[next_page : next_page + page_count, 0, :, 0] = page_scales + value_scale[next_page : next_page + page_count, 0, :, 0] = page_scales + expected_values.append(scales[:length].mean()) + next_page += page_count + + context_lengths = torch.tensor(lengths, dtype=dtypes.i32) + output = torch.full_like(query, float("nan")) + pa_decode( + output, + query, + key_cache, + value_cache, + context_lengths, + block_tables, + head_dim**-0.5, + 1, + num_partitions, + KV_COMPUTE_BLOCK, + quant_dtype, + None, + key_scale, + value_scale, + ) + + expected = ( + torch.stack(expected_values) + .reshape(batch_size, 1, 1) + .expand(batch_size, num_query_heads, head_dim) + ) + assert torch.isfinite(output).all() + torch.testing.assert_close(output.float(), expected, rtol=0.005, atol=0.005) + + +@pytest.mark.parametrize("trans_v", [False, True]) +@pytest.mark.parametrize("block_size", [16, 128]) +def test_per_token_sparse_pages_with_odd_partition_start(block_size, trans_v): + """Page and scale prefetches stay aligned across odd-starting partitions. + + Ten compute tiles split four ways start at 0, 3, 6 and 9. The partition + starting at 3 therefore both starts in the odd scale buffer and reuses it. + Nontrivial K/V and independently varying scales expose mistakes hidden by + uniform attention. Selected physical pages are shuffled and separated by + unused pages, so a logical-page/physical-page mixup cannot pass unnoticed. + """ + _require_gpu() + context_length, num_partitions = 2305, 4 + num_query_heads, head_dim = 16, 128 + logical_pages = (context_length + block_size - 1) // block_size + num_pages = 2 * logical_pages + 1 + quant_dtype = _quant_dtype() + generator = torch.Generator(device="cuda").manual_seed(23) + pages = 2 * torch.randperm(logical_pages, generator=generator) + 1 + token = torch.arange(logical_pages * block_size) + tile = token // KV_COMPUTE_BLOCK + wave = (token % KV_COMPUTE_BLOCK) // 64 + dim = torch.arange(head_dim) + signs = (1 - 2 * (dim % 2)).float() + query_factors = 0.5 + 0.25 * (torch.arange(num_query_heads) % 3).float() + query = (query_factors[:, None] * signs[None, :]).unsqueeze(0).to(dtypes.bf16) + + key = torch.rand(num_pages, 1, block_size, head_dim, generator=generator) * 0.25 + value = torch.rand(num_pages, 1, block_size, head_dim, generator=generator) - 0.5 + # Correlated logits and tile-dependent values make a wrong K scale affect + # the answer, while variation within pages exercises every token lane. + logical_key = (0.125 + 0.0625 * ((tile + wave) % 4))[:, None] * signs[None, :] + logical_key += 0.015625 * ((token[:, None] + dim[None, :]) % 3 - 1) + logical_value = ( + 0.25 * (tile % 3 - 1)[:, None] + + 0.0625 * (token // 16 % 3 - 1)[:, None] + + 0.03125 * (dim % 5 - 2)[None, :] + ) + key[pages] = logical_key.reshape(logical_pages, 1, block_size, head_dim) + value[pages] = logical_value.reshape(logical_pages, 1, block_size, head_dim) + key = key.to(quant_dtype) + value = value.to(quant_dtype) + key_scale = torch.ones(num_pages, 1, block_size, 1, dtype=dtypes.fp32) + value_scale = torch.ones_like(key_scale) + key_scale[pages, 0, :, 0] = torch.exp2( + ((2 * tile + wave + token % 3) % 4 - 2).float() + ).reshape(logical_pages, block_size) + value_scale[pages, 0, :, 0] = torch.exp2( + ((tile + 2 * wave + token % 5) % 4 - 3).float() + ).reshape(logical_pages, block_size) + + key_cache = ( + key.view(num_pages, 1, block_size, head_dim // 16, 16) + .permute(0, 1, 3, 2, 4) + .contiguous() + ) + value_plain = value.permute(0, 1, 3, 2).contiguous() + value_cache = ( + value.view(num_pages, 1, block_size // 16, 16, head_dim) + .permute(0, 1, 2, 4, 3) + .contiguous() + if trans_v + else value_plain + ) + block_tables = pages.to(dtypes.i32).unsqueeze(0) + context_lengths = torch.tensor([context_length], dtype=dtypes.i32) + reference = run_torch( + query.float(), + key, + value_plain, + block_tables, + context_lengths, + key_scale, + value_scale, + ) + output = torch.full_like(query, float("nan")) + torch.ops.aiter.pa_decode_flydsl( + output, + query, + key_cache, + value_cache, + context_lengths, + block_tables, + head_dim**-0.5, + 1, + num_partitions, + KV_COMPUTE_BLOCK, + quant_dtype, + None, + key_scale, + value_scale, + ) + assert torch.isfinite(output).all() + torch.testing.assert_close(output.float(), reference, rtol=0.005, atol=0.005) + + +def _run_mtp4_fused_reference_case( + lengths, num_kv_heads, block_size, trans_v, num_partitions +): + """Check MTP4 with sparse pages, varying scales, and a causal FP32 reference.""" + torch.manual_seed(37) + query_length, query_group_size, head_dim = 4, 16, 128 + batch_size = len(lengths) + num_query_heads = num_kv_heads * query_group_size + pages_per_sequence = [(length + block_size - 1) // block_size for length in lengths] + max_pages = max(pages_per_sequence) + num_pages = sum(pages_per_sequence) + quant_dtype = _quant_dtype() + + query = torch.empty( + batch_size * query_length, + num_query_heads, + head_dim, + dtype=dtypes.bf16, + ).uniform_(-0.5, 0.5) + key = torch.empty( + num_pages, num_kv_heads, block_size, head_dim, dtype=dtypes.bf16 + ).uniform_(-0.5, 0.5) + value = torch.empty( + num_pages, num_kv_heads, head_dim, block_size, dtype=dtypes.bf16 + ).uniform_(-0.5, 0.5) + key_quant, key_scale = pertoken_quant(key, quant_dtype=quant_dtype) + value_token_quant, value_scale = pertoken_quant( + value.permute(0, 1, 3, 2).contiguous(), quant_dtype=quant_dtype + ) + value_quant = value_token_quant.permute(0, 1, 3, 2).contiguous() + + token = torch.arange(num_pages * block_size).reshape(num_pages, 1, block_size, 1) + kv_head = torch.arange(num_kv_heads).reshape(1, num_kv_heads, 1, 1) + key_scale *= torch.exp2(((2 * token + kv_head) % 4 - 2).float()) + value_scale *= torch.exp2(((token + 2 * kv_head) % 5 - 3).float()) + + # Make C1--C4 analytically exact while still detecting causal off-by-one: + # K=0 gives uniform attention, V progresses as [.25, .5, .75, 1], and + # scale=1. A C4 row therefore produces [.25, .375, .5, .625]. + data_page = 0 + for length, page_count in zip(lengths, pages_per_sequence): + if 0 < length <= query_length: + capacity = page_count * block_size + token_values = ( + torch.arange(capacity, dtype=dtypes.fp32) % query_length + 1 + ) * 0.25 + key_quant[data_page : data_page + page_count].zero_() + value_quant[data_page : data_page + page_count] = ( + token_values.reshape(page_count, 1, 1, block_size) + .expand(page_count, num_kv_heads, head_dim, block_size) + .to(quant_dtype) + ) + key_scale[data_page : data_page + page_count].fill_(1.0) + value_scale[data_page : data_page + page_count].fill_(1.0) + data_page += page_count + + # Scatter logical pages into a permuted set of odd physical pages. Leave + # holes and page zero finite so masked/padded token reads remain valid. + selected_pages = 2 * torch.randperm(num_pages) + 1 + physical_page_count = 2 * num_pages + 1 + + def scatter_pages(tensor): + sparse = torch.zeros( + (physical_page_count, *tensor.shape[1:]), dtype=tensor.dtype + ) + sparse[selected_pages] = tensor + return sparse + + key_quant = scatter_pages(key_quant) + value_quant = scatter_pages(value_quant) + key_scale = scatter_pages(key_scale) + value_scale = scatter_pages(value_scale) + num_pages = physical_page_count + + key_cache = ( + key_quant.view( + num_pages, + num_kv_heads, + block_size, + head_dim // 16, + 16, + ) + .permute(0, 1, 3, 2, 4) + .contiguous() + ) + if trans_v: + value_cache = ( + value_quant.permute(0, 1, 3, 2) + .contiguous() + .view( + num_pages, + num_kv_heads, + block_size // 16, + 16, + head_dim, + ) + .permute(0, 1, 2, 4, 3) + .contiguous() + ) + else: + value_cache = value_quant + + block_tables = torch.zeros(batch_size, max_pages, dtype=dtypes.i32) + next_page = 0 + for seq, page_count in enumerate(pages_per_sequence): + block_tables[seq, :page_count] = selected_pages[ + next_page : next_page + page_count + ] + next_page += page_count + context_lengths = torch.tensor(lengths, dtype=dtypes.i32) + + reference = torch.zeros_like(query, dtype=dtypes.fp32) + for seq, length in enumerate(lengths): + if length == 0: + continue + token_ids = torch.arange(length) + logical_pages = token_ids // block_size + token_offsets = token_ids % block_size + physical_pages = block_tables[seq, logical_pages].long() + keys = key_quant[physical_pages, :, token_offsets, :].float() + values = value_quant[physical_pages, :, :, token_offsets].float() + keys *= key_scale[physical_pages, :, token_offsets, 0].float().unsqueeze(-1) + values *= value_scale[physical_pages, :, token_offsets, 0].float().unsqueeze(-1) + keys = keys.repeat_interleave(query_group_size, dim=1) + values = values.repeat_interleave(query_group_size, dim=1) + for position in range(query_length): + visible = max(0, length - (query_length - 1) + position) + if visible == 0: + continue + row = seq * query_length + position + scores = ( + torch.einsum("hd,khd->hk", query[row].float(), keys[:visible]) + * head_dim**-0.5 + ) + probs = torch.softmax(scores, dim=-1) + reference[row] = torch.einsum("hk,khd->hd", probs, values[:visible]) + + output = torch.full_like(query, float("nan")) + torch.ops.aiter.pa_decode_flydsl( + output, + query, + key_cache, + value_cache, + context_lengths, + block_tables, + head_dim**-0.5, + query_length, + num_partitions, + KV_COMPUTE_BLOCK, + quant_dtype, + None, + key_scale, + value_scale, + ) + assert torch.isfinite(output).all() + torch.testing.assert_close(output.float(), reference, rtol=0.005, atol=0.005) + + +@pytest.mark.parametrize("num_partitions", [1, 8]) +@pytest.mark.parametrize("trans_v", [False, True]) +@pytest.mark.parametrize("block_size", [16, 128]) +@pytest.mark.parametrize( + "lengths,num_kv_heads", + [ + pytest.param([4], 1, id="c4-hkv1"), + pytest.param([257], 1, id="c257-hkv1"), + pytest.param([259], 1, id="c259-hkv1"), + pytest.param([2305], 1, id="c2305-hkv1"), + pytest.param([2307], 1, id="c2307-hkv1"), + pytest.param( + [0, 1, 2, 3, 4, 255, 256, 257, 258, 259, 2303, 2304, 2305, 2306, 2307], + 1, + id="mixed-hkv1", + ), + pytest.param([257], 2, id="c257-hkv2"), + pytest.param([259], 2, id="c259-hkv2"), + ], +) +def test_mtp4_fused_phase_b_accuracy( + monkeypatch, lengths, num_kv_heads, block_size, trans_v, num_partitions +): + """Force the fused path and cover its MTP causal/partition boundaries.""" + _require_gpu() + pa_decode_module = importlib.import_module("aiter.ops.flydsl.pa_decode") + compile_tile = pa_decode_module.compile_pa_decode_tile + + def compile_fused_mtp4(**kwargs): + kwargs["query_splits"] = 1 + return compile_tile(**kwargs) + + monkeypatch.setattr(pa_decode_module, "compile_pa_decode_tile", compile_fused_mtp4) + _run_mtp4_fused_reference_case( + lengths, num_kv_heads, block_size, trans_v, num_partitions + ) + + +@pytest.mark.parametrize("trans_v", [False, True]) +@pytest.mark.parametrize("block_size", [16, 128]) +def test_mtp4_public_selector_fused_path(monkeypatch, block_size, trans_v): + """B32/NP8 selects the fused one-CTA-per-sequence MTP4 path.""" + _require_gpu() + pa_decode_module = importlib.import_module("aiter.ops.flydsl.pa_decode") + compile_tile = pa_decode_module.compile_pa_decode_tile + selected_splits = [] + + def capture_query_splits(**kwargs): + selected_splits.append(kwargs.get("query_splits", 1)) + return compile_tile(**kwargs) + + monkeypatch.setattr( + pa_decode_module, "compile_pa_decode_tile", capture_query_splits + ) + _run_mtp4_fused_reference_case([257, 259] * 16, 1, block_size, trans_v, 8) + assert selected_splits and set(selected_splits) == {1} + + @pytest.mark.parametrize( "head_dim,num_partitions", [(64, 1), (256, 1), (1024, 1), (1024, 4)], @@ -1407,6 +1826,241 @@ def test_page128_reduction_rows_across_m_tiles( _assert_matches(output, reference) +@pytest.mark.parametrize("query_splits", [2, 4]) +@pytest.mark.parametrize("num_partitions", [1, 8]) +@pytest.mark.parametrize("block_size", [16, 128]) +@pytest.mark.parametrize("trans_v", [False, True]) +@pytest.mark.parametrize("context_length", [4, 257, 2305]) +def test_per_token_mtp_query_splits( + monkeypatch, query_splits, num_partitions, block_size, trans_v, context_length +): + """Query-split CTAs preserve MTP causality and the unsplit partial layout.""" + _require_gpu() + pa_decode_module = importlib.import_module("aiter.ops.flydsl.pa_decode") + compile_tile = pa_decode_module.compile_pa_decode_tile + + def compile_split_tile(**kwargs): + kwargs["query_splits"] = query_splits + kwargs["prefetch_v"] = query_splits == 4 + return compile_tile(**kwargs) + + monkeypatch.setattr(pa_decode_module, "compile_pa_decode_tile", compile_split_tile) + if context_length == 4: + # An exact short-prefix case isolates the causal/query mapping. With + # random V scales, even the unsplit FP8 kernel has >0.005 absolute + # quantization error at four tokens; do not weaken the assertion. + quant_dtype = _quant_dtype() + query = torch.ones(4, 16, 128, dtype=dtypes.bf16) + key_cache = torch.zeros(1, 1, 8, block_size, 16, dtype=quant_dtype) + value = torch.zeros(1, 1, block_size, 128, dtype=dtypes.fp32) + value[0, 0, :4] = ( + (torch.arange(4).float() + 1)[:, None] + * 0.125 + * (1 + 0.25 * (torch.arange(128) % 4).float())[None, :] + ) + value = value.to(quant_dtype) + value_cache = ( + value.view(1, 1, block_size // 16, 16, 128) + .permute(0, 1, 2, 4, 3) + .contiguous() + if trans_v + else value.permute(0, 1, 3, 2).contiguous() + ) + key_scale = torch.ones(1, 1, block_size, 1, dtype=dtypes.fp32) + value_scale = torch.ones_like(key_scale) + value_scale[0, 0, :4, 0] = torch.tensor([0.5, 1.0, 0.25, 1.0]) + prefix_values = (value[0, 0, :4].float() * value_scale[0, 0, :4]).cumsum(0) + reference = (prefix_values / (torch.arange(4).float() + 1)[:, None])[ + :, None, : + ].expand(4, 16, 128) + output = torch.full_like(query, float("nan")) + torch.ops.aiter.pa_decode_flydsl( + output, + query, + key_cache, + value_cache, + torch.tensor([4], dtype=dtypes.i32), + torch.zeros(1, 1, dtype=dtypes.i32), + 128**-0.5, + 4, + num_partitions, + KV_COMPUTE_BLOCK, + quant_dtype, + None, + key_scale, + value_scale, + ) + output = output.float() + else: + output, reference = _adversarial_case( + head_dim=128, + query_group_size=16, + query_length=4, + context_length=context_length, + block_size=block_size, + trans_v=trans_v, + per_token=True, + num_partitions=num_partitions, + ) + assert torch.isfinite(output).all() + torch.testing.assert_close(output, reference, rtol=0.005, atol=0.005) + + +@pytest.mark.parametrize("num_partitions", [1, 8]) +@pytest.mark.parametrize("block_size", [16, 128]) +@pytest.mark.parametrize("trans_v", [False, True]) +def test_per_token_mtp_query_split_mixed_contexts(num_partitions, block_size, trans_v): + """Automatic MTP splitting keeps empty and nonempty sequences independent.""" + _require_gpu() + lengths = [0, 4, 2305] + batch_size, query_length, query_heads, head_dim = 3, 4, 16, 128 + pages_per_sequence = [(length + block_size - 1) // block_size for length in lengths] + page_values = [ + (seq + 1) * 0.125 + for seq, count in enumerate(pages_per_sequence) + for _ in range(count) + ] + num_pages = len(page_values) + quant_dtype = _quant_dtype() + query = torch.ones( + batch_size * query_length, query_heads, head_dim, dtype=dtypes.bf16 + ) + key_cache = torch.zeros( + num_pages, 1, head_dim // 16, block_size, 16, dtype=quant_dtype + ) + value_plain = ( + torch.tensor(page_values, dtype=dtypes.fp32) + .reshape(num_pages, 1, 1, 1) + .expand(num_pages, 1, head_dim, block_size) + .contiguous() + .to(quant_dtype) + ) + value_cache = ( + value_plain.view(num_pages, 1, head_dim, block_size // 16, 16) + .permute(0, 1, 3, 2, 4) + .contiguous() + if trans_v + else value_plain + ) + block_tables = torch.full( + (batch_size, max(pages_per_sequence)), num_pages + 1024, dtype=dtypes.i32 + ) + next_page = 0 + for seq, count in enumerate(pages_per_sequence): + block_tables[seq, :count] = torch.arange(next_page, next_page + count) + next_page += count + scales = torch.ones(num_pages, 1, block_size, 1, dtype=dtypes.fp32) + output = torch.full_like(query, float("nan")) + torch.ops.aiter.pa_decode_flydsl( + output, + query, + key_cache, + value_cache, + torch.tensor(lengths, dtype=dtypes.i32), + block_tables, + head_dim**-0.5, + query_length, + num_partitions, + KV_COMPUTE_BLOCK, + quant_dtype, + None, + scales, + scales, + ) + expected = ( + torch.tensor([0.0, 0.25, 0.375], dtype=dtypes.fp32) + .repeat_interleave(query_length) + .reshape(batch_size * query_length, 1, 1) + .expand_as(output) + ) + assert torch.isfinite(output).all() + torch.testing.assert_close(output.float(), expected, rtol=0.005, atol=0.005) + + +@pytest.mark.parametrize( + "batch_size,overrides,expected_splits", + [ + pytest.param(1, {}, 4, id="small"), + pytest.param(16, {}, 4, id="cu-boundary"), + pytest.param(17, {}, 1, id="past-cu-boundary"), + pytest.param(32, {}, 1, id="batch32-fused"), + pytest.param(16, {"parts": 16}, 1, id="more-partitions"), + pytest.param(1, {"query_length": 1}, 1, id="decode-unchanged"), + pytest.param(1, {"head_dim": 64}, 1, id="other-head-dim"), + pytest.param(1, {"kv_heads": 2}, 1, id="multiple-kv-heads"), + pytest.param(1, {"query_group": 8}, 1, id="other-gqa"), + pytest.param(1, {"block_size": 64}, 1, id="other-page-size"), + pytest.param(1, {"dtype": dtypes.fp16}, 1, id="fp16"), + pytest.param(1, {"per_token": False}, 1, id="scalar-scales"), + ], +) +def test_mtp_query_split_dispatch(monkeypatch, batch_size, overrides, expected_splits): + """Query splitting does not widen the supported shape or partition policy.""" + _require_gpu() + module = importlib.import_module("aiter.ops.flydsl.pa_decode") + compile_arguments = [] + + def capture_compile(**kwargs): + compile_arguments.append(kwargs) + return {"launch": object()} + + monkeypatch.setattr(module, "get_gfx_runtime", lambda: "gfx950") + monkeypatch.setattr(module, "compile_pa_decode_tile", capture_compile) + monkeypatch.setattr(module, "_run_compiled", lambda *args: None) + monkeypatch.setattr(module, "ptr_arg", lambda tensor, dtype=None: tensor) + monkeypatch.setattr( + module, "launch_pa_decode_ps_reduce", lambda *args, **kwargs: None + ) + monkeypatch.setattr( + torch.cuda, + "get_device_properties", + lambda *_args, **_kwargs: type("Props", (), {"multi_processor_count": 256})(), + ) + query_length = overrides.get("query_length", 4) + head_dim = overrides.get("head_dim", 128) + kv_heads = overrides.get("kv_heads", 1) + query_group = overrides.get("query_group", 16) + block_size = overrides.get("block_size", 16) + parts = overrides.get("parts", 8) + dtype = overrides.get("dtype", dtypes.bf16) + query = torch.empty( + batch_size * query_length, kv_heads * query_group, head_dim, dtype=dtype + ) + key_cache = torch.empty( + batch_size, kv_heads, head_dim // 16, block_size, 16, dtype=torch.float8_e4m3fn + ) + value_cache = torch.empty( + batch_size, kv_heads, head_dim, block_size, dtype=torch.float8_e4m3fn + ) + scale_shape = ( + (batch_size, kv_heads, block_size, 1) + if overrides.get("per_token", True) + else (1,) + ) + scales = torch.ones(scale_shape, dtype=dtypes.fp32) + pa_decode( + torch.empty_like(query), + query, + key_cache, + value_cache, + torch.full((batch_size,), block_size, dtype=dtypes.i32), + torch.arange(batch_size, dtype=dtypes.i32).reshape(batch_size, 1), + head_dim**-0.5, + query_length, + parts, + KV_COMPUTE_BLOCK, + torch.float8_e4m3fn, + None, + scales, + scales, + ) + assert len(compile_arguments) == 1 + assert compile_arguments[0]["query_splits"] == expected_splits + assert compile_arguments[0]["num_partitions"] == parts + if expected_splits > 1: + assert compile_arguments[0]["prefetch_v"] is True + + def main(): torch.set_default_device("cuda") if not torch.cuda.is_available(): From ae7bb728bc48b3701a536823e009c22ef59eb92a Mon Sep 17 00:00:00 2001 From: fsx950223 Date: Tue, 15 Sep 2026 06:06:52 +0000 Subject: [PATCH 43/66] refactor(flydsl): remove PA prefetch environment overrides Preserve the default scalar page-16 schedule and automatic page-128, per-token, and MTP4 dispatch. Remove the three legacy environment switches and their unused reader without changing the tile kernel or reducers. Add 26 dispatch regression cases covering unset and enabled legacy variables, and include them in the CI test selector. Validated on MI355 with FlyDSL 0.3.2: 421 passed, 1 skipped; Black, Ruff, Python compilation, and shell syntax checks passed. --- .github/scripts/aiter_test.sh | 1 + aiter/ops/flydsl/pa_decode.py | 48 +++----------- op_tests/test_flydsl_pa_decode.py | 104 ++++++++++++++++++++++++++++++ 3 files changed, 114 insertions(+), 39 deletions(-) diff --git a/.github/scripts/aiter_test.sh b/.github/scripts/aiter_test.sh index 1ef643c1f54..46c09b5ed7f 100755 --- a/.github/scripts/aiter_test.sh +++ b/.github/scripts/aiter_test.sh @@ -136,6 +136,7 @@ for file in "${sharded_files[@]}"; do "${test_file}::test_recommended_splits_has_configurable_upper_clamp" \ "${test_file}::test_recommended_splits_rejects_invalid_upper_clamp" \ "${test_file}::test_v_prefetch_workgroup_interval" \ + "${test_file}::test_v_prefetch_dispatch_ignores_removed_environment_overrides" \ "${test_file}::test_pa_decode_fixed_length_accuracy" \ "${test_file}::test_pa_decode_variable_length_accuracy" \ "${test_file}::test_large_negative_logits_preserve_online_softmax" \ diff --git a/aiter/ops/flydsl/pa_decode.py b/aiter/ops/flydsl/pa_decode.py index abf1451df01..da491a7a794 100644 --- a/aiter/ops/flydsl/pa_decode.py +++ b/aiter/ops/flydsl/pa_decode.py @@ -33,8 +33,6 @@ LDS round-trip on P transposing ownership between the two MMAs. """ -import os - import flydsl.compiler as flyc import flydsl.expr as fx import torch @@ -49,19 +47,6 @@ from .kernels.tensor_shim import _run_compiled, ptr_arg from .kernels.utils import cdiv -_PAGE16_VPIPE_ENV = "AITER_FLYDSL_PA_PAGE16_VPIPE" -_PAGE16_VPIPE_IGLP_ENV = "AITER_FLYDSL_PA_PAGE16_VPIPE_IGLP" -_PAGE128_EARLY_V_ENV = "AITER_FLYDSL_PA_PAGE128_EARLY_V" - - -def _env_enabled(name: str) -> bool: - return os.getenv(name, "0").strip().lower() in { - "1", - "true", - "yes", - "on", - } - def get_recommended_splits( num_sequences: int, @@ -543,24 +528,11 @@ def normalize_scale(scale, name): # if either does. wide_kv_addressing = max(key_cache.numel(), value_cache.numel()) >= 2**31 - # Keep the scalar page-16 pipeline opt-in. Both environment spellings - # select its tuned schedule; IGLP is now derived from that specialization. - page16_vpipe = ( - (_env_enabled(_PAGE16_VPIPE_ENV) or _env_enabled(_PAGE16_VPIPE_IGLP_ENV)) - and arch == "gfx950" - and head_dim == 128 - and block_size == 16 - and trans_v - and not per_token_kv - and query.dtype == torch.bfloat16 - and num_kv_heads == 1 - and query_length == 1 - and query_group_size in (8, 16) - ) + # Scalar page-16 decode keeps its non-prefetch schedule. Per-token paths + # select their tuned pipelines below without environment overrides. + prefetch_v = False # Early V loads help the measured one-to-two-workgroup-per-CU regime. - # Keep other grids on the existing schedule by default; the opt-in override - # lets long-context under-filled grids be measured without changing policy. - prefetch_v = page16_vpipe + # Keep other scalar page-128 grids on the non-prefetch schedule. if ( arch == "gfx950" and head_dim == 128 @@ -570,13 +542,11 @@ def normalize_scale(scale, name): and query_length * query_group_size <= 16 ): num_cus = torch.cuda.get_device_properties(dev).multi_processor_count - prefetch_v = _env_enabled(_PAGE128_EARLY_V_ENV) or ( - _workgroup_count_enables_v_prefetch( - num_seqs, - num_kv_heads, - num_partitions, - num_cus, - ) + prefetch_v = _workgroup_count_enables_v_prefetch( + num_seqs, + num_kv_heads, + num_partitions, + num_cus, ) # Per-token decode benefits from overlapping V with QK. Page 16 also diff --git a/op_tests/test_flydsl_pa_decode.py b/op_tests/test_flydsl_pa_decode.py index 5bb1906e139..a994ab50597 100644 --- a/op_tests/test_flydsl_pa_decode.py +++ b/op_tests/test_flydsl_pa_decode.py @@ -2061,6 +2061,110 @@ def capture_compile(**kwargs): assert compile_arguments[0]["prefetch_v"] is True +@pytest.mark.parametrize("trans_v", [False, True]) +@pytest.mark.parametrize( + "batch_size,block_size,query_length,per_token,expected_prefetch,expected_splits", + [ + pytest.param(1, 16, 1, False, (False, False), 1, id="scalar16-small"), + pytest.param(33, 16, 1, False, (False, False), 1, id="scalar16-large"), + pytest.param(32, 128, 1, False, (False, False), 1, id="scalar128-cu"), + pytest.param(33, 128, 1, False, (False, True), 1, id="scalar128-past-cu"), + pytest.param(64, 128, 1, False, (False, True), 1, id="scalar128-2cu"), + pytest.param(65, 128, 1, False, (False, False), 1, id="scalar128-past-2cu"), + pytest.param(33, 16, 1, True, (True, True), 1, id="per-token16"), + pytest.param(32, 128, 1, True, (True, True), 1, id="per-token128-cu"), + pytest.param(33, 128, 1, True, (True, False), 1, id="per-token128-past-cu"), + pytest.param(16, 16, 4, True, (True, True), 4, id="mtp4-page16-split"), + pytest.param(17, 16, 4, True, (False, False), 1, id="mtp4-page16-fused"), + pytest.param(16, 128, 4, True, (True, True), 4, id="mtp4-page128-split"), + pytest.param(17, 128, 4, True, (False, False), 1, id="mtp4-page128-fused"), + ], +) +def test_v_prefetch_dispatch_ignores_removed_environment_overrides( + monkeypatch, + trans_v, + batch_size, + block_size, + query_length, + per_token, + expected_prefetch, + expected_splits, +): + """Legacy overrides cannot change the default automatic specialization.""" + _require_gpu() + module = importlib.import_module("aiter.ops.flydsl.pa_decode") + compile_arguments = [] + + def capture_compile(**kwargs): + compile_arguments.append(kwargs) + return {"launch": object()} + + monkeypatch.setattr(module, "get_gfx_runtime", lambda: "gfx950") + monkeypatch.setattr(module, "compile_pa_decode_tile", capture_compile) + monkeypatch.setattr(module, "_run_compiled", lambda *args: None) + monkeypatch.setattr(module, "ptr_arg", lambda tensor, dtype=None: tensor) + monkeypatch.setattr( + module, "launch_pa_decode_ps_reduce", lambda *args, **kwargs: None + ) + monkeypatch.setattr( + torch.cuda, + "get_device_properties", + lambda *_args, **_kwargs: type("Props", (), {"multi_processor_count": 256})(), + ) + parts, head_dim, query_group = 8, 128, 16 + query = torch.empty( + batch_size * query_length, query_group, head_dim, dtype=dtypes.bf16 + ) + output = torch.empty_like(query) + key_cache = torch.empty( + batch_size, 1, head_dim // 16, block_size, 16, dtype=torch.float8_e4m3fn + ) + value_shape = ( + (batch_size, 1, block_size // 16, head_dim, 16) + if trans_v + else (batch_size, 1, head_dim, block_size) + ) + value_cache = torch.empty(value_shape, dtype=torch.float8_e4m3fn) + scale_shape = (batch_size, 1, block_size, 1) if per_token else (1,) + scales = torch.ones(scale_shape, dtype=dtypes.fp32) + context_lengths = torch.full((batch_size,), block_size, dtype=dtypes.i32) + block_tables = torch.arange(batch_size, dtype=dtypes.i32).reshape(batch_size, 1) + legacy_environment_names = ( + "AITER_FLYDSL_PA_PAGE16_VPIPE", + "AITER_FLYDSL_PA_PAGE16_VPIPE_IGLP", + "AITER_FLYDSL_PA_PAGE128_EARLY_V", + ) + for enabled in (False, True): + for name in legacy_environment_names: + if enabled: + monkeypatch.setenv(name, "1") + else: + monkeypatch.delenv(name, raising=False) + pa_decode( + output, + query, + key_cache, + value_cache, + context_lengths, + block_tables, + head_dim**-0.5, + query_length, + parts, + KV_COMPUTE_BLOCK, + torch.float8_e4m3fn, + None, + scales, + scales, + ) + + assert len(compile_arguments) == 2 + assert compile_arguments[0] == compile_arguments[1] + # Expectations are indexed by the V layout: rank-4 plain, then rank-5 transposed. + assert compile_arguments[0]["prefetch_v"] is expected_prefetch[int(trans_v)] + assert compile_arguments[0]["query_splits"] == expected_splits + assert compile_arguments[0]["num_partitions"] == parts + + def main(): torch.set_default_device("cuda") if not torch.cuda.is_available(): From 60e63316e437baf94dff0be2ef2cf6d67193d941 Mon Sep 17 00:00:00 2001 From: fsx950223 Date: Tue, 15 Sep 2026 06:06:04 +0000 Subject: [PATCH 44/66] refactor(flydsl): address PA kernel cleanup feedback Signed-off-by: fsx950223 --- ...{pa_decode_tile.py => pa_decode_kernel.py} | 46 ++++++++++++------- aiter/ops/flydsl/kernels/utils.py | 43 +---------------- aiter/ops/flydsl/pa_decode.py | 25 +++------- 3 files changed, 38 insertions(+), 76 deletions(-) rename aiter/ops/flydsl/kernels/{pa_decode_tile.py => pa_decode_kernel.py} (98%) diff --git a/aiter/ops/flydsl/kernels/pa_decode_tile.py b/aiter/ops/flydsl/kernels/pa_decode_kernel.py similarity index 98% rename from aiter/ops/flydsl/kernels/pa_decode_tile.py rename to aiter/ops/flydsl/kernels/pa_decode_kernel.py index ed5749eda9e..d8bbf983ac5 100644 --- a/aiter/ops/flydsl/kernels/pa_decode_tile.py +++ b/aiter/ops/flydsl/kernels/pa_decode_kernel.py @@ -46,7 +46,7 @@ from . import dpp_utils from .tensor_shim import ptr_buf_tensor -from .utils import cdiv, exp2_amdgcn_scalar, exp2_f32_fast, rcp_f32 +from .utils import rcp_f32 MFMA_MNK = ( 16 # M = N = 16 for the MMA atom; also query rows handled per CTA (padded to 16) @@ -143,7 +143,7 @@ def compile_pa_decode_tile( # Flattened query-row axis (MTP outer, GQA head inner), tiled into 16-row M-tiles. TOTAL_ROWS = query_length * query_group_size CTA_ROWS = QUERIES_PER_CTA * query_group_size - M_TILES = cdiv(CTA_ROWS, MFMA_MNK) + M_TILES = (CTA_ROWS + MFMA_MNK - 1) // MFMA_MNK ROWS_PADDED = M_TILES * MFMA_MNK # CDNA4 contracts four legacy fp8 K=32 groups in one K=128 instruction. # Concatenating the existing packs preserves cache/LDS layouts for both @@ -203,7 +203,7 @@ def compile_pa_decode_tile( TOK_PER_WARP = 64 # tokens each warp owns per compute block (matches production KV_COMPUTE_BLOCK) TILE_TOK = NWARP * TOK_PER_WARP # 256 tokens / compute block # A warp owns 64 tokens: four page-16s, one page-64, or half a page-128. - PAGES_PER_CHUNK = cdiv(TOK_PER_WARP, block_size) + PAGES_PER_CHUNK = (TOK_PER_WARP + block_size - 1) // block_size assert ( head_dim % (NWARP * MFMA_MNK) == 0 ), "head_dim must split across the 4 warps for PV" @@ -420,8 +420,8 @@ def _v_load16(byte_off): # returns page 0 for that out-of-range read instead of faulting (those # tail tokens are masked out anyway). bt_num_records_bytes = ( - fx.Index(gpu.grid_dim.x) * fx.Index(max_blocks_per_seq) * 4 - ) # int32 entries + fx.Int64(gpu.grid_dim.x) * 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, @@ -438,9 +438,11 @@ def _v_load16(byte_off): key_scale = fx.Float32(key_scale_buf[0]) value_scale = fx.Float32(value_scale_buf[0]) - num_tiles = cdiv(context_len, TILE_TOK) - num_pages = cdiv(context_len, block_size) # pages this sequence really owns - tiles_per_part = cdiv(num_tiles, NP) + num_tiles = (context_len + TILE_TOK - 1) // TILE_TOK + num_pages = ( + context_len + block_size - 1 + ) // block_size # pages this sequence really owns + 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) @@ -1224,8 +1226,9 @@ def _lmax_off_m(m): v_sc = _load_scale_vec(sVScale_off, a, cur_kv_buf) for m in range_constexpr(M_TILES): Pa = fx.Vector( - exp2_f32_fast( - masked_chunks_saved[m][a] - safe_max_saved[m] + fx.exp2( + masked_chunks_saved[m][a] - safe_max_saved[m], + fastmath="fast", ) ) ls_saved[m] = ls_saved[m] + Pa.reduce(ReductionOp.ADD) @@ -1274,7 +1277,9 @@ def _lmax_off_m(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 = fx.Float32(exp2_amdgcn_scalar(m_prev - safe_max)) + corr_reg = fx.Float32( + fx.exp2(m_prev - safe_max, fastmath="fast") + ) gsum = _ld_lw_row(lsum_base, lane16).reduce(ReductionOp.ADD) l_new = l_prev * corr_reg + gsum p_ops = _lds_load( @@ -1360,7 +1365,9 @@ def _lmax_off_m(m): ls = fx.Float32(0.0) words = [] for a in range_constexpr(NCHUNK): - Pa = fx.Vector(exp2_f32_fast(masked_chunks[a] - m_new_b)) + 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 = ( @@ -1394,7 +1401,9 @@ def _lmax_off_m(m): # swap, so correction/denominator are per-lane scalars (no sCorr). # Empty history must contribute zero: exp2(-inf - safe_max). # Replacing m_prev with 0 can overflow for negative logits. - corr_reg = fx.Float32(exp2_amdgcn_scalar(m_prev - safe_max)) + corr_reg = fx.Float32( + fx.exp2(m_prev - safe_max, fastmath="fast") + ) if rgroup == 0: _st_lw(lsum_base, lane16, warp, ls) gpu.barrier() @@ -1618,7 +1627,9 @@ def _lmax_off_m(m): if const_expr(M1_SCALE_BEFORE_MASK): # Invalid lanes are -inf, hence exp2(-inf-safe_max)=0 # without a second validity compare/select. - Pa = fx.Vector(exp2_f32_fast(masked_chunks[a] - m_new_b)) + Pa = fx.Vector( + fx.exp2(masked_chunks[a] - m_new_b, fastmath="fast") + ) else: # Legacy path: re-mask Pa so a fully-masked chunk # contributes exactly 0. @@ -1627,7 +1638,10 @@ def _lmax_off_m(m): ) Pa = valid_a.select( fx.Vector( - exp2_f32_fast(masked_chunks[a] * scale - m_new_b) + fx.exp2( + masked_chunks[a] * scale - m_new_b, + fastmath="fast", + ) ), zero4_p, ) @@ -1659,7 +1673,7 @@ def _lmax_off_m(m): ls = ls + ls.shuffle_xor(sh, WAVE) # PV (V=A, P=B) -> output [head-dim, query-row=lane16]; same as # the phase-split path. - corr_reg = fx.Float32(exp2_amdgcn_scalar(m_prev - softmax_max)) + corr_reg = fx.Float32(fx.exp2(m_prev - softmax_max, fastmath="fast")) if rgroup == 0: _st_lw(sLsum_off, lane16, warp, ls) gpu.barrier() diff --git a/aiter/ops/flydsl/kernels/utils.py b/aiter/ops/flydsl/kernels/utils.py index 5ae043e1e24..19ac9d4cfe4 100644 --- a/aiter/ops/flydsl/kernels/utils.py +++ b/aiter/ops/flydsl/kernels/utils.py @@ -2,9 +2,7 @@ # Copyright (C) 2025-2026 FlyDSL Project Contributors import flydsl.expr as fx -from flydsl._mlir import ir -from flydsl._mlir.dialects import llvm -from flydsl.expr import arith, const_expr, rocdl +from flydsl.expr import const_expr, rocdl from flydsl.expr.typing import T @@ -12,45 +10,6 @@ def rcp_f32(value): return rocdl.rcp(T.f32, value) -def exp2_amdgcn_scalar(scalar_value): - raw = ( - arith.unwrap(scalar_value) - if hasattr(scalar_value, "ir_value") or hasattr(scalar_value, "type") - else scalar_value - ) - f32_ty = ir.F32Type.get() - return llvm.call_intrinsic(f32_ty, "llvm.amdgcn.exp2.f32", [raw], [], []) - - -def exp2_f32_fast(value): - from flydsl._mlir.dialects import vector as _vector_dialect - - raw = ( - arith.unwrap(value) - if hasattr(value, "ir_value") or hasattr(value, "type") - else value - ) - ty = raw.type - if isinstance(ty, ir.VectorType): - n = ty.shape[0] - elems = [] - for i in range(n): - scalar = _vector_dialect.extract( - raw, static_position=[i], dynamic_position=[] - ) - elems.append(exp2_amdgcn_scalar(scalar)) - return _vector_dialect.from_elements(ty, elems) - return exp2_amdgcn_scalar(raw) - - -def cdiv(numer: int, denom: int) -> int: - return (numer + denom - 1) // denom - - -# Alias: several kernels historically spelled this ``ceildiv``. -ceildiv = cdiv - - def align_up(value: int, align: int) -> int: """Round *value* up to the next multiple of *align* (static ints).""" return ((int(value) + int(align) - 1) // int(align)) * int(align) diff --git a/aiter/ops/flydsl/pa_decode.py b/aiter/ops/flydsl/pa_decode.py index da491a7a794..cb821fcc94b 100644 --- a/aiter/ops/flydsl/pa_decode.py +++ b/aiter/ops/flydsl/pa_decode.py @@ -11,7 +11,7 @@ without changing this normalization policy. The tuned gfx950 BF16 scalar-scale decode path uses direct fp8 Q/P conversion instead. It trades dynamic range for speed and can truncate small probabilities; -see ``kernels.pa_decode_tile`` for its scope. Per-token scales keep normalization. +see ``kernels.pa_decode_kernel`` for its scope. Per-token scales keep normalization. ``key_scale``/``value_scale`` are either a ``[1]`` per-tensor scalar or a ``[num_blocks, num_kv_heads, block_size, 1]`` per-token tensor. @@ -39,13 +39,12 @@ 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_reduce import ( MAX_CONTEXT_PARTITIONS, compile_pa_decode_ps_reduce, ) -from .kernels.pa_decode_tile import KV_COMPUTE_BLOCK, compile_pa_decode_tile -from .kernels.tensor_shim import _run_compiled, ptr_arg -from .kernels.utils import cdiv +from .kernels.tensor_shim import _run_compiled, get_dtype_str, ptr_arg def get_recommended_splits( @@ -68,7 +67,7 @@ def get_recommended_splits( 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 = cdiv(num_sm, denom) * split_kv_blocks + n = ((num_sm + denom - 1) // denom) * split_kv_blocks return max(4, min(n, max_partitions)) @@ -88,16 +87,6 @@ def _workgroup_count_enables_v_prefetch( return num_compute_units < workgroups <= 2 * num_compute_units -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}") - - def _flydsl_pointer_dtype(dtype: torch.dtype): return { torch.float32: fx.Float32, @@ -142,9 +131,9 @@ def launch_pa_decode_ps_reduce( compiled = compile_pa_decode_ps_reduce( max_context_partition_num=context_partition_num, head_size=head_size, - output_dtype_str=_flydsl_dtype_str(output.dtype), - logits_dtype_str=_flydsl_dtype_str(logits.dtype), - sink_dtype_str=_flydsl_dtype_str( + 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, From b4d34f6deccc9839fc47a404267e9d665087fc30 Mon Sep 17 00:00:00 2001 From: fsx950223 Date: Tue, 15 Sep 2026 07:32:15 +0000 Subject: [PATCH 45/66] test(flydsl): extend and streamline PA benchmark --- .github/scripts/aiter_test.sh | 13 +- op_tests/test_flydsl_pa_decode.py | 874 +++++++++++++----------------- 2 files changed, 392 insertions(+), 495 deletions(-) diff --git a/.github/scripts/aiter_test.sh b/.github/scripts/aiter_test.sh index 46c09b5ed7f..e944789d2ed 100755 --- a/.github/scripts/aiter_test.sh +++ b/.github/scripts/aiter_test.sh @@ -132,19 +132,14 @@ for file in "${sharded_files[@]}"; do set -euo pipefail test_file=$1 python3 -m pytest -q \ - "${test_file}::test_pa_decode_maps_buffers_and_scale_layout" \ "${test_file}::test_recommended_splits_has_configurable_upper_clamp" \ - "${test_file}::test_recommended_splits_rejects_invalid_upper_clamp" \ - "${test_file}::test_v_prefetch_workgroup_interval" \ - "${test_file}::test_v_prefetch_dispatch_ignores_removed_environment_overrides" \ + "${test_file}::test_pa_decode_benchmark_reference_causal_mtp" \ + "${test_file}::test_pa_decode_benchmark_cli_defaults_and_bs200" \ + "${test_file}::test_pa_decode_benchmark_entry" \ "${test_file}::test_pa_decode_fixed_length_accuracy" \ "${test_file}::test_pa_decode_variable_length_accuracy" \ "${test_file}::test_large_negative_logits_preserve_online_softmax" \ - "${test_file}::test_additional_supported_head_dims_are_accurate" \ - "${test_file}::test_supported_head_dim_boundaries_are_accepted" \ - "${test_file}::test_unsupported_head_dim_is_rejected" \ - "${test_file}::test_ps_false_is_rejected" \ - "${test_file}::test_invalid_tensor_structure_is_rejected" + "${test_file}::test_additional_supported_head_dims_are_accurate" python3 "$test_file" ' _ "$file" diff --git a/op_tests/test_flydsl_pa_decode.py b/op_tests/test_flydsl_pa_decode.py index a994ab50597..9d5ffa4826c 100644 --- a/op_tests/test_flydsl_pa_decode.py +++ b/op_tests/test_flydsl_pa_decode.py @@ -1,7 +1,18 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. -"""Correctness and performance sweep for FlyDSL paged-attention Tile.""" +"""Correctness and performance sweep for FlyDSL paged-attention Tile. + +For example, BS200/MTP4 with FP8 per-token KV scales and explicit split counts:: + + python3 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 have equal lengths and include the MTP query tokens. Query position +``p`` attends to ``max(0, context_length - query_length + 1 + p)`` KV tokens. +Timing includes the FlyDSL attention kernel and its native FlyDSL reduction. +""" import argparse import importlib @@ -14,7 +25,6 @@ import aiter from aiter import dtypes, per_tensor_quant, pertoken_quant from aiter.jit.utils.chip_info import get_gfx_runtime -from aiter.ops.attention import pa_decode_flydsl as public_pa_decode from aiter.test_common import benchmark, checkAllclose, run_perftest @@ -51,10 +61,12 @@ def _default_cuda_device(): try: from aiter.ops.flydsl.pa_decode import ( + MAX_CONTEXT_PARTITIONS, get_recommended_splits, pa_decode, ) except (ImportError, AttributeError, RuntimeError, OSError): + MAX_CONTEXT_PARTITIONS = 256 get_recommended_splits = None pa_decode = None @@ -74,145 +86,6 @@ def _require_gpu(): pytest.skip(f"pa_decode is unsupported on {get_gfx_runtime()}") -def test_pa_decode_maps_buffers_and_scale_layout(monkeypatch): - _require_gpu() - - pa_decode_module = importlib.import_module("aiter.ops.flydsl.pa_decode") - attention_module = importlib.import_module("aiter.ops.attention") - captured = {} - - monkeypatch.setattr( - pa_decode_module, - "compile_pa_decode_tile", - lambda **kwargs: {"launch": object()}, - ) - - def capture_launch( - launch, - output, - max_logits, - exp_sums, - temporary_output, - query, - key_cache, - value_cache, - block_tables, - context_lengths, - key_scale, - value_scale, - *args, - ): - captured.update( - max_logits=max_logits, - exp_sums=exp_sums, - temporary_output=temporary_output, - key_scale=key_scale, - value_scale=value_scale, - ) - - monkeypatch.setattr(pa_decode_module, "_run_compiled", capture_launch) - monkeypatch.setattr(pa_decode_module, "ptr_arg", lambda tensor, dtype=None: tensor) - - def capture_reduce( - output_5d, reduce_exp_sums, reduce_max_logits, logits, *args, **kwargs - ): - captured.update( - reduce_output=output_5d, - reduce_exp_sums=reduce_exp_sums, - reduce_max_logits=reduce_max_logits, - reduce_logits=logits, - reduce_kwargs=kwargs, - ) - - monkeypatch.setattr(pa_decode_module, "launch_pa_decode_ps_reduce", capture_reduce) - - query = torch.empty(1, 8, 128, dtype=torch.bfloat16) - output = torch.empty_like(query) - key_cache = torch.empty(1, 1, 8, 16, 16, dtype=_quant_dtype()) - value_cache = torch.empty(1, 1, 128, 16, dtype=_quant_dtype()) - context_lengths = torch.tensor([16], dtype=torch.int32) - block_tables = torch.tensor([[0]], dtype=torch.int32) - key_scale = torch.ones(1, 1, 16, 1, dtype=torch.float32) - value_scale = torch.ones_like(key_scale) - max_logits = torch.empty(1, 1, 2, 8, dtype=torch.float32) - exp_sums = torch.empty_like(max_logits) - temporary_output = torch.empty(1, 1, 2, 8, 128, dtype=query.dtype) - - pa_decode( - output, - query, - key_cache, - value_cache, - context_lengths, - block_tables, - 128**-0.5, - 1, - 2, - compute_type=key_cache.dtype, - key_scale=key_scale, - value_scale=value_scale, - exp_sums=exp_sums, - max_logits=max_logits, - temporary_output=temporary_output, - ) - - assert captured["key_scale"].shape == (1, 1, 16) - assert captured["value_scale"].shape == (1, 1, 16) - assert captured["max_logits"].data_ptr() == max_logits.data_ptr() - assert captured["exp_sums"].data_ptr() == exp_sums.data_ptr() - assert captured["temporary_output"].data_ptr() == temporary_output.data_ptr() - assert captured["reduce_output"].data_ptr() == output.data_ptr() - assert captured["reduce_exp_sums"].data_ptr() == exp_sums.data_ptr() - assert captured["reduce_max_logits"].data_ptr() == max_logits.data_ptr() - assert captured["reduce_logits"].data_ptr() == temporary_output.data_ptr() - assert captured["reduce_kwargs"]["context_partition_num"] == 2 - - dispatches = [] - monkeypatch.setattr( - attention_module, - "_pa_decode_flydsl", - lambda *args, **kwargs: dispatches.append("flydsl"), - ) - public_pa_decode( - output, - query, - key_cache, - value_cache, - context_lengths, - block_tables, - 128**-0.5, - 1, - 2, - compute_type=key_cache.dtype, - key_scale=key_scale, - value_scale=value_scale, - exp_sums=exp_sums, - max_logits=max_logits, - temporary_output=temporary_output, - ) - assert dispatches == ["flydsl"] - - -@pytest.mark.parametrize( - "batch_size,expected", - [(32, False), (33, True), (64, True), (65, False)], -) -def test_v_prefetch_workgroup_interval(batch_size, expected): - """Warmup boundaries cover both compile-time prefetch variants.""" - if pa_decode is None: - pytest.skip("FlyDSL is not available") - pa_decode_module = importlib.import_module("aiter.ops.flydsl.pa_decode") - assert ( - pa_decode_module._workgroup_count_enables_v_prefetch( - batch_size, - num_kv_heads=1, - num_partitions=8, - num_compute_units=256, - ) - is expected - ) - - @pytest.mark.parametrize( "batch_size,max_partitions,expected", [(1, 8, 8), (1, 256, 256), (16, 256, 32), (64, 256, 8)], @@ -239,14 +112,6 @@ def test_recommended_splits_has_configurable_upper_clamp( ) -@pytest.mark.parametrize("max_partitions", [0, 3, 257]) -def test_recommended_splits_rejects_invalid_upper_clamp(max_partitions): - if get_recommended_splits is None: - pytest.skip("FlyDSL is not available") - with pytest.raises(ValueError, match="max_partitions"): - get_recommended_splits(1, 1, max_partitions=max_partitions) - - def run_torch( query: torch.Tensor, key_cache: torch.Tensor, @@ -255,16 +120,27 @@ def run_torch( context_lengths: torch.Tensor, key_scale: torch.Tensor, value_scale: torch.Tensor, + query_length: int = 1, ) -> torch.Tensor: - batch_size, num_query_heads, head_dim = query.shape + """Dequantized FP32 reference with sequence-major, causal MTP queries.""" + num_queries, num_query_heads, head_dim = query.shape + batch_size = context_lengths.numel() + if query_length < 1 or num_queries != batch_size * query_length: + raise ValueError("query must have batch_size * query_length rows") block_size = key_cache.shape[2] num_kv_heads = key_cache.shape[1] query_group_size = num_query_heads // num_kv_heads softmax_scale = head_dim**-0.5 - output = torch.empty_like(query) + output = torch.zeros_like(query) + queries = query.reshape( + batch_size, query_length, num_kv_heads, query_group_size, head_dim + ) + positions = torch.arange(query_length, device=query.device) for seq_idx in range(batch_size): context_length = int(context_lengths[seq_idx].item()) + if context_length == 0: + continue token_ids = torch.arange(context_length, device=query.device) logical_pages = token_ids // block_size token_offsets = token_ids % block_size @@ -280,13 +156,22 @@ def run_torch( token_value_scale = value_scale[physical_pages, :, token_offsets, 0].float() keys = keys * token_key_scale.unsqueeze(-1) values = values * token_value_scale.unsqueeze(-1) - keys = keys.repeat_interleave(query_group_size, dim=1) - values = values.repeat_interleave(query_group_size, dim=1) + # Keep KV heads grouped instead of materializing Hq/Hkv copies of the + # long-context cache just for the reference (notably BS200/C200k). scores = ( - torch.einsum("hd,khd->hk", query[seq_idx].float(), keys) * softmax_scale + torch.einsum("qhgd,khd->qhgk", queries[seq_idx].float(), keys) + * softmax_scale ) + visible = context_length - query_length + 1 + positions + masked = token_ids.unsqueeze(0) >= visible.unsqueeze(1) + scores.masked_fill_(masked[:, None, None, :], float("-inf")) probs = torch.softmax(scores, dim=-1) - output[seq_idx] = torch.einsum("hk,khd->hd", probs, values).to(query.dtype) + # Contexts shorter than QL have leading queries with no visible tokens. + probs.masked_fill_(visible[:, None, None, None] <= 0, 0) + first_row = seq_idx * query_length + output[first_row : first_row + query_length] = torch.einsum( + "qhgk,khd->qhgd", probs, values + ).reshape(query_length, num_query_heads, head_dim) return output @@ -340,20 +225,38 @@ def run_pa_decode_tile_case( trans_v, max_partitions=8, per_token=False, + query_length=1, + num_partitions=None, ): + if query_length < 1: + raise ValueError("query_length must be positive") + 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}]") + if batch_size < 1 or context_length < 1: + raise ValueError("batch_size and context_length must be positive") if pa_decode is None or get_recommended_splits is None: raise RuntimeError("FlyDSL is not available") if dtype not in (dtypes.fp16, dtypes.bf16): raise ValueError(f"pa_decode only supports fp16/bf16, got {dtype}") - if num_query_heads % num_kv_heads != 0: + if num_query_heads < 1 or num_kv_heads < 1 or num_query_heads % num_kv_heads != 0: raise ValueError("num_query_heads must be divisible by num_kv_heads") + # An explicit NP is an exact override, not an upper clamp. In particular, + # NP=3 and NP=5 must not be rounded by the automatic recommendation. + if num_partitions is None: + num_partitions = get_recommended_splits( + batch_size, + num_kv_heads, + split_kv_blocks=KV_COMPUTE_BLOCK // block_size, + max_partitions=max_partitions, + ) + torch.manual_seed(0) blocks_per_sequence = (context_length + block_size - 1) // block_size num_blocks = batch_size * blocks_per_sequence query = torch.empty( - batch_size, + batch_size * query_length, num_query_heads, head_dim, dtype=dtype, @@ -381,9 +284,11 @@ def run_pa_decode_tile_case( value_token_major, quant_dtype=quant_dtype ) value_quant = value_token_quant.permute(0, 1, 3, 2).contiguous() + del value_token_major, value_token_quant else: key_quant, key_scale = per_tensor_quant(key, quant_dtype=quant_dtype) value_quant, value_scale = per_tensor_quant(value, quant_dtype=quant_dtype) + del key, value key_cache = ( key_quant.view( num_blocks, @@ -419,20 +324,17 @@ def run_pa_decode_tile_case( context_lengths, key_scale, value_scale, + query_length=query_length, ) + # Only the kernel's cache layouts need to survive into allocation rotation. + del key_quant, value_quant query_group_size = num_query_heads // num_kv_heads - num_partitions = get_recommended_splits( - batch_size, - num_kv_heads, - split_kv_blocks=KV_COMPUTE_BLOCK // block_size, - max_partitions=max_partitions, - ) partial_shape = ( batch_size, num_kv_heads, num_partitions, - query_group_size, + query_length * query_group_size, ) pmax = torch.empty(partial_shape, dtype=dtypes.fp32) psum = torch.empty_like(pmax) @@ -441,10 +343,17 @@ def run_pa_decode_tile_case( candidates = {"flydsl": _run_flydsl} - # QK and PV each perform one multiply-add per query-head/context pair. - flops = 4 * batch_size * num_query_heads * context_length * head_dim + # QK and PV each perform one multiply-add per visible query/KV-token pair. + attended_tokens = sum( + max(0, context_length - query_length + 1 + position) + for position in range(query_length) + ) + flops = 4 * batch_size * num_query_heads * attended_tokens * head_dim # Effective bandwidth counts Q + O + referenced K/V tokens and metadata once. # It excludes padded tokens, repeated loads, and partition scratch traffic. + # All MTP positions share the KV cache and its scales: count them once, + # including only valid tokens rather than a partially padded final page. + scale_elements = batch_size * num_kv_heads * context_length if per_token else 1 nbytes = ( 2 * query.numel() * query.element_size() + 2 @@ -455,8 +364,7 @@ def run_pa_decode_tile_case( * key_cache.element_size() + block_tables.numel() * block_tables.element_size() + context_lengths.numel() * context_lengths.element_size() - + key_scale.numel() * key_scale.element_size() - + value_scale.numel() * value_scale.element_size() + + scale_elements * (key_scale.element_size() + value_scale.element_size()) ) ret = { @@ -491,13 +399,251 @@ def run_pa_decode_tile_case( tol_err_ratio=0.0, msg=f"{name}: pa_decode", ) + if err: + raise AssertionError(f"{name}: pa_decode mismatch ratio {err}") ret[f"{name} us"] = us + # Amortized time per output token, not the latency of an MTP step. + ret[f"{name} us/token"] = us / (batch_size * query_length) ret[f"{name} TFLOPS"] = flops / us / 1e6 ret[f"{name} TB/s"] = nbytes / us / 1e6 ret[f"{name} err"] = err return ret +def test_pa_decode_benchmark_reference_causal_mtp(): + """An analytic oracle catches flattened-row, KV-head, and causal mistakes.""" + if not torch.cuda.is_available(): + pytest.skip("ROCm is not available") + dtype = dtypes.bf16 + query_length, query_group_size, head_dim, block_size = 4, 2, 4, 16 + lengths, pages = [4, 2, 0], [1, 0, -1] + query = torch.zeros(12, 4, head_dim, dtype=dtype) + key = torch.zeros(2, 2, block_size, head_dim, dtype=dtypes.fp32) + value = ( + 100 * torch.arange(2).reshape(2, 1, 1, 1) + + 10 * torch.arange(2).reshape(1, 2, 1, 1) + + torch.arange(head_dim).reshape(1, 1, head_dim, 1) + + torch.arange(1, block_size + 1).reshape(1, 1, 1, block_size) + ).float() + value_scale = ( + torch.arange(1, block_size + 1, dtype=dtypes.fp32) + .reshape(1, 1, block_size, 1) + .expand(2, 2, block_size, 1) + .contiguous() + ) + key_scale = torch.ones_like(value_scale) + actual = run_torch( + query, + key, + value, + torch.tensor(pages, dtype=dtypes.i32).reshape(3, 1), + torch.tensor(lengths, dtype=dtypes.i32), + key_scale, + value_scale, + query_length=query_length, + ) + expected = [] + for length, page in zip(lengths, pages): + for position in range(query_length): + visible = max(0, length - query_length + 1 + position) + expected.append( + [ + [ + sum( + (100 * page + 10 * (head // query_group_size) + dim + t + 1) + * (t + 1) + for t in range(visible) + ) + / max(visible, 1) + for dim in range(head_dim) + ] + for head in range(4) + ] + ) + assert actual.dtype == dtype + assert torch.isfinite(actual).all() + torch.testing.assert_close( + actual, torch.tensor(expected, dtype=dtype), rtol=0, atol=0 + ) + + +def test_pa_decode_benchmark_cli_defaults_and_bs200(): + defaults = _parse_args([]) + assert defaults.query_length == [1] + assert defaults.num_partitions == [None] + assert defaults.max_partitions == 8 + args = _parse_args( + [ + "-d", + "bf16", + "-b", + "200", + "-s", + "16,1,128,200000", + "-q", + "4", + "--block-size", + "16", + "128", + "--trans-v", + "0", + "1", + "--per-token", + "1", + "--num-partitions", + "3", + "5", + "8", + ] + ) + assert args.dtype == [dtypes.bf16] + assert args.batch == [200] + assert args.shapes == [(16, 1, 128, 200000)] + assert args.query_length == [4] + assert args.block_size == [16, 128] + assert args.trans_v == [0, 1] + assert args.per_token == [1] + assert args.num_partitions == [3, 5, 8] + boundaries = _parse_args( + ["-q", "1", "4", "--num-partitions", "1", "256", "--max-partitions", "256"] + ) + assert boundaries.query_length == [1, 4] + assert boundaries.num_partitions == [1, 256] + assert boundaries.max_partitions == 256 + + +@pytest.mark.parametrize( + "batch_size,query_length,num_partitions,block_size,trans_v,dtype,per_token", + [ + pytest.param( + 2, 1, None, 16, False, dtypes.bf16, False, id="legacy-auto-decode" + ), + pytest.param(2, 4, 3, 16, True, dtypes.bf16, True, id="mtp4-page16-np3"), + pytest.param(2, 4, 3, 128, True, dtypes.fp16, False, id="mtp4-fp16-scalar"), + pytest.param(200, 4, 5, 128, False, dtypes.bf16, True, id="bs200-mtp4-np5"), + ], +) +def test_pa_decode_benchmark_entry( + monkeypatch, + batch_size, + query_length, + num_partitions, + block_size, + trans_v, + dtype, + per_token, +): + """Exercise the real public kernel once, without a 101-iteration benchmark.""" + _require_gpu() + context_length, num_query_heads, num_kv_heads, head_dim = 257, 16, 1, 128 + expected_partitions = 4 if num_partitions is None else num_partitions + recommendations, timed_calls = [], [] + + def recommend(*args, **kwargs): + assert num_partitions is None, "an explicit NP must bypass auto selection" + recommendations.append((args, kwargs)) + return expected_partitions + + def run_once(fn, *args, **kwargs): + assert fn is _run_flydsl + assert not kwargs + ( + output, + query, + key_cache, + value_cache, + block_tables, + context_lengths, + key_scale, + value_scale, + partitions, + softmax_scale, + pmax, + psum, + pout, + ) = args + assert output.shape == query.shape == (batch_size * query_length, 16, 128) + assert output.dtype == query.dtype == pout.dtype == dtype + assert key_cache.dtype == value_cache.dtype == _quant_dtype() + assert value_cache.ndim == (5 if trans_v else 4) + assert partitions == expected_partitions + assert softmax_scale == head_dim**-0.5 + partial_shape = (batch_size, 1, partitions, query_length * 16) + assert pmax.shape == psum.shape == partial_shape + assert pmax.dtype == psum.dtype == dtypes.fp32 + assert pout.shape == (*partial_shape, 128) + blocks_per_sequence = (context_length + block_size - 1) // block_size + assert block_tables.shape == (batch_size, blocks_per_sequence) + assert context_lengths.shape == (batch_size,) + assert bool((context_lengths == context_length).all()) + scale_elements = ( + batch_size * blocks_per_sequence * block_size if per_token else 1 + ) + assert key_scale.numel() == value_scale.numel() == scale_elements + timed_calls.append(partitions) + return fn(*args), 100.0 + + monkeypatch.setitem(globals(), "get_recommended_splits", recommend) + monkeypatch.setitem(globals(), "run_perftest", run_once) + overrides = {} + if query_length != 1: + overrides["query_length"] = query_length + if num_partitions is not None: + overrides["num_partitions"] = num_partitions + if num_partitions == 5: + # The exact override must also bypass a smaller automatic clamp. + overrides["max_partitions"] = 4 + result = run_pa_decode_tile_case( + batch_size, + num_query_heads, + num_kv_heads, + head_dim, + context_length, + block_size, + dtype, + trans_v, + per_token=per_token, + **overrides, + ) + assert timed_calls == [expected_partitions] + assert recommendations == ( + [ + ( + (batch_size, num_kv_heads), + { + "split_kv_blocks": KV_COMPUTE_BLOCK // block_size, + "max_partitions": 8, + }, + ) + ] + if num_partitions is None + else [] + ) + assert result["query_length"] == query_length + assert result["partitions"] == expected_partitions + assert result["flydsl err"] == 0 + assert result["flydsl us"] == 100.0 + assert result["flydsl us/token"] == pytest.approx( + 100.0 / (batch_size * query_length) + ) + attended_tokens = ( + query_length * context_length - query_length * (query_length - 1) // 2 + ) + flops = 4 * batch_size * num_query_heads * head_dim * attended_tokens + blocks_per_sequence = (context_length + block_size - 1) // block_size + # BF16/FP16 Q/O, FP8 KV, int32 metadata, FP32 scales. KV/scales are not + # multiplied by MTP length, and the padded part of page128 is excluded. + logical_bytes = ( + 2 * batch_size * query_length * num_query_heads * head_dim * 2 + + 2 * batch_size * num_kv_heads * context_length * head_dim + + batch_size * blocks_per_sequence * 4 + + batch_size * 4 + + (batch_size * num_kv_heads * context_length if per_token else 1) * 2 * 4 + ) + assert result["flydsl TFLOPS"] == pytest.approx(flops / 100.0 / 1e6) + assert result["flydsl TB/s"] == pytest.approx(logical_bytes / 100.0 / 1e6) + + def _run_accuracy_case( lengths, *, @@ -639,86 +785,6 @@ def test_pa_decode_variable_length_accuracy(query_dtype): ) -def _valid_pa_decode_arguments(head_dim=128, num_partitions=1): - """Small valid call used by host-contract tests that fail before compile.""" - quant_dtype = _quant_dtype() - query = torch.empty(1, 8, head_dim, dtype=dtypes.bf16, device="cuda") - scale = torch.ones(1, dtype=dtypes.fp32, device=query.device) - return { - "output": torch.empty_like(query), - "query": query, - "key_cache": torch.empty( - 1, 1, head_dim // 16, 16, 16, dtype=quant_dtype, device=query.device - ), - "value_cache": torch.empty( - 1, 1, head_dim, 16, dtype=quant_dtype, device=query.device - ), - "context_lengths": torch.tensor([16], dtype=dtypes.i32, device=query.device), - "block_tables": torch.tensor([[0]], dtype=dtypes.i32, device=query.device), - "softmax_scale": head_dim**-0.5, - "query_length": 1, - "max_context_partition_num": num_partitions, - "compute_type": quant_dtype, - "key_scale": scale, - "value_scale": scale, - } - - -def test_ps_false_is_rejected(): - _require_gpu() - with pytest.raises(NotImplementedError, match="ps=False"): - pa_decode(**_valid_pa_decode_arguments(), ps=False) - - -@pytest.mark.parametrize( - "invalid_input,error_match", - [ - pytest.param("query_rank", "query", id="query-rank"), - pytest.param("output_rank", "output", id="output-rank"), - pytest.param("key_cache_rank", "key_cache", id="key-cache-rank"), - pytest.param("value_cache_rank", "value_cache", id="value-cache-rank"), - pytest.param("context_lengths_rank", "context_lengths", id="context-rank"), - pytest.param("block_tables_rank", "block_tables", id="block-table-rank"), - pytest.param("query_device", "CUDA device", id="query-device"), - pytest.param("zero_query_heads", "at least one head", id="zero-query-heads"), - pytest.param("block_table_rows", "block_tables", id="block-table-rows"), - pytest.param("kv_block_count", "block", id="kv-block-count"), - pytest.param("kv_head_count", "head", id="kv-head-count"), - ], -) -def test_invalid_tensor_structure_is_rejected(invalid_input, error_match): - _require_gpu() - arguments = _valid_pa_decode_arguments() - if invalid_input == "query_rank": - arguments["query"] = arguments["query"].squeeze(0) - elif invalid_input == "output_rank": - arguments["output"] = arguments["output"].squeeze(0) - elif invalid_input == "key_cache_rank": - arguments["key_cache"] = arguments["key_cache"].squeeze(0) - elif invalid_input == "value_cache_rank": - arguments["value_cache"] = arguments["value_cache"].squeeze(0) - elif invalid_input == "context_lengths_rank": - arguments["context_lengths"] = arguments["context_lengths"].reshape(1, 1) - elif invalid_input == "block_tables_rank": - arguments["block_tables"] = arguments["block_tables"].flatten() - elif invalid_input == "query_device": - arguments["query"] = arguments["query"].cpu() - elif invalid_input == "zero_query_heads": - arguments["query"] = arguments["query"][:, :0] - arguments["output"] = arguments["output"][:, :0] - elif invalid_input == "block_table_rows": - arguments["block_tables"] = arguments["block_tables"].repeat(2, 1) - elif invalid_input == "kv_block_count": - arguments["value_cache"] = arguments["value_cache"].repeat(2, 1, 1, 1) - elif invalid_input == "kv_head_count": - arguments["value_cache"] = arguments["value_cache"].repeat(1, 2, 1, 1) - else: - raise AssertionError(f"unknown invalid input: {invalid_input}") - - with pytest.raises(ValueError, match=error_match): - pa_decode(**arguments) - - def _adversarial_case( head_dim=128, context_length=257, @@ -1412,42 +1478,6 @@ def test_additional_supported_head_dims_are_accurate(head_dim, num_partitions): _assert_matches(output, reference) -@pytest.mark.parametrize("num_partitions", [1, 4]) -@pytest.mark.parametrize("head_dim", [64, 128, 256, 1024]) -def test_supported_head_dim_boundaries_are_accepted( - monkeypatch, head_dim, num_partitions -): - """The public support check is identical before direct and reduced paths.""" - _require_gpu() - pa_decode_module = importlib.import_module("aiter.ops.flydsl.pa_decode") - compile_arguments = [] - - def capture_compile(**kwargs): - compile_arguments.append(kwargs) - return {"launch": object()} - - monkeypatch.setattr(pa_decode_module, "compile_pa_decode_tile", capture_compile) - monkeypatch.setattr(pa_decode_module, "_run_compiled", lambda *args: None) - monkeypatch.setattr(pa_decode_module, "ptr_arg", lambda tensor, dtype=None: tensor) - monkeypatch.setattr( - pa_decode_module, "launch_pa_decode_ps_reduce", lambda *args, **kwargs: None - ) - - pa_decode(**_valid_pa_decode_arguments(head_dim, num_partitions)) - assert len(compile_arguments) == 1 - assert compile_arguments[0]["head_dim"] == head_dim - assert compile_arguments[0]["num_partitions"] == num_partitions - - -@pytest.mark.parametrize("num_partitions", [1, 4]) -@pytest.mark.parametrize("head_dim", [96, 192, 1152]) -def test_unsupported_head_dim_is_rejected(head_dim, num_partitions): - """NP=1 and NP>1 expose the same public head-dimension support domain.""" - _require_gpu() - with pytest.raises(NotImplementedError, match="head_dim"): - pa_decode(**_valid_pa_decode_arguments(head_dim, num_partitions)) - - def _constant_page_decode_case( lengths, block_size, @@ -1977,206 +2007,14 @@ def test_per_token_mtp_query_split_mixed_contexts(num_partitions, block_size, tr torch.testing.assert_close(output.float(), expected, rtol=0.005, atol=0.005) -@pytest.mark.parametrize( - "batch_size,overrides,expected_splits", - [ - pytest.param(1, {}, 4, id="small"), - pytest.param(16, {}, 4, id="cu-boundary"), - pytest.param(17, {}, 1, id="past-cu-boundary"), - pytest.param(32, {}, 1, id="batch32-fused"), - pytest.param(16, {"parts": 16}, 1, id="more-partitions"), - pytest.param(1, {"query_length": 1}, 1, id="decode-unchanged"), - pytest.param(1, {"head_dim": 64}, 1, id="other-head-dim"), - pytest.param(1, {"kv_heads": 2}, 1, id="multiple-kv-heads"), - pytest.param(1, {"query_group": 8}, 1, id="other-gqa"), - pytest.param(1, {"block_size": 64}, 1, id="other-page-size"), - pytest.param(1, {"dtype": dtypes.fp16}, 1, id="fp16"), - pytest.param(1, {"per_token": False}, 1, id="scalar-scales"), - ], -) -def test_mtp_query_split_dispatch(monkeypatch, batch_size, overrides, expected_splits): - """Query splitting does not widen the supported shape or partition policy.""" - _require_gpu() - module = importlib.import_module("aiter.ops.flydsl.pa_decode") - compile_arguments = [] - - def capture_compile(**kwargs): - compile_arguments.append(kwargs) - return {"launch": object()} - - monkeypatch.setattr(module, "get_gfx_runtime", lambda: "gfx950") - monkeypatch.setattr(module, "compile_pa_decode_tile", capture_compile) - monkeypatch.setattr(module, "_run_compiled", lambda *args: None) - monkeypatch.setattr(module, "ptr_arg", lambda tensor, dtype=None: tensor) - monkeypatch.setattr( - module, "launch_pa_decode_ps_reduce", lambda *args, **kwargs: None - ) - monkeypatch.setattr( - torch.cuda, - "get_device_properties", - lambda *_args, **_kwargs: type("Props", (), {"multi_processor_count": 256})(), - ) - query_length = overrides.get("query_length", 4) - head_dim = overrides.get("head_dim", 128) - kv_heads = overrides.get("kv_heads", 1) - query_group = overrides.get("query_group", 16) - block_size = overrides.get("block_size", 16) - parts = overrides.get("parts", 8) - dtype = overrides.get("dtype", dtypes.bf16) - query = torch.empty( - batch_size * query_length, kv_heads * query_group, head_dim, dtype=dtype - ) - key_cache = torch.empty( - batch_size, kv_heads, head_dim // 16, block_size, 16, dtype=torch.float8_e4m3fn - ) - value_cache = torch.empty( - batch_size, kv_heads, head_dim, block_size, dtype=torch.float8_e4m3fn - ) - scale_shape = ( - (batch_size, kv_heads, block_size, 1) - if overrides.get("per_token", True) - else (1,) - ) - scales = torch.ones(scale_shape, dtype=dtypes.fp32) - pa_decode( - torch.empty_like(query), - query, - key_cache, - value_cache, - torch.full((batch_size,), block_size, dtype=dtypes.i32), - torch.arange(batch_size, dtype=dtypes.i32).reshape(batch_size, 1), - head_dim**-0.5, - query_length, - parts, - KV_COMPUTE_BLOCK, - torch.float8_e4m3fn, - None, - scales, - scales, - ) - assert len(compile_arguments) == 1 - assert compile_arguments[0]["query_splits"] == expected_splits - assert compile_arguments[0]["num_partitions"] == parts - if expected_splits > 1: - assert compile_arguments[0]["prefetch_v"] is True - - -@pytest.mark.parametrize("trans_v", [False, True]) -@pytest.mark.parametrize( - "batch_size,block_size,query_length,per_token,expected_prefetch,expected_splits", - [ - pytest.param(1, 16, 1, False, (False, False), 1, id="scalar16-small"), - pytest.param(33, 16, 1, False, (False, False), 1, id="scalar16-large"), - pytest.param(32, 128, 1, False, (False, False), 1, id="scalar128-cu"), - pytest.param(33, 128, 1, False, (False, True), 1, id="scalar128-past-cu"), - pytest.param(64, 128, 1, False, (False, True), 1, id="scalar128-2cu"), - pytest.param(65, 128, 1, False, (False, False), 1, id="scalar128-past-2cu"), - pytest.param(33, 16, 1, True, (True, True), 1, id="per-token16"), - pytest.param(32, 128, 1, True, (True, True), 1, id="per-token128-cu"), - pytest.param(33, 128, 1, True, (True, False), 1, id="per-token128-past-cu"), - pytest.param(16, 16, 4, True, (True, True), 4, id="mtp4-page16-split"), - pytest.param(17, 16, 4, True, (False, False), 1, id="mtp4-page16-fused"), - pytest.param(16, 128, 4, True, (True, True), 4, id="mtp4-page128-split"), - pytest.param(17, 128, 4, True, (False, False), 1, id="mtp4-page128-fused"), - ], -) -def test_v_prefetch_dispatch_ignores_removed_environment_overrides( - monkeypatch, - trans_v, - batch_size, - block_size, - query_length, - per_token, - expected_prefetch, - expected_splits, -): - """Legacy overrides cannot change the default automatic specialization.""" - _require_gpu() - module = importlib.import_module("aiter.ops.flydsl.pa_decode") - compile_arguments = [] - - def capture_compile(**kwargs): - compile_arguments.append(kwargs) - return {"launch": object()} - - monkeypatch.setattr(module, "get_gfx_runtime", lambda: "gfx950") - monkeypatch.setattr(module, "compile_pa_decode_tile", capture_compile) - monkeypatch.setattr(module, "_run_compiled", lambda *args: None) - monkeypatch.setattr(module, "ptr_arg", lambda tensor, dtype=None: tensor) - monkeypatch.setattr( - module, "launch_pa_decode_ps_reduce", lambda *args, **kwargs: None - ) - monkeypatch.setattr( - torch.cuda, - "get_device_properties", - lambda *_args, **_kwargs: type("Props", (), {"multi_processor_count": 256})(), - ) - parts, head_dim, query_group = 8, 128, 16 - query = torch.empty( - batch_size * query_length, query_group, head_dim, dtype=dtypes.bf16 - ) - output = torch.empty_like(query) - key_cache = torch.empty( - batch_size, 1, head_dim // 16, block_size, 16, dtype=torch.float8_e4m3fn - ) - value_shape = ( - (batch_size, 1, block_size // 16, head_dim, 16) - if trans_v - else (batch_size, 1, head_dim, block_size) - ) - value_cache = torch.empty(value_shape, dtype=torch.float8_e4m3fn) - scale_shape = (batch_size, 1, block_size, 1) if per_token else (1,) - scales = torch.ones(scale_shape, dtype=dtypes.fp32) - context_lengths = torch.full((batch_size,), block_size, dtype=dtypes.i32) - block_tables = torch.arange(batch_size, dtype=dtypes.i32).reshape(batch_size, 1) - legacy_environment_names = ( - "AITER_FLYDSL_PA_PAGE16_VPIPE", - "AITER_FLYDSL_PA_PAGE16_VPIPE_IGLP", - "AITER_FLYDSL_PA_PAGE128_EARLY_V", - ) - for enabled in (False, True): - for name in legacy_environment_names: - if enabled: - monkeypatch.setenv(name, "1") - else: - monkeypatch.delenv(name, raising=False) - pa_decode( - output, - query, - key_cache, - value_cache, - context_lengths, - block_tables, - head_dim**-0.5, - query_length, - parts, - KV_COMPUTE_BLOCK, - torch.float8_e4m3fn, - None, - scales, - scales, - ) - - assert len(compile_arguments) == 2 - assert compile_arguments[0] == compile_arguments[1] - # Expectations are indexed by the V layout: rank-4 plain, then rank-5 transposed. - assert compile_arguments[0]["prefetch_v"] is expected_prefetch[int(trans_v)] - assert compile_arguments[0]["query_splits"] == expected_splits - assert compile_arguments[0]["num_partitions"] == parts +def _positive_int(value): + value = int(value) + if value < 1: + raise argparse.ArgumentTypeError("must be a positive integer") + return value -def main(): - torch.set_default_device("cuda") - 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: - aiter.logger.warning("pa_decode unsupported on %s; skipping", get_gfx_runtime()) - return - if pa_decode is None: - aiter.logger.warning("flydsl is unavailable; skipping pa_decode") - return - +def _parse_args(argv=None): parser = argparse.ArgumentParser( formatter_class=argparse.RawTextHelpFormatter, description="FlyDSL pa_decode correctness + perf sweep", @@ -2193,12 +2031,21 @@ def main(): parser.add_argument( "-b", "--batch", - type=int, + type=_positive_int, nargs="*", default=DEFAULT_BATCH_SIZES, help="""Batch sizes. e.g.: -b 1 3 16""", ) + parser.add_argument( + "-q", + "--query-length", + type=_positive_int, + nargs="+", + default=[1], + help="""Query tokens per sequence: 1 for decode, 4 for MTP4. + e.g.: -q 1 4. MTP positions use dense causal masking.""", + ) parser.add_argument( "-s", "--shapes", @@ -2206,6 +2053,7 @@ def main(): nargs="*", default=DEFAULT_SHAPES, help="""(num_query_heads,num_kv_heads,head_dim,context_length). + Contexts are equal-length and include the MTP query tokens. e.g.: -s 8,1,128,257""", ) parser.add_argument( @@ -2229,7 +2077,16 @@ def main(): "--max-partitions", type=int, default=8, - help="""Upper clamp passed to get_recommended_splits (4..256).""", + help="""Upper clamp passed to get_recommended_splits (4..256). + Only used when --num-partitions is omitted.""", + ) + parser.add_argument( + "--num-partitions", + type=_positive_int, + nargs="+", + default=[None], + help="""Exact partition counts (1..256), bypassing auto selection. + e.g.: --num-partitions 3 5. Overrides --max-partitions.""", ) parser.add_argument( "--per-token", @@ -2239,16 +2096,59 @@ def main(): default=[0], help="""KV scale layout: 0 for per-tensor, 1 for per-token.""", ) - args = parser.parse_args() + args = parser.parse_args(argv) + if not 4 <= args.max_partitions <= MAX_CONTEXT_PARTITIONS: + parser.error(f"--max-partitions must be in [4, {MAX_CONTEXT_PARTITIONS}]") + if any( + count is not None and count > MAX_CONTEXT_PARTITIONS + for count 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(value < 1 for value 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(): + # Parse first so --help and invalid options do not initialize the 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: + aiter.logger.warning("pa_decode unsupported on %s; skipping", get_gfx_runtime()) + return + if pa_decode is None: + aiter.logger.warning("flydsl is unavailable; skipping pa_decode") + return + torch.set_default_device("cuda") rows = [] - for dtype, batch_size, shape, block_size, trans_v, per_token in itertools.product( + for ( + dtype, + batch_size, + shape, + block_size, + trans_v, + per_token, + query_length, + num_partitions, + ) in itertools.product( args.dtype, args.batch, args.shapes, args.block_size, args.trans_v, args.per_token, + args.query_length, + args.num_partitions, ): num_query_heads, num_kv_heads, head_dim, context_length = shape rows.append( @@ -2263,6 +2163,8 @@ def main(): bool(trans_v), args.max_partitions, bool(per_token), + query_length=query_length, + num_partitions=num_partitions, ) ) From b7b549c3e4f749e85fa4d17381aabb93d8ca3e1c Mon Sep 17 00:00:00 2001 From: fsx950223 Date: Tue, 15 Sep 2026 09:11:59 +0000 Subject: [PATCH 46/66] perf(flydsl): optimize MTP2 and MTP3 decode --- .github/scripts/aiter_test.sh | 12 +- aiter/ops/flydsl/kernels/pa_decode_kernel.py | 6 +- aiter/ops/flydsl/pa_decode.py | 11 +- op_tests/test_flydsl_pa_decode.py | 1816 ++---------------- 4 files changed, 218 insertions(+), 1627 deletions(-) diff --git a/.github/scripts/aiter_test.sh b/.github/scripts/aiter_test.sh index e944789d2ed..e71e0f6954c 100755 --- a/.github/scripts/aiter_test.sh +++ b/.github/scripts/aiter_test.sh @@ -125,21 +125,13 @@ for file in "${sharded_files[@]}"; do test_cmd=(env AITER_MLA_DECODE_PERSISTENT_MAX_BATCH=0 timeout 60m python3 "$file") ;; op_tests/test_flydsl_pa_decode.py) - # The CLI sweep does not collect parametrized pytest regressions. + # 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}::test_recommended_splits_has_configurable_upper_clamp" \ - "${test_file}::test_pa_decode_benchmark_reference_causal_mtp" \ - "${test_file}::test_pa_decode_benchmark_cli_defaults_and_bs200" \ - "${test_file}::test_pa_decode_benchmark_entry" \ - "${test_file}::test_pa_decode_fixed_length_accuracy" \ - "${test_file}::test_pa_decode_variable_length_accuracy" \ - "${test_file}::test_large_negative_logits_preserve_online_softmax" \ - "${test_file}::test_additional_supported_head_dims_are_accurate" + python3 -m pytest -q "${test_file}::test_pa_decode" python3 "$test_file" ' _ "$file" diff --git a/aiter/ops/flydsl/kernels/pa_decode_kernel.py b/aiter/ops/flydsl/kernels/pa_decode_kernel.py index d8bbf983ac5..f8f2da6405b 100644 --- a/aiter/ops/flydsl/kernels/pa_decode_kernel.py +++ b/aiter/ops/flydsl/kernels/pa_decode_kernel.py @@ -146,15 +146,15 @@ def compile_pa_decode_tile( M_TILES = (CTA_ROWS + MFMA_MNK - 1) // MFMA_MNK ROWS_PADDED = M_TILES * MFMA_MNK # CDNA4 contracts four legacy fp8 K=32 groups in one K=128 instruction. - # Concatenating the existing packs preserves cache/LDS layouts for both - # fused and query-split MTP4 CTAs. + # Concatenating the existing packs preserves cache/LDS layouts for fused + # MTP3/MTP4 and query-split MTP4 CTAs. WIDE_FP8_MFMA = ( is_gfx950 and per_token_kv and query_dtype == "bf16" and head_dim == 128 and block_size in (16, 128) - and query_length == 4 + and query_length in (3, 4) and query_group_size == 16 ) # Small grids split query positions into separate CTAs; otherwise retain diff --git a/aiter/ops/flydsl/pa_decode.py b/aiter/ops/flydsl/pa_decode.py index cb821fcc94b..7449d48ef62 100644 --- a/aiter/ops/flydsl/pa_decode.py +++ b/aiter/ops/flydsl/pa_decode.py @@ -558,9 +558,10 @@ def normalize_scale(scale, name): num_cus = torch.cuda.get_device_properties(dev).multi_processor_count prefetch_v = num_seqs * num_kv_heads * num_partitions <= num_cus - # Split small MTP4 grids into one query per CTA, retaining the original + # Split small MTP2/MTP4 grids into one query per CTA, retaining the original # partial layout and partition count. Larger grids amortize KV loads by - # processing all four queries together. Keep QL1's selector above intact. + # processing all queries together. MTP3 stays fused because the kernel's + # supported split counts (1, 2, 4) must divide query_length. query_splits = 1 if ( arch == "gfx950" @@ -568,14 +569,14 @@ def normalize_scale(scale, name): and head_dim == 128 and query.dtype == torch.bfloat16 and num_kv_heads == 1 - and query_length == 4 + and query_length in (2, 4) and query_group_size == 16 and block_size in (16, 128) ): num_cus = torch.cuda.get_device_properties(dev).multi_processor_count base_workgroups = num_seqs * num_kv_heads * num_partitions - if 4 * base_workgroups <= 2 * num_cus: - query_splits = 4 + if query_length * base_workgroups <= 2 * num_cus: + query_splits = query_length prefetch_v = True with torch.cuda.device(dev): diff --git a/op_tests/test_flydsl_pa_decode.py b/op_tests/test_flydsl_pa_decode.py index 9d5ffa4826c..8c059963f2f 100644 --- a/op_tests/test_flydsl_pa_decode.py +++ b/op_tests/test_flydsl_pa_decode.py @@ -41,10 +41,8 @@ def _default_cuda_device(): SUPPORTED_GFX = ["gfx942", "gfx950"] KV_COMPUTE_BLOCK = 256 -# Fixed-length outputs use the tighter bound, while independently selected -# variable-length rows use the established varlen bound. +# All parametrized cases use equal context lengths and the strict bound. FIXED_LENGTH_ACCURACY_TOLERANCE = 5e-3 -VARIABLE_LENGTH_ACCURACY_TOLERANCE = 5e-2 # Pairwise coverage of the normal-accuracy axes in FlyDSL's PA regression test: # batches {3, 81, 128}, Q/KV heads {(4,1), (8,1), (16,1)}, head dims @@ -86,32 +84,6 @@ def _require_gpu(): pytest.skip(f"pa_decode is unsupported on {get_gfx_runtime()}") -@pytest.mark.parametrize( - "batch_size,max_partitions,expected", - [(1, 8, 8), (1, 256, 256), (16, 256, 32), (64, 256, 8)], -) -def test_recommended_splits_has_configurable_upper_clamp( - monkeypatch, batch_size, max_partitions, expected -): - """The default stays at eight while long-context callers may opt into 256.""" - if get_recommended_splits is None: - pytest.skip("FlyDSL is not available") - monkeypatch.setattr( - torch.cuda, - "get_device_properties", - lambda *_args, **_kwargs: type("Props", (), {"multi_processor_count": 256})(), - ) - assert ( - get_recommended_splits( - batch_size, - num_kv_heads=1, - split_kv_blocks=2, - max_partitions=max_partitions, - ) - == expected - ) - - def run_torch( query: torch.Tensor, key_cache: torch.Tensor, @@ -410,1601 +382,227 @@ def run_pa_decode_tile_case( return ret -def test_pa_decode_benchmark_reference_causal_mtp(): - """An analytic oracle catches flattened-row, KV-head, and causal mistakes.""" - if not torch.cuda.is_available(): - pytest.skip("ROCm is not available") - dtype = dtypes.bf16 - query_length, query_group_size, head_dim, block_size = 4, 2, 4, 16 - lengths, pages = [4, 2, 0], [1, 0, -1] - query = torch.zeros(12, 4, head_dim, dtype=dtype) - key = torch.zeros(2, 2, block_size, head_dim, dtype=dtypes.fp32) - value = ( - 100 * torch.arange(2).reshape(2, 1, 1, 1) - + 10 * torch.arange(2).reshape(1, 2, 1, 1) - + torch.arange(head_dim).reshape(1, 1, head_dim, 1) - + torch.arange(1, block_size + 1).reshape(1, 1, 1, block_size) - ).float() - value_scale = ( - torch.arange(1, block_size + 1, dtype=dtypes.fp32) - .reshape(1, 1, block_size, 1) - .expand(2, 2, block_size, 1) - .contiguous() - ) - key_scale = torch.ones_like(value_scale) - actual = run_torch( - query, - key, - value, - torch.tensor(pages, dtype=dtypes.i32).reshape(3, 1), - torch.tensor(lengths, dtype=dtypes.i32), - key_scale, - value_scale, - query_length=query_length, - ) - expected = [] - for length, page in zip(lengths, pages): - for position in range(query_length): - visible = max(0, length - query_length + 1 + position) - expected.append( - [ - [ - sum( - (100 * page + 10 * (head // query_group_size) + dim + t + 1) - * (t + 1) - for t in range(visible) - ) - / max(visible, 1) - for dim in range(head_dim) - ] - for head in range(4) - ] - ) - assert actual.dtype == dtype - assert torch.isfinite(actual).all() - torch.testing.assert_close( - actual, torch.tensor(expected, dtype=dtype), rtol=0, atol=0 - ) +PA_DECODE_TEST_CASES = [ + pytest.param( + { + "batch_size": 3, + "num_query_heads": 8, + "num_kv_heads": 1, + "head_dim": 128, + "context_length": 257, + "block_size": 16, + "dtype": dtypes.bf16, + "trans_v": False, + "per_token": False, + "query_length": 1, + "num_partitions": None, + }, + id="decode-auto-page16-bf16-scalar", + ), + pytest.param( + { + "batch_size": 3, + "num_query_heads": 8, + "num_kv_heads": 2, + "head_dim": 128, + "context_length": 1027, + "block_size": 64, + "dtype": dtypes.fp16, + "trans_v": True, + "per_token": False, + "query_length": 1, + "num_partitions": 1, + }, + id="decode-np1-page64-fp16-trans-v", + ), + pytest.param( + { + "batch_size": 3, + "num_query_heads": 8, + "num_kv_heads": 2, + "head_dim": 128, + "context_length": 257, + "block_size": 128, + "dtype": dtypes.bf16, + "trans_v": False, + "per_token": True, + "query_length": 1, + "num_partitions": 4, + }, + id="decode-reduce-page128-per-token", + ), + pytest.param( + { + "batch_size": 2, + "num_query_heads": 16, + "num_kv_heads": 1, + "head_dim": 128, + "context_length": 257, + "block_size": 128, + "dtype": dtypes.bf16, + "trans_v": True, + "per_token": True, + "query_length": 2, + "num_partitions": 3, + "gfx950_query_splits": 2, + }, + id="mtp2-query-split-page128-per-token", + ), + pytest.param( + { + "batch_size": 2, + "num_query_heads": 16, + "num_kv_heads": 1, + "head_dim": 128, + "context_length": 1027, + "block_size": 128, + "dtype": dtypes.bf16, + "trans_v": False, + "per_token": True, + "query_length": 3, + "num_partitions": 4, + "gfx950_query_splits": 1, + }, + id="mtp3-fused-page128-per-token", + ), + pytest.param( + { + "batch_size": 2, + "num_query_heads": 16, + "num_kv_heads": 1, + "head_dim": 128, + "context_length": 257, + "block_size": 16, + "dtype": dtypes.bf16, + "trans_v": True, + "per_token": True, + "query_length": 4, + "num_partitions": 3, + "gfx950_query_splits": 4, + }, + id="mtp4-query-split-page16-per-token", + ), + pytest.param( + { + "batch_size": 2, + "num_query_heads": 16, + "num_kv_heads": 1, + "head_dim": 128, + "context_length": 257, + "block_size": 128, + "dtype": dtypes.fp16, + "trans_v": True, + "per_token": False, + "query_length": 4, + "num_partitions": 3, + "gfx950_query_splits": 1, + }, + id="mtp4-reduce-page128-fp16-scalar", + ), + pytest.param( + { + "batch_size": 2, + "num_query_heads": 8, + "num_kv_heads": 1, + "head_dim": 256, + "context_length": 64, + "block_size": 16, + "dtype": dtypes.bf16, + "trans_v": False, + "per_token": False, + "query_length": 1, + "num_partitions": 1, + }, + id="decode-head256-np1", + ), + pytest.param( + { + "batch_size": 1, + "num_query_heads": 16, + "num_kv_heads": 1, + "head_dim": 128, + "context_length": 200000, + "block_size": 128, + "dtype": dtypes.bf16, + "trans_v": True, + "per_token": True, + "query_length": 1, + "num_partitions": 8, + "max_partitions": 4, + }, + id="decode-long-context-200k", + ), + pytest.param( + { + "batch_size": 200, + "num_query_heads": 16, + "num_kv_heads": 1, + "head_dim": 128, + "context_length": 257, + "block_size": 128, + "dtype": dtypes.bf16, + "trans_v": False, + "per_token": True, + "query_length": 4, + "num_partitions": None, + "gfx950_query_splits": 1, + }, + id="mtp4-bs200-auto-smoke", + ), +] -def test_pa_decode_benchmark_cli_defaults_and_bs200(): - defaults = _parse_args([]) - assert defaults.query_length == [1] - assert defaults.num_partitions == [None] - assert defaults.max_partitions == 8 - args = _parse_args( - [ - "-d", - "bf16", - "-b", - "200", - "-s", - "16,1,128,200000", - "-q", - "4", - "--block-size", - "16", - "128", - "--trans-v", - "0", - "1", - "--per-token", - "1", - "--num-partitions", - "3", - "5", - "8", - ] - ) - assert args.dtype == [dtypes.bf16] - assert args.batch == [200] - assert args.shapes == [(16, 1, 128, 200000)] - assert args.query_length == [4] - assert args.block_size == [16, 128] - assert args.trans_v == [0, 1] - assert args.per_token == [1] - assert args.num_partitions == [3, 5, 8] - boundaries = _parse_args( - ["-q", "1", "4", "--num-partitions", "1", "256", "--max-partitions", "256"] - ) - assert boundaries.query_length == [1, 4] - assert boundaries.num_partitions == [1, 256] - assert boundaries.max_partitions == 256 +@pytest.mark.parametrize("case", PA_DECODE_TEST_CASES) +def test_pa_decode(case, monkeypatch): + """Run one shared real-kernel/reference path across the supported PA axes.""" + _require_gpu() + case = case.copy() + expected_query_splits = case.pop("gfx950_query_splits", None) + selected_query_splits = [] + if expected_query_splits is not None: + module = importlib.import_module("aiter.ops.flydsl.pa_decode") + compile_tile = module.compile_pa_decode_tile -@pytest.mark.parametrize( - "batch_size,query_length,num_partitions,block_size,trans_v,dtype,per_token", - [ - pytest.param( - 2, 1, None, 16, False, dtypes.bf16, False, id="legacy-auto-decode" - ), - pytest.param(2, 4, 3, 16, True, dtypes.bf16, True, id="mtp4-page16-np3"), - pytest.param(2, 4, 3, 128, True, dtypes.fp16, False, id="mtp4-fp16-scalar"), - pytest.param(200, 4, 5, 128, False, dtypes.bf16, True, id="bs200-mtp4-np5"), - ], -) -def test_pa_decode_benchmark_entry( - monkeypatch, - batch_size, - query_length, - num_partitions, - block_size, - trans_v, - dtype, - per_token, -): - """Exercise the real public kernel once, without a 101-iteration benchmark.""" - _require_gpu() - context_length, num_query_heads, num_kv_heads, head_dim = 257, 16, 1, 128 - expected_partitions = 4 if num_partitions is None else num_partitions - recommendations, timed_calls = [], [] + def capture_compile(**kwargs): + selected_query_splits.append(kwargs["query_splits"]) + return compile_tile(**kwargs) - def recommend(*args, **kwargs): - assert num_partitions is None, "an explicit NP must bypass auto selection" - recommendations.append((args, kwargs)) - return expected_partitions + monkeypatch.setattr(module, "compile_pa_decode_tile", capture_compile) def run_once(fn, *args, **kwargs): assert fn is _run_flydsl assert not kwargs - ( - output, - query, - key_cache, - value_cache, - block_tables, - context_lengths, - key_scale, - value_scale, - partitions, - softmax_scale, - pmax, - psum, - pout, - ) = args - assert output.shape == query.shape == (batch_size * query_length, 16, 128) - assert output.dtype == query.dtype == pout.dtype == dtype - assert key_cache.dtype == value_cache.dtype == _quant_dtype() - assert value_cache.ndim == (5 if trans_v else 4) - assert partitions == expected_partitions - assert softmax_scale == head_dim**-0.5 - partial_shape = (batch_size, 1, partitions, query_length * 16) - assert pmax.shape == psum.shape == partial_shape - assert pmax.dtype == psum.dtype == dtypes.fp32 - assert pout.shape == (*partial_shape, 128) - blocks_per_sequence = (context_length + block_size - 1) // block_size - assert block_tables.shape == (batch_size, blocks_per_sequence) - assert context_lengths.shape == (batch_size,) - assert bool((context_lengths == context_length).all()) - scale_elements = ( - batch_size * blocks_per_sequence * block_size if per_token else 1 - ) - assert key_scale.numel() == value_scale.numel() == scale_elements - timed_calls.append(partitions) - return fn(*args), 100.0 + return fn(*args), 1.0 - monkeypatch.setitem(globals(), "get_recommended_splits", recommend) monkeypatch.setitem(globals(), "run_perftest", run_once) - overrides = {} - if query_length != 1: - overrides["query_length"] = query_length - if num_partitions is not None: - overrides["num_partitions"] = num_partitions - if num_partitions == 5: - # The exact override must also bypass a smaller automatic clamp. - overrides["max_partitions"] = 4 - result = run_pa_decode_tile_case( - batch_size, - num_query_heads, - num_kv_heads, - head_dim, - context_length, - block_size, - dtype, - trans_v, - per_token=per_token, - **overrides, - ) - assert timed_calls == [expected_partitions] - assert recommendations == ( - [ - ( - (batch_size, num_kv_heads), - { - "split_kv_blocks": KV_COMPUTE_BLOCK // block_size, - "max_partitions": 8, - }, - ) - ] - if num_partitions is None - else [] - ) - assert result["query_length"] == query_length - assert result["partitions"] == expected_partitions - assert result["flydsl err"] == 0 - assert result["flydsl us"] == 100.0 - assert result["flydsl us/token"] == pytest.approx( - 100.0 / (batch_size * query_length) - ) - attended_tokens = ( - query_length * context_length - query_length * (query_length - 1) // 2 - ) - flops = 4 * batch_size * num_query_heads * head_dim * attended_tokens - blocks_per_sequence = (context_length + block_size - 1) // block_size - # BF16/FP16 Q/O, FP8 KV, int32 metadata, FP32 scales. KV/scales are not - # multiplied by MTP length, and the padded part of page128 is excluded. - logical_bytes = ( - 2 * batch_size * query_length * num_query_heads * head_dim * 2 - + 2 * batch_size * num_kv_heads * context_length * head_dim - + batch_size * blocks_per_sequence * 4 - + batch_size * 4 - + (batch_size * num_kv_heads * context_length if per_token else 1) * 2 * 4 - ) - assert result["flydsl TFLOPS"] == pytest.approx(flops / 100.0 / 1e6) - assert result["flydsl TB/s"] == pytest.approx(logical_bytes / 100.0 / 1e6) - - -def _run_accuracy_case( - lengths, - *, - num_query_heads, - num_kv_heads, - head_dim, - block_size, - query_dtype, - num_partitions, - tolerance, - trans_v=False, -): - """Run PA against a dequantized-FP8 torch reference without benchmarking.""" - _require_gpu() - if num_query_heads % num_kv_heads: - raise ValueError("num_query_heads must be divisible by num_kv_heads") - - torch.manual_seed(0) - batch_size = len(lengths) - pages_per_sequence = [ - (context_length + block_size - 1) // block_size for context_length in lengths - ] - max_pages = max(pages_per_sequence) - num_blocks = sum(pages_per_sequence) - - query = torch.empty( - batch_size, num_query_heads, head_dim, dtype=query_dtype - ).uniform_(-0.5, 0.5) - key = torch.empty( - num_blocks, num_kv_heads, block_size, head_dim, dtype=query_dtype - ).uniform_(-0.5, 0.5) - value = torch.empty( - num_blocks, num_kv_heads, head_dim, block_size, dtype=query_dtype - ).uniform_(-0.5, 0.5) - - quant_dtype = _quant_dtype() - key_quant, key_scale = per_tensor_quant(key, quant_dtype=quant_dtype) - value_quant, value_scale = per_tensor_quant(value, quant_dtype=quant_dtype) - key_cache = ( - key_quant.view(num_blocks, num_kv_heads, block_size, head_dim // 16, 16) - .permute(0, 1, 3, 2, 4) - .contiguous() - ) - if trans_v: - value_cache = ( - value_quant.permute(0, 1, 3, 2) - .contiguous() - .view(num_blocks, num_kv_heads, block_size // 16, 16, head_dim) - .permute(0, 1, 2, 4, 3) - .contiguous() - ) - else: - value_cache = value_quant.contiguous() - - block_tables = torch.zeros( - batch_size, max_pages, dtype=dtypes.i32, device=query.device - ) - next_page = 0 - for seq_idx, page_count in enumerate(pages_per_sequence): - block_tables[seq_idx, :page_count] = torch.arange( - next_page, - next_page + page_count, - dtype=dtypes.i32, - device=query.device, - ) - next_page += page_count - context_lengths = torch.tensor(lengths, dtype=dtypes.i32, device=query.device) - - reference = run_torch( - query, - key_quant, - value_quant, - block_tables, - context_lengths, - key_scale, - value_scale, - ) - output = torch.empty_like(query) - torch.ops.aiter.pa_decode_flydsl( - output, - query, - key_cache, - value_cache, - context_lengths, - block_tables, - head_dim**-0.5, - 1, - num_partitions, - KV_COMPUTE_BLOCK, - quant_dtype, - None, - key_scale, - value_scale, - ) - - torch.testing.assert_close( - output.float(), - reference.float(), - atol=tolerance, - rtol=tolerance, - ) - - -@pytest.mark.parametrize( - "block_size,query_dtype,num_kv_heads,num_partitions", - [ - pytest.param(16, dtypes.bf16, 1, 1, id="page16-bf16-hkv1-np1"), - pytest.param(64, dtypes.fp16, 2, 1, id="page64-fp16-hkv2-np1"), - pytest.param(128, dtypes.bf16, 2, 4, id="page128-bf16-hkv2-np4"), - pytest.param(128, dtypes.fp16, 1, 4, id="page128-fp16-hkv1-np4"), - ], -) -def test_pa_decode_fixed_length_accuracy( - block_size, query_dtype, num_kv_heads, num_partitions -): - _run_accuracy_case( - [257, 257, 257], - num_query_heads=8, - num_kv_heads=num_kv_heads, - head_dim=128, - block_size=block_size, - query_dtype=query_dtype, - num_partitions=num_partitions, - tolerance=FIXED_LENGTH_ACCURACY_TOLERANCE, - ) - -@pytest.mark.parametrize("query_dtype", [dtypes.bf16, dtypes.fp16]) -def test_pa_decode_variable_length_accuracy(query_dtype): - _run_accuracy_case( - [1, 127, 257, 1027], - num_query_heads=8, - num_kv_heads=2, - head_dim=128, - block_size=128, - query_dtype=query_dtype, - num_partitions=4, - tolerance=VARIABLE_LENGTH_ACCURACY_TOLERANCE, - ) - - -def _adversarial_case( - head_dim=128, - context_length=257, - block_size=16, - query_group_size=8, - query_length=1, - per_token=False, - zero_query=False, - poison_padding_blocks=False, - tail_value_scale=None, - seed=0, - trans_v=False, - num_partitions=1, -): - """Build one case plus its torch reference over the same dequantised fp8 KV. - - ``poison_padding_blocks`` fills every block past the sequence's real extent - with NaN and pads the block table out to reach them -- the entries a caller - leaves behind, which the kernel must resolve to block 0 rather than follow. - """ - quant_dtype = _quant_dtype() - generator = torch.Generator(device="cuda").manual_seed(seed) - owned = (context_length + block_size - 1) // block_size - # A 256-token tile always walks 256/block_size pages, so pad past `owned` to - # give the block table entries that must not be dereferenced. - blocks = owned + 256 // block_size if poison_padding_blocks else owned - tokens = blocks * block_size - - def rand(*shape): - return torch.rand(*shape, generator=generator, device="cuda") - 0.5 - - query = ( - torch.zeros(query_length, query_group_size, head_dim, dtype=dtypes.bf16) - if zero_query - else rand(query_length, query_group_size, head_dim).to(dtypes.bf16) - ) - key = rand(blocks, 1, block_size, head_dim).to(quant_dtype) - value = rand(blocks, 1, block_size, head_dim).to(quant_dtype) - - invalid = (torch.arange(tokens, device="cuda") >= context_length).view( - blocks, 1, block_size, 1 - ) - if poison_padding_blocks: - padding = (torch.arange(tokens, device="cuda") >= owned * block_size).view( - blocks, 1, block_size, 1 - ) - poison = torch.full_like(value, float("nan"), dtype=dtypes.fp32).to(quant_dtype) - value = torch.where(padding.expand_as(value), poison, value) - key = torch.where(padding.expand_as(key), poison, key) - - if per_token: - key_scale = ( - torch.rand(blocks, 1, block_size, 1, generator=generator, device="cuda") - * 0.5 - + 0.75 - ) - value_scale = ( - torch.rand(blocks, 1, block_size, 1, generator=generator, device="cuda") - * 0.5 - + 0.75 - ) - if tail_value_scale is not None: - value_scale = torch.where( - invalid, torch.full_like(value_scale, tail_value_scale), value_scale - ) - else: - key_scale = torch.ones(1, dtype=dtypes.fp32) - value_scale = torch.ones(1, dtype=dtypes.fp32) - - key_cache = ( - key.view(blocks, 1, block_size, head_dim // 16, 16) - .permute(0, 1, 3, 2, 4) - .contiguous() - ) - if trans_v: - value_cache = ( - value.view(blocks, 1, block_size // 16, 16, head_dim) - .permute(0, 1, 2, 4, 3) - .contiguous() - ) - else: - value_cache = value.permute(0, 1, 3, 2).contiguous() - block_tables = torch.arange(blocks, dtype=dtypes.i32).reshape(1, blocks) - context_lengths = torch.full((1,), context_length, dtype=dtypes.i32) - - keys = key.float().reshape(tokens, head_dim) - values = value.float().reshape(tokens, head_dim) - if per_token: - keys = keys * key_scale.float().reshape(tokens, 1) - values = values * value_scale.float().reshape(tokens, 1) - reference = torch.empty( - query_length, query_group_size, head_dim, dtype=dtypes.fp32, device="cuda" - ) - for position in range(query_length): - # MTP position `p` sees context_length - (query_length - 1) + p tokens. - visible = context_length - (query_length - 1) + position - scores = query[position].float() @ keys[:visible].T * head_dim**-0.5 - reference[position] = torch.softmax(scores, dim=-1) @ values[:visible] - - output = torch.empty_like(query) - torch.ops.aiter.pa_decode_flydsl( - output, - query, - key_cache, - value_cache, - context_lengths, - block_tables, - head_dim**-0.5, - query_length, - num_partitions, - 256, - quant_dtype, - None, - key_scale, - value_scale, - ) - return output.float(), reference - - -def _assert_matches(output, reference, tolerance=0.06): - assert torch.isfinite(output).all(), ( - f"{int((~torch.isfinite(output)).sum())} of {output.numel()} outputs " - "are NaN or inf" - ) - scale = float(reference.abs().max()) or 1.0 - error = float((output - reference).abs().max()) / scale - assert error < tolerance, f"relative error {error:.3%} exceeds {tolerance:.1%}" - - -@pytest.mark.parametrize("query_length", [1, 4]) -@pytest.mark.parametrize("per_token", [False, True]) -def test_zero_query_stays_finite(query_length, per_token): - """An all-zero query row quantises to scale 0; -inf * 0 must not reach exp2.""" - _require_gpu() - output, reference = _adversarial_case( - query_length=query_length, per_token=per_token, zero_query=True - ) - _assert_matches(output, reference) - - -@pytest.mark.parametrize("query_dtype", [dtypes.bf16, dtypes.fp16]) -@pytest.mark.parametrize("per_token", [False, True]) -@pytest.mark.parametrize( - "query_group_size,query_length,block_size,context_length,num_partitions,trans_v,split_values", - [ - pytest.param(32, 1, 64, 128, 1, False, False, id="decode"), - pytest.param(16, 3, 128, 257, 4, True, False, id="mtp-partitions"), - pytest.param(8, 3, 128, 257, 1, True, False, id="padded-m-tile"), - pytest.param(32, 1, 64, 512, 1, False, True, id="history"), - ], -) -def test_large_negative_logits_preserve_online_softmax( - query_group_size, - query_length, - block_size, - context_length, - num_partitions, - trans_v, - split_values, - per_token, - query_dtype, -): - """Large negative logits must preserve both initial and accumulated state.""" - _require_gpu() - head_dim = 128 - blocks = (context_length + block_size - 1) // block_size - quant_dtype = _quant_dtype() - query = torch.full( - (query_length, query_group_size, head_dim), -3.0, dtype=query_dtype - ) - key = torch.full((blocks, 1, block_size, head_dim), 3.0).to(quant_dtype) - value = torch.ones(blocks, 1, block_size, head_dim) - if split_values: - # With equal logits, both KV tiles contribute equally. Resetting the - # correction to zero on every tile would incorrectly return 1, not 0.5. - value[: KV_COMPUTE_BLOCK // block_size] = 0 - value = value.to(quant_dtype) - key_cache = ( - key.view(blocks, 1, block_size, head_dim // 16, 16) - .permute(0, 1, 3, 2, 4) - .contiguous() - ) - if trans_v: - value_cache = ( - value.view(blocks, 1, block_size // 16, 16, head_dim) - .permute(0, 1, 2, 4, 3) - .contiguous() - ) - else: - value_cache = value.permute(0, 1, 3, 2).contiguous() - block_tables = torch.arange(blocks, dtype=dtypes.i32).reshape(1, blocks) - context_lengths = torch.full((1,), context_length, dtype=dtypes.i32) - scale = torch.ones( - (blocks, 1, block_size, 1) if per_token else (1,), dtype=dtypes.fp32 - ) - output = torch.full_like(query, float("nan")) - torch.ops.aiter.pa_decode_flydsl( - output, - query, - key_cache, - value_cache, - context_lengths, - block_tables, - head_dim**-0.5, - query_length, - num_partitions, - KV_COMPUTE_BLOCK, - quant_dtype, - None, - scale, - scale, - ) - reference = torch.full_like(output, 0.5 if split_values else 1.0, dtype=dtypes.fp32) - _assert_matches(output.float(), reference, tolerance=0.02) - - -@pytest.mark.parametrize("block_size", [16, 64, 128]) -@pytest.mark.parametrize("context_length", [1, 257, 1027]) -def test_block_table_padding_is_not_dereferenced(block_size, context_length): - """Entries past the sequence's extent must resolve to block 0. - - A 256-token tile always walks 256/block_size pages, so it reads block-table - entries the sequence does not own. Those hold whatever the caller left -- - often a stale id pointing at another sequence's live page. The tokens behind - them get a zero probability, but the PV matmul still multiplies their V - bytes (``0 * NaN == NaN``), so the page index itself has to be pinned. - - Not covered here, by design: the unwritten tail of the last page the - sequence does own. Those slots are inside a real block and the caller is - expected to leave them finite. - """ - _require_gpu() - output, reference = _adversarial_case( - context_length=context_length, - block_size=block_size, - poison_padding_blocks=True, - ) - _assert_matches(output, reference) - - -@pytest.mark.parametrize("query_length", [1, 4]) -@pytest.mark.parametrize("context_length", [1, 257]) -def test_extreme_tail_value_scale_is_ignored(context_length, query_length): - """A huge scale on an unwritten slot must not underflow valid probabilities. - - Only finite tails are covered: scale tensors are allocated zeroed, so a NaN - scale past ``context_len`` is not a case the kernel has to survive. - """ - _require_gpu() - if context_length < query_length: - pytest.skip("MTP position 0 would see a non-positive causal bound") - output, reference = _adversarial_case( - context_length=context_length, - query_length=query_length, - per_token=True, - tail_value_scale=1e9, - ) - _assert_matches(output, reference) - - -@pytest.mark.parametrize("num_partitions", [1, 4, 8]) -@pytest.mark.parametrize("trans_v", [False, True]) -@pytest.mark.parametrize("block_size", [16, 128]) -def test_per_token_scale_extremes_with_mixed_lengths( - block_size, trans_v, num_partitions -): - """Per-token normalization ignores tails and survives odd split starts. - - K is zero, so attention is uniform and the expected output is simply the - mean of the valid per-token V scales. Valid scales vary by 256-token tile - and 64-token wave, the first tile has an exactly-zero maximum, and unused - slots in the final physical page carry a hostile 1e9 scale. Lengths 769 - and 2305 put a one-token tail after three and nine full tiles respectively; - with four/eight partitions they also exercise odd global tile indices in - the ping-pong scale buffer. The empty row checks the neutral result. - """ - _require_gpu() - lengths = [0, 769, 2305] - batch_size, num_query_heads, num_kv_heads, head_dim = 3, 16, 1, 128 - pages_per_sequence = [(length + block_size - 1) // block_size for length in lengths] - max_pages = max(pages_per_sequence) - num_pages = sum(pages_per_sequence) - quant_dtype = _quant_dtype() - - query = torch.ones(batch_size, num_query_heads, head_dim, dtype=dtypes.bf16) - key_cache = torch.zeros( - num_pages, - num_kv_heads, - head_dim // 16, - block_size, - 16, - dtype=quant_dtype, - ) - value_plain = torch.ones( - num_pages, num_kv_heads, head_dim, block_size, dtype=quant_dtype - ) - if trans_v: - value_cache = ( - value_plain.view( - num_pages, - num_kv_heads, - head_dim, - block_size // 16, - 16, - ) - .permute(0, 1, 3, 2, 4) - .contiguous() + explicit_partitions = case["num_partitions"] + max_partitions = case.get("max_partitions", 8) + if explicit_partitions is None: + expected_partitions = get_recommended_splits( + case["batch_size"], + case["num_kv_heads"], + split_kv_blocks=KV_COMPUTE_BLOCK // case["block_size"], + max_partitions=max_partitions, ) else: - value_cache = value_plain + expected_partitions = explicit_partitions - block_tables = torch.zeros(batch_size, max_pages, dtype=dtypes.i32) - key_scale = torch.full( - (num_pages, num_kv_heads, block_size, 1), 1e9, dtype=dtypes.fp32 + result = run_pa_decode_tile_case(**case) + assert result["partitions"] == expected_partitions + assert result["flydsl err"] == 0 + assert result["flydsl us"] == 1.0 + assert result["flydsl us/token"] == pytest.approx( + 1.0 / (case["batch_size"] * case["query_length"]) ) - value_scale = torch.full_like(key_scale, 1e9) - expected_values = [] - next_page = 0 - for seq, (length, page_count) in enumerate(zip(lengths, pages_per_sequence)): - if length == 0: - expected_values.append(torch.zeros((), dtype=dtypes.fp32)) - continue - block_tables[seq, :page_count] = torch.arange( - next_page, next_page + page_count, dtype=dtypes.i32 - ) - capacity = page_count * block_size - token = torch.arange(capacity, dtype=dtypes.i32) - tile = token // KV_COMPUTE_BLOCK - wave = (token % KV_COMPUTE_BLOCK) // 64 - exponent = ((tile + wave) % 5 - 4).to(dtypes.fp32) - scales = torch.exp2(exponent) - scales = torch.where(token < KV_COMPUTE_BLOCK, 0.0, scales) - scales = torch.where(token < length, scales, 1e9) - page_scales = scales.reshape(page_count, block_size) - key_scale[next_page : next_page + page_count, 0, :, 0] = page_scales - value_scale[next_page : next_page + page_count, 0, :, 0] = page_scales - expected_values.append(scales[:length].mean()) - next_page += page_count + assert result["flydsl TFLOPS"] > 0 + assert result["flydsl TB/s"] > 0 - context_lengths = torch.tensor(lengths, dtype=dtypes.i32) - output = torch.full_like(query, float("nan")) - pa_decode( - output, - query, - key_cache, - value_cache, - context_lengths, - block_tables, - head_dim**-0.5, - 1, - num_partitions, - KV_COMPUTE_BLOCK, - quant_dtype, - None, - key_scale, - value_scale, - ) - - expected = ( - torch.stack(expected_values) - .reshape(batch_size, 1, 1) - .expand(batch_size, num_query_heads, head_dim) - ) - assert torch.isfinite(output).all() - torch.testing.assert_close(output.float(), expected, rtol=0.005, atol=0.005) - - -@pytest.mark.parametrize("trans_v", [False, True]) -@pytest.mark.parametrize("block_size", [16, 128]) -def test_per_token_sparse_pages_with_odd_partition_start(block_size, trans_v): - """Page and scale prefetches stay aligned across odd-starting partitions. - - Ten compute tiles split four ways start at 0, 3, 6 and 9. The partition - starting at 3 therefore both starts in the odd scale buffer and reuses it. - Nontrivial K/V and independently varying scales expose mistakes hidden by - uniform attention. Selected physical pages are shuffled and separated by - unused pages, so a logical-page/physical-page mixup cannot pass unnoticed. - """ - _require_gpu() - context_length, num_partitions = 2305, 4 - num_query_heads, head_dim = 16, 128 - logical_pages = (context_length + block_size - 1) // block_size - num_pages = 2 * logical_pages + 1 - quant_dtype = _quant_dtype() - generator = torch.Generator(device="cuda").manual_seed(23) - pages = 2 * torch.randperm(logical_pages, generator=generator) + 1 - token = torch.arange(logical_pages * block_size) - tile = token // KV_COMPUTE_BLOCK - wave = (token % KV_COMPUTE_BLOCK) // 64 - dim = torch.arange(head_dim) - signs = (1 - 2 * (dim % 2)).float() - query_factors = 0.5 + 0.25 * (torch.arange(num_query_heads) % 3).float() - query = (query_factors[:, None] * signs[None, :]).unsqueeze(0).to(dtypes.bf16) - - key = torch.rand(num_pages, 1, block_size, head_dim, generator=generator) * 0.25 - value = torch.rand(num_pages, 1, block_size, head_dim, generator=generator) - 0.5 - # Correlated logits and tile-dependent values make a wrong K scale affect - # the answer, while variation within pages exercises every token lane. - logical_key = (0.125 + 0.0625 * ((tile + wave) % 4))[:, None] * signs[None, :] - logical_key += 0.015625 * ((token[:, None] + dim[None, :]) % 3 - 1) - logical_value = ( - 0.25 * (tile % 3 - 1)[:, None] - + 0.0625 * (token // 16 % 3 - 1)[:, None] - + 0.03125 * (dim % 5 - 2)[None, :] - ) - key[pages] = logical_key.reshape(logical_pages, 1, block_size, head_dim) - value[pages] = logical_value.reshape(logical_pages, 1, block_size, head_dim) - key = key.to(quant_dtype) - value = value.to(quant_dtype) - key_scale = torch.ones(num_pages, 1, block_size, 1, dtype=dtypes.fp32) - value_scale = torch.ones_like(key_scale) - key_scale[pages, 0, :, 0] = torch.exp2( - ((2 * tile + wave + token % 3) % 4 - 2).float() - ).reshape(logical_pages, block_size) - value_scale[pages, 0, :, 0] = torch.exp2( - ((tile + 2 * wave + token % 5) % 4 - 3).float() - ).reshape(logical_pages, block_size) - - key_cache = ( - key.view(num_pages, 1, block_size, head_dim // 16, 16) - .permute(0, 1, 3, 2, 4) - .contiguous() - ) - value_plain = value.permute(0, 1, 3, 2).contiguous() - value_cache = ( - value.view(num_pages, 1, block_size // 16, 16, head_dim) - .permute(0, 1, 2, 4, 3) - .contiguous() - if trans_v - else value_plain - ) - block_tables = pages.to(dtypes.i32).unsqueeze(0) - context_lengths = torch.tensor([context_length], dtype=dtypes.i32) - reference = run_torch( - query.float(), - key, - value_plain, - block_tables, - context_lengths, - key_scale, - value_scale, - ) - output = torch.full_like(query, float("nan")) - torch.ops.aiter.pa_decode_flydsl( - output, - query, - key_cache, - value_cache, - context_lengths, - block_tables, - head_dim**-0.5, - 1, - num_partitions, - KV_COMPUTE_BLOCK, - quant_dtype, - None, - key_scale, - value_scale, - ) - assert torch.isfinite(output).all() - torch.testing.assert_close(output.float(), reference, rtol=0.005, atol=0.005) - - -def _run_mtp4_fused_reference_case( - lengths, num_kv_heads, block_size, trans_v, num_partitions -): - """Check MTP4 with sparse pages, varying scales, and a causal FP32 reference.""" - torch.manual_seed(37) - query_length, query_group_size, head_dim = 4, 16, 128 - batch_size = len(lengths) - num_query_heads = num_kv_heads * query_group_size - pages_per_sequence = [(length + block_size - 1) // block_size for length in lengths] - max_pages = max(pages_per_sequence) - num_pages = sum(pages_per_sequence) - quant_dtype = _quant_dtype() - - query = torch.empty( - batch_size * query_length, - num_query_heads, - head_dim, - dtype=dtypes.bf16, - ).uniform_(-0.5, 0.5) - key = torch.empty( - num_pages, num_kv_heads, block_size, head_dim, dtype=dtypes.bf16 - ).uniform_(-0.5, 0.5) - value = torch.empty( - num_pages, num_kv_heads, head_dim, block_size, dtype=dtypes.bf16 - ).uniform_(-0.5, 0.5) - key_quant, key_scale = pertoken_quant(key, quant_dtype=quant_dtype) - value_token_quant, value_scale = pertoken_quant( - value.permute(0, 1, 3, 2).contiguous(), quant_dtype=quant_dtype - ) - value_quant = value_token_quant.permute(0, 1, 3, 2).contiguous() - - token = torch.arange(num_pages * block_size).reshape(num_pages, 1, block_size, 1) - kv_head = torch.arange(num_kv_heads).reshape(1, num_kv_heads, 1, 1) - key_scale *= torch.exp2(((2 * token + kv_head) % 4 - 2).float()) - value_scale *= torch.exp2(((token + 2 * kv_head) % 5 - 3).float()) - - # Make C1--C4 analytically exact while still detecting causal off-by-one: - # K=0 gives uniform attention, V progresses as [.25, .5, .75, 1], and - # scale=1. A C4 row therefore produces [.25, .375, .5, .625]. - data_page = 0 - for length, page_count in zip(lengths, pages_per_sequence): - if 0 < length <= query_length: - capacity = page_count * block_size - token_values = ( - torch.arange(capacity, dtype=dtypes.fp32) % query_length + 1 - ) * 0.25 - key_quant[data_page : data_page + page_count].zero_() - value_quant[data_page : data_page + page_count] = ( - token_values.reshape(page_count, 1, 1, block_size) - .expand(page_count, num_kv_heads, head_dim, block_size) - .to(quant_dtype) - ) - key_scale[data_page : data_page + page_count].fill_(1.0) - value_scale[data_page : data_page + page_count].fill_(1.0) - data_page += page_count - - # Scatter logical pages into a permuted set of odd physical pages. Leave - # holes and page zero finite so masked/padded token reads remain valid. - selected_pages = 2 * torch.randperm(num_pages) + 1 - physical_page_count = 2 * num_pages + 1 - - def scatter_pages(tensor): - sparse = torch.zeros( - (physical_page_count, *tensor.shape[1:]), dtype=tensor.dtype - ) - sparse[selected_pages] = tensor - return sparse - - key_quant = scatter_pages(key_quant) - value_quant = scatter_pages(value_quant) - key_scale = scatter_pages(key_scale) - value_scale = scatter_pages(value_scale) - num_pages = physical_page_count - - key_cache = ( - key_quant.view( - num_pages, - num_kv_heads, - block_size, - head_dim // 16, - 16, - ) - .permute(0, 1, 3, 2, 4) - .contiguous() - ) - if trans_v: - value_cache = ( - value_quant.permute(0, 1, 3, 2) - .contiguous() - .view( - num_pages, - num_kv_heads, - block_size // 16, - 16, - head_dim, - ) - .permute(0, 1, 2, 4, 3) - .contiguous() - ) - else: - value_cache = value_quant - - block_tables = torch.zeros(batch_size, max_pages, dtype=dtypes.i32) - next_page = 0 - for seq, page_count in enumerate(pages_per_sequence): - block_tables[seq, :page_count] = selected_pages[ - next_page : next_page + page_count - ] - next_page += page_count - context_lengths = torch.tensor(lengths, dtype=dtypes.i32) - - reference = torch.zeros_like(query, dtype=dtypes.fp32) - for seq, length in enumerate(lengths): - if length == 0: - continue - token_ids = torch.arange(length) - logical_pages = token_ids // block_size - token_offsets = token_ids % block_size - physical_pages = block_tables[seq, logical_pages].long() - keys = key_quant[physical_pages, :, token_offsets, :].float() - values = value_quant[physical_pages, :, :, token_offsets].float() - keys *= key_scale[physical_pages, :, token_offsets, 0].float().unsqueeze(-1) - values *= value_scale[physical_pages, :, token_offsets, 0].float().unsqueeze(-1) - keys = keys.repeat_interleave(query_group_size, dim=1) - values = values.repeat_interleave(query_group_size, dim=1) - for position in range(query_length): - visible = max(0, length - (query_length - 1) + position) - if visible == 0: - continue - row = seq * query_length + position - scores = ( - torch.einsum("hd,khd->hk", query[row].float(), keys[:visible]) - * head_dim**-0.5 - ) - probs = torch.softmax(scores, dim=-1) - reference[row] = torch.einsum("hk,khd->hd", probs, values[:visible]) - - output = torch.full_like(query, float("nan")) - torch.ops.aiter.pa_decode_flydsl( - output, - query, - key_cache, - value_cache, - context_lengths, - block_tables, - head_dim**-0.5, - query_length, - num_partitions, - KV_COMPUTE_BLOCK, - quant_dtype, - None, - key_scale, - value_scale, - ) - assert torch.isfinite(output).all() - torch.testing.assert_close(output.float(), reference, rtol=0.005, atol=0.005) - - -@pytest.mark.parametrize("num_partitions", [1, 8]) -@pytest.mark.parametrize("trans_v", [False, True]) -@pytest.mark.parametrize("block_size", [16, 128]) -@pytest.mark.parametrize( - "lengths,num_kv_heads", - [ - pytest.param([4], 1, id="c4-hkv1"), - pytest.param([257], 1, id="c257-hkv1"), - pytest.param([259], 1, id="c259-hkv1"), - pytest.param([2305], 1, id="c2305-hkv1"), - pytest.param([2307], 1, id="c2307-hkv1"), - pytest.param( - [0, 1, 2, 3, 4, 255, 256, 257, 258, 259, 2303, 2304, 2305, 2306, 2307], - 1, - id="mixed-hkv1", - ), - pytest.param([257], 2, id="c257-hkv2"), - pytest.param([259], 2, id="c259-hkv2"), - ], -) -def test_mtp4_fused_phase_b_accuracy( - monkeypatch, lengths, num_kv_heads, block_size, trans_v, num_partitions -): - """Force the fused path and cover its MTP causal/partition boundaries.""" - _require_gpu() - pa_decode_module = importlib.import_module("aiter.ops.flydsl.pa_decode") - compile_tile = pa_decode_module.compile_pa_decode_tile - - def compile_fused_mtp4(**kwargs): - kwargs["query_splits"] = 1 - return compile_tile(**kwargs) - - monkeypatch.setattr(pa_decode_module, "compile_pa_decode_tile", compile_fused_mtp4) - _run_mtp4_fused_reference_case( - lengths, num_kv_heads, block_size, trans_v, num_partitions - ) - - -@pytest.mark.parametrize("trans_v", [False, True]) -@pytest.mark.parametrize("block_size", [16, 128]) -def test_mtp4_public_selector_fused_path(monkeypatch, block_size, trans_v): - """B32/NP8 selects the fused one-CTA-per-sequence MTP4 path.""" - _require_gpu() - pa_decode_module = importlib.import_module("aiter.ops.flydsl.pa_decode") - compile_tile = pa_decode_module.compile_pa_decode_tile - selected_splits = [] - - def capture_query_splits(**kwargs): - selected_splits.append(kwargs.get("query_splits", 1)) - return compile_tile(**kwargs) - - monkeypatch.setattr( - pa_decode_module, "compile_pa_decode_tile", capture_query_splits - ) - _run_mtp4_fused_reference_case([257, 259] * 16, 1, block_size, trans_v, 8) - assert selected_splits and set(selected_splits) == {1} - - -@pytest.mark.parametrize( - "head_dim,num_partitions", - [(64, 1), (256, 1), (1024, 1), (1024, 4)], -) -def test_additional_supported_head_dims_are_accurate(head_dim, num_partitions): - _require_gpu() - output, reference = _adversarial_case( - head_dim=head_dim, - context_length=64, - num_partitions=num_partitions, - ) - _assert_matches(output, reference) - - -def _constant_page_decode_case( - lengths, - block_size, - num_partitions, - block_table_width, - poison_padding=False, - per_token=False, - trans_v=False, -): - """Each sequence's V pages are constant, so its attention output is known.""" - _require_gpu() - batch_size, query_heads, head_dim = len(lengths), 8, 128 - pages_per_sequence = [(length + block_size - 1) // block_size for length in lengths] - page_values = [ - float(seq + 1) - for seq, pages in enumerate(pages_per_sequence) - for _ in range(pages) - ] - num_pages = len(page_values) - quant_dtype = _quant_dtype() - query = torch.ones(batch_size, query_heads, head_dim, dtype=dtypes.bf16) - key_cache = torch.zeros( - num_pages, 1, head_dim // 16, block_size, 16, dtype=quant_dtype - ) - value_cache = ( - torch.tensor(page_values, dtype=dtypes.fp32) - .reshape(num_pages, 1, 1, 1) - .expand(num_pages, 1, head_dim, block_size) - .contiguous() - .to(quant_dtype) - ) - if trans_v: - value_cache = ( - value_cache.view(num_pages, 1, head_dim, block_size // 16, 16) - .permute(0, 1, 3, 2, 4) - .contiguous() - ) - block_tables = torch.full( - (batch_size, block_table_width), - num_pages + 1024 if poison_padding else 0, - dtype=dtypes.i32, - ) - next_page = 0 - for seq, pages in enumerate(pages_per_sequence): - block_tables[seq, :pages] = torch.arange( - next_page, next_page + pages, dtype=dtypes.i32 - ) - next_page += pages - context_lengths = torch.tensor(lengths, dtype=dtypes.i32) - scale = torch.ones( - (num_pages, 1, block_size, 1) if per_token else (1,), dtype=dtypes.fp32 - ) - output = torch.full_like(query, float("nan")) - torch.ops.aiter.pa_decode_flydsl( - output, - query, - key_cache, - value_cache, - context_lengths, - block_tables, - head_dim**-0.5, - 1, - num_partitions, - 256, - quant_dtype, - None, - scale, - scale, - ) - # The reference uses the values actually representable in the FP8 cache. - expected = ( - torch.tensor( - [float(seq + 1) if length else 0.0 for seq, length in enumerate(lengths)], - dtype=dtypes.fp32, - ) - .to(quant_dtype) - .to(dtypes.fp32) - .reshape(batch_size, 1, 1) - .expand_as(output) - ) - _assert_matches(output.float(), expected, tolerance=0.02) - - -@pytest.mark.parametrize("num_partitions", [1, 4]) -@pytest.mark.parametrize("block_table_width", [1, 2, 3, 4, 5]) -def test_block_table_row_alignment(block_table_width, num_partitions): - """A contiguous block table need not have four-entry-aligned row starts.""" - _constant_page_decode_case([16, 16, 16], 16, num_partitions, block_table_width) - - -@pytest.mark.parametrize("num_partitions", [1, 4]) -@pytest.mark.parametrize("block_size", [16, 64, 128]) -@pytest.mark.parametrize( - "lengths,per_token", - [ - pytest.param([64, 0], False, id="mixed"), - pytest.param([64, 0], True, id="mixed-per-token"), - pytest.param([0, 0], False, id="all-empty"), - ], -) -def test_empty_contexts_do_not_read_kv(lengths, per_token, block_size, num_partitions): - """Empty rows must not follow padding or read an empty KV allocation.""" - width = 2 * KV_COMPUTE_BLOCK // block_size if any(lengths) else 0 - _constant_page_decode_case( - lengths, - block_size, - num_partitions, - width, - poison_padding=True, - per_token=per_token, - ) - - -@pytest.mark.parametrize("context_length", [63, 64, 65, 127, 128, 129, 257]) -@pytest.mark.parametrize("trans_v", [False, True]) -@pytest.mark.parametrize("per_token", [False, True]) -def test_page128_chunk_boundaries(context_length, trans_v, per_token): - """Two waves share a page-128, but must load different K/V/scales halves.""" - _require_gpu() - output, reference = _adversarial_case( - block_size=128, - context_length=context_length, - query_group_size=16, - trans_v=trans_v, - per_token=per_token, - num_partitions=4, - ) - _assert_matches(output, reference) - - -@pytest.mark.parametrize("block_size", [16, 128]) -@pytest.mark.parametrize("context_length", [2048, 100000, 200000]) -@pytest.mark.parametrize("query_length", [1, 3]) -@pytest.mark.parametrize("per_token", [False, True]) -def test_tp4_contexts(block_size, context_length, query_length, per_token): - """64 Q / 4 KV heads under TP4: Hq16, Hkv1, D128, including dense MTP. - - Page size decides how many block-table entries a 256-token tile walks -- - two at page-128, sixteen at page-16 -- so both served page sizes are run - against every context. - """ - _require_gpu() - output, reference = _adversarial_case( - block_size=block_size, - query_group_size=16, - context_length=context_length, - query_length=query_length, - trans_v=True, - per_token=per_token, - num_partitions=8, - ) - _assert_matches(output, reference) - - -@pytest.mark.parametrize("batch_size", [1, 2, 4, 8, 16, 32, 64]) -@pytest.mark.parametrize("block_size", [16, 128]) -@pytest.mark.parametrize("context_length", [2048, 100000, 200000]) -def test_tp4_decode_batches(batch_size, block_size, context_length): - """Served TP4 decode grid: Hq16, Hkv1, D128, transposed V. - - Batch is the only axis that moves the CTA count, and at NP=8 the sweep - straddles the tuned V-prefetch window (256 < batch * 8 <= 512 on a 256-CU - part), so page-128 batch 64 compiles the prefetch specialization while the - smaller batches compile the default schedule. Page 16 never enters that - window -- the gate also requires block_size == 128 -- so it covers the - sixteen-pages-per-tile block-table walk over the same grids instead. - """ - _run_accuracy_case( - [context_length] * batch_size, - num_query_heads=16, - num_kv_heads=1, - head_dim=128, - block_size=block_size, - query_dtype=dtypes.bf16, - num_partitions=8, - tolerance=FIXED_LENGTH_ACCURACY_TOLERANCE, - trans_v=True, - ) - - -@pytest.mark.parametrize("trans_v", [False, True]) -def test_page128_cache_above_2gib(trans_v): - """Large batches at 200k context cross the signed-i32 cache offset limit.""" - _require_gpu() - block_size, head_dim = 128, 128 - boundary_page = 2**31 // (block_size * head_dim) - pages = boundary_page + 2 - quant_dtype = _quant_dtype() - key_cache = torch.empty(pages, 1, 8, block_size, 16, dtype=quant_dtype) - value_shape = ( - (pages, 1, 8, head_dim, 16) if trans_v else (pages, 1, head_dim, block_size) - ) - value_cache = torch.empty(value_shape, dtype=quant_dtype) - key_cache[0].zero_() - value_cache[0].zero_() - page_ids = [boundary_page - 1, boundary_page, boundary_page + 1] - for index, page in enumerate(page_ids): - key_cache[page].zero_() - value_cache[page].fill_(index + 1) - query = torch.ones(1, 16, head_dim, dtype=dtypes.bf16) - output = torch.empty_like(query) - scale = torch.ones(1, dtype=dtypes.fp32) - torch.ops.aiter.pa_decode_flydsl( - output, - query, - key_cache, - value_cache, - torch.tensor([257], dtype=dtypes.i32), - torch.tensor([page_ids], dtype=dtypes.i32), - head_dim**-0.5, - 1, - 8, - 256, - quant_dtype, - None, - scale, - scale, - ) - expected = torch.full_like(output, (128 + 2 * 128 + 3) / 257, dtype=dtypes.fp32) - _assert_matches(output.float(), expected, tolerance=0.02) - - -@pytest.mark.parametrize( - "block_size,per_token,packed", - [ - (16, False, False), - (16, True, False), - (16, False, True), - (128, False, False), - (128, True, False), - ], -) -def test_prepared_sparse_page_tables(block_size, per_token, packed): - """PA consumes selected pages; each MTP query already has its own table. - - Select 2048 tokens from a 200k-token cache. In the page-16 adapter each - selected logical page-128 expands to eight physical pages. Packed K/V - sides share a block and V has a different base pointer and cache extent. - """ - _require_gpu() - torch.manual_seed(7) - rows, query_heads, head_dim = 6, 16, 128 # two requests with three MTP queries - logical_blocks = (200000 + 127) // 128 - pages_per_logical = 128 // block_size - page_stride = pages_per_logical * (2 if packed else 1) - num_pages = logical_blocks * page_stride - quant_dtype = _quant_dtype() - key = (torch.rand(num_pages, 1, block_size, head_dim) - 0.5).to(quant_dtype) - value = (torch.rand_like(key, dtype=dtypes.fp32) - 0.5).to(quant_dtype) - key_cache = ( - key.view(num_pages, 1, block_size, head_dim // 16, 16) - .permute(0, 1, 3, 2, 4) - .contiguous() - ) - value_cache = ( - value.view(num_pages, 1, block_size // 16, 16, head_dim) - .permute(0, 1, 2, 4, 3) - .contiguous() - ) - if packed: - # Mirror the adapter's contiguous, overlapping views into interleaved - # K/V storage. Only K-side page ids are put into the block table. - value_cache = key_cache.flatten()[ - pages_per_logical * block_size * head_dim : - ].view(num_pages - pages_per_logical, 1, block_size // 16, head_dim, 16) - value = value_cache.permute(0, 1, 2, 4, 3).reshape( - num_pages - pages_per_logical, 1, block_size, head_dim - ) - - query = (torch.rand(rows, query_heads, head_dim) - 0.5).to(dtypes.bf16) - lengths = [2046, 2047, 2048] * 2 - # Spare entry deliberately leaves page-16 rows non-vector-aligned. - block_tables = torch.full( - (rows, 16 * pages_per_logical + 1), num_pages + 1024, dtype=dtypes.i32 - ) - if per_token: - key_scale = 0.5 + torch.rand(num_pages, 1, block_size, 1) - value_scale = 0.5 + torch.rand_like(key_scale) - else: - key_scale = torch.ones(1, dtype=dtypes.fp32) - value_scale = torch.ones(1, dtype=dtypes.fp32) - reference = torch.empty_like(query, dtype=dtypes.fp32) - for row, length in enumerate(lengths): - selected = torch.randperm(logical_blocks - 1)[:15] - # Partial current block comes last; earlier selections differ per query. - selected = torch.cat((selected, selected.new_tensor([logical_blocks - 1]))) - pages = ( - selected[:, None] * page_stride + torch.arange(pages_per_logical)[None, :] - ).flatten() - block_tables[row, : pages.numel()] = pages.to(dtypes.i32) - keys = key[pages].float() - values = value[pages].float() - if per_token: - keys = keys * key_scale[pages] - values = values * value_scale[pages] - keys = keys.reshape(-1, head_dim)[:length] - values = values.reshape(-1, head_dim)[:length] - scores = query[row].float() @ keys.T * head_dim**-0.5 - reference[row] = torch.softmax(scores, dim=-1) @ values - output = torch.empty_like(query) - torch.ops.aiter.pa_decode_flydsl( - output, - query, - key_cache, - value_cache, - torch.tensor(lengths, dtype=dtypes.i32), - block_tables, - head_dim**-0.5, - 1, - 8, - 256, - quant_dtype, - None, - key_scale, - value_scale, - sliding_window=-1, - ) - _assert_matches(output.float(), reference) - - -@pytest.mark.parametrize("context_length", [3, 257, 2051]) -@pytest.mark.parametrize("num_partitions", [1, 4]) -@pytest.mark.parametrize("trans_v", [False, True]) -@pytest.mark.parametrize("zero_query", [False, True]) -def test_mtp3_page128_partition_boundaries( - context_length, num_partitions, trans_v, zero_query -): - """MTP rows retain independent softmax state across partial/empty partitions.""" - _require_gpu() - output, reference = _adversarial_case( - head_dim=128, - context_length=context_length, - block_size=128, - query_group_size=16, - query_length=3, - num_partitions=num_partitions, - trans_v=trans_v, - zero_query=zero_query, - ) - _assert_matches(output, reference) - - -@pytest.mark.parametrize("all_empty", [False, True]) -def test_page128_decode_prefetch_mixed_contexts(all_empty): - """A two-CTA-per-CU decode grid handles empty partitions and poisoned padding.""" - _require_gpu() - lengths = [0] * 64 if all_empty else [0, 1, 63, 127, 128, 129, 257, 2051] * 8 - _constant_page_decode_case( - lengths, - block_size=128, - num_partitions=8, - block_table_width=0 if all_empty else 17, - poison_padding=True, - trans_v=True, - ) - - -@pytest.mark.parametrize( - "query_group_size,query_length", [(8, 3), (16, 2), (16, 4), (16, 5)] -) -@pytest.mark.parametrize("context_length", [129, 1027]) -def test_page128_reduction_rows_across_m_tiles( - query_group_size, query_length, context_length -): - """Aligned scratch rows stay independent across padded M-tiles and partitions.""" - _require_gpu() - output, reference = _adversarial_case( - head_dim=128, - context_length=context_length, - block_size=128, - query_group_size=query_group_size, - query_length=query_length, - num_partitions=4, - trans_v=True, - ) - _assert_matches(output, reference) - - -@pytest.mark.parametrize("query_splits", [2, 4]) -@pytest.mark.parametrize("num_partitions", [1, 8]) -@pytest.mark.parametrize("block_size", [16, 128]) -@pytest.mark.parametrize("trans_v", [False, True]) -@pytest.mark.parametrize("context_length", [4, 257, 2305]) -def test_per_token_mtp_query_splits( - monkeypatch, query_splits, num_partitions, block_size, trans_v, context_length -): - """Query-split CTAs preserve MTP causality and the unsplit partial layout.""" - _require_gpu() - pa_decode_module = importlib.import_module("aiter.ops.flydsl.pa_decode") - compile_tile = pa_decode_module.compile_pa_decode_tile - - def compile_split_tile(**kwargs): - kwargs["query_splits"] = query_splits - kwargs["prefetch_v"] = query_splits == 4 - return compile_tile(**kwargs) - - monkeypatch.setattr(pa_decode_module, "compile_pa_decode_tile", compile_split_tile) - if context_length == 4: - # An exact short-prefix case isolates the causal/query mapping. With - # random V scales, even the unsplit FP8 kernel has >0.005 absolute - # quantization error at four tokens; do not weaken the assertion. - quant_dtype = _quant_dtype() - query = torch.ones(4, 16, 128, dtype=dtypes.bf16) - key_cache = torch.zeros(1, 1, 8, block_size, 16, dtype=quant_dtype) - value = torch.zeros(1, 1, block_size, 128, dtype=dtypes.fp32) - value[0, 0, :4] = ( - (torch.arange(4).float() + 1)[:, None] - * 0.125 - * (1 + 0.25 * (torch.arange(128) % 4).float())[None, :] - ) - value = value.to(quant_dtype) - value_cache = ( - value.view(1, 1, block_size // 16, 16, 128) - .permute(0, 1, 2, 4, 3) - .contiguous() - if trans_v - else value.permute(0, 1, 3, 2).contiguous() - ) - key_scale = torch.ones(1, 1, block_size, 1, dtype=dtypes.fp32) - value_scale = torch.ones_like(key_scale) - value_scale[0, 0, :4, 0] = torch.tensor([0.5, 1.0, 0.25, 1.0]) - prefix_values = (value[0, 0, :4].float() * value_scale[0, 0, :4]).cumsum(0) - reference = (prefix_values / (torch.arange(4).float() + 1)[:, None])[ - :, None, : - ].expand(4, 16, 128) - output = torch.full_like(query, float("nan")) - torch.ops.aiter.pa_decode_flydsl( - output, - query, - key_cache, - value_cache, - torch.tensor([4], dtype=dtypes.i32), - torch.zeros(1, 1, dtype=dtypes.i32), - 128**-0.5, - 4, - num_partitions, - KV_COMPUTE_BLOCK, - quant_dtype, - None, - key_scale, - value_scale, - ) - output = output.float() - else: - output, reference = _adversarial_case( - head_dim=128, - query_group_size=16, - query_length=4, - context_length=context_length, - block_size=block_size, - trans_v=trans_v, - per_token=True, - num_partitions=num_partitions, - ) - assert torch.isfinite(output).all() - torch.testing.assert_close(output, reference, rtol=0.005, atol=0.005) - - -@pytest.mark.parametrize("num_partitions", [1, 8]) -@pytest.mark.parametrize("block_size", [16, 128]) -@pytest.mark.parametrize("trans_v", [False, True]) -def test_per_token_mtp_query_split_mixed_contexts(num_partitions, block_size, trans_v): - """Automatic MTP splitting keeps empty and nonempty sequences independent.""" - _require_gpu() - lengths = [0, 4, 2305] - batch_size, query_length, query_heads, head_dim = 3, 4, 16, 128 - pages_per_sequence = [(length + block_size - 1) // block_size for length in lengths] - page_values = [ - (seq + 1) * 0.125 - for seq, count in enumerate(pages_per_sequence) - for _ in range(count) - ] - num_pages = len(page_values) - quant_dtype = _quant_dtype() - query = torch.ones( - batch_size * query_length, query_heads, head_dim, dtype=dtypes.bf16 - ) - key_cache = torch.zeros( - num_pages, 1, head_dim // 16, block_size, 16, dtype=quant_dtype - ) - value_plain = ( - torch.tensor(page_values, dtype=dtypes.fp32) - .reshape(num_pages, 1, 1, 1) - .expand(num_pages, 1, head_dim, block_size) - .contiguous() - .to(quant_dtype) - ) - value_cache = ( - value_plain.view(num_pages, 1, head_dim, block_size // 16, 16) - .permute(0, 1, 3, 2, 4) - .contiguous() - if trans_v - else value_plain - ) - block_tables = torch.full( - (batch_size, max(pages_per_sequence)), num_pages + 1024, dtype=dtypes.i32 - ) - next_page = 0 - for seq, count in enumerate(pages_per_sequence): - block_tables[seq, :count] = torch.arange(next_page, next_page + count) - next_page += count - scales = torch.ones(num_pages, 1, block_size, 1, dtype=dtypes.fp32) - output = torch.full_like(query, float("nan")) - torch.ops.aiter.pa_decode_flydsl( - output, - query, - key_cache, - value_cache, - torch.tensor(lengths, dtype=dtypes.i32), - block_tables, - head_dim**-0.5, - query_length, - num_partitions, - KV_COMPUTE_BLOCK, - quant_dtype, - None, - scales, - scales, - ) - expected = ( - torch.tensor([0.0, 0.25, 0.375], dtype=dtypes.fp32) - .repeat_interleave(query_length) - .reshape(batch_size * query_length, 1, 1) - .expand_as(output) - ) - assert torch.isfinite(output).all() - torch.testing.assert_close(output.float(), expected, rtol=0.005, atol=0.005) + if expected_query_splits is not None: + expected = expected_query_splits if get_gfx_runtime() == "gfx950" else 1 + assert selected_query_splits + assert set(selected_query_splits) == {expected} def _positive_int(value): @@ -2043,8 +641,8 @@ def _parse_args(argv=None): type=_positive_int, nargs="+", default=[1], - help="""Query tokens per sequence: 1 for decode, 4 for MTP4. - e.g.: -q 1 4. MTP positions use dense causal masking.""", + help="""Query tokens per sequence: 1 for decode, >1 for MTP. + e.g.: -q 1 2 3 4. MTP positions use dense causal masking.""", ) parser.add_argument( "-s", From cecd335940ff298d969a02f12dc5f67dd748f706 Mon Sep 17 00:00:00 2001 From: Elton Date: Wed, 16 Sep 2026 15:34:06 +0800 Subject: [PATCH 47/66] [FlyDSL] perf(pa_decode): add dynamic KV work planning and tune partition scheduling (#5546) * perf(flydsl): balance paged-attention KV work with GPU plans * perf(flydsl): skip unused planned PA work * test(flydsl): consolidate paged attention coverage * test(flydsl): remove standalone PA plan benchmark --------- Co-authored-by: Elton <211304267+zhiding512@users.noreply.github.com> Co-authored-by: fsx950223 --- .github/scripts/aiter_test.sh | 2 +- aiter/ops/flydsl/kernels/pa_decode_kernel.py | 186 ++++-- aiter/ops/flydsl/kernels/pa_decode_plan.py | 166 +++++ aiter/ops/flydsl/kernels/pa_decode_reduce.py | 338 +++++++--- aiter/ops/flydsl/pa_decode.py | 92 ++- docs/flydsl_pa_decode_plan.md | 92 +++ op_tests/test_flydsl_pa_decode.py | 622 ++++++++++++++++++- 7 files changed, 1341 insertions(+), 157 deletions(-) create mode 100644 aiter/ops/flydsl/kernels/pa_decode_plan.py create mode 100644 docs/flydsl_pa_decode_plan.md diff --git a/.github/scripts/aiter_test.sh b/.github/scripts/aiter_test.sh index e71e0f6954c..330c789d69b 100755 --- a/.github/scripts/aiter_test.sh +++ b/.github/scripts/aiter_test.sh @@ -131,7 +131,7 @@ for file in "${sharded_files[@]}"; do bash -c ' set -euo pipefail test_file=$1 - python3 -m pytest -q "${test_file}::test_pa_decode" + python3 -m pytest -q "$test_file" python3 "$test_file" ' _ "$file" diff --git a/aiter/ops/flydsl/kernels/pa_decode_kernel.py b/aiter/ops/flydsl/kernels/pa_decode_kernel.py index f8f2da6405b..056a44c4ed3 100644 --- a/aiter/ops/flydsl/kernels/pa_decode_kernel.py +++ b/aiter/ops/flydsl/kernels/pa_decode_kernel.py @@ -74,6 +74,7 @@ def compile_pa_decode_tile( wide_kv_addressing: bool = False, prefetch_v: bool = False, query_splits: int = 1, + use_work_plan: bool = False, ): """Build the tile-programming PA-decode kernel + launch wrapper. @@ -160,6 +161,10 @@ def compile_pa_decode_tile( # Small grids split query positions into separate CTAs; otherwise retain # all four queries here and share KV loads, scales, and P publication. MTP4_FUSED = WIDE_FP8_MFMA and M_TILES == 4 + # With 64-bit cache addresses, page-128 V carried through QK increases + # register pressure: load plain V during P packing and transposed V after + # P publication. Page-16 gathers and narrow addresses retain prefetching. + MTP4_PREFETCH_V = MTP4_FUSED and (block_size == 16 or not wide_kv_addressing) PER_TOKEN_M1 = ( per_token_kv and is_gfx950 @@ -284,8 +289,8 @@ def compile_pa_decode_tile( class SharedStorage: buf: fx.Array[fx.Int32, total_bytes // 4, 16] - @flyc.kernel(known_block_size=(BLOCK_THREADS, 1, 1)) - def pa_decode_tile_kernel( + @flyc.jit + def _pa_decode_tile_task( output_ptr: fx.Pointer, # [num_seqs*query_length, num_q_heads, head_dim] (written directly when NP==1) # per-partition partial outputs (combined by the reduce kernel when NP>1): pmax_ptr: fx.Pointer, # [num_seqs, num_kv_heads, num_partitions, query_length*query_group_size] row max @@ -305,6 +310,11 @@ def pa_decode_tile_kernel( 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 @@ -315,6 +325,12 @@ def pa_decode_tile_kernel( 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): + 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) @@ -399,28 +415,27 @@ def _load_q_row(qi, gs_head): gs_head = flat_idx - qi * query_group_size q_units_prefetched.append(_load_q_row(qi, gs_head)) - 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) - - def _ctxlen_load(elem_idx): - fx.copy(ctx_copy_atom, fx.slice(ctx_tiled, (None, elem_idx)), ctx_reg) - return fx.Vector(fx.memref_load_vec(ctx_reg)) - 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) - context_len = fx.Int32(_ctxlen_load(seq)[0]) + 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]) # Bound block_tables to its real extent: the last (partial) 256-token tile # can index a page past ceil(context/block_size); the bounded resource # returns page 0 for that out-of-range read instead of faulting (those # tail tokens are masked out anyway). bt_num_records_bytes = ( - fx.Int64(gpu.grid_dim.x) * fx.Int64(max_blocks_per_seq) * 4 + 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( @@ -438,14 +453,18 @@ def _v_load16(byte_off): key_scale = fx.Float32(key_scale_buf[0]) value_scale = fx.Float32(value_scale_buf[0]) - num_tiles = (context_len + TILE_TOK - 1) // TILE_TOK num_pages = ( context_len + block_size - 1 ) // block_size # pages this sequence really owns - 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) + 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) # One i8 blob carved into typed byte-offset pointers. `lds_base` is an # ir.Value pointer (safe inside scf control flow); the Python `lds` @@ -951,7 +970,7 @@ def _l_slot(m): # Fused pipeline: carry next V in registers, reusing each chunk only # after its final query has consumed the current tile's values. V_DATA_SLOT = 2 + STATE_PER_M * M_TILES - if const_expr(MTP4_FUSED): + if const_expr(MTP4_PREFETCH_V): v_pf0 = fx.Vector.filled(VHE_CHUNKS * NVOPS, 0, fx.Int64) if part_start < part_end: v_flat0 = [] @@ -1008,13 +1027,16 @@ def _mask_v_scale(vec, a, thr=ctx_thr, zero=zero4_scale): # V is independent of QK/softmax. Load it early for the tuned # decode path and reuse it across M-tiles in the phase-split path. v_vh_shared = None - if const_expr(MTP4_FUSED): + 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 or (EARLY_V and not SCALES_BEFORE_CURRENT_V)): + elif const_expr( + (M_TILES > 1 and not MTP4_FUSED) + or (EARLY_V and not SCALES_BEFORE_CURRENT_V) + ): v_vh_shared = [ _v_ops(v_page_cur, vh) for vh in range_constexpr(VHE_CHUNKS) ] @@ -1189,6 +1211,11 @@ def _lmax_off_m(m): ] 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 independent of the query row. # Compute it once instead of once per M tile. v_max_global = _ld_lw_row(sVScaleMax_off, 0).reduce(ReductionOp.MAX) @@ -1266,6 +1293,12 @@ def _lmax_off_m(m): gpu.barrier() + if const_expr(not MTP4_PREFETCH_V and trans_v): + # Contiguous transposed V can wait until saved scores + # are dead, avoiding overlap with their live registers. + 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 @@ -1298,7 +1331,7 @@ def _lmax_off_m(m): [v_max_scaled], dtype=fx.Float32 ).broadcast_to(OP_ELEMS) o_acc[vh] = o_acc[vh] * corr_b + op - if const_expr(m == M_TILES - 1): + if const_expr(MTP4_PREFETCH_V and m == M_TILES - 1): # The current chunk is dead; reuse its registers # for the next tile while the final PV finishes. v_next_chunk = fx.Vector.filled(NVOPS, 0, fx.Int64) @@ -1310,14 +1343,15 @@ def _lmax_off_m(m): next_state.extend([*o_acc, m_new, l_new]) if const_expr(m < M_TILES - 1): fx.rocdl.sched_barrier(0) - 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, - ) + if const_expr(MTP4_PREFETCH_V): + 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 @@ -1711,7 +1745,7 @@ def _lmax_off_m(m): ).broadcast_to(OP_ELEMS) o_acc[vh] = o_acc[vh] * corr_b + op next_state.extend([*o_acc, m_new, l_new]) - if const_expr(MTP4_FUSED): + if const_expr(MTP4_PREFETCH_V): next_state.append(v_next) results = yield next_state o_final = results @@ -1741,7 +1775,7 @@ def _lmax_off_m(m): qh = kv_h * query_group_size + gs_head_e def _emit(o_norm, sub, query_idx, query_head): - if const_expr(NP == 1): + 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 @@ -1749,9 +1783,7 @@ def _emit(o_norm, sub, query_idx, query_head): ) fx.ptr_store(o_norm, fx.add_offset(output, out_offset)) else: - base = ( - (seq * n_kv + kv_h) * NP + part - ) * TOTAL_ROWS + global_row # noqa: B023 + base = partial_slot * TOTAL_ROWS + global_row # noqa: B023 pout_offset = base * head_dim + sub * OP_ELEMS fx.ptr_store( o_norm, @@ -1769,9 +1801,9 @@ def _emit(o_norm, sub, query_idx, query_head): if row < CTA_ROWS: _emit(o_norm, sub, qi_e, qh) - if const_expr(NP > 1): # noqa: SIM102 + if const_expr(NP > 1 or use_work_plan): # noqa: SIM102 if warp == 0 and rgroup == 0: - base = ((seq * n_kv + kv_h) * NP + part) * TOTAL_ROWS + global_row + base = partial_slot * TOTAL_ROWS + global_row if row < CTA_ROWS: # Convert the running max from log2 units (scale_qk folds # in LOG2E) to natural-log units: the shared reduce @@ -1779,6 +1811,76 @@ def _emit(o_norm, sub, query_idx, query_head): 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, + 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, + 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): + # Capacity is fixed for graph capture, so the planner clears any + # unused tail records. The reducer never consumes those slots; + # guard the whole task body to avoid Q processing and scratch + # writes for an empty record. The condition is CTA-uniform, so + # all barriers in the task body remain convergent. + 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, @@ -1801,6 +1903,8 @@ def pa_decode_tile_launch( 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 ): # `contract` lets the mul/add pairs in the softmax rescale and the @@ -1826,8 +1930,14 @@ def pa_decode_tile_launch( stride_o_head, stride_q_row, stride_q_head, + work_info, + num_seqs, ).launch( - grid=(num_seqs, num_kv_heads * query_splits, NP), + grid=( + 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, ) 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..fdd10b4abfd --- /dev/null +++ b/aiter/ops/flydsl/kernels/pa_decode_plan.py @@ -0,0 +1,166 @@ +# 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 + +from .pa_decode_reduce import MAX_CONTEXT_PARTITIONS + + +@triton.jit +def _plan_pa_decode( + lengths, + work, + reduce_info, + B: tl.constexpr, + CAPACITY: tl.constexpr, + MAX_PARTS: tl.constexpr, + BLOCK_B: tl.constexpr, + BLOCK_P: tl.constexpr, +): + seq = tl.program_id(0) + b = tl.arange(0, BLOCK_B) + ctx = tl.load(lengths + b, b < B, other=0) + tiles = tl.where(b < B, tl.maximum(ctx, 0).to(tl.int64) + 255, 0) // 256 + 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) + # Integer prefix apportionment telescopes to the budget, even when the + # lengths differ by orders of magnitude. Each nonempty request gets one + # task first; clamp extras to useful tiles and the reducer's per-row limit. + 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_tiles = (tl.maximum(seq_ctx, 0).to(tl.int64) + 255) // 256 + 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 = part.to(tl.int64) * seq_tiles // tl.maximum(count, 1) + end = (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) + + # The launch capacity is static for graph capture. Clear every padded work + # record so the attention kernel can skip it without touching stale scratch. + # These stores and the active stores are disjoint, including when every + # context is empty. + 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, first/last 256-token tile, length. + ``reduce_info`` is [batch, 2]: first packed task, actual task count. + Scratch uses [KV heads, capacity, query rows (, head dim)]. + """ + + work_info: torch.Tensor + reduce_info: torch.Tensor + num_kv_heads: int + max_partitions: int + + @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") + if not 1 <= self.max_partitions <= MAX_CONTEXT_PARTITIONS: + raise ValueError("invalid plan max_partitions") + 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 = MAX_CONTEXT_PARTITIONS, + workgroup_budget: int | None = None, + plan: PADecodePlan | None = None, +) -> PADecodePlan: + """Build/update a plan on the current stream without GPU-to-CPU readback. + + Allocate once outside graph capture, then pass ``plan=...`` to refresh the + same metadata in place. Include this refresh in end-to-end measurements. + The budget counts CTAs over all KV heads, with fused query positions. + This is opt-in: uniform or short-context workloads may favor static splits. + """ + 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") + if not 1 <= max_partitions <= MAX_CONTEXT_PARTITIONS: + raise ValueError(f"max_partitions must be in [1, {MAX_CONTEXT_PARTITIONS}]") + dev = context_lengths.device + if plan is None: + if workgroup_budget is None: + workgroup_budget = ( + 2 * torch.cuda.get_device_properties(dev).multi_processor_count + ) + 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, + ) + 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") + 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, + 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 index c8f889704c9..aa602493436 100644 --- a/aiter/ops/flydsl/kernels/pa_decode_reduce.py +++ b/aiter/ops/flydsl/kernels/pa_decode_reduce.py @@ -70,6 +70,7 @@ def compile_pa_decode_ps_reduce( logits_dtype_str: str, sink_dtype_str: str, use_sinks: bool, + use_work_plan: bool = False, ): """Build the partitioned-softmax reduction used by ``pa_decode``. @@ -152,6 +153,7 @@ def pa_decode_ps_reduce_kernel( 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 @@ -174,6 +176,15 @@ def pa_decode_ps_reduce_kernel( 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]) + 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 c_reduce_width = fx.Int32(reduce_width) c_four = fx.Int32(4) c_qgs = query_group_size @@ -219,9 +230,9 @@ def _wave_reduce_sum(value): 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): + if fx.const_expr(chunk_size == warp_size and not use_work_plan): stats_offset = ( - batch_idx * stride_exp_sums_seq + stats_seq_offset + kv_head_idx * stride_exp_sums_head + part_idx * stride_exp_sums_part + eqgs_idx @@ -229,9 +240,11 @@ def _wave_reduce_sum(value): 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) + lane_in_range = (lane < fx.Int32(chunk_size)) & ( + part_idx < c_part_num + ) stats_offset = ( - batch_idx * stride_exp_sums_seq + stats_seq_offset + kv_head_idx * stride_exp_sums_head + part_idx * stride_exp_sums_part + eqgs_idx @@ -251,13 +264,25 @@ def _wave_reduce_sum(value): lane_exp_sum = zero_f for chunk_idx in fx.range_constexpr(partitions_per_lane): part_max = part_maxes[chunk_idx] - part_scale = (part_max > neg_inf).select( - fx.exp2( - (part_max - safe_global_max) * c_log2e, - fastmath="fast", - ), - zero_f, - ) + if fx.const_expr(use_work_plan): + # ``select`` evaluates its exp2 operand even for an + # inactive planned partition. Keep the static path + # branch-free, but predicate that expensive operation + # when the per-request count is dynamic. + 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 @@ -281,13 +306,21 @@ def _wave_reduce_sum(value): chunk_base = chunk_idx * warp_size chunk_size = min(warp_size, max_context_partition_num - chunk_base) part_idx = lane + fx.Int32(chunk_base) - weight = scaled_sums[chunk_idx] / safe_global_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 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] / safe_global_exp_sum + lds_weights[part_idx] = weight + else: + weight = scaled_sums[chunk_idx] / safe_global_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() @@ -296,42 +329,75 @@ def _wave_reduce_sum(value): output_element = head_wave * c_warp_size + lane group_part_begin = partition_group * fx.Int32(parts_per_group) acc = zero_f - 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 = ( - batch_idx * stride_logits_seq - + 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 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: - partial_offset = (partition_group - fx.Int32(1)) * fx.Int32( - head_size - ) + output_element - lds_partials[partial_offset] = acc + 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): - partial_offset = ( - fx.Int32((other_group - 1) * head_size) + output_element - ) - acc = acc + fx.Float32(lds_partials[partial_offset]) + 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): # Exact powers of two have no inactive lanes inside their reduction # subgroup. Keep that original path unchanged; only partial # subgroups need an EXEC-masked load to avoid carrying a predicate # across the shuffle sequence. - if fx.const_expr(max_context_partition_num == reduce_width): + if fx.const_expr( + max_context_partition_num == reduce_width and not use_work_plan + ): lane_in_range = lane < c_part_num lane_in_reduce = lane < c_reduce_width part_sum = zero_f @@ -339,7 +405,7 @@ def _wave_reduce_sum(value): if lane_in_reduce: part_idx = lane_in_range.select(lane, zero_i) stats_offset = ( - batch_idx * stride_exp_sums_seq + stats_seq_offset + kv_head_idx * stride_exp_sums_head + part_idx * stride_exp_sums_part + eqgs_idx @@ -351,7 +417,7 @@ def _wave_reduce_sum(value): else: lane_in_range = lane < c_part_num stats_offset = ( - batch_idx * stride_exp_sums_seq + stats_seq_offset + kv_head_idx * stride_exp_sums_head + lane * stride_exp_sums_part + eqgs_idx @@ -364,10 +430,18 @@ def _wave_reduce_sum(value): global_max = _wave_reduce_max(part_max) 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, - ) + 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_sum * part_scale global_exp_sum = _wave_reduce_sum(scaled_sum) if fx.const_expr(use_sinks): @@ -383,29 +457,56 @@ def _wave_reduce_sum(value): safe_global_exp_sum = (global_exp_sum > zero_f).select( global_exp_sum, one_f ) - weight_local = scaled_sum / safe_global_exp_sum + if fx.const_expr(use_work_plan): + weight_local = zero_f + if lane_in_range: + weight_local = scaled_sum / safe_global_exp_sum + else: + weight_local = scaled_sum / safe_global_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, + if fx.const_expr(use_work_plan): + for part_idx in fx.range_constexpr(max_context_partition_num): + c_part_idx = fx.Int32(part_idx) + if c_part_idx < c_part_num: + 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: + 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 = ( - batch_idx * stride_logits_seq - + 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 + 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: # A wave covers several 64-partition chunks. Lane ``l`` owns # partitions l, l+64, l+128, and l+192 (as present). Reduce the @@ -422,9 +523,9 @@ def _wave_reduce_sum(value): 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): + if fx.const_expr(chunk_size == warp_size and not use_work_plan): stats_offset = ( - batch_idx * stride_exp_sums_seq + stats_seq_offset + kv_head_idx * stride_exp_sums_head + part_idx * stride_exp_sums_part + eqgs_idx @@ -432,9 +533,11 @@ def _wave_reduce_sum(value): 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) + lane_in_range = (lane < fx.Int32(chunk_size)) & ( + part_idx < c_part_num + ) stats_offset = ( - batch_idx * stride_exp_sums_seq + stats_seq_offset + kv_head_idx * stride_exp_sums_head + part_idx * stride_exp_sums_part + eqgs_idx @@ -454,13 +557,21 @@ def _wave_reduce_sum(value): lane_exp_sum = zero_f for chunk_idx in fx.range_constexpr(partitions_per_lane): part_max = part_maxes[chunk_idx] - part_scale = (part_max > neg_inf).select( - fx.exp2( - (part_max - safe_global_max) * c_log2e, - fastmath="fast", - ), - zero_f, - ) + 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 @@ -482,31 +593,62 @@ def _wave_reduce_sum(value): acc = zero_f for chunk_idx in fx.range_constexpr(partitions_per_lane): - weight_local_i32 = ( - scaled_sums[chunk_idx] / safe_global_exp_sum - ).bitcast(fx.Int32) chunk_base = chunk_idx * warp_size chunk_size = min(warp_size, max_context_partition_num - chunk_base) - 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, + if fx.const_expr(use_work_plan): + # Initialize in the enclosing constexpr-loop scope so the + # FlyDSL dynamic-if rewriter never observes a stale value + # from a previous unrolled chunk. + weight_local_i32 = zero_f.bitcast(fx.Int32) + if fx.Int32(chunk_base) < c_part_num: + weight_local_i32 = ( + scaled_sums[chunk_idx] / safe_global_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] / safe_global_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 = ( - batch_idx * stride_logits_seq - + 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 + 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 query_idx = eqgs_idx // c_qgs if fx.const_expr(use_parallel_lds): @@ -551,6 +693,7 @@ def launch_pa_decode_ps_reduce_kernel( 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( @@ -571,6 +714,7 @@ def launch_pa_decode_ps_reduce_kernel( 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), diff --git a/aiter/ops/flydsl/pa_decode.py b/aiter/ops/flydsl/pa_decode.py index 7449d48ef62..b3adc361046 100644 --- a/aiter/ops/flydsl/pa_decode.py +++ b/aiter/ops/flydsl/pa_decode.py @@ -40,6 +40,7 @@ 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, plan_pa_decode as plan_pa_decode from .kernels.pa_decode_reduce import ( MAX_CONTEXT_PARTITIONS, compile_pa_decode_ps_reduce, @@ -51,14 +52,30 @@ def get_recommended_splits( num_sequences: int, num_kv_heads: int, split_kv_blocks: int = 1, - max_partitions: int = 8, + max_partitions: int | None = None, + *, + max_context_length: int | None = None, ) -> int: - """Choose a split count, with a caller-configurable upper clamp. - - The default of eight preserves the existing policy. Long-context callers - can opt into the FlyDSL reducer's larger domain with, for example, - ``max_partitions=256``. + """Choose a split count before allocating partition scratch buffers. + + Without a host-side ``max_context_length`` hint, retain the legacy policy + and default upper clamp of eight. With a hint, target two CTAs per CU, + capped by the number of 256-token compute tiles and the reducer's limit. + This lets small, long-context batches use more KV partitions without + multiplying the partition count by the number of physical pages per tile. + Short contexts and already large grids retain the legacy recommendation. + + An explicit ``max_partitions`` remains an upper clamp in either mode. + The hint is a scheduling bound, not a replacement for the GPU-resident + per-sequence lengths; no tensor readback or planning kernel is needed. + Callers must allocate scratch for the returned count and pass that exact + count to ``pa_decode``. For heterogeneous lengths this still uses the same + partition count per sequence; it is not a variable-work scheduler. """ + 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}], " @@ -68,6 +85,16 @@ def get_recommended_splits( 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)) + # KV partitions consist of compute tiles, regardless of physical page + # size. Keep a short-context floor of eight to preserve the existing + # query-split launch and avoid perturbing already tuned small grids. + 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)) @@ -121,6 +148,7 @@ def launch_pa_decode_ps_reduce( head_size: int, context_partition_num: int, stream: torch.cuda.Stream, + reduce_info: torch.Tensor | None = None, ) -> None: if context_partition_num > MAX_CONTEXT_PARTITIONS: raise ImportError( @@ -137,6 +165,7 @@ def launch_pa_decode_ps_reduce( output.dtype if sink_token is None else sink_token.dtype ), use_sinks=use_sinks, + use_work_plan=reduce_info is not None, ) sink_ptr = ( ptr_arg(sink_token, _flydsl_pointer_dtype(sink_token.dtype)) @@ -165,6 +194,11 @@ def launch_pa_decode_ps_reduce( 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, ) @@ -191,6 +225,7 @@ def pa_decode( sinks: torch.Tensor = None, sliding_window: int = 0, ps: bool = True, + work_plan: PADecodePlan | None = None, ) -> None: """FlyDSL paged-attention fp8 decode. @@ -211,6 +246,12 @@ def pa_decode( that every block-table entry used by a sequence is a physical block index in ``[0, min(key_cache.shape[0], value_cache.shape[0]))`` -- a packed cache reaches V through a shifted view that spans fewer blocks than K. + + ``work_plan`` opts into GPU-planned variable partition counts. Build or + refresh it with ``plan_pa_decode`` on the current stream after updating + lengths. Its ``max_partitions`` must equal ``max_context_partition_num``. + Planned scratch is packed as [KV heads, plan.capacity, query rows (, D)]; + the static API's dense per-sequence scratch layout remains unchanged. """ if context_partition_size != KV_COMPUTE_BLOCK: raise NotImplementedError( @@ -507,6 +548,14 @@ def normalize_scale(scale, name): ) 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" + ) pmax = max_logits psum = exp_sums pout = temporary_output @@ -564,7 +613,8 @@ def normalize_scale(scale, name): # supported split counts (1, 2, 4) must divide query_length. query_splits = 1 if ( - arch == "gfx950" + work_plan is None + and arch == "gfx950" and per_token_kv and head_dim == 128 and query.dtype == torch.bfloat16 @@ -593,15 +643,18 @@ def normalize_scale(scale, name): wide_kv_addressing=wide_kv_addressing, prefetch_v=prefetch_v, query_splits=query_splits, + use_work_plan=work_plan is not None, ) - if num_partitions == 1: + if num_partitions == 1 and work_plan is None: # NP==1 writes output directly; partials unused (caller buffers ignored). 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( @@ -682,9 +735,15 @@ def normalize_scale(scale, name): 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: + 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 ) @@ -698,16 +757,17 @@ def normalize_scale(scale, name): output_5d.stride(1), output_5d.stride(2), output_5d.stride(3), - pmax.stride(0), - pmax.stride(1), - pmax.stride(2), - pout.stride(0), - pout.stride(1), - pout.stride(2), - pout.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/docs/flydsl_pa_decode_plan.md b/docs/flydsl_pa_decode_plan.md new file mode 100644 index 00000000000..ec86f4e7e90 --- /dev/null +++ b/docs/flydsl_pa_decode_plan.md @@ -0,0 +1,92 @@ +# GPU work plans for FlyDSL PA decode + +`plan_pa_decode` assigns different KV partition counts to requests according to +their GPU-resident context lengths. It is intended for batches with unequal work, +such as one long context mixed with many short contexts. It is opt-in: a planning +launch and packed reduction can cost more than static splits for uniform or short +contexts. Compare **plan + attention + reduction** against the static path. + +The attention math, per-token scaling and causal MTP masking are shared with the +static FlyDSL kernel. The planner is a Triton GPU kernel, following the existing +variable-context scheduling approach in `mqa_logits/pa_mqa_logits_fp4.py`. + +## Scheduling and storage + +Each nonempty request first receives one task. Remaining task slots are apportioned +by its number of 256-token compute tiles using integer prefix sums. Counts are +capped by useful tiles and `max_partitions` (1 through 256). The workgroup budget is +a target across all KV heads; capacity is rounded up to at least one slot per +request/head. Clamping can leave unused slots. Physical page size is independent +of this scheduling tile size. + +The planner writes two contiguous int32 CUDA tensors: + +- `work_info[capacity, 4]`: request index, first tile, exclusive last tile, context length. +- `reduce_info[batch, 2]`: first packed task slot and actual partition count. + +The launch grid has a fixed capacity for graph capture. Padded work records are +zeroed and skipped by the attention kernel, leaving their scratch slots untouched. +Each real task processes one contiguous KV range; ranges cover each request exactly +once. Reduction reads only the request's valid partitions and produces zero for +empty contexts. No GPU-to-CPU readback, dynamic allocation, or CPU length +inspection is required when refreshing a preallocated plan. + +Planned scratch is packed as `[num_kv_heads, capacity, query_rows]` for max/sum and +`[num_kv_heads, capacity, query_rows, head_dim]` for partial outputs, where +`query_rows = query_length * num_query_heads // num_kv_heads`. +The static API continues to use its dense per-request scratch layout. + +## Usage + +Use the Python FlyDSL entry point for planned execution. The shared +`torch.ops.aiter.pa_decode_flydsl` entry point retains its existing static interface. + +```python +import torch +from aiter.ops.flydsl.pa_decode import pa_decode, plan_pa_decode + +# query/output: [batch * query_length, num_query_heads, head_dim] +# K/V cache layouts and scales are the same as for static pa_decode. +num_kv_heads = key_cache.shape[1] +plan = plan_pa_decode(context_lengths, num_kv_heads, max_partitions=256) +rows = query_length * query.shape[1] // num_kv_heads +shape = (num_kv_heads, plan.capacity, rows) +exp_sums = torch.empty(shape, dtype=torch.float32, device=query.device) +max_logits = torch.empty_like(exp_sums) +partials = torch.empty((*shape, query.shape[2]), dtype=query.dtype, device=query.device) + +def step(): + # Refresh after changing lengths, on the same stream as attention. + plan_pa_decode(context_lengths, num_kv_heads, max_partitions=256, plan=plan) + pa_decode( + output, query, key_cache, value_cache, context_lengths, block_tables, + softmax_scale, query_length, plan.max_partitions, + compute_type=key_cache.dtype, + key_scale=key_scale, value_scale=value_scale, + exp_sums=exp_sums, max_logits=max_logits, temporary_output=partials, + work_plan=plan, + ) + +step() # Warm up compilation before graph capture. +graph = torch.cuda.CUDAGraph() +with torch.cuda.graph(graph): + step() +# Update existing input buffers, then graph.replay(). +``` + +Reuse an already built plan without refreshing only when lengths are unchanged. +Changing batch size, KV-head count, the partition limit or launch capacity requires +a compatible new plan and scratch buffers. Metadata values should be produced by +the planner, not edited independently. The planner accepts contiguous int32 CUDA +length vectors with 1 through 4096 requests. Lengths and block-table indices must +satisfy the same validity requirements as static attention. + +## Verification + +```bash +python -m pytest -q op_tests/test_flydsl_pa_decode.py +``` + +Tests cover exact work ownership, sparse pages, empty requests, causal tails, +non-power-of-two partition caps, multiple KV heads, supported reducer branches, +NaN-filled scratch, and graph replay after changing lengths. diff --git a/op_tests/test_flydsl_pa_decode.py b/op_tests/test_flydsl_pa_decode.py index 8c059963f2f..6207d339b1c 100644 --- a/op_tests/test_flydsl_pa_decode.py +++ b/op_tests/test_flydsl_pa_decode.py @@ -12,6 +12,9 @@ Contexts have equal lengths and include the MTP query tokens. Query position ``p`` attends to ``max(0, context_length - query_length + 1 + p)`` KV tokens. Timing includes the FlyDSL attention kernel and its native FlyDSL reduction. +Automatic splits use the host-known context length and GPU occupancy, up to +256 partitions. Use ``--max-partitions 8`` for the legacy clamp, or +``--num-partitions`` to bypass the recommendation with exact counts. """ import argparse @@ -33,6 +36,7 @@ def _default_cuda_device(): # Scoped rather than set at import time: this module now runs in the shared # Standard Tests shard, where a global default-device switch would follow # every later test in the session. + _require_gpu() previous = torch.get_default_device() torch.set_default_device("cuda") yield @@ -62,11 +66,13 @@ def _default_cuda_device(): MAX_CONTEXT_PARTITIONS, get_recommended_splits, pa_decode, + plan_pa_decode, ) except (ImportError, AttributeError, RuntimeError, OSError): MAX_CONTEXT_PARTITIONS = 256 get_recommended_splits = None pa_decode = None + plan_pa_decode = None def _quant_dtype() -> torch.dtype: @@ -147,6 +153,267 @@ def run_torch( return output +def _run_accuracy_case( + lengths, + *, + num_query_heads, + num_kv_heads, + head_dim, + block_size, + query_dtype, + num_partitions, + tolerance, + trans_v=False, +): + """Run PA against a dequantized-FP8 torch reference without benchmarking.""" + _require_gpu() + if num_query_heads % num_kv_heads: + raise ValueError("num_query_heads must be divisible by num_kv_heads") + + torch.manual_seed(0) + batch_size = len(lengths) + pages_per_sequence = [ + (context_length + block_size - 1) // block_size for context_length in lengths + ] + max_pages = max(pages_per_sequence) + num_blocks = sum(pages_per_sequence) + + query = torch.empty( + batch_size, num_query_heads, head_dim, dtype=query_dtype + ).uniform_(-0.5, 0.5) + key = torch.empty( + num_blocks, num_kv_heads, block_size, head_dim, dtype=query_dtype + ).uniform_(-0.5, 0.5) + value = torch.empty( + num_blocks, num_kv_heads, head_dim, block_size, dtype=query_dtype + ).uniform_(-0.5, 0.5) + + quant_dtype = _quant_dtype() + key_quant, key_scale = per_tensor_quant(key, quant_dtype=quant_dtype) + value_quant, value_scale = per_tensor_quant(value, quant_dtype=quant_dtype) + key_cache = ( + key_quant.view(num_blocks, num_kv_heads, block_size, head_dim // 16, 16) + .permute(0, 1, 3, 2, 4) + .contiguous() + ) + if trans_v: + value_cache = ( + value_quant.permute(0, 1, 3, 2) + .contiguous() + .view(num_blocks, num_kv_heads, block_size // 16, 16, head_dim) + .permute(0, 1, 2, 4, 3) + .contiguous() + ) + else: + value_cache = value_quant.contiguous() + + block_tables = torch.zeros( + batch_size, max_pages, dtype=dtypes.i32, device=query.device + ) + next_page = 0 + for seq_idx, page_count in enumerate(pages_per_sequence): + block_tables[seq_idx, :page_count] = torch.arange( + next_page, + next_page + page_count, + dtype=dtypes.i32, + device=query.device, + ) + next_page += page_count + context_lengths = torch.tensor(lengths, dtype=dtypes.i32, device=query.device) + + reference = run_torch( + query, + key_quant, + value_quant, + block_tables, + context_lengths, + key_scale, + value_scale, + ) + output = torch.empty_like(query) + torch.ops.aiter.pa_decode_flydsl( + output, + query, + key_cache, + value_cache, + context_lengths, + block_tables, + head_dim**-0.5, + 1, + num_partitions, + KV_COMPUTE_BLOCK, + quant_dtype, + None, + key_scale, + value_scale, + ) + + torch.testing.assert_close( + output.float(), + reference.float(), + atol=tolerance, + rtol=tolerance, + ) + + +def _run_mtp4_fused_reference_case( + lengths, num_kv_heads, block_size, trans_v, num_partitions +): + """Check MTP4 with sparse pages, varying scales, and a causal FP32 reference.""" + torch.manual_seed(37) + query_length, query_group_size, head_dim = 4, 16, 128 + batch_size = len(lengths) + num_query_heads = num_kv_heads * query_group_size + pages_per_sequence = [(length + block_size - 1) // block_size for length in lengths] + max_pages = max(pages_per_sequence) + num_pages = sum(pages_per_sequence) + quant_dtype = _quant_dtype() + + query = torch.empty( + batch_size * query_length, + num_query_heads, + head_dim, + dtype=dtypes.bf16, + ).uniform_(-0.5, 0.5) + key = torch.empty( + num_pages, num_kv_heads, block_size, head_dim, dtype=dtypes.bf16 + ).uniform_(-0.5, 0.5) + value = torch.empty( + num_pages, num_kv_heads, head_dim, block_size, dtype=dtypes.bf16 + ).uniform_(-0.5, 0.5) + key_quant, key_scale = pertoken_quant(key, quant_dtype=quant_dtype) + value_token_quant, value_scale = pertoken_quant( + value.permute(0, 1, 3, 2).contiguous(), quant_dtype=quant_dtype + ) + value_quant = value_token_quant.permute(0, 1, 3, 2).contiguous() + + token = torch.arange(num_pages * block_size).reshape(num_pages, 1, block_size, 1) + kv_head = torch.arange(num_kv_heads).reshape(1, num_kv_heads, 1, 1) + key_scale *= torch.exp2(((2 * token + kv_head) % 4 - 2).float()) + value_scale *= torch.exp2(((token + 2 * kv_head) % 5 - 3).float()) + + # Make C1--C4 analytically exact while still detecting causal off-by-one: + # K=0 gives uniform attention, V progresses as [.25, .5, .75, 1], and + # scale=1. A C4 row therefore produces [.25, .375, .5, .625]. + data_page = 0 + for length, page_count in zip(lengths, pages_per_sequence): + if 0 < length <= query_length: + capacity = page_count * block_size + token_values = ( + torch.arange(capacity, dtype=dtypes.fp32) % query_length + 1 + ) * 0.25 + key_quant[data_page : data_page + page_count].zero_() + value_quant[data_page : data_page + page_count] = ( + token_values.reshape(page_count, 1, 1, block_size) + .expand(page_count, num_kv_heads, head_dim, block_size) + .to(quant_dtype) + ) + key_scale[data_page : data_page + page_count].fill_(1.0) + value_scale[data_page : data_page + page_count].fill_(1.0) + data_page += page_count + + # Scatter logical pages into a permuted set of odd physical pages. Leave + # holes and page zero finite so masked/padded token reads remain valid. + selected_pages = 2 * torch.randperm(num_pages) + 1 + physical_page_count = 2 * num_pages + 1 + + def scatter_pages(tensor): + sparse = torch.zeros( + (physical_page_count, *tensor.shape[1:]), dtype=tensor.dtype + ) + sparse[selected_pages] = tensor + return sparse + + key_quant = scatter_pages(key_quant) + value_quant = scatter_pages(value_quant) + key_scale = scatter_pages(key_scale) + value_scale = scatter_pages(value_scale) + num_pages = physical_page_count + + key_cache = ( + key_quant.view( + num_pages, + num_kv_heads, + block_size, + head_dim // 16, + 16, + ) + .permute(0, 1, 3, 2, 4) + .contiguous() + ) + if trans_v: + value_cache = ( + value_quant.permute(0, 1, 3, 2) + .contiguous() + .view( + num_pages, + num_kv_heads, + block_size // 16, + 16, + head_dim, + ) + .permute(0, 1, 2, 4, 3) + .contiguous() + ) + else: + value_cache = value_quant + + block_tables = torch.zeros(batch_size, max_pages, dtype=dtypes.i32) + next_page = 0 + for seq, page_count in enumerate(pages_per_sequence): + block_tables[seq, :page_count] = selected_pages[ + next_page : next_page + page_count + ] + next_page += page_count + context_lengths = torch.tensor(lengths, dtype=dtypes.i32) + + reference = torch.zeros_like(query, dtype=dtypes.fp32) + for seq, length in enumerate(lengths): + if length == 0: + continue + token_ids = torch.arange(length) + logical_pages = token_ids // block_size + token_offsets = token_ids % block_size + physical_pages = block_tables[seq, logical_pages].long() + keys = key_quant[physical_pages, :, token_offsets, :].float() + values = value_quant[physical_pages, :, :, token_offsets].float() + keys *= key_scale[physical_pages, :, token_offsets, 0].float().unsqueeze(-1) + values *= value_scale[physical_pages, :, token_offsets, 0].float().unsqueeze(-1) + keys = keys.repeat_interleave(query_group_size, dim=1) + values = values.repeat_interleave(query_group_size, dim=1) + for position in range(query_length): + visible = max(0, length - (query_length - 1) + position) + if visible == 0: + continue + row = seq * query_length + position + scores = ( + torch.einsum("hd,khd->hk", query[row].float(), keys[:visible]) + * head_dim**-0.5 + ) + probs = torch.softmax(scores, dim=-1) + reference[row] = torch.einsum("hk,khd->hd", probs, values[:visible]) + + output = torch.full_like(query, float("nan")) + torch.ops.aiter.pa_decode_flydsl( + output, + query, + key_cache, + value_cache, + context_lengths, + block_tables, + head_dim**-0.5, + query_length, + num_partitions, + KV_COMPUTE_BLOCK, + quant_dtype, + None, + key_scale, + value_scale, + ) + assert torch.isfinite(output).all() + torch.testing.assert_close(output.float(), reference, rtol=0.005, atol=0.005) + + def _run_flydsl( output, query, @@ -195,7 +462,7 @@ def run_pa_decode_tile_case( block_size, dtype, trans_v, - max_partitions=8, + max_partitions=None, per_token=False, query_length=1, num_partitions=None, @@ -221,6 +488,7 @@ def run_pa_decode_tile_case( num_kv_heads, split_kv_blocks=KV_COMPUTE_BLOCK // block_size, max_partitions=max_partitions, + max_context_length=context_length, ) torch.manual_seed(0) @@ -578,13 +846,14 @@ def run_once(fn, *args, **kwargs): monkeypatch.setitem(globals(), "run_perftest", run_once) explicit_partitions = case["num_partitions"] - max_partitions = case.get("max_partitions", 8) + max_partitions = case.get("max_partitions") if explicit_partitions is None: expected_partitions = get_recommended_splits( case["batch_size"], case["num_kv_heads"], split_kv_blocks=KV_COMPUTE_BLOCK // case["block_size"], max_partitions=max_partitions, + max_context_length=case["context_length"], ) else: expected_partitions = explicit_partitions @@ -605,6 +874,345 @@ def run_once(fn, *args, **kwargs): assert set(selected_query_splits) == {expected} +def _assert_plan(plan, lengths): + work = plan.work_info.cpu().tolist() + reductions = plan.reduce_info.cpu().tolist() + tiles_per_sequence = [(max(length, 0) + 255) // 256 for length in lengths] + nonempty = [int(tiles > 0) for tiles in tiles_per_sequence] + total_tiles = max(sum(tiles_per_sequence), 1) + remaining = plan.capacity - sum(nonempty) + cumulative_tiles = 0 + expected_counts = [] + for tiles, is_nonempty in zip(tiles_per_sequence, nonempty): + lower = cumulative_tiles * remaining // total_tiles + cumulative_tiles += tiles + upper = cumulative_tiles * remaining // total_tiles + expected_counts.append( + min(is_nonempty + upper - lower, tiles, plan.max_partitions) + ) + + expected_reductions = [] + expected_work = [] + offset = 0 + for seq, (length, tiles, count) in enumerate( + zip(lengths, tiles_per_sequence, expected_counts) + ): + start = offset + expected_reductions.append([start, count]) + for part in range(count): + expected_work.append( + [ + seq, + part * tiles // count, + (part + 1) * tiles // count, + length, + ] + ) + + actual_start, actual_count = reductions[seq] + assert (actual_start, actual_count) == (start, count) + assert 0 <= actual_count <= min(tiles, plan.max_partitions) + assert (actual_count > 0) == (length > 0) + previous_end = 0 + for task in work[actual_start : actual_start + actual_count]: + assert task[0] == seq and task[3] == length + assert task[1] == previous_end and task[1] < task[2] <= tiles + previous_end = task[2] + assert previous_end == tiles + offset += actual_count + expected_work.extend([[0, 0, 0, 0]] * (plan.capacity - offset)) + assert reductions == expected_reductions + assert work == expected_work + assert offset <= plan.capacity + assert all(row == [0, 0, 0, 0] for row in work[offset:]) + + +@pytest.mark.parametrize( + "heads,max_parts,budget", [(1, 1, 512), (1, 7, 17), (1, 256, 512), (2, 256, 512)] +) +@pytest.mark.parametrize( + "lengths", + [ + [0], + [1], + [257], + [200003], + [0] * 8, + [0, 1, 3, 4, 257, 4095, 16385, 200003], + [100000] * 8, + [i * 7919 % 200003 for i in range(200)], + ], +) +def test_plan_covers_each_token_once(lengths, heads, max_parts, budget): + context = torch.tensor(lengths, dtype=torch.int32) + plan = plan_pa_decode( + context, heads, max_partitions=max_parts, workgroup_budget=budget + ) + _assert_plan(plan, lengths) + + +def test_plan_graph_refresh_overwrites_old_metadata(): + lengths = [200003] * 8 + context = torch.tensor(lengths, dtype=torch.int32) + plan = plan_pa_decode(context, 1) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + plan_pa_decode(context, 1, plan=plan) + for lengths in ([0] * 8, [0, 1, 3, 4, 257, 4096, 16385, 200003], [200003] * 8): + context.copy_(torch.tensor(lengths, dtype=torch.int32)) + graph.replay() + _assert_plan(plan, lengths) + + +def _planned_call(*args, **kwargs): + plan = plan_pa_decode(args[4], args[2].shape[1], max_partitions=args[8]) + pa_decode(*args, work_plan=plan, **kwargs) + + +@pytest.mark.parametrize("block_size", [16, 128]) +@pytest.mark.parametrize("trans_v", [False, True]) +@pytest.mark.parametrize("max_parts", [1, 7, 64, 86, 256]) +@pytest.mark.parametrize("heads", [1, 2]) +def test_planned_mtp4_sparse_causal_reference( + monkeypatch, block_size, trans_v, max_parts, heads +): + monkeypatch.setattr(torch.ops.aiter, "pa_decode_flydsl", _planned_call) + _run_mtp4_fused_reference_case( + [0, 1, 2, 3, 4, 255, 256, 257, 2305, 16387], + heads, + block_size, + trans_v, + max_parts, + ) + + +@pytest.mark.parametrize( + "head_dim,block_size,parts,dtype", + [ + (64, 16, 8, torch.bfloat16), + (256, 64, 86, torch.float16), + (1024, 128, 256, torch.bfloat16), + (128, 128, 256, torch.float16), + ], +) +def test_planned_decode_other_reducer_paths( + monkeypatch, head_dim, block_size, parts, dtype +): + monkeypatch.setattr(torch.ops.aiter, "pa_decode_flydsl", _planned_call) + _run_accuracy_case( + [0, 1, 257, 4099], + num_query_heads=8, + num_kv_heads=2, + head_dim=head_dim, + block_size=block_size, + query_dtype=dtype, + num_partitions=parts, + tolerance=0.005, + trans_v=True, + ) + + +@pytest.mark.parametrize("long_context", [16384, 16385]) +def test_planned_decode_register_reducer_chunk_boundary(monkeypatch, long_context): + """Exercise the generic reducer with 64 and 65 active partitions.""" + monkeypatch.setattr(torch.ops.aiter, "pa_decode_flydsl", _planned_call) + _run_accuracy_case( + [0, 1, 257, long_context], + num_query_heads=8, + num_kv_heads=2, + head_dim=256, + block_size=64, + query_dtype=torch.float16, + num_partitions=86, + tolerance=0.005, + trans_v=True, + ) + + +def test_planned_decode_graph_replay_with_poisoned_scratch(monkeypatch): + def capture_and_replay(*args, **kwargs): + output, query, key, _value, context = args[:5] + heads, rows = ( + key.shape[1], + query.shape[0] // context.numel() * query.shape[1] // key.shape[1], + ) + plan = plan_pa_decode(context, heads, max_partitions=args[8]) + shape = (heads, plan.capacity, rows) + psum = torch.full(shape, float("nan"), dtype=torch.float32) + pmax = torch.full_like(psum, float("nan")) + pout = torch.full((*shape, query.shape[2]), float("nan"), dtype=query.dtype) + + def run(): + plan_pa_decode(context, heads, max_partitions=args[8], plan=plan) + pa_decode( + *args, + **kwargs, + work_plan=plan, + exp_sums=psum, + max_logits=pmax, + temporary_output=pout, + ) + + run() # Compile before capture. + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + active_tasks = int(plan.reduce_info[:, 1].sum().item()) + assert active_tasks < plan.capacity + for scratch in (pmax, psum, pout): + assert torch.isnan(scratch[:, active_tasks:]).all() + original = context.clone() + for scratch in (pmax, psum, pout): + scratch.fill_(float("nan")) + context.zero_() + graph.replay() + assert torch.equal(output, torch.zeros_like(output)) + for scratch in (pmax, psum, pout): + assert torch.isnan(scratch).all() + context.copy_(original) + graph.replay() + for scratch in (pmax, psum, pout): + assert torch.isnan(scratch[:, active_tasks:]).all() + + monkeypatch.setattr(torch.ops.aiter, "pa_decode_flydsl", capture_and_replay) + _run_mtp4_fused_reference_case([0, 257, 4096, 200003], 1, 128, True, 256) + + +def test_plan_rejects_incompatible_reuse(): + context = torch.tensor([1, 257], dtype=torch.int32) + plan = plan_pa_decode(context, 1, max_partitions=7) + with pytest.raises(ValueError, match="max_partitions"): + plan_pa_decode(context, 1, plan=plan) + with pytest.raises(ValueError, match="KV head"): + plan_pa_decode(context, 2, max_partitions=7, plan=plan) + with pytest.raises(ValueError, match="shape"): + plan_pa_decode(context[:1], 1, max_partitions=7, plan=plan) + + +@pytest.mark.parametrize("split_kv_blocks", [2, 16]) +@pytest.mark.parametrize( + "batch_size,context_length,max_partitions,expected", + [ + (1, None, None, 8), + (1, 0, None, 8), + (1, 257, None, 8), + (1, 4096, None, 16), + (1, 16384, None, 64), + (1, 100000, None, 256), + (2, 200000, None, 256), + (3, 200000, None, 171), + (5, 200000, None, 103), + (6, 200000, None, 86), + (7, 200000, None, 74), + (4, 200000, None, 128), + (8, 200000, None, 64), + (1, 200000, 8, 8), + (4, 200000, 5, 5), + (8, 200000, 32, 32), + ], +) +def test_recommended_splits_uses_context_work( + monkeypatch, split_kv_blocks, batch_size, context_length, max_partitions, expected +): + if get_recommended_splits is None: + pytest.skip("FlyDSL is not available") + monkeypatch.setattr( + torch.cuda, + "get_device_properties", + lambda *_args, **_kwargs: type("Props", (), {"multi_processor_count": 256})(), + ) + assert ( + get_recommended_splits( + batch_size, + 1, + split_kv_blocks, + max_partitions, + max_context_length=context_length, + ) + == expected + ) + + +@pytest.mark.parametrize("split_kv_blocks,expected", [(2, 4), (16, 8)]) +def test_recommended_splits_keeps_large_batch_policy( + monkeypatch, split_kv_blocks, expected +): + if get_recommended_splits is None: + pytest.skip("FlyDSL is not available") + monkeypatch.setattr( + torch.cuda, + "get_device_properties", + lambda *_args, **_kwargs: type("Props", (), {"multi_processor_count": 256})(), + ) + assert ( + get_recommended_splits(200, 1, split_kv_blocks, max_context_length=200000) + == expected + ) + + +@pytest.mark.parametrize("num_partitions", [64, 74, 86, 103, 128, 171, 256]) +@pytest.mark.parametrize("trans_v", [False, True]) +@pytest.mark.parametrize("block_size", [16, 128]) +def test_mtp4_large_partition_mixed_contexts(num_partitions, trans_v, block_size): + """Long, sparse MTP contexts and empty partitions share the high-NP reducer.""" + _require_gpu() + _run_mtp4_fused_reference_case( + [0, 1, 3, 4, 255, 256, 257, 200003], + 1, + block_size, + trans_v, + num_partitions, + ) + + +@pytest.mark.parametrize("wide_kv_addressing", [False, True]) +@pytest.mark.parametrize("num_partitions", [1, 8]) +@pytest.mark.parametrize("trans_v", [False, True]) +@pytest.mark.parametrize("block_size", [16, 128]) +@pytest.mark.parametrize( + "lengths,num_kv_heads", + [ + pytest.param([4], 1, id="c4-hkv1"), + pytest.param([257], 1, id="c257-hkv1"), + pytest.param([259], 1, id="c259-hkv1"), + pytest.param([2305], 1, id="c2305-hkv1"), + pytest.param([2307], 1, id="c2307-hkv1"), + pytest.param( + [0, 1, 2, 3, 4, 255, 256, 257, 258, 259, 2303, 2304, 2305, 2306, 2307], + 1, + id="mixed-hkv1", + ), + pytest.param([257], 2, id="c257-hkv2"), + pytest.param([259], 2, id="c259-hkv2"), + ], +) +def test_mtp4_fused_phase_b_accuracy( + monkeypatch, + lengths, + num_kv_heads, + block_size, + trans_v, + num_partitions, + wide_kv_addressing, +): + """Cover fused MTP boundaries under both narrow and wide V-load schedules.""" + _require_gpu() + pa_decode_module = importlib.import_module("aiter.ops.flydsl.pa_decode") + compile_tile = pa_decode_module.compile_pa_decode_tile + + def compile_fused_mtp4(**kwargs): + kwargs["query_splits"] = 1 + # Exercise wide-address scheduling with sparse/empty/tail cases without + # allocating a multi-GiB cache for every boundary in this matrix. + kwargs["wide_kv_addressing"] = wide_kv_addressing + return compile_tile(**kwargs) + + monkeypatch.setattr(pa_decode_module, "compile_pa_decode_tile", compile_fused_mtp4) + _run_mtp4_fused_reference_case( + lengths, num_kv_heads, block_size, trans_v, num_partitions + ) + + def _positive_int(value): value = int(value) if value < 1: @@ -674,9 +1282,10 @@ def _parse_args(argv=None): parser.add_argument( "--max-partitions", type=int, - default=8, + default=None, help="""Upper clamp passed to get_recommended_splits (4..256). - Only used when --num-partitions is omitted.""", + By default, use context length and GPU occupancy, up to 256 partitions. + Set 8 to retain the legacy clamp. Only used without --num-partitions.""", ) parser.add_argument( "--num-partitions", @@ -695,7 +1304,10 @@ def _parse_args(argv=None): help="""KV scale layout: 0 for per-tensor, 1 for per-token.""", ) args = parser.parse_args(argv) - if not 4 <= args.max_partitions <= MAX_CONTEXT_PARTITIONS: + 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( count is not None and count > MAX_CONTEXT_PARTITIONS From d699fb2a97479b1068bf5783a93d9a38025f009f Mon Sep 17 00:00:00 2001 From: fsx950223 Date: Wed, 16 Sep 2026 07:43:50 +0000 Subject: [PATCH 48/66] fix(flydsl): satisfy lint for plan re-export --- aiter/ops/flydsl/pa_decode.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/aiter/ops/flydsl/pa_decode.py b/aiter/ops/flydsl/pa_decode.py index b3adc361046..1fce23c1f5e 100644 --- a/aiter/ops/flydsl/pa_decode.py +++ b/aiter/ops/flydsl/pa_decode.py @@ -40,7 +40,8 @@ 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, plan_pa_decode as plan_pa_decode +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 939d64a56544b3e0db89fdde12bd146ceadcf625 Mon Sep 17 00:00:00 2001 From: fsx950223 Date: Wed, 16 Sep 2026 10:01:00 +0000 Subject: [PATCH 49/66] refactor(flydsl): consolidate PA tuning constants --- aiter/ops/flydsl/kernels/pa_decode_kernel.py | 54 ++++++++++---------- 1 file changed, 27 insertions(+), 27 deletions(-) diff --git a/aiter/ops/flydsl/kernels/pa_decode_kernel.py b/aiter/ops/flydsl/kernels/pa_decode_kernel.py index 056a44c4ed3..49b13d9c10d 100644 --- a/aiter/ops/flydsl/kernels/pa_decode_kernel.py +++ b/aiter/ops/flydsl/kernels/pa_decode_kernel.py @@ -179,34 +179,31 @@ def compile_pa_decode_tile( # The wrapper selects earlier V loads for tuned gfx950 decode shapes. # Page 16 additionally delays next K until PV and uses IGLP to overlap # groups of V loads with QK MFMA instructions. - tune_page128 = ( + TUNED_V_PATH = ( is_gfx950 and head_dim == 128 - and block_size == 128 - and (trans_v or PER_TOKEN_M1) - and (not per_token_kv or PER_TOKEN_M1) - ) - tune_page16_vpipe = ( - is_gfx950 - and head_dim == 128 - and block_size == 16 - and (trans_v or PER_TOKEN_M1) - and (not per_token_kv or PER_TOKEN_M1) + and block_size in (16, 128) + and (PER_TOKEN_M1 or (trans_v and not per_token_kv)) ) - EARLY_V = prefetch_v and (tune_page128 or tune_page16_vpipe) and M_TILES == 1 - PAGE16_VPIPE = prefetch_v and tune_page16_vpipe and M_TILES == 1 + TUNE_PAGE128 = TUNED_V_PATH and block_size == 128 + EARLY_V = prefetch_v and TUNED_V_PATH and M_TILES == 1 + PAGE16_VPIPE = EARLY_V and block_size == 16 SCALES_BEFORE_CURRENT_V = WIDE_FP8_MFMA and ( PAGE16_VPIPE or (PER_TOKEN_M1 and not trans_v) ) # Apply score scales before the -inf mask on the tuned single-M-tile # paths. This removes the second mask without introducing -inf * 0. M1_SCALE_BEFORE_MASK = PAGE16_VPIPE or SCALAR_FP8_DECODE or PER_TOKEN_M1 - P_BUFFERS = M_TILES if MTP4_FUSED else 2 if tune_page128 and M_TILES == 3 else 1 + P_BUFFERS = M_TILES if MTP4_FUSED else 2 if TUNE_PAGE128 and M_TILES == 3 else 1 # PV layout: V=A, P=B -> output [head-dim (row), query-row (col=lane16)], # generalized over head_dim via the VHE_CHUNKS loop. NWARP = 4 # 4 waves / CTA - TOK_PER_WARP = 64 # tokens each warp owns per compute block (matches production KV_COMPUTE_BLOCK) - TILE_TOK = NWARP * TOK_PER_WARP # 256 tokens / compute block + 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" # 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 assert ( @@ -262,7 +259,7 @@ def compile_pa_decode_tile( sQscale_bytes = 0 if SCALAR_FP8_DECODE else ROWS_PADDED * f32 # Keep tuned cross-wave reduction rows 16-byte aligned for vector reads. # Other paths retain the original bank-conflict padding. - NWARP_PAD = NWARP if tune_page128 or PER_TOKEN_M1 or MTP4_FUSED else NWARP + 1 + 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 @@ -489,8 +486,7 @@ def _lds_load(byte_off, elem_ty, n): def _lds_store(byte_off, elem_ty, vec): fx.ptr_store(vec, _lds_ptr(byte_off, elem_ty)) - TOK_CHUNK = NWARP * MFMA_MNK # 64 - NCHUNK = TILE_TOK // TOK_CHUNK # 4 + NCHUNK = TOK_PER_WARP // MFMA_MNK # 4 if const_expr(per_token_kv): scale_load_width = NCHUNK if block_size >= 64 else 1 @@ -620,7 +616,7 @@ def _stage_kv_scale_to_lds(phys_vec, buf_off=0): def _load_scale_vec(base_off, a, buf_off=0): # This lane's 4 per-token scales for chunk `a` from an LDS scale region. - slot = (warp * TOK_CHUNK + a * c16 + rgroup * 4) * f32 + slot = (warp * TOK_PER_WARP + a * c16 + rgroup * 4) * f32 return _lds_load(base_off + buf_off + slot, fx.Float32, 4) def _load_kv_scale_vecs(a, buf_off=0): @@ -662,7 +658,7 @@ def _mfma_fp8(a_ops, b_ops, a_base, b_base, k_steps, acc): return acc # -- raw dwordx4 K load (A operand) -- - # token = warp*TOK_CHUNK + a*c16 + lane16 (the softmax mask and P-pack + # token = warp*TOK_PER_WARP + a*c16 + lane16 (the softmax mask and P-pack # write position below must encode this same formula). def _k_ops(phys, a): within_page_tok = (warp * TOK_PER_WARP + a * c16 + lane16) % block_size @@ -1015,7 +1011,7 @@ def _l_slot(m): ctx_thr = fx.Vector.from_elements( [ tile_valid.to(fx.Float32) - - fx.Int32(warp * TOK_CHUNK + rgroup * 4).to(fx.Float32) + - fx.Int32(warp * TOK_PER_WARP + rgroup * 4).to(fx.Float32) ], dtype=fx.Float32, ).broadcast_to(4) @@ -1098,7 +1094,9 @@ def _lmax_off_m(m): scale = scale_qk * fx.Float32(q_scale_vec[m]) n_valid_tile = (causal_bound[m] - tok0).to(fx.Float32) - base_tok_f = fx.Int32(warp * TOK_CHUNK + rgroup * 4).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) @@ -1265,7 +1263,7 @@ def _lmax_off_m(m): sP_off + m * MFMA_MNK * SP_ROW_BYTES + lane16 * SP_ROW_BYTES - + warp * TOK_CHUNK + + warp * TOK_PER_WARP + rgroup * 4 + a * (c16 // 4) * f32 ) @@ -1417,7 +1415,7 @@ def _lmax_off_m(m): p_off0 = ( p_base + lane16 * SP_ROW_BYTES - + warp * TOK_CHUNK + + warp * TOK_PER_WARP + rgroup * 4 ) # The NCHUNK P words scatter across the row at stride c16//4 @@ -1524,7 +1522,7 @@ def _lmax_off_m(m): else scale_qk * _ld1(sQscale_off, lane16) ) # per-qhead positive score scale n_valid_tile = (causal_bound[0] - tok0).to(fx.Float32) - base_tok_f = fx.Int32(warp * TOK_CHUNK + rgroup * 4).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) @@ -1693,7 +1691,9 @@ def _lmax_off_m(m): 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_CHUNK + rgroup * 4 + p_off0 = ( + sP_off + lane16 * SP_ROW_BYTES + warp * TOK_PER_WARP + rgroup * 4 + ) # NCHUNK P words scatter at stride c16//4 i32 (see phase-split). for a in range_constexpr(NCHUNK): _lds_store( From ed84952322157f409d6f732e6eb4b0d4da6a67e0 Mon Sep 17 00:00:00 2001 From: fsx950223 Date: Thu, 17 Sep 2026 05:32:00 +0000 Subject: [PATCH 50/66] perf(flydsl): unify per-token V prefetch policy --- aiter/ops/flydsl/pa_decode.py | 129 ++++++++++++++---------------- op_tests/test_flydsl_pa_decode.py | 87 +++++++++++++++++++- 2 files changed, 148 insertions(+), 68 deletions(-) diff --git a/aiter/ops/flydsl/pa_decode.py b/aiter/ops/flydsl/pa_decode.py index 1fce23c1f5e..ff2326e5577 100644 --- a/aiter/ops/flydsl/pa_decode.py +++ b/aiter/ops/flydsl/pa_decode.py @@ -99,22 +99,6 @@ def get_recommended_splits( return max(4, min(n, max_partitions)) -def _workgroup_count_enables_v_prefetch( - num_sequences: int, - num_kv_heads: int, - num_partitions: int, - num_compute_units: int, -) -> bool: - """Return whether the launch is in the tuned one-to-two-CTA-per-CU window. - - This result is a kernel specialization and therefore part of the compile - cache key. Applications serving shapes on both sides of this interval must - materialize both variants before graph capture or latency-sensitive use. - """ - workgroups = num_sequences * num_kv_heads * num_partitions - return num_compute_units < workgroups <= 2 * num_compute_units - - def _flydsl_pointer_dtype(dtype: torch.dtype): return { torch.float32: fx.Float32, @@ -567,68 +551,79 @@ def normalize_scale(scale, name): # if either does. wide_kv_addressing = max(key_cache.numel(), value_cache.numel()) >= 2**31 - # Scalar page-16 decode keeps its non-prefetch schedule. Per-token paths - # select their tuned pipelines below without environment overrides. - prefetch_v = False - # Early V loads help the measured one-to-two-workgroup-per-CU regime. - # Keep other scalar page-128 grids on the non-prefetch schedule. - if ( - arch == "gfx950" - and head_dim == 128 + dense_partition_workgroups = num_seqs * num_kv_heads * num_partitions + gfx950_d128 = arch == "gfx950" and head_dim == 128 + scalar_prefetch_candidate = ( + gfx950_d128 and block_size == 128 and trans_v and not per_token_kv and query_length * query_group_size <= 16 - ): - num_cus = torch.cuda.get_device_properties(dev).multi_processor_count - prefetch_v = _workgroup_count_enables_v_prefetch( - num_seqs, - num_kv_heads, - num_partitions, - num_cus, - ) - - # Per-token decode benefits from overlapping V with QK. Page 16 also - # delays next K until PV; page 128 stages next scales before issuing K. - # With transposed page-128 V, more than one workgroup per CU already hides - # latency: retain the lower-register schedule in that regime. - if ( - arch == "gfx950" + ) + tuned_per_token_shape = ( + gfx950_d128 and per_token_kv - and head_dim == 128 - and query.dtype == torch.bfloat16 - and num_kv_heads == 1 + and query_dtype == "bf16" + and block_size in (16, 128) + ) + # Hkv2 static decode is covered by the same workgroup/layout policy. + # Planned Hkv2 and larger KV-head counts retain their existing schedule. + tuned_m1_heads = num_kv_heads == 1 or (work_plan is None and num_kv_heads == 2) + per_token_single_query = ( + tuned_per_token_shape + and tuned_m1_heads and query_length == 1 and query_group_size in (8, 16) - and block_size in (16, 128) - ): - if block_size == 16 or not trans_v: - prefetch_v = True - else: - num_cus = torch.cuda.get_device_properties(dev).multi_processor_count - prefetch_v = num_seqs * num_kv_heads * num_partitions <= num_cus - - # Split small MTP2/MTP4 grids into one query per CTA, retaining the original - # partial layout and partition count. Larger grids amortize KV loads by - # processing all queries together. MTP3 stays fused because the kernel's - # supported split counts (1, 2, 4) must divide query_length. - query_splits = 1 - if ( + ) + query_split_candidate = ( work_plan is None - and arch == "gfx950" - and per_token_kv - and head_dim == 128 - and query.dtype == torch.bfloat16 + and tuned_per_token_shape and num_kv_heads == 1 and query_length in (2, 4) and query_group_size == 16 - and block_size in (16, 128) - ): - num_cus = torch.cuda.get_device_properties(dev).multi_processor_count - base_workgroups = num_seqs * num_kv_heads * num_partitions - if query_length * base_workgroups <= 2 * num_cus: - query_splits = query_length - prefetch_v = True + ) + needs_num_cus = ( + scalar_prefetch_candidate + or query_split_candidate + or (per_token_single_query and block_size == 128 and trans_v) + ) + num_cus = ( + torch.cuda.get_device_properties(dev).multi_processor_count + if needs_num_cus + else 0 + ) + + # Split small MTP2/MTP4 grids into one query per CTA. MTP3 stays fused + # because the supported split counts (1, 2, 4) must divide query_length. + query_splits = ( + query_length + if query_split_candidate + and query_length * dense_partition_workgroups <= 2 * num_cus + else 1 + ) + queries_per_cta = query_length // query_splits + total_workgroups = dense_partition_workgroups * query_splits + per_token_m1 = ( + tuned_per_token_shape + and queries_per_cta == 1 + and query_group_size in (8, 16) + and (query_splits > 1 or per_token_single_query) + ) + + # Use one policy for the single-M-tile per-token pipelines. Page16 and + # plain V tolerate high occupancy; unsplit transposed page128 retains its + # lower-register schedule above one workgroup per CU. + per_token_prefetch = per_token_m1 and ( + query_splits > 1 + or block_size == 16 + or not trans_v + or total_workgroups <= num_cus + ) + scalar_prefetch = ( + scalar_prefetch_candidate + and num_cus < dense_partition_workgroups <= 2 * num_cus + ) + prefetch_v = scalar_prefetch or per_token_prefetch with torch.cuda.device(dev): compiled = compile_pa_decode_tile( diff --git a/op_tests/test_flydsl_pa_decode.py b/op_tests/test_flydsl_pa_decode.py index 6207d339b1c..e994e0416c3 100644 --- a/op_tests/test_flydsl_pa_decode.py +++ b/op_tests/test_flydsl_pa_decode.py @@ -699,6 +699,74 @@ def run_pa_decode_tile_case( }, id="decode-reduce-page128-per-token", ), + pytest.param( + { + "batch_size": 16, + "num_query_heads": 16, + "num_kv_heads": 2, + "head_dim": 128, + "context_length": 257, + "block_size": 128, + "dtype": dtypes.bf16, + "trans_v": True, + "per_token": True, + "query_length": 1, + "num_partitions": 8, + "gfx950_prefetch_v": "occupancy", + }, + id="decode-hkv2-one-wg-per-cu-prefetch", + ), + pytest.param( + { + "batch_size": 32, + "num_query_heads": 16, + "num_kv_heads": 2, + "head_dim": 128, + "context_length": 257, + "block_size": 128, + "dtype": dtypes.bf16, + "trans_v": True, + "per_token": True, + "query_length": 1, + "num_partitions": 8, + "gfx950_prefetch_v": "occupancy", + }, + id="decode-hkv2-two-wg-per-cu-no-prefetch", + ), + pytest.param( + { + "batch_size": 32, + "num_query_heads": 16, + "num_kv_heads": 2, + "head_dim": 128, + "context_length": 257, + "block_size": 16, + "dtype": dtypes.bf16, + "trans_v": True, + "per_token": True, + "query_length": 1, + "num_partitions": 8, + "gfx950_prefetch_v": True, + }, + id="decode-hkv2-page16-prefetch", + ), + pytest.param( + { + "batch_size": 32, + "num_query_heads": 32, + "num_kv_heads": 2, + "head_dim": 128, + "context_length": 257, + "block_size": 128, + "dtype": dtypes.bf16, + "trans_v": False, + "per_token": True, + "query_length": 1, + "num_partitions": 8, + "gfx950_prefetch_v": True, + }, + id="decode-hkv2-plain-v-prefetch", + ), pytest.param( { "batch_size": 2, @@ -713,6 +781,7 @@ def run_pa_decode_tile_case( "query_length": 2, "num_partitions": 3, "gfx950_query_splits": 2, + "gfx950_prefetch_v": True, }, id="mtp2-query-split-page128-per-token", ), @@ -747,6 +816,7 @@ def run_pa_decode_tile_case( "query_length": 4, "num_partitions": 3, "gfx950_query_splits": 4, + "gfx950_prefetch_v": True, }, id="mtp4-query-split-page16-per-token", ), @@ -826,14 +896,17 @@ def test_pa_decode(case, monkeypatch): _require_gpu() case = case.copy() expected_query_splits = case.pop("gfx950_query_splits", None) + expected_prefetch_v = case.pop("gfx950_prefetch_v", None) selected_query_splits = [] + selected_prefetch_v = [] - if expected_query_splits is not None: + if expected_query_splits is not None or expected_prefetch_v is not None: module = importlib.import_module("aiter.ops.flydsl.pa_decode") compile_tile = module.compile_pa_decode_tile def capture_compile(**kwargs): selected_query_splits.append(kwargs["query_splits"]) + selected_prefetch_v.append(kwargs["prefetch_v"]) return compile_tile(**kwargs) monkeypatch.setattr(module, "compile_pa_decode_tile", capture_compile) @@ -872,6 +945,18 @@ def run_once(fn, *args, **kwargs): expected = expected_query_splits if get_gfx_runtime() == "gfx950" else 1 assert selected_query_splits assert set(selected_query_splits) == {expected} + if expected_prefetch_v is not None: + if expected_prefetch_v == "occupancy": + workgroups = case["batch_size"] * case["num_kv_heads"] * expected_partitions + expected_prefetch_v = ( + workgroups + <= torch.cuda.get_device_properties( + torch.cuda.current_device() + ).multi_processor_count + ) + expected = expected_prefetch_v if get_gfx_runtime() == "gfx950" else False + assert selected_prefetch_v + assert set(selected_prefetch_v) == {expected} def _assert_plan(plan, lengths): From 2f295c6dcc4fb82bc2c4837061f7f3001d6f772b Mon Sep 17 00:00:00 2001 From: fsx950223 Date: Thu, 17 Sep 2026 06:18:57 +0000 Subject: [PATCH 51/66] feat(flydsl): add sliding windows and attention sinks to PA decode --- aiter/ops/flydsl/kernels/pa_decode_kernel.py | 153 +++- aiter/ops/flydsl/kernels/pa_decode_plan.py | 57 +- aiter/ops/flydsl/kernels/pa_decode_reduce.py | 44 +- aiter/ops/flydsl/pa_decode.py | 52 +- docs/flydsl_pa_decode_plan.md | 80 +- op_tests/test_flydsl_pa_decode.py | 896 ++++++++++++++++++- 6 files changed, 1182 insertions(+), 100 deletions(-) diff --git a/aiter/ops/flydsl/kernels/pa_decode_kernel.py b/aiter/ops/flydsl/kernels/pa_decode_kernel.py index 49b13d9c10d..f86fc719c73 100644 --- a/aiter/ops/flydsl/kernels/pa_decode_kernel.py +++ b/aiter/ops/flydsl/kernels/pa_decode_kernel.py @@ -75,6 +75,9 @@ def compile_pa_decode_tile( prefetch_v: bool = False, query_splits: int = 1, use_work_plan: bool = False, + sliding_window: int = 0, + use_sinks: bool = False, + sink_dtype_str: str = "f32", ): """Build the tile-programming PA-decode kernel + launch wrapper. @@ -86,6 +89,12 @@ def compile_pa_decode_tile( output layout remains unsplit so the existing FlyDSL reducer can combine context partitions without an intermediate transpose. + A positive ``sliding_window`` includes the query token itself. Tile ranges + cover the union of the MTP windows, and scores are masked per query row. + ``use_sinks`` adds a per-query-head zero-value attention logit in the final + epilogue for direct NP=1 output. Partitioned/planned execution instead adds + it exactly once in the reducer, leaving partial statistics unchanged. + Tokens at/after ``context_len`` still take part in the PV matmul with a zero probability, so their V bytes must be finite: the MFMA does not treat a zero operand as an annihilator (``0 * NaN == NaN``). Whole pages past the @@ -101,6 +110,9 @@ def compile_pa_decode_tile( neutral at ``block_size=16``. The wrapper turns it on from the cache size. """ is_gfx950 = "gfx95" in get_rocm_arch() + # Context lengths are int32, so larger windows have identical visibility. + # Bound the device constant while retaining the caller's value in the plan. + 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 @@ -238,6 +250,14 @@ def compile_pa_decode_tile( 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 @@ -300,6 +320,7 @@ def _pa_decode_tile_task( context_lengths_ptr: fx.Pointer, # [num_seqs] key_scale_ptr: fx.Pointer, # [1] per-tensor OR [num_blocks, num_kv_heads, block_size] per-token value_scale_ptr: fx.Pointer, # same shape as key_scale_ptr + sinks_ptr: fx.Pointer, # [num_q_heads], used only for direct NP=1 output max_blocks_per_seq: fx.Int32, stride_ks_block: fx.Int32, stride_ks_head: fx.Int32, @@ -333,6 +354,8 @@ def _pa_decode_tile_task( 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(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): @@ -458,8 +481,19 @@ def _v_load16(byte_off): part_end = planned_end else: num_tiles = (context_len + TILE_TOK - 1) // TILE_TOK - tiles_per_part = (num_tiles + NP - 1) // NP + active_tiles = num_tiles + if const_expr(sliding_window > 0): + window_start = ( + fx.Int64(context_len) - (query_length - 1) - sliding_window + ) + first_tile = fx.Int32( + (window_start > 0).select(window_start, 0) // TILE_TOK + ) + active_tiles = num_tiles - first_tile + tiles_per_part = (active_tiles + NP - 1) // NP part_start = part * tiles_per_part + if const_expr(sliding_window > 0): + part_start = first_tile + part_start part_end_raw = part_start + tiles_per_part part_end = (part_end_raw < num_tiles).select(part_end_raw, num_tiles) @@ -879,6 +913,13 @@ def _quant_q_row(m, q_row_off, q_units): fx.Vector.from_elements([float(a * c16 + 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 + # P.V is loop-tiled over head-dim (like production's VHELOOP): each # step computes O[:, vh*VHE_SIZE:+VHE_SIZE] instead of materializing # the full [16, head_dim] at once. @@ -1015,10 +1056,30 @@ def _l_slot(m): ], dtype=fx.Float32, ).broadcast_to(4) + window_scale_thr = None + if const_expr(sliding_window > 0): + # Scales outside every query's window must not shrink + # visible probabilities to zero during FP8 normalization. + # Keep this bound query-independent for shared MTP scales. + first_visible = ( + fx.Int64(context_len) + - (query_length - 1) + - sliding_window + - fx.Int64(tok0) + ).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, zero=zero4_scale): - return (_ct[a] < thr).select(vec, zero) + 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) # V is independent of QK/softmax. Load it early for the tuned # decode path and reuse it across M-tiles in the phase-split path. @@ -1100,6 +1161,16 @@ def _lmax_off_m(m): 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 in integer space first so a large window + # cannot round its left edge before the comparison. + first_valid_tile = ( + fx.Int64(causal_bound[m]) - sliding_window - fx.Int64(tok0) + ).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) # Fold the per-row score scale in BEFORE the -inf mask. A @@ -1121,8 +1192,9 @@ def _lmax_off_m(m): ] if const_expr(MTP4_FUSED): - # All four causal rows fully cover non-tail tiles. - # Keep one uniform branch around the sixteen masks. + # All four causal windows fully cover interior tiles. + # Keep one uniform branch around the sixteen masks, + # including both the left window edge and causal tail. masked_all = fx.Vector.from_elements( [ scaled_frags[a][r] @@ -1131,10 +1203,17 @@ def _lmax_off_m(m): ], dtype=fx.Float32, ) - if tile_valid < TILE_TOK + query_length - 1: + needs_mask = tile_valid < TILE_TOK + query_length - 1 + if const_expr(sliding_window > 0): + needs_mask = needs_mask | ( + tok0 < context_len - sliding_window + ) + if needs_mask: masked_all = fx.Vector.from_elements( [ - (_ct[a] < thr).select(scaled_frags[a], neg4)[r] + _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) ], @@ -1152,7 +1231,9 @@ def _lmax_off_m(m): ] else: masked_chunks = [ - (_ct[a] < thr).select(scaled_frags[a], neg4) + _score_mask(a, thr, window_thr).select( + scaled_frags[a], neg4 + ) for a in range_constexpr(NCHUNK) ] @@ -1526,6 +1607,14 @@ def _lmax_off_m(m): 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 = ( + fx.Int64(causal_bound[0]) - sliding_window - fx.Int64(tok0) + ).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, @@ -1554,14 +1643,18 @@ def _lmax_off_m(m): if const_expr(M1_SCALE_BEFORE_MASK): scaled_frag = scaled_frag * scale_b masked_chunks.append( - (_ct[a] < thr).select(scaled_frag, neg4) + _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 = [ - (_ct[a] < thr).select(frag_Ss[a] * scale_b, neg4) + _score_mask(a, thr, window_thr).select( + frag_Ss[a] * scale_b, neg4 + ) for a in range_constexpr(NCHUNK) ] else: @@ -1569,7 +1662,7 @@ def _lmax_off_m(m): # Reused in pass 2 below, halving the mask instruction count. if const_expr(not M1_SCALE_BEFORE_MASK): masked_chunks = [ - (_ct[a] < thr).select(scaled_frags[a], neg4) + _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 @@ -1598,8 +1691,8 @@ def _lmax_off_m(m): # Slots at/after context_len are uninitialised, so an arbitrarily # large one there would shrink every valid probability to 0 in the # fp8 conversion below -- restrict the reduction to real tokens. - # (Causally-masked tokens inside the context keep real scales and - # stay in, which also keeps pv_max independent of the query row.) + # Causally-masked tokens within the MTP window union keep real + # scales and stay in, keeping pv_max independent of the query row. if const_expr(per_token_kv): pv_max = fx.Float32(0.0) for a in range_constexpr(NCHUNK): @@ -1756,9 +1849,36 @@ def _lmax_off_m(m): 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): + # Keep the sink out of the online Q/P quantization. It only + # changes the final denominator, with zero numerator mass. + # Compare in natural-logit units before multiplying by LOG2E, + # so even a very large finite f32 sink cannot overflow here. + 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): @@ -1770,9 +1890,6 @@ def _lmax_off_m(m): o_scale_b = fx.Vector.from_elements( [o_scale], dtype=fx.Float32 ).broadcast_to(OP_ELEMS) - qi_e = row // query_group_size - gs_head_e = row - qi_e * query_group_size - qh = kv_h * query_group_size + gs_head_e def _emit(o_norm, sub, query_idx, query_head): if const_expr(NP == 1 and not use_work_plan): @@ -1824,6 +1941,7 @@ def pa_decode_tile_kernel( 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, @@ -1847,6 +1965,7 @@ def _run_task(seq, start, end, context): context_lengths_ptr, key_scale_ptr, value_scale_ptr, + sinks_ptr, max_blocks_per_seq, stride_ks_block, stride_ks_head, @@ -1894,6 +2013,7 @@ def pa_decode_tile_launch( context_lengths: fx.Pointer, key_scale: fx.Pointer, # [1] per-tensor OR [num_blocks, num_kv_heads, block_size] per-token value_scale: fx.Pointer, # same shape as key_scale + sinks: fx.Pointer, max_blocks_per_seq: fx.Int32, num_seqs: fx.Int32, num_kv_heads: fx.Int32, @@ -1923,6 +2043,7 @@ def pa_decode_tile_launch( context_lengths, key_scale, value_scale, + sinks, max_blocks_per_seq, stride_ks_block, stride_ks_head, diff --git a/aiter/ops/flydsl/kernels/pa_decode_plan.py b/aiter/ops/flydsl/kernels/pa_decode_plan.py index fdd10b4abfd..836f02b9c1b 100644 --- a/aiter/ops/flydsl/kernels/pa_decode_plan.py +++ b/aiter/ops/flydsl/kernels/pa_decode_plan.py @@ -20,13 +20,20 @@ def _plan_pa_decode( 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.load(lengths + b, b < B, other=0) - tiles = tl.where(b < B, tl.maximum(ctx, 0).to(tl.int64) + 255, 0) // 256 + 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: + # Fused MTP queries have different left edges. Schedule their union, + # rounding down to an absolute tile so page/scaling offsets stay valid. + 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) @@ -40,15 +47,21 @@ def _plan_pa_decode( 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_tiles = (tl.maximum(seq_ctx, 0).to(tl.int64) + 255) // 256 + 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 = part.to(tl.int64) * seq_tiles // tl.maximum(count, 1) - end = (part.to(tl.int64) + 1) * seq_tiles // tl.maximum(count, 1) + 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) @@ -69,7 +82,8 @@ def _plan_pa_decode( class PADecodePlan: """Reusable GPU metadata; refresh it whenever context lengths change. - ``work_info`` is [capacity, 4]: sequence, first/last 256-token tile, length. + ``work_info`` is [capacity, 4]: sequence, first/last absolute 256-token tile, + original context length. Sliding windows cover the union of MTP queries. ``reduce_info`` is [batch, 2]: first packed task, actual task count. Scratch uses [KV heads, capacity, query rows (, head dim)]. """ @@ -78,6 +92,8 @@ class PADecodePlan: reduce_info: torch.Tensor num_kv_heads: int max_partitions: int + sliding_window: int = 0 + query_length: int = 1 @property def capacity(self) -> int: @@ -88,6 +104,8 @@ def validate(self, batch_size: int, num_kv_heads: int, device: torch.device): raise ValueError("plan KV head count does not match the cache") if not 1 <= self.max_partitions <= MAX_CONTEXT_PARTITIONS: raise ValueError("invalid plan max_partitions") + 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: @@ -107,6 +125,8 @@ def plan_pa_decode( *, max_partitions: int = MAX_CONTEXT_PARTITIONS, workgroup_budget: int | None = None, + sliding_window: int = 0, + query_length: int = 1, plan: PADecodePlan | None = None, ) -> PADecodePlan: """Build/update a plan on the current stream without GPU-to-CPU readback. @@ -115,7 +135,22 @@ def plan_pa_decode( same metadata in place. Include this refresh in end-to-end measurements. The budget counts CTAs over all KV heads, with fused query positions. This is opt-in: uniform or short-context workloads may favor static splits. + + A positive ``sliding_window`` counts visible tokens including the query's + own position; 0 and -1 disable it. Context lengths include the MTP tokens, + so the planned range is the union of ``query_length`` causal windows. + Pass the same window and (when enabled) query length to ``pa_decode`` and + when refreshing the plan. Dense plans remain independent of 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(): @@ -144,12 +179,18 @@ def plan_pa_decode( 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,)]( @@ -159,6 +200,10 @@ def plan_pa_decode( batch, plan.capacity, plan.max_partitions, + # Larger windows cover every int32 context. Dense planning is + # query-length independent, including its compilation cache key. + 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, diff --git a/aiter/ops/flydsl/kernels/pa_decode_reduce.py b/aiter/ops/flydsl/kernels/pa_decode_reduce.py index aa602493436..0a2cbf34ec4 100644 --- a/aiter/ops/flydsl/kernels/pa_decode_reduce.py +++ b/aiter/ops/flydsl/kernels/pa_decode_reduce.py @@ -78,6 +78,8 @@ def compile_pa_decode_ps_reduce( thread. For D=128 and larger counts, a 2-D workgroup materializes weights once in LDS and splits each output element's partition chain over several waves. Other head sizes retain the register-only lane-striped fallback. + A sink is a per-query-head zero-value logit: include it in the shared max + for stability and add its mass to the denominator once after summing KV. """ _validate_pa_decode_ps_reduce_config( max_context_partition_num=max_context_partition_num, @@ -210,6 +212,12 @@ def _wave_reduce_sum(value): ) 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") + if fx.const_expr(use_parallel_lds): # One wave materializes the normalized partition weights once. # All output waves then reuse those weights from LDS and split the @@ -259,6 +267,9 @@ def _wave_reduce_sum(value): 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 @@ -289,14 +300,7 @@ def _wave_reduce_sum(value): global_exp_sum = _wave_reduce_sum(lane_exp_sum) if fx.const_expr(use_sinks): - sink_value = fx.Float32(sink_token[kv_head_idx * c_qgs + group_idx]) - sink_scale = (global_max > neg_inf).select( - fx.exp2( - (sink_value - safe_global_max) * c_log2e, - fastmath="fast", - ), - zero_f, - ) + 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 @@ -429,6 +433,9 @@ def _wave_reduce_sum(value): 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) if fx.const_expr(use_work_plan): part_scale = zero_f @@ -445,14 +452,7 @@ def _wave_reduce_sum(value): scaled_sum = part_sum * part_scale global_exp_sum = _wave_reduce_sum(scaled_sum) if fx.const_expr(use_sinks): - sink_value = fx.Float32(sink_token[kv_head_idx * c_qgs + group_idx]) - sink_scale = (global_max > neg_inf).select( - fx.exp2( - (sink_value - safe_global_max) * c_log2e, - fastmath="fast", - ), - zero_f, - ) + 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 @@ -552,6 +552,9 @@ def _wave_reduce_sum(value): 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 @@ -578,14 +581,7 @@ def _wave_reduce_sum(value): global_exp_sum = _wave_reduce_sum(lane_exp_sum) if fx.const_expr(use_sinks): - sink_value = fx.Float32(sink_token[kv_head_idx * c_qgs + group_idx]) - sink_scale = (global_max > neg_inf).select( - fx.exp2( - (sink_value - safe_global_max) * c_log2e, - fastmath="fast", - ), - zero_f, - ) + 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 diff --git a/aiter/ops/flydsl/pa_decode.py b/aiter/ops/flydsl/pa_decode.py index ff2326e5577..c525b220e6b 100644 --- a/aiter/ops/flydsl/pa_decode.py +++ b/aiter/ops/flydsl/pa_decode.py @@ -221,9 +221,14 @@ def pa_decode( Sparse attention uses caller-prepared block tables and selected context lengths. Each independently selected MTP query must have its own table row and use query_length=1; query_length>1 applies dense causal masking. - ALiBi, attention sinks, sliding-window attention, and externally quantized - FP8 queries are not supported. The ``ps=False`` partitioning policy is also - not supported. + A positive ``sliding_window`` limits each query to that many causal tokens, + including its own position; 0 and -1 disable the window. Optional ``sinks`` + is a contiguous [num_query_heads] BF16/FP16/FP32 tensor on the query device. + Each entry is an unscaled, zero-value attention logit shared across batch + and MTP positions. It contributes to the denominator once, independently + of the window; -inf disables a head's sink and +inf suppresses its output. + ALiBi and externally quantized FP8 queries are not supported. The ``ps=False`` + partitioning policy is also not supported. ``context_lengths`` and ``block_tables`` are GPU-resident, so their values are not inspected here (which would synchronize the device). Callers must @@ -235,6 +240,8 @@ def pa_decode( ``work_plan`` opts into GPU-planned variable partition counts. Build or refresh it with ``plan_pa_decode`` on the current stream after updating lengths. Its ``max_partitions`` must equal ``max_context_partition_num``. + Its window must match ``sliding_window``, and a windowed plan must be built + with the same ``query_length`` so it covers every MTP query's window. Planned scratch is packed as [KV heads, plan.capacity, query rows (, D)]; the static API's dense per-sequence scratch layout remains unchanged. """ @@ -249,14 +256,17 @@ def pa_decode( ) if alibi_slopes is not None: raise NotImplementedError("pa_decode does not support ALiBi") - if sinks is not None: - raise NotImplementedError("pa_decode does not support attention sinks") - if sliding_window not in (0, -1): - raise NotImplementedError("pa_decode does not support sliding-window attention") + 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(ps, bool): raise TypeError(f"ps must be a bool, got {type(ps).__name__}") if not ps: raise NotImplementedError("pa_decode does not support ps=False") + 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 not 1 <= max_context_partition_num <= MAX_CONTEXT_PARTITIONS: @@ -442,6 +452,17 @@ def pa_decode( ) 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), @@ -541,6 +562,10 @@ def normalize_scale(scale, name): 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 @@ -625,6 +650,9 @@ def normalize_scale(scale, name): ) prefetch_v = scalar_prefetch or per_token_prefetch + # Partial outputs must exclude the sink: reduction adds it once across all + # partitions. Keep direct NP=1 output fused, without extra scratch/launches. + 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, @@ -640,6 +668,9 @@ def normalize_scale(scale, name): prefetch_v=prefetch_v, query_splits=query_splits, use_work_plan=work_plan is not 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: @@ -722,6 +753,11 @@ def normalize_scale(scale, name): 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), @@ -748,7 +784,7 @@ def normalize_scale(scale, name): psum, pmax, pout, - None, + sinks, output_5d.stride(0), output_5d.stride(1), output_5d.stride(2), diff --git a/docs/flydsl_pa_decode_plan.md b/docs/flydsl_pa_decode_plan.md index ec86f4e7e90..2cae1d4e230 100644 --- a/docs/flydsl_pa_decode_plan.md +++ b/docs/flydsl_pa_decode_plan.md @@ -6,9 +6,16 @@ such as one long context mixed with many short contexts. It is opt-in: a plannin launch and packed reduction can cost more than static splits for uniform or short contexts. Compare **plan + attention + reduction** against the static path. -The attention math, per-token scaling and causal MTP masking are shared with the -static FlyDSL kernel. The planner is a Triton GPU kernel, following the existing -variable-context scheduling approach in `mqa_logits/pa_mqa_logits_fp4.py`. +The attention math, per-token scaling, causal MTP masking and sliding windows are +shared with the static FlyDSL kernel. The planner is a Triton GPU kernel, following +the existing variable-context scheduling approach in +`mqa_logits/pa_mqa_logits_fp4.py`. + +A positive `sliding_window=W` retains at most W causal tokens **including the +query's own position**; W=1 attends only to that token. Values 0 and -1 disable the +window. Context lengths include the MTP query tokens: for context length C, query +length Q and query position p, the visible range is +`[max(0, C - Q + 1 + p - W), max(0, C - Q + 1 + p))`. ## Scheduling and storage @@ -19,16 +26,23 @@ a target across all KV heads; capacity is rounded up to at least one slot per request/head. Clamping can leave unused slots. Physical page size is independent of this scheduling tile size. +With a window, only tiles intersecting the union of the Q query windows count +toward the allocation: from `max(0, C - (Q - 1) - W) // 256` through +`ceil(C / 256)` (exclusive). Unaligned windows can touch an extra tile. The +attention kernel applies the exact per-query mask within those tiles; it does +not load the skipped prefix tiles. + The planner writes two contiguous int32 CUDA tensors: -- `work_info[capacity, 4]`: request index, first tile, exclusive last tile, context length. +- `work_info[capacity, 4]`: request index, absolute first tile, exclusive absolute + last tile, original context length. - `reduce_info[batch, 2]`: first packed task slot and actual partition count. The launch grid has a fixed capacity for graph capture. Padded work records are zeroed and skipped by the attention kernel, leaving their scratch slots untouched. -Each real task processes one contiguous KV range; ranges cover each request exactly -once. Reduction reads only the request's valid partitions and produces zero for -empty contexts. No GPU-to-CPU readback, dynamic allocation, or CPU length +Each real task processes one contiguous KV range; ranges cover each request's +active tiles exactly once. Reduction reads only the valid partitions and produces +zero for empty contexts. No GPU-to-CPU readback, dynamic allocation, or CPU length inspection is required when refreshing a preallocated plan. Planned scratch is packed as `[num_kv_heads, capacity, query_rows]` for max/sum and @@ -36,6 +50,26 @@ Planned scratch is packed as `[num_kv_heads, capacity, query_rows]` for max/sum `query_rows = query_length * num_query_heads // num_kv_heads`. The static API continues to use its dense per-request scratch layout. +## Attention sinks + +Both static and planned `pa_decode` accept optional `sinks`: a contiguous +`[num_query_heads]` BF16, FP16 or FP32 tensor on the query device. Each entry is +a zero-value attention logit shared across sequences and MTP query positions. +It is not multiplied by `softmax_scale` and remains visible independently of +the sliding window. For head h, the output is mathematically +`sum(exp(score_i) * V_i) / (sum(exp(score_i)) + exp(sinks[h]))`. + +The sink participates in the stable softmax maximum and contributes to the +denominator exactly once, not once per partition. Static NP=1 handles it in +the compute epilogue without an extra launch; multi-partition and planned calls +add it in the final reducer. KV partials and FP8 Q/P quantization are unchanged. +Empty contexts or query rows with no visible KV tokens still produce zero. +`-inf` disables an individual head's sink; `+inf` produces zero for that head. + +Sinks do not change work ownership, so they require no planner arguments or new +metadata. Their values can be updated in place between graph replays without +rebuilding a plan, provided the context lengths remain unchanged. + ## Usage Use the Python FlyDSL entry point for planned execution. The shared @@ -48,23 +82,32 @@ from aiter.ops.flydsl.pa_decode import pa_decode, plan_pa_decode # query/output: [batch * query_length, num_query_heads, head_dim] # K/V cache layouts and scales are the same as for static pa_decode. num_kv_heads = key_cache.shape[1] -plan = plan_pa_decode(context_lengths, num_kv_heads, max_partitions=256) +sliding_window = 4096 # Use 0 (the default) or -1 for full causal attention. +plan = plan_pa_decode( + context_lengths, num_kv_heads, max_partitions=256, + sliding_window=sliding_window, query_length=query_length, +) rows = query_length * query.shape[1] // num_kv_heads shape = (num_kv_heads, plan.capacity, rows) exp_sums = torch.empty(shape, dtype=torch.float32, device=query.device) max_logits = torch.empty_like(exp_sums) partials = torch.empty((*shape, query.shape[2]), dtype=query.dtype, device=query.device) +sinks = torch.zeros(query.shape[1], dtype=torch.float32, device=query.device) +# Use sinks=None for attention without the extra normalization logit. def step(): # Refresh after changing lengths, on the same stream as attention. - plan_pa_decode(context_lengths, num_kv_heads, max_partitions=256, plan=plan) + plan_pa_decode( + context_lengths, num_kv_heads, max_partitions=256, plan=plan, + sliding_window=sliding_window, query_length=query_length, + ) pa_decode( output, query, key_cache, value_cache, context_lengths, block_tables, softmax_scale, query_length, plan.max_partitions, compute_type=key_cache.dtype, key_scale=key_scale, value_scale=value_scale, exp_sums=exp_sums, max_logits=max_logits, temporary_output=partials, - work_plan=plan, + work_plan=plan, sliding_window=sliding_window, sinks=sinks, ) step() # Warm up compilation before graph capture. @@ -76,8 +119,12 @@ with torch.cuda.graph(graph): Reuse an already built plan without refreshing only when lengths are unchanged. Changing batch size, KV-head count, the partition limit or launch capacity requires -a compatible new plan and scratch buffers. Metadata values should be produced by -the planner, not edited independently. The planner accepts contiguous int32 CUDA +a compatible new plan and scratch buffers. A plan's window is fixed, as is its +query length when the window is enabled; changing either requires a new plan. +Dense plans are independent of query length. The attention call checks these +settings to avoid silently using a plan that omits visible tokens. +Metadata values should be produced by the planner, not edited independently. +The planner accepts contiguous int32 CUDA length vectors with 1 through 4096 requests. Lengths and block-table indices must satisfy the same validity requirements as static attention. @@ -87,6 +134,11 @@ satisfy the same validity requirements as static attention. python -m pytest -q op_tests/test_flydsl_pa_decode.py ``` -Tests cover exact work ownership, sparse pages, empty requests, causal tails, +Tests cover exact absolute-tile ownership, sparse pages, empty requests, causal tails, non-power-of-two partition caps, multiple KV heads, supported reducer branches, -NaN-filled scratch, and graph replay after changing lengths. +NaN-filled scratch, and graph replay after changing lengths. Window tests include +W=1, non-aligned boundaries, windows larger than the context, QL1/2/3/4, per-token +and per-tensor scales, static/query-split execution and incompatible plan reuse. +Sink tests cover per-head logits, direct output and all reducer paths, mixed +input/sink dtypes, large and infinite logits, window composition, and graph +replay after updating sink values. diff --git a/op_tests/test_flydsl_pa_decode.py b/op_tests/test_flydsl_pa_decode.py index e994e0416c3..85058ea31c2 100644 --- a/op_tests/test_flydsl_pa_decode.py +++ b/op_tests/test_flydsl_pa_decode.py @@ -99,8 +99,10 @@ def run_torch( key_scale: torch.Tensor, value_scale: torch.Tensor, query_length: int = 1, + sliding_window: int = 0, + sinks: torch.Tensor | None = None, ) -> torch.Tensor: - """Dequantized FP32 reference with sequence-major, causal MTP queries.""" + """Dequantized FP32 reference with sequence-major, windowed causal queries.""" num_queries, num_query_heads, head_dim = query.shape batch_size = context_lengths.numel() if query_length < 1 or num_queries != batch_size * query_length: @@ -142,8 +144,17 @@ def run_torch( ) visible = context_length - query_length + 1 + positions masked = token_ids.unsqueeze(0) >= visible.unsqueeze(1) + if sliding_window > 0: + masked |= token_ids.unsqueeze(0) < (visible - sliding_window).unsqueeze(1) scores.masked_fill_(masked[:, None, None, :], float("-inf")) - probs = torch.softmax(scores, dim=-1) + if sinks is None: + probs = torch.softmax(scores, dim=-1) + else: + sink_logits = sinks.float().reshape(1, num_kv_heads, query_group_size, 1) + sink_logits = sink_logits.expand(query_length, -1, -1, -1) + # The virtual token's V is zero, and its logit is not QK-scaled. + probs = torch.softmax(torch.cat((scores, sink_logits), dim=-1), dim=-1) + probs = probs[..., :-1] # Contexts shorter than QL have leading queries with no visible tokens. probs.masked_fill_(visible[:, None, None, None] <= 0, 0) first_row = seq_idx * query_length @@ -164,6 +175,8 @@ def _run_accuracy_case( num_partitions, tolerance, trans_v=False, + sliding_window=0, + sinks=None, ): """Run PA against a dequantized-FP8 torch reference without benchmarking.""" _require_gpu() @@ -229,6 +242,8 @@ def _run_accuracy_case( context_lengths, key_scale, value_scale, + sliding_window=sliding_window, + sinks=sinks, ) output = torch.empty_like(query) torch.ops.aiter.pa_decode_flydsl( @@ -246,6 +261,8 @@ def _run_accuracy_case( None, key_scale, value_scale, + sliding_window=sliding_window, + sinks=sinks, ) torch.testing.assert_close( @@ -257,11 +274,19 @@ def _run_accuracy_case( def _run_mtp4_fused_reference_case( - lengths, num_kv_heads, block_size, trans_v, num_partitions + lengths, + num_kv_heads, + block_size, + trans_v, + num_partitions, + *, + query_length=4, + sliding_window=0, + sinks=None, ): - """Check MTP4 with sparse pages, varying scales, and a causal FP32 reference.""" + """Check sparse MTP pages and varying scales against a causal FP32 reference.""" torch.manual_seed(37) - query_length, query_group_size, head_dim = 4, 16, 128 + query_group_size, head_dim = 16, 128 batch_size = len(lengths) num_query_heads = num_kv_heads * query_group_size pages_per_sequence = [(length + block_size - 1) // block_size for length in lengths] @@ -290,18 +315,25 @@ def _run_mtp4_fused_reference_case( token = torch.arange(num_pages * block_size).reshape(num_pages, 1, block_size, 1) kv_head = torch.arange(num_kv_heads).reshape(1, num_kv_heads, 1, 1) key_scale *= torch.exp2(((2 * token + kv_head) % 4 - 2).float()) - value_scale *= torch.exp2(((token + 2 * kv_head) % 5 - 3).float()) + value_scale_factors = torch.exp2(((token + 2 * kv_head) % 5 - 3).float()) + if sliding_window == 1: + # Exact binary ratios make a single visible token's normalized P + # representable in FP8 even when MTP queries share the V-scale max. + # This isolates mask/scale addressing from P-rounding error; wider + # windows retain the random quantization scales below. + value_scale = value_scale_factors * 2**-10 + else: + value_scale *= value_scale_factors - # Make C1--C4 analytically exact while still detecting causal off-by-one: - # K=0 gives uniform attention, V progresses as [.25, .5, .75, 1], and - # scale=1. A C4 row therefore produces [.25, .375, .5, .625]. + # Preserve the original exact C1--C4 cases independently of QL: unrestricted + # random short contexts hit FP8 probability-rounding error even without SW. + # K=0, V=[.25, .5, .75, 1], and scale=1 still detect causal/window off-by-one: + # dense QL4/C4 produces [.25, .375, .5, .625], and W=1 returns V directly. data_page = 0 for length, page_count in zip(lengths, pages_per_sequence): - if 0 < length <= query_length: + if 0 < length <= 4: capacity = page_count * block_size - token_values = ( - torch.arange(capacity, dtype=dtypes.fp32) % query_length + 1 - ) * 0.25 + token_values = (torch.arange(capacity, dtype=dtypes.fp32) % 4 + 1) * 0.25 key_quant[data_page : data_page + page_count].zero_() value_quant[data_page : data_page + page_count] = ( token_values.reshape(page_count, 1, 1, block_size) @@ -385,13 +417,19 @@ def scatter_pages(tensor): visible = max(0, length - (query_length - 1) + position) if visible == 0: continue + first = max(0, visible - sliding_window) if sliding_window > 0 else 0 row = seq * query_length + position scores = ( - torch.einsum("hd,khd->hk", query[row].float(), keys[:visible]) + torch.einsum("hd,khd->hk", query[row].float(), keys[first:visible]) * head_dim**-0.5 ) - probs = torch.softmax(scores, dim=-1) - reference[row] = torch.einsum("hk,khd->hd", probs, values[:visible]) + if sinks is None: + probs = torch.softmax(scores, dim=-1) + else: + probs = torch.softmax( + torch.cat((scores, sinks.float().unsqueeze(-1)), dim=-1), dim=-1 + )[:, :-1] + reference[row] = torch.einsum("hk,khd->hd", probs, values[first:visible]) output = torch.full_like(query, float("nan")) torch.ops.aiter.pa_decode_flydsl( @@ -409,6 +447,8 @@ def scatter_pages(tensor): None, key_scale, value_scale, + sliding_window=sliding_window, + sinks=sinks, ) assert torch.isfinite(output).all() torch.testing.assert_close(output.float(), reference, rtol=0.005, atol=0.005) @@ -962,7 +1002,16 @@ def run_once(fn, *args, **kwargs): def _assert_plan(plan, lengths): work = plan.work_info.cpu().tolist() reductions = plan.reduce_info.cpu().tolist() - tiles_per_sequence = [(max(length, 0) + 255) // 256 for length in lengths] + last_tiles = [(max(length, 0) + 255) // 256 for length in lengths] + first_tiles = [ + ( + max(0, length - (plan.query_length - 1) - plan.sliding_window) // 256 + if plan.sliding_window > 0 + else 0 + ) + for length in lengths + ] + tiles_per_sequence = [last - first for first, last in zip(first_tiles, last_tiles)] nonempty = [int(tiles > 0) for tiles in tiles_per_sequence] total_tiles = max(sum(tiles_per_sequence), 1) remaining = plan.capacity - sum(nonempty) @@ -979,8 +1028,8 @@ def _assert_plan(plan, lengths): expected_reductions = [] expected_work = [] offset = 0 - for seq, (length, tiles, count) in enumerate( - zip(lengths, tiles_per_sequence, expected_counts) + for seq, (length, first, last, tiles, count) in enumerate( + zip(lengths, first_tiles, last_tiles, tiles_per_sequence, expected_counts) ): start = offset expected_reductions.append([start, count]) @@ -988,8 +1037,8 @@ def _assert_plan(plan, lengths): expected_work.append( [ seq, - part * tiles // count, - (part + 1) * tiles // count, + first + part * tiles // count, + first + (part + 1) * tiles // count, length, ] ) @@ -998,12 +1047,12 @@ def _assert_plan(plan, lengths): assert (actual_start, actual_count) == (start, count) assert 0 <= actual_count <= min(tiles, plan.max_partitions) assert (actual_count > 0) == (length > 0) - previous_end = 0 + previous_end = first for task in work[actual_start : actual_start + actual_count]: assert task[0] == seq and task[3] == length - assert task[1] == previous_end and task[1] < task[2] <= tiles + assert task[1] == previous_end and task[1] < task[2] <= last previous_end = task[2] - assert previous_end == tiles + assert previous_end == last offset += actual_count expected_work.extend([[0, 0, 0, 0]] * (plan.capacity - offset)) assert reductions == expected_reductions @@ -1036,13 +1085,76 @@ def test_plan_covers_each_token_once(lengths, heads, max_parts, budget): _assert_plan(plan, lengths) -def test_plan_graph_refresh_overwrites_old_metadata(): +@pytest.mark.parametrize( + "sliding_window,query_length", + [ + (1, 1), + (1, 4), + (255, 2), + (256, 3), + (257, 4), + (509, 4), + (8192, 4), + (2**40, 4), + (0, 4), + (-1, 2), + ], +) +@pytest.mark.parametrize( + "heads,max_parts,budget", [(1, 1, 32), (1, 7, 17), (2, 256, 512)] +) +def test_sliding_window_plan_covers_absolute_tiles( + sliding_window, query_length, heads, max_parts, budget +): + # Include MTP unions crossing a 256-token boundary, negative/empty contexts, + # and int32's maximum length without allocating a corresponding KV cache. + lengths = [ + -1, + 0, + 1, + 2, + 3, + 4, + 255, + 256, + 257, + 258, + 259, + 511, + 512, + 513, + 514, + 515, + 8193, + 200003, + 2**31 - 1, + ] + context = torch.tensor(lengths, dtype=torch.int32) + plan = plan_pa_decode( + context, + heads, + max_partitions=max_parts, + workgroup_budget=budget, + query_length=query_length, + sliding_window=sliding_window, + ) + assert plan.sliding_window == max(sliding_window, 0) + assert plan.query_length == query_length + _assert_plan(plan, lengths) + + +@pytest.mark.parametrize("sliding_window,query_length", [(0, 1), (1, 4), (257, 3)]) +def test_plan_graph_refresh_overwrites_old_metadata(sliding_window, query_length): lengths = [200003] * 8 context = torch.tensor(lengths, dtype=torch.int32) - plan = plan_pa_decode(context, 1) + plan_options = { + "sliding_window": sliding_window, + "query_length": query_length, + } + plan = plan_pa_decode(context, 1, **plan_options) graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): - plan_pa_decode(context, 1, plan=plan) + plan_pa_decode(context, 1, plan=plan, **plan_options) for lengths in ([0] * 8, [0, 1, 3, 4, 257, 4096, 16385, 200003], [200003] * 8): context.copy_(torch.tensor(lengths, dtype=torch.int32)) graph.replay() @@ -1050,10 +1162,179 @@ def test_plan_graph_refresh_overwrites_old_metadata(): def _planned_call(*args, **kwargs): - plan = plan_pa_decode(args[4], args[2].shape[1], max_partitions=args[8]) + plan = plan_pa_decode( + args[4], + args[2].shape[1], + max_partitions=args[8], + query_length=args[7], + sliding_window=kwargs.get("sliding_window", 0), + ) pa_decode(*args, work_plan=plan, **kwargs) +@pytest.mark.parametrize( + "query_length,sliding_window,block_size,trans_v,max_parts,heads", + [ + (1, 1, 16, False, 1, 1), + (1, 257, 128, True, 7, 2), + (2, 1, 16, True, 7, 1), + (2, 255, 64, False, 64, 2), + (2, 8192, 128, True, 256, 1), + (3, 1, 128, False, 1, 2), + (3, 257, 16, True, 7, 1), + (3, 513, 64, True, 86, 1), + (4, 1, 16, False, 7, 1), + (4, 256, 128, True, 7, 2), + (4, 257, 128, False, 256, 1), + (4, 8192, 16, True, 1, 1), + (4, 2**40, 16, False, 7, 1), + ], +) +def test_planned_sliding_window_sparse_causal_reference( + monkeypatch, query_length, sliding_window, block_size, trans_v, max_parts, heads +): + """Each query has its own window even when its tile is shared with MTP peers.""" + monkeypatch.setattr(torch.ops.aiter, "pa_decode_flydsl", _planned_call) + _run_mtp4_fused_reference_case( + [0, 1, 2, 3, 4, 255, 256, 257, 258, 259, 513, 514, 515, 4099], + heads, + block_size, + trans_v, + max_parts, + query_length=query_length, + sliding_window=sliding_window, + ) + + +@pytest.mark.parametrize( + "query_length,sliding_window,block_size,trans_v,parts", + [ + (1, 1, 16, False, 1), + (2, 257, 128, True, 7), + (3, 1, 64, False, 7), + (4, 257, 128, True, 7), + (4, 8192, 16, False, 1), + (4, 2**31 - 1, 16, False, 1), + (4, 2**31, 16, False, 1), + ], +) +def test_static_sliding_window_sparse_causal_reference( + query_length, sliding_window, block_size, trans_v, parts +): + """The unplanned path, including query-split kernels, uses the same mask.""" + _run_mtp4_fused_reference_case( + [0, 1, 2, 3, 4, 255, 256, 257, 259, 513, 514, 4099], + 1, + block_size, + trans_v, + parts, + query_length=query_length, + sliding_window=sliding_window, + ) + + +@pytest.mark.parametrize( + "sliding_window,head_dim,block_size,trans_v,parts,dtype", + [ + (1, 64, 16, False, 1, torch.bfloat16), + (1, 128, 16, True, 7, torch.bfloat16), + (257, 128, 128, True, 7, torch.float16), + (257, 128, 128, True, 7, torch.bfloat16), + (513, 256, 64, True, 86, torch.float16), + (8192, 1024, 128, False, 256, torch.bfloat16), + ], +) +def test_planned_sliding_window_per_tensor_scales( + monkeypatch, sliding_window, head_dim, block_size, trans_v, parts, dtype +): + # GQA8 + BF16 + D128 + transposed V also exercises scalar FP8 decode on gfx950. + monkeypatch.setattr(torch.ops.aiter, "pa_decode_flydsl", _planned_call) + _run_accuracy_case( + [0, 1, 255, 256, 257, 513, 4099], + num_query_heads=16, + num_kv_heads=2, + head_dim=head_dim, + block_size=block_size, + query_dtype=dtype, + num_partitions=parts, + tolerance=0.005, + trans_v=trans_v, + sliding_window=sliding_window, + ) + + +@pytest.mark.parametrize("query_length", [1, 4]) +@pytest.mark.parametrize("planned", [False, True]) +def test_sliding_window_ignores_outside_value_scales(query_length, planned): + """Masked prefix scales cannot affect FP8 P normalization in the last tile.""" + context_length, block_size, head_dim = 511, 128, 128 + num_pages, num_kv_heads, query_group_size = 4, 1, 16 + quant_dtype = _quant_dtype() + context = torch.tensor([context_length], dtype=torch.int32) + block_tables = torch.arange(num_pages, dtype=torch.int32).reshape(1, num_pages) + query = torch.ones((query_length, query_group_size, head_dim), dtype=torch.bfloat16) + output = torch.full_like(query, float("nan")) + key_cache = torch.zeros( + (num_pages, num_kv_heads, head_dim // 16, block_size, 16), dtype=quant_dtype + ) + value_tokens = torch.zeros( + (num_pages, num_kv_heads, block_size, head_dim), dtype=quant_dtype + ) + scale_shape = (num_pages, num_kv_heads, block_size, 1) + key_scale = torch.ones(scale_shape, dtype=torch.float32) + value_scale = torch.full(scale_shape, 1e9, dtype=torch.float32) + first_visible = context_length - query_length + value_scale.reshape(-1)[first_visible:].fill_(1.0) + # Every query sees exactly its corresponding tail token. Prefix V=0 keeps + # dequantized values finite even though the excluded prefix scales are huge. + for position in range(query_length): + token = first_visible + position + value_tokens[token // block_size, 0, token % block_size].fill_( + (position + 1) / query_length + ) + value_cache = ( + value_tokens.view(num_pages, num_kv_heads, block_size // 16, 16, head_dim) + .permute(0, 1, 2, 4, 3) + .contiguous() + ) + plan = ( + plan_pa_decode( + context, + num_kv_heads, + sliding_window=1, + query_length=query_length, + max_partitions=7, + ) + if planned + else None + ) + pa_decode( + output, + query, + key_cache, + value_cache, + context, + block_tables, + head_dim**-0.5, + query_length, + 7, + KV_COMPUTE_BLOCK, + quant_dtype, + None, + key_scale, + value_scale, + sliding_window=1, + work_plan=plan, + ) + expected = ( + (torch.arange(1, query_length + 1, dtype=query.dtype) / query_length) + .reshape(query_length, 1, 1) + .expand_as(query) + ) + assert torch.isfinite(output).all() + torch.testing.assert_close(output, expected, atol=0.005, rtol=0.005) + + @pytest.mark.parametrize("block_size", [16, 128]) @pytest.mark.parametrize("trans_v", [False, True]) @pytest.mark.parametrize("max_parts", [1, 7, 64, 86, 256]) @@ -1114,21 +1395,31 @@ def test_planned_decode_register_reducer_chunk_boundary(monkeypatch, long_contex ) -def test_planned_decode_graph_replay_with_poisoned_scratch(monkeypatch): +@pytest.mark.parametrize( + "sliding_window,query_length", [(0, 4), (1, 4), (257, 3), (8193, 4)] +) +def test_planned_decode_graph_replay_with_poisoned_scratch( + monkeypatch, sliding_window, query_length +): def capture_and_replay(*args, **kwargs): output, query, key, _value, context = args[:5] heads, rows = ( key.shape[1], query.shape[0] // context.numel() * query.shape[1] // key.shape[1], ) - plan = plan_pa_decode(context, heads, max_partitions=args[8]) + plan_options = { + "max_partitions": args[8], + "sliding_window": sliding_window, + "query_length": query_length, + } + plan = plan_pa_decode(context, heads, **plan_options) shape = (heads, plan.capacity, rows) psum = torch.full(shape, float("nan"), dtype=torch.float32) pmax = torch.full_like(psum, float("nan")) pout = torch.full((*shape, query.shape[2]), float("nan"), dtype=query.dtype) def run(): - plan_pa_decode(context, heads, max_partitions=args[8], plan=plan) + plan_pa_decode(context, heads, plan=plan, **plan_options) pa_decode( *args, **kwargs, @@ -1151,16 +1442,26 @@ def run(): scratch.fill_(float("nan")) context.zero_() graph.replay() + _assert_plan(plan, [0] * context.numel()) assert torch.equal(output, torch.zeros_like(output)) for scratch in (pmax, psum, pout): assert torch.isnan(scratch).all() context.copy_(original) graph.replay() + _assert_plan(plan, original.cpu().tolist()) for scratch in (pmax, psum, pout): assert torch.isnan(scratch[:, active_tasks:]).all() monkeypatch.setattr(torch.ops.aiter, "pa_decode_flydsl", capture_and_replay) - _run_mtp4_fused_reference_case([0, 257, 4096, 200003], 1, 128, True, 256) + _run_mtp4_fused_reference_case( + [0, 257, 4096, 200003], + 1, + 128, + True, + 256, + query_length=query_length, + sliding_window=sliding_window, + ) def test_plan_rejects_incompatible_reuse(): @@ -1174,6 +1475,537 @@ def test_plan_rejects_incompatible_reuse(): plan_pa_decode(context[:1], 1, max_partitions=7, plan=plan) +def test_sliding_window_plan_rejects_incompatible_reuse(): + context = torch.tensor([1, 257, 4099], dtype=torch.int32) + plan = plan_pa_decode( + context, 1, sliding_window=257, query_length=4, max_partitions=7 + ) + for window in (0, 256): + with pytest.raises(ValueError, match="sliding_window"): + plan_pa_decode( + context, + 1, + sliding_window=window, + query_length=4, + max_partitions=7, + plan=plan, + ) + with pytest.raises(ValueError, match="query_length"): + plan_pa_decode(context, 1, sliding_window=257, max_partitions=7, plan=plan) + reused = plan_pa_decode( + context, 1, sliding_window=257, query_length=4, max_partitions=7, plan=plan + ) + assert reused is plan + _assert_plan(plan, context.cpu().tolist()) + + +@pytest.mark.parametrize( + "sliding_window,query_length,error,match", + [ + (-2, 1, ValueError, "sliding_window"), + (1.5, 1, TypeError, "sliding_window"), + (257, 0, ValueError, "query_length"), + (257, 1.5, TypeError, "query_length"), + ], +) +def test_sliding_window_plan_rejects_invalid_arguments( + sliding_window, query_length, error, match +): + context = torch.tensor([257], dtype=torch.int32) + with pytest.raises(error, match=match): + plan_pa_decode( + context, 1, sliding_window=sliding_window, query_length=query_length + ) + + +@pytest.mark.parametrize( + "plan_window,plan_query_length,sliding_window,query_length,match", + [ + (0, 4, 257, 4, "sliding_window"), + (257, 4, 0, 4, "sliding_window"), + (257, 4, 256, 4, "sliding_window"), + (257, 3, 257, 4, "query_length"), + (257, 4, 257, 1, "query_length"), + ], +) +def test_planned_decode_rejects_sliding_window_mismatch( + monkeypatch, plan_window, plan_query_length, sliding_window, query_length, match +): + def incompatible_plan(*args, **kwargs): + plan = plan_pa_decode( + args[4], + args[2].shape[1], + max_partitions=args[8], + sliding_window=plan_window, + query_length=plan_query_length, + ) + pa_decode(*args, work_plan=plan, **kwargs) + + monkeypatch.setattr(torch.ops.aiter, "pa_decode_flydsl", incompatible_plan) + with pytest.raises(ValueError, match=match): + _run_mtp4_fused_reference_case( + [0, 1, 2, 3, 4, 257], + 1, + 16, + False, + 7, + query_length=query_length, + sliding_window=sliding_window, + ) + + +@pytest.mark.parametrize( + "plan_window,sliding_window,query_length", [(-1, 0, 4), (0, -1, 3)] +) +def test_planned_decode_accepts_disabled_sliding_window_aliases( + monkeypatch, plan_window, sliding_window, query_length +): + def dense_plan(*args, **kwargs): + # A default dense plan is still valid for any number of causal queries. + plan = plan_pa_decode( + args[4], + args[2].shape[1], + max_partitions=args[8], + sliding_window=plan_window, + ) + assert plan.sliding_window == 0 and plan.query_length == 1 + assert ( + plan_pa_decode( + args[4], + args[2].shape[1], + max_partitions=args[8], + plan=plan, + sliding_window=sliding_window, + query_length=query_length, + ) + is plan + ) + pa_decode(*args, work_plan=plan, **kwargs) + + monkeypatch.setattr(torch.ops.aiter, "pa_decode_flydsl", dense_plan) + _run_mtp4_fused_reference_case( + [0, 1, 2, 3, 4, 257, 513], + 1, + 16, + False, + 7, + query_length=query_length, + sliding_window=sliding_window, + ) + + +def _make_sink_case( + lengths, + *, + query_length=4, + head_dim=128, + num_kv_heads=1, + query_group_size=16, + query_dtype=torch.bfloat16, + per_token=True, + num_partitions=7, + sliding_window=0, +): + """Build exact zero-Q/K data and a length-aware, overflow-safe sink oracle.""" + block_size = 128 + num_pages = (max(max(lengths), 1) + block_size - 1) // block_size + num_query_heads = num_kv_heads * query_group_size + context = torch.tensor(lengths, dtype=torch.int32) + query = torch.zeros( + (len(lengths) * query_length, num_query_heads, head_dim), dtype=query_dtype + ) + output = torch.full_like(query, float("nan")) + quant_dtype = _quant_dtype() + key_cache = torch.zeros( + (num_pages, num_kv_heads, head_dim // 16, block_size, 16), dtype=quant_dtype + ) + token_values = ( + torch.arange(num_pages * block_size, dtype=torch.float32) % 4 + 1 + ) * 0.25 + kv_factors = torch.arange(1, num_kv_heads + 1, dtype=torch.float32) + values = ( + ( + token_values.reshape(num_pages, 1, block_size, 1) + * kv_factors.reshape(1, num_kv_heads, 1, 1) + ) + .expand(num_pages, num_kv_heads, block_size, head_dim) + .to(quant_dtype) + ) + value_cache = ( + values.reshape(num_pages, num_kv_heads, block_size // 16, 16, head_dim) + .permute(0, 1, 2, 4, 3) + .contiguous() + ) + scale_shape = (num_pages, num_kv_heads, block_size, 1) if per_token else (1,) + key_scale = torch.ones(scale_shape, dtype=torch.float32) + value_scale = torch.ones_like(key_scale) + block_tables = ( + torch.arange(num_pages, dtype=torch.int32) + .expand(len(lengths), num_pages) + .contiguous() + ) + args = ( + output, + query, + key_cache, + value_cache, + context, + block_tables, + head_dim**-0.5, + query_length, + num_partitions, + KV_COMPUTE_BLOCK, + quant_dtype, + None, + key_scale, + value_scale, + ) + + def reference(sinks): + # The shared pages encode [.25, .5, .75, 1] * (KV head + 1). + # Unlike exp(sinks), logaddexp stays finite for large finite logits. + positions = torch.arange(query_length, dtype=torch.int32) + visible = (context[:, None] - query_length + 1 + positions).clamp_min(0) + first = (visible - sliding_window).clamp_min(0) if sliding_window > 0 else 0 + count = (visible - first).float() + + def prefix_sum(tokens): + remainder = (tokens % 4).float() + return (tokens // 4).float() * 2.5 + remainder * (remainder + 1) * 0.125 + + value_sum = prefix_sum(visible) + if sliding_window > 0: + value_sum -= prefix_sum(first) + value_sum = value_sum[..., None] * kv_factors.repeat_interleave( + query_group_size + ) + logits = ( + torch.full((num_query_heads,), float("-inf"), dtype=torch.float32) + if sinks is None + else sinks.float() + ) + log_denominator = torch.logaddexp(count.log()[..., None], logits) + expected = value_sum * torch.exp(-log_denominator) + expected.masked_fill_(count[..., None] == 0, 0) + return expected.reshape(-1, num_query_heads, 1).expand_as(query) + + return args, reference + + +def _analytic_sinks(num_heads, dtype): + values = [float("-inf"), float("inf"), -1000, 1000, -2, 0, 0.5, 5] + if dtype == torch.float32: + values += [-3e38, 3e38] + return ( + torch.tensor(values, dtype=dtype) + .repeat((num_heads + len(values) - 1) // len(values))[:num_heads] + .contiguous() + ) + + +@pytest.mark.parametrize( + "planned,query_length,parts,window,block_size,trans_v,heads,sink_dtype", + [ + (False, 1, 1, 0, 16, False, 1, torch.float32), + (False, 4, 1, 257, 128, True, 1, torch.bfloat16), + (True, 4, 1, 1, 16, False, 1, torch.float16), + (True, 2, 7, 257, 128, True, 2, torch.float32), + (False, 3, 7, 0, 64, False, 2, torch.bfloat16), + (True, 4, 7, 0, 128, True, 1, torch.float32), + (False, 2, 86, 1, 16, True, 1, torch.float16), + (True, 1, 86, 257, 64, False, 2, torch.bfloat16), + (False, 4, 256, 257, 128, False, 1, torch.float32), + (True, 3, 256, 0, 16, True, 1, torch.float16), + (True, 1, 7, 1, 128, True, 1, torch.float32), + (False, 4, 1, 0, 16, False, 1, torch.float16), + ], +) +def test_sinks_sparse_causal_reference( + monkeypatch, + planned, + query_length, + parts, + window, + block_size, + trans_v, + heads, + sink_dtype, +): + if planned: + monkeypatch.setattr(torch.ops.aiter, "pa_decode_flydsl", _planned_call) + sinks = torch.linspace(-3, 6, heads * 16, dtype=sink_dtype) + sinks[0] = float("-inf") + _run_mtp4_fused_reference_case( + [0, 1, 2, 3, 4, 255, 256, 257, 259, 513, 4099], + heads, + block_size, + trans_v, + parts, + query_length=query_length, + sliding_window=window, + sinks=sinks, + ) + + +@pytest.mark.parametrize( + "planned,parts,head_dim,block_size,trans_v,query_dtype,sink_dtype,window", + [ + (False, 1, 64, 16, False, torch.float16, torch.float32, 0), + (True, 7, 256, 64, True, torch.float16, torch.bfloat16, 257), + (False, 86, 128, 128, True, torch.bfloat16, torch.float16, 1), + (True, 256, 256, 128, False, torch.bfloat16, torch.float32, 0), + (True, 1, 128, 16, True, torch.bfloat16, torch.float32, 257), + (False, 256, 64, 128, True, torch.float16, torch.float16, 257), + ], +) +def test_sinks_per_tensor_reference( + monkeypatch, + planned, + parts, + head_dim, + block_size, + trans_v, + query_dtype, + sink_dtype, + window, +): + if planned: + monkeypatch.setattr(torch.ops.aiter, "pa_decode_flydsl", _planned_call) + sinks = torch.linspace(-3, 7, 16, dtype=sink_dtype) + _run_accuracy_case( + [0, 1, 257, 16385], + num_query_heads=16, + num_kv_heads=2, + head_dim=head_dim, + block_size=block_size, + query_dtype=query_dtype, + num_partitions=parts, + tolerance=0.005, + trans_v=trans_v, + sliding_window=window, + sinks=sinks, + ) + + +@pytest.mark.parametrize("planned", [False, True]) +@pytest.mark.parametrize( + "query_length,parts,head_dim,window,query_dtype,sink_dtype,heads,group,per_token", + [ + (1, 1, 128, 0, torch.bfloat16, torch.float32, 2, 8, False), + (4, 1, 128, 1, torch.bfloat16, torch.bfloat16, 1, 16, True), + (2, 7, 64, 257, torch.float16, torch.float32, 2, 4, False), + (3, 7, 128, 0, torch.bfloat16, torch.float16, 2, 8, True), + (4, 86, 256, 257, torch.float16, torch.bfloat16, 1, 8, False), + (1, 256, 128, 0, torch.bfloat16, torch.float16, 1, 16, True), + (4, 256, 64, 1, torch.bfloat16, torch.float32, 2, 4, True), + (2, 86, 256, 0, torch.bfloat16, torch.bfloat16, 2, 8, False), + ], +) +def test_sinks_counted_once_per_query( + monkeypatch, + planned, + query_length, + parts, + head_dim, + window, + query_dtype, + sink_dtype, + heads, + group, + per_token, +): + """Different head logits, including extremes, add one unscaled zero-V token.""" + args, reference = _make_sink_case( + [0, 1, 3, 4, 257, 513], + query_length=query_length, + head_dim=head_dim, + num_kv_heads=heads, + query_group_size=group, + query_dtype=query_dtype, + per_token=per_token, + num_partitions=parts, + sliding_window=window, + ) + sinks = _analytic_sinks(heads * group, sink_dtype) + module = importlib.import_module("aiter.ops.flydsl.pa_decode") + compile_tile = module.compile_pa_decode_tile + direct_sink_flags = [] + + def capture_compile(**kwargs): + direct_sink_flags.append(kwargs["use_sinks"]) + return compile_tile(**kwargs) + + monkeypatch.setattr(module, "compile_pa_decode_tile", capture_compile) + call = _planned_call if planned else pa_decode + call(*args, sinks=sinks, sliding_window=window) + assert direct_sink_flags and set(direct_sink_flags) == {not planned and parts == 1} + assert torch.isfinite(args[0]).all() + torch.testing.assert_close( + args[0].float(), reference(sinks), atol=0.005, rtol=0.005 + ) + + +@pytest.mark.parametrize("query_splits", [1, 4]) +def test_sinks_np1_fused_and_query_split(monkeypatch, query_splits): + """Static NP1 applies sinks in the compute epilogue without a reduction launch.""" + args, reference = _make_sink_case( + [0, 1, 3, 4, 257, 513], num_partitions=1, sliding_window=1 + ) + sinks = _analytic_sinks(args[1].shape[1], torch.float32) + module = importlib.import_module("aiter.ops.flydsl.pa_decode") + compile_tile = module.compile_pa_decode_tile + + def compile_query_split(**kwargs): + kwargs["query_splits"] = query_splits + assert kwargs["use_sinks"] and kwargs["sink_dtype_str"] == "f32" + return compile_tile(**kwargs) + + def unexpected_reducer(*_args, **_kwargs): + raise AssertionError("static NP1 sinks must not launch a reducer") + + monkeypatch.setattr(module, "compile_pa_decode_tile", compile_query_split) + monkeypatch.setattr(module, "launch_pa_decode_ps_reduce", unexpected_reducer) + pa_decode(*args, sinks=sinks, sliding_window=1) + torch.testing.assert_close( + args[0].float(), reference(sinks), atol=0.005, rtol=0.005 + ) + + +@pytest.mark.parametrize( + "planned,parts,query_length,head_dim,window", + [ + (False, 1, 4, 128, 257), + (False, 7, 1, 64, 0), + (True, 1, 4, 128, 1), + (True, 256, 3, 256, 257), + ], +) +def test_sinks_negative_infinity_matches_none( + planned, parts, query_length, head_dim, window +): + args, reference = _make_sink_case( + [0, 1, 3, 4, 257, 513], + query_length=query_length, + head_dim=head_dim, + num_partitions=parts, + sliding_window=window, + ) + call = _planned_call if planned else pa_decode + call(*args, sinks=None, sliding_window=window) + without_sinks = args[0].clone() + args[0].fill_(float("nan")) + sinks = torch.full((args[1].shape[1],), float("-inf"), dtype=torch.float32) + call(*args, sinks=sinks, sliding_window=window) + torch.testing.assert_close(args[0], without_sinks, atol=0.005, rtol=0.005) + torch.testing.assert_close(args[0].float(), reference(None), atol=0.005, rtol=0.005) + + +@pytest.mark.parametrize( + "planned,head_dim,query_length,window,parts,sink_dtype", + [ + (True, 128, 4, 257, 7, torch.float32), + (True, 256, 3, 0, 86, torch.bfloat16), + (False, 128, 4, 1, 7, torch.float16), + (False, 64, 2, 0, 86, torch.float32), + ], +) +def test_sinks_graph_replay_with_updated_logits_and_poisoned_scratch( + planned, head_dim, query_length, window, parts, sink_dtype +): + args, reference = _make_sink_case( + [513, 4, 1, 0], + query_length=query_length, + head_dim=head_dim, + num_partitions=parts, + sliding_window=window, + ) + output, query, key_cache, _, context = args[:5] + heads = key_cache.shape[1] + rows = query_length * query.shape[1] // heads + sinks = torch.linspace(-3, 6, query.shape[1], dtype=sink_dtype) + plan_options = { + "max_partitions": parts, + "query_length": query_length, + "sliding_window": window, + } + plan = plan_pa_decode(context, heads, **plan_options) if planned else None + shape = ( + (heads, plan.capacity, rows) + if planned + else (context.numel(), heads, parts, rows) + ) + psum = torch.full(shape, float("nan"), dtype=torch.float32) + pmax = torch.full_like(psum, float("nan")) + pout = torch.full((*shape, head_dim), float("nan"), dtype=query.dtype) + + def run(): + if planned: + plan_pa_decode(context, heads, plan=plan, **plan_options) + pa_decode( + *args, + sinks=sinks, + sliding_window=window, + work_plan=plan, + exp_sums=psum, + max_logits=pmax, + temporary_output=pout, + ) + + run() # Compile before capture; all captured pointers remain fixed. + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + for lengths, updated_sinks in ( + ([0, 1, 4, 513], torch.linspace(7, -4, query.shape[1], dtype=sink_dtype)), + ([0, 0, 0, 0], torch.full_like(sinks, 1000)), + ([513, 4, 1, 0], _analytic_sinks(query.shape[1], sink_dtype)), + ): + context.copy_(torch.tensor(lengths, dtype=torch.int32)) + sinks.copy_(updated_sinks) + output.fill_(float("nan")) + for scratch in (pmax, psum, pout): + scratch.fill_(float("nan")) + graph.replay() + assert torch.isfinite(output).all() + torch.testing.assert_close( + output.float(), reference(sinks), atol=0.005, rtol=0.005 + ) + if planned: + _assert_plan(plan, lengths) + active = int(plan.reduce_info[:, 1].sum().item()) + assert active < plan.capacity + for scratch in (pmax, psum, pout): + assert torch.isnan(scratch[:, active:]).all() + + +@pytest.mark.parametrize( + "invalid,error", + [ + ("not_tensor", TypeError), + ("rank", ValueError), + ("head_count", ValueError), + ("integer_dtype", TypeError), + ("float64_dtype", TypeError), + ("cpu", ValueError), + ("noncontiguous", ValueError), + ], +) +def test_sinks_rejects_invalid_tensor(invalid, error): + args, _ = _make_sink_case([257], query_length=1) + heads = args[1].shape[1] + invalid_sinks = { + "not_tensor": lambda: [0.0] * heads, + "rank": lambda: torch.zeros((1, heads)), + "head_count": lambda: torch.zeros(heads - 1), + "integer_dtype": lambda: torch.zeros(heads, dtype=torch.int32), + "float64_dtype": lambda: torch.zeros(heads, dtype=torch.float64), + "cpu": lambda: torch.zeros(heads, device="cpu"), + "noncontiguous": lambda: torch.zeros(heads * 2)[::2], + }[invalid]() + with pytest.raises(error, match="sinks"): + pa_decode(*args, sinks=invalid_sinks) + + @pytest.mark.parametrize("split_kv_blocks", [2, 16]) @pytest.mark.parametrize( "batch_size,context_length,max_partitions,expected", From fd825332fae0aac33b5e5b0123c3e58c3776de2b Mon Sep 17 00:00:00 2001 From: fsx950223 Date: Thu, 17 Sep 2026 06:57:38 +0000 Subject: [PATCH 52/66] test(flydsl): consolidate PA decode into one parametrized test Share input generation, the FP32 reference, and launch helpers with the CLI while retaining sliding-window, sinks, plan, and graph-replay coverage. Validated on dzm-mi355-gpu34-dev: 72 pytest cases and 10 CLI smoke configurations passed. --- op_tests/test_flydsl_pa_decode.py | 2833 ++++++++--------------------- 1 file changed, 745 insertions(+), 2088 deletions(-) diff --git a/op_tests/test_flydsl_pa_decode.py b/op_tests/test_flydsl_pa_decode.py index 85058ea31c2..1849b4189d9 100644 --- a/op_tests/test_flydsl_pa_decode.py +++ b/op_tests/test_flydsl_pa_decode.py @@ -1,25 +1,23 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. -"""Correctness and performance sweep for FlyDSL paged-attention Tile. +"""One parametrized correctness test and a CLI performance sweep for FlyDSL PA. -For example, BS200/MTP4 with FP8 per-token KV scales and explicit split counts:: - - python3 op_tests/test_flydsl_pa_decode.py -d bf16 -b 200 -q 4 \ + 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 have equal lengths and include the MTP query tokens. Query position -``p`` attends to ``max(0, context_length - query_length + 1 + p)`` KV tokens. -Timing includes the FlyDSL attention kernel and its native FlyDSL reduction. -Automatic splits use the host-known context length and GPU occupancy, up to -256 partitions. Use ``--max-partitions 8`` for the legacy clamp, or -``--num-partitions`` to bypass the recommendation with exact counts. +Both entry points share input generation, the FP32 reference and kernel launch. +Context lengths include the MTP query tokens. Explicit partition counts override +automatic recommendations; the CLI times attention plus the native reducer. """ import argparse import importlib import itertools +from dataclasses import dataclass +from functools import partial import pandas as pd import pytest @@ -28,30 +26,22 @@ 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, checkAllclose, run_perftest - - -@pytest.fixture(autouse=True) -def _default_cuda_device(): - # Scoped rather than set at import time: this module now runs in the shared - # Standard Tests shard, where a global default-device switch would follow - # every later test in the session. - _require_gpu() - previous = torch.get_default_device() - torch.set_default_device("cuda") - yield - torch.set_default_device(previous) +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"] +SUPPORTED_GFX = ("gfx942", "gfx950") KV_COMPUTE_BLOCK = 256 -# All parametrized cases use equal context lengths and the strict bound. -FIXED_LENGTH_ACCURACY_TOLERANCE = 5e-3 - -# Pairwise coverage of the normal-accuracy axes in FlyDSL's PA regression test: -# batches {3, 81, 128}, Q/KV heads {(4,1), (8,1), (16,1)}, head dims -# {128, 256}, and contexts {1027, 8192}. Keep the original 257-token boundary -# case as well. All supported block sizes are crossed with every case in main(). +ACCURACY_TOLERANCE = 5e-3 DEFAULT_BATCH_SIZES = [3, 81, 128] DEFAULT_SHAPES = [ (8, 1, 128, 257), @@ -60,25 +50,53 @@ def _default_cuda_device(): (8, 1, 256, 1027), (16, 1, 128, 8192), ] - -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 = None - pa_decode = None - plan_pa_decode = None +BOUNDARY_LENGTHS = ( + 0, + 1, + 2, + 3, + 4, + 255, + 256, + 257, + 258, + 259, + 511, + 512, + 513, + 514, + 515, + 2303, + 2304, + 2305, + 2306, + 2307, + 4099, +) -def _quant_dtype() -> torch.dtype: - return ( - torch.float8_e4m3fn if get_gfx_runtime() == "gfx950" else torch.float8_e4m3fnuz - ) +@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 + expected_prefetch: bool | str | None = None + expected_query_splits: int | None = None def _require_gpu(): @@ -90,372 +108,19 @@ def _require_gpu(): pytest.skip(f"pa_decode is unsupported on {get_gfx_runtime()}") -def run_torch( - query: torch.Tensor, - key_cache: torch.Tensor, - value_cache: torch.Tensor, - block_tables: torch.Tensor, - context_lengths: torch.Tensor, - key_scale: torch.Tensor, - value_scale: torch.Tensor, - query_length: int = 1, - sliding_window: int = 0, - sinks: torch.Tensor | None = None, -) -> torch.Tensor: - """Dequantized FP32 reference with sequence-major, windowed causal queries.""" - num_queries, num_query_heads, head_dim = query.shape - batch_size = context_lengths.numel() - if query_length < 1 or num_queries != batch_size * query_length: - raise ValueError("query must have batch_size * query_length rows") - block_size = key_cache.shape[2] - num_kv_heads = key_cache.shape[1] - query_group_size = num_query_heads // num_kv_heads - softmax_scale = head_dim**-0.5 - output = torch.zeros_like(query) - queries = query.reshape( - batch_size, query_length, num_kv_heads, query_group_size, head_dim - ) - positions = torch.arange(query_length, device=query.device) - - for seq_idx in range(batch_size): - context_length = int(context_lengths[seq_idx].item()) - if context_length == 0: - continue - token_ids = torch.arange(context_length, device=query.device) - logical_pages = token_ids // block_size - token_offsets = token_ids % block_size - physical_pages = block_tables[seq_idx, logical_pages].long() - - keys = key_cache[physical_pages, :, token_offsets, :].float() - values = value_cache[physical_pages, :, :, token_offsets].float() - if key_scale.numel() == 1: - keys = keys * key_scale.float() - values = values * value_scale.float() - else: - token_key_scale = key_scale[physical_pages, :, token_offsets, 0].float() - token_value_scale = value_scale[physical_pages, :, token_offsets, 0].float() - keys = keys * token_key_scale.unsqueeze(-1) - values = values * token_value_scale.unsqueeze(-1) - # Keep KV heads grouped instead of materializing Hq/Hkv copies of the - # long-context cache just for the reference (notably BS200/C200k). - scores = ( - torch.einsum("qhgd,khd->qhgk", queries[seq_idx].float(), keys) - * softmax_scale - ) - visible = context_length - query_length + 1 + positions - masked = token_ids.unsqueeze(0) >= visible.unsqueeze(1) - if sliding_window > 0: - masked |= token_ids.unsqueeze(0) < (visible - sliding_window).unsqueeze(1) - scores.masked_fill_(masked[:, None, None, :], float("-inf")) - if sinks is None: - probs = torch.softmax(scores, dim=-1) - else: - sink_logits = sinks.float().reshape(1, num_kv_heads, query_group_size, 1) - sink_logits = sink_logits.expand(query_length, -1, -1, -1) - # The virtual token's V is zero, and its logit is not QK-scaled. - probs = torch.softmax(torch.cat((scores, sink_logits), dim=-1), dim=-1) - probs = probs[..., :-1] - # Contexts shorter than QL have leading queries with no visible tokens. - probs.masked_fill_(visible[:, None, None, None] <= 0, 0) - first_row = seq_idx * query_length - output[first_row : first_row + query_length] = torch.einsum( - "qhgk,khd->qhgd", probs, values - ).reshape(query_length, num_query_heads, head_dim) - return output - - -def _run_accuracy_case( - lengths, - *, - num_query_heads, - num_kv_heads, - head_dim, - block_size, - query_dtype, - num_partitions, - tolerance, - trans_v=False, - sliding_window=0, - sinks=None, -): - """Run PA against a dequantized-FP8 torch reference without benchmarking.""" +@pytest.fixture(autouse=True) +def _default_cuda_device(): + # Do not leak the default device into other files in the shared CI shard. _require_gpu() - if num_query_heads % num_kv_heads: - raise ValueError("num_query_heads must be divisible by num_kv_heads") - - torch.manual_seed(0) - batch_size = len(lengths) - pages_per_sequence = [ - (context_length + block_size - 1) // block_size for context_length in lengths - ] - max_pages = max(pages_per_sequence) - num_blocks = sum(pages_per_sequence) - - query = torch.empty( - batch_size, num_query_heads, head_dim, dtype=query_dtype - ).uniform_(-0.5, 0.5) - key = torch.empty( - num_blocks, num_kv_heads, block_size, head_dim, dtype=query_dtype - ).uniform_(-0.5, 0.5) - value = torch.empty( - num_blocks, num_kv_heads, head_dim, block_size, dtype=query_dtype - ).uniform_(-0.5, 0.5) - - quant_dtype = _quant_dtype() - key_quant, key_scale = per_tensor_quant(key, quant_dtype=quant_dtype) - value_quant, value_scale = per_tensor_quant(value, quant_dtype=quant_dtype) - key_cache = ( - key_quant.view(num_blocks, num_kv_heads, block_size, head_dim // 16, 16) - .permute(0, 1, 3, 2, 4) - .contiguous() - ) - if trans_v: - value_cache = ( - value_quant.permute(0, 1, 3, 2) - .contiguous() - .view(num_blocks, num_kv_heads, block_size // 16, 16, head_dim) - .permute(0, 1, 2, 4, 3) - .contiguous() - ) - else: - value_cache = value_quant.contiguous() - - block_tables = torch.zeros( - batch_size, max_pages, dtype=dtypes.i32, device=query.device - ) - next_page = 0 - for seq_idx, page_count in enumerate(pages_per_sequence): - block_tables[seq_idx, :page_count] = torch.arange( - next_page, - next_page + page_count, - dtype=dtypes.i32, - device=query.device, - ) - next_page += page_count - context_lengths = torch.tensor(lengths, dtype=dtypes.i32, device=query.device) - - reference = run_torch( - query, - key_quant, - value_quant, - block_tables, - context_lengths, - key_scale, - value_scale, - sliding_window=sliding_window, - sinks=sinks, - ) - output = torch.empty_like(query) - torch.ops.aiter.pa_decode_flydsl( - output, - query, - key_cache, - value_cache, - context_lengths, - block_tables, - head_dim**-0.5, - 1, - num_partitions, - KV_COMPUTE_BLOCK, - quant_dtype, - None, - key_scale, - value_scale, - sliding_window=sliding_window, - sinks=sinks, - ) - - torch.testing.assert_close( - output.float(), - reference.float(), - atol=tolerance, - rtol=tolerance, - ) - - -def _run_mtp4_fused_reference_case( - lengths, - num_kv_heads, - block_size, - trans_v, - num_partitions, - *, - query_length=4, - sliding_window=0, - sinks=None, -): - """Check sparse MTP pages and varying scales against a causal FP32 reference.""" - torch.manual_seed(37) - query_group_size, head_dim = 16, 128 - batch_size = len(lengths) - num_query_heads = num_kv_heads * query_group_size - pages_per_sequence = [(length + block_size - 1) // block_size for length in lengths] - max_pages = max(pages_per_sequence) - num_pages = sum(pages_per_sequence) - quant_dtype = _quant_dtype() - - query = torch.empty( - batch_size * query_length, - num_query_heads, - head_dim, - dtype=dtypes.bf16, - ).uniform_(-0.5, 0.5) - key = torch.empty( - num_pages, num_kv_heads, block_size, head_dim, dtype=dtypes.bf16 - ).uniform_(-0.5, 0.5) - value = torch.empty( - num_pages, num_kv_heads, head_dim, block_size, dtype=dtypes.bf16 - ).uniform_(-0.5, 0.5) - key_quant, key_scale = pertoken_quant(key, quant_dtype=quant_dtype) - value_token_quant, value_scale = pertoken_quant( - value.permute(0, 1, 3, 2).contiguous(), quant_dtype=quant_dtype - ) - value_quant = value_token_quant.permute(0, 1, 3, 2).contiguous() - - token = torch.arange(num_pages * block_size).reshape(num_pages, 1, block_size, 1) - kv_head = torch.arange(num_kv_heads).reshape(1, num_kv_heads, 1, 1) - key_scale *= torch.exp2(((2 * token + kv_head) % 4 - 2).float()) - value_scale_factors = torch.exp2(((token + 2 * kv_head) % 5 - 3).float()) - if sliding_window == 1: - # Exact binary ratios make a single visible token's normalized P - # representable in FP8 even when MTP queries share the V-scale max. - # This isolates mask/scale addressing from P-rounding error; wider - # windows retain the random quantization scales below. - value_scale = value_scale_factors * 2**-10 - else: - value_scale *= value_scale_factors - - # Preserve the original exact C1--C4 cases independently of QL: unrestricted - # random short contexts hit FP8 probability-rounding error even without SW. - # K=0, V=[.25, .5, .75, 1], and scale=1 still detect causal/window off-by-one: - # dense QL4/C4 produces [.25, .375, .5, .625], and W=1 returns V directly. - data_page = 0 - for length, page_count in zip(lengths, pages_per_sequence): - if 0 < length <= 4: - capacity = page_count * block_size - token_values = (torch.arange(capacity, dtype=dtypes.fp32) % 4 + 1) * 0.25 - key_quant[data_page : data_page + page_count].zero_() - value_quant[data_page : data_page + page_count] = ( - token_values.reshape(page_count, 1, 1, block_size) - .expand(page_count, num_kv_heads, head_dim, block_size) - .to(quant_dtype) - ) - key_scale[data_page : data_page + page_count].fill_(1.0) - value_scale[data_page : data_page + page_count].fill_(1.0) - data_page += page_count - - # Scatter logical pages into a permuted set of odd physical pages. Leave - # holes and page zero finite so masked/padded token reads remain valid. - selected_pages = 2 * torch.randperm(num_pages) + 1 - physical_page_count = 2 * num_pages + 1 - - def scatter_pages(tensor): - sparse = torch.zeros( - (physical_page_count, *tensor.shape[1:]), dtype=tensor.dtype - ) - sparse[selected_pages] = tensor - return sparse - - key_quant = scatter_pages(key_quant) - value_quant = scatter_pages(value_quant) - key_scale = scatter_pages(key_scale) - value_scale = scatter_pages(value_scale) - num_pages = physical_page_count - - key_cache = ( - key_quant.view( - num_pages, - num_kv_heads, - block_size, - head_dim // 16, - 16, - ) - .permute(0, 1, 3, 2, 4) - .contiguous() - ) - if trans_v: - value_cache = ( - value_quant.permute(0, 1, 3, 2) - .contiguous() - .view( - num_pages, - num_kv_heads, - block_size // 16, - 16, - head_dim, - ) - .permute(0, 1, 2, 4, 3) - .contiguous() - ) - else: - value_cache = value_quant - - block_tables = torch.zeros(batch_size, max_pages, dtype=dtypes.i32) - next_page = 0 - for seq, page_count in enumerate(pages_per_sequence): - block_tables[seq, :page_count] = selected_pages[ - next_page : next_page + page_count - ] - next_page += page_count - context_lengths = torch.tensor(lengths, dtype=dtypes.i32) - - reference = torch.zeros_like(query, dtype=dtypes.fp32) - for seq, length in enumerate(lengths): - if length == 0: - continue - token_ids = torch.arange(length) - logical_pages = token_ids // block_size - token_offsets = token_ids % block_size - physical_pages = block_tables[seq, logical_pages].long() - keys = key_quant[physical_pages, :, token_offsets, :].float() - values = value_quant[physical_pages, :, :, token_offsets].float() - keys *= key_scale[physical_pages, :, token_offsets, 0].float().unsqueeze(-1) - values *= value_scale[physical_pages, :, token_offsets, 0].float().unsqueeze(-1) - keys = keys.repeat_interleave(query_group_size, dim=1) - values = values.repeat_interleave(query_group_size, dim=1) - for position in range(query_length): - visible = max(0, length - (query_length - 1) + position) - if visible == 0: - continue - first = max(0, visible - sliding_window) if sliding_window > 0 else 0 - row = seq * query_length + position - scores = ( - torch.einsum("hd,khd->hk", query[row].float(), keys[first:visible]) - * head_dim**-0.5 - ) - if sinks is None: - probs = torch.softmax(scores, dim=-1) - else: - probs = torch.softmax( - torch.cat((scores, sinks.float().unsqueeze(-1)), dim=-1), dim=-1 - )[:, :-1] - reference[row] = torch.einsum("hk,khd->hd", probs, values[first:visible]) - - output = torch.full_like(query, float("nan")) - torch.ops.aiter.pa_decode_flydsl( - output, - query, - key_cache, - value_cache, - context_lengths, - block_tables, - head_dim**-0.5, - query_length, - num_partitions, - KV_COMPUTE_BLOCK, - quant_dtype, - None, - key_scale, - value_scale, - sliding_window=sliding_window, - sinks=sinks, - ) - assert torch.isfinite(output).all() - torch.testing.assert_close(output.float(), reference, rtol=0.005, atol=0.005) + previous = torch.get_default_device() + torch.set_default_device("cuda") + try: + yield + finally: + torch.set_default_device(previous) -def _run_flydsl( - output, +def run_torch( query, key_cache, value_cache, @@ -463,1671 +128,704 @@ def _run_flydsl( context_lengths, key_scale, value_scale, - num_partitions, - softmax_scale, - pmax, - psum, - pout, -): - torch.ops.aiter.pa_decode_flydsl( - output, - query, - key_cache, - value_cache, - context_lengths, - block_tables, - softmax_scale, - query.shape[0] // context_lengths.shape[0], - num_partitions, - 256, - key_cache.dtype, - None, - key_scale, - value_scale, - exp_sums=psum, - max_logits=pmax, - temporary_output=pout, - ps=True, - ) - return output - - -@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 query_length < 1: - raise ValueError("query_length must be positive") - 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}]") - if batch_size < 1 or context_length < 1: - raise ValueError("batch_size and context_length must be positive") - if pa_decode is None or get_recommended_splits is None: - raise RuntimeError("FlyDSL is not available") - if dtype not in (dtypes.fp16, dtypes.bf16): - raise ValueError(f"pa_decode only supports fp16/bf16, got {dtype}") - if num_query_heads < 1 or num_kv_heads < 1 or num_query_heads % num_kv_heads != 0: - raise ValueError("num_query_heads must be divisible by num_kv_heads") - - # An explicit NP is an exact override, not an upper clamp. In particular, - # NP=3 and NP=5 must not be rounded by the automatic recommendation. - if num_partitions is None: - num_partitions = get_recommended_splits( - batch_size, - num_kv_heads, - split_kv_blocks=KV_COMPUTE_BLOCK // block_size, - max_partitions=max_partitions, - max_context_length=context_length, - ) - - torch.manual_seed(0) - blocks_per_sequence = (context_length + block_size - 1) // block_size - num_blocks = batch_size * blocks_per_sequence - - query = torch.empty( - batch_size * query_length, - num_query_heads, - head_dim, - dtype=dtype, - ).uniform_(-0.5, 0.5) - key = torch.empty( - num_blocks, - num_kv_heads, - block_size, - head_dim, - dtype=dtype, - ).uniform_(-0.5, 0.5) - value = torch.empty( - num_blocks, - num_kv_heads, - head_dim, - block_size, - dtype=dtype, - ).uniform_(-0.5, 0.5) - - quant_dtype = _quant_dtype() - if per_token: - key_quant, key_scale = pertoken_quant(key, quant_dtype=quant_dtype) - value_token_major = value.permute(0, 1, 3, 2).contiguous() - value_token_quant, value_scale = pertoken_quant( - value_token_major, quant_dtype=quant_dtype - ) - value_quant = value_token_quant.permute(0, 1, 3, 2).contiguous() - del value_token_major, value_token_quant - else: - key_quant, key_scale = per_tensor_quant(key, quant_dtype=quant_dtype) - value_quant, value_scale = per_tensor_quant(value, quant_dtype=quant_dtype) - del key, value - key_cache = ( - key_quant.view( - num_blocks, - num_kv_heads, - block_size, - head_dim // 16, - 16, - ) - .permute(0, 1, 3, 2, 4) - .contiguous() - ) - if trans_v: - value_cache = ( - value_quant.permute(0, 1, 3, 2) - .contiguous() - .view(num_blocks, num_kv_heads, block_size // 16, 16, head_dim) - .permute(0, 1, 2, 4, 3) - .contiguous() - ) - else: - value_cache = value_quant.contiguous() - block_tables = torch.arange(num_blocks, dtype=torch.int32).reshape( - batch_size, blocks_per_sequence - ) - context_lengths = torch.full((batch_size,), context_length, dtype=torch.int32) - output = torch.empty_like(query) - - reference = run_torch( - query, - key_quant, - value_quant, - block_tables, - context_lengths, - key_scale, - value_scale, - query_length=query_length, - ) - # Only the kernel's cache layouts need to survive into allocation rotation. - del key_quant, value_quant - - query_group_size = num_query_heads // num_kv_heads - partial_shape = ( - batch_size, - num_kv_heads, - num_partitions, - query_length * query_group_size, - ) - pmax = torch.empty(partial_shape, dtype=dtypes.fp32) - psum = torch.empty_like(pmax) - pout = torch.empty(*partial_shape, head_dim, dtype=dtype) - softmax_scale = head_dim**-0.5 - - candidates = {"flydsl": _run_flydsl} - - # QK and PV each perform one multiply-add per visible query/KV-token pair. - attended_tokens = sum( - max(0, context_length - query_length + 1 + position) - for position in range(query_length) - ) - flops = 4 * batch_size * num_query_heads * attended_tokens * head_dim - # Effective bandwidth counts Q + O + referenced K/V tokens and metadata once. - # It excludes padded tokens, repeated loads, and partition scratch traffic. - # All MTP positions share the KV cache and its scales: count them once, - # including only valid tokens rather than a partially padded final page. - 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 - * key_cache.element_size() - + block_tables.numel() * block_tables.element_size() - + context_lengths.numel() * context_lengths.element_size() - + scale_elements * (key_scale.element_size() + value_scale.element_size()) - ) - - ret = { - "gfx": get_gfx_runtime(), - "partitions": num_partitions, - "trans_v": trans_v, - "per_token": per_token, - } - for name, fn in candidates.items(): - # Pass tensors explicitly so perftest can rotate their allocations. - out, us = run_perftest( - fn, - output, - query, - key_cache, - value_cache, - block_tables, - context_lengths, - key_scale, - value_scale, - num_partitions, - softmax_scale, - pmax, - psum, - pout, - ) - err = checkAllclose( - reference.to(dtypes.fp32), - out.to(dtypes.fp32), - rtol=FIXED_LENGTH_ACCURACY_TOLERANCE, - atol=FIXED_LENGTH_ACCURACY_TOLERANCE, - tol_err_ratio=0.0, - msg=f"{name}: pa_decode", - ) - if err: - raise AssertionError(f"{name}: pa_decode mismatch ratio {err}") - ret[f"{name} us"] = us - # Amortized time per output token, not the latency of an MTP step. - ret[f"{name} us/token"] = us / (batch_size * query_length) - ret[f"{name} TFLOPS"] = flops / us / 1e6 - ret[f"{name} TB/s"] = nbytes / us / 1e6 - ret[f"{name} err"] = err - return ret - - -PA_DECODE_TEST_CASES = [ - pytest.param( - { - "batch_size": 3, - "num_query_heads": 8, - "num_kv_heads": 1, - "head_dim": 128, - "context_length": 257, - "block_size": 16, - "dtype": dtypes.bf16, - "trans_v": False, - "per_token": False, - "query_length": 1, - "num_partitions": None, - }, - id="decode-auto-page16-bf16-scalar", - ), - pytest.param( - { - "batch_size": 3, - "num_query_heads": 8, - "num_kv_heads": 2, - "head_dim": 128, - "context_length": 1027, - "block_size": 64, - "dtype": dtypes.fp16, - "trans_v": True, - "per_token": False, - "query_length": 1, - "num_partitions": 1, - }, - id="decode-np1-page64-fp16-trans-v", - ), - pytest.param( - { - "batch_size": 3, - "num_query_heads": 8, - "num_kv_heads": 2, - "head_dim": 128, - "context_length": 257, - "block_size": 128, - "dtype": dtypes.bf16, - "trans_v": False, - "per_token": True, - "query_length": 1, - "num_partitions": 4, - }, - id="decode-reduce-page128-per-token", - ), - pytest.param( - { - "batch_size": 16, - "num_query_heads": 16, - "num_kv_heads": 2, - "head_dim": 128, - "context_length": 257, - "block_size": 128, - "dtype": dtypes.bf16, - "trans_v": True, - "per_token": True, - "query_length": 1, - "num_partitions": 8, - "gfx950_prefetch_v": "occupancy", - }, - id="decode-hkv2-one-wg-per-cu-prefetch", - ), - pytest.param( - { - "batch_size": 32, - "num_query_heads": 16, - "num_kv_heads": 2, - "head_dim": 128, - "context_length": 257, - "block_size": 128, - "dtype": dtypes.bf16, - "trans_v": True, - "per_token": True, - "query_length": 1, - "num_partitions": 8, - "gfx950_prefetch_v": "occupancy", - }, - id="decode-hkv2-two-wg-per-cu-no-prefetch", - ), - pytest.param( - { - "batch_size": 32, - "num_query_heads": 16, - "num_kv_heads": 2, - "head_dim": 128, - "context_length": 257, - "block_size": 16, - "dtype": dtypes.bf16, - "trans_v": True, - "per_token": True, - "query_length": 1, - "num_partitions": 8, - "gfx950_prefetch_v": True, - }, - id="decode-hkv2-page16-prefetch", - ), - pytest.param( - { - "batch_size": 32, - "num_query_heads": 32, - "num_kv_heads": 2, - "head_dim": 128, - "context_length": 257, - "block_size": 128, - "dtype": dtypes.bf16, - "trans_v": False, - "per_token": True, - "query_length": 1, - "num_partitions": 8, - "gfx950_prefetch_v": True, - }, - id="decode-hkv2-plain-v-prefetch", - ), - pytest.param( - { - "batch_size": 2, - "num_query_heads": 16, - "num_kv_heads": 1, - "head_dim": 128, - "context_length": 257, - "block_size": 128, - "dtype": dtypes.bf16, - "trans_v": True, - "per_token": True, - "query_length": 2, - "num_partitions": 3, - "gfx950_query_splits": 2, - "gfx950_prefetch_v": True, - }, - id="mtp2-query-split-page128-per-token", - ), - pytest.param( - { - "batch_size": 2, - "num_query_heads": 16, - "num_kv_heads": 1, - "head_dim": 128, - "context_length": 1027, - "block_size": 128, - "dtype": dtypes.bf16, - "trans_v": False, - "per_token": True, - "query_length": 3, - "num_partitions": 4, - "gfx950_query_splits": 1, - }, - id="mtp3-fused-page128-per-token", - ), - pytest.param( - { - "batch_size": 2, - "num_query_heads": 16, - "num_kv_heads": 1, - "head_dim": 128, - "context_length": 257, - "block_size": 16, - "dtype": dtypes.bf16, - "trans_v": True, - "per_token": True, - "query_length": 4, - "num_partitions": 3, - "gfx950_query_splits": 4, - "gfx950_prefetch_v": True, - }, - id="mtp4-query-split-page16-per-token", - ), - pytest.param( - { - "batch_size": 2, - "num_query_heads": 16, - "num_kv_heads": 1, - "head_dim": 128, - "context_length": 257, - "block_size": 128, - "dtype": dtypes.fp16, - "trans_v": True, - "per_token": False, - "query_length": 4, - "num_partitions": 3, - "gfx950_query_splits": 1, - }, - id="mtp4-reduce-page128-fp16-scalar", - ), - pytest.param( - { - "batch_size": 2, - "num_query_heads": 8, - "num_kv_heads": 1, - "head_dim": 256, - "context_length": 64, - "block_size": 16, - "dtype": dtypes.bf16, - "trans_v": False, - "per_token": False, - "query_length": 1, - "num_partitions": 1, - }, - id="decode-head256-np1", - ), - pytest.param( - { - "batch_size": 1, - "num_query_heads": 16, - "num_kv_heads": 1, - "head_dim": 128, - "context_length": 200000, - "block_size": 128, - "dtype": dtypes.bf16, - "trans_v": True, - "per_token": True, - "query_length": 1, - "num_partitions": 8, - "max_partitions": 4, - }, - id="decode-long-context-200k", - ), - pytest.param( - { - "batch_size": 200, - "num_query_heads": 16, - "num_kv_heads": 1, - "head_dim": 128, - "context_length": 257, - "block_size": 128, - "dtype": dtypes.bf16, - "trans_v": False, - "per_token": True, - "query_length": 4, - "num_partitions": None, - "gfx950_query_splits": 1, - }, - id="mtp4-bs200-auto-smoke", - ), -] - - -@pytest.mark.parametrize("case", PA_DECODE_TEST_CASES) -def test_pa_decode(case, monkeypatch): - """Run one shared real-kernel/reference path across the supported PA axes.""" - _require_gpu() - case = case.copy() - expected_query_splits = case.pop("gfx950_query_splits", None) - expected_prefetch_v = case.pop("gfx950_prefetch_v", None) - selected_query_splits = [] - selected_prefetch_v = [] - - if expected_query_splits is not None or expected_prefetch_v is not None: - module = importlib.import_module("aiter.ops.flydsl.pa_decode") - compile_tile = module.compile_pa_decode_tile - - def capture_compile(**kwargs): - selected_query_splits.append(kwargs["query_splits"]) - selected_prefetch_v.append(kwargs["prefetch_v"]) - return compile_tile(**kwargs) - - monkeypatch.setattr(module, "compile_pa_decode_tile", capture_compile) - - def run_once(fn, *args, **kwargs): - assert fn is _run_flydsl - assert not kwargs - return fn(*args), 1.0 - - monkeypatch.setitem(globals(), "run_perftest", run_once) - - explicit_partitions = case["num_partitions"] - max_partitions = case.get("max_partitions") - if explicit_partitions is None: - expected_partitions = get_recommended_splits( - case["batch_size"], - case["num_kv_heads"], - split_kv_blocks=KV_COMPUTE_BLOCK // case["block_size"], - max_partitions=max_partitions, - max_context_length=case["context_length"], - ) - else: - expected_partitions = explicit_partitions - - result = run_pa_decode_tile_case(**case) - assert result["partitions"] == expected_partitions - assert result["flydsl err"] == 0 - assert result["flydsl us"] == 1.0 - assert result["flydsl us/token"] == pytest.approx( - 1.0 / (case["batch_size"] * case["query_length"]) - ) - assert result["flydsl TFLOPS"] > 0 - assert result["flydsl TB/s"] > 0 - - if expected_query_splits is not None: - expected = expected_query_splits if get_gfx_runtime() == "gfx950" else 1 - assert selected_query_splits - assert set(selected_query_splits) == {expected} - if expected_prefetch_v is not None: - if expected_prefetch_v == "occupancy": - workgroups = case["batch_size"] * case["num_kv_heads"] * expected_partitions - expected_prefetch_v = ( - workgroups - <= torch.cuda.get_device_properties( - torch.cuda.current_device() - ).multi_processor_count - ) - expected = expected_prefetch_v if get_gfx_runtime() == "gfx950" else False - assert selected_prefetch_v - assert set(selected_prefetch_v) == {expected} - - -def _assert_plan(plan, lengths): - work = plan.work_info.cpu().tolist() - reductions = plan.reduce_info.cpu().tolist() - last_tiles = [(max(length, 0) + 255) // 256 for length in lengths] - first_tiles = [ - ( - max(0, length - (plan.query_length - 1) - plan.sliding_window) // 256 - if plan.sliding_window > 0 - else 0 - ) - for length in lengths - ] - tiles_per_sequence = [last - first for first, last in zip(first_tiles, last_tiles)] - nonempty = [int(tiles > 0) for tiles in tiles_per_sequence] - total_tiles = max(sum(tiles_per_sequence), 1) - remaining = plan.capacity - sum(nonempty) - cumulative_tiles = 0 - expected_counts = [] - for tiles, is_nonempty in zip(tiles_per_sequence, nonempty): - lower = cumulative_tiles * remaining // total_tiles - cumulative_tiles += tiles - upper = cumulative_tiles * remaining // total_tiles - expected_counts.append( - min(is_nonempty + upper - lower, tiles, plan.max_partitions) - ) - - expected_reductions = [] - expected_work = [] - offset = 0 - for seq, (length, first, last, tiles, count) in enumerate( - zip(lengths, first_tiles, last_tiles, tiles_per_sequence, expected_counts) - ): - start = offset - expected_reductions.append([start, count]) - for part in range(count): - expected_work.append( - [ - seq, - first + part * tiles // count, - first + (part + 1) * tiles // count, - length, - ] - ) - - actual_start, actual_count = reductions[seq] - assert (actual_start, actual_count) == (start, count) - assert 0 <= actual_count <= min(tiles, plan.max_partitions) - assert (actual_count > 0) == (length > 0) - previous_end = first - for task in work[actual_start : actual_start + actual_count]: - assert task[0] == seq and task[3] == length - assert task[1] == previous_end and task[1] < task[2] <= last - previous_end = task[2] - assert previous_end == last - offset += actual_count - expected_work.extend([[0, 0, 0, 0]] * (plan.capacity - offset)) - assert reductions == expected_reductions - assert work == expected_work - assert offset <= plan.capacity - assert all(row == [0, 0, 0, 0] for row in work[offset:]) - - -@pytest.mark.parametrize( - "heads,max_parts,budget", [(1, 1, 512), (1, 7, 17), (1, 256, 512), (2, 256, 512)] -) -@pytest.mark.parametrize( - "lengths", - [ - [0], - [1], - [257], - [200003], - [0] * 8, - [0, 1, 3, 4, 257, 4095, 16385, 200003], - [100000] * 8, - [i * 7919 % 200003 for i in range(200)], - ], -) -def test_plan_covers_each_token_once(lengths, heads, max_parts, budget): - context = torch.tensor(lengths, dtype=torch.int32) - plan = plan_pa_decode( - context, heads, max_partitions=max_parts, workgroup_budget=budget - ) - _assert_plan(plan, lengths) - - -@pytest.mark.parametrize( - "sliding_window,query_length", - [ - (1, 1), - (1, 4), - (255, 2), - (256, 3), - (257, 4), - (509, 4), - (8192, 4), - (2**40, 4), - (0, 4), - (-1, 2), - ], -) -@pytest.mark.parametrize( - "heads,max_parts,budget", [(1, 1, 32), (1, 7, 17), (2, 256, 512)] -) -def test_sliding_window_plan_covers_absolute_tiles( - sliding_window, query_length, heads, max_parts, budget -): - # Include MTP unions crossing a 256-token boundary, negative/empty contexts, - # and int32's maximum length without allocating a corresponding KV cache. - lengths = [ - -1, - 0, - 1, - 2, - 3, - 4, - 255, - 256, - 257, - 258, - 259, - 511, - 512, - 513, - 514, - 515, - 8193, - 200003, - 2**31 - 1, - ] - context = torch.tensor(lengths, dtype=torch.int32) - plan = plan_pa_decode( - context, - heads, - max_partitions=max_parts, - workgroup_budget=budget, - query_length=query_length, - sliding_window=sliding_window, - ) - assert plan.sliding_window == max(sliding_window, 0) - assert plan.query_length == query_length - _assert_plan(plan, lengths) - - -@pytest.mark.parametrize("sliding_window,query_length", [(0, 1), (1, 4), (257, 3)]) -def test_plan_graph_refresh_overwrites_old_metadata(sliding_window, query_length): - lengths = [200003] * 8 - context = torch.tensor(lengths, dtype=torch.int32) - plan_options = { - "sliding_window": sliding_window, - "query_length": query_length, - } - plan = plan_pa_decode(context, 1, **plan_options) - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - plan_pa_decode(context, 1, plan=plan, **plan_options) - for lengths in ([0] * 8, [0, 1, 3, 4, 257, 4096, 16385, 200003], [200003] * 8): - context.copy_(torch.tensor(lengths, dtype=torch.int32)) - graph.replay() - _assert_plan(plan, lengths) - - -def _planned_call(*args, **kwargs): - plan = plan_pa_decode( - args[4], - args[2].shape[1], - max_partitions=args[8], - query_length=args[7], - sliding_window=kwargs.get("sliding_window", 0), - ) - pa_decode(*args, work_plan=plan, **kwargs) - - -@pytest.mark.parametrize( - "query_length,sliding_window,block_size,trans_v,max_parts,heads", - [ - (1, 1, 16, False, 1, 1), - (1, 257, 128, True, 7, 2), - (2, 1, 16, True, 7, 1), - (2, 255, 64, False, 64, 2), - (2, 8192, 128, True, 256, 1), - (3, 1, 128, False, 1, 2), - (3, 257, 16, True, 7, 1), - (3, 513, 64, True, 86, 1), - (4, 1, 16, False, 7, 1), - (4, 256, 128, True, 7, 2), - (4, 257, 128, False, 256, 1), - (4, 8192, 16, True, 1, 1), - (4, 2**40, 16, False, 7, 1), - ], -) -def test_planned_sliding_window_sparse_causal_reference( - monkeypatch, query_length, sliding_window, block_size, trans_v, max_parts, heads -): - """Each query has its own window even when its tile is shared with MTP peers.""" - monkeypatch.setattr(torch.ops.aiter, "pa_decode_flydsl", _planned_call) - _run_mtp4_fused_reference_case( - [0, 1, 2, 3, 4, 255, 256, 257, 258, 259, 513, 514, 515, 4099], - heads, - block_size, - trans_v, - max_parts, - query_length=query_length, - sliding_window=sliding_window, - ) - - -@pytest.mark.parametrize( - "query_length,sliding_window,block_size,trans_v,parts", - [ - (1, 1, 16, False, 1), - (2, 257, 128, True, 7), - (3, 1, 64, False, 7), - (4, 257, 128, True, 7), - (4, 8192, 16, False, 1), - (4, 2**31 - 1, 16, False, 1), - (4, 2**31, 16, False, 1), - ], -) -def test_static_sliding_window_sparse_causal_reference( - query_length, sliding_window, block_size, trans_v, parts -): - """The unplanned path, including query-split kernels, uses the same mask.""" - _run_mtp4_fused_reference_case( - [0, 1, 2, 3, 4, 255, 256, 257, 259, 513, 514, 4099], - 1, - block_size, - trans_v, - parts, - query_length=query_length, - sliding_window=sliding_window, - ) - - -@pytest.mark.parametrize( - "sliding_window,head_dim,block_size,trans_v,parts,dtype", - [ - (1, 64, 16, False, 1, torch.bfloat16), - (1, 128, 16, True, 7, torch.bfloat16), - (257, 128, 128, True, 7, torch.float16), - (257, 128, 128, True, 7, torch.bfloat16), - (513, 256, 64, True, 86, torch.float16), - (8192, 1024, 128, False, 256, torch.bfloat16), - ], -) -def test_planned_sliding_window_per_tensor_scales( - monkeypatch, sliding_window, head_dim, block_size, trans_v, parts, dtype -): - # GQA8 + BF16 + D128 + transposed V also exercises scalar FP8 decode on gfx950. - monkeypatch.setattr(torch.ops.aiter, "pa_decode_flydsl", _planned_call) - _run_accuracy_case( - [0, 1, 255, 256, 257, 513, 4099], - num_query_heads=16, - num_kv_heads=2, - head_dim=head_dim, - block_size=block_size, - query_dtype=dtype, - num_partitions=parts, - tolerance=0.005, - trans_v=trans_v, - sliding_window=sliding_window, - ) - - -@pytest.mark.parametrize("query_length", [1, 4]) -@pytest.mark.parametrize("planned", [False, True]) -def test_sliding_window_ignores_outside_value_scales(query_length, planned): - """Masked prefix scales cannot affect FP8 P normalization in the last tile.""" - context_length, block_size, head_dim = 511, 128, 128 - num_pages, num_kv_heads, query_group_size = 4, 1, 16 - quant_dtype = _quant_dtype() - context = torch.tensor([context_length], dtype=torch.int32) - block_tables = torch.arange(num_pages, dtype=torch.int32).reshape(1, num_pages) - query = torch.ones((query_length, query_group_size, head_dim), dtype=torch.bfloat16) - output = torch.full_like(query, float("nan")) - key_cache = torch.zeros( - (num_pages, num_kv_heads, head_dim // 16, block_size, 16), dtype=quant_dtype - ) - value_tokens = torch.zeros( - (num_pages, num_kv_heads, block_size, head_dim), dtype=quant_dtype - ) - scale_shape = (num_pages, num_kv_heads, block_size, 1) - key_scale = torch.ones(scale_shape, dtype=torch.float32) - value_scale = torch.full(scale_shape, 1e9, dtype=torch.float32) - first_visible = context_length - query_length - value_scale.reshape(-1)[first_visible:].fill_(1.0) - # Every query sees exactly its corresponding tail token. Prefix V=0 keeps - # dequantized values finite even though the excluded prefix scales are huge. - for position in range(query_length): - token = first_visible + position - value_tokens[token // block_size, 0, token % block_size].fill_( - (position + 1) / query_length - ) - value_cache = ( - value_tokens.view(num_pages, num_kv_heads, block_size // 16, 16, head_dim) - .permute(0, 1, 2, 4, 3) - .contiguous() - ) - plan = ( - plan_pa_decode( - context, - num_kv_heads, - sliding_window=1, - query_length=query_length, - max_partitions=7, - ) - if planned - else None - ) - pa_decode( - output, - query, - key_cache, - value_cache, - context, - block_tables, - head_dim**-0.5, - query_length, - 7, - KV_COMPUTE_BLOCK, - quant_dtype, - None, - key_scale, - value_scale, - sliding_window=1, - work_plan=plan, - ) - expected = ( - (torch.arange(1, query_length + 1, dtype=query.dtype) / query_length) - .reshape(query_length, 1, 1) - .expand_as(query) - ) - assert torch.isfinite(output).all() - torch.testing.assert_close(output, expected, atol=0.005, rtol=0.005) - - -@pytest.mark.parametrize("block_size", [16, 128]) -@pytest.mark.parametrize("trans_v", [False, True]) -@pytest.mark.parametrize("max_parts", [1, 7, 64, 86, 256]) -@pytest.mark.parametrize("heads", [1, 2]) -def test_planned_mtp4_sparse_causal_reference( - monkeypatch, block_size, trans_v, max_parts, heads -): - monkeypatch.setattr(torch.ops.aiter, "pa_decode_flydsl", _planned_call) - _run_mtp4_fused_reference_case( - [0, 1, 2, 3, 4, 255, 256, 257, 2305, 16387], - heads, - block_size, - trans_v, - max_parts, - ) - - -@pytest.mark.parametrize( - "head_dim,block_size,parts,dtype", - [ - (64, 16, 8, torch.bfloat16), - (256, 64, 86, torch.float16), - (1024, 128, 256, torch.bfloat16), - (128, 128, 256, torch.float16), - ], -) -def test_planned_decode_other_reducer_paths( - monkeypatch, head_dim, block_size, parts, dtype -): - monkeypatch.setattr(torch.ops.aiter, "pa_decode_flydsl", _planned_call) - _run_accuracy_case( - [0, 1, 257, 4099], - num_query_heads=8, - num_kv_heads=2, - head_dim=head_dim, - block_size=block_size, - query_dtype=dtype, - num_partitions=parts, - tolerance=0.005, - trans_v=True, - ) - - -@pytest.mark.parametrize("long_context", [16384, 16385]) -def test_planned_decode_register_reducer_chunk_boundary(monkeypatch, long_context): - """Exercise the generic reducer with 64 and 65 active partitions.""" - monkeypatch.setattr(torch.ops.aiter, "pa_decode_flydsl", _planned_call) - _run_accuracy_case( - [0, 1, 257, long_context], - num_query_heads=8, - num_kv_heads=2, - head_dim=256, - block_size=64, - query_dtype=torch.float16, - num_partitions=86, - tolerance=0.005, - trans_v=True, - ) - - -@pytest.mark.parametrize( - "sliding_window,query_length", [(0, 4), (1, 4), (257, 3), (8193, 4)] -) -def test_planned_decode_graph_replay_with_poisoned_scratch( - monkeypatch, sliding_window, query_length + sliding_window=0, + sinks=None, ): - def capture_and_replay(*args, **kwargs): - output, query, key, _value, context = args[:5] - heads, rows = ( - key.shape[1], - query.shape[0] // context.numel() * query.shape[1] // key.shape[1], - ) - plan_options = { - "max_partitions": args[8], - "sliding_window": sliding_window, - "query_length": query_length, - } - plan = plan_pa_decode(context, heads, **plan_options) - shape = (heads, plan.capacity, rows) - psum = torch.full(shape, float("nan"), dtype=torch.float32) - pmax = torch.full_like(psum, float("nan")) - pout = torch.full((*shape, query.shape[2]), float("nan"), dtype=query.dtype) - - def run(): - plan_pa_decode(context, heads, plan=plan, **plan_options) - pa_decode( - *args, - **kwargs, - work_plan=plan, - exp_sums=psum, - max_logits=pmax, - temporary_output=pout, - ) - - run() # Compile before capture. - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - run() - active_tasks = int(plan.reduce_info[:, 1].sum().item()) - assert active_tasks < plan.capacity - for scratch in (pmax, psum, pout): - assert torch.isnan(scratch[:, active_tasks:]).all() - original = context.clone() - for scratch in (pmax, psum, pout): - scratch.fill_(float("nan")) - context.zero_() - graph.replay() - _assert_plan(plan, [0] * context.numel()) - assert torch.equal(output, torch.zeros_like(output)) - for scratch in (pmax, psum, pout): - assert torch.isnan(scratch).all() - context.copy_(original) - graph.replay() - _assert_plan(plan, original.cpu().tolist()) - for scratch in (pmax, psum, pout): - assert torch.isnan(scratch[:, active_tasks:]).all() - - monkeypatch.setattr(torch.ops.aiter, "pa_decode_flydsl", capture_and_replay) - _run_mtp4_fused_reference_case( - [0, 257, 4096, 200003], - 1, - 128, - True, - 256, - query_length=query_length, - sliding_window=sliding_window, - ) - - -def test_plan_rejects_incompatible_reuse(): - context = torch.tensor([1, 257], dtype=torch.int32) - plan = plan_pa_decode(context, 1, max_partitions=7) - with pytest.raises(ValueError, match="max_partitions"): - plan_pa_decode(context, 1, plan=plan) - with pytest.raises(ValueError, match="KV head"): - plan_pa_decode(context, 2, max_partitions=7, plan=plan) - with pytest.raises(ValueError, match="shape"): - plan_pa_decode(context[:1], 1, max_partitions=7, plan=plan) - + """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) -def test_sliding_window_plan_rejects_incompatible_reuse(): - context = torch.tensor([1, 257, 4099], dtype=torch.int32) - plan = plan_pa_decode( - context, 1, sliding_window=257, query_length=4, max_partitions=7 - ) - for window in (0, 256): - with pytest.raises(ValueError, match="sliding_window"): - plan_pa_decode( - context, - 1, - sliding_window=window, - query_length=4, - max_partitions=7, - plan=plan, + 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) ) - with pytest.raises(ValueError, match="query_length"): - plan_pa_decode(context, 1, sliding_window=257, max_partitions=7, plan=plan) - reused = plan_pa_decode( - context, 1, sliding_window=257, query_length=4, max_partitions=7, plan=plan + # A +inf sink suppresses all finite KV logits. Fully masked rows need + # an explicit zero because -inf - -inf is undefined without a sink. + 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): + """Build one sparse/dense paged cache, launch arguments and reference call.""" + 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 ) - assert reused is plan - _assert_plan(plan, context.cpu().tolist()) - - -@pytest.mark.parametrize( - "sliding_window,query_length,error,match", - [ - (-2, 1, ValueError, "sliding_window"), - (1.5, 1, TypeError, "sliding_window"), - (257, 0, ValueError, "query_length"), - (257, 1.5, TypeError, "query_length"), - ], -) -def test_sliding_window_plan_rejects_invalid_arguments( - sliding_window, query_length, error, match -): - context = torch.tensor([257], dtype=torch.int32) - with pytest.raises(error, match=match): - plan_pa_decode( - context, 1, sliding_window=sliding_window, query_length=query_length - ) - - -@pytest.mark.parametrize( - "plan_window,plan_query_length,sliding_window,query_length,match", - [ - (0, 4, 257, 4, "sliding_window"), - (257, 4, 0, 4, "sliding_window"), - (257, 4, 256, 4, "sliding_window"), - (257, 3, 257, 4, "query_length"), - (257, 4, 257, 1, "query_length"), - ], -) -def test_planned_decode_rejects_sliding_window_mismatch( - monkeypatch, plan_window, plan_query_length, sliding_window, query_length, match -): - def incompatible_plan(*args, **kwargs): - plan = plan_pa_decode( - args[4], - args[2].shape[1], - max_partitions=args[8], - sliding_window=plan_window, - query_length=plan_query_length, - ) - pa_decode(*args, work_plan=plan, **kwargs) + query = torch.empty((batch * ql, heads, dim), dtype=case.dtype).uniform_(-0.5, 0.5) + if case.masked_scale: + query.zero_() # Also exercise online Q quantization with zero Q scale. + 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 - monkeypatch.setattr(torch.ops.aiter, "pa_decode_flydsl", incompatible_plan) - with pytest.raises(ValueError, match=match): - _run_mtp4_fused_reference_case( - [0, 1, 2, 3, 4, 257], - 1, - 16, - False, - 7, - query_length=query_length, - sliding_window=sliding_window, + 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()) + # Exact binary ratios isolate W=1 masking from existing FP8 P rounding. + value_scale = ( + factors * 2**-10 if case.sliding_window == 1 else value_scale * factors ) - -@pytest.mark.parametrize( - "plan_window,sliding_window,query_length", [(-1, 0, 4), (0, -1, 3)] -) -def test_planned_decode_accepts_disabled_sliding_window_aliases( - monkeypatch, plan_window, sliding_window, query_length -): - def dense_plan(*args, **kwargs): - # A default dense plan is still valid for any number of causal queries. - plan = plan_pa_decode( - args[4], - args[2].shape[1], - max_partitions=args[8], - sliding_window=plan_window, - ) - assert plan.sliding_window == 0 and plan.query_length == 1 - assert ( - plan_pa_decode( - args[4], - args[2].shape[1], - max_partitions=args[8], - plan=plan, - sliding_window=sliding_window, - query_length=query_length, + 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 covers short rows and makes repeated sink + # mass across multiple partitions observable without cancellation. + # Per-tensor cases keep their 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 ) - is plan - ) - pa_decode(*args, work_plan=plan, **kwargs) + 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 - monkeypatch.setattr(torch.ops.aiter, "pa_decode_flydsl", dense_plan) - _run_mtp4_fused_reference_case( - [0, 1, 2, 3, 4, 257, 513], - 1, - 16, - False, - 7, - query_length=query_length, - sliding_window=sliding_window, + 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 -def _make_sink_case( - lengths, - *, - query_length=4, - head_dim=128, - num_kv_heads=1, - query_group_size=16, - query_dtype=torch.bfloat16, - per_token=True, - num_partitions=7, - sliding_window=0, -): - """Build exact zero-Q/K data and a length-aware, overflow-safe sink oracle.""" - block_size = 128 - num_pages = (max(max(lengths), 1) + block_size - 1) // block_size - num_query_heads = num_kv_heads * query_group_size - context = torch.tensor(lengths, dtype=torch.int32) - query = torch.zeros( - (len(lengths) * query_length, num_query_heads, head_dim), dtype=query_dtype - ) - output = torch.full_like(query, float("nan")) - quant_dtype = _quant_dtype() - key_cache = torch.zeros( - (num_pages, num_kv_heads, head_dim // 16, block_size, 16), dtype=quant_dtype - ) - token_values = ( - torch.arange(num_pages * block_size, dtype=torch.float32) % 4 + 1 - ) * 0.25 - kv_factors = torch.arange(1, num_kv_heads + 1, dtype=torch.float32) - values = ( - ( - token_values.reshape(num_pages, 1, block_size, 1) - * kv_factors.reshape(1, num_kv_heads, 1, 1) - ) - .expand(num_pages, num_kv_heads, block_size, head_dim) - .to(quant_dtype) + 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 = ( - values.reshape(num_pages, num_kv_heads, block_size // 16, 16, head_dim) + 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: + 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 ) - scale_shape = (num_pages, num_kv_heads, block_size, 1) if per_token else (1,) - key_scale = torch.ones(scale_shape, dtype=torch.float32) - value_scale = torch.ones_like(key_scale) - block_tables = ( - torch.arange(num_pages, dtype=torch.int32) - .expand(len(lengths), num_pages) - .contiguous() + 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 = ( - output, + torch.full_like(query, float("nan")), query, key_cache, value_cache, context, - block_tables, - head_dim**-0.5, - query_length, - num_partitions, + 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 reference(sinks): - # The shared pages encode [.25, .5, .75, 1] * (KV head + 1). - # Unlike exp(sinks), logaddexp stays finite for large finite logits. - positions = torch.arange(query_length, dtype=torch.int32) - visible = (context[:, None] - query_length + 1 + positions).clamp_min(0) - first = (visible - sliding_window).clamp_min(0) if sliding_window > 0 else 0 - count = (visible - first).float() - - def prefix_sum(tokens): - remainder = (tokens % 4).float() - return (tokens // 4).float() * 2.5 + remainder * (remainder + 1) * 0.125 - value_sum = prefix_sum(visible) - if sliding_window > 0: - value_sum -= prefix_sum(first) - value_sum = value_sum[..., None] * kv_factors.repeat_interleave( - query_group_size - ) - logits = ( - torch.full((num_query_heads,), float("-inf"), dtype=torch.float32) - if sinks is None - else sinks.float() +def _run_flydsl(*args, sliding_window=0, work_plan=None): + if work_plan is None: + torch.ops.aiter.pa_decode_flydsl(*args, sliding_window=sliding_window) + else: + plan_pa_decode( + args[4], + args[2].shape[1], + max_partitions=args[8], + query_length=args[7], + sliding_window=sliding_window, + plan=work_plan, ) - log_denominator = torch.logaddexp(count.log()[..., None], logits) - expected = value_sum * torch.exp(-log_denominator) - expected.masked_fill_(count[..., None] == 0, 0) - return expected.reshape(-1, num_query_heads, 1).expand_as(query) - - return args, reference - - -def _analytic_sinks(num_heads, dtype): - values = [float("-inf"), float("inf"), -1000, 1000, -2, 0, 0.5, 5] - if dtype == torch.float32: - values += [-3e38, 3e38] - return ( - torch.tensor(values, dtype=dtype) - .repeat((num_heads + len(values) - 1) // len(values))[:num_heads] - .contiguous() - ) + pa_decode(*args, sliding_window=sliding_window, work_plan=work_plan) + return args[0] -@pytest.mark.parametrize( - "planned,query_length,parts,window,block_size,trans_v,heads,sink_dtype", - [ - (False, 1, 1, 0, 16, False, 1, torch.float32), - (False, 4, 1, 257, 128, True, 1, torch.bfloat16), - (True, 4, 1, 1, 16, False, 1, torch.float16), - (True, 2, 7, 257, 128, True, 2, torch.float32), - (False, 3, 7, 0, 64, False, 2, torch.bfloat16), - (True, 4, 7, 0, 128, True, 1, torch.float32), - (False, 2, 86, 1, 16, True, 1, torch.float16), - (True, 1, 86, 257, 64, False, 2, torch.bfloat16), - (False, 4, 256, 257, 128, False, 1, torch.float32), - (True, 3, 256, 0, 16, True, 1, torch.float16), - (True, 1, 7, 1, 128, True, 1, torch.float32), - (False, 4, 1, 0, 16, False, 1, torch.float16), - ], -) -def test_sinks_sparse_causal_reference( - monkeypatch, - planned, - query_length, - parts, - window, - block_size, - trans_v, - heads, - sink_dtype, -): - if planned: - monkeypatch.setattr(torch.ops.aiter, "pa_decode_flydsl", _planned_call) - sinks = torch.linspace(-3, 6, heads * 16, dtype=sink_dtype) - sinks[0] = float("-inf") - _run_mtp4_fused_reference_case( - [0, 1, 2, 3, 4, 255, 256, 257, 259, 513, 4099], - heads, - block_size, - trans_v, - parts, - query_length=query_length, - sliding_window=window, - sinks=sinks, +def _assert_close(output, reference): + assert torch.isfinite(output).all() + torch.testing.assert_close( + output.float(), + reference.float(), + atol=ACCURACY_TOLERANCE, + rtol=ACCURACY_TOLERANCE, ) -@pytest.mark.parametrize( - "planned,parts,head_dim,block_size,trans_v,query_dtype,sink_dtype,window", - [ - (False, 1, 64, 16, False, torch.float16, torch.float32, 0), - (True, 7, 256, 64, True, torch.float16, torch.bfloat16, 257), - (False, 86, 128, 128, True, torch.bfloat16, torch.float16, 1), - (True, 256, 256, 128, False, torch.bfloat16, torch.float32, 0), - (True, 1, 128, 16, True, torch.bfloat16, torch.float32, 257), - (False, 256, 64, 128, True, torch.float16, torch.float16, 257), - ], -) -def test_sinks_per_tensor_reference( - monkeypatch, - planned, - parts, - head_dim, - block_size, - trans_v, - query_dtype, - sink_dtype, - window, -): - if planned: - monkeypatch.setattr(torch.ops.aiter, "pa_decode_flydsl", _planned_call) - sinks = torch.linspace(-3, 7, 16, dtype=sink_dtype) - _run_accuracy_case( - [0, 1, 257, 16385], - num_query_heads=16, - num_kv_heads=2, - head_dim=head_dim, - block_size=block_size, - query_dtype=query_dtype, - num_partitions=parts, - tolerance=0.005, - trans_v=trans_v, - sliding_window=window, - sinks=sinks, - ) +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 -@pytest.mark.parametrize("planned", [False, True]) -@pytest.mark.parametrize( - "query_length,parts,head_dim,window,query_dtype,sink_dtype,heads,group,per_token", - [ - (1, 1, 128, 0, torch.bfloat16, torch.float32, 2, 8, False), - (4, 1, 128, 1, torch.bfloat16, torch.bfloat16, 1, 16, True), - (2, 7, 64, 257, torch.float16, torch.float32, 2, 4, False), - (3, 7, 128, 0, torch.bfloat16, torch.float16, 2, 8, True), - (4, 86, 256, 257, torch.float16, torch.bfloat16, 1, 8, False), - (1, 256, 128, 0, torch.bfloat16, torch.float16, 1, 16, True), - (4, 256, 64, 1, torch.bfloat16, torch.float32, 2, 4, True), - (2, 86, 256, 0, torch.bfloat16, torch.bfloat16, 2, 8, False), - ], -) -def test_sinks_counted_once_per_query( - monkeypatch, - planned, - query_length, - parts, - head_dim, - window, - query_dtype, - sink_dtype, - heads, - group, - per_token, +def _assert_contracts(args, options): + """Reuse the valid inputs for API validation instead of separate fixtures.""" + 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) + plan = options["work_plan"] + if plan is not None: + context, heads = args[4], args[2].shape[1] + reuse = { + "max_partitions": plan.max_partitions, + "sliding_window": plan.sliding_window, + "query_length": plan.query_length, + "plan": plan, + } + 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, ): - """Different head logits, including extremes, add one unscaled zero-V token.""" - args, reference = _make_sink_case( - [0, 1, 3, 4, 257, 513], - query_length=query_length, - head_dim=head_dim, - num_kv_heads=heads, - query_group_size=group, - query_dtype=query_dtype, - per_token=per_token, - num_partitions=parts, - sliding_window=window, - ) - sinks = _analytic_sinks(heads * group, sink_dtype) - module = importlib.import_module("aiter.ops.flydsl.pa_decode") - compile_tile = module.compile_pa_decode_tile - direct_sink_flags = [] - - def capture_compile(**kwargs): - direct_sink_flags.append(kwargs["use_sinks"]) - return compile_tile(**kwargs) - - monkeypatch.setattr(module, "compile_pa_decode_tile", capture_compile) - call = _planned_call if planned else pa_decode - call(*args, sinks=sinks, sliding_window=window) - assert direct_sink_flags and set(direct_sink_flags) == {not planned and parts == 1} - assert torch.isfinite(args[0]).all() - torch.testing.assert_close( - args[0].float(), reference(sinks), atol=0.005, rtol=0.005 - ) + # 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, + ) + + +# Each row uses the same eager/graph checks with static and packed plans. +BF16, FP16, FP32 = torch.bfloat16, torch.float16, torch.float32 +CASES = [ + _case("scalar-direct", (1, 2, 8, 128), (16, 1, 0), 1), + _case("scalar-window-sinks", (1, 2, 8, 128), (128, 1, 0), 7, 257, FP32), + _case( + "register-64-65", + (1, 2, 4, 256), + (64, 1, 0), + 86, + sink=FP16, + dtype=FP16, + lengths=(0, 1, 257, 16384, 16385), + ), + _case("head64-fp16", (2, 2, 4, 64), (16, 0, 0), 7, 257, BF16, dtype=FP16), + _case("head1024", (1, 2, 4, 1024), (128, 1, 0), 256, 8192, FP32), + _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("mtp3-window", (3, 1, 16, 128), (128, 0, 1), 7, 257, FP16), + _case("mtp2-odd-parts", (2, 2, 16, 128), (16, 0, 1), 3), + _case( + "mtp2-query-split", + (2, 1, 16, 128), + parts=3, + lengths=(257, 259), + expected_query_splits=2, + expected_prefetch=True, + ), + _case( + "hkv2-prefetch-1wg", + (1, 2, 8, 128), + parts=8, + window=1, + sink=FP32, + lengths=(257,) * 16, + expected_prefetch="occupancy", + ), + _case( + "hkv2-prefetch-2wg", + (1, 2, 8, 128), + parts=8, + lengths=(257,) * 32, + expected_prefetch="occupancy", + ), + _case( + "hkv2-prefetch-page16", + (1, 2, 8, 128), + (16, 1, 1), + 8, + 257, + FP16, + lengths=(257,) * 32, + expected_prefetch=True, + ), + _case( + "hkv2-prefetch-plain-v", + (1, 2, 16, 128), + (128, 0, 1), + 8, + sink=BF16, + lengths=(257,) * 32, + expected_prefetch=True, + ), + _case("hkv2-direct-sinks", (1, 2, 8, 128), parts=1, window=1, sink=BF16), + _case( + "long-200k", + cache=(128, 0, 1), + parts=256, + sink=FP32, + lengths=(0, 1, 3, 4, 255, 256, 257, 200003), + ), + _case("window-int64", cache=(16, 1, 1), window=2**40, wide_kv_addressing=True), + _case( + "masked-scale-decode", + (1, 1, 16, 128), + window=1, + sink=FP32, + lengths=(511,), + masked_scale=True, + ), + _case("masked-scale-mtp", window=1, lengths=(511,), masked_scale=True), + _case( + "disabled-window-wide", + cache=(16, 0, 1), + parts=256, + window=-1, + sink=FP32, + query_splits=1, + wide_kv_addressing=True, + ), + _case( + "large-batch-auto", + cache=(128, 0, 1), + parts=None, + lengths=(257,) * 200, + expected_query_splits=1, + ), + _case("small-batch-auto", (1, 1, 8, 128), (16, 0, 0), None, lengths=(257,) * 3), + _case( + "exact-parts-override", + (1, 1, 16, 128), + parts=5, + lengths=(200000,), + max_partitions=4, + ), + _case( + "window255-budget", + (2, 2, 16, 128), + (64, 0, 1), + 64, + 255, + BF16, + workgroup_budget=17, + ), + _case("window256", (3, 2, 8, 128), (128, 0, 1), 7, 256), + _case("window509", cache=(64, 0, 1), parts=86, window=509, sink=FP16), + _case("window8192", (2, 1, 16, 128), parts=256, window=8192), + _case("window-int32-max", cache=(16, 0, 1), parts=1, window=2**31 - 1), + _case("window-int32-overflow", cache=(16, 0, 1), parts=1, window=2**31, sink=FP32), + _case( + "long-np64", + cache=(16, 1, 1), + parts=64, + lengths=(0, 3, 257, 16384, 16385, 65537), + ), + _case("head64-mtp-window", (4, 2, 4, 64), (16, 1, 1), 256, 1, FP32), + _case("scalar-fp16-mtp", cache=(128, 1, 0), parts=3, dtype=FP16), + _case( + "fused-narrow", + cache=(16, 0, 1), + parts=1, + 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("empty", window=1, sink=FP16, lengths=(0,) * 4), +] -@pytest.mark.parametrize("query_splits", [1, 4]) -def test_sinks_np1_fused_and_query_split(monkeypatch, query_splits): - """Static NP1 applies sinks in the compute epilogue without a reduction launch.""" - args, reference = _make_sink_case( - [0, 1, 3, 4, 257, 513], num_partitions=1, sliding_window=1 - ) - sinks = _analytic_sinks(args[1].shape[1], torch.float32) +@pytest.mark.parametrize("planned", [False, True], ids=["static", "planned"]) +@pytest.mark.parametrize("case", CASES) +def test_pa_decode(case, planned, monkeypatch): + """One real-kernel path for all layouts, masks, sinks, plans and replays.""" + args, options, reference = _make_inputs(case, planned) + output, query, _, _, context = args[:5] + scratch, sinks, plan = args[14:17], args[-1], options["work_plan"] module = importlib.import_module("aiter.ops.flydsl.pa_decode") compile_tile = module.compile_pa_decode_tile - - def compile_query_split(**kwargs): - kwargs["query_splits"] = query_splits - assert kwargs["use_sinks"] and kwargs["sink_dtype_str"] == "f32" + selected = [] + + def compile_checked(**kwargs): + if case.query_splits is not None and not planned: + kwargs["query_splits"] = case.query_splits + if case.wide_kv_addressing is not None: + kwargs["wide_kv_addressing"] = case.wide_kv_addressing + assert kwargs["use_sinks"] == ( + sinks is not None and args[8] == 1 and not planned + ) + selected.append(kwargs) return compile_tile(**kwargs) def unexpected_reducer(*_args, **_kwargs): - raise AssertionError("static NP1 sinks must not launch a reducer") - - monkeypatch.setattr(module, "compile_pa_decode_tile", compile_query_split) - monkeypatch.setattr(module, "launch_pa_decode_ps_reduce", unexpected_reducer) - pa_decode(*args, sinks=sinks, sliding_window=1) - torch.testing.assert_close( - args[0].float(), reference(sinks), atol=0.005, rtol=0.005 - ) - - -@pytest.mark.parametrize( - "planned,parts,query_length,head_dim,window", - [ - (False, 1, 4, 128, 257), - (False, 7, 1, 64, 0), - (True, 1, 4, 128, 1), - (True, 256, 3, 256, 257), - ], -) -def test_sinks_negative_infinity_matches_none( - planned, parts, query_length, head_dim, window -): - args, reference = _make_sink_case( - [0, 1, 3, 4, 257, 513], - query_length=query_length, - head_dim=head_dim, - num_partitions=parts, - sliding_window=window, - ) - call = _planned_call if planned else pa_decode - call(*args, sinks=None, sliding_window=window) - without_sinks = args[0].clone() - args[0].fill_(float("nan")) - sinks = torch.full((args[1].shape[1],), float("-inf"), dtype=torch.float32) - call(*args, sinks=sinks, sliding_window=window) - torch.testing.assert_close(args[0], without_sinks, atol=0.005, rtol=0.005) - torch.testing.assert_close(args[0].float(), reference(None), atol=0.005, rtol=0.005) - - -@pytest.mark.parametrize( - "planned,head_dim,query_length,window,parts,sink_dtype", - [ - (True, 128, 4, 257, 7, torch.float32), - (True, 256, 3, 0, 86, torch.bfloat16), - (False, 128, 4, 1, 7, torch.float16), - (False, 64, 2, 0, 86, torch.float32), - ], -) -def test_sinks_graph_replay_with_updated_logits_and_poisoned_scratch( - planned, head_dim, query_length, window, parts, sink_dtype -): - args, reference = _make_sink_case( - [513, 4, 1, 0], - query_length=query_length, - head_dim=head_dim, - num_partitions=parts, - sliding_window=window, - ) - output, query, key_cache, _, context = args[:5] - heads = key_cache.shape[1] - rows = query_length * query.shape[1] // heads - sinks = torch.linspace(-3, 6, query.shape[1], dtype=sink_dtype) - plan_options = { - "max_partitions": parts, - "query_length": query_length, - "sliding_window": window, - } - plan = plan_pa_decode(context, heads, **plan_options) if planned else None - shape = ( - (heads, plan.capacity, rows) - if planned - else (context.numel(), heads, parts, rows) - ) - psum = torch.full(shape, float("nan"), dtype=torch.float32) - pmax = torch.full_like(psum, float("nan")) - pout = torch.full((*shape, head_dim), float("nan"), dtype=query.dtype) - - def run(): - if planned: - plan_pa_decode(context, heads, plan=plan, **plan_options) - pa_decode( - *args, - sinks=sinks, - sliding_window=window, - work_plan=plan, - exp_sums=psum, - max_logits=pmax, - temporary_output=pout, - ) - - run() # Compile before capture; all captured pointers remain fixed. + raise AssertionError("static NP=1 must not launch a reducer") + + monkeypatch.setattr(module, "compile_pa_decode_tile", compile_checked) + if not planned and args[8] == 1: + 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()) + 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() - for lengths, updated_sinks in ( - ([0, 1, 4, 513], torch.linspace(7, -4, query.shape[1], dtype=sink_dtype)), - ([0, 0, 0, 0], torch.full_like(sinks, 1000)), - ([513, 4, 1, 0], _analytic_sinks(query.shape[1], sink_dtype)), + _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)) - sinks.copy_(updated_sinks) - output.fill_(float("nan")) - for scratch in (pmax, psum, pout): - scratch.fill_(float("nan")) + 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() - assert torch.isfinite(output).all() - torch.testing.assert_close( - output.float(), reference(sinks), atol=0.005, rtol=0.005 - ) - if planned: - _assert_plan(plan, lengths) - active = int(plan.reduce_info[:, 1].sum().item()) - assert active < plan.capacity - for scratch in (pmax, psum, pout): - assert torch.isnan(scratch[:, active:]).all() - - -@pytest.mark.parametrize( - "invalid,error", - [ - ("not_tensor", TypeError), - ("rank", ValueError), - ("head_count", ValueError), - ("integer_dtype", TypeError), - ("float64_dtype", TypeError), - ("cpu", ValueError), - ("noncontiguous", ValueError), - ], -) -def test_sinks_rejects_invalid_tensor(invalid, error): - args, _ = _make_sink_case([257], query_length=1) - heads = args[1].shape[1] - invalid_sinks = { - "not_tensor": lambda: [0.0] * heads, - "rank": lambda: torch.zeros((1, heads)), - "head_count": lambda: torch.zeros(heads - 1), - "integer_dtype": lambda: torch.zeros(heads, dtype=torch.int32), - "float64_dtype": lambda: torch.zeros(heads, dtype=torch.float64), - "cpu": lambda: torch.zeros(heads, device="cpu"), - "noncontiguous": lambda: torch.zeros(heads * 2)[::2], - }[invalid]() - with pytest.raises(error, match="sinks"): - pa_decode(*args, sinks=invalid_sinks) + check(disabled_sink=step == 2) - -@pytest.mark.parametrize("split_kv_blocks", [2, 16]) -@pytest.mark.parametrize( - "batch_size,context_length,max_partitions,expected", - [ - (1, None, None, 8), - (1, 0, None, 8), - (1, 257, None, 8), - (1, 4096, None, 16), - (1, 16384, None, 64), - (1, 100000, None, 256), - (2, 200000, None, 256), - (3, 200000, None, 171), - (5, 200000, None, 103), - (6, 200000, None, 86), - (7, 200000, None, 74), - (4, 200000, None, 128), - (8, 200000, None, 64), - (1, 200000, 8, 8), - (4, 200000, 5, 5), - (8, 200000, 32, 32), - ], -) -def test_recommended_splits_uses_context_work( - monkeypatch, split_kv_blocks, batch_size, context_length, max_partitions, expected -): - if get_recommended_splits is None: - pytest.skip("FlyDSL is not available") - monkeypatch.setattr( - torch.cuda, - "get_device_properties", - lambda *_args, **_kwargs: type("Props", (), {"multi_processor_count": 256})(), - ) - assert ( - get_recommended_splits( - batch_size, - 1, - split_kv_blocks, - max_partitions, - max_context_length=context_length, + if case.expected_prefetch is not None: + expected = case.expected_prefetch + if expected == "occupancy": + workgroups = len(case.lengths) * case.num_kv_heads * args[8] + expected = ( + workgroups + <= torch.cuda.get_device_properties(query.device).multi_processor_count + ) + expected = bool(expected) and not planned and get_gfx_runtime() == "gfx950" + assert selected and all(config["prefetch_v"] == expected for config in selected) + if case.expected_query_splits is not None: + expected = ( + case.expected_query_splits + if not planned and get_gfx_runtime() == "gfx950" + else 1 + ) + assert selected and all( + config["query_splits"] == expected for config in selected ) - == expected - ) - - -@pytest.mark.parametrize("split_kv_blocks,expected", [(2, 4), (16, 8)]) -def test_recommended_splits_keeps_large_batch_policy( - monkeypatch, split_kv_blocks, expected -): - if get_recommended_splits is None: - pytest.skip("FlyDSL is not available") - monkeypatch.setattr( - torch.cuda, - "get_device_properties", - lambda *_args, **_kwargs: type("Props", (), {"multi_processor_count": 256})(), - ) - assert ( - get_recommended_splits(200, 1, split_kv_blocks, max_context_length=200000) - == expected - ) - -@pytest.mark.parametrize("num_partitions", [64, 74, 86, 103, 128, 171, 256]) -@pytest.mark.parametrize("trans_v", [False, True]) -@pytest.mark.parametrize("block_size", [16, 128]) -def test_mtp4_large_partition_mixed_contexts(num_partitions, trans_v, block_size): - """Long, sparse MTP contexts and empty partitions share the high-NP reducer.""" - _require_gpu() - _run_mtp4_fused_reference_case( - [0, 1, 3, 4, 255, 256, 257, 200003], - 1, - block_size, - trans_v, - num_partitions, - ) + _assert_contracts(args, options) + if plan is not None: + # Exercise planner integer limits without allocating an INT32_MAX 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) -@pytest.mark.parametrize("wide_kv_addressing", [False, True]) -@pytest.mark.parametrize("num_partitions", [1, 8]) -@pytest.mark.parametrize("trans_v", [False, True]) -@pytest.mark.parametrize("block_size", [16, 128]) -@pytest.mark.parametrize( - "lengths,num_kv_heads", - [ - pytest.param([4], 1, id="c4-hkv1"), - pytest.param([257], 1, id="c257-hkv1"), - pytest.param([259], 1, id="c259-hkv1"), - pytest.param([2305], 1, id="c2305-hkv1"), - pytest.param([2307], 1, id="c2307-hkv1"), - pytest.param( - [0, 1, 2, 3, 4, 255, 256, 257, 258, 259, 2303, 2304, 2305, 2306, 2307], - 1, - id="mixed-hkv1", - ), - pytest.param([257], 2, id="c257-hkv2"), - pytest.param([259], 2, id="c259-hkv2"), - ], -) -def test_mtp4_fused_phase_b_accuracy( - monkeypatch, - 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, - num_partitions, - wide_kv_addressing, + max_partitions=None, + per_token=False, + query_length=1, + num_partitions=None, ): - """Cover fused MTP boundaries under both narrow and wide V-load schedules.""" - _require_gpu() - pa_decode_module = importlib.import_module("aiter.ops.flydsl.pa_decode") - compile_tile = pa_decode_module.compile_pa_decode_tile - - def compile_fused_mtp4(**kwargs): - kwargs["query_splits"] = 1 - # Exercise wide-address scheduling with sparse/empty/tail cases without - # allocating a multi-GiB cache for every boundary in this matrix. - kwargs["wide_kv_addressing"] = wide_kv_addressing - return compile_tile(**kwargs) - - monkeypatch.setattr(pa_decode_module, "compile_pa_decode_tile", compile_fused_mtp4) - _run_mtp4_fused_reference_case( - lengths, num_kv_heads, block_size, trans_v, num_partitions + """CLI-only timing wrapper around the same input/reference/launch helpers.""" + 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): @@ -2139,26 +837,13 @@ def _positive_int(value): def _parse_args(argv=None): parser = argparse.ArgumentParser( - formatter_class=argparse.RawTextHelpFormatter, - description="FlyDSL pa_decode correctness + perf sweep", + description="FlyDSL pa_decode correctness + perf sweep" ) parser.add_argument( - "-d", - "--dtype", - type=dtypes.str2Dtype, - nargs="*", - default=[dtypes.bf16], - help="""Query/output data type. - e.g.: -d bf16 fp16""", + "-d", "--dtype", type=dtypes.str2Dtype, nargs="*", default=[dtypes.bf16] ) parser.add_argument( - "-b", - "--batch", - type=_positive_int, - nargs="*", - default=DEFAULT_BATCH_SIZES, - help="""Batch sizes. - e.g.: -b 1 3 16""", + "-b", "--batch", type=_positive_int, nargs="*", default=DEFAULT_BATCH_SIZES ) parser.add_argument( "-q", @@ -2166,8 +851,7 @@ def _parse_args(argv=None): type=_positive_int, nargs="+", default=[1], - help="""Query tokens per sequence: 1 for decode, >1 for MTP. - e.g.: -q 1 2 3 4. MTP positions use dense causal masking.""", + help="Query tokens per sequence; context lengths include them.", ) parser.add_argument( "-s", @@ -2175,9 +859,7 @@ def _parse_args(argv=None): type=dtypes.str2tuple, nargs="*", default=DEFAULT_SHAPES, - help="""(num_query_heads,num_kv_heads,head_dim,context_length). - Contexts are equal-length and include the MTP query tokens. - e.g.: -s 8,1,128,257""", + help="num_query_heads,num_kv_heads,head_dim,context_length", ) parser.add_argument( "--block-size", @@ -2185,7 +867,6 @@ def _parse_args(argv=None): nargs="*", choices=[16, 64, 128], default=[16, 64, 128], - help="""KV-cache block sizes.""", ) parser.add_argument( "--trans-v", @@ -2193,32 +874,28 @@ def _parse_args(argv=None): nargs="*", choices=[0, 1], default=[0, 1], - help="""V-cache layouts: 0 is the plain 4-D cache, 1 the transposed 5-D - cache production serves.""", + 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 passed to get_recommended_splits (4..256). - By default, use context length and GPU occupancy, up to 256 partitions. - Set 8 to retain the legacy clamp. Only used without --num-partitions.""", + 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 partition counts (1..256), bypassing auto selection. - e.g.: --num-partitions 3 5. Overrides --max-partitions.""", - ) - parser.add_argument( - "--per-token", - type=int, - nargs="*", - choices=[0, 1], - default=[0], - help="""KV scale layout: 0 for per-tensor, 1 for per-token.""", + help="Exact split counts (1..256), overriding --max-partitions.", ) args = parser.parse_args(argv) if ( @@ -2226,17 +903,10 @@ def _parse_args(argv=None): and not 4 <= args.max_partitions <= MAX_CONTEXT_PARTITIONS ): parser.error(f"--max-partitions must be in [4, {MAX_CONTEXT_PARTITIONS}]") - if any( - count is not None and count > MAX_CONTEXT_PARTITIONS - for count in args.num_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(value < 1 for value in shape) - ): + 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") @@ -2244,30 +914,17 @@ def _parse_args(argv=None): def main(): - # Parse first so --help and invalid options do not initialize the GPU. + # Parse before GPU checks so --help and invalid options remain usable. 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: - aiter.logger.warning("pa_decode unsupported on %s; skipping", get_gfx_runtime()) - return - if pa_decode is None: - aiter.logger.warning("flydsl is unavailable; skipping pa_decode") + 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_size, - shape, - block_size, - trans_v, - per_token, - query_length, - num_partitions, - ) in itertools.product( + for dtype, batch, shape, page, trans_v, per_token, ql, parts in itertools.product( args.dtype, args.batch, args.shapes, @@ -2277,26 +934,26 @@ def main(): args.query_length, args.num_partitions, ): - num_query_heads, num_kv_heads, head_dim, context_length = shape + heads, kv_heads, dim, context = shape rows.append( run_pa_decode_tile_case( - batch_size, - num_query_heads, - num_kv_heads, - head_dim, - context_length, - block_size, + batch, + heads, + kv_heads, + dim, + context, + page, dtype, bool(trans_v), args.max_partitions, bool(per_token), - query_length=query_length, - num_partitions=num_partitions, + query_length=ql, + num_partitions=parts, ) ) - - df = pd.DataFrame(rows) - aiter.logger.info("pa_decode summary (markdown):\n%s", df.to_markdown(index=False)) + aiter.logger.info( + "pa_decode summary (markdown):\n%s", pd.DataFrame(rows).to_markdown(index=False) + ) if __name__ == "__main__": From 380b9493c6e17f7ab07fc712b196d4bae284738a Mon Sep 17 00:00:00 2001 From: fsx950223 Date: Thu, 17 Sep 2026 06:56:11 +0000 Subject: [PATCH 53/66] refactor(flydsl): consolidate PA decode constants and conditions --- aiter/ops/flydsl/kernels/pa_decode_kernel.py | 149 ++++++++----------- 1 file changed, 63 insertions(+), 86 deletions(-) diff --git a/aiter/ops/flydsl/kernels/pa_decode_kernel.py b/aiter/ops/flydsl/kernels/pa_decode_kernel.py index f86fc719c73..119568e829b 100644 --- a/aiter/ops/flydsl/kernels/pa_decode_kernel.py +++ b/aiter/ops/flydsl/kernels/pa_decode_kernel.py @@ -132,20 +132,6 @@ def compile_pa_decode_tile( ), f"pa_decode_tile only supports query_dtype in ('f16', 'bf16'), got {query_dtype}" Q_DTYPE = fx.BFloat16 if query_dtype == "bf16" else fx.Float16 - # Fix the tuned scalar-scale shape to the direct-fp8 fast path. This - # removes the numerical-policy switch, not the fp8 range/precision tradeoff - # documented above. Per-token and other shapes keep normalized conversion. - SCALAR_FP8_DECODE = ( - is_gfx950 - and head_dim == 128 - and block_size in (16, 128) - and trans_v - and not per_token_kv - and query_dtype == "bf16" - and query_length == 1 - and query_group_size in (8, 16) - ) - assert ( head_dim % 64 == 0 ), f"pa_decode_tile only supports head_dim that's a multiple of 64, got {head_dim}" @@ -158,15 +144,27 @@ def compile_pa_decode_tile( CTA_ROWS = QUERIES_PER_CTA * query_group_size M_TILES = (CTA_ROWS + MFMA_MNK - 1) // MFMA_MNK ROWS_PADDED = M_TILES * MFMA_MNK + + # Share the device/shape gate; scalar V scheduling also supports f16 and + # MTP, while the numerical and MFMA specializations below require bf16. + TUNED_SHAPE = is_gfx950 and head_dim == 128 and block_size in (16, 128) + TUNED_BF16 = TUNED_SHAPE and query_dtype == "bf16" + # Fix the tuned scalar-scale shape to the direct-fp8 fast path. This + # removes the numerical-policy switch, not the fp8 range/precision tradeoff + # documented above. Per-token and other shapes keep normalized conversion. + SCALAR_FP8_DECODE = ( + TUNED_BF16 + and trans_v + and not per_token_kv + and query_length == 1 + and query_group_size in (8, 16) + ) # CDNA4 contracts four legacy fp8 K=32 groups in one K=128 instruction. # Concatenating the existing packs preserves cache/LDS layouts for fused # MTP3/MTP4 and query-split MTP4 CTAs. WIDE_FP8_MFMA = ( - is_gfx950 + TUNED_BF16 and per_token_kv - and query_dtype == "bf16" - and head_dim == 128 - and block_size in (16, 128) and query_length in (3, 4) and query_group_size == 16 ) @@ -178,34 +176,26 @@ def compile_pa_decode_tile( # P publication. Page-16 gathers and narrow addresses retain prefetching. MTP4_PREFETCH_V = MTP4_FUSED and (block_size == 16 or not wide_kv_addressing) PER_TOKEN_M1 = ( - per_token_kv - and is_gfx950 - and head_dim == 128 - and block_size in (16, 128) - and query_dtype == "bf16" + TUNED_BF16 + and per_token_kv and QUERIES_PER_CTA == 1 and query_group_size in (8, 16) - and M_TILES == 1 and prefetch_v ) # The wrapper selects earlier V loads for tuned gfx950 decode shapes. # Page 16 additionally delays next K until PV and uses IGLP to overlap # groups of V loads with QK MFMA instructions. - TUNED_V_PATH = ( - is_gfx950 - and head_dim == 128 - and block_size in (16, 128) - and (PER_TOKEN_M1 or (trans_v and not per_token_kv)) - ) + TUNED_V_PATH = PER_TOKEN_M1 or (TUNED_SHAPE and trans_v and not per_token_kv) TUNE_PAGE128 = TUNED_V_PATH and block_size == 128 EARLY_V = prefetch_v and TUNED_V_PATH and M_TILES == 1 PAGE16_VPIPE = EARLY_V and block_size == 16 - SCALES_BEFORE_CURRENT_V = WIDE_FP8_MFMA and ( - PAGE16_VPIPE or (PER_TOKEN_M1 and not trans_v) + REUSE_KV_PAGES = PAGE16_VPIPE or PER_TOKEN_M1 + SCALES_BEFORE_CURRENT_V = ( + WIDE_FP8_MFMA and PER_TOKEN_M1 and (block_size == 16 or not trans_v) ) # Apply score scales before the -inf mask on the tuned single-M-tile # paths. This removes the second mask without introducing -inf * 0. - M1_SCALE_BEFORE_MASK = PAGE16_VPIPE or SCALAR_FP8_DECODE or PER_TOKEN_M1 + 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 layout: V=A, P=B -> output [head-dim (row), query-row (col=lane16)], # generalized over head_dim via the VHE_CHUNKS loop. @@ -216,8 +206,10 @@ def compile_pa_decode_tile( 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" @@ -244,8 +236,17 @@ def compile_pa_decode_tile( f"head_dim {head_dim} is unsupported: head_dim//{NQCHUNK} ({QCHUNK}) must " f"be <= 8 or a multiple of 8" ) + # Each per-lane Q chunk uses at most 128 bits per buffer load. + 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 + # PV processes one VHE_SIZE-wide slice of the output head dimension. + VHE_SIZE = head_dim // VHE_CHUNKS + OP_ELEMS = MFMA_ACC_ELEMS # PV C-fragment elements/lane/chunk + NVOPS = TILE_TOK // MFMA_K # 8 PV k_steps (256 tokens / K=32) + 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) @@ -261,6 +262,11 @@ def compile_pa_decode_tile( 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 layout (shared across the 4 warps) -- # sQ: fp8[ROWS_PADDED,head_dim] staged+quantized query. sP: fp8[16,TILE_TOK] # quantized probs. sQscale: f32[ROWS_PADDED] except when Q is converted @@ -288,11 +294,10 @@ def compile_pa_decode_tile( sVPage_bytes = NWARP * PAGES_PER_CHUNK * 4 # i32 # Per-token KV-scale staging (per_token only), double-buffered so the tt+1 # prefetch (into the other buffer) doesn't clobber the current tile's scales. - KV_BUFS = 2 if per_token_kv else 1 KV_BUF_STRIDE = 2 * NWARP * TOK_PER_WARP * f32 # k-region + v-region, one buffer sKScale_off = sVPage_off + sVPage_bytes sVScale_off = sKScale_off + NWARP * TOK_PER_WARP * f32 - sKVScale_bytes = KV_BUFS * KV_BUF_STRIDE if per_token_kv else 0 + sKVScale_bytes = 2 * 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 @@ -376,7 +381,6 @@ def _load(elem_idx): # A cache above 2 GiB overflows an i32 element offset, so the flat view # and the offsets below widen together. - KV_EXTENT = (1 << 42) if wide_kv_addressing else (1 << 30) _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) @@ -389,8 +393,6 @@ def _kv_addr(phys, page_elems, rest): # Per-lane Q chunk (QCHUNK 16-bit elems) fetched in QLOAD_UNIT-wide # pieces (128b max per buffer load): head_dim=256 needs 2 pieces. - QLOAD_UNIT = min(8, QCHUNK) - N_QLOADS = QCHUNK // QLOAD_UNIT _q_copy_op = ( fx.rocdl.BufferCopy128b() if QLOAD_UNIT == 8 else fx.rocdl.BufferCopy64b() ) @@ -403,9 +405,8 @@ 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)) - c16 = 16 - rgroup = lane // c16 # 0..3: which quarter-wave (paired with warp -> query row) - lane16 = lane - rgroup * c16 # 0..15: this row's head-dim chunk index + 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): @@ -520,8 +521,6 @@ def _lds_load(byte_off, elem_ty, n): def _lds_store(byte_off, elem_ty, vec): fx.ptr_store(vec, _lds_ptr(byte_off, elem_ty)) - NCHUNK = TOK_PER_WARP // MFMA_MNK # 4 - if const_expr(per_token_kv): scale_load_width = NCHUNK if block_size >= 64 else 1 scale_copy_op = ( @@ -636,7 +635,7 @@ def _stage_kv_scale_to_lds(phys_vec, buf_off=0): 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 * c16 + lane16) * f32 + slot = (warp * TOK_PER_WARP + rgroup * MFMA_MNK + lane16) * f32 _lds_store( sKScale_off + buf_off + slot, fx.Float32, @@ -650,7 +649,7 @@ def _stage_kv_scale_to_lds(phys_vec, buf_off=0): def _load_scale_vec(base_off, a, buf_off=0): # This lane's 4 per-token scales for chunk `a` from an LDS scale region. - slot = (warp * TOK_PER_WARP + a * c16 + rgroup * 4) * f32 + 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): @@ -692,10 +691,10 @@ def _mfma_fp8(a_ops, b_ops, a_base, b_base, k_steps, acc): return acc # -- raw dwordx4 K load (A operand) -- - # token = warp*TOK_PER_WARP + a*c16 + lane16 (the softmax mask and P-pack + # token = warp*TOK_PER_WARP + a*MFMA_MNK + lane16 (softmax mask and P-pack # write position below must encode this same formula). def _k_ops(phys, a): - within_page_tok = (warp * TOK_PER_WARP + a * c16 + lane16) % block_size + 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 @@ -715,7 +714,7 @@ def _k_ops(phys, a): def _k_ops_from_phys(phys_vec): flat = [] for a in range_constexpr(NCHUNK): - phys = fx.Int32(phys_vec[(a * c16) // block_size]) + 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) @@ -741,7 +740,7 @@ def _k_ops_flat(tt_i32): k_pf0, phys_vec0 = _k_ops_flat(part_start) # Issue the V page-index prefetch alongside K; the LDS write is # visible after the barrier below. - if const_expr(PAGE16_VPIPE or PER_TOKEN_M1): + if const_expr(REUSE_KV_PAGES): # The K page vector is already resident. Reuse it for the V # page matrix instead of issuing a second block-table load. _stage_v_page_row(phys_vec0) @@ -910,7 +909,9 @@ def _quant_q_row(m, q_row_off, q_units): # QK in NCHUNK chunks of 4 tokens: each chunk yields a f32x4 # C-fragment, so softmax processes 4 scores at a time (low VGPR peak). _ct = [ - fx.Vector.from_elements([float(a * c16 + r) for r in range_constexpr(4)]) + fx.Vector.from_elements( + [float(a * MFMA_MNK + r) for r in range_constexpr(4)] + ) for a in range_constexpr(NCHUNK) ] @@ -920,21 +921,9 @@ def _score_mask(a, upper, lower): valid = valid & (_ct[a] >= lower) return valid - # P.V is loop-tiled over head-dim (like production's VHELOOP): each - # step computes O[:, vh*VHE_SIZE:+VHE_SIZE] instead of materializing - # the full [16, head_dim] at once. - VHE_SIZE = head_dim // VHE_CHUNKS - OP_ELEMS = ( - MFMA_MNK * VHE_SIZE // (NWARP * WAVE) - ) # PV C-fragment elements/lane/chunk (probed = 4) - # -- raw dwordx4 V load (B operand) -- # One dwordx4 load per (16-token sub-block, head_elem); trans_v only # changes the offset formula. `sub`/`step` walk pages/16-token sub-blocks. - NVOPS = TILE_TOK // MFMA_K # 8 PV k_steps (256 tokens / K=32) - STEPS_PER_PAGE = block_size // 16 - STEPS_PER_CHUNK = min(block_size, TOK_PER_WARP) // 16 - def _v_ops(phys_row, vh): head_group = ((vh * VHE_SIZE) // 16) + warp head_element = head_group * 16 + lane16 @@ -985,12 +974,6 @@ def _v_ops(phys_row, vh): for m in range_constexpr(M_TILES) ] - K_SLOT, V_SLOT = 0, 1 - # Per-M-tile loop-carried state after the K/V slots: one output chunk - # per VHE_CHUNKS (2 for head_dim=128, 4 for head_dim=256, 1 for 64), - # then running-max + running-denom. - STATE_PER_M = VHE_CHUNKS + 2 - def _o_slot(m, vh): return 2 + STATE_PER_M * m + vh @@ -1006,7 +989,6 @@ def _l_slot(m): init_state.extend([o_zero] * VHE_CHUNKS + [NEG_INF, ZERO_F]) # Fused pipeline: carry next V in registers, reusing each chunk only # after its final query has consumed the current tile's values. - V_DATA_SLOT = 2 + STATE_PER_M * M_TILES if const_expr(MTP4_PREFETCH_V): v_pf0 = fx.Vector.filled(VHE_CHUNKS * NVOPS, 0, fx.Int64) if part_start < part_end: @@ -1107,9 +1089,7 @@ def _mask_v_scale( ) def _lmax_off_m(m): - return sLmax_off + ( - m * MFMA_MNK * NWARP_PAD * f32 if const_expr(M_TILES > 1) else 0 - ) + return sLmax_off + m * MFMA_MNK * NWARP_PAD * f32 # Phase-split (M_TILES>1): pass-1 (QK+mask+max) loops all M-tiles into # per-M-tile LDS slices so they share ONE barrier. M_TILES==1: else below. @@ -1346,7 +1326,7 @@ def _lmax_off_m(m): + lane16 * SP_ROW_BYTES + warp * TOK_PER_WARP + rgroup * 4 - + a * (c16 // 4) * f32 + + a * (MFMA_MNK // 4) * f32 ) _lds_store( p_off, @@ -1499,12 +1479,12 @@ def _lmax_off_m(m): + warp * TOK_PER_WARP + rgroup * 4 ) - # The NCHUNK P words scatter across the row at stride c16//4 + # The NCHUNK P words scatter at stride MFMA_MNK//4 # i32 (the token->fp8-lane interleave the PV ds_read_b128 # expects); one strided store per chunk. for a in range_constexpr(NCHUNK): _lds_store( - p_off0 + a * (c16 // 4) * f32, + p_off0 + a * (MFMA_MNK // 4) * f32, fx.Int32, fx.Vector.from_elements([words[a]], dtype=fx.Int32), ) @@ -1572,18 +1552,15 @@ def _lmax_off_m(m): # Page 16 defers K until PV; page 128 overlaps K with softmax. k_next = k_cur if tt1 < part_end: - if const_expr(PAGE16_VPIPE): - # Publish next page ids/scales here. K itself is issued - # after the P barrier to overlap the current tile's PV. + if const_expr(REUSE_KV_PAGES): + # Publish next pages/scales once. Page 16 defers K + # until the P barrier; per-token page 128 issues it + # after scales so their VMEM wait cannot drain K. 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)) - elif const_expr(PER_TOKEN_M1): - # Finish the scale VMEM-to-LDS dependency before - # issuing K, so its wait cannot drain the K prefetch. - phys_vec1 = _v_page_fetch_and_stage(tt1) - _stage_kv_scale_to_lds(phys_vec1, _kv_buf_off(tt1)) - k_next = _k_ops_from_phys(phys_vec1) + 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) @@ -1787,10 +1764,10 @@ def _lmax_off_m(m): p_off0 = ( sP_off + lane16 * SP_ROW_BYTES + warp * TOK_PER_WARP + rgroup * 4 ) - # NCHUNK P words scatter at stride c16//4 i32 (see phase-split). + # NCHUNK P words scatter at stride MFMA_MNK//4 i32 (see phase-split). for a in range_constexpr(NCHUNK): _lds_store( - p_off0 + a * (c16 // 4) * f32, + p_off0 + a * (MFMA_MNK // 4) * f32, fx.Int32, fx.Vector.from_elements([words[a]], dtype=fx.Int32), ) From 21008937f39fed0107e9d5880d09e001ec7ff32c Mon Sep 17 00:00:00 2001 From: fsx950223 Date: Thu, 17 Sep 2026 07:37:29 +0000 Subject: [PATCH 54/66] refactor(flydsl): centralize PA decode scheduling Select query splitting and V prefetching in the compiler entry point while caching only the effective specialization. Clarify FP8 operand-pack versus MFMA instruction K and verify scheduling boundaries and cache reuse in the generic decode test. --- aiter/ops/flydsl/kernels/pa_decode_kernel.py | 184 +++++++++++++++---- aiter/ops/flydsl/pa_decode.py | 81 +------- op_tests/test_flydsl_pa_decode.py | 102 +++++++++- 3 files changed, 251 insertions(+), 116 deletions(-) diff --git a/aiter/ops/flydsl/kernels/pa_decode_kernel.py b/aiter/ops/flydsl/kernels/pa_decode_kernel.py index 119568e829b..809b274055c 100644 --- a/aiter/ops/flydsl/kernels/pa_decode_kernel.py +++ b/aiter/ops/flydsl/kernels/pa_decode_kernel.py @@ -3,12 +3,13 @@ """Readable tile-programming reference for paged-attention fp8 decode. -K/V are fp8 e4m3 (FNUZ on gfx942, OCP on gfx950) fed straight into -``mfma_f32_16x16x32_fp8_fp8``; Q (bf16/f16) and the softmax probabilities P are -quantized to fp8 too. Scales fold out of the matmuls (q/key scale into the QK -score, value scale + 1/FP8_MAX into the epilogue); softmax max/sum stay f32. -The gfx950 BF16 per-token MTP4 specialization uses K=128 FP8 MFMA atoms, -with the same normalized Q/P values and operand layouts. +K/V are fp8 e4m3 (FNUZ on gfx942, OCP on gfx950) fed straight into FP8 MFMA; +Q (bf16/f16) and the softmax probabilities P are quantized to fp8 too. Scales +fold out of the matmuls (q/key scale into the QK score, value scale + 1/FP8_MAX +into the epilogue); softmax max/sum stay f32. +The default instruction is ``mfma_f32_16x16x32_fp8_fp8``. Tuned gfx950 BF16 +per-token MTP3/MTP4 shapes use K=128 FP8 MFMA atoms, with the same normalized +Q/P values and operand layouts, including query-split MTP4 CTAs. The tuned gfx950 BF16 scalar-scale decode specialization instead casts Q/P directly to fp8, without range normalization or the compensating 1/FP8_MAX. This is a speed/precision tradeoff: small Q values and long probability tails @@ -51,16 +52,119 @@ MFMA_MNK = ( 16 # M = N = 16 for the MMA atom; also query rows handled per CTA (padded to 16) ) -MFMA_K = 32 # fp8 MFMA contracts K = 32 per instruction (mfma_f32_16x16x32_fp8_fp8) +# Logical K coverage across a wave of one i64 FP8 operand pack per lane. +# Operand layouts retain this unit for both K32 and K128 instructions. +FP8_PACK_K = 32 WAVE = 64 -# f32 C-fragment elements each lane holds for one 16x16x32 atom (16*16 / 64). +# f32 C-fragment elements per lane for a 16x16 atom, independent of its K. MFMA_ACC_ELEMS = MFMA_MNK * MFMA_MNK // WAVE LOG2E = 1.4426950408889634 KV_COMPUTE_BLOCK = 256 -@functools.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, + sliding_window: int = 0, + use_sinks: bool = False, + sink_dtype_str: str = "f32", +): + """Select the schedule and return a cached kernel plus launch wrapper. + + Batch/head/CU counts only select query splitting and V prefetching; the + builder caches the resulting specialization, not those runtime sizes. + ``query_splits=None`` selects automatically; an explicit count overrides + splitting while retaining the matching prefetch policy. + """ + gfx950_d128 = "gfx95" in get_rocm_arch() and head_dim == 128 + tuned_per_token_shape = ( + gfx950_d128 + and per_token_kv + and query_dtype == "bf16" + and block_size in (16, 128) + ) + # Keep the dense-grid estimate for planned execution too. The plan's + # capacity and active token windows do not change the scheduling policy. + dense_workgroups = num_seqs * num_kv_heads * num_partitions + assert query_length >= 1, f"query_length must be >= 1, got {query_length}" + if query_splits is None: + # Small static MTP2/MTP4 grids use one query per CTA. MTP3 stays fused. + query_splits = ( + query_length + if tuned_per_token_shape + and not use_work_plan + and num_kv_heads == 1 + and query_length in (2, 4) + and query_group_size == 16 + and query_length * dense_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" + + # Split queries always prefetch. Unsplit decode supports Hkv1 and static + # Hkv2; transposed page128 keeps its lower-register schedule above one + # workgroup per CU. One query with GQA8/16 already implies one M-tile. + per_token_m1 = ( + tuned_per_token_shape + and query_length == query_splits + and query_group_size in (8, 16) + and ( + query_splits > 1 + or ( + (num_kv_heads == 1 or (not use_work_plan and num_kv_heads == 2)) + and ( + block_size == 16 + or not trans_v + or dense_workgroups <= num_compute_units + ) + ) + ) + ) + scalar_prefetch = ( + gfx950_d128 + and block_size == 128 + and trans_v + and not per_token_kv + and query_length * query_group_size <= MFMA_MNK + and num_compute_units < dense_workgroups <= 2 * num_compute_units + ) + return _compile_pa_decode_tile( + head_dim=head_dim, + query_group_size=query_group_size, + block_size=block_size, + 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, + prefetch_v=per_token_m1 or scalar_prefetch, + query_splits=query_splits, + use_work_plan=use_work_plan, + sliding_window=sliding_window, + use_sinks=use_sinks, + sink_dtype_str=sink_dtype_str, + ) + + +@functools.cache +def _compile_pa_decode_tile( *, head_dim: int, query_group_size: int, @@ -135,9 +239,6 @@ def compile_pa_decode_tile( assert ( head_dim % 64 == 0 ), f"pa_decode_tile only supports head_dim that's a multiple of 64, got {head_dim}" - assert query_length >= 1, f"query_length must be >= 1, got {query_length}" - 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 # Flattened query-row axis (MTP outer, GQA head inner), tiled into 16-row M-tiles. TOTAL_ROWS = query_length * query_group_size @@ -168,6 +269,8 @@ def compile_pa_decode_tile( 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 # Small grids split query positions into separate CTAs; otherwise retain # all four queries here and share KV loads, scales, and P publication. MTP4_FUSED = WIDE_FP8_MFMA and M_TILES == 4 @@ -175,21 +278,16 @@ def compile_pa_decode_tile( # register pressure: load plain V during P packing and transposed V after # P publication. Page-16 gathers and narrow addresses retain prefetching. MTP4_PREFETCH_V = MTP4_FUSED and (block_size == 16 or not wide_kv_addressing) - PER_TOKEN_M1 = ( - TUNED_BF16 - and per_token_kv - and QUERIES_PER_CTA == 1 - and query_group_size in (8, 16) - and prefetch_v - ) - # The wrapper selects earlier V loads for tuned gfx950 decode shapes. + # The selector only enables prefetching for supported single-M-tile shapes. + PER_TOKEN_M1 = per_token_kv and prefetch_v # Page 16 additionally delays next K until PV and uses IGLP to overlap # groups of V loads with QK MFMA instructions. - TUNED_V_PATH = PER_TOKEN_M1 or (TUNED_SHAPE and trans_v and not per_token_kv) - TUNE_PAGE128 = TUNED_V_PATH and block_size == 128 - EARLY_V = prefetch_v and TUNED_V_PATH and M_TILES == 1 - PAGE16_VPIPE = EARLY_V and block_size == 16 - REUSE_KV_PAGES = PAGE16_VPIPE or PER_TOKEN_M1 + TUNE_PAGE128 = block_size == 128 and ( + PER_TOKEN_M1 or (TUNED_SHAPE and trans_v and not per_token_kv) + ) + PAGE16_VPIPE = prefetch_v and block_size == 16 + # The selected page16 pipeline is always per-token as well. + 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) ) @@ -222,7 +320,10 @@ def compile_pa_decode_tile( assert ( QKHE_LOOP >= 1 ), f"head_dim {head_dim} must be at least {RGROUP_QUARTERS * QK_CHUNK_ELEMS}" - N_SUBCHUNKS = QKHE_LOOP * (QK_CHUNK_ELEMS // 8) + # 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" # Q-quant chunk width: NQCHUNK stays fixed at 16 (tied to `lane16`'s role # as the absmax butterfly width); QCHUNK scales with head_dim instead. @@ -244,7 +345,8 @@ def compile_pa_decode_tile( # PV processes one VHE_SIZE-wide slice of the output head dimension. VHE_SIZE = head_dim // VHE_CHUNKS OP_ELEMS = MFMA_ACC_ELEMS # PV C-fragment elements/lane/chunk - NVOPS = TILE_TOK // MFMA_K # 8 PV k_steps (256 tokens / K=32) + # 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 @@ -657,15 +759,23 @@ def _load_kv_scale_vecs(a, buf_off=0): sVScale_off, a, buf_off ) - def _mfma_fp8(a_ops, b_ops, a_base, b_base, k_steps, acc): + def _mfma_fp8(a_ops, b_ops, a_base, b_base, k_packs, acc): + # Counts and offsets use i64 operand packs per lane. Each instruction + # consumes one pack for K32 or PACKS_PER_MFMA packs for K128. if const_expr(WIDE_FP8_MFMA): - for ks in range_constexpr(k_steps // 4): + for inst in range_constexpr(k_packs // PACKS_PER_MFMA): a_pack = fx.Vector.from_elements( - [a_ops[a_base + ks * 4 + i] for i in range_constexpr(4)], + [ + 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 + ks * 4 + i] for i in range_constexpr(4)], + [ + 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( @@ -683,10 +793,10 @@ def _mfma_fp8(a_ops, b_ops, a_base, b_base, k_steps, acc): ], ) else: - for ks in range_constexpr(k_steps): + for pack in range_constexpr(k_packs): acc = fx.rocdl.mfma_f32_16x16x32_fp8_fp8( T.f32x4, - [a_ops[a_base + ks], b_ops[b_base + ks], acc, 0, 0, 0], + [a_ops[a_base + pack], b_ops[b_base + pack], acc, 0, 0, 0], ) return acc @@ -704,7 +814,7 @@ def _k_ops(phys, a): ((kv_h * QCHUNK + he_idx) * block_size + within_page_tok) * QK_CHUNK_ELEMS, ) - w = _k_load16(base) # head[he_idx*16 : +16] -> k_step 2*qkhe, 2*qkhe+1 + w = _k_load16(base) # head[he_idx*16 : +16] -> two K32 operand packs if const_expr(block_size == 16): # help the scheduler overlap the PAGES_PER_CHUNK gathered loads fx.rocdl.sched_barrier(fx.rocdl.mask_vmem_rd) @@ -1074,7 +1184,7 @@ def _mask_v_scale( ] elif const_expr( (M_TILES > 1 and not MTP4_FUSED) - or (EARLY_V and not SCALES_BEFORE_CURRENT_V) + 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) @@ -1540,7 +1650,7 @@ def _lmax_off_m(m): 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 - # QK: each NCHUNK chunk accumulates N_SUBCHUNKS k_steps into an f32x4. + # QK: each NCHUNK chunk accumulates N_SUBCHUNKS packs into an f32x4. frag_Ss = [] for a in range_constexpr(NCHUNK): acc = fx.Vector.filled(MFMA_ACC_ELEMS, 0.0, fx.Float32) @@ -1798,7 +1908,7 @@ def _lmax_off_m(m): ).broadcast_to(OP_ELEMS) # Use the early V loads where enabled; other shapes retain # the existing batched loads before PV. - if const_expr(EARLY_V): + if const_expr(prefetch_v): v_vh_batch = v_vh_shared else: v_vh_batch = [ diff --git a/aiter/ops/flydsl/pa_decode.py b/aiter/ops/flydsl/pa_decode.py index c525b220e6b..f8bee740f16 100644 --- a/aiter/ops/flydsl/pa_decode.py +++ b/aiter/ops/flydsl/pa_decode.py @@ -576,80 +576,6 @@ def normalize_scale(scale, name): # if either does. wide_kv_addressing = max(key_cache.numel(), value_cache.numel()) >= 2**31 - dense_partition_workgroups = num_seqs * num_kv_heads * num_partitions - gfx950_d128 = arch == "gfx950" and head_dim == 128 - scalar_prefetch_candidate = ( - gfx950_d128 - and block_size == 128 - and trans_v - and not per_token_kv - and query_length * query_group_size <= 16 - ) - tuned_per_token_shape = ( - gfx950_d128 - and per_token_kv - and query_dtype == "bf16" - and block_size in (16, 128) - ) - # Hkv2 static decode is covered by the same workgroup/layout policy. - # Planned Hkv2 and larger KV-head counts retain their existing schedule. - tuned_m1_heads = num_kv_heads == 1 or (work_plan is None and num_kv_heads == 2) - per_token_single_query = ( - tuned_per_token_shape - and tuned_m1_heads - and query_length == 1 - and query_group_size in (8, 16) - ) - query_split_candidate = ( - work_plan is None - and tuned_per_token_shape - and num_kv_heads == 1 - and query_length in (2, 4) - and query_group_size == 16 - ) - needs_num_cus = ( - scalar_prefetch_candidate - or query_split_candidate - or (per_token_single_query and block_size == 128 and trans_v) - ) - num_cus = ( - torch.cuda.get_device_properties(dev).multi_processor_count - if needs_num_cus - else 0 - ) - - # Split small MTP2/MTP4 grids into one query per CTA. MTP3 stays fused - # because the supported split counts (1, 2, 4) must divide query_length. - query_splits = ( - query_length - if query_split_candidate - and query_length * dense_partition_workgroups <= 2 * num_cus - else 1 - ) - queries_per_cta = query_length // query_splits - total_workgroups = dense_partition_workgroups * query_splits - per_token_m1 = ( - tuned_per_token_shape - and queries_per_cta == 1 - and query_group_size in (8, 16) - and (query_splits > 1 or per_token_single_query) - ) - - # Use one policy for the single-M-tile per-token pipelines. Page16 and - # plain V tolerate high occupancy; unsplit transposed page128 retains its - # lower-register schedule above one workgroup per CU. - per_token_prefetch = per_token_m1 and ( - query_splits > 1 - or block_size == 16 - or not trans_v - or total_workgroups <= num_cus - ) - scalar_prefetch = ( - scalar_prefetch_candidate - and num_cus < dense_partition_workgroups <= 2 * num_cus - ) - prefetch_v = scalar_prefetch or per_token_prefetch - # Partial outputs must exclude the sink: reduction adds it once across all # partitions. Keep direct NP=1 output fused, without extra scratch/launches. use_direct_sinks = sinks is not None and num_partitions == 1 and work_plan is None @@ -658,6 +584,11 @@ def normalize_scale(scale, name): 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, @@ -665,8 +596,6 @@ def normalize_scale(scale, name): query_length=query_length, trans_v=trans_v, wide_kv_addressing=wide_kv_addressing, - prefetch_v=prefetch_v, - query_splits=query_splits, use_work_plan=work_plan is not None, sliding_window=sliding_window, use_sinks=use_direct_sinks, diff --git a/op_tests/test_flydsl_pa_decode.py b/op_tests/test_flydsl_pa_decode.py index 1849b4189d9..77a337efa55 100644 --- a/op_tests/test_flydsl_pa_decode.py +++ b/op_tests/test_flydsl_pa_decode.py @@ -97,6 +97,7 @@ class DecodeCase: wide_kv_addressing: bool | None = None expected_prefetch: bool | str | None = None expected_query_splits: int | None = None + check_schedule: bool = False def _require_gpu(): @@ -474,6 +475,66 @@ def _assert_contracts(args, options): pa_decode(*args, **{**options, "work_plan": wrong_plan}) +def _assert_schedule_cache(compile_tile, kwargs, selected): + """Check deterministic CU boundaries without launching hypothetical grids.""" + workgroups = kwargs["num_seqs"] * kwargs["num_kv_heads"] * kwargs["num_partitions"] + planned, ql = kwargs["use_work_plan"], kwargs["query_length"] + cached = {} + + def check(cus, expected, **changes): + previous_count = len(selected) + compiled = compile_tile(**{**kwargs, "num_compute_units": cus, **changes}) + assert len(selected) == previous_count + 1 + actual = (selected[-1]["query_splits"], selected[-1]["prefetch_v"]) + assert actual == expected + previous = cached.setdefault(expected, compiled) + assert compiled["kernel"] is previous["kernel"] + assert compiled["launch"] is previous["launch"] + + if not kwargs["per_token_kv"]: + # Scalar prefetch requires CU < workgroups <= 2 * CU. + half = (workgroups + 1) // 2 + points = [ + (workgroups, (1, False)), + (workgroups - 1, (1, True)), + (half, (1, True)), + (half - 1, (1, False)), + ] + elif ql == 1: + # Transposed page128 prefetch stops above one workgroup per CU. + points = [ + (workgroups - 1, (1, False)), + (workgroups, (1, not planned)), + (workgroups + 1, (1, not planned)), + ] + else: + # The selected MTP2 case has an exact QL * workgroups == 2 * CU boundary. + boundary = ql * workgroups // 2 + split = (1, False) if planned else (ql, True) + points = [(boundary - 1, (1, False)), (boundary, split), (boundary + 1, split)] + for cus, expected in points: + check(cus, expected) + + if kwargs["per_token_kv"] and ql == 1: + # Hkv1 remains tuned for plans; Hkv2 is static-only; Hkv>2 is untuned. + for heads, prefetch in [(1, True), (2, not planned), (3, False)]: + cus = kwargs["num_seqs"] * heads * kwargs["num_partitions"] + check(cus, (1, prefetch), num_kv_heads=heads) + elif kwargs["per_token_kv"] and not planned: + # Explicit split overrides must also select matching prefetching. + check(boundary - 1, (ql, True), query_splits=ql) + check(boundary + 1, (1, False), query_splits=1) + + # Runtime size/CU metadata must not fragment the private specialization cache. + cus, expected = points[-1] + changes = {"num_seqs": kwargs["num_seqs"] * 2} + multiplier = 2 + if not kwargs["per_token_kv"]: + changes["num_kv_heads"] = kwargs["num_kv_heads"] * 2 + multiplier = 4 + check(cus * multiplier, expected, **changes) + + def _case( name, shape=(4, 1, 16, 128), @@ -508,7 +569,15 @@ def _case( BF16, FP16, FP32 = torch.bfloat16, torch.float16, torch.float32 CASES = [ _case("scalar-direct", (1, 2, 8, 128), (16, 1, 0), 1), - _case("scalar-window-sinks", (1, 2, 8, 128), (128, 1, 0), 7, 257, FP32), + _case( + "scalar-window-sinks", + (1, 2, 8, 128), + (128, 1, 0), + 7, + 257, + FP32, + check_schedule=True, + ), _case( "register-64-65", (1, 2, 4, 256), @@ -533,6 +602,7 @@ def _case( lengths=(257, 259), expected_query_splits=2, expected_prefetch=True, + check_schedule=True, ), _case( "hkv2-prefetch-1wg", @@ -542,6 +612,7 @@ def _case( sink=FP32, lengths=(257,) * 16, expected_prefetch="occupancy", + check_schedule=True, ), _case( "hkv2-prefetch-2wg", @@ -663,10 +734,25 @@ def test_pa_decode(case, planned, monkeypatch): output, query, _, _, context = args[:5] scratch, sinks, plan = args[14:17], args[-1], options["work_plan"] module = importlib.import_module("aiter.ops.flydsl.pa_decode") + kernel_module = importlib.import_module("aiter.ops.flydsl.kernels.pa_decode_kernel") compile_tile = module.compile_pa_decode_tile + build_tile = kernel_module._compile_pa_decode_tile + num_compute_units = torch.cuda.get_device_properties( + query.device + ).multi_processor_count selected = [] + compile_args = None + + def capture_build(**kwargs): + # Observe the final schedule before the real cache lookup, including hits. + selected.append(kwargs) + return build_tile(**kwargs) def compile_checked(**kwargs): + nonlocal compile_args + assert kwargs["num_seqs"] == context.numel() + assert kwargs["num_kv_heads"] == case.num_kv_heads + assert kwargs["num_compute_units"] == num_compute_units if case.query_splits is not None and not planned: kwargs["query_splits"] = case.query_splits if case.wide_kv_addressing is not None: @@ -674,12 +760,18 @@ def compile_checked(**kwargs): assert kwargs["use_sinks"] == ( sinks is not None and args[8] == 1 and not planned ) - selected.append(kwargs) - return compile_tile(**kwargs) + compile_args = kwargs + previous_count = len(selected) + compiled = compile_tile(**kwargs) + assert len(selected) == previous_count + 1 + if case.query_splits is not None and not planned: + assert selected[-1]["query_splits"] == case.query_splits + return compiled def unexpected_reducer(*_args, **_kwargs): raise AssertionError("static NP=1 must not launch a reducer") + monkeypatch.setattr(kernel_module, "_compile_pa_decode_tile", capture_build) monkeypatch.setattr(module, "compile_pa_decode_tile", compile_checked) if not planned and args[8] == 1: monkeypatch.setattr(module, "launch_pa_decode_ps_reduce", unexpected_reducer) @@ -737,6 +829,10 @@ def check(disabled_sink=False): config["query_splits"] == expected for config in selected ) + if case.check_schedule and get_gfx_runtime() == "gfx950": + # Compile-only checks stay outside capture and never launch fake-CU results. + _assert_schedule_cache(compile_tile, compile_args, selected) + _assert_contracts(args, options) if plan is not None: # Exercise planner integer limits without allocating an INT32_MAX KV cache. From a4902d0023e64ed3b59d2bf580cc1febdec664b6 Mon Sep 17 00:00:00 2001 From: fsx950223 Date: Thu, 17 Sep 2026 08:13:32 +0000 Subject: [PATCH 55/66] refactor(flydsl): unify PA decode scheduling and compilation --- aiter/ops/flydsl/kernels/pa_decode_kernel.py | 197 ++++++++----------- op_tests/test_flydsl_pa_decode.py | 16 +- 2 files changed, 94 insertions(+), 119 deletions(-) diff --git a/aiter/ops/flydsl/kernels/pa_decode_kernel.py b/aiter/ops/flydsl/kernels/pa_decode_kernel.py index 809b274055c..7aa536e3240 100644 --- a/aiter/ops/flydsl/kernels/pa_decode_kernel.py +++ b/aiter/ops/flydsl/kernels/pa_decode_kernel.py @@ -60,8 +60,14 @@ MFMA_ACC_ELEMS = MFMA_MNK * MFMA_MNK // WAVE LOG2E = 1.4426950408889634 KV_COMPUTE_BLOCK = 256 +# Cache only the selected specialization, not the batch/head/CU counts used +# to choose it. Entries are published after both JIT wrappers are constructed. +_PA_DECODE_TILE_CACHE = {} +# Repeated calls bypass the builder's closure setup. Keep this metadata cache +# bounded; the specialization cache below shares kernels across runtime sizes. +@functools.lru_cache(maxsize=128) def compile_pa_decode_tile( *, head_dim: int, @@ -83,20 +89,49 @@ def compile_pa_decode_tile( use_sinks: bool = False, sink_dtype_str: str = "f32", ): - """Select the schedule and return a cached kernel plus launch wrapper. + """Select the schedule and cache the PA-decode kernel + launch wrapper. - Batch/head/CU counts only select query splitting and V prefetching; the - builder caches the resulting specialization, not those runtime sizes. + Batch/head/CU counts select query splitting and V prefetching. Calls that + select the same specialization share one kernel and launch wrapper. ``query_splits=None`` selects automatically; an explicit count overrides splitting while retaining the matching prefetch policy. + + ``block_size``, ``head_dim``, and ``query_dtype`` are compile-time + constants. ``query_length`` (MTP) and ``query_group_size`` flatten into + ``TOTAL_ROWS = query_length * query_group_size``. ``query_splits`` assigns + equal groups of query positions to separate CTAs, each processing + ``M_TILES = ceil(TOTAL_ROWS / query_splits / 16)`` MFMA tiles. The partial + output layout remains unsplit so the existing FlyDSL reducer can combine + context partitions without an intermediate transpose. + + A positive ``sliding_window`` includes the query token itself. Tile ranges + cover the union of the MTP windows, and scores are masked per query row. + ``use_sinks`` adds a per-query-head zero-value attention logit in the final + epilogue for direct NP=1 output. Partitioned/planned execution instead adds + it exactly once in the reducer, leaving partial statistics unchanged. + + Tokens at/after ``context_len`` still take part in the PV matmul with a zero + probability, so their V bytes must be finite: the MFMA does not treat a zero + operand as an annihilator (``0 * NaN == NaN``). Whole pages past the + sequence's extent are pinned to block 0 (see ``_load_phys_scalar``); the + remaining slots are the unwritten tail of the last page the sequence owns, + which the caller is expected to leave finite. + + ``wide_kv_addressing`` computes K/V element offsets in 64-bit. It is only + correct-critical once a single cache tensor passes 2 GiB, where the i32 + page-stride product wraps, and it is not free: measured +18% at + ``block_size=64`` (one page index feeds four step-loads there, so i32 lets + the page-base multiply be hoisted and the offsets stay 32-bit adds) and + neutral at ``block_size=16``. The wrapper turns it on from the cache size. """ - gfx950_d128 = "gfx95" in get_rocm_arch() and head_dim == 128 - tuned_per_token_shape = ( - gfx950_d128 - and per_token_kv - and query_dtype == "bf16" - and block_size in (16, 128) - ) + 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 V scheduling also supports f16 and MTP; the numerical fast path + # below further restricts this gate to bf16 single-query decode. + TUNED_SCALAR = TUNED_SHAPE and trans_v and not per_token_kv + # Keep the dense-grid estimate for planned execution too. The plan's # capacity and active token windows do not change the scheduling policy. dense_workgroups = num_seqs * num_kv_heads * num_partitions @@ -105,7 +140,7 @@ def compile_pa_decode_tile( # Small static MTP2/MTP4 grids use one query per CTA. MTP3 stays fused. query_splits = ( query_length - if tuned_per_token_shape + if TUNED_PER_TOKEN and not use_work_plan and num_kv_heads == 1 and query_length in (2, 4) @@ -115,13 +150,15 @@ def compile_pa_decode_tile( ) 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 # Split queries always prefetch. Unsplit decode supports Hkv1 and static # Hkv2; transposed page128 keeps its lower-register schedule above one # workgroup per CU. One query with GQA8/16 already implies one M-tile. - per_token_m1 = ( - tuned_per_token_shape - and query_length == query_splits + PER_TOKEN_M1 = ( + TUNED_PER_TOKEN + and QUERIES_PER_CTA == 1 and query_group_size in (8, 16) and ( query_splits > 1 @@ -135,85 +172,36 @@ def compile_pa_decode_tile( ) ) ) - scalar_prefetch = ( - gfx950_d128 + prefetch_v = PER_TOKEN_M1 or ( + TUNED_SCALAR and block_size == 128 - and trans_v - and not per_token_kv - and query_length * query_group_size <= MFMA_MNK + and TOTAL_ROWS <= MFMA_MNK and num_compute_units < dense_workgroups <= 2 * num_compute_units ) - return _compile_pa_decode_tile( - head_dim=head_dim, - query_group_size=query_group_size, - block_size=block_size, - 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, - prefetch_v=per_token_m1 or scalar_prefetch, - query_splits=query_splits, - use_work_plan=use_work_plan, - sliding_window=sliding_window, - use_sinks=use_sinks, - sink_dtype_str=sink_dtype_str, + # Normalize runtime scheduling metadata to the final specialization before + # looking up the cache, including explicit query-split overrides. + 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, + sliding_window, + use_sinks, + sink_dtype_str, ) + cached = _PA_DECODE_TILE_CACHE.get(cache_key) + if cached is not None: + return cached - -@functools.cache -def _compile_pa_decode_tile( - *, - head_dim: int, - query_group_size: int, - block_size: 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, - prefetch_v: bool = False, - query_splits: int = 1, - use_work_plan: bool = False, - sliding_window: int = 0, - use_sinks: bool = False, - sink_dtype_str: str = "f32", -): - """Build the tile-programming PA-decode kernel + launch wrapper. - - ``block_size``, ``head_dim``, and ``query_dtype`` are compile-time - constants. ``query_length`` (MTP) and ``query_group_size`` flatten into - ``TOTAL_ROWS = query_length * query_group_size``. ``query_splits`` assigns - equal groups of query positions to separate CTAs, each processing - ``M_TILES = ceil(TOTAL_ROWS / query_splits / 16)`` MFMA tiles. The partial - output layout remains unsplit so the existing FlyDSL reducer can combine - context partitions without an intermediate transpose. - - A positive ``sliding_window`` includes the query token itself. Tile ranges - cover the union of the MTP windows, and scores are masked per query row. - ``use_sinks`` adds a per-query-head zero-value attention logit in the final - epilogue for direct NP=1 output. Partitioned/planned execution instead adds - it exactly once in the reducer, leaving partial statistics unchanged. - - Tokens at/after ``context_len`` still take part in the PV matmul with a zero - probability, so their V bytes must be finite: the MFMA does not treat a zero - operand as an annihilator (``0 * NaN == NaN``). Whole pages past the - sequence's extent are pinned to block 0 (see ``_load_phys_scalar``); the - remaining slots are the unwritten tail of the last page the sequence owns, - which the caller is expected to leave finite. - - ``wide_kv_addressing`` computes K/V element offsets in 64-bit. It is only - correct-critical once a single cache tensor passes 2 GiB, where the i32 - page-stride product wraps, and it is not free: measured +18% at - ``block_size=64`` (one page index feeds four step-loads there, so i32 lets - the page-base multiply be hoisted and the offsets stay 32-bit adds) and - neutral at ``block_size=16``. The wrapper turns it on from the cache size. - """ - is_gfx950 = "gfx95" in get_rocm_arch() # Context lengths are int32, so larger windows have identical visibility. # Bound the device constant while retaining the caller's value in the plan. sliding_window = min(sliding_window, 2**31 - 1) @@ -234,40 +222,27 @@ def _compile_pa_decode_tile( "f16", "bf16", ), f"pa_decode_tile only supports query_dtype in ('f16', 'bf16'), got {query_dtype}" - Q_DTYPE = fx.BFloat16 if query_dtype == "bf16" else fx.Float16 + 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}" - QUERIES_PER_CTA = query_length // query_splits # Flattened query-row axis (MTP outer, GQA head inner), tiled into 16-row M-tiles. - TOTAL_ROWS = query_length * query_group_size CTA_ROWS = QUERIES_PER_CTA * query_group_size M_TILES = (CTA_ROWS + MFMA_MNK - 1) // MFMA_MNK ROWS_PADDED = M_TILES * MFMA_MNK - # Share the device/shape gate; scalar V scheduling also supports f16 and - # MTP, while the numerical and MFMA specializations below require bf16. - TUNED_SHAPE = is_gfx950 and head_dim == 128 and block_size in (16, 128) - TUNED_BF16 = TUNED_SHAPE and query_dtype == "bf16" # Fix the tuned scalar-scale shape to the direct-fp8 fast path. This # removes the numerical-policy switch, not the fp8 range/precision tradeoff # documented above. Per-token and other shapes keep normalized conversion. SCALAR_FP8_DECODE = ( - TUNED_BF16 - and trans_v - and not per_token_kv - and query_length == 1 - and query_group_size in (8, 16) + TUNED_SCALAR and IS_BF16 and query_length == 1 and query_group_size in (8, 16) ) # CDNA4 contracts four legacy fp8 K=32 groups in one K=128 instruction. # Concatenating the existing packs preserves cache/LDS layouts for fused # MTP3/MTP4 and query-split MTP4 CTAs. WIDE_FP8_MFMA = ( - TUNED_BF16 - and per_token_kv - and query_length in (3, 4) - and query_group_size == 16 + 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 @@ -278,13 +253,9 @@ def _compile_pa_decode_tile( # register pressure: load plain V during P packing and transposed V after # P publication. Page-16 gathers and narrow addresses retain prefetching. MTP4_PREFETCH_V = MTP4_FUSED and (block_size == 16 or not wide_kv_addressing) - # The selector only enables prefetching for supported single-M-tile shapes. - PER_TOKEN_M1 = per_token_kv and prefetch_v # Page 16 additionally delays next K until PV and uses IGLP to overlap # groups of V loads with QK MFMA instructions. - TUNE_PAGE128 = block_size == 128 and ( - PER_TOKEN_M1 or (TUNED_SHAPE and trans_v and not per_token_kv) - ) + TUNE_PAGE128 = block_size == 128 and (PER_TOKEN_M1 or TUNED_SCALAR) PAGE16_VPIPE = prefetch_v and block_size == 16 # The selected page16 pipeline is always per-token as well. REUSE_KV_PAGES = PER_TOKEN_M1 @@ -2150,4 +2121,10 @@ def pa_decode_tile_launch( stream=stream, ) - return {"launch": pa_decode_tile_launch, "kernel": pa_decode_tile_kernel} + compiled = { + "launch": pa_decode_tile_launch, + "kernel": pa_decode_tile_kernel, + "query_splits": query_splits, + "prefetch_v": prefetch_v, + } + return _PA_DECODE_TILE_CACHE.setdefault(cache_key, compiled) diff --git a/op_tests/test_flydsl_pa_decode.py b/op_tests/test_flydsl_pa_decode.py index 77a337efa55..66c498435d2 100644 --- a/op_tests/test_flydsl_pa_decode.py +++ b/op_tests/test_flydsl_pa_decode.py @@ -525,7 +525,7 @@ def check(cus, expected, **changes): check(boundary - 1, (ql, True), query_splits=ql) check(boundary + 1, (1, False), query_splits=1) - # Runtime size/CU metadata must not fragment the private specialization cache. + # Runtime size/CU metadata must not fragment the kernel specialization cache. cus, expected = points[-1] changes = {"num_seqs": kwargs["num_seqs"] * 2} multiplier = 2 @@ -734,19 +734,18 @@ def test_pa_decode(case, planned, monkeypatch): output, query, _, _, context = args[:5] scratch, sinks, plan = args[14:17], args[-1], options["work_plan"] module = importlib.import_module("aiter.ops.flydsl.pa_decode") - kernel_module = importlib.import_module("aiter.ops.flydsl.kernels.pa_decode_kernel") - compile_tile = module.compile_pa_decode_tile - build_tile = kernel_module._compile_pa_decode_tile + build_tile = module.compile_pa_decode_tile num_compute_units = torch.cuda.get_device_properties( query.device ).multi_processor_count selected = [] compile_args = None - def capture_build(**kwargs): - # Observe the final schedule before the real cache lookup, including hits. - selected.append(kwargs) - return build_tile(**kwargs) + def compile_tile(**kwargs): + # Observe the builder's actual schedule, including cached specializations. + compiled = build_tile(**kwargs) + selected.append(compiled) + return compiled def compile_checked(**kwargs): nonlocal compile_args @@ -771,7 +770,6 @@ def compile_checked(**kwargs): def unexpected_reducer(*_args, **_kwargs): raise AssertionError("static NP=1 must not launch a reducer") - monkeypatch.setattr(kernel_module, "_compile_pa_decode_tile", capture_build) monkeypatch.setattr(module, "compile_pa_decode_tile", compile_checked) if not planned and args[8] == 1: monkeypatch.setattr(module, "launch_pa_decode_ps_reduce", unexpected_reducer) From 466e94c3754e52f91b4b263db74641dcaff8b8b6 Mon Sep 17 00:00:00 2001 From: fsx950223 Date: Thu, 17 Sep 2026 09:56:59 +0000 Subject: [PATCH 56/66] refactor(flydsl): align PA decode API and test arguments Remove ps from the core decode API and forward work_plan through the Python wrapper. Place the wrapper's ignored ps argument before sinks and align the shared unit-test launch helper with that order. Update the API documentation. Validation: - Ruff and Black pass for the changed Python files. - Eight CPU helper checks preserve tensor, sink, window, and plan forwarding. - MI355 pytest remains blocked during collection: Torch infer_schema rejects the wrapper's PADecodePlan | None parameter. No GPU numerical tests passed for this API revision. --- aiter/ops/attention.py | 12 ++++++++++-- aiter/ops/flydsl/pa_decode.py | 9 ++------- docs/flydsl_pa_decode_plan.md | 8 ++++++-- op_tests/test_flydsl_pa_decode.py | 5 ++++- 4 files changed, 22 insertions(+), 12 deletions(-) diff --git a/aiter/ops/attention.py b/aiter/ops/attention.py index eae49925f4d..ad9d6ec0365 100644 --- a/aiter/ops/attention.py +++ b/aiter/ops/attention.py @@ -22,6 +22,7 @@ from ..jit.utils.chip_info import get_cu_num, get_gfx try: + from aiter.ops.flydsl.pa_decode import PADecodePlan from aiter.ops.flydsl.pa_decode import pa_decode as _pa_decode_flydsl except (ImportError, AttributeError, RuntimeError, OSError): _pa_decode_flydsl = None @@ -48,10 +49,17 @@ def pa_decode_flydsl( max_logits: torch.Tensor = None, temporary_output: torch.Tensor = None, alibi_slopes: torch.Tensor = None, + ps: bool = True, sinks: torch.Tensor = None, sliding_window: int = 0, - ps: bool = True, + work_plan: "PADecodePlan | None" = None, ) -> None: + """FlyDSL decode with an optional work plan. + + Pass ``work_plan`` by keyword. The ``ps`` compatibility argument precedes + ``sinks`` and is accepted and ignored. + """ + del ps if _pa_decode_flydsl is None: raise RuntimeError("pa_decode_flydsl requires the `flydsl` package") _pa_decode_flydsl( @@ -75,7 +83,7 @@ def pa_decode_flydsl( alibi_slopes, sinks, sliding_window, - ps, + work_plan=work_plan, ) diff --git a/aiter/ops/flydsl/pa_decode.py b/aiter/ops/flydsl/pa_decode.py index f8bee740f16..9299f236963 100644 --- a/aiter/ops/flydsl/pa_decode.py +++ b/aiter/ops/flydsl/pa_decode.py @@ -209,7 +209,6 @@ def pa_decode( alibi_slopes: torch.Tensor = None, sinks: torch.Tensor = None, sliding_window: int = 0, - ps: bool = True, work_plan: PADecodePlan | None = None, ) -> None: """FlyDSL paged-attention fp8 decode. @@ -227,8 +226,8 @@ def pa_decode( Each entry is an unscaled, zero-value attention logit shared across batch and MTP positions. It contributes to the denominator once, independently of the window; -inf disables a head's sink and +inf suppresses its output. - ALiBi and externally quantized FP8 queries are not supported. The ``ps=False`` - partitioning policy is also not supported. + ALiBi and externally quantized FP8 queries are not supported. Partitioning + is controlled by ``max_context_partition_num`` and ``work_plan``. ``context_lengths`` and ``block_tables`` are GPU-resident, so their values are not inspected here (which would synchronize the device). Callers must @@ -261,10 +260,6 @@ def pa_decode( if sliding_window < -1: raise ValueError("sliding_window must be -1, 0, or positive") sliding_window = max(sliding_window, 0) - if not isinstance(ps, bool): - raise TypeError(f"ps must be a bool, got {type(ps).__name__}") - if not ps: - raise NotImplementedError("pa_decode does not support ps=False") if not isinstance(query_length, int): raise TypeError("query_length must be an int") if query_length < 1: diff --git a/docs/flydsl_pa_decode_plan.md b/docs/flydsl_pa_decode_plan.md index 2cae1d4e230..2a44492a650 100644 --- a/docs/flydsl_pa_decode_plan.md +++ b/docs/flydsl_pa_decode_plan.md @@ -72,8 +72,12 @@ rebuilding a plan, provided the context lengths remain unchanged. ## Usage -Use the Python FlyDSL entry point for planned execution. The shared -`torch.ops.aiter.pa_decode_flydsl` entry point retains its existing static interface. +Use the Python FlyDSL entry point for planned execution. Its signature ends +with `sliding_window, work_plan` and does not accept `ps`. The Python wrapper +`aiter.ops.attention.pa_decode_flydsl` also accepts `work_plan` by keyword. +Its trailing arguments are `alibi_slopes, ps, sinks, sliding_window, work_plan`; +`ps` is accepted and ignored. Prefer keyword arguments for `sinks`, +`sliding_window`, and `work_plan` when switching between these entry points. ```python import torch diff --git a/op_tests/test_flydsl_pa_decode.py b/op_tests/test_flydsl_pa_decode.py index 66c498435d2..425a6ea1513 100644 --- a/op_tests/test_flydsl_pa_decode.py +++ b/op_tests/test_flydsl_pa_decode.py @@ -26,6 +26,7 @@ import aiter from aiter import dtypes, per_tensor_quant, pertoken_quant from aiter.jit.utils.chip_info import get_gfx_runtime +from aiter.ops.attention import pa_decode_flydsl from aiter.test_common import benchmark, run_perftest try: @@ -350,6 +351,8 @@ def scatter(tensor): def _run_flydsl(*args, sliding_window=0, work_plan=None): + # Fixtures use the core signature; the wrapper inserts ps before sinks. + args = (*args[:-1], True, args[-1]) if work_plan is None: torch.ops.aiter.pa_decode_flydsl(*args, sliding_window=sliding_window) else: @@ -361,7 +364,7 @@ def _run_flydsl(*args, sliding_window=0, work_plan=None): sliding_window=sliding_window, plan=work_plan, ) - pa_decode(*args, sliding_window=sliding_window, work_plan=work_plan) + pa_decode_flydsl(*args, sliding_window=sliding_window, work_plan=work_plan) return args[0] From 0b9f100d30349446278ed096b822b91cc1745f0f Mon Sep 17 00:00:00 2001 From: fsx950223 Date: Thu, 17 Sep 2026 10:12:54 +0000 Subject: [PATCH 57/66] fix(flydsl): exclude work_plan from the Torch operator schema Torch schema inference rejects PADecodePlan and prevents importing aiter. Keep work_plan on the Python PA decode API and omit it when registering the static torch.ops entry point. Add an opt-in registration option for optional trailing Python arguments. Infer the schema from a separate signature while preserving annotation globals, defaults, and the existing dispatcher implementation. Validation: - 72 PA decode tests passed on dzm-mi355-gpu34-dev (gfx950, Torch 2.10, FlyDSL 0.3.2), including static/planned execution and graph replay. - Torch 2.8 CPU schema/dispatcher checks passed. - Actual aiter import passed with a simulated FlyDSL backend import failure. - Ruff, Black, and git diff --check passed for the changed files. --- aiter/ops/attention.py | 3 +++ csrc/cpp_itfs/torch_utils.py | 37 ++++++++++++++++++++++++++++++++++-- 2 files changed, 38 insertions(+), 2 deletions(-) diff --git a/aiter/ops/attention.py b/aiter/ops/attention.py index ad9d6ec0365..7315a9f25c0 100644 --- a/aiter/ops/attention.py +++ b/aiter/ops/attention.py @@ -58,6 +58,8 @@ def pa_decode_flydsl( Pass ``work_plan`` by keyword. The ``ps`` compatibility argument precedes ``sinks`` and is accepted and ignored. + ``work_plan`` is available through this Python API; the registered + ``torch.ops.aiter.pa_decode_flydsl`` operator uses static scheduling. """ del ps if _pa_decode_flydsl is None: @@ -97,6 +99,7 @@ def pa_decode_flydsl( "pa_decode_flydsl", pa_decode_flydsl, ["output", "exp_sums", "max_logits", "temporary_output"], + python_only_args=("work_plan",), ) diff --git a/csrc/cpp_itfs/torch_utils.py b/csrc/cpp_itfs/torch_utils.py index 39b485fb12f..c59424121d8 100644 --- a/csrc/cpp_itfs/torch_utils.py +++ b/csrc/cpp_itfs/torch_utils.py @@ -3,7 +3,9 @@ import ctypes +import inspect from collections.abc import Callable +from types import FunctionType import torch from torch.library import Library @@ -90,6 +92,8 @@ def direct_register_custom_op( target_lib: Library | None = None, dispatch_key: str = "CUDA", tags: tuple[torch.Tag, ...] = (), + *, + python_only_args: tuple[str, ...] = (), ): """ `torch.library.custom_op` can have significant overhead because it @@ -102,6 +106,10 @@ def direct_register_custom_op( want to register it to a different library, you can pass the library object to the `target_lib` argument. + `python_only_args` names optional trailing parameters to omit from the + Torch schema, in signature order. Calls through torch.ops use their + defaults; direct Python calls can supply them. + IMPORTANT: the lifetime of the operator is tied to the lifetime of the library object. If you want to bind the operator to a different library, make sure the library object is alive when the operator is used. @@ -113,13 +121,38 @@ def _op_func(*args, **kwargs): log_args(op_func, *args, **kwargs) return op_func(*args, **kwargs) + schema_func = op_func + if python_only_args: + signature = inspect.signature(op_func) + parameters = list(signature.parameters.values()) + excluded = parameters[-len(python_only_args) :] + if tuple(param.name for param in excluded) != python_only_args or any( + param.default is inspect.Parameter.empty for param in excluded + ): + raise ValueError( + "python_only_args must name optional trailing parameters " + "in signature order" + ) + # Preserve the original globals for forward-referenced annotations + # while keeping the public Python signature unchanged. + schema_func = FunctionType( + op_func.__code__, + op_func.__globals__, + op_func.__name__, + op_func.__defaults__, + op_func.__closure__, + ) + schema_func.__signature__ = signature.replace( + parameters=parameters[: -len(python_only_args)] + ) + if hasattr(torch.library, "infer_schema"): - schema_str = torch.library.infer_schema(op_func, mutates_args=mutates_args) + schema_str = torch.library.infer_schema(schema_func, mutates_args=mutates_args) else: # for pytorch 2.4 import torch._custom_op.impl - schema_str = torch._custom_op.impl.infer_schema(op_func, mutates_args) + schema_str = torch._custom_op.impl.infer_schema(schema_func, mutates_args) my_lib = target_lib or aiter_lib my_lib.define(op_name + schema_str, tags=tags) my_lib.impl(op_name, _op_func, dispatch_key=dispatch_key) From 98dd157dd93e31083c81a4ffff9f5cf5a1cb46bf Mon Sep 17 00:00:00 2001 From: fsx950223 Date: Fri, 18 Sep 2026 09:58:12 +0000 Subject: [PATCH 58/66] refactor(flydsl): simplify and broaden PA decode plan scheduling Derive planned specializations from capacity and visible tiles, and widen the M1 and batch-first selectors using continuous ranges. Simplify compiler caching and returned metadata, refine planned scratch and reduction paths, expand the generic unit test, and remove the obsolete plan document. Validation: 176 pytest cases passed on dzm-mi355-gpu34-dev (gfx950). Ruff and staged diff checks passed. Committed Python sources match the GPU-tested source archive byte-for-byte. Known performance caveat: B12/W8192 batch-first remains about 13.3% slower than flat for decode+reduce in the repeat measurement; plan refresh is excluded from those timings. --- aiter/ops/attention.py | 1 + aiter/ops/flydsl/kernels/pa_decode_kernel.py | 513 ++++++++---- aiter/ops/flydsl/kernels/pa_decode_plan.py | 3 +- aiter/ops/flydsl/kernels/pa_decode_reduce.py | 374 ++++++--- aiter/ops/flydsl/pa_decode.py | 41 +- docs/flydsl_pa_decode_plan.md | 148 ---- op_tests/test_flydsl_pa_decode.py | 804 +++++++++++++++++-- 7 files changed, 1418 insertions(+), 466 deletions(-) delete mode 100644 docs/flydsl_pa_decode_plan.md diff --git a/aiter/ops/attention.py b/aiter/ops/attention.py index 7315a9f25c0..b36887d42e3 100644 --- a/aiter/ops/attention.py +++ b/aiter/ops/attention.py @@ -60,6 +60,7 @@ def pa_decode_flydsl( ``sinks`` and is accepted and ignored. ``work_plan`` is available through this Python API; the registered ``torch.ops.aiter.pa_decode_flydsl`` operator uses static scheduling. + Positive ``sliding_window`` requires a work plan through the Python API. """ del ps if _pa_decode_flydsl is None: diff --git a/aiter/ops/flydsl/kernels/pa_decode_kernel.py b/aiter/ops/flydsl/kernels/pa_decode_kernel.py index 7aa536e3240..3f7bb28c14a 100644 --- a/aiter/ops/flydsl/kernels/pa_decode_kernel.py +++ b/aiter/ops/flydsl/kernels/pa_decode_kernel.py @@ -34,8 +34,6 @@ an LDS round-trip on P transposing ownership between the two MMAs. """ -import functools - import flydsl.compiler as flyc import flydsl.expr as fx from flydsl.compiler.kernel_function import CompilationContext @@ -46,7 +44,7 @@ from flydsl.runtime.device import get_rocm_arch from . import dpp_utils -from .tensor_shim import ptr_buf_tensor +from .tensor_shim import buf_base_i64, buf_copy_store, ptr_buf_tensor from .utils import rcp_f32 MFMA_MNK = ( @@ -65,9 +63,6 @@ _PA_DECODE_TILE_CACHE = {} -# Repeated calls bypass the builder's closure setup. Keep this metadata cache -# bounded; the specialization cache below shares kernels across runtime sizes. -@functools.lru_cache(maxsize=128) def compile_pa_decode_tile( *, head_dim: int, @@ -85,16 +80,20 @@ def compile_pa_decode_tile( 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 the schedule and cache the PA-decode kernel + launch wrapper. - Batch/head/CU counts select query splitting and V prefetching. Calls that + Batch/head/CU counts and plan capacity select the schedule. Calls that select the same specialization share one kernel and launch wrapper. ``query_splits=None`` selects automatically; an explicit count overrides - splitting while retaining the matching prefetch policy. + splitting while retaining the matching prefetch policy. Planned splitting + uses a host-known task-count upper bound, without reading GPU metadata. + Sufficient window capacity also proves that every active planned task has + one KV tile, selecting a single-iteration specialization. ``block_size``, ``head_dim``, and ``query_dtype`` are compile-time constants. ``query_length`` (MTP) and ``query_group_size`` flatten into @@ -104,8 +103,28 @@ def compile_pa_decode_tile( output layout remains unsplit so the existing FlyDSL reducer can combine context partitions without an intermediate transpose. - A positive ``sliding_window`` includes the query token itself. Tile ranges - cover the union of the MTP windows, and scores are masked per query row. + The single-tile specialization requires a standard work plan whose + host-known capacity proves every nonempty record contains one KV tile. + Padding records still skip the whole task, and partial outputs keep the + same packed layout and reducer as multi-tile plans. + + Selected split-query, single-tile plans map packed slots over a (B, C/B) + physical grid. Physical x is not a sequence index: each active task still + takes its sequence from the plan record. This specialization requires + positive launch B, capacity C divisible by B, and launch sizes consistent + with the schedule selection. These runtime size preconditions are checked + here, not dynamically in the JIT launch. Neither B nor C/B enters the + kernel specialization-cache key. + + Planned BF16 D128 partial output automatically uses cached 64-bit buffer + stores on gfx950 when a packed slot's byte span fits a positive signed i32. + Each descriptor is rebased in i64 to one slot. The view retains only the + input dtype's 2-byte alignment; offsets inside the slot are multiples of + 8 bytes. Partial statistics and direct output are unchanged. + + A positive ``sliding_window`` requires a work plan and includes the query + token itself. The plan covers the union of the MTP windows; scores are + masked per query row within those tiles. ``use_sinks`` adds a per-query-head zero-value attention logit in the final epilogue for direct NP=1 output. Partitioned/planned execution instead adds it exactly once in the reducer, leaving partial statistics unchanged. @@ -124,6 +143,8 @@ def compile_pa_decode_tile( the page-base multiply be hoisted and the offsets stay 32-bit adds) and neutral at ``block_size=16``. The wrapper turns it on from the cache size. """ + 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) @@ -132,20 +153,50 @@ def compile_pa_decode_tile( # below further restricts this gate to bf16 single-query decode. TUNED_SCALAR = TUNED_SHAPE and trans_v and not per_token_kv - # Keep the dense-grid estimate for planned execution too. The plan's - # capacity and active token windows do not change the scheduling policy. + # Preserve the dense-grid estimate for the existing V-prefetch policies. 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: + # An unaligned union of Q causal windows spans at most this many + # tiles per sequence. Padding slots do not provide useful 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 + ) + # With n nonempty sequences, T <= n * window_tiles and the + # remaining budget is at least n * (window_tiles - 1). Prefix + # apportionment therefore grants each sequence at least tiles-1 + # extras, so clamping leaves exactly one task per visible tile. + # This holds for variable lengths and every in-place plan 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: - # Small static MTP2/MTP4 grids use one query per CTA. MTP3 stays fused. + # Small MTP2/MTP4 grids use one query per CTA. Proven one-tile MTP4 + # page128 plans favor several lightweight M1 CTAs per CU over fused + # MTP CTAs, so allow larger grids there. Preserve the existing + # threshold for static, MTP2 and multi-tile paths. + 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 not use_work_plan and num_kv_heads == 1 and query_length in (2, 4) and query_group_size == 16 - and query_length * dense_workgroups <= 2 * num_compute_units + 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" @@ -153,23 +204,20 @@ def compile_pa_decode_tile( QUERIES_PER_CTA = query_length // query_splits TOTAL_ROWS = query_length * query_group_size - # Split queries always prefetch. Unsplit decode supports Hkv1 and static - # Hkv2; transposed page128 keeps its lower-register schedule above one - # workgroup per CU. One query with GQA8/16 already implies one M-tile. + # One query with GQA8..16 implies one M-tile. Split queries, page16 and + # plain V prefetch unconditionally; transposed page128 uses the grid load + # across all KV heads. One-tile plans allow two effective tasks per CU; + # other grids keep the lower-register schedule above one task per CU. PER_TOKEN_M1 = ( TUNED_PER_TOKEN and QUERIES_PER_CTA == 1 - and query_group_size in (8, 16) + and 8 <= query_group_size <= 16 and ( query_splits > 1 - or ( - (num_kv_heads == 1 or (not use_work_plan and num_kv_heads == 2)) - and ( - block_size == 16 - or not trans_v - or dense_workgroups <= num_compute_units - ) - ) + 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 ( @@ -178,6 +226,22 @@ def compile_pa_decode_tile( and TOTAL_ROWS <= MFMA_MNK and num_compute_units < dense_workgroups <= 2 * num_compute_units ) + # Bound batch/window sizes and keep the exact one-tile task budget + # within the small planned reducer. Its cap may exceed the window bound + # without increasing the packed grid capacity. + 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 + ) # Normalize runtime scheduling metadata to the final specialization before # looking up the cache, including explicit query-split overrides. cache_key = ( @@ -194,6 +258,8 @@ def compile_pa_decode_tile( prefetch_v, query_splits, use_work_plan, + single_tile_plan, + batch_first_plan_grid, sliding_window, use_sinks, sink_dtype_str, @@ -202,6 +268,13 @@ def compile_pa_decode_tile( 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 + ) # Context lengths are int32, so larger windows have identical visibility. # Bound the device constant while retaining the caller's value in the plan. sliding_window = min(sliding_window, 2**31 - 1) @@ -210,9 +283,9 @@ def compile_pa_decode_tile( 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 head_dim % MFMA_MNK == 0, ( + f"head_dim {head_dim} must be a multiple of {MFMA_MNK}" + ) assert block_size in ( 16, 64, @@ -231,7 +304,22 @@ def compile_pa_decode_tile( CTA_ROWS = QUERIES_PER_CTA * query_group_size M_TILES = (CTA_ROWS + MFMA_MNK - 1) // MFMA_MNK ROWS_PADDED = M_TILES * MFMA_MNK - + # Keep local BF16 maxima in f32 instead of packing/unpacking every max. + # Other dtypes and non-windowed/static paths retain their input-type reduction. + Q_ABSMAX_F32 = sliding_window > 0 and TUNED_PER_TOKEN + # Cooperatively stage each token's scales once instead of having all four + # rgroups load/store the same 64 values. Reuse this producer for the legacy + # multi-tile QL1 path and prefetched single-tile QL1 CTAs; keep + # the LDS layout, consumers and publication barriers unchanged. + 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 + ) # Fix the tuned scalar-scale shape to the direct-fp8 fast path. This # removes the numerical-policy switch, not the fp8 range/precision tradeoff # documented above. Per-token and other shapes keep normalized conversion. @@ -272,25 +360,25 @@ def compile_pa_decode_tile( 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" + 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" + assert head_dim % (NWARP * MFMA_MNK) == 0, ( + "head_dim must split across the 4 warps for PV" + ) # head_dim splits into 16-element chunks (QK_CHUNK_ELEMS, one dwordx4 load); # RGROUP_QUARTERS of them make a 64-element fetch group, QKHE_LOOP groups total. 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}" + 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" @@ -342,19 +430,24 @@ def compile_pa_decode_tile( # -- LDS layout (shared across the 4 warps) -- # sQ: fp8[ROWS_PADDED,head_dim] staged+quantized query. sP: fp8[16,TILE_TOK] - # quantized probs. sQscale: f32[ROWS_PADDED] except when Q is converted - # directly to fp8. sLmax/sLsum: cross-warp scratch. + # quantized probs, sharing sQ's storage after Q has been read into registers. + # sQscale: f32[ROWS_PADDED] except when Q is converted directly to fp8. + # sLmax/sLsum: cross-warp scratch. # sVPage: V page-index broadcast. sKScale/sVScale/sVScaleMax: per-token K/V # scale staging. No sO/sM/sL/sCorr: PV output is register-resident/loop-carried # and stored straight to global (V=A/P=B swap). f32 = 4 sQ_bytes = ROWS_PADDED * head_dim * 1 # fp8 - sP_off = sQ_bytes + # Every path finishes all M-tiles' first QK and executes the max-publication + # barrier before packing P. No later KV tile reads sQ, so P may overwrite + # it without another barrier. Keep the still-live Q scales outside both + # regions; both sizes are already multiples of 16 bytes. + sP_off = 0 # +16B row padding breaks a 32-bank LDS conflict on the P-pack writes while # keeping the row 16B-aligned for PV's ds_read_b128. SP_ROW_BYTES = TILE_TOK + 16 sP_bytes = P_BUFFERS * MFMA_MNK * SP_ROW_BYTES # fp8, padded rows - sQscale_off = sP_off + sP_bytes + sQscale_off = max(sQ_bytes, sP_bytes) sQscale_bytes = 0 if SCALAR_FP8_DECODE else ROWS_PADDED * f32 # Keep tuned cross-wave reduction rows 16-byte aligned for vector reads. # Other paths retain the original bank-conflict padding. @@ -365,12 +458,13 @@ def compile_pa_decode_tile( # V page-index broadcast (V's page depends on `rgroup`, shared across warps). sVPage_off = sLsum_off + P_BUFFERS * MFMA_MNK * NWARP_PAD * f32 sVPage_bytes = NWARP * PAGES_PER_CHUNK * 4 # i32 - # Per-token KV-scale staging (per_token only), double-buffered so the tt+1 - # prefetch (into the other buffer) doesn't clobber the current tile's scales. + # Generic per-token KV-scale staging is double-buffered so tt+1 prefetch + # cannot clobber current scales. Proven single-tile tasks never prefetch. 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 = 2 * KV_BUF_STRIDE if per_token_kv else 0 + 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 @@ -422,9 +516,22 @@ def _pa_decode_tile_task( 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): - part = seq # Packed work slot, shared by all KV heads. - seq = planned_seq - partial_slot = kv_h * fx.Int32(gpu.grid_dim.x) + part + 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")) + ) + # Packed variable-length requests need not own C/B tasks. + # Physical x only groups slots; the plan owns sequence IDs. + 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 @@ -432,6 +539,29 @@ def _pa_decode_tile_task( 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 each factor before multiplication: the packed allocation + # may exceed 2 GiB even though one slot's offsets fit signed i32. + if const_expr(batch_first_plan_grid): + # The two-dimensional grid still addresses capacity C, not B. + # Keep the complete packed offset correct for every KV head. + 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, + # A contiguous BF16 view may start at an odd storage offset: + # do not strengthen its base alignment from 2 to 8 bytes. + unit_stride=1, + num_records_bytes=pout_slot_bytes, + ) if const_expr(DIRECT_SINKS): sink_token = fx.recast_iter(SINK_DTYPE, sinks_ptr) @@ -555,19 +685,8 @@ def _v_load16(byte_off): part_end = planned_end else: num_tiles = (context_len + TILE_TOK - 1) // TILE_TOK - active_tiles = num_tiles - if const_expr(sliding_window > 0): - window_start = ( - fx.Int64(context_len) - (query_length - 1) - sliding_window - ) - first_tile = fx.Int32( - (window_start > 0).select(window_start, 0) // TILE_TOK - ) - active_tiles = num_tiles - first_tile - tiles_per_part = (active_tiles + NP - 1) // NP + tiles_per_part = (num_tiles + NP - 1) // NP part_start = part * tiles_per_part - if const_expr(sliding_window > 0): - part_start = first_tile + part_start part_end_raw = part_start + tiles_per_part part_end = (part_end_raw < num_tiles).select(part_end_raw, num_tiles) @@ -595,11 +714,13 @@ 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 else 1 + scale_load_width = ( + NCHUNK if block_size >= 64 and not UNIQUE_SCALE_STAGING else 1 + ) scale_copy_op = ( - fx.rocdl.BufferCopy128b() - if block_size >= 64 - else fx.rocdl.BufferCopy32b() + 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) @@ -685,14 +806,16 @@ def _k_page_read_warp(): def _kv_buf_off(tt_val): # ping-pong buffer byte offset (0 when single-buffered). - if const_expr(per_token_kv): + 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 = lane16 * NCHUNK + 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) @@ -891,6 +1014,12 @@ def _f32_to_fp8_words(vf32): 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) + # Each M-tile quantizes 16 rows of the flattened (MTP position, GQA # head) axis: `flat_idx = m*16 + qh_local`, decomposed as # `qi = flat_idx // query_group_size`, `gs_head = flat_idx % @@ -909,11 +1038,11 @@ def _quant_q_row(m, q_row_off, q_units): _f32_to_fp8_words(q_units[u].to(fx.Float32)), ) else: - absmax = fmath.absf(q_units[0]).reduce(ReductionOp.MAX).to(fx.Float32) + absmax = _q_local_absmax(q_units[0]) for u in range_constexpr(1, N_QLOADS): absmax = fx.maxnumf( absmax, - fmath.absf(q_units[u]).reduce(ReductionOp.MAX).to(fx.Float32), + _q_local_absmax(q_units[u]), ) for sh in (8, 4, 2, 1): absmax = fx.maxnumf(absmax, dpp_utils.dpp_xor_f32(absmax, sh)) @@ -1039,21 +1168,27 @@ def _v_ops(phys_row, vh): fx.rocdl.sched_vmem(len(ops) // 2) return ops # NVOPS i64, the 64-token contiguous run for this head - # One query per CTA has a wave-uniform causal bound, including split - # MTP. Keep the original query position when selecting visible tokens. + # Distance from each query to the newest query, including split MTP. + # Subtract from tile_valid inside the loop, keeping bounds tile-relative. if const_expr(QUERIES_PER_CTA == 1): - causal_bound = [ - context_len - (query_length - 1) + query_begin + causal_offset = [ + query_length - 1 - query_begin for _m in range_constexpr(M_TILES) ] else: - causal_bound = [ - context_len - - (query_length - 1) - + query_begin - + (m * MFMA_MNK + lane16) // query_group_size + 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 are not stored; keep their bounds in the same + # nonnegative-offset domain as real query rows. + causal_offset = [ + (offset > 0).select(offset, 0) for offset in causal_offset + ] def _o_slot(m, vh): return 2 + STATE_PER_M * m + vh @@ -1078,14 +1213,23 @@ def _l_slot(m): 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) - for tt, ostate in range(part_start, part_end, 1, init=init_state): + # Fixed bounds let the compiler remove the loop/history for proven + # single-tile plans. Keep tt absolute for cache addressing and masks; + # the outer work-record guard still excludes every padded task. + 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(tt) + 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 in the scalar phase-split path, # and with V loads in the tuned page-16 decode pipeline. @@ -1110,6 +1254,14 @@ def _l_slot(m): # context hold real data, and using it keeps both masks independent # of the query row. tile_valid = context_len - tok0 + window_left = None + if const_expr(sliding_window > 0): + # A planned tile has 1 <= tile_valid <= INT32_MAX, and W is + # clamped to INT32_MAX. Clamp a negative tile-local left edge + # before subtracting MTP offsets: all tile tokens are >= 0, + # so this preserves the mask without underflow for huge W. + 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( @@ -1124,12 +1276,9 @@ def _l_slot(m): # Scales outside every query's window must not shrink # visible probabilities to zero during FP8 normalization. # Keep this bound query-independent for shared MTP scales. - first_visible = ( - fx.Int64(context_len) - - (query_length - 1) - - sliding_window - - fx.Int64(tok0) - ).to(fx.Float32) + first_visible = (window_left - (query_length - 1)).to( + fx.Float32 + ) window_scale_thr = fx.Vector.from_elements( [ first_visible @@ -1215,7 +1364,7 @@ def _lmax_off_m(m): frag_Ss.append(fx.Vector(acc)) scale = scale_qk * fx.Float32(q_scale_vec[m]) - n_valid_tile = (causal_bound[m] - tok0).to(fx.Float32) + 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 ) @@ -1226,9 +1375,9 @@ def _lmax_off_m(m): if const_expr(sliding_window > 0): # Subtract in integer space first so a large window # cannot round its left edge before the comparison. - first_valid_tile = ( - fx.Int64(causal_bound[m]) - sliding_window - fx.Int64(tok0) - ).to(fx.Float32) + 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) @@ -1266,9 +1415,7 @@ def _lmax_off_m(m): ) needs_mask = tile_valid < TILE_TOK + query_length - 1 if const_expr(sliding_window > 0): - needs_mask = needs_mask | ( - tok0 < context_len - sliding_window - ) + needs_mask = needs_mask | (window_left > 0) if needs_mask: masked_all = fx.Vector.from_elements( [ @@ -1320,7 +1467,7 @@ def _lmax_off_m(m): if const_expr(MTP4_FUSED) else k_cur ) - if tt1 < part_end: + if const_expr(not single_tile_plan) and tt1 < part_end: if const_expr(MTP4_FUSED): # Publish next pages/scales now; defer next K until P # packing releases the saved score registers. @@ -1336,7 +1483,7 @@ def _lmax_off_m(m): gpu.barrier() v_page_next = v_page_cur - if tt1 < part_end: + 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 @@ -1372,12 +1519,15 @@ def _lmax_off_m(m): m_new_saved = [] safe_max_saved = [] for m in range_constexpr(M_TILES): - m_new = fx.maxnumf( - ostate[_m_slot(m)], - _ld_lw_row(_lmax_off_m(m), lane16).reduce( - ReductionOp.MAX, fastmath=fm_nnan - ), - fastmath=fm_nnan, + 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) @@ -1425,7 +1575,7 @@ def _lmax_off_m(m): warp, ls, ) - if tt1 < part_end: + if const_expr(not single_tile_plan) and tt1 < part_end: # Let next K overlap the P-publication barrier and PV, # after saved scores have been consumed by P packing. k_next = _k_ops_from_phys(_k_page_read_warp()) @@ -1450,11 +1600,17 @@ def _lmax_off_m(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 = fx.Float32( - fx.exp2(m_prev - safe_max, fastmath="fast") + 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 = l_prev * corr_reg + gsum + 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, @@ -1470,12 +1626,20 @@ def _lmax_off_m(m): op = fx.Vector(acc) * fx.Vector.from_elements( [v_max_scaled], dtype=fx.Float32 ).broadcast_to(OP_ELEMS) - o_acc[vh] = o_acc[vh] * corr_b + op - if const_expr(MTP4_PREFETCH_V and m == M_TILES - 1): + 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 + ): # The current chunk is dead; reuse its registers # for the next tile while the final PV finishes. v_next_chunk = fx.Vector.filled(NVOPS, 0, fx.Int64) - if tt1 < part_end: + 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 ) @@ -1484,14 +1648,19 @@ def _lmax_off_m(m): if const_expr(m < M_TILES - 1): fx.rocdl.sched_barrier(0) if const_expr(MTP4_PREFETCH_V): - 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, - ) + if const_expr(single_tile_plan): + # Preserve the loop-state shape; no next V is + # loaded, and the epilogue never consumes this slot. + 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 @@ -1523,12 +1692,13 @@ def _lmax_off_m(m): [norm_factor], dtype=fx.Float32 ).broadcast_to(4) - m_new = fx.maxnumf( - m_prev, - _ld_lw_row(_lmax_off_m(m), lane16).reduce( - ReductionOp.MAX, fastmath=fm_nnan - ), - fastmath=fm_nnan, + 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) ) # Fully-invalid row: use 0 as the effective max so masked lanes # give exp2(-inf-0)==0 (avoids the -inf-(-inf) cancellation). @@ -1575,14 +1745,20 @@ def _lmax_off_m(m): # swap, so correction/denominator are per-lane scalars (no sCorr). # Empty history must contribute zero: exp2(-inf - safe_max). # Replacing m_prev with 0 can overflow for negative logits. - corr_reg = fx.Float32( - fx.exp2(m_prev - safe_max, fastmath="fast") + 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 = l_prev * corr_reg + gsum + 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, @@ -1604,7 +1780,11 @@ def _lmax_off_m(m): op = op * fx.Vector.from_elements( [v_max_scaled], dtype=fx.Float32 ).broadcast_to(OP_ELEMS) - o_acc[vh] = o_acc[vh] * corr_b + op + 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]) # A shared P/Lsum slot needs all reads to finish before # the next M-tile overwrites it. With alternating slots, @@ -1632,7 +1812,7 @@ def _lmax_off_m(m): # Publish next page ids and scales before the pass-1 barrier. # Page 16 defers K until PV; page 128 overlaps K with softmax. k_next = k_cur - if tt1 < part_end: + if const_expr(not single_tile_plan) and tt1 < part_end: if const_expr(REUSE_KV_PAGES): # Publish next pages/scales once. Page 16 defers K # until the P barrier; per-token page 128 issues it @@ -1660,16 +1840,14 @@ def _lmax_off_m(m): if const_expr(SCALAR_FP8_DECODE) else scale_qk * _ld1(sQscale_off, lane16) ) # per-qhead positive score scale - n_valid_tile = (causal_bound[0] - tok0).to(fx.Float32) + 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 = ( - fx.Int64(causal_bound[0]) - sliding_window - fx.Int64(tok0) - ).to(fx.Float32) + 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) @@ -1764,7 +1942,7 @@ def _lmax_off_m(m): gpu.barrier() # V page-index row for next tile, now visible after the barrier. v_page_next = v_page_cur - if tt1 < part_end: + 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 # per_token_kv: combine warps' max V-scale into the tile's @@ -1781,17 +1959,20 @@ def _lmax_off_m(m): ).broadcast_to(4) # pass 2: global max over warps -> exp -> fp8 P pack (-> sP) -> sum if const_expr(M1_SCALE_BEFORE_MASK): - m_new = fx.maxnumf( - m_prev, - _ld_lw_row(sLmax_off, lane16).reduce( - ReductionOp.MAX, fastmath=fm_nnan - ), - fastmath=fm_nnan, + 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: - m_new = fx.maxnumf( - m_prev, - _ld_lw_row(sLmax_off, lane16).reduce(ReductionOp.MAX), + 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 -inf in the loop-carried/persisted max for an empty # partition, but use zero as the exponent reference so @@ -1858,7 +2039,11 @@ def _lmax_off_m(m): ls = ls + ls.shuffle_xor(sh, WAVE) # PV (V=A, P=B) -> output [head-dim, query-row=lane16]; same as # the phase-split path. - corr_reg = fx.Float32(fx.exp2(m_prev - softmax_max, fastmath="fast")) + 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() @@ -1866,11 +2051,13 @@ def _lmax_off_m(m): # The page matrix is now visible. Issue next K before the # current P read/PV so those independent operations can # cover its VMEM latency. - if tt1 < part_end: + 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 = l_prev * corr_reg + gsum + 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 ) @@ -1894,7 +2081,9 @@ def _lmax_off_m(m): op = op * fx.Vector.from_elements( [v_max_scaled], dtype=fx.Float32 ).broadcast_to(OP_ELEMS) - o_acc[vh] = o_acc[vh] * corr_b + op + 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) @@ -1957,6 +2146,18 @@ def _emit(o_norm, sub, query_idx, query_head): + sub * OP_ELEMS ) fx.ptr_store(o_norm, fx.add_offset(output, out_offset)) + elif const_expr(buffer_plan_output): + # Existing row guards keep all four BF16 elements inside + # this slot. Preserve the 8-byte width and cached policy. + 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 @@ -2044,7 +2245,13 @@ def _run_task(seq, start, end, context): # guard the whole task body to avoid Q processing and scratch # writes for an empty record. The condition is CTA-uniform, so # all barriers in the task body remain convergent. - slot = fx.Int32(gpu.block_id("x")) + 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), @@ -2113,9 +2320,13 @@ def pa_decode_tile_launch( num_seqs, ).launch( grid=( - work_capacity if use_work_plan else num_seqs, - num_kv_heads * query_splits, - 1 if use_work_plan else NP, + (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, @@ -2124,7 +2335,5 @@ def pa_decode_tile_launch( compiled = { "launch": pa_decode_tile_launch, "kernel": pa_decode_tile_kernel, - "query_splits": query_splits, - "prefetch_v": prefetch_v, } 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 index 836f02b9c1b..3c3f6f316e6 100644 --- a/aiter/ops/flydsl/kernels/pa_decode_plan.py +++ b/aiter/ops/flydsl/kernels/pa_decode_plan.py @@ -133,7 +133,8 @@ def plan_pa_decode( Allocate once outside graph capture, then pass ``plan=...`` to refresh the same metadata in place. Include this refresh in end-to-end measurements. - The budget counts CTAs over all KV heads, with fused query positions. + The budget counts task slots over all KV heads before query splitting; + splitting queries can launch multiple CTAs per slot without extra scratch. This is opt-in: uniform or short-context workloads may favor static splits. A positive ``sliding_window`` counts visible tokens including the query's diff --git a/aiter/ops/flydsl/kernels/pa_decode_reduce.py b/aiter/ops/flydsl/kernels/pa_decode_reduce.py index 0a2cbf34ec4..12596fb5f8e 100644 --- a/aiter/ops/flydsl/kernels/pa_decode_reduce.py +++ b/aiter/ops/flydsl/kernels/pa_decode_reduce.py @@ -9,6 +9,9 @@ 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 + MAX_CONTEXT_PARTITIONS = 256 _DTYPE_MAP = { "f32": fx.Float32, @@ -71,6 +74,9 @@ def compile_pa_decode_ps_reduce( 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 the partitioned-softmax reduction used by ``pa_decode``. @@ -80,6 +86,15 @@ def compile_pa_decode_ps_reduce( waves. Other head sizes retain the register-only lane-striped fallback. A sink is a per-query-head zero-value logit: include it in the shared max for stability and add its mass to the denominator once after summing KV. + ``query_group_size=None`` retains runtime GQA addressing; a positive + specialization constant must match the GQA passed to the launch wrapper. + ``bounded_plan_logits`` requires packed, contiguous logits rows and a + positive NP-part byte span no larger than INT32_MAX. A per-sequence buffer + bound then zero-fills inactive parts without reading adjacent sequences. + ``vectorize_plan_logits`` additionally requires D=128, bf16/f16 logits, + a four-byte-aligned logits base and even head/part/group strides. One wave + then assigns two consecutive output elements to each lane, sharing the + unchanged partition weights. Output stores retain scalar alignment. """ _validate_pa_decode_ps_reduce_config( max_context_partition_num=max_context_partition_num, @@ -88,6 +103,20 @@ def compile_pa_decode_ps_reduce( logits_dtype_str=logits_dtype_str, sink_dtype_str=sink_dtype_str, ) + 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] @@ -120,7 +149,9 @@ def compile_pa_decode_ps_reduce( 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, 1, 1] + [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 @@ -182,28 +213,58 @@ def pa_decode_ps_reduce_kernel( 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 sequences may point one past the packed allocation. The + # <=64 path uses clamped loads, so give them a valid base slot; + # every loaded value is masked before it contributes to a result. + 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 before narrowing offsets: the whole packed workspace may + # exceed 2 GiB even when this sequence's NP-part span is small. + 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) - c_qgs = query_group_size lane = tid & c_wave_mask - group_idx = eqgs_idx % c_qgs + 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): + def _wave_reduce_max(value, offsets=reduce_shuffle_offsets): reduced = value - for offset in reduce_shuffle_offsets: + for offset in offsets: reduced = reduced.maximumf( reduced.shuffle_xor(fx.Int32(offset), c_warp_size) ) return reduced - def _wave_reduce_sum(value): + def _wave_reduce_sum(value, offsets=reduce_shuffle_offsets): reduced = value - for offset in reduce_shuffle_offsets: + for offset in offsets: reduced = fx.Float32( reduced.addf( reduced.shuffle_xor(fx.Int32(offset), c_warp_size), @@ -218,6 +279,124 @@ def _sink_exp(sink_value, safe_max): 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): + # Both widths use the same packed-scratch and softmax rules. + # Invalid lanes read this sequence's first partition; count=0 + # has a safe global slot 0 base. Select out NaNs before math. + 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) + + # Preload before softmax so independent loads can overlap. The + # caller's uniform branch encloses these loads, not just the FMAs. + 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): + # The descriptor ends at count*part_stride, before the + # next sequence or NaN padding. Even count=0 is a valid + # zero-sized resource; raw buffer loads return zero. + planned_logits_offset = ( + planned_part_idx * stride_logits_part + + eqgs_idx * stride_logits_group + + tid * fx.Int32(planned_elements_per_thread) + ) + # With unit_stride=1 this remains an element offset, not + # a vector-unit index. A D=128 pair ends within its row; + # even count=0 suppresses the complete 32-bit buffer 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): + # Keep the prefetched pair packed in 32 bits across + # softmax; widen only the part being accumulated below. + 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] + # Elementwise mul then add, with the same increasing-part + # order as the scalar path: no FMA or reassociation hint. + planned_acc = planned_acc + planned_part_logits * planned_weight + return planned_acc + if fx.const_expr(use_parallel_lds): # One wave materializes the normalized partition weights once. # All output waves then reuse those weights from LDS and split the @@ -305,6 +484,7 @@ def _sink_exp(sink_value, safe_max): 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 @@ -315,10 +495,10 @@ def _sink_exp(sink_value, safe_max): part_idx < c_part_num ) if lane_in_range: - weight = scaled_sums[chunk_idx] / safe_global_exp_sum + weight = scaled_sums[chunk_idx] * parallel_inv_exp_sum lds_weights[part_idx] = weight else: - weight = scaled_sums[chunk_idx] / safe_global_exp_sum + weight = scaled_sums[chunk_idx] * parallel_inv_exp_sum if fx.const_expr(chunk_size == warp_size): lds_weights[part_idx] = weight else: @@ -395,99 +575,92 @@ def _sink_exp(sink_value, safe_max): acc = acc + fx.Float32(lds_partials[partial_offset]) elif fx.const_expr(max_context_partition_num <= warp_size): - # Exact powers of two have no inactive lanes inside their reduction - # subgroup. Keep that original path unchanged; only partial - # subgroups need an EXEC-masked load to avoid carrying a predicate - # across the shuffle sequence. - if fx.const_expr( - max_context_partition_num == reduce_width and not use_work_plan - ): - 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) + if fx.const_expr(use_work_plan): + # Both arms of the dynamic short-count branch, and its + # initial carried value, must have 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): + # The count is uniform across this CTA. Keep the complete + # short reduction under the branch so large NP does not + # force short requests to preload and accumulate padding. + # Small NP stays straight-line: the branch costs more than + # the few clamped loads it removes on those specializations. + if c_part_num <= c_four: + acc = _reduce_planned_wave(4, (2, 1)) + elif c_part_num <= fx.Int32(8): + # Medium-count requests need only the first eight + # lanes/partials, even when NP reserves up to 64. + # Keep the preload, softmax and accumulation inside + # this uniform branch, including empty-row handling. + 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): + # Preserve the static mapping: exact powers of two have no + # inactive lanes inside their reduction subgroup, while + # partial subgroups use EXEC-masked loads. + 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 - + part_idx * stride_exp_sums_part + + lane * 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]) + 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) - 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: + 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 - ) - if fx.const_expr(use_work_plan): - weight_local = zero_f - if lane_in_range: - weight_local = scaled_sum / safe_global_exp_sum - else: - weight_local = scaled_sum / safe_global_exp_sum - weight_local_i32 = weight_local.bitcast(fx.Int32) + 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 - if fx.const_expr(use_work_plan): - for part_idx in fx.range_constexpr(max_context_partition_num): - c_part_idx = fx.Int32(part_idx) - if c_part_idx < c_part_num: - 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: + 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( @@ -586,6 +759,7 @@ def _sink_exp(sink_value, safe_max): 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): @@ -598,7 +772,7 @@ def _sink_exp(sink_value, safe_max): weight_local_i32 = zero_f.bitcast(fx.Int32) if fx.Int32(chunk_base) < c_part_num: weight_local_i32 = ( - scaled_sums[chunk_idx] / safe_global_exp_sum + 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 @@ -623,7 +797,7 @@ def _sink_exp(sink_value, safe_max): acc = acc + part_logits * weight else: weight_local_i32 = ( - scaled_sums[chunk_idx] / safe_global_exp_sum + 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 @@ -646,7 +820,10 @@ def _sink_exp(sink_value, safe_max): part_logits = fx.Float32(logits[logits_offset]) acc = acc + part_logits * weight - query_idx = eqgs_idx // c_qgs + 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 = ( @@ -663,9 +840,14 @@ def _sink_exp(sink_value, safe_max): + query_idx * stride_output_len + kv_head_idx * stride_output_kv_head + group_idx * stride_output_group_size - + tid + + tid * fx.Int32(planned_elements_per_thread) ) - output[output_offset] = acc.to(output_dtype) + 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( diff --git a/aiter/ops/flydsl/pa_decode.py b/aiter/ops/flydsl/pa_decode.py index 9299f236963..4868928ddd4 100644 --- a/aiter/ops/flydsl/pa_decode.py +++ b/aiter/ops/flydsl/pa_decode.py @@ -141,6 +141,34 @@ def launch_pa_decode_ps_reduce( f"{MAX_CONTEXT_PARTITIONS} partitions" ) use_sinks = sink_token is not None + # Buffer offsets are byte-sized i32 values. Bound the largest attempted + # access, not only the current count: an inactive part must not wrap back + # into valid data. Nonstandard standalone reducer strides keep the old path. + bounded_plan_logits = ( + reduce_info is not None + 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 + ) + # The pair load is one aligned dword. Packed rows are complete D=128 + # vectors, and even strides preserve alignment after the 64-bit rebase. + # The R9 maximum-span gate also bounds the last byte of every pair load. + # Output remains scalar-stored, so its alignment need not match logits. + 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, @@ -151,6 +179,9 @@ def launch_pa_decode_ps_reduce( ), use_sinks=use_sinks, use_work_plan=reduce_info is not None, + 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)) @@ -220,9 +251,10 @@ def pa_decode( Sparse attention uses caller-prepared block tables and selected context lengths. Each independently selected MTP query must have its own table row and use query_length=1; query_length>1 applies dense causal masking. - A positive ``sliding_window`` limits each query to that many causal tokens, - including its own position; 0 and -1 disable the window. Optional ``sinks`` - is a contiguous [num_query_heads] BF16/FP16/FP32 tensor on the query device. + A positive ``sliding_window`` requires ``work_plan`` and limits each query + to that many causal tokens, including its own position; 0 and -1 disable + the window. Optional ``sinks`` is a contiguous [num_query_heads] + BF16/FP16/FP32 tensor on the query device. Each entry is an unscaled, zero-value attention logit shared across batch and MTP positions. It contributes to the denominator once, independently of the window; -inf disables a head's sink and +inf suppresses its output. @@ -260,6 +292,8 @@ def pa_decode( 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: @@ -592,6 +626,7 @@ def normalize_scale(scale, name): 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", diff --git a/docs/flydsl_pa_decode_plan.md b/docs/flydsl_pa_decode_plan.md deleted file mode 100644 index 2a44492a650..00000000000 --- a/docs/flydsl_pa_decode_plan.md +++ /dev/null @@ -1,148 +0,0 @@ -# GPU work plans for FlyDSL PA decode - -`plan_pa_decode` assigns different KV partition counts to requests according to -their GPU-resident context lengths. It is intended for batches with unequal work, -such as one long context mixed with many short contexts. It is opt-in: a planning -launch and packed reduction can cost more than static splits for uniform or short -contexts. Compare **plan + attention + reduction** against the static path. - -The attention math, per-token scaling, causal MTP masking and sliding windows are -shared with the static FlyDSL kernel. The planner is a Triton GPU kernel, following -the existing variable-context scheduling approach in -`mqa_logits/pa_mqa_logits_fp4.py`. - -A positive `sliding_window=W` retains at most W causal tokens **including the -query's own position**; W=1 attends only to that token. Values 0 and -1 disable the -window. Context lengths include the MTP query tokens: for context length C, query -length Q and query position p, the visible range is -`[max(0, C - Q + 1 + p - W), max(0, C - Q + 1 + p))`. - -## Scheduling and storage - -Each nonempty request first receives one task. Remaining task slots are apportioned -by its number of 256-token compute tiles using integer prefix sums. Counts are -capped by useful tiles and `max_partitions` (1 through 256). The workgroup budget is -a target across all KV heads; capacity is rounded up to at least one slot per -request/head. Clamping can leave unused slots. Physical page size is independent -of this scheduling tile size. - -With a window, only tiles intersecting the union of the Q query windows count -toward the allocation: from `max(0, C - (Q - 1) - W) // 256` through -`ceil(C / 256)` (exclusive). Unaligned windows can touch an extra tile. The -attention kernel applies the exact per-query mask within those tiles; it does -not load the skipped prefix tiles. - -The planner writes two contiguous int32 CUDA tensors: - -- `work_info[capacity, 4]`: request index, absolute first tile, exclusive absolute - last tile, original context length. -- `reduce_info[batch, 2]`: first packed task slot and actual partition count. - -The launch grid has a fixed capacity for graph capture. Padded work records are -zeroed and skipped by the attention kernel, leaving their scratch slots untouched. -Each real task processes one contiguous KV range; ranges cover each request's -active tiles exactly once. Reduction reads only the valid partitions and produces -zero for empty contexts. No GPU-to-CPU readback, dynamic allocation, or CPU length -inspection is required when refreshing a preallocated plan. - -Planned scratch is packed as `[num_kv_heads, capacity, query_rows]` for max/sum and -`[num_kv_heads, capacity, query_rows, head_dim]` for partial outputs, where -`query_rows = query_length * num_query_heads // num_kv_heads`. -The static API continues to use its dense per-request scratch layout. - -## Attention sinks - -Both static and planned `pa_decode` accept optional `sinks`: a contiguous -`[num_query_heads]` BF16, FP16 or FP32 tensor on the query device. Each entry is -a zero-value attention logit shared across sequences and MTP query positions. -It is not multiplied by `softmax_scale` and remains visible independently of -the sliding window. For head h, the output is mathematically -`sum(exp(score_i) * V_i) / (sum(exp(score_i)) + exp(sinks[h]))`. - -The sink participates in the stable softmax maximum and contributes to the -denominator exactly once, not once per partition. Static NP=1 handles it in -the compute epilogue without an extra launch; multi-partition and planned calls -add it in the final reducer. KV partials and FP8 Q/P quantization are unchanged. -Empty contexts or query rows with no visible KV tokens still produce zero. -`-inf` disables an individual head's sink; `+inf` produces zero for that head. - -Sinks do not change work ownership, so they require no planner arguments or new -metadata. Their values can be updated in place between graph replays without -rebuilding a plan, provided the context lengths remain unchanged. - -## Usage - -Use the Python FlyDSL entry point for planned execution. Its signature ends -with `sliding_window, work_plan` and does not accept `ps`. The Python wrapper -`aiter.ops.attention.pa_decode_flydsl` also accepts `work_plan` by keyword. -Its trailing arguments are `alibi_slopes, ps, sinks, sliding_window, work_plan`; -`ps` is accepted and ignored. Prefer keyword arguments for `sinks`, -`sliding_window`, and `work_plan` when switching between these entry points. - -```python -import torch -from aiter.ops.flydsl.pa_decode import pa_decode, plan_pa_decode - -# query/output: [batch * query_length, num_query_heads, head_dim] -# K/V cache layouts and scales are the same as for static pa_decode. -num_kv_heads = key_cache.shape[1] -sliding_window = 4096 # Use 0 (the default) or -1 for full causal attention. -plan = plan_pa_decode( - context_lengths, num_kv_heads, max_partitions=256, - sliding_window=sliding_window, query_length=query_length, -) -rows = query_length * query.shape[1] // num_kv_heads -shape = (num_kv_heads, plan.capacity, rows) -exp_sums = torch.empty(shape, dtype=torch.float32, device=query.device) -max_logits = torch.empty_like(exp_sums) -partials = torch.empty((*shape, query.shape[2]), dtype=query.dtype, device=query.device) -sinks = torch.zeros(query.shape[1], dtype=torch.float32, device=query.device) -# Use sinks=None for attention without the extra normalization logit. - -def step(): - # Refresh after changing lengths, on the same stream as attention. - plan_pa_decode( - context_lengths, num_kv_heads, max_partitions=256, plan=plan, - sliding_window=sliding_window, query_length=query_length, - ) - pa_decode( - output, query, key_cache, value_cache, context_lengths, block_tables, - softmax_scale, query_length, plan.max_partitions, - compute_type=key_cache.dtype, - key_scale=key_scale, value_scale=value_scale, - exp_sums=exp_sums, max_logits=max_logits, temporary_output=partials, - work_plan=plan, sliding_window=sliding_window, sinks=sinks, - ) - -step() # Warm up compilation before graph capture. -graph = torch.cuda.CUDAGraph() -with torch.cuda.graph(graph): - step() -# Update existing input buffers, then graph.replay(). -``` - -Reuse an already built plan without refreshing only when lengths are unchanged. -Changing batch size, KV-head count, the partition limit or launch capacity requires -a compatible new plan and scratch buffers. A plan's window is fixed, as is its -query length when the window is enabled; changing either requires a new plan. -Dense plans are independent of query length. The attention call checks these -settings to avoid silently using a plan that omits visible tokens. -Metadata values should be produced by the planner, not edited independently. -The planner accepts contiguous int32 CUDA -length vectors with 1 through 4096 requests. Lengths and block-table indices must -satisfy the same validity requirements as static attention. - -## Verification - -```bash -python -m pytest -q op_tests/test_flydsl_pa_decode.py -``` - -Tests cover exact absolute-tile ownership, sparse pages, empty requests, causal tails, -non-power-of-two partition caps, multiple KV heads, supported reducer branches, -NaN-filled scratch, and graph replay after changing lengths. Window tests include -W=1, non-aligned boundaries, windows larger than the context, QL1/2/3/4, per-token -and per-tensor scales, static/query-split execution and incompatible plan reuse. -Sink tests cover per-head logits, direct output and all reducer paths, mixed -input/sink dtypes, large and infinite logits, window composition, and graph -replay after updating sink values. diff --git a/op_tests/test_flydsl_pa_decode.py b/op_tests/test_flydsl_pa_decode.py index 425a6ea1513..02739473dc3 100644 --- a/op_tests/test_flydsl_pa_decode.py +++ b/op_tests/test_flydsl_pa_decode.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. -"""One parametrized correctness test and a CLI performance sweep for FlyDSL PA. +"""One parametrized correctness/contract test and a CLI performance sweep. 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 \ @@ -11,6 +11,9 @@ Both entry points share input generation, the FP32 reference and kernel launch. Context lengths include the MTP query tokens. Explicit partition counts override automatic recommendations; the CLI times attention plus the native reducer. +Positive sliding windows require a work plan: planned cases check numerics and +graph replays, while static cases check rejection by the core and wrapper APIs. +Disabled windows retain both static and planned numerical coverage. """ import argparse @@ -96,8 +99,6 @@ class DecodeCase: masked_scale: bool = False query_splits: int | None = None wide_kv_addressing: bool | None = None - expected_prefetch: bool | str | None = None - expected_query_splits: int | None = None check_schedule: bool = False @@ -478,21 +479,120 @@ def _assert_contracts(args, options): pa_decode(*args, **{**options, "work_plan": wrong_plan}) -def _assert_schedule_cache(compile_tile, kwargs, selected): - """Check deterministic CU boundaries without launching hypothetical grids.""" +def _max_window_tiles(window, query_length): + # The union has W + QL - 1 tokens and can start at any offset. + union_tokens = window + query_length - 1 + return 1 + (union_tokens + KV_COMPUTE_BLOCK - 2) // KV_COMPUTE_BLOCK + + +def _expected_single_tile_plan(kwargs): + capacity, window = kwargs["work_capacity"], kwargs["sliding_window"] + if not kwargs["use_work_plan"] or capacity is None or window <= 0: + return False + tiles = _max_window_tiles(window, kwargs["query_length"]) + return kwargs["num_partitions"] >= tiles and capacity >= kwargs["num_seqs"] * tiles + + +def _expected_batch_first_plan_grid(kwargs, query_splits): + """Independent shape/occupancy oracle for the measured grid domain.""" + return ( + get_gfx_runtime() == "gfx950" + and kwargs["head_dim"] == 128 + and kwargs["query_dtype"] == "bf16" + and kwargs["per_token_kv"] + and _expected_single_tile_plan(kwargs) + and kwargs["num_partitions"] <= 64 + and kwargs["query_length"] == query_splits == 4 + and kwargs["num_kv_heads"] == 1 + and kwargs["query_group_size"] == 16 + and kwargs["block_size"] == 128 + and 4096 <= kwargs["sliding_window"] <= 8192 + and 4 <= kwargs["num_seqs"] <= 24 + and kwargs["work_capacity"] + == kwargs["num_seqs"] + * _max_window_tiles(kwargs["sliding_window"], kwargs["query_length"]) + and kwargs["num_compute_units"] >= (kwargs["work_capacity"] + 1) // 2 + ) + + +def _assert_schedule_cache(compile_tile, kwargs): + """Check schedule equivalence through handles, without launching fake grids.""" workgroups = kwargs["num_seqs"] * kwargs["num_kv_heads"] * kwargs["num_partitions"] planned, ql = kwargs["use_work_plan"], kwargs["query_length"] cached = {} def check(cus, expected, **changes): - previous_count = len(selected) - compiled = compile_tile(**{**kwargs, "num_compute_units": cus, **changes}) - assert len(selected) == previous_count + 1 - actual = (selected[-1]["query_splits"], selected[-1]["prefetch_v"]) - assert actual == expected - previous = cached.setdefault(expected, compiled) - assert compiled["kernel"] is previous["kernel"] - assert compiled["launch"] is previous["launch"] + config = {**kwargs, "num_compute_units": cus, **changes} + compiled = compile_tile(**config) + expected = ( + *expected, + _expected_single_tile_plan(config), + _expected_batch_first_plan_grid(config, expected[0]), + ) + # Only these compile constants vary within this helper. Runtime sizes + # must share handles iff they select the same expected schedule. + compile_key = ( + config["num_partitions"], + config["query_length"], + config["trans_v"], + ) + schedules = cached.setdefault(compile_key, {}) + for schedule, previous in schedules.items(): + same_schedule = schedule == expected + assert (compiled is previous) == same_schedule + assert (compiled["kernel"] is previous["kernel"]) == same_schedule + assert (compiled["launch"] is previous["launch"]) == same_schedule + schedules.setdefault(expected, compiled) + + def split_boundary(capacity, **changes): + config = {**kwargs, **changes, "work_capacity": capacity} + tasks = config["num_seqs"] * config["num_partitions"] + if config["use_work_plan"]: + if capacity is not None: + tasks = capacity + if config["sliding_window"] > 0: + max_tiles = _max_window_tiles(config["sliding_window"], ql) + tasks = min(tasks, config["num_seqs"] * max_tiles) + queries_per_task = ql + if ( + _expected_single_tile_plan(config) + and ql == 4 + and config["block_size"] == 128 + and config["trans_v"] + ): + queries_per_task = 1 + return (queries_per_task * config["num_kv_heads"] * tasks + 1) // 2 + + def ql1_points(**changes): + config = {**kwargs, **changes} + dense = config["num_seqs"] * config["num_kv_heads"] * config["num_partitions"] + eligible = ( + config["head_dim"] == 128 + and config["query_dtype"] == "bf16" + and config["block_size"] in (16, 128) + and 8 <= config["query_group_size"] <= 16 + ) + # Express the independent expectation as the minimum CU count. A + # proven one-tile plan has at most B * Hkv * window_tiles workgroups, + # regardless of excess capacity or the nominal B * Hkv * NP grid. + threshold = None + if eligible: + threshold = ( + dense if config["block_size"] == 128 and config["trans_v"] else 0 + ) + if _expected_single_tile_plan(config): + tasks = ( + config["num_seqs"] + * config["num_kv_heads"] + * _max_window_tiles(config["sliding_window"], 1) + ) + threshold = min(threshold, (tasks + 1) // 2) + cus = {max(1, dense - 1), dense, dense + 1} + if threshold is not None: + cus.update((max(1, threshold - 1), max(1, threshold), threshold + 1)) + return [ + (cu, (1, threshold is not None and cu >= threshold)) for cu in sorted(cus) + ] if not kwargs["per_token_kv"]: # Scalar prefetch requires CU < workgroups <= 2 * CU. @@ -504,33 +604,132 @@ def check(cus, expected, **changes): (half - 1, (1, False)), ] elif ql == 1: - # Transposed page128 prefetch stops above one workgroup per CU. - points = [ - (workgroups - 1, (1, False)), - (workgroups, (1, not planned)), - (workgroups + 1, (1, not planned)), - ] + points = ql1_points() else: - # The selected MTP2 case has an exact QL * workgroups == 2 * CU boundary. - boundary = ql * workgroups // 2 - split = (1, False) if planned else (ql, True) + # Ordinary grids count QL CTAs per task; one-tile QL4/page128/transV + # plans permit two tasks per CU. Round up for an odd workgroup count. + boundary = split_boundary(kwargs["work_capacity"]) + split = (ql, True) points = [(boundary - 1, (1, False)), (boundary, split), (boundary + 1, split)] for cus, expected in points: check(cus, expected) if kwargs["per_token_kv"] and ql == 1: - # Hkv1 remains tuned for plans; Hkv2 is static-only; Hkv>2 is untuned. - for heads, prefetch in [(1, True), (2, not planned), (3, False)]: - cus = kwargs["num_seqs"] * heads * kwargs["num_partitions"] - check(cus, (1, prefetch), num_kv_heads=heads) - elif kwargs["per_token_kv"] and not planned: + # Head counts affect occupancy, not single-query M1 eligibility. + for heads in (1, 2, 3): + for cus, expected in ql1_points(num_kv_heads=heads): + check(cus, expected, num_kv_heads=heads) + elif kwargs["per_token_kv"]: # Explicit split overrides must also select matching prefetching. check(boundary - 1, (ql, True), query_splits=ql) check(boundary + 1, (1, False), query_splits=1) + # Capacity selects splitting and the semantic one-tile guarantee, not a + # raw cache key. The guarantee uses full QL even when queries are split. + capacities = [ + None, + kwargs["num_seqs"], + kwargs["num_seqs"] * kwargs["num_partitions"], + ] + if kwargs["sliding_window"] > 0: + full_window = kwargs["num_seqs"] * _max_window_tiles( + kwargs["sliding_window"], ql + ) + capacities += [full_window - 1, full_window, full_window + 1] + for capacity in dict.fromkeys(capacities): + capacity_points = points + if kwargs["per_token_kv"] and ql > 1: + boundary = split_boundary(capacity) + capacity_points = [(boundary - 1, (1, False)), (boundary, (ql, True))] + elif kwargs["per_token_kv"] and ql == 1: + capacity_points = ql1_points(work_capacity=capacity) + for cus, expected in capacity_points: + check(cus, expected, work_capacity=capacity) + + if kwargs["sliding_window"] > 0 and kwargs["per_token_kv"] and ql == 1: + # Isolate NP's single-tile guarantee from the already sufficient + # capacity; hypothetical metadata is compiled but never launched. + tiles = _max_window_tiles(kwargs["sliding_window"], 1) + for parts in dict.fromkeys((max(1, tiles - 1), tiles)): + for cus, expected in ql1_points( + num_partitions=parts, work_capacity=full_window + ): + check(cus, expected, num_partitions=parts, work_capacity=full_window) + + if kwargs["sliding_window"] > 0 and kwargs["per_token_kv"] and ql > 1: + # Hypothetical metadata isolates the NP guard from the capacity guard. + # Real plans cannot exceed B * NP; these compile-only grids never launch. + tiles = _max_window_tiles(kwargs["sliding_window"], ql) + for parts in (tiles - 1, tiles, tiles + 1): + boundary = split_boundary(full_window, num_partitions=parts) + for cus, expected in [(boundary - 1, (1, False)), (boundary, (ql, True))]: + check( + cus, + expected, + num_partitions=parts, + work_capacity=full_window, + ) + if planned and kwargs["block_size"] == 128 and kwargs["trans_v"]: + # Enough capacity alone must not widen plain-V or QL2 thresholds. + wider_boundary = (full_window * kwargs["num_kv_heads"] + 1) // 2 + check( + wider_boundary, + (1, False), + trans_v=False, + work_capacity=full_window, + ) + check( + wider_boundary, + (1, False), + query_length=2, + work_capacity=full_window, + ) + + if ( + planned + and kwargs["per_token_kv"] + and ql == 4 + and kwargs["num_kv_heads"] == 1 + and kwargs["block_size"] == 128 + and kwargs["sliding_window"] > 0 + ): + # Probe both batch endpoints and just outside them for every window. + # Out-of-range windows must stay flat even for an in-range batch. + # Plain V retains the full QL split threshold in either address width. + # These hypothetical grids compile but never launch against this plan. + tiles = _max_window_tiles(kwargs["sliding_window"], ql) + for batch in (3, 4, 5, 12, 24, 25): + capacity = batch * tiles + boundary = split_boundary(capacity, num_seqs=batch, num_partitions=tiles) + for cus in (boundary, boundary + 1): + check( + cus, + (4, True), + num_seqs=batch, + num_partitions=tiles, + work_capacity=capacity, + ) + # NP is itself a compile constant. Compare B4/B25 within each cap: + # only in-range windows with NP64 separate handles; NP65 stays flat. + for parts in (64, 65): + for batch in (4, 25): + capacity = batch * tiles + boundary = split_boundary( + capacity, num_seqs=batch, num_partitions=parts + ) + check( + boundary, + (4, True), + num_seqs=batch, + num_partitions=parts, + work_capacity=capacity, + ) + # Runtime size/CU metadata must not fragment the kernel specialization cache. cus, expected = points[-1] changes = {"num_seqs": kwargs["num_seqs"] * 2} + if planned and kwargs["work_capacity"] is not None: + changes["work_capacity"] = kwargs["work_capacity"] * 2 multiplier = 2 if not kwargs["per_token_kv"]: changes["num_kv_heads"] = kwargs["num_kv_heads"] * 2 @@ -568,7 +767,8 @@ def _case( ) -# Each row uses the same eager/graph checks with static and packed plans. +# Positive-window rows reject static calls and run eager/graph checks with plans. +# Disabled-window rows retain the same numerical checks in both modes. BF16, FP16, FP32 = torch.bfloat16, torch.float16, torch.float32 CASES = [ _case("scalar-direct", (1, 2, 8, 128), (16, 1, 0), 1), @@ -592,10 +792,43 @@ def _case( ), _case("head64-fp16", (2, 2, 4, 64), (16, 0, 0), 7, 257, BF16, dtype=FP16), _case("head1024", (1, 2, 4, 1024), (128, 1, 0), 256, 8192, FP32), + # QL4/W1 can span two tiles, including when explicitly split into four CTAs. + # NP1 therefore retains the multi-tile fallback even with a full plan budget. _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), + parts=34, + window=4096, + sink=BF16, + lengths=(1, 257, 4353, 0), + query_splits=4, + ), + _case( + "full-m1-fp16-d256", + (2, 2, 8, 256), + (64, 1, 1), + parts=5, + window=257, + sink=FP16, + dtype=FP16, + lengths=(0, 1, 257, 513), + query_splits=1, + ), _case("mtp3-window", (3, 1, 16, 128), (128, 0, 1), 7, 257, FP16), _case("mtp2-odd-parts", (2, 2, 16, 128), (16, 0, 1), 3), _case( @@ -603,10 +836,369 @@ def _case( (2, 1, 16, 128), parts=3, lengths=(257, 259), - expected_query_splits=2, - expected_prefetch=True, check_schedule=True, ), + _case( + "mtp4-window-query-split", + cache=(16, 1, 1), + window=1024, + lengths=(1023, 1024, 1025, 4099), + check_schedule=True, + ), + _case( + "mtp4-window-capacity", + parts=256, + window=1024, + lengths=(0, 3, 1024, 4099), + check_schedule=True, + ), + _case( + "mtp4-dense-capacity", + cache=(128, 0, 1), + parts=256, + lengths=(257, 259, 1027, 4099), + workgroup_budget=17, + check_schedule=True, + ), + _case( + "mtp4-window-large-grid", + window=1024, + lengths=(1027,) * 200, + check_schedule=True, + ), + # Exact one-tile budgets enter the measured (B, H*QS, C/B) grid on gfx950. + # Ragged/empty owners are not physical x; each row also runs poisoned + # scratch, changing sinks, minus7/all-empty/restored graph plan refreshes. + _case( + "batch-first-window4096", + parts=18, + window=4096, + sink=FP32, + lengths=(0, 1, 4097, 4353), + workgroup_budget=72, + check_schedule=True, + ), + _case( + "batch-first-window8192", + parts=34, + window=8192, + sink=BF16, + lengths=(0, 1, 3, 4, 257, 8193, 8449, 0), + workgroup_budget=272, + check_schedule=True, + ), + _case( + "batch-first-window8192-batch12", + parts=34, + window=8192, + sink=FP32, + lengths=(0, 1, 8193, 8449) * 3, + workgroup_budget=408, + check_schedule=True, + ), + # Interior window/odd batch: C125 reaches the QS4 threshold at CU63. + _case( + "batch-first-window6000-batch5", + parts=25, + window=6000, + sink=FP32, + lengths=(0, 1, 6001, 6145, 6257), + workgroup_budget=125, + check_schedule=True, + ), + # These plans satisfy the tile/capacity guards but miss the window range. + _case( + "batch-first-window4095-fallback", + parts=18, + window=4095, + sink=FP32, + lengths=(0, 1, 4097, 4353), + workgroup_budget=72, + check_schedule=True, + ), + _case( + "batch-first-window8193-fallback", + parts=34, + window=8193, + sink=FP32, + lengths=(0, 1, 8193, 8449), + workgroup_budget=136, + check_schedule=True, + ), + _case( + "batch-first-window1024-rejected", + parts=6, + window=1024, + sink=FP32, + lengths=(0, 1, 1025, 1281) * 2, + workgroup_budget=48, + check_schedule=True, + ), + _case( + "batch-first-excess-partitions", + parts=35, + window=8192, + sink=FP32, + lengths=(0, 1, 8193, 8449), + workgroup_budget=140, + check_schedule=True, + ), + _case( + "batch-first-excess-capacity", + parts=35, + window=8192, + sink=BF16, + lengths=(0, 1, 8193, 8449), + workgroup_budget=137, + check_schedule=True, + ), + _case( + "batch-first-tight-capacity", + parts=34, + window=8192, + sink=FP32, + lengths=(0, 1, 8193, 8449), + workgroup_budget=135, + check_schedule=True, + ), + _case( + "batch-first-window4096-wide", + parts=18, + window=4096, + sink=FP32, + lengths=(0, 1, 4097, 4353), + workgroup_budget=72, + wide_kv_addressing=True, + check_schedule=True, + ), + # Plain V needs CU >= QL * C / 2 = 144 to split this B4/C72 plan. + # Exercise both address widths and a larger partition cap with real inputs. + _case( + "batch-first-window4096-plain-v", + cache=(128, 0, 1), + parts=18, + window=4096, + sink=FP32, + lengths=(0, 1, 4097, 4353), + workgroup_budget=72, + check_schedule=True, + ), + _case( + "batch-first-window4096-plain-v-wide", + cache=(128, 0, 1), + parts=18, + window=4096, + sink=FP32, + lengths=(0, 1, 4097, 4353), + workgroup_budget=72, + wide_kv_addressing=True, + check_schedule=True, + ), + _case( + "batch-first-window4096-np64-plain-v-wide", + cache=(128, 0, 1), + parts=64, + window=4096, + sink=FP32, + lengths=(0, 1, 4097, 4353), + workgroup_budget=72, + wide_kv_addressing=True, + check_schedule=True, + ), + # The unchanged plain-V split threshold is CU272: CU256 uses QS1/flat. + _case( + "batch-first-window8192-plain-v-auto", + cache=(128, 0, 1), + parts=34, + window=8192, + sink=FP32, + lengths=(0, 1, 8193, 8449), + workgroup_budget=136, + check_schedule=True, + ), + _case( + "batch-first-window4096-batch24", + parts=18, + window=4096, + sink=FP32, + lengths=(0, 1, 3, 4, 257, 1025, 2049, 4096, 4097, 4098, 4353, 4354) * 2, + workgroup_budget=432, + check_schedule=True, + ), + # NP need not equal the window bound: exact useful capacity still proves + # one task per tile. Counts span 0/1/4/8/16/17/18 and poisoned padding. + _case( + "batch-first-window4096-np64", + parts=64, + window=4096, + sink=FP32, + lengths=(0, 1, 769, 1793, 4096, 4097, 4353, 4354), + workgroup_budget=144, + check_schedule=True, + ), + # Large NP keeps the flat grid. The 34-task rows still cross the NP256 + # reducer's 32-partition group boundary, including on graph replays. + _case( + "batch-first-window8192-np256-fallback", + parts=256, + window=8192, + sink=FP32, + lengths=(0, 1, 8449, 8450), + workgroup_budget=136, + check_schedule=True, + ), + _case( + "ql1-window-prefetch", + (1, 1, 16, 128), + parts=64, + window=8192, + lengths=(8193,) * 8, + workgroup_budget=512, + check_schedule=True, + ), + _case( + "ql1-window-prefetch-small-capacity", + (1, 1, 16, 128), + parts=64, + window=8192, + lengths=(8193,) * 8, + workgroup_budget=263, + check_schedule=True, + ), + _case( + "ql1-window-prefetch-wide", + (1, 1, 16, 128), + parts=64, + window=8192, + lengths=(8193,) * 8, + workgroup_budget=512, + wide_kv_addressing=True, + check_schedule=True, + ), + # G8 pads half an M-tile. Distinct per-head data/scales also cover planned + # Hkv2; the Hkv1/NP256 row isolates G8 with a large nominal dense grid. + _case( + "ql1-hkv2-g8-window-prefetch", + (1, 2, 8, 128), + parts=64, + window=8192, + sink=FP32, + lengths=(0, 1, 8193, 8449), + workgroup_budget=264, + check_schedule=True, + ), + _case( + "ql1-hkv1-g8-large-np-prefetch", + (1, 1, 8, 128), + parts=256, + window=8192, + sink=FP32, + lengths=(0, 1, 8193, 8449), + workgroup_budget=132, + check_schedule=True, + ), + # Hkv4 exercises dense M1 prefetch with multi-tile tasks in each KV layout. + _case( + "ql1-hkv4-page128-multi", + (1, 4, 16, 128), + parts=2, + window=1024, + sink=FP32, + lengths=(0, 1, 1025, 1281), + workgroup_budget=32, + check_schedule=True, + ), + _case( + "ql1-hkv4-page16-multi", + (1, 4, 16, 128), + (16, 1, 1), + parts=2, + window=1024, + sink=FP32, + lengths=(0, 1, 1025, 1281), + workgroup_budget=32, + check_schedule=True, + ), + _case( + "ql1-hkv4-plain-v-multi", + (1, 4, 16, 128), + (128, 0, 1), + parts=2, + window=1024, + sink=FP32, + lengths=(0, 1, 1025, 1281), + workgroup_budget=32, + check_schedule=True, + ), + _case( + "ql1-hkv4-page16-direct-sinks", + (1, 4, 16, 128), + (16, 1, 1), + parts=1, + sink=FP32, + lengths=(0, 1, 257, 769), + workgroup_budget=16, + check_schedule=True, + ), + # Interior GQA sizes pad to one M-tile; G7/G17 exercise both range misses. + # Fixed groups per fixture keep the handle-only cache checks unambiguous. + *[ + _case( + f"ql1-hkv2-g{group}-window", + (1, 2, group, 128), + parts=8, + window=1024, + sink=FP32, + lengths=(0, 1, 1025, 1281), + workgroup_budget=40, + check_schedule=True, + ) + for group in (7, *range(9, 16), 17) + ], + _case( + "ql1-hkv2-g9-page16-direct-sinks", + (1, 2, 9, 128), + (16, 1, 1), + parts=1, + sink=FP32, + lengths=(0, 1, 257, 769), + workgroup_budget=8, + check_schedule=True, + ), + _case( + "ql1-hkv2-g15-plain-v-multi", + (1, 2, 15, 128), + (128, 0, 1), + parts=2, + window=1024, + sink=FP32, + lengths=(0, 1, 1025, 1281), + workgroup_budget=16, + check_schedule=True, + ), + # Explicit splits also use one padded M-tile for interior GQA sizes. + _case( + "split2-hkv2-g9-page16-direct-sinks", + (2, 2, 9, 128), + (16, 1, 1), + parts=1, + sink=FP32, + lengths=(0, 1, 257, 769), + workgroup_budget=8, + query_splits=2, + ), + _case( + "split4-hkv1-g15-plain-v-window-wide", + (4, 1, 15, 128), + (128, 0, 1), + parts=8, + window=1024, + sink=FP32, + lengths=(0, 1, 1025, 1281), + workgroup_budget=24, + query_splits=4, + wide_kv_addressing=True, + ), _case( "hkv2-prefetch-1wg", (1, 2, 8, 128), @@ -614,7 +1206,6 @@ def _case( window=1, sink=FP32, lengths=(257,) * 16, - expected_prefetch="occupancy", check_schedule=True, ), _case( @@ -622,7 +1213,6 @@ def _case( (1, 2, 8, 128), parts=8, lengths=(257,) * 32, - expected_prefetch="occupancy", ), _case( "hkv2-prefetch-page16", @@ -632,7 +1222,6 @@ def _case( 257, FP16, lengths=(257,) * 32, - expected_prefetch=True, ), _case( "hkv2-prefetch-plain-v", @@ -641,8 +1230,8 @@ def _case( 8, sink=BF16, lengths=(257,) * 32, - expected_prefetch=True, ), + # QL1/W1 has M=1: NP=M and capacity=B*M both meet the exact boundary. _case("hkv2-direct-sinks", (1, 2, 8, 128), parts=1, window=1, sink=BF16), _case( "long-200k", @@ -652,6 +1241,7 @@ def _case( lengths=(0, 1, 3, 4, 255, 256, 257, 200003), ), _case("window-int64", cache=(16, 1, 1), window=2**40, wide_kv_addressing=True), + # The single active tile is absolute tile 1; padding stays NaN across replay. _case( "masked-scale-decode", (1, 1, 16, 128), @@ -675,7 +1265,6 @@ def _case( cache=(128, 0, 1), parts=None, lengths=(257,) * 200, - expected_query_splits=1, ), _case("small-batch-auto", (1, 1, 8, 128), (16, 0, 0), None, lengths=(257,) * 3), _case( @@ -699,18 +1288,72 @@ def _case( _case("window8192", (2, 1, 16, 128), parts=256, window=8192), _case("window-int32-max", cache=(16, 0, 1), parts=1, window=2**31 - 1), _case("window-int32-overflow", cache=(16, 0, 1), parts=1, window=2**31, sink=FP32), + # QL3/G8 pads 24 query rows to 32; C 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, query, _, _, context = args[:5] scratch, sinks, plan = args[14:17], args[-1], options["work_plan"] + if plan is not None and case.workgroup_budget is not None: + budget_slots = ( + case.workgroup_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") build_tile = module.compile_pa_decode_tile + compile_reduce = module.compile_pa_decode_ps_reduce num_compute_units = torch.cuda.get_device_properties( query.device ).multi_processor_count - selected = [] compile_args = None + cached_reducer = None def compile_tile(**kwargs): - # Observe the builder's actual schedule, including cached specializations. + # Schedule details remain internal; callers receive only launch handles. + assert ( + not { + "single_tile_plan", + "batch_first_plan_grid", + "buffer_plan_output", + } + & kwargs.keys() + ) compiled = build_tile(**kwargs) - selected.append(compiled) + assert set(compiled) == {"launch", "kernel"} return compiled def compile_checked(**kwargs): @@ -755,7 +1420,8 @@ def compile_checked(**kwargs): assert kwargs["num_seqs"] == context.numel() assert kwargs["num_kv_heads"] == case.num_kv_heads assert kwargs["num_compute_units"] == num_compute_units - if case.query_splits is not None and not planned: + assert kwargs["work_capacity"] == (plan.capacity if planned else None) + if case.query_splits is not None: kwargs["query_splits"] = case.query_splits if case.wide_kv_addressing is not None: kwargs["wide_kv_addressing"] = case.wide_kv_addressing @@ -763,24 +1429,50 @@ def compile_checked(**kwargs): sinks is not None and args[8] == 1 and not planned ) compile_args = kwargs - previous_count = len(selected) - compiled = compile_tile(**kwargs) - assert len(selected) == previous_count + 1 - if case.query_splits is not None and not planned: - assert selected[-1]["query_splits"] == case.query_splits + return compile_tile(**kwargs) + + def compile_reduce_checked(**kwargs): + nonlocal cached_reducer + assert kwargs["query_group_size"] == case.query_group_size + assert kwargs["bounded_plan_logits"] == (planned and args[8] <= 64) + assert kwargs["vectorize_plan_logits"] == ( + planned + and args[8] <= 64 + and case.head_dim == 128 + and case.dtype in (BF16, FP16) + and scratch[2].data_ptr() % 4 == 0 + and all(stride % 2 == 0 for stride in scratch[2].stride()[:-1]) + ) + compiled = compile_reduce(**kwargs) + if cached_reducer is None: + cached_reducer = compiled + assert compiled is cached_reducer return compiled def unexpected_reducer(*_args, **_kwargs): raise AssertionError("static NP=1 must not launch a reducer") monkeypatch.setattr(module, "compile_pa_decode_tile", compile_checked) + monkeypatch.setattr(module, "compile_pa_decode_ps_reduce", compile_reduce_checked) if not planned and args[8] == 1: 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()) + if _expected_single_tile_plan(compile_args): + records = plan.work_info[:active] + assert (records[:, 2] - records[:, 1] == 1).all() for tensor in scratch: assert torch.isnan(tensor[:, active:]).all() elif args[8] == 1: @@ -810,29 +1502,9 @@ def check(disabled_sink=False): graph.replay() check(disabled_sink=step == 2) - if case.expected_prefetch is not None: - expected = case.expected_prefetch - if expected == "occupancy": - workgroups = len(case.lengths) * case.num_kv_heads * args[8] - expected = ( - workgroups - <= torch.cuda.get_device_properties(query.device).multi_processor_count - ) - expected = bool(expected) and not planned and get_gfx_runtime() == "gfx950" - assert selected and all(config["prefetch_v"] == expected for config in selected) - if case.expected_query_splits is not None: - expected = ( - case.expected_query_splits - if not planned and get_gfx_runtime() == "gfx950" - else 1 - ) - assert selected and all( - config["query_splits"] == expected for config in selected - ) - if case.check_schedule and get_gfx_runtime() == "gfx950": # Compile-only checks stay outside capture and never launch fake-CU results. - _assert_schedule_cache(compile_tile, compile_args, selected) + _assert_schedule_cache(compile_tile, compile_args) _assert_contracts(args, options) if plan is not None: From c5911aaa9633bd4320e1068c13f2e85a5d8247c7 Mon Sep 17 00:00:00 2001 From: fsx950223 Date: Fri, 18 Sep 2026 10:33:48 +0000 Subject: [PATCH 59/66] style(flydsl): fix PA decode Black formatting Match Black 26.5.1 used by the Checks workflow. Only assertion layout and line wrapping change; the Python AST is unchanged. Validation: Black --check --diff across the worktree, Ruff check for the changed file, and git diff --check passed. --- aiter/ops/flydsl/kernels/pa_decode_kernel.py | 31 +++++++++----------- 1 file changed, 14 insertions(+), 17 deletions(-) diff --git a/aiter/ops/flydsl/kernels/pa_decode_kernel.py b/aiter/ops/flydsl/kernels/pa_decode_kernel.py index 3f7bb28c14a..32628e0d0eb 100644 --- a/aiter/ops/flydsl/kernels/pa_decode_kernel.py +++ b/aiter/ops/flydsl/kernels/pa_decode_kernel.py @@ -283,9 +283,9 @@ def compile_pa_decode_tile( 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 ( + head_dim % MFMA_MNK == 0 + ), f"head_dim {head_dim} must be a multiple of {MFMA_MNK}" assert block_size in ( 16, 64, @@ -360,25 +360,25 @@ def compile_pa_decode_tile( 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" - ) + 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" - ) + assert ( + head_dim % (NWARP * MFMA_MNK) == 0 + ), "head_dim must split across the 4 warps for PV" # head_dim splits into 16-element chunks (QK_CHUNK_ELEMS, one dwordx4 load); # RGROUP_QUARTERS of them make a 64-element fetch group, QKHE_LOOP groups total. 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}" - ) + 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" @@ -1172,8 +1172,7 @@ def _v_ops(phys_row, vh): # Subtract from tile_valid inside the loop, keeping bounds tile-relative. if const_expr(QUERIES_PER_CTA == 1): causal_offset = [ - query_length - 1 - query_begin - for _m in range_constexpr(M_TILES) + query_length - 1 - query_begin for _m in range_constexpr(M_TILES) ] else: causal_offset = [ @@ -1276,9 +1275,7 @@ def _l_slot(m): # Scales outside every query's window must not shrink # visible probabilities to zero during FP8 normalization. # Keep this bound query-independent for shared MTP scales. - first_visible = (window_left - (query_length - 1)).to( - fx.Float32 - ) + first_visible = (window_left - (query_length - 1)).to(fx.Float32) window_scale_thr = fx.Vector.from_elements( [ first_visible From aa7ac40afd6c967c3b479ff6b77b1912322b4aaa Mon Sep 17 00:00:00 2001 From: fsx950223 Date: Fri, 18 Sep 2026 10:46:46 +0000 Subject: [PATCH 60/66] test(flydsl): consolidate PA decode unit cases Generate the existing 88 cases from grouped parameter tables and keep one shared correctness/contract flow for static and planned decode. Remove fake-CU selector/cache mirrors and retain only explicit query split/address-width overrides. Preserve the reference, plan oracle, sinks/empty-row checks, poisoned scratch, graph replays and CLI behavior. Validation: all 176 pytest cases passed on dzm-mi355-gpu34-dev (gfx950); the CLI help smoke test, Black 26.5.1 and Ruff checks passed. Independent CPU checks confirmed all case parameters and common helpers are unchanged. --- op_tests/test_flydsl_pa_decode.py | 1026 +++++++---------------------- 1 file changed, 232 insertions(+), 794 deletions(-) diff --git a/op_tests/test_flydsl_pa_decode.py b/op_tests/test_flydsl_pa_decode.py index 02739473dc3..98d8942148a 100644 --- a/op_tests/test_flydsl_pa_decode.py +++ b/op_tests/test_flydsl_pa_decode.py @@ -99,7 +99,6 @@ class DecodeCase: masked_scale: bool = False query_splits: int | None = None wide_kv_addressing: bool | None = None - check_schedule: bool = False def _require_gpu(): @@ -479,264 +478,6 @@ def _assert_contracts(args, options): pa_decode(*args, **{**options, "work_plan": wrong_plan}) -def _max_window_tiles(window, query_length): - # The union has W + QL - 1 tokens and can start at any offset. - union_tokens = window + query_length - 1 - return 1 + (union_tokens + KV_COMPUTE_BLOCK - 2) // KV_COMPUTE_BLOCK - - -def _expected_single_tile_plan(kwargs): - capacity, window = kwargs["work_capacity"], kwargs["sliding_window"] - if not kwargs["use_work_plan"] or capacity is None or window <= 0: - return False - tiles = _max_window_tiles(window, kwargs["query_length"]) - return kwargs["num_partitions"] >= tiles and capacity >= kwargs["num_seqs"] * tiles - - -def _expected_batch_first_plan_grid(kwargs, query_splits): - """Independent shape/occupancy oracle for the measured grid domain.""" - return ( - get_gfx_runtime() == "gfx950" - and kwargs["head_dim"] == 128 - and kwargs["query_dtype"] == "bf16" - and kwargs["per_token_kv"] - and _expected_single_tile_plan(kwargs) - and kwargs["num_partitions"] <= 64 - and kwargs["query_length"] == query_splits == 4 - and kwargs["num_kv_heads"] == 1 - and kwargs["query_group_size"] == 16 - and kwargs["block_size"] == 128 - and 4096 <= kwargs["sliding_window"] <= 8192 - and 4 <= kwargs["num_seqs"] <= 24 - and kwargs["work_capacity"] - == kwargs["num_seqs"] - * _max_window_tiles(kwargs["sliding_window"], kwargs["query_length"]) - and kwargs["num_compute_units"] >= (kwargs["work_capacity"] + 1) // 2 - ) - - -def _assert_schedule_cache(compile_tile, kwargs): - """Check schedule equivalence through handles, without launching fake grids.""" - workgroups = kwargs["num_seqs"] * kwargs["num_kv_heads"] * kwargs["num_partitions"] - planned, ql = kwargs["use_work_plan"], kwargs["query_length"] - cached = {} - - def check(cus, expected, **changes): - config = {**kwargs, "num_compute_units": cus, **changes} - compiled = compile_tile(**config) - expected = ( - *expected, - _expected_single_tile_plan(config), - _expected_batch_first_plan_grid(config, expected[0]), - ) - # Only these compile constants vary within this helper. Runtime sizes - # must share handles iff they select the same expected schedule. - compile_key = ( - config["num_partitions"], - config["query_length"], - config["trans_v"], - ) - schedules = cached.setdefault(compile_key, {}) - for schedule, previous in schedules.items(): - same_schedule = schedule == expected - assert (compiled is previous) == same_schedule - assert (compiled["kernel"] is previous["kernel"]) == same_schedule - assert (compiled["launch"] is previous["launch"]) == same_schedule - schedules.setdefault(expected, compiled) - - def split_boundary(capacity, **changes): - config = {**kwargs, **changes, "work_capacity": capacity} - tasks = config["num_seqs"] * config["num_partitions"] - if config["use_work_plan"]: - if capacity is not None: - tasks = capacity - if config["sliding_window"] > 0: - max_tiles = _max_window_tiles(config["sliding_window"], ql) - tasks = min(tasks, config["num_seqs"] * max_tiles) - queries_per_task = ql - if ( - _expected_single_tile_plan(config) - and ql == 4 - and config["block_size"] == 128 - and config["trans_v"] - ): - queries_per_task = 1 - return (queries_per_task * config["num_kv_heads"] * tasks + 1) // 2 - - def ql1_points(**changes): - config = {**kwargs, **changes} - dense = config["num_seqs"] * config["num_kv_heads"] * config["num_partitions"] - eligible = ( - config["head_dim"] == 128 - and config["query_dtype"] == "bf16" - and config["block_size"] in (16, 128) - and 8 <= config["query_group_size"] <= 16 - ) - # Express the independent expectation as the minimum CU count. A - # proven one-tile plan has at most B * Hkv * window_tiles workgroups, - # regardless of excess capacity or the nominal B * Hkv * NP grid. - threshold = None - if eligible: - threshold = ( - dense if config["block_size"] == 128 and config["trans_v"] else 0 - ) - if _expected_single_tile_plan(config): - tasks = ( - config["num_seqs"] - * config["num_kv_heads"] - * _max_window_tiles(config["sliding_window"], 1) - ) - threshold = min(threshold, (tasks + 1) // 2) - cus = {max(1, dense - 1), dense, dense + 1} - if threshold is not None: - cus.update((max(1, threshold - 1), max(1, threshold), threshold + 1)) - return [ - (cu, (1, threshold is not None and cu >= threshold)) for cu in sorted(cus) - ] - - if not kwargs["per_token_kv"]: - # Scalar prefetch requires CU < workgroups <= 2 * CU. - half = (workgroups + 1) // 2 - points = [ - (workgroups, (1, False)), - (workgroups - 1, (1, True)), - (half, (1, True)), - (half - 1, (1, False)), - ] - elif ql == 1: - points = ql1_points() - else: - # Ordinary grids count QL CTAs per task; one-tile QL4/page128/transV - # plans permit two tasks per CU. Round up for an odd workgroup count. - boundary = split_boundary(kwargs["work_capacity"]) - split = (ql, True) - points = [(boundary - 1, (1, False)), (boundary, split), (boundary + 1, split)] - for cus, expected in points: - check(cus, expected) - - if kwargs["per_token_kv"] and ql == 1: - # Head counts affect occupancy, not single-query M1 eligibility. - for heads in (1, 2, 3): - for cus, expected in ql1_points(num_kv_heads=heads): - check(cus, expected, num_kv_heads=heads) - elif kwargs["per_token_kv"]: - # Explicit split overrides must also select matching prefetching. - check(boundary - 1, (ql, True), query_splits=ql) - check(boundary + 1, (1, False), query_splits=1) - - # Capacity selects splitting and the semantic one-tile guarantee, not a - # raw cache key. The guarantee uses full QL even when queries are split. - capacities = [ - None, - kwargs["num_seqs"], - kwargs["num_seqs"] * kwargs["num_partitions"], - ] - if kwargs["sliding_window"] > 0: - full_window = kwargs["num_seqs"] * _max_window_tiles( - kwargs["sliding_window"], ql - ) - capacities += [full_window - 1, full_window, full_window + 1] - for capacity in dict.fromkeys(capacities): - capacity_points = points - if kwargs["per_token_kv"] and ql > 1: - boundary = split_boundary(capacity) - capacity_points = [(boundary - 1, (1, False)), (boundary, (ql, True))] - elif kwargs["per_token_kv"] and ql == 1: - capacity_points = ql1_points(work_capacity=capacity) - for cus, expected in capacity_points: - check(cus, expected, work_capacity=capacity) - - if kwargs["sliding_window"] > 0 and kwargs["per_token_kv"] and ql == 1: - # Isolate NP's single-tile guarantee from the already sufficient - # capacity; hypothetical metadata is compiled but never launched. - tiles = _max_window_tiles(kwargs["sliding_window"], 1) - for parts in dict.fromkeys((max(1, tiles - 1), tiles)): - for cus, expected in ql1_points( - num_partitions=parts, work_capacity=full_window - ): - check(cus, expected, num_partitions=parts, work_capacity=full_window) - - if kwargs["sliding_window"] > 0 and kwargs["per_token_kv"] and ql > 1: - # Hypothetical metadata isolates the NP guard from the capacity guard. - # Real plans cannot exceed B * NP; these compile-only grids never launch. - tiles = _max_window_tiles(kwargs["sliding_window"], ql) - for parts in (tiles - 1, tiles, tiles + 1): - boundary = split_boundary(full_window, num_partitions=parts) - for cus, expected in [(boundary - 1, (1, False)), (boundary, (ql, True))]: - check( - cus, - expected, - num_partitions=parts, - work_capacity=full_window, - ) - if planned and kwargs["block_size"] == 128 and kwargs["trans_v"]: - # Enough capacity alone must not widen plain-V or QL2 thresholds. - wider_boundary = (full_window * kwargs["num_kv_heads"] + 1) // 2 - check( - wider_boundary, - (1, False), - trans_v=False, - work_capacity=full_window, - ) - check( - wider_boundary, - (1, False), - query_length=2, - work_capacity=full_window, - ) - - if ( - planned - and kwargs["per_token_kv"] - and ql == 4 - and kwargs["num_kv_heads"] == 1 - and kwargs["block_size"] == 128 - and kwargs["sliding_window"] > 0 - ): - # Probe both batch endpoints and just outside them for every window. - # Out-of-range windows must stay flat even for an in-range batch. - # Plain V retains the full QL split threshold in either address width. - # These hypothetical grids compile but never launch against this plan. - tiles = _max_window_tiles(kwargs["sliding_window"], ql) - for batch in (3, 4, 5, 12, 24, 25): - capacity = batch * tiles - boundary = split_boundary(capacity, num_seqs=batch, num_partitions=tiles) - for cus in (boundary, boundary + 1): - check( - cus, - (4, True), - num_seqs=batch, - num_partitions=tiles, - work_capacity=capacity, - ) - # NP is itself a compile constant. Compare B4/B25 within each cap: - # only in-range windows with NP64 separate handles; NP65 stays flat. - for parts in (64, 65): - for batch in (4, 25): - capacity = batch * tiles - boundary = split_boundary( - capacity, num_seqs=batch, num_partitions=parts - ) - check( - boundary, - (4, True), - num_seqs=batch, - num_partitions=parts, - work_capacity=capacity, - ) - - # Runtime size/CU metadata must not fragment the kernel specialization cache. - cus, expected = points[-1] - changes = {"num_seqs": kwargs["num_seqs"] * 2} - if planned and kwargs["work_capacity"] is not None: - changes["work_capacity"] = kwargs["work_capacity"] * 2 - multiplier = 2 - if not kwargs["per_token_kv"]: - changes["num_kv_heads"] = kwargs["num_kv_heads"] * 2 - multiplier = 4 - check(cus * multiplier, expected, **changes) - - def _case( name, shape=(4, 1, 16, 128), @@ -767,20 +508,42 @@ def _case( ) -# Positive-window rows reject static calls and run eager/graph checks with plans. -# Disabled-window rows retain the same numerical checks in both modes. +def _cases(columns, rows, *, prefix="", **shared): + """Expand a small parameter table using the same defaults as _case.""" + return [ + _case( + prefix + name, + **{**shared, **dict(zip(columns.split(), values, strict=True))}, + ) + for name, *values in rows + ] + + +# Every row uses the same correctness/contract flow in static and planned modes. +# Positive windows reject static calls; disabled windows run numerics in both. 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 = [ - _case("scalar-direct", (1, 2, 8, 128), (16, 1, 0), 1), - _case( - "scalar-window-sinks", - (1, 2, 8, 128), - (128, 1, 0), - 7, - 257, - FP32, - check_schedule=True, + # Scalar scales, head dimensions and ordinary MTP/window boundaries. + *_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), @@ -790,10 +553,7 @@ def _case( dtype=FP16, lengths=(0, 1, 257, 16384, 16385), ), - _case("head64-fp16", (2, 2, 4, 64), (16, 0, 0), 7, 257, BF16, dtype=FP16), - _case("head1024", (1, 2, 4, 1024), (128, 1, 0), 256, 8192, FP32), - # QL4/W1 can span two tiles, including when explicitly split into four CTAs. - # NP1 therefore retains the multi-tile fallback even with a full plan budget. + # Explicit splits also cover padded query rows and direct/partitioned sinks. _case( "np1-fused-sink", cache=(16, 1, 1), parts=1, window=1, sink=FP32, query_splits=1 ), @@ -812,9 +572,9 @@ def _case( "split4-hkv2-full-m1", (4, 2, 16, 128), (16, 0, 1), - parts=34, - window=4096, - sink=BF16, + 34, + 4096, + BF16, lengths=(1, 257, 4353, 0), query_splits=4, ), @@ -822,417 +582,204 @@ def _case( "full-m1-fp16-d256", (2, 2, 8, 256), (64, 1, 1), - parts=5, - window=257, - sink=FP16, + 5, + 257, + FP16, dtype=FP16, lengths=(0, 1, 257, 513), query_splits=1, ), - _case("mtp3-window", (3, 1, 16, 128), (128, 0, 1), 7, 257, FP16), - _case("mtp2-odd-parts", (2, 2, 16, 128), (16, 0, 1), 3), _case( - "mtp2-query-split", - (2, 1, 16, 128), - parts=3, - lengths=(257, 259), - check_schedule=True, - ), - _case( - "mtp4-window-query-split", - cache=(16, 1, 1), - window=1024, - lengths=(1023, 1024, 1025, 4099), - check_schedule=True, - ), - _case( - "mtp4-window-capacity", - parts=256, - window=1024, - lengths=(0, 3, 1024, 4099), - check_schedule=True, + "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( - "mtp4-dense-capacity", - cache=(128, 0, 1), + "disabled-window-wide", + cache=(16, 0, 1), parts=256, - lengths=(257, 259, 1027, 4099), - workgroup_budget=17, - check_schedule=True, - ), - _case( - "mtp4-window-large-grid", - window=1024, - lengths=(1027,) * 200, - check_schedule=True, - ), - # Exact one-tile budgets enter the measured (B, H*QS, C/B) grid on gfx950. - # Ragged/empty owners are not physical x; each row also runs poisoned - # scratch, changing sinks, minus7/all-empty/restored graph plan refreshes. - _case( - "batch-first-window4096", - parts=18, - window=4096, - sink=FP32, - lengths=(0, 1, 4097, 4353), - workgroup_budget=72, - check_schedule=True, - ), - _case( - "batch-first-window8192", - parts=34, - window=8192, - sink=BF16, - lengths=(0, 1, 3, 4, 257, 8193, 8449, 0), - workgroup_budget=272, - check_schedule=True, - ), - _case( - "batch-first-window8192-batch12", - parts=34, - window=8192, - sink=FP32, - lengths=(0, 1, 8193, 8449) * 3, - workgroup_budget=408, - check_schedule=True, - ), - # Interior window/odd batch: C125 reaches the QS4 threshold at CU63. - _case( - "batch-first-window6000-batch5", - parts=25, - window=6000, + window=-1, sink=FP32, - lengths=(0, 1, 6001, 6145, 6257), - workgroup_budget=125, - check_schedule=True, + query_splits=1, + wide_kv_addressing=True, ), - # These plans satisfy the tile/capacity guards but miss the window range. _case( - "batch-first-window4095-fallback", - parts=18, - window=4095, + "fused-narrow", + cache=(16, 0, 1), + parts=1, sink=FP32, - lengths=(0, 1, 4097, 4353), - workgroup_budget=72, - check_schedule=True, + query_splits=1, + wide_kv_addressing=False, ), + _case("fused-wide", parts=8, query_splits=1, wide_kv_addressing=True), _case( - "batch-first-window8193-fallback", - parts=34, - window=8193, - sink=FP32, - lengths=(0, 1, 8193, 8449), - workgroup_budget=136, - check_schedule=True, + "fused-hkv2-window", + (4, 2, 16, 128), + (128, 0, 1), + 8, + 257, + BF16, + query_splits=1, + wide_kv_addressing=False, ), - _case( - "batch-first-window1024-rejected", - parts=6, - window=1024, - sink=FP32, - lengths=(0, 1, 1025, 1281) * 2, - workgroup_budget=48, - check_schedule=True, + # Small/large batches and overprovisioned or tight plan capacities. + _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-", ), - _case( - "batch-first-excess-partitions", - parts=35, - window=8192, + # Ragged/empty owners, both window endpoints and their immediate neighbors. + *_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, - lengths=(0, 1, 8193, 8449), - workgroup_budget=140, - check_schedule=True, ), - _case( - "batch-first-excess-capacity", - parts=35, + *_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, - lengths=(0, 1, 8193, 8449), - workgroup_budget=137, - check_schedule=True, ), - _case( - "batch-first-tight-capacity", - parts=34, - window=8192, - sink=FP32, - lengths=(0, 1, 8193, 8449), - workgroup_budget=135, - check_schedule=True, - ), - _case( - "batch-first-window4096-wide", - parts=18, + *_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=(0, 1, 4097, 4353), + lengths=LENS_4096, workgroup_budget=72, - wide_kv_addressing=True, - check_schedule=True, ), - # Plain V needs CU >= QL * C / 2 = 144 to split this B4/C72 plan. - # Exercise both address widths and a larger partition cap with real inputs. - _case( - "batch-first-window4096-plain-v", - cache=(128, 0, 1), - parts=18, - window=4096, - sink=FP32, - lengths=(0, 1, 4097, 4353), - workgroup_budget=72, - check_schedule=True, - ), - _case( - "batch-first-window4096-plain-v-wide", - cache=(128, 0, 1), - parts=18, - window=4096, - sink=FP32, - lengths=(0, 1, 4097, 4353), - workgroup_budget=72, - wide_kv_addressing=True, - check_schedule=True, - ), - _case( - "batch-first-window4096-np64-plain-v-wide", - cache=(128, 0, 1), - parts=64, - window=4096, - sink=FP32, - lengths=(0, 1, 4097, 4353), - workgroup_budget=72, - wide_kv_addressing=True, - check_schedule=True, - ), - # The unchanged plain-V split threshold is CU272: CU256 uses QS1/flat. _case( "batch-first-window8192-plain-v-auto", cache=(128, 0, 1), parts=34, window=8192, sink=FP32, - lengths=(0, 1, 8193, 8449), - workgroup_budget=136, - check_schedule=True, - ), - _case( - "batch-first-window4096-batch24", - parts=18, - window=4096, - sink=FP32, - lengths=(0, 1, 3, 4, 257, 1025, 2049, 4096, 4097, 4098, 4353, 4354) * 2, - workgroup_budget=432, - check_schedule=True, - ), - # NP need not equal the window bound: exact useful capacity still proves - # one task per tile. Counts span 0/1/4/8/16/17/18 and poisoned padding. - _case( - "batch-first-window4096-np64", - parts=64, - window=4096, - sink=FP32, - lengths=(0, 1, 769, 1793, 4096, 4097, 4353, 4354), - workgroup_budget=144, - check_schedule=True, - ), - # Large NP keeps the flat grid. The 34-task rows still cross the NP256 - # reducer's 32-partition group boundary, including on graph replays. - _case( - "batch-first-window8192-np256-fallback", - parts=256, - window=8192, - sink=FP32, - lengths=(0, 1, 8449, 8450), + lengths=LENS_8192, workgroup_budget=136, - check_schedule=True, - ), - _case( - "ql1-window-prefetch", - (1, 1, 16, 128), - parts=64, - window=8192, - lengths=(8193,) * 8, - workgroup_budget=512, - check_schedule=True, ), - _case( - "ql1-window-prefetch-small-capacity", - (1, 1, 16, 128), + # QL1: layouts, address widths, head counts and padded GQA sizes. + *_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, - workgroup_budget=263, - check_schedule=True, ), - _case( - "ql1-window-prefetch-wide", - (1, 1, 16, 128), - parts=64, - window=8192, - lengths=(8193,) * 8, - workgroup_budget=512, - wide_kv_addressing=True, - check_schedule=True, - ), - # G8 pads half an M-tile. Distinct per-head data/scales also cover planned - # Hkv2; the Hkv1/NP256 row isolates G8 with a large nominal dense grid. - _case( - "ql1-hkv2-g8-window-prefetch", - (1, 2, 8, 128), - parts=64, - window=8192, - sink=FP32, - lengths=(0, 1, 8193, 8449), - workgroup_budget=264, - check_schedule=True, - ), - _case( - "ql1-hkv1-g8-large-np-prefetch", - (1, 1, 8, 128), - parts=256, + *_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=(0, 1, 8193, 8449), - workgroup_budget=132, - check_schedule=True, - ), - # Hkv4 exercises dense M1 prefetch with multi-tile tasks in each KV layout. - _case( - "ql1-hkv4-page128-multi", - (1, 4, 16, 128), - parts=2, - window=1024, - sink=FP32, - lengths=(0, 1, 1025, 1281), - workgroup_budget=32, - check_schedule=True, - ), - _case( - "ql1-hkv4-page16-multi", - (1, 4, 16, 128), - (16, 1, 1), - parts=2, - window=1024, - sink=FP32, - lengths=(0, 1, 1025, 1281), - workgroup_budget=32, - check_schedule=True, - ), - _case( - "ql1-hkv4-plain-v-multi", - (1, 4, 16, 128), - (128, 0, 1), - parts=2, - window=1024, - sink=FP32, - lengths=(0, 1, 1025, 1281), - workgroup_budget=32, - check_schedule=True, - ), - _case( - "ql1-hkv4-page16-direct-sinks", - (1, 4, 16, 128), - (16, 1, 1), - parts=1, - sink=FP32, - lengths=(0, 1, 257, 769), - workgroup_budget=16, - check_schedule=True, + lengths=LENS_8192, ), - # Interior GQA sizes pad to one M-tile; G7/G17 exercise both range misses. - # Fixed groups per fixture keep the handle-only cache checks unambiguous. - *[ - _case( - f"ql1-hkv2-g{group}-window", - (1, 2, group, 128), - parts=8, - window=1024, - sink=FP32, - lengths=(0, 1, 1025, 1281), - workgroup_budget=40, - check_schedule=True, - ) - for group in (7, *range(9, 16), 17) - ], - _case( - "ql1-hkv2-g9-page16-direct-sinks", - (1, 2, 9, 128), - (16, 1, 1), - parts=1, - sink=FP32, - lengths=(0, 1, 257, 769), - workgroup_budget=8, - check_schedule=True, - ), - _case( - "ql1-hkv2-g15-plain-v-multi", - (1, 2, 15, 128), - (128, 0, 1), + *_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=(0, 1, 1025, 1281), - workgroup_budget=16, - check_schedule=True, + lengths=LENS_1024, ), - # Explicit splits also use one padded M-tile for interior GQA sizes. - _case( - "split2-hkv2-g9-page16-direct-sinks", - (2, 2, 9, 128), - (16, 1, 1), + *_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), - workgroup_budget=8, - query_splits=2, ), - _case( - "split4-hkv1-g15-plain-v-window-wide", - (4, 1, 15, 128), - (128, 0, 1), + *_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=(0, 1, 1025, 1281), - workgroup_budget=24, - query_splits=4, - wide_kv_addressing=True, - ), - _case( - "hkv2-prefetch-1wg", - (1, 2, 8, 128), - parts=8, - window=1, - sink=FP32, - lengths=(257,) * 16, - check_schedule=True, + lengths=LENS_1024, + workgroup_budget=40, ), - _case( - "hkv2-prefetch-2wg", - (1, 2, 8, 128), + *_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, - lengths=(257,) * 32, ), - _case( - "hkv2-prefetch-page16", - (1, 2, 8, 128), - (16, 1, 1), - 8, - 257, - FP16, - lengths=(257,) * 32, - ), - _case( - "hkv2-prefetch-plain-v", - (1, 2, 16, 128), - (128, 0, 1), - 8, - sink=BF16, - lengths=(257,) * 32, - ), - # QL1/W1 has M=1: NP=M and capacity=B*M both meet the exact boundary. - _case("hkv2-direct-sinks", (1, 2, 8, 128), parts=1, window=1, sink=BF16), + # Large/automatic partition counts and integer-limit windows. _case( "long-200k", cache=(128, 0, 1), @@ -1240,32 +787,13 @@ def _case( sink=FP32, lengths=(0, 1, 3, 4, 255, 256, 257, 200003), ), - _case("window-int64", cache=(16, 1, 1), window=2**40, wide_kv_addressing=True), - # The single active tile is absolute tile 1; padding stays NaN across replay. _case( - "masked-scale-decode", - (1, 1, 16, 128), - window=1, - sink=FP32, - lengths=(511,), - masked_scale=True, - ), - _case("masked-scale-mtp", window=1, lengths=(511,), masked_scale=True), - _case( - "disabled-window-wide", - cache=(16, 0, 1), - parts=256, - window=-1, - sink=FP32, - query_splits=1, - wide_kv_addressing=True, - ), - _case( - "large-batch-auto", - cache=(128, 0, 1), - parts=None, - lengths=(257,) * 200, + "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( "exact-parts-override", @@ -1274,22 +802,14 @@ def _case( lengths=(200000,), max_partitions=4, ), - _case( - "window255-budget", - (2, 2, 16, 128), - (64, 0, 1), - 64, - 255, - BF16, - workgroup_budget=17, + _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("window256", (3, 2, 8, 128), (128, 0, 1), 7, 256), - _case("window509", cache=(64, 0, 1), parts=86, window=509, sink=FP16), - _case("window8192", (2, 1, 16, 128), parts=256, window=8192), - _case("window-int32-max", cache=(16, 0, 1), parts=1, window=2**31 - 1), - _case("window-int32-overflow", cache=(16, 0, 1), parts=1, window=2**31, sink=FP32), - # QL3/G8 pads 24 query rows to 32; C Date: Mon, 21 Sep 2026 03:16:48 +0000 Subject: [PATCH 61/66] fix(flydsl): use Numeric bitcast in DPP XOR helper --- aiter/ops/flydsl/kernels/dpp_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/aiter/ops/flydsl/kernels/dpp_utils.py b/aiter/ops/flydsl/kernels/dpp_utils.py index ea966bf318c..3b93e430bad 100644 --- a/aiter/ops/flydsl/kernels/dpp_utils.py +++ b/aiter/ops/flydsl/kernels/dpp_utils.py @@ -44,7 +44,7 @@ def dpp_xor_f32(src, offset: int, **kw): from flydsl._mlir.dialects import arith as _arith_dialect from flydsl.expr.typing import T - src_i32 = _to_ir(src).bitcast(T.i32) + 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) From 408aed7e32e3e22bb9642eab87485c10b16caf11 Mon Sep 17 00:00:00 2001 From: fsx950223 Date: Mon, 21 Sep 2026 03:39:22 +0000 Subject: [PATCH 62/66] fix(flydsl): cap planned PA decode partitions at device CU count Default the per-sequence plan limit to the context device's CU count and preserve explicit smaller limits when refreshing a plan. Validate planned decode and reduction against this device limit while retaining the static 256-partition bound. Extend the generic PA decode unit test with default-CU plans, out-of-range limits, refresh reuse, and long-context reduction coverage. Verified 180 test points on gfx950 and two numerical cases with a simulated host CU304 limit on the same real GPU. Signed-off-by: fsx950223 --- aiter/ops/flydsl/kernels/pa_decode_plan.py | 26 +++++--- aiter/ops/flydsl/kernels/pa_decode_reduce.py | 13 +++- aiter/ops/flydsl/pa_decode.py | 22 +++++-- op_tests/test_flydsl_pa_decode.py | 69 +++++++++++++++++--- 4 files changed, 103 insertions(+), 27 deletions(-) diff --git a/aiter/ops/flydsl/kernels/pa_decode_plan.py b/aiter/ops/flydsl/kernels/pa_decode_plan.py index 3c3f6f316e6..f9df016f8f1 100644 --- a/aiter/ops/flydsl/kernels/pa_decode_plan.py +++ b/aiter/ops/flydsl/kernels/pa_decode_plan.py @@ -9,8 +9,6 @@ import triton import triton.language as tl -from .pa_decode_reduce import MAX_CONTEXT_PARTITIONS - @triton.jit def _plan_pa_decode( @@ -102,8 +100,11 @@ def capacity(self) -> int: 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") - if not 1 <= self.max_partitions <= MAX_CONTEXT_PARTITIONS: - raise ValueError("invalid plan max_partitions") + 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): @@ -123,7 +124,7 @@ def plan_pa_decode( context_lengths: torch.Tensor, num_kv_heads: int, *, - max_partitions: int = MAX_CONTEXT_PARTITIONS, + max_partitions: int | None = None, workgroup_budget: int | None = None, sliding_window: int = 0, query_length: int = 1, @@ -137,6 +138,10 @@ def plan_pa_decode( splitting queries can launch multiple CTAs per slot without extra scratch. This is opt-in: uniform or short-context workloads may favor static splits. + The per-sequence partition limit defaults to the context device's CU count; + an explicit limit must be in [1, CU count]. When refreshing an existing + plan, omitting ``max_partitions`` preserves its original limit. + A positive ``sliding_window`` counts visible tokens including the query's own position; 0 and -1 disable it. Context lengths include the MTP tokens, so the planned range is the union of ``query_length`` causal windows. @@ -161,14 +166,15 @@ def plan_pa_decode( raise ValueError("plan supports batches in [1, 4096]") if num_kv_heads < 1: raise ValueError("num_kv_heads must be positive") - if not 1 <= max_partitions <= MAX_CONTEXT_PARTITIONS: - raise ValueError(f"max_partitions must be in [1, {MAX_CONTEXT_PARTITIONS}]") 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 * torch.cuda.get_device_properties(dev).multi_processor_count - ) + workgroup_budget = 2 * num_compute_units if workgroup_budget < 1: raise ValueError("workgroup_budget must be positive") capacity = min( diff --git a/aiter/ops/flydsl/kernels/pa_decode_reduce.py b/aiter/ops/flydsl/kernels/pa_decode_reduce.py index 12596fb5f8e..6de6121b309 100644 --- a/aiter/ops/flydsl/kernels/pa_decode_reduce.py +++ b/aiter/ops/flydsl/kernels/pa_decode_reduce.py @@ -12,6 +12,7 @@ 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, @@ -27,10 +28,15 @@ def _validate_pa_decode_ps_reduce_config( output_dtype_str: str, logits_dtype_str: str, sink_dtype_str: str, + use_work_plan: bool = False, ) -> None: - if not 1 <= max_context_partition_num <= MAX_CONTEXT_PARTITIONS: + # 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 in [1, {MAX_CONTEXT_PARTITIONS}], " + 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: @@ -49,6 +55,7 @@ def is_pa_decode_ps_reduce_supported( 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: @@ -58,6 +65,7 @@ def is_pa_decode_ps_reduce_supported( 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 @@ -102,6 +110,7 @@ def compile_pa_decode_ps_reduce( 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") diff --git a/aiter/ops/flydsl/pa_decode.py b/aiter/ops/flydsl/pa_decode.py index 4868928ddd4..33261a11bf1 100644 --- a/aiter/ops/flydsl/pa_decode.py +++ b/aiter/ops/flydsl/pa_decode.py @@ -135,17 +135,22 @@ def launch_pa_decode_ps_reduce( stream: torch.cuda.Stream, reduce_info: torch.Tensor | None = None, ) -> None: - if context_partition_num > MAX_CONTEXT_PARTITIONS: + 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( - "FlyDSL pa_decode reduce supports at most " - f"{MAX_CONTEXT_PARTITIONS} partitions" + f"FlyDSL pa_decode reduce supports at most {partition_limit} partitions" ) use_sinks = sink_token is not None # Buffer offsets are byte-sized i32 values. Bound the largest attempted # access, not only the current count: an inactive part must not wrap back # into valid data. Nonstandard standalone reducer strides keep the old path. bounded_plan_logits = ( - reduce_info is not None + use_work_plan and context_partition_num <= 64 and query_seq_len > 0 and query_group_size > 0 @@ -178,7 +183,7 @@ def launch_pa_decode_ps_reduce( output.dtype if sink_token is None else sink_token.dtype ), use_sinks=use_sinks, - use_work_plan=reduce_info is not None, + use_work_plan=use_work_plan, query_group_size=query_group_size, bounded_plan_logits=bounded_plan_logits, vectorize_plan_logits=vectorize_plan_logits, @@ -271,6 +276,8 @@ def pa_decode( ``work_plan`` opts into GPU-planned variable partition counts. Build or refresh it with ``plan_pa_decode`` on the current stream after updating lengths. Its ``max_partitions`` must equal ``max_context_partition_num``. + The plan's partition limit is bounded by the query device's CU count; + static scheduling retains the fixed limit of 256. Its window must match ``sliding_window``, and a windowed plan must be built with the same ``query_length`` so it covers every MTP query's window. Planned scratch is packed as [KV heads, plan.capacity, query rows (, D)]; @@ -298,7 +305,10 @@ def pa_decode( raise TypeError("query_length must be an int") if query_length < 1: raise ValueError(f"query_length must be positive, got {query_length}") - if not 1 <= max_context_partition_num <= MAX_CONTEXT_PARTITIONS: + 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}" diff --git a/op_tests/test_flydsl_pa_decode.py b/op_tests/test_flydsl_pa_decode.py index 98d8942148a..df97ccbc097 100644 --- a/op_tests/test_flydsl_pa_decode.py +++ b/op_tests/test_flydsl_pa_decode.py @@ -11,6 +11,7 @@ Both entry points share input generation, the FP32 reference and kernel launch. Context lengths include the MTP query tokens. Explicit partition counts override automatic recommendations; the CLI times attention plus the native reducer. +Plans default to a per-sequence partition cap equal to the device's CU count. Positive sliding windows require a work plan: planned cases check numerics and graph replays, while static cases check rejection by the core and wrapper APIs. Disabled windows retain both static and planned numerical coverage. @@ -19,7 +20,7 @@ import argparse import importlib import itertools -from dataclasses import dataclass +from dataclasses import dataclass, replace from functools import partial import pandas as pd @@ -285,7 +286,7 @@ def scatter(tensor): 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: + if parts is None and not planned: parts = get_recommended_splits( batch, kv_heads, @@ -305,6 +306,8 @@ def scatter(tensor): 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) @@ -359,7 +362,6 @@ def _run_flydsl(*args, sliding_window=0, work_plan=None): plan_pa_decode( args[4], args[2].shape[1], - max_partitions=args[8], query_length=args[7], sliding_window=sliding_window, plan=work_plan, @@ -439,15 +441,32 @@ def _assert_contracts(args, options): 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) - plan = options["work_plan"] + 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 = { - "max_partitions": plan.max_partitions, "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}, @@ -795,6 +814,17 @@ def _cases(columns, rows, *, prefix="", **shared): ), _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), + # With CU > 256, the long row uses 257 tasks and exercises the larger reducer. + _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), @@ -864,6 +894,17 @@ def _cases(columns, rows, *, prefix="", **shared): @pytest.mark.parametrize("case", CASES) def test_pa_decode(case, planned, monkeypatch): """Reject static windows; otherwise check real kernels, plans and 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"): @@ -873,10 +914,20 @@ def test_pa_decode(case, planned, monkeypatch): return output, context = args[0], args[4] scratch, sinks, plan = args[14:17], args[-1], options["work_plan"] - if plan is not None and case.workgroup_budget is not None: - budget_slots = ( - case.workgroup_budget + case.num_kv_heads - 1 - ) // case.num_kv_heads + 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) ) From 436ea3ab0df259ee372255ef619e5e246ec25ce2 Mon Sep 17 00:00:00 2001 From: fsx950223 Date: Mon, 21 Sep 2026 07:27:06 +0000 Subject: [PATCH 63/66] docs(flydsl): trim PA decode comments Condense repetitive kernel, planner, reducer, API, and test documentation while retaining numerical, masking, synchronization, and calling-contract notes. No executable logic or test parameters changed. Signed-off-by: fsx950223 --- aiter/ops/attention.py | 9 +- aiter/ops/flydsl/kernels/pa_decode_kernel.py | 589 ++++++------------- aiter/ops/flydsl/kernels/pa_decode_plan.py | 50 +- aiter/ops/flydsl/kernels/pa_decode_reduce.py | 111 ++-- aiter/ops/flydsl/pa_decode.py | 153 ++--- op_tests/test_flydsl_pa_decode.py | 53 +- 6 files changed, 288 insertions(+), 677 deletions(-) diff --git a/aiter/ops/attention.py b/aiter/ops/attention.py index 5c8f4f908f7..8ec5f8d2640 100644 --- a/aiter/ops/attention.py +++ b/aiter/ops/attention.py @@ -56,13 +56,10 @@ def pa_decode_flydsl( sliding_window: int = 0, work_plan: "PADecodePlan | None" = None, ) -> None: - """FlyDSL decode with an optional work plan. + """FlyDSL decode; ``ps`` is ignored for API compatibility. - Pass ``work_plan`` by keyword. The ``ps`` compatibility argument precedes - ``sinks`` and is accepted and ignored. - ``work_plan`` is available through this Python API; the registered - ``torch.ops.aiter.pa_decode_flydsl`` operator uses static scheduling. - Positive ``sliding_window`` requires a work plan through the Python API. + Pass ``work_plan`` by keyword to enable planning and positive sliding windows. + The registered ``torch.ops.aiter.pa_decode_flydsl`` operator is static-only. """ del ps if _pa_decode_flydsl is None: diff --git a/aiter/ops/flydsl/kernels/pa_decode_kernel.py b/aiter/ops/flydsl/kernels/pa_decode_kernel.py index 32628e0d0eb..f5c3b5da333 100644 --- a/aiter/ops/flydsl/kernels/pa_decode_kernel.py +++ b/aiter/ops/flydsl/kernels/pa_decode_kernel.py @@ -1,25 +1,17 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (C) 2025-2026 FlyDSL Project Contributors -"""Readable tile-programming reference for paged-attention fp8 decode. - -K/V are fp8 e4m3 (FNUZ on gfx942, OCP on gfx950) fed straight into FP8 MFMA; -Q (bf16/f16) and the softmax probabilities P are quantized to fp8 too. Scales -fold out of the matmuls (q/key scale into the QK score, value scale + 1/FP8_MAX -into the epilogue); softmax max/sum stay f32. -The default instruction is ``mfma_f32_16x16x32_fp8_fp8``. Tuned gfx950 BF16 -per-token MTP3/MTP4 shapes use K=128 FP8 MFMA atoms, with the same normalized -Q/P values and operand layouts, including query-split MTP4 CTAs. -The tuned gfx950 BF16 scalar-scale decode specialization instead casts Q/P -directly to fp8, without range normalization or the compensating 1/FP8_MAX. -This is a speed/precision tradeoff: small Q values and long probability tails -can round to zero, and Q must fit the fp8 range. Per-token scales retain the -range-normalized Q/P policy. -``key_scale``/``value_scale`` are either a ``[1]`` per-tensor scalar or a -``[num_blocks, num_kv_heads, block_size]`` per-token tensor (chosen by rank). - -``block_size`` (16/64/128) and ``head_dim`` (multiple of 64) are compile-time -constants. Layouts are logical, not production's preshuffle. +"""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 @@ -28,10 +20,10 @@ * ``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 -One CTA (4 waves) per (seq, kv_head, partition, query group) runs online softmax over -256-token blocks; the 4 waves split tokens for Q.KT and head-dim for P.V, with -an LDS round-trip on P transposing ownership between the two MMAs. +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 @@ -47,19 +39,15 @@ from .tensor_shim import buf_base_i64, buf_copy_store, ptr_buf_tensor from .utils import rcp_f32 -MFMA_MNK = ( - 16 # M = N = 16 for the MMA atom; also query rows handled per CTA (padded to 16) -) -# Logical K coverage across a wave of one i64 FP8 operand pack per lane. -# Operand layouts retain this unit for both K32 and K128 instructions. +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 -# f32 C-fragment elements per lane for a 16x16 atom, independent of its K. +# Accumulator elements/lane are independent of MFMA K. MFMA_ACC_ELEMS = MFMA_MNK * MFMA_MNK // WAVE LOG2E = 1.4426950408889634 KV_COMPUTE_BLOCK = 256 -# Cache only the selected specialization, not the batch/head/CU counts used -# to choose it. Entries are published after both JIT wrappers are constructed. +# Key by selected specialization, not batch/head/CU scheduling inputs. _PA_DECODE_TILE_CACHE = {} @@ -85,63 +73,30 @@ def compile_pa_decode_tile( use_sinks: bool = False, sink_dtype_str: str = "f32", ): - """Select the schedule and cache the PA-decode kernel + launch wrapper. - - Batch/head/CU counts and plan capacity select the schedule. Calls that - select the same specialization share one kernel and launch wrapper. - ``query_splits=None`` selects automatically; an explicit count overrides - splitting while retaining the matching prefetch policy. Planned splitting - uses a host-known task-count upper bound, without reading GPU metadata. - Sufficient window capacity also proves that every active planned task has - one KV tile, selecting a single-iteration specialization. - - ``block_size``, ``head_dim``, and ``query_dtype`` are compile-time - constants. ``query_length`` (MTP) and ``query_group_size`` flatten into - ``TOTAL_ROWS = query_length * query_group_size``. ``query_splits`` assigns - equal groups of query positions to separate CTAs, each processing - ``M_TILES = ceil(TOTAL_ROWS / query_splits / 16)`` MFMA tiles. The partial - output layout remains unsplit so the existing FlyDSL reducer can combine - context partitions without an intermediate transpose. - - The single-tile specialization requires a standard work plan whose - host-known capacity proves every nonempty record contains one KV tile. - Padding records still skip the whole task, and partial outputs keep the - same packed layout and reducer as multi-tile plans. - - Selected split-query, single-tile plans map packed slots over a (B, C/B) - physical grid. Physical x is not a sequence index: each active task still - takes its sequence from the plan record. This specialization requires - positive launch B, capacity C divisible by B, and launch sizes consistent - with the schedule selection. These runtime size preconditions are checked - here, not dynamically in the JIT launch. Neither B nor C/B enters the - kernel specialization-cache key. - - Planned BF16 D128 partial output automatically uses cached 64-bit buffer - stores on gfx950 when a packed slot's byte span fits a positive signed i32. - Each descriptor is rebased in i64 to one slot. The view retains only the - input dtype's 2-byte alignment; offsets inside the slot are multiples of - 8 bytes. Partial statistics and direct output are unchanged. - - A positive ``sliding_window`` requires a work plan and includes the query - token itself. The plan covers the union of the MTP windows; scores are - masked per query row within those tiles. - ``use_sinks`` adds a per-query-head zero-value attention logit in the final - epilogue for direct NP=1 output. Partitioned/planned execution instead adds - it exactly once in the reducer, leaving partial statistics unchanged. - - Tokens at/after ``context_len`` still take part in the PV matmul with a zero - probability, so their V bytes must be finite: the MFMA does not treat a zero - operand as an annihilator (``0 * NaN == NaN``). Whole pages past the - sequence's extent are pinned to block 0 (see ``_load_phys_scalar``); the - remaining slots are the unwritten tail of the last page the sequence owns, - which the caller is expected to leave finite. - - ``wide_kv_addressing`` computes K/V element offsets in 64-bit. It is only - correct-critical once a single cache tensor passes 2 GiB, where the i32 - page-stride product wraps, and it is not free: measured +18% at - ``block_size=64`` (one page index feeds four step-loads there, so i32 lets - the page-base multiply be hoisted and the offsets stay 32-bit adds) and - neutral at ``block_size=16``. The wrapper turns it on from the cache size. + """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") @@ -149,11 +104,9 @@ def compile_pa_decode_tile( 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 V scheduling also supports f16 and MTP; the numerical fast path - # below further restricts this gate to bf16 single-query decode. + # Scalar scheduling is broader than the BF16 single-query numerical fast path. TUNED_SCALAR = TUNED_SHAPE and trans_v and not per_token_kv - # Preserve the dense-grid estimate for the existing V-prefetch policies. 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 @@ -162,29 +115,23 @@ def compile_pa_decode_tile( if work_capacity is not None: split_workgroups = work_capacity * num_kv_heads if sliding_window > 0: - # An unaligned union of Q causal windows spans at most this many - # tiles per sequence. Padding slots do not provide useful CTAs. + # 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 ) - # With n nonempty sequences, T <= n * window_tiles and the - # remaining budget is at least n * (window_tiles - 1). Prefix - # apportionment therefore grants each sequence at least tiles-1 - # extras, so clamping leaves exactly one task per visible tile. - # This holds for variable lengths and every in-place plan refresh. + # 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: - # Small MTP2/MTP4 grids use one query per CTA. Proven one-tile MTP4 - # page128 plans favor several lightweight M1 CTAs per CU over fused - # MTP CTAs, so allow larger grids there. Preserve the existing - # threshold for static, MTP2 and multi-tile paths. + # 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 @@ -204,10 +151,7 @@ def compile_pa_decode_tile( QUERIES_PER_CTA = query_length // query_splits TOTAL_ROWS = query_length * query_group_size - # One query with GQA8..16 implies one M-tile. Split queries, page16 and - # plain V prefetch unconditionally; transposed page128 uses the grid load - # across all KV heads. One-tile plans allow two effective tasks per CU; - # other grids keep the lower-register schedule above one task per CU. + # Prefetch single-M-tile queries; large multi-tile grids favor fewer registers. PER_TOKEN_M1 = ( TUNED_PER_TOKEN and QUERIES_PER_CTA == 1 @@ -226,9 +170,7 @@ def compile_pa_decode_tile( and TOTAL_ROWS <= MFMA_MNK and num_compute_units < dense_workgroups <= 2 * num_compute_units ) - # Bound batch/window sizes and keep the exact one-tile task budget - # within the small planned reducer. Its cap may exceed the window bound - # without increasing the packed grid capacity. + # Require an exact one-tile task budget and the small planned reducer. batch_first_plan_grid = ( TUNED_PER_TOKEN and single_tile_plan @@ -242,8 +184,6 @@ def compile_pa_decode_tile( and work_capacity == num_seqs * window_tiles and split_workgroups <= 2 * num_compute_units ) - # Normalize runtime scheduling metadata to the final specialization before - # looking up the cache, including explicit query-split overrides. cache_key = ( head_dim, query_group_size, @@ -275,8 +215,7 @@ def compile_pa_decode_tile( and IS_BF16 and 0 < TOTAL_ROWS * head_dim * 2 <= 0x7FFFFFFF ) - # Context lengths are int32, so larger windows have identical visibility. - # Bound the device constant while retaining the caller's value in the plan. + # 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 = ( @@ -300,17 +239,13 @@ def compile_pa_decode_tile( assert ( head_dim % 64 == 0 ), f"pa_decode_tile only supports head_dim that's a multiple of 64, got {head_dim}" - # Flattened query-row axis (MTP outer, GQA head inner), tiled into 16-row M-tiles. + # 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 - # Keep local BF16 maxima in f32 instead of packing/unpacking every max. - # Other dtypes and non-windowed/static paths retain their input-type reduction. + # Avoid repeated BF16 max packing/unpacking. Q_ABSMAX_F32 = sliding_window > 0 and TUNED_PER_TOKEN - # Cooperatively stage each token's scales once instead of having all four - # rgroups load/store the same 64 values. Reuse this producer for the legacy - # multi-tile QL1 path and prefetched single-tile QL1 CTAs; keep - # the LDS layout, consumers and publication barriers unchanged. + # Stage each token's scales once rather than once per rgroup. UNIQUE_SCALE_STAGING = ( sliding_window > 0 and TUNED_PER_TOKEN @@ -320,42 +255,28 @@ def compile_pa_decode_tile( and query_length == 1 and single_tile_plan == prefetch_v ) - # Fix the tuned scalar-scale shape to the direct-fp8 fast path. This - # removes the numerical-policy switch, not the fp8 range/precision tradeoff - # documented above. Per-token and other shapes keep normalized conversion. SCALAR_FP8_DECODE = ( TUNED_SCALAR and IS_BF16 and query_length == 1 and query_group_size in (8, 16) ) - # CDNA4 contracts four legacy fp8 K=32 groups in one K=128 instruction. - # Concatenating the existing packs preserves cache/LDS layouts for fused - # MTP3/MTP4 and query-split MTP4 CTAs. + # 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 - # Small grids split query positions into separate CTAs; otherwise retain - # all four queries here and share KV loads, scales, and P publication. MTP4_FUSED = WIDE_FP8_MFMA and M_TILES == 4 - # With 64-bit cache addresses, page-128 V carried through QK increases - # register pressure: load plain V during P packing and transposed V after - # P publication. Page-16 gathers and narrow addresses retain prefetching. + # 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) - # Page 16 additionally delays next K until PV and uses IGLP to overlap - # groups of V loads with QK MFMA instructions. TUNE_PAGE128 = block_size == 128 and (PER_TOKEN_M1 or TUNED_SCALAR) PAGE16_VPIPE = prefetch_v and block_size == 16 - # The selected page16 pipeline is always per-token as well. 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) ) - # Apply score scales before the -inf mask on the tuned single-M-tile - # paths. This removes the second mask without introducing -inf * 0. + # 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 layout: V=A, P=B -> output [head-dim (row), query-row (col=lane16)], - # generalized over head_dim via the VHE_CHUNKS loop. + # 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 @@ -371,8 +292,7 @@ def compile_pa_decode_tile( head_dim % (NWARP * MFMA_MNK) == 0 ), "head_dim must split across the 4 warps for PV" - # head_dim splits into 16-element chunks (QK_CHUNK_ELEMS, one dwordx4 load); - # RGROUP_QUARTERS of them make a 64-element fetch group, QKHE_LOOP groups total. + # 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) @@ -384,24 +304,20 @@ def compile_pa_decode_tile( 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" - # Q-quant chunk width: NQCHUNK stays fixed at 16 (tied to `lane16`'s role - # as the absmax butterfly width); QCHUNK scales with head_dim instead. + # 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) - # The chunk is fetched in min(8, QCHUNK)-wide pieces (128b max per load), so - # a QCHUNK above 8 that isn't a multiple of 8 would leave its tail unloaded. + # 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" ) - # Each per-lane Q chunk uses at most 128 bits per buffer load. 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 - # PV processes one VHE_SIZE-wide slice of the output head dimension. 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. @@ -428,38 +344,26 @@ def compile_pa_decode_tile( STATE_PER_M = VHE_CHUNKS + 2 V_DATA_SLOT = 2 + STATE_PER_M * M_TILES - # -- LDS layout (shared across the 4 warps) -- - # sQ: fp8[ROWS_PADDED,head_dim] staged+quantized query. sP: fp8[16,TILE_TOK] - # quantized probs, sharing sQ's storage after Q has been read into registers. - # sQscale: f32[ROWS_PADDED] except when Q is converted directly to fp8. - # sLmax/sLsum: cross-warp scratch. - # sVPage: V page-index broadcast. sKScale/sVScale/sVScaleMax: per-token K/V - # scale staging. No sO/sM/sL/sCorr: PV output is register-resident/loop-carried - # and stored straight to global (V=A/P=B swap). + # 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 - # Every path finishes all M-tiles' first QK and executes the max-publication - # barrier before packing P. No later KV tile reads sQ, so P may overwrite - # it without another barrier. Keep the still-live Q scales outside both - # regions; both sizes are already multiples of 16 bytes. + # 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 - # +16B row padding breaks a 32-bank LDS conflict on the P-pack writes while - # keeping the row 16B-aligned for PV's ds_read_b128. + # 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 - # Keep tuned cross-wave reduction rows 16-byte aligned for vector reads. - # Other paths retain the original bank-conflict padding. + # 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 - # V page-index broadcast (V's page depends on `rgroup`, shared across warps). sVPage_off = sLsum_off + P_BUFFERS * MFMA_MNK * NWARP_PAD * f32 sVPage_bytes = NWARP * PAGES_PER_CHUNK * 4 # i32 - # Generic per-token KV-scale staging is double-buffered so tt+1 prefetch - # cannot clobber current scales. Proven single-tile tasks never prefetch. + # 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 @@ -471,28 +375,26 @@ def compile_pa_decode_tile( ) # m-independent: one cross-warp slot total_bytes = sVScaleMax_off + sVScaleMax_bytes - # LDS blob: one i32 array carved into typed byte-offset views (see the - # `_lds_*` helpers in the kernel). Every region size above is 4-byte - # aligned, so `total_bytes // 4` covers the blob exactly. + # 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, # [num_seqs*query_length, num_q_heads, head_dim] (written directly when NP==1) - # per-partition partial outputs (combined by the reduce kernel when NP>1): - pmax_ptr: fx.Pointer, # [num_seqs, num_kv_heads, num_partitions, query_length*query_group_size] row max - psum_ptr: fx.Pointer, # [num_seqs, num_kv_heads, num_partitions, query_length*query_group_size] row sum - pout_ptr: fx.Pointer, # same shape plus head_dim; Q_DTYPE normalized O_p/l_p - query_ptr: fx.Pointer, # [num_seqs*query_length, num_q_heads, head_dim] -- row = seq*query_length + qi (MTP position) - key_cache_ptr: fx.Pointer, # [num_blocks, num_kv_heads, head_dim//16, block_size, 16] (blocked, see module docstring) - value_cache_ptr: fx.Pointer, # [num_blocks, num_kv_heads, block_size//16, head_dim, 16] (blocked, see module docstring) - block_tables_ptr: fx.Pointer, # [num_seqs, max_blocks_per_seq] - context_lengths_ptr: fx.Pointer, # [num_seqs] - key_scale_ptr: fx.Pointer, # [1] per-tensor OR [num_blocks, num_kv_heads, block_size] per-token - value_scale_ptr: fx.Pointer, # same shape as key_scale_ptr - sinks_ptr: fx.Pointer, # [num_q_heads], used only for direct NP=1 output + 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, @@ -521,8 +423,7 @@ def _pa_decode_tile_task( fx.Uint32(gpu.block_id("x")) * fx.Uint32(gpu.grid_dim.z) + fx.Uint32(gpu.block_id("z")) ) - # Packed variable-length requests need not own C/B tasks. - # Physical x only groups slots; the plan owns sequence IDs. + # 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) @@ -540,11 +441,8 @@ def _pa_decode_tile_task( psum = fx.recast_iter(fx.Float32, psum_ptr) pout = fx.recast_iter(Q_DTYPE, pout_ptr) if const_expr(buffer_plan_output): - # Widen each factor before multiplication: the packed allocation - # may exceed 2 GiB even though one slot's offsets fit signed i32. + # Widen before multiplying: packed storage can exceed 2 GiB. if const_expr(batch_first_plan_grid): - # The two-dimensional grid still addresses capacity C, not B. - # Keep the complete packed offset correct for every KV head. pout_slot = fx.Int64(kv_h) * fx.Int64(gpu.grid_dim.x) * fx.Int64( gpu.grid_dim.z ) + fx.Int64(part) @@ -557,8 +455,7 @@ def _pa_decode_tile_task( Q_DTYPE, n=TOTAL_ROWS * head_dim, unit_elems=OP_ELEMS, - # A contiguous BF16 view may start at an odd storage offset: - # do not strengthen its base alignment from 2 to 8 bytes. + # BF16 views may be only 2-byte aligned, even for 8-byte stores. unit_stride=1, num_records_bytes=pout_slot_bytes, ) @@ -582,20 +479,15 @@ def _load(elem_idx): return _load - # A cache above 2 GiB overflows an i32 element offset, so the flat view - # and the offsets below widen together. _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): - # `phys * page_elems` is the term that overflows first: it reaches - # 2^31 elements at a 2 GiB cache. `rest` stays inside one page. + # 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 - # Per-lane Q chunk (QCHUNK 16-bit elems) fetched in QLOAD_UNIT-wide - # pieces (128b max per buffer load): head_dim=256 needs 2 pieces. _q_copy_op = ( fx.rocdl.BufferCopy128b() if QLOAD_UNIT == 8 else fx.rocdl.BufferCopy64b() ) @@ -621,15 +513,12 @@ def _load_q_row(qi, gs_head): base_elem = ( row_byte0 + lane16 * (QCHUNK * 2) ) // 2 # byte offset -> element index - # QCHUNK elements, loaded as N_QLOADS contiguous QLOAD_UNIT pieces - # (a buffer load is 128b max); head_dim=256 splits into 2 pieces. return [ _q_load_chunk(base_elem + u * QLOAD_UNIT) for u in range_constexpr(N_QLOADS) ] - # Start all four independent Q loads before the initial K/page/scale - # prefetch. Their quantization and LDS publication remain below. + # Issue independent Q loads ahead of K/page/scale prefetch. q_units_prefetched = None if const_expr(MTP4_FUSED): q_units_prefetched = [] @@ -654,10 +543,8 @@ def _v_load16(byte_off): 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]) - # Bound block_tables to its real extent: the last (partial) 256-token tile - # can index a page past ceil(context/block_size); the bounded resource - # returns page 0 for that out-of-range read instead of faulting (those - # tail tokens are masked out anyway). + # 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 ) @@ -669,8 +556,6 @@ def _v_load16(byte_off): unit_stride=1, num_records_bytes=bt_num_records_bytes, ) - # Per-tensor: a single global scale, read once. Per-token: read - # per-token instead (see _kv_scale_ops/_stage_kv_scale_to_lds below). 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) @@ -690,15 +575,13 @@ def _v_load16(byte_off): part_end_raw = part_start + tiles_per_part part_end = (part_end_raw < num_tiles).select(part_end_raw, num_tiles) - # One i8 blob carved into typed byte-offset pointers. `lds_base` is an - # ir.Value pointer (safe inside scf control flow); the Python `lds` - # handle is not. + # 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): - # Pin the element-size alignment explicitly (a bare recast off the - # byte base would gcd the alignment down to 1). + # 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) @@ -751,10 +634,8 @@ def _v_scale_load(elem_idx): return fx.Vector(fx.memref_load_vec(v_scale_reg)) def _load_phys_scalar(page, vec_width=1): - # Past `num_pages` the block-table entry is padding: whatever the - # caller left there, often a stale id pointing at another sequence's - # live page. Pin those to block 0 so the tokens the softmax masks - # out always resolve to one known, in-bounds page instead. + # 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] @@ -782,8 +663,7 @@ def _stage_v_page_row(phys_vec): ) def _v_page_fetch_and_stage(tt_i32): - # V's page depends on `rgroup` (shared across warps): warp w fetches - # its rgroup row and broadcasts via LDS (read back by _v_page_read_row). + # 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 @@ -805,7 +685,6 @@ def _k_page_read_warp(): return _lds_load(off, fx.Int32, PAGES_PER_CHUNK) def _kv_buf_off(tt_val): - # ping-pong buffer byte offset (0 when single-buffered). if const_expr(per_token_kv and not single_tile_plan): return (tt_val & fx.Int32(1)) * KV_BUF_STRIDE return 0 @@ -824,8 +703,7 @@ def _stage_kv_scale_to_lds(phys_vec, buf_off=0): _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: - # block_size==16: each lane stages its own rgroup's page/token, - # so the 4 sub-blocks stage in parallel across rgroup-groups. + # 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]) @@ -844,7 +722,6 @@ def _stage_kv_scale_to_lds(phys_vec, buf_off=0): ) def _load_scale_vec(base_off, a, buf_off=0): - # This lane's 4 per-token scales for chunk `a` from an LDS scale region. slot = (warp * TOK_PER_WARP + a * MFMA_MNK + rgroup * 4) * f32 return _lds_load(base_off + buf_off + slot, fx.Float32, 4) @@ -854,8 +731,7 @@ def _load_kv_scale_vecs(a, buf_off=0): ) def _mfma_fp8(a_ops, b_ops, a_base, b_base, k_packs, acc): - # Counts and offsets use i64 operand packs per lane. Each instruction - # consumes one pack for K32 or PACKS_PER_MFMA packs for K128. + # 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( @@ -894,9 +770,8 @@ def _mfma_fp8(a_ops, b_ops, a_base, b_base, k_packs, acc): ) return acc - # -- raw dwordx4 K load (A operand) -- - # token = warp*TOK_PER_WARP + a*MFMA_MNK + lane16 (softmax mask and P-pack - # write position below must encode this same formula). + # 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 = [] @@ -910,7 +785,7 @@ def _k_ops(phys, a): ) w = _k_load16(base) # head[he_idx*16 : +16] -> two K32 operand packs if const_expr(block_size == 16): - # help the scheduler overlap the PAGES_PER_CHUNK gathered loads + # Overlap page16 gathers. fx.rocdl.sched_barrier(fx.rocdl.mask_vmem_rd) ops.extend([w[0], w[1]]) return ops # N_SUBCHUNKS i64 operands @@ -937,16 +812,12 @@ def _k_ops_flat(tt_i32): ) return _k_ops_from_phys(phys_vec), phys_vec - # Empty partitions retain neutral state and must not read K/V or the - # block table. In particular, an empty context has no valid first tile. + # 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) - # Issue the V page-index prefetch alongside K; the LDS write is - # visible after the barrier below. if const_expr(REUSE_KV_PAGES): - # The K page vector is already resident. Reuse it for the V - # page matrix instead of issuing a second block-table load. + # Reuse K page IDs for V's LDS broadcast. _stage_v_page_row(phys_vec0) else: _v_page_fetch_and_stage(part_start) @@ -959,9 +830,7 @@ def _k_ops_flat(tt_i32): fx.Vector.filled(PAGES_PER_CHUNK, 0, fx.Int32), ) - # per_token_kv has no single global key_scale/value_scale: scale_qk - # drops the key_scale factor (folded in per-token, see masked_chunks - # below) and v_scale_f is unused (replaced by v_max_scaled). + # 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: @@ -969,9 +838,8 @@ def _k_ops_flat(tt_i32): v_scale_f = fx.Float32(value_scale) NEG_INF = fx.Float32(float("-inf")) ZERO_F = fx.Float32(0.0) - # Softmax scores are finite or the -inf mask sentinel -- never NaN -- so - # nnan lets maxnum lower to a bare v_max (no v_cmp_u NaN check + its s_nop - # hazard) and fuse to v_max3. (ninf must NOT be set: -inf is load-bearing.) + # 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): @@ -987,7 +855,7 @@ def _st1(byte_off, m_idx, val): fx.Vector.from_elements([val], dtype=fx.Float32), ) - # f32[16, NWARP] cross-warp scratch: scalar write at (row, warp), vec read of a row's NWARP valid slots. + # 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( @@ -999,8 +867,7 @@ def _ld_lw_row(base_off, row): return _lds_load(off, fx.Float32, NWARP) def _f32_to_fp8_words(vf32): - # f32 -> fp8 must use the HW cvt (arith.truncf to fp8 doesn't lower); - # pack 4 f32 -> 1 i32 (4 fp8) via two cvt_pk_fp8_f32 calls. + # Use hardware FP8 conversion; arith.truncf cannot lower it. n = vf32.shape[0] words = [] for i in range_constexpr(n // 4): @@ -1020,15 +887,9 @@ def _q_local_absmax(q_unit): else: return fmath.absf(q_unit).reduce(ReductionOp.MAX).to(fx.Float32) - # Each M-tile quantizes 16 rows of the flattened (MTP position, GQA - # head) axis: `flat_idx = m*16 + qh_local`, decomposed as - # `qi = flat_idx // query_group_size`, `gs_head = flat_idx % - # query_group_size` (same convention as `_mtp_groups`). No - # cross-M-tile dependency, so no barriers needed between iterations. + # 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): - # Cast Q directly to the fp8 MFMA type; no per-row absmax - # normalization or Q scale is needed on this specialization. for u in range_constexpr(N_QLOADS): _st_words( q_row_off @@ -1063,10 +924,7 @@ def _quant_q_row(m, q_row_off, q_units): _f32_to_fp8_words(q_scaled_unit), ) if lane16 == 0: - # Transposed [qh][m] (not [m][qh]) so the whole M_TILES-wide - # row for a fixed qh is contiguous, letting the KV-loop read - # it back in one wide load instead of M_TILES separate - # narrow ones -- see the read site below. + # 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): @@ -1074,10 +932,7 @@ def _quant_q_row(m, q_row_off, q_units): qi = flat_idx // query_group_size gs_head = flat_idx - qi * query_group_size q_row_off = m * MFMA_MNK * head_dim - # `flat_idx < CTA_ROWS` is statically true for every lane except - # possibly on the last M-tile (only it can be a partial tile), so - # skip the runtime branch (and its EXEC-mask overhead) entirely - # for every other M-tile. + # 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] @@ -1097,13 +952,9 @@ def _quant_q_row(m, q_row_off, q_units): gpu.barrier() - # First tile's V page-index row, now visible after the barrier above - # (the fetch+LDS-store was issued earlier, alongside k_pf0). v_page_pf0 = _v_page_read_row() - # Q is the B operand, read raw from sQ once per M-tile and held in - # registers. MUST use the exact same (qkhe, rgroup, qkr) -> head_dim - # permutation as K's `_k_ops`. + # 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 @@ -1113,11 +964,6 @@ def _quant_q_row(m, q_row_off, q_units): q_row_off + lane16 * head_dim + he_idx * QK_CHUNK_ELEMS, fx.Int64, 2 ) q_ops_all.extend([chunk[0], chunk[1]]) - # q_ops_all[m*N_SUBCHUNKS+s] for s=0..N_SUBCHUNKS-1, s = qkhe*2+qkr, - # = M-tile m's head[he_idx*16+qkr*8 : +8] of qhead=lane16 - - # QK in NCHUNK chunks of 4 tokens: each chunk yields a f32x4 - # C-fragment, so softmax processes 4 scores at a time (low VGPR peak). _ct = [ fx.Vector.from_elements( [float(a * MFMA_MNK + r) for r in range_constexpr(4)] @@ -1131,9 +977,7 @@ def _score_mask(a, upper, lower): valid = valid & (_ct[a] >= lower) return valid - # -- raw dwordx4 V load (B operand) -- - # One dwordx4 load per (16-token sub-block, head_elem); trans_v only - # changes the offset formula. `sub`/`step` walk pages/16-token sub-blocks. + # 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 @@ -1161,15 +1005,13 @@ def _v_ops(phys_row, vh): ) w = _v_load16(base) if const_expr(block_size == 16): - # help the scheduler overlap the per-page gathered loads (see _k_ops) 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 from each query to the newest query, including split MTP. - # Subtract from tile_valid inside the loop, keeping bounds tile-relative. + # 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) @@ -1183,8 +1025,7 @@ def _v_ops(phys_row, vh): for m in range_constexpr(M_TILES) ] if const_expr(CTA_ROWS % MFMA_MNK != 0): - # Padded rows are not stored; keep their bounds in the same - # nonnegative-offset domain as real query rows. + # Padded rows still need nonnegative causal offsets. causal_offset = [ (offset > 0).select(offset, 0) for offset in causal_offset ] @@ -1202,8 +1043,7 @@ def _l_slot(m): 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]) - # Fused pipeline: carry next V in registers, reusing each chunk only - # after its final query has consumed the current tile's values. + # 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: @@ -1212,9 +1052,8 @@ def _l_slot(m): 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) - # Fixed bounds let the compiler remove the loop/history for proven - # single-tile plans. Keep tt absolute for cache addressing and masks; - # the outer work-record guard still excludes every padded task. + # 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): @@ -1230,8 +1069,7 @@ def _l_slot(m): else fx.Int32(loop_i) ) tok0 = tt * TILE_TOK - # Interleave MFMA with VALU/LDS in the scalar phase-split path, - # and with V loads in the tuned page-16 decode pipeline. + # 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) @@ -1242,23 +1080,15 @@ def _l_slot(m): None, ] # slots 0/1 (K_SLOT/V_SLOT) filled in at m==0 below - # This tile's ping-pong scale buffer; the tt+1 prefetch stages the other. cur_kv_buf = _kv_buf_off(tt) - # Tokens at/after context_len index cache slots the caller never - # wrote. They are masked out of the scores, but their V bytes and V - # scales still reach the PV MFMA / the fp8 normalization, so both - # have to be neutralised here. `context_len` (not the per-row causal - # bound) is the right cutoff: causally-masked tokens inside the - # context hold real data, and using it keeps both masks independent - # of the query row. + # 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): - # A planned tile has 1 <= tile_valid <= INT32_MAX, and W is - # clamped to INT32_MAX. Clamp a negative tile-local left edge - # before subtracting MTP offsets: all tile tokens are >= 0, - # so this preserves the mask without underflow for huge W. + # 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) @@ -1272,9 +1102,8 @@ def _l_slot(m): ).broadcast_to(4) window_scale_thr = None if const_expr(sliding_window > 0): - # Scales outside every query's window must not shrink - # visible probabilities to zero during FP8 normalization. - # Keep this bound query-independent for shared MTP scales. + # 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( [ @@ -1290,8 +1119,7 @@ def _mask_v_scale( ): return _score_mask(a, thr, lower).select(vec, zero) - # V is independent of QK/softmax. Load it early for the tuned - # decode path and reuse it across M-tiles in the phase-split path. + # 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] @@ -1307,8 +1135,6 @@ def _mask_v_scale( _v_ops(v_page_cur, vh) for vh in range_constexpr(VHE_CHUNKS) ] - # q_scale doesn't depend on `m`; read the whole M_TILES-wide row once - # (contiguous via the transposed [qh][m] sQscale layout). q_scale_vec = None if const_expr(M_TILES > 1): q_scale_vec = _lds_load( @@ -1318,14 +1144,11 @@ def _mask_v_scale( def _lmax_off_m(m): return sLmax_off + m * MFMA_MNK * NWARP_PAD * f32 - # Phase-split (M_TILES>1): pass-1 (QK+mask+max) loops all M-tiles into - # per-M-tile LDS slices so they share ONE barrier. M_TILES==1: else below. + # Phase A publishes all M-tiles' QK maxima with one barrier. if const_expr(M_TILES > 1): masked_chunks_saved = [None] * M_TILES - # per_token: compute the m-independent pv_max once, hold only - # k_scale across Phase A, re-read v_scale for Phase B after the - # barrier (peak scale liveness 16, not 32). + # Reread V scales after Phase A to reduce peak register liveness. k_scale_shared = None if const_expr(per_token_kv): v_scale_A = [ @@ -1370,8 +1193,7 @@ def _lmax_off_m(m): ).broadcast_to(4) window_thr = None if const_expr(sliding_window > 0): - # Subtract in integer space first so a large window - # cannot round its left edge before the comparison. + # Subtract before converting to avoid rounding large edges. first_valid_tile = (window_left - causal_offset[m]).to( fx.Float32 ) @@ -1380,11 +1202,8 @@ def _lmax_off_m(m): ).broadcast_to(4) neg4 = fx.Vector.filled(4, float("-inf"), fx.Float32) - # Fold the per-row score scale in BEFORE the -inf mask. A - # zero q_scale (all-zero Q row) would otherwise reach - # `-inf * 0 == NaN` in pass 2 and poison the whole row. - # max() commutes with a non-negative scale, so `pm` below is - # still the correctly scaled row max. + # 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) @@ -1399,9 +1218,7 @@ def _lmax_off_m(m): ] if const_expr(MTP4_FUSED): - # All four causal windows fully cover interior tiles. - # Keep one uniform branch around the sixteen masks, - # including both the left window edge and causal tail. + # Interior tiles need no mask for any of the four queries. masked_all = fx.Vector.from_elements( [ scaled_frags[a][r] @@ -1451,14 +1268,11 @@ def _lmax_off_m(m): ) for sh in (16, 32): pm = fx.maxnumf(pm, pm.shuffle_xor(sh, WAVE), fastmath=fm_nnan) - # `pm` is already scaled (see scaled_frags above); a fully - # masked row keeps -inf, which `safe_max` turns into 0. _st_lw(_lmax_off_m(m), lane16, warp, pm) masked_chunks_saved[m] = masked_chunks - # tt+1 K/V/scale prefetch, issued once here so the V-page read - # reuses this barrier. + # 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) @@ -1466,8 +1280,7 @@ def _lmax_off_m(m): ) if const_expr(not single_tile_plan) and tt1 < part_end: if const_expr(MTP4_FUSED): - # Publish next pages/scales now; defer next K until P - # packing releases the saved score registers. + # 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: @@ -1484,9 +1297,7 @@ def _lmax_off_m(m): v_page_next = _v_page_read_row() next_state[V_SLOT] = v_page_next - # Small generic M-tile groups retain the whole V-scale tile. - # Fused MTP4 instead shares one chunk across all four queries; - # other large groups reread per chunk to bound VGPR liveness. + # 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 = [ @@ -1500,8 +1311,7 @@ def _lmax_off_m(m): v_vh_shared = [ _v_ops(v_page_cur, vh) for vh in range_constexpr(VHE_CHUNKS) ] - # V-scale normalization is independent of the query row. - # Compute it once instead of once per M tile. + # 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) @@ -1510,9 +1320,7 @@ def _lmax_off_m(m): [norm_factor], dtype=fx.Float32 ).broadcast_to(4) - # Give every M tile private P/Lsum storage. Publish all four - # rows with one barrier, then consume them without any - # inter-M overwrite barrier. + # Private P/Lsum slots share one barrier without overwrite races. m_new_saved = [] safe_max_saved = [] for m in range_constexpr(M_TILES): @@ -1534,8 +1342,7 @@ def _lmax_off_m(m): ).broadcast_to(4) ) ls_saved = [ZERO_F for _ in range_constexpr(M_TILES)] - # A four-value V-scale fragment is shared by all four - # query tiles without retaining the full scale tile. + # 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): @@ -1573,16 +1380,14 @@ def _lmax_off_m(m): ls, ) if const_expr(not single_tile_plan) and tt1 < part_end: - # Let next K overlap the P-publication barrier and PV, - # after saved scores have been consumed by P packing. + # 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): - # Contiguous transposed V can wait until saved scores - # are dead, avoiding overlap with their live registers. + # 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) ] @@ -1633,8 +1438,7 @@ def _lmax_off_m(m): and m == M_TILES - 1 and not single_tile_plan ): - # The current chunk is dead; reuse its registers - # for the next tile while the final PV finishes. + # 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( @@ -1646,8 +1450,7 @@ def _lmax_off_m(m): fx.rocdl.sched_barrier(0) if const_expr(MTP4_PREFETCH_V): if const_expr(single_tile_plan): - # Preserve the loop-state shape; no next V is - # loaded, and the epilogue never consumes this slot. + # Preserve the unused next-V slot in the loop state. v_next = ostate[V_DATA_SLOT] else: v_next = fx.Vector.from_elements( @@ -1697,8 +1500,7 @@ def _lmax_off_m(m): if const_expr(single_tile_plan) else fx.maxnumf(m_prev, tile_max, fastmath=fm_nnan) ) - # Fully-invalid row: use 0 as the effective max so masked lanes - # give exp2(-inf-0)==0 (avoids the -inf-(-inf) cancellation). + # 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 @@ -1727,9 +1529,7 @@ def _lmax_off_m(m): + warp * TOK_PER_WARP + rgroup * 4 ) - # The NCHUNK P words scatter at stride MFMA_MNK//4 - # i32 (the token->fp8-lane interleave the PV ds_read_b128 - # expects); one strided store per chunk. + # 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, @@ -1738,10 +1538,8 @@ def _lmax_off_m(m): ) for sh in (16, 32): ls = ls + ls.shuffle_xor(sh, WAVE) - # PV output is [head-dim, query-row=lane16] after the operand - # swap, so correction/denominator are per-lane scalars (no sCorr). - # Empty history must contribute zero: exp2(-inf - safe_max). - # Replacing m_prev with 0 can overflow for negative logits. + # 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) @@ -1769,8 +1567,6 @@ def _lmax_off_m(m): for vh in range_constexpr(VHE_CHUNKS): v_vh = v_vh_shared[vh] acc = fx.Vector.filled(MFMA_ACC_ELEMS, 0.0, fx.Float32) - # SWAPPED operands (V=A, P=B): output row = - # head-dim, output col = query-row=lane16. acc = _mfma_fp8(v_vh, p_ops, 0, 0, NVOPS, acc) op = fx.Vector(acc) if const_expr(per_token_kv): @@ -1783,22 +1579,17 @@ def _lmax_off_m(m): else o_acc[vh] * corr_b + op ) next_state.extend([*o_acc, m_new, l_new]) - # A shared P/Lsum slot needs all reads to finish before - # the next M-tile overwrites it. With alternating slots, - # the next tile's write barrier retires those reads before - # that slot is reused. The Phase A barrier protects reuse - # across loop iterations. Keep the scheduler fence to - # bound register liveness across M-tiles. + # 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: - # M_TILES==1 single tile (m==0 for the _o_slot/_m_slot/_l_slot helpers). 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 - # QK: each NCHUNK chunk accumulates N_SUBCHUNKS packs into an f32x4. frag_Ss = [] for a in range_constexpr(NCHUNK): acc = fx.Vector.filled(MFMA_ACC_ELEMS, 0.0, fx.Float32) @@ -1806,14 +1597,11 @@ def _lmax_off_m(m): k_cur, q_ops_all, a * N_SUBCHUNKS, 0, N_SUBCHUNKS, acc ) frag_Ss.append(fx.Vector(acc)) - # Publish next page ids and scales before the pass-1 barrier. - # Page 16 defers K until PV; page 128 overlaps K with softmax. k_next = k_cur if const_expr(not single_tile_plan) and tt1 < part_end: if const_expr(REUSE_KV_PAGES): - # Publish next pages/scales once. Page 16 defers K - # until the P barrier; per-token page 128 issues it - # after scales so their VMEM wait cannot drain K. + # 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)) @@ -1825,13 +1613,10 @@ def _lmax_off_m(m): if const_expr(per_token_kv): _stage_kv_scale_to_lds(phys_vec1, _kv_buf_off(tt1)) if const_expr(SCALES_BEFORE_CURRENT_V): - # Next scales must reach LDS before their barrier. Issue - # current V afterward so that scale VMEM waits do not - # drain V; softmax can then hide the current V latency. + # 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) ] - # Softmax: each lane owns one qhead (lane%16); register reduce + shuffle_xor. scale = ( scale_qk if const_expr(SCALAR_FP8_DECODE) @@ -1853,12 +1638,7 @@ def _lmax_off_m(m): float("-inf") if M1_SCALE_BEFORE_MASK else -1e30, fx.Float32, ) - # The tuned path follows the M_TILES>1 numerical ordering: - # fold every positive score scale into finite logits before - # selecting -inf for invalid tokens. This makes the reduced - # max directly usable by pass 2 and avoids `-inf * 0` for an - # all-zero Q row. The per-token branch is retained here so - # this ordering stays correct if the gate is widened later. + # 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( @@ -1892,7 +1672,6 @@ def _lmax_off_m(m): ] else: scaled_frags = frag_Ss - # Reused in pass 2 below, halving the mask instruction count. if const_expr(not M1_SCALE_BEFORE_MASK): masked_chunks = [ _score_mask(a, thr, window_thr).select(scaled_frags[a], neg4) @@ -1920,12 +1699,6 @@ def _lmax_off_m(m): warp, pm if const_expr(M1_SCALE_BEFORE_MASK) else pm * scale, ) # redundant across the 4 lanes sharing this qhead - # per_token_kv: max V-scale for the per-tile fp8 normalization. - # Slots at/after context_len are uninitialised, so an arbitrarily - # large one there would shrink every valid probability to 0 in the - # fp8 conversion below -- restrict the reduction to real tokens. - # Causally-masked tokens within the MTP window union keep real - # scales and stay in, keeping pv_max independent of the query row. if const_expr(per_token_kv): pv_max = fx.Float32(0.0) for a in range_constexpr(NCHUNK): @@ -1937,13 +1710,10 @@ def _lmax_off_m(m): pv_max = fx.maxnumf(pv_max, pv_max.shuffle_xor(sh, WAVE)) _st_lw(sVScaleMax_off, 0, warp, pv_max) gpu.barrier() - # V page-index row for next tile, now visible after the 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 - # per_token_kv: combine warps' max V-scale into the tile's - # normalization factor (also the PV correction below). v_max_scaled = None norm_factor_b = None if const_expr(per_token_kv): @@ -1971,9 +1741,8 @@ def _lmax_off_m(m): if const_expr(single_tile_plan) else fx.maxnumf(m_prev, tile_max) ) - # Keep -inf in the loop-carried/persisted max for an empty - # partition, but use zero as the exponent reference so - # -inf-(-inf) cannot manufacture NaNs in P or the correction. + # 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) @@ -1986,14 +1755,11 @@ def _lmax_off_m(m): zero4_p = fx.Vector.filled(4, 0.0, fx.Float32) for a in range_constexpr(NCHUNK): if const_expr(M1_SCALE_BEFORE_MASK): - # Invalid lanes are -inf, hence exp2(-inf-safe_max)=0 - # without a second validity compare/select. Pa = fx.Vector( fx.exp2(masked_chunks[a] - m_new_b, fastmath="fast") ) else: - # Legacy path: re-mask Pa so a fully-masked chunk - # contributes exactly 0. + # Finite mask sentinels require an explicit zero probability. valid_a = masked_chunks[a] > fx.Vector.filled( 4, -1e29, fx.Float32 ) @@ -2015,7 +1781,6 @@ def _lmax_off_m(m): ) p_scaled = Pa * v_scale_this * norm_factor_b elif const_expr(SCALAR_FP8_DECODE): - # This scalar specialization converts P directly. p_scaled = Pa else: p_scaled = Pa * fx.Vector.filled(4, FP8_MAX, fx.Float32) @@ -2023,7 +1788,6 @@ def _lmax_off_m(m): p_off0 = ( sP_off + lane16 * SP_ROW_BYTES + warp * TOK_PER_WARP + rgroup * 4 ) - # NCHUNK P words scatter at stride MFMA_MNK//4 i32 (see phase-split). for a in range_constexpr(NCHUNK): _lds_store( p_off0 + a * (MFMA_MNK // 4) * f32, @@ -2034,8 +1798,6 @@ def _lmax_off_m(m): fx.rocdl.sched_dswr(NCHUNK) for sh in (16, 32): ls = ls + ls.shuffle_xor(sh, WAVE) - # PV (V=A, P=B) -> output [head-dim, query-row=lane16]; same as - # the phase-split path. corr_reg = ( ZERO_F if const_expr(single_tile_plan) @@ -2045,9 +1807,7 @@ def _lmax_off_m(m): _st_lw(sLsum_off, lane16, warp, ls) gpu.barrier() if const_expr(PAGE16_VPIPE): # noqa: SIM102 - # The page matrix is now visible. Issue next K before the - # current P read/PV so those independent operations can - # cover its VMEM latency. + # 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 @@ -2061,8 +1821,6 @@ def _lmax_off_m(m): corr_b = fx.Vector.from_elements( [corr_reg], dtype=fx.Float32 ).broadcast_to(OP_ELEMS) - # Use the early V loads where enabled; other shapes retain - # the existing batched loads before PV. if const_expr(prefetch_v): v_vh_batch = v_vh_shared else: @@ -2087,8 +1845,7 @@ def _lmax_off_m(m): results = yield next_state o_final = results - # Direct-store epilogue: after the PV swap each lane holds its 4 head-dim - # values for one query-row and writes them straight to global. + # 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 @@ -2100,10 +1857,8 @@ def _lmax_off_m(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): - # Keep the sink out of the online Q/P quantization. It only - # changes the final denominator, with zero numerator mass. - # Compare in natural-logit units before multiplying by LOG2E, - # so even a very large finite f32 sink cannot overflow here. + # 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) @@ -2126,8 +1881,7 @@ def _lmax_off_m(m): if const_expr(per_token_kv): o_scale = inv_l elif const_expr(SCALAR_FP8_DECODE): - # Direct P conversion did not scale by FP8_MAX, so preserve - # value_scale but omit FlyDSL's compensating 1/FP8_MAX. + # 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) @@ -2144,8 +1898,7 @@ def _emit(o_norm, sub, query_idx, query_head): ) fx.ptr_store(o_norm, fx.add_offset(output, out_offset)) elif const_expr(buffer_plan_output): - # Existing row guards keep all four BF16 elements inside - # this slot. Preserve the 8-byte width and cached policy. + # 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, @@ -2170,7 +1923,6 @@ def _emit(o_norm, sub, query_idx, query_head): vh * (NWARP * MFMA_MNK) + warp * MFMA_MNK + rgroup * OP_ELEMS ) sub = head_base // OP_ELEMS - # Guard the partial last tile's out-of-range rows (folded away for full tiles). if row < CTA_ROWS: _emit(o_norm, sub, qi_e, qh) @@ -2178,9 +1930,7 @@ def _emit(o_norm, sub, query_idx, query_head): if warp == 0 and rgroup == 0: base = partial_slot * TOTAL_ROWS + global_row if row < CTA_ROWS: - # Convert the running max from log2 units (scale_qk folds - # in LOG2E) to natural-log units: the shared reduce - # re-applies LOG2E itself when combining partitions. + # 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 @@ -2237,11 +1987,8 @@ def _run_task(seq, start, end, context): ) if const_expr(use_work_plan): - # Capacity is fixed for graph capture, so the planner clears any - # unused tail records. The reducer never consumes those slots; - # guard the whole task body to avoid Q processing and scratch - # writes for an empty record. The condition is CTA-uniform, so - # all barriers in the task body remain convergent. + # 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) @@ -2273,8 +2020,8 @@ def pa_decode_tile_launch( value_cache: fx.Pointer, block_tables: fx.Pointer, context_lengths: fx.Pointer, - key_scale: fx.Pointer, # [1] per-tensor OR [num_blocks, num_kv_heads, block_size] per-token - value_scale: fx.Pointer, # same shape as key_scale + key_scale: fx.Pointer, + value_scale: fx.Pointer, sinks: fx.Pointer, max_blocks_per_seq: fx.Int32, num_seqs: fx.Int32, @@ -2289,9 +2036,7 @@ def pa_decode_tile_launch( work_capacity: fx.Int32, stream: fx.Stream = fx.Stream(None), # noqa: B008 ): - # `contract` lets the mul/add pairs in the softmax rescale and the - # epilogue fuse into FMAs. Set as an ambient hint rather than per op; - # an explicit `fastmath=` on an op still wins (see fm_nnan). + # Ambient contract permits FMAs; explicit per-op fastmath still wins. with CompilationContext.compile_hints({"fastmath": "contract"}): pa_decode_tile_kernel( output, diff --git a/aiter/ops/flydsl/kernels/pa_decode_plan.py b/aiter/ops/flydsl/kernels/pa_decode_plan.py index f9df016f8f1..d0cf319b862 100644 --- a/aiter/ops/flydsl/kernels/pa_decode_plan.py +++ b/aiter/ops/flydsl/kernels/pa_decode_plan.py @@ -28,17 +28,15 @@ def _plan_pa_decode( 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: - # Fused MTP queries have different left edges. Schedule their union, - # rounding down to an absolute tile so page/scaling offsets stay valid. + # 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) - # Integer prefix apportionment telescopes to the budget, even when the - # lengths differ by orders of magnitude. Each nonempty request gets one - # task first; clamp extras to useful tiles and the reducer's per-row limit. + # 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) @@ -66,10 +64,8 @@ def _plan_pa_decode( tl.store(work + slot * 4 + 2, end, active) tl.store(work + slot * 4 + 3, seq_ctx, active) - # The launch capacity is static for graph capture. Clear every padded work - # record so the attention kernel can skip it without touching stale scratch. - # These stores and the active stores are disjoint, including when every - # context is empty. + # 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): @@ -80,10 +76,9 @@ def _plan_pa_decode( class PADecodePlan: """Reusable GPU metadata; refresh it whenever context lengths change. - ``work_info`` is [capacity, 4]: sequence, first/last absolute 256-token tile, - original context length. Sliding windows cover the union of MTP queries. + ``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. - Scratch uses [KV heads, capacity, query rows (, head dim)]. """ work_info: torch.Tensor @@ -130,23 +125,17 @@ def plan_pa_decode( query_length: int = 1, plan: PADecodePlan | None = None, ) -> PADecodePlan: - """Build/update a plan on the current stream without GPU-to-CPU readback. - - Allocate once outside graph capture, then pass ``plan=...`` to refresh the - same metadata in place. Include this refresh in end-to-end measurements. - The budget counts task slots over all KV heads before query splitting; - splitting queries can launch multiple CTAs per slot without extra scratch. - This is opt-in: uniform or short-context workloads may favor static splits. - - The per-sequence partition limit defaults to the context device's CU count; - an explicit limit must be in [1, CU count]. When refreshing an existing - plan, omitting ``max_partitions`` preserves its original limit. - - A positive ``sliding_window`` counts visible tokens including the query's - own position; 0 and -1 disable it. Context lengths include the MTP tokens, - so the planned range is the union of ``query_length`` causal windows. - Pass the same window and (when enabled) query length to ``pa_decode`` and - when refreshing the plan. Dense plans remain independent of query length. + """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") @@ -207,8 +196,7 @@ def plan_pa_decode( batch, plan.capacity, plan.max_partitions, - # Larger windows cover every int32 context. Dense planning is - # query-length independent, including its compilation cache key. + # 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), diff --git a/aiter/ops/flydsl/kernels/pa_decode_reduce.py b/aiter/ops/flydsl/kernels/pa_decode_reduce.py index 6de6121b309..1c1f6738c7c 100644 --- a/aiter/ops/flydsl/kernels/pa_decode_reduce.py +++ b/aiter/ops/flydsl/kernels/pa_decode_reduce.py @@ -86,23 +86,18 @@ def compile_pa_decode_ps_reduce( bounded_plan_logits: bool = False, vectorize_plan_logits: bool = False, ): - """Build the partitioned-softmax reduction used by ``pa_decode``. - - Counts up to one wave use one partition per lane and one output element per - thread. For D=128 and larger counts, a 2-D workgroup materializes weights - once in LDS and splits each output element's partition chain over several - waves. Other head sizes retain the register-only lane-striped fallback. - A sink is a per-query-head zero-value logit: include it in the shared max - for stability and add its mass to the denominator once after summing KV. - ``query_group_size=None`` retains runtime GQA addressing; a positive - specialization constant must match the GQA passed to the launch wrapper. - ``bounded_plan_logits`` requires packed, contiguous logits rows and a - positive NP-part byte span no larger than INT32_MAX. A per-sequence buffer - bound then zero-fills inactive parts without reading adjacent sequences. - ``vectorize_plan_logits`` additionally requires D=128, bf16/f16 logits, - a four-byte-aligned logits base and even head/part/group strides. One wave - then assigns two consecutive output elements to each lane, sharing the - unchanged partition weights. Output stores retain scalar alignment. + """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, @@ -142,18 +137,10 @@ def compile_pa_decode_ps_reduce( offset for offset in (32, 16, 8, 4, 2, 1) if offset < reduce_width ] - # The original mapping gives each output element to one thread, so every - # thread walks every partition. That is a good fit for <=1 wave of - # partitions, but it leaves NP=160..256 as a long dependent load/FMA - # chain. For the decode shape used by PA (D=128), split that chain over - # two or eight independent wave pairs. A pair covers the two 64-element - # halves of the output vector, while its y-coordinate selects a disjoint - # contiguous range of partitions. + # 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: - # Eight groups win from NP=128 onward on gfx950; two avoid excessive - # synchronization/thread overhead for the small >64 tail. parallel_groups = 2 if max_context_partition_num <= 96 else 8 head_waves = head_size // warp_size worker_waves = head_waves * parallel_groups @@ -166,8 +153,7 @@ def compile_pa_decode_ps_reduce( max_context_partition_num + parallel_groups - 1 ) // parallel_groups - # Keep the legacy specializations effectively LDS-free. The fields are - # only allocated from the compile-time parallel branch below. + # 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 @@ -222,9 +208,8 @@ def pa_decode_ps_reduce_kernel( 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 sequences may point one past the packed allocation. The - # <=64 path uses clamped loads, so give them a valid base slot; - # every loaded value is masked before it contributes to a result. + # 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 @@ -232,8 +217,7 @@ def pa_decode_ps_reduce_kernel( 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 before narrowing offsets: the whole packed workspace may - # exceed 2 GiB even when this sequence's NP-part span is small. + # 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) @@ -289,9 +273,8 @@ def _sink_exp(sink_value, safe_max): return fx.exp2(shift * c_log2e, fastmath="fast") def _reduce_planned_wave(part_count, shuffle_offsets): - # Both widths use the same packed-scratch and softmax rules. - # Invalid lanes read this sequence's first partition; count=0 - # has a safe global slot 0 base. Select out NaNs before math. + # 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 = ( @@ -305,23 +288,19 @@ def _reduce_planned_wave(part_count, shuffle_offsets): 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) - # Preload before softmax so independent loads can overlap. The - # caller's uniform branch encloses these loads, not just the FMAs. 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): - # The descriptor ends at count*part_stride, before the - # next sequence or NaN padding. Even count=0 is a valid - # zero-sized resource; raw buffer loads return zero. + # 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) ) - # With unit_stride=1 this remains an element offset, not - # a vector-unit index. A D=128 pair ends within its row; - # even count=0 suppresses the complete 32-bit buffer load. + # 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, @@ -329,8 +308,6 @@ def _reduce_planned_wave(part_count, shuffle_offsets): unit_elems=planned_elements_per_thread, ) if fx.const_expr(vectorize_plan_logits): - # Keep the prefetched pair packed in 32 bits across - # softmax; widen only the part being accumulated below. planned_logits.append(planned_loaded_logits) else: planned_logits.append(fx.Float32(planned_loaded_logits)) @@ -401,16 +378,13 @@ def _reduce_planned_wave(part_count, shuffle_offsets): ) else: planned_part_logits = planned_logits[planned_acc_part] - # Elementwise mul then add, with the same increasing-part - # order as the scalar path: no FMA or reassociation hint. + # 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): - # One wave materializes the normalized partition weights once. - # All output waves then reuse those weights from LDS and split the - # long partition loop. This avoids both duplicated exp2 work and - # a ds_bpermute for every output FMA. + # 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 @@ -464,10 +438,7 @@ def _reduce_planned_wave(part_count, shuffle_offsets): 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 its exp2 operand even for an - # inactive planned partition. Keep the static path - # branch-free, but predicate that expensive operation - # when the per-request count is dynamic. + # select evaluates exp2 eagerly; branch around inactive parts. part_scale = zero_f if part_max > neg_inf: part_scale = fx.exp2( @@ -585,25 +556,16 @@ def _reduce_planned_wave(part_count, shuffle_offsets): elif fx.const_expr(max_context_partition_num <= warp_size): if fx.const_expr(use_work_plan): - # Both arms of the dynamic short-count branch, and its - # initial carried value, must have the same scalar/vector type. + # 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): - # The count is uniform across this CTA. Keep the complete - # short reduction under the branch so large NP does not - # force short requests to preload and accumulate padding. - # Small NP stays straight-line: the branch costs more than - # the few clamped loads it removes on those specializations. + # 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): - # Medium-count requests need only the first eight - # lanes/partials, even when NP reserves up to 64. - # Keep the preload, softmax and accumulation inside - # this uniform branch, including empty-row handling. acc = _reduce_planned_wave(8, (4, 2, 1)) else: acc = _reduce_planned_wave( @@ -615,9 +577,7 @@ def _reduce_planned_wave(part_count, shuffle_offsets): ) else: if fx.const_expr(max_context_partition_num == reduce_width): - # Preserve the static mapping: exact powers of two have no - # inactive lanes inside their reduction subgroup, while - # partial subgroups use EXEC-masked loads. + # 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 @@ -690,11 +650,7 @@ def _reduce_planned_wave(part_count, shuffle_offsets): part_logits = fx.Float32(logits[logits_offset]) acc = acc + part_logits * weight else: - # A wave covers several 64-partition chunks. Lane ``l`` owns - # partitions l, l+64, l+128, and l+192 (as present). Reduce the - # lane-local maxima/sums before the usual wave reduction; later, - # select the corresponding local weight and broadcast from - # ``part_idx % 64``. This stays register-only through NP=256. + # Lane l owns l + 64*k; reduce locally, then broadcast each weight. partitions_per_lane = ( max_context_partition_num + warp_size - 1 ) // warp_size @@ -775,9 +731,8 @@ def _reduce_planned_wave(part_count, shuffle_offsets): 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 in the enclosing constexpr-loop scope so the - # FlyDSL dynamic-if rewriter never observes a stale value - # from a previous unrolled chunk. + # 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 = ( diff --git a/aiter/ops/flydsl/pa_decode.py b/aiter/ops/flydsl/pa_decode.py index 33261a11bf1..f9aaf134d96 100644 --- a/aiter/ops/flydsl/pa_decode.py +++ b/aiter/ops/flydsl/pa_decode.py @@ -1,36 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (C) 2025-2026 FlyDSL Project Contributors -"""Readable tile-programming reference for paged-attention fp8 decode. - -K/V are fp8 e4m3 (FNUZ on gfx942, OCP on gfx950) fed straight into -``mfma_f32_16x16x32_fp8_fp8``; Q (bf16/f16) and the softmax probabilities P are -quantized to fp8 too. Scales fold out of the matmuls (q/key scale into the QK -score, value scale + 1/FP8_MAX into the epilogue); softmax max/sum stay f32. -The gfx950 BF16 per-token MTP4 specialization uses K=128 FP8 MFMA atoms -without changing this normalization policy. -The tuned gfx950 BF16 scalar-scale decode path uses direct fp8 Q/P conversion -instead. It trades dynamic range for speed and can truncate small probabilities; -see ``kernels.pa_decode_kernel`` for its scope. Per-token scales keep normalization. -``key_scale``/``value_scale`` are either a ``[1]`` per-tensor scalar or a -``[num_blocks, num_kv_heads, block_size, 1]`` per-token tensor. - -``block_size`` (16/64/128) and ``head_dim`` (64, or a multiple of 128 up to -1024) are compile-time constants. Layouts are logical, not production's -preshuffle. - -* ``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 - -Each CTA (4 waves) processes a KV head and context partition over 256-token -blocks. Small MTP4 grids split query positions among CTAs; larger grids keep -queries fused. The waves split tokens for Q.KT and head-dim for P.V, with an -LDS round-trip on P transposing ownership between the two MMAs. +"""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 @@ -57,21 +32,13 @@ def get_recommended_splits( *, max_context_length: int | None = None, ) -> int: - """Choose a split count before allocating partition scratch buffers. - - Without a host-side ``max_context_length`` hint, retain the legacy policy - and default upper clamp of eight. With a hint, target two CTAs per CU, - capped by the number of 256-token compute tiles and the reducer's limit. - This lets small, long-context batches use more KV partitions without - multiplying the partition count by the number of physical pages per tile. - Short contexts and already large grids retain the legacy recommendation. - - An explicit ``max_partitions`` remains an upper clamp in either mode. - The hint is a scheduling bound, not a replacement for the GPU-resident - per-sequence lengths; no tensor readback or planning kernel is needed. - Callers must allocate scratch for the returned count and pass that exact - count to ``pa_decode``. For heterogeneous lengths this still uses the same - partition count per sequence; it is not a variable-work scheduler. + """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") @@ -88,9 +55,7 @@ def get_recommended_splits( n = ((num_sm + denom - 1) // denom) * split_kv_blocks if max_context_length is not None: legacy = max(4, min(n, 8)) - # KV partitions consist of compute tiles, regardless of physical page - # size. Keep a short-context floor of eight to preserve the existing - # query-split launch and avoid perturbing already tuned small grids. + # 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) @@ -146,9 +111,8 @@ def launch_pa_decode_ps_reduce( f"FlyDSL pa_decode reduce supports at most {partition_limit} partitions" ) use_sinks = sink_token is not None - # Buffer offsets are byte-sized i32 values. Bound the largest attempted - # access, not only the current count: an inactive part must not wrap back - # into valid data. Nonstandard standalone reducer strides keep the old path. + # 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 @@ -161,10 +125,7 @@ def launch_pa_decode_ps_reduce( < context_partition_num * stride_logits_part * logits.element_size() <= 2**31 - 1 ) - # The pair load is one aligned dword. Packed rows are complete D=128 - # vectors, and even strides preserve alignment after the 64-bit rebase. - # The R9 maximum-span gate also bounds the last byte of every pair load. - # Output remains scalar-stored, so its alignment need not match logits. + # Paired loads need dword alignment; scalar output stores do not. vectorize_plan_logits = ( bounded_plan_logits and head_size == 128 @@ -236,7 +197,7 @@ def pa_decode( max_context_partition_num: int, context_partition_size: int = 256, compute_type: torch.dtype = torch.bfloat16, - query_scale: torch.Tensor = None, # [num_seqs * query_length, num_query_heads, 1] or [1] + 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] @@ -247,41 +208,31 @@ def pa_decode( sliding_window: int = 0, work_plan: PADecodePlan | None = None, ) -> None: - """FlyDSL paged-attention fp8 decode. - - The call signature and intermediate-buffer layouts follow the shared - aiter paged-attention decode API. This kernel currently supports FP8 K/V caches, - BF16/FP16 queries, a 256-token compute tile, block sizes 16/64/128, and - ``head_dim=64`` or a multiple of 128 up to 1024. - Sparse attention uses caller-prepared block tables and selected context - lengths. Each independently selected MTP query must have its own table row - and use query_length=1; query_length>1 applies dense causal masking. - A positive ``sliding_window`` requires ``work_plan`` and limits each query - to that many causal tokens, including its own position; 0 and -1 disable - the window. Optional ``sinks`` is a contiguous [num_query_heads] - BF16/FP16/FP32 tensor on the query device. - Each entry is an unscaled, zero-value attention logit shared across batch - and MTP positions. It contributes to the denominator once, independently - of the window; -inf disables a head's sink and +inf suppresses its output. - ALiBi and externally quantized FP8 queries are not supported. Partitioning - is controlled by ``max_context_partition_num`` and ``work_plan``. - - ``context_lengths`` and ``block_tables`` are GPU-resident, so their values - are not inspected here (which would synchronize the device). Callers must - ensure ``0 <= context_lengths[i] <= block_tables.shape[1] * block_size`` and - that every block-table entry used by a sequence is a physical block index in - ``[0, min(key_cache.shape[0], value_cache.shape[0]))`` -- a packed cache - reaches V through a shifted view that spans fewer blocks than K. - - ``work_plan`` opts into GPU-planned variable partition counts. Build or - refresh it with ``plan_pa_decode`` on the current stream after updating - lengths. Its ``max_partitions`` must equal ``max_context_partition_num``. - The plan's partition limit is bounded by the query device's CU count; - static scheduling retains the fixed limit of 256. - Its window must match ``sliding_window``, and a windowed plan must be built - with the same ``query_length`` so it covers every MTP query's window. - Planned scratch is packed as [KV heads, plan.capacity, query rows (, D)]; - the static API's dense per-sequence scratch layout remains unchanged. + """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( @@ -371,9 +322,8 @@ def pa_decode( f"key_cache must contain at least one KV head, got {num_kv_heads}" ) - # Keep NP==1 and NP>1 on the same domain. The tile requires a multiple of - # 64 and complete <=8-element Q-load pieces; the NP>1 reducer additionally - # caps its thread block (head_dim) at 1024. + # 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 @@ -415,10 +365,7 @@ def pa_decode( f"got {tuple(value_cache.shape)} for block_size={block_size}, " f"head_dim={head_dim}" ) - # A packed cache interleaves K and V in one allocation, so the V view starts - # part-way into it and legitimately spans fewer blocks than K. Block ids live - # on the device, so V's usable extent is a caller contract either way; only - # the inverted-argument direction is worth rejecting here. + # 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, " @@ -609,14 +556,10 @@ def normalize_scale(scale, name): psum = exp_sums pout = temporary_output - # An i32 element offset wraps once a single cache tensor passes 2^31 - # elements (2 GiB at fp8). The wider math costs ~18% at block_size=64, so - # pay it only where it is needed. Both caches share the code path, so widen - # if either does. + # Widen before either cache's i32 element offsets can wrap. wide_kv_addressing = max(key_cache.numel(), value_cache.numel()) >= 2**31 - # Partial outputs must exclude the sink: reduction adds it once across all - # partitions. Keep direct NP=1 output fused, without extra scratch/launches. + # 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( @@ -643,7 +586,7 @@ def normalize_scale(scale, name): ) if num_partitions == 1 and work_plan is None: - # NP==1 writes output directly; partials unused (caller buffers ignored). + # Direct output ignores caller scratch; supply unused kernel arguments. dummy = torch.empty(1, dtype=torch.float32, device=dev) pmax = psum = pout = dummy else: diff --git a/op_tests/test_flydsl_pa_decode.py b/op_tests/test_flydsl_pa_decode.py index df97ccbc097..7a9da0ebf2e 100644 --- a/op_tests/test_flydsl_pa_decode.py +++ b/op_tests/test_flydsl_pa_decode.py @@ -1,20 +1,16 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. -"""One parametrized correctness/contract test and a CLI performance sweep. +"""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 -Both entry points share input generation, the FP32 reference and kernel launch. -Context lengths include the MTP query tokens. Explicit partition counts override -automatic recommendations; the CLI times attention plus the native reducer. -Plans default to a per-sequence partition cap equal to the device's CU count. -Positive sliding windows require a work plan: planned cases check numerics and -graph replays, while static cases check rejection by the core and wrapper APIs. -Disabled windows retain both static and planned numerical coverage. +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 @@ -113,7 +109,7 @@ def _require_gpu(): @pytest.fixture(autouse=True) def _default_cuda_device(): - # Do not leak the default device into other files in the shared CI shard. + # Restore the default device for other tests. _require_gpu() previous = torch.get_default_device() torch.set_default_device("cuda") @@ -169,8 +165,7 @@ def run_torch( log_denominator = torch.logaddexp( log_denominator, sinks.float().reshape(1, kv_heads, group, 1) ) - # A +inf sink suppresses all finite KV logits. Fully masked rows need - # an explicit zero because -inf - -inf is undefined without a sink. + # +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) @@ -178,7 +173,6 @@ def run_torch( def _make_inputs(case, planned=False): - """Build one sparse/dense paged cache, launch arguments and reference call.""" 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 @@ -190,7 +184,7 @@ def _make_inputs(case, planned=False): ) query = torch.empty((batch * ql, heads, dim), dtype=case.dtype).uniform_(-0.5, 0.5) if case.masked_scale: - query.zero_() # Also exercise online Q quantization with zero Q 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 ) @@ -205,7 +199,7 @@ def _make_inputs(case, planned=False): 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()) - # Exact binary ratios isolate W=1 masking from existing FP8 P rounding. + # 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 ) @@ -214,9 +208,8 @@ def _make_inputs(case, planned=False): 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 covers short rows and makes repeated sink - # mass across multiple partitions observable without cancellation. - # Per-tensor cases keep their shared scale and use larger FP8 values. + # 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_() @@ -354,7 +347,7 @@ def scatter(tensor): def _run_flydsl(*args, sliding_window=0, work_plan=None): - # Fixtures use the core signature; the wrapper inserts ps before sinks. + # The wrapper needs ps=True before sinks. args = (*args[:-1], True, args[-1]) if work_plan is None: torch.ops.aiter.pa_decode_flydsl(*args, sliding_window=sliding_window) @@ -418,7 +411,7 @@ def _assert_plan(plan, lengths): def _assert_contracts(args, options): - """Reuse the valid inputs for API validation instead of separate fixtures.""" + """Reuse numerical inputs for API validation.""" heads = args[1].shape[1] for invalid, error in [ ([0.0] * heads, TypeError), @@ -528,7 +521,7 @@ def _case( def _cases(columns, rows, *, prefix="", **shared): - """Expand a small parameter table using the same defaults as _case.""" + """Expand parameter rows over shared defaults.""" return [ _case( prefix + name, @@ -538,14 +531,11 @@ def _cases(columns, rows, *, prefix="", **shared): ] -# Every row uses the same correctness/contract flow in static and planned modes. -# Positive windows reject static calls; disabled windows run numerics in both. 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 = [ - # Scalar scales, head dimensions and ordinary MTP/window boundaries. *_cases( "shape cache parts window sink", [ @@ -572,7 +562,6 @@ def _cases(columns, rows, *, prefix="", **shared): dtype=FP16, lengths=(0, 1, 257, 16384, 16385), ), - # Explicit splits also cover padded query rows and direct/partitioned sinks. _case( "np1-fused-sink", cache=(16, 1, 1), parts=1, window=1, sink=FP32, query_splits=1 ), @@ -648,7 +637,6 @@ def _cases(columns, rows, *, prefix="", **shared): query_splits=1, wide_kv_addressing=False, ), - # Small/large batches and overprovisioned or tight plan capacities. _case("mtp2-query-split", (2, 1, 16, 128), parts=3, lengths=(257, 259)), *_cases( "cache parts window lengths workgroup_budget", @@ -660,7 +648,6 @@ def _cases(columns, rows, *, prefix="", **shared): ], prefix="mtp4-", ), - # Ragged/empty owners, both window endpoints and their immediate neighbors. *_cases( "parts window lengths workgroup_budget", [ @@ -724,7 +711,6 @@ def _cases(columns, rows, *, prefix="", **shared): lengths=LENS_8192, workgroup_budget=136, ), - # QL1: layouts, address widths, head counts and padded GQA sizes. *_cases( "workgroup_budget wide_kv_addressing", [ @@ -798,7 +784,6 @@ def _cases(columns, rows, *, prefix="", **shared): prefix="hkv2-prefetch-", parts=8, ), - # Large/automatic partition counts and integer-limit windows. _case( "long-200k", cache=(128, 0, 1), @@ -815,7 +800,7 @@ def _cases(columns, rows, *, prefix="", **shared): _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), - # With CU > 256, the long row uses 257 tasks and exercises the larger reducer. + # CU > 256 gives the long row 257 reducer tasks. _case( "default-cu-head256", (2, 1, 8, 256), @@ -857,7 +842,7 @@ def _cases(columns, rows, *, prefix="", **shared): BF16, workgroup_budget=17, ), - # Masked per-token scales and poisoned reducer padding/tail boundaries. + # 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)], @@ -893,7 +878,7 @@ def _cases(columns, rows, *, prefix="", **shared): @pytest.mark.parametrize("planned", [False, True], ids=["static", "planned"]) @pytest.mark.parametrize("case", CASES) def test_pa_decode(case, planned, monkeypatch): - """Reject static windows; otherwise check real kernels, plans and replays.""" + """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( @@ -932,7 +917,6 @@ def test_pa_decode(case, planned, monkeypatch): context.numel() * args[8], max(context.numel(), budget_slots) ) - # Force only explicitly requested variants; ordinary cases use production defaults. module = importlib.import_module("aiter.ops.flydsl.pa_decode") overrides = { name: getattr(case, name) @@ -997,7 +981,7 @@ def check(disabled_sink=False): _assert_contracts(args, options) if plan is not None: - # Exercise planner integer limits without allocating an INT32_MAX KV cache. + # 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), @@ -1024,7 +1008,6 @@ def run_pa_decode_tile_case( query_length=1, num_partitions=None, ): - """CLI-only timing wrapper around the same input/reference/launch helpers.""" 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" @@ -1172,7 +1155,7 @@ def _parse_args(argv=None): def main(): - # Parse before GPU checks so --help and invalid options remain usable. + # 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") From 12f43f884038728ff027a66d8aca25525c0d0606 Mon Sep 17 00:00:00 2001 From: fsx950223 Date: Mon, 21 Sep 2026 07:46:31 +0000 Subject: [PATCH 64/66] refactor(flydsl): defer PA decode custom-op registration --- aiter/ops/attention.py | 8 ------- csrc/cpp_itfs/torch_utils.py | 37 ++----------------------------- op_tests/test_flydsl_pa_decode.py | 6 ++--- 3 files changed, 4 insertions(+), 47 deletions(-) diff --git a/aiter/ops/attention.py b/aiter/ops/attention.py index 8ec5f8d2640..2879c51dd17 100644 --- a/aiter/ops/attention.py +++ b/aiter/ops/attention.py @@ -59,7 +59,6 @@ def pa_decode_flydsl( """FlyDSL decode; ``ps`` is ignored for API compatibility. Pass ``work_plan`` by keyword to enable planning and positive sliding windows. - The registered ``torch.ops.aiter.pa_decode_flydsl`` operator is static-only. """ del ps if _pa_decode_flydsl is None: @@ -95,13 +94,6 @@ def pa_decode_flydsl( ["output", "exp_sums", "max_logits", "temporary_output"], ) -direct_register_custom_op( - "pa_decode_flydsl", - pa_decode_flydsl, - ["output", "exp_sums", "max_logits", "temporary_output"], - python_only_args=("work_plan",), -) - def gen_pa_fwd_native_fake( # [num_seqs, num_heads, head_size] diff --git a/csrc/cpp_itfs/torch_utils.py b/csrc/cpp_itfs/torch_utils.py index c59424121d8..39b485fb12f 100644 --- a/csrc/cpp_itfs/torch_utils.py +++ b/csrc/cpp_itfs/torch_utils.py @@ -3,9 +3,7 @@ import ctypes -import inspect from collections.abc import Callable -from types import FunctionType import torch from torch.library import Library @@ -92,8 +90,6 @@ def direct_register_custom_op( target_lib: Library | None = None, dispatch_key: str = "CUDA", tags: tuple[torch.Tag, ...] = (), - *, - python_only_args: tuple[str, ...] = (), ): """ `torch.library.custom_op` can have significant overhead because it @@ -106,10 +102,6 @@ def direct_register_custom_op( want to register it to a different library, you can pass the library object to the `target_lib` argument. - `python_only_args` names optional trailing parameters to omit from the - Torch schema, in signature order. Calls through torch.ops use their - defaults; direct Python calls can supply them. - IMPORTANT: the lifetime of the operator is tied to the lifetime of the library object. If you want to bind the operator to a different library, make sure the library object is alive when the operator is used. @@ -121,38 +113,13 @@ def _op_func(*args, **kwargs): log_args(op_func, *args, **kwargs) return op_func(*args, **kwargs) - schema_func = op_func - if python_only_args: - signature = inspect.signature(op_func) - parameters = list(signature.parameters.values()) - excluded = parameters[-len(python_only_args) :] - if tuple(param.name for param in excluded) != python_only_args or any( - param.default is inspect.Parameter.empty for param in excluded - ): - raise ValueError( - "python_only_args must name optional trailing parameters " - "in signature order" - ) - # Preserve the original globals for forward-referenced annotations - # while keeping the public Python signature unchanged. - schema_func = FunctionType( - op_func.__code__, - op_func.__globals__, - op_func.__name__, - op_func.__defaults__, - op_func.__closure__, - ) - schema_func.__signature__ = signature.replace( - parameters=parameters[: -len(python_only_args)] - ) - if hasattr(torch.library, "infer_schema"): - schema_str = torch.library.infer_schema(schema_func, mutates_args=mutates_args) + schema_str = torch.library.infer_schema(op_func, mutates_args=mutates_args) else: # for pytorch 2.4 import torch._custom_op.impl - schema_str = torch._custom_op.impl.infer_schema(schema_func, mutates_args) + schema_str = torch._custom_op.impl.infer_schema(op_func, mutates_args) my_lib = target_lib or aiter_lib my_lib.define(op_name + schema_str, tags=tags) my_lib.impl(op_name, _op_func, dispatch_key=dispatch_key) diff --git a/op_tests/test_flydsl_pa_decode.py b/op_tests/test_flydsl_pa_decode.py index 7a9da0ebf2e..83ff828299b 100644 --- a/op_tests/test_flydsl_pa_decode.py +++ b/op_tests/test_flydsl_pa_decode.py @@ -349,9 +349,7 @@ def scatter(tensor): def _run_flydsl(*args, sliding_window=0, work_plan=None): # The wrapper needs ps=True before sinks. args = (*args[:-1], True, args[-1]) - if work_plan is None: - torch.ops.aiter.pa_decode_flydsl(*args, sliding_window=sliding_window) - else: + if work_plan is not None: plan_pa_decode( args[4], args[2].shape[1], @@ -359,7 +357,7 @@ def _run_flydsl(*args, sliding_window=0, work_plan=None): sliding_window=sliding_window, plan=work_plan, ) - pa_decode_flydsl(*args, sliding_window=sliding_window, work_plan=work_plan) + pa_decode_flydsl(*args, sliding_window=sliding_window, work_plan=work_plan) return args[0] From d9d566396423d45a7377ec6af3c729a4ecbc92c0 Mon Sep 17 00:00:00 2001 From: fsx950223 Date: Mon, 21 Sep 2026 07:50:22 +0000 Subject: [PATCH 65/66] refactor(flydsl): remove unused kernel helpers --- aiter/ops/flydsl/kernels/utils.py | 10 ---------- 1 file changed, 10 deletions(-) diff --git a/aiter/ops/flydsl/kernels/utils.py b/aiter/ops/flydsl/kernels/utils.py index 19ac9d4cfe4..48a3a117bb3 100644 --- a/aiter/ops/flydsl/kernels/utils.py +++ b/aiter/ops/flydsl/kernels/utils.py @@ -10,11 +10,6 @@ def rcp_f32(value): return rocdl.rcp(T.f32, value) -def align_up(value: int, align: int) -> int: - """Round *value* up to the next multiple of *align* (static ints).""" - return ((int(value) + int(align) - 1) // int(align)) * int(align) - - def pow2_shift(value: int) -> int: assert value > 0 and (value & (value - 1)) == 0 return value.bit_length() - 1 @@ -42,8 +37,3 @@ def urem_const(value, divisor: int): if const_expr(is_pow2(divisor)): return urem_pow2(value, divisor) return value % fx.Int32(divisor) - - -def unflatten_k(k_flat, qkhe_loop: int = 2): - n = qkhe_loop * 2 - return [[k_flat[td * n + j] for j in range(n)] for td in range(len(k_flat) // n)] From 452a91d98eb7f246d4e35fbac2472b0d66e37009 Mon Sep 17 00:00:00 2001 From: fsx950223 Date: Mon, 21 Sep 2026 08:05:56 +0000 Subject: [PATCH 66/66] refactor(flydsl): simplify PA decode entry and DPP bitcast --- aiter/ops/attention.py | 63 --------------------------- aiter/ops/flydsl/kernels/dpp_utils.py | 3 +- op_tests/test_flydsl_pa_decode.py | 5 +-- 3 files changed, 2 insertions(+), 69 deletions(-) diff --git a/aiter/ops/attention.py b/aiter/ops/attention.py index 2879c51dd17..2781a46ba79 100644 --- a/aiter/ops/attention.py +++ b/aiter/ops/attention.py @@ -23,71 +23,8 @@ from ..jit.utils.asm_guard import require_gfx1250_asm from ..jit.utils.chip_info import get_cu_num, get_gfx -try: - from aiter.ops.flydsl.pa_decode import PADecodePlan - from aiter.ops.flydsl.pa_decode import pa_decode as _pa_decode_flydsl -except (ImportError, AttributeError, RuntimeError, OSError): - _pa_decode_flydsl = None - MD_NAME = "module_attention" - -def pa_decode_flydsl( - output: torch.Tensor, - query: torch.Tensor, - key_cache: torch.Tensor, - value_cache: torch.Tensor, - context_lengths: torch.Tensor, - block_tables: torch.Tensor, - 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, - value_scale: torch.Tensor = None, - exp_sums: torch.Tensor = None, - max_logits: torch.Tensor = None, - temporary_output: torch.Tensor = None, - alibi_slopes: torch.Tensor = None, - ps: bool = True, - sinks: torch.Tensor = None, - sliding_window: int = 0, - work_plan: "PADecodePlan | None" = None, -) -> None: - """FlyDSL decode; ``ps`` is ignored for API compatibility. - - Pass ``work_plan`` by keyword to enable planning and positive sliding windows. - """ - del ps - if _pa_decode_flydsl is None: - raise RuntimeError("pa_decode_flydsl requires the `flydsl` package") - _pa_decode_flydsl( - output, - query, - key_cache, - value_cache, - context_lengths, - block_tables, - softmax_scale, - query_length, - max_context_partition_num, - context_partition_size, - compute_type, - query_scale, - key_scale, - value_scale, - exp_sums, - max_logits, - temporary_output, - alibi_slopes, - sinks, - sliding_window, - work_plan=work_plan, - ) - - direct_register_custom_op( "pa_decode_gluon", pa_decode_gluon, diff --git a/aiter/ops/flydsl/kernels/dpp_utils.py b/aiter/ops/flydsl/kernels/dpp_utils.py index 3b93e430bad..d93b3b4e65d 100644 --- a/aiter/ops/flydsl/kernels/dpp_utils.py +++ b/aiter/ops/flydsl/kernels/dpp_utils.py @@ -41,7 +41,6 @@ def update_dpp_i32( def dpp_xor_f32(src, offset: int, **kw): """Return ``src`` from the lane selected by a 16-lane XOR DPP pattern.""" - from flydsl._mlir.dialects import arith as _arith_dialect from flydsl.expr.typing import T src_i32 = fx.Float32(src).bitcast(fx.Int32) @@ -59,4 +58,4 @@ def dpp_xor_f32(src, offset: int, **kw): raise ValueError( f"dpp_xor_f32 only supports 16-lane offsets 1, 2, 4, 8; got {offset}" ) - return _arith_dialect.BitcastOp(T.f32, out_i32).result + return out_i32.bitcast(T.f32) diff --git a/op_tests/test_flydsl_pa_decode.py b/op_tests/test_flydsl_pa_decode.py index 83ff828299b..eb41734f223 100644 --- a/op_tests/test_flydsl_pa_decode.py +++ b/op_tests/test_flydsl_pa_decode.py @@ -26,7 +26,6 @@ import aiter from aiter import dtypes, per_tensor_quant, pertoken_quant from aiter.jit.utils.chip_info import get_gfx_runtime -from aiter.ops.attention import pa_decode_flydsl from aiter.test_common import benchmark, run_perftest try: @@ -347,8 +346,6 @@ def scatter(tensor): def _run_flydsl(*args, sliding_window=0, work_plan=None): - # The wrapper needs ps=True before sinks. - args = (*args[:-1], True, args[-1]) if work_plan is not None: plan_pa_decode( args[4], @@ -357,7 +354,7 @@ def _run_flydsl(*args, sliding_window=0, work_plan=None): sliding_window=sliding_window, plan=work_plan, ) - pa_decode_flydsl(*args, sliding_window=sliding_window, work_plan=work_plan) + pa_decode(*args, sliding_window=sliding_window, work_plan=work_plan) return args[0]