From 45f1bb98d423cec6291413229499215d0629d76a Mon Sep 17 00:00:00 2001 From: Xiake Sun Date: Fri, 21 Aug 2026 13:33:25 +0000 Subject: [PATCH 01/59] feat(fmha): support paged FP8 D192 with asymmetric values Signed-off-by: Xiake Sun --- .../attention/flash_attn_fp8_paged_gfx950.py | 964 ++++++++++++++++++ kernels/attention/flash_attn_interface.py | 161 ++- kernels/attention/flash_attn_utils.py | 846 ++++++++++++++- tests/kernels/test_flash_attn_fwd.py | 129 +++ 4 files changed, 2054 insertions(+), 46 deletions(-) create mode 100644 kernels/attention/flash_attn_fp8_paged_gfx950.py diff --git a/kernels/attention/flash_attn_fp8_paged_gfx950.py b/kernels/attention/flash_attn_fp8_paged_gfx950.py new file mode 100644 index 0000000000..05b909fe62 --- /dev/null +++ b/kernels/attention/flash_attn_fp8_paged_gfx950.py @@ -0,0 +1,964 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2025 FlyDSL Project Contributors + +"""gfx950 page-64 FP8 flash attention with independent value dimensions.""" + +import flydsl.compiler as flyc +import flydsl.expr as fx +from flydsl.compiler.kernel_function import CompilationContext +from flydsl.expr import const_expr, range_constexpr, rocdl +from flydsl.expr.typing import T +from flydsl.expr.utils.arith import ArithValue +from flydsl.expr.utils.arith import _to_raw as _raw +from flydsl.runtime.device import get_rocm_arch as get_hip_arch +from kernels.attention.flash_attn_utils import ( + DualwaveFp8GemmHelper, + DualwaveFp8KernelContext, + DualwaveFp8KvGmemToLdsLoader, + DualwaveFp8KvLdsToVgprLoader, + DualwaveFp8QLoader, + DualwaveFp8SoftmaxHelper, + DualwaveFp8StoreHelper, + DualwaveSplitKCombineContext, + DualwaveSplitKCombineHelper, + _make_paged_dualwave_swp_fp8_traits, + _sched_barrier_exp_pairs, + _sched_barrier_pairs, + _stagger_extra_barrier_if_one, + _stagger_extra_barrier_if_zero, + _waitcnt_vm_n, + dualwave_splitk_workspace_elems, # noqa: F401 +) +from kernels.common.kernels_common import dtype_to_elem_type + + +def build_flash_attn_paged_fp8_module( + num_heads, + head_dim, + value_head_dim=None, + causal=True, + dtype_str="bf16", + num_kv_heads=None, + waves_per_eu=2, + daz=True, + dualwave_swp_lazy_rescale=True, + rescale_threshold=8.0, + dualwave_swp_setprio=True, + dualwave_swp_debug_lazy_counts=False, + dualwave_swp_enable_stagger=True, + num_kv_splits=1, + varlen=False, + cross_seqlen=False, + paged=False, + kv_cache_layout="linear", +): + """Build the gfx950 packed-varlen page-64 FP8 attention launcher.""" + gpu_arch = get_hip_arch() + if value_head_dim is None: + value_head_dim = head_dim + + if not gpu_arch.startswith("gfx950"): + raise RuntimeError(f"flash_attn_dualwave_swp requires gfx950+ (uses ds_read_tr16_b64), got {gpu_arch}") + if ( + not paged + or dtype_str != "fp8" + or not causal + or not varlen + or not cross_seqlen + or head_dim != 192 + or value_head_dim not in (128, 192) + or kv_cache_layout != "vectorized" + or int(num_kv_splits) != 1 + ): + raise RuntimeError( + "paged FP8 flash_attn requires gfx950, causal packed-varlen cross-attention, " + "page-64 vectorized KV, head_dim=192, value_head_dim in {128,192}, " + "and num_kv_splits=1" + ) + + if num_kv_heads is None: + num_kv_heads = num_heads + assert num_heads % num_kv_heads == 0 + NUM_KV_SPLITS = int(num_kv_splits) + assert NUM_KV_SPLITS >= 1 + if varlen and num_kv_splits and int(num_kv_splits) > 1: + raise ValueError("varlen is not supported together with num_kv_splits > 1") + + # All compile-time tile/layout constants live in the fp8 traits object. + traits = _make_paged_dualwave_swp_fp8_traits( + num_heads, + num_kv_heads, + head_dim, + value_head_dim=value_head_dim, + causal=causal, + waves_per_eu=waves_per_eu, + daz=daz, + dualwave_swp_lazy_rescale=dualwave_swp_lazy_rescale, + rescale_threshold=rescale_threshold, + dualwave_swp_setprio=dualwave_swp_setprio, + dualwave_swp_debug_lazy_counts=dualwave_swp_debug_lazy_counts, + dualwave_swp_enable_stagger=dualwave_swp_enable_stagger, + num_kv_splits=num_kv_splits, + varlen=varlen, + cross_seqlen=cross_seqlen, + paged=paged, + kv_cache_layout=kv_cache_layout, + ) + # Builder-level aliases used by SharedStorage and the launch/compile wrappers. + SPLITK = traits.SPLITK + BLOCK_M = traits.BLOCK_M + BLOCK_SIZE = traits.BLOCK_SIZE + HEAD_DIM = traits.HEAD_DIM + NUM_HEADS_Q = traits.NUM_HEADS_Q + PAGED = traits.PAGED + DEFAULT_STRIDE_Q_N = traits.DEFAULT_STRIDE_Q_N + DEFAULT_STRIDE_O_N = traits.NUM_HEADS_Q * traits.V_HEAD_DIM + DEFAULT_STRIDE_KV_N = traits.DEFAULT_STRIDE_KV_N + _dualwave_swp_fp8_cache_tag = traits.cache_tag + _lds_elem_dtype = dtype_to_elem_type(traits.DTYPE_STR) + + @fx.struct + class SharedStorage: + kv: fx.Array[_lds_elem_dtype, traits.LDS_KV_TOTAL_SIZE, 16] + vt: fx.Array[fx.BFloat16, traits.VT_BF16_TOTAL, 16] + + @flyc.kernel(known_block_size=[BLOCK_SIZE, 1, 1]) + def flash_attn_dualwave_swp_fp8_gfx950_kernel( + Q: fx.Tensor, + K: fx.Tensor, + V: fx.Tensor, + O: fx.Tensor, # noqa: E741 + DebugCounts: fx.Tensor, + CuSeqQ: fx.Tensor, + CuSeqKv: fx.Tensor, + BlockTable: fx.Tensor, + block_table_stride: fx.Int32, + QDescale: fx.Tensor, + KDescale: fx.Tensor, + VDescale: fx.Tensor, + seq_len: fx.Int32, + seq_len_kv: fx.Int32, + stride_q_n: fx.Int32, + stride_o_n: fx.Int32, + stride_kv_n: fx.Int32, + head_dim_runtime: fx.Int32, + ): + # Per-kernel setup lives in the fp8 context; the inline pipeline helpers below + # bind the ctx fields to local names so the schedule reads unchanged. + ctx = DualwaveFp8KernelContext( + traits, + Q, + K, + V, + O, + DebugCounts, + CuSeqQ, + CuSeqKv, + QDescale, + KDescale, + VDescale, + seq_len, + seq_len_kv, + stride_q_n, + stride_kv_n, + head_dim_runtime, + stride_o_n=stride_o_n, + BlockTable=BlockTable, + block_table_stride=block_table_stride, + ) + ctx.init_types_and_constants() + ctx.init_runtime_indices() + ctx.init_lds(SharedStorage) + ctx.init_thread_mapping() + ctx.init_sequence_lengths() + ctx.init_descriptors() + ctx.init_atoms_and_lds_ptrs() + ctx.init_dma_thread_offsets() + ctx.init_descale() + ctx.init_tile_bounds() + ctx.init_workspace_io() + + # fp8 pipeline helpers (logic lives in flash_attn_utils; the kernel drives the + # software-pipeline schedule below and calls into these). + q_loader = DualwaveFp8QLoader(ctx) + gemm_helper = DualwaveFp8GemmHelper(ctx) + softmax_helper = DualwaveFp8SoftmaxHelper(ctx) + kv_gmem_to_lds = DualwaveFp8KvGmemToLdsLoader(ctx) + kv_lds_to_regs = DualwaveFp8KvLdsToVgprLoader(ctx) + output_store = DualwaveFp8StoreHelper(ctx) + + # Skip packed-varlen q-blocks beyond this request's query length. The + # condition is uniform across the workgroup, so barriers stay balanced. + @flyc.jit + def _run_q_block(): + # Prologue: load K tile split_t0 -> LDS buf0, wait, and sync the workgroup. + kv_gmem_to_lds.load_k(ctx.split_t0 * traits.BLOCK_N, 0) + rocdl.s_waitcnt(0) + rocdl.sched_barrier(0) + rocdl.s_barrier() + + # Load this wave's raw fp8 Q rows for wide QK; q/k descale is applied to + # the fp32 logits after the MFMA. init_q_row sets q_row/q_row_i32/ + # q_start_pos_i32 on ctx for the causal-mask helpers. + ctx.init_q_row() + q_row = ctx.q_row + q_all_wide = q_loader.load_all_wide(ctx.q_row_in_block) + + # Pipeline ahead: prefetch K tile1 (buf1) + V tile0 (buf0) as background + kv_gmem_to_lds.load_k((ctx.split_t0 + 1) * traits.BLOCK_N, 1) + kv_gmem_to_lds.load_v(ctx.split_t0 * traits.BLOCK_N, 0) + v_k = kv_lds_to_regs.load_k(0) + rocdl.sched_barrier(0) + rocdl.s_waitcnt(traits.LGKMCNT_0_ONLY) + _waitcnt_vm_n(ctx.NUM_DMA_V) + + # OPEN the wave-group phase shift: one extra s_barrier on group B + if const_expr(traits.DUALWAVE_SWP_ENABLE_STAGGER): + _stagger_extra_barrier_if_one(ctx.stagger_i32) # group B: +1 s_barrier -> open the shift + else: + rocdl.sched_barrier(0) + rocdl.s_barrier() + + # Prologue scores + first softmax pass for KV tile 0 + v_s_0 = gemm_helper.qk(v_k, q_all_wide) + rocdl.sched_barrier(0) + if const_expr(traits.CAUSAL): + if const_expr(SPLITK): + v_s_0 = softmax_helper.causal_mask_prologue_if_needed( + v_s_0, ctx.split_t0, (ctx.split_t0 + 1) * traits.BLOCK_N + ) + else: + v_s_0 = softmax_helper.causal_mask_prologue_if_needed(v_s_0) + else: + # Non-causal padding mask for the prologue tile too: for tiny seq_len + # tile 0 is the only real tile, so its keys >= seq_len must be masked + # here. Gated -> no-op once tile 0 is full (seq_len >= BLOCK_N). + if const_expr(SPLITK): + v_s_0 = softmax_helper.seq_pad_mask_if_needed(v_s_0, ctx.split_t0) + else: + v_s_0 = softmax_helper.seq_pad_mask_if_needed(v_s_0) + m_row_pro = softmax_helper.reduce_max(v_s_0) + if const_expr(traits.CAUSAL): + # Floor fully-masked rows (-inf) to finite so exp2 yields 0, not NaN. + m_row_pro = softmax_helper.floor_masked_max(m_row_pro) + v_s_0 = softmax_helper.sub_m(v_s_0, m_row_pro) + v_p_0 = softmax_helper.exp2(v_s_0, 0, 16) + rocdl.sched_barrier(0) + rocdl.s_barrier() + rocdl.sched_barrier(0) + + # Prefetch K tile 2 into buf0, keeping the K double-buffer one step ahead + kv_gmem_to_lds.load_k((ctx.split_t0 + 2) * traits.BLOCK_N, 0) + + # Loop-carried state (scf.for init args): m_row, l_row(=0), D_CHUNKS zero + l_row_init = ctx.c_zero_f + init_args = [m_row_pro, l_row_init] + for _ in range_constexpr(traits.D_CHUNKS): + init_args.append(ctx.c_zero_v16f32) + init_args.append(ctx.v_pair_to_vec32(v_p_0)) + + # ============================= Main loop ============================= + # Software-pipelined inner loop + if const_expr(SPLITK): + loop_lb = ctx.split_t0 + 3 + else: + loop_lb = fx.Index(3) + loop_results = init_args + for j, loop_args in range( + loop_lb, + ctx.split_t_end - fx.Index(1), + fx.Index(2), + init=init_args, + ): + m_row = loop_args[0] + l_row = loop_args[1] + v_o = [loop_args[2 + i] for i in range_constexpr(traits.D_CHUNKS)] + v_p_0 = ctx.v_vec32_to_pair(loop_args[2 + traits.D_CHUNKS]) + j_idx = j + + # Cluster 0 (memory): prefetch next V (buf1), read resident K from LDS + # (v_k) for MMA0, wait + sync. + kv_gmem_to_lds.load_v((j_idx - 2) * traits.BLOCK_N, 1) + v_k = kv_lds_to_regs.load_k(1) + rocdl.s_waitcnt(traits.LGKMCNT_0_ONLY) + _waitcnt_vm_n(ctx.NUM_DMA_K + ctx.NUM_DMA_V) + rocdl.sched_barrier(0) + rocdl.s_barrier() + rocdl.sched_barrier(0) + + # Cluster 1 (compute): MMA0 -> v_s_1; finish v_p_0's 2nd-half exp2, + # sum into l_row, cast to bf16 for P*V. + v_s_1 = gemm_helper.qk(v_k, q_all_wide) + v_p_0 = softmax_helper.exp2(v_p_0, 16, 16) + l_row = softmax_helper.reduce_sum(l_row, v_p_0) + v_p_0 = softmax_helper.cast_p(v_p_0) + v_p_0 = softmax_helper.anchor_v_p(v_p_0) + _sched_barrier_exp_pairs(traits, 6, 3, 1) + _sched_barrier_pairs(traits, 10, 5, 1) + rocdl.sched_barrier(0) + rocdl.s_barrier() + rocdl.sched_barrier(0) + + # Cluster 2 (memory): prefetch next K (buf1), read this tile's V from + # LDS (v_v) for P*V, wait + sync. + kv_gmem_to_lds.load_k(j_idx * traits.BLOCK_N, 1) + v_v = kv_lds_to_regs.load_v(0) + rocdl.s_waitcnt(traits.LGKMCNT_0_ONLY) + _waitcnt_vm_n(ctx.NUM_DMA_K + ctx.NUM_DMA_V) + rocdl.sched_barrier(0) + rocdl.s_barrier() + rocdl.sched_barrier(0) + + # Cluster 3 (compute): first P*V step + row max of v_s_1, lazy + # rescale, remaining 3 P*V steps, sub row + 1st-half exp2 of v_s_1. + if const_expr(traits.DUALWAVE_SWP_SETPRIO): + rocdl.s_setprio(1) + v_o = gemm_helper.pv_step_k(0, v_p_0, v_v, v_o) + # Cross-length causal can put a diagonal tile in v_s_1; mask it here. + # Self-attention skips this to keep the existing schedule. + if const_expr(traits.CAUSAL and traits.CROSS_SEQLEN): + v_s_1 = softmax_helper.causal_mask_prologue_if_needed( + v_s_1, j_idx - 2, (j_idx - 1) * traits.BLOCK_N + ) + else: + v_s_1 = softmax_helper.v_s_vec_to_lists(v_s_1) + m_tile_max_a = softmax_helper.reduce_max(v_s_1) + + _sched_barrier_pairs(traits, 4, 6, 2) + + if const_expr(traits.DUALWAVE_SWP_LAZY_RESCALE): + v_o, m_row, l_row, v_p_0 = softmax_helper.lazy_rescale_o(v_o, m_row, l_row, m_tile_max_a, v_p_0) + else: + v_o, m_row, l_row, v_p_0 = softmax_helper.rescale_o(v_o, m_row, l_row, m_tile_max_a, v_p_0) + v_o = gemm_helper.pv_step_k(1, v_p_0, v_v, v_o) + v_o = gemm_helper.pv_step_k(2, v_p_0, v_v, v_o) + v_o = gemm_helper.pv_step_k(3, v_p_0, v_v, v_o) + v_s_1 = softmax_helper.sub_m(v_s_1, m_row) + v_p_1 = softmax_helper.exp2(v_s_1, 0, 16) + + _sched_barrier_pairs(traits, 6, 6, 2) + # IGroupLP hint (group 2): 6 MFMA each paired with 3 EXP/TRANS (mask + # 0x400) so the new softmax exp2 stays near its MFMA window. + _sched_barrier_exp_pairs(traits, 6, 3, 2) + if const_expr(traits.DUALWAVE_SWP_SETPRIO): + rocdl.s_setprio(0) + # sched_barrier(0): compiler scheduling fence (mask 0 = nothing + # crosses), pinning s_setprio(0) and the closing s_barrier at the + # cluster boundary. Emits no ISA; the real sync is s_barrier(). + rocdl.sched_barrier(0) + rocdl.s_barrier() + rocdl.sched_barrier(0) + + # Cluster 4 (memory, mirror of C0): prefetch V (buf0), read K from + # buf0 into v_k, wait + sync. + kv_gmem_to_lds.load_v((j_idx - 1) * traits.BLOCK_N, 0) + v_k = kv_lds_to_regs.load_k(0) + rocdl.s_waitcnt(traits.LGKMCNT_0_ONLY) + _waitcnt_vm_n(ctx.NUM_DMA_K + ctx.NUM_DMA_V) + rocdl.sched_barrier(0) + rocdl.s_barrier() + rocdl.sched_barrier(0) + + # Cluster 5 (compute, mirror of C1): MMA0 -> v_s_0; finish v_p_1's + # 2nd-half exp2, sum into l_row, cast to bf16. + v_s_0 = gemm_helper.qk(v_k, q_all_wide) + v_p_1 = softmax_helper.exp2(v_p_1, 16, 16) + l_row = softmax_helper.reduce_sum(l_row, v_p_1) + v_p_1 = softmax_helper.cast_p(v_p_1) + v_p_1 = softmax_helper.anchor_v_p(v_p_1) + _sched_barrier_exp_pairs(traits, 6, 3, 3) + _sched_barrier_pairs(traits, 10, 5, 3) + rocdl.sched_barrier(0) + rocdl.s_barrier() + rocdl.sched_barrier(0) + + # Cluster 6 (memory): prefetch next K (buf0), read V packs (buf1), + # apply causal mask to v_s_0 (if causal), wait + sync. + kv_gmem_to_lds.load_k((j_idx + 1) * traits.BLOCK_N, 0) + v_packs_b = kv_lds_to_regs.load_v(1) + if const_expr(traits.CAUSAL): + v_s_0 = softmax_helper.causal_mask_prologue_if_needed( + v_s_0, + j_idx - 1, + j_idx * traits.BLOCK_N, + ) + else: + v_s_0 = softmax_helper.v_s_vec_to_lists(v_s_0) + rocdl.s_waitcnt(traits.LGKMCNT_0_ONLY) + _waitcnt_vm_n(ctx.NUM_DMA_K + ctx.NUM_DMA_V) + rocdl.sched_barrier(0) + rocdl.s_barrier() + rocdl.sched_barrier(0) + + # Cluster 7 (compute, mirror of C3 for v_p_1/v_s_0): closes the iter, + # yield_args carries (m_row, l_row, v_o, packed v_p_0) to the next. + if const_expr(traits.DUALWAVE_SWP_SETPRIO): + rocdl.s_setprio(1) + v_v = v_packs_b + v_o = gemm_helper.pv_step_k(0, v_p_1, v_v, v_o) + m_tile_max_b = softmax_helper.reduce_max(v_s_0) + _sched_barrier_pairs(traits, 4, 6, 4) + + if const_expr(traits.DUALWAVE_SWP_LAZY_RESCALE): + v_o, m_row, l_row, v_p_1 = softmax_helper.lazy_rescale_o(v_o, m_row, l_row, m_tile_max_b, v_p_1) + else: + v_o, m_row, l_row, v_p_1 = softmax_helper.rescale_o(v_o, m_row, l_row, m_tile_max_b, v_p_1) + v_v = v_packs_b + v_o = gemm_helper.pv_step_k(1, v_p_1, v_v, v_o) + v_o = gemm_helper.pv_step_k(2, v_p_1, v_v, v_o) + v_o = gemm_helper.pv_step_k(3, v_p_1, v_v, v_o) + v_s_0 = softmax_helper.sub_m(v_s_0, m_row) + v_p_0 = softmax_helper.exp2(v_s_0, 0, 16) + _sched_barrier_pairs(traits, 6, 5, 4) + _sched_barrier_exp_pairs(traits, 6, 3, 4) + if const_expr(traits.DUALWAVE_SWP_SETPRIO): + rocdl.s_setprio(0) + rocdl.sched_barrier(0) + rocdl.s_barrier() + rocdl.sched_barrier(0) + + yield_args = [m_row, l_row] + v_o + [ctx.v_pair_to_vec32(v_p_0)] + loop_results = yield yield_args + + # Epilogue: drain the pipeline for the final tiles the loop left in + # flight. Mirrors the main-loop clusters but with no further + # prefetch-ahead. Unpack the loop-carried state: + m_row = loop_results[0] + l_row = loop_results[1] + v_o = [loop_results[2 + i] for i in range_constexpr(traits.D_CHUNKS)] + v_p_0 = ctx.v_vec32_to_pair(loop_results[2 + traits.D_CHUNKS]) + + # Tile indices for the last three tiles handled by the epilogue. + max_m3 = ctx.split_t_end - 3 + max_m2 = ctx.split_t_end - 2 + max_m1 = ctx.split_t_end - 1 + + # Epilogue C0 (memory): prefetch V max_m3 (buf1), read K from buf1, sync. + kv_gmem_to_lds.load_v(max_m3 * traits.BLOCK_N, 1) + v_k = kv_lds_to_regs.load_k(1) + rocdl.s_waitcnt(traits.LGKMCNT_0_ONLY) + _waitcnt_vm_n(ctx.NUM_DMA_K + ctx.NUM_DMA_V) + rocdl.sched_barrier(0) + rocdl.s_barrier() + rocdl.sched_barrier(0) + + # Epilogue C1 (compute): MMA0 -> v_s_1; finish v_p_0 softmax (like C1). + v_s_1 = gemm_helper.qk(v_k, q_all_wide) + v_p_0 = softmax_helper.exp2(v_p_0, 16, 16) + l_row = softmax_helper.reduce_sum(l_row, v_p_0) + v_p_0 = softmax_helper.cast_p(v_p_0) + v_p_0 = softmax_helper.anchor_v_p(v_p_0) + _sched_barrier_exp_pairs(traits, 6, 3, 5) + _sched_barrier_pairs(traits, 10, 5, 5) + rocdl.sched_barrier(0) + rocdl.s_barrier() + rocdl.sched_barrier(0) + + # Epilogue C2 (memory): prefetch K max_m1, read V packs (buf0), causal mask v_s_1, sync. + kv_gmem_to_lds.load_k(max_m1 * traits.BLOCK_N, 1) + v_packs_e3 = kv_lds_to_regs.load_v(0) + if const_expr(traits.CAUSAL): + v_s_1 = softmax_helper.causal_mask_prologue_if_needed( + v_s_1, + max_m3, + max_m2 * traits.BLOCK_N, + ) + else: + v_s_1 = softmax_helper.seq_pad_mask_if_needed(v_s_1, max_m3) + rocdl.s_waitcnt(traits.LGKMCNT_0_ONLY) + _waitcnt_vm_n(ctx.NUM_DMA_K + ctx.NUM_DMA_V) + rocdl.sched_barrier(0) + rocdl.s_barrier() + rocdl.sched_barrier(0) + + # Epilogue C3 (compute): full P*V + unconditional rescale + if const_expr(traits.DUALWAVE_SWP_SETPRIO): + rocdl.s_setprio(1) + v_o = gemm_helper.pv(v_p_0, v_packs_e3, v_o) + m_tile_max_e3 = softmax_helper.reduce_max(v_s_1) + row_max_e3, rescale_e3 = softmax_helper.rescale_from_tile_max(m_row, m_tile_max_e3) + m_row = row_max_e3 + v_s_1 = softmax_helper.sub_m(v_s_1, row_max_e3) + v_p_1 = softmax_helper.exp2(v_s_1, 0, 16) + _sched_barrier_pairs(traits, 10, 5, 6) + _sched_barrier_exp_pairs(traits, 6, 3, 6) + rocdl.sched_barrier(0) + softmax_helper.scale_o(v_o, rescale_e3) + v_o = softmax_helper.anchor_v_o(v_o) + + if const_expr(traits.DUALWAVE_SWP_SETPRIO): + rocdl.s_setprio(0) + rocdl.sched_barrier(0) + rocdl.s_barrier() + rocdl.sched_barrier(0) + + # Epilogue C4 (memory): prefetch V max_m2 (buf0), read K from buf0, sync. + kv_gmem_to_lds.load_v(max_m2 * traits.BLOCK_N, 0) + v_k = kv_lds_to_regs.load_k(0) + rocdl.s_waitcnt(traits.LGKMCNT_0_ONLY) + _waitcnt_vm_n(ctx.NUM_DMA_K + ctx.NUM_DMA_V) + rocdl.sched_barrier(0) + rocdl.s_barrier() + rocdl.sched_barrier(0) + + # Epilogue C5 (compute): MMA0 -> v_s_0; fold rescale_e3 into l_row, finish + # v_p_1 softmax. + v_s_0 = gemm_helper.qk(v_k, q_all_wide) + l_row = softmax_helper.apply_l_rescale(l_row, rescale_e3) + v_p_1 = softmax_helper.exp2(v_p_1, 16, 16) + l_row = softmax_helper.reduce_sum(l_row, v_p_1) + v_p_1 = softmax_helper.cast_p(v_p_1) + v_p_1 = softmax_helper.anchor_v_p(v_p_1) + _sched_barrier_exp_pairs(traits, 6, 3, 7) + _sched_barrier_pairs(traits, 10, 5, 7) + rocdl.sched_barrier(0) + rocdl.s_barrier() + rocdl.sched_barrier(0) + + # Epilogue C6 (memory): read V packs (buf1), causal mask v_s_0, sync. + v_packs_e7 = kv_lds_to_regs.load_v(1) + if const_expr(traits.CAUSAL): + v_s_0 = softmax_helper.causal_mask_prologue_if_needed( + v_s_0, + max_m2, + max_m1 * traits.BLOCK_N, + ) + else: + v_s_0 = softmax_helper.seq_pad_mask_if_needed(v_s_0, max_m2) + rocdl.s_waitcnt(traits.LGKMCNT_0_ONLY) + _waitcnt_vm_n(ctx.NUM_DMA_V) + rocdl.sched_barrier(0) + rocdl.s_barrier() + rocdl.sched_barrier(0) + + # Epilogue C7 (compute, mirror of C3): full P*V + unconditional rescale. + if const_expr(traits.DUALWAVE_SWP_SETPRIO): + rocdl.s_setprio(1) + v_o = gemm_helper.pv(v_p_1, v_packs_e7, v_o) + m_tile_max_e7 = softmax_helper.reduce_max(v_s_0) + row_max_e7, rescale_e7 = softmax_helper.rescale_from_tile_max(m_row, m_tile_max_e7) + m_row = row_max_e7 + v_s_0 = softmax_helper.sub_m(v_s_0, row_max_e7) + v_p_0 = softmax_helper.exp2(v_s_0, 0, 16) + _sched_barrier_pairs(traits, 10, 5, 8) + _sched_barrier_exp_pairs(traits, 6, 3, 8) + rocdl.sched_barrier(0) + softmax_helper.scale_o(v_o, rescale_e7) + v_o = softmax_helper.anchor_v_o(v_o) + if const_expr(traits.DUALWAVE_SWP_SETPRIO): + rocdl.s_setprio(0) + rocdl.sched_barrier(0) + rocdl.s_barrier() + rocdl.sched_barrier(0) + + # Epilogue C8 (memory): prefetch V max_m1 (buf1), read K from buf1, sync. + kv_gmem_to_lds.load_v(max_m1 * traits.BLOCK_N, 1) + v_k = kv_lds_to_regs.load_k(1) + rocdl.s_waitcnt(traits.LGKMCNT_0_ONLY) + _waitcnt_vm_n(ctx.NUM_DMA_V) + rocdl.sched_barrier(0) + rocdl.s_barrier() + rocdl.sched_barrier(0) + + # Epilogue C9 (compute): MMA0 -> v_s_1 (last tile); fold rescale_e7 into + # l_row, finish v_p_0 softmax. + v_s_1 = gemm_helper.qk(v_k, q_all_wide) + l_row = softmax_helper.apply_l_rescale(l_row, rescale_e7) + v_p_0 = softmax_helper.exp2(v_p_0, 16, 16) + l_row = softmax_helper.reduce_sum(l_row, v_p_0) + v_p_0 = softmax_helper.cast_p(v_p_0) + v_p_0 = softmax_helper.anchor_v_p(v_p_0) + _sched_barrier_exp_pairs(traits, 6, 3, 9) + _sched_barrier_pairs(traits, 10, 5, 9) + rocdl.sched_barrier(0) + rocdl.s_barrier() + rocdl.sched_barrier(0) + + # Epilogue C10 (memory): read last V packs (buf0), causal mask v_s_1, + # drain all DMAs (vmcnt 0), sync. + v_packs_e11 = kv_lds_to_regs.load_v(0) + if const_expr(traits.CAUSAL): + v_s_1 = softmax_helper.causal_mask_prologue_if_needed( + v_s_1, + max_m1, + ctx.split_t_end * traits.BLOCK_N, + ) + else: + v_s_1 = softmax_helper.seq_pad_mask_if_needed(v_s_1, max_m1) + rocdl.s_waitcnt(traits.LGKMCNT_0_ONLY) + _waitcnt_vm_n(0) + rocdl.sched_barrier(0) + rocdl.s_barrier() + rocdl.sched_barrier(0) + + # Epilogue C11 (compute): full P*V + rescale for v_p_0, then complete the + # last tile's softmax in-place (both exp2 halves, sum, cast) since no + # further pass follows. + v_o = gemm_helper.pv(v_p_0, v_packs_e11, v_o) + m_tile_max_e11 = softmax_helper.reduce_max(v_s_1) + row_max_e11, rescale_e11 = softmax_helper.rescale_from_tile_max(m_row, m_tile_max_e11) + m_row = row_max_e11 + v_s_1 = softmax_helper.sub_m(v_s_1, row_max_e11) + v_p_1 = softmax_helper.exp2(v_s_1, 0, 16) + _sched_barrier_pairs(traits, 9, 6, 10) + _sched_barrier_exp_pairs(traits, 7, 3, 10) + rocdl.sched_barrier(0) + v_p_1 = softmax_helper.exp2(v_p_1, 16, 16) + l_row = softmax_helper.apply_l_rescale(l_row, rescale_e11) + l_row = softmax_helper.reduce_sum(l_row, v_p_1) + v_p_1 = softmax_helper.cast_p(v_p_1) + v_p_1 = softmax_helper.anchor_v_p(v_p_1) + rocdl.sched_barrier(0) + softmax_helper.scale_o(v_o, rescale_e11) + v_o = softmax_helper.anchor_v_o(v_o) + rocdl.s_barrier() + rocdl.sched_barrier(0) + + # Epilogue C12 (memory): read the final V packs for the closing P*V. + v_packs_e13 = kv_lds_to_regs.load_v(1) + rocdl.s_waitcnt(traits.LGKMCNT_0_ONLY) + rocdl.sched_barrier(0) + rocdl.s_barrier() + rocdl.sched_barrier(0) + + # Epilogue C13 (compute): final P*V -> v_o holds the unnormalized output. + v_o = gemm_helper.pv(v_p_1, v_packs_e13, v_o) + + # Normalize by l_row; zero rows become zero instead of NaN. + # Split-K normalizes before packing so O_partial keeps useful mantissa + # range; the combine kernel later applies w_s*l_s. + # HIPREC already folds v_descale into the bf16 vt scratch, so O only needs + # the 1/l normalization here. + inv_l_rcp = rocdl.rcp(T.f32, _raw(l_row)) + inv_l = ArithValue(fx.Float32(l_row) > ctx.c_zero_f).select(inv_l_rcp, ctx.c_zero_f) + softmax_helper.scale_o(v_o, inv_l) + + # CLOSE the phase shift: one extra s_barrier on group A (complement of + # the prologue's group-B barrier) realigns the two groups before the + # store. Disabled -> one plain barrier. + if const_expr(traits.DUALWAVE_SWP_ENABLE_STAGGER): + _stagger_extra_barrier_if_zero(ctx.stagger_i32) # group A: +1 s_barrier -> close the shift + else: + rocdl.s_barrier() + + # 128b stores fuse this lane and its half-wave partner, so each pair + # covers 8 contiguous columns instead of two 64b stores. + if const_expr(not SPLITK): + if const_expr(traits.VARLEN): + output_store.store_final_o_if_valid(v_o, q_row) + else: + output_store.store_final_o(v_o, q_row) + else: + output_store.store_splitk_partial_o(v_o, m_row, l_row, q_row) + + if ctx.q_start < ctx.seqlen_q_v: + _run_q_block() + + if const_expr(SPLITK): + output_store.store_empty_split() + + # Combine kernel: out = sum_s w_s * O_s / sum_s w_s * l_s, w_s = exp2(m_s - m_max). + # One wave row of 32 lanes covers a (b, h, s) row, 4 contiguous cols/lane. + COMBINE_BLOCK = 256 + COMBINE_LANES_PER_ROW = traits.HEAD_DIM // 4 + COMBINE_ROWS_PER_BLOCK = COMBINE_BLOCK // COMBINE_LANES_PER_ROW + + @flyc.kernel(known_block_size=[COMBINE_BLOCK, 1, 1]) + def flash_attn_splitk_combine_kernel( + O: fx.Tensor, # noqa: E741 + WS: fx.Tensor, + batch_size: fx.Int32, + seq_len: fx.Int32, + stride_q_n: fx.Int32, + ): + ctx = DualwaveSplitKCombineContext(traits, O, WS, batch_size, seq_len, stride_q_n) + ctx.init_types_and_constants() + ctx.init_runtime_indices() + ctx.init_thread_mapping(COMBINE_ROWS_PER_BLOCK, COMBINE_LANES_PER_ROW) + ctx.init_workspace() + ctx.init_descriptors() + + combine = DualwaveSplitKCombineHelper(ctx) + m_s, l_s = combine.load_ml_rows() + m_max = combine.reduce_m_max(m_s) + acc, den = combine.accumulate_splits(m_s, l_s, m_max) + o_pack = combine.pack_output(acc, den) + combine.store_output(o_pack) + + @flyc.jit + def launch_flash_attn_dualwave_swp( + Q: fx.Tensor, + K: fx.Tensor, + V: fx.Tensor, + O: fx.Tensor, # noqa: E741 + DebugCounts: fx.Tensor, + CuSeqQ: fx.Tensor, + CuSeqKv: fx.Tensor, + BlockTable: fx.Tensor, + block_table_stride: fx.Int32, + QDescale: fx.Tensor, + KDescale: fx.Tensor, + VDescale: fx.Tensor, + batch_size: fx.Int32, + seq_len: fx.Int32, + seq_len_kv: fx.Int32, + stride_q_n: fx.Int32, + stride_o_n: fx.Int32, + stride_kv_n: fx.Int32, + head_dim_runtime: fx.Int32, + stream: fx.Stream = fx.Stream(None), + ): + # Make shape/mode traits visible to the JIT cache key. + _ = _dualwave_swp_fp8_cache_tag + bs_idx = fx.Index(batch_size) + sl_idx = fx.Index(seq_len) + num_q_blocks = (sl_idx + BLOCK_M - 1) // BLOCK_M + if const_expr(SPLITK): + grid_z = bs_idx * NUM_KV_SPLITS + else: + grid_z = bs_idx + + passthrough_entries = ( + [ + ["denormal-fp-math-f32", "preserve-sign,preserve-sign"], + ["no-nans-fp-math", "true"], + ["unsafe-fp-math", "true"], + ] + if const_expr(daz) + else None + ) + flash_attn_dualwave_swp_fp8_gfx950_kernel( + Q, + K, + V, + O, + DebugCounts, + CuSeqQ, + CuSeqKv, + BlockTable, + block_table_stride, + QDescale, + KDescale, + VDescale, + seq_len, + seq_len_kv, + stride_q_n, + stride_o_n, + stride_kv_n, + head_dim_runtime, + value_attrs={ + "rocdl.waves_per_eu": waves_per_eu, + "rocdl.flat_work_group_size": f"{BLOCK_SIZE},{BLOCK_SIZE}", + "passthrough": passthrough_entries, + }, + ).launch( + grid=(NUM_HEADS_Q, num_q_blocks, grid_z), + block=(BLOCK_SIZE, 1, 1), + stream=stream, + ) + if const_expr(SPLITK): + combine_rows = bs_idx * NUM_HEADS_Q * sl_idx + flash_attn_splitk_combine_kernel(O, DebugCounts, batch_size, seq_len, stride_o_n).launch( + grid=(combine_rows // COMBINE_ROWS_PER_BLOCK, 1, 1), + block=(COMBINE_BLOCK, 1, 1), + stream=stream, + ) + + _dualwave_swp_compile_hints = { + "fast_fp_math": True, + "unsafe_fp_math": True, + "llvm_options": { + "enable-post-misched": False, + "lsr-drop-solution": True, + }, + } + + def _launch( + Q, + K, + V, + O, # noqa: E741 + batch_size, + seq_len, + stride_kv_n=None, + stride_q_n=None, + stride_o_n=None, + head_dim_runtime=None, + debug_counts=None, + *, + seq_len_kv=None, + workspace=None, + cu_seqlens_q=None, + cu_seqlens_kv=None, + block_table=None, + block_table_stride=None, + q_descale=None, + k_descale=None, + v_descale=None, + stream=None, + ): + if stride_kv_n is None: + stride_kv_n = DEFAULT_STRIDE_KV_N + if stride_q_n is None: + stride_q_n = DEFAULT_STRIDE_Q_N + if stride_o_n is None: + stride_o_n = DEFAULT_STRIDE_O_N + if head_dim_runtime is None: + head_dim_runtime = HEAD_DIM + # seq_len_kv defaults to seq_len (self-attention / equal Q,KV lengths). + if seq_len_kv is None: + seq_len_kv = seq_len + if SPLITK: + if workspace is None: + raise ValueError("num_kv_splits > 1 requires a fp32 workspace (see dualwave_splitk_workspace_elems)") + debug_counts = workspace + if debug_counts is None: + debug_counts = O + # Dense launches still pass valid tensors for the (unused) cu_seqlens slots; + # the kernel only reads them under const_expr(VARLEN). Use O as a placeholder. + if cu_seqlens_q is None: + cu_seqlens_q = O + if cu_seqlens_kv is None: + cu_seqlens_kv = O + if block_table is None: + block_table = O + if block_table_stride is None: + block_table_stride = 0 + if PAGED and block_table is O: + raise ValueError("paged fp8 flash_attn requires block_table") + # Per-tensor fp8 descales (shape-[1] fp32). The kernel only reads them on + # the fp8 path; bf16/f16 launches pass O as an unused placeholder. + if q_descale is None: + q_descale = O + if k_descale is None: + k_descale = O + if v_descale is None: + v_descale = O + with CompilationContext.compile_hints(_dualwave_swp_compile_hints): + if stream is None: + return launch_flash_attn_dualwave_swp( + Q, + K, + V, + O, + debug_counts, + cu_seqlens_q, + cu_seqlens_kv, + block_table, + block_table_stride, + q_descale, + k_descale, + v_descale, + batch_size, + seq_len, + seq_len_kv, + stride_q_n, + stride_o_n, + stride_kv_n, + head_dim_runtime, + ) + return launch_flash_attn_dualwave_swp( + Q, + K, + V, + O, + debug_counts, + cu_seqlens_q, + cu_seqlens_kv, + block_table, + block_table_stride, + q_descale, + k_descale, + v_descale, + batch_size, + seq_len, + seq_len_kv, + stride_q_n, + stride_o_n, + stride_kv_n, + head_dim_runtime, + stream=stream, + ) + + def _compile( + Q, + K, + V, + O, # noqa: E741 + batch_size, + seq_len, + stride_kv_n=None, + stride_q_n=None, + stride_o_n=None, + head_dim_runtime=None, + debug_counts=None, + *, + seq_len_kv=None, + workspace=None, + cu_seqlens_q=None, + cu_seqlens_kv=None, + block_table=None, + block_table_stride=None, + q_descale=None, + k_descale=None, + v_descale=None, + stream=None, + ): + if stride_kv_n is None: + stride_kv_n = DEFAULT_STRIDE_KV_N + if stride_q_n is None: + stride_q_n = DEFAULT_STRIDE_Q_N + if stride_o_n is None: + stride_o_n = DEFAULT_STRIDE_O_N + if head_dim_runtime is None: + head_dim_runtime = HEAD_DIM + if seq_len_kv is None: + seq_len_kv = seq_len + if SPLITK: + if workspace is None: + raise ValueError("num_kv_splits > 1 requires a fp32 workspace (see dualwave_splitk_workspace_elems)") + debug_counts = workspace + if debug_counts is None: + debug_counts = O + if cu_seqlens_q is None: + cu_seqlens_q = O + if cu_seqlens_kv is None: + cu_seqlens_kv = O + if block_table is None: + block_table = O + if block_table_stride is None: + block_table_stride = 0 + if q_descale is None: + q_descale = O + if k_descale is None: + k_descale = O + if v_descale is None: + v_descale = O + with CompilationContext.compile_hints(_dualwave_swp_compile_hints): + return flyc.compile( + launch_flash_attn_dualwave_swp, + Q, + K, + V, + O, + debug_counts, + cu_seqlens_q, + cu_seqlens_kv, + block_table, + block_table_stride, + q_descale, + k_descale, + v_descale, + batch_size, + seq_len, + seq_len_kv, + stride_q_n, + stride_o_n, + stride_kv_n, + head_dim_runtime, + fx.Stream(stream), + ) + + _launch.compile = _compile + + return _launch diff --git a/kernels/attention/flash_attn_interface.py b/kernels/attention/flash_attn_interface.py index 13bccd3e76..3cc95c6f8d 100644 --- a/kernels/attention/flash_attn_interface.py +++ b/kernels/attention/flash_attn_interface.py @@ -464,6 +464,41 @@ def _build_paged( ) +@functools.lru_cache(maxsize=64) +def _build_paged_fp8( + num_heads: int, + num_kv_heads: int, + head_dim: int, + value_head_dim: int, + waves_per_eu: int, + daz: bool, + lazy_rescale: bool, + setprio: bool, + enable_stagger: bool, +): + """Build the gfx950 packed-varlen, vectorized page-64 FP8 launcher.""" + from kernels.attention.flash_attn_fp8_paged_gfx950 import build_flash_attn_paged_fp8_module + + return build_flash_attn_paged_fp8_module( + num_heads=num_heads, + num_kv_heads=num_kv_heads, + head_dim=head_dim, + value_head_dim=value_head_dim, + causal=True, + dtype_str="fp8", + waves_per_eu=waves_per_eu, + daz=daz, + dualwave_swp_lazy_rescale=lazy_rescale, + dualwave_swp_setprio=setprio, + dualwave_swp_enable_stagger=enable_stagger, + num_kv_splits=1, + varlen=True, + cross_seqlen=True, + paged=True, + kv_cache_layout="vectorized", + ) + + # ── paged-KV native path ──────────────────────────────────────────────────── # gfx950 dualwave paged-KV currently supports exactly one configuration. @@ -488,6 +523,9 @@ def _flydsl_flash_attn_paged( max_seqlen_q: Optional[int], cross_seqlen: Optional[bool], num_kv_splits: int, + q_descale: Optional[torch.Tensor], + k_descale: Optional[torch.Tensor], + v_descale: Optional[torch.Tensor], out: Optional[torch.Tensor], waves_per_eu: int, daz: bool, @@ -499,8 +537,9 @@ def _flydsl_flash_attn_paged( """Native paged-KV attention on the gfx950 dualwave kernel. Supported config ONLY (anything else raises): linear/vectorized cache layout - [NumBlocks, PageSize=64, NumKVHeads, HeadDim], vLLM lookup (block_table + - seqlen_k), causal, D=64/128, dtype bf16/f16. + with page size 64 and vLLM ``block_table`` / ``seqlen_k`` metadata. BF16/F16 + support D64/D128. gfx950 FP8 additionally supports packed-varlen causal + Q/K D192 with vectorized V/output D128 or D192. - Dense 4D Q ``[B, Sq, H, D]``: split-K (num_kv_splits>1) supported (seq_len>=384). - Varlen packed Q ``[total_q, H, D]`` (cu_seqlens_q given): paged K/V looked up per kv-tile via block_table; split-K not supported (matches dense varlen). @@ -522,8 +561,9 @@ def _flydsl_flash_attn_paged( f"flydsl_flash_attn_func: linear paged K/V must be 4D [NumBlocks,PageSize,Hkv,D], got {k.dim()}D" ) - varlen = cu_seqlens_q is not None dtype_str = _dtype_str(q) + paged_fp8 = dtype_str == "fp8" + varlen = cu_seqlens_q is not None if varlen: # Packed varlen Q: [total_q, H, D]. Per-batch ranges come from cu_seqlens # inside the kernel; grid_y is sized by max_seqlen_q. @@ -543,24 +583,62 @@ def _flydsl_flash_attn_paged( raise ValueError(f"flydsl_flash_attn_func: paged dense q must be 4D [B,Sq,H,D], got {q.dim()}D") B, Sq, H, D = q.shape if vectorized: - kvs = 16 // q.element_size() + kvs = 16 // k.element_size() Hkv = int(k.shape[1]) page_size = int(k.shape[3]) k_head_dim = int(k.shape[2]) * int(k.shape[4]) # (D/kVS) * kVS if int(k.shape[4]) != kvs: raise ValueError(f"flydsl_flash_attn_func: vectorized K last dim ({k.shape[4]}) must equal kVS={kvs}") + value_head_dim = int(v.shape[3]) + expected_v_tail = (Hkv, page_size // kvs, value_head_dim, kvs) + if tuple(v.shape[1:]) != expected_v_tail: + raise ValueError( + f"flydsl_flash_attn_func: vectorized V tail must be {expected_v_tail}, got {tuple(v.shape[1:])}" + ) else: page_size = int(k.shape[1]) Hkv = int(k.shape[2]) k_head_dim = int(k.shape[3]) + value_head_dim = int(v.shape[3]) + if tuple(v.shape[1:3]) != (page_size, Hkv): + raise ValueError( + f"flydsl_flash_attn_func: linear V must match K page/head axes, got K{tuple(k.shape)} V{tuple(v.shape)}" + ) if page_size != _PAGED_PAGE_SIZE: raise NotImplementedError( f"flydsl_flash_attn_func: native paged KV supports page_size={_PAGED_PAGE_SIZE} only, got {page_size}" ) - if D not in (64, 128): - raise NotImplementedError(f"flydsl_flash_attn_func: native paged KV supports head_dim=64 or 128, got {D}") if k_head_dim != D: raise ValueError(f"flydsl_flash_attn_func: paged K head_dim ({k_head_dim}) must match q head_dim ({D})") + if paged_fp8: + _arch = _gpu_arch(q.device) + if not _arch.startswith("gfx950"): + raise ValueError(f"flydsl_flash_attn_func: paged FP8 requires gfx950, got '{_arch or 'unknown'}'") + if not ( + causal and varlen and cross_seqlen is not False and vectorized and D == 192 and value_head_dim in (128, 192) + ): + raise NotImplementedError( + "flydsl_flash_attn_func: paged FP8 requires causal packed-varlen vectorized KV, " + f"Q/K D192, and V D128/D192; got causal={causal}, varlen={varlen}, " + f"layout={kv_cache_layout!r}, Q/K D{D}, V D{value_head_dim}" + ) + if num_kv_splits != 1: + raise NotImplementedError("flydsl_flash_attn_func: paged FP8 does not support split-K") + if bias is not None: + raise NotImplementedError("flydsl_flash_attn_func: paged FP8 does not support bias") + if any(x is None for x in (q_descale, k_descale, v_descale)): + raise ValueError("flydsl_flash_attn_func: paged FP8 requires q_descale, k_descale, and v_descale") + for name, scale in (("q_descale", q_descale), ("k_descale", k_descale), ("v_descale", v_descale)): + if scale.device != q.device or scale.dtype != torch.float32 or scale.numel() != 1: + raise ValueError( + f"flydsl_flash_attn_func: {name} must be one float32 value on {q.device}, " + f"got shape={tuple(scale.shape)} dtype={scale.dtype} device={scale.device}" + ) + elif D not in (64, 128) or value_head_dim != D: + raise NotImplementedError( + "flydsl_flash_attn_func: BF16/F16 paged KV requires matching K/V head_dim 64 or 128, " + f"got Q/K D{D}, V D{value_head_dim}" + ) if num_kv_heads is None: num_kv_heads = Hkv @@ -583,7 +661,7 @@ def _flydsl_flash_attn_paged( skv = int(max_seqlen_kv) if max_seqlen_kv is not None else int(max(_kv_lens)) max_kv_pages = (skv + page_size - 1) // page_size max_pages_per_split = (max_kv_pages + int(num_kv_splits) - 1) // int(num_kv_splits) - if max_pages_per_split > _PAGED_BT_LDS_SIZE: + if not paged_fp8 and max_pages_per_split > _PAGED_BT_LDS_SIZE: max_supported_kv = _PAGED_BT_LDS_SIZE * int(num_kv_splits) * page_size raise NotImplementedError( f"flydsl_flash_attn_func: paged KV length {skv} exceeds block-table LDS window " @@ -618,6 +696,16 @@ def _flydsl_flash_attn_paged( f"flydsl_flash_attn_func: paged bias needs >= max_seqlen_kv={skv} columns, got {bias.shape[1]}" ) + if block_table.dim() != 2 or block_table.device != q.device: + raise ValueError( + f"flydsl_flash_attn_func: block_table must be 2D on {q.device}, " + f"got shape={tuple(block_table.shape)} device={block_table.device}" + ) + if seqlen_k.dtype != torch.int32 or seqlen_k.device != q.device or seqlen_k.numel() != B: + raise ValueError( + f"flydsl_flash_attn_func: seqlen_k must be int32 [{B}] on {q.device}, " + f"got shape={tuple(seqlen_k.shape)} dtype={seqlen_k.dtype} device={seqlen_k.device}" + ) block_table_stride = int(block_table.shape[1]) # Flatten so the kernel's flat row-major index addresses block_table correctly. block_table_i32 = ( @@ -635,7 +723,19 @@ def _flydsl_flash_attn_paged( and dtype_str in ("bf16", "f16") and (not _arch.startswith("gfx950") or Sq <= _VARLEN_LIGHT_MAX_SEQ) ) - if _paged_light_ok: + if paged_fp8: + exe = _build_paged_fp8( + num_heads=H, + num_kv_heads=num_kv_heads, + head_dim=D, + value_head_dim=value_head_dim, + waves_per_eu=waves_per_eu, + daz=daz, + lazy_rescale=dualwave_swp_lazy_rescale, + setprio=dualwave_swp_setprio, + enable_stagger=dualwave_swp_enable_stagger, + ) + elif _paged_light_ok: exe = _build_paged_light( num_heads=H, num_kv_heads=num_kv_heads, @@ -670,15 +770,31 @@ def _flydsl_flash_attn_paged( kv_cache_layout=kv_cache_layout, has_bias=bias is not None, ) + expected_out_shape = (*q.shape[:-1], value_head_dim) if out is None: - out = torch.empty_like(q) - # Keep tensors in natural shape; flattening can overflow int32 C-ABI dims. - # The paged kernel rebuilds per-batch/page descriptors from base pointers. - q_flat = q.contiguous() + out_dtype = torch.bfloat16 if paged_fp8 else q.dtype + out = torch.empty(expected_out_shape, dtype=out_dtype, device=q.device) + elif out.shape != expected_out_shape or not out.is_contiguous(): + raise ValueError( + f"flydsl_flash_attn_func: paged output must be contiguous with shape {expected_out_shape}, " + f"got shape={tuple(out.shape)} strides={out.stride()}" + ) + elif paged_fp8 and out.dtype != torch.bfloat16: + raise ValueError(f"flydsl_flash_attn_func: paged FP8 output must be bf16, got {out.dtype}") + elif not paged_fp8 and out.dtype != q.dtype: + raise ValueError( + f"flydsl_flash_attn_func: paged output dtype must match q dtype {q.dtype}, got {out.dtype}" + ) + # Keep serving-sized physical K/V caches rank-5 because flattening their + # dynamic memref shape can exceed signed int32. The FP8 schedule consumes + # Q/O as flat token-major buffers, matching its explicit runtime strides. + q_flat = q.contiguous().view(-1) if paged_fp8 else q.contiguous() k_flat = k.contiguous() v_flat = v.contiguous() - o_flat = out.contiguous() + o_flat = out.contiguous().view(-1) if paged_fp8 else out.contiguous() kwargs = dict(block_table=block_table_i32, block_table_stride=block_table_stride, stream=launch_stream) + if paged_fp8: + kwargs.update(q_descale=q_descale, k_descale=k_descale, v_descale=v_descale) if bias is not None: kwargs["bias"] = bias if varlen: @@ -773,7 +889,7 @@ def flydsl_flash_attn_func( # Per-head attention-sink logit: one extra softmax denominator term with no # matching V row. Same path and restrictions as `alibi_slopes`; combinable. sink: Optional[torch.Tensor] = None, - # fp8 dense ABI: per-tensor descales for pre-quantized e4m3fn Q/K/V. + # fp8 ABI: per-tensor descales for pre-quantized e4m3fn Q/K/V. q_descale: Optional[torch.Tensor] = None, k_descale: Optional[torch.Tensor] = None, v_descale: Optional[torch.Tensor] = None, @@ -806,7 +922,7 @@ def flydsl_flash_attn_func( k: Key tensor. Dense: ``[B, Skv, Hkv, D]``. Varlen: ``[total_kv, Hkv, D]``. v: Value tensor, same shape as k. - Paged KV cache (future ABI): physical K/V cache tensors. Supported + Paged KV cache: physical K/V cache tensors. Supported ``kv_cache_layout`` values: - ``linear``: 4D paged K/V, ``[NumBlocks, PageSize, NumKVHeads, HeadDim]``. - ``linear3d``: page_size=1 special case, @@ -828,7 +944,8 @@ def flydsl_flash_attn_func( dense mode infers it from ``q.shape[1] != k.shape[1]``. block_table / seqlen_k: vLLM-style 2D block table metadata. Enables the native paged-KV path, which supports ``bias`` but not - ``alibi_slopes``, ``sink``, ``return_lse``, or fp8. + ``alibi_slopes``, ``sink``, or ``return_lse``. gfx950 FP8 supports + the causal packed-varlen vectorized page-64 D192 path with V D128/D192. num_kv_splits: Split-K factor (>1: gfx950 only, D=64/128, bf16/f16, seq>=384). ``None`` lets fp8 autotune it; ``1`` keeps the kernel unsplit. fp8_block_m: Pin the fp8 tile height to 128 or 256. fp8 only. @@ -872,7 +989,7 @@ def flydsl_flash_attn_func( with paged KV or fp8 -- but freely combinable with ``bias`` and ``alibi_slopes``. q_descale / k_descale / v_descale: fp32 shape-[1] descales required - for dense fp8 e4m3fn inputs. + for dense or paged fp8 e4m3fn inputs. out: Optional pre-allocated output tensor. For fp8, output is bf16; otherwise it has the same dtype as q. waves_per_eu: Kernel occupancy hint. @@ -885,8 +1002,9 @@ def flydsl_flash_attn_func( stream: CUDA/HIP stream to launch on. Returns: - Output tensor with the same shape as q. The dtype is bf16 for fp8 - inputs, otherwise the same dtype as q. When ``return_lse=True`` returns + Output tensor with the same shape as q except that paged asymmetric-V + attention uses V's final dimension. The dtype is bf16 for fp8 inputs, + otherwise the same dtype as q. When ``return_lse=True`` returns ``(out, lse)`` where ``lse`` is fp32 ``[B, num_heads, Sq]`` (varlen: ``[B, num_heads, max_seqlen_q]``, padded) holding the per-row natural-log, scale-folded log-sum-exp. @@ -906,8 +1024,6 @@ def flydsl_flash_attn_func( if return_lse and dtype_str == "fp8": raise NotImplementedError("flydsl_flash_attn_func: return_lse is not supported for fp8") paged_kv = any(x is not None for x in (block_table, seqlen_k)) - if dtype_str == "fp8" and paged_kv: - raise NotImplementedError("flydsl_flash_attn_func: fp8 flash_attn does not support paged KV") if return_lse and paged_kv: raise NotImplementedError("flydsl_flash_attn_func: return_lse is not supported for paged KV") has_bias = bias is not None @@ -1006,6 +1122,9 @@ def flydsl_flash_attn_func( max_seqlen_q=max_seqlen_q, cross_seqlen=cross_seqlen, num_kv_splits=num_kv_splits, + q_descale=q_descale, + k_descale=k_descale, + v_descale=v_descale, out=out, waves_per_eu=waves_per_eu, daz=daz, diff --git a/kernels/attention/flash_attn_utils.py b/kernels/attention/flash_attn_utils.py index 10ed3c4be3..2d4fe8a24b 100644 --- a/kernels/attention/flash_attn_utils.py +++ b/kernels/attention/flash_attn_utils.py @@ -1827,6 +1827,15 @@ class DualwaveSwpFp8Traits: LGKMCNT_0_ONLY: int XCD_SWIZZLE: bool = False BATCH_INTERLEAVE_GROUP: int = 1 + PAGED: bool = False + KV_VECTORIZED: bool = False + QREG: bool = True + FP8_PV: bool = True + FP8_PV_DIRECT: bool = True + + @property + def V_HEAD_DIM(self): + return self.HEAD_DIM_V @property def cache_tag(self): @@ -2048,6 +2057,333 @@ def dualwave_fp8_dma_per_iter(traits): return 2 * k_instr + 2 * v_instr_min +@dataclass(frozen=True) +class PagedDualwaveSwpFp8Traits: + """Pure compile-time tile/layout constants for the gfx950 DUALWAVE_SWP fp8 kernel. + + fp8 runs a single path: WIDE QK (32x32x64 mfma_scale) feeding HIPREC PV (fp8 V + dequantized into a bf16 ``vt`` LDS scratch, then a bf16 PV MMA). The ``*_BF`` + fields describe that bf16 vt layout; ``ELEM_BYTES`` is 1 (Q/K/V are fp8).""" + + BLOCK_M: int + BLOCK_N: int + K_SUB_N: int + WARP_SIZE: int + NUM_WAVES: int + BLOCK_SIZE: int + ROWS_PER_WAVE: int + HEAD_DIM: int + V_HEAD_DIM: int + D_CHUNK: int + D_CHUNKS: int + PV_K_STEPS: int + NUM_HEADS_Q: int + NUM_HEADS_KV: int + GQA_GROUP_SIZE: int + CAUSAL: bool + DTYPE_STR: str + WAVES_PER_EU: int + DAZ: bool + DUALWAVE_SWP_LAZY_RESCALE: bool + DUALWAVE_SWP_SETPRIO: bool + DUALWAVE_SWP_DEBUG_LAZY_COUNTS: bool + DUALWAVE_SWP_ENABLE_STAGGER: bool + NUM_KV_SPLITS: int + SPLITK: bool + VARLEN: bool + CROSS_SEQLEN: bool + FP8_PV: bool + FP8_PV_DIRECT: bool + BN128: bool + BN128_PF: bool + QREG: bool + VDMA: bool + PAGED: bool + KV_CACHE_LAYOUT: str + KV_VECTORIZED: bool + PAGE_SIZE: int + DEFAULT_STRIDE_Q_N: int + DEFAULT_STRIDE_KV_N: int + DMA_BYTES: int + ELEM_BYTES: int + OUT_ELEM_BYTES: int + D_128B_SIZE: int + VEC_KV: int + KV_VEC_SIZE: int + LANE_SPLIT_KV: int + SMEM_N_PER_WAVE: int + SMEM_N_RPT: int + SMEM_D_RPT: int + SMEM_K_LINE_STRIDE: int + SMEM_K_TILE_ELEMS: int + NUM_PREFETCH_K: int + DUALWAVE_SWP_KV_PER_BUFFER: int + LDS_KV_TOTAL_SIZE: int + DUALWAVE_SWP_K_BUF_BASE: tuple[int, int] + DUALWAVE_SWP_V_BUF_BASE: tuple[int, int] + # bf16 vt scratch layout (HIPREC V dequant target + transpose read strides). + EB_BF: int + D128_BF: int + VEC_BF: int + SDRPT_BF: int + SNRPT_BF: int + VLS_BF: int + VEC_V_ROW_STRIDE: int + VT_BF16_ELEMS: int + VT_BF16_TOTAL: int + URV_GRPK_BF: int + URV_GRP_N_BF: int + URV_LANE_LO_BF: int + URV_LANE_HI_BF: int + URV_STEPK_BF: int + URV_DC_AXIS0_BF: int + URV_DC_AXIS1_BF: int + URV_I5_BF: int + DUALWAVE_SWP_RESCALE_THRESHOLD: float + SCHED_MFMA_MASK: int + SCHED_VALU_MASK: int + SCHED_EXP_MASK: int + SCHED_DS_READ_MASK: int + LDS_SCOPE_NAMES: tuple[str, str, str, str] + NEG_INF_F32_BITS: int + LGKMCNT_0_ONLY: int + XCD_SWIZZLE: bool = False + + @property + def HEAD_DIM_V(self): + return self.V_HEAD_DIM + + @property + def QLDS(self): + return self.QREG + + @property + def cache_tag(self): + return ( + self.NUM_HEADS_Q, + self.NUM_HEADS_KV, + self.HEAD_DIM, + self.V_HEAD_DIM, + self.CAUSAL, + self.DTYPE_STR, + self.WAVES_PER_EU, + self.DAZ, + self.DUALWAVE_SWP_LAZY_RESCALE, + self.DUALWAVE_SWP_RESCALE_THRESHOLD, + self.DUALWAVE_SWP_SETPRIO, + self.DUALWAVE_SWP_DEBUG_LAZY_COUNTS, + self.DUALWAVE_SWP_ENABLE_STAGGER, + self.NUM_KV_SPLITS, + self.SPLITK, + self.VARLEN, + self.CROSS_SEQLEN, + self.PAGED, + self.KV_CACHE_LAYOUT, + self.KV_VECTORIZED, + "fp8_wide_qk_hiprec_pv_paged_vec5d_v2" if self.PAGED else "fp8_wide_qk_hiprec_pv", + self.ELEM_BYTES, + self.OUT_ELEM_BYTES, + self.LANE_SPLIT_KV, + self.VT_BF16_ELEMS, + self.VT_BF16_TOTAL, + self.FP8_PV, + self.FP8_PV_DIRECT, + self.NUM_PREFETCH_K, + self.BN128, + self.BN128_PF, + self.QREG, + self.VDMA, + self.XCD_SWIZZLE, + ) + + +def _make_paged_dualwave_swp_fp8_traits( + num_heads, + num_kv_heads, + head_dim, + rescale_threshold, + value_head_dim=None, + causal=True, + waves_per_eu=2, + daz=True, + dualwave_swp_lazy_rescale=True, + dualwave_swp_setprio=True, + dualwave_swp_debug_lazy_counts=False, + dualwave_swp_enable_stagger=True, + num_kv_splits=1, + varlen=False, + cross_seqlen=False, + xcd_swizzle=False, + paged=False, + kv_cache_layout="linear", +): + """Build gfx950 DUALWAVE_SWP fp8 compile-time layout traits (dtype fixed to fp8).""" + # Tile shape and wave geometry follow the gfx950 dual-wave 8-wave CTA. + block_m = 256 + block_n = 64 + k_sub_n = 32 + warp_size = 64 + num_waves = 8 + block_size = num_waves * warp_size + rows_per_wave = 32 + + if value_head_dim is None: + value_head_dim = head_dim + d_chunk = 32 + d_chunks = value_head_dim // d_chunk + pv_k_step = 16 + pv_k_steps = k_sub_n // pv_k_step + + gqa_group_size = num_heads // num_kv_heads + default_stride_q_n = num_heads * head_dim + default_stride_kv_n = num_kv_heads * head_dim + + # fp8: Q/K/V are 1B; O is bf16 (2B). ELEM_BYTES=1 drives the fp8 address math. + elem_bytes = 1 + out_elem_bytes = 2 + d_128b_size = 128 // elem_bytes + vec_kv = 16 // elem_bytes + lane_split_kv = 8 + smem_linear_wave = warp_size * 16 // elem_bytes + smem_n_per_wave = smem_linear_wave // d_128b_size + smem_n_rpt = block_n // smem_n_per_wave + # D192 needs a second 128-byte K repeat. In vectorized page-64 mode the + # final four 16-byte lane buckets naturally land beyond the page descriptor + # and read zero, giving the QK MFMA a correctly padded D256 LDS tile without + # ever crossing into another physical page. + smem_d_rpt = (head_dim + d_128b_size - 1) // d_128b_size + smem_k_pad = 16 // elem_bytes + smem_v_pad = 64 // elem_bytes + smem_k_line_stride = smem_linear_wave + smem_k_pad + smem_v_line_stride = smem_linear_wave + smem_v_pad + smem_k_tile_elems = smem_n_rpt * smem_d_rpt * smem_k_line_stride + smem_v_tile_elems = smem_n_rpt * smem_d_rpt * smem_v_line_stride + bn128 = (num_kv_splits <= 1) and (not varlen) + bn128_pf = bn128 + qreg = bn128_pf + vdma = bn128_pf + deep_ring = bn128 + num_prefetch_k = (6 if bn128_pf else 4) if deep_ring else 2 + if bn128_pf: + dualwave_swp_kv_per_buffer = smem_k_tile_elems + else: + dualwave_swp_kv_per_buffer = smem_k_tile_elems + smem_v_tile_elems + lds_kv_total_size = num_prefetch_k * dualwave_swp_kv_per_buffer + dualwave_swp_k_buf_base = tuple(i * dualwave_swp_kv_per_buffer for i in range(num_prefetch_k)) + dualwave_swp_v_buf_base = tuple(smem_k_tile_elems + i * dualwave_swp_kv_per_buffer for i in range(num_prefetch_k)) + + # bf16 vt scratch layout: HIPREC dequantizes fp8 V into these positions so the + # proven bf16 V transpose read (ds_read_tr16) + bf16 PV MMA are reused unchanged. + eb_bf = 2 + d128_bf = 128 // eb_bf + vec_bf = 16 // eb_bf + slw_bf = warp_size * 16 // eb_bf + snrpt_bf = block_n // (slw_bf // d128_bf) + sdrpt_bf = value_head_dim // d128_bf + vls_bf = slw_bf + 64 // eb_bf + vt_bf16_elems = snrpt_bf * sdrpt_bf * vls_bf + fp8_v_tile_bytes = (block_n // 8) * (head_dim // 16) * 128 + if bn128_pf: + vt_bf16_total = num_prefetch_k * (fp8_v_tile_bytes // eb_bf) + 128 + else: + vt_bf16_total = (2 if deep_ring else num_prefetch_k) * vt_bf16_elems + + splitk = num_kv_splits > 1 + paged = bool(paged) + kv_vectorized = paged and kv_cache_layout == "vectorized" + page_size = block_n + + fp8_pv = not paged and os.getenv("FLYDSL_FA_FP8_PV", "0") == "1" + fp8_pv_direct = bn128 and not paged + if fp8_pv_direct: + fp8_pv = True + + return PagedDualwaveSwpFp8Traits( + BLOCK_M=block_m, + BLOCK_N=block_n, + K_SUB_N=k_sub_n, + WARP_SIZE=warp_size, + NUM_WAVES=num_waves, + BLOCK_SIZE=block_size, + ROWS_PER_WAVE=rows_per_wave, + HEAD_DIM=head_dim, + V_HEAD_DIM=value_head_dim, + D_CHUNK=d_chunk, + D_CHUNKS=d_chunks, + PV_K_STEPS=pv_k_steps, + NUM_HEADS_Q=num_heads, + NUM_HEADS_KV=num_kv_heads, + GQA_GROUP_SIZE=gqa_group_size, + CAUSAL=causal, + DTYPE_STR="fp8", + WAVES_PER_EU=waves_per_eu, + DAZ=bool(daz), + DUALWAVE_SWP_LAZY_RESCALE=bool(dualwave_swp_lazy_rescale), + DUALWAVE_SWP_SETPRIO=bool(dualwave_swp_setprio), + DUALWAVE_SWP_DEBUG_LAZY_COUNTS=bool(dualwave_swp_debug_lazy_counts), + DUALWAVE_SWP_ENABLE_STAGGER=bool(dualwave_swp_enable_stagger), + NUM_KV_SPLITS=num_kv_splits, + SPLITK=splitk, + VARLEN=bool(varlen), + CROSS_SEQLEN=bool(cross_seqlen), + FP8_PV=fp8_pv, + FP8_PV_DIRECT=bool(fp8_pv_direct), + BN128=bool(bn128), + BN128_PF=bool(bn128_pf), + QREG=bool(qreg), + VDMA=bool(vdma), + PAGED=paged, + KV_CACHE_LAYOUT=kv_cache_layout, + KV_VECTORIZED=kv_vectorized, + PAGE_SIZE=page_size, + DEFAULT_STRIDE_Q_N=default_stride_q_n, + DEFAULT_STRIDE_KV_N=default_stride_kv_n, + DMA_BYTES=16, + ELEM_BYTES=elem_bytes, + OUT_ELEM_BYTES=out_elem_bytes, + D_128B_SIZE=d_128b_size, + VEC_KV=vec_kv, + KV_VEC_SIZE=vec_kv, + LANE_SPLIT_KV=lane_split_kv, + SMEM_N_PER_WAVE=smem_n_per_wave, + SMEM_N_RPT=smem_n_rpt, + SMEM_D_RPT=smem_d_rpt, + SMEM_K_LINE_STRIDE=smem_k_line_stride, + SMEM_K_TILE_ELEMS=smem_k_tile_elems, + NUM_PREFETCH_K=num_prefetch_k, + DUALWAVE_SWP_KV_PER_BUFFER=dualwave_swp_kv_per_buffer, + LDS_KV_TOTAL_SIZE=lds_kv_total_size, + DUALWAVE_SWP_K_BUF_BASE=dualwave_swp_k_buf_base, + DUALWAVE_SWP_V_BUF_BASE=dualwave_swp_v_buf_base, + EB_BF=eb_bf, + D128_BF=d128_bf, + VEC_BF=vec_bf, + SDRPT_BF=sdrpt_bf, + SNRPT_BF=snrpt_bf, + VLS_BF=vls_bf, + VEC_V_ROW_STRIDE=vls_bf, + VT_BF16_ELEMS=vt_bf16_elems, + VT_BF16_TOTAL=vt_bf16_total, + URV_GRPK_BF=4 * vls_bf, + URV_GRP_N_BF=16, + URV_LANE_LO_BF=4, + URV_LANE_HI_BF=vls_bf, + URV_STEPK_BF=128, + URV_DC_AXIS0_BF=snrpt_bf * vls_bf, + URV_DC_AXIS1_BF=32, + URV_I5_BF=d128_bf, + DUALWAVE_SWP_RESCALE_THRESHOLD=rescale_threshold, + SCHED_MFMA_MASK=0x008, + SCHED_VALU_MASK=0x002, + SCHED_EXP_MASK=0x400, + SCHED_DS_READ_MASK=0x100, + LDS_SCOPE_NAMES=("lds_k0", "lds_k1", "lds_v0", "lds_v1"), + NEG_INF_F32_BITS=0xFF800000, + LGKMCNT_0_ONLY=0xC07F, + XCD_SWIZZLE=bool(xcd_swizzle), + ) + + # Kernel context @@ -4365,6 +4701,9 @@ def __init__( stride_q_n=None, stride_kv_n=None, head_dim_runtime=None, + stride_o_n=None, + BlockTable=None, + block_table_stride=None, ): if isinstance(traits_or_ctx, DualwaveFp8KernelContext): self.__dict__.update(traits_or_ctx.__dict__) @@ -4385,8 +4724,11 @@ def __init__( self.seq_len = seq_len self.seq_len_kv = seq_len_kv self.stride_q_n = stride_q_n + self.stride_o_n = stride_q_n if stride_o_n is None else stride_o_n self.stride_kv_n = stride_kv_n self.head_dim_runtime = head_dim_runtime + self.BlockTable = BlockTable + self.block_table_stride = block_table_stride def init_types_and_constants(self): traits = self.traits @@ -4398,7 +4740,12 @@ def init_types_and_constants(self): self.v2i32_type = Vec.make_type(2, fx.Int32) self.p_elem = fx.BFloat16 self.v4bf16_type = Vec.make_type(4, fx.BFloat16) - self.NUM_DMA_K = len(traits.K_BAND_CHUNK) + self.v8bf16_type = Vec.make_type(8, fx.BFloat16) + if const_expr(traits.PAGED): + self.NUM_DMA_K = traits.SMEM_D_RPT + self.NUM_DMA_V = traits.SMEM_D_RPT + else: + self.NUM_DMA_K = len(traits.K_BAND_CHUNK) self.c_neg_inf = fx.Float32(float("-inf")) self.c_neg_floor = fx.Float32(-3.0e38) self.c_zero_f = fx.Float32(0.0) @@ -4410,8 +4757,11 @@ def init_runtime_indices(self): self.seq_len_v = fx.Index(self.seq_len) self.seq_len_kv_v = fx.Index(self.seq_len_kv) self.stride_q_n_v = fx.Index(self.stride_q_n) + self.stride_o_n_v = fx.Index(self.stride_o_n) self.stride_kv_n_v = fx.Index(self.stride_kv_n) - if traits.HEAD_DIM_V == traits.HEAD_DIM: + if const_expr(traits.PAGED): + self.stride_v_n_v = fx.Index(traits.NUM_HEADS_KV * traits.V_HEAD_DIM) + elif traits.HEAD_DIM_V == traits.HEAD_DIM: self.stride_v_n_v = self.stride_kv_n_v self.stride_o_n_v = self.stride_q_n_v else: @@ -4435,11 +4785,21 @@ def init_lds(self, shared_storage): lds = fx.SharedAllocator().allocate(shared_storage).peek() self.lds = lds self.lds_kv_base_idx = fx.Index(fx.ptrtoint(lds.kv.ptr)) - self.lds_kv_base_ptr = lds.kv.ptr.llvm_ptr + if const_expr(self.traits.PAGED): + self.lds_kv_base_ptr = buffer_ops.create_llvm_ptr(self.lds_kv_base_idx, address_space=3) + else: + self.lds_kv_base_ptr = lds.kv.ptr.llvm_ptr self.lds_vt_base_idx = fx.Index(fx.ptrtoint(lds.vt.ptr)) - self.lds_vt_base_ptr = lds.vt.ptr.llvm_ptr - self.lds_q_base_idx = fx.Index(fx.ptrtoint(lds.q.ptr)) - self.lds_q_base_ptr = lds.q.ptr.llvm_ptr + if const_expr(self.traits.PAGED): + self.lds_vt_base_ptr = buffer_ops.create_llvm_ptr(self.lds_vt_base_idx, address_space=3) + else: + self.lds_vt_base_ptr = lds.vt.ptr.llvm_ptr + if hasattr(lds, "q"): + self.lds_q_base_idx = fx.Index(fx.ptrtoint(lds.q.ptr)) + self.lds_q_base_ptr = lds.q.ptr.llvm_ptr + else: + self.lds_q_base_idx = None + self.lds_q_base_ptr = None def init_thread_mapping(self): _init_dualwave_thread_mapping(self) @@ -4486,8 +4846,6 @@ def init_descriptors(self): traits = self.traits eb = traits.ELEM_BYTES q_nrec_bytes = as_mlir_value(self.q_tok_end * self.stride_q_n_v * eb) - kv_nrec_bytes = as_mlir_value(self.kv_tok_end * self.stride_kv_n_v * eb) - v_nrec_bytes = as_mlir_value(self.kv_tok_end * self.stride_v_n_v * eb) o_nrec_bytes = as_mlir_value(self.q_tok_end * self.stride_o_n_v * traits.OUT_ELEM_BYTES) def _make_buf_div(tensor, nrec_bytes): @@ -4504,11 +4862,50 @@ def _make_buf_div(tensor, nrec_bytes): return fx.logical_divide(bt, fx.make_layout(1, 1)) self.q_div = _make_buf_div(self.Q, q_nrec_bytes) - self.k_div = _make_buf_div(self.K, kv_nrec_bytes) - self.v_div = _make_buf_div(self.V, v_nrec_bytes) self.o_div = fx.logical_divide( fx.rocdl.make_buffer_tensor(self.O, num_records_bytes=o_nrec_bytes), fx.make_layout(1, 1) ) + if const_expr(traits.PAGED): + # K/V are one-byte OCP FP8 pages. Rebase one descriptor per tile + # with 64-bit pointer arithmetic so serving-sized pools beyond the + # 4-GiB boundary retain the exact physical page address. + self.buf_flags_i32 = fx.Int32(buffer_ops._get_buffer_flags()) + self.k_base_iter = fx.get_iter(self.K) + self.v_base_iter = fx.get_iter(self.V) + k_page_elems = fx.Index(traits.PAGE_SIZE * traits.NUM_HEADS_KV * traits.HEAD_DIM) + v_page_elems = fx.Index(traits.PAGE_SIZE * traits.NUM_HEADS_KV * traits.V_HEAD_DIM) + self.k_page_byte_stride = k_page_elems * fx.Index(traits.ELEM_BYTES) + self.v_page_byte_stride = v_page_elems * fx.Index(traits.ELEM_BYTES) + self.k_page_nrec_bytes = fx.Int64(self.k_page_byte_stride) + self.v_page_nrec_bytes = fx.Int64(self.v_page_byte_stride) + self.k_page_layout = fx.make_layout(fx.Int32(k_page_elems), fx.Int32(1)) + self.v_page_layout = fx.make_layout(fx.Int32(v_page_elems), fx.Int32(1)) + self.page_indices_div = fx.logical_divide( + fx.rocdl.make_buffer_tensor(self.BlockTable), fx.make_layout(1, 1) + ) + self.page_i32_atom = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), fx.Int32) + self.page_v1i32 = Vec.make_type(1, fx.Int32) + self.page_base = self.batch_idx * fx.Index(self.block_table_stride) + self.k_div = None + self.v_div = None + else: + kv_nrec_bytes = as_mlir_value(self.kv_tok_end * self.stride_kv_n_v * eb) + v_nrec_bytes = as_mlir_value(self.kv_tok_end * self.stride_v_n_v * eb) + self.k_div = _make_buf_div(self.K, kv_nrec_bytes) + self.v_div = _make_buf_div(self.V, v_nrec_bytes) + self.buf_flags_i32 = None + self.k_base_iter = None + self.v_base_iter = None + self.k_page_byte_stride = None + self.v_page_byte_stride = None + self.k_page_nrec_bytes = None + self.v_page_nrec_bytes = None + self.k_page_layout = None + self.v_page_layout = None + self.page_indices_div = None + self.page_i32_atom = None + self.page_v1i32 = None + self.page_base = None def init_atoms_and_lds_ptrs(self): traits = self.traits @@ -4521,6 +4918,9 @@ def init_atoms_and_lds_ptrs(self): self.o_store_reg_128 = fx.make_rmem_tensor(fx.make_layout(4, 1), fx.Int32) # fp8 global->LDS DMA uses i8 destination typing; K/V LDS reads are byte-addressed. self.lds_ptr_ty = fx.PointerType.get(fx.Int8.ir_type, 2, traits.DMA_BYTES) + if const_expr(traits.PAGED): + self.bf16_mma_atom = fx.make_mma_atom(fx.rocdl.MFMA(32, 32, 16, fx.BFloat16)) + self.v_fp8_load64_atom = fx.make_copy_atom(fx.rocdl.BufferCopy64b(), fx.Int32) def init_descale(self): def _load_scale_scalar(tensor): @@ -4543,6 +4943,7 @@ def init_tile_bounds(self): traits = self.traits kv_tile_size = traits.BLOCK_N num_kv_tiles = (self.seqlen_kv_v + kv_tile_size - 1) // kv_tile_size + self.num_kv_tiles = num_kv_tiles if const_expr(traits.CAUSAL): causal_end_raw_i32 = fx.Int32(self.q_start + traits.BLOCK_M) + self.delta_i32 causal_end_i32 = fx.Int32((causal_end_raw_i32 > fx.Int32(0)).select(causal_end_raw_i32, fx.Int32(0))) @@ -4630,6 +5031,11 @@ def bf16_trunc_pack_v8(self, f32_vals): def buffer_load_128(self, elem_index): return _buffer_load_128(elem_index, self.load_atom_128, self.q_div, self.v4i32_type) + def buffer_load_fp8x16(self, src_div, elem_index): + return fly.copy_atom_call_ssa( + [self.v4i32_type], self.load_atom_128, fx.slice(src_div, (None, fx.Int32(elem_index))) + ) + def buffer_load_lds_128(self, src_div, lds_byte_addr, src_elem, soffset_elems): _buffer_load_lds_128( src_div, lds_byte_addr, src_elem, soffset_elems, _dma_atom=self.dma_atom, _lds_ptr_ty=self.lds_ptr_ty @@ -4642,16 +5048,65 @@ def global_idx_q(self, token_idx, col): return (self.q_tok_base + token_idx) * self.stride_q_n_v + self.q_head_idx * self.traits.HEAD_DIM + col def global_idx_o(self, token_idx, col): - """Element index into O, which is HEAD_DIM_V wide (not HEAD_DIM).""" - return (self.q_tok_base + token_idx) * self.stride_o_n_v + self.q_head_idx * self.traits.HEAD_DIM_V + col + return (self.q_tok_base + token_idx) * self.stride_o_n_v + self.q_head_idx * self.traits.V_HEAD_DIM + col + + def load_page_id(self, tile_start): + """Load this request's physical page for a 64-token logical KV tile.""" + local_page = fx.Index(tile_start) // fx.Index(self.traits.PAGE_SIZE) + # The software pipeline rounds to an even tile count and at least four + # tiles. Redirect those masked padding tiles to entry zero instead of + # reading past the final request's page-index allocation. + page_idx = self.page_base + local_page + safe_idx = fx.Index((local_page < self.num_kv_tiles).select(page_idx, fx.Index(0))) + v = fly.copy_atom_call_ssa( + [self.page_v1i32], + self.page_i32_atom, + fx.slice(self.page_indices_div, (None, fx.Int32(safe_idx))), + ) + return fx.Index(fx.Int32(rocdl.readfirstlane(T.i32, as_mlir_value(fx.Int32(Vec(v, (1,), fx.Int32)[0]))))) + + def make_page_view(self, tensor_iter, page_id, *, is_value=False): + if const_expr(is_value): + page_byte_stride = self.v_page_byte_stride + page_nrec_bytes = self.v_page_nrec_bytes + page_layout = self.v_page_layout + else: + page_byte_stride = self.k_page_byte_stride + page_nrec_bytes = self.k_page_nrec_bytes + page_layout = self.k_page_layout + return _make_page_view( + tensor_iter, + tensor_iter.type, + tensor_iter.alignment, + page_id, + page_byte_stride, + page_nrec_bytes, + page_layout, + fx.Int8.ir_type, + self.buf_flags_i32, + ) def read_i32x8_lds(self, base_ptr, byte_row): + if const_expr(self.traits.PAGED): + words = [] + for w in range_constexpr(8): + p = buffer_ops.get_element_ptr( + base_ptr, + byte_offset=fx.Int32(byte_row + w * 4), + elem_type=T.i8, + ) + words.append(fx.Int32(llvm.LoadOp(T.i32, p, alignment=1).result)) + return Vec.from_elements(words, fx.Int32).ir_value() halves = [] for h in range_constexpr(2): p = buffer_ops.get_element_ptr(base_ptr, byte_offset=fx.Int32(byte_row + h * 16), elem_type=T.i8) halves.append(Vec(llvm.LoadOp(Vec.make_type(4, fx.Int32), p, alignment=16).result)) return halves[0].shuffle(halves[1], [0, 1, 2, 3, 4, 5, 6, 7]).ir_value() + def read_i32x4_lds(self, base_ptr, byte_row): + ptr = buffer_ops.get_element_ptr(base_ptr, byte_offset=fx.Int32(byte_row), elem_type=T.i8) + return llvm.LoadOp(self.v4i32_type, ptr, alignment=16).result + class DualwaveFp8QLoader(DualwaveFp8KernelContext): def __init__(self, ctx): @@ -4674,8 +5129,15 @@ def load_all_wide(self, q_row_in_block): d_base = self.lane_div_32 * 32 packs = [] for ws in range_constexpr(traits.HEAD_DIM // 64): - byte_row = q_row_in_block * fx.Index(traits.HEAD_DIM) + fx.Index(ws * 64) + d_base - packs.append(self.read_i32x8_lds(self.lds_q_base_ptr, fx.Int32(byte_row))) + q_col = ws * 64 + d_base + if const_expr(self.lds_q_base_ptr is None): + g_idx = self.q_gmem_elem_offset + q_row_in_block * self.stride_q_n_v + q_col + lo = Vec(self.buffer_load_128(g_idx), (4,), fx.Int32) + hi = Vec(self.buffer_load_128(g_idx + 16), (4,), fx.Int32) + packs.append(lo.shuffle(hi, [0, 1, 2, 3, 4, 5, 6, 7]).ir_value()) + else: + byte_row = q_row_in_block * fx.Index(traits.HEAD_DIM) + q_col + packs.append(self.read_i32x8_lds(self.lds_q_base_ptr, fx.Int32(byte_row))) return packs @@ -4683,6 +5145,54 @@ class DualwaveFp8GemmHelper(DualwaveFp8KernelContext): def __init__(self, ctx): super().__init__(ctx) + def _mfma_acc_bf16(self, a_v8, b_v8, c_v16): + return fly.mma_atom_call_ssa([self.v16f32_type], self.bf16_mma_atom, a_v8, b_v8, c_v16) + + def _v8bf16_to_f32(self, v8): + f32 = Vec(llvm.FPExtOp(Vec.make_type(8, fx.Float32), as_mlir_value(v8)).result, (8,), fx.Float32) + return [f32[i] for i in range_constexpr(8)] + + def _p_to_fp8_i32x8(self, v_p): + p_lo, p_hi = v_p + f32 = [] + for pk in (p_lo[0], p_lo[1], p_hi[0], p_hi[1]): + f32 += self._v8bf16_to_f32(pk) + return self._pack_fp8_i32x8(f32) + + def _v_to_fp8_i32x8(self, v_v, dc): + f32 = [] + for step in range_constexpr(4): + f32 += self._v8bf16_to_f32(v_v[step][dc]) + return self._pack_fp8_i32x8(f32) + + def _pv_fp8(self, v_p, v_v, v_o): + p_fp8 = self._p_to_fp8_i32x8(v_p) + for dc in range_constexpr(self.traits.D_CHUNKS): + v_op = self._v_concat_i32x8(v_v, dc) + v_o[dc] = self._mfma_acc_fp8_wide(v_op, p_fp8, v_o[dc]) + return v_o + + def _pv_step_fp8(self, step, v_p, v_v, v_o): + if const_expr(step == 0): + self._pv_p_fp8_cache = self._p_to_fp8_i32x8(v_p) + v_op = self._v_concat_i32x8(v_v, step) + v_o[step] = self._mfma_acc_fp8_wide(v_op, self._pv_p_fp8_cache, v_o[step]) + return v_o + + def pv_step_k(self, step, v_p, v_v, v_o): + if const_expr(self.traits.FP8_PV): + return self._pv_step_fp8(step, v_p, v_v, v_o) + # HIPREC PV: P and V are both v8 bf16, accumulated by a bf16 MMA. + v_p_lo, v_p_hi = v_p + v_pk = v_v[step] + if const_expr(step < 2): + p_pk = v_p_lo[step] + else: + p_pk = v_p_hi[step - 2] + for dc in range_constexpr(self.traits.D_CHUNKS): + v_o[dc] = self._mfma_acc_bf16(v_pk[dc], p_pk, v_o[dc]) + return v_o + def _mfma_acc_fp8_wide(self, a_i32x8, b_i32x8, c_v16): # Wide fp8 QK: mfma_scale (32x32x64) with unit E8M0 scales, i32x8 operands. return rocdl.mfma_scale_f32_32x32x64_f8f6f4( @@ -4755,12 +5265,17 @@ def qk(self, v_k, q_wide=None): q_w = q_all_wide[ws] v_s_lo = self._mfma_acc_fp8_wide(k_lo[ws], q_w, v_s_lo) v_s_hi = self._mfma_acc_fp8_wide(k_hi[ws], q_w, v_s_hi) - n_ds = const_expr(traits.HEAD_DIM // 64 * 4) - n_mfma = const_expr(traits.HEAD_DIM // 64 * 2) - rocdl.sched_group_barrier(traits.SCHED_DS_READ_MASK, n_ds // 2, 12) - rocdl.sched_group_barrier(traits.SCHED_MFMA_MASK, 1, 12) - rocdl.sched_group_barrier(traits.SCHED_DS_READ_MASK, n_ds // 2, 12) - rocdl.sched_group_barrier(traits.SCHED_MFMA_MASK, n_mfma - 1, 12) + if const_expr(traits.QREG): + n_ds = const_expr(traits.HEAD_DIM // 64 * 4) + n_mfma = const_expr(traits.HEAD_DIM // 64 * 2) + rocdl.sched_group_barrier(traits.SCHED_DS_READ_MASK, n_ds // 2, 12) + rocdl.sched_group_barrier(traits.SCHED_MFMA_MASK, 1, 12) + rocdl.sched_group_barrier(traits.SCHED_DS_READ_MASK, n_ds // 2, 12) + rocdl.sched_group_barrier(traits.SCHED_MFMA_MASK, n_mfma - 1, 12) + if const_expr(traits.PAGED): + scale_vec = Vec.from_elements([self.c_logit_scale], fx.Float32).broadcast_to(16) + v_s_lo = as_mlir_value(Vec(v_s_lo) * scale_vec) + v_s_hi = as_mlir_value(Vec(v_s_hi) * scale_vec) return (v_s_lo, v_s_hi) def cast_p_fp8_direct(self, v_p): @@ -4782,7 +5297,13 @@ def _pv_fp8_direct(self, p_fp8, v_v, v_o): return v_o def pv(self, v_p, v_v, v_o): - return self._pv_fp8_direct(v_p, v_v, v_o) + if const_expr(self.traits.FP8_PV_DIRECT): + return self._pv_fp8_direct(v_p, v_v, v_o) + if const_expr(self.traits.FP8_PV): + return self._pv_fp8(v_p, v_v, v_o) + for step in range_constexpr(4): + v_o = self.pv_step_k(step, v_p, v_v, v_o) + return v_o class DualwaveFp8KvGmemToLdsLoader(DualwaveFp8KernelContext): @@ -4799,6 +5320,30 @@ def load_k(self, tile_start, buf_id): traits = self.traits eb = traits.ELEM_BYTES k_lds_byte_base = self.lds_kv_base_idx + self.k_buf_base(buf_id) * eb + if const_expr(traits.KV_VECTORIZED): + page_id = self.load_page_id(tile_start) + src_div = self.make_page_view(self.k_base_iter, page_id) + # Each wave owns two 64-lane octet rows at D192. The physical K + # page has only 12 d-groups; the final four groups address beyond + # the page descriptor and the gfx950 buffer load supplies zeros. + for d in range_constexpr(self.NUM_DMA_K): + oct_idx = ( + self.wave_id_uni * (traits.WARP_SIZE * traits.SMEM_D_RPT) + d * traits.WARP_SIZE + self.lane_in_warp + ) + ni = oct_idx % traits.BLOCK_N + dg = oct_idx // traits.BLOCK_N + src_oct = dg * traits.BLOCK_N + _sigma_k_tile_n(ni) + src_elem = ( + self.kv_head_idx * (traits.HEAD_DIM // traits.KV_VEC_SIZE) * traits.BLOCK_N * traits.KV_VEC_SIZE + + src_oct * traits.KV_VEC_SIZE + ) + lds_addr = ( + k_lds_byte_base + (self.wave_id_uni * traits.SMEM_D_RPT + d) * traits.WARP_SIZE * traits.KV_VEC_SIZE + ) + lds_addr = rocdl.readfirstlane(T.i32, as_mlir_value(fx.Int32(lds_addr))) + self.buffer_load_lds_128(src_div, lds_addr, src_elem, 0) + return + rows_per_wave = -(-traits.BLOCK_N // traits.NUM_WAVES) for d in range_constexpr(self.NUM_DMA_K): lanes_per_row = traits.K_BAND_CHUNK[d] // traits.VEC_KV @@ -4832,7 +5377,49 @@ def _run(): _run() def load_v(self, tile_start, buf_id): - self._stage_v_fp8_block_dma(tile_start, buf_id) + if const_expr(self.traits.KV_VECTORIZED): + self._stage_vt_dequant_fp8_vectorized(tile_start, buf_id) + else: + self._stage_v_fp8_block_dma(tile_start, buf_id) + + def zero_v_fp8_lds(self): + traits = self.traits + v_tile_bytes = (traits.BLOCK_N // 8) * (traits.HEAD_DIM // 16) * 128 + total = 2 * v_tile_bytes + aligned_base = ((self.lds_vt_base_idx + fx.Index(127)) // fx.Index(128)) * fx.Index(128) + zero = Vec.from_elements([fx.Int32(0) for _ in range_constexpr(4)], fx.Int32) + per = total // (traits.BLOCK_SIZE) # bytes per thread + for i in range_constexpr(per // 16): + off = aligned_base + self.tid * fx.Index(per) + fx.Index(i * 16) + p = buffer_ops.create_llvm_ptr(off, address_space=3) + llvm.StoreOp(as_mlir_value(zero), p, alignment=16) + + def _stage_v_fp8_block(self, tile_start, buf_id): + traits = self.traits + if const_expr(traits.VDMA): + return self._stage_v_fp8_block_dma(tile_start, buf_id) + v_tile_bytes = (traits.BLOCK_N // 8) * (traits.HEAD_DIM // 16) * 128 + buf_off = buf_id * v_tile_bytes + n = self.wave_id * fx.Index(8) + self.lane // fx.Index(8) + d_block = self.lane % fx.Index(8) + src_elem = ( + self.kv_gmem_elem_offset + n * self.stride_kv_n_v + d_block * fx.Index(16) + tile_start * self.stride_kv_n_v + ) + v16 = fly.copy_atom_call_ssa( + [Vec.make_type(4, fx.Int32)], self.load_atom_128, fx.slice(self.v_div, (None, fx.Int32(src_elem))) + ) + n_i = fx.Int32(n) + w16 = n_i % fx.Int32(16) + c_add = (w16 >= fx.Int32(4)) & (w16 < fx.Int32(8)) + c_sub = (w16 >= fx.Int32(8)) & (w16 < fx.Int32(12)) + dest_n = n_i + c_add.select(fx.Int32(4), fx.Int32(0)) - c_sub.select(fx.Int32(4), fx.Int32(0)) + dest_wave = fx.Index(dest_n // fx.Int32(8)) + dest_m = fx.Index(dest_n % fx.Int32(8)) + block = dest_wave * fx.Index(8) + self.lane % fx.Index(8) + aligned_base = ((self.lds_vt_base_idx + fx.Index(127)) // fx.Index(128)) * fx.Index(128) + byte_off = aligned_base + fx.Index(buf_off) + block * fx.Index(128) + fx.Index(16) * dest_m + lds_ptr = buffer_ops.create_llvm_ptr(byte_off, address_space=3) + llvm.StoreOp(as_mlir_value(Vec(v16)), lds_ptr, alignment=16) def _stage_v_fp8_block_dma(self, tile_start, buf_id): traits = self.traits @@ -4878,6 +5465,94 @@ def _run(): _run() + def _stage_vt_dequant_fp8_vectorized(self, tile_start, buf_id): + """Dequantize SHUFFLE V [H,4,D,16] into the bf16 vectorized LDS tile.""" + traits = self.traits + page_id = self.load_page_id(tile_start) + src_div = self.make_page_view(self.v_base_iter, page_id, is_value=True) + vt_buf = buf_id * traits.VT_BF16_ELEMS + token_groups16 = traits.BLOCK_N // traits.KV_VEC_SIZE + total_v16 = traits.V_HEAD_DIM * token_groups16 + passes = (total_v16 + traits.BLOCK_SIZE - 1) // traits.BLOCK_SIZE + + for pass_id in range_constexpr(passes): + flat = self.tid + fx.Index(pass_id * traits.BLOCK_SIZE) + + def _load_store_active(): + n_group16 = flat % fx.Index(token_groups16) + d_col = flat // fx.Index(token_groups16) + src_elem = ( + self.kv_head_idx * token_groups16 * traits.V_HEAD_DIM * traits.KV_VEC_SIZE + + n_group16 * traits.V_HEAD_DIM * traits.KV_VEC_SIZE + + d_col * traits.KV_VEC_SIZE + ) + v_i32x4 = self.buffer_load_fp8x16(src_div, src_elem) + v_words = Vec(v_i32x4, (4,), fx.Int32) + bf = [] + for w in range_constexpr(4): + word = as_mlir_value(fx.Int32(v_words[w])) + lo2 = Vec(rocdl.cvt_pk_f32_fp8(Vec.make_type(2, fx.Float32), word, False), (2,), fx.Float32) + hi2 = Vec(rocdl.cvt_pk_f32_fp8(Vec.make_type(2, fx.Float32), word, True), (2,), fx.Float32) + for e in (lo2[0], lo2[1], hi2[0], hi2[1]): + bf.append(fx.Float32(e) * self.vd_fp8) + + # Vectorized K applies sigma(n) while staging each 16-token + # group, so MFMA score slot n represents physical key + # sigma(n). Apply the same (involutive) permutation to V; + # otherwise selective attention pairs K[sigma(n)] with V[n] + # even though near-uniform low-amplitude tests appear valid. + bf = [bf[(i & 3) | ((i & 8) >> 1) | ((i & 4) << 1)] for i in range_constexpr(16)] + + # The bf16 vectorized LDS layout groups eight tokens. Split + # the 16-token SHUFFLE vector into two aligned v8 stores. + row = d_col // fx.Index(traits.VEC_BF) + d_in_row = d_col % fx.Index(traits.VEC_BF) + for half in range_constexpr(2): + no = n_group16 * fx.Index(2) + fx.Index(half) + lane_slot = no * fx.Index(traits.VEC_BF) + d_in_row + dst_elem = vt_buf + row * fx.Index(traits.VLS_BF) + lane_slot * fx.Index(traits.VEC_BF) + v8bf = self.bf16_trunc_pack_v8(bf[half * 8 : (half + 1) * 8]) + lds_ptr = buffer_ops.get_element_ptr( + self.lds_vt_base_ptr, + byte_offset=fx.Int32(dst_elem * fx.Index(traits.EB_BF)), + elem_type=T.i8, + ) + llvm.StoreOp(as_mlir_value(v8bf), lds_ptr, alignment=16) + + scf_if_dispatch(flat < fx.Index(total_v16), _load_store_active) + + def _stage_vt_dequant_fp8(self, tile_start, buf_id): + # Dequantize fp8 V into the exact bf16 V staging positions. The two d-iters + # load 8 fp8 at D offsets 64 apart; a contiguous 16B load would gather wrong. + traits = self.traits + vt_buf = buf_id * traits.VT_BF16_ELEMS + n_in_tile = self.n_in_warp * traits.NUM_WAVES + self.wave_id + for d in range_constexpr(traits.SDRPT_BF): + global_d = self.d_bucket * traits.VEC_BF + (d * traits.D128_BF) + src_elem = ( + self.kv_gmem_elem_offset + n_in_tile * self.stride_kv_n_v + global_d + tile_start * self.stride_kv_n_v + ) + v_i32x2 = fly.copy_atom_call_ssa( + [self.v2i32_type], self.v_fp8_load64_atom, fx.slice(self.v_div, (None, fx.Int32(src_elem))) + ) + v_words = Vec(v_i32x2, (2,), fx.Int32) + bf = [] + for w in range_constexpr(2): + word = as_mlir_value(fx.Int32(v_words[w])) + lo2 = Vec(rocdl.cvt_pk_f32_fp8(Vec.make_type(2, fx.Float32), word, False), (2,), fx.Float32) + hi2 = Vec(rocdl.cvt_pk_f32_fp8(Vec.make_type(2, fx.Float32), word, True), (2,), fx.Float32) + for e in (lo2[0], lo2[1], hi2[0], hi2[1]): + bf.append(fx.Float32(e) * self.vd_fp8) + v8bf = self.bf16_trunc_pack_v8(bf) + byte_off = ( + vt_buf + + self.wave_id_uni * traits.VLS_BF + + d * traits.SNRPT_BF * traits.VLS_BF + + self.lane * traits.VEC_BF + ) * traits.EB_BF + lds_ptr = buffer_ops.get_element_ptr(self.lds_vt_base_ptr, byte_offset=byte_off, elem_type=T.i8) + llvm.StoreOp(as_mlir_value(v8bf), lds_ptr, alignment=16) + class DualwaveFp8KvLdsToVgprLoader(DualwaveFp8KernelContext): def __init__(self, ctx): @@ -4892,6 +5567,22 @@ def load_k(self, buf_id): n_lo = self.lane_mod_32 n_hi = self.lane_mod_32 + 32 + if const_expr(traits.KV_VECTORIZED): + + def _read_vec_key(key): + key_sigma = _sigma_k_tile_n(key) + packs = [] + for ws in range_constexpr(traits.HEAD_DIM // 64): + dg0 = ws * 4 + self.lane_div_32 * 2 + off0 = (k_base + (dg0 * traits.BLOCK_N + key_sigma) * traits.KV_VEC_SIZE) * traits.ELEM_BYTES + off1 = (k_base + ((dg0 + 1) * traits.BLOCK_N + key_sigma) * traits.KV_VEC_SIZE) * traits.ELEM_BYTES + lo = Vec(self.read_i32x4_lds(self.lds_kv_base_ptr, off0), (4,), fx.Int32) + hi = Vec(self.read_i32x4_lds(self.lds_kv_base_ptr, off1), (4,), fx.Int32) + packs.append(lo.shuffle(hi, [0, 1, 2, 3, 4, 5, 6, 7]).ir_value()) + return packs + + return (_read_vec_key(n_lo), _read_vec_key(n_hi)) + rows_per_line = traits.NUM_WAVES def _read_strip(key): @@ -4907,7 +5598,52 @@ def _read_strip(key): return (_read_strip(n_lo), _read_strip(n_hi)) def load_v(self, buf_id): - return self._load_v_fp8_block(buf_id) + if const_expr(self.traits.FP8_PV): + return self._load_v_fp8_block(buf_id) + # Read all V packs from the bf16 vt scratch for buffer `buf_id`. + traits = self.traits + if const_expr(traits.KV_VECTORIZED): + v_base = buf_id * traits.VT_BF16_ELEMS + lane_base = ( + (self.lane_mod_32 // traits.SMEM_N_PER_WAVE) * traits.VEC_V_ROW_STRIDE + + self.lane_div_32 * traits.D128_BF + + (self.lane_mod_32 % traits.SMEM_N_PER_WAVE) * traits.VEC_BF + ) + v_base_ptr = buffer_ops.get_element_ptr( + self.lds_vt_base_ptr, + byte_offset=fx.Int32((v_base + lane_base) * traits.EB_BF), + elem_type=T.i8, + ) + packs = [[None] * traits.D_CHUNKS for _ in range(4)] + for dc in range_constexpr(traits.D_CHUNKS): + for k_substep in range_constexpr(4): + const_off = dc * ( + traits.D_CHUNK // traits.SMEM_N_PER_WAVE + ) * traits.VEC_V_ROW_STRIDE + k_substep * (2 * traits.D128_BF) + ptr = buffer_ops.get_element_ptr( + v_base_ptr, + byte_offset=fx.Int32(const_off * traits.EB_BF), + elem_type=T.i8, + ) + packs[k_substep][dc] = llvm.LoadOp(self.v8bf16_type, ptr, alignment=16).result + return packs + + urv = ( + self.lane_div_32 * traits.URV_GRPK_BF + + ((self.lane % 16) // 4) * traits.URV_LANE_HI_BF + + ((self.lane // 16) % 2) * traits.URV_GRP_N_BF + + (self.lane % 4) * traits.URV_LANE_LO_BF + ) + packs = [[None] * traits.D_CHUNKS for _ in range(4)] + for dc in range_constexpr(traits.D_CHUNKS): + dc_off = (dc // 2) * traits.URV_DC_AXIS0_BF + (dc % 2) * traits.URV_DC_AXIS1_BF + for k_substep in range_constexpr(4): + imm_lo = (k_substep * traits.URV_STEPK_BF + dc_off) * traits.EB_BF + byte0 = (urv + buf_id * traits.VT_BF16_ELEMS) * traits.EB_BF + self.lds_vt_base_idx + a = _ds_read_tr16_b64_imm(self.v4bf16_type, fx.Int32(byte0), imm_lo) + b = _ds_read_tr16_b64_imm(self.v4bf16_type, fx.Int32(byte0), imm_lo + traits.URV_I5_BF * traits.EB_BF) + packs[k_substep][dc] = Vec(a).shuffle(Vec(b), [0, 1, 2, 3, 4, 5, 6, 7]).ir_value() + return packs def _load_v_fp8_block(self, buf_id): traits = self.traits @@ -4949,6 +5685,11 @@ def _causal_mask_inplace(self, v_s, tile_idx): s_lo, s_hi = v_s kv_tile_start = tile_idx * traits.BLOCK_N kv_start_i32 = fx.Int32(kv_tile_start) + # The wide fp8 MFMA exposes eight score elements per half-wave group, + # but after vectorized-K sigma their physical key ids are arranged in + # 4-token groups (0..3, 8..11, ...), exactly like the dense mapping. + # Keep this local to the fp8 path; the bf16 vectorized kernel has a + # different operand/output layout. lane_off_i32 = fx.Int32(self.lane_div_32) * fx.Int32(4) # q_row_i32 is set by init_q_row (called after helper construction), so read # it from the live ctx. @@ -5012,8 +5753,7 @@ def _run(v_s_a, v_s_b, tile_a=tile_a, kv_end_pos=kv_end_pos): def _seq_pad_mask_inplace(self, v_s_lists, tile_idx): traits = self.traits s_lo, s_hi = v_s_lists - kv_tile_start = tile_idx * traits.BLOCK_N - col_base = fx.Int32(kv_tile_start) + fx.Int32(self.lane_div_32) * fx.Int32(4) + col_base = fx.Int32(tile_idx * traits.BLOCK_N) + fx.Int32(self.lane_div_32) * fx.Int32(4) for r in range_constexpr(16): thr = (r // 4) * 8 + (r % 4) col_lo = col_base + fx.Int32(thr) @@ -5044,12 +5784,16 @@ def max2(self, a, b): return fx.maxnumf(a, b) def floor_masked_max(self, row_max): + if const_expr(self.traits.PAGED): + return fx.maxnumf(row_max, self.c_neg_floor) return fx.maxnumf(row_max, self.c_neg_floor) # log2 of e4m3's largest finite value, 448. _P_HEADROOM_LOG2 = 8.807354922057604 def sub_m(self, v_s, row_max): + if const_expr(self.traits.PAGED): + return _sub_score_pair(v_s, row_max, self.fm_fast) # P is cast to e4m3, whose smallest subnormal is 2**-9, so a softmax # over thousands of keys loses its tail to flush-to-zero -- while l_row, # summed before the cast, still counts it. Scaling P up first uses the @@ -5072,6 +5816,8 @@ def tile_sum(self, v_p): return _score_pair_sum(v_p, self.c_zero_f, self.fm_fast) def reduce_sum(self, l_row, v_p): + if const_expr(self.traits.PAGED): + return l_row + self.tile_sum(v_p) return l_row + self.tile_sum(v_p) def cast_p(self, v_p): @@ -5105,12 +5851,16 @@ def safe_l_inv(self, l_row): return _safe_l_inv(l_row, self.c_zero_f) def rescale_from_tile_max(self, m_row, m_tile_max): + if const_expr(self.traits.PAGED): + return _rescale_from_tile_max(m_row, m_tile_max, self.fm_fast) row_max = fx.maxnumf(m_row, m_tile_max) diff_scaled = (m_row - row_max) * self.c_logit_scale rescale = rocdl.exp2(T.f32, as_mlir_value(diff_scaled)) return row_max, rescale def apply_l_rescale(self, l_row, rescale): + if const_expr(self.traits.PAGED): + return l_row * rescale return l_row * rescale def rescale_o(self, v_o, m_row, l_row, m_tile_max, v_p): @@ -5137,6 +5887,9 @@ def _lazy_correction(self, v_o, m_row, m_tile_max): return m_new, corr, [as_mlir_value(scaled_accs[i]) for i in range(self.traits.D_CHUNKS)] def lazy_rescale_o(self, v_o, m_row, l_row, m_tile_max, v_p): + if const_expr(self.traits.PAGED): + return self._lazy_rescale_o_paged(v_o, m_row, l_row, m_tile_max, v_p) + @flyc.jit def _run(v_o, m_row, l_row, m_tile_max, v_p): m_diff = m_tile_max - m_row @@ -5162,6 +5915,44 @@ def _run(v_o, m_row, l_row, m_tile_max, v_p): return _run(v_o, m_row, l_row, m_tile_max, v_p) + def _lazy_rescale_o_paged(self, v_o, m_row, l_row, m_tile_max, v_p): + traits = self.traits + + @flyc.jit + def _run(v_o, m_row, l_row, m_tile_max, v_p): + m_diff = m_tile_max - m_row + below = fx.Float32(m_diff) <= self.c_rescale_thr_f + ballot = rocdl.ballot(T.i64, as_mlir_value(below)) + all_below = arith.cmpi(arith.CmpIPredicate.eq, as_mlir_value(ballot), _read_exec_i64()) + all_below = llvm.intr_expect(all_below, arith.constant(1, type=ir.IntegerType.get_signless(1))) + + state = [as_mlir_value(v_o[dc]) for dc in range(traits.D_CHUNKS)] + state += [self.v_p_to_vec32(v_p), as_mlir_value(l_row), as_mlir_value(m_row)] + names = tuple(f"_fp8_paged_lr{i}" for i in range(traits.D_CHUNKS + 3)) + + def _rescale(_names, *_state): + m_new, corr, scaled_accs = self._lazy_correction(v_o, m_row, m_tile_max) + result = list(scaled_accs) + result.append(self.v_p_to_vec32(self.scale_v_p(v_p, corr))) + result.append(as_mlir_value(l_row * corr)) + result.append(self.anchor_scalar_f32(m_new)) + return result + + result = scf_if_dispatch( + all_below, + lambda *_args: None, + _rescale, + state_names=names, + state_values=state, + ) + o_out = list(result[: traits.D_CHUNKS]) + vp_out = result[traits.D_CHUNKS] + l_out = result[traits.D_CHUNKS + 1] + m_out = result[traits.D_CHUNKS + 2] + return (o_out, m_out, l_out, self.v_vec32_to_p(vp_out)) + + return _run(v_o, m_row, l_row, m_tile_max, v_p) + def lazy_correct_o(self, v_o, m_row, l_row, m_tile_max): @flyc.jit def _run(v_o, m_row, l_row, m_tile_max): @@ -5227,6 +6018,11 @@ def store_final_o(self, v_o, q_row): o_global = self.global_idx_o(q_row, d_col) self.buffer_store_128(o_pack, o_global) + @flyc.jit + def store_final_o_if_valid(self, v_o, q_row): + if q_row < self.seqlen_q_v: + self.store_final_o(v_o, q_row) + def store_splitk_partial_o(self, v_o, m_row, l_row, q_row): m_row = fx.Float32(m_row) * self.c_logit_scale split_z = self.batch_idx * self.traits.NUM_KV_SPLITS + self.split_idx diff --git a/tests/kernels/test_flash_attn_fwd.py b/tests/kernels/test_flash_attn_fwd.py index 818ed3f1e8..5e14499da4 100644 --- a/tests/kernels/test_flash_attn_fwd.py +++ b/tests/kernels/test_flash_attn_fwd.py @@ -4575,6 +4575,135 @@ def test_return_lse_rejects_fp8(): ) +@_requires_gfx950 +@pytest.mark.parametrize("value_head_dim", [128, 192]) +@pytest.mark.parametrize("use_non_default_stream", [False, True]) +def test_paged_fp8_d192_asymmetric_value_matches_torch(value_head_dim, use_non_default_stream): + """Packed causal FP8 page-64 attention supports native V/output width.""" + torch.manual_seed(17) + query_lengths = [64, 32] + kv_lengths = [128, 96] + q_indptr = torch.tensor([0, 64, 96], device="cuda", dtype=torch.int32) + kv_indptr = torch.tensor([0, 128, 224], device="cuda", dtype=torch.int32) + block_table = torch.tensor([[2, 0], [3, 1]], device="cuda", dtype=torch.int32) + seqlen_k = torch.tensor(kv_lengths, device="cuda", dtype=torch.int32) + num_pages = 4 + + query, query_descale = quantize_per_tensor_fp8(torch.randn(sum(query_lengths), 16, 192, device="cuda") * 0.2) + key, key_descale = quantize_per_tensor_fp8(torch.randn(num_pages, 1, 12, 64, 16, device="cuda") * 0.2) + value, value_descale = quantize_per_tensor_fp8( + torch.randn(num_pages, 1, 4, value_head_dim, 16, device="cuda") * 0.2 + ) + + stream = torch.cuda.Stream() if use_non_default_stream else None + if stream is not None: + torch.cuda.synchronize() + actual = flydsl_flash_attn_func( + query, + key, + value, + causal=True, + num_kv_heads=1, + cu_seqlens_q=q_indptr, + cu_seqlens_kv=kv_indptr, + max_seqlen_q=max(query_lengths), + max_seqlen_kv=max(kv_lengths), + cross_seqlen=True, + block_table=block_table, + seqlen_k=seqlen_k, + kv_cache_layout="vectorized", + q_descale=query_descale, + k_descale=key_descale, + v_descale=value_descale, + stream=stream, + ) + if stream is not None: + stream.synchronize() + else: + torch.cuda.synchronize() + + expected = [] + for batch_idx, (query_length, kv_length) in enumerate(zip(query_lengths, kv_lengths)): + query_batch = query[q_indptr[batch_idx] : q_indptr[batch_idx + 1]].float() * query_descale + physical_pages = block_table[batch_idx].long() + key_batch = ( + key[physical_pages].permute(0, 3, 1, 2, 4).reshape(-1, 1, 192)[:kv_length].float() * key_descale + ).expand(-1, 16, -1) + value_batch = ( + value[physical_pages].permute(0, 2, 4, 1, 3).reshape(-1, 1, value_head_dim)[:kv_length].float() + * value_descale + ).expand(-1, 16, -1) + logits = torch.einsum("qhd,khd->hqk", query_batch, key_batch) / math.sqrt(192) + query_positions = torch.arange(query_length, device="cuda") + key_positions = torch.arange(kv_length, device="cuda") + causal_mask = key_positions[None, :] <= (kv_length - query_length + query_positions)[:, None] + logits.masked_fill_(~causal_mask[None, :, :], float("-inf")) + expected.append(torch.einsum("hqk,khd->qhd", torch.softmax(logits, dim=-1), value_batch).to(torch.bfloat16)) + + expected = torch.cat(expected) + assert actual.shape == (sum(query_lengths), 16, value_head_dim) + assert bool(torch.isfinite(actual).all().item()) + torch.testing.assert_close(actual, expected, rtol=2.0e-2, atol=2.0e-2) + + +@_requires_gfx950 +@pytest.mark.large_shape +def test_paged_fp8_d192_cache_offsets_above_4gib(): + """Physical K and V page rebasing remains 64-bit beyond 4 GiB.""" + head_dim = 192 + value_head_dim = 128 + page_size = 64 + key_page_bytes = page_size * head_dim + value_page_bytes = page_size * value_head_dim + high_page = math.ceil(2**32 / min(key_page_bytes, value_page_bytes)) + num_pages = high_page + 1 + + query, query_descale = quantize_per_tensor_fp8(torch.randn(1, 16, head_dim, device="cuda") * 0.2) + key = torch.empty(num_pages, 1, head_dim // 16, page_size, 16, dtype=FP8_DTYPE, device="cuda") + value = torch.empty( + num_pages, + 1, + page_size // 16, + value_head_dim, + 16, + dtype=FP8_DTYPE, + device="cuda", + ) + key_page, key_descale = quantize_per_tensor_fp8(torch.randn_like(key[high_page], dtype=torch.float32) * 0.2) + value_page, value_descale = quantize_per_tensor_fp8(torch.randn_like(value[high_page], dtype=torch.float32) * 0.2) + key[high_page].copy_(key_page) + value[high_page].copy_(value_page) + + indptr = torch.tensor([0, 1], device="cuda", dtype=torch.int32) + block_table = torch.tensor([[high_page]], device="cuda", dtype=torch.int32) + seqlen_k = torch.ones(1, device="cuda", dtype=torch.int32) + actual = flydsl_flash_attn_func( + query, + key, + value, + causal=True, + num_kv_heads=1, + cu_seqlens_q=indptr, + cu_seqlens_kv=indptr, + max_seqlen_q=1, + max_seqlen_kv=1, + cross_seqlen=True, + block_table=block_table, + seqlen_k=seqlen_k, + kv_cache_layout="vectorized", + q_descale=query_descale, + k_descale=key_descale, + v_descale=value_descale, + ) + torch.cuda.synchronize() + + expected = (value_page[0, 0, :, 0].float() * value_descale).expand(16, -1).unsqueeze(0).to(torch.bfloat16) + assert high_page * key_page_bytes >= 2**32 + assert high_page * value_page_bytes >= 2**32 + assert bool(torch.isfinite(actual).all().item()) + torch.testing.assert_close(actual, expected, rtol=2.0e-2, atol=2.0e-2) + + @_requires_gfx950 @pytest.mark.parametrize("H", [8, 16, 32, 64]) def test_xcd_swizzle_is_bit_identical(H): From ef39d59b0ea326648fa80bcfdefd226385182f48 Mon Sep 17 00:00:00 2001 From: Xiake Sun Date: Wed, 26 Aug 2026 15:06:53 +0000 Subject: [PATCH 02/59] test(fa): guard large offset test by available VRAM - Calculate the paged K/V cache footprint and require 2 GiB of additional headroom before running the >4 GiB addressing regression. - Skip with a clear diagnostic when the selected GPU has insufficient free memory. Signed-off-by: Xiake Sun --- tests/kernels/test_flash_attn_fwd.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/tests/kernels/test_flash_attn_fwd.py b/tests/kernels/test_flash_attn_fwd.py index 5e14499da4..be2b12b811 100644 --- a/tests/kernels/test_flash_attn_fwd.py +++ b/tests/kernels/test_flash_attn_fwd.py @@ -4657,6 +4657,15 @@ def test_paged_fp8_d192_cache_offsets_above_4gib(): value_page_bytes = page_size * value_head_dim high_page = math.ceil(2**32 / min(key_page_bytes, value_page_bytes)) num_pages = high_page + 1 + cache_bytes = num_pages * (key_page_bytes + value_page_bytes) + required_free_bytes = cache_bytes + 2 * 2**30 + torch.cuda.empty_cache() + free_bytes, _ = torch.cuda.mem_get_info() + if free_bytes < required_free_bytes: + pytest.skip( + f"requires {required_free_bytes / 2**30:.1f} GiB free VRAM for " + f">4 GiB K/V offsets, got {free_bytes / 2**30:.1f} GiB" + ) query, query_descale = quantize_per_tensor_fp8(torch.randn(1, 16, head_dim, device="cuda") * 0.2) key = torch.empty(num_pages, 1, head_dim // 16, page_size, 16, dtype=FP8_DTYPE, device="cuda") From 490bb168b8c7696b94fb5be8a6a63d45a90bebaf Mon Sep 17 00:00:00 2001 From: Xiake Sun Date: Mon, 31 Aug 2026 11:35:31 +0000 Subject: [PATCH 03/59] fix(fmha): order paged copies on the launch stream - Run wrapper-owned casts, contiguous copies, output allocation, and workspace allocation on the selected stream. - Reuse validated caller output directly without creating a throwaway contiguous copy. - Add delayed side-stream coverage for non-contiguous Q/K/V and block-table conversion. Signed-off-by: Xiake Sun --- kernels/attention/flash_attn_interface.py | 93 +++++++++++++---------- tests/kernels/test_flash_attn_fwd.py | 64 +++++++++++++--- 2 files changed, 108 insertions(+), 49 deletions(-) diff --git a/kernels/attention/flash_attn_interface.py b/kernels/attention/flash_attn_interface.py index 3cc95c6f8d..e10c89be1f 100644 --- a/kernels/attention/flash_attn_interface.py +++ b/kernels/attention/flash_attn_interface.py @@ -707,10 +707,20 @@ def _flydsl_flash_attn_paged( f"got shape={tuple(seqlen_k.shape)} dtype={seqlen_k.dtype} device={seqlen_k.device}" ) block_table_stride = int(block_table.shape[1]) - # Flatten so the kernel's flat row-major index addresses block_table correctly. - block_table_i32 = ( - (block_table if block_table.dtype == torch.int32 else block_table.to(torch.int32)).contiguous().reshape(-1) - ) + expected_out_shape = (*q.shape[:-1], value_head_dim) + + if out is not None: + if out.shape != expected_out_shape or not out.is_contiguous(): + raise ValueError( + f"flydsl_flash_attn_func: paged output must be contiguous with shape {expected_out_shape}, " + f"got shape={tuple(out.shape)} strides={out.stride()}" + ) + if paged_fp8 and out.dtype != torch.bfloat16: + raise ValueError(f"flydsl_flash_attn_func: paged FP8 output must be bf16, got {out.dtype}") + if not paged_fp8 and out.dtype != q.dtype: + raise ValueError( + f"flydsl_flash_attn_func: paged output dtype must match q dtype {q.dtype}, got {out.dtype}" + ) with torch.cuda.device(q.device.index): launch_stream = torch.cuda.current_stream(q.device) if stream is None else stream @@ -770,43 +780,48 @@ def _flydsl_flash_attn_paged( kv_cache_layout=kv_cache_layout, has_bias=bias is not None, ) - expected_out_shape = (*q.shape[:-1], value_head_dim) - if out is None: - out_dtype = torch.bfloat16 if paged_fp8 else q.dtype - out = torch.empty(expected_out_shape, dtype=out_dtype, device=q.device) - elif out.shape != expected_out_shape or not out.is_contiguous(): - raise ValueError( - f"flydsl_flash_attn_func: paged output must be contiguous with shape {expected_out_shape}, " - f"got shape={tuple(out.shape)} strides={out.stride()}" + with torch.cuda.stream(launch_stream): + # Keep wrapper-owned casts and copies ordered with an explicit + # non-current launch stream. + block_table_i32 = ( + (block_table if block_table.dtype == torch.int32 else block_table.to(torch.int32)) + .contiguous() + .reshape(-1) ) - elif paged_fp8 and out.dtype != torch.bfloat16: - raise ValueError(f"flydsl_flash_attn_func: paged FP8 output must be bf16, got {out.dtype}") - elif not paged_fp8 and out.dtype != q.dtype: - raise ValueError( - f"flydsl_flash_attn_func: paged output dtype must match q dtype {q.dtype}, got {out.dtype}" + if out is None: + out_dtype = torch.bfloat16 if paged_fp8 else q.dtype + out = torch.empty(expected_out_shape, dtype=out_dtype, device=q.device) + # Keep serving-sized physical K/V caches rank-5 because flattening + # their dynamic memref shape can exceed signed int32. The FP8 + # schedule consumes Q/O as flat token-major buffers, matching its + # explicit runtime strides. + q_flat = q.contiguous().view(-1) if paged_fp8 else q.contiguous() + k_flat = k.contiguous() + v_flat = v.contiguous() + o_flat = out.view(-1) if paged_fp8 else out + kwargs = dict( + block_table=block_table_i32, + block_table_stride=block_table_stride, + stream=launch_stream, ) - # Keep serving-sized physical K/V caches rank-5 because flattening their - # dynamic memref shape can exceed signed int32. The FP8 schedule consumes - # Q/O as flat token-major buffers, matching its explicit runtime strides. - q_flat = q.contiguous().view(-1) if paged_fp8 else q.contiguous() - k_flat = k.contiguous() - v_flat = v.contiguous() - o_flat = out.contiguous().view(-1) if paged_fp8 else out.contiguous() - kwargs = dict(block_table=block_table_i32, block_table_stride=block_table_stride, stream=launch_stream) - if paged_fp8: - kwargs.update(q_descale=q_descale, k_descale=k_descale, v_descale=v_descale) - if bias is not None: - kwargs["bias"] = bias - if varlen: - kwargs["cu_seqlens_q"] = cu_seqlens_q - kwargs["cu_seqlens_kv"] = cu_seqlens_kv - if cross: - kwargs["seq_len_kv"] = skv - if splitk: - ws_elems = dualwave_splitk_workspace_elems(B, H, Sq, int(num_kv_splits), head_dim=D) - _ws = torch.empty(ws_elems, dtype=torch.float32, device=q.device) - kwargs["workspace"] = _ws - exe(q_flat, k_flat, v_flat, o_flat, B, Sq, **kwargs) + if paged_fp8: + kwargs.update( + q_descale=q_descale, + k_descale=k_descale, + v_descale=v_descale, + ) + if bias is not None: + kwargs["bias"] = bias + if varlen: + kwargs["cu_seqlens_q"] = cu_seqlens_q + kwargs["cu_seqlens_kv"] = cu_seqlens_kv + if cross: + kwargs["seq_len_kv"] = skv + if splitk: + ws_elems = dualwave_splitk_workspace_elems(B, H, Sq, int(num_kv_splits), head_dim=D) + _ws = torch.empty(ws_elems, dtype=torch.float32, device=q.device) + kwargs["workspace"] = _ws + exe(q_flat, k_flat, v_flat, o_flat, B, Sq, **kwargs) return out diff --git a/tests/kernels/test_flash_attn_fwd.py b/tests/kernels/test_flash_attn_fwd.py index be2b12b811..79cba03a1a 100644 --- a/tests/kernels/test_flash_attn_fwd.py +++ b/tests/kernels/test_flash_attn_fwd.py @@ -4577,8 +4577,12 @@ def test_return_lse_rejects_fp8(): @_requires_gfx950 @pytest.mark.parametrize("value_head_dim", [128, 192]) -@pytest.mark.parametrize("use_non_default_stream", [False, True]) -def test_paged_fp8_d192_asymmetric_value_matches_torch(value_head_dim, use_non_default_stream): +@pytest.mark.parametrize( + ("use_non_default_stream", "force_internal_copies"), + [(False, False), (True, False), (True, True)], + ids=["default", "side", "side-copies"], +) +def test_paged_fp8_d192_asymmetric_value_matches_torch(value_head_dim, use_non_default_stream, force_internal_copies): """Packed causal FP8 page-64 attention supports native V/output width.""" torch.manual_seed(17) query_lengths = [64, 32] @@ -4595,13 +4599,24 @@ def test_paged_fp8_d192_asymmetric_value_matches_torch(value_head_dim, use_non_d torch.randn(num_pages, 1, 4, value_head_dim, 16, device="cuda") * 0.2 ) - stream = torch.cuda.Stream() if use_non_default_stream else None - if stream is not None: - torch.cuda.synchronize() - actual = flydsl_flash_attn_func( - query, - key, - value, + if force_internal_copies: + + def noncontiguous_copy(tensor): + storage = torch.empty( + (*tensor.shape[:-1], tensor.shape[-1] * 2), + dtype=tensor.dtype, + device=tensor.device, + ) + view = storage[..., ::2] + view.copy_(tensor) + return view + + query = noncontiguous_copy(query) + key = noncontiguous_copy(key) + value = noncontiguous_copy(value) + block_table = block_table.to(torch.int64) + + call_kwargs = dict( causal=True, num_kv_heads=1, cu_seqlens_q=q_indptr, @@ -4609,18 +4624,47 @@ def test_paged_fp8_d192_asymmetric_value_matches_torch(value_head_dim, use_non_d max_seqlen_q=max(query_lengths), max_seqlen_kv=max(kv_lengths), cross_seqlen=True, - block_table=block_table, seqlen_k=seqlen_k, kv_cache_layout="vectorized", q_descale=query_descale, k_descale=key_descale, v_descale=value_descale, + ) + copy_reference = None + if force_internal_copies: + copy_reference = flydsl_flash_attn_func( + query.contiguous(), + key.contiguous(), + value.contiguous(), + block_table=block_table.to(torch.int32).contiguous(), + **call_kwargs, + ) + torch.cuda.synchronize() + + stream = torch.cuda.Stream() if use_non_default_stream else None + if stream is not None: + torch.cuda.synchronize() + default_done = None + if force_internal_copies: + torch.cuda._sleep(3_000_000_000) + default_done = torch.cuda.Event() + default_done.record() + actual = flydsl_flash_attn_func( + query, + key, + value, + block_table=block_table, stream=stream, + **call_kwargs, ) if stream is not None: stream.synchronize() else: torch.cuda.synchronize() + if default_done is not None: + assert not default_done.query(), "side stream unexpectedly waited for the blocked default stream" + torch.cuda.synchronize() + torch.testing.assert_close(actual, copy_reference, rtol=0, atol=0) expected = [] for batch_idx, (query_length, kv_length) in enumerate(zip(query_lengths, kv_lengths)): From ff770715f61568f0f9a43ba864fb44bab1c89de9 Mon Sep 17 00:00:00 2001 From: Xiake Sun Date: Mon, 31 Aug 2026 11:36:49 +0000 Subject: [PATCH 04/59] fix(fmha): guard paged FP8 flattened dimensions - Reject flattened Q or output dimensions at the signed-int32 C ABI boundary. - Raise a descriptive error before kernel build or launch. - Add lowered-limit V128 and V192 regression coverage. Signed-off-by: Xiake Sun --- kernels/attention/flash_attn_interface.py | 8 +++++ tests/kernels/test_flash_attn_fwd.py | 39 +++++++++++++++++++++++ 2 files changed, 47 insertions(+) diff --git a/kernels/attention/flash_attn_interface.py b/kernels/attention/flash_attn_interface.py index e10c89be1f..aa20b10ae5 100644 --- a/kernels/attention/flash_attn_interface.py +++ b/kernels/attention/flash_attn_interface.py @@ -708,6 +708,14 @@ def _flydsl_flash_attn_paged( ) block_table_stride = int(block_table.shape[1]) expected_out_shape = (*q.shape[:-1], value_head_dim) + q_flat_elems = q.numel() + out_flat_elems = q_flat_elems // D * value_head_dim + if paged_fp8 and max(q_flat_elems, out_flat_elems) >= _FP8_MAX_FLAT_ELEMS: + raise NotImplementedError( + "flydsl_flash_attn_func: paged FP8 flattens Q/O and packs the dynamic " + f"dimension as int32, so each must contain fewer than {_FP8_MAX_FLAT_ELEMS} " + f"elements; got q={q_flat_elems}, out={out_flat_elems}. Shorten the packed query." + ) if out is not None: if out.shape != expected_out_shape or not out.is_contiguous(): diff --git a/tests/kernels/test_flash_attn_fwd.py b/tests/kernels/test_flash_attn_fwd.py index 79cba03a1a..b83fba91a6 100644 --- a/tests/kernels/test_flash_attn_fwd.py +++ b/tests/kernels/test_flash_attn_fwd.py @@ -4690,6 +4690,45 @@ def noncontiguous_copy(tensor): torch.testing.assert_close(actual, expected, rtol=2.0e-2, atol=2.0e-2) +@_requires_gfx950 +@pytest.mark.parametrize("value_head_dim", [128, 192]) +def test_paged_fp8_d192_rejects_flattened_int32_overflow(monkeypatch, value_head_dim): + """Reject flattened Q/O dimensions before entering the signed-int32 C ABI.""" + query = torch.zeros((1, 16, 192), device="cuda", dtype=FP8_DTYPE) + key = torch.zeros((1, 1, 12, 64, 16), device="cuda", dtype=FP8_DTYPE) + value = torch.zeros((1, 1, 4, value_head_dim, 16), device="cuda", dtype=FP8_DTYPE) + indptr = torch.tensor([0, 1], device="cuda", dtype=torch.int32) + block_table = torch.zeros((1, 1), device="cuda", dtype=torch.int32) + seqlen_k = torch.ones((1,), device="cuda", dtype=torch.int32) + scale = torch.ones((1,), device="cuda", dtype=torch.float32) + expected_out_elems = query.numel() // query.shape[-1] * value_head_dim + monkeypatch.setattr( + flash_attn_interface, + "_FP8_MAX_FLAT_ELEMS", + max(query.numel(), expected_out_elems), + ) + + with pytest.raises(NotImplementedError, match="paged FP8.*int32"): + flydsl_flash_attn_func( + query, + key, + value, + causal=True, + num_kv_heads=1, + cu_seqlens_q=indptr, + cu_seqlens_kv=indptr, + max_seqlen_q=1, + max_seqlen_kv=1, + cross_seqlen=True, + block_table=block_table, + seqlen_k=seqlen_k, + kv_cache_layout="vectorized", + q_descale=scale, + k_descale=scale, + v_descale=scale, + ) + + @_requires_gfx950 @pytest.mark.large_shape def test_paged_fp8_d192_cache_offsets_above_4gib(): From e40644e5f42f38c452dd14a02819d792f91fc1e0 Mon Sep 17 00:00:00 2001 From: Xiake Sun Date: Mon, 31 Aug 2026 12:19:27 +0000 Subject: [PATCH 05/59] refactor(fmha): clean up paged FP8 address helpers - Remove redundant paged conditionals. - Reuse shared K-address and sigma-permutation helpers. - Correct the D192 padding comments without changing the schedule. Signed-off-by: Xiake Sun --- kernels/attention/flash_attn_utils.py | 35 +++++++++------------------ 1 file changed, 12 insertions(+), 23 deletions(-) diff --git a/kernels/attention/flash_attn_utils.py b/kernels/attention/flash_attn_utils.py index 2d4fe8a24b..5ea9fe0eea 100644 --- a/kernels/attention/flash_attn_utils.py +++ b/kernels/attention/flash_attn_utils.py @@ -2247,10 +2247,9 @@ def _make_paged_dualwave_swp_fp8_traits( smem_linear_wave = warp_size * 16 // elem_bytes smem_n_per_wave = smem_linear_wave // d_128b_size smem_n_rpt = block_n // smem_n_per_wave - # D192 needs a second 128-byte K repeat. In vectorized page-64 mode the - # final four 16-byte lane buckets naturally land beyond the page descriptor - # and read zero, giving the QK MFMA a correctly padded D256 LDS tile without - # ever crossing into another physical page. + # D192 needs a second 128-byte K repeat for the D256-shaped LDS tile. The + # QK reader consumes only the 12 physical D192 groups; the remaining four + # DMA groups are unused padding. smem_d_rpt = (head_dim + d_128b_size - 1) // d_128b_size smem_k_pad = 16 // elem_bytes smem_v_pad = 64 // elem_bytes @@ -5323,19 +5322,15 @@ def load_k(self, tile_start, buf_id): if const_expr(traits.KV_VECTORIZED): page_id = self.load_page_id(tile_start) src_div = self.make_page_view(self.k_base_iter, page_id) - # Each wave owns two 64-lane octet rows at D192. The physical K - # page has only 12 d-groups; the final four groups address beyond - # the page descriptor and the gfx950 buffer load supplies zeros. + # Each wave owns two 64-lane octet rows for the D256-shaped LDS + # tile. The QK reader ignores the four groups beyond physical D192. for d in range_constexpr(self.NUM_DMA_K): - oct_idx = ( - self.wave_id_uni * (traits.WARP_SIZE * traits.SMEM_D_RPT) + d * traits.WARP_SIZE + self.lane_in_warp - ) - ni = oct_idx % traits.BLOCK_N - dg = oct_idx // traits.BLOCK_N - src_oct = dg * traits.BLOCK_N + _sigma_k_tile_n(ni) - src_elem = ( - self.kv_head_idx * (traits.HEAD_DIM // traits.KV_VEC_SIZE) * traits.BLOCK_N * traits.KV_VEC_SIZE - + src_oct * traits.KV_VEC_SIZE + src_elem = _vec_k_src_elem( + traits, + d, + self.wave_id_uni, + self.lane_in_warp, + self.kv_head_idx, ) lds_addr = ( k_lds_byte_base + (self.wave_id_uni * traits.SMEM_D_RPT + d) * traits.WARP_SIZE * traits.KV_VEC_SIZE @@ -5501,7 +5496,7 @@ def _load_store_active(): # sigma(n). Apply the same (involutive) permutation to V; # otherwise selective attention pairs K[sigma(n)] with V[n] # even though near-uniform low-amplitude tests appear valid. - bf = [bf[(i & 3) | ((i & 8) >> 1) | ((i & 4) << 1)] for i in range_constexpr(16)] + bf = [bf[_sigma_k_tile_n(i)] for i in range_constexpr(16)] # The bf16 vectorized LDS layout groups eight tokens. Split # the 16-token SHUFFLE vector into two aligned v8 stores. @@ -5784,8 +5779,6 @@ def max2(self, a, b): return fx.maxnumf(a, b) def floor_masked_max(self, row_max): - if const_expr(self.traits.PAGED): - return fx.maxnumf(row_max, self.c_neg_floor) return fx.maxnumf(row_max, self.c_neg_floor) # log2 of e4m3's largest finite value, 448. @@ -5816,8 +5809,6 @@ def tile_sum(self, v_p): return _score_pair_sum(v_p, self.c_zero_f, self.fm_fast) def reduce_sum(self, l_row, v_p): - if const_expr(self.traits.PAGED): - return l_row + self.tile_sum(v_p) return l_row + self.tile_sum(v_p) def cast_p(self, v_p): @@ -5859,8 +5850,6 @@ def rescale_from_tile_max(self, m_row, m_tile_max): return row_max, rescale def apply_l_rescale(self, l_row, rescale): - if const_expr(self.traits.PAGED): - return l_row * rescale return l_row * rescale def rescale_o(self, v_o, m_row, l_row, m_tile_max, v_p): From a3c0d34be5ddf678ade0269e58c863a5a95cb07a Mon Sep 17 00:00:00 2001 From: Xiake Sun Date: Mon, 31 Aug 2026 12:21:21 +0000 Subject: [PATCH 06/59] fix(fmha): preserve paged BF16 seqlen compatibility - Apply the strict int32 [B] seqlen_k contract only to paged FP8. - Preserve existing BF16/F16 acceptance of host-only length metadata. - Retain common block-table and output-buffer safety validation. Signed-off-by: Xiake Sun --- kernels/attention/flash_attn_interface.py | 4 ++-- tests/kernels/test_flash_attn_fwd.py | 26 +++++++++++++++++++++++ 2 files changed, 28 insertions(+), 2 deletions(-) diff --git a/kernels/attention/flash_attn_interface.py b/kernels/attention/flash_attn_interface.py index aa20b10ae5..ab7483d5d0 100644 --- a/kernels/attention/flash_attn_interface.py +++ b/kernels/attention/flash_attn_interface.py @@ -701,9 +701,9 @@ def _flydsl_flash_attn_paged( f"flydsl_flash_attn_func: block_table must be 2D on {q.device}, " f"got shape={tuple(block_table.shape)} device={block_table.device}" ) - if seqlen_k.dtype != torch.int32 or seqlen_k.device != q.device or seqlen_k.numel() != B: + if paged_fp8 and (seqlen_k.dtype != torch.int32 or seqlen_k.device != q.device or seqlen_k.numel() != B): raise ValueError( - f"flydsl_flash_attn_func: seqlen_k must be int32 [{B}] on {q.device}, " + f"flydsl_flash_attn_func: paged FP8 seqlen_k must be int32 [{B}] on {q.device}, " f"got shape={tuple(seqlen_k.shape)} dtype={seqlen_k.dtype} device={seqlen_k.device}" ) block_table_stride = int(block_table.shape[1]) diff --git a/tests/kernels/test_flash_attn_fwd.py b/tests/kernels/test_flash_attn_fwd.py index b83fba91a6..6cbc93b530 100644 --- a/tests/kernels/test_flash_attn_fwd.py +++ b/tests/kernels/test_flash_attn_fwd.py @@ -4729,6 +4729,32 @@ def test_paged_fp8_d192_rejects_flattened_int32_overflow(monkeypatch, value_head ) +@_requires_gfx950 +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +def test_paged_legacy_dtype_accepts_host_only_seqlen_metadata(dtype): + """BF16/F16 paged attention retains its permissive seqlen_k ABI.""" + query = torch.zeros((1, 64, 1, 128), device="cuda", dtype=dtype) + key = torch.zeros((1, 64, 1, 128), device="cuda", dtype=dtype) + value = torch.zeros_like(key) + block_table = torch.zeros((1, 1), device="cuda", dtype=torch.int32) + seqlen_k = torch.tensor([64, 0], device="cuda", dtype=torch.int64) + + output = flydsl_flash_attn_func( + query, + key, + value, + causal=True, + num_kv_heads=1, + block_table=block_table, + seqlen_k=seqlen_k, + kv_cache_layout="linear", + ) + torch.cuda.synchronize() + + assert output.shape == query.shape + assert torch.count_nonzero(output) == 0 + + @_requires_gfx950 @pytest.mark.large_shape def test_paged_fp8_d192_cache_offsets_above_4gib(): From 2175c2f5d432cb7120b8239c6c6141155d2304a7 Mon Sep 17 00:00:00 2001 From: Xiake Sun Date: Mon, 31 Aug 2026 13:33:03 +0000 Subject: [PATCH 07/59] fix(fmha): validate paged output device - Reject caller-provided paged outputs that are not colocated with Q/K/V. - Prevent cross-device output buffers from reaching the kernel launch. - Add host-side regression coverage for a CPU output tensor. Signed-off-by: Xiake Sun --- kernels/attention/flash_attn_interface.py | 2 ++ tests/kernels/test_flash_attn_fwd.py | 34 +++++++++++++++++++++++ 2 files changed, 36 insertions(+) diff --git a/kernels/attention/flash_attn_interface.py b/kernels/attention/flash_attn_interface.py index ab7483d5d0..541563ee7f 100644 --- a/kernels/attention/flash_attn_interface.py +++ b/kernels/attention/flash_attn_interface.py @@ -718,6 +718,8 @@ def _flydsl_flash_attn_paged( ) if out is not None: + if out.device != q.device: + raise ValueError(f"flydsl_flash_attn_func: paged output must be on {q.device}, " f"got {out.device}") if out.shape != expected_out_shape or not out.is_contiguous(): raise ValueError( f"flydsl_flash_attn_func: paged output must be contiguous with shape {expected_out_shape}, " diff --git a/tests/kernels/test_flash_attn_fwd.py b/tests/kernels/test_flash_attn_fwd.py index 6cbc93b530..fb55948911 100644 --- a/tests/kernels/test_flash_attn_fwd.py +++ b/tests/kernels/test_flash_attn_fwd.py @@ -4729,6 +4729,40 @@ def test_paged_fp8_d192_rejects_flattened_int32_overflow(monkeypatch, value_head ) +@_requires_gfx950 +def test_paged_fp8_d192_rejects_output_on_wrong_device(): + """A caller-owned output must reside on the same device as Q/K/V.""" + query = torch.zeros((1, 16, 192), device="cuda", dtype=FP8_DTYPE) + key = torch.zeros((1, 1, 12, 64, 16), device="cuda", dtype=FP8_DTYPE) + value = torch.zeros((1, 1, 4, 128, 16), device="cuda", dtype=FP8_DTYPE) + output = torch.empty((1, 16, 128), device="cpu", dtype=torch.bfloat16) + indptr = torch.tensor([0, 1], device="cuda", dtype=torch.int32) + block_table = torch.zeros((1, 1), device="cuda", dtype=torch.int32) + seqlen_k = torch.ones((1,), device="cuda", dtype=torch.int32) + scale = torch.ones((1,), device="cuda", dtype=torch.float32) + + with pytest.raises(ValueError, match="paged output must be on cuda"): + flydsl_flash_attn_func( + query, + key, + value, + causal=True, + num_kv_heads=1, + cu_seqlens_q=indptr, + cu_seqlens_kv=indptr, + max_seqlen_q=1, + max_seqlen_kv=1, + cross_seqlen=True, + block_table=block_table, + seqlen_k=seqlen_k, + kv_cache_layout="vectorized", + q_descale=scale, + k_descale=scale, + v_descale=scale, + out=output, + ) + + @_requires_gfx950 @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) def test_paged_legacy_dtype_accepts_host_only_seqlen_metadata(dtype): From 4b040455f78234759d926eb88af37661e91353d5 Mon Sep 17 00:00:00 2001 From: Xiake Sun Date: Tue, 1 Sep 2026 16:03:55 +0000 Subject: [PATCH 08/59] perf(fmha): reduce paged FP8 V192 register pressure - Load V192 LDS fragments in two smaller groups. - Overlap tail V reads with softmax computation. - Enable post-RA scheduling and disable machine sinking. - Tune D192 setprio and long-context V192 staggering. - Preserve D192/V128 correctness and performance. Signed-off-by: Xiake Sun --- .../attention/flash_attn_fp8_paged_gfx950.py | 105 ++++++++++++++---- kernels/attention/flash_attn_interface.py | 6 +- kernels/attention/flash_attn_utils.py | 63 +++++++---- 3 files changed, 129 insertions(+), 45 deletions(-) diff --git a/kernels/attention/flash_attn_fp8_paged_gfx950.py b/kernels/attention/flash_attn_fp8_paged_gfx950.py index 05b909fe62..5b40e5dfc6 100644 --- a/kernels/attention/flash_attn_fp8_paged_gfx950.py +++ b/kernels/attention/flash_attn_fp8_paged_gfx950.py @@ -302,7 +302,10 @@ def _run_q_block(): # Cluster 2 (memory): prefetch next K (buf1), read this tile's V from # LDS (v_v) for P*V, wait + sync. kv_gmem_to_lds.load_k(j_idx * traits.BLOCK_N, 1) - v_v = kv_lds_to_regs.load_v(0) + if const_expr(traits.D_CHUNKS > 4): + v_v = kv_lds_to_regs.load_v_steps(0, 0, 2) + else: + v_v = kv_lds_to_regs.load_v(0) rocdl.s_waitcnt(traits.LGKMCNT_0_ONLY) _waitcnt_vm_n(ctx.NUM_DMA_K + ctx.NUM_DMA_V) rocdl.sched_barrier(0) @@ -331,10 +334,18 @@ def _run_q_block(): else: v_o, m_row, l_row, v_p_0 = softmax_helper.rescale_o(v_o, m_row, l_row, m_tile_max_a, v_p_0) v_o = gemm_helper.pv_step_k(1, v_p_0, v_v, v_o) - v_o = gemm_helper.pv_step_k(2, v_p_0, v_v, v_o) - v_o = gemm_helper.pv_step_k(3, v_p_0, v_v, v_o) - v_s_1 = softmax_helper.sub_m(v_s_1, m_row) - v_p_1 = softmax_helper.exp2(v_s_1, 0, 16) + if const_expr(traits.D_CHUNKS > 4): + v_v_tail = kv_lds_to_regs.load_v_steps(0, 2, 2) + v_s_1 = softmax_helper.sub_m(v_s_1, m_row) + v_p_1 = softmax_helper.exp2(v_s_1, 0, 16) + rocdl.s_waitcnt(traits.LGKMCNT_0_ONLY) + v_o = gemm_helper.pv_step_k(2, v_p_0, v_v_tail, v_o) + v_o = gemm_helper.pv_step_k(3, v_p_0, v_v_tail, v_o) + else: + v_o = gemm_helper.pv_step_k(2, v_p_0, v_v, v_o) + v_o = gemm_helper.pv_step_k(3, v_p_0, v_v, v_o) + v_s_1 = softmax_helper.sub_m(v_s_1, m_row) + v_p_1 = softmax_helper.exp2(v_s_1, 0, 16) _sched_barrier_pairs(traits, 6, 6, 2) # IGroupLP hint (group 2): 6 MFMA each paired with 3 EXP/TRANS (mask @@ -375,7 +386,10 @@ def _run_q_block(): # Cluster 6 (memory): prefetch next K (buf0), read V packs (buf1), # apply causal mask to v_s_0 (if causal), wait + sync. kv_gmem_to_lds.load_k((j_idx + 1) * traits.BLOCK_N, 0) - v_packs_b = kv_lds_to_regs.load_v(1) + if const_expr(traits.D_CHUNKS > 4): + v_packs_b = kv_lds_to_regs.load_v_steps(1, 0, 2) + else: + v_packs_b = kv_lds_to_regs.load_v(1) if const_expr(traits.CAUSAL): v_s_0 = softmax_helper.causal_mask_prologue_if_needed( v_s_0, @@ -405,10 +419,18 @@ def _run_q_block(): v_o, m_row, l_row, v_p_1 = softmax_helper.rescale_o(v_o, m_row, l_row, m_tile_max_b, v_p_1) v_v = v_packs_b v_o = gemm_helper.pv_step_k(1, v_p_1, v_v, v_o) - v_o = gemm_helper.pv_step_k(2, v_p_1, v_v, v_o) - v_o = gemm_helper.pv_step_k(3, v_p_1, v_v, v_o) - v_s_0 = softmax_helper.sub_m(v_s_0, m_row) - v_p_0 = softmax_helper.exp2(v_s_0, 0, 16) + if const_expr(traits.D_CHUNKS > 4): + v_v_tail = kv_lds_to_regs.load_v_steps(1, 2, 2) + v_s_0 = softmax_helper.sub_m(v_s_0, m_row) + v_p_0 = softmax_helper.exp2(v_s_0, 0, 16) + rocdl.s_waitcnt(traits.LGKMCNT_0_ONLY) + v_o = gemm_helper.pv_step_k(2, v_p_1, v_v_tail, v_o) + v_o = gemm_helper.pv_step_k(3, v_p_1, v_v_tail, v_o) + else: + v_o = gemm_helper.pv_step_k(2, v_p_1, v_v, v_o) + v_o = gemm_helper.pv_step_k(3, v_p_1, v_v, v_o) + v_s_0 = softmax_helper.sub_m(v_s_0, m_row) + v_p_0 = softmax_helper.exp2(v_s_0, 0, 16) _sched_barrier_pairs(traits, 6, 5, 4) _sched_barrier_exp_pairs(traits, 6, 3, 4) if const_expr(traits.DUALWAVE_SWP_SETPRIO): @@ -456,7 +478,10 @@ def _run_q_block(): # Epilogue C2 (memory): prefetch K max_m1, read V packs (buf0), causal mask v_s_1, sync. kv_gmem_to_lds.load_k(max_m1 * traits.BLOCK_N, 1) - v_packs_e3 = kv_lds_to_regs.load_v(0) + if const_expr(traits.D_CHUNKS > 4): + v_packs_e3 = kv_lds_to_regs.load_v_steps(0, 0, 2) + else: + v_packs_e3 = kv_lds_to_regs.load_v(0) if const_expr(traits.CAUSAL): v_s_1 = softmax_helper.causal_mask_prologue_if_needed( v_s_1, @@ -474,7 +499,15 @@ def _run_q_block(): # Epilogue C3 (compute): full P*V + unconditional rescale if const_expr(traits.DUALWAVE_SWP_SETPRIO): rocdl.s_setprio(1) - v_o = gemm_helper.pv(v_p_0, v_packs_e3, v_o) + if const_expr(traits.D_CHUNKS > 4): + v_o = gemm_helper.pv_step_k(0, v_p_0, v_packs_e3, v_o) + v_o = gemm_helper.pv_step_k(1, v_p_0, v_packs_e3, v_o) + v_packs_e3 = kv_lds_to_regs.load_v_steps(0, 2, 2) + rocdl.s_waitcnt(traits.LGKMCNT_0_ONLY) + v_o = gemm_helper.pv_step_k(2, v_p_0, v_packs_e3, v_o) + v_o = gemm_helper.pv_step_k(3, v_p_0, v_packs_e3, v_o) + else: + v_o = gemm_helper.pv(v_p_0, v_packs_e3, v_o) m_tile_max_e3 = softmax_helper.reduce_max(v_s_1) row_max_e3, rescale_e3 = softmax_helper.rescale_from_tile_max(m_row, m_tile_max_e3) m_row = row_max_e3 @@ -516,7 +549,10 @@ def _run_q_block(): rocdl.sched_barrier(0) # Epilogue C6 (memory): read V packs (buf1), causal mask v_s_0, sync. - v_packs_e7 = kv_lds_to_regs.load_v(1) + if const_expr(traits.D_CHUNKS > 4): + v_packs_e7 = kv_lds_to_regs.load_v_steps(1, 0, 2) + else: + v_packs_e7 = kv_lds_to_regs.load_v(1) if const_expr(traits.CAUSAL): v_s_0 = softmax_helper.causal_mask_prologue_if_needed( v_s_0, @@ -534,7 +570,15 @@ def _run_q_block(): # Epilogue C7 (compute, mirror of C3): full P*V + unconditional rescale. if const_expr(traits.DUALWAVE_SWP_SETPRIO): rocdl.s_setprio(1) - v_o = gemm_helper.pv(v_p_1, v_packs_e7, v_o) + if const_expr(traits.D_CHUNKS > 4): + v_o = gemm_helper.pv_step_k(0, v_p_1, v_packs_e7, v_o) + v_o = gemm_helper.pv_step_k(1, v_p_1, v_packs_e7, v_o) + v_packs_e7 = kv_lds_to_regs.load_v_steps(1, 2, 2) + rocdl.s_waitcnt(traits.LGKMCNT_0_ONLY) + v_o = gemm_helper.pv_step_k(2, v_p_1, v_packs_e7, v_o) + v_o = gemm_helper.pv_step_k(3, v_p_1, v_packs_e7, v_o) + else: + v_o = gemm_helper.pv(v_p_1, v_packs_e7, v_o) m_tile_max_e7 = softmax_helper.reduce_max(v_s_0) row_max_e7, rescale_e7 = softmax_helper.rescale_from_tile_max(m_row, m_tile_max_e7) m_row = row_max_e7 @@ -576,7 +620,10 @@ def _run_q_block(): # Epilogue C10 (memory): read last V packs (buf0), causal mask v_s_1, # drain all DMAs (vmcnt 0), sync. - v_packs_e11 = kv_lds_to_regs.load_v(0) + if const_expr(traits.D_CHUNKS > 4): + v_packs_e11 = kv_lds_to_regs.load_v_steps(0, 0, 2) + else: + v_packs_e11 = kv_lds_to_regs.load_v(0) if const_expr(traits.CAUSAL): v_s_1 = softmax_helper.causal_mask_prologue_if_needed( v_s_1, @@ -594,7 +641,15 @@ def _run_q_block(): # Epilogue C11 (compute): full P*V + rescale for v_p_0, then complete the # last tile's softmax in-place (both exp2 halves, sum, cast) since no # further pass follows. - v_o = gemm_helper.pv(v_p_0, v_packs_e11, v_o) + if const_expr(traits.D_CHUNKS > 4): + v_o = gemm_helper.pv_step_k(0, v_p_0, v_packs_e11, v_o) + v_o = gemm_helper.pv_step_k(1, v_p_0, v_packs_e11, v_o) + v_packs_e11 = kv_lds_to_regs.load_v_steps(0, 2, 2) + rocdl.s_waitcnt(traits.LGKMCNT_0_ONLY) + v_o = gemm_helper.pv_step_k(2, v_p_0, v_packs_e11, v_o) + v_o = gemm_helper.pv_step_k(3, v_p_0, v_packs_e11, v_o) + else: + v_o = gemm_helper.pv(v_p_0, v_packs_e11, v_o) m_tile_max_e11 = softmax_helper.reduce_max(v_s_1) row_max_e11, rescale_e11 = softmax_helper.rescale_from_tile_max(m_row, m_tile_max_e11) m_row = row_max_e11 @@ -615,14 +670,25 @@ def _run_q_block(): rocdl.sched_barrier(0) # Epilogue C12 (memory): read the final V packs for the closing P*V. - v_packs_e13 = kv_lds_to_regs.load_v(1) + if const_expr(traits.D_CHUNKS > 4): + v_packs_e13 = kv_lds_to_regs.load_v_steps(1, 0, 2) + else: + v_packs_e13 = kv_lds_to_regs.load_v(1) rocdl.s_waitcnt(traits.LGKMCNT_0_ONLY) rocdl.sched_barrier(0) rocdl.s_barrier() rocdl.sched_barrier(0) # Epilogue C13 (compute): final P*V -> v_o holds the unnormalized output. - v_o = gemm_helper.pv(v_p_1, v_packs_e13, v_o) + if const_expr(traits.D_CHUNKS > 4): + v_o = gemm_helper.pv_step_k(0, v_p_1, v_packs_e13, v_o) + v_o = gemm_helper.pv_step_k(1, v_p_1, v_packs_e13, v_o) + v_packs_e13 = kv_lds_to_regs.load_v_steps(1, 2, 2) + rocdl.s_waitcnt(traits.LGKMCNT_0_ONLY) + v_o = gemm_helper.pv_step_k(2, v_p_1, v_packs_e13, v_o) + v_o = gemm_helper.pv_step_k(3, v_p_1, v_packs_e13, v_o) + else: + v_o = gemm_helper.pv(v_p_1, v_packs_e13, v_o) # Normalize by l_row; zero rows become zero instead of NaN. # Split-K normalizes before packing so O_partial keeps useful mantissa @@ -768,8 +834,9 @@ def launch_flash_attn_dualwave_swp( "fast_fp_math": True, "unsafe_fp_math": True, "llvm_options": { - "enable-post-misched": False, + "enable-post-misched": True, "lsr-drop-solution": True, + "disable-machine-sink": True, }, } diff --git a/kernels/attention/flash_attn_interface.py b/kernels/attention/flash_attn_interface.py index 541563ee7f..e6dd0cd1b4 100644 --- a/kernels/attention/flash_attn_interface.py +++ b/kernels/attention/flash_attn_interface.py @@ -744,6 +744,8 @@ def _flydsl_flash_attn_paged( and (not _arch.startswith("gfx950") or Sq <= _VARLEN_LIGHT_MAX_SEQ) ) if paged_fp8: + paged_setprio = dualwave_swp_setprio and D != 192 + paged_stagger = dualwave_swp_enable_stagger and not ((D, value_head_dim) == (192, 192) and skv >= 65536) exe = _build_paged_fp8( num_heads=H, num_kv_heads=num_kv_heads, @@ -752,8 +754,8 @@ def _flydsl_flash_attn_paged( waves_per_eu=waves_per_eu, daz=daz, lazy_rescale=dualwave_swp_lazy_rescale, - setprio=dualwave_swp_setprio, - enable_stagger=dualwave_swp_enable_stagger, + setprio=paged_setprio, + enable_stagger=paged_stagger, ) elif _paged_light_ok: exe = _build_paged_light( diff --git a/kernels/attention/flash_attn_utils.py b/kernels/attention/flash_attn_utils.py index 5ea9fe0eea..1e433abd02 100644 --- a/kernels/attention/flash_attn_utils.py +++ b/kernels/attention/flash_attn_utils.py @@ -5598,30 +5598,7 @@ def load_v(self, buf_id): # Read all V packs from the bf16 vt scratch for buffer `buf_id`. traits = self.traits if const_expr(traits.KV_VECTORIZED): - v_base = buf_id * traits.VT_BF16_ELEMS - lane_base = ( - (self.lane_mod_32 // traits.SMEM_N_PER_WAVE) * traits.VEC_V_ROW_STRIDE - + self.lane_div_32 * traits.D128_BF - + (self.lane_mod_32 % traits.SMEM_N_PER_WAVE) * traits.VEC_BF - ) - v_base_ptr = buffer_ops.get_element_ptr( - self.lds_vt_base_ptr, - byte_offset=fx.Int32((v_base + lane_base) * traits.EB_BF), - elem_type=T.i8, - ) - packs = [[None] * traits.D_CHUNKS for _ in range(4)] - for dc in range_constexpr(traits.D_CHUNKS): - for k_substep in range_constexpr(4): - const_off = dc * ( - traits.D_CHUNK // traits.SMEM_N_PER_WAVE - ) * traits.VEC_V_ROW_STRIDE + k_substep * (2 * traits.D128_BF) - ptr = buffer_ops.get_element_ptr( - v_base_ptr, - byte_offset=fx.Int32(const_off * traits.EB_BF), - elem_type=T.i8, - ) - packs[k_substep][dc] = llvm.LoadOp(self.v8bf16_type, ptr, alignment=16).result - return packs + return self.load_v_steps(buf_id, 0, 4) urv = ( self.lane_div_32 * traits.URV_GRPK_BF @@ -5640,6 +5617,44 @@ def load_v(self, buf_id): packs[k_substep][dc] = Vec(a).shuffle(Vec(b), [0, 1, 2, 3, 4, 5, 6, 7]).ir_value() return packs + def load_v_steps(self, buf_id, first_step, num_steps): + """Read a compile-time subset of vectorized HIPREC V K-steps. + + V192 otherwise keeps all 24 v8bf16 packs live together with six output + accumulators and spills. Loading two K-steps at a time bounds that + transient register set while preserving the existing LDS layout. + """ + traits = self.traits + if const_expr(not traits.KV_VECTORIZED or traits.FP8_PV): + raise RuntimeError("load_v_steps requires vectorized HIPREC V") + if const_expr(first_step < 0 or num_steps < 1 or first_step + num_steps > 4): + raise RuntimeError("invalid V K-step range") + + v_base = buf_id * traits.VT_BF16_ELEMS + lane_base = ( + (self.lane_mod_32 // traits.SMEM_N_PER_WAVE) * traits.VEC_V_ROW_STRIDE + + self.lane_div_32 * traits.D128_BF + + (self.lane_mod_32 % traits.SMEM_N_PER_WAVE) * traits.VEC_BF + ) + v_base_ptr = buffer_ops.get_element_ptr( + self.lds_vt_base_ptr, + byte_offset=fx.Int32((v_base + lane_base) * traits.EB_BF), + elem_type=T.i8, + ) + packs = [[None] * traits.D_CHUNKS for _ in range(4)] + for dc in range_constexpr(traits.D_CHUNKS): + for k_substep in range_constexpr(first_step, first_step + num_steps): + const_off = dc * (traits.D_CHUNK // traits.SMEM_N_PER_WAVE) * traits.VEC_V_ROW_STRIDE + k_substep * ( + 2 * traits.D128_BF + ) + ptr = buffer_ops.get_element_ptr( + v_base_ptr, + byte_offset=fx.Int32(const_off * traits.EB_BF), + elem_type=T.i8, + ) + packs[k_substep][dc] = llvm.LoadOp(self.v8bf16_type, ptr, alignment=16).result + return packs + def _load_v_fp8_block(self, buf_id): traits = self.traits v_tile_bytes = (traits.BLOCK_N // 8) * (traits.HEAD_DIM_V // 16) * 128 From 5b8fe46c472e4c98b8d6c75de50af325cd897634 Mon Sep 17 00:00:00 2001 From: Xiake Sun Date: Tue, 1 Sep 2026 16:21:14 +0000 Subject: [PATCH 09/59] feat(fmha): enable paged FP8 D128 attention - Accept vectorized page-64 D128/V128 inputs on gfx950. - Use native FP8 P*V for D128 paged attention. - Apply V descale during final output normalization. - Convert vectorized V into the FP8 MFMA operand layout. - Cover ragged batches and explicit-stream execution. Signed-off-by: Xiake Sun --- .../attention/flash_attn_fp8_paged_gfx950.py | 12 +-- kernels/attention/flash_attn_interface.py | 17 ++-- kernels/attention/flash_attn_utils.py | 81 ++++++++++++++++--- tests/kernels/test_flash_attn_fwd.py | 33 ++++++-- 4 files changed, 115 insertions(+), 28 deletions(-) diff --git a/kernels/attention/flash_attn_fp8_paged_gfx950.py b/kernels/attention/flash_attn_fp8_paged_gfx950.py index 5b40e5dfc6..a1b7f60673 100644 --- a/kernels/attention/flash_attn_fp8_paged_gfx950.py +++ b/kernels/attention/flash_attn_fp8_paged_gfx950.py @@ -65,14 +65,14 @@ def build_flash_attn_paged_fp8_module( or not causal or not varlen or not cross_seqlen - or head_dim != 192 - or value_head_dim not in (128, 192) + or (head_dim, value_head_dim) not in ((128, 128), (192, 128), (192, 192)) or kv_cache_layout != "vectorized" or int(num_kv_splits) != 1 ): raise RuntimeError( "paged FP8 flash_attn requires gfx950, causal packed-varlen cross-attention, " - "page-64 vectorized KV, head_dim=192, value_head_dim in {128,192}, " + "page-64 vectorized KV, (head_dim,value_head_dim) in " + "{(128,128),(192,128),(192,192)}, " "and num_kv_splits=1" ) @@ -693,10 +693,12 @@ def _run_q_block(): # Normalize by l_row; zero rows become zero instead of NaN. # Split-K normalizes before packing so O_partial keeps useful mantissa # range; the combine kernel later applies w_s*l_s. - # HIPREC already folds v_descale into the bf16 vt scratch, so O only needs - # the 1/l normalization here. + # HIPREC folds v_descale into the bf16 vt scratch. The direct FP8 + # D128 path keeps raw V and applies its descale once at the end. inv_l_rcp = rocdl.rcp(T.f32, _raw(l_row)) inv_l = ArithValue(fx.Float32(l_row) > ctx.c_zero_f).select(inv_l_rcp, ctx.c_zero_f) + if const_expr(traits.FP8_PV): + inv_l = ArithValue(inv_l) * ctx.vd_fp8 softmax_helper.scale_o(v_o, inv_l) # CLOSE the phase shift: one extra s_barrier on group A (complement of diff --git a/kernels/attention/flash_attn_interface.py b/kernels/attention/flash_attn_interface.py index e6dd0cd1b4..92abd29935 100644 --- a/kernels/attention/flash_attn_interface.py +++ b/kernels/attention/flash_attn_interface.py @@ -539,7 +539,8 @@ def _flydsl_flash_attn_paged( Supported config ONLY (anything else raises): linear/vectorized cache layout with page size 64 and vLLM ``block_table`` / ``seqlen_k`` metadata. BF16/F16 support D64/D128. gfx950 FP8 additionally supports packed-varlen causal - Q/K D192 with vectorized V/output D128 or D192. + Q/K D128 with V/output D128, or Q/K D192 with vectorized V/output + D128 or D192. - Dense 4D Q ``[B, Sq, H, D]``: split-K (num_kv_splits>1) supported (seq_len>=384). - Varlen packed Q ``[total_q, H, D]`` (cu_seqlens_q given): paged K/V looked up per kv-tile via block_table; split-K not supported (matches dense varlen). @@ -614,12 +615,17 @@ def _flydsl_flash_attn_paged( _arch = _gpu_arch(q.device) if not _arch.startswith("gfx950"): raise ValueError(f"flydsl_flash_attn_func: paged FP8 requires gfx950, got '{_arch or 'unknown'}'") + fp8_head_dims = (D, value_head_dim) if not ( - causal and varlen and cross_seqlen is not False and vectorized and D == 192 and value_head_dim in (128, 192) + causal + and varlen + and cross_seqlen is not False + and vectorized + and fp8_head_dims in ((128, 128), (192, 128), (192, 192)) ): raise NotImplementedError( "flydsl_flash_attn_func: paged FP8 requires causal packed-varlen vectorized KV, " - f"Q/K D192, and V D128/D192; got causal={causal}, varlen={varlen}, " + f"Q/K-V D128-D128, D192-D128, or D192-D192; got causal={causal}, varlen={varlen}, " f"layout={kv_cache_layout!r}, Q/K D{D}, V D{value_head_dim}" ) if num_kv_splits != 1: @@ -745,7 +751,7 @@ def _flydsl_flash_attn_paged( ) if paged_fp8: paged_setprio = dualwave_swp_setprio and D != 192 - paged_stagger = dualwave_swp_enable_stagger and not ((D, value_head_dim) == (192, 192) and skv >= 65536) + paged_stagger = dualwave_swp_enable_stagger and not (fp8_head_dims == (192, 192) and skv >= 65536) exe = _build_paged_fp8( num_heads=H, num_kv_heads=num_kv_heads, @@ -972,7 +978,8 @@ def flydsl_flash_attn_func( block_table / seqlen_k: vLLM-style 2D block table metadata. Enables the native paged-KV path, which supports ``bias`` but not ``alibi_slopes``, ``sink``, or ``return_lse``. gfx950 FP8 supports - the causal packed-varlen vectorized page-64 D192 path with V D128/D192. + causal packed-varlen vectorized page-64 D128/V128 and + D192/V128-or-V192 paths. num_kv_splits: Split-K factor (>1: gfx950 only, D=64/128, bf16/f16, seq>=384). ``None`` lets fp8 autotune it; ``1`` keeps the kernel unsplit. fp8_block_m: Pin the fp8 tile height to 128 or 256. fp8 only. diff --git a/kernels/attention/flash_attn_utils.py b/kernels/attention/flash_attn_utils.py index 1e433abd02..e9fc8c7c71 100644 --- a/kernels/attention/flash_attn_utils.py +++ b/kernels/attention/flash_attn_utils.py @@ -2061,9 +2061,11 @@ def dualwave_fp8_dma_per_iter(traits): class PagedDualwaveSwpFp8Traits: """Pure compile-time tile/layout constants for the gfx950 DUALWAVE_SWP fp8 kernel. - fp8 runs a single path: WIDE QK (32x32x64 mfma_scale) feeding HIPREC PV (fp8 V - dequantized into a bf16 ``vt`` LDS scratch, then a bf16 PV MMA). The ``*_BF`` - fields describe that bf16 vt layout; ``ELEM_BYTES`` is 1 (Q/K/V are fp8).""" + QK uses WIDE 32x32x64 ``mfma_scale``. P*V either uses the HIPREC path + (FP8 V dequantized into BF16 LDS scratch) or the D128 native-FP8 path. The + ``*_BF`` fields describe the HIPREC layout; ``ELEM_BYTES`` is 1 because + Q/K/V are FP8. + """ BLOCK_M: int BLOCK_N: int @@ -2292,8 +2294,10 @@ def _make_paged_dualwave_swp_fp8_traits( kv_vectorized = paged and kv_cache_layout == "vectorized" page_size = block_n - fp8_pv = not paged and os.getenv("FLYDSL_FA_FP8_PV", "0") == "1" - fp8_pv_direct = bn128 and not paged + # D128 paged attention can consume its native FP8 V cache directly once + # the vectorized cache is transposed into the existing FP8 PV LDS layout. + fp8_pv = (paged and head_dim == value_head_dim == 128) or (not paged and os.getenv("FLYDSL_FA_FP8_PV", "0") == "1") + fp8_pv_direct = bn128 and (not paged or fp8_pv) if fp8_pv_direct: fp8_pv = True @@ -5172,7 +5176,10 @@ def _pv_fp8(self, v_p, v_v, v_o): return v_o def _pv_step_fp8(self, step, v_p, v_v, v_o): - if const_expr(step == 0): + # The paged pipeline performs its online-softmax correction between + # steps 0 and 1. Repack P after that correction before the remaining + # output chunks consume it. + if const_expr(step == 0 or (self.traits.PAGED and step == 1)): self._pv_p_fp8_cache = self._p_to_fp8_i32x8(v_p) v_op = self._v_concat_i32x8(v_v, step) v_o[step] = self._mfma_acc_fp8_wide(v_op, self._pv_p_fp8_cache, v_o[step]) @@ -5371,9 +5378,13 @@ def _run(): _run() - def load_v(self, tile_start, buf_id): - if const_expr(self.traits.KV_VECTORIZED): - self._stage_vt_dequant_fp8_vectorized(tile_start, buf_id) + def load_v(self, tile_start, buf_id, page_id=None): + if const_expr(self.traits.KV_VECTORIZED and self.traits.FP8_PV_DIRECT): + self._stage_v_fp8_vectorized_bn128(tile_start, buf_id, page_id=page_id) + elif const_expr(self.traits.KV_VECTORIZED and self.traits.FP8_PV): + self._stage_v_fp8_vectorized_dense_layout(tile_start, buf_id, page_id=page_id) + elif const_expr(self.traits.KV_VECTORIZED): + self._stage_vt_dequant_fp8_vectorized(tile_start, buf_id, page_id=page_id) else: self._stage_v_fp8_block_dma(tile_start, buf_id) @@ -5460,10 +5471,11 @@ def _run(): _run() - def _stage_vt_dequant_fp8_vectorized(self, tile_start, buf_id): + def _stage_vt_dequant_fp8_vectorized(self, tile_start, buf_id, page_id=None): """Dequantize SHUFFLE V [H,4,D,16] into the bf16 vectorized LDS tile.""" traits = self.traits - page_id = self.load_page_id(tile_start) + if page_id is None: + page_id = self.load_page_id(tile_start) src_div = self.make_page_view(self.v_base_iter, page_id, is_value=True) vt_buf = buf_id * traits.VT_BF16_ELEMS token_groups16 = traits.BLOCK_N // traits.KV_VEC_SIZE @@ -5516,6 +5528,53 @@ def _load_store_active(): scf_if_dispatch(flat < fx.Index(total_v16), _load_store_active) + def _stage_v_fp8_vectorized_dense_layout(self, tile_start, buf_id, page_id=None): + """Transpose vectorized V into the LDS layout used by direct FP8 P*V.""" + traits = self.traits + if page_id is None: + page_id = self.load_page_id(tile_start) + src_div = self.make_page_view(self.v_base_iter, page_id, is_value=True) + v_tile_bytes = (traits.BLOCK_N // 8) * (traits.V_HEAD_DIM // 16) * 128 + aligned_base = ((self.lds_vt_base_idx + fx.Index(127)) // fx.Index(128)) * fx.Index(128) + dst_base = aligned_base + fx.Index(buf_id * v_tile_bytes) + token_groups16 = traits.BLOCK_N // traits.KV_VEC_SIZE + total_v16 = traits.V_HEAD_DIM * token_groups16 + passes = (total_v16 + traits.BLOCK_SIZE - 1) // traits.BLOCK_SIZE + + for pass_id in range_constexpr(passes): + flat = self.tid + fx.Index(pass_id * traits.BLOCK_SIZE) + + def _load_store_active(): + n_group16 = flat % fx.Index(token_groups16) + d_col = flat // fx.Index(token_groups16) + src_elem = ( + self.kv_head_idx * token_groups16 * traits.V_HEAD_DIM * traits.KV_VEC_SIZE + + n_group16 * traits.V_HEAD_DIM * traits.KV_VEC_SIZE + + d_col * traits.KV_VEC_SIZE + ) + src_i32x4 = self.buffer_load_fp8x16(src_div, src_elem) + src_bytes = Vec(src_i32x4, (4,), fx.Int32).bitcast(fx.Int8) + d_block = d_col // fx.Index(16) + d_in_block = d_col % fx.Index(16) + for n_in_group in range_constexpr(16): + n = n_group16 * fx.Index(16) + fx.Index(n_in_group) + dest_wave = n // fx.Index(8) + dest_n = n % fx.Index(8) + dst_byte = ( + dst_base + + (dest_wave * fx.Index(traits.V_HEAD_DIM // 16) + d_block) * fx.Index(128) + + dest_n * fx.Index(16) + + d_in_block + ) + ptr = buffer_ops.get_element_ptr( + self.lds_vt_base_ptr, + byte_offset=fx.Int32(dst_byte - self.lds_vt_base_idx), + elem_type=T.i8, + ) + llvm.StoreOp(as_mlir_value(src_bytes[_sigma_k_tile_n(n_in_group)]), ptr, alignment=1) + + scf_if_dispatch(flat < fx.Index(total_v16), _load_store_active) + def _stage_vt_dequant_fp8(self, tile_start, buf_id): # Dequantize fp8 V into the exact bf16 V staging positions. The two d-iters # load 8 fp8 at D offsets 64 apart; a contiguous 16B load would gather wrong. diff --git a/tests/kernels/test_flash_attn_fwd.py b/tests/kernels/test_flash_attn_fwd.py index fb55948911..d6d89f0730 100644 --- a/tests/kernels/test_flash_attn_fwd.py +++ b/tests/kernels/test_flash_attn_fwd.py @@ -4576,14 +4576,23 @@ def test_return_lse_rejects_fp8(): @_requires_gfx950 -@pytest.mark.parametrize("value_head_dim", [128, 192]) +@pytest.mark.parametrize( + ("head_dim", "value_head_dim"), + [(128, 128), (192, 128), (192, 192)], + ids=["d128-v128", "d192-v128", "d192-v192"], +) @pytest.mark.parametrize( ("use_non_default_stream", "force_internal_copies"), [(False, False), (True, False), (True, True)], ids=["default", "side", "side-copies"], ) -def test_paged_fp8_d192_asymmetric_value_matches_torch(value_head_dim, use_non_default_stream, force_internal_copies): - """Packed causal FP8 page-64 attention supports native V/output width.""" +def test_paged_fp8_asymmetric_value_matches_torch( + head_dim, + value_head_dim, + use_non_default_stream, + force_internal_copies, +): + """Packed causal FP8 page-64 attention supports native Q/K and V widths.""" torch.manual_seed(17) query_lengths = [64, 32] kv_lengths = [128, 96] @@ -4593,8 +4602,8 @@ def test_paged_fp8_d192_asymmetric_value_matches_torch(value_head_dim, use_non_d seqlen_k = torch.tensor(kv_lengths, device="cuda", dtype=torch.int32) num_pages = 4 - query, query_descale = quantize_per_tensor_fp8(torch.randn(sum(query_lengths), 16, 192, device="cuda") * 0.2) - key, key_descale = quantize_per_tensor_fp8(torch.randn(num_pages, 1, 12, 64, 16, device="cuda") * 0.2) + query, query_descale = quantize_per_tensor_fp8(torch.randn(sum(query_lengths), 16, head_dim, device="cuda") * 0.2) + key, key_descale = quantize_per_tensor_fp8(torch.randn(num_pages, 1, head_dim // 16, 64, 16, device="cuda") * 0.2) value, value_descale = quantize_per_tensor_fp8( torch.randn(num_pages, 1, 4, value_head_dim, 16, device="cuda") * 0.2 ) @@ -4644,6 +4653,16 @@ def noncontiguous_copy(tensor): stream = torch.cuda.Stream() if use_non_default_stream else None if stream is not None: torch.cuda.synchronize() + if force_internal_copies and stream is not None: + flydsl_flash_attn_func( + query.contiguous(), + key.contiguous(), + value.contiguous(), + block_table=block_table.to(torch.int32).contiguous(), + stream=stream, + **call_kwargs, + ) + stream.synchronize() default_done = None if force_internal_copies: torch.cuda._sleep(3_000_000_000) @@ -4671,13 +4690,13 @@ def noncontiguous_copy(tensor): query_batch = query[q_indptr[batch_idx] : q_indptr[batch_idx + 1]].float() * query_descale physical_pages = block_table[batch_idx].long() key_batch = ( - key[physical_pages].permute(0, 3, 1, 2, 4).reshape(-1, 1, 192)[:kv_length].float() * key_descale + key[physical_pages].permute(0, 3, 1, 2, 4).reshape(-1, 1, head_dim)[:kv_length].float() * key_descale ).expand(-1, 16, -1) value_batch = ( value[physical_pages].permute(0, 2, 4, 1, 3).reshape(-1, 1, value_head_dim)[:kv_length].float() * value_descale ).expand(-1, 16, -1) - logits = torch.einsum("qhd,khd->hqk", query_batch, key_batch) / math.sqrt(192) + logits = torch.einsum("qhd,khd->hqk", query_batch, key_batch) / math.sqrt(head_dim) query_positions = torch.arange(query_length, device="cuda") key_positions = torch.arange(kv_length, device="cuda") causal_mask = key_positions[None, :] <= (kv_length - query_length + query_positions)[:, None] From 2fc285e4885668a1b8529992515bc4efa8ecf64e Mon Sep 17 00:00:00 2001 From: Xiake Sun Date: Tue, 1 Sep 2026 16:32:09 +0000 Subject: [PATCH 10/59] perf(fmha): add paired-page D128 fast path - Process adjacent page-64 tiles through a BN128 softmax iteration. - Reuse paired 64-bit block-table loads for K and V. - Stage V through an 80-byte bank-padded D-major LDS layout. - Use DPP lane exchange instead of LDS cross-lane shuffles. - Validate direct BN128 launch assumptions. - Cover even-page fast dispatch, odd-page fallback, partial tails, and graph replay. Signed-off-by: Xiake Sun --- .../attention/flash_attn_fp8_paged_gfx950.py | 321 ++++++++++++++++-- kernels/attention/flash_attn_interface.py | 5 + kernels/attention/flash_attn_utils.py | 117 ++++++- tests/kernels/test_flash_attn_fwd.py | 210 +++++++++++- 4 files changed, 611 insertions(+), 42 deletions(-) diff --git a/kernels/attention/flash_attn_fp8_paged_gfx950.py b/kernels/attention/flash_attn_fp8_paged_gfx950.py index a1b7f60673..c1d5017ebf 100644 --- a/kernels/attention/flash_attn_fp8_paged_gfx950.py +++ b/kernels/attention/flash_attn_fp8_paged_gfx950.py @@ -51,6 +51,7 @@ def build_flash_attn_paged_fp8_module( cross_seqlen=False, paged=False, kv_cache_layout="linear", + paged_bn128=False, ): """Build the gfx950 packed-varlen page-64 FP8 attention launcher.""" gpu_arch = get_hip_arch() @@ -78,6 +79,8 @@ def build_flash_attn_paged_fp8_module( if num_kv_heads is None: num_kv_heads = num_heads + if paged_bn128 and (head_dim, value_head_dim) != (128, 128): + raise RuntimeError("paged BN128 currently requires head_dim=value_head_dim=128") assert num_heads % num_kv_heads == 0 NUM_KV_SPLITS = int(num_kv_splits) assert NUM_KV_SPLITS >= 1 @@ -99,7 +102,7 @@ def build_flash_attn_paged_fp8_module( dualwave_swp_debug_lazy_counts=dualwave_swp_debug_lazy_counts, dualwave_swp_enable_stagger=dualwave_swp_enable_stagger, num_kv_splits=num_kv_splits, - varlen=varlen, + varlen=False if paged_bn128 else varlen, cross_seqlen=cross_seqlen, paged=paged, kv_cache_layout=kv_cache_layout, @@ -111,16 +114,219 @@ def build_flash_attn_paged_fp8_module( HEAD_DIM = traits.HEAD_DIM NUM_HEADS_Q = traits.NUM_HEADS_Q PAGED = traits.PAGED + PAGED_BN128 = bool(paged_bn128) DEFAULT_STRIDE_Q_N = traits.DEFAULT_STRIDE_Q_N DEFAULT_STRIDE_O_N = traits.NUM_HEADS_Q * traits.V_HEAD_DIM DEFAULT_STRIDE_KV_N = traits.DEFAULT_STRIDE_KV_N _dualwave_swp_fp8_cache_tag = traits.cache_tag _lds_elem_dtype = dtype_to_elem_type(traits.DTYPE_STR) - @fx.struct - class SharedStorage: - kv: fx.Array[_lds_elem_dtype, traits.LDS_KV_TOTAL_SIZE, 16] - vt: fx.Array[fx.BFloat16, traits.VT_BF16_TOTAL, 16] + if PAGED_BN128: + + @fx.struct + class SharedStorage: + kv: fx.Array[_lds_elem_dtype, traits.LDS_KV_TOTAL_SIZE, 16] + vt: fx.Array[fx.BFloat16, traits.VT_BF16_TOTAL, 16] + q: fx.Array[_lds_elem_dtype, BLOCK_M * HEAD_DIM, 16] + + else: + + @fx.struct + class SharedStorage: + kv: fx.Array[_lds_elem_dtype, traits.LDS_KV_TOTAL_SIZE, 16] + vt: fx.Array[fx.BFloat16, traits.VT_BF16_TOTAL, 16] + + # BN128: two BLOCK_N=64 KV tiles per iteration, one merged softmax correction. + @flyc.kernel(known_block_size=[BLOCK_SIZE, 1, 1]) + def flash_attn_paged_fp8_bn128_kernel( + Q: fx.Tensor, + K: fx.Tensor, + V: fx.Tensor, + O: fx.Tensor, # noqa: E741 + DebugCounts: fx.Tensor, + CuSeqQ: fx.Tensor, + CuSeqKv: fx.Tensor, + BlockTable: fx.Tensor, + block_table_stride: fx.Int32, + QDescale: fx.Tensor, + KDescale: fx.Tensor, + VDescale: fx.Tensor, + seq_len: fx.Int32, + seq_len_kv: fx.Int32, + stride_q_n: fx.Int32, + stride_kv_n: fx.Int32, + head_dim_runtime: fx.Int32, + ): + ctx = DualwaveFp8KernelContext( + traits, + Q, + K, + V, + O, + DebugCounts, + CuSeqQ, + CuSeqKv, + QDescale, + KDescale, + VDescale, + seq_len, + seq_len_kv, + stride_q_n, + stride_kv_n, + head_dim_runtime, + stride_o_n=DEFAULT_STRIDE_O_N, + BlockTable=BlockTable, + block_table_stride=block_table_stride, + ) + ctx.init_types_and_constants() + ctx.init_runtime_indices() + ctx.init_lds(SharedStorage) + ctx.init_thread_mapping() + if const_expr(traits.CAUSAL): + ctx.init_causal_lpt_order() + ctx.init_sequence_lengths() + ctx.init_descriptors() + ctx.init_atoms_and_lds_ptrs() + ctx.init_dma_thread_offsets() + ctx.init_descale() + ctx.init_tile_bounds() + ctx.init_workspace_io() + + q_loader = DualwaveFp8QLoader(ctx) + gemm_helper = DualwaveFp8GemmHelper(ctx) + softmax_helper = DualwaveFp8SoftmaxHelper(ctx) + kv_gmem_to_lds = DualwaveFp8KvGmemToLdsLoader(ctx) + kv_lds_to_regs = DualwaveFp8KvLdsToVgprLoader(ctx) + output_store = DualwaveFp8StoreHelper(ctx) + + BN = traits.BLOCK_N + D_CHUNKS = traits.D_CHUNKS + NPF = const_expr(traits.NUM_PREFETCH_K) + t0 = ctx.split_t0 + t_end = ctx.split_t_end + + def _softmax_part(v_s, l_row, m_new): + v_s = softmax_helper.sub_m(v_s, m_new) + v_p = softmax_helper.exp2(v_s, 0, 16) + v_p = softmax_helper.exp2(v_p, 16, 16) + l_row = softmax_helper.reduce_sum(l_row, v_p) + v_p = gemm_helper.cast_p_fp8_direct(v_p) + return v_p, l_row + + def _pv_part(v_p, v_v, v_o): + v_o = gemm_helper.pv(v_p, v_v, v_o) + return softmax_helper.anchor_v_o(v_o) + + def _subtile_tail(v_s, v_v, v_o, l_row, m_new): + v_p, l_row = _softmax_part(v_s, l_row, m_new) + v_o = _pv_part(v_p, v_v, v_o) + return v_o, l_row + + def _mask_pair(v_s_a, v_s_b, j): + if const_expr(traits.CAUSAL): + return softmax_helper.causal_mask_pair_if_needed(v_s_a, v_s_b, j) + return v_s_a, v_s_b + + def _correct_o(v_o, m_row, l_row, m_tile): + if const_expr(traits.DUALWAVE_SWP_LAZY_RESCALE): + return softmax_helper.lazy_correct_o(v_o, m_row, l_row, m_tile) + m_new, corr = softmax_helper.rescale_from_tile_max(m_row, m_tile) + softmax_helper.scale_o(v_o, corr) + return v_o, m_new, softmax_helper.apply_l_rescale(l_row, corr) + + def _merge_tile_max(v_s_a, v_s_b): + m_tile = softmax_helper.max2(softmax_helper.reduce_max(v_s_a), softmax_helper.reduce_max(v_s_b)) + if const_expr(traits.CAUSAL): + m_tile = softmax_helper.floor_masked_max(m_tile) + return m_tile + + page_t0, page_t1 = ctx.load_page_id_pair(t0 * BN) + kv_gmem_to_lds.load_k(t0 * BN, t0 % fx.Index(NPF), page_id=page_t0) + if const_expr(traits.QREG): + q_loader.stage_q_to_lds() + rocdl.s_waitcnt(0) + rocdl.sched_barrier(0) + rocdl.s_barrier() + + ctx.init_q_row() + q_row = ctx.q_row + + q_wide = gemm_helper.load_q_wide() if const_expr(traits.QREG) else None + + page_t2, page_t3 = ctx.load_page_id_pair((t0 + 2) * BN) + kv_gmem_to_lds.load_k((t0 + 1) * BN, (t0 + 1) % fx.Index(NPF), page_id=page_t1) + kv_gmem_to_lds.load_v(t0 * BN, t0 % fx.Index(NPF), page_id=page_t0) + kv_gmem_to_lds.load_v((t0 + 1) * BN, (t0 + 1) % fx.Index(NPF), page_id=page_t1) + kv_gmem_to_lds.load_k((t0 + 2) * BN, (t0 + 2) % fx.Index(NPF), page_id=page_t2) + kv_gmem_to_lds.load_k((t0 + 3) * BN, (t0 + 3) % fx.Index(NPF), page_id=page_t3) + kv_gmem_to_lds.load_v((t0 + 2) * BN, (t0 + 2) % fx.Index(NPF), page_id=page_t2) + kv_gmem_to_lds.load_v((t0 + 3) * BN, (t0 + 3) % fx.Index(NPF), page_id=page_t3) + rocdl.s_waitcnt(0) + rocdl.sched_barrier(0) + rocdl.s_barrier() + rocdl.sched_barrier(0) + + m_row = ctx.c_neg_inf + l_row = ctx.c_zero_f + v_o = [ctx.c_zero_v16f32 for _ in range_constexpr(D_CHUNKS)] + + NPF_I = const_expr(fx.Index(NPF)) + + def _ring_wrap(x): + return (x >= NPF_I).select(x - NPF_I, x) + + init_args = [m_row, l_row] + v_o + [t0 % fx.Index(NPF)] + loop_results = init_args + for j, loop_args in range(fx.Index(t0), t_end, fx.Index(2), init=init_args): + m_row = loop_args[0] + l_row = loop_args[1] + v_o = [loop_args[2 + i] for i in range_constexpr(D_CHUNKS)] + + a_buf = loop_args[2 + D_CHUNKS] + b_buf = _ring_wrap(a_buf + fx.Index(1)) + nn_a_buf = _ring_wrap(a_buf + fx.Index(2)) + f_a_buf = _ring_wrap(a_buf + fx.Index(4)) + f_b_buf = _ring_wrap(a_buf + fx.Index(5)) + + v_k_a = kv_lds_to_regs.load_k(a_buf) + v_k_b = kv_lds_to_regs.load_k(b_buf) + v_v_a = kv_lds_to_regs.load_v(a_buf) + + page_f_a, page_f_b = ctx.load_page_id_pair((j + fx.Index(4)) * BN) + kv_gmem_to_lds.load_k((j + fx.Index(4)) * BN, f_a_buf, page_id=page_f_a) + kv_gmem_to_lds.load_k((j + fx.Index(5)) * BN, f_b_buf, page_id=page_f_b) + kv_gmem_to_lds.load_v((j + fx.Index(4)) * BN, f_a_buf, page_id=page_f_a) + kv_gmem_to_lds.load_v((j + fx.Index(5)) * BN, f_b_buf, page_id=page_f_b) + + v_s_a = gemm_helper.qk(v_k_a, q_wide) + v_s_b = gemm_helper.qk(v_k_b, q_wide) + v_s_a, v_s_b = _mask_pair(v_s_a, v_s_b, j) + m_tile = _merge_tile_max(v_s_a, v_s_b) + v_o, m_new, l_row = _correct_o(v_o, m_row, l_row, m_tile) + v_o = softmax_helper.anchor_v_o(v_o) + + v_o, l_row = _subtile_tail(v_s_a, v_v_a, v_o, l_row, m_new) + v_v_b = kv_lds_to_regs.load_v(b_buf) + v_o, l_row = _subtile_tail(v_s_b, v_v_b, v_o, l_row, m_new) + m_row = m_new + + rocdl.s_waitcnt(0) + rocdl.sched_barrier(0) + rocdl.s_barrier() + rocdl.sched_barrier(0) + + loop_results = yield [m_row, l_row] + v_o + [nn_a_buf] + m_row = loop_results[0] + l_row = loop_results[1] + v_o = [loop_results[2 + i] for i in range_constexpr(D_CHUNKS)] + + inv_l_rcp = rocdl.rcp(T.f32, _raw(l_row)) + inv_l = ArithValue(fx.Float32(l_row) > ctx.c_zero_f).select(inv_l_rcp, ctx.c_zero_f) + if const_expr(traits.FP8_PV): + inv_l = ArithValue(inv_l) * ctx.vd_fp8 + softmax_helper.scale_o(v_o, inv_l) + rocdl.s_barrier() + output_store.store_final_o(v_o, q_row) @flyc.kernel(known_block_size=[BLOCK_SIZE, 1, 1]) def flash_attn_dualwave_swp_fp8_gfx950_kernel( @@ -795,35 +1001,62 @@ def launch_flash_attn_dualwave_swp( if const_expr(daz) else None ) - flash_attn_dualwave_swp_fp8_gfx950_kernel( - Q, - K, - V, - O, - DebugCounts, - CuSeqQ, - CuSeqKv, - BlockTable, - block_table_stride, - QDescale, - KDescale, - VDescale, - seq_len, - seq_len_kv, - stride_q_n, - stride_o_n, - stride_kv_n, - head_dim_runtime, - value_attrs={ - "rocdl.waves_per_eu": waves_per_eu, - "rocdl.flat_work_group_size": f"{BLOCK_SIZE},{BLOCK_SIZE}", - "passthrough": passthrough_entries, - }, - ).launch( - grid=(NUM_HEADS_Q, num_q_blocks, grid_z), - block=(BLOCK_SIZE, 1, 1), - stream=stream, - ) + kernel_attrs = { + "rocdl.waves_per_eu": waves_per_eu, + "rocdl.flat_work_group_size": f"{BLOCK_SIZE},{BLOCK_SIZE}", + "passthrough": passthrough_entries, + } + if const_expr(PAGED_BN128): + flash_attn_paged_fp8_bn128_kernel( + Q, + K, + V, + O, + DebugCounts, + CuSeqQ, + CuSeqKv, + BlockTable, + block_table_stride, + QDescale, + KDescale, + VDescale, + seq_len, + seq_len_kv, + stride_q_n, + stride_kv_n, + head_dim_runtime, + value_attrs=kernel_attrs, + ).launch( + grid=(NUM_HEADS_Q, num_q_blocks, grid_z), + block=(BLOCK_SIZE, 1, 1), + stream=stream, + ) + else: + flash_attn_dualwave_swp_fp8_gfx950_kernel( + Q, + K, + V, + O, + DebugCounts, + CuSeqQ, + CuSeqKv, + BlockTable, + block_table_stride, + QDescale, + KDescale, + VDescale, + seq_len, + seq_len_kv, + stride_q_n, + stride_o_n, + stride_kv_n, + head_dim_runtime, + value_attrs=kernel_attrs, + ).launch( + grid=(NUM_HEADS_Q, num_q_blocks, grid_z), + block=(BLOCK_SIZE, 1, 1), + stream=stream, + ) if const_expr(SPLITK): combine_rows = bs_idx * NUM_HEADS_Q * sl_idx flash_attn_splitk_combine_kernel(O, DebugCounts, batch_size, seq_len, stride_o_n).launch( @@ -842,6 +1075,24 @@ def launch_flash_attn_dualwave_swp( }, } + def _validate_paged_bn128_launch(batch_size, seq_len_kv, block_table_stride): + if not PAGED_BN128: + return + batch_size = int(batch_size) + seq_len_kv = int(seq_len_kv) + block_table_stride = int(block_table_stride) + num_kv_pages = (seq_len_kv + traits.PAGE_SIZE - 1) // traits.PAGE_SIZE + if batch_size != 1 or num_kv_pages < 2 or num_kv_pages % 2 != 0: + raise ValueError( + "paged BN128 requires batch_size=1 and a positive even number " + f"of KV pages; got batch_size={batch_size}, seq_len_kv={seq_len_kv}, " + f"page_size={traits.PAGE_SIZE}" + ) + if block_table_stride < num_kv_pages: + raise ValueError( + "paged BN128 block table has too few entries: " f"need {num_kv_pages}, got stride {block_table_stride}" + ) + def _launch( Q, K, @@ -895,6 +1146,7 @@ def _launch( block_table_stride = 0 if PAGED and block_table is O: raise ValueError("paged fp8 flash_attn requires block_table") + _validate_paged_bn128_launch(batch_size, seq_len_kv, block_table_stride) # Per-tensor fp8 descales (shape-[1] fp32). The kernel only reads them on # the fp8 path; bf16/f16 launches pass O as an unused placeholder. if q_descale is None: @@ -997,6 +1249,7 @@ def _compile( block_table = O if block_table_stride is None: block_table_stride = 0 + _validate_paged_bn128_launch(batch_size, seq_len_kv, block_table_stride) if q_descale is None: q_descale = O if k_descale is None: diff --git a/kernels/attention/flash_attn_interface.py b/kernels/attention/flash_attn_interface.py index 92abd29935..6ba4e76ca5 100644 --- a/kernels/attention/flash_attn_interface.py +++ b/kernels/attention/flash_attn_interface.py @@ -475,6 +475,7 @@ def _build_paged_fp8( lazy_rescale: bool, setprio: bool, enable_stagger: bool, + use_bn128: bool, ): """Build the gfx950 packed-varlen, vectorized page-64 FP8 launcher.""" from kernels.attention.flash_attn_fp8_paged_gfx950 import build_flash_attn_paged_fp8_module @@ -496,6 +497,7 @@ def _build_paged_fp8( cross_seqlen=True, paged=True, kv_cache_layout="vectorized", + paged_bn128=use_bn128, ) @@ -750,6 +752,8 @@ def _flydsl_flash_attn_paged( and (not _arch.startswith("gfx950") or Sq <= _VARLEN_LIGHT_MAX_SEQ) ) if paged_fp8: + num_kv_pages = (skv + page_size - 1) // page_size + use_bn128 = fp8_head_dims == (128, 128) and B == 1 and num_kv_pages % 2 == 0 paged_setprio = dualwave_swp_setprio and D != 192 paged_stagger = dualwave_swp_enable_stagger and not (fp8_head_dims == (192, 192) and skv >= 65536) exe = _build_paged_fp8( @@ -762,6 +766,7 @@ def _flydsl_flash_attn_paged( lazy_rescale=dualwave_swp_lazy_rescale, setprio=paged_setprio, enable_stagger=paged_stagger, + use_bn128=use_bn128, ) elif _paged_light_ok: exe = _build_paged_light( diff --git a/kernels/attention/flash_attn_utils.py b/kernels/attention/flash_attn_utils.py index e9fc8c7c71..26d623f2d4 100644 --- a/kernels/attention/flash_attn_utils.py +++ b/kernels/attention/flash_attn_utils.py @@ -2133,6 +2133,7 @@ class PagedDualwaveSwpFp8Traits: VEC_V_ROW_STRIDE: int VT_BF16_ELEMS: int VT_BF16_TOTAL: int + FP8_V_ROW_STRIDE: int URV_GRPK_BF: int URV_GRP_N_BF: int URV_LANE_LO_BF: int @@ -2188,6 +2189,7 @@ def cache_tag(self): self.LANE_SPLIT_KV, self.VT_BF16_ELEMS, self.VT_BF16_TOTAL, + self.FP8_V_ROW_STRIDE, self.FP8_PV, self.FP8_PV_DIRECT, self.NUM_PREFETCH_K, @@ -2283,7 +2285,10 @@ def _make_paged_dualwave_swp_fp8_traits( sdrpt_bf = value_head_dim // d128_bf vls_bf = slw_bf + 64 // eb_bf vt_bf16_elems = snrpt_bf * sdrpt_bf * vls_bf + fp8_v_row_stride = block_n + 16 fp8_v_tile_bytes = (block_n // 8) * (head_dim // 16) * 128 + if paged and bn128 and head_dim == value_head_dim == 128: + fp8_v_tile_bytes = value_head_dim * fp8_v_row_stride if bn128_pf: vt_bf16_total = num_prefetch_k * (fp8_v_tile_bytes // eb_bf) + 128 else: @@ -2367,6 +2372,7 @@ def _make_paged_dualwave_swp_fp8_traits( VEC_V_ROW_STRIDE=vls_bf, VT_BF16_ELEMS=vt_bf16_elems, VT_BF16_TOTAL=vt_bf16_total, + FP8_V_ROW_STRIDE=fp8_v_row_stride, URV_GRPK_BF=4 * vls_bf, URV_GRP_N_BF=16, URV_LANE_LO_BF=4, @@ -4888,6 +4894,8 @@ def _make_buf_div(tensor, nrec_bytes): ) self.page_i32_atom = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), fx.Int32) self.page_v1i32 = Vec.make_type(1, fx.Int32) + self.page_i32x2_atom = fx.make_copy_atom(fx.rocdl.BufferCopy64b(), fx.Int32) + self.page_v2i32 = Vec.make_type(2, fx.Int32) self.page_base = self.batch_idx * fx.Index(self.block_table_stride) self.k_div = None self.v_div = None @@ -4908,6 +4916,8 @@ def _make_buf_div(tensor, nrec_bytes): self.page_indices_div = None self.page_i32_atom = None self.page_v1i32 = None + self.page_i32x2_atom = None + self.page_v2i32 = None self.page_base = None def init_atoms_and_lds_ptrs(self): @@ -5068,6 +5078,21 @@ def load_page_id(self, tile_start): ) return fx.Index(fx.Int32(rocdl.readfirstlane(T.i32, as_mlir_value(fx.Int32(Vec(v, (1,), fx.Int32)[0]))))) + def load_page_id_pair(self, tile_start): + """Load two adjacent physical page IDs for a BN128 tile pair.""" + local_page = fx.Index(tile_start) // fx.Index(self.traits.PAGE_SIZE) + page_idx = self.page_base + local_page + safe_idx = fx.Index((local_page < self.num_kv_tiles).select(page_idx, self.page_base)) + v = fly.copy_atom_call_ssa( + [self.page_v2i32], + self.page_i32x2_atom, + fx.slice(self.page_indices_div, (None, fx.Int32(safe_idx))), + ) + pages = Vec(v, (2,), fx.Int32) + page0 = rocdl.readfirstlane(T.i32, as_mlir_value(fx.Int32(pages[0]))) + page1 = rocdl.readfirstlane(T.i32, as_mlir_value(fx.Int32(pages[1]))) + return fx.Index(fx.Int32(page0)), fx.Index(fx.Int32(page1)) + def make_page_view(self, tensor_iter, page_id, *, is_value=False): if const_expr(is_value): page_byte_stride = self.v_page_byte_stride @@ -5316,7 +5341,7 @@ class DualwaveFp8KvGmemToLdsLoader(DualwaveFp8KernelContext): def __init__(self, ctx): super().__init__(ctx) - def load_k(self, tile_start, buf_id): + def load_k(self, tile_start, buf_id, page_id=None): """DMA one K tile into LDS, one pass per head-dim band. A band's LDS line is this wave's n-rows of `chunk` bytes, row-contiguous, @@ -5327,7 +5352,8 @@ def load_k(self, tile_start, buf_id): eb = traits.ELEM_BYTES k_lds_byte_base = self.lds_kv_base_idx + self.k_buf_base(buf_id) * eb if const_expr(traits.KV_VECTORIZED): - page_id = self.load_page_id(tile_start) + if page_id is None: + page_id = self.load_page_id(tile_start) src_div = self.make_page_view(self.k_base_iter, page_id) # Each wave owns two 64-lane octet rows for the D256-shaped LDS # tile. The QK reader ignores the four groups beyond physical D192. @@ -5528,6 +5554,70 @@ def _load_store_active(): scf_if_dispatch(flat < fx.Index(total_v16), _load_store_active) + def _stage_v_fp8_vectorized_bn128(self, tile_start, buf_id, page_id=None): + """Stage V in a bank-padded D-major layout for paged BN128 P*V.""" + traits = self.traits + if page_id is None: + page_id = self.load_page_id(tile_start) + src_div = self.make_page_view(self.v_base_iter, page_id, is_value=True) + row_stride = traits.FP8_V_ROW_STRIDE + v_tile_bytes = traits.V_HEAD_DIM * row_stride + aligned_base = ((self.lds_vt_base_idx + fx.Index(127)) // fx.Index(128)) * fx.Index(128) + dst_base = aligned_base + fx.Index(buf_id * v_tile_bytes) + token_groups16 = traits.BLOCK_N // traits.KV_VEC_SIZE + total_v16 = traits.V_HEAD_DIM * token_groups16 + passes = (total_v16 + traits.BLOCK_SIZE - 1) // traits.BLOCK_SIZE + + for pass_id in range_constexpr(passes): + flat = self.tid + fx.Index(pass_id * traits.BLOCK_SIZE) + + def _load_store_active(): + n_group16 = flat % fx.Index(token_groups16) + d_col = flat // fx.Index(token_groups16) + src_elem = ( + self.kv_head_idx * token_groups16 * traits.V_HEAD_DIM * traits.KV_VEC_SIZE + + n_group16 * traits.V_HEAD_DIM * traits.KV_VEC_SIZE + + d_col * traits.KV_VEC_SIZE + ) + src_i32x4 = self.buffer_load_fp8x16(src_div, src_elem) + src_words = Vec(src_i32x4, (4,), fx.Int32) + # gfx9 DPP quad_perm [1,0,3,2] is lane^1 without using the LDS + # cross-lane pipe, which remains busy serving the P*V fragments. + peer_words = [ + fx.Int32( + rocdl.update_dpp( + T.i32, + as_mlir_value(fx.Int32(src_words[i])), + as_mlir_value(fx.Int32(src_words[i])), + 0xB1, + 0xF, + 0xF, + True, + ) + ) + for i in range_constexpr(4) + ] + odd_group = (fx.Int32(n_group16) & fx.Int32(1)) != fx.Int32(0) + reordered = Vec.from_elements( + [ + odd_group.select(peer_words[1], src_words[0]), + odd_group.select(peer_words[3], src_words[2]), + odd_group.select(src_words[1], peer_words[0]), + odd_group.select(src_words[3], peer_words[2]), + ], + fx.Int32, + ) + dest_group16 = (n_group16 % fx.Index(2)) * fx.Index(2) + n_group16 // fx.Index(2) + dst_byte = dst_base + d_col * fx.Index(row_stride) + dest_group16 * fx.Index(16) + ptr = buffer_ops.get_element_ptr( + self.lds_vt_base_ptr, + byte_offset=fx.Int32(dst_byte - self.lds_vt_base_idx), + elem_type=T.i8, + ) + llvm.StoreOp(as_mlir_value(reordered), ptr, alignment=16) + + scf_if_dispatch(flat < fx.Index(total_v16), _load_store_active) + def _stage_v_fp8_vectorized_dense_layout(self, tile_start, buf_id, page_id=None): """Transpose vectorized V into the LDS layout used by direct FP8 P*V.""" traits = self.traits @@ -5653,6 +5743,8 @@ def _read_strip(key): def load_v(self, buf_id): if const_expr(self.traits.FP8_PV): + if const_expr(self.traits.KV_VECTORIZED and self.traits.FP8_PV_DIRECT): + return self._load_v_fp8_vectorized_bn128(buf_id) return self._load_v_fp8_block(buf_id) # Read all V packs from the bf16 vt scratch for buffer `buf_id`. traits = self.traits @@ -5676,6 +5768,27 @@ def load_v(self, buf_id): packs[k_substep][dc] = Vec(a).shuffle(Vec(b), [0, 1, 2, 3, 4, 5, 6, 7]).ir_value() return packs + def _load_v_fp8_vectorized_bn128(self, buf_id): + """Read the padded D-major V tile into the FP8 MFMA A fragment.""" + traits = self.traits + row_stride = traits.FP8_V_ROW_STRIDE + v_tile_bytes = traits.V_HEAD_DIM * row_stride + aligned_base = ((self.lds_vt_base_idx + fx.Index(127)) // fx.Index(128)) * fx.Index(128) + tile_base = aligned_base + fx.Index(buf_id * v_tile_bytes) + token_base = self.lane_div_32 * fx.Index(32) + packs = [[None] * traits.D_CHUNKS for _ in range(4)] + for dc in range_constexpr(traits.D_CHUNKS): + d_row = fx.Index(dc * traits.D_CHUNK) + self.lane_mod_32 + row_base = tile_base + d_row * fx.Index(row_stride) + token_base + for ks in range_constexpr(4): + ptr = buffer_ops.get_element_ptr( + self.lds_vt_base_ptr, + byte_offset=fx.Int32(row_base + fx.Index(ks * 8) - self.lds_vt_base_idx), + elem_type=T.i8, + ) + packs[ks][dc] = llvm.LoadOp(T.i64, ptr, alignment=8).result + return packs + def load_v_steps(self, buf_id, first_step, num_steps): """Read a compile-time subset of vectorized HIPREC V K-steps. diff --git a/tests/kernels/test_flash_attn_fwd.py b/tests/kernels/test_flash_attn_fwd.py index d6d89f0730..567017fde0 100644 --- a/tests/kernels/test_flash_attn_fwd.py +++ b/tests/kernels/test_flash_attn_fwd.py @@ -4586,21 +4586,38 @@ def test_return_lse_rejects_fp8(): [(False, False), (True, False), (True, True)], ids=["default", "side", "side-copies"], ) +@pytest.mark.parametrize( + ("query_lengths", "kv_lengths", "block_table_rows"), + [ + ([64, 32], [128, 96], [[2, 0], [3, 1]]), + ([64], [128], [[1, 0]]), + ([64], [192], [[2, 0, 1]]), + ], + ids=["ragged", "single-sequence-even-pages", "single-sequence-odd-pages"], +) def test_paged_fp8_asymmetric_value_matches_torch( head_dim, value_head_dim, use_non_default_stream, force_internal_copies, + query_lengths, + kv_lengths, + block_table_rows, ): """Packed causal FP8 page-64 attention supports native Q/K and V widths.""" + if len(query_lengths) == 1 and force_internal_copies and head_dim == 192: + pytest.skip("D192 wrapper-copy stream ordering is covered by the ragged cases") torch.manual_seed(17) - query_lengths = [64, 32] - kv_lengths = [128, 96] - q_indptr = torch.tensor([0, 64, 96], device="cuda", dtype=torch.int32) - kv_indptr = torch.tensor([0, 128, 224], device="cuda", dtype=torch.int32) - block_table = torch.tensor([[2, 0], [3, 1]], device="cuda", dtype=torch.int32) + q_offsets = [0] + kv_offsets = [0] + for query_length, kv_length in zip(query_lengths, kv_lengths): + q_offsets.append(q_offsets[-1] + query_length) + kv_offsets.append(kv_offsets[-1] + kv_length) + q_indptr = torch.tensor(q_offsets, device="cuda", dtype=torch.int32) + kv_indptr = torch.tensor(kv_offsets, device="cuda", dtype=torch.int32) + block_table = torch.tensor(block_table_rows, device="cuda", dtype=torch.int32) seqlen_k = torch.tensor(kv_lengths, device="cuda", dtype=torch.int32) - num_pages = 4 + num_pages = max(max(row) for row in block_table_rows) + 1 query, query_descale = quantize_per_tensor_fp8(torch.randn(sum(query_lengths), 16, head_dim, device="cuda") * 0.2) key, key_descale = quantize_per_tensor_fp8(torch.randn(num_pages, 1, head_dim // 16, 64, 16, device="cuda") * 0.2) @@ -4709,6 +4726,187 @@ def noncontiguous_copy(tensor): torch.testing.assert_close(actual, expected, rtol=2.0e-2, atol=2.0e-2) +@_requires_gfx950 +@pytest.mark.parametrize( + ("head_dim", "value_head_dim"), + [(128, 128), (192, 128), (192, 192)], + ids=["d128-v128", "d192-v128", "d192-v192"], +) +def test_paged_fp8_long_context_random_pages_matches_torch(head_dim, value_head_dim): + """Long context, shuffled pages, and a partial tail preserve accuracy.""" + torch.manual_seed(23) + query_length = 63 + kv_length = 8161 + num_pages = (kv_length + 63) // 64 + + query, query_descale = quantize_per_tensor_fp8(torch.randn(query_length, 16, head_dim, device="cuda") * 0.2) + key, key_descale = quantize_per_tensor_fp8(torch.randn(num_pages, 1, head_dim // 16, 64, 16, device="cuda") * 0.2) + value, value_descale = quantize_per_tensor_fp8( + torch.randn(num_pages, 1, 4, value_head_dim, 16, device="cuda") * 0.2 + ) + block_table = torch.randperm(num_pages, device="cuda", dtype=torch.int32).reshape(1, num_pages) + q_indptr = torch.tensor([0, query_length], device="cuda", dtype=torch.int32) + kv_indptr = torch.tensor([0, kv_length], device="cuda", dtype=torch.int32) + seqlen_k = torch.tensor([kv_length], device="cuda", dtype=torch.int32) + + actual = flydsl_flash_attn_func( + query, + key, + value, + causal=True, + num_kv_heads=1, + cu_seqlens_q=q_indptr, + cu_seqlens_kv=kv_indptr, + max_seqlen_q=query_length, + max_seqlen_kv=kv_length, + cross_seqlen=True, + block_table=block_table, + seqlen_k=seqlen_k, + kv_cache_layout="vectorized", + q_descale=query_descale, + k_descale=key_descale, + v_descale=value_descale, + ) + torch.cuda.synchronize() + + physical_pages = block_table[0].long() + query_ref = query.float() * query_descale + key_ref = ( + key[physical_pages].permute(0, 3, 1, 2, 4).reshape(-1, 1, head_dim)[:kv_length].float() * key_descale + ).expand(-1, 16, -1) + value_ref = ( + value[physical_pages].permute(0, 2, 4, 1, 3).reshape(-1, 1, value_head_dim)[:kv_length].float() * value_descale + ).expand(-1, 16, -1) + logits = torch.einsum("qhd,khd->hqk", query_ref, key_ref) / math.sqrt(head_dim) + query_positions = torch.arange(query_length, device="cuda") + key_positions = torch.arange(kv_length, device="cuda") + causal_mask = key_positions[None, :] <= (kv_length - query_length + query_positions)[:, None] + logits.masked_fill_(~causal_mask[None, :, :], float("-inf")) + expected = torch.einsum("hqk,khd->qhd", torch.softmax(logits, dim=-1), value_ref).to(torch.bfloat16) + + assert bool(torch.isfinite(actual).all().item()) + torch.testing.assert_close(actual, expected, rtol=2.0e-2, atol=2.0e-2) + + +@_requires_gfx950 +@pytest.mark.parametrize( + ("head_dim", "value_head_dim"), + [(128, 128), (192, 128), (192, 192)], + ids=["d128-v128", "d192-v128", "d192-v192"], +) +def test_paged_fp8_graph_replay_matches_torch(head_dim, value_head_dim): + """The optimized paged variants preserve caller-owned graph output.""" + torch.manual_seed(29) + query_length = 63 + kv_length = 128 + num_pages = kv_length // 64 + + query, query_descale = quantize_per_tensor_fp8(torch.randn(query_length, 16, head_dim, device="cuda") * 0.2) + key, key_descale = quantize_per_tensor_fp8(torch.randn(num_pages, 1, head_dim // 16, 64, 16, device="cuda") * 0.2) + value, value_descale = quantize_per_tensor_fp8( + torch.randn(num_pages, 1, 4, value_head_dim, 16, device="cuda") * 0.2 + ) + output = torch.empty(query_length, 16, value_head_dim, device="cuda", dtype=torch.bfloat16) + block_table = torch.tensor([[1, 0]], device="cuda", dtype=torch.int32) + q_indptr = torch.tensor([0, query_length], device="cuda", dtype=torch.int32) + kv_indptr = torch.tensor([0, kv_length], device="cuda", dtype=torch.int32) + seqlen_k = torch.tensor([kv_length], device="cuda", dtype=torch.int32) + kwargs = dict( + causal=True, + num_kv_heads=1, + cu_seqlens_q=q_indptr, + cu_seqlens_kv=kv_indptr, + max_seqlen_q=query_length, + max_seqlen_kv=kv_length, + cross_seqlen=True, + block_table=block_table, + seqlen_k=seqlen_k, + kv_cache_layout="vectorized", + q_descale=query_descale, + k_descale=key_descale, + v_descale=value_descale, + out=output, + ) + + warmup_stream = torch.cuda.Stream() + with torch.cuda.stream(warmup_stream): + flydsl_flash_attn_func(query, key, value, **kwargs) + torch.cuda.current_stream().wait_stream(warmup_stream) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = flydsl_flash_attn_func(query, key, value, **kwargs) + graph.replay() + torch.cuda.synchronize() + + physical_pages = block_table[0].long() + query_ref = query.float() * query_descale + key_ref = ( + key[physical_pages].permute(0, 3, 1, 2, 4).reshape(-1, 1, head_dim)[:kv_length].float() * key_descale + ).expand(-1, 16, -1) + value_ref = ( + value[physical_pages].permute(0, 2, 4, 1, 3).reshape(-1, 1, value_head_dim)[:kv_length].float() * value_descale + ).expand(-1, 16, -1) + logits = torch.einsum("qhd,khd->hqk", query_ref, key_ref) / math.sqrt(head_dim) + query_positions = torch.arange(query_length, device="cuda") + key_positions = torch.arange(kv_length, device="cuda") + causal_mask = key_positions[None, :] <= (kv_length - query_length + query_positions)[:, None] + logits.masked_fill_(~causal_mask[None, :, :], float("-inf")) + expected = torch.einsum("hqk,khd->qhd", torch.softmax(logits, dim=-1), value_ref).to(torch.bfloat16) + + assert captured.data_ptr() == output.data_ptr() + assert bool(torch.isfinite(output).all().item()) + torch.testing.assert_close(output, expected, rtol=2.0e-2, atol=2.0e-2) + + +@_requires_gfx950 +def test_paged_fp8_bn128_launcher_rejects_unsupported_shapes(): + """Direct BN128 launches must preserve the single-sequence page-pair contract.""" + from kernels.attention.flash_attn_fp8_paged_gfx950 import build_flash_attn_paged_fp8_module + + launch = build_flash_attn_paged_fp8_module( + num_heads=16, + num_kv_heads=1, + head_dim=128, + value_head_dim=128, + causal=True, + dtype_str="fp8", + varlen=True, + cross_seqlen=True, + paged=True, + kv_cache_layout="vectorized", + paged_bn128=True, + ) + tensor = torch.empty(1, dtype=FP8_DTYPE, device="cuda") + output = torch.empty(1, dtype=torch.bfloat16, device="cuda") + block_table = torch.zeros(3, dtype=torch.int32, device="cuda") + + with pytest.raises(ValueError, match="batch_size=1"): + launch( + tensor, + tensor, + tensor, + output, + batch_size=2, + seq_len=64, + seq_len_kv=128, + block_table=block_table, + block_table_stride=2, + ) + with pytest.raises(ValueError, match="positive even number"): + launch( + tensor, + tensor, + tensor, + output, + batch_size=1, + seq_len=64, + seq_len_kv=192, + block_table=block_table, + block_table_stride=3, + ) + + @_requires_gfx950 @pytest.mark.parametrize("value_head_dim", [128, 192]) def test_paged_fp8_d192_rejects_flattened_int32_overflow(monkeypatch, value_head_dim): From f56ab1ff9c03577475ee801d13d5e3a79f0796f0 Mon Sep 17 00:00:00 2001 From: Xiake Sun Date: Tue, 1 Sep 2026 17:36:38 +0000 Subject: [PATCH 11/59] perf(fmha): relax paged D128 P*V accumulator anchors - Keep post-P*V accumulators in SSA for the paired-page path. - Reduce wait, MFMA, and packed-add stalls without increasing resources. - Improve D128 throughput across the 4K, 16K, and 64K targets. Signed-off-by: Xiake Sun --- kernels/attention/flash_attn_fp8_paged_gfx950.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/kernels/attention/flash_attn_fp8_paged_gfx950.py b/kernels/attention/flash_attn_fp8_paged_gfx950.py index c1d5017ebf..a650cc1da4 100644 --- a/kernels/attention/flash_attn_fp8_paged_gfx950.py +++ b/kernels/attention/flash_attn_fp8_paged_gfx950.py @@ -214,8 +214,9 @@ def _softmax_part(v_s, l_row, m_new): return v_p, l_row def _pv_part(v_p, v_v, v_o): - v_o = gemm_helper.pv(v_p, v_v, v_o) - return softmax_helper.anchor_v_o(v_o) + # Keep the post-MFMA accumulators in SSA. Pinning them after each + # subtile lengthens the paged schedule without reducing registers. + return gemm_helper.pv(v_p, v_v, v_o) def _subtile_tail(v_s, v_v, v_o, l_row, m_new): v_p, l_row = _softmax_part(v_s, l_row, m_new) From aa3126c237f0ea5596cd03ccec087758357dee03 Mon Sep 17 00:00:00 2001 From: Xiake Sun Date: Tue, 1 Sep 2026 17:55:18 +0000 Subject: [PATCH 12/59] perf(fmha): vectorize paged D128 V fragment loads - Load each 32-byte V fragment with two aligned 128-bit LDS operations. - Split vector loads into the four i64 FP8 MFMA K-steps. - Reduce V-fragment address and wait overhead without increasing resources. - Improve D128 throughput across the 4K, 16K, and 64K targets. Signed-off-by: Xiake Sun --- kernels/attention/flash_attn_utils.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/kernels/attention/flash_attn_utils.py b/kernels/attention/flash_attn_utils.py index 26d623f2d4..bf832fbac5 100644 --- a/kernels/attention/flash_attn_utils.py +++ b/kernels/attention/flash_attn_utils.py @@ -5780,13 +5780,15 @@ def _load_v_fp8_vectorized_bn128(self, buf_id): for dc in range_constexpr(traits.D_CHUNKS): d_row = fx.Index(dc * traits.D_CHUNK) + self.lane_mod_32 row_base = tile_base + d_row * fx.Index(row_stride) + token_base - for ks in range_constexpr(4): + for half in range_constexpr(2): ptr = buffer_ops.get_element_ptr( self.lds_vt_base_ptr, - byte_offset=fx.Int32(row_base + fx.Index(ks * 8) - self.lds_vt_base_idx), + byte_offset=fx.Int32(row_base + fx.Index(half * 16) - self.lds_vt_base_idx), elem_type=T.i8, ) - packs[ks][dc] = llvm.LoadOp(T.i64, ptr, alignment=8).result + pair = Vec(llvm.LoadOp(Vec.make_type(2, fx.Int64), ptr, alignment=16).result, (2,), fx.Int64) + packs[half * 2][dc] = as_mlir_value(fx.Int64(pair[0])) + packs[half * 2 + 1][dc] = as_mlir_value(fx.Int64(pair[1])) return packs def load_v_steps(self, buf_id, first_step, num_steps): From 3db264f9d1be1256c132b2a2df056192ae8408bb Mon Sep 17 00:00:00 2001 From: Xiake Sun Date: Tue, 1 Sep 2026 19:31:31 +0000 Subject: [PATCH 13/59] perf(fmha): use permlane16 for paged D128 V staging - Map each 16-lane row to one vectorized V token group. - Replace four DPP exchanges and selects with two permlane16 swaps. - Require a full EXEC mask through the static D128/V128 tile shape. - Improve D128 throughput across the 4K, 16K, and 64K targets. Signed-off-by: Xiake Sun --- kernels/attention/flash_attn_utils.py | 96 ++++++++++++--------------- 1 file changed, 44 insertions(+), 52 deletions(-) diff --git a/kernels/attention/flash_attn_utils.py b/kernels/attention/flash_attn_utils.py index bf832fbac5..20f8d6f4f1 100644 --- a/kernels/attention/flash_attn_utils.py +++ b/kernels/attention/flash_attn_utils.py @@ -5565,58 +5565,50 @@ def _stage_v_fp8_vectorized_bn128(self, tile_start, buf_id, page_id=None): aligned_base = ((self.lds_vt_base_idx + fx.Index(127)) // fx.Index(128)) * fx.Index(128) dst_base = aligned_base + fx.Index(buf_id * v_tile_bytes) token_groups16 = traits.BLOCK_N // traits.KV_VEC_SIZE - total_v16 = traits.V_HEAD_DIM * token_groups16 - passes = (total_v16 + traits.BLOCK_SIZE - 1) // traits.BLOCK_SIZE - - for pass_id in range_constexpr(passes): - flat = self.tid + fx.Index(pass_id * traits.BLOCK_SIZE) - - def _load_store_active(): - n_group16 = flat % fx.Index(token_groups16) - d_col = flat // fx.Index(token_groups16) - src_elem = ( - self.kv_head_idx * token_groups16 * traits.V_HEAD_DIM * traits.KV_VEC_SIZE - + n_group16 * traits.V_HEAD_DIM * traits.KV_VEC_SIZE - + d_col * traits.KV_VEC_SIZE - ) - src_i32x4 = self.buffer_load_fp8x16(src_div, src_elem) - src_words = Vec(src_i32x4, (4,), fx.Int32) - # gfx9 DPP quad_perm [1,0,3,2] is lane^1 without using the LDS - # cross-lane pipe, which remains busy serving the P*V fragments. - peer_words = [ - fx.Int32( - rocdl.update_dpp( - T.i32, - as_mlir_value(fx.Int32(src_words[i])), - as_mlir_value(fx.Int32(src_words[i])), - 0xB1, - 0xF, - 0xF, - True, - ) - ) - for i in range_constexpr(4) - ] - odd_group = (fx.Int32(n_group16) & fx.Int32(1)) != fx.Int32(0) - reordered = Vec.from_elements( - [ - odd_group.select(peer_words[1], src_words[0]), - odd_group.select(peer_words[3], src_words[2]), - odd_group.select(src_words[1], peer_words[0]), - odd_group.select(src_words[3], peer_words[2]), - ], - fx.Int32, - ) - dest_group16 = (n_group16 % fx.Index(2)) * fx.Index(2) + n_group16 // fx.Index(2) - dst_byte = dst_base + d_col * fx.Index(row_stride) + dest_group16 * fx.Index(16) - ptr = buffer_ops.get_element_ptr( - self.lds_vt_base_ptr, - byte_offset=fx.Int32(dst_byte - self.lds_vt_base_idx), - elem_type=T.i8, - ) - llvm.StoreOp(as_mlir_value(reordered), ptr, alignment=16) - - scf_if_dispatch(flat < fx.Index(total_v16), _load_store_active) + # permlane16 requires a full EXEC mask. D128/V128 has exactly one + # 16-byte V vector per CTA thread, so no per-lane bounds branch is needed. + assert traits.V_HEAD_DIM * token_groups16 == traits.BLOCK_SIZE + n_group16 = self.lane_in_warp // fx.Index(16) + d_col = self.wave_id * fx.Index(16) + self.lane_in_warp % fx.Index(16) + src_elem = ( + self.kv_head_idx * token_groups16 * traits.V_HEAD_DIM * traits.KV_VEC_SIZE + + n_group16 * traits.V_HEAD_DIM * traits.KV_VEC_SIZE + + d_col * traits.KV_VEC_SIZE + ) + src_i32x4 = self.buffer_load_fp8x16(src_div, src_elem) + src_words = Vec(src_i32x4, (4,), fx.Int32) + pair_ty = ir.Type.parse("!llvm.struct<(i32, i32)>") + pair_lo = rocdl.permlane16_swap( + pair_ty, + as_mlir_value(fx.Int32(src_words[0])), + as_mlir_value(fx.Int32(src_words[1])), + False, + False, + ) + pair_hi = rocdl.permlane16_swap( + pair_ty, + as_mlir_value(fx.Int32(src_words[2])), + as_mlir_value(fx.Int32(src_words[3])), + False, + False, + ) + reordered = Vec.from_elements( + [ + fx.Int32(llvm.extractvalue(T.i32, pair_lo, [0])), + fx.Int32(llvm.extractvalue(T.i32, pair_hi, [0])), + fx.Int32(llvm.extractvalue(T.i32, pair_lo, [1])), + fx.Int32(llvm.extractvalue(T.i32, pair_hi, [1])), + ], + fx.Int32, + ) + dest_group16 = (n_group16 % fx.Index(2)) * fx.Index(2) + n_group16 // fx.Index(2) + dst_byte = dst_base + d_col * fx.Index(row_stride) + dest_group16 * fx.Index(16) + ptr = buffer_ops.get_element_ptr( + self.lds_vt_base_ptr, + byte_offset=fx.Int32(dst_byte - self.lds_vt_base_idx), + elem_type=T.i8, + ) + llvm.StoreOp(as_mlir_value(reordered), ptr, alignment=16) def _stage_v_fp8_vectorized_dense_layout(self, tile_start, buf_id, page_id=None): """Transpose vectorized V into the LDS layout used by direct FP8 P*V.""" From c554b1d4eb5b90cf061746357be6e93fe5ebc4db Mon Sep 17 00:00:00 2001 From: Xiake Sun Date: Tue, 1 Sep 2026 19:51:00 +0000 Subject: [PATCH 14/59] perf(fmha): keep paged D128 permlane results in native order - Store permlane16 swap results without a staging-register reshuffle. - Apply the matching fixed word permutation to packed P. - Preserve the FP8 contraction mapping and full-EXEC requirement. - Improve D128 throughput at the 16K and 64K targets. Signed-off-by: Xiake Sun --- kernels/attention/flash_attn_utils.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/kernels/attention/flash_attn_utils.py b/kernels/attention/flash_attn_utils.py index 20f8d6f4f1..f99e3da167 100644 --- a/kernels/attention/flash_attn_utils.py +++ b/kernels/attention/flash_attn_utils.py @@ -5318,7 +5318,14 @@ def cast_p_fp8_direct(self, v_p): for pks in range_constexpr(self.traits.PV_K_STEPS): p_base = pks * 8 f32 += [hi_full[p_base + s] for s in range_constexpr(8)] - return self._pack_fp8_i32x8(f32) + packed = self._pack_fp8_i32x8(f32) + if const_expr(self.traits.PAGED and self.traits.FP8_PV_DIRECT): + words = Vec(packed, (8,), fx.Int32) + return Vec.from_elements( + [words[i] for i in (0, 2, 1, 3, 4, 6, 5, 7)], + fx.Int32, + ).ir_value() + return packed def _pv_fp8_direct(self, p_fp8, v_v, v_o): v_o = _anchor_v_o(self.traits, v_o) @@ -5595,8 +5602,8 @@ def _stage_v_fp8_vectorized_bn128(self, tile_start, buf_id, page_id=None): reordered = Vec.from_elements( [ fx.Int32(llvm.extractvalue(T.i32, pair_lo, [0])), - fx.Int32(llvm.extractvalue(T.i32, pair_hi, [0])), fx.Int32(llvm.extractvalue(T.i32, pair_lo, [1])), + fx.Int32(llvm.extractvalue(T.i32, pair_hi, [0])), fx.Int32(llvm.extractvalue(T.i32, pair_hi, [1])), ], fx.Int32, From b48cb40504d79d171b882377ae3c99dcb9e1c6d8 Mon Sep 17 00:00:00 2001 From: Xiake Sun Date: Tue, 1 Sep 2026 20:17:15 +0000 Subject: [PATCH 15/59] perf(fmha): overlap paged D128 V prefetch - Split vectorized V loads from permutation and LDS staging. - Hide V4 latency behind K prefetch and V5 latency behind QK MFMAs. - Keep one prefetched V vector live to preserve register and scratch usage. - Improve random-input D128 latency across the 4K, 16K, and 64K targets. Signed-off-by: Xiake Sun --- .../attention/flash_attn_fp8_paged_gfx950.py | 8 +++++-- kernels/attention/flash_attn_utils.py | 21 ++++++++++++++----- 2 files changed, 22 insertions(+), 7 deletions(-) diff --git a/kernels/attention/flash_attn_fp8_paged_gfx950.py b/kernels/attention/flash_attn_fp8_paged_gfx950.py index a650cc1da4..8a3a77624f 100644 --- a/kernels/attention/flash_attn_fp8_paged_gfx950.py +++ b/kernels/attention/flash_attn_fp8_paged_gfx950.py @@ -294,13 +294,17 @@ def _ring_wrap(x): v_v_a = kv_lds_to_regs.load_v(a_buf) page_f_a, page_f_b = ctx.load_page_id_pair((j + fx.Index(4)) * BN) + # Keep only one prefetched V vector live at a time. K4/K5 address + # generation hides V4 latency, while the two QK MFMAs hide V5. + v_f_a = kv_gmem_to_lds._load_v_fp8_vectorized_bn128((j + fx.Index(4)) * BN, page_id=page_f_a) kv_gmem_to_lds.load_k((j + fx.Index(4)) * BN, f_a_buf, page_id=page_f_a) kv_gmem_to_lds.load_k((j + fx.Index(5)) * BN, f_b_buf, page_id=page_f_b) - kv_gmem_to_lds.load_v((j + fx.Index(4)) * BN, f_a_buf, page_id=page_f_a) - kv_gmem_to_lds.load_v((j + fx.Index(5)) * BN, f_b_buf, page_id=page_f_b) + kv_gmem_to_lds._store_v_fp8_vectorized_bn128(v_f_a, f_a_buf) + v_f_b = kv_gmem_to_lds._load_v_fp8_vectorized_bn128((j + fx.Index(5)) * BN, page_id=page_f_b) v_s_a = gemm_helper.qk(v_k_a, q_wide) v_s_b = gemm_helper.qk(v_k_b, q_wide) + kv_gmem_to_lds._store_v_fp8_vectorized_bn128(v_f_b, f_b_buf) v_s_a, v_s_b = _mask_pair(v_s_a, v_s_b, j) m_tile = _merge_tile_max(v_s_a, v_s_b) v_o, m_new, l_row = _correct_o(v_o, m_row, l_row, m_tile) diff --git a/kernels/attention/flash_attn_utils.py b/kernels/attention/flash_attn_utils.py index f99e3da167..9d8d77492a 100644 --- a/kernels/attention/flash_attn_utils.py +++ b/kernels/attention/flash_attn_utils.py @@ -5563,14 +5563,15 @@ def _load_store_active(): def _stage_v_fp8_vectorized_bn128(self, tile_start, buf_id, page_id=None): """Stage V in a bank-padded D-major layout for paged BN128 P*V.""" + src_i32x4 = self._load_v_fp8_vectorized_bn128(tile_start, page_id=page_id) + self._store_v_fp8_vectorized_bn128(src_i32x4, buf_id) + + def _load_v_fp8_vectorized_bn128(self, tile_start, page_id=None): + """Issue one vectorized paged-V load for the BN128 direct-FP8 path.""" traits = self.traits if page_id is None: page_id = self.load_page_id(tile_start) src_div = self.make_page_view(self.v_base_iter, page_id, is_value=True) - row_stride = traits.FP8_V_ROW_STRIDE - v_tile_bytes = traits.V_HEAD_DIM * row_stride - aligned_base = ((self.lds_vt_base_idx + fx.Index(127)) // fx.Index(128)) * fx.Index(128) - dst_base = aligned_base + fx.Index(buf_id * v_tile_bytes) token_groups16 = traits.BLOCK_N // traits.KV_VEC_SIZE # permlane16 requires a full EXEC mask. D128/V128 has exactly one # 16-byte V vector per CTA thread, so no per-lane bounds branch is needed. @@ -5582,7 +5583,17 @@ def _stage_v_fp8_vectorized_bn128(self, tile_start, buf_id, page_id=None): + n_group16 * traits.V_HEAD_DIM * traits.KV_VEC_SIZE + d_col * traits.KV_VEC_SIZE ) - src_i32x4 = self.buffer_load_fp8x16(src_div, src_elem) + return self.buffer_load_fp8x16(src_div, src_elem) + + def _store_v_fp8_vectorized_bn128(self, src_i32x4, buf_id): + """Permute a prefetched BN128 V vector and store it into LDS.""" + traits = self.traits + row_stride = traits.FP8_V_ROW_STRIDE + v_tile_bytes = traits.V_HEAD_DIM * row_stride + aligned_base = ((self.lds_vt_base_idx + fx.Index(127)) // fx.Index(128)) * fx.Index(128) + dst_base = aligned_base + fx.Index(buf_id * v_tile_bytes) + n_group16 = self.lane_in_warp // fx.Index(16) + d_col = self.wave_id * fx.Index(16) + self.lane_in_warp % fx.Index(16) src_words = Vec(src_i32x4, (4,), fx.Int32) pair_ty = ir.Type.parse("!llvm.struct<(i32, i32)>") pair_lo = rocdl.permlane16_swap( From fe33810dfb111ebf612e13bb1b764f5bf86e45e0 Mon Sep 17 00:00:00 2001 From: Xiake Sun Date: Tue, 1 Sep 2026 20:31:29 +0000 Subject: [PATCH 16/59] perf(fmha): defer paged D128 logit scaling - Keep paged BN128 QK accumulators unscaled through row-max reduction. - Fold the runtime logit scale into score subtraction and max correction. - Remove eager per-score multiplies and reduce VGPR use from 246 to 218. - Reach 92.57% of dense throughput at the 64K/128K target. Signed-off-by: Xiake Sun --- kernels/attention/flash_attn_utils.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/kernels/attention/flash_attn_utils.py b/kernels/attention/flash_attn_utils.py index 9d8d77492a..6be653dfa4 100644 --- a/kernels/attention/flash_attn_utils.py +++ b/kernels/attention/flash_attn_utils.py @@ -5303,7 +5303,12 @@ def qk(self, v_k, q_wide=None): rocdl.sched_group_barrier(traits.SCHED_MFMA_MASK, 1, 12) rocdl.sched_group_barrier(traits.SCHED_DS_READ_MASK, n_ds // 2, 12) rocdl.sched_group_barrier(traits.SCHED_MFMA_MASK, n_mfma - 1, 12) - if const_expr(traits.PAGED): + # Generic paged D192 keeps logits in scaled log2 units because its + # HIPREC softmax path was built around that representation. The + # paired D128 path follows dense BN128 instead: defer scaling to the + # score-minus-max FMA so it does not materialize a separate multiply + # for every QK accumulator element. + if const_expr(traits.PAGED and not traits.BN128): scale_vec = Vec.from_elements([self.c_logit_scale], fx.Float32).broadcast_to(16) v_s_lo = as_mlir_value(Vec(v_s_lo) * scale_vec) v_s_hi = as_mlir_value(Vec(v_s_hi) * scale_vec) @@ -5984,7 +5989,7 @@ def floor_masked_max(self, row_max): _P_HEADROOM_LOG2 = 8.807354922057604 def sub_m(self, v_s, row_max): - if const_expr(self.traits.PAGED): + if const_expr(self.traits.PAGED and not self.traits.BN128): return _sub_score_pair(v_s, row_max, self.fm_fast) # P is cast to e4m3, whose smallest subnormal is 2**-9, so a softmax # over thousands of keys loses its tail to flush-to-zero -- while l_row, @@ -6041,7 +6046,7 @@ def safe_l_inv(self, l_row): return _safe_l_inv(l_row, self.c_zero_f) def rescale_from_tile_max(self, m_row, m_tile_max): - if const_expr(self.traits.PAGED): + if const_expr(self.traits.PAGED and not self.traits.BN128): return _rescale_from_tile_max(m_row, m_tile_max, self.fm_fast) row_max = fx.maxnumf(m_row, m_tile_max) diff_scaled = (m_row - row_max) * self.c_logit_scale From 88ff0b9d175be9b28871578318ae0a7f4198132f Mon Sep 17 00:00:00 2001 From: Xiake Sun Date: Tue, 1 Sep 2026 20:37:47 +0000 Subject: [PATCH 17/59] perf(fmha): overlap paged D128 V loads with QK - Pair each future V load with an independent QK contraction. - Keep only one prefetched V vector live at a time. - Preserve 218-VGPR, zero-scratch kernel resource usage. - Improve the 16K and 64K random-input attention targets. Signed-off-by: Xiake Sun --- kernels/attention/flash_attn_fp8_paged_gfx950.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/kernels/attention/flash_attn_fp8_paged_gfx950.py b/kernels/attention/flash_attn_fp8_paged_gfx950.py index 8a3a77624f..283fefb1ab 100644 --- a/kernels/attention/flash_attn_fp8_paged_gfx950.py +++ b/kernels/attention/flash_attn_fp8_paged_gfx950.py @@ -294,15 +294,15 @@ def _ring_wrap(x): v_v_a = kv_lds_to_regs.load_v(a_buf) page_f_a, page_f_b = ctx.load_page_id_pair((j + fx.Index(4)) * BN) - # Keep only one prefetched V vector live at a time. K4/K5 address - # generation hides V4 latency, while the two QK MFMAs hide V5. + # Keep only one prefetched V vector live at a time and overlap each + # load with one of the two independent QK contractions. v_f_a = kv_gmem_to_lds._load_v_fp8_vectorized_bn128((j + fx.Index(4)) * BN, page_id=page_f_a) kv_gmem_to_lds.load_k((j + fx.Index(4)) * BN, f_a_buf, page_id=page_f_a) kv_gmem_to_lds.load_k((j + fx.Index(5)) * BN, f_b_buf, page_id=page_f_b) - kv_gmem_to_lds._store_v_fp8_vectorized_bn128(v_f_a, f_a_buf) - v_f_b = kv_gmem_to_lds._load_v_fp8_vectorized_bn128((j + fx.Index(5)) * BN, page_id=page_f_b) v_s_a = gemm_helper.qk(v_k_a, q_wide) + kv_gmem_to_lds._store_v_fp8_vectorized_bn128(v_f_a, f_a_buf) + v_f_b = kv_gmem_to_lds._load_v_fp8_vectorized_bn128((j + fx.Index(5)) * BN, page_id=page_f_b) v_s_b = gemm_helper.qk(v_k_b, q_wide) kv_gmem_to_lds._store_v_fp8_vectorized_bn128(v_f_b, f_b_buf) v_s_a, v_s_b = _mask_pair(v_s_a, v_s_b, j) From 9232c1729b551fd44d10c3d470c58864e918957e Mon Sep 17 00:00:00 2001 From: Xiake Sun Date: Tue, 1 Sep 2026 20:45:58 +0000 Subject: [PATCH 18/59] perf(fmha): defer paged D192 logit scaling - Keep generic paged QK accumulators in their raw domain. - Fold logit scaling into score subtraction and max correction. - Compare lazy-rescale thresholds in scaled log2 units. - Improve D192 steady-state latency and reduce V192 scratch usage. Signed-off-by: Xiake Sun --- kernels/attention/flash_attn_utils.py | 16 +++------------- 1 file changed, 3 insertions(+), 13 deletions(-) diff --git a/kernels/attention/flash_attn_utils.py b/kernels/attention/flash_attn_utils.py index 6be653dfa4..832f64b202 100644 --- a/kernels/attention/flash_attn_utils.py +++ b/kernels/attention/flash_attn_utils.py @@ -5303,15 +5303,6 @@ def qk(self, v_k, q_wide=None): rocdl.sched_group_barrier(traits.SCHED_MFMA_MASK, 1, 12) rocdl.sched_group_barrier(traits.SCHED_DS_READ_MASK, n_ds // 2, 12) rocdl.sched_group_barrier(traits.SCHED_MFMA_MASK, n_mfma - 1, 12) - # Generic paged D192 keeps logits in scaled log2 units because its - # HIPREC softmax path was built around that representation. The - # paired D128 path follows dense BN128 instead: defer scaling to the - # score-minus-max FMA so it does not materialize a separate multiply - # for every QK accumulator element. - if const_expr(traits.PAGED and not traits.BN128): - scale_vec = Vec.from_elements([self.c_logit_scale], fx.Float32).broadcast_to(16) - v_s_lo = as_mlir_value(Vec(v_s_lo) * scale_vec) - v_s_hi = as_mlir_value(Vec(v_s_hi) * scale_vec) return (v_s_lo, v_s_hi) def cast_p_fp8_direct(self, v_p): @@ -5990,7 +5981,7 @@ def floor_masked_max(self, row_max): def sub_m(self, v_s, row_max): if const_expr(self.traits.PAGED and not self.traits.BN128): - return _sub_score_pair(v_s, row_max, self.fm_fast) + return _scale_sub_score_pair(v_s, row_max, self.c_logit_scale, self.c_zero_f, self.fm_fast) # P is cast to e4m3, whose smallest subnormal is 2**-9, so a softmax # over thousands of keys loses its tail to flush-to-zero -- while l_row, # summed before the cast, still counts it. Scaling P up first uses the @@ -6046,8 +6037,6 @@ def safe_l_inv(self, l_row): return _safe_l_inv(l_row, self.c_zero_f) def rescale_from_tile_max(self, m_row, m_tile_max): - if const_expr(self.traits.PAGED and not self.traits.BN128): - return _rescale_from_tile_max(m_row, m_tile_max, self.fm_fast) row_max = fx.maxnumf(m_row, m_tile_max) diff_scaled = (m_row - row_max) * self.c_logit_scale rescale = rocdl.exp2(T.f32, as_mlir_value(diff_scaled)) @@ -6114,7 +6103,8 @@ def _lazy_rescale_o_paged(self, v_o, m_row, l_row, m_tile_max, v_p): @flyc.jit def _run(v_o, m_row, l_row, m_tile_max, v_p): m_diff = m_tile_max - m_row - below = fx.Float32(m_diff) <= self.c_rescale_thr_f + m_diff_scaled = m_diff * self.c_logit_scale + below = fx.Float32(m_diff_scaled) <= self.c_rescale_thr_f ballot = rocdl.ballot(T.i64, as_mlir_value(below)) all_below = arith.cmpi(arith.CmpIPredicate.eq, as_mlir_value(ballot), _read_exec_i64()) all_below = llvm.intr_expect(all_below, arith.constant(1, type=ir.IntegerType.get_signless(1))) From 7f4f70949d73cc70335e867b49054bcd0744995c Mon Sep 17 00:00:00 2001 From: Xiake Sun Date: Tue, 1 Sep 2026 20:58:59 +0000 Subject: [PATCH 19/59] perf(fmha): use direct FP8 P*V for paged V128 - Add an explicit bank-padded FP8 V-layout trait. - Reuse vectorized V staging for D192/V128 and generic D128 fallback. - Pack softmax probabilities in the matching MFMA word order. - Improve D192/V128 latency while preserving accuracy and zero scratch. Signed-off-by: Xiake Sun --- .../attention/flash_attn_fp8_paged_gfx950.py | 8 +-- kernels/attention/flash_attn_utils.py | 57 ++++++++++++------- 2 files changed, 39 insertions(+), 26 deletions(-) diff --git a/kernels/attention/flash_attn_fp8_paged_gfx950.py b/kernels/attention/flash_attn_fp8_paged_gfx950.py index 283fefb1ab..956c7efbd1 100644 --- a/kernels/attention/flash_attn_fp8_paged_gfx950.py +++ b/kernels/attention/flash_attn_fp8_paged_gfx950.py @@ -296,15 +296,15 @@ def _ring_wrap(x): page_f_a, page_f_b = ctx.load_page_id_pair((j + fx.Index(4)) * BN) # Keep only one prefetched V vector live at a time and overlap each # load with one of the two independent QK contractions. - v_f_a = kv_gmem_to_lds._load_v_fp8_vectorized_bn128((j + fx.Index(4)) * BN, page_id=page_f_a) + v_f_a = kv_gmem_to_lds._load_v_fp8_vectorized_bankpad_source((j + fx.Index(4)) * BN, page_id=page_f_a) kv_gmem_to_lds.load_k((j + fx.Index(4)) * BN, f_a_buf, page_id=page_f_a) kv_gmem_to_lds.load_k((j + fx.Index(5)) * BN, f_b_buf, page_id=page_f_b) v_s_a = gemm_helper.qk(v_k_a, q_wide) - kv_gmem_to_lds._store_v_fp8_vectorized_bn128(v_f_a, f_a_buf) - v_f_b = kv_gmem_to_lds._load_v_fp8_vectorized_bn128((j + fx.Index(5)) * BN, page_id=page_f_b) + kv_gmem_to_lds._store_v_fp8_vectorized_bankpad(v_f_a, f_a_buf) + v_f_b = kv_gmem_to_lds._load_v_fp8_vectorized_bankpad_source((j + fx.Index(5)) * BN, page_id=page_f_b) v_s_b = gemm_helper.qk(v_k_b, q_wide) - kv_gmem_to_lds._store_v_fp8_vectorized_bn128(v_f_b, f_b_buf) + kv_gmem_to_lds._store_v_fp8_vectorized_bankpad(v_f_b, f_b_buf) v_s_a, v_s_b = _mask_pair(v_s_a, v_s_b, j) m_tile = _merge_tile_max(v_s_a, v_s_b) v_o, m_new, l_row = _correct_o(v_o, m_row, l_row, m_tile) diff --git a/kernels/attention/flash_attn_utils.py b/kernels/attention/flash_attn_utils.py index 832f64b202..39b8f04574 100644 --- a/kernels/attention/flash_attn_utils.py +++ b/kernels/attention/flash_attn_utils.py @@ -2096,6 +2096,7 @@ class PagedDualwaveSwpFp8Traits: CROSS_SEQLEN: bool FP8_PV: bool FP8_PV_DIRECT: bool + FP8_V_BANKPAD: bool BN128: bool BN128_PF: bool QREG: bool @@ -2192,6 +2193,7 @@ def cache_tag(self): self.FP8_V_ROW_STRIDE, self.FP8_PV, self.FP8_PV_DIRECT, + self.FP8_V_BANKPAD, self.NUM_PREFETCH_K, self.BN128, self.BN128_PF, @@ -2285,23 +2287,26 @@ def _make_paged_dualwave_swp_fp8_traits( sdrpt_bf = value_head_dim // d128_bf vls_bf = slw_bf + 64 // eb_bf vt_bf16_elems = snrpt_bf * sdrpt_bf * vls_bf + paged = bool(paged) + fp8_v_bankpad = paged and value_head_dim == 128 fp8_v_row_stride = block_n + 16 fp8_v_tile_bytes = (block_n // 8) * (head_dim // 16) * 128 - if paged and bn128 and head_dim == value_head_dim == 128: + if fp8_v_bankpad: fp8_v_tile_bytes = value_head_dim * fp8_v_row_stride - if bn128_pf: + if fp8_v_bankpad: + vt_bf16_total = num_prefetch_k * (fp8_v_tile_bytes // eb_bf) + 128 + elif bn128_pf: vt_bf16_total = num_prefetch_k * (fp8_v_tile_bytes // eb_bf) + 128 else: vt_bf16_total = (2 if deep_ring else num_prefetch_k) * vt_bf16_elems splitk = num_kv_splits > 1 - paged = bool(paged) kv_vectorized = paged and kv_cache_layout == "vectorized" page_size = block_n - # D128 paged attention can consume its native FP8 V cache directly once - # the vectorized cache is transposed into the existing FP8 PV LDS layout. - fp8_pv = (paged and head_dim == value_head_dim == 128) or (not paged and os.getenv("FLYDSL_FA_FP8_PV", "0") == "1") + # Paged V128 attention can consume its native FP8 V cache directly once + # the vectorized cache is transposed into the bank-padded FP8 P*V layout. + fp8_pv = (paged and value_head_dim == 128) or (not paged and os.getenv("FLYDSL_FA_FP8_PV", "0") == "1") fp8_pv_direct = bn128 and (not paged or fp8_pv) if fp8_pv_direct: fp8_pv = True @@ -2336,6 +2341,7 @@ def _make_paged_dualwave_swp_fp8_traits( CROSS_SEQLEN=bool(cross_seqlen), FP8_PV=fp8_pv, FP8_PV_DIRECT=bool(fp8_pv_direct), + FP8_V_BANKPAD=bool(fp8_v_bankpad), BN128=bool(bn128), BN128_PF=bool(bn128_pf), QREG=bool(qreg), @@ -5185,7 +5191,14 @@ def _p_to_fp8_i32x8(self, v_p): f32 = [] for pk in (p_lo[0], p_lo[1], p_hi[0], p_hi[1]): f32 += self._v8bf16_to_f32(pk) - return self._pack_fp8_i32x8(f32) + packed = self._pack_fp8_i32x8(f32) + if const_expr(self.traits.PAGED and self.traits.FP8_V_BANKPAD): + words = Vec(packed, (8,), fx.Int32) + return Vec.from_elements( + [words[i] for i in (0, 2, 1, 3, 4, 6, 5, 7)], + fx.Int32, + ).ir_value() + return packed def _v_to_fp8_i32x8(self, v_v, dc): f32 = [] @@ -5408,8 +5421,8 @@ def _run(): _run() def load_v(self, tile_start, buf_id, page_id=None): - if const_expr(self.traits.KV_VECTORIZED and self.traits.FP8_PV_DIRECT): - self._stage_v_fp8_vectorized_bn128(tile_start, buf_id, page_id=page_id) + if const_expr(self.traits.KV_VECTORIZED and self.traits.FP8_V_BANKPAD): + self._stage_v_fp8_vectorized_bankpad(tile_start, buf_id, page_id=page_id) elif const_expr(self.traits.KV_VECTORIZED and self.traits.FP8_PV): self._stage_v_fp8_vectorized_dense_layout(tile_start, buf_id, page_id=page_id) elif const_expr(self.traits.KV_VECTORIZED): @@ -5557,20 +5570,20 @@ def _load_store_active(): scf_if_dispatch(flat < fx.Index(total_v16), _load_store_active) - def _stage_v_fp8_vectorized_bn128(self, tile_start, buf_id, page_id=None): - """Stage V in a bank-padded D-major layout for paged BN128 P*V.""" - src_i32x4 = self._load_v_fp8_vectorized_bn128(tile_start, page_id=page_id) - self._store_v_fp8_vectorized_bn128(src_i32x4, buf_id) + def _stage_v_fp8_vectorized_bankpad(self, tile_start, buf_id, page_id=None): + """Stage vectorized paged V in the bank-padded direct-FP8 layout.""" + src_i32x4 = self._load_v_fp8_vectorized_bankpad_source(tile_start, page_id=page_id) + self._store_v_fp8_vectorized_bankpad(src_i32x4, buf_id) - def _load_v_fp8_vectorized_bn128(self, tile_start, page_id=None): - """Issue one vectorized paged-V load for the BN128 direct-FP8 path.""" + def _load_v_fp8_vectorized_bankpad_source(self, tile_start, page_id=None): + """Issue one vectorized paged-V load for the bank-padded path.""" traits = self.traits if page_id is None: page_id = self.load_page_id(tile_start) src_div = self.make_page_view(self.v_base_iter, page_id, is_value=True) token_groups16 = traits.BLOCK_N // traits.KV_VEC_SIZE - # permlane16 requires a full EXEC mask. D128/V128 has exactly one - # 16-byte V vector per CTA thread, so no per-lane bounds branch is needed. + # permlane16 requires a full EXEC mask. V128 has exactly one vector per + # CTA thread, so no per-lane bounds branch is needed. assert traits.V_HEAD_DIM * token_groups16 == traits.BLOCK_SIZE n_group16 = self.lane_in_warp // fx.Index(16) d_col = self.wave_id * fx.Index(16) + self.lane_in_warp % fx.Index(16) @@ -5581,8 +5594,8 @@ def _load_v_fp8_vectorized_bn128(self, tile_start, page_id=None): ) return self.buffer_load_fp8x16(src_div, src_elem) - def _store_v_fp8_vectorized_bn128(self, src_i32x4, buf_id): - """Permute a prefetched BN128 V vector and store it into LDS.""" + def _store_v_fp8_vectorized_bankpad(self, src_i32x4, buf_id): + """Permute a prefetched V vector and store it into bank-padded LDS.""" traits = self.traits row_stride = traits.FP8_V_ROW_STRIDE v_tile_bytes = traits.V_HEAD_DIM * row_stride @@ -5749,8 +5762,8 @@ def _read_strip(key): def load_v(self, buf_id): if const_expr(self.traits.FP8_PV): - if const_expr(self.traits.KV_VECTORIZED and self.traits.FP8_PV_DIRECT): - return self._load_v_fp8_vectorized_bn128(buf_id) + if const_expr(self.traits.KV_VECTORIZED and self.traits.FP8_V_BANKPAD): + return self._load_v_fp8_vectorized_bankpad(buf_id) return self._load_v_fp8_block(buf_id) # Read all V packs from the bf16 vt scratch for buffer `buf_id`. traits = self.traits @@ -5774,7 +5787,7 @@ def load_v(self, buf_id): packs[k_substep][dc] = Vec(a).shuffle(Vec(b), [0, 1, 2, 3, 4, 5, 6, 7]).ir_value() return packs - def _load_v_fp8_vectorized_bn128(self, buf_id): + def _load_v_fp8_vectorized_bankpad(self, buf_id): """Read the padded D-major V tile into the FP8 MFMA A fragment.""" traits = self.traits row_stride = traits.FP8_V_ROW_STRIDE From 7c02fdb5bf33a7b1e7bf6b69dcd2da0c31f2308a Mon Sep 17 00:00:00 2001 From: Xiake Sun Date: Wed, 2 Sep 2026 06:02:20 +0000 Subject: [PATCH 20/59] perf(fmha): cache compiled fp8 page attention - Reuse the compiled paged launcher instead of resolving the FlyDSL JIT cache key on every invocation. - Cache per-device architecture detection and reuse it during paged dispatch. - Avoid redundant current-stream context setup while preserving explicit-stream ordering. - Reduce 4K host submission overhead to reach over 90% of D128 throughput compare with dense attention Signed-off-by: Xiake Sun --- .../attention/flash_attn_fp8_paged_gfx950.py | 28 +++---------------- kernels/attention/flash_attn_interface.py | 17 +++++++---- 2 files changed, 15 insertions(+), 30 deletions(-) diff --git a/kernels/attention/flash_attn_fp8_paged_gfx950.py b/kernels/attention/flash_attn_fp8_paged_gfx950.py index 956c7efbd1..1ff8d55741 100644 --- a/kernels/attention/flash_attn_fp8_paged_gfx950.py +++ b/kernels/attention/flash_attn_fp8_paged_gfx950.py @@ -30,6 +30,7 @@ dualwave_splitk_workspace_elems, # noqa: F401 ) from kernels.common.kernels_common import dtype_to_elem_type +from kernels.common.tensor_shim import _run_compiled def build_flash_attn_paged_fp8_module( @@ -1161,29 +1162,8 @@ def _launch( if v_descale is None: v_descale = O with CompilationContext.compile_hints(_dualwave_swp_compile_hints): - if stream is None: - return launch_flash_attn_dualwave_swp( - Q, - K, - V, - O, - debug_counts, - cu_seqlens_q, - cu_seqlens_kv, - block_table, - block_table_stride, - q_descale, - k_descale, - v_descale, - batch_size, - seq_len, - seq_len_kv, - stride_q_n, - stride_o_n, - stride_kv_n, - head_dim_runtime, - ) - return launch_flash_attn_dualwave_swp( + return _run_compiled( + launch_flash_attn_dualwave_swp, Q, K, V, @@ -1203,7 +1183,7 @@ def _launch( stride_o_n, stride_kv_n, head_dim_runtime, - stream=stream, + fx.Stream(stream), ) def _compile( diff --git a/kernels/attention/flash_attn_interface.py b/kernels/attention/flash_attn_interface.py index 6ba4e76ca5..daf804c0a8 100644 --- a/kernels/attention/flash_attn_interface.py +++ b/kernels/attention/flash_attn_interface.py @@ -22,6 +22,7 @@ from __future__ import annotations +import contextlib import functools from typing import Optional @@ -127,6 +128,7 @@ def _dtype_str(t: torch.Tensor) -> str: return s +@functools.lru_cache(maxsize=16) def _gpu_arch(device: torch.device) -> str: try: return torch.cuda.get_device_properties(device.index).gcnArchName.split(":")[0] @@ -566,6 +568,7 @@ def _flydsl_flash_attn_paged( dtype_str = _dtype_str(q) paged_fp8 = dtype_str == "fp8" + arch = _gpu_arch(q.device) varlen = cu_seqlens_q is not None if varlen: # Packed varlen Q: [total_q, H, D]. Per-batch ranges come from cu_seqlens @@ -614,9 +617,8 @@ def _flydsl_flash_attn_paged( if k_head_dim != D: raise ValueError(f"flydsl_flash_attn_func: paged K head_dim ({k_head_dim}) must match q head_dim ({D})") if paged_fp8: - _arch = _gpu_arch(q.device) - if not _arch.startswith("gfx950"): - raise ValueError(f"flydsl_flash_attn_func: paged FP8 requires gfx950, got '{_arch or 'unknown'}'") + if not arch.startswith("gfx950"): + raise ValueError(f"flydsl_flash_attn_func: paged FP8 requires gfx950, got '{arch or 'unknown'}'") fp8_head_dims = (D, value_head_dim) if not ( causal @@ -743,13 +745,12 @@ def _flydsl_flash_attn_paged( with torch.cuda.device(q.device.index): launch_stream = torch.cuda.current_stream(q.device) if stream is None else stream # Short paged attention uses generic light; unsupported cases stay on dualwave. - _arch = _gpu_arch(q.device) _paged_light_ok = ( (num_kv_splits <= 1) and bias is None # the light paged kernel has no bias path and D in (64, 128) and dtype_str in ("bf16", "f16") - and (not _arch.startswith("gfx950") or Sq <= _VARLEN_LIGHT_MAX_SEQ) + and (not arch.startswith("gfx950") or Sq <= _VARLEN_LIGHT_MAX_SEQ) ) if paged_fp8: num_kv_pages = (skv + page_size - 1) // page_size @@ -803,7 +804,11 @@ def _flydsl_flash_attn_paged( kv_cache_layout=kv_cache_layout, has_bias=bias is not None, ) - with torch.cuda.stream(launch_stream): + # The ambient current stream is already ordered correctly. Enter a + # stream context only when the caller explicitly requests a different + # stream for wrapper-owned casts, copies, and allocations. + stream_context = contextlib.nullcontext() if stream is None else torch.cuda.stream(launch_stream) + with stream_context: # Keep wrapper-owned casts and copies ordered with an explicit # non-current launch stream. block_table_i32 = ( From fc7f9b8e3f603e2c484985da4365e72d30cb2fe1 Mon Sep 17 00:00:00 2001 From: Xiake Sun Date: Wed, 2 Sep 2026 17:13:33 +0000 Subject: [PATCH 21/59] perf(fmha): schedule paged BF16 V192 fragments on demand - Load and consume each BF16 V192 P*V fragment immediately to reduce register pressure and scratch usage. - Enable the schedule for paged FP8 D192/V192 at KV lengths of 64K and above. - Share the long-context threshold with stagger selection and cover its dispatch boundary. - Validate random pages, partial tails, explicit streams, and graph replay. Signed-off-by: Xiake Sun --- .../attention/flash_attn_fp8_paged_gfx950.py | 2 + kernels/attention/flash_attn_interface.py | 11 ++- kernels/attention/flash_attn_utils.py | 69 +++++++++++++++---- tests/kernels/test_flash_attn_fwd.py | 21 +++++- 4 files changed, 86 insertions(+), 17 deletions(-) diff --git a/kernels/attention/flash_attn_fp8_paged_gfx950.py b/kernels/attention/flash_attn_fp8_paged_gfx950.py index 1ff8d55741..b64db44fd5 100644 --- a/kernels/attention/flash_attn_fp8_paged_gfx950.py +++ b/kernels/attention/flash_attn_fp8_paged_gfx950.py @@ -53,6 +53,7 @@ def build_flash_attn_paged_fp8_module( paged=False, kv_cache_layout="linear", paged_bn128=False, + v192_bf16_on_demand=False, ): """Build the gfx950 packed-varlen page-64 FP8 attention launcher.""" gpu_arch = get_hip_arch() @@ -107,6 +108,7 @@ def build_flash_attn_paged_fp8_module( cross_seqlen=cross_seqlen, paged=paged, kv_cache_layout=kv_cache_layout, + bf16_pv_on_demand=v192_bf16_on_demand, ) # Builder-level aliases used by SharedStorage and the launch/compile wrappers. SPLITK = traits.SPLITK diff --git a/kernels/attention/flash_attn_interface.py b/kernels/attention/flash_attn_interface.py index daf804c0a8..1b31baa59b 100644 --- a/kernels/attention/flash_attn_interface.py +++ b/kernels/attention/flash_attn_interface.py @@ -478,6 +478,7 @@ def _build_paged_fp8( setprio: bool, enable_stagger: bool, use_bn128: bool, + v192_bf16_on_demand: bool, ): """Build the gfx950 packed-varlen, vectorized page-64 FP8 launcher.""" from kernels.attention.flash_attn_fp8_paged_gfx950 import build_flash_attn_paged_fp8_module @@ -500,6 +501,7 @@ def _build_paged_fp8( paged=True, kv_cache_layout="vectorized", paged_bn128=use_bn128, + v192_bf16_on_demand=v192_bf16_on_demand, ) @@ -508,6 +510,11 @@ def _build_paged_fp8( # gfx950 dualwave paged-KV currently supports exactly one configuration. _PAGED_PAGE_SIZE = 64 _PAGED_BT_LDS_SIZE = 2048 +_PAGED_FP8_V192_ON_DEMAND_MIN_KV = 65536 + + +def _use_paged_fp8_v192_on_demand(head_dims: tuple[int, int], max_seqlen_kv: int) -> bool: + return head_dims == (192, 192) and max_seqlen_kv >= _PAGED_FP8_V192_ON_DEMAND_MIN_KV def _flydsl_flash_attn_paged( @@ -756,7 +763,8 @@ def _flydsl_flash_attn_paged( num_kv_pages = (skv + page_size - 1) // page_size use_bn128 = fp8_head_dims == (128, 128) and B == 1 and num_kv_pages % 2 == 0 paged_setprio = dualwave_swp_setprio and D != 192 - paged_stagger = dualwave_swp_enable_stagger and not (fp8_head_dims == (192, 192) and skv >= 65536) + v192_bf16_on_demand = _use_paged_fp8_v192_on_demand(fp8_head_dims, skv) + paged_stagger = dualwave_swp_enable_stagger and not v192_bf16_on_demand exe = _build_paged_fp8( num_heads=H, num_kv_heads=num_kv_heads, @@ -768,6 +776,7 @@ def _flydsl_flash_attn_paged( setprio=paged_setprio, enable_stagger=paged_stagger, use_bn128=use_bn128, + v192_bf16_on_demand=v192_bf16_on_demand, ) elif _paged_light_ok: exe = _build_paged_light( diff --git a/kernels/attention/flash_attn_utils.py b/kernels/attention/flash_attn_utils.py index 39b8f04574..be08128de1 100644 --- a/kernels/attention/flash_attn_utils.py +++ b/kernels/attention/flash_attn_utils.py @@ -2097,6 +2097,7 @@ class PagedDualwaveSwpFp8Traits: FP8_PV: bool FP8_PV_DIRECT: bool FP8_V_BANKPAD: bool + BF16_PV_ON_DEMAND: bool BN128: bool BN128_PF: bool QREG: bool @@ -2194,6 +2195,7 @@ def cache_tag(self): self.FP8_PV, self.FP8_PV_DIRECT, self.FP8_V_BANKPAD, + self.BF16_PV_ON_DEMAND, self.NUM_PREFETCH_K, self.BN128, self.BN128_PF, @@ -2222,6 +2224,7 @@ def _make_paged_dualwave_swp_fp8_traits( xcd_swizzle=False, paged=False, kv_cache_layout="linear", + bf16_pv_on_demand=False, ): """Build gfx950 DUALWAVE_SWP fp8 compile-time layout traits (dtype fixed to fp8).""" # Tile shape and wave geometry follow the gfx950 dual-wave 8-wave CTA. @@ -2310,6 +2313,7 @@ def _make_paged_dualwave_swp_fp8_traits( fp8_pv_direct = bn128 and (not paged or fp8_pv) if fp8_pv_direct: fp8_pv = True + bf16_pv_on_demand = bool(bf16_pv_on_demand and paged and head_dim == 192 and value_head_dim == 192) return PagedDualwaveSwpFp8Traits( BLOCK_M=block_m, @@ -2342,6 +2346,7 @@ def _make_paged_dualwave_swp_fp8_traits( FP8_PV=fp8_pv, FP8_PV_DIRECT=bool(fp8_pv_direct), FP8_V_BANKPAD=bool(fp8_v_bankpad), + BF16_PV_ON_DEMAND=bf16_pv_on_demand, BN128=bool(bn128), BN128_PF=bool(bn128_pf), QREG=bool(qreg), @@ -5223,20 +5228,6 @@ def _pv_step_fp8(self, step, v_p, v_v, v_o): v_o[step] = self._mfma_acc_fp8_wide(v_op, self._pv_p_fp8_cache, v_o[step]) return v_o - def pv_step_k(self, step, v_p, v_v, v_o): - if const_expr(self.traits.FP8_PV): - return self._pv_step_fp8(step, v_p, v_v, v_o) - # HIPREC PV: P and V are both v8 bf16, accumulated by a bf16 MMA. - v_p_lo, v_p_hi = v_p - v_pk = v_v[step] - if const_expr(step < 2): - p_pk = v_p_lo[step] - else: - p_pk = v_p_hi[step - 2] - for dc in range_constexpr(self.traits.D_CHUNKS): - v_o[dc] = self._mfma_acc_bf16(v_pk[dc], p_pk, v_o[dc]) - return v_o - def _mfma_acc_fp8_wide(self, a_i32x8, b_i32x8, c_v16): # Wide fp8 QK: mfma_scale (32x32x64) with unit E8M0 scales, i32x8 operands. return rocdl.mfma_scale_f32_32x32x64_f8f6f4( @@ -5272,6 +5263,37 @@ def _v_concat_i32x8(self, v_v, dc): words.append(fx.Int32(v2[1])) return Vec.from_elements(words, fx.Int32).ir_value() + def _pv_step_bf16_on_demand(self, step, v_p, buf_id, v_o): + traits = self.traits + v_base = buf_id * traits.VT_BF16_ELEMS + lane_base = ( + (self.lane_mod_32 // traits.SMEM_N_PER_WAVE) * traits.VEC_V_ROW_STRIDE + + self.lane_div_32 * traits.D128_BF + + (self.lane_mod_32 % traits.SMEM_N_PER_WAVE) * traits.VEC_BF + ) + v_base_ptr = buffer_ops.get_element_ptr( + self.lds_vt_base_ptr, + byte_offset=fx.Int32((v_base + lane_base) * traits.EB_BF), + elem_type=T.i8, + ) + v_values = [] + for dc in range_constexpr(traits.D_CHUNKS): + const_off = dc * (traits.D_CHUNK // traits.SMEM_N_PER_WAVE) * traits.VEC_V_ROW_STRIDE + step * ( + 2 * traits.D128_BF + ) + ptr = buffer_ops.get_element_ptr( + v_base_ptr, + byte_offset=fx.Int32(const_off * traits.EB_BF), + elem_type=T.i8, + ) + v_values.append(llvm.LoadOp(self.v8bf16_type, ptr, alignment=16).result) + rocdl.s_waitcnt(traits.LGKMCNT_0_ONLY) + + v_p_lo, v_p_hi = v_p + p_pk = v_p_lo[step] if const_expr(step < 2) else v_p_hi[step - 2] + for dc in range_constexpr(traits.D_CHUNKS): + v_o[dc] = self._mfma_acc_bf16(v_values[dc], p_pk, v_o[dc]) + return v_o def _load_q_wide_lds(self): traits = self.traits q_row_in_block = self.ctx_ref.q_row_in_block @@ -5318,6 +5340,21 @@ def qk(self, v_k, q_wide=None): rocdl.sched_group_barrier(traits.SCHED_MFMA_MASK, n_mfma - 1, 12) return (v_s_lo, v_s_hi) + def pv_step_k(self, step, v_p, v_v, v_o): + if const_expr(self.traits.BF16_PV_ON_DEMAND): + return self._pv_step_bf16_on_demand(step, v_p, v_v, v_o) + if const_expr(self.traits.FP8_PV): + return self._pv_step_fp8(step, v_p, v_v, v_o) + # HIPREC PV: P and V are both v8 bf16, accumulated by a bf16 MMA. + v_p_lo, v_p_hi = v_p + v_pk = v_v[step] + if const_expr(step < 2): + p_pk = v_p_lo[step] + else: + p_pk = v_p_hi[step - 2] + for dc in range_constexpr(self.traits.D_CHUNKS): + v_o[dc] = self._mfma_acc_bf16(v_pk[dc], p_pk, v_o[dc]) + return v_o def cast_p_fp8_direct(self, v_p): lo_partial_list, hi_full = v_p f32 = [] @@ -5761,6 +5798,8 @@ def _read_strip(key): return (_read_strip(n_lo), _read_strip(n_hi)) def load_v(self, buf_id): + if const_expr(self.traits.BF16_PV_ON_DEMAND): + return fx.Index(buf_id) if const_expr(self.traits.FP8_PV): if const_expr(self.traits.KV_VECTORIZED and self.traits.FP8_V_BANKPAD): return self._load_v_fp8_vectorized_bankpad(buf_id) @@ -5818,6 +5857,8 @@ def load_v_steps(self, buf_id, first_step, num_steps): transient register set while preserving the existing LDS layout. """ traits = self.traits + if const_expr(traits.BF16_PV_ON_DEMAND): + return fx.Index(buf_id) if const_expr(not traits.KV_VECTORIZED or traits.FP8_PV): raise RuntimeError("load_v_steps requires vectorized HIPREC V") if const_expr(first_step < 0 or num_steps < 1 or first_step + num_steps > 4): diff --git a/tests/kernels/test_flash_attn_fwd.py b/tests/kernels/test_flash_attn_fwd.py index 567017fde0..da6f38919c 100644 --- a/tests/kernels/test_flash_attn_fwd.py +++ b/tests/kernels/test_flash_attn_fwd.py @@ -4732,12 +4732,14 @@ def noncontiguous_copy(tensor): [(128, 128), (192, 128), (192, 192)], ids=["d128-v128", "d192-v128", "d192-v192"], ) -def test_paged_fp8_long_context_random_pages_matches_torch(head_dim, value_head_dim): +def test_paged_fp8_long_context_random_pages_matches_torch(head_dim, value_head_dim, monkeypatch): """Long context, shuffled pages, and a partial tail preserve accuracy.""" torch.manual_seed(23) query_length = 63 kv_length = 8161 num_pages = (kv_length + 63) // 64 + if (head_dim, value_head_dim) == (192, 192): + monkeypatch.setattr(flash_attn_interface, "_PAGED_FP8_V192_ON_DEMAND_MIN_KV", kv_length) query, query_descale = quantize_per_tensor_fp8(torch.randn(query_length, 16, head_dim, device="cuda") * 0.2) key, key_descale = quantize_per_tensor_fp8(torch.randn(num_pages, 1, head_dim // 16, 64, 16, device="cuda") * 0.2) @@ -4794,12 +4796,14 @@ def test_paged_fp8_long_context_random_pages_matches_torch(head_dim, value_head_ [(128, 128), (192, 128), (192, 192)], ids=["d128-v128", "d192-v128", "d192-v192"], ) -def test_paged_fp8_graph_replay_matches_torch(head_dim, value_head_dim): +def test_paged_fp8_graph_replay_matches_torch(head_dim, value_head_dim, monkeypatch): """The optimized paged variants preserve caller-owned graph output.""" torch.manual_seed(29) query_length = 63 kv_length = 128 num_pages = kv_length // 64 + if (head_dim, value_head_dim) == (192, 192): + monkeypatch.setattr(flash_attn_interface, "_PAGED_FP8_V192_ON_DEMAND_MIN_KV", kv_length) query, query_descale = quantize_per_tensor_fp8(torch.randn(query_length, 16, head_dim, device="cuda") * 0.2) key, key_descale = quantize_per_tensor_fp8(torch.randn(num_pages, 1, head_dim // 16, 64, 16, device="cuda") * 0.2) @@ -5517,6 +5521,19 @@ def test_fp8_rescale_threshold_drops_past_the_long_sequence_bound(): assert set(f(s) for s in (1, 1024, 4096, 4097, 8192, 131072)) == {6.0, 4.0} +@pytest.mark.parametrize( + ("head_dims", "max_seqlen_kv", "expected"), + [ + ((192, 192), 65535, False), + ((192, 192), 65536, True), + ((192, 128), 65536, False), + ((128, 128), 65536, False), + ], +) +def test_paged_fp8_v192_on_demand_dispatch_boundary(head_dims, max_seqlen_kv, expected): + assert flash_attn_interface._use_paged_fp8_v192_on_demand(head_dims, max_seqlen_kv) is expected + + @_requires_gfx950 def test_fp8_default_is_the_lazy_rescale(): """Every other fp8 case passes the flag, so the default would go untested. From a3a4c4fcbc0ddf19a84640ff5b53968fb8a5d8a2 Mon Sep 17 00:00:00 2001 From: Xiake Sun Date: Thu, 3 Sep 2026 08:10:11 +0000 Subject: [PATCH 22/59] perf(fmha): use segmented FP8 P*V for paged V192 - Factor value dimensions into generic H1/H2 segments and avoid padded D256 P*V work. - Reuse full-CTA H1 staging and stage H2 with complete active waves in the bank-padded FP8 layout. - Distribute FP8 P*V chunks across the existing pipeline and disable the BF16-specific wave stagger. - Preserve dense FP8 and other paged dimensions while covering factorization, dispatch, streams, tails, and graph replay. Signed-off-by: Xiake Sun --- .../attention/flash_attn_fp8_paged_gfx950.py | 4 +- kernels/attention/flash_attn_interface.py | 17 +- kernels/attention/flash_attn_utils.py | 227 +++++++++++++----- tests/kernels/test_flash_attn_fwd.py | 33 ++- 4 files changed, 190 insertions(+), 91 deletions(-) diff --git a/kernels/attention/flash_attn_fp8_paged_gfx950.py b/kernels/attention/flash_attn_fp8_paged_gfx950.py index b64db44fd5..1483ffcc82 100644 --- a/kernels/attention/flash_attn_fp8_paged_gfx950.py +++ b/kernels/attention/flash_attn_fp8_paged_gfx950.py @@ -53,7 +53,7 @@ def build_flash_attn_paged_fp8_module( paged=False, kv_cache_layout="linear", paged_bn128=False, - v192_bf16_on_demand=False, + fp8_pv_segmented=False, ): """Build the gfx950 packed-varlen page-64 FP8 attention launcher.""" gpu_arch = get_hip_arch() @@ -108,7 +108,7 @@ def build_flash_attn_paged_fp8_module( cross_seqlen=cross_seqlen, paged=paged, kv_cache_layout=kv_cache_layout, - bf16_pv_on_demand=v192_bf16_on_demand, + fp8_pv_segmented=fp8_pv_segmented, ) # Builder-level aliases used by SharedStorage and the launch/compile wrappers. SPLITK = traits.SPLITK diff --git a/kernels/attention/flash_attn_interface.py b/kernels/attention/flash_attn_interface.py index 1b31baa59b..1b2089c02e 100644 --- a/kernels/attention/flash_attn_interface.py +++ b/kernels/attention/flash_attn_interface.py @@ -478,7 +478,7 @@ def _build_paged_fp8( setprio: bool, enable_stagger: bool, use_bn128: bool, - v192_bf16_on_demand: bool, + fp8_pv_segmented: bool, ): """Build the gfx950 packed-varlen, vectorized page-64 FP8 launcher.""" from kernels.attention.flash_attn_fp8_paged_gfx950 import build_flash_attn_paged_fp8_module @@ -501,7 +501,7 @@ def _build_paged_fp8( paged=True, kv_cache_layout="vectorized", paged_bn128=use_bn128, - v192_bf16_on_demand=v192_bf16_on_demand, + fp8_pv_segmented=fp8_pv_segmented, ) @@ -510,11 +510,10 @@ def _build_paged_fp8( # gfx950 dualwave paged-KV currently supports exactly one configuration. _PAGED_PAGE_SIZE = 64 _PAGED_BT_LDS_SIZE = 2048 -_PAGED_FP8_V192_ON_DEMAND_MIN_KV = 65536 -def _use_paged_fp8_v192_on_demand(head_dims: tuple[int, int], max_seqlen_kv: int) -> bool: - return head_dims == (192, 192) and max_seqlen_kv >= _PAGED_FP8_V192_ON_DEMAND_MIN_KV +def _use_paged_fp8_segmented_pv(head_dims: tuple[int, int]) -> bool: + return head_dims == (192, 192) def _flydsl_flash_attn_paged( @@ -763,8 +762,10 @@ def _flydsl_flash_attn_paged( num_kv_pages = (skv + page_size - 1) // page_size use_bn128 = fp8_head_dims == (128, 128) and B == 1 and num_kv_pages % 2 == 0 paged_setprio = dualwave_swp_setprio and D != 192 - v192_bf16_on_demand = _use_paged_fp8_v192_on_demand(fp8_head_dims, skv) - paged_stagger = dualwave_swp_enable_stagger and not v192_bf16_on_demand + fp8_pv_segmented = _use_paged_fp8_segmented_pv(fp8_head_dims) + # The BF16 phase shift exposes H2 staging and probability-pack + # latency because segmented FP8 has only six wide P*V MFMAs. + paged_stagger = dualwave_swp_enable_stagger and not fp8_pv_segmented exe = _build_paged_fp8( num_heads=H, num_kv_heads=num_kv_heads, @@ -776,7 +777,7 @@ def _flydsl_flash_attn_paged( setprio=paged_setprio, enable_stagger=paged_stagger, use_bn128=use_bn128, - v192_bf16_on_demand=v192_bf16_on_demand, + fp8_pv_segmented=fp8_pv_segmented, ) elif _paged_light_ok: exe = _build_paged_light( diff --git a/kernels/attention/flash_attn_utils.py b/kernels/attention/flash_attn_utils.py index be08128de1..e32c20fe35 100644 --- a/kernels/attention/flash_attn_utils.py +++ b/kernels/attention/flash_attn_utils.py @@ -2097,7 +2097,9 @@ class PagedDualwaveSwpFp8Traits: FP8_PV: bool FP8_PV_DIRECT: bool FP8_V_BANKPAD: bool - BF16_PV_ON_DEMAND: bool + FP8_PV_SEGMENTED: bool + FP8_V_H1: int + FP8_V_H2: int BN128: bool BN128_PF: bool QREG: bool @@ -2195,7 +2197,9 @@ def cache_tag(self): self.FP8_PV, self.FP8_PV_DIRECT, self.FP8_V_BANKPAD, - self.BF16_PV_ON_DEMAND, + self.FP8_PV_SEGMENTED, + self.FP8_V_H1, + self.FP8_V_H2, self.NUM_PREFETCH_K, self.BN128, self.BN128_PF, @@ -2205,6 +2209,17 @@ def cache_tag(self): ) +def _factor_fp8_pv_head_dim(value_head_dim, segment_capacity, alignment): + """Split a vector width into at most two independently staged segments.""" + if value_head_dim <= 0 or value_head_dim % alignment != 0: + raise ValueError(f"FP8 P*V head dim must be a positive multiple of {alignment}, got {value_head_dim}") + h1 = min(1 << (value_head_dim.bit_length() - 1), segment_capacity) + h2 = value_head_dim - h1 + if h2 > segment_capacity: + raise ValueError(f"FP8 P*V head dim {value_head_dim} exceeds two segment capacities of {segment_capacity}") + return h1, h2 + + def _make_paged_dualwave_swp_fp8_traits( num_heads, num_kv_heads, @@ -2224,7 +2239,7 @@ def _make_paged_dualwave_swp_fp8_traits( xcd_swizzle=False, paged=False, kv_cache_layout="linear", - bf16_pv_on_demand=False, + fp8_pv_segmented=False, ): """Build gfx950 DUALWAVE_SWP fp8 compile-time layout traits (dtype fixed to fp8).""" # Tile shape and wave geometry follow the gfx950 dual-wave 8-wave CTA. @@ -2291,7 +2306,20 @@ def _make_paged_dualwave_swp_fp8_traits( vls_bf = slw_bf + 64 // eb_bf vt_bf16_elems = snrpt_bf * sdrpt_bf * vls_bf paged = bool(paged) - fp8_v_bankpad = paged and value_head_dim == 128 + token_groups16 = block_n // vec_kv + fp8_v_segment_capacity = block_size // token_groups16 + fp8_pv_segmented = bool( + fp8_pv_segmented + and paged + and value_head_dim > fp8_v_segment_capacity + and value_head_dim <= 2 * fp8_v_segment_capacity + and value_head_dim % d_chunk == 0 + ) + if fp8_pv_segmented: + fp8_v_h1, fp8_v_h2 = _factor_fp8_pv_head_dim(value_head_dim, fp8_v_segment_capacity, d_chunk) + else: + fp8_v_h1, fp8_v_h2 = value_head_dim, 0 + fp8_v_bankpad = paged and (value_head_dim == fp8_v_segment_capacity or fp8_pv_segmented) fp8_v_row_stride = block_n + 16 fp8_v_tile_bytes = (block_n // 8) * (head_dim // 16) * 128 if fp8_v_bankpad: @@ -2309,11 +2337,14 @@ def _make_paged_dualwave_swp_fp8_traits( # Paged V128 attention can consume its native FP8 V cache directly once # the vectorized cache is transposed into the bank-padded FP8 P*V layout. - fp8_pv = (paged and value_head_dim == 128) or (not paged and os.getenv("FLYDSL_FA_FP8_PV", "0") == "1") + fp8_pv = ( + (paged and value_head_dim == fp8_v_segment_capacity) + or fp8_pv_segmented + or (not paged and os.getenv("FLYDSL_FA_FP8_PV", "0") == "1") + ) fp8_pv_direct = bn128 and (not paged or fp8_pv) if fp8_pv_direct: fp8_pv = True - bf16_pv_on_demand = bool(bf16_pv_on_demand and paged and head_dim == 192 and value_head_dim == 192) return PagedDualwaveSwpFp8Traits( BLOCK_M=block_m, @@ -2346,7 +2377,9 @@ def _make_paged_dualwave_swp_fp8_traits( FP8_PV=fp8_pv, FP8_PV_DIRECT=bool(fp8_pv_direct), FP8_V_BANKPAD=bool(fp8_v_bankpad), - BF16_PV_ON_DEMAND=bf16_pv_on_demand, + FP8_PV_SEGMENTED=fp8_pv_segmented, + FP8_V_H1=fp8_v_h1, + FP8_V_H2=fp8_v_h2, BN128=bool(bn128), BN128_PF=bool(bn128_pf), QREG=bool(qreg), @@ -5263,36 +5296,47 @@ def _v_concat_i32x8(self, v_v, dc): words.append(fx.Int32(v2[1])) return Vec.from_elements(words, fx.Int32).ir_value() - def _pv_step_bf16_on_demand(self, step, v_p, buf_id, v_o): + def _load_v_fp8_segmented_chunk(self, buf_id, dc): traits = self.traits - v_base = buf_id * traits.VT_BF16_ELEMS - lane_base = ( - (self.lane_mod_32 // traits.SMEM_N_PER_WAVE) * traits.VEC_V_ROW_STRIDE - + self.lane_div_32 * traits.D128_BF - + (self.lane_mod_32 % traits.SMEM_N_PER_WAVE) * traits.VEC_BF - ) - v_base_ptr = buffer_ops.get_element_ptr( - self.lds_vt_base_ptr, - byte_offset=fx.Int32((v_base + lane_base) * traits.EB_BF), - elem_type=T.i8, - ) - v_values = [] - for dc in range_constexpr(traits.D_CHUNKS): - const_off = dc * (traits.D_CHUNK // traits.SMEM_N_PER_WAVE) * traits.VEC_V_ROW_STRIDE + step * ( - 2 * traits.D128_BF - ) + row_stride = traits.FP8_V_ROW_STRIDE + tile_bytes = traits.V_HEAD_DIM * row_stride + aligned_base = ((self.lds_vt_base_idx + fx.Index(127)) // fx.Index(128)) * fx.Index(128) + tile_base = aligned_base + fx.Index(buf_id * tile_bytes) + d_global = dc * traits.D_CHUNK + if const_expr(d_global < traits.FP8_V_H1): + segment_base = 0 + d_local = d_global + else: + segment_base = traits.FP8_V_H1 * row_stride + d_local = d_global - traits.FP8_V_H1 + token_base = self.lane_div_32 * fx.Index(32) + d_row = fx.Index(d_local) + self.lane_mod_32 + row_base = tile_base + fx.Index(segment_base) + d_row * fx.Index(row_stride) + token_base + halves = [] + for half in range_constexpr(2): ptr = buffer_ops.get_element_ptr( - v_base_ptr, - byte_offset=fx.Int32(const_off * traits.EB_BF), + self.lds_vt_base_ptr, + byte_offset=fx.Int32(row_base + fx.Index(half * 16) - self.lds_vt_base_idx), elem_type=T.i8, ) - v_values.append(llvm.LoadOp(self.v8bf16_type, ptr, alignment=16).result) + halves.append(llvm.LoadOp(Vec.make_type(2, fx.Int64), ptr, alignment=16).result) rocdl.s_waitcnt(traits.LGKMCNT_0_ONLY) + return _concat_vectors(halves[0], halves[1]).bitcast(fx.Int32).ir_value() - v_p_lo, v_p_hi = v_p - p_pk = v_p_lo[step] if const_expr(step < 2) else v_p_hi[step - 2] - for dc in range_constexpr(traits.D_CHUNKS): - v_o[dc] = self._mfma_acc_bf16(v_values[dc], p_pk, v_o[dc]) + def _pv_step_fp8_segmented(self, step, v_p, buf_id, v_o): + traits = self.traits + if const_expr(step == 0 or (traits.PAGED and step == 1)): + self._pv_segmented_p_cache = self._p_to_fp8_i32x8(v_p) + + chunks_per_step = traits.D_CHUNKS // 4 + extra_chunks = traits.D_CHUNKS % 4 + step_extra = 1 if step < extra_chunks else 0 + dc_start = step * chunks_per_step + min(step, extra_chunks) + dc_end = dc_start + chunks_per_step + step_extra + for dc in range_constexpr(dc_start, dc_end): + v_op = self._load_v_fp8_segmented_chunk(buf_id, dc) + v_o[dc] = self._mfma_acc_fp8_wide(v_op, self._pv_segmented_p_cache, v_o[dc]) + rocdl.sched_barrier(0) return v_o def _load_q_wide_lds(self): traits = self.traits @@ -5341,8 +5385,8 @@ def qk(self, v_k, q_wide=None): return (v_s_lo, v_s_hi) def pv_step_k(self, step, v_p, v_v, v_o): - if const_expr(self.traits.BF16_PV_ON_DEMAND): - return self._pv_step_bf16_on_demand(step, v_p, v_v, v_o) + if const_expr(self.traits.FP8_PV_SEGMENTED): + return self._pv_step_fp8_segmented(step, v_p, v_v, v_o) if const_expr(self.traits.FP8_PV): return self._pv_step_fp8(step, v_p, v_v, v_o) # HIPREC PV: P and V are both v8 bf16, accumulated by a bf16 MMA. @@ -5458,7 +5502,9 @@ def _run(): _run() def load_v(self, tile_start, buf_id, page_id=None): - if const_expr(self.traits.KV_VECTORIZED and self.traits.FP8_V_BANKPAD): + if const_expr(self.traits.KV_VECTORIZED and self.traits.FP8_PV_SEGMENTED): + self._stage_v_fp8_vectorized_segmented(tile_start, buf_id, page_id=page_id) + elif const_expr(self.traits.KV_VECTORIZED and self.traits.FP8_V_BANKPAD): self._stage_v_fp8_vectorized_bankpad(tile_start, buf_id, page_id=page_id) elif const_expr(self.traits.KV_VECTORIZED and self.traits.FP8_PV): self._stage_v_fp8_vectorized_dense_layout(tile_start, buf_id, page_id=page_id) @@ -5607,6 +5653,82 @@ def _load_store_active(): scf_if_dispatch(flat < fx.Index(total_v16), _load_store_active) + def _stage_v_fp8_vectorized_segmented(self, tile_start, buf_id, page_id=None): + if page_id is None: + page_id = self.load_page_id(tile_start) + prefix = self._load_v_fp8_vectorized_bankpad_source(tile_start, page_id=page_id) + self._store_v_fp8_vectorized_bankpad(prefix, buf_id) + if const_expr(self.traits.FP8_V_H2 > 0): + self._stage_v_fp8_vectorized_segment( + tile_start, + buf_id, + self.traits.FP8_V_H1, + self.traits.FP8_V_H2, + page_id, + ) + + def _permute_v_fp8_vectorized(self, src_i32x4): + src_words = Vec(src_i32x4, (4,), fx.Int32) + pair_ty = ir.Type.parse("!llvm.struct<(i32, i32)>") + pair_lo = rocdl.permlane16_swap( + pair_ty, + as_mlir_value(fx.Int32(src_words[0])), + as_mlir_value(fx.Int32(src_words[1])), + False, + False, + ) + pair_hi = rocdl.permlane16_swap( + pair_ty, + as_mlir_value(fx.Int32(src_words[2])), + as_mlir_value(fx.Int32(src_words[3])), + False, + False, + ) + return Vec.from_elements( + [ + fx.Int32(llvm.extractvalue(T.i32, pair_lo, [0])), + fx.Int32(llvm.extractvalue(T.i32, pair_lo, [1])), + fx.Int32(llvm.extractvalue(T.i32, pair_hi, [0])), + fx.Int32(llvm.extractvalue(T.i32, pair_hi, [1])), + ], + fx.Int32, + ) + + def _stage_v_fp8_vectorized_segment(self, tile_start, buf_id, d_offset, segment_dim, page_id): + traits = self.traits + src_div = self.make_page_view(self.v_base_iter, page_id, is_value=True) + token_groups16 = traits.BLOCK_N // traits.KV_VEC_SIZE + flat = self.tid + + def _stage_active(): + n_group16 = self.lane_in_warp // fx.Index(16) + d_local = self.wave_id * fx.Index(16) + self.lane_in_warp % fx.Index(16) + d_col = fx.Index(d_offset) + d_local + src_elem = ( + self.kv_head_idx * token_groups16 * traits.V_HEAD_DIM * traits.KV_VEC_SIZE + + n_group16 * traits.V_HEAD_DIM * traits.KV_VEC_SIZE + + d_col * traits.KV_VEC_SIZE + ) + reordered = self._permute_v_fp8_vectorized(self.buffer_load_fp8x16(src_div, src_elem)) + dest_group16 = (n_group16 % fx.Index(2)) * fx.Index(2) + n_group16 // fx.Index(2) + aligned_base = ((self.lds_vt_base_idx + fx.Index(127)) // fx.Index(128)) * fx.Index(128) + tile_bytes = traits.V_HEAD_DIM * traits.FP8_V_ROW_STRIDE + dst_byte = ( + aligned_base + + fx.Index(buf_id * tile_bytes) + + fx.Index(d_offset * traits.FP8_V_ROW_STRIDE) + + d_local * fx.Index(traits.FP8_V_ROW_STRIDE) + + dest_group16 * fx.Index(16) + ) + ptr = buffer_ops.get_element_ptr( + self.lds_vt_base_ptr, + byte_offset=fx.Int32(dst_byte - self.lds_vt_base_idx), + elem_type=T.i8, + ) + llvm.StoreOp(as_mlir_value(reordered), ptr, alignment=16) + + scf_if_dispatch(flat < fx.Index(segment_dim * token_groups16), _stage_active) + def _stage_v_fp8_vectorized_bankpad(self, tile_start, buf_id, page_id=None): """Stage vectorized paged V in the bank-padded direct-FP8 layout.""" src_i32x4 = self._load_v_fp8_vectorized_bankpad_source(tile_start, page_id=page_id) @@ -5619,9 +5741,10 @@ def _load_v_fp8_vectorized_bankpad_source(self, tile_start, page_id=None): page_id = self.load_page_id(tile_start) src_div = self.make_page_view(self.v_base_iter, page_id, is_value=True) token_groups16 = traits.BLOCK_N // traits.KV_VEC_SIZE - # permlane16 requires a full EXEC mask. V128 has exactly one vector per - # CTA thread, so no per-lane bounds branch is needed. - assert traits.V_HEAD_DIM * token_groups16 == traits.BLOCK_SIZE + # permlane16 requires a full EXEC mask. The first segment has exactly + # one vector per CTA thread, so no per-lane bounds branch is needed. + prefix_dim = traits.FP8_V_H1 if traits.FP8_PV_SEGMENTED else traits.V_HEAD_DIM + assert prefix_dim * token_groups16 == traits.BLOCK_SIZE n_group16 = self.lane_in_warp // fx.Index(16) d_col = self.wave_id * fx.Index(16) + self.lane_in_warp % fx.Index(16) src_elem = ( @@ -5640,31 +5763,7 @@ def _store_v_fp8_vectorized_bankpad(self, src_i32x4, buf_id): dst_base = aligned_base + fx.Index(buf_id * v_tile_bytes) n_group16 = self.lane_in_warp // fx.Index(16) d_col = self.wave_id * fx.Index(16) + self.lane_in_warp % fx.Index(16) - src_words = Vec(src_i32x4, (4,), fx.Int32) - pair_ty = ir.Type.parse("!llvm.struct<(i32, i32)>") - pair_lo = rocdl.permlane16_swap( - pair_ty, - as_mlir_value(fx.Int32(src_words[0])), - as_mlir_value(fx.Int32(src_words[1])), - False, - False, - ) - pair_hi = rocdl.permlane16_swap( - pair_ty, - as_mlir_value(fx.Int32(src_words[2])), - as_mlir_value(fx.Int32(src_words[3])), - False, - False, - ) - reordered = Vec.from_elements( - [ - fx.Int32(llvm.extractvalue(T.i32, pair_lo, [0])), - fx.Int32(llvm.extractvalue(T.i32, pair_lo, [1])), - fx.Int32(llvm.extractvalue(T.i32, pair_hi, [0])), - fx.Int32(llvm.extractvalue(T.i32, pair_hi, [1])), - ], - fx.Int32, - ) + reordered = self._permute_v_fp8_vectorized(src_i32x4) dest_group16 = (n_group16 % fx.Index(2)) * fx.Index(2) + n_group16 // fx.Index(2) dst_byte = dst_base + d_col * fx.Index(row_stride) + dest_group16 * fx.Index(16) ptr = buffer_ops.get_element_ptr( @@ -5798,7 +5897,7 @@ def _read_strip(key): return (_read_strip(n_lo), _read_strip(n_hi)) def load_v(self, buf_id): - if const_expr(self.traits.BF16_PV_ON_DEMAND): + if const_expr(self.traits.FP8_PV_SEGMENTED): return fx.Index(buf_id) if const_expr(self.traits.FP8_PV): if const_expr(self.traits.KV_VECTORIZED and self.traits.FP8_V_BANKPAD): @@ -5857,7 +5956,7 @@ def load_v_steps(self, buf_id, first_step, num_steps): transient register set while preserving the existing LDS layout. """ traits = self.traits - if const_expr(traits.BF16_PV_ON_DEMAND): + if const_expr(traits.FP8_PV_SEGMENTED): return fx.Index(buf_id) if const_expr(not traits.KV_VECTORIZED or traits.FP8_PV): raise RuntimeError("load_v_steps requires vectorized HIPREC V") diff --git a/tests/kernels/test_flash_attn_fwd.py b/tests/kernels/test_flash_attn_fwd.py index da6f38919c..2797e516bb 100644 --- a/tests/kernels/test_flash_attn_fwd.py +++ b/tests/kernels/test_flash_attn_fwd.py @@ -4732,15 +4732,12 @@ def noncontiguous_copy(tensor): [(128, 128), (192, 128), (192, 192)], ids=["d128-v128", "d192-v128", "d192-v192"], ) -def test_paged_fp8_long_context_random_pages_matches_torch(head_dim, value_head_dim, monkeypatch): +def test_paged_fp8_long_context_random_pages_matches_torch(head_dim, value_head_dim): """Long context, shuffled pages, and a partial tail preserve accuracy.""" torch.manual_seed(23) query_length = 63 kv_length = 8161 num_pages = (kv_length + 63) // 64 - if (head_dim, value_head_dim) == (192, 192): - monkeypatch.setattr(flash_attn_interface, "_PAGED_FP8_V192_ON_DEMAND_MIN_KV", kv_length) - query, query_descale = quantize_per_tensor_fp8(torch.randn(query_length, 16, head_dim, device="cuda") * 0.2) key, key_descale = quantize_per_tensor_fp8(torch.randn(num_pages, 1, head_dim // 16, 64, 16, device="cuda") * 0.2) value, value_descale = quantize_per_tensor_fp8( @@ -4796,15 +4793,12 @@ def test_paged_fp8_long_context_random_pages_matches_torch(head_dim, value_head_ [(128, 128), (192, 128), (192, 192)], ids=["d128-v128", "d192-v128", "d192-v192"], ) -def test_paged_fp8_graph_replay_matches_torch(head_dim, value_head_dim, monkeypatch): +def test_paged_fp8_graph_replay_matches_torch(head_dim, value_head_dim): """The optimized paged variants preserve caller-owned graph output.""" torch.manual_seed(29) query_length = 63 kv_length = 128 num_pages = kv_length // 64 - if (head_dim, value_head_dim) == (192, 192): - monkeypatch.setattr(flash_attn_interface, "_PAGED_FP8_V192_ON_DEMAND_MIN_KV", kv_length) - query, query_descale = quantize_per_tensor_fp8(torch.randn(query_length, 16, head_dim, device="cuda") * 0.2) key, key_descale = quantize_per_tensor_fp8(torch.randn(num_pages, 1, head_dim // 16, 64, 16, device="cuda") * 0.2) value, value_descale = quantize_per_tensor_fp8( @@ -5522,16 +5516,21 @@ def test_fp8_rescale_threshold_drops_past_the_long_sequence_bound(): @pytest.mark.parametrize( - ("head_dims", "max_seqlen_kv", "expected"), - [ - ((192, 192), 65535, False), - ((192, 192), 65536, True), - ((192, 128), 65536, False), - ((128, 128), 65536, False), - ], + ("value_head_dim", "expected"), + [(128, (128, 0)), (160, (128, 32)), (192, (128, 64)), (224, (128, 96)), (256, (128, 128))], +) +def test_fp8_pv_head_dim_uses_two_generic_segments(value_head_dim, expected): + from kernels.attention.flash_attn_utils import _factor_fp8_pv_head_dim + + assert _factor_fp8_pv_head_dim(value_head_dim, segment_capacity=128, alignment=32) == expected + + +@pytest.mark.parametrize( + ("head_dims", "expected"), + [((128, 128), False), ((192, 128), False), ((192, 192), True)], ) -def test_paged_fp8_v192_on_demand_dispatch_boundary(head_dims, max_seqlen_kv, expected): - assert flash_attn_interface._use_paged_fp8_v192_on_demand(head_dims, max_seqlen_kv) is expected +def test_paged_fp8_segmented_pv_dispatch(head_dims, expected): + assert flash_attn_interface._use_paged_fp8_segmented_pv(head_dims) is expected @_requires_gfx950 From 9d07ca5935fed9c883e3aaef5ad1b9d7f012cfaf Mon Sep 17 00:00:00 2001 From: Xiake Sun Date: Thu, 3 Sep 2026 17:40:13 +0000 Subject: [PATCH 23/59] perf(fmha): pipeline paged FP8 V prefetch - Carry two future V fragments across BN128 loop iterations. - Materialize prefetched fragments into bank-padded LDS after QK work. - Keep paired page-table indexing in the int32 kernel ABI. - Preserve D192 paths and the generic ragged/odd-page fallback. Signed-off-by: Xiake Sun --- .../attention/flash_attn_fp8_paged_gfx950.py | 36 +++++++++++++------ kernels/attention/flash_attn_utils.py | 10 +++--- 2 files changed, 32 insertions(+), 14 deletions(-) diff --git a/kernels/attention/flash_attn_fp8_paged_gfx950.py b/kernels/attention/flash_attn_fp8_paged_gfx950.py index 1483ffcc82..43f8aeb7e9 100644 --- a/kernels/attention/flash_attn_fp8_paged_gfx950.py +++ b/kernels/attention/flash_attn_fp8_paged_gfx950.py @@ -263,8 +263,14 @@ def _merge_tile_max(v_s_a, v_s_b): kv_gmem_to_lds.load_v((t0 + 1) * BN, (t0 + 1) % fx.Index(NPF), page_id=page_t1) kv_gmem_to_lds.load_k((t0 + 2) * BN, (t0 + 2) % fx.Index(NPF), page_id=page_t2) kv_gmem_to_lds.load_k((t0 + 3) * BN, (t0 + 3) % fx.Index(NPF), page_id=page_t3) - kv_gmem_to_lds.load_v((t0 + 2) * BN, (t0 + 2) % fx.Index(NPF), page_id=page_t2) - kv_gmem_to_lds.load_v((t0 + 3) * BN, (t0 + 3) % fx.Index(NPF), page_id=page_t3) + next_v_a = kv_gmem_to_lds._load_v_fp8_vectorized_bankpad_source( + (t0 + 2) * BN, + page_id=page_t2, + ) + next_v_b = kv_gmem_to_lds._load_v_fp8_vectorized_bankpad_source( + (t0 + 3) * BN, + page_id=page_t3, + ) rocdl.s_waitcnt(0) rocdl.sched_barrier(0) rocdl.s_barrier() @@ -279,8 +285,9 @@ def _merge_tile_max(v_s_a, v_s_b): def _ring_wrap(x): return (x >= NPF_I).select(x - NPF_I, x) - init_args = [m_row, l_row] + v_o + [t0 % fx.Index(NPF)] + init_args = [m_row, l_row] + v_o + [t0 % fx.Index(NPF), next_v_a, next_v_b] loop_results = init_args + next_v_arg_idx = 3 + D_CHUNKS for j, loop_args in range(fx.Index(t0), t_end, fx.Index(2), init=init_args): m_row = loop_args[0] l_row = loop_args[1] @@ -289,25 +296,33 @@ def _ring_wrap(x): a_buf = loop_args[2 + D_CHUNKS] b_buf = _ring_wrap(a_buf + fx.Index(1)) nn_a_buf = _ring_wrap(a_buf + fx.Index(2)) + nn_b_buf = _ring_wrap(a_buf + fx.Index(3)) f_a_buf = _ring_wrap(a_buf + fx.Index(4)) f_b_buf = _ring_wrap(a_buf + fx.Index(5)) + next_v_a = loop_args[next_v_arg_idx] + next_v_b = loop_args[next_v_arg_idx + 1] + v_k_a = kv_lds_to_regs.load_k(a_buf) v_k_b = kv_lds_to_regs.load_k(b_buf) v_v_a = kv_lds_to_regs.load_v(a_buf) page_f_a, page_f_b = ctx.load_page_id_pair((j + fx.Index(4)) * BN) - # Keep only one prefetched V vector live at a time and overlap each - # load with one of the two independent QK contractions. - v_f_a = kv_gmem_to_lds._load_v_fp8_vectorized_bankpad_source((j + fx.Index(4)) * BN, page_id=page_f_a) kv_gmem_to_lds.load_k((j + fx.Index(4)) * BN, f_a_buf, page_id=page_f_a) kv_gmem_to_lds.load_k((j + fx.Index(5)) * BN, f_b_buf, page_id=page_f_b) v_s_a = gemm_helper.qk(v_k_a, q_wide) - kv_gmem_to_lds._store_v_fp8_vectorized_bankpad(v_f_a, f_a_buf) - v_f_b = kv_gmem_to_lds._load_v_fp8_vectorized_bankpad_source((j + fx.Index(5)) * BN, page_id=page_f_b) + kv_gmem_to_lds._store_v_fp8_vectorized_bankpad(next_v_a, nn_a_buf) + v_f_a = kv_gmem_to_lds._load_v_fp8_vectorized_bankpad_source( + (j + fx.Index(4)) * BN, + page_id=page_f_a, + ) + v_f_b = kv_gmem_to_lds._load_v_fp8_vectorized_bankpad_source( + (j + fx.Index(5)) * BN, + page_id=page_f_b, + ) v_s_b = gemm_helper.qk(v_k_b, q_wide) - kv_gmem_to_lds._store_v_fp8_vectorized_bankpad(v_f_b, f_b_buf) + kv_gmem_to_lds._store_v_fp8_vectorized_bankpad(next_v_b, nn_b_buf) v_s_a, v_s_b = _mask_pair(v_s_a, v_s_b, j) m_tile = _merge_tile_max(v_s_a, v_s_b) v_o, m_new, l_row = _correct_o(v_o, m_row, l_row, m_tile) @@ -323,7 +338,8 @@ def _ring_wrap(x): rocdl.s_barrier() rocdl.sched_barrier(0) - loop_results = yield [m_row, l_row] + v_o + [nn_a_buf] + next_args = [m_row, l_row] + v_o + [nn_a_buf, v_f_a, v_f_b] + loop_results = yield next_args m_row = loop_results[0] l_row = loop_results[1] v_o = [loop_results[2 + i] for i in range_constexpr(D_CHUNKS)] diff --git a/kernels/attention/flash_attn_utils.py b/kernels/attention/flash_attn_utils.py index e32c20fe35..e586164bd3 100644 --- a/kernels/attention/flash_attn_utils.py +++ b/kernels/attention/flash_attn_utils.py @@ -5124,13 +5124,15 @@ def load_page_id(self, tile_start): def load_page_id_pair(self, tile_start): """Load two adjacent physical page IDs for a BN128 tile pair.""" - local_page = fx.Index(tile_start) // fx.Index(self.traits.PAGE_SIZE) - page_idx = self.page_base + local_page - safe_idx = fx.Index((local_page < self.num_kv_tiles).select(page_idx, self.page_base)) + local_page = fx.Int32(fx.Index(tile_start) // fx.Index(self.traits.PAGE_SIZE)) + page_base = fx.Int32(self.page_base) + num_kv_tiles = fx.Int32(self.num_kv_tiles) + page_idx = page_base + local_page + safe_idx = (local_page < num_kv_tiles).select(page_idx, page_base) v = fly.copy_atom_call_ssa( [self.page_v2i32], self.page_i32x2_atom, - fx.slice(self.page_indices_div, (None, fx.Int32(safe_idx))), + fx.slice(self.page_indices_div, (None, safe_idx)), ) pages = Vec(v, (2,), fx.Int32) page0 = rocdl.readfirstlane(T.i32, as_mlir_value(fx.Int32(pages[0]))) From 688884bdc7f7674af1276c804f316fec6883c456 Mon Sep 17 00:00:00 2001 From: Xiake Sun Date: Thu, 3 Sep 2026 19:20:00 +0000 Subject: [PATCH 24/59] perf(fmha): extend paged FP8 BN128 to varlen batches - Decouple BN128 tile selection from packed-varlen metadata. - Load per-sequence Q/KV offsets and causal lengths for B>1. - Preserve active q-block-aware longest-first scheduling. - Maintain B=1-relative performance through batch size 32. - Cover ragged multi-q-block batches with distinct causal offsets. Signed-off-by: Xiake Sun --- .../attention/flash_attn_fp8_paged_gfx950.py | 28 +++++++++++++++---- kernels/attention/flash_attn_interface.py | 6 +++- kernels/attention/flash_attn_utils.py | 3 +- tests/kernels/test_flash_attn_fwd.py | 14 ++++++++++ 4 files changed, 43 insertions(+), 8 deletions(-) diff --git a/kernels/attention/flash_attn_fp8_paged_gfx950.py b/kernels/attention/flash_attn_fp8_paged_gfx950.py index 43f8aeb7e9..65fd06d1e4 100644 --- a/kernels/attention/flash_attn_fp8_paged_gfx950.py +++ b/kernels/attention/flash_attn_fp8_paged_gfx950.py @@ -53,6 +53,7 @@ def build_flash_attn_paged_fp8_module( paged=False, kv_cache_layout="linear", paged_bn128=False, + paged_bn128_varlen=False, fp8_pv_segmented=False, ): """Build the gfx950 packed-varlen page-64 FP8 attention launcher.""" @@ -104,11 +105,12 @@ def build_flash_attn_paged_fp8_module( dualwave_swp_debug_lazy_counts=dualwave_swp_debug_lazy_counts, dualwave_swp_enable_stagger=dualwave_swp_enable_stagger, num_kv_splits=num_kv_splits, - varlen=False if paged_bn128 else varlen, + varlen=varlen if paged_bn128_varlen else False if paged_bn128 else varlen, cross_seqlen=cross_seqlen, paged=paged, kv_cache_layout=kv_cache_layout, fp8_pv_segmented=fp8_pv_segmented, + force_bn128=paged_bn128, ) # Builder-level aliases used by SharedStorage and the launch/compile wrappers. SPLITK = traits.SPLITK @@ -118,6 +120,7 @@ def build_flash_attn_paged_fp8_module( NUM_HEADS_Q = traits.NUM_HEADS_Q PAGED = traits.PAGED PAGED_BN128 = bool(paged_bn128) + PAGED_BN128_VARLEN = bool(paged_bn128_varlen) DEFAULT_STRIDE_Q_N = traits.DEFAULT_STRIDE_Q_N DEFAULT_STRIDE_O_N = traits.NUM_HEADS_Q * traits.V_HEAD_DIM DEFAULT_STRIDE_KV_N = traits.DEFAULT_STRIDE_KV_N @@ -185,9 +188,21 @@ def flash_attn_paged_fp8_bn128_kernel( ctx.init_runtime_indices() ctx.init_lds(SharedStorage) ctx.init_thread_mapping() - if const_expr(traits.CAUSAL): - ctx.init_causal_lpt_order() - ctx.init_sequence_lengths() + if const_expr(PAGED_BN128_VARLEN): + ctx.init_sequence_lengths() + if const_expr(traits.CAUSAL): + num_q_blocks = (ctx.seqlen_q_v + traits.BLOCK_M - 1) // traits.BLOCK_M + active_q_block = ctx.q_block_idx < num_q_blocks + reversed_q_block = num_q_blocks - fx.Index(1) - ctx.q_block_idx + ctx.q_block_idx = active_q_block.select(reversed_q_block, ctx.q_block_idx) + ctx.q_start = ctx.q_block_idx * traits.BLOCK_M + ctx.q_gmem_elem_offset = ( + ctx.q_tok_base + ctx.q_start + ) * ctx.stride_q_n_v + ctx.q_head_idx * traits.HEAD_DIM + else: + if const_expr(traits.CAUSAL): + ctx.init_causal_lpt_order() + ctx.init_sequence_lengths() ctx.init_descriptors() ctx.init_atoms_and_lds_ptrs() ctx.init_dma_thread_offsets() @@ -1106,9 +1121,10 @@ def _validate_paged_bn128_launch(batch_size, seq_len_kv, block_table_stride): seq_len_kv = int(seq_len_kv) block_table_stride = int(block_table_stride) num_kv_pages = (seq_len_kv + traits.PAGE_SIZE - 1) // traits.PAGE_SIZE - if batch_size != 1 or num_kv_pages < 2 or num_kv_pages % 2 != 0: + if (not PAGED_BN128_VARLEN and batch_size != 1) or num_kv_pages < 2 or num_kv_pages % 2 != 0: raise ValueError( - "paged BN128 requires batch_size=1 and a positive even number " + "paged BN128 requires batch_size=1 unless compiled for packed varlen, " + "and a positive even number " f"of KV pages; got batch_size={batch_size}, seq_len_kv={seq_len_kv}, " f"page_size={traits.PAGE_SIZE}" ) diff --git a/kernels/attention/flash_attn_interface.py b/kernels/attention/flash_attn_interface.py index 1b2089c02e..bdb801eded 100644 --- a/kernels/attention/flash_attn_interface.py +++ b/kernels/attention/flash_attn_interface.py @@ -478,6 +478,7 @@ def _build_paged_fp8( setprio: bool, enable_stagger: bool, use_bn128: bool, + paged_bn128_varlen: bool, fp8_pv_segmented: bool, ): """Build the gfx950 packed-varlen, vectorized page-64 FP8 launcher.""" @@ -501,6 +502,7 @@ def _build_paged_fp8( paged=True, kv_cache_layout="vectorized", paged_bn128=use_bn128, + paged_bn128_varlen=paged_bn128_varlen, fp8_pv_segmented=fp8_pv_segmented, ) @@ -760,7 +762,8 @@ def _flydsl_flash_attn_paged( ) if paged_fp8: num_kv_pages = (skv + page_size - 1) // page_size - use_bn128 = fp8_head_dims == (128, 128) and B == 1 and num_kv_pages % 2 == 0 + use_bn128 = fp8_head_dims == (128, 128) and num_kv_pages % 2 == 0 + paged_bn128_varlen = use_bn128 and B > 1 paged_setprio = dualwave_swp_setprio and D != 192 fp8_pv_segmented = _use_paged_fp8_segmented_pv(fp8_head_dims) # The BF16 phase shift exposes H2 staging and probability-pack @@ -777,6 +780,7 @@ def _flydsl_flash_attn_paged( setprio=paged_setprio, enable_stagger=paged_stagger, use_bn128=use_bn128, + paged_bn128_varlen=paged_bn128_varlen, fp8_pv_segmented=fp8_pv_segmented, ) elif _paged_light_ok: diff --git a/kernels/attention/flash_attn_utils.py b/kernels/attention/flash_attn_utils.py index e586164bd3..645a6422d0 100644 --- a/kernels/attention/flash_attn_utils.py +++ b/kernels/attention/flash_attn_utils.py @@ -2240,6 +2240,7 @@ def _make_paged_dualwave_swp_fp8_traits( paged=False, kv_cache_layout="linear", fp8_pv_segmented=False, + force_bn128=None, ): """Build gfx950 DUALWAVE_SWP fp8 compile-time layout traits (dtype fixed to fp8).""" # Tile shape and wave geometry follow the gfx950 dual-wave 8-wave CTA. @@ -2281,7 +2282,7 @@ def _make_paged_dualwave_swp_fp8_traits( smem_v_line_stride = smem_linear_wave + smem_v_pad smem_k_tile_elems = smem_n_rpt * smem_d_rpt * smem_k_line_stride smem_v_tile_elems = smem_n_rpt * smem_d_rpt * smem_v_line_stride - bn128 = (num_kv_splits <= 1) and (not varlen) + bn128 = bool(force_bn128) if force_bn128 is not None else (num_kv_splits <= 1) and (not varlen) bn128_pf = bn128 qreg = bn128_pf vdma = bn128_pf diff --git a/tests/kernels/test_flash_attn_fwd.py b/tests/kernels/test_flash_attn_fwd.py index 2797e516bb..5a5f3d29e9 100644 --- a/tests/kernels/test_flash_attn_fwd.py +++ b/tests/kernels/test_flash_attn_fwd.py @@ -4726,6 +4726,20 @@ def noncontiguous_copy(tensor): torch.testing.assert_close(actual, expected, rtol=2.0e-2, atol=2.0e-2) +@_requires_gfx950 +def test_paged_fp8_d128_bn128_ragged_multiblock_matches_torch(): + """The multi-batch BN128 path handles distinct causal offsets and inactive q-blocks.""" + test_paged_fp8_asymmetric_value_matches_torch( + head_dim=128, + value_head_dim=128, + use_non_default_stream=False, + force_internal_copies=False, + query_lengths=[512, 128], + kv_lengths=[1024, 512], + block_table_rows=[list(range(15, -1, -1)), list(range(23, 15, -1)) + [0] * 8], + ) + + @_requires_gfx950 @pytest.mark.parametrize( ("head_dim", "value_head_dim"), From c388f3b0d0a1b7b95651d2d28f75d7f601ca9d61 Mon Sep 17 00:00:00 2001 From: Xiake Sun Date: Thu, 3 Sep 2026 20:08:35 +0000 Subject: [PATCH 25/59] perf(fmha): balance paged FP8 D192 batches - Interleave up to eight MiMo batches on the head-fast grid axis. - Issue active causal q-blocks longest-first within each batch group. - Preserve the cache-locality-favored V192 mapping above batch size 16. - Keep divisor-safe fallbacks for arbitrary batch counts. - Cover dispatch policy and ragged multi-q-block D192 accuracy. Signed-off-by: Xiake Sun --- .../attention/flash_attn_fp8_paged_gfx950.py | 37 ++++++++++++++++++- kernels/attention/flash_attn_interface.py | 20 ++++++++++ kernels/attention/flash_attn_utils.py | 15 ++++++-- tests/kernels/test_flash_attn_fwd.py | 35 ++++++++++++++++++ 4 files changed, 102 insertions(+), 5 deletions(-) diff --git a/kernels/attention/flash_attn_fp8_paged_gfx950.py b/kernels/attention/flash_attn_fp8_paged_gfx950.py index 65fd06d1e4..85b3e09d62 100644 --- a/kernels/attention/flash_attn_fp8_paged_gfx950.py +++ b/kernels/attention/flash_attn_fp8_paged_gfx950.py @@ -55,6 +55,7 @@ def build_flash_attn_paged_fp8_module( paged_bn128=False, paged_bn128_varlen=False, fp8_pv_segmented=False, + batch_interleave_group=1, ): """Build the gfx950 packed-varlen page-64 FP8 attention launcher.""" gpu_arch = get_hip_arch() @@ -82,6 +83,11 @@ def build_flash_attn_paged_fp8_module( if num_kv_heads is None: num_kv_heads = num_heads + batch_interleave_group = int(batch_interleave_group) + if batch_interleave_group < 1: + raise ValueError(f"batch_interleave_group must be positive, got {batch_interleave_group}") + if batch_interleave_group > 1 and (paged_bn128 or head_dim != 192): + raise ValueError("batch interleaving is supported only by the generic paged D192 kernel") if paged_bn128 and (head_dim, value_head_dim) != (128, 128): raise RuntimeError("paged BN128 currently requires head_dim=value_head_dim=128") assert num_heads % num_kv_heads == 0 @@ -111,6 +117,7 @@ def build_flash_attn_paged_fp8_module( kv_cache_layout=kv_cache_layout, fp8_pv_segmented=fp8_pv_segmented, force_bn128=paged_bn128, + batch_interleave_group=batch_interleave_group, ) # Builder-level aliases used by SharedStorage and the launch/compile wrappers. SPLITK = traits.SPLITK @@ -121,6 +128,7 @@ def build_flash_attn_paged_fp8_module( PAGED = traits.PAGED PAGED_BN128 = bool(paged_bn128) PAGED_BN128_VARLEN = bool(paged_bn128_varlen) + BATCH_INTERLEAVE_GROUP = traits.BATCH_INTERLEAVE_GROUP DEFAULT_STRIDE_Q_N = traits.DEFAULT_STRIDE_Q_N DEFAULT_STRIDE_O_N = traits.NUM_HEADS_Q * traits.V_HEAD_DIM DEFAULT_STRIDE_KV_N = traits.DEFAULT_STRIDE_KV_N @@ -416,6 +424,16 @@ def flash_attn_dualwave_swp_fp8_gfx950_kernel( ctx.init_lds(SharedStorage) ctx.init_thread_mapping() ctx.init_sequence_lengths() + if const_expr(traits.HEAD_DIM == 192 and traits.CAUSAL and traits.BATCH_INTERLEAVE_GROUP > 1): + # Issue the longest active q-blocks first within each batch group. + num_q_blocks = (ctx.seqlen_q_v + traits.BLOCK_M - 1) // traits.BLOCK_M + active_q_block = ctx.q_block_idx < num_q_blocks + reversed_q_block = num_q_blocks - fx.Index(1) - ctx.q_block_idx + ctx.q_block_idx = active_q_block.select(reversed_q_block, ctx.q_block_idx) + ctx.q_start = ctx.q_block_idx * traits.BLOCK_M + ctx.q_gmem_elem_offset = ( + ctx.q_tok_base + ctx.q_start + ) * ctx.stride_q_n_v + ctx.q_head_idx * traits.HEAD_DIM ctx.init_descriptors() ctx.init_atoms_and_lds_ptrs() ctx.init_dma_thread_offsets() @@ -1092,7 +1110,15 @@ def launch_flash_attn_dualwave_swp( head_dim_runtime, value_attrs=kernel_attrs, ).launch( - grid=(NUM_HEADS_Q, num_q_blocks, grid_z), + grid=( + ( + NUM_HEADS_Q * traits.BATCH_INTERLEAVE_GROUP, + num_q_blocks, + grid_z // traits.BATCH_INTERLEAVE_GROUP, + ) + if const_expr(traits.BATCH_INTERLEAVE_GROUP > 1) + else (NUM_HEADS_Q, num_q_blocks, grid_z) + ), block=(BLOCK_SIZE, 1, 1), stream=stream, ) @@ -1133,6 +1159,13 @@ def _validate_paged_bn128_launch(batch_size, seq_len_kv, block_table_stride): "paged BN128 block table has too few entries: " f"need {num_kv_pages}, got stride {block_table_stride}" ) + def _validate_batch_interleave_launch(batch_size): + if int(batch_size) % BATCH_INTERLEAVE_GROUP != 0: + raise ValueError( + "paged D192 batch size must be divisible by its interleave group: " + f"batch_size={int(batch_size)}, group={BATCH_INTERLEAVE_GROUP}" + ) + def _launch( Q, K, @@ -1187,6 +1220,7 @@ def _launch( if PAGED and block_table is O: raise ValueError("paged fp8 flash_attn requires block_table") _validate_paged_bn128_launch(batch_size, seq_len_kv, block_table_stride) + _validate_batch_interleave_launch(batch_size) # Per-tensor fp8 descales (shape-[1] fp32). The kernel only reads them on # the fp8 path; bf16/f16 launches pass O as an unused placeholder. if q_descale is None: @@ -1269,6 +1303,7 @@ def _compile( if block_table_stride is None: block_table_stride = 0 _validate_paged_bn128_launch(batch_size, seq_len_kv, block_table_stride) + _validate_batch_interleave_launch(batch_size) if q_descale is None: q_descale = O if k_descale is None: diff --git a/kernels/attention/flash_attn_interface.py b/kernels/attention/flash_attn_interface.py index bdb801eded..41734e4053 100644 --- a/kernels/attention/flash_attn_interface.py +++ b/kernels/attention/flash_attn_interface.py @@ -480,6 +480,7 @@ def _build_paged_fp8( use_bn128: bool, paged_bn128_varlen: bool, fp8_pv_segmented: bool, + batch_interleave_group: int, ): """Build the gfx950 packed-varlen, vectorized page-64 FP8 launcher.""" from kernels.attention.flash_attn_fp8_paged_gfx950 import build_flash_attn_paged_fp8_module @@ -504,6 +505,7 @@ def _build_paged_fp8( paged_bn128=use_bn128, paged_bn128_varlen=paged_bn128_varlen, fp8_pv_segmented=fp8_pv_segmented, + batch_interleave_group=batch_interleave_group, ) @@ -512,12 +514,27 @@ def _build_paged_fp8( # gfx950 dualwave paged-KV currently supports exactly one configuration. _PAGED_PAGE_SIZE = 64 _PAGED_BT_LDS_SIZE = 2048 +_PAGED_FP8_BATCH_INTERLEAVE_MAX_GROUP = 8 +_PAGED_FP8_V192_BATCH_INTERLEAVE_MAX_BATCH = 16 def _use_paged_fp8_segmented_pv(head_dims: tuple[int, int]) -> bool: return head_dims == (192, 192) +def _paged_fp8_batch_interleave_group(batch_size: int, head_dims: tuple[int, int]) -> int: + """Choose a divisor-sized batch group for the tuned MiMo D192 launch.""" + if head_dims[0] != 192 or batch_size <= 1: + return 1 + # V192 crosses back to the cache-locality-favored original grid above B=16. + if head_dims == (192, 192) and batch_size > _PAGED_FP8_V192_BATCH_INTERLEAVE_MAX_BATCH: + return 1 + for group_size in (_PAGED_FP8_BATCH_INTERLEAVE_MAX_GROUP, 4, 2): + if batch_size % group_size == 0: + return group_size + return 1 + + def _flydsl_flash_attn_paged( q: torch.Tensor, k: torch.Tensor, @@ -782,6 +799,9 @@ def _flydsl_flash_attn_paged( use_bn128=use_bn128, paged_bn128_varlen=paged_bn128_varlen, fp8_pv_segmented=fp8_pv_segmented, + batch_interleave_group=( + _paged_fp8_batch_interleave_group(B, fp8_head_dims) if H == 16 and num_kv_heads == 1 else 1 + ), ) elif _paged_light_ok: exe = _build_paged_light( diff --git a/kernels/attention/flash_attn_utils.py b/kernels/attention/flash_attn_utils.py index 645a6422d0..ab4e76bf4b 100644 --- a/kernels/attention/flash_attn_utils.py +++ b/kernels/attention/flash_attn_utils.py @@ -1001,10 +1001,13 @@ def _init_dualwave_thread_mapping(ctx): linear_wg = fx.Index(gpu.block_idx.x) + fx.Index(gpu.block_idx.y) * fx.Index(traits.NUM_HEADS_Q) ctx.h_idx = linear_wg // num_q_blocks ctx.q_block_idx = linear_wg % num_q_blocks - elif const_expr(batch_interleave_group > 1): + elif const_expr(batch_interleave_group > 1 and not traits.SPLITK): + # Keep a bounded batch group head-fast so the final causal q-block + # occupies more CUs without abandoning K/V locality across all batches. linear_head_batch = fx.Index(gpu.block_idx.x) ctx.h_idx = linear_head_batch % traits.NUM_HEADS_Q - ctx.batch_idx = fx.Index(gpu.block_idx.z) * batch_interleave_group + linear_head_batch // traits.NUM_HEADS_Q + batch_in_group = linear_head_batch // traits.NUM_HEADS_Q + ctx.batch_idx = fx.Index(gpu.block_idx.z) * batch_interleave_group + batch_in_group ctx.q_block_idx = fx.Index(gpu.block_idx.y) else: ctx.h_idx = fx.Index(gpu.block_idx.x) @@ -1013,10 +1016,10 @@ def _init_dualwave_thread_mapping(ctx): ctx.bz_idx = fx.Index(gpu.block_idx.z) ctx.batch_idx = ctx.bz_idx // traits.NUM_KV_SPLITS ctx.split_idx = ctx.bz_idx % traits.NUM_KV_SPLITS - elif const_expr(batch_interleave_group > 1): + elif const_expr(batch_interleave_group <= 1): + ctx.batch_idx = fx.Index(gpu.block_idx.z) ctx.split_idx = None else: - ctx.batch_idx = fx.Index(gpu.block_idx.z) ctx.split_idx = None ctx.tid = fx.Index(gpu.thread_idx.x) @@ -2155,6 +2158,7 @@ class PagedDualwaveSwpFp8Traits: NEG_INF_F32_BITS: int LGKMCNT_0_ONLY: int XCD_SWIZZLE: bool = False + BATCH_INTERLEAVE_GROUP: int = 1 @property def HEAD_DIM_V(self): @@ -2206,6 +2210,7 @@ def cache_tag(self): self.QREG, self.VDMA, self.XCD_SWIZZLE, + self.BATCH_INTERLEAVE_GROUP, ) @@ -2241,6 +2246,7 @@ def _make_paged_dualwave_swp_fp8_traits( kv_cache_layout="linear", fp8_pv_segmented=False, force_bn128=None, + batch_interleave_group=1, ): """Build gfx950 DUALWAVE_SWP fp8 compile-time layout traits (dtype fixed to fp8).""" # Tile shape and wave geometry follow the gfx950 dual-wave 8-wave CTA. @@ -2435,6 +2441,7 @@ def _make_paged_dualwave_swp_fp8_traits( NEG_INF_F32_BITS=0xFF800000, LGKMCNT_0_ONLY=0xC07F, XCD_SWIZZLE=bool(xcd_swizzle), + BATCH_INTERLEAVE_GROUP=int(batch_interleave_group), ) diff --git a/tests/kernels/test_flash_attn_fwd.py b/tests/kernels/test_flash_attn_fwd.py index 5a5f3d29e9..9a5586e9f6 100644 --- a/tests/kernels/test_flash_attn_fwd.py +++ b/tests/kernels/test_flash_attn_fwd.py @@ -4740,6 +4740,41 @@ def test_paged_fp8_d128_bn128_ragged_multiblock_matches_torch(): ) +@pytest.mark.parametrize( + ("batch_size", "head_dims", "expected"), + [ + (1, (192, 128), 1), + (2, (192, 128), 2), + (8, (192, 128), 8), + (16, (192, 128), 8), + (32, (192, 128), 8), + (16, (192, 192), 8), + (32, (192, 192), 1), + (24, (192, 192), 1), + (12, (192, 128), 4), + (3, (192, 128), 1), + (8, (128, 128), 1), + ], +) +def test_paged_fp8_d192_batch_interleave_group(batch_size, head_dims, expected): + assert flash_attn_interface._paged_fp8_batch_interleave_group(batch_size, head_dims) == expected + + +@_requires_gfx950 +@pytest.mark.parametrize("value_head_dim", [128, 192]) +def test_paged_fp8_d192_batch_interleave_ragged_multiblock_matches_torch(value_head_dim): + """The D192 batch-interleaved grid preserves ragged causal q-block mapping.""" + test_paged_fp8_asymmetric_value_matches_torch( + head_dim=192, + value_head_dim=value_head_dim, + use_non_default_stream=False, + force_internal_copies=False, + query_lengths=[512, 128], + kv_lengths=[1024, 512], + block_table_rows=[list(range(15, -1, -1)), list(range(23, 15, -1)) + [0] * 8], + ) + + @_requires_gfx950 @pytest.mark.parametrize( ("head_dim", "value_head_dim"), From e82ca92ff992492c7ea6d704c3b973190e408b69 Mon Sep 17 00:00:00 2001 From: Xiake Sun Date: Thu, 3 Sep 2026 20:22:47 +0000 Subject: [PATCH 26/59] refactor(fmha): use typed paged FP8 helpers - Reuse the shared safe reciprocal helper for output normalization. - Remove private arithmetic conversion imports from the paged kernel. - Express paged lazy rescaling with typed FlyDSL if/else state. - Preserve generated ISA, resources, numerics, and performance. Signed-off-by: Xiake Sun --- .../attention/flash_attn_fp8_paged_gfx950.py | 12 +++---- kernels/attention/flash_attn_utils.py | 33 +++++++------------ 2 files changed, 15 insertions(+), 30 deletions(-) diff --git a/kernels/attention/flash_attn_fp8_paged_gfx950.py b/kernels/attention/flash_attn_fp8_paged_gfx950.py index 85b3e09d62..94f18be78e 100644 --- a/kernels/attention/flash_attn_fp8_paged_gfx950.py +++ b/kernels/attention/flash_attn_fp8_paged_gfx950.py @@ -8,8 +8,6 @@ from flydsl.compiler.kernel_function import CompilationContext from flydsl.expr import const_expr, range_constexpr, rocdl from flydsl.expr.typing import T -from flydsl.expr.utils.arith import ArithValue -from flydsl.expr.utils.arith import _to_raw as _raw from flydsl.runtime.device import get_rocm_arch as get_hip_arch from kernels.attention.flash_attn_utils import ( DualwaveFp8GemmHelper, @@ -367,10 +365,9 @@ def _ring_wrap(x): l_row = loop_results[1] v_o = [loop_results[2 + i] for i in range_constexpr(D_CHUNKS)] - inv_l_rcp = rocdl.rcp(T.f32, _raw(l_row)) - inv_l = ArithValue(fx.Float32(l_row) > ctx.c_zero_f).select(inv_l_rcp, ctx.c_zero_f) + inv_l = softmax_helper.safe_l_inv(l_row) if const_expr(traits.FP8_PV): - inv_l = ArithValue(inv_l) * ctx.vd_fp8 + inv_l = inv_l * ctx.vd_fp8 softmax_helper.scale_o(v_o, inv_l) rocdl.s_barrier() output_store.store_final_o(v_o, q_row) @@ -958,10 +955,9 @@ def _run_q_block(): # range; the combine kernel later applies w_s*l_s. # HIPREC folds v_descale into the bf16 vt scratch. The direct FP8 # D128 path keeps raw V and applies its descale once at the end. - inv_l_rcp = rocdl.rcp(T.f32, _raw(l_row)) - inv_l = ArithValue(fx.Float32(l_row) > ctx.c_zero_f).select(inv_l_rcp, ctx.c_zero_f) + inv_l = softmax_helper.safe_l_inv(l_row) if const_expr(traits.FP8_PV): - inv_l = ArithValue(inv_l) * ctx.vd_fp8 + inv_l = inv_l * ctx.vd_fp8 softmax_helper.scale_o(v_o, inv_l) # CLOSE the phase shift: one extra s_barrier on group A (complement of diff --git a/kernels/attention/flash_attn_utils.py b/kernels/attention/flash_attn_utils.py index ab4e76bf4b..69c467f532 100644 --- a/kernels/attention/flash_attn_utils.py +++ b/kernels/attention/flash_attn_utils.py @@ -6272,29 +6272,18 @@ def _run(v_o, m_row, l_row, m_tile_max, v_p): all_below = arith.cmpi(arith.CmpIPredicate.eq, as_mlir_value(ballot), _read_exec_i64()) all_below = llvm.intr_expect(all_below, arith.constant(1, type=ir.IntegerType.get_signless(1))) - state = [as_mlir_value(v_o[dc]) for dc in range(traits.D_CHUNKS)] - state += [self.v_p_to_vec32(v_p), as_mlir_value(l_row), as_mlir_value(m_row)] - names = tuple(f"_fp8_paged_lr{i}" for i in range(traits.D_CHUNKS + 3)) - - def _rescale(_names, *_state): + o_out = [as_mlir_value(v_o[dc]) for dc in range(traits.D_CHUNKS)] + vp_out = self.v_p_to_vec32(v_p) + l_out = as_mlir_value(l_row) + m_out = as_mlir_value(m_row) + if fx.Boolean(all_below): + pass + else: m_new, corr, scaled_accs = self._lazy_correction(v_o, m_row, m_tile_max) - result = list(scaled_accs) - result.append(self.v_p_to_vec32(self.scale_v_p(v_p, corr))) - result.append(as_mlir_value(l_row * corr)) - result.append(self.anchor_scalar_f32(m_new)) - return result - - result = scf_if_dispatch( - all_below, - lambda *_args: None, - _rescale, - state_names=names, - state_values=state, - ) - o_out = list(result[: traits.D_CHUNKS]) - vp_out = result[traits.D_CHUNKS] - l_out = result[traits.D_CHUNKS + 1] - m_out = result[traits.D_CHUNKS + 2] + o_out = [as_mlir_value(scaled_accs[dc]) for dc in range(traits.D_CHUNKS)] + vp_out = self.v_p_to_vec32(self.scale_v_p(v_p, corr)) + l_out = as_mlir_value(l_row * corr) + m_out = self.anchor_scalar_f32(m_new) return (o_out, m_out, l_out, self.v_vec32_to_p(vp_out)) return _run(v_o, m_row, l_row, m_tile_max, v_p) From fc99957fd3b3e4c40cceaa6f09b7fd014bd89406 Mon Sep 17 00:00:00 2001 From: Xiake Sun Date: Thu, 3 Sep 2026 20:32:46 +0000 Subject: [PATCH 27/59] refactor(fmha): remove unreachable paged FP8 split-K code - Delete split-K combine imports, kernel, and launch branches. - Remove dead workspace and partial-output plumbing. - Keep the existing num_kv_splits=1 validation contract. - Preserve generated ISA, resources, numerics, and performance. Signed-off-by: Xiake Sun --- .../attention/flash_attn_fp8_paged_gfx950.py | 93 ++----------------- 1 file changed, 7 insertions(+), 86 deletions(-) diff --git a/kernels/attention/flash_attn_fp8_paged_gfx950.py b/kernels/attention/flash_attn_fp8_paged_gfx950.py index 94f18be78e..679a51193f 100644 --- a/kernels/attention/flash_attn_fp8_paged_gfx950.py +++ b/kernels/attention/flash_attn_fp8_paged_gfx950.py @@ -7,7 +7,6 @@ import flydsl.expr as fx from flydsl.compiler.kernel_function import CompilationContext from flydsl.expr import const_expr, range_constexpr, rocdl -from flydsl.expr.typing import T from flydsl.runtime.device import get_rocm_arch as get_hip_arch from kernels.attention.flash_attn_utils import ( DualwaveFp8GemmHelper, @@ -17,15 +16,12 @@ DualwaveFp8QLoader, DualwaveFp8SoftmaxHelper, DualwaveFp8StoreHelper, - DualwaveSplitKCombineContext, - DualwaveSplitKCombineHelper, _make_paged_dualwave_swp_fp8_traits, _sched_barrier_exp_pairs, _sched_barrier_pairs, _stagger_extra_barrier_if_one, _stagger_extra_barrier_if_zero, _waitcnt_vm_n, - dualwave_splitk_workspace_elems, # noqa: F401 ) from kernels.common.kernels_common import dtype_to_elem_type from kernels.common.tensor_shim import _run_compiled @@ -89,11 +85,6 @@ def build_flash_attn_paged_fp8_module( if paged_bn128 and (head_dim, value_head_dim) != (128, 128): raise RuntimeError("paged BN128 currently requires head_dim=value_head_dim=128") assert num_heads % num_kv_heads == 0 - NUM_KV_SPLITS = int(num_kv_splits) - assert NUM_KV_SPLITS >= 1 - if varlen and num_kv_splits and int(num_kv_splits) > 1: - raise ValueError("varlen is not supported together with num_kv_splits > 1") - # All compile-time tile/layout constants live in the fp8 traits object. traits = _make_paged_dualwave_swp_fp8_traits( num_heads, @@ -118,7 +109,6 @@ def build_flash_attn_paged_fp8_module( batch_interleave_group=batch_interleave_group, ) # Builder-level aliases used by SharedStorage and the launch/compile wrappers. - SPLITK = traits.SPLITK BLOCK_M = traits.BLOCK_M BLOCK_SIZE = traits.BLOCK_SIZE HEAD_DIM = traits.HEAD_DIM @@ -214,7 +204,6 @@ def flash_attn_paged_fp8_bn128_kernel( ctx.init_dma_thread_offsets() ctx.init_descale() ctx.init_tile_bounds() - ctx.init_workspace_io() q_loader = DualwaveFp8QLoader(ctx) gemm_helper = DualwaveFp8GemmHelper(ctx) @@ -436,7 +425,6 @@ def flash_attn_dualwave_swp_fp8_gfx950_kernel( ctx.init_dma_thread_offsets() ctx.init_descale() ctx.init_tile_bounds() - ctx.init_workspace_io() # fp8 pipeline helpers (logic lives in flash_attn_utils; the kernel drives the # software-pipeline schedule below and calls into these). @@ -483,20 +471,12 @@ def _run_q_block(): v_s_0 = gemm_helper.qk(v_k, q_all_wide) rocdl.sched_barrier(0) if const_expr(traits.CAUSAL): - if const_expr(SPLITK): - v_s_0 = softmax_helper.causal_mask_prologue_if_needed( - v_s_0, ctx.split_t0, (ctx.split_t0 + 1) * traits.BLOCK_N - ) - else: - v_s_0 = softmax_helper.causal_mask_prologue_if_needed(v_s_0) + v_s_0 = softmax_helper.causal_mask_prologue_if_needed(v_s_0) else: # Non-causal padding mask for the prologue tile too: for tiny seq_len # tile 0 is the only real tile, so its keys >= seq_len must be masked # here. Gated -> no-op once tile 0 is full (seq_len >= BLOCK_N). - if const_expr(SPLITK): - v_s_0 = softmax_helper.seq_pad_mask_if_needed(v_s_0, ctx.split_t0) - else: - v_s_0 = softmax_helper.seq_pad_mask_if_needed(v_s_0) + v_s_0 = softmax_helper.seq_pad_mask_if_needed(v_s_0) m_row_pro = softmax_helper.reduce_max(v_s_0) if const_expr(traits.CAUSAL): # Floor fully-masked rows (-inf) to finite so exp2 yields 0, not NaN. @@ -519,10 +499,7 @@ def _run_q_block(): # ============================= Main loop ============================= # Software-pipelined inner loop - if const_expr(SPLITK): - loop_lb = ctx.split_t0 + 3 - else: - loop_lb = fx.Index(3) + loop_lb = fx.Index(3) loop_results = init_args for j, loop_args in range( loop_lb, @@ -951,8 +928,6 @@ def _run_q_block(): v_o = gemm_helper.pv(v_p_1, v_packs_e13, v_o) # Normalize by l_row; zero rows become zero instead of NaN. - # Split-K normalizes before packing so O_partial keeps useful mantissa - # range; the combine kernel later applies w_s*l_s. # HIPREC folds v_descale into the bf16 vt scratch. The direct FP8 # D128 path keeps raw V and applies its descale once at the end. inv_l = softmax_helper.safe_l_inv(l_row) @@ -970,48 +945,14 @@ def _run_q_block(): # 128b stores fuse this lane and its half-wave partner, so each pair # covers 8 contiguous columns instead of two 64b stores. - if const_expr(not SPLITK): - if const_expr(traits.VARLEN): - output_store.store_final_o_if_valid(v_o, q_row) - else: - output_store.store_final_o(v_o, q_row) + if const_expr(traits.VARLEN): + output_store.store_final_o_if_valid(v_o, q_row) else: - output_store.store_splitk_partial_o(v_o, m_row, l_row, q_row) + output_store.store_final_o(v_o, q_row) if ctx.q_start < ctx.seqlen_q_v: _run_q_block() - if const_expr(SPLITK): - output_store.store_empty_split() - - # Combine kernel: out = sum_s w_s * O_s / sum_s w_s * l_s, w_s = exp2(m_s - m_max). - # One wave row of 32 lanes covers a (b, h, s) row, 4 contiguous cols/lane. - COMBINE_BLOCK = 256 - COMBINE_LANES_PER_ROW = traits.HEAD_DIM // 4 - COMBINE_ROWS_PER_BLOCK = COMBINE_BLOCK // COMBINE_LANES_PER_ROW - - @flyc.kernel(known_block_size=[COMBINE_BLOCK, 1, 1]) - def flash_attn_splitk_combine_kernel( - O: fx.Tensor, # noqa: E741 - WS: fx.Tensor, - batch_size: fx.Int32, - seq_len: fx.Int32, - stride_q_n: fx.Int32, - ): - ctx = DualwaveSplitKCombineContext(traits, O, WS, batch_size, seq_len, stride_q_n) - ctx.init_types_and_constants() - ctx.init_runtime_indices() - ctx.init_thread_mapping(COMBINE_ROWS_PER_BLOCK, COMBINE_LANES_PER_ROW) - ctx.init_workspace() - ctx.init_descriptors() - - combine = DualwaveSplitKCombineHelper(ctx) - m_s, l_s = combine.load_ml_rows() - m_max = combine.reduce_m_max(m_s) - acc, den = combine.accumulate_splits(m_s, l_s, m_max) - o_pack = combine.pack_output(acc, den) - combine.store_output(o_pack) - @flyc.jit def launch_flash_attn_dualwave_swp( Q: fx.Tensor, @@ -1040,10 +981,7 @@ def launch_flash_attn_dualwave_swp( bs_idx = fx.Index(batch_size) sl_idx = fx.Index(seq_len) num_q_blocks = (sl_idx + BLOCK_M - 1) // BLOCK_M - if const_expr(SPLITK): - grid_z = bs_idx * NUM_KV_SPLITS - else: - grid_z = bs_idx + grid_z = bs_idx passthrough_entries = ( [ @@ -1118,13 +1056,6 @@ def launch_flash_attn_dualwave_swp( block=(BLOCK_SIZE, 1, 1), stream=stream, ) - if const_expr(SPLITK): - combine_rows = bs_idx * NUM_HEADS_Q * sl_idx - flash_attn_splitk_combine_kernel(O, DebugCounts, batch_size, seq_len, stride_o_n).launch( - grid=(combine_rows // COMBINE_ROWS_PER_BLOCK, 1, 1), - block=(COMBINE_BLOCK, 1, 1), - stream=stream, - ) _dualwave_swp_compile_hints = { "fast_fp_math": True, @@ -1176,7 +1107,6 @@ def _launch( debug_counts=None, *, seq_len_kv=None, - workspace=None, cu_seqlens_q=None, cu_seqlens_kv=None, block_table=None, @@ -1197,10 +1127,6 @@ def _launch( # seq_len_kv defaults to seq_len (self-attention / equal Q,KV lengths). if seq_len_kv is None: seq_len_kv = seq_len - if SPLITK: - if workspace is None: - raise ValueError("num_kv_splits > 1 requires a fp32 workspace (see dualwave_splitk_workspace_elems)") - debug_counts = workspace if debug_counts is None: debug_counts = O # Dense launches still pass valid tensors for the (unused) cu_seqlens slots; @@ -1264,7 +1190,6 @@ def _compile( debug_counts=None, *, seq_len_kv=None, - workspace=None, cu_seqlens_q=None, cu_seqlens_kv=None, block_table=None, @@ -1284,10 +1209,6 @@ def _compile( head_dim_runtime = HEAD_DIM if seq_len_kv is None: seq_len_kv = seq_len - if SPLITK: - if workspace is None: - raise ValueError("num_kv_splits > 1 requires a fp32 workspace (see dualwave_splitk_workspace_elems)") - debug_counts = workspace if debug_counts is None: debug_counts = O if cu_seqlens_q is None: From 484646f4186516ba43221b3b9d8557c46b0738cc Mon Sep 17 00:00:00 2001 From: Xiake Sun Date: Mon, 7 Sep 2026 01:22:20 +0000 Subject: [PATCH 28/59] refactor(fmha): remove unreachable paged FP8 non-causal paths - Specialize the dedicated paged FP8 kernels to their enforced causal contract. - Remove dead non-causal initialization, masking, and epilogue branches. - Preserve generated ISA, resources, accuracy, and performance. Signed-off-by: Xiake Sun --- .../attention/flash_attn_fp8_paged_gfx950.py | 97 +++++++------------ 1 file changed, 36 insertions(+), 61 deletions(-) diff --git a/kernels/attention/flash_attn_fp8_paged_gfx950.py b/kernels/attention/flash_attn_fp8_paged_gfx950.py index 679a51193f..f19b62b009 100644 --- a/kernels/attention/flash_attn_fp8_paged_gfx950.py +++ b/kernels/attention/flash_attn_fp8_paged_gfx950.py @@ -186,18 +186,16 @@ def flash_attn_paged_fp8_bn128_kernel( ctx.init_thread_mapping() if const_expr(PAGED_BN128_VARLEN): ctx.init_sequence_lengths() - if const_expr(traits.CAUSAL): - num_q_blocks = (ctx.seqlen_q_v + traits.BLOCK_M - 1) // traits.BLOCK_M - active_q_block = ctx.q_block_idx < num_q_blocks - reversed_q_block = num_q_blocks - fx.Index(1) - ctx.q_block_idx - ctx.q_block_idx = active_q_block.select(reversed_q_block, ctx.q_block_idx) - ctx.q_start = ctx.q_block_idx * traits.BLOCK_M - ctx.q_gmem_elem_offset = ( - ctx.q_tok_base + ctx.q_start - ) * ctx.stride_q_n_v + ctx.q_head_idx * traits.HEAD_DIM + num_q_blocks = (ctx.seqlen_q_v + traits.BLOCK_M - 1) // traits.BLOCK_M + active_q_block = ctx.q_block_idx < num_q_blocks + reversed_q_block = num_q_blocks - fx.Index(1) - ctx.q_block_idx + ctx.q_block_idx = active_q_block.select(reversed_q_block, ctx.q_block_idx) + ctx.q_start = ctx.q_block_idx * traits.BLOCK_M + ctx.q_gmem_elem_offset = ( + ctx.q_tok_base + ctx.q_start + ) * ctx.stride_q_n_v + ctx.q_head_idx * traits.HEAD_DIM else: - if const_expr(traits.CAUSAL): - ctx.init_causal_lpt_order() + ctx.init_causal_lpt_order() ctx.init_sequence_lengths() ctx.init_descriptors() ctx.init_atoms_and_lds_ptrs() @@ -237,9 +235,7 @@ def _subtile_tail(v_s, v_v, v_o, l_row, m_new): return v_o, l_row def _mask_pair(v_s_a, v_s_b, j): - if const_expr(traits.CAUSAL): - return softmax_helper.causal_mask_pair_if_needed(v_s_a, v_s_b, j) - return v_s_a, v_s_b + return softmax_helper.causal_mask_pair_if_needed(v_s_a, v_s_b, j) def _correct_o(v_o, m_row, l_row, m_tile): if const_expr(traits.DUALWAVE_SWP_LAZY_RESCALE): @@ -250,9 +246,7 @@ def _correct_o(v_o, m_row, l_row, m_tile): def _merge_tile_max(v_s_a, v_s_b): m_tile = softmax_helper.max2(softmax_helper.reduce_max(v_s_a), softmax_helper.reduce_max(v_s_b)) - if const_expr(traits.CAUSAL): - m_tile = softmax_helper.floor_masked_max(m_tile) - return m_tile + return softmax_helper.floor_masked_max(m_tile) page_t0, page_t1 = ctx.load_page_id_pair(t0 * BN) kv_gmem_to_lds.load_k(t0 * BN, t0 % fx.Index(NPF), page_id=page_t0) @@ -410,7 +404,7 @@ def flash_attn_dualwave_swp_fp8_gfx950_kernel( ctx.init_lds(SharedStorage) ctx.init_thread_mapping() ctx.init_sequence_lengths() - if const_expr(traits.HEAD_DIM == 192 and traits.CAUSAL and traits.BATCH_INTERLEAVE_GROUP > 1): + if const_expr(traits.HEAD_DIM == 192 and traits.BATCH_INTERLEAVE_GROUP > 1): # Issue the longest active q-blocks first within each batch group. num_q_blocks = (ctx.seqlen_q_v + traits.BLOCK_M - 1) // traits.BLOCK_M active_q_block = ctx.q_block_idx < num_q_blocks @@ -470,17 +464,10 @@ def _run_q_block(): # Prologue scores + first softmax pass for KV tile 0 v_s_0 = gemm_helper.qk(v_k, q_all_wide) rocdl.sched_barrier(0) - if const_expr(traits.CAUSAL): - v_s_0 = softmax_helper.causal_mask_prologue_if_needed(v_s_0) - else: - # Non-causal padding mask for the prologue tile too: for tiny seq_len - # tile 0 is the only real tile, so its keys >= seq_len must be masked - # here. Gated -> no-op once tile 0 is full (seq_len >= BLOCK_N). - v_s_0 = softmax_helper.seq_pad_mask_if_needed(v_s_0) + v_s_0 = softmax_helper.causal_mask_prologue_if_needed(v_s_0) m_row_pro = softmax_helper.reduce_max(v_s_0) - if const_expr(traits.CAUSAL): - # Floor fully-masked rows (-inf) to finite so exp2 yields 0, not NaN. - m_row_pro = softmax_helper.floor_masked_max(m_row_pro) + # Floor fully-masked rows (-inf) to finite so exp2 yields 0, not NaN. + m_row_pro = softmax_helper.floor_masked_max(m_row_pro) v_s_0 = softmax_helper.sub_m(v_s_0, m_row_pro) v_p_0 = softmax_helper.exp2(v_s_0, 0, 16) rocdl.sched_barrier(0) @@ -556,7 +543,7 @@ def _run_q_block(): v_o = gemm_helper.pv_step_k(0, v_p_0, v_v, v_o) # Cross-length causal can put a diagonal tile in v_s_1; mask it here. # Self-attention skips this to keep the existing schedule. - if const_expr(traits.CAUSAL and traits.CROSS_SEQLEN): + if const_expr(traits.CROSS_SEQLEN): v_s_1 = softmax_helper.causal_mask_prologue_if_needed( v_s_1, j_idx - 2, (j_idx - 1) * traits.BLOCK_N ) @@ -627,14 +614,11 @@ def _run_q_block(): v_packs_b = kv_lds_to_regs.load_v_steps(1, 0, 2) else: v_packs_b = kv_lds_to_regs.load_v(1) - if const_expr(traits.CAUSAL): - v_s_0 = softmax_helper.causal_mask_prologue_if_needed( - v_s_0, - j_idx - 1, - j_idx * traits.BLOCK_N, - ) - else: - v_s_0 = softmax_helper.v_s_vec_to_lists(v_s_0) + v_s_0 = softmax_helper.causal_mask_prologue_if_needed( + v_s_0, + j_idx - 1, + j_idx * traits.BLOCK_N, + ) rocdl.s_waitcnt(traits.LGKMCNT_0_ONLY) _waitcnt_vm_n(ctx.NUM_DMA_K + ctx.NUM_DMA_V) rocdl.sched_barrier(0) @@ -719,14 +703,11 @@ def _run_q_block(): v_packs_e3 = kv_lds_to_regs.load_v_steps(0, 0, 2) else: v_packs_e3 = kv_lds_to_regs.load_v(0) - if const_expr(traits.CAUSAL): - v_s_1 = softmax_helper.causal_mask_prologue_if_needed( - v_s_1, - max_m3, - max_m2 * traits.BLOCK_N, - ) - else: - v_s_1 = softmax_helper.seq_pad_mask_if_needed(v_s_1, max_m3) + v_s_1 = softmax_helper.causal_mask_prologue_if_needed( + v_s_1, + max_m3, + max_m2 * traits.BLOCK_N, + ) rocdl.s_waitcnt(traits.LGKMCNT_0_ONLY) _waitcnt_vm_n(ctx.NUM_DMA_K + ctx.NUM_DMA_V) rocdl.sched_barrier(0) @@ -790,14 +771,11 @@ def _run_q_block(): v_packs_e7 = kv_lds_to_regs.load_v_steps(1, 0, 2) else: v_packs_e7 = kv_lds_to_regs.load_v(1) - if const_expr(traits.CAUSAL): - v_s_0 = softmax_helper.causal_mask_prologue_if_needed( - v_s_0, - max_m2, - max_m1 * traits.BLOCK_N, - ) - else: - v_s_0 = softmax_helper.seq_pad_mask_if_needed(v_s_0, max_m2) + v_s_0 = softmax_helper.causal_mask_prologue_if_needed( + v_s_0, + max_m2, + max_m1 * traits.BLOCK_N, + ) rocdl.s_waitcnt(traits.LGKMCNT_0_ONLY) _waitcnt_vm_n(ctx.NUM_DMA_V) rocdl.sched_barrier(0) @@ -861,14 +839,11 @@ def _run_q_block(): v_packs_e11 = kv_lds_to_regs.load_v_steps(0, 0, 2) else: v_packs_e11 = kv_lds_to_regs.load_v(0) - if const_expr(traits.CAUSAL): - v_s_1 = softmax_helper.causal_mask_prologue_if_needed( - v_s_1, - max_m1, - ctx.split_t_end * traits.BLOCK_N, - ) - else: - v_s_1 = softmax_helper.seq_pad_mask_if_needed(v_s_1, max_m1) + v_s_1 = softmax_helper.causal_mask_prologue_if_needed( + v_s_1, + max_m1, + ctx.split_t_end * traits.BLOCK_N, + ) rocdl.s_waitcnt(traits.LGKMCNT_0_ONLY) _waitcnt_vm_n(0) rocdl.sched_barrier(0) From e92b34be236b6b528df37035446b622b7da0702a Mon Sep 17 00:00:00 2001 From: Xiake Sun Date: Mon, 7 Sep 2026 01:26:18 +0000 Subject: [PATCH 29/59] docs(fmha): trim paged FP8 schedule narration - Remove comments that restate individual pipeline loads and MFMAs. - Keep masking, scheduling, ABI, and layout rationale. - Correct stale dense/BF16 placeholder wording. Signed-off-by: Xiake Sun --- .../attention/flash_attn_fp8_paged_gfx950.py | 62 ++----------------- 1 file changed, 6 insertions(+), 56 deletions(-) diff --git a/kernels/attention/flash_attn_fp8_paged_gfx950.py b/kernels/attention/flash_attn_fp8_paged_gfx950.py index f19b62b009..9d3ee09126 100644 --- a/kernels/attention/flash_attn_fp8_paged_gfx950.py +++ b/kernels/attention/flash_attn_fp8_paged_gfx950.py @@ -85,7 +85,6 @@ def build_flash_attn_paged_fp8_module( if paged_bn128 and (head_dim, value_head_dim) != (128, 128): raise RuntimeError("paged BN128 currently requires head_dim=value_head_dim=128") assert num_heads % num_kv_heads == 0 - # All compile-time tile/layout constants live in the fp8 traits object. traits = _make_paged_dualwave_swp_fp8_traits( num_heads, num_kv_heads, @@ -108,7 +107,6 @@ def build_flash_attn_paged_fp8_module( force_bn128=paged_bn128, batch_interleave_group=batch_interleave_group, ) - # Builder-level aliases used by SharedStorage and the launch/compile wrappers. BLOCK_M = traits.BLOCK_M BLOCK_SIZE = traits.BLOCK_SIZE HEAD_DIM = traits.HEAD_DIM @@ -376,8 +374,6 @@ def flash_attn_dualwave_swp_fp8_gfx950_kernel( stride_kv_n: fx.Int32, head_dim_runtime: fx.Int32, ): - # Per-kernel setup lives in the fp8 context; the inline pipeline helpers below - # bind the ctx fields to local names so the schedule reads unchanged. ctx = DualwaveFp8KernelContext( traits, Q, @@ -420,8 +416,6 @@ def flash_attn_dualwave_swp_fp8_gfx950_kernel( ctx.init_descale() ctx.init_tile_bounds() - # fp8 pipeline helpers (logic lives in flash_attn_utils; the kernel drives the - # software-pipeline schedule below and calls into these). q_loader = DualwaveFp8QLoader(ctx) gemm_helper = DualwaveFp8GemmHelper(ctx) softmax_helper = DualwaveFp8SoftmaxHelper(ctx) @@ -433,7 +427,6 @@ def flash_attn_dualwave_swp_fp8_gfx950_kernel( # condition is uniform across the workgroup, so barriers stay balanced. @flyc.jit def _run_q_block(): - # Prologue: load K tile split_t0 -> LDS buf0, wait, and sync the workgroup. kv_gmem_to_lds.load_k(ctx.split_t0 * traits.BLOCK_N, 0) rocdl.s_waitcnt(0) rocdl.sched_barrier(0) @@ -446,7 +439,6 @@ def _run_q_block(): q_row = ctx.q_row q_all_wide = q_loader.load_all_wide(ctx.q_row_in_block) - # Pipeline ahead: prefetch K tile1 (buf1) + V tile0 (buf0) as background kv_gmem_to_lds.load_k((ctx.split_t0 + 1) * traits.BLOCK_N, 1) kv_gmem_to_lds.load_v(ctx.split_t0 * traits.BLOCK_N, 0) v_k = kv_lds_to_regs.load_k(0) @@ -461,7 +453,6 @@ def _run_q_block(): rocdl.sched_barrier(0) rocdl.s_barrier() - # Prologue scores + first softmax pass for KV tile 0 v_s_0 = gemm_helper.qk(v_k, q_all_wide) rocdl.sched_barrier(0) v_s_0 = softmax_helper.causal_mask_prologue_if_needed(v_s_0) @@ -474,18 +465,15 @@ def _run_q_block(): rocdl.s_barrier() rocdl.sched_barrier(0) - # Prefetch K tile 2 into buf0, keeping the K double-buffer one step ahead kv_gmem_to_lds.load_k((ctx.split_t0 + 2) * traits.BLOCK_N, 0) - # Loop-carried state (scf.for init args): m_row, l_row(=0), D_CHUNKS zero l_row_init = ctx.c_zero_f init_args = [m_row_pro, l_row_init] for _ in range_constexpr(traits.D_CHUNKS): init_args.append(ctx.c_zero_v16f32) init_args.append(ctx.v_pair_to_vec32(v_p_0)) - # ============================= Main loop ============================= - # Software-pipelined inner loop + # Software-pipelined KV loop. loop_lb = fx.Index(3) loop_results = init_args for j, loop_args in range( @@ -500,8 +488,6 @@ def _run_q_block(): v_p_0 = ctx.v_vec32_to_pair(loop_args[2 + traits.D_CHUNKS]) j_idx = j - # Cluster 0 (memory): prefetch next V (buf1), read resident K from LDS - # (v_k) for MMA0, wait + sync. kv_gmem_to_lds.load_v((j_idx - 2) * traits.BLOCK_N, 1) v_k = kv_lds_to_regs.load_k(1) rocdl.s_waitcnt(traits.LGKMCNT_0_ONLY) @@ -510,8 +496,6 @@ def _run_q_block(): rocdl.s_barrier() rocdl.sched_barrier(0) - # Cluster 1 (compute): MMA0 -> v_s_1; finish v_p_0's 2nd-half exp2, - # sum into l_row, cast to bf16 for P*V. v_s_1 = gemm_helper.qk(v_k, q_all_wide) v_p_0 = softmax_helper.exp2(v_p_0, 16, 16) l_row = softmax_helper.reduce_sum(l_row, v_p_0) @@ -523,8 +507,6 @@ def _run_q_block(): rocdl.s_barrier() rocdl.sched_barrier(0) - # Cluster 2 (memory): prefetch next K (buf1), read this tile's V from - # LDS (v_v) for P*V, wait + sync. kv_gmem_to_lds.load_k(j_idx * traits.BLOCK_N, 1) if const_expr(traits.D_CHUNKS > 4): v_v = kv_lds_to_regs.load_v_steps(0, 0, 2) @@ -536,8 +518,6 @@ def _run_q_block(): rocdl.s_barrier() rocdl.sched_barrier(0) - # Cluster 3 (compute): first P*V step + row max of v_s_1, lazy - # rescale, remaining 3 P*V steps, sub row + 1st-half exp2 of v_s_1. if const_expr(traits.DUALWAVE_SWP_SETPRIO): rocdl.s_setprio(1) v_o = gemm_helper.pv_step_k(0, v_p_0, v_v, v_o) @@ -584,8 +564,6 @@ def _run_q_block(): rocdl.s_barrier() rocdl.sched_barrier(0) - # Cluster 4 (memory, mirror of C0): prefetch V (buf0), read K from - # buf0 into v_k, wait + sync. kv_gmem_to_lds.load_v((j_idx - 1) * traits.BLOCK_N, 0) v_k = kv_lds_to_regs.load_k(0) rocdl.s_waitcnt(traits.LGKMCNT_0_ONLY) @@ -594,8 +572,6 @@ def _run_q_block(): rocdl.s_barrier() rocdl.sched_barrier(0) - # Cluster 5 (compute, mirror of C1): MMA0 -> v_s_0; finish v_p_1's - # 2nd-half exp2, sum into l_row, cast to bf16. v_s_0 = gemm_helper.qk(v_k, q_all_wide) v_p_1 = softmax_helper.exp2(v_p_1, 16, 16) l_row = softmax_helper.reduce_sum(l_row, v_p_1) @@ -607,8 +583,6 @@ def _run_q_block(): rocdl.s_barrier() rocdl.sched_barrier(0) - # Cluster 6 (memory): prefetch next K (buf0), read V packs (buf1), - # apply causal mask to v_s_0 (if causal), wait + sync. kv_gmem_to_lds.load_k((j_idx + 1) * traits.BLOCK_N, 0) if const_expr(traits.D_CHUNKS > 4): v_packs_b = kv_lds_to_regs.load_v_steps(1, 0, 2) @@ -625,8 +599,6 @@ def _run_q_block(): rocdl.s_barrier() rocdl.sched_barrier(0) - # Cluster 7 (compute, mirror of C3 for v_p_1/v_s_0): closes the iter, - # yield_args carries (m_row, l_row, v_o, packed v_p_0) to the next. if const_expr(traits.DUALWAVE_SWP_SETPRIO): rocdl.s_setprio(1) v_v = v_packs_b @@ -663,20 +635,16 @@ def _run_q_block(): yield_args = [m_row, l_row] + v_o + [ctx.v_pair_to_vec32(v_p_0)] loop_results = yield yield_args - # Epilogue: drain the pipeline for the final tiles the loop left in - # flight. Mirrors the main-loop clusters but with no further - # prefetch-ahead. Unpack the loop-carried state: + # Drain the final three tiles without further prefetch-ahead. m_row = loop_results[0] l_row = loop_results[1] v_o = [loop_results[2 + i] for i in range_constexpr(traits.D_CHUNKS)] v_p_0 = ctx.v_vec32_to_pair(loop_results[2 + traits.D_CHUNKS]) - # Tile indices for the last three tiles handled by the epilogue. max_m3 = ctx.split_t_end - 3 max_m2 = ctx.split_t_end - 2 max_m1 = ctx.split_t_end - 1 - # Epilogue C0 (memory): prefetch V max_m3 (buf1), read K from buf1, sync. kv_gmem_to_lds.load_v(max_m3 * traits.BLOCK_N, 1) v_k = kv_lds_to_regs.load_k(1) rocdl.s_waitcnt(traits.LGKMCNT_0_ONLY) @@ -685,7 +653,6 @@ def _run_q_block(): rocdl.s_barrier() rocdl.sched_barrier(0) - # Epilogue C1 (compute): MMA0 -> v_s_1; finish v_p_0 softmax (like C1). v_s_1 = gemm_helper.qk(v_k, q_all_wide) v_p_0 = softmax_helper.exp2(v_p_0, 16, 16) l_row = softmax_helper.reduce_sum(l_row, v_p_0) @@ -697,7 +664,6 @@ def _run_q_block(): rocdl.s_barrier() rocdl.sched_barrier(0) - # Epilogue C2 (memory): prefetch K max_m1, read V packs (buf0), causal mask v_s_1, sync. kv_gmem_to_lds.load_k(max_m1 * traits.BLOCK_N, 1) if const_expr(traits.D_CHUNKS > 4): v_packs_e3 = kv_lds_to_regs.load_v_steps(0, 0, 2) @@ -714,7 +680,6 @@ def _run_q_block(): rocdl.s_barrier() rocdl.sched_barrier(0) - # Epilogue C3 (compute): full P*V + unconditional rescale if const_expr(traits.DUALWAVE_SWP_SETPRIO): rocdl.s_setprio(1) if const_expr(traits.D_CHUNKS > 4): @@ -743,7 +708,6 @@ def _run_q_block(): rocdl.s_barrier() rocdl.sched_barrier(0) - # Epilogue C4 (memory): prefetch V max_m2 (buf0), read K from buf0, sync. kv_gmem_to_lds.load_v(max_m2 * traits.BLOCK_N, 0) v_k = kv_lds_to_regs.load_k(0) rocdl.s_waitcnt(traits.LGKMCNT_0_ONLY) @@ -752,8 +716,6 @@ def _run_q_block(): rocdl.s_barrier() rocdl.sched_barrier(0) - # Epilogue C5 (compute): MMA0 -> v_s_0; fold rescale_e3 into l_row, finish - # v_p_1 softmax. v_s_0 = gemm_helper.qk(v_k, q_all_wide) l_row = softmax_helper.apply_l_rescale(l_row, rescale_e3) v_p_1 = softmax_helper.exp2(v_p_1, 16, 16) @@ -766,7 +728,6 @@ def _run_q_block(): rocdl.s_barrier() rocdl.sched_barrier(0) - # Epilogue C6 (memory): read V packs (buf1), causal mask v_s_0, sync. if const_expr(traits.D_CHUNKS > 4): v_packs_e7 = kv_lds_to_regs.load_v_steps(1, 0, 2) else: @@ -782,7 +743,6 @@ def _run_q_block(): rocdl.s_barrier() rocdl.sched_barrier(0) - # Epilogue C7 (compute, mirror of C3): full P*V + unconditional rescale. if const_expr(traits.DUALWAVE_SWP_SETPRIO): rocdl.s_setprio(1) if const_expr(traits.D_CHUNKS > 4): @@ -810,7 +770,6 @@ def _run_q_block(): rocdl.s_barrier() rocdl.sched_barrier(0) - # Epilogue C8 (memory): prefetch V max_m1 (buf1), read K from buf1, sync. kv_gmem_to_lds.load_v(max_m1 * traits.BLOCK_N, 1) v_k = kv_lds_to_regs.load_k(1) rocdl.s_waitcnt(traits.LGKMCNT_0_ONLY) @@ -819,8 +778,6 @@ def _run_q_block(): rocdl.s_barrier() rocdl.sched_barrier(0) - # Epilogue C9 (compute): MMA0 -> v_s_1 (last tile); fold rescale_e7 into - # l_row, finish v_p_0 softmax. v_s_1 = gemm_helper.qk(v_k, q_all_wide) l_row = softmax_helper.apply_l_rescale(l_row, rescale_e7) v_p_0 = softmax_helper.exp2(v_p_0, 16, 16) @@ -833,8 +790,6 @@ def _run_q_block(): rocdl.s_barrier() rocdl.sched_barrier(0) - # Epilogue C10 (memory): read last V packs (buf0), causal mask v_s_1, - # drain all DMAs (vmcnt 0), sync. if const_expr(traits.D_CHUNKS > 4): v_packs_e11 = kv_lds_to_regs.load_v_steps(0, 0, 2) else: @@ -850,9 +805,6 @@ def _run_q_block(): rocdl.s_barrier() rocdl.sched_barrier(0) - # Epilogue C11 (compute): full P*V + rescale for v_p_0, then complete the - # last tile's softmax in-place (both exp2 halves, sum, cast) since no - # further pass follows. if const_expr(traits.D_CHUNKS > 4): v_o = gemm_helper.pv_step_k(0, v_p_0, v_packs_e11, v_o) v_o = gemm_helper.pv_step_k(1, v_p_0, v_packs_e11, v_o) @@ -881,7 +833,6 @@ def _run_q_block(): rocdl.s_barrier() rocdl.sched_barrier(0) - # Epilogue C12 (memory): read the final V packs for the closing P*V. if const_expr(traits.D_CHUNKS > 4): v_packs_e13 = kv_lds_to_regs.load_v_steps(1, 0, 2) else: @@ -891,7 +842,6 @@ def _run_q_block(): rocdl.s_barrier() rocdl.sched_barrier(0) - # Epilogue C13 (compute): final P*V -> v_o holds the unnormalized output. if const_expr(traits.D_CHUNKS > 4): v_o = gemm_helper.pv_step_k(0, v_p_1, v_packs_e13, v_o) v_o = gemm_helper.pv_step_k(1, v_p_1, v_packs_e13, v_o) @@ -1104,8 +1054,8 @@ def _launch( seq_len_kv = seq_len if debug_counts is None: debug_counts = O - # Dense launches still pass valid tensors for the (unused) cu_seqlens slots; - # the kernel only reads them under const_expr(VARLEN). Use O as a placeholder. + # Non-varlen B=1 BN128 ignores the cu_seqlens slots; use O as a placeholder + # for direct launcher calls that omit them. if cu_seqlens_q is None: cu_seqlens_q = O if cu_seqlens_kv is None: @@ -1118,8 +1068,8 @@ def _launch( raise ValueError("paged fp8 flash_attn requires block_table") _validate_paged_bn128_launch(batch_size, seq_len_kv, block_table_stride) _validate_batch_interleave_launch(batch_size) - # Per-tensor fp8 descales (shape-[1] fp32). The kernel only reads them on - # the fp8 path; bf16/f16 launches pass O as an unused placeholder. + # Direct launcher calls must supply shape-[1] fp32 descales; O keeps the + # compiled signature valid when a placeholder is needed. if q_descale is None: q_descale = O if k_descale is None: From fcb549ebfeeaeca4201f9a9080f5c6321bf79e98 Mon Sep 17 00:00:00 2001 From: Xiake Sun Date: Mon, 7 Sep 2026 02:39:11 +0000 Subject: [PATCH 30/59] refactor(fmha): attach compile hints to paged launcher - Store the validated compiler hints on the private JIT launcher. - Remove per-call CompilationContext scopes and the direct import. - Preserve hint-aware cache keys, generated ISA, resources, and performance. Signed-off-by: Xiake Sun --- .../attention/flash_attn_fp8_paged_gfx950.py | 96 +++++++++---------- 1 file changed, 47 insertions(+), 49 deletions(-) diff --git a/kernels/attention/flash_attn_fp8_paged_gfx950.py b/kernels/attention/flash_attn_fp8_paged_gfx950.py index 9d3ee09126..5d69eb7f3a 100644 --- a/kernels/attention/flash_attn_fp8_paged_gfx950.py +++ b/kernels/attention/flash_attn_fp8_paged_gfx950.py @@ -5,7 +5,6 @@ import flydsl.compiler as flyc import flydsl.expr as fx -from flydsl.compiler.kernel_function import CompilationContext from flydsl.expr import const_expr, range_constexpr, rocdl from flydsl.runtime.device import get_rocm_arch as get_hip_arch from kernels.attention.flash_attn_utils import ( @@ -991,6 +990,7 @@ def launch_flash_attn_dualwave_swp( "disable-machine-sink": True, }, } + launch_flash_attn_dualwave_swp.compile_hints = dict(_dualwave_swp_compile_hints) def _validate_paged_bn128_launch(batch_size, seq_len_kv, block_table_stride): if not PAGED_BN128: @@ -1076,30 +1076,29 @@ def _launch( k_descale = O if v_descale is None: v_descale = O - with CompilationContext.compile_hints(_dualwave_swp_compile_hints): - return _run_compiled( - launch_flash_attn_dualwave_swp, - Q, - K, - V, - O, - debug_counts, - cu_seqlens_q, - cu_seqlens_kv, - block_table, - block_table_stride, - q_descale, - k_descale, - v_descale, - batch_size, - seq_len, - seq_len_kv, - stride_q_n, - stride_o_n, - stride_kv_n, - head_dim_runtime, - fx.Stream(stream), - ) + return _run_compiled( + launch_flash_attn_dualwave_swp, + Q, + K, + V, + O, + debug_counts, + cu_seqlens_q, + cu_seqlens_kv, + block_table, + block_table_stride, + q_descale, + k_descale, + v_descale, + batch_size, + seq_len, + seq_len_kv, + stride_q_n, + stride_o_n, + stride_kv_n, + head_dim_runtime, + fx.Stream(stream), + ) def _compile( Q, @@ -1152,30 +1151,29 @@ def _compile( k_descale = O if v_descale is None: v_descale = O - with CompilationContext.compile_hints(_dualwave_swp_compile_hints): - return flyc.compile( - launch_flash_attn_dualwave_swp, - Q, - K, - V, - O, - debug_counts, - cu_seqlens_q, - cu_seqlens_kv, - block_table, - block_table_stride, - q_descale, - k_descale, - v_descale, - batch_size, - seq_len, - seq_len_kv, - stride_q_n, - stride_o_n, - stride_kv_n, - head_dim_runtime, - fx.Stream(stream), - ) + return flyc.compile( + launch_flash_attn_dualwave_swp, + Q, + K, + V, + O, + debug_counts, + cu_seqlens_q, + cu_seqlens_kv, + block_table, + block_table_stride, + q_descale, + k_descale, + v_descale, + batch_size, + seq_len, + seq_len_kv, + stride_q_n, + stride_o_n, + stride_kv_n, + head_dim_runtime, + fx.Stream(stream), + ) _launch.compile = _compile From 4c0912b66f4b8300b58706b3d1ef1eb35cfb6f92 Mon Sep 17 00:00:00 2001 From: Xiake Sun Date: Mon, 7 Sep 2026 03:20:31 +0000 Subject: [PATCH 31/59] refactor(fmha): use architecture-dispatched paged FP8 waits - Replace raw wait-all immediates with named VM/LGKM/EXP counters. - Replace the gfx950 LGKMCNT magic value with the public target-aware API. - Preserve wait placement, generated ISA, resources, accuracy, and performance. Signed-off-by: Xiake Sun --- .../attention/flash_attn_fp8_paged_gfx950.py | 44 +++++++++---------- 1 file changed, 22 insertions(+), 22 deletions(-) diff --git a/kernels/attention/flash_attn_fp8_paged_gfx950.py b/kernels/attention/flash_attn_fp8_paged_gfx950.py index 5d69eb7f3a..9652b5ae32 100644 --- a/kernels/attention/flash_attn_fp8_paged_gfx950.py +++ b/kernels/attention/flash_attn_fp8_paged_gfx950.py @@ -249,7 +249,7 @@ def _merge_tile_max(v_s_a, v_s_b): kv_gmem_to_lds.load_k(t0 * BN, t0 % fx.Index(NPF), page_id=page_t0) if const_expr(traits.QREG): q_loader.stage_q_to_lds() - rocdl.s_waitcnt(0) + fx.rocdl.s_waitcnt(vmcnt=0, lgkmcnt=0, expcnt=0) rocdl.sched_barrier(0) rocdl.s_barrier() @@ -272,7 +272,7 @@ def _merge_tile_max(v_s_a, v_s_b): (t0 + 3) * BN, page_id=page_t3, ) - rocdl.s_waitcnt(0) + fx.rocdl.s_waitcnt(vmcnt=0, lgkmcnt=0, expcnt=0) rocdl.sched_barrier(0) rocdl.s_barrier() rocdl.sched_barrier(0) @@ -334,7 +334,7 @@ def _ring_wrap(x): v_o, l_row = _subtile_tail(v_s_b, v_v_b, v_o, l_row, m_new) m_row = m_new - rocdl.s_waitcnt(0) + fx.rocdl.s_waitcnt(vmcnt=0, lgkmcnt=0, expcnt=0) rocdl.sched_barrier(0) rocdl.s_barrier() rocdl.sched_barrier(0) @@ -427,7 +427,7 @@ def flash_attn_dualwave_swp_fp8_gfx950_kernel( @flyc.jit def _run_q_block(): kv_gmem_to_lds.load_k(ctx.split_t0 * traits.BLOCK_N, 0) - rocdl.s_waitcnt(0) + fx.rocdl.s_waitcnt(vmcnt=0, lgkmcnt=0, expcnt=0) rocdl.sched_barrier(0) rocdl.s_barrier() @@ -442,7 +442,7 @@ def _run_q_block(): kv_gmem_to_lds.load_v(ctx.split_t0 * traits.BLOCK_N, 0) v_k = kv_lds_to_regs.load_k(0) rocdl.sched_barrier(0) - rocdl.s_waitcnt(traits.LGKMCNT_0_ONLY) + fx.rocdl.s_waitcnt(lgkmcnt=0) _waitcnt_vm_n(ctx.NUM_DMA_V) # OPEN the wave-group phase shift: one extra s_barrier on group B @@ -489,7 +489,7 @@ def _run_q_block(): kv_gmem_to_lds.load_v((j_idx - 2) * traits.BLOCK_N, 1) v_k = kv_lds_to_regs.load_k(1) - rocdl.s_waitcnt(traits.LGKMCNT_0_ONLY) + fx.rocdl.s_waitcnt(lgkmcnt=0) _waitcnt_vm_n(ctx.NUM_DMA_K + ctx.NUM_DMA_V) rocdl.sched_barrier(0) rocdl.s_barrier() @@ -511,7 +511,7 @@ def _run_q_block(): v_v = kv_lds_to_regs.load_v_steps(0, 0, 2) else: v_v = kv_lds_to_regs.load_v(0) - rocdl.s_waitcnt(traits.LGKMCNT_0_ONLY) + fx.rocdl.s_waitcnt(lgkmcnt=0) _waitcnt_vm_n(ctx.NUM_DMA_K + ctx.NUM_DMA_V) rocdl.sched_barrier(0) rocdl.s_barrier() @@ -541,7 +541,7 @@ def _run_q_block(): v_v_tail = kv_lds_to_regs.load_v_steps(0, 2, 2) v_s_1 = softmax_helper.sub_m(v_s_1, m_row) v_p_1 = softmax_helper.exp2(v_s_1, 0, 16) - rocdl.s_waitcnt(traits.LGKMCNT_0_ONLY) + fx.rocdl.s_waitcnt(lgkmcnt=0) v_o = gemm_helper.pv_step_k(2, v_p_0, v_v_tail, v_o) v_o = gemm_helper.pv_step_k(3, v_p_0, v_v_tail, v_o) else: @@ -565,7 +565,7 @@ def _run_q_block(): kv_gmem_to_lds.load_v((j_idx - 1) * traits.BLOCK_N, 0) v_k = kv_lds_to_regs.load_k(0) - rocdl.s_waitcnt(traits.LGKMCNT_0_ONLY) + fx.rocdl.s_waitcnt(lgkmcnt=0) _waitcnt_vm_n(ctx.NUM_DMA_K + ctx.NUM_DMA_V) rocdl.sched_barrier(0) rocdl.s_barrier() @@ -592,7 +592,7 @@ def _run_q_block(): j_idx - 1, j_idx * traits.BLOCK_N, ) - rocdl.s_waitcnt(traits.LGKMCNT_0_ONLY) + fx.rocdl.s_waitcnt(lgkmcnt=0) _waitcnt_vm_n(ctx.NUM_DMA_K + ctx.NUM_DMA_V) rocdl.sched_barrier(0) rocdl.s_barrier() @@ -615,7 +615,7 @@ def _run_q_block(): v_v_tail = kv_lds_to_regs.load_v_steps(1, 2, 2) v_s_0 = softmax_helper.sub_m(v_s_0, m_row) v_p_0 = softmax_helper.exp2(v_s_0, 0, 16) - rocdl.s_waitcnt(traits.LGKMCNT_0_ONLY) + fx.rocdl.s_waitcnt(lgkmcnt=0) v_o = gemm_helper.pv_step_k(2, v_p_1, v_v_tail, v_o) v_o = gemm_helper.pv_step_k(3, v_p_1, v_v_tail, v_o) else: @@ -646,7 +646,7 @@ def _run_q_block(): kv_gmem_to_lds.load_v(max_m3 * traits.BLOCK_N, 1) v_k = kv_lds_to_regs.load_k(1) - rocdl.s_waitcnt(traits.LGKMCNT_0_ONLY) + fx.rocdl.s_waitcnt(lgkmcnt=0) _waitcnt_vm_n(ctx.NUM_DMA_K + ctx.NUM_DMA_V) rocdl.sched_barrier(0) rocdl.s_barrier() @@ -673,7 +673,7 @@ def _run_q_block(): max_m3, max_m2 * traits.BLOCK_N, ) - rocdl.s_waitcnt(traits.LGKMCNT_0_ONLY) + fx.rocdl.s_waitcnt(lgkmcnt=0) _waitcnt_vm_n(ctx.NUM_DMA_K + ctx.NUM_DMA_V) rocdl.sched_barrier(0) rocdl.s_barrier() @@ -685,7 +685,7 @@ def _run_q_block(): v_o = gemm_helper.pv_step_k(0, v_p_0, v_packs_e3, v_o) v_o = gemm_helper.pv_step_k(1, v_p_0, v_packs_e3, v_o) v_packs_e3 = kv_lds_to_regs.load_v_steps(0, 2, 2) - rocdl.s_waitcnt(traits.LGKMCNT_0_ONLY) + fx.rocdl.s_waitcnt(lgkmcnt=0) v_o = gemm_helper.pv_step_k(2, v_p_0, v_packs_e3, v_o) v_o = gemm_helper.pv_step_k(3, v_p_0, v_packs_e3, v_o) else: @@ -709,7 +709,7 @@ def _run_q_block(): kv_gmem_to_lds.load_v(max_m2 * traits.BLOCK_N, 0) v_k = kv_lds_to_regs.load_k(0) - rocdl.s_waitcnt(traits.LGKMCNT_0_ONLY) + fx.rocdl.s_waitcnt(lgkmcnt=0) _waitcnt_vm_n(ctx.NUM_DMA_K + ctx.NUM_DMA_V) rocdl.sched_barrier(0) rocdl.s_barrier() @@ -736,7 +736,7 @@ def _run_q_block(): max_m2, max_m1 * traits.BLOCK_N, ) - rocdl.s_waitcnt(traits.LGKMCNT_0_ONLY) + fx.rocdl.s_waitcnt(lgkmcnt=0) _waitcnt_vm_n(ctx.NUM_DMA_V) rocdl.sched_barrier(0) rocdl.s_barrier() @@ -748,7 +748,7 @@ def _run_q_block(): v_o = gemm_helper.pv_step_k(0, v_p_1, v_packs_e7, v_o) v_o = gemm_helper.pv_step_k(1, v_p_1, v_packs_e7, v_o) v_packs_e7 = kv_lds_to_regs.load_v_steps(1, 2, 2) - rocdl.s_waitcnt(traits.LGKMCNT_0_ONLY) + fx.rocdl.s_waitcnt(lgkmcnt=0) v_o = gemm_helper.pv_step_k(2, v_p_1, v_packs_e7, v_o) v_o = gemm_helper.pv_step_k(3, v_p_1, v_packs_e7, v_o) else: @@ -771,7 +771,7 @@ def _run_q_block(): kv_gmem_to_lds.load_v(max_m1 * traits.BLOCK_N, 1) v_k = kv_lds_to_regs.load_k(1) - rocdl.s_waitcnt(traits.LGKMCNT_0_ONLY) + fx.rocdl.s_waitcnt(lgkmcnt=0) _waitcnt_vm_n(ctx.NUM_DMA_V) rocdl.sched_barrier(0) rocdl.s_barrier() @@ -798,7 +798,7 @@ def _run_q_block(): max_m1, ctx.split_t_end * traits.BLOCK_N, ) - rocdl.s_waitcnt(traits.LGKMCNT_0_ONLY) + fx.rocdl.s_waitcnt(lgkmcnt=0) _waitcnt_vm_n(0) rocdl.sched_barrier(0) rocdl.s_barrier() @@ -808,7 +808,7 @@ def _run_q_block(): v_o = gemm_helper.pv_step_k(0, v_p_0, v_packs_e11, v_o) v_o = gemm_helper.pv_step_k(1, v_p_0, v_packs_e11, v_o) v_packs_e11 = kv_lds_to_regs.load_v_steps(0, 2, 2) - rocdl.s_waitcnt(traits.LGKMCNT_0_ONLY) + fx.rocdl.s_waitcnt(lgkmcnt=0) v_o = gemm_helper.pv_step_k(2, v_p_0, v_packs_e11, v_o) v_o = gemm_helper.pv_step_k(3, v_p_0, v_packs_e11, v_o) else: @@ -836,7 +836,7 @@ def _run_q_block(): v_packs_e13 = kv_lds_to_regs.load_v_steps(1, 0, 2) else: v_packs_e13 = kv_lds_to_regs.load_v(1) - rocdl.s_waitcnt(traits.LGKMCNT_0_ONLY) + fx.rocdl.s_waitcnt(lgkmcnt=0) rocdl.sched_barrier(0) rocdl.s_barrier() rocdl.sched_barrier(0) @@ -845,7 +845,7 @@ def _run_q_block(): v_o = gemm_helper.pv_step_k(0, v_p_1, v_packs_e13, v_o) v_o = gemm_helper.pv_step_k(1, v_p_1, v_packs_e13, v_o) v_packs_e13 = kv_lds_to_regs.load_v_steps(1, 2, 2) - rocdl.s_waitcnt(traits.LGKMCNT_0_ONLY) + fx.rocdl.s_waitcnt(lgkmcnt=0) v_o = gemm_helper.pv_step_k(2, v_p_1, v_packs_e13, v_o) v_o = gemm_helper.pv_step_k(3, v_p_1, v_packs_e13, v_o) else: From dded5be64a16bd01bcf3bfe66c84e0835d7b19ca Mon Sep 17 00:00:00 2001 From: Xiake Sun Date: Mon, 7 Sep 2026 03:33:27 +0000 Subject: [PATCH 32/59] refactor(fmha): simplify static paged FP8 contracts - Remove compile-time branches excluded by the dedicated builder contract. - Use direct Q residency, causal masking, FP8 descale, and varlen output paths. - Collapse redundant bank-padded V trait initialization. - Use the architecture-dispatched wait API for segmented FP8 V loads. - Preserve kernel resources, accuracy, and performance. Signed-off-by: Xiake Sun --- .../attention/flash_attn_fp8_paged_gfx950.py | 26 +++++-------------- kernels/attention/flash_attn_utils.py | 3 +-- 2 files changed, 8 insertions(+), 21 deletions(-) diff --git a/kernels/attention/flash_attn_fp8_paged_gfx950.py b/kernels/attention/flash_attn_fp8_paged_gfx950.py index 9652b5ae32..e522147a74 100644 --- a/kernels/attention/flash_attn_fp8_paged_gfx950.py +++ b/kernels/attention/flash_attn_fp8_paged_gfx950.py @@ -110,7 +110,6 @@ def build_flash_attn_paged_fp8_module( BLOCK_SIZE = traits.BLOCK_SIZE HEAD_DIM = traits.HEAD_DIM NUM_HEADS_Q = traits.NUM_HEADS_Q - PAGED = traits.PAGED PAGED_BN128 = bool(paged_bn128) PAGED_BN128_VARLEN = bool(paged_bn128_varlen) BATCH_INTERLEAVE_GROUP = traits.BATCH_INTERLEAVE_GROUP @@ -247,8 +246,7 @@ def _merge_tile_max(v_s_a, v_s_b): page_t0, page_t1 = ctx.load_page_id_pair(t0 * BN) kv_gmem_to_lds.load_k(t0 * BN, t0 % fx.Index(NPF), page_id=page_t0) - if const_expr(traits.QREG): - q_loader.stage_q_to_lds() + q_loader.stage_q_to_lds() fx.rocdl.s_waitcnt(vmcnt=0, lgkmcnt=0, expcnt=0) rocdl.sched_barrier(0) rocdl.s_barrier() @@ -256,7 +254,7 @@ def _merge_tile_max(v_s_a, v_s_b): ctx.init_q_row() q_row = ctx.q_row - q_wide = gemm_helper.load_q_wide() if const_expr(traits.QREG) else None + q_wide = gemm_helper.load_q_wide() page_t2, page_t3 = ctx.load_page_id_pair((t0 + 2) * BN) kv_gmem_to_lds.load_k((t0 + 1) * BN, (t0 + 1) % fx.Index(NPF), page_id=page_t1) @@ -346,8 +344,7 @@ def _ring_wrap(x): v_o = [loop_results[2 + i] for i in range_constexpr(D_CHUNKS)] inv_l = softmax_helper.safe_l_inv(l_row) - if const_expr(traits.FP8_PV): - inv_l = inv_l * ctx.vd_fp8 + inv_l = inv_l * ctx.vd_fp8 softmax_helper.scale_o(v_o, inv_l) rocdl.s_barrier() output_store.store_final_o(v_o, q_row) @@ -522,12 +519,7 @@ def _run_q_block(): v_o = gemm_helper.pv_step_k(0, v_p_0, v_v, v_o) # Cross-length causal can put a diagonal tile in v_s_1; mask it here. # Self-attention skips this to keep the existing schedule. - if const_expr(traits.CROSS_SEQLEN): - v_s_1 = softmax_helper.causal_mask_prologue_if_needed( - v_s_1, j_idx - 2, (j_idx - 1) * traits.BLOCK_N - ) - else: - v_s_1 = softmax_helper.v_s_vec_to_lists(v_s_1) + v_s_1 = softmax_helper.causal_mask_prologue_if_needed(v_s_1, j_idx - 2, (j_idx - 1) * traits.BLOCK_N) m_tile_max_a = softmax_helper.reduce_max(v_s_1) _sched_barrier_pairs(traits, 4, 6, 2) @@ -855,8 +847,7 @@ def _run_q_block(): # HIPREC folds v_descale into the bf16 vt scratch. The direct FP8 # D128 path keeps raw V and applies its descale once at the end. inv_l = softmax_helper.safe_l_inv(l_row) - if const_expr(traits.FP8_PV): - inv_l = inv_l * ctx.vd_fp8 + inv_l = inv_l * ctx.vd_fp8 softmax_helper.scale_o(v_o, inv_l) # CLOSE the phase shift: one extra s_barrier on group A (complement of @@ -869,10 +860,7 @@ def _run_q_block(): # 128b stores fuse this lane and its half-wave partner, so each pair # covers 8 contiguous columns instead of two 64b stores. - if const_expr(traits.VARLEN): - output_store.store_final_o_if_valid(v_o, q_row) - else: - output_store.store_final_o(v_o, q_row) + output_store.store_final_o_if_valid(v_o, q_row) if ctx.q_start < ctx.seqlen_q_v: _run_q_block() @@ -1064,7 +1052,7 @@ def _launch( block_table = O if block_table_stride is None: block_table_stride = 0 - if PAGED and block_table is O: + if block_table is O: raise ValueError("paged fp8 flash_attn requires block_table") _validate_paged_bn128_launch(batch_size, seq_len_kv, block_table_stride) _validate_batch_interleave_launch(batch_size) diff --git a/kernels/attention/flash_attn_utils.py b/kernels/attention/flash_attn_utils.py index 69c467f532..78e9a62391 100644 --- a/kernels/attention/flash_attn_utils.py +++ b/kernels/attention/flash_attn_utils.py @@ -2331,7 +2331,6 @@ def _make_paged_dualwave_swp_fp8_traits( fp8_v_tile_bytes = (block_n // 8) * (head_dim // 16) * 128 if fp8_v_bankpad: fp8_v_tile_bytes = value_head_dim * fp8_v_row_stride - if fp8_v_bankpad: vt_bf16_total = num_prefetch_k * (fp8_v_tile_bytes // eb_bf) + 128 elif bn128_pf: vt_bf16_total = num_prefetch_k * (fp8_v_tile_bytes // eb_bf) + 128 @@ -5330,7 +5329,7 @@ def _load_v_fp8_segmented_chunk(self, buf_id, dc): elem_type=T.i8, ) halves.append(llvm.LoadOp(Vec.make_type(2, fx.Int64), ptr, alignment=16).result) - rocdl.s_waitcnt(traits.LGKMCNT_0_ONLY) + fx.rocdl.s_waitcnt(lgkmcnt=0) return _concat_vectors(halves[0], halves[1]).bitcast(fx.Int32).ir_value() def _pv_step_fp8_segmented(self, step, v_p, buf_id, v_o): From 50d98a1796860ef90cf4a3f907f894767cd902bd Mon Sep 17 00:00:00 2001 From: Xiake Sun Date: Tue, 8 Sep 2026 15:05:36 +0000 Subject: [PATCH 33/59] fix(fmha): rebase and fix dense and paged FP8 trait integration Signed-off-by: Xiake Sun --- kernels/attention/flash_attn_interface.py | 4 ++-- kernels/attention/flash_attn_utils.py | 3 +++ 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/kernels/attention/flash_attn_interface.py b/kernels/attention/flash_attn_interface.py index 41734e4053..e3a545042f 100644 --- a/kernels/attention/flash_attn_interface.py +++ b/kernels/attention/flash_attn_interface.py @@ -1080,8 +1080,8 @@ def flydsl_flash_attn_func( stream: CUDA/HIP stream to launch on. Returns: - Output tensor with the same shape as q except that paged asymmetric-V - attention uses V's final dimension. The dtype is bf16 for fp8 inputs, + Output tensor with q's leading dimensions and the V head width for + both dense and paged attention. The dtype is bf16 for fp8 inputs, otherwise the same dtype as q. When ``return_lse=True`` returns ``(out, lse)`` where ``lse`` is fp32 ``[B, num_heads, Sq]`` (varlen: ``[B, num_heads, max_seqlen_q]``, padded) holding the per-row diff --git a/kernels/attention/flash_attn_utils.py b/kernels/attention/flash_attn_utils.py index 78e9a62391..d443393ed7 100644 --- a/kernels/attention/flash_attn_utils.py +++ b/kernels/attention/flash_attn_utils.py @@ -1835,6 +1835,7 @@ class DualwaveSwpFp8Traits: QREG: bool = True FP8_PV: bool = True FP8_PV_DIRECT: bool = True + FP8_PV_SEGMENTED: bool = False @property def V_HEAD_DIM(self): @@ -5347,6 +5348,7 @@ def _pv_step_fp8_segmented(self, step, v_p, buf_id, v_o): v_o[dc] = self._mfma_acc_fp8_wide(v_op, self._pv_segmented_p_cache, v_o[dc]) rocdl.sched_barrier(0) return v_o + def _load_q_wide_lds(self): traits = self.traits q_row_in_block = self.ctx_ref.q_row_in_block @@ -5408,6 +5410,7 @@ def pv_step_k(self, step, v_p, v_v, v_o): for dc in range_constexpr(self.traits.D_CHUNKS): v_o[dc] = self._mfma_acc_bf16(v_pk[dc], p_pk, v_o[dc]) return v_o + def cast_p_fp8_direct(self, v_p): lo_partial_list, hi_full = v_p f32 = [] From c9e72fa1b736b12c0daef3f68c03f45c1dc5c44d Mon Sep 17 00:00:00 2001 From: Xiake Sun Date: Tue, 8 Sep 2026 15:43:32 +0000 Subject: [PATCH 34/59] perf(fmha): extend paged BN128 to D192/V128 - Keep D192 queries resident in registers without allocating a full Q LDS tile. - Reuse the paired-page native FP8 pipeline for D192/V128 and packed batches. - Preserve the existing generic fallback and vectorized page-cache ABI. Signed-off-by: Xiake Sun --- kernels/attention/flash_attn_fp8_paged_gfx950.py | 10 ++++++---- kernels/attention/flash_attn_interface.py | 6 ++++-- kernels/attention/flash_attn_utils.py | 2 +- 3 files changed, 11 insertions(+), 7 deletions(-) diff --git a/kernels/attention/flash_attn_fp8_paged_gfx950.py b/kernels/attention/flash_attn_fp8_paged_gfx950.py index e522147a74..c24ca4d548 100644 --- a/kernels/attention/flash_attn_fp8_paged_gfx950.py +++ b/kernels/attention/flash_attn_fp8_paged_gfx950.py @@ -81,8 +81,8 @@ def build_flash_attn_paged_fp8_module( raise ValueError(f"batch_interleave_group must be positive, got {batch_interleave_group}") if batch_interleave_group > 1 and (paged_bn128 or head_dim != 192): raise ValueError("batch interleaving is supported only by the generic paged D192 kernel") - if paged_bn128 and (head_dim, value_head_dim) != (128, 128): - raise RuntimeError("paged BN128 currently requires head_dim=value_head_dim=128") + if paged_bn128 and (head_dim, value_head_dim) not in ((128, 128), (192, 128)): + raise RuntimeError("paged BN128 currently requires value_head_dim=128") assert num_heads % num_kv_heads == 0 traits = _make_paged_dualwave_swp_fp8_traits( num_heads, @@ -118,6 +118,7 @@ def build_flash_attn_paged_fp8_module( DEFAULT_STRIDE_KV_N = traits.DEFAULT_STRIDE_KV_N _dualwave_swp_fp8_cache_tag = traits.cache_tag _lds_elem_dtype = dtype_to_elem_type(traits.DTYPE_STR) + _q_lds_elems = BLOCK_M * HEAD_DIM if traits.QLDS else 16 if PAGED_BN128: @@ -125,7 +126,7 @@ def build_flash_attn_paged_fp8_module( class SharedStorage: kv: fx.Array[_lds_elem_dtype, traits.LDS_KV_TOTAL_SIZE, 16] vt: fx.Array[fx.BFloat16, traits.VT_BF16_TOTAL, 16] - q: fx.Array[_lds_elem_dtype, BLOCK_M * HEAD_DIM, 16] + q: fx.Array[_lds_elem_dtype, _q_lds_elems, 16] else: @@ -246,7 +247,8 @@ def _merge_tile_max(v_s_a, v_s_b): page_t0, page_t1 = ctx.load_page_id_pair(t0 * BN) kv_gmem_to_lds.load_k(t0 * BN, t0 % fx.Index(NPF), page_id=page_t0) - q_loader.stage_q_to_lds() + if const_expr(traits.QLDS): + q_loader.stage_q_to_lds() fx.rocdl.s_waitcnt(vmcnt=0, lgkmcnt=0, expcnt=0) rocdl.sched_barrier(0) rocdl.s_barrier() diff --git a/kernels/attention/flash_attn_interface.py b/kernels/attention/flash_attn_interface.py index e3a545042f..e1d96f99f3 100644 --- a/kernels/attention/flash_attn_interface.py +++ b/kernels/attention/flash_attn_interface.py @@ -779,7 +779,7 @@ def _flydsl_flash_attn_paged( ) if paged_fp8: num_kv_pages = (skv + page_size - 1) // page_size - use_bn128 = fp8_head_dims == (128, 128) and num_kv_pages % 2 == 0 + use_bn128 = fp8_head_dims in ((128, 128), (192, 128)) and num_kv_pages % 2 == 0 paged_bn128_varlen = use_bn128 and B > 1 paged_setprio = dualwave_swp_setprio and D != 192 fp8_pv_segmented = _use_paged_fp8_segmented_pv(fp8_head_dims) @@ -800,7 +800,9 @@ def _flydsl_flash_attn_paged( paged_bn128_varlen=paged_bn128_varlen, fp8_pv_segmented=fp8_pv_segmented, batch_interleave_group=( - _paged_fp8_batch_interleave_group(B, fp8_head_dims) if H == 16 and num_kv_heads == 1 else 1 + _paged_fp8_batch_interleave_group(B, fp8_head_dims) + if not use_bn128 and H == 16 and num_kv_heads == 1 + else 1 ), ) elif _paged_light_ok: diff --git a/kernels/attention/flash_attn_utils.py b/kernels/attention/flash_attn_utils.py index d443393ed7..26610fd089 100644 --- a/kernels/attention/flash_attn_utils.py +++ b/kernels/attention/flash_attn_utils.py @@ -2167,7 +2167,7 @@ def HEAD_DIM_V(self): @property def QLDS(self): - return self.QREG + return self.QREG and self.HEAD_DIM <= 128 @property def cache_tag(self): From c15f035e532e8cd7dfa8ca6d730f51ef6f432a52 Mon Sep 17 00:00:00 2001 From: Xiake Sun Date: Tue, 8 Sep 2026 16:19:24 +0000 Subject: [PATCH 35/59] perf(fmha): use an eight-slot paged D128 prefetch ring - Load paged Q directly into registers to free LDS capacity. - Use eight K/V slots and power-of-two wrapping for the paired D128 pipeline. - Preserve the generic fallback for odd-page configurations. - Reach at least 95 percent of matched dense throughput across the single-sequence and batch-scaling gates. Signed-off-by: Xiake Sun --- kernels/attention/flash_attn_fp8_paged_gfx950.py | 2 ++ kernels/attention/flash_attn_utils.py | 4 +++- 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/kernels/attention/flash_attn_fp8_paged_gfx950.py b/kernels/attention/flash_attn_fp8_paged_gfx950.py index c24ca4d548..b949a30ae4 100644 --- a/kernels/attention/flash_attn_fp8_paged_gfx950.py +++ b/kernels/attention/flash_attn_fp8_paged_gfx950.py @@ -284,6 +284,8 @@ def _merge_tile_max(v_s_a, v_s_b): NPF_I = const_expr(fx.Index(NPF)) def _ring_wrap(x): + if const_expr(NPF == 8): + return x & fx.Index(7) return (x >= NPF_I).select(x - NPF_I, x) init_args = [m_row, l_row] + v_o + [t0 % fx.Index(NPF), next_v_a, next_v_b] diff --git a/kernels/attention/flash_attn_utils.py b/kernels/attention/flash_attn_utils.py index 26610fd089..e4ce75c9bf 100644 --- a/kernels/attention/flash_attn_utils.py +++ b/kernels/attention/flash_attn_utils.py @@ -2167,7 +2167,7 @@ def HEAD_DIM_V(self): @property def QLDS(self): - return self.QREG and self.HEAD_DIM <= 128 + return False @property def cache_tag(self): @@ -2295,6 +2295,8 @@ def _make_paged_dualwave_swp_fp8_traits( vdma = bn128_pf deep_ring = bn128 num_prefetch_k = (6 if bn128_pf else 4) if deep_ring else 2 + if paged and bn128 and head_dim == 128: + num_prefetch_k = 8 if bn128_pf: dualwave_swp_kv_per_buffer = smem_k_tile_elems else: From f537e4eef951057ab9bc692ef421ca0b7e752ec2 Mon Sep 17 00:00:00 2001 From: Xiake Sun Date: Tue, 8 Sep 2026 16:19:24 +0000 Subject: [PATCH 36/59] perf(fmha): add compact paired-page FP8 V192 pipeline - Compact K staging and XOR V addresses to fit six native FP8 K/V slots in LDS. - Stage future V directly into LDS and load current fragments after softmax. - Retain the public page-cache ABI and generic odd-page fallback. - Add ragged multi-KV-head accuracy coverage and pass the expanded attention suite. Signed-off-by: Xiake Sun --- .../attention/flash_attn_fp8_paged_gfx950.py | 121 +++++++++++------- kernels/attention/flash_attn_interface.py | 2 +- kernels/attention/flash_attn_utils.py | 91 +++++++++---- tests/kernels/test_flash_attn_fwd.py | 39 +++++- 4 files changed, 175 insertions(+), 78 deletions(-) diff --git a/kernels/attention/flash_attn_fp8_paged_gfx950.py b/kernels/attention/flash_attn_fp8_paged_gfx950.py index b949a30ae4..0b514b8a9e 100644 --- a/kernels/attention/flash_attn_fp8_paged_gfx950.py +++ b/kernels/attention/flash_attn_fp8_paged_gfx950.py @@ -81,8 +81,8 @@ def build_flash_attn_paged_fp8_module( raise ValueError(f"batch_interleave_group must be positive, got {batch_interleave_group}") if batch_interleave_group > 1 and (paged_bn128 or head_dim != 192): raise ValueError("batch interleaving is supported only by the generic paged D192 kernel") - if paged_bn128 and (head_dim, value_head_dim) not in ((128, 128), (192, 128)): - raise RuntimeError("paged BN128 currently requires value_head_dim=128") + if paged_bn128 and (head_dim, value_head_dim) not in ((128, 128), (192, 128), (192, 192)): + raise RuntimeError("paged BN128 requires supported FP8 head dimensions") assert num_heads % num_kv_heads == 0 traits = _make_paged_dualwave_swp_fp8_traits( num_heads, @@ -264,14 +264,18 @@ def _merge_tile_max(v_s_a, v_s_b): kv_gmem_to_lds.load_v((t0 + 1) * BN, (t0 + 1) % fx.Index(NPF), page_id=page_t1) kv_gmem_to_lds.load_k((t0 + 2) * BN, (t0 + 2) % fx.Index(NPF), page_id=page_t2) kv_gmem_to_lds.load_k((t0 + 3) * BN, (t0 + 3) % fx.Index(NPF), page_id=page_t3) - next_v_a = kv_gmem_to_lds._load_v_fp8_vectorized_bankpad_source( - (t0 + 2) * BN, - page_id=page_t2, - ) - next_v_b = kv_gmem_to_lds._load_v_fp8_vectorized_bankpad_source( - (t0 + 3) * BN, - page_id=page_t3, - ) + if const_expr(traits.FP8_PV_SEGMENTED): + kv_gmem_to_lds.load_v((t0 + 2) * BN, (t0 + 2) % fx.Index(NPF), page_id=page_t2) + kv_gmem_to_lds.load_v((t0 + 3) * BN, (t0 + 3) % fx.Index(NPF), page_id=page_t3) + else: + next_v_a = kv_gmem_to_lds._load_v_fp8_vectorized_bankpad_source( + (t0 + 2) * BN, + page_id=page_t2, + ) + next_v_b = kv_gmem_to_lds._load_v_fp8_vectorized_bankpad_source( + (t0 + 3) * BN, + page_id=page_t3, + ) fx.rocdl.s_waitcnt(vmcnt=0, lgkmcnt=0, expcnt=0) rocdl.sched_barrier(0) rocdl.s_barrier() @@ -288,7 +292,9 @@ def _ring_wrap(x): return x & fx.Index(7) return (x >= NPF_I).select(x - NPF_I, x) - init_args = [m_row, l_row] + v_o + [t0 % fx.Index(NPF), next_v_a, next_v_b] + init_args = [m_row, l_row] + v_o + [t0 % fx.Index(NPF)] + if const_expr(not traits.FP8_PV_SEGMENTED): + init_args += [next_v_a, next_v_b] loop_results = init_args next_v_arg_idx = 3 + D_CHUNKS for j, loop_args in range(fx.Index(t0), t_end, fx.Index(2), init=init_args): @@ -303,45 +309,70 @@ def _ring_wrap(x): f_a_buf = _ring_wrap(a_buf + fx.Index(4)) f_b_buf = _ring_wrap(a_buf + fx.Index(5)) - next_v_a = loop_args[next_v_arg_idx] - next_v_b = loop_args[next_v_arg_idx + 1] - - v_k_a = kv_lds_to_regs.load_k(a_buf) - v_k_b = kv_lds_to_regs.load_k(b_buf) - v_v_a = kv_lds_to_regs.load_v(a_buf) - - page_f_a, page_f_b = ctx.load_page_id_pair((j + fx.Index(4)) * BN) - kv_gmem_to_lds.load_k((j + fx.Index(4)) * BN, f_a_buf, page_id=page_f_a) - kv_gmem_to_lds.load_k((j + fx.Index(5)) * BN, f_b_buf, page_id=page_f_b) - - v_s_a = gemm_helper.qk(v_k_a, q_wide) - kv_gmem_to_lds._store_v_fp8_vectorized_bankpad(next_v_a, nn_a_buf) - v_f_a = kv_gmem_to_lds._load_v_fp8_vectorized_bankpad_source( - (j + fx.Index(4)) * BN, - page_id=page_f_a, - ) - v_f_b = kv_gmem_to_lds._load_v_fp8_vectorized_bankpad_source( - (j + fx.Index(5)) * BN, - page_id=page_f_b, - ) - v_s_b = gemm_helper.qk(v_k_b, q_wide) - kv_gmem_to_lds._store_v_fp8_vectorized_bankpad(next_v_b, nn_b_buf) - v_s_a, v_s_b = _mask_pair(v_s_a, v_s_b, j) - m_tile = _merge_tile_max(v_s_a, v_s_b) - v_o, m_new, l_row = _correct_o(v_o, m_row, l_row, m_tile) - v_o = softmax_helper.anchor_v_o(v_o) - - v_o, l_row = _subtile_tail(v_s_a, v_v_a, v_o, l_row, m_new) - v_v_b = kv_lds_to_regs.load_v(b_buf) - v_o, l_row = _subtile_tail(v_s_b, v_v_b, v_o, l_row, m_new) - m_row = m_new + if const_expr(traits.FP8_PV_SEGMENTED): + v_k_a = kv_lds_to_regs.load_k(a_buf) + v_k_b = kv_lds_to_regs.load_k(b_buf) + v_s_a = gemm_helper.qk(v_k_a, q_wide) + v_s_b = gemm_helper.qk(v_k_b, q_wide) + + page_f_a, page_f_b = ctx.load_page_id_pair((j + fx.Index(4)) * BN) + kv_gmem_to_lds.load_k((j + fx.Index(4)) * BN, f_a_buf, page_id=page_f_a) + kv_gmem_to_lds.load_k((j + fx.Index(5)) * BN, f_b_buf, page_id=page_f_b) + kv_gmem_to_lds.load_v((j + fx.Index(4)) * BN, f_a_buf, page_id=page_f_a) + kv_gmem_to_lds.load_v((j + fx.Index(5)) * BN, f_b_buf, page_id=page_f_b) + + v_s_a, v_s_b = _mask_pair(v_s_a, v_s_b, j) + m_tile = _merge_tile_max(v_s_a, v_s_b) + v_o, m_new, l_row = _correct_o(v_o, m_row, l_row, m_tile) + v_o = softmax_helper.anchor_v_o(v_o) + + v_p_a, l_row = _softmax_part(v_s_a, l_row, m_new) + v_v_a = kv_lds_to_regs.load_v(a_buf) + v_o = _pv_part(v_p_a, v_v_a, v_o) + v_p_b, l_row = _softmax_part(v_s_b, l_row, m_new) + v_v_b = kv_lds_to_regs.load_v(b_buf) + v_o = _pv_part(v_p_b, v_v_b, v_o) + m_row = m_new + next_args = [m_row, l_row] + v_o + [nn_a_buf] + else: + next_v_a = loop_args[next_v_arg_idx] + next_v_b = loop_args[next_v_arg_idx + 1] + + v_k_a = kv_lds_to_regs.load_k(a_buf) + v_k_b = kv_lds_to_regs.load_k(b_buf) + v_v_a = kv_lds_to_regs.load_v(a_buf) + + page_f_a, page_f_b = ctx.load_page_id_pair((j + fx.Index(4)) * BN) + kv_gmem_to_lds.load_k((j + fx.Index(4)) * BN, f_a_buf, page_id=page_f_a) + kv_gmem_to_lds.load_k((j + fx.Index(5)) * BN, f_b_buf, page_id=page_f_b) + + v_s_a = gemm_helper.qk(v_k_a, q_wide) + kv_gmem_to_lds._store_v_fp8_vectorized_bankpad(next_v_a, nn_a_buf) + v_f_a = kv_gmem_to_lds._load_v_fp8_vectorized_bankpad_source( + (j + fx.Index(4)) * BN, + page_id=page_f_a, + ) + v_f_b = kv_gmem_to_lds._load_v_fp8_vectorized_bankpad_source( + (j + fx.Index(5)) * BN, + page_id=page_f_b, + ) + v_s_b = gemm_helper.qk(v_k_b, q_wide) + kv_gmem_to_lds._store_v_fp8_vectorized_bankpad(next_v_b, nn_b_buf) + v_s_a, v_s_b = _mask_pair(v_s_a, v_s_b, j) + m_tile = _merge_tile_max(v_s_a, v_s_b) + v_o, m_new, l_row = _correct_o(v_o, m_row, l_row, m_tile) + v_o = softmax_helper.anchor_v_o(v_o) + + v_o, l_row = _subtile_tail(v_s_a, v_v_a, v_o, l_row, m_new) + v_v_b = kv_lds_to_regs.load_v(b_buf) + v_o, l_row = _subtile_tail(v_s_b, v_v_b, v_o, l_row, m_new) + m_row = m_new + next_args = [m_row, l_row] + v_o + [nn_a_buf, v_f_a, v_f_b] fx.rocdl.s_waitcnt(vmcnt=0, lgkmcnt=0, expcnt=0) rocdl.sched_barrier(0) rocdl.s_barrier() rocdl.sched_barrier(0) - - next_args = [m_row, l_row] + v_o + [nn_a_buf, v_f_a, v_f_b] loop_results = yield next_args m_row = loop_results[0] l_row = loop_results[1] diff --git a/kernels/attention/flash_attn_interface.py b/kernels/attention/flash_attn_interface.py index e1d96f99f3..888880d0ee 100644 --- a/kernels/attention/flash_attn_interface.py +++ b/kernels/attention/flash_attn_interface.py @@ -779,7 +779,7 @@ def _flydsl_flash_attn_paged( ) if paged_fp8: num_kv_pages = (skv + page_size - 1) // page_size - use_bn128 = fp8_head_dims in ((128, 128), (192, 128)) and num_kv_pages % 2 == 0 + use_bn128 = fp8_head_dims in ((128, 128), (192, 128), (192, 192)) and num_kv_pages % 2 == 0 paged_bn128_varlen = use_bn128 and B > 1 paged_setprio = dualwave_swp_setprio and D != 192 fp8_pv_segmented = _use_paged_fp8_segmented_pv(fp8_head_dims) diff --git a/kernels/attention/flash_attn_utils.py b/kernels/attention/flash_attn_utils.py index e4ce75c9bf..df7c8417d4 100644 --- a/kernels/attention/flash_attn_utils.py +++ b/kernels/attention/flash_attn_utils.py @@ -2290,6 +2290,10 @@ def _make_paged_dualwave_swp_fp8_traits( smem_k_tile_elems = smem_n_rpt * smem_d_rpt * smem_k_line_stride smem_v_tile_elems = smem_n_rpt * smem_d_rpt * smem_v_line_stride bn128 = bool(force_bn128) if force_bn128 is not None else (num_kv_splits <= 1) and (not varlen) + compact_v192 = paged and bn128 and value_head_dim == 192 + if compact_v192: + # Six K/V slots fit in 160 KiB only when K omits its unused D256 tail. + smem_k_tile_elems = block_n * head_dim bn128_pf = bn128 qreg = bn128_pf vdma = bn128_pf @@ -2330,7 +2334,7 @@ def _make_paged_dualwave_swp_fp8_traits( else: fp8_v_h1, fp8_v_h2 = value_head_dim, 0 fp8_v_bankpad = paged and (value_head_dim == fp8_v_segment_capacity or fp8_pv_segmented) - fp8_v_row_stride = block_n + 16 + fp8_v_row_stride = block_n if compact_v192 else block_n + 16 fp8_v_tile_bytes = (block_n // 8) * (head_dim // 16) * 128 if fp8_v_bankpad: fp8_v_tile_bytes = value_head_dim * fp8_v_row_stride @@ -5466,21 +5470,30 @@ def load_k(self, tile_start, buf_id, page_id=None): if page_id is None: page_id = self.load_page_id(tile_start) src_div = self.make_page_view(self.k_base_iter, page_id) - # Each wave owns two 64-lane octet rows for the D256-shaped LDS - # tile. The QK reader ignores the four groups beyond physical D192. - for d in range_constexpr(self.NUM_DMA_K): - src_elem = _vec_k_src_elem( - traits, - d, - self.wave_id_uni, - self.lane_in_warp, - self.kv_head_idx, - ) - lds_addr = ( - k_lds_byte_base + (self.wave_id_uni * traits.SMEM_D_RPT + d) * traits.WARP_SIZE * traits.KV_VEC_SIZE - ) - lds_addr = rocdl.readfirstlane(T.i32, as_mlir_value(fx.Int32(lds_addr))) - self.buffer_load_lds_128(src_div, lds_addr, src_elem, 0) + + def _load_vectorized(): + for d in range_constexpr(self.NUM_DMA_K): + src_elem = _vec_k_src_elem(traits, d, self.wave_id_uni, self.lane_in_warp, self.kv_head_idx) + lds_addr = ( + k_lds_byte_base + + (self.wave_id_uni * traits.SMEM_D_RPT + d) * traits.WARP_SIZE * traits.KV_VEC_SIZE + ) + lds_addr = rocdl.readfirstlane(T.i32, as_mlir_value(fx.Int32(lds_addr))) + self.buffer_load_lds_128(src_div, lds_addr, src_elem, 0) + + if const_expr(traits.BN128 and traits.V_HEAD_DIM == 192): + + @flyc.jit + def _load_compact(): + copy_waves = (traits.BLOCK_N * traits.HEAD_DIM) // ( + self.NUM_DMA_K * traits.WARP_SIZE * traits.KV_VEC_SIZE + ) + if self.wave_id_uni < fx.Index(copy_waves): + _load_vectorized() + + _load_compact() + else: + _load_vectorized() return rows_per_wave = -(-traits.BLOCK_N // traits.NUM_WAVES) @@ -5516,7 +5529,9 @@ def _run(): _run() def load_v(self, tile_start, buf_id, page_id=None): - if const_expr(self.traits.KV_VECTORIZED and self.traits.FP8_PV_SEGMENTED): + if const_expr(self.traits.KV_VECTORIZED and self.traits.FP8_PV_SEGMENTED and self.traits.BN128): + self._stage_v_fp8_vectorized_bankpad(tile_start, buf_id, page_id=page_id) + elif const_expr(self.traits.KV_VECTORIZED and self.traits.FP8_PV_SEGMENTED): self._stage_v_fp8_vectorized_segmented(tile_start, buf_id, page_id=page_id) elif const_expr(self.traits.KV_VECTORIZED and self.traits.FP8_V_BANKPAD): self._stage_v_fp8_vectorized_bankpad(tile_start, buf_id, page_id=page_id) @@ -5749,7 +5764,7 @@ def _stage_v_fp8_vectorized_bankpad(self, tile_start, buf_id, page_id=None): self._store_v_fp8_vectorized_bankpad(src_i32x4, buf_id) def _load_v_fp8_vectorized_bankpad_source(self, tile_start, page_id=None): - """Issue one vectorized paged-V load for the bank-padded path.""" + """Issue vectorized V loads for the paged FP8 LDS layout.""" traits = self.traits if page_id is None: page_id = self.load_page_id(tile_start) @@ -5766,19 +5781,42 @@ def _load_v_fp8_vectorized_bankpad_source(self, tile_start, page_id=None): + n_group16 * traits.V_HEAD_DIM * traits.KV_VEC_SIZE + d_col * traits.KV_VEC_SIZE ) - return self.buffer_load_fp8x16(src_div, src_elem) + source = self.buffer_load_fp8x16(src_div, src_elem) + if const_expr(traits.BN128 and traits.FP8_PV_SEGMENTED): + tail = self.buffer_load_fp8x16(src_div, src_elem + fx.Index(traits.FP8_V_H1 * traits.KV_VEC_SIZE)) + source = Vec(source).shuffle(Vec(tail), [0, 1, 2, 3, 4, 5, 6, 7]).ir_value() + return source def _store_v_fp8_vectorized_bankpad(self, src_i32x4, buf_id): - """Permute a prefetched V vector and store it into bank-padded LDS.""" + """Permute prefetched V vectors into the paged FP8 LDS layout.""" + if const_expr(self.traits.BN128 and self.traits.FP8_PV_SEGMENTED): + source = Vec(src_i32x4) + prefix = source.shuffle(source, [0, 1, 2, 3]).ir_value() + tail = source.shuffle(source, [4, 5, 6, 7]).ir_value() + self._store_v_fp8_vectorized_bankpad_segment(prefix, buf_id, 0) + + @flyc.jit + def _store_tail(): + if self.wave_id_uni < fx.Index(self.traits.FP8_V_H2 // self.traits.KV_VEC_SIZE): + self._store_v_fp8_vectorized_bankpad_segment(tail, buf_id, self.traits.FP8_V_H1) + + _store_tail() + else: + self._store_v_fp8_vectorized_bankpad_segment(src_i32x4, buf_id, 0) + + def _store_v_fp8_vectorized_bankpad_segment(self, src_i32x4, buf_id, d_offset): traits = self.traits row_stride = traits.FP8_V_ROW_STRIDE v_tile_bytes = traits.V_HEAD_DIM * row_stride aligned_base = ((self.lds_vt_base_idx + fx.Index(127)) // fx.Index(128)) * fx.Index(128) dst_base = aligned_base + fx.Index(buf_id * v_tile_bytes) n_group16 = self.lane_in_warp // fx.Index(16) - d_col = self.wave_id * fx.Index(16) + self.lane_in_warp % fx.Index(16) + d_col = fx.Index(d_offset) + self.wave_id * fx.Index(16) + self.lane_in_warp % fx.Index(16) reordered = self._permute_v_fp8_vectorized(src_i32x4) dest_group16 = (n_group16 % fx.Index(2)) * fx.Index(2) + n_group16 // fx.Index(2) + if const_expr(traits.BN128 and traits.V_HEAD_DIM == 192): + # XOR 16-byte groups instead of padding each row in the six-slot ring. + dest_group16 = dest_group16 ^ ((d_col // fx.Index(4)) % fx.Index(4)) dst_byte = dst_base + d_col * fx.Index(row_stride) + dest_group16 * fx.Index(16) ptr = buffer_ops.get_element_ptr( self.lds_vt_base_ptr, @@ -5911,7 +5949,7 @@ def _read_strip(key): return (_read_strip(n_lo), _read_strip(n_hi)) def load_v(self, buf_id): - if const_expr(self.traits.FP8_PV_SEGMENTED): + if const_expr(self.traits.FP8_PV_SEGMENTED and not self.traits.BN128): return fx.Index(buf_id) if const_expr(self.traits.FP8_PV): if const_expr(self.traits.KV_VECTORIZED and self.traits.FP8_V_BANKPAD): @@ -5940,7 +5978,7 @@ def load_v(self, buf_id): return packs def _load_v_fp8_vectorized_bankpad(self, buf_id): - """Read the padded D-major V tile into the FP8 MFMA A fragment.""" + """Read the D-major V tile into the FP8 MFMA A fragment.""" traits = self.traits row_stride = traits.FP8_V_ROW_STRIDE v_tile_bytes = traits.V_HEAD_DIM * row_stride @@ -5950,11 +5988,14 @@ def _load_v_fp8_vectorized_bankpad(self, buf_id): packs = [[None] * traits.D_CHUNKS for _ in range(4)] for dc in range_constexpr(traits.D_CHUNKS): d_row = fx.Index(dc * traits.D_CHUNK) + self.lane_mod_32 - row_base = tile_base + d_row * fx.Index(row_stride) + token_base + row_base = tile_base + d_row * fx.Index(row_stride) for half in range_constexpr(2): + token_offset = token_base + fx.Index(half * 16) + if const_expr(traits.BN128 and traits.V_HEAD_DIM == 192): + token_offset = token_offset ^ (((d_row // fx.Index(4)) % fx.Index(4)) * fx.Index(16)) ptr = buffer_ops.get_element_ptr( self.lds_vt_base_ptr, - byte_offset=fx.Int32(row_base + fx.Index(half * 16) - self.lds_vt_base_idx), + byte_offset=fx.Int32(row_base + token_offset - self.lds_vt_base_idx), elem_type=T.i8, ) pair = Vec(llvm.LoadOp(Vec.make_type(2, fx.Int64), ptr, alignment=16).result, (2,), fx.Int64) diff --git a/tests/kernels/test_flash_attn_fwd.py b/tests/kernels/test_flash_attn_fwd.py index 9a5586e9f6..be97f8ec5f 100644 --- a/tests/kernels/test_flash_attn_fwd.py +++ b/tests/kernels/test_flash_attn_fwd.py @@ -4603,6 +4603,7 @@ def test_paged_fp8_asymmetric_value_matches_torch( query_lengths, kv_lengths, block_table_rows, + num_kv_heads=1, ): """Packed causal FP8 page-64 attention supports native Q/K and V widths.""" if len(query_lengths) == 1 and force_internal_copies and head_dim == 192: @@ -4620,9 +4621,11 @@ def test_paged_fp8_asymmetric_value_matches_torch( num_pages = max(max(row) for row in block_table_rows) + 1 query, query_descale = quantize_per_tensor_fp8(torch.randn(sum(query_lengths), 16, head_dim, device="cuda") * 0.2) - key, key_descale = quantize_per_tensor_fp8(torch.randn(num_pages, 1, head_dim // 16, 64, 16, device="cuda") * 0.2) + key, key_descale = quantize_per_tensor_fp8( + torch.randn(num_pages, num_kv_heads, head_dim // 16, 64, 16, device="cuda") * 0.2 + ) value, value_descale = quantize_per_tensor_fp8( - torch.randn(num_pages, 1, 4, value_head_dim, 16, device="cuda") * 0.2 + torch.randn(num_pages, num_kv_heads, 4, value_head_dim, 16, device="cuda") * 0.2 ) if force_internal_copies: @@ -4644,7 +4647,7 @@ def noncontiguous_copy(tensor): call_kwargs = dict( causal=True, - num_kv_heads=1, + num_kv_heads=num_kv_heads, cu_seqlens_q=q_indptr, cu_seqlens_kv=kv_indptr, max_seqlen_q=max(query_lengths), @@ -4707,12 +4710,13 @@ def noncontiguous_copy(tensor): query_batch = query[q_indptr[batch_idx] : q_indptr[batch_idx + 1]].float() * query_descale physical_pages = block_table[batch_idx].long() key_batch = ( - key[physical_pages].permute(0, 3, 1, 2, 4).reshape(-1, 1, head_dim)[:kv_length].float() * key_descale - ).expand(-1, 16, -1) + key[physical_pages].permute(0, 3, 1, 2, 4).reshape(-1, num_kv_heads, head_dim)[:kv_length].float() + * key_descale + ).repeat_interleave(16 // num_kv_heads, dim=1) value_batch = ( - value[physical_pages].permute(0, 2, 4, 1, 3).reshape(-1, 1, value_head_dim)[:kv_length].float() + value[physical_pages].permute(0, 2, 4, 1, 3).reshape(-1, num_kv_heads, value_head_dim)[:kv_length].float() * value_descale - ).expand(-1, 16, -1) + ).repeat_interleave(16 // num_kv_heads, dim=1) logits = torch.einsum("qhd,khd->hqk", query_batch, key_batch) / math.sqrt(head_dim) query_positions = torch.arange(query_length, device="cuda") key_positions = torch.arange(kv_length, device="cuda") @@ -4726,6 +4730,27 @@ def noncontiguous_copy(tensor): torch.testing.assert_close(actual, expected, rtol=2.0e-2, atol=2.0e-2) +@_requires_gfx950 +@pytest.mark.parametrize("value_head_dim", [128, 192]) +@pytest.mark.parametrize("num_kv_heads", [2, 4]) +def test_paged_fp8_d192_bn128_multiple_kv_heads(value_head_dim, num_kv_heads): + """Compact K tiles and the V192 tail preserve head and ragged-page boundaries.""" + test_paged_fp8_asymmetric_value_matches_torch( + head_dim=192, + value_head_dim=value_head_dim, + use_non_default_stream=False, + force_internal_copies=False, + query_lengths=[300, 256, 63], + kv_lengths=[512, 384, 256], + block_table_rows=[ + [7, 0, 6, 1, 5, 2, 4, 3], + [13, 8, 12, 9, 11, 10, 0, 0], + [17, 14, 16, 15, 0, 0, 0, 0], + ], + num_kv_heads=num_kv_heads, + ) + + @_requires_gfx950 def test_paged_fp8_d128_bn128_ragged_multiblock_matches_torch(): """The multi-batch BN128 path handles distinct causal offsets and inactive q-blocks.""" From cfbdfb5657c5e58a1da02dcc38f1143fc11c1b20 Mon Sep 17 00:00:00 2001 From: Xiake Sun Date: Wed, 9 Sep 2026 03:04:25 +0000 Subject: [PATCH 37/59] refactor(fmha): specialize paged FP8 layout contracts - Derive native value staging from the supported head dimensions. - Remove obsolete BF16 V modes, unreferenced staging helpers, and copied trait state. - Drop redundant dimension guards and unused Q LDS scaffolding. - Preserve public attention behavior and native instruction schedules. Signed-off-by: Xiake Sun --- .../attention/flash_attn_fp8_paged_gfx950.py | 35 +- kernels/attention/flash_attn_interface.py | 12 +- kernels/attention/flash_attn_utils.py | 631 +++--------------- tests/kernels/test_flash_attn_fwd.py | 8 - 4 files changed, 83 insertions(+), 603 deletions(-) diff --git a/kernels/attention/flash_attn_fp8_paged_gfx950.py b/kernels/attention/flash_attn_fp8_paged_gfx950.py index 0b514b8a9e..edd9f06f16 100644 --- a/kernels/attention/flash_attn_fp8_paged_gfx950.py +++ b/kernels/attention/flash_attn_fp8_paged_gfx950.py @@ -47,7 +47,6 @@ def build_flash_attn_paged_fp8_module( kv_cache_layout="linear", paged_bn128=False, paged_bn128_varlen=False, - fp8_pv_segmented=False, batch_interleave_group=1, ): """Build the gfx950 packed-varlen page-64 FP8 attention launcher.""" @@ -81,15 +80,12 @@ def build_flash_attn_paged_fp8_module( raise ValueError(f"batch_interleave_group must be positive, got {batch_interleave_group}") if batch_interleave_group > 1 and (paged_bn128 or head_dim != 192): raise ValueError("batch interleaving is supported only by the generic paged D192 kernel") - if paged_bn128 and (head_dim, value_head_dim) not in ((128, 128), (192, 128), (192, 192)): - raise RuntimeError("paged BN128 requires supported FP8 head dimensions") assert num_heads % num_kv_heads == 0 traits = _make_paged_dualwave_swp_fp8_traits( num_heads, num_kv_heads, head_dim, value_head_dim=value_head_dim, - causal=causal, waves_per_eu=waves_per_eu, daz=daz, dualwave_swp_lazy_rescale=dualwave_swp_lazy_rescale, @@ -97,13 +93,8 @@ def build_flash_attn_paged_fp8_module( dualwave_swp_setprio=dualwave_swp_setprio, dualwave_swp_debug_lazy_counts=dualwave_swp_debug_lazy_counts, dualwave_swp_enable_stagger=dualwave_swp_enable_stagger, - num_kv_splits=num_kv_splits, - varlen=varlen if paged_bn128_varlen else False if paged_bn128 else varlen, - cross_seqlen=cross_seqlen, - paged=paged, - kv_cache_layout=kv_cache_layout, - fp8_pv_segmented=fp8_pv_segmented, - force_bn128=paged_bn128, + varlen=not paged_bn128 or paged_bn128_varlen, + bn128=paged_bn128, batch_interleave_group=batch_interleave_group, ) BLOCK_M = traits.BLOCK_M @@ -118,22 +109,11 @@ def build_flash_attn_paged_fp8_module( DEFAULT_STRIDE_KV_N = traits.DEFAULT_STRIDE_KV_N _dualwave_swp_fp8_cache_tag = traits.cache_tag _lds_elem_dtype = dtype_to_elem_type(traits.DTYPE_STR) - _q_lds_elems = BLOCK_M * HEAD_DIM if traits.QLDS else 16 - if PAGED_BN128: - - @fx.struct - class SharedStorage: - kv: fx.Array[_lds_elem_dtype, traits.LDS_KV_TOTAL_SIZE, 16] - vt: fx.Array[fx.BFloat16, traits.VT_BF16_TOTAL, 16] - q: fx.Array[_lds_elem_dtype, _q_lds_elems, 16] - - else: - - @fx.struct - class SharedStorage: - kv: fx.Array[_lds_elem_dtype, traits.LDS_KV_TOTAL_SIZE, 16] - vt: fx.Array[fx.BFloat16, traits.VT_BF16_TOTAL, 16] + @fx.struct + class SharedStorage: + kv: fx.Array[_lds_elem_dtype, traits.LDS_KV_TOTAL_SIZE, 16] + vt: fx.Array[fx.BFloat16, traits.VT_BF16_TOTAL, 16] # BN128: two BLOCK_N=64 KV tiles per iteration, one merged softmax correction. @flyc.kernel(known_block_size=[BLOCK_SIZE, 1, 1]) @@ -200,7 +180,6 @@ def flash_attn_paged_fp8_bn128_kernel( ctx.init_descale() ctx.init_tile_bounds() - q_loader = DualwaveFp8QLoader(ctx) gemm_helper = DualwaveFp8GemmHelper(ctx) softmax_helper = DualwaveFp8SoftmaxHelper(ctx) kv_gmem_to_lds = DualwaveFp8KvGmemToLdsLoader(ctx) @@ -247,8 +226,6 @@ def _merge_tile_max(v_s_a, v_s_b): page_t0, page_t1 = ctx.load_page_id_pair(t0 * BN) kv_gmem_to_lds.load_k(t0 * BN, t0 % fx.Index(NPF), page_id=page_t0) - if const_expr(traits.QLDS): - q_loader.stage_q_to_lds() fx.rocdl.s_waitcnt(vmcnt=0, lgkmcnt=0, expcnt=0) rocdl.sched_barrier(0) rocdl.s_barrier() diff --git a/kernels/attention/flash_attn_interface.py b/kernels/attention/flash_attn_interface.py index 888880d0ee..2c4d9b82cd 100644 --- a/kernels/attention/flash_attn_interface.py +++ b/kernels/attention/flash_attn_interface.py @@ -479,7 +479,6 @@ def _build_paged_fp8( enable_stagger: bool, use_bn128: bool, paged_bn128_varlen: bool, - fp8_pv_segmented: bool, batch_interleave_group: int, ): """Build the gfx950 packed-varlen, vectorized page-64 FP8 launcher.""" @@ -504,7 +503,6 @@ def _build_paged_fp8( kv_cache_layout="vectorized", paged_bn128=use_bn128, paged_bn128_varlen=paged_bn128_varlen, - fp8_pv_segmented=fp8_pv_segmented, batch_interleave_group=batch_interleave_group, ) @@ -518,10 +516,6 @@ def _build_paged_fp8( _PAGED_FP8_V192_BATCH_INTERLEAVE_MAX_BATCH = 16 -def _use_paged_fp8_segmented_pv(head_dims: tuple[int, int]) -> bool: - return head_dims == (192, 192) - - def _paged_fp8_batch_interleave_group(batch_size: int, head_dims: tuple[int, int]) -> int: """Choose a divisor-sized batch group for the tuned MiMo D192 launch.""" if head_dims[0] != 192 or batch_size <= 1: @@ -779,13 +773,12 @@ def _flydsl_flash_attn_paged( ) if paged_fp8: num_kv_pages = (skv + page_size - 1) // page_size - use_bn128 = fp8_head_dims in ((128, 128), (192, 128), (192, 192)) and num_kv_pages % 2 == 0 + use_bn128 = num_kv_pages % 2 == 0 paged_bn128_varlen = use_bn128 and B > 1 paged_setprio = dualwave_swp_setprio and D != 192 - fp8_pv_segmented = _use_paged_fp8_segmented_pv(fp8_head_dims) # The BF16 phase shift exposes H2 staging and probability-pack # latency because segmented FP8 has only six wide P*V MFMAs. - paged_stagger = dualwave_swp_enable_stagger and not fp8_pv_segmented + paged_stagger = dualwave_swp_enable_stagger and value_head_dim == 128 exe = _build_paged_fp8( num_heads=H, num_kv_heads=num_kv_heads, @@ -798,7 +791,6 @@ def _flydsl_flash_attn_paged( enable_stagger=paged_stagger, use_bn128=use_bn128, paged_bn128_varlen=paged_bn128_varlen, - fp8_pv_segmented=fp8_pv_segmented, batch_interleave_group=( _paged_fp8_batch_interleave_group(B, fp8_head_dims) if not use_bn128 and H == 16 and num_kv_heads == 1 diff --git a/kernels/attention/flash_attn_utils.py b/kernels/attention/flash_attn_utils.py index df7c8417d4..41165adf89 100644 --- a/kernels/attention/flash_attn_utils.py +++ b/kernels/attention/flash_attn_utils.py @@ -1833,7 +1833,6 @@ class DualwaveSwpFp8Traits: PAGED: bool = False KV_VECTORIZED: bool = False QREG: bool = True - FP8_PV: bool = True FP8_PV_DIRECT: bool = True FP8_PV_SEGMENTED: bool = False @@ -2063,121 +2062,82 @@ def dualwave_fp8_dma_per_iter(traits): @dataclass(frozen=True) class PagedDualwaveSwpFp8Traits: - """Pure compile-time tile/layout constants for the gfx950 DUALWAVE_SWP fp8 kernel. + """Layouts for native page-64 FP8 QK/PV with register-resident Q.""" - QK uses WIDE 32x32x64 ``mfma_scale``. P*V either uses the HIPREC path - (FP8 V dequantized into BF16 LDS scratch) or the D128 native-FP8 path. The - ``*_BF`` fields describe the HIPREC layout; ``ELEM_BYTES`` is 1 because - Q/K/V are FP8. - """ - - BLOCK_M: int - BLOCK_N: int - K_SUB_N: int - WARP_SIZE: int - NUM_WAVES: int - BLOCK_SIZE: int - ROWS_PER_WAVE: int HEAD_DIM: int V_HEAD_DIM: int - D_CHUNK: int D_CHUNKS: int - PV_K_STEPS: int NUM_HEADS_Q: int NUM_HEADS_KV: int GQA_GROUP_SIZE: int - CAUSAL: bool - DTYPE_STR: str WAVES_PER_EU: int DAZ: bool DUALWAVE_SWP_LAZY_RESCALE: bool DUALWAVE_SWP_SETPRIO: bool DUALWAVE_SWP_DEBUG_LAZY_COUNTS: bool DUALWAVE_SWP_ENABLE_STAGGER: bool - NUM_KV_SPLITS: int - SPLITK: bool VARLEN: bool - CROSS_SEQLEN: bool - FP8_PV: bool FP8_PV_DIRECT: bool - FP8_V_BANKPAD: bool FP8_PV_SEGMENTED: bool FP8_V_H1: int FP8_V_H2: int BN128: bool - BN128_PF: bool QREG: bool - VDMA: bool - PAGED: bool - KV_CACHE_LAYOUT: str - KV_VECTORIZED: bool - PAGE_SIZE: int DEFAULT_STRIDE_Q_N: int DEFAULT_STRIDE_KV_N: int - DMA_BYTES: int - ELEM_BYTES: int - OUT_ELEM_BYTES: int - D_128B_SIZE: int - VEC_KV: int - KV_VEC_SIZE: int - LANE_SPLIT_KV: int - SMEM_N_PER_WAVE: int - SMEM_N_RPT: int SMEM_D_RPT: int - SMEM_K_LINE_STRIDE: int SMEM_K_TILE_ELEMS: int NUM_PREFETCH_K: int DUALWAVE_SWP_KV_PER_BUFFER: int LDS_KV_TOTAL_SIZE: int - DUALWAVE_SWP_K_BUF_BASE: tuple[int, int] - DUALWAVE_SWP_V_BUF_BASE: tuple[int, int] - # bf16 vt scratch layout (HIPREC V dequant target + transpose read strides). - EB_BF: int - D128_BF: int - VEC_BF: int - SDRPT_BF: int - SNRPT_BF: int - VLS_BF: int - VEC_V_ROW_STRIDE: int - VT_BF16_ELEMS: int + DUALWAVE_SWP_K_BUF_BASE: tuple[int, ...] + DUALWAVE_SWP_V_BUF_BASE: tuple[int, ...] VT_BF16_TOTAL: int FP8_V_ROW_STRIDE: int - URV_GRPK_BF: int - URV_GRP_N_BF: int - URV_LANE_LO_BF: int - URV_LANE_HI_BF: int - URV_STEPK_BF: int - URV_DC_AXIS0_BF: int - URV_DC_AXIS1_BF: int - URV_I5_BF: int DUALWAVE_SWP_RESCALE_THRESHOLD: float - SCHED_MFMA_MASK: int - SCHED_VALU_MASK: int - SCHED_EXP_MASK: int - SCHED_DS_READ_MASK: int - LDS_SCOPE_NAMES: tuple[str, str, str, str] - NEG_INF_F32_BITS: int - LGKMCNT_0_ONLY: int - XCD_SWIZZLE: bool = False BATCH_INTERLEAVE_GROUP: int = 1 + BLOCK_M: int = 256 + BLOCK_N: int = 64 + WARP_SIZE: int = 64 + NUM_WAVES: int = 8 + BLOCK_SIZE: int = 512 + ROWS_PER_WAVE: int = 32 + D_CHUNK: int = 32 + PV_K_STEPS: int = 2 + CAUSAL: bool = True + DTYPE_STR: str = "fp8" + NUM_KV_SPLITS: int = 1 + SPLITK: bool = False + CROSS_SEQLEN: bool = True + PAGED: bool = True + KV_VECTORIZED: bool = True + PAGE_SIZE: int = 64 + DMA_BYTES: int = 16 + ELEM_BYTES: int = 1 + OUT_ELEM_BYTES: int = 2 + VEC_KV: int = 16 + KV_VEC_SIZE: int = 16 + LANE_SPLIT_KV: int = 8 + SCHED_MFMA_MASK: int = 0x008 + SCHED_VALU_MASK: int = 0x002 + SCHED_EXP_MASK: int = 0x400 + SCHED_DS_READ_MASK: int = 0x100 + NEG_INF_F32_BITS: int = 0xFF800000 + XCD_SWIZZLE: bool = False + QLDS: bool = False @property def HEAD_DIM_V(self): return self.V_HEAD_DIM - @property - def QLDS(self): - return False - @property def cache_tag(self): return ( + "paged_fp8_native", self.NUM_HEADS_Q, self.NUM_HEADS_KV, self.HEAD_DIM, self.V_HEAD_DIM, - self.CAUSAL, - self.DTYPE_STR, self.WAVES_PER_EU, self.DAZ, self.DUALWAVE_SWP_LAZY_RESCALE, @@ -2185,32 +2145,10 @@ def cache_tag(self): self.DUALWAVE_SWP_SETPRIO, self.DUALWAVE_SWP_DEBUG_LAZY_COUNTS, self.DUALWAVE_SWP_ENABLE_STAGGER, - self.NUM_KV_SPLITS, - self.SPLITK, self.VARLEN, - self.CROSS_SEQLEN, - self.PAGED, - self.KV_CACHE_LAYOUT, - self.KV_VECTORIZED, - "fp8_wide_qk_hiprec_pv_paged_vec5d_v2" if self.PAGED else "fp8_wide_qk_hiprec_pv", - self.ELEM_BYTES, - self.OUT_ELEM_BYTES, - self.LANE_SPLIT_KV, - self.VT_BF16_ELEMS, - self.VT_BF16_TOTAL, - self.FP8_V_ROW_STRIDE, - self.FP8_PV, - self.FP8_PV_DIRECT, - self.FP8_V_BANKPAD, - self.FP8_PV_SEGMENTED, - self.FP8_V_H1, - self.FP8_V_H2, - self.NUM_PREFETCH_K, self.BN128, - self.BN128_PF, - self.QREG, - self.VDMA, - self.XCD_SWIZZLE, + self.NUM_PREFETCH_K, + self.FP8_V_ROW_STRIDE, self.BATCH_INTERLEAVE_GROUP, ) @@ -2230,223 +2168,66 @@ def _make_paged_dualwave_swp_fp8_traits( num_heads, num_kv_heads, head_dim, + value_head_dim, rescale_threshold, - value_head_dim=None, - causal=True, waves_per_eu=2, daz=True, dualwave_swp_lazy_rescale=True, dualwave_swp_setprio=True, dualwave_swp_debug_lazy_counts=False, dualwave_swp_enable_stagger=True, - num_kv_splits=1, varlen=False, - cross_seqlen=False, - xcd_swizzle=False, - paged=False, - kv_cache_layout="linear", - fp8_pv_segmented=False, - force_bn128=None, + bn128=False, batch_interleave_group=1, ): - """Build gfx950 DUALWAVE_SWP fp8 compile-time layout traits (dtype fixed to fp8).""" - # Tile shape and wave geometry follow the gfx950 dual-wave 8-wave CTA. - block_m = 256 + """Build layouts after the dedicated builder validates the paged contract.""" block_n = 64 - k_sub_n = 32 - warp_size = 64 - num_waves = 8 - block_size = num_waves * warp_size - rows_per_wave = 32 - - if value_head_dim is None: - value_head_dim = head_dim - d_chunk = 32 - d_chunks = value_head_dim // d_chunk - pv_k_step = 16 - pv_k_steps = k_sub_n // pv_k_step - - gqa_group_size = num_heads // num_kv_heads - default_stride_q_n = num_heads * head_dim - default_stride_kv_n = num_kv_heads * head_dim - - # fp8: Q/K/V are 1B; O is bf16 (2B). ELEM_BYTES=1 drives the fp8 address math. - elem_bytes = 1 - out_elem_bytes = 2 - d_128b_size = 128 // elem_bytes - vec_kv = 16 // elem_bytes - lane_split_kv = 8 - smem_linear_wave = warp_size * 16 // elem_bytes - smem_n_per_wave = smem_linear_wave // d_128b_size - smem_n_rpt = block_n // smem_n_per_wave - # D192 needs a second 128-byte K repeat for the D256-shaped LDS tile. The - # QK reader consumes only the 12 physical D192 groups; the remaining four - # DMA groups are unused padding. - smem_d_rpt = (head_dim + d_128b_size - 1) // d_128b_size - smem_k_pad = 16 // elem_bytes - smem_v_pad = 64 // elem_bytes - smem_k_line_stride = smem_linear_wave + smem_k_pad - smem_v_line_stride = smem_linear_wave + smem_v_pad - smem_k_tile_elems = smem_n_rpt * smem_d_rpt * smem_k_line_stride - smem_v_tile_elems = smem_n_rpt * smem_d_rpt * smem_v_line_stride - bn128 = bool(force_bn128) if force_bn128 is not None else (num_kv_splits <= 1) and (not varlen) - compact_v192 = paged and bn128 and value_head_dim == 192 + smem_d_rpt = (head_dim + 127) // 128 + # Preserve the generic kernel's D256-shaped K slot and spacing. + smem_k_tile_elems = 8 * smem_d_rpt * 1040 + smem_v_tile_elems = 8 * smem_d_rpt * 1088 + compact_v192 = bn128 and value_head_dim == 192 if compact_v192: # Six K/V slots fit in 160 KiB only when K omits its unused D256 tail. smem_k_tile_elems = block_n * head_dim - bn128_pf = bn128 - qreg = bn128_pf - vdma = bn128_pf - deep_ring = bn128 - num_prefetch_k = (6 if bn128_pf else 4) if deep_ring else 2 - if paged and bn128 and head_dim == 128: - num_prefetch_k = 8 - if bn128_pf: - dualwave_swp_kv_per_buffer = smem_k_tile_elems - else: - dualwave_swp_kv_per_buffer = smem_k_tile_elems + smem_v_tile_elems - lds_kv_total_size = num_prefetch_k * dualwave_swp_kv_per_buffer - dualwave_swp_k_buf_base = tuple(i * dualwave_swp_kv_per_buffer for i in range(num_prefetch_k)) - dualwave_swp_v_buf_base = tuple(smem_k_tile_elems + i * dualwave_swp_kv_per_buffer for i in range(num_prefetch_k)) - - # bf16 vt scratch layout: HIPREC dequantizes fp8 V into these positions so the - # proven bf16 V transpose read (ds_read_tr16) + bf16 PV MMA are reused unchanged. - eb_bf = 2 - d128_bf = 128 // eb_bf - vec_bf = 16 // eb_bf - slw_bf = warp_size * 16 // eb_bf - snrpt_bf = block_n // (slw_bf // d128_bf) - sdrpt_bf = value_head_dim // d128_bf - vls_bf = slw_bf + 64 // eb_bf - vt_bf16_elems = snrpt_bf * sdrpt_bf * vls_bf - paged = bool(paged) - token_groups16 = block_n // vec_kv - fp8_v_segment_capacity = block_size // token_groups16 - fp8_pv_segmented = bool( - fp8_pv_segmented - and paged - and value_head_dim > fp8_v_segment_capacity - and value_head_dim <= 2 * fp8_v_segment_capacity - and value_head_dim % d_chunk == 0 - ) - if fp8_pv_segmented: - fp8_v_h1, fp8_v_h2 = _factor_fp8_pv_head_dim(value_head_dim, fp8_v_segment_capacity, d_chunk) - else: - fp8_v_h1, fp8_v_h2 = value_head_dim, 0 - fp8_v_bankpad = paged and (value_head_dim == fp8_v_segment_capacity or fp8_pv_segmented) + num_prefetch_k = (8 if head_dim == 128 else 6) if bn128 else 2 + slot_elems = smem_k_tile_elems if bn128 else smem_k_tile_elems + smem_v_tile_elems + fp8_v_h1, fp8_v_h2 = _factor_fp8_pv_head_dim(value_head_dim, 128, 32) fp8_v_row_stride = block_n if compact_v192 else block_n + 16 - fp8_v_tile_bytes = (block_n // 8) * (head_dim // 16) * 128 - if fp8_v_bankpad: - fp8_v_tile_bytes = value_head_dim * fp8_v_row_stride - vt_bf16_total = num_prefetch_k * (fp8_v_tile_bytes // eb_bf) + 128 - elif bn128_pf: - vt_bf16_total = num_prefetch_k * (fp8_v_tile_bytes // eb_bf) + 128 - else: - vt_bf16_total = (2 if deep_ring else num_prefetch_k) * vt_bf16_elems - - splitk = num_kv_splits > 1 - kv_vectorized = paged and kv_cache_layout == "vectorized" - page_size = block_n - - # Paged V128 attention can consume its native FP8 V cache directly once - # the vectorized cache is transposed into the bank-padded FP8 P*V layout. - fp8_pv = ( - (paged and value_head_dim == fp8_v_segment_capacity) - or fp8_pv_segmented - or (not paged and os.getenv("FLYDSL_FA_FP8_PV", "0") == "1") - ) - fp8_pv_direct = bn128 and (not paged or fp8_pv) - if fp8_pv_direct: - fp8_pv = True + fp8_v_tile_bytes = value_head_dim * fp8_v_row_stride return PagedDualwaveSwpFp8Traits( - BLOCK_M=block_m, - BLOCK_N=block_n, - K_SUB_N=k_sub_n, - WARP_SIZE=warp_size, - NUM_WAVES=num_waves, - BLOCK_SIZE=block_size, - ROWS_PER_WAVE=rows_per_wave, HEAD_DIM=head_dim, V_HEAD_DIM=value_head_dim, - D_CHUNK=d_chunk, - D_CHUNKS=d_chunks, - PV_K_STEPS=pv_k_steps, + D_CHUNKS=value_head_dim // 32, NUM_HEADS_Q=num_heads, NUM_HEADS_KV=num_kv_heads, - GQA_GROUP_SIZE=gqa_group_size, - CAUSAL=causal, - DTYPE_STR="fp8", + GQA_GROUP_SIZE=num_heads // num_kv_heads, WAVES_PER_EU=waves_per_eu, DAZ=bool(daz), DUALWAVE_SWP_LAZY_RESCALE=bool(dualwave_swp_lazy_rescale), DUALWAVE_SWP_SETPRIO=bool(dualwave_swp_setprio), DUALWAVE_SWP_DEBUG_LAZY_COUNTS=bool(dualwave_swp_debug_lazy_counts), DUALWAVE_SWP_ENABLE_STAGGER=bool(dualwave_swp_enable_stagger), - NUM_KV_SPLITS=num_kv_splits, - SPLITK=splitk, + DUALWAVE_SWP_RESCALE_THRESHOLD=rescale_threshold, VARLEN=bool(varlen), - CROSS_SEQLEN=bool(cross_seqlen), - FP8_PV=fp8_pv, - FP8_PV_DIRECT=bool(fp8_pv_direct), - FP8_V_BANKPAD=bool(fp8_v_bankpad), - FP8_PV_SEGMENTED=fp8_pv_segmented, + FP8_PV_DIRECT=bool(bn128), + FP8_PV_SEGMENTED=fp8_v_h2 > 0, FP8_V_H1=fp8_v_h1, FP8_V_H2=fp8_v_h2, BN128=bool(bn128), - BN128_PF=bool(bn128_pf), - QREG=bool(qreg), - VDMA=bool(vdma), - PAGED=paged, - KV_CACHE_LAYOUT=kv_cache_layout, - KV_VECTORIZED=kv_vectorized, - PAGE_SIZE=page_size, - DEFAULT_STRIDE_Q_N=default_stride_q_n, - DEFAULT_STRIDE_KV_N=default_stride_kv_n, - DMA_BYTES=16, - ELEM_BYTES=elem_bytes, - OUT_ELEM_BYTES=out_elem_bytes, - D_128B_SIZE=d_128b_size, - VEC_KV=vec_kv, - KV_VEC_SIZE=vec_kv, - LANE_SPLIT_KV=lane_split_kv, - SMEM_N_PER_WAVE=smem_n_per_wave, - SMEM_N_RPT=smem_n_rpt, + QREG=bool(bn128), + DEFAULT_STRIDE_Q_N=num_heads * head_dim, + DEFAULT_STRIDE_KV_N=num_kv_heads * head_dim, SMEM_D_RPT=smem_d_rpt, - SMEM_K_LINE_STRIDE=smem_k_line_stride, SMEM_K_TILE_ELEMS=smem_k_tile_elems, NUM_PREFETCH_K=num_prefetch_k, - DUALWAVE_SWP_KV_PER_BUFFER=dualwave_swp_kv_per_buffer, - LDS_KV_TOTAL_SIZE=lds_kv_total_size, - DUALWAVE_SWP_K_BUF_BASE=dualwave_swp_k_buf_base, - DUALWAVE_SWP_V_BUF_BASE=dualwave_swp_v_buf_base, - EB_BF=eb_bf, - D128_BF=d128_bf, - VEC_BF=vec_bf, - SDRPT_BF=sdrpt_bf, - SNRPT_BF=snrpt_bf, - VLS_BF=vls_bf, - VEC_V_ROW_STRIDE=vls_bf, - VT_BF16_ELEMS=vt_bf16_elems, - VT_BF16_TOTAL=vt_bf16_total, + DUALWAVE_SWP_KV_PER_BUFFER=slot_elems, + LDS_KV_TOTAL_SIZE=num_prefetch_k * slot_elems, + DUALWAVE_SWP_K_BUF_BASE=tuple(i * slot_elems for i in range(num_prefetch_k)), + DUALWAVE_SWP_V_BUF_BASE=tuple(smem_k_tile_elems + i * slot_elems for i in range(num_prefetch_k)), + VT_BF16_TOTAL=num_prefetch_k * (fp8_v_tile_bytes // 2) + 128, FP8_V_ROW_STRIDE=fp8_v_row_stride, - URV_GRPK_BF=4 * vls_bf, - URV_GRP_N_BF=16, - URV_LANE_LO_BF=4, - URV_LANE_HI_BF=vls_bf, - URV_STEPK_BF=128, - URV_DC_AXIS0_BF=snrpt_bf * vls_bf, - URV_DC_AXIS1_BF=32, - URV_I5_BF=d128_bf, - DUALWAVE_SWP_RESCALE_THRESHOLD=rescale_threshold, - SCHED_MFMA_MASK=0x008, - SCHED_VALU_MASK=0x002, - SCHED_EXP_MASK=0x400, - SCHED_DS_READ_MASK=0x100, - LDS_SCOPE_NAMES=("lds_k0", "lds_k1", "lds_v0", "lds_v1"), - NEG_INF_F32_BITS=0xFF800000, - LGKMCNT_0_ONLY=0xC07F, - XCD_SWIZZLE=bool(xcd_swizzle), BATCH_INTERLEAVE_GROUP=int(batch_interleave_group), ) @@ -4807,7 +4588,6 @@ def init_types_and_constants(self): self.v2i32_type = Vec.make_type(2, fx.Int32) self.p_elem = fx.BFloat16 self.v4bf16_type = Vec.make_type(4, fx.BFloat16) - self.v8bf16_type = Vec.make_type(8, fx.BFloat16) if const_expr(traits.PAGED): self.NUM_DMA_K = traits.SMEM_D_RPT self.NUM_DMA_V = traits.SMEM_D_RPT @@ -4989,9 +4769,6 @@ def init_atoms_and_lds_ptrs(self): self.o_store_reg_128 = fx.make_rmem_tensor(fx.make_layout(4, 1), fx.Int32) # fp8 global->LDS DMA uses i8 destination typing; K/V LDS reads are byte-addressed. self.lds_ptr_ty = fx.PointerType.get(fx.Int8.ir_type, 2, traits.DMA_BYTES) - if const_expr(traits.PAGED): - self.bf16_mma_atom = fx.make_mma_atom(fx.rocdl.MFMA(32, 32, 16, fx.BFloat16)) - self.v_fp8_load64_atom = fx.make_copy_atom(fx.rocdl.BufferCopy64b(), fx.Int32) def init_descale(self): def _load_scale_scalar(tensor): @@ -5233,9 +5010,6 @@ class DualwaveFp8GemmHelper(DualwaveFp8KernelContext): def __init__(self, ctx): super().__init__(ctx) - def _mfma_acc_bf16(self, a_v8, b_v8, c_v16): - return fly.mma_atom_call_ssa([self.v16f32_type], self.bf16_mma_atom, a_v8, b_v8, c_v16) - def _v8bf16_to_f32(self, v8): f32 = Vec(llvm.FPExtOp(Vec.make_type(8, fx.Float32), as_mlir_value(v8)).result, (8,), fx.Float32) return [f32[i] for i in range_constexpr(8)] @@ -5246,7 +5020,7 @@ def _p_to_fp8_i32x8(self, v_p): for pk in (p_lo[0], p_lo[1], p_hi[0], p_hi[1]): f32 += self._v8bf16_to_f32(pk) packed = self._pack_fp8_i32x8(f32) - if const_expr(self.traits.PAGED and self.traits.FP8_V_BANKPAD): + if const_expr(self.traits.PAGED): words = Vec(packed, (8,), fx.Int32) return Vec.from_elements( [words[i] for i in (0, 2, 1, 3, 4, 6, 5, 7)], @@ -5254,12 +5028,6 @@ def _p_to_fp8_i32x8(self, v_p): ).ir_value() return packed - def _v_to_fp8_i32x8(self, v_v, dc): - f32 = [] - for step in range_constexpr(4): - f32 += self._v8bf16_to_f32(v_v[step][dc]) - return self._pack_fp8_i32x8(f32) - def _pv_fp8(self, v_p, v_v, v_o): p_fp8 = self._p_to_fp8_i32x8(v_p) for dc in range_constexpr(self.traits.D_CHUNKS): @@ -5404,18 +5172,7 @@ def qk(self, v_k, q_wide=None): def pv_step_k(self, step, v_p, v_v, v_o): if const_expr(self.traits.FP8_PV_SEGMENTED): return self._pv_step_fp8_segmented(step, v_p, v_v, v_o) - if const_expr(self.traits.FP8_PV): - return self._pv_step_fp8(step, v_p, v_v, v_o) - # HIPREC PV: P and V are both v8 bf16, accumulated by a bf16 MMA. - v_p_lo, v_p_hi = v_p - v_pk = v_v[step] - if const_expr(step < 2): - p_pk = v_p_lo[step] - else: - p_pk = v_p_hi[step - 2] - for dc in range_constexpr(self.traits.D_CHUNKS): - v_o[dc] = self._mfma_acc_bf16(v_pk[dc], p_pk, v_o[dc]) - return v_o + return self._pv_step_fp8(step, v_p, v_v, v_o) def cast_p_fp8_direct(self, v_p): lo_partial_list, hi_full = v_p @@ -5445,11 +5202,7 @@ def _pv_fp8_direct(self, p_fp8, v_v, v_o): def pv(self, v_p, v_v, v_o): if const_expr(self.traits.FP8_PV_DIRECT): return self._pv_fp8_direct(v_p, v_v, v_o) - if const_expr(self.traits.FP8_PV): - return self._pv_fp8(v_p, v_v, v_o) - for step in range_constexpr(4): - v_o = self.pv_step_k(step, v_p, v_v, v_o) - return v_o + return self._pv_fp8(v_p, v_v, v_o) class DualwaveFp8KvGmemToLdsLoader(DualwaveFp8KernelContext): @@ -5529,58 +5282,14 @@ def _run(): _run() def load_v(self, tile_start, buf_id, page_id=None): - if const_expr(self.traits.KV_VECTORIZED and self.traits.FP8_PV_SEGMENTED and self.traits.BN128): - self._stage_v_fp8_vectorized_bankpad(tile_start, buf_id, page_id=page_id) - elif const_expr(self.traits.KV_VECTORIZED and self.traits.FP8_PV_SEGMENTED): - self._stage_v_fp8_vectorized_segmented(tile_start, buf_id, page_id=page_id) - elif const_expr(self.traits.KV_VECTORIZED and self.traits.FP8_V_BANKPAD): - self._stage_v_fp8_vectorized_bankpad(tile_start, buf_id, page_id=page_id) - elif const_expr(self.traits.KV_VECTORIZED and self.traits.FP8_PV): - self._stage_v_fp8_vectorized_dense_layout(tile_start, buf_id, page_id=page_id) - elif const_expr(self.traits.KV_VECTORIZED): - self._stage_vt_dequant_fp8_vectorized(tile_start, buf_id, page_id=page_id) + if const_expr(self.traits.PAGED): + if const_expr(self.traits.FP8_PV_SEGMENTED and not self.traits.BN128): + self._stage_v_fp8_vectorized_segmented(tile_start, buf_id, page_id=page_id) + else: + self._stage_v_fp8_vectorized_bankpad(tile_start, buf_id, page_id=page_id) else: self._stage_v_fp8_block_dma(tile_start, buf_id) - def zero_v_fp8_lds(self): - traits = self.traits - v_tile_bytes = (traits.BLOCK_N // 8) * (traits.HEAD_DIM // 16) * 128 - total = 2 * v_tile_bytes - aligned_base = ((self.lds_vt_base_idx + fx.Index(127)) // fx.Index(128)) * fx.Index(128) - zero = Vec.from_elements([fx.Int32(0) for _ in range_constexpr(4)], fx.Int32) - per = total // (traits.BLOCK_SIZE) # bytes per thread - for i in range_constexpr(per // 16): - off = aligned_base + self.tid * fx.Index(per) + fx.Index(i * 16) - p = buffer_ops.create_llvm_ptr(off, address_space=3) - llvm.StoreOp(as_mlir_value(zero), p, alignment=16) - - def _stage_v_fp8_block(self, tile_start, buf_id): - traits = self.traits - if const_expr(traits.VDMA): - return self._stage_v_fp8_block_dma(tile_start, buf_id) - v_tile_bytes = (traits.BLOCK_N // 8) * (traits.HEAD_DIM // 16) * 128 - buf_off = buf_id * v_tile_bytes - n = self.wave_id * fx.Index(8) + self.lane // fx.Index(8) - d_block = self.lane % fx.Index(8) - src_elem = ( - self.kv_gmem_elem_offset + n * self.stride_kv_n_v + d_block * fx.Index(16) + tile_start * self.stride_kv_n_v - ) - v16 = fly.copy_atom_call_ssa( - [Vec.make_type(4, fx.Int32)], self.load_atom_128, fx.slice(self.v_div, (None, fx.Int32(src_elem))) - ) - n_i = fx.Int32(n) - w16 = n_i % fx.Int32(16) - c_add = (w16 >= fx.Int32(4)) & (w16 < fx.Int32(8)) - c_sub = (w16 >= fx.Int32(8)) & (w16 < fx.Int32(12)) - dest_n = n_i + c_add.select(fx.Int32(4), fx.Int32(0)) - c_sub.select(fx.Int32(4), fx.Int32(0)) - dest_wave = fx.Index(dest_n // fx.Int32(8)) - dest_m = fx.Index(dest_n % fx.Int32(8)) - block = dest_wave * fx.Index(8) + self.lane % fx.Index(8) - aligned_base = ((self.lds_vt_base_idx + fx.Index(127)) // fx.Index(128)) * fx.Index(128) - byte_off = aligned_base + fx.Index(buf_off) + block * fx.Index(128) + fx.Index(16) * dest_m - lds_ptr = buffer_ops.create_llvm_ptr(byte_off, address_space=3) - llvm.StoreOp(as_mlir_value(Vec(v16)), lds_ptr, alignment=16) - def _stage_v_fp8_block_dma(self, tile_start, buf_id): traits = self.traits nbands = traits.HEAD_DIM_V // 16 @@ -5625,63 +5334,6 @@ def _run(): _run() - def _stage_vt_dequant_fp8_vectorized(self, tile_start, buf_id, page_id=None): - """Dequantize SHUFFLE V [H,4,D,16] into the bf16 vectorized LDS tile.""" - traits = self.traits - if page_id is None: - page_id = self.load_page_id(tile_start) - src_div = self.make_page_view(self.v_base_iter, page_id, is_value=True) - vt_buf = buf_id * traits.VT_BF16_ELEMS - token_groups16 = traits.BLOCK_N // traits.KV_VEC_SIZE - total_v16 = traits.V_HEAD_DIM * token_groups16 - passes = (total_v16 + traits.BLOCK_SIZE - 1) // traits.BLOCK_SIZE - - for pass_id in range_constexpr(passes): - flat = self.tid + fx.Index(pass_id * traits.BLOCK_SIZE) - - def _load_store_active(): - n_group16 = flat % fx.Index(token_groups16) - d_col = flat // fx.Index(token_groups16) - src_elem = ( - self.kv_head_idx * token_groups16 * traits.V_HEAD_DIM * traits.KV_VEC_SIZE - + n_group16 * traits.V_HEAD_DIM * traits.KV_VEC_SIZE - + d_col * traits.KV_VEC_SIZE - ) - v_i32x4 = self.buffer_load_fp8x16(src_div, src_elem) - v_words = Vec(v_i32x4, (4,), fx.Int32) - bf = [] - for w in range_constexpr(4): - word = as_mlir_value(fx.Int32(v_words[w])) - lo2 = Vec(rocdl.cvt_pk_f32_fp8(Vec.make_type(2, fx.Float32), word, False), (2,), fx.Float32) - hi2 = Vec(rocdl.cvt_pk_f32_fp8(Vec.make_type(2, fx.Float32), word, True), (2,), fx.Float32) - for e in (lo2[0], lo2[1], hi2[0], hi2[1]): - bf.append(fx.Float32(e) * self.vd_fp8) - - # Vectorized K applies sigma(n) while staging each 16-token - # group, so MFMA score slot n represents physical key - # sigma(n). Apply the same (involutive) permutation to V; - # otherwise selective attention pairs K[sigma(n)] with V[n] - # even though near-uniform low-amplitude tests appear valid. - bf = [bf[_sigma_k_tile_n(i)] for i in range_constexpr(16)] - - # The bf16 vectorized LDS layout groups eight tokens. Split - # the 16-token SHUFFLE vector into two aligned v8 stores. - row = d_col // fx.Index(traits.VEC_BF) - d_in_row = d_col % fx.Index(traits.VEC_BF) - for half in range_constexpr(2): - no = n_group16 * fx.Index(2) + fx.Index(half) - lane_slot = no * fx.Index(traits.VEC_BF) + d_in_row - dst_elem = vt_buf + row * fx.Index(traits.VLS_BF) + lane_slot * fx.Index(traits.VEC_BF) - v8bf = self.bf16_trunc_pack_v8(bf[half * 8 : (half + 1) * 8]) - lds_ptr = buffer_ops.get_element_ptr( - self.lds_vt_base_ptr, - byte_offset=fx.Int32(dst_elem * fx.Index(traits.EB_BF)), - elem_type=T.i8, - ) - llvm.StoreOp(as_mlir_value(v8bf), lds_ptr, alignment=16) - - scf_if_dispatch(flat < fx.Index(total_v16), _load_store_active) - def _stage_v_fp8_vectorized_segmented(self, tile_start, buf_id, page_id=None): if page_id is None: page_id = self.load_page_id(tile_start) @@ -5825,85 +5477,6 @@ def _store_v_fp8_vectorized_bankpad_segment(self, src_i32x4, buf_id, d_offset): ) llvm.StoreOp(as_mlir_value(reordered), ptr, alignment=16) - def _stage_v_fp8_vectorized_dense_layout(self, tile_start, buf_id, page_id=None): - """Transpose vectorized V into the LDS layout used by direct FP8 P*V.""" - traits = self.traits - if page_id is None: - page_id = self.load_page_id(tile_start) - src_div = self.make_page_view(self.v_base_iter, page_id, is_value=True) - v_tile_bytes = (traits.BLOCK_N // 8) * (traits.V_HEAD_DIM // 16) * 128 - aligned_base = ((self.lds_vt_base_idx + fx.Index(127)) // fx.Index(128)) * fx.Index(128) - dst_base = aligned_base + fx.Index(buf_id * v_tile_bytes) - token_groups16 = traits.BLOCK_N // traits.KV_VEC_SIZE - total_v16 = traits.V_HEAD_DIM * token_groups16 - passes = (total_v16 + traits.BLOCK_SIZE - 1) // traits.BLOCK_SIZE - - for pass_id in range_constexpr(passes): - flat = self.tid + fx.Index(pass_id * traits.BLOCK_SIZE) - - def _load_store_active(): - n_group16 = flat % fx.Index(token_groups16) - d_col = flat // fx.Index(token_groups16) - src_elem = ( - self.kv_head_idx * token_groups16 * traits.V_HEAD_DIM * traits.KV_VEC_SIZE - + n_group16 * traits.V_HEAD_DIM * traits.KV_VEC_SIZE - + d_col * traits.KV_VEC_SIZE - ) - src_i32x4 = self.buffer_load_fp8x16(src_div, src_elem) - src_bytes = Vec(src_i32x4, (4,), fx.Int32).bitcast(fx.Int8) - d_block = d_col // fx.Index(16) - d_in_block = d_col % fx.Index(16) - for n_in_group in range_constexpr(16): - n = n_group16 * fx.Index(16) + fx.Index(n_in_group) - dest_wave = n // fx.Index(8) - dest_n = n % fx.Index(8) - dst_byte = ( - dst_base - + (dest_wave * fx.Index(traits.V_HEAD_DIM // 16) + d_block) * fx.Index(128) - + dest_n * fx.Index(16) - + d_in_block - ) - ptr = buffer_ops.get_element_ptr( - self.lds_vt_base_ptr, - byte_offset=fx.Int32(dst_byte - self.lds_vt_base_idx), - elem_type=T.i8, - ) - llvm.StoreOp(as_mlir_value(src_bytes[_sigma_k_tile_n(n_in_group)]), ptr, alignment=1) - - scf_if_dispatch(flat < fx.Index(total_v16), _load_store_active) - - def _stage_vt_dequant_fp8(self, tile_start, buf_id): - # Dequantize fp8 V into the exact bf16 V staging positions. The two d-iters - # load 8 fp8 at D offsets 64 apart; a contiguous 16B load would gather wrong. - traits = self.traits - vt_buf = buf_id * traits.VT_BF16_ELEMS - n_in_tile = self.n_in_warp * traits.NUM_WAVES + self.wave_id - for d in range_constexpr(traits.SDRPT_BF): - global_d = self.d_bucket * traits.VEC_BF + (d * traits.D128_BF) - src_elem = ( - self.kv_gmem_elem_offset + n_in_tile * self.stride_kv_n_v + global_d + tile_start * self.stride_kv_n_v - ) - v_i32x2 = fly.copy_atom_call_ssa( - [self.v2i32_type], self.v_fp8_load64_atom, fx.slice(self.v_div, (None, fx.Int32(src_elem))) - ) - v_words = Vec(v_i32x2, (2,), fx.Int32) - bf = [] - for w in range_constexpr(2): - word = as_mlir_value(fx.Int32(v_words[w])) - lo2 = Vec(rocdl.cvt_pk_f32_fp8(Vec.make_type(2, fx.Float32), word, False), (2,), fx.Float32) - hi2 = Vec(rocdl.cvt_pk_f32_fp8(Vec.make_type(2, fx.Float32), word, True), (2,), fx.Float32) - for e in (lo2[0], lo2[1], hi2[0], hi2[1]): - bf.append(fx.Float32(e) * self.vd_fp8) - v8bf = self.bf16_trunc_pack_v8(bf) - byte_off = ( - vt_buf - + self.wave_id_uni * traits.VLS_BF - + d * traits.SNRPT_BF * traits.VLS_BF - + self.lane * traits.VEC_BF - ) * traits.EB_BF - lds_ptr = buffer_ops.get_element_ptr(self.lds_vt_base_ptr, byte_offset=byte_off, elem_type=T.i8) - llvm.StoreOp(as_mlir_value(v8bf), lds_ptr, alignment=16) - class DualwaveFp8KvLdsToVgprLoader(DualwaveFp8KernelContext): def __init__(self, ctx): @@ -5951,31 +5524,9 @@ def _read_strip(key): def load_v(self, buf_id): if const_expr(self.traits.FP8_PV_SEGMENTED and not self.traits.BN128): return fx.Index(buf_id) - if const_expr(self.traits.FP8_PV): - if const_expr(self.traits.KV_VECTORIZED and self.traits.FP8_V_BANKPAD): - return self._load_v_fp8_vectorized_bankpad(buf_id) - return self._load_v_fp8_block(buf_id) - # Read all V packs from the bf16 vt scratch for buffer `buf_id`. - traits = self.traits - if const_expr(traits.KV_VECTORIZED): - return self.load_v_steps(buf_id, 0, 4) - - urv = ( - self.lane_div_32 * traits.URV_GRPK_BF - + ((self.lane % 16) // 4) * traits.URV_LANE_HI_BF - + ((self.lane // 16) % 2) * traits.URV_GRP_N_BF - + (self.lane % 4) * traits.URV_LANE_LO_BF - ) - packs = [[None] * traits.D_CHUNKS for _ in range(4)] - for dc in range_constexpr(traits.D_CHUNKS): - dc_off = (dc // 2) * traits.URV_DC_AXIS0_BF + (dc % 2) * traits.URV_DC_AXIS1_BF - for k_substep in range_constexpr(4): - imm_lo = (k_substep * traits.URV_STEPK_BF + dc_off) * traits.EB_BF - byte0 = (urv + buf_id * traits.VT_BF16_ELEMS) * traits.EB_BF + self.lds_vt_base_idx - a = _ds_read_tr16_b64_imm(self.v4bf16_type, fx.Int32(byte0), imm_lo) - b = _ds_read_tr16_b64_imm(self.v4bf16_type, fx.Int32(byte0), imm_lo + traits.URV_I5_BF * traits.EB_BF) - packs[k_substep][dc] = Vec(a).shuffle(Vec(b), [0, 1, 2, 3, 4, 5, 6, 7]).ir_value() - return packs + if const_expr(self.traits.PAGED): + return self._load_v_fp8_vectorized_bankpad(buf_id) + return self._load_v_fp8_block(buf_id) def _load_v_fp8_vectorized_bankpad(self, buf_id): """Read the D-major V tile into the FP8 MFMA A fragment.""" @@ -6004,44 +5555,12 @@ def _load_v_fp8_vectorized_bankpad(self, buf_id): return packs def load_v_steps(self, buf_id, first_step, num_steps): - """Read a compile-time subset of vectorized HIPREC V K-steps. - - V192 otherwise keeps all 24 v8bf16 packs live together with six output - accumulators and spills. Loading two K-steps at a time bounds that - transient register set while preserving the existing LDS layout. - """ - traits = self.traits - if const_expr(traits.FP8_PV_SEGMENTED): - return fx.Index(buf_id) - if const_expr(not traits.KV_VECTORIZED or traits.FP8_PV): - raise RuntimeError("load_v_steps requires vectorized HIPREC V") + """Carry the segmented V buffer for the generic on-demand P*V steps.""" + if const_expr(not self.traits.FP8_PV_SEGMENTED): + raise RuntimeError("load_v_steps requires segmented paged FP8 V") if const_expr(first_step < 0 or num_steps < 1 or first_step + num_steps > 4): raise RuntimeError("invalid V K-step range") - - v_base = buf_id * traits.VT_BF16_ELEMS - lane_base = ( - (self.lane_mod_32 // traits.SMEM_N_PER_WAVE) * traits.VEC_V_ROW_STRIDE - + self.lane_div_32 * traits.D128_BF - + (self.lane_mod_32 % traits.SMEM_N_PER_WAVE) * traits.VEC_BF - ) - v_base_ptr = buffer_ops.get_element_ptr( - self.lds_vt_base_ptr, - byte_offset=fx.Int32((v_base + lane_base) * traits.EB_BF), - elem_type=T.i8, - ) - packs = [[None] * traits.D_CHUNKS for _ in range(4)] - for dc in range_constexpr(traits.D_CHUNKS): - for k_substep in range_constexpr(first_step, first_step + num_steps): - const_off = dc * (traits.D_CHUNK // traits.SMEM_N_PER_WAVE) * traits.VEC_V_ROW_STRIDE + k_substep * ( - 2 * traits.D128_BF - ) - ptr = buffer_ops.get_element_ptr( - v_base_ptr, - byte_offset=fx.Int32(const_off * traits.EB_BF), - elem_type=T.i8, - ) - packs[k_substep][dc] = llvm.LoadOp(self.v8bf16_type, ptr, alignment=16).result - return packs + return fx.Index(buf_id) def _load_v_fp8_block(self, buf_id): traits = self.traits diff --git a/tests/kernels/test_flash_attn_fwd.py b/tests/kernels/test_flash_attn_fwd.py index be97f8ec5f..abf5eec17b 100644 --- a/tests/kernels/test_flash_attn_fwd.py +++ b/tests/kernels/test_flash_attn_fwd.py @@ -5599,14 +5599,6 @@ def test_fp8_pv_head_dim_uses_two_generic_segments(value_head_dim, expected): assert _factor_fp8_pv_head_dim(value_head_dim, segment_capacity=128, alignment=32) == expected -@pytest.mark.parametrize( - ("head_dims", "expected"), - [((128, 128), False), ((192, 128), False), ((192, 192), True)], -) -def test_paged_fp8_segmented_pv_dispatch(head_dims, expected): - assert flash_attn_interface._use_paged_fp8_segmented_pv(head_dims) is expected - - @_requires_gfx950 def test_fp8_default_is_the_lazy_rescale(): """Every other fp8 case passes the flag, so the default would go untested. From 3791131efc6e3c9edf485041c126e7ee0f46d8c6 Mon Sep 17 00:00:00 2001 From: Xiake Sun Date: Wed, 9 Sep 2026 03:24:12 +0000 Subject: [PATCH 38/59] refactor(fmha): use explicit-width paged FP8 arithmetic - Replace PR-introduced index wrappers with explicit 64-bit integer values. - Keep the ring state typed across loop boundaries and retain 32-bit ABI casts. - Use typed vector conversion and remove redundant scalar wrappers. Signed-off-by: Xiake Sun --- .../attention/flash_attn_fp8_paged_gfx950.py | 72 +++++------ kernels/attention/flash_attn_utils.py | 116 +++++++++--------- 2 files changed, 94 insertions(+), 94 deletions(-) diff --git a/kernels/attention/flash_attn_fp8_paged_gfx950.py b/kernels/attention/flash_attn_fp8_paged_gfx950.py index edd9f06f16..65a396ddcb 100644 --- a/kernels/attention/flash_attn_fp8_paged_gfx950.py +++ b/kernels/attention/flash_attn_fp8_paged_gfx950.py @@ -165,7 +165,7 @@ def flash_attn_paged_fp8_bn128_kernel( ctx.init_sequence_lengths() num_q_blocks = (ctx.seqlen_q_v + traits.BLOCK_M - 1) // traits.BLOCK_M active_q_block = ctx.q_block_idx < num_q_blocks - reversed_q_block = num_q_blocks - fx.Index(1) - ctx.q_block_idx + reversed_q_block = num_q_blocks - 1 - ctx.q_block_idx ctx.q_block_idx = active_q_block.select(reversed_q_block, ctx.q_block_idx) ctx.q_start = ctx.q_block_idx * traits.BLOCK_M ctx.q_gmem_elem_offset = ( @@ -188,7 +188,7 @@ def flash_attn_paged_fp8_bn128_kernel( BN = traits.BLOCK_N D_CHUNKS = traits.D_CHUNKS - NPF = const_expr(traits.NUM_PREFETCH_K) + NPF = traits.NUM_PREFETCH_K t0 = ctx.split_t0 t_end = ctx.split_t_end @@ -225,7 +225,7 @@ def _merge_tile_max(v_s_a, v_s_b): return softmax_helper.floor_masked_max(m_tile) page_t0, page_t1 = ctx.load_page_id_pair(t0 * BN) - kv_gmem_to_lds.load_k(t0 * BN, t0 % fx.Index(NPF), page_id=page_t0) + kv_gmem_to_lds.load_k(t0 * BN, t0 % NPF, page_id=page_t0) fx.rocdl.s_waitcnt(vmcnt=0, lgkmcnt=0, expcnt=0) rocdl.sched_barrier(0) rocdl.s_barrier() @@ -236,14 +236,14 @@ def _merge_tile_max(v_s_a, v_s_b): q_wide = gemm_helper.load_q_wide() page_t2, page_t3 = ctx.load_page_id_pair((t0 + 2) * BN) - kv_gmem_to_lds.load_k((t0 + 1) * BN, (t0 + 1) % fx.Index(NPF), page_id=page_t1) - kv_gmem_to_lds.load_v(t0 * BN, t0 % fx.Index(NPF), page_id=page_t0) - kv_gmem_to_lds.load_v((t0 + 1) * BN, (t0 + 1) % fx.Index(NPF), page_id=page_t1) - kv_gmem_to_lds.load_k((t0 + 2) * BN, (t0 + 2) % fx.Index(NPF), page_id=page_t2) - kv_gmem_to_lds.load_k((t0 + 3) * BN, (t0 + 3) % fx.Index(NPF), page_id=page_t3) + kv_gmem_to_lds.load_k((t0 + 1) * BN, (t0 + 1) % NPF, page_id=page_t1) + kv_gmem_to_lds.load_v(t0 * BN, t0 % NPF, page_id=page_t0) + kv_gmem_to_lds.load_v((t0 + 1) * BN, (t0 + 1) % NPF, page_id=page_t1) + kv_gmem_to_lds.load_k((t0 + 2) * BN, (t0 + 2) % NPF, page_id=page_t2) + kv_gmem_to_lds.load_k((t0 + 3) * BN, (t0 + 3) % NPF, page_id=page_t3) if const_expr(traits.FP8_PV_SEGMENTED): - kv_gmem_to_lds.load_v((t0 + 2) * BN, (t0 + 2) % fx.Index(NPF), page_id=page_t2) - kv_gmem_to_lds.load_v((t0 + 3) * BN, (t0 + 3) % fx.Index(NPF), page_id=page_t3) + kv_gmem_to_lds.load_v((t0 + 2) * BN, (t0 + 2) % NPF, page_id=page_t2) + kv_gmem_to_lds.load_v((t0 + 3) * BN, (t0 + 3) % NPF, page_id=page_t3) else: next_v_a = kv_gmem_to_lds._load_v_fp8_vectorized_bankpad_source( (t0 + 2) * BN, @@ -262,29 +262,29 @@ def _merge_tile_max(v_s_a, v_s_b): l_row = ctx.c_zero_f v_o = [ctx.c_zero_v16f32 for _ in range_constexpr(D_CHUNKS)] - NPF_I = const_expr(fx.Index(NPF)) + NPF_I = fx.Int64(NPF) def _ring_wrap(x): if const_expr(NPF == 8): - return x & fx.Index(7) + return x & 7 return (x >= NPF_I).select(x - NPF_I, x) - init_args = [m_row, l_row] + v_o + [t0 % fx.Index(NPF)] + init_args = [m_row, l_row] + v_o + [fx.Int64(t0) % NPF] if const_expr(not traits.FP8_PV_SEGMENTED): init_args += [next_v_a, next_v_b] loop_results = init_args next_v_arg_idx = 3 + D_CHUNKS - for j, loop_args in range(fx.Index(t0), t_end, fx.Index(2), init=init_args): + for j, loop_args in range(fx.Int64(t0), t_end, fx.Int64(2), init=init_args): m_row = loop_args[0] l_row = loop_args[1] v_o = [loop_args[2 + i] for i in range_constexpr(D_CHUNKS)] - a_buf = loop_args[2 + D_CHUNKS] - b_buf = _ring_wrap(a_buf + fx.Index(1)) - nn_a_buf = _ring_wrap(a_buf + fx.Index(2)) - nn_b_buf = _ring_wrap(a_buf + fx.Index(3)) - f_a_buf = _ring_wrap(a_buf + fx.Index(4)) - f_b_buf = _ring_wrap(a_buf + fx.Index(5)) + a_buf = fx.Int64(loop_args[2 + D_CHUNKS]) + b_buf = _ring_wrap(a_buf + 1) + nn_a_buf = _ring_wrap(a_buf + 2) + nn_b_buf = _ring_wrap(a_buf + 3) + f_a_buf = _ring_wrap(a_buf + 4) + f_b_buf = _ring_wrap(a_buf + 5) if const_expr(traits.FP8_PV_SEGMENTED): v_k_a = kv_lds_to_regs.load_k(a_buf) @@ -292,11 +292,11 @@ def _ring_wrap(x): v_s_a = gemm_helper.qk(v_k_a, q_wide) v_s_b = gemm_helper.qk(v_k_b, q_wide) - page_f_a, page_f_b = ctx.load_page_id_pair((j + fx.Index(4)) * BN) - kv_gmem_to_lds.load_k((j + fx.Index(4)) * BN, f_a_buf, page_id=page_f_a) - kv_gmem_to_lds.load_k((j + fx.Index(5)) * BN, f_b_buf, page_id=page_f_b) - kv_gmem_to_lds.load_v((j + fx.Index(4)) * BN, f_a_buf, page_id=page_f_a) - kv_gmem_to_lds.load_v((j + fx.Index(5)) * BN, f_b_buf, page_id=page_f_b) + page_f_a, page_f_b = ctx.load_page_id_pair((j + 4) * BN) + kv_gmem_to_lds.load_k((j + 4) * BN, f_a_buf, page_id=page_f_a) + kv_gmem_to_lds.load_k((j + 5) * BN, f_b_buf, page_id=page_f_b) + kv_gmem_to_lds.load_v((j + 4) * BN, f_a_buf, page_id=page_f_a) + kv_gmem_to_lds.load_v((j + 5) * BN, f_b_buf, page_id=page_f_b) v_s_a, v_s_b = _mask_pair(v_s_a, v_s_b, j) m_tile = _merge_tile_max(v_s_a, v_s_b) @@ -319,18 +319,18 @@ def _ring_wrap(x): v_k_b = kv_lds_to_regs.load_k(b_buf) v_v_a = kv_lds_to_regs.load_v(a_buf) - page_f_a, page_f_b = ctx.load_page_id_pair((j + fx.Index(4)) * BN) - kv_gmem_to_lds.load_k((j + fx.Index(4)) * BN, f_a_buf, page_id=page_f_a) - kv_gmem_to_lds.load_k((j + fx.Index(5)) * BN, f_b_buf, page_id=page_f_b) + page_f_a, page_f_b = ctx.load_page_id_pair((j + 4) * BN) + kv_gmem_to_lds.load_k((j + 4) * BN, f_a_buf, page_id=page_f_a) + kv_gmem_to_lds.load_k((j + 5) * BN, f_b_buf, page_id=page_f_b) v_s_a = gemm_helper.qk(v_k_a, q_wide) kv_gmem_to_lds._store_v_fp8_vectorized_bankpad(next_v_a, nn_a_buf) v_f_a = kv_gmem_to_lds._load_v_fp8_vectorized_bankpad_source( - (j + fx.Index(4)) * BN, + (j + 4) * BN, page_id=page_f_a, ) v_f_b = kv_gmem_to_lds._load_v_fp8_vectorized_bankpad_source( - (j + fx.Index(5)) * BN, + (j + 5) * BN, page_id=page_f_b, ) v_s_b = gemm_helper.qk(v_k_b, q_wide) @@ -412,7 +412,7 @@ def flash_attn_dualwave_swp_fp8_gfx950_kernel( # Issue the longest active q-blocks first within each batch group. num_q_blocks = (ctx.seqlen_q_v + traits.BLOCK_M - 1) // traits.BLOCK_M active_q_block = ctx.q_block_idx < num_q_blocks - reversed_q_block = num_q_blocks - fx.Index(1) - ctx.q_block_idx + reversed_q_block = num_q_blocks - 1 - ctx.q_block_idx ctx.q_block_idx = active_q_block.select(reversed_q_block, ctx.q_block_idx) ctx.q_start = ctx.q_block_idx * traits.BLOCK_M ctx.q_gmem_elem_offset = ( @@ -482,12 +482,12 @@ def _run_q_block(): init_args.append(ctx.v_pair_to_vec32(v_p_0)) # Software-pipelined KV loop. - loop_lb = fx.Index(3) + loop_lb = fx.Int64(3) loop_results = init_args for j, loop_args in range( loop_lb, - ctx.split_t_end - fx.Index(1), - fx.Index(2), + ctx.split_t_end - 1, + fx.Int64(2), init=init_args, ): m_row = loop_args[0] @@ -902,8 +902,8 @@ def launch_flash_attn_dualwave_swp( ): # Make shape/mode traits visible to the JIT cache key. _ = _dualwave_swp_fp8_cache_tag - bs_idx = fx.Index(batch_size) - sl_idx = fx.Index(seq_len) + bs_idx = fx.Int64(batch_size) + sl_idx = fx.Int64(seq_len) num_q_blocks = (sl_idx + BLOCK_M - 1) // BLOCK_M grid_z = bs_idx diff --git a/kernels/attention/flash_attn_utils.py b/kernels/attention/flash_attn_utils.py index 41165adf89..c6a624877e 100644 --- a/kernels/attention/flash_attn_utils.py +++ b/kernels/attention/flash_attn_utils.py @@ -1007,7 +1007,7 @@ def _init_dualwave_thread_mapping(ctx): linear_head_batch = fx.Index(gpu.block_idx.x) ctx.h_idx = linear_head_batch % traits.NUM_HEADS_Q batch_in_group = linear_head_batch // traits.NUM_HEADS_Q - ctx.batch_idx = fx.Index(gpu.block_idx.z) * batch_interleave_group + batch_in_group + ctx.batch_idx = fx.Int64(gpu.block_idx.z) * batch_interleave_group + batch_in_group ctx.q_block_idx = fx.Index(gpu.block_idx.y) else: ctx.h_idx = fx.Index(gpu.block_idx.x) @@ -1017,7 +1017,7 @@ def _init_dualwave_thread_mapping(ctx): ctx.batch_idx = ctx.bz_idx // traits.NUM_KV_SPLITS ctx.split_idx = ctx.bz_idx % traits.NUM_KV_SPLITS elif const_expr(batch_interleave_group <= 1): - ctx.batch_idx = fx.Index(gpu.block_idx.z) + ctx.batch_idx = fx.Int64(gpu.block_idx.z) ctx.split_idx = None else: ctx.split_idx = None @@ -4604,10 +4604,10 @@ def init_runtime_indices(self): self.seq_len_v = fx.Index(self.seq_len) self.seq_len_kv_v = fx.Index(self.seq_len_kv) self.stride_q_n_v = fx.Index(self.stride_q_n) - self.stride_o_n_v = fx.Index(self.stride_o_n) + self.stride_o_n_v = fx.Int64(self.stride_o_n) self.stride_kv_n_v = fx.Index(self.stride_kv_n) if const_expr(traits.PAGED): - self.stride_v_n_v = fx.Index(traits.NUM_HEADS_KV * traits.V_HEAD_DIM) + self.stride_v_n_v = fx.Int64(traits.NUM_HEADS_KV * traits.V_HEAD_DIM) elif traits.HEAD_DIM_V == traits.HEAD_DIM: self.stride_v_n_v = self.stride_kv_n_v self.stride_o_n_v = self.stride_q_n_v @@ -4642,7 +4642,7 @@ def init_lds(self, shared_storage): else: self.lds_vt_base_ptr = lds.vt.ptr.llvm_ptr if hasattr(lds, "q"): - self.lds_q_base_idx = fx.Index(fx.ptrtoint(lds.q.ptr)) + self.lds_q_base_idx = fx.Int64(fx.ptrtoint(lds.q.ptr)) self.lds_q_base_ptr = lds.q.ptr.llvm_ptr else: self.lds_q_base_idx = None @@ -4719,10 +4719,10 @@ def _make_buf_div(tensor, nrec_bytes): self.buf_flags_i32 = fx.Int32(buffer_ops._get_buffer_flags()) self.k_base_iter = fx.get_iter(self.K) self.v_base_iter = fx.get_iter(self.V) - k_page_elems = fx.Index(traits.PAGE_SIZE * traits.NUM_HEADS_KV * traits.HEAD_DIM) - v_page_elems = fx.Index(traits.PAGE_SIZE * traits.NUM_HEADS_KV * traits.V_HEAD_DIM) - self.k_page_byte_stride = k_page_elems * fx.Index(traits.ELEM_BYTES) - self.v_page_byte_stride = v_page_elems * fx.Index(traits.ELEM_BYTES) + k_page_elems = fx.Int64(traits.PAGE_SIZE * traits.NUM_HEADS_KV * traits.HEAD_DIM) + v_page_elems = fx.Int64(traits.PAGE_SIZE * traits.NUM_HEADS_KV * traits.V_HEAD_DIM) + self.k_page_byte_stride = k_page_elems * fx.Int64(traits.ELEM_BYTES) + self.v_page_byte_stride = v_page_elems * fx.Int64(traits.ELEM_BYTES) self.k_page_nrec_bytes = fx.Int64(self.k_page_byte_stride) self.v_page_nrec_bytes = fx.Int64(self.v_page_byte_stride) self.k_page_layout = fx.make_layout(fx.Int32(k_page_elems), fx.Int32(1)) @@ -4734,7 +4734,7 @@ def _make_buf_div(tensor, nrec_bytes): self.page_v1i32 = Vec.make_type(1, fx.Int32) self.page_i32x2_atom = fx.make_copy_atom(fx.rocdl.BufferCopy64b(), fx.Int32) self.page_v2i32 = Vec.make_type(2, fx.Int32) - self.page_base = self.batch_idx * fx.Index(self.block_table_stride) + self.page_base = self.batch_idx * fx.Int64(self.block_table_stride) self.k_div = None self.v_div = None else: @@ -4900,22 +4900,22 @@ def global_idx_o(self, token_idx, col): def load_page_id(self, tile_start): """Load this request's physical page for a 64-token logical KV tile.""" - local_page = fx.Index(tile_start) // fx.Index(self.traits.PAGE_SIZE) + local_page = fx.Int64(tile_start) // fx.Int64(self.traits.PAGE_SIZE) # The software pipeline rounds to an even tile count and at least four # tiles. Redirect those masked padding tiles to entry zero instead of # reading past the final request's page-index allocation. page_idx = self.page_base + local_page - safe_idx = fx.Index((local_page < self.num_kv_tiles).select(page_idx, fx.Index(0))) + safe_idx = fx.Int64((local_page < self.num_kv_tiles).select(page_idx, fx.Int64(0))) v = fly.copy_atom_call_ssa( [self.page_v1i32], self.page_i32_atom, fx.slice(self.page_indices_div, (None, fx.Int32(safe_idx))), ) - return fx.Index(fx.Int32(rocdl.readfirstlane(T.i32, as_mlir_value(fx.Int32(Vec(v, (1,), fx.Int32)[0]))))) + return fx.Int64(fx.Int32(rocdl.readfirstlane(T.i32, as_mlir_value(fx.Int32(Vec(v, (1,), fx.Int32)[0]))))) def load_page_id_pair(self, tile_start): """Load two adjacent physical page IDs for a BN128 tile pair.""" - local_page = fx.Int32(fx.Index(tile_start) // fx.Index(self.traits.PAGE_SIZE)) + local_page = fx.Int32(fx.Int64(tile_start) // fx.Int64(self.traits.PAGE_SIZE)) page_base = fx.Int32(self.page_base) num_kv_tiles = fx.Int32(self.num_kv_tiles) page_idx = page_base + local_page @@ -4928,7 +4928,7 @@ def load_page_id_pair(self, tile_start): pages = Vec(v, (2,), fx.Int32) page0 = rocdl.readfirstlane(T.i32, as_mlir_value(fx.Int32(pages[0]))) page1 = rocdl.readfirstlane(T.i32, as_mlir_value(fx.Int32(pages[1]))) - return fx.Index(fx.Int32(page0)), fx.Index(fx.Int32(page1)) + return fx.Int64(fx.Int32(page0)), fx.Int64(fx.Int32(page1)) def make_page_view(self, tensor_iter, page_id, *, is_value=False): if const_expr(is_value): @@ -5001,7 +5001,7 @@ def load_all_wide(self, q_row_in_block): hi = Vec(self.buffer_load_128(g_idx + 16), (4,), fx.Int32) packs.append(lo.shuffle(hi, [0, 1, 2, 3, 4, 5, 6, 7]).ir_value()) else: - byte_row = q_row_in_block * fx.Index(traits.HEAD_DIM) + q_col + byte_row = q_row_in_block * fx.Int64(traits.HEAD_DIM) + q_col packs.append(self.read_i32x8_lds(self.lds_q_base_ptr, fx.Int32(byte_row))) return packs @@ -5011,7 +5011,7 @@ def __init__(self, ctx): super().__init__(ctx) def _v8bf16_to_f32(self, v8): - f32 = Vec(llvm.FPExtOp(Vec.make_type(8, fx.Float32), as_mlir_value(v8)).result, (8,), fx.Float32) + f32 = Vec(v8).to(fx.Float32) return [f32[i] for i in range_constexpr(8)] def _p_to_fp8_i32x8(self, v_p): @@ -5084,8 +5084,8 @@ def _load_v_fp8_segmented_chunk(self, buf_id, dc): traits = self.traits row_stride = traits.FP8_V_ROW_STRIDE tile_bytes = traits.V_HEAD_DIM * row_stride - aligned_base = ((self.lds_vt_base_idx + fx.Index(127)) // fx.Index(128)) * fx.Index(128) - tile_base = aligned_base + fx.Index(buf_id * tile_bytes) + aligned_base = ((self.lds_vt_base_idx + fx.Int64(127)) // fx.Int64(128)) * fx.Int64(128) + tile_base = aligned_base + fx.Int64(buf_id * tile_bytes) d_global = dc * traits.D_CHUNK if const_expr(d_global < traits.FP8_V_H1): segment_base = 0 @@ -5093,14 +5093,14 @@ def _load_v_fp8_segmented_chunk(self, buf_id, dc): else: segment_base = traits.FP8_V_H1 * row_stride d_local = d_global - traits.FP8_V_H1 - token_base = self.lane_div_32 * fx.Index(32) - d_row = fx.Index(d_local) + self.lane_mod_32 - row_base = tile_base + fx.Index(segment_base) + d_row * fx.Index(row_stride) + token_base + token_base = self.lane_div_32 * fx.Int64(32) + d_row = fx.Int64(d_local) + self.lane_mod_32 + row_base = tile_base + fx.Int64(segment_base) + d_row * fx.Int64(row_stride) + token_base halves = [] for half in range_constexpr(2): ptr = buffer_ops.get_element_ptr( self.lds_vt_base_ptr, - byte_offset=fx.Int32(row_base + fx.Index(half * 16) - self.lds_vt_base_idx), + byte_offset=fx.Int32(row_base + fx.Int64(half * 16) - self.lds_vt_base_idx), elem_type=T.i8, ) halves.append(llvm.LoadOp(Vec.make_type(2, fx.Int64), ptr, alignment=16).result) @@ -5241,7 +5241,7 @@ def _load_compact(): copy_waves = (traits.BLOCK_N * traits.HEAD_DIM) // ( self.NUM_DMA_K * traits.WARP_SIZE * traits.KV_VEC_SIZE ) - if self.wave_id_uni < fx.Index(copy_waves): + if self.wave_id_uni < fx.Int64(copy_waves): _load_vectorized() _load_compact() @@ -5353,15 +5353,15 @@ def _permute_v_fp8_vectorized(self, src_i32x4): pair_ty = ir.Type.parse("!llvm.struct<(i32, i32)>") pair_lo = rocdl.permlane16_swap( pair_ty, - as_mlir_value(fx.Int32(src_words[0])), - as_mlir_value(fx.Int32(src_words[1])), + src_words[0].ir_value(), + src_words[1].ir_value(), False, False, ) pair_hi = rocdl.permlane16_swap( pair_ty, - as_mlir_value(fx.Int32(src_words[2])), - as_mlir_value(fx.Int32(src_words[3])), + src_words[2].ir_value(), + src_words[3].ir_value(), False, False, ) @@ -5382,24 +5382,24 @@ def _stage_v_fp8_vectorized_segment(self, tile_start, buf_id, d_offset, segment_ flat = self.tid def _stage_active(): - n_group16 = self.lane_in_warp // fx.Index(16) - d_local = self.wave_id * fx.Index(16) + self.lane_in_warp % fx.Index(16) - d_col = fx.Index(d_offset) + d_local + n_group16 = self.lane_in_warp // fx.Int64(16) + d_local = self.wave_id * fx.Int64(16) + self.lane_in_warp % fx.Int64(16) + d_col = fx.Int64(d_offset) + d_local src_elem = ( self.kv_head_idx * token_groups16 * traits.V_HEAD_DIM * traits.KV_VEC_SIZE + n_group16 * traits.V_HEAD_DIM * traits.KV_VEC_SIZE + d_col * traits.KV_VEC_SIZE ) reordered = self._permute_v_fp8_vectorized(self.buffer_load_fp8x16(src_div, src_elem)) - dest_group16 = (n_group16 % fx.Index(2)) * fx.Index(2) + n_group16 // fx.Index(2) - aligned_base = ((self.lds_vt_base_idx + fx.Index(127)) // fx.Index(128)) * fx.Index(128) + dest_group16 = (n_group16 % fx.Int64(2)) * fx.Int64(2) + n_group16 // fx.Int64(2) + aligned_base = ((self.lds_vt_base_idx + fx.Int64(127)) // fx.Int64(128)) * fx.Int64(128) tile_bytes = traits.V_HEAD_DIM * traits.FP8_V_ROW_STRIDE dst_byte = ( aligned_base - + fx.Index(buf_id * tile_bytes) - + fx.Index(d_offset * traits.FP8_V_ROW_STRIDE) - + d_local * fx.Index(traits.FP8_V_ROW_STRIDE) - + dest_group16 * fx.Index(16) + + fx.Int64(buf_id * tile_bytes) + + fx.Int64(d_offset * traits.FP8_V_ROW_STRIDE) + + d_local * fx.Int64(traits.FP8_V_ROW_STRIDE) + + dest_group16 * fx.Int64(16) ) ptr = buffer_ops.get_element_ptr( self.lds_vt_base_ptr, @@ -5408,7 +5408,7 @@ def _stage_active(): ) llvm.StoreOp(as_mlir_value(reordered), ptr, alignment=16) - scf_if_dispatch(flat < fx.Index(segment_dim * token_groups16), _stage_active) + scf_if_dispatch(flat < fx.Int64(segment_dim * token_groups16), _stage_active) def _stage_v_fp8_vectorized_bankpad(self, tile_start, buf_id, page_id=None): """Stage vectorized paged V in the bank-padded direct-FP8 layout.""" @@ -5426,8 +5426,8 @@ def _load_v_fp8_vectorized_bankpad_source(self, tile_start, page_id=None): # one vector per CTA thread, so no per-lane bounds branch is needed. prefix_dim = traits.FP8_V_H1 if traits.FP8_PV_SEGMENTED else traits.V_HEAD_DIM assert prefix_dim * token_groups16 == traits.BLOCK_SIZE - n_group16 = self.lane_in_warp // fx.Index(16) - d_col = self.wave_id * fx.Index(16) + self.lane_in_warp % fx.Index(16) + n_group16 = self.lane_in_warp // fx.Int64(16) + d_col = self.wave_id * fx.Int64(16) + self.lane_in_warp % fx.Int64(16) src_elem = ( self.kv_head_idx * token_groups16 * traits.V_HEAD_DIM * traits.KV_VEC_SIZE + n_group16 * traits.V_HEAD_DIM * traits.KV_VEC_SIZE @@ -5435,7 +5435,7 @@ def _load_v_fp8_vectorized_bankpad_source(self, tile_start, page_id=None): ) source = self.buffer_load_fp8x16(src_div, src_elem) if const_expr(traits.BN128 and traits.FP8_PV_SEGMENTED): - tail = self.buffer_load_fp8x16(src_div, src_elem + fx.Index(traits.FP8_V_H1 * traits.KV_VEC_SIZE)) + tail = self.buffer_load_fp8x16(src_div, src_elem + fx.Int64(traits.FP8_V_H1 * traits.KV_VEC_SIZE)) source = Vec(source).shuffle(Vec(tail), [0, 1, 2, 3, 4, 5, 6, 7]).ir_value() return source @@ -5449,7 +5449,7 @@ def _store_v_fp8_vectorized_bankpad(self, src_i32x4, buf_id): @flyc.jit def _store_tail(): - if self.wave_id_uni < fx.Index(self.traits.FP8_V_H2 // self.traits.KV_VEC_SIZE): + if self.wave_id_uni < fx.Int64(self.traits.FP8_V_H2 // self.traits.KV_VEC_SIZE): self._store_v_fp8_vectorized_bankpad_segment(tail, buf_id, self.traits.FP8_V_H1) _store_tail() @@ -5460,16 +5460,16 @@ def _store_v_fp8_vectorized_bankpad_segment(self, src_i32x4, buf_id, d_offset): traits = self.traits row_stride = traits.FP8_V_ROW_STRIDE v_tile_bytes = traits.V_HEAD_DIM * row_stride - aligned_base = ((self.lds_vt_base_idx + fx.Index(127)) // fx.Index(128)) * fx.Index(128) - dst_base = aligned_base + fx.Index(buf_id * v_tile_bytes) - n_group16 = self.lane_in_warp // fx.Index(16) - d_col = fx.Index(d_offset) + self.wave_id * fx.Index(16) + self.lane_in_warp % fx.Index(16) + aligned_base = ((self.lds_vt_base_idx + fx.Int64(127)) // fx.Int64(128)) * fx.Int64(128) + dst_base = aligned_base + fx.Int64(buf_id * v_tile_bytes) + n_group16 = self.lane_in_warp // fx.Int64(16) + d_col = fx.Int64(d_offset) + self.wave_id * fx.Int64(16) + self.lane_in_warp % fx.Int64(16) reordered = self._permute_v_fp8_vectorized(src_i32x4) - dest_group16 = (n_group16 % fx.Index(2)) * fx.Index(2) + n_group16 // fx.Index(2) + dest_group16 = (n_group16 % fx.Int64(2)) * fx.Int64(2) + n_group16 // fx.Int64(2) if const_expr(traits.BN128 and traits.V_HEAD_DIM == 192): # XOR 16-byte groups instead of padding each row in the six-slot ring. - dest_group16 = dest_group16 ^ ((d_col // fx.Index(4)) % fx.Index(4)) - dst_byte = dst_base + d_col * fx.Index(row_stride) + dest_group16 * fx.Index(16) + dest_group16 = dest_group16 ^ ((d_col // fx.Int64(4)) % fx.Int64(4)) + dst_byte = dst_base + d_col * fx.Int64(row_stride) + dest_group16 * fx.Int64(16) ptr = buffer_ops.get_element_ptr( self.lds_vt_base_ptr, byte_offset=fx.Int32(dst_byte - self.lds_vt_base_idx), @@ -5523,7 +5523,7 @@ def _read_strip(key): def load_v(self, buf_id): if const_expr(self.traits.FP8_PV_SEGMENTED and not self.traits.BN128): - return fx.Index(buf_id) + return fx.Int64(buf_id) if const_expr(self.traits.PAGED): return self._load_v_fp8_vectorized_bankpad(buf_id) return self._load_v_fp8_block(buf_id) @@ -5533,17 +5533,17 @@ def _load_v_fp8_vectorized_bankpad(self, buf_id): traits = self.traits row_stride = traits.FP8_V_ROW_STRIDE v_tile_bytes = traits.V_HEAD_DIM * row_stride - aligned_base = ((self.lds_vt_base_idx + fx.Index(127)) // fx.Index(128)) * fx.Index(128) - tile_base = aligned_base + fx.Index(buf_id * v_tile_bytes) - token_base = self.lane_div_32 * fx.Index(32) + aligned_base = ((self.lds_vt_base_idx + fx.Int64(127)) // fx.Int64(128)) * fx.Int64(128) + tile_base = aligned_base + fx.Int64(buf_id * v_tile_bytes) + token_base = self.lane_div_32 * fx.Int64(32) packs = [[None] * traits.D_CHUNKS for _ in range(4)] for dc in range_constexpr(traits.D_CHUNKS): - d_row = fx.Index(dc * traits.D_CHUNK) + self.lane_mod_32 - row_base = tile_base + d_row * fx.Index(row_stride) + d_row = fx.Int64(dc * traits.D_CHUNK) + self.lane_mod_32 + row_base = tile_base + d_row * fx.Int64(row_stride) for half in range_constexpr(2): - token_offset = token_base + fx.Index(half * 16) + token_offset = token_base + fx.Int64(half * 16) if const_expr(traits.BN128 and traits.V_HEAD_DIM == 192): - token_offset = token_offset ^ (((d_row // fx.Index(4)) % fx.Index(4)) * fx.Index(16)) + token_offset = token_offset ^ (((d_row // fx.Int64(4)) % fx.Int64(4)) * fx.Int64(16)) ptr = buffer_ops.get_element_ptr( self.lds_vt_base_ptr, byte_offset=fx.Int32(row_base + token_offset - self.lds_vt_base_idx), @@ -5560,7 +5560,7 @@ def load_v_steps(self, buf_id, first_step, num_steps): raise RuntimeError("load_v_steps requires segmented paged FP8 V") if const_expr(first_step < 0 or num_steps < 1 or first_step + num_steps > 4): raise RuntimeError("invalid V K-step range") - return fx.Index(buf_id) + return fx.Int64(buf_id) def _load_v_fp8_block(self, buf_id): traits = self.traits From fa3cee06cc451a121ebd7990c65f1397e5881f17 Mon Sep 17 00:00:00 2001 From: Xiake Sun Date: Wed, 9 Sep 2026 03:24:12 +0000 Subject: [PATCH 39/59] test(fmha): share the paged FP8 reference calculation - Gather vectorized page caches in one independent reference helper. - Reuse the existing cross-length PyTorch attention oracle. - Preserve ragged, graph, long-context, and multi-KV-head coverage. Signed-off-by: Xiake Sun --- tests/kernels/test_flash_attn_fwd.py | 99 ++++++++++++++-------------- 1 file changed, 51 insertions(+), 48 deletions(-) diff --git a/tests/kernels/test_flash_attn_fwd.py b/tests/kernels/test_flash_attn_fwd.py index abf5eec17b..7048e62273 100644 --- a/tests/kernels/test_flash_attn_fwd.py +++ b/tests/kernels/test_flash_attn_fwd.py @@ -4559,6 +4559,30 @@ def test_return_lse_false_returns_only_out(): assert out.shape == q.shape +def _paged_fp8_torch_reference(query, key, value, block_table, query_lengths, kv_lengths, descales): + """Independent PyTorch reference for packed page-64 FP8 inputs.""" + q_scale, k_scale, v_scale = descales + head_dim = query.shape[-1] + num_kv_heads = key.shape[1] + value_head_dim = value.shape[3] + expected = [] + q_offset = 0 + for batch_idx, (query_length, kv_length) in enumerate(zip(query_lengths, kv_lengths)): + query_batch = query[q_offset : q_offset + query_length].float() * q_scale + q_offset += query_length + physical_pages = block_table[batch_idx].long() + key_batch = ( + key[physical_pages].permute(0, 3, 1, 2, 4).reshape(-1, num_kv_heads, head_dim)[:kv_length].float() * k_scale + ) + value_batch = ( + value[physical_pages].permute(0, 2, 4, 1, 3).reshape(-1, num_kv_heads, value_head_dim)[:kv_length].float() + * v_scale + ) + result = pytorch_ref_attention_qkv_diff(query_batch[None], key_batch[None], value_batch[None]) + expected.append(result[0].to(torch.bfloat16)) + return torch.cat(expected) + + def test_return_lse_rejects_fp8(): dtype = torch.bfloat16 q = _rand_lse(2, 128, 4, 128, dtype=dtype) @@ -4705,26 +4729,15 @@ def noncontiguous_copy(tensor): torch.cuda.synchronize() torch.testing.assert_close(actual, copy_reference, rtol=0, atol=0) - expected = [] - for batch_idx, (query_length, kv_length) in enumerate(zip(query_lengths, kv_lengths)): - query_batch = query[q_indptr[batch_idx] : q_indptr[batch_idx + 1]].float() * query_descale - physical_pages = block_table[batch_idx].long() - key_batch = ( - key[physical_pages].permute(0, 3, 1, 2, 4).reshape(-1, num_kv_heads, head_dim)[:kv_length].float() - * key_descale - ).repeat_interleave(16 // num_kv_heads, dim=1) - value_batch = ( - value[physical_pages].permute(0, 2, 4, 1, 3).reshape(-1, num_kv_heads, value_head_dim)[:kv_length].float() - * value_descale - ).repeat_interleave(16 // num_kv_heads, dim=1) - logits = torch.einsum("qhd,khd->hqk", query_batch, key_batch) / math.sqrt(head_dim) - query_positions = torch.arange(query_length, device="cuda") - key_positions = torch.arange(kv_length, device="cuda") - causal_mask = key_positions[None, :] <= (kv_length - query_length + query_positions)[:, None] - logits.masked_fill_(~causal_mask[None, :, :], float("-inf")) - expected.append(torch.einsum("hqk,khd->qhd", torch.softmax(logits, dim=-1), value_batch).to(torch.bfloat16)) - - expected = torch.cat(expected) + expected = _paged_fp8_torch_reference( + query, + key, + value, + block_table, + query_lengths, + kv_lengths, + (query_descale, key_descale, value_descale), + ) assert actual.shape == (sum(query_lengths), 16, value_head_dim) assert bool(torch.isfinite(actual).all().item()) torch.testing.assert_close(actual, expected, rtol=2.0e-2, atol=2.0e-2) @@ -4842,20 +4855,15 @@ def test_paged_fp8_long_context_random_pages_matches_torch(head_dim, value_head_ ) torch.cuda.synchronize() - physical_pages = block_table[0].long() - query_ref = query.float() * query_descale - key_ref = ( - key[physical_pages].permute(0, 3, 1, 2, 4).reshape(-1, 1, head_dim)[:kv_length].float() * key_descale - ).expand(-1, 16, -1) - value_ref = ( - value[physical_pages].permute(0, 2, 4, 1, 3).reshape(-1, 1, value_head_dim)[:kv_length].float() * value_descale - ).expand(-1, 16, -1) - logits = torch.einsum("qhd,khd->hqk", query_ref, key_ref) / math.sqrt(head_dim) - query_positions = torch.arange(query_length, device="cuda") - key_positions = torch.arange(kv_length, device="cuda") - causal_mask = key_positions[None, :] <= (kv_length - query_length + query_positions)[:, None] - logits.masked_fill_(~causal_mask[None, :, :], float("-inf")) - expected = torch.einsum("hqk,khd->qhd", torch.softmax(logits, dim=-1), value_ref).to(torch.bfloat16) + expected = _paged_fp8_torch_reference( + query, + key, + value, + block_table, + [query_length], + [kv_length], + (query_descale, key_descale, value_descale), + ) assert bool(torch.isfinite(actual).all().item()) torch.testing.assert_close(actual, expected, rtol=2.0e-2, atol=2.0e-2) @@ -4911,20 +4919,15 @@ def test_paged_fp8_graph_replay_matches_torch(head_dim, value_head_dim): graph.replay() torch.cuda.synchronize() - physical_pages = block_table[0].long() - query_ref = query.float() * query_descale - key_ref = ( - key[physical_pages].permute(0, 3, 1, 2, 4).reshape(-1, 1, head_dim)[:kv_length].float() * key_descale - ).expand(-1, 16, -1) - value_ref = ( - value[physical_pages].permute(0, 2, 4, 1, 3).reshape(-1, 1, value_head_dim)[:kv_length].float() * value_descale - ).expand(-1, 16, -1) - logits = torch.einsum("qhd,khd->hqk", query_ref, key_ref) / math.sqrt(head_dim) - query_positions = torch.arange(query_length, device="cuda") - key_positions = torch.arange(kv_length, device="cuda") - causal_mask = key_positions[None, :] <= (kv_length - query_length + query_positions)[:, None] - logits.masked_fill_(~causal_mask[None, :, :], float("-inf")) - expected = torch.einsum("hqk,khd->qhd", torch.softmax(logits, dim=-1), value_ref).to(torch.bfloat16) + expected = _paged_fp8_torch_reference( + query, + key, + value, + block_table, + [query_length], + [kv_length], + (query_descale, key_descale, value_descale), + ) assert captured.data_ptr() == output.data_ptr() assert bool(torch.isfinite(output).all().item()) From 752c52e28f9221e13480055c4c149b2437ff71fe Mon Sep 17 00:00:00 2001 From: Xiake Sun Date: Wed, 9 Sep 2026 04:26:12 +0000 Subject: [PATCH 40/59] refactor(fmha): use public paged FP8 memory views - Construct bounded page resources with 64-bit rebasing and make_buffer_tensor. - Use aligned 16-byte LDS tiles, typed vector loads, and fx.copy stores. - Replace manual paged control flow and pointer casts with DSL operations. - Share canonical value-width traits and the register-resident Q loader. Signed-off-by: Xiake Sun --- .../attention/flash_attn_fp8_paged_gfx950.py | 6 +- kernels/attention/flash_attn_utils.py | 276 +++++++----------- 2 files changed, 104 insertions(+), 178 deletions(-) diff --git a/kernels/attention/flash_attn_fp8_paged_gfx950.py b/kernels/attention/flash_attn_fp8_paged_gfx950.py index 65a396ddcb..0906877b15 100644 --- a/kernels/attention/flash_attn_fp8_paged_gfx950.py +++ b/kernels/attention/flash_attn_fp8_paged_gfx950.py @@ -12,7 +12,6 @@ DualwaveFp8KernelContext, DualwaveFp8KvGmemToLdsLoader, DualwaveFp8KvLdsToVgprLoader, - DualwaveFp8QLoader, DualwaveFp8SoftmaxHelper, DualwaveFp8StoreHelper, _make_paged_dualwave_swp_fp8_traits, @@ -105,7 +104,7 @@ def build_flash_attn_paged_fp8_module( PAGED_BN128_VARLEN = bool(paged_bn128_varlen) BATCH_INTERLEAVE_GROUP = traits.BATCH_INTERLEAVE_GROUP DEFAULT_STRIDE_Q_N = traits.DEFAULT_STRIDE_Q_N - DEFAULT_STRIDE_O_N = traits.NUM_HEADS_Q * traits.V_HEAD_DIM + DEFAULT_STRIDE_O_N = traits.NUM_HEADS_Q * traits.HEAD_DIM_V DEFAULT_STRIDE_KV_N = traits.DEFAULT_STRIDE_KV_N _dualwave_swp_fp8_cache_tag = traits.cache_tag _lds_elem_dtype = dtype_to_elem_type(traits.DTYPE_STR) @@ -424,7 +423,6 @@ def flash_attn_dualwave_swp_fp8_gfx950_kernel( ctx.init_descale() ctx.init_tile_bounds() - q_loader = DualwaveFp8QLoader(ctx) gemm_helper = DualwaveFp8GemmHelper(ctx) softmax_helper = DualwaveFp8SoftmaxHelper(ctx) kv_gmem_to_lds = DualwaveFp8KvGmemToLdsLoader(ctx) @@ -445,7 +443,7 @@ def _run_q_block(): # q_start_pos_i32 on ctx for the causal-mask helpers. ctx.init_q_row() q_row = ctx.q_row - q_all_wide = q_loader.load_all_wide(ctx.q_row_in_block) + q_all_wide = gemm_helper.load_q_wide() kv_gmem_to_lds.load_k((ctx.split_t0 + 1) * traits.BLOCK_N, 1) kv_gmem_to_lds.load_v(ctx.split_t0 * traits.BLOCK_N, 0) diff --git a/kernels/attention/flash_attn_utils.py b/kernels/attention/flash_attn_utils.py index c6a624877e..287eafd5ee 100644 --- a/kernels/attention/flash_attn_utils.py +++ b/kernels/attention/flash_attn_utils.py @@ -1832,14 +1832,8 @@ class DualwaveSwpFp8Traits: BATCH_INTERLEAVE_GROUP: int = 1 PAGED: bool = False KV_VECTORIZED: bool = False - QREG: bool = True - FP8_PV_DIRECT: bool = True FP8_PV_SEGMENTED: bool = False - @property - def V_HEAD_DIM(self): - return self.HEAD_DIM_V - @property def cache_tag(self): return ( @@ -2065,7 +2059,7 @@ class PagedDualwaveSwpFp8Traits: """Layouts for native page-64 FP8 QK/PV with register-resident Q.""" HEAD_DIM: int - V_HEAD_DIM: int + HEAD_DIM_V: int D_CHUNKS: int NUM_HEADS_Q: int NUM_HEADS_KV: int @@ -2077,12 +2071,10 @@ class PagedDualwaveSwpFp8Traits: DUALWAVE_SWP_DEBUG_LAZY_COUNTS: bool DUALWAVE_SWP_ENABLE_STAGGER: bool VARLEN: bool - FP8_PV_DIRECT: bool FP8_PV_SEGMENTED: bool FP8_V_H1: int FP8_V_H2: int BN128: bool - QREG: bool DEFAULT_STRIDE_Q_N: int DEFAULT_STRIDE_KV_N: int SMEM_D_RPT: int @@ -2126,10 +2118,6 @@ class PagedDualwaveSwpFp8Traits: XCD_SWIZZLE: bool = False QLDS: bool = False - @property - def HEAD_DIM_V(self): - return self.V_HEAD_DIM - @property def cache_tag(self): return ( @@ -2137,7 +2125,7 @@ def cache_tag(self): self.NUM_HEADS_Q, self.NUM_HEADS_KV, self.HEAD_DIM, - self.V_HEAD_DIM, + self.HEAD_DIM_V, self.WAVES_PER_EU, self.DAZ, self.DUALWAVE_SWP_LAZY_RESCALE, @@ -2198,7 +2186,7 @@ def _make_paged_dualwave_swp_fp8_traits( return PagedDualwaveSwpFp8Traits( HEAD_DIM=head_dim, - V_HEAD_DIM=value_head_dim, + HEAD_DIM_V=value_head_dim, D_CHUNKS=value_head_dim // 32, NUM_HEADS_Q=num_heads, NUM_HEADS_KV=num_kv_heads, @@ -2211,12 +2199,10 @@ def _make_paged_dualwave_swp_fp8_traits( DUALWAVE_SWP_ENABLE_STAGGER=bool(dualwave_swp_enable_stagger), DUALWAVE_SWP_RESCALE_THRESHOLD=rescale_threshold, VARLEN=bool(varlen), - FP8_PV_DIRECT=bool(bn128), FP8_PV_SEGMENTED=fp8_v_h2 > 0, FP8_V_H1=fp8_v_h1, FP8_V_H2=fp8_v_h2, BN128=bool(bn128), - QREG=bool(bn128), DEFAULT_STRIDE_Q_N=num_heads * head_dim, DEFAULT_STRIDE_KV_N=num_kv_heads * head_dim, SMEM_D_RPT=smem_d_rpt, @@ -4607,7 +4593,7 @@ def init_runtime_indices(self): self.stride_o_n_v = fx.Int64(self.stride_o_n) self.stride_kv_n_v = fx.Index(self.stride_kv_n) if const_expr(traits.PAGED): - self.stride_v_n_v = fx.Int64(traits.NUM_HEADS_KV * traits.V_HEAD_DIM) + self.stride_v_n_v = fx.Int64(traits.NUM_HEADS_KV * traits.HEAD_DIM_V) elif traits.HEAD_DIM_V == traits.HEAD_DIM: self.stride_v_n_v = self.stride_kv_n_v self.stride_o_n_v = self.stride_q_n_v @@ -4632,15 +4618,20 @@ def init_lds(self, shared_storage): lds = fx.SharedAllocator().allocate(shared_storage).peek() self.lds = lds self.lds_kv_base_idx = fx.Index(fx.ptrtoint(lds.kv.ptr)) - if const_expr(self.traits.PAGED): - self.lds_kv_base_ptr = buffer_ops.create_llvm_ptr(self.lds_kv_base_idx, address_space=3) - else: - self.lds_kv_base_ptr = lds.kv.ptr.llvm_ptr + self.lds_kv_base_ptr = lds.kv.ptr.llvm_ptr self.lds_vt_base_idx = fx.Index(fx.ptrtoint(lds.vt.ptr)) + self.lds_vt_base_ptr = lds.vt.ptr.llvm_ptr if const_expr(self.traits.PAGED): - self.lds_vt_base_ptr = buffer_ops.create_llvm_ptr(self.lds_vt_base_idx, address_space=3) - else: - self.lds_vt_base_ptr = lds.vt.ptr.llvm_ptr + self.k_lds_i32_tiles = fx.logical_divide( + fx.make_view( + fx.recast_iter(fx.Int32, lds.kv.ptr), fx.make_layout(self.traits.LDS_KV_TOTAL_SIZE // 4, 1) + ), + fx.make_layout(4, 1), + ) + self.v_lds_i32_tiles = fx.logical_divide( + fx.make_view(fx.recast_iter(fx.Int32, lds.vt.ptr), fx.make_layout(self.traits.VT_BF16_TOTAL // 2, 1)), + fx.make_layout(4, 1), + ) if hasattr(lds, "q"): self.lds_q_base_idx = fx.Int64(fx.ptrtoint(lds.q.ptr)) self.lds_q_base_ptr = lds.q.ptr.llvm_ptr @@ -4716,17 +4707,10 @@ def _make_buf_div(tensor, nrec_bytes): # K/V are one-byte OCP FP8 pages. Rebase one descriptor per tile # with 64-bit pointer arithmetic so serving-sized pools beyond the # 4-GiB boundary retain the exact physical page address. - self.buf_flags_i32 = fx.Int32(buffer_ops._get_buffer_flags()) - self.k_base_iter = fx.get_iter(self.K) - self.v_base_iter = fx.get_iter(self.V) - k_page_elems = fx.Int64(traits.PAGE_SIZE * traits.NUM_HEADS_KV * traits.HEAD_DIM) - v_page_elems = fx.Int64(traits.PAGE_SIZE * traits.NUM_HEADS_KV * traits.V_HEAD_DIM) - self.k_page_byte_stride = k_page_elems * fx.Int64(traits.ELEM_BYTES) - self.v_page_byte_stride = v_page_elems * fx.Int64(traits.ELEM_BYTES) - self.k_page_nrec_bytes = fx.Int64(self.k_page_byte_stride) - self.v_page_nrec_bytes = fx.Int64(self.v_page_byte_stride) - self.k_page_layout = fx.make_layout(fx.Int32(k_page_elems), fx.Int32(1)) - self.v_page_layout = fx.make_layout(fx.Int32(v_page_elems), fx.Int32(1)) + self.k_base_iter = fx.recast_iter(fx.Int8, fx.get_iter(self.K)) + self.v_base_iter = fx.recast_iter(fx.Int8, fx.get_iter(self.V)) + self.k_page_bytes = traits.PAGE_SIZE * traits.NUM_HEADS_KV * traits.HEAD_DIM + self.v_page_bytes = traits.PAGE_SIZE * traits.NUM_HEADS_KV * traits.HEAD_DIM_V self.page_indices_div = fx.logical_divide( fx.rocdl.make_buffer_tensor(self.BlockTable), fx.make_layout(1, 1) ) @@ -4734,29 +4718,14 @@ def _make_buf_div(tensor, nrec_bytes): self.page_v1i32 = Vec.make_type(1, fx.Int32) self.page_i32x2_atom = fx.make_copy_atom(fx.rocdl.BufferCopy64b(), fx.Int32) self.page_v2i32 = Vec.make_type(2, fx.Int32) - self.page_base = self.batch_idx * fx.Int64(self.block_table_stride) + self.page_base = fx.Int64(self.batch_idx) * fx.Int64(self.block_table_stride) self.k_div = None self.v_div = None else: - kv_nrec_bytes = as_mlir_value(self.kv_tok_end * self.stride_kv_n_v * eb) - v_nrec_bytes = as_mlir_value(self.kv_tok_end * self.stride_v_n_v * eb) + kv_nrec_bytes = (self.kv_tok_end * self.stride_kv_n_v * eb).ir_value() + v_nrec_bytes = (self.kv_tok_end * self.stride_v_n_v * eb).ir_value() self.k_div = _make_buf_div(self.K, kv_nrec_bytes) self.v_div = _make_buf_div(self.V, v_nrec_bytes) - self.buf_flags_i32 = None - self.k_base_iter = None - self.v_base_iter = None - self.k_page_byte_stride = None - self.v_page_byte_stride = None - self.k_page_nrec_bytes = None - self.v_page_nrec_bytes = None - self.k_page_layout = None - self.v_page_layout = None - self.page_indices_div = None - self.page_i32_atom = None - self.page_v1i32 = None - self.page_i32x2_atom = None - self.page_v2i32 = None - self.page_base = None def init_atoms_and_lds_ptrs(self): traits = self.traits @@ -4896,26 +4865,26 @@ def global_idx_q(self, token_idx, col): return (self.q_tok_base + token_idx) * self.stride_q_n_v + self.q_head_idx * self.traits.HEAD_DIM + col def global_idx_o(self, token_idx, col): - return (self.q_tok_base + token_idx) * self.stride_o_n_v + self.q_head_idx * self.traits.V_HEAD_DIM + col + return (self.q_tok_base + token_idx) * self.stride_o_n_v + self.q_head_idx * self.traits.HEAD_DIM_V + col def load_page_id(self, tile_start): """Load this request's physical page for a 64-token logical KV tile.""" - local_page = fx.Int64(tile_start) // fx.Int64(self.traits.PAGE_SIZE) + local_page = fx.Int64(tile_start) // self.traits.PAGE_SIZE # The software pipeline rounds to an even tile count and at least four # tiles. Redirect those masked padding tiles to entry zero instead of # reading past the final request's page-index allocation. page_idx = self.page_base + local_page - safe_idx = fx.Int64((local_page < self.num_kv_tiles).select(page_idx, fx.Int64(0))) + safe_idx = (local_page < self.num_kv_tiles).select(page_idx, 0) v = fly.copy_atom_call_ssa( [self.page_v1i32], self.page_i32_atom, fx.slice(self.page_indices_div, (None, fx.Int32(safe_idx))), ) - return fx.Int64(fx.Int32(rocdl.readfirstlane(T.i32, as_mlir_value(fx.Int32(Vec(v, (1,), fx.Int32)[0]))))) + return fx.Int64(rocdl.readfirstlane(T.i32, Vec(v, (1,), fx.Int32)[0].ir_value())) def load_page_id_pair(self, tile_start): """Load two adjacent physical page IDs for a BN128 tile pair.""" - local_page = fx.Int32(fx.Int64(tile_start) // fx.Int64(self.traits.PAGE_SIZE)) + local_page = fx.Int32(fx.Int64(tile_start) // self.traits.PAGE_SIZE) page_base = fx.Int32(self.page_base) num_kv_tiles = fx.Int32(self.num_kv_tiles) page_idx = page_base + local_page @@ -4926,51 +4895,27 @@ def load_page_id_pair(self, tile_start): fx.slice(self.page_indices_div, (None, safe_idx)), ) pages = Vec(v, (2,), fx.Int32) - page0 = rocdl.readfirstlane(T.i32, as_mlir_value(fx.Int32(pages[0]))) - page1 = rocdl.readfirstlane(T.i32, as_mlir_value(fx.Int32(pages[1]))) - return fx.Int64(fx.Int32(page0)), fx.Int64(fx.Int32(page1)) + page0 = rocdl.readfirstlane(T.i32, pages[0].ir_value()) + page1 = rocdl.readfirstlane(T.i32, pages[1].ir_value()) + return fx.Int64(page0), fx.Int64(page1) def make_page_view(self, tensor_iter, page_id, *, is_value=False): - if const_expr(is_value): - page_byte_stride = self.v_page_byte_stride - page_nrec_bytes = self.v_page_nrec_bytes - page_layout = self.v_page_layout - else: - page_byte_stride = self.k_page_byte_stride - page_nrec_bytes = self.k_page_nrec_bytes - page_layout = self.k_page_layout - return _make_page_view( - tensor_iter, - tensor_iter.type, - tensor_iter.alignment, - page_id, - page_byte_stride, - page_nrec_bytes, - page_layout, - fx.Int8.ir_type, - self.buf_flags_i32, - ) + page_bytes = self.v_page_bytes if is_value else self.k_page_bytes + # Rebase in 64 bits before constructing the bounded per-page resource. + page_ptr = fx.add_offset(tensor_iter, page_id * page_bytes) + page = fx.make_view(page_ptr, fx.make_layout(page_bytes, 1)) + return fx.logical_divide(fx.rocdl.make_buffer_tensor(page, num_records_bytes=page_bytes), fx.make_layout(1, 1)) def read_i32x8_lds(self, base_ptr, byte_row): - if const_expr(self.traits.PAGED): - words = [] - for w in range_constexpr(8): - p = buffer_ops.get_element_ptr( - base_ptr, - byte_offset=fx.Int32(byte_row + w * 4), - elem_type=T.i8, - ) - words.append(fx.Int32(llvm.LoadOp(T.i32, p, alignment=1).result)) - return Vec.from_elements(words, fx.Int32).ir_value() halves = [] for h in range_constexpr(2): p = buffer_ops.get_element_ptr(base_ptr, byte_offset=fx.Int32(byte_row + h * 16), elem_type=T.i8) halves.append(Vec(llvm.LoadOp(Vec.make_type(4, fx.Int32), p, alignment=16).result)) return halves[0].shuffle(halves[1], [0, 1, 2, 3, 4, 5, 6, 7]).ir_value() - def read_i32x4_lds(self, base_ptr, byte_row): - ptr = buffer_ops.get_element_ptr(base_ptr, byte_offset=fx.Int32(byte_row), elem_type=T.i8) - return llvm.LoadOp(self.v4i32_type, ptr, alignment=16).result + def read_i32x4_lds(self, byte_row): + tile = fx.slice(self.k_lds_i32_tiles, (None, fx.Uint32(byte_row) // 16)) + return fx.memref_load_vec(tile).ir_value() class DualwaveFp8QLoader(DualwaveFp8KernelContext): @@ -4994,15 +4939,8 @@ def load_all_wide(self, q_row_in_block): d_base = self.lane_div_32 * 32 packs = [] for ws in range_constexpr(traits.HEAD_DIM // 64): - q_col = ws * 64 + d_base - if const_expr(self.lds_q_base_ptr is None): - g_idx = self.q_gmem_elem_offset + q_row_in_block * self.stride_q_n_v + q_col - lo = Vec(self.buffer_load_128(g_idx), (4,), fx.Int32) - hi = Vec(self.buffer_load_128(g_idx + 16), (4,), fx.Int32) - packs.append(lo.shuffle(hi, [0, 1, 2, 3, 4, 5, 6, 7]).ir_value()) - else: - byte_row = q_row_in_block * fx.Int64(traits.HEAD_DIM) + q_col - packs.append(self.read_i32x8_lds(self.lds_q_base_ptr, fx.Int32(byte_row))) + byte_row = q_row_in_block * fx.Index(traits.HEAD_DIM) + fx.Index(ws * 64) + d_base + packs.append(self.read_i32x8_lds(self.lds_q_base_ptr, fx.Int32(byte_row))) return packs @@ -5083,7 +5021,7 @@ def _v_concat_i32x8(self, v_v, dc): def _load_v_fp8_segmented_chunk(self, buf_id, dc): traits = self.traits row_stride = traits.FP8_V_ROW_STRIDE - tile_bytes = traits.V_HEAD_DIM * row_stride + tile_bytes = traits.HEAD_DIM_V * row_stride aligned_base = ((self.lds_vt_base_idx + fx.Int64(127)) // fx.Int64(128)) * fx.Int64(128) tile_base = aligned_base + fx.Int64(buf_id * tile_bytes) d_global = dc * traits.D_CHUNK @@ -5098,12 +5036,8 @@ def _load_v_fp8_segmented_chunk(self, buf_id, dc): row_base = tile_base + fx.Int64(segment_base) + d_row * fx.Int64(row_stride) + token_base halves = [] for half in range_constexpr(2): - ptr = buffer_ops.get_element_ptr( - self.lds_vt_base_ptr, - byte_offset=fx.Int32(row_base + fx.Int64(half * 16) - self.lds_vt_base_idx), - elem_type=T.i8, - ) - halves.append(llvm.LoadOp(Vec.make_type(2, fx.Int64), ptr, alignment=16).result) + view = fx.slice(self.v_lds_i32_tiles, (None, fx.Uint32(row_base + half * 16 - self.lds_vt_base_idx) // 16)) + halves.append(fx.memref_load_vec(view)) fx.rocdl.s_waitcnt(lgkmcnt=0) return _concat_vectors(halves[0], halves[1]).bitcast(fx.Int32).ir_value() @@ -5160,7 +5094,7 @@ def qk(self, v_k, q_wide=None): q_w = q_all_wide[ws] v_s_lo = self._mfma_acc_fp8_wide(k_lo[ws], q_w, v_s_lo) v_s_hi = self._mfma_acc_fp8_wide(k_hi[ws], q_w, v_s_hi) - if const_expr(traits.QREG): + if const_expr(not traits.PAGED or traits.BN128): n_ds = const_expr(traits.HEAD_DIM // 64 * 4) n_mfma = const_expr(traits.HEAD_DIM // 64 * 2) rocdl.sched_group_barrier(traits.SCHED_DS_READ_MASK, n_ds // 2, 12) @@ -5184,7 +5118,7 @@ def cast_p_fp8_direct(self, v_p): p_base = pks * 8 f32 += [hi_full[p_base + s] for s in range_constexpr(8)] packed = self._pack_fp8_i32x8(f32) - if const_expr(self.traits.PAGED and self.traits.FP8_PV_DIRECT): + if const_expr(self.traits.PAGED): words = Vec(packed, (8,), fx.Int32) return Vec.from_elements( [words[i] for i in (0, 2, 1, 3, 4, 6, 5, 7)], @@ -5200,7 +5134,7 @@ def _pv_fp8_direct(self, p_fp8, v_v, v_o): return v_o def pv(self, v_p, v_v, v_o): - if const_expr(self.traits.FP8_PV_DIRECT): + if const_expr(not self.traits.PAGED or self.traits.BN128): return self._pv_fp8_direct(v_p, v_v, v_o) return self._pv_fp8(v_p, v_v, v_o) @@ -5231,10 +5165,10 @@ def _load_vectorized(): k_lds_byte_base + (self.wave_id_uni * traits.SMEM_D_RPT + d) * traits.WARP_SIZE * traits.KV_VEC_SIZE ) - lds_addr = rocdl.readfirstlane(T.i32, as_mlir_value(fx.Int32(lds_addr))) + lds_addr = rocdl.readfirstlane(T.i32, fx.Int32(lds_addr).ir_value()) self.buffer_load_lds_128(src_div, lds_addr, src_elem, 0) - if const_expr(traits.BN128 and traits.V_HEAD_DIM == 192): + if const_expr(traits.BN128 and traits.HEAD_DIM_V == 192): @flyc.jit def _load_compact(): @@ -5381,40 +5315,43 @@ def _stage_v_fp8_vectorized_segment(self, tile_start, buf_id, d_offset, segment_ token_groups16 = traits.BLOCK_N // traits.KV_VEC_SIZE flat = self.tid + @flyc.jit def _stage_active(): - n_group16 = self.lane_in_warp // fx.Int64(16) - d_local = self.wave_id * fx.Int64(16) + self.lane_in_warp % fx.Int64(16) - d_col = fx.Int64(d_offset) + d_local - src_elem = ( - self.kv_head_idx * token_groups16 * traits.V_HEAD_DIM * traits.KV_VEC_SIZE - + n_group16 * traits.V_HEAD_DIM * traits.KV_VEC_SIZE - + d_col * traits.KV_VEC_SIZE - ) - reordered = self._permute_v_fp8_vectorized(self.buffer_load_fp8x16(src_div, src_elem)) - dest_group16 = (n_group16 % fx.Int64(2)) * fx.Int64(2) + n_group16 // fx.Int64(2) - aligned_base = ((self.lds_vt_base_idx + fx.Int64(127)) // fx.Int64(128)) * fx.Int64(128) - tile_bytes = traits.V_HEAD_DIM * traits.FP8_V_ROW_STRIDE - dst_byte = ( - aligned_base - + fx.Int64(buf_id * tile_bytes) - + fx.Int64(d_offset * traits.FP8_V_ROW_STRIDE) - + d_local * fx.Int64(traits.FP8_V_ROW_STRIDE) - + dest_group16 * fx.Int64(16) - ) - ptr = buffer_ops.get_element_ptr( - self.lds_vt_base_ptr, - byte_offset=fx.Int32(dst_byte - self.lds_vt_base_idx), - elem_type=T.i8, - ) - llvm.StoreOp(as_mlir_value(reordered), ptr, alignment=16) + if flat < fx.Int64(segment_dim * token_groups16): + n_group16 = self.lane_in_warp // fx.Int64(16) + d_local = self.wave_id * fx.Int64(16) + self.lane_in_warp % fx.Int64(16) + d_col = fx.Int64(d_offset) + d_local + src_elem = ( + self.kv_head_idx * token_groups16 * traits.HEAD_DIM_V * traits.KV_VEC_SIZE + + n_group16 * traits.HEAD_DIM_V * traits.KV_VEC_SIZE + + d_col * traits.KV_VEC_SIZE + ) + reordered = self._permute_v_fp8_vectorized(self.buffer_load_fp8x16(src_div, src_elem)) + dest_group16 = (n_group16 % fx.Int64(2)) * fx.Int64(2) + n_group16 // fx.Int64(2) + aligned_base = ((self.lds_vt_base_idx + fx.Int64(127)) // fx.Int64(128)) * fx.Int64(128) + tile_bytes = traits.HEAD_DIM_V * traits.FP8_V_ROW_STRIDE + dst_byte = ( + aligned_base + + fx.Int64(buf_id * tile_bytes) + + fx.Int64(d_offset * traits.FP8_V_ROW_STRIDE) + + d_local * fx.Int64(traits.FP8_V_ROW_STRIDE) + + dest_group16 * fx.Int64(16) + ) + self._store_v_fp8_lds(reordered, dst_byte - self.lds_vt_base_idx) - scf_if_dispatch(flat < fx.Int64(segment_dim * token_groups16), _stage_active) + _stage_active() def _stage_v_fp8_vectorized_bankpad(self, tile_start, buf_id, page_id=None): """Stage vectorized paged V in the bank-padded direct-FP8 layout.""" src_i32x4 = self._load_v_fp8_vectorized_bankpad_source(tile_start, page_id=page_id) self._store_v_fp8_vectorized_bankpad(src_i32x4, buf_id) + def _store_v_fp8_lds(self, data, byte_offset): + dst = fx.slice(self.v_lds_i32_tiles, (None, fx.Uint32(byte_offset) // 16)) + src = fx.make_rmem_tensor(fx.make_layout(4, 1), fx.Int32) + src.store(data) + fx.copy(fx.make_copy_atom(fx.UniversalCopy128b(), fx.Int32), src, dst) + def _load_v_fp8_vectorized_bankpad_source(self, tile_start, page_id=None): """Issue vectorized V loads for the paged FP8 LDS layout.""" traits = self.traits @@ -5424,13 +5361,13 @@ def _load_v_fp8_vectorized_bankpad_source(self, tile_start, page_id=None): token_groups16 = traits.BLOCK_N // traits.KV_VEC_SIZE # permlane16 requires a full EXEC mask. The first segment has exactly # one vector per CTA thread, so no per-lane bounds branch is needed. - prefix_dim = traits.FP8_V_H1 if traits.FP8_PV_SEGMENTED else traits.V_HEAD_DIM + prefix_dim = traits.FP8_V_H1 if traits.FP8_PV_SEGMENTED else traits.HEAD_DIM_V assert prefix_dim * token_groups16 == traits.BLOCK_SIZE n_group16 = self.lane_in_warp // fx.Int64(16) d_col = self.wave_id * fx.Int64(16) + self.lane_in_warp % fx.Int64(16) src_elem = ( - self.kv_head_idx * token_groups16 * traits.V_HEAD_DIM * traits.KV_VEC_SIZE - + n_group16 * traits.V_HEAD_DIM * traits.KV_VEC_SIZE + self.kv_head_idx * token_groups16 * traits.HEAD_DIM_V * traits.KV_VEC_SIZE + + n_group16 * traits.HEAD_DIM_V * traits.KV_VEC_SIZE + d_col * traits.KV_VEC_SIZE ) source = self.buffer_load_fp8x16(src_div, src_elem) @@ -5459,23 +5396,18 @@ def _store_tail(): def _store_v_fp8_vectorized_bankpad_segment(self, src_i32x4, buf_id, d_offset): traits = self.traits row_stride = traits.FP8_V_ROW_STRIDE - v_tile_bytes = traits.V_HEAD_DIM * row_stride + v_tile_bytes = traits.HEAD_DIM_V * row_stride aligned_base = ((self.lds_vt_base_idx + fx.Int64(127)) // fx.Int64(128)) * fx.Int64(128) dst_base = aligned_base + fx.Int64(buf_id * v_tile_bytes) n_group16 = self.lane_in_warp // fx.Int64(16) d_col = fx.Int64(d_offset) + self.wave_id * fx.Int64(16) + self.lane_in_warp % fx.Int64(16) reordered = self._permute_v_fp8_vectorized(src_i32x4) dest_group16 = (n_group16 % fx.Int64(2)) * fx.Int64(2) + n_group16 // fx.Int64(2) - if const_expr(traits.BN128 and traits.V_HEAD_DIM == 192): + if const_expr(traits.BN128 and traits.HEAD_DIM_V == 192): # XOR 16-byte groups instead of padding each row in the six-slot ring. dest_group16 = dest_group16 ^ ((d_col // fx.Int64(4)) % fx.Int64(4)) dst_byte = dst_base + d_col * fx.Int64(row_stride) + dest_group16 * fx.Int64(16) - ptr = buffer_ops.get_element_ptr( - self.lds_vt_base_ptr, - byte_offset=fx.Int32(dst_byte - self.lds_vt_base_idx), - elem_type=T.i8, - ) - llvm.StoreOp(as_mlir_value(reordered), ptr, alignment=16) + self._store_v_fp8_lds(reordered, dst_byte - self.lds_vt_base_idx) class DualwaveFp8KvLdsToVgprLoader(DualwaveFp8KernelContext): @@ -5500,8 +5432,8 @@ def _read_vec_key(key): dg0 = ws * 4 + self.lane_div_32 * 2 off0 = (k_base + (dg0 * traits.BLOCK_N + key_sigma) * traits.KV_VEC_SIZE) * traits.ELEM_BYTES off1 = (k_base + ((dg0 + 1) * traits.BLOCK_N + key_sigma) * traits.KV_VEC_SIZE) * traits.ELEM_BYTES - lo = Vec(self.read_i32x4_lds(self.lds_kv_base_ptr, off0), (4,), fx.Int32) - hi = Vec(self.read_i32x4_lds(self.lds_kv_base_ptr, off1), (4,), fx.Int32) + lo = Vec(self.read_i32x4_lds(off0), (4,), fx.Int32) + hi = Vec(self.read_i32x4_lds(off1), (4,), fx.Int32) packs.append(lo.shuffle(hi, [0, 1, 2, 3, 4, 5, 6, 7]).ir_value()) return packs @@ -5532,7 +5464,7 @@ def _load_v_fp8_vectorized_bankpad(self, buf_id): """Read the D-major V tile into the FP8 MFMA A fragment.""" traits = self.traits row_stride = traits.FP8_V_ROW_STRIDE - v_tile_bytes = traits.V_HEAD_DIM * row_stride + v_tile_bytes = traits.HEAD_DIM_V * row_stride aligned_base = ((self.lds_vt_base_idx + fx.Int64(127)) // fx.Int64(128)) * fx.Int64(128) tile_base = aligned_base + fx.Int64(buf_id * v_tile_bytes) token_base = self.lane_div_32 * fx.Int64(32) @@ -5542,16 +5474,14 @@ def _load_v_fp8_vectorized_bankpad(self, buf_id): row_base = tile_base + d_row * fx.Int64(row_stride) for half in range_constexpr(2): token_offset = token_base + fx.Int64(half * 16) - if const_expr(traits.BN128 and traits.V_HEAD_DIM == 192): + if const_expr(traits.BN128 and traits.HEAD_DIM_V == 192): token_offset = token_offset ^ (((d_row // fx.Int64(4)) % fx.Int64(4)) * fx.Int64(16)) - ptr = buffer_ops.get_element_ptr( - self.lds_vt_base_ptr, - byte_offset=fx.Int32(row_base + token_offset - self.lds_vt_base_idx), - elem_type=T.i8, + view = fx.slice( + self.v_lds_i32_tiles, (None, fx.Uint32(row_base + token_offset - self.lds_vt_base_idx) // 16) ) - pair = Vec(llvm.LoadOp(Vec.make_type(2, fx.Int64), ptr, alignment=16).result, (2,), fx.Int64) - packs[half * 2][dc] = as_mlir_value(fx.Int64(pair[0])) - packs[half * 2 + 1][dc] = as_mlir_value(fx.Int64(pair[1])) + pair = fx.memref_load_vec(view).bitcast(fx.Int64) + packs[half * 2][dc] = pair[0].ir_value() + packs[half * 2 + 1][dc] = pair[1].ir_value() return packs def load_v_steps(self, buf_id, first_step, num_steps): @@ -5825,28 +5755,26 @@ def _run(v_o, m_row, l_row, m_tile_max, v_p): return _run(v_o, m_row, l_row, m_tile_max, v_p) def _lazy_rescale_o_paged(self, v_o, m_row, l_row, m_tile_max, v_p): - traits = self.traits - @flyc.jit def _run(v_o, m_row, l_row, m_tile_max, v_p): m_diff = m_tile_max - m_row m_diff_scaled = m_diff * self.c_logit_scale below = fx.Float32(m_diff_scaled) <= self.c_rescale_thr_f - ballot = rocdl.ballot(T.i64, as_mlir_value(below)) - all_below = arith.cmpi(arith.CmpIPredicate.eq, as_mlir_value(ballot), _read_exec_i64()) - all_below = llvm.intr_expect(all_below, arith.constant(1, type=ir.IntegerType.get_signless(1))) + ballot = fx.Int64(rocdl.ballot(T.i64, below.ir_value())) + all_below = ballot == fx.Int64(_read_exec_i64()) + all_below = fx.Boolean(llvm.intr_expect(all_below.ir_value(), fx.Boolean(True).ir_value())) - o_out = [as_mlir_value(v_o[dc]) for dc in range(traits.D_CHUNKS)] + o_out = list(v_o) vp_out = self.v_p_to_vec32(v_p) - l_out = as_mlir_value(l_row) - m_out = as_mlir_value(m_row) - if fx.Boolean(all_below): + l_out = l_row + m_out = m_row + if all_below: pass else: m_new, corr, scaled_accs = self._lazy_correction(v_o, m_row, m_tile_max) - o_out = [as_mlir_value(scaled_accs[dc]) for dc in range(traits.D_CHUNKS)] + o_out = scaled_accs vp_out = self.v_p_to_vec32(self.scale_v_p(v_p, corr)) - l_out = as_mlir_value(l_row * corr) + l_out = l_row * corr m_out = self.anchor_scalar_f32(m_new) return (o_out, m_out, l_out, self.v_vec32_to_p(vp_out)) From 7148042e593c2525bf238c017102d06346a6e2d0 Mon Sep 17 00:00:00 2001 From: Xiake Sun Date: Wed, 9 Sep 2026 04:26:12 +0000 Subject: [PATCH 41/59] refactor(fmha): share paged launch and compile preparation - Normalize launch arguments once for cached launches and explicit compilation. - Preserve the compiled launch fast path and existing argument validation. - Check explicit compilation against PyTorch for both paged schedules and all supported widths. Signed-off-by: Xiake Sun --- .../attention/flash_attn_fp8_paged_gfx950.py | 82 ++----------------- tests/kernels/test_flash_attn_fwd.py | 25 +++++- 2 files changed, 30 insertions(+), 77 deletions(-) diff --git a/kernels/attention/flash_attn_fp8_paged_gfx950.py b/kernels/attention/flash_attn_fp8_paged_gfx950.py index 0906877b15..6452d89d37 100644 --- a/kernels/attention/flash_attn_fp8_paged_gfx950.py +++ b/kernels/attention/flash_attn_fp8_paged_gfx950.py @@ -1038,6 +1038,7 @@ def _launch( k_descale=None, v_descale=None, stream=None, + _compile_only=False, ): if stride_kv_n is None: stride_kv_n = DEFAULT_STRIDE_KV_N @@ -1062,7 +1063,7 @@ def _launch( block_table = O if block_table_stride is None: block_table_stride = 0 - if block_table is O: + if block_table is O and not _compile_only: raise ValueError("paged fp8 flash_attn requires block_table") _validate_paged_bn128_launch(batch_size, seq_len_kv, block_table_stride) _validate_batch_interleave_launch(batch_size) @@ -1074,7 +1075,8 @@ def _launch( k_descale = O if v_descale is None: v_descale = O - return _run_compiled( + dispatch = flyc.compile if _compile_only else _run_compiled + return dispatch( launch_flash_attn_dualwave_swp, Q, K, @@ -1098,80 +1100,8 @@ def _launch( fx.Stream(stream), ) - def _compile( - Q, - K, - V, - O, # noqa: E741 - batch_size, - seq_len, - stride_kv_n=None, - stride_q_n=None, - stride_o_n=None, - head_dim_runtime=None, - debug_counts=None, - *, - seq_len_kv=None, - cu_seqlens_q=None, - cu_seqlens_kv=None, - block_table=None, - block_table_stride=None, - q_descale=None, - k_descale=None, - v_descale=None, - stream=None, - ): - if stride_kv_n is None: - stride_kv_n = DEFAULT_STRIDE_KV_N - if stride_q_n is None: - stride_q_n = DEFAULT_STRIDE_Q_N - if stride_o_n is None: - stride_o_n = DEFAULT_STRIDE_O_N - if head_dim_runtime is None: - head_dim_runtime = HEAD_DIM - if seq_len_kv is None: - seq_len_kv = seq_len - if debug_counts is None: - debug_counts = O - if cu_seqlens_q is None: - cu_seqlens_q = O - if cu_seqlens_kv is None: - cu_seqlens_kv = O - if block_table is None: - block_table = O - if block_table_stride is None: - block_table_stride = 0 - _validate_paged_bn128_launch(batch_size, seq_len_kv, block_table_stride) - _validate_batch_interleave_launch(batch_size) - if q_descale is None: - q_descale = O - if k_descale is None: - k_descale = O - if v_descale is None: - v_descale = O - return flyc.compile( - launch_flash_attn_dualwave_swp, - Q, - K, - V, - O, - debug_counts, - cu_seqlens_q, - cu_seqlens_kv, - block_table, - block_table_stride, - q_descale, - k_descale, - v_descale, - batch_size, - seq_len, - seq_len_kv, - stride_q_n, - stride_o_n, - stride_kv_n, - head_dim_runtime, - fx.Stream(stream), - ) + def _compile(*args, **kwargs): + return _launch(*args, _compile_only=True, **kwargs) _launch.compile = _compile diff --git a/tests/kernels/test_flash_attn_fwd.py b/tests/kernels/test_flash_attn_fwd.py index 7048e62273..df980e53d6 100644 --- a/tests/kernels/test_flash_attn_fwd.py +++ b/tests/kernels/test_flash_attn_fwd.py @@ -4800,7 +4800,7 @@ def test_paged_fp8_d192_batch_interleave_group(batch_size, head_dims, expected): @_requires_gfx950 @pytest.mark.parametrize("value_head_dim", [128, 192]) -def test_paged_fp8_d192_batch_interleave_ragged_multiblock_matches_torch(value_head_dim): +def test_paged_fp8_d192_ragged_multiblock_matches_torch(value_head_dim): """The D192 batch-interleaved grid preserves ragged causal q-block mapping.""" test_paged_fp8_asymmetric_value_matches_torch( head_dim=192, @@ -4934,6 +4934,29 @@ def test_paged_fp8_graph_replay_matches_torch(head_dim, value_head_dim): torch.testing.assert_close(output, expected, rtol=2.0e-2, atol=2.0e-2) +@_requires_gfx950 +@pytest.mark.parametrize("head_dim,value_head_dim", [(128, 128), (192, 128), (192, 192)]) +@pytest.mark.parametrize("kv_length", [128, 192]) +def test_paged_fp8_explicit_compile_matches_torch(monkeypatch, head_dim, value_head_dim, kv_length): + """Explicit compilation honors the launch arguments for both paged schedules.""" + build = flash_attn_interface._build_paged_fp8 + + def compile_launcher(**kwargs): + return build(**kwargs).compile + + monkeypatch.setattr(flash_attn_interface, "_build_paged_fp8", compile_launcher) + pages = kv_length // 64 + test_paged_fp8_asymmetric_value_matches_torch( + head_dim=head_dim, + value_head_dim=value_head_dim, + use_non_default_stream=True, + force_internal_copies=False, + query_lengths=[64, 32], + kv_lengths=[kv_length, kv_length - 32], + block_table_rows=[list(reversed(range(pages))), list(range(pages, 2 * pages))], + ) + + @_requires_gfx950 def test_paged_fp8_bn128_launcher_rejects_unsupported_shapes(): """Direct BN128 launches must preserve the single-sequence page-pair contract.""" From 44302e9b6d3024795989289da88ccdb8700d0941 Mon Sep 17 00:00:00 2001 From: Xiake Sun Date: Wed, 9 Sep 2026 04:37:40 +0000 Subject: [PATCH 42/59] docs(fmha): clean up paged FP8 implementation comments - Remove stale BF16 V, self-attention, and per-step schedule narration. - Document aligned LDS views and the remaining intrinsic boundaries. - Keep cache ABI, synchronization, and numerical rationale. Signed-off-by: Xiake Sun --- .../attention/flash_attn_fp8_paged_gfx950.py | 22 ++++------------- kernels/attention/flash_attn_interface.py | 8 +++---- kernels/attention/flash_attn_utils.py | 24 ++++++++++--------- tests/kernels/test_flash_attn_fwd.py | 2 +- 4 files changed, 23 insertions(+), 33 deletions(-) diff --git a/kernels/attention/flash_attn_fp8_paged_gfx950.py b/kernels/attention/flash_attn_fp8_paged_gfx950.py index 6452d89d37..e75c5680f3 100644 --- a/kernels/attention/flash_attn_fp8_paged_gfx950.py +++ b/kernels/attention/flash_attn_fp8_paged_gfx950.py @@ -438,9 +438,7 @@ def _run_q_block(): rocdl.sched_barrier(0) rocdl.s_barrier() - # Load this wave's raw fp8 Q rows for wide QK; q/k descale is applied to - # the fp32 logits after the MFMA. init_q_row sets q_row/q_row_i32/ - # q_start_pos_i32 on ctx for the causal-mask helpers. + # Keep Q in registers; apply Q/K descales in softmax. ctx.init_q_row() q_row = ctx.q_row q_all_wide = gemm_helper.load_q_wide() @@ -479,7 +477,6 @@ def _run_q_block(): init_args.append(ctx.c_zero_v16f32) init_args.append(ctx.v_pair_to_vec32(v_p_0)) - # Software-pipelined KV loop. loop_lb = fx.Int64(3) loop_results = init_args for j, loop_args in range( @@ -528,7 +525,6 @@ def _run_q_block(): rocdl.s_setprio(1) v_o = gemm_helper.pv_step_k(0, v_p_0, v_v, v_o) # Cross-length causal can put a diagonal tile in v_s_1; mask it here. - # Self-attention skips this to keep the existing schedule. v_s_1 = softmax_helper.causal_mask_prologue_if_needed(v_s_1, j_idx - 2, (j_idx - 1) * traits.BLOCK_N) m_tile_max_a = softmax_helper.reduce_max(v_s_1) @@ -553,14 +549,11 @@ def _run_q_block(): v_p_1 = softmax_helper.exp2(v_s_1, 0, 16) _sched_barrier_pairs(traits, 6, 6, 2) - # IGroupLP hint (group 2): 6 MFMA each paired with 3 EXP/TRANS (mask - # 0x400) so the new softmax exp2 stays near its MFMA window. + # Keep softmax EXP groups near their MFMA window. _sched_barrier_exp_pairs(traits, 6, 3, 2) if const_expr(traits.DUALWAVE_SWP_SETPRIO): rocdl.s_setprio(0) - # sched_barrier(0): compiler scheduling fence (mask 0 = nothing - # crosses), pinning s_setprio(0) and the closing s_barrier at the - # cluster boundary. Emits no ISA; the real sync is s_barrier(). + # Fence the closing priority/barrier pair at the cluster boundary. rocdl.sched_barrier(0) rocdl.s_barrier() rocdl.sched_barrier(0) @@ -853,16 +846,12 @@ def _run_q_block(): else: v_o = gemm_helper.pv(v_p_1, v_packs_e13, v_o) - # Normalize by l_row; zero rows become zero instead of NaN. - # HIPREC folds v_descale into the bf16 vt scratch. The direct FP8 - # D128 path keeps raw V and applies its descale once at the end. + # Apply V descale once after normalizing the native FP8 P*V result. inv_l = softmax_helper.safe_l_inv(l_row) inv_l = inv_l * ctx.vd_fp8 softmax_helper.scale_o(v_o, inv_l) - # CLOSE the phase shift: one extra s_barrier on group A (complement of - # the prologue's group-B barrier) realigns the two groups before the - # store. Disabled -> one plain barrier. + # Group A's extra barrier closes the prologue's phase shift before stores. if const_expr(traits.DUALWAVE_SWP_ENABLE_STAGGER): _stagger_extra_barrier_if_zero(ctx.stagger_i32) # group A: +1 s_barrier -> close the shift else: @@ -1048,7 +1037,6 @@ def _launch( stride_o_n = DEFAULT_STRIDE_O_N if head_dim_runtime is None: head_dim_runtime = HEAD_DIM - # seq_len_kv defaults to seq_len (self-attention / equal Q,KV lengths). if seq_len_kv is None: seq_len_kv = seq_len if debug_counts is None: diff --git a/kernels/attention/flash_attn_interface.py b/kernels/attention/flash_attn_interface.py index 2c4d9b82cd..e351f91e5f 100644 --- a/kernels/attention/flash_attn_interface.py +++ b/kernels/attention/flash_attn_interface.py @@ -509,7 +509,7 @@ def _build_paged_fp8( # ── paged-KV native path ──────────────────────────────────────────────────── -# gfx950 dualwave paged-KV currently supports exactly one configuration. +# Native page geometry and generic D192 batch-interleave tuning. _PAGED_PAGE_SIZE = 64 _PAGED_BT_LDS_SIZE = 2048 _PAGED_FP8_BATCH_INTERLEAVE_MAX_GROUP = 8 @@ -517,10 +517,10 @@ def _build_paged_fp8( def _paged_fp8_batch_interleave_group(batch_size: int, head_dims: tuple[int, int]) -> int: - """Choose a divisor-sized batch group for the tuned MiMo D192 launch.""" + """Choose a divisor-sized batch group for generic paged D192.""" if head_dims[0] != 192 or batch_size <= 1: return 1 - # V192 crosses back to the cache-locality-favored original grid above B=16. + # Generic V192 favors the original cache-local grid above B=16. if head_dims == (192, 192) and batch_size > _PAGED_FP8_V192_BATCH_INTERLEAVE_MAX_BATCH: return 1 for group_size in (_PAGED_FP8_BATCH_INTERLEAVE_MAX_GROUP, 4, 2): @@ -566,7 +566,7 @@ def _flydsl_flash_attn_paged( D128 or D192. - Dense 4D Q ``[B, Sq, H, D]``: split-K (num_kv_splits>1) supported (seq_len>=384). - Varlen packed Q ``[total_q, H, D]`` (cu_seqlens_q given): paged K/V looked up - per kv-tile via block_table; split-K not supported (matches dense varlen). + per kv-tile via block_table; paged split-K is not supported. """ if kv_cache_layout not in ("linear", "vectorized"): raise NotImplementedError( diff --git a/kernels/attention/flash_attn_utils.py b/kernels/attention/flash_attn_utils.py index 287eafd5ee..316c4528dc 100644 --- a/kernels/attention/flash_attn_utils.py +++ b/kernels/attention/flash_attn_utils.py @@ -4513,9 +4513,9 @@ def _store_splitk_partial_if_qrow(): class DualwaveFp8KernelContext: """Shared per-kernel state for the gfx950 dualwave fp8 attention helpers. - Mirrors ``DualwaveKernelContext`` but for the fp8 single path: raw fp8 Q/K/V - (i8 buffer views), per-tensor Q/K/V descale scalars applied to the fp32 logits, - and a bf16 ``vt`` LDS scratch for HIPREC PV.""" + Q/K/V remain byte-typed for DMA. Descales apply to accumulated QK scores + and the final P*V output; the ``vt`` allocation holds packed FP8 V bytes. + """ def __init__( self, @@ -4622,6 +4622,8 @@ def init_lds(self, shared_storage): self.lds_vt_base_idx = fx.Index(fx.ptrtoint(lds.vt.ptr)) self.lds_vt_base_ptr = lds.vt.ptr.llvm_ptr if const_expr(self.traits.PAGED): + # The 16-byte tile stride preserves wide LDS loads; i64-element + # views can instead lower to slower paired 64-bit reads. self.k_lds_i32_tiles = fx.logical_divide( fx.make_view( fx.recast_iter(fx.Int32, lds.kv.ptr), fx.make_layout(self.traits.LDS_KV_TOTAL_SIZE // 4, 1) @@ -4837,9 +4839,8 @@ def v_vec32_to_pair(self, v): return _v_vec32_to_pair(v) def bf16_trunc_pack_v8(self, f32_vals): - # HIPREC carries P/V as v8 bf16 regardless of the fp8 element dtype: pack - # 8 f32 -> 4 cvt_pk_bf16 dwords. (The generic _bf16_trunc_pack_v8 branches on - # DTYPE_STR and would take the fp16 path for fp8, so keep this bf16-only pack.) + # Generic paged softmax carries P in BF16 between correction steps. + # The dtype-dispatched pack would incorrectly select FP16 for FP8 input. pairs = [] for j in range_constexpr(4): pairs.append(rocdl.cvt_pk_bf16_f32(f32_vals[j * 2], f32_vals[j * 2 + 1])) @@ -5068,7 +5069,7 @@ def _load_q_wide_lds(self): return packs def _load_q_wide_global(self): - """Pull this lane's Q operands straight from global into VGPRs (head_dim > 128).""" + """Pull this lane's Q operands straight from global into VGPRs.""" traits = self.traits d_base = self.lane_div_32 * 32 packs = [] @@ -5144,11 +5145,10 @@ def __init__(self, ctx): super().__init__(ctx) def load_k(self, tile_start, buf_id, page_id=None): - """DMA one K tile into LDS, one pass per head-dim band. + """DMA K using the vectorized page layout or dense head-dimension bands. - A band's LDS line is this wave's n-rows of `chunk` bytes, row-contiguous, - so the QK read indexes it as (band, row, 64-byte slice). The 64-byte tail - band at head_dim 192 runs on the low 32 lanes and moves no padding. + Dense bands are row-contiguous within each wave. Their D192 tail + uses the low 32 lanes; paged tiles retain the physical D-group mapping. """ traits = self.traits eb = traits.ELEM_BYTES @@ -5284,6 +5284,7 @@ def _stage_v_fp8_vectorized_segmented(self, tile_start, buf_id, page_id=None): def _permute_v_fp8_vectorized(self, src_i32x4): src_words = Vec(src_i32x4, (4,), fx.Int32) + # permlane16_swap returns an LLVM pair; unpack it at this intrinsic boundary. pair_ty = ir.Type.parse("!llvm.struct<(i32, i32)>") pair_lo = rocdl.permlane16_swap( pair_ty, @@ -5762,6 +5763,7 @@ def _run(v_o, m_row, l_row, m_tile_max, v_p): below = fx.Float32(m_diff_scaled) <= self.c_rescale_thr_f ballot = fx.Int64(rocdl.ballot(T.i64, below.ir_value())) all_below = ballot == fx.Int64(_read_exec_i64()) + # Most tiles reuse the running maximum; retain the likely-branch hint. all_below = fx.Boolean(llvm.intr_expect(all_below.ir_value(), fx.Boolean(True).ir_value())) o_out = list(v_o) diff --git a/tests/kernels/test_flash_attn_fwd.py b/tests/kernels/test_flash_attn_fwd.py index df980e53d6..d4c55b3c94 100644 --- a/tests/kernels/test_flash_attn_fwd.py +++ b/tests/kernels/test_flash_attn_fwd.py @@ -4801,7 +4801,7 @@ def test_paged_fp8_d192_batch_interleave_group(batch_size, head_dims, expected): @_requires_gfx950 @pytest.mark.parametrize("value_head_dim", [128, 192]) def test_paged_fp8_d192_ragged_multiblock_matches_torch(value_head_dim): - """The D192 batch-interleaved grid preserves ragged causal q-block mapping.""" + """Paged D192 preserves ragged causal q-block mapping.""" test_paged_fp8_asymmetric_value_matches_torch( head_dim=192, value_head_dim=value_head_dim, From 4a72f6e7f377198b1f4ee6de375afe1857675389 Mon Sep 17 00:00:00 2001 From: Xiake Sun Date: Wed, 9 Sep 2026 05:47:54 +0000 Subject: [PATCH 43/59] refactor(fmha): keep segmented LDS offsets in tile units - Simplify the segmented V chunk address calculation and remove redundant segment offsets. - Keep lane, buffer, and chunk offsets in 16-byte tile units so LLVM shares LDS addresses. - Reduce generic V192 spilling while preserving wide loads and odd-page accuracy. Signed-off-by: Xiake Sun --- kernels/attention/flash_attn_utils.py | 17 +++++------------ 1 file changed, 5 insertions(+), 12 deletions(-) diff --git a/kernels/attention/flash_attn_utils.py b/kernels/attention/flash_attn_utils.py index 316c4528dc..8ad363e0b2 100644 --- a/kernels/attention/flash_attn_utils.py +++ b/kernels/attention/flash_attn_utils.py @@ -5024,20 +5024,13 @@ def _load_v_fp8_segmented_chunk(self, buf_id, dc): row_stride = traits.FP8_V_ROW_STRIDE tile_bytes = traits.HEAD_DIM_V * row_stride aligned_base = ((self.lds_vt_base_idx + fx.Int64(127)) // fx.Int64(128)) * fx.Int64(128) - tile_base = aligned_base + fx.Int64(buf_id * tile_bytes) - d_global = dc * traits.D_CHUNK - if const_expr(d_global < traits.FP8_V_H1): - segment_base = 0 - d_local = d_global - else: - segment_base = traits.FP8_V_H1 * row_stride - d_local = d_global - traits.FP8_V_H1 - token_base = self.lane_div_32 * fx.Int64(32) - d_row = fx.Int64(d_local) + self.lane_mod_32 - row_base = tile_base + fx.Int64(segment_base) + d_row * fx.Int64(row_stride) + token_base + tile_index = fx.Uint32(aligned_base - self.lds_vt_base_idx) // 16 + tile_index = tile_index + fx.Uint32(self.lane_mod_32) * (row_stride // 16) + fx.Uint32(self.lane_div_32) * 2 + tile_index = tile_index + fx.Uint32(buf_id) * (tile_bytes // 16) + chunk_offset = dc * traits.D_CHUNK * row_stride // 16 halves = [] for half in range_constexpr(2): - view = fx.slice(self.v_lds_i32_tiles, (None, fx.Uint32(row_base + half * 16 - self.lds_vt_base_idx) // 16)) + view = fx.slice(self.v_lds_i32_tiles, (None, tile_index + chunk_offset + half)) halves.append(fx.memref_load_vec(view)) fx.rocdl.s_waitcnt(lgkmcnt=0) return _concat_vectors(halves[0], halves[1]).bitcast(fx.Int32).ir_value() From e6e7f86b7b997e3a526b9ef24a992aa2543a2143 Mon Sep 17 00:00:00 2001 From: Xiake Sun Date: Wed, 9 Sep 2026 07:57:27 +0000 Subject: [PATCH 44/59] perf(fmha): share a balanced maximum across paged FP8 tile pairs - Reduce both score tiles through a balanced max tree before one half-wave exchange. - Preserve causal masking, lazy rescaling, and the native FP8 cache layouts. Signed-off-by: Xiake Sun --- kernels/attention/flash_attn_fp8_paged_gfx950.py | 2 +- kernels/attention/flash_attn_utils.py | 8 ++++++++ 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/kernels/attention/flash_attn_fp8_paged_gfx950.py b/kernels/attention/flash_attn_fp8_paged_gfx950.py index e75c5680f3..3d58f85121 100644 --- a/kernels/attention/flash_attn_fp8_paged_gfx950.py +++ b/kernels/attention/flash_attn_fp8_paged_gfx950.py @@ -220,7 +220,7 @@ def _correct_o(v_o, m_row, l_row, m_tile): return v_o, m_new, softmax_helper.apply_l_rescale(l_row, corr) def _merge_tile_max(v_s_a, v_s_b): - m_tile = softmax_helper.max2(softmax_helper.reduce_max(v_s_a), softmax_helper.reduce_max(v_s_b)) + m_tile = softmax_helper.reduce_max_pair(v_s_a, v_s_b) return softmax_helper.floor_masked_max(m_tile) page_t0, page_t1 = ctx.load_page_id_pair(t0 * BN) diff --git a/kernels/attention/flash_attn_utils.py b/kernels/attention/flash_attn_utils.py index 8ad363e0b2..a7f4cb9482 100644 --- a/kernels/attention/flash_attn_utils.py +++ b/kernels/attention/flash_attn_utils.py @@ -5515,6 +5515,14 @@ class DualwaveFp8SoftmaxHelper(DualwaveFp8KernelContext): def __init__(self, ctx): super().__init__(ctx) + def reduce_max_pair(self, v_s_a, v_s_b): + values = [self.c_neg_inf] + for scores in (*v_s_a, *v_s_b): + values += [Vec(scores)[i] for i in range_constexpr(16)] + local_max = _tree_reduce(values, fx.maxnumf) + lhs, rhs = _reduction_pair(local_max) + return fx.maxnumf(lhs, rhs) + def _attn_mask_vec2_imm(self, rel_i32, neg_inf_i32, thr_x, thr_y, x_ref_i32, y_ref_i32): return _attn_mask_vec2_imm(rel_i32, neg_inf_i32, thr_x, thr_y, x_ref_i32, y_ref_i32) From cf79b387fcae3be2d39888de1c62079341f8fb7d Mon Sep 17 00:00:00 2001 From: Xiake Sun Date: Wed, 9 Sep 2026 08:39:50 +0000 Subject: [PATCH 45/59] perf(fmha): split the D192 V128 unmasked attention prefix - Separate below-diagonal pairs from the masked tail to expose QK and max-reduction overlap. - Carry the same K/V ring state and preserve one cross-wave rendezvous per pair. - Keep D128 and V192 on their compact single-loop schedule. Signed-off-by: Xiake Sun --- .../attention/flash_attn_fp8_paged_gfx950.py | 28 ++++++++++++++++--- 1 file changed, 24 insertions(+), 4 deletions(-) diff --git a/kernels/attention/flash_attn_fp8_paged_gfx950.py b/kernels/attention/flash_attn_fp8_paged_gfx950.py index 3d58f85121..50c9b0f913 100644 --- a/kernels/attention/flash_attn_fp8_paged_gfx950.py +++ b/kernels/attention/flash_attn_fp8_paged_gfx950.py @@ -273,7 +273,8 @@ def _ring_wrap(x): init_args += [next_v_a, next_v_b] loop_results = init_args next_v_arg_idx = 3 + D_CHUNKS - for j, loop_args in range(fx.Int64(t0), t_end, fx.Int64(2), init=init_args): + + def _iterate(j, loop_args, do_mask): m_row = loop_args[0] l_row = loop_args[1] v_o = [loop_args[2 + i] for i in range_constexpr(D_CHUNKS)] @@ -297,7 +298,8 @@ def _ring_wrap(x): kv_gmem_to_lds.load_v((j + 4) * BN, f_a_buf, page_id=page_f_a) kv_gmem_to_lds.load_v((j + 5) * BN, f_b_buf, page_id=page_f_b) - v_s_a, v_s_b = _mask_pair(v_s_a, v_s_b, j) + if const_expr(do_mask): + v_s_a, v_s_b = _mask_pair(v_s_a, v_s_b, j) m_tile = _merge_tile_max(v_s_a, v_s_b) v_o, m_new, l_row = _correct_o(v_o, m_row, l_row, m_tile) v_o = softmax_helper.anchor_v_o(v_o) @@ -334,7 +336,8 @@ def _ring_wrap(x): ) v_s_b = gemm_helper.qk(v_k_b, q_wide) kv_gmem_to_lds._store_v_fp8_vectorized_bankpad(next_v_b, nn_b_buf) - v_s_a, v_s_b = _mask_pair(v_s_a, v_s_b, j) + if const_expr(do_mask): + v_s_a, v_s_b = _mask_pair(v_s_a, v_s_b, j) m_tile = _merge_tile_max(v_s_a, v_s_b) v_o, m_new, l_row = _correct_o(v_o, m_row, l_row, m_tile) v_o = softmax_helper.anchor_v_o(v_o) @@ -349,7 +352,24 @@ def _ring_wrap(x): rocdl.sched_barrier(0) rocdl.s_barrier() rocdl.sched_barrier(0) - loop_results = yield next_args + return next_args + + if const_expr(traits.HEAD_DIM == 192 and traits.HEAD_DIM_V == 128): + # The prefix boundary is wave-uniform, not CTA-uniform. Both loops + # must keep the same pair order and one rendezvous per pair. + prefix_end = fx.Int64(ctx.q_start_pos_i32 + ctx.delta_i32) // (2 * BN) * 2 + prefix_end = fx.min(fx.Int64(t_end), fx.max(fx.Int64(t0), prefix_end)) + for j, loop_args in range(fx.Int64(t0), prefix_end, fx.Int64(2), init=init_args): + next_args = _iterate(j, loop_args, False) + loop_results = yield next_args + tail_init = loop_results + for j, loop_args in range(prefix_end, fx.Int64(t_end), fx.Int64(2), init=tail_init): + next_args = _iterate(j, loop_args, True) + loop_results = yield next_args + else: + for j, loop_args in range(fx.Int64(t0), t_end, fx.Int64(2), init=init_args): + next_args = _iterate(j, loop_args, True) + loop_results = yield next_args m_row = loop_results[0] l_row = loop_results[1] v_o = [loop_results[2 + i] for i in range_constexpr(D_CHUNKS)] From 9df3371ada69e7b97813282079f8a683598cc35d Mon Sep 17 00:00:00 2001 From: Xiake Sun Date: Wed, 9 Sep 2026 09:28:09 +0000 Subject: [PATCH 46/59] perf(fmha): select scalar centering for native paged FP8 - Use scalar FP32 FMA for single-sequence D128 and D192 softmax centering. - Keep packed centering for batched D128 and all dense and generic paths. - Localize instruction selection while preserving probability headroom scaling. Signed-off-by: Xiake Sun --- kernels/attention/flash_attn_utils.py | 44 +++++++++++++++++++++++---- 1 file changed, 38 insertions(+), 6 deletions(-) diff --git a/kernels/attention/flash_attn_utils.py b/kernels/attention/flash_attn_utils.py index a7f4cb9482..51e4a00ece 100644 --- a/kernels/attention/flash_attn_utils.py +++ b/kernels/attention/flash_attn_utils.py @@ -394,7 +394,7 @@ def _sub_score_pair(v_s, row_max, fm_fast): return Vec.from_elements(lo_sub, fx.Float32).ir_value(), Vec.from_elements(hi_sub, fx.Float32).ir_value() -def _scale_sub_score_pair(v_s, row_max_raw, scale, zero_f, fm_fast, bias=None): +def _scale_sub_score_pair(v_s, row_max_raw, scale, zero_f, fm_fast, bias=None, scalar_fma=False): """Fused softmax-scale + row-max subtraction (optimization 1-A). Returns ``scale * (v_s - row_max_raw) + bias`` per element via a single FMA @@ -412,10 +412,31 @@ def _scale_sub_score_pair(v_s, row_max_raw, scale, zero_f, fm_fast, bias=None): # exp2 lands on 2**bias * P instead of P, at no extra instruction: the # FMA's addend absorbs it. neg_scaled_max = neg_scaled_max + bias - scale_v = Vec.from_elements([scale], fx.Float32).broadcast_to(16) - nsm_v = Vec.from_elements([neg_scaled_max], fx.Float32).broadcast_to(16) - lo = fx.fma(Vec(s_lo), scale_v, nsm_v, fastmath=fm_fast) - hi = fx.fma(Vec(s_hi), scale_v, nsm_v, fastmath=fm_fast) + if const_expr(scalar_fma): + # Preserve scalar instruction selection for MFMA/VALU co-issue. + + def _center(values): + centered = [] + for i in range_constexpr(16): + centered.append( + fx.Float32( + llvm.inline_asm( + T.f32, + [Vec(values)[i].ir_value(), scale.ir_value(), neg_scaled_max.ir_value()], + "v_fma_f32 $0, $1, $2, $3", + "=v,v,v,v", + has_side_effects=False, + ) + ) + ) + return Vec.from_elements(centered, fx.Float32) + + lo, hi = _center(s_lo), _center(s_hi) + else: + scale_v = Vec.from_elements([scale], fx.Float32).broadcast_to(16) + nsm_v = Vec.from_elements([neg_scaled_max], fx.Float32).broadcast_to(16) + lo = fx.fma(Vec(s_lo), scale_v, nsm_v, fastmath=fm_fast) + hi = fx.fma(Vec(s_hi), scale_v, nsm_v, fastmath=fm_fast) return as_mlir_value(lo), as_mlir_value(hi) @@ -5654,7 +5675,18 @@ def sub_m(self, v_s, row_max): if const_expr(self.traits.DUALWAVE_SWP_LAZY_RESCALE): headroom -= self.traits.DUALWAVE_SWP_RESCALE_THRESHOLD bias = fx.Float32(headroom) if headroom > 0.0 else None - return _scale_sub_score_pair(v_s, row_max, self.c_logit_scale, self.c_zero_f, self.fm_fast, bias) + # Batched D128 retains its tuned packed-FMA schedule. + return _scale_sub_score_pair( + v_s, + row_max, + self.c_logit_scale, + self.c_zero_f, + self.fm_fast, + bias, + scalar_fma=self.traits.PAGED + and self.traits.BN128 + and (self.traits.HEAD_DIM != 128 or not self.traits.VARLEN), + ) def exp2(self, v_s, start, length): return _exp2_score_slice(v_s, start, length) From 9b291b2f59ae93a6a067581d49720b05983e08f0 Mon Sep 17 00:00:00 2001 From: Xiake Sun Date: Wed, 9 Sep 2026 11:29:49 +0000 Subject: [PATCH 47/59] perf(fmha): interleave paired D192/V128 batches in grid X - Enable batch-fast paired-page ordering for the two-request D192/V128 configuration. - Preserve head-fast ordering in the generic fallback and validate packed-varlen group launches. - Cover ragged multiple-query-block batches with shuffled pages on a non-default stream. Signed-off-by: Xiake Sun --- kernels/attention/flash_attn_fp8_paged_gfx950.py | 6 +++--- kernels/attention/flash_attn_interface.py | 2 +- kernels/attention/flash_attn_utils.py | 14 ++++++++++---- tests/kernels/test_flash_attn_fwd.py | 13 +++++++++++++ 4 files changed, 27 insertions(+), 8 deletions(-) diff --git a/kernels/attention/flash_attn_fp8_paged_gfx950.py b/kernels/attention/flash_attn_fp8_paged_gfx950.py index 50c9b0f913..d24722d28d 100644 --- a/kernels/attention/flash_attn_fp8_paged_gfx950.py +++ b/kernels/attention/flash_attn_fp8_paged_gfx950.py @@ -77,8 +77,8 @@ def build_flash_attn_paged_fp8_module( batch_interleave_group = int(batch_interleave_group) if batch_interleave_group < 1: raise ValueError(f"batch_interleave_group must be positive, got {batch_interleave_group}") - if batch_interleave_group > 1 and (paged_bn128 or head_dim != 192): - raise ValueError("batch interleaving is supported only by the generic paged D192 kernel") + if batch_interleave_group > 1 and (head_dim != 192 or (paged_bn128 and not paged_bn128_varlen)): + raise ValueError("batch interleaving requires paged D192 and packed-varlen mode for BN128") assert num_heads % num_kv_heads == 0 traits = _make_paged_dualwave_swp_fp8_traits( num_heads, @@ -949,7 +949,7 @@ def launch_flash_attn_dualwave_swp( head_dim_runtime, value_attrs=kernel_attrs, ).launch( - grid=(NUM_HEADS_Q, num_q_blocks, grid_z), + grid=(NUM_HEADS_Q * BATCH_INTERLEAVE_GROUP, num_q_blocks, grid_z // BATCH_INTERLEAVE_GROUP), block=(BLOCK_SIZE, 1, 1), stream=stream, ) diff --git a/kernels/attention/flash_attn_interface.py b/kernels/attention/flash_attn_interface.py index e351f91e5f..499caf3646 100644 --- a/kernels/attention/flash_attn_interface.py +++ b/kernels/attention/flash_attn_interface.py @@ -793,7 +793,7 @@ def _flydsl_flash_attn_paged( paged_bn128_varlen=paged_bn128_varlen, batch_interleave_group=( _paged_fp8_batch_interleave_group(B, fp8_head_dims) - if not use_bn128 and H == 16 and num_kv_heads == 1 + if H == 16 and num_kv_heads == 1 and (not use_bn128 or (fp8_head_dims == (192, 128) and B == 2)) else 1 ), ) diff --git a/kernels/attention/flash_attn_utils.py b/kernels/attention/flash_attn_utils.py index 51e4a00ece..5830b7e983 100644 --- a/kernels/attention/flash_attn_utils.py +++ b/kernels/attention/flash_attn_utils.py @@ -1023,12 +1023,18 @@ def _init_dualwave_thread_mapping(ctx): ctx.h_idx = linear_wg // num_q_blocks ctx.q_block_idx = linear_wg % num_q_blocks elif const_expr(batch_interleave_group > 1 and not traits.SPLITK): - # Keep a bounded batch group head-fast so the final causal q-block + # Keep a bounded batch group in grid X so the final causal q-block # occupies more CUs without abandoning K/V locality across all batches. linear_head_batch = fx.Index(gpu.block_idx.x) - ctx.h_idx = linear_head_batch % traits.NUM_HEADS_Q - batch_in_group = linear_head_batch // traits.NUM_HEADS_Q - ctx.batch_idx = fx.Int64(gpu.block_idx.z) * batch_interleave_group + batch_in_group + if const_expr(getattr(traits, "PAGED", False) and getattr(traits, "BN128", False)): + ctx.h_idx = linear_head_batch // batch_interleave_group + ctx.batch_idx = ( + fx.Int64(gpu.block_idx.z) * batch_interleave_group + linear_head_batch % batch_interleave_group + ) + else: + ctx.h_idx = linear_head_batch % traits.NUM_HEADS_Q + batch_in_group = linear_head_batch // traits.NUM_HEADS_Q + ctx.batch_idx = fx.Int64(gpu.block_idx.z) * batch_interleave_group + batch_in_group ctx.q_block_idx = fx.Index(gpu.block_idx.y) else: ctx.h_idx = fx.Index(gpu.block_idx.x) diff --git a/tests/kernels/test_flash_attn_fwd.py b/tests/kernels/test_flash_attn_fwd.py index d4c55b3c94..1a5d0e825e 100644 --- a/tests/kernels/test_flash_attn_fwd.py +++ b/tests/kernels/test_flash_attn_fwd.py @@ -4813,6 +4813,19 @@ def test_paged_fp8_d192_ragged_multiblock_matches_torch(value_head_dim): ) +@_requires_gfx950 +def test_paged_fp8_bn128_batch2_ragged_multiblock_matches_torch(): + test_paged_fp8_asymmetric_value_matches_torch( + head_dim=192, + value_head_dim=128, + use_non_default_stream=True, + force_internal_copies=False, + query_lengths=[513, 257], + kv_lengths=[1024, 768], + block_table_rows=[list(reversed(range(16))), list(reversed(range(16, 28))) + [0] * 4], + ) + + @_requires_gfx950 @pytest.mark.parametrize( ("head_dim", "value_head_dim"), From 1c2a27a340c91a96dabc0dc5aa6adeea34b87b2e Mon Sep 17 00:00:00 2001 From: Xiake Sun Date: Thu, 10 Sep 2026 05:42:58 +0000 Subject: [PATCH 48/59] perf(fmha): prune bounded paged FP8 rescaling - Bound future FP8 scores with the query L1 norm and the OCP E4M3 finite range. - Select a max-free loop only when every lane proves lazy rescaling unnecessary; retain the original fallback and cache ABI. - Initialize the first pair without infinite-seed fast math and keep proof comparisons strict. - Add ragged, multi-KV-head, mixed-wave, and rescaling regressions; validate 332 tests and 45 of 48 dense-parity comparisons. Signed-off-by: Xiake Sun --- .../attention/flash_attn_fp8_paged_gfx950.py | 61 +++++++++++++++-- tests/kernels/test_flash_attn_fwd.py | 66 +++++++++++++++++++ 2 files changed, 122 insertions(+), 5 deletions(-) diff --git a/kernels/attention/flash_attn_fp8_paged_gfx950.py b/kernels/attention/flash_attn_fp8_paged_gfx950.py index d24722d28d..22768dcb68 100644 --- a/kernels/attention/flash_attn_fp8_paged_gfx950.py +++ b/kernels/attention/flash_attn_fp8_paged_gfx950.py @@ -25,6 +25,15 @@ from kernels.common.tensor_shim import _run_compiled +def _query_bound_is_safe(ctx, upper_bound, maximum): + # A proof must reject unordered comparisons and preserve overflow behavior. + with fx.fastmath(None): + safe = ((upper_bound - maximum) * ctx.c_logit_scale <= ctx.c_rescale_thr_f) & (ctx.c_logit_scale > 0) + unsafe = safe == fx.Boolean(False) + unsafe_lanes = fx.Int64(rocdl.ballot(fx.Int64.ir_type, unsafe.ir_value())) + return unsafe_lanes == 0 + + def build_flash_attn_paged_fp8_module( num_heads, head_dim, @@ -188,6 +197,9 @@ def flash_attn_paged_fp8_bn128_kernel( BN = traits.BLOCK_N D_CHUNKS = traits.D_CHUNKS NPF = traits.NUM_PREFETCH_K + BOUNDED_MAX = ( + traits.HEAD_DIM == 128 and traits.DUALWAVE_SWP_LAZY_RESCALE and not traits.DUALWAVE_SWP_DEBUG_LAZY_COUNTS + ) t0 = ctx.split_t0 t_end = ctx.split_t_end @@ -234,6 +246,18 @@ def _merge_tile_max(v_s_a, v_s_b): q_wide = gemm_helper.load_q_wide() + upper_bound = ctx.c_zero_f + if const_expr(BOUNDED_MAX): + # E4M3FN finite keys satisfy |K_i| <= 448. Bound every score by + # 2 * 448 * sum(abs(Q_i)); factor two dominates FP32 summation error. + max_key = fx.Vector.filled(8, 0x7E7E7E7E, fx.Int32) + bound_tile = ctx.c_zero_v16f32 + for ws in range_constexpr(2): + absolute_query = fx.Vector(q_wide[ws]) & fx.Int32(0x7F7F7F7F) + bound_tile = gemm_helper._mfma_acc_fp8_wide(max_key, absolute_query, bound_tile) + with fx.fastmath(None): + upper_bound = fx.Float32(softmax_helper.anchor_scalar_f32(fx.Vector(bound_tile)[0] * fx.Float32(2.0))) + page_t2, page_t3 = ctx.load_page_id_pair((t0 + 2) * BN) kv_gmem_to_lds.load_k((t0 + 1) * BN, (t0 + 1) % NPF, page_id=page_t1) kv_gmem_to_lds.load_v(t0 * BN, t0 % NPF, page_id=page_t0) @@ -274,7 +298,7 @@ def _ring_wrap(x): loop_results = init_args next_v_arg_idx = 3 + D_CHUNKS - def _iterate(j, loop_args, do_mask): + def _iterate(j, loop_args, do_mask, skip_max=False, initialize=False): m_row = loop_args[0] l_row = loop_args[1] v_o = [loop_args[2 + i] for i in range_constexpr(D_CHUNKS)] @@ -338,8 +362,14 @@ def _iterate(j, loop_args, do_mask): kv_gmem_to_lds._store_v_fp8_vectorized_bankpad(next_v_b, nn_b_buf) if const_expr(do_mask): v_s_a, v_s_b = _mask_pair(v_s_a, v_s_b, j) - m_tile = _merge_tile_max(v_s_a, v_s_b) - v_o, m_new, l_row = _correct_o(v_o, m_row, l_row, m_tile) + m_new = m_row + if const_expr(initialize): + # Zero O/l need no rescaling; avoid folding -inf through fast math. + with fx.fastmath(None): + m_new = _merge_tile_max(v_s_a, v_s_b) + elif const_expr(not skip_max): + m_tile = _merge_tile_max(v_s_a, v_s_b) + v_o, m_new, l_row = _correct_o(v_o, m_row, l_row, m_tile) v_o = softmax_helper.anchor_v_o(v_o) v_o, l_row = _subtile_tail(v_s_a, v_v_a, v_o, l_row, m_new) @@ -354,7 +384,24 @@ def _iterate(j, loop_args, do_mask): rocdl.sched_barrier(0) return next_args - if const_expr(traits.HEAD_DIM == 192 and traits.HEAD_DIM_V == 128): + if const_expr(BOUNDED_MAX): + first_end = fx.min(fx.Int64(t_end), fx.Int64(t0) + 2) + for j, loop_args in range(fx.Int64(t0), first_end, fx.Int64(2), init=init_args): + next_args = _iterate(j, loop_args, True, initialize=True) + loop_results = yield next_args + sealed = _query_bound_is_safe(ctx, upper_bound, fx.Float32(loop_results[0])) + fast_end = sealed.select(fx.Int64(t_end), first_end) + first_state = loop_results + # Waves can choose different loops, but pair order and barrier count agree. + for j, loop_args in range(first_end, fast_end, fx.Int64(2), init=first_state): + next_args = _iterate(j, loop_args, True, skip_max=True) + loop_results = yield next_args + slow_start = sealed.select(fx.Int64(t_end), first_end) + slow_state = loop_results + for j, loop_args in range(slow_start, fx.Int64(t_end), fx.Int64(2), init=slow_state): + next_args = _iterate(j, loop_args, True) + loop_results = yield next_args + elif const_expr(traits.HEAD_DIM == 192 and traits.HEAD_DIM_V == 128): # The prefix boundary is wave-uniform, not CTA-uniform. Both loops # must keep the same pair order and one rendezvous per pair. prefix_end = fx.Int64(ctx.q_start_pos_i32 + ctx.delta_i32) // (2 * BN) * 2 @@ -949,7 +996,11 @@ def launch_flash_attn_dualwave_swp( head_dim_runtime, value_attrs=kernel_attrs, ).launch( - grid=(NUM_HEADS_Q * BATCH_INTERLEAVE_GROUP, num_q_blocks, grid_z // BATCH_INTERLEAVE_GROUP), + grid=( + NUM_HEADS_Q * BATCH_INTERLEAVE_GROUP, + num_q_blocks, + grid_z // BATCH_INTERLEAVE_GROUP, + ), block=(BLOCK_SIZE, 1, 1), stream=stream, ) diff --git a/tests/kernels/test_flash_attn_fwd.py b/tests/kernels/test_flash_attn_fwd.py index 1a5d0e825e..61a26d1a99 100644 --- a/tests/kernels/test_flash_attn_fwd.py +++ b/tests/kernels/test_flash_attn_fwd.py @@ -4778,6 +4778,72 @@ def test_paged_fp8_d128_bn128_ragged_multiblock_matches_torch(): ) +@_requires_gfx950 +@pytest.mark.parametrize("num_kv_heads", [1, 2]) +@pytest.mark.parametrize("mode", ["bounded", "escape", "mixed-waves", "negative"]) +def test_paged_fp8_d128_query_bound_preserves_rescaling(mode, num_kv_heads): + """A query bound may prune max checks only for every active lane of a wave.""" + torch.manual_seed(29) + query_lengths, kv_lengths = [300, 65], [1024, 512] + query_offsets, kv_offsets = [0, 300, 365], [0, 1024, 1536] + num_pages = sum(kv_lengths) // 64 + physical_pages = torch.randperm(num_pages, device="cuda") + table = torch.zeros(2, 16, device="cuda", dtype=torch.int32) + query = torch.zeros(sum(query_lengths), 16, 128, device="cuda") + key = torch.zeros(num_pages, num_kv_heads, 8, 64, 16, device="cuda") + head_sign = torch.where(torch.arange(num_kv_heads, device="cuda") % 2 == 0, 1.0, -1.0) + page_offset = 0 + for batch, (query_length, kv_length) in enumerate(zip(query_lengths, kv_lengths)): + rows = torch.arange(query_length, device="cuda") + coefficients = torch.ones(query_length, device="cuda") + if mode == "mixed-waves": + coefficients = torch.where((rows // 32) % 2 == 0, 1.0, 4.0) + query[query_offsets[batch] : query_offsets[batch + 1], :, 0] = coefficients[:, None] + pages = kv_length // 64 + selected = physical_pages[page_offset : page_offset + pages] + table[batch, :pages] = selected.to(torch.int32) + pair = torch.arange(pages, device="cuda") // 2 + levels = torch.where(pair % 2 == 1, 1.0, -1.0) + levels[:2] = 0.0 + if mode == "negative": + levels[:] = -1.0 + key[selected, :, 0, :, 0] = levels[:, None, None] * head_sign[None, :, None] * 448.0 + page_offset += pages + + query = query.to(torch.float8_e4m3fn) + key = key.to(torch.float8_e4m3fn) + value, value_descale = quantize_per_tensor_fp8( + torch.randn(num_pages, num_kv_heads, 4, 128, 16, device="cuda") * 0.2 + 0.25 + ) + peak_log2 = {"bounded": 3.0, "escape": 16.0, "mixed-waves": 3.0, "negative": 2.5}[mode] + query_descale = torch.ones(1, device="cuda") + key_descale = torch.tensor([peak_log2 * math.sqrt(128) / (448.0 * math.log2(math.e))], device="cuda") + actual = flydsl_flash_attn_func( + query, + key, + value, + causal=True, + num_kv_heads=num_kv_heads, + cu_seqlens_q=torch.tensor(query_offsets, device="cuda", dtype=torch.int32), + cu_seqlens_kv=torch.tensor(kv_offsets, device="cuda", dtype=torch.int32), + max_seqlen_q=max(query_lengths), + max_seqlen_kv=max(kv_lengths), + cross_seqlen=True, + block_table=table, + seqlen_k=torch.tensor(kv_lengths, device="cuda", dtype=torch.int32), + kv_cache_layout="vectorized", + q_descale=query_descale, + k_descale=key_descale, + v_descale=value_descale, + ) + torch.cuda.synchronize() + expected = _paged_fp8_torch_reference( + query, key, value, table, query_lengths, kv_lengths, (query_descale, key_descale, value_descale) + ) + assert bool(torch.isfinite(actual).all().item()) + torch.testing.assert_close(actual, expected, rtol=2.0e-2, atol=5.0e-3) + + @pytest.mark.parametrize( ("batch_size", "head_dims", "expected"), [ From eb85926bd72be2f5cd5fc74585197651f7b5de98 Mon Sep 17 00:00:00 2001 From: Xiake Sun Date: Thu, 10 Sep 2026 07:36:54 +0000 Subject: [PATCH 49/59] perf(fmha): interleave small paired D128 batches - Use batch-fast grid-X ordering for paired D128 at B=2/3/5 while preserving other dispatch policies and the native cache ABI. - Expand ragged, multi-KV-head, bounded/fallback, mixed-wave, and side-stream-copy regression coverage. - Isolate the stream concurrency test from HIP hardware-queue sharing without changing benchmark priority or assertions. - Validate all 48 matched comparisons at or above dense throughput, with 359 tests passing and four existing skips. Signed-off-by: Xiake Sun --- .../attention/flash_attn_fp8_paged_gfx950.py | 8 ++-- kernels/attention/flash_attn_interface.py | 16 +++++-- tests/kernels/test_flash_attn_fwd.py | 47 +++++++++++++++++-- 3 files changed, 59 insertions(+), 12 deletions(-) diff --git a/kernels/attention/flash_attn_fp8_paged_gfx950.py b/kernels/attention/flash_attn_fp8_paged_gfx950.py index 22768dcb68..c865487f3e 100644 --- a/kernels/attention/flash_attn_fp8_paged_gfx950.py +++ b/kernels/attention/flash_attn_fp8_paged_gfx950.py @@ -86,8 +86,10 @@ def build_flash_attn_paged_fp8_module( batch_interleave_group = int(batch_interleave_group) if batch_interleave_group < 1: raise ValueError(f"batch_interleave_group must be positive, got {batch_interleave_group}") - if batch_interleave_group > 1 and (head_dim != 192 or (paged_bn128 and not paged_bn128_varlen)): - raise ValueError("batch interleaving requires paged D192 and packed-varlen mode for BN128") + if batch_interleave_group > 1 and ( + (head_dim == 128 and not paged_bn128) or (paged_bn128 and not paged_bn128_varlen) + ): + raise ValueError("batch interleaving requires generic D192 or packed-varlen BN128") assert num_heads % num_kv_heads == 0 traits = _make_paged_dualwave_swp_fp8_traits( num_heads, @@ -1072,7 +1074,7 @@ def _validate_paged_bn128_launch(batch_size, seq_len_kv, block_table_stride): def _validate_batch_interleave_launch(batch_size): if int(batch_size) % BATCH_INTERLEAVE_GROUP != 0: raise ValueError( - "paged D192 batch size must be divisible by its interleave group: " + "paged FP8 batch size must be divisible by its interleave group: " f"batch_size={int(batch_size)}, group={BATCH_INTERLEAVE_GROUP}" ) diff --git a/kernels/attention/flash_attn_interface.py b/kernels/attention/flash_attn_interface.py index 499caf3646..aac08d6a5b 100644 --- a/kernels/attention/flash_attn_interface.py +++ b/kernels/attention/flash_attn_interface.py @@ -516,8 +516,10 @@ def _build_paged_fp8( _PAGED_FP8_V192_BATCH_INTERLEAVE_MAX_BATCH = 16 -def _paged_fp8_batch_interleave_group(batch_size: int, head_dims: tuple[int, int]) -> int: - """Choose a divisor-sized batch group for generic paged D192.""" +def _paged_fp8_batch_interleave_group(batch_size: int, head_dims: tuple[int, int], *, paired: bool = False) -> int: + """Choose a divisor-sized group for paged D192 or small paired D128 batches.""" + if paired and head_dims == (128, 128) and batch_size in (2, 3, 5): + return batch_size if head_dims[0] != 192 or batch_size <= 1: return 1 # Generic V192 favors the original cache-local grid above B=16. @@ -792,8 +794,14 @@ def _flydsl_flash_attn_paged( use_bn128=use_bn128, paged_bn128_varlen=paged_bn128_varlen, batch_interleave_group=( - _paged_fp8_batch_interleave_group(B, fp8_head_dims) - if H == 16 and num_kv_heads == 1 and (not use_bn128 or (fp8_head_dims == (192, 128) and B == 2)) + _paged_fp8_batch_interleave_group(B, fp8_head_dims, paired=use_bn128) + if H == 16 + and num_kv_heads == 1 + and ( + not use_bn128 + or (fp8_head_dims == (192, 128) and B == 2) + or (fp8_head_dims == (128, 128) and B in (2, 3, 5)) + ) else 1 ), ) diff --git a/tests/kernels/test_flash_attn_fwd.py b/tests/kernels/test_flash_attn_fwd.py index 61a26d1a99..b583ddcaa9 100644 --- a/tests/kernels/test_flash_attn_fwd.py +++ b/tests/kernels/test_flash_attn_fwd.py @@ -4694,7 +4694,9 @@ def noncontiguous_copy(tensor): ) torch.cuda.synchronize() - stream = torch.cuda.Stream() if use_non_default_stream else None + # HIP streams can share a hardware queue. The blocked-default-stream check + # needs a distinct priority pool so queue sharing cannot serialize the test. + stream = torch.cuda.Stream(priority=-1 if force_internal_copies else 0) if use_non_default_stream else None if stream is not None: torch.cuda.synchronize() if force_internal_copies and stream is not None: @@ -4779,16 +4781,21 @@ def test_paged_fp8_d128_bn128_ragged_multiblock_matches_torch(): @_requires_gfx950 +@pytest.mark.parametrize("batch_size", [2, 3, 5]) @pytest.mark.parametrize("num_kv_heads", [1, 2]) @pytest.mark.parametrize("mode", ["bounded", "escape", "mixed-waves", "negative"]) -def test_paged_fp8_d128_query_bound_preserves_rescaling(mode, num_kv_heads): +def test_paged_fp8_d128_query_bound_preserves_rescaling(mode, num_kv_heads, batch_size): """A query bound may prune max checks only for every active lane of a wave.""" torch.manual_seed(29) - query_lengths, kv_lengths = [300, 65], [1024, 512] - query_offsets, kv_offsets = [0, 300, 365], [0, 1024, 1536] + query_lengths = [300, 65, 257, 33, 127][:batch_size] + kv_lengths = [1024, 512, 768, 256, 384][:batch_size] + query_offsets, kv_offsets = [0], [0] + for query_length, kv_length in zip(query_lengths, kv_lengths): + query_offsets.append(query_offsets[-1] + query_length) + kv_offsets.append(kv_offsets[-1] + kv_length) num_pages = sum(kv_lengths) // 64 physical_pages = torch.randperm(num_pages, device="cuda") - table = torch.zeros(2, 16, device="cuda", dtype=torch.int32) + table = torch.zeros(batch_size, 16, device="cuda", dtype=torch.int32) query = torch.zeros(sum(query_lengths), 16, 128, device="cuda") key = torch.zeros(num_pages, num_kv_heads, 8, 64, 16, device="cuda") head_sign = torch.where(torch.arange(num_kv_heads, device="cuda") % 2 == 0, 1.0, -1.0) @@ -4864,6 +4871,36 @@ def test_paged_fp8_d192_batch_interleave_group(batch_size, head_dims, expected): assert flash_attn_interface._paged_fp8_batch_interleave_group(batch_size, head_dims) == expected +@pytest.mark.parametrize("batch_size", [1, 2, 3, 4, 5, 7, 8, 16, 33]) +def test_paged_fp8_d128_small_batch_interleave_group(batch_size): + choose = flash_attn_interface._paged_fp8_batch_interleave_group + assert choose(batch_size, (128, 128), paired=False) == 1 + expected = batch_size if batch_size in (2, 3, 5) else 1 + assert choose(batch_size, (128, 128), paired=True) == expected + + +@_requires_gfx950 +@pytest.mark.parametrize("batch_size", [3, 5]) +def test_paged_fp8_d128_interleaved_ragged_side_copies(batch_size): + query_lengths = [513, 257, 65, 300, 33][:batch_size] + kv_lengths = [1024, 768, 256, 512, 128][:batch_size] + physical_pages = list(reversed(range(sum(kv_lengths) // 64))) + page_offset, rows = 0, [] + for length in kv_lengths: + count = length // 64 + rows.append(physical_pages[page_offset : page_offset + count] + [0] * (16 - count)) + page_offset += count + test_paged_fp8_asymmetric_value_matches_torch( + head_dim=128, + value_head_dim=128, + use_non_default_stream=True, + force_internal_copies=True, + query_lengths=query_lengths, + kv_lengths=kv_lengths, + block_table_rows=rows, + ) + + @_requires_gfx950 @pytest.mark.parametrize("value_head_dim", [128, 192]) def test_paged_fp8_d192_ragged_multiblock_matches_torch(value_head_dim): From 3cc57af39eb640e282851825877cfe5fe113b9ee Mon Sep 17 00:00:00 2001 From: Xiake Sun Date: Thu, 10 Sep 2026 08:30:33 +0000 Subject: [PATCH 50/59] refactor(fmha): simplify paged FP8 pipeline helpers - Share probability packing and active query-block ordering between paged schedules. - Remove no-op segmented V load-step wrappers while preserving all loads, waits, and rescaling order. - Retain scheduler-sensitive address expressions and vector construction; nine checked ISA files remain byte-identical. - Validate the complete paged subset with 124 passing tests and four existing skips. Signed-off-by: Xiake Sun --- .../attention/flash_attn_fp8_paged_gfx950.py | 96 +++++-------------- kernels/attention/flash_attn_utils.py | 35 ++++--- 2 files changed, 39 insertions(+), 92 deletions(-) diff --git a/kernels/attention/flash_attn_fp8_paged_gfx950.py b/kernels/attention/flash_attn_fp8_paged_gfx950.py index c865487f3e..13ea744f3f 100644 --- a/kernels/attention/flash_attn_fp8_paged_gfx950.py +++ b/kernels/attention/flash_attn_fp8_paged_gfx950.py @@ -173,14 +173,7 @@ def flash_attn_paged_fp8_bn128_kernel( ctx.init_thread_mapping() if const_expr(PAGED_BN128_VARLEN): ctx.init_sequence_lengths() - num_q_blocks = (ctx.seqlen_q_v + traits.BLOCK_M - 1) // traits.BLOCK_M - active_q_block = ctx.q_block_idx < num_q_blocks - reversed_q_block = num_q_blocks - 1 - ctx.q_block_idx - ctx.q_block_idx = active_q_block.select(reversed_q_block, ctx.q_block_idx) - ctx.q_start = ctx.q_block_idx * traits.BLOCK_M - ctx.q_gmem_elem_offset = ( - ctx.q_tok_base + ctx.q_start - ) * ctx.stride_q_n_v + ctx.q_head_idx * traits.HEAD_DIM + ctx.init_varlen_causal_lpt_order() else: ctx.init_causal_lpt_order() ctx.init_sequence_lengths() @@ -213,19 +206,13 @@ def _softmax_part(v_s, l_row, m_new): v_p = gemm_helper.cast_p_fp8_direct(v_p) return v_p, l_row - def _pv_part(v_p, v_v, v_o): - # Keep the post-MFMA accumulators in SSA. Pinning them after each - # subtile lengthens the paged schedule without reducing registers. - return gemm_helper.pv(v_p, v_v, v_o) - def _subtile_tail(v_s, v_v, v_o, l_row, m_new): v_p, l_row = _softmax_part(v_s, l_row, m_new) - v_o = _pv_part(v_p, v_v, v_o) + # Keep the post-MFMA accumulators in SSA. Pinning them after each + # subtile lengthens the paged schedule without reducing registers. + v_o = gemm_helper.pv(v_p, v_v, v_o) return v_o, l_row - def _mask_pair(v_s_a, v_s_b, j): - return softmax_helper.causal_mask_pair_if_needed(v_s_a, v_s_b, j) - def _correct_o(v_o, m_row, l_row, m_tile): if const_expr(traits.DUALWAVE_SWP_LAZY_RESCALE): return softmax_helper.lazy_correct_o(v_o, m_row, l_row, m_tile) @@ -325,17 +312,17 @@ def _iterate(j, loop_args, do_mask, skip_max=False, initialize=False): kv_gmem_to_lds.load_v((j + 5) * BN, f_b_buf, page_id=page_f_b) if const_expr(do_mask): - v_s_a, v_s_b = _mask_pair(v_s_a, v_s_b, j) + v_s_a, v_s_b = softmax_helper.causal_mask_pair_if_needed(v_s_a, v_s_b, j) m_tile = _merge_tile_max(v_s_a, v_s_b) v_o, m_new, l_row = _correct_o(v_o, m_row, l_row, m_tile) v_o = softmax_helper.anchor_v_o(v_o) v_p_a, l_row = _softmax_part(v_s_a, l_row, m_new) v_v_a = kv_lds_to_regs.load_v(a_buf) - v_o = _pv_part(v_p_a, v_v_a, v_o) + v_o = gemm_helper.pv(v_p_a, v_v_a, v_o) v_p_b, l_row = _softmax_part(v_s_b, l_row, m_new) v_v_b = kv_lds_to_regs.load_v(b_buf) - v_o = _pv_part(v_p_b, v_v_b, v_o) + v_o = gemm_helper.pv(v_p_b, v_v_b, v_o) m_row = m_new next_args = [m_row, l_row] + v_o + [nn_a_buf] else: @@ -363,7 +350,7 @@ def _iterate(j, loop_args, do_mask, skip_max=False, initialize=False): v_s_b = gemm_helper.qk(v_k_b, q_wide) kv_gmem_to_lds._store_v_fp8_vectorized_bankpad(next_v_b, nn_b_buf) if const_expr(do_mask): - v_s_a, v_s_b = _mask_pair(v_s_a, v_s_b, j) + v_s_a, v_s_b = softmax_helper.causal_mask_pair_if_needed(v_s_a, v_s_b, j) m_new = m_row if const_expr(initialize): # Zero O/l need no rescaling; avoid folding -inf through fast math. @@ -478,14 +465,7 @@ def flash_attn_dualwave_swp_fp8_gfx950_kernel( ctx.init_sequence_lengths() if const_expr(traits.HEAD_DIM == 192 and traits.BATCH_INTERLEAVE_GROUP > 1): # Issue the longest active q-blocks first within each batch group. - num_q_blocks = (ctx.seqlen_q_v + traits.BLOCK_M - 1) // traits.BLOCK_M - active_q_block = ctx.q_block_idx < num_q_blocks - reversed_q_block = num_q_blocks - 1 - ctx.q_block_idx - ctx.q_block_idx = active_q_block.select(reversed_q_block, ctx.q_block_idx) - ctx.q_start = ctx.q_block_idx * traits.BLOCK_M - ctx.q_gmem_elem_offset = ( - ctx.q_tok_base + ctx.q_start - ) * ctx.stride_q_n_v + ctx.q_head_idx * traits.HEAD_DIM + ctx.init_varlen_causal_lpt_order() ctx.init_descriptors() ctx.init_atoms_and_lds_ptrs() ctx.init_dma_thread_offsets() @@ -580,10 +560,7 @@ def _run_q_block(): rocdl.sched_barrier(0) kv_gmem_to_lds.load_k(j_idx * traits.BLOCK_N, 1) - if const_expr(traits.D_CHUNKS > 4): - v_v = kv_lds_to_regs.load_v_steps(0, 0, 2) - else: - v_v = kv_lds_to_regs.load_v(0) + v_v = kv_lds_to_regs.load_v(0) fx.rocdl.s_waitcnt(lgkmcnt=0) _waitcnt_vm_n(ctx.NUM_DMA_K + ctx.NUM_DMA_V) rocdl.sched_barrier(0) @@ -605,12 +582,11 @@ def _run_q_block(): v_o, m_row, l_row, v_p_0 = softmax_helper.rescale_o(v_o, m_row, l_row, m_tile_max_a, v_p_0) v_o = gemm_helper.pv_step_k(1, v_p_0, v_v, v_o) if const_expr(traits.D_CHUNKS > 4): - v_v_tail = kv_lds_to_regs.load_v_steps(0, 2, 2) v_s_1 = softmax_helper.sub_m(v_s_1, m_row) v_p_1 = softmax_helper.exp2(v_s_1, 0, 16) fx.rocdl.s_waitcnt(lgkmcnt=0) - v_o = gemm_helper.pv_step_k(2, v_p_0, v_v_tail, v_o) - v_o = gemm_helper.pv_step_k(3, v_p_0, v_v_tail, v_o) + v_o = gemm_helper.pv_step_k(2, v_p_0, v_v, v_o) + v_o = gemm_helper.pv_step_k(3, v_p_0, v_v, v_o) else: v_o = gemm_helper.pv_step_k(2, v_p_0, v_v, v_o) v_o = gemm_helper.pv_step_k(3, v_p_0, v_v, v_o) @@ -647,10 +623,7 @@ def _run_q_block(): rocdl.sched_barrier(0) kv_gmem_to_lds.load_k((j_idx + 1) * traits.BLOCK_N, 0) - if const_expr(traits.D_CHUNKS > 4): - v_packs_b = kv_lds_to_regs.load_v_steps(1, 0, 2) - else: - v_packs_b = kv_lds_to_regs.load_v(1) + v_v = kv_lds_to_regs.load_v(1) v_s_0 = softmax_helper.causal_mask_prologue_if_needed( v_s_0, j_idx - 1, @@ -664,7 +637,6 @@ def _run_q_block(): if const_expr(traits.DUALWAVE_SWP_SETPRIO): rocdl.s_setprio(1) - v_v = v_packs_b v_o = gemm_helper.pv_step_k(0, v_p_1, v_v, v_o) m_tile_max_b = softmax_helper.reduce_max(v_s_0) _sched_barrier_pairs(traits, 4, 6, 4) @@ -673,15 +645,13 @@ def _run_q_block(): v_o, m_row, l_row, v_p_1 = softmax_helper.lazy_rescale_o(v_o, m_row, l_row, m_tile_max_b, v_p_1) else: v_o, m_row, l_row, v_p_1 = softmax_helper.rescale_o(v_o, m_row, l_row, m_tile_max_b, v_p_1) - v_v = v_packs_b v_o = gemm_helper.pv_step_k(1, v_p_1, v_v, v_o) if const_expr(traits.D_CHUNKS > 4): - v_v_tail = kv_lds_to_regs.load_v_steps(1, 2, 2) v_s_0 = softmax_helper.sub_m(v_s_0, m_row) v_p_0 = softmax_helper.exp2(v_s_0, 0, 16) fx.rocdl.s_waitcnt(lgkmcnt=0) - v_o = gemm_helper.pv_step_k(2, v_p_1, v_v_tail, v_o) - v_o = gemm_helper.pv_step_k(3, v_p_1, v_v_tail, v_o) + v_o = gemm_helper.pv_step_k(2, v_p_1, v_v, v_o) + v_o = gemm_helper.pv_step_k(3, v_p_1, v_v, v_o) else: v_o = gemm_helper.pv_step_k(2, v_p_1, v_v, v_o) v_o = gemm_helper.pv_step_k(3, v_p_1, v_v, v_o) @@ -728,10 +698,7 @@ def _run_q_block(): rocdl.sched_barrier(0) kv_gmem_to_lds.load_k(max_m1 * traits.BLOCK_N, 1) - if const_expr(traits.D_CHUNKS > 4): - v_packs_e3 = kv_lds_to_regs.load_v_steps(0, 0, 2) - else: - v_packs_e3 = kv_lds_to_regs.load_v(0) + v_packs_e3 = kv_lds_to_regs.load_v(0) v_s_1 = softmax_helper.causal_mask_prologue_if_needed( v_s_1, max_m3, @@ -748,7 +715,6 @@ def _run_q_block(): if const_expr(traits.D_CHUNKS > 4): v_o = gemm_helper.pv_step_k(0, v_p_0, v_packs_e3, v_o) v_o = gemm_helper.pv_step_k(1, v_p_0, v_packs_e3, v_o) - v_packs_e3 = kv_lds_to_regs.load_v_steps(0, 2, 2) fx.rocdl.s_waitcnt(lgkmcnt=0) v_o = gemm_helper.pv_step_k(2, v_p_0, v_packs_e3, v_o) v_o = gemm_helper.pv_step_k(3, v_p_0, v_packs_e3, v_o) @@ -791,10 +757,7 @@ def _run_q_block(): rocdl.s_barrier() rocdl.sched_barrier(0) - if const_expr(traits.D_CHUNKS > 4): - v_packs_e7 = kv_lds_to_regs.load_v_steps(1, 0, 2) - else: - v_packs_e7 = kv_lds_to_regs.load_v(1) + v_packs_e7 = kv_lds_to_regs.load_v(1) v_s_0 = softmax_helper.causal_mask_prologue_if_needed( v_s_0, max_m2, @@ -811,7 +774,6 @@ def _run_q_block(): if const_expr(traits.D_CHUNKS > 4): v_o = gemm_helper.pv_step_k(0, v_p_1, v_packs_e7, v_o) v_o = gemm_helper.pv_step_k(1, v_p_1, v_packs_e7, v_o) - v_packs_e7 = kv_lds_to_regs.load_v_steps(1, 2, 2) fx.rocdl.s_waitcnt(lgkmcnt=0) v_o = gemm_helper.pv_step_k(2, v_p_1, v_packs_e7, v_o) v_o = gemm_helper.pv_step_k(3, v_p_1, v_packs_e7, v_o) @@ -853,10 +815,7 @@ def _run_q_block(): rocdl.s_barrier() rocdl.sched_barrier(0) - if const_expr(traits.D_CHUNKS > 4): - v_packs_e11 = kv_lds_to_regs.load_v_steps(0, 0, 2) - else: - v_packs_e11 = kv_lds_to_regs.load_v(0) + v_packs_e11 = kv_lds_to_regs.load_v(0) v_s_1 = softmax_helper.causal_mask_prologue_if_needed( v_s_1, max_m1, @@ -871,7 +830,6 @@ def _run_q_block(): if const_expr(traits.D_CHUNKS > 4): v_o = gemm_helper.pv_step_k(0, v_p_0, v_packs_e11, v_o) v_o = gemm_helper.pv_step_k(1, v_p_0, v_packs_e11, v_o) - v_packs_e11 = kv_lds_to_regs.load_v_steps(0, 2, 2) fx.rocdl.s_waitcnt(lgkmcnt=0) v_o = gemm_helper.pv_step_k(2, v_p_0, v_packs_e11, v_o) v_o = gemm_helper.pv_step_k(3, v_p_0, v_packs_e11, v_o) @@ -896,10 +854,7 @@ def _run_q_block(): rocdl.s_barrier() rocdl.sched_barrier(0) - if const_expr(traits.D_CHUNKS > 4): - v_packs_e13 = kv_lds_to_regs.load_v_steps(1, 0, 2) - else: - v_packs_e13 = kv_lds_to_regs.load_v(1) + v_packs_e13 = kv_lds_to_regs.load_v(1) fx.rocdl.s_waitcnt(lgkmcnt=0) rocdl.sched_barrier(0) rocdl.s_barrier() @@ -908,7 +863,6 @@ def _run_q_block(): if const_expr(traits.D_CHUNKS > 4): v_o = gemm_helper.pv_step_k(0, v_p_1, v_packs_e13, v_o) v_o = gemm_helper.pv_step_k(1, v_p_1, v_packs_e13, v_o) - v_packs_e13 = kv_lds_to_regs.load_v_steps(1, 2, 2) fx.rocdl.s_waitcnt(lgkmcnt=0) v_o = gemm_helper.pv_step_k(2, v_p_1, v_packs_e13, v_o) v_o = gemm_helper.pv_step_k(3, v_p_1, v_packs_e13, v_o) @@ -1029,13 +983,9 @@ def launch_flash_attn_dualwave_swp( value_attrs=kernel_attrs, ).launch( grid=( - ( - NUM_HEADS_Q * traits.BATCH_INTERLEAVE_GROUP, - num_q_blocks, - grid_z // traits.BATCH_INTERLEAVE_GROUP, - ) - if const_expr(traits.BATCH_INTERLEAVE_GROUP > 1) - else (NUM_HEADS_Q, num_q_blocks, grid_z) + NUM_HEADS_Q * BATCH_INTERLEAVE_GROUP, + num_q_blocks, + grid_z // BATCH_INTERLEAVE_GROUP, ), block=(BLOCK_SIZE, 1, 1), stream=stream, @@ -1068,7 +1018,7 @@ def _validate_paged_bn128_launch(batch_size, seq_len_kv, block_table_stride): ) if block_table_stride < num_kv_pages: raise ValueError( - "paged BN128 block table has too few entries: " f"need {num_kv_pages}, got stride {block_table_stride}" + f"paged BN128 block table has too few entries: need {num_kv_pages}, got stride {block_table_stride}" ) def _validate_batch_interleave_launch(batch_size): diff --git a/kernels/attention/flash_attn_utils.py b/kernels/attention/flash_attn_utils.py index 5830b7e983..2b4c6815d7 100644 --- a/kernels/attention/flash_attn_utils.py +++ b/kernels/attention/flash_attn_utils.py @@ -1043,10 +1043,9 @@ def _init_dualwave_thread_mapping(ctx): ctx.bz_idx = fx.Index(gpu.block_idx.z) ctx.batch_idx = ctx.bz_idx // traits.NUM_KV_SPLITS ctx.split_idx = ctx.bz_idx % traits.NUM_KV_SPLITS - elif const_expr(batch_interleave_group <= 1): - ctx.batch_idx = fx.Int64(gpu.block_idx.z) - ctx.split_idx = None else: + if const_expr(batch_interleave_group <= 1): + ctx.batch_idx = fx.Int64(gpu.block_idx.z) ctx.split_idx = None ctx.tid = fx.Index(gpu.thread_idx.x) @@ -4709,6 +4708,16 @@ def init_sequence_lengths(self): self.kv_gmem_elem_offset = self.kv_tok_base * self.stride_kv_n_v + self.kv_head_idx * traits.HEAD_DIM self.v_gmem_elem_offset = self.kv_tok_base * self.stride_v_n_v + self.kv_head_idx * traits.HEAD_DIM_V + def init_varlen_causal_lpt_order(self): + num_q_blocks = (self.seqlen_q_v + self.traits.BLOCK_M - 1) // self.traits.BLOCK_M + active_q_block = self.q_block_idx < num_q_blocks + reversed_q_block = num_q_blocks - 1 - self.q_block_idx + self.q_block_idx = active_q_block.select(reversed_q_block, self.q_block_idx) + self.q_start = self.q_block_idx * self.traits.BLOCK_M + self.q_gmem_elem_offset = ( + self.q_tok_base + self.q_start + ) * self.stride_q_n_v + self.q_head_idx * self.traits.HEAD_DIM + def init_descriptors(self): traits = self.traits eb = traits.ELEM_BYTES @@ -4985,6 +4994,9 @@ def _p_to_fp8_i32x8(self, v_p): f32 = [] for pk in (p_lo[0], p_lo[1], p_hi[0], p_hi[1]): f32 += self._v8bf16_to_f32(pk) + return self._pack_p_fp8(f32) + + def _pack_p_fp8(self, f32): packed = self._pack_fp8_i32x8(f32) if const_expr(self.traits.PAGED): words = Vec(packed, (8,), fx.Int32) @@ -5138,14 +5150,7 @@ def cast_p_fp8_direct(self, v_p): for pks in range_constexpr(self.traits.PV_K_STEPS): p_base = pks * 8 f32 += [hi_full[p_base + s] for s in range_constexpr(8)] - packed = self._pack_fp8_i32x8(f32) - if const_expr(self.traits.PAGED): - words = Vec(packed, (8,), fx.Int32) - return Vec.from_elements( - [words[i] for i in (0, 2, 1, 3, 4, 6, 5, 7)], - fx.Int32, - ).ir_value() - return packed + return self._pack_p_fp8(f32) def _pv_fp8_direct(self, p_fp8, v_v, v_o): v_o = _anchor_v_o(self.traits, v_o) @@ -5505,14 +5510,6 @@ def _load_v_fp8_vectorized_bankpad(self, buf_id): packs[half * 2 + 1][dc] = pair[1].ir_value() return packs - def load_v_steps(self, buf_id, first_step, num_steps): - """Carry the segmented V buffer for the generic on-demand P*V steps.""" - if const_expr(not self.traits.FP8_PV_SEGMENTED): - raise RuntimeError("load_v_steps requires segmented paged FP8 V") - if const_expr(first_step < 0 or num_steps < 1 or first_step + num_steps > 4): - raise RuntimeError("invalid V K-step range") - return fx.Int64(buf_id) - def _load_v_fp8_block(self, buf_id): traits = self.traits v_tile_bytes = (traits.BLOCK_N // 8) * (traits.HEAD_DIM_V // 16) * 128 From 9e8c5efb91895602ff1b0f6f597844efd4983b63 Mon Sep 17 00:00:00 2001 From: Xiake Sun Date: Thu, 10 Sep 2026 08:30:33 +0000 Subject: [PATCH 51/59] refactor(fmha): centralize paged dispatch policy and test helpers - Centralize paired-batch interleave selection and reuse the validated page count. - Separate reusable numerical checks from pytest entry points and parameterize duplicate ragged cases. - Add 18 paired-policy cases without changing numerical assertions or stream concurrency checks. - Verify dispatch equivalence across 3840 shape, batch, and head combinations. Signed-off-by: Xiake Sun --- kernels/attention/flash_attn_interface.py | 19 +++---- tests/kernels/test_flash_attn_fwd.py | 64 +++++++++++++---------- 2 files changed, 44 insertions(+), 39 deletions(-) diff --git a/kernels/attention/flash_attn_interface.py b/kernels/attention/flash_attn_interface.py index aac08d6a5b..b8c6aa01a8 100644 --- a/kernels/attention/flash_attn_interface.py +++ b/kernels/attention/flash_attn_interface.py @@ -518,8 +518,10 @@ def _build_paged_fp8( def _paged_fp8_batch_interleave_group(batch_size: int, head_dims: tuple[int, int], *, paired: bool = False) -> int: """Choose a divisor-sized group for paged D192 or small paired D128 batches.""" - if paired and head_dims == (128, 128) and batch_size in (2, 3, 5): - return batch_size + if paired: + if head_dims == (128, 128) and batch_size in (2, 3, 5): + return batch_size + return 2 if head_dims == (192, 128) and batch_size == 2 else 1 if head_dims[0] != 192 or batch_size <= 1: return 1 # Generic V192 favors the original cache-local grid above B=16. @@ -750,7 +752,7 @@ def _flydsl_flash_attn_paged( if out is not None: if out.device != q.device: - raise ValueError(f"flydsl_flash_attn_func: paged output must be on {q.device}, " f"got {out.device}") + raise ValueError(f"flydsl_flash_attn_func: paged output must be on {q.device}, got {out.device}") if out.shape != expected_out_shape or not out.is_contiguous(): raise ValueError( f"flydsl_flash_attn_func: paged output must be contiguous with shape {expected_out_shape}, " @@ -774,8 +776,7 @@ def _flydsl_flash_attn_paged( and (not arch.startswith("gfx950") or Sq <= _VARLEN_LIGHT_MAX_SEQ) ) if paged_fp8: - num_kv_pages = (skv + page_size - 1) // page_size - use_bn128 = num_kv_pages % 2 == 0 + use_bn128 = max_kv_pages % 2 == 0 paged_bn128_varlen = use_bn128 and B > 1 paged_setprio = dualwave_swp_setprio and D != 192 # The BF16 phase shift exposes H2 staging and probability-pack @@ -795,13 +796,7 @@ def _flydsl_flash_attn_paged( paged_bn128_varlen=paged_bn128_varlen, batch_interleave_group=( _paged_fp8_batch_interleave_group(B, fp8_head_dims, paired=use_bn128) - if H == 16 - and num_kv_heads == 1 - and ( - not use_bn128 - or (fp8_head_dims == (192, 128) and B == 2) - or (fp8_head_dims == (128, 128) and B in (2, 3, 5)) - ) + if H == 16 and num_kv_heads == 1 else 1 ), ) diff --git a/tests/kernels/test_flash_attn_fwd.py b/tests/kernels/test_flash_attn_fwd.py index b583ddcaa9..1830b09dc6 100644 --- a/tests/kernels/test_flash_attn_fwd.py +++ b/tests/kernels/test_flash_attn_fwd.py @@ -4627,6 +4627,26 @@ def test_paged_fp8_asymmetric_value_matches_torch( query_lengths, kv_lengths, block_table_rows, +): + _check_paged_fp8_matches_torch( + head_dim, + value_head_dim, + use_non_default_stream, + force_internal_copies, + query_lengths, + kv_lengths, + block_table_rows, + ) + + +def _check_paged_fp8_matches_torch( + head_dim, + value_head_dim, + use_non_default_stream, + force_internal_copies, + query_lengths, + kv_lengths, + block_table_rows, num_kv_heads=1, ): """Packed causal FP8 page-64 attention supports native Q/K and V widths.""" @@ -4750,7 +4770,7 @@ def noncontiguous_copy(tensor): @pytest.mark.parametrize("num_kv_heads", [2, 4]) def test_paged_fp8_d192_bn128_multiple_kv_heads(value_head_dim, num_kv_heads): """Compact K tiles and the V192 tail preserve head and ragged-page boundaries.""" - test_paged_fp8_asymmetric_value_matches_torch( + _check_paged_fp8_matches_torch( head_dim=192, value_head_dim=value_head_dim, use_non_default_stream=False, @@ -4767,11 +4787,12 @@ def test_paged_fp8_d192_bn128_multiple_kv_heads(value_head_dim, num_kv_heads): @_requires_gfx950 -def test_paged_fp8_d128_bn128_ragged_multiblock_matches_torch(): +@pytest.mark.parametrize("head_dim,value_head_dim", [(128, 128), (192, 128), (192, 192)]) +def test_paged_fp8_bn128_ragged_multiblock_matches_torch(head_dim, value_head_dim): """The multi-batch BN128 path handles distinct causal offsets and inactive q-blocks.""" - test_paged_fp8_asymmetric_value_matches_torch( - head_dim=128, - value_head_dim=128, + _check_paged_fp8_matches_torch( + head_dim=head_dim, + value_head_dim=value_head_dim, use_non_default_stream=False, force_internal_copies=False, query_lengths=[512, 128], @@ -4872,11 +4893,15 @@ def test_paged_fp8_d192_batch_interleave_group(batch_size, head_dims, expected): @pytest.mark.parametrize("batch_size", [1, 2, 3, 4, 5, 7, 8, 16, 33]) -def test_paged_fp8_d128_small_batch_interleave_group(batch_size): +@pytest.mark.parametrize("head_dims", [(128, 128), (192, 128), (192, 192)]) +def test_paged_fp8_paired_batch_interleave_group(batch_size, head_dims): choose = flash_attn_interface._paged_fp8_batch_interleave_group - assert choose(batch_size, (128, 128), paired=False) == 1 - expected = batch_size if batch_size in (2, 3, 5) else 1 - assert choose(batch_size, (128, 128), paired=True) == expected + if head_dims == (128, 128): + assert choose(batch_size, head_dims, paired=False) == 1 + expected = batch_size if batch_size in (2, 3, 5) else 1 + else: + expected = 2 if head_dims == (192, 128) and batch_size == 2 else 1 + assert choose(batch_size, head_dims, paired=True) == expected @_requires_gfx950 @@ -4890,7 +4915,7 @@ def test_paged_fp8_d128_interleaved_ragged_side_copies(batch_size): count = length // 64 rows.append(physical_pages[page_offset : page_offset + count] + [0] * (16 - count)) page_offset += count - test_paged_fp8_asymmetric_value_matches_torch( + _check_paged_fp8_matches_torch( head_dim=128, value_head_dim=128, use_non_default_stream=True, @@ -4901,24 +4926,9 @@ def test_paged_fp8_d128_interleaved_ragged_side_copies(batch_size): ) -@_requires_gfx950 -@pytest.mark.parametrize("value_head_dim", [128, 192]) -def test_paged_fp8_d192_ragged_multiblock_matches_torch(value_head_dim): - """Paged D192 preserves ragged causal q-block mapping.""" - test_paged_fp8_asymmetric_value_matches_torch( - head_dim=192, - value_head_dim=value_head_dim, - use_non_default_stream=False, - force_internal_copies=False, - query_lengths=[512, 128], - kv_lengths=[1024, 512], - block_table_rows=[list(range(15, -1, -1)), list(range(23, 15, -1)) + [0] * 8], - ) - - @_requires_gfx950 def test_paged_fp8_bn128_batch2_ragged_multiblock_matches_torch(): - test_paged_fp8_asymmetric_value_matches_torch( + _check_paged_fp8_matches_torch( head_dim=192, value_head_dim=128, use_non_default_stream=True, @@ -5062,7 +5072,7 @@ def compile_launcher(**kwargs): monkeypatch.setattr(flash_attn_interface, "_build_paged_fp8", compile_launcher) pages = kv_length // 64 - test_paged_fp8_asymmetric_value_matches_torch( + _check_paged_fp8_matches_torch( head_dim=head_dim, value_head_dim=value_head_dim, use_non_default_stream=True, From cacbae012e7e69ab0eed5807071ecc7748884eba Mon Sep 17 00:00:00 2001 From: Xiake Sun Date: Thu, 10 Sep 2026 08:32:50 +0000 Subject: [PATCH 52/59] docs(fmha): clarify paged FP8 layout and scheduling invariants - Trim duplicated stream and phase-shift narration while retaining the barrier and cache-layout rationale. - Clarify asymmetric V shapes, packed probability order, and on-demand segmented V reads. - Document the intentionally separate tail address expression that avoids added spills. - Verify that the executable AST is unchanged; only comments and docstrings differ. Signed-off-by: Xiake Sun --- kernels/attention/flash_attn_fp8_paged_gfx950.py | 9 ++++----- kernels/attention/flash_attn_interface.py | 15 ++++++--------- kernels/attention/flash_attn_utils.py | 11 ++++++----- 3 files changed, 16 insertions(+), 19 deletions(-) diff --git a/kernels/attention/flash_attn_fp8_paged_gfx950.py b/kernels/attention/flash_attn_fp8_paged_gfx950.py index 13ea744f3f..d5bbca6d90 100644 --- a/kernels/attention/flash_attn_fp8_paged_gfx950.py +++ b/kernels/attention/flash_attn_fp8_paged_gfx950.py @@ -478,8 +478,7 @@ def flash_attn_dualwave_swp_fp8_gfx950_kernel( kv_lds_to_regs = DualwaveFp8KvLdsToVgprLoader(ctx) output_store = DualwaveFp8StoreHelper(ctx) - # Skip packed-varlen q-blocks beyond this request's query length. The - # condition is uniform across the workgroup, so barriers stay balanced. + # The active-query guard is workgroup-uniform, keeping barriers balanced. @flyc.jit def _run_q_block(): kv_gmem_to_lds.load_k(ctx.split_t0 * traits.BLOCK_N, 0) @@ -499,9 +498,9 @@ def _run_q_block(): fx.rocdl.s_waitcnt(lgkmcnt=0) _waitcnt_vm_n(ctx.NUM_DMA_V) - # OPEN the wave-group phase shift: one extra s_barrier on group B + # Group B's extra barrier opens the wave-group phase shift. if const_expr(traits.DUALWAVE_SWP_ENABLE_STAGGER): - _stagger_extra_barrier_if_one(ctx.stagger_i32) # group B: +1 s_barrier -> open the shift + _stagger_extra_barrier_if_one(ctx.stagger_i32) else: rocdl.sched_barrier(0) rocdl.s_barrier() @@ -876,7 +875,7 @@ def _run_q_block(): # Group A's extra barrier closes the prologue's phase shift before stores. if const_expr(traits.DUALWAVE_SWP_ENABLE_STAGGER): - _stagger_extra_barrier_if_zero(ctx.stagger_i32) # group A: +1 s_barrier -> close the shift + _stagger_extra_barrier_if_zero(ctx.stagger_i32) else: rocdl.s_barrier() diff --git a/kernels/attention/flash_attn_interface.py b/kernels/attention/flash_attn_interface.py index b8c6aa01a8..84c2a440c4 100644 --- a/kernels/attention/flash_attn_interface.py +++ b/kernels/attention/flash_attn_interface.py @@ -509,7 +509,7 @@ def _build_paged_fp8( # ── paged-KV native path ──────────────────────────────────────────────────── -# Native page geometry and generic D192 batch-interleave tuning. +# Native page geometry and batch-interleave limits. _PAGED_PAGE_SIZE = 64 _PAGED_BT_LDS_SIZE = 2048 _PAGED_FP8_BATCH_INTERLEAVE_MAX_GROUP = 8 @@ -517,7 +517,7 @@ def _build_paged_fp8( def _paged_fp8_batch_interleave_group(batch_size: int, head_dims: tuple[int, int], *, paired: bool = False) -> int: - """Choose a divisor-sized group for paged D192 or small paired D128 batches.""" + """Choose a batch divisor for the generic or paired-page schedule.""" if paired: if head_dims == (128, 128) and batch_size in (2, 3, 5): return batch_size @@ -835,13 +835,10 @@ def _flydsl_flash_attn_paged( kv_cache_layout=kv_cache_layout, has_bias=bias is not None, ) - # The ambient current stream is already ordered correctly. Enter a - # stream context only when the caller explicitly requests a different - # stream for wrapper-owned casts, copies, and allocations. + # Wrapper-owned copies must follow an explicit launch stream; the ambient + # current stream needs no extra context. stream_context = contextlib.nullcontext() if stream is None else torch.cuda.stream(launch_stream) with stream_context: - # Keep wrapper-owned casts and copies ordered with an explicit - # non-current launch stream. block_table_i32 = ( (block_table if block_table.dtype == torch.int32 else block_table.to(torch.int32)) .contiguous() @@ -995,7 +992,7 @@ def flydsl_flash_attn_func( Varlen: ``[total_q, H, D]`` (packed, cu_seqlens_q required). k: Key tensor. Dense: ``[B, Skv, Hkv, D]``. Varlen: ``[total_kv, Hkv, D]``. - v: Value tensor, same shape as k. + v: Value tensor with k's leading dimensions and its own head width. Paged KV cache: physical K/V cache tensors. Supported ``kv_cache_layout`` values: - ``linear``: 4D paged K/V, ``[NumBlocks, PageSize, NumKVHeads, HeadDim]``. @@ -1004,7 +1001,7 @@ def flydsl_flash_attn_func( - ``vectorized``: aiter-style 5D K/V, where ``K = [NumBlocks, NumKVHeads, HeadDim / kVectorSize, PageSize, kVectorSize]`` and - ``V = [NumBlocks, NumKVHeads, PageSize / kVectorSize, HeadDim, kVectorSize]``. + ``V = [NumBlocks, NumKVHeads, PageSize / kVectorSize, ValueHeadDim, kVectorSize]``. Here ``kVectorSize = 16 / element_size`` (bf16/fp16: 8, fp8: 16); page_size and head_dim must be divisible by it. causal: Bottom-right aligned causal mask when True. diff --git a/kernels/attention/flash_attn_utils.py b/kernels/attention/flash_attn_utils.py index 2b4c6815d7..62a1b5350f 100644 --- a/kernels/attention/flash_attn_utils.py +++ b/kernels/attention/flash_attn_utils.py @@ -4709,6 +4709,7 @@ def init_sequence_lengths(self): self.v_gmem_elem_offset = self.kv_tok_base * self.stride_v_n_v + self.kv_head_idx * traits.HEAD_DIM_V def init_varlen_causal_lpt_order(self): + """Reverse active query blocks without making padded workgroups valid.""" num_q_blocks = (self.seqlen_q_v + self.traits.BLOCK_M - 1) // self.traits.BLOCK_M active_q_block = self.q_block_idx < num_q_blocks reversed_q_block = num_q_blocks - 1 - self.q_block_idx @@ -4997,6 +4998,7 @@ def _p_to_fp8_i32x8(self, v_p): return self._pack_p_fp8(f32) def _pack_p_fp8(self, f32): + # P words must follow vectorized K's four-token groups for paged P*V. packed = self._pack_fp8_i32x8(f32) if const_expr(self.traits.PAGED): words = Vec(packed, (8,), fx.Int32) @@ -5336,6 +5338,7 @@ def _permute_v_fp8_vectorized(self, src_i32x4): ) def _stage_v_fp8_vectorized_segment(self, tile_start, buf_id, d_offset, segment_dim, page_id): + # Sharing the prefix store's address expression increases tail VGPR spills. traits = self.traits src_div = self.make_page_view(self.v_base_iter, page_id, is_value=True) token_groups16 = traits.BLOCK_N // traits.KV_VEC_SIZE @@ -5481,6 +5484,7 @@ def _read_strip(key): def load_v(self, buf_id): if const_expr(self.traits.FP8_PV_SEGMENTED and not self.traits.BN128): + # P*V reads segmented chunks on demand; carry only their LDS slot. return fx.Int64(buf_id) if const_expr(self.traits.PAGED): return self._load_v_fp8_vectorized_bankpad(buf_id) @@ -5558,11 +5562,8 @@ def _causal_mask_inplace(self, v_s, tile_idx): s_lo, s_hi = v_s kv_tile_start = tile_idx * traits.BLOCK_N kv_start_i32 = fx.Int32(kv_tile_start) - # The wide fp8 MFMA exposes eight score elements per half-wave group, - # but after vectorized-K sigma their physical key ids are arranged in - # 4-token groups (0..3, 8..11, ...), exactly like the dense mapping. - # Keep this local to the fp8 path; the bf16 vectorized kernel has a - # different operand/output layout. + # Vectorized-K sigma arranges wide FP8 scores in four-token groups + # (0..3, 8..11, ...), matching dense FP8 but not vectorized BF16. lane_off_i32 = fx.Int32(self.lane_div_32) * fx.Int32(4) # q_row_i32 is set by init_q_row (called after helper construction), so read # it from the live ctx. From 3bbbeef6fd6f40d9647c8bd17049f7d410f4e66f Mon Sep 17 00:00:00 2001 From: Xiake Sun Date: Thu, 10 Sep 2026 14:37:41 +0000 Subject: [PATCH 53/59] refactor(fmha): adopt official LLVM memory primitives - Use generic_load and generic_store for the four typed, 16-byte-aligned paged FP8 LDS access sites. - Remove the temporary register tensor used solely to issue the V LDS store while preserving bounded global loads and all synchronization. - Retain the exact qualified prototype implementation: nine ISA files are byte-identical and all 48 comparisons remain above dense throughput. - Recheck the integrated attention suite with 377 passing tests and four existing skips. Signed-off-by: Xiake Sun --- kernels/attention/flash_attn_utils.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/kernels/attention/flash_attn_utils.py b/kernels/attention/flash_attn_utils.py index 62a1b5350f..6e8b26941f 100644 --- a/kernels/attention/flash_attn_utils.py +++ b/kernels/attention/flash_attn_utils.py @@ -4953,7 +4953,7 @@ def read_i32x8_lds(self, base_ptr, byte_row): def read_i32x4_lds(self, byte_row): tile = fx.slice(self.k_lds_i32_tiles, (None, fx.Uint32(byte_row) // 16)) - return fx.memref_load_vec(tile).ir_value() + return fx.generic_load(fx.get_iter(tile), dtype=fx.Int32, count=4).ir_value() class DualwaveFp8QLoader(DualwaveFp8KernelContext): @@ -5072,7 +5072,7 @@ def _load_v_fp8_segmented_chunk(self, buf_id, dc): halves = [] for half in range_constexpr(2): view = fx.slice(self.v_lds_i32_tiles, (None, tile_index + chunk_offset + half)) - halves.append(fx.memref_load_vec(view)) + halves.append(fx.generic_load(fx.get_iter(view), dtype=fx.Int32, count=4)) fx.rocdl.s_waitcnt(lgkmcnt=0) return _concat_vectors(halves[0], halves[1]).bitcast(fx.Int32).ir_value() @@ -5377,9 +5377,7 @@ def _stage_v_fp8_vectorized_bankpad(self, tile_start, buf_id, page_id=None): def _store_v_fp8_lds(self, data, byte_offset): dst = fx.slice(self.v_lds_i32_tiles, (None, fx.Uint32(byte_offset) // 16)) - src = fx.make_rmem_tensor(fx.make_layout(4, 1), fx.Int32) - src.store(data) - fx.copy(fx.make_copy_atom(fx.UniversalCopy128b(), fx.Int32), src, dst) + fx.generic_store(fx.get_iter(dst), data) def _load_v_fp8_vectorized_bankpad_source(self, tile_start, page_id=None): """Issue vectorized V loads for the paged FP8 LDS layout.""" @@ -5509,7 +5507,7 @@ def _load_v_fp8_vectorized_bankpad(self, buf_id): view = fx.slice( self.v_lds_i32_tiles, (None, fx.Uint32(row_base + token_offset - self.lds_vt_base_idx) // 16) ) - pair = fx.memref_load_vec(view).bitcast(fx.Int64) + pair = fx.generic_load(fx.get_iter(view), dtype=fx.Int32, count=4).bitcast(fx.Int64) packs[half * 2][dc] = pair[0].ir_value() packs[half * 2 + 1][dc] = pair[1].ir_value() return packs From 497c76f980fdc4d8d4fbd9cd37c22535c2764519 Mon Sep 17 00:00:00 2001 From: Xiake Sun Date: Fri, 11 Sep 2026 02:08:01 +0000 Subject: [PATCH 54/59] refactor(fmha): unify paged FP8 attention on official APIs - Replace the staggered odd-page fallback with one bounds-safe paired-page pipeline and remove obsolete helpers, scheduling traits, and dispatch options. - Use fx.gemm with MFMA_Scale, typed LLVM memory primitives, causal selects, and BF16 conversions; replace paged inline assembly with standard constrained-FMA and arithmetic-fence intrinsics. - Preserve the native cache ABI, stream behavior, and strict query-bound proof while adding odd-page, multi-KV-head, eager/lazy, and no-inline-assembly regressions.. Signed-off-by: Xiake Sun --- .../attention/flash_attn_fp8_paged_gfx950.py | 592 ++---------------- kernels/attention/flash_attn_interface.py | 10 - kernels/attention/flash_attn_utils.py | 279 ++------- tests/kernels/test_flash_attn_fwd.py | 34 + 4 files changed, 146 insertions(+), 769 deletions(-) diff --git a/kernels/attention/flash_attn_fp8_paged_gfx950.py b/kernels/attention/flash_attn_fp8_paged_gfx950.py index d5bbca6d90..f8a76301c9 100644 --- a/kernels/attention/flash_attn_fp8_paged_gfx950.py +++ b/kernels/attention/flash_attn_fp8_paged_gfx950.py @@ -15,11 +15,6 @@ DualwaveFp8SoftmaxHelper, DualwaveFp8StoreHelper, _make_paged_dualwave_swp_fp8_traits, - _sched_barrier_exp_pairs, - _sched_barrier_pairs, - _stagger_extra_barrier_if_one, - _stagger_extra_barrier_if_zero, - _waitcnt_vm_n, ) from kernels.common.kernels_common import dtype_to_elem_type from kernels.common.tensor_shim import _run_compiled @@ -57,13 +52,17 @@ def build_flash_attn_paged_fp8_module( paged_bn128_varlen=False, batch_interleave_group=1, ): - """Build the gfx950 packed-varlen page-64 FP8 attention launcher.""" + """Build the gfx950 packed-varlen page-64 FP8 attention launcher. + + Priority/stagger options are accepted for caller compatibility; this + paired-page pipeline uses neither wave-priority nor staggered phases. + """ gpu_arch = get_hip_arch() if value_head_dim is None: value_head_dim = head_dim if not gpu_arch.startswith("gfx950"): - raise RuntimeError(f"flash_attn_dualwave_swp requires gfx950+ (uses ds_read_tr16_b64), got {gpu_arch}") + raise RuntimeError(f"paged FP8 flash attention requires gfx950, got {gpu_arch}") if ( not paged or dtype_str != "fp8" @@ -89,7 +88,7 @@ def build_flash_attn_paged_fp8_module( if batch_interleave_group > 1 and ( (head_dim == 128 and not paged_bn128) or (paged_bn128 and not paged_bn128_varlen) ): - raise ValueError("batch interleaving requires generic D192 or packed-varlen BN128") + raise ValueError("batch interleaving requires general D192 or packed-varlen paired page IDs") assert num_heads % num_kv_heads == 0 traits = _make_paged_dualwave_swp_fp8_traits( num_heads, @@ -100,19 +99,17 @@ def build_flash_attn_paged_fp8_module( daz=daz, dualwave_swp_lazy_rescale=dualwave_swp_lazy_rescale, rescale_threshold=rescale_threshold, - dualwave_swp_setprio=dualwave_swp_setprio, dualwave_swp_debug_lazy_counts=dualwave_swp_debug_lazy_counts, - dualwave_swp_enable_stagger=dualwave_swp_enable_stagger, varlen=not paged_bn128 or paged_bn128_varlen, - bn128=paged_bn128, + paired_page_ids=paged_bn128, batch_interleave_group=batch_interleave_group, ) BLOCK_M = traits.BLOCK_M BLOCK_SIZE = traits.BLOCK_SIZE HEAD_DIM = traits.HEAD_DIM NUM_HEADS_Q = traits.NUM_HEADS_Q - PAGED_BN128 = bool(paged_bn128) - PAGED_BN128_VARLEN = bool(paged_bn128_varlen) + PAIRED_PAGE_IDS = bool(paged_bn128) + USE_CU_SEQLENS = not paged_bn128 or paged_bn128_varlen BATCH_INTERLEAVE_GROUP = traits.BATCH_INTERLEAVE_GROUP DEFAULT_STRIDE_Q_N = traits.DEFAULT_STRIDE_Q_N DEFAULT_STRIDE_O_N = traits.NUM_HEADS_Q * traits.HEAD_DIM_V @@ -143,7 +140,7 @@ def flash_attn_paged_fp8_bn128_kernel( seq_len: fx.Int32, seq_len_kv: fx.Int32, stride_q_n: fx.Int32, - stride_kv_n: fx.Int32, + stride_o_n: fx.Int32, head_dim_runtime: fx.Int32, ): ctx = DualwaveFp8KernelContext( @@ -161,9 +158,9 @@ def flash_attn_paged_fp8_bn128_kernel( seq_len, seq_len_kv, stride_q_n, - stride_kv_n, + DEFAULT_STRIDE_KV_N, head_dim_runtime, - stride_o_n=DEFAULT_STRIDE_O_N, + stride_o_n=DEFAULT_STRIDE_O_N if PAIRED_PAGE_IDS else stride_o_n, BlockTable=BlockTable, block_table_stride=block_table_stride, ) @@ -171,7 +168,7 @@ def flash_attn_paged_fp8_bn128_kernel( ctx.init_runtime_indices() ctx.init_lds(SharedStorage) ctx.init_thread_mapping() - if const_expr(PAGED_BN128_VARLEN): + if const_expr(USE_CU_SEQLENS): ctx.init_sequence_lengths() ctx.init_varlen_causal_lpt_order() else: @@ -416,476 +413,6 @@ def _iterate(j, loop_args, do_mask, skip_max=False, initialize=False): rocdl.s_barrier() output_store.store_final_o(v_o, q_row) - @flyc.kernel(known_block_size=[BLOCK_SIZE, 1, 1]) - def flash_attn_dualwave_swp_fp8_gfx950_kernel( - Q: fx.Tensor, - K: fx.Tensor, - V: fx.Tensor, - O: fx.Tensor, # noqa: E741 - DebugCounts: fx.Tensor, - CuSeqQ: fx.Tensor, - CuSeqKv: fx.Tensor, - BlockTable: fx.Tensor, - block_table_stride: fx.Int32, - QDescale: fx.Tensor, - KDescale: fx.Tensor, - VDescale: fx.Tensor, - seq_len: fx.Int32, - seq_len_kv: fx.Int32, - stride_q_n: fx.Int32, - stride_o_n: fx.Int32, - stride_kv_n: fx.Int32, - head_dim_runtime: fx.Int32, - ): - ctx = DualwaveFp8KernelContext( - traits, - Q, - K, - V, - O, - DebugCounts, - CuSeqQ, - CuSeqKv, - QDescale, - KDescale, - VDescale, - seq_len, - seq_len_kv, - stride_q_n, - stride_kv_n, - head_dim_runtime, - stride_o_n=stride_o_n, - BlockTable=BlockTable, - block_table_stride=block_table_stride, - ) - ctx.init_types_and_constants() - ctx.init_runtime_indices() - ctx.init_lds(SharedStorage) - ctx.init_thread_mapping() - ctx.init_sequence_lengths() - if const_expr(traits.HEAD_DIM == 192 and traits.BATCH_INTERLEAVE_GROUP > 1): - # Issue the longest active q-blocks first within each batch group. - ctx.init_varlen_causal_lpt_order() - ctx.init_descriptors() - ctx.init_atoms_and_lds_ptrs() - ctx.init_dma_thread_offsets() - ctx.init_descale() - ctx.init_tile_bounds() - - gemm_helper = DualwaveFp8GemmHelper(ctx) - softmax_helper = DualwaveFp8SoftmaxHelper(ctx) - kv_gmem_to_lds = DualwaveFp8KvGmemToLdsLoader(ctx) - kv_lds_to_regs = DualwaveFp8KvLdsToVgprLoader(ctx) - output_store = DualwaveFp8StoreHelper(ctx) - - # The active-query guard is workgroup-uniform, keeping barriers balanced. - @flyc.jit - def _run_q_block(): - kv_gmem_to_lds.load_k(ctx.split_t0 * traits.BLOCK_N, 0) - fx.rocdl.s_waitcnt(vmcnt=0, lgkmcnt=0, expcnt=0) - rocdl.sched_barrier(0) - rocdl.s_barrier() - - # Keep Q in registers; apply Q/K descales in softmax. - ctx.init_q_row() - q_row = ctx.q_row - q_all_wide = gemm_helper.load_q_wide() - - kv_gmem_to_lds.load_k((ctx.split_t0 + 1) * traits.BLOCK_N, 1) - kv_gmem_to_lds.load_v(ctx.split_t0 * traits.BLOCK_N, 0) - v_k = kv_lds_to_regs.load_k(0) - rocdl.sched_barrier(0) - fx.rocdl.s_waitcnt(lgkmcnt=0) - _waitcnt_vm_n(ctx.NUM_DMA_V) - - # Group B's extra barrier opens the wave-group phase shift. - if const_expr(traits.DUALWAVE_SWP_ENABLE_STAGGER): - _stagger_extra_barrier_if_one(ctx.stagger_i32) - else: - rocdl.sched_barrier(0) - rocdl.s_barrier() - - v_s_0 = gemm_helper.qk(v_k, q_all_wide) - rocdl.sched_barrier(0) - v_s_0 = softmax_helper.causal_mask_prologue_if_needed(v_s_0) - m_row_pro = softmax_helper.reduce_max(v_s_0) - # Floor fully-masked rows (-inf) to finite so exp2 yields 0, not NaN. - m_row_pro = softmax_helper.floor_masked_max(m_row_pro) - v_s_0 = softmax_helper.sub_m(v_s_0, m_row_pro) - v_p_0 = softmax_helper.exp2(v_s_0, 0, 16) - rocdl.sched_barrier(0) - rocdl.s_barrier() - rocdl.sched_barrier(0) - - kv_gmem_to_lds.load_k((ctx.split_t0 + 2) * traits.BLOCK_N, 0) - - l_row_init = ctx.c_zero_f - init_args = [m_row_pro, l_row_init] - for _ in range_constexpr(traits.D_CHUNKS): - init_args.append(ctx.c_zero_v16f32) - init_args.append(ctx.v_pair_to_vec32(v_p_0)) - - loop_lb = fx.Int64(3) - loop_results = init_args - for j, loop_args in range( - loop_lb, - ctx.split_t_end - 1, - fx.Int64(2), - init=init_args, - ): - m_row = loop_args[0] - l_row = loop_args[1] - v_o = [loop_args[2 + i] for i in range_constexpr(traits.D_CHUNKS)] - v_p_0 = ctx.v_vec32_to_pair(loop_args[2 + traits.D_CHUNKS]) - j_idx = j - - kv_gmem_to_lds.load_v((j_idx - 2) * traits.BLOCK_N, 1) - v_k = kv_lds_to_regs.load_k(1) - fx.rocdl.s_waitcnt(lgkmcnt=0) - _waitcnt_vm_n(ctx.NUM_DMA_K + ctx.NUM_DMA_V) - rocdl.sched_barrier(0) - rocdl.s_barrier() - rocdl.sched_barrier(0) - - v_s_1 = gemm_helper.qk(v_k, q_all_wide) - v_p_0 = softmax_helper.exp2(v_p_0, 16, 16) - l_row = softmax_helper.reduce_sum(l_row, v_p_0) - v_p_0 = softmax_helper.cast_p(v_p_0) - v_p_0 = softmax_helper.anchor_v_p(v_p_0) - _sched_barrier_exp_pairs(traits, 6, 3, 1) - _sched_barrier_pairs(traits, 10, 5, 1) - rocdl.sched_barrier(0) - rocdl.s_barrier() - rocdl.sched_barrier(0) - - kv_gmem_to_lds.load_k(j_idx * traits.BLOCK_N, 1) - v_v = kv_lds_to_regs.load_v(0) - fx.rocdl.s_waitcnt(lgkmcnt=0) - _waitcnt_vm_n(ctx.NUM_DMA_K + ctx.NUM_DMA_V) - rocdl.sched_barrier(0) - rocdl.s_barrier() - rocdl.sched_barrier(0) - - if const_expr(traits.DUALWAVE_SWP_SETPRIO): - rocdl.s_setprio(1) - v_o = gemm_helper.pv_step_k(0, v_p_0, v_v, v_o) - # Cross-length causal can put a diagonal tile in v_s_1; mask it here. - v_s_1 = softmax_helper.causal_mask_prologue_if_needed(v_s_1, j_idx - 2, (j_idx - 1) * traits.BLOCK_N) - m_tile_max_a = softmax_helper.reduce_max(v_s_1) - - _sched_barrier_pairs(traits, 4, 6, 2) - - if const_expr(traits.DUALWAVE_SWP_LAZY_RESCALE): - v_o, m_row, l_row, v_p_0 = softmax_helper.lazy_rescale_o(v_o, m_row, l_row, m_tile_max_a, v_p_0) - else: - v_o, m_row, l_row, v_p_0 = softmax_helper.rescale_o(v_o, m_row, l_row, m_tile_max_a, v_p_0) - v_o = gemm_helper.pv_step_k(1, v_p_0, v_v, v_o) - if const_expr(traits.D_CHUNKS > 4): - v_s_1 = softmax_helper.sub_m(v_s_1, m_row) - v_p_1 = softmax_helper.exp2(v_s_1, 0, 16) - fx.rocdl.s_waitcnt(lgkmcnt=0) - v_o = gemm_helper.pv_step_k(2, v_p_0, v_v, v_o) - v_o = gemm_helper.pv_step_k(3, v_p_0, v_v, v_o) - else: - v_o = gemm_helper.pv_step_k(2, v_p_0, v_v, v_o) - v_o = gemm_helper.pv_step_k(3, v_p_0, v_v, v_o) - v_s_1 = softmax_helper.sub_m(v_s_1, m_row) - v_p_1 = softmax_helper.exp2(v_s_1, 0, 16) - - _sched_barrier_pairs(traits, 6, 6, 2) - # Keep softmax EXP groups near their MFMA window. - _sched_barrier_exp_pairs(traits, 6, 3, 2) - if const_expr(traits.DUALWAVE_SWP_SETPRIO): - rocdl.s_setprio(0) - # Fence the closing priority/barrier pair at the cluster boundary. - rocdl.sched_barrier(0) - rocdl.s_barrier() - rocdl.sched_barrier(0) - - kv_gmem_to_lds.load_v((j_idx - 1) * traits.BLOCK_N, 0) - v_k = kv_lds_to_regs.load_k(0) - fx.rocdl.s_waitcnt(lgkmcnt=0) - _waitcnt_vm_n(ctx.NUM_DMA_K + ctx.NUM_DMA_V) - rocdl.sched_barrier(0) - rocdl.s_barrier() - rocdl.sched_barrier(0) - - v_s_0 = gemm_helper.qk(v_k, q_all_wide) - v_p_1 = softmax_helper.exp2(v_p_1, 16, 16) - l_row = softmax_helper.reduce_sum(l_row, v_p_1) - v_p_1 = softmax_helper.cast_p(v_p_1) - v_p_1 = softmax_helper.anchor_v_p(v_p_1) - _sched_barrier_exp_pairs(traits, 6, 3, 3) - _sched_barrier_pairs(traits, 10, 5, 3) - rocdl.sched_barrier(0) - rocdl.s_barrier() - rocdl.sched_barrier(0) - - kv_gmem_to_lds.load_k((j_idx + 1) * traits.BLOCK_N, 0) - v_v = kv_lds_to_regs.load_v(1) - v_s_0 = softmax_helper.causal_mask_prologue_if_needed( - v_s_0, - j_idx - 1, - j_idx * traits.BLOCK_N, - ) - fx.rocdl.s_waitcnt(lgkmcnt=0) - _waitcnt_vm_n(ctx.NUM_DMA_K + ctx.NUM_DMA_V) - rocdl.sched_barrier(0) - rocdl.s_barrier() - rocdl.sched_barrier(0) - - if const_expr(traits.DUALWAVE_SWP_SETPRIO): - rocdl.s_setprio(1) - v_o = gemm_helper.pv_step_k(0, v_p_1, v_v, v_o) - m_tile_max_b = softmax_helper.reduce_max(v_s_0) - _sched_barrier_pairs(traits, 4, 6, 4) - - if const_expr(traits.DUALWAVE_SWP_LAZY_RESCALE): - v_o, m_row, l_row, v_p_1 = softmax_helper.lazy_rescale_o(v_o, m_row, l_row, m_tile_max_b, v_p_1) - else: - v_o, m_row, l_row, v_p_1 = softmax_helper.rescale_o(v_o, m_row, l_row, m_tile_max_b, v_p_1) - v_o = gemm_helper.pv_step_k(1, v_p_1, v_v, v_o) - if const_expr(traits.D_CHUNKS > 4): - v_s_0 = softmax_helper.sub_m(v_s_0, m_row) - v_p_0 = softmax_helper.exp2(v_s_0, 0, 16) - fx.rocdl.s_waitcnt(lgkmcnt=0) - v_o = gemm_helper.pv_step_k(2, v_p_1, v_v, v_o) - v_o = gemm_helper.pv_step_k(3, v_p_1, v_v, v_o) - else: - v_o = gemm_helper.pv_step_k(2, v_p_1, v_v, v_o) - v_o = gemm_helper.pv_step_k(3, v_p_1, v_v, v_o) - v_s_0 = softmax_helper.sub_m(v_s_0, m_row) - v_p_0 = softmax_helper.exp2(v_s_0, 0, 16) - _sched_barrier_pairs(traits, 6, 5, 4) - _sched_barrier_exp_pairs(traits, 6, 3, 4) - if const_expr(traits.DUALWAVE_SWP_SETPRIO): - rocdl.s_setprio(0) - rocdl.sched_barrier(0) - rocdl.s_barrier() - rocdl.sched_barrier(0) - - yield_args = [m_row, l_row] + v_o + [ctx.v_pair_to_vec32(v_p_0)] - loop_results = yield yield_args - - # Drain the final three tiles without further prefetch-ahead. - m_row = loop_results[0] - l_row = loop_results[1] - v_o = [loop_results[2 + i] for i in range_constexpr(traits.D_CHUNKS)] - v_p_0 = ctx.v_vec32_to_pair(loop_results[2 + traits.D_CHUNKS]) - - max_m3 = ctx.split_t_end - 3 - max_m2 = ctx.split_t_end - 2 - max_m1 = ctx.split_t_end - 1 - - kv_gmem_to_lds.load_v(max_m3 * traits.BLOCK_N, 1) - v_k = kv_lds_to_regs.load_k(1) - fx.rocdl.s_waitcnt(lgkmcnt=0) - _waitcnt_vm_n(ctx.NUM_DMA_K + ctx.NUM_DMA_V) - rocdl.sched_barrier(0) - rocdl.s_barrier() - rocdl.sched_barrier(0) - - v_s_1 = gemm_helper.qk(v_k, q_all_wide) - v_p_0 = softmax_helper.exp2(v_p_0, 16, 16) - l_row = softmax_helper.reduce_sum(l_row, v_p_0) - v_p_0 = softmax_helper.cast_p(v_p_0) - v_p_0 = softmax_helper.anchor_v_p(v_p_0) - _sched_barrier_exp_pairs(traits, 6, 3, 5) - _sched_barrier_pairs(traits, 10, 5, 5) - rocdl.sched_barrier(0) - rocdl.s_barrier() - rocdl.sched_barrier(0) - - kv_gmem_to_lds.load_k(max_m1 * traits.BLOCK_N, 1) - v_packs_e3 = kv_lds_to_regs.load_v(0) - v_s_1 = softmax_helper.causal_mask_prologue_if_needed( - v_s_1, - max_m3, - max_m2 * traits.BLOCK_N, - ) - fx.rocdl.s_waitcnt(lgkmcnt=0) - _waitcnt_vm_n(ctx.NUM_DMA_K + ctx.NUM_DMA_V) - rocdl.sched_barrier(0) - rocdl.s_barrier() - rocdl.sched_barrier(0) - - if const_expr(traits.DUALWAVE_SWP_SETPRIO): - rocdl.s_setprio(1) - if const_expr(traits.D_CHUNKS > 4): - v_o = gemm_helper.pv_step_k(0, v_p_0, v_packs_e3, v_o) - v_o = gemm_helper.pv_step_k(1, v_p_0, v_packs_e3, v_o) - fx.rocdl.s_waitcnt(lgkmcnt=0) - v_o = gemm_helper.pv_step_k(2, v_p_0, v_packs_e3, v_o) - v_o = gemm_helper.pv_step_k(3, v_p_0, v_packs_e3, v_o) - else: - v_o = gemm_helper.pv(v_p_0, v_packs_e3, v_o) - m_tile_max_e3 = softmax_helper.reduce_max(v_s_1) - row_max_e3, rescale_e3 = softmax_helper.rescale_from_tile_max(m_row, m_tile_max_e3) - m_row = row_max_e3 - v_s_1 = softmax_helper.sub_m(v_s_1, row_max_e3) - v_p_1 = softmax_helper.exp2(v_s_1, 0, 16) - _sched_barrier_pairs(traits, 10, 5, 6) - _sched_barrier_exp_pairs(traits, 6, 3, 6) - rocdl.sched_barrier(0) - softmax_helper.scale_o(v_o, rescale_e3) - v_o = softmax_helper.anchor_v_o(v_o) - - if const_expr(traits.DUALWAVE_SWP_SETPRIO): - rocdl.s_setprio(0) - rocdl.sched_barrier(0) - rocdl.s_barrier() - rocdl.sched_barrier(0) - - kv_gmem_to_lds.load_v(max_m2 * traits.BLOCK_N, 0) - v_k = kv_lds_to_regs.load_k(0) - fx.rocdl.s_waitcnt(lgkmcnt=0) - _waitcnt_vm_n(ctx.NUM_DMA_K + ctx.NUM_DMA_V) - rocdl.sched_barrier(0) - rocdl.s_barrier() - rocdl.sched_barrier(0) - - v_s_0 = gemm_helper.qk(v_k, q_all_wide) - l_row = softmax_helper.apply_l_rescale(l_row, rescale_e3) - v_p_1 = softmax_helper.exp2(v_p_1, 16, 16) - l_row = softmax_helper.reduce_sum(l_row, v_p_1) - v_p_1 = softmax_helper.cast_p(v_p_1) - v_p_1 = softmax_helper.anchor_v_p(v_p_1) - _sched_barrier_exp_pairs(traits, 6, 3, 7) - _sched_barrier_pairs(traits, 10, 5, 7) - rocdl.sched_barrier(0) - rocdl.s_barrier() - rocdl.sched_barrier(0) - - v_packs_e7 = kv_lds_to_regs.load_v(1) - v_s_0 = softmax_helper.causal_mask_prologue_if_needed( - v_s_0, - max_m2, - max_m1 * traits.BLOCK_N, - ) - fx.rocdl.s_waitcnt(lgkmcnt=0) - _waitcnt_vm_n(ctx.NUM_DMA_V) - rocdl.sched_barrier(0) - rocdl.s_barrier() - rocdl.sched_barrier(0) - - if const_expr(traits.DUALWAVE_SWP_SETPRIO): - rocdl.s_setprio(1) - if const_expr(traits.D_CHUNKS > 4): - v_o = gemm_helper.pv_step_k(0, v_p_1, v_packs_e7, v_o) - v_o = gemm_helper.pv_step_k(1, v_p_1, v_packs_e7, v_o) - fx.rocdl.s_waitcnt(lgkmcnt=0) - v_o = gemm_helper.pv_step_k(2, v_p_1, v_packs_e7, v_o) - v_o = gemm_helper.pv_step_k(3, v_p_1, v_packs_e7, v_o) - else: - v_o = gemm_helper.pv(v_p_1, v_packs_e7, v_o) - m_tile_max_e7 = softmax_helper.reduce_max(v_s_0) - row_max_e7, rescale_e7 = softmax_helper.rescale_from_tile_max(m_row, m_tile_max_e7) - m_row = row_max_e7 - v_s_0 = softmax_helper.sub_m(v_s_0, row_max_e7) - v_p_0 = softmax_helper.exp2(v_s_0, 0, 16) - _sched_barrier_pairs(traits, 10, 5, 8) - _sched_barrier_exp_pairs(traits, 6, 3, 8) - rocdl.sched_barrier(0) - softmax_helper.scale_o(v_o, rescale_e7) - v_o = softmax_helper.anchor_v_o(v_o) - if const_expr(traits.DUALWAVE_SWP_SETPRIO): - rocdl.s_setprio(0) - rocdl.sched_barrier(0) - rocdl.s_barrier() - rocdl.sched_barrier(0) - - kv_gmem_to_lds.load_v(max_m1 * traits.BLOCK_N, 1) - v_k = kv_lds_to_regs.load_k(1) - fx.rocdl.s_waitcnt(lgkmcnt=0) - _waitcnt_vm_n(ctx.NUM_DMA_V) - rocdl.sched_barrier(0) - rocdl.s_barrier() - rocdl.sched_barrier(0) - - v_s_1 = gemm_helper.qk(v_k, q_all_wide) - l_row = softmax_helper.apply_l_rescale(l_row, rescale_e7) - v_p_0 = softmax_helper.exp2(v_p_0, 16, 16) - l_row = softmax_helper.reduce_sum(l_row, v_p_0) - v_p_0 = softmax_helper.cast_p(v_p_0) - v_p_0 = softmax_helper.anchor_v_p(v_p_0) - _sched_barrier_exp_pairs(traits, 6, 3, 9) - _sched_barrier_pairs(traits, 10, 5, 9) - rocdl.sched_barrier(0) - rocdl.s_barrier() - rocdl.sched_barrier(0) - - v_packs_e11 = kv_lds_to_regs.load_v(0) - v_s_1 = softmax_helper.causal_mask_prologue_if_needed( - v_s_1, - max_m1, - ctx.split_t_end * traits.BLOCK_N, - ) - fx.rocdl.s_waitcnt(lgkmcnt=0) - _waitcnt_vm_n(0) - rocdl.sched_barrier(0) - rocdl.s_barrier() - rocdl.sched_barrier(0) - - if const_expr(traits.D_CHUNKS > 4): - v_o = gemm_helper.pv_step_k(0, v_p_0, v_packs_e11, v_o) - v_o = gemm_helper.pv_step_k(1, v_p_0, v_packs_e11, v_o) - fx.rocdl.s_waitcnt(lgkmcnt=0) - v_o = gemm_helper.pv_step_k(2, v_p_0, v_packs_e11, v_o) - v_o = gemm_helper.pv_step_k(3, v_p_0, v_packs_e11, v_o) - else: - v_o = gemm_helper.pv(v_p_0, v_packs_e11, v_o) - m_tile_max_e11 = softmax_helper.reduce_max(v_s_1) - row_max_e11, rescale_e11 = softmax_helper.rescale_from_tile_max(m_row, m_tile_max_e11) - m_row = row_max_e11 - v_s_1 = softmax_helper.sub_m(v_s_1, row_max_e11) - v_p_1 = softmax_helper.exp2(v_s_1, 0, 16) - _sched_barrier_pairs(traits, 9, 6, 10) - _sched_barrier_exp_pairs(traits, 7, 3, 10) - rocdl.sched_barrier(0) - v_p_1 = softmax_helper.exp2(v_p_1, 16, 16) - l_row = softmax_helper.apply_l_rescale(l_row, rescale_e11) - l_row = softmax_helper.reduce_sum(l_row, v_p_1) - v_p_1 = softmax_helper.cast_p(v_p_1) - v_p_1 = softmax_helper.anchor_v_p(v_p_1) - rocdl.sched_barrier(0) - softmax_helper.scale_o(v_o, rescale_e11) - v_o = softmax_helper.anchor_v_o(v_o) - rocdl.s_barrier() - rocdl.sched_barrier(0) - - v_packs_e13 = kv_lds_to_regs.load_v(1) - fx.rocdl.s_waitcnt(lgkmcnt=0) - rocdl.sched_barrier(0) - rocdl.s_barrier() - rocdl.sched_barrier(0) - - if const_expr(traits.D_CHUNKS > 4): - v_o = gemm_helper.pv_step_k(0, v_p_1, v_packs_e13, v_o) - v_o = gemm_helper.pv_step_k(1, v_p_1, v_packs_e13, v_o) - fx.rocdl.s_waitcnt(lgkmcnt=0) - v_o = gemm_helper.pv_step_k(2, v_p_1, v_packs_e13, v_o) - v_o = gemm_helper.pv_step_k(3, v_p_1, v_packs_e13, v_o) - else: - v_o = gemm_helper.pv(v_p_1, v_packs_e13, v_o) - - # Apply V descale once after normalizing the native FP8 P*V result. - inv_l = softmax_helper.safe_l_inv(l_row) - inv_l = inv_l * ctx.vd_fp8 - softmax_helper.scale_o(v_o, inv_l) - - # Group A's extra barrier closes the prologue's phase shift before stores. - if const_expr(traits.DUALWAVE_SWP_ENABLE_STAGGER): - _stagger_extra_barrier_if_zero(ctx.stagger_i32) - else: - rocdl.s_barrier() - - # 128b stores fuse this lane and its half-wave partner, so each pair - # covers 8 contiguous columns instead of two 64b stores. - output_store.store_final_o_if_valid(v_o, q_row) - - if ctx.q_start < ctx.seqlen_q_v: - _run_q_block() - @flyc.jit def launch_flash_attn_dualwave_swp( Q: fx.Tensor, @@ -930,65 +457,34 @@ def launch_flash_attn_dualwave_swp( "rocdl.flat_work_group_size": f"{BLOCK_SIZE},{BLOCK_SIZE}", "passthrough": passthrough_entries, } - if const_expr(PAGED_BN128): - flash_attn_paged_fp8_bn128_kernel( - Q, - K, - V, - O, - DebugCounts, - CuSeqQ, - CuSeqKv, - BlockTable, - block_table_stride, - QDescale, - KDescale, - VDescale, - seq_len, - seq_len_kv, - stride_q_n, - stride_kv_n, - head_dim_runtime, - value_attrs=kernel_attrs, - ).launch( - grid=( - NUM_HEADS_Q * BATCH_INTERLEAVE_GROUP, - num_q_blocks, - grid_z // BATCH_INTERLEAVE_GROUP, - ), - block=(BLOCK_SIZE, 1, 1), - stream=stream, - ) - else: - flash_attn_dualwave_swp_fp8_gfx950_kernel( - Q, - K, - V, - O, - DebugCounts, - CuSeqQ, - CuSeqKv, - BlockTable, - block_table_stride, - QDescale, - KDescale, - VDescale, - seq_len, - seq_len_kv, - stride_q_n, - stride_o_n, - stride_kv_n, - head_dim_runtime, - value_attrs=kernel_attrs, - ).launch( - grid=( - NUM_HEADS_Q * BATCH_INTERLEAVE_GROUP, - num_q_blocks, - grid_z // BATCH_INTERLEAVE_GROUP, - ), - block=(BLOCK_SIZE, 1, 1), - stream=stream, - ) + flash_attn_paged_fp8_bn128_kernel( + Q, + K, + V, + O, + DebugCounts, + CuSeqQ, + CuSeqKv, + BlockTable, + block_table_stride, + QDescale, + KDescale, + VDescale, + seq_len, + seq_len_kv, + stride_q_n, + stride_o_n, + head_dim_runtime, + value_attrs=kernel_attrs, + ).launch( + grid=( + NUM_HEADS_Q * BATCH_INTERLEAVE_GROUP, + num_q_blocks, + grid_z // BATCH_INTERLEAVE_GROUP, + ), + block=(BLOCK_SIZE, 1, 1), + stream=stream, + ) _dualwave_swp_compile_hints = { "fast_fp_math": True, @@ -1002,13 +498,13 @@ def launch_flash_attn_dualwave_swp( launch_flash_attn_dualwave_swp.compile_hints = dict(_dualwave_swp_compile_hints) def _validate_paged_bn128_launch(batch_size, seq_len_kv, block_table_stride): - if not PAGED_BN128: + if not PAIRED_PAGE_IDS: return batch_size = int(batch_size) seq_len_kv = int(seq_len_kv) block_table_stride = int(block_table_stride) num_kv_pages = (seq_len_kv + traits.PAGE_SIZE - 1) // traits.PAGE_SIZE - if (not PAGED_BN128_VARLEN and batch_size != 1) or num_kv_pages < 2 or num_kv_pages % 2 != 0: + if (not USE_CU_SEQLENS and batch_size != 1) or num_kv_pages < 2 or num_kv_pages % 2 != 0: raise ValueError( "paged BN128 requires batch_size=1 unless compiled for packed varlen, " "and a positive even number " diff --git a/kernels/attention/flash_attn_interface.py b/kernels/attention/flash_attn_interface.py index 84c2a440c4..0d7830a240 100644 --- a/kernels/attention/flash_attn_interface.py +++ b/kernels/attention/flash_attn_interface.py @@ -475,8 +475,6 @@ def _build_paged_fp8( waves_per_eu: int, daz: bool, lazy_rescale: bool, - setprio: bool, - enable_stagger: bool, use_bn128: bool, paged_bn128_varlen: bool, batch_interleave_group: int, @@ -494,8 +492,6 @@ def _build_paged_fp8( waves_per_eu=waves_per_eu, daz=daz, dualwave_swp_lazy_rescale=lazy_rescale, - dualwave_swp_setprio=setprio, - dualwave_swp_enable_stagger=enable_stagger, num_kv_splits=1, varlen=True, cross_seqlen=True, @@ -778,10 +774,6 @@ def _flydsl_flash_attn_paged( if paged_fp8: use_bn128 = max_kv_pages % 2 == 0 paged_bn128_varlen = use_bn128 and B > 1 - paged_setprio = dualwave_swp_setprio and D != 192 - # The BF16 phase shift exposes H2 staging and probability-pack - # latency because segmented FP8 has only six wide P*V MFMAs. - paged_stagger = dualwave_swp_enable_stagger and value_head_dim == 128 exe = _build_paged_fp8( num_heads=H, num_kv_heads=num_kv_heads, @@ -790,8 +782,6 @@ def _flydsl_flash_attn_paged( waves_per_eu=waves_per_eu, daz=daz, lazy_rescale=dualwave_swp_lazy_rescale, - setprio=paged_setprio, - enable_stagger=paged_stagger, use_bn128=use_bn128, paged_bn128_varlen=paged_bn128_varlen, batch_interleave_group=( diff --git a/kernels/attention/flash_attn_utils.py b/kernels/attention/flash_attn_utils.py index 6e8b26941f..ba6271a78f 100644 --- a/kernels/attention/flash_attn_utils.py +++ b/kernels/attention/flash_attn_utils.py @@ -394,7 +394,7 @@ def _sub_score_pair(v_s, row_max, fm_fast): return Vec.from_elements(lo_sub, fx.Float32).ir_value(), Vec.from_elements(hi_sub, fx.Float32).ir_value() -def _scale_sub_score_pair(v_s, row_max_raw, scale, zero_f, fm_fast, bias=None, scalar_fma=False): +def _scale_sub_score_pair(v_s, row_max_raw, scale, zero_f, fm_fast, bias=None, constrained_fma=False): """Fused softmax-scale + row-max subtraction (optimization 1-A). Returns ``scale * (v_s - row_max_raw) + bias`` per element via a single FMA @@ -412,20 +412,21 @@ def _scale_sub_score_pair(v_s, row_max_raw, scale, zero_f, fm_fast, bias=None, s # exp2 lands on 2**bias * P instead of P, at no extra instruction: the # FMA's addend absorbs it. neg_scaled_max = neg_scaled_max + bias - if const_expr(scalar_fma): - # Preserve scalar instruction selection for MFMA/VALU co-issue. + if const_expr(constrained_fma): + # fx.fma has no constrained-rounding option. Keep this standard LLVM + # intrinsic boundary local; current codegen retains unpacked FP32 FMA. def _center(values): centered = [] for i in range_constexpr(16): centered.append( fx.Float32( - llvm.inline_asm( - T.f32, - [Vec(values)[i].ir_value(), scale.ir_value(), neg_scaled_max.ir_value()], - "v_fma_f32 $0, $1, $2, $3", - "=v,v,v,v", - has_side_effects=False, + llvm.intr_experimental_constrained_fma( + Vec(values)[i].ir_value(), + scale.ir_value(), + neg_scaled_max.ir_value(), + ir.IntegerAttr.get(T.i64, int(llvm.RoundingMode.Dynamic)), + llvm.FPExceptionBehavior.Ignore, ) ) ) @@ -1026,7 +1027,7 @@ def _init_dualwave_thread_mapping(ctx): # Keep a bounded batch group in grid X so the final causal q-block # occupies more CUs without abandoning K/V locality across all batches. linear_head_batch = fx.Index(gpu.block_idx.x) - if const_expr(getattr(traits, "PAGED", False) and getattr(traits, "BN128", False)): + if const_expr(getattr(traits, "PAGED", False) and getattr(traits, "PAIRED_PAGE_IDS", False)): ctx.h_idx = linear_head_batch // batch_interleave_group ctx.batch_idx = ( fx.Int64(gpu.block_idx.z) * batch_interleave_group + linear_head_batch % batch_interleave_group @@ -2093,14 +2094,11 @@ class PagedDualwaveSwpFp8Traits: WAVES_PER_EU: int DAZ: bool DUALWAVE_SWP_LAZY_RESCALE: bool - DUALWAVE_SWP_SETPRIO: bool DUALWAVE_SWP_DEBUG_LAZY_COUNTS: bool - DUALWAVE_SWP_ENABLE_STAGGER: bool VARLEN: bool FP8_PV_SEGMENTED: bool FP8_V_H1: int FP8_V_H2: int - BN128: bool DEFAULT_STRIDE_Q_N: int DEFAULT_STRIDE_KV_N: int SMEM_D_RPT: int @@ -2109,10 +2107,10 @@ class PagedDualwaveSwpFp8Traits: DUALWAVE_SWP_KV_PER_BUFFER: int LDS_KV_TOTAL_SIZE: int DUALWAVE_SWP_K_BUF_BASE: tuple[int, ...] - DUALWAVE_SWP_V_BUF_BASE: tuple[int, ...] VT_BF16_TOTAL: int FP8_V_ROW_STRIDE: int DUALWAVE_SWP_RESCALE_THRESHOLD: float + PAIRED_PAGE_IDS: bool = True BATCH_INTERLEAVE_GROUP: int = 1 BLOCK_M: int = 256 BLOCK_N: int = 64 @@ -2137,8 +2135,6 @@ class PagedDualwaveSwpFp8Traits: KV_VEC_SIZE: int = 16 LANE_SPLIT_KV: int = 8 SCHED_MFMA_MASK: int = 0x008 - SCHED_VALU_MASK: int = 0x002 - SCHED_EXP_MASK: int = 0x400 SCHED_DS_READ_MASK: int = 0x100 NEG_INF_F32_BITS: int = 0xFF800000 XCD_SWIZZLE: bool = False @@ -2156,14 +2152,12 @@ def cache_tag(self): self.DAZ, self.DUALWAVE_SWP_LAZY_RESCALE, self.DUALWAVE_SWP_RESCALE_THRESHOLD, - self.DUALWAVE_SWP_SETPRIO, self.DUALWAVE_SWP_DEBUG_LAZY_COUNTS, - self.DUALWAVE_SWP_ENABLE_STAGGER, self.VARLEN, - self.BN128, self.NUM_PREFETCH_K, self.FP8_V_ROW_STRIDE, self.BATCH_INTERLEAVE_GROUP, + self.PAIRED_PAGE_IDS, ) @@ -2187,25 +2181,22 @@ def _make_paged_dualwave_swp_fp8_traits( waves_per_eu=2, daz=True, dualwave_swp_lazy_rescale=True, - dualwave_swp_setprio=True, dualwave_swp_debug_lazy_counts=False, - dualwave_swp_enable_stagger=True, varlen=False, - bn128=False, batch_interleave_group=1, + paired_page_ids=True, ): """Build layouts after the dedicated builder validates the paged contract.""" block_n = 64 smem_d_rpt = (head_dim + 127) // 128 - # Preserve the generic kernel's D256-shaped K slot and spacing. + # Retain padded K-slot spacing; V192 below compacts the unused D256 tail. smem_k_tile_elems = 8 * smem_d_rpt * 1040 - smem_v_tile_elems = 8 * smem_d_rpt * 1088 - compact_v192 = bn128 and value_head_dim == 192 + compact_v192 = value_head_dim == 192 if compact_v192: # Six K/V slots fit in 160 KiB only when K omits its unused D256 tail. smem_k_tile_elems = block_n * head_dim - num_prefetch_k = (8 if head_dim == 128 else 6) if bn128 else 2 - slot_elems = smem_k_tile_elems if bn128 else smem_k_tile_elems + smem_v_tile_elems + num_prefetch_k = 8 if head_dim == 128 else 6 + slot_elems = smem_k_tile_elems fp8_v_h1, fp8_v_h2 = _factor_fp8_pv_head_dim(value_head_dim, 128, 32) fp8_v_row_stride = block_n if compact_v192 else block_n + 16 fp8_v_tile_bytes = value_head_dim * fp8_v_row_stride @@ -2220,15 +2211,13 @@ def _make_paged_dualwave_swp_fp8_traits( WAVES_PER_EU=waves_per_eu, DAZ=bool(daz), DUALWAVE_SWP_LAZY_RESCALE=bool(dualwave_swp_lazy_rescale), - DUALWAVE_SWP_SETPRIO=bool(dualwave_swp_setprio), DUALWAVE_SWP_DEBUG_LAZY_COUNTS=bool(dualwave_swp_debug_lazy_counts), - DUALWAVE_SWP_ENABLE_STAGGER=bool(dualwave_swp_enable_stagger), DUALWAVE_SWP_RESCALE_THRESHOLD=rescale_threshold, VARLEN=bool(varlen), FP8_PV_SEGMENTED=fp8_v_h2 > 0, FP8_V_H1=fp8_v_h1, FP8_V_H2=fp8_v_h2, - BN128=bool(bn128), + PAIRED_PAGE_IDS=bool(paired_page_ids), DEFAULT_STRIDE_Q_N=num_heads * head_dim, DEFAULT_STRIDE_KV_N=num_kv_heads * head_dim, SMEM_D_RPT=smem_d_rpt, @@ -2237,7 +2226,6 @@ def _make_paged_dualwave_swp_fp8_traits( DUALWAVE_SWP_KV_PER_BUFFER=slot_elems, LDS_KV_TOTAL_SIZE=num_prefetch_k * slot_elems, DUALWAVE_SWP_K_BUF_BASE=tuple(i * slot_elems for i in range(num_prefetch_k)), - DUALWAVE_SWP_V_BUF_BASE=tuple(smem_k_tile_elems + i * slot_elems for i in range(num_prefetch_k)), VT_BF16_TOTAL=num_prefetch_k * (fp8_v_tile_bytes // 2) + 128, FP8_V_ROW_STRIDE=fp8_v_row_stride, BATCH_INTERLEAVE_GROUP=int(batch_interleave_group), @@ -4602,7 +4590,6 @@ def init_types_and_constants(self): self.v4bf16_type = Vec.make_type(4, fx.BFloat16) if const_expr(traits.PAGED): self.NUM_DMA_K = traits.SMEM_D_RPT - self.NUM_DMA_V = traits.SMEM_D_RPT else: self.NUM_DMA_K = len(traits.K_BAND_CHUNK) self.c_neg_inf = fx.Float32(float("-inf")) @@ -4876,7 +4863,6 @@ def v_vec32_to_pair(self, v): return _v_vec32_to_pair(v) def bf16_trunc_pack_v8(self, f32_vals): - # Generic paged softmax carries P in BF16 between correction steps. # The dtype-dispatched pack would incorrectly select FP16 for FP8 input. pairs = [] for j in range_constexpr(4): @@ -4922,6 +4908,8 @@ def load_page_id(self, tile_start): def load_page_id_pair(self, tile_start): """Load two adjacent physical page IDs for a BN128 tile pair.""" + if const_expr(not self.traits.PAIRED_PAGE_IDS): + return self.load_page_id(tile_start), self.load_page_id(tile_start + self.traits.PAGE_SIZE) local_page = fx.Int32(fx.Int64(tile_start) // self.traits.PAGE_SIZE) page_base = fx.Int32(self.page_base) num_kv_tiles = fx.Int32(self.num_kv_tiles) @@ -4955,6 +4943,12 @@ def read_i32x4_lds(self, byte_row): tile = fx.slice(self.k_lds_i32_tiles, (None, fx.Uint32(byte_row) // 16)) return fx.generic_load(fx.get_iter(tile), dtype=fx.Int32, count=4).ir_value() + def preserve_accumulators(self, v_o): + if const_expr(self.traits.PAGED): + # Preserve FP expression boundaries without inline-assembly pins. + return [llvm.intr_arithmetic_fence(fx.as_ir_value(acc)) for acc in v_o] + return _anchor_v_o(self.traits, v_o) + class DualwaveFp8QLoader(DualwaveFp8KernelContext): def __init__(self, ctx): @@ -4985,17 +4979,8 @@ def load_all_wide(self, q_row_in_block): class DualwaveFp8GemmHelper(DualwaveFp8KernelContext): def __init__(self, ctx): super().__init__(ctx) - - def _v8bf16_to_f32(self, v8): - f32 = Vec(v8).to(fx.Float32) - return [f32[i] for i in range_constexpr(8)] - - def _p_to_fp8_i32x8(self, v_p): - p_lo, p_hi = v_p - f32 = [] - for pk in (p_lo[0], p_lo[1], p_hi[0], p_hi[1]): - f32 += self._v8bf16_to_f32(pk) - return self._pack_p_fp8(f32) + if const_expr(self.traits.PAGED): + self.fp8_mma = fx.make_mma_atom(fx.rocdl.cdna4.MFMA_Scale(32, 32, 64, fx.Float8E4M3FN)) def _pack_p_fp8(self, f32): # P words must follow vectorized K's four-token groups for paged P*V. @@ -5008,24 +4993,16 @@ def _pack_p_fp8(self, f32): ).ir_value() return packed - def _pv_fp8(self, v_p, v_v, v_o): - p_fp8 = self._p_to_fp8_i32x8(v_p) - for dc in range_constexpr(self.traits.D_CHUNKS): - v_op = self._v_concat_i32x8(v_v, dc) - v_o[dc] = self._mfma_acc_fp8_wide(v_op, p_fp8, v_o[dc]) - return v_o - - def _pv_step_fp8(self, step, v_p, v_v, v_o): - # The paged pipeline performs its online-softmax correction between - # steps 0 and 1. Repack P after that correction before the remaining - # output chunks consume it. - if const_expr(step == 0 or (self.traits.PAGED and step == 1)): - self._pv_p_fp8_cache = self._p_to_fp8_i32x8(v_p) - v_op = self._v_concat_i32x8(v_v, step) - v_o[step] = self._mfma_acc_fp8_wide(v_op, self._pv_p_fp8_cache, v_o[step]) - return v_o - def _mfma_acc_fp8_wide(self, a_i32x8, b_i32x8, c_v16): + if const_expr(self.traits.PAGED): + a = fx.make_rmem_tensor(8, fx.Int32) + b = fx.make_rmem_tensor(8, fx.Int32) + c = fx.make_rmem_tensor(16, fx.Float32) + a.store(Vec(a_i32x8)) + b.store(Vec(b_i32x8)) + c.store(Vec(c_v16)) + fx.gemm(self.fp8_mma, c, a, b, c, scale_a=fx.Int32(0x7F7F7F7F), scale_b=fx.Int32(0x7F7F7F7F)) + return c.load().ir_value() # Wide fp8 QK: mfma_scale (32x32x64) with unit E8M0 scales, i32x8 operands. return rocdl.mfma_scale_f32_32x32x64_f8f6f4( self.v16f32_type, @@ -5060,38 +5037,6 @@ def _v_concat_i32x8(self, v_v, dc): words.append(fx.Int32(v2[1])) return Vec.from_elements(words, fx.Int32).ir_value() - def _load_v_fp8_segmented_chunk(self, buf_id, dc): - traits = self.traits - row_stride = traits.FP8_V_ROW_STRIDE - tile_bytes = traits.HEAD_DIM_V * row_stride - aligned_base = ((self.lds_vt_base_idx + fx.Int64(127)) // fx.Int64(128)) * fx.Int64(128) - tile_index = fx.Uint32(aligned_base - self.lds_vt_base_idx) // 16 - tile_index = tile_index + fx.Uint32(self.lane_mod_32) * (row_stride // 16) + fx.Uint32(self.lane_div_32) * 2 - tile_index = tile_index + fx.Uint32(buf_id) * (tile_bytes // 16) - chunk_offset = dc * traits.D_CHUNK * row_stride // 16 - halves = [] - for half in range_constexpr(2): - view = fx.slice(self.v_lds_i32_tiles, (None, tile_index + chunk_offset + half)) - halves.append(fx.generic_load(fx.get_iter(view), dtype=fx.Int32, count=4)) - fx.rocdl.s_waitcnt(lgkmcnt=0) - return _concat_vectors(halves[0], halves[1]).bitcast(fx.Int32).ir_value() - - def _pv_step_fp8_segmented(self, step, v_p, buf_id, v_o): - traits = self.traits - if const_expr(step == 0 or (traits.PAGED and step == 1)): - self._pv_segmented_p_cache = self._p_to_fp8_i32x8(v_p) - - chunks_per_step = traits.D_CHUNKS // 4 - extra_chunks = traits.D_CHUNKS % 4 - step_extra = 1 if step < extra_chunks else 0 - dc_start = step * chunks_per_step + min(step, extra_chunks) - dc_end = dc_start + chunks_per_step + step_extra - for dc in range_constexpr(dc_start, dc_end): - v_op = self._load_v_fp8_segmented_chunk(buf_id, dc) - v_o[dc] = self._mfma_acc_fp8_wide(v_op, self._pv_segmented_p_cache, v_o[dc]) - rocdl.sched_barrier(0) - return v_o - def _load_q_wide_lds(self): traits = self.traits q_row_in_block = self.ctx_ref.q_row_in_block @@ -5129,20 +5074,14 @@ def qk(self, v_k, q_wide=None): q_w = q_all_wide[ws] v_s_lo = self._mfma_acc_fp8_wide(k_lo[ws], q_w, v_s_lo) v_s_hi = self._mfma_acc_fp8_wide(k_hi[ws], q_w, v_s_hi) - if const_expr(not traits.PAGED or traits.BN128): - n_ds = const_expr(traits.HEAD_DIM // 64 * 4) - n_mfma = const_expr(traits.HEAD_DIM // 64 * 2) - rocdl.sched_group_barrier(traits.SCHED_DS_READ_MASK, n_ds // 2, 12) - rocdl.sched_group_barrier(traits.SCHED_MFMA_MASK, 1, 12) - rocdl.sched_group_barrier(traits.SCHED_DS_READ_MASK, n_ds // 2, 12) - rocdl.sched_group_barrier(traits.SCHED_MFMA_MASK, n_mfma - 1, 12) + n_ds = const_expr(traits.HEAD_DIM // 64 * 4) + n_mfma = const_expr(traits.HEAD_DIM // 64 * 2) + rocdl.sched_group_barrier(traits.SCHED_DS_READ_MASK, n_ds // 2, 12) + rocdl.sched_group_barrier(traits.SCHED_MFMA_MASK, 1, 12) + rocdl.sched_group_barrier(traits.SCHED_DS_READ_MASK, n_ds // 2, 12) + rocdl.sched_group_barrier(traits.SCHED_MFMA_MASK, n_mfma - 1, 12) return (v_s_lo, v_s_hi) - def pv_step_k(self, step, v_p, v_v, v_o): - if const_expr(self.traits.FP8_PV_SEGMENTED): - return self._pv_step_fp8_segmented(step, v_p, v_v, v_o) - return self._pv_step_fp8(step, v_p, v_v, v_o) - def cast_p_fp8_direct(self, v_p): lo_partial_list, hi_full = v_p f32 = [] @@ -5155,16 +5094,14 @@ def cast_p_fp8_direct(self, v_p): return self._pack_p_fp8(f32) def _pv_fp8_direct(self, p_fp8, v_v, v_o): - v_o = _anchor_v_o(self.traits, v_o) + v_o = self.preserve_accumulators(v_o) for dc in range_constexpr(self.traits.D_CHUNKS): v_op = self._v_concat_i32x8(v_v, dc) v_o[dc] = self._mfma_acc_fp8_wide(v_op, p_fp8, v_o[dc]) return v_o def pv(self, v_p, v_v, v_o): - if const_expr(not self.traits.PAGED or self.traits.BN128): - return self._pv_fp8_direct(v_p, v_v, v_o) - return self._pv_fp8(v_p, v_v, v_o) + return self._pv_fp8_direct(v_p, v_v, v_o) class DualwaveFp8KvGmemToLdsLoader(DualwaveFp8KernelContext): @@ -5195,7 +5132,7 @@ def _load_vectorized(): lds_addr = rocdl.readfirstlane(T.i32, fx.Int32(lds_addr).ir_value()) self.buffer_load_lds_128(src_div, lds_addr, src_elem, 0) - if const_expr(traits.BN128 and traits.HEAD_DIM_V == 192): + if const_expr(traits.HEAD_DIM_V == 192): @flyc.jit def _load_compact(): @@ -5244,10 +5181,7 @@ def _run(): def load_v(self, tile_start, buf_id, page_id=None): if const_expr(self.traits.PAGED): - if const_expr(self.traits.FP8_PV_SEGMENTED and not self.traits.BN128): - self._stage_v_fp8_vectorized_segmented(tile_start, buf_id, page_id=page_id) - else: - self._stage_v_fp8_vectorized_bankpad(tile_start, buf_id, page_id=page_id) + self._stage_v_fp8_vectorized_bankpad(tile_start, buf_id, page_id=page_id) else: self._stage_v_fp8_block_dma(tile_start, buf_id) @@ -5295,20 +5229,6 @@ def _run(): _run() - def _stage_v_fp8_vectorized_segmented(self, tile_start, buf_id, page_id=None): - if page_id is None: - page_id = self.load_page_id(tile_start) - prefix = self._load_v_fp8_vectorized_bankpad_source(tile_start, page_id=page_id) - self._store_v_fp8_vectorized_bankpad(prefix, buf_id) - if const_expr(self.traits.FP8_V_H2 > 0): - self._stage_v_fp8_vectorized_segment( - tile_start, - buf_id, - self.traits.FP8_V_H1, - self.traits.FP8_V_H2, - page_id, - ) - def _permute_v_fp8_vectorized(self, src_i32x4): src_words = Vec(src_i32x4, (4,), fx.Int32) # permlane16_swap returns an LLVM pair; unpack it at this intrinsic boundary. @@ -5337,39 +5257,6 @@ def _permute_v_fp8_vectorized(self, src_i32x4): fx.Int32, ) - def _stage_v_fp8_vectorized_segment(self, tile_start, buf_id, d_offset, segment_dim, page_id): - # Sharing the prefix store's address expression increases tail VGPR spills. - traits = self.traits - src_div = self.make_page_view(self.v_base_iter, page_id, is_value=True) - token_groups16 = traits.BLOCK_N // traits.KV_VEC_SIZE - flat = self.tid - - @flyc.jit - def _stage_active(): - if flat < fx.Int64(segment_dim * token_groups16): - n_group16 = self.lane_in_warp // fx.Int64(16) - d_local = self.wave_id * fx.Int64(16) + self.lane_in_warp % fx.Int64(16) - d_col = fx.Int64(d_offset) + d_local - src_elem = ( - self.kv_head_idx * token_groups16 * traits.HEAD_DIM_V * traits.KV_VEC_SIZE - + n_group16 * traits.HEAD_DIM_V * traits.KV_VEC_SIZE - + d_col * traits.KV_VEC_SIZE - ) - reordered = self._permute_v_fp8_vectorized(self.buffer_load_fp8x16(src_div, src_elem)) - dest_group16 = (n_group16 % fx.Int64(2)) * fx.Int64(2) + n_group16 // fx.Int64(2) - aligned_base = ((self.lds_vt_base_idx + fx.Int64(127)) // fx.Int64(128)) * fx.Int64(128) - tile_bytes = traits.HEAD_DIM_V * traits.FP8_V_ROW_STRIDE - dst_byte = ( - aligned_base - + fx.Int64(buf_id * tile_bytes) - + fx.Int64(d_offset * traits.FP8_V_ROW_STRIDE) - + d_local * fx.Int64(traits.FP8_V_ROW_STRIDE) - + dest_group16 * fx.Int64(16) - ) - self._store_v_fp8_lds(reordered, dst_byte - self.lds_vt_base_idx) - - _stage_active() - def _stage_v_fp8_vectorized_bankpad(self, tile_start, buf_id, page_id=None): """Stage vectorized paged V in the bank-padded direct-FP8 layout.""" src_i32x4 = self._load_v_fp8_vectorized_bankpad_source(tile_start, page_id=page_id) @@ -5398,14 +5285,14 @@ def _load_v_fp8_vectorized_bankpad_source(self, tile_start, page_id=None): + d_col * traits.KV_VEC_SIZE ) source = self.buffer_load_fp8x16(src_div, src_elem) - if const_expr(traits.BN128 and traits.FP8_PV_SEGMENTED): + if const_expr(traits.FP8_PV_SEGMENTED): tail = self.buffer_load_fp8x16(src_div, src_elem + fx.Int64(traits.FP8_V_H1 * traits.KV_VEC_SIZE)) source = Vec(source).shuffle(Vec(tail), [0, 1, 2, 3, 4, 5, 6, 7]).ir_value() return source def _store_v_fp8_vectorized_bankpad(self, src_i32x4, buf_id): """Permute prefetched V vectors into the paged FP8 LDS layout.""" - if const_expr(self.traits.BN128 and self.traits.FP8_PV_SEGMENTED): + if const_expr(self.traits.FP8_PV_SEGMENTED): source = Vec(src_i32x4) prefix = source.shuffle(source, [0, 1, 2, 3]).ir_value() tail = source.shuffle(source, [4, 5, 6, 7]).ir_value() @@ -5430,7 +5317,7 @@ def _store_v_fp8_vectorized_bankpad_segment(self, src_i32x4, buf_id, d_offset): d_col = fx.Int64(d_offset) + self.wave_id * fx.Int64(16) + self.lane_in_warp % fx.Int64(16) reordered = self._permute_v_fp8_vectorized(src_i32x4) dest_group16 = (n_group16 % fx.Int64(2)) * fx.Int64(2) + n_group16 // fx.Int64(2) - if const_expr(traits.BN128 and traits.HEAD_DIM_V == 192): + if const_expr(traits.HEAD_DIM_V == 192): # XOR 16-byte groups instead of padding each row in the six-slot ring. dest_group16 = dest_group16 ^ ((d_col // fx.Int64(4)) % fx.Int64(4)) dst_byte = dst_base + d_col * fx.Int64(row_stride) + dest_group16 * fx.Int64(16) @@ -5481,9 +5368,6 @@ def _read_strip(key): return (_read_strip(n_lo), _read_strip(n_hi)) def load_v(self, buf_id): - if const_expr(self.traits.FP8_PV_SEGMENTED and not self.traits.BN128): - # P*V reads segmented chunks on demand; carry only their LDS slot. - return fx.Int64(buf_id) if const_expr(self.traits.PAGED): return self._load_v_fp8_vectorized_bankpad(buf_id) return self._load_v_fp8_block(buf_id) @@ -5502,7 +5386,7 @@ def _load_v_fp8_vectorized_bankpad(self, buf_id): row_base = tile_base + d_row * fx.Int64(row_stride) for half in range_constexpr(2): token_offset = token_base + fx.Int64(half * 16) - if const_expr(traits.BN128 and traits.HEAD_DIM_V == 192): + if const_expr(traits.HEAD_DIM_V == 192): token_offset = token_offset ^ (((d_row // fx.Int64(4)) % fx.Int64(4)) * fx.Int64(16)) view = fx.slice( self.v_lds_i32_tiles, (None, fx.Uint32(row_base + token_offset - self.lds_vt_base_idx) // 16) @@ -5569,8 +5453,14 @@ def _causal_mask_inplace(self, v_s, tile_idx): rel_hi_i32 = fx.Int32(rel_lo_i32 - fx.Int32(32)) neg_inf_i32 = fx.Int32(traits.NEG_INF_F32_BITS) pair_thresholds = _causal_pair_thresholds(False) - _apply_dualwave_causal_mask_pair(s_lo, rel_lo_i32, neg_inf_i32, pair_thresholds) - _apply_dualwave_causal_mask_pair(s_hi, rel_hi_i32, neg_inf_i32, pair_thresholds) + if const_expr(self.traits.PAGED): + for r in range_constexpr(16): + threshold = (r // 4) * 8 + r % 4 + s_lo[r] = (rel_lo_i32 < threshold).select(self.c_neg_inf, s_lo[r]) + s_hi[r] = (rel_hi_i32 < threshold).select(self.c_neg_inf, s_hi[r]) + else: + _apply_dualwave_causal_mask_pair(s_lo, rel_lo_i32, neg_inf_i32, pair_thresholds) + _apply_dualwave_causal_mask_pair(s_hi, rel_hi_i32, neg_inf_i32, pair_thresholds) def causal_mask_prologue_if_needed(self, v_s, tile_idx=None, kv_end_pos=None): if tile_idx is None: @@ -5662,8 +5552,6 @@ def floor_masked_max(self, row_max): _P_HEADROOM_LOG2 = 8.807354922057604 def sub_m(self, v_s, row_max): - if const_expr(self.traits.PAGED and not self.traits.BN128): - return _scale_sub_score_pair(v_s, row_max, self.c_logit_scale, self.c_zero_f, self.fm_fast) # P is cast to e4m3, whose smallest subnormal is 2**-9, so a softmax # over thousands of keys loses its tail to flush-to-zero -- while l_row, # summed before the cast, still counts it. Scaling P up first uses the @@ -5685,9 +5573,7 @@ def sub_m(self, v_s, row_max): self.c_zero_f, self.fm_fast, bias, - scalar_fma=self.traits.PAGED - and self.traits.BN128 - and (self.traits.HEAD_DIM != 128 or not self.traits.VARLEN), + constrained_fma=self.traits.PAGED and (self.traits.HEAD_DIM != 128 or not self.traits.VARLEN), ) def exp2(self, v_s, start, length): @@ -5721,9 +5607,11 @@ def anchor_v_p(self, v_p): return _anchor_v_p(self.traits, v_p, elem_dtype=self.p_elem) def anchor_v_o(self, v_o): - return _anchor_v_o(self.traits, v_o) + return self.preserve_accumulators(v_o) def anchor_scalar_f32(self, x): + if const_expr(self.traits.PAGED): + return llvm.intr_arithmetic_fence(fx.as_ir_value(x)) return _anchor_scalar_f32(x) def safe_l_inv(self, l_row): @@ -5762,9 +5650,6 @@ def _lazy_correction(self, v_o, m_row, m_tile_max): return m_new, corr, [as_mlir_value(scaled_accs[i]) for i in range(self.traits.D_CHUNKS)] def lazy_rescale_o(self, v_o, m_row, l_row, m_tile_max, v_p): - if const_expr(self.traits.PAGED): - return self._lazy_rescale_o_paged(v_o, m_row, l_row, m_tile_max, v_p) - @flyc.jit def _run(v_o, m_row, l_row, m_tile_max, v_p): m_diff = m_tile_max - m_row @@ -5790,33 +5675,6 @@ def _run(v_o, m_row, l_row, m_tile_max, v_p): return _run(v_o, m_row, l_row, m_tile_max, v_p) - def _lazy_rescale_o_paged(self, v_o, m_row, l_row, m_tile_max, v_p): - @flyc.jit - def _run(v_o, m_row, l_row, m_tile_max, v_p): - m_diff = m_tile_max - m_row - m_diff_scaled = m_diff * self.c_logit_scale - below = fx.Float32(m_diff_scaled) <= self.c_rescale_thr_f - ballot = fx.Int64(rocdl.ballot(T.i64, below.ir_value())) - all_below = ballot == fx.Int64(_read_exec_i64()) - # Most tiles reuse the running maximum; retain the likely-branch hint. - all_below = fx.Boolean(llvm.intr_expect(all_below.ir_value(), fx.Boolean(True).ir_value())) - - o_out = list(v_o) - vp_out = self.v_p_to_vec32(v_p) - l_out = l_row - m_out = m_row - if all_below: - pass - else: - m_new, corr, scaled_accs = self._lazy_correction(v_o, m_row, m_tile_max) - o_out = scaled_accs - vp_out = self.v_p_to_vec32(self.scale_v_p(v_p, corr)) - l_out = l_row * corr - m_out = self.anchor_scalar_f32(m_new) - return (o_out, m_out, l_out, self.v_vec32_to_p(vp_out)) - - return _run(v_o, m_row, l_row, m_tile_max, v_p) - def lazy_correct_o(self, v_o, m_row, l_row, m_tile_max): @flyc.jit def _run(v_o, m_row, l_row, m_tile_max): @@ -5848,6 +5706,10 @@ def __init__(self, ctx): def _o_pack_2dw(self, v_o, dc, store_group): r_base = store_group * 4 + if const_expr(self.traits.PAGED): + values = Vec(v_o[dc]) + packed = values.shuffle(values, list(range(r_base, r_base + 4))).to(fx.BFloat16).bitcast(fx.Int32) + return packed[0].ir_value(), packed[1].ir_value() lo = rocdl.cvt_pk_bf16_f32(Vec(v_o[dc])[r_base], Vec(v_o[dc])[r_base + 1]) hi = rocdl.cvt_pk_bf16_f32(Vec(v_o[dc])[r_base + 2], Vec(v_o[dc])[r_base + 3]) return lo, hi @@ -5882,11 +5744,6 @@ def store_final_o(self, v_o, q_row): o_global = self.global_idx_o(q_row, d_col) self.buffer_store_128(o_pack, o_global) - @flyc.jit - def store_final_o_if_valid(self, v_o, q_row): - if q_row < self.seqlen_q_v: - self.store_final_o(v_o, q_row) - def store_splitk_partial_o(self, v_o, m_row, l_row, q_row): m_row = fx.Float32(m_row) * self.c_logit_scale split_z = self.batch_idx * self.traits.NUM_KV_SPLITS + self.split_idx diff --git a/tests/kernels/test_flash_attn_fwd.py b/tests/kernels/test_flash_attn_fwd.py index 1830b09dc6..c2a5d55a48 100644 --- a/tests/kernels/test_flash_attn_fwd.py +++ b/tests/kernels/test_flash_attn_fwd.py @@ -4648,6 +4648,7 @@ def _check_paged_fp8_matches_torch( kv_lengths, block_table_rows, num_kv_heads=1, + lazy_rescale=True, ): """Packed causal FP8 page-64 attention supports native Q/K and V widths.""" if len(query_lengths) == 1 and force_internal_copies and head_dim == 192: @@ -4702,6 +4703,7 @@ def noncontiguous_copy(tensor): q_descale=query_descale, k_descale=key_descale, v_descale=value_descale, + dualwave_swp_lazy_rescale=lazy_rescale, ) copy_reference = None if force_internal_copies: @@ -4765,6 +4767,38 @@ def noncontiguous_copy(tensor): torch.testing.assert_close(actual, expected, rtol=2.0e-2, atol=2.0e-2) +@_requires_gfx950 +@pytest.mark.parametrize("head_dim,value_head_dim", [(128, 128), (192, 128), (192, 192)]) +@pytest.mark.parametrize("num_kv_heads", [1, 2, 4]) +@pytest.mark.parametrize("lazy_rescale", [True, False]) +def test_paged_fp8_odd_page_ragged_multiblock_matches_torch(head_dim, value_head_dim, num_kv_heads, lazy_rescale): + _check_paged_fp8_matches_torch( + head_dim=head_dim, + value_head_dim=value_head_dim, + use_non_default_stream=True, + force_internal_copies=False, + query_lengths=[300, 257, 33], + kv_lengths=[513, 385, 129], + block_table_rows=[list(reversed(range(9))), list(reversed(range(9, 16))) + [0] * 2, [18, 16, 17] + [0] * 6], + num_kv_heads=num_kv_heads, + lazy_rescale=lazy_rescale, + ) + + +@_requires_gfx950 +@pytest.mark.parametrize("head_dim,value_head_dim", [(128, 128), (192, 128), (192, 192)]) +def test_paged_fp8_single_partial_page_matches_torch(head_dim, value_head_dim): + _check_paged_fp8_matches_torch( + head_dim=head_dim, + value_head_dim=value_head_dim, + use_non_default_stream=False, + force_internal_copies=False, + query_lengths=[17], + kv_lengths=[63], + block_table_rows=[[0]], + ) + + @_requires_gfx950 @pytest.mark.parametrize("value_head_dim", [128, 192]) @pytest.mark.parametrize("num_kv_heads", [2, 4]) From 9251049e7402ee034be00e13eed2ff50fbe8ca62 Mon Sep 17 00:00:00 2001 From: Xiake Sun Date: Fri, 11 Sep 2026 09:31:48 +0000 Subject: [PATCH 55/59] perf(fmha): optimize native small-page FP8 attention - Support native physical page sizes 1, 16, 64, and 1024 with independent QK/V widths and bounds-safe cache addressing. - Replace page-1 byte gathers with vector loads and an exact wave transpose; specialize page-16 K/V staging and reuse D128 page translations. - Use bounded cache descriptors for small pools while retaining the validated 64-bit fallback for larger caches. - Preserve page-64/page-1024 ISA and qualify accuracy, stream/graph behavior, large offsets, and the complete performance matrix. Signed-off-by: Xiake Sun --- .../attention/flash_attn_fp8_paged_gfx950.py | 26 +- kernels/attention/flash_attn_interface.py | 75 +++- kernels/attention/flash_attn_utils.py | 287 +++++++++++++- tests/kernels/test_flash_attn_fwd.py | 362 +++++++++++++++--- 4 files changed, 662 insertions(+), 88 deletions(-) diff --git a/kernels/attention/flash_attn_fp8_paged_gfx950.py b/kernels/attention/flash_attn_fp8_paged_gfx950.py index f8a76301c9..c231c06ef4 100644 --- a/kernels/attention/flash_attn_fp8_paged_gfx950.py +++ b/kernels/attention/flash_attn_fp8_paged_gfx950.py @@ -1,13 +1,14 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2025 FlyDSL Project Contributors -"""gfx950 page-64 FP8 flash attention with independent value dimensions.""" +"""gfx950 paged FP8 flash attention with independent value dimensions.""" import flydsl.compiler as flyc import flydsl.expr as fx from flydsl.expr import const_expr, range_constexpr, rocdl from flydsl.runtime.device import get_rocm_arch as get_hip_arch from kernels.attention.flash_attn_utils import ( + PAGED_FP8_BUFFER_LIMIT_BYTES, DualwaveFp8GemmHelper, DualwaveFp8KernelContext, DualwaveFp8KvGmemToLdsLoader, @@ -51,8 +52,10 @@ def build_flash_attn_paged_fp8_module( paged_bn128=False, paged_bn128_varlen=False, batch_interleave_group=1, + page_size=64, + cache_buffered=False, ): - """Build the gfx950 packed-varlen page-64 FP8 attention launcher. + """Build the gfx950 packed-varlen paged FP8 attention launcher. Priority/stagger options are accepted for caller compatibility; this paired-page pipeline uses neither wave-priority nor staggered phases. @@ -63,6 +66,9 @@ def build_flash_attn_paged_fp8_module( if not gpu_arch.startswith("gfx950"): raise RuntimeError(f"paged FP8 flash attention requires gfx950, got {gpu_arch}") + native_cache_layout = (page_size == 1 and kv_cache_layout in ("linear", "linear3d")) or ( + page_size in (16, 64, 1024) and kv_cache_layout == "vectorized" + ) if ( not paged or dtype_str != "fp8" @@ -70,18 +76,22 @@ def build_flash_attn_paged_fp8_module( or not varlen or not cross_seqlen or (head_dim, value_head_dim) not in ((128, 128), (192, 128), (192, 192)) - or kv_cache_layout != "vectorized" + or not native_cache_layout or int(num_kv_splits) != 1 ): raise RuntimeError( "paged FP8 flash_attn requires gfx950, causal packed-varlen cross-attention, " - "page-64 vectorized KV, (head_dim,value_head_dim) in " + "page-16/64/1024 vectorized or page-1 linear/linear3d KV, (head_dim,value_head_dim) in " "{(128,128),(192,128),(192,192)}, " "and num_kv_splits=1" ) if num_kv_heads is None: num_kv_heads = num_heads + if paged_bn128 and page_size != 64: + raise ValueError("paired page-ID loads require physical page size 64") + if cache_buffered and page_size not in (1, 16): + raise ValueError("whole-cache buffer descriptors are reserved for page sizes 1/16") batch_interleave_group = int(batch_interleave_group) if batch_interleave_group < 1: raise ValueError(f"batch_interleave_group must be positive, got {batch_interleave_group}") @@ -103,6 +113,9 @@ def build_flash_attn_paged_fp8_module( varlen=not paged_bn128 or paged_bn128_varlen, paired_page_ids=paged_bn128, batch_interleave_group=batch_interleave_group, + page_size=page_size, + kv_cache_layout=kv_cache_layout, + cache_buffered=cache_buffered, ) BLOCK_M = traits.BLOCK_M BLOCK_SIZE = traits.BLOCK_SIZE @@ -547,6 +560,11 @@ def _launch( stream=None, _compile_only=False, ): + if ( + cache_buffered + and max(K.numel() * K.element_size(), V.numel() * V.element_size()) > PAGED_FP8_BUFFER_LIMIT_BYTES + ): + raise ValueError("paged FP8 whole-cache buffer descriptor exceeds its byte limit") if stride_kv_n is None: stride_kv_n = DEFAULT_STRIDE_KV_N if stride_q_n is None: diff --git a/kernels/attention/flash_attn_interface.py b/kernels/attention/flash_attn_interface.py index 0d7830a240..8b84312db5 100644 --- a/kernels/attention/flash_attn_interface.py +++ b/kernels/attention/flash_attn_interface.py @@ -34,6 +34,7 @@ DUALWAVE_SWP_BLOCK_M, MIN_Q_BLOCKS_XCD_SWIZZLE, NUM_XCD_GFX950, + PAGED_FP8_BUFFER_LIMIT_BYTES, bias_addressing_error, dualwave_splitk_workspace_elems, ) @@ -478,8 +479,11 @@ def _build_paged_fp8( use_bn128: bool, paged_bn128_varlen: bool, batch_interleave_group: int, + page_size: int = 64, + kv_cache_layout: str = "vectorized", + cache_buffered: bool = False, ): - """Build the gfx950 packed-varlen, vectorized page-64 FP8 launcher.""" + """Build the gfx950 packed-varlen paged FP8 launcher for a physical cache ABI.""" from kernels.attention.flash_attn_fp8_paged_gfx950 import build_flash_attn_paged_fp8_module return build_flash_attn_paged_fp8_module( @@ -496,10 +500,12 @@ def _build_paged_fp8( varlen=True, cross_seqlen=True, paged=True, - kv_cache_layout="vectorized", + kv_cache_layout=kv_cache_layout, paged_bn128=use_bn128, paged_bn128_varlen=paged_bn128_varlen, batch_interleave_group=batch_interleave_group, + page_size=page_size, + cache_buffered=cache_buffered, ) @@ -559,18 +565,18 @@ def _flydsl_flash_attn_paged( ) -> torch.Tensor: """Native paged-KV attention on the gfx950 dualwave kernel. - Supported config ONLY (anything else raises): linear/vectorized cache layout - with page size 64 and vLLM ``block_table`` / ``seqlen_k`` metadata. BF16/F16 - support D64/D128. gfx950 FP8 additionally supports packed-varlen causal - Q/K D128 with V/output D128, or Q/K D192 with vectorized V/output - D128 or D192. + BF16/F16 support page-64 linear/vectorized caches and D64/D128. gfx950 + FP8 supports packed-varlen causal QK/V widths 128/128, 192/128, 192/192, + with vectorized page-16/64/1024 caches or linear/linear3d page-1 caches. + All paths use vLLM ``block_table`` / ``seqlen_k`` metadata. - Dense 4D Q ``[B, Sq, H, D]``: split-K (num_kv_splits>1) supported (seq_len>=384). - Varlen packed Q ``[total_q, H, D]`` (cu_seqlens_q given): paged K/V looked up per kv-tile via block_table; paged split-K is not supported. """ - if kv_cache_layout not in ("linear", "vectorized"): + if kv_cache_layout not in ("linear", "linear3d", "vectorized"): raise NotImplementedError( - f"flydsl_flash_attn_func: native paged KV supports kv_cache_layout in ('linear','vectorized'), " + "flydsl_flash_attn_func: native paged KV supports kv_cache_layout in " + "('linear','linear3d','vectorized'), " f"got {kv_cache_layout!r}" ) if block_table is None or seqlen_k is None: @@ -580,7 +586,10 @@ def _flydsl_flash_attn_paged( # aiter 5D: K [NumBlocks, Hkv, D/kVS, PageSize, kVS], V [NumBlocks, Hkv, PageSize/kVS, D, kVS]. if k.dim() != 5 or v.dim() != 5: raise ValueError(f"flydsl_flash_attn_func: vectorized paged K/V must be 5D, got K{k.dim()}D V{v.dim()}D") - elif k.dim() != 4: + elif kv_cache_layout == "linear3d": + if k.dim() != 3 or v.dim() != 3: + raise ValueError("flydsl_flash_attn_func: linear3d paged K/V must be 3D [NumBlocks,Hkv,D]") + elif k.dim() != 4 or v.dim() != 4: raise ValueError( f"flydsl_flash_attn_func: linear paged K/V must be 4D [NumBlocks,PageSize,Hkv,D], got {k.dim()}D" ) @@ -611,6 +620,8 @@ def _flydsl_flash_attn_paged( kvs = 16 // k.element_size() Hkv = int(k.shape[1]) page_size = int(k.shape[3]) + if page_size % kvs: + raise ValueError(f"flydsl_flash_attn_func: vectorized page size must be divisible by kVS={kvs}") k_head_dim = int(k.shape[2]) * int(k.shape[4]) # (D/kVS) * kVS if int(k.shape[4]) != kvs: raise ValueError(f"flydsl_flash_attn_func: vectorized K last dim ({k.shape[4]}) must equal kVS={kvs}") @@ -620,6 +631,13 @@ def _flydsl_flash_attn_paged( raise ValueError( f"flydsl_flash_attn_func: vectorized V tail must be {expected_v_tail}, got {tuple(v.shape[1:])}" ) + elif kv_cache_layout == "linear3d": + page_size = 1 + Hkv = int(k.shape[1]) + k_head_dim = int(k.shape[2]) + value_head_dim = int(v.shape[2]) + if int(v.shape[1]) != Hkv: + raise ValueError("flydsl_flash_attn_func: linear3d K/V must have matching KV head counts") else: page_size = int(k.shape[1]) Hkv = int(k.shape[2]) @@ -629,25 +647,32 @@ def _flydsl_flash_attn_paged( raise ValueError( f"flydsl_flash_attn_func: linear V must match K page/head axes, got K{tuple(k.shape)} V{tuple(v.shape)}" ) - if page_size != _PAGED_PAGE_SIZE: + supported_page_sizes = (1, 16, 64, 1024) if paged_fp8 else (_PAGED_PAGE_SIZE,) + if page_size not in supported_page_sizes: raise NotImplementedError( - f"flydsl_flash_attn_func: native paged KV supports page_size={_PAGED_PAGE_SIZE} only, got {page_size}" + f"flydsl_flash_attn_func: native paged KV supports page sizes {supported_page_sizes}, got {page_size}" ) if k_head_dim != D: raise ValueError(f"flydsl_flash_attn_func: paged K head_dim ({k_head_dim}) must match q head_dim ({D})") if paged_fp8: + if k.shape[0] != v.shape[0]: + raise ValueError("flydsl_flash_attn_func: paged FP8 K/V must have matching physical page counts") + if num_kv_heads is not None and num_kv_heads != Hkv: + raise ValueError("flydsl_flash_attn_func: num_kv_heads must match the paged FP8 cache") if not arch.startswith("gfx950"): raise ValueError(f"flydsl_flash_attn_func: paged FP8 requires gfx950, got '{arch or 'unknown'}'") fp8_head_dims = (D, value_head_dim) + native_layout = (vectorized and page_size != 1) or (not vectorized and page_size == 1) if not ( causal and varlen and cross_seqlen is not False - and vectorized + and native_layout and fp8_head_dims in ((128, 128), (192, 128), (192, 192)) ): raise NotImplementedError( - "flydsl_flash_attn_func: paged FP8 requires causal packed-varlen vectorized KV, " + "flydsl_flash_attn_func: paged FP8 requires causal packed-varlen KV, " + "vectorized page-16/64/1024 or linear/linear3d page-1 caches, " f"Q/K-V D128-D128, D192-D128, or D192-D192; got causal={causal}, varlen={varlen}, " f"layout={kv_cache_layout!r}, Q/K D{D}, V D{value_head_dim}" ) @@ -736,6 +761,11 @@ def _flydsl_flash_attn_paged( f"got shape={tuple(seqlen_k.shape)} dtype={seqlen_k.dtype} device={seqlen_k.device}" ) block_table_stride = int(block_table.shape[1]) + if paged_fp8 and (block_table.shape[0] != B or block_table_stride < max_kv_pages): + raise ValueError( + f"flydsl_flash_attn_func: paged FP8 block_table must have {B} rows and at least {max_kv_pages} " + f"physical page entries per row, got {tuple(block_table.shape)}" + ) expected_out_shape = (*q.shape[:-1], value_head_dim) q_flat_elems = q.numel() out_flat_elems = q_flat_elems // D * value_head_dim @@ -772,7 +802,7 @@ def _flydsl_flash_attn_paged( and (not arch.startswith("gfx950") or Sq <= _VARLEN_LIGHT_MAX_SEQ) ) if paged_fp8: - use_bn128 = max_kv_pages % 2 == 0 + use_bn128 = page_size == 64 and max_kv_pages % 2 == 0 paged_bn128_varlen = use_bn128 and B > 1 exe = _build_paged_fp8( num_heads=H, @@ -789,6 +819,12 @@ def _flydsl_flash_attn_paged( if H == 16 and num_kv_heads == 1 else 1 ), + page_size=page_size, + kv_cache_layout=kv_cache_layout, + cache_buffered=( + page_size in (1, 16) + and max(k.numel() * k.element_size(), v.numel() * v.element_size()) <= PAGED_FP8_BUFFER_LIMIT_BYTES + ), ) elif _paged_light_ok: exe = _build_paged_light( @@ -837,7 +873,7 @@ def _flydsl_flash_attn_paged( if out is None: out_dtype = torch.bfloat16 if paged_fp8 else q.dtype out = torch.empty(expected_out_shape, dtype=out_dtype, device=q.device) - # Keep serving-sized physical K/V caches rank-5 because flattening + # Keep serving-sized physical K/V caches in their native rank because flattening # their dynamic memref shape can exceed signed int32. The FP8 # schedule consumes Q/O as flat token-major buffers, matching its # explicit runtime strides. @@ -987,7 +1023,7 @@ def flydsl_flash_attn_func( ``kv_cache_layout`` values: - ``linear``: 4D paged K/V, ``[NumBlocks, PageSize, NumKVHeads, HeadDim]``. - ``linear3d``: page_size=1 special case, - ``[NumBlocks, NumKVHeads, HeadDim]``. + ``[NumBlocks, NumKVHeads, HeadDim]`` (gfx950 FP8 only). - ``vectorized``: aiter-style 5D K/V, where ``K = [NumBlocks, NumKVHeads, HeadDim / kVectorSize, PageSize, kVectorSize]`` and @@ -1006,8 +1042,9 @@ def flydsl_flash_attn_func( block_table / seqlen_k: vLLM-style 2D block table metadata. Enables the native paged-KV path, which supports ``bias`` but not ``alibi_slopes``, ``sink``, or ``return_lse``. gfx950 FP8 supports - causal packed-varlen vectorized page-64 D128/V128 and - D192/V128-or-V192 paths. + causal packed-varlen D128/V128 and D192/V128-or-V192 paths, + with vectorized page sizes 16/64/1024 or linear/linear3d page 1. + BF16/F16 native paged paths require page size 64. num_kv_splits: Split-K factor (>1: gfx950 only, D=64/128, bf16/f16, seq>=384). ``None`` lets fp8 autotune it; ``1`` keeps the kernel unsplit. fp8_block_m: Pin the fp8 tile height to 128 or 256. fp8 only. diff --git a/kernels/attention/flash_attn_utils.py b/kernels/attention/flash_attn_utils.py index ba6271a78f..c3cbbd9d6d 100644 --- a/kernels/attention/flash_attn_utils.py +++ b/kernels/attention/flash_attn_utils.py @@ -25,6 +25,8 @@ # The dual-wave 8-wave CTA fixes the q-block height; callers need it to count # q-blocks before any traits object exists. DUALWAVE_SWP_BLOCK_M = 256 +# Reserve an aligned out-of-bounds vector offset below the signed-i32 limit. +PAGED_FP8_BUFFER_LIMIT_BYTES = (1 << 31) - 16 _LDS_ALIAS_DOMAIN = '#llvm.alias_scope_domain' @@ -2083,7 +2085,7 @@ def dualwave_fp8_dma_per_iter(traits): @dataclass(frozen=True) class PagedDualwaveSwpFp8Traits: - """Layouts for native page-64 FP8 QK/PV with register-resident Q.""" + """Layouts for native paged FP8 QK/PV with register-resident Q.""" HEAD_DIM: int HEAD_DIM_V: int @@ -2128,6 +2130,8 @@ class PagedDualwaveSwpFp8Traits: PAGED: bool = True KV_VECTORIZED: bool = True PAGE_SIZE: int = 64 + KV_CACHE_LAYOUT: str = "vectorized" + CACHE_BUFFERED: bool = False DMA_BYTES: int = 16 ELEM_BYTES: int = 1 OUT_ELEM_BYTES: int = 2 @@ -2158,6 +2162,9 @@ def cache_tag(self): self.FP8_V_ROW_STRIDE, self.BATCH_INTERLEAVE_GROUP, self.PAIRED_PAGE_IDS, + self.PAGE_SIZE, + self.KV_CACHE_LAYOUT, + self.CACHE_BUFFERED, ) @@ -2185,6 +2192,9 @@ def _make_paged_dualwave_swp_fp8_traits( varlen=False, batch_interleave_group=1, paired_page_ids=True, + page_size=64, + kv_cache_layout="vectorized", + cache_buffered=False, ): """Build layouts after the dedicated builder validates the paged contract.""" block_n = 64 @@ -2218,6 +2228,9 @@ def _make_paged_dualwave_swp_fp8_traits( FP8_V_H1=fp8_v_h1, FP8_V_H2=fp8_v_h2, PAIRED_PAGE_IDS=bool(paired_page_ids), + PAGE_SIZE=page_size, + KV_CACHE_LAYOUT=kv_cache_layout, + CACHE_BUFFERED=bool(cache_buffered), DEFAULT_STRIDE_Q_N=num_heads * head_dim, DEFAULT_STRIDE_KV_N=num_kv_heads * head_dim, SMEM_D_RPT=smem_d_rpt, @@ -4747,6 +4760,18 @@ def _make_buf_div(tensor, nrec_bytes): self.page_base = fx.Int64(self.batch_idx) * fx.Int64(self.block_table_stride) self.k_div = None self.v_div = None + if const_expr(traits.CACHE_BUFFERED): + + def _cache_div(tensor, base_iter): + num_bytes = fx.Int64(fx.get_scalar(fx.cosize(fx.get_layout(tensor)))) + flat = fx.make_view(base_iter, fx.make_layout(num_bytes, 1)) + # make_buffer_tensor defaults to an unbounded resource; + # supply the exact byte count for zero-filling OOB copies. + bounded = fx.rocdl.make_buffer_tensor(flat, num_records_bytes=num_bytes.ir_value()) + return fx.logical_divide(bounded, fx.make_layout(1, 1)) + + self.k_div = _cache_div(self.K, self.k_base_iter) + self.v_div = _cache_div(self.V, self.v_base_iter) else: kv_nrec_bytes = (self.kv_tok_end * self.stride_kv_n_v * eb).ir_value() v_nrec_bytes = (self.kv_tok_end * self.stride_v_n_v * eb).ir_value() @@ -4877,6 +4902,16 @@ def buffer_load_fp8x16(self, src_div, elem_index): [self.v4i32_type], self.load_atom_128, fx.slice(src_div, (None, fx.Int32(elem_index))) ) + def global_load_fp8x16(self, base_iter, byte_offset, valid, *, is_value=False): + """Load a native FP8 vector through a bounded cache or a 64-bit pointer.""" + if const_expr(self.traits.CACHE_BUFFERED): + offset = valid.select(fx.Int32(byte_offset), fx.Int32(PAGED_FP8_BUFFER_LIMIT_BYTES)) + return self.buffer_load_fp8x16(self.v_div if is_value else self.k_div, offset) + # Do not issue an out-of-allocation load before selecting zero padding. + safe_offset = valid.select(fx.Int64(byte_offset), fx.Int64(0)) + loaded = fx.generic_load(fx.add_offset(base_iter, safe_offset), dtype=fx.Int32, count=4) + return valid.select(loaded, fx.Vector.filled(4, 0, fx.Int32)).ir_value() + def buffer_load_lds_128(self, src_div, lds_byte_addr, src_elem, soffset_elems): _buffer_load_lds_128( src_div, lds_byte_addr, src_elem, soffset_elems, _dma_atom=self.dma_atom, _lds_ptr_ty=self.lds_ptr_ty @@ -4891,25 +4926,40 @@ def global_idx_q(self, token_idx, col): def global_idx_o(self, token_idx, col): return (self.q_tok_base + token_idx) * self.stride_o_n_v + self.q_head_idx * self.traits.HEAD_DIM_V + col - def load_page_id(self, tile_start): - """Load this request's physical page for a 64-token logical KV tile.""" + def load_page_id(self, tile_start, *, uniform=True): + """Load a logical token's physical page, optionally uniform across the wave.""" local_page = fx.Int64(tile_start) // self.traits.PAGE_SIZE # The software pipeline rounds to an even tile count and at least four # tiles. Redirect those masked padding tiles to entry zero instead of # reading past the final request's page-index allocation. page_idx = self.page_base + local_page - safe_idx = (local_page < self.num_kv_tiles).select(page_idx, 0) + if const_expr(self.traits.PAGE_SIZE == self.traits.BLOCK_N): + num_pages = self.num_kv_tiles + else: + num_pages = (self.seqlen_kv_v + self.traits.PAGE_SIZE - 1) // self.traits.PAGE_SIZE + safe_idx = (local_page < num_pages).select(page_idx, 0) v = fly.copy_atom_call_ssa( [self.page_v1i32], self.page_i32_atom, fx.slice(self.page_indices_div, (None, fx.Int32(safe_idx))), ) - return fx.Int64(rocdl.readfirstlane(T.i32, Vec(v, (1,), fx.Int32)[0].ir_value())) + page = Vec(v, (1,), fx.Int32)[0] + if uniform: + return fx.Int64(rocdl.readfirstlane(T.i32, page.ir_value())) + return fx.Int64(page) def load_page_id_pair(self, tile_start): - """Load two adjacent physical page IDs for a BN128 tile pair.""" + """Return reusable page IDs for a compute-tile pair, or defer to producers.""" + if const_expr(self.traits.PAGE_SIZE == 16 and self.traits.HEAD_DIM == 128 and self.traits.CACHE_BUFFERED): + # K and V use the same four lane-group pages. Share their metadata + # loads across the two producers instead of looking them up twice. + token = fx.Int64(tile_start) + (self.lane_in_warp // fx.Int64(16)) * fx.Int64(16) + return self.load_page_id(token, uniform=False), self.load_page_id(token + 64, uniform=False) + if const_expr(self.traits.PAGE_SIZE < self.traits.BLOCK_N): + # A compute tile spans several pages; its loaders gather IDs per lane. + return None, None if const_expr(not self.traits.PAIRED_PAGE_IDS): - return self.load_page_id(tile_start), self.load_page_id(tile_start + self.traits.PAGE_SIZE) + return self.load_page_id(tile_start), self.load_page_id(tile_start + self.traits.BLOCK_N) local_page = fx.Int32(fx.Int64(tile_start) // self.traits.PAGE_SIZE) page_base = fx.Int32(self.page_base) num_kv_tiles = fx.Int32(self.num_kv_tiles) @@ -5104,6 +5154,32 @@ def pv(self, v_p, v_v, v_o): return self._pv_fp8_direct(v_p, v_v, v_o) +def _transpose_v_fp8_16x16(source, lane): + """Exchange row/column bits within each 16-lane group; full EXEC is required.""" + words = [fx.Uint32(Vec(source)[i]) for i in range_constexpr(4)] + lane = fx.Int32(lane) + for bit in range_constexpr(4): + peer_addr = (lane ^ (1 << bit)) * 4 + lane_bit = (lane & (1 << bit)) != 0 + updated = [] + for word in range_constexpr(4): + peer_word = word if bit < 2 else word ^ (1 << (bit - 2)) + peer = fx.Uint32(rocdl.ds_bpermute(T.i32, peer_addr, words[peer_word])) + if const_expr(bit < 2): + shift = 8 << bit + low_mask = fx.Uint32(0x00FF00FF if bit == 0 else 0x0000FFFF) + high_mask = fx.Uint32(0xFFFFFFFF) ^ low_mask + low_row = (words[word] & low_mask) | ((peer & low_mask) << shift) + high_row = (words[word] & high_mask) | ((peer & high_mask) >> shift) + updated.append(lane_bit.select(high_row, low_row)) + elif const_expr(word & (1 << (bit - 2))): + updated.append(lane_bit.select(words[word], peer)) + else: + updated.append(lane_bit.select(peer, words[word])) + words = updated + return Vec.from_elements([fx.Int32(word) for word in words], fx.Int32) + + class DualwaveFp8KvGmemToLdsLoader(DualwaveFp8KernelContext): def __init__(self, ctx): super().__init__(ctx) @@ -5115,22 +5191,61 @@ def load_k(self, tile_start, buf_id, page_id=None): uses the low 32 lanes; paged tiles retain the physical D-group mapping. """ traits = self.traits + if const_expr(traits.PAGED and traits.PAGE_SIZE == 16 and traits.HEAD_DIM_V == 192): + self._load_k_page16(tile_start, buf_id) + return eb = traits.ELEM_BYTES k_lds_byte_base = self.lds_kv_base_idx + self.k_buf_base(buf_id) * eb if const_expr(traits.KV_VECTORIZED): - if page_id is None: - page_id = self.load_page_id(tile_start) - src_div = self.make_page_view(self.k_base_iter, page_id) + src_div = None + if const_expr(traits.PAGE_SIZE >= traits.BLOCK_N): + if page_id is None: + page_id = self.load_page_id(tile_start) + src_div = self.make_page_view(self.k_base_iter, page_id) def _load_vectorized(): for d in range_constexpr(self.NUM_DMA_K): - src_elem = _vec_k_src_elem(traits, d, self.wave_id_uni, self.lane_in_warp, self.kv_head_idx) - lds_addr = ( - k_lds_byte_base - + (self.wave_id_uni * traits.SMEM_D_RPT + d) * traits.WARP_SIZE * traits.KV_VEC_SIZE - ) - lds_addr = rocdl.readfirstlane(T.i32, fx.Int32(lds_addr).ir_value()) - self.buffer_load_lds_128(src_div, lds_addr, src_elem, 0) + if const_expr(traits.PAGE_SIZE < traits.BLOCK_N): + oct_idx = _vec_k_dma_oct_idx(traits, d, self.wave_id_uni, self.lane_in_warp) + token = fx.Int64(tile_start) + _sigma_k_tile_n(oct_idx % traits.BLOCK_N) + dim_group = oct_idx // traits.BLOCK_N + physical_page = self.load_page_id(token, uniform=False) if page_id is None else page_id + src_elem = ( + physical_page * self.k_page_bytes + + self.kv_head_idx * traits.HEAD_DIM * traits.PAGE_SIZE + + (dim_group * traits.PAGE_SIZE + token % traits.PAGE_SIZE) * traits.KV_VEC_SIZE + ) + valid = (token < self.seqlen_kv_v) & (dim_group * traits.KV_VEC_SIZE < traits.HEAD_DIM) + if const_expr(traits.CACHE_BUFFERED): + offset = valid.select(fx.Int32(src_elem), fx.Int32(PAGED_FP8_BUFFER_LIMIT_BYTES)) + dst_byte = ( + k_lds_byte_base + + (self.wave_id_uni * traits.SMEM_D_RPT + d) * traits.WARP_SIZE * traits.KV_VEC_SIZE + ) + dst_byte = rocdl.readfirstlane(T.i32, fx.Int32(dst_byte).ir_value()) + self.buffer_load_lds_128(self.k_div, dst_byte, offset, 0) + else: + source = self.global_load_fp8x16(self.k_base_iter, src_elem, valid) + dst_byte = self.k_buf_base(buf_id) + oct_idx * traits.KV_VEC_SIZE + dst = fx.slice(self.k_lds_i32_tiles, (None, fx.Uint32(dst_byte) // 16)) + fx.generic_store(fx.get_iter(dst), fx.Vector(source)) + elif const_expr(traits.PAGE_SIZE == traits.BLOCK_N): + src_elem = _vec_k_src_elem(traits, d, self.wave_id_uni, self.lane_in_warp, self.kv_head_idx) + else: + oct_idx = _vec_k_dma_oct_idx(traits, d, self.wave_id_uni, self.lane_in_warp) + token = fx.Int64(tile_start) % traits.PAGE_SIZE + _sigma_k_tile_n(oct_idx % traits.BLOCK_N) + dim_group = oct_idx // traits.BLOCK_N + src_elem = ( + self.kv_head_idx * traits.HEAD_DIM * traits.PAGE_SIZE + + (dim_group * traits.PAGE_SIZE + token) * traits.KV_VEC_SIZE + ) + if const_expr(traits.PAGE_SIZE >= traits.BLOCK_N): + lds_addr = ( + k_lds_byte_base + + (self.wave_id_uni * traits.SMEM_D_RPT + d) * traits.WARP_SIZE * traits.KV_VEC_SIZE + ) + lds_addr = rocdl.readfirstlane(T.i32, fx.Int32(lds_addr).ir_value()) + self.buffer_load_lds_128(src_div, lds_addr, src_elem, 0) if const_expr(traits.HEAD_DIM_V == 192): @@ -5168,6 +5283,38 @@ def _load_compact(): else: self._load_k_band_partial_wave(lds_addr, src_elem, tile_start, active) + def _load_k_page16(self, tile_start, buf_id): + """Each wave copies four D-groups from one physical page directly to LDS.""" + traits = self.traits + page_in_tile = self.wave_id_uni % fx.Int64(4) + page_id = self.load_page_id(fx.Int64(tile_start) + page_in_tile * fx.Int64(16)) + src_div = self.make_page_view(self.k_base_iter, page_id) + num_dma = traits.HEAD_DIM // traits.KV_VEC_SIZE + for part in range_constexpr((num_dma + traits.NUM_WAVES - 1) // traits.NUM_WAVES): + dma_id = self.wave_id_uni + part * traits.NUM_WAVES + + def _copy_page_groups(): + d_group = (dma_id // fx.Int64(4)) * fx.Int64(4) + self.lane_in_warp // fx.Int64(16) + token = _sigma_k_tile_n(self.lane_in_warp % fx.Int64(16)) + src_elem = ( + self.kv_head_idx * traits.HEAD_DIM * traits.PAGE_SIZE + + (d_group * traits.PAGE_SIZE + token) * traits.KV_VEC_SIZE + ) + dst_byte = self.lds_kv_base_idx + self.k_buf_base(buf_id) + dma_id * fx.Int64(1024) + dst_byte = rocdl.readfirstlane(T.i32, fx.Int32(dst_byte).ir_value()) + self.buffer_load_lds_128(src_div, dst_byte, src_elem, 0) + + if const_expr(num_dma % traits.NUM_WAVES == 0 or part < num_dma // traits.NUM_WAVES): + _copy_page_groups() + else: + + @flyc.jit + def _copy_tail_waves(): + if dma_id < fx.Int64(num_dma): + _copy_page_groups() + + _copy_tail_waves() + def _load_k_band_partial_wave(self, lds_addr, src_elem, tile_start, active_lanes): soffset = tile_start * self.stride_kv_n_v k_div = self.k_div @@ -5269,6 +5416,8 @@ def _store_v_fp8_lds(self, data, byte_offset): def _load_v_fp8_vectorized_bankpad_source(self, tile_start, page_id=None): """Issue vectorized V loads for the paged FP8 LDS layout.""" traits = self.traits + if const_expr(traits.PAGE_SIZE < traits.BLOCK_N): + return self._load_v_fp8_small_pages_source(tile_start, page_id=page_id) if page_id is None: page_id = self.load_page_id(tile_start) src_div = self.make_page_view(self.v_base_iter, page_id, is_value=True) @@ -5280,18 +5429,76 @@ def _load_v_fp8_vectorized_bankpad_source(self, tile_start, page_id=None): n_group16 = self.lane_in_warp // fx.Int64(16) d_col = self.wave_id * fx.Int64(16) + self.lane_in_warp % fx.Int64(16) src_elem = ( - self.kv_head_idx * token_groups16 * traits.HEAD_DIM_V * traits.KV_VEC_SIZE + self.kv_head_idx * traits.PAGE_SIZE * traits.HEAD_DIM_V + n_group16 * traits.HEAD_DIM_V * traits.KV_VEC_SIZE + d_col * traits.KV_VEC_SIZE ) + if const_expr(traits.PAGE_SIZE != traits.BLOCK_N): + src_elem += (fx.Int64(tile_start) % traits.PAGE_SIZE) * traits.HEAD_DIM_V source = self.buffer_load_fp8x16(src_div, src_elem) if const_expr(traits.FP8_PV_SEGMENTED): tail = self.buffer_load_fp8x16(src_div, src_elem + fx.Int64(traits.FP8_V_H1 * traits.KV_VEC_SIZE)) source = Vec(source).shuffle(Vec(tail), [0, 1, 2, 3, 4, 5, 6, 7]).ir_value() return source + def _load_v_fp8_small_pages_source(self, tile_start, page_id=None): + """Load small-page V, grouping D192's page-16 prefetch by physical page.""" + traits = self.traits + if const_expr(traits.PAGE_SIZE == 1): + return self._load_v_fp8_page1_source(tile_start) + if const_expr(traits.HEAD_DIM == 128): + token = fx.Int64(tile_start) + (self.lane_in_warp // fx.Int64(16)) * fx.Int64(16) + d_col = self.wave_id * fx.Int64(16) + self.lane_in_warp % fx.Int64(16) + if page_id is None: + page_id = self.load_page_id(token, uniform=False) + src_elem = ( + page_id * self.v_page_bytes + + self.kv_head_idx * traits.PAGE_SIZE * traits.HEAD_DIM_V + + d_col * traits.KV_VEC_SIZE + ) + return self.global_load_fp8x16(self.v_base_iter, src_elem, token < self.seqlen_kv_v, is_value=True) + token = fx.Int64(tile_start) + (self.wave_id_uni % fx.Int64(4)) * fx.Int64(16) + d_col = (self.wave_id // fx.Int64(4)) * fx.Int64(64) + self.lane_in_warp + page_id = self.load_page_id(token) + src_div = self.make_page_view(self.v_base_iter, page_id, is_value=True) + src_elem = self.kv_head_idx * traits.PAGE_SIZE * traits.HEAD_DIM_V + d_col * traits.KV_VEC_SIZE + valid = token < self.seqlen_kv_v + src_elem = valid.select(src_elem, fx.Int64(self.v_page_bytes)) + source = self.buffer_load_fp8x16(src_div, src_elem) + if const_expr(traits.FP8_PV_SEGMENTED): + tail_elem = (valid & (d_col + traits.FP8_V_H1 < traits.HEAD_DIM_V)).select( + src_elem + fx.Int64(traits.FP8_V_H1 * traits.KV_VEC_SIZE), fx.Int64(self.v_page_bytes) + ) + tail = self.buffer_load_fp8x16(src_div, tail_elem) + source = Vec(source).shuffle(Vec(tail), [0, 1, 2, 3, 4, 5, 6, 7]).ir_value() + return source + + def _load_v_fp8_page1_source(self, tile_start): + """Load native token-major vectors and transpose each 16-lane byte tile.""" + traits = self.traits + token = fx.Int64(tile_start) + self.lane_in_warp + d_base = self.wave_id * fx.Int64(16) + page_id = self.load_page_id(token, uniform=False) + byte_offset = page_id * self.v_page_bytes + self.kv_head_idx * traits.HEAD_DIM_V + d_base + valid = token < self.seqlen_kv_v + prefix = self.global_load_fp8x16(self.v_base_iter, byte_offset, valid, is_value=True) + source = _transpose_v_fp8_16x16(prefix, self.lane_in_warp) + if const_expr(traits.FP8_PV_SEGMENTED): + tail = self.global_load_fp8x16( + self.v_base_iter, + byte_offset + traits.FP8_V_H1, + valid & (d_base + traits.FP8_V_H1 < traits.HEAD_DIM_V), + is_value=True, + ) + tail_source = _transpose_v_fp8_16x16(tail, self.lane_in_warp) + source = source.shuffle(tail_source, [0, 1, 2, 3, 4, 5, 6, 7]) + return source.ir_value() + def _store_v_fp8_vectorized_bankpad(self, src_i32x4, buf_id): """Permute prefetched V vectors into the paged FP8 LDS layout.""" + if const_expr(self.traits.PAGE_SIZE == 16 and self.traits.HEAD_DIM == 192): + self._store_v_fp8_page16(src_i32x4, buf_id) + return if const_expr(self.traits.FP8_PV_SEGMENTED): source = Vec(src_i32x4) prefix = source.shuffle(source, [0, 1, 2, 3]).ir_value() @@ -5307,6 +5514,36 @@ def _store_tail(): else: self._store_v_fp8_vectorized_bankpad_segment(src_i32x4, buf_id, 0) + def _store_v_fp8_page16(self, source, buf_id): + """Scatter native page vectors into the retained MFMA byte order.""" + traits = self.traits + words = Vec(source) + page = self.wave_id % fx.Int64(4) + group_offset = (page % fx.Int64(2)) * fx.Int64(4) + (page // fx.Int64(2)) * fx.Int64(16) + aligned_base = ((self.lds_vt_base_idx + fx.Int64(127)) // fx.Int64(128)) * fx.Int64(128) + tile_base = aligned_base - self.lds_vt_base_idx + buf_id * (traits.HEAD_DIM_V * traits.FP8_V_ROW_STRIDE) + d_col = (self.wave_id // fx.Int64(4)) * fx.Int64(64) + self.lane_in_warp + + def _store_segment(first_word, d_offset): + row = d_col + fx.Int64(d_offset) + row_base = tile_base + row * traits.FP8_V_ROW_STRIDE + for word in range_constexpr(4): + token_offset = group_offset + (word % 2) * 32 + (word // 2) * 8 + if const_expr(traits.HEAD_DIM_V == 192): + token_offset = token_offset ^ (((row // fx.Int64(4)) % fx.Int64(4)) * fx.Int64(16)) + ptr = fx.add_offset(fx.get_iter(self.v_lds_i32_tiles), fx.Uint32(row_base + token_offset) // 4) + fx.generic_store(ptr, fx.Int32(words[first_word + word])) + + _store_segment(0, 0) + if const_expr(traits.FP8_PV_SEGMENTED): + + @flyc.jit + def _store_tail_pages(): + if self.wave_id_uni < fx.Int64(4): + _store_segment(4, traits.FP8_V_H1) + + _store_tail_pages() + def _store_v_fp8_vectorized_bankpad_segment(self, src_i32x4, buf_id, d_offset): traits = self.traits row_stride = traits.FP8_V_ROW_STRIDE @@ -5344,8 +5581,18 @@ def _read_vec_key(key): packs = [] for ws in range_constexpr(traits.HEAD_DIM // 64): dg0 = ws * 4 + self.lane_div_32 * 2 - off0 = (k_base + (dg0 * traits.BLOCK_N + key_sigma) * traits.KV_VEC_SIZE) * traits.ELEM_BYTES - off1 = (k_base + ((dg0 + 1) * traits.BLOCK_N + key_sigma) * traits.KV_VEC_SIZE) * traits.ELEM_BYTES + if const_expr(traits.PAGE_SIZE == 16 and traits.HEAD_DIM_V == 192): + # [D-group/4, page, D-group%4, token%16, byte]. The + # offset delta from page-64 layout is divisible by 256, + # preserving the gfx950 LDS bank index. + oct_idx = (ws * 4 + key_sigma // 16) * 64 + self.lane_div_32 * 32 + key_sigma % 16 + off0 = k_base + oct_idx * traits.KV_VEC_SIZE + off1 = off0 + 16 * traits.KV_VEC_SIZE + else: + off0 = (k_base + (dg0 * traits.BLOCK_N + key_sigma) * traits.KV_VEC_SIZE) * traits.ELEM_BYTES + off1 = ( + k_base + ((dg0 + 1) * traits.BLOCK_N + key_sigma) * traits.KV_VEC_SIZE + ) * traits.ELEM_BYTES lo = Vec(self.read_i32x4_lds(off0), (4,), fx.Int32) hi = Vec(self.read_i32x4_lds(off1), (4,), fx.Int32) packs.append(lo.shuffle(hi, [0, 1, 2, 3, 4, 5, 6, 7]).ir_value()) diff --git a/tests/kernels/test_flash_attn_fwd.py b/tests/kernels/test_flash_attn_fwd.py index c2a5d55a48..4ced2237f3 100644 --- a/tests/kernels/test_flash_attn_fwd.py +++ b/tests/kernels/test_flash_attn_fwd.py @@ -4559,25 +4559,39 @@ def test_return_lse_false_returns_only_out(): assert out.shape == q.shape -def _paged_fp8_torch_reference(query, key, value, block_table, query_lengths, kv_lengths, descales): - """Independent PyTorch reference for packed page-64 FP8 inputs.""" +def _paged_fp8_cache_shapes(num_pages, num_kv_heads, head_dim, value_head_dim, page_size, kv_cache_layout): + if kv_cache_layout == "vectorized": + return ( + (num_pages, num_kv_heads, head_dim // 16, page_size, 16), + (num_pages, num_kv_heads, page_size // 16, value_head_dim, 16), + ) + if kv_cache_layout == "linear3d": + return (num_pages, num_kv_heads, head_dim), (num_pages, num_kv_heads, value_head_dim) + return (num_pages, page_size, num_kv_heads, head_dim), (num_pages, page_size, num_kv_heads, value_head_dim) + + +def _paged_fp8_torch_reference( + query, key, value, block_table, query_lengths, kv_lengths, descales, kv_cache_layout="vectorized" +): + """Materialize logical tokens independently of the kernel's paged loaders.""" q_scale, k_scale, v_scale = descales head_dim = query.shape[-1] - num_kv_heads = key.shape[1] - value_head_dim = value.shape[3] + vectorized = kv_cache_layout == "vectorized" + num_kv_heads = key.shape[1] if vectorized else key.shape[-2] + value_head_dim = value.shape[3] if vectorized else value.shape[-1] expected = [] q_offset = 0 for batch_idx, (query_length, kv_length) in enumerate(zip(query_lengths, kv_lengths)): query_batch = query[q_offset : q_offset + query_length].float() * q_scale q_offset += query_length physical_pages = block_table[batch_idx].long() - key_batch = ( - key[physical_pages].permute(0, 3, 1, 2, 4).reshape(-1, num_kv_heads, head_dim)[:kv_length].float() * k_scale - ) - value_batch = ( - value[physical_pages].permute(0, 2, 4, 1, 3).reshape(-1, num_kv_heads, value_head_dim)[:kv_length].float() - * v_scale - ) + key_pages = key[physical_pages] + value_pages = value[physical_pages] + if vectorized: + key_pages = key_pages.permute(0, 3, 1, 2, 4) + value_pages = value_pages.permute(0, 2, 4, 1, 3) + key_batch = key_pages.reshape(-1, num_kv_heads, head_dim)[:kv_length].float() * k_scale + value_batch = value_pages.reshape(-1, num_kv_heads, value_head_dim)[:kv_length].float() * v_scale result = pytorch_ref_attention_qkv_diff(query_batch[None], key_batch[None], value_batch[None]) expected.append(result[0].to(torch.bfloat16)) return torch.cat(expected) @@ -4649,8 +4663,10 @@ def _check_paged_fp8_matches_torch( block_table_rows, num_kv_heads=1, lazy_rescale=True, + page_size=64, + kv_cache_layout="vectorized", ): - """Packed causal FP8 page-64 attention supports native Q/K and V widths.""" + """Packed causal FP8 attention supports native Q/K and V widths.""" if len(query_lengths) == 1 and force_internal_copies and head_dim == 192: pytest.skip("D192 wrapper-copy stream ordering is covered by the ragged cases") torch.manual_seed(17) @@ -4666,12 +4682,11 @@ def _check_paged_fp8_matches_torch( num_pages = max(max(row) for row in block_table_rows) + 1 query, query_descale = quantize_per_tensor_fp8(torch.randn(sum(query_lengths), 16, head_dim, device="cuda") * 0.2) - key, key_descale = quantize_per_tensor_fp8( - torch.randn(num_pages, num_kv_heads, head_dim // 16, 64, 16, device="cuda") * 0.2 - ) - value, value_descale = quantize_per_tensor_fp8( - torch.randn(num_pages, num_kv_heads, 4, value_head_dim, 16, device="cuda") * 0.2 + key_shape, value_shape = _paged_fp8_cache_shapes( + num_pages, num_kv_heads, head_dim, value_head_dim, page_size, kv_cache_layout ) + key, key_descale = quantize_per_tensor_fp8(torch.randn(key_shape, device="cuda") * 0.2) + value, value_descale = quantize_per_tensor_fp8(torch.randn(value_shape, device="cuda") * 0.2) if force_internal_copies: @@ -4699,7 +4714,7 @@ def noncontiguous_copy(tensor): max_seqlen_kv=max(kv_lengths), cross_seqlen=True, seqlen_k=seqlen_k, - kv_cache_layout="vectorized", + kv_cache_layout=kv_cache_layout, q_descale=query_descale, k_descale=key_descale, v_descale=value_descale, @@ -4761,10 +4776,266 @@ def noncontiguous_copy(tensor): query_lengths, kv_lengths, (query_descale, key_descale, value_descale), + kv_cache_layout=kv_cache_layout, ) assert actual.shape == (sum(query_lengths), 16, value_head_dim) assert bool(torch.isfinite(actual).all().item()) torch.testing.assert_close(actual, expected, rtol=2.0e-2, atol=2.0e-2) + return actual + + +_PAGED_FP8_PHYSICAL_LAYOUTS = [(1, "linear"), (1, "linear3d"), (16, "vectorized"), (1024, "vectorized")] + + +@_requires_gfx950 +@pytest.mark.parametrize("page_size,kv_cache_layout", [(1, "linear"), (1, "linear3d"), (16, "vectorized")]) +@pytest.mark.parametrize("head_dim,value_head_dim", [(128, 128), (192, 128), (192, 192)]) +def test_paged_fp8_cache_buffer_matches_wide_fallback( + monkeypatch, page_size, kv_cache_layout, head_dim, value_head_dim +): + kv_lengths = [257, 81] + counts = [(length + page_size - 1) // page_size for length in kv_lengths] + pages = list(reversed(range(sum(counts)))) + kwargs = dict( + head_dim=head_dim, + value_head_dim=value_head_dim, + use_non_default_stream=False, + force_internal_copies=False, + query_lengths=[65, 17], + kv_lengths=kv_lengths, + block_table_rows=[pages[: counts[0]], pages[counts[0] :] + [0] * (counts[0] - counts[1])], + num_kv_heads=2, + page_size=page_size, + kv_cache_layout=kv_cache_layout, + ) + buffered = _check_paged_fp8_matches_torch(**kwargs) + build = flash_attn_interface._build_paged_fp8 + calls = [] + + def build_wide(**options): + assert options["cache_buffered"], "small caches should select bounded buffer descriptors" + options["cache_buffered"] = False + calls.append(options) + return build(**options) + + monkeypatch.setattr(flash_attn_interface, "_build_paged_fp8", build_wide) + wide = _check_paged_fp8_matches_torch(**kwargs) + assert calls + torch.testing.assert_close(buffered, wide, rtol=0, atol=0) + + +@_requires_gfx950 +@pytest.mark.parametrize("page_size", [1, 16]) +@pytest.mark.parametrize("oversized", ["key", "value"]) +def test_paged_fp8_cache_buffer_rejects_oversized_direct_launch(page_size, oversized): + from kernels.attention.flash_attn_fp8_paged_gfx950 import build_flash_attn_paged_fp8_module + from kernels.attention.flash_attn_utils import PAGED_FP8_BUFFER_LIMIT_BYTES + + class MetadataOnlyTensor: + def __init__(self, num_bytes): + self.num_bytes = num_bytes + + def numel(self): + return self.num_bytes + + def element_size(self): + return 1 + + launch = build_flash_attn_paged_fp8_module( + num_heads=16, + num_kv_heads=1, + head_dim=128, + value_head_dim=128, + dtype_str="fp8", + causal=True, + varlen=True, + cross_seqlen=True, + paged=True, + page_size=page_size, + kv_cache_layout="linear3d" if page_size == 1 else "vectorized", + cache_buffered=True, + ) + key = MetadataOnlyTensor(PAGED_FP8_BUFFER_LIMIT_BYTES + 16 if oversized == "key" else 128) + value = MetadataOnlyTensor(PAGED_FP8_BUFFER_LIMIT_BYTES + 16 if oversized == "value" else 128) + with pytest.raises(ValueError, match="buffer descriptor exceeds its byte limit"): + launch(None, key, value, None, batch_size=1, seq_len=1) + + +@_requires_gfx950 +@pytest.mark.parametrize("active_rows", [1, 15, 16, 17, 63, 64]) +def test_paged_fp8_page1_transpose_is_byte_exact(active_rows): + import flydsl.compiler as flyc + import flydsl.expr as fx + from kernels.attention.flash_attn_utils import _transpose_v_fp8_16x16 + + @flyc.kernel + def transpose_bytes(source: fx.Tensor, output: fx.Tensor): + lane = fx.Int32(fx.thread_idx.x) + values = fx.generic_load(fx.add_offset(fx.get_iter(source), lane * 4), dtype=fx.Int32, count=4) + transposed = _transpose_v_fp8_16x16(values, lane) + fx.generic_store(fx.add_offset(fx.get_iter(output), lane * 4), transposed) + + @flyc.jit + def launch(source: fx.Tensor, output: fx.Tensor): + transpose_bytes(source, output).launch(grid=(1, 1, 1), block=(64, 1, 1)) + + torch.manual_seed(47) + source = torch.randint(0, 256, (64, 16), device="cuda", dtype=torch.uint8) + source[active_rows:].zero_() + output = torch.empty_like(source) + launch(source.view(torch.int32), output.view(torch.int32)) + torch.cuda.synchronize() + expected = source.reshape(4, 16, 16).transpose(1, 2) + torch.testing.assert_close(output.reshape(4, 16, 16), expected, rtol=0, atol=0) + + +@_requires_gfx950 +@pytest.mark.parametrize( + "page_size,kv_cache_layout,error_type", + [(size, "linear", NotImplementedError) for size in (0, 2, 8)] + + [(size, "vectorized", NotImplementedError) for size in (32, 128, 256, 2048)] + + [(size, "linear", NotImplementedError) for size in (16, 64, 1024)] + + [(1, "vectorized", ValueError)], +) +def test_paged_fp8_rejects_unsupported_page_layout(page_size, kv_cache_layout, error_type): + query = torch.zeros((1, 16, 128), dtype=FP8_DTYPE, device="cuda") + key_shape, value_shape = _paged_fp8_cache_shapes(1, 1, 128, 128, page_size, kv_cache_layout) + key = torch.zeros(key_shape, dtype=FP8_DTYPE, device="cuda") + value = torch.zeros(value_shape, dtype=FP8_DTYPE, device="cuda") + indptr = torch.tensor([0, 1], dtype=torch.int32, device="cuda") + scale = torch.ones(1, dtype=torch.float32, device="cuda") + with pytest.raises(error_type, match="page|paged"): + flydsl_flash_attn_func( + query, + key, + value, + causal=True, + num_kv_heads=1, + cu_seqlens_q=indptr, + cu_seqlens_kv=indptr, + max_seqlen_q=1, + max_seqlen_kv=1, + cross_seqlen=True, + block_table=torch.zeros((1, 1), dtype=torch.int32, device="cuda"), + seqlen_k=torch.ones(1, dtype=torch.int32, device="cuda"), + kv_cache_layout=kv_cache_layout, + q_descale=scale, + k_descale=scale, + v_descale=scale, + ) + + +@_requires_gfx950 +@pytest.mark.parametrize("violation", ["page_count", "head_count", "table_rows", "table_width"]) +def test_paged_fp8_rejects_inconsistent_page_metadata(violation): + query = torch.zeros((1, 16, 128), dtype=FP8_DTYPE, device="cuda") + key = torch.zeros((5, 1, 8, 16, 16), dtype=FP8_DTYPE, device="cuda") + value = torch.zeros((4 if violation == "page_count" else 5, 1, 1, 128, 16), dtype=FP8_DTYPE, device="cuda") + table_shape = (2 if violation == "table_rows" else 1, 4 if violation == "table_width" else 5) + scale = torch.ones(1, dtype=torch.float32, device="cuda") + with pytest.raises(ValueError, match="page|paged"): + flydsl_flash_attn_func( + query, + key, + value, + causal=True, + num_kv_heads=2 if violation == "head_count" else 1, + cu_seqlens_q=torch.tensor([0, 1], dtype=torch.int32, device="cuda"), + cu_seqlens_kv=torch.tensor([0, 65], dtype=torch.int32, device="cuda"), + max_seqlen_q=1, + max_seqlen_kv=65, + cross_seqlen=True, + block_table=torch.zeros(table_shape, dtype=torch.int32, device="cuda"), + seqlen_k=torch.tensor([65], dtype=torch.int32, device="cuda"), + kv_cache_layout="vectorized", + q_descale=scale, + k_descale=scale, + v_descale=scale, + ) + + +@_requires_gfx950 +@pytest.mark.parametrize("page_size,kv_cache_layout", _PAGED_FP8_PHYSICAL_LAYOUTS) +@pytest.mark.parametrize("head_dim,value_head_dim", [(128, 128), (192, 128), (192, 192)]) +def test_paged_fp8_physical_pages_side_stream_copies(page_size, kv_cache_layout, head_dim, value_head_dim): + kv_lengths = [page_size + 65, page_size + 17, 33] + counts = [(length + page_size - 1) // page_size for length in kv_lengths] + pages = list(reversed(range(sum(counts)))) + rows = [] + offset = 0 + for count in counts: + rows.append(pages[offset : offset + count] + [0] * (max(counts) - count)) + offset += count + _check_paged_fp8_matches_torch( + head_dim=head_dim, + value_head_dim=value_head_dim, + use_non_default_stream=True, + force_internal_copies=True, + query_lengths=[65, 17, 1], + kv_lengths=kv_lengths, + block_table_rows=rows, + page_size=page_size, + kv_cache_layout=kv_cache_layout, + ) + + +@_requires_gfx950 +@pytest.mark.parametrize("page_size,kv_cache_layout", _PAGED_FP8_PHYSICAL_LAYOUTS) +@pytest.mark.parametrize("head_dim,value_head_dim", [(128, 128), (192, 128), (192, 192)]) +@pytest.mark.parametrize("num_kv_heads", [1, 2, 4]) +@pytest.mark.parametrize("lazy_rescale", [True, False]) +def test_paged_fp8_physical_page_size_ragged_matches_torch( + page_size, kv_cache_layout, head_dim, value_head_dim, num_kv_heads, lazy_rescale +): + query_lengths = [300, 257, 33, 7, 1] + kv_lengths = [max(513, 2 * page_size + 1), max(385, page_size + 1), 65, 16, 1] + counts = [(length + page_size - 1) // page_size for length in kv_lengths] + physical_pages = list(reversed(range(sum(counts)))) + rows = [] + offset = 0 + for count in counts: + rows.append(physical_pages[offset : offset + count] + [0] * (max(counts) - count)) + offset += count + _check_paged_fp8_matches_torch( + head_dim=head_dim, + value_head_dim=value_head_dim, + use_non_default_stream=True, + force_internal_copies=False, + query_lengths=query_lengths, + kv_lengths=kv_lengths, + block_table_rows=rows, + num_kv_heads=num_kv_heads, + lazy_rescale=lazy_rescale, + page_size=page_size, + kv_cache_layout=kv_cache_layout, + ) + + +@_requires_gfx950 +@pytest.mark.parametrize("head_dim,value_head_dim", [(128, 128), (192, 128), (192, 192)]) +@pytest.mark.parametrize( + "page_size,kv_cache_layout,kv_length", + [ + (size, layout, length) + for size, layout in _PAGED_FP8_PHYSICAL_LAYOUTS + for length in sorted({63, 64, 65, max(1, size - 1), size, size + 1}) + ], +) +def test_paged_fp8_physical_page_boundary_matches_torch( + head_dim, value_head_dim, page_size, kv_cache_layout, kv_length +): + pages = (kv_length + page_size - 1) // page_size + _check_paged_fp8_matches_torch( + head_dim=head_dim, + value_head_dim=value_head_dim, + use_non_default_stream=False, + force_internal_copies=False, + query_lengths=[min(17, kv_length)], + kv_lengths=[kv_length], + block_table_rows=[list(reversed(range(pages)))], + page_size=page_size, + kv_cache_layout=kv_cache_layout, + ) @_requires_gfx950 @@ -5035,19 +5306,19 @@ def test_paged_fp8_long_context_random_pages_matches_torch(head_dim, value_head_ [(128, 128), (192, 128), (192, 192)], ids=["d128-v128", "d192-v128", "d192-v192"], ) -def test_paged_fp8_graph_replay_matches_torch(head_dim, value_head_dim): +@pytest.mark.parametrize("page_size,kv_cache_layout", [(64, "vectorized")] + _PAGED_FP8_PHYSICAL_LAYOUTS) +def test_paged_fp8_graph_replay_matches_torch(head_dim, value_head_dim, page_size, kv_cache_layout): """The optimized paged variants preserve caller-owned graph output.""" torch.manual_seed(29) query_length = 63 - kv_length = 128 - num_pages = kv_length // 64 + kv_length = max(128, page_size + 1) + num_pages = (kv_length + page_size - 1) // page_size query, query_descale = quantize_per_tensor_fp8(torch.randn(query_length, 16, head_dim, device="cuda") * 0.2) - key, key_descale = quantize_per_tensor_fp8(torch.randn(num_pages, 1, head_dim // 16, 64, 16, device="cuda") * 0.2) - value, value_descale = quantize_per_tensor_fp8( - torch.randn(num_pages, 1, 4, value_head_dim, 16, device="cuda") * 0.2 - ) + key_shape, value_shape = _paged_fp8_cache_shapes(num_pages, 1, head_dim, value_head_dim, page_size, kv_cache_layout) + key, key_descale = quantize_per_tensor_fp8(torch.randn(key_shape, device="cuda") * 0.2) + value, value_descale = quantize_per_tensor_fp8(torch.randn(value_shape, device="cuda") * 0.2) output = torch.empty(query_length, 16, value_head_dim, device="cuda", dtype=torch.bfloat16) - block_table = torch.tensor([[1, 0]], device="cuda", dtype=torch.int32) + block_table = torch.arange(num_pages - 1, -1, -1, device="cuda", dtype=torch.int32).reshape(1, -1) q_indptr = torch.tensor([0, query_length], device="cuda", dtype=torch.int32) kv_indptr = torch.tensor([0, kv_length], device="cuda", dtype=torch.int32) seqlen_k = torch.tensor([kv_length], device="cuda", dtype=torch.int32) @@ -5061,7 +5332,7 @@ def test_paged_fp8_graph_replay_matches_torch(head_dim, value_head_dim): cross_seqlen=True, block_table=block_table, seqlen_k=seqlen_k, - kv_cache_layout="vectorized", + kv_cache_layout=kv_cache_layout, q_descale=query_descale, k_descale=key_descale, v_descale=value_descale, @@ -5087,6 +5358,7 @@ def test_paged_fp8_graph_replay_matches_torch(head_dim, value_head_dim): [query_length], [kv_length], (query_descale, key_descale, value_descale), + kv_cache_layout=kv_cache_layout, ) assert captured.data_ptr() == output.data_ptr() @@ -5097,7 +5369,10 @@ def test_paged_fp8_graph_replay_matches_torch(head_dim, value_head_dim): @_requires_gfx950 @pytest.mark.parametrize("head_dim,value_head_dim", [(128, 128), (192, 128), (192, 192)]) @pytest.mark.parametrize("kv_length", [128, 192]) -def test_paged_fp8_explicit_compile_matches_torch(monkeypatch, head_dim, value_head_dim, kv_length): +@pytest.mark.parametrize("page_size,kv_cache_layout", [(64, "vectorized")] + _PAGED_FP8_PHYSICAL_LAYOUTS) +def test_paged_fp8_explicit_compile_matches_torch( + monkeypatch, head_dim, value_head_dim, kv_length, page_size, kv_cache_layout +): """Explicit compilation honors the launch arguments for both paged schedules.""" build = flash_attn_interface._build_paged_fp8 @@ -5105,7 +5380,7 @@ def compile_launcher(**kwargs): return build(**kwargs).compile monkeypatch.setattr(flash_attn_interface, "_build_paged_fp8", compile_launcher) - pages = kv_length // 64 + pages = (kv_length + page_size - 1) // page_size _check_paged_fp8_matches_torch( head_dim=head_dim, value_head_dim=value_head_dim, @@ -5114,6 +5389,8 @@ def compile_launcher(**kwargs): query_lengths=[64, 32], kv_lengths=[kv_length, kv_length - 32], block_table_rows=[list(reversed(range(pages))), list(range(pages, 2 * pages))], + page_size=page_size, + kv_cache_layout=kv_cache_layout, ) @@ -5266,11 +5543,10 @@ def test_paged_legacy_dtype_accepts_host_only_seqlen_metadata(dtype): @_requires_gfx950 @pytest.mark.large_shape -def test_paged_fp8_d192_cache_offsets_above_4gib(): +@pytest.mark.parametrize("head_dim,value_head_dim", [(128, 128), (192, 128), (192, 192)]) +@pytest.mark.parametrize("page_size,kv_cache_layout", [(64, "vectorized")] + _PAGED_FP8_PHYSICAL_LAYOUTS) +def test_paged_fp8_cache_offsets_above_4gib(head_dim, value_head_dim, page_size, kv_cache_layout): """Physical K and V page rebasing remains 64-bit beyond 4 GiB.""" - head_dim = 192 - value_head_dim = 128 - page_size = 64 key_page_bytes = page_size * head_dim value_page_bytes = page_size * value_head_dim high_page = math.ceil(2**32 / min(key_page_bytes, value_page_bytes)) @@ -5286,16 +5562,9 @@ def test_paged_fp8_d192_cache_offsets_above_4gib(): ) query, query_descale = quantize_per_tensor_fp8(torch.randn(1, 16, head_dim, device="cuda") * 0.2) - key = torch.empty(num_pages, 1, head_dim // 16, page_size, 16, dtype=FP8_DTYPE, device="cuda") - value = torch.empty( - num_pages, - 1, - page_size // 16, - value_head_dim, - 16, - dtype=FP8_DTYPE, - device="cuda", - ) + key_shape, value_shape = _paged_fp8_cache_shapes(num_pages, 1, head_dim, value_head_dim, page_size, kv_cache_layout) + key = torch.empty(key_shape, dtype=FP8_DTYPE, device="cuda") + value = torch.empty(value_shape, dtype=FP8_DTYPE, device="cuda") key_page, key_descale = quantize_per_tensor_fp8(torch.randn_like(key[high_page], dtype=torch.float32) * 0.2) value_page, value_descale = quantize_per_tensor_fp8(torch.randn_like(value[high_page], dtype=torch.float32) * 0.2) key[high_page].copy_(key_page) @@ -5317,14 +5586,17 @@ def test_paged_fp8_d192_cache_offsets_above_4gib(): cross_seqlen=True, block_table=block_table, seqlen_k=seqlen_k, - kv_cache_layout="vectorized", + kv_cache_layout=kv_cache_layout, q_descale=query_descale, k_descale=key_descale, v_descale=value_descale, ) torch.cuda.synchronize() - expected = (value_page[0, 0, :, 0].float() * value_descale).expand(16, -1).unsqueeze(0).to(torch.bfloat16) + first_value = ( + value_page[0, 0, :, 0] if kv_cache_layout == "vectorized" else value_page.reshape(-1, value_head_dim)[0] + ) + expected = (first_value.float() * value_descale).expand(16, -1).unsqueeze(0).to(torch.bfloat16) assert high_page * key_page_bytes >= 2**32 assert high_page * value_page_bytes >= 2**32 assert bool(torch.isfinite(actual).all().item()) From 5196a225e13d8a472afe3f6de6a939a33d4bde24 Mon Sep 17 00:00:00 2001 From: Xiake Sun Date: Fri, 11 Sep 2026 15:54:58 +0000 Subject: [PATCH 56/59] perf(fmha): streamline small-page FP8 data movement - Use DPP byte transposes and width-specific LDS word placement for native page-1 attention. - Reuse small-page K/V translations and bounded cache resources while preserving the 64-bit fallback. - Shorten page-16 fragment lifetimes and stage V192 segments with uniform scalar K offsets. - Simplify unused launch plumbing and supported V-width traits while retaining official FlyDSL APIs. Signed-off-by: Xiake Sun --- .../attention/flash_attn_fp8_paged_gfx950.py | 24 +- kernels/attention/flash_attn_utils.py | 177 +++++++++++---- tests/kernels/test_flash_attn_fwd.py | 212 +++++++++++++++++- 3 files changed, 345 insertions(+), 68 deletions(-) diff --git a/kernels/attention/flash_attn_fp8_paged_gfx950.py b/kernels/attention/flash_attn_fp8_paged_gfx950.py index c231c06ef4..c9a378294c 100644 --- a/kernels/attention/flash_attn_fp8_paged_gfx950.py +++ b/kernels/attention/flash_attn_fp8_paged_gfx950.py @@ -17,7 +17,6 @@ DualwaveFp8StoreHelper, _make_paged_dualwave_swp_fp8_traits, ) -from kernels.common.kernels_common import dtype_to_elem_type from kernels.common.tensor_shim import _run_compiled @@ -128,11 +127,10 @@ def build_flash_attn_paged_fp8_module( DEFAULT_STRIDE_O_N = traits.NUM_HEADS_Q * traits.HEAD_DIM_V DEFAULT_STRIDE_KV_N = traits.DEFAULT_STRIDE_KV_N _dualwave_swp_fp8_cache_tag = traits.cache_tag - _lds_elem_dtype = dtype_to_elem_type(traits.DTYPE_STR) @fx.struct class SharedStorage: - kv: fx.Array[_lds_elem_dtype, traits.LDS_KV_TOTAL_SIZE, 16] + kv: fx.Array[fx.Float8E4M3FN, traits.LDS_KV_TOTAL_SIZE, 16] vt: fx.Array[fx.BFloat16, traits.VT_BF16_TOTAL, 16] # BN128: two BLOCK_N=64 KV tiles per iteration, one merged softmax correction. @@ -205,6 +203,7 @@ def flash_attn_paged_fp8_bn128_kernel( BOUNDED_MAX = ( traits.HEAD_DIM == 128 and traits.DUALWAVE_SWP_LAZY_RESCALE and not traits.DUALWAVE_SWP_DEBUG_LAZY_COUNTS ) + STREAM_PAGE16_V128 = traits.PAGE_SIZE == 16 and traits.HEAD_DIM == 192 and traits.HEAD_DIM_V == 128 t0 = ctx.split_t0 t_end = ctx.split_t_end @@ -340,14 +339,17 @@ def _iterate(j, loop_args, do_mask, skip_max=False, initialize=False): next_v_b = loop_args[next_v_arg_idx + 1] v_k_a = kv_lds_to_regs.load_k(a_buf) - v_k_b = kv_lds_to_regs.load_k(b_buf) - v_v_a = kv_lds_to_regs.load_v(a_buf) + if const_expr(not STREAM_PAGE16_V128): + v_k_b = kv_lds_to_regs.load_k(b_buf) + v_v_a = kv_lds_to_regs.load_v(a_buf) page_f_a, page_f_b = ctx.load_page_id_pair((j + 4) * BN) kv_gmem_to_lds.load_k((j + 4) * BN, f_a_buf, page_id=page_f_a) kv_gmem_to_lds.load_k((j + 5) * BN, f_b_buf, page_id=page_f_b) v_s_a = gemm_helper.qk(v_k_a, q_wide) + if const_expr(STREAM_PAGE16_V128): + v_k_b = kv_lds_to_regs.load_k(b_buf) kv_gmem_to_lds._store_v_fp8_vectorized_bankpad(next_v_a, nn_a_buf) v_f_a = kv_gmem_to_lds._load_v_fp8_vectorized_bankpad_source( (j + 4) * BN, @@ -359,6 +361,10 @@ def _iterate(j, loop_args, do_mask, skip_max=False, initialize=False): ) v_s_b = gemm_helper.qk(v_k_b, q_wide) kv_gmem_to_lds._store_v_fp8_vectorized_bankpad(next_v_b, nn_b_buf) + if const_expr(STREAM_PAGE16_V128): + # The current V slot is distinct from both next-V stores. + # Delay its fragment until K operands have been consumed. + v_v_a = kv_lds_to_regs.load_v(a_buf) if const_expr(do_mask): v_s_a, v_s_b = softmax_helper.causal_mask_pair_if_needed(v_s_a, v_s_b, j) m_new = m_row @@ -445,7 +451,6 @@ def launch_flash_attn_dualwave_swp( seq_len_kv: fx.Int32, stride_q_n: fx.Int32, stride_o_n: fx.Int32, - stride_kv_n: fx.Int32, head_dim_runtime: fx.Int32, stream: fx.Stream = fx.Stream(None), ): @@ -499,7 +504,7 @@ def launch_flash_attn_dualwave_swp( stream=stream, ) - _dualwave_swp_compile_hints = { + launch_flash_attn_dualwave_swp.compile_hints = { "fast_fp_math": True, "unsafe_fp_math": True, "llvm_options": { @@ -508,7 +513,6 @@ def launch_flash_attn_dualwave_swp( "disable-machine-sink": True, }, } - launch_flash_attn_dualwave_swp.compile_hints = dict(_dualwave_swp_compile_hints) def _validate_paged_bn128_launch(batch_size, seq_len_kv, block_table_stride): if not PAIRED_PAGE_IDS: @@ -565,8 +569,7 @@ def _launch( and max(K.numel() * K.element_size(), V.numel() * V.element_size()) > PAGED_FP8_BUFFER_LIMIT_BYTES ): raise ValueError("paged FP8 whole-cache buffer descriptor exceeds its byte limit") - if stride_kv_n is None: - stride_kv_n = DEFAULT_STRIDE_KV_N + # stride_kv_n is accepted for compatibility; native cache layouts fix it. if stride_q_n is None: stride_q_n = DEFAULT_STRIDE_Q_N if stride_o_n is None: @@ -619,7 +622,6 @@ def _launch( seq_len_kv, stride_q_n, stride_o_n, - stride_kv_n, head_dim_runtime, fx.Stream(stream), ) diff --git a/kernels/attention/flash_attn_utils.py b/kernels/attention/flash_attn_utils.py index c3cbbd9d6d..07e7897f10 100644 --- a/kernels/attention/flash_attn_utils.py +++ b/kernels/attention/flash_attn_utils.py @@ -2168,17 +2168,6 @@ def cache_tag(self): ) -def _factor_fp8_pv_head_dim(value_head_dim, segment_capacity, alignment): - """Split a vector width into at most two independently staged segments.""" - if value_head_dim <= 0 or value_head_dim % alignment != 0: - raise ValueError(f"FP8 P*V head dim must be a positive multiple of {alignment}, got {value_head_dim}") - h1 = min(1 << (value_head_dim.bit_length() - 1), segment_capacity) - h2 = value_head_dim - h1 - if h2 > segment_capacity: - raise ValueError(f"FP8 P*V head dim {value_head_dim} exceeds two segment capacities of {segment_capacity}") - return h1, h2 - - def _make_paged_dualwave_swp_fp8_traits( num_heads, num_kv_heads, @@ -2207,7 +2196,8 @@ def _make_paged_dualwave_swp_fp8_traits( smem_k_tile_elems = block_n * head_dim num_prefetch_k = 8 if head_dim == 128 else 6 slot_elems = smem_k_tile_elems - fp8_v_h1, fp8_v_h2 = _factor_fp8_pv_head_dim(value_head_dim, 128, 32) + # The builder accepts V128/V192: one 128-channel prefix and an optional V64 tail. + fp8_v_h1, fp8_v_h2 = 128, value_head_dim - 128 fp8_v_row_stride = block_n if compact_v192 else block_n + 16 fp8_v_tile_bytes = value_head_dim * fp8_v_row_stride @@ -4950,8 +4940,14 @@ def load_page_id(self, tile_start, *, uniform=True): def load_page_id_pair(self, tile_start): """Return reusable page IDs for a compute-tile pair, or defer to producers.""" - if const_expr(self.traits.PAGE_SIZE == 16 and self.traits.HEAD_DIM == 128 and self.traits.CACHE_BUFFERED): - # K and V use the same four lane-group pages. Share their metadata + if const_expr(self.traits.PAGE_SIZE == 1 and self.traits.HEAD_DIM_V == 192): + token = fx.Int64(tile_start) + self.lane_in_warp + return self.load_page_id(token, uniform=False), self.load_page_id(token + 64, uniform=False) + if const_expr(self.traits.PAGE_SIZE == 16 and self.traits.CACHE_BUFFERED): + if const_expr(self.traits.HEAD_DIM_V == 192): + token = fx.Int64(tile_start) + (self.wave_id_uni % fx.Int64(4)) * fx.Int64(16) + return self.load_page_id(token), self.load_page_id(token + 64) + # V128 K and V use the same four lane-group pages. Share metadata # loads across the two producers instead of looking them up twice. token = fx.Int64(tile_start) + (self.lane_in_warp // fx.Int64(16)) * fx.Int64(16) return self.load_page_id(token, uniform=False), self.load_page_id(token + 64, uniform=False) @@ -5154,28 +5150,41 @@ def pv(self, v_p, v_v, v_o): return self._pv_fp8_direct(v_p, v_v, v_o) -def _transpose_v_fp8_16x16(source, lane): - """Exchange row/column bits within each 16-lane group; full EXEC is required.""" +def _page1_k_page_ids(page_ids): + """Convert token-major page IDs to K's sigma lanes without another lookup.""" + # Swap lane bits 2/3: move lanes 4..7 forward and 8..11 backward by four. + value = fx.Int32(page_ids).ir_value() + swapped = rocdl.update_dpp(T.i32, value, value, 0x104, 15, 2, False) + result = rocdl.update_dpp(T.i32, swapped, value, 0x114, 15, 4, False) + return fx.Int64(fx.Int32(result)) + + +def _transpose_v_fp8_16x16(source, lane, stages=4): + """Transpose byte quads (two stages) or 16-lane tiles (four), with full EXEC.""" + assert stages in (2, 4) words = [fx.Uint32(Vec(source)[i]) for i in range_constexpr(4)] lane = fx.Int32(lane) - for bit in range_constexpr(4): - peer_addr = (lane ^ (1 << bit)) * 4 + for bit in range_constexpr(stages): lane_bit = (lane & (1 << bit)) != 0 + if const_expr(bit < 2): + # Select alternating bytes, then halfwords, from self/peer words. + low_selector = 0x06020400 if bit == 0 else 0x05040100 + high_selector = 0x03070105 if bit == 0 else 0x03020706 + selector = lane_bit.select(fx.Int32(high_selector), fx.Int32(low_selector)) updated = [] for word in range_constexpr(4): peer_word = word if bit < 2 else word ^ (1 << (bit - 2)) - peer = fx.Uint32(rocdl.ds_bpermute(T.i32, peer_addr, words[peer_word])) + peer_value = words[peer_word].ir_value() if const_expr(bit < 2): - shift = 8 << bit - low_mask = fx.Uint32(0x00FF00FF if bit == 0 else 0x0000FFFF) - high_mask = fx.Uint32(0xFFFFFFFF) ^ low_mask - low_row = (words[word] & low_mask) | ((peer & low_mask) << shift) - high_row = (words[word] & high_mask) | ((peer & high_mask) >> shift) - updated.append(lane_bit.select(high_row, low_row)) - elif const_expr(word & (1 << (bit - 2))): - updated.append(lane_bit.select(words[word], peer)) + dpp_ctrl = 0xB1 if bit == 0 else 0x4E + peer = rocdl.update_dpp(T.i32, peer_value, peer_value, dpp_ctrl, 15, 15, False) + updated.append(fx.Uint32(rocdl.perm_b32(peer, words[word], selector))) else: - updated.append(lane_bit.select(peer, words[word])) + peer = fx.Uint32(rocdl.ds_swizzle(T.i32, peer_value, fx.Int32(31 | (1 << (bit + 10))).ir_value())) + if const_expr(word & (1 << (bit - 2))): + updated.append(lane_bit.select(words[word], peer)) + else: + updated.append(lane_bit.select(peer, words[word])) words = updated return Vec.from_elements([fx.Int32(word) for word in words], fx.Int32) @@ -5192,12 +5201,15 @@ def load_k(self, tile_start, buf_id, page_id=None): """ traits = self.traits if const_expr(traits.PAGED and traits.PAGE_SIZE == 16 and traits.HEAD_DIM_V == 192): - self._load_k_page16(tile_start, buf_id) + self._load_k_page16(tile_start, buf_id, page_id=page_id) return eb = traits.ELEM_BYTES k_lds_byte_base = self.lds_kv_base_idx + self.k_buf_base(buf_id) * eb if const_expr(traits.KV_VECTORIZED): src_div = None + if const_expr(traits.PAGE_SIZE == 1): + if page_id is not None: + page_id = _page1_k_page_ids(page_id) if const_expr(traits.PAGE_SIZE >= traits.BLOCK_N): if page_id is None: page_id = self.load_page_id(tile_start) @@ -5283,12 +5295,20 @@ def _load_compact(): else: self._load_k_band_partial_wave(lds_addr, src_elem, tile_start, active) - def _load_k_page16(self, tile_start, buf_id): + def _load_k_page16(self, tile_start, buf_id, page_id=None): """Each wave copies four D-groups from one physical page directly to LDS.""" traits = self.traits page_in_tile = self.wave_id_uni % fx.Int64(4) - page_id = self.load_page_id(fx.Int64(tile_start) + page_in_tile * fx.Int64(16)) - src_div = self.make_page_view(self.k_base_iter, page_id) + if page_id is None: + page_id = self.load_page_id(fx.Int64(tile_start) + page_in_tile * fx.Int64(16)) + if const_expr(traits.CACHE_BUFFERED): + # The byte bound makes the uniform displacement fit i32; the + # shared resource retains the full 64-bit cache base pointer. + src_div = self.k_div + page_offset = fx.Int32(page_id * self.k_page_bytes) + else: + src_div = self.make_page_view(self.k_base_iter, page_id) + page_offset = fx.Int32(0) num_dma = traits.HEAD_DIM // traits.KV_VEC_SIZE for part in range_constexpr((num_dma + traits.NUM_WAVES - 1) // traits.NUM_WAVES): dma_id = self.wave_id_uni + part * traits.NUM_WAVES @@ -5302,7 +5322,7 @@ def _copy_page_groups(): ) dst_byte = self.lds_kv_base_idx + self.k_buf_base(buf_id) + dma_id * fx.Int64(1024) dst_byte = rocdl.readfirstlane(T.i32, fx.Int32(dst_byte).ir_value()) - self.buffer_load_lds_128(src_div, dst_byte, src_elem, 0) + self.buffer_load_lds_128(src_div, dst_byte, src_elem, page_offset) if const_expr(num_dma % traits.NUM_WAVES == 0 or part < num_dma // traits.NUM_WAVES): _copy_page_groups() @@ -5406,9 +5426,33 @@ def _permute_v_fp8_vectorized(self, src_i32x4): def _stage_v_fp8_vectorized_bankpad(self, tile_start, buf_id, page_id=None): """Stage vectorized paged V in the bank-padded direct-FP8 layout.""" + if const_expr(self.traits.PAGE_SIZE == 16 and self.traits.HEAD_DIM_V == 192 and self.traits.CACHE_BUFFERED): + self._stage_v_fp8_page16_segments(tile_start, buf_id, page_id=page_id) + return src_i32x4 = self._load_v_fp8_vectorized_bankpad_source(tile_start, page_id=page_id) self._store_v_fp8_vectorized_bankpad(src_i32x4, buf_id) + def _stage_v_fp8_page16_segments(self, tile_start, buf_id, page_id=None): + """Consume each V192 segment before loading the next segment's words.""" + traits = self.traits + token = fx.Int64(tile_start) + (self.wave_id_uni % fx.Int64(4)) * fx.Int64(16) + d_col = (self.wave_id // fx.Int64(4)) * fx.Int64(64) + self.lane_in_warp + if page_id is None: + page_id = self.load_page_id(token) + src_div = self.make_page_view(self.v_base_iter, page_id, is_value=True) + src_elem = self.kv_head_idx * traits.PAGE_SIZE * traits.HEAD_DIM_V + d_col * traits.KV_VEC_SIZE + src_elem = (token < self.seqlen_kv_v).select(src_elem, fx.Int64(self.v_page_bytes)) + prefix = self.buffer_load_fp8x16(src_div, src_elem) + self._store_v_fp8_page16(prefix, buf_id, include_tail=False) + + @flyc.jit + def _stage_tail(): + if self.wave_id_uni < fx.Int64(4): + tail = self.buffer_load_fp8x16(src_div, src_elem + fx.Int64(traits.FP8_V_H1 * traits.KV_VEC_SIZE)) + self._store_v_fp8_page16(tail, buf_id, segment_offset=traits.FP8_V_H1, include_tail=False) + + _stage_tail() + def _store_v_fp8_lds(self, data, byte_offset): dst = fx.slice(self.v_lds_i32_tiles, (None, fx.Uint32(byte_offset) // 16)) fx.generic_store(fx.get_iter(dst), data) @@ -5442,11 +5486,11 @@ def _load_v_fp8_vectorized_bankpad_source(self, tile_start, page_id=None): return source def _load_v_fp8_small_pages_source(self, tile_start, page_id=None): - """Load small-page V, grouping D192's page-16 prefetch by physical page.""" + """Load native small-page V in the selected lane- or wave-grouped layout.""" traits = self.traits if const_expr(traits.PAGE_SIZE == 1): - return self._load_v_fp8_page1_source(tile_start) - if const_expr(traits.HEAD_DIM == 128): + return self._load_v_fp8_page1_source(tile_start, page_id=page_id) + if const_expr(traits.HEAD_DIM == 128 or (traits.CACHE_BUFFERED and traits.HEAD_DIM_V == 128)): token = fx.Int64(tile_start) + (self.lane_in_warp // fx.Int64(16)) * fx.Int64(16) d_col = self.wave_id * fx.Int64(16) + self.lane_in_warp % fx.Int64(16) if page_id is None: @@ -5459,7 +5503,8 @@ def _load_v_fp8_small_pages_source(self, tile_start, page_id=None): return self.global_load_fp8x16(self.v_base_iter, src_elem, token < self.seqlen_kv_v, is_value=True) token = fx.Int64(tile_start) + (self.wave_id_uni % fx.Int64(4)) * fx.Int64(16) d_col = (self.wave_id // fx.Int64(4)) * fx.Int64(64) + self.lane_in_warp - page_id = self.load_page_id(token) + if page_id is None: + page_id = self.load_page_id(token) src_div = self.make_page_view(self.v_base_iter, page_id, is_value=True) src_elem = self.kv_head_idx * traits.PAGE_SIZE * traits.HEAD_DIM_V + d_col * traits.KV_VEC_SIZE valid = token < self.seqlen_kv_v @@ -5473,16 +5518,18 @@ def _load_v_fp8_small_pages_source(self, tile_start, page_id=None): source = Vec(source).shuffle(Vec(tail), [0, 1, 2, 3, 4, 5, 6, 7]).ir_value() return source - def _load_v_fp8_page1_source(self, tile_start): - """Load native token-major vectors and transpose each 16-lane byte tile.""" + def _load_v_fp8_page1_source(self, tile_start, page_id=None): + """Load token-major vectors for the selected full/scattered V transpose.""" traits = self.traits token = fx.Int64(tile_start) + self.lane_in_warp d_base = self.wave_id * fx.Int64(16) - page_id = self.load_page_id(token, uniform=False) + if page_id is None: + page_id = self.load_page_id(token, uniform=False) byte_offset = page_id * self.v_page_bytes + self.kv_head_idx * traits.HEAD_DIM_V + d_base valid = token < self.seqlen_kv_v prefix = self.global_load_fp8x16(self.v_base_iter, byte_offset, valid, is_value=True) - source = _transpose_v_fp8_16x16(prefix, self.lane_in_warp) + stages = 4 if traits.HEAD_DIM == 192 and traits.HEAD_DIM_V == 128 else 2 + source = _transpose_v_fp8_16x16(prefix, self.lane_in_warp, stages=stages) if const_expr(traits.FP8_PV_SEGMENTED): tail = self.global_load_fp8x16( self.v_base_iter, @@ -5490,13 +5537,20 @@ def _load_v_fp8_page1_source(self, tile_start): valid & (d_base + traits.FP8_V_H1 < traits.HEAD_DIM_V), is_value=True, ) - tail_source = _transpose_v_fp8_16x16(tail, self.lane_in_warp) + tail_source = _transpose_v_fp8_16x16(tail, self.lane_in_warp, stages=2) source = source.shuffle(tail_source, [0, 1, 2, 3, 4, 5, 6, 7]) return source.ir_value() def _store_v_fp8_vectorized_bankpad(self, src_i32x4, buf_id): """Permute prefetched V vectors into the paged FP8 LDS layout.""" - if const_expr(self.traits.PAGE_SIZE == 16 and self.traits.HEAD_DIM == 192): + if const_expr(self.traits.PAGE_SIZE == 1 and (self.traits.HEAD_DIM == 128 or self.traits.HEAD_DIM_V == 192)): + self._store_v_fp8_page1(src_i32x4, buf_id) + return + if const_expr( + self.traits.PAGE_SIZE == 16 + and self.traits.HEAD_DIM == 192 + and (self.traits.HEAD_DIM_V == 192 or not self.traits.CACHE_BUFFERED) + ): self._store_v_fp8_page16(src_i32x4, buf_id) return if const_expr(self.traits.FP8_PV_SEGMENTED): @@ -5514,7 +5568,38 @@ def _store_tail(): else: self._store_v_fp8_vectorized_bankpad_segment(src_i32x4, buf_id, 0) - def _store_v_fp8_page16(self, source, buf_id): + def _store_v_fp8_page1(self, source, buf_id): + """Complete the byte transpose by scattering four-token words to D rows.""" + traits = self.traits + words = Vec(source) + lane = self.lane_in_warp + # Retained V token-bit order is [0, 1, 4, 3, 5, 2]; low bits stay in each word. + token_word = ((lane & 16) >> 2) + ((lane & 32) >> 1) + ((lane & 4) << 3) + (lane & 8) + aligned_base = ((self.lds_vt_base_idx + fx.Int64(127)) // fx.Int64(128)) * fx.Int64(128) + tile_base = aligned_base - self.lds_vt_base_idx + buf_id * (traits.HEAD_DIM_V * traits.FP8_V_ROW_STRIDE) + d_col = self.wave_id * fx.Int64(16) + (lane & 3) + + def _store_segment(first_word, d_offset): + for word in range_constexpr(4): + row = d_col + fx.Int64(d_offset + word * 4) + token_offset = token_word + if const_expr(traits.HEAD_DIM_V == 192): + token_offset = token_offset ^ (((row // fx.Int64(4)) % fx.Int64(4)) * fx.Int64(16)) + dst = tile_base + row * traits.FP8_V_ROW_STRIDE + token_offset + ptr = fx.add_offset(fx.get_iter(self.v_lds_i32_tiles), fx.Uint32(dst) // 4) + fx.generic_store(ptr, fx.Int32(words[first_word + word])) + + _store_segment(0, 0) + if const_expr(traits.FP8_PV_SEGMENTED): + + @flyc.jit + def _store_tail(): + if self.wave_id_uni < fx.Int64(traits.FP8_V_H2 // traits.KV_VEC_SIZE): + _store_segment(4, traits.FP8_V_H1) + + _store_tail() + + def _store_v_fp8_page16(self, source, buf_id, segment_offset=0, include_tail=True): """Scatter native page vectors into the retained MFMA byte order.""" traits = self.traits words = Vec(source) @@ -5534,8 +5619,8 @@ def _store_segment(first_word, d_offset): ptr = fx.add_offset(fx.get_iter(self.v_lds_i32_tiles), fx.Uint32(row_base + token_offset) // 4) fx.generic_store(ptr, fx.Int32(words[first_word + word])) - _store_segment(0, 0) - if const_expr(traits.FP8_PV_SEGMENTED): + _store_segment(0, segment_offset) + if const_expr(traits.FP8_PV_SEGMENTED and include_tail): @flyc.jit def _store_tail_pages(): diff --git a/tests/kernels/test_flash_attn_fwd.py b/tests/kernels/test_flash_attn_fwd.py index 4ced2237f3..3157b135c3 100644 --- a/tests/kernels/test_flash_attn_fwd.py +++ b/tests/kernels/test_flash_attn_fwd.py @@ -4824,6 +4824,83 @@ def build_wide(**options): torch.testing.assert_close(buffered, wide, rtol=0, atol=0) +@_requires_gfx950 +@pytest.mark.large_shape +def test_paged_fp8_page16_scalar_offset_near_buffer_limit(monkeypatch): + """Exercise scalar K DMA offsets near 2 GiB, not just small-cache offsets.""" + from kernels.attention.flash_attn_utils import PAGED_FP8_BUFFER_LIMIT_BYTES + + page_size = 16 + head_dim = value_head_dim = 192 + page_bytes = page_size * head_dim + num_pages = PAGED_FP8_BUFFER_LIMIT_BYTES // page_bytes + cache_bytes = num_pages * page_bytes + assert cache_bytes <= PAGED_FP8_BUFFER_LIMIT_BYTES + assert (num_pages - 4) * page_bytes > 2**31 - 8 * page_bytes + torch.cuda.empty_cache() + free_bytes, _ = torch.cuda.mem_get_info() + if free_bytes < 2 * cache_bytes + 2 * 2**30: + pytest.skip("near-limit cache-offset regression requires 6 GiB of free GPU memory") + + key_shape, value_shape = _paged_fp8_cache_shapes(num_pages, 1, head_dim, value_head_dim, page_size, "vectorized") + key = torch.empty(key_shape, device="cuda", dtype=FP8_DTYPE) + value = torch.empty(value_shape, device="cuda", dtype=FP8_DTYPE) + for logical_page in range(4): + physical_page = num_pages - 4 + logical_page + key[physical_page].fill_(logical_page) + value[physical_page].fill_(logical_page + 1) + query = torch.ones((32, 16, head_dim), device="cuda", dtype=FP8_DTYPE) + block_table = torch.arange(num_pages - 4, num_pages, device="cuda", dtype=torch.int32)[None] + cu_q = torch.tensor([0, 32], device="cuda", dtype=torch.int32) + cu_kv = torch.tensor([0, 64], device="cuda", dtype=torch.int32) + seqlen_k = torch.tensor([64], device="cuda", dtype=torch.int32) + scale = torch.ones((1,), device="cuda", dtype=torch.float32) + + def call(): + return flydsl_flash_attn_func( + query, + key, + value, + causal=True, + num_kv_heads=1, + cu_seqlens_q=cu_q, + cu_seqlens_kv=cu_kv, + max_seqlen_q=32, + max_seqlen_kv=64, + cross_seqlen=True, + block_table=block_table, + seqlen_k=seqlen_k, + kv_cache_layout="vectorized", + q_descale=scale, + k_descale=scale, + v_descale=scale, + ) + + buffered = call() + build = flash_attn_interface._build_paged_fp8 + calls = [] + + def build_wide(**options): + assert options["cache_buffered"], "both near-limit caches still fit the bounded specialization" + options["cache_buffered"] = False + calls.append(options) + return build(**options) + + monkeypatch.setattr(flash_attn_interface, "_build_paged_fp8", build_wide) + wide = call() + torch.cuda.synchronize() + assert calls + torch.testing.assert_close(buffered, wide, rtol=0, atol=0) + logical_key = torch.arange(4, device="cuda", dtype=torch.float32).repeat_interleave(16) + logical_value = logical_key + 1 + expected = pytorch_ref_attention_qkv_diff( + query[None].float(), + logical_key[None, :, None, None].expand(1, 64, 1, head_dim), + logical_value[None, :, None, None].expand(1, 64, 1, value_head_dim), + )[0].to(torch.bfloat16) + torch.testing.assert_close(buffered, expected, rtol=2.0e-2, atol=2.0e-2) + + @_requires_gfx950 @pytest.mark.parametrize("page_size", [1, 16]) @pytest.mark.parametrize("oversized", ["key", "value"]) @@ -4863,7 +4940,37 @@ def element_size(self): @_requires_gfx950 @pytest.mark.parametrize("active_rows", [1, 15, 16, 17, 63, 64]) -def test_paged_fp8_page1_transpose_is_byte_exact(active_rows): +def test_paged_fp8_page1_shared_page_ids_are_byte_exact(active_rows): + import flydsl.compiler as flyc + import flydsl.expr as fx + from kernels.attention.flash_attn_utils import _page1_k_page_ids + + @flyc.kernel + def permute_ids(source: fx.Tensor, output: fx.Tensor): + lane = fx.Int32(fx.thread_idx.x) + value = fx.generic_load(fx.add_offset(fx.get_iter(source), lane), dtype=fx.Int32) + permuted = _page1_k_page_ids(value) + fx.generic_store(fx.add_offset(fx.get_iter(output), lane), fx.Int32(permuted)) + + @flyc.jit + def launch(source: fx.Tensor, output: fx.Tensor): + permute_ids(source, output).launch(grid=(1, 1, 1), block=(64, 1, 1)) + + torch.manual_seed(314) + source = torch.randint(0, 2**30, (64,), device="cuda", dtype=torch.int32) + source[active_rows:].zero_() + output = torch.empty_like(source) + launch(source, output) + torch.cuda.synchronize() + lane = torch.arange(64, device="cuda") + sigma = (lane & 3) | ((lane & 8) >> 1) | ((lane & 4) << 1) | (lane & ~15) + torch.testing.assert_close(output, source.index_select(0, sigma), rtol=0, atol=0) + + +@_requires_gfx950 +@pytest.mark.parametrize("active_rows", [1, 15, 16, 17, 63, 64]) +@pytest.mark.parametrize("stages", [2, 4]) +def test_paged_fp8_page1_transpose_is_byte_exact(active_rows, stages): import flydsl.compiler as flyc import flydsl.expr as fx from kernels.attention.flash_attn_utils import _transpose_v_fp8_16x16 @@ -4872,7 +4979,7 @@ def test_paged_fp8_page1_transpose_is_byte_exact(active_rows): def transpose_bytes(source: fx.Tensor, output: fx.Tensor): lane = fx.Int32(fx.thread_idx.x) values = fx.generic_load(fx.add_offset(fx.get_iter(source), lane * 4), dtype=fx.Int32, count=4) - transposed = _transpose_v_fp8_16x16(values, lane) + transposed = _transpose_v_fp8_16x16(values, lane, stages=stages) fx.generic_store(fx.add_offset(fx.get_iter(output), lane * 4), transposed) @flyc.jit @@ -4885,8 +4992,80 @@ def launch(source: fx.Tensor, output: fx.Tensor): output = torch.empty_like(source) launch(source.view(torch.int32), output.view(torch.int32)) torch.cuda.synchronize() - expected = source.reshape(4, 16, 16).transpose(1, 2) - torch.testing.assert_close(output.reshape(4, 16, 16), expected, rtol=0, atol=0) + if stages == 2: + expected = source.reshape(16, 4, 4, 4).permute(0, 3, 2, 1) + else: + expected = source.reshape(4, 16, 16).transpose(1, 2) + torch.testing.assert_close(output, expected.reshape(64, 16), rtol=0, atol=0) + + +@_requires_gfx950 +@pytest.mark.parametrize("value_dim", [128, 192]) +@pytest.mark.parametrize("active_rows", [1, 15, 16, 17, 63, 64]) +def test_paged_fp8_page1_word_scatter_matches_lds_layout(value_dim, active_rows): + from types import SimpleNamespace + + import flydsl.compiler as flyc + import flydsl.expr as fx + from kernels.attention.flash_attn_utils import DualwaveFp8KvGmemToLdsLoader, _transpose_v_fp8_16x16 + + stride = 80 if value_dim == 128 else 64 + + @flyc.kernel + def scatter(source: fx.Tensor, output: fx.Tensor): + tid = fx.Int32(fx.thread_idx.x) + lane = tid % 64 + wave = tid // 64 + offset = lane * (value_dim // 4) + wave * 4 + prefix = fx.generic_load(fx.add_offset(fx.get_iter(source), offset), dtype=fx.Int32, count=4) + words = _transpose_v_fp8_16x16(prefix, lane, stages=2) + if value_dim == 192: + tail_offset = (wave < 4).select(offset + 32, 0) + tail = fx.generic_load(fx.add_offset(fx.get_iter(source), tail_offset), dtype=fx.Int32, count=4) + tail_words = _transpose_v_fp8_16x16(tail, lane, stages=2) + words = words.shuffle(tail_words, [0, 1, 2, 3, 4, 5, 6, 7]) + # Use a guarded global buffer to inspect every byte/address produced by + # the same store helper used for LDS, including untouched row padding. + ctx = SimpleNamespace( + traits=SimpleNamespace( + HEAD_DIM_V=value_dim, + FP8_V_ROW_STRIDE=stride, + FP8_PV_SEGMENTED=value_dim == 192, + FP8_V_H1=128, + FP8_V_H2=value_dim - 128, + KV_VEC_SIZE=16, + ), + lane_in_warp=fx.Int64(lane), + wave_id=fx.Int64(wave), + wave_id_uni=fx.Int64(wave), + lds_vt_base_idx=fx.Int64(0), + v_lds_i32_tiles=output, + ) + DualwaveFp8KvGmemToLdsLoader._store_v_fp8_page1(ctx, words.ir_value(), fx.Int64(0)) + + @flyc.jit + def launch(source: fx.Tensor, output: fx.Tensor): + scatter(source, output).launch(grid=(1, 1, 1), block=(512, 1, 1)) + + torch.manual_seed(777) + source = torch.randint(0, 256, (64, value_dim), device="cuda", dtype=torch.uint8) + source[active_rows:].zero_() + storage = torch.full((value_dim * stride + 128,), 0xAB, dtype=torch.uint8, device="cuda") + output = storage[64:-64] + launch(source.view(torch.int32).reshape(-1), output.view(torch.int32)) + torch.cuda.synchronize() + expected = torch.full_like(storage, 0xAB) + expected_tile = expected[64:-64].reshape(value_dim, stride) + token = torch.arange(64, device="cuda") + page = token // 16 + word = (token % 16) // 4 + offset = (page % 2) * 4 + (page // 2) * 16 + (word % 2) * 32 + (word // 2) * 8 + token % 4 + rows = torch.arange(value_dim, device="cuda")[:, None] + offsets = offset[None, :].expand(value_dim, -1) + if value_dim == 192: + offsets = offsets ^ (((rows // 4) % 4) * 16) + expected_tile[rows, offsets] = source.T + torch.testing.assert_close(storage, expected, rtol=0, atol=0) @_requires_gfx950 @@ -6047,14 +6226,25 @@ def test_fp8_rescale_threshold_drops_past_the_long_sequence_bound(): assert set(f(s) for s in (1, 1024, 4096, 4097, 8192, 131072)) == {6.0, 4.0} -@pytest.mark.parametrize( - ("value_head_dim", "expected"), - [(128, (128, 0)), (160, (128, 32)), (192, (128, 64)), (224, (128, 96)), (256, (128, 128))], -) -def test_fp8_pv_head_dim_uses_two_generic_segments(value_head_dim, expected): - from kernels.attention.flash_attn_utils import _factor_fp8_pv_head_dim +@pytest.mark.parametrize("head_dim,value_head_dim", [(128, 128), (192, 128), (192, 192)]) +@pytest.mark.parametrize("page_size", [1, 16, 64, 1024]) +def test_paged_fp8_traits_match_supported_value_segments(head_dim, value_head_dim, page_size): + from kernels.attention.flash_attn_utils import LDS_BYTES_GFX950, _make_paged_dualwave_swp_fp8_traits - assert _factor_fp8_pv_head_dim(value_head_dim, segment_capacity=128, alignment=32) == expected + traits = _make_paged_dualwave_swp_fp8_traits( + 16, + 1, + head_dim, + value_head_dim, + rescale_threshold=8.0, + page_size=page_size, + kv_cache_layout="linear3d" if page_size == 1 else "vectorized", + cache_buffered=page_size in (1, 16), + ) + assert (traits.FP8_V_H1, traits.FP8_V_H2) == (128, value_head_dim - 128) + assert traits.FP8_PV_SEGMENTED == (value_head_dim == 192) + assert traits.VT_BF16_TOTAL * 2 >= traits.NUM_PREFETCH_K * value_head_dim * traits.FP8_V_ROW_STRIDE + assert traits.LDS_KV_TOTAL_SIZE + traits.VT_BF16_TOTAL * 2 <= LDS_BYTES_GFX950 @_requires_gfx950 From 3579e323f3d9327c05cf812ac4890643a8e53d4c Mon Sep 17 00:00:00 2001 From: Xiake Sun Date: Sat, 12 Sep 2026 11:22:03 +0000 Subject: [PATCH 57/59] fix(fmha): harden paged FP8 contracts and use public FMA - Honor packed cu-seqlens for every batch, including B=1, and validate launch metadata, head counts and supported public options. - Bound per-request page-table lookups and clear inactive V bytes, preserve full-packet fast path. - Normalize strided cumulative metadata and single-value descales, preserve graph replay and native page-1/16/64/1024 cache layouts. - Use public fx.fma with explicit nearest-even rounding and load the scalar V descale through fx.generic_load at the paged epilogue. - Remove unused paged debug counters and obsolete B=1/varlen plumbing; preserve lazy rescaling, non-FP8 debugging and dense split-K handling. Signed-off-by: Xiake Sun --- .../attention/flash_attn_fp8_paged_gfx950.py | 147 ++-- kernels/attention/flash_attn_interface.py | 110 ++- kernels/attention/flash_attn_utils.py | 149 ++-- tests/kernels/test_flash_attn_fwd.py | 751 +++++++++++++++++- 4 files changed, 975 insertions(+), 182 deletions(-) diff --git a/kernels/attention/flash_attn_fp8_paged_gfx950.py b/kernels/attention/flash_attn_fp8_paged_gfx950.py index c9a378294c..8c2fe196b0 100644 --- a/kernels/attention/flash_attn_fp8_paged_gfx950.py +++ b/kernels/attention/flash_attn_fp8_paged_gfx950.py @@ -41,7 +41,6 @@ def build_flash_attn_paged_fp8_module( dualwave_swp_lazy_rescale=True, rescale_threshold=8.0, dualwave_swp_setprio=True, - dualwave_swp_debug_lazy_counts=False, dualwave_swp_enable_stagger=True, num_kv_splits=1, varlen=False, @@ -49,7 +48,6 @@ def build_flash_attn_paged_fp8_module( paged=False, kv_cache_layout="linear", paged_bn128=False, - paged_bn128_varlen=False, batch_interleave_group=1, page_size=64, cache_buffered=False, @@ -58,6 +56,10 @@ def build_flash_attn_paged_fp8_module( Priority/stagger options are accepted for caller compatibility; this paired-page pipeline uses neither wave-priority nor staggered phases. + All page-ID paths consume packed cu-seqlens, including B=1. Direct + callers supply native cache tensors, flat contiguous Q/O, int32 metadata + and one-element fp32 descales for both launch and explicit compilation; + the public interface owns shape validation, copies and empty outputs. """ gpu_arch = get_hip_arch() if value_head_dim is None: @@ -87,6 +89,8 @@ def build_flash_attn_paged_fp8_module( if num_kv_heads is None: num_kv_heads = num_heads + if num_heads <= 0 or num_kv_heads <= 0 or num_heads % num_kv_heads != 0: + raise ValueError("paged FP8 query and KV head counts must be positive, with query heads divisible by KV heads") if paged_bn128 and page_size != 64: raise ValueError("paired page-ID loads require physical page size 64") if cache_buffered and page_size not in (1, 16): @@ -94,11 +98,6 @@ def build_flash_attn_paged_fp8_module( batch_interleave_group = int(batch_interleave_group) if batch_interleave_group < 1: raise ValueError(f"batch_interleave_group must be positive, got {batch_interleave_group}") - if batch_interleave_group > 1 and ( - (head_dim == 128 and not paged_bn128) or (paged_bn128 and not paged_bn128_varlen) - ): - raise ValueError("batch interleaving requires general D192 or packed-varlen paired page IDs") - assert num_heads % num_kv_heads == 0 traits = _make_paged_dualwave_swp_fp8_traits( num_heads, num_kv_heads, @@ -108,8 +107,6 @@ def build_flash_attn_paged_fp8_module( daz=daz, dualwave_swp_lazy_rescale=dualwave_swp_lazy_rescale, rescale_threshold=rescale_threshold, - dualwave_swp_debug_lazy_counts=dualwave_swp_debug_lazy_counts, - varlen=not paged_bn128 or paged_bn128_varlen, paired_page_ids=paged_bn128, batch_interleave_group=batch_interleave_group, page_size=page_size, @@ -121,7 +118,6 @@ def build_flash_attn_paged_fp8_module( HEAD_DIM = traits.HEAD_DIM NUM_HEADS_Q = traits.NUM_HEADS_Q PAIRED_PAGE_IDS = bool(paged_bn128) - USE_CU_SEQLENS = not paged_bn128 or paged_bn128_varlen BATCH_INTERLEAVE_GROUP = traits.BATCH_INTERLEAVE_GROUP DEFAULT_STRIDE_Q_N = traits.DEFAULT_STRIDE_Q_N DEFAULT_STRIDE_O_N = traits.NUM_HEADS_Q * traits.HEAD_DIM_V @@ -140,7 +136,6 @@ def flash_attn_paged_fp8_bn128_kernel( K: fx.Tensor, V: fx.Tensor, O: fx.Tensor, # noqa: E741 - DebugCounts: fx.Tensor, CuSeqQ: fx.Tensor, CuSeqKv: fx.Tensor, BlockTable: fx.Tensor, @@ -160,17 +155,16 @@ def flash_attn_paged_fp8_bn128_kernel( K, V, O, - DebugCounts, - CuSeqQ, - CuSeqKv, - QDescale, - KDescale, - VDescale, - seq_len, - seq_len_kv, - stride_q_n, - DEFAULT_STRIDE_KV_N, - head_dim_runtime, + CuSeqQ=CuSeqQ, + CuSeqKv=CuSeqKv, + QDescale=QDescale, + KDescale=KDescale, + VDescale=VDescale, + seq_len=seq_len, + seq_len_kv=seq_len_kv, + stride_q_n=stride_q_n, + stride_kv_n=DEFAULT_STRIDE_KV_N, + head_dim_runtime=head_dim_runtime, stride_o_n=DEFAULT_STRIDE_O_N if PAIRED_PAGE_IDS else stride_o_n, BlockTable=BlockTable, block_table_stride=block_table_stride, @@ -179,12 +173,8 @@ def flash_attn_paged_fp8_bn128_kernel( ctx.init_runtime_indices() ctx.init_lds(SharedStorage) ctx.init_thread_mapping() - if const_expr(USE_CU_SEQLENS): - ctx.init_sequence_lengths() - ctx.init_varlen_causal_lpt_order() - else: - ctx.init_causal_lpt_order() - ctx.init_sequence_lengths() + ctx.init_sequence_lengths() + ctx.init_varlen_causal_lpt_order() ctx.init_descriptors() ctx.init_atoms_and_lds_ptrs() ctx.init_dma_thread_offsets() @@ -200,9 +190,7 @@ def flash_attn_paged_fp8_bn128_kernel( BN = traits.BLOCK_N D_CHUNKS = traits.D_CHUNKS NPF = traits.NUM_PREFETCH_K - BOUNDED_MAX = ( - traits.HEAD_DIM == 128 and traits.DUALWAVE_SWP_LAZY_RESCALE and not traits.DUALWAVE_SWP_DEBUG_LAZY_COUNTS - ) + BOUNDED_MAX = traits.HEAD_DIM == 128 and traits.DUALWAVE_SWP_LAZY_RESCALE STREAM_PAGE16_V128 = traits.PAGE_SIZE == 16 and traits.HEAD_DIM == 192 and traits.HEAD_DIM_V == 128 t0 = ctx.split_t0 t_end = ctx.split_t_end @@ -296,7 +284,7 @@ def _ring_wrap(x): loop_results = init_args next_v_arg_idx = 3 + D_CHUNKS - def _iterate(j, loop_args, do_mask, skip_max=False, initialize=False): + def _iterate(j, loop_args, do_mask, skip_max=False, initialize=False, mask_v=True): m_row = loop_args[0] l_row = loop_args[1] v_o = [loop_args[2 + i] for i in range_constexpr(D_CHUNKS)] @@ -317,8 +305,8 @@ def _iterate(j, loop_args, do_mask, skip_max=False, initialize=False): page_f_a, page_f_b = ctx.load_page_id_pair((j + 4) * BN) kv_gmem_to_lds.load_k((j + 4) * BN, f_a_buf, page_id=page_f_a) kv_gmem_to_lds.load_k((j + 5) * BN, f_b_buf, page_id=page_f_b) - kv_gmem_to_lds.load_v((j + 4) * BN, f_a_buf, page_id=page_f_a) - kv_gmem_to_lds.load_v((j + 5) * BN, f_b_buf, page_id=page_f_b) + kv_gmem_to_lds.load_v((j + 4) * BN, f_a_buf, page_id=page_f_a, mask_padding=mask_v) + kv_gmem_to_lds.load_v((j + 5) * BN, f_b_buf, page_id=page_f_b, mask_padding=mask_v) if const_expr(do_mask): v_s_a, v_s_b = softmax_helper.causal_mask_pair_if_needed(v_s_a, v_s_b, j) @@ -350,7 +338,7 @@ def _iterate(j, loop_args, do_mask, skip_max=False, initialize=False): v_s_a = gemm_helper.qk(v_k_a, q_wide) if const_expr(STREAM_PAGE16_V128): v_k_b = kv_lds_to_regs.load_k(b_buf) - kv_gmem_to_lds._store_v_fp8_vectorized_bankpad(next_v_a, nn_a_buf) + kv_gmem_to_lds._store_v_fp8_vectorized_bankpad(next_v_a, nn_a_buf, (j + 2) * BN, mask_padding=mask_v) v_f_a = kv_gmem_to_lds._load_v_fp8_vectorized_bankpad_source( (j + 4) * BN, page_id=page_f_a, @@ -360,7 +348,7 @@ def _iterate(j, loop_args, do_mask, skip_max=False, initialize=False): page_id=page_f_b, ) v_s_b = gemm_helper.qk(v_k_b, q_wide) - kv_gmem_to_lds._store_v_fp8_vectorized_bankpad(next_v_b, nn_b_buf) + kv_gmem_to_lds._store_v_fp8_vectorized_bankpad(next_v_b, nn_b_buf, (j + 3) * BN, mask_padding=mask_v) if const_expr(STREAM_PAGE16_V128): # The current V slot is distinct from both next-V stores. # Delay its fragment until K operands have been consumed. @@ -389,6 +377,14 @@ def _iterate(j, loop_args, do_mask, skip_max=False, initialize=False): rocdl.sched_barrier(0) return next_args + # V128 stores packets j+2/j+3; segmented V stores j+4/j+5. Keep + # byte masking out of iterations whose furthest packet is fully valid. + if const_expr(traits.PAGE_SIZE == 1): + v_prefix_end = fx.Int64(t_end) + else: + v_ahead = 4 if traits.FP8_PV_SEGMENTED else 2 + v_prefix_end = fx.max(fx.Int64(0), fx.Int64(ctx.seqlen_kv_v) // (2 * BN) * 2 - v_ahead) + if const_expr(BOUNDED_MAX): first_end = fx.min(fx.Int64(t_end), fx.Int64(t0) + 2) for j, loop_args in range(fx.Int64(t0), first_end, fx.Int64(2), init=init_args): @@ -396,30 +392,45 @@ def _iterate(j, loop_args, do_mask, skip_max=False, initialize=False): loop_results = yield next_args sealed = _query_bound_is_safe(ctx, upper_bound, fx.Float32(loop_results[0])) fast_end = sealed.select(fx.Int64(t_end), first_end) + fast_prefix_end = fx.min(fast_end, fx.max(first_end, v_prefix_end)) first_state = loop_results # Waves can choose different loops, but pair order and barrier count agree. - for j, loop_args in range(first_end, fast_end, fx.Int64(2), init=first_state): + for j, loop_args in range(first_end, fast_prefix_end, fx.Int64(2), init=first_state): + next_args = _iterate(j, loop_args, True, skip_max=True, mask_v=False) + loop_results = yield next_args + fast_tail_state = loop_results + for j, loop_args in range(fast_prefix_end, fast_end, fx.Int64(2), init=fast_tail_state): next_args = _iterate(j, loop_args, True, skip_max=True) loop_results = yield next_args slow_start = sealed.select(fx.Int64(t_end), first_end) + slow_prefix_end = fx.min(fx.Int64(t_end), fx.max(slow_start, v_prefix_end)) slow_state = loop_results - for j, loop_args in range(slow_start, fx.Int64(t_end), fx.Int64(2), init=slow_state): + for j, loop_args in range(slow_start, slow_prefix_end, fx.Int64(2), init=slow_state): + next_args = _iterate(j, loop_args, True, mask_v=False) + loop_results = yield next_args + slow_tail_state = loop_results + for j, loop_args in range(slow_prefix_end, fx.Int64(t_end), fx.Int64(2), init=slow_tail_state): next_args = _iterate(j, loop_args, True) loop_results = yield next_args elif const_expr(traits.HEAD_DIM == 192 and traits.HEAD_DIM_V == 128): # The prefix boundary is wave-uniform, not CTA-uniform. Both loops # must keep the same pair order and one rendezvous per pair. prefix_end = fx.Int64(ctx.q_start_pos_i32 + ctx.delta_i32) // (2 * BN) * 2 - prefix_end = fx.min(fx.Int64(t_end), fx.max(fx.Int64(t0), prefix_end)) + prefix_end = fx.min(fx.Int64(t_end), fx.max(fx.Int64(t0), fx.min(prefix_end, v_prefix_end))) for j, loop_args in range(fx.Int64(t0), prefix_end, fx.Int64(2), init=init_args): - next_args = _iterate(j, loop_args, False) + next_args = _iterate(j, loop_args, False, mask_v=False) loop_results = yield next_args tail_init = loop_results for j, loop_args in range(prefix_end, fx.Int64(t_end), fx.Int64(2), init=tail_init): next_args = _iterate(j, loop_args, True) loop_results = yield next_args else: - for j, loop_args in range(fx.Int64(t0), t_end, fx.Int64(2), init=init_args): + prefix_end = fx.min(fx.Int64(t_end), fx.max(fx.Int64(t0), v_prefix_end)) + for j, loop_args in range(fx.Int64(t0), prefix_end, fx.Int64(2), init=init_args): + next_args = _iterate(j, loop_args, True, mask_v=False) + loop_results = yield next_args + tail_init = loop_results + for j, loop_args in range(prefix_end, fx.Int64(t_end), fx.Int64(2), init=tail_init): next_args = _iterate(j, loop_args, True) loop_results = yield next_args m_row = loop_results[0] @@ -427,7 +438,8 @@ def _iterate(j, loop_args, do_mask, skip_max=False, initialize=False): v_o = [loop_results[2 + i] for i in range_constexpr(D_CHUNKS)] inv_l = softmax_helper.safe_l_inv(l_row) - inv_l = inv_l * ctx.vd_fp8 + value_descale = fx.generic_load(fx.get_iter(ctx.VDescale), dtype=fx.Float32, count=1) + inv_l = inv_l * value_descale softmax_helper.scale_o(v_o, inv_l) rocdl.s_barrier() output_store.store_final_o(v_o, q_row) @@ -438,7 +450,6 @@ def launch_flash_attn_dualwave_swp( K: fx.Tensor, V: fx.Tensor, O: fx.Tensor, # noqa: E741 - DebugCounts: fx.Tensor, CuSeqQ: fx.Tensor, CuSeqKv: fx.Tensor, BlockTable: fx.Tensor, @@ -480,7 +491,6 @@ def launch_flash_attn_dualwave_swp( K, V, O, - DebugCounts, CuSeqQ, CuSeqKv, BlockTable, @@ -514,18 +524,16 @@ def launch_flash_attn_dualwave_swp( }, } - def _validate_paged_bn128_launch(batch_size, seq_len_kv, block_table_stride): + def _validate_paged_bn128_launch(seq_len_kv, block_table_stride): if not PAIRED_PAGE_IDS: return - batch_size = int(batch_size) seq_len_kv = int(seq_len_kv) block_table_stride = int(block_table_stride) num_kv_pages = (seq_len_kv + traits.PAGE_SIZE - 1) // traits.PAGE_SIZE - if (not USE_CU_SEQLENS and batch_size != 1) or num_kv_pages < 2 or num_kv_pages % 2 != 0: + if num_kv_pages < 2 or num_kv_pages % 2 != 0: raise ValueError( - "paged BN128 requires batch_size=1 unless compiled for packed varlen, " - "and a positive even number " - f"of KV pages; got batch_size={batch_size}, seq_len_kv={seq_len_kv}, " + "paged BN128 requires a positive even number " + f"of KV pages; got seq_len_kv={seq_len_kv}, " f"page_size={traits.PAGE_SIZE}" ) if block_table_stride < num_kv_pages: @@ -551,7 +559,6 @@ def _launch( stride_q_n=None, stride_o_n=None, head_dim_runtime=None, - debug_counts=None, *, seq_len_kv=None, cu_seqlens_q=None, @@ -569,6 +576,19 @@ def _launch( and max(K.numel() * K.element_size(), V.numel() * V.element_size()) > PAGED_FP8_BUFFER_LIMIT_BYTES ): raise ValueError("paged FP8 whole-cache buffer descriptor exceeds its byte limit") + if ( + cu_seqlens_q is None + or cu_seqlens_kv is None + or block_table is None + or block_table_stride is None + or q_descale is None + or k_descale is None + or v_descale is None + ): + raise ValueError( + "paged FP8 flash_attn requires cu_seqlens_q, cu_seqlens_kv, block_table, " + "block_table_stride, q_descale, k_descale, and v_descale" + ) # stride_kv_n is accepted for compatibility; native cache layouts fix it. if stride_q_n is None: stride_q_n = DEFAULT_STRIDE_Q_N @@ -578,30 +598,8 @@ def _launch( head_dim_runtime = HEAD_DIM if seq_len_kv is None: seq_len_kv = seq_len - if debug_counts is None: - debug_counts = O - # Non-varlen B=1 BN128 ignores the cu_seqlens slots; use O as a placeholder - # for direct launcher calls that omit them. - if cu_seqlens_q is None: - cu_seqlens_q = O - if cu_seqlens_kv is None: - cu_seqlens_kv = O - if block_table is None: - block_table = O - if block_table_stride is None: - block_table_stride = 0 - if block_table is O and not _compile_only: - raise ValueError("paged fp8 flash_attn requires block_table") - _validate_paged_bn128_launch(batch_size, seq_len_kv, block_table_stride) + _validate_paged_bn128_launch(seq_len_kv, block_table_stride) _validate_batch_interleave_launch(batch_size) - # Direct launcher calls must supply shape-[1] fp32 descales; O keeps the - # compiled signature valid when a placeholder is needed. - if q_descale is None: - q_descale = O - if k_descale is None: - k_descale = O - if v_descale is None: - v_descale = O dispatch = flyc.compile if _compile_only else _run_compiled return dispatch( launch_flash_attn_dualwave_swp, @@ -609,7 +607,6 @@ def _launch( K, V, O, - debug_counts, cu_seqlens_q, cu_seqlens_kv, block_table, diff --git a/kernels/attention/flash_attn_interface.py b/kernels/attention/flash_attn_interface.py index 8b84312db5..a5080f20dc 100644 --- a/kernels/attention/flash_attn_interface.py +++ b/kernels/attention/flash_attn_interface.py @@ -477,7 +477,6 @@ def _build_paged_fp8( daz: bool, lazy_rescale: bool, use_bn128: bool, - paged_bn128_varlen: bool, batch_interleave_group: int, page_size: int = 64, kv_cache_layout: str = "vectorized", @@ -502,7 +501,6 @@ def _build_paged_fp8( paged=True, kv_cache_layout=kv_cache_layout, paged_bn128=use_bn128, - paged_bn128_varlen=paged_bn128_varlen, batch_interleave_group=batch_interleave_group, page_size=page_size, cache_buffered=cache_buffered, @@ -573,6 +571,7 @@ def _flydsl_flash_attn_paged( - Varlen packed Q ``[total_q, H, D]`` (cu_seqlens_q given): paged K/V looked up per kv-tile via block_table; paged split-K is not supported. """ + device = q.device if kv_cache_layout not in ("linear", "linear3d", "vectorized"): raise NotImplementedError( "flydsl_flash_attn_func: native paged KV supports kv_cache_layout in " @@ -596,7 +595,7 @@ def _flydsl_flash_attn_paged( dtype_str = _dtype_str(q) paged_fp8 = dtype_str == "fp8" - arch = _gpu_arch(q.device) + arch = _gpu_arch(device) varlen = cu_seqlens_q is not None if varlen: # Packed varlen Q: [total_q, H, D]. Per-batch ranges come from cu_seqlens @@ -611,7 +610,17 @@ def _flydsl_flash_attn_paged( raise ValueError(f"flydsl_flash_attn_func: varlen paged q must be 3D [total_q,H,D], got {q.dim()}D") _total_q, H, D = q.shape B = cu_seqlens_q.numel() - 1 + if paged_fp8: + for name, lengths in (("cu_seqlens_q", cu_seqlens_q), ("cu_seqlens_kv", cu_seqlens_kv)): + if lengths.shape != (B + 1,) or lengths.dtype != torch.int32 or lengths.device != device or B < 0: + raise ValueError( + f"flydsl_flash_attn_func: paged FP8 {name} must be int32 [B+1] on {device} " + f"with at least one boundary and matching batch counts, got " + f"shape={tuple(lengths.shape)} dtype={lengths.dtype} device={lengths.device}" + ) Sq = int(max_seqlen_q) + if paged_fp8 and Sq < 0: + raise ValueError("flydsl_flash_attn_func: paged FP8 max_seqlen_q must be nonnegative") else: if q.dim() != 4: raise ValueError(f"flydsl_flash_attn_func: paged dense q must be 4D [B,Sq,H,D], got {q.dim()}D") @@ -655,7 +664,8 @@ def _flydsl_flash_attn_paged( if k_head_dim != D: raise ValueError(f"flydsl_flash_attn_func: paged K head_dim ({k_head_dim}) must match q head_dim ({D})") if paged_fp8: - if k.shape[0] != v.shape[0]: + num_cache_pages = int(k.shape[0]) + if num_cache_pages != v.shape[0]: raise ValueError("flydsl_flash_attn_func: paged FP8 K/V must have matching physical page counts") if num_kv_heads is not None and num_kv_heads != Hkv: raise ValueError("flydsl_flash_attn_func: num_kv_heads must match the paged FP8 cache") @@ -683,9 +693,9 @@ def _flydsl_flash_attn_paged( if any(x is None for x in (q_descale, k_descale, v_descale)): raise ValueError("flydsl_flash_attn_func: paged FP8 requires q_descale, k_descale, and v_descale") for name, scale in (("q_descale", q_descale), ("k_descale", k_descale), ("v_descale", v_descale)): - if scale.device != q.device or scale.dtype != torch.float32 or scale.numel() != 1: + if scale.device != device or scale.dtype != torch.float32 or scale.numel() != 1: raise ValueError( - f"flydsl_flash_attn_func: {name} must be one float32 value on {q.device}, " + f"flydsl_flash_attn_func: {name} must be one float32 value on {device}, " f"got shape={tuple(scale.shape)} dtype={scale.dtype} device={scale.device}" ) elif D not in (64, 128) or value_head_dim != D: @@ -696,6 +706,8 @@ def _flydsl_flash_attn_paged( if num_kv_heads is None: num_kv_heads = Hkv + if H <= 0 or num_kv_heads <= 0: + raise ValueError("flydsl_flash_attn_func: paged query and KV head counts must be positive") if H % num_kv_heads != 0: raise ValueError(f"flydsl_flash_attn_func: num_heads ({H}) must be divisible by num_kv_heads ({num_kv_heads})") @@ -711,8 +723,13 @@ def _flydsl_flash_attn_paged( # Per-batch KV lengths differ in general → bottom-right cross-length masking. Varlen # paged always uses cross masking (per-batch seqlen_q/seqlen_kv come from cu_seqlens). - _kv_lens = seqlen_k.reshape(-1).tolist() if max_seqlen_kv is None or (bias is not None and not varlen) else None - skv = int(max_seqlen_kv) if max_seqlen_kv is not None else int(max(_kv_lens)) + _kv_lens = None + if max_seqlen_kv is None or (bias is not None and not varlen): + with torch.cuda.stream(stream): + _kv_lens = seqlen_k.reshape(-1).tolist() + skv = int(max_seqlen_kv) if max_seqlen_kv is not None else int(max(_kv_lens, default=0)) + if paged_fp8 and skv < 0: + raise ValueError("flydsl_flash_attn_func: paged FP8 max_seqlen_kv must be nonnegative") max_kv_pages = (skv + page_size - 1) // page_size max_pages_per_split = (max_kv_pages + int(num_kv_splits) - 1) // int(num_kv_splits) if not paged_fp8 and max_pages_per_split > _PAGED_BT_LDS_SIZE: @@ -750,14 +767,14 @@ def _flydsl_flash_attn_paged( f"flydsl_flash_attn_func: paged bias needs >= max_seqlen_kv={skv} columns, got {bias.shape[1]}" ) - if block_table.dim() != 2 or block_table.device != q.device: + if block_table.dim() != 2 or block_table.device != device: raise ValueError( - f"flydsl_flash_attn_func: block_table must be 2D on {q.device}, " + f"flydsl_flash_attn_func: block_table must be 2D on {device}, " f"got shape={tuple(block_table.shape)} device={block_table.device}" ) - if paged_fp8 and (seqlen_k.dtype != torch.int32 or seqlen_k.device != q.device or seqlen_k.numel() != B): + if paged_fp8 and (seqlen_k.dtype != torch.int32 or seqlen_k.device != device or seqlen_k.numel() != B): raise ValueError( - f"flydsl_flash_attn_func: paged FP8 seqlen_k must be int32 [{B}] on {q.device}, " + f"flydsl_flash_attn_func: paged FP8 seqlen_k must be int32 [{B}] on {device}, " f"got shape={tuple(seqlen_k.shape)} dtype={seqlen_k.dtype} device={seqlen_k.device}" ) block_table_stride = int(block_table.shape[1]) @@ -777,8 +794,8 @@ def _flydsl_flash_attn_paged( ) if out is not None: - if out.device != q.device: - raise ValueError(f"flydsl_flash_attn_func: paged output must be on {q.device}, got {out.device}") + if out.device != device: + raise ValueError(f"flydsl_flash_attn_func: paged output must be on {device}, got {out.device}") if out.shape != expected_out_shape or not out.is_contiguous(): raise ValueError( f"flydsl_flash_attn_func: paged output must be contiguous with shape {expected_out_shape}, " @@ -791,8 +808,16 @@ def _flydsl_flash_attn_paged( f"flydsl_flash_attn_func: paged output dtype must match q dtype {q.dtype}, got {out.dtype}" ) - with torch.cuda.device(q.device.index): - launch_stream = torch.cuda.current_stream(q.device) if stream is None else stream + with torch.cuda.device(device.index): + launch_stream = torch.cuda.current_stream(device) if stream is None else stream + if paged_fp8 and (q_flat_elems == 0 or skv == 0 or num_cache_pages == 0): + # No physical page is available when KV is empty. Do not enter the + # speculative K/V prefetch pipeline just to produce zero output. + empty_stream = contextlib.nullcontext() if stream is None else torch.cuda.stream(launch_stream) + with empty_stream: + if out is None: + out = torch.empty(expected_out_shape, dtype=torch.bfloat16, device=device) + return out.zero_() # Short paged attention uses generic light; unsupported cases stay on dualwave. _paged_light_ok = ( (num_kv_splits <= 1) @@ -803,7 +828,6 @@ def _flydsl_flash_attn_paged( ) if paged_fp8: use_bn128 = page_size == 64 and max_kv_pages % 2 == 0 - paged_bn128_varlen = use_bn128 and B > 1 exe = _build_paged_fp8( num_heads=H, num_kv_heads=num_kv_heads, @@ -813,7 +837,6 @@ def _flydsl_flash_attn_paged( daz=daz, lazy_rescale=dualwave_swp_lazy_rescale, use_bn128=use_bn128, - paged_bn128_varlen=paged_bn128_varlen, batch_interleave_group=( _paged_fp8_batch_interleave_group(B, fp8_head_dims, paired=use_bn128) if H == 16 and num_kv_heads == 1 @@ -872,7 +895,7 @@ def _flydsl_flash_attn_paged( ) if out is None: out_dtype = torch.bfloat16 if paged_fp8 else q.dtype - out = torch.empty(expected_out_shape, dtype=out_dtype, device=q.device) + out = torch.empty(expected_out_shape, dtype=out_dtype, device=device) # Keep serving-sized physical K/V caches in their native rank because flattening # their dynamic memref shape can exceed signed int32. The FP8 # schedule consumes Q/O as flat token-major buffers, matching its @@ -888,20 +911,20 @@ def _flydsl_flash_attn_paged( ) if paged_fp8: kwargs.update( - q_descale=q_descale, - k_descale=k_descale, - v_descale=v_descale, + q_descale=q_descale if q_descale.stride() == (1,) else q_descale.as_strided((1,), (1,)), + k_descale=k_descale if k_descale.stride() == (1,) else k_descale.as_strided((1,), (1,)), + v_descale=v_descale if v_descale.stride() == (1,) else v_descale.as_strided((1,), (1,)), ) if bias is not None: kwargs["bias"] = bias if varlen: - kwargs["cu_seqlens_q"] = cu_seqlens_q - kwargs["cu_seqlens_kv"] = cu_seqlens_kv + kwargs["cu_seqlens_q"] = cu_seqlens_q.contiguous() if paged_fp8 else cu_seqlens_q + kwargs["cu_seqlens_kv"] = cu_seqlens_kv.contiguous() if paged_fp8 else cu_seqlens_kv if cross: kwargs["seq_len_kv"] = skv if splitk: ws_elems = dualwave_splitk_workspace_elems(B, H, Sq, int(num_kv_splits), head_dim=D) - _ws = torch.empty(ws_elems, dtype=torch.float32, device=q.device) + _ws = torch.empty(ws_elems, dtype=torch.float32, device=device) kwargs["workspace"] = _ws exe(q_flat, k_flat, v_flat, o_flat, B, Sq, **kwargs) @@ -1031,12 +1054,16 @@ def flydsl_flash_attn_func( Here ``kVectorSize = 16 / element_size`` (bf16/fp16: 8, fp8: 16); page_size and head_dim must be divisible by it. causal: Bottom-right aligned causal mask when True. - num_kv_heads: KV head count for GQA/MQA; defaults to q num_heads (MHA). + num_kv_heads: KV head count for GQA/MQA; defaults to q num_heads (MHA), + or the physical cache's head count for paged KV. Both head counts + must be positive, with Q heads divisible by KV heads. cu_seqlens_q: Int32 ``[B+1]`` cumulative Q token counts (varlen). cu_seqlens_kv: Int32 ``[B+1]`` cumulative KV token counts (varlen). max_seqlen_q: Maximum per-batch Q seqlen (varlen). Required in varlen mode. max_seqlen_kv: Maximum per-batch KV seqlen (varlen cross-attn). Required when - seqlen_q != seqlen_kv per batch. + seqlen_q != seqlen_kv per batch for non-paged attention. Paged KV can + infer it from ``seqlen_k``, synchronizing the launch stream; supply + it explicitly for graph capture and to avoid that synchronization. cross_seqlen: Whether seqlen_q and seqlen_kv differ. Required in varlen mode; dense mode infers it from ``q.shape[1] != k.shape[1]``. block_table / seqlen_k: vLLM-style 2D block table metadata. Enables the @@ -1044,10 +1071,19 @@ def flydsl_flash_attn_func( ``alibi_slopes``, ``sink``, or ``return_lse``. gfx950 FP8 supports causal packed-varlen D128/V128 and D192/V128-or-V192 paths, with vectorized page sizes 16/64/1024 or linear/linear3d page 1. - BF16/F16 native paged paths require page size 64. + FP8 cu-seqlens must be int32 on Q's device; strided views are copied + on the launch stream. Prefix values must start at zero, be + nondecreasing, and describe the packed Q tokens and logical KV + lengths; ``seqlen_k`` must agree with the KV differences. Maxima + are upper bounds, not actual lengths. Empty requests are allowed. + Active page IDs must address the cache; unused table slots and + inactive cache tokens are ignored and need not be initialized. + These value invariants are caller-owned to avoid device + synchronization. BF16/F16 native paged paths require page size 64. num_kv_splits: Split-K factor (>1: gfx950 only, D=64/128, bf16/f16, seq>=384). ``None`` lets fp8 autotune it; ``1`` keeps the kernel unsplit. - fp8_block_m: Pin the fp8 tile height to 128 or 256. fp8 only. + fp8_block_m: Pin the fp8 tile height to 128 or 256. Paged FP8 supports + only its fixed 256-row tile (``None`` or ``256``). bias: Additive attention bias with the same dtype as q, folded in as ``softmax(q @ k^T * sm_scale + bias)`` -- after the scale, before the causal/padding mask. Dense: ``[Sq, Skv]``, broadcast over batch and @@ -1088,16 +1124,20 @@ def flydsl_flash_attn_func( with paged KV or fp8 -- but freely combinable with ``bias`` and ``alibi_slopes``. q_descale / k_descale / v_descale: fp32 shape-[1] descales required - for dense or paged fp8 e4m3fn inputs. + for dense or paged fp8 e4m3fn inputs. Paged FP8 also accepts scalar + and other single-element tensors. out: Optional pre-allocated output tensor. For fp8, output is bf16; otherwise it has the same dtype as q. waves_per_eu: Kernel occupancy hint. daz: Enable denormals-are-zero. dualwave_swp_lazy_rescale: Enable lazy online softmax rescale. - dualwave_swp_setprio: Enable s_setprio scheduling hints. - dualwave_swp_enable_stagger: Enable wave-group phase stagger. + dualwave_swp_setprio: Enable s_setprio scheduling hints. Accepted for + compatibility but unused by the paged FP8 pipeline. + dualwave_swp_enable_stagger: Enable wave-group phase stagger. Accepted + for compatibility but unused by the paged FP8 pipeline. debug_counts: Float32[2] tensor; when given, counts lazy-rescale branches (debug_counts[0] = all-below-true, debug_counts[1] = all-below-false). + Non-FP8 dense attention only. stream: CUDA/HIP stream to launch on. Returns: @@ -1122,6 +1162,8 @@ def flydsl_flash_attn_func( num_kv_splits = 1 if return_lse and dtype_str == "fp8": raise NotImplementedError("flydsl_flash_attn_func: return_lse is not supported for fp8") + if dtype_str == "fp8" and debug_counts is not None: + raise NotImplementedError("flydsl_flash_attn_func: fp8 flash_attn does not support debug_counts") paged_kv = any(x is not None for x in (block_table, seqlen_k)) if return_lse and paged_kv: raise NotImplementedError("flydsl_flash_attn_func: return_lse is not supported for paged KV") @@ -1205,6 +1247,8 @@ def flydsl_flash_attn_func( if sink.dim() != 1: raise ValueError(f"flydsl_flash_attn_func: sink must be 1D [H], got {sink.dim()}D") if paged_kv: + if dtype_str == "fp8" and fp8_block_m not in (None, 256): + raise NotImplementedError("flydsl_flash_attn_func: paged FP8 requires fp8_block_m=256") return _flydsl_flash_attn_paged( q, k, @@ -1236,8 +1280,6 @@ def flydsl_flash_attn_func( varlen = cu_seqlens_q is not None if dtype_str == "fp8": - if debug_counts is not None: - raise NotImplementedError("flydsl_flash_attn_func: fp8 flash_attn does not support debug_counts") if any(x is None for x in (q_descale, k_descale, v_descale)): raise ValueError("flydsl_flash_attn_func: fp8 requires q_descale, k_descale, and v_descale") for name, scale in (("q_descale", q_descale), ("k_descale", k_descale), ("v_descale", v_descale)): diff --git a/kernels/attention/flash_attn_utils.py b/kernels/attention/flash_attn_utils.py index 07e7897f10..e2893a9c25 100644 --- a/kernels/attention/flash_attn_utils.py +++ b/kernels/attention/flash_attn_utils.py @@ -396,7 +396,7 @@ def _sub_score_pair(v_s, row_max, fm_fast): return Vec.from_elements(lo_sub, fx.Float32).ir_value(), Vec.from_elements(hi_sub, fx.Float32).ir_value() -def _scale_sub_score_pair(v_s, row_max_raw, scale, zero_f, fm_fast, bias=None, constrained_fma=False): +def _scale_sub_score_pair(v_s, row_max_raw, scale, zero_f, fm_fast, bias=None, explicit_rounding=False): """Fused softmax-scale + row-max subtraction (optimization 1-A). Returns ``scale * (v_s - row_max_raw) + bias`` per element via a single FMA @@ -414,22 +414,20 @@ def _scale_sub_score_pair(v_s, row_max_raw, scale, zero_f, fm_fast, bias=None, c # exp2 lands on 2**bias * P instead of P, at no extra instruction: the # FMA's addend absorbs it. neg_scaled_max = neg_scaled_max + bias - if const_expr(constrained_fma): - # fx.fma has no constrained-rounding option. Keep this standard LLVM - # intrinsic boundary local; current codegen retains unpacked FP32 FMA. + if const_expr(explicit_rounding): + # Match the kernel descriptor's FP32 nearest-even mode through the public + # math API. Explicit rounding retains unpacked FMA in gfx950 codegen. def _center(values): centered = [] for i in range_constexpr(16): centered.append( - fx.Float32( - llvm.intr_experimental_constrained_fma( - Vec(values)[i].ir_value(), - scale.ir_value(), - neg_scaled_max.ir_value(), - ir.IntegerAttr.get(T.i64, int(llvm.RoundingMode.Dynamic)), - llvm.FPExceptionBehavior.Ignore, - ) + fx.fma( + Vec(values)[i], + scale, + neg_scaled_max, + fastmath=fx.arith.FastMathFlags.none, + roundingmode=fx.RoundingMode.to_nearest_even, ) ) return Vec.from_elements(centered, fx.Float32) @@ -2096,8 +2094,6 @@ class PagedDualwaveSwpFp8Traits: WAVES_PER_EU: int DAZ: bool DUALWAVE_SWP_LAZY_RESCALE: bool - DUALWAVE_SWP_DEBUG_LAZY_COUNTS: bool - VARLEN: bool FP8_PV_SEGMENTED: bool FP8_V_H1: int FP8_V_H2: int @@ -2112,6 +2108,7 @@ class PagedDualwaveSwpFp8Traits: VT_BF16_TOTAL: int FP8_V_ROW_STRIDE: int DUALWAVE_SWP_RESCALE_THRESHOLD: float + VARLEN: bool = True PAIRED_PAGE_IDS: bool = True BATCH_INTERLEAVE_GROUP: int = 1 BLOCK_M: int = 256 @@ -2156,7 +2153,6 @@ def cache_tag(self): self.DAZ, self.DUALWAVE_SWP_LAZY_RESCALE, self.DUALWAVE_SWP_RESCALE_THRESHOLD, - self.DUALWAVE_SWP_DEBUG_LAZY_COUNTS, self.VARLEN, self.NUM_PREFETCH_K, self.FP8_V_ROW_STRIDE, @@ -2177,8 +2173,6 @@ def _make_paged_dualwave_swp_fp8_traits( waves_per_eu=2, daz=True, dualwave_swp_lazy_rescale=True, - dualwave_swp_debug_lazy_counts=False, - varlen=False, batch_interleave_group=1, paired_page_ids=True, page_size=64, @@ -2211,9 +2205,7 @@ def _make_paged_dualwave_swp_fp8_traits( WAVES_PER_EU=waves_per_eu, DAZ=bool(daz), DUALWAVE_SWP_LAZY_RESCALE=bool(dualwave_swp_lazy_rescale), - DUALWAVE_SWP_DEBUG_LAZY_COUNTS=bool(dualwave_swp_debug_lazy_counts), DUALWAVE_SWP_RESCALE_THRESHOLD=rescale_threshold, - VARLEN=bool(varlen), FP8_PV_SEGMENTED=fp8_v_h2 > 0, FP8_V_H1=fp8_v_h1, FP8_V_H2=fp8_v_h2, @@ -4675,11 +4667,10 @@ def init_sequence_lengths(self): _cuk_div = fx.logical_divide(fx.rocdl.make_buffer_tensor(self.CuSeqKv), fx.make_layout(1, 1)) _cu_atom = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), fx.Int32) _cu_v1i32 = Vec.make_type(1, fx.Int32) - self.q_tok_base = _cu_load(_cuq_div, self.batch_idx, _cu_atom, _cu_v1i32) - self.q_tok_end = _cu_load(_cuq_div, self.batch_idx + fx.Index(1), _cu_atom, _cu_v1i32) + self.q_tok_end = _cu_load(_cuq_div, self.batch_idx + 1, _cu_atom, _cu_v1i32) self.kv_tok_base = _cu_load(_cuk_div, self.batch_idx, _cu_atom, _cu_v1i32) - self.kv_tok_end = _cu_load(_cuk_div, self.batch_idx + fx.Index(1), _cu_atom, _cu_v1i32) + self.kv_tok_end = _cu_load(_cuk_div, self.batch_idx + 1, _cu_atom, _cu_v1i32) self.seqlen_q_v = self.q_tok_end - self.q_tok_base self.seqlen_kv_v = self.kv_tok_end - self.kv_tok_base self.seqlen_kv_i32 = fx.Int32(self.seqlen_kv_v) @@ -4740,14 +4731,7 @@ def _make_buf_div(tensor, nrec_bytes): self.v_base_iter = fx.recast_iter(fx.Int8, fx.get_iter(self.V)) self.k_page_bytes = traits.PAGE_SIZE * traits.NUM_HEADS_KV * traits.HEAD_DIM self.v_page_bytes = traits.PAGE_SIZE * traits.NUM_HEADS_KV * traits.HEAD_DIM_V - self.page_indices_div = fx.logical_divide( - fx.rocdl.make_buffer_tensor(self.BlockTable), fx.make_layout(1, 1) - ) - self.page_i32_atom = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), fx.Int32) - self.page_v1i32 = Vec.make_type(1, fx.Int32) - self.page_i32x2_atom = fx.make_copy_atom(fx.rocdl.BufferCopy64b(), fx.Int32) - self.page_v2i32 = Vec.make_type(2, fx.Int32) - self.page_base = fx.Int64(self.batch_idx) * fx.Int64(self.block_table_stride) + self.init_page_table() self.k_div = None self.v_div = None if const_expr(traits.CACHE_BUFFERED): @@ -4780,19 +4764,20 @@ def init_atoms_and_lds_ptrs(self): # fp8 global->LDS DMA uses i8 destination typing; K/V LDS reads are byte-addressed. self.lds_ptr_ty = fx.PointerType.get(fx.Int8.ir_type, 2, traits.DMA_BYTES) - def init_descale(self): - def _load_scale_scalar(tensor): - _div = fx.logical_divide(fx.rocdl.make_buffer_tensor(tensor), fx.make_layout(1, 1)) - _atom = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), fx.Float32) - _v = fly.copy_atom_call_ssa([Vec.make_type(1, fx.Float32)], _atom, fx.slice(_div, (None, fx.Int32(0)))) - return fx.Float32(Vec(_v, (1,), fx.Float32)[0]) + def _load_scale_scalar(self, tensor): + _div = fx.logical_divide(fx.rocdl.make_buffer_tensor(tensor), fx.make_layout(1, 1)) + _atom = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), fx.Float32) + _v = fly.copy_atom_call_ssa([Vec.make_type(1, fx.Float32)], _atom, fx.slice(_div, (None, fx.Int32(0)))) + return Vec(_v, (1,), fx.Float32)[0] + def init_descale(self): head_dim_f32 = fx.Float32(fx.Int32(self.head_dim_runtime)) c_log2e_f = fx.Float32(_LOG2E) c_sm_scale_log2e = fx.rsqrt(head_dim_f32, fastmath=self.fm_fast) * c_log2e_f - _qd = _load_scale_scalar(self.QDescale) - _kd = _load_scale_scalar(self.KDescale) - self.vd_fp8 = _load_scale_scalar(self.VDescale) + _qd = self._load_scale_scalar(self.QDescale) + _kd = self._load_scale_scalar(self.KDescale) + if const_expr(not self.traits.PAGED): + self.vd_fp8 = self._load_scale_scalar(self.VDescale) # fp8 feeds raw Q/K into the MFMA, so q/k descale * softmax scale multiplies # the fp32 logits after QK. self.c_logit_scale = c_sm_scale_log2e * (_qd * _kd) @@ -4916,24 +4901,38 @@ def global_idx_q(self, token_idx, col): def global_idx_o(self, token_idx, col): return (self.q_tok_base + token_idx) * self.stride_o_n_v + self.q_head_idx * self.traits.HEAD_DIM_V + col + def init_page_table(self): + self.page_indices_div = fx.logical_divide(fx.rocdl.make_buffer_tensor(self.BlockTable), fx.make_layout(1, 1)) + self.page_i32_atom = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), fx.Int32) + self.page_v1i32 = Vec.make_type(1, fx.Int32) + self.page_i32x2_atom = fx.make_copy_atom(fx.rocdl.BufferCopy64b(), fx.Int32) + self.page_v2i32 = Vec.make_type(2, fx.Int32) + self.page_base = fx.Int64(self.batch_idx) * fx.Int64(self.block_table_stride) + if const_expr(self.traits.PAIRED_PAGE_IDS): + request = fx.make_view( + fx.add_offset(fx.get_iter(self.BlockTable), self.page_base), + fx.make_layout(fx.Int64(self.block_table_stride), 1), + ) + num_pages = (self.seqlen_kv_v + self.traits.PAGE_SIZE - 1) // self.traits.PAGE_SIZE + bounded = fx.rocdl.make_buffer_tensor(request, num_records_bytes=(num_pages * 4).ir_value()) + self.page_pair_indices_div = fx.logical_divide(bounded, fx.make_layout(1, 1)) + def load_page_id(self, tile_start, *, uniform=True): - """Load a logical token's physical page, optionally uniform across the wave.""" + """Load a physical page ID, zero-filling lookups past this request's pages.""" local_page = fx.Int64(tile_start) // self.traits.PAGE_SIZE - # The software pipeline rounds to an even tile count and at least four - # tiles. Redirect those masked padding tiles to entry zero instead of - # reading past the final request's page-index allocation. page_idx = self.page_base + local_page if const_expr(self.traits.PAGE_SIZE == self.traits.BLOCK_N): num_pages = self.num_kv_tiles else: num_pages = (self.seqlen_kv_v + self.traits.PAGE_SIZE - 1) // self.traits.PAGE_SIZE - safe_idx = (local_page < num_pages).select(page_idx, 0) + valid = local_page < num_pages + safe_idx = valid.select(page_idx, 0) v = fly.copy_atom_call_ssa( [self.page_v1i32], self.page_i32_atom, fx.slice(self.page_indices_div, (None, fx.Int32(safe_idx))), ) - page = Vec(v, (1,), fx.Int32)[0] + page = valid.select(Vec(v, (1,), fx.Int32)[0], fx.Int32(0)) if uniform: return fx.Int64(rocdl.readfirstlane(T.i32, page.ir_value())) return fx.Int64(page) @@ -4957,19 +4956,17 @@ def load_page_id_pair(self, tile_start): if const_expr(not self.traits.PAIRED_PAGE_IDS): return self.load_page_id(tile_start), self.load_page_id(tile_start + self.traits.BLOCK_N) local_page = fx.Int32(fx.Int64(tile_start) // self.traits.PAGE_SIZE) - page_base = fx.Int32(self.page_base) - num_kv_tiles = fx.Int32(self.num_kv_tiles) - page_idx = page_base + local_page - safe_idx = (local_page < num_kv_tiles).select(page_idx, page_base) + # gfx950 zero-fills each out-of-range dword, including an odd pair tail. v = fly.copy_atom_call_ssa( [self.page_v2i32], self.page_i32x2_atom, - fx.slice(self.page_indices_div, (None, safe_idx)), + fx.slice(self.page_pair_indices_div, (None, local_page)), ) pages = Vec(v, (2,), fx.Int32) - page0 = rocdl.readfirstlane(T.i32, pages[0].ir_value()) - page1 = rocdl.readfirstlane(T.i32, pages[1].ir_value()) - return fx.Int64(page0), fx.Int64(page1) + return ( + fx.Int64(rocdl.readfirstlane(T.i32, pages[0].ir_value())), + fx.Int64(rocdl.readfirstlane(T.i32, pages[1].ir_value())), + ) def make_page_view(self, tensor_iter, page_id, *, is_value=False): page_bytes = self.v_page_bytes if is_value else self.k_page_bytes @@ -5346,9 +5343,9 @@ def _run(): _run() - def load_v(self, tile_start, buf_id, page_id=None): + def load_v(self, tile_start, buf_id, page_id=None, *, mask_padding=True): if const_expr(self.traits.PAGED): - self._stage_v_fp8_vectorized_bankpad(tile_start, buf_id, page_id=page_id) + self._stage_v_fp8_vectorized_bankpad(tile_start, buf_id, page_id=page_id, mask_padding=mask_padding) else: self._stage_v_fp8_block_dma(tile_start, buf_id) @@ -5424,15 +5421,30 @@ def _permute_v_fp8_vectorized(self, src_i32x4): fx.Int32, ) - def _stage_v_fp8_vectorized_bankpad(self, tile_start, buf_id, page_id=None): + def _stage_v_fp8_vectorized_bankpad(self, tile_start, buf_id, page_id=None, *, mask_padding=True): """Stage vectorized paged V in the bank-padded direct-FP8 layout.""" if const_expr(self.traits.PAGE_SIZE == 16 and self.traits.HEAD_DIM_V == 192 and self.traits.CACHE_BUFFERED): - self._stage_v_fp8_page16_segments(tile_start, buf_id, page_id=page_id) + self._stage_v_fp8_page16_segments(tile_start, buf_id, page_id=page_id, mask_padding=mask_padding) return src_i32x4 = self._load_v_fp8_vectorized_bankpad_source(tile_start, page_id=page_id) - self._store_v_fp8_vectorized_bankpad(src_i32x4, buf_id) + self._store_v_fp8_vectorized_bankpad(src_i32x4, buf_id, tile_start, mask_padding=mask_padding) + + def _mask_v_fp8_group(self, source, token_start): + """Zero inactive tokens before PV: a zero probability cannot mask a NaN V.""" + packed = Vec(source) + full_words = fx.Int32(self.seqlen_kv_v) >> 2 + tail_bits = (fx.Uint32(self.seqlen_kv_v) & 3) * 8 + tail_mask = (fx.Uint32(1) << tail_bits) - 1 + group_word = fx.Int32(fx.Int64(token_start) >> 2) + words = [] + for word in range_constexpr(packed.numel): + index = group_word + word % 4 + last_word = (index == full_words).select(tail_mask, fx.Uint32(0)) + mask = (index < full_words).select(fx.Uint32(0xFFFFFFFF), last_word) + words.append(fx.Int32(fx.Uint32(packed[word]) & mask)) + return Vec.from_elements(words, fx.Int32).ir_value() - def _stage_v_fp8_page16_segments(self, tile_start, buf_id, page_id=None): + def _stage_v_fp8_page16_segments(self, tile_start, buf_id, page_id=None, *, mask_padding=True): """Consume each V192 segment before loading the next segment's words.""" traits = self.traits token = fx.Int64(tile_start) + (self.wave_id_uni % fx.Int64(4)) * fx.Int64(16) @@ -5443,12 +5455,16 @@ def _stage_v_fp8_page16_segments(self, tile_start, buf_id, page_id=None): src_elem = self.kv_head_idx * traits.PAGE_SIZE * traits.HEAD_DIM_V + d_col * traits.KV_VEC_SIZE src_elem = (token < self.seqlen_kv_v).select(src_elem, fx.Int64(self.v_page_bytes)) prefix = self.buffer_load_fp8x16(src_div, src_elem) + if const_expr(mask_padding): + prefix = self._mask_v_fp8_group(prefix, token) self._store_v_fp8_page16(prefix, buf_id, include_tail=False) @flyc.jit def _stage_tail(): if self.wave_id_uni < fx.Int64(4): tail = self.buffer_load_fp8x16(src_div, src_elem + fx.Int64(traits.FP8_V_H1 * traits.KV_VEC_SIZE)) + if const_expr(mask_padding): + tail = self._mask_v_fp8_group(tail, token) self._store_v_fp8_page16(tail, buf_id, segment_offset=traits.FP8_V_H1, include_tail=False) _stage_tail() @@ -5541,16 +5557,22 @@ def _load_v_fp8_page1_source(self, tile_start, page_id=None): source = source.shuffle(tail_source, [0, 1, 2, 3, 4, 5, 6, 7]) return source.ir_value() - def _store_v_fp8_vectorized_bankpad(self, src_i32x4, buf_id): + def _store_v_fp8_vectorized_bankpad(self, src_i32x4, buf_id, tile_start, *, mask_padding=True): """Permute prefetched V vectors into the paged FP8 LDS layout.""" if const_expr(self.traits.PAGE_SIZE == 1 and (self.traits.HEAD_DIM == 128 or self.traits.HEAD_DIM_V == 192)): self._store_v_fp8_page1(src_i32x4, buf_id) return - if const_expr( + page16_scatter = ( self.traits.PAGE_SIZE == 16 and self.traits.HEAD_DIM == 192 and (self.traits.HEAD_DIM_V == 192 or not self.traits.CACHE_BUFFERED) - ): + ) + + if const_expr(mask_padding and self.traits.PAGE_SIZE != 1): + group = self.wave_id_uni % 4 if page16_scatter else self.lane_in_warp // 16 + token = fx.Int64(tile_start) + fx.Int64(group) * 16 + src_i32x4 = self._mask_v_fp8_group(src_i32x4, token) + if const_expr(page16_scatter): self._store_v_fp8_page16(src_i32x4, buf_id) return if const_expr(self.traits.FP8_PV_SEGMENTED): @@ -5897,7 +5919,6 @@ def sub_m(self, v_s, row_max): if const_expr(self.traits.DUALWAVE_SWP_LAZY_RESCALE): headroom -= self.traits.DUALWAVE_SWP_RESCALE_THRESHOLD bias = fx.Float32(headroom) if headroom > 0.0 else None - # Batched D128 retains its tuned packed-FMA schedule. return _scale_sub_score_pair( v_s, row_max, @@ -5905,7 +5926,7 @@ def sub_m(self, v_s, row_max): self.c_zero_f, self.fm_fast, bias, - constrained_fma=self.traits.PAGED and (self.traits.HEAD_DIM != 128 or not self.traits.VARLEN), + explicit_rounding=self.traits.PAGED, ) def exp2(self, v_s, start, length): diff --git a/tests/kernels/test_flash_attn_fwd.py b/tests/kernels/test_flash_attn_fwd.py index 3157b135c3..eb924f6b79 100644 --- a/tests/kernels/test_flash_attn_fwd.py +++ b/tests/kernels/test_flash_attn_fwd.py @@ -4579,12 +4579,14 @@ def _paged_fp8_torch_reference( vectorized = kv_cache_layout == "vectorized" num_kv_heads = key.shape[1] if vectorized else key.shape[-2] value_head_dim = value.shape[3] if vectorized else value.shape[-1] + page_size = key.shape[3] if vectorized else (key.shape[1] if key.dim() == 4 else 1) expected = [] q_offset = 0 for batch_idx, (query_length, kv_length) in enumerate(zip(query_lengths, kv_lengths)): query_batch = query[q_offset : q_offset + query_length].float() * q_scale q_offset += query_length - physical_pages = block_table[batch_idx].long() + num_pages = (kv_length + page_size - 1) // page_size + physical_pages = block_table[batch_idx, :num_pages].long() key_pages = key[physical_pages] value_pages = value[physical_pages] if vectorized: @@ -4665,10 +4667,10 @@ def _check_paged_fp8_matches_torch( lazy_rescale=True, page_size=64, kv_cache_layout="vectorized", + num_query_heads=16, + strided_metadata=False, ): """Packed causal FP8 attention supports native Q/K and V widths.""" - if len(query_lengths) == 1 and force_internal_copies and head_dim == 192: - pytest.skip("D192 wrapper-copy stream ordering is covered by the ragged cases") torch.manual_seed(17) q_offsets = [0] kv_offsets = [0] @@ -4677,11 +4679,16 @@ def _check_paged_fp8_matches_torch( kv_offsets.append(kv_offsets[-1] + kv_length) q_indptr = torch.tensor(q_offsets, device="cuda", dtype=torch.int32) kv_indptr = torch.tensor(kv_offsets, device="cuda", dtype=torch.int32) + if strided_metadata: + q_indptr = torch.stack((q_indptr, torch.zeros_like(q_indptr)), dim=1)[:, 0] + kv_indptr = torch.stack((kv_indptr, torch.zeros_like(kv_indptr)), dim=1)[:, 0] block_table = torch.tensor(block_table_rows, device="cuda", dtype=torch.int32) seqlen_k = torch.tensor(kv_lengths, device="cuda", dtype=torch.int32) num_pages = max(max(row) for row in block_table_rows) + 1 - query, query_descale = quantize_per_tensor_fp8(torch.randn(sum(query_lengths), 16, head_dim, device="cuda") * 0.2) + query, query_descale = quantize_per_tensor_fp8( + torch.randn(sum(query_lengths), num_query_heads, head_dim, device="cuda") * 0.2 + ) key_shape, value_shape = _paged_fp8_cache_shapes( num_pages, num_kv_heads, head_dim, value_head_dim, page_size, kv_cache_layout ) @@ -4778,7 +4785,7 @@ def noncontiguous_copy(tensor): (query_descale, key_descale, value_descale), kv_cache_layout=kv_cache_layout, ) - assert actual.shape == (sum(query_lengths), 16, value_head_dim) + assert actual.shape == (sum(query_lengths), num_query_heads, value_head_dim) assert bool(torch.isfinite(actual).all().item()) torch.testing.assert_close(actual, expected, rtol=2.0e-2, atol=2.0e-2) return actual @@ -4787,6 +4794,630 @@ def noncontiguous_copy(tensor): _PAGED_FP8_PHYSICAL_LAYOUTS = [(1, "linear"), (1, "linear3d"), (16, "vectorized"), (1024, "vectorized")] +@_requires_gfx950 +@pytest.mark.parametrize( + "page_size,kv_cache_layout,longest_kv", + [(64, "vectorized", 193)] + [(p, layout, 129) for p, layout in [(64, "vectorized")] + _PAGED_FP8_PHYSICAL_LAYOUTS], +) +@pytest.mark.parametrize("head_dim,value_head_dim", [(128, 128), (192, 128), (192, 192)]) +@pytest.mark.parametrize("num_query_heads,num_kv_heads", [(1, 1), (3, 1), (6, 2), (6, 3), (8, 8), (24, 3), (16, 1)]) +@pytest.mark.parametrize("lazy_rescale", [False, True]) +def test_paged_fp8_general_head_counts_and_ragged_lengths( + page_size, kv_cache_layout, longest_kv, head_dim, value_head_dim, num_query_heads, num_kv_heads, lazy_rescale +): + """MHA/GQA and masked rows must not depend on the Hq16/Hkv1 benchmark.""" + query_lengths = [0, 65, 300, 17] + kv_lengths = [33, 0, longest_kv, 65] + counts = [(length + page_size - 1) // page_size for length in kv_lengths] + pages = list(reversed(range(sum(counts)))) + rows = [] + offset = 0 + for count in counts: + rows.append(pages[offset : offset + count] + [0] * (max(counts) - count)) + offset += count + _check_paged_fp8_matches_torch( + head_dim=head_dim, + value_head_dim=value_head_dim, + use_non_default_stream=False, + force_internal_copies=False, + query_lengths=query_lengths, + kv_lengths=kv_lengths, + block_table_rows=rows, + num_query_heads=num_query_heads, + num_kv_heads=num_kv_heads, + lazy_rescale=lazy_rescale, + page_size=page_size, + kv_cache_layout=kv_cache_layout, + ) + + +@_requires_gfx950 +@pytest.mark.parametrize( + "page_size,kv_cache_layout,mode", + [(64, "vectorized", "paired"), (64, "vectorized", "scalar"), (1024, "vectorized", "native")] + + [(p, layout, mode) for p, layout in _PAGED_FP8_PHYSICAL_LAYOUTS if p < 64 for mode in ("buffered", "wide")], +) +@pytest.mark.parametrize("head_dim,value_head_dim", [(128, 128), (192, 128), (192, 192)]) +@pytest.mark.parametrize("kv_length", [1, 2, 15, 16, 17, 63, 64, 65, 66, 127, 128, 129, 513, 1025]) +def test_paged_fp8_inactive_cache_values_do_not_contaminate_output( + monkeypatch, page_size, kv_cache_layout, mode, head_dim, value_head_dim, kv_length +): + """Only active K/V values matter, including empty requests and partial pages.""" + if mode == "wide": + build = flash_attn_interface._build_paged_fp8 + + def build_wide(**kwargs): + kwargs["cache_buffered"] = False + return build(**kwargs) + + monkeypatch.setattr(flash_attn_interface, "_build_paged_fp8", build_wide) + query_length, num_heads, num_kv_heads = 17, 6, 2 + active_pages = (kv_length + page_size - 1) // page_size + num_pages = active_pages + 1 + physical_ids = torch.arange(active_pages, 0, -1, device="cuda", dtype=torch.int32) + query = torch.zeros((2 * query_length, num_heads, head_dim), device="cuda", dtype=FP8_DTYPE) + key = torch.full((num_pages, page_size, num_kv_heads, head_dim), float("nan"), device="cuda") + value = torch.full((num_pages, page_size, num_kv_heads, value_head_dim), float("nan"), device="cuda") + logical_key = torch.full_like(key[1:], float("nan")) + logical_value = torch.full_like(value[1:], float("nan")) + logical_key.view(-1, num_kv_heads, head_dim)[:kv_length] = 0.0 + logical_value.view(-1, num_kv_heads, value_head_dim)[:kv_length] = 1.0 + key[physical_ids.long()] = logical_key + value[physical_ids.long()] = logical_value + if kv_cache_layout == "vectorized": + key = key.view(num_pages, page_size, num_kv_heads, head_dim // 16, 16).permute(0, 2, 3, 1, 4).contiguous() + value = ( + value.view(num_pages, page_size // 16, 16, num_kv_heads, value_head_dim).permute(0, 3, 1, 4, 2).contiguous() + ) + elif kv_cache_layout == "linear3d": + key, value = key[:, 0], value[:, 0] + key, value = key.to(FP8_DTYPE), value.to(FP8_DTYPE) + max_kv = ((active_pages | 1) * 64) if mode == "scalar" else max(512, ((kv_length + 127) // 128) * 128) + block_table = torch.full((2, (max_kv + page_size - 1) // page_size), -1, device="cuda", dtype=torch.int32) + block_table[0, :active_pages] = physical_ids + scale = torch.ones(1, device="cuda", dtype=torch.float32) + actual = flydsl_flash_attn_func( + query, + key, + value, + num_kv_heads=num_kv_heads, + cu_seqlens_q=torch.tensor([0, query_length, 2 * query_length], device="cuda", dtype=torch.int32), + cu_seqlens_kv=torch.tensor([0, kv_length, kv_length], device="cuda", dtype=torch.int32), + max_seqlen_q=query_length, + max_seqlen_kv=max_kv, + cross_seqlen=True, + block_table=block_table, + seqlen_k=torch.tensor([kv_length, 0], device="cuda", dtype=torch.int32), + kv_cache_layout=kv_cache_layout, + q_descale=scale, + k_descale=scale, + v_descale=scale, + ) + torch.cuda.synchronize() + expected = torch.zeros_like(actual) + expected[max(0, query_length - kv_length) : query_length] = 1.0 + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + +@_requires_gfx950 +@pytest.mark.parametrize("num_words", [4, 8]) +@pytest.mark.parametrize("token_base", [0, 2**31 - 128]) +def test_paged_fp8_value_mask_is_byte_exact(num_words, token_base): + """Preserve every active byte and clear all tails, including signed-int32 boundaries.""" + from types import SimpleNamespace + + import flydsl.compiler as flyc + import flydsl.expr as fx + from flydsl.expr import gpu + from kernels.attention.flash_attn_utils import DualwaveFp8KvGmemToLdsLoader + + bit_patterns = [0xFE807F01, 0x12345678, 0xFFFFFFFF, 0x01020304] * (num_words // 4) + signed_words = [word if word < 2**31 else word - 2**32 for word in bit_patterns] + + @flyc.kernel + def mask(output: fx.Tensor): + lane = fx.Int64(gpu.thread_idx.x) + group = fx.Int64(gpu.block_idx.x) + ctx = SimpleNamespace(seqlen_kv_v=fx.Int64(token_base) + lane) + source = fx.Vector.from_elements([fx.Int32(word) for word in signed_words], fx.Int32) + result = DualwaveFp8KvGmemToLdsLoader._mask_v_fp8_group(ctx, source, fx.Int64(token_base) + group * 16) + offset = (group * 128 + lane) * num_words + fx.generic_store(fx.add_offset(fx.get_iter(output), offset), fx.Vector(result)) + + @flyc.jit + def launch(output: fx.Tensor): + mask(output).launch(grid=(9, 1, 1), block=(128, 1, 1)) + + output = torch.empty((9, 128, num_words), device="cuda", dtype=torch.int32) + launch(output.view(-1)) + torch.cuda.synchronize() + expected = [ + [ + [ + word & ((1 << (min(max(length - group * 16 - (i % 4) * 4, 0), 4) * 8)) - 1) + for i, word in enumerate(bit_patterns) + ] + for length in range(128) + ] + for group in range(9) + ] + torch.testing.assert_close( + output, torch.tensor(expected, device="cuda", dtype=torch.int64).to(torch.int32), rtol=0, atol=0 + ) + + +@_requires_gfx950 +@pytest.mark.parametrize("page_size,kv_cache_layout", [(64, "vectorized")] + _PAGED_FP8_PHYSICAL_LAYOUTS) +@pytest.mark.parametrize("head_dim,value_head_dim", [(128, 128), (192, 128), (192, 192)]) +def test_paged_fp8_shared_pages_and_strided_metadata(page_size, kv_cache_layout, head_dim, value_head_dim): + """Page IDs may be shared; cumulative lengths need not be contiguous views.""" + kv_lengths = [256, 65] + counts = [(n + page_size - 1) // page_size for n in kv_lengths] + pages = list(reversed(range(max(counts)))) + _check_paged_fp8_matches_torch( + head_dim=head_dim, + value_head_dim=value_head_dim, + use_non_default_stream=True, + force_internal_copies=False, + query_lengths=[65, 17], + kv_lengths=kv_lengths, + block_table_rows=[pages, pages[: counts[1]] + [-1] * (max(counts) - counts[1])], + num_query_heads=6, + num_kv_heads=2, + page_size=page_size, + kv_cache_layout=kv_cache_layout, + strided_metadata=True, + ) + + +@_requires_gfx950 +@pytest.mark.parametrize("page_size,kv_cache_layout", [(64, "vectorized")] + _PAGED_FP8_PHYSICAL_LAYOUTS) +@pytest.mark.parametrize("head_dim,value_head_dim", [(128, 128), (192, 128), (192, 192)]) +@pytest.mark.parametrize("strided_metadata", [False, True]) +def test_paged_fp8_graph_replay_updates_lengths_and_ignores_padding( + page_size, kv_cache_layout, head_dim, value_head_dim, strided_metadata +): + """One captured graph honors changing device lengths and poisoned unused IDs.""" + torch.manual_seed(981) + query_length, max_query, max_kv = 17, 512, 512 + num_heads, num_kv_heads = 6, 2 + num_pages = (max_kv + page_size - 1) // page_size + query, q_scale = quantize_per_tensor_fp8(torch.randn(query_length, num_heads, head_dim, device="cuda") * 0.2) + key_shape, value_shape = _paged_fp8_cache_shapes( + num_pages, num_kv_heads, head_dim, value_head_dim, page_size, kv_cache_layout + ) + key, k_scale = quantize_per_tensor_fp8(torch.randn(key_shape, device="cuda") * 0.2) + value, v_scale = quantize_per_tensor_fp8(torch.randn(value_shape, device="cuda") * 0.2) + pages = torch.randperm(num_pages, dtype=torch.int32, device="cuda") + table = torch.empty((1, num_pages), dtype=torch.int32, device="cuda") + cu_q = torch.tensor([0, query_length], dtype=torch.int32, device="cuda") + cu_kv = torch.tensor([0, 65], dtype=torch.int32, device="cuda") + if strided_metadata: + cu_q = cu_q.repeat_interleave(2)[::2] + cu_kv = cu_kv.repeat_interleave(2)[::2] + seqlen_k = torch.tensor([65], dtype=torch.int32, device="cuda") + output = torch.empty((query_length, num_heads, value_head_dim), dtype=torch.bfloat16, device="cuda") + + def update_metadata(length): + table.fill_(-1) + count = (length + page_size - 1) // page_size + table[0, :count].copy_(pages[:count]) + cu_kv[1] = length + seqlen_k[0] = length + + def call(): + return flydsl_flash_attn_func( + query, + key, + value, + causal=True, + num_kv_heads=num_kv_heads, + cu_seqlens_q=cu_q, + cu_seqlens_kv=cu_kv, + max_seqlen_q=max_query, + max_seqlen_kv=max_kv, + cross_seqlen=True, + block_table=table, + seqlen_k=seqlen_k, + kv_cache_layout=kv_cache_layout, + q_descale=q_scale, + k_descale=k_scale, + v_descale=v_scale, + out=output, + ) + + update_metadata(65) + call() + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = call() + assert captured is output + for length in (65, 129, 0, 33): + update_metadata(length) + output.fill_(float("nan")) + graph.replay() + torch.cuda.synchronize() + expected = _paged_fp8_torch_reference( + query, key, value, table, [query_length], [length], (q_scale, k_scale, v_scale), kv_cache_layout + ) + assert bool(torch.isfinite(output).all()) + torch.testing.assert_close(output, expected, rtol=2.0e-2, atol=2.0e-2) + + +@_requires_gfx950 +@pytest.mark.parametrize("page_size,kv_cache_layout", [(64, "vectorized")] + _PAGED_FP8_PHYSICAL_LAYOUTS) +@pytest.mark.parametrize("head_dim,value_head_dim", [(128, 128), (192, 128), (192, 192)]) +@pytest.mark.parametrize("max_query", [256, 512]) +def test_paged_fp8_maxima_are_upper_bounds(page_size, kv_cache_layout, head_dim, value_head_dim, max_query): + """Loose launch bounds must not replace device-side lengths or overwrite O.""" + query_length, kv_length = 17, 65 + max_kv = 512 + num_heads, num_kv_heads = 3, 1 + num_pages = (max_kv + page_size - 1) // page_size + # Backing allocations cover the loose bounds: a broken specialization + # fails the canary/reference checks without an illegal GPU memory access. + query_storage = torch.zeros((max_query, num_heads, head_dim), dtype=FP8_DTYPE, device="cuda") + query = query_storage[:query_length] + output_storage = torch.full((max_query, num_heads, value_head_dim), 123.0, dtype=torch.bfloat16, device="cuda") + output = output_storage[:query_length] + key_shape, _ = _paged_fp8_cache_shapes( + num_pages, num_kv_heads, head_dim, value_head_dim, page_size, kv_cache_layout + ) + key = torch.zeros(key_shape, dtype=FP8_DTYPE, device="cuda") + logical_value = torch.full((num_pages, page_size, num_kv_heads, value_head_dim), 8.0, device="cuda") + logical_value.view(-1, num_kv_heads, value_head_dim)[:kv_length] = 1.0 + if kv_cache_layout == "vectorized": + logical_value = ( + logical_value.view(num_pages, page_size // 16, 16, num_kv_heads, value_head_dim) + .permute(0, 3, 1, 4, 2) + .contiguous() + ) + elif kv_cache_layout == "linear3d": + logical_value = logical_value[:, 0] + logical_value = logical_value.to(FP8_DTYPE) + pages = torch.arange(num_pages - 1, -1, -1, dtype=torch.int32, device="cuda") + value = torch.empty_like(logical_value) + value[pages.long()] = logical_value + scale = torch.ones(1, dtype=torch.float32, device="cuda") + actual = flydsl_flash_attn_func( + query, + key, + value, + causal=True, + num_kv_heads=num_kv_heads, + cu_seqlens_q=torch.tensor([0, query_length], dtype=torch.int32, device="cuda"), + cu_seqlens_kv=torch.tensor([0, kv_length], dtype=torch.int32, device="cuda"), + max_seqlen_q=max_query, + max_seqlen_kv=max_kv, + cross_seqlen=True, + block_table=pages[None], + seqlen_k=torch.tensor([kv_length], dtype=torch.int32, device="cuda"), + kv_cache_layout=kv_cache_layout, + q_descale=scale, + k_descale=scale, + v_descale=scale, + out=output, + ) + torch.cuda.synchronize() + assert actual.data_ptr() == output.data_ptr() + torch.testing.assert_close(output_storage[query_length:], torch.full_like(output_storage[query_length:], 123.0)) + torch.testing.assert_close(actual, torch.ones_like(actual), rtol=2.0e-2, atol=2.0e-2) + + +@_requires_gfx950 +@pytest.mark.parametrize("page_size,kv_cache_layout", [(64, "vectorized")] + _PAGED_FP8_PHYSICAL_LAYOUTS) +@pytest.mark.parametrize("head_dim,value_head_dim", [(128, 128), (192, 128), (192, 192)]) +@pytest.mark.parametrize("query_lengths,kv_lengths", [([0, 0], [65, 0]), ([17, 0, 65], [0, 0, 0]), ([], [])]) +@pytest.mark.parametrize("preallocated", [False, True]) +@pytest.mark.parametrize("kv_max_mode", ["exact", "upper_bound", "infer"]) +def test_paged_fp8_empty_attention_skips_kernel( + monkeypatch, + page_size, + kv_cache_layout, + head_dim, + value_head_dim, + query_lengths, + kv_lengths, + preallocated, + kv_max_mode, +): + """Empty queries or KV produce the empty/zero result without speculative DMA.""" + num_heads, num_kv_heads = 3, 1 + batch = len(query_lengths) + counts = [(n + page_size - 1) // page_size for n in kv_lengths] + num_pages = sum(counts) + key_shape, value_shape = _paged_fp8_cache_shapes( + num_pages, num_kv_heads, head_dim, value_head_dim, page_size, kv_cache_layout + ) + query = torch.zeros((sum(query_lengths), num_heads, head_dim), dtype=FP8_DTYPE, device="cuda") + key = torch.zeros(key_shape, dtype=FP8_DTYPE, device="cuda") + value = torch.zeros(value_shape, dtype=FP8_DTYPE, device="cuda") + q_indptr = torch.tensor([0] + query_lengths, dtype=torch.int32, device="cuda").cumsum(0, dtype=torch.int32) + kv_indptr = torch.tensor([0] + kv_lengths, dtype=torch.int32, device="cuda").cumsum(0, dtype=torch.int32) + max_kv = 512 if kv_max_mode == "upper_bound" else max(kv_lengths, default=0) + block_table = torch.zeros((batch, (max_kv + page_size - 1) // page_size), dtype=torch.int32, device="cuda") + scale = torch.ones(1, dtype=torch.float32, device="cuda") + expected = torch.zeros((sum(query_lengths), num_heads, value_head_dim), dtype=torch.bfloat16, device="cuda") + output = torch.full_like(expected, 123.0) if preallocated else None + + def unexpected_build(**kwargs): + raise AssertionError("empty attention must not build or launch a paged kernel") + + monkeypatch.setattr(flash_attn_interface, "_build_paged_fp8", unexpected_build) + actual = flydsl_flash_attn_func( + query, + key, + value, + causal=True, + num_kv_heads=num_kv_heads, + cu_seqlens_q=q_indptr, + cu_seqlens_kv=kv_indptr, + max_seqlen_q=max(query_lengths, default=0), + max_seqlen_kv=None if kv_max_mode == "infer" else max_kv, + cross_seqlen=True, + block_table=block_table, + seqlen_k=torch.tensor(kv_lengths, dtype=torch.int32, device="cuda"), + kv_cache_layout=kv_cache_layout, + q_descale=scale, + k_descale=scale, + v_descale=scale, + out=output, + ) + torch.cuda.synchronize() + if output is not None: + assert actual is output + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + +def _paged_fp8_validation_inputs(num_heads=6, num_kv_heads=2, page_size=64, kv_cache_layout="vectorized"): + query = torch.zeros((17, num_heads, 128), dtype=FP8_DTYPE, device="cuda") + num_pages = (256 + page_size - 1) // page_size + key_shape, value_shape = _paged_fp8_cache_shapes(num_pages, num_kv_heads, 128, 128, page_size, kv_cache_layout) + key = torch.zeros(key_shape, dtype=FP8_DTYPE, device="cuda") + value = torch.zeros(value_shape, dtype=FP8_DTYPE, device="cuda") + scale = torch.ones(1, dtype=torch.float32, device="cuda") + kwargs = dict( + causal=True, + num_kv_heads=num_kv_heads, + cu_seqlens_q=torch.tensor([0, 17], dtype=torch.int32, device="cuda"), + cu_seqlens_kv=torch.tensor([0, 129], dtype=torch.int32, device="cuda"), + max_seqlen_q=17, + max_seqlen_kv=256, + cross_seqlen=True, + block_table=torch.arange(num_pages, dtype=torch.int32, device="cuda")[None], + seqlen_k=torch.tensor([129], dtype=torch.int32, device="cuda"), + kv_cache_layout=kv_cache_layout, + q_descale=scale, + k_descale=scale, + v_descale=scale, + ) + return (query, key, value), kwargs + + +def _unexpected_paged_fp8_build(**kwargs): + raise AssertionError("invalid arguments reached paged kernel construction") + + +@_requires_gfx950 +@pytest.mark.parametrize("page_size,kv_cache_layout", [(64, "vectorized")] + _PAGED_FP8_PHYSICAL_LAYOUTS) +@pytest.mark.parametrize( + "scale_shape,scale_stride", [((), None), ((1,), None), ((1, 1), None), ((1,), (0,)), ((1,), (2,))] +) +def test_paged_fp8_single_value_descale_shapes(page_size, kv_cache_layout, scale_shape, scale_stride): + (query, key, value), kwargs = _paged_fp8_validation_inputs(page_size=page_size, kv_cache_layout=kv_cache_layout) + value.fill_(1.0) + for name in ("q_descale", "k_descale", "v_descale"): + kwargs[name] = torch.full(scale_shape, 0.5, dtype=torch.float32, device="cuda") + if scale_stride is not None: + kwargs[name] = kwargs[name].as_strided(scale_shape, scale_stride) + flash_attn_interface._build_paged_fp8.cache_clear() + actual = flydsl_flash_attn_func(query, key, value, **kwargs) + torch.cuda.synchronize() + torch.testing.assert_close(actual, torch.full_like(actual, 0.5), rtol=0, atol=0) + + +@_requires_gfx950 +def test_paged_fp8_infers_kv_bound_on_launch_stream(monkeypatch): + (query, key, value), kwargs = _paged_fp8_validation_inputs() + kwargs["max_seqlen_kv"] = None + value.fill_(1.0) + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + tolist = torch.Tensor.tolist + length_ptr = kwargs["seqlen_k"].data_ptr() + observed = [] + + def checked_tolist(tensor): + if tensor.is_cuda and tensor.data_ptr() == length_ptr: + observed.append(True) + assert torch.cuda.current_stream(query.device) == side + return tolist(tensor) + + monkeypatch.setattr(torch.Tensor, "tolist", checked_tolist) + actual = flydsl_flash_attn_func(query, key, value, stream=side, **kwargs) + side.synchronize() + assert observed + torch.testing.assert_close(actual, torch.ones_like(actual), rtol=0, atol=0) + + +@_requires_gfx950 +@pytest.mark.parametrize( + "name,violation", + [ + ("cu_seqlens_q", "dtype"), + ("cu_seqlens_q", "device"), + ("cu_seqlens_q", "rank"), + ("cu_seqlens_q", "empty"), + ("cu_seqlens_kv", "dtype"), + ("cu_seqlens_kv", "device"), + ("cu_seqlens_kv", "rank"), + ("cu_seqlens_kv", "length"), + ], +) +def test_paged_fp8_rejects_invalid_cumulative_metadata(monkeypatch, name, violation): + tensors, kwargs = _paged_fp8_validation_inputs() + original = kwargs[name] + if violation == "dtype": + kwargs[name] = original.to(torch.int64) + elif violation == "device": + kwargs[name] = original.cpu() + elif violation == "rank": + kwargs[name] = original[None] + elif violation == "empty": + kwargs[name] = original[:0] + else: + kwargs[name] = torch.tensor([0, 65, 129], dtype=torch.int32, device="cuda") + monkeypatch.setattr(flash_attn_interface, "_build_paged_fp8", _unexpected_paged_fp8_build) + with pytest.raises(ValueError, match="cu_seqlens"): + flydsl_flash_attn_func(*tensors, **kwargs) + + +@_requires_gfx950 +@pytest.mark.parametrize("num_heads,num_kv_heads", [(0, 1), (3, 0), (3, 2)]) +def test_paged_fp8_rejects_invalid_head_counts(monkeypatch, num_heads, num_kv_heads): + tensors, kwargs = _paged_fp8_validation_inputs(num_heads, num_kv_heads) + monkeypatch.setattr(flash_attn_interface, "_build_paged_fp8", _unexpected_paged_fp8_build) + with pytest.raises(ValueError, match="head"): + flydsl_flash_attn_func(*tensors, **kwargs) + + +@_requires_gfx950 +@pytest.mark.parametrize("option", ["debug_counts", "fp8_block_m"]) +def test_paged_fp8_rejects_unimplemented_public_options(monkeypatch, option): + tensors, kwargs = _paged_fp8_validation_inputs() + kwargs[option] = torch.zeros(2, dtype=torch.float32, device="cuda") if option == "debug_counts" else 128 + monkeypatch.setattr(flash_attn_interface, "_build_paged_fp8", _unexpected_paged_fp8_build) + with pytest.raises(NotImplementedError, match=option): + flydsl_flash_attn_func(*tensors, **kwargs) + + +@_requires_gfx950 +@pytest.mark.parametrize("name", ["max_seqlen_q", "max_seqlen_kv"]) +def test_paged_fp8_rejects_negative_launch_bounds(monkeypatch, name): + tensors, kwargs = _paged_fp8_validation_inputs() + kwargs[name] = -1 + monkeypatch.setattr(flash_attn_interface, "_build_paged_fp8", _unexpected_paged_fp8_build) + with pytest.raises(ValueError, match=name): + flydsl_flash_attn_func(*tensors, **kwargs) + + +@_requires_gfx950 +@pytest.mark.parametrize("page_size,kv_cache_layout", [(64, "vectorized")] + _PAGED_FP8_PHYSICAL_LAYOUTS) +@pytest.mark.parametrize("preallocated", [False, True]) +def test_paged_fp8_empty_cache_on_stream_and_graph(monkeypatch, page_size, kv_cache_layout, preallocated): + (query, key, value), kwargs = _paged_fp8_validation_inputs(page_size=page_size, kv_cache_layout=kv_cache_layout) + key, value = key[:0], value[:0] + kwargs["cu_seqlens_kv"].zero_() + kwargs["seqlen_k"].zero_() + output = torch.full(query.shape, 123.0, dtype=torch.bfloat16, device="cuda") if preallocated else None + monkeypatch.setattr(flash_attn_interface, "_build_paged_fp8", _unexpected_paged_fp8_build) + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + actual = flydsl_flash_attn_func(query, key, value, out=output, stream=side, **kwargs) + with torch.cuda.stream(side): + consumed = actual.clone() + side.synchronize() + torch.testing.assert_close(consumed, torch.zeros_like(consumed), rtol=0, atol=0) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=side): + captured = flydsl_flash_attn_func(query, key, value, out=output, stream=side, **kwargs) + if output is not None: + assert captured is output + captured.fill_(float("nan")) + graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close(captured, torch.zeros_like(captured), rtol=0, atol=0) + + +@_requires_gfx950 +@pytest.mark.parametrize("page_size", [1, 16, 64, 1024]) +@pytest.mark.parametrize("page_base", [0, 5]) +@pytest.mark.parametrize( + "num_pages,uniform,start_page", [(0, True, 0), (1, True, 0), (3, True, 4), (3, False, 0), (3, False, 3)] +) +def test_paged_fp8_scalar_page_ids_ignore_padding(page_size, page_base, num_pages, uniform, start_page): + """Uniform and per-lane lookups never expose unused table contents.""" + from types import SimpleNamespace + + import flydsl.compiler as flyc + import flydsl.expr as fx + from kernels.attention.flash_attn_utils import DualwaveFp8KernelContext + + kv_length = (num_pages - 1) * page_size + 1 if num_pages else 0 + + @flyc.kernel + def lookup(table: fx.Tensor, output: fx.Tensor): + lane = fx.Int64(fx.thread_idx.x) + local_page = fx.Int64(start_page) + if fx.const_expr(not uniform): + local_page = local_page + lane % 5 + ctx = SimpleNamespace( + traits=SimpleNamespace(PAGE_SIZE=page_size, BLOCK_N=64), + page_base=fx.Int64(page_base), + num_kv_tiles=fx.Int64((kv_length + 63) // 64), + seqlen_kv_v=fx.Int64(kv_length), + page_indices_div=fx.logical_divide(fx.rocdl.make_buffer_tensor(table), fx.make_layout(1, 1)), + page_i32_atom=fx.make_copy_atom(fx.rocdl.BufferCopy32b(), fx.Int32), + page_v1i32=fx.Vector.make_type(1, fx.Int32), + ) + page = DualwaveFp8KernelContext.load_page_id(ctx, local_page * page_size, uniform=uniform) + fx.generic_store(fx.add_offset(fx.get_iter(output), lane), fx.Vector.from_elements([fx.Int32(page)], fx.Int32)) + + @flyc.jit + def launch(table: fx.Tensor, output: fx.Tensor): + lookup(table, output).launch(grid=(1, 1, 1), block=(64, 1, 1)) + + table = torch.full((page_base + 8,), -1, device="cuda", dtype=torch.int32) + table[page_base : page_base + num_pages] = torch.arange(6, 6 + num_pages, device="cuda", dtype=torch.int32) + output = torch.empty(64, device="cuda", dtype=torch.int32) + launch(table, output) + torch.cuda.synchronize() + pages = torch.full_like(output, start_page) + if not uniform: + pages += torch.arange(64, device="cuda", dtype=torch.int32) % 5 + expected = torch.where(pages < num_pages, pages + 6, 0) + torch.testing.assert_close(output, expected, rtol=0, atol=0) + + +@_requires_gfx950 +@pytest.mark.parametrize("page_base", [0, 5]) +@pytest.mark.parametrize("num_tiles,tile", [(0, 0), (1, 0), (2, 0), (3, 2), (3, 4), (4, 2), (4, 4)]) +def test_paged_fp8_paired_page_ids_ignore_padding(page_base, num_tiles, tile): + """Inspect IDs without dereferencing caches, including garbage unused slots.""" + from types import SimpleNamespace + + import flydsl.compiler as flyc + import flydsl.expr as fx + from kernels.attention.flash_attn_utils import DualwaveFp8KernelContext + + @flyc.kernel + def lookup(table: fx.Tensor, output: fx.Tensor): + ctx = SimpleNamespace( + traits=SimpleNamespace(PAGE_SIZE=64, BLOCK_N=64, PAIRED_PAGE_IDS=True), + BlockTable=table, + batch_idx=fx.Int64(page_base // 5), + block_table_stride=fx.Int32(5), + seqlen_kv_v=fx.Int64(num_tiles * 64), + ) + DualwaveFp8KernelContext.init_page_table(ctx) + first, second = DualwaveFp8KernelContext.load_page_id_pair(ctx, fx.Int64(tile * 64)) + result = fx.Vector.from_elements([fx.Int32(first), fx.Int32(second)], fx.Int32) + fx.generic_store(fx.add_offset(fx.get_iter(output), fx.Int32(fx.thread_idx.x) * 2), result) + + @flyc.jit + def launch(table: fx.Tensor, output: fx.Tensor): + lookup(table, output).launch(grid=(1, 1, 1), block=(64, 1, 1)) + + table = torch.full((page_base + 5,), -1, device="cuda", dtype=torch.int32) + table[page_base : page_base + num_tiles] = torch.arange(6, 6 + num_tiles, device="cuda", dtype=torch.int32) + output = torch.empty((64, 2), device="cuda", dtype=torch.int32) + launch(table, output) + torch.cuda.synchronize() + expected = torch.tensor([6 + i if i < num_tiles else 0 for i in (tile, tile + 1)], device="cuda") + torch.testing.assert_close(output, expected.to(torch.int32).expand_as(output), rtol=0, atol=0) + + @_requires_gfx950 @pytest.mark.parametrize("page_size,kv_cache_layout", [(1, "linear"), (1, "linear3d"), (16, "vectorized")]) @pytest.mark.parametrize("head_dim,value_head_dim", [(128, 128), (192, 128), (192, 192)]) @@ -5573,9 +6204,106 @@ def compile_launcher(**kwargs): ) +@_requires_gfx950 +@pytest.mark.parametrize( + "page_size,kv_cache_layout,longest_kv", + [(64, "vectorized", 193)] + [(p, layout, 129) for p, layout in [(64, "vectorized")] + _PAGED_FP8_PHYSICAL_LAYOUTS], +) +@pytest.mark.parametrize("head_dim,value_head_dim", [(128, 128), (192, 128), (192, 192)]) +def test_paged_fp8_interleave_is_not_limited_to_benchmark_heads( + monkeypatch, page_size, kv_cache_layout, longest_kv, head_dim, value_head_dim +): + """The grid permutation is valid for generic D128 and non-benchmark GQA.""" + build = flash_attn_interface._build_paged_fp8 + + def interleaved(**kwargs): + kwargs["batch_interleave_group"] = 2 + return build(**kwargs) + + monkeypatch.setattr(flash_attn_interface, "_build_paged_fp8", interleaved) + kv_lengths = [33, 0, longest_kv, 65] + counts = [(n + page_size - 1) // page_size for n in kv_lengths] + pages = list(reversed(range(max(counts)))) + _check_paged_fp8_matches_torch( + head_dim=head_dim, + value_head_dim=value_head_dim, + use_non_default_stream=True, + force_internal_copies=False, + query_lengths=[0, 65, 300, 17], + kv_lengths=kv_lengths, + block_table_rows=[pages[:n] + [-1] * (max(counts) - n) for n in counts], + num_query_heads=6, + num_kv_heads=2, + page_size=page_size, + kv_cache_layout=kv_cache_layout, + ) + + +@_requires_gfx950 +@pytest.mark.parametrize("page_size,kv_cache_layout", [(64, "vectorized")] + _PAGED_FP8_PHYSICAL_LAYOUTS) +@pytest.mark.parametrize("compile_only", [False, True]) +@pytest.mark.parametrize( + "missing", + ["cu_seqlens_q", "cu_seqlens_kv", "block_table", "block_table_stride", "q_descale", "k_descale", "v_descale"], +) +def test_paged_fp8_direct_launcher_requires_metadata(monkeypatch, page_size, kv_cache_layout, compile_only, missing): + from kernels.attention import flash_attn_fp8_paged_gfx950 as paged_module + + (query, key, value), public_kwargs = _paged_fp8_validation_inputs( + page_size=page_size, kv_cache_layout=kv_cache_layout + ) + output = torch.empty(query.shape, dtype=torch.bfloat16, device="cuda") + launch = paged_module.build_flash_attn_paged_fp8_module( + num_heads=6, + num_kv_heads=2, + head_dim=128, + dtype_str="fp8", + varlen=True, + cross_seqlen=True, + paged=True, + kv_cache_layout=kv_cache_layout, + page_size=page_size, + paged_bn128=page_size == 64, + ) + kwargs = { + name: public_kwargs[name] + for name in ("cu_seqlens_q", "cu_seqlens_kv", "block_table", "q_descale", "k_descale", "v_descale") + } + kwargs["block_table"] = kwargs["block_table"].flatten() + kwargs["block_table_stride"] = public_kwargs["block_table"].shape[1] + del kwargs[missing] + + def unexpected_dispatch(*args, **kwargs): + raise AssertionError("missing required metadata reached device dispatch") + + monkeypatch.setattr(paged_module, "_run_compiled", unexpected_dispatch) + monkeypatch.setattr(paged_module.flyc, "compile", unexpected_dispatch) + call = launch.compile if compile_only else launch + with pytest.raises(ValueError, match=missing): + call(query.flatten(), key, value, output.flatten(), 1, 17, seq_len_kv=256, **kwargs) + + +@_requires_gfx950 +@pytest.mark.parametrize("num_heads,num_kv_heads", [(0, 1), (3, 0), (3, 2)]) +def test_paged_fp8_builder_rejects_invalid_head_counts(num_heads, num_kv_heads): + from kernels.attention.flash_attn_fp8_paged_gfx950 import build_flash_attn_paged_fp8_module + + with pytest.raises(ValueError, match="head"): + build_flash_attn_paged_fp8_module( + num_heads=num_heads, + num_kv_heads=num_kv_heads, + head_dim=128, + dtype_str="fp8", + varlen=True, + cross_seqlen=True, + paged=True, + kv_cache_layout="vectorized", + ) + + @_requires_gfx950 def test_paged_fp8_bn128_launcher_rejects_unsupported_shapes(): - """Direct BN128 launches must preserve the single-sequence page-pair contract.""" + """Direct paired-ID loads require an even page bound and sufficient table stride.""" from kernels.attention.flash_attn_fp8_paged_gfx950 import build_flash_attn_paged_fp8_module launch = build_flash_attn_paged_fp8_module( @@ -5594,18 +6322,22 @@ def test_paged_fp8_bn128_launcher_rejects_unsupported_shapes(): tensor = torch.empty(1, dtype=FP8_DTYPE, device="cuda") output = torch.empty(1, dtype=torch.bfloat16, device="cuda") block_table = torch.zeros(3, dtype=torch.int32, device="cuda") + metadata = torch.zeros(2, dtype=torch.int32, device="cuda") + scale = torch.ones(1, dtype=torch.float32, device="cuda") + kwargs = dict(cu_seqlens_q=metadata, cu_seqlens_kv=metadata, q_descale=scale, k_descale=scale, v_descale=scale) - with pytest.raises(ValueError, match="batch_size=1"): + with pytest.raises(ValueError, match="too few entries"): launch( tensor, tensor, tensor, output, - batch_size=2, + batch_size=1, seq_len=64, seq_len_kv=128, block_table=block_table, - block_table_stride=2, + block_table_stride=1, + **kwargs, ) with pytest.raises(ValueError, match="positive even number"): launch( @@ -5618,6 +6350,7 @@ def test_paged_fp8_bn128_launcher_rejects_unsupported_shapes(): seq_len_kv=192, block_table=block_table, block_table_stride=3, + **kwargs, ) From 08ac7cbbb2da4122a27d2fdbfec9029e614e372d Mon Sep 17 00:00:00 2001 From: Xiake Sun Date: Sat, 12 Sep 2026 17:22:46 +0000 Subject: [PATCH 58/59] refactor(fmha): simplify paged FP8 scheduling and tests - Centralize KV layouts and simplify prefix/tail handling. - Stream V fragments for a common D192/V128 page schedule. - Deduplicate tests while preserving boundary and fallback coverage. Signed-off-by: Xiake Sun --- .../attention/flash_attn_fp8_paged_gfx950.py | 86 ++-- kernels/attention/flash_attn_utils.py | 135 ++++--- tests/kernels/test_flash_attn_fwd.py | 379 ++++++++---------- 3 files changed, 295 insertions(+), 305 deletions(-) diff --git a/kernels/attention/flash_attn_fp8_paged_gfx950.py b/kernels/attention/flash_attn_fp8_paged_gfx950.py index 8c2fe196b0..96093fe6b1 100644 --- a/kernels/attention/flash_attn_fp8_paged_gfx950.py +++ b/kernels/attention/flash_attn_fp8_paged_gfx950.py @@ -191,10 +191,15 @@ def flash_attn_paged_fp8_bn128_kernel( D_CHUNKS = traits.D_CHUNKS NPF = traits.NUM_PREFETCH_K BOUNDED_MAX = traits.HEAD_DIM == 128 and traits.DUALWAVE_SWP_LAZY_RESCALE - STREAM_PAGE16_V128 = traits.PAGE_SIZE == 16 and traits.HEAD_DIM == 192 and traits.HEAD_DIM_V == 128 + FRAGMENT_PV = traits.HEAD_DIM == 192 and traits.HEAD_DIM_V == 128 t0 = ctx.split_t0 t_end = ctx.split_t_end + def _load_v_first(buf_id): + if const_expr(FRAGMENT_PV): + return kv_lds_to_regs.load_v_fragment(buf_id, 0) + return kv_lds_to_regs.load_v(buf_id) + def _softmax_part(v_s, l_row, m_new): v_s = softmax_helper.sub_m(v_s, m_new) v_p = softmax_helper.exp2(v_s, 0, 16) @@ -203,11 +208,22 @@ def _softmax_part(v_s, l_row, m_new): v_p = gemm_helper.cast_p_fp8_direct(v_p) return v_p, l_row - def _subtile_tail(v_s, v_v, v_o, l_row, m_new): + def _subtile_tail(v_s, v_v, v_o, l_row, m_new, buf_id): v_p, l_row = _softmax_part(v_s, l_row, m_new) # Keep the post-MFMA accumulators in SSA. Pinning them after each # subtile lengthens the paged schedule without reducing registers. - v_o = gemm_helper.pv(v_p, v_v, v_o) + if const_expr(FRAGMENT_PV): + v_o = gemm_helper.preserve_accumulators(v_o) + next_v = v_v + # Keep K early on every page layout; retain only one current V + # fragment and overlap the next LDS read with this fragment's PV. + for dc in range_constexpr(D_CHUNKS): + current = next_v + if const_expr(dc + 1 < D_CHUNKS): + next_v = kv_lds_to_regs.load_v_fragment(buf_id, dc + 1) + v_o[dc] = gemm_helper._mfma_acc_fp8_wide(current, v_p, v_o[dc]) + else: + v_o = gemm_helper.pv(v_p, v_v, v_o) return v_o, l_row def _correct_o(v_o, m_row, l_row, m_tile): @@ -254,14 +270,8 @@ def _merge_tile_max(v_s_a, v_s_b): kv_gmem_to_lds.load_v((t0 + 2) * BN, (t0 + 2) % NPF, page_id=page_t2) kv_gmem_to_lds.load_v((t0 + 3) * BN, (t0 + 3) % NPF, page_id=page_t3) else: - next_v_a = kv_gmem_to_lds._load_v_fp8_vectorized_bankpad_source( - (t0 + 2) * BN, - page_id=page_t2, - ) - next_v_b = kv_gmem_to_lds._load_v_fp8_vectorized_bankpad_source( - (t0 + 3) * BN, - page_id=page_t3, - ) + next_v_a = kv_gmem_to_lds.load_v_source((t0 + 2) * BN, page_id=page_t2) + next_v_b = kv_gmem_to_lds.load_v_source((t0 + 3) * BN, page_id=page_t3) fx.rocdl.s_waitcnt(vmcnt=0, lgkmcnt=0, expcnt=0) rocdl.sched_barrier(0) rocdl.s_barrier() @@ -327,32 +337,19 @@ def _iterate(j, loop_args, do_mask, skip_max=False, initialize=False, mask_v=Tru next_v_b = loop_args[next_v_arg_idx + 1] v_k_a = kv_lds_to_regs.load_k(a_buf) - if const_expr(not STREAM_PAGE16_V128): - v_k_b = kv_lds_to_regs.load_k(b_buf) - v_v_a = kv_lds_to_regs.load_v(a_buf) + v_k_b = kv_lds_to_regs.load_k(b_buf) + v_v_a = _load_v_first(a_buf) page_f_a, page_f_b = ctx.load_page_id_pair((j + 4) * BN) kv_gmem_to_lds.load_k((j + 4) * BN, f_a_buf, page_id=page_f_a) kv_gmem_to_lds.load_k((j + 5) * BN, f_b_buf, page_id=page_f_b) v_s_a = gemm_helper.qk(v_k_a, q_wide) - if const_expr(STREAM_PAGE16_V128): - v_k_b = kv_lds_to_regs.load_k(b_buf) - kv_gmem_to_lds._store_v_fp8_vectorized_bankpad(next_v_a, nn_a_buf, (j + 2) * BN, mask_padding=mask_v) - v_f_a = kv_gmem_to_lds._load_v_fp8_vectorized_bankpad_source( - (j + 4) * BN, - page_id=page_f_a, - ) - v_f_b = kv_gmem_to_lds._load_v_fp8_vectorized_bankpad_source( - (j + 5) * BN, - page_id=page_f_b, - ) + kv_gmem_to_lds.store_v_source(next_v_a, nn_a_buf, (j + 2) * BN, mask_padding=mask_v) + v_f_a = kv_gmem_to_lds.load_v_source((j + 4) * BN, page_id=page_f_a) + v_f_b = kv_gmem_to_lds.load_v_source((j + 5) * BN, page_id=page_f_b) v_s_b = gemm_helper.qk(v_k_b, q_wide) - kv_gmem_to_lds._store_v_fp8_vectorized_bankpad(next_v_b, nn_b_buf, (j + 3) * BN, mask_padding=mask_v) - if const_expr(STREAM_PAGE16_V128): - # The current V slot is distinct from both next-V stores. - # Delay its fragment until K operands have been consumed. - v_v_a = kv_lds_to_regs.load_v(a_buf) + kv_gmem_to_lds.store_v_source(next_v_b, nn_b_buf, (j + 3) * BN, mask_padding=mask_v) if const_expr(do_mask): v_s_a, v_s_b = softmax_helper.causal_mask_pair_if_needed(v_s_a, v_s_b, j) m_new = m_row @@ -365,9 +362,9 @@ def _iterate(j, loop_args, do_mask, skip_max=False, initialize=False, mask_v=Tru v_o, m_new, l_row = _correct_o(v_o, m_row, l_row, m_tile) v_o = softmax_helper.anchor_v_o(v_o) - v_o, l_row = _subtile_tail(v_s_a, v_v_a, v_o, l_row, m_new) - v_v_b = kv_lds_to_regs.load_v(b_buf) - v_o, l_row = _subtile_tail(v_s_b, v_v_b, v_o, l_row, m_new) + v_o, l_row = _subtile_tail(v_s_a, v_v_a, v_o, l_row, m_new, a_buf) + v_v_b = _load_v_first(b_buf) + v_o, l_row = _subtile_tail(v_s_b, v_v_b, v_o, l_row, m_new, b_buf) m_row = m_new next_args = [m_row, l_row] + v_o + [nn_a_buf, v_f_a, v_f_b] @@ -412,22 +409,17 @@ def _iterate(j, loop_args, do_mask, skip_max=False, initialize=False, mask_v=Tru for j, loop_args in range(slow_prefix_end, fx.Int64(t_end), fx.Int64(2), init=slow_tail_state): next_args = _iterate(j, loop_args, True) loop_results = yield next_args - elif const_expr(traits.HEAD_DIM == 192 and traits.HEAD_DIM_V == 128): - # The prefix boundary is wave-uniform, not CTA-uniform. Both loops - # must keep the same pair order and one rendezvous per pair. - prefix_end = fx.Int64(ctx.q_start_pos_i32 + ctx.delta_i32) // (2 * BN) * 2 - prefix_end = fx.min(fx.Int64(t_end), fx.max(fx.Int64(t0), fx.min(prefix_end, v_prefix_end))) - for j, loop_args in range(fx.Int64(t0), prefix_end, fx.Int64(2), init=init_args): - next_args = _iterate(j, loop_args, False, mask_v=False) - loop_results = yield next_args - tail_init = loop_results - for j, loop_args in range(prefix_end, fx.Int64(t_end), fx.Int64(2), init=tail_init): - next_args = _iterate(j, loop_args, True) - loop_results = yield next_args else: - prefix_end = fx.min(fx.Int64(t_end), fx.max(fx.Int64(t0), v_prefix_end)) + mask_prefix = (traits.HEAD_DIM, traits.HEAD_DIM_V) != (192, 128) + prefix_end = v_prefix_end + if const_expr(not mask_prefix): + # This boundary is wave-uniform, not CTA-uniform. Both loops + # must keep the same pair order and one rendezvous per pair. + causal_prefix_end = fx.Int64(ctx.q_start_pos_i32 + ctx.delta_i32) // (2 * BN) * 2 + prefix_end = fx.min(causal_prefix_end, v_prefix_end) + prefix_end = fx.min(fx.Int64(t_end), fx.max(fx.Int64(t0), prefix_end)) for j, loop_args in range(fx.Int64(t0), prefix_end, fx.Int64(2), init=init_args): - next_args = _iterate(j, loop_args, True, mask_v=False) + next_args = _iterate(j, loop_args, mask_prefix, mask_v=False) loop_results = yield next_args tail_init = loop_results for j, loop_args in range(prefix_end, fx.Int64(t_end), fx.Int64(2), init=tail_init): diff --git a/kernels/attention/flash_attn_utils.py b/kernels/attention/flash_attn_utils.py index e2893a9c25..1629dda407 100644 --- a/kernels/attention/flash_attn_utils.py +++ b/kernels/attention/flash_attn_utils.py @@ -2163,6 +2163,28 @@ def cache_tag(self): self.CACHE_BUFFERED, ) + @property + def K_LDS_PAGE_GROUPED(self): + return self.PAGE_SIZE == 16 and self.HEAD_DIM_V == 192 + + @property + def V_LOAD_LAYOUT(self): + """Producer word order; the matching store and byte mask must agree.""" + if self.PAGE_SIZE == 1: + return "token_transpose" if (self.HEAD_DIM, self.HEAD_DIM_V) == (192, 128) else "token_words" + if self.PAGE_SIZE == 16 and self.HEAD_DIM == 192: + if self.HEAD_DIM_V == 192 or not self.CACHE_BUFFERED: + return "page_waves" + return "lane_groups" + + @property + def SHARED_PAGE_ID_LAYOUT(self): + if self.PAGE_SIZE == 1 and self.HEAD_DIM_V == 192: + return "tokens" + if self.PAGE_SIZE == 16 and self.CACHE_BUFFERED: + return self.V_LOAD_LAYOUT + return None + def _make_paged_dualwave_swp_fp8_traits( num_heads, @@ -4939,20 +4961,18 @@ def load_page_id(self, tile_start, *, uniform=True): def load_page_id_pair(self, tile_start): """Return reusable page IDs for a compute-tile pair, or defer to producers.""" - if const_expr(self.traits.PAGE_SIZE == 1 and self.traits.HEAD_DIM_V == 192): - token = fx.Int64(tile_start) + self.lane_in_warp - return self.load_page_id(token, uniform=False), self.load_page_id(token + 64, uniform=False) - if const_expr(self.traits.PAGE_SIZE == 16 and self.traits.CACHE_BUFFERED): - if const_expr(self.traits.HEAD_DIM_V == 192): - token = fx.Int64(tile_start) + (self.wave_id_uni % fx.Int64(4)) * fx.Int64(16) - return self.load_page_id(token), self.load_page_id(token + 64) - # V128 K and V use the same four lane-group pages. Share metadata - # loads across the two producers instead of looking them up twice. - token = fx.Int64(tile_start) + (self.lane_in_warp // fx.Int64(16)) * fx.Int64(16) - return self.load_page_id(token, uniform=False), self.load_page_id(token + 64, uniform=False) if const_expr(self.traits.PAGE_SIZE < self.traits.BLOCK_N): - # A compute tile spans several pages; its loaders gather IDs per lane. - return None, None + layout = self.traits.SHARED_PAGE_ID_LAYOUT + if const_expr(layout is None): + return None, None + if const_expr(layout == "tokens"): + token = fx.Int64(tile_start) + self.lane_in_warp + elif const_expr(layout == "page_waves"): + token = fx.Int64(tile_start) + (self.wave_id_uni % fx.Int64(4)) * fx.Int64(16) + else: + token = fx.Int64(tile_start) + (self.lane_in_warp // fx.Int64(16)) * fx.Int64(16) + uniform = layout == "page_waves" + return self.load_page_id(token, uniform=uniform), self.load_page_id(token + 64, uniform=uniform) if const_expr(not self.traits.PAIRED_PAGE_IDS): return self.load_page_id(tile_start), self.load_page_id(tile_start + self.traits.BLOCK_N) local_page = fx.Int32(fx.Int64(tile_start) // self.traits.PAGE_SIZE) @@ -5197,7 +5217,7 @@ def load_k(self, tile_start, buf_id, page_id=None): uses the low 32 lanes; paged tiles retain the physical D-group mapping. """ traits = self.traits - if const_expr(traits.PAGED and traits.PAGE_SIZE == 16 and traits.HEAD_DIM_V == 192): + if const_expr(traits.PAGED and traits.K_LDS_PAGE_GROUPED): self._load_k_page16(tile_start, buf_id, page_id=page_id) return eb = traits.ELEM_BYTES @@ -5238,11 +5258,11 @@ def _load_vectorized(): dst_byte = self.k_buf_base(buf_id) + oct_idx * traits.KV_VEC_SIZE dst = fx.slice(self.k_lds_i32_tiles, (None, fx.Uint32(dst_byte) // 16)) fx.generic_store(fx.get_iter(dst), fx.Vector(source)) - elif const_expr(traits.PAGE_SIZE == traits.BLOCK_N): - src_elem = _vec_k_src_elem(traits, d, self.wave_id_uni, self.lane_in_warp, self.kv_head_idx) else: oct_idx = _vec_k_dma_oct_idx(traits, d, self.wave_id_uni, self.lane_in_warp) - token = fx.Int64(tile_start) % traits.PAGE_SIZE + _sigma_k_tile_n(oct_idx % traits.BLOCK_N) + token = _sigma_k_tile_n(oct_idx % traits.BLOCK_N) + if const_expr(traits.PAGE_SIZE != traits.BLOCK_N): + token = fx.Int64(tile_start) % traits.PAGE_SIZE + token dim_group = oct_idx // traits.BLOCK_N src_elem = ( self.kv_head_idx * traits.HEAD_DIM * traits.PAGE_SIZE @@ -5423,11 +5443,11 @@ def _permute_v_fp8_vectorized(self, src_i32x4): def _stage_v_fp8_vectorized_bankpad(self, tile_start, buf_id, page_id=None, *, mask_padding=True): """Stage vectorized paged V in the bank-padded direct-FP8 layout.""" - if const_expr(self.traits.PAGE_SIZE == 16 and self.traits.HEAD_DIM_V == 192 and self.traits.CACHE_BUFFERED): + if const_expr(self.traits.K_LDS_PAGE_GROUPED and self.traits.CACHE_BUFFERED): self._stage_v_fp8_page16_segments(tile_start, buf_id, page_id=page_id, mask_padding=mask_padding) return - src_i32x4 = self._load_v_fp8_vectorized_bankpad_source(tile_start, page_id=page_id) - self._store_v_fp8_vectorized_bankpad(src_i32x4, buf_id, tile_start, mask_padding=mask_padding) + source = self.load_v_source(tile_start, page_id=page_id) + self.store_v_source(source, buf_id, tile_start, mask_padding=mask_padding) def _mask_v_fp8_group(self, source, token_start): """Zero inactive tokens before PV: a zero probability cannot mask a NaN V.""" @@ -5444,8 +5464,8 @@ def _mask_v_fp8_group(self, source, token_start): words.append(fx.Int32(fx.Uint32(packed[word]) & mask)) return Vec.from_elements(words, fx.Int32).ir_value() - def _stage_v_fp8_page16_segments(self, tile_start, buf_id, page_id=None, *, mask_padding=True): - """Consume each V192 segment before loading the next segment's words.""" + def _page16_v_source(self, tile_start, page_id): + """Coordinates and bounded resource for the wave-grouped V producer.""" traits = self.traits token = fx.Int64(tile_start) + (self.wave_id_uni % fx.Int64(4)) * fx.Int64(16) d_col = (self.wave_id // fx.Int64(4)) * fx.Int64(64) + self.lane_in_warp @@ -5453,6 +5473,12 @@ def _stage_v_fp8_page16_segments(self, tile_start, buf_id, page_id=None, *, mask page_id = self.load_page_id(token) src_div = self.make_page_view(self.v_base_iter, page_id, is_value=True) src_elem = self.kv_head_idx * traits.PAGE_SIZE * traits.HEAD_DIM_V + d_col * traits.KV_VEC_SIZE + return token, d_col, src_div, src_elem + + def _stage_v_fp8_page16_segments(self, tile_start, buf_id, page_id=None, *, mask_padding=True): + """Consume each V192 segment before loading the next segment's words.""" + traits = self.traits + token, _, src_div, src_elem = self._page16_v_source(tile_start, page_id) src_elem = (token < self.seqlen_kv_v).select(src_elem, fx.Int64(self.v_page_bytes)) prefix = self.buffer_load_fp8x16(src_div, src_elem) if const_expr(mask_padding): @@ -5473,8 +5499,8 @@ def _store_v_fp8_lds(self, data, byte_offset): dst = fx.slice(self.v_lds_i32_tiles, (None, fx.Uint32(byte_offset) // 16)) fx.generic_store(fx.get_iter(dst), data) - def _load_v_fp8_vectorized_bankpad_source(self, tile_start, page_id=None): - """Issue vectorized V loads for the paged FP8 LDS layout.""" + def load_v_source(self, tile_start, page_id=None): + """Prefetch a native V packet; store_v_source consumes its selected layout.""" traits = self.traits if const_expr(traits.PAGE_SIZE < traits.BLOCK_N): return self._load_v_fp8_small_pages_source(tile_start, page_id=page_id) @@ -5506,7 +5532,7 @@ def _load_v_fp8_small_pages_source(self, tile_start, page_id=None): traits = self.traits if const_expr(traits.PAGE_SIZE == 1): return self._load_v_fp8_page1_source(tile_start, page_id=page_id) - if const_expr(traits.HEAD_DIM == 128 or (traits.CACHE_BUFFERED and traits.HEAD_DIM_V == 128)): + if const_expr(traits.V_LOAD_LAYOUT == "lane_groups"): token = fx.Int64(tile_start) + (self.lane_in_warp // fx.Int64(16)) * fx.Int64(16) d_col = self.wave_id * fx.Int64(16) + self.lane_in_warp % fx.Int64(16) if page_id is None: @@ -5517,12 +5543,7 @@ def _load_v_fp8_small_pages_source(self, tile_start, page_id=None): + d_col * traits.KV_VEC_SIZE ) return self.global_load_fp8x16(self.v_base_iter, src_elem, token < self.seqlen_kv_v, is_value=True) - token = fx.Int64(tile_start) + (self.wave_id_uni % fx.Int64(4)) * fx.Int64(16) - d_col = (self.wave_id // fx.Int64(4)) * fx.Int64(64) + self.lane_in_warp - if page_id is None: - page_id = self.load_page_id(token) - src_div = self.make_page_view(self.v_base_iter, page_id, is_value=True) - src_elem = self.kv_head_idx * traits.PAGE_SIZE * traits.HEAD_DIM_V + d_col * traits.KV_VEC_SIZE + token, d_col, src_div, src_elem = self._page16_v_source(tile_start, page_id) valid = token < self.seqlen_kv_v src_elem = valid.select(src_elem, fx.Int64(self.v_page_bytes)) source = self.buffer_load_fp8x16(src_div, src_elem) @@ -5544,7 +5565,7 @@ def _load_v_fp8_page1_source(self, tile_start, page_id=None): byte_offset = page_id * self.v_page_bytes + self.kv_head_idx * traits.HEAD_DIM_V + d_base valid = token < self.seqlen_kv_v prefix = self.global_load_fp8x16(self.v_base_iter, byte_offset, valid, is_value=True) - stages = 4 if traits.HEAD_DIM == 192 and traits.HEAD_DIM_V == 128 else 2 + stages = 2 if traits.V_LOAD_LAYOUT == "token_words" else 4 source = _transpose_v_fp8_16x16(prefix, self.lane_in_warp, stages=stages) if const_expr(traits.FP8_PV_SEGMENTED): tail = self.global_load_fp8x16( @@ -5557,16 +5578,12 @@ def _load_v_fp8_page1_source(self, tile_start, page_id=None): source = source.shuffle(tail_source, [0, 1, 2, 3, 4, 5, 6, 7]) return source.ir_value() - def _store_v_fp8_vectorized_bankpad(self, src_i32x4, buf_id, tile_start, *, mask_padding=True): - """Permute prefetched V vectors into the paged FP8 LDS layout.""" - if const_expr(self.traits.PAGE_SIZE == 1 and (self.traits.HEAD_DIM == 128 or self.traits.HEAD_DIM_V == 192)): + def store_v_source(self, src_i32x4, buf_id, tile_start, *, mask_padding=True): + """Store a load_v_source packet in the common MFMA byte order.""" + if const_expr(self.traits.V_LOAD_LAYOUT == "token_words"): self._store_v_fp8_page1(src_i32x4, buf_id) return - page16_scatter = ( - self.traits.PAGE_SIZE == 16 - and self.traits.HEAD_DIM == 192 - and (self.traits.HEAD_DIM_V == 192 or not self.traits.CACHE_BUFFERED) - ) + page16_scatter = self.traits.V_LOAD_LAYOUT == "page_waves" if const_expr(mask_padding and self.traits.PAGE_SIZE != 1): group = self.wave_id_uni % 4 if page16_scatter else self.lane_in_warp // 16 @@ -5688,7 +5705,7 @@ def _read_vec_key(key): packs = [] for ws in range_constexpr(traits.HEAD_DIM // 64): dg0 = ws * 4 + self.lane_div_32 * 2 - if const_expr(traits.PAGE_SIZE == 16 and traits.HEAD_DIM_V == 192): + if const_expr(traits.PAGED and traits.K_LDS_PAGE_GROUPED): # [D-group/4, page, D-group%4, token%16, byte]. The # offset delta from page-64 layout is divisible by 256, # preserving the gfx950 LDS bank index. @@ -5728,27 +5745,33 @@ def load_v(self, buf_id): def _load_v_fp8_vectorized_bankpad(self, buf_id): """Read the D-major V tile into the FP8 MFMA A fragment.""" + packs = [[None] * self.traits.D_CHUNKS for _ in range(4)] + for dc in range_constexpr(self.traits.D_CHUNKS): + fragment = Vec(self.load_v_fragment(buf_id, dc)).bitcast(fx.Int64) + for ks in range_constexpr(4): + packs[ks][dc] = fragment[ks].ir_value() + return packs + + def load_v_fragment(self, buf_id, dc): + """Read one 32-channel V fragment in the retained MFMA byte order.""" traits = self.traits row_stride = traits.FP8_V_ROW_STRIDE v_tile_bytes = traits.HEAD_DIM_V * row_stride aligned_base = ((self.lds_vt_base_idx + fx.Int64(127)) // fx.Int64(128)) * fx.Int64(128) tile_base = aligned_base + fx.Int64(buf_id * v_tile_bytes) token_base = self.lane_div_32 * fx.Int64(32) - packs = [[None] * traits.D_CHUNKS for _ in range(4)] - for dc in range_constexpr(traits.D_CHUNKS): - d_row = fx.Int64(dc * traits.D_CHUNK) + self.lane_mod_32 - row_base = tile_base + d_row * fx.Int64(row_stride) - for half in range_constexpr(2): - token_offset = token_base + fx.Int64(half * 16) - if const_expr(traits.HEAD_DIM_V == 192): - token_offset = token_offset ^ (((d_row // fx.Int64(4)) % fx.Int64(4)) * fx.Int64(16)) - view = fx.slice( - self.v_lds_i32_tiles, (None, fx.Uint32(row_base + token_offset - self.lds_vt_base_idx) // 16) - ) - pair = fx.generic_load(fx.get_iter(view), dtype=fx.Int32, count=4).bitcast(fx.Int64) - packs[half * 2][dc] = pair[0].ir_value() - packs[half * 2 + 1][dc] = pair[1].ir_value() - return packs + d_row = fx.Int64(dc * traits.D_CHUNK) + self.lane_mod_32 + row_base = tile_base + d_row * fx.Int64(row_stride) + halves = [] + for half in range_constexpr(2): + token_offset = token_base + fx.Int64(half * 16) + if const_expr(traits.HEAD_DIM_V == 192): + token_offset = token_offset ^ (((d_row // fx.Int64(4)) % fx.Int64(4)) * fx.Int64(16)) + view = fx.slice( + self.v_lds_i32_tiles, (None, fx.Uint32(row_base + token_offset - self.lds_vt_base_idx) // 16) + ) + halves.append(fx.generic_load(fx.get_iter(view), dtype=fx.Int32, count=4)) + return halves[0].shuffle(halves[1], [0, 1, 2, 3, 4, 5, 6, 7]).ir_value() def _load_v_fp8_block(self, buf_id): traits = self.traits diff --git a/tests/kernels/test_flash_attn_fwd.py b/tests/kernels/test_flash_attn_fwd.py index eb924f6b79..ed00db16db 100644 --- a/tests/kernels/test_flash_attn_fwd.py +++ b/tests/kernels/test_flash_attn_fwd.py @@ -3832,16 +3832,16 @@ def _rand_lse(*shape, dtype, device="cuda"): return torch.randn(*shape, device=device, dtype=dtype) -@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) @pytest.mark.parametrize("causal", [False, True]) @pytest.mark.parametrize( - "B,S,H,Hkv,D", + "B,S,H,Hkv,D,dtype", [ - (2, 256, 8, 8, 128), # MHA - (2, 256, 8, 2, 128), # GQA - (1, 256, 4, 1, 128), # MQA - (2, 192, 4, 4, 64), # D=64 - (3, 128, 6, 3, 128), # GQA, odd batch + (2, 256, 8, 8, 128, torch.bfloat16), # MHA + (2, 256, 8, 2, 128, torch.float16), # GQA + (1, 256, 4, 1, 128, torch.bfloat16), # MQA + (2, 192, 4, 4, 64, torch.bfloat16), + (2, 192, 4, 4, 64, torch.float16), + (3, 128, 6, 3, 128, torch.bfloat16), # GQA, odd batch ], ) def test_lse_dense(dtype, causal, B, S, H, Hkv, D): @@ -3859,8 +3859,7 @@ def test_lse_dense(dtype, causal, B, S, H, Hkv, D): @_requires_gfx950 @pytest.mark.parametrize("causal", [False, True]) -@pytest.mark.parametrize("num_kv_splits", [2, 3]) -@pytest.mark.parametrize("D", [64, 128]) +@pytest.mark.parametrize("D,num_kv_splits", [(64, 2), (128, 3)]) def test_lse_splitk(causal, num_kv_splits, D): # Split-K requires seq_len >= 384, D in {64,128}, bf16/f16. B, S, H = 1, 512, 8 @@ -4559,6 +4558,11 @@ def test_return_lse_false_returns_only_out(): assert out.shape == q.shape +_PAGED_FP8_DIMS = [(128, 128), (192, 128), (192, 192)] +_PAGED_FP8_PHYSICAL_LAYOUTS = [(1, "linear"), (1, "linear3d"), (16, "vectorized"), (1024, "vectorized")] +_PAGED_FP8_LAYOUTS = [(64, "vectorized")] + _PAGED_FP8_PHYSICAL_LAYOUTS + + def _paged_fp8_cache_shapes(num_pages, num_kv_heads, head_dim, value_head_dim, page_size, kv_cache_layout): if kv_cache_layout == "vectorized": return ( @@ -4617,41 +4621,30 @@ def test_return_lse_rejects_fp8(): @_requires_gfx950 @pytest.mark.parametrize( - ("head_dim", "value_head_dim"), - [(128, 128), (192, 128), (192, 192)], - ids=["d128-v128", "d192-v128", "d192-v192"], -) -@pytest.mark.parametrize( - ("use_non_default_stream", "force_internal_copies"), - [(False, False), (True, False), (True, True)], - ids=["default", "side", "side-copies"], -) -@pytest.mark.parametrize( - ("query_lengths", "kv_lengths", "block_table_rows"), + ("head_dim", "value_head_dim", "query_lengths", "kv_lengths", "block_table_rows"), [ - ([64, 32], [128, 96], [[2, 0], [3, 1]]), - ([64], [128], [[1, 0]]), - ([64], [192], [[2, 0, 1]]), + (128, 128, [64, 32], [128, 96], [[2, 0], [3, 1]]), + (192, 192, [64], [128], [[1, 0]]), + (192, 128, [64], [192], [[2, 0, 1]]), ], ids=["ragged", "single-sequence-even-pages", "single-sequence-odd-pages"], ) -def test_paged_fp8_asymmetric_value_matches_torch( +def test_paged_fp8_page64_side_stream_copies( head_dim, value_head_dim, - use_non_default_stream, - force_internal_copies, query_lengths, kv_lengths, block_table_rows, ): + """Copy ordering for B=1/B>1 and paired/scalar page-ID schedules.""" _check_paged_fp8_matches_torch( - head_dim, - value_head_dim, - use_non_default_stream, - force_internal_copies, - query_lengths, - kv_lengths, - block_table_rows, + head_dim=head_dim, + value_head_dim=value_head_dim, + use_non_default_stream=True, + force_internal_copies=True, + query_lengths=query_lengths, + kv_lengths=kv_lengths, + block_table_rows=block_table_rows, ) @@ -4791,21 +4784,26 @@ def noncontiguous_copy(tensor): return actual -_PAGED_FP8_PHYSICAL_LAYOUTS = [(1, "linear"), (1, "linear3d"), (16, "vectorized"), (1024, "vectorized")] - - @_requires_gfx950 @pytest.mark.parametrize( - "page_size,kv_cache_layout,longest_kv", - [(64, "vectorized", 193)] + [(p, layout, 129) for p, layout in [(64, "vectorized")] + _PAGED_FP8_PHYSICAL_LAYOUTS], + "page_size,kv_cache_layout,longest_kv,num_query_heads,num_kv_heads,head_dim,value_head_dim,lazy_rescale", + [ + (1, "linear3d", 129, 1, 1, 128, 128, True), + (16, "vectorized", 129, 3, 1, 192, 128, True), + (64, "vectorized", 129, 6, 2, 192, 192, True), + (1024, "vectorized", 129, 6, 3, 128, 128, True), + (1, "linear", 129, 8, 8, 192, 192, True), + (64, "vectorized", 193, 24, 3, 192, 128, True), + (64, "vectorized", 129, 16, 1, 128, 128, True), + *[(64, "vectorized", 193, 16, 1, d, v, False) for d, v in _PAGED_FP8_DIMS], + ], ) -@pytest.mark.parametrize("head_dim,value_head_dim", [(128, 128), (192, 128), (192, 192)]) -@pytest.mark.parametrize("num_query_heads,num_kv_heads", [(1, 1), (3, 1), (6, 2), (6, 3), (8, 8), (24, 3), (16, 1)]) -@pytest.mark.parametrize("lazy_rescale", [False, True]) def test_paged_fp8_general_head_counts_and_ragged_lengths( page_size, kv_cache_layout, longest_kv, head_dim, value_head_dim, num_query_heads, num_kv_heads, lazy_rescale ): """MHA/GQA and masked rows must not depend on the Hq16/Hkv1 benchmark.""" + # Vary head ratios separately from the full loader-width matrix below. + # Keep eager paired-page coverage at every width; multiblock covers scalar IDs. query_lengths = [0, 65, 300, 17] kv_lengths = [33, 0, longest_kv, 65] counts = [(length + page_size - 1) // page_size for length in kv_lengths] @@ -4837,12 +4835,14 @@ def test_paged_fp8_general_head_counts_and_ragged_lengths( [(64, "vectorized", "paired"), (64, "vectorized", "scalar"), (1024, "vectorized", "native")] + [(p, layout, mode) for p, layout in _PAGED_FP8_PHYSICAL_LAYOUTS if p < 64 for mode in ("buffered", "wide")], ) -@pytest.mark.parametrize("head_dim,value_head_dim", [(128, 128), (192, 128), (192, 192)]) -@pytest.mark.parametrize("kv_length", [1, 2, 15, 16, 17, 63, 64, 65, 66, 127, 128, 129, 513, 1025]) +@pytest.mark.parametrize("head_dim,value_head_dim", _PAGED_FP8_DIMS) +@pytest.mark.parametrize("kv_length", [1, 16, 17, 63, 64, 65, 127, 128, 129, 513, 1025]) def test_paged_fp8_inactive_cache_values_do_not_contaminate_output( monkeypatch, page_size, kv_cache_layout, mode, head_dim, value_head_dim, kv_length ): """Only active K/V values matter, including empty requests and partial pages.""" + # Byte-residue combinations are exhaustive in test_paged_fp8_value_mask_is_byte_exact; + # this integration matrix targets token-group, tile, page-pair and pipeline boundaries. if mode == "wide": build = flash_attn_interface._build_paged_fp8 @@ -4947,8 +4947,8 @@ def launch(output: fx.Tensor): @_requires_gfx950 -@pytest.mark.parametrize("page_size,kv_cache_layout", [(64, "vectorized")] + _PAGED_FP8_PHYSICAL_LAYOUTS) -@pytest.mark.parametrize("head_dim,value_head_dim", [(128, 128), (192, 128), (192, 192)]) +@pytest.mark.parametrize("page_size,kv_cache_layout", _PAGED_FP8_LAYOUTS) +@pytest.mark.parametrize("head_dim,value_head_dim", _PAGED_FP8_DIMS) def test_paged_fp8_shared_pages_and_strided_metadata(page_size, kv_cache_layout, head_dim, value_head_dim): """Page IDs may be shared; cumulative lengths need not be contiguous views.""" kv_lengths = [256, 65] @@ -4971,15 +4971,15 @@ def test_paged_fp8_shared_pages_and_strided_metadata(page_size, kv_cache_layout, @_requires_gfx950 -@pytest.mark.parametrize("page_size,kv_cache_layout", [(64, "vectorized")] + _PAGED_FP8_PHYSICAL_LAYOUTS) -@pytest.mark.parametrize("head_dim,value_head_dim", [(128, 128), (192, 128), (192, 192)]) +@pytest.mark.parametrize("page_size,kv_cache_layout", _PAGED_FP8_LAYOUTS) +@pytest.mark.parametrize("head_dim,value_head_dim", _PAGED_FP8_DIMS) @pytest.mark.parametrize("strided_metadata", [False, True]) def test_paged_fp8_graph_replay_updates_lengths_and_ignores_padding( page_size, kv_cache_layout, head_dim, value_head_dim, strided_metadata ): """One captured graph honors changing device lengths and poisoned unused IDs.""" torch.manual_seed(981) - query_length, max_query, max_kv = 17, 512, 512 + query_length, max_query, max_kv = 17, 512, max(512, page_size + 1) num_heads, num_kv_heads = 6, 2 num_pages = (max_kv + page_size - 1) // page_size query, q_scale = quantize_per_tensor_fp8(torch.randn(query_length, num_heads, head_dim, device="cuda") * 0.2) @@ -5033,7 +5033,7 @@ def call(): with torch.cuda.graph(graph): captured = call() assert captured is output - for length in (65, 129, 0, 33): + for length in (65, 129, max_kv, 0, 33): update_metadata(length) output.fill_(float("nan")) graph.replay() @@ -5046,14 +5046,14 @@ def call(): @_requires_gfx950 -@pytest.mark.parametrize("page_size,kv_cache_layout", [(64, "vectorized")] + _PAGED_FP8_PHYSICAL_LAYOUTS) -@pytest.mark.parametrize("head_dim,value_head_dim", [(128, 128), (192, 128), (192, 192)]) +@pytest.mark.parametrize("page_size,kv_cache_layout", _PAGED_FP8_LAYOUTS) +@pytest.mark.parametrize("head_dim,value_head_dim", _PAGED_FP8_DIMS) @pytest.mark.parametrize("max_query", [256, 512]) def test_paged_fp8_maxima_are_upper_bounds(page_size, kv_cache_layout, head_dim, value_head_dim, max_query): """Loose launch bounds must not replace device-side lengths or overwrite O.""" query_length, kv_length = 17, 65 max_kv = 512 - num_heads, num_kv_heads = 3, 1 + num_heads, num_kv_heads = 6, 2 num_pages = (max_kv + page_size - 1) // page_size # Backing allocations cover the loose bounds: a broken specialization # fails the canary/reference checks without an illegal GPU memory access. @@ -5105,13 +5105,7 @@ def test_paged_fp8_maxima_are_upper_bounds(page_size, kv_cache_layout, head_dim, torch.testing.assert_close(actual, torch.ones_like(actual), rtol=2.0e-2, atol=2.0e-2) -@_requires_gfx950 -@pytest.mark.parametrize("page_size,kv_cache_layout", [(64, "vectorized")] + _PAGED_FP8_PHYSICAL_LAYOUTS) -@pytest.mark.parametrize("head_dim,value_head_dim", [(128, 128), (192, 128), (192, 192)]) -@pytest.mark.parametrize("query_lengths,kv_lengths", [([0, 0], [65, 0]), ([17, 0, 65], [0, 0, 0]), ([], [])]) -@pytest.mark.parametrize("preallocated", [False, True]) -@pytest.mark.parametrize("kv_max_mode", ["exact", "upper_bound", "infer"]) -def test_paged_fp8_empty_attention_skips_kernel( +def _check_paged_fp8_empty_attention( monkeypatch, page_size, kv_cache_layout, @@ -5170,6 +5164,34 @@ def unexpected_build(**kwargs): torch.testing.assert_close(actual, expected, rtol=0, atol=0) +@_requires_gfx950 +@pytest.mark.parametrize("query_lengths,kv_lengths", [([0, 0], [65, 0]), ([17, 0, 65], [0, 0, 0]), ([], [])]) +@pytest.mark.parametrize("preallocated", [False, True]) +@pytest.mark.parametrize("kv_max_mode", ["exact", "upper_bound", "infer"]) +def test_paged_fp8_empty_attention_skips_kernel(monkeypatch, query_lengths, kv_lengths, preallocated, kv_max_mode): + """Empty-result control flow does not depend on physical cache geometry.""" + _check_paged_fp8_empty_attention( + monkeypatch, 64, "vectorized", 192, 128, query_lengths, kv_lengths, preallocated, kv_max_mode + ) + + +@_requires_gfx950 +@pytest.mark.parametrize( + "page_size,kv_cache_layout,head_dim,value_head_dim", + [ + (page, layout, head_dim, value_dim) + for page, layout in _PAGED_FP8_LAYOUTS + for head_dim, value_dim in _PAGED_FP8_DIMS + if (page, layout, head_dim, value_dim) != (64, "vectorized", 192, 128) + ], +) +def test_paged_fp8_empty_output_layouts(monkeypatch, page_size, kv_cache_layout, head_dim, value_head_dim): + """All native cache/output shapes support the same empty-KV fast path.""" + _check_paged_fp8_empty_attention( + monkeypatch, page_size, kv_cache_layout, head_dim, value_head_dim, [17, 0, 65], [0, 0, 0], True, "upper_bound" + ) + + def _paged_fp8_validation_inputs(num_heads=6, num_kv_heads=2, page_size=64, kv_cache_layout="vectorized"): query = torch.zeros((17, num_heads, 128), dtype=FP8_DTYPE, device="cuda") num_pages = (256 + page_size - 1) // page_size @@ -5200,18 +5222,22 @@ def _unexpected_paged_fp8_build(**kwargs): @_requires_gfx950 -@pytest.mark.parametrize("page_size,kv_cache_layout", [(64, "vectorized")] + _PAGED_FP8_PHYSICAL_LAYOUTS) @pytest.mark.parametrize( - "scale_shape,scale_stride", [((), None), ((1,), None), ((1, 1), None), ((1,), (0,)), ((1,), (2,))] + "page_size,kv_cache_layout,scale_shape,scale_stride", + [ + (64, "vectorized", shape, stride) + for shape, stride in [((), None), ((1,), None), ((1, 1), None), ((1,), (0,)), ((1,), (2,))] + ] + + [(page, layout, (), None) for page, layout in _PAGED_FP8_PHYSICAL_LAYOUTS], ) def test_paged_fp8_single_value_descale_shapes(page_size, kv_cache_layout, scale_shape, scale_stride): + """Normalize every singleton form once, and accept scalars on every layout.""" (query, key, value), kwargs = _paged_fp8_validation_inputs(page_size=page_size, kv_cache_layout=kv_cache_layout) value.fill_(1.0) for name in ("q_descale", "k_descale", "v_descale"): kwargs[name] = torch.full(scale_shape, 0.5, dtype=torch.float32, device="cuda") if scale_stride is not None: kwargs[name] = kwargs[name].as_strided(scale_shape, scale_stride) - flash_attn_interface._build_paged_fp8.cache_clear() actual = flydsl_flash_attn_func(query, key, value, **kwargs) torch.cuda.synchronize() torch.testing.assert_close(actual, torch.full_like(actual, 0.5), rtol=0, atol=0) @@ -5303,7 +5329,7 @@ def test_paged_fp8_rejects_negative_launch_bounds(monkeypatch, name): @_requires_gfx950 -@pytest.mark.parametrize("page_size,kv_cache_layout", [(64, "vectorized")] + _PAGED_FP8_PHYSICAL_LAYOUTS) +@pytest.mark.parametrize("page_size,kv_cache_layout", _PAGED_FP8_LAYOUTS) @pytest.mark.parametrize("preallocated", [False, True]) def test_paged_fp8_empty_cache_on_stream_and_graph(monkeypatch, page_size, kv_cache_layout, preallocated): (query, key, value), kwargs = _paged_fp8_validation_inputs(page_size=page_size, kv_cache_layout=kv_cache_layout) @@ -5420,7 +5446,7 @@ def launch(table: fx.Tensor, output: fx.Tensor): @_requires_gfx950 @pytest.mark.parametrize("page_size,kv_cache_layout", [(1, "linear"), (1, "linear3d"), (16, "vectorized")]) -@pytest.mark.parametrize("head_dim,value_head_dim", [(128, 128), (192, 128), (192, 192)]) +@pytest.mark.parametrize("head_dim,value_head_dim", _PAGED_FP8_DIMS) def test_paged_fp8_cache_buffer_matches_wide_fallback( monkeypatch, page_size, kv_cache_layout, head_dim, value_head_dim ): @@ -5435,6 +5461,7 @@ def test_paged_fp8_cache_buffer_matches_wide_fallback( query_lengths=[65, 17], kv_lengths=kv_lengths, block_table_rows=[pages[: counts[0]], pages[counts[0] :] + [0] * (counts[0] - counts[1])], + num_query_heads=6, num_kv_heads=2, page_size=page_size, kv_cache_layout=kv_cache_layout, @@ -5765,9 +5792,12 @@ def test_paged_fp8_rejects_inconsistent_page_metadata(violation): @_requires_gfx950 -@pytest.mark.parametrize("page_size,kv_cache_layout", _PAGED_FP8_PHYSICAL_LAYOUTS) -@pytest.mark.parametrize("head_dim,value_head_dim", [(128, 128), (192, 128), (192, 192)]) +@pytest.mark.parametrize( + "page_size,kv_cache_layout,head_dim,value_head_dim", + [(1, "linear", 128, 128), (1, "linear3d", 192, 128), (16, "vectorized", 192, 192), (1024, "vectorized", 192, 128)], +) def test_paged_fp8_physical_pages_side_stream_copies(page_size, kv_cache_layout, head_dim, value_head_dim): + """Exercise each native cache rank/layout; page-64 copies cover all widths.""" kv_lengths = [page_size + 65, page_size + 17, 33] counts = [(length + page_size - 1) // page_size for length in kv_lengths] pages = list(reversed(range(sum(counts)))) @@ -5790,15 +5820,16 @@ def test_paged_fp8_physical_pages_side_stream_copies(page_size, kv_cache_layout, @_requires_gfx950 -@pytest.mark.parametrize("page_size,kv_cache_layout", _PAGED_FP8_PHYSICAL_LAYOUTS) -@pytest.mark.parametrize("head_dim,value_head_dim", [(128, 128), (192, 128), (192, 192)]) -@pytest.mark.parametrize("num_kv_heads", [1, 2, 4]) -@pytest.mark.parametrize("lazy_rescale", [True, False]) -def test_paged_fp8_physical_page_size_ragged_matches_torch( - page_size, kv_cache_layout, head_dim, value_head_dim, num_kv_heads, lazy_rescale -): +@pytest.mark.parametrize( + "page_size,kv_cache_layout,head_dim,value_head_dim,lazy_rescale", + [(p, layout, d, v, True) for p, layout in _PAGED_FP8_LAYOUTS for d, v in _PAGED_FP8_DIMS] + + [(64, "vectorized", d, v, False) for d, v in _PAGED_FP8_DIMS] + + [(1, "linear3d", 192, 192, False), (16, "vectorized", 192, 128, False), (1024, "vectorized", 128, 128, False)], +) +def test_paged_fp8_ragged_multiblock_matches_torch(page_size, kv_cache_layout, head_dim, value_head_dim, lazy_rescale): + """All loader/width pairs run multiblock GQA; eager cases span pages and widths.""" query_lengths = [300, 257, 33, 7, 1] - kv_lengths = [max(513, 2 * page_size + 1), max(385, page_size + 1), 65, 16, 1] + kv_lengths = [max(513, 2 * page_size + 1), max(385, page_size + 1), 129, 16, 1] counts = [(length + page_size - 1) // page_size for length in kv_lengths] physical_pages = list(reversed(range(sum(counts)))) rows = [] @@ -5814,7 +5845,8 @@ def test_paged_fp8_physical_page_size_ragged_matches_torch( query_lengths=query_lengths, kv_lengths=kv_lengths, block_table_rows=rows, - num_kv_heads=num_kv_heads, + num_query_heads=6, + num_kv_heads=2, lazy_rescale=lazy_rescale, page_size=page_size, kv_cache_layout=kv_cache_layout, @@ -5822,10 +5854,11 @@ def test_paged_fp8_physical_page_size_ragged_matches_torch( @_requires_gfx950 -@pytest.mark.parametrize("head_dim,value_head_dim", [(128, 128), (192, 128), (192, 192)]) +@pytest.mark.parametrize("head_dim,value_head_dim", _PAGED_FP8_DIMS) @pytest.mark.parametrize( "page_size,kv_cache_layout,kv_length", - [ + [(64, "vectorized", 63)] + + [ (size, layout, length) for size, layout in _PAGED_FP8_PHYSICAL_LAYOUTS for length in sorted({63, 64, 65, max(1, size - 1), size, size + 1}) @@ -5843,43 +5876,13 @@ def test_paged_fp8_physical_page_boundary_matches_torch( query_lengths=[min(17, kv_length)], kv_lengths=[kv_length], block_table_rows=[list(reversed(range(pages)))], + num_query_heads=6, + num_kv_heads=2, page_size=page_size, kv_cache_layout=kv_cache_layout, ) -@_requires_gfx950 -@pytest.mark.parametrize("head_dim,value_head_dim", [(128, 128), (192, 128), (192, 192)]) -@pytest.mark.parametrize("num_kv_heads", [1, 2, 4]) -@pytest.mark.parametrize("lazy_rescale", [True, False]) -def test_paged_fp8_odd_page_ragged_multiblock_matches_torch(head_dim, value_head_dim, num_kv_heads, lazy_rescale): - _check_paged_fp8_matches_torch( - head_dim=head_dim, - value_head_dim=value_head_dim, - use_non_default_stream=True, - force_internal_copies=False, - query_lengths=[300, 257, 33], - kv_lengths=[513, 385, 129], - block_table_rows=[list(reversed(range(9))), list(reversed(range(9, 16))) + [0] * 2, [18, 16, 17] + [0] * 6], - num_kv_heads=num_kv_heads, - lazy_rescale=lazy_rescale, - ) - - -@_requires_gfx950 -@pytest.mark.parametrize("head_dim,value_head_dim", [(128, 128), (192, 128), (192, 192)]) -def test_paged_fp8_single_partial_page_matches_torch(head_dim, value_head_dim): - _check_paged_fp8_matches_torch( - head_dim=head_dim, - value_head_dim=value_head_dim, - use_non_default_stream=False, - force_internal_copies=False, - query_lengths=[17], - kv_lengths=[63], - block_table_rows=[[0]], - ) - - @_requires_gfx950 @pytest.mark.parametrize("value_head_dim", [128, 192]) @pytest.mark.parametrize("num_kv_heads", [2, 4]) @@ -5902,7 +5905,7 @@ def test_paged_fp8_d192_bn128_multiple_kv_heads(value_head_dim, num_kv_heads): @_requires_gfx950 -@pytest.mark.parametrize("head_dim,value_head_dim", [(128, 128), (192, 128), (192, 192)]) +@pytest.mark.parametrize("head_dim,value_head_dim", _PAGED_FP8_DIMS) def test_paged_fp8_bn128_ragged_multiblock_matches_torch(head_dim, value_head_dim): """The multi-batch BN128 path handles distinct causal offsets and inactive q-blocks.""" _check_paged_fp8_matches_torch( @@ -6008,7 +6011,7 @@ def test_paged_fp8_d192_batch_interleave_group(batch_size, head_dims, expected): @pytest.mark.parametrize("batch_size", [1, 2, 3, 4, 5, 7, 8, 16, 33]) -@pytest.mark.parametrize("head_dims", [(128, 128), (192, 128), (192, 192)]) +@pytest.mark.parametrize("head_dims", _PAGED_FP8_DIMS) def test_paged_fp8_paired_batch_interleave_group(batch_size, head_dims): choose = flash_attn_interface._paged_fp8_batch_interleave_group if head_dims == (128, 128): @@ -6057,7 +6060,7 @@ def test_paged_fp8_bn128_batch2_ragged_multiblock_matches_torch(): @_requires_gfx950 @pytest.mark.parametrize( ("head_dim", "value_head_dim"), - [(128, 128), (192, 128), (192, 192)], + _PAGED_FP8_DIMS, ids=["d128-v128", "d192-v128", "d192-v192"], ) def test_paged_fp8_long_context_random_pages_matches_torch(head_dim, value_head_dim): @@ -6111,79 +6114,15 @@ def test_paged_fp8_long_context_random_pages_matches_torch(head_dim, value_head_ @_requires_gfx950 +@pytest.mark.parametrize("head_dim,value_head_dim", _PAGED_FP8_DIMS) @pytest.mark.parametrize( - ("head_dim", "value_head_dim"), - [(128, 128), (192, 128), (192, 192)], - ids=["d128-v128", "d192-v128", "d192-v192"], + "page_size,kv_cache_layout,kv_length", + [(page, layout, 128) for page, layout in _PAGED_FP8_LAYOUTS] + [(64, "vectorized", 192)], ) -@pytest.mark.parametrize("page_size,kv_cache_layout", [(64, "vectorized")] + _PAGED_FP8_PHYSICAL_LAYOUTS) -def test_paged_fp8_graph_replay_matches_torch(head_dim, value_head_dim, page_size, kv_cache_layout): - """The optimized paged variants preserve caller-owned graph output.""" - torch.manual_seed(29) - query_length = 63 - kv_length = max(128, page_size + 1) - num_pages = (kv_length + page_size - 1) // page_size - query, query_descale = quantize_per_tensor_fp8(torch.randn(query_length, 16, head_dim, device="cuda") * 0.2) - key_shape, value_shape = _paged_fp8_cache_shapes(num_pages, 1, head_dim, value_head_dim, page_size, kv_cache_layout) - key, key_descale = quantize_per_tensor_fp8(torch.randn(key_shape, device="cuda") * 0.2) - value, value_descale = quantize_per_tensor_fp8(torch.randn(value_shape, device="cuda") * 0.2) - output = torch.empty(query_length, 16, value_head_dim, device="cuda", dtype=torch.bfloat16) - block_table = torch.arange(num_pages - 1, -1, -1, device="cuda", dtype=torch.int32).reshape(1, -1) - q_indptr = torch.tensor([0, query_length], device="cuda", dtype=torch.int32) - kv_indptr = torch.tensor([0, kv_length], device="cuda", dtype=torch.int32) - seqlen_k = torch.tensor([kv_length], device="cuda", dtype=torch.int32) - kwargs = dict( - causal=True, - num_kv_heads=1, - cu_seqlens_q=q_indptr, - cu_seqlens_kv=kv_indptr, - max_seqlen_q=query_length, - max_seqlen_kv=kv_length, - cross_seqlen=True, - block_table=block_table, - seqlen_k=seqlen_k, - kv_cache_layout=kv_cache_layout, - q_descale=query_descale, - k_descale=key_descale, - v_descale=value_descale, - out=output, - ) - - warmup_stream = torch.cuda.Stream() - with torch.cuda.stream(warmup_stream): - flydsl_flash_attn_func(query, key, value, **kwargs) - torch.cuda.current_stream().wait_stream(warmup_stream) - - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - captured = flydsl_flash_attn_func(query, key, value, **kwargs) - graph.replay() - torch.cuda.synchronize() - - expected = _paged_fp8_torch_reference( - query, - key, - value, - block_table, - [query_length], - [kv_length], - (query_descale, key_descale, value_descale), - kv_cache_layout=kv_cache_layout, - ) - - assert captured.data_ptr() == output.data_ptr() - assert bool(torch.isfinite(output).all().item()) - torch.testing.assert_close(output, expected, rtol=2.0e-2, atol=2.0e-2) - - -@_requires_gfx950 -@pytest.mark.parametrize("head_dim,value_head_dim", [(128, 128), (192, 128), (192, 192)]) -@pytest.mark.parametrize("kv_length", [128, 192]) -@pytest.mark.parametrize("page_size,kv_cache_layout", [(64, "vectorized")] + _PAGED_FP8_PHYSICAL_LAYOUTS) def test_paged_fp8_explicit_compile_matches_torch( monkeypatch, head_dim, value_head_dim, kv_length, page_size, kv_cache_layout ): - """Explicit compilation honors the launch arguments for both paged schedules.""" + """Compile every layout/width; only page 64 switches paired/scalar IDs at 192.""" build = flash_attn_interface._build_paged_fp8 def compile_launcher(**kwargs): @@ -6199,6 +6138,8 @@ def compile_launcher(**kwargs): query_lengths=[64, 32], kv_lengths=[kv_length, kv_length - 32], block_table_rows=[list(reversed(range(pages))), list(range(pages, 2 * pages))], + num_query_heads=6, + num_kv_heads=2, page_size=page_size, kv_cache_layout=kv_cache_layout, ) @@ -6206,14 +6147,20 @@ def compile_launcher(**kwargs): @_requires_gfx950 @pytest.mark.parametrize( - "page_size,kv_cache_layout,longest_kv", - [(64, "vectorized", 193)] + [(p, layout, 129) for p, layout in [(64, "vectorized")] + _PAGED_FP8_PHYSICAL_LAYOUTS], + "page_size,kv_cache_layout,longest_kv,head_dim,value_head_dim", + [ + (64, "vectorized", 193, 192, 128), + (64, "vectorized", 129, 128, 128), + (1, "linear", 129, 192, 192), + (1, "linear3d", 129, 192, 128), + (16, "vectorized", 129, 128, 128), + (1024, "vectorized", 129, 192, 192), + ], ) -@pytest.mark.parametrize("head_dim,value_head_dim", [(128, 128), (192, 128), (192, 192)]) def test_paged_fp8_interleave_is_not_limited_to_benchmark_heads( monkeypatch, page_size, kv_cache_layout, longest_kv, head_dim, value_head_dim ): - """The grid permutation is valid for generic D128 and non-benchmark GQA.""" + """Force non-benchmark GQA interleave on each layout and both page-64 ID paths.""" build = flash_attn_interface._build_paged_fp8 def interleaved(**kwargs): @@ -6240,15 +6187,16 @@ def interleaved(**kwargs): @_requires_gfx950 -@pytest.mark.parametrize("page_size,kv_cache_layout", [(64, "vectorized")] + _PAGED_FP8_PHYSICAL_LAYOUTS) @pytest.mark.parametrize("compile_only", [False, True]) @pytest.mark.parametrize( "missing", ["cu_seqlens_q", "cu_seqlens_kv", "block_table", "block_table_stride", "q_descale", "k_descale", "v_descale"], ) -def test_paged_fp8_direct_launcher_requires_metadata(monkeypatch, page_size, kv_cache_layout, compile_only, missing): +def test_paged_fp8_direct_launcher_requires_metadata(monkeypatch, compile_only, missing): + """The common host guard precedes page-specific compilation and launch.""" from kernels.attention import flash_attn_fp8_paged_gfx950 as paged_module + page_size, kv_cache_layout = 64, "vectorized" (query, key, value), public_kwargs = _paged_fp8_validation_inputs( page_size=page_size, kv_cache_layout=kv_cache_layout ) @@ -6455,12 +6403,13 @@ def test_paged_legacy_dtype_accepts_host_only_seqlen_metadata(dtype): @_requires_gfx950 @pytest.mark.large_shape -@pytest.mark.parametrize("head_dim,value_head_dim", [(128, 128), (192, 128), (192, 192)]) -@pytest.mark.parametrize("page_size,kv_cache_layout", [(64, "vectorized")] + _PAGED_FP8_PHYSICAL_LAYOUTS) +@pytest.mark.parametrize("head_dim,value_head_dim", _PAGED_FP8_DIMS) +@pytest.mark.parametrize("page_size,kv_cache_layout", _PAGED_FP8_LAYOUTS) def test_paged_fp8_cache_offsets_above_4gib(head_dim, value_head_dim, page_size, kv_cache_layout): """Physical K and V page rebasing remains 64-bit beyond 4 GiB.""" - key_page_bytes = page_size * head_dim - value_page_bytes = page_size * value_head_dim + num_heads, num_kv_heads = 6, 2 + key_page_bytes = page_size * num_kv_heads * head_dim + value_page_bytes = page_size * num_kv_heads * value_head_dim high_page = math.ceil(2**32 / min(key_page_bytes, value_page_bytes)) num_pages = high_page + 1 cache_bytes = num_pages * (key_page_bytes + value_page_bytes) @@ -6473,8 +6422,10 @@ def test_paged_fp8_cache_offsets_above_4gib(head_dim, value_head_dim, page_size, f">4 GiB K/V offsets, got {free_bytes / 2**30:.1f} GiB" ) - query, query_descale = quantize_per_tensor_fp8(torch.randn(1, 16, head_dim, device="cuda") * 0.2) - key_shape, value_shape = _paged_fp8_cache_shapes(num_pages, 1, head_dim, value_head_dim, page_size, kv_cache_layout) + query, query_descale = quantize_per_tensor_fp8(torch.randn(1, num_heads, head_dim, device="cuda") * 0.2) + key_shape, value_shape = _paged_fp8_cache_shapes( + num_pages, num_kv_heads, head_dim, value_head_dim, page_size, kv_cache_layout + ) key = torch.empty(key_shape, dtype=FP8_DTYPE, device="cuda") value = torch.empty(value_shape, dtype=FP8_DTYPE, device="cuda") key_page, key_descale = quantize_per_tensor_fp8(torch.randn_like(key[high_page], dtype=torch.float32) * 0.2) @@ -6490,7 +6441,7 @@ def test_paged_fp8_cache_offsets_above_4gib(head_dim, value_head_dim, page_size, key, value, causal=True, - num_kv_heads=1, + num_kv_heads=num_kv_heads, cu_seqlens_q=indptr, cu_seqlens_kv=indptr, max_seqlen_q=1, @@ -6506,9 +6457,16 @@ def test_paged_fp8_cache_offsets_above_4gib(head_dim, value_head_dim, page_size, torch.cuda.synchronize() first_value = ( - value_page[0, 0, :, 0] if kv_cache_layout == "vectorized" else value_page.reshape(-1, value_head_dim)[0] + value_page[:, 0, :, 0] + if kv_cache_layout == "vectorized" + else value_page.reshape(-1, num_kv_heads, value_head_dim)[0] + ) + expected = ( + (first_value.float() * value_descale) + .repeat_interleave(num_heads // num_kv_heads, dim=0) + .unsqueeze(0) + .to(torch.bfloat16) ) - expected = (first_value.float() * value_descale).expand(16, -1).unsqueeze(0).to(torch.bfloat16) assert high_page * key_page_bytes >= 2**32 assert high_page * value_page_bytes >= 2**32 assert bool(torch.isfinite(actual).all().item()) @@ -6516,7 +6474,7 @@ def test_paged_fp8_cache_offsets_above_4gib(head_dim, value_head_dim, page_size, @_requires_gfx950 -@pytest.mark.parametrize("H", [8, 16, 32, 64]) +@pytest.mark.parametrize("H", [8, 64]) def test_xcd_swizzle_is_bit_identical(H): """The head-slow remap must not change a single bit of the output. @@ -6959,9 +6917,11 @@ def test_fp8_rescale_threshold_drops_past_the_long_sequence_bound(): assert set(f(s) for s in (1, 1024, 4096, 4097, 8192, 131072)) == {6.0, 4.0} -@pytest.mark.parametrize("head_dim,value_head_dim", [(128, 128), (192, 128), (192, 192)]) +@pytest.mark.parametrize("head_dim,value_head_dim", _PAGED_FP8_DIMS) @pytest.mark.parametrize("page_size", [1, 16, 64, 1024]) def test_paged_fp8_traits_match_supported_value_segments(head_dim, value_head_dim, page_size): + from dataclasses import replace + from kernels.attention.flash_attn_utils import LDS_BYTES_GFX950, _make_paged_dualwave_swp_fp8_traits traits = _make_paged_dualwave_swp_fp8_traits( @@ -6978,6 +6938,19 @@ def test_paged_fp8_traits_match_supported_value_segments(head_dim, value_head_di assert traits.FP8_PV_SEGMENTED == (value_head_dim == 192) assert traits.VT_BF16_TOTAL * 2 >= traits.NUM_PREFETCH_K * value_head_dim * traits.FP8_V_ROW_STRIDE assert traits.LDS_KV_TOTAL_SIZE + traits.VT_BF16_TOTAL * 2 <= LDS_BYTES_GFX950 + # Source layout, shared page IDs and LDS consumers must select one policy. + for buffered in (False, True): + layout = replace(traits, CACHE_BUFFERED=buffered) + assert layout.K_LDS_PAGE_GROUPED == (page_size == 16 and value_head_dim == 192) + if page_size == 1: + expected = "token_transpose" if (head_dim, value_head_dim) == (192, 128) else "token_words" + shared = "tokens" if value_head_dim == 192 else None + elif page_size == 16: + expected = "page_waves" if head_dim == 192 and (value_head_dim == 192 or not buffered) else "lane_groups" + shared = expected if buffered else None + else: + expected, shared = "lane_groups", None + assert (layout.V_LOAD_LAYOUT, layout.SHARED_PAGE_ID_LAYOUT) == (expected, shared) @_requires_gfx950 @@ -7109,10 +7082,10 @@ def test_fp8_split_kv(causal, num_kv_splits, head_dim, head_dim_v): @_requires_gfx950 @pytest.mark.parametrize("causal", [False, True]) -@pytest.mark.parametrize("batch", (1, 2, 3, 4)) +@pytest.mark.parametrize("batch", (2, 3, 4)) def test_fp8_split_kv_batched(causal, batch): """Split-KV over batches: the grid is B * num_kv_splits deep, so the batch is - what decides whether splitting still fills the GPU.""" + what decides whether splitting still fills the GPU. B=1 is in test_fp8_split_kv.""" _assert_fp8_shape(causal, batch=batch, seq_len=8192, num_kv_splits=4) @@ -7126,9 +7099,12 @@ def test_fp8_split_kv_cross_length(causal, seq_len, seqlen_kv, num_kv_splits): @_requires_gfx950 @pytest.mark.parametrize("causal", [False, True]) -@pytest.mark.parametrize("seq_len,num_kv_splits", list(zip(FP8_SPLITKV_SEQLENS, FP8_SPLITKV_SPLITS))) +@pytest.mark.parametrize( + "seq_len,num_kv_splits", + [(seq, splits) for seq, splits in zip(FP8_SPLITKV_SEQLENS, FP8_SPLITKV_SPLITS) if seq != 8192], +) def test_fp8_split_kv_long_sequence(causal, seq_len, num_kv_splits): - """Split count scaled with the KV length, 4k/2 through 32k/16.""" + """Scaled split counts; the 8K/4-split case is already in test_fp8_split_kv.""" _assert_fp8_shape(causal, batch=1, seq_len=seq_len, num_kv_splits=num_kv_splits) @@ -7244,7 +7220,6 @@ def test_fp8_dense_ragged_seq_lens(seq_len): (8, 32, 2048, 2048, True, 256), (16, 32, 1024, 1024, True, 256), (32, 32, 512, 512, True, 256), - (1, 8, 2048, 2048, True, 128), (1, 8, 4096, 4096, True, 256), (1, 8, 4096, 16384, True, 256), (2, 8, 4096, 4096, True, 256), From 8a662b0b19d0b564fe01fc230640d7a3c5602bf1 Mon Sep 17 00:00:00 2001 From: Xiake Sun Date: Mon, 14 Sep 2026 03:25:23 +0000 Subject: [PATCH 59/59] increase CI test timeout to 70 min Signed-off-by: Xiake Sun --- .github/workflows/flydsl.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/flydsl.yaml b/.github/workflows/flydsl.yaml index 519aed2bee..33b149f620 100644 --- a/.github/workflows/flydsl.yaml +++ b/.github/workflows/flydsl.yaml @@ -617,7 +617,7 @@ jobs: docker exec flydsl_test bash -c "python3 -c 'import torch; import triton; assert tuple(map(int, triton.__version__.split(\"+\", 1)[0].split(\".\")[:3])) >= (3, 6, 0), triton.__version__; print(\"Using image Triton\", triton.__version__, \"with PyTorch\", torch.__version__)'" - name: Run tests - timeout-minutes: 60 + timeout-minutes: 70 run: | docker exec -e ROCM_PATH flydsl_test bash -c 'export PYTHONPATH=/tmp/aiter:${PYTHONPATH:-} && export AITER_REPO=/tmp/aiter && cd /flydsl-test && RUN_TESTS_FULL=1 bash scripts/run_tests.sh'