diff --git a/aiter/ops/flydsl/kernels/chunk_gated_delta_h_mfma16x16x16.py b/aiter/ops/flydsl/kernels/chunk_gated_delta_h_mfma16x16x16.py new file mode 100644 index 00000000000..59c5f07fae3 --- /dev/null +++ b/aiter/ops/flydsl/kernels/chunk_gated_delta_h_mfma16x16x16.py @@ -0,0 +1,1203 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +""" +Gated Delta Net K5 hidden-state recurrence kernel using the @flyc.kernel API. + +mfma16 / HIP-aligned fork (formerly the "vk" fork): the compute path uses the +16x16x16 bf16 MFMA (``mfma_f32_16x16x16bf16_1k``) -- the SAME instruction as the +hand-tuned HIP/C++ K5 kernel -- and the SAME warp partition (BT split-M, K split +across waves, V not split across warps). This fork is NON-VWARP / split-M ONLY +(the alternative OPT-VWARP layout has been removed). It writes the public VK +layout [..., V, K] via a [V][K] transpose buffer + b128 store (HIP-aligned). + +For each chunk t (serial over NT chunks): + 1. Store h snapshot for downstream K6 + 2. v_new = u - w @ h (delta correction via MFMA) + 3. Gated decay + state update: + v_new *= exp(g_last - g_cumsum) + h = h * exp(g_last) + k^T @ v_new +""" + +import math + +import flydsl.compiler as flyc +import flydsl.expr as fx +from flydsl._mlir import ir +from flydsl.compiler.kernel_function import CompilationContext +from flydsl.expr import arith, const_expr, gpu, range_constexpr, rocdl, vector +from flydsl.expr.typing import T +from flydsl.runtime.device import get_rocm_arch +from flydsl.utils.smem_allocator import SmemAllocator, SmemPtr + +from .tensor_shim import GTensor, STensor, _to_raw + +_LOG2E = math.log2(math.e) # 1.4426950408889634 +_LLVM_GEP_DYNAMIC = -2147483648 + + +def _llvm_lds_ptr_ty(): + return ir.Type.parse("!llvm.ptr<3>") + + +def _make_fast_exp(g_is_log2_scaled: bool): + """Return the ``exp`` helper for this kernel compile. + + If ``g_is_log2_scaled`` is False (default), ``g_cumsum`` is in the natural + log domain (matches upstream K12) and we lower ``exp(x)`` as + ``exp2(x * log2(e))`` so the multiplier merges into one ``v_exp_f32`` plus + one ``v_mul_f32`` on AMD. + + If True, the caller has pre-scaled ``g_cumsum`` by ``log2(e)`` already + (the K12 prescale optimization), so we can drop the per-call ``* LOG2E`` + multiply and lower directly to a single ``v_exp_f32``. NOTE: enabling + this flag without the matching K12 prescale produces incorrect outputs; + it exists for ISA-level perf probing of the prescale upper bound. + """ + if g_is_log2_scaled: + + def _fast_exp(x): + return rocdl.exp2(T.f32, x) + + else: + + def _fast_exp(x): + return rocdl.exp2(T.f32, x * _LOG2E) + + return _fast_exp + + +def _mfma_bf16_16x16x32(a_bf16x8, b_bf16x8, acc_f32x4): + """Single mfma_f32_16x16x32_bf16 instruction.""" + return rocdl.mfma_f32_16x16x32_bf16( + T.f32x4, [a_bf16x8, b_bf16x8, acc_f32x4, 0, 0, 0] + ) + + +def _mfma_bf16_16x16x16(a_bf16x4, b_bf16x4, acc_f32x4): + """Single mfma_f32_16x16x16_bf16 instruction (gfx950 / CDNA bf16 1k form). + + K-tile is 16 (half of the 16x16x32 variant), so A and B are bf16x4 + (one ds_read_tr16_b64 worth) instead of bf16x8. Output C is f32x4, + identical layout to the 16x16x32 form. + + The ``rocdl.mfma.f32.16x16x16bf16.1k`` op takes its A/B operands as + ``vector<4xi16>`` (bf16 bit-pattern as signless i16), so bitcast the + bf16x4 fragments before handing them over. + """ + a_i16x4 = a_bf16x4.bitcast(fx.Int16) + b_i16x4 = b_bf16x4.bitcast(fx.Int16) + return rocdl.mfma_f32_16x16x16bf16_1k( + T.f32x4, [a_i16x4, b_i16x4, acc_f32x4, 0, 0, 0] + ) + + +def _f32x4_to_bf16x4_rne_gfx950(vec_f32x4): + """Round-to-nearest-even f32x4 -> bf16x4 via the gfx950 native convert. + + Mirrors HIP's ``float_to_bf16`` which uses the gfx950 native RNE convert + ``v_cvt_pk_bf16_f32`` (``static_cast<__bf16>``), NOT the FlyDSL default + ``arith.truncf`` (bit-truncation, differs from torch/HIP by ~1 ulp). Packs + the 4 lanes with two ``cvt_pk_bf16_f32`` (each 2xf32 -> i32 holding 2xbf16), + then bitcasts the 2xi32 back to a vector<4xbf16>. Returns a raw + vector<4xbf16> ir.Value (drop-in for the previous ``.truncf(vec4 bf16)``). + + gfx950 (CDNA4) ONLY -- ``v_cvt_pk_bf16_f32`` does not exist on gfx942. + """ + lo = rocdl.cvt_pk_bf16_f32(vec_f32x4[0], vec_f32x4[1]) # i32: [bf16(0), bf16(1)] + hi = rocdl.cvt_pk_bf16_f32(vec_f32x4[2], vec_f32x4[3]) # i32: [bf16(2), bf16(3)] + packed = vector.from_elements(T.vec(2, T.i32), [lo, hi]) + return vector.bitcast(T.vec(4, T.bf16), packed) + + +def _f32x4_to_bf16x4_rne_portable(vec_f32x4): + """Round-to-nearest-even f32x4 -> bf16x4 without ``v_cvt_pk_bf16_f32``. + + Portable software RNE for architectures lacking the gfx950 native + ``v_cvt_pk_bf16_f32`` convert (e.g. gfx942 / CDNA3). Emulates HIP's + ``__float2bfloat16_rn`` software fallback with the standard "add a + round-to-nearest-even bias, then keep the high 16 bits" integer + sequence, applied lane-wise over the whole vector<4xf32>: + + x = bitcast(f) + rounding_bias = 0x7FFF + ((x >> 16) & 1) # even-tie -> +0x7FFF, odd -> +0x8000 + bf16_bits = (x + rounding_bias) >> 16 + + This matches torch/HIP RNE (not FlyDSL's default ``truncf`` truncation). + Inf/NaN survive: the bias never carries a finite value up to Inf, and a + NaN keeps a non-zero mantissa. Returns a vector<4xbf16> ir.Value, a + drop-in replacement for the gfx950 path. + """ + i32x4 = T.vec(4, T.i32) + x = vector.bitcast(i32x4, vec_f32x4) + c16 = arith.constant_vector(16, i32x4) + c1 = arith.constant_vector(1, i32x4) + c7fff = arith.constant_vector(0x7FFF, i32x4) + lsb = arith.andi(arith.shrui(x, c16), c1) + rounding_bias = arith.addi(lsb, c7fff) + rounded = arith.addi(x, rounding_bias) + hi = arith.shrui(rounded, c16) + hi16 = arith.trunci(T.vec(4, T.i16), hi) + return vector.bitcast(T.vec(4, T.bf16), hi16) + + +def _f32x4_to_bf16x4_trunc(vec_f32x4): + """Truncating f32x4 -> bf16x4: keep the high 16 bits, NO rounding bias. + + Matches HIP's ``float_to_bf16`` (``__builtin_bit_cast(x) >> 16``) + exactly, so outputs are bit-identical to the HIP/C++ K5 kernel. Arch- + independent (pure integer shift; no ``v_cvt_pk_bf16_f32``). + """ + i32x4 = T.vec(4, T.i32) + x = vector.bitcast(i32x4, vec_f32x4) + c16 = arith.constant_vector(16, i32x4) + hi = arith.shrui(x, c16) + hi16 = arith.trunci(T.vec(4, T.i16), hi) + return vector.bitcast(T.vec(4, T.bf16), hi16) + + +# fp32->bf16 output-conversion mode. When True (default), use bit-truncation to +# match HIP's ``float_to_bf16`` (``bit_cast(x) >> 16``) so flydsl-hip +# outputs are bit-identical to the HIP/C++ K5 kernel. When False, use RNE +# (~0.5 ulp closer to the FP32 reference, but NOT bit-matching HIP). +_BF16_CONVERT_TRUNC = True + + +def _f32x4_to_bf16x4_rne(vec_f32x4): + """Arch-aware fp32x4 -> bf16x4 output conversion. + + With ``_BF16_CONVERT_TRUNC`` (default) this truncates to match HIP exactly. + Otherwise it uses the native ``v_cvt_pk_bf16_f32`` RNE convert on gfx950 + (CDNA4) and a portable integer software RNE everywhere else (gfx942 / CDNA3 + has no ``v_cvt_pk_bf16_f32``). Called at trace time, so dispatching on + ``get_rocm_arch()`` here selects the right lowering per compile. + """ + if _BF16_CONVERT_TRUNC: + return _f32x4_to_bf16x4_trunc(vec_f32x4) + if "gfx950" in get_rocm_arch(): + return _f32x4_to_bf16x4_rne_gfx950(vec_f32x4) + return _f32x4_to_bf16x4_rne_portable(vec_f32x4) + + +# -- Compile the kernel --------------------------------------------------- + + +def compile_chunk_gated_delta_h_mfma16_hip( + *, + K: int, + V: int, + BT: int = 64, + BV: int = 32, + H: int, + Hg: int, + USE_G: bool = True, + USE_GK: bool = False, + USE_INITIAL_STATE: bool = True, + STORE_FINAL_STATE: bool = True, + SAVE_NEW_VALUE: bool = True, + IS_VARLEN: bool = True, + WU_CONTIGUOUS: bool = True, + STATE_DTYPE_BF16: bool = False, + G_IS_LOG2_SCALED: bool = False, + USE_STATE_INDICES: bool = False, + SCHED_GFX942: bool = False, +): + """Compile the GDN K5 kernel. + + Returns a @flyc.jit function: + launch_fn(k, v, w, v_new, g, gk, h, h0, ht, + cu_seqlens, chunk_offsets, + T_val, T_flat, N_val, stream) + + When ``STATE_DTYPE_BF16=False`` (default) the SSM state tensors ``h0`` / + ``ht`` are ``float32``. When ``STATE_DTYPE_BF16=True`` they are + ``bfloat16``: ``h0`` is ``extf``-promoted to f32 right after each load, + and ``ht`` is ``truncf``-demoted to bf16 right before each store. The + f32 accumulator (``h_accs``) and all intermediate LDS layouts are + unchanged, so this only affects HBM bandwidth / footprint of the SSM + state. Mirrors the pattern used by ``kernels/gdr_decode.py``. + """ + assert K <= 256 + assert K % 64 == 0 + assert BV % 16 == 0 + NUM_K_BLOCKS = K // 64 + + # Tensor slots use the ``fx.Pointer`` ABI (raw data pointer). The kernel + # body wraps every slot as ``GTensor(..., shape=(-1,))`` and never reads + # the FlyDSL-injected memref shape/stride, so passing a bare pointer + # produces identical device code while skipping the per-launch DLPack + # export + layout-buffer packing that the default layout-dynamic + # ``fx.Tensor`` memref incurs under flydsl >=0.2.0. The host side wraps + # each tensor with ``flyc.from_c_void_p`` (see ``_as_ptr`` in the host + # wrapper module), which requires flydsl >=0.2.0. + + _fast_exp = _make_fast_exp(G_IS_LOG2_SCALED) + + WARP_SIZE = 64 + NUM_WARPS = 4 + BLOCK_THREADS = NUM_WARPS * WARP_SIZE + + WMMA_N = 16 + WMMA_K = 32 + N_REPEAT = BV // WMMA_N + + NUM_H_ACCS = NUM_K_BLOCKS * N_REPEAT + + # HIP-ALIGNED (non-VWARP / split-M only): this fork uses ONLY the "wid owns + # 16 BT rows, all V" split-M warp partition -- the SAME warp scheme as the + # hand-tuned HIP/C++ K5 kernel (BT split-M in GEMM1, K split across waves in + # GEMM2, V not split across warps) with the SAME 16x16x16 bf16 MFMA. The + # alternative OPT-VWARP layout has been removed entirely from this fork. + + # HIP-ALIGN 2a: w panels (w_panel0/1), one [BT][64] panel per 64-K block, + # written with HIP's ``w_panel_swizzle`` (bank-conflict-free) and read by + # GEMM1 via ``load_a_w_fragment_swizzled`` (plain b64, contiguous 16-K). + LDS_WP_PANEL_ELEMS = BT * 64 + LDS_WP_ELEMS = NUM_K_BLOCKS * LDS_WP_PANEL_ELEMS + LDS_WP_BYTES = LDS_WP_ELEMS * 2 + + # HIP-ALIGN 2b: k in rotating-pair swizzled panels (one per 64-K block). + # Mirrors HIP's k_panel0[64*BT] + k_panel_rotating_pair_addr_bytes layout + # exactly: global load reads b128 (8 bf16) for adjacent token pairs (t0,t1), + # then scatters b16x2 packed writes to swizzled LDS addresses; GEMM2 reads + # the MFMA A frag via load_a_k_fragment_rotating (plain b64 from the + # swizzled address, no ds_read_tr16). Panel size = 64 K-rows * BT tokens + # = 64*BT bf16 elements per panel. + LDS_KP_PANEL_ELEMS = 64 * BT + LDS_KP_ELEMS = NUM_K_BLOCKS * LDS_KP_PANEL_ELEMS + LDS_KP_BYTES = LDS_KP_ELEMS * 2 + + # HIP-ALIGN 3: gated v_new in a shared2 panel (like h / HIP gated_v_panel); + # GEMM2 B read via load_shared2 (contiguous BT, matches the k A frag). + LDS_GV_ELEMS = (BT // 4) * BV * 4 # 16 bt_groups x BV cols x 4 BT + + # HIP-ALIGN phase 1b: h_state panels (shared2 layout) for the GEMM1 B + # operand. One [row_block][V] panel per 64-K block (like HIP's + # h_state_panel0/1); GEMM1 reads the MFMA B fragment with a plain b64 load + # (load_b_shared2) instead of the hardware-transpose ds_read_tr16. Each + # cell holds 4 K (a k_group): 64/4=16 row_blocks x BV cols x 4 K. + LDS_HP_PANEL_ELEMS = (64 // 4) * BV * 4 # = 64 * BV per 64-K block + LDS_HP_ELEMS = NUM_K_BLOCKS * LDS_HP_PANEL_ELEMS + LDS_HP_BYTES = LDS_HP_ELEMS * 2 + + # HIP-ALIGN phase 1a: [V][K] transpose buffer for the h snapshot HBM store. + # h_accs (f32) is written here in V-major / K-innermost order so 8 adjacent + # K land contiguously -> a single b128 (ds_read_b128 + buffer_store b128), + # matching HIP's ``h_transpose_buf`` + ``coalesced_vk_store_from_transpose`` + # (fewer store instructions than the per-element bf16 readout). K-contiguous + # (no pad) so the 8-wide vec load/store stays 16 B aligned. + LDS_HT_STRIDE = K + LDS_HT_ELEMS = BV * LDS_HT_STRIDE + LDS_HT_BYTES = LDS_HT_ELEMS * 2 + + # Bump revision so the FlyDSL JIT disk cache (~/.flydsl/cache/) invalidates + # on revision change (port of FlyDSL commit d4643e0e). + _K5_KERNEL_REVISION = ( + 122 # fp32->bf16 输出改用截断(_BF16_CONVERT_TRUNC)对齐 hip float_to_bf16 + ) + + GPU_ARCH = get_rocm_arch() + allocator = SmemAllocator( + None, + arch=GPU_ARCH, + global_sym_name=f"gdn_h_mfma16_hip_smem_v{_K5_KERNEL_REVISION}", + ) + lds_wp_offset = allocator._align(allocator.ptr, 16) + allocator.ptr = lds_wp_offset + LDS_WP_BYTES + lds_kp_offset = allocator._align(allocator.ptr, 16) + allocator.ptr = lds_kp_offset + LDS_KP_BYTES + # HIP-ALIGN 1b: h_state panels (GEMM1 B operand, shared2 layout). + lds_hp_offset = allocator._align(allocator.ptr, 16) + allocator.ptr = lds_hp_offset + LDS_HP_BYTES + # HIP-ALIGN 3: gated_v ALIASES h_state_panel1 (the kb=1 h_state panel), like + # HIP's ``gated_v_panel = h_state_panel1``. LDS_GV_ELEMS == LDS_HP_PANEL_ELEMS + # (both 64*BV) so it fits exactly, and no separate buffer is allocated + # (saves LDS_GV_BYTES). Correct because gated_v is written only AFTER GEMM1 + # has finished reading the h_state panels -- a WAR barrier is inserted right + # before the gated_v store to enforce that ordering across warps. + lds_gv_offset = lds_hp_offset + LDS_HP_PANEL_ELEMS * 2 # panel 1 (byte offset) + # HIP-ALIGN phase 1a: [V][K] transpose buffer for the b128 h store. + lds_ht_offset = allocator._align(allocator.ptr, 16) + allocator.ptr = lds_ht_offset + LDS_HT_BYTES + + # Cooperative load parameters + LOAD_VEC_WIDTH = 8 # 8 bf16 = 16 bytes = buffer_load_dwordx4 + THREADS_PER_ROW_64 = 64 // LOAD_VEC_WIDTH # 8 + ROWS_PER_BATCH_64 = BLOCK_THREADS // THREADS_PER_ROW_64 # 32 + NUM_LOAD_BATCHES_64 = BT // ROWS_PER_BATCH_64 # 2 + + # ---- OPT-VC: precompute the GEMM1 prefetch interleaving schedule. + # All quantities here depend ONLY on compile-time constants + # (K, BV, USE_G, USE_GK) and live in the outer compile_*-function + # scope so they are pure Python ints/lists -- the FlyDSL AST rewriter + # only touches the @flyc.kernel body below, so any control flow here + # is safe to mix as ordinary Python. + # OPT-VC enablement gate: only spread prefetch into GEMM1 when N_REPEAT + # == 1 (i.e. BV == WMMA_N == 16). When OPT_VC_ENABLED is False (BV>=32), + # emit all g/gk/u prefetch in a BATCH BEFORE GEMM1 starts (via + # PROLOGUE_EMITTER_CT), exactly matching the pre-OPT-VC (rev5) layout -- + # this leaves the full GEMM1 MFMA chain to overlap the HBM latency. + # An earlier attempt (rev21) routed disabled-BV prefetch to the GEMM1 + # tail (TAIL_EMITTER_CT), which empirically lost 9-14% on BV>=32 shapes + # because the prefetched values had no MFMA to hide behind before being + # consumed by the gating / vn = u - bv computation. + K_STEPS_PER_BLOCK = 64 // WMMA_K + OPT_VC_ENABLED = N_REPEAT == 1 + # OPT-W is gated together with OPT-VC. On BV>=32 (N_REPEAT>=2) the GEMM2 + # inner loop is also thin enough that interleaving w_next vec_loads into + # it causes the SIMD's single VMEM port to bottleneck on certain varlen + # shapes. Disabling the interleave on BV>=32 falls back to the rev5-style + # batched issue right before GEMM2, where the full MFMA chain hides the + # HBM latency. + NUM_INNER_SLOTS = NUM_K_BLOCKS * K_STEPS_PER_BLOCK * N_REPEAT + NUM_GK_LOADS_CT = (NUM_K_BLOCKS * 4) if USE_GK else 0 + # g_last + g_row are batched through the emitter queue (5 loads). + NUM_G_LOADS_CT = 5 if USE_G else 0 + NUM_U_LOADS_CT = N_REPEAT * 4 + _NUM_U_QUEUE_CT = NUM_U_LOADS_CT + NUM_EXTRA_LOADS_CT = NUM_GK_LOADS_CT + NUM_G_LOADS_CT + _NUM_U_QUEUE_CT + if OPT_VC_ENABLED and NUM_INNER_SLOTS > 0 and NUM_EXTRA_LOADS_CT > 0: + EXTRAS_PER_SLOT_CT = ( + NUM_EXTRA_LOADS_CT + NUM_INNER_SLOTS - 1 + ) // NUM_INNER_SLOTS + else: + EXTRAS_PER_SLOT_CT = 0 + # Map each emitter idx (0..NUM_EXTRA_LOADS_CT-1) to one of three buckets: + # * SLOT_ASSIGN_CT[slot_idx] -- emitted inside GEMM1 at (kb,ks,nr) slot + # (used when OPT_VC_ENABLED is True, BV=16 path) + # * PROLOGUE_EMITTER_CT -- emitted right BEFORE GEMM1 main loop + # (used when OPT_VC_ENABLED is False, BV>=32 path; matches rev5) + # * TAIL_EMITTER_CT -- emitted AFTER GEMM1 (kept as future- + # facing safety net; not used by the current schedule). + SLOT_ASSIGN_CT: list[list[int]] = [[] for _ in range(NUM_INNER_SLOTS)] + PROLOGUE_EMITTER_CT: list[int] = [] + for _e_idx in range(NUM_EXTRA_LOADS_CT): + if OPT_VC_ENABLED and NUM_INNER_SLOTS > 0: + _slot = min(_e_idx // max(EXTRAS_PER_SLOT_CT, 1), NUM_INNER_SLOTS - 1) + SLOT_ASSIGN_CT[_slot].append(_e_idx) + else: + PROLOGUE_EMITTER_CT.append(_e_idx) + + @flyc.kernel(name="chunk_gdn_fwd_h_flydsl_mfma16_hip") + def gdn_h_kernel( + k_tensor: fx.Pointer, + v_tensor: fx.Pointer, + w_tensor: fx.Pointer, + v_new_tensor: fx.Pointer, + g_tensor: fx.Pointer, + gk_tensor: fx.Pointer, + h_tensor: fx.Pointer, + h0_tensor: fx.Pointer, + ht_tensor: fx.Pointer, + cu_seqlens_tensor: fx.Pointer, + chunk_offsets_tensor: fx.Pointer, + state_indices_tensor: fx.Pointer, + T_val: fx.Int32, + T_flat: fx.Int32, + N_val: fx.Int32, + ): + i_v = fx.block_idx.x + i_nh = fx.block_idx.y + i_n = i_nh // fx.Int32(H) + i_h = i_nh % fx.Int32(H) + + # Indexed state-pool gather: when USE_STATE_INDICES, the SSM state slot + # for sequence ``i_n`` is ``state_indices[i_n]`` (addressing a pool + # ``[pool_size, H, V, K]``) rather than ``i_n`` itself (dense + # ``[N, H, V, K]``). Only h0 (read) and ht (in-place write-back) use this + # slot; the per-chunk h snapshot stays dense (i_n-indexed). + if const_expr(USE_STATE_INDICES): + si_ = GTensor(state_indices_tensor, dtype=T.i32, shape=(-1,)) + state_n = si_[fx.Index(i_n)] + else: + state_n = i_n + state_nh = state_n * fx.Int32(H) + i_h + + tid = fx.thread_idx.x + wid = tid // fx.Int32(WARP_SIZE) + lane = tid % fx.Int32(WARP_SIZE) + + k_ = GTensor(k_tensor, dtype=T.bf16, shape=(-1,)) + v_ = GTensor(v_tensor, dtype=T.bf16, shape=(-1,)) + w_ = GTensor(w_tensor, dtype=T.bf16, shape=(-1,)) + h_ = GTensor(h_tensor, dtype=T.bf16, shape=(-1,)) + g_ = GTensor(g_tensor, dtype=T.f32, shape=(-1,)) + if const_expr(USE_GK): + gk_ = GTensor(gk_tensor, dtype=T.f32, shape=(-1,)) + + vn_ = GTensor(v_new_tensor, dtype=T.bf16, shape=(-1,)) + # SSM-state dtype is selected by the compile-time flag; ``T.f32`` / + # ``T.bf16`` must be evaluated *inside* the kernel body where an MLIR + # context is active (mirrors how ``gdr_decode.py`` resolves + # ``state_dtype_`` from inside its kernel function). + state_t = T.bf16 if STATE_DTYPE_BF16 else T.f32 + if const_expr(USE_INITIAL_STATE): + h0_ = GTensor(h0_tensor, dtype=state_t, shape=(-1,)) + if const_expr(STORE_FINAL_STATE): + ht_ = GTensor(ht_tensor, dtype=state_t, shape=(-1,)) + + if const_expr(IS_VARLEN): + cu_ = GTensor(cu_seqlens_tensor, dtype=T.i32, shape=(-1,)) + co_ = GTensor(chunk_offsets_tensor, dtype=T.i32, shape=(-1,)) + + # -- LDS views -- + lds_base_ptr = allocator.get_base() + + # w panels (bf16, HIP w_panel swizzle) -- separate from k + lds_wp_ptr = SmemPtr(lds_base_ptr, lds_wp_offset, T.bf16, shape=(LDS_WP_ELEMS,)) + lds_wp = STensor(lds_wp_ptr, dtype=T.bf16, shape=(LDS_WP_ELEMS,)) + lds_wp_memref = lds_wp_ptr.get() + + # k panels (bf16, rotating-pair swizzle) -- separate from w + lds_kp_ptr = SmemPtr(lds_base_ptr, lds_kp_offset, T.bf16, shape=(LDS_KP_ELEMS,)) + lds_kp = STensor(lds_kp_ptr, dtype=T.bf16, shape=(LDS_KP_ELEMS,)) + lds_kp_memref = lds_kp_ptr.get() + + # gated v_new (bf16, shared2) + lds_gv_ptr = SmemPtr(lds_base_ptr, lds_gv_offset, T.bf16, shape=(LDS_GV_ELEMS,)) + lds_gv = STensor(lds_gv_ptr, dtype=T.bf16, shape=(LDS_GV_ELEMS,)) + lds_gv_memref = lds_gv_ptr.get() + + # h_state panels (shared2 layout) for the GEMM1 B operand + lds_hp_ptr = SmemPtr(lds_base_ptr, lds_hp_offset, T.bf16, shape=(LDS_HP_ELEMS,)) + lds_hp = STensor(lds_hp_ptr, dtype=T.bf16, shape=(LDS_HP_ELEMS,)) + lds_hp_memref = lds_hp_ptr.get() + + # h snapshot transpose buffer [V][K] (bf16) for the b128 HBM store + lds_ht_ptr = SmemPtr(lds_base_ptr, lds_ht_offset, T.bf16, shape=(LDS_HT_ELEMS,)) + lds_ht = STensor(lds_ht_ptr, dtype=T.bf16, shape=(LDS_HT_ELEMS,)) + lds_ht_memref = lds_ht_ptr.get() + + # HIP-ALIGN 1b: plain b64 read of a shared2 panel cell (GEMM1 B frag). + v4bf16_hp_type = T.vec(4, T.bf16) + + def _lds_read_hp_bf16x4(elem_idx): + return vector.load_op(v4bf16_hp_type, lds_hp_memref, [fx.Index(elem_idx)]) + + # HIP-ALIGN 2b: rotating-pair swizzled reads of k panel / plain b64 gated_v. + def _lds_read_kp_bf16x4(elem_idx): + return vector.load_op(v4bf16_hp_type, lds_kp_memref, [fx.Index(elem_idx)]) + + def _lds_read_gv_bf16x4(elem_idx): + return vector.load_op(v4bf16_hp_type, lds_gv_memref, [fx.Index(elem_idx)]) + + # -- Cooperative load decomposition -- + load_row_in_batch = tid // fx.Int32(THREADS_PER_ROW_64) + load_col_base = (tid % fx.Int32(THREADS_PER_ROW_64)) * fx.Int32(LOAD_VEC_WIDTH) + + # -- LDS vector read helpers (generates ds_read_b128 for 8xbf16) -- + + # HIP-ALIGN 2a: w_panel swizzle (returns ELEMENT offset within a panel). + # Port of w_panel_swizzle_base_bytes >> 1 (all bf16 = 2 B). + # row_in_half = row & 31; col_group = col_base >> 3 + # tid = row_in_half*8 + col_group + # base_bytes = ((tid<<4)&4080) ^ (tid&120); if row&32: base|=4096 + def _w_panel_swz_elems(row, col_base): + row_in_half = row & fx.Int32(31) + col_group = col_base >> fx.Int32(3) + tid_like = row_in_half * fx.Int32(8) + col_group + base = ((tid_like << fx.Int32(4)) & fx.Int32(4080)) ^ ( + tid_like & fx.Int32(120) + ) + base = base | ((row & fx.Int32(32)) << fx.Int32(7)) # 32<<7 = 4096 + return base >> fx.Int32(1) # bytes -> bf16 elements + + v4bf16_w_type = T.vec(4, T.bf16) + + # HIP-ALIGN 2a: load_a_w_fragment_swizzled (contiguous 16-K A frag). + # Caller passes row = row_base + lane&15 and k0 = k_base + (lane>>4)*4; + # the ^4-elem (^8-byte) toggle picks the low/high 4 within an 8-group. + def _load_a_w_swizzled(panel_base_elems, row, k0): + col_base = k0 & fx.Int32(~7) + elem = _w_panel_swz_elems(row, col_base) ^ (k0 & fx.Int32(4)) + return vector.load_op( + v4bf16_w_type, + lds_wp_memref, + [fx.Index(panel_base_elems + elem)], + ) + + # HIP-ALIGN 2b: k_panel rotating-pair swizzle address computation. + # Port of HIP's k_panel_rotating_pair_base_bytes / _addr_bytes. + # Returns a BYTE offset within a panel; caller converts to element + # offset (>> 1) before indexing the bf16 memref. + def _k_panel_rotating_pair_addr_bytes(row, pair_col): + row_block = row >> fx.Int32(3) + row_in_block = row & fx.Int32(7) + # k_panel_rotating_pair_base_bytes(row_block, pair_col) + tid_like = (pair_col << fx.Int32(3)) | row_block + lane_1_2 = tid_like & fx.Int32(6) + base = lane_1_2 << fx.Int32(10) + low = (lane_1_2 << fx.Int32(2)) ^ ( + (tid_like & fx.Int32(0xF8)) >> fx.Int32(1) + ) + toggle = (tid_like & fx.Int32(1)) * fx.Int32(0x440) + low = low ^ toggle + base_bytes = base | low + return (base_bytes ^ (row_in_block << fx.Int32(3))) + ( + row_in_block << fx.Int32(7) + ) + + # HIP-ALIGN 2b: load_a_k_fragment_rotating -- GEMM2 A operand read. + # Mirrors HIP's load_a_k_fragment_rotating(base, row_base, t_base, lane). + def _load_a_k_rotating(panel_base_elems, row_base, t_base): + row = row_base + lane_n + t0 = t_base + lane_m_base * fx.Int32(4) + pair_col = t0 >> fx.Int32(1) + byte_off = _k_panel_rotating_pair_addr_bytes(row, pair_col) + elem_off = byte_off >> fx.Int32(1) + return vector.load_op( + v4bf16_w_type, + lds_kp_memref, + [fx.Index(panel_base_elems + elem_off)], + ) + + # -- Prologue: compute bos, T_local, NT, boh -- + if const_expr(IS_VARLEN): + bos = cu_[fx.Index(i_n)] + eos = cu_[fx.Index(i_n) + fx.Index(1)] + T_local = eos - bos + NT = (T_local + fx.Int32(BT - 1)) // fx.Int32(BT) + boh = co_[fx.Index(i_n)] + else: + bos = i_n * T_val + T_local = T_val + NT = (T_local + fx.Int32(BT - 1)) // fx.Int32(BT) + boh = i_n * NT + + # -- Base pointer offsets (element counts) -- + # h: [B, NT, H, V, K] (VK) -- base = (boh*H + i_h) * V * K + h_base = (boh * fx.Int32(H) + i_h) * fx.Int32(V * K) + stride_h = fx.Int32(H * V * K) + + # k: [B, T, Hg, K] -- base = (bos*Hg + i_h//(H//Hg)) * K + gqa_ratio = H // Hg + k_base = (bos * fx.Int32(Hg) + i_h // fx.Int32(gqa_ratio)) * fx.Int32(K) + stride_k = fx.Int32(Hg * K) + + if const_expr(WU_CONTIGUOUS): + if const_expr(IS_VARLEN): + v_base = (i_h * T_flat + bos) * fx.Int32(V) + w_base = (i_h * T_flat + bos) * fx.Int32(K) + else: + v_base = ((i_n * fx.Int32(H) + i_h) * T_flat) * fx.Int32(V) + w_base = ((i_n * fx.Int32(H) + i_h) * T_flat) * fx.Int32(K) + stride_v = fx.Int32(V) + stride_w = fx.Int32(K) + else: + v_base = (bos * fx.Int32(H) + i_h) * fx.Int32(V) + w_base = (bos * fx.Int32(H) + i_h) * fx.Int32(K) + stride_v = fx.Int32(H * V) + stride_w = fx.Int32(H * K) + + if const_expr(IS_VARLEN): + vn_base = (i_h * T_flat + bos) * fx.Int32(V) + else: + vn_base = ((i_n * fx.Int32(H) + i_h) * T_flat) * fx.Int32(V) + + if const_expr(USE_INITIAL_STATE): + h0_base = state_nh * fx.Int32(V * K) + if const_expr(STORE_FINAL_STATE): + ht_base = state_nh * fx.Int32(V * K) + + # -- MFMA lane mapping for 16x16 tiles -- + lane_n = lane % fx.Int32(16) + lane_m_base = lane // fx.Int32(16) + + # -- Initialize h accumulators -- + acc_zero = fx.full(4, 0.0, fx.Float32) + + # h_accs[kb][nr] = f32x4 accumulator for k-block kb, v-repeat nr + h_accs = [] + for _kb in range_constexpr(NUM_K_BLOCKS): + for _nr in range_constexpr(N_REPEAT): + h_accs.append(acc_zero) + + # -- Load initial state if provided -- + # OPT-F: 4 x scalar f32 load -> 1 x buffer_load_dwordx4 (16 B). + # h0 is [V, K] so K is innermost; 4 consecutive K positions are + # contiguous in memory -> a single vec_load(4) covers them. + if const_expr(USE_INITIAL_STATE): + for kb in range_constexpr(NUM_K_BLOCKS): + for slot in range_constexpr(N_REPEAT): + h0_col = i_v * fx.Int32(BV) + fx.Int32(slot * 16) + lane_n + h0_row_base = ( + fx.Int32(kb * 64) + + wid * fx.Int32(16) + + lane_m_base * fx.Int32(4) + ) + h0_off_base = h0_base + h0_col * fx.Int32(K) + h0_row_base + loaded_vec = h0_.vec_load((fx.Index(h0_off_base),), 4) + if const_expr(STATE_DTYPE_BF16): + loaded_vec = loaded_vec.extf(T.f32x4) + acc_idx = kb * N_REPEAT + slot + h_accs[acc_idx] = h_accs[acc_idx] + loaded_vec + + # -- HIP-aligned pipelined main chunk loop -- + # Prologue loads w/k for chunk 0 to LDS; each iteration prefetches + # NEXT chunk's w (during GEMM1) and k (after gating) into VGPRs, + # then writes them to LDS at GEMM2 end (mirrors HIP .cu:1125-1194). + GEMM1_PF_SPLIT = min(1, NUM_K_BLOCKS) + k_row_base_pf = (lane & fx.Int32(7)) * fx.Int32(8) + k_pair_col_pf = wid * fx.Int32(8) + (lane >> fx.Int32(3)) + k_t0_pf = k_pair_col_pf * fx.Int32(2) + k_t1_pf = k_t0_pf + fx.Int32(1) + + init_state = [_to_raw(v) for v in h_accs] + c_zero = fx.Index(0) + c_one = fx.Index(1) + nt_idx = fx.Index(NT) + + # -- PROLOGUE: load w/k for chunk 0 to LDS (HIP .cu:1125-1141) -- + _prol_it = fx.Int32(0) + for kb in range_constexpr(NUM_K_BLOCKS): + wp_panel_base = fx.Int32(kb * LDS_WP_PANEL_ELEMS) + for batch in range_constexpr(NUM_LOAD_BATCHES_64): + row = fx.Int32(batch * ROWS_PER_BATCH_64) + load_row_in_batch + abs_row = _prol_it * fx.Int32(BT) + row + safe_row = (abs_row < T_local).select(abs_row, fx.Int32(0)) + w_g_off = ( + w_base + safe_row * stride_w + fx.Int32(kb * 64) + load_col_base + ) + wvec = w_.vec_load((fx.Index(w_g_off),), LOAD_VEC_WIDTH) + swz = wp_panel_base + _w_panel_swz_elems(row, load_col_base) + lds_wp.vec_store((fx.Index(swz),), wvec.shuffle(wvec, [0, 1, 2, 3]), 4) + lds_wp.vec_store( + (fx.Index(swz ^ fx.Int32(4)),), + wvec.shuffle(wvec, [4, 5, 6, 7]), + 4, + ) + k_abs_t0_prol = _prol_it * fx.Int32(BT) + k_t0_pf + k_abs_t1_prol = _prol_it * fx.Int32(BT) + k_t1_pf + k_safe_t0_prol = (k_abs_t0_prol < T_local).select(k_abs_t0_prol, fx.Int32(0)) + k_safe_t1_prol = (k_abs_t1_prol < T_local).select(k_abs_t1_prol, fx.Int32(0)) + for kb in range_constexpr(NUM_K_BLOCKS): + kp_pbase = fx.Int32(kb * LDS_KP_PANEL_ELEMS) + k_col_off = fx.Int32(kb * 64) + k_row_base_pf + k_g_off_t0 = k_base + k_safe_t0_prol * stride_k + k_col_off + k_g_off_t1 = k_base + k_safe_t1_prol * stride_k + k_col_off + kvec_t0 = k_.vec_load((fx.Index(k_g_off_t0),), LOAD_VEC_WIDTH) + kvec_t1 = k_.vec_load((fx.Index(k_g_off_t1),), LOAD_VEC_WIDTH) + for i in range_constexpr(LOAD_VEC_WIDTH): + row_i = k_row_base_pf + fx.Int32(i) + byte_off = _k_panel_rotating_pair_addr_bytes(row_i, k_pair_col_pf) + elem_off = byte_off >> fx.Int32(1) + lds_kp[fx.Index(kp_pbase + elem_off)] = kvec_t0[i] + lds_kp[fx.Index(kp_pbase + elem_off + fx.Int32(1))] = kvec_t1[i] + gpu.barrier() + + for i_t, state in range(c_zero, nt_idx, c_one, init=init_state): + h_accs_in = list(state[:NUM_H_ACCS]) + i_t_i32 = fx.Int32(i_t) + + # Stage h_accs into (a) the h_state panels [row_block][V] for the + # GEMM1 B operand (HIP shared2 layout, plain b64 read) and (b) the + # [V][K] transpose buffer for the b128 HBM store. split-M mapping: + # wid -> K sub-tile, acc_j(nr) -> V-tile. + for kb in range_constexpr(NUM_K_BLOCKS): + for acc_j in range_constexpr(N_REPEAT): + acc_idx = kb * N_REPEAT + acc_j + acc_val = h_accs_in[acc_idx] + # acc_j == nr (V-tile); K sub-tile = wid*16. + hp_col = fx.Int32(acc_j * 16) + lane_n + + # HIP-ALIGN 1b: write the h_state panel cell (shared2). This + # lane owns k_group (row_block = wid*4+lane_m_base) at V-col + # = nr*16+lane_n; 4 warps together fill all 16 row_blocks. + hp_row_block = wid * fx.Int32(4) + lane_m_base + hp_cell = fx.Int32(kb * LDS_HP_PANEL_ELEMS) + ( + hp_row_block * fx.Int32(BV) + hp_col + ) * fx.Int32(4) + lds_hp.vec_store( + (fx.Index(hp_cell),), + _f32x4_to_bf16x4_rne(acc_val), + 4, + ) + + # HIP-ALIGN 1a: [V][K/4-group] transpose buffer, XOR-swizzled + # (k_group ^ (v & 0xF)) to break bank conflicts on the + # scatter write -- mirrors HIP ``h_transpose_buf_offset``. + # The 4 elem_i are one k_group (4 contiguous K) -> one b64. + ht_kg = fx.Int32(kb * 16) + wid * fx.Int32(4) + lane_m_base + ht_idx = ( + hp_col * fx.Int32(K // 4) + (ht_kg ^ (hp_col & fx.Int32(0xF))) + ) * fx.Int32(4) + lds_ht.vec_store( + (fx.Index(ht_idx),), + _f32x4_to_bf16x4_rne(acc_val), + 4, + ) + + # w/k for this chunk already in LDS (prologue or prev GEMM2 end). + gpu.barrier() + + # last_idx for the current chunk (gating). + next_chunk_end = (i_t_i32 + fx.Int32(1)) * fx.Int32(BT) + last_idx_raw = (next_chunk_end < T_local).select( + next_chunk_end, T_local + ) - fx.Int32(1) + + # >>> PREFETCH u + g BEFORE GEMM1: issue all HBM loads for u + # (N_REPEAT × 4 ushort) and g (4 rows + 1 g_last = 5 dword) now, + # so the full 64-MFMA GEMM1 chain hides their HBM latency. + # Without this, LLVM hoists the loads into the GEMM1 middle where + # only ~2 MFMA can hide them → 4.2M cycle vmcnt stall (39% of stalls). + u_prefetch = [] # N_REPEAT × 4 bf16 scalars + for idx in range_constexpr(N_REPEAT): + u_col_pf = i_v * fx.Int32(BV) + fx.Int32(idx * 16) + lane_n + for elem_i in range_constexpr(4): + u_bt_row_raw = ( + i_t_i32 * fx.Int32(BT) + + wid * fx.Int32(16) + + lane_m_base * fx.Int32(4) + + fx.Int32(elem_i) + ) + safe_u_row = (u_bt_row_raw < T_local).select( + u_bt_row_raw, fx.Int32(0) + ) + u_off = v_base + safe_u_row * stride_v + u_col_pf + u_prefetch.append(v_[fx.Index(u_off)]) + + if const_expr(USE_G): + g_last = g_[fx.Index(i_h * T_flat + (bos + last_idx_raw))] + g_row_pf = [] + for elem_i in range_constexpr(4): + abs_row = ( + i_t_i32 * fx.Int32(BT) + + wid * fx.Int32(16) + + lane_m_base * fx.Int32(4) + + fx.Int32(elem_i) + ) + in_bounds = abs_row < T_local + safe_row = in_bounds.select(abs_row, fx.Int32(0)) + g_row_pf.append( + (g_[fx.Index(i_h * T_flat + (bos + safe_row))], in_bounds) + ) + + # -- GEMM1: b_v = w @ h_state, with w_next prefetch interleaved. + # u/g/w_next loads are all in flight; 64 MFMA hide HBM latency. + bv_accs = [] + for _i in range_constexpr(N_REPEAT): + bv_accs.append(fx.full(4, 0.0, fx.Float32)) + + # GEMM1 first K-block(s) -- before w prefetch. + for kb in range_constexpr(GEMM1_PF_SPLIT): + for ks in range_constexpr(K_STEPS_PER_BLOCK): + wp_pbase = fx.Int32(kb * LDS_WP_PANEL_ELEMS) + a_row = wid * fx.Int32(16) + lane_n + a_frag_lo = _load_a_w_swizzled( + wp_pbase, + a_row, + fx.Int32(ks * WMMA_K) + lane_m_base * fx.Int32(4), + ) + a_frag_hi = _load_a_w_swizzled( + wp_pbase, + a_row, + fx.Int32(ks * WMMA_K + 16) + lane_m_base * fx.Int32(4), + ) + for nr in range_constexpr(N_REPEAT): + hp_base = fx.Int32(kb * LDS_HP_PANEL_ELEMS) + hp_col_b = fx.Int32(nr * 16) + lane_n + rb_lo = fx.Int32((ks * WMMA_K) >> 2) + lane_m_base + rb_hi = fx.Int32(((ks * WMMA_K) + 16) >> 2) + lane_m_base + idx_lo = hp_base + (rb_lo * fx.Int32(BV) + hp_col_b) * fx.Int32( + 4 + ) + idx_hi = hp_base + (rb_hi * fx.Int32(BV) + hp_col_b) * fx.Int32( + 4 + ) + b_frag_lo = _lds_read_hp_bf16x4(idx_lo) + b_frag_hi = _lds_read_hp_bf16x4(idx_hi) + bv_accs[nr] = _mfma_bf16_16x16x16( + a_frag_lo, b_frag_lo, bv_accs[nr] + ) + bv_accs[nr] = _mfma_bf16_16x16x16( + a_frag_hi, b_frag_hi, bv_accs[nr] + ) + # gfx942 调度：在每个 ks 边界插 sched_barrier(mask_mfma)，阻止 LLVM + # 把各 ks 的 b_frag ds_read 跨迭代 hoist 聚成一大簇（LDS 端口背压 + # 主因），但放行 MFMA 跨 ks 重叠以保留流水隐藏延迟。sched_barrier + # 只约束指令排布、不改地址/数值，正确性安全。 + if const_expr(SCHED_GFX942): + rocdl.sched_barrier(rocdl.mask_mfma) + + # >>> PREFETCH w_next: HBM loads for next chunk (HIP .cu:670-672). + next_i_t = i_t_i32 + fx.Int32(1) + w_next_vecs = [] + w_next_swz = [] + for kb in range_constexpr(NUM_K_BLOCKS): + wp_pb_next = fx.Int32(kb * LDS_WP_PANEL_ELEMS) + for batch in range_constexpr(NUM_LOAD_BATCHES_64): + row = fx.Int32(batch * ROWS_PER_BATCH_64) + load_row_in_batch + abs_row_next = next_i_t * fx.Int32(BT) + row + safe_row_next = (abs_row_next < T_local).select( + abs_row_next, fx.Int32(0) + ) + w_g_off_next = ( + w_base + + safe_row_next * stride_w + + fx.Int32(kb * 64) + + load_col_base + ) + w_next_vecs.append( + w_.vec_load((fx.Index(w_g_off_next),), LOAD_VEC_WIDTH) + ) + w_next_swz.append( + wp_pb_next + _w_panel_swz_elems(row, load_col_base) + ) + + # GEMM1 remaining K-blocks -- MFMA hides u/g/w_next HBM latency. + for kb in range_constexpr(GEMM1_PF_SPLIT, NUM_K_BLOCKS): + for ks in range_constexpr(K_STEPS_PER_BLOCK): + wp_pbase = fx.Int32(kb * LDS_WP_PANEL_ELEMS) + a_row = wid * fx.Int32(16) + lane_n + a_frag_lo = _load_a_w_swizzled( + wp_pbase, + a_row, + fx.Int32(ks * WMMA_K) + lane_m_base * fx.Int32(4), + ) + a_frag_hi = _load_a_w_swizzled( + wp_pbase, + a_row, + fx.Int32(ks * WMMA_K + 16) + lane_m_base * fx.Int32(4), + ) + for nr in range_constexpr(N_REPEAT): + hp_base = fx.Int32(kb * LDS_HP_PANEL_ELEMS) + hp_col_b = fx.Int32(nr * 16) + lane_n + rb_lo = fx.Int32((ks * WMMA_K) >> 2) + lane_m_base + rb_hi = fx.Int32(((ks * WMMA_K) + 16) >> 2) + lane_m_base + idx_lo = hp_base + (rb_lo * fx.Int32(BV) + hp_col_b) * fx.Int32( + 4 + ) + idx_hi = hp_base + (rb_hi * fx.Int32(BV) + hp_col_b) * fx.Int32( + 4 + ) + b_frag_lo = _lds_read_hp_bf16x4(idx_lo) + b_frag_hi = _lds_read_hp_bf16x4(idx_hi) + bv_accs[nr] = _mfma_bf16_16x16x16( + a_frag_lo, b_frag_lo, bv_accs[nr] + ) + bv_accs[nr] = _mfma_bf16_16x16x16( + a_frag_hi, b_frag_hi, bv_accs[nr] + ) + # gfx942 调度：同上，切断 remaining K-block 跨 ks 的 ds_read 聚簇，放行 MFMA。 + if const_expr(SCHED_GFX942): + rocdl.sched_barrier(rocdl.mask_mfma) + + # WAR barrier (.cu:692). + gpu.barrier() + + # -- FUSED v_new + gating + gated_v store -- + # Consume prefetched u/g (already in VGPRs from before GEMM1). + if const_expr(USE_G): + exp_g_last = _fast_exp(g_last) + gate_elems = [] + for elem_i in range_constexpr(4): + g_row_val, in_bounds = g_row_pf[elem_i] + gate = _fast_exp(g_last - g_row_val) + gate_elems.append(in_bounds.select(gate, fx.Float32(0.0))) + gate_vec = vector.from_elements(T.f32x4, gate_elems) + + if const_expr(SAVE_NEW_VALUE): + + def _emit_vn_store(off, value): + vn_[fx.Index(off)] = value + + for idx in range_constexpr(N_REPEAT): + bv_val = bv_accs[idx] + u_col = i_v * fx.Int32(BV) + fx.Int32(idx * 16) + lane_n + u_f32_elems = [] + for elem_i in range_constexpr(4): + u_raw = u_prefetch[idx * 4 + elem_i] + u_bf16 = fx.BFloat16(u_raw) + u_f32_elems.append(u_bf16.to(fx.Float32)) + u_f32 = vector.from_elements(T.f32x4, u_f32_elems) + vn_val = u_f32 - bv_val + + if const_expr(SAVE_NEW_VALUE): + vn_bf16 = fx.Vector(_f32x4_to_bf16x4_rne(vn_val), (4,), fx.BFloat16) + bt_tile_base = wid * fx.Int32(16) + for elem_i in range_constexpr(4): + vn_bt_row = ( + i_t_i32 * fx.Int32(BT) + + bt_tile_base + + lane_m_base * fx.Int32(4) + + fx.Int32(elem_i) + ) + if (vn_bt_row < T_local).ir_value(): + bf16_v = vn_bf16[elem_i] + vn_off = vn_base + vn_bt_row * fx.Int32(V) + u_col + _emit_vn_store(vn_off, bf16_v) + + if const_expr(USE_G): + gated_val = vn_val * gate_vec + else: + gated_val = vn_val + gv_col = fx.Int32(idx * 16) + lane_n + gv_row_block = wid * fx.Int32(4) + lane_m_base + gv_cell = (gv_row_block * fx.Int32(BV) + gv_col) * fx.Int32(4) + lds_gv.vec_store( + (fx.Index(gv_cell),), + _f32x4_to_bf16x4_rne(gated_val), + 4, + ) + + # >>> PREFETCH k_next: HBM loads for next chunk (HIP .cu:727-731). + # Overlaps with the barrier + h store below. + k_abs_t0_next = next_i_t * fx.Int32(BT) + k_t0_pf + k_abs_t1_next = next_i_t * fx.Int32(BT) + k_t1_pf + k_safe_t0_next = (k_abs_t0_next < T_local).select( + k_abs_t0_next, fx.Int32(0) + ) + k_safe_t1_next = (k_abs_t1_next < T_local).select( + k_abs_t1_next, fx.Int32(0) + ) + k_next_vecs_t0 = [] + k_next_vecs_t1 = [] + for kb in range_constexpr(NUM_K_BLOCKS): + k_col_off_pf = fx.Int32(kb * 64) + k_row_base_pf + k_next_vecs_t0.append( + k_.vec_load( + (fx.Index(k_base + k_safe_t0_next * stride_k + k_col_off_pf),), + LOAD_VEC_WIDTH, + ) + ) + k_next_vecs_t1.append( + k_.vec_load( + (fx.Index(k_base + k_safe_t1_next * stride_k + k_col_off_pf),), + LOAD_VEC_WIDTH, + ) + ) + + # Apply exp(g_last) decay to h_accs (scalar broadcast). + if const_expr(USE_G): + exp_g_last_s = fx.Float32(exp_g_last) + for kb in range_constexpr(NUM_K_BLOCKS): + for nr in range_constexpr(N_REPEAT): + acc_idx = kb * N_REPEAT + nr + h_accs_in[acc_idx] = h_accs_in[acc_idx] * exp_g_last_s + + # Per-K decay: h[v, k] *= exp(gk_last[k]) at chunk end. + if const_expr(USE_GK): + gk_chunk_base = (bos + last_idx_raw) * fx.Int32(H * K) + i_h * fx.Int32( + K + ) + for kb in range_constexpr(NUM_K_BLOCKS): + gk_elems = [] + for elem_i in range_constexpr(4): + global_k = ( + fx.Int32(kb * 64) + + wid * fx.Int32(16) + + lane_m_base * fx.Int32(4) + + fx.Int32(elem_i) + ) + gk_raw = gk_[fx.Index(gk_chunk_base + global_k)] + gk_elems.append(_fast_exp(gk_raw)) + gk_vec = vector.from_elements(T.f32x4, gk_elems) + for nr in range_constexpr(N_REPEAT): + acc_idx = kb * N_REPEAT + nr + h_accs_in[acc_idx] = h_accs_in[acc_idx] * gk_vec + + gpu.barrier() + + # -- 5. h store from XOR-swizzled transpose buffer (HIP-aligned: + # after GEMM1, before GEMM2 -- mirrors HIP's + # coalesced_vk_store_from_transpose called between run_gemm1 and + # run_gemm2). The transpose buffer was populated during staging + # above and is read-only hereafter; gated_v was written to a + # different LDS region, so no conflict. + K_VECS = K // LOAD_VEC_WIDTH + NUM_HT_VECS = BV * K_VECS + for vbase in range_constexpr(0, NUM_HT_VECS, BLOCK_THREADS): + vec_idx = fx.Int32(vbase) + tid + kv = vec_idx % fx.Int32(K_VECS) + v_loc = vec_idx // fx.Int32(K_VECS) + k8 = kv * fx.Int32(LOAD_VEC_WIDTH) + v_xor = v_loc & fx.Int32(0xF) + kg_lo = kv * fx.Int32(2) + kg_hi = kg_lo + fx.Int32(1) + off_lo = (v_loc * fx.Int32(K // 4) + (kg_lo ^ v_xor)) * fx.Int32(4) + off_hi = (v_loc * fx.Int32(K // 4) + (kg_hi ^ v_xor)) * fx.Int32(4) + val_lo = fx.Vector( + vector.load_op(v4bf16_w_type, lds_ht_memref, [fx.Index(off_lo)]), + (4,), + fx.BFloat16, + ) + val_hi = fx.Vector( + vector.load_op(v4bf16_w_type, lds_ht_memref, [fx.Index(off_hi)]), + (4,), + fx.BFloat16, + ) + vec8 = val_lo.shuffle(val_hi, [0, 1, 2, 3, 4, 5, 6, 7]) + v_global = i_v * fx.Int32(BV) + v_loc + h_off = h_base + i_t_i32 * stride_h + v_global * fx.Int32(K) + k8 + h_.vec_store((fx.Index(h_off),), vec8, LOAD_VEC_WIDTH) + + # -- 6. GEMM2: h += k^T @ v_new_gated (no w prefetch/interleave). + BT_STEPS = BT // WMMA_K + for kb in range_constexpr(NUM_K_BLOCKS): + for bt_s in range_constexpr(BT_STEPS): + # HIP-ALIGN 2b: A = k load_a_k_fragment_rotating. + # row_base = wid*16 within the panel; t_base lo/hi split. + kp_pbase = fx.Int32(kb * LDS_KP_PANEL_ELEMS) + k_a_lo = _load_a_k_rotating( + kp_pbase, wid * fx.Int32(16), fx.Int32(bt_s * WMMA_K) + ) + k_a_hi = _load_a_k_rotating( + kp_pbase, wid * fx.Int32(16), fx.Int32(bt_s * WMMA_K + 16) + ) + + # gated_v B: shared2 layout (unchanged). + gv_rb_lo = fx.Int32((bt_s * WMMA_K) >> 2) + lane_m_base + gv_rb_hi = fx.Int32(((bt_s * WMMA_K) + 16) >> 2) + lane_m_base + for nr in range_constexpr(N_REPEAT): + gv_col = fx.Int32(nr * 16) + lane_n + vn_b_lo = _lds_read_gv_bf16x4( + (gv_rb_lo * fx.Int32(BV) + gv_col) * fx.Int32(4) + ) + vn_b_hi = _lds_read_gv_bf16x4( + (gv_rb_hi * fx.Int32(BV) + gv_col) * fx.Int32(4) + ) + + acc_idx = kb * N_REPEAT + nr + h_accs_in[acc_idx] = _mfma_bf16_16x16x16( + k_a_lo, vn_b_lo, h_accs_in[acc_idx] + ) + h_accs_in[acc_idx] = _mfma_bf16_16x16x16( + k_a_hi, vn_b_hi, h_accs_in[acc_idx] + ) + + # >>> WRITE prefetched w_next/k_next to LDS for next iteration. + # Barrier ensures GEMM2 done reading old panels (HIP .cu:794-798). + # Closures hide lds_wp/lds_kp (STensor) from the FlyDSL AST + # rewriter which otherwise tries to yield them through scf.if. + # Closures hide lds_wp/lds_kp (STensor) from the FlyDSL AST + # rewriter which otherwise tries to yield them through scf.if. + def _emit_wp_vec_store(idx_tuple, val, width): + lds_wp.vec_store(idx_tuple, val, width) + + def _emit_kp_scalar_store(idx, val): + lds_kp[fx.Index(idx)] = val + + has_next = next_i_t * fx.Int32(BT) < T_local + if has_next.ir_value(): + gpu.barrier() + _pf_idx = 0 + for kb in range_constexpr(NUM_K_BLOCKS): + for batch in range_constexpr(NUM_LOAD_BATCHES_64): + wvec_pf = w_next_vecs[_pf_idx] + swz_pf = w_next_swz[_pf_idx] + _emit_wp_vec_store( + (fx.Index(swz_pf),), + wvec_pf.shuffle(wvec_pf, [0, 1, 2, 3]), + 4, + ) + _emit_wp_vec_store( + (fx.Index(swz_pf ^ fx.Int32(4)),), + wvec_pf.shuffle(wvec_pf, [4, 5, 6, 7]), + 4, + ) + _pf_idx += 1 + for kb in range_constexpr(NUM_K_BLOCKS): + kp_pbase = fx.Int32(kb * LDS_KP_PANEL_ELEMS) + kvec_t0_pf = k_next_vecs_t0[kb] + kvec_t1_pf = k_next_vecs_t1[kb] + for i in range_constexpr(LOAD_VEC_WIDTH): + row_i = k_row_base_pf + fx.Int32(i) + byte_off = _k_panel_rotating_pair_addr_bytes( + row_i, k_pair_col_pf + ) + elem_off = byte_off >> fx.Int32(1) + _emit_kp_scalar_store(kp_pbase + elem_off, kvec_t0_pf[i]) + _emit_kp_scalar_store( + kp_pbase + elem_off + fx.Int32(1), kvec_t1_pf[i] + ) + + results = yield [_to_raw(v) for v in h_accs_in] + + h_accs_final = list(results[:NUM_H_ACCS]) + + # -- Epilogue: store final state -- + # OPT-7: 4 x scalar f32 store -> 1 x buffer_store_dwordx4 (16 B). + # acc_val is already f32x4 with element i at K offset i -> vec_store + # directly (no extract + from_elements needed). + if const_expr(STORE_FINAL_STATE): + for kb in range_constexpr(NUM_K_BLOCKS): + for slot in range_constexpr(N_REPEAT): + acc_idx = kb * N_REPEAT + slot + acc_val = h_accs_final[acc_idx] + + ht_col = i_v * fx.Int32(BV) + fx.Int32(slot * 16) + lane_n + ht_row_base = ( + fx.Int32(kb * 64) + + wid * fx.Int32(16) + + lane_m_base * fx.Int32(4) + ) + ht_off_base = ht_base + ht_col * fx.Int32(K) + ht_row_base + if const_expr(STATE_DTYPE_BF16): + out_vec = _f32x4_to_bf16x4_rne(acc_val) + else: + out_vec = acc_val + ht_.vec_store((fx.Index(ht_off_base),), out_vec, 4) + + # -- Host launcher ------------------------------------------------------ + @flyc.jit + def launch_gdn_h( + k_tensor: fx.Pointer, + v_tensor: fx.Pointer, + w_tensor: fx.Pointer, + v_new_tensor: fx.Pointer, + g_tensor: fx.Pointer, + gk_tensor: fx.Pointer, + h_tensor: fx.Pointer, + h0_tensor: fx.Pointer, + ht_tensor: fx.Pointer, + cu_seqlens_tensor: fx.Pointer, + chunk_offsets_tensor: fx.Pointer, + state_indices_tensor: fx.Pointer, + T_val: fx.Int32, + T_flat: fx.Int32, + N_val: fx.Int32, + grid_v: fx.Int32, + grid_nh: fx.Int32, + stream: fx.Stream, + ): + allocator.finalized = False + ctx = CompilationContext.get_current() + with ir.InsertionPoint(ctx.gpu_module_body): + allocator.finalize() + + launcher = gdn_h_kernel( + k_tensor, + v_tensor, + w_tensor, + v_new_tensor, + g_tensor, + gk_tensor, + h_tensor, + h0_tensor, + ht_tensor, + cu_seqlens_tensor, + chunk_offsets_tensor, + state_indices_tensor, + T_val, + T_flat, + N_val, + ) + launcher.launch( + grid=(grid_v, grid_nh, 1), + block=(BLOCK_THREADS, 1, 1), + stream=stream, + ) + + return launch_gdn_h + + +# NOTE: The Python host wrapper, BV autotune, and kernel cache live in +# ``aiter.ops.flydsl.linear_attention_prefill_kernels`` to keep this module +# free of any ``torch`` / ``triton`` dependency (mirrors the layering used +# by ``aiter.ops.flydsl.kernels.gdr_decode``). + + +__all__ = [ + "compile_chunk_gated_delta_h_mfma16_hip", +] diff --git a/aiter/ops/flydsl/linear_attention_prefill_kernels.py b/aiter/ops/flydsl/linear_attention_prefill_kernels.py index 4564e9d1c9d..19af5911036 100644 --- a/aiter/ops/flydsl/linear_attention_prefill_kernels.py +++ b/aiter/ops/flydsl/linear_attention_prefill_kernels.py @@ -18,18 +18,30 @@ from __future__ import annotations +import functools import math - +import os + +# NOTE (mfma16_hip fork): ``flydsl.compiler`` / ``flydsl.expr`` / +# ``get_rocm_arch`` are imported here for the additive HIP-aligned fork below. +# They are side-effect-free (``flydsl`` is already a hard dependency of the +# baseline ``compile_chunk_gated_delta_h``) and do NOT raise on flydsl <0.2.0 -- +# the mfma16_hip-only ``>=0.2.0`` requirement (fx.Pointer ABI) is enforced +# lazily in ``_get_or_compile_mfma16_hip`` so the baseline path keeps its +# original ``>=0.1.8`` compatibility. +import flydsl.compiler as flyc +import flydsl.expr as fx import torch import triton +from flydsl.runtime.device import get_rocm_arch -from .kernels.chunk_gated_delta_h import compile_chunk_gated_delta_h -from .kernels.tensor_shim import _run_compiled from ..triton._triton_kernels.gated_delta_rule.utils import ( prepare_chunk_offsets, prepare_num_chunks, prepare_rebased_cu_seqlens, ) +from .kernels.chunk_gated_delta_h import compile_chunk_gated_delta_h +from .kernels.tensor_shim import _run_compiled # log2(e); g pre-scaled by this constant lets the kernel use exp2(g) in # place of exp(g) (matches the Triton VK / HIP K5 convention). @@ -38,6 +50,7 @@ __all__ = [ "chunk_gated_delta_rule_fwd_h_flydsl", + "chunk_gated_delta_rule_fwd_h_flydsl_mfma16_hip", ] @@ -532,3 +545,421 @@ def chunk_gated_delta_rule_fwd_h_flydsl( ) return h, v_new, final_state + + +# ========================================================================== +# mfma16_hip fork (additive) -- HIP-aligned FlyDSL K5 implementation. +# +# Everything below is self-contained and does NOT touch the baseline wrapper +# above: it has its own compiled-kernel cache, BV selection (reusing the hip +# K5 selector), pointer-ABI launch, and public entry point +# ``chunk_gated_delta_rule_fwd_h_flydsl_mfma16_hip``. The baseline path keeps +# its original behaviour / flydsl>=0.1.8 compatibility; the mfma16_hip fork's +# fx.Pointer ABI requires flydsl>=0.2.0, enforced lazily below. +# ========================================================================== + +# mfma16_hip fork requires the fx.Pointer argument ABI (flyc.from_c_void_p + +# PointerAdaptor fast dispatch), which only exists from flydsl 0.2.0. Enforced +# lazily (in ``_get_or_compile_mfma16_hip``) so importing this module and using +# the baseline wrapper keeps working on flydsl>=0.1.8. +_MFMA16_HIP_MIN_FLYDSL_VERSION = "0.2.0" + +# gfx942 gate: only the mfma16_hip fork toggles the gfx942 GEMM1 ds-scheduling +# (SCHED_GFX942). ``get_rocm_arch()`` may return a feature-suffixed string like +# ``gfx942:sramecc+:xnack-``; normalize before matching. +_IS_GFX942 = get_rocm_arch().split(":")[0].startswith("gfx942") + +_INT32_ATTR = "_flydsl_int32_view" +_PROLOGUE_ATTR = "_flydsl_prologue_cache" + + +def _as_ptr(t: torch.Tensor): + """Wrap a torch tensor as a flydsl ``Pointer`` argument (raw data ptr). + + Uses ``fx.Uint8`` element type: the mfma16_hip kernel re-types every slot + inside the body via ``GTensor(ptr, dtype=...)``, so the host-side element + type is irrelevant to codegen and only needs to be a valid 1-byte unit. + """ + return flyc.from_c_void_p(fx.Uint8, t.data_ptr()) + + +def _as_int32(t: torch.Tensor) -> torch.Tensor: + """Return an int32 narrowing of ``t``, cached on the tensor itself. + + ``t`` is expected to come from one of the ``@tensor_cache``-decorated + prologue helpers (so its identity is stable across forwards). The cached + int32 result lives as an attribute on ``t`` itself, keeping cache + invalidation trivially correct. + """ + if t.dtype == torch.int32: + return t + cached = getattr(t, _INT32_ATTR, None) + if cached is None: + cached = t.to(torch.int32) + try: + object.__setattr__(t, _INT32_ATTR, cached) + except (AttributeError, TypeError): + pass + return cached + + +def _resolve_prologue( + cu_seqlens: torch.Tensor, + BT: int, + num_decodes: int, + num_decode_tokens: int, +): + """Resolve the per-shape varlen prologue in one cached lookup. + + Collapses the three ``@tensor_cache``-decorated prologue helpers into a + single tuple attached to ``cu_seqlens`` (keyed by ``(BT, num_decodes, + num_decode_tokens)``), so repeat forwards on the same ``cu_seqlens`` tensor + are one ``getattr`` + one dict get. + + Returns ``(NT, chunk_offsets, kernel_cu_seqlens, N, min_seqlen)``. + """ + cache_key = (BT, num_decodes, num_decode_tokens) + cache = getattr(cu_seqlens, _PROLOGUE_ATTR, None) + if cache is None: + cache = {} + try: + object.__setattr__(cu_seqlens, _PROLOGUE_ATTR, cache) + except (AttributeError, TypeError): + cache = None + if cache is not None: + hit = cache.get(cache_key) + if hit is not None: + return hit + + chunk_offsets = prepare_chunk_offsets( + cu_seqlens, BT, num_decodes, num_decode_tokens + ) + NT = prepare_num_chunks(cu_seqlens, BT, num_decodes, num_decode_tokens) + kernel_cu_seqlens = prepare_rebased_cu_seqlens( + cu_seqlens, num_decodes, num_decode_tokens + ) + N = len(kernel_cu_seqlens) - 1 + if N >= 1: + seg_lens = kernel_cu_seqlens[1:] - kernel_cu_seqlens[:-1] + min_seqlen = int(seg_lens.min().item()) + else: + min_seqlen = None + result = (NT, chunk_offsets, kernel_cu_seqlens, N, min_seqlen) + if cache is not None: + cache[cache_key] = result + return result + + +def _resolve_state_dtype(initial_state, state_dtype): + """Resolve/validate the SSM state dtype (float32 or bfloat16).""" + if initial_state is not None: + resolved = initial_state.dtype + if state_dtype is not None and state_dtype != resolved: + raise ValueError( + f"state_dtype={state_dtype} conflicts with " + f"initial_state.dtype={initial_state.dtype}; pass them " + f"consistently or omit state_dtype." + ) + elif state_dtype is not None: + resolved = state_dtype + else: + resolved = torch.float32 + if resolved not in (torch.float32, torch.bfloat16): + raise ValueError( + f"SSM state dtype must be float32 or bfloat16, got {resolved}." + ) + return resolved + + +@functools.cache +def _get_or_compile_mfma16_hip( + K, + V, + BT, + BV, + H, + Hg, + use_g, + use_gk, + use_h0, + store_fs, + save_vn, + is_varlen, + wu_contig, + state_bf16=False, + g_log2_scaled=False, + use_state_indices=False, + sched_gfx942=False, +): + """Compile (and cache) the mfma16 / HIP-aligned K5 kernel: 16x16x16 bf16 + MFMA + HIP-matching warp partition, writing the public VK layout [..., V, K]. + + ``use_state_indices`` compiles the indexed state-pool variant: the SSM + ``initial_state`` is a pool ``[pool_size, H, V, K]`` and each sequence's slot + is gathered from an ``initial_state_indices[N]`` int32 array (with in-place + final-state write-back into the same pool slot), mirroring the HIP kernel. + + The hip compile module + its flydsl>=0.2.0 requirement are imported lazily + here so the baseline path is unaffected. + """ + import flydsl + from packaging.version import Version + + installed = Version(getattr(flydsl, "__version__", "0").split("+")[0]) + if installed < Version(_MFMA16_HIP_MIN_FLYDSL_VERSION): + raise ImportError( + "FlyDSL K5 mfma16_hip fork requires `flydsl` " + f">=`{_MFMA16_HIP_MIN_FLYDSL_VERSION}` (for the fx.Pointer argument " + f"ABI), but got `{getattr(flydsl, '__version__', 'unknown')}`." + ) + + from .kernels.chunk_gated_delta_h_mfma16x16x16 import ( + compile_chunk_gated_delta_h_mfma16_hip, + ) + + return compile_chunk_gated_delta_h_mfma16_hip( + K=K, + V=V, + BT=BT, + BV=BV, + H=H, + Hg=Hg, + USE_G=use_g, + USE_GK=use_gk, + USE_INITIAL_STATE=use_h0, + STORE_FINAL_STATE=store_fs, + SAVE_NEW_VALUE=save_vn, + IS_VARLEN=is_varlen, + WU_CONTIGUOUS=wu_contig, + STATE_DTYPE_BF16=state_bf16, + G_IS_LOG2_SCALED=g_log2_scaled, + USE_STATE_INDICES=use_state_indices, + SCHED_GFX942=sched_gfx942, + ) + + +def chunk_gated_delta_rule_fwd_h_flydsl_mfma16_hip( + k: torch.Tensor, + w: torch.Tensor, + u: torch.Tensor, + g: torch.Tensor | None = None, + gk: torch.Tensor | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + chunk_size: int = 64, + save_new_value: bool = True, + cu_seqlens: torch.LongTensor | None = None, + state_dtype: torch.dtype | None = None, + use_exp2: bool = True, + num_decodes: int = 0, + num_decode_tokens: int = 0, + initial_state_indices: torch.Tensor | None = None, + inplace_final_state: bool | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: + """mfma16 / HIP-aligned K5 implementation: NON-VWARP only -- uses the + 16x16x16 bf16 MFMA and the SAME split-M warp partition (BT split-M, K split + across waves, V not split across warps) as the hand-tuned HIP/C++ K5 kernel, + writing the public VK layout [..., V, K]. API-compatible with + ``chunk_gated_delta_rule_fwd_h_flydsl`` (plus the indexed state-pool + contract via ``initial_state_indices`` / ``inplace_final_state``, matching + ``chunk_gated_delta_rule_fwd_h_hip_fn``). + + Unlike the baseline wrapper, BV is chosen by the hip K5 selector + (``_select_bv_for_varlen`` / ``_select_bv_for_dense``) so it matches + ``chunk_gated_delta_rule_fwd_h_hip_fn`` exactly; ``FLYDSL_K5_MFMA16HIP_BV`` + (in {16,32,64}) overrides it for A/B sweeps. + """ + use_g = g is not None + use_gk = gk is not None + use_h0 = initial_state is not None + g_log2_scaled = bool(use_exp2) + + # Indexed state-pool support: when ``initial_state_indices`` is given, + # ``initial_state`` is a pool ``[pool_size, H, V, K]`` and each sequence + # gathers its slot from the index array; the final state is written back + # in place into that same pool. ``inplace_final_state`` defaults to True + # whenever indices are given. + use_state_indices = initial_state_indices is not None + inplace = use_state_indices if inplace_final_state is None else inplace_final_state + if use_state_indices: + if initial_state is None: + raise ValueError( + "FlyDSL K5: initial_state_indices requires initial_state (the " + "state pool)." + ) + if not inplace: + raise ValueError( + "FlyDSL K5: initial_state_indices requires in-place final-state " + "write-back; leave inplace_final_state unset or set it to True." + ) + if not output_final_state: + raise ValueError( + "FlyDSL K5: initial_state_indices requires output_final_state=True " + "(the indexed path writes the final state back into the pool)." + ) + elif inplace and initial_state is None: + raise ValueError("FlyDSL K5: inplace_final_state requires initial_state.") + + resolved_state_dtype = _resolve_state_dtype(initial_state, state_dtype) + state_bf16 = resolved_state_dtype is torch.bfloat16 + + # mfma16_hip keeps the token-major [B, T_flat, Hg, K] k layout (no + # host-side pre-transpose), matching the Triton VK convention. + B, T, Hg, K = k.shape + H = w.shape[1] + V = u.shape[-1] + T_flat = w.shape[2] + BT = chunk_size + assert K <= 256 + + if cu_seqlens is None: + N = B + NT = triton.cdiv(T, BT) + chunk_offsets = None + kernel_cu_seqlens = None + is_varlen = False + else: + NT, chunk_offsets, kernel_cu_seqlens, N, _min_seqlen = _resolve_prologue( + cu_seqlens, BT, num_decodes, num_decode_tokens + ) + is_varlen = True + + # BV selection: reuse the hip K5 analytic selector so the fork picks the + # same BV as ``chunk_gated_delta_rule_fwd_h_hip_fn`` (lazy import avoids a + # potential circular import; both sides key on the same chunk_offsets). + from aiter.ops.chunk_gated_delta_rule_fwd_h import ( + _select_bv_for_dense as _hip_select_bv_for_dense, + ) + from aiter.ops.chunk_gated_delta_rule_fwd_h import ( + _select_bv_for_varlen as _hip_select_bv_for_varlen, + ) + + if is_varlen: + BV = _hip_select_bv_for_varlen(chunk_offsets, H) + else: + BV = _hip_select_bv_for_dense(B, T_flat, chunk_size, H, k.device) + + # Env override for A/B BV sweeps; the hand-tuned HIP K5 reference is fixed + # at BV=16 (FLYDSL_K5_MFMA16HIP_BV=16 reproduces it). + _bv_env = os.environ.get("FLYDSL_K5_MFMA16HIP_BV") + if _bv_env: + BV = int(_bv_env) + assert BV in (16, 32, 64), "mfma16_hip BV must be in {16,32,64}" + if V % BV != 0: + raise ValueError( + f"FlyDSL K5 mfma16_hip: requires V % BV == 0; got V={V}, BV={BV}." + ) + + # SCHED_GFX942 is only enabled on gfx942; other arches (incl. gfx950) pass + # False, keeping their emitted code byte-identical, and it joins the + # lru_cache key as a distinct compiled product. + launch_fn = _get_or_compile_mfma16_hip( + K, + V, + BT, + BV, + H, + Hg, + use_g, + use_gk, + use_h0, + output_final_state, + save_new_value, + is_varlen, + True, + state_bf16=state_bf16, + g_log2_scaled=g_log2_scaled, + use_state_indices=use_state_indices, + sched_gfx942=_IS_GFX942, + ) + + # Null-arg placeholder for the @flyc.jit slots ignored on this path. Sized + # 1 (not 0) so its ``data_ptr()`` is always a valid non-null device address. + dummy = torch.empty(1, device=k.device, dtype=torch.float32) + int32_dummy = dummy.to(torch.int32) if not is_varlen else None + cu_arg = ( + _as_int32(kernel_cu_seqlens) if kernel_cu_seqlens is not None else int32_dummy + ) + co_arg = _as_int32(chunk_offsets) if chunk_offsets is not None else int32_dummy + stream = torch.cuda.current_stream(k.device) + + grid_v = triton.cdiv(V, BV) + grid_nh = N * H + + # mfma16_hip writes the public VK layout ([..., V, K]) directly. + h_shape = (B, NT, H, V, K) + vn_shape = (B, H, T_flat, V) + vn_dtype = u.dtype + fs_shape = (N, H, V, K) if output_final_state else None + fs_dtype = resolved_state_dtype if output_final_state else None + save_vn = save_new_value + + # G contiguity guard (stride-1 along T; head-major [B, H, T_flat]). + if g is not None and not g.is_contiguous(): + raise AssertionError( + "FlyDSL K5: ``g`` must be contiguous (head-major [B, H, T_flat] " + f"or [H, T_flat]); got strides={g.stride()}, shape={tuple(g.shape)}." + ) + + # gk pre-scaling to log2 space (mirrors the Triton VK wrapper). + if gk is not None: + gk = gk.contiguous() + if g_log2_scaled: + gk = gk * _RCP_LN2 + + h = k.new_empty(h_shape) + v_new_buf = k.new_empty(vn_shape, dtype=vn_dtype) + if fs_shape is None: + final_state = None + elif inplace: + # In-place write-back: the final state aliases the ``initial_state`` + # buffer (the pool when indexed, or the dense [N,H,V,K] state + # otherwise), so no separate output tensor is allocated. + final_state = initial_state + else: + final_state = k.new_empty(fs_shape, dtype=fs_dtype) + + # The 11 tensor slots, wrapped as raw fx.Pointer args. Keep the torch + # tensors referenced as locals so the storage stays alive across the + # (synchronous) launch -- ``from_c_void_p`` only captures the data pointer. + tensor_args = tuple( + _as_ptr(t) + for t in ( + k, + u, + w, + v_new_buf, + g if g is not None else dummy, + gk if gk is not None else dummy, + h, + initial_state if initial_state is not None else dummy, + final_state if final_state is not None else dummy, + cu_arg, + co_arg, + ) + ) + + # The mfma16_hip kernel carries an extra ``state_indices`` slot (12th tensor + # arg): a real int32 [N] index array when indexed, else a 1-elem int32 dummy. + if use_state_indices: + si_i32 = initial_state_indices.to(torch.int32).contiguous() + if si_i32.numel() != N: + raise ValueError( + "FlyDSL K5: initial_state_indices length " + f"({si_i32.numel()}) must equal the number of sequences N={N}." + ) + else: + si_i32 = dummy.to(torch.int32) + tensor_args = tensor_args + (_as_ptr(si_i32),) + + launch_fn( + *tensor_args, + T, + T_flat, + N, + grid_v, + grid_nh, + stream, + ) + + return h, (v_new_buf if save_vn else None), final_state diff --git a/aiter/ops/triton/_triton_kernels/gated_delta_rule/prefill/chunk.py b/aiter/ops/triton/_triton_kernels/gated_delta_rule/prefill/chunk.py index 854018e9e3d..422482e878d 100644 --- a/aiter/ops/triton/_triton_kernels/gated_delta_rule/prefill/chunk.py +++ b/aiter/ops/triton/_triton_kernels/gated_delta_rule/prefill/chunk.py @@ -322,10 +322,10 @@ def chunk_gated_delta_rule_fwd_opt_vk( # head-major, so pass it through directly. The wrapper accepts # ``use_exp2`` as a kwarg and pre-scales ``gk`` internally. from aiter.ops.flydsl.linear_attention_prefill_kernels import ( - chunk_gated_delta_rule_fwd_h_flydsl, + chunk_gated_delta_rule_fwd_h_flydsl_mfma16_hip, ) - h, v_new, final_state = chunk_gated_delta_rule_fwd_h_flydsl( + h, v_new, final_state = chunk_gated_delta_rule_fwd_h_flydsl_mfma16_hip( k=k, w=w, u=u, diff --git a/op_tests/flydsl_tests/test_flydsl_linear_attention_prefill.py b/op_tests/flydsl_tests/test_flydsl_linear_attention_prefill.py index 31b6fb7b26e..b01c5bef7c2 100644 --- a/op_tests/flydsl_tests/test_flydsl_linear_attention_prefill.py +++ b/op_tests/flydsl_tests/test_flydsl_linear_attention_prefill.py @@ -4,18 +4,20 @@ """Unit tests for FlyDSL Linear Attention Prefill (chunk_gated_delta_h) regressions. Usage: - HIP_VISIBLE_DEVICES=7 pytest -sv aiter/ops/flydsl/test_flydsl_linear_attention_prefill.py::TestPerformance -s - HIP_VISIBLE_DEVICES=7 python -m pytest aiter/ops/flydsl/test_flydsl_linear_attention_prefill.py::TestPerformance -k "varlen-16k-aws" -v -s + rm -rf ~/.triton/cache + export GATED_DELTA_RULE_TRITON_AUTOTUNE=1 + FLYDSL_RUNTIME_ENABLE_CACHE=0 HIP_VISIBLE_DEVICES=7 pytest -sv op_tests/flydsl_tests/test_flydsl_linear_attention_prefill.py::TestPerformance -s + FLYDSL_RUNTIME_ENABLE_CACHE=0 HIP_VISIBLE_DEVICES=7 python -m pytest op_tests/flydsl_tests/test_flydsl_linear_attention_prefill.py::TestPerformance -k "varlen-64k-qwen-ptpc-ali" -v -s """ from __future__ import annotations +import math from dataclasses import dataclass import pytest import torch import triton -import triton.language as tl from torch.profiler import ProfilerActivity, profile from aiter.ops.flydsl.utils import is_flydsl_available @@ -30,7 +32,7 @@ try: from aiter.ops.flydsl.linear_attention_prefill_kernels import ( - chunk_gated_delta_rule_fwd_h_flydsl, + chunk_gated_delta_rule_fwd_h_flydsl_mfma16_hip, ) from aiter.ops.triton._triton_kernels.gated_delta_rule.prefill.chunk_delta_h import ( chunk_gated_delta_rule_fwd_h_opt_vk, @@ -47,10 +49,30 @@ ) _HAS_VLLM_K5 = True -except Exception: +except ImportError: chunk_gated_delta_rule_fwd_h_vllm = None _HAS_VLLM_K5 = False +# HIP/C++ K5 (chunk_gated_delta_rule_fwd_h.cu). JIT-compiled on first call. +# Same public VK outputs as the FlyDSL / Triton opt_vk backends, but it +# requires K=V=128 + bf16 inputs, so cases that violate that are skipped +# in the correctness test and excluded from the perf launch. +try: + from aiter.ops.chunk_gated_delta_rule_fwd_h import ( + chunk_gated_delta_rule_fwd_h_hip_fn, + ) + + _HAS_HIP_K5 = True +except ImportError: + chunk_gated_delta_rule_fwd_h_hip_fn = None + _HAS_HIP_K5 = False + +# When True, ``test_consistency_flydsl_mfma16_hip_vs_hip`` requires the +# flydsl-hip fork to match HIP/C++ BIT-FOR-BIT (torch.equal). Until the LDS / +# layout / (optionally) numeric alignment work fully lands it stays False and +# the test only records the gap + asserts a loose same-algorithm band. +_MFMA16_HIP_VS_HIP_BITEXACT = False + torch.set_default_device("cuda") @@ -88,6 +110,11 @@ class PrefillArgs: # shapes share the same ``(T, num_seqs)``. Typical values are a log # count or a hex digest of cu_seqlens. trace_tag: str = "" + # Appended to the display id when a group sweeps multiple + # ``max_num_batched_tokens`` values, so a fixed (tp, full_prompt_len) stays + # unique across the batched-token sweep. Empty for single-value groups, so + # their ids are unchanged. + bt_tag: str = "" @property def Hg(self): @@ -123,6 +150,8 @@ def __repr__(self): tag = self.model_name + "_" if self.model_name else "" tag += f"K{self.K}_V{self.V}_Hk{self.Hk}_Hv{self.Hv}" tag += f"_TP{self.tp}_T{self.full_prompt_len}" + if self.bt_tag: + tag += f"_{self.bt_tag}" if not self.is_varlen: tag += "_novarlen" if not self.output_final_state: @@ -167,7 +196,13 @@ class PrefillGroup: is_varlen: bool = True output_final_state: bool = True ssm_state_dtype: torch.dtype = torch.float32 - # Three semantics for ``max_num_batched_tokens``: + # Semantics for ``max_num_batched_tokens``: + # - list/tuple : sweep -- materialise one case per element (Cartesian with + # tps x full_prompt_lens). Each element is itself one of the specs + # below (int / "full_prompt_len" / None). For the varlen path this + # sweeps the batch size N = mnbt // full_prompt_len. ids get an + # ``mnbt{value}`` suffix so a fixed (tp, full_prompt_len) stays + # unique. Example: ``max_num_batched_tokens=[16384, 32768, 65536]``. # - int : use this exact value for every expanded case (e.g. you want # a fixed scheduler budget across a sweep of full_prompt_len). # - "full_prompt_len" : tie it to each case's full_prompt_len. The @@ -181,36 +216,134 @@ class PrefillGroup: # length 1024. Preserving that behavior is what keeps the # varlen path's per-case shape unchanged across this refactor. max_num_batched_tokens: object = None + # Optional "trace-derived 3-segment" expansion knob. When set, each + # expanded case overrides ``_build_context_lens`` with the explicit + # 3-segment layout ``[head, mid_seqlen, full_prompt_len - head - mid_seqlen]``, + # i.e. cu_seqlens = [0, head, head + mid_seqlen, full_prompt_len]. + # This reproduces the worst K5 regression family found in bench + # results 20260603 (n=3, T ~= 16384, middle segment == 10000): the + # K5 kernel exhibits a near-constant ~543us cost across this whole + # cluster regardless of head_seqlen, while triton K5 varies with the + # head split between ~460-495us. Sweeping head_seqlens lets us probe + # the kernel's sensitivity (or lack thereof) to the head boundary. + # Group is materialised as the (tps x full_prompt_lens x head_seqlens) + # Cartesian product when this is not None. + head_seqlens: object = None # list[int] | None + mid_seqlen: int = 10000 + # Number of segments per expanded case when ``head_seqlens`` is set: + # num_segments=3 (default): context_lens = [head, mid_seqlen, full_len-head-mid_seqlen] + # -> cu_seqlens = [0, head, head+mid_seqlen, full_len] (n=3) + # num_segments=2 : context_lens = [head, full_len-head] + # -> cu_seqlens = [0, head, full_len] (n=2) + # ``mid_seqlen`` is ignored in this mode; the tail length is whatever + # remains after ``head``. Used to cover the n=2 T=16384 regression + # clusters (head near 6400 / 8192 / 9912 / 10000) found in the + # bench_gdr 20260604 trace. + num_segments: int = 3 def expand_groups(groups): out = [] for g in groups: + # ``max_num_batched_tokens`` may be a single spec (int / "full_prompt_len" + # / None) OR a list/tuple of such specs. A list materialises one case per + # value (Cartesian with tps x full_prompt_lens) -- e.g. to sweep the + # scheduler token budget, which for the varlen path sweeps the batch size + # (N = mnbt // full_prompt_len). When more than one value is present, ids + # gain an ``mnbt{value}`` suffix so a fixed (tp, full_prompt_len) stays + # unique; a single value keeps the original ids unchanged. + mnbt_specs = g.max_num_batched_tokens + if not isinstance(mnbt_specs, (list, tuple)): + mnbt_specs = [mnbt_specs] + _sweep_mnbt = len(mnbt_specs) > 1 for tp in g.tps: for full_len in g.full_prompt_lens: - if g.max_num_batched_tokens == "full_prompt_len": - mnbt = full_len - elif g.max_num_batched_tokens is None: - mnbt = 32768 # PrefillArgs dataclass default - else: - mnbt = g.max_num_batched_tokens - out.append( - PrefillArgs( - K=g.K, - V=g.V, - Hk=g.Hk, - Hv=g.Hv, - tp=tp, - full_prompt_len=full_len, - model_name=g.model_name, - BT=g.BT, - max_num_batched_tokens=mnbt, - dtype=g.dtype, - is_varlen=g.is_varlen, - output_final_state=g.output_final_state, - ssm_state_dtype=g.ssm_state_dtype, - ) - ) + for mnbt_spec in mnbt_specs: + if mnbt_spec == "full_prompt_len": + mnbt = full_len + elif mnbt_spec is None: + mnbt = 32768 # PrefillArgs dataclass default + else: + mnbt = mnbt_spec + bt_tag = f"mnbt{mnbt}" if _sweep_mnbt else "" + + # head_seqlens=None : preserve the original "equal split via + # _build_context_lens" behavior. Otherwise materialise one + # PrefillArgs per (tp, full_len, head) triple with an + # explicit 3-segment cu_seqlens layout + # [head, mid_seqlen, full_len - head - mid_seqlen]. + if g.head_seqlens is None: + out.append( + PrefillArgs( + K=g.K, + V=g.V, + Hk=g.Hk, + Hv=g.Hv, + tp=tp, + full_prompt_len=full_len, + model_name=g.model_name, + BT=g.BT, + max_num_batched_tokens=mnbt, + dtype=g.dtype, + is_varlen=g.is_varlen, + output_final_state=g.output_final_state, + ssm_state_dtype=g.ssm_state_dtype, + bt_tag=bt_tag, + ) + ) + else: + for head in g.head_seqlens: + if g.num_segments == 2: + tail = full_len - head + if tail <= 0: + raise ValueError( + f"head_seqlens (num_segments=2) produced " + f"non-positive tail ({tail}) for " + f"group={g.model_name!r} " + f"full_prompt_len={full_len} head={head}." + ) + context_lens = [head, tail] + tag = f"head{head}_tail{tail}" + elif g.num_segments == 3: + tail = full_len - head - g.mid_seqlen + if tail <= 0: + raise ValueError( + f"head_seqlens (num_segments=3) produced " + f"non-positive tail ({tail}) for " + f"group={g.model_name!r} " + f"full_prompt_len={full_len} head={head} " + f"mid_seqlen={g.mid_seqlen}. Drop this " + f"(full_len, head) combo or raise " + f"full_prompt_len." + ) + context_lens = [head, g.mid_seqlen, tail] + tag = f"head{head}_mid{g.mid_seqlen}" + else: + raise ValueError( + f"num_segments={g.num_segments} unsupported; " + f"only 2 or 3 are implemented." + ) + if _sweep_mnbt: + tag = f"{tag}_mnbt{mnbt}" + out.append( + PrefillArgs( + K=g.K, + V=g.V, + Hk=g.Hk, + Hv=g.Hv, + tp=tp, + full_prompt_len=full_len, + model_name=g.model_name, + BT=g.BT, + max_num_batched_tokens=mnbt, + dtype=g.dtype, + is_varlen=g.is_varlen, + output_final_state=g.output_final_state, + ssm_state_dtype=g.ssm_state_dtype, + context_lens=context_lens, + trace_tag=tag, + ) + ) return out @@ -237,6 +370,14 @@ def expand_groups(groups): output_final_state=False, max_num_batched_tokens="full_prompt_len", ), + PrefillGroup( + model_name="Qwen3.5-397B-ptpc-ali", + Hv=64, + tps=[8], + full_prompt_lens=[1024, 2048, 4096, 8192], + is_varlen=False, + max_num_batched_tokens="full_prompt_len", + ), # varlen + final_state (default path), TP=4 / TP=8 share everything # else, so they collapse into a single group. Original rows left # max_num_batched_tokens at the PrefillArgs default of 32768, which @@ -250,11 +391,18 @@ def expand_groups(groups): full_prompt_lens=[1024, 2048, 4096, 8192], max_num_batched_tokens=32768, ), + PrefillGroup( + model_name="varlen-64k-qwen-ptpc-ali", + Hv=64, + tps=[8], + full_prompt_lens=[8192], + max_num_batched_tokens=[8192, 16384, 24576, 32768, 40960, 49152, 57344, 65536], + # max_num_batched_tokens=[65536], + ), PrefillGroup( model_name="varlen-16k-aws", Hv=32, tps=[1], - # full_prompt_lens=[1000, 2000, 3000, 4000, 5000, 6000, 7000, 8000, 9000, 10000], full_prompt_lens=[1000, 5000, 10000], max_num_batched_tokens=16384, ), @@ -262,39 +410,41 @@ def expand_groups(groups): model_name="varlen-32k-aws", Hv=32, tps=[1], - # full_prompt_lens=[1000, 2000, 3000, 4000, 5000, 6000, 7000, 8000, 9000, 10000], full_prompt_lens=[1000, 5000, 10000], + # full_prompt_lens=[1000], max_num_batched_tokens=32768, ), + PrefillGroup( + model_name="flydsl-k5-n1", + Hv=32, + tps=[1], + full_prompt_lens=[5000, 10000], + max_num_batched_tokens="full_prompt_len", + ), + PrefillGroup( + model_name="flydsl-k5-n3-mid10k", + Hv=32, + tps=[1], + full_prompt_lens=[16384], + max_num_batched_tokens=16384, + head_seqlens=[5, 10, 65, 704, 936, 1820, 4467, 5508], + mid_seqlen=10000, + ), + PrefillGroup( + model_name="flydsl-k5-n2-16k", + Hv=32, + tps=[1], + full_prompt_lens=[16384], + max_num_batched_tokens=16384, + head_seqlens=[4000, 6396, 8192, 9912, 10000], + num_segments=2, + ), ] PREFILL_PARAMS = expand_groups(_PREFILL_GROUPS) -# Perf-test parametrization is identical to the correctness one; trace- -# derived shapes have been removed from this file. -PERF_PARAMS = list(PREFILL_PARAMS) -PERF_TEST_IDS = [repr(p) for p in PERF_PARAMS] - - -# Mirror every base shape with a bf16-SSM-state variant. The bf16 vs f32 -# kernel paths only differ in two ``if const_expr`` branches: -# - h0 load (gated by USE_INITIAL_STATE) -# - ht store (gated by STORE_FINAL_STATE) -# The bf16 mirror keeps ``output_final_state`` from the base shape, so: -# - ``_nofs`` shapes (use_h0=True, store_fs=False) cover the h0 load path -# - default shapes (use_h0=True, store_fs=True) cover both paths -# Only ``(use_h0=False, store_fs=False)`` would generate IR identical to -# the f32 path; none of the current PREFILL_PARAMS hits that combo, so we -# do not filter here. If you add such a case later, gate the mirror with -# ``if _base.output_final_state or _make_inputs(...) provides h0``. -# NOTE: bf16 SSM-state mirrors disabled for focused perf profiling. -# PREFILL_PARAMS.extend( -# [ -# _dataclass_replace(_base, ssm_state_dtype=torch.bfloat16) -# for _base in list(PREFILL_PARAMS) -# ] -# ) +PREFILL_TEST_IDS = [repr(p) for p in PREFILL_PARAMS] # -- bf16 SSM-state params (paired with TestStateDtypeBF16 below) ------ @@ -517,18 +667,111 @@ def _normalize_opt_v_new(vn_opt): return vn_opt.permute(0, 2, 1, 3).contiguous() +def _is_gfx950() -> bool: + """Whether the current GPU is CDNA4 / gfx950 (MI350). + + The baseline / ``naive`` / ``naive_opt`` FlyDSL K5 forks emit the + ``mfma_f32_16x16x32_bf16`` (K=32 bf16) MFMA and ``mfma32_vk`` emits + ``mfma_f32_32x32x16_bf16`` -- both are gfx950-only instructions. On gfx942 + (CDNA3 / MI300) they fail to compile with an LLVM ``Cannot select`` + abort, so the perf harness skips them there. The remaining forks + (``kv`` / ``mfma16_hip`` / ``mfma16_2wave_opt1`` / ``mfma16_3wave_opt2``) + use the K=16 ``mfma_f32_16x16x16bf16_1k`` and run on both. + """ + try: + arch = torch.cuda.get_device_properties(0).gcnArchName + except Exception: # noqa: BLE001 + return False + return "gfx950" in arch + + +def _hip_k5_supported(args: PrefillArgs) -> bool: + """The HIP K5 kernel only handles K=V=128, bf16 inputs, chunk_size=64.""" + return ( + _HAS_HIP_K5 + and args.K == 128 + and args.V == 128 + and args.dtype == torch.bfloat16 + and args.BT == 64 + ) + + +def chunk_gated_delta_rule_fwd_h_hip_k5( + k, + w, + u, + g=None, + initial_state=None, + output_final_state=False, + cu_seqlens=None, +): + """HIP/C++ K5 host wrapper, adapted to this file's K5 calling convention. + + Mirrors the FlyDSL / Triton ``opt_vk`` backends: takes the GQA-layout + ``k`` ([B, T, Hg, K]), head-major ``w`` / ``u`` ([B, H, T, K/V]), and a + head-major cumulative-gate ``g`` ([H, T_total] or [B, H, T_total]) in + natural-log space, and returns VK-ordered ``h`` ([B, NT, H, V, K]), + head-major ``v_new`` ([B, H, T, V]), and VK ``final_state`` + ([N, H, V, K]) -- identical public outputs to the other backends, so + the shared ``_assert_k5_outputs_match_ref`` comparator applies directly. + + The underlying kernel's ``USE_EXP2`` path expects log2-space gates, so we + pass ``use_exp2=False`` here to keep the natural-log-space ``g`` contract + shared with the PyTorch reference (the kernel then applies the LOG2E + scale internally). + """ + H = w.shape[1] + T_flat = w.shape[2] + + # The HIP wrapper wants a 3-D head-major g [B, H, T_flat]. This file + # produces a 2-D [H, T_total] gate for the B=1 varlen / dense cases. + if g is not None: + if g.dim() == 2: + g_hip = g.reshape(1, H, T_flat).contiguous() + else: + g_hip = g.contiguous() + else: + g_hip = None + + return chunk_gated_delta_rule_fwd_h_hip_fn( + k, + w, + u, + g=g_hip, + initial_state=initial_state, + output_final_state=output_final_state, + chunk_size=64, + cu_seqlens=cu_seqlens, + use_exp2=False, + g_head_major=True, + ) + + # -- Performance benchmark ---------------------------------------------- _K5_KERNEL_PREFIXES = [ "chunk_gdn_fwd_h_flydsl_vk", + "chunk_gdn_fwd_h_flydsl_kv", + "chunk_gdn_fwd_h_flydsl_mfma16", + "chunk_gdn_fwd_h_flydsl_naive", "chunk_gated_delta_rule_fwd_kernel_h", ] +# The HIP/C++ K5 kernel is a templated __global__ whose profiler symbol is +# either the demangled ``...chunk_gated_delta_rule_fwd_h_hip_kernel<...>`` or +# a mangled ``_ZN...`` form. Match it as a substring (the templated name never +# appears at offset 0 after demangling because of the leading return type). +_K5_KERNEL_SUBSTRINGS = [ + "chunk_gated_delta_rule_fwd_h_hip_kernel", +] + def _is_k5_kernel(name: str) -> bool: """Return True if *name* is a K5 hidden-state recurrence kernel.""" - return any(name.startswith(p) for p in _K5_KERNEL_PREFIXES) + if any(name.startswith(p) for p in _K5_KERNEL_PREFIXES): + return True + return any(s in name for s in _K5_KERNEL_SUBSTRINGS) def _bench_fn(fn, *args, **kwargs): @@ -560,7 +803,69 @@ def _bench_fn(fn, *args, **kwargs): # -- Correctness tests --------------------------------------------------- -PREFILL_TEST_IDS = [repr(p) for p in PREFILL_PARAMS] + +def _assert_mean_abs_within(out, ref, *, mean_atol, label): + """Guard the *mean* absolute error, not just the per-element worst case. + + ``torch.testing.assert_close``'s ``atol`` only bounds the single worst + element, which for bf16 K5 over a long chunked recurrence sits close to + the 5e-2 element tolerance (a few outlier tokens at the tail of the + 33-segment chain). The mean abs error, by contrast, is ~2-3e-3 in + practice and is what actually moves when an implementation regresses the + *whole* distribution (e.g. a gating / accumulation bug) without yet + tripping any single element past 5e-2. Bound it here. + """ + mean_abs = (out.float() - ref.float()).abs().mean().item() + assert mean_abs <= mean_atol, ( + f"{label}: mean abs error {mean_abs:.3e} exceeds mean_atol " + f"{mean_atol:.3e} (per-element atol may still pass; this guards " + f"whole-distribution drift)" + ) + + +def _assert_close_lowmem(a, b, *, atol, rtol, msg, chunk_rows=1 << 22): + """Memory-frugal elementwise ``|a-b| <= atol + rtol*|b|`` check. + + Equivalent in semantics to ``torch.testing.assert_close(a, b, atol, rtol)`` + but streams over a flattened view in row chunks so it never materialises + more than one chunk-sized fp32 temporary. Used for the mfma16_2wave_opt1/triton_vk + consistency check, where the h / v_new tensors at long context are + multi-GiB and the stock ``assert_close`` (which up-casts both whole tensors + and builds a full mismatch report) OOMs on a 256 GiB card. On mismatch it + reports the worst element's abs error / allowed tol rather than dumping the + entire tensor. + """ + assert a.shape == b.shape, f"{msg}: shape {tuple(a.shape)} vs {tuple(b.shape)}" + af = a.reshape(-1) + bf = b.reshape(-1) + n = af.numel() + worst_abs = 0.0 + worst_allowed = 0.0 + worst_idx = -1 + n_bad = 0 + for s in range(0, n, chunk_rows): + e = min(s + chunk_rows, n) + ac = af[s:e].float() + bc = bf[s:e].float() + abs_e = (ac - bc).abs() + allowed = atol + rtol * bc.abs() + bad = abs_e > allowed + nb = int(bad.sum().item()) + if nb: + n_bad += nb + # track the single worst (abs - allowed) margin in this chunk + margin = abs_e - allowed + mi = int(margin.argmax().item()) + if abs_e[mi].item() - allowed[mi].item() > worst_abs - worst_allowed: + worst_abs = abs_e[mi].item() + worst_allowed = allowed[mi].item() + worst_idx = s + mi + del ac, bc, abs_e, allowed, bad + assert n_bad == 0, ( + f"{msg}: {n_bad}/{n} elements exceed atol={atol:g}+rtol={rtol:g}*|b|. " + f"Worst @ flat idx {worst_idx}: abs_err={worst_abs:.3e} > " + f"allowed={worst_allowed:.3e}." + ) def _assert_k5_outputs_match_ref( @@ -575,6 +880,7 @@ def _assert_k5_outputs_match_ref( label, atol=5e-2, rtol=5e-2, + mean_atol=5e-3, ): """Compare a K5 backend's outputs against the PyTorch FP32 reference. @@ -587,29 +893,46 @@ def _assert_k5_outputs_match_ref( f32-state is one ``truncf`` on the final_state, which stays well within bf16 ULP for sane inputs and never exceeds the historical f32-state margins. + + Two complementary bounds are enforced per output: + * ``atol`` / ``rtol`` (5e-2): the per-element worst case. + * ``mean_atol`` (5e-3): the mean abs error, which catches a regression + that shifts the whole distribution before any single element trips + the looser element tolerance. Measured mean abs is ~2-3e-3 on the + varlen-32k-aws shape, so 5e-3 leaves ~2x headroom over bf16 noise. """ + h_out_f = h_out.float() + vn_out_f = _normalize_opt_v_new(vn_out).float() torch.testing.assert_close( - h_out.float(), + h_out_f, h_ref.float(), atol=atol, rtol=rtol, msg=f"{label}: h mismatch", ) + _assert_mean_abs_within(h_out_f, h_ref, mean_atol=mean_atol, label=f"{label} h") torch.testing.assert_close( - _normalize_opt_v_new(vn_out).float(), + vn_out_f, vn_ref.float(), atol=atol, rtol=rtol, msg=f"{label}: v_new mismatch", ) + _assert_mean_abs_within( + vn_out_f, vn_ref, mean_atol=mean_atol, label=f"{label} v_new" + ) if output_final_state: + fs_out_f = fs_out.float() torch.testing.assert_close( - fs_out.float(), + fs_out_f, fs_ref.float(), atol=atol, rtol=rtol, msg=f"{label}: final_state mismatch", ) + _assert_mean_abs_within( + fs_out_f, fs_ref, mean_atol=mean_atol, label=f"{label} final_state" + ) else: assert fs_out is None, f"{label}: expected None final_state" assert fs_ref is None @@ -619,13 +942,16 @@ class TestCorrectness: """Correctness against PyTorch FP32 reference for all three K5 backends.""" @pytest.mark.parametrize("args", PREFILL_PARAMS, ids=PREFILL_TEST_IDS) - def test_correctness_flydsl(self, args: PrefillArgs): + def test_correctness_flydsl_mfma16_hip(self, args: PrefillArgs): + """mfma16 / HIP-aligned FlyDSL K5 impl (formerly the "vk" fork): 16x16x16 + MFMA + HIP warp partition. Same VK public outputs as the baseline flydsl + path; only the BV==64 configs exercise the kernel, others fall back.""" context_lens = args.resolve_context_lens() k, w_orig, u_orig, w_c, u_c, g, h0, cu, _ = _make_inputs( context_lens, args=args ) - h_fly, vn_fly, fs_fly = chunk_gated_delta_rule_fwd_h_flydsl( + h_fly, vn_fly, fs_fly = chunk_gated_delta_rule_fwd_h_flydsl_mfma16_hip( k, w_c, u_c, @@ -652,1062 +978,142 @@ def test_correctness_flydsl(self, args: PrefillArgs): vn_ref, fs_ref, output_final_state=args.output_final_state, - label="flydsl", - ) - - @pytest.mark.parametrize("args", PREFILL_PARAMS, ids=PREFILL_TEST_IDS) - def test_correctness_triton_vk(self, args: PrefillArgs): - """Triton VK K5 (h: [V, K]) -- same input/output layout as FlyDSL.""" - if args.ssm_state_dtype != torch.float32: - pytest.skip("Triton VK reference only supports f32 SSM state.") - context_lens = args.resolve_context_lens() - k, w_orig, u_orig, w_c, u_c, g, h0, cu, _ = _make_inputs( - context_lens, args=args - ) - - h_vk, vn_vk, fs_vk = chunk_gated_delta_rule_fwd_h_opt_vk( - k, - w_c, - u_c, - g=g, - initial_state=h0, - output_final_state=args.output_final_state, - cu_seqlens=cu, - ) - h_ref, vn_ref, fs_ref = ref_chunk_gated_delta_rule_fwd_h( - k, - w_orig, - u_orig, - g=g, - initial_state=h0, - output_final_state=args.output_final_state, - cu_seqlens=cu, - ) - - _assert_k5_outputs_match_ref( - h_vk, - vn_vk, - fs_vk, - h_ref, - vn_ref, - fs_ref, - output_final_state=args.output_final_state, - label="triton_vk", - ) - - @pytest.mark.skipif( - not _HAS_VLLM_K5, - reason="vllm.model_executor.layers.fla.ops.chunk_delta_h not importable", - ) - @pytest.mark.parametrize("args", PREFILL_PARAMS, ids=PREFILL_TEST_IDS) - def test_correctness_vllm(self, args: PrefillArgs): - """vLLM upstream K5 (h: [V, K]) -- same input/output layout as FlyDSL. - - vLLM's ``chunk_gated_delta_rule_fwd_h`` is the FLA upstream port that - powers ``vllm.model_executor.layers.fla.ops.chunk_gated_delta_rule``, - and shares the ``chunk_gated_delta_rule_fwd_kernel_h_blockdim64`` - kernel source with aiter's ``opt_vk``. We still cover it explicitly - to catch upstream version drift (e.g. signature, scaling factors, - or default chunk_size changes that would not show up in the aiter - ``opt_vk`` test above). - """ - if args.ssm_state_dtype != torch.float32: - pytest.skip("vLLM K5 reference only supports f32 SSM state.") - context_lens = args.resolve_context_lens() - k, w_orig, u_orig, w_c, u_c, g, h0, cu, _ = _make_inputs( - context_lens, args=args - ) - - # vLLM's host wrapper infers ``H = u.shape[-2]`` and ``T = k.shape[1]`` - # (T-major), so it needs the un-permuted ``w_orig`` / ``u_orig`` of - # shape ``[B, T, H, *]``, NOT the H-major ``w_c`` / ``u_c`` that - # aiter's ``opt_vk`` consumes. Feeding ``u_c`` would make vLLM think - # ``H = T`` and try to allocate ``(B, NT, T, V, K)`` for ``h`` - # (terabytes for long contexts). vLLM supports GQA natively, so we - # also do NOT repeat_interleave ``k``. - # vLLM's upstream K5 host wrapper indexes ``g`` as token-major - # ``[T_total, H]`` (FLA's original layout). Our test fixture now - # produces ``g`` in head-major ``[H, T_total]`` (matching the - # aiter / FlyDSL K5 convention), so transpose+contiguous it back - # to token-major before launching vLLM. - g_token_major = g.transpose(0, 1).contiguous() - h_vllm, vn_vllm, fs_vllm = chunk_gated_delta_rule_fwd_h_vllm( - k, - w_orig, - u_orig, - g=g_token_major, - initial_state=h0, - output_final_state=args.output_final_state, - cu_seqlens=cu, - ) - h_ref, vn_ref, fs_ref = ref_chunk_gated_delta_rule_fwd_h( - k, - w_orig, - u_orig, - g=g, - initial_state=h0, - output_final_state=args.output_final_state, - cu_seqlens=cu, - ) - - # vLLM's ``v_new = empty_like(u_orig)`` is already T-major [B, T, H, V], - # matching ``vn_ref`` directly -- we must NOT route through - # ``_assert_k5_outputs_match_ref`` because that helper calls - # ``_normalize_opt_v_new`` (permute 0,2,1,3) on the assumption that the - # K5 returned ``v_new`` in head-major [B, H, T, V] layout (aiter's - # convention). ``h`` and ``final_state`` follow the same vk layout as - # the aiter ``opt_vk`` path, so we compare them directly. - atol, rtol = 5e-2, 5e-2 - torch.testing.assert_close( - h_vllm.float(), - h_ref.float(), - atol=atol, - rtol=rtol, - msg="vllm: h mismatch", - ) - torch.testing.assert_close( - vn_vllm.float(), - vn_ref.float(), - atol=atol, - rtol=rtol, - msg="vllm: v_new mismatch", - ) - if args.output_final_state: - torch.testing.assert_close( - fs_vllm.float(), - fs_ref.float(), - atol=atol, - rtol=rtol, - msg="vllm: final_state mismatch", - ) - else: - assert fs_vllm is None, "vllm: expected None final_state" - - @pytest.mark.parametrize("args", PREFILL_PARAMS, ids=PREFILL_TEST_IDS) - def test_correctness_triton_origin_opt(self, args: PrefillArgs): - """triton_origin_opt K5: standalone fwd_h (BV=16 + exp2) variant. - - Same kernel as triton_origin but with BV=16 + USE_EXP2 (the - ``new pipeline`` config from ``0423_gdr_prefill_bench_standalone.py``). - Unlike triton_origin this K5 supports GQA natively (the kernel - takes ``Hg`` as a constexpr), so the GQA expand step is skipped. - """ - if args.ssm_state_dtype != torch.float32: - pytest.skip("triton_origin_opt reference only supports f32 SSM state.") - context_lens = args.resolve_context_lens() - k, w_orig, u_orig, w_c, u_c, g, h0, cu, _ = _make_inputs( - context_lens, args=args - ) - - # GQA-aware K5: no repeat_interleave needed. Hidden state in [K,V] - # layout (same as triton_origin), so use h0 transposed from the - # VK reference layout. - h0_kv = h0.transpose(-2, -1).contiguous() if h0 is not None else None - - h_origin_opt, vn_origin_opt, fs_origin_opt = ( - chunk_gated_delta_rule_fwd_h_origin_opt( - k, - w_orig, - u_orig, - g=g, - initial_state=h0_kv, - output_final_state=args.output_final_state, - cu_seqlens=cu, - ) - ) - h_ref, vn_ref, fs_ref = ref_chunk_gated_delta_rule_fwd_h( - k, - w_orig, - u_orig, - g=g, - initial_state=h0, - output_final_state=args.output_final_state, - cu_seqlens=cu, + label="flydsl_mfma16_hip", ) - h_origin_opt_vk = h_origin_opt.transpose(-2, -1).contiguous() - fs_origin_opt_vk = ( - fs_origin_opt.transpose(-2, -1).contiguous() - if fs_origin_opt is not None - else None - ) - - atol, rtol = 5e-2, 5e-2 - torch.testing.assert_close( - h_origin_opt_vk.float(), - h_ref.float(), - atol=atol, - rtol=rtol, - msg="triton_origin_opt: h mismatch", - ) - torch.testing.assert_close( - vn_origin_opt.float(), - vn_ref.float(), - atol=atol, - rtol=rtol, - msg="triton_origin_opt: v_new mismatch", - ) - if args.output_final_state: - torch.testing.assert_close( - fs_origin_opt_vk.float(), - fs_ref.float(), - atol=atol, - rtol=rtol, - msg="triton_origin_opt: final_state mismatch", - ) +# -- Performance benchmark (flydsl-hip vs hip vs triton) ----------------- _perf_results: list[dict] = [] def _run_perf_comparison(args: PrefillArgs): - """Per-shape K5 perf body used by ``TestPerformance``. - - Bench 4 backends (FlyDSL, Triton opt_vk, Triton origin_opt, vLLM) - under identical inputs and append a row to ``_perf_results``. The - summary table is printed once at session teardown by - ``_print_summary_table``. - """ + """Bench the same shape on flydsl-hip / hip(C++) / triton(opt_vk) and record + a row into ``_perf_results``; the session-scoped ``_print_summary_table`` + fixture prints an aligned table after all tests finish. hip/triton are + mainline backends used only as references; hip is skipped for shapes it does + not support (needs K=V=128, bf16, chunk_size=64).""" context_lens = args.resolve_context_lens() - k, w_orig, u_orig, w_c, u_c, g, h0, cu, _ = _make_inputs(context_lens, args=args) - total_tokens = sum(context_lens) - - # Triton K5 host wrappers only accept f32 ``initial_state`` and always - # produce an f32 ``final_state``. When FlyDSL is benched with a bf16 - # SSM state, we still want a Triton baseline for comparison, so we - # promote h0 to f32 once (outside the timed window) and feed it to - # the Triton closures. The resulting "Triton(f32) vs FlyDSL(bf16)" - # row answers the practical question "how much does enabling - # bf16-state win against the existing Triton baseline?". - h0_triton_vk = h0.float() if (h0 is not None and h0.dtype != torch.float32) else h0 - - # triton_origin_opt uses a [K, V] hidden-state layout, so its h0 - # is the VK reference h0 transposed on the last two dims. - h0_origin_kv = ( - h0_triton_vk.transpose(-2, -1).contiguous() - if h0_triton_vk is not None - else None + k, _w_orig, _u_orig, w_c, u_c, g, h0, cu, _ = _make_inputs(context_lens, args=args) + ofs = args.output_final_state + total_tokens = int(cu[-1].item()) + + us_fly = _bench_fn( + chunk_gated_delta_rule_fwd_h_flydsl_mfma16_hip, + k, + w_c, + u_c, + g=g, + initial_state=h0, + output_final_state=ofs, + cu_seqlens=cu, ) - - # K5 launch closures: each invokes the K5 host wrapper of its backend. - def flydsl_launch(): - chunk_gated_delta_rule_fwd_h_flydsl( - k=k, - w=w_c, - u=u_c, + us_tri = _bench_fn( + chunk_gated_delta_rule_fwd_h_opt_vk, + k, + w_c, + u_c, + g=g, + initial_state=h0, + output_final_state=ofs, + cu_seqlens=cu, + ) + if _HAS_HIP_K5 and _hip_k5_supported(args): + us_hip = _bench_fn( + chunk_gated_delta_rule_fwd_h_hip_k5, + k, + w_c, + u_c, g=g, initial_state=h0, - output_final_state=args.output_final_state, - cu_seqlens=cu, - ) - - def triton_vk_launch(): - chunk_gated_delta_rule_fwd_h_opt_vk( - k=k, - w=w_c, - u=u_c, - g=g, - initial_state=h0_triton_vk, - output_final_state=args.output_final_state, - cu_seqlens=cu, - ) - - # vLLM upstream K5 (chunk_delta_h.chunk_gated_delta_rule_fwd_h). Uses - # the same vk hidden-state layout and same final_state dtype as the - # aiter ``opt_vk`` wrapper, but unlike ``opt_vk`` its host wrapper - # infers ``H = u.shape[-2]`` (T-major), so it requires the un-permuted - # ``w_orig`` / ``u_orig`` -- feeding the H-major ``w_c`` / ``u_c`` - # would make vLLM compute ``H = T`` and try to allocate a - # terabyte-sized ``h``. vLLM supports GQA natively, so ``k`` is also - # passed un-expanded. We still feed the same h0_triton_vk (fp32) as - # ``triton_vk_launch`` because vk hidden-state layout is identical. - def vllm_launch(): - chunk_gated_delta_rule_fwd_h_vllm( - k=k, - w=w_orig, - u=u_orig, - g=g, - initial_state=h0_triton_vk, - output_final_state=args.output_final_state, + output_final_state=ofs, cu_seqlens=cu, ) - - def triton_origin_opt_launch(): - # GQA-aware (uses unexpanded k) BV=16 + exp2 variant of fwd_h - # from the standalone bench's new pipeline. Hidden state in - # [K,V] layout, so it needs h0_origin_kv (h0 transposed from - # the VK reference layout). - chunk_gated_delta_rule_fwd_h_origin_opt( - k=k, - w=w_orig, - u=u_orig, - g=g, - initial_state=h0_origin_kv, - output_final_state=args.output_final_state, - cu_seqlens=cu, - ) - - # Warmup FlyDSL once so its internal BV-autotune sweep does not - # leak into the timed window. Triton's own ``triton.autotune`` is - # already absorbed by ``_bench_fn``'s NUM_WARMUP=5 prelude, except - # for vLLM upstream's K5 -- its first call also runs a BV/warps/ - # stages sweep and the 5-iter prelude is not always enough to - # converge the autotuner on long-T shapes. Pre-warm it once here - # for parity with FlyDSL so that ``us_vllm`` reflects steady-state - # kernel time, not the autotune sweep. - flydsl_launch() - if _HAS_VLLM_K5: - vllm_launch() - torch.cuda.synchronize() - - us_fly = _bench_fn(flydsl_launch) - us_triton_vk = _bench_fn(triton_vk_launch) - us_triton_origin_opt = _bench_fn(triton_origin_opt_launch) - us_vllm = _bench_fn(vllm_launch) if _HAS_VLLM_K5 else float("nan") - - fly_vs_vk = us_triton_vk / us_fly if us_fly > 0 else float("inf") - fly_vs_origin_opt = us_triton_origin_opt / us_fly if us_fly > 0 else float("inf") - fly_vs_vllm = ( - us_vllm / us_fly if (us_fly > 0 and us_vllm == us_vllm) else float("nan") - ) - - # bench333 cases carry trace-derived structural features (head/mid/ - # tail seqlen, log_count); compute these once so the summary table - # can display them in a bench333-specific sub-table. Non-bench333 - # rows leave these fields at their defaults (None / 0). - head_seqlen = mid_seqlen = tail_seqlen = 0 - log_count = 0 - if args.context_lens is not None and args.model_name == "prefill-bench333": - lens = list(args.context_lens) - if len(lens) == 1: - head_seqlen, tail_seqlen = int(lens[0]), 0 - elif len(lens) == 2: - head_seqlen, tail_seqlen = int(lens[0]), int(lens[1]) - elif len(lens) >= 3: - mids = lens[1:-1] - head_seqlen = int(lens[0]) - tail_seqlen = int(lens[-1]) - mid_seqlen = int(mids[0]) if all(m == mids[0] for m in mids) else 0 - # trace_tag is set to "cnt{log_count}" by _build_bench407_params. - if args.trace_tag.startswith("cnt"): - try: - log_count = int(args.trace_tag[3:]) - except ValueError: - log_count = 0 - - # For bench333 trace shapes the model name is the same string for - # all 333 rows ("prefill-bench333") which is useless in the per-row - # summary table. Use the trailing "T{T}_n{N}_cnt{log_count}" suffix - # of the pytest id instead, so each row's Model cell uniquely - # identifies the case. - if args.model_name == "prefill-bench333" and args.context_lens is not None: - n_seqs = len(args.context_lens) - T_total = sum(args.context_lens) - model_label = f"T{T_total}_n{n_seqs}_{args.trace_tag}" else: - model_label = args.model_name or "-" + us_hip = float("nan") + has_hip = not math.isnan(us_hip) # not NaN _perf_results.append( { - "Model": model_label, + "Model": args.model_name or "-", "TP": args.tp, - "K": args.K, - "V": args.V, "Hg": args.Hg, "H": args.H, "SeqLen": args.full_prompt_len, "T": total_tokens, "varlen": args.is_varlen, - "final_st": args.output_final_state, - "state": "bf16" if args.ssm_state_dtype == torch.bfloat16 else "fp32", - # bench333-only fields (0 for main-table rows) - "T_prefill": total_tokens if args.model_name == "prefill-bench333" else 0, - "num_seqs_prefill": ( - len(args.context_lens) - if args.context_lens is not None - and args.model_name == "prefill-bench333" - else 0 - ), - "head_seqlen": head_seqlen, - "mid_seqlen": mid_seqlen, - "tail_seqlen": tail_seqlen, - "log_count": log_count, - # Perf columns (same for all rows) - "FlyDSL_vk(us)": us_fly, - "Triton_vk(us)": us_triton_vk, - "Triton_origin_opt(us)": us_triton_origin_opt, - "vLLM_vk(us)": us_vllm, - "flydsl_vs_vk": fly_vs_vk, - "flydsl_vs_origin_opt": fly_vs_origin_opt, - "flydsl_vs_vllm": fly_vs_vllm, + "final_st": ofs, + "fly_hip": us_fly, + "HIP": us_hip, + "Triton": us_tri, + # speedup vs hip (hip is the baseline): >1 faster than hip, <1 slower. + "fly/hip": (us_hip / us_fly) if has_hip else float("nan"), + "tri/hip": (us_hip / us_tri) if has_hip else float("nan"), } ) -class TestPerformance: - """Kernel-only performance comparison: FlyDSL vs Triton opt_vk vs Triton opt3_kv.""" - - @pytest.mark.parametrize("args", PERF_PARAMS, ids=PERF_TEST_IDS) - def test_perf_comparison(self, args: PrefillArgs): - _run_perf_comparison(args) - - def _print_perf_table(): if not _perf_results: return - - # Two column layouts: - # - main_cols : for hand-written PrefillGroup shapes (Qwen3.5-35B - # / 397B / varlen-fs / bench333-varlen) showing - # SeqLen + T + varlen + final_st. - # - bench_cols : for bench333 trace-derived shapes showing - # T_prefill + num_seqs_prefill + head/mid/tail - # + log_count instead. SeqLen / T are redundant - # here because T_prefill carries the same info. - perf_tail_cols = [ - ("FlyDSL", "FlyDSL_vk(us)", 8), - ("Tri_vk", "Triton_vk(us)", 8), - ("Tri_orig_opt", "Triton_origin_opt(us)", 12), - ("vLLM", "vLLM_vk(us)", 8), - ("fly/vk", "flydsl_vs_vk", 7), - ("fly/o_opt", "flydsl_vs_origin_opt", 9), - ("fly/vllm", "flydsl_vs_vllm", 8), + _model_w = max([len("Model")] + [len(str(r["Model"])) for r in _perf_results]) + # (header_display, row_key, width): header uses the 1st, cell lookup the 2nd. + cols = [ + ("Model", "Model", _model_w), + ("TP", "TP", 2), + ("Hg", "Hg", 2), + ("H", "H", 2), + ("SeqLen", "SeqLen", 6), + ("T", "T", 6), + ("varlen", "varlen", 6), + ("final_st", "final_st", 8), + ("FlyDSL_hip(us)", "fly_hip", 14), + ("HIP(us)", "HIP", 8), + ("Triton(us)", "Triton", 10), + ("fly/hip", "fly/hip", 7), + ("tri/hip", "tri/hip", 7), ] - main_cols = [ - ("Model", "Model", 16), - ("TP", "TP", 3), - ("Hg", "Hg", 3), - ("H", "H", 3), - ("SeqLen", "SeqLen", 7), - ("T", "T", 7), - ("var", "varlen", 3), - ("fs", "final_st", 3), - ] + perf_tail_cols - bench_cols = [ - ("Model", "Model", 22), - ("TP", "TP", 3), - ("Hg", "Hg", 3), - ("H", "H", 3), - ("T_prefill", "T_prefill", 9), - ("n_pref", "num_seqs_prefill", 6), - ("head", "head_seqlen", 6), - ("mid", "mid_seqlen", 6), - ("tail", "tail_seqlen", 6), - ("log_cnt", "log_count", 7), - ] + perf_tail_cols - - def _build_header_sep(cols): - header = " | ".join(display.rjust(width) for display, _, width in cols) - sep = "-+-".join("-" * width for _, _, width in cols) - return header, sep - - def _fmt_row(row, cols): - cells = [] - for display, key, width in cols: - val = row[key] - if isinstance(val, bool): - cells.append(("Y" if val else "N").rjust(width)) - elif isinstance(val, float): - if val != val: # NaN (e.g. vLLM column when vllm not installed) - cells.append("-".rjust(width)) - elif "_vs_" in key: - cells.append(f"{val:.2f}x".rjust(width)) - else: - cells.append(f"{val:.1f}".rjust(width)) - else: - cells.append(str(val).rjust(width)) - return " | ".join(cells) - - main_header, main_sep = _build_header_sep(main_cols) - bench_header, bench_sep = _build_header_sep(bench_cols) - border = "=" * max(len(main_header), len(bench_header)) - - # Bucket rows by SSM-state dtype and by main-vs-bench333. Keep each - # bucket's order consistent with ``_perf_results`` insertion so rows - # line up with the parametrize id order. bench333 trace rows are - # tagged by ``log_count > 0`` (main groups all leave log_count at 0). - def _split(rows): - main = [r for r in rows if r.get("log_count", 0) == 0] - bench = [r for r in rows if r.get("log_count", 0) > 0] - return main, bench - - rows_fp32_main, rows_fp32_bench = _split( - [r for r in _perf_results if r["state"] == "fp32"] - ) - rows_bf16_main, rows_bf16_bench = _split( - [r for r in _perf_results if r["state"] == "bf16"] - ) - lines = ["", border] - lines.append( - "K5 Prefill Performance Summary " - "(K5 device kernel time only, via torch.profiler)" - ) - lines.append( - " Triton K5 references always use fp32 SSM state; only FlyDSL's " - "SSM-state dtype changes between the sub-tables below." - ) - lines.append(border) - - def _emit_subtable(title, rows, cols, header, sep): - if not rows: - return - lines.append("") - lines.append(title) - lines.append(sep) - lines.append(header) - lines.append(sep) - for row in rows: - lines.append(_fmt_row(row, cols)) - lines.append(sep) - - _emit_subtable( - "[FlyDSL SSM state = fp32] -- main groups", - rows_fp32_main, - main_cols, - main_header, - main_sep, - ) - _emit_subtable( - "[FlyDSL SSM state = fp32] -- bench333 trace shapes", - rows_fp32_bench, - bench_cols, - bench_header, - bench_sep, - ) - _emit_subtable( - "[FlyDSL SSM state = bf16] -- main groups", - rows_bf16_main, - main_cols, - main_header, - main_sep, - ) - _emit_subtable( - "[FlyDSL SSM state = bf16] -- bench333 trace shapes", - rows_bf16_bench, - bench_cols, - bench_header, - bench_sep, - ) + def _fmt_cell(val, key, width): + if isinstance(val, bool): + return ("Y" if val else "N").rjust(width) + if isinstance(val, float): + if math.isnan(val): # NaN (hip skipped for unsupported shapes) + return "-".rjust(width) + return (f"{val:.2f}x" if "/" in key else f"{val:.1f}").rjust(width) + return str(val).rjust(width) + + header = "|".join(disp.rjust(w) for disp, _, w in cols) + sep = "+".join("-" * w for _, _, w in cols) + border = "=" * len(header) + lines = [ + "", + border, + ( + "K5 Prefill Perf Summary (mfma16_hip vs hip vs triton; K5 device kernel us via " + "torch.profiler; fly/hip & tri/hip = speedup vs hip, >1 faster / <1 slower)" + ), + border, + "", + sep, + header, + sep, + ] + for row in _perf_results: + lines.append("|".join(_fmt_cell(row[k], k, w) for _, k, w in cols)) + lines.append(sep) lines.append("") print("\n".join(lines)) @pytest.fixture(scope="session", autouse=True) def _print_summary_table(request): - """Print the summary performance table after all tests finish.""" + """Print the perf summary table after all tests in the session finish.""" yield _print_perf_table() -# -- bf16 SSM-state correctness ---------------------------------------- - - -class TestStateDtypeBF16: - """Validate that ``state_dtype=bfloat16`` matches the ``float32`` path. - - The bf16-state kernel keeps the f32 accumulator unchanged and only - rounds h0 (extf) and final_state (truncf) at the HBM boundary, so its - output should agree with the f32-state kernel up to one bf16 trunc - error on the SSM state plus accumulated round-off through the chunk - loop. We compare against the *flydsl f32-state* path on the exact same - shape rather than the PyTorch reference, which gives the tightest - regression signal for this specific feature. - """ - - @pytest.mark.parametrize("args", STATE_BF16_PARAMS, ids=STATE_BF16_TEST_IDS) - def test_state_bf16_matches_state_f32(self, args: PrefillArgs): - context_lens = args.resolve_context_lens() - k, _, _, w_c, u_c, g, h0_f32, cu, _ = _make_inputs(context_lens, args=args) - h0_bf16 = h0_f32.to(torch.bfloat16) - - h_f32, vn_f32, fs_f32 = chunk_gated_delta_rule_fwd_h_flydsl( - k, - w_c, - u_c, - g=g, - initial_state=h0_f32, - output_final_state=args.output_final_state, - cu_seqlens=cu, - ) - h_bf16, vn_bf16, fs_bf16 = chunk_gated_delta_rule_fwd_h_flydsl( - k, - w_c, - u_c, - g=g, - initial_state=h0_bf16, - output_final_state=args.output_final_state, - cu_seqlens=cu, - ) - - # final_state dtype must follow the input dtype. - if args.output_final_state: - assert ( - fs_f32 is not None and fs_f32.dtype == torch.float32 - ), f"f32 path produced {fs_f32.dtype} final_state" - assert ( - fs_bf16 is not None and fs_bf16.dtype == torch.bfloat16 - ), f"bf16 path produced {fs_bf16.dtype} final_state" - else: - assert fs_f32 is None and fs_bf16 is None - - # h and v_new are bf16 in both paths (decoupled from state dtype). - assert h_f32.dtype == h_bf16.dtype == k.dtype - if vn_f32 is not None: - assert vn_f32.dtype == vn_bf16.dtype == u_c.dtype - - # The two paths diverge only by the rounding applied to h0/ht. With - # f32 accumulation this stays well within bf16 ULP * (1 + chunk - # length) for sane inputs. - atol = 5e-2 - rtol = 5e-2 - torch.testing.assert_close( - h_bf16.float(), - h_f32.float(), - atol=atol, - rtol=rtol, - msg="bf16-state vs f32-state: h mismatch", - ) - if vn_f32 is not None: - torch.testing.assert_close( - vn_bf16.float(), - vn_f32.float(), - atol=atol, - rtol=rtol, - msg="bf16-state vs f32-state: v_new mismatch", - ) - if args.output_final_state: - torch.testing.assert_close( - fs_bf16.float(), - fs_f32.float(), - atol=atol, - rtol=rtol, - msg="bf16-state vs f32-state: final_state mismatch", - ) - - @pytest.mark.parametrize("args", STATE_BF16_PARAMS, ids=STATE_BF16_TEST_IDS) - def test_state_dtype_kwarg_no_initial_state(self, args: PrefillArgs): - """``state_dtype`` kwarg controls final_state dtype when h0 is None.""" - if not args.output_final_state: - pytest.skip("kwarg only meaningful when final_state is requested") - context_lens = args.resolve_context_lens() - k, _, _, w_c, u_c, g, _, cu, _ = _make_inputs( - context_lens, args=args, with_initial_state=False - ) - - _, _, fs_f32 = chunk_gated_delta_rule_fwd_h_flydsl( - k, - w_c, - u_c, - g=g, - initial_state=None, - output_final_state=True, - cu_seqlens=cu, - # default -> f32 - ) - assert fs_f32 is not None and fs_f32.dtype == torch.float32 - - _, _, fs_bf16 = chunk_gated_delta_rule_fwd_h_flydsl( - k, - w_c, - u_c, - g=g, - initial_state=None, - output_final_state=True, - cu_seqlens=cu, - state_dtype=torch.bfloat16, - ) - assert fs_bf16 is not None and fs_bf16.dtype == torch.bfloat16 - - def test_state_dtype_conflict_raises(self): - """Mismatched ``state_dtype`` and ``initial_state.dtype`` must raise.""" - args = STATE_BF16_PARAMS[0] - context_lens = args.resolve_context_lens() - k, _, _, w_c, u_c, g, h0, cu, _ = _make_inputs(context_lens, args=args) - with pytest.raises(ValueError): - chunk_gated_delta_rule_fwd_h_flydsl( - k, - w_c, - u_c, - g=g, - initial_state=h0, # f32 - output_final_state=args.output_final_state, - cu_seqlens=cu, - state_dtype=torch.bfloat16, # conflict - ) - - def test_state_dtype_unsupported_raises(self): - """Unsupported state dtypes must raise (e.g. fp16).""" - args = STATE_BF16_PARAMS[0] - context_lens = args.resolve_context_lens() - k, _, _, w_c, u_c, g, _, cu, _ = _make_inputs( - context_lens, args=args, with_initial_state=False - ) - with pytest.raises(ValueError): - chunk_gated_delta_rule_fwd_h_flydsl( - k, - w_c, - u_c, - g=g, - initial_state=None, - output_final_state=True, - cu_seqlens=cu, - state_dtype=torch.float16, - ) - - -# -- triton_origin_opt: BV=16 + exp2 variant of fwd_h -------------------- -# -# Inlined from the (now-deleted) standalone benchmark script -# ``0423_gdr_prefill_bench_standalone.py``. Launches the same recurrence -# kernel as the existing ``chunk_gated_delta_rule_fwd_h`` (``triton_origin`` -# in this file), but with two changes that the standalone bench's new -# pipeline applies on top of the original RTP config: -# -# - BV = 16 (was 32): smaller V-tile -> more (V/BV) blocks per (B*H) -# program-id pair -> better occupancy on MI355X. -# - USE_EXP2 = True (was False): emits a single ``v_exp_f32`` per -# gate evaluation instead of the ``v_log + v_mul + v_exp`` chain that -# ``tl.exp`` lowers to. -# -# Because USE_EXP2 expects gates pre-scaled by ``1/ln(2)``, the wrapper -# multiplies the supplied ``g`` (already a per-chunk cumsum, as produced -# by ``_make_inputs``) by RCP_LN2 before launching. That scale step is -# excluded from the K5 kernel time -- we time only the kernel itself, in -# line with how the other K5 wrappers in this file are benchmarked. -# -# The kernel itself is wrapped in a ``triton.autotune`` sweep over -# (BV, num_warps, num_stages); the standalone version pinned BV=16 -# only, but exposing the sweep here matches what aiter's own -# ``chunk_gated_delta_rule_fwd_kernel_h_blockdim64`` does internally -# and lets each shape pick its own best config on first run. - -_RCP_LN2 = 1.0 / 0.6931471805599453 - -_exp = tl.exp -_exp2 = tl.math.exp2 - - -# Decorator stack mirrors FLA's K5 kernels (Heuristics outer, Autotune -# inner) so that Triton 3.x writes the sweep result to its persistent -# autotune cache (`~/.triton/autotune`) via ``cache_results=True``. After -# the first run each (H, K, V, BT, IS_VARLEN) key is served from disk and -# subsequent runs no longer launch the full BV/warps/stages sweep -- the -# rocprof kernel-stats CSV then reports the same ~56 calls as the other -# K5 kernels, instead of the 9000+ that an un-cached sweep produces. -# -# ``Hg`` is intentionally excluded from ``key``: it only affects host-side -# address arithmetic (``H // Hg`` divisor for the K block-ptr), not the -# compiled binary or tile shape, so adding it would just multiply the -# number of unique keys and force redundant sweeps. -@triton.heuristics( - { - "USE_G": lambda args: args["g"] is not None, - "USE_GK": lambda args: args["gk"] is not None, - "USE_INITIAL_STATE": lambda args: args["h0"] is not None, - "STORE_FINAL_STATE": lambda args: args["ht"] is not None, - "SAVE_NEW_VALUE": lambda args: args["v_new"] is not None, - "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, - } -) -@triton.autotune( - configs=[ - triton.Config({"BV": BV}, num_warps=num_warps, num_stages=num_stages) - for BV in (16, 32, 64) - for num_warps in (2, 4) - for num_stages in (2, 3) - ], - key=["H", "K", "V", "BT", "IS_VARLEN"], - cache_results=True, -) -@triton.jit(do_not_specialize=["T"]) -def chunk_gated_delta_rule_fwd_kernel_h_origin_opt( - k, - v, - w, - v_new, - g, - gk, - h, - h0, - ht, - cu_seqlens, - chunk_offsets, - T, - H: tl.constexpr, - Hg: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BV: tl.constexpr, - USE_G: tl.constexpr, - USE_GK: tl.constexpr, - USE_INITIAL_STATE: tl.constexpr, - STORE_FINAL_STATE: tl.constexpr, - SAVE_NEW_VALUE: tl.constexpr, - USE_EXP2: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_v, i_nh = tl.program_id(0), tl.program_id(1) - i_n, i_h = i_nh // H, i_nh % H - if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load( - cu_seqlens + i_n + 1 - ).to(tl.int32) - T = eos - bos - NT = tl.cdiv(T, BT) - boh = tl.load(chunk_offsets + i_n).to(tl.int32) - else: - bos, eos = i_n * T, i_n * T + T - NT = tl.cdiv(T, BT) - boh = i_n * NT - - b_h1 = tl.zeros([64, BV], dtype=tl.float32) - if K > 64: - b_h2 = tl.zeros([64, BV], dtype=tl.float32) - if K > 128: - b_h3 = tl.zeros([64, BV], dtype=tl.float32) - if K > 192: - b_h4 = tl.zeros([64, BV], dtype=tl.float32) - - h += ((boh * H + i_h) * K * V).to(tl.int64) - v += ((bos * H + i_h) * V).to(tl.int64) - k += ((bos * Hg + i_h // (H // Hg)) * K).to(tl.int64) - w += ((bos * H + i_h) * K).to(tl.int64) - if SAVE_NEW_VALUE: - v_new += ((bos * H + i_h) * V).to(tl.int64) - stride_v = H * V - stride_h = H * K * V - stride_k = Hg * K - stride_w = H * K - if USE_INITIAL_STATE: - h0 = h0 + i_nh * K * V - if STORE_FINAL_STATE: - ht = ht + i_nh * K * V - - if USE_INITIAL_STATE: - p_h0_1 = tl.make_block_ptr(h0, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)) - b_h1 += tl.load(p_h0_1, boundary_check=(0, 1)).to(tl.float32) - if K > 64: - p_h0_2 = tl.make_block_ptr( - h0, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0) - ) - b_h2 += tl.load(p_h0_2, boundary_check=(0, 1)).to(tl.float32) - if K > 128: - p_h0_3 = tl.make_block_ptr( - h0, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0) - ) - b_h3 += tl.load(p_h0_3, boundary_check=(0, 1)).to(tl.float32) - if K > 192: - p_h0_4 = tl.make_block_ptr( - h0, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0) - ) - b_h4 += tl.load(p_h0_4, boundary_check=(0, 1)).to(tl.float32) - - for i_t in range(NT): - p_h1 = tl.make_block_ptr( - h + i_t * stride_h, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0) - ) - tl.store(p_h1, b_h1.to(p_h1.dtype.element_ty), boundary_check=(0, 1)) - if K > 64: - p_h2 = tl.make_block_ptr( - h + i_t * stride_h, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0) - ) - tl.store(p_h2, b_h2.to(p_h2.dtype.element_ty), boundary_check=(0, 1)) - if K > 128: - p_h3 = tl.make_block_ptr( - h + i_t * stride_h, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0) - ) - tl.store(p_h3, b_h3.to(p_h3.dtype.element_ty), boundary_check=(0, 1)) - if K > 192: - p_h4 = tl.make_block_ptr( - h + i_t * stride_h, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0) - ) - tl.store(p_h4, b_h4.to(p_h4.dtype.element_ty), boundary_check=(0, 1)) - - p_w = tl.make_block_ptr( - w, (T, K), (stride_w, 1), (i_t * BT, 0), (BT, 64), (1, 0) - ) - b_w = tl.load(p_w, boundary_check=(0, 1)) - b_v = tl.dot(b_w, b_h1.to(b_w.dtype)) - if K > 64: - p_w = tl.make_block_ptr( - w, (T, K), (stride_w, 1), (i_t * BT, 64), (BT, 64), (1, 0) - ) - b_w = tl.load(p_w, boundary_check=(0, 1)) - b_v += tl.dot(b_w, b_h2.to(b_w.dtype)) - if K > 128: - p_w = tl.make_block_ptr( - w, (T, K), (stride_w, 1), (i_t * BT, 128), (BT, 64), (1, 0) - ) - b_w = tl.load(p_w, boundary_check=(0, 1)) - b_v += tl.dot(b_w, b_h3.to(b_w.dtype)) - if K > 192: - p_w = tl.make_block_ptr( - w, (T, K), (stride_w, 1), (i_t * BT, 192), (BT, 64), (1, 0) - ) - b_w = tl.load(p_w, boundary_check=(0, 1)) - b_v += tl.dot(b_w, b_h4.to(b_w.dtype)) - p_v = tl.make_block_ptr( - v, (T, V), (stride_v, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0) - ) - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_v - - if SAVE_NEW_VALUE: - p_v2 = tl.make_block_ptr( - v_new, (T, V), (stride_v, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0) - ) - tl.store(p_v2, b_v.to(p_v2.dtype.element_ty), boundary_check=(0, 1)) - - last_idx = min((i_t + 1) * BT, T) - 1 - if USE_G: - m_t = (i_t * BT + tl.arange(0, BT)) < T - b_g_last = tl.load(g + (bos * H + last_idx * H + i_h).to(tl.int64)).to( - tl.float32 - ) - p_g = tl.make_block_ptr( - g + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,) - ) - b_g = tl.load(p_g, boundary_check=(0,)).to(tl.float32) - if USE_EXP2: - b_v = b_v * tl.where(m_t, _exp2(b_g_last - b_g), 0)[:, None] - b_g_last = _exp2(b_g_last) - else: - b_v = b_v * tl.where(m_t, _exp(b_g_last - b_g), 0)[:, None] - b_g_last = _exp(b_g_last) - b_h1 = b_h1 * b_g_last - if K > 64: - b_h2 = b_h2 * b_g_last - if K > 128: - b_h3 = b_h3 * b_g_last - if K > 192: - b_h4 = b_h4 * b_g_last - - b_v = b_v.to(k.dtype.element_ty) - p_k = tl.make_block_ptr( - k, (K, T), (1, stride_k), (0, i_t * BT), (64, BT), (0, 1) - ) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_h1 += tl.dot(b_k, b_v) - if K > 64: - p_k = tl.make_block_ptr( - k, (K, T), (1, stride_k), (64, i_t * BT), (64, BT), (0, 1) - ) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_h2 += tl.dot(b_k, b_v) - if K > 128: - p_k = tl.make_block_ptr( - k, (K, T), (1, stride_k), (128, i_t * BT), (64, BT), (0, 1) - ) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_h3 += tl.dot(b_k, b_v) - if K > 192: - p_k = tl.make_block_ptr( - k, (K, T), (1, stride_k), (192, i_t * BT), (64, BT), (0, 1) - ) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_h4 += tl.dot(b_k, b_v) - - if STORE_FINAL_STATE: - p_ht = tl.make_block_ptr(ht, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)) - tl.store(p_ht, b_h1.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) - if K > 64: - p_ht = tl.make_block_ptr( - ht, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0) - ) - tl.store(p_ht, b_h2.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) - - -_TRITON_ORIGIN_OPT_KERNEL = chunk_gated_delta_rule_fwd_kernel_h_origin_opt - - -def chunk_gated_delta_rule_fwd_h_origin_opt( - k, - w, - u, - g=None, - initial_state=None, - output_final_state=False, - cu_seqlens=None, -): - """``triton_origin_opt`` K5: USE_EXP2 + autotuned BV/warps/stages variant. - - Mirrors the standalone bench's ``fwd_h`` host wrapper but adds a - Triton autotune sweep over ``BV ? {16, 32, 64}``, ``num_warps ? - {2, 4}``, ``num_stages ? {1, 2, 3}``. Keyed on ``(H, Hg, K, V, BT, - IS_VARLEN)`` so each shape picks its own best config on first run. - - Inputs use the GQA layout from PREFILL_PARAMS unchanged -- since this - K5 kernel accepts ``Hg`` directly, no ``repeat_interleave`` is needed - (unlike the original ``chunk_gated_delta_rule_fwd_h``, which is - MHA-only). - - NOTE: The RCP_LN2 scale required by USE_EXP2=True is applied here so - that callers can pass the same per-chunk-cumsum ``g`` as the other - K5 wrappers. This scale is a cheap elementwise multiply and is - excluded from the kernel-time measurement when ``_bench_fn`` profiles - only the kernel launch. - """ - import triton as _triton - - B, T, Hg, K = k.shape - V = u.shape[-1] - H = u.shape[-2] - BT = 64 - if cu_seqlens is None: - N, NT, chunk_offsets = B, _triton.cdiv(T, BT), None - else: - from aiter.ops.triton._triton_kernels.gated_delta_rule.prefill.chunk_delta_h import ( - prepare_chunk_indices, - prepare_chunk_offsets, - ) - - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - N = len(cu_seqlens) - 1 - NT = len(chunk_indices) - chunk_offsets = prepare_chunk_offsets(cu_seqlens, BT) - - h = k.new_empty(B, NT, H, K, V) - v_new = torch.empty_like(u) - final_state = ( - k.new_empty(N, H, K, V, dtype=torch.float32) if output_final_state else None - ) - - # USE_EXP2=True expects gates pre-scaled by 1/ln(2). Cheap elementwise - # op; excluded from kernel time when profiled via torch.profiler. - # ``triton_origin_opt`` reads g with token-major offsets - # ``(bos*H + row*H + i_h)`` so transpose head-major [H, T_total] back - # to token-major [T_total, H] before launch. - if g is not None: - if g.dim() == 2: - g_for_kernel = g.transpose(0, 1).contiguous() # [T_total, H] - else: - g_for_kernel = g.transpose(-2, -1).contiguous() - g_scaled = g_for_kernel * _RCP_LN2 - else: - g_scaled = None - - def grid(meta): - return (_triton.cdiv(V, meta["BV"]), N * H) - - _TRITON_ORIGIN_OPT_KERNEL[grid]( - k=k, - v=u, - w=w, - v_new=v_new, - g=g_scaled, - gk=None, - h=h, - h0=initial_state, - ht=final_state, - cu_seqlens=cu_seqlens, - chunk_offsets=chunk_offsets, - T=T, - H=H, - Hg=Hg, - K=K, - V=V, - BT=BT, - USE_EXP2=True, - ) - return h, v_new, final_state +class TestPerformance: + @pytest.mark.parametrize("args", PREFILL_PARAMS, ids=PREFILL_TEST_IDS) + def test_perf_comparison(self, args: PrefillArgs): + _run_perf_comparison(args) diff --git a/op_tests/test_gated_delta_rule.py b/op_tests/test_gated_delta_rule.py index 54024098e26..5c4d520ff79 100644 --- a/op_tests/test_gated_delta_rule.py +++ b/op_tests/test_gated_delta_rule.py @@ -10,17 +10,11 @@ import torch.nn.functional as F from einops import rearrange, repeat -from aiter.ops.triton.gated_delta_net import ( - fused_recurrent_gated_delta_rule, - chunk_gated_delta_rule, - chunk_gated_delta_rule_opt, - chunk_gated_delta_rule_opt_vk, -) from aiter.ops.chunk_gated_delta_rule_fwd_h import ( chunk_gated_delta_rule_fwd_h_hip_fn, ) -from aiter.ops.triton._triton_kernels.gated_delta_rule.prefill import ( - chunk_gated_delta_rule_fwd_h_opt_vk, +from aiter.ops.flydsl.linear_attention_prefill_kernels import ( + chunk_gated_delta_rule_fwd_h_flydsl_mfma16_hip, ) from aiter.ops.triton._triton_kernels.gated_delta_rule.decode.fused_sigmoid_gating_recurrent import ( fused_sigmoid_gating_delta_rule_update, @@ -31,6 +25,15 @@ assert_close, device, ) +from aiter.ops.triton._triton_kernels.gated_delta_rule.prefill import ( + chunk_gated_delta_rule_fwd_h_opt_vk, +) +from aiter.ops.triton.gated_delta_net import ( + chunk_gated_delta_rule, + chunk_gated_delta_rule_opt, + chunk_gated_delta_rule_opt_vk, + fused_recurrent_gated_delta_rule, +) def _is_gfx12_runtime() -> bool: @@ -959,6 +962,19 @@ def test_chunk_opt_varlen_hip( ), ], ), + pytest.param( + "flydsl", + id="flydsl", + marks=[ + pytest.mark.skipif( + not IS_AMD, reason="FlyDSL backend requires an AMD device" + ), + pytest.mark.skipif( + _is_gfx12_runtime(), + reason="FlyDSL mfma16_hip K5 kernel does not support gfx12!", + ), + ], + ), ], ) @pytest.mark.parametrize( @@ -1009,6 +1025,13 @@ def test_chunk_opt_vk_indice( if D != 128: pytest.skip(reason="HIP kernel requires D=128 and bfloat16") extra_kwargs = {"g_head_major": True} + elif backend == "flydsl": + fwd_h = chunk_gated_delta_rule_fwd_h_flydsl_mfma16_hip + # FlyDSL mfma16_hip is likewise specialized for K=V=128 / bf16 and takes + # head-major [B, H, T] gates directly (no g_head_major flag needed). + if D != 128: + pytest.skip(reason="FlyDSL mfma16_hip kernel requires D=128 and bfloat16") + extra_kwargs = {} else: fwd_h = chunk_gated_delta_rule_fwd_h_opt_vk extra_kwargs = {}